<s> from __future__ import print_function <EOL> from copy import deepcopy <EOL> from math import sqrt , ceil <EOL> import numpy as np <EOL> from scipy import linalg <EOL> from . mxne_debiasing import compute_bias <EOL> from . . utils import logger , verbose , sum_squared , warn <EOL> from . . time_frequency . stft import stft_norm2 , stft , istft <EOL> from . . externals . six . moves import xrange as range <EOL> def groups_norm2 ( A , n_orient ) : <EOL> """<STR_LIT>""" <EOL> n_positions = A . shape [ <NUM_LIT:0> ] // n_orient <EOL> return np . sum ( np . power ( A , <NUM_LIT:2> , A ) . reshape ( n_positions , - <NUM_LIT:1> ) , axis = <NUM_LIT:1> ) <EOL> def norm_l2inf ( A , n_orient , copy = True ) : <EOL> """<STR_LIT>""" <EOL> if A . size == <NUM_LIT:0> : <EOL> return <NUM_LIT:0.0> <EOL> if copy : <EOL> A = A . copy ( ) <EOL> return sqrt ( np . max ( groups_norm2 ( A , n_orient ) ) ) <EOL> def norm_l21 ( A , n_orient , copy = True ) : <EOL> """<STR_LIT>""" <EOL> if A . size == <NUM_LIT:0> : <EOL> return <NUM_LIT:0.0> <EOL> if copy : <EOL> A = A . copy ( ) <EOL> return np . sum ( np . sqrt ( groups_norm2 ( A , n_orient ) ) ) <EOL> def prox_l21 ( Y , alpha , n_orient , shape = None , is_stft = False ) : <EOL> """<STR_LIT>""" <EOL> if len ( Y ) == <NUM_LIT:0> : <EOL> return np . zeros_like ( Y ) , np . zeros ( ( <NUM_LIT:0> , ) , dtype = np . bool ) <EOL> if shape is not None : <EOL> shape_init = Y . shape <EOL> Y = Y . reshape ( * shape ) <EOL> n_positions = Y . shape [ <NUM_LIT:0> ] // n_orient <EOL> if is_stft : <EOL> rows_norm = np . sqrt ( stft_norm2 ( Y ) . reshape ( n_positions , - <NUM_LIT:1> ) . sum ( axis = <NUM_LIT:1> ) ) <EOL> else : <EOL> rows_norm = np . sqrt ( ( Y * Y . conj ( ) ) . real . reshape ( n_positions , <EOL> - <NUM_LIT:1> ) . sum ( axis = <NUM_LIT:1> ) ) <EOL> shrink = np . maximum ( <NUM_LIT:1.0> - alpha / np . maximum ( rows_norm , alpha ) , <NUM_LIT:0.0> ) <EOL> active_set = shrink > <NUM_LIT:0.0> <EOL> if n_orient > <NUM_LIT:1> : <EOL> active_set = np . tile ( active_set [ : , None ] , [ <NUM_LIT:1> , n_orient ] ) . ravel ( ) <EOL> shrink = np . tile ( shrink [ : , None ] , [ <NUM_LIT:1> , n_orient ] ) . ravel ( ) <EOL> Y = Y [ active_set ] <EOL> if shape is None : <EOL> Y *= shrink [ active_set ] [ : , np . newaxis ] <EOL> else : <EOL> Y *= shrink [ active_set ] [ : , np . newaxis , np . newaxis ] <EOL> Y = Y . reshape ( - <NUM_LIT:1> , * shape_init [ <NUM_LIT:1> : ] ) <EOL> return Y , active_set <EOL> def prox_l1 ( Y , alpha , n_orient ) : <EOL> """<STR_LIT>""" <EOL> n_positions = Y . shape [ <NUM_LIT:0> ] // n_orient <EOL> norms = np . sqrt ( ( Y * Y . conj ( ) ) . real . T . reshape ( - <NUM_LIT:1> , n_orient ) . sum ( axis = <NUM_LIT:1> ) ) <EOL> shrink = np . maximum ( <NUM_LIT:1.0> - alpha / np . maximum ( norms , alpha ) , <NUM_LIT:0.0> ) <EOL> shrink = shrink . reshape ( - <NUM_LIT:1> , n_positions ) . T <EOL> active_set = np . any ( shrink > <NUM_LIT:0.0> , axis = <NUM_LIT:1> ) <EOL> shrink = shrink [ active_set ] <EOL> if n_orient > <NUM_LIT:1> : <EOL> active_set = np . tile ( active_set [ : , None ] , [ <NUM_LIT:1> , n_orient ] ) . ravel ( ) <EOL> Y = Y [ active_set ] <EOL> if len ( Y ) > <NUM_LIT:0> : <EOL> for o in range ( n_orient ) : <EOL> Y [ o : : n_orient ] *= shrink <EOL> return Y , active_set <EOL> def dgap_l21 ( M , G , X , active_set , alpha , n_orient ) : <EOL> """<STR_LIT>""" <EOL> GX = np . dot ( G [ : , active_set ] , X ) <EOL> R = M - GX <EOL> penalty = norm_l21 ( X , n_orient , copy = True ) <EOL> nR2 = sum_squared ( R ) <EOL> pobj = <NUM_LIT:0.5> * nR2 + alpha * penalty <EOL> dual_norm = norm_l2inf ( np . dot ( G . T , R ) , n_orient , copy = False ) <EOL> scaling = alpha / dual_norm <EOL> scaling = min ( scaling , <NUM_LIT:1.0> ) <EOL> dobj = <NUM_LIT:0.5> * ( scaling ** <NUM_LIT:2> ) * nR2 + scaling * np . sum ( R * GX ) <EOL> gap = pobj - dobj <EOL> return gap , pobj , dobj , R <EOL> @ verbose <EOL> def _mixed_norm_solver_prox ( M , G , alpha , lipschitz_constant , maxit = <NUM_LIT:200> , <EOL> tol = <NUM_LIT> , verbose = None , init = None , n_orient = <NUM_LIT:1> ) : <EOL> """<STR_LIT>""" <EOL> n_sensors , n_times = M . shape <EOL> n_sensors , n_sources = G . shape <EOL> if n_sources < n_sensors : <EOL> gram = np . dot ( G . T , G ) <EOL> GTM = np . dot ( G . T , M ) <EOL> else : <EOL> gram = None <EOL> if init is None : <EOL> X = <NUM_LIT:0.0> <EOL> R = M . copy ( ) <EOL> if gram is not None : <EOL> R = np . dot ( G . T , R ) <EOL> else : <EOL> X = init <EOL> if gram is None : <EOL> R = M - np . dot ( G , X ) <EOL> else : <EOL> R = GTM - np . dot ( gram , X ) <EOL> t = <NUM_LIT:1.0> <EOL> Y = np . zeros ( ( n_sources , n_times ) ) <EOL> E = [ ] <EOL> active_set = np . ones ( n_sources , dtype = np . bool ) <EOL> for i in range ( maxit ) : <EOL> X0 , active_set_0 = X , active_set <EOL> if gram is None : <EOL> Y += np . dot ( G . T , R ) / lipschitz_constant <EOL> else : <EOL> Y += R / lipschitz_constant <EOL> X , active_set = prox_l21 ( Y , alpha / lipschitz_constant , n_orient ) <EOL> t0 = t <EOL> t = <NUM_LIT:0.5> * ( <NUM_LIT:1.0> + sqrt ( <NUM_LIT:1.0> + <NUM_LIT> * t ** <NUM_LIT:2> ) ) <EOL> Y . fill ( <NUM_LIT:0.0> ) <EOL> dt = ( ( t0 - <NUM_LIT:1.0> ) / t ) <EOL> Y [ active_set ] = ( <NUM_LIT:1.0> + dt ) * X <EOL> Y [ active_set_0 ] -= dt * X0 <EOL> Y_as = active_set_0 | active_set <EOL> if gram is None : <EOL> R = M - np . dot ( G [ : , Y_as ] , Y [ Y_as ] ) <EOL> else : <EOL> R = GTM - np . dot ( gram [ : , Y_as ] , Y [ Y_as ] ) <EOL> gap , pobj , dobj , _ = dgap_l21 ( M , G , X , active_set , alpha , n_orient ) <EOL> E . append ( pobj ) <EOL> logger . debug ( "<STR_LIT>" % ( pobj , gap ) ) <EOL> if gap < tol : <EOL> logger . debug ( '<STR_LIT>' % ( gap , tol ) ) <EOL> break <EOL> return X , active_set , E <EOL> @ verbose <EOL> def _mixed_norm_solver_cd ( M , G , alpha , lipschitz_constant , maxit = <NUM_LIT> , <EOL> tol = <NUM_LIT> , verbose = None , init = None , n_orient = <NUM_LIT:1> ) : <EOL> """<STR_LIT>""" <EOL> from sklearn . linear_model . coordinate_descent import MultiTaskLasso <EOL> n_sensors , n_times = M . shape <EOL> n_sensors , n_sources = G . shape <EOL> if init is not None : <EOL> init = init . T <EOL> clf = MultiTaskLasso ( alpha = alpha / len ( M ) , tol = tol , normalize = False , <EOL> fit_intercept = False , max_iter = maxit , <EOL> warm_start = True ) <EOL> clf . coef_ = init <EOL> clf . fit ( G , M ) <EOL> X = clf . coef_ . T <EOL> active_set = np . any ( X , axis = <NUM_LIT:1> ) <EOL> X = X [ active_set ] <EOL> gap , pobj , dobj , _ = dgap_l21 ( M , G , X , active_set , alpha , n_orient ) <EOL> return X , active_set , pobj <EOL> @ verbose <EOL> def _mixed_norm_solver_bcd ( M , G , alpha , lipschitz_constant , maxit = <NUM_LIT:200> , <EOL> tol = <NUM_LIT> , verbose = None , init = None , n_orient = <NUM_LIT:1> ) : <EOL> """<STR_LIT>""" <EOL> G = np . asfortranarray ( G ) <EOL> n_sensors , n_times = M . shape <EOL> n_sensors , n_sources = G . shape <EOL> n_positions = n_sources // n_orient <EOL> if init is None : <EOL> X = np . zeros ( ( n_sources , n_times ) ) <EOL> R = M . copy ( ) <EOL> else : <EOL> X = init <EOL> R = M - np . dot ( G , X ) <EOL> E = [ ] <EOL> active_set = np . zeros ( n_sources , dtype = np . bool ) <EOL> alpha_lc = alpha / lipschitz_constant <EOL> for i in range ( maxit ) : <EOL> for j in range ( n_positions ) : <EOL> idx = slice ( j * n_orient , ( j + <NUM_LIT:1> ) * n_orient ) <EOL> G_j = G [ : , idx ] <EOL> X_j = X [ idx ] <EOL> X_j_new = np . dot ( G_j . T , R ) / lipschitz_constant [ j ] <EOL> was_non_zero = np . any ( X_j ) <EOL> if was_non_zero : <EOL> R += np . dot ( G_j , X_j ) <EOL> X_j_new += X_j <EOL> block_norm = linalg . norm ( X_j_new , '<STR_LIT>' ) <EOL> if block_norm <= alpha_lc [ j ] : <EOL> X_j . fill ( <NUM_LIT:0.> ) <EOL> active_set [ idx ] = False <EOL> else : <EOL> shrink = np . maximum ( <NUM_LIT:1.0> - alpha_lc [ j ] / block_norm , <NUM_LIT:0.0> ) <EOL> X_j_new *= shrink <EOL> R -= np . dot ( G_j , X_j_new ) <EOL> X_j [ : ] = X_j_new <EOL> active_set [ idx ] = True <EOL> gap , pobj , dobj , _ = dgap_l21 ( M , G , X [ active_set ] , active_set , alpha , <EOL> n_orient ) <EOL> E . append ( pobj ) <EOL> logger . debug ( "<STR_LIT>" % ( <EOL> i + <NUM_LIT:1> , pobj , gap , np . sum ( active_set ) / n_orient ) ) <EOL> if gap < tol : <EOL> logger . debug ( '<STR_LIT>' % ( gap , tol ) ) <EOL> break <EOL> X = X [ active_set ] <EOL> return X , active_set , E <EOL> @ verbose <EOL> def mixed_norm_solver ( M , G , alpha , maxit = <NUM_LIT> , tol = <NUM_LIT> , verbose = None , <EOL> active_set_size = <NUM_LIT:50> , debias = True , n_orient = <NUM_LIT:1> , <EOL> solver = '<STR_LIT>' ) : <EOL> """<STR_LIT>""" <EOL> n_dipoles = G . shape [ <NUM_LIT:1> ] <EOL> n_positions = n_dipoles // n_orient <EOL> n_sensors , n_times = M . shape <EOL> alpha_max = norm_l2inf ( np . dot ( G . T , M ) , n_orient , copy = False ) <EOL> logger . info ( "<STR_LIT>" % alpha_max ) <EOL> alpha = float ( alpha ) <EOL> has_sklearn = True <EOL> try : <EOL> from sklearn . linear_model . coordinate_descent import MultiTaskLasso <EOL> except ImportError : <EOL> has_sklearn = False <EOL> if solver == '<STR_LIT>' : <EOL> if has_sklearn and ( n_orient == <NUM_LIT:1> ) : <EOL> solver = '<STR_LIT>' <EOL> else : <EOL> solver = '<STR_LIT>' <EOL> if solver == '<STR_LIT>' : <EOL> if n_orient == <NUM_LIT:1> and not has_sklearn : <EOL> warn ( '<STR_LIT>' <EOL> '<STR_LIT>' ) <EOL> solver = '<STR_LIT>' <EOL> if n_orient > <NUM_LIT:1> : <EOL> warn ( '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' ) <EOL> solver = '<STR_LIT>' <EOL> if solver == '<STR_LIT>' : <EOL> logger . info ( "<STR_LIT>" ) <EOL> l21_solver = _mixed_norm_solver_cd <EOL> lc = None <EOL> elif solver == '<STR_LIT>' : <EOL> logger . info ( "<STR_LIT>" ) <EOL> l21_solver = _mixed_norm_solver_bcd <EOL> G = np . asfortranarray ( G ) <EOL> if n_orient == <NUM_LIT:1> : <EOL> lc = np . sum ( G * G , axis = <NUM_LIT:0> ) <EOL> else : <EOL> lc = np . empty ( n_positions ) <EOL> for j in range ( n_positions ) : <EOL> G_tmp = G [ : , ( j * n_orient ) : ( ( j + <NUM_LIT:1> ) * n_orient ) ] <EOL> lc [ j ] = linalg . norm ( np . dot ( G_tmp . T , G_tmp ) , ord = <NUM_LIT:2> ) <EOL> else : <EOL> logger . info ( "<STR_LIT>" ) <EOL> l21_solver = _mixed_norm_solver_prox <EOL> lc = <NUM_LIT> * linalg . norm ( G , ord = <NUM_LIT:2> ) ** <NUM_LIT:2> <EOL> if active_set_size is not None : <EOL> E = list ( ) <EOL> X_init = None <EOL> active_set = np . zeros ( n_dipoles , dtype = np . bool ) <EOL> idx_large_corr = np . argsort ( groups_norm2 ( np . dot ( G . T , M ) , n_orient ) ) <EOL> new_active_idx = idx_large_corr [ - active_set_size : ] <EOL> if n_orient > <NUM_LIT:1> : <EOL> new_active_idx = ( n_orient * new_active_idx [ : , None ] + <EOL> np . arange ( n_orient ) [ None , : ] ) . ravel ( ) <EOL> active_set [ new_active_idx ] = True <EOL> as_size = np . sum ( active_set ) <EOL> for k in range ( maxit ) : <EOL> if solver == '<STR_LIT>' : <EOL> lc_tmp = lc [ active_set [ : : n_orient ] ] <EOL> elif solver == '<STR_LIT>' : <EOL> lc_tmp = None <EOL> else : <EOL> lc_tmp = <NUM_LIT> * linalg . norm ( G [ : , active_set ] , ord = <NUM_LIT:2> ) ** <NUM_LIT:2> <EOL> X , as_ , _ = l21_solver ( M , G [ : , active_set ] , alpha , lc_tmp , <EOL> maxit = maxit , tol = tol , init = X_init , <EOL> n_orient = n_orient ) <EOL> active_set [ active_set ] = as_ . copy ( ) <EOL> idx_old_active_set = np . where ( active_set ) [ <NUM_LIT:0> ] <EOL> gap , pobj , dobj , R = dgap_l21 ( M , G , X , active_set , alpha , <EOL> n_orient ) <EOL> E . append ( pobj ) <EOL> logger . info ( "<STR_LIT>" <EOL> "<STR_LIT>" % ( <EOL> k + <NUM_LIT:1> , pobj , gap , as_size // n_orient , <EOL> np . sum ( active_set ) // n_orient ) ) <EOL> if gap < tol : <EOL> logger . info ( '<STR_LIT>' <EOL> % ( gap , tol ) ) <EOL> break <EOL> if k < ( maxit - <NUM_LIT:1> ) : <EOL> idx_large_corr = np . argsort ( groups_norm2 ( np . dot ( G . T , R ) , <EOL> n_orient ) ) <EOL> new_active_idx = idx_large_corr [ - active_set_size : ] <EOL> if n_orient > <NUM_LIT:1> : <EOL> new_active_idx = ( n_orient * new_active_idx [ : , None ] + <EOL> np . arange ( n_orient ) [ None , : ] ) <EOL> new_active_idx = new_active_idx . ravel ( ) <EOL> active_set [ new_active_idx ] = True <EOL> idx_active_set = np . where ( active_set ) [ <NUM_LIT:0> ] <EOL> as_size = np . sum ( active_set ) <EOL> X_init = np . zeros ( ( as_size , n_times ) , dtype = X . dtype ) <EOL> idx = np . searchsorted ( idx_active_set , idx_old_active_set ) <EOL> X_init [ idx ] = X <EOL> else : <EOL> warn ( '<STR_LIT>' % ( gap , tol ) ) <EOL> else : <EOL> X , active_set , E = l21_solver ( M , G , alpha , lc , maxit = maxit , <EOL> tol = tol , n_orient = n_orient , init = None ) <EOL> if np . any ( active_set ) and debias : <EOL> bias = compute_bias ( M , G [ : , active_set ] , X , n_orient = n_orient ) <EOL> X *= bias [ : , np . newaxis ] <EOL> logger . info ( '<STR_LIT>' % ( np . sum ( active_set ) // n_orient ) ) <EOL> return X , active_set , E <EOL> @ verbose <EOL> def iterative_mixed_norm_solver ( M , G , alpha , n_mxne_iter , maxit = <NUM_LIT> , <EOL> tol = <NUM_LIT> , verbose = None , active_set_size = <NUM_LIT:50> , <EOL> debias = True , n_orient = <NUM_LIT:1> , solver = '<STR_LIT>' ) : <EOL> """<STR_LIT>""" <EOL> def g ( w ) : <EOL> return np . sqrt ( np . sqrt ( groups_norm2 ( w . copy ( ) , n_orient ) ) ) <EOL> def gprime ( w ) : <EOL> return <NUM_LIT> * np . repeat ( g ( w ) , n_orient ) . ravel ( ) <EOL> E = list ( ) <EOL> active_set = np . ones ( G . shape [ <NUM_LIT:1> ] , dtype = np . bool ) <EOL> weights = np . ones ( G . shape [ <NUM_LIT:1> ] ) <EOL> X = np . zeros ( ( G . shape [ <NUM_LIT:1> ] , M . shape [ <NUM_LIT:1> ] ) ) <EOL> for k in range ( n_mxne_iter ) : <EOL> X0 = X . copy ( ) <EOL> active_set_0 = active_set . copy ( ) <EOL> G_tmp = G [ : , active_set ] * weights [ np . newaxis , : ] <EOL> if active_set_size is not None : <EOL> if np . sum ( active_set ) > ( active_set_size * n_orient ) : <EOL> X , _active_set , _ = mixed_norm_solver ( <EOL> M , G_tmp , alpha , debias = False , n_orient = n_orient , <EOL> maxit = maxit , tol = tol , active_set_size = active_set_size , <EOL> solver = solver , verbose = verbose ) <EOL> else : <EOL> X , _active_set , _ = mixed_norm_solver ( <EOL> M , G_tmp , alpha , debias = False , n_orient = n_orient , <EOL> maxit = maxit , tol = tol , active_set_size = None , solver = solver , <EOL> verbose = verbose ) <EOL> else : <EOL> X , _active_set , _ = mixed_norm_solver ( <EOL> M , G_tmp , alpha , debias = False , n_orient = n_orient , <EOL> maxit = maxit , tol = tol , active_set_size = None , solver = solver , <EOL> verbose = verbose ) <EOL> logger . info ( '<STR_LIT>' % ( _active_set . sum ( ) / n_orient ) ) <EOL> if _active_set . sum ( ) > <NUM_LIT:0> : <EOL> active_set [ active_set ] = _active_set <EOL> X *= weights [ _active_set ] [ : , np . newaxis ] <EOL> weights = gprime ( X ) <EOL> p_obj = <NUM_LIT:0.5> * linalg . norm ( M - np . dot ( G [ : , active_set ] , X ) , <EOL> '<STR_LIT>' ) ** <NUM_LIT> + alpha * np . sum ( g ( X ) ) <EOL> E . append ( p_obj ) <EOL> if ( ( k >= <NUM_LIT:1> ) and np . all ( active_set == active_set_0 ) and <EOL> np . all ( np . abs ( X - X0 ) < tol ) ) : <EOL> print ( '<STR_LIT>' % k ) <EOL> break <EOL> else : <EOL> active_set = np . zeros_like ( active_set ) <EOL> p_obj = <NUM_LIT:0.5> * linalg . norm ( M ) ** <NUM_LIT> <EOL> E . append ( p_obj ) <EOL> break <EOL> if np . any ( active_set ) and debias : <EOL> bias = compute_bias ( M , G [ : , active_set ] , X , n_orient = n_orient ) <EOL> X *= bias [ : , np . newaxis ] <EOL> return X , active_set , E <EOL> @ verbose <EOL> def tf_lipschitz_constant ( M , G , phi , phiT , tol = <NUM_LIT> , verbose = None ) : <EOL> """<STR_LIT>""" <EOL> n_times = M . shape [ <NUM_LIT:1> ] <EOL> n_points = G . shape [ <NUM_LIT:1> ] <EOL> iv = np . ones ( ( n_points , n_times ) , dtype = np . float ) <EOL> v = phi ( iv ) <EOL> L = <NUM_LIT> <EOL> for it in range ( <NUM_LIT:100> ) : <EOL> L_old = L <EOL> logger . info ( '<STR_LIT>' % it ) <EOL> iv = np . real ( phiT ( v ) ) <EOL> Gv = np . dot ( G , iv ) <EOL> GtGv = np . dot ( G . T , Gv ) <EOL> w = phi ( GtGv ) <EOL> L = np . max ( np . abs ( w ) ) <EOL> v = w / L <EOL> if abs ( ( L - L_old ) / L_old ) < tol : <EOL> break <EOL> return L <EOL> def safe_max_abs ( A , ia ) : <EOL> """<STR_LIT>""" <EOL> if np . sum ( ia ) : <EOL> return np . max ( np . abs ( A [ ia ] ) ) <EOL> else : <EOL> return <NUM_LIT:0.> <EOL> def safe_max_abs_diff ( A , ia , B , ib ) : <EOL> """<STR_LIT>""" <EOL> A = A [ ia ] if np . sum ( ia ) else <NUM_LIT:0.0> <EOL> B = B [ ib ] if np . sum ( ia ) else <NUM_LIT:0.0> <EOL> return np . max ( np . abs ( A - B ) ) <EOL> class _Phi ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , wsize , tstep , n_coefs ) : <EOL> self . wsize = wsize <EOL> self . tstep = tstep <EOL> self . n_coefs = n_coefs <EOL> def __call__ ( self , x ) : <EOL> return stft ( x , self . wsize , self . tstep , <EOL> verbose = False ) . reshape ( - <NUM_LIT:1> , self . n_coefs ) <EOL> class _PhiT ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , tstep , n_freq , n_step , n_times ) : <EOL> self . tstep = tstep <EOL> self . n_freq = n_freq <EOL> self . n_step = n_step <EOL> self . n_times = n_times <EOL> def __call__ ( self , z ) : <EOL> return istft ( z . reshape ( - <NUM_LIT:1> , self . n_freq , self . n_step ) , self . tstep , <EOL> self . n_times ) <EOL> def norm_l21_tf ( Z , shape , n_orient ) : <EOL> if Z . shape [ <NUM_LIT:0> ] : <EOL> Z2 = Z . reshape ( * shape ) <EOL> l21_norm = np . sqrt ( stft_norm2 ( Z2 ) . reshape ( - <NUM_LIT:1> , n_orient ) . sum ( axis = <NUM_LIT:1> ) ) <EOL> l21_norm = l21_norm . sum ( ) <EOL> else : <EOL> l21_norm = <NUM_LIT:0.> <EOL> return l21_norm <EOL> def norm_l1_tf ( Z , shape , n_orient ) : <EOL> if Z . shape [ <NUM_LIT:0> ] : <EOL> n_positions = Z . shape [ <NUM_LIT:0> ] // n_orient <EOL> Z_ = np . sqrt ( np . sum ( ( np . abs ( Z ) ** <NUM_LIT> ) . reshape ( ( n_orient , - <NUM_LIT:1> ) , <EOL> order = '<STR_LIT:F>' ) , axis = <NUM_LIT:0> ) ) <EOL> Z_ = Z_ . reshape ( ( n_positions , - <NUM_LIT:1> ) , order = '<STR_LIT:F>' ) . reshape ( * shape ) <EOL> l1_norm = ( <NUM_LIT> * Z_ . sum ( axis = <NUM_LIT:2> ) . sum ( axis = <NUM_LIT:1> ) - np . sum ( Z_ [ : , <NUM_LIT:0> , : ] , <EOL> axis = <NUM_LIT:1> ) - np . sum ( Z_ [ : , - <NUM_LIT:1> , : ] , axis = <NUM_LIT:1> ) ) <EOL> l1_norm = l1_norm . sum ( ) <EOL> else : <EOL> l1_norm = <NUM_LIT:0.> <EOL> return l1_norm <EOL> @ verbose <EOL> def _tf_mixed_norm_solver_bcd_ ( M , G , Z , active_set , alpha_space , alpha_time , <EOL> lipschitz_constant , phi , phiT , <EOL> wsize = <NUM_LIT:64> , tstep = <NUM_LIT:4> , n_orient = <NUM_LIT:1> , <EOL> maxit = <NUM_LIT:200> , tol = <NUM_LIT> , log_objective = True , <EOL> perc = None , verbose = None ) : <EOL> G = np . asfortranarray ( G ) <EOL> n_sensors , n_times = M . shape <EOL> n_sources = G . shape [ <NUM_LIT:1> ] <EOL> n_positions = n_sources // n_orient <EOL> n_step = int ( ceil ( n_times / float ( tstep ) ) ) <EOL> n_freq = wsize // <NUM_LIT:2> + <NUM_LIT:1> <EOL> shape = ( - <NUM_LIT:1> , n_freq , n_step ) <EOL> G = dict ( zip ( np . arange ( n_positions ) , np . hsplit ( G , n_positions ) ) ) <EOL> R = M . copy ( ) <EOL> active = np . where ( active_set ) [ <NUM_LIT:0> ] [ : : n_orient ] // n_orient <EOL> for idx in active : <EOL> R -= np . dot ( G [ idx ] , phiT ( Z [ idx ] ) ) <EOL> E = [ ] <EOL> alpha_time_lc = alpha_time / lipschitz_constant <EOL> alpha_space_lc = alpha_space / lipschitz_constant <EOL> converged = False <EOL> for i in range ( maxit ) : <EOL> val_norm_l21_tf = <NUM_LIT:0.0> <EOL> val_norm_l1_tf = <NUM_LIT:0.0> <EOL> max_diff = <NUM_LIT:0.0> <EOL> active_set_0 = active_set . copy ( ) <EOL> for j in range ( n_positions ) : <EOL> ids = j * n_orient <EOL> ide = ids + n_orient <EOL> G_j = G [ j ] <EOL> Z_j = Z [ j ] <EOL> active_set_j = active_set [ ids : ide ] <EOL> Z0 = deepcopy ( Z_j ) <EOL> was_active = np . any ( active_set_j ) <EOL> GTR = np . dot ( G_j . T , R ) / lipschitz_constant [ j ] <EOL> X_j_new = GTR . copy ( ) <EOL> if was_active : <EOL> X_j = phiT ( Z_j ) <EOL> R += np . dot ( G_j , X_j ) <EOL> X_j_new += X_j <EOL> rows_norm = linalg . norm ( X_j_new , '<STR_LIT>' ) <EOL> if rows_norm <= alpha_space_lc [ j ] : <EOL> if was_active : <EOL> Z [ j ] = <NUM_LIT:0.0> <EOL> active_set_j [ : ] = False <EOL> else : <EOL> if was_active : <EOL> Z_j_new = Z_j + phi ( GTR ) <EOL> else : <EOL> Z_j_new = phi ( GTR ) <EOL> col_norm = np . sqrt ( np . sum ( np . abs ( Z_j_new ) ** <NUM_LIT:2> , axis = <NUM_LIT:0> ) ) <EOL> if np . all ( col_norm <= alpha_time_lc [ j ] ) : <EOL> Z [ j ] = <NUM_LIT:0.0> <EOL> active_set_j [ : ] = False <EOL> else : <EOL> shrink = np . maximum ( <NUM_LIT:1.0> - alpha_time_lc [ j ] / np . maximum ( <EOL> col_norm , alpha_time_lc [ j ] ) , <NUM_LIT:0.0> ) <EOL> Z_j_new *= shrink [ np . newaxis , : ] <EOL> shape_init = Z_j_new . shape <EOL> Z_j_new = Z_j_new . reshape ( * shape ) <EOL> row_norm = np . sqrt ( stft_norm2 ( Z_j_new ) . sum ( ) ) <EOL> if row_norm <= alpha_space_lc [ j ] : <EOL> Z [ j ] = <NUM_LIT:0.0> <EOL> active_set_j [ : ] = False <EOL> else : <EOL> shrink = np . maximum ( <NUM_LIT:1.0> - alpha_space_lc [ j ] / <EOL> np . maximum ( row_norm , <EOL> alpha_space_lc [ j ] ) , <NUM_LIT:0.0> ) <EOL> Z_j_new *= shrink <EOL> Z [ j ] = Z_j_new . reshape ( - <NUM_LIT:1> , * shape_init [ <NUM_LIT:1> : ] ) . copy ( ) <EOL> active_set_j [ : ] = True <EOL> R -= np . dot ( G_j , phiT ( Z [ j ] ) ) <EOL> if log_objective : <EOL> val_norm_l21_tf += norm_l21_tf ( <EOL> Z [ j ] , shape , n_orient ) <EOL> val_norm_l1_tf += norm_l1_tf ( <EOL> Z [ j ] , shape , n_orient ) <EOL> max_diff = np . maximum ( max_diff , np . max ( np . abs ( Z [ j ] - Z0 ) ) ) <EOL> if log_objective : <EOL> pobj = ( <NUM_LIT:0.5> * ( R ** <NUM_LIT> ) . sum ( ) + alpha_space * val_norm_l21_tf + <EOL> alpha_time * val_norm_l1_tf ) <EOL> E . append ( pobj ) <EOL> logger . info ( "<STR_LIT>" % ( i + <NUM_LIT:1> , <EOL> pobj , np . sum ( active_set ) / n_orient ) ) <EOL> else : <EOL> logger . info ( "<STR_LIT>" % ( i + <NUM_LIT:1> ) ) <EOL> if perc is not None : <EOL> if np . sum ( active_set ) / float ( n_orient ) <= perc * n_positions : <EOL> break <EOL> if np . array_equal ( active_set , active_set_0 ) : <EOL> if max_diff < tol : <EOL> logger . info ( "<STR_LIT>" ) <EOL> converged = True <EOL> break <EOL> return Z , active_set , E , converged <EOL> @ verbose <EOL> def _tf_mixed_norm_solver_bcd_active_set ( <EOL> M , G , alpha_space , alpha_time , lipschitz_constant , phi , phiT , <EOL> Z_init = None , wsize = <NUM_LIT:64> , tstep = <NUM_LIT:4> , n_orient = <NUM_LIT:1> , maxit = <NUM_LIT:200> , tol = <NUM_LIT> , <EOL> log_objective = True , perc = None , verbose = None ) : <EOL> """<STR_LIT>""" <EOL> n_sources = G . shape [ <NUM_LIT:1> ] <EOL> n_positions = n_sources // n_orient <EOL> if Z_init is None : <EOL> Z = dict . fromkeys ( range ( n_positions ) , <NUM_LIT:0.0> ) <EOL> active_set = np . zeros ( n_sources , dtype = np . bool ) <EOL> else : <EOL> active_set = np . zeros ( n_sources , dtype = np . bool ) <EOL> active = list ( ) <EOL> for i in range ( n_positions ) : <EOL> if np . any ( Z_init [ i * n_orient : ( i + <NUM_LIT:1> ) * n_orient ] ) : <EOL> active_set [ i * n_orient : ( i + <NUM_LIT:1> ) * n_orient ] = True <EOL> active . append ( i ) <EOL> Z = dict . fromkeys ( range ( n_positions ) , <NUM_LIT:0.0> ) <EOL> if len ( active ) : <EOL> Z . update ( dict ( zip ( active , np . vsplit ( Z_init [ active_set ] , <EOL> len ( active ) ) ) ) ) <EOL> Z , active_set , E , _ = _tf_mixed_norm_solver_bcd_ ( <EOL> M , G , Z , active_set , alpha_space , alpha_time , lipschitz_constant , <EOL> phi , phiT , wsize = wsize , tstep = tstep , n_orient = n_orient , maxit = <NUM_LIT:1> , <EOL> tol = tol , log_objective = log_objective , perc = None , verbose = verbose ) <EOL> while active_set . sum ( ) : <EOL> active = np . where ( active_set ) [ <NUM_LIT:0> ] [ : : n_orient ] // n_orient <EOL> Z_init = dict ( zip ( range ( len ( active ) ) , [ Z [ idx ] for idx in active ] ) ) <EOL> Z , as_ , E_tmp , converged = _tf_mixed_norm_solver_bcd_ ( <EOL> M , G [ : , active_set ] , Z_init , <EOL> np . ones ( len ( active ) * n_orient , dtype = np . bool ) , <EOL> alpha_space , alpha_time , <EOL> lipschitz_constant [ active_set [ : : n_orient ] ] , <EOL> phi , phiT , wsize = wsize , tstep = tstep , n_orient = n_orient , <EOL> maxit = maxit , tol = tol , log_objective = log_objective , <EOL> perc = <NUM_LIT:0.5> , verbose = verbose ) <EOL> E += E_tmp <EOL> active = np . where ( active_set ) [ <NUM_LIT:0> ] [ : : n_orient ] // n_orient <EOL> Z_init = dict . fromkeys ( range ( n_positions ) , <NUM_LIT:0.0> ) <EOL> Z_init . update ( dict ( zip ( active , Z . values ( ) ) ) ) <EOL> active_set [ active_set ] = as_ <EOL> active_set_0 = active_set . copy ( ) <EOL> Z , active_set , E_tmp , _ = _tf_mixed_norm_solver_bcd_ ( <EOL> M , G , Z_init , active_set , alpha_space , alpha_time , <EOL> lipschitz_constant , phi , phiT , wsize = wsize , tstep = tstep , <EOL> n_orient = n_orient , maxit = <NUM_LIT:1> , tol = tol , log_objective = log_objective , <EOL> perc = None , verbose = verbose ) <EOL> E += E_tmp <EOL> if converged : <EOL> if np . array_equal ( active_set_0 , active_set ) : <EOL> break <EOL> if active_set . sum ( ) : <EOL> Z = np . vstack ( [ Z_ for Z_ in list ( Z . values ( ) ) if np . any ( Z_ ) ] ) <EOL> X = phiT ( Z ) <EOL> else : <EOL> n_sensors , n_times = M . shape <EOL> n_step = int ( ceil ( n_times / float ( tstep ) ) ) <EOL> n_freq = wsize // <NUM_LIT:2> + <NUM_LIT:1> <EOL> Z = np . zeros ( ( <NUM_LIT:0> , n_step * n_freq ) , dtype = np . complex ) <EOL> X = np . zeros ( ( <NUM_LIT:0> , n_times ) ) <EOL> return X , Z , active_set , E <EOL> @ verbose <EOL> def tf_mixed_norm_solver ( M , G , alpha_space , alpha_time , wsize = <NUM_LIT:64> , tstep = <NUM_LIT:4> , <EOL> n_orient = <NUM_LIT:1> , maxit = <NUM_LIT:200> , tol = <NUM_LIT> , log_objective = True , <EOL> debias = True , verbose = None ) : <EOL> """<STR_LIT>""" <EOL> n_sensors , n_times = M . shape <EOL> n_sensors , n_sources = G . shape <EOL> n_positions = n_sources // n_orient <EOL> n_step = int ( ceil ( n_times / float ( tstep ) ) ) <EOL> n_freq = wsize // <NUM_LIT:2> + <NUM_LIT:1> <EOL> n_coefs = n_step * n_freq <EOL> phi = _Phi ( wsize , tstep , n_coefs ) <EOL> phiT = _PhiT ( tstep , n_freq , n_step , n_times ) <EOL> if n_orient == <NUM_LIT:1> : <EOL> lc = np . sum ( G * G , axis = <NUM_LIT:0> ) <EOL> else : <EOL> lc = np . empty ( n_positions ) <EOL> for j in range ( n_positions ) : <EOL> G_tmp = G [ : , ( j * n_orient ) : ( ( j + <NUM_LIT:1> ) * n_orient ) ] <EOL> lc [ j ] = linalg . norm ( np . dot ( G_tmp . T , G_tmp ) , ord = <NUM_LIT:2> ) <EOL> logger . info ( "<STR_LIT>" ) <EOL> X , Z , active_set , E = _tf_mixed_norm_solver_bcd_active_set ( <EOL> M , G , alpha_space , alpha_time , lc , phi , phiT , Z_init = None , <EOL> wsize = wsize , tstep = tstep , n_orient = n_orient , maxit = maxit , tol = tol , <EOL> log_objective = log_objective , verbose = None ) <EOL> if np . any ( active_set ) and debias : <EOL> bias = compute_bias ( M , G [ : , active_set ] , X , n_orient = n_orient ) <EOL> X *= bias [ : , np . newaxis ] <EOL> return X , active_set , E </s>
