ROGER gives nan loss when running with EPI

[mfaizan-ai]

@mylyu thank you so much for your assistance and the incredible work.

With reference to #1 you mentioned that I should stack the data on readout axis and add CAIPI shift, so I did the same when I run the ROGER inference with the weights that you shared, it gives NaN with loss values and the reconstructed images are blank and does not show anything after running the inference.

Here is what I do:

1. I have the single band reference data
2. correcting the phase / nyquist ghost for both sms and single band
3. estimating coil sensitivity maps as you did in fastMRI_step1.ipynb
4. simulating single band data to my EPI data (teh simulated data should look identical to my EPI data). So I first do CAIPI shifts in phase encoding axis (PE) with [0, 0.5, 1, 1.5]-->[0, 32, 64, 96) using np.roll() function in image space (shifts are in pixels as my image spatial dims are 64x64). Then convert to kspace to simulate sms data and added. This matches my sms data and successfully reconstructed it and EPI using slice GRAPPA and spit slice grappa.
5. I followed the procedure in the fastMRI_step2.ipynb to create caipi shifts and stack from the data in step 3, but I had to customise it to match my sms data. So the result image that I got from this step is shown below.

so my sms data looks like this in a scan.

this is how my single band reference data looks like:

So Roger does not work with it? What could be wrong? It works with your fastMRI simulated data. so it gives the loss like this

loss: nan; mean: nan; min: nan; max: nan; :  14%|███████████▎                                                                        | 135/1000 [00:14<01:31,  9.49it/s]loss: nan; mean: nan; min: nan; max: nan; :  14%|███████████▎                                                                        | 135/1000 [00:14<01:35,  9.08it/s]

[mfaizan-ai]

@mylyu I figured it. I used the following details to work Roger with EPI.

1. estimating coil sensitivity maps with the same thresholds and stacking all the caipi shifted single band reference slices in readout and estimating coils form it the slices that corresponds to sms slice.
2. created mask
3. loaded sms slice and stacked it also in readout for large fov
4. ran Roger with sms slice, coil sensitivity maps, and masks
   so I sat masks to one as there is no undersampling in kspace in my case. Is it ok? I am using the pretrained weights that you provided on google drive, I think I see some slice leakage. I will attach the code and images of getting Roger work.

import numpy as np
import os
from numpy.fft import fftn, ifftn, fftshift, ifftshift
from spreco.common import utils


def ifft2c(kspace, axes=(-2, -1)):
    return fftshift(
        ifftn(ifftshift(kspace, axes=axes), axes=axes, norm="ortho"),
        axes=axes
    )

def fft2c(img, axes=(-2, -1)):
    return fftshift(
        fftn(ifftshift(img, axes=axes), axes=axes, norm="ortho"),
        axes=axes
    )

def sos(x, axis=-1):
    return np.sqrt(np.sum(np.abs(x) ** 2, axis=axis))


# ------------------------------------------------------------
# Parameters
# ------------------------------------------------------------
MB = 4
R  = 1
pe = 64

# Siemens CMRR: lFOVShiftFactor = 2
# => shifts = [0, 0.5, 1.0, 1.5] FOV
shifts = [i * 0.5 for i in range(MB)]
print("CAIPI shifts (FOV fractions):", shifts)


# ------------------------------------------------------------
# Input SBREF k-space
# sbref shape: (t, sli, lin, cha, col) = (1,36,64,38,64)
# ------------------------------------------------------------
# sbref must already be ghost-corrected k-space
sbref_kspace = fft2c(sbref_img[0], axes=(1, 3)) # (36, 64, 38, 64)                          # (36,64,38,64)
sbref_kspace = sbref_kspace.transpose(0, 1, 3, 2)  # (sli, pe, ro, coils)

Nslices, pe, ro, nc = sbref_kspace.shape
Ngroups = Nslices // MB

print(f"SBREF k-space shape: {sbref_kspace.shape}")


# ------------------------------------------------------------
# Slice groups (CMRR ordering)
# ------------------------------------------------------------
slice_groups = [
    [g + k * Ngroups for k in range(MB)]
    for g in range(Ngroups)
]


