Add support for retro recon and ggrappa at center of k-space

src/mri/cli/reconstruct.py

@@ -3,9 +3,10 @@
 from mri.io.output import save_data
 from mri.cli.utils import raw_config, traj_config, setup_hydra_config, get_outdir_path
 from mri.operators.fourier.utils import discard_frequency_outliers
-from mrinufft.io.utils import add_phase_to_kspace_with_shifts
+from mrinufft.io.utils import add_phase_to_kspace_with_shifts, remove_extra_kspace_samples
 from pymrt.recipes.coils import compress_svd
 from mri.reconstructors import SelfCalibrationReconstructor
+from mri.reconstructors.ggrappa import do_grappa_and_append_data, GRAPPA_RECON_AVAILABLE
 
 import json
 import numpy as np
@@ -20,7 +21,8 @@
 
 
 def dc_adjoint(obs_file: str|np.ndarray, traj_file: str, coil_compress: str|int, debug: int,
-               obs_reader, traj_reader, fourier, output_filename: str = "dc_adjoint.nii"):
+               obs_reader, traj_reader, fourier, grappa_recon=None, output_filename: str = "dc_adjoint.nii",
+               return_data=False):
     """
     Reconstructs an image using the adjoint operator.
 
@@ -48,6 +50,8 @@ def dc_adjoint(obs_file: str|np.ndarray, traj_file: str, coil_compress: str|int,
         It can be:
         1) *.pkl / *.mat: Holds the reconstructed results saved in dictionary as `recon`.
         2) *.nii : NIFTI file holding the reconstructed images.
+    grappa_af: Union[list[int], tuple[int, ...]], optional default 1
+        The acceleration factor for the GRAPPA reconstruction.
 
     Returns
     -------
@@ -86,33 +90,55 @@ def dc_adjoint(obs_file: str|np.ndarray, traj_file: str, coil_compress: str|int,
         for size in traj_params['img_size']
     ])
     log.info(f"Trajectory Parameters: {traj_params}")
-    kspace_loc = shots.reshape(-1, traj_params["dimension"])
     data_header["shifts"] = data_header['shifts'][:traj_params["dimension"]]
     normalized_shifts = (
         np.array(data_header["shifts"])
         / np.array(traj_params["FOV"])
         * np.array(traj_params["img_size"])
         / 1000
     )
-    if kspace_loc.max() > 0.5 or kspace_loc.min() < 0.5:
-        log.warn(f"K-space locations are above the unity range, discarding the outlier data")
-        if data_header["type"] == "retro_recon":
-            kspace_loc = discard_frequency_outliers(kspace_loc)
-            kspace_data = np.squeeze(raw_data)
-        else:
-            kspace_loc, kspace_data = discard_frequency_outliers(kspace_loc, np.squeeze(raw_data))
-    kspace_data = kspace_data.astype(np.complex64)
-    kspace_loc = kspace_loc.astype(np.float32)
+    kspace_data = np.squeeze(raw_data).astype(np.complex64)
+    kspace_loc = shots.reshape(-1, traj_params["dimension"]).astype(np.float32)
+    kspace_data = remove_extra_kspace_samples(kspace_data, shots.shape[1])
+    kspace_data = kspace_data.reshape(kspace_data.shape[0], -1)
     log.info(f"Phase shifting raw data for Normalized shifts: {normalized_shifts}")
     kspace_data = add_phase_to_kspace_with_shifts(
-        kspace_data, kspace_loc, normalized_shifts
+        kspace_data, kspace_loc.reshape(-1, traj_params["dimension"]), normalized_shifts
     )
+    try:
+        af_string = data_header['trajectory_name'].split('_G')[1].split('_')[0].split('x')
+        if len(af_string) > 1 and 'd' in af_string[1]:
+            af_caipi = af_string[1].split('d')
+            af_string[1] = af_caipi[0]
+            grappa_recon.keywords['delta'] = int(af_caipi[1])
+        grappa_recon.keywords['af'] = tuple([int(float(af)) for af in af_string])
+    except:
+        grappa_recon.keywords['af'] = (1, )
+        grappa_recon.keywords['delta'] = 0
+    if grappa_recon is not None and np.prod(grappa_recon.keywords['af'])>1:
+        log.info("Performing GRAPPA Reconstruction: AF: %s", af_string)
+        log.info("GRAPPA AF: %s", grappa_recon.keywords['af'])
+        kspace_loc, kspace_data = do_grappa_and_append_data(
+            kspace_loc,
+            kspace_data,
+            traj_params,
+            grappa_recon,
+            acs=data_header["acs"], # Pass ACS if read in data (external)
+        )
     if coil_compress != -1:
+        log.info("Compressing coils")
         kspace_data = np.ascontiguousarray(compress_svd(
             kspace_data,
             k_svd=coil_compress,
             coil_axis=0
-        ))
+        )).astype(np.complex64)
+    if kspace_loc.max() > 0.5 or kspace_loc.min() < 0.5:
+        log.warn(f"K-space locations are above the unity range, discarding the outlier data")
+        if data_header["type"] == "retro_recon":
+            kspace_loc = discard_frequency_outliers(kspace_loc)
+            kspace_data = np.squeeze(raw_data)
+        else:
+            kspace_loc, kspace_data = discard_frequency_outliers(kspace_loc, kspace_data)
     fourier.keywords['smaps'] = partial(
         fourier.keywords['smaps'],
         kspace_data=kspace_data,
@@ -127,6 +153,7 @@ def dc_adjoint(obs_file: str|np.ndarray, traj_file: str, coil_compress: str|int,
             'density_comp': fourier_op.impl.density,
             'traj_params': traj_params,
             'data_header': data_header,
+            'kspace_loc': kspace_loc,
         }
         save_data_hydra('smaps.nii', fourier_op.impl.smaps)
         if coil_compress != -1:
@@ -140,13 +167,15 @@ def dc_adjoint(obs_file: str|np.ndarray, traj_file: str, coil_compress: str|int,
     log.info("Saving DC Adjoint")
     data_header['traj_params'] = traj_params
     save_data_hydra(output_filename, dc_adjoint, data_header)
-    return dc_adjoint, (fourier_op, kspace_data, traj_params, data_header)
+    if return_data:
+        return dc_adjoint, (fourier_op, kspace_data, traj_params, data_header)
 
 
 
 def recon(obs_file: str, traj_file: str, mu: float, num_iterations: int, coil_compress: str|int, 
           algorithm: str, debug: int, obs_reader, traj_reader, fourier, linear, sparsity,
-          output_filename: str = "recon.nii", remove_dc_for_recon: bool = True, validation_recon: np.ndarray = None, metrics: dict = None):
+          output_filename: str = "recon.nii", remove_dc_for_recon: bool = True, validation_recon: np.ndarray = None, metrics: dict = None, 
+          grappa_recon=None):
     """Reconstructs an MRI image using the given parameters.
 
     Parameters
