Internal rework around spams.lassoWeighted

CHANGELOG.md

@@ -1,5 +1,12 @@
 # Changelog
 
+## [1.0] - 2025-xx-xx
+
+- Some internal rework for faster processing by an order of magnitude
+	- spams-bin 2.6.12 or higher is now *required* as it fixes a positivity bug needed in the solver
+- Some cleanups of old deprecated command line options
+- Frozen builds are now using the threading interface
+
 ## [0.7.3] - 2025-10-06
 
 - Removed `PIESNO` as a noise estimation method. Use `auto` instead, which is the new default since 0.7, as it will automatically estimate `N` for you.

docs/conf.py

@@ -59,9 +59,9 @@
 # built documents.
 #
 # The short X.Y version.
-version = '0.7.3'
+version = '1.0'
 # The full version, including alpha/beta/rc tags.
-release = '0.7.3'
+release = '1.0'
 
 # List of patterns, relative to source directory, that match files and
 # directories to ignore when looking for source files.

nlsam/denoiser.py

@@ -189,43 +189,39 @@ def local_denoise(data, block_size, overlap, variance, n_iter=10, mask=None,
     if mask is None:
         mask = np.ones(data.shape[:-1], dtype=bool)
 
-    X = extract_patches(data, block_size, [1, 1, 1, block_size[-1]]).reshape(-1, np.prod(block_size)).T
+    m = np.prod(block_size)
+    X = extract_patches(data, block_size, [1, 1, 1, block_size[-1]]).reshape(-1, m).T
 
     # Solving for D
     param_alpha = {}
     param_alpha['pos'] = True
     param_alpha['mode'] = 1
+    param_alpha['numThreads'] = 1
 
     param_D = {}
     param_D['verbose'] = False
     param_D['posAlpha'] = True
     param_D['posD'] = True
     param_D['mode'] = 2
-    param_D['lambda1'] = 1.2 / np.sqrt(np.prod(block_size))
-    param_D['K'] = int(2 * np.prod(block_size))
+    param_D['lambda1'] = 1.2 / np.sqrt(m)
+    param_D['K'] = int(2 * m)
     param_D['iter'] = 150
     param_D['batchsize'] = 500
     param_D['numThreads'] = n_cores
 
-    if 'D' in param_alpha:
-        param_D['D'] = param_alpha['D']
-
     mask_col = extract_patches(mask, block_size[:-1], (1, 1, 1), flatten=False)
     axis = tuple(range(mask_col.ndim//2, mask_col.ndim))
     train_idx = np.sum(mask_col, axis=axis).ravel() > (np.prod(block_size[:-1]) / 2)
 
     train_data = np.asfortranarray(X[:, train_idx])
-    train_data /= np.sqrt(np.sum(train_data**2, axis=0, keepdims=True), dtype=dtype)
+    train_data /= np.linalg.norm(train_data, axis=0, keepdims=True).astype(dtype, copy=False)
 
-    param_alpha['D'] = spams.trainDL(train_data, **param_D)
-    param_alpha['D'] /= np.sqrt(np.sum(param_alpha['D']**2, axis=0, keepdims=True, dtype=dtype))
-    param_D['D'] = param_alpha['D']
+    D = spams.trainDL(train_data, **param_D)
+    D /= np.linalg.norm(D, axis=0, keepdims=True).astype(dtype, copy=False)
+    param_alpha['D'] = D
 
     del train_idx, train_data, X, mask_col
 
-    param_alpha['numThreads'] = 1
-    param_D['numThreads'] = 1
-
     slicer = [np.index_exp[:, :, k:k + block_size[2]] for k in range((data.shape[2] - block_size[2] + 1))]
 
     if verbose:
@@ -241,13 +237,12 @@ def local_denoise(data, block_size, overlap, variance, n_iter=10, mask=None,
                                                                 block_size,
                                                                 overlap,
                                                                 param_alpha,
-                                                                param_D,
                                                                 current_slice,
                                                                 dtype,
                                                                 n_iter)
                                                                 for current_slice in progress_slicer)
 
-    logger.info(f'Multiprocessing done in {(time() - time_multi) / 60:.2f} mins.')
+    logger.info(f'Multiprocessing done in {int(time() - time_multi)}s')
 
     # Put together the multiprocessed results
     data_subset = np.zeros_like(data, dtype=np.float32)
@@ -261,7 +256,7 @@ def local_denoise(data, block_size, overlap, variance, n_iter=10, mask=None,
     return data_subset
 
 
-def processer(data, mask, variance, block_size, overlap, param_alpha, param_D, current_slice,
+def processer(data, mask, variance, block_size, overlap, param_alpha, current_slice,
               dtype=np.float64, n_iter=10, gamma=3, tau=1, tolerance=1e-5):
 
     # Fetch the current slice for parallel processing since now the arrays are dumped and read from disk
@@ -283,54 +278,49 @@ def processer(data, mask, variance, block_size, overlap, param_alpha, param_D, c
     var_mat = np.median(im2col_nd(variance, block_size[:-1], overlap[:-1])[:, train_idx], axis=0)
     X_full_shape = X.shape
     X = X[:, train_idx].astype(dtype)
-
-    param_alpha['L'] = int(0.5 * X.shape[0])
-
     D = param_alpha['D']
 
-    alpha = np.zeros((D.shape[1], X.shape[1]), dtype=dtype)
-    W = np.ones(alpha.shape, dtype=dtype)
-    temp = np.zeros([alpha.shape[0], 1], dtype=dtype)
+    m, n = X.shape
+    p = D.shape[1]
+    param_alpha['L'] = m // 2
 
-    DtD = np.asfortranarray(D.T @ D)
-    DtX = np.asfortranarray(D.T @ X)
-    DtXW = np.empty_like(DtX, order='F')
-    DtDW = np.empty_like(DtD, order='F')
+    alpha = np.zeros((p, n), dtype=dtype)
+    alpha_old = np.ones_like(alpha)
+    W = np.ones_like(alpha, order='F', dtype=dtype)
 
-    alpha_old = np.ones(alpha.shape, dtype=dtype)
-    not_converged = np.ones(alpha.shape[1], dtype=bool)
+    not_converged = np.ones(n, dtype=bool)
     nonzero_ind = np.zeros(alpha.shape, dtype=bool)
-    arr = np.empty(alpha.shape)
 
-    xi = np.random.randn(X.shape[0], X.shape[1]) * var_mat
-    var_mat *= (X.shape[0] + gamma * np.sqrt(2 * X.shape[0]))
+    xi = np.random.randn(m, n) * var_mat
+    var_mat *= (m + gamma * np.sqrt(2 * m))
     eps = np.max(np.abs(D.T @ xi), axis=0)
 
-    for _ in range(n_iter):
-        DtXW[:, not_converged] = DtX[:, not_converged] / W[:, not_converged]
+    # Rescale to ensure lambda1=1 as a constant, this allows passing everything to spams.lassoWeighted
+    # but the output coefficients are scaled by an extra np.sqrt(var_mat) factor since we do not rescale D to keep it constant also
+    scale = np.sqrt(var_mat)
+    X = np.asfortranarray(X / scale, dtype=dtype)
+    param_alpha['lambda1'] = 1
 
-        for i in range(alpha.shape[1]):
-            if not_converged[i]:
-                param_alpha['lambda1'] = var_mat[i]
-                DtDW[:] = (1 / W[..., None, i]) * DtD * (1 / W[:, i])
-                spams.lasso(X[:, i:i+1], Q=DtDW, q=DtXW[:, i:i+1], **param_alpha).todense(out=temp)
-                alpha[:, i:i+1] = temp
+    for _ in range(n_iter):
+        Xi = X[:, not_converged]
+        Wi = W[:, not_converged]
+        alpha[:, not_converged] = spams.lassoWeighted(X=Xi, W=Wi, **param_alpha).toarray() * scale[not_converged]
 
-        arr[:] = alpha
-        nonzero_ind[:] = arr != 0
-        arr[nonzero_ind] /= W[nonzero_ind]
-        not_converged[:] = np.max(np.abs(alpha_old - arr), axis=0) > tolerance
+        nonzero_ind[:] = alpha != 0
+        not_converged[:] = np.max(np.abs(alpha_old - alpha), axis=0) > tolerance
 
         if not np.any(not_converged):
             break
 
-        alpha_old[:] = arr
+        alpha_old[:] = alpha
         W[:] = 1 / (np.abs(alpha_old**tau) + eps)
 
     weights = np.ones(X_full_shape[1], dtype=dtype)
     weights[train_idx] = 1 / (np.sum(alpha != 0, axis=0) + 1)
     X = np.zeros(X_full_shape, dtype=dtype, order='F')
-    X[:, train_idx] = D @ arr
+    X[:, train_idx] = D @ alpha
+
     out = col2im_nd(X, block_size, orig_shape, overlap, weights)
-    del X, W, alpha, alpha_old, DtX, DtXW, DtDW
+    del X, W, alpha, alpha_old
+
     return out

nlsam/script.py

@@ -203,19 +203,19 @@ def buildArgsParser():
     # Old stuff
     ############
 
-    deprecated = p.add_argument_group('Deprecated options')
+    # deprecated = p.add_argument_group('Deprecated options')
 
-    deprecated.add_argument('--fix_implausible', action='store_true', dest='implausible_signal_fix',
-                            help='This option has been removed and has no effect.')
+    # deprecated.add_argument('--fix_implausible', action='store_true', dest='implausible_signal_fix',
+    #                         help='This option has been removed and has no effect.')
 
-    deprecated.add_argument('--sh_order', metavar='int', default=0, type=int, choices=[0, 2, 4, 6, 8],
-                            help='This option has been removed and has no effect.')
+    # deprecated.add_argument('--sh_order', metavar='int', default=0, type=int, choices=[0, 2, 4, 6, 8],
+    #                         help='This option has been removed and has no effect.')
 
-    deprecated.add_argument('--mp_method', metavar='string',
-                            help='This option has been removed and has no effect.')
+    # deprecated.add_argument('--mp_method', metavar='string',
+    #                         help='This option has been removed and has no effect.')
 
-    deprecated.add_argument('--use_threading', action='store_true',
-                            help='This option has been removed and has no effect.')
+    # deprecated.add_argument('--use_threading', action='store_true',
+    #                         help='This option has been removed and has no effect.')
 
     return p
 
@@ -263,18 +263,6 @@ def main():
     else:
         dtype = np.float64
 
-    if args.implausible_signal_fix:
-        logger.warning('Option --implausible_signal_fix has been deprecated')
-
-    if args.sh_order:
-        logger.warning('Option --sh_order has been deprecated')
-
-    if args.mp_method:
-        logger.warning('Option --mp_method has been deprecated')
-
-    if args.use_threading:
-        logger.warning('Option --use_threading has been deprecated')
-
     ##########################################
     #  Load up data and do some sanity checks
     ##########################################

nlsam/tests/test_main.py

@@ -24,20 +24,20 @@ def unzip(zip, file):
 
 
 commands = [
-    ('nlsam_denoising', 'dwi_crop.nii.gz', 'dwi_nlsam.nii.gz', 'bvals', 'bvecs', '-f', '-N', '1', '--noise_est', 'local_std', '--sh_order', '0', '--cores', '1', '-m', 'mask_crop.nii.gz', '-v'),
-    ('nlsam_denoising', 'dwi_crop.nii.gz', 'dwi_nlsam.nii.gz', 'bvals', 'bvecs', '-m', 'mask_crop.nii.gz', '-f', '-N', '1', '--noise_est', 'local_std', '--sh_order', '6', '--iterations', '5', '--verbose', '--save_sigma', 'sigma.nii.gz', '--log', 'log.txt'),
-    ('nlsam_denoising', 'dwi_crop.nii.gz', 'dwi_nlsam.nii.gz', 'bvals', 'bvecs', '-m', 'mask_crop.nii.gz', '-f', '--noise_est', 'auto', '--sh_order', '6', '--iterations', '5', '--verbose', '--save_sigma', 'sigma.nii.gz', '--save_N', 'N.nii.gz'),
-    ('nlsam_denoising', 'dwi_crop.nii.gz', 'dwi_nlsam.nii.gz', 'bvals', 'bvecs', '-m', 'mask_crop.nii.gz', '-f', '--verbose', '--sh_order', '0', '--b0_threshold', '10', '--noise_mask', 'pmask.nii.gz'),
-    ('nlsam_denoising', 'dwi_crop.nii', 'dwi_nlsam.nii', 'bvals', 'bvecs', '-m', 'mask_crop.nii', '-f', '--verbose', '--sh_order', '0', '-N', '1', '--no_stabilization', '--load_sigma', 'sigma.nii', '--is_symmetric', '--use_threading', '--save_difference', 'diff.nii'),
-    ('nlsam_denoising', 'dwi_crop.nii', 'dwi_nlsam.nii', 'bvals', 'bvecs', '-m', 'mask_crop.nii', '-f', '--verbose', '--sh_order', '0', '--no_denoising', '--save_sigma', 'sigma.nii', '--save_stab', 'stab.nii', '--load_mhat', 'dwi_crop.nii', '--save_eta', 'eta.nii'),
-    ('nlsam_denoising', 'dwi_crop.nii.gz', 'dwi_nlsam.nii.gz', 'bvals', 'bvecs', '-m', 'mask_crop.nii.gz', '-f', '--verbose', '-N', '1', '--noise_est', 'local_std', '--sh_order', '0', '--block_size', '2,2,2', '--save_sigma', 'sigma.nii.gz', '--cores', '1'),
-    ('nlsam_denoising', 'dwi_crop.nii.gz', 'dwi_nlsam.nii.gz', 'bvals', 'bvecs', '-m', 'mask_crop.nii.gz', '-f', '--verbose', '--sh_order', '0', '-N', '1', '--load_sigma', 'sigma.nii.gz', '--no_subsample', '--fix_implausible', '--no_clip_eta'),
-    ('nlsam_denoising', 'dwi_crop.nii.gz', 'dwi_nlsam.nii.gz', 'bvals', 'bvecs', '-m', 'mask_crop.nii.gz', '-f', '--verbose', '--sh_order', '0', '--noise_map', 'noise.nii.gz', '--noise_mask', 'pmask.nii.gz', '--use_f32'),
-    ('nlsam_denoising', 'dwi_crop.nii.gz', 'dwi_nlsam.nii.gz', 'bvals', 'bvecs', '-m', 'mask_crop.nii.gz', '-f', '--verbose', '--sh_order', '0', '-N', '1', '--load_sigma', 'sigma.nii.gz', '--no_stabilization', '--no_denoising'),
-    ('nlsam_denoising', 'dwi_crop.nii.gz', 'dwi_nlsam.nii.gz', 'bvals', 'bvecs', '-m', 'mask_crop.nii.gz', '-f', '--verbose', '--sh_order', '0', '-N', '1', '--load_sigma', 'sigma.nii.gz', '--no_denoising', '--no_clip_eta'),
-    ('nlsam_denoising', 'dwi_crop.nii.gz', 'dwi_nlsam.nii.gz', 'bvals', 'bvecs', '-m', 'mask_crop.nii.gz', '-f', '--verbose', '--sh_order', '0', '-N', '1', '--noise_est', 'local_std', '--no_denoising'),
-    ('nlsam_denoising', 'dwi_crop.nii.gz', 'dwi_nlsam.nii.gz', 'bvals', 'bvecs', '-m', 'mask_crop.nii.gz', '-f', '--verbose', '--sh_order', '0', '--noise_est', 'auto', '--no_denoising', '--cores', '4'),
-    ('nlsam_denoising', 'dwi_crop.nii.gz', 'dwi_nlsam.nii.gz', 'bvals', 'bvecs', '-m', 'mask_crop.nii.gz', '-f', '--verbose', '--sh_order', '0', '--no_denoising')
+    ('nlsam_denoising', 'dwi_crop.nii.gz', 'dwi_nlsam.nii.gz', 'bvals', 'bvecs', '-f', '-N', '1', '--noise_est', 'local_std', '--cores', '1', '-m', 'mask_crop.nii.gz', '-v'),
+    ('nlsam_denoising', 'dwi_crop.nii.gz', 'dwi_nlsam.nii.gz', 'bvals', 'bvecs', '-m', 'mask_crop.nii.gz', '-f', '-N', '1', '--noise_est', 'local_std', '--iterations', '5', '--verbose', '--save_sigma', 'sigma.nii.gz', '--log', 'log.txt'),
+    ('nlsam_denoising', 'dwi_crop.nii.gz', 'dwi_nlsam.nii.gz', 'bvals', 'bvecs', '-m', 'mask_crop.nii.gz', '-f', '--noise_est', 'auto', '--iterations', '5', '--verbose', '--save_sigma', 'sigma.nii.gz', '--save_N', 'N.nii.gz'),
+    ('nlsam_denoising', 'dwi_crop.nii.gz', 'dwi_nlsam.nii.gz', 'bvals', 'bvecs', '-m', 'mask_crop.nii.gz', '-f', '--verbose', '--b0_threshold', '10', '--noise_mask', 'pmask.nii.gz'),
+    ('nlsam_denoising', 'dwi_crop.nii', 'dwi_nlsam.nii', 'bvals', 'bvecs', '-m', 'mask_crop.nii', '-f', '--verbose', '-N', '1', '--no_stabilization', '--load_sigma', 'sigma.nii', '--is_symmetric', '--save_difference', 'diff.nii'),
+    ('nlsam_denoising', 'dwi_crop.nii', 'dwi_nlsam.nii', 'bvals', 'bvecs', '-m', 'mask_crop.nii', '-f', '--verbose', '--no_denoising', '--save_sigma', 'sigma.nii', '--save_stab', 'stab.nii', '--load_mhat', 'dwi_crop.nii', '--save_eta', 'eta.nii'),
+    ('nlsam_denoising', 'dwi_crop.nii.gz', 'dwi_nlsam.nii.gz', 'bvals', 'bvecs', '-m', 'mask_crop.nii.gz', '-f', '--verbose', '-N', '1', '--noise_est', 'local_std', '--block_size', '2,2,2', '--save_sigma', 'sigma.nii.gz', '--cores', '1'),
+    ('nlsam_denoising', 'dwi_crop.nii.gz', 'dwi_nlsam.nii.gz', 'bvals', 'bvecs', '-m', 'mask_crop.nii.gz', '-f', '--verbose', '-N', '1', '--load_sigma', 'sigma.nii.gz', '--no_subsample', '--no_clip_eta'),
+    ('nlsam_denoising', 'dwi_crop.nii.gz', 'dwi_nlsam.nii.gz', 'bvals', 'bvecs', '-m', 'mask_crop.nii.gz', '-f', '--verbose', '--noise_map', 'noise.nii.gz', '--noise_mask', 'pmask.nii.gz', '--use_f32'),
+    ('nlsam_denoising', 'dwi_crop.nii.gz', 'dwi_nlsam.nii.gz', 'bvals', 'bvecs', '-m', 'mask_crop.nii.gz', '-f', '--verbose', '-N', '1', '--load_sigma', 'sigma.nii.gz', '--no_stabilization', '--no_denoising'),
+    ('nlsam_denoising', 'dwi_crop.nii.gz', 'dwi_nlsam.nii.gz', 'bvals', 'bvecs', '-m', 'mask_crop.nii.gz', '-f', '--verbose', '-N', '1', '--load_sigma', 'sigma.nii.gz', '--no_denoising', '--no_clip_eta'),
+    ('nlsam_denoising', 'dwi_crop.nii.gz', 'dwi_nlsam.nii.gz', 'bvals', 'bvecs', '-m', 'mask_crop.nii.gz', '-f', '--verbose', '-N', '1', '--noise_est', 'local_std', '--no_denoising'),
+    ('nlsam_denoising', 'dwi_crop.nii.gz', 'dwi_nlsam.nii.gz', 'bvals', 'bvecs', '-m', 'mask_crop.nii.gz', '-f', '--verbose', '--noise_est', 'auto', '--no_denoising', '--cores', '4'),
+    ('nlsam_denoising', 'dwi_crop.nii.gz', 'dwi_nlsam.nii.gz', 'bvals', 'bvecs', '-m', 'mask_crop.nii.gz', '-f', '--verbose', '--no_denoising')
     ]
 
 cwd = Path(__file__).parent / Path("datasets")

pyproject.toml

@@ -2,13 +2,12 @@
 requires = ["Cython>=3.0",
             "scipy>=1.5",
             "numpy>=2.0",
-            "setuptools>=61.0",
-            "wheel"]
+            "setuptools>=78.0"]
 build-backend = "setuptools.build_meta"
 
 [project]
 name = "nlsam"
-version = '0.7.3'
+version = '1.0'
 authors = [{name = "Samuel St-Jean"}]
 description='Implementation of "Non Local Spatial and Angular Matching : Enabling higher spatial resolution diffusion MRI datasets through adaptive denoising"'
 readme = "README.md"
@@ -19,10 +18,10 @@ keywords = ["MRI", "diffusion", "dmri", "denoising"]
 dependencies = [
     'numpy>=1.21.3',
     'scipy>=1.5',
-    'nibabel>=2.0',
+    'nibabel>=4.0',
     'joblib>=1.3.0',
     'autodmri>=0.2.1',
-    'spams-bin>=2.6.2',
+    'spams-bin>=2.6.12',
     'tqdm>=4.56']
 
 [project.urls]