From 479a4a5b91073d27b9e377139c6b7a1975d8b2e5 Mon Sep 17 00:00:00 2001 From: Felix Zimmermann Date: Tue, 15 Apr 2025 00:07:12 +0200 Subject: [PATCH 1/2] add circulant preconditoner --- .../operators/CirculantPreconditioner.py | 65 +++++++++++++++++++ 1 file changed, 65 insertions(+) create mode 100644 src/mrpro/operators/CirculantPreconditioner.py diff --git a/src/mrpro/operators/CirculantPreconditioner.py b/src/mrpro/operators/CirculantPreconditioner.py new file mode 100644 index 000000000..cd99c1b9d --- /dev/null +++ b/src/mrpro/operators/CirculantPreconditioner.py @@ -0,0 +1,65 @@ +"""A preconditioner for non-Cartesian iterative SENSE reconstruction.""" + +import torch + +from mrpro.data.DcfData import DcfData +from mrpro.operators.LinearOperator import LinearOperator +from mrpro.operators.NonUniformFastFourierOp import NonUniformFastFourierOp + + +class CirculantPreconditioner(LinearOperator): + """A preconditioner for a non-Cartesian SENSE reconstruction.""" + + def __init__(self, nufft_operator: NonUniformFastFourierOp, dcf: DcfData): + """Initialize a circulant preconditioner for a non-Cartesian SENSE reconstruction. + + This operator acts as a preconditioner P for iterative algorithms + solving Ax=b, where A involves NUFFT operations (F) andoptionally + coil sensitivities (C), e.g., A = C^H F^H F C. + + The preconditioner approximates the inverse of the density-compensated + operator, P ≈ (F^H W F)^(-1), where W represents the density + compensation factors (DCF). It is constructed by: + 1. Simulating the density-compensated Point Spread Function (PSF) + h_w = F^H W F(delta). + 2. Computing the FFT of the PSF. These are the + eigenvalues of the circulant approximation. + 3. Regularizing and inverting these eigenvalues to get the k-space + kernel + + This preconditioner is suitable for accelerating solvers for the + *unweighted* least-squares problem (where A does not include W), + as it compensates for density variations internally. + + Parameters + ---------- + nufft_operator + The non-uniform fast fourier transform operator. + dcf + density compensation weights + """ + super().__init__() + delta_image = torch.zeros(nufft_operator._im_size, dtype=torch.complex64) + center_indices = tuple(d // 2 for d in nufft_operator._im_size) + delta_image[center_indices] = 1.0 + delta_image = delta_image.to(dcf.device) + (k,) = nufft_operator(delta_image) + k = k * dcf.data + (psf,) = nufft_operator.adjoint(k) + kernel = torch.fft.fftn(psf, dim=(-1, -2, -3)) + kernel = torch.polar(kernel.abs().clamp(min=1e-5), kernel.angle()).reciprocal() + self.kernel = kernel + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor,]: + """Apply the inverse of the preconditioner.""" + x = torch.fft.fftn(x, dim=(-1, -2, -3)) + x = x * self.kernel + x = torch.fft.ifftn(x, dim=(-1, -2, -3)) + return (x,) + + def adjoint(self, y: torch.Tensor) -> tuple[torch.Tensor,]: + """Apply the adjoint of the inverse of the preconditioner.""" + y = torch.fft.fftn(y, dim=(-1, -2, -3)) + y = y * self.kernel.conj() + y = torch.fft.ifftn(y, dim=(-1, -2, -3)) + return (y,) From 4685bf3001b8015fad29a3519b240baea598a81c Mon Sep 17 00:00:00 2001 From: Felix Zimmermann Date: Tue, 10 Feb 2026 15:53:50 +0100 Subject: [PATCH 2/2] add test --- .../operators/CirculantPreconditioner.py | 24 +-- src/mrpro/operators/__init__.py | 2 + .../test_circulant_preconditioner.py | 144 ++++++++++++++++++ 3 files changed, 160 insertions(+), 10 deletions(-) create mode 100644 tests/operators/test_circulant_preconditioner.py diff --git a/src/mrpro/operators/CirculantPreconditioner.py b/src/mrpro/operators/CirculantPreconditioner.py index cd99c1b9d..9217ae4c6 100644 --- a/src/mrpro/operators/CirculantPreconditioner.py +++ b/src/mrpro/operators/CirculantPreconditioner.py @@ -39,14 +39,18 @@ def __init__(self, nufft_operator: NonUniformFastFourierOp, dcf: DcfData): density compensation weights """ super().__init__() - delta_image = torch.zeros(nufft_operator._im_size, dtype=torch.complex64) - center_indices = tuple(d // 2 for d in nufft_operator._im_size) - delta_image[center_indices] = 1.0 - delta_image = delta_image.to(dcf.device) + device = dcf.device if dcf.device is not None else nufft_operator._omega.device + im_shape_zyx = [1, 1, 1] + for dim, size in zip(nufft_operator._direction_zyx, nufft_operator._im_size, strict=True): + im_shape_zyx[dim + 3] = size + + delta_image = torch.zeros((1, *im_shape_zyx), dtype=torch.complex64, device=device) + center_indices = tuple(size // 2 for size in im_shape_zyx) + delta_image[(0, *center_indices)] = 1.0 (k,) = nufft_operator(delta_image) k = k * dcf.data (psf,) = nufft_operator.adjoint(k) - kernel = torch.fft.fftn(psf, dim=(-1, -2, -3)) + kernel = torch.fft.fftn(torch.fft.ifftshift(psf, dim=(-1, -2, -3)), dim=(-1, -2, -3)) kernel = torch.polar(kernel.abs().clamp(min=1e-5), kernel.angle()).reciprocal() self.kernel = kernel @@ -57,9 +61,9 @@ def forward(self, x: torch.Tensor) -> tuple[torch.Tensor,]: x = torch.fft.ifftn(x, dim=(-1, -2, -3)) return (x,) - def adjoint(self, y: torch.Tensor) -> tuple[torch.Tensor,]: + def adjoint(self, x: torch.Tensor) -> tuple[torch.Tensor,]: """Apply the adjoint of the inverse of the preconditioner.""" - y = torch.fft.fftn(y, dim=(-1, -2, -3)) - y = y * self.kernel.conj() - y = torch.fft.ifftn(y, dim=(-1, -2, -3)) - return (y,) + x = torch.fft.fftn(x, dim=(-1, -2, -3)) + x = x * self.kernel.conj() + x = torch.fft.ifftn(x, dim=(-1, -2, -3)) + return (x,) diff --git a/src/mrpro/operators/__init__.py b/src/mrpro/operators/__init__.py index 3c357dfef..cbe7a6a8c 100644 --- a/src/mrpro/operators/__init__.py +++ b/src/mrpro/operators/__init__.py @@ -6,6 +6,7 @@ from mrpro.operators import functionals, models from mrpro.operators.AveragingOp import AveragingOp from mrpro.operators.CartesianSamplingOp import CartesianSamplingOp, CartesianMaskingOp +from mrpro.operators.CirculantPreconditioner import CirculantPreconditioner from mrpro.operators.ConjugateGradientOp import ConjugateGradientOp from mrpro.operators.ConstraintsOp import ConstraintsOp from mrpro.operators.DensityCompensationOp import DensityCompensationOp @@ -39,6 +40,7 @@ "AveragingOp", "CartesianMaskingOp", "CartesianSamplingOp", + "CirculantPreconditioner", "ConjugateGradientOp", "ConstraintsOp", "DensityCompensationOp", diff --git a/tests/operators/test_circulant_preconditioner.py b/tests/operators/test_circulant_preconditioner.py new file mode 100644 index 000000000..f32a69cd8 --- /dev/null +++ b/tests/operators/test_circulant_preconditioner.py @@ -0,0 +1,144 @@ +"""Tests for circulant preconditioner operator.""" + +import pytest +import torch +from mrpro.algorithms.optimizers import cg +from mrpro.data import DcfData +from mrpro.data.SpatialDimension import SpatialDimension +from mrpro.operators import CirculantPreconditioner, NonUniformFastFourierOp +from mrpro.utils import RandomGenerator + +from tests import dotproduct_adjointness_test +from tests.conftest import create_traj + + +def create_circulant_preconditioner_and_domain_range() -> tuple[ + CirculantPreconditioner, + NonUniformFastFourierOp, + DcfData, + torch.Tensor, + torch.Tensor, +]: + """Create circulant preconditioner and random elements from domain and range.""" + rng = RandomGenerator(seed=0) + + img_shape = (1, 1, 1, 24, 24) + nkx = (1, 1, 1, 12, 24) + nky = (1, 1, 1, 12, 24) + nkz = (1, 1, 1, 1, 1) + traj = create_traj(nkx, nky, nkz, 'non-uniform', 'non-uniform', 'zero') + + recon_matrix = SpatialDimension(img_shape[-3], img_shape[-2], img_shape[-1]) + encoding_matrix = SpatialDimension( + int(traj.kz.max() - traj.kz.min() + 1), + int(traj.ky.max() - traj.ky.min() + 1), + int(traj.kx.max() - traj.kx.min() + 1), + ) + direction = [d for d, e in zip(('z', 'y', 'x'), encoding_matrix.zyx, strict=False) if e > 1] + nufft_op = NonUniformFastFourierOp( + direction=direction, # type: ignore[arg-type] + recon_matrix=recon_matrix, + encoding_matrix=encoding_matrix, + traj=traj, + ) + + dcf = DcfData.from_traj_voronoi(traj) + circulant_preconditioner = CirculantPreconditioner(nufft_op, dcf) + + u = rng.complex64_tensor(size=img_shape) + v = rng.complex64_tensor(size=img_shape) + + return circulant_preconditioner, nufft_op, dcf, u, v + + +def test_circulant_preconditioner_adjointness() -> None: + """Test adjoint property of circulant preconditioner.""" + circulant_preconditioner, _, _, u, v = create_circulant_preconditioner_and_domain_range() + dotproduct_adjointness_test(circulant_preconditioner, u, v) + + +def test_circulant_preconditioner_cg_iteration() -> None: + """Test circulant preconditioner in CG iterations.""" + circulant_preconditioner, nufft_op, dcf, _, _ = create_circulant_preconditioner_and_domain_range() + + operator = nufft_op.H @ dcf.as_operator() @ nufft_op + rng = RandomGenerator(seed=1) + true_solution = rng.complex64_tensor(size=(1, 1, 1, 24, 24)) + (right_hand_side,) = operator(true_solution) + + initial_value = torch.zeros_like(true_solution) + initial_residual = torch.linalg.vector_norm(right_hand_side.flatten()) + + (solution_without_preconditioner,) = cg( + operator, + right_hand_side, + initial_value=initial_value, + max_iterations=3, + tolerance=0, + ) + residual_without_preconditioner = torch.linalg.vector_norm( + (operator(solution_without_preconditioner)[0] - right_hand_side).flatten() + ) + + (solution_with_preconditioner,) = cg( + operator, + right_hand_side, + initial_value=initial_value, + preconditioner_inverse=circulant_preconditioner, + max_iterations=3, + tolerance=0, + ) + residual_with_preconditioner = torch.linalg.vector_norm( + (operator(solution_with_preconditioner)[0] - right_hand_side).flatten() + ) + + assert residual_without_preconditioner < initial_residual + assert residual_with_preconditioner < initial_residual + + +@pytest.mark.cuda +def test_circulant_preconditioner_cuda() -> None: + """Test circulant preconditioner works on CUDA devices.""" + rng = RandomGenerator(seed=2) + x = rng.complex64_tensor(size=(1, 1, 1, 24, 24)) + + # Create on CPU, transfer to GPU, run on GPU + preconditioner, _, _, _, _ = create_circulant_preconditioner_and_domain_range() + preconditioner.cuda() + (result,) = preconditioner(x.cuda()) + assert result.is_cuda + + # Create on CPU, run on CPU + preconditioner, _, _, _, _ = create_circulant_preconditioner_and_domain_range() + (result,) = preconditioner(x) + assert result.is_cpu + + # Create on GPU, run on GPU + img_shape = (1, 1, 1, 24, 24) + nkx = (1, 1, 1, 12, 24) + nky = (1, 1, 1, 12, 24) + nkz = (1, 1, 1, 1, 1) + traj = create_traj(nkx, nky, nkz, 'non-uniform', 'non-uniform', 'zero').cuda() + + recon_matrix = SpatialDimension(img_shape[-3], img_shape[-2], img_shape[-1]) + encoding_matrix = SpatialDimension( + int(traj.kz.max() - traj.kz.min() + 1), + int(traj.ky.max() - traj.ky.min() + 1), + int(traj.kx.max() - traj.kx.min() + 1), + ) + direction = [d for d, e in zip(('z', 'y', 'x'), encoding_matrix.zyx, strict=False) if e > 1] + nufft_op = NonUniformFastFourierOp( + direction=direction, # type: ignore[arg-type] + recon_matrix=recon_matrix, + encoding_matrix=encoding_matrix, + traj=traj, + ) + dcf = DcfData.from_traj_voronoi(traj) + preconditioner = CirculantPreconditioner(nufft_op, dcf) + (result,) = preconditioner(x.cuda()) + assert result.is_cuda + + # Create on GPU, transfer to CPU, run on CPU + preconditioner.cpu() + (result,) = preconditioner(x) + assert result.is_cpu