feat: organize import and typing

Author: paquiteau

src/snake/core/engine/base.py

@@ -9,6 +9,7 @@
 from collections.abc import Mapping, Sequence
 from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
 from multiprocessing.managers import SharedMemoryManager
+from pathlib import Path
 from tempfile import TemporaryDirectory
 from typing import Any, ClassVar
 
@@ -17,21 +18,25 @@
 from numpy.typing import NDArray
 from tqdm.auto import tqdm
 
+from typing_extensions import dataclass_transform
+
 from snake._meta import EnvConfig, MetaDCRegister, batched
 
 from ...mrd_utils import MRDLoader, make_base_mrd
+from ..handlers import AbstractHandler, HandlerList
 from ..parallel import ArrayProps
 from ..phantom import DynamicData, Phantom, PropTissueEnum
-from ..simulation import SimConfig
 from ..sampling import BaseSampler
-from ..handlers import AbstractHandler, HandlerList
+from ..simulation import SimConfig
 from .utils import get_ideal_phantom, get_noise
 
+AnyPath =  str | Path
 
+@dataclass_transform(kw_only_default=True)
 class MetaEngine(MetaDCRegister):
     """MetaClass for engines."""
 
-    dunder_name = "engine"
+    dunder_name: ClassVar[str] = "engine"
 
 
 class BaseAcquisitionEngine(metaclass=MetaEngine):
@@ -113,7 +118,7 @@ def _write_chunk_data(
 
     def _acquire_ksp_job(
         self,
-        filename: os.PathLike,
+        filename: AnyPath,
         chunk: Sequence[int],
         tmp_dir: str,
         shared_phantom_props: (
@@ -160,7 +165,7 @@ def _acquire_ksp_job(
 
     def __call__(
         self,
-        filename: os.PathLike,
+        filename: AnyPath,
         sampler: BaseSampler,
         phantom: Phantom,
         sim_conf: SimConfig,
@@ -175,7 +180,7 @@ def __call__(
 
         Parameters
         ----------
-        filename : os.PathLike
+        filename : AnyPath
             The path to the MRD file.
         sampler : BaseSampler
             The sampler to use.

src/snake/core/handlers/__init__.py

@@ -4,12 +4,14 @@
 from .activations import BlockActivationHandler
 from .noise import NoiseHandler
 from .motion import RandomMotionImageHandler
+from .fov import FOVHandler
 
 __all__ = [
     "AbstractHandler",
     "HandlerList",
     "get_handler",
     "H",
+    "FOVHandler",
     "BlockActivationHandler",
     "NoiseHandler",
     "RandomMotionImageHandler",

src/snake/core/handlers/base.py

@@ -1,23 +1,25 @@
 """Base handler module."""
 
 from __future__ import annotations
-import yaml
+
 import dataclasses
 from collections import UserList
+from typing import Any, ClassVar, TypeVar
 
-from ..._meta import MetaDCRegister
-from typing import ClassVar, TypeVar, Any
+import yaml
+from typing_extensions import dataclass_transform
 
+from ..._meta import MetaDCRegister
+from ..phantom import DynamicData, KspaceDynamicData, Phantom
 from ..simulation import SimConfig
-from ..phantom import Phantom, DynamicData, KspaceDynamicData
 
 T = TypeVar("T")
 
-
+@dataclass_transform(kw_only_default=True)
 class MetaHandler(MetaDCRegister):
     """MetaClass for Handlers."""
 
-    dunder_name = "handler"
+    dunder_name: ClassVar[str] = "handler"
 
 
 class AbstractHandler(metaclass=MetaHandler):

src/snake/core/handlers/fov.py

@@ -1,4 +1,9 @@
-"""Handler for modifying the Field of View of the Phantom."""
+"""Handler for modifying the Field of View of the Phantom.
+
+TODO: Add a FOV-motion handler that combines FOV and motion (moving the head is
+equivalent to changing the center point of FOV + angles).
+
+"""
 
 from __future__ import annotations
 import warnings

src/snake/core/handlers/motion/image.py

@@ -1,4 +1,8 @@
-"""Motion in the image domain."""
+"""Motion in the image domain.
+
+TODO: Add constraints on maximal displacements
+TODO: Use FOV displacements
+"""
 
 from copy import deepcopy
 

src/snake/core/sampling/base.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 import logging
 from typing import ClassVar
+from typing_extensions import dataclass_transform
 from numpy.typing import NDArray
 
 from snake._meta import MetaDCRegister
@@ -11,6 +12,7 @@
 import ismrmrd as mrd
 
 
+@dataclass_transform(kw_only_default=True) # Required here for pyright to work.
 class MetaSampler(MetaDCRegister):
     """MetaClass for Samplers."""
 

src/snake/core/sampling/samplers.py

@@ -249,10 +249,10 @@ class StackOfSpiralSampler(NonCartesianAcquisitionSampler):
 
     acsz: float | int
     accelz: int
-    orderz: VDSorder = VDSorder.TOP_DOWN
+    orderz: str | VDSorder = VDSorder.TOP_DOWN
     nb_revolutions: int = 10
     spiral_name: str = "archimedes"
-    pdfz: VDSpdf = VDSpdf.GAUSSIAN
+    pdfz: str | VDSpdf = VDSpdf.GAUSSIAN
     constant: bool = False
     in_out: bool = True
     rotate_angle: AngleRotation = AngleRotation.ZERO

src/snake/mrd_utils/loader.py

@@ -15,6 +15,7 @@
 from .._meta import LogMixin
 
 if TYPE_CHECKING:
+    from _typeshed import AnyPath
     from ..core import Phantom, DynamicData
     from ..core import SimConfig
 
@@ -23,7 +24,7 @@
 log = logging.getLogger(__name__)
 
 
-def read_mrd_header(filename: os.PathLike | mrd.Dataset) -> mrd.xsd.ismrmrdHeader:
+def read_mrd_header(filename: AnyPath | mrd.Dataset) -> mrd.xsd.ismrmrdHeader:
     """Read the header of the MRD file."""
     if isinstance(filename, mrd.Dataset):
         dataset = filename
@@ -50,7 +51,7 @@ class MRDLoader(LogMixin):
 
     def __init__(
         self,
-        filename: os.PathLike,
+        filename: AnyPath,
         dataset_name: str = "dataset",
         writeable: bool = False,
         swmr: bool = False,
@@ -85,7 +86,7 @@ def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any):
 
     def iter_frames(
         self, start: int | None = None, stop: int | None = None, step: int | None = None
-    ) -> Generator[tuple[int, NDArray, NDArray], None, None]:
+    ) -> Generator[tuple[int, NDArray[np.float32], NDArray[np.complex64]], None, None]:
         """Iterate over kspace frames of the dataset."""
         if start is None:
             start = 0
@@ -97,7 +98,7 @@ def iter_frames(
             for i in np.arange(start, stop, step):
                 yield i, *self.get_kspace_frame(i)
 
-    def get_kspace_frame(self, idx: int) -> tuple[NDArray, NDArray]:
+    def get_kspace_frame(self, idx: int) -> tuple[NDArray[np.float32], NDArray[np.complex64]]:
         """Get k-space frame trajectory/mask and data."""
         raise NotImplementedError()
 
@@ -275,15 +276,15 @@ def get_sim_conf(self) -> SimConfig:
         """Parse the sim config."""
         return parse_sim_conf(self.header)
 
-    def _get_image_data(self, name: str, idx: int = 0) -> NDArray | None:
+    def _get_image_data(self, name: str, idx: int = 0) -> NDArray[np.complex64] | None:
         try:
-            image = self._read_image(name, idx).data
+            image = self._read_image(name, idx).data.astype(np.complex64)
         except LookupError:
             log.warning(f"No {name} found in the dataset.")
             return None
         return image
 
-    def get_smaps(self) -> NDArray | None:
+    def get_smaps(self) -> NDArray[np.complex64] | None:
         """Load the sensitivity maps from the dataset."""
         return self._get_image_data("smaps")
 
@@ -344,7 +345,7 @@ class NonCartesianFrameDataLoader(MRDLoader):
 
     def get_kspace_frame(
         self, idx: int, shot_dim: bool = False
-    ) -> tuple[np.ndarray, np.ndarray]:
+    ) -> tuple[NDArray[np.float32], NDArray[np.complex64]]:
         """Get the k-space frame and the associated trajectory.
 
         Parameters

src/snake/toolkit/plotting.py

@@ -198,8 +198,8 @@ def plot_frames_activ(
 
 
 def axis3dcut(
-    background: NDArray,
-    z_score: NDArray,
+    background: NDArray[np.float32],
+    z_score: NDArray[np.float32] | None,
     gt_roi: NDArray | None = None,
     width_inches: float = 7,
     cbar: bool = True,
@@ -312,17 +312,19 @@ def axis3dcut(
             slices_[i],
             bbox_[i],
             bg_cmap=bg_cmap,
+            z_thresh=z_thresh,
+            z_max=z_max
         )
 
     if cbar:
         cax = type(ax)(fig, ax.get_position(original=True))
         cax.set_axes_locator(divider.new_locator(nx=3, ny=0, ny1=-1))
         if z_score is not None:
             im = ScalarMappable(norm="linear", cmap=get_coolgraywarm())
-            im.set_clim(-11, 11)
+            im.set_clim(-z_max, z_max)
             matplotlib.colorbar.Colorbar(cax, im, orientation="vertical")
             cax.set_ylabel("z-scores", labelpad=-20)
-            cax.set_yticks(np.concatenate([-np.arange(3, 12, 2), np.arange(3, 12, 2)]))
+            cax.set_yticks(np.concatenate([-np.arange(z_thresh, z_max+1, 2), np.arange(z_thresh, z_max+1, 2)]))
         else:
             # use the background image
             if vmin_vmax is None:

src/snake/toolkit/reconstructors/base.py

@@ -1,19 +1,23 @@
 """Base Class for Reconstructors."""
 
-from numpy.typing import NDArray
-
 import logging
 from dataclasses import field
 from typing import Any, ClassVar
-from ..._meta import MetaDCRegister
-from snake.mrd_utils import MRDLoader
+
+from numpy.typing import NDArray
+from typing_extensions import dataclass_transform
+
 from snake.core.simulation import SimConfig
+from snake.mrd_utils import MRDLoader
+
+from ..._meta import MetaDCRegister
 
 
+@dataclass_transform(kw_only_default=True)
 class MetaReconstructor(MetaDCRegister):
     """MetaClass Reconstructor."""
 
-    dunder_name = "reconstructor"
+    dunder_name: ClassVar[str] = "reconstructor"
 
 
 class BaseReconstructor(metaclass=MetaReconstructor):
@@ -34,7 +38,7 @@ def setup(self, sim_conf: SimConfig) -> None:
         """Set up the reconstructor."""
         self.log.info(f"Setup reconstructor {self.__class__.__name__}")
 
-    def reconstruct(self, data_loader: MRDLoader, sim_conf: SimConfig) -> NDArray:
+    def reconstruct(self, data_loader: MRDLoader) -> NDArray:
         """Reconstruct the kspace data to image space."""
         raise NotImplementedError