# ------------------------------------------------------------
# Loop over SMS slice groups
# ------------------------------------------------------------
for slice_group in slice_groups:

    print(f"\nProcessing slice group {slice_group}")

    # --------------------------------------------------------
    # Extract SB k-space for this group
    # (MB, pe, ro, coils)
    # --------------------------------------------------------
    k_sb = np.stack([sbref_kspace[s] for s in slice_group], axis=0)

    # --------------------------------------------------------
    # Image space
    # --------------------------------------------------------
    img_sb = ifft2c(k_sb, axes=(1, 2))   # (MB, pe, ro, coils)

    # --------------------------------------------------------
    # Apply CAIPI shifts along PE
    # --------------------------------------------------------
    img_shifted = np.zeros_like(img_sb)

    for i, fshift in enumerate(shifts):
        shift_pix = int(fshift * pe)     # (0,32,64,96)
        slice_max = np.abs(img_sb[i]).max()
        if slice_max > 0:
            slice_normalized = img_sb[i] / slice_max
        else:
            slice_normalized = img_sb[i].copy()
            
        img_shifted[i] = np.roll(
            slice_normalized,
            shift=shift_pix,
            axis=0                       # PE axis
        )
        print(f"  slice {slice_group[i]} → shift {shift_pix} px")
    
    # --------------------------------------------------------
    # RO-concatenation (extended FOV)
    # Result: (pe, ro*MB, coils)
    # --------------------------------------------------------
    ro_ext = ro * MB
    img_ro = np.zeros((pe, ro_ext, nc), dtype=np.complex64)

    for i in range(MB):
        img_ro[:, i*ro:(i+1)*ro, :] = img_shifted[i]

    # Center for even MB (matches fastMRI + Siemens convention)
    if MB % 2 == 0:
        img_ro = np.roll(img_ro, -(ro // 2), axis=1)

    # --------------------------------------------------------
    # Back to k-space (RO-concat space)
    # --------------------------------------------------------
    k_ro = fft2c(img_ro, axes=(0, 1))    # (pe, ro_ext, coils)

    # --------------------------------------------------------
    # Convert to bart comptabile axis 
    # --------------------------------------------------------
    k_bart = np.transpose(k_ro, (1, 0, 2))   # (ro_ext, pe, coils)

    # --------------------------------------------------------
    # ESPIRiT coil sensitivity estimation, calibration region and threshold can be changed 
    # for ehanced results
    # TODO: study on threshold and calibraiton region, how it effect the output results 
    # --------------------------------------------------------
    csm_001 = utils.bart(
        1, "ecalib -m1 -r20 -c0.001", k_bart[None]
    )[0]

    csm_985 = utils.bart(
        1, "ecalib -m1 -r20 -c0.985", k_bart[None]
    )[0]

    # --------------------------------------------------------
    # Mask + RSS (for sanity checks)
    # --------------------------------------------------------
    mask = np.abs(csm_985[..., 0]) > 0
    rss  = sos(img_ro, axis=-1)

    # --------------------------------------------------------
    # Save this slice group as .npy file 
    # --------------------------------------------------------
    out = {
        "kspace": k_ro.astype(np.complex64),
        "csm_001": csm_001.astype(np.complex64),
        "csm_985": csm_985.astype(np.complex64),
        "mask": mask.astype(np.uint8),
        "rss": rss.astype(np.float32),
        "shifts": np.array(shifts),
        "slices": np.array(slice_group),
        "MB": MB,
        "R": R,
    }

    out_dir = f"roger_data/coil_maps_estimation/sbref_MB{MB}_R{R}"
    os.makedirs(out_dir, exist_ok=True)
    np.save(
        f"{out_dir}/sbref_MB{MB}_group_{'_'.join(map(str, slice_group))}.npy",
        out
    )
    # break
print("\n✅ Coil sensitivity estimation COMPLETE (RO-concat, CAIPI-consistent)")

then I prepare Roger input like the following code:

import os
import numpy as np
from numpy.fft import fftshift, ifftshift, fftn, ifftn

# ============================================================
# FFT helpers (centered, RO/PE axes = (0,1))
# ============================================================
def fft2c(img, axes=(0, 1)):
    return fftshift(
        fftn(ifftshift(img, axes=axes), axes=axes, norm="ortho"),
        axes=axes
    )

def ifft2c(kspace, axes=(0, 1)):
    return fftshift(
        ifftn(ifftshift(kspace, axes=axes), axes=axes, norm="ortho"),
        axes=axes
    )

def sos(x, axis=-1):
    return np.sqrt(np.sum(np.abs(x) ** 2, axis=axis))



# One coil-map file from previous step
coilmap_npy = (
    "roger_data/coil_maps_estimation/sbref_MB4_R1/"
    "sbref_MB4_group_3_12_21_30.npy"
)

# Output directory
out_dir = "roger_data/input"
os.makedirs(out_dir, exist_ok=True)

# ============================================================
# LOAD COIL MAP DATA (already RO-concat)
# ============================================================
cm = np.load(coilmap_npy, allow_pickle=True).item()

csm_001 = cm["csm_001"]    # (RO_ext, PE, coils)
csm_985 = cm["csm_985"]    # (RO_ext, PE, coils)
mask_in = cm["mask"]       # (RO_ext, PE, coils)
slice_group = cm["slices"] # i.e, [0,9,18,27] or other 
shifts_fov = cm["shifts"]  # [0,0.5,1.0,1.5]
MB = int(cm["MB"])
R  = int(cm["R"])

RO_ext, PE, Nc = csm_001.shape
RO = RO_ext // MB

print("Loaded coil maps:")
print("  csm_001:", csm_001.shape)
print("  mask   :", mask_in.shape)
print("  slices :", slice_group)
print("  shifts :", shifts_fov)
print("  MB,R  :", MB, R)

# ============================================================
# FIX MASK → (RO_ext, PE)
# ============================================================
# mask was saved per-coil → collapse coil dimension
mask = np.any(mask_in >0, axis=-1).astype(np.uint8)
# shape: (RO_ext, PE)

# ============================================================
# SIMULATE SMS FROM SBREF (IMAGE SPACE)
# ============================================================

# sbref_img: (Nslices, PE, RO, coils)
simulate = False  # set true if simulate sbref to sms for roger input 
 
if simulate:
    sbref_img = sbref_img[0]
    Nslices, PE_sb, Nc_sb, RO_sb= sbref_img.shape
    assert PE_sb == PE and RO_sb == RO and Nc_sb == Nc

    # ---- 1) extract SBREF slices for group
    img_sb = np.stack([sbref_img[s] for s in slice_group], axis=0)
    # shape: (MB, PE, RO, coils)

    # ---- 2) CAIPI shifts in PE (image space)
    img_shifted = np.zeros_like(img_sb)

    for i, f in enumerate(shifts_fov):
        shift_pix = int(round(f * PE))
        img_shifted[i] = np.roll(img_sb[i], shift_pix, axis=0)
        print(f"Slice {slice_group[i]} → CAIPI shift {shift_pix} px (PE)")

    # ---- 3) collapse slices (simulate SMS acquisition)
    img_sms_collapsed = np.sum(img_shifted, axis=0).transpose(0, 2, 1)
    # print(img_sms_collapsed.shape)
    
else:
    repetition, slice = 0, 3
    img_sms_collapsed = mb_img[repetition][slice].transpose(0, 2, 1)

# ---- 4) stack collapsed slice MB times along RO
img_ro = np.zeros((PE, ro_ext,  Nc), dtype=np.complex64)

for i in range(MB):
    img_ro[:, i*RO:(i+1)*RO] = img_sms_collapsed
img_ro = img_ro.transpose(1, 0, 2) # transpose puts RO first

# ---- 5) MB centering (even MB)
if MB % 2 == 0:
    img_ro = np.roll(img_ro, -(RO // 2), axis=0)

# ============================================================
# FFT → ALIASED SMS K-SPACE (RO-CONCAT)
# ============================================================
kspace_sms = fft2c(img_ro, axes=(0, 1))
# shape: (RO_ext, PE, coils)

# ============================================================
# RSS FOR VISUAL CHECK
# ============================================================
rss = sos(img_ro, axis=-1)   # (RO_ext, PE)

# ============================================================
# SAMPLING / CALIBRATION MASKS (R=1)
# ============================================================
acc_mask   = np.ones((RO_ext, PE), dtype=np.float32)
calib_mask = np.ones((RO_ext, PE), dtype=np.float32)

# ============================================================
# PACKAGE FOR ROGER
# ============================================================
samples = {
    "readout_data":        (kspace_sms * acc_mask[..., None]).astype(np.complex64),
    "readout_calibration": (kspace_sms * calib_mask[..., None]).astype(np.complex64),
    "readout_mask":        mask.astype(np.uint8),
    "readout_csm":         csm_001.astype(np.complex64),
    "readout_csm_985":     csm_985.astype(np.complex64),
    "rss":                 rss.astype(np.float32),
    "shifts":              np.array(shifts_fov, dtype=np.float32),
    "slices":              np.array(slice_group, dtype=np.int32),
    "MB":                  MB,
    "R":                   R,
}

out_path = os.path.join(
    out_dir,
    f"roger_realSMS_MB{MB}_group_{'_'.join(map(str, slice_group))}.npz"
)

np.savez(out_path, data=samples)
print("\n✅ Saved ROGER input:", out_path)

and then I run teh Roger inference. and I get this output, I will add a sample reconstruction.

is it the correct preprocessing? How do I reduce the noise or slice leakage that I see in teh first slice? How can I further improve the reconstruction result? To me it seems like slice GRAPPA remonstration is better for now, but I think we can improve it with ROGER, so if you have any tips, please let me know.

[mylyu]

@mfaizan-ai You can try fine-tuning the model we released on some clean single-band EPI data from your scanner.