From b1d0218d09992fcaf5fe984193d2dcd053038320 Mon Sep 17 00:00:00 2001 From: mcencini Date: Mon, 5 Oct 2026 18:17:16 +0200 Subject: [PATCH] Fit moba on a card, and check pics there through the reference command apps.moba made its starting point, its reference and its coil coefficients on the host whatever device the k-space was on, so a fit given CUDA k-space failed on its first Gauss-Newton step with a device mismatch. They are now made on the k-space's device. The two moba tests that fit with and without known coils run on each device. scripts/check_device.py called bartorch.tools.pics, which the apps assemble now; its three pics checks run the reference command instead. Co-Authored-By: Claude Opus 5.5 --- scripts/check_device.py | 11 ++++++----- src/bartorch/apps/_moba.py | 4 ++-- tests/test_apps.py | 15 ++++++++++----- 3 files changed, 18 insertions(+), 12 deletions(-) diff --git a/scripts/check_device.py b/scripts/check_device.py index 940ffcea..12e45408 100644 --- a/scripts/check_device.py +++ b/scripts/check_device.py @@ -17,6 +17,7 @@ import bartorch import bartorch.tools as bt +from bartorch import _reference from bartorch.linop import NUFFT CHECKS: list = [] @@ -151,10 +152,10 @@ def _pics_on_device(): kspace = bartorch.nufft(image, traj) bartorch._finufft.reset_counters() - toeplitz = bt.pics(kspace, maps, t=traj) + toeplitz = _reference.pics(kspace, maps, traj=traj) normals = bartorch._finufft.normals_built() - fast = _time(lambda: bt.pics(kspace, maps, t=traj), reps=2) - pair = _time(lambda: bt.pics(kspace, maps, t=traj, no_toeplitz=True), reps=2) + fast = _time(lambda: _reference.pics(kspace, maps, traj=traj), reps=2) + pair = _time(lambda: _reference.pics(kspace, maps, traj=traj, toeplitz=False), reps=2) if toeplitz.device.type != "cuda": return False, f"the reconstruction came back on {toeplitz.device}" @@ -175,7 +176,7 @@ def _streams(): bartorch._cuda.set_streams(streams) except ValueError: continue - timings[streams] = _time(lambda: bt.pics(kspace, maps, t=traj), reps=2) + timings[streams] = _time(lambda: _reference.pics(kspace, maps, traj=traj), reps=2) bartorch._cuda.set_streams(1) if not timings: return False, "set_streams was refused for every count" @@ -200,7 +201,7 @@ def _memcache(): idle = bartorch._cuda.free_memory() ballast = torch.empty(1 << 27, dtype=torch.complex64, device="cuda") with_torch = bartorch._cuda.free_memory() - out = bt.pics(kspace, maps, t=traj) + out = _reference.pics(kspace, maps, traj=traj) del ballast, out torch.cuda.empty_cache() after_tool = bartorch._cuda.free_memory() diff --git a/src/bartorch/apps/_moba.py b/src/bartorch/apps/_moba.py index 649d88bf..bc06c6ce 100644 --- a/src/bartorch/apps/_moba.py +++ b/src/bartorch/apps/_moba.py @@ -320,7 +320,7 @@ def packed(named): amplitude = named.get("amplitude") return model.initial( **(named if amplitude is None else {**named, "amplitude": amplitude / scale}) - ) + ).to(kspace.device) x0 = lifted(packed(values) if start is None else start) xref = x0 if reference is None else lifted(packed(_inside(model, reference))) @@ -383,7 +383,7 @@ def packed(named): _joint(sense, solved), data, lambda: sense.transform.adjoint(data), - torch.zeros(sense.ishapes[1], dtype=torch.complex64), + torch.zeros(sense.ishapes[1], dtype=torch.complex64, device=kspace.device), ) if not return_sensitivities: return maps diff --git a/tests/test_apps.py b/tests/test_apps.py index 9b72d3a0..42e72bea 100644 --- a/tests/test_apps.py +++ b/tests/test_apps.py @@ -469,20 +469,22 @@ def _relative(fitted, reference, support): return (fitted - reference)[support].abs() / reference[support] -def test_moba_fits_the_decay_behind_known_coils(): +def test_moba_fits_the_decay_behind_known_coils(device): from bartorch import nlop images, coils, t2, support, _ = _moba_phantom() model = nlop.MultiEcho(MOBA_ECHO_TIMES, (MOBA_SIZE, MOBA_SIZE)) + kspace = _moba_kspace(images, coils).to(device) - fitted = apps.moba(_moba_kspace(images, coils), model, coils, **MOBA_SETTINGS, T2=80.0) + fitted = apps.moba(kspace, model, coils.to(device), **MOBA_SETTINGS, T2=80.0) - error = _relative(fitted["T2"], t2, support) + assert fitted["T2"].device.type == device + error = _relative(fitted["T2"].cpu(), t2, support) assert float(error.median()) < 0.02 assert float(error.max()) < 0.08 -def test_moba_estimates_the_coils_with_the_decay(): +def test_moba_estimates_the_coils_with_the_decay(device): """Without sensitivities the coils are a second unknown; the T2 is still the one the images were made from, and the product of the fitted amplitude and coils -- which the data does fix, unlike either factor -- is the one the @@ -491,10 +493,13 @@ def test_moba_estimates_the_coils_with_the_decay(): images, coils, t2, support, amplitude = _moba_phantom() model = nlop.MultiEcho(MOBA_ECHO_TIMES, (MOBA_SIZE, MOBA_SIZE)) + kspace = _moba_kspace(images, coils).to(device) fitted, estimated = apps.moba( - _moba_kspace(images, coils), model, return_sensitivities=True, **MOBA_SETTINGS, T2=80.0 + kspace, model, return_sensitivities=True, **MOBA_SETTINGS, T2=80.0 ) + assert estimated.device.type == device + fitted, estimated = {name: m.cpu() for name, m in fitted.items()}, estimated.cpu() error = _relative(fitted["T2"], t2, support) assert float(error.median()) < 0.03