"""The brain Space's asset (disco-space#8 step 4): a MASiVar scan's multi-shell DWI, reconstructed by dmipy-fit *alone* (three-tissue responses, Dhollander 2016; multi-shell multi-tissue CSD, Jeurissen 2014, ``CsdCvxpyOptimizer``) into a WM FOD field and WM/GM/CSF fractions, plus the parcellation of ``build_parcellation.py``, written to an asset directory the Space's ``Brain`` source reads. The FOD and the fractions are dmipy-fit's ``MultiCompartmentSphericalHarmonicsModel`` end to end; MRtrix appears nowhere. dipy is used by this offline tool for what is not reconstruction: the median-Otsu brain mask and the peak count of the sanity check (and inside dmipy-fit's own tissue-response estimator until dmipy-fit#39 replaces it). Nothing here runs in the Space. OMP_NUM_THREADS=4 MKL_NUM_THREADS=4 python tools/build_brain_asset.py --dwi-dir DIR --parcellation-dir PARC_DIR \\ --out-dir OUT [--n-test-voxels 200] [--solver csd_cvxpy] ``--dwi-dir`` holds the PreQual-preprocessed single-shell NIfTI + bval + bvec triples of one scan (the derivative keeps every shell separate; this script is what merges them into one multi-shell volume). ``--parcellation-dir`` is ``build_parcellation.py``'s output (``labels.npy``, ``stop_mask.npy``, ``regions.json``); it is resampled here onto the DWI's own grid only if its own grid differs (it does not, by construction: both scripts are handed the same mean b=0 as the reference grid). With ``--n-test-voxels`` set (no ``--out-dir`` write of the whole-brain arrays), fits only a random subset of masked voxels and reports the measured per-voxel time and the whole-brain extrapolation, for a dry run before the (long) full fit. """ from __future__ import annotations import argparse import glob import json import os import subprocess import time import numpy as np SH_ORDER = 8 B0_THRESHOLD_S_MM2 = 50.0 # a row below this b-value (s/mm^2) is a b = 0 measurement RESPONSE_ERODE = 3 # erosion passes of the brain mask the responses are selected from RESPONSES = ("three_tissue_response_dhollander16 (mask eroded 3 passes, as MRtrix dwi2response dhollander -erode 3; " "white-matter selection tournier13)") def merge_shells(dwi_dir): """The PreQual per-shell NIfTI + bval + bvec triples of ``dwi_dir`` (one file per shell, each already carrying its own b = 0 volume), concatenated into one multi-shell acquisition: ``(data, bvals, bvecs, affine, shells)``. ``data`` is ``(X, Y, Z, N)`` float32, ``bvals`` (s/mm^2) and ``bvecs`` (image-frame unit vectors, BIDS/FSL convention) are ``(N,)`` / ``(N, 3)``; ``shells`` lists ``(b_value, n_directions)`` as found, ascending.""" import nibabel as nib niis = sorted(glob.glob(os.path.join(dwi_dir, "*_dwi.nii.gz"))) if not niis: raise FileNotFoundError(f"no *_dwi.nii.gz in {dwi_dir}") datas, bvals, bvecs, shells, affine, ref_shape = [], [], [], [], None, None for f in niis: stem = f[: -len(".nii.gz")] im = nib.load(f) if affine is None: affine, ref_shape = im.affine, im.shape[:3] elif not np.allclose(im.affine, affine, atol=1e-4) or im.shape[:3] != ref_shape: raise ValueError(f"{f} is not on the other shells' grid") d = np.asarray(im.dataobj, dtype=np.float32) bv = np.loadtxt(stem + ".bval") vec = np.loadtxt(stem + ".bvec").T datas.append(d) bvals.append(bv) bvecs.append(vec) n_dwi = int((bv >= B0_THRESHOLD_S_MM2).sum()) if n_dwi: shells.append((float(np.median(bv[bv >= B0_THRESHOLD_S_MM2])), n_dwi)) data = np.concatenate(datas, axis=-1) bvals = np.concatenate(bvals) bvecs = np.concatenate(bvecs, axis=0) shells = sorted(shells) return data, bvals, bvecs, affine, shells def mean_b0(data, bvals): """The mean of every b = 0 volume (``bvals < B0_THRESHOLD_S_MM2``), ``(X, Y, Z)`` float32.""" b0 = bvals < B0_THRESHOLD_S_MM2 if not np.any(b0): raise ValueError("no b = 0 measurement in this acquisition") return data[..., b0].mean(axis=-1).astype(np.float32) def brain_mask(mean_b0_volume): """The brain mask from the mean b = 0 image (dipy's median-Otsu, the standard tool; not dmipy-fit's job).""" from dipy.segment.mask import median_otsu _, mask = median_otsu(mean_b0_volume, median_radius=4, numpass=4) return mask.astype(bool) def fit_msmt_csd(data, bvals, bvecs, mask, sh_order=SH_ORDER, solver="csd_cvxpy", voxel_positions=None): """dmipy-fit end to end: :func:`three_tissue_response_dhollander16` estimates the WM / GM / CSF response kernels from the data itself (selected from ``mask`` eroded by ``RESPONSE_ERODE`` passes, the WM by tournier13), then ``MultiCompartmentSphericalHarmonicsModel`` (``S0_tissue_responses`` set, volume fractions free) fits every voxel of ``mask`` with ``solver`` (``CsdCvxpyOptimizer`` for ``'csd_cvxpy'``). If ``voxel_positions`` is given (a ``(mask.sum(),)``-shaped boolean subset, for a dry run) only those voxels are fit; the rest of ``mask`` is still used to estimate the responses. Always fits serially (``use_parallel_processing=False``): dmipy-fit's parallel path forks a ``ProcessPoolExecutor``, which JAX warns is unsafe once JAX is imported. Thread-level parallelism (BLAS/OSQP) is capped by ``OMP_NUM_THREADS``/``MKL_NUM_THREADS`` in the environment, per this lab's shared-box rule. Returns ``(sh, fractions, S0_responses, response_models, scheme, seconds)``: ``sh`` is ``(X, Y, Z, n_coef)`` (the WM FOD, zero outside the fitted voxels), ``fractions`` is ``(X, Y, Z, 3)`` (wm, gm, csf, zero outside), the rest are for the manifest / checks.""" from dmipy_fit.core.acquisition_scheme import acquisition_scheme_from_bvalues from dmipy_fit.core.modeling_framework import MultiCompartmentSphericalHarmonicsModel from dmipy_fit.tissue_response.three_tissue_response import three_tissue_response_dhollander16 scheme = acquisition_scheme_from_bvalues(bvals * 1e6, bvecs, b0_threshold=B0_THRESHOLD_S_MM2 * 1e6) data = np.nan_to_num(np.asarray(data, dtype=np.float64), nan=0.0) t0 = time.perf_counter() # The responses are selected from the brain mask eroded by 3 passes (the function's default, MRtrix's # dwi2response dhollander -erode 3), the white matter by Tournier 2013's iterative peak-ratio selection (its # default wm_algorithm): RESPONSES names both in the manifest. (s0_wm, s0_gm, s0_csf), (tr2_wm, tr1_gm, tr1_csf), _selection = three_tissue_response_dhollander16( scheme, data, mask=mask, erode=RESPONSE_ERODE) response_seconds = time.perf_counter() - t0 mc = MultiCompartmentSphericalHarmonicsModel( models=[tr2_wm, tr1_gm, tr1_csf], S0_tissue_responses=[s0_wm, s0_gm, s0_csf], sh_order=sh_order) fit_mask = mask if voxel_positions is None else voxel_positions t0 = time.perf_counter() fitted = mc.fit(scheme, data, mask=fit_mask, solver=solver, verbose=False, use_parallel_processing=False) fit_seconds = time.perf_counter() - t0 sh = np.asarray(fitted.fitted_parameters["sh_coeff"], dtype=np.float32) fractions = np.stack([ np.asarray(fitted.fitted_parameters["partial_volume_0"], dtype=np.float32), np.asarray(fitted.fitted_parameters["partial_volume_1"], dtype=np.float32), np.asarray(fitted.fitted_parameters["partial_volume_2"], dtype=np.float32), ], axis=-1) responses = {"S0_wm": float(s0_wm), "S0_gm": float(s0_gm), "S0_csf": float(s0_csf), "response_seconds": response_seconds, "fit_seconds": fit_seconds, "n_voxels_fit": int(np.sum(fit_mask))} return sh, fractions, responses, scheme def _git_sha(repo_dir): try: return subprocess.check_output(["git", "-C", repo_dir, "rev-parse", "HEAD"], text=True).strip() except Exception: return None def _dmipy_fit_commit(): """The git commit of whichever ``dmipy-fit`` is actually importable (found from the package's own file path, not a hard-coded location: ``PYTHONPATH`` picks which checkout this is).""" import dmipy_fit return _git_sha(os.path.dirname(os.path.dirname(os.path.abspath(dmipy_fit.__file__)))) def build_manifest(*, subject, scan, source_doi, grid_shape, voxel_size_mm, affine, bvals, bvecs, te_s, delta_s, delta_cap_s, dmipy_fit_commit, disco_space_commit, solver, n_regions): """The asset's ``manifest.json`` payload (see disco-space#8's schema in the issue body).""" return { "subject": subject, "source": {"dataset": "OpenNeuro ds003416", "doi": source_doi, "paper_doi": "10.1002/mrm.28926", "license": "CC0", "scan": scan}, "grid": {"shape": list(int(s) for s in grid_shape), "voxel_size_mm": list(float(v) for v in voxel_size_mm), "affine": np.asarray(affine, dtype=float).tolist()}, "protocol": {"bvals_s_mm2": np.asarray(bvals, dtype=float).tolist(), "bvecs": np.asarray(bvecs, dtype=float).tolist(), "bvec_frame": "image", "TE_s": te_s, "delta_s": delta_s, "Delta_s": delta_cap_s}, "fod": {"basis": "dmipy-fit tournier07 (dipy real_sh_tournier real-SH ordering)", "lmax": SH_ORDER, "frame": "image"}, "fractions": ["wm", "gm", "csf"], "reconstruction": {"tool": "dmipy-fit", "commit": dmipy_fit_commit, "responses": RESPONSES, "solver": solver, "sh_order": SH_ORDER}, "parcellation": {"tool": "SynthSeg --parc --robust (BBillot/SynthSeg) + dipy rigid registration " "(mutual information) + majority vote onto the diffusion grid", "atlas": "Desikan-Killiany + aseg (fs_default 84)", "n_regions": n_regions}, "built_by": {"disco_space_commit": disco_space_commit, "date": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())}, } def save_checks_sample(out_dir, data, bvals, sh, fractions, mask, n_sample=2000, seed=0): """``checks/msmt_sample.npz``: a fixed-seed random sample of the input signal + fitted SH coefficients + fractions, so dmipy-fit#39's fast solver can be checked against this (slow, cvxpy) reference.""" rng = np.random.default_rng(seed) idx = np.argwhere(mask) n_sample = min(n_sample, len(idx)) sel = idx[rng.choice(len(idx), size=n_sample, replace=False)] pos = tuple(sel.T) os.makedirs(os.path.join(out_dir, "checks"), exist_ok=True) np.savez(os.path.join(out_dir, "checks", "msmt_sample.npz"), voxel_index=sel.astype(np.int32), signal=data[pos].astype(np.float32), bvals_s_mm2=np.asarray(bvals, dtype=np.float32), sh_coeff=sh[pos].astype(np.float32), fractions=fractions[pos].astype(np.float32), seed=seed) def fraction_sanity(fractions, mask): """The fraction-sum sanity numbers to report: the median and 99th percentile of ``wm + gm + csf`` in-mask (it is not constrained to 1 by the unity constraint being off for ``S0_tissue_responses``-scaled fits -- see ``MultiCompartmentSphericalHarmonicsModel.fit``'s ``unity_constraint='kernel_dependent'`` default), and the count of voxels with WM fraction > 0.5.""" s = fractions[mask].sum(axis=-1).astype(np.float64) return { "sum_median": float(np.median(s)), "sum_p99": float(np.percentile(s, 99)), "n_wm_gt_half": int(np.sum(fractions[..., 0][mask] > 0.5)), "n_mask": int(mask.sum()), } def fod_peak_counts(sh, wm_mask, sh_order=SH_ORDER, relative_peak_threshold=0.5, min_separation_angle=25): """The number of FOD peaks per WM voxel (dipy's standard ``peak_directions`` on a fixed discrete sphere, the tournier07 SH basis evaluated at its vertices): ``{n_peaks: count}``, the single-vs-crossing distribution.""" from collections import Counter from dipy.data import get_sphere from dipy.direction.peaks import peak_directions from dmipy_tract import sh_matrix sphere = get_sphere(name="repulsion724") basis = sh_matrix(sh_order, sphere.vertices) coef = np.nan_to_num(np.asarray(sh, np.float64))[wm_mask] amplitudes = coef @ basis.T counts = Counter() for odf in amplitudes: odf = np.clip(odf, 0, None) if not np.any(odf > 0): counts[0] += 1 continue _, vals, _ = peak_directions(odf, sphere, relative_peak_threshold=relative_peak_threshold, min_separation_angle=min_separation_angle) counts[len(vals)] += 1 return dict(counts) def save_mean_b0_fod_png(out_path, mean_b0_volume, sh, wm_mask, z=None, sh_order=SH_ORDER, step=2): """A PNG of one axial slice of the mean b = 0 image with the WM FOD's principal direction overlaid as a short line segment per voxel (every ``step``-th voxel, to keep the plot legible); Matplotlib's ``Agg`` backend, no display needed.""" import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt from dmipy_tract import hemisphere, sh_matrix z = mean_b0_volume.shape[2] // 2 if z is None else z dirs = hemisphere(362) basis = sh_matrix(sh_order, dirs) coef = np.nan_to_num(np.asarray(sh, np.float64)).reshape(-1, sh.shape[-1]) amp = coef @ basis.T k = np.argmax(amp, axis=1) peak = np.take_along_axis(amp, k[:, None], axis=1)[:, 0] principal = np.where(peak[:, None] > 0, dirs[k], 0.0).reshape(sh.shape[:-1] + (3,)) fig, ax = plt.subplots(figsize=(6, 6)) ax.imshow(mean_b0_volume[:, :, z].T, cmap="gray", origin="lower") xs, ys, us, vs = [], [], [], [] for x in range(0, mean_b0_volume.shape[0], step): for y in range(0, mean_b0_volume.shape[1], step): if not wm_mask[x, y, z]: continue d = principal[x, y, z] if np.linalg.norm(d) < 1e-6: continue xs.append(x); ys.append(y); us.append(d[0]); vs.append(d[1]) ax.quiver(xs, ys, us, vs, angles="xy", scale_units="xy", scale=1.2, headaxislength=0, headlength=0, width=0.003, color="red") ax.set_title(f"mean b=0, axial z={z}, WM FOD principal direction") ax.set_xticks([]); ax.set_yticks([]) fig.tight_layout() fig.savefig(out_path, dpi=130) plt.close(fig) def run_checks(out_dir): """Loads a finished asset directory's arrays and writes ``checks/sanity_summary.json`` and ``checks/mean_b0_fod_slice.png`` (the peak-count / fraction-sum / WM-volume numbers this issue asks to report, decoupled from the -- long -- fit itself, so a completed fit's checks can be (re)computed on their own).""" fractions = np.load(os.path.join(out_dir, "fractions.npy")).astype(np.float32) mask = np.load(os.path.join(out_dir, "mask.npy")) sh = np.load(os.path.join(out_dir, "fod_wm.npy")) mean_b0_volume = np.load(os.path.join(out_dir, "mean_b0.npy")) wm_mask = mask & (fractions[..., 0] > 0.5) summary = {"fractions": fraction_sanity(fractions, mask), "fod_peak_counts_in_wm": fod_peak_counts(sh, wm_mask)} os.makedirs(os.path.join(out_dir, "checks"), exist_ok=True) with open(os.path.join(out_dir, "checks", "sanity_summary.json"), "w") as f: json.dump(summary, f, indent=2) save_mean_b0_fod_png(os.path.join(out_dir, "checks", "mean_b0_fod_slice.png"), mean_b0_volume, sh, wm_mask) print(json.dumps(summary, indent=2)) return summary def main(): ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("--dwi-dir", required=False, default=None) ap.add_argument("--parcellation-dir", required=False, default=None) ap.add_argument("--out-dir", required=True) ap.add_argument("--subject", default="MASiVar sub-cIs1 ses-s1Ax1 run-1xx") ap.add_argument("--source-doi", default="10.18112/openneuro.ds003416.v2.0.2") ap.add_argument("--solver", default="csd_cvxpy") ap.add_argument("--n-test-voxels", type=int, default=None, help="fit only this many random masked voxels (a dry run; nothing written under --out-dir " "except the printed timing)") ap.add_argument("--seed", type=int, default=0) ap.add_argument("--checks-only", action="store_true", help="skip everything above and just (re)compute checks/sanity_summary.json + the FOD PNG " "from an --out-dir a previous run already wrote") a = ap.parse_args() if a.checks_only: run_checks(a.out_dir) return if not a.dwi_dir or not a.parcellation_dir: ap.error("--dwi-dir and --parcellation-dir are required unless --checks-only") import nibabel as nib data, bvals, bvecs, affine, shells = merge_shells(a.dwi_dir) b0 = mean_b0(data, bvals) mask = brain_mask(b0) voxel_size_mm = np.linalg.norm(affine[:3, :3], axis=0) print(f"grid {data.shape[:3]}, voxel {voxel_size_mm} mm, {data.shape[-1]} measurements, shells {shells}, " f"mask voxels {mask.sum()}") if a.n_test_voxels: rng = np.random.default_rng(a.seed) idx = np.argwhere(mask) sel = idx[rng.choice(len(idx), size=min(a.n_test_voxels, len(idx)), replace=False)] voxel_positions = np.zeros(mask.shape, dtype=bool) voxel_positions[tuple(sel.T)] = True sh, fractions, responses, _scheme = fit_msmt_csd( data, bvals, bvecs, mask, solver=a.solver, voxel_positions=voxel_positions) per_voxel = responses["fit_seconds"] / responses["n_voxels_fit"] print(f"responses {responses['response_seconds']:.1f} s; fit {responses['n_voxels_fit']} voxels in " f"{responses['fit_seconds']:.1f} s ({per_voxel * 1e3:.1f} ms/voxel); whole brain " f"({int(mask.sum())} voxels) extrapolates to {per_voxel * int(mask.sum()) / 60:.1f} min serial") return labels = np.load(os.path.join(a.parcellation_dir, "labels.npy")) stop_mask = np.load(os.path.join(a.parcellation_dir, "stop_mask.npy")) with open(os.path.join(a.parcellation_dir, "regions.json")) as f: regions = json.load(f) if labels.shape != data.shape[:3]: raise ValueError(f"parcellation grid {labels.shape} does not match the DWI grid {data.shape[:3]}") sh, fractions, responses, _scheme = fit_msmt_csd(data, bvals, bvecs, mask, solver=a.solver) print(f"responses {responses['response_seconds']:.1f} s; whole-brain fit " f"{responses['n_voxels_fit']} voxels in {responses['fit_seconds']:.1f} s") os.makedirs(a.out_dir, exist_ok=True) np.save(os.path.join(a.out_dir, "fod_wm.npy"), sh) np.save(os.path.join(a.out_dir, "fractions.npy"), np.clip(fractions, 0, 1).astype(np.float16)) np.save(os.path.join(a.out_dir, "mask.npy"), mask) np.save(os.path.join(a.out_dir, "mean_b0.npy"), b0) np.save(os.path.join(a.out_dir, "labels.npy"), labels) np.save(os.path.join(a.out_dir, "stop_mask.npy"), stop_mask) with open(os.path.join(a.out_dir, "regions.json"), "w") as f: json.dump(regions, f, indent=2) manifest = build_manifest( subject=a.subject, scan="ses-s1Ax1 run-1xx (b1000/1500/2000/2500/3000, n=96 each, PreQual-preprocessed)", source_doi=a.source_doi, grid_shape=data.shape[:3], voxel_size_mm=voxel_size_mm, affine=affine, bvals=bvals, bvecs=bvecs, te_s=None, delta_s=None, delta_cap_s=None, dmipy_fit_commit=_dmipy_fit_commit(), disco_space_commit=_git_sha(os.path.join(os.path.dirname(__file__), "..")), solver=a.solver, n_regions=len(regions["regions"])) with open(os.path.join(a.out_dir, "manifest.json"), "w") as f: json.dump(manifest, f, indent=2) save_checks_sample(a.out_dir, data, bvals, sh, fractions, mask, seed=a.seed) run_checks(a.out_dir) print(f"asset written to {a.out_dir}") if __name__ == "__main__": main()