diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 22fb2c3f..8dcd71bf 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -52,6 +52,8 @@ jobs: python-version: "3.14" cc: clang runs-on: ${{ matrix.os }} + # A hung test otherwise holds the runner for GitHub's six-hour default. + timeout-minutes: 60 defaults: run: shell: bash @@ -231,6 +233,7 @@ jobs: # without a card allows: it compiles, links, loads, and runs the host suite. cuda-build: runs-on: ubuntu-latest + timeout-minutes: 60 steps: - uses: actions/checkout@v4 with: diff --git a/AGENTS.md b/AGENTS.md index e27ff629..b90b32e8 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -45,6 +45,7 @@ would otherwise have been linked against. | `src/csrc/substitute/fftw_bind.c` | The FFTW3 functions a FINUFFT module on oneMKL calls, forwarded to the ones the host hands over. | | `src/csrc/substitute/finufft.c`, `nufft_finufft.c` | FINUFFT's and cuFINUFFT's plans through their C API, and BART's NUFFT operator built out of a pair of their plans -- and the normal, which stores one of those in BART's operator through `noncart/nufft_priv.h` rather than letting it grid one. | | `src/csrc/substitute/openmp.c` | `__kmpc_dispatch_deinit` for an OpenMP runtime that lacks it, on macOS and Windows. | +| `src/csrc/substitute/critical.c` | OpenMP's critical-section entry points, forwarded to the runtime with a record per thread of the sections held, so that a section BART's error jumped out of is left when the catch returns. | | `src/csrc/substitute/psf.c` | The three `compute_psf*` entry points, so that the adjoint transform a point spread function is comes from the substitution. | | `src/bartorch/` | The package. Every public interface is flat: `bartorch` and each subpackage export their names from `__init__.py`, and every module beneath them is private, named with a leading underscore. Public: the functions in `_fourier.py`, `_wavelet.py`, `_thresh.py`, `_util.py`, `_interp.py`, `_kspace.py` (apodization windows and readout-oversampling removal) and `_settings.py`, re-exported as `bartorch.*`; `linop/` and `nlop/` (a class per operator); `optim/` (a class per BART iteration); `priors/` (BART's regularization terms, and its denoisers); `learning/` (networks for complex images, patchwise execution, self-supervised splitting, uncertainty, and in `_training.py` the Lightning training stages, which `learning/__init__.py` imports the first time `Reconstruction` or `RandomGain` is asked for); `apps/` (BART's reconstruction pipelines, assembled from this package rather than run as commands); `cli/` (BART's command line, served by this package); `tools/` (BART's applications that have no pipeline or operator counterpart, in five sections, and two with no BART command behind them: `_correct/`, corrections of data and images outside the reconstruction, and `_motion/`, rigid motion from navigators); `io/` (CFL files, ISMRMRD raw data, DICOM and NIfTI images, with the `io` extra); `interop.py` (the deepinv adapter). Private: `_abi.py` (the ctypes signatures, generated from the header), `_lib.py` (finding and loading the library), `_marshal.py` (what an ABI argument looks like), `_backend.py` (which library serves BLAS and LAPACK), `_buffer.py` (a tensor over one of BART's buffers, host or device), `_dispatch.py` (running a command on tensors), `_operator.py` (what every operator shares), `_grid.py` (what the operations on a grid share, including BART's motion layout), `_finufft.py` and `_cuda.py` (the substitution's and the card's controls), `_catalogue.py` and `_options.py` (what BART declares, and what each option is called here), `_call.py` (the mark on a hand-written wrapper, and wrappers built from the catalogue), `_coverage.py` (where each command is exposed, or why not), `_reference.py` (the reconstruction commands the apps are tested against); inside `linop/`, `_form.py` (the encoding form and the plan it reports) and `_plan.py` (matching a composition against that form). | | `scripts/gen_abi.py` | Generates `_abi.py` from `src/csrc/include/bartorch.h`. Run after changing the header; `tests/test_abi.py` fails when the checked-in file is not what it writes. | @@ -692,6 +693,23 @@ glibc's `assert` calls `abort()` and a wrong shape takes the interpreter down. nor the libraries, and the symbol is hidden so it binds inside this library alone. +The catch is a long jump, and a jump out of an OpenMP critical section skips +the call that leaves it. BART parses options in one (`bart_getopt`, where an +unknown option or regularization term is raised) and loads every file in +another (`bart_file_access2`, where a missing input is); left held, either one +makes the next command to enter it wait for ever. +`src/csrc/substitute/critical.c` defines the runtime's entry and exit points +(`__kmpc_critical` and `__kmpc_end_critical` for clang, the four +`GOMP_critical_*` for GCC) hidden inside the library, forwards each to the +runtime the library is bound to, and keeps per thread the sections entered and +not yet left. Both catch points -- `bartorch_command` and `guarded` in +`ops.c` -- leave whatever is still held once the catch returns. A thread +entering a section it already holds leaves the stale one first, because BART's +own cleanup after a failed command runs on that thread before control is back +in the library. Explicit `omp_lock_t`s are not tracked: BART destroys them +with the objects that own them. `tests/test_core.py` runs a failed command +and then two good ones in a child process with a timeout. + ## The operator layer `LinearOperator` is the one operator class: two shapes, a forward, an adjoint, @@ -925,7 +943,9 @@ same six through the command route with `numpy.array_equal`. `mobafit` and TorchSim's bounded parameterisation: the same minimum reached by a different path. `cli/_apps.py` writes BART's closed forms in TorchSim's variables -- `-T` `(M0, R2)` as a multi-echo decay, `-I` `(M0, R1, c)` and `-L` -`(Mss, M0, R1s)` as an inversion recovery with its efficiency free -- and +`(Mss, M0, R1s)` as an inversion recovery with its efficiency free, `-G`'s +water and fat as a complex amplitude with a fat fraction and phase, `-D` as +TorchSim's diffusion decay and `-M` as its Lorentzian lines -- and converts the named maps back into rates in 1/s from times in seconds, and `moba`'s amplitudes into the units of the data scaling the command applies. Measured on the test phantoms, `mobafit` agrees with the command to 1e-05 of @@ -938,16 +958,16 @@ because the command declared it. That includes `-i`, which counts Gauss-Newton steps over BART's own coefficients and would stop a bounded fit short; `moba`'s `-C`, which counts FISTA iterations rather than conjugate gradients; `moba`'s default `-l1`, a wavelet term on the maps the app does not -carry; the gradient-echo, diffusion and simulation models; and anything off a -grid. +carry; `mobafit`'s separate water and fat decays (`-m 2`), whose two rates the +bounded fit does not separate, its phase and simulation models, and `-M` +without a start; `moba`'s gradient-echo and simulation models; and anything +off a grid. An input file that is not there is the one thing the command line names itself. -A BART command that fails while loading its arguments leaves the library unable -to serve the next call in the same process -- `ecalib` and `nufft` -handed a name with no file behind it both spin the call after them, while `fft` -does not -- so a caller who runs `main` twice would hang rather than see the -second answer. `cli._missing` checks the names against the filesystem before -BART is asked. +An app reads its inputs in Python, where a missing one is an exception rather +than an exit code, so `cli._missing` checks the names against the filesystem +before either route runs, and a missing input is reported the same way +whichever route would have served the command. Help is the catalogue's, because BART answers its own by calling `exit`, which in this process ends the interpreter. @@ -1010,7 +1030,7 @@ proximal operator and the transform BART makes of it. The solver is handed those, not a description to rebuild from, so solving twice with a term builds nothing the second time. The letters are BART's own, and a test holds every one this package offers against `grecon/optreg.c`, because a term BART does -not know is answered with `error()` and that leaves the library spinning. +not know is answered with `error()`. Three of BART's terms cannot be built alone: TGV and the two infimal convolutions extend the optimisation variable, and what they add is counted diff --git a/CMakeLists.txt b/CMakeLists.txt index 741787ae..0ce2701d 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -300,6 +300,7 @@ set(BARTORCH_SOURCES src/csrc/substitute/backend.c src/csrc/substitute/ref_blas.c src/csrc/substitute/cblas_shim.c + src/csrc/substitute/critical.c src/csrc/substitute/lapacke_shim.c src/csrc/substitute/fft.cpp src/csrc/substitute/finufft.c diff --git a/docs/api/cli.md b/docs/api/cli.md index df2034d2..26acd2ee 100644 --- a/docs/api/cli.md +++ b/docs/api/cli.md @@ -19,10 +19,12 @@ the BART command. Either route writes the output files of the BART command, in its array layout and units. A `pics` command line routed to the app writes the same values bit for bit. `mobafit` and `moba` are routed for the models whose parameters -the TorchSim model of the app represents exactly -- `-T`, `-I` and `-L` for -`mobafit`, and `-T` and `-L` with `-l2` on a Cartesian grid for `moba` -- and -the fitted maps are converted to BART's coefficients: relaxation rates in 1/s -for times in seconds, stacked along `COEFF_DIM`. Their values agree with the +the TorchSim model of the app represents exactly -- `-T`, `-I`, `-L`, `-D`, +`-M` with `--init`, and `-G` with `-m 0`, `1`, `3` or `4` for `mobafit`, and +`-T` and `-L` with `-l2` on a Cartesian grid for `moba` -- and the fitted maps +are converted to BART's coefficients: relaxation rates in 1/s for times in +seconds, off-resonance in Hz, water and fat as complex amplitudes, stacked +along `COEFF_DIM`. Their values agree with the BART command to the accuracy of the fit, since the two minimize the same objective in different parameterizations. `-i`, which counts Gauss-Newton steps in BART's parameterization, sends either command to BART. diff --git a/docs/guides/user/prerequisites.md b/docs/guides/user/prerequisites.md index 0674ec69..ed53e133 100644 --- a/docs/guides/user/prerequisites.md +++ b/docs/guides/user/prerequisites.md @@ -7,7 +7,7 @@ | Python | 3.10 or later; each wheel is tested on 3.10 and 3.14 | | PyTorch | 2.2 or later (2.3 on macOS, the first to carry `libomp.dylib`; 2.7.1 on Linux, the first whose `libgomp` is named `libgomp.so.1`), CPU or CUDA build | | NumPy, SciPy | NumPy 1.24 and SciPy 1.10 or later; MRI-NUFFT, a dependency, raises these to NumPy 2.2 and SciPy 1.13 | -| TorchSim, MRI-NUFFT | TorchSim 0.0.5 and MRI-NUFFT 1.0 or later, installed as dependencies | +| TorchSim, MRI-NUFFT | TorchSim 0.0.7 and MRI-NUFFT 1.0 or later, installed as dependencies | ## Platforms diff --git a/external/bart b/external/bart index 7fab96d7..dadb1ef7 160000 --- a/external/bart +++ b/external/bart @@ -1 +1 @@ -Subproject commit 7fab96d7369e0be35c2146533476f59c3ba14d53 +Subproject commit dadb1ef7842c6b297d09089a77a93b8b0779e188 diff --git a/pyproject.toml b/pyproject.toml index 84bd0792..3644b849 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,7 +48,7 @@ dependencies = [ # # It asks for numpy, scipy and torch, all of which are here already, and # nothing else. - "torchsim>=0.0.5", + "torchsim>=0.0.7", ] classifiers = [ "Development Status :: 2 - Pre-Alpha", diff --git a/src/bartorch/cli/__init__.py b/src/bartorch/cli/__init__.py index 833b2f79..3d6dbb3e 100644 --- a/src/bartorch/cli/__init__.py +++ b/src/bartorch/cli/__init__.py @@ -11,9 +11,8 @@ An argument the reader does not express is not an error: the whole command line goes to BART instead, which is what declared it. An input file that is not -there is the one thing named here rather than by BART, because a command that -fails while loading its arguments leaves the library unable to serve the next -call in the same process. +there is the one thing named here rather than by BART, so that it is reported +the same way whichever route would have served the command. """ from __future__ import annotations @@ -108,11 +107,8 @@ def route(name: str, argv: list[str]) -> tuple[str, dict | None]: def _missing(name: str, argv: list[str]) -> list[str]: """The input files this command line names and the filesystem does not have. - BART reports a missing input itself, and would be the one to ask -- except - that a command which fails while loading its arguments leaves the library - unable to serve the next call in the same process. A caller who runs - `main` twice would then hang rather than see the second answer, so the - names are checked before BART is asked. + Checked before either route runs: an app reads its inputs in Python, where + a missing one is an exception rather than an exit code. """ command = COMMANDS[name] try: diff --git a/src/bartorch/cli/_apps.py b/src/bartorch/cli/_apps.py index 0ef807a7..3c6f6698 100644 --- a/src/bartorch/cli/_apps.py +++ b/src/bartorch/cli/_apps.py @@ -159,8 +159,7 @@ def _start(mode: str, given: tuple[float, ...] | None) -> dict[str, Any]: """ if given is None: return {} if mode == "T" else {"inv_efficiency": 1.0} - values = list(given[: _COEFFICIENTS[mode]]) - values += [0.0] * (_COEFFICIENTS[mode] - len(values)) + values = _padded(given, _COEFFICIENTS[mode]) if mode == "T": amplitude, rate = values start = {"amplitude": amplitude, "T2": _time(rate, (1.0, 1000.0))} @@ -180,6 +179,12 @@ def _start(mode: str, given: tuple[float, ...] | None) -> dict[str, Any]: return start +def _padded(given: tuple[float, ...], count: int) -> list[float]: + """``--init`` as the command reads it: zeros past the values given.""" + values = [float(v) for v in given[:count]] + return values + [0.0] * (count - len(values)) + + def _time(rate: float, bounds: tuple[float, float]) -> float: """A rate in 1/s as a time in ms, strictly inside the model's bounds.""" if rate <= 0.0 or not bounds[0] < 1e3 / rate < bounds[1]: @@ -199,6 +204,157 @@ def _coefficients(mode: str, maps: dict[str, torch.Tensor], scale: float = 1.0) return torch.stack([m.to(torch.complex64) for m in made]) +# --- water and fat ----------------------------------------------------------- +# +# `-G` fits BART's multi-echo models in the variables of TorchSim's +# MultiGradientEchoSimulator. BART's echo times are in seconds -- its fat +# spectrum is -- so R2* comes back in 1/s and fB0 in Hz. +# +# m BART TorchSim +# 0 (W, F, fB0) fat_fraction, fat_phase, B0 +# 1 (W, F, R2*, fB0) ... and T2star +# 2 (W, R2*W, F, R2*F, fB0) (goes to the command) +# 3 (rho, R2*, fB0) T2star, B0 +# 4 (rho, fB0) B0 +# +# with W = A (1 - f), F = A f exp(i phase), R2* = 1000 / T2star and +# fB0 = -B0: the off-resonance turns the other way round in TorchSim. A +# bounded T2* keeps each step physical; a thousand milliseconds is an R2* of +# 1/s, below anything a gradient-echo train resolves. + +_WATER_FAT = { + 0: ("fat_fraction", "fat_phase", "B0"), + 1: ("fat_fraction", "fat_phase", "T2star", "B0"), + 3: ("T2star", "B0"), + 4: ("B0",), +} +_T2STAR = (1.0, 1000.0) +#: The command's five steps are over coefficients it scales to order one; four +#: bounded and periodic variables take this many to settle. +_WATER_FAT_ITERATIONS = 60 + + +def _water_fat_model(times, voxels, which: int, field_strength: float, spectrum: str): + from torchsim.simulators import MultiGradientEchoSimulator + + from bartorch import nlop + + unknown = _WATER_FAT[which] + bounds = {"fat_fraction": (0.0, 1.0), "T2star": _T2STAR} + scale = {"fat_phase": 90.0, "B0": 50.0} + sequence = MultiGradientEchoSimulator( + TE=[1e3 * float(t) for t in times], + field_strength=field_strength, + fat_spectrum=spectrum, + ) + return nlop.Bloch( + sequence, + *unknown, + shape=voxels, + bounds={name: bounds[name] for name in unknown if name in bounds}, + contrasts=len(times), + **{name: scale[name] for name in unknown if name in scale}, + ) + + +def _water_fat_start(which: int, given: tuple[float, ...] | None) -> dict[str, Any]: + """The app's start, or BART's ``--init`` in the app's variables. + + Raises + ------ + Unsupported + A start the model's bounds do not contain. + """ + unknown = _WATER_FAT[which] + if given is None: + start = {"fat_fraction": 0.5, "fat_phase": 0.0, "T2star": 30.0} + return {name: value for name, value in start.items() if name in unknown} | {"B0": 0.0} + values = _padded(given, len(unknown) + (0 if which < 3 else 1)) + if which == 0: + water, fat, frequency = values + rates = [] + elif which == 1: + water, fat, rate, frequency = values + rates = [rate] + else: + water, fat, *rates, frequency = values[0], 0.0, *values[1:] + start: dict[str, Any] = {"B0": -frequency} + if which < 3: + total = abs(water) + abs(fat) + fraction = 0.5 if total == 0.0 else abs(fat) / total + if not 0.0 < fraction < 1.0: + raise Unsupported("a start with no water or no fat has no phase between them") + start |= { + "amplitude": water / (1 - fraction), + "fat_fraction": fraction, + "fat_phase": 0.0 if fat * water >= 0 else 180.0, + } + else: + start["amplitude"] = water + for rate in rates: + start["T2star"] = _time(rate, _T2STAR) + return start + + +def _water_fat_coefficients(which: int, maps: dict[str, torch.Tensor]) -> torch.Tensor: + amplitude = maps["amplitude"] + frequency = -maps["B0"] + if which >= 3: + rates = [1e3 / maps["T2star"]] if which == 3 else [] + made = [amplitude, *rates, frequency] + else: + fraction = maps["fat_fraction"] + water = amplitude * (1 - fraction) + fat = amplitude * fraction * torch.exp(1j * torch.deg2rad(maps["fat_phase"])) + rates = [1e3 / maps["T2star"]] if which == 1 else [] + made = [water, fat, *rates, frequency] + return torch.stack([m.to(torch.complex64) for m in made]) + + +# --- diffusion and Z-spectra --------------------------------------------------- +# +# `-D` is M0 exp(enc x) over an encoding with one column, which is -b: TorchSim's +# DiffusionSimulator is M0 exp(-1e-3 b D), so b = -1000 enc and D is BART's +# coefficient in whatever units the encoding's reciprocal carries. `-M` is the +# Lorentzian Z-spectrum over the offsets the encoding lists, (M0, then amplitude, +# width and shift per pool), in the same order as TorchSim's properties. + + +def _diffusion_model(times, voxels): + from torchsim.simulators import DiffusionSimulator + + from bartorch import nlop + + b = [-1e3 * float(t) for t in times] + largest = max(abs(value) for value in b) or 1.0 + return nlop.Bloch(DiffusionSimulator(b=b), "D", shape=voxels, contrasts=len(b), D=1e3 / largest) + + +def _lorentzian_names(pools: int) -> tuple[str, ...]: + return tuple( + f"pool{index}_{what}" + for index in range(1, pools + 1) + for what in ("amplitude", "width", "shift") + ) + + +def _lorentzian_model(offsets, voxels, pools: int): + from torchsim.simulators import LorentzianSimulator + + from bartorch import nlop + + # A line's width enters squared; bounding it positive keeps the fit on + # the side the command's start puts it. + names = _lorentzian_names(pools) + return nlop.Bloch( + LorentzianSimulator(pools, offsets=[float(t) for t in offsets]), + *names, + shape=voxels, + contrasts=len(offsets), + bounds={name: (0.0, None) for name in names if name.endswith("_width")}, + ) + + def _mode(options: dict[str, list[Any]], name: str, modes: tuple[str, ...]) -> str: chosen = [mode for mode in modes if mode in options] if len(chosen) != 1: @@ -250,16 +406,29 @@ def _mobafit(options: dict[str, list[Any]], inputs: list, outputs: int) -> Call: not something the app takes; the app's own count applies. ``--liniter`` counts conjugate-gradient steps in both. Without ``--init`` the app starts from its own values, since the command's start of zero is no - relaxation time. + relaxation time; ``-M`` has no such start and is routed with one. """ - mode = _mode(options, "mobafit", ("T", "I", "L")) + pools = _only(options.pop("M")) if "M" in options else 0 + # With no model named the command fits the multi-echo one. + named = [flag for flag in ("T", "I", "L", "G", "D") if flag in options] + mode = "M" if pools else _mode(options, "mobafit", ("T", "I", "L", "G", "D")) if named else "G" + which = _only(options.pop("m", [1])) + if mode == "G" and which == 2: + # Two decays that differ only by fat's share of the signal are nearly + # one; the Gauss-Newton weight towards the middle of each bound walks + # the pair off where the command's unbounded rates stay. + raise Unsupported("mobafit -m 2 does not converge in the app's bounded variables") + if mode == "G" and which not in _WATER_FAT: + raise Unsupported(f"mobafit -m {which} is not a multi-echo model") + field_strength = _only(options.pop("field_strength", [3.0])) + spectrum = "middleton2009" if options.pop("fat_spec_0", False) else "hamilton2011" made: dict[str, Any] = {} if "a" in options: options.pop("a") made["magnitude"] = True if "liniter" in options: made["cg_maxiter"] = _only(options.pop("liniter")) - start = _start(mode, _only(options.pop("init")) if "init" in options else None) + given = _only(options.pop("init")) if "init" in options else None if options: raise Unsupported(f"mobafit {sorted(options)} is not something the app takes") if outputs != 1: @@ -274,15 +443,48 @@ def _mobafit(options: dict[str, list[Any]], inputs: list, outputs: int) -> Call: rest = np.moveaxis(y, _TE, 0) data = torch.as_tensor(rest.reshape(contrasts, -1)).to(torch.complex64) empty = data.abs().amax(dim=0) == 0 + voxels = (data.shape[1],) + + if mode == "G": + model = _water_fat_model(times, voxels, which, field_strength, spectrum) + start = _water_fat_start(which, given) + made.setdefault("iterations", _WATER_FAT_ITERATIONS) + + def coefficients(maps): + return _water_fat_coefficients(which, maps) + elif mode == "D": + model = _diffusion_model(times, voxels) + amplitude, rate = _padded(given, 2) if given is not None else (1.0, 0.0) + start = {"amplitude": amplitude, "D": rate} + + def coefficients(maps): + return torch.stack([maps["amplitude"], maps["D"].to(torch.complex64)]) + elif mode == "M": + if given is None: + raise Unsupported("mobafit -M is routed with --init; a Z-spectrum has no neutral start") + model = _lorentzian_model(times, voxels, pools) + names = _lorentzian_names(pools) + values = _padded(given, 1 + len(names)) + start = {"amplitude": values[0], **dict(zip(names, values[1:], strict=True))} + + def coefficients(maps): + return torch.stack( + [maps["amplitude"], *(maps[name].to(torch.complex64) for name in names)] + ) + else: + model = _model(mode, times, voxels) + start = _start(mode, given) + + def coefficients(maps): + return _coefficients(mode, maps) def written(maps: dict[str, torch.Tensor]) -> list[torch.Tensor]: # A voxel with no data is written as zero, which is what the command # writes for a slice with none. - made = torch.where(empty, 0, _coefficients(mode, maps)).numpy() + made = torch.where(empty, 0, coefficients(maps)).numpy() made = made.reshape(made.shape[:1] + rest.shape[1:]) return [_ours(np.moveaxis(made, 0, _COEFF))] - model = _model(mode, times, (data.shape[1],)) return Call((data, model), {**made, **start}, written) diff --git a/src/csrc/abi/api.c b/src/csrc/abi/api.c index 8f5bab11..d8010ec1 100644 --- a/src/csrc/abi/api.c +++ b/src/csrc/abi/api.c @@ -18,6 +18,7 @@ #include "include/bartorch.h" #include "substitute/backend.h" +#include "substitute/critical.h" /* The ABI spells BART's dimension, stride and flag types by their widths, so * a dimension vector crosses as it stands, without a copy. */ @@ -255,6 +256,8 @@ int bartorch_command(int argc, const char* const* argv, char* out, size_t outlen int ret = bart_command(have_out ? (int)outlen : 0, have_out ? out : NULL, argc + 1, av); + bartorch_leave_held_criticals(); + for (int i = 0; i < argc + 1; i++) free(av[i]); diff --git a/src/csrc/ops/ops.c b/src/csrc/ops/ops.c index 6539f73b..ac10ae4a 100644 --- a/src/csrc/ops/ops.c +++ b/src/csrc/ops/ops.c @@ -55,6 +55,7 @@ #include "include/bartorch.h" #include "substitute/backend.h" +#include "substitute/critical.h" struct bartorch_linop_s { const struct linop_s* op; }; @@ -88,8 +89,9 @@ static int guard_shim(int argc, char* argv[argc]) /* Run fn under BART's error catcher; an error inside returns -1. * * The message is cleared first so that what the host reads afterwards belongs - * to this call. A nested call leaves the outer one's catcher and message in - * place, which is what carries the reason out. */ + * to this call, and the critical sections an error jumped out of are left once + * the catch returns. A nested call leaves the outer one's catcher and message + * in place, which is what carries the reason out. */ static int guarded(int (*fn)(void*), void* arg) { if (error_jumper.initialized) @@ -99,7 +101,11 @@ static int guarded(int (*fn)(void*), void* arg) struct guard_call c = { fn, arg }; char* argv[1] = { (char*)&c }; - return error_catcher(guard_shim, 1, argv); + int ret = error_catcher(guard_shim, 1, argv); + + bartorch_leave_held_criticals(); + + return ret; } /* An operator as a handle the ABI can hand out. Not static: `iter.c` wraps diff --git a/src/csrc/substitute/critical.c b/src/csrc/substitute/critical.c new file mode 100644 index 00000000..d6104e16 --- /dev/null +++ b/src/csrc/substitute/critical.c @@ -0,0 +1,253 @@ +/* + * OpenMP critical sections, entered and left through the runtime, with a + * record per thread of the sections it holds. + * + * BART's `error` leaves a command by a long jump. A jump out of a critical + * section skips the call that leaves it, so the section stays locked and the + * next thread to enter it waits for ever: a bad option leaves `bart_getopt` + * held and every command after it hangs, a missing input does the same to + * the section every file is loaded in. So every critical section compiled + * into the library is entered through the definitions below, which forward to + * the runtime and keep, per thread, the sections entered and not yet left; + * the two places that catch BART's errors leave whatever is still held once + * the catch returns (`bartorch_leave_held_criticals`). + * + * A thread may not enter a section it is already inside, so one it enters + * while holding can only be held from before a jump, and is left first. That + * is the case of BART's cleanup after a failed command, which runs on the + * same thread before the catch has returned to the library. + * + * clang calls `__kmpc_critical` and GCC `GOMP_critical_start` and + * `GOMP_critical_name_start`; both sets are defined, and each forwards to the + * runtime the library is bound to -- on Windows the one torch loads + * (cmake/openmp.cmake), elsewhere the image `omp_get_max_threads` resolved + * to. The definitions are hidden, so they bind the library's own objects, + * BART's and FINUFFT's, and nothing outside it. + */ +#ifndef _GNU_SOURCE +#define _GNU_SOURCE /* dladdr, on glibc */ +#endif + +#include "substitute/critical.h" + +#ifdef _OPENMP + +#include +#include +#include +#include +#include + +#include + +#ifdef _WIN32 +#include +#else +#include +#endif + +#define HIDDEN __attribute__((visibility("hidden"))) + +typedef void (*any_fn)(void); +typedef void (*kmpc_fn)(void* loc, int32_t gtid, void* crit); +typedef void (*kmpc_hint_fn)(void* loc, int32_t gtid, void* crit, uint32_t hint); +typedef void (*gomp_fn)(void); +typedef void (*gomp_name_fn)(void** name); + +enum entry { + KMPC_CRITICAL, KMPC_CRITICAL_WITH_HINT, KMPC_END_CRITICAL, + GOMP_START, GOMP_END, GOMP_NAME_START, GOMP_NAME_END, + ENTRIES +}; + +static const char* const entry_names[ENTRIES] = { + "__kmpc_critical", "__kmpc_critical_with_hint", "__kmpc_end_critical", + "GOMP_critical_start", "GOMP_critical_end", + "GOMP_critical_name_start", "GOMP_critical_name_end", +}; + +static any_fn lookup(const char* name) +{ + any_fn fn = NULL; +#ifdef _WIN32 + HMODULE rt = GetModuleHandleA("libiomp5md.dll"); + + if (NULL != rt) + fn = (any_fn)(void (*)(void))GetProcAddress(rt, name); +#else + Dl_info info; + + if (dladdr((void*)omp_get_max_threads, &info) && (NULL != info.dli_fname)) { + + void* rt = dlopen(info.dli_fname, RTLD_LAZY | RTLD_NOLOAD); + + if (NULL != rt) { + + fn = (any_fn)dlsym(rt, name); + dlclose(rt); + } + } + + if (NULL == fn) + fn = (any_fn)dlsym(RTLD_DEFAULT, name); +#endif + if (NULL == fn) { + + fprintf(stderr, "bartorch: the OpenMP runtime does not export %s\n", name); + abort(); + } + + return fn; +} + +/* Resolved once; threads racing here all find the same answer. */ +static any_fn runtime(enum entry which) +{ + static any_fn cached[ENTRIES]; + + any_fn fn = __atomic_load_n(&cached[which], __ATOMIC_ACQUIRE); + + if (NULL == fn) { + + fn = lookup(entry_names[which]); + __atomic_store_n(&cached[which], fn, __ATOMIC_RELEASE); + } + + return fn; +} + + +enum kind { KMPC, GOMP, GOMP_NAME }; + +struct held { + + enum kind kind; + void* lock; + void* loc; + int32_t gtid; +}; + +/* Deeper than any nesting of BART's sections; one past it goes unrecorded, + * which leaves it as the runtime alone would. */ +#define MAX_HELD 32 + +static _Thread_local struct held held[MAX_HELD]; +static _Thread_local int nheld; + +/* The lock GCC's unnamed section is, as far as the record is concerned. */ +static char gomp_unnamed; + +static void leave(const struct held* h) +{ + switch (h->kind) { + + case KMPC: + ((kmpc_fn)runtime(KMPC_END_CRITICAL))(h->loc, h->gtid, h->lock); + break; + + case GOMP: + ((gomp_fn)runtime(GOMP_END))(); + break; + + case GOMP_NAME: + ((gomp_name_fn)runtime(GOMP_NAME_END))((void**)h->lock); + break; + } +} + +static void remember(enum kind kind, void* lock, void* loc, int32_t gtid) +{ + if (nheld < MAX_HELD) + held[nheld++] = (struct held){ kind, lock, loc, gtid }; +} + +static bool forget(const void* lock, struct held* out) +{ + for (int i = nheld - 1; i >= 0; i--) { + + if (held[i].lock != lock) + continue; + + if (NULL != out) + *out = held[i]; + + memmove(&held[i], &held[i + 1], (size_t)(nheld - i - 1) * sizeof held[0]); + nheld--; + + return true; + } + + return false; +} + +static void leave_if_held(const void* lock) +{ + struct held stale; + + if (forget(lock, &stale)) + leave(&stale); +} + +void bartorch_leave_held_criticals(void) +{ + while (0 < nheld) { + + nheld--; + leave(&held[nheld]); + } +} + + +HIDDEN void __kmpc_critical(void* loc, int32_t gtid, void* crit) +{ + leave_if_held(crit); + ((kmpc_fn)runtime(KMPC_CRITICAL))(loc, gtid, crit); + remember(KMPC, crit, loc, gtid); +} + +HIDDEN void __kmpc_critical_with_hint(void* loc, int32_t gtid, void* crit, uint32_t hint) +{ + leave_if_held(crit); + ((kmpc_hint_fn)runtime(KMPC_CRITICAL_WITH_HINT))(loc, gtid, crit, hint); + remember(KMPC, crit, loc, gtid); +} + +HIDDEN void __kmpc_end_critical(void* loc, int32_t gtid, void* crit) +{ + forget(crit, NULL); + ((kmpc_fn)runtime(KMPC_END_CRITICAL))(loc, gtid, crit); +} + +HIDDEN void GOMP_critical_start(void) +{ + leave_if_held(&gomp_unnamed); + ((gomp_fn)runtime(GOMP_START))(); + remember(GOMP, &gomp_unnamed, NULL, 0); +} + +HIDDEN void GOMP_critical_end(void) +{ + forget(&gomp_unnamed, NULL); + ((gomp_fn)runtime(GOMP_END))(); +} + +HIDDEN void GOMP_critical_name_start(void** name) +{ + leave_if_held(name); + ((gomp_name_fn)runtime(GOMP_NAME_START))(name); + remember(GOMP_NAME, name, NULL, 0); +} + +HIDDEN void GOMP_critical_name_end(void** name) +{ + forget(name, NULL); + ((gomp_name_fn)runtime(GOMP_NAME_END))(name); +} + +#else + +void bartorch_leave_held_criticals(void) +{ +} + +#endif diff --git a/src/csrc/substitute/critical.h b/src/csrc/substitute/critical.h new file mode 100644 index 00000000..a32124cf --- /dev/null +++ b/src/csrc/substitute/critical.h @@ -0,0 +1,10 @@ +/* Internal: the OpenMP critical sections each thread holds. */ +#ifndef BARTORCH_CRITICAL_H +#define BARTORCH_CRITICAL_H + +/* Leaves every critical section the calling thread entered and has not left. + * Called where BART's error catcher returns: an error raised inside a section + * jumps past the call that would have left it. */ +void bartorch_leave_held_criticals(void); + +#endif diff --git a/tests/test_cli.py b/tests/test_cli.py index 845566a8..4dbd678c 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -121,6 +121,101 @@ def _series(times: np.ndarray, series: np.ndarray) -> None: } +#: Multi-gradient-echo series. BART's fat is Hamilton et al.'s six peaks at +#: 3 T, with the echo times in seconds. +GRADIENT_ECHO_TIMES = 1.1e-3 * (np.arange(8) + 1) +_HAMILTON = [(-3.80, 0.086), (-3.40, 0.537), (-2.60, 0.165), (-1.94, 0.046), (-0.39, 0.052), + (0.60, 0.114)] # fmt: skip +_FAT = sum( + a * np.exp(2j * np.pi * 42.57747892e6 * 3.0 * p * 1e-6 * GRADIENT_ECHO_TIMES) + for p, a in _HAMILTON +) +WATER, FAT = _halves(0.9, 0.4) * PHASE, _halves(0.1, 0.5) * PHASE * np.exp(0.5j) +R2S, FAT_R2S, FB0 = _halves(30.0, 50.0), _halves(60.0, 80.0), _halves(20.0, -10.0) + + +def _gradient_echo(water, fat, water_rate, fat_rate, frequency): + t = GRADIENT_ECHO_TIMES + return ( + water[..., None] * np.exp(-t * water_rate[..., None]) + + fat[..., None] * _FAT * np.exp(-t * fat_rate[..., None]) + ) * np.exp(2j * np.pi * frequency[..., None] * t) + + +_NONE = np.zeros((FIT, FIT)) +_FITS |= { + "G0": ( + ["-G", "-m", "0"], + GRADIENT_ECHO_TIMES, + _gradient_echo(WATER, FAT, _NONE, _NONE, FB0), + [WATER, FAT, FB0], + ), + "G1": ( + ["-G"], + GRADIENT_ECHO_TIMES, + _gradient_echo(WATER, FAT, R2S, R2S, FB0), + [WATER, FAT, R2S, FB0], + ), + "G3": ( + ["-G", "-m", "3"], + GRADIENT_ECHO_TIMES, + _gradient_echo(WATER, _NONE, R2S, R2S, FB0), + [WATER, R2S, FB0], + ), + "G4": ( + ["-G", "-m", "4"], + GRADIENT_ECHO_TIMES, + _gradient_echo(WATER, _NONE, _NONE, _NONE, FB0), + [WATER, FB0], + ), + "G0 from a start": ( + ["-G", "-m", "0", "--init", "1:0.3:0"], + GRADIENT_ECHO_TIMES, + _gradient_echo(WATER, FAT, _NONE, _NONE, FB0), + [WATER, FAT, FB0], + ), + "G1 from a start": ( + ["-G", "--init", "1:0.3:30:0"], + GRADIENT_ECHO_TIMES, + _gradient_echo(WATER, FAT, R2S, R2S, FB0), + [WATER, FAT, R2S, FB0], + ), + "G3 from a start": ( + ["-G", "-m", "3", "--init", "1:30:0"], + GRADIENT_ECHO_TIMES, + _gradient_echo(WATER, _NONE, R2S, R2S, FB0), + [WATER, R2S, FB0], + ), +} + +#: Diffusion: the encoding is -b in s/mm^2, D in mm^2/s. +B_VALUES = np.array([0.0, 200.0, 500.0, 1000.0, 1500.0, 2000.0]) +DIFFUSIVITY = _halves(0.8e-3, 1.6e-3) +#: One Lorentzian pool over offsets in ppm. +OFFSETS = np.linspace(-5.0, 5.0, 21) +DEPTH, WIDTH, SHIFT = _halves(0.8, 0.6), _halves(1.5, 2.5), _halves(0.1, -0.2) +_FITS |= { + "D": ( + ["-D"], + -B_VALUES, + M0 * np.exp(-B_VALUES * DIFFUSIVITY[..., None]), + [np.full(DIFFUSIVITY.shape, M0), DIFFUSIVITY], + ), + "M": ( + ["-M", "1", "--init", "1:0.7:2:0"], + OFFSETS, + M0 + * ( + 1 + - DEPTH[..., None] + * (WIDTH[..., None] / 2) ** 2 + / ((WIDTH[..., None] / 2) ** 2 + (OFFSETS - SHIFT[..., None]) ** 2) + ), + [np.full(DEPTH.shape, M0), DEPTH, WIDTH, SHIFT], + ), +} + + def _relative(ours: np.ndarray, reference: np.ndarray) -> float: return float(np.abs(ours - reference).max() / np.abs(reference).max()) @@ -153,14 +248,29 @@ def test_mobafit_writes_the_commands_coefficients_to_the_tolerance_of_a_fit( @pytest.mark.parametrize( "line", [ - ["mobafit", "t", "y", "x"], - ["mobafit", "-G", "t", "y", "x"], + ["mobafit", "-P", "t", "y", "x"], + ["mobafit", "-G", "--fB0-init", "t", "y", "x"], + ["mobafit", "-G", "-m", "2", "t", "y", "x"], + ["mobafit", "-G", "-m", "5", "t", "y", "x"], + ["mobafit", "-G", "--init", "1:0:30:0", "t", "y", "x"], + ["mobafit", "-M", "1", "t", "y", "x"], ["mobafit", "-T", "-i", "5", "t", "y", "x"], ["mobafit", "-T", "--init", "1:0", "t", "y", "x"], ["mobafit", "-T", "--scale", "1:10", "t", "y", "x"], ["mobafit", "-T", "t", "y", "x", "covariance"], ], - ids=["default MGRE", "-G", "-i", "R2 start of zero", "--scale", "covariance"], + ids=[ + "-P", + "--fB0-init", + "-m 2", + "-m 5", + "water and no fat", + "Z-spectrum with no start", + "-i", + "R2 start of zero", + "--scale", + "covariance", + ], ) def test_a_mobafit_the_app_cannot_express_goes_to_the_command(line, tmp_path, monkeypatch): """The multi-echo gradient-echo models, BART's step count over its own @@ -454,17 +564,27 @@ def test_a_command_with_no_app_runs_as_itself(_dataset): assert readcfl("img").shape == (SIZE, SIZE, 1, COILS) -def test_an_input_that_is_not_there_is_named_before_bart_is_asked(_dataset, capsys): - """BART would report it, and would be the one to ask -- except that a - command which fails while loading its arguments leaves the library unable - to serve the next call in the same process.""" +def test_one_process_writes_the_headers_of_many_files(tmp_path, monkeypatch): + """A header is a few dozen formatted writes to the file's descriptor, and + the Windows C runtime holds a fixed number of streams per process, so a + write that opened one over the descriptor would fail a few dozen files in.""" + monkeypatch.chdir(tmp_path) + for index in range(64): + code, _, failure = run_command(["ones", "3", "2", "2", "2", f"x{index}"]) + assert code == 0, f"file {index}: {failure}" + assert readcfl("x63").shape == (2, 2, 2) + + +def test_an_input_that_is_not_there_is_named_before_either_route_runs(_dataset, capsys): + """The first line would reach the app, which reads its inputs in Python, + and the second BART, which reports a missing input itself; both are named + the same way.""" assert main(["pics", "nosuchfile", "maps", "out"]) == 1 assert "no such input: nosuchfile" in capsys.readouterr().err assert main(["pics", "-t", "nosuchtraj", "ksp", "maps", "out"]) == 1 assert "no such input: nosuchtraj" in capsys.readouterr().err - # And the library still answers, which is the whole point. assert main(["fft", "-i", "6", "ksp", "img"]) == 0 diff --git a/tests/test_core.py b/tests/test_core.py index 34d2cb4a..22a61482 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -4,6 +4,8 @@ import ctypes import os +import subprocess +import sys import numpy as np import pytest @@ -97,6 +99,34 @@ def test_registry_is_empty_after_a_failure(): assert library().bartorch_unlink_all() == 0 +_FAILURES = { + "an unknown option": ["ones", "-q", "3", "2", "2", "2", "x"], + "an unknown regularization term": ["pics", "-R", "Z:7:0:0.01", "ksp", "maps", "x"], + "a missing input": ["fft", "-i", "6", "nosuchfile", "x"], +} + + +@pytest.mark.parametrize("failing", _FAILURES.values(), ids=_FAILURES.keys()) +def test_a_failed_command_leaves_the_next_one_free_to_run(failing, tmp_path): + """BART's error leaves a command by a long jump, and each of these is raised + inside an OpenMP critical section: the options are parsed in one and every + file is loaded in another. A jump out of a section leaves it locked unless + the library leaves it. The commands run in a child process, so that a + section left locked fails this test on its timeout rather than hanging the + suite.""" + script = f""" +import os +from bartorch._dispatch import run_command +os.chdir({str(tmp_path)!r}) +code, _, _ = run_command({failing!r}) +assert code != 0, "the command was meant to fail" +for line in (["ones", "3", "2", "2", "2", "y"], ["fft", "-i", "6", "y", "z"]): + code, _, failure = run_command(line) + assert code == 0, failure +""" + subprocess.run([sys.executable, "-c", script], timeout=120, check=True) + + def test_rss_matches_numpy(): x = torch.randn(8, 16, 16, dtype=torch.complex64) y = bartorch.rss(x, axes=0) diff --git a/tests/test_signal_models.py b/tests/test_signal_models.py new file mode 100644 index 00000000..9de3d92a --- /dev/null +++ b/tests/test_signal_models.py @@ -0,0 +1,211 @@ +"""TorchSim's closed forms against the signal models BART computes. + +Every closed-form signal model BART ships -- the ``signal`` tool's and the +nonlinear operators ``mobasig`` evaluates for ``mobafit`` and ``moba`` -- has a +TorchSim simulator, and these tests hold each pair to single-precision +round-off. The units are each library's own: BART takes seconds, rates in +1/s and ``fB0`` in Hz with a positive phase, TorchSim milliseconds and ``B0`` +with the opposite sign. BART's Look-Locker curves are the longitudinal +magnetization; TorchSim's carry the readout's ``sin(flip)``. +""" + +import numpy as np +import pytest +import torch + +from bartorch import _call + +torchsim = pytest.importorskip("torchsim") +simulators = pytest.importorskip("torchsim.simulators") +if not hasattr(simulators, "MultiGradientEchoSimulator"): + pytest.skip("this TorchSim has no closed forms for BART's models", allow_module_level=True) + +#: Both commands are private: what they return is a curve, not a model a fit +#: can be built on. They are reachable for exactly this. +signal = _call.build("signal", __name__) +mobasig = _call.build("mobasig", __name__) + +T1_S, T2_S = 0.9, 0.08 + + +def _agree(ours, theirs, rtol=2e-6): + ours = np.asarray(ours).reshape(-1) + theirs = np.asarray(theirs).reshape(-1) + assert np.max(np.abs(ours - theirs)) <= rtol * np.max(np.abs(theirs)) + + +def _coefficients(*values): + return torch.tensor(values, dtype=torch.complex64).reshape(len(values), 1, 1, 1, 1, 1, 1) + + +def _encoding(values): + return torch.tensor(np.asarray(values), dtype=torch.complex64).reshape(-1, 1, 1, 1, 1, 1) + + +# --- the signal tool ------------------------------------------------------------ + + +@pytest.mark.parametrize( + "flags", + [{"I": True}, {"I": True, "short_TR_LL_approx": True}, {"s": True}], + ids=["from-equilibrium", "short-TR", "from-steady-state"], +) +def test_the_look_locker_recovery_is_barts(flags): + tr, flip, n = 0.005, 8.0, 60 + theirs = signal(F=True, r=tr, f=flip, n=n, flag_1=(T1_S, T1_S, 1), **flags) + sequence = simulators.LookLockerSimulator( + flip=flip, + TR=tr * 1e3, + TI=tr * 1e3 * np.arange(n), + short_TR=flags.get("short_TR_LL_approx", False), + from_steady_state=flags.get("s", False), + ) + ours = sequence.simulate(T1=T1_S * 1e3) / np.sin(np.deg2rad(flip)) + _agree(ours, theirs) + + +def test_the_molli_blocks_are_barts(): + tr, flip = 0.005, 30.0 + theirs = signal(M=True, I=True, r=tr, f=flip, n=60, b=3, t=0.5, flag_1=(T1_S, T1_S, 1)) + sequence = simulators.MOLLISimulator( + flip=flip, TR=tr * 1e3, readouts=(20, 20, 20), recovery=500.0 + ) + _agree(sequence.simulate(T1=T1_S * 1e3) / np.sin(np.deg2rad(flip)), theirs) + + +def test_the_ir_bssfp_recovery_is_barts(): + tr, flip, n = 0.0045, 45.0, 60 + theirs = signal( + B=True, I=True, r=tr, f=flip, n=n, flag_1=(T1_S, T1_S, 1), flag_2=(T2_S, T2_S, 1) + ) + sequence = simulators.IRbSSFPSimulator(flip=flip, TI=tr * 1e3 * np.arange(n)) + _agree(sequence.simulate(T1=T1_S * 1e3, T2=T2_S * 1e3), theirs) + + +def test_the_spin_echo_is_barts(): + theirs = signal(S=True, r=2.0, e=0.01, n=6, flag_1=(T1_S, T1_S, 1), flag_2=(T2_S, T2_S, 1)) + sequence = simulators.SpinEchoSimulator(TE=10.0 * np.arange(6), TR=2000.0) + _agree(sequence.simulate(T1=T1_S * 1e3, T2=T2_S * 1e3), theirs) + + +def test_the_inversion_recovery_spin_echo_is_barts(): + # With -I, BART steps the inversion time and holds the echo time. + theirs = signal( + S=True, I=True, r=5.0, e=0.02, i=0.1, n=6, + flag_1=(T1_S, T1_S, 1), flag_2=(T2_S, T2_S, 1), + ) # fmt: skip + sequence = simulators.SpinEchoSimulator(TE=20.0, TR=5000.0, TI=100.0 * np.arange(6)) + _agree(sequence.simulate(T1=T1_S * 1e3, T2=T2_S * 1e3), theirs) + + +@pytest.mark.parametrize("refocusing", [180.0, 120.0]) +def test_the_stimulated_echo_train_is_barts_generating_function(refocusing): + # BART's first sample is the excitation itself; the echoes follow it. + theirs = signal( + T=True, e=0.01, n=13, f=refocusing, flag_1=(1.0, 1.0, 1), flag_2=(T2_S, T2_S, 1) + ) + sequence = simulators.FSESimulator(flip=refocusing * torch.ones(12), ESP=10.0, TR=1e6) + ours = sequence.simulate(T1=1000.0, T2=T2_S * 1e3) + _agree(np.abs(np.asarray(ours)), np.abs(np.asarray(theirs).reshape(-1)[1:]), rtol=1e-4) + + +def test_water_and_fat_are_barts(): + te, n = 0.0012, 8 + theirs = signal( + G=True, fat=True, e=te, n=n, d=0.25, flag_2=(0.03, 0.03, 1), flag_0=(20.0, 20.0, 1) + ) + sequence = simulators.MultiGradientEchoSimulator(TE=te * 1e3 * np.arange(n)) + _agree(sequence.simulate(fat_fraction=0.25, T2star=30.0, B0=-20.0), theirs) + + +@pytest.mark.parametrize("labelling", ["continuous", "pulsed"]) +def test_the_arterial_spin_labelling_curve_is_barts(labelling): + pulsed = labelling == "pulsed" + theirs = signal( + A=True, pulsed=pulsed, r=0.1, n=50, flag_1=(1.4, 1.4, 1), flag_6=(60.0, 60.0, 1) + ) + sequence = simulators.ASLSimulator( + t=100.0 * np.arange(50), tau=800.0 if pulsed else 1800.0, labelling=labelling + ) + _agree(sequence.simulate(CBF=60.0, ATT=1800.0, T1=1400.0), theirs, rtol=2e-6) + + +# --- the operators mobafit and moba fit ------------------------------------------- + +_W, _F = 0.7 - 0.2j, 0.3 + 0.1j +_TE_S = 1e-3 * np.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6]) + + +def _water_and_fat(): + """BART's complex W and F as a complex scale, a fat fraction and a fat phase.""" + fraction = abs(_F) / (abs(_W) + abs(_F)) + scale = _W / (1 - fraction) + phase = np.rad2deg(np.angle(_F / (scale * fraction))) + return scale, {"fat_fraction": fraction, "fat_phase": phase} + + +@pytest.mark.parametrize( + ("model", "coefficients", "properties"), + [ + (0, (_W, _F, 20.0), {"B0": -20.0}), + (1, (_W, _F, 30.0, 20.0), {"T2star": 1e3 / 30.0, "B0": -20.0}), + ( + 2, + (_W, 30.0, _F, 50.0, 20.0), + {"T2star": 1e3 / 30.0, "fat_T2star": 1e3 / 50.0, "B0": -20.0}, + ), + ], + ids=["WF", "WFR2S", "WF2R2S"], +) # fmt: skip +def test_the_water_fat_models_are_barts(model, coefficients, properties): + theirs = mobasig(_coefficients(*coefficients), _encoding(_TE_S), G=True, m=model) + scale, fat = _water_and_fat() + sequence = simulators.MultiGradientEchoSimulator(TE=_TE_S * 1e3) + _agree(scale * np.asarray(sequence.simulate(**properties, **fat)), theirs) + + +@pytest.mark.parametrize( + ("model", "coefficients", "properties"), + [(3, (_W, 30.0, 20.0), {"T2star": 1e3 / 30.0, "B0": -20.0}), (4, (_W, 20.0), {"B0": -20.0})], + ids=["R2S", "PHASEDIFF"], +) +def test_the_water_only_models_are_barts(model, coefficients, properties): + theirs = mobasig(_coefficients(*coefficients), _encoding(_TE_S), G=True, m=model) + sequence = simulators.MultiGradientEchoSimulator(TE=_TE_S * 1e3) + _agree(_W * np.asarray(sequence.simulate(**properties)), theirs) + + +def test_the_other_fat_spectrum_at_another_field_is_barts(): + theirs = mobasig( + _coefficients(_W, _F, 30.0, 20.0), _encoding(_TE_S), + G=True, m=1, fat_spec_0=True, field_strength=1.5, + ) # fmt: skip + scale, fat = _water_and_fat() + sequence = simulators.MultiGradientEchoSimulator( + TE=_TE_S * 1e3, fat_spectrum="middleton2009", field_strength=1.5 + ) + _agree(scale * np.asarray(sequence.simulate(T2star=1e3 / 30.0, B0=-20.0, **fat)), theirs) + + +def test_the_diffusion_decay_is_barts(): + # mobasig -D is M0 exp(enc . x): the encoding is -b and the coefficient D. + b = np.array([0.0, 500.0, 1000.0, 2000.0]) + theirs = mobasig(_coefficients(2.0, 0.8e-3), _encoding(-b), D=True) + _agree(2.0 * np.asarray(simulators.DiffusionSimulator(b=b).simulate(D=0.8)), theirs) + + +def test_the_lorentzian_z_spectrum_is_barts(): + offsets = np.linspace(-5.0, 5.0, 21) + pools = ((0.8, 1.5, 0.1), (0.1, 2.0, 3.5)) + theirs = mobasig( + _coefficients(1.2, *(value for pool in pools for value in pool)), + _encoding(offsets), + M=True, + ) + named = {} + for index, (amplitude, width, shift) in enumerate(pools, start=1): + named[f"pool{index}_amplitude"] = amplitude + named[f"pool{index}_width"] = width + named[f"pool{index}_shift"] = shift + sequence = simulators.LorentzianSimulator(len(pools), offsets=offsets) + _agree(1.2 * np.asarray(sequence.simulate(**named)), theirs)