diff --git a/AFQ/_fixes.py b/AFQ/_fixes.py index d38488a8..f9177584 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() @@ -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): 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/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]]) 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/data.py b/AFQ/tasks/data.py index de21a667..0a91408f 100644 --- a/AFQ/tasks/data.py +++ b/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, 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, 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 diff --git a/AFQ/tasks/tissue.py b/AFQ/tasks/tissue.py index 61ebb06e..012f7bd2 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 @@ -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'"