Fixups2

.github/workflows/python-package.yml

@@ -14,7 +14,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13", "3.14", "3.14t"]
         os: [ubuntu-latest, macos-latest, windows-latest]
 
     steps:
@@ -27,8 +27,8 @@ jobs:
 
     - name: Install pytest
       run: |
-        python -m pip install --upgrade pip pytest
+        python -m pip install --upgrade pip pytest pytest-xdist
         pip install .
 
     - name: Run pytest
-      run: pytest .
+      run: pytest . -n auto

autodmri/estimator.py

@@ -1,6 +1,10 @@
 import numpy as np
 
-from scipy.ndimage.interpolation import zoom
+try:
+    from scipy.ndimage import zoom
+except ImportError:
+    from scipy.ndimage.interpolation import zoom
+
 from scipy.special import gammaincinv
 
 from autodmri.gamma import get_noise_distribution
@@ -103,13 +107,11 @@ def lambda_cdf(N, alpha_prob):
         return out
 
     def get_mask(data, N_min, N_max, phi, alpha_prob=0.05):
-        K = data.shape[-1]
-        sum_data2 = np.sum(data**2, axis=-1)
+        if alpha_prob <= 0 or alpha_prob >= 1:
+            raise ValueError(f'alpha_prob is a parameter in ]0, 1[ with default value 0.05, but has value {alpha_prob}')
 
-        data[data == 0] = np.nan
-        sum_data2 = np.nansum(data**2, axis=-1)
-        K = np.sum(np.isfinite(data), axis=-1)
-        data[np.isnan(data)] = 0
+        K = np.sum(data > 0, axis=-1)
+        sum_data2 = np.sum(data**2, axis=-1)
 
         lambda_minus = lambda_cdf(N_min*K, alpha_prob/2)
         lambda_plus = lambda_cdf(N_max*K, 1 - alpha_prob/2)
@@ -131,7 +133,9 @@ def get_mask(data, N_min, N_max, phi, alpha_prob=0.05):
         exclude_mask = np.zeros(data.shape[:-1], dtype=bool)
 
     # we don't know N, so guess parameters iteratively
-    data = data.astype(np.float64)  # prevent data**4 overflow
+    data = data.astype(np.float64) # prevent data**4 overflow
+    data = np.nan_to_num(data).clip(min=0) # cleanup in case we had bad data somewhere to prevent nan propagation
+
     sigma_prev = -1
     N_prev = -1
     sigma_init = median / np.sqrt(2 * lambda_cdf(N_max, 0.5))

autodmri/gamma.py

@@ -1,4 +1,7 @@
+from __future__ import annotations # Needed for | annotations on python 3.9
+
 import numpy as np
+import numpy.typing as npt
 
 from scipy.special import digamma, polygamma
 
@@ -10,6 +13,7 @@ def get_noise_distribution(data, method='moments'):
     -----
     data
         A numpy array of gamma distributed values
+
     method='moments' or method='maxlk'
         Use either the moments or maximum likelihood equations to estimate the parameters.
 
@@ -30,16 +34,13 @@ def get_noise_distribution(data, method='moments'):
     if method == 'moments':
         mdata2 = np.mean(data**2)
         mdata4 = np.mean(data**4)
-
-        p1 = mdata4 / mdata2
-        p2 = mdata2
-        sigma = np.sqrt(p1 - p2) / np.sqrt(2)
+        sigma = np.sqrt(mdata4 / mdata2 - mdata2) / np.sqrt(2)
     elif method == 'maxlk':
         sigma = maxlk_sigma(data)
     else:
         raise ValueError(f'Invalid method name {method}')
 
-    t = data**2 / (2*sigma**2)
+    t = 1/2 * (data / sigma)**2
 
     # Now compute N
     if method == 'moments':
@@ -53,18 +54,17 @@ def get_noise_distribution(data, method='moments'):
     return sigma, N
 
 
-def maxlk_sigma(m, xold=None, eps=1e-8, max_iter=100):
+def maxlk_sigma(m: npt.NDArray, xold: float | None = None, eps: float=1e-8, max_iter: int=100):
     '''Maximum likelihood equation to estimate sigma from gamma distributed values'''
 
