Merge branch 'motion-pac'

src/conf/scenario2-image-motion.yaml

@@ -0,0 +1,115 @@
+# This files contains the configuration to reproduce the scenario 2 of the Simfmri paper.
+
+defaults:
+  - handlers:
+    - phantom-brainweb
+    - activation-block
+    - motion-image
+    - noise-gaussian
+    - acquisition-sos
+  - reconstructors:
+    - adjoint
+    - sequential
+  - _self_
+
+force_sim: false
+cache_dir: ${oc.env:PWD}/cache
+result_dir: ${oc.env:PWD}/results
+ignore_patterns:
+  - "n_jobs"
+
+sim_params:
+  sim_tr: 0.1
+  sim_time: 300
+  shape: [-1,-1,-1] # shape will be determined with the bbox and the brainweb phantom.
+  fov: [-1,-1,-1]
+  n_coils: 8
+  rng: 19980408
+  lazy: true
+
+handlers:
+  phantom-brainweb:
+    sub_id: 5
+    bbox: [0.225,-0.07, 0.06, -0.055, null, null]
+    brainweb_folder: ${cache_dir}/brainweb
+    res: [3.0, 3.0, 3.0]
+  activation-block:
+    event_name: block_on
+    block_on: 20
+    block_off: 20
+    duration: 300
+    bold_strength: 0.02
+
+  motion-image:
+    ts_std_mm: [1.0, 1.0, 1.0]
+    rs_std_mm: [0.1, 0.1, 0.1]
+
+  noise-gaussian:
+    snr: 30
+  acquisition-sos:
+    shot_time_ms: 50
+    acsz: 0.125
+    accelz: 4
+    n_samples: 25000
+    constant: false
+    orderz: CENTER_OUT
+    backend: stacked-cufinufft
+    n_jobs: 2
+
+reconstructors:
+  adjoint:
+    nufft_kwargs:
+      backend_name: stacked-cufinufft
+      density: voronoi
+      z_index: auto
+  sequential:
+    nufft_kwargs:
+      backend_name: stacked-cufinufft
+      density: voronoi
+      z_index: auto
+
+  # lr_f:
+  #   nufft_kwargs:
+  #     backend_name: stacked-cufinufft
+  #     density: voronoi
+  #     z_index: auto
+  #   max_iter: 40
+  #   lambda_l: 0.01
+  #   lambda_s: sure
+  #   algorithm: otazo
+
+stats:
+  contrast_name: block_on
+
+
+hydra:
+  job:
+    chdir: true
+
+  run:
+    dir: ${result_dir}/outputs/scenario2-motion/${now:%Y-%m-%d_%H-%M-%S}
+  sweep:
+    dir: ${result_dir}/multirun/scenario2-motion/${now:%Y-%m-%d_%H-%M-%S}
+    subdir: ${hydra.job.num}
+
+  callbacks:
+    gather_files:
+      _target_: hydra_callbacks.MultiRunGatherer
+      aggregator:
+        _partial_: true
+        _target_: snkf.cli.utils.aggregate_results
+
+    latest_run:
+      _target_: hydra_callbacks.LatestRunLink
+      run_base_dir:  ${result_dir}/outputs
+      multirun_base_dir:  ${result_dir}/multirun
+
+    register_run:
+      _target_: hydra_callbacks.RegisterRunCallback
+      run_base_dir:  ${result_dir}/outputs/scenario2-motion
+      multirun_base_dir:  ${result_dir}/multirun/scenario2-motion
+
+    resource_monitor:
+      _target_: hydra_callbacks.ResourceMonitor
+      enabled: true
+      sample_interval: 1

src/snkf/cli/main.py

@@ -119,10 +119,10 @@ def main_app(cfg: ConfigSnakeFMRI) -> None:
         del zscore
         gc.collect()
 
-    # 3. Clean up and saving
-    with PerfLogger(logger, name="Cleanup"):
-        with open("results.json", "w") as f:
-            json.dump(results, f, default=lambda o: str(o))
+        # 3. Clean up and saving
+        with PerfLogger(logger, name="Cleanup"):
+            with open("results.json", "w") as f:
+                json.dump(results, f, default=lambda o: str(o))
 
     PerfLogger.recap(logger)
 

src/snkf/handlers/motion/__init__.py

@@ -0,0 +1,3 @@
+"""Motion Module."""
+
+from .image import RandomMotionImageHandler

src/snkf/handlers/motion/image.py

