[FIX/ENH] VOF HBN changes

AFQ/_fixes.py

@@ -272,17 +272,17 @@ def gaussian_weights(
     n_sls, n_nodes, _ = sls.shape
 
     def _weighting_failed():
-        weights = np.ones((n_sls, n_nodes))
         logger.warning(
             (
-                "Not enough streamlines for weight calculation, "
+                f"Not enough or too many streamlines (n_sls: {n_sls}) "
+                "for weight calculation, "
                 "weighting everything evenly"
             )
         )
         if return_mahalanobis:
             return np.full((n_sls, n_nodes), np.nan)
         else:
-            return weights / np.sum(weights, 0)
+            return np.full((n_sls, n_nodes), 1.0 / n_sls)
 
     if n_sls < 15:  # Cov^-1 unstable under this amount
         return _weighting_failed()
@@ -340,26 +340,44 @@ def _weighting_failed():
                 cov = eigenvectors @ np.diag(eigenvalues) @ eigenvectors.T
 
             if np.any(cov > 0):
-                weights.ravel()[mask] = np.sqrt(
-                    np.einsum("ij,jk,ik->i", diff, pinvh(cov), diff)
-                )
+                m = np.einsum("ij,jk,ik->i", diff, pinvh(cov), diff)
+                np.clip(m, 0, None, out=m)
+                weights.ravel()[mask] = np.sqrt(m)
 
     if return_mahalanobis:
         return weights
 
-    with np.errstate(divide="ignore"):
+    with np.errstate(divide="ignore", invalid="ignore"):
         w_inv = 1.0 / weights
-    w_inv[np.isinf(w_inv)] = 0
+    w_inv[~np.isfinite(w_inv)] = 0
+
+    denom = np.sum(w_inv, axis=0, dtype=np.float64)
+    w = np.divide(
+        w_inv,
+        denom,
+        out=np.zeros_like(w_inv, dtype=np.float64),
+        where=denom != 0,
+        dtype=np.float64,
+    )
 
-    denom = np.sum(w_inv, axis=0)
-    w = np.divide(w_inv, denom, out=np.zeros_like(w_inv), where=denom != 0)
-    col_sums = w.sum(axis=0)
-    if np.max(np.abs(col_sums - 1)) > 1e-3:
+    col_sums = w.sum(axis=0, dtype=np.float64)
+    if not np.all(np.abs(col_sums - 1) <= 1e-3):
         return _weighting_failed()
     else:
-        final_sums = w.sum(axis=0, keepdims=True)
-        np.divide(w, final_sums, out=w, where=final_sums != 0)
-        return w
+        final_sums = w.sum(axis=0, keepdims=True, dtype=np.float64)
+        np.divide(w, final_sums, out=w, where=final_sums != 0, dtype=np.float64)
+
+        weight_sum = np.sum(w, axis=0, dtype=np.float64)
+        if not np.allclose(weight_sum, np.ones(n_nodes)):
+            bad = ~np.isclose(weight_sum, 1.0, rtol=1e-5, atol=1e-8)
+            logger.warning(f"weights not normalized: {w.shape}, {w.dtype}")
+            logger.warning(f"n_sls: {n_sls}")
+            logger.warning(f"offending nodes: {np.where(bad)[0]}")
+            logger.warning(f"their sums: {weight_sum[bad]}")
+            logger.warning(f"any nan: {np.isnan(w).any()}")
+            del w, w_inv
+            return _weighting_failed()
+    return w
 
 
 def make_mp4(show_m, out_path, n_frames=720, az_ang=-0.5, fps=30, crf=35, verbose=True):

AFQ/nn/brainchop.py

@@ -41,7 +41,7 @@ def run_brainchop(ort, t1_img, model_name, onnx_kwargs):
         https://doi.org/10.21105/joss.05098
     """
     model = _get_model(model_name)
-    t1_data, conformed_affine = prepare_t1_for_nn(t1_img)
+    t1_data, conformed_affine = prepare_t1_for_nn(t1_img, orientation="LIA")
 
     image = t1_data.astype(np.float32)[None, None, ...]
 

AFQ/nn/tests/test_nn.py

@@ -0,0 +1,21 @@
+import numpy as np
+
+from AFQ.nn.utils import merge_PVEs
+
+
+def test_wm_is_elementwise_maximum():
+    a = np.array([[0.5, 0.3, 0.2], [0.1, 0.1, 0.8]])
+    b = np.array([[0.1, 0.2, 0.7], [0.4, 0.3, 0.3]])
+    np.testing.assert_allclose(merge_PVEs(a, b)[..., 2], [0.7, 0.8])
+
+
+def test_larger_gm_source_wins():
+    a = np.array([[0.6, 0.2, 0.2]])
+    b = np.array([[0.2, 0.5, 0.3]])
+    np.testing.assert_allclose(merge_PVEs(a, b), b)
+
+
+def test_tie_prefers_a_and_renormalizes():
+    a = np.array([[0.6, 0.2, 0.2]])
+    b = np.array([[0.4, 0.2, 0.4]])
+    np.testing.assert_allclose(merge_PVEs(a, b), [[0.45, 0.15, 0.4]])

AFQ/nn/utils.py

@@ -23,14 +23,14 @@ def crop_to_nonzero(img):
     return nib.Nifti1Image(cropped_data, new_affine)
 
 
-def prepare_t1_for_nn(t1_img):
+def prepare_t1_for_nn(t1_img, orientation="RAS"):
     t1_img_cropped = crop_to_nonzero(t1_img)
 
     t1_img_conformed = nbp.conform(
         t1_img_cropped,
         out_shape=(256, 256, 256),
         voxel_size=(1.0, 1.0, 1.0),
-        orientation="RAS",
+        orientation=orientation,
         order=1,
     )
 
@@ -49,3 +49,38 @@ def resample_output(output, conformed_affine, t1_img):
         t1_img,
         order=0,
     )
+
+
+def merge_PVEs(PVE_a, PVE_b):
+    """
+    Merge two PVE (Partial Volume Estimation) images, PVE_a and PVE_b,
+    into a single PVE image. The merging is done by taking the maximum
+    WM estimate from both images, then the maximum GM estimate, and
+    finally normalizing.
+    """
+
+    csf_a, gm_a, wm_a = PVE_a[..., 0], PVE_a[..., 1], PVE_a[..., 2]
+    csf_b, gm_b, wm_b = PVE_b[..., 0], PVE_b[..., 1], PVE_b[..., 2]
+
+    new_wm = np.maximum(wm_a, wm_b)
+
+    # 2. Choose the source (a or b) with the larger GM estimate, per voxel
+    use_a = gm_a >= gm_b
+    chosen_csf = np.where(use_a, csf_a, csf_b)
+    chosen_gm = np.where(use_a, gm_a, gm_b)
+
+    # 3. Normalize chosen csf+gm to fill the remaining mass (1 - new_wm)
+    remaining = 1.0 - new_wm
+    chosen_sum = chosen_csf + chosen_gm
+
+    # avoid divide-by-zero where chosen_sum is 0 (e.g. pure-WM voxel)
+    scale = np.divide(
+        remaining, chosen_sum, out=np.zeros_like(chosen_sum), where=chosen_sum > 0
+    )
+
+    new_csf = chosen_csf * scale
+    new_gm = chosen_gm * scale
+
+    PVE = np.stack([new_csf, new_gm, new_wm], axis=-1)
+
+    return PVE

AFQ/tasks/data.py

@@ -147,6 +147,54 @@ def t1w_over_b0(structural_imap, b0, citations, min_b0_for_r1_approximation=1e-2
     return data, meta
 
 
+@immlib.calc("t1w_over_log_b0")
+@as_file("_desc-T1wOverLogB0.nii.gz")
+@as_img
+def t1w_over_log_b0(
+    structural_imap, b0, citations, min_b0_for_logr1_approximation=1.05
+):
+    """
+    full path to a nifti file containing the T1w over log(mean b0)
+    which is a proxy for R1 [1]_
+
+    Parameters
+    ----------
+    min_b0_for_logr1_approximation : float, optional
+        The minimum value of b0 to consider when doing the division.
+        This is to avoid dividing by small numbers.
+        Default: 1.05
+
+    References
+    ----------
+    .. [1] Moskovich, S., Shtangel, O., & Mezer, A. (2024).
+             Approximating R1 and R2: A quantitative approach to
+             clinical weighted MRI. Hum. Brain Mapp., 45(18), e70102.
+
+    """
+    citations.add("Moskovich2024-qd")
+
+    t1_img = nib.load(structural_imap["t1_masked"])
+    b0_img = nib.load(b0)
+    resampled_t1 = resample(t1_img, b0_img)
+
+    b0_data = b0_img.get_fdata()
+    mask = b0_data >= min_b0_for_logr1_approximation
+
+    log_b0 = np.full_like(b0_data, np.nan, dtype=float)
+    np.log(b0_data, out=log_b0, where=mask)
+
+    data = np.zeros_like(resampled_t1.get_fdata(), dtype=float)
+    np.divide(
+        resampled_t1.get_fdata(),
+        log_b0,
+        out=data,
+        where=mask,
+    )
+
+    meta = dict(T1w=structural_imap["t1_masked"], b0=b0)
+    return data, meta
+
+
 @immlib.calc("masked_b0")
 @as_file("_desc-masked_b0ref.nii.gz")
 @as_img
@@ -1429,6 +1477,7 @@ def get_data_plan(kwargs):
             b0,
             b0_mask,
             t1w_over_b0,
+            t1w_over_log_b0,
             brain_mask,
             dti_fit,
             dki_fit,

AFQ/tasks/segmentation.py

@@ -338,7 +338,9 @@ def _median_weight(bundle):
                         "low variance in the scalar data."
                     )
                 )
-                this_prof_weights = np.ones_like(this_prof_weights)
+                this_prof_weights = np.ones_like(this_prof_weights) / len(
+                    this_prof_weights
+                )
             this_profile[ii] = afq_profile(
                 scalar_data,
                 this_sl,

AFQ/tasks/structural.py

@@ -34,7 +34,7 @@ def configure_ncpus_nthreads(numba_n_threads=None, low_memory=False):
         Default: False
     """
     if numba_n_threads is None:
-        numba_n_threads = max(get_num_threads() - 1, 1)
+        numba_n_threads = min(max(get_num_threads() - 1, 1), 32)
 
     return numba_n_threads, low_memory
 

AFQ/tasks/tissue.py

@@ -25,6 +25,7 @@
 from AFQ.nn.brainchop import pve_from_subcortex
 from AFQ.nn.multiaxial import extract_pve
 from AFQ.nn.synthseg import pve_from_synthseg
+from AFQ.nn.utils import merge_PVEs
 from AFQ.tasks.decorators import as_file, as_fit_deriv, as_img
 from AFQ.tasks.utils import with_name
 
@@ -85,7 +86,7 @@ def pve_internal(structural_imap, pve="synthseg"):
                 SynthsegParcellation=structural_imap["synthseg_model"],
                 labels=["csf", "gm", "wm"],
             )
