From 5d83ef2618d1f9b411cff0860821aceb77f33025 Mon Sep 17 00:00:00 2001 From: 36000 Date: Thu, 9 Jul 2026 13:22:44 -0700 Subject: [PATCH 1/9] WIP/ENH add t1w/logb0 metric --- AFQ/tasks/data.py | 47 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/AFQ/tasks/data.py b/AFQ/tasks/data.py index de21a667..d866419d 100644 --- a/AFQ/tasks/data.py +++ b/AFQ/tasks/data.py @@ -147,6 +147,52 @@ 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_r1_approximation=1e-2): + """ + full path to a nifti file containing the T1w over log(mean b0) + which is a proxy for R1 [1]_ + + Parameters + ---------- + min_b0_for_r1_approximation : float, optional + The minimum value of b0 to consider when doing the division. + This is to avoid dividing by small numbers. + Default: 1e-2 + + 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_r1_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 +1475,7 @@ def get_data_plan(kwargs): b0, b0_mask, t1w_over_b0, + t1w_over_log_b0, brain_mask, dti_fit, dki_fit, From 1842bdd39a483d385ca7dfacb1e4b276f6724e3b Mon Sep 17 00:00:00 2001 From: 36000 Date: Thu, 9 Jul 2026 19:39:09 -0700 Subject: [PATCH 2/9] set max numba n threads --- AFQ/tasks/structural.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/AFQ/tasks/structural.py b/AFQ/tasks/structural.py index 13c5ec76..aa85286d 100644 --- a/AFQ/tasks/structural.py +++ b/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 From dd167721cebc990e8b409fb0b9cfc5351fe2bed7 Mon Sep 17 00:00:00 2001 From: 36000 Date: Thu, 9 Jul 2026 19:42:35 -0700 Subject: [PATCH 3/9] better bound logb0 --- AFQ/tasks/data.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/AFQ/tasks/data.py b/AFQ/tasks/data.py index d866419d..0a91408f 100644 --- a/AFQ/tasks/data.py +++ b/AFQ/tasks/data.py @@ -150,17 +150,19 @@ def t1w_over_b0(structural_imap, b0, citations, min_b0_for_r1_approximation=1e-2 @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_r1_approximation=1e-2): +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_r1_approximation : float, optional + 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: 1e-2 + Default: 1.05 References ---------- @@ -176,7 +178,7 @@ def t1w_over_log_b0(structural_imap, b0, citations, min_b0_for_r1_approximation= resampled_t1 = resample(t1_img, b0_img) b0_data = b0_img.get_fdata() - mask = b0_data >= min_b0_for_r1_approximation + 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) From 08ba023ca984f905a38571e46f48290764da6f9c Mon Sep 17 00:00:00 2001 From: 36000 Date: Tue, 14 Jul 2026 12:41:51 -0700 Subject: [PATCH 4/9] Debug gaussian weights error --- AFQ/_fixes.py | 20 +++++++++++++++----- AFQ/tasks/segmentation.py | 4 +++- 2 files changed, 18 insertions(+), 6 deletions(-) diff --git a/AFQ/_fixes.py b/AFQ/_fixes.py index d38488a8..7b8e790e 100644 --- a/AFQ/_fixes.py +++ b/AFQ/_fixes.py @@ -340,25 +340,35 @@ 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"): 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) 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: + 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) + + weight_sum = np.sum(w, axis=0) + if not np.allclose(weight_sum, np.ones(n_points)): + 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"bundle sls: {len(bundle)}, weight rows: {w.shape[0]}") + 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()}") + return _weighting_failed() return w diff --git a/AFQ/tasks/segmentation.py b/AFQ/tasks/segmentation.py index 709f85c9..f6dbd2a6 100644 --- a/AFQ/tasks/segmentation.py +++ b/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, From 7861b4071bd8d38e63f51cc06c73630661fc8bf8 Mon Sep 17 00:00:00 2001 From: 36000 Date: Wed, 15 Jul 2026 13:13:40 -0700 Subject: [PATCH 5/9] try float64 --- AFQ/_fixes.py | 26 +++++++++++++++----------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/AFQ/_fixes.py b/AFQ/_fixes.py index 7b8e790e..44e4e9bd 100644 --- a/AFQ/_fixes.py +++ b/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() @@ -347,29 +347,33 @@ def _weighting_failed(): 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.isfinite(w_inv)] = 0 - 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) + denom = np.sum(w_inv, axis=0, dtype=np.float64) + w = np.divide( + w_inv, denom, out=np.zeros_like(w_inv), where=denom != 0, dtype=np.float64 + ) + + 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) + np.divide(w, final_sums, out=w, where=final_sums != 0, dtype=np.float64) weight_sum = np.sum(w, axis=0) - if not np.allclose(weight_sum, np.ones(n_points)): + 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"bundle sls: {len(bundle)}, weight rows: {w.shape[0]}") + 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 + return w def make_mp4(show_m, out_path, n_frames=720, az_ang=-0.5, fps=30, crf=35, verbose=True): From bd3b11e0be7839099339067d973db8e1856dcbe7 Mon Sep 17 00:00:00 2001 From: 36000 Date: Wed, 15 Jul 2026 14:24:40 -0700 Subject: [PATCH 6/9] fix brainchop --- AFQ/nn/brainchop.py | 2 +- AFQ/nn/utils.py | 39 +++++++++++++++++++++++++++++++++++++-- AFQ/tasks/tissue.py | 3 ++- 3 files changed, 40 insertions(+), 4 deletions(-) diff --git a/AFQ/nn/brainchop.py b/AFQ/nn/brainchop.py index a6edc0b7..91c96bf7 100644 --- a/AFQ/nn/brainchop.py +++ b/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, ...] diff --git a/AFQ/nn/utils.py b/AFQ/nn/utils.py index 7a01fb96..07f4dd65 100644 --- a/AFQ/nn/utils.py +++ b/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 diff --git a/AFQ/tasks/tissue.py b/AFQ/tasks/tissue.py index 61ebb06e..74476097 100644 --- a/AFQ/tasks/tissue.py +++ b/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 @@ -104,7 +105,7 @@ 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( SubCortexParcellation=structural_imap["t1_subcortex"], From 129414a9a224d0d97b962caa324ffc74c2c532cb Mon Sep 17 00:00:00 2001 From: 36000 Date: Wed, 15 Jul 2026 20:11:32 -0700 Subject: [PATCH 7/9] return float64 weights --- AFQ/_fixes.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/AFQ/_fixes.py b/AFQ/_fixes.py index 44e4e9bd..e1d71cce 100644 --- a/AFQ/_fixes.py +++ b/AFQ/_fixes.py @@ -353,17 +353,21 @@ def _weighting_failed(): denom = np.sum(w_inv, axis=0, dtype=np.float64) w = np.divide( - w_inv, denom, out=np.zeros_like(w_inv), where=denom != 0, dtype=np.float64 + w_inv, + denom, + out=np.zeros_like(w_inv, dtype=np.float64), + where=denom != 0, + dtype=np.float64, ) 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) + 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) + 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}") From 5a26f3f881548279dee1ae003b08b3f6962fd903 Mon Sep 17 00:00:00 2001 From: 36000 Date: Thu, 16 Jul 2026 16:18:15 -0700 Subject: [PATCH 8/9] combination model --- AFQ/tasks/tissue.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/AFQ/tasks/tissue.py b/AFQ/tasks/tissue.py index 74476097..012f7bd2 100644 --- a/AFQ/tasks/tissue.py +++ b/AFQ/tasks/tissue.py @@ -86,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. " @@ -107,11 +107,19 @@ def pve_internal(structural_imap, pve="synthseg"): # Use predictions from both to get final estimates 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'" From 631c7948ea5d6d1d5f0f124c36f83120f5e6db35 Mon Sep 17 00:00:00 2001 From: 36000 Date: Tue, 21 Jul 2026 19:33:16 -0700 Subject: [PATCH 9/9] add nn test --- AFQ/_fixes.py | 2 +- AFQ/nn/tests/test_nn.py | 21 +++++++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) create mode 100644 AFQ/nn/tests/test_nn.py diff --git a/AFQ/_fixes.py b/AFQ/_fixes.py index e1d71cce..f9177584 100644 --- a/AFQ/_fixes.py +++ b/AFQ/_fixes.py @@ -274,7 +274,7 @@ def gaussian_weights( def _weighting_failed(): logger.warning( ( - f"Not enough or too many streamlines (n_sls: {n_sls})" + f"Not enough or too many streamlines (n_sls: {n_sls}) " "for weight calculation, " "weighting everything evenly" ) diff --git a/AFQ/nn/tests/test_nn.py b/AFQ/nn/tests/test_nn.py new file mode 100644 index 00000000..c8b32be4 --- /dev/null +++ b/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]])