load from mrd + 2d adjoint reconstructor

Author: alineyyy

src/cli-conf/scenario2-2d.yaml

@@ -0,0 +1,93 @@
+# This files contains the configuration to reproduce the scenario 1 of the Snkf paper.
+
+defaults:
+  - base_config
+  - handlers:
+    - activation-block
+  - sampler:
+    - stack-of-spiral
+  - reconstructors:
+      - adjoint
+      #- sequential
+  - _self_
+
+cache_dir: ${oc.env:PWD}/cache
+result_dir: results/scenario2
+filename:  ${cache_dir}/scenario2_2d_${engine.model}_${engine.snr}_${sampler.stack-of-spiral.constant}_${sampler.stack-of-spiral.accelz}.mrd
+
+sim_conf:
+  max_sim_time: 360
+  seq: {TR: 50, TE: 25,  FA: 12}
+  hardware:
+    n_coils: 1
+    dwell_time_ms: 0.001
+  shape: [60, 72, 60]
+  fov_mm: [181.0, 217.0, 181.0]
+
+phantom:
+  name: brainweb
+  sub_id: 4
+  tissue_file: "tissue_7T"
+
+handlers:
+  activation-block:
+    event_name: block_on
+    block_on: 20 # seconds
+    block_off: 20 #seconds
+    duration: 360 # seconds
+    #delta_r2s: 1000 # millisecond^-1
+
+sampler:
+  stack-of-spiral:
+    acsz: 1
+    accelz: 1
+    nb_revolutions: 10
+    constant: true
+    spiral_name: "galilean"
+
+engine:
+  n_jobs: 1
+  chunk_size: 10
+  model: "simple"
+  snr: 10000
+  nufft_backend: "gpuNUFFT"
+  slice_2d: true
+
+reconstructors:
+  adjoint:
+    nufft_backend: "gpuNUFFT"
+    density_compensation: "pipe"
+  # sequential:
+  #   nufft_backend: "gpuNUFFT"
+  #   density_compensation: false
+  #   restart_strategy: WARM
+  #   max_iter_per_frame: 50
+  #   wavelet: "sym4"
+
+
+
+
+
+hydra:
+  job:
+    chdir: true
+
+  run:
+    dir: ${result_dir}/outputs/${hydra.job.name}/${now:%Y-%m-%d_%H-%M-%S}
+  sweep:
+    dir: ${result_dir}/multirun/${hydra.job.name}/${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
+
+    log_job:
+      _target_: hydra.experimental.callbacks.LogJobReturnCallback
+    latest_run:
+      _target_: hydra_callbacks.LatestRunLink
+      run_base_dir:  ${result_dir}/outputs
+      multirun_base_dir:  ${result_dir}/multirun

src/snake/core/engine/base.py

@@ -209,7 +209,17 @@ def __call__(
         `_job_model_simple` methods.
         """
         # Create the base dataset
-        make_base_mrd(filename, sampler, phantom, sim_conf, handlers, smaps, coil_cov, self.model, self.slice_2d)
+        make_base_mrd(
+            filename,
+            sampler,
+            phantom,
+            sim_conf,
+            handlers,
+            smaps,
+            coil_cov,
+            self.model,
+            self.slice_2d,
+        )
 
         # Guesstimate the workload
         if worker_chunk_size <= 0:

src/snake/mrd_utils/loader.py

@@ -211,6 +211,16 @@ def n_shots(self) -> int:
         """
         return self.header.encoding[0].encodingLimits.kspace_encoding_step_1.maximum
 
+    @property
+    def engine_model(self) -> str:
+        """Get the engine model."""
+        return self.header.userParameters.userParameterString[0].value
+
+    @property
+    def slice_2d(self) -> bool:
+        """Is the slice 2D."""
+        return bool(self.header.userParameters.userParameterString[1].value)
+
     #############
     # Get data  #
     #############

src/snake/mrd_utils/writer.py

@@ -22,7 +22,9 @@
 log = logging.getLogger(__name__)
 
 
-def get_mrd_header(sim_conf: SimConfig, engine: str, model: str, slice_2d: bool) -> mrd.xsd.ismrmrdHeader:
+def get_mrd_header(
+    sim_conf: SimConfig, engine: str, model: str, slice_2d: bool
+) -> mrd.xsd.ismrmrdHeader:
     """Create a MRD Header for snake-fmri data."""
     H = mrd.xsd.ismrmrdHeader()
     # Experimental conditions
@@ -85,7 +87,7 @@ def get_mrd_header(sim_conf: SimConfig, engine: str, model: str, slice_2d: bool)
                 ("engine_model", model),
                 ("slice_2d", str(slice_2d)),
             ]
-        ]
+        ],
     )
 
     return H

src/snake/toolkit/cli/__init__.py

