Tests for signal models on gpu

src/mrpro/operators/WaveletOp.py

@@ -3,7 +3,6 @@
 from collections.abc import Sequence
 from typing import Literal
 
-import numpy as np
 import torch
 from ptwt.conv_transform import wavedec, waverec
 from ptwt.conv_transform_2 import wavedec2, waverec2
@@ -366,7 +365,7 @@ def _stacked_tensor_to_coeff(self, coefficients_stack: torch.Tensor) -> list[tor
             3D: [aaa, aad_n, ada_n, add_n, ..., ..., aad_1, ada_1, add_1, ...]
         """
         coefficients = torch.split(
-            coefficients_stack, [int(np.prod(shape)) for shape in self.coefficients_shape], dim=-1
+            coefficients_stack, [int(torch.prod(torch.as_tensor(shape))) for shape in self.coefficients_shape], dim=-1
         )
         return [
             torch.reshape(coeff, (*coeff.shape[:-1], *shape))

tests/__init__.py

@@ -1,2 +1,10 @@
 from ._RandomGenerator import RandomGenerator
-from .helper import relative_image_difference, dotproduct_adjointness_test, operator_isometry_test, linear_operator_unitary_test, autodiff_test
+from .helper import (
+    relative_image_difference,
+    dotproduct_adjointness_test,
+    operator_isometry_test,
+    linear_operator_unitary_test,
+    autodiff_test,
+    gradient_of_linear_operator_test,
+    forward_mode_autodiff_of_linear_operator_test
+)

tests/helper.py

@@ -170,3 +170,84 @@ def autodiff_test(
     with torch.autograd.detect_anomaly():
         (_, vjpfunc) = torch.func.vjp(operator.forward, *u)
         vjpfunc(v_range)
+
+
+def gradient_of_linear_operator_test(
+    operator: LinearOperator,
+    u: torch.Tensor,
+    v: torch.Tensor,
+    relative_tolerance: float = 1e-3,
+    absolute_tolerance=1e-5,
+):
+    """Test the gradient of a linear operator is the adjoint.
+    Note: This property should hold for all u and v.
+    Commonly, this function is called with two random vectors u and v.
+    Parameters
+    ----------
+    operator
+        linear operator
+    u
+        element of the domain of the operator
+    v
+        element of the range of the operator
+    relative_tolerance
+        default is pytorch's default for float16
+    absolute_tolerance
+        default is pytorch's default for float16
+    Raises
+    ------
+    AssertionError
+        if the gradient is not the adjoint
+    """
+    # Gradient of the forward via vjp
+    (_, vjpfunc) = torch.func.vjp(operator.forward, u)
+    assert torch.allclose(vjpfunc((v,))[0], operator.adjoint(v)[0], rtol=relative_tolerance, atol=absolute_tolerance)
+
+    # Gradient of the adjoint via vjp
+    (_, vjpfunc) = torch.func.vjp(operator.adjoint, v)
+    assert torch.allclose(vjpfunc((u,))[0], operator.forward(u)[0], rtol=relative_tolerance, atol=absolute_tolerance)
+
+
+def forward_mode_autodiff_of_linear_operator_test(
+    operator: LinearOperator,
+    u: torch.Tensor,
+    v: torch.Tensor,
+    relative_tolerance: float = 1e-3,
+    absolute_tolerance=1e-5,
+):
+    """Test the forward-mode autodiff calculation.
+    Verifies that the Jacobian-vector product (jvp) is equivalent to applying the operator.
+    Note: This property should hold for all u and v.
+    Commonly, this function is called with two random vectors u and v.
+    Parameters
+    ----------
+    operator
+        linear operator
+    u
+        element of the domain of the operator
+    v
+        element of the range of the operator
+    relative_tolerance
+        default is pytorch's default for float16
+    absolute_tolerance
+        default is pytorch's default for float16
+    Raises
+    ------
+    AssertionError
+        if the jvp yields different results than applying the operator
+    """
+    # jvp of the forward
+    assert torch.allclose(
+        torch.func.jvp(operator.forward, (u,), (u,))[0][0],
+        operator.forward(u)[0],
+        rtol=relative_tolerance,
+        atol=absolute_tolerance,
+    )
+
+    # jvp of the adjoint
+    assert torch.allclose(
+        torch.func.jvp(operator.adjoint, (v,), (v,))[0][0],
+        operator.adjoint(v)[0],
+        rtol=relative_tolerance,
+        atol=absolute_tolerance,
+    )

tests/operators/models/test_inversion_recovery.py

@@ -47,3 +47,26 @@ def test_autodiff_inversion_recovery():
     model = InversionRecovery(ti=10)
     m0, t1 = create_parameter_tensor_tuples(parameter_shape=(2, 5, 10, 10, 10), number_of_tensors=2)
     autodiff_test(model, m0, t1)
+
+
+@pytest.mark.cuda
+def test_inversion_recovery_cuda():
+    """Test the inversion recovery model works on cuda devices."""
+    m0, t1 = create_parameter_tensor_tuples(parameter_shape=(2, 5, 10, 10, 10), number_of_tensors=2)
+
+    # Create on CPU, transfer to GPU and run on GPU
+    model = InversionRecovery(ti=10)
+    model.cuda()
+    (signal,) = model(m0.cuda(), t1.cuda())
+    assert signal.is_cuda
+
+    # Create on GPU and run on GPU
+    model = InversionRecovery(ti=torch.as_tensor(10).cuda())
+    (signal,) = model(m0.cuda(), t1.cuda())
+    assert signal.is_cuda
+
+    # Create on GPU, transfer to CPU and run on CPU
+    model = InversionRecovery(ti=torch.as_tensor(10).cuda())
+    model.cpu()
+    (signal,) = model(m0, t1)
+    assert signal.is_cpu

tests/operators/models/test_molli.py

@@ -52,3 +52,26 @@ def test_autodiff_molli():
     model = MOLLI(ti=10)
     a, b, t1 = create_parameter_tensor_tuples((2, 5, 10, 10, 10), number_of_tensors=3)
     autodiff_test(model, a, b, t1)
+
+
+@pytest.mark.cuda
+def test_molli_cuda():
+    """Test the molli model works on cuda devices."""
+    a, b, t1 = create_parameter_tensor_tuples((2, 5, 10, 10, 10), number_of_tensors=3)
+
+    # Create on CPU, transfer to GPU and run on GPU
+    model = MOLLI(ti=10)
+    model.cuda()
+    (signal,) = model(a.cuda(), b.cuda(), t1.cuda())
+    assert signal.is_cuda
+
+    # Create on GPU and run on GPU
+    model = MOLLI(ti=torch.as_tensor(10).cuda())
+    (signal,) = model(a.cuda(), b.cuda(), t1.cuda())
+    assert signal.is_cuda
+
+    # Create on GPU, transfer to CPU and run on CPU
+    model = MOLLI(ti=torch.as_tensor(10).cuda())
+    model.cpu()
+    (signal,) = model(a, b, t1)
+    assert signal.is_cpu

tests/operators/models/test_mono_exponential_decay.py

@@ -44,8 +44,31 @@ def test_mono_exponential_decay_shape(parameter_shape, contrast_dim_shape, signa
     assert signal.shape == signal_shape
 
 
-def test_autodiff_exponential_decay():
+def test_autodiff_mono_exponential_decay():
     """Test autodiff works for mono-exponential decay model."""
     model = MonoExponentialDecay(decay_time=20)
     m0, decay_constant = create_parameter_tensor_tuples(parameter_shape=(2, 5, 10, 10, 10), number_of_tensors=2)
     autodiff_test(model, m0, decay_constant)
+
+
+@pytest.mark.cuda
+def test_mono_exponential_deay_cuda():
+    """Test the mono-exponential decay model works on cuda devices."""
+    m0, decay_constant = create_parameter_tensor_tuples(parameter_shape=(2, 5, 10, 10, 10), number_of_tensors=2)
+
+    # Create on CPU, transfer to GPU and run on GPU
+    model = MonoExponentialDecay(decay_time=10)
+    model.cuda()
+    (signal,) = model(m0.cuda(), decay_constant.cuda())
+    assert signal.is_cuda
+
+    # Create on GPU and run on GPU
+    model = MonoExponentialDecay(decay_time=torch.as_tensor(10).cuda())
+    (signal,) = model(m0.cuda(), decay_constant.cuda())
+    assert signal.is_cuda
+
+    # Create on GPU, transfer to CPU and run on CPU
+    model = MonoExponentialDecay(decay_time=torch.as_tensor(10).cuda())
+    model.cpu()
+    (signal,) = model(m0, decay_constant)
+    assert signal.is_cpu

tests/operators/models/test_saturation_recovery.py

@@ -49,3 +49,26 @@ def test_autodiff_aturation_recovery():
     model = SaturationRecovery(ti=10)
     m0, t1 = create_parameter_tensor_tuples((2, 5, 10, 10, 10), number_of_tensors=2)
     autodiff_test(model, m0, t1)
+
+
+@pytest.mark.cuda
+def test_saturation_recovery_cuda():
+    """Test the saturation recovery model works on cuda devices."""
+    m0, t1 = create_parameter_tensor_tuples(parameter_shape=(2, 5, 10, 10, 10), number_of_tensors=2)
+
+    # Create on CPU, transfer to GPU and run on GPU
+    model = SaturationRecovery(ti=10)
+    model.cuda()
+    (signal,) = model(m0.cuda(), t1.cuda())
+    assert signal.is_cuda
+
+    # Create on GPU and run on GPU
+    model = SaturationRecovery(ti=torch.as_tensor(10).cuda())
+    (signal,) = model(m0.cuda(), t1.cuda())
+    assert signal.is_cuda
+
+    # Create on GPU, transfer to CPU and run on CPU
+    model = SaturationRecovery(ti=torch.as_tensor(10).cuda())
+    model.cpu()
+    (signal,) = model(m0, t1)
+    assert signal.is_cpu

tests/operators/models/test_transient_steady_state_with_preparation.py

@@ -94,3 +94,36 @@ def test_autodiff_transient_steady_state():
     )
     m0, t1, flip_angle = create_parameter_tensor_tuples(parameter_shape=(2, 5, 10, 10, 10), number_of_tensors=3)
     autodiff_test(model, m0, t1, flip_angle)
+
+
+@pytest.mark.cuda
+def test_transient_steady_state_cuda():
+    """Test the transient steady state model works on cuda devices."""
+    m0, t1, flip_angle = create_parameter_tensor_tuples(parameter_shape=(2, 5, 10, 10, 10), number_of_tensors=3)
+    contrast_dim_shape = (6,)
+    (sampling_time,) = create_parameter_tensor_tuples(contrast_dim_shape, number_of_tensors=1)
+    repetition_time, m0_scaling_preparation, delay_after_preparation = create_parameter_tensor_tuples(
+        contrast_dim_shape[1:], number_of_tensors=3
+    )
+    # Create on CPU, transfer to GPU and run on GPU
+    model = TransientSteadyStateWithPreparation(
+        sampling_time, repetition_time, m0_scaling_preparation, delay_after_preparation
+    )
+    model.cuda()
+    (signal,) = model(m0.cuda(), t1.cuda(), flip_angle.cuda())
+    assert signal.is_cuda
+
+    # Create on GPU and run on GPU
+    model = TransientSteadyStateWithPreparation(
+        sampling_time.cuda(), repetition_time, m0_scaling_preparation, delay_after_preparation
+    )
+    (signal,) = model(m0.cuda(), t1.cuda(), flip_angle.cuda())
+    assert signal.is_cuda
+
+    # Create on GPU, transfer to CPU and run on CPU
+    model = TransientSteadyStateWithPreparation(
+        sampling_time.cuda(), repetition_time, m0_scaling_preparation, delay_after_preparation
+    )
+    model.cpu()
+    (signal,) = model(m0, t1, flip_angle)
+    assert signal.is_cpu

tests/operators/models/test_wasabi.py

@@ -1,5 +1,6 @@
 """Tests for the WASABI signal model."""
 
+import pytest
 import torch
 from mrpro.operators.models import WASABI
 from tests import autodiff_test
@@ -53,3 +54,26 @@ def test_autodiff_WASABI():
     offset, b0_shift, rb1, c, d = create_data(offset_max=300, n_offsets=2)
     wasabi_model = WASABI(offsets=offset)
     autodiff_test(wasabi_model, b0_shift, rb1, c, d)
+
+
+@pytest.mark.cuda
+def test_wasabi_cuda():
+    """Test the WASABI model works on cuda devices."""
+    offset, b0_shift, rb1, c, d = create_data(offset_max=300, n_offsets=2)
+
+    # Create on CPU, transfer to GPU and run on GPU
+    model = WASABI(offset)
+    model.cuda()
+    (signal,) = model(b0_shift.cuda(), rb1.cuda(), c.cuda(), d.cuda())
+    assert signal.is_cuda
+
+    # Create on GPU and run on GPU
+    model = WASABI(offset.cuda())
+    (signal,) = model(b0_shift.cuda(), rb1.cuda(), c.cuda(), d.cuda())
+    assert signal.is_cuda
+
+    # Create on GPU, transfer to CPU and run on CPU
+    model = WASABI(offset.cuda())
+    model.cpu()
+    (signal,) = model(b0_shift, rb1, c, d)
+    assert signal.is_cpu

tests/operators/models/test_wasabiti.py

@@ -85,3 +85,27 @@ def test_autodiff_WASABITI():
     trec = torch.ones_like(offset) * t1
     wasabiti_model = WASABITI(offsets=offset, trec=trec)
     autodiff_test(wasabiti_model, b0_shift, rb1, t1)
+
+
+@pytest.mark.cuda
+def test_wasabiti_cuda():
+    """Test the WASABITI model works on cuda devices."""
+    offset, b0_shift, rb1, t1 = create_data(offset_max=300, n_offsets=2)
+    trec = torch.ones_like(offset) * t1
+
+    # Create on CPU, transfer to GPU and run on GPU
+    model = WASABITI(offset, trec)
+    model.cuda()
+    (signal,) = model(b0_shift.cuda(), rb1.cuda(), t1.cuda())
+    assert signal.is_cuda
+
+    # Create on GPU and run on GPU
+    model = WASABITI(offset.cuda(), trec)
+    (signal,) = model(b0_shift.cuda(), rb1.cuda(), t1.cuda())
+    assert signal.is_cuda
+
+    # Create on GPU, transfer to CPU and run on CPU
+    model = WASABITI(offset.cuda(), trec)
+    model.cpu()
+    (signal,) = model(b0_shift, rb1, t1)
+    assert signal.is_cpu