feat!: smaps is now part of phantom.

This is a breaking change, but for the best: the smaps can now be
modified by handlers and are also shared in memory.

Author: paquiteau

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."""

src/snake/core/engine/base.py

@@ -30,7 +30,8 @@
 from ..simulation import SimConfig
 from .utils import get_ideal_phantom, get_noise
 
-AnyPath =  str | Path
+AnyPath = str | Path
+
 
 @dataclass_transform(kw_only_default=True)
 class MetaEngine(MetaDCRegister):
@@ -93,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:
@@ -105,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:
@@ -122,7 +121,7 @@ def _acquire_ksp_job(
         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:
@@ -149,15 +148,12 @@ 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)
@@ -170,7 +166,6 @@ def __call__(
         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,
@@ -190,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
@@ -220,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

@@ -91,7 +91,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 +104,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 +120,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 = phantom.smaps[..., slice_loc]
                 phantom_state = phantom_state[:, None, ..., slice_loc]
             else:
                 phantom_state = phantom_state[:, None, ...]
@@ -140,7 +139,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 +151,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 +163,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

src/snake/core/handlers/fov.py

@@ -16,13 +16,10 @@
 from snake.core.parallel import run_parallel
 from snake.core.phantom import Phantom
 from snake.core.simulation import SimConfig
-
+from snake._meta import ThreeInts, ThreeFloats
 
 # TODO allow to use cupy for faster computation (if available)
 
-ThreeInts = tuple[int, int, int]
-ThreeFloats = tuple[float, float, float]
-
 
 def extract_rotated_3d_region(
     volume: NDArray,
@@ -134,7 +131,13 @@ def get_static(self, phantom: Phantom, sim_conf: SimConfig) -> Phantom:
             ),
             dtype=phantom.masks.dtype,
         )
-        print("=======", size_vox, zoom_factor)
+        new_smaps = np.zeros(
+            (
+                phantom.smaps.shape[0],
+                *tuple(round(size_vox[i] / zoom_factor[i]) for i in range(3)),
+            ),
+            dtype=phantom.smaps.dtype,
+        )
 
         run_parallel(
             _apply_transform,
@@ -146,9 +149,22 @@ def get_static(self, phantom: Phantom, sim_conf: SimConfig) -> Phantom:
             angles=self.angles,
             zoom_factor=zoom_factor,
         )
+
+        run_parallel(
+            _apply_transform,
+            phantom.smaps,
+            new_smaps,
+            parallel_axis=0,
+            center=center_vox,
+            size=size_vox,
+            angles=self.angles,
+            zoom_factor=zoom_factor,
+        )
+
         # Create a new phantom with updated masks
         new_phantom = phantom.copy()
         new_phantom.masks = new_masks
+        new_phantom.smaps = new_smaps
 
         # update the sim_config
         new_shape = new_phantom.anat_shape

src/snake/core/parallel.py

@@ -101,6 +101,9 @@ def array_from_shm(
     shms = []
     arrays: list[NDArray] = []
     for prop in array_props:
+        if prop is None:  # optional arrays are ignored
+            arrays.append(None)
+            continue
         nbytes = int(np.dtype(prop.dtype).itemsize * np.prod(prop.shape))
         shms.append(SharedMemory(name=prop.name, size=nbytes))
         arrays.append(