@@ -0,0 +1,87 @@
+"""Add motion in the image domain."""
+
+import numpy as np
+from numpy.typing import NDArray
+
+from snkf.base import validate_rng
+
+from ...simulation import SimData
+from ..base import AbstractHandler, requires_field
+from .utils import apply_shift, apply_rotation_at_center, motion_generator
+
+requires_data_acq = requires_field("data_acq", lambda sim: sim.data_ref.copy())
+
+
+@requires_data_acq
+class RandomMotionImageHandler(AbstractHandler):
+    """Add Random Motion in Image.
+
+    Parameters
+    ----------
+    ts_std_mm
+        Translation standard deviation, in mm/s.
+    rs_std_mm
+        Rotation standard deviation, in radians/s.
+
+    Notes
+    -----
+    The motion is generated by drawing from a normal distribution with standard
+    deviation for the 6 motion parameters (3 translations and 3 rotations, in
+    this order). Then the cumulative motion is computed by summing the motion
+    at each frame.
+
+    The handlers is parametrized with speed in mm/s and rad/s, as these values
+    provides an independent control of the motion amplitude regardless of the
+    time resolution for the simulation.
+    """
+
+    __handler_name__ = "motion-image"
+
+    ts_std_mm: tuple[float, float, float]
+    rs_std_mm: tuple[float, float, float]
+
+    def _handle(self, sim: SimData) -> SimData:
+        n_frames = sim.data_acq.shape[0]
+        # update the translation to be in voxel units
+        ts_std_pix = np.array(self.ts_std_mm) / np.array(sim.res_mm)
+        motion = motion_generator(
+            n_frames,
+            ts_std_pix,
+            self.rs_std_mm,
+            sim.sim_tr,
+            sim.rng,
+        )
+        sim.extra_infos["motion_params"] = motion
+
+        if sim.lazy:
+            sim.data_acq.apply(add_motion_to_frame, motion)
+        else:
+            for i in range(len(sim.data_acq)):
+                sim.data_acq[i] = add_motion_to_frame(sim.data_acq[i], motion, i)
+        return sim
+
+
+def add_motion_to_frame(
+    data: NDArray[np.complexfloating] | NDArray[np.floating],
+    motion: NDArray[np.floating],
+    frame_idx: int = 0,
+) -> np.ndarray:
+    """Add motion to a base array.
+
+    Parameters
+    ----------
+    data: np.ndarray
+        The data to which motion is added.
+    motion: np.ndarray
+        The N_frames x 6 motion trajectory.
+    frame_idx: int
+        The frame index used to compute the motion at that frame.
+
+    Returns
+    -------
+    np.ndarray
+        The data with motion added.
+    """
+    rotated = apply_rotation_at_center(data, motion[frame_idx, 3:])
+    rotated_and_translated = apply_shift(rotated, motion[frame_idx, :3])
+    return rotated_and_translated

src/snkf/handlers/motion/kspace.py

@@ -0,0 +1,110 @@
+"""Add Motion in the k-space.
+
+TODO: 2D Support
+TODO: Actually modify the k-space data
+
+"""
+
+import numpy as np
+from np.typing import NDArray
+
+from ...simulation import SimData
+from ..base import AbstractHandler, requires_field
+from .utils import motion_generator, rotation3D
+
+
+@requires_field(["kspace_data", "kspace_mask"])
+class RandomMotionKspaceHandler(AbstractHandler):
+    """Add Random Motion in the K-space.
+
+    Parameters
+    ----------
+    ts_std_mm
+        Translation standard deviation, in mm/s.
+    rs_std_mm
+        Rotation standard deviation, in radians/s.
+
+    Notes
+    -----
+    The motion is generated by drawing from a normal distribution with standard
+    deviation for the 6 motion parameters (3 translations and 3 rotations, in
+    this order). Then the cumulative motion is computed by summing the motion
+    at each frame.
+
+    The handlers is parametrized with speed in mm/s and rad/s, as these values
+    provides an independent control of the motion amplitude regardless of the
+    time resolution for the simulation.
+    """
+
+    __handler_name__ = "motion-image"
+
+    ts_std_mm: tuple[float, float, float]
+    rs_std_mm: tuple[float, float, float]
+
+    def _handle(self, sim: SimData) -> SimData:
+        n_frames = sim.data_acq.shape[0]
+        # update the translation to be in voxel units
+        ts_std_pix = np.array(self.ts_std_mm) / np.array(sim.res_mm)
+        motion = motion_generator(
+            n_frames,
+            ts_std_pix,
+            self.rs_std_mm,
+            sim.sim_tr,
+            sim.rng,
+        )
+        sim.extra_infos["motion_params"] = motion
+
+        ...
+
+
+def translate(
+    kspace_data: NDArray,
+    kspace_loc: NDArray,
+    translation_Matrix: NDArray,
+) -> NDArray:
+    """Translate the image in the kspace data.
+
+    Parameters
+    ----------
+    kspace_data :(5, length,N)
+        the data in the Fourier domain
+    kspace_loc : (length,N,3)
+        the kspace locations
+    translation_Matrix : (length,3)-array
+        defines the translation vector in 3d space
+
+    Returns
+    -------
+    A new kspace_data: (5,N)array
+    """
+    phi = np.zeros_like(kspace_data)
+
+    for t in range(translation_Matrix.shape[0]):
+        for i in range(kspace_loc.shape[2]):
+            phi[:, t, :] += kspace_loc[t, :, i] * translation_Matrix[t, i]
+    phi = np.reshape(phi, (kspace_data.shape[0], phi.shape[1] * phi.shape[2]))
+    return kspace_data * np.exp(-2j * np.pi * phi)
+
+
+def rotate(kspace_loc_to_corrupt: NDArray, rotation: NDArray) -> NDArray:
+    """Rotate the image in the kspace.
+
+    kspace_data :(5, length,N)
+        the data in the Fourier domain
+    kspace_loc : (length,N,3)
+        the kspace locations
+    translation : (3,1)-array
+        defines the translation vector in 3d space
+    td : int
+        beginning shot of motion
+    tf : int
+        final shot of the motion
+    center: (3,1) array
+        the center of the rotation
+    """
+    new_loc = np.zeros_like(kspace_loc_to_corrupt)
+    for t in range(kspace_loc_to_corrupt.shape[0]):
+        R = rotation3D(rotation[:, t])
+        new_loc[t, :, :] = np.matmul(kspace_loc_to_corrupt[t, :, :], R)
+
+    return new_loc