Skip to content

Commit

Permalink
chore: improve output from jupyter
Browse files Browse the repository at this point in the history
  • Loading branch information
@nightwnvol
nightwnvol committed Jun 28, 2024
1 parent 9ecacf3 commit dd9796d
Show file tree
Hide file tree
Showing 3 changed files with 80 additions and 58 deletions.
87 changes: 51 additions & 36 deletions commit/core.pyx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -223,8 +223,9 @@ cdef class Evaluation :

if self.get_config('doNormalizeSignal') :
if self.scheme.b0_count > 0:
logger.subinfo('Normalizing to b0:', with_progress=True, indent_char='*', indent_lvl=1)
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=True) as pbar:
log_list = []
ret_subinfo = logger.subinfo('Normalizing to b0:', with_progress=True, indent_char='*', indent_lvl=1)
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=ret_subinfo, log_list=log_list):
b0 = np.mean( self.niiDWI_img[:,:,:,self.scheme.b0_idx], axis=3 )
idx = b0 <= b0_min_signal * b0[b0>0].mean()
b0[ idx ] = 1
Expand Down Expand Up @@ -352,8 +353,9 @@ cdef class Evaluation :
tic = time.time()
logger.subinfo('')
logger.info( 'Loading the kernels' )
logger.subinfo( 'Resampling LUT for subject "%s":' % self.get_config('subject'), indent_char='*', indent_lvl=1, with_progress=True ) # TODO: check why not printed
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=True) as pbar:
log_list = []
ret_subinfo = logger.subinfo( 'Resampling LUT for subject "%s":' % self.get_config('subject'), indent_char='*', indent_lvl=1, with_progress=True ) # TODO: check why not printed
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=ret_subinfo, log_list=log_list):
# auxiliary data structures
idx_OUT, Ylm_OUT = amico.lut.aux_structures_resample( self.scheme, self.get_config('lmax') )

Expand All @@ -375,8 +377,9 @@ cdef class Evaluation :

# De-mean kernels
if self.get_config('doDemean') :
logger.subinfo('Demeaning signal', with_progress=True, indent_lvl=2, indent_char='-' )
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=True) as pbar:
log_list = []
ret_subinfo = logger.subinfo('Demeaning signal', with_progress=True, indent_lvl=2, indent_char='-' )
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=ret_subinfo, log_list=log_list):
for j in xrange(self.get_config('ndirs')) :
for i in xrange(nIC) :
self.KERNELS['wmr'][i,j,:] -= self.KERNELS['wmr'][i,j,:].mean()
Expand All @@ -387,8 +390,9 @@ cdef class Evaluation :

# Normalize atoms
if self.get_config('doNormalizeKernels') :
logger.subinfo('Normalizing kernels', with_progress=True, indent_lvl=2, indent_char='-' )
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=True) as pbar:
log_list = []
ret_subinfo = logger.subinfo('Normalizing kernels', with_progress=True, indent_lvl=2, indent_char='-' )
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=ret_subinfo, log_list=log_list):
self.KERNELS['wmr_norm'] = np.zeros( nIC )
for i in xrange(nIC) :
self.KERNELS['wmr_norm'][i] = np.linalg.norm( self.KERNELS['wmr'][i,0,:] )
Expand Down Expand Up @@ -458,8 +462,9 @@ cdef class Evaluation :

# segments from the tracts
# ------------------------
logger.subinfo('Segments from the tracts:', indent_char='*', indent_lvl=1, with_progress=True )
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=True) as pbar:
log_list = []
ret_subinfo = logger.subinfo('Segments from the tracts:', indent_char='*', indent_lvl=1, with_progress=True )
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=ret_subinfo, log_list=log_list):
self.DICTIONARY['TRK'] = {}
self.DICTIONARY['TRK']['kept'] = np.fromfile( pjoin(self.get_config('TRACKING_path'),'dictionary_TRK_kept.dict'), dtype=np.uint8 )
self.DICTIONARY['TRK']['norm'] = np.fromfile( pjoin(self.get_config('TRACKING_path'),'dictionary_TRK_norm.dict'), dtype=np.float32 )
Expand Down Expand Up @@ -498,8 +503,9 @@ cdef class Evaluation :

