From dd9796ddd58325d02b9f3ebd58d6b91307d53d86 Mon Sep 17 00:00:00 2001 From: nightwnvol Date: Fri, 28 Jun 2024 10:50:11 +0200 Subject: [PATCH] chore: improve output from jupyter --- commit/core.pyx | 87 ++++++++++++++---------- commit/solvers.py | 2 +- commit/trk2dictionary/trk2dictionary.pyx | 49 +++++++------ 3 files changed, 80 insertions(+), 58 deletions(-) diff --git a/commit/core.pyx b/commit/core.pyx index 3454d46..8f263f0 100755 --- a/commit/core.pyx +++ b/commit/core.pyx @@ -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 @@ -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') ) @@ -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() @@ -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,:] ) @@ -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 ) @@ -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 ) @@ -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 ) @@ -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' ) @@ -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'] @@ -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 : @@ -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 : @@ -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.' ) @@ -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 @@ -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] @@ -1526,8 +1536,9 @@ 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] @@ -1535,8 +1546,9 @@ cdef class Evaluation : 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] @@ -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) ) @@ -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') : @@ -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 diff --git a/commit/solvers.py b/commit/solvers.py index d45325c..760b5db 100755 --- a/commit/solvers.py +++ b/commit/solvers.py @@ -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" diff --git a/commit/trk2dictionary/trk2dictionary.pyx b/commit/trk2dictionary/trk2dictionary.pyx index ab7c8ee..9e8ae5b 100755 --- a/commit/trk2dictionary/trk2dictionary.pyx +++ b/commit/trk2dictionary/trk2dictionary.pyx @@ -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 ) @@ -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)} ]' ) \ No newline at end of file