draft of 2d SliceDataloader

Author: alineyyy

src/snake/mrd_utils/__init__.py

@@ -4,6 +4,7 @@
     CartesianFrameDataLoader,
     MRDLoader,
     NonCartesianFrameDataLoader,
+    SliceDataloader,
     parse_sim_conf,
     parse_waveform_information,
     read_mrd_header,
@@ -16,6 +17,7 @@
     "MRDLoader",
     "CartesianFrameDataLoader",
     "NonCartesianFrameDataLoader",
+    "SliceDataloader",
     "parse_sim_conf",
     "parse_waveform_information",
     "make_base_mrd",

src/snake/mrd_utils/loader.py

@@ -366,6 +366,7 @@ class NonCartesianFrameDataLoader(MRDLoader):
     ...     image = nufft.adj_op(kspace)
     """
 
+
     def get_kspace_frame(
         self, idx: int, shot_dim: bool = False
     ) -> tuple[NDArray[np.float32], NDArray[np.complex64]]:
@@ -402,7 +403,48 @@ def get_kspace_frame(
             )
         return traj, data
 
+class SliceDataloader(MRDLoader):
+    """Load slice MRD files k-space frames iteratively."""
+
+    def __init__(self,
+                 frame_dl: MRDLoader,
+                 ): 
+        super().__init__(
+            filename=frame_dl._filename,
+            dataset_name=frame_dl._dataset_name,
+            writeable=frame_dl._writeable,
+            swmr=frame_dl._swmr,
+        )
+        self.__class__ = type("SliceDataloader", (type(frame_dl), SliceDataloader), {}) 
+        #replace by adding an attribute of if_cartesian, also need to change on pysap
+        self.frame = frame_dl
+        self.is_cartesian = isinstance(frame_dl, CartesianFrameDataLoader)
+        self.get_kspace_frame = self._get_kspace_frame
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self.frame, name)
+
+    @property
+    def n_frames(self) -> int:
+        """Number of frames."""
+        return self.frame.n_acquisition 
+
+    @property
+    def shape(self) -> tuple[int, int]:
+        """Shape of the volume."""
+        return self.frame.shape[:2]
+    
+    def _get_kspace_frame(
+        self, idx: int, 
+         ) -> tuple[NDArray[np.float32], NDArray[np.complex64]]:
+        """Get the k-space frame."""
+        n_acq_per_frame = self.frame.n_acquisition // self.frame.n_frames
+        traj, data = self.frame.get_kspace_frame(idx // self.frame.n_shots)
+        traj =  traj.reshape(n_acq_per_frame, -1, 3)
+        data = data.reshape(self.frame.n_coils, n_acq_per_frame, -1)
 
+        return traj[idx%self.frame.n_shots,:,:2], data[:,idx%self.frame.n_shots,:]
+    
 def parse_sim_conf(header: mrd.xsd.ismrmrdHeader) -> SimConfig:
     """Parse the header to populate SimConfig from an MRD Header."""
     from ..core import GreConfig, HardwareConfig, SimConfig

src/snake/toolkit/cli/reconstruction.py

@@ -49,9 +49,13 @@ def reconstruction(cfg: DictConfig) -> None:
     elif engine == "NUFFT":
         DataLoader = NonCartesianFrameDataLoader
 
+    
     # Reconstructor.setup(sim_conf) # initialize operators
     # array = Reconstructor.reconstruct(dataloader, sim_conf)
     with DataLoader(cfg.filename) as data_loader:
+        phantom = data_loader.get_phantom()
+        dyn_datas = data_loader.get_all_dynamic()
+        #data_loader = SliceDataloader(data_loader) if cfg.engine.slice_2d else data_loader
         for name, rec in cfg.reconstructors.items():
             rec_str = str(rec)  # FIXME Also use parameters  of reconstructors
             data_rec_file = Path(f"data_rec_{rec_str}.npy")
@@ -63,9 +67,7 @@ def reconstruction(cfg: DictConfig) -> None:
             np.save(data_rec_file, rec_data)
             log.info(f"Saved to {data_rec_file.resolve()}")
 
-        phantom = data_loader.get_phantom()
         roi_mask = phantom.masks[phantom.labels == cfg.stats.roi_tissue_name]
-        dyn_datas = data_loader.get_all_dynamic()
         waveform_name = f"activation-{cfg.stats.event_name}"
         good_d = None
         for d in dyn_datas:

src/snake/toolkit/reconstructors/fourier.py

@@ -2,7 +2,7 @@
 
 from numpy.typing import NDArray
 import scipy as sp
-from snake.mrd_utils.loader import CartesianFrameDataLoader, NonCartesianFrameDataLoader
+from snake.mrd_utils.loader import CartesianFrameDataLoader, NonCartesianFrameDataLoader, MRDLoader
 
 
 def fft(image: NDArray, axis: int | tuple[int] = -1) -> NDArray:
@@ -53,15 +53,13 @@ def init_nufft(
     density_compensation: bool = False,
 ):
     from mrinufft import get_operator
-
-    smaps = data_loader.get_smaps()
+    if data_loader.n_coils > 1:
+        smaps = data_loader.get_smaps().squeeze()
+    else:
+        smaps = None
     shape = data_loader.shape
     traj, _ = data_loader.get_kspace_frame(0)
 
-    if data_loader.slice_2d:
-        shape = data_loader.shape[:2]
-        traj = traj.reshape(data_loader.n_shots, -1, traj.shape[-1])[0, :, :2]
-
     kwargs = dict(
         shape=shape,
         n_coils=data_loader.n_coils,

src/snake/toolkit/reconstructors/pysap.py

@@ -27,7 +27,7 @@
 
 from .base import BaseReconstructor
 from .fourier import ifft, init_nufft
-
+from mrinufft.density.utils import get_density
 
 def _reconstruct_cartesian_frame(
     filename: os.PathLike,
@@ -42,7 +42,7 @@ def _reconstruct_cartesian_frame(
     ):
         mask, kspace = data_loader.get_kspace_frame(idx)
         if data_loader.slice_2d:
-            axes = (-2,-1)
+            axes = (-2, -1)
         else:
             axes = tuple(range(len(data_loader.shape), 0, -1))
         adj_data = ifft(kspace, axis=axes)
@@ -140,22 +140,37 @@ def _reconstruct_nufft(
         nufft_operator = init_nufft(
             data_loader, self.nufft_backend, self.density_compensation
         )
+    
+        method = self.density_compensation
+
+        if isinstance(method, str):
+            method = get_density(method)
+        if not callable(method):
+            raise ValueError(f"Unknown density method: {method}")  
+        
         final_images = np.empty(
             (data_loader.n_frames, *data_loader.shape), dtype=np.float32
         )
+        smaps = data_loader.get_smaps()
 
         for i in tqdm(range(data_loader.n_frames)):
             traj, data = data_loader.get_kspace_frame(i)
-            if data_loader.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))
+
+
+            # if self.density_compensation is True : #fix: it should update every frame when rotating
+            #     nufft_operator.density = method(
+            #     traj, data_loader.shape, backend=self.nufft_backend
+            #     )
+           
+            if smaps is not None:
+                nufft_operator.smaps = smaps[...,i % data_loader.frame.n_shots]
+
+            nufft_operator.samples = traj
+            final_images[i] = abs(nufft_operator.adj_op(data))
+        if data_loader.slice_2d:
+            final_images = np.moveaxis(final_images.reshape(
+                (data_loader.frame.n_frames, -1, *final_images.shape[-2:])
+            ), 1, -1)
         return final_images
 
 
@@ -212,7 +227,7 @@ def setup(self, sim_conf: SimConfig = None, shape: tuple[int] = None) -> None:
 
         self.space_linear_op = WaveletTransform(
             self.wavelet,
-            shape=shape,  ##### shape[:2] if data_loader.slice_2d
+            shape=shape,  
             level=3,
             mode="zero",
             compute_backend=self.compute_backend,
@@ -233,11 +248,10 @@ def setup(self, sim_conf: SimConfig = None, shape: tuple[int] = None) -> None:
                 linear=Identity(), weights=self.threshold
             )
 
-    def reconstruct(self, data_loader: MRDLoader) -> np.ndarray:
+    def reconstruct(self, data_loader: NonCartesianFrameDataLoader) -> np.ndarray:
         """Reconstruct with Sequential."""
         shape = data_loader.shape
-        if data_loader.slice_2d:
-            shape = shape[:2]
+
         self.setup(shape=shape)
         from fmri.operators.gradient import (
             GradAnalysis,
@@ -250,13 +264,9 @@ def reconstruct(self, data_loader: MRDLoader) -> np.ndarray:
         xp, _ = get_backend(self.compute_backend)
 
         traj, data = data_loader.get_kspace_frame(0)
-        if data_loader.slice_2d:
-            traj = np.reshape(
-                traj,(data_loader.n_shots, -1, traj.shape[-1])
-                )[0, :, :2]
-            data = np.reshape(data, (data_loader.n_shots, -1))
-            
-        smaps = data_loader.get_smaps()
+
+
+        smaps = data_loader.get_smaps() 
 
         density_compensation = self.density_compensation
         if (
@@ -273,10 +283,10 @@ def reconstruct(self, data_loader: MRDLoader) -> np.ndarray:
 
         fourier_op = get_operator(
             self.nufft_backend,
-            samples=traj,  
-            shape=shape,##### shape[:2] if 2D
+            samples=traj,
+            shape=shape,  ##### shape[:2] if 2D
             n_coils=data_loader.n_coils,
-            smaps=smaps,
+            smaps=smaps.squeeze(),
             # smaps=xp.array(smaps) if smaps is not None else None,
             density=density_compensation,
             **kwargs,
@@ -307,32 +317,33 @@ def reconstruct(self, data_loader: MRDLoader) -> np.ndarray:
             x_init = fourier_op.adj_op(xp.array(data * density_comp_vector, copy=False))
         else:
             if data_loader.slice_2d:
-               for i in range (data.shape[0]):
-                    x_init[...,i] = fourier_op.adj_op(
-                        xp.array(data[i][None,...], copy=False)
-                        )
+                data = np.reshape(data, (data_loader.n_shots, data_loader.n_coils, -1))
+                for i in range(data.shape[0]):
+                    x_init[..., i] = fourier_op.adj_op(
+                        xp.array(data[i][None, ...], copy=False)
+                    )
 
         pbar_frames = tqdm(total=data_loader.n_frames, position=0)
         pbar_iter = tqdm(total=self.max_iter_per_frame, position=1)
-        x_iter = x_init.copy() 
-        for i, traj, data in data_loader.iter_frames():
-            #grad_op.fourier_op.samples = traj
+        x_iter = x_init.copy()
+        for i, traj, data in data_loader.iter_frames(): #iter_slice
+            # grad_op.fourier_op.samples = traj
             spec_rad = grad_op.fourier_op.get_lipschitz_cst(20)
-            #grad_op._obs_data = xp.array(data)
+            # grad_op._obs_data = xp.array(data)
             grad_op.spec_rad = spec_rad
             grad_op.inv_spec_rad = 1 / spec_rad
             # pseudo code
-            if data_loader.slice_2d: 
-                traj = np.reshape(
-                        traj,(data_loader.n_shots, -1, traj.shape[-1])
-                         )[0, :, :2]
-                data = np.reshape(data, ( data_loader.n_shots, -1))
+            if data_loader.slice_2d:
+                traj = np.reshape(traj, (data_loader.n_shots, -1, traj.shape[-1]))[
+                    0, :, :2
+                ]
+                data = np.reshape(data, (data_loader.n_shots, data_loader.n_coils, -1))
                 grad_op.fourier_op.samples = traj
                 for j in range(data.shape[0]):
-                    grad_op._obs_data = xp.array(data[j][None,...])
-                    x_iter[...,j] = self._reconstruct_frame(
-                        grad_op, x_init[...,j], n_iter=self.max_iter_per_frame
-                        )   
+                    grad_op._obs_data = xp.array(data[j][None, ...])
+                    x_iter[..., j] = self._reconstruct_frame(
+                        grad_op, x_init[..., j], n_iter=self.max_iter_per_frame
+                    )
             # Prepare for next iteration and save results
             x_init = (
                 x_iter.copy()
@@ -349,26 +360,26 @@ def reconstruct(self, data_loader: MRDLoader) -> np.ndarray:
             return final_estimate
         # else, we do a second pass on the data using the last iteration as a solution.
         pbar_frames.reset()
-        #pbar_iter.reset()
+        # pbar_iter.reset()
         x_init = x_iter.copy()  # last iteration results.
         for i, traj, data in data_loader.iter_frames():
-            #grad_op.fourier_op.samples = traj
+            # grad_op.fourier_op.samples = traj
             spec_rad = grad_op.fourier_op.get_lipschitz_cst()
-            #grad_op._obs_data = xp.array(data)
+            # grad_op._obs_data = xp.array(data)
             grad_op.spec_rad = spec_rad
             grad_op.inv_spec_rad = 1 / spec_rad
             # pseudo code
-            if data_loader.slice_2d:  
-                traj = np.reshape(
-                        traj,(data_loader.n_shots, -1, traj.shape[-1])
-                         )[0, :, :2]
-                data = np.reshape(data, (data_loader.n_shots, -1))
+            if data_loader.slice_2d:
+                traj = np.reshape(traj, (data_loader.n_shots, -1, traj.shape[-1]))[
+                    0, :, :2
+                ]
+                data = np.reshape(data, (data_loader.n_shots, data_loader.n_coils, -1))
                 grad_op.fourier_op.samples = traj
                 for i in range(data.shape[0]):
                     grad_op._obs_data = xp.array(data[i])
-                    x_iter[...,i] = self._reconstruct_frame(
-                        grad_op, x_init[:,:,i], n_iter=self.max_iter_per_frame
-                        )          
+                    x_iter[..., i] = self._reconstruct_frame(
+                        grad_op, x_init[:, :, i], n_iter=self.max_iter_per_frame
+                    )
             if self.compute_backend == "cupy":
                 final_estimate[i, ...] = abs(x_iter).get()  # type: ignore
             else:
@@ -411,4 +422,4 @@ def _reconstruct_frame(
             img = grad_op.linear_op.adj_op(opt.x_final)
         else:
             img = opt.x_final
-        return img
+        return img