Spaces:
Running on Zero
Running on Zero
Download space/sources/brain.py from rfick/brain-zero: direct link, hf CLI and curl.
- Browser
- Download file 94.4 kB
-
https://huggingface.co/spaces/rfick/brain-zero/resolve/main/space/sources/brain.py
- Command line
-
hf download hf://spaces/rfick/brain-zero/space/sources/brain.py
-
curl -L -o brain.py https://huggingface.co/spaces/rfick/brain-zero/resolve/main/space/sources/brain.py
94.4 kB
| """The brain source: a brain composed, voxel by voxel, from a white-matter FOD field, WM / GM / CSF fractions and | |
| replay packs, replayed on the device, reconstructed, tracked and scored against the connectome of its own input. | |
| The input is an **asset** (the directory layout of ``SubstrateCommons/masivar-brain``: ``manifest.json``, | |
| ``fod_wm.npy``, ``fractions.npy``, ``mask.npy``, ``labels.npy``, ``regions.json``, ``stop_mask.npy``, | |
| ``mean_b0.npy``): the multi-tissue CSD of a real scan, made once offline, with its parcellation. The replay is the | |
| kernel route of dmipy-sim's phantom (RPH.md 6): per voxel and measurement | |
| S = | f_wm m0_wm FOD . C_wm + f_gm m0_gm E_gm + f_csf m0_csf E_csf | | |
| where ``C_wm`` is the WM pack's pose response contracted to the FOD's harmonics | |
| (:meth:`dmipy_sim.replay.ReplayPack.pose_responses` at ``keep=(8, 0)`` through :func:`dmipy_sim.replay.so3.axis_density_coeffs`), | |
| ``E_gm`` the GM pack's response to an isotropic orientation distribution, ``E_csf`` free water in closed form | |
| (:class:`dmipy_sim.phantom.FreeWater`), each at the run's tissue (T2 / T1 per pool at the echo, the WM walls' surface | |
| relaxivity, the myelin sheath's susceptibility at the scanner's field and its direction), and ``m0`` the proton | |
| density per tissue. The pose responses (:meth:`Brain.prepare`: seconds on the CPU) are cached in the page's process per | |
| pack, protocol, tissue and field; a run looks them up on the page and computes the missing ones inside the GPU call, | |
| so the call is entered as soon as the request arrives. The contraction over voxels x | |
| measurements x 45 runs in torch on the device. It is exactly what ``Phantom.compose(...).replay(seq, pose=...)`` | |
| computes, which the tests check. The acquisition is in the image frame; the field's direction is given in the | |
| scanner frame and reaches the packs through the specimen's pose (:func:`specimen_pose`). | |
| """ | |
| from __future__ import annotations | |
| import functools | |
| import json | |
| import os | |
| import time | |
| from dataclasses import dataclass, replace | |
| from typing import Optional | |
| import numpy as np | |
| from .. import pipeline as P | |
| ASSET_FILES = ("manifest.json", "fod_wm.npy", "fractions.npy", "mask.npy", "labels.npy", "regions.json", "stop_mask.npy", "mean_b0.npy") | |
| TISSUES = ("wm", "gm", "csf") | |
| WM_POOLS = ("intra", "extra", "myelin") # the catalogue's white-matter pools, by the packs' pool names | |
| M0_KEYS = {"wm": "proton_density_white_matter", "gm": "proton_density_grey_matter", "csf": "proton_density_csf"} | |
| RELAXATION_KEYS = {"gm": ("T2_grey_matter", "T1_grey_matter"), "csf": ("T2_csf", "T1_csf")} | |
| SCAN = "scan" # the timing class of the asset's own protocol | |
| B0_MODES = {"along the bore (z)": P.B0_ALONG_Z, "transverse (x): 90° from the bore, as in a biplanar magnet like the Swoop": P.B0_TRANSVERSE} | |
| FREE_B0 = "free (polar and azimuth angles below)" | |
| LMAX = 8 | |
| N_COEF = 45 | |
| REACH_MM = 250.0 # the longest streamline the tracker must be able to reach | |
| PROFILE_ANGLES = np.linspace(0.0, 90.0, 19) # degrees between the gradient and a fibre, for the response plot | |
| C00 = 1.0 / (2.0 * np.sqrt(np.pi)) # the l = 0 coefficient of a unit-integral density | |
| #: why a magnet with its own gradient (a field law: the Swoop) is not on the brain's scanner menu, as the page says it | |
| FIELD_LAW_REFUSAL = ("Its own gradient (the permanent magnet's field law, 1.4 mT/m at 8 cm) is on through every pulse and dead " | |
| "time, so every voxel plays a two-direction waveform of its own. dmipy-sim's closed form expands that " | |
| "waveform exactly; the cost keeps it off this menu. Binned to the menu's 1 % of b, this head needs 9,752 " | |
| "encoding classes for the scan's protocol (665 at 3 %), and one class's pose responses take 2-3 minutes " | |
| "(WM 104-166 s, GM 4-21 s at 485 measurements on a GH200), so one protocol at one tissue and field is about " | |
| "two weeks of one device. The DiSCo Space plays the Swoop exactly.") | |
| # ---- the asset --------------------------------------------------------------------------------------------------------- | |
| def asset_location(cfg): | |
| """Where the asset is: ``DISCO_BRAIN_ASSET`` in the environment (a local directory), else the configuration's | |
| ``[asset]``: ``local`` (a directory) or ``hub`` + ``revision`` (a public dataset downloaded into the Hub cache).""" | |
| env = os.environ.get("DISCO_BRAIN_ASSET") | |
| if env: | |
| return env | |
| a = cfg["asset"] | |
| return a["local"] if a.get("local") else f"hf://{a['hub']}@{a['revision']}" | |
| def asset_dir(where): | |
| """The asset's local directory: ``where`` itself, or the dataset ``hf://repo@revision`` downloaded (once per | |
| container: the Hub cache holds it).""" | |
| if not where.startswith("hf://"): | |
| return where | |
| from huggingface_hub import snapshot_download | |
| repo, _, rev = where[len("hf://"):].partition("@") | |
| return snapshot_download(repo, repo_type="dataset", revision=rev or None, allow_patterns=list(ASSET_FILES)) | |
| def manifest(where): | |
| """The asset's ``manifest.json`` (downloaded alone for a Hub asset).""" | |
| if where.startswith("hf://"): | |
| from huggingface_hub import hf_hub_download | |
| repo, _, rev = where[len("hf://"):].partition("@") | |
| path = hf_hub_download(repo, "manifest.json", repo_type="dataset", revision=rev or None) | |
| else: | |
| path = os.path.join(where, "manifest.json") | |
| with open(path) as f: | |
| return json.load(f) | |
| def declared_segments(uri): | |
| """The segment table ``{n, n_t, T}`` (RPK.md 4.3) a pack declares, read without its arrays: a Hub pack's | |
| record in its dataset's ``manifest.json`` (``hf://owner/name/path``), a local pack's header.""" | |
| from dmipy_sim.replay.publish import MANIFEST, header_of, is_hub_uri, parse_uri | |
| if is_hub_uri(uri): | |
| from huggingface_hub import hf_hub_download | |
| repo, path = parse_uri(uri) | |
| with open(hf_hub_download(repo, MANIFEST, repo_type="dataset")) as f: | |
| rows = [r for r in json.load(f).get("packs") or [] if r.get("path") == path] | |
| if not rows: | |
| raise ValueError(f"{uri}: the dataset {repo} holds no record of this pack (its manifest.json)") | |
| seg = rows[0].get("segments") | |
| else: | |
| seg = header_of(uri)["walk_params"].get("segments") | |
| if not seg: | |
| raise ValueError(f"{uri}: no declared segment table (walk_params.segments, RPK.md 4.3) in its " | |
| f"{'record' if is_hub_uri(uri) else 'header'}") | |
| return dict(n=int(seg["n"]), n_t=int(seg["n_t"]), T=float(seg["T"])) | |
| def image_rotation(man): | |
| """``R`` (image -> scanner) of the asset's grid (:meth:`dmipy_sim.phantom.Grid.from_oblique_affine`).""" | |
| from dmipy_sim.phantom import Grid | |
| _, R = Grid.from_oblique_affine(np.asarray(man["grid"]["affine"], np.float64), tuple(man["grid"]["shape"])) | |
| return np.asarray(R, np.float64) | |
| def read_regions(path): | |
| """``(regions, groups)`` from ``regions.json``: the regions ``[{id, name, hemisphere, lobe}]`` (a list, or the | |
| ``regions`` of an object) sorted by id, and the lobar groups ``((name, (id, ...)), ...)`` made from each region's | |
| hemisphere and lobe in order of first appearance; an object's ``groups`` must agree with them.""" | |
| with open(path) as f: | |
| doc = json.load(f) | |
| regions = sorted(doc["regions"] if isinstance(doc, dict) else doc, key=lambda r: int(r["id"])) | |
| groups = {} | |
| for r in regions: | |
| groups.setdefault(f"{str(r['hemisphere'])[0].upper()} {r['lobe']}", []).append(int(r["id"])) | |
| groups = tuple((name, tuple(ids)) for name, ids in groups.items()) | |
| if isinstance(doc, dict) and doc.get("groups"): | |
| stated = {tuple(sorted(int(i) for i in g["ids"])) for g in doc["groups"]} | |
| if stated != {tuple(sorted(ids)) for _, ids in groups}: | |
| raise ValueError("regions.json: its groups are not the regions' hemisphere x lobe") | |
| return regions, groups | |
| class Asset: | |
| """The asset, cropped to the brain's bounding box (one voxel of margin): ``fod`` ``(X, Y, Z, 45)`` in the image | |
| frame, ``fractions`` ``(X, Y, Z, 3)`` WM / GM / CSF (a voxel whose fractions sum above one is divided by its sum, a | |
| WM fraction without an FOD is not replayed; ``notes`` counts both), ``mask``, ``labels``, ``stop``, ``b0``, ``affine`` the | |
| cropped grid's voxel -> scanner mm, ``R`` image -> scanner, ``regions`` / ``groups`` from ``regions.json``, | |
| ``manifest``, ``crop`` the slices into the full grid.""" | |
| manifest: dict | |
| fod: np.ndarray | |
| fractions: np.ndarray | |
| mask: np.ndarray | |
| labels: np.ndarray | |
| stop: np.ndarray | |
| b0: np.ndarray | |
| affine: np.ndarray | |
| R: np.ndarray | |
| regions: list | |
| groups: tuple | |
| crop: tuple | |
| notes: tuple = () | |
| def load(cls, where): | |
| d = asset_dir(where) | |
| missing = [f for f in ASSET_FILES if not os.path.exists(os.path.join(d, f))] | |
| if missing: | |
| raise FileNotFoundError(f"the brain asset at {where} lacks {missing}") | |
| man = manifest(where) | |
| shape = tuple(man["grid"]["shape"]) | |
| fod = np.load(os.path.join(d, "fod_wm.npy"), mmap_mode="r") | |
| mask = np.load(os.path.join(d, "mask.npy")).astype(bool) | |
| if fod.shape != shape + (N_COEF,) or mask.shape != shape: | |
| raise ValueError(f"the asset's FOD {fod.shape} and mask {mask.shape} are not on its grid {shape} with {N_COEF} coefficients") | |
| from dmipy_sim.phantom.orientation import SH_BASES | |
| basis = str(man["fod"]["basis"]).split()[0] | |
| if SH_BASES.get(basis) != ("tournier07", False) or int(man["fod"]["lmax"]) != LMAX or man["fod"]["frame"] != "image": | |
| raise ValueError(f"the asset's FOD is {man['fod']}: this source reads dmipy-sim's required basis (a basis whose " | |
| f"name, the first word, dmipy_sim.phantom.orientation.SH_BASES maps to orthonormal tournier07: " | |
| f"tournier07, mrtrix3, dmipy-fit) at lmax {LMAX} in the image frame") | |
| if list(man["fractions"]) != list(TISSUES): | |
| raise ValueError(f"the asset's fractions are {man['fractions']}, this source reads {list(TISSUES)}") | |
| ijk = np.argwhere(mask) | |
| lo = np.maximum(ijk.min(0) - 1, 0); hi = np.minimum(ijk.max(0) + 2, shape) | |
| crop = tuple(slice(int(a), int(b)) for a, b in zip(lo, hi)) | |
| affine = np.asarray(man["grid"]["affine"], np.float64).copy() | |
| affine[:3, 3] = affine[:3, :3] @ lo + affine[:3, 3] | |
| regions, groups = read_regions(os.path.join(d, "regions.json")) | |
| labels = np.load(os.path.join(d, "labels.npy"))[crop].astype(np.int32) | |
| n = int(man["parcellation"]["n_regions"]) | |
| if [int(r["id"]) for r in regions] != list(range(1, n + 1)) or labels.min() < 0 or labels.max() > n: | |
| raise ValueError(f"the asset's regions are not 1..{n} (regions.json ids {[r['id'] for r in regions][:5]}..., labels up to {labels.max()})") | |
| m = mask[crop] | |
| fodc = np.asarray(fod[crop], np.float32) * m[..., None] | |
| fr = np.clip(np.load(os.path.join(d, "fractions.npy"))[crop].astype(np.float32), 0.0, None) * m[..., None] | |
| notes = [] | |
| no_fod = (fr[..., 0] > 0) & ~(fodc[..., 0] > 0) | |
| if no_fod.any(): # a WM fraction needs an orientation to compose with | |
| notes.append(("voxels with a WM fraction and no FOD (their WM fraction is not replayed)", | |
| f"{int(no_fod.sum()):,}, WM fraction up to {float(fr[..., 0][no_fod].max()):.3f}")) | |
| fr[..., 0] = np.where(no_fod, 0.0, fr[..., 0]) | |
| total = fr.sum(-1, keepdims=True) | |
| over = total[..., 0] > 1.0 | |
| if over.any(): # shares of a voxel's volume: float16 rounding, an unconstrained fit | |
| notes.append(("voxels whose fractions sum above one (divided by their sum)", f"{int(over.sum()):,}, largest sum {float(total.max()):.3f}")) | |
| fr = fr / np.maximum(total, 1.0) | |
| return cls(man, fodc, fr, m, labels, np.load(os.path.join(d, "stop_mask.npy"))[crop].astype(bool) & m, | |
| np.load(os.path.join(d, "mean_b0.npy"))[crop].astype(np.float32), affine, image_rotation(man), regions, groups, crop, tuple(notes)) | |
| def scan_protocol(man, *, name): | |
| """The asset's own protocol: its rows' b-values (s/mm^2) and directions turned into the image frame (the frame | |
| of the FOD the replay composes), one shell per b rounded to 50 s/mm^2, on the timing class :data:`SCAN`.""" | |
| b = np.asarray(man["protocol"]["bvals_s_mm2"], np.float64) | |
| g = np.asarray(man["protocol"]["bvecs"], np.float64) | |
| if g.shape == (3, len(b)): | |
| g = g.T | |
| frame = man["protocol"]["bvec_frame"] | |
| if frame not in ("image", "scanner"): | |
| raise ValueError(f"bvec_frame is 'image' or 'scanner', got {frame!r}") | |
| if frame == "scanner": | |
| g = g @ image_rotation(man) # R.T g per row | |
| group = np.round(b / 50.0).astype(int) | |
| return P.protocol_from_rows(b, g, group, lambda k, rows: P.Shell(SCAN, float(np.round(b[rows].mean())), int(len(rows))), name=name)[0] | |
| # ---- the physics -------------------------------------------------------------------------------------------------------- | |
| class BrainPhysics: | |
| """The tissue and scanner a brain replay is evaluated at: ``field_T`` (the scanner's field, T; the catalogue's | |
| relaxation is cited at the nearest of 1.5 / 3 / 7 T), ``b0_direction`` (its unit direction in the scanner frame), | |
| ``T2`` / ``T1`` as ``{"wm": {pool: s}, "gm": s, "csf": s}``, ``rho2`` the WM walls' surface relaxivity (m/s), | |
| ``chi_iso`` / ``chi_aniso`` the myelin sheath's susceptibility (SI, the WM pack's field source), ``D_csf`` free | |
| water's diffusivity (m^2/s), ``m0`` the proton density per tissue relative to CSF, ``wm_pack`` / ``gm_pack`` the | |
| packs by their configuration label; ``relaxation``, ``contact`` and ``field`` switch the three tiers; ``scanner`` | |
| the machine (a key of dmipy-sim's catalogue): its slew rate shapes every pulse's ramps, its field and direction are | |
| the physics', and every voxel is played the gradient it delivers there with the head centre at isocentre | |
| (:meth:`Brain.prepare`); None is the ideal scanner, ideal pulses with vertical ramps. All off is bare diffusion at | |
| the same proton densities and packs.""" | |
| field_T: float | |
| T2: dict | |
| T1: dict | |
| rho2: float | |
| chi_iso: float | |
| chi_aniso: float | |
| D_csf: float | |
| m0: dict | |
| wm_pack: str | |
| gm_pack: str | |
| b0_direction: tuple = P.B0_ALONG_Z | |
| relaxation: bool = True | |
| contact: bool = True | |
| field: bool = True | |
| scanner: Optional[str] = None | |
| def __post_init__(self): | |
| if self.scanner is not None: | |
| L = P.limits(self.scanner) | |
| if abs(float(L.field_T) - float(self.field_T)) > 1e-9: | |
| raise ValueError(f"{self.scanner} is a {L.field_T:g} T magnet; the physics says {self.field_T:g} T") | |
| if L.has_field_law: | |
| raise ValueError(f"{self.scanner} is not on the brain's scanner menu. {FIELD_LAW_REFUSAL}") | |
| if not (0 < self.field_T < 30): | |
| raise ValueError(f"the field is in tesla, got {self.field_T}") | |
| if set(self.m0) != set(TISSUES) or any(float(v) < 0 for v in self.m0.values()): | |
| raise ValueError(f"m0 is a non-negative proton density per tissue {TISSUES}, got {self.m0}") | |
| if self.rho2 < 0: | |
| raise ValueError("the surface relaxivity is non-negative") | |
| u = np.asarray(self.b0_direction, np.float64) | |
| if u.shape != (3,) or not np.isclose(np.linalg.norm(u), 1.0, atol=1e-6): | |
| raise ValueError("b0_direction is a unit vector") | |
| def bare(self): | |
| return not (self.relaxation or self.contact or self.field) | |
| def scanner_field(self): | |
| """The replay's ``scanner=``: the field in tesla with the field tier, else None (relaxation at a field is | |
| already in T2 and T1).""" | |
| return self.field_T if self.field else None | |
| def slew_rate(self): | |
| """The slew rate every pulse's ramps are played at (T/m/s): the catalogue class's ``slew_max``, infinite | |
| (vertical ramps) on the ideal scanner.""" | |
| if self.scanner is None: | |
| return np.inf | |
| from dmipy_sim.acquisition.scanners import ScannerLimits | |
| return float(ScannerLimits.of(self.scanner).slew_max) | |
| def label(self): | |
| if self.bare: | |
| return "bare diffusion" + (f" on the {self.scanner}" if self.scanner else "") | |
| tiers = "+".join(t for t in ("relaxation", "contact", "field") if getattr(self, t)) | |
| u = self.b0_direction | |
| grad = f", on the {self.scanner} (its slew {self.slew_rate:g} T/m/s and every catalogued term, head centre at isocentre)" if self.scanner else "" | |
| return f"{self.field_T:g} T along ({u[0]:.2f}, {u[1]:.2f}, {u[2]:.2f}), tiers {tiers}{grad}" | |
| def tissues(self, wm_pack, gm_pack): | |
| """``(wm, gm, csf)`` :class:`dmipy_sim.spec.Tissue` for the packs ``wm_pack`` and ``gm_pack``: the WM pack's | |
| pools with their T2 / T1, the walls' rho2 and, on a WM pack whose spec declares its susceptibility field | |
| present (a magnetic myelin pool), the sheath's susceptibility; the GM pack's pools at grey matter's T2 / T1; | |
| free water's D, T2 and T1. A tier off leaves its values out (``None`` for a pack with nothing on: its bare | |
| diffusion). A WM pack whose field is absent has no sheath, so its tissue carries no susceptibility and its | |
| field tier is zero at any field.""" | |
| from dmipy_sim.spec import Tissue | |
| wm_pools, gm_pools = Brain.pools(wm_pack), Brain.pools(gm_pack) | |
| relax, contact = self.relaxation, self.contact | |
| field = self.field and wm_pack.susceptibility_field == "present" | |
| wm = Tissue(T2={p: self.T2["wm"][p] for p in wm_pools} if relax else None, T1={p: self.T1["wm"][p] for p in wm_pools} if relax else None, | |
| rho2=self.rho2 if contact else None, chi_iso=self.chi_iso if field else None, | |
| chi_aniso=self.chi_aniso if field else 0.0) if (relax or contact or field) else None | |
| gm = Tissue(T2={p: self.T2["gm"] for p in gm_pools}, T1={p: self.T1["gm"] for p in gm_pools}) if relax else None | |
| csf = Tissue(D=self.D_csf, T2=self.T2["csf"] if relax else None, T1=self.T1["csf"] if relax else None) | |
| return wm, gm, csf | |
| def catalogue(field_T): | |
| """The catalogue at the cited field nearest ``field_T`` (in log distance): the white matter's pools | |
| (:func:`space.pipeline.catalogue`), grey matter's and CSF's T2 and T1, free water's diffusivity, the proton | |
| densities, all from ``dmipy_sim.substrate.biophysical_constants`` by key; ``catalogue_field`` says which field.""" | |
| import warnings | |
| from dmipy_sim.substrate.biophysical_constants import get_value | |
| wm = P.catalogue(field_T, WM_POOLS) | |
| near = wm["catalogue_field"] | |
| with warnings.catch_warnings(): | |
| warnings.simplefilter("ignore") | |
| T2 = {"wm": wm["T2"], **{t: float(get_value(k[0], near, allow_nearest=True)) for t, k in RELAXATION_KEYS.items()}} | |
| T1 = {"wm": wm["T1"], **{t: float(get_value(k[1], near, allow_nearest=True)) for t, k in RELAXATION_KEYS.items()}} | |
| return dict(catalogue_field=near, T2=T2, T1=T1, rho2=wm["rho2"], chi_iso=wm["chi_iso"], chi_aniso=wm["chi_aniso"], | |
| D_csf=float(get_value("D_csf")), m0={t: float(get_value(k)) for t, k in M0_KEYS.items()}) | |
| def specimen_pose(R, b0_direction): | |
| """The specimen's pose in the bore for a replay whose acquisition is in the image frame: ``W R`` with ``R`` the | |
| image -> scanner rotation and ``W`` the rotation taking the field's scanner-frame direction onto z (the bore's | |
| field axis), so the packs see the field at ``R.T b0_direction`` in the image frame. A sequence composed at this | |
| pose carries the gradients ``pose @ g`` (:meth:`Brain.sequence`).""" | |
| from dmipy_sim.replay.so3 import rotation_of | |
| return rotation_of(np.asarray(b0_direction, np.float64)).T @ np.asarray(R, np.float64) | |
| class Kernels: | |
| """A replay's pose responses (:meth:`Brain.prepare`): per encoding class and measurement, the WM response contracted | |
| to the FOD's harmonics ``wm`` ``(n_classes, n_meas, 45)`` complex, the GM response to an isotropic distribution ``gm`` | |
| and free water's ``csf`` ``(n_classes, n_meas)`` complex (a class fills the rows of its timing class), ``classes`` | |
| ``(n_voxels, n_timing)`` each voxel's class per timing class and ``timing_rows`` the measurement rows of each timing | |
| class (the ideal scanner: one class per timing class, every voxel in it), the proton densities ``m0``, the packs' | |
| certified floors, the seconds it took; | |
| ``profile`` the true response curves at the run's shells (``b``, ``wm`` ``(n_shells, n_angles)``, ``gm``, ``csf`` | |
| ``(n_shells,)``, each over its b = 0 value), ``b0`` the b = 0 signal per tissue under each tier alone, and | |
| ``stages`` the seconds of each stage of :meth:`Brain.prepare` (``load_wm``, ``load_gm``, ``pose_wm``, ``pose_gm``, | |
| ``weights``).""" | |
| wm: np.ndarray | |
| gm: np.ndarray | |
| csf: np.ndarray | |
| classes: np.ndarray | |
| timing_rows: tuple | |
| m0: dict | |
| floors: dict | |
| seconds: float | |
| profile: Optional[dict] = None | |
| b0: Optional[dict] = None | |
| stages: Optional[dict] = None | |
| # ---- the source --------------------------------------------------------------------------------------------------------- | |
| class Brain(P.Source): | |
| """A brain from its asset: the regions of its parcellation, the tracker seeded in its white matter and stopped | |
| where the white matter and the regions end, the connectome scored against the connectome of the asset's own FOD | |
| tracked with the same settings, seeds and key.""" | |
| plays_slew = True | |
| TIERS = ("relaxation", "contact", "field") | |
| STAGES = ("replay", "noise", "csd", "track", "truth", "score") | |
| def __init__(self, cfg, *, asset=None): | |
| super().__init__(cfg) | |
| self.where = asset or asset_location(cfg) | |
| self.asset = a = Asset.load(self.where) | |
| self.shapes = self.shape_table(cfg, self.where) | |
| self.affine = a.affine | |
| self.mask = a.mask | |
| self.regions = P.Regions(a.labels, a.affine, tuple(r["name"] for r in a.regions), a.groups) | |
| self.vox = np.argwhere(a.mask) # the replayed voxels | |
| f = a.fractions[a.mask].astype(np.float64) | |
| c0 = a.fod[a.mask][:, 0].astype(np.float64) | |
| unit = np.where((c0 > 0)[:, None], a.fod[a.mask] * (C00 / np.where(c0 > 0, c0, 1.0))[:, None], 0.0) | |
| self.fod_unit = unit.astype(np.float32) # every voxel's FOD as a unit-integral density | |
| self.f = f.astype(np.float32) | |
| self.packs = {} | |
| self._reference = {} | |
| self._kernels = {} # prepare()'s results, by response_key | |
| self._compiled = set() # the _band()s prepare() ran in this process | |
| self._encodings = {} # encoding()'s results, by machine and rows | |
| def load(cls, cfg, *, local=None): | |
| return cls(cfg, asset=local) | |
| def warm_entries(self): | |
| """The ``(meas, physics)`` a run can ask for without a custom acquisition, one per response key, in the order | |
| the warm-up computes them: every preset at the panel's defaults with its ladder rungs, then B of every knob | |
| on the default protocol, the field presets last in ascending field along both B0 directions (the band to | |
| compile grows with the field: 7 T and 11.7 T cost minutes each, so they follow everything cheaper, the | |
| pulse-timing classes among it), with the scan's protocol on each machine of the scanner menu that can play it, | |
| at the system's field, with its slew and its delivered gradient, before the field presets. A knob that leaves the response unchanged (noise, M0, the default pack | |
| or direction) repeats no entry; an entry the packs cannot play is left out (a run asking for it is refused | |
| by name).""" | |
| cfg = self.cfg | |
| panel = self.panel(cfg) | |
| values = {c.name: c.value for row in panel.rows for c in row if c.name in panel.fields} | |
| physics = self.physics_from(cfg, values) | |
| entries = [] | |
| for name in self.presets(cfg): | |
| prot = self.protocol(cfg, name) | |
| entries += [(prot, physics)] + [(prot, rung) for _, rung in self.ladder_steps(physics)] | |
| scan = self.protocol(cfg, self.presets(cfg)[0]) | |
| knobs = [c for c in self.knobs(cfg).values() if c is not None] | |
| fields = sorted((c for c in knobs if c[0] == "field"), key=lambda c: c[1]) | |
| def knob_entries(changes): | |
| out = [] | |
| for change in changes: | |
| prot, _, _, v = self.apply_knob(cfg, change, scan, True, None, values) | |
| modes = list(B0_MODES) if change[0] == "field" else [v["b0_mode"]] | |
| out += [(prot, self.physics_from(cfg, {**v, "b0_mode": mode})) for mode in modes] | |
| return out | |
| entries += knob_entries([c for c in knobs if c[0] != "field"]) | |
| for key in P.machines(cfg).values(): # the machines that can play the scan, at their field | |
| if not all(r[-1] for r in P.playable(scan, self.shapes, key)): | |
| continue | |
| _, _, _, v = self.apply_knob(cfg, ("field", float(P.limits(key).field_T)), scan, True, None, values) | |
| entries.append((scan, self.physics_from(cfg, v, scanner=key))) | |
| entries += knob_entries(fields) | |
| out = []; keys = set() | |
| for prot, ph in entries: | |
| try: | |
| meas = self.validate(prot, ph) | |
| except ValueError: | |
| continue | |
| key = self.response_key(meas, ph) | |
| if key not in keys: | |
| keys.add(key) | |
| out.append((meas, ph)) | |
| return out | |
| def warm_in_background(self): | |
| """:meth:`warm` in a process of its own while this one serves: the entries' responses arrive one by one and | |
| are kept here (:meth:`keep`), so a run finds the cache filling from the first entry on and pays, inside its | |
| reservation, only for an entry not yet there. A spawned process, not a thread: the GPU worker is a fork of | |
| this process, and a fork taken while a thread runs JAX can deadlock the child; the spawned process compiles | |
| its own bands, so an entry of a band this process never compiled is priced cold (:meth:`responses`) until a | |
| run computes it here. Returns the thread that drains the results; the packs are loaded here first (the | |
| windows the entries reach), so a run arriving before the warm-up ends has them. Both sides print their | |
| progress to stdout (the Space's run log): the process its start (CPUs, memory limit), each entry with its | |
| seconds and RSS; this side each entry kept and the process's exit code when it ends, with or without its | |
| end marker (a process killed by the container's memory limit sends none).""" | |
| import multiprocessing as mp | |
| import threading | |
| from queue import Empty | |
| entries = self.warm_entries() | |
| for meas, ph in entries: | |
| for tissue, label in (("wm", ph.wm_pack), ("gm", ph.gm_pack)): | |
| self.pack(tissue, label, self.windows_needed(tissue, label, meas)) | |
| ctx = mp.get_context("spawn") | |
| queue = ctx.Queue() | |
| proc = ctx.Process(target=_warm_worker, args=(self.cfg, self.where, entries, queue), daemon=True, name="brain-warm") | |
| proc.start() | |
| def drain(): | |
| n = 0 | |
| while True: | |
| try: | |
| item = queue.get(timeout=30) | |
| except Empty: | |
| if proc.is_alive(): | |
| continue | |
| print(f"warm-up: the process ended with exit code {proc.exitcode} after {n} of {len(entries)} entries, " | |
| "without its end marker", flush=True) | |
| return | |
| if item is None: | |
| break | |
| key, kernels = item | |
| self.keep({key: kernels}) | |
| n += 1 | |
| print(f"warm-up: the page keeps entry {n}/{len(entries)}", flush=True) | |
| proc.join() | |
| print(f"warm-up: the process ended with exit code {proc.exitcode} after {n} of {len(entries)} entries", flush=True) | |
| t = threading.Thread(target=drain, daemon=True, name="brain-warm-drain") | |
| t.start() | |
| return t | |
| # ---- the packs ---- | |
| def refused_machines(cls, cfg): | |
| """``{label: reason}`` for ``[scanners] refused`` of the configuration: each a catalogued magnet with its own | |
| gradient, kept off the menu for :data:`FIELD_LAW_REFUSAL`.""" | |
| out = {} | |
| for m in (cfg.get("scanners") or {}).get("refused", []): | |
| if not P.limits(m["key"]).has_field_law: | |
| raise ValueError(f"[scanners] refused names {m['key']!r}, which has no field law: the brain plays it, so it belongs on the menu") | |
| out[m["label"]] = FIELD_LAW_REFUSAL | |
| return out | |
| def pack_menu(cls, cfg, tissue): | |
| """The configuration's packs for ``tissue`` (``wm`` / ``gm``): ``{label: uri}`` in the menu's order.""" | |
| return {p["label"]: p["uri"] for p in cfg["packs"][tissue]} | |
| def segments(self, tissue, label): | |
| """The declared segment table ``{n, n_t, T}`` of the pack ``label`` of ``tissue``'s menu | |
| (:func:`declared_segments`: read without the pack's arrays).""" | |
| menu = self.pack_menu(self.cfg, tissue) | |
| if label not in menu: | |
| raise ValueError(f"no {tissue.upper()} pack {label!r}; the menu is {list(menu)}") | |
| return declared_segments(menu[label]) | |
| def windows_reached(self, segments, meas): | |
| """``{class: k}``: per timing class of ``meas``, the number of the pack's windows its acquisition reaches. | |
| ``segments`` is a declared table ``{n, n_t, T}`` (:func:`declared_segments`) or a pack's meta (its | |
| ``walk_params.segments``). The acquisition lasts ``T_acq = (G.shape[1] - 1) dt`` of the class's sequence, | |
| as :meth:`dmipy_sim.replay.ReplayPack._compile` reads it, and window ``i`` starts at ``i T``; a window whose | |
| first save sits at or beyond ``T_acq`` is not reached, so the save at a window's end belongs to that window | |
| alone. A pack of one window (``n = 1``) is reached as one window whatever the class; ``k`` | |
| above ``n`` is an acquisition beyond the walk (:meth:`windows_needed` refuses it).""" | |
| seg = (segments.get("walk_params") or {}).get("segments") if "walk_params" in segments else segments | |
| out = {} | |
| for name in map(str, np.unique(meas.shape)): | |
| if int(seg["n"]) <= 1: | |
| out[name] = 1 | |
| continue | |
| s = self._b0_sequence(name, np.eye(3)) | |
| T_acq = (np.shape(s.G)[1] - 1) * float(s.dt) | |
| out[name] = max(1, int(np.ceil(T_acq * (1.0 - 1e-12) / float(seg["T"])))) | |
| return out | |
| def windows_needed(self, tissue, label, meas): | |
| """The windows of the pack ``label`` of ``tissue``'s menu the classes of ``meas`` reach (the largest | |
| :meth:`windows_reached`), refused by name for a class that reaches beyond the windows the pack declares.""" | |
| seg = self.segments(tissue, label) | |
| reach = self.windows_reached(seg, meas) | |
| for name, k in reach.items(): | |
| if k > int(seg["n"]): | |
| raise ValueError(f"the {tissue.upper()} pack {label!r} declares {seg['n']} window(s) of {float(seg['T']) * 1e3:g} ms: " | |
| f"the class {name!r} (TE {float(self.shapes[name]['TE']) * 1e3:g} ms) reaches window {k - 1}") | |
| return max(reach.values()) | |
| def saves_spanned(segments, k): | |
| """The saves of the first ``k`` windows of a declared table: ``k (n_t - 1) + 1`` (consecutive windows share | |
| their boundary save; one window is its ``n_t``).""" | |
| return int(k) * (int(segments["n_t"]) - 1) + 1 | |
| def pack(self, tissue, label, k): | |
| """The pack ``label`` of ``tissue``'s menu holding at least its first ``k`` windows: a Hub URI whose | |
| declared table has ``n > 1`` loads ``windows=range(k)`` (:class:`dmipy_sim.phantom.PackSubstrate`; the pack | |
| records them as ``windows_present``); a local path or a one-window pack loads whole. A later request for | |
| more windows than it holds reloads with the larger range (the library fetches only the new ranges) and | |
| replaces the cached pack.""" | |
| from dmipy_sim.phantom import PackSubstrate | |
| from dmipy_sim.replay.publish import is_hub_uri | |
| key = (tissue, label) | |
| seg = self.segments(tissue, label) | |
| if key in self.packs and self.windows_held(self.packs[key], seg) >= int(k): | |
| return self.packs[key] | |
| uri = self.pack_menu(self.cfg, tissue)[label] | |
| windows = {"windows": range(int(k))} if is_hub_uri(uri) and int(seg["n"]) > 1 else {} | |
| self.packs[key] = PackSubstrate(uri, m0=1.0, name=f"{tissue}:{label}", **windows).pack | |
| return self.packs[key] | |
| def windows_held(pack, segments): | |
| """The windows ``pack`` holds: its ``windows_present`` when loaded by window, else every window ``segments`` | |
| declares.""" | |
| return int(segments["n"]) if pack.windows_present is None else int(pack.windows_present) | |
| def pools(pack): | |
| """The pack's pool names (its embedded spec), or None for a pack without a spec.""" | |
| spec = pack.substrate | |
| return None if spec is None else tuple(p.name for p in spec.pools) | |
| def walk_seconds(pack): | |
| return float(pack.meta["walk_params"]["T_max"]) | |
| # ---- the page's side ---- | |
| def shape_table(cls, cfg, where=None): | |
| """The timing classes: the configuration's ``[shapes]`` and :data:`SCAN`, the scan's own pulse timing from the | |
| asset's manifest, or the configuration's ``[scan]`` timing when the manifest states none (the label says | |
| which).""" | |
| pr = manifest(where or asset_location(cfg))["protocol"] | |
| stated = all(pr.get(k) is not None for k in ("delta_s", "Delta_s", "TE_s")) | |
| t = dict(delta=float(pr["delta_s"]), Delta=float(pr["Delta_s"]), TE=float(pr["TE_s"])) if stated else \ | |
| {k: float(cfg["scan"][k]) for k in ("delta", "Delta", "TE")} | |
| why = "the scan's own timing" if stated else "the scan's b-values and directions at the configuration's timing (the asset states none)" | |
| return {**cfg["shapes"], SCAN: dict(label=f"{why}: δ {t['delta'] * 1e3:g} / Δ {t['Delta'] * 1e3:g} ms, TE {t['TE'] * 1e3:g} ms", **t)} | |
| def describe(cls, cfg): | |
| d = cfg["describe"] | |
| return dict( | |
| title=d["title"], | |
| heading=(f"# {d['title']}\n" + d["heading"]), | |
| acquisition=("A shell's **pulse timing** (δ, Δ, TE; square pulses) is one of the classes below, each played by the packs' " | |
| "walks, stored in windows (125 ms for the windowed packs): a class reads the windows its acquisition reaches, and a timing beyond the " | |
| "walk is refused by name; the b-values, the directions and their number, the " | |
| "tissue, the proton densities, the packs, the SNR and the tracker are free. The first preset is the scan's own " | |
| "protocol, its directions turned into the image frame. Timing classes: " | |
| + "; ".join(f"`{n}` = {t['label']}" for n, t in cls.shape_table(cfg).items()) + "."), | |
| tissue=("The replay composes three tissues per voxel: the WM pack contracted with the voxel's FOD, the GM pack as an " | |
| "isotropic distribution, free water in closed form, each weighted by its fraction and proton density M0 " | |
| "(relative to CSF). **SNR is defined at M0 = 1**, the b = 0 signal of a pure-CSF voxel before relaxation, so " | |
| "lowering a tissue's M0 lowers its SNR. The **field** tier is the WM pack's myelin-sheath susceptibility at the " | |
| "scanner's field and direction (given in the scanner frame; the head's own susceptibility is not modelled); the " | |
| "GM pack and free water carry no field source. The packs' responses to the presets and the field presets are " | |
| "computed once when the Space starts; any other (a custom protocol, a field or tissue not seen before) is " | |
| "computed inside the GPU call and is part of what the run reserves. The pack menu is the " | |
| "configuration's list; the WM pack's contact tier uses the catalogue's white-matter ρ₂, the GM pack's walls take " | |
| "none (no cited value)."), | |
| explorer=("What the replay made, before the noise and the tractography. **A** is the run you configured (its replay " | |
| "before the noise); **B** is the same run with the one knob you chose changed. **Ingredients**: the proton-density " | |
| "map Σ f·M0, the relaxation weight of every voxel's b = 0 signal at the echo, and the contact tier's survival " | |
| "(the field tier's effect is its rung of the ladder). " | |
| "**Layers**: A replayed with its tiers switched on one at a time (bare → + relaxation → + contact → + field), then B, as the " | |
| "DWI or the tensor's MD and FA. **Estimated vs true response**: the responses the reconstruction estimated from the " | |
| "replayed data against the packs' exact response at the run's tissue, per shell: what the response heuristic " | |
| "gets right is itself a result. There is no per-voxel split-half floor here: the accuracy is the packs' " | |
| "certified floors (the accuracy table)."), | |
| results_tab="4 · brain results", | |
| truth_tab="2 · the input", | |
| truth_labels=("the regions (centroids, by lobe)", "the input: WM FOD peaks, fractions and regions"), | |
| stages={"replay": "replaying the brain from the packs' responses", "noise": "adding Rician noise", | |
| "csd": "estimating the responses and fitting CSD", "track": "tracking from the white matter", | |
| "truth": "tracking the input FOD with the same settings (the truth)", "score": "scoring the connectome"}, | |
| fixed=[f"the subject: {d['subject']}; its FOD field, fractions and parcellation (made once, offline)", | |
| "one pack per tissue class per run, from the configuration's menu (Hub records later)", | |
| "the head's own susceptibility and RF effects are not modelled", | |
| "the head centre at the scanner's isocentre (a machine of the scanner menu plays every voxel the gradient it delivers there; docs/scanner.md)", | |
| "no per-voxel split-half floor: the packs' certified floors are the accuracy"], | |
| accuracy=("A brain replay is exact composition (the tests check it against dmipy-sim's own `Phantom.compose(...).replay`) of " | |
| "the packs' responses, so its accuracy is the packs' certified Monte-Carlo floors; there is no per-voxel split-half " | |
| "floor, because every voxel reads the same two walks."), | |
| ingredients=("proton density: Σ f·M0 (relative to CSF)", "relaxation: the b = 0 signal's T2 / T1 weight at the echo", | |
| "contact: the WM walls' surface-relaxivity survival at b = 0"), | |
| truth=("the connectome of the input FOD tracked directly with the same tracker, settings, seeds and key: Pearson of " | |
| "log(1 + count) over the region pairs, the pairs connected in one only, and the same on the lobar groups"), | |
| ladder="Replay DWI Explorer: replay A's tier ladder too (bare, + relaxation, + contact; noise-free; a few seconds of GPU time)", | |
| views=dict(truth=True, fractions=True, lobar=True, roundtrip=True, response=True, floor=False), | |
| files="brain", step_unit="mm", | |
| volumes_label="DWI (.nii.gz, bvecs in the image frame) with bvals/bvecs, the FOD SH field (tournier07 order 8), the connectome (.csv)", | |
| run_label="replay → responses + CSD → track → connectome") | |
| def presets(cls, cfg): | |
| return [cfg["describe"]["scan_preset"]] + list(cfg["presets"]) | |
| def protocol(cls, cfg, name): | |
| if name == cfg["describe"]["scan_preset"]: | |
| return scan_protocol(manifest(asset_location(cfg)), name=name) | |
| return P.preset_protocol(cfg, name) | |
| def panel(cls, cfg): | |
| phys = cfg["physics"]; f0 = float(phys["default_field"]); c0 = cls.catalogue_numbers(cfg, f0) | |
| cat = catalogue(f0) | |
| C = P.Control | |
| wm_menu = list(cls.pack_menu(cfg, "wm")); gm_menu = list(cls.pack_menu(cfg, "gm")) | |
| rows = ( | |
| (C("on", "checkbox", "evaluate in tissue (off: bare diffusion at the same proton densities)", bool(phys["default_on"])),), | |
| (C("field_preset", "field_preset", "field preset", f"{f0:g} T", tuple(f"{float(f):g} T" for f in phys["fields"])), | |
| C("field_T", "slider", "B0 (T)", f0, minimum=0.05, maximum=12.0, step=0.001)), | |
| (C("b0_mode", "dropdown", "B0 direction (scanner frame)", list(B0_MODES)[0], tuple(B0_MODES) + (FREE_B0,)), | |
| C("theta", "slider", "polar angle from z (°)", 0, minimum=0, maximum=180, step=1), | |
| C("phi", "slider", "azimuth from x (°)", 0, minimum=0, maximum=360, step=1)), | |
| (C("relaxation", "checkbox", "relaxation (T2, T1 per tissue)", True), C("contact", "checkbox", "contact (the WM walls' ρ₂)", True), | |
| C("field", "checkbox", "field (the myelin sheath's susceptibility)", True)), | |
| tuple(C(f"T2_wm_{p}", "number", f"T2 WM {p} (ms)", c0[k]) for k, p in enumerate(WM_POOLS)), | |
| (C("T2_gm", "number", "T2 GM (ms)", c0[3]), C("T2_csf", "number", "T2 CSF (ms)", c0[4])), | |
| (C("rho2", "number", "ρ₂ of the WM walls, transverse (µm/s)", c0[5]), C("chi_iso", "number", "χ_iso of the sheath (ppm)", c0[6]), | |
| C("chi_aniso", "number", "Δχ_a of the sheath (ppm)", c0[7])), | |
| (C("m0_wm", "slider", "M0 WM (relative to CSF)", cat["m0"]["wm"], minimum=0.0, maximum=1.0, step=0.01), | |
| C("m0_gm", "slider", "M0 GM", cat["m0"]["gm"], minimum=0.0, maximum=1.0, step=0.01), | |
| C("m0_csf", "slider", "M0 CSF", cat["m0"]["csf"], minimum=0.0, maximum=1.0, step=0.01)), | |
| (C("wm_pack", "dropdown", "WM pack", wm_menu[0], tuple(wm_menu)), C("gm_pack", "dropdown", "GM pack", gm_menu[0], tuple(gm_menu))), | |
| (C("catalogue_note", "catalogue_note", value=c0[-1]), C("reset", "reset", "reset to the catalogue at this field")), | |
| ) | |
| fields = ("on", "field_T", "b0_mode", "theta", "phi", "relaxation", "contact", "field", *[f"T2_wm_{p}" for p in WM_POOLS], "T2_gm", "T2_csf", | |
| "rho2", "chi_iso", "chi_aniso", "m0_wm", "m0_gm", "m0_csf", "wm_pack", "gm_pack") | |
| return P.Panel(rows, fields, (*[f"T2_wm_{p}" for p in WM_POOLS], "T2_gm", "T2_csf", "rho2", "chi_iso", "chi_aniso"), | |
| tuple(float(f) for f in phys["fields"])) | |
| def catalogue_numbers(cls, cfg, field_T): | |
| """The catalogue's T2 per WM pool, GM and CSF (ms), the WM walls' rho2 (µm/s) and the sheath's susceptibility | |
| (ppm) at ``field_T``, and the note saying which cited field they came from.""" | |
| c = catalogue(float(field_T)) | |
| note = (f"catalogue values at {c['catalogue_field']:g} T" if abs(c["catalogue_field"] - float(field_T)) < 1e-9 | |
| else f"the catalogue has no cited relaxation at {float(field_T):g} T: nearest is {c['catalogue_field']:g} T, edit as you see fit") | |
| return [c["T2"]["wm"][p] * 1e3 for p in WM_POOLS] + [c["T2"]["gm"] * 1e3, c["T2"]["csf"] * 1e3, c["rho2"] * 1e6, | |
| c["chi_iso"] * 1e6, c["chi_aniso"] * 1e6, note] | |
| def physics_from(cls, cfg, values, scanner=None): | |
| """The panel (page units: ms, µm/s, ppm) as a :class:`BrainPhysics`: T1 and free water's D from the catalogue at | |
| the physics' field, the tiers off (bare diffusion at the same M0 and packs) when the panel is off; ``scanner`` | |
| the machine of the page's menu (a catalogue key; None: the ideal scanner), whose field and direction replace | |
| the panel's and whose slew shapes every pulse.""" | |
| on = bool(values["on"]) | |
| field_T = float(values["field_T"]) | |
| u = B0_MODES[values["b0_mode"]] if values["b0_mode"] in B0_MODES else P.b0_direction(values["theta"], values["phi"]) | |
| if scanner is not None: | |
| L = P.limits(scanner) | |
| field_T, u = float(L.field_T), tuple(float(x) for x in (L.b0_axis or P.B0_ALONG_Z)) | |
| c = catalogue(field_T) | |
| return BrainPhysics(field_T=field_T, b0_direction=tuple(float(x) for x in u), | |
| T2={"wm": {p: float(values[f"T2_wm_{p}"]) * 1e-3 for p in WM_POOLS}, "gm": float(values["T2_gm"]) * 1e-3, "csf": float(values["T2_csf"]) * 1e-3}, | |
| T1=c["T1"], rho2=float(values["rho2"]) * 1e-6, chi_iso=float(values["chi_iso"]) * 1e-6, chi_aniso=float(values["chi_aniso"]) * 1e-6, | |
| D_csf=c["D_csf"], m0={t: float(values[f"m0_{t}"]) for t in TISSUES}, wm_pack=str(values["wm_pack"]), gm_pack=str(values["gm_pack"]), | |
| relaxation=on and bool(values["relaxation"]), contact=on and bool(values["contact"]), field=on and bool(values["field"]), | |
| scanner=scanner) | |
| def knobs(cls, cfg): | |
| """B = A with one change: the field (with the catalogue's tissue at it), the field's direction, a tier off, the | |
| tissue off, the WM or GM pack, grey matter's M0, the noise, or every shell's pulse timing (a stimulated echo | |
| among them).""" | |
| out = {P.NO_KNOB: None} | |
| for f in cfg["physics"]["fields"]: | |
| out[f"field → {float(f):g} T (catalogue tissue at that field)"] = ("field", float(f)) | |
| for label in B0_MODES: | |
| out[f"B0 direction → {label}"] = ("b0", label) | |
| for tier in cls.TIERS: | |
| out[f"{tier} tier → off"] = ("tier", tier) | |
| out["tissue → off (bare diffusion)"] = ("bare", None) | |
| for tissue in ("wm", "gm"): | |
| for label in cls.pack_menu(cfg, tissue): | |
| out[f"{tissue.upper()} pack → {label}"] = ("pack", (tissue, label)) | |
| for v in cfg["physics"]["m0_gm_knob"]: | |
| out[f"M0 of GM → {float(v):g}"] = ("m0_gm", float(v)) | |
| for v in (10, 100): | |
| out[f"SNR → {v}"] = ("snr", float(v)) | |
| out["noise → off"] = ("snr", None) | |
| for name, sh in cfg["shapes"].items(): | |
| out[f"every shell's pulse timing → {name} ({sh['label']})"] = ("shape", name) | |
| return out | |
| def apply_knob(cls, cfg, change, protocol, snr_on, snr, values): | |
| kind, value = change | |
| v = dict(values) | |
| if kind in ("field", "b0", "tier") and not v["on"]: | |
| raise ValueError(f"the knob {kind!r} changes the tissue panel, which is off for A: switch it on, or choose another knob") | |
| if kind == "field": | |
| v["field_T"] = value; v.update(zip(cls.panel(cfg).catalogue, cls.catalogue_numbers(cfg, value))) | |
| elif kind == "b0": | |
| v["b0_mode"] = value | |
| elif kind == "tier": | |
| v[value] = False | |
| elif kind == "bare": | |
| v["on"] = False | |
| elif kind == "pack": | |
| v[f"{value[0]}_pack"] = value[1] | |
| elif kind == "m0_gm": | |
| v["m0_gm"] = value | |
| elif kind == "snr": | |
| snr_on = value is not None; snr = value if value is not None else snr | |
| elif kind == "shape": | |
| protocol = P.retime(protocol, value) | |
| else: | |
| raise ValueError(f"unknown knob {kind!r}") | |
| return protocol, snr_on, snr, v | |
| def tracking_controls(cls, cfg): | |
| t = cfg["tracking"] | |
| return dict(density=(1, 4, t["density"], 1, "seeds per WM voxel (density³)"), max_angle=(10, 60, t["max_angle"], 1, "max angle (°)"), | |
| step=(0.5, 2.5, t["step_mm"], 0.05, "step (mm)")) | |
| def tracking(cls, cfg, *, density, max_angle, step, key): | |
| """The tracker with :meth:`max_steps` for the step: a streamline can reach :data:`REACH_MM`.""" | |
| return P.Tracking(density=int(density), step_mm=float(step), max_angle=float(max_angle), max_steps=cls.max_steps(step), key=int(key)) | |
| def shapes_of(cls, cfg): | |
| return cls.shape_table(cfg) | |
| def max_steps(cls, step_mm): | |
| """Points per half-streamline such that :data:`REACH_MM` is reachable from a seed at either end.""" | |
| return int(np.ceil(REACH_MM / float(step_mm))) + 1 | |
| def estimated_seconds(cls, cfg, protocol, *, density, knob, n_keys, ladder, responses=(), scanner=None): | |
| """The GPU seconds a run reserves: ``[budget]`` of the configuration, measured on an L40S with the BATMAN | |
| fixture (tools/measure_brain.py): a fixed part (the worker's start, noise, the reconstruction's fixed cost, the | |
| handoff), a part per measurement (the contraction and the reconstruction scale with the rows), the tracking | |
| and the truth's tracking per density, the ladder's two rungs of contraction, B (a second run whose truth is A's), | |
| each further tracker key, and the pose responses the worker computes because the page has them not cached | |
| (``responses``: ``(state, n_meas, saves)`` per entry, :meth:`responses`, each priced ``fixed + per_meas x n_meas + | |
| per_save x saves``; a machine's encoding classes are in ``n_meas``); times the margin, within 30 and 480 s.""" | |
| b = cfg["budget"] | |
| n = protocol.n_meas; d = str(int(density)) | |
| track = float(b["track"][d]) | |
| run = b["fixed"] + b["per_meas"] * n + track | |
| secs = run + track + (b["ladder_per_meas"] * n if ladder else 0.0) + ((run - b["worker"]) if knob != P.NO_KNOB else 0.0) + track * (int(n_keys) - 1) | |
| secs += sum(b[f"response_{state}"]["fixed"] + b[f"response_{state}"]["per_meas"] * m + b[f"response_{state}"]["per_save"] * saves | |
| for state, m, saves in responses) | |
| return int(min(480, max(30, b["margin"] * secs))) | |
| # ---- the run's side ---- | |
| def validate(self, protocol, physics): | |
| """The measurements, refused by name when a class reaches beyond the windows a pack declares | |
| (:meth:`windows_needed`, from the declared table, loaded or not), when a class's echo lies beyond a loaded | |
| pack's walk (:meth:`prepare` loads the chosen ones), or when a loaded pack's embedded spec does not declare its | |
| susceptibility field (a pack without a spec, or whose magnetic pools, field tier and stored field disagree).""" | |
| meas = P.measurements(protocol, self.shapes) | |
| if physics is None: | |
| raise ValueError("a brain replay needs its physics (the proton densities and the packs), even for bare diffusion") | |
| for tissue, label in (("wm", physics.wm_pack), ("gm", physics.gm_pack)): | |
| self.windows_needed(tissue, label, meas) | |
| pk = self.packs.get((tissue, label)) | |
| if pk is None: | |
| continue | |
| T = self.walk_seconds(pk) | |
| for name in map(str, np.unique(meas.shape)): | |
| TE = float(self.shapes[name]["TE"]) | |
| if TE > T * (1 + 1e-9): | |
| raise ValueError(f"the {tissue.upper()} pack {label!r} is a {T * 1e3:g} ms walk: it cannot play the class {name!r} " | |
| f"at TE {TE * 1e3:g} ms") | |
| pk.susceptibility_field # refuses a pack whose embedded spec does not declare its field, by name | |
| return meas | |
| def sequence(self, meas, rows, pose, slew_rate=np.inf): | |
| """The ScannerSequence of the measurement rows ``rows`` (one timing class) played at the specimen's ``pose``: | |
| the image-frame directions turned by ``pose``, trapezoid pulses at the class's delta / Delta / TE with ramps at | |
| ``slew_rate`` (vertical at the default, infinite), a stimulated echo storing for TM = Delta - delta when the | |
| class is one.""" | |
| import dmipy_sim as d | |
| names = np.unique(meas.shape[rows]) | |
| if len(names) != 1: | |
| raise ValueError(f"one timing class per sequence, got {list(names)}") | |
| t = self.shapes[str(names[0])] | |
| dirs = (meas.dirs[rows] @ np.asarray(pose).T).tolist(); b = meas.bvals[rows] * 1e6 | |
| if t.get("kind", "pgse") == "pgste": | |
| return d.pgste(dirs, float(t["delta"]), float(t["Delta"]) - float(t["delta"]), bvalues=b, TE=float(t["TE"]), n_t=1000, | |
| slew_rate=float(slew_rate), ste_flip_angles=(90.0, 90.0, 90.0)) | |
| return d.pgse(dirs, float(t["delta"]), float(t["Delta"]), bvalues=b, TE=float(t["TE"]), n_t=1000, slew_rate=float(slew_rate)) | |
| def _b0_sequence(self, name, pose, slew_rate=np.inf): | |
| """The class ``name``'s b = 0 measurement as a sequence at ``pose`` (:meth:`sequence`): what the b = 0 signal | |
| and the acquisition's duration are read from.""" | |
| one = P.Measurements(np.zeros(1), np.array([[0.0, 0.0, 1.0]]), np.array([name]), np.zeros(1), np.zeros(1), np.zeros(1)) | |
| return self.sequence(one, np.arange(1), pose, slew_rate) | |
| def response_key(self, meas, physics=None): | |
| """The key of :meth:`prepare`'s result in this process's cache: the measurements (b-values, directions, | |
| timing classes) and the physics but its proton densities, which the device applies.""" | |
| import hashlib | |
| h = hashlib.sha256() | |
| for x in (meas.bvals, meas.dirs, meas.shape.astype(str)): | |
| h.update(np.ascontiguousarray(x).tobytes()) | |
| return h.hexdigest(), repr(replace(physics, m0={t: 0.0 for t in TISSUES})) | |
| def _band(meas, physics): | |
| """What the pose route compiles for: the packs, the rows per timing class and, with the field tier, the field | |
| and the sheath's susceptibility (the phase amplitude sets the expansion's band, whose shapes compile; measured: | |
| a new tissue or field direction at a compiled band costs no compilation, a run without the field tier none | |
| at all).""" | |
| names, counts = np.unique(meas.shape.astype(str), return_counts=True) | |
| field = (physics.field_T, physics.chi_iso, physics.chi_aniso) if physics.field else None | |
| return physics.wm_pack, physics.gm_pack, tuple(zip(names.tolist(), counts.tolist())), field | |
| def cached(self, meas, physics=None): | |
| """:meth:`prepare`'s result for ``meas`` at ``physics`` when this process has it (at ``physics``'s proton | |
| densities, zero seconds), else None.""" | |
| k = self._kernels.get(self.response_key(meas, physics)) | |
| return None if k is None else replace(k, m0=dict(physics.m0), seconds=0.0) | |
| def keep(self, kernels): | |
| """``{response_key: Kernels}`` computed in another process (the GPU worker) kept in this one's cache.""" | |
| for key, k in kernels.items(): | |
| self._store(key, k) | |
| def _store(self, key, k): | |
| self._kernels[key] = k | |
| while len(self._kernels) > 128: # a bounded memory: the oldest goes first | |
| self._kernels.pop(next(iter(self._kernels))) | |
| def bore_grid(self): | |
| """``(grid, index)``: the brain's voxels as they sit in the bore, the head centre (the mask's centroid) at | |
| isocentre: :meth:`dmipy_sim.phantom.Grid.from_oblique_affine` of the cropped asset's affine, whose rotation | |
| is :attr:`Asset.R` (the frame the page composes in), and ``index`` the brain voxels' own indices, each at | |
| ``affine @ [ijk, 1]`` in the scanner for a left- or right-handed affine.""" | |
| from dmipy_sim.phantom import Grid | |
| A = np.asarray(self.affine, np.float64) | |
| centre = (A[:3, :3] @ self.vox.mean(0) + A[:3, 3]) * 1e-3 | |
| grid, R = Grid.from_oblique_affine(A, self.mask.shape, isocenter_m=centre) | |
| if not np.allclose(R, self.asset.R, atol=1e-6): | |
| raise ValueError("the asset's rotation is not its affine's: the bore placement would be wrong") | |
| return grid, self.vox.copy() | |
| def encoding(self, meas, physics, rows): | |
| """``(class_of_voxel, played)``: the timing class ``rows`` of ``meas`` as the run's machine plays it at every brain | |
| voxel, binned to ``[scanners] tolerance`` of b (:func:`dmipy_sim.phantom.bore.encoding_classes` on | |
| :attr:`bore_grid`), the played acquisitions in the image frame; computed once per machine, slew and rows.""" | |
| from dmipy_sim.phantom.bore import encoding_classes | |
| import hashlib | |
| h = hashlib.sha256() | |
| for x in (meas.bvals[rows], meas.dirs[rows], meas.shape[rows].astype(str)): | |
| h.update(np.ascontiguousarray(x).tobytes()) | |
| key = (physics.scanner, physics.slew_rate, h.hexdigest()) | |
| if key not in self._encodings: | |
| grid, idx = self.bore_grid | |
| tol = float(self.cfg["scanners"]["tolerance"]) | |
| seq = self.sequence(meas, rows, np.eye(3), physics.slew_rate) | |
| out = encoding_classes(P.limits(physics.scanner), grid, seq, idx, tolerance=tol) | |
| if out is None: | |
| out = (np.zeros(len(idx), int), [seq]) | |
| self._encodings[key] = (np.asarray(out[0], int), list(out[1])) | |
| while len(self._encodings) > 16: | |
| self._encodings.pop(next(iter(self._encodings))) | |
| return self._encodings[key] | |
| def n_classes(self, meas, physics): | |
| """The encoding classes the run's acquisition needs over its timing classes (one per timing class on the ideal | |
| scanner): what its pose responses cost, in units of the commanded acquisition.""" | |
| groups = [np.flatnonzero(meas.shape == name) for name in np.unique(meas.shape)] | |
| if physics is None or physics.scanner is None: | |
| return 1 | |
| return max(len(self.encoding(meas, physics, rows)[1]) for rows in groups) | |
| def responses(self, entries): | |
| """What :meth:`prepare` costs for a run's ``entries`` ``[(meas, physics)]`` in a process forked from this one, | |
| in order: ``(state, n_meas, saves)`` for each entry not cached here, once per key; ``state`` is ``no_field`` | |
| (the field tier off: nothing to compile), ``warm`` (a band compiled here, or by an earlier entry of the run) | |
| or ``cold`` (a band to compile first); ``saves`` the WM pack's saves the entry's classes span | |
| (:meth:`saves_spanned` of :meth:`windows_needed`). ``[budget]`` prices them (``response_<state>``); on a machine | |
| an entry counts its rows once per encoding class times ``response_per_class`` (a class's rows each play their own | |
| amplitude, so the shells no longer share the expansion's bodies).""" | |
| out = []; keys = set(); bands = set(self._compiled) | |
| for meas, ph in entries: | |
| key = self.response_key(meas, ph) | |
| if key in self._kernels or key in keys: | |
| continue | |
| keys.add(key) | |
| band = self._band(meas, ph) | |
| saves = self.saves_spanned(self.segments("wm", ph.wm_pack), self.windows_needed("wm", ph.wm_pack, meas)) | |
| n = len(meas.bvals) | |
| if ph.scanner is not None: # every class's rows are their own bodies | |
| n = int(np.ceil(n * self.n_classes(meas, ph) * float(self.cfg["budget"]["response_per_class"]))) | |
| out.append(("no_field" if not ph.field else "warm" if band in bands else "cold", n, saves)) | |
| bands.add(band) | |
| return out | |
| def prepare(self, meas, physics=None): | |
| """The pose responses of a replay, on the CPU: per timing class of ``meas``, the WM and GM packs' pose | |
| responses at the run's tissue, field and specimen pose (dmipy-sim's closed form) contracted to the FOD's | |
| harmonics, free water's closed form; the true response curves for the explorer and the b = 0 weight of each | |
| tier alone (both at the first class's shells, within the windows ``meas`` reaches). Loads the windows of the | |
| chosen packs the classes reach (:meth:`windows_needed`, :meth:`pack`); computed every call and kept in this | |
| process's cache under :meth:`response_key` (:meth:`cached` reads it).""" | |
| from dmipy_sim.phantom import FreeWater | |
| from dmipy_sim.replay import so3 | |
| t0 = time.perf_counter(); stages = {} | |
| wm_pack = self.pack("wm", physics.wm_pack, self.windows_needed("wm", physics.wm_pack, meas)) | |
| stages["load_wm"] = time.perf_counter() - t0 | |
| gm_pack = self.pack("gm", physics.gm_pack, self.windows_needed("gm", physics.gm_pack, meas)) | |
| stages["load_gm"] = time.perf_counter() - t0 - stages["load_wm"] | |
| for name in np.unique(meas.shape): # the walks and the WM pack's tiers, now loaded | |
| self.validate(P.Protocol((P.Shell(str(name), 1.0, 1),), name="check"), physics) | |
| t_wm, t_gm, t_csf = physics.tissues(wm_pack, gm_pack) | |
| pose = specimen_pose(self.asset.R, physics.b0_direction) | |
| A = axis_map() | |
| n = len(meas.bvals) | |
| classes = [np.flatnonzero(meas.shape == name) for name in np.unique(meas.shape)] | |
| # per timing class, the acquisitions the voxels are played: the commanded one, or one per encoding class of the | |
| # machine (its delivered gradient, in the image frame, turned to the specimen's pose) | |
| plan, seqs = [], [] | |
| from dmipy_sim.acquisition.waveforms import rotate_waveform | |
| for rows in classes: | |
| if physics.scanner is None: | |
| cls, played = np.zeros(len(self.vox), int), [self.sequence(meas, rows, pose, physics.slew_rate)] | |
| else: | |
| cls, img = self.encoding(meas, physics, rows) | |
| played = [rotate_waveform(q, pose) for q in img] | |
| plan.append((rows, cls, len(seqs))); seqs += played | |
| prof = self._profile_sequence(meas, classes[0], pose, physics.slew_rate) # the explorer's curves ride in the same pass | |
| batch = seqs + ([prof[0]] if prof is not None else []) | |
| t = time.perf_counter() | |
| r_wm = wm_pack.pose_responses(batch, tissue=t_wm, scanner=physics.scanner_field, pose=pose, keep=(LMAX, 0)) | |
| stages["pose_wm"] = time.perf_counter() - t; t = time.perf_counter() | |
| r_gm = gm_pack.pose_responses(batch, tissue=t_gm, scanner=physics.scanner_field, pose=pose, keep=(0, 0)) | |
| stages["pose_gm"] = time.perf_counter() - t | |
| free = FreeWater(m0=1.0, tissue=t_csf) | |
| K = len(seqs) | |
| C_wm = np.zeros((K, n, N_COEF), np.complex128); e_gm = np.zeros((K, n), np.complex128); e_csf = np.zeros((K, n), np.complex128) | |
| voxel_class = np.zeros((len(self.vox), len(classes)), np.int32) | |
| for t_, (rows, cls, off) in enumerate(plan): | |
| for j in range(int(cls.max()) + 1): | |
| k_ = off + j | |
| C_wm[k_, rows] = r_wm[k_].retained(LMAX, 0) @ A | |
| e_gm[k_, rows] = r_gm[k_].retained(0, 0)[:, 0] * A[0, 0] * C00 | |
| e_csf[k_, rows] = free.response(seqs[k_]) | |
| voxel_class[:, t_] = off + cls | |
| profile = None | |
| if prof is not None: | |
| seq, shells = prof | |
| axis = so3.Distribution.axis((0.0, 0.0, 1.0), lmax=LMAX, nmax=0) | |
| wm = np.abs(r_wm[-1].compose(axis)); gm = np.abs(r_gm[-1].retained(0, 0)[:, 0]); csf = np.abs(free.response(seq)) | |
| k_ = len(PROFILE_ANGLES) | |
| profile = dict(b=shells, angles=PROFILE_ANGLES, wm=(wm[1:] / wm[0]).reshape(len(shells), k_), | |
| gm=(gm[1:] / gm[0]).reshape(len(shells), k_)[:, 0], csf=(csf[1:] / csf[0]).reshape(len(shells), k_)[:, 0]) | |
| t = time.perf_counter() | |
| b0 = self._tier_weights(meas, physics, wm_pack, gm_pack, pose, classes[0]) | |
| stages["weights"] = time.perf_counter() - t | |
| floors = {"wm": _floor(wm_pack), "gm": _floor(gm_pack)} | |
| k = Kernels(C_wm, e_gm, e_csf, voxel_class, tuple(classes), dict(physics.m0), floors, time.perf_counter() - t0, profile, b0, stages) | |
| self._store(self.response_key(meas, physics), k) | |
| self._compiled.add(self._band(meas, physics)) | |
| return k | |
| def _profile_sequence(self, meas, rows, pose, slew_rate=np.inf): | |
| """``(seq, shells)``: the explorer's profile at the shells of ``meas``'s first class (``rows``): b = 0, then | |
| per shell a gradient at each of :data:`PROFILE_ANGLES` from the image's z, where the profile's fibre lies; None | |
| without a diffusion-weighted shell.""" | |
| shells = np.unique(np.round(meas.bvals[rows][~meas.b0[rows]])) | |
| if not shells.size: | |
| return None | |
| th = np.radians(PROFILE_ANGLES) | |
| dirs = np.stack([np.sin(th), np.zeros_like(th), np.cos(th)], 1) | |
| b = np.concatenate([[0.0], np.repeat(shells, len(th))]) | |
| g = np.concatenate([[[0.0, 0.0, 1.0]], np.tile(dirs, (len(shells), 1))]) | |
| prof = P.Measurements(b, g, np.full(len(b), meas.shape[rows[0]]), np.zeros(len(b)), np.zeros(len(b)), np.zeros(len(b))) | |
| return self.sequence(prof, np.arange(len(b)), pose, slew_rate), shells | |
| def _tier_weights(self, meas, physics, wm_pack, gm_pack, pose, rows): | |
| """The b = 0 signal per tissue with each tier alone, over bare: ``relaxation`` (T2 at the echo, T1 over a | |
| stimulated echo's storage) per tissue, ``contact`` (the WM walls' survival) for WM; a tier that is off is | |
| absent. The b = 0 signal has no orientation, so one pose of the pack answers it.""" | |
| from dmipy_sim.phantom import FreeWater | |
| seq = self._b0_sequence(str(meas.shape[rows[0]]), pose, physics.slew_rate) | |
| out = {} | |
| alone = dict(relaxation=False, contact=False, field=False) | |
| if physics.relaxation: | |
| t_wm, t_gm, t_csf = replace(physics, **{**alone, "relaxation": True}).tissues(wm_pack, gm_pack) | |
| out["relaxation"] = {"wm": float(np.abs(wm_pack.replay(seq, tissue=t_wm))[0]), "gm": float(np.abs(gm_pack.replay(seq, tissue=t_gm))[0]), | |
| "csf": float(np.abs(FreeWater(m0=1.0, tissue=t_csf).response(seq))[0])} | |
| if physics.contact: | |
| t_wm, _, _ = replace(physics, **{**alone, "contact": True}).tissues(wm_pack, gm_pack) | |
| out["contact"] = {"wm": float(np.abs(wm_pack.replay(seq, tissue=t_wm))[0])} | |
| return out | |
| def replay(self, meas, physics=None, prepared=None): | |
| """The kernel route on the device: per timing class and encoding class, ``S = |F_wm C_wm^T + f_gm m0_gm e_gm + | |
| f_csf m0_csf e_csf|`` over the class's voxels in float32 torch (the device's GPU when there is one), | |
| S0-normalised; ``s0_factor`` is every voxel's b = 0 signal in M0 units (M0 = 1: pure CSF before relaxation); | |
| the floor is zero (the packs' floors are the accuracy).""" | |
| import torch | |
| if prepared is None: | |
| raise ValueError("a brain replay needs its pose responses: Brain.prepare(meas, physics) first") | |
| t0 = time.perf_counter() | |
| dev = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| k = prepared | |
| m0 = k.m0 | |
| f = torch.as_tensor(self.f, device=dev) | |
| F = torch.as_tensor(self.fod_unit, device=dev) * (f[:, 0] * m0["wm"])[:, None] | |
| g = (f[:, 1] * m0["gm"])[:, None]; c = (f[:, 2] * m0["csf"])[:, None] | |
| n_vox, n = F.shape[0], len(meas.bvals) | |
| Sr = torch.zeros((n_vox, n), dtype=torch.float32, device=dev); Si = torch.zeros_like(Sr) | |
| for t_, rows in enumerate(k.timing_rows): | |
| r = torch.as_tensor(np.asarray(rows), device=dev) | |
| for cls in np.unique(k.classes[:, t_]): | |
| sel = np.flatnonzero(k.classes[:, t_] == cls) | |
| v = torch.as_tensor(sel, device=dev) if len(sel) < n_vox else None | |
| C = k.wm[cls][rows] | |
| Cr = torch.as_tensor(C.real.astype(np.float32), device=dev); Ci = torch.as_tensor(C.imag.astype(np.float32), device=dev) | |
| eg = torch.as_tensor(np.stack([k.gm[cls][rows].real, k.gm[cls][rows].imag]).astype(np.float32), device=dev) | |
| ec = torch.as_tensor(np.stack([k.csf[cls][rows].real, k.csf[cls][rows].imag]).astype(np.float32), device=dev) | |
| Fv, gv, cv = (F, g, c) if v is None else (F[v], g[v], c[v]) | |
| re = Fv @ Cr.T + gv * eg[0][None, :] + cv * ec[0][None, :] | |
| im = Fv @ Ci.T + gv * eg[1][None, :] + cv * ec[1][None, :] | |
| if v is None: | |
| Sr[:, r] = re; Si[:, r] = im | |
| else: | |
| Sr[v[:, None], r[None, :]] = re; Si[v[:, None], r[None, :]] = im | |
| S = torch.sqrt(Sr * Sr + Si * Si) | |
| S0 = S[:, torch.as_tensor(meas.b0, device=dev)].mean(1) | |
| live = S0 > 0 | |
| dwi = torch.where(live[:, None], S / torch.where(live, S0, 1.0)[:, None], torch.nan) | |
| out = np.full(self.mask.shape + (len(meas.bvals),), np.nan, np.float32) | |
| out[self.mask] = dwi.cpu().numpy() | |
| factor = np.full(self.mask.shape, np.nan, np.float32) | |
| factor[self.mask] = torch.where(live, S0, torch.nan).cpu().numpy() | |
| return out, np.zeros(self.mask.shape, np.float32), factor, time.perf_counter() - t0 | |
| def reconstruct(self, dwi, s0_factor, meas, mask): | |
| """The configuration's reconstruction (:func:`space.pipeline.reconstruction`) and the round trip against the | |
| asset: ``msmt``, three-tissue responses and multi-tissue CSD on the signal in M0 units, or ``tournier07``, | |
| the single-fibre response and single-tissue CSD on the S0-normalised signal; each step's time under a name | |
| that says which ran. ``extras``: the path, the recovered fractions (multi-tissue only), the estimated response | |
| curves at the explorer's profile, and the round-trip numbers (:meth:`roundtrip`).""" | |
| method = P.reconstruction(self.cfg) | |
| if method == "tournier07": | |
| sh, t, (S0, model) = P.csd(dwi, meas, mask, backend=self.backend) | |
| seconds = {f"csd · single-tissue: tournier07 response + csd_tournier07_{self.backend}": t} | |
| responses = {"wm": (S0, model)}; fractions = None | |
| path = f"single-tissue Tournier 2007 (DISCO_RECONSTRUCTION=tournier07): white_matter_response_tournier07 + csd_tournier07_{self.backend}" | |
| else: | |
| sh, fractions, (S0s, models), (t_resp, t_fit) = P.csd_msmt(dwi * s0_factor[..., None], meas, mask, backend=self.backend) | |
| seconds = {"csd · responses: three_tissue_response_dhollander16": t_resp, f"csd · MT-CSD: csd_msmt_{self.backend}": t_fit} | |
| responses = dict(zip(TISSUES, zip(S0s, models))); path = f"multi-tissue: three_tissue_response_dhollander16 + csd_msmt_{self.backend}" | |
| t0 = time.perf_counter() | |
| extras = dict(reconstruction=path, fractions=None if fractions is None else fractions.astype(np.float32), | |
| responses=estimated_profile(responses, meas), roundtrip=self.roundtrip(sh, fractions)) | |
| seconds["round trip · the FOD and fractions against the input"] = time.perf_counter() - t0 | |
| return sh, seconds, extras | |
| def roundtrip(self, sh, fractions): | |
| """The recovered FOD against the asset's in the WM (the stop mask's voxels): the principal peaks' angle | |
| (median, 95 %), the Pearson of the l = 0 coefficient (the apparent fibre density), the share of voxels whose | |
| peak count (up to three) agrees; with fractions, their Pearson per tissue and median absolute difference | |
| over the brain.""" | |
| wm = self.asset.stop | |
| rec = P.fod_peaks(np.asarray(sh)[wm]); ref = P.fod_peaks(self.asset.fod[wm]) | |
| cosang = np.abs(np.einsum("ij,ij->i", rec[0][:, 0], ref[0][:, 0])) | |
| both = (rec[1][:, 0] > 0) & (ref[1][:, 0] > 0) | |
| ang = np.degrees(np.arccos(np.clip(cosang[both], 0, 1))) | |
| afd_rec = np.asarray(sh)[wm][:, 0] * (1.0 if fractions is None else fractions[wm][:, 0]) # the WM signal's l = 0 amplitude | |
| afd_in = self.asset.fod[wm][:, 0] * self.asset.fractions[wm][:, 0] | |
| out = dict(peak_angle_median=float(np.median(ang)) if ang.size else float("nan"), peak_angle_p95=float(np.quantile(ang, 0.95)) if ang.size else float("nan"), | |
| afd_pearson=P.pearson(afd_rec, afd_in), | |
| peak_count_agreement=float(np.mean(rec[2] == ref[2])), n_wm=int(wm.sum()), | |
| peak_counts_recovered=np.bincount(rec[2], minlength=4)[:4].tolist(), peak_counts_input=np.bincount(ref[2], minlength=4)[:4].tolist()) | |
| if fractions is not None: | |
| m = self.mask | |
| for k, t in enumerate(TISSUES): | |
| out[f"fraction_{t}_pearson"] = P.pearson(fractions[m][:, k], self.asset.fractions[m][:, k]) | |
| out[f"fraction_{t}_mad"] = float(np.median(np.abs(fractions[m][:, k] - self.asset.fractions[m][:, k]))) | |
| return out | |
| def tracking_inputs(self, sh, density): | |
| """``(field, seeds)``: the FOD turned from the image frame into the scanner frame on the asset's affine, the | |
| domain the white matter and the regions (leaving it ends a streamline), ``density^3`` seeds per WM voxel.""" | |
| from dmipy_sim.replay.so3 import rotate_sh | |
| from dmipy_tract import FODField, seeds_from_mask | |
| a = self.asset | |
| world = np.zeros(np.shape(sh), np.float32) | |
| world[self.mask] = rotate_sh(np.asarray(sh)[self.mask], a.R).astype(np.float32) | |
| return FODField(world, a.affine, a.stop | (a.labels > 0)), seeds_from_mask(a.stop, a.affine, density=density) | |
| def reference(self, tracking, seeds): | |
| """The truth: the asset's own FOD tracked with ``tracking`` from ``seeds`` (the same tracker, settings, seeds | |
| and key as the run), its connectome and lobar sum; computed once per (settings, seeds) and reused (B shares | |
| A's).""" | |
| key = (tracking, len(seeds), float(np.asarray(seeds).sum())) | |
| if key in self._reference: | |
| return self._reference[key], 0.0 | |
| t0 = time.perf_counter() | |
| fld, _ = self.tracking_inputs(self.asset.fod, tracking.density) | |
| tg, _ = P.track(fld, seeds, tracking, self.backend) | |
| M = self.regions.matrix(tg) | |
| ref = dict(matrix=M, lobar=self.regions.grouped(M), streamlines=len(tg)) | |
| self._reference = {key: ref} | |
| return ref, time.perf_counter() - t0 | |
| def score(self, M, reference): | |
| """Against the truth's connectome: Pearson of log(1 + count) over the region pairs (``pearson_log``), the | |
| pairs connected in the run only and in the truth only, and the same Pearson over the lobar groups' upper | |
| triangle with the diagonal (the within-group counts).""" | |
| T = reference["matrix"]; pairs = self.regions.pairs | |
| a, b = M[pairs], T[pairs] | |
| G, Gt = self.regions.grouped(M), reference["lobar"] | |
| gp = np.triu_indices(len(G)) | |
| return dict(pearson_log=P.pearson(np.log1p(a), np.log1p(b)), connected_pairs=int((a > 0).sum()), truth_pairs=int((b > 0).sum()), | |
| only_run=int(((a > 0) & (b == 0)).sum()), only_truth=int(((a == 0) & (b > 0)).sum()), | |
| lobar_pearson_log=P.pearson(np.log1p(G[gp]), np.log1p(Gt[gp])), n_pairs=len(pairs[0])) | |
| def compare(self, a, b): | |
| """A against B over the region pairs: the Pearson of log(1 + count) between the two, the pairs connected in | |
| one only, and B's score minus A's.""" | |
| pairs = self.regions.pairs | |
| ma, mb = a.matrix[pairs], b.matrix[pairs] | |
| return dict(pearson_ab=P.pearson(np.log1p(ma), np.log1p(mb)), only_a=int(((ma > 0) & (mb == 0)).sum()), | |
| only_b=int(((mb > 0) & (ma == 0)).sum()), delta_log=b.score["pearson_log"] - a.score["pearson_log"]) | |
| def score_key(self): | |
| return "pearson_log" | |
| def ladder_steps(self, physics): | |
| """The rungs below ``physics``: bare diffusion (the tiers off, the same proton densities and packs), then the | |
| tiers on one at a time.""" | |
| if physics is None or physics.bare: | |
| return [] | |
| steps = [("bare diffusion", replace(physics, **{q: False for q in self.TIERS}))] | |
| wanted = [q for q in self.TIERS if getattr(physics, q)] | |
| on = [] | |
| for tier in wanted[:-1]: | |
| on.append(tier) | |
| steps.append(("+ " + " + ".join(on), replace(physics, **{q: q in on for q in self.TIERS}))) | |
| return steps | |
| def ingredients(self, meas, physics=None, prepared=None): | |
| """The per-voxel maps of what the tissue multiplies into the b = 0 signal: ``m0`` (Σ f·M0), ``relaxation`` | |
| (the voxel's b = 0 weight from T2 / T1 alone), ``contact`` (the share the WM walls' surface relaxivity leaves | |
| of it), and ``tiers``, the per-tissue weights they are made of.""" | |
| if prepared is None: | |
| return None | |
| f = self.asset.fractions; m0 = prepared.m0 | |
| w = {t: f[..., k] * m0[t] for k, t in enumerate(TISSUES)} | |
| total = sum(w.values()) | |
| out = dict(m0=np.where(self.mask, total, np.nan).astype(np.float32), relaxation=None, contact=None, tiers=prepared.b0) | |
| with np.errstate(invalid="ignore", divide="ignore"): | |
| if "relaxation" in (prepared.b0 or {}): | |
| r = prepared.b0["relaxation"] | |
| relaxed = sum(w[t] * r[t] for t in TISSUES) | |
| out["relaxation"] = np.where(self.mask & (total > 0), relaxed / total, np.nan).astype(np.float32) | |
| if "contact" in (prepared.b0 or {}): | |
| c = prepared.b0["contact"]["wm"] | |
| out["contact"] = np.where(self.mask & (total > 0), (w["wm"] * c + w["gm"] + w["csf"]) / total, np.nan).astype(np.float32) | |
| return out | |
| def accuracy(self, res=None): | |
| a = self.asset.manifest | |
| rows = [["subject", str(a["subject"])], ["source", f"{a['source'].get('dataset')} ({a['source'].get('license')})"], | |
| ["asset", self.where], ["grid (cropped to the brain)", f"{self.mask.shape}, {int(self.mask.sum()):,} voxels, " | |
| f"{int(self.asset.stop.sum()):,} in the WM stop mask"], | |
| ["regions", f"{self.regions.n} ({len(self.regions.groups)} lobar groups)"], | |
| ["the scan's timing class", self.shapes[SCAN]["label"]], | |
| *[list(n) for n in self.asset.notes], | |
| ["replay", "exact composition of the packs' pose responses (dmipy-sim's phantom route), float32 on the device"], | |
| ["per-voxel split-half floor", "none: every voxel reads the same two walks; the packs' certified floors below are the accuracy"]] | |
| for (tissue, label), pk in self.packs.items(): | |
| seg = self.segments(tissue, label); held = self.windows_held(pk, seg) | |
| rows.append([f"{tissue.upper()} pack {label}", f"{pk.meta.get('id')}, {pk.n_walkers:,} walkers, {self.walk_seconds(pk) * 1e3:g} ms walk, " | |
| f"windows held / declared {held} / {seg['n']} of {float(seg['T']) * 1e3:g} ms " | |
| f"({self.saves_spanned(seg, held):,} of {self.saves_spanned(seg, seg['n']):,} saves at " | |
| f"{float(seg['T']) / (int(seg['n_t']) - 1) * 1e6:.4g} µs), certified floor (max) {_floor(pk):.4g}"]) | |
| return rows + super().accuracy(res) | |
| def score_text(self, tag, res): | |
| s = res.score; physics = res.physics; rt = res.extras.get("roundtrip") or {} | |
| snr = P.b0_snr(res) | |
| noise = "" if snr is None else f", SNR {res.snr:g} at M0 = 1 ({snr['median']:.1f} at b = 0 in the median voxel)" | |
| return (f"**{tag}: connectome vs the input's, Pearson log(1 + count) {s['pearson_log']:.3f}** (lobar {s['lobar_pearson_log']:.3f}; " | |
| f"{s['only_run']} pairs in the run only, {s['only_truth']} in the input's only); FOD round trip: principal peak " | |
| f"{rt.get('peak_angle_median', float('nan')):.1f}° median / {rt.get('peak_angle_p95', float('nan')):.1f}° 95 %, AFD r " | |
| f"{rt.get('afd_pearson', float('nan')):.3f} ({res.protocol.n_meas} measurements, {physics.label() if physics else 'bare'}{noise}, " | |
| f"{len(res.tractogram):,} streamlines, {res.seconds['total']:.1f} s; {res.extras.get('reconstruction', '')})") | |
| def compare_text(self, c, knob): | |
| return (f"**A vs B: connectome Pearson log(1 + count) {c['pearson_ab']:.3f}**, {c['only_a']} pairs in A only, {c['only_b']} in B only, " | |
| f"B − A {c['delta_log']:+.3f} vs the input's connectome; B = A with {knob}.") | |
| def matrices(self, M, score, reference): | |
| from .. import viewers as V | |
| return V.pair_panel([(f"this run (Pearson log {score['pearson_log']:.3f})", M), ("the input FOD tracked the same way", reference["matrix"])], | |
| f"{score['connected_pairs']} of {score['n_pairs']} pairs connected, the input's {score['truth_pairs']}; " | |
| f"{score['only_run']} in the run only, {score['only_truth']} in the input's only", names=self.regions.names) | |
| def lobar(self, M, score, reference): | |
| from .. import viewers as V | |
| return V.pair_panel([(f"lobar groups, this run (Pearson log {score['lobar_pearson_log']:.3f})", self.regions.grouped(M)), | |
| ("lobar groups, the input", reference["lobar"])], "the connectome summed over the lobar groups", | |
| names=tuple(n for n, _ in self.regions.groups)) | |
| def ingredient_layers(ing): | |
| if not ing: | |
| return [None, None, None] | |
| return [("proton density Σ f·M0 (relative to CSF)", ing["m0"], dict(vmin=0, vmax=1)), | |
| ("relaxation: the voxel's b = 0 weight from T2 / T1 at the echo", ing["relaxation"], dict(cmap="magma", vmin=0, vmax=1)) if ing.get("relaxation") is not None else None, | |
| ("contact: the share of the b = 0 signal the WM walls' ρ₂ leaves", ing["contact"], dict(cmap="magma", vmax=1)) if ing.get("contact") is not None else None] | |
| def input_peaks(self): | |
| """The asset's WM FOD principal peaks and amplitudes (image frame), for the slices.""" | |
| return P.peaks(self.asset.fod) | |
| def truth_view(self, dwi, meas, z, m): | |
| from .. import viewers as V | |
| pk, amp = self.input_peaks | |
| return V.dwi_slice(dwi, meas, z, m, peaks=pk, peak_amp=amp, overlay=True, label="the input FOD's peaks · ") | |
| def fractions_view(self, extras, z): | |
| from .. import viewers as V | |
| return V.fractions_pair(extras.get("fractions"), self.asset.fractions, z) | |
| def roundtrip_rows(self, res): | |
| rt = res.extras.get("roundtrip") or {} | |
| rows = [["reconstruction", res.extras.get("reconstruction", "")], | |
| ["principal peak vs the input's, WM voxels: median / 95 % (°)", f"{rt.get('peak_angle_median', float('nan')):.2f} / {rt.get('peak_angle_p95', float('nan')):.2f}"], | |
| ["AFD (the WM signal's l = 0 amplitude: the FOD's, times the WM fraction where one is fitted) Pearson, WM voxels", f"{rt.get('afd_pearson', float('nan')):.4f}"], | |
| ["peak count agreement (up to 3), WM voxels", f"{100 * rt.get('peak_count_agreement', float('nan')):.1f} % of {rt.get('n_wm', 0):,}"], | |
| ["voxels with 0 / 1 / 2 / 3 peaks: recovered", " / ".join(str(x) for x in rt.get("peak_counts_recovered", []))], | |
| ["voxels with 0 / 1 / 2 / 3 peaks: input", " / ".join(str(x) for x in rt.get("peak_counts_input", []))]] | |
| for t in TISSUES: | |
| if f"fraction_{t}_pearson" in rt: | |
| rows.append([f"{t.upper()} fraction vs the input's: Pearson / median |Δ|", f"{rt[f'fraction_{t}_pearson']:.4f} / {rt[f'fraction_{t}_mad']:.4f}"]) | |
| if "fraction_wm_pearson" not in rt: | |
| rows.append(["fractions", "not estimated by the single-tissue reconstruction"]) | |
| s = res.score | |
| rows += [["connectome vs the input's: Pearson log(1 + count), region pairs", f"{s['pearson_log']:.4f}"], | |
| ["the same over the lobar groups", f"{s['lobar_pearson_log']:.4f}"], | |
| ["pairs connected: run / input / run only / input only", f"{s['connected_pairs']} / {s['truth_pairs']} / {s['only_run']} / {s['only_truth']} of {s['n_pairs']}"], | |
| ["streamlines: run / input", f"{len(res.tractogram):,} / {res.reference['streamlines']:,}"]] | |
| return rows | |
| def response_view(self, prepared, extras): | |
| from .. import viewers as V | |
| est = extras.get("responses") | |
| note = "" if est is not None and est.get("gm") is not None else "estimated: WM only, the single-tissue reconstruction" | |
| return V.response_curves(prepared.profile if prepared is not None else None, est, note=note) | |
| def extra_files(self, res, out_dir, stem): | |
| """The connectome and its lobar sum as CSV, the regions' names as the header.""" | |
| import csv | |
| paths = [] | |
| for what, M, names in (("connectome", res.matrix, self.regions.names), ("lobar", self.regions.grouped(res.matrix), tuple(n for n, _ in self.regions.groups))): | |
| path = os.path.join(out_dir, f"{stem}_{what}.csv") | |
| with open(path, "w", newline="") as f: | |
| w = csv.writer(f); w.writerow([""] + list(names)) | |
| for name, row in zip(names, M): | |
| w.writerow([name] + [f"{v:g}" for v in row]) | |
| paths.append(path) | |
| return paths | |
| def truth_views(self): | |
| """The input tab: the region centroids in 3-D coloured by lobe, and a slice of the input (the WM FOD's principal | |
| peaks over the mean b = 0, the fractions, the labels).""" | |
| from .. import viewers as V | |
| a = self.asset | |
| pk, amp = P.peaks(a.fod) | |
| z = self.mask.shape[2] // 2 | |
| return (V.regions3d(self.regions, self.mask.shape), V.input_slices(a.b0, pk, amp, a.fractions, a.labels, z)) | |
| def _warm_worker(cfg, where, entries, queue): | |
| """The warm-up's process (:meth:`Brain.warm_in_background`): its own :class:`Brain` on the same configuration | |
| and asset, every ``(meas, physics)`` of ``entries`` computed on the CPU and put on ``queue`` as | |
| ``(response_key, Kernels)``, then ``None``. One line per entry on stdout (the Space's run log) with its index, | |
| the seconds it took, :attr:`Kernels.stages`, the pose expansion's phases (dmipy-sim's run record, which the | |
| expansion joins), the process's user and system CPU seconds and page faults over the entry (a system-bound entry | |
| is paying the kernel, for pages or I/O, not computing) and the seconds since the start.""" | |
| os.environ.setdefault("JAX_PLATFORMS", "cpu") | |
| t0 = time.perf_counter() | |
| from ..limits import cpu_quota, THREAD_VARS | |
| print(f"warm-up: process {os.getpid()} starts on {len(entries)} entries, {len(os.sched_getaffinity(0))} CPUs visible, " | |
| f"cgroup CPU quota {cpu_quota()}, cpuset {_cgroup_cpuset()}, threads {os.environ.get(THREAD_VARS[0], 'unset')}, " | |
| f"memory limit {_cgroup_memory_limit()}, memory.high {_cgroup_memory_high()}, memory.events {_cgroup_memory_events()}", flush=True) | |
| try: | |
| src = Brain(cfg, asset=where) | |
| print(f"warm-up: the source loaded, RSS {_rss_gb():.1f} GB ({time.perf_counter() - t0:.0f} s since the start)", flush=True) | |
| from dmipy_sim.run import Run | |
| import resource | |
| for i, (meas, ph) in enumerate(entries, 1): | |
| t = time.perf_counter(); r0 = resource.getrusage(resource.RUSAGE_SELF) | |
| with Run("brain-warm", params=dict(entry=i)) as run: # the pose expansion's phases join this record | |
| kernels = src.prepare(meas, ph) | |
| r1 = resource.getrusage(resource.RUSAGE_SELF) | |
| phases = ", ".join(f"{n} {v:.1f}" for n, v in run.phase_seconds().items()) | |
| usage = (f"user {r1.ru_utime - r0.ru_utime:.0f} s, system {r1.ru_stime - r0.ru_stime:.0f} s, " | |
| f"minor faults {r1.ru_minflt - r0.ru_minflt}, major faults {r1.ru_majflt - r0.ru_majflt}") | |
| queue.put((src.response_key(meas, ph), kernels)) | |
| windows = {t: src.windows_needed(t, label, meas) for t, label in (("wm", ph.wm_pack), ("gm", ph.gm_pack))} | |
| print(f"warm-up {i}/{len(entries)}: {len(meas.bvals)} measurements, WM {ph.wm_pack!r} ({windows['wm']} window(s)), " | |
| f"GM {ph.gm_pack!r} ({windows['gm']} window(s)), {ph.field_T:g} T along {ph.b0_direction}, " | |
| f"tiers {'on' if ph.field else 'off'}: {time.perf_counter() - t:.1f} s " | |
| f"({', '.join(f'{n} {v:.1f}' for n, v in (kernels.stages or {}).items())}; phases {phases}; {usage}), RSS {_rss_gb():.1f} GB, " | |
| f"memory.events {_cgroup_memory_events()} ({time.perf_counter() - t0:.0f} s since the start)", flush=True) | |
| finally: | |
| queue.put(None) | |
| def _rss_gb(): | |
| """This process's resident set in GB (VmRSS of /proc/self/status; nan where there is no procfs).""" | |
| try: | |
| with open("/proc/self/status") as f: | |
| for line in f: | |
| if line.startswith("VmRSS:"): | |
| return int(line.split()[1]) / 1e6 | |
| except OSError: | |
| pass | |
| return float("nan") | |
| def _cgroup_cpuset(): | |
| """The CPUs the container may run on as its cgroup states them (v2 ``cpuset.cpus.effective``), ``'unknown'`` | |
| where the file does not exist.""" | |
| try: | |
| with open("/sys/fs/cgroup/cpuset.cpus.effective") as f: | |
| return f.read().strip() or "unknown" | |
| except OSError: | |
| return "unknown" | |
| def _cgroup_memory_high(): | |
| """The cgroup's memory.high (v2), the throttling threshold below memory.max: ``'max'`` when unset, in GB | |
| otherwise, ``'unknown'`` where the file does not exist.""" | |
| try: | |
| with open("/sys/fs/cgroup/memory.high") as f: | |
| v = f.read().strip() | |
| return v if v == "max" else f"{int(v) / 1e9:.1f} GB" | |
| except OSError: | |
| return "unknown" | |
| def _cgroup_memory_events(): | |
| """The cgroup's memory.events (v2) as ``{name: count}`` (``high`` counts the times memory.high throttled the | |
| group, ``max`` the times memory.max was hit); empty where the file does not exist.""" | |
| try: | |
| with open("/sys/fs/cgroup/memory.events") as f: | |
| return {k: int(v) for k, v in (line.split() for line in f if line.strip())} | |
| except OSError: | |
| return {} | |
| def _cgroup_memory_limit(): | |
| """The container's memory limit as the cgroup states it (v2 ``memory.max``, v1 ``memory.limit_in_bytes``), in GB, | |
| ``'max'`` when unlimited, ``'unknown'`` where neither file exists.""" | |
| for path in ("/sys/fs/cgroup/memory.max", "/sys/fs/cgroup/memory/memory.limit_in_bytes"): | |
| try: | |
| with open(path) as f: | |
| v = f.read().strip() | |
| except OSError: | |
| continue | |
| return v if v == "max" else f"{int(v) / 1e9:.1f} GB" | |
| return "unknown" | |
| def _floor(pack): | |
| return float((pack.meta.get("fidelity") or {}).get("floor_max") or float("nan")) | |
| def axis_map(): | |
| """``(n_so3(8, 0), 45)``: a compact even-order density's harmonics to its SO(3) coefficients | |
| (:func:`dmipy_sim.replay.so3.axis_density_coeffs`, which is linear), so ``resp.retained(8, 0) @ axis_map()`` is the | |
| response contracted to the FOD's harmonics.""" | |
| from dmipy_sim.replay import so3 | |
| return np.stack([so3.axis_density_coeffs(e, LMAX, 0) for e in np.eye(N_COEF)], 1) | |
| def estimated_profile(responses, meas): | |
| """The estimated response curves at the explorer's profile (the shells of ``meas``, :data:`PROFILE_ANGLES`): the | |
| WM response a fibre along z gives each gradient angle, GM's and CSF's isotropic signal, each over its b = 0 | |
| value, from the models the reconstruction estimated (``{tissue: (S0, model)}``); ``None`` for a tissue it did not | |
| estimate.""" | |
| from dmipy_fit.core.acquisition_scheme import acquisition_scheme_from_bvalues | |
| rows = meas.shape == meas.shape[0] | |
| shells = np.unique(np.round(meas.bvals[rows & ~meas.b0])) | |
| if not shells.size: | |
| return None | |
| th = np.radians(PROFILE_ANGLES) | |
| dirs = np.stack([np.sin(th), np.zeros_like(th), np.cos(th)], 1) | |
| b = np.concatenate([[0.0], np.repeat(shells, len(th))]); g = np.concatenate([[[0.0, 0.0, 1.0]], np.tile(dirs, (len(shells), 1))]) | |
| i = int(np.flatnonzero(rows)[0]) | |
| sch = acquisition_scheme_from_bvalues(b * 1e6, g, delta=np.full(len(b), meas.delta[i]), Delta=np.full(len(b), meas.Delta[i]), | |
| TE=np.full(len(b), meas.TE[i]), b0_threshold=P.B0_MAX * 1e6) | |
| out = dict(b=shells, angles=PROFILE_ANGLES) | |
| for t in TISSUES: | |
| if t not in responses: | |
| out[t] = None | |
| continue | |
| _, model = responses[t] | |
| E = np.asarray(model(sch, mu=[0.0, 0.0]) if t == "wm" else model(sch), np.float64) | |
| E = (E[1:] / E[0]).reshape(len(shells), len(th)) | |
| out[t] = E if t == "wm" else E[:, 0] | |
| return out | |