typing and deps

Author: paquiteau

pyproject.toml

@@ -4,7 +4,7 @@ description = "Python Sparse data Analysis Package external fMRI plugin."
 authors = [{name="Pierre-antoine Comby", email="pierre-antoine.comby@crans.org"}]
 readme = "README.rst"
 
-dependencies = ["modopt", "numpy", "tqdm", "joblib", "numba", "scipy", "pywavelets", "mri-nufft"]
+dependencies = ["modopt", "numpy", "tqdm", "joblib", "numba", "scipy", "pywavelets", "mri-nufft[cufinufft,gpunufft]"]
 dynamic = ["version"]
 
 [project.optional-dependencies]
@@ -13,7 +13,7 @@ doc = ["pydata-sphinx-theme", "numpydoc", "sphinx_gallery", "sphinx", "sphinx-au
 dev = ["black", "isort", "ruff"]
 
 [build-system]
-requires = ["setuptools", "setuptools-scm[toml]","wheel"]
+requires = ["setuptools", "setuptools-scm[toml]", "wheel"]
 
 ######################
 # Tool configuration #

src/fmri/operators/fourier.py

@@ -7,19 +7,17 @@
 import itertools
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
+from tracemalloc import BaseFilter
 
 import numpy as np
-from mrinufft import get_operator
 from modopt.base.backend import get_array_module
+from mrinufft import get_operator
+from numpy.typing import NDArray
 
-try:
-    from mrinufft.operators.interfaces.gpunufft import make_pinned_smaps
-except ImportError:
-    make_pinned_smaps = None
+from mrinufft.operators.interfaces.gpunufft import make_pinned_smaps
 
 from .utils.fft import fft, ifft
 
-MRINUFFT_AVAILABLE = True
 CUPY_AVAILABLE = True
 
 try:
@@ -51,15 +49,19 @@ def __init__(self):
         self.smaps = None
 
     @abstractmethod
-    def op(self, img):
+    def op(self, img: NDArray) -> NDArray:
         """Forward operator."""
         pass
 
     @abstractmethod
-    def adj_op(self, data):
+    def adj_op(self, data: NDArray) -> NDArray:
         """Adjoint operator."""
         pass
 
+    def data_consistency(self, data, obs_data):
+        """Data consistency operation"""
+        return self.adj_op(self.op(data) - obs_data)
+
 
 class CartesianSpaceFourier(SpaceFourierBase):
     """A Fourier Operator in space."""
@@ -131,7 +133,7 @@ def op(self, img):
             ksp = fft(img, axis=axes)
             return ksp * self.mask
 
-    def adj_op(self, kspace_data):
+    def adj_op(self, data):
         """Apply the adjoint operator.
 
         Parameters
@@ -146,12 +148,12 @@ def adj_op(self, kspace_data):
         """
         axes = tuple(range(-len(self.shape), 0))
         if self.n_coils > 1:
-            img = ifft(kspace_data, axis=axes)
+            img = ifft(data, axis=axes)
             if self.smaps is None:
                 return img
             return np.sum(img * np.conj(self.smaps), axis=1)
         else:
-            return ifft(kspace_data, axis=axes)
+            return ifft(data, axis=axes)
 
 
 class RepeatOperator(SpaceFourierBase):
@@ -160,22 +162,22 @@ class RepeatOperator(SpaceFourierBase):
     def __init__(self, fourier_ops):
         self.fourier_ops = list(fourier_ops)
 
-    def op(self, images):
+    def op(self, img):
         """Apply the forward operator."""
         final_ksp = np.empty(
-            (len(images), self.n_coils, self.n_samples), dtype=np.complex64
+            (len(img), self.n_coils, self.n_samples), dtype=np.complex64
         )
-        for i in range(len(images)):
-            final_ksp[i] = self.fourier_ops[i].op(images[i])
+        for i in range(len(img)):
+            final_ksp[i] = self.fourier_ops[i].op(img[i])
         return final_ksp
 
-    def adj_op(self, coeffs):
+    def adj_op(self, data):
         """Apply Adjoint Operator."""
         c = 1 if self.uses_sense else self.n_coils
-        xp = get_array_module(coeffs)
+        xp = get_array_module(data)
         final_image = xp.empty((self.n_frames, c, *self.shape), dtype=np.complex64)
-        for i in range(len(coeffs)):
-            final_image[i] = self.fourier_ops[i].adj_op(coeffs[i])
+        for i in range(len(data)):
+            final_image[i] = self.fourier_ops[i].adj_op(data[i])
         return final_image.squeeze()
 
     def __getattr__(self, attrName):
@@ -327,7 +329,7 @@ def _init_density(self, density):
 
     def _init_operators(self, **kwargs):
         # initialize all the operators
-        factory = get_operator("gpunufft")
+        factory: SpaceFourierBase = get_operator("gpunufft")
         self.fourier_ops = [None] * self.n_frames
         for i, p_img, p_ksp in zip(
             range(self.n_frames),
@@ -346,13 +348,13 @@ def _init_operators(self, **kwargs):
                 **kwargs,
             )
 
-    def op(self, images):
+    def op(self, img):
         """Apply the forward operator."""
         final_ksp = np.empty(
-            (len(images), self.n_coils, self.n_samples), dtype=np.complex64
+            (len(img), self.n_coils, self.n_samples), dtype=np.complex64
         )
-        for i in range(len(images)):
-            final_ksp[i] = self.fourier_ops[i].op(images[i])
+        for i in range(len(img)):
+            final_ksp[i] = self.fourier_ops[i].op(img[i])
         return final_ksp
 
     def adj_op(self, coeffs):

src/fmri/operators/gradient.py

@@ -1,12 +1,18 @@
-""""Gradient operators for MRI reconstruction.
+""" "Gradient operators for MRI reconstruction.
 
 Adapted from pysap-mri and Modopt libraries.
 """
 
 from functools import cached_property
 
 import numpy as np
-import cupy as cp
+
+CUPY_AVAILABLE = True
+try:
+    import cupy as cp
+except ImportError:
+    CUPY_AVAILABLE = False
+
 from modopt.math.matrix import PowerMethod
 from modopt.opt.gradient import GradBasic, GradParent
 from modopt.base.backend import get_backend, get_array_module

src/fmri/operators/utils/fft.py

@@ -3,7 +3,7 @@
 import scipy as sp
 
 
-def fft(image, axis=-1):
+def fft(image, axis: int | tuple[int, ...] = -1):
     """Apply the FFT operator.
 
     Parameters
@@ -24,7 +24,7 @@ def fft(image, axis=-1):
     )
 
 
-def ifft(kspace_data, axis=-1):
+def ifft(kspace_data, axis: int | tuple[int, ...] = -1):
     """Apply the inverse FFT operator."""
     return sp.fft.fftshift(
         sp.fft.ifftn(sp.fft.ifftshift(kspace_data, axes=axis), norm="ortho", axes=axis),