Fixes #258: Enable NUFFT operator to support batched sensitivity maps for batched inputs

Author: Asma TANABENE

src/mrinufft/operators/autodiff.py

@@ -152,3 +152,94 @@ def samples(self, value):
     def __getattr__(self, name):
         """Forward all other attributes to the nufft_op."""
         return getattr(self.nufft_op, name)
+
+
+class BatchedNufftAutoGrad(MRINufftAutoGrad):
+    """
+    A batched wrapper for NUFFT operator with support for autodifferentiation
+    and varying sensitivity maps (smaps) across batches.
+
+    Parameters
+    ----------
+    nufft : object
+        An instance of a standard NUFFT operator.
+    batch_size : int
+        Number of batches to process simultaneously.
+    **kwargs : dict, optional
+        Additional arguments.
+
+    Notes
+    -----
+    #TODO Future improvements may include support for varying trajectories across batches.
+    """
+
+    def __init__(self, nufft_op, wrt_data=True, wrt_traj=False, batch_size=1, **kwargs):
+        super().__init__(
+            nufft_op=nufft_op, wrt_data=wrt_data, wrt_traj=wrt_traj
+        )
+        self.batch_size = batch_size
+
+    def op(self, batched_smaps, batched_imgs):
+        """Compute the forward batched_imgs -> batched_kspace."""
+        self._check_input_shape(imgs=batched_imgs)
+        self._check_input_shape(imgs=batched_smaps)
+        batched_kspace = []
+        for i in range(self.batch_size):
+            try:
+                # update smaps for proper backward computation
+                self.nufft_op.smaps = batched_smaps[i]
+                batched_kspace.append(
+                    _NUFFT_OP.apply(batched_imgs[i], self.samples, self.nufft_op)
+                )
+            except Exception as e:
+                raise RuntimeError(
+                    f"Failed at batch index {i+1}: {e}"
+                )  # For an easier debugging
+        return torch.stack(batched_kspace, dim=0)
+
+    def adj_op(self, batched_smaps, batched_kspace):
+        """Compute the adjoint batched_kspace -> batched_imgs."""
+        self._check_input_shape(ksps=batched_kspace)
+        self._check_input_shape(imgs=batched_smaps)
+        batched_imgs = []
+        for i in range(self.batch_size):
+            try:
+                self.nufft_op.smaps = batched_smaps[i]
+                batched_imgs.append(
+                    _NUFFT_ADJOP.apply(batched_kspace[i], self.samples, self.nufft_op)
+                )
+            except Exception as e:
+                raise RuntimeError(f"Failed at batch index {i+1}: {e}")
+        return torch.stack(batched_imgs, dim=0)
+
+    def _check_input_shape(self, *, imgs=None, ksps=None):
+        """
+        Validates the batch size of either image or k-space input against the expected batch size.
+
+        Parameters
+        ----------
+        imgs : np.ndarray, optional
+            Image data array. If provided, its batch dimension will be validated.
+
+        ksps : np.ndarray or object, optional
+            K-space data array or compatible object. If provided, its batch dimension will be validated.
+
+        Raises
+        ------
+        ValueError
+            If the batch size of the image or k-space input does not match the expected batch size.
+        """
+        if imgs is not None:
+            if imgs.shape[0] != self.batch_size:
+                raise ValueError(
+                    f"Image batch size mismatch: got {imgs.shape[0]}, expected {self.batch_size}. "
+                    f"Image shape: {imgs.shape}"
+                )
+        if ksps is not None:
+            if ksps.shape[0] != self.batch_size:
+                raise ValueError(
+                    f"K-space batch size mismatch: got {ksps.shape[0]}, expected {self.batch_size}. "
+                    f"K-space shape: {ksps.shape}"
+                )
+        if imgs is None and ksps is None:
+            raise ValueError("Provide either `imgs` or `ksps` as input.")

src/mrinufft/operators/base.py

@@ -50,7 +50,13 @@ def list_backends(available_only=False):
 
 
 def get_operator(
-    backend_name: str, wrt_data: bool = False, wrt_traj: bool = False, *args, **kwargs
+    backend_name: str,
+    wrt_data: bool = False,
+    wrt_traj: bool = False,
+    use_batched_mode: bool = False,
+    batch_size: int = 1,
+    *args,
+    **kwargs,
 ):
     """Return an MRI Fourier operator interface using the correct backend.
 
@@ -62,6 +68,10 @@ def get_operator(
         if set gradients wrt to data and images will be available.
     wrt_traj: bool, default False
         if set gradients wrt to trajectory will be available.
+    use_batched_mode : bool, optional
+        If True, uses a batched version of the NUFFT operator that supports varying data/smaps pairs.
+    batch_size : int, optional
+        Batch size to be used in batched mode. Only relevant if `use_batched_mode=True`. Default is 1.
     *args, **kwargs:
         Arguments to pass to the operator constructor.
 
@@ -97,10 +107,14 @@ class or instance of class if args or kwargs are given.
     # if autograd:
     if wrt_data or wrt_traj:
         if isinstance(operator, FourierOperatorBase):
-            operator = operator.make_autograd(wrt_data, wrt_traj)
+            operator = operator.make_autograd(
+                wrt_data, wrt_traj, use_batched_mode, batch_size
+            )
         else:
             # instance will be created later
-            operator = partial(operator.with_autograd, wrt_data, wrt_traj)
+            operator = partial(
+                operator.with_autograd, wrt_data, wrt_traj, use_batched_mode, batch_size
+            )
 
     return operator
 
@@ -257,7 +271,9 @@ def compute_smaps(self, method: NDArray | Callable | str | dict | None = None):
             **kwargs,
         )
 
-    def make_autograd(self, wrt_data=True, wrt_traj=False):
+    def make_autograd(
+        self, wrt_data=True, wrt_traj=False, use_batched_mode=False, batch_size=1
+    ):
         """Make a new Operator with autodiff support.
 
         Parameters
