Add tests for autodiff of operators (#405)

Co-authored-by: Felix F Zimmermann <[email protected]>

tests/__init__.py

@@ -1 +1,2 @@
 from ._RandomGenerator import RandomGenerator
+from .helper import relative_image_difference, dotproduct_adjointness_test, operator_isometry_test, linear_operator_unitary_test, autodiff_test

tests/algorithms/csm/test_inati.py

@@ -3,8 +3,8 @@
 import torch
 from mrpro.algorithms.csm import inati
 from mrpro.data import SpatialDimension
+from tests import relative_image_difference
 from tests.algorithms.csm.conftest import multi_coil_image
-from tests.helper import relative_image_difference
 
 
 def test_inati(ellipse_phantom, random_kheader):

tests/algorithms/csm/test_walsh.py

@@ -3,8 +3,8 @@
 import torch
 from mrpro.algorithms.csm import walsh
 from mrpro.data import SpatialDimension
+from tests import relative_image_difference
 from tests.algorithms.csm.conftest import multi_coil_image
-from tests.helper import relative_image_difference
 
 
 def test_walsh(ellipse_phantom, random_kheader):

tests/data/test_csm_data.py

@@ -6,8 +6,8 @@
 import torch
 from mrpro.data import CsmData, SpatialDimension
 
+from tests import relative_image_difference
 from tests.algorithms.csm.test_walsh import multi_coil_image
-from tests.helper import relative_image_difference
 
 
 def test_CsmData_is_frozen_dataclass(random_test_data, random_kheader):

tests/data/test_kdata.py

@@ -10,9 +10,9 @@
 from mrpro.operators import FastFourierOp
 from mrpro.utils import split_idx
 
+from tests import relative_image_difference
 from tests.conftest import RandomGenerator, generate_random_data
 from tests.data import IsmrmrdRawTestData
-from tests.helper import relative_image_difference
 from tests.phantoms import EllipsePhantomTestData
 
 

tests/helper.py

@@ -1,7 +1,8 @@
 """Helper/Utilities for test functions."""
 
 import torch
-from mrpro.operators import Operator
+from mrpro.operators import LinearOperator, Operator
+from typing_extensions import TypeVarTuple, Unpack
 
 
 def relative_image_difference(img1: torch.Tensor, img2: torch.Tensor) -> torch.Tensor:
@@ -26,9 +27,13 @@ def relative_image_difference(img1: torch.Tensor, img2: torch.Tensor) -> torch.T
 
 
 def dotproduct_adjointness_test(
-    operator: Operator, u: torch.Tensor, v: torch.Tensor, relative_tolerance: float = 1e-3, absolute_tolerance=1e-5
+    operator: LinearOperator,
+    u: torch.Tensor,
+    v: torch.Tensor,
+    relative_tolerance: float = 1e-3,
+    absolute_tolerance=1e-5,
 ):
-    """Test the adjointness of operator and operator.H.
+    """Test the adjointness of linear operator and operator.H.
 
     Test if
          <Operator(u),v> == <u, Operator^H(v)>
@@ -42,7 +47,7 @@ def dotproduct_adjointness_test(
     Parameters
     ----------
     operator
-        operator
+        linear operator
     u
         element of the domain of the operator
     v
@@ -74,9 +79,12 @@ def dotproduct_adjointness_test(
 
 
 def operator_isometry_test(
-    operator: Operator, u: torch.Tensor, relative_tolerance: float = 1e-3, absolute_tolerance=1e-5
+    operator: Operator[torch.Tensor, tuple[torch.Tensor]],
+    u: torch.Tensor,
+    relative_tolerance: float = 1e-3,
+    absolute_tolerance=1e-5,
 ):
-    """Test the isometry of an operator.
+    """Test the isometry of a operator.
 
     Test if
          ||Operator(u)|| == ||u||
@@ -103,10 +111,10 @@ def operator_isometry_test(
     )
 
 
-def operator_unitary_test(
-    operator: Operator, u: torch.Tensor, relative_tolerance: float = 1e-3, absolute_tolerance=1e-5
+def linear_operator_unitary_test(
+    operator: LinearOperator, u: torch.Tensor, relative_tolerance: float = 1e-3, absolute_tolerance=1e-5
 ):
-    """Test if an operator is unitary.
+    """Test if a linear operator is unitary.
 
     Test if
          Operator.adjoint(Operator(u)) == u
@@ -115,7 +123,7 @@ def operator_unitary_test(
     Parameters
     ----------
     operator
-        operator
+        linear operator
     u
         element of the domain of the operator
     relative_tolerance
@@ -129,3 +137,36 @@ def operator_unitary_test(
         if the adjointness property does not hold
     """
     torch.testing.assert_close(u, operator.adjoint(operator(u)[0])[0], rtol=relative_tolerance, atol=absolute_tolerance)
+
+
+Tin = TypeVarTuple('Tin')
+
+
+def autodiff_test(
+    operator: Operator[Unpack[Tin], tuple[torch.Tensor, ...]],
+    *u: Unpack[Tin],
+):
+    """Test if autodiff of an operator is working.
+    This test does not check that the gradient is correct but simply that it can be calculated using both torch.func.jvp
+    and torch.func.vjp.
+
+    Parameters
+    ----------
+    operator
+        operator
+    u
+        element(s) of the domain of the operator
+
+    Raises
+    ------
+    AssertionError
+        if autodiff fails
+    """
+    # Forward-mode autodiff using jvp
+    with torch.autograd.detect_anomaly():
+        v_range, _ = torch.func.jvp(operator.forward, u, u)
+
+    # Backward-mode autodiff using vjp
+    with torch.autograd.detect_anomaly():
+        (_, vjpfunc) = torch.func.vjp(operator.forward, *u)
+        vjpfunc(v_range)

tests/operators/models/test_inversion_recovery.py

@@ -3,6 +3,7 @@
 import pytest
 import torch
 from mrpro.operators.models import InversionRecovery
+from tests import autodiff_test
 from tests.operators.models.conftest import SHAPE_VARIATIONS_SIGNAL_MODELS, create_parameter_tensor_tuples
 
 
@@ -39,3 +40,10 @@ def test_inversion_recovery_shape(parameter_shape, contrast_dim_shape, signal_sh
     m0, t1 = create_parameter_tensor_tuples(parameter_shape, number_of_tensors=2)
     (signal,) = model_op(m0, t1)
     assert signal.shape == signal_shape
+
+
+def test_autodiff_inversion_recovery():
+    """Test autodiff works for inversion_recovery model."""
+    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)

tests/operators/models/test_molli.py

@@ -3,7 +3,7 @@
 import pytest
 import torch
 from mrpro.operators.models import MOLLI
-from tests import RandomGenerator
+from tests import RandomGenerator, autodiff_test
 from tests.operators.models.conftest import SHAPE_VARIATIONS_SIGNAL_MODELS, create_parameter_tensor_tuples
 
 
@@ -45,3 +45,10 @@ def test_molli_shape(parameter_shape, contrast_dim_shape, signal_shape):
     a, c, t1 = create_parameter_tensor_tuples(parameter_shape, number_of_tensors=3)
     (signal,) = model_op(a, c, t1)
     assert signal.shape == signal_shape
+
+
+def test_autodiff_molli():
+    """Test autodiff works for molli model."""
+    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)

tests/operators/models/test_mono_exponential_decay.py

@@ -3,6 +3,7 @@
 import pytest
 import torch
 from mrpro.operators.models import MonoExponentialDecay
+from tests import autodiff_test
 from tests.operators.models.conftest import SHAPE_VARIATIONS_SIGNAL_MODELS, create_parameter_tensor_tuples
 
 
@@ -41,3 +42,10 @@ def test_mono_exponential_decay_shape(parameter_shape, contrast_dim_shape, signa
     m0, decay_constant = create_parameter_tensor_tuples(parameter_shape, number_of_tensors=2)
     (signal,) = model_op(m0, decay_constant)
     assert signal.shape == signal_shape
+
+
+def test_autodiff_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)

tests/operators/models/test_saturation_recovery.py

@@ -3,6 +3,7 @@
 import pytest
 import torch
 from mrpro.operators.models import SaturationRecovery
+from tests import autodiff_test
 from tests.operators.models.conftest import SHAPE_VARIATIONS_SIGNAL_MODELS, create_parameter_tensor_tuples
 
 
@@ -41,3 +42,10 @@ def test_saturation_recovery_shape(parameter_shape, contrast_dim_shape, signal_s
     m0, t1 = create_parameter_tensor_tuples(parameter_shape, number_of_tensors=2)
     (signal,) = model_op(m0, t1)
     assert signal.shape == signal_shape
+
+
+def test_autodiff_aturation_recovery():
+    """Test autodiff works for aturation recovery model."""
+    model = SaturationRecovery(ti=10)
+    m0, t1 = create_parameter_tensor_tuples((2, 5, 10, 10, 10), number_of_tensors=2)
+    autodiff_test(model, m0, t1)

tests/operators/models/test_transient_steady_state_with_preparation.py

@@ -4,6 +4,7 @@
 import torch
 from einops import repeat
 from mrpro.operators.models import TransientSteadyStateWithPreparation
+from tests import autodiff_test
 from tests.operators.models.conftest import SHAPE_VARIATIONS_SIGNAL_MODELS, create_parameter_tensor_tuples
 
 
@@ -79,3 +80,17 @@ def test_transient_steady_state_shape(parameter_shape, contrast_dim_shape, signa
     m0, t1, flip_angle = create_parameter_tensor_tuples(parameter_shape, number_of_tensors=3)
     (signal,) = model_op(m0, t1, flip_angle)
     assert signal.shape == signal_shape
+
+
+def test_autodiff_transient_steady_state():
+    """Test autodiff works for transient steady state model."""
+    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
+    )
+    model = TransientSteadyStateWithPreparation(
+        sampling_time, repetition_time, m0_scaling_preparation, delay_after_preparation
+    )
+    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)

tests/operators/models/test_wasabi.py

@@ -2,6 +2,7 @@
 
 import torch
 from mrpro.operators.models import WASABI
+from tests import autodiff_test
 from tests.operators.models.conftest import SHAPE_VARIATIONS_SIGNAL_MODELS, create_parameter_tensor_tuples
 
 
@@ -45,3 +46,10 @@ def test_WASABI_shape(parameter_shape, contrast_dim_shape, signal_shape):
     b0_shift, rb1, c, d = create_parameter_tensor_tuples(parameter_shape, number_of_tensors=4)
     (signal,) = model_op(b0_shift, rb1, c, d)
     assert signal.shape == signal_shape
+
+
+def test_autodiff_WASABI():
+    """Test autodiff works for WASABI model."""
+    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)

tests/operators/models/test_wasabiti.py

@@ -3,6 +3,7 @@
 import pytest
 import torch
 from mrpro.operators.models import WASABITI
+from tests import autodiff_test
 from tests.operators.models.conftest import SHAPE_VARIATIONS_SIGNAL_MODELS, create_parameter_tensor_tuples
 
 
@@ -76,3 +77,11 @@ def test_WASABITI_shape(parameter_shape, contrast_dim_shape, signal_shape):
     b0_shift, rb1, t1 = create_parameter_tensor_tuples(parameter_shape, number_of_tensors=3)
     (signal,) = model_op(b0_shift, rb1, t1)
     assert signal.shape == signal_shape
+
+
+def test_autodiff_WASABITI():
+    """Test autodiff works for WASABITI model."""
+    offset, b0_shift, rb1, t1 = create_data(offset_max=300, n_offsets=2)
+    trec = torch.ones_like(offset) * t1
+    wasabiti_model = WASABITI(offsets=offset, trec=trec)
+    autodiff_test(wasabiti_model, b0_shift, rb1, t1)

tests/operators/test_autograd_linop.py

@@ -5,8 +5,7 @@
 from mrpro.operators import LinearOperator
 from torch.autograd.gradcheck import GradcheckError
 
-from tests import RandomGenerator
-from tests.helper import dotproduct_adjointness_test
+from tests import RandomGenerator, dotproduct_adjointness_test
 
 
 class NonDifferentiableOperator(LinearOperator, adjoint_as_backward=False):

tests/operators/test_cartesian_sampling_op.py

@@ -6,9 +6,8 @@
 from mrpro.data import KTrajectory, SpatialDimension
 from mrpro.operators import CartesianSamplingOp
 
-from tests import RandomGenerator
+from tests import RandomGenerator, dotproduct_adjointness_test
 from tests.conftest import create_traj
-from tests.helper import dotproduct_adjointness_test
 
 
 def test_cart_sampling_op_data_match():

tests/operators/test_constraints_op.py

@@ -4,7 +4,7 @@
 import torch
 from mrpro.operators import ConstraintsOp
 
-from tests import RandomGenerator
+from tests import RandomGenerator, autodiff_test
 
 
 @pytest.mark.parametrize(
@@ -141,3 +141,15 @@ def test_constraints_operator_multiple_inputs(bounds):
 def test_constraints_operator_illegal_bounds(bounds):
     with pytest.raises(ValueError, match='invalid'):
         ConstraintsOp(bounds)
+
+
+def test_autodiff_constraints_operator():
+    """Test autodiff works for constraints operator."""
+    # random tensors with arbitrary values
+    random_generator = RandomGenerator(seed=0)
+    x1 = random_generator.float32_tensor(size=(36, 72), low=-1, high=1)
+    x2 = random_generator.float32_tensor(size=(36, 72), low=-1, high=1)
+    x3 = random_generator.float32_tensor(size=(36, 72), low=-1, high=1)
+
+    constraints_op = ConstraintsOp(bounds=((None, None), (1.0, None), (None, 1.0)))
+    autodiff_test(constraints_op, x1, x2, x3)

tests/operators/test_density_compensation_op.py

@@ -4,8 +4,7 @@
 from mrpro.data import DcfData
 from mrpro.operators import DensityCompensationOp
 
-from tests import RandomGenerator
-from tests.helper import dotproduct_adjointness_test
+from tests import RandomGenerator, dotproduct_adjointness_test
 
 
 def test_density_compensation_op_adjointness():

tests/operators/test_einsum_op.py

@@ -4,8 +4,7 @@
 import torch
 from mrpro.operators.EinsumOp import EinsumOp
 
-from tests import RandomGenerator
-from tests.helper import dotproduct_adjointness_test
+from tests import RandomGenerator, dotproduct_adjointness_test
 
 
 @pytest.mark.parametrize('dtype', ['float32', 'complex128'])

tests/operators/test_fast_fourier_op.py

@@ -6,8 +6,7 @@
 from mrpro.data import SpatialDimension
 from mrpro.operators import FastFourierOp
 
-from tests import RandomGenerator
-from tests.helper import dotproduct_adjointness_test
+from tests import RandomGenerator, dotproduct_adjointness_test
 
 
 @pytest.mark.parametrize(('npoints', 'a'), [(100, 20), (300, 20)])

tests/operators/test_finite_difference_op.py

@@ -5,8 +5,7 @@
 from einops import repeat
 from mrpro.operators import FiniteDifferenceOp
 
-from tests import RandomGenerator
-from tests.helper import dotproduct_adjointness_test
+from tests import RandomGenerator, dotproduct_adjointness_test
 
 
 @pytest.mark.parametrize('mode', ['central', 'forward', 'backward'])

tests/operators/test_fourier_op.py

@@ -6,9 +6,8 @@
 from mrpro.data.traj_calculators import KTrajectoryCartesian
 from mrpro.operators import FourierOp
 
-from tests import RandomGenerator
+from tests import RandomGenerator, dotproduct_adjointness_test
 from tests.conftest import COMMON_MR_TRAJECTORIES, create_traj
-from tests.helper import dotproduct_adjointness_test
 
 
 def create_data(im_shape, k_shape, nkx, nky, nkz, type_kx, type_ky, type_kz):