operators.py

"""MRI encoding operators used for BASS reconstructions.

The main operator here, `Emat_xyt`, implements a centered 2D FFT SENSE forward
and adjoint for data shaped as:
- mask: (T, H, W) or (H, W)
- coil_maps: (C, H, W)
- k-space: (C, T, H, W)
- image: (T, H, W)
"""

import numpy as np
from utils import fft2c, ifft2c


class Emat_xyt:
    """SENSE encoding operator for multi-coil Cartesian MRI."""

    def __init__(self, mask, coil_maps):
        self.mask = self._normalize_mask(mask)
        self.coil_maps = np.asarray(coil_maps, dtype=np.complex128)
        self.adjoint = False

    @staticmethod
    def _normalize_mask(mask: np.ndarray) -> np.ndarray:
        mask_arr = np.asarray(mask)
        if mask_arr.ndim == 2:
            mask_arr = mask_arr[np.newaxis, ...]
        if mask_arr.ndim != 3:
            raise ValueError("mask must have shape (T, H, W) or (H, W)")
        return mask_arr.astype(float, copy=False)

    @property
    def H(self):
        adjoint_op = Emat_xyt(self.mask, self.coil_maps)
        adjoint_op.adjoint = not self.adjoint
        return adjoint_op

    def __mul__(self, b):
        if self.adjoint:
            if b.ndim != 4:
                raise ValueError("Expected k-space with shape (coils, frames, height, width)")
            mask_broadcast = self.mask[np.newaxis, ...]
            masked = b.astype(np.complex128, copy=False) * mask_broadcast
            coil_images = ifft2c(masked, axes=(2, 3))
            conj_coils = np.conj(self.coil_maps).astype(np.complex128, copy=False)
            combined_image = np.sum(coil_images * conj_coils[:, np.newaxis, :, :], axis=0)
            return combined_image

        else:
            # --- Forward: H * image ---
            # Forward operation can remain vectorized as it's less problematic.
            # b is an image of shape (T, H, W)
            coils_c128 = self.coil_maps
            multi_coil_images = b[np.newaxis, ...] * coils_c128[:, np.newaxis, :, :]
            k_space = fft2c(multi_coil_images, axes=(2, 3))
            return k_space * self.mask[np.newaxis, ...]