# segments from the peaks
# -----------------------
logger.subinfo('Segments from the peaks:', indent_char='*', indent_lvl=1, with_progress=True )
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=True) as pbar:
log_list = []
ret_subinfo = logger.subinfo('Segments from the peaks:', indent_char='*', indent_lvl=1, with_progress=True )
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=ret_subinfo, log_list=log_list):
self.DICTIONARY['EC'] = {}
self.DICTIONARY['EC']['v'] = np.fromfile( pjoin(self.get_config('TRACKING_path'),'dictionary_EC_v.dict'), dtype=np.uint32 )
self.DICTIONARY['EC']['o'] = np.fromfile( pjoin(self.get_config('TRACKING_path'),'dictionary_EC_o.dict'), dtype=np.uint16 )
Expand All @@ -515,8 +521,9 @@ cdef class Evaluation :

# isotropic compartments
# ----------------------
logger.subinfo('Isotropic contributions:', indent_char='*', indent_lvl=1, with_progress=True )
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=True) as pbar:
log_list = []
ret_subinfo = logger.subinfo('Isotropic contributions:', indent_char='*', indent_lvl=1, with_progress=True )
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=ret_subinfo, log_list=log_list):
self.DICTIONARY['ISO'] = {}

self.DICTIONARY['ISO']['v'] = np.fromfile( pjoin(self.get_config('TRACKING_path'),'dictionary_ISO_v.dict'), dtype=np.uint32 )
Expand All @@ -533,8 +540,9 @@ cdef class Evaluation :

# post-processing
# ---------------
logger.subinfo('Post-processing', indent_char='*', indent_lvl=1, with_progress=True )
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=True) as pbar:
log_list = []
ret_subinfo = logger.subinfo('Post-processing', indent_char='*', indent_lvl=1, with_progress=True )
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=ret_subinfo, log_list=log_list):
# get the indices to extract the VOI as in MATLAB (in place of DICTIONARY.MASKidx)
idx = self.DICTIONARY['MASK'].ravel(order='F').nonzero()[0]
self.DICTIONARY['MASK_ix'], self.DICTIONARY['MASK_iy'], self.DICTIONARY['MASK_iz'] = np.unravel_index( idx, self.DICTIONARY['MASK'].shape, order='F' )
Expand All @@ -557,7 +565,7 @@ cdef class Evaluation :
----------
n : integer
Number of threads to use (default : number of CPUs in the system)
"""
"""
if n is None :
# Use the same number of threads used in trk2dictionary
n = self.DICTIONARY['n_threads']
Expand All @@ -584,8 +592,9 @@ cdef class Evaluation :
logger.subinfo('Number of threads: %d' % n , indent_char='*', indent_lvl=1 )

# Distribute load for the computation of A*x product
logger.subinfo('A operator ', indent_char='*', indent_lvl=1, with_progress=True )
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=True) as pbar:
log_list = []
ret_subinfo = logger.subinfo('A operator ', indent_char='*', indent_lvl=1, with_progress=True )
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=ret_subinfo, log_list=log_list):
if self.DICTIONARY['IC']['n'] > 0 :
self.THREADS['IC'] = np.zeros( n+1, dtype=np.uint32 )
if n > 1 :
Expand Down Expand Up @@ -636,8 +645,9 @@ cdef class Evaluation :


# Distribute load for the computation of At*y product
logger.subinfo('A\' operator', indent_char='*', indent_lvl=1, with_progress=True )
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=True) as pbar:
log_list = []
ret_subinfo = logger.subinfo('A\' operator', indent_char='*', indent_lvl=1, with_progress=True )
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=ret_subinfo, log_list=log_list):
if self.DICTIONARY['IC']['n'] > 0 :
self.THREADS['ICt'] = np.full( self.DICTIONARY['IC']['n'], n-1, dtype=np.uint8 )
if n > 1 :
Expand Down Expand Up @@ -741,7 +751,6 @@ cdef class Evaluation :
compilation_is_needed = True