-        elif pve == "multiaxial+brainchop":
+        elif pve == "multiaxial+brainchop" or pve == "multiaxial+brainchop+synthseg":
             logger.warning(
                 (
                     "Using MultiAxial+BrainChop for PVE estimation. "
@@ -104,13 +105,21 @@ def pve_internal(structural_imap, pve="synthseg"):
             PVE_multiaxial = extract_pve(mx_model_img).get_fdata().astype(np.float32)
 
             # Use predictions from both to get final estimates
-            PVE = (PVE_brainchop + PVE_multiaxial) / 2
+            PVE = merge_PVEs(PVE_brainchop, PVE_multiaxial)
 
-            return nib.Nifti1Image(PVE, t1_subcortex_img.affine), dict(
+            meta = dict(
                 SubCortexParcellation=structural_imap["t1_subcortex"],
                 MultiAxialSegmentation=structural_imap["mx_model"],
                 labels=["csf", "gm", "wm"],
             )
+            if pve == "multiaxial+brainchop+synthseg":
+                synthseg_seg = nib.load(structural_imap["synthseg_model"])
+                PVE_synthseg = pve_from_synthseg(synthseg_seg.get_fdata())
+                PVE = merge_PVEs(PVE, PVE_synthseg)
+
+                meta["SynthsegParcellation"] = structural_imap["synthseg_model"]
+
+            return nib.Nifti1Image(PVE, t1_subcortex_img.affine), meta
 
     raise ValueError(
         "pve must be a PVEImage, PVEImages, 'synthseg', or 'multiaxial+brainchop'"