@@ -1 +1 @@
-""""Command Line Interface for SNAKE."""
+""" "Command Line Interface for SNAKE."""

src/snake/toolkit/cli/config.py

@@ -125,5 +125,5 @@ def cleanup_cuda() -> None:
 def make_hydra_cli(fun: callable) -> callable:
     """Create a Hydra CLI for the function."""
     return hydra.main(
-        version_base=None, config_path="../../../cli-conf", config_name="config"
+        version_base=None, config_path="../../../cli-conf", config_name="scenario2-2d"
     )(fun)

src/snake/toolkit/cli/reconstruction.py

@@ -57,7 +57,7 @@ def reconstruction(cfg: DictConfig) -> None:
             data_rec_file = Path(f"data_rec_{rec_str}.npy")
             log.info(f"Using {name} reconstructor")
             rec.setup(sim_conf)
-            rec_data = rec.reconstruct(data_loader, sim_conf)
+            rec_data = rec.reconstruct(data_loader, sim_conf, data_loader.slice_2d)
             log.info(f"Reconstruction done with {name}")
             # Save the reconstruction
             np.save(data_rec_file, rec_data)

src/snake/toolkit/reconstructors/pysap.py

@@ -73,18 +73,20 @@ def setup(self, sim_conf: SimConfig) -> None:
         """Initialize Reconstructor."""
         pass
 
-    def reconstruct(self, data_loader: MRDLoader, sim_conf: SimConfig) -> NDArray:
+    def reconstruct(
+        self, data_loader: MRDLoader, sim_conf: SimConfig, slice_2d: bool
+    ) -> NDArray:
         """Reconstruct data with zero-filled method."""
         with data_loader:
             if isinstance(data_loader, CartesianFrameDataLoader):
-                return self._reconstruct_cartesian(data_loader, sim_conf)
+                return self._reconstruct_cartesian(data_loader, sim_conf, slice_2d)
             elif isinstance(data_loader, NonCartesianFrameDataLoader):
-                return self._reconstruct_nufft(data_loader, sim_conf)
+                return self._reconstruct_nufft(data_loader, sim_conf, slice_2d)
             else:
                 raise ValueError("Unknown dataloader")
 
     def _reconstruct_cartesian(
-        self, data_loader: CartesianFrameDataLoader, sim_conf: SimConfig
+        self, data_loader: CartesianFrameDataLoader, sim_conf: SimConfig, slice_2d
     ) -> NDArray:
         smaps = data_loader.get_smaps()
         if smaps is None and data_loader.n_coils > 1:
@@ -114,7 +116,6 @@ def _reconstruct_cartesian(
                 ): idx
                 for idx in range(data_loader.n_frames)
             }
-
             for future in as_completed(futures):
                 future.result()
                 pbar.update(1)
@@ -126,16 +127,21 @@ def _reconstruct_cartesian(
         return final_images
 
     def _reconstruct_nufft(
-        self, data_loader: NonCartesianFrameDataLoader, sim_conf: SimConfig
+        self, data_loader: NonCartesianFrameDataLoader, sim_conf: SimConfig, slice_2d
     ) -> NDArray:
         """Reconstruct data with nufft method."""
         from mrinufft import get_operator
 
         smaps = data_loader.get_smaps()
-
+        shape = data_loader.shape
         traj, kspace_data = data_loader.get_kspace_frame(0)
+
+        if slice_2d:
+            shape = data_loader.shape[:2]
+            traj = traj.reshape(data_loader.n_shots, -1, traj.shape[-1])[0, :, :2]
+
         kwargs = dict(
-            shape=data_loader.shape,
+            shape=shape,
             n_coils=data_loader.n_coils,
             smaps=smaps,
         )
@@ -146,6 +152,7 @@ def _reconstruct_nufft(
             kwargs["density"] = self.density_compensation
         if "stacked" in self.nufft_backend:
             kwargs["z_index"] = "auto"
+
         nufft_operator = get_operator(
             self.nufft_backend,
             samples=traj,
@@ -158,8 +165,16 @@ def _reconstruct_nufft(
 
         for i in tqdm(range(data_loader.n_frames)):
             traj, data = data_loader.get_kspace_frame(i)
-            nufft_operator.samples = traj
-            final_images[i] = abs(nufft_operator.adj_op(data))
+            if slice_2d:
+                nufft_operator.samples = traj.reshape(
+                    data_loader.n_shots, -1, traj.shape[-1]
+                )[0, :, :2]
+                data = np.reshape(data, (data.shape[0], data_loader.n_shots, -1))
+                for j in range(data.shape[1]):
+                    final_images[i, :, :, j] = abs(nufft_operator.adj_op(data[:, j]))
+            else:
+                nufft_operator.samples = traj
+                final_images[i] = abs(nufft_operator.adj_op(data))
         return final_images