if compilation_is_needed or not 'commit.operator.operator' in sys.modules :

if build_dir is not None:
if isdir(build_dir) and not len(listdir(build_dir)) == 0:
logger.error( '\nbuild_dir is not empty, unsafe build option.' )
Expand Down Expand Up @@ -1353,7 +1362,7 @@ cdef class Evaluation :

# run solver
t = time.time()
with ProgressBar(disable=self.verbose!=3, hide_on_exit=True) as pb:
with ProgressBar(disable=self.verbose!=3, hide_on_exit=True):
self.x, opt_details = commit.solvers.solve(self.get_y(), self.A, self.A.T, tol_fun=tol_fun, tol_x=tol_x, max_iter=max_iter, verbose=self.verbose, x0=x0, regularisation=self.regularisation_params, confidence_array=confidence_array)

self.CONFIG['optimization']['fit_details'] = opt_details
Expand Down Expand Up @@ -1515,8 +1524,9 @@ cdef class Evaluation :
# Map of compartment contributions
logger.subinfo('Voxelwise contributions:', indent_char='*', indent_lvl=1)

logger.subinfo('Intra-axonal', indent_lvl=2, indent_char='-', with_progress=True)
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=True) as pbar:
log_list = []
ret_subinfo = logger.subinfo('Intra-axonal', indent_lvl=2, indent_char='-', with_progress=True)
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=ret_subinfo, log_list=log_list):
niiIC_img = np.zeros( self.get_config('dim'), dtype=np.float32 )
if len(self.KERNELS['wmr']) > 0 :
offset = nF * self.KERNELS['wmr'].shape[0]
Expand All @@ -1526,17 +1536,19 @@ cdef class Evaluation :
).astype(np.float32)
niiIC_img[ self.DICTIONARY['MASK_ix'], self.DICTIONARY['MASK_iy'], self.DICTIONARY['MASK_iz'] ] = xv

logger.subinfo('Extra-axonal', indent_lvl=2, indent_char='-', with_progress=True)
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=True) as pbar:
log_list = []
ret_subinfo = logger.subinfo('Extra-axonal', indent_lvl=2, indent_char='-', with_progress=True)
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=ret_subinfo, log_list=log_list):
niiEC_img = np.zeros( self.get_config('dim'), dtype=np.float32 )
if len(self.KERNELS['wmh']) > 0 :
offset = nF * self.KERNELS['wmr'].shape[0]
tmp = x[offset:offset+nE*len(self.KERNELS['wmh'])].reshape( (-1,nE) ).sum( axis=0 )
xv = np.bincount( self.DICTIONARY['EC']['v'], weights=tmp, minlength=nV ).astype(np.float32)
niiEC_img[ self.DICTIONARY['MASK_ix'], self.DICTIONARY['MASK_iy'], self.DICTIONARY['MASK_iz'] ] = xv

logger.subinfo('Isotropic ', indent_lvl=2, indent_char='-', with_progress=True)
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=True) as pbar:
log_list = []
ret_subinfo = logger.subinfo('Isotropic ', indent_lvl=2, indent_char='-', with_progress=True)
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=ret_subinfo, log_list=log_list):
niiISO_img = np.zeros( self.get_config('dim'), dtype=np.float32 )
if len(self.KERNELS['iso']) > 0 :
offset = nF * self.KERNELS['wmr'].shape[0] + nE * self.KERNELS['wmh'].shape[0]
Expand All @@ -1560,8 +1572,9 @@ cdef class Evaluation :

