3333from .rank import compute_rank
3434from .utils import (check_fname , logger , verbose , check_version , _time_mask ,
3535 warn , copy_function_doc_to_method_doc , _pl ,
36- _undo_scaling_cov , _scaled_array )
36+ _undo_scaling_cov , _scaled_array , _validate_type )
3737from . import viz
3838
3939from .fixes import BaseEstimator , EmpiricalCovariance , _logdet
@@ -1303,36 +1303,6 @@ def _get_ch_whitener(A, pca, ch_type, rank):
13031303 return eig , eigvec , mask
13041304
13051305
1306- def _get_whitener (noise_cov , info = None , ch_names = None , rank = None ,
1307- pca = False , scalings = None , prepared = False ):
1308- #
1309- # Handle noise cov
1310- #
1311- if not prepared :
1312- noise_cov = prepare_noise_cov (noise_cov , info , ch_names , rank ,
1313- scalings )
1314- n_chan = len (noise_cov ['eig' ])
1315-
1316- # Omit the zeroes due to projection
1317- eig = noise_cov ['eig' ].copy ()
1318- nzero = (eig > 0 )
1319- eig [~ nzero ] = 0. # get rid of numerical noise (negative) ones
1320- n_nzero = np .sum (nzero )
1321-
1322- whitener = np .zeros ((n_chan , 1 ), dtype = np .float )
1323- whitener [nzero , 0 ] = 1.0 / np .sqrt (eig [nzero ])
1324- # Rows of eigvec are the eigenvectors
1325- whitener = whitener * noise_cov ['eigvec' ] # C ** -0.5
1326- colorer = np .sqrt (eig ) * noise_cov ['eigvec' ].T # C ** 0.5
1327- if pca :
1328- whitener = whitener [nzero ]
1329- colorer = colorer [:, nzero ]
1330- logger .info (' Created the whitener using a noise covariance matrix '
1331- 'with rank %d (%d small eigenvalues omitted)'
1332- % (n_nzero , noise_cov ['dim' ] - np .sum (nzero )))
1333- return whitener , colorer , noise_cov , n_nzero
1334-
1335-
13361306@verbose
13371307def prepare_noise_cov (noise_cov , info , ch_names = None , rank = None ,
13381308 scalings = None , verbose = None ):
@@ -1366,8 +1336,18 @@ def prepare_noise_cov(noise_cov, info, ch_names=None, rank=None,
13661336 and parameters updated.
13671337 """
13681338 # reorder C and info to match ch_names order
1339+ noise_cov_idx = list ()
1340+ missing = list ()
13691341 ch_names = info ['ch_names' ] if ch_names is None else ch_names
1370- noise_cov_idx = [noise_cov .ch_names .index (c ) for c in ch_names ]
1342+ for c in ch_names :
1343+ # this could be try/except ValueError, but it is not the preferred way
1344+ if c in noise_cov .ch_names :
1345+ noise_cov_idx .append (noise_cov .ch_names .index (c ))
1346+ else :
1347+ missing .append (c )
1348+ if len (missing ):
1349+ raise RuntimeError ('Not all channels present in noise covariance:\n %s'
1350+ % missing )
13711351 if not noise_cov ['diag' ]:
13721352 C = noise_cov .data [np .ix_ (noise_cov_idx , noise_cov_idx )]
13731353 else :
@@ -1647,16 +1627,18 @@ def _regularized_covariance(data, reg=None, method_params=None, info=None,
16471627
16481628
16491629@verbose
1650- def compute_whitener (noise_cov , info , picks = None , rank = None , scalings = None ,
1651- return_rank = False , pca = False , verbose = None ):
1630+ def compute_whitener (noise_cov , info = None , picks = None , rank = None ,
1631+ scalings = None , return_rank = False , pca = False ,
1632+ return_colorer = False , verbose = None ):
16521633 """Compute whitening matrix.
16531634
16541635 Parameters
16551636 ----------
16561637 noise_cov : Covariance
16571638 The noise covariance.
1658- info : dict
1659- The measurement info.
1639+ info : dict | None
1640+ The measurement info. Can be None if `noise_cov` has already been
1641+ prepared with :func:`prepare_noise_cov`.
16601642 %(picks_good_data_noref)s
16611643 %(rank_None)s
16621644
@@ -1669,45 +1651,84 @@ def compute_whitener(noise_cov, info, picks=None, rank=None, scalings=None,
16691651 If True, return the rank used to compute the whitener.
16701652
16711653 .. versionadded:: 0.15
1672- pca : bool
1673- If True, return the rank-reduced whitener will be returned.
1654+ pca : bool | str
1655+ Space to project the data into. Options:
1656+
1657+ :data:`python:True`
1658+ Whitener will be shape (n_nonzero, n_channels).
1659+ ``'white'``
1660+ Whitener will be shape (n_channels, n_channels), potentially rank
1661+ deficient, and have the first ``n_channels - n_nonzero`` rows and
1662+ columns set to zero.
1663+ :data:`python:False` (default)
1664+ Whitener will be shape (n_channels, n_channels), potentially rank
1665+ deficient, and rotated back to the space of the original data.
16741666
16751667 .. versionadded:: 0.18
1668+ return_colorer : bool
1669+ If True, return the colorer as well.
16761670 %(verbose)s
16771671
16781672 Returns
16791673 -------
16801674 W : ndarray, shape (n_channels, n_channels) or (n_nonzero, n_channels)
1681- The whitening matrix, backprojected or not based on whether `pca` is
1682- False or True, respectively.
1675+ The whitening matrix.
16831676 ch_names : list
16841677 The channel names.
16851678 rank : int
16861679 Rank reduction of the whitener. Returned only if return_rank is True.
1687- """
1688- picks = _picks_to_idx (info , picks , with_ref_meg = False )
1680+ colorer : ndarray, shape (n_channels, n_channels) or (n_channels, n_nonzero)
1681+ The coloring matrix.
1682+ """ # noqa: E501
1683+ _validate_type (pca , (str , bool ), 'space' )
1684+ _valid_pcas = (True , 'white' , False )
1685+ if pca not in _valid_pcas :
1686+ raise ValueError ('space must be one of %s, got %s'
1687+ % (_valid_pcas , pca ))
1688+ if info is None :
1689+ if 'eig' not in noise_cov :
1690+ raise ValueError ('info can only be None if the noise cov has '
1691+ 'already been prepared with prepare_noise_cov' )
1692+ ch_names = deepcopy (noise_cov ['names' ])
1693+ else :
1694+ picks = _picks_to_idx (info , picks , with_ref_meg = False )
1695+ ch_names = [info ['ch_names' ][k ] for k in picks ]
1696+ del picks
1697+ noise_cov = prepare_noise_cov (
1698+ noise_cov , info , ch_names , rank , scalings )
16891699
1690- ch_names = [info ['ch_names' ][k ] for k in picks ]
1700+ n_chan = len (ch_names )
1701+ assert n_chan == len (noise_cov ['eig' ])
16911702
1692- # XXX this relies on pick_channels, which does not respect order,
1693- # so this could create problems if users have reordered their data
1694- noise_cov = pick_channels_cov (noise_cov , include = ch_names , exclude = [])
1695- if len (noise_cov ['data' ]) != len (ch_names ):
1696- missing = list (set (ch_names ) - set (noise_cov ['names' ]))
1697- raise RuntimeError ('Not all channels present in noise covariance:\n %s'
1698- % missing )
1703+ # Omit the zeroes due to projection
1704+ eig = noise_cov ['eig' ].copy ()
1705+ nzero = (eig > 0 )
1706+ eig [~ nzero ] = 0. # get rid of numerical noise (negative) ones
16991707
1700- W , _ , noise_cov , n_nzero = _get_whitener (
1701- noise_cov , info , ch_names , rank , pca = pca , scalings = scalings )
1708+ W = np .zeros ((n_chan , 1 ), dtype = np .float )
1709+ W [nzero , 0 ] = 1.0 / np .sqrt (eig [nzero ])
1710+ # Rows of eigvec are the eigenvectors
1711+ W = W * noise_cov ['eigvec' ] # C ** -0.5
1712+ C = np .sqrt (eig ) * noise_cov ['eigvec' ].T # C ** 0.5
1713+ n_nzero = nzero .sum ()
1714+ logger .info (' Created the whitener using a noise covariance matrix '
1715+ 'with rank %d (%d small eigenvalues omitted)'
1716+ % (n_nzero , noise_cov ['dim' ] - n_nzero ))
17021717
1703- # Do the back projection
1704- if not pca :
1718+ # Do the requested projection
1719+ if pca is True :
1720+ W = W [nzero ]
1721+ C = C [:, nzero ]
1722+ elif pca is False :
17051723 W = np .dot (noise_cov ['eigvec' ].T , W )
1706- else :
1707- assert W .shape [0 ] == n_nzero
1724+ C = np .dot (C , noise_cov ['eigvec' ])
1725+
1726+ # Triage return
17081727 out = W , ch_names
17091728 if return_rank :
17101729 out += (n_nzero ,)
1730+ if return_colorer :
1731+ out += (C ,)
17111732 return out
17121733
17131734
0 commit comments