from fastmri.data.subsample import create_mask_for_mask_type
from fastmri.data.transforms import apply_mask, to_tensor, center_crop
import numpy as np

mask_func =create_mask_for_mask_type(
    mask_type_str="equispaced",
    center_fractions=[0.37],
    accelerations=[4]
)

kspace = np.load("data/prostate1_kspace.npy")
print(kspace.shape) # (34, 14, 640, 451)
kspace = to_tensor(kspace)
print(kspace.shape) # torch.Size([34, 14, 640, 451, 2])
subsampled_kspace, mask, num_low_frequencies = apply_mask(
    kspace,
    mask_func,
    seed=1
)