# Configuration and results
logger.subinfo('Configuration and results:', indent_char='*', indent_lvl=1)
logger.subinfo('streamline_weights.txt', indent_lvl=2, indent_char='-', with_progress=True)
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=True) as pbar:
log_list = []
ret_subinfo = logger.subinfo('streamline_weights.txt', indent_lvl=2, indent_char='-', with_progress=True)
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=ret_subinfo, log_list=log_list):
xic, _, _ = self.get_coeffs()
if stat_coeffs != 'all' and xic.size > 0 :
xic = np.reshape( xic, (-1,self.DICTIONARY['TRK']['kept'].size) )
Expand Down Expand Up @@ -1620,8 +1633,9 @@ cdef class Evaluation :
# item 0: dictionary with all the configuration details
# item 1: np.array obtained through the optimisation process with the normalised kernels
# item 2: np.array renormalisation of coeffs in item 1
logger.subinfo('results.pickle', indent_char='-', indent_lvl=2, with_progress=True)
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=True) as pbar:
log_list = []
ret_subinfo = logger.subinfo('results.pickle', indent_char='-', indent_lvl=2, with_progress=True)
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=ret_subinfo, log_list=log_list):
xic, xec, xiso = self.get_coeffs()
x = self.x
if self.get_config('doNormalizeKernels') :
Expand All @@ -1631,9 +1645,10 @@ cdef class Evaluation :
self.CONFIG['optimization']['regularisation'].pop('prox', None)
pickle.dump( [self.CONFIG, x, self.x], fid, protocol=2 )