-    sum_m2 = np.sum(m**2)
-    K = m.size
-    sum_log_m2 = np.sum(np.log(m**2))
+    mean_m2 = np.mean(m**2)
+    mean_log_m2 = 2 * np.mean(np.log(m))
 
     def f(sigma):
-        return digamma(sum_m2/(2*K*sigma**2)) - sum_log_m2/K + np.log(2*sigma**2)
+        return digamma(mean_m2 / (2*sigma**2)) - mean_log_m2 + np.log(2*sigma**2)
 
     def fprime(sigma):
-        return -sum_m2 * polygamma(1, sum_m2/(2*K*sigma**2)) / (K*sigma**3) + 2/sigma
+        return -mean_m2 * polygamma(1, mean_m2 / (2*sigma**2)) / sigma**3 + 2/sigma
 
     if xold is None:
         xold = m.std()
@@ -81,13 +81,13 @@ def fprime(sigma):
     return xnew
 
 
-def inv_digamma(y, eps=1e-8, max_iter=100):
+def inv_digamma(y: float, eps: float=1e-8, max_iter: int=100) -> float:
     '''Numerical inverse to the digamma function by root finding'''
 
     if y >= -2.22:
-        xold = np.exp(y) + 0.5
+        xold = np.exp(y) + 1/2
     else:
-        xold = -1 / (y - digamma(1))
+        xold = -1 / (y + np.euler_gamma)
 
     for _ in range(max_iter):
 

autodmri/scripts.py

@@ -192,7 +192,3 @@ def main():
     nib.Nifti1Image(sigma, aff).to_filename(args.sigma)
     nib.Nifti1Image(N, aff).to_filename(args.N)
     nib.Nifti1Image(mask, aff).to_filename(args.mask)
-
-
-if __name__ == "__main__":
-    main()

autodmri/tests/test_blocks.py

@@ -31,15 +31,12 @@ def test_extract_patches_strided():
     expected_views = expected_views_1D + expected_views_2D + expected_views_3D
     last_patches = last_patch_1D + last_patch_2D + last_patch_3D
 
-    for (image_shape, patch_size, patch_step, expected_view,
-         last_patch) in zip(image_shapes, patch_sizes, patch_steps,
-                            expected_views, last_patches):
+    for (image_shape, patch_size, patch_step, expected_view, last_patch) in zip(image_shapes, patch_sizes, patch_steps, expected_views, last_patches):
         image = np.arange(np.prod(image_shape)).reshape(image_shape)
         patches = extract_patches(image, patch_size, patch_step, flatten=False)
 
         ndim = len(image_shape)
 
         assert patches.shape[:ndim] == expected_view
-        last_patch_slices = tuple(slice(i, i + j, None) for i, j in
-                                  zip(last_patch, patch_size))
+        last_patch_slices = tuple(slice(i, i + j, None) for i, j in zip(last_patch, patch_size))
         assert (patches[(-1, None, None) * ndim] == image[last_patch_slices].squeeze()).all()

autodmri/tests/test_estimator.py

