From 9199134459d1afe85c673dea90a55d2ef09cb3ba Mon Sep 17 00:00:00 2001 From: Felix Zimmermann Date: Wed, 6 May 2026 14:02:52 +0200 Subject: [PATCH 1/3] Add subspace NUFFT Toeplitz Gram support --- .../operators/NonUniformFastFourierOp.py | 243 +++++++++++++++++- src/mrpro/operators/__init__.py | 6 +- .../test_non_uniform_fast_fourier_op.py | 143 ++++++++++- 3 files changed, 379 insertions(+), 13 deletions(-) diff --git a/src/mrpro/operators/NonUniformFastFourierOp.py b/src/mrpro/operators/NonUniformFastFourierOp.py index 44a40a8a8..e72c9488d 100644 --- a/src/mrpro/operators/NonUniformFastFourierOp.py +++ b/src/mrpro/operators/NonUniformFastFourierOp.py @@ -6,6 +6,7 @@ from types import EllipsisType from typing import Literal +import einops import numpy as np import torch from pytorch_finufft.functional import finufft_type1, finufft_type2 @@ -15,6 +16,8 @@ from mrpro.data.KTrajectory import KTrajectory from mrpro.data.SpatialDimension import SpatialDimension from mrpro.operators.LinearOperator import LinearOperator +from mrpro.operators.PCACompressionOp import PCACompressionOp +from mrpro.utils import normalize_index, unsqueeze_right class NonUniformFastFourierOp(LinearOperator, adjoint_as_backward=True): @@ -200,12 +203,14 @@ def __call__(self, x: torch.Tensor) -> tuple[torch.Tensor,]: Parameters ---------- x - Coil image data, typically with shape `(..., coils, z, y, x)`. + Image-space data with selected spatial dimensions as trailing axes, e.g. + `(..., coils, z, y, x)` for 3D or `(..., coils, y, x)` for 2D. Returns ------- - Coil k-space data at non-uniform locations, - with shape `(..., coils, k2, k1, k0)`. + K-space data at the non-uniform trajectory locations with the same leading + dimensions as the input and trailing sampled dimensions, e.g. + `(..., coils, k2, k1, k0)` for 3D or `(..., coils, k1, k0)` for 2D. """ return super().__call__(x) @@ -282,10 +287,35 @@ def adjoint(self, x: torch.Tensor) -> tuple[torch.Tensor,]: x = x.permute(*unpermute_zyx) return (x,) + def toeplitz( + self, + weight: torch.Tensor | DcfData | None = None, + subspace: torch.Tensor | PCACompressionOp | None = None, + subspace_dim: int = 0, + ) -> LinearOperator: + """Return the Toeplitz Gram operator. + + Parameters + ---------- + weight + Optional density compensation or k-space weights. If provided, calculates + the normal operator corresponding to :math:`F^H W F`. + subspace + Optional temporal subspace basis of shape `(n_timepoints, n_coefficients)`, or + a `PCACompressionOp` whose adjoint defines the temporal expansion basis. + subspace_dim + Dimension of the coefficient channel in the input/output tensors when `subspace` is provided. + """ + if subspace is None: + return NonUniformFastFourierOpGramOp(self, weight=weight) + return SubspaceNonUniformFastFourierOpGramOp( + self, subspace_basis=subspace, weight=weight, subspace_dim=subspace_dim + ) + @property def gram(self) -> LinearOperator: - """Return the gram operator.""" - return NonUniformFastFourierOpGramOp(self) + """Return the unweighted Toeplitz Gram operator.""" + return self.toeplitz() def __repr__(self) -> str: """Representation method for NUFFT operator.""" @@ -404,17 +434,14 @@ def __init__(self, nufft_op: NonUniformFastFourierOp, weight: torch.Tensor | Dcf if isinstance(weight, DcfData): weight = weight.data - weight_shape = [*nufft_op._traj_broadcast_shape[:-4], 1, *nufft_op._traj_broadcast_shape[-3:]] - if weight is None: - weight = torch.ones(weight_shape, device=nufft_op._omega.device) + weight = torch.ones(weight_shape, device=nufft_op._omega.device, dtype=nufft_op._omega.dtype.to_real()) else: weight = weight.to(nufft_op._omega.device).broadcast_to(weight_shape) # We rearrange weight into (sep_dims, joint_dims, nufft_directions) _, permute_zyx, sep_dims_210, permute_210 = nufft_op._separate_joint_dimensions(weight.ndim) unpermute_zyx = torch.tensor(permute_zyx).argsort().tolist() - weight = weight.permute(*permute_210) unflatten_other_shape = weight.shape[: -len(nufft_op._dimension_210) - 1] # -1 for coil if sep_dims_210: # combine sep_dims @@ -428,6 +455,14 @@ def __init__(self, nufft_op: NonUniformFastFourierOp, weight: torch.Tensor | Dcf kernel = kernel.reshape(*unflatten_other_shape, -1, *kernel.shape[-len(nufft_op._direction_zyx) :]) self._kernel = kernel.permute(*unpermute_zyx) * nufft_op.scale.item() ** 2 + @property + def recon_matrix(self) -> SpatialDimension[int]: + """Expected reconstructed image shape as a spatial dimension.""" + zyx = [1, 1, 1] + for direction, size in zip(self._dim, self._recon_shape, strict=True): + zyx[direction] = size + return SpatialDimension(*zyx) + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor,]: """Apply forward of NonUniformFastFourierOpGramOp. @@ -472,7 +507,7 @@ def __call__(self, x: torch.Tensor) -> tuple[torch.Tensor,]: Returns ------- - Output tensor, image-space data after NUFFT.H @ NUFFT has been applied. + Image-space data after applying `NUFFT.H @ NUFFT`, with the same shape as `x`. """ return super().__call__(x) @@ -489,7 +524,193 @@ def adjoint(self, x: torch.Tensor) -> tuple[torch.Tensor,]: Returns ------- - Output tensor, same shape as the input. + Image-space data after applying the adjoint Gram operator, with the same shape as `x`. + """ + return super().__call__(x) + + @property + def H(self) -> Self: # noqa: N802 + """Adjoint operator of the gram operator.""" + return self + + +class SubspaceNonUniformFastFourierOpGramOp(LinearOperator, adjoint_as_backward=True): + """Subspace Gram operator for `NonUniformFastFourierOp`. + + This operator acts on temporally compressed coefficient images and applies the + compressed normal operator corresponding to a time-varying non-uniform Fourier + encoding. The Toeplitz kernel is precomputed once during initialization. + + The current implementation is intentionally restricted to the common MRF case: + exactly one varying separate trajectory dimension, interpreted as time, and no + additional non-singleton joint trajectory dimensions. + """ + + _kernel: torch.Tensor | None + + def __init__( + self, + nufft_op: NonUniformFastFourierOp, + subspace_basis: torch.Tensor | PCACompressionOp, + weight: torch.Tensor | DcfData | None = None, + subspace_dim: int = 0, + ) -> None: + """Initialize the subspace Gram operator. + + Parameters + ---------- + nufft_op + The non-uniform Fourier operator defining the trajectory and spatial geometry. + subspace_basis + Temporal subspace basis of shape `(n_timepoints, n_coefficients)`, or + a `PCACompressionOp` whose adjoint defines the temporal expansion basis. + weight + Optional density compensation or k-space weights. If provided, calculates + the compressed normal operator corresponding to :math:`F^H W F`. + subspace_dim + Dimension of the coefficient channel in the input/output tensors. Spatial + dimensions are assumed to remain trailing, matching the usual MR2 image layout. + """ + super().__init__() + self._dim = nufft_op._direction_zyx + self._recon_shape = nufft_op._im_size + self._subspace_dim = subspace_dim + + if not nufft_op._dimension_210: + self._kernel = None + return + + if isinstance(subspace_basis, PCACompressionOp): + basis = subspace_basis._compression_matrix.mH + else: + basis = subspace_basis + basis = basis.squeeze() + if basis.ndim > 2: + raise ValueError(f'Basis cannot contain non-singleton batch dimensions; got squeezed shape {basis.shape}.') + elif basis.ndim == 1: # rank-1 special case, we squeezed the singleton subspace dimension + basis = basis.unsqueeze(-1) + basis = basis.to(device=nufft_op._omega.device) + + if isinstance(weight, DcfData): + weight = weight.data + weight_shape = [*nufft_op._traj_broadcast_shape[:-4], 1, *nufft_op._traj_broadcast_shape[-3:]] + if weight is None: + weight = torch.ones(weight_shape, device=nufft_op._omega.device, dtype=nufft_op._omega.dtype.to_real()) + else: + weight = weight.to(nufft_op._omega.device).broadcast_to(weight_shape) + _, _permute_zyx, sep_dims_210, permute_210 = nufft_op._separate_joint_dimensions(weight.ndim) + weight = weight.permute(*permute_210) + if sep_dims_210: + weight = weight.flatten(end_dim=len(sep_dims_210) - 1) + else: + weight = weight[None, :] + weight = weight.flatten(start_dim=1, end_dim=-len(nufft_op._dimension_210) - 1).flatten(start_dim=2) + + omega = nufft_op._omega + + if len(sep_dims_210) > 1: + raise NotImplementedError( + 'SubspaceNonUniformFastFourierOpGramOp currently only supports a single varying separate dimension.' + ) + if basis.shape[0] != (n_timepoints := weight.shape[0]): + raise ValueError( + f'subspace_basis has {basis.shape[0]} time points, but the trajectory varies over {n_timepoints}.' + ) + + subspace_kernel: torch.Tensor | None = None + for basis_row, weight_row, omega_row in zip(basis, weight, omega, strict=True): + current = gram_nufft_kernel(weight_row[None], omega_row[None], nufft_op._im_size)[0, 0] + coeff_outer = basis_row.conj()[:, None] * basis_row[None, :] + term = unsqueeze_right(coeff_outer, current.ndim) * current + subspace_kernel = term if subspace_kernel is None else subspace_kernel + term + + assert subspace_kernel is not None # noqa: S101 + self._kernel = subspace_kernel * nufft_op.scale.item() ** 2 + + @property + def n_coefficients(self) -> int: + """Number of compressed coefficient channels.""" + if self._kernel is None: + raise RuntimeError('n_coefficients is undefined before the Toeplitz kernel is initialized.') + return self._kernel.shape[0] + + @property + def recon_matrix(self) -> SpatialDimension[int]: + """Expected reconstructed image shape as a spatial dimension.""" + zyx = [1, 1, 1] + for direction, size in zip(self._dim, self._recon_shape, strict=True): + zyx[direction] = size + return SpatialDimension(*zyx) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor,]: + """Apply forward of SubspaceNonUniformFastFourierOpGramOp. + + .. note:: + Prefer calling the instance of the SubspaceNonUniformFastFourierOpGramOp operator as ``operator(x)`` over + directly calling this method. See this PyTorch `discussion `_. + """ + # We do the fft, coefficient mixing, and cropping here directly, as it is faster and more memory efficient + # than using separate operators. + + # This function on its own is also not autograd save (in-place ops), but we rely on the + # adjoint-as-backward trick for linear operators to make it work. + if self._kernel is None: + return (x,) + subspace_dim = normalize_index(x.ndim, self._subspace_dim) + if x.shape[subspace_dim] != self.n_coefficients: + raise ValueError( + f'Input has {x.shape[self._subspace_dim]} coefficients along subspace_dim={self._subspace_dim}, ' + f'expected {self.n_coefficients}.' + ) + x = x.movedim(subspace_dim, 0) + + spatial_dims = tuple(range(x.ndim - len(self._dim), x.ndim)) + padded_shape = [2 * s for s in self._recon_shape] + spatial_crop: list[slice] = [slice(None)] * x.ndim + for d, s in zip(spatial_dims, self._recon_shape, strict=True): + spatial_crop[d] = slice(0, s) + + x = torch.fft.fftn(x, s=padded_shape, dim=spatial_dims) + x = einops.einsum(self._kernel.to(x.dtype), x, 'coeff_out coeff_in ..., coeff_in ... -> coeff_out ...') + x = torch.fft.ifftn(x, dim=spatial_dims) + x = x[tuple(spatial_crop)] + out = x.clone().movedim(0, subspace_dim) + return (out,) + + def __call__(self, x: torch.Tensor) -> tuple[torch.Tensor,]: + """Apply the compressed Toeplitz Gram operator. + + Parameters + ---------- + x + Subspace coefficient images. The coefficient axis is given by `subspace_dim`; + the selected spatial dimensions are trailing axes, e.g. `(coeff, z, y, x)` + for 3D or `(coeff, y, x)` for 2D when `subspace_dim=0`. + + Returns + ------- + Subspace coefficient images after applying the compressed Gram operator, + with the same shape as `x`. + """ + return super().__call__(x) + + def adjoint(self, x: torch.Tensor) -> tuple[torch.Tensor,]: + """Apply the adjoint of the compressed Toeplitz Gram operator. + + Since the compressed Gram operator is self-adjoint, this method calls the + forward operation. + + Parameters + ---------- + x + Subspace coefficient images. The coefficient axis is given by `subspace_dim`; + the selected spatial dimensions are trailing axes, e.g. `(coeff, z, y, x)` + for 3D or `(coeff, y, x)` for 2D when `subspace_dim=0`. + + Returns + ------- + Subspace coefficient images after applying the adjoint compressed Gram operator, + with the same shape as `x`. """ return super().__call__(x) diff --git a/src/mrpro/operators/__init__.py b/src/mrpro/operators/__init__.py index b8e64333b..489dc58fc 100644 --- a/src/mrpro/operators/__init__.py +++ b/src/mrpro/operators/__init__.py @@ -33,7 +33,10 @@ from mrpro.operators.LinearOperatorMatrix import LinearOperatorMatrix from mrpro.operators.MagnitudeOp import MagnitudeOp from mrpro.operators.MultiIdentityOp import MultiIdentityOp -from mrpro.operators.NonUniformFastFourierOp import NonUniformFastFourierOp +from mrpro.operators.NonUniformFastFourierOp import ( + NonUniformFastFourierOp, + SubspaceNonUniformFastFourierOpGramOp, +) from mrpro.operators.OptimizerOp import OptimizerOp from mrpro.operators.PatchOp import PatchOp from mrpro.operators.PCACompressionOp import PCACompressionOp @@ -91,6 +94,7 @@ "SensitivityOp", "SignalModel", "SliceProjectionOp", + "SubspaceNonUniformFastFourierOpGramOp", "TimeSegmentedFourierOp", "WaveletOp", "ZeroOp", diff --git a/tests/operators/test_non_uniform_fast_fourier_op.py b/tests/operators/test_non_uniform_fast_fourier_op.py index 5a29bb6f0..8848554f7 100644 --- a/tests/operators/test_non_uniform_fast_fourier_op.py +++ b/tests/operators/test_non_uniform_fast_fourier_op.py @@ -1,12 +1,19 @@ """Tests for Non-Uniform Fast Fourier operator.""" +import einops import pytest import torch from mrpro.data import KData, KTrajectory from mrpro.data.SpatialDimension import SpatialDimension from mrpro.data.traj_calculators import KTrajectoryIsmrmrd -from mrpro.operators import FastFourierOp, NonUniformFastFourierOp +from mrpro.operators import ( + FastFourierOp, + NonUniformFastFourierOp, + PCACompressionOp, + SubspaceNonUniformFastFourierOpGramOp, +) from mrpro.utils import RandomGenerator +from torch.autograd.gradcheck import gradcheck from tests.conftest import COMMON_MR_TRAJECTORIES, create_traj from tests.helper import dotproduct_adjointness_test, relative_image_difference @@ -20,6 +27,29 @@ def create_data(img_shape, nkx, nky, nkz, type_kx, type_ky, type_kz) -> tuple[to return img, trajectory +def create_time_varying_2d_nufft_op( + n_timepoints: int = 4, + image_shape: tuple[int, int] = (12, 10), + dtype: torch.dtype = torch.float32, + device: torch.device | str = 'cpu', +) -> NonUniformFastFourierOp: + """Create a small 2D NUFFT with one radial spoke per time point.""" + angles = torch.linspace(0, torch.pi, n_timepoints + 1, dtype=dtype, device=device)[:-1] + readout = torch.linspace(-image_shape[-1] / 2, image_shape[-1] / 2, image_shape[-1], dtype=dtype, device=device) + kx = (torch.cos(angles)[:, None] * readout[None, :])[:, None, None, :] + ky = (torch.sin(angles)[:, None] * readout[None, :])[:, None, None, :] + kx = kx[:, None, :, :] + ky = ky[:, None, :, :] + kz = torch.zeros_like(kx) + trajectory = KTrajectory(kz=kz, ky=ky, kx=kx) + return NonUniformFastFourierOp( + direction=(-2, -1), + recon_matrix=SpatialDimension(z=1, y=image_shape[0], x=image_shape[1]), + encoding_matrix=SpatialDimension(z=1, y=image_shape[0], x=image_shape[1]), + traj=trajectory, + ) + + @COMMON_MR_TRAJECTORIES def test_non_uniform_fast_fourier_op_fwd_adj_property( img_shape, k_shape, nkx, nky, nkz, type_kx, type_ky, type_kz, type_k0, type_k1, type_k2 @@ -74,8 +104,10 @@ def test_non_uniform_fast_fourier_op_gram( (expected,) = (nufft_op.H @ nufft_op)(img) (actual,) = nufft_op.gram(img) + (actual_toeplitz,) = nufft_op.toeplitz()(img) torch.testing.assert_close(actual, expected, rtol=1e-3, atol=1e-3) + torch.testing.assert_close(actual_toeplitz, expected, rtol=1e-3, atol=1e-3) def test_non_uniform_fast_fourier_op_equal_to_fft(ismrmrd_cart_high_res) -> None: @@ -239,6 +271,115 @@ def test_non_uniform_fast_fourier_op_repr(): assert 'device' in repr_str +def test_subspace_non_uniform_fast_fourier_op_gram() -> None: + """Test subspace Toeplitz Gram against explicit expand-apply-compress reference.""" + rng = RandomGenerator(seed=1) + n_timepoints, n_coefficients = 4, 2 + image_shape = (12, 10) + + nufft_op = create_time_varying_2d_nufft_op(n_timepoints=n_timepoints, image_shape=image_shape) + basis = rng.complex64_tensor((n_timepoints, n_coefficients)) + alpha = rng.complex64_tensor((n_coefficients, 1, 1, *image_shape)) + + subspace_gram = nufft_op.toeplitz(subspace=basis) + assert isinstance(subspace_gram, SubspaceNonUniformFastFourierOpGramOp) + + expanded = einops.einsum(basis, alpha, 'time coeff, coeff joint coil ... -> time joint coil ...') + (kspace,) = nufft_op(expanded) + (backprojected,) = nufft_op.H(kspace) + expected = einops.einsum(basis.conj(), backprojected, 'time coeff, time joint coil ... -> coeff joint coil ...') + (actual,) = subspace_gram(alpha) + + torch.testing.assert_close(actual, expected, rtol=2e-3, atol=2e-3) + + +def test_subspace_non_uniform_fast_fourier_op_gram_accepts_pca_operator() -> None: + """Test subspace Gram accepts a PCACompressionOp as basis input.""" + rng = RandomGenerator(seed=2) + n_timepoints, n_coefficients = 4, 2 + image_shape = (12, 10) + + nufft_op = create_time_varying_2d_nufft_op(n_timepoints=n_timepoints, image_shape=image_shape) + training_signals = rng.complex64_tensor((1, 32, n_timepoints)) + pca_op = PCACompressionOp(training_signals, n_components=n_coefficients, centering=False) + basis = pca_op._compression_matrix.squeeze(0).mH + alpha = rng.complex64_tensor((n_coefficients, 1, 1, *image_shape)) + + subspace_gram_from_pca = SubspaceNonUniformFastFourierOpGramOp(nufft_op, pca_op) + subspace_gram_from_basis = SubspaceNonUniformFastFourierOpGramOp(nufft_op, basis) + + (actual_from_pca,) = subspace_gram_from_pca(alpha) + (actual_from_basis,) = subspace_gram_from_basis(alpha) + + torch.testing.assert_close(actual_from_pca, actual_from_basis) + + +def test_non_uniform_fast_fourier_op_gram_autograd() -> None: + """Test autograd of the Toeplitz Gram operator.""" + rng = RandomGenerator(seed=3) + n_timepoints = 2 + image_shape = (5, 4) + nufft_op = create_time_varying_2d_nufft_op( + n_timepoints=n_timepoints, + image_shape=image_shape, + dtype=torch.float64, + ) + operator = nufft_op.toeplitz().double() + image = rng.complex128_tensor((n_timepoints, 1, 1, *image_shape)).requires_grad_(True) + + gradcheck(operator, (image,), fast_mode=True) + + +def test_subspace_non_uniform_fast_fourier_op_gram_autograd() -> None: + """Test autograd of the subspace Toeplitz Gram operator.""" + rng = RandomGenerator(seed=4) + n_timepoints, n_coefficients = 3, 2 + image_shape = (5, 4) + nufft_op = create_time_varying_2d_nufft_op( + n_timepoints=n_timepoints, + image_shape=image_shape, + dtype=torch.float64, + ) + basis = rng.complex128_tensor((n_timepoints, n_coefficients)) + operator = nufft_op.toeplitz(subspace=basis).double() + coefficients = rng.complex128_tensor((n_coefficients, 1, 1, *image_shape)).requires_grad_(True) + + gradcheck(operator, (coefficients,), fast_mode=True) + + +@pytest.mark.cuda +def test_non_uniform_fast_fourier_op_gram_cuda() -> None: + """Test Toeplitz Gram operators work on CUDA devices.""" + rng = RandomGenerator(seed=5) + n_timepoints, n_coefficients = 4, 2 + image_shape = (12, 10) + + nufft_op = create_time_varying_2d_nufft_op(n_timepoints=n_timepoints, image_shape=image_shape) + image = rng.complex64_tensor((n_timepoints, 1, 1, *image_shape)) + coefficients = rng.complex64_tensor((n_coefficients, 1, 1, *image_shape)) + basis = rng.complex64_tensor((n_timepoints, n_coefficients)) + + gram = nufft_op.toeplitz() + gram.cuda() + (result,) = gram(image.cuda()) + assert result.is_cuda + + subspace_gram = nufft_op.toeplitz(subspace=basis) + subspace_gram.cuda() + (result,) = subspace_gram(coefficients.cuda()) + assert result.is_cuda + + nufft_op = create_time_varying_2d_nufft_op( + n_timepoints=n_timepoints, + image_shape=image_shape, + device='cuda', + ) + (result,) = nufft_op.toeplitz()(image.cuda()) + assert result.is_cuda + (result,) = nufft_op.toeplitz(subspace=basis.cuda())(coefficients.cuda()) + assert result.is_cuda + + @pytest.mark.cuda def test_non_uniform_fast_fourier_op_cuda() -> None: """Test non-uniform fast Fourier operator works on CUDA devices.""" From c5f6f04e0be31d38bd0532d2302d04ffbc9e42bc Mon Sep 17 00:00:00 2001 From: Felix F Zimmermann Date: Fri, 8 May 2026 14:02:59 +0200 Subject: [PATCH 2/3] Update non-uniform fast Fourier operator tests Refactor non-uniform fast Fourier operator tests to include density compensation and ensure compatibility with CUDA. --- .../test_non_uniform_fast_fourier_op.py | 73 ++++++++++++------- 1 file changed, 48 insertions(+), 25 deletions(-) diff --git a/tests/operators/test_non_uniform_fast_fourier_op.py b/tests/operators/test_non_uniform_fast_fourier_op.py index 8848554f7..f7ce812c2 100644 --- a/tests/operators/test_non_uniform_fast_fourier_op.py +++ b/tests/operators/test_non_uniform_fast_fourier_op.py @@ -3,10 +3,11 @@ import einops import pytest import torch -from mrpro.data import KData, KTrajectory -from mrpro.data.SpatialDimension import SpatialDimension +from mrpro.data import DcfData, KData, KTrajectory +from mrpro.data import SpatialDimension from mrpro.data.traj_calculators import KTrajectoryIsmrmrd from mrpro.operators import ( + DensityCompensationOp, FastFourierOp, NonUniformFastFourierOp, PCACompressionOp, @@ -31,11 +32,10 @@ def create_time_varying_2d_nufft_op( n_timepoints: int = 4, image_shape: tuple[int, int] = (12, 10), dtype: torch.dtype = torch.float32, - device: torch.device | str = 'cpu', ) -> NonUniformFastFourierOp: """Create a small 2D NUFFT with one radial spoke per time point.""" - angles = torch.linspace(0, torch.pi, n_timepoints + 1, dtype=dtype, device=device)[:-1] - readout = torch.linspace(-image_shape[-1] / 2, image_shape[-1] / 2, image_shape[-1], dtype=dtype, device=device) + angles = torch.linspace(0, torch.pi, n_timepoints + 1, dtype=dtype)[:-1] + readout = torch.linspace(-image_shape[-1] / 2, image_shape[-1] / 2, image_shape[-1], dtype=dtype) kx = (torch.cos(angles)[:, None] * readout[None, :])[:, None, None, :] ky = (torch.sin(angles)[:, None] * readout[None, :])[:, None, None, :] kx = kx[:, None, :, :] @@ -104,10 +104,8 @@ def test_non_uniform_fast_fourier_op_gram( (expected,) = (nufft_op.H @ nufft_op)(img) (actual,) = nufft_op.gram(img) - (actual_toeplitz,) = nufft_op.toeplitz()(img) torch.testing.assert_close(actual, expected, rtol=1e-3, atol=1e-3) - torch.testing.assert_close(actual_toeplitz, expected, rtol=1e-3, atol=1e-3) def test_non_uniform_fast_fourier_op_equal_to_fft(ismrmrd_cart_high_res) -> None: @@ -314,6 +312,25 @@ def test_subspace_non_uniform_fast_fourier_op_gram_accepts_pca_operator() -> Non torch.testing.assert_close(actual_from_pca, actual_from_basis) +def test_non_uniform_fast_fourier_op_weighted_toeplitz_matches_explicit_dcf_normal_operator() -> None: + """Test Toeplitz(weight=dcf) matches the explicit weighted normal operator F^H DCF F.""" + rng = RandomGenerator(seed=6) + n_timepoints = 4 + image_shape = (12, 10) + + nufft_op = create_time_varying_2d_nufft_op(n_timepoints=n_timepoints, image_shape=image_shape) + image = rng.complex64_tensor((n_timepoints, 1, 1, *image_shape)) + dcf = DcfData(data=rng.float32_tensor((n_timepoints, 1, 1, 1, image_shape[-1]))) + + explicit_operator = nufft_op.H @ DensityCompensationOp(dcf) @ nufft_op + weighted_toeplitz = nufft_op.toeplitz(weight=dcf) + + (expected,) = explicit_operator(image) + (actual,) = weighted_toeplitz(image) + + torch.testing.assert_close(actual, expected, rtol=2e-3, atol=2e-3) + + def test_non_uniform_fast_fourier_op_gram_autograd() -> None: """Test autograd of the Toeplitz Gram operator.""" rng = RandomGenerator(seed=3) @@ -351,32 +368,38 @@ def test_subspace_non_uniform_fast_fourier_op_gram_autograd() -> None: def test_non_uniform_fast_fourier_op_gram_cuda() -> None: """Test Toeplitz Gram operators work on CUDA devices.""" rng = RandomGenerator(seed=5) - n_timepoints, n_coefficients = 4, 2 + n_timepoints = 4 image_shape = (12, 10) + image = rng.complex64_tensor((n_timepoints, 1, 1, *image_shape)).cuda() + + nufft_op = create_time_varying_2d_nufft_op(n_timepoints=n_timepoints, image_shape=image_shape).cuda() + gram = nufft_op.gram.cuda() + (result,) = gram(image) + assert result.is_cuda nufft_op = create_time_varying_2d_nufft_op(n_timepoints=n_timepoints, image_shape=image_shape) - image = rng.complex64_tensor((n_timepoints, 1, 1, *image_shape)) - coefficients = rng.complex64_tensor((n_coefficients, 1, 1, *image_shape)) + gram = nufft_op.gram + (result,) = gram(image) + assert result.is_cuda + + +@pytest.mark.cuda +def test_non_uniform_fast_fourier_op_subspace_gram_cuda() -> None: + """Test subspace Toeplitz Gram operators work on CUDA devices.""" + rng = RandomGenerator(seed=5) + n_timepoints, n_coefficients = 4, 2 + image_shape = (12, 10) + coefficients = rng.complex64_tensor((n_coefficients, 1, 1, *image_shape)).cuda() basis = rng.complex64_tensor((n_timepoints, n_coefficients)) - gram = nufft_op.toeplitz() - gram.cuda() - (result,) = gram(image.cuda()) + nufft_op = create_time_varying_2d_nufft_op(n_timepoints=n_timepoints, image_shape=image_shape) + subspace_gram = nufft_op.toeplitz(subspace=basis).cuda() + (result,) = subspace_gram(coefficients) assert result.is_cuda + nufft_op = create_time_varying_2d_nufft_op(n_timepoints=n_timepoints, image_shape=image_shape).cuda() subspace_gram = nufft_op.toeplitz(subspace=basis) - subspace_gram.cuda() - (result,) = subspace_gram(coefficients.cuda()) - assert result.is_cuda - - nufft_op = create_time_varying_2d_nufft_op( - n_timepoints=n_timepoints, - image_shape=image_shape, - device='cuda', - ) - (result,) = nufft_op.toeplitz()(image.cuda()) - assert result.is_cuda - (result,) = nufft_op.toeplitz(subspace=basis.cuda())(coefficients.cuda()) + (result,) = subspace_gram(coefficients) assert result.is_cuda From 25a2139031ea5db56cbac9c0c4d6aa717347b1b8 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 8 May 2026 12:04:12 +0000 Subject: [PATCH 3/3] [pre-commit] auto fixes from pre-commit hooks --- tests/operators/test_non_uniform_fast_fourier_op.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/operators/test_non_uniform_fast_fourier_op.py b/tests/operators/test_non_uniform_fast_fourier_op.py index f7ce812c2..1489cf02f 100644 --- a/tests/operators/test_non_uniform_fast_fourier_op.py +++ b/tests/operators/test_non_uniform_fast_fourier_op.py @@ -3,8 +3,7 @@ import einops import pytest import torch -from mrpro.data import DcfData, KData, KTrajectory -from mrpro.data import SpatialDimension +from mrpro.data import DcfData, KData, KTrajectory, SpatialDimension from mrpro.data.traj_calculators import KTrajectoryIsmrmrd from mrpro.operators import ( DensityCompensationOp,