@@ -193,7 +222,9 @@ def recon(obs_file: str, traj_file: str, mu: float, num_iterations: int, coil_co
         obs_reader,
         traj_reader,
         fourier,
+        grappa_recon=grappa_recon,
         output_filename='dc_adj_' + output_filename,
+        return_data=True,
     )
     fourier_op, kspace_data, traj_params, data_header = additional_data
     if remove_dc_for_recon:
@@ -246,6 +277,7 @@ def recon(obs_file: str, traj_file: str, mu: float, num_iterations: int, coil_co
         "_self_",
         {"fourier": "gpu"},
         {"fourier/density_comp": "pipe"},
+        {"grappa_recon": "disable"} if GRAPPA_RECON_AVAILABLE else {},
         {"fourier/smaps": "low_frequency"},
     ],
     name="dc_adjoint",
@@ -257,11 +289,13 @@ def recon(obs_file: str, traj_file: str, mu: float, num_iterations: int, coil_co
     algorithm="pogm",
     num_iterations=30,
     coil_compress=10,
+    mu=1e-7,
     debug=1,
     hydra_defaults=[
         "_self_",
         {"fourier": "gpu"},
         {"fourier/density_comp": "pipe"},
+        {"grappa_recon": "disable"} if GRAPPA_RECON_AVAILABLE else {},
         {"fourier/smaps": "low_frequency"},
         {"linear": "gpu"},
         {"sparsity": "weighted_sparse"},
@@ -276,10 +310,12 @@ def recon(obs_file: str, traj_file: str, mu: float, num_iterations: int, coil_co
     algorithm="pogm",
     num_iterations=30,
     coil_compress=5,
+    mu=1e-7,
     debug=1,
     hydra_defaults=[
         "_self_",
         {"fourier": "gpu_lowmem"},
+        {"grappa_recon": "disable"} if GRAPPA_RECON_AVAILABLE else {},
         {"fourier/density_comp": "pipe_lowmem"},
         {"fourier/smaps": "low_frequency"},
     ],

src/mri/cli/utils.py

@@ -12,6 +12,18 @@
 from modopt.opt.linear import Identity
 import os
 
+try:
+    from ggrappa.grappaND import GRAPPA_Recon
+    grappa_config = builds(
+        GRAPPA_Recon,
+        zen_exclude=["sig", "acs", "isGolfSparks", "quiet", "af"],
+        zen_partial=True,        
+        populate_full_signature=True,
+    )
+    grappa_store = store(group="grappa_recon")
+    grappa_store(grappa_config, name="disable")
+except:
+    pass
 
 raw_config = builds(read_arbgrad_rawdat, populate_full_signature=True, zen_partial=True)
 
@@ -31,6 +43,7 @@
 smaps_config = builds(
     get_smaps("low_frequency"),
     populate_full_signature=True,
+    blurr_factor=30.0,
     # We estimate density, with separate args. It is passed by compute_smaps in mri-nufft
     zen_exclude=["density"],
     zen_partial=True,

src/mri/reconstructors/ggrappa.py

@@ -0,0 +1,27 @@
+import numpy as np
+
+GRAPPA_RECON_AVAILABLE = False
+try:
+    from ggrappa.grappaND import GRAPPA_Recon
+    import torch
+    from ggrappa.utils import get_cart_portion_sparkling, get_grappa_filled_data_and_loc
+    GRAPPA_RECON_AVAILABLE = True
+except:
+    pass
+
+
+def do_grappa_and_append_data(kspace_loc, kspace_data, traj_params, grappa_maker, acs=None):
+    kspace_shots = kspace_loc.reshape(traj_params['num_shots'], -1, traj_params['dimension'])
+    if not GRAPPA_RECON_AVAILABLE:
+        raise ValueError("GRAPPA is not available")
+    gridded_center = get_cart_portion_sparkling(kspace_shots, traj_params, kspace_data)
+    grappa_recon, grappa_kernel = grappa_maker(
+        sig=torch.tensor(gridded_center).permute(0, 2, 3, 1),
+        acs=torch.tensor(acs).permute(0, 2, 3, 1) if acs is not None else None,
+        isGolfSparks=True,
+    )
+    grappa_recon = grappa_recon.permute(0, 3, 1, 2).numpy()
+    extra_loc, extra_data = get_grappa_filled_data_and_loc(gridded_center, grappa_recon, traj_params)
+    kspace_loc = np.concatenate([kspace_loc, extra_loc], axis=0)
+    kspace_data = np.hstack([kspace_data, extra_data])
+    return kspace_loc, kspace_data