@@ -0,0 +1,51 @@
+from __future__ import annotations # Needed for | annotations on python 3.9
+
+import numpy as np
+import numpy.typing as npt
+import pytest
+
+from autodmri.estimator import estimate_from_dwis, estimate_from_nmaps
+from itertools import product
+
+
+sigma = 1, 10, 50, 100, 200
+N = 1, 4, 8, 12, 24
+methods = ['moments', 'maxlk']
+
+all_items = product(sigma, N, methods)
+
+@pytest.mark.parametrize('sigma, N, method', all_items)
+def test_estimators(sigma, N, method):
+    data = np.zeros([25,25,25,30])
+    data[10:20, 10:20, 10:20] = 1000
+    empty = np.zeros_like(data)
+
+    noisy = _make_noise(data, sigma, N)
+    noise_maps = _make_noise(empty, sigma, N)
+
+    dwis_sigma, dwis_N = estimate_from_dwis(noisy, method=method)
+    nmaps_sigma, nmaps_N = estimate_from_nmaps(noise_maps, method=method, return_mask=False, full=False)
+    nmaps_sigma_full, nmaps_N_full = estimate_from_nmaps(noise_maps, method=method, return_mask=False, full=True)
+
+    for (estimated_sigma, estimated_N) in zip([dwis_sigma, nmaps_sigma, nmaps_sigma_full],
+                                              [dwis_N, nmaps_N, nmaps_N_full]):
+        # less than 5% error on average estimate?
+        assert np.abs(sigma - estimated_sigma.mean()) / sigma < 0.05
+        assert np.abs(N - estimated_N.mean()) / N < 0.05
+
+
+def _make_noise(data: npt.NDArray, sigma: float, N: int, seed: int | None = None) -> np.ndarray:
+    size = data.shape[:-1]
+    out = np.zeros_like(data, dtype=np.float32)
+    n1 = np.zeros(size, dtype=np.float32)
+    n2 = np.zeros(size, dtype=np.float32)
+
+    rng = np.random.RandomState(seed)
+
+    for _ in range(N):
+        for i in range(data.shape[-1]):
+            n1[:] = rng.normal(0, sigma, size)
+            n2[:] = rng.normal(0, sigma, size)
+            out[..., i] += (data[..., i] + n1)**2 + n2**2
+
+    return np.sqrt(out)

autodmri/tests/test_script.py

@@ -4,11 +4,11 @@
 from pathlib import Path
 
 cwd = Path(__file__).parents[2] / Path("datasets")
-commands = [['autodmri_get_distribution', 'data_SENSE3_MB3_noisemap.nii.gz', 'sigma_maxlk_nmaps.nii.gz', 'N_maxlk_nmaps.nii.gz', 'mask_maxlk_nmaps.nii.gz', '-m', 'maxlk', '--noise_maps'],
-            ['autodmri_get_distribution', 'data_SENSE3_MB3_noisemap.nii.gz', 'sigma_nmaps.nii.gz', 'N_nmaps.nii.gz', 'mask_nmaps.nii.gz', '--noise_maps'],
+commands = [['autodmri_get_distribution', 'data_SENSE3_MB3_noisemap.nii.gz', 'sigma_maxlk_nmaps.nii.gz', 'N_maxlk_nmaps.nii.gz', 'mask_maxlk_nmaps.nii.gz', '-m', 'maxlk', '--noise_maps', '--force'],
+            ['autodmri_get_distribution', 'data_SENSE3_MB3_noisemap.nii.gz', 'sigma_nmaps.nii.gz', 'N_nmaps.nii.gz', 'mask_nmaps.nii.gz', '--noise_maps', '--force'],
             ['autodmri_get_distribution', 'data_SENSE3_MB3_noisemap.nii.gz', 'sigma_nmaps.nii.gz', 'N_nmaps.nii.gz', 'mask_nmaps.nii.gz', '--noise_maps', '-f', '--subsample'],
             ['autodmri_get_distribution', 'data_SENSE3_MB3_noisemap.nii.gz', 'sigma_nmaps.nii.gz', 'N_nmaps.nii.gz', 'mask_nmaps.nii.gz', '--noise_maps', '-f', '--fast_median', '-m', 'maxlk'],
-            ['autodmri_get_distribution', 'dwi_1_8.nii.gz', 'sigma.nii.gz', 'N.nii.gz', 'mask.nii.gz', '-v'],
+            ['autodmri_get_distribution', 'dwi_1_8.nii.gz', 'sigma.nii.gz', 'N.nii.gz', 'mask.nii.gz', '-v', '--force'],
             ['autodmri_get_distribution', 'dwi_1_8.nii.gz', 'sigma.nii.gz', 'N.nii.gz', 'mask.nii.gz', '-m', 'maxlk', '-f', '--ncores', '4'],
             ['autodmri_get_distribution', 'dwi_1_8.nii.gz', 'sigma_maxlk.nii.gz', 'N_maxlk.nii.gz', 'mask_maxlk.nii.gz', '-m', 'maxlk', '--size', '3', '-f', '-v', '--axis', '0']]