Fit moba on a card, and check pics there through the reference command

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()

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

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