@@ -271,6 +287,12 @@ def make_autograd(self, wrt_data=True, wrt_traj=False):
         wrt_traj : bool, optional
             If the gradient with respect to the trajectory is computed, default is false
 
+        use_batched_mode : bool, optional
+            If True, uses a batched version of the NUFFT operator that supports varying smaps
+
+        batch_size : int, optional
+            Batch size to be used in batched mode. Only relevant if `use_batched_mode=True`. Default is 1.
+
         Returns
         -------
         torch.nn.module
@@ -286,9 +308,20 @@ def make_autograd(self, wrt_data=True, wrt_traj=False):
         if not self.autograd_available:
             raise ValueError("Backend does not support auto-differentiation.")
 
-        from mrinufft.operators.autodiff import MRINufftAutoGrad
+        if use_batched_mode:
+            if batch_size < 1:
+                raise ValueError(
+                    "Provide a valid batch size." f"Batch size : {batch_size}"
+                )
+            from mrinufft.operators.autodiff import BatchedNufftAutoGrad
+
+            return BatchedNufftAutoGrad(
+                self, wrt_data=wrt_data, wrt_traj=wrt_traj, batch_size=batch_size
+            )
+        else:
+            from mrinufft.operators.autodiff import MRINufftAutoGrad
 
-        return MRINufftAutoGrad(self, wrt_data=wrt_data, wrt_traj=wrt_traj)
+            return MRINufftAutoGrad(self, wrt_data=wrt_data, wrt_traj=wrt_traj)
 
     def compute_density(self, method=None):
         """Compute the density compensation weights and set it.
@@ -476,9 +509,19 @@ def __repr__(self):
         )
 
     @classmethod
-    def with_autograd(cls, wrt_data=True, wrt_traj=False, *args, **kwargs):
+    def with_autograd(
+        cls,
+        wrt_data=True,
+        wrt_traj=False,
+        use_batched_mode=False,
+        batch_size=1,
+        *args,
+        **kwargs,
+    ):
         """Return a Fourier operator with autograd capabilities."""
-        return cls(*args, **kwargs).make_autograd(wrt_data, wrt_traj)
+        return cls(*args, **kwargs).make_autograd(
+            wrt_data, wrt_traj, use_batched_mode, batch_size
+        )
 
 
 class FourierOperatorCPU(FourierOperatorBase):

tests/helpers/__init__.py

@@ -4,6 +4,7 @@
 from .factories import (
     kspace_from_op,
     image_from_op,
+    batchedSmpas_from_op,
     to_interface,
     from_interface,
     CUPY_AVAILABLE,
@@ -21,4 +22,5 @@
     "CUPY_AVAILABLE",
     "TORCH_AVAILABLE",
     "param_array_interface",
+    "batchedSmpas_from_op",
 ]

tests/helpers/factories.py

@@ -19,28 +19,42 @@
 
 def image_from_op(operator):
     """Generate a random image."""
+    batch_dim = (operator.batch_size,) if hasattr(operator, "batch_size") else ()
     if operator.smaps is None:
-        img = np.random.randn(operator.n_coils, *operator.shape).astype(
+        img = np.random.randn(*batch_dim, operator.n_coils, *operator.shape).astype(
             operator.cpx_dtype
         )
     elif operator.smaps is not None and operator.n_coils > 1:
-        img = np.random.randn(*operator.shape).astype(operator.cpx_dtype)
+        img = np.random.randn(*batch_dim, *operator.shape).astype(operator.cpx_dtype)
 
     img += 1j * np.random.randn(*img.shape).astype(operator.cpx_dtype)
     return img
 
 
 def kspace_from_op(operator):
     """Generate a random kspace data."""
-    kspace = (1j * np.random.randn(operator.n_coils, operator.n_samples)).astype(
-        operator.cpx_dtype
-    )
-    kspace += np.random.randn(operator.n_coils, operator.n_samples).astype(
+    batch_dim = (operator.batch_size,) if hasattr(operator, "batch_size") else ()
+    kspace = (
+        1j * np.random.randn(*batch_dim, operator.n_coils, operator.n_samples)
+    ).astype(operator.cpx_dtype)
+    kspace += np.random.randn(*batch_dim, operator.n_coils, operator.n_samples).astype(
         operator.cpx_dtype
     )
     return kspace
 
 
+def batchedSmpas_from_op(operator):
+    """Generate random batched smaps."""
+    smaps = 1j * np.random.randn(
+        operator.batch_size, operator.n_coils, *operator.shape
+    ).astype(np.complex64)
+    smaps += np.random.randn(
+        operator.batch_size, operator.n_coils, *operator.shape
+    ).astype(np.complex64)
+    smaps /= np.linalg.norm(smaps, axis=1, keepdims=True)
+    return smaps
+
+
 def to_interface(data, interface):
     """Make DATA an array from INTERFACE."""
     if interface == "cupy":