if save_est_dwi :
logger.subinfo('Estimated signal:', indent_char='-', indent_lvl=2, with_progress=True)
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=True) as pbar:
if save_est_dwi:
log_list = []
ret_subinfo = logger.subinfo('Estimated signal:', indent_char='-', indent_lvl=2, with_progress=True)
with ProgressBar(disable=self.verbose < 3, hide_on_exit=True, subinfo=ret_subinfo, log_list=log_list):
self.niiDWI_img[ self.DICTIONARY['MASK_ix'], self.DICTIONARY['MASK_iy'], self.DICTIONARY['MASK_iz'], : ] = y_est
nibabel.save( nibabel.Nifti1Image( self.niiDWI_img , affine ), pjoin(RESULTS_path,'fit_signal_estimated.nii.gz') )
self.niiDWI_img[ self.DICTIONARY['MASK_ix'], self.DICTIONARY['MASK_iy'], self.DICTIONARY['MASK_iz'], : ] = y_mea
Expand Down
2 changes: 1 addition & 1 deletion commit/solvers.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -341,7 +341,7 @@ def fista(y, A, At, omega, prox, sqrt_W=None, tol_fun=1e-4, tol_x=1e-6, max_iter
abs_x = np.linalg.norm(x - prev_x)
rel_x = abs_x / ( np.linalg.norm(x) + eps )
if verbose==4 :
print( " %13.7e %13.7e | %13.7e %13.7e %13.7e | %13.7e %13.7e" % ( 0.5 * res_norm**2, reg_term_x, curr_obj, abs_obj, rel_obj, abs_x, rel_x ) )
print( " %13.7e %13.7e | %13.7e %13.7e %13.7e | %13.7e %13.7e" % ( 0.5 * res_norm**2, reg_term_x, curr_obj, abs_obj, rel_obj, abs_x, rel_x ), flush=True )

if abs_obj < eps :
criterion = "Absolute tolerance on the objective"
Expand Down
49 changes: 28 additions & 21 deletions commit/trk2dictionary/trk2dictionary.pyx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -543,10 +543,15 @@ cpdef run( filename_tractogram=None, path_out=None, filename_peaks=None, filenam
logger.warning( 'DICTIONARY not generated' )
return None

# NOTE: this is to ensure flushing all the output from the cpp code
print(end='', flush=True)
time.sleep(0.5)

# Concatenate files together
logger.subinfo( 'Saving data structure', indent_char='*', indent_lvl=1, with_progress=True )
log_list = []
ret_subinfo = logger.subinfo( 'Saving data structure', indent_char='*', indent_lvl=1, with_progress=verbose>2 )
cdef int discarded = 0
with ProgressBar(disable=verbose<3, hide_on_exit=True, subinfo=True) as pbar:
with ProgressBar(disable=verbose<3, hide_on_exit=True, subinfo=ret_subinfo, log_list=log_list):
for j in range(n_threads-1):
path_IC_f = path_temp + f'/dictionary_IC_f_{j+1}.dict'
kept = np.fromfile( path_temp + f'/dictionary_TRK_kept_{j}.dict', dtype=np.uint8 )
Expand Down Expand Up @@ -619,29 +624,31 @@ cpdef run( filename_tractogram=None, path_out=None, filename_peaks=None, filenam
TDI_affine = np.diag( [Px, Py, Pz, 1] )

# save TDI and MASK maps
logger.subinfo( 'Saving TDI and MASK maps', indent_char='*', indent_lvl=1 )
v = np.fromfile( join(path_out, 'dictionary_IC_v.dict'), dtype=np.uint32 )
l = np.fromfile( join(path_out, 'dictionary_IC_len.dict'), dtype=np.float32 )

cdef np.float32_t [::1] niiTDI_mem = np.zeros( Nx*Ny*Nz, dtype=np.float32 )
niiTDI_mem = compute_tdi( v, l, Nx, Ny, Nz, verbose )
niiTDI_img_save = np.reshape( niiTDI_mem, (Nx,Ny,Nz), order='F' )
log_list = []
ret_subinfo = logger.subinfo( 'Saving TDI and MASK maps', indent_char='*', indent_lvl=1, with_progress=verbose>2)
with ProgressBar(disable=verbose<3, hide_on_exit=True, subinfo=ret_subinfo, log_list=log_list):
v = np.fromfile( join(path_out, 'dictionary_IC_v.dict'), dtype=np.uint32 )
l = np.fromfile( join(path_out, 'dictionary_IC_len.dict'), dtype=np.float32 )

niiTDI_mem = compute_tdi( v, l, Nx, Ny, Nz, verbose )
niiTDI_img_save = np.reshape( niiTDI_mem, (Nx,Ny,Nz), order='F' )

niiTDI = nibabel.Nifti1Image( niiTDI_img_save, TDI_affine )
niiTDI_hdr = _get_header( niiTDI )
niiTDI_hdr['descrip'] = f'Created with COMMIT {get_distribution("dmri-commit").version}'
nibabel.save( niiTDI, join(path_out,'dictionary_tdi.nii.gz') )
niiTDI = nibabel.Nifti1Image( niiTDI_img_save, TDI_affine )
niiTDI_hdr = _get_header( niiTDI )
niiTDI_hdr['descrip'] = f'Created with COMMIT {get_distribution("dmri-commit").version}'
nibabel.save( niiTDI, join(path_out,'dictionary_tdi.nii.gz') )

if filename_mask is not None :
niiMASK = nibabel.Nifti1Image( niiMASK_img, TDI_affine )
else :
niiMASK = nibabel.Nifti1Image( (np.asarray(niiTDI_img_save)>0).astype(np.float32), TDI_affine )
if filename_mask is not None :
niiMASK = nibabel.Nifti1Image( niiMASK_img, TDI_affine )
else :
niiMASK = nibabel.Nifti1Image( (np.asarray(niiTDI_img_save)>0).astype(np.float32), TDI_affine )

niiMASK_hdr = _get_header( niiMASK )
niiMASK_hdr['descrip'] = niiTDI_hdr['descrip']
nibabel.save( niiMASK, join(path_out,'dictionary_mask.nii.gz') )
niiMASK_hdr = _get_header( niiMASK )
niiMASK_hdr['descrip'] = niiTDI_hdr['descrip']
nibabel.save( niiMASK, join(path_out,'dictionary_mask.nii.gz') )

if not keep_temp:
shutil.rmtree(path_temp)
if not keep_temp:
shutil.rmtree(path_temp)

logger.info( f'[ {format_time(time.time() - tic)} ]' )

0 comments on commit dd9796d

Please sign in to comment.