antspymm 1.5.5__tar.gz → 1.5.6__tar.gz

{antspymm-1.5.5 → antspymm-1.5.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: antspymm
-Version: 1.5.5
+Version: 1.5.6
 Summary: multi-channel/time-series medical image processing with antspyx
 Author-email: "Avants, Gosselin, Tustison, Reardon" <stnava@gmail.com>
 License: Apache-2.0

{antspymm-1.5.5 → antspymm-1.5.6}/antspymm/__init__.py

@@ -65,6 +65,8 @@ from .mm import mc_denoise
 from .mm import mc_reg
 from .mm import dti_reg
 from .mm import timeseries_reg
+from .mm import timeseries_transform
+from .mm import copy_spatial_metadata_from_3d_to_4d
 from .mm import concat_dewarp
 from .mm import mc_resample_image_to_target
 from .mm import trim_dti_mask
@@ -136,5 +138,7 @@ from .mm import segment_timeseries_by_bvalue
 from .mm import shorten_pymm_names
 from .mm import pet3d_summary
 from .mm import deformation_gradient_optimized
+from .mm import efficient_dwi_fit_voxelwise
+from .mm import generate_voxelwise_bvecs
 
 

{antspymm-1.5.5 → antspymm-1.5.6}/antspymm/mm.py

@@ -1794,6 +1794,89 @@ def merge_timeseries_data( img_LR, img_RL, allow_resample=True ):
         mimg.append( temp )
     return ants.list_to_ndimage( img_LR, mimg )
 
+def copy_spatial_metadata_from_3d_to_4d(spatial_img, timeseries_img):
+    """
+    Copy spatial metadata (origin, spacing, direction) from a 3D image to the
+    spatial dimensions (first 3) of a 4D image, preserving the 4th dimension's metadata.
+
+    Parameters
+    ----------
+    spatial_img : ants.ANTsImage
+        A 3D ANTsImage with the desired spatial metadata.
+    timeseries_img : ants.ANTsImage
+        A 4D ANTsImage to update.
+
+    Returns
+    -------
+    ants.ANTsImage
+        A 4D ANTsImage with updated spatial metadata.
+    """
+    if spatial_img.dimension != 3:
+        raise ValueError("spatial_img must be a 3D ANTsImage.")
+    if timeseries_img.dimension != 4:
+        raise ValueError("timeseries_img must be a 4D ANTsImage.")
+    # Get 3D metadata
+    spatial_origin = list(spatial_img.origin)
+    spatial_spacing = list(spatial_img.spacing)
+    spatial_direction = spatial_img.direction  # 3x3
+    # Get original 4D metadata
+    ts_spacing = list(timeseries_img.spacing)
+    ts_origin = list(timeseries_img.origin)
+    ts_direction = timeseries_img.direction  # 4x4
+    # Replace only the first 3 entries for origin and spacing
+    new_origin = spatial_origin + [ts_origin[3]]
+    new_spacing = spatial_spacing + [ts_spacing[3]]
+    # Replace top-left 3x3 block of direction matrix, preserve last row/column
+    new_direction = ts_direction.copy()
+    new_direction[:3, :3] = spatial_direction
+    # Create updated image
+    updated_img = ants.from_numpy(
+        timeseries_img.numpy(),
+        origin=new_origin,
+        spacing=new_spacing,
+        direction=new_direction
+    )
+    return updated_img
+
+def timeseries_transform(transform, image, reference, interpolation='linear'):
+    """
+    Apply a spatial transform to each 3D volume in a 4D time series image.
+
+    Parameters
+    ----------
+    transform : ants transform object
+        Path(s) to ANTs-compatible transform(s) to apply.
+    image : ants.ANTsImage
+        4D input image with shape (X, Y, Z, T).
+    reference : ants.ANTsImage
+        Reference image to match in space.
+    interpolation : str
+        Interpolation method: 'linear', 'nearestNeighbor', etc.
+
+    Returns
+    -------
+    ants.ANTsImage
+        4D transformed image.
+    """
+    if image.dimension != 4:
+        raise ValueError("Input image must be 4D (X, Y, Z, T).")
+    n_volumes = image.shape[3]
+    transformed_volumes = []
+    for t in range(n_volumes):
+        vol = ants.slice_image( image, 3, t )
+        transformed = ants.apply_ants_transform_to_image(
+            transform=transform,
+            image=vol,
+            reference=reference,
+            interpolation=interpolation
+        )
+        transformed_volumes.append(transformed.numpy())
+    # Stack along time axis and convert to ANTsImage
+    transformed_array = np.stack(transformed_volumes, axis=-1)
+    out_image = ants.from_numpy(transformed_array)
+    out_image = ants.copy_image_info(image, out_image)
+    out_image = copy_spatial_metadata_from_3d_to_4d(reference, out_image)
+    return out_image
 
 def timeseries_reg(
     image,
@@ -3932,6 +4015,135 @@ def efficient_dwi_fit(gtab, diffusion_model, imagein, maskin,
     return full_fit, FA_img, MD_img, RGB_img
 
 
+def efficient_dwi_fit_voxelwise(imagein, maskin, bvals, bvecs_5d, model_params=None,
+                                bvals_to_use=None, num_threads=1, verbose=True):
+    """
+    Voxel-wise diffusion model fitting with individual b-vectors per voxel.
+
+    Parameters
+    ----------
+    imagein : ants.ANTsImage
+        4D DWI image (X, Y, Z, N).
+    maskin : ants.ANTsImage
+        3D binary mask.
+    bvals : (N,) array-like
+        Common b-values across volumes.
+    bvecs_5d : (X, Y, Z, N, 3) ndarray
+        Voxel-specific b-vectors.
+    model_params : dict
+        Extra arguments for model.
+    bvals_to_use : list[int]
+        Subset of b-values to include.
+    num_threads : int
+        Number of threads to use.
+    verbose : bool
+        Whether to print status.
+
+    Returns
+    -------
+    FA_img : ants.ANTsImage
+        Fractional anisotropy.
+    MD_img : ants.ANTsImage
+        Mean diffusivity.
+    RGB_img : ants.ANTsImage
+        RGB FA image.
+    """
+    import numpy as np
+    import ants
+    import dipy.reconst.dti as dti
+    from dipy.core.gradients import gradient_table
+    from dipy.reconst.dti import fractional_anisotropy, color_fa, mean_diffusivity
+    from concurrent.futures import ThreadPoolExecutor
+    from tqdm import tqdm
+
+    model_params = model_params or {}
+    img = imagein.numpy()
+    mask = maskin.numpy().astype(bool)
+    X, Y, Z, N = img.shape
+
+    if bvals_to_use is not None:
+        sel = np.isin(bvals, bvals_to_use)
+        img = img[..., sel]
+        bvals = bvals[sel]
+        bvecs_5d = bvecs_5d[..., sel, :]
+
+    FA = np.zeros((X, Y, Z), dtype=np.float32)
+    MD = np.zeros((X, Y, Z), dtype=np.float32)
+    RGB = np.zeros((X, Y, Z, 3), dtype=np.float32)
+
+    def fit_voxel(ix, iy, iz):
+        if not mask[ix, iy, iz]:
+            return
+        sig = img[ix, iy, iz, :]
+        if np.all(sig == 0):
+            return
+        bv = bvecs_5d[ix, iy, iz, :, :]
+        gtab = gradient_table(bvals, bv)
+        try:
+            model = dti.TensorModel(gtab, **model_params)
+            fit = model.fit(sig)
+            evals = fit.evals
+            evecs = fit.evecs
+            FA[ix, iy, iz] = fractional_anisotropy(evals)
+            MD[ix, iy, iz] = mean_diffusivity(evals)
+            RGB[ix, iy, iz, :] = color_fa(FA[ix, iy, iz], evecs)
+        except Exception as e:
+            if verbose:
+                print(f"Voxel ({ix},{iy},{iz}) fit failed: {e}")
+
+    coords = np.argwhere(mask)
+    if verbose:
+        print(f"[INFO] Fitting {len(coords)} voxels using {num_threads} threads...")
+
+    if num_threads > 1:
+        with ThreadPoolExecutor(max_workers=num_threads) as executor:
+            list(tqdm(executor.map(lambda c: fit_voxel(*c), coords), total=len(coords)))
+    else:
+        for c in tqdm(coords):
+            fit_voxel(*c)
+
+    ref = ants.slice_image(imagein, axis=3, idx=0)
+    return (
+        ants.copy_image_info(ref, ants.from_numpy(FA)),
+        ants.copy_image_info(ref, ants.from_numpy(MD)),
+        ants.merge_channels([ants.copy_image_info(ref, ants.from_numpy(RGB[..., i])) for i in range(3)])
+    )
+
+
+def generate_voxelwise_bvecs(global_bvecs, voxel_rotations, transpose=False):
+    """
+    Generate voxel-wise b-vectors from a global bvec and voxel-wise rotation field.
+
+    Parameters
+    ----------
+    global_bvecs : ndarray of shape (N, 3)
+        Global diffusion gradient directions.
+    voxel_rotations : ndarray of shape (X, Y, Z, 3, 3)
+        3x3 rotation matrix for each voxel (can come from Jacobian of deformation field).
+    transpose : bool, optional
+        If True, transpose the rotation matrices before applying them to the b-vectors.
+
+
+    Returns
+    -------
+    bvecs_5d : ndarray of shape (X, Y, Z, N, 3)
+        Voxel-specific b-vectors.
+    """
+    X, Y, Z, _, _ = voxel_rotations.shape
+    N = global_bvecs.shape[0]
+    bvecs_5d = np.zeros((X, Y, Z, N, 3), dtype=np.float32)
+
+    for n in range(N):
+        bvec = global_bvecs[n]
+        for i in range(X):
+            for j in range(Y):
+                for k in range(Z):
+                    R = voxel_rotations[i, j, k]
+                    if transpose:
+                        R = R.T  # Use transpose if needed
+                    bvecs_5d[i, j, k, n, :] = R @ bvec
+    return bvecs_5d
+
 def dipy_dti_recon(
     image,
     bvalsfn,

{antspymm-1.5.5 → antspymm-1.5.6}/antspymm.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: antspymm
-Version: 1.5.5
+Version: 1.5.6
 Summary: multi-channel/time-series medical image processing with antspyx
 Author-email: "Avants, Gosselin, Tustison, Reardon" <stnava@gmail.com>
 License: Apache-2.0

{antspymm-1.5.5 → antspymm-1.5.6}/antspymm.egg-info/SOURCES.txt

@@ -24,7 +24,9 @@ docs/deepnbm.jpg
 docs/deformation_gradient_reo.py
 docs/describe_mm_data.R
 docs/dipy_dti_recon.py
+docs/dti_distortion_correction_voxelwise_varying_bvectors_example_WIP.py
 docs/dti_recon.py
+docs/dti_reconstruction_voxelwise_varying_bvectors_localint.py
 docs/dti_reg.py
 docs/dwi_rebasing.py
 docs/dwi_run.py
@@ -62,4 +64,5 @@ docs/ukbb_to_nrg_processing2.py
 docs/visualize_tractogram.py
 tests/test_loop.py
 tests/test_nrg_validation.py
-tests/test_reference_run.py
+tests/test_reference_run.py
+tests/voxelwise_bvec_dti_recon_test.py

antspymm-1.5.6/docs/dti_distortion_correction_voxelwise_varying_bvectors_example_WIP.py

@@ -0,0 +1,238 @@
+import numpy as np
+import ants
+import os
+from dipy.io.gradients import read_bvals_bvecs
+from scipy.stats import pearsonr
+import antspymm
+nt = 2
+
+import numpy as np
+from scipy.stats import pearsonr
+
+def read_bvecs_rotated(bvec_file, rotmat):
+    bvecs = np.loadtxt(bvec_file)
+    if bvecs.shape[0] != 3:
+        bvecs = bvecs.T
+    rotated_bvecs = (rotmat @ bvecs).T
+    return rotated_bvecs
+
+def mean_rgb_correlation(img1, img2, mask):
+    """
+    Compute the mean correlation between two RGB images.
+
+    Parameters
+    ----------
+    img1 : np.ndarray
+        First RGB image as a (H, W, 3) NumPy array.
+    img2 : np.ndarray
+        Second RGB image as a (H, W, 3) NumPy array.
+
+    Returns
+    -------
+    float
+        Mean Pearson correlation across the three RGB channels.
+    """
+    if img1.shape != img2.shape:
+        raise ValueError("Input images must have the same shape.")
+    correlations = []
+    img1c = ants.split_channels(img1)
+    img2c = ants.split_channels(img2)   
+    for c in range(3):  # R, G, B
+        x = extract_masked_values( img1c[c], mask)
+        y = extract_masked_values( img2c[c], mask)
+        if np.std(x) == 0 or np.std(y) == 0:
+            corr = 0.0  # Handle flat images
+        else:
+            corr, _ = pearsonr(x, y)
+        correlations.append(corr)
+    return np.mean(correlations)
+
+import numpy as np
+import ants
+
+def mean_rgb_mae(img1, img2, mask):
+    """
+    Compute the mean absolute error (MAE) between two RGB images.
+
+    Parameters
+    ----------
+    img1 : np.ndarray
+        First RGB image as a (H, W, 3) NumPy array.
+    img2 : np.ndarray
+        Second RGB image as a (H, W, 3) NumPy array.
+    mask : ants.ANTsImage
+        Binary mask defining valid pixels for error calculation.
+
+    Returns
+    -------
+    float
+        Mean absolute error across the three RGB channels.
+    """
+    if img1.shape != img2.shape:
+        raise ValueError("Input images must have the same shape.")
+    
+    mae_values = []
+    img1c = ants.split_channels(img1)
+    img2c = ants.split_channels(img2)   
+    
+    for c in range(3):  # R, G, B
+        x = extract_masked_values(img1c[c], mask)
+        y = extract_masked_values(img2c[c], mask)
+        mae = np.mean(np.abs(x - y))
+        mae_values.append(mae)
+    
+    return np.mean(mae_values)
+
+def extract_masked_values(image, mask):
+    return image.numpy()[mask.numpy() > 0]
+
+import numpy as np
+
+def verify_unit_bvecs(bvecs, tol=1e-5):
+    """
+    Verifies that each b-vector has unit norm within a tolerance.
+    
+    Parameters
+    ----------
+    bvecs : array-like, shape (N, 3)
+        Array of b-vectors, one per diffusion direction.
+    tol : float
+        Tolerance for unit norm.
+
+    Returns
+    -------
+    is_unit : np.ndarray, shape (N,)
+        Boolean array indicating if each b-vector is unit norm.
+    norms : np.ndarray, shape (N,)
+        Norms of each b-vector.
+    """
+    bvecs = np.asarray(bvecs)
+    norms = np.linalg.norm(bvecs, axis=1)
+    is_unit = np.abs(norms - 1) < tol
+    return is_unit, norms
+
+# def test_efficient_dwi_fit_voxelwise_distortion_correction():
+if True:
+    print("simple test for distortion_correction consistency of efficient_dwi_fit_voxelwise")
+    ex_path = os.path.expanduser( "~/.antspyt1w/" )
+    ex_path_mm = os.path.expanduser( "~/.antspymm/" )
+    #### Load in data ####
+    print("Load in subject data ...")
+    lrid = "I1499279_Anon_20210819142214_5"
+    rlid = "I1499337_Anon_20210819142214_6"
+    # Load paths
+    print("📁 Loading subject LR data...")
+    lrid = os.path.join(ex_path_mm, lrid )
+    img_LR_in = ants.image_read(lrid + '.nii.gz')
+    img_LR_in_avg = ants.get_average_of_timeseries( img_LR_in )
+    mask = img_LR_in_avg.get_mask()
+    bvals, bvecs = read_bvals_bvecs(lrid + '.bval', lrid + '.bvec')
+    bvecs = np.asarray(bvecs)
+    shape = img_LR_in.shape[:3]
+
+    print("📁 Loading subject RL data...")
+    rlid = os.path.join(ex_path_mm, rlid )
+    img_RL_in = ants.image_read(rlid + '.nii.gz')
+    bvalsRL, bvecsRL = read_bvals_bvecs(rlid + '.bval', rlid + '.bvec')
+
+    img_RL_in_avg = ants.get_average_of_timeseries( img_RL_in )
+    maskRL = img_RL_in_avg.get_mask()
+
+    print("🧠 Running baseline LR fit...")
+    bvecs_5d_orig = np.broadcast_to(bvecs, shape + bvecs.shape).copy()
+    if not "FA_orig" in globals():
+        FA_orig, MD_orig, RGB_orig = antspymm.efficient_dwi_fit_voxelwise(
+            imagein=img_LR_in,
+            maskin=mask,
+            bvals=bvals,
+            bvecs_5d=bvecs_5d_orig,
+            model_params={},
+            bvals_to_use=None,
+            num_threads=nt,
+            verbose=False
+        )
+
+    bvecs_5d_origRL = np.broadcast_to(np.asarray(bvecsRL), shape + bvecsRL.shape).copy()
+    if not "FA_origRL" in globals():
+        FA_origRL, MD_origRL, RGB_origRL = antspymm.efficient_dwi_fit_voxelwise(
+            imagein=img_RL_in,
+            maskin=maskRL,
+            bvals=bvalsRL,
+            bvecs_5d=bvecs_5d_origRL,
+            model_params={},
+            bvals_to_use=None,
+            num_threads=nt,
+            verbose=False
+        )
+
+    print("dist corr")
+    if not "mytx" in globals():
+        mytx = ants.registration( FA_orig, FA_origRL, 'SyNBold' )
+        mytx2 = ants.apply_transforms(FA_orig, FA_origRL, mytx['fwdtransforms'],
+            interpolator='linear', imagetype=0, compose='/tmp/comptx' )
+        print( mytx2 )
+
+    print("🔄 now with distortion correction...")
+    mydef = ants.image_read( mytx2 )
+    mywarp = ants.transform_from_displacement_field( mydef )
+    img_w = antspymm.timeseries_transform(mywarp, img_RL_in, reference=img_LR_in_avg)
+    mask_w = ants.apply_ants_transform_to_image(mywarp, maskRL, reference=img_LR_in_avg, interpolation='nearestNeighbor')
+    print("🧠 Running warped fit...")
+    if not "FA_w" in globals():
+        bvecsRL = np.asarray(bvecsRL)
+        mydefgrad = antspymm.deformation_gradient_optimized( mydef, 
+            to_rotation=False, to_inverse_rotation=True )
+        bvecsRLw = antspymm.generate_voxelwise_bvecs( bvecsRL, mydefgrad, transpose=False )
+        FA_w, MD_w, RGB_w = antspymm.efficient_dwi_fit_voxelwise(
+            imagein=img_w,
+            maskin=mask_w,
+            bvals=bvalsRL,
+            bvecs_5d=bvecsRLw,
+            model_params={},
+            bvals_to_use=None,
+            num_threads=nt,
+            verbose=False
+        )
+
+    if not "FA_w2" in globals():
+        bvecsRL = np.asarray(bvecsRL)
+        FA_w2, MD_w2, RGB_w2 = antspymm.efficient_dwi_fit_voxelwise(
+            imagein=img_w,
+            maskin=mask_w,
+            bvals=bvalsRL,
+            bvecs_5d=bvecs_5d_origRL,
+            model_params={},
+            bvals_to_use=None,
+            num_threads=nt,
+            verbose=False
+        )
+
+    fff = mean_rgb_correlation
+    print("📊 Comparing results...")
+    maskJoined = ants.threshold_image( mask + mask_w, 1.05, 2.0 )
+    maske=ants.iMath(maskJoined,'ME',3)
+    fa_corr = fff( RGB_orig, RGB_w, maske )
+    print(f"✅ FA correlation (original vs distortion corrected): {fa_corr:.4f}")
+
+    fa_corrX = fff( RGB_orig, RGB_w2, maske )
+    print(f"✅ FA correlation (original vs distortion corrected global recon): {fa_corrX:.4f}")
+
+    RGB_origRLc = ants.split_channels(RGB_origRL)
+    for c in range(3):
+        RGB_origRLc[c] = ants.apply_ants_transform_to_image(
+            mywarp, RGB_origRLc[c], 
+            reference=img_LR_in_avg, interpolation='linear'
+        )
+
+    fa_corrY = fff( RGB_orig, ants.merge_channels(RGB_origRLc), maske )
+    print(f"✅ FA correlation (original vs warped RGB global recon): {fa_corrY:.4f}")
+
+#    assert fa_corr > 0.80, "FA correlation too low"
+
+    print("🎉 Test passed: model is distortion-consistent.")
+
+# ants.image_write( FA_orig, '/tmp/xxx.nii.gz' )
+# ants.image_write( FA_w, '/tmp/yyy.nii.gz' )
+#
+# Example usage:
+# test_efficient_dwi_fit_voxelwise_distortion_correction()

antspymm-1.5.6/docs/dti_reconstruction_voxelwise_varying_bvectors_localint.py

@@ -0,0 +1,281 @@
+import numpy as np
+import ants
+import os
+from dipy.io.gradients import read_bvals_bvecs
+from scipy.stats import pearsonr
+import antspymm
+import matplotlib.pyplot as plt
+nt = 2
+# amount or rotation around x, y, z-axis
+degrotx = 15
+degroty = 15
+degrotz = 15
+
+
+def mean_rgb_correlation(img1, img2, mask):
+    """
+    Compute the mean correlation between two RGB images.
+
+    Parameters
+    ----------
+    img1 : np.ndarray
+        First RGB image as a (H, W, 3) NumPy array.
+    img2 : np.ndarray
+        Second RGB image as a (H, W, 3) NumPy array.
+
+    Returns
+    -------
+    float
+        Mean Pearson correlation across the three RGB channels.
+    """
+    if img1.shape != img2.shape:
+        raise ValueError("Input images must have the same shape.")
+    correlations = []
+    img1c = ants.split_channels(img1)
+    img2c = ants.split_channels(img2)   
+    for c in range(3):  # R, G, B
+        x = extract_masked_values( img1c[c], mask)
+        y = extract_masked_values( img2c[c], mask)
+        if np.std(x) == 0 or np.std(y) == 0:
+            corr = 0.0  # Handle flat images
+        else:
+            corr, _ = pearsonr(x, y)
+        correlations.append(corr)
+    return np.mean(correlations)
+
+
+def plot_correlation(x, y, xlabel="Image 1", ylabel="Image 2", title_prefix="Correlation", point_alpha=0.3):
+    """
+    Plot the correlation between two 1D arrays and display Pearson r.
+
+    Parameters
+    ----------
+    x : array-like
+        First set of values.
+    y : array-like
+        Second set of values.
+    xlabel : str
+        Label for the x-axis.
+    ylabel : str
+        Label for the y-axis.
+    title_prefix : str
+        Title prefix, will append r value.
+    point_alpha : float
+        Transparency of scatter plot points.
+
+    Returns
+    -------
+    float
+        Pearson correlation coefficient (r).
+    """
+    x = np.asarray(x)
+    y = np.asarray(y)
+
+    if x.shape != y.shape:
+        raise ValueError("Inputs must have the same shape.")
+
+    r, _ = pearsonr(x, y)
+
+    plt.figure(figsize=(6, 6))
+    plt.scatter(x, y, s=1, alpha=point_alpha, color='blue')
+    plt.xlabel(xlabel)
+    plt.ylabel(ylabel)
+    plt.title(f"{title_prefix} (r = {r:.3f})")
+    plt.plot([x.min(), x.max()], [x.min(), x.max()], 'r--', linewidth=1)  # Identity line
+    plt.grid(True)
+    plt.axis('equal')
+    plt.tight_layout()
+    plt.show()
+
+    return r
+
+def read_bvecs_rotated(bvec_file, rotmat):
+    bvecs = np.loadtxt(bvec_file)
+    if bvecs.shape[0] != 3:
+        bvecs = bvecs.T
+    rotated_bvecs = (rotmat @ bvecs).T
+    return rotated_bvecs
+
+def broadcast_bvecs_voxelwise(rotated_bvecs, shape):
+    return np.broadcast_to(rotated_bvecs, shape + rotated_bvecs.shape).copy()
+
+def extract_masked_values(image, mask):
+    return image.numpy()[mask.numpy() > 0]
+
+# def test_efficient_dwi_fit_voxelwise_rotation_consistency():
+if True:
+    print("simple test for rotation consistency of efficient_dwi_fit_voxelwise")
+    ex_path = os.path.expanduser( "~/.antspyt1w/" )
+    ex_path_mm = os.path.expanduser( "~/.antspymm/" )
+    JHU_atlas = ants.image_read( ex_path + 'JHU-ICBM-FA-1mm.nii.gz' ) # Read in JHU atlas
+    JHU_labels = ants.image_read( ex_path + 'JHU-ICBM-labels-1mm.nii.gz' ) # Read in JHU labels
+    #### Load in data ####
+    print("Load in subject data ...")
+    lrid = ex_path_mm + "I1499279_Anon_20210819142214_5"
+    rlid = ex_path_mm + "I1499337_Anon_20210819142214_6"
+    t1id = ex_path_mm + "t1_rand.nii.gz"
+    # Load paths
+    print("📁 Loading subject data...")
+    lrid = os.path.join(ex_path_mm, "I1499279_Anon_20210819142214_5")
+    img_LR_in = ants.image_read(lrid + '.nii.gz')
+    img_LR_in_avg = ants.get_average_of_timeseries( img_LR_in )
+    mask = img_LR_in_avg.get_mask()
+    
+    bvals, bvecs = read_bvals_bvecs(lrid + '.bval', lrid + '.bvec')
+    bvecs = np.asarray(bvecs)
+    shape = img_LR_in.shape[:3]
+
+    print("🧠 Running baseline (unrotated) fit...")
+    bvecs_5d_orig = np.broadcast_to(bvecs, shape + bvecs.shape).copy()
+    if not "FA_orig" in globals():
+        FA_orig, MD_orig, RGB_orig = antspymm.efficient_dwi_fit_voxelwise(
+            imagein=img_LR_in,
+            maskin=mask,
+            bvals=bvals,
+            bvecs_5d=bvecs_5d_orig,
+            model_params={},
+            bvals_to_use=None,
+            num_threads=nt,
+            verbose=False
+        )
+
+    print("🔄 Applying known rotation...")
+    # maxtrans = 10.0
+    # rotator = ants.contrib.RandomRotate3D((-maxtrans, maxtrans), reference=img_LR_in_avg)
+    rotator=ants.contrib.Rotate3D(rotation=(degrotx,degroty,degrotz), reference=img_LR_in_avg )
+    rotation = rotator.transform()
+    img_rotated_ref = ants.apply_ants_transform_to_image(rotation, img_LR_in_avg, reference=img_LR_in_avg)
+    img_rotated = antspymm.timeseries_transform(rotation, img_LR_in, reference=img_LR_in_avg)
+    mask_rotated = ants.apply_ants_transform_to_image(rotation, mask, reference=img_LR_in_avg, interpolation='nearestNeighbor')
+    # note: we apply the inverse rotation to the bvecs
+    # i.e. if we register A to B and get R then apply R_inv to bvecs
+    rotmat = ants.get_ants_transform_parameters(rotation.invert()).reshape((4, 3))[:3, :3]
+    bvecs_rotated = read_bvecs_rotated(lrid + '.bvec', rotmat)
+    bvecs_5d_rot = broadcast_bvecs_voxelwise(bvecs_rotated, shape)
+
+    print("🧠 Running rotated fit...")
+    if not "FA_rot" in globals():
+        FA_rot, MD_rot, RGB_rot = antspymm.efficient_dwi_fit_voxelwise(
+            imagein=img_rotated,
+            maskin=mask_rotated,
+            bvals=bvals,
+            bvecs_5d=bvecs_5d_rot,
+            model_params={},
+            bvals_to_use=None,
+            num_threads=nt,
+            verbose=False
+        )
+
+    FA_rot_back = ants.apply_ants_transform_to_image(
+        rotation.invert(),
+        FA_rot,
+        reference = img_LR_in_avg,
+        )
+    MD_rot_back = ants.apply_ants_transform_to_image(
+        rotation.invert(),
+        MD_rot,
+        reference = img_LR_in_avg,
+        )
+    print("📊 Comparing results...")
+    maske=ants.iMath(mask,'ME',2)
+    # smoothing simulates double interpolation
+    FA_origs = ants.smooth_image(FA_orig, 1.25)
+    MD_origs = ants.smooth_image(MD_orig, 1.25)
+    fa1 = extract_masked_values(FA_origs, maske)
+    fa2 = extract_masked_values(FA_rot_back, maske)
+    md1 = extract_masked_values(MD_origs, maske)
+    md2 = extract_masked_values(MD_rot_back, maske)
+
+    fa_corr, _ = pearsonr(fa1, fa2)
+    md_corr, _ = pearsonr(md1, md2)
+
+#    plot_correlation( fa1, fa2 )
+
+    print(f"✅ FA correlation (original vs rotated): {fa_corr:.4f}")
+    print(f"✅ MD correlation (original vs rotated): {md_corr:.4f}")
+    
+    assert fa_corr > 0.80, "FA correlation too low"
+    assert md_corr > 0.80, "MD correlation too low"
+
+    print("🎉 Test passed: model is rotation-consistent with voxelwise bvecs.")
+
+    print("This shows that the simulation is effective.")
+    print("Now use the simulated data to test the distortion correction...")
+#    ants.image_write( RGB_orig, '/tmp/temp0rgb.nii.gz' )
+#    ants.image_write( RGB_rot, '/tmp/temp1rgb.nii.gz' )
+
+
+    # now we have to map the img_rotated and its bvec_rotated partner 
+    # as we would with a generic distortion correction framework. 
+    # first --- write the transform to a file
+    rotationinv = rotation.invert()
+
+    import tempfile
+    with tempfile.TemporaryDirectory() as tempdir:
+        rotmatfile = os.path.join(tempdir, "rotation.mat")
+        compositefile = os.path.join(tempdir, "composite")
+
+        # Write the rotation transform
+        ants.write_transform(rotationinv, rotmatfile)
+
+        # Run registration if not already done
+        if "reg" not in globals():
+            reg = ants.registration(FA_orig, FA_rot, 'SyN', initial_transform=rotmatfile)
+
+        # Apply the transformation using a temporary composite path
+        comptx = ants.apply_transforms(
+            img_LR_in_avg,
+            img_rotated_ref,
+            reg['fwdtransforms'],
+            interpolator='linear',
+            compose=compositefile,
+            verbose=True
+        )
+        mydef=ants.image_read(comptx)
+        mydefgrad = antspymm.deformation_gradient_optimized( mydef, 
+                to_rotation=False, to_inverse_rotation=True )
+        bvecsRLw = antspymm.generate_voxelwise_bvecs( bvecs_rotated, mydefgrad, transpose=False )
+        mywarp = ants.transform_from_displacement_field( mydef )
+        img_rotated_avg = ants.get_average_of_timeseries(  img_rotated )
+        img_w = antspymm.timeseries_transform(mywarp, img_rotated, reference=img_rotated_avg )
+
+        correlations = []
+        labels = ["NoBvecReo", "BvecReo"]
+
+        RGB_origs = ants.smooth_image(RGB_orig, 1.0)
+
+        for label, bv in zip(labels, [bvecs_5d_rot, bvecsRLw]):
+            FA_w, MD_w, RGB_w = antspymm.efficient_dwi_fit_voxelwise(
+                imagein=img_w,
+                maskin=ants.get_mask(ants.get_average_of_timeseries(img_w)),
+                bvals=bvals,
+                bvecs_5d=bv,
+                model_params={},
+                bvals_to_use=None,
+                num_threads=nt,
+                verbose=False
+            )
+            # Optional: write RGB images to temporary files for inspection
+            rgb_orig_path = os.path.join(tempdir, "RGB_orig.nii.gz")
+            rgb_warped_path = os.path.join(tempdir, f"RGB_warped_{label}.nii.gz")
+
+            ants.image_write(RGB_origs, rgb_orig_path)
+            ants.image_write(RGB_w, rgb_warped_path)
+
+            print(f"📝 Saved RGB_orig to: {rgb_orig_path}")
+            print(f"📝 Saved RGB_warped ({label}) to: {rgb_warped_path}")
+
+            compmask = ants.get_mask(ants.get_average_of_timeseries(img_w))
+            compmask = ants.iMath(compmask, 'ME', 2)
+
+            corr = mean_rgb_correlation(RGB_origs, RGB_w, compmask)
+            print(f"{label:10s} correlation: {corr:.4f}")
+            correlations.append(corr)
+
+# Compare the two correlations
+if len(correlations) == 2:
+    print("\nComparison Result:")
+    if correlations[1] > correlations[0]:
+        print("✅ BvecReo gives higher correlation than NoBvecReo.")
+    else:
+        print("❌ NoBvecReo gives equal or higher correlation than BvecReo.")

{antspymm-1.5.5 → antspymm-1.5.6}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "antspymm"
-version = "1.5.5"
+version = "1.5.6"
 description = "multi-channel/time-series medical image processing with antspyx"
 readme = "README.md"
 requires-python = ">=3.9"

antspymm-1.5.6/tests/voxelwise_bvec_dti_recon_test.py

@@ -0,0 +1,118 @@
+import numpy as np
+from scipy.spatial.transform import Rotation as R
+import antspymm
+
+def test_generate_voxelwise_bvecs():
+    # Define synthetic bvecs (N = 3 directions)
+    bvecs = np.array([
+        [1, 0, 0],   # x-axis
+        [0, 1, 0],   # y-axis
+        [0, 0, 1],   # z-axis
+    ])
+    
+    # Define image shape (small synthetic 4D volume)
+    X, Y, Z = 2, 2, 2
+    N = bvecs.shape[0]
+    
+    # Define a known rotation (e.g., 90-degree rotation around z-axis)
+    rot = R.from_euler('z', 33, degrees=True).as_matrix()  # shape (3, 3)
+    
+    # Create a voxelwise rotation field (X, Y, Z, 3, 3), same rot everywhere
+    voxel_rotations = np.tile(rot, (X, Y, Z, 1, 1))
+    voxel_rotations = voxel_rotations.reshape(X, Y, Z, 3, 3)
+    
+    # Expected rotated bvecs
+    expected = np.dot(voxel_rotations[0, 0, 0], bvecs.T).T  # (N, 3)
+
+    # Call the function under test
+    bvecs_5d = antspymm.generate_voxelwise_bvecs(bvecs, voxel_rotations)  # shape (X, Y, Z, N, 3)
+    
+    # Check that all voxel outputs match expected result
+    for i in range(X):
+        for j in range(Y):
+            for k in range(Z):
+                actual = bvecs_5d[i, j, k]
+                assert np.allclose(actual, expected, atol=1e-6), \
+                    f"Mismatch at voxel {(i, j, k)}: {actual} vs {expected}"
+
+    print("✅ test_generate_voxelwise_bvecs passed!")
+
+# Run the test
+print("Running test for generate_voxelwise_bvecs...")
+test_generate_voxelwise_bvecs()
+
+
+import numpy as np
+import ants
+from scipy.spatial.transform import Rotation as R
+
+def generate_dummy_dwi_data(shape, n_volumes):
+    """
+    Generate synthetic DWI-like 4D data and a brain mask.
+    """
+    np.random.seed(42)
+    dwi = np.random.rand(*shape, n_volumes).astype(np.float32)
+    mask = np.ones(shape, dtype=np.uint8)
+    return dwi, mask
+
+def create_voxelwise_bvecs(shape, bvecs, rotation_matrix=None):
+    """
+    Create voxelwise (5D) bvecs array (X, Y, Z, N, 3), optionally rotated.
+    """
+    if rotation_matrix is not None:
+        rotated_bvecs = (rotation_matrix @ bvecs.T).T
+    else:
+        rotated_bvecs = bvecs
+
+    # Broadcast to 5D shape
+    bvecs_5d = np.broadcast_to(rotated_bvecs, shape + rotated_bvecs.shape)
+    return bvecs_5d.copy()
+
+def test_efficient_dwi_fit_voxelwise():
+
+    # Parameters
+    shape = (3, 3, 3)
+    n_vols = 6
+
+    # Synthetic data
+    dwi_data, mask_data = generate_dummy_dwi_data(shape, n_vols)
+    ants_dwi = ants.from_numpy(dwi_data)
+    ants_mask = ants.from_numpy(mask_data)
+
+    # Define bvals and bvecs
+    bvals = np.array([0, 1000, 1000, 1000, 1000, 1000])
+    bvecs = np.array([
+        [0, 0, 0],
+        [1, 0, 0],
+        [0, 1, 0],
+        [0, 0, 1],
+        [1/np.sqrt(2), 1/np.sqrt(2), 0],
+        [1/np.sqrt(2), 0, 1/np.sqrt(2)],
+    ])
+
+    # Optional: apply rotation
+    rotation_matrix = R.from_euler('z', 45, degrees=True).as_matrix()
+    bvecs_5d = create_voxelwise_bvecs(shape, bvecs, rotation_matrix)
+
+    # Call function under test
+    FA_img, MD_img, RGB_img = antspymm.efficient_dwi_fit_voxelwise(
+        imagein=ants_dwi,
+        maskin=ants_mask,
+        bvals=bvals,
+        bvecs_5d=bvecs_5d,
+        model_params={},
+        bvals_to_use=None,
+        num_threads=1,
+        verbose=False
+    )
+
+    # Tests
+    assert isinstance(FA_img, ants.ANTsImage), "FA_img should be an ANTsImage"
+    assert FA_img.shape == shape, f"FA image has shape {FA_img.shape}, expected {shape}"
+    assert np.all((FA_img.numpy() >= 0) & (FA_img.numpy() <= 1)), "FA values should be in [0, 1]"
+    assert MD_img.shape == shape, "MD image has incorrect shape"
+    
+    print("✅ test_efficient_dwi_fit_voxelwise passed!")
+
+# Run the test
+test_efficient_dwi_fit_voxelwise()