Merge remote-tracking branch 'upstream/main'

Author: alineyyy

.github/workflows/test.yml

@@ -42,8 +42,8 @@ jobs:
           pip install -e .[dev]
       - name: Linters
         run: |
-          black --check src tests
-          ruff check .
+          ruff check src
+          ruff format --check src
       - name: Annotate locations with typos
         if: always()
         uses: codespell-project/codespell-problem-matcher@v1

examples/anatomical/example_fov_select.py

@@ -37,14 +37,14 @@
 
 
 # the center and size target_res are all in millimeters.
-fov_handlers = FOVHandler(
+fov_handler = FOVHandler(
     center=(90, 110, 100),
     size=(192, 192, 128),
     angles=(5, 0, 0),
     target_res=(2.0, 2.0, 2.0),
 )
 
-new_phantom = fov_handlers.get_static(phantom, sim_conf)
+new_phantom = fov_handler.get_static(phantom, sim_conf)
 
 # %%
 new_contrast = get_ideal_phantom(new_phantom, sim_conf)

pyproject.toml

@@ -131,17 +131,15 @@ convention = "numpy"
 "tests/test_*.py" = ["D", "ANN"]
 
 
-[tool.pylsp-mypy]
-enabled = true
-live_mode = true
-strict = true
-#overrides = ["--ignore-missing-imports"]
+[tool.basedpyright]
 
-[tool.mypy]
-ignore_missing_imports=true
-exclude = ["examples/"]
-#overrides = ["--ignore-missing-imports"]
+typeCheckingMode = "off"
 
 [tool.codespell]
 ignore-words-list = ["TE","fpr"]
 skip = ["docs/generated", "*.ipynb"]
+
+[dependency-groups]
+dev = [
+    "ipykernel>=6.29.5",
+]

src/snake/_meta.py

@@ -17,6 +17,9 @@
 
 T = TypeVar("T")
 
+ThreeInts = tuple[int, int, int]
+ThreeFloats = tuple[float, float, float]
+
 
 def make_log_property(dunder_name: str) -> Callable:
     """Create a property logger."""
@@ -155,7 +158,6 @@ class ENVCONFIG(metaclass=Singleton):
 
     @classmethod
     def __getitem__(cls, key: str) -> Any:
-
         if key in os.environ:
             return os.environ[key]
         return getattr(cls, key)

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,26 @@
 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):
@@ -88,7 +94,6 @@ def _job_model_T2s(
         dyn_datas: list[DynamicData],
         sim_conf: SimConfig,
         trajectories: NDArray,  # (Chunksize, N, 3)
-        smaps: NDArray,
         *args: Any,
         **kwargs: Any,
     ) -> NDArray:
@@ -100,7 +105,6 @@ def _job_model_simple(
         dyn_datas: list[DynamicData],
         sim_conf: SimConfig,
         trajectories: NDArray,  # (Chunksize, N, 3)
-        smaps: NDArray,
         *args: Any,
         **kwargs: Any,
     ) -> NDArray:
@@ -113,11 +117,11 @@ def _write_chunk_data(
 
     def _acquire_ksp_job(
         self,
-        filename: os.PathLike,
+        filename: AnyPath,
         chunk: Sequence[int],
         tmp_dir: str,
         shared_phantom_props: (
-            tuple[str, ArrayProps, ArrayProps, ArrayProps] | None
+            tuple[str, ArrayProps, ArrayProps, ArrayProps, ArrayProps | None] | None
         ) = None,
         **kwargs: Mapping[str, Any],
     ) -> str:
@@ -144,28 +148,24 @@ def _acquire_ksp_job(
             trajs = self._job_trajectories(data_loader, hdr, sim_conf, chunk)
 
             _job_model = getattr(self, f"_job_model_{self.model}")
-            smaps = None
-            if sim_conf.hardware.n_coils > 1:
-                smaps = data_loader.get_smaps()
             if shared_phantom_props is None:
                 phantom = data_loader.get_phantom()
-                ksp = _job_model(phantom, ddatas, sim_conf, trajs, smaps, **kwargs)
+                ksp = _job_model(phantom, ddatas, sim_conf, trajs, **kwargs)
             else:
                 with Phantom.from_shared_memory(*shared_phantom_props) as phantom:
-                    ksp = _job_model(phantom, ddatas, sim_conf, trajs, smaps, **kwargs)
+                    ksp = _job_model(phantom, ddatas, sim_conf, trajs, **kwargs)
 
         chunk_file = os.path.join(tmp_dir, f"partial_{chunk[0]}-{chunk[-1]}.npy")
         np.save(chunk_file, ksp)
         return chunk_file
 
     def __call__(
         self,
-        filename: os.PathLike,
+        filename: AnyPath,
         sampler: BaseSampler,
         phantom: Phantom,
         sim_conf: SimConfig,
         handlers: list[AbstractHandler] | HandlerList | None = None,
-        smaps: NDArray | None = None,
         coil_cov: NDArray | None = None,
         worker_chunk_size: int = 0,
         n_workers: int = 0,
@@ -175,7 +175,7 @@ def __call__(
 
         Parameters
         ----------
-        filename : os.PathLike
+        filename : AnyPath
             The path to the MRD file.
         sampler : BaseSampler
             The sampler to use.
@@ -185,8 +185,6 @@ def __call__(
             The simulation configuration.
         handlers : list[AbstractHandler] | HandlerList | None, optional
             The handlers to use, by default None.
-        smaps : NDArray | None, optional
-            The sensitivity maps, by default None.
         coil_cov : NDArray | None, optional
             The coil covariance matrix, by default None.
         worker_chunk_size : int, optional
@@ -215,7 +213,6 @@ def __call__(
             phantom,
             sim_conf,
             handlers,
-            smaps,
             coil_cov,
             self.model,
             self.slice_2d,

src/snake/core/engine/cartesian.py

@@ -63,7 +63,6 @@ def _job_model_T2s(
         dyn_datas: list[DynamicData],
         sim_conf: SimConfig,
         trajectories: NDArray,  # (Chunksize, N, 3)
-        smaps: NDArray,
         slice_2d: bool = False,
     ) -> np.ndarray:
         """Acquire k-space data. With T2s decay."""
@@ -107,18 +106,20 @@ def _job_model_T2s(
                 slice_location = flat_epi[0, 0]  # FIXME: the slice is always axial.
                 flat_epi = flat_epi[:, 1:]
                 phantom_slice = phantom_state[:, slice_location]
-                if smaps is None:
+                if phantom.smaps is None:
                     phantom_slice = phantom_slice[:, None, ...]
                 else:
-                    smaps_ = smaps[:, slice_location]
+                    smaps_ = phantom.smaps[:, slice_location]
                     phantom_slice = phantom_slice[:, None, ...] * smaps_
 
                 ksp = fft(phantom_slice, axis=(-2, -1))
             else:
-                if smaps is None:
+                if phantom.smaps is None:
                     ksp = fft(phantom_state[:, None, ...], axis=(-3, -2, -1))
                 else:
-                    ksp = fft(phantom_state[:, None, ...] * smaps, axis=(-3, -2, -1))
+                    ksp = fft(
+                        phantom_state[:, None, ...] * phantom.smaps, axis=(-3, -2, -1)
+                    )
 
             for c in range(sim_conf.hardware.n_coils):
                 ksp_coil_sum = np.zeros(
@@ -135,7 +136,6 @@ def _job_model_simple(
         dyn_datas: list[DynamicData],
         sim_conf: SimConfig,
         trajectories: NDArray,  # (Chunksize, N, 3)
-        smaps: NDArray,
         slice_2d: bool = False,
     ) -> np.ndarray:
         """Acquire k-space data. No T2s decay."""
@@ -157,17 +157,19 @@ def _job_model_simple(
                 slice_location = flat_epi[0, 0]  # FIXME: the slice is always axial.
                 flat_epi = flat_epi[:, 1:]  # Reduced to 2D.
                 phantom_slice = phantom_state[slice_location]
-                if smaps is None:
+                if phantom.smaps is None:
                     phantom_slice = phantom_slice[None, ...]
                 else:
-                    smaps_ = smaps[:, slice_location]
+                    smaps_ = phantom.smaps[:, slice_location]
                     phantom_slice = phantom_slice[None, ...] * smaps_
                 ksp = fft(phantom_slice, axis=(-2, -1))
             else:
-                if smaps is None:
+                if phantom.smaps is None:
                     ksp = fft(phantom_state[None, ...], axis=(-3, -2, -1))
                 else:
-                    ksp = fft(phantom_state[None, ...] * smaps, axis=(-3, -2, -1))
+                    ksp = fft(
+                        phantom_state[None, ...] * phantom.smaps, axis=(-3, -2, -1)
+                    )
             for c in range(sim_conf.hardware.n_coils):
                 ksp_coil = ksp[c]
                 a = ksp_coil[tuple(flat_epi.T)]
@@ -247,7 +249,6 @@ def _job_model_T2s(
         dyn_datas: list[DynamicData],
         sim_conf: SimConfig,
         trajectories: NDArray,  # (Chunksize, N, 3)
-        smaps: NDArray,
     ) -> np.ndarray:
         """Acquire k-space data. With T2s decay."""
         readout_length = trajectories.shape[-2]
@@ -299,10 +300,12 @@ def _job_model_T2s(
                 * frame_phantom.masks
             )
 
-            if smaps is None:
+            if phantom.smaps is None:
                 ksp = fft(phantom_state[:, None, ...], axis=(-3, -2, -1))
             else:
-                ksp = fft(phantom_state[:, None, ...] * smaps, axis=(-3, -2, -1))
+                ksp = fft(
+                    phantom_state[:, None, ...] * phantom.smaps, axis=(-3, -2, -1)
+                )
             flat_evi = evi.reshape(-1, 3)
             for c in range(sim_conf.hardware.n_coils):
                 ksp_coil_sum = np.zeros(
@@ -322,7 +325,6 @@ def _job_model_simple(
         dyn_datas: list[DynamicData],
         sim_conf: SimConfig,
         trajectories: NDArray,  # (Chunksize, N, 3)
-        smaps: NDArray,
     ) -> np.ndarray:
         """Acquire k-space data. No T2s decay."""
         final_ksp = np.zeros(
@@ -351,10 +353,10 @@ def _job_model_simple(
                 * frame_phantom.masks,
                 axis=0,
             )
-            if smaps is None:
+            if phantom.smaps is None:
                 ksp = fft(phantom_state[None, ...], axis=(-3, -2, -1))
             else:
-                ksp = fft(phantom_state[None, ...] * smaps, axis=(-3, -2, -1))
+                ksp = fft(phantom_state[None, ...] * phantom.smaps, axis=(-3, -2, -1))
             flat_epi = epi_2d.reshape(-1, 3)
             for c in range(sim_conf.hardware.n_coils):
                 ksp_coil = ksp[c]

src/snake/core/engine/nufft.py

@@ -28,7 +28,7 @@ def _job_trajectories(
         dataset: mrd.Dataset,
         hdr: mrd.xsd.ismrmrdHeader,
         sim_conf: SimConfig,
-        shot_idx: Sequence[int],
+        shot_idx: Sequence[int] | int,
     ) -> NDArray:
         """Get Non Cartesian trajectories from the dataset.
 
@@ -72,10 +72,12 @@ def _init_model_nufft(
         if slice_2d:
             shape_ = sim_conf.shape[:-1]
             if smaps is not None:
-                smaps_ = smaps[..., 0]
+                smaps_ = np.ascontiguousarray(
+                    smaps[..., 0]
+                )  # will be updated in the loop
 
         nufft = get_operator(backend)(
-            samples,  # dummy samples locs
+            samples,  # will be updated in the loop
             shape=shape_,
             n_coils=n_coils,
             smaps=smaps_,
@@ -91,7 +93,6 @@ def _job_model_T2s(
         dyn_datas: list[DynamicData],
         sim_conf: SimConfig,
         trajectories: NDArray,
-        smaps: NDArray,
         nufft_backend: str,
         slice_2d: bool = False,
     ) -> np.ndarray:
@@ -105,7 +106,7 @@ def _job_model_T2s(
         nufft = NufftAcquisitionEngine._init_model_nufft(
             trajectories[0],
             sim_conf,
-            smaps,
+            phantom.smaps,
             backend=nufft_backend,
             slice_2d=slice_2d,
         )
@@ -121,8 +122,8 @@ def _job_model_T2s(
             if slice_2d:
                 slice_loc = round((traj[0, -1] + 0.5) * sim_conf.shape[-1])
                 nufft.samples = traj[:, :2]
-                if smaps is not None:
-                    nufft.smaps = smaps[..., slice_loc]
+                if phantom.smaps is not None:
+                    nufft.smaps = np.ascontiguousarray(phantom.smaps[..., slice_loc])
                 phantom_state = phantom_state[:, None, ..., slice_loc]
             else:
                 phantom_state = phantom_state[:, None, ...]
@@ -140,7 +141,6 @@ def _job_model_simple(
         dyn_datas: list[DynamicData],
         sim_conf: SimConfig,
         trajectories: NDArray,
-        smaps: NDArray,
         nufft_backend: str,
         slice_2d: bool = False,
     ) -> np.ndarray:
@@ -153,7 +153,7 @@ def _job_model_simple(
         nufft = NufftAcquisitionEngine._init_model_nufft(
             trajectories[0],
             sim_conf,
-            smaps,
+            phantom.smaps,
             backend=nufft_backend,
             slice_2d=slice_2d,
         )
@@ -165,8 +165,8 @@ def _job_model_simple(
             if slice_2d:
                 slice_loc = int((traj[0, -1] + 0.5) * sim_conf.shape[-1])
                 nufft.samples = traj[:, :2]
-                if smaps is not None:
-                    nufft.smaps = smaps[..., slice_loc]
+                if phantom.smaps is not None:
+                    nufft.smaps = phantom.smaps[..., slice_loc]
                 phantom_state = phantom_state[None, ..., slice_loc]
             else:
                 nufft.samples = traj