init

.gitattributes

@@ -0,0 +1,61 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.lz4 filter=lfs diff=lfs merge=lfs -text
+*.mds filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+# Audio files - uncompressed
+*.pcm filter=lfs diff=lfs merge=lfs -text
+*.sam filter=lfs diff=lfs merge=lfs -text
+*.raw filter=lfs diff=lfs merge=lfs -text
+# Audio files - compressed
+*.aac filter=lfs diff=lfs merge=lfs -text
+*.flac filter=lfs diff=lfs merge=lfs -text
+*.mp3 filter=lfs diff=lfs merge=lfs -text
+*.ogg filter=lfs diff=lfs merge=lfs -text
+*.wav filter=lfs diff=lfs merge=lfs -text
+# Image files - uncompressed
+*.bmp filter=lfs diff=lfs merge=lfs -text
+*.gif filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.tiff filter=lfs diff=lfs merge=lfs -text
+# Image files - compressed
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text
+*.webp filter=lfs diff=lfs merge=lfs -text
+# Video files - compressed
+*.mp4 filter=lfs diff=lfs merge=lfs -text
+*.webm filter=lfs diff=lfs merge=lfs -text
+dataset/kspace/data.mdb filter=lfs diff=lfs merge=lfs -text
+dataset/rss/data.mdb filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,124 @@
+# Config files
+fastmri.yaml
+
+# Python specific
+__pycache__/
+.pytest_cache/
+*.py[cod]
+*.so
+*.egg-info/
+*.pyo
+*.pyd
+
+# Virtual environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Code in development
+ignore_**.py
+
+# Hidden / ignore folders
+hidden/
+ignore/
+hidden_**
+
+# Jupyter Notebook Checkpoints
+.ipynb_checkpoints
+
+# Data files
+data/
+datasets/
+dataset
+*.csv
+*.tsv
+*.h5
+*.json
+*.xml
+*.parquet
+*.pkl
+
+# Model files
+*.ckpt
+*.h5
+*.tflite
+*.onnx
+*.pb
+*.pth
+*.pt
+*.joblib
+*.pkl
+
+# Logs and outputs
+logs/
+wandb/
+*.log
+*.out
+*.txt
+*.csv
+
+# Test dir
+!tests/**/*.txt
+!tests/datasets
+
+# Results
+results/
+output/
+runs/
+outfig/
+figs/*.png
+
+# SLURM
+slurm/
+
+# Ignore files related to experiments
+experiments/
+
+# Temporary files
+*.tmp
+*.temp
+*.swp
+*.swo
+
+# VS Code specific
+.vscode/
+*.code-workspace
+
+# System files
+.DS_Store
+Thumbs.db
+
+# Environment files
+*.env
+
+# Ignore files from data processing tools
+*.dvc
+.dvc/
+
+# PyTorch Lightning Logs
+lightning_logs/
+
+# Ignore files generated by package managers
+Pipfile
+Pipfile.lock
+poetry.lock
+
+# TensorBoard logs
+logs/
+events.out.tfevents.*
+
+# Checkpoints and weights
+checkpoints/
+weights/
+
+# Large file extensions
+*.tar.gz
+*.zip
+*.tar
+*.gz
+
+

.gradio/certificate.pem

@@ -0,0 +1,31 @@
+-----BEGIN CERTIFICATE-----
+MIIFazCCA1OgAwIBAgIRAIIQz7DSQONZRGPgu2OCiwAwDQYJKoZIhvcNAQELBQAw
+TzELMAkGA1UEBhMCVVMxKTAnBgNVBAoTIEludGVybmV0IFNlY3VyaXR5IFJlc2Vh
+cmNoIEdyb3VwMRUwEwYDVQQDEwxJU1JHIFJvb3QgWDEwHhcNMTUwNjA0MTEwNDM4
+WhcNMzUwNjA0MTEwNDM4WjBPMQswCQYDVQQGEwJVUzEpMCcGA1UEChMgSW50ZXJu
+ZXQgU2VjdXJpdHkgUmVzZWFyY2ggR3JvdXAxFTATBgNVBAMTDElTUkcgUm9vdCBY
+MTCCAiIwDQYJKoZIhvcNAQEBBQADggIPADCCAgoCggIBAK3oJHP0FDfzm54rVygc
+h77ct984kIxuPOZXoHj3dcKi/vVqbvYATyjb3miGbESTtrFj/RQSa78f0uoxmyF+
+0TM8ukj13Xnfs7j/EvEhmkvBioZxaUpmZmyPfjxwv60pIgbz5MDmgK7iS4+3mX6U
+A5/TR5d8mUgjU+g4rk8Kb4Mu0UlXjIB0ttov0DiNewNwIRt18jA8+o+u3dpjq+sW
+T8KOEUt+zwvo/7V3LvSye0rgTBIlDHCNAymg4VMk7BPZ7hm/ELNKjD+Jo2FR3qyH
+B5T0Y3HsLuJvW5iB4YlcNHlsdu87kGJ55tukmi8mxdAQ4Q7e2RCOFvu396j3x+UC
+B5iPNgiV5+I3lg02dZ77DnKxHZu8A/lJBdiB3QW0KtZB6awBdpUKD9jf1b0SHzUv
+KBds0pjBqAlkd25HN7rOrFleaJ1/ctaJxQZBKT5ZPt0m9STJEadao0xAH0ahmbWn
+OlFuhjuefXKnEgV4We0+UXgVCwOPjdAvBbI+e0ocS3MFEvzG6uBQE3xDk3SzynTn
+jh8BCNAw1FtxNrQHusEwMFxIt4I7mKZ9YIqioymCzLq9gwQbooMDQaHWBfEbwrbw
+qHyGO0aoSCqI3Haadr8faqU9GY/rOPNk3sgrDQoo//fb4hVC1CLQJ13hef4Y53CI
+rU7m2Ys6xt0nUW7/vGT1M0NPAgMBAAGjQjBAMA4GA1UdDwEB/wQEAwIBBjAPBgNV
+HRMBAf8EBTADAQH/MB0GA1UdDgQWBBR5tFnme7bl5AFzgAiIyBpY9umbbjANBgkq
+hkiG9w0BAQsFAAOCAgEAVR9YqbyyqFDQDLHYGmkgJykIrGF1XIpu+ILlaS/V9lZL
+ubhzEFnTIZd+50xx+7LSYK05qAvqFyFWhfFQDlnrzuBZ6brJFe+GnY+EgPbk6ZGQ
+3BebYhtF8GaV0nxvwuo77x/Py9auJ/GpsMiu/X1+mvoiBOv/2X/qkSsisRcOj/KK
+NFtY2PwByVS5uCbMiogziUwthDyC3+6WVwW6LLv3xLfHTjuCvjHIInNzktHCgKQ5
+ORAzI4JMPJ+GslWYHb4phowim57iaztXOoJwTdwJx4nLCgdNbOhdjsnvzqvHu7Ur
+TkXWStAmzOVyyghqpZXjFaH3pO3JLF+l+/+sKAIuvtd7u+Nxe5AW0wdeRlN8NwdC
+jNPElpzVmbUq4JUagEiuTDkHzsxHpFKVK7q4+63SM1N95R1NbdWhscdCb+ZAJzVc
+oyi3B43njTOQ5yOf+1CceWxG1bQVs5ZufpsMljq4Ui0/1lvh+wjChP4kqKOJ2qxq
+4RgqsahDYVvTH9w7jXbyLeiNdd8XM2w9U/t7y0Ff/9yi0GE44Za4rF2LN9d11TPA
+mRGunUHBcnWEvgJBQl9nJEiU0Zsnvgc/ubhPgXRR4Xq37Z0j4r7g1SgEEzwxA57d
+emyPxgcYxn/eR44/KJ4EBs+lVDR3veyJm+kXQ99b21/+jh5Xos1AnX5iItreGCc=
+-----END CERTIFICATE-----

README.md

@@ -0,0 +1,33 @@
+---
+tags:
+- medical
+- mri
+- neuraloperator
+- fastmri
+pretty_name: fastMRI Tiny
+---
+# A Unified Model for Compressed Sensing MRI Across Undersampling Patterns
+
+> [**A Unified Model for Compressed Sensing MRI Across Undersampling Patterns**](https://arxiv.org/abs/2410.16290)  
+> Armeet Singh Jatyani, Jiayun Wang, Aditi Chandrashekar, Zihui Wu, Miguel Liu-Schiaffini, Bahareh Tolooshams, Anima Anandkumar  
+> *Paper at [CVPR 2025](https://cvpr.thecvf.com/Conferences/2025/AcceptedPapers)*
+
+This is a tiny subset of 230 fastMRI samples, used in the demo for the above [paper](https://huggingface.co/armeet/nomri) at CVPR 2025!
+
+
+## Citation
+
+If you found our work helpful or used any of our models (UDNO), please cite the following:
+```bibtex
+@inproceedings{jatyani2025nomri,
+  author    = {Armeet Singh Jatyani* and Jiayun Wang* and Aditi Chandrashekar and Zihui Wu and Miguel Liu-Schiaffini and Bahareh Tolooshams and Anima Anandkumar},
+  title     = {A Unified Model for Compressed Sensing MRI Across Undersampling Patterns},
+  booktitle = {Conference on Computer Vision and Pattern Recognition (CVPR) Proceedings},
+  abbr      = {CVPR},
+  year      = {2025}
+}
+```
+
+![paper_preview](https://github.com/user-attachments/assets/7e6adaa5-a5fa-4b68-bd8c-5279f6f643d7)
+
+https://arxiv.org/abs/2410.16290

app.py

@@ -0,0 +1,264 @@
+import io
+import os
+import sys
+
+import gradio as gr
+import numpy as np
+
+# import spaces
+# from huggingface_hub import hf_hub_download
+from huggingface_hub import snapshot_download
+from PIL import Image, ImageDraw, ImageFont
+
+# Set the working directory to the root directory
+# root_dir = os.path.abspath("..")
+# os.chdir(root_dir)
+# sys.path.insert(0, root_dir)
+
+# download dataset & weights
+snapshot_download(repo_id="armeet/fastmri-tiny", repo_type="dataset", local_dir=".")
+
+
+device = "cuda"
+# dataset_path = "/global/homes/p/peterwg/pscratch/datasets/mri_knee_dummy"
+dataset_path = "dataset"
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+import fastmri
+from fastmri.datasets import SliceDatasetLMDB, SliceSample
+from fastmri.subsample import create_mask_for_mask_type
+from models.lightning.no_varnet_module import NOVarnetModule
+from models.lightning.varnet_module import VarNetModule
+
+acceleration_to_fractions = {
+    1: 1,
+    2: 0.16,
+    4: 0.08,
+    6: 0.06,
+    8: 0.04,
+    16: 0.02,
+    32: 0.01,
+}
+
+
+def create_mask_fn(center_fraction, acceleration):
+    mask_fn = create_mask_for_mask_type(
+        "equispaced_fraction",
+        [center_fraction],
+        [acceleration],
+    )
+    return mask_fn
+
+
+mask_4x = create_mask_fn(acceleration_to_fractions[4], 4)
+mask_6x = create_mask_fn(acceleration_to_fractions[6], 6)
+mask_8x = create_mask_fn(acceleration_to_fractions[8], 8)
+mask_16x = create_mask_fn(acceleration_to_fractions[16], 16)
+
+val_dataset_4x = SliceDatasetLMDB(
+    "knee",
+    partition="val",
+    mask_fns=[mask_4x],
+    complex=False,
+    root=dataset_path,
+    crop_shape=(320, 320),
+    coils=15,
+)
+
+val_dataset_6x = SliceDatasetLMDB(
+    "knee",
+    partition="val",
+    mask_fns=[mask_6x],
+    complex=False,
+    root=dataset_path,
+    crop_shape=(320, 320),
+    coils=15,
+)
+
+val_dataset_8x = SliceDatasetLMDB(
+    "knee",
+    partition="val",
+    mask_fns=[mask_8x],
+    complex=False,
+    root=dataset_path,
+    crop_shape=(320, 320),
+    coils=15,
+)
+val_dataset_16x = SliceDatasetLMDB(
+    "knee",
+    partition="val",
+    mask_fns=[mask_16x],
+    complex=False,
+    root=dataset_path,
+    crop_shape=(320, 320),
+    coils=15,
+)
+
+vn = VarNetModule.load_from_checkpoint(
+    "vn.ckpt",
+)
+no = NOVarnetModule.load_from_checkpoint(
+    "no.ckpt",
+)
+no.eval()
+vn.eval()
+
+bright_samples = [42, 69, 80, 137, 139, 226, 229]
+
+
+def v(x):
+    return x.detach().cpu().numpy().squeeze()
+
+
+def viz(x, cmap="gray", vmin=0, vmax=1):
+    processed_data = v(x)
+    fig, ax = plt.subplots()
+    ax.imshow(processed_data, cmap=cmap, vmin=vmin, vmax=vmax)
+    ax.axis("off")  # Turn off axes
+    fig.subplots_adjust(left=0, right=1, top=1, bottom=0)  # Adjust margins
+    buf = io.BytesIO()
+    plt.savefig(buf, format="png", bbox_inches="tight", pad_inches=0)
+    buf.seek(0)  # Rewind the buffer to the beginning
+    plt.show()
+    try:
+        img = Image.open(buf)
+        img_array = np.array(img)
+    except Exception as e:
+        print(f"Error converting image buffer to NumPy array: {e}")
+        img_array = None
+    finally:
+        plt.close(fig)
+        buf.close()
+    return img_array
+
+
+def forward(model, idx, rate):
+    if rate == 4:
+        dataset = val_dataset_4x
+    elif rate == 6:
+        dataset = val_dataset_6x
+    elif rate == 8:
+        dataset = val_dataset_8x
+    elif rate == 16:
+        dataset = val_dataset_16x
+    else:
+        raise ValueError("Invalid rate")
+
+    sample = dataset[idx]
+    mask, k, target = (
+        sample.mask.to(device),
+        sample.masked_kspace.to(device),
+        sample.target.to(device),
+    )
+    pred = model(k.unsqueeze(0), mask.unsqueeze(0), None)
+
+    return mask, k, target, pred[0]
+
+
+def update_interface(sample_id, sample_rate):
+    n = [None] * 6
+    if sample_id is None or sample_rate is None or sample_id not in bright_samples:
+        return n
+
+    mask, k, target, pred_vn = forward(vn, sample_id, sample_rate)
+    _, _, _, pred_no = forward(no, sample_id, sample_rate)
+
+    k = viz(mask[0, :, :, 0], cmap="gray", vmin=0, vmax=1)
+    target_res = viz(target, cmap="gray", vmin=None, vmax=None)
+
+    pred_no_res = viz(pred_no, cmap="gray", vmin=None, vmax=None)
+    pred_vn_res = viz(pred_vn, cmap="gray", vmin=None, vmax=None)
+
+    diff_no_res = viz(torch.abs(pred_no - target), cmap=None, vmin=None, vmax=None)
+    diff_vn_res = viz(torch.abs(pred_vn - target), cmap=None, vmin=None, vmax=None)
+
+    return k, target_res, pred_no_res, pred_vn_res, diff_no_res, diff_vn_res
+
+
+with gr.Blocks(theme=gr.themes.Monochrome(), fill_width=True) as demo:
+    gr.Markdown(
+        "# A Unified Model for Compressed Sensing MRI Across Undersampling Patterns [CPVR 2025]"
+    )
+    gr.Markdown("""  
+> Armeet Singh Jatyani, Jiayun Wang, Aditi Chandrashekar, Zihui Wu, Miguel Liu-Schiaffini, Bahareh Tolooshams, Anima Anandkumar  
+                """)
+    gr.Markdown(
+        "[![arXiv](https://img.shields.io/badge/arXiv-2410.16290-b31b1b.svg?style=flat-square&logo=arxiv)](https://arxiv.org/abs/2410.16290)"
+    )
+    gr.Markdown(
+        "[![](https://img.shields.io/badge/Blog-armeet.ca%2Fnomri-yellow?style=flat-square)](https://armeet.ca/nomri)"
+    )
+
+    gr.Markdown(
+        "This demo showcases the performance of our unified model for compressed sensing MRI across different acceleration rates."
+    )
+
+    with gr.Row():
+        dropdown_sample = gr.Dropdown(
+            choices=bright_samples,
+            label="Select a Sample",
+            info="Choose one of the available samples.",
+            filterable=False,
+            value=229,
+        )
+    with gr.Row():
+        dropdown_rate = gr.Radio(
+            choices=[16, 8, 6, 4],
+            value=16,
+            label="Select an Acceleration Rate",
+            info="Ex: 4x means the model is trained to reconstruct from 4x undersampled k-space data",
+            # filterable=False,
+        )
+
+    with gr.Row():
+        with gr.Column():
+            gr.Label("Undersampling Mask")
+            k = gr.Image(label=None, interactive=False)
+        with gr.Column():
+            gr.Label("Ground Truth")
+            target = gr.Image(label=None, interactive=False)
+        with gr.Column():
+            gr.Label("NO (ours)")
+            pred_no = gr.Image(label="Reconstruction", interactive=False)
+        with gr.Column():
+            gr.Label("VN (existing)")
+            pred_vn = gr.Image(label="Reconstruction", interactive=False)
+    with gr.Row():
+        with gr.Column():
+            pass
+        with gr.Column():
+            pass
+        with gr.Column():
+            diff_no = gr.Image(label="| Recon - GT |", interactive=False)
+        with gr.Column():
+            diff_vn = gr.Image(label="| Recon - GT |", interactive=False)
+
+    gr.Markdown("""
+```
+@inproceedings{jatyani2025nomri,
+  author    = {Armeet Singh Jatyani* and Jiayun Wang* and Aditi Chandrashekar and Zihui Wu and Miguel Liu-Schiaffini and Bahareh Tolooshams and Anima Anandkumar},
+  title     = {A Unified Model for Compressed Sensing MRI Across Undersampling Patterns},
+  booktitle = {Conference on Computer Vision and Pattern Recognition (CVPR) Proceedings},
+  abbr      = {CVPR},
+  year      = {2025}
+}
+```
+                """)
+
+    update_inputs = [dropdown_sample, dropdown_rate]
+    update_outputs = [k, target, pred_no, pred_vn, diff_no, diff_vn]
+
+    dropdown_sample.change(
+        fn=update_interface, inputs=update_inputs, outputs=update_outputs
+    )
+    dropdown_rate.change(
+        fn=update_interface, inputs=update_inputs, outputs=update_outputs
+    )
+
+if __name__ == "__main__":
+    demo.launch(share=True)

environment.yml

@@ -0,0 +1,158 @@
+name: no-med
+channels:
+  - conda-forge
+  - defaults
+dependencies:
+  - _libgcc_mutex=0.1=conda_forge
+  - _openmp_mutex=4.5=2_gnu
+  - asttokens=2.4.1=pyhd8ed1ab_0
+  - bzip2=1.0.8=h5eee18b_6
+  - ca-certificates=2024.8.30=hbcca054_0
+  - comm=0.2.2=pyhd8ed1ab_0
+  - debugpy=1.6.7=py312h6a678d5_0
+  - decorator=5.1.1=pyhd8ed1ab_0
+  - exceptiongroup=1.2.2=pyhd8ed1ab_0
+  - executing=2.1.0=pyhd8ed1ab_0
+  - expat=2.6.3=h6a678d5_0
+  - importlib-metadata=8.5.0=pyha770c72_0
+  - ipykernel=6.29.5=pyh3099207_0
+  - ipython=8.27.0=pyh707e725_0
+  - jedi=0.19.1=pyhd8ed1ab_0
+  - jupyter_client=8.6.3=pyhd8ed1ab_0
+  - jupyter_core=5.7.2=py312h06a4308_0
+  - krb5=1.21.3=h143b758_0
+  - ld_impl_linux-64=2.38=h1181459_1
+  - libedit=3.1.20230828=h5eee18b_0
+  - libffi=3.4.4=h6a678d5_1
+  - libgcc=14.1.0=h77fa898_1
+  - libgcc-ng=14.1.0=h69a702a_1
+  - libgomp=14.1.0=h77fa898_1
+  - libsodium=1.0.20=h4ab18f5_0
+  - libstdcxx=14.1.0=hc0a3c3a_1
+  - libstdcxx-ng=11.2.0=h1234567_1
+  - libuuid=1.41.5=h5eee18b_0
+  - matplotlib-inline=0.1.7=pyhd8ed1ab_0
+  - ncurses=6.4=h6a678d5_0
+  - nest-asyncio=1.6.0=pyhd8ed1ab_0
+  - openssl=3.3.2=hb9d3cd8_0
+  - packaging=24.1=pyhd8ed1ab_0
+  - parso=0.8.4=pyhd8ed1ab_0
+  - pexpect=4.9.0=pyhd8ed1ab_0
+  - pickleshare=0.7.5=py_1003
+  - pip=24.2=py312h06a4308_0
+  - prompt-toolkit=3.0.47=pyha770c72_0
+  - ptyprocess=0.7.0=pyhd3deb0d_0
+  - pure_eval=0.2.3=pyhd8ed1ab_0
+  - pygments=2.18.0=pyhd8ed1ab_0
+  - python=3.12.4=h5148396_1
+  - pyzmq=25.1.2=py312h6a678d5_0
+  - readline=8.2=h5eee18b_0
+  - setuptools=72.1.0=py312h06a4308_0
+  - six=1.16.0=pyh6c4a22f_0
+  - sqlite=3.45.3=h5eee18b_0
+  - stack_data=0.6.2=pyhd8ed1ab_0
+  - tk=8.6.14=h39e8969_0
+  - tornado=6.4.1=py312h5eee18b_0
+  - traitlets=5.14.3=pyhd8ed1ab_0
+  - typing_extensions=4.12.2=pyha770c72_0
+  - wcwidth=0.2.13=pyhd8ed1ab_0
+  - wheel=0.43.0=py312h06a4308_0
+  - xz=5.4.6=h5eee18b_1
+  - zeromq=4.3.5=ha4adb4c_5
+  - zipp=3.20.2=pyhd8ed1ab_0
+  - zlib=1.2.13=h5eee18b_1
+  - pip:
+      - aiohappyeyeballs==2.4.0
+      - aiohttp==3.10.5
+      - aiosignal==1.3.1
+      - antlr4-python3-runtime==4.9.3
+      - attrs==24.2.0
+      - black==24.10.0
+      - certifi==2024.8.30
+      - charset-normalizer==3.3.2
+      - click==8.1.7
+      - cloudpickle==3.0.0
+      - contourpy==1.3.0
+      - cycler==0.12.1
+      - docker-pycreds==0.4.0
+      - filelock==3.16.0
+      - fonttools==4.53.1
+      - frozenlist==1.4.1
+      - fsspec==2024.9.0
+      - gitdb==4.0.11
+      - gitpython==3.1.43
+      - h5py==3.11.0
+      - hydra-core==1.3.2
+      - hydra-submitit-launcher==1.2.0
+      - idna==3.8
+      - imageio==2.35.1
+      - iniconfig==2.0.0
+      - isort==5.13.2
+      - jinja2==3.1.4
+      - joblib==1.4.2
+      - kiwisolver==1.4.7
+      - lazy-loader==0.4
+      - lightning==2.4.0
+      - lightning-utilities==0.11.7
+      - llvmlite==0.43.0
+      - lmdb==1.5.1
+      - markupsafe==2.1.5
+      - matplotlib==3.9.2
+      - mpmath==1.3.0
+      - multidict==6.0.5
+      - mypy-extensions==1.0.0
+      - networkx==3.3
+      - no-med==0.0.0
+      - numba==0.60.0
+      - numpy==2.0.2
+      - nvidia-cublas-cu12==12.1.3.1
+      - nvidia-cuda-cupti-cu12==12.1.105
+      - nvidia-cuda-nvrtc-cu12==12.1.105
+      - nvidia-cuda-runtime-cu12==12.1.105
+      - nvidia-cudnn-cu12==9.1.0.70
+      - nvidia-cufft-cu12==11.0.2.54
+      - nvidia-curand-cu12==10.3.2.106
+      - nvidia-cusolver-cu12==11.4.5.107
+      - nvidia-cusparse-cu12==12.1.0.106
+      - nvidia-nccl-cu12==2.20.5
+      - nvidia-nvjitlink-cu12==12.6.68
+      - nvidia-nvtx-cu12==12.1.105
+      - omegaconf==2.3.0
+      - opencv-python==4.10.0.84
+      - pandas==2.2.2
+      - pathspec==0.12.1
+      - pillow==10.4.0
+      - platformdirs==4.3.2
+      - pluggy==1.5.0
+      - protobuf==5.28.0
+      - psutil==6.0.0
+      - pyparsing==3.1.4
+      - pytest==8.3.3
+      - python-dateutil==2.9.0.post0
+      - pytorch-lightning==2.4.0
+      - pytz==2024.1
+      - pywavelets==1.7.0
+      - pyyaml==6.0.2
+      - requests==2.32.3
+      - runstats==2.0.0
+      - scikit-image==0.24.0
+      - scipy==1.14.1
+      - sentry-sdk==2.13.0
+      - setproctitle==1.3.3
+      - sigpy==0.1.26
+      - smmap==5.0.1
+      - submitit==1.5.1
+      - sympy==1.13.2
+      - tabulate==0.9.0
+      - tifffile==2024.8.30
+      - toolz==1.0.0
+      - torch==2.4.1
+      - torchmetrics==1.4.1
+      - torchvision==0.19.1
+      - tqdm==4.66.5
+      - triton==3.0.0
+      - tzdata==2024.1
+      - urllib3==2.2.2
+      - wandb==0.17.9
+      - yarl==1.11.0
+prefix: /global/homes/p/peterwg/local/miniconda3/envs/no-med

fastmri/__init__.py

@@ -0,0 +1,20 @@
+"""
+Copyright (c) Facebook, Inc. and its affiliates.
+
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+from .coil_combine import rss, rss_complex, mvue
+from .fftc import fft2c_new as fft2c
+from .fftc import fftshift
+from .fftc import ifft2c_new as ifft2c
+from .fftc import ifftshift, roll
+from .losses import SSIMLoss
+from .math_utils import (
+    complex_abs,
+    complex_abs_sq,
+    complex_conj,
+    complex_mul,
+    tensor_to_complex_np,
+)

fastmri/coil_combine.py

@@ -0,0 +1,67 @@
+"""
+Copyright (c) Facebook, Inc. and its affiliates.
+
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+import torch
+
+import fastmri
+import sigpy as sp
+import numpy as np
+
+
+def rss(data: torch.Tensor, dim: int = 0) -> torch.Tensor:
+    """
+    Compute the Root Sum of Squares (RSS).
+
+    The RSS is computed assuming that `dim` is the coil dimension.
+
+    Parameters
+    ----------
+    data : torch.Tensor
+        The input tensor.
+    dim : int, optional
+        The dimension along which to apply the RSS transform (default is 0).
+
+    Returns
+    -------
+    torch.Tensor
+        The computed RSS value.
+    """
+    return torch.sqrt((data**2).sum(dim))
+
+
+def mvue(spatial_pred, sens_maps, dim: int = 0) -> torch.Tensor:
+    spatial_pred = torch.view_as_complex(spatial_pred)
+    sens_maps = torch.view_as_complex(sens_maps)
+    
+    numerator = torch.sum(spatial_pred * torch.conj(sens_maps), dim=dim)
+    denominator = torch.sqrt(
+        torch.sum(torch.square(torch.abs(sens_maps)), dim=dim)
+    )
+    res = numerator / denominator
+    res = torch.abs(res)
+    return res
+
+
+def rss_complex(data: torch.Tensor, dim: int = 0) -> torch.Tensor:
+    """
+    Compute the Root Sum of Squares (RSS) for complex inputs.
+
+    The RSS is computed assuming that `dim` is the coil dimension.
+
+    Parameters
+    ----------
+    data : torch.Tensor
+        The input tensor containing complex values.
+    dim : int, optional
+        The dimension along which to apply the RSS transform (default is 0).
+
+    Returns
+    -------
+    torch.Tensor
+        The computed RSS value for complex inputs.
+    """
+    return torch.sqrt(fastmri.complex_abs_sq(data).sum(dim))

fastmri/datasets.py

@@ -0,0 +1,583 @@
+import random
+import xml.etree.ElementTree as etree
+from pathlib import Path
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    Literal,
+    NamedTuple,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
+)
+
+import h5py
+import lmdb
+import numpy as np
+import torch
+import yaml
+import sigpy as sp
+import pandas as pd
+
+import fastmri
+import fastmri.transforms as T
+
+
+class RawSample(NamedTuple):
+    fname: Path
+    slice_num: int
+    metadata: Dict[str, Any]
+
+
+class SliceSample(NamedTuple):
+    masked_kspace: torch.Tensor
+    mask: torch.Tensor
+    num_low_frequencies: int
+    target: torch.Tensor
+    max_value: float
+    # attrs: Dict[str, Any]
+    fname: str
+    slice_num: int
+
+class SliceSampleMVUE(NamedTuple):
+    masked_kspace: torch.Tensor
+    mask: torch.Tensor
+    num_low_frequencies: int
+    target: torch.Tensor
+    rss: torch.Tensor
+    max_value: float
+    # attrs: Dict[str, Any]
+    fname: str
+    slice_num: int
+
+def et_query(
+    root: etree.Element,
+    qlist: Sequence[str],
+    namespace: str = "http://www.ismrm.org/ISMRMRD",
+) -> str:
+    """
+    Query an XML document using ElementTree.
+
+    This function allows querying an XML document by specifying a root and a list of nested queries.
+    It supports optional XML namespaces.
+
+    Parameters
+    ----------
+    root : ElementTree.Element
+        The root element of the XML to search through.
+    qlist : list of str
+        A list of strings for nested searches, e.g., ["Encoding", "matrixSize"].
+    namespace : str, optional
+        An optional XML namespace to prepend to the query (default is None).
+
+    Returns
+    -------
+    str
+        The retrieved data as a string.
+    """
+
+    s = "."
+    prefix = "ismrmrd_namespace"
+
+    ns = {prefix: namespace}
+
+    for el in qlist:
+        s = s + f"//{prefix}:{el}"
+
+    value = root.find(s, ns)
+    if value is None:
+        raise RuntimeError("Element not found")
+
+    return str(value.text)
+
+
+class SliceDataset(torch.utils.data.Dataset):
+    """
+    A simplified PyTorch Dataset that provides access to multicoil MR image
+    slices from the fastMRI dataset.
+    """
+
+    def __init__(
+        self,
+        # root: Optional[Path | str],
+        body_part: Literal["knee", "brain"],
+        partition: Literal["train", "val", "test"],
+        mask_fns: Optional[List[Callable]] = None,
+        sample_rate: float = 1.0,
+        complex: bool = False,
+        crop_shape: Tuple[int, int] = (320, 320),
+        slug: str = "",
+        contrast: Optional[Literal["T1", "T2"]] = None,
+        coils: Optional[int] = None,
+    ):
+        """
+        Initializes the fastMRI multi-coil challenge dataset.
+
+        Samples are individual 2D slices taken from k-space volume data.
+
+        Parameters
+        ----------
+        body_part : {'knee', 'brain'}
+            The body part to analyze.
+        partition : {'train', 'val', 'test'}
+            The data partition type.
+        mask_fns : list of callable, optional
+            A list of masking functions to apply to samples.
+            If multiple are given, a mask is randomly chosen for each sample.
+        sample_rate : float, optional
+            Fraction of data to sample, by default 1.0.
+        complex : bool, optional
+            Whether the $k$-space data should return complex-valued, by default False.
+            If True, kspace values will be complex.
+            If False, kspace values will be real (shape, 2).
+        crop_shape : tuple of two ints, optional
+            The shape to center crop the k-space data, by default (320, 320).
+        slug : string
+            dataset slug name
+        contrast :  {'T1', 'T2'}
+            If partition is brain, the contrast of images to use.
+        """
+
+        with open("fastmri.yaml", "r") as file:
+            config = yaml.safe_load(file)
+        self.contrast = contrast
+        self.slug = slug
+        self.partition = partition
+        self.body_part = body_part
+        self.root = (
+            Path(config.get(f"{body_part}_path")) / f"multicoil_{partition}"
+        )
+        self.mask_fns = mask_fns
+        self.sample_rate = sample_rate
+        self.raw_samples: List[RawSample] = self._load_samples()
+        self.complex = complex
+        self.crop_shape = crop_shape
+        self.coils = coils
+
+    def _load_samples(self):
+        # Gather all files in the root directory
+        if self.body_part == "brain" and self.contrast:
+            files = list(self.root.glob(f"*{self.contrast}*.h5"))
+        else:
+            files = list(self.root.glob("*.h5"))
+        raw_samples = []
+
+        # Load and process metadata from each file
+        for fname in sorted(files):
+            with h5py.File(fname, "r") as hf:
+                metadata, num_slices = self._retrieve_metadata(fname)
+
+                # Collect samples for each slice, discard first c slices, and last c slices
+                c = 6
+                for slice_num in range(num_slices):
+                    if c <= slice_num <= num_slices - c - 1:
+                        raw_samples.append(
+                            RawSample(fname, slice_num, metadata)
+                        )
+
+        # Subsample if desired
+        if self.sample_rate < 1.0:
+            raw_samples = random.sample(
+                raw_samples, int(len(raw_samples) * self.sample_rate)
+            )
+
+        return raw_samples
+
+    def _retrieve_metadata(self, fname):
+        with h5py.File(fname, "r") as hf:
+            et_root = etree.fromstring(hf["ismrmrd_header"][()])
+
+            enc = ["encoding", "encodedSpace", "matrixSize"]
+            enc_size = (
+                int(et_query(et_root, enc + ["x"])),
+                int(et_query(et_root, enc + ["y"])),
+                int(et_query(et_root, enc + ["z"])),
+            )
+            rec = ["encoding", "reconSpace", "matrixSize"]
+            recon_size = (
+                int(et_query(et_root, rec + ["x"])),
+                int(et_query(et_root, rec + ["y"])),
+                int(et_query(et_root, rec + ["z"])),
+            )
+
+            lims = ["encoding", "encodingLimits", "kspace_encoding_step_1"]
+            enc_limits_center = int(et_query(et_root, lims + ["center"]))
+            enc_limits_max = int(et_query(et_root, lims + ["maximum"])) + 1
+
+            padding_left = enc_size[1] // 2 - enc_limits_center
+            padding_right = padding_left + enc_limits_max
+
+            num_slices = hf["kspace"].shape[0]
+
+            metadata = {
+                "padding_left": padding_left,
+                "padding_right": padding_right,
+                "encoding_size": enc_size,
+                "recon_size": recon_size,
+                **hf.attrs,
+            }
+
+        return metadata, num_slices
+
+    def __len__(self):
+        return len(self.raw_samples)
+
+    def __getitem__(self, idx) -> SliceSample:
+        try:
+            raw_sample: RawSample = self.raw_samples[idx]
+            fname, slice_num, metadata = raw_sample
+
+            # load kspace and target
+            with h5py.File(fname, "r") as hf:
+                kspace = torch.tensor(hf["kspace"][()][slice_num])
+                if not self.complex:
+                    kspace = torch.view_as_real(kspace)
+                if self.coils:
+                    if kspace.shape[0] < self.coils:
+                        return None
+                    kspace = kspace[: self.coils, :, :, :]
+                target_key = (
+                    "reconstruction_rss"
+                    if self.partition in ["train", "val"]
+                    else "reconstruction_esc"
+                )
+                target = hf.get(target_key, None)
+                if target is not None:
+                    target = torch.tensor(target[()][slice_num])
+                if self.body_part == "brain":
+                    target = T.center_crop(target, self.crop_shape)
+
+            # center crop to enable collating for batching
+            if self.complex:
+                # if complex, crop across dims: -2 and -1 (last 2)
+                raise NotImplementedError("Not implemented for complex native")
+            else:
+                # crop in image space, to not lose high-frequency information
+                image = fastmri.ifft2c(kspace)
+                image_cropped = T.complex_center_crop(image, self.crop_shape)
+                kspace = fastmri.fft2c(image_cropped)
+
+            # apply transform mask if there is one
+            if self.mask_fns:
+                # choose a random mask
+                mask_fn = random.choice(self.mask_fns)
+                kspace, mask, num_low_frequencies = T.apply_mask(
+                    kspace,
+                    mask_fn,
+                    # seed=seed,
+                )
+                mask = mask.bool()
+            else:
+                mask = torch.ones_like(kspace, dtype=torch.bool)
+                num_low_frequencies = 0
+            sample = SliceSample(
+                kspace,
+                mask,
+                num_low_frequencies,
+                target,
+                metadata["max"],
+                fname.name,
+                slice_num,
+            )
+            return sample
+        except:
+            return None
+
+
+class SliceDatasetLMDB(torch.utils.data.Dataset):
+    """
+    A simplified PyTorch Dataset that provides access to multicoil MR image
+    slices from the fastMRI dataset. Loads from LMDB saved samples.
+    """
+
+    def __init__(
+        self,
+        body_part: Literal["knee", "brain"],
+        partition: Literal["train", "val", "test"],
+        root: Optional[Path | str] = None,
+        mask_fns: Optional[List[Callable]] = None,
+        sample_rate: float = 1.0,
+        complex: bool = False,
+        crop_shape: Tuple[int, int] = (320, 320),
+        slug: str = "",
+        coils: int = 15,
+    ):
+        """
+        Initializes the fastMRI multi-coil challenge dataset.
+
+        Samples are individual 2D slices taken from k-space volume data.
+
+        Parameters
+        ----------
+        body_part : {'knee', 'brain'}
+            The body part to analyze.
+        root : Path or str, optional
+            Root to lmdb dataset. If not provided, the root is automatically
+            loaded directly from fastmri.yaml config
+        partition : {'train', 'val', 'test'}
+            The data partition type.
+        mask_fns : list of callable, optional
+            A list of masking functions to apply to samples.
+            If multiple are given, a mask is randomly chosen for each sample.
+        sample_rate : float, optional
+            Fraction of data to sample, by default 1.0.
+        complex : bool, optional
+            Whether the $k$-space data should return complex-valued, by default False.
+            If True, kspace values will be complex.
+            If False, kspace values will be real (shape, 2).
+        crop_shape : tuple of two ints, optional
+            The shape to center crop the k-space data, by default (320, 320).
+        slug : string
+            dataset slug name
+        """
+
+        # set attrs
+        self.coils = coils
+        self.slug = slug
+        self.partition = partition
+        self.mask_fns = mask_fns
+        self.sample_rate = sample_rate
+        self.complex = complex
+        self.crop_shape = crop_shape
+
+        # load lmdb info
+        if root:
+            if isinstance(root, str):
+                root = Path(root)
+            assert root.exists(), "Provided root doesn't exist."
+            self.root = root
+        else:
+            with open("fastmri.yaml", "r") as file:
+                config = yaml.safe_load(file)
+            self.root = Path(config["lmdb"][f"{body_part}_{partition}_path"])
+        self.meta = np.load(self.root / "meta.npy")
+        self.kspace_env = lmdb.open(
+            str(self.root / "kspace"),
+            readonly=True,
+            lock=False,
+            create=False,
+        )
+        self.kspace_txn = self.kspace_env.begin(write=False)
+        self.rss_env = lmdb.open(
+            str(self.root / "rss"),
+            readonly=True,
+            lock=False,
+            create=False,
+        )
+        self.rss_txn = self.rss_env.begin(write=False)
+        self.length = self.kspace_txn.stat()["entries"]
+
+    def __len__(self):
+        return int(self.sample_rate * self.length)
+
+    def __getitem__(self, idx) -> SliceSample:
+        idx_key = str(idx).encode("utf-8")
+
+        # load sample data
+        kspace = torch.from_numpy(
+            np.frombuffer(self.kspace_txn.get(idx_key), dtype=np.float32)
+            .reshape(self.coils, 320, 320, 2)
+            .copy()
+        )
+        rss = torch.from_numpy(
+            np.frombuffer(self.rss_txn.get(idx_key), dtype=np.float32)
+            .reshape(320, 320)
+            .copy()
+        )
+
+        # crop in image space, to not lose high-frequency information
+        if self.crop_shape and self.crop_shape != (320, 320):
+            image = fastmri.ifft2c(kspace)
+            image_cropped = T.complex_center_crop(image, self.crop_shape)
+            kspace = fastmri.fft2c(image_cropped)
+            rss = T.center_crop(rss, self.crop_shape)
+
+        # load and apply mask
+        if self.mask_fns:
+            # choose a random mask
+            mask_fn = random.choice(self.mask_fns)
+            kspace, mask, num_low_frequencies = T.apply_mask(
+                kspace,
+                mask_fn,  # type: ignore
+            )
+            mask = mask.bool()
+        else:
+            mask = torch.ones_like(kspace, dtype=torch.bool)
+            num_low_frequencies = 0
+
+        # load metadata
+        fname, slice_num, max_value = self.meta[idx]
+        fname = str(fname)
+        slice_num = int(slice_num)
+        max_value = float(max_value)
+
+        return SliceSample(
+            kspace,
+            mask,
+            num_low_frequencies,
+            rss,
+            max_value,
+            fname,
+            slice_num,
+        )
+
+
+class SliceDatasetLMDB_MVUE(torch.utils.data.Dataset):
+    """
+    Loads from LMDB brain saved samples.
+
+    Modified to have MVUE targets
+    """
+
+    def __init__(
+        self,
+        root: Path | str,
+        mask_fns: Optional[List[Callable]] = None,
+        sample_rate: float = 1.0,
+        crop_shape: Tuple[int, int] = (320, 320),
+        slug: str = "",
+        coils: int = 15,
+    ):
+
+        # set attrs
+        self.coils = coils
+        self.slug = slug
+        self.mask_fns = mask_fns
+        self.sample_rate = sample_rate
+        self.complex = complex
+        self.crop_shape = crop_shape
+
+        # load lmdb info
+        if isinstance(root, str):
+            root = Path(root)
+        assert root.exists(), "Provided root doesn't exist."
+        self.root = root
+        self.meta = np.load(self.root / "meta.npy")
+        self.mapping = pd.read_csv("brain_mvue_map.csv")
+        self.kspace_env = lmdb.open(
+            str(self.root / "kspace"),
+            readonly=True,
+            lock=False,
+            create=False,
+        )
+        self.kspace_txn = self.kspace_env.begin(write=False)
+        self.rss_env = lmdb.open(
+            str(self.root / "rss"),
+            readonly=True,
+            lock=False,
+            create=False,
+        )
+        self.rss_txn = self.rss_env.begin(write=False)
+       
+        # ray mvue dataset 
+        self.mvue_env = lmdb.open(
+            str("/pscratch/sd/p/peterwg/datasets/raytemp"),
+            readonly=True,
+            lock=False,
+            create=False,
+        )
+        self.mvue_txn = self.mvue_env.begin(write=False)
+        
+        self.length = len(self.mapping) 
+        # self.length = self.kspace_txn.stat()["entries"]
+
+    def __len__(self):
+        return int(self.sample_rate * self.length)
+
+    def __getitem__(self, idx) -> SliceSampleMVUE:
+        # ray's index: 0-n
+        ray_idx = idx
+        # my index: lookup(ray index)
+        idx = int(self.mapping.iloc[ray_idx].my_index)
+        ray_idx_key = str(ray_idx).encode("utf-8")
+        idx_key = str(idx).encode("utf-8")
+
+        # load sample data
+        kspace = torch.from_numpy(
+            np.frombuffer(self.kspace_txn.get(idx_key), dtype=np.float32)
+            .reshape(self.coils, 320, 320, 2)
+            .copy()
+        )
+
+        # mvue_target = np.sum(
+        #     sp.ifft(kspace, axes=(-1, -2)) * np.conj(s_maps), axis=1
+        # ) / np.sqrt(np.sum(np.square(np.abs(s_maps)), axis=1))
+        rss = torch.from_numpy(
+            np.frombuffer(self.rss_txn.get(idx_key), dtype=np.float32)
+            .reshape(320, 320)
+            .copy()
+        )
+        
+        # load mvue from ray dataset 
+        mvue = torch.from_numpy(
+            np.frombuffer(self.mvue_txn.get(ray_idx_key), dtype=np.complex64)
+            .reshape(320, 320)
+            .copy()
+        )
+        mvue = torch.abs(mvue)
+
+        # crop in image space, to not lose high-frequency information
+        if self.crop_shape and self.crop_shape != (320, 320):
+            image = fastmri.ifft2c(kspace)
+            image_cropped = T.complex_center_crop(image, self.crop_shape)
+            kspace = fastmri.fft2c(image_cropped)
+            rss = T.center_crop(rss, self.crop_shape)
+
+        # load and apply mask
+        if self.mask_fns:
+            # choose a random mask
+            mask_fn = random.choice(self.mask_fns)
+            kspace, mask, num_low_frequencies = T.apply_mask(
+                kspace,
+                mask_fn,  # type: ignore
+            )
+            mask = mask.bool()
+        else:
+            mask = torch.ones_like(kspace, dtype=torch.bool)
+            num_low_frequencies = 0
+
+        # load metadata
+        fname, slice_num, max_value = self.meta[idx]
+        fname = str(fname)
+        slice_num = int(slice_num)
+        max_value = float(max_value)
+
+        return SliceSampleMVUE(
+            kspace,
+            mask,
+            num_low_frequencies,
+            mvue,
+            rss,
+            max_value,
+            fname,
+            slice_num,
+        )
+
+
+# d = SliceDatasetLMDB("knee", "val", None, 1, True, (320, 320), "testdataset")
+# print(len(d))
+# breakpoint()
+
+# ds = SuperSliceDatasetLMDB(
+#     "brain",  # body_part
+#     "val",  # partition
+#     None,  # root
+#     None,  # mask_fns
+#     1.0,  # sample_rate
+#     True,  # complex
+#     (320, 320),  # crop_shape
+#     "test-superres",  # slug
+#     coils=16,  # coils
+# )
+# breakpoint()
+
+# d = SliceDataset("brain", "train", None, contrast="T2")
+
+
+# # TESTING MVUE
+# d = SliceDatasetLMDB_MVUE("/pscratch/sd/p/peterwg/datasets/mri_brain_train_lmdb", coils=16)
+# x = d[0]
+# d = SliceDatasetLMDB_MVUE("/pscratch/sd/p/peterwg/datasets/raytemp/", coils=16)

fastmri/evaluate.py

@@ -0,0 +1,174 @@
+"""
+Copyright (c) Facebook, Inc. and its affiliates.
+
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+import argparse
+import pathlib
+from argparse import ArgumentParser
+from typing import Optional
+
+import h5py
+import numpy as np
+from runstats import Statistics
+from skimage.metrics import peak_signal_noise_ratio, structural_similarity
+
+from fastmri import transforms
+
+
+def mse(gt: np.ndarray, pred: np.ndarray) -> np.ndarray:
+    """Compute Mean Squared Error (MSE)"""
+    return np.mean((gt - pred) ** 2)
+
+
+def nmse(gt: np.ndarray, pred: np.ndarray) -> np.ndarray:
+    """Compute Normalized Mean Squared Error (NMSE)"""
+    return np.array(np.linalg.norm(gt - pred) ** 2 / np.linalg.norm(gt) ** 2)
+
+
+def psnr(
+    gt: np.ndarray, pred: np.ndarray, maxval: Optional[float] = None
+) -> np.ndarray:
+    """Compute Peak Signal to Noise Ratio metric (PSNR)"""
+    if maxval is None:
+        maxval = gt.max()
+    return peak_signal_noise_ratio(gt, pred, data_range=maxval)
+
+
+def ssim(
+    gt: np.ndarray, pred: np.ndarray, maxval: Optional[float] = None
+) -> np.ndarray:
+    """Compute Structural Similarity Index Metric (SSIM)"""
+    if not gt.ndim == 3:
+        raise ValueError("Unexpected number of dimensions in ground truth.")
+    if not gt.ndim == pred.ndim:
+        raise ValueError("Ground truth dimensions does not match pred.")
+
+    maxval = gt.max() if maxval is None else maxval
+
+    ssim = np.array([0])
+    for slice_num in range(gt.shape[0]):
+        ssim = ssim + structural_similarity(
+            gt[slice_num], pred[slice_num], data_range=maxval
+        )
+
+    return ssim / gt.shape[0]
+
+
+METRIC_FUNCS = dict(
+    MSE=mse,
+    NMSE=nmse,
+    PSNR=psnr,
+    SSIM=ssim,
+)
+
+
+class Metrics:
+    """
+    Maintains running statistics for a given collection of metrics.
+    """
+
+    def __init__(self, metric_funcs):
+        """
+        Parameters
+        ----------
+        metric_funcs : dict
+            A dictionary where the keys are metric names (as strings) and the values
+            are Python functions for evaluating the corresponding metrics.
+        """
+
+        self.metrics = {metric: Statistics() for metric in metric_funcs}
+
+    def push(self, target, recons):
+        for metric, func in METRIC_FUNCS.items():
+            self.metrics[metric].push(func(target, recons))
+
+    def means(self):
+        return {metric: stat.mean() for metric, stat in self.metrics.items()}
+
+    def stddevs(self):
+        return {metric: stat.stddev() for metric, stat in self.metrics.items()}
+
+    def __repr__(self):
+        means = self.means()
+        stddevs = self.stddevs()
+        metric_names = sorted(list(means))
+        return " ".join(
+            f"{name} = {means[name]:.4g} +/- {2 * stddevs[name]:.4g}"
+            for name in metric_names
+        )
+
+
+def evaluate(args, recons_key):
+    metrics = Metrics(METRIC_FUNCS)
+
+    for tgt_file in args.target_path.iterdir():
+        with h5py.File(tgt_file, "r") as target, h5py.File(
+            args.predictions_path / tgt_file.name, "r"
+        ) as recons:
+            if args.acquisition and args.acquisition != target.attrs["acquisition"]:
+                continue
+
+            if args.acceleration and target.attrs["acceleration"] != args.acceleration:
+                continue
+
+            target = target[recons_key][()]
+            recons = recons["reconstruction"][()]
+            target = transforms.center_crop(
+                target, (target.shape[-1], target.shape[-1])
+            )
+            recons = transforms.center_crop(
+                recons, (target.shape[-1], target.shape[-1])
+            )
+            metrics.push(target, recons)
+
+    return metrics
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument(
+        "--target-path",
+        type=pathlib.Path,
+        required=True,
+        help="Path to the ground truth data",
+    )
+    parser.add_argument(
+        "--predictions-path",
+        type=pathlib.Path,
+        required=True,
+        help="Path to reconstructions",
+    )
+    parser.add_argument(
+        "--challenge",
+        choices=["singlecoil", "multicoil"],
+        required=True,
+        help="Which challenge",
+    )
+    parser.add_argument("--acceleration", type=int, default=None)
+    parser.add_argument(
+        "--acquisition",
+        choices=[
+            "CORPD_FBK",
+            "CORPDFS_FBK",
+            "AXT1",
+            "AXT1PRE",
+            "AXT1POST",
+            "AXT2",
+            "AXFLAIR",
+        ],
+        default=None,
+        help=(
+            "If set, only volumes of the specified acquisition type are used "
+            "for evaluation. By default, all volumes are included."
+        ),
+    )
+    args = parser.parse_args()
+
+    recons_key = (
+        "reconstruction_rss" if args.challenge == "multicoil" else "reconstruction_esc"
+    )
+    metrics = evaluate(args, recons_key)
+    print(metrics)

fastmri/fftc.py

@@ -0,0 +1,203 @@
+"""
+Copyright (c) Facebook, Inc. and its affiliates.
+
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+from typing import List, Optional
+
+import torch
+import torch.fft
+
+
+def fft2c_new(data: torch.Tensor, norm: str = "ortho") -> torch.Tensor:
+    """
+    Apply a centered 2-dimensional Fast Fourier Transform (FFT).
+
+    Parameters
+    ----------
+    data : torch.Tensor
+        Complex-valued input data containing at least 3 dimensions.
+        Dimensions -3 and -2 are spatial dimensions, and dimension -1 has size 2.
+        All other dimensions are assumed to be batch dimensions.
+    norm : str
+        Normalization mode. Refer to `torch.fft.fft` for details on normalization options.
+
+    Returns
+    -------
+    torch.Tensor
+        The FFT of the input data.
+    """
+
+    if not data.shape[-1] == 2:
+        raise ValueError("Tensor does not have separate complex dim.")
+
+    data = ifftshift(data, dim=[-3, -2])
+    data = torch.view_as_real(
+        torch.fft.fftn(  # type: ignore
+            torch.view_as_complex(data), dim=(-2, -1), norm=norm
+        )
+    )
+    data = fftshift(data, dim=[-3, -2])
+
+    return data
+
+
+def ifft2c_new(data: torch.Tensor, norm: str = "ortho") -> torch.Tensor:
+    """
+    Apply a centered 2-dimensional Inverse Fast Fourier Transform (IFFT).
+
+    Parameters
+    ----------
+    data : torch.Tensor
+        Complex-valued input data containing at least 3 dimensions.
+        Dimensions -3 and -2 are spatial dimensions, and dimension -1 has size 2.
+        All other dimensions are assumed to be batch dimensions.
+    norm : str
+        Normalization mode. Refer to `torch.fft.ifft` for details on normalization options.
+
+    Returns
+    -------
+    torch.Tensor
+        The IFFT of the input data.
+    """
+
+    if not data.shape[-1] == 2:
+        raise ValueError("Tensor does not have separate complex dim.")
+
+    data = ifftshift(data, dim=[-3, -2])
+    data = torch.view_as_real(
+        torch.fft.ifftn(  # type: ignore
+            torch.view_as_complex(data), dim=(-2, -1), norm=norm
+        )
+    )
+    data = fftshift(data, dim=[-3, -2])
+
+    return data
+
+
+# Helper functions
+
+
+def roll_one_dim(x: torch.Tensor, shift: int, dim: int) -> torch.Tensor:
+    """
+    Roll a PyTorch tensor along a specified dimension.
+
+    This function is similar to `torch.roll` but operates on a single dimension.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        The input tensor to be rolled.
+    shift : int
+        Amount to roll.
+    dim : int
+        The dimension along which to roll the tensor.
+
+    Returns
+    -------
+    torch.Tensor
+        A tensor with the same shape as `x`, but rolled along the specified dimension.
+    """
+
+    shift = shift % x.size(dim)
+    if shift == 0:
+        return x
+
+    left = x.narrow(dim, 0, x.size(dim) - shift)
+    right = x.narrow(dim, x.size(dim) - shift, shift)
+
+    return torch.cat((right, left), dim=dim)
+
+
+def roll(
+    x: torch.Tensor,
+    shift: List[int],
+    dim: List[int],
+) -> torch.Tensor:
+    """
+    Similar to np.roll but applies to PyTorch Tensors.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        A PyTorch tensor.
+    shift : int
+        Amount to roll.
+    dim : int
+        Which dimension to roll.
+
+    Returns
+    -------
+    torch.Tensor
+        Rolled version of x.
+    """
+
+    if len(shift) != len(dim):
+        raise ValueError("len(shift) must match len(dim)")
+
+    for s, d in zip(shift, dim):
+        x = roll_one_dim(x, s, d)
+
+    return x
+
+
+def fftshift(x: torch.Tensor, dim: Optional[List[int]] = None) -> torch.Tensor:
+    """
+    Similar to np.fft.fftshift but applies to PyTorch Tensors.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        A PyTorch tensor.
+    dim : list of int, optional
+        Which dimension to apply fftshift. If None, the shift is applied to all dimensions (default is None).
+
+    Returns
+    -------
+    torch.Tensor
+        fftshifted version of x.
+    """
+    if dim is None:
+        # this weird code is necessary for torch.jit.script typing
+        dim = [0] * (x.dim())
+        for i in range(1, x.dim()):
+            dim[i] = i
+
+    # also necessary for torch.jit.script
+    shift = [0] * len(dim)
+    for i, dim_num in enumerate(dim):
+        shift[i] = x.shape[dim_num] // 2
+
+    return roll(x, shift, dim)
+
+
+def ifftshift(x: torch.Tensor, dim: Optional[List[int]] = None) -> torch.Tensor:
+    """
+    Similar to np.fft.ifftshift but applies to PyTorch Tensors.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        A PyTorch tensor.
+    dim : list of int, optional
+        Which dimension to apply ifftshift. If None, the shift is applied to all dimensions (default is None).
+
+    Returns
+    -------
+    torch.Tensor
+        ifftshifted version of x.
+    """
+    if dim is None:
+        # this weird code is necessary for torch.jit.script typing
+        dim = [0] * (x.dim())
+        for i in range(1, x.dim()):
+            dim[i] = i
+
+    # also necessary for torch.jit.script
+    shift = [0] * len(dim)
+    for i, dim_num in enumerate(dim):
+        shift[i] = (x.shape[dim_num] + 1) // 2
+
+    return roll(x, shift, dim)

fastmri/losses.py

@@ -0,0 +1,91 @@
+"""
+Copyright (c) Facebook, Inc. and its affiliates.
+
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class SSIMLoss(nn.Module):
+    """
+    SSIM loss module.
+    """
+
+    def __init__(self, win_size: int = 7, k1: float = 0.01, k2: float = 0.03):
+        """
+        Initialize the Losses class.
+
+        Parameters
+        ----------
+        win_size : int, optional
+            Window size for SSIM calculation.
+        k1 : float, optional
+            k1 parameter for SSIM calculation.
+        k2 : float, optional
+            k2 parameter for SSIM calculation.
+        """
+        super().__init__()
+        self.win_size = win_size
+        self.k1, self.k2 = k1, k2
+        self.register_buffer("w", torch.ones(1, 1, win_size, win_size) / win_size**2)
+        NP = win_size**2
+        self.cov_norm = NP / (NP - 1)
+
+    def forward(
+        self,
+        X: torch.Tensor,
+        Y: torch.Tensor,
+        data_range: torch.Tensor,
+        reduced: bool = True,
+    ):
+        assert isinstance(self.w, torch.Tensor)
+
+        data_range = data_range[:, None, None, None].to(X.device)
+        C1 = (self.k1 * data_range) ** 2
+        C2 = (self.k2 * data_range) ** 2
+
+        # Compute means
+        ux = F.conv2d(X, self.w)
+        uy = F.conv2d(Y, self.w)
+
+        # Compute variances
+        uxx = F.conv2d(X * X, self.w)
+        uyy = F.conv2d(Y * Y, self.w)
+        uxy = F.conv2d(X * Y, self.w)
+
+        # Compute covariances
+        vx = self.cov_norm * (uxx - ux * ux)
+        vy = self.cov_norm * (uyy - uy * uy)
+        vxy = self.cov_norm * (uxy - ux * uy)
+
+        # Compute SSIM components
+        A1, A2 = 2 * ux * uy + C1, 2 * vxy + C2
+        B1, B2 = ux**2 + uy**2 + C1, vx + vy + C2
+        D = B1 * B2
+        S = (A1 * A2) / D
+
+        if reduced:
+            return 1 - S.mean()
+        else:
+            return 1 - S
+
+
+if __name__ == "__main__":
+    # Example usage
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # Create the SSIMLoss module and move it to the GPU
+    ssim_loss = SSIMLoss().to(device)
+
+    # Create example tensors and move them to the GPU
+    X = torch.randn(4, 1, 256, 256).to(device)
+    Y = torch.randn(4, 1, 256, 256).to(device)
+    data_range = torch.rand(4).to(device)
+
+    # Compute the loss
+    loss = ssim_loss(X, Y, data_range)
+    print(loss)

fastmri/math_utils.py

@@ -0,0 +1,121 @@
+"""
+Copyright (c) Facebook, Inc. and its affiliates.
+
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+import numpy as np
+import torch
+
+
+def complex_mul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    """
+    Complex multiplication.
+
+    Multiplies two complex tensors assuming that they are both stored as
+    real arrays with the last dimension being the complex dimension.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        A PyTorch tensor with the last dimension of size 2.
+    y : torch.Tensor
+        A PyTorch tensor with the last dimension of size 2.
+
+    Returns
+    -------
+    torch.Tensor
+        A PyTorch tensor with the last dimension of size 2, representing
+        the result of the complex multiplication.
+    """
+    if not x.shape[-1] == y.shape[-1] == 2:
+        raise ValueError("Tensors do not have separate complex dim.")
+
+    re = x[..., 0] * y[..., 0] - x[..., 1] * y[..., 1]
+    im = x[..., 0] * y[..., 1] + x[..., 1] * y[..., 0]
+
+    return torch.stack((re, im), dim=-1)
+
+
+def complex_conj(x: torch.Tensor) -> torch.Tensor:
+    """
+    Complex conjugate.
+
+    Applies the complex conjugate assuming that the input array has the
+    last dimension as the complex dimension.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        A PyTorch tensor with the last dimension of size 2.
+
+    Returns
+    -------
+    torch.Tensor
+        A PyTorch tensor with the last dimension of size 2, representing
+        the complex conjugate of the input tensor.
+    """
+    if not x.shape[-1] == 2:
+        raise ValueError("Tensor does not have separate complex dim.")
+
+    return torch.stack((x[..., 0], -x[..., 1]), dim=-1)
+
+
+def complex_abs(data: torch.Tensor) -> torch.Tensor:
+    """
+    Compute the absolute value of a complex-valued input tensor.
+
+    Parameters
+    ----------
+    data : torch.Tensor
+        A complex-valued tensor, where the size of the final dimension
+        should be 2.
+
+    Returns
+    -------
+    torch.Tensor
+        Absolute value of the input tensor.
+    """
+    if not data.shape[-1] == 2:
+        raise ValueError("Tensor does not have separate complex dim.")
+
+    return (data**2).sum(dim=-1).sqrt()
+
+
+def complex_abs_sq(data: torch.Tensor) -> torch.Tensor:
+    """
+    Compute the squared absolute value of a complex tensor.
+
+    Parameters
+    ----------
+    data : torch.Tensor
+        A complex-valued tensor, where the size of the final dimension
+        should be 2.
+
+    Returns
+    -------
+    torch.Tensor
+        Squared absolute value of the input tensor.
+    """
+    if not data.shape[-1] == 2:
+        raise ValueError("Tensor does not have separate complex dim.")
+
+    return (data**2).sum(dim=-1)
+
+
+def tensor_to_complex_np(data: torch.Tensor) -> np.ndarray:
+    """
+    Convert a complex PyTorch tensor to a NumPy array.
+
+    Parameters
+    ----------
+    data : torch.Tensor
+        Input data to be converted to a NumPy array.
+
+    Returns
+    -------
+    np.ndarray
+        A complex NumPy array version of the input tensor.
+    """
+    return torch.view_as_complex(data).numpy()

fastmri/poisson_cache/poisson_16x.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:22199ef1c9b045b6f747e57c4effa9a6667cfce864773ee035df8c3d2a28138f
+size 819328

fastmri/poisson_cache/poisson_2x.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0d6908dbd90fda83085cc9e7d3ab35f7ed215ab3f735fd39023a091a9f1632df
+size 819328

fastmri/poisson_cache/poisson_32x.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7e316852431f100b1c5c6749b672dfbf2dffc4f23ba4d118eddc64956b8c22f4
+size 819328

fastmri/poisson_cache/poisson_4x.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4c9f0c9b2c3be534b7c94b8398c2aa66c8e634d343e9b45b842450137266cbc8
+size 819328

fastmri/poisson_cache/poisson_6x.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f7eeed6d470af6ef2b7594da388ea952dad44e74a33080a1bf59faf9e8973ca8
+size 819328

fastmri/poisson_cache/poisson_8x.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:28ce774dc182798a3fd4358c70cb86c81f936f12336fee104c93e89765727462
+size 819328

fastmri/subsample.py

@@ -0,0 +1,818 @@
+"""
+Copyright (c) Facebook, Inc. and its affiliates.
+
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+import os
+from typing import Dict, Optional, Sequence, Tuple, Union
+
+import numpy as np
+import torch
+import torch.distributions as D
+from networkx import center
+from sigpy.mri import poisson, radial, spiral
+
+
+class MaskFunc:
+    """
+    An object for GRAPPA-style sampling masks.
+
+    This crates a sampling mask that densely samples the center while
+    subsampling outer k-space regions based on the undersampling factor.
+
+    When called, ``MaskFunc`` uses internal functions create mask by 1)
+    creating a mask for the k-space center, 2) create a mask outside of the
+    k-space center, and 3) combining them into a total mask. The internals are
+    handled by ``sample_mask``, which calls ``calculate_center_mask`` for (1)
+    and ``calculate_acceleration_mask`` for (2). The combination is executed
+    in the ``MaskFunc`` ``__call__`` function.
+
+    If you would like to implement a new mask, simply subclass ``MaskFunc``
+    and overwrite the ``sample_mask`` logic. See examples in ``RandomMaskFunc``
+    and ``EquispacedMaskFunc``.
+    """
+
+    def __init__(
+        self,
+        center_fractions: Sequence[float],
+        accelerations: Sequence[int],
+        allow_any_combination: bool = False,
+        seed: Optional[int] = None,
+    ):
+        """
+        Args:
+            center_fractions: Fraction of low-frequency columns to be retained.
+                If multiple values are provided, then one of these numbers is
+                chosen uniformly each time.
+            accelerations: Amount of under-sampling. This should have the same
+                length as center_fractions. If multiple values are provided,
+                then one of these is chosen uniformly each time.
+            allow_any_combination: Whether to allow cross combinations of
+                elements from ``center_fractions`` and ``accelerations``.
+            seed: Seed for starting the internal random number generator of the
+                ``MaskFunc``.
+        """
+        if (
+            len(center_fractions) != len(accelerations)
+            and not allow_any_combination
+        ):
+            raise ValueError(
+                "Number of center fractions should match number of"
+                " accelerations if allow_any_combination is False."
+            )
+
+        self.center_fractions = center_fractions
+        self.accelerations = accelerations
+        self.allow_any_combination = allow_any_combination
+        self.rng = np.random.RandomState(seed)
+
+    def __call__(
+        self,
+        shape: Sequence[int],
+        offset: Optional[int] = None,
+        seed: Optional[Union[int, Tuple[int, ...]]] = None,
+    ) -> Tuple[torch.Tensor, int]:
+        """
+        Sample and return a k-space mask.
+
+        Args:
+            shape: Shape of k-space.
+            offset: Offset from 0 to begin mask (for equispaced masks). If no
+                offset is given, then one is selected randomly.
+            seed: Seed for random number generator for reproducibility.
+
+        Returns:
+            A 2-tuple containing 1) the k-space mask and 2) the number of
+            center frequency lines.
+        """
+        if len(shape) < 3:
+            raise ValueError("Shape should have 3 or more dimensions")
+
+        center_mask, accel_mask, num_low_frequencies = self.sample_mask(
+            shape, offset
+        )
+        # combine masks together
+        return torch.max(center_mask, accel_mask), num_low_frequencies
+
+    def sample_mask(
+        self,
+        shape: Sequence[int],
+        offset: Optional[int],
+    ) -> Tuple[torch.Tensor, torch.Tensor, int]:
+        """
+        Sample a new k-space mask.
+
+        This function samples and returns two components of a k-space mask: 1)
+        the center mask (e.g., for sensitivity map calculation) and 2) the
+        acceleration mask (for the edge of k-space). Both of these masks, as
+        well as the integer of low frequency samples, are returned.
+
+        Args:
+            shape: Shape of the k-space to subsample.
+            offset: Offset from 0 to begin mask (for equispaced masks).
+
+        Returns:
+            A 3-tuple contaiing 1) the mask for the center of k-space, 2) the
+            mask for the high frequencies of k-space, and 3) the integer count
+            of low frequency samples.
+        """
+        num_cols = shape[-2]
+        center_fraction, acceleration = self.choose_acceleration()
+        num_low_frequencies = round(num_cols * center_fraction)
+        center_mask = self.reshape_mask(
+            self.calculate_center_mask(shape, num_low_frequencies), shape
+        )
+        acceleration_mask = self.reshape_mask(
+            self.calculate_acceleration_mask(
+                num_cols, acceleration, offset, num_low_frequencies
+            ),
+            shape,
+        )
+        return center_mask, acceleration_mask, num_low_frequencies
+
+    def reshape_mask(
+        self, mask: torch.Tensor, shape: Sequence[int]
+    ) -> torch.Tensor:
+        """Reshape mask to desired output shape."""
+        if len(mask.shape) == 1:
+            mask = torch.tensor(mask)
+            mask_num_freqs = len(mask)
+            mask = mask.reshape(1, 1, mask_num_freqs, 1)
+            mask = mask.expand(shape)
+        return mask.expand(shape)
+
+    def reshape_mask_old(
+        self, mask: np.ndarray, shape: Sequence[int]
+    ) -> torch.Tensor:
+        """Reshape mask to desired output shape."""
+        num_cols = shape[-2]
+        mask_shape = [1 for s in shape]
+        mask_shape[-2] = num_cols
+
+        return torch.from_numpy(mask.reshape(*mask_shape).astype(np.float32))
+
+    def calculate_acceleration_mask(
+        self,
+        num_cols: int,
+        acceleration: int,
+        offset: Optional[int],
+        num_low_frequencies: int,
+    ) -> np.ndarray:
+        """
+        Produce mask for non-central acceleration lines.
+
+        Args:
+            num_cols: Number of columns of k-space (2D subsampling).
+            acceleration: Desired acceleration rate.
+            offset: Offset from 0 to begin masking (for equispaced masks).
+            num_low_frequencies: Integer count of low-frequency lines sampled.
+
+        Returns:
+            A mask for the high spatial frequencies of k-space.
+        """
+        raise NotImplementedError
+
+    def calculate_center_mask(
+        self, shape: Sequence[int], num_low_freqs: int
+    ) -> np.ndarray:
+        """
+        Build center mask based on number of low frequencies.
+
+        Args:
+            shape: Shape of k-space to mask.
+            num_low_freqs: Number of low-frequency lines to sample.
+
+        Returns:
+            A mask for hte low spatial frequencies of k-space.
+        """
+        num_cols = shape[-2]
+        mask = np.zeros(num_cols, dtype=np.float32)
+        pad = (num_cols - num_low_freqs + 1) // 2
+        mask[pad : pad + num_low_freqs] = 1
+        assert mask.sum() == num_low_freqs
+
+        return mask
+
+    def choose_acceleration(self):
+        """Choose acceleration based on class parameters."""
+        if self.allow_any_combination:
+            return self.rng.choice(self.center_fractions), self.rng.choice(
+                self.accelerations
+            )
+        else:
+            choice = self.rng.randint(len(self.center_fractions))
+            return self.center_fractions[choice], self.accelerations[choice]
+
+
+class RandomMaskFunc(MaskFunc):
+    """
+    Creates a random sub-sampling mask of a given shape.
+
+    The mask selects a subset of columns from the input k-space data. If the
+    k-space data has N columns, the mask picks out:
+        1. N_low_freqs = (N * center_fraction) columns in the center
+           corresponding to low-frequencies.
+        2. The other columns are selected uniformly at random with a
+        probability equal to: prob = (N / acceleration - N_low_freqs) /
+        (N - N_low_freqs). This ensures that the expected number of columns
+        selected is equal to (N / acceleration).
+
+    It is possible to use multiple center_fractions and accelerations, in which
+    case one possible (center_fraction, acceleration) is chosen uniformly at
+    random each time the ``RandomMaskFunc`` object is called.
+
+    For example, if accelerations = [4, 8] and center_fractions = [0.08, 0.04],
+    then there is a 50% probability that 4-fold acceleration with 8% center
+    fraction is selected and a 50% probability that 8-fold acceleration with 4%
+    center fraction is selected.
+    """
+
+    def calculate_acceleration_mask(
+        self,
+        num_cols: int,
+        acceleration: int,
+        offset: Optional[int],
+        num_low_frequencies: int,
+    ) -> np.ndarray:
+        prob = (num_cols / acceleration - num_low_frequencies) / (
+            num_cols - num_low_frequencies
+        )
+
+        return self.rng.uniform(size=num_cols) < prob
+
+
+class EquiSpacedMaskFunc(MaskFunc):
+    """
+    Sample data with equally-spaced k-space lines.
+
+    The lines are spaced exactly evenly, as is done in standard GRAPPA-style
+    acquisitions. This means that with a densely-sampled center,
+    ``acceleration`` will be greater than the true acceleration rate.
+    """
+
+    def calculate_acceleration_mask(
+        self,
+        num_cols: int,
+        acceleration: int,
+        offset: Optional[int],
+        num_low_frequencies: int,
+    ) -> np.ndarray:
+        """
+        Produce mask for non-central acceleration lines.
+
+        Args:
+            num_cols: Number of columns of k-space (2D subsampling).
+            acceleration: Desired acceleration rate.
+            offset: Offset from 0 to begin masking. If no offset is specified,
+                then one is selected randomly.
+            num_low_frequencies: Not used.
+
+        Returns:
+            A mask for the high spatial frequencies of k-space.
+        """
+        if offset is None:
+            offset = self.rng.randint(0, high=round(acceleration))
+
+        mask = np.zeros(num_cols, dtype=np.float32)
+        mask[offset::acceleration] = 1
+
+        return mask
+
+
+class EquispacedMaskFractionFunc(MaskFunc):
+    """
+    Equispaced mask with approximate acceleration matching.
+
+    The mask selects a subset of columns from the input k-space data. If the
+    k-space data has N columns, the mask picks out:
+        1. N_low_freqs = (N * center_fraction) columns in the center
+           corresponding to low-frequencies.
+        2. The other columns are selected with equal spacing at a proportion
+           that reaches the desired acceleration rate taking into consideration
+           the number of low frequencies. This ensures that the expected number
+           of columns selected is equal to (N / acceleration)
+
+    It is possible to use multiple center_fractions and accelerations, in which
+    case one possible (center_fraction, acceleration) is chosen uniformly at
+    random each time the EquispacedMaskFunc object is called.
+
+    Note that this function may not give equispaced samples (documented in
+    https://github.com/facebookresearch/fastMRI/issues/54), which will require
+    modifications to standard GRAPPA approaches. Nonetheless, this aspect of
+    the function has been preserved to match the public multicoil data.
+    """
+
+    def calculate_acceleration_mask(
+        self,
+        num_cols: int,
+        acceleration: int,
+        offset: Optional[int],
+        num_low_frequencies: int,
+    ) -> np.ndarray:
+        """
+        Produce mask for non-central acceleration lines.
+
+        Args:
+            num_cols: Number of columns of k-space (2D subsampling).
+            acceleration: Desired acceleration rate.
+            offset: Offset from 0 to begin masking. If no offset is specified,
+                then one is selected randomly.
+            num_low_frequencies: Number of low frequencies. Used to adjust mask
+                to exactly match the target acceleration.
+
+        Returns:
+            A mask for the high spatial frequencies of k-space.
+        """
+        # determine acceleration rate by adjusting for the number of low frequencies
+        adjusted_accel = (acceleration * (num_low_frequencies - num_cols)) / (
+            num_low_frequencies * acceleration - num_cols
+        )
+        if offset is None:
+            offset = self.rng.randint(0, high=round(adjusted_accel))
+
+        mask = np.zeros(num_cols, dtype=np.float32)
+        accel_samples = np.arange(offset, num_cols - 1, adjusted_accel)
+        accel_samples = np.around(accel_samples).astype(np.uint)
+        mask[accel_samples] = 1.0
+
+        return mask
+
+
+class MagicMaskFunc(MaskFunc):
+    """
+    Masking function for exploiting conjugate symmetry via offset-sampling.
+
+    This function applies the mask described in the following paper:
+
+    Defazio, A. (2019). Offset Sampling Improves Deep Learning based
+    Accelerated MRI Reconstructions by Exploiting Symmetry. arXiv preprint,
+    arXiv:1912.01101.
+
+    It is essentially an equispaced mask with an offset for the opposite site
+    of k-space. Since MRI images often exhibit approximate conjugate k-space
+    symmetry, this mask is generally more efficient than a standard equispaced
+    mask.
+
+    Similarly to ``EquispacedMaskFunc``, this mask will usually undereshoot the
+    target acceleration rate.
+    """
+
+    def calculate_acceleration_mask(
+        self,
+        num_cols: int,
+        acceleration: int,
+        offset: Optional[int],
+        num_low_frequencies: int,
+    ) -> np.ndarray:
+        """
+        Produce mask for non-central acceleration lines.
+
+        Args:
+            num_cols: Number of columns of k-space (2D subsampling).
+            acceleration: Desired acceleration rate.
+            offset: Offset from 0 to begin masking. If no offset is specified,
+                then one is selected randomly.
+            num_low_frequencies: Not used.
+
+        Returns:
+            A mask for the high spatial frequencies of k-space.
+        """
+        if offset is None:
+            offset = self.rng.randint(0, high=acceleration)
+
+        if offset % 2 == 0:
+            offset_pos = offset + 1
+            offset_neg = offset + 2
+        else:
+            offset_pos = offset - 1 + 3
+            offset_neg = offset - 1 + 0
+
+        poslen = (num_cols + 1) // 2
+        neglen = num_cols - (num_cols + 1) // 2
+        mask_positive = np.zeros(poslen, dtype=np.float32)
+        mask_negative = np.zeros(neglen, dtype=np.float32)
+
+        mask_positive[offset_pos::acceleration] = 1
+        mask_negative[offset_neg::acceleration] = 1
+        mask_negative = np.flip(mask_negative)
+
+        mask = np.concatenate((mask_positive, mask_negative))
+
+        return np.fft.fftshift(mask)  # shift mask and return
+
+
+class MagicMaskFractionFunc(MagicMaskFunc):
+    """
+    Masking function for exploiting conjugate symmetry via offset-sampling.
+
+    This function applies the mask described in the following paper:
+
+    Defazio, A. (2019). Offset Sampling Improves Deep Learning based
+    Accelerated MRI Reconstructions by Exploiting Symmetry. arXiv preprint,
+    arXiv:1912.01101.
+
+    It is essentially an equispaced mask with an offset for the opposite site
+    of k-space. Since MRI images often exhibit approximate conjugate k-space
+    symmetry, this mask is generally more efficient than a standard equispaced
+    mask.
+
+    Similarly to ``EquispacedMaskFractionFunc``, this method exactly matches
+    the target acceleration by adjusting the offsets.
+    """
+
+    def sample_mask(
+        self,
+        shape: Sequence[int],
+        offset: Optional[int],
+    ) -> Tuple[torch.Tensor, torch.Tensor, int]:
+        """
+        Sample a new k-space mask.
+
+        This function samples and returns two components of a k-space mask: 1)
+        the center mask (e.g., for sensitivity map calculation) and 2) the
+        acceleration mask (for the edge of k-space). Both of these masks, as
+        well as the integer of low frequency samples, are returned.
+
+        Args:
+            shape: Shape of the k-space to subsample.
+            offset: Offset from 0 to begin mask (for equispaced masks).
+
+        Returns:
+            A 3-tuple contaiing 1) the mask for the center of k-space, 2) the
+            mask for the high frequencies of k-space, and 3) the integer count
+            of low frequency samples.
+        """
+        num_cols = shape[-2]
+        fraction_low_freqs, acceleration = self.choose_acceleration()
+        num_cols = shape[-2]
+        num_low_frequencies = round(num_cols * fraction_low_freqs)
+
+        # bound the number of low frequencies between 1 and target columns
+        target_columns_to_sample = round(num_cols / acceleration)
+        num_low_frequencies = max(
+            min(num_low_frequencies, target_columns_to_sample), 1
+        )
+
+        # adjust acceleration rate based on target acceleration.
+        adjusted_target_columns_to_sample = (
+            target_columns_to_sample - num_low_frequencies
+        )
+        adjusted_acceleration = 0
+        if adjusted_target_columns_to_sample > 0:
+            adjusted_acceleration = round(
+                num_cols / adjusted_target_columns_to_sample
+            )
+
+        center_mask = self.reshape_mask(
+            self.calculate_center_mask(shape, num_low_frequencies), shape
+        )
+        accel_mask = self.reshape_mask(
+            self.calculate_acceleration_mask(
+                num_cols, adjusted_acceleration, offset, num_low_frequencies
+            ),
+            shape,
+        )
+
+        return center_mask, accel_mask, num_low_frequencies
+
+
+class Gaussian2DMaskFunc(MaskFunc):
+    """Gaussian 2D Masking
+
+    Args:
+        MaskFunc (_type_): _description_
+    """
+
+    def __init__(
+        self,
+        accelerations: Sequence[int],
+        stds: Sequence[float],
+        seed: Optional[int] = None,
+    ):
+        """initialize Gaussian 2D Mask
+
+        Args:
+            accelerations (Sequence[int]): list of acceleration factors, when
+                generating a mask, an acceleration factor from this list will be chosen
+            stds (Sequence[float]): list of torch.Normal scale (~std) to choose from
+            seed (Optional[int], optional): Seed for selecting mask parameters. Defaults to None.
+        """
+        self.rng = np.random.RandomState(seed)
+        self.accelerations = accelerations
+        self.stds = stds
+
+    def __call__(
+        self,
+        shape: Sequence[int],
+        offset: Optional[int] = None,
+        seed: Optional[Union[int, Tuple[int, ...]]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, int]:
+        if len(shape) < 3:
+            raise ValueError("Shape should have 3 or more dimensions")
+
+        acceleration = self.rng.choice(self.accelerations)
+        std = self.rng.choice(self.stds)
+
+        x, y = shape[-3], shape[-2]
+        mean_x = x // 2
+        mean_y = y // 2
+        num_samples_collected = 0
+
+        dist = D.Normal(
+            loc=torch.tensor([mean_x, mean_y], dtype=torch.float32),
+            scale=std,
+        )
+
+        N = (
+            int(1 / acceleration * x * y) + 10000
+        )  # add constant or won't reach desired subsampling rate
+        sample_x, sample_y = (
+            torch.zeros(N, dtype=torch.int),
+            torch.zeros(N, dtype=torch.int),
+        )
+
+        while num_samples_collected < N:
+            samples = dist.sample((N,))  # type: ignore
+            valid_samples = (
+                (samples[:, 0] >= 0)
+                & (samples[:, 0] < x)
+                & (samples[:, 1] >= 0)
+                & (samples[:, 1] < y)
+            )
+
+            valid_x = samples[valid_samples, 0].int()
+            valid_y = samples[valid_samples, 1].int()
+
+            num_to_take = min(N - num_samples_collected, valid_x.size(0))
+            sample_x[
+                num_samples_collected : num_samples_collected + num_to_take
+            ] = valid_x[:num_to_take]
+            sample_y[
+                num_samples_collected : num_samples_collected + num_to_take
+            ] = valid_y[:num_to_take]
+            num_samples_collected += num_to_take
+
+        mask = torch.zeros((x, y))
+        mask[sample_x, sample_y] = 1.0
+
+        # broadcasting mask (x, y) --> (N, x, y, C) C=2, N=batch_size
+        mask = mask.unsqueeze(-1)  # (x, y, 1)
+        mask = mask.unsqueeze(0)  # (1, x, y, 1)
+        mask = mask.expand((1, mask.shape[1], mask.shape[2], 2)).clone()
+
+        # num_low_freqs doesn't make sense so just return std (a number)
+        # returning None doesn't work since we can't stack for multiple batches
+        return mask, std
+
+
+class Poisson2DMaskFunc(MaskFunc):
+    """
+    Variable Density Poisson Disk Sampling
+    https://sigpy.readthedocs.io/en/latest/generated/sigpy.mri.poisson.html#sigpy.mri.poisson
+    """
+
+    def __init__(
+        self,
+        accelerations: Sequence[int],
+        stds: None,
+        seed: Optional[int] = None,
+        use_cache: bool = True,
+    ):
+        """initialize VDPD (Poisson) mask
+
+        Args:
+            accelerations (Sequence[int]): list of acceleration factors to
+                choose from
+            stds: Dummy param. Do not pass value. Defaults to None.
+            seed (Optional[int], optional): Seed for selecting mask params.
+                Defaults to None.
+        """
+        self.rng = np.random.RandomState(seed)
+        self.accelerations = accelerations
+        self.use_cache = use_cache
+        if use_cache:
+            self.cache: Dict[int, np.ndarray] = dict()
+            for acc in accelerations:
+                # assert os.path.exists(
+                #     f"fastmri/poisson_cache/poisson_{acc}x.npy"
+                # )
+                # self.cache[acc] = np.load(
+                #     f"fastmri/poisson_cache/poisson_{acc}x.npy"
+                # )
+                self.cache[acc] = np.load(
+                    f"/global/homes/p/peterwg/more/medical-imaging/fastmri/poisson_cache/poisson_{acc}x.npy"
+                )
+
+    def __call__(
+        self,
+        shape: Sequence[int],
+        offset: Optional[int] = None,
+        seed: Optional[Union[int, Tuple[int, ...]]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, int]:
+        if self.use_cache:
+            acceleration = self.rng.choice(self.accelerations)
+            return torch.from_numpy(self.cache[acceleration]), 1.0  # type: ignore
+        if len(shape) < 3:
+            raise ValueError("Shape should have 3 or more dimensions")
+
+        acceleration = self.rng.choice(self.accelerations)
+        x, y = shape[-3], shape[-2]
+
+        mask = poisson(img_shape=(x, y), accel=acceleration, dtype=np.float32)
+        mask = torch.from_numpy(mask)
+
+        # broadcasting mask (x, y) --> (N, x, y, C) C=2, N=batch_size
+        mask = mask.unsqueeze(-1)  # (x, y, 1e
+        mask = mask.unsqueeze(0)  # (1, x, y, 1)
+        mask = mask.expand((1, mask.shape[1], mask.shape[2], 2)).clone()
+
+        # num low freqs doesn't make sense here, so we return arbitrary value 1.0
+        return mask, 100.0
+
+
+class Radial2DMaskFunc(MaskFunc):
+    """
+    Radial trajectory MRI masking method.
+    https://sigpy.readthedocs.io/en/latest/generated/sigpy.mri.radial.html#sigpy.mri.radial
+    """
+
+    def __init__(
+        self,
+        accelerations: Sequence[int],
+        arms: Optional[Sequence[int]],
+        seed: Optional[int] = None,
+    ):
+        """
+        initialize Radial mask
+
+        Args:
+            accelerations (Sequence[int]): list of acceleration factors to
+                choose from
+            arms: Number of radial arms.
+            seed (Optional[int], optional): Seed for selecting mask params.
+                Defaults to None.
+        """
+        self.rng = np.random.RandomState(seed)
+        self.accelerations = accelerations
+        self.arms = arms
+
+    def __call__(
+        self,
+        shape: Sequence[int],
+        offset: Optional[int] = None,
+        seed: Optional[Union[int, Tuple[int, ...]]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, int]:
+        if len(shape) < 3:
+            raise ValueError("Shape should have 3 or more dimensions")
+
+        acceleration = self.rng.choice(self.accelerations)
+        x, y = shape[-3], shape[-2]
+        npoints = int(x * y * (1 / acceleration))
+        if self.arms:
+            arms = self.rng.choice(self.arms)
+        else:
+            points_per_arm = x // 3
+            arms = npoints // points_per_arm
+
+        # calculate radial parameters to satisfy acceleration factor
+        ntr = arms  # num radial lines
+        nro = npoints // arms  # num points on each radial line
+        ndim = 2  # 2D
+
+        # gen trajectory w/ shape (ntr, nro, ndim)
+        traj = radial(
+            coord_shape=[ntr, nro, ndim],
+            img_shape=(x, y),
+            golden=True,
+            dtype=int,
+        )
+
+        mask = torch.zeros(x, y, dtype=torch.float32)
+        x_coords = traj[..., 0].flatten() + (x // 2)
+        y_coords = traj[..., 1].flatten() + (y // 2)
+        mask[x_coords, y_coords] = 1.0
+
+        # broadcasting mask (x, y) --> (N, x, y, C) C=2, N=batch_size
+        mask = mask.unsqueeze(-1)  # (x, y, 1)
+        mask = mask.unsqueeze(0)  # (1, x, y, 1)
+        mask = mask.expand((1, mask.shape[1], mask.shape[2], 2)).clone()
+
+        # num low freqs doesn't make sense here, so we return arbitrary value 1.0
+        return mask, 100.0
+
+
+class Spiral2DMaskFunc(MaskFunc):
+    """
+    Spiral trajectory MRI masking method.
+    https://sigpy.readthedocs.io/en/latest/generated/sigpy.mri.spiral.html#sigpy.mri.spiral
+    """
+
+    def __init__(
+        self,
+        accelerations: Sequence[int],
+        arms: Sequence[int],
+        seed: Optional[int] = None,
+    ):
+        """
+        initialize Radial mask
+
+        Args:
+            accelerations (Sequence[int]): list of acceleration factors to
+                choose from
+            arms: Number of radial arms.
+            seed (Optional[int], optional): Seed for selecting mask params.
+                Defaults to None.
+        """
+        self.rng = np.random.RandomState(seed)
+        self.accelerations = accelerations
+        self.arms = arms
+
+    def __call__(
+        self,
+        shape: Sequence[int],
+        offset: Optional[int] = None,
+        seed: Optional[Union[int, Tuple[int, ...]]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, int]:
+        # TODO: implement
+        raise (NotImplementedError("Spiral2D not implemented"))
+        if len(shape) < 3:
+            raise ValueError("Shape should have 3 or more dimensions")
+        acceleration = self.rng.choice(self.accelerations)
+        arms = self.rng.choice(self.arms)
+        x, y = shape[-3], shape[-2]
+
+        # calculate radial parameters to satisfy acceleration factor
+        npoints = int(x * y * (1 / acceleration))
+
+        # gen trajectory w/ shape (ntr, nro, ndim)
+        traj = spiral(
+            N=npoints,
+            img_shape=(x, y),
+            golden=True,
+            dtype=int,
+        )
+
+        mask = torch.zeros(x, y, dtype=float)
+        x_coords = traj[..., 0].flatten() + (x // 2)
+        y_coords = traj[..., 1].flatten() + (y // 2)
+        mask[x_coords, y_coords] = 1.0
+
+        # broadcasting mask (x, y) --> (N, x, y, C) C=2, N=batch_size
+        mask = mask.unsqueeze(-1)  # (x, y, 1)
+        mask = mask.unsqueeze(0)  # (1, x, y, 1)
+        mask = mask.expand((1, mask.shape[1], mask.shape[2], 2)).clone()
+
+        # num low freqs doesn't make sense here, so we return arbitrary value 1.0
+        return mask, 100.0
+
+
+def create_mask_for_mask_type(
+    mask_type_str: str,
+    center_fractions: Optional[Sequence],
+    accelerations: Sequence[int],
+) -> MaskFunc:
+    """
+    Creates a mask of the specified type.
+
+    Args:
+        center_fractions: What fraction of the center of k-space to include.
+        accelerations: What accelerations to apply.
+
+    Returns:
+        A mask func for the target mask type.
+    """
+    if mask_type_str == "random":
+        return RandomMaskFunc(center_fractions, accelerations)
+    elif mask_type_str == "equispaced":
+        return EquiSpacedMaskFunc(center_fractions, accelerations)
+    elif mask_type_str == "equispaced_fraction":
+        return EquispacedMaskFractionFunc(center_fractions, accelerations)
+    elif mask_type_str == "magic":
+        return MagicMaskFunc(center_fractions, accelerations)
+    elif mask_type_str == "magic_fraction":
+        return MagicMaskFractionFunc(center_fractions, accelerations)
+    elif mask_type_str == "gaussian_2d":
+        return Gaussian2DMaskFunc(
+            stds=center_fractions,
+            accelerations=accelerations,
+        )
+    elif mask_type_str == "poisson_2d":
+        return Poisson2DMaskFunc(
+            accelerations=accelerations,
+            stds=None,
+        )
+    elif mask_type_str == "radial_2d":
+        return Radial2DMaskFunc(
+            accelerations=accelerations,
+            arms=(
+                [int(arm) for arm in center_fractions]
+                if center_fractions
+                else None
+            ),
+        )
+    elif mask_type_str == "spiral_2d":
+        raise NotImplementedError("spiral_2d not implemented")
+    else:
+        raise ValueError(f"{mask_type_str} not supported")

fastmri/transforms.py

@@ -0,0 +1,974 @@
+"""
+Copyright (c) Facebook, Inc. and its affiliates.
+
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+import random
+from typing import Dict, List, NamedTuple, Optional, Sequence, Tuple, Union
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+import fastmri
+
+from .subsample import MaskFunc
+
+
+def to_tensor(data: np.ndarray) -> torch.Tensor:
+    """
+    Convert numpy array to PyTorch tensor.
+
+    For complex arrays, the real and imaginary parts are stacked along the last
+    dimension.
+
+    Args:
+        data: Input numpy array.
+
+    Returns:
+        PyTorch version of data.
+    """
+    if np.iscomplexobj(data):
+        data = np.stack((data.real, data.imag), axis=-1)
+
+    return torch.from_numpy(data)
+
+
+def tensor_to_complex_np(data: torch.Tensor) -> np.ndarray:
+    """
+    Converts a complex torch tensor to numpy array.
+
+    Args:
+        data: Input data to be converted to numpy.
+
+    Returns:
+        Complex numpy version of data.
+    """
+    return torch.view_as_complex(data).numpy()
+
+
+def apply_mask(
+    data: torch.Tensor,
+    mask_func: MaskFunc,
+    offset: Optional[int] = None,
+    seed: Optional[Union[int, Tuple[int, ...]]] = None,
+    padding: Optional[Sequence[int]] = None,
+) -> Tuple[torch.Tensor, torch.Tensor, int]:
+    """
+    Subsample given k-space by multiplying with a mask.
+
+    Args:
+        data: The input k-space data. This should have at least 3 dimensions,
+            where dimensions -3 and -2 are the spatial dimensions, and the
+            final dimension has size 2 (for complex values).
+        mask_func: A function that takes a shape (tuple of ints) and a random
+            number seed and returns a mask.
+        seed: Seed for the random number generator.
+        padding: Padding value to apply for mask.
+
+    Returns:
+        tuple containing:
+            masked data: Subsampled k-space data.
+            mask: The generated mask.
+            num_low_frequencies: The number of low-resolution frequency samples
+                in the mask.
+    """
+    shape = (1,) * len(data.shape[:-3]) + tuple(data.shape[-3:])
+    mask, num_low_frequencies = mask_func(shape, offset, seed)
+    if padding is not None:
+        mask[..., : padding[0], :] = 0
+        mask[..., padding[1] :, :] = (
+            0  # padding value inclusive on right of zeros
+        )
+
+    masked_data = data * mask + 0.0  # the + 0.0 removes the sign of the zeros
+
+    return masked_data, mask, num_low_frequencies
+
+
+def mask_center(x: torch.Tensor, mask_from: int, mask_to: int) -> torch.Tensor:
+    """
+    Initializes a mask with the center filled in.
+
+    Args:
+        mask_from: Part of center to start filling.
+        mask_to: Part of center to end filling.
+
+    Returns:
+        A mask with the center filled.
+    """
+    mask = torch.zeros_like(x)
+    mask[:, :, :, mask_from:mask_to] = x[:, :, :, mask_from:mask_to]
+
+    return mask
+
+
+def batched_mask_center(
+    x: torch.Tensor, mask_from: torch.Tensor, mask_to: torch.Tensor
+) -> torch.Tensor:
+    """
+    Initializes a mask with the center filled in.
+
+    Can operate with different masks for each batch element.
+
+    Args:
+        mask_from: Part of center to start filling.
+        mask_to: Part of center to end filling.
+
+    Returns:
+        A mask with the center filled.
+    """
+    if not mask_from.shape == mask_to.shape:
+        raise ValueError("mask_from and mask_to must match shapes.")
+    if not mask_from.ndim == 1:
+        raise ValueError("mask_from and mask_to must have 1 dimension.")
+    if not mask_from.shape[0] == 1:
+        if (not x.shape[0] == mask_from.shape[0]) or (
+            not x.shape[0] == mask_to.shape[0]
+        ):
+            raise ValueError(
+                "mask_from and mask_to must have batch_size length."
+            )
+
+    if mask_from.shape[0] == 1:
+        mask = mask_center(x, int(mask_from), int(mask_to))
+    else:
+        mask = torch.zeros_like(x)
+        for i, (start, end) in enumerate(zip(mask_from, mask_to)):
+            mask[i, :, :, start:end] = x[i, :, :, start:end]
+
+    return mask
+
+
+def center_crop(data: torch.Tensor, shape: Tuple[int, int]) -> torch.Tensor:
+    """
+    Apply a center crop to the input real image or batch of real images.
+
+    Args:
+        data: The input tensor to be center cropped. It should
+            have at least 2 dimensions and the cropping is applied along the
+            last two dimensions.
+        shape: The output shape. The shape should be smaller
+            than the corresponding dimensions of data.
+
+    Returns:
+        The center cropped image.
+    """
+    if not (0 < shape[0] <= data.shape[-2] and 0 < shape[1] <= data.shape[-1]):
+        raise ValueError("Invalid shapes.")
+
+    w_from = (data.shape[-2] - shape[0]) // 2
+    h_from = (data.shape[-1] - shape[1]) // 2
+    w_to = w_from + shape[0]
+    h_to = h_from + shape[1]
+
+    return data[..., w_from:w_to, h_from:h_to]
+
+
+def complex_center_crop(
+    data: torch.Tensor, shape: Tuple[int, int]
+) -> torch.Tensor:
+    """
+    Apply a center crop to the input image or batch of complex images.
+
+    Args:
+        data: The complex input tensor to be center cropped. It should have at
+            least 3 dimensions and the cropping is applied along dimensions -3
+            and -2 and the last dimensions should have a size of 2.
+        shape: The output shape. The shape should be smaller than the
+            corresponding dimensions of data.
+    Returns:
+        The center cropped image
+    """
+    if not (0 < shape[0] <= data.shape[-3] and 0 < shape[1] <= data.shape[-2]):
+        raise ValueError("Invalid shapes.")
+
+    w_from = (data.shape[-3] - shape[0]) // 2
+    h_from = (data.shape[-2] - shape[1]) // 2
+    w_to = w_from + shape[0]
+    h_to = h_from + shape[1]
+
+    return data[..., w_from:w_to, h_from:h_to, :]
+
+
+def center_crop_to_smallest(
+    x: torch.Tensor, y: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Apply a center crop on the larger image to the size of the smaller.
+
+    The minimum is taken over dim=-1 and dim=-2. If x is smaller than y at
+    dim=-1 and y is smaller than x at dim=-2, then the returned dimension will
+    be a mixture of the two.
+
+    Args:
+        x: The first image.
+        y: The second image.
+
+    Returns:
+        tuple of tensors x and y, each cropped to the minimim size.
+    """
+    smallest_width = min(x.shape[-1], y.shape[-1])
+    smallest_height = min(x.shape[-2], y.shape[-2])
+    x = center_crop(x, (smallest_height, smallest_width))
+    y = center_crop(y, (smallest_height, smallest_width))
+
+    return x, y
+
+
+def normalize(
+    data: torch.Tensor,
+    mean: Union[float, torch.Tensor],
+    stddev: Union[float, torch.Tensor],
+    eps: Union[float, torch.Tensor] = 0.0,
+) -> torch.Tensor:
+    """
+    Normalize the given tensor.
+
+    Applies the formula (data - mean) / (stddev + eps).
+
+    Args:
+        data: Input data to be normalized.
+        mean: Mean value.
+        stddev: Standard deviation.
+        eps: Added to stddev to prevent dividing by zero.
+
+    Returns:
+        Normalized tensor.
+    """
+    return (data - mean) / (stddev + eps)
+
+
+def normalize_instance(
+    data: torch.Tensor, eps: Union[float, torch.Tensor] = 0.0
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Normalize the given tensor  with instance norm/
+
+    Applies the formula (data - mean) / (stddev + eps), where mean and stddev
+    are computed from the data itself.
+
+    Args:
+        data: Input data to be normalized
+        eps: Added to stddev to prevent dividing by zero.
+
+    Returns:
+        torch.Tensor: Normalized tensor
+    """
+    mean = data.mean()
+    std = data.std()
+
+    return normalize(data, mean, std, eps), mean, std
+
+
+class UnetSample(NamedTuple):
+    """
+    A subsampled image for U-Net reconstruction.
+
+    Args:
+        image: Subsampled image after inverse FFT.
+        target: The target image (if applicable).
+        mean: Per-channel mean values used for normalization.
+        std: Per-channel standard deviations used for normalization.
+        fname: File name.
+        slice_num: The slice index.
+        max_value: Maximum image value.
+    """
+
+    image: torch.Tensor
+    target: torch.Tensor
+    mean: torch.Tensor
+    std: torch.Tensor
+    fname: str
+    slice_num: int
+    max_value: float
+
+
+class UnetDataTransform:
+    """
+    Data Transformer for training U-Net models.
+    """
+
+    def __init__(
+        self,
+        which_challenge: str,
+        mask_func: Optional[MaskFunc] = None,
+        use_seed: bool = True,
+    ):
+        """
+        Args:
+            which_challenge: Challenge from ("singlecoil", "multicoil").
+            mask_func: Optional; A function that can create a mask of
+                appropriate shape.
+            use_seed: If true, this class computes a pseudo random number
+                generator seed from the filename. This ensures that the same
+                mask is used for all the slices of a given volume every time.
+        """
+        if which_challenge not in ("singlecoil", "multicoil"):
+            raise ValueError(
+                "Challenge should either be 'singlecoil' or 'multicoil'"
+            )
+
+        self.mask_func = mask_func
+        self.which_challenge = which_challenge
+        self.use_seed = use_seed
+
+    def __call__(
+        self,
+        kspace: np.ndarray,
+        mask: np.ndarray,
+        target: np.ndarray,
+        attrs: Dict,
+        fname: str,
+        slice_num: int,
+    ) -> Tuple[
+        torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, str, int, float
+    ]:
+        """
+        Args:
+            kspace: Input k-space of shape (num_coils, rows, cols) for
+                multi-coil data or (rows, cols) for single coil data.
+            mask: Mask from the test dataset.
+            target: Target image.
+            attrs: Acquisition related information stored in the HDF5 object.
+            fname: File name.
+            slice_num: Serial number of the slice.
+
+        Returns:
+            A tuple containing, zero-filled input image, the reconstruction
+            target, the mean used for normalization, the standard deviations
+            used for normalization, the filename, and the slice number.
+        """
+        kspace_torch = to_tensor(kspace)
+
+        # check for max value
+        max_value = attrs["max"] if "max" in attrs.keys() else 0.0
+
+        # apply mask
+        if self.mask_func:
+            seed = None if not self.use_seed else tuple(map(ord, fname))
+            # we only need first element, which is k-space after masking
+            masked_kspace = apply_mask(kspace_torch, self.mask_func, seed=seed)[
+                0
+            ]
+        else:
+            masked_kspace = kspace_torch
+
+        # inverse Fourier transform to get zero filled solution
+        image = fastmri.ifft2c(masked_kspace)
+
+        # crop input to correct size
+        if target is not None:
+            crop_size = (target.shape[-2], target.shape[-1])
+        else:
+            crop_size = (attrs["recon_size"][0], attrs["recon_size"][1])
+
+        # check for FLAIR 203
+        if image.shape[-2] < crop_size[1]:
+            crop_size = (image.shape[-2], image.shape[-2])
+
+        image = complex_center_crop(image, crop_size)
+
+        # absolute value
+        image = fastmri.complex_abs(image)
+
+        # apply Root-Sum-of-Squares if multicoil data
+        if self.which_challenge == "multicoil":
+            image = fastmri.rss(image)
+
+        # normalize input
+        image, mean, std = normalize_instance(image, eps=1e-11)
+        image = image.clamp(-6, 6)
+
+        # normalize target
+        if target is not None:
+            target_torch = to_tensor(target)
+            target_torch = center_crop(target_torch, crop_size)
+            target_torch = normalize(target_torch, mean, std, eps=1e-11)
+            target_torch = target_torch.clamp(-6, 6)
+        else:
+            target_torch = torch.Tensor([0])
+
+        return UnetSample(
+            image=image,
+            target=target_torch,
+            mean=mean,
+            std=std,
+            fname=fname,
+            slice_num=slice_num,
+            max_value=max_value,
+        )
+
+
+class VarNetSample(NamedTuple):
+    """
+    A sample of masked k-space for variational network reconstruction.
+
+    Args:
+        masked_kspace: k-space after applying sampling mask.
+        mask: The applied sampling mask.
+        num_low_frequencies: The number of samples for the densely-sampled
+            center.
+        target: The target image (if applicable).
+        fname: File name.
+        slice_num: The slice index.
+        max_value: Maximum image value.
+        crop_size: The size to crop the final image.
+    """
+
+    masked_kspace: torch.Tensor
+    mask: torch.Tensor
+    num_low_frequencies: Optional[int]
+    target: torch.Tensor
+    fname: str
+    slice_num: int
+    max_value: float
+    crop_size: Tuple[int, int]
+
+
+class VarNetDataTransform:
+    """
+    Data Transformer for training VarNet models.
+    """
+
+    def __init__(
+        self, mask_func: Optional[MaskFunc] = None, use_seed: bool = True
+    ):
+        """
+        Args:
+            mask_func: Optional; A function that can create a mask of
+                appropriate shape. Defaults to None.
+            use_seed: If True, this class computes a pseudo random number
+                generator seed from the filename. This ensures that the same
+                mask is used for all the slices of a given volume every time.
+        """
+        self.mask_func = mask_func
+        self.use_seed = use_seed
+
+    def __call__(
+        self,
+        kspace: np.ndarray,
+        mask: np.ndarray,
+        target: Optional[np.ndarray],
+        attrs: Dict,
+        fname: str,
+        slice_num: int,
+    ) -> VarNetSample:
+        """
+        Args:
+            kspace: Input k-space of shape (num_coils, rows, cols) for
+                multi-coil data.
+            mask: Mask from the test dataset.
+            target: Target image.
+            attrs: Acquisition related information stored in the HDF5 object.
+            fname: File name.
+            slice_num: Serial number of the slice.
+
+        Returns:
+            A VarNetSample with the masked k-space, sampling mask, target
+            image, the filename, the slice number, the maximum image value
+            (from target), the target crop size, and the number of low
+            frequency lines sampled.
+        """
+        if target is not None:
+            target_torch = to_tensor(target)
+            max_value = attrs["max"]
+        else:
+            target_torch = torch.tensor(0)
+            max_value = 0.0
+
+        kspace_torch = to_tensor(kspace)
+        seed = None if not self.use_seed else tuple(map(ord, fname))
+        acq_start = attrs["padding_left"]
+        acq_end = attrs["padding_right"]
+
+        crop_size = (attrs["recon_size"][0], attrs["recon_size"][1])
+
+        if self.mask_func is not None:
+            masked_kspace, mask_torch, num_low_frequencies = apply_mask(
+                kspace_torch,
+                self.mask_func,
+                seed=seed,
+                padding=(acq_start, acq_end),
+            )
+
+            sample = VarNetSample(
+                masked_kspace=masked_kspace,
+                mask=mask_torch.to(torch.bool),
+                num_low_frequencies=num_low_frequencies,
+                target=target_torch,
+                fname=fname,
+                slice_num=slice_num,
+                max_value=max_value,
+                crop_size=crop_size,
+            )
+        else:
+            masked_kspace = kspace_torch
+            shape = np.array(kspace_torch.shape)
+            num_cols = shape[-2]
+            shape[:-3] = 1
+            mask_shape = [1] * len(shape)
+            mask_shape[-2] = num_cols
+            mask_torch = torch.from_numpy(
+                mask.reshape(*mask_shape).astype(np.float32)
+            )
+            mask_torch = mask_torch.reshape(*mask_shape)
+            mask_torch[:, :, :acq_start] = 0
+            mask_torch[:, :, acq_end:] = 0
+
+            sample = VarNetSample(
+                masked_kspace=masked_kspace,
+                mask=mask_torch.to(torch.bool),
+                num_low_frequencies=0,
+                target=target_torch,
+                fname=fname,
+                slice_num=slice_num,
+                max_value=max_value,
+                crop_size=crop_size,
+            )
+
+        # whether to crop samples for batch processing
+        batch_crop = False
+
+        def save_img(x, fname):
+            slice_kspace2 = x
+            slice_image = fastmri.ifft2c(
+                slice_kspace2
+            )  # Apply Inverse Fourier Transform to get the complex image
+            slice_image_abs = fastmri.complex_abs(
+                slice_image
+            )  # Compute absolute value to get a real image
+            slice_image_rss = fastmri.rss(slice_image_abs, dim=0)
+            plt.imsave(f"{fname}.png", torch.abs(slice_image_rss), cmap="gray")
+
+        def save_raw_img(x, fname):
+            # slice_kspace2 = x
+            # slice_image = fastmri.ifft2c(
+            #     slice_kspace2
+            # )  # Apply Inverse Fourier Transform to get the complex image
+            # slice_image_abs = fastmri.complex_abs(
+            #     slice_image
+            # )  # Compute absolute value to get a real image
+            x = fastmri.rss(x, dim=0)[:, :, 0]
+
+            plt.imsave(f"{fname}.png", torch.abs(x))
+
+        if batch_crop:
+            # crop kspace data to minx, miny size (640, 320 cols)
+            square_crop = (attrs["recon_size"][0], attrs["recon_size"][1])
+            # print(square_crop)
+            cropped_kspace = fastmri.fft2c(
+                complex_center_crop(
+                    fastmri.ifft2c(sample.masked_kspace), square_crop
+                )
+            )
+            cropped_kspace = complex_center_crop(cropped_kspace, (320, 320))
+            # print(cropped_kspace.shape)
+            # exit(0)
+
+            # CHECK: debugging purposes
+            # save_img(sample.masked_kspace, "og")
+            # save_img(cropped_kspace, "cropped")
+            # save_raw_img(sample.masked_kspace, "og_kspace")
+            # save_raw_img(cropped_kspace, "cropped_kspace")
+
+            # exit(0)
+
+            # crop mask shape
+            h_from = (mask_torch.shape[-2] - 320) // 2
+            h_to = h_from + 320
+            cropped_mask = mask_torch[..., :, h_from:h_to, :]
+
+            sample = VarNetSample(
+                masked_kspace=cropped_kspace,
+                mask=cropped_mask.to(torch.bool),
+                num_low_frequencies=0,
+                target=target_torch,
+                fname=fname,
+                slice_num=slice_num,
+                max_value=max_value,
+                crop_size=crop_size,
+            )
+        return sample
+
+
+class EnhancedVarNetDataTransform(VarNetDataTransform):
+    """
+    Enhanced Data Transformer for training VarNet models with additional functionality.
+        - allows for training on multiple patterns
+    """
+
+    def __init__(
+        self, mask_funcs: List[MaskFunc] = None, use_seed: bool = True
+    ):
+        self.mask_funcs = mask_funcs
+        self.use_seed = use_seed
+
+    def __call__(
+        self,
+        kspace: np.ndarray,
+        mask: np.ndarray,
+        target: Optional[np.ndarray],
+        attrs: Dict,
+        fname: str,
+        slice_num: int,
+    ) -> VarNetSample:
+        """
+        Args:
+            kspace: Input k-space of shape (num_coils, rows, cols) for
+                multi-coil data.
+            mask: Mask from the test dataset.
+                use mask for test data see og VarNetDataTransform __call__
+            target: Target image.
+            attrs: Acquisition related information stored in the HDF5 object.
+            fname: File name.
+            slice_num: Serial number of the slice.
+
+        Returns:
+            A VarNetSample with the masked k-space, sampling mask, target
+            image, the filename, the slice number, the maximum image value
+            (from target), the target crop size, and the number of low
+            frequency lines sampled.
+        """
+        if target is not None:
+            target_torch = to_tensor(target)
+            max_value = attrs["max"]
+        else:
+            target_torch = torch.tensor(0)
+            max_value = 0.0
+
+        kspace_torch = to_tensor(kspace)
+        seed = None if not self.use_seed else tuple(map(ord, fname))
+        acq_start = attrs["padding_left"]
+        acq_end = attrs["padding_right"]
+
+        crop_size = (attrs["recon_size"][0], attrs["recon_size"][1])
+
+        # choose one of the masking functions provided randomly
+        mask_func = random.choice(self.mask_funcs)
+
+        masked_kspace, mask_torch, num_low_frequencies = apply_mask(
+            kspace_torch,
+            mask_func,
+            seed=seed,
+            padding=(acq_start, acq_end),
+        )
+
+        # print(masked_kspace.shape)
+        # print(mask_torch.shape)
+
+        # torch.save(masked_kspace, f"masked_kspace_{slice_num}.pkl")
+        # torch.save(mask_torch, f"mask_torch_{slice_num}.pkl")
+
+        sample = VarNetSample(
+            masked_kspace=masked_kspace,
+            mask=mask_torch.to(torch.bool),
+            num_low_frequencies=num_low_frequencies,
+            target=target_torch,
+            fname=fname,
+            slice_num=slice_num,
+            max_value=max_value,
+            crop_size=crop_size,
+        )
+
+        # whether to crop samples for batch processing
+        batch_crop = False
+
+        if batch_crop:
+            # crop kspace data to minx, miny size (640, 320 cols)
+            square_crop = (attrs["recon_size"][0], attrs["recon_size"][1])
+            # print(square_crop)
+            cropped_kspace = fastmri.fft2c(
+                complex_center_crop(
+                    fastmri.ifft2c(sample.masked_kspace), square_crop
+                )
+            )
+            # cropped_kspace = complex_center_crop(cropped_kspace, (640, 320))
+
+            # exit(0)
+
+            # crop mask shape
+            h_from = (mask_torch.shape[-2] - 320) // 2
+            h_to = h_from + 320
+            cropped_mask = mask_torch[..., :, h_from:h_to, :]
+
+            sample = VarNetSample(
+                masked_kspace=cropped_kspace,
+                mask=cropped_mask.to(torch.bool),
+                num_low_frequencies=0,
+                target=target_torch,
+                fname=fname,
+                slice_num=slice_num,
+                max_value=max_value,
+                crop_size=crop_size,
+            )
+
+        return sample
+
+
+class MiniCoilSample(NamedTuple):
+    """
+    A sample of masked coil-compressed k-space for reconstruction.
+
+    Args:
+        kspace: the original k-space before masking.
+        masked_kspace: k-space after applying sampling mask.
+        mask: The applied sampling mask.
+        num_low_frequencies: The number of samples for the densely-sampled
+            center.
+        target: The target image (if applicable).
+        fname: File name.
+        slice_num: The slice index.
+        max_value: Maximum image value.
+        crop_size: The size to crop the final image.
+    """
+
+    kspace: torch.Tensor
+    masked_kspace: torch.Tensor
+    mask: torch.Tensor
+    target: torch.Tensor
+    fname: str
+    slice_num: int
+    max_value: float
+    crop_size: Tuple[int, int]
+
+
+class MiniCoilTransform:
+    """
+    Multi-coil compressed transform, for faster prototyping.
+    """
+
+    def __init__(
+        self,
+        mask_func: Optional[MaskFunc] = None,
+        use_seed: Optional[bool] = True,
+        crop_size: Optional[tuple] = None,
+        num_compressed_coils: Optional[int] = None,
+    ):
+        """
+        Args:
+            mask_func: Optional; A function that can create a mask of
+                appropriate shape. Defaults to None.
+            use_seed: If True, this class computes a pseudo random number
+                generator seed from the filename. This ensures that the same
+                mask is used for all the slices of a given volume every time.
+            crop_size: Image dimensions for mini MR images.
+            num_compressed_coils: Number of coils to output from coil
+                compression.
+        """
+        self.mask_func = mask_func
+        self.use_seed = use_seed
+        self.crop_size = crop_size
+        self.num_compressed_coils = num_compressed_coils
+
+    def __call__(self, kspace, mask, target, attrs, fname, slice_num):
+        """
+        Args:
+            kspace: Input k-space of shape (num_coils, rows, cols) for
+                multi-coil data.
+            mask: Mask from the test dataset. Not used if mask_func is defined.
+            target: Target image.
+            attrs: Acquisition related information stored in the HDF5 object.
+            fname: File name.
+            slice_num: Serial number of the slice.
+
+        Returns:
+            tuple containing:
+                kspace: original kspace (used for active acquisition only).
+                masked_kspace: k-space after applying sampling mask. If there
+                    is no mask or mask_func, returns same as kspace.
+                mask: The applied sampling mask
+                target: The target image (if applicable). The target is built
+                    from the RSS opp of all coils pre-compression.
+                fname: File name.
+                slice_num: The slice index.
+                max_value: Maximum image value.
+                crop_size: The size to crop the final image.
+        """
+        if target is not None:
+            target = to_tensor(target)
+            max_value = attrs["max"]
+        else:
+            target = torch.tensor(0)
+            max_value = 0.0
+
+        if self.crop_size is None:
+            crop_size = torch.tensor(
+                [attrs["recon_size"][0], attrs["recon_size"][1]]
+            )
+        else:
+            if isinstance(self.crop_size, tuple) or isinstance(
+                self.crop_size, list
+            ):
+                assert len(self.crop_size) == 2
+                if self.crop_size[0] is None or self.crop_size[1] is None:
+                    crop_size = torch.tensor(
+                        [attrs["recon_size"][0], attrs["recon_size"][1]]
+                    )
+                else:
+                    crop_size = torch.tensor(self.crop_size)
+            elif isinstance(self.crop_size, int):
+                crop_size = torch.tensor((self.crop_size, self.crop_size))
+            else:
+                raise ValueError(
+                    "`crop_size` should be None, tuple, list, or int, not:"
+                    f" {type(self.crop_size)}"
+                )
+
+        if self.num_compressed_coils is None:
+            num_compressed_coils = kspace.shape[0]
+        else:
+            num_compressed_coils = self.num_compressed_coils
+
+        seed = None if not self.use_seed else tuple(map(ord, fname))
+        acq_start = 0
+        acq_end = crop_size[1]
+
+        # new cropping section
+        square_crop = (attrs["recon_size"][0], attrs["recon_size"][1])
+        kspace = fastmri.fft2c(
+            complex_center_crop(fastmri.ifft2c(to_tensor(kspace)), square_crop)
+        ).numpy()
+        kspace = complex_center_crop(kspace, crop_size)
+
+        # we calculate the target before coil compression. This causes the mini
+        # simulation to be one where we have a 15-coil, low-resolution image
+        # and our reconstructor has an SVD coil approximation. This is a little
+        # bit more realistic than doing the target after SVD compression
+        target = fastmri.rss_complex(fastmri.ifft2c(to_tensor(kspace)))
+        max_value = target.max()
+
+        # apply coil compression
+        new_shape = (num_compressed_coils,) + kspace.shape[1:]
+        kspace = np.reshape(kspace, (kspace.shape[0], -1))
+        left_vec, _, _ = np.linalg.svd(
+            kspace, compute_uv=True, full_matrices=False
+        )
+        kspace = np.reshape(
+            np.array(np.matrix(left_vec[:, :num_compressed_coils]).H @ kspace),
+            new_shape,
+        )
+        kspace = to_tensor(kspace)
+
+        # Mask kspace
+        if self.mask_func:
+            masked_kspace, mask, _ = apply_mask(
+                kspace, self.mask_func, seed, (acq_start, acq_end)
+            )
+            mask = mask.byte()
+        elif mask is not None:
+            masked_kspace = kspace
+            shape = np.array(kspace.shape)
+            num_cols = shape[-2]
+            shape[:-3] = 1
+            mask_shape = [1] * len(shape)
+            mask_shape[-2] = num_cols
+            mask = torch.from_numpy(
+                mask.reshape(*mask_shape).astype(np.float32)
+            )
+            mask = mask.reshape(*mask_shape)
+            mask = mask.byte()
+        else:
+            masked_kspace = kspace
+            shape = np.array(kspace.shape)
+            num_cols = shape[-2]
+
+        return MiniCoilSample(
+            kspace,
+            masked_kspace,
+            mask,
+            target,
+            fname,
+            slice_num,
+            max_value,
+            crop_size,
+        )
+
+
+"""
+sens maps & feature transformations
+- expand
+- reduce
+- batch -> chan
+- chan -> batch
+"""
+
+
+def sens_expand(x: torch.Tensor, sens_maps: torch.Tensor) -> torch.Tensor:
+    """
+    Calculates F (x sens_maps)
+
+    Parameters
+    ----------
+    x : ndarray
+        Single-channel image of shape (..., H, W, 2)
+    sens_maps : ndarray
+        Sensitivity maps (image space)
+
+    Returns
+    -------
+    ndarray
+        Result of the operation F (x sens_maps)
+    """
+    return fastmri.fft2c(fastmri.complex_mul(x, sens_maps))
+
+
+def sens_reduce(k: torch.Tensor, sens_maps: torch.Tensor) -> torch.Tensor:
+    """
+    Calculates F^{-1}(k) * conj(sens_maps)
+    where conj(sens_maps) is the element-wise applied complex conjugate
+
+    Parameters
+    ----------
+    k : ndarray
+        Multi-channel k-space of shape (B, C, H, W, 2)
+    sens_maps : ndarray
+        Sensitivity maps (image space)
+
+    Returns
+    -------
+    ndarray
+        Result of the operation F^{-1}(k) * conj(sens_maps)
+    """
+    return fastmri.complex_mul(
+        fastmri.ifft2c(k), fastmri.complex_conj(sens_maps)
+    ).sum(dim=1, keepdim=True)
+
+
+def chans_to_batch_dim(x: torch.Tensor) -> Tuple[torch.Tensor, int]:
+    """Reshapes batched multi-channel samples into multiple single channel samples.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        x has shape (b, c, h, w, 2)
+
+    Returns
+    -------
+    Tuple[torch.Tensor, int]
+        tensor of shape (b * c, 1, h, w, 2), b
+    """
+    b, c, h, w, comp = x.shape
+    return x.view(b * c, 1, h, w, comp), b
+
+
+def batch_chans_to_chan_dim(x: torch.Tensor, batch_size: int) -> torch.Tensor:
+    """Reshapes batched independent samples into original multi-channel samples.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        tensor of shape (b * c, 1, h, w, 2)
+    batch_size : int
+        batch size
+
+    Returns
+    -------
+    torch.Tensor
+        original multi-channel tensor of shape (b, c, h, w, 2)
+    """
+    bc, _, h, w, comp = x.shape
+    c = bc // batch_size
+    return x.view(batch_size, c, h, w, comp)

models/lightning/mri_module.py

@@ -0,0 +1,402 @@
+"""
+Modified for use in <TODO: paper name>
+- minified and removed extraneous abstractions
+- updated to latest version of lightning
+
+Copyright (c) Facebook, Inc. and its affiliates.
+
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+from collections import defaultdict
+from io import BytesIO
+import pathlib
+import os
+from argparse import ArgumentParser
+from collections import defaultdict
+
+import numpy as np
+import wandb
+import lightning as L
+import torch
+from torchmetrics.metric import Metric
+import matplotlib
+import matplotlib.pyplot as plt
+from PIL import Image
+
+matplotlib.use("Agg")
+
+from fastmri import evaluate
+
+
+class DistributedMetricSum(Metric):
+    def __init__(self, dist_sync_on_step=True):
+        super().__init__(dist_sync_on_step=dist_sync_on_step)
+
+        self.add_state(
+            "quantity", default=torch.tensor(0.0), dist_reduce_fx="sum"
+        )
+
+    def update(self, batch: torch.Tensor):  # type: ignore
+        self.quantity += batch
+
+    def compute(self):
+        return self.quantity
+
+
+class MriModule(L.LightningModule):
+    """
+    Abstract super class for deep learning reconstruction models.
+
+    This is a subclass of the LightningModule class from lightning,
+    with some additional functionality specific to fastMRI:
+        - Evaluating reconstructions
+        - Visualization
+
+    To implement a new reconstruction model, inherit from this class and
+    implement the following methods:
+        - training_step: Define what happens in one step of training
+        - validation_step: Define what happens in one step of validation
+        - test_step: Define what happens in one step of testing
+        - configure_optimizers: Create and return the optimizers
+
+    Other methods from LightningModule can be overridden as needed.
+    """
+
+    def __init__(self, num_log_images: int = 16):
+        """
+        Initialize the MRI module.
+
+        Parameters
+        ----------
+        num_log_images : int, optional
+            Number of images to log. Defaults to 16.
+        """
+        super().__init__()
+
+        self.num_log_images = num_log_images
+        self.val_log_indices = [1, 2, 3, 4, 5]
+        self.val_batch_results = []
+
+        self.NMSE = DistributedMetricSum()
+        self.SSIM = DistributedMetricSum()
+        self.PSNR = DistributedMetricSum()
+        self.ValLoss = DistributedMetricSum()
+        self.TotExamples = DistributedMetricSum()
+        self.TotSliceExamples = DistributedMetricSum()
+
+    def log_image(self, name, image):
+        if self.logger != None:
+            self.logger.log_image(
+                key=f"{name}", images=[image], caption=[{self.global_step}]
+            )
+
+    def on_validation_batch_end(
+        self, outputs, batch, batch_idx, dataloader_idx=0
+    ):
+        # breakpoint()
+        val_logs = outputs
+
+        mse_vals = defaultdict(dict)
+        target_norms = defaultdict(dict)
+        ssim_vals = defaultdict(dict)
+        max_vals = dict()
+
+        for i, fname in enumerate(val_logs["fname"]):
+            if i == 0 and batch_idx in self.val_log_indices:
+                key = f"val_images_idx_{batch_idx}"
+                target = val_logs["target"][i].unsqueeze(0)
+                output = val_logs["output"][i].unsqueeze(0)
+                error = torch.abs(target - output)
+                output = output / output.max()
+                target = target / target.max()
+                error = error / error.max()
+                self.log_image(f"{key}/target", target)
+                self.log_image(f"{key}/reconstruction", output)
+                self.log_image(f"{key}/error", error)
+            slice_num = int(val_logs["slice_num"][i].cpu())
+
+            maxval = val_logs["max_value"][i].cpu().numpy()
+            output = val_logs["output"][i].cpu().numpy()
+            target = val_logs["target"][i].cpu().numpy()
+            mse_vals[fname][slice_num] = torch.tensor(
+                evaluate.mse(target, output)
+            ).view(1)
+            target_norms[fname][slice_num] = torch.tensor(
+                evaluate.mse(target, np.zeros_like(target))
+            ).view(1)
+            ssim_vals[fname][slice_num] = torch.tensor(
+                evaluate.ssim(
+                    target[None, ...], output[None, ...], maxval=maxval
+                )
+            ).view(1)
+            max_vals[fname] = maxval
+
+        self.val_batch_results.append(
+            {
+                "slug": val_logs["slug"],
+                "val_loss": val_logs["val_loss"],
+                "mse_vals": dict(mse_vals),
+                "target_norms": dict(target_norms),
+                "ssim_vals": dict(ssim_vals),
+                "max_vals": max_vals,
+            }
+        )
+
+    def on_validation_epoch_end(self):
+        val_logs = self.val_batch_results
+
+        dataset_metrics = defaultdict(
+            lambda: {
+                "losses": [],
+                "mse_vals": defaultdict(dict),
+                "target_norms": defaultdict(dict),
+                "ssim_vals": defaultdict(dict),
+                "max_vals": dict(),
+            }
+        )
+
+        # use dict updates to handle duplicate slices
+        for val_log in val_logs:
+            slug = val_log["slug"]
+            dataset_metrics[slug]["losses"].append(val_log["val_loss"].view(-1))
+
+            for k in val_log["mse_vals"].keys():
+                dataset_metrics[slug]["mse_vals"][k].update(
+                    val_log["mse_vals"][k]
+                )
+            for k in val_log["target_norms"].keys():
+                dataset_metrics[slug]["target_norms"][k].update(
+                    val_log["target_norms"][k]
+                )
+            for k in val_log["ssim_vals"].keys():
+                dataset_metrics[slug]["ssim_vals"][k].update(
+                    val_log["ssim_vals"][k]
+                )
+            for k in val_log["max_vals"]:
+                dataset_metrics[slug]["max_vals"][k] = val_log["max_vals"][k]
+
+        metrics_to_plot = {"psnr": [], "ssim": [], "nmse": []}
+        slugs = []
+
+        for slug, metrics_data in dataset_metrics.items():
+            mse_vals, target_norms, ssim_vals, max_vals, losses = (
+                metrics_data["mse_vals"],
+                metrics_data["target_norms"],
+                metrics_data["ssim_vals"],
+                metrics_data["max_vals"],
+                metrics_data["losses"],
+            )
+            # check to make sure we have all files in all metrics
+            assert (
+                mse_vals.keys()
+                == target_norms.keys()
+                == ssim_vals.keys()
+                == max_vals.keys()
+            )
+
+            # apply means across image volumes
+            metrics = {"nmse": 0, "ssim": 0, "psnr": 0}
+            metric_values = {
+                "nmse": [],
+                "ssim": [],
+                "psnr": [],
+            }  # to store individual values for std
+            local_examples = 0
+
+            for fname in mse_vals.keys():
+                local_examples = local_examples + 1
+                mse_val = torch.mean(
+                    torch.cat([v.view(-1) for _, v in mse_vals[fname].items()])
+                )
+                target_norm = torch.mean(
+                    torch.cat(
+                        [v.view(-1) for _, v in target_norms[fname].items()]
+                    )
+                )
+                nmse = mse_val / target_norm
+                psnr = 20 * torch.log10(
+                    torch.tensor(
+                        max_vals[fname],
+                        dtype=mse_val.dtype,
+                        device=mse_val.device,
+                    )
+                ) - 10 * torch.log10(mse_val)
+                ssim = torch.mean(
+                    torch.cat([v.view(-1) for _, v in ssim_vals[fname].items()])
+                )
+
+                # Accumulate metric values
+                metrics["nmse"] += nmse
+                metrics["psnr"] += psnr
+                metrics["ssim"] += ssim
+
+                # Store individual metric values for std calculation
+                metric_values["nmse"].append(nmse)
+                metric_values["psnr"].append(psnr)
+                metric_values["ssim"].append(ssim)
+
+            # reduce across ddp via sum
+            metrics["nmse"] = self.NMSE(metrics["nmse"])
+            metrics["ssim"] = self.SSIM(metrics["ssim"])
+            metrics["psnr"] = self.PSNR(metrics["psnr"])
+
+            tot_examples = self.TotExamples(torch.tensor(local_examples))
+            val_loss = self.ValLoss(torch.sum(torch.cat(losses))) # type: ignore
+            tot_slice_examples = self.TotSliceExamples(
+                torch.tensor(len(losses), dtype=torch.float)
+            )
+
+            metrics_to_plot["nmse"].append(
+                (
+                    (metrics["nmse"] / tot_examples).item(),
+                    torch.std(torch.stack(metric_values["nmse"])).item(),
+                )
+            )
+            metrics_to_plot["psnr"].append(
+                (
+                    (metrics["psnr"] / tot_examples).item(),
+                    torch.std(torch.stack(metric_values["psnr"])).item(),
+                )
+            )
+            metrics_to_plot["ssim"].append(
+                (
+                    (metrics["ssim"] / tot_examples).item(),
+                    torch.std(torch.stack(metric_values["ssim"])).item(),
+                )
+            )
+            slugs.append(slug)
+
+            # Log the mean values
+            self.log(
+                f"{slug}--validation_loss",
+                val_loss / tot_slice_examples,
+                prog_bar=True,
+            )
+            for metric, value in metrics.items():
+                self.log(f"{slug}--val_metrics_{metric}", value / tot_examples)
+
+            # Calculate and log the standard deviation for each metric
+            for metric, values in metric_values.items():
+                std_value = torch.std(torch.stack(values))
+                self.log(f"{slug}--val_metrics_{metric}_std", std_value)
+
+        # generate graph
+        # breakpoint()
+        for metric_name, values in metrics_to_plot.items():
+            scores = [val[0] for val in values]
+            std_devs = [val[1] for val in values]
+
+            plt.figure(figsize=(10, 6))
+            plt.bar(slugs, scores, yerr=std_devs, capsize=5)
+            plt.xlabel("Dataset Slug")
+            plt.ylabel(f"{metric_name.upper()} Score")
+            plt.title(
+                f"{metric_name.upper()} per Dataset with Standard Deviation"
+            )
+            plt.xticks(rotation=45)
+            plt.tight_layout()
+
+            # Save the plot
+            buf = BytesIO()
+            plt.savefig(buf, format="png")
+            buf.seek(0)
+            image = Image.open(buf)
+            image_array = np.array(image)
+            self.log_image(f"summary_plot_{metric_name}", image_array)
+            buf.close()
+            plt.close()
+
+    def OLD_on_validation_epoch_end(self):
+        val_logs = self.val_batch_results
+
+        # aggregate losses
+        losses = []
+        mse_vals = defaultdict(dict)
+        target_norms = defaultdict(dict)
+        ssim_vals = defaultdict(dict)
+        max_vals = dict()
+
+        # use dict updates to handle duplicate slices
+        for val_log in val_logs:
+            losses.append(val_log["val_loss"].view(-1))
+
+            for k in val_log["mse_vals"].keys():
+                mse_vals[k].update(val_log["mse_vals"][k])
+            for k in val_log["target_norms"].keys():
+                target_norms[k].update(val_log["target_norms"][k])
+            for k in val_log["ssim_vals"].keys():
+                ssim_vals[k].update(val_log["ssim_vals"][k])
+            for k in val_log["max_vals"]:
+                max_vals[k] = val_log["max_vals"][k]
+
+        # check to make sure we have all files in all metrics
+        assert (
+            mse_vals.keys()
+            == target_norms.keys()
+            == ssim_vals.keys()
+            == max_vals.keys()
+        )
+
+        # apply means across image volumes
+        metrics = {"nmse": 0, "ssim": 0, "psnr": 0}
+        local_examples = 0
+        for fname in mse_vals.keys():
+            local_examples = local_examples + 1
+            mse_val = torch.mean(
+                torch.cat([v.view(-1) for _, v in mse_vals[fname].items()])
+            )
+            target_norm = torch.mean(
+                torch.cat([v.view(-1) for _, v in target_norms[fname].items()])
+            )
+            metrics["nmse"] = metrics["nmse"] + mse_val / target_norm
+            metrics["psnr"] = (
+                metrics["psnr"]
+                + 20
+                * torch.log10(
+                    torch.tensor(
+                        max_vals[fname],
+                        dtype=mse_val.dtype,
+                        device=mse_val.device,
+                    )
+                )
+                - 10 * torch.log10(mse_val)
+            )
+            metrics["ssim"] = metrics["ssim"] + torch.mean(
+                torch.cat([v.view(-1) for _, v in ssim_vals[fname].items()])
+            )
+
+        # reduce across ddp via sum
+        metrics["nmse"] = self.NMSE(metrics["nmse"])
+        metrics["ssim"] = self.SSIM(metrics["ssim"])
+        metrics["psnr"] = self.PSNR(metrics["psnr"])
+
+        tot_examples = self.TotExamples(torch.tensor(local_examples))
+        val_loss = self.ValLoss(torch.sum(torch.cat(losses)))
+        tot_slice_examples = self.TotSliceExamples(
+            torch.tensor(len(losses), dtype=torch.float)
+        )
+
+        self.log(
+            "validation_loss", val_loss / tot_slice_examples, prog_bar=True
+        )
+        for metric, value in metrics.items():
+            self.log(f"val_metrics_{metric}", value / tot_examples)
+
+    @staticmethod
+    def add_model_specific_args(parent_parser):  # pragma: no-cover
+        """
+        Define parameters that only apply to this model
+        """
+        parser = ArgumentParser(parents=[parent_parser], add_help=False)
+
+        # logging params
+        parser.add_argument(
+            "--num_log_images",
+            default=16,
+            type=int,
+            help="Number of images to log to Tensorboard",
+        )
+
+        return parser

models/lightning/no_shared_module.py

@@ -0,0 +1,274 @@
+from argparse import ArgumentParser
+from typing import Tuple
+
+import torch
+
+import fastmri
+from fastmri import transforms
+from models.no_shared import NOShared
+
+from models.lightning.mri_module import MriModule
+from type_utils import tuple_type
+
+
+class NOSharedModule(MriModule):
+    """
+    NO-Shared training module.
+    """
+
+    def __init__(
+        self,
+        num_cascades: int = 12,
+        pools: int = 4,
+        chans: int = 18,
+        sens_pools: int = 4,
+        sens_chans: int = 8,
+        gno_pools: int = 4,
+        gno_chans: int = 16,
+        gno_radius_cutoff: float = 0.02,
+        gno_kernel_shape: Tuple[int, int] = (6, 7),
+        radius_cutoff: float = 0.02,
+        kernel_shape: Tuple[int, int] = (6, 7),
+        in_shape: Tuple[int, int] = (320, 320),
+        use_dc_term: bool = True,
+        lr: float = 0.0003,
+        lr_step_size: int = 40,
+        lr_gamma: float = 0.1,
+        weight_decay: float = 0.0,
+        **kwargs,
+    ):
+        """
+        Parameters
+        ----------
+        num_cascades : int
+            Number of cascades (i.e., layers) for the variational network.
+        pools : int
+            Number of downsampling and upsampling layers for the cascade U-Net.
+        chans : int
+            Number of channels for the cascade U-Net.
+        sens_pools : int
+            Number of downsampling and upsampling layers for the sensitivity map U-Net.
+        sens_chans : int
+            Number of channels for the sensitivity map U-Net.
+        lr : float
+            Learning rate.
+        lr_step_size : int
+            Learning rate step size.
+        lr_gamma : float
+            Learning rate gamma decay.
+        weight_decay : float
+            Parameter for penalizing weights norm.
+        """
+        super().__init__(**kwargs)
+        self.save_hyperparameters()
+
+        self.num_cascades = num_cascades
+        self.pools = pools
+        self.chans = chans
+        self.sens_pools = sens_pools
+        self.sens_chans = sens_chans
+        self.gno_pools = gno_pools
+        self.gno_chans = gno_chans
+        self.gno_radius_cutoff = gno_radius_cutoff
+        self.gno_kernel_shape = gno_kernel_shape
+        self.radius_cutoff = radius_cutoff
+        self.kernel_shape = kernel_shape
+        self.in_shape = in_shape
+        self.use_dc_term = use_dc_term
+        self.lr = lr
+        self.lr_step_size = lr_step_size
+        self.lr_gamma = lr_gamma
+        self.weight_decay = weight_decay
+
+        self.model = NOShared(
+            num_cascades=self.num_cascades,
+            sens_chans=self.sens_chans,
+            sens_pools=self.sens_pools,
+            chans=self.chans,
+            pools=self.pools,
+            gno_chans=self.gno_chans,
+            gno_pools=self.gno_pools,
+            gno_radius_cutoff=self.gno_radius_cutoff,
+            gno_kernel_shape=self.gno_kernel_shape,
+            radius_cutoff=radius_cutoff,
+            kernel_shape=kernel_shape,
+            in_shape=in_shape,
+            use_dc_term=use_dc_term,
+        )
+
+        self.criterion = fastmri.SSIMLoss()
+        self.num_params = sum(p.numel() for p in self.parameters())
+
+    def forward(self, masked_kspace, mask, num_low_frequencies):
+        return self.model(masked_kspace, mask, num_low_frequencies)
+
+    def training_step(self, batch, batch_idx):
+        output = self.forward(
+            batch.masked_kspace, batch.mask, batch.num_low_frequencies
+        )
+
+        target, output = transforms.center_crop_to_smallest(
+            batch.target, output
+        )
+        loss = self.criterion(
+            output.unsqueeze(1), target.unsqueeze(1), data_range=batch.max_value
+        )
+
+        self.log("train_loss", loss, on_step=True, on_epoch=True)
+        self.log("epoch", int(self.current_epoch), on_step=True, on_epoch=True)
+
+        return loss
+
+    def validation_step(self, batch, batch_idx, dataloader_idx=0):
+        dataloaders = self.trainer.val_dataloaders
+        slug = list(dataloaders.keys())[dataloader_idx]
+
+        output = self.forward(
+            batch.masked_kspace, batch.mask, batch.num_low_frequencies
+        )
+
+        target, output = transforms.center_crop_to_smallest(
+            batch.target, output
+        )
+
+        loss = self.criterion(
+            output.unsqueeze(1),
+            target.unsqueeze(1),
+            data_range=batch.max_value,
+        )
+
+        return {
+            "slug": slug,
+            "fname": batch.fname,
+            "slice_num": batch.slice_num,
+            "max_value": batch.max_value,
+            "output": output,
+            "target": target,
+            "val_loss": loss,
+        }
+
+    def configure_optimizers(self):
+        optim = torch.optim.Adam(
+            self.parameters(), lr=self.lr, weight_decay=self.weight_decay
+        )
+        scheduler = torch.optim.lr_scheduler.StepLR(
+            optim, self.lr_step_size, self.lr_gamma
+        )
+
+        return [optim], [scheduler]
+
+    @staticmethod
+    def add_model_specific_args(parent_parser):
+        """
+        Define parameters that only apply to this model
+        """
+        parser = ArgumentParser(parents=[parent_parser], add_help=False)
+        parser = MriModule.add_model_specific_args(parser)
+
+        # network params
+        parser.add_argument(
+            "--num_cascades",
+            default=12,
+            type=int,
+            help="Number of VarNet cascades",
+        )
+        parser.add_argument(
+            "--pools",
+            default=4,
+            type=int,
+            help="Number of U-Net pooling layers in VarNet blocks",
+        )
+        parser.add_argument(
+            "--chans",
+            default=18,
+            type=int,
+            help="Number of channels for U-Net in VarNet blocks",
+        )
+        parser.add_argument(
+            "--sens_pools",
+            default=4,
+            type=int,
+            help=(
+                "Number of pooling layers for sense map estimation U-Net in"
+                " VarNet"
+            ),
+        )
+        parser.add_argument(
+            "--sens_chans",
+            default=8,
+            type=float,
+            help="Number of channels for sense map estimation U-Net in VarNet",
+        )
+        parser.add_argument(
+            "--gno_pools",
+            default=4,
+            type=int,
+            help=("Number of pooling layers for GNO"),
+        )
+        parser.add_argument(
+            "--gno_chans",
+            default=16,
+            type=int,
+            help="Number of channels for GNO",
+        )
+        parser.add_argument(
+            "--gno_radius_cutoff",
+            default=0.02,
+            type=float,
+            help="GNO module radius_cutoff",
+        )
+        parser.add_argument(
+            "--gno_kernel_shape",
+            default=(6, 7),
+            type=tuple_type,
+            help="GNO module kernel_shape. Ex: (6, 7)",
+        )
+        parser.add_argument(
+            "--radius_cutoff",
+            default=0.02,
+            type=float,
+            help="DISCO module radius_cutoff",
+        )
+        parser.add_argument(
+            "--kernel_shape",
+            default=(6, 7),
+            type=tuple_type,
+            help="DISCO module kernel_shape. Ex: (6, 7)",
+        )
+        parser.add_argument(
+            "--in_shape",
+            default=(320, 320),
+            type=tuple_type,
+            help="Spatial dimensions of masked_kspace samples. Ex: (640, 320)",
+        )
+        parser.add_argument(
+            "--use_dc_term",
+            default=True,
+            type=bool,
+            help="Whether to use the DC term in the unrolled iterative update step",
+        )
+
+        # training params (opt)
+        parser.add_argument(
+            "--lr", default=0.0003, type=float, help="Adam learning rate"
+        )
+        parser.add_argument(
+            "--lr_step_size",
+            default=40,
+            type=int,
+            help="Epoch at which to decrease step size",
+        )
+        parser.add_argument(
+            "--lr_gamma",
+            default=0.1,
+            type=float,
+            help="Extent to which step size should be decreased",
+        )
+        parser.add_argument(
+            "--weight_decay",
+            default=0.0,
+            type=float,
+            help="Strength of weight decay regularization",
+        )
+
+        return parser

models/lightning/no_varnet_module.py

@@ -0,0 +1,299 @@
+from argparse import ArgumentParser
+from typing import Tuple
+
+import torch
+
+import fastmri
+from fastmri import transforms
+from models.no_varnet import NOVarnet
+
+from models.lightning.mri_module import MriModule
+# from type_utils import tuple_type
+
+def tuple_type(strings):
+    strings = strings.replace("(", "").replace(")", "").replace(" ", "")
+    mapped_int = map(int, strings.split(","))
+    return tuple(mapped_int)
+
+
+class NOVarnetModule(MriModule):
+    """
+    NO-Varnet training module.
+    """
+
+    def __init__(
+        self,
+        num_cascades: int = 12,
+        pools: int = 4,
+        chans: int = 18,
+        sens_pools: int = 4,
+        sens_chans: int = 8,
+        gno_pools: int = 4,
+        gno_chans: int = 16,
+        gno_radius_cutoff: float = 0.02,
+        gno_kernel_shape: Tuple[int, int] = (6, 7),
+        radius_cutoff: float = 0.02,
+        kernel_shape: Tuple[int, int] = (6, 7),
+        in_shape: Tuple[int, int] = (320, 320),
+        use_dc_term: bool = True,
+        lr: float = 0.0003,
+        lr_step_size: int = 40,
+        lr_gamma: float = 0.1,
+        weight_decay: float = 0.0,
+        reduction_method: str = "rss",
+        skip_method: str = "add",
+        **kwargs,
+    ):
+        """
+        Parameters
+        ----------
+        num_cascades : int
+            Number of cascades (i.e., layers) for the variational network.
+        pools : int
+            Number of downsampling and upsampling layers for the cascade U-Net.
+        chans : int
+            Number of channels for the cascade U-Net.
+        sens_pools : int
+            Number of downsampling and upsampling layers for the sensitivity map U-Net.
+        sens_chans : int
+            Number of channels for the sensitivity map U-Net.
+        lr : float
+            Learning rate.
+        lr_step_size : int
+            Learning rate step size.
+        lr_gamma : float
+            Learning rate gamma decay.
+        weight_decay : float
+            Parameter for penalizing weights norm.
+        """
+        super().__init__(**kwargs)
+        self.save_hyperparameters()
+
+        self.num_cascades = num_cascades
+        self.pools = pools
+        self.chans = chans
+        self.sens_pools = sens_pools
+        self.sens_chans = sens_chans
+        self.gno_pools = gno_pools
+        self.gno_chans = gno_chans
+        self.gno_radius_cutoff = gno_radius_cutoff
+        self.gno_kernel_shape = gno_kernel_shape
+        self.radius_cutoff = radius_cutoff
+        self.kernel_shape = kernel_shape
+        self.in_shape = in_shape
+        self.use_dc_term = use_dc_term
+        self.lr = lr
+        self.lr_step_size = lr_step_size
+        self.lr_gamma = lr_gamma
+        self.weight_decay = weight_decay
+        self.reduction_method = reduction_method
+        self.skip_method = skip_method
+
+        self.model = NOVarnet(
+            num_cascades=self.num_cascades,
+            sens_chans=self.sens_chans,
+            sens_pools=self.sens_pools,
+            chans=self.chans,
+            pools=self.pools,
+            gno_chans=self.gno_chans,
+            gno_pools=self.gno_pools,
+            gno_radius_cutoff=self.gno_radius_cutoff,
+            gno_kernel_shape=self.gno_kernel_shape,
+            radius_cutoff=radius_cutoff,
+            kernel_shape=kernel_shape,
+            in_shape=in_shape,
+            use_dc_term=use_dc_term,
+            reduction_method=reduction_method,
+            skip_method=skip_method,
+        )
+
+        self.criterion = fastmri.SSIMLoss()
+        self.num_params = sum(p.numel() for p in self.parameters())
+
+    def forward(self, masked_kspace, mask, num_low_frequencies):
+        return self.model(masked_kspace, mask, num_low_frequencies)
+
+    def training_step(self, batch, batch_idx):
+        output = self.forward(
+            batch.masked_kspace, batch.mask, batch.num_low_frequencies
+        )
+
+        target, output = transforms.center_crop_to_smallest(batch.target, output)
+        loss = self.criterion(
+            output.unsqueeze(1), target.unsqueeze(1), data_range=batch.max_value
+        )
+
+        self.log("train_loss", loss, on_step=True, on_epoch=True)
+        self.log("epoch", int(self.current_epoch), on_step=True, on_epoch=True)
+
+        return loss
+
+    def validation_step(self, batch, batch_idx, dataloader_idx=0):
+        dataloaders = self.trainer.val_dataloaders
+        slug = list(dataloaders.keys())[dataloader_idx]
+
+        output = self.forward(
+            batch.masked_kspace, batch.mask, batch.num_low_frequencies
+        )
+
+        target, output = transforms.center_crop_to_smallest(batch.target, output)
+
+        loss = self.criterion(
+            output.unsqueeze(1),
+            target.unsqueeze(1),
+            data_range=batch.max_value,
+        )
+
+        return {
+            "slug": slug,
+            "fname": batch.fname,
+            "slice_num": batch.slice_num,
+            "max_value": batch.max_value,
+            "output": output,
+            "target": target,
+            "val_loss": loss,
+        }
+
+    def configure_optimizers(self):
+        optim = torch.optim.Adam(
+            self.parameters(), lr=self.lr, weight_decay=self.weight_decay
+        )
+        scheduler = torch.optim.lr_scheduler.StepLR(
+            optim, self.lr_step_size, self.lr_gamma
+        )
+
+        return [optim], [scheduler]
+
+    @staticmethod
+    def add_model_specific_args(parent_parser):
+        """
+        Define parameters that only apply to this model
+        """
+        parser = ArgumentParser(parents=[parent_parser], add_help=False)
+        parser = MriModule.add_model_specific_args(parser)
+
+        # network params
+        parser.add_argument(
+            "--num_cascades",
+            default=12,
+            type=int,
+            help="Number of VarNet cascades",
+        )
+        parser.add_argument(
+            "--pools",
+            default=4,
+            type=int,
+            help="Number of U-Net pooling layers in VarNet blocks",
+        )
+        parser.add_argument(
+            "--chans",
+            default=18,
+            type=int,
+            help="Number of channels for U-Net in VarNet blocks",
+        )
+        parser.add_argument(
+            "--sens_pools",
+            default=4,
+            type=int,
+            help=(
+                "Number of pooling layers for sense map estimation U-Net in" " VarNet"
+            ),
+        )
+        parser.add_argument(
+            "--sens_chans",
+            default=8,
+            type=float,
+            help="Number of channels for sense map estimation U-Net in VarNet",
+        )
+        parser.add_argument(
+            "--gno_pools",
+            default=4,
+            type=int,
+            help=("Number of pooling layers for GNO"),
+        )
+        parser.add_argument(
+            "--gno_chans",
+            default=16,
+            type=int,
+            help="Number of channels for GNO",
+        )
+        parser.add_argument(
+            "--gno_radius_cutoff",
+            default=0.02,
+            type=float,
+            required=True,
+            help="GNO module radius_cutoff",
+        )
+        parser.add_argument(
+            "--gno_kernel_shape",
+            default=(6, 7),
+            type=tuple_type,
+            required=True,
+            help="GNO module kernel_shape. Ex: (6, 7)",
+        )
+        parser.add_argument(
+            "--radius_cutoff",
+            default=0.01,
+            type=float,
+            required=True,
+            help="DISCO module radius_cutoff",
+        )
+        parser.add_argument(
+            "--kernel_shape",
+            default=(6, 7),
+            type=tuple_type,
+            required=True,
+            help="DISCO module kernel_shape. Ex: (6, 7)",
+        )
+        parser.add_argument(
+            "--in_shape",
+            default=(640, 320),
+            type=tuple_type,
+            required=True,
+            help="Spatial dimensions of masked_kspace samples. Ex: (640, 320)",
+        )
+        parser.add_argument(
+            "--use_dc_term",
+            default=True,
+            type=bool,
+            help="Whether to use the DC term in the unrolled iterative update step",
+        )
+
+        # training params (opt)
+        parser.add_argument(
+            "--lr", default=0.0003, type=float, help="Adam learning rate"
+        )
+        parser.add_argument(
+            "--lr_step_size",
+            default=40,
+            type=int,
+            help="Epoch at which to decrease step size",
+        )
+        parser.add_argument(
+            "--lr_gamma",
+            default=0.1,
+            type=float,
+            help="Extent to which step size should be decreased",
+        )
+        parser.add_argument(
+            "--weight_decay",
+            default=0.0,
+            type=float,
+            help="Strength of weight decay regularization",
+        )
+        parser.add_argument(
+            "--reduction_method",
+            default="rss",
+            type=str,
+            choices=["rss", "batch"],
+            help="Reduction method used to reduce multi-channel k-space data before inpainting module. Read documentation of GNO for more information.",
+        )
+        parser.add_argument(
+            "--skip_method",
+            default="add_inv",
+            type=str,
+            choices=["add_inv", "add", "concat", "replace"],
+            help="Method for skip connection around inpainting module.",
+        )
+
+        return parser

models/lightning/no_varnet_nokno_module.py

@@ -0,0 +1,294 @@
+from argparse import ArgumentParser
+from typing import Tuple
+
+import torch
+
+import fastmri
+from fastmri import transforms
+
+from models.lightning.mri_module import MriModule
+from models.no_varnet_nokno import NOVarnet_no_KNO
+from type_utils import tuple_type
+
+
+class NOVarnet_no_KNOModule(MriModule):
+    """
+    NO-Varnet training module.
+    """
+
+    def __init__(
+        self,
+        num_cascades: int = 12,
+        pools: int = 4,
+        chans: int = 18,
+        sens_pools: int = 4,
+        sens_chans: int = 8,
+        gno_pools: int = 4,
+        gno_chans: int = 16,
+        gno_radius_cutoff: float = 0.02,
+        gno_kernel_shape: Tuple[int, int] = (6, 7),
+        radius_cutoff: float = 0.02,
+        kernel_shape: Tuple[int, int] = (6, 7),
+        in_shape: Tuple[int, int] = (320, 320),
+        use_dc_term: bool = True,
+        lr: float = 0.0003,
+        lr_step_size: int = 40,
+        lr_gamma: float = 0.1,
+        weight_decay: float = 0.0,
+        reduction_method: str = "rss",
+        skip_method: str = "add",
+        **kwargs,
+    ):
+        """
+        Parameters
+        ----------
+        num_cascades : int
+            Number of cascades (i.e., layers) for the variational network.
+        pools : int
+            Number of downsampling and upsampling layers for the cascade U-Net.
+        chans : int
+            Number of channels for the cascade U-Net.
+        sens_pools : int
+            Number of downsampling and upsampling layers for the sensitivity map U-Net.
+        sens_chans : int
+            Number of channels for the sensitivity map U-Net.
+        lr : float
+            Learning rate.
+        lr_step_size : int
+            Learning rate step size.
+        lr_gamma : float
+            Learning rate gamma decay.
+        weight_decay : float
+            Parameter for penalizing weights norm.
+        """
+        super().__init__(**kwargs)
+        self.save_hyperparameters()
+
+        self.num_cascades = num_cascades
+        self.pools = pools
+        self.chans = chans
+        self.sens_pools = sens_pools
+        self.sens_chans = sens_chans
+        self.gno_pools = gno_pools
+        self.gno_chans = gno_chans
+        self.gno_radius_cutoff = gno_radius_cutoff
+        self.gno_kernel_shape = gno_kernel_shape
+        self.radius_cutoff = radius_cutoff
+        self.kernel_shape = kernel_shape
+        self.in_shape = in_shape
+        self.use_dc_term = use_dc_term
+        self.lr = lr
+        self.lr_step_size = lr_step_size
+        self.lr_gamma = lr_gamma
+        self.weight_decay = weight_decay
+        self.reduction_method = reduction_method
+        self.skip_method = skip_method
+
+        self.model = NOVarnet_no_KNO(
+            num_cascades=self.num_cascades,
+            sens_chans=self.sens_chans,
+            sens_pools=self.sens_pools,
+            chans=self.chans,
+            pools=self.pools,
+            gno_chans=self.gno_chans,
+            gno_pools=self.gno_pools,
+            gno_radius_cutoff=self.gno_radius_cutoff,
+            gno_kernel_shape=self.gno_kernel_shape,
+            radius_cutoff=radius_cutoff,
+            kernel_shape=kernel_shape,
+            in_shape=in_shape,
+            use_dc_term=use_dc_term,
+            reduction_method=reduction_method,
+            skip_method=skip_method,
+        )
+
+        self.criterion = fastmri.SSIMLoss()
+        self.num_params = sum(p.numel() for p in self.parameters())
+
+    def forward(self, masked_kspace, mask, num_low_frequencies):
+        return self.model(masked_kspace, mask, num_low_frequencies)
+
+    def training_step(self, batch, batch_idx):
+        output = self.forward(
+            batch.masked_kspace, batch.mask, batch.num_low_frequencies
+        )
+
+        target, output = transforms.center_crop_to_smallest(batch.target, output)
+        loss = self.criterion(
+            output.unsqueeze(1), target.unsqueeze(1), data_range=batch.max_value
+        )
+
+        self.log("train_loss", loss, on_step=True, on_epoch=True)
+        self.log("epoch", int(self.current_epoch), on_step=True, on_epoch=True)
+
+        return loss
+
+    def validation_step(self, batch, batch_idx, dataloader_idx=0):
+        dataloaders = self.trainer.val_dataloaders
+        slug = list(dataloaders.keys())[dataloader_idx]
+
+        output = self.forward(
+            batch.masked_kspace, batch.mask, batch.num_low_frequencies
+        )
+
+        target, output = transforms.center_crop_to_smallest(batch.target, output)
+
+        loss = self.criterion(
+            output.unsqueeze(1),
+            target.unsqueeze(1),
+            data_range=batch.max_value,
+        )
+
+        return {
+            "slug": slug,
+            "fname": batch.fname,
+            "slice_num": batch.slice_num,
+            "max_value": batch.max_value,
+            "output": output,
+            "target": target,
+            "val_loss": loss,
+        }
+
+    def configure_optimizers(self):
+        optim = torch.optim.Adam(
+            self.parameters(), lr=self.lr, weight_decay=self.weight_decay
+        )
+        scheduler = torch.optim.lr_scheduler.StepLR(
+            optim, self.lr_step_size, self.lr_gamma
+        )
+
+        return [optim], [scheduler]
+
+    @staticmethod
+    def add_model_specific_args(parent_parser):
+        """
+        Define parameters that only apply to this model
+        """
+        parser = ArgumentParser(parents=[parent_parser], add_help=False)
+        parser = MriModule.add_model_specific_args(parser)
+
+        # network params
+        parser.add_argument(
+            "--num_cascades",
+            default=12,
+            type=int,
+            help="Number of VarNet cascades",
+        )
+        parser.add_argument(
+            "--pools",
+            default=4,
+            type=int,
+            help="Number of U-Net pooling layers in VarNet blocks",
+        )
+        parser.add_argument(
+            "--chans",
+            default=18,
+            type=int,
+            help="Number of channels for U-Net in VarNet blocks",
+        )
+        parser.add_argument(
+            "--sens_pools",
+            default=4,
+            type=int,
+            help=(
+                "Number of pooling layers for sense map estimation U-Net in" " VarNet"
+            ),
+        )
+        parser.add_argument(
+            "--sens_chans",
+            default=8,
+            type=float,
+            help="Number of channels for sense map estimation U-Net in VarNet",
+        )
+        parser.add_argument(
+            "--gno_pools",
+            default=4,
+            type=int,
+            help=("Number of pooling layers for GNO"),
+        )
+        parser.add_argument(
+            "--gno_chans",
+            default=16,
+            type=int,
+            help="Number of channels for GNO",
+        )
+        parser.add_argument(
+            "--gno_radius_cutoff",
+            default=0.02,
+            type=float,
+            required=True,
+            help="GNO module radius_cutoff",
+        )
+        parser.add_argument(
+            "--gno_kernel_shape",
+            default=(6, 7),
+            type=tuple_type,
+            required=True,
+            help="GNO module kernel_shape. Ex: (6, 7)",
+        )
+        parser.add_argument(
+            "--radius_cutoff",
+            default=0.01,
+            type=float,
+            required=True,
+            help="DISCO module radius_cutoff",
+        )
+        parser.add_argument(
+            "--kernel_shape",
+            default=(6, 7),
+            type=tuple_type,
+            required=True,
+            help="DISCO module kernel_shape. Ex: (6, 7)",
+        )
+        parser.add_argument(
+            "--in_shape",
+            default=(640, 320),
+            type=tuple_type,
+            required=True,
+            help="Spatial dimensions of masked_kspace samples. Ex: (640, 320)",
+        )
+        parser.add_argument(
+            "--use_dc_term",
+            default=True,
+            type=bool,
+            help="Whether to use the DC term in the unrolled iterative update step",
+        )
+
+        # training params (opt)
+        parser.add_argument(
+            "--lr", default=0.0003, type=float, help="Adam learning rate"
+        )
+        parser.add_argument(
+            "--lr_step_size",
+            default=40,
+            type=int,
+            help="Epoch at which to decrease step size",
+        )
+        parser.add_argument(
+            "--lr_gamma",
+            default=0.1,
+            type=float,
+            help="Extent to which step size should be decreased",
+        )
+        parser.add_argument(
+            "--weight_decay",
+            default=0.0,
+            type=float,
+            help="Strength of weight decay regularization",
+        )
+        parser.add_argument(
+            "--reduction_method",
+            default="rss",
+            type=str,
+            choices=["rss", "batch"],
+            help="Reduction method used to reduce multi-channel k-space data before inpainting module. Read documentation of GNO for more information.",
+        )
+        parser.add_argument(
+            "--skip_method",
+            default="add_inv",
+            type=str,
+            choices=["add_inv", "add", "concat", "replace"],
+            help="Method for skip connection around inpainting module.",
+        )
+
+        return parser

models/lightning/varnet_module.py

@@ -0,0 +1,224 @@
+"""
+Copyright (c) Facebook, Inc. and its affiliates.
+
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+from argparse import ArgumentParser
+
+import torch
+
+import fastmri
+from fastmri import transforms
+from ..varnet import VarNet
+import wandb
+
+from .mri_module import MriModule
+
+
+class VarNetModule(MriModule):
+    """
+    VarNet training module.
+
+    This can be used to train variational networks from the paper:
+
+    A. Sriram et al. End-to-end variational networks for accelerated MRI
+    reconstruction. In International Conference on Medical Image Computing and
+    Computer-Assisted Intervention, 2020.
+
+    which was inspired by the earlier paper:
+
+    K. Hammernik et al. Learning a variational network for reconstruction of
+    accelerated MRI data. Magnetic Resonance inMedicine, 79(6):3055–3071, 2018.
+    """
+
+    def __init__(
+        self,
+        num_cascades: int = 12,
+        pools: int = 4,
+        chans: int = 18,
+        sens_pools: int = 4,
+        sens_chans: int = 8,
+        lr: float = 0.0003,
+        lr_step_size: int = 40,
+        lr_gamma: float = 0.1,
+        weight_decay: float = 0.0,
+        **kwargs,
+    ):
+        """
+        Parameters
+        ----------
+        num_cascades : int
+            Number of cascades (i.e., layers) for the variational network.
+        pools : int
+            Number of downsampling and upsampling layers for the cascade U-Net.
+        chans : int
+            Number of channels for the cascade U-Net.
+        sens_pools : int
+            Number of downsampling and upsampling layers for the sensitivity map U-Net.
+        sens_chans : int
+            Number of channels for the sensitivity map U-Net.
+        lr : float
+            Learning rate.
+        lr_step_size : int
+            Learning rate step size.
+        lr_gamma : float
+            Learning rate gamma decay.
+        weight_decay : float
+            Parameter for penalizing weights norm.
+        num_sense_lines : int, optional
+            Number of low-frequency lines to use for sensitivity map computation.
+            Must be even or `None`. Default `None` will automatically compute the number
+            from masks. Default behavior may cause some slices to use more low-frequency
+            lines than others, when used in conjunction with e.g. the EquispacedMaskFunc
+            defaults. To prevent this, either set `num_sense_lines`, or set
+            `skip_low_freqs` and `skip_around_low_freqs` to `True` in the EquispacedMaskFunc.
+            Note that setting this value may lead to undesired behavior when training on
+            multiple accelerations simultaneously.
+        """
+        super().__init__(**kwargs)
+        self.save_hyperparameters()
+
+        self.num_cascades = num_cascades
+        self.pools = pools
+        self.chans = chans
+        self.sens_pools = sens_pools
+        self.sens_chans = sens_chans
+        self.lr = lr
+        self.lr_step_size = lr_step_size
+        self.lr_gamma = lr_gamma
+        self.weight_decay = weight_decay
+
+        self.varnet = VarNet(
+            num_cascades=self.num_cascades,
+            sens_chans=self.sens_chans,
+            sens_pools=self.sens_pools,
+            chans=self.chans,
+            pools=self.pools,
+        )
+
+        self.criterion = fastmri.SSIMLoss()
+        self.num_params = sum(p.numel() for p in self.parameters())
+
+    def forward(self, masked_kspace, mask, num_low_frequencies):
+        return self.varnet(masked_kspace, mask, num_low_frequencies)
+
+    def training_step(self, batch, batch_idx):
+        output = self.forward(
+            batch.masked_kspace, batch.mask, batch.num_low_frequencies
+        )
+
+        target, output = transforms.center_crop_to_smallest(batch.target, output)
+        loss = self.criterion(
+            output.unsqueeze(1), target.unsqueeze(1), data_range=batch.max_value
+        )
+
+        self.log("train_loss", loss, on_step=True, on_epoch=True)
+        self.log("epoch", int(self.current_epoch), on_step=True, on_epoch=True)
+
+        return loss
+
+    def validation_step(self, batch, batch_idx, dataloader_idx=0):
+        dataloaders = self.trainer.val_dataloaders
+        slug = list(dataloaders.keys())[dataloader_idx]
+
+        # breakpoint()
+        output = self.forward(
+            batch.masked_kspace, batch.mask, batch.num_low_frequencies
+        )
+
+        target, output = transforms.center_crop_to_smallest(batch.target, output)
+
+        loss = self.criterion(
+            output.unsqueeze(1),
+            target.unsqueeze(1),
+            data_range=batch.max_value,
+        )
+
+        return {
+            "slug": slug,
+            "fname": batch.fname,
+            "slice_num": batch.slice_num,
+            "max_value": batch.max_value,
+            "output": output,
+            "target": target,
+            "val_loss": loss,
+        }
+
+    def configure_optimizers(self):
+        optim = torch.optim.Adam(
+            self.parameters(), lr=self.lr, weight_decay=self.weight_decay
+        )
+        scheduler = torch.optim.lr_scheduler.StepLR(
+            optim, self.lr_step_size, self.lr_gamma
+        )
+
+        return [optim], [scheduler]
+
+    @staticmethod
+    def add_model_specific_args(parent_parser):  # pragma: no-cover
+        """
+        Define parameters that only apply to this model
+        """
+        parser = ArgumentParser(parents=[parent_parser], add_help=False)
+        parser = MriModule.add_model_specific_args(parser)
+
+        # network params
+        parser.add_argument(
+            "--num_cascades",
+            default=12,
+            type=int,
+            help="Number of VarNet cascades",
+        )
+        parser.add_argument(
+            "--pools",
+            default=4,
+            type=int,
+            help="Number of U-Net pooling layers in VarNet blocks",
+        )
+        parser.add_argument(
+            "--chans",
+            default=18,
+            type=int,
+            help="Number of channels for U-Net in VarNet blocks",
+        )
+        parser.add_argument(
+            "--sens_pools",
+            default=4,
+            type=int,
+            help=(
+                "Number of pooling layers for sense map estimation U-Net in" " VarNet"
+            ),
+        )
+        parser.add_argument(
+            "--sens_chans",
+            default=8,
+            type=float,
+            help="Number of channels for sense map estimation U-Net in VarNet",
+        )
+
+        # training params (opt)
+        parser.add_argument(
+            "--lr", default=0.0003, type=float, help="Adam learning rate"
+        )
+        parser.add_argument(
+            "--lr_step_size",
+            default=40,
+            type=int,
+            help="Epoch at which to decrease step size",
+        )
+        parser.add_argument(
+            "--lr_gamma",
+            default=0.1,
+            type=float,
+            help="Extent to which step size should be decreased",
+        )
+        parser.add_argument(
+            "--weight_decay",
+            default=0.0,
+            type=float,
+            help="Strength of weight decay regularization",
+        )
+
+        return parser

models/no_shared.py

@@ -0,0 +1,467 @@
+import math
+from typing import List, Literal, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import fastmri
+from fastmri.transforms import (
+    batched_mask_center,
+    batch_chans_to_chan_dim,
+    chans_to_batch_dim,
+    sens_reduce,
+    sens_expand,
+)
+from models.udno import UDNO
+
+
+class NormUDNO(nn.Module):
+    """
+    Normalized UDNO model.
+
+    Inputs are normalized before the UDNO for numerically stable training.
+    """
+
+    def __init__(
+        self,
+        chans: int,
+        num_pool_layers: int,
+        radius_cutoff: float,
+        in_shape: Tuple[int, int],
+        kernel_shape: Tuple[int, int],
+        in_chans: int = 2,
+        out_chans: int = 2,
+        drop_prob: float = 0.0,
+    ):
+        """
+        Initialize the VarNet model.
+
+        Parameters
+        ----------
+        chans : int
+            Number of output channels of the first convolution layer.
+        num_pools : int
+            Number of down-sampling and up-sampling layers.
+        in_chans : int, optional
+            Number of channels in the input to the U-Net model. Default is 2.
+        out_chans : int, optional
+            Number of channels in the output to the U-Net model. Default is 2.
+        drop_prob : float, optional
+            Dropout probability. Default is 0.0.
+        """
+        super().__init__()
+
+        self.udno = UDNO(
+            in_chans=in_chans,
+            out_chans=out_chans,
+            radius_cutoff=radius_cutoff,
+            chans=chans,
+            num_pool_layers=num_pool_layers,
+            drop_prob=drop_prob,
+            in_shape=in_shape,
+            kernel_shape=kernel_shape,
+        )
+
+    def complex_to_chan_dim(self, x: torch.Tensor) -> torch.Tensor:
+        b, c, h, w, two = x.shape
+        assert two == 2
+        return x.permute(0, 4, 1, 2, 3).reshape(b, 2 * c, h, w)
+
+    def chan_complex_to_last_dim(self, x: torch.Tensor) -> torch.Tensor:
+        b, c2, h, w = x.shape
+        assert c2 % 2 == 0
+        c = c2 // 2
+        return x.view(b, 2, c, h, w).permute(0, 2, 3, 4, 1).contiguous()
+
+    def norm(
+        self, x: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        # group norm
+        b, c, h, w = x.shape
+        x = x.view(b, 2, c // 2 * h * w)
+
+        mean = x.mean(dim=2).view(b, 2, 1, 1)
+        std = x.std(dim=2).view(b, 2, 1, 1)
+
+        x = x.view(b, c, h, w)
+
+        return (x - mean) / std, mean, std
+
+    def norm_new(
+        self, x: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        # group norm
+        b, c, h, w = x.shape
+        num_groups = 2
+        assert (
+            c % num_groups == 0
+        ), f"Number of channels ({c}) must be divisible by number of groups ({num_groups})."
+
+        x = x.view(b, num_groups, c // num_groups * h * w)
+
+        mean = x.mean(dim=2).view(b, num_groups, 1, 1)
+        std = x.std(dim=2).view(b, num_groups, 1, 1)
+        print(x.shape, mean.shape, std.shape)
+
+        x = x.view(b, c, h, w)
+        mean = (
+            mean.view(b, num_groups, 1, 1)
+            .repeat(1, c // num_groups, h, w)
+            .view(b, c, h, w)
+        )
+        std = (
+            std.view(b, num_groups, 1, 1)
+            .repeat(1, c // num_groups, h, w)
+            .view(b, c, h, w)
+        )
+
+        return (x - mean) / std, mean, std
+
+    def unnorm(
+        self, x: torch.Tensor, mean: torch.Tensor, std: torch.Tensor
+    ) -> torch.Tensor:
+        return x * std + mean
+
+    def pad(
+        self, x: torch.Tensor
+    ) -> Tuple[torch.Tensor, Tuple[List[int], List[int], int, int]]:
+        _, _, h, w = x.shape
+        w_mult = ((w - 1) | 15) + 1
+        h_mult = ((h - 1) | 15) + 1
+        w_pad = [math.floor((w_mult - w) / 2), math.ceil((w_mult - w) / 2)]
+        h_pad = [math.floor((h_mult - h) / 2), math.ceil((h_mult - h) / 2)]
+        # TODO: fix this type when PyTorch fixes theirs
+        # the documentation lies - this actually takes a list
+        # https://github.com/pytorch/pytorch/blob/master/torch/nn/functional.py#L3457
+        # https://github.com/pytorch/pytorch/pull/16949
+        x = F.pad(x, w_pad + h_pad)
+
+        return x, (h_pad, w_pad, h_mult, w_mult)
+
+    def unpad(
+        self,
+        x: torch.Tensor,
+        h_pad: List[int],
+        w_pad: List[int],
+        h_mult: int,
+        w_mult: int,
+    ) -> torch.Tensor:
+        return x[
+            ..., h_pad[0] : h_mult - h_pad[1], w_pad[0] : w_mult - w_pad[1]
+        ]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if not x.shape[-1] == 2:
+            raise ValueError("Last dimension must be 2 for complex.")
+
+        chans = x.shape[1]
+        if chans == 2:
+            # FIXME: hard coded skip norm/pad temporarily to avoid group norm bug
+            x = self.complex_to_chan_dim(x)
+            x = self.udno(x)
+            return self.chan_complex_to_last_dim(x)
+
+        # get shapes for unet and normalize
+        x = self.complex_to_chan_dim(x)
+        x, mean, std = self.norm(x)
+        x, pad_sizes = self.pad(x)
+
+        x = self.udno(x)
+
+        # get shapes back and unnormalize
+        x = self.unpad(x, *pad_sizes)
+        x = self.unnorm(x, mean, std)
+        x = self.chan_complex_to_last_dim(x)
+
+        return x
+
+
+class SensitivityModel(nn.Module):
+    """
+    Learn sensitivity maps
+    """
+
+    def __init__(
+        self,
+        chans: int,
+        num_pools: int,
+        radius_cutoff: float,
+        in_shape: Tuple[int, int],
+        kernel_shape: Tuple[int, int],
+        in_chans: int = 2,
+        out_chans: int = 2,
+        drop_prob: float = 0.0,
+        mask_center: bool = True,
+    ):
+        """
+        Parameters
+        ----------
+        chans : int
+            Number of output channels of the first convolution layer.
+        num_pools : int
+            Number of down-sampling and up-sampling layers.
+        in_chans : int, optional
+            Number of channels in the input to the U-Net model. Default is 2.
+        out_chans : int, optional
+            Number of channels in the output to the U-Net model. Default is 2.
+        drop_prob : float, optional
+            Dropout probability. Default is 0.0.
+        mask_center : bool, optional
+            Whether to mask center of k-space for sensitivity map calculation.
+            Default is True.
+        """
+        super().__init__()
+        self.mask_center = mask_center
+        self.norm_udno = NormUDNO(
+            chans,
+            num_pools,
+            radius_cutoff,
+            in_shape,
+            kernel_shape,
+            in_chans=in_chans,
+            out_chans=out_chans,
+            drop_prob=drop_prob,
+        )
+
+    def divide_root_sum_of_squares(self, x: torch.Tensor) -> torch.Tensor:
+        return x / fastmri.rss_complex(x, dim=1).unsqueeze(-1).unsqueeze(1)
+
+    def get_pad_and_num_low_freqs(
+        self, mask: torch.Tensor, num_low_frequencies=None
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if num_low_frequencies is None or any(
+            torch.any(t == 0) for t in num_low_frequencies
+        ):
+            # get low frequency line locations and mask them out
+            squeezed_mask = mask[:, 0, 0, :, 0].to(torch.int8)
+            cent = squeezed_mask.shape[1] // 2
+            # running argmin returns the first non-zero
+            left = torch.argmin(squeezed_mask[:, :cent].flip(1), dim=1)
+            right = torch.argmin(squeezed_mask[:, cent:], dim=1)
+            num_low_frequencies_tensor = torch.max(
+                2 * torch.min(left, right), torch.ones_like(left)
+            )  # force a symmetric center unless 1
+        else:
+            num_low_frequencies_tensor = num_low_frequencies * torch.ones(
+                mask.shape[0], dtype=mask.dtype, device=mask.device
+            )
+
+        pad = (mask.shape[-2] - num_low_frequencies_tensor + 1) // 2
+
+        return pad.type(torch.long), num_low_frequencies_tensor.type(torch.long)
+
+    def forward(
+        self,
+        masked_kspace: torch.Tensor,
+        mask: torch.Tensor,
+        num_low_frequencies: Optional[int] = None,
+    ) -> torch.Tensor:
+        if self.mask_center:
+            pad, num_low_freqs = self.get_pad_and_num_low_freqs(
+                mask, num_low_frequencies
+            )
+            masked_kspace = batched_mask_center(
+                masked_kspace, pad, pad + num_low_freqs
+            )
+
+        # convert to image space
+        images, batches = chans_to_batch_dim(fastmri.ifft2c(masked_kspace))
+
+        # estimate sensitivities
+        return self.divide_root_sum_of_squares(
+            batch_chans_to_chan_dim(self.norm_udno(images), batches)
+        )
+
+
+class VarNetBlock(nn.Module):
+    """
+    Model block for iterative refinement of k-space data.
+
+    This model applies a combination of soft data consistency with the input
+    model as a regularizer. A series of these blocks can be stacked to form
+    the full variational network.
+
+    aka Refinement Module in Fig 1
+    """
+
+    def __init__(self, kno: nn.Module, ino: nn.Module):
+        """
+        Args:
+            model: Module for "regularization" component of variational
+                network.
+        """
+        super().__init__()
+        self.kno = kno
+        self.ino = ino
+        self.dc_weight = nn.Parameter(torch.ones(1))
+
+    def forward(
+        self,
+        current_kspace: torch.Tensor,
+        ref_kspace: torch.Tensor,
+        mask: torch.Tensor,
+        sens_maps: torch.Tensor,
+        use_dc_term: bool = True,
+    ) -> torch.Tensor:
+        """
+        Args:
+            current_kspace: The current k-space data (frequency domain data)
+                            being processed by the network. (torch.Tensor)
+            ref_kspace: Original subsampled k-space data (from which we are
+                reconstrucintg the image (reference k-space). (torch.Tensor)
+            mask: A binary mask indicating the locations in k-space where
+                data consistency should be enforced. (torch.Tensor)
+            sens_maps: Sensitivity maps for the different coils in parallel
+                    imaging. (torch.Tensor)
+        """
+
+        # model-term see orange box of Fig 1 in E2E-VarNet paper!
+        # multi channel k-space -> single channel image-space
+        b, c, h, w, _ = current_kspace.shape
+
+        # ======= kNO in measurement (k) space ========
+        current_kspace, b = chans_to_batch_dim(current_kspace)  # reduce
+        current_kspace = self.kno(current_kspace)  # inpaint
+        current_kspace = batch_chans_to_chan_dim(current_kspace, b)  # expand
+
+        # ======= iNO in image (i) space ========
+        reduced_image = sens_reduce(current_kspace, sens_maps)
+        # single channel image-space
+        refined_image = self.ino(reduced_image)
+        # single channel image-space -> multi channel k-space
+        model_term = sens_expand(refined_image, sens_maps)
+
+        # only use first 15 channels (masked_kspace) in the update
+        # current_kspace = current_kspace[:, :15, :, :, :]
+
+        if not use_dc_term:
+            return current_kspace - model_term
+
+        """
+        Soft data consistency term:
+            - Calculates the difference between current k-space and reference k-space where the mask is true.
+            - Multiplies this difference by the data consistency weight.
+        """
+        # dc_term: see green box of Fig 1 in E2E-VarNet paper!
+        zero = torch.zeros(1, 1, 1, 1, 1).to(current_kspace)
+        soft_dc = (
+            torch.where(mask, current_kspace - ref_kspace, zero)
+            * self.dc_weight
+        )
+        return current_kspace - soft_dc - model_term
+
+
+class NOShared(nn.Module):
+    """
+    Neural Operator model with shared cascade parameters for MRI reconstruction.
+
+    Uses a variational architecture (iterative updates) with a learned sensitivity
+    model. All operations are resolution invariant employing neural operator
+    modules.
+    """
+
+    def __init__(
+        self,
+        num_cascades: int = 12,
+        sens_chans: int = 8,
+        sens_pools: int = 4,
+        chans: int = 18,
+        pools: int = 4,
+        gno_chans: int = 16,
+        gno_pools: int = 4,
+        gno_radius_cutoff: float = 0.02,
+        gno_kernel_shape: Tuple[int, int] = (6, 7),
+        radius_cutoff: float = 0.01,
+        kernel_shape: Tuple[int, int] = (3, 4),
+        in_shape: Tuple[int, int] = (320, 320),
+        mask_center: bool = True,
+        use_dc_term: bool = True,
+    ):
+        """
+        Parameters
+        ----------
+        num_cascades : int
+            Number of cascades (i.e., layers) for variational network.
+        sens_chans : int
+            Number of channels for sensitivity map U-Net.
+        sens_pools : int
+            Number of downsampling and upsampling layers for sensitivity map U-Net.
+        chans : int
+            Number of channels for cascade U-Net.
+        pools : int
+            Number of downsampling and upsampling layers for cascade U-Net.
+        mask_center : bool
+            Whether to mask center of k-space for sensitivity map calculation.
+        use_dc_term : bool
+            Whether to use the data consistency term.
+        """
+
+        super().__init__()
+
+        self.num_cascades = num_cascades
+
+        self.sens_net = SensitivityModel(
+            sens_chans,
+            sens_pools,
+            radius_cutoff,
+            in_shape,
+            kernel_shape,
+            mask_center=False,
+        )
+        self.kno = NormUDNO(
+            gno_chans,
+            gno_pools,
+            in_shape=in_shape,
+            radius_cutoff=gno_radius_cutoff,
+            kernel_shape=gno_kernel_shape,
+            in_chans=2,
+            out_chans=2,
+        )
+        self.ino = NormUDNO(
+            chans,
+            pools,
+            radius_cutoff,
+            in_shape,
+            kernel_shape,
+            in_chans=2,
+            out_chans=2,
+        )
+        self.cascade = VarNetBlock(self.kno, self.ino)
+        self.use_dc_term = use_dc_term
+
+    def forward(
+        self,
+        masked_kspace: torch.Tensor,
+        mask: torch.Tensor,
+        num_low_frequencies: Optional[int] = None,
+    ) -> torch.Tensor:
+
+        # (B, C, X, Y, 2)
+        kspace_pred = masked_kspace
+        # iterative update
+        for _ in range(self.num_cascades):
+            # sens model
+            sens_maps = self.sens_net(kspace_pred, mask, num_low_frequencies)
+
+            # kno + ino (cascade)
+            kspace_pred = self.cascade(
+                kspace_pred, masked_kspace, mask, sens_maps, self.use_dc_term
+            )
+
+        spatial_pred = fastmri.ifft2c(kspace_pred)
+        spatial_pred_abs = fastmri.complex_abs(spatial_pred)
+        combined_spatial = fastmri.rss(spatial_pred_abs, dim=1)
+
+        return combined_spatial
+
+
+if __name__ == "__main__":
+    model = NOShared(
+        num_cascades=4,
+        radius_cutoff=0.02,
+        kernel_shape=(6, 7),
+    )
+
+    x = torch.rand((2, 15, 320, 320, 2))
+    o = model(x, x.bool(), None)

models/no_varnet.py

@@ -0,0 +1,598 @@
+import math
+from typing import List, Literal, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import fastmri
+from fastmri import transforms
+from models.udno import UDNO
+
+
+def sens_expand(x: torch.Tensor, sens_maps: torch.Tensor) -> torch.Tensor:
+    """
+    Calculates F (x sens_maps)
+
+    Parameters
+    ----------
+    x : ndarray
+        Single-channel image of shape (..., H, W, 2)
+    sens_maps : ndarray
+        Sensitivity maps (image space)
+
+    Returns
+    -------
+    ndarray
+        Result of the operation F (x sens_maps)
+    """
+    return fastmri.fft2c(fastmri.complex_mul(x, sens_maps))
+
+
+def sens_reduce(k: torch.Tensor, sens_maps: torch.Tensor) -> torch.Tensor:
+    """
+    Calculates F^{-1}(k) * conj(sens_maps)
+    where conj(sens_maps) is the element-wise applied complex conjugate
+
+    Parameters
+    ----------
+    k : ndarray
+        Multi-channel k-space of shape (B, C, H, W, 2)
+    sens_maps : ndarray
+        Sensitivity maps (image space)
+
+    Returns
+    -------
+    ndarray
+        Result of the operation F^{-1}(k) * conj(sens_maps)
+    """
+    return fastmri.complex_mul(fastmri.ifft2c(k), fastmri.complex_conj(sens_maps)).sum(
+        dim=1, keepdim=True
+    )
+
+
+def chans_to_batch_dim(x: torch.Tensor) -> Tuple[torch.Tensor, int]:
+    """Reshapes batched multi-channel samples into multiple single channel samples.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        x has shape (b, c, h, w, 2)
+
+    Returns
+    -------
+    Tuple[torch.Tensor, int]
+        tensor of shape (b * c, 1, h, w, 2), b
+    """
+    b, c, h, w, comp = x.shape
+    return x.view(b * c, 1, h, w, comp), b
+
+
+def batch_chans_to_chan_dim(x: torch.Tensor, batch_size: int) -> torch.Tensor:
+    """Reshapes batched independent samples into original multi-channel samples.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        tensor of shape (b * c, 1, h, w, 2)
+    batch_size : int
+        batch size
+
+    Returns
+    -------
+    torch.Tensor
+        original multi-channel tensor of shape (b, c, h, w, 2)
+    """
+    bc, _, h, w, comp = x.shape
+    c = bc // batch_size
+    return x.view(batch_size, c, h, w, comp)
+
+
+class NormUDNO(nn.Module):
+    """
+    Normalized UDNO model.
+
+    Inputs are normalized before the UDNO for numerically stable training.
+    """
+
+    def __init__(
+        self,
+        chans: int,
+        num_pool_layers: int,
+        radius_cutoff: float,
+        in_shape: Tuple[int, int],
+        kernel_shape: Tuple[int, int],
+        in_chans: int = 2,
+        out_chans: int = 2,
+        drop_prob: float = 0.0,
+    ):
+        """
+        Initialize the VarNet model.
+
+        Parameters
+        ----------
+        chans : int
+            Number of output channels of the first convolution layer.
+        num_pools : int
+            Number of down-sampling and up-sampling layers.
+        in_chans : int, optional
+            Number of channels in the input to the U-Net model. Default is 2.
+        out_chans : int, optional
+            Number of channels in the output to the U-Net model. Default is 2.
+        drop_prob : float, optional
+            Dropout probability. Default is 0.0.
+        """
+        super().__init__()
+
+        self.udno = UDNO(
+            in_chans=in_chans,
+            out_chans=out_chans,
+            radius_cutoff=radius_cutoff,
+            chans=chans,
+            num_pool_layers=num_pool_layers,
+            drop_prob=drop_prob,
+            in_shape=in_shape,
+            kernel_shape=kernel_shape,
+        )
+
+    def complex_to_chan_dim(self, x: torch.Tensor) -> torch.Tensor:
+        b, c, h, w, two = x.shape
+        assert two == 2
+        return x.permute(0, 4, 1, 2, 3).reshape(b, 2 * c, h, w)
+
+    def chan_complex_to_last_dim(self, x: torch.Tensor) -> torch.Tensor:
+        b, c2, h, w = x.shape
+        assert c2 % 2 == 0
+        c = c2 // 2
+        return x.view(b, 2, c, h, w).permute(0, 2, 3, 4, 1).contiguous()
+
+    def norm(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        # group norm
+        b, c, h, w = x.shape
+        x = x.view(b, 2, c // 2 * h * w)
+
+        mean = x.mean(dim=2).view(b, 2, 1, 1)
+        std = x.std(dim=2).view(b, 2, 1, 1)
+
+        x = x.view(b, c, h, w)
+
+        return (x - mean) / std, mean, std
+
+    def norm_new(
+        self, x: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        # FIXME: not working, wip
+        # group norm
+        b, c, h, w = x.shape
+        num_groups = 2
+        assert (
+            c % num_groups == 0
+        ), f"Number of channels ({c}) must be divisible by number of groups ({num_groups})."
+
+        x = x.view(b, num_groups, c // num_groups * h * w)
+
+        mean = x.mean(dim=2).view(b, num_groups, 1, 1)
+        std = x.std(dim=2).view(b, num_groups, 1, 1)
+        print(x.shape, mean.shape, std.shape)
+
+        x = x.view(b, c, h, w)
+        mean = (
+            mean.view(b, num_groups, 1, 1)
+            .repeat(1, c // num_groups, h, w)
+            .view(b, c, h, w)
+        )
+        std = (
+            std.view(b, num_groups, 1, 1)
+            .repeat(1, c // num_groups, h, w)
+            .view(b, c, h, w)
+        )
+
+        return (x - mean) / std, mean, std
+
+    def unnorm(
+        self, x: torch.Tensor, mean: torch.Tensor, std: torch.Tensor
+    ) -> torch.Tensor:
+        return x * std + mean
+
+    def pad(
+        self, x: torch.Tensor
+    ) -> Tuple[torch.Tensor, Tuple[List[int], List[int], int, int]]:
+        _, _, h, w = x.shape
+        w_mult = ((w - 1) | 15) + 1
+        h_mult = ((h - 1) | 15) + 1
+        w_pad = [math.floor((w_mult - w) / 2), math.ceil((w_mult - w) / 2)]
+        h_pad = [math.floor((h_mult - h) / 2), math.ceil((h_mult - h) / 2)]
+        # TODO: fix this type when PyTorch fixes theirs
+        # the documentation lies - this actually takes a list
+        # https://github.com/pytorch/pytorch/blob/master/torch/nn/functional.py#L3457
+        # https://github.com/pytorch/pytorch/pull/16949
+        x = F.pad(x, w_pad + h_pad)
+
+        return x, (h_pad, w_pad, h_mult, w_mult)
+
+    def unpad(
+        self,
+        x: torch.Tensor,
+        h_pad: List[int],
+        w_pad: List[int],
+        h_mult: int,
+        w_mult: int,
+    ) -> torch.Tensor:
+        return x[..., h_pad[0] : h_mult - h_pad[1], w_pad[0] : w_mult - w_pad[1]]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if not x.shape[-1] == 2:
+            raise ValueError("Last dimension must be 2 for complex.")
+
+        chans = x.shape[1]
+        if chans == 2:
+            # FIXME: hard coded skip norm/pad temporarily to avoid group norm bug
+            x = self.complex_to_chan_dim(x)
+            x = self.udno(x)
+            return self.chan_complex_to_last_dim(x)
+
+        # get shapes for unet and normalize
+        x = self.complex_to_chan_dim(x)
+        x, mean, std = self.norm(x)
+        x, pad_sizes = self.pad(x)
+
+        x = self.udno(x)
+
+        # get shapes back and unnormalize
+        x = self.unpad(x, *pad_sizes)
+        x = self.unnorm(x, mean, std)
+        x = self.chan_complex_to_last_dim(x)
+
+        return x
+
+
+class SensitivityModel(nn.Module):
+    """
+    Learn sensitivity maps
+    """
+
+    def __init__(
+        self,
+        chans: int,
+        num_pools: int,
+        radius_cutoff: float,
+        in_shape: Tuple[int, int],
+        kernel_shape: Tuple[int, int],
+        in_chans: int = 2,
+        out_chans: int = 2,
+        drop_prob: float = 0.0,
+        mask_center: bool = True,
+    ):
+        """
+        Parameters
+        ----------
+        chans : int
+            Number of output channels of the first convolution layer.
+        num_pools : int
+            Number of down-sampling and up-sampling layers.
+        in_chans : int, optional
+            Number of channels in the input to the U-Net model. Default is 2.
+        out_chans : int, optional
+            Number of channels in the output to the U-Net model. Default is 2.
+        drop_prob : float, optional
+            Dropout probability. Default is 0.0.
+        mask_center : bool, optional
+            Whether to mask center of k-space for sensitivity map calculation.
+            Default is True.
+        """
+        super().__init__()
+        self.mask_center = mask_center
+        self.norm_udno = NormUDNO(
+            chans,
+            num_pools,
+            radius_cutoff,
+            in_shape,
+            kernel_shape,
+            in_chans=in_chans,
+            out_chans=out_chans,
+            drop_prob=drop_prob,
+        )
+
+    def divide_root_sum_of_squares(self, x: torch.Tensor) -> torch.Tensor:
+        return x / fastmri.rss_complex(x, dim=1).unsqueeze(-1).unsqueeze(1)
+
+    def get_pad_and_num_low_freqs(
+        self, mask: torch.Tensor, num_low_frequencies: Optional[int] = None
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if num_low_frequencies is None or any(
+            torch.any(t == 0) for t in num_low_frequencies
+        ):
+            # get low frequency line locations and mask them out
+            squeezed_mask = mask[:, 0, 0, :, 0].to(torch.int8)
+            cent = squeezed_mask.shape[1] // 2
+            # running argmin returns the first non-zero
+            left = torch.argmin(squeezed_mask[:, :cent].flip(1), dim=1)
+            right = torch.argmin(squeezed_mask[:, cent:], dim=1)
+            num_low_frequencies_tensor = torch.max(
+                2 * torch.min(left, right), torch.ones_like(left)
+            )  # force a symmetric center unless 1
+        else:
+            num_low_frequencies_tensor = num_low_frequencies * torch.ones(
+                mask.shape[0], dtype=mask.dtype, device=mask.device
+            )
+
+        pad = (mask.shape[-2] - num_low_frequencies_tensor + 1) // 2
+
+        return pad.type(torch.long), num_low_frequencies_tensor.type(torch.long)
+
+    def forward(
+        self,
+        masked_kspace: torch.Tensor,
+        mask: torch.Tensor,
+        num_low_frequencies: Optional[int] = None,
+    ) -> torch.Tensor:
+        if self.mask_center:
+            pad, num_low_freqs = self.get_pad_and_num_low_freqs(
+                mask, num_low_frequencies
+            )
+            masked_kspace = transforms.batched_mask_center(
+                masked_kspace, pad, pad + num_low_freqs
+            )
+
+        # convert to image space
+        images, batches = chans_to_batch_dim(fastmri.ifft2c(masked_kspace))
+
+        # estimate sensitivities
+        return self.divide_root_sum_of_squares(
+            batch_chans_to_chan_dim(self.norm_udno(images), batches)
+        )
+
+
+class VarNetBlock(nn.Module):
+    """
+    Model block for iterative refinement of k-space data.
+
+    This model applies a combination of soft data consistency with the input
+    model as a regularizer. A series of these blocks can be stacked to form
+    the full variational network.
+
+    aka Refinement Module in Fig 1
+    """
+
+    def __init__(self, model: nn.Module):
+        """
+        Args:
+            model: Module for "regularization" component of variational
+                network.
+        """
+        super().__init__()
+
+        self.model = model
+        self.dc_weight = nn.Parameter(torch.ones(1))
+
+    def forward(
+        self,
+        current_kspace: torch.Tensor,
+        ref_kspace: torch.Tensor,
+        mask: torch.Tensor,
+        sens_maps: torch.Tensor,
+        use_dc_term: bool = True,
+    ) -> torch.Tensor:
+        """
+        Args:
+            current_kspace: The current k-space data (frequency domain data)
+                            being processed by the network. (torch.Tensor)
+            ref_kspace: Original subsampled k-space data (from which we are
+                reconstrucintg the image (reference k-space). (torch.Tensor)
+            mask: A binary mask indicating the locations in k-space where
+                data consistency should be enforced. (torch.Tensor)
+            sens_maps: Sensitivity maps for the different coils in parallel
+                    imaging. (torch.Tensor)
+        """
+
+        # model-term see orange box of Fig 1 in E2E-VarNet paper!
+        # multi channel k-space -> single channel image-space
+        b, c, h, w, _ = current_kspace.shape
+
+        if c == 30:
+            # get kspace and inpainted kspace
+            kspace = current_kspace[:, :15, :, :, :]
+            in_kspace = current_kspace[:, 15:, :, :, :]
+            # convert to image space
+            image = sens_reduce(kspace, sens_maps)
+            in_image = sens_reduce(in_kspace, sens_maps)
+            # concatenate both onto each other
+            reduced_image = torch.cat([image, in_image], dim=1)
+        else:
+            reduced_image = sens_reduce(current_kspace, sens_maps)
+
+        # single channel image-space
+        refined_image = self.model(reduced_image)
+
+        # single channel image-space -> multi channel k-space
+        model_term = sens_expand(refined_image, sens_maps)
+
+        # only use first 15 channels (masked_kspace) in the update
+        # current_kspace = current_kspace[:, :15, :, :, :]
+
+        if not use_dc_term:
+            return current_kspace - model_term
+
+        """
+        Soft data consistency term:
+            - Calculates the difference between current k-space and reference k-space where the mask is true.
+            - Multiplies this difference by the data consistency weight.
+        """
+        # dc_term: see green box of Fig 1 in E2E-VarNet paper!
+        zero = torch.zeros(1, 1, 1, 1, 1).to(current_kspace)
+        soft_dc = torch.where(mask, current_kspace - ref_kspace, zero) * self.dc_weight
+        return current_kspace - soft_dc - model_term
+
+
+class NOVarnet(nn.Module):
+    """
+    Neural Operator model for MRI reconstruction.
+
+    Uses a variational architecture (iterative updates) with a learned sensitivity
+    model. All operations are resolution invariant employing neural operator
+    modules (GNO, UDNO).
+    """
+
+    def __init__(
+        self,
+        num_cascades: int = 12,
+        sens_chans: int = 8,
+        sens_pools: int = 4,
+        chans: int = 18,
+        pools: int = 4,
+        gno_chans: int = 16,
+        gno_pools: int = 4,
+        gno_radius_cutoff: float = 0.02,
+        gno_kernel_shape: Tuple[int, int] = (6, 7),
+        radius_cutoff: float = 0.01,
+        kernel_shape: Tuple[int, int] = (3, 4),
+        in_shape: Tuple[int, int] = (640, 320),
+        mask_center: bool = True,
+        use_dc_term: bool = True,
+        reduction_method: Literal["batch", "rss"] = "rss",
+        skip_method: Literal["replace", "add", "add_inv", "concat"] = "add",
+    ):
+        """
+        Parameters
+        ----------
+        num_cascades : int
+            Number of cascades (i.e., layers) for variational network.
+        sens_chans : int
+            Number of channels for sensitivity map U-Net.
+        sens_pools : int
+            Number of downsampling and upsampling layers for sensitivity map U-Net.
+        chans : int
+            Number of channels for cascade U-Net.
+        pools : int
+            Number of downsampling and upsampling layers for cascade U-Net.
+        mask_center : bool
+            Whether to mask center of k-space for sensitivity map calculation.
+        use_dc_term : bool
+            Whether to use the data consistency term.
+        reduction_method : "batch" or "rss"
+            Method for reducing sensitivity maps to single channel.
+            "batch" reduces to single channel by stacking channels.
+            "rss" reduces to single channel by root sum of squares.
+        skip_method : "replace" or "add" or "add_inv" or "concat"
+            "replace" replaces the input with the output of the GNO
+            "add" adds the output of the GNO to the input
+            "add_inv" adds the output of the GNO to the input (only where samples are missing)
+            "concat" concatenates the output of the GNO to the input
+        """
+
+        super().__init__()
+
+        self.sens_net = SensitivityModel(
+            sens_chans,
+            sens_pools,
+            radius_cutoff,
+            in_shape,
+            kernel_shape,
+            mask_center=mask_center,
+        )
+        self.gno = NormUDNO(
+            gno_chans,
+            gno_pools,
+            in_shape=in_shape,
+            radius_cutoff=radius_cutoff,
+            kernel_shape=kernel_shape,
+            # radius_cutoff=gno_radius_cutoff,
+            # kernel_shape=gno_kernel_shape,
+            in_chans=2,
+            out_chans=2,
+        )
+        self.cascades = nn.ModuleList(
+            [
+                VarNetBlock(
+                    NormUDNO(
+                        chans,
+                        pools,
+                        radius_cutoff,
+                        in_shape,
+                        kernel_shape,
+                        in_chans=(
+                            4 if skip_method == "concat" and cascade_idx == 0 else 2
+                        ),
+                        out_chans=2,
+                    )
+                )
+                for cascade_idx in range(num_cascades)
+            ]
+        )
+        self.use_dc_term = use_dc_term
+        self.reduction_method = reduction_method
+        self.skip_method = skip_method
+
+    def forward(
+        self,
+        masked_kspace: torch.Tensor,
+        mask: torch.Tensor,
+        num_low_frequencies: Optional[int] = None,
+    ) -> torch.Tensor:
+
+        # (B, C, X, Y, 2)
+        sens_maps = self.sens_net(masked_kspace, mask, num_low_frequencies)
+
+        # reduce before inpainting
+        if self.reduction_method == "rss":
+            # (B, 1, H, W, 2) single channel image space
+            x_reduced = sens_reduce(masked_kspace, sens_maps)
+            # (B, 1, H, W, 2)
+            k_reduced = fastmri.fft2c(x_reduced)
+        elif self.reduction_method == "batch":
+            k_reduced, b = chans_to_batch_dim(masked_kspace)
+
+        # inpainting
+        if self.skip_method == "replace":
+            kspace_pred = self.gno(k_reduced)
+        elif self.skip_method == "add_inv":
+            # FIXME: this is not correct (mask has shape B, 1, H, W, 2 and self.gno(k_reduced) has shape B*C, 1, H, W, 2)
+            kspace_pred = k_reduced.clone() + (~mask * self.gno(k_reduced))
+        elif self.skip_method == "add":
+            kspace_pred = k_reduced.clone() + self.gno(k_reduced)
+        elif self.skip_method == "concat":
+            kspace_pred = torch.cat([k_reduced.clone(), self.gno(k_reduced)], dim=1)
+        else:
+            raise NotImplementedError("skip_method not implemented")
+
+        # expand after inpainting
+        if self.reduction_method == "rss":
+            if self.skip_method == "concat":
+                # kspace_pred is (B, 2, H, W, 2)
+                kspace = kspace_pred[:, :1, :, :, :]
+                in_kspace = kspace_pred[:, 1:, :, :, :]
+                # B, 2C, H, W, 2
+                kspace_pred = torch.cat(
+                    [sens_expand(kspace, sens_maps), sens_expand(in_kspace, sens_maps)],
+                    dim=1,
+                )
+            else:
+                # (B, 1, H, W, 2) -> (B, C, H, W, 2) multi-channel k space
+                kspace_pred = sens_expand(kspace_pred, sens_maps)
+        elif self.reduction_method == "batch":
+            # (B, C, H, W, 2) multi-channel k space
+            if self.skip_method == "concat":
+                kspace = kspace_pred[:, :1, :, :, :]
+                in_kspace = kspace_pred[:, 1:, :, :, :]
+                # B, 2C, H, W, 2
+                kspace_pred = torch.cat(
+                    [
+                        batch_chans_to_chan_dim(kspace, b),
+                        batch_chans_to_chan_dim(in_kspace, b),
+                    ],
+                    dim=1,
+                )
+            else:
+                kspace_pred = batch_chans_to_chan_dim(kspace_pred, b)
+
+        # iterative update
+        for cascade in self.cascades:
+            kspace_pred = cascade(
+                kspace_pred, masked_kspace, mask, sens_maps, self.use_dc_term
+            )
+
+        spatial_pred = fastmri.ifft2c(kspace_pred)
+        spatial_pred_abs = fastmri.complex_abs(spatial_pred)
+        combined_spatial = fastmri.rss(spatial_pred_abs, dim=1)
+
+        return combined_spatial

models/no_varnet_nokno.py

@@ -0,0 +1,581 @@
+"""
+NO Varnet WITHOUT KNO for ablation
+"""
+
+import math
+from typing import Iterable, List, Literal, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import fastmri
+from fastmri import transforms
+from fastmri.datasets import SliceDatasetLMDB, SliceSample
+from models.udno import UDNO
+
+
+def sens_expand(x: torch.Tensor, sens_maps: torch.Tensor) -> torch.Tensor:
+    """
+    Calculates F (x sens_maps)
+
+    Parameters
+    ----------
+    x : ndarray
+        Single-channel image of shape (..., H, W, 2)
+    sens_maps : ndarray
+        Sensitivity maps (image space)
+
+    Returns
+    -------
+    ndarray
+        Result of the operation F (x sens_maps)
+    """
+    return fastmri.fft2c(fastmri.complex_mul(x, sens_maps))
+
+
+def sens_reduce(k: torch.Tensor, sens_maps: torch.Tensor) -> torch.Tensor:
+    """
+    Calculates F^{-1}(k) * conj(sens_maps)
+    where conj(sens_maps) is the element-wise applied complex conjugate
+
+    Parameters
+    ----------
+    k : ndarray
+        Multi-channel k-space of shape (B, C, H, W, 2)
+    sens_maps : ndarray
+        Sensitivity maps (image space)
+
+    Returns
+    -------
+    ndarray
+        Result of the operation F^{-1}(k) * conj(sens_maps)
+    """
+    return fastmri.complex_mul(
+        fastmri.ifft2c(k), fastmri.complex_conj(sens_maps)
+    ).sum(dim=1, keepdim=True)
+
+
+def chans_to_batch_dim(x: torch.Tensor) -> Tuple[torch.Tensor, int]:
+    """Reshapes batched multi-channel samples into multiple single channel samples.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        x has shape (b, c, h, w, 2)
+
+    Returns
+    -------
+    Tuple[torch.Tensor, int]
+        tensor of shape (b * c, 1, h, w, 2), b
+    """
+    b, c, h, w, comp = x.shape
+    return x.view(b * c, 1, h, w, comp), b
+
+
+def batch_chans_to_chan_dim(x: torch.Tensor, batch_size: int) -> torch.Tensor:
+    """Reshapes batched independent samples into original multi-channel samples.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        tensor of shape (b * c, 1, h, w, 2)
+    batch_size : int
+        batch size
+
+    Returns
+    -------
+    torch.Tensor
+        original multi-channel tensor of shape (b, c, h, w, 2)
+    """
+    bc, _, h, w, comp = x.shape
+    c = bc // batch_size
+    return x.view(batch_size, c, h, w, comp)
+
+
+class NormUDNO(nn.Module):
+    """
+    Normalized UDNO model.
+
+    Inputs are normalized before the UDNO for numerically stable training.
+    """
+
+    def __init__(
+        self,
+        chans: int,
+        num_pool_layers: int,
+        radius_cutoff: float,
+        in_shape: Tuple[int, int],
+        kernel_shape: Tuple[int, int],
+        in_chans: int = 2,
+        out_chans: int = 2,
+        drop_prob: float = 0.0,
+    ):
+        """
+        Initialize the VarNet model.
+
+        Parameters
+        ----------
+        chans : int
+            Number of output channels of the first convolution layer.
+        num_pools : int
+            Number of down-sampling and up-sampling layers.
+        in_chans : int, optional
+            Number of channels in the input to the U-Net model. Default is 2.
+        out_chans : int, optional
+            Number of channels in the output to the U-Net model. Default is 2.
+        drop_prob : float, optional
+            Dropout probability. Default is 0.0.
+        """
+        super().__init__()
+
+        self.udno = UDNO(
+            in_chans=in_chans,
+            out_chans=out_chans,
+            radius_cutoff=radius_cutoff,
+            chans=chans,
+            num_pool_layers=num_pool_layers,
+            drop_prob=drop_prob,
+            in_shape=in_shape,
+            kernel_shape=kernel_shape,
+        )
+
+    def complex_to_chan_dim(self, x: torch.Tensor) -> torch.Tensor:
+        b, c, h, w, two = x.shape
+        assert two == 2
+        return x.permute(0, 4, 1, 2, 3).reshape(b, 2 * c, h, w)
+
+    def chan_complex_to_last_dim(self, x: torch.Tensor) -> torch.Tensor:
+        b, c2, h, w = x.shape
+        assert c2 % 2 == 0
+        c = c2 // 2
+        return x.view(b, 2, c, h, w).permute(0, 2, 3, 4, 1).contiguous()
+
+    def norm(
+        self, x: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        # group norm
+        b, c, h, w = x.shape
+        x = x.view(b, 2, c // 2 * h * w)
+
+        mean = x.mean(dim=2).view(b, 2, 1, 1)
+        std = x.std(dim=2).view(b, 2, 1, 1)
+
+        x = x.view(b, c, h, w)
+
+        return (x - mean) / std, mean, std
+
+    def norm_new(
+        self, x: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        # FIXME: not working, wip
+        # group norm
+        b, c, h, w = x.shape
+        num_groups = 2
+        assert (
+            c % num_groups == 0
+        ), f"Number of channels ({c}) must be divisible by number of groups ({num_groups})."
+
+        x = x.view(b, num_groups, c // num_groups * h * w)
+
+        mean = x.mean(dim=2).view(b, num_groups, 1, 1)
+        std = x.std(dim=2).view(b, num_groups, 1, 1)
+        print(x.shape, mean.shape, std.shape)
+
+        x = x.view(b, c, h, w)
+        mean = (
+            mean.view(b, num_groups, 1, 1)
+            .repeat(1, c // num_groups, h, w)
+            .view(b, c, h, w)
+        )
+        std = (
+            std.view(b, num_groups, 1, 1)
+            .repeat(1, c // num_groups, h, w)
+            .view(b, c, h, w)
+        )
+
+        return (x - mean) / std, mean, std
+
+    def unnorm(
+        self, x: torch.Tensor, mean: torch.Tensor, std: torch.Tensor
+    ) -> torch.Tensor:
+        return x * std + mean
+
+    def pad(
+        self, x: torch.Tensor
+    ) -> Tuple[torch.Tensor, Tuple[List[int], List[int], int, int]]:
+        _, _, h, w = x.shape
+        w_mult = ((w - 1) | 15) + 1
+        h_mult = ((h - 1) | 15) + 1
+        w_pad = [math.floor((w_mult - w) / 2), math.ceil((w_mult - w) / 2)]
+        h_pad = [math.floor((h_mult - h) / 2), math.ceil((h_mult - h) / 2)]
+        # TODO: fix this type when PyTorch fixes theirs
+        # the documentation lies - this actually takes a list
+        # https://github.com/pytorch/pytorch/blob/master/torch/nn/functional.py#L3457
+        # https://github.com/pytorch/pytorch/pull/16949
+        x = F.pad(x, w_pad + h_pad)
+
+        return x, (h_pad, w_pad, h_mult, w_mult)
+
+    def unpad(
+        self,
+        x: torch.Tensor,
+        h_pad: List[int],
+        w_pad: List[int],
+        h_mult: int,
+        w_mult: int,
+    ) -> torch.Tensor:
+        return x[
+            ..., h_pad[0] : h_mult - h_pad[1], w_pad[0] : w_mult - w_pad[1]
+        ]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if not x.shape[-1] == 2:
+            raise ValueError("Last dimension must be 2 for complex.")
+
+        chans = x.shape[1]
+        if chans == 2:
+            # FIXME: hard coded skip norm/pad temporarily to avoid group norm bug
+            x = self.complex_to_chan_dim(x)
+            x = self.udno(x)
+            return self.chan_complex_to_last_dim(x)
+
+        # get shapes for unet and normalize
+        x = self.complex_to_chan_dim(x)
+        x, mean, std = self.norm(x)
+        x, pad_sizes = self.pad(x)
+
+        x = self.udno(x)
+
+        # get shapes back and unnormalize
+        x = self.unpad(x, *pad_sizes)
+        x = self.unnorm(x, mean, std)
+        x = self.chan_complex_to_last_dim(x)
+
+        return x
+
+
+class SensitivityModel(nn.Module):
+    """
+    Learn sensitivity maps
+    """
+
+    def __init__(
+        self,
+        chans: int,
+        num_pools: int,
+        radius_cutoff: float,
+        in_shape: Tuple[int, int],
+        kernel_shape: Tuple[int, int],
+        in_chans: int = 2,
+        out_chans: int = 2,
+        drop_prob: float = 0.0,
+        mask_center: bool = True,
+    ):
+        """
+        Parameters
+        ----------
+        chans : int
+            Number of output channels of the first convolution layer.
+        num_pools : int
+            Number of down-sampling and up-sampling layers.
+        in_chans : int, optional
+            Number of channels in the input to the U-Net model. Default is 2.
+        out_chans : int, optional
+            Number of channels in the output to the U-Net model. Default is 2.
+        drop_prob : float, optional
+            Dropout probability. Default is 0.0.
+        mask_center : bool, optional
+            Whether to mask center of k-space for sensitivity map calculation.
+            Default is True.
+        """
+        super().__init__()
+        self.mask_center = mask_center
+        self.norm_udno = NormUDNO(
+            chans,
+            num_pools,
+            radius_cutoff,
+            in_shape,
+            kernel_shape,
+            in_chans=in_chans,
+            out_chans=out_chans,
+            drop_prob=drop_prob,
+        )
+
+    def divide_root_sum_of_squares(self, x: torch.Tensor) -> torch.Tensor:
+        return x / fastmri.rss_complex(x, dim=1).unsqueeze(-1).unsqueeze(1)
+
+    def get_pad_and_num_low_freqs(
+        self, mask: torch.Tensor, num_low_frequencies: Optional[int] = None
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if num_low_frequencies is None or (isinstance(num_low_frequencies, Iterable) and any(
+            torch.any(t == 0) for t in num_low_frequencies
+        )):
+            # get low frequency line locations and mask them out
+            squeezed_mask = mask[:, 0, 0, :, 0].to(torch.int8)
+            cent = squeezed_mask.shape[1] // 2
+            # running argmin returns the first non-zero
+            left = torch.argmin(squeezed_mask[:, :cent].flip(1), dim=1)
+            right = torch.argmin(squeezed_mask[:, cent:], dim=1)
+            num_low_frequencies_tensor = torch.max(
+                2 * torch.min(left, right), torch.ones_like(left)
+            )  # force a symmetric center unless 1
+        else:
+            num_low_frequencies_tensor = num_low_frequencies * torch.ones(
+                mask.shape[0], dtype=mask.dtype, device=mask.device
+            )
+
+        pad = (mask.shape[-2] - num_low_frequencies_tensor + 1) // 2
+
+        return pad.type(torch.long), num_low_frequencies_tensor.type(torch.long)
+
+    def forward(
+        self,
+        masked_kspace: torch.Tensor,
+        mask: torch.Tensor,
+        num_low_frequencies: Optional[int] = None,
+    ) -> torch.Tensor:
+        if self.mask_center:
+            pad, num_low_freqs = self.get_pad_and_num_low_freqs(
+                mask, num_low_frequencies
+            )
+            masked_kspace = transforms.batched_mask_center(
+                masked_kspace, pad, pad + num_low_freqs
+            )
+
+        # convert to image space
+        images, batches = chans_to_batch_dim(fastmri.ifft2c(masked_kspace))
+
+        # estimate sensitivities
+        return self.divide_root_sum_of_squares(
+            batch_chans_to_chan_dim(self.norm_udno(images), batches)
+        )
+
+
+class VarNetBlock(nn.Module):
+    """
+    Model block for iterative refinement of k-space data.
+
+    This model applies a combination of soft data consistency with the input
+    model as a regularizer. A series of these blocks can be stacked to form
+    the full variational network.
+
+    aka Refinement Module in Fig 1
+    """
+
+    def __init__(self, model: nn.Module):
+        """
+        Args:
+            model: Module for "regularization" component of variational
+                network.
+        """
+        super().__init__()
+
+        self.model = model
+        self.dc_weight = nn.Parameter(torch.ones(1))
+
+    def forward(
+        self,
+        current_kspace: torch.Tensor,
+        ref_kspace: torch.Tensor,
+        mask: torch.Tensor,
+        sens_maps: torch.Tensor,
+        use_dc_term: bool = True,
+    ) -> torch.Tensor:
+        """
+        Args:
+            current_kspace: The current k-space data (frequency domain data)
+                            being processed by the network. (torch.Tensor)
+            ref_kspace: Original subsampled k-space data (from which we are
+                reconstrucintg the image (reference k-space). (torch.Tensor)
+            mask: A binary mask indicating the locations in k-space where
+                data consistency should be enforced. (torch.Tensor)
+            sens_maps: Sensitivity maps for the different coils in parallel
+                    imaging. (torch.Tensor)
+        """
+
+        # model-term see orange box of Fig 1 in E2E-VarNet paper!
+        # multi channel k-space -> single channel image-space
+        b, c, h, w, _ = current_kspace.shape
+
+        if c == 30:
+            # get kspace and inpainted kspace
+            kspace = current_kspace[:, :15, :, :, :]
+            in_kspace = current_kspace[:, 15:, :, :, :]
+            # convert to image space
+            image = sens_reduce(kspace, sens_maps)
+            in_image = sens_reduce(in_kspace, sens_maps)
+            # concatenate both onto each other
+            reduced_image = torch.cat([image, in_image], dim=1)
+        else:
+            reduced_image = sens_reduce(current_kspace, sens_maps)
+
+        # single channel image-space
+        refined_image = self.model(reduced_image)
+
+        # single channel image-space -> multi channel k-space
+        model_term = sens_expand(refined_image, sens_maps)
+
+        # only use first 15 channels (masked_kspace) in the update
+        # current_kspace = current_kspace[:, :15, :, :, :]
+
+        if not use_dc_term:
+            return current_kspace - model_term
+
+        """
+        Soft data consistency term:
+            - Calculates the difference between current k-space and reference k-space where the mask is true.
+            - Multiplies this difference by the data consistency weight.
+        """
+        # dc_term: see green box of Fig 1 in E2E-VarNet paper!
+        zero = torch.zeros(1, 1, 1, 1, 1).to(current_kspace)
+        soft_dc = (
+            torch.where(mask, current_kspace - ref_kspace, zero)
+            * self.dc_weight
+        )
+        return current_kspace - soft_dc - model_term
+
+
+class NOVarnet_no_KNO(nn.Module):
+    """
+    Neural Operator model for MRI reconstruction.
+
+    Uses a variational architecture (iterative updates) with a learned sensitivity
+    model. All operations are resolution invariant employing neural operator
+    modules (GNO, UDNO).
+    """
+
+    def __init__(
+        self,
+        num_cascades: int = 12,
+        sens_chans: int = 8,
+        sens_pools: int = 4,
+        chans: int = 18,
+        pools: int = 4,
+        gno_chans: int = 16,
+        gno_pools: int = 4,
+        gno_radius_cutoff: float = 0.02,
+        gno_kernel_shape: Tuple[int, int] = (6, 7),
+        radius_cutoff: float = 0.01,
+        kernel_shape: Tuple[int, int] = (3, 4),
+        in_shape: Tuple[int, int] = (640, 320),
+        mask_center: bool = True,
+        use_dc_term: bool = True,
+        reduction_method: Literal["batch", "rss"] = "rss",
+        skip_method: Literal["replace", "add", "add_inv", "concat"] = "add",
+    ):
+        """
+        Parameters
+        ----------
+        num_cascades : int
+            Number of cascades (i.e., layers) for variational network.
+        sens_chans : int
+            Number of channels for sensitivity map U-Net.
+        sens_pools : int
+            Number of downsampling and upsampling layers for sensitivity map U-Net.
+        chans : int
+            Number of channels for cascade U-Net.
+        pools : int
+            Number of downsampling and upsampling layers for cascade U-Net.
+        mask_center : bool
+            Whether to mask center of k-space for sensitivity map calculation.
+        use_dc_term : bool
+            Whether to use the data consistency term.
+        reduction_method : "batch" or "rss"
+            Method for reducing sensitivity maps to single channel.
+            "batch" reduces to single channel by stacking channels.
+            "rss" reduces to single channel by root sum of squares.
+        skip_method : "replace" or "add" or "add_inv" or "concat"
+            "replace" replaces the input with the output of the GNO
+            "add" adds the output of the GNO to the input
+            "add_inv" adds the output of the GNO to the input (only where samples are missing)
+            "concat" concatenates the output of the GNO to the input
+        """
+
+        super().__init__()
+
+        self.sens_net = SensitivityModel(
+            sens_chans,
+            sens_pools,
+            radius_cutoff,
+            in_shape,
+            kernel_shape,
+            mask_center=mask_center,
+        )
+        # self.gno = NormUDNO(
+        #     gno_chans,
+        #     gno_pools,
+        #     in_shape=in_shape,
+        #     radius_cutoff=radius_cutoff,
+        #     kernel_shape=kernel_shape,
+        #     # radius_cutoff=gno_radius_cutoff,
+        #     # kernel_shape=gno_kernel_shape,
+        #     in_chans=2,
+        #     out_chans=2,
+        # )
+        self.cascades = nn.ModuleList(
+            [
+                VarNetBlock(
+                    NormUDNO(
+                        chans,
+                        pools,
+                        radius_cutoff,
+                        in_shape,
+                        kernel_shape,
+                        in_chans=(
+                            4
+                            if skip_method == "concat" and cascade_idx == 0
+                            else 2
+                        ),
+                        out_chans=2,
+                    )
+                )
+                for cascade_idx in range(num_cascades)
+            ]
+        )
+        self.use_dc_term = use_dc_term
+        self.reduction_method = reduction_method # not used anywhere anymore
+        self.skip_method = skip_method # not used anywhere anymore
+
+    def forward(
+        self,
+        masked_kspace: torch.Tensor,
+        mask: torch.Tensor,
+        num_low_frequencies: Optional[int] = None,
+    ) -> torch.Tensor:
+
+        # (B, C, X, Y, 2)
+        sens_maps = self.sens_net(masked_kspace, mask, num_low_frequencies)
+
+        kspace_pred = masked_kspace.clone()
+
+        # iterative update
+        for cascade in self.cascades:
+            kspace_pred = cascade(
+                kspace_pred, masked_kspace, mask, sens_maps, self.use_dc_term
+            )
+
+        spatial_pred = fastmri.ifft2c(kspace_pred)
+        spatial_pred_abs = fastmri.complex_abs(spatial_pred)
+        combined_spatial = fastmri.rss(spatial_pred_abs, dim=1)
+
+        return combined_spatial
+
+
+if __name__ == "__main__":
+    ds = SliceDatasetLMDB(
+        "knee",
+        partition="train",
+        mask_fns=None,  # type: ignore
+        complex=False,
+        sample_rate=0.5,
+        crop_shape=(320, 320),
+        coils=15,
+    )
+
+    sample: SliceSample = ds[0]
+    kspace = sample.masked_kspace
+    target = sample.target
+
+    model = NOVarnet_no_KNO(1)
+    res = model.forward(sample.masked_kspace.unsqueeze(0), sample.mask.unsqueeze(0), torch.tensor(sample.num_low_frequencies).unsqueeze(0))

models/temp/no_repeatk.py

@@ -0,0 +1,562 @@
+import math
+import pdb
+from typing import List, Optional, Tuple, Literal
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import fastmri
+from fastmri import transforms
+
+from models.udno import UDNO
+
+
+def sens_expand(x: torch.Tensor, sens_maps: torch.Tensor) -> torch.Tensor:
+    """
+    Calculates F (x sens_maps)
+
+    Parameters
+    ----------
+    x : ndarray
+        Single-channel image of shape (..., H, W, 2)
+    sens_maps : ndarray
+        Sensitivity maps (image space)
+
+    Returns
+    -------
+    ndarray
+        Result of the operation F (x sens_maps)
+    """
+    return fastmri.fft2c(fastmri.complex_mul(x, sens_maps))
+
+
+def sens_reduce(k: torch.Tensor, sens_maps: torch.Tensor) -> torch.Tensor:
+    """
+    Calculates F^{-1}(k) * conj(sens_maps)
+    where conj(sens_maps) is the element-wise applied complex conjugate
+
+    Parameters
+    ----------
+    k : ndarray
+        Multi-channel k-space of shape (B, C, H, W, 2)
+    sens_maps : ndarray
+        Sensitivity maps (image space)
+
+    Returns
+    -------
+    ndarray
+        Result of the operation F^{-1}(k) * conj(sens_maps)
+    """
+    return fastmri.complex_mul(fastmri.ifft2c(k), fastmri.complex_conj(sens_maps)).sum(
+        dim=1, keepdim=True
+    )
+
+
+def chans_to_batch_dim(x: torch.Tensor) -> Tuple[torch.Tensor, int]:
+    """Reshapes batched multi-channel samples into multiple single channel samples.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        x has shape (b, c, h, w, 2)
+
+    Returns
+    -------
+    Tuple[torch.Tensor, int]
+        tensor of shape (b * c, 1, h, w, 2), b
+    """
+    b, c, h, w, comp = x.shape
+    return x.view(b * c, 1, h, w, comp), b
+
+
+def batch_chans_to_chan_dim(x: torch.Tensor, batch_size: int) -> torch.Tensor:
+    """Reshapes batched independent samples into original multi-channel samples.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        tensor of shape (b * c, 1, h, w, 2)
+    batch_size : int
+        batch size
+
+    Returns
+    -------
+    torch.Tensor
+        original multi-channel tensor of shape (b, c, h, w, 2)
+    """
+    bc, _, h, w, comp = x.shape
+    c = bc // batch_size
+    return x.view(batch_size, c, h, w, comp)
+
+
+class NormUDNO(nn.Module):
+    """
+    Normalized UDNO model.
+
+    Inputs are normalized before the UDNO for numerically stable training.
+    """
+
+    def __init__(
+        self,
+        chans: int,
+        num_pool_layers: int,
+        radius_cutoff: float,
+        in_shape: Tuple[int, int],
+        kernel_shape: Tuple[int, int],
+        in_chans: int = 2,
+        out_chans: int = 2,
+        drop_prob: float = 0.0,
+    ):
+        """
+        Initialize the VarNet model.
+
+        Parameters
+        ----------
+        chans : int
+            Number of output channels of the first convolution layer.
+        num_pools : int
+            Number of down-sampling and up-sampling layers.
+        in_chans : int, optional
+            Number of channels in the input to the U-Net model. Default is 2.
+        out_chans : int, optional
+            Number of channels in the output to the U-Net model. Default is 2.
+        drop_prob : float, optional
+            Dropout probability. Default is 0.0.
+        """
+        super().__init__()
+
+        self.udno = UDNO(
+            in_chans=in_chans,
+            out_chans=out_chans,
+            radius_cutoff=radius_cutoff,
+            chans=chans,
+            num_pool_layers=num_pool_layers,
+            drop_prob=drop_prob,
+            in_shape=in_shape,
+            kernel_shape=kernel_shape,
+        )
+
+    def complex_to_chan_dim(self, x: torch.Tensor) -> torch.Tensor:
+        b, c, h, w, two = x.shape
+        assert two == 2
+        return x.permute(0, 4, 1, 2, 3).reshape(b, 2 * c, h, w)
+
+    def chan_complex_to_last_dim(self, x: torch.Tensor) -> torch.Tensor:
+        b, c2, h, w = x.shape
+        assert c2 % 2 == 0
+        c = c2 // 2
+        return x.view(b, 2, c, h, w).permute(0, 2, 3, 4, 1).contiguous()
+
+    def norm(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        # group norm
+        b, c, h, w = x.shape
+        x = x.view(b, 2, c // 2 * h * w)
+
+        mean = x.mean(dim=2).view(b, 2, 1, 1)
+        std = x.std(dim=2).view(b, 2, 1, 1)
+
+        x = x.view(b, c, h, w)
+
+        return (x - mean) / std, mean, std
+
+    def norm_new(
+        self, x: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        # FIXME: not working, wip
+        # group norm
+        b, c, h, w = x.shape
+        num_groups = 2
+        assert (
+            c % num_groups == 0
+        ), f"Number of channels ({c}) must be divisible by number of groups ({num_groups})."
+
+        x = x.view(b, num_groups, c // num_groups * h * w)
+
+        mean = x.mean(dim=2).view(b, num_groups, 1, 1)
+        std = x.std(dim=2).view(b, num_groups, 1, 1)
+        print(x.shape, mean.shape, std.shape)
+
+        x = x.view(b, c, h, w)
+        mean = (
+            mean.view(b, num_groups, 1, 1)
+            .repeat(1, c // num_groups, h, w)
+            .view(b, c, h, w)
+        )
+        std = (
+            std.view(b, num_groups, 1, 1)
+            .repeat(1, c // num_groups, h, w)
+            .view(b, c, h, w)
+        )
+
+        return (x - mean) / std, mean, std
+
+    def unnorm(
+        self, x: torch.Tensor, mean: torch.Tensor, std: torch.Tensor
+    ) -> torch.Tensor:
+        return x * std + mean
+
+    def pad(
+        self, x: torch.Tensor
+    ) -> Tuple[torch.Tensor, Tuple[List[int], List[int], int, int]]:
+        _, _, h, w = x.shape
+        w_mult = ((w - 1) | 15) + 1
+        h_mult = ((h - 1) | 15) + 1
+        w_pad = [math.floor((w_mult - w) / 2), math.ceil((w_mult - w) / 2)]
+        h_pad = [math.floor((h_mult - h) / 2), math.ceil((h_mult - h) / 2)]
+        # TODO: fix this type when PyTorch fixes theirs
+        # the documentation lies - this actually takes a list
+        # https://github.com/pytorch/pytorch/blob/master/torch/nn/functional.py#L3457
+        # https://github.com/pytorch/pytorch/pull/16949
+        x = F.pad(x, w_pad + h_pad)
+
+        return x, (h_pad, w_pad, h_mult, w_mult)
+
+    def unpad(
+        self,
+        x: torch.Tensor,
+        h_pad: List[int],
+        w_pad: List[int],
+        h_mult: int,
+        w_mult: int,
+    ) -> torch.Tensor:
+        return x[..., h_pad[0] : h_mult - h_pad[1], w_pad[0] : w_mult - w_pad[1]]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if not x.shape[-1] == 2:
+            raise ValueError("Last dimension must be 2 for complex.")
+
+        chans = x.shape[1]
+        if chans == 2:
+            # FIXME: hard coded skip norm/pad temporarily to avoid group norm bug
+            x = self.complex_to_chan_dim(x)
+            x = self.udno(x)
+            return self.chan_complex_to_last_dim(x)
+
+        # get shapes for unet and normalize
+        x = self.complex_to_chan_dim(x)
+        x, mean, std = self.norm(x)
+        x, pad_sizes = self.pad(x)
+
+        x = self.udno(x)
+
+        # get shapes back and unnormalize
+        x = self.unpad(x, *pad_sizes)
+        x = self.unnorm(x, mean, std)
+        x = self.chan_complex_to_last_dim(x)
+
+        return x
+
+
+class SensitivityModel(nn.Module):
+    """
+    Learn sensitivity maps
+    """
+
+    def __init__(
+        self,
+        chans: int,
+        num_pools: int,
+        radius_cutoff: float,
+        in_shape: Tuple[int, int],
+        kernel_shape: Tuple[int, int],
+        in_chans: int = 2,
+        out_chans: int = 2,
+        drop_prob: float = 0.0,
+        mask_center: bool = True,
+    ):
+        """
+        Parameters
+        ----------
+        chans : int
+            Number of output channels of the first convolution layer.
+        num_pools : int
+            Number of down-sampling and up-sampling layers.
+        in_chans : int, optional
+            Number of channels in the input to the U-Net model. Default is 2.
+        out_chans : int, optional
+            Number of channels in the output to the U-Net model. Default is 2.
+        drop_prob : float, optional
+            Dropout probability. Default is 0.0.
+        mask_center : bool, optional
+            Whether to mask center of k-space for sensitivity map calculation.
+            Default is True.
+        """
+        super().__init__()
+        self.mask_center = mask_center
+        self.norm_udno = NormUDNO(
+            chans,
+            num_pools,
+            radius_cutoff,
+            in_shape,
+            kernel_shape,
+            in_chans=in_chans,
+            out_chans=out_chans,
+            drop_prob=drop_prob,
+        )
+
+    def divide_root_sum_of_squares(self, x: torch.Tensor) -> torch.Tensor:
+        return x / fastmri.rss_complex(x, dim=1).unsqueeze(-1).unsqueeze(1)
+
+    def get_pad_and_num_low_freqs(
+        self, mask: torch.Tensor, num_low_frequencies: Optional[int] = None
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if num_low_frequencies is None or any(
+            torch.any(t == 0) for t in num_low_frequencies
+        ):
+            # get low frequency line locations and mask them out
+            squeezed_mask = mask[:, 0, 0, :, 0].to(torch.int8)
+            cent = squeezed_mask.shape[1] // 2
+            # running argmin returns the first non-zero
+            left = torch.argmin(squeezed_mask[:, :cent].flip(1), dim=1)
+            right = torch.argmin(squeezed_mask[:, cent:], dim=1)
+            num_low_frequencies_tensor = torch.max(
+                2 * torch.min(left, right), torch.ones_like(left)
+            )  # force a symmetric center unless 1
+        else:
+            num_low_frequencies_tensor = num_low_frequencies * torch.ones(
+                mask.shape[0], dtype=mask.dtype, device=mask.device
+            )
+
+        pad = (mask.shape[-2] - num_low_frequencies_tensor + 1) // 2
+
+        return pad.type(torch.long), num_low_frequencies_tensor.type(torch.long)
+
+    def forward(
+        self,
+        masked_kspace: torch.Tensor,
+        mask: torch.Tensor,
+        num_low_frequencies: Optional[int] = None,
+    ) -> torch.Tensor:
+        if self.mask_center:
+            pad, num_low_freqs = self.get_pad_and_num_low_freqs(
+                mask, num_low_frequencies
+            )
+            masked_kspace = transforms.batched_mask_center(
+                masked_kspace, pad, pad + num_low_freqs
+            )
+
+        # convert to image space
+        images, batches = chans_to_batch_dim(fastmri.ifft2c(masked_kspace))
+
+        # estimate sensitivities
+        return self.divide_root_sum_of_squares(
+            batch_chans_to_chan_dim(self.norm_udno(images), batches)
+        )
+
+
+class VarNetBlock(nn.Module):
+    """
+    Model block for iterative refinement of k-space data.
+
+    This model applies a combination of soft data consistency with the input
+    model as a regularizer. A series of these blocks can be stacked to form
+    the full variational network.
+
+    aka Refinement Module in Fig 1
+    """
+
+    def __init__(self, model: nn.Module):
+        """
+        Args:
+            model: Module for "regularization" component of variational
+                network.
+        """
+        super().__init__()
+
+        self.model = model
+        self.dc_weight = nn.Parameter(torch.ones(1))
+
+    def forward(
+        self,
+        current_kspace: torch.Tensor,
+        ref_kspace: torch.Tensor,
+        mask: torch.Tensor,
+        sens_maps: torch.Tensor,
+        use_dc_term: bool = True,
+    ) -> torch.Tensor:
+        """
+        Args:
+            current_kspace: The current k-space data (frequency domain data)
+                            being processed by the network. (torch.Tensor)
+            ref_kspace: Original subsampled k-space data (from which we are
+                reconstrucintg the image (reference k-space). (torch.Tensor)
+            mask: A binary mask indicating the locations in k-space where
+                data consistency should be enforced. (torch.Tensor)
+            sens_maps: Sensitivity maps for the different coils in parallel
+                    imaging. (torch.Tensor)
+        """
+
+        # model-term see orange box of Fig 1 in E2E-VarNet paper!
+        # multi channel k-space -> single channel image-space
+        b, c, h, w, _ = current_kspace.shape
+
+        if c == 30:
+            # get kspace and inpainted kspace
+            kspace = current_kspace[:, :15, :, :, :]
+            in_kspace = current_kspace[:, 15:, :, :, :]
+            # convert to image space
+            image = sens_reduce(kspace, sens_maps)
+            in_image = sens_reduce(in_kspace, sens_maps)
+            # concatenate both onto each other
+            reduced_image = torch.cat([image, in_image], dim=1)
+        else:
+            reduced_image = sens_reduce(current_kspace, sens_maps)
+
+        # single channel image-space
+        refined_image = self.model(reduced_image)
+
+        # single channel image-space -> multi channel k-space
+        model_term = sens_expand(refined_image, sens_maps)
+
+        # only use first 15 channels (masked_kspace) in the update
+        current_kspace = current_kspace[:, :15, :, :, :]
+
+        if not use_dc_term:
+            return current_kspace - model_term
+
+        """
+        Soft data consistency term:
+            - Calculates the difference between current k-space and reference k-space where the mask is true.
+            - Multiplies this difference by the data consistency weight.
+        """
+        # dc_term: see green box of Fig 1 in E2E-VarNet paper!
+        zero = torch.zeros(1, 1, 1, 1, 1).to(current_kspace)
+        soft_dc = torch.where(mask, current_kspace - ref_kspace, zero) * self.dc_weight
+        return current_kspace - soft_dc - model_term
+
+
+class NOVarnet(nn.Module):
+    """
+    Neural Operator model for MRI reconstruction.
+
+    Uses a variational architecture (iterative updates) with a learned sensitivity
+    model. All operations are resolution invariant employing neural operator
+    modules (GNO, UDNO).
+    """
+
+    def __init__(
+        self,
+        num_cascades: int = 12,
+        sens_chans: int = 8,
+        sens_pools: int = 4,
+        chans: int = 18,
+        pools: int = 4,
+        gno_chans: int = 16,
+        gno_pools: int = 4,
+        gno_radius_cutoff: float = 0.02,
+        gno_kernel_shape: Tuple[int, int] = (6, 7),
+        radius_cutoff: float = 0.01,
+        kernel_shape: Tuple[int, int] = (3, 4),
+        in_shape: Tuple[int, int] = (640, 320),
+        mask_center: bool = True,
+        use_dc_term: bool = True,
+        reduction_method: Literal["batch", "rss"] = "rss",
+        skip_method: Literal["replace", "add", "add_inv", "concat"] = "add",
+    ):
+        """
+        Parameters
+        ----------
+        num_cascades : int
+            Number of cascades (i.e., layers) for variational network.
+        sens_chans : int
+            Number of channels for sensitivity map U-Net.
+        sens_pools : int
+            Number of downsampling and upsampling layers for sensitivity map U-Net.
+        chans : int
+            Number of channels for cascade U-Net.
+        pools : int
+            Number of downsampling and upsampling layers for cascade U-Net.
+        mask_center : bool
+            Whether to mask center of k-space for sensitivity map calculation.
+        use_dc_term : bool
+            Whether to use the data consistency term.
+        reduction_method : "batch" or "rss"
+            Method for reducing sensitivity maps to single channel.
+            "batch" reduces to single channel by stacking channels.
+            "rss" reduces to single channel by root sum of squares.
+        skip_method : "replace" or "add" or "add_inv" or "concat"
+            "replace" replaces the input with the output of the GNO
+            "add" adds the output of the GNO to the input
+            "add_inv" adds the output of the GNO to the input (only where samples are missing)
+            "concat" concatenates the output of the GNO to the input
+        """
+
+        super().__init__()
+
+        self.sens_net = SensitivityModel(
+            sens_chans,
+            sens_pools,
+            radius_cutoff,
+            in_shape,
+            kernel_shape,
+            mask_center=mask_center,
+        )
+        self.gno = NormUDNO(
+            gno_chans,
+            gno_pools,
+            in_shape=in_shape,
+            radius_cutoff=radius_cutoff,
+            kernel_shape=kernel_shape,
+            in_chans=2,
+            out_chans=2,
+        )
+        self.cascades = nn.ModuleList(
+            [
+                VarNetBlock(
+                    NormUDNO(
+                        chans,
+                        pools,
+                        radius_cutoff,
+                        in_shape,
+                        kernel_shape,
+                        in_chans=(
+                            4 if skip_method == "concat" and cascade_idx == 0 else 2
+                        ),
+                        out_chans=2,
+                    )
+                )
+                for cascade_idx in range(num_cascades)
+            ]
+        )
+        self.use_dc_term = use_dc_term
+        self.reduction_method = reduction_method
+        self.skip_method = skip_method
+
+        print("===================================")
+        print("===================================")
+        print("initialized no repeat k ")
+        print("===================================")
+        print("===================================")
+
+    def forward(
+        self,
+        masked_kspace: torch.Tensor,
+        mask: torch.Tensor,
+        num_low_frequencies: Optional[int] = None,
+    ) -> torch.Tensor:
+
+        # (B, C, X, Y, 2)
+        sens_maps = self.sens_net(masked_kspace, mask, num_low_frequencies)
+
+        kspace_pred = masked_kspace
+        # iterative update
+        for cascade in self.cascades:
+            # breakpoint()
+            sens_maps = self.sens_net(kspace_pred, mask, num_low_frequencies)
+            # reduce before inpainting
+            kspace_pred, b = chans_to_batch_dim(kspace_pred)
+            # inpainting
+            kspace_pred = self.gno(kspace_pred)
+            kspace_pred = batch_chans_to_chan_dim(kspace_pred, b)
+
+            # image
+            kspace_pred = cascade(
+                kspace_pred, masked_kspace, mask, sens_maps, self.use_dc_term
+            )
+
+        spatial_pred = fastmri.ifft2c(kspace_pred)
+        spatial_pred_abs = fastmri.complex_abs(spatial_pred)
+        combined_spatial = fastmri.rss(spatial_pred_abs, dim=1)
+
+        return combined_spatial

models/temp/no_repeatk_module.py

@@ -0,0 +1,303 @@
+from argparse import ArgumentParser
+from typing import Tuple
+
+import torch
+
+import fastmri
+from fastmri import transforms
+from models.temp.no_repeatk import NOVarnet
+
+from models.lightning.mri_module import MriModule
+from type_utils import tuple_type
+
+
+class NORepeatKModule(MriModule):
+    """
+    NO-Varnet repeat-k (temp) training module.
+    """
+
+    def __init__(
+        self,
+        num_cascades: int = 12,
+        pools: int = 4,
+        chans: int = 18,
+        sens_pools: int = 4,
+        sens_chans: int = 8,
+        gno_pools: int = 4,
+        gno_chans: int = 16,
+        gno_radius_cutoff: float = 0.02,
+        gno_kernel_shape: Tuple[int, int] = (6, 7),
+        radius_cutoff: float = 0.02,
+        kernel_shape: Tuple[int, int] = (6, 7),
+        in_shape: Tuple[int, int] = (320, 320),
+        use_dc_term: bool = True,
+        lr: float = 0.0003,
+        lr_step_size: int = 40,
+        lr_gamma: float = 0.1,
+        weight_decay: float = 0.0,
+        reduction_method: str = "rss",
+        skip_method: str = "add",
+        **kwargs,
+    ):
+        """
+        Parameters
+        ----------
+        num_cascades : int
+            Number of cascades (i.e., layers) for the variational network.
+        pools : int
+            Number of downsampling and upsampling layers for the cascade U-Net.
+        chans : int
+            Number of channels for the cascade U-Net.
+        sens_pools : int
+            Number of downsampling and upsampling layers for the sensitivity map U-Net.
+        sens_chans : int
+            Number of channels for the sensitivity map U-Net.
+        lr : float
+            Learning rate.
+        lr_step_size : int
+            Learning rate step size.
+        lr_gamma : float
+            Learning rate gamma decay.
+        weight_decay : float
+            Parameter for penalizing weights norm.
+        num_sense_lines : int, optional
+            Number of low-frequency lines to use for sensitivity map computation.
+            Must be even or `None`. Default `None` will automatically compute the number
+            from masks. Default behavior may cause some slices to use more low-frequency
+            lines than others, when used in conjunction with e.g. the EquispacedMaskFunc
+            defaults. To prevent this, either set `num_sense_lines`, or set
+            `skip_low_freqs` and `skip_around_low_freqs` to `True` in the EquispacedMaskFunc.
+            Note that setting this value may lead to undesired behavior when training on
+            multiple accelerations simultaneously.
+        """
+        super().__init__(**kwargs)
+        self.save_hyperparameters()
+
+        self.num_cascades = num_cascades
+        self.pools = pools
+        self.chans = chans
+        self.sens_pools = sens_pools
+        self.sens_chans = sens_chans
+        self.gno_pools = gno_pools
+        self.gno_chans = gno_chans
+        self.gno_radius_cutoff = gno_radius_cutoff
+        self.gno_kernel_shape = gno_kernel_shape
+        self.radius_cutoff = radius_cutoff
+        self.kernel_shape = kernel_shape
+        self.in_shape = in_shape
+        self.use_dc_term = use_dc_term
+        self.lr = lr
+        self.lr_step_size = lr_step_size
+        self.lr_gamma = lr_gamma
+        self.weight_decay = weight_decay
+        self.reduction_method = reduction_method
+        self.skip_method = skip_method
+
+        self.model = NOVarnet(
+            num_cascades=self.num_cascades,
+            sens_chans=self.sens_chans,
+            sens_pools=self.sens_pools,
+            chans=self.chans,
+            pools=self.pools,
+            gno_chans=self.gno_chans,
+            gno_pools=self.gno_pools,
+            gno_radius_cutoff=self.gno_radius_cutoff,
+            gno_kernel_shape=self.gno_kernel_shape,
+            radius_cutoff=radius_cutoff,
+            kernel_shape=kernel_shape,
+            in_shape=in_shape,
+            use_dc_term=use_dc_term,
+            reduction_method=reduction_method,
+            skip_method=skip_method,
+        )
+
+        self.criterion = fastmri.SSIMLoss()
+        self.num_params = sum(p.numel() for p in self.parameters())
+
+    def forward(self, masked_kspace, mask, num_low_frequencies):
+        return self.model(masked_kspace, mask, num_low_frequencies)
+
+    def training_step(self, batch, batch_idx):
+        output = self.forward(
+            batch.masked_kspace, batch.mask, batch.num_low_frequencies
+        )
+
+        target, output = transforms.center_crop_to_smallest(batch.target, output)
+        loss = self.criterion(
+            output.unsqueeze(1), target.unsqueeze(1), data_range=batch.max_value
+        )
+
+        self.log("train_loss", loss, on_step=True, on_epoch=True)
+        self.log("epoch", int(self.current_epoch), on_step=True, on_epoch=True)
+
+        return loss
+
+    def validation_step(self, batch, batch_idx, dataloader_idx=0):
+        dataloaders = self.trainer.val_dataloaders
+        slug = list(dataloaders.keys())[dataloader_idx]
+
+        output = self.forward(
+            batch.masked_kspace, batch.mask, batch.num_low_frequencies
+        )
+
+        target, output = transforms.center_crop_to_smallest(batch.target, output)
+
+        loss = self.criterion(
+            output.unsqueeze(1),
+            target.unsqueeze(1),
+            data_range=batch.max_value,
+        )
+
+        return {
+            "slug": slug,
+            "fname": batch.fname,
+            "slice_num": batch.slice_num,
+            "max_value": batch.max_value,
+            "output": output,
+            "target": target,
+            "val_loss": loss,
+        }
+
+    def configure_optimizers(self):
+        optim = torch.optim.Adam(
+            self.parameters(), lr=self.lr, weight_decay=self.weight_decay
+        )
+        scheduler = torch.optim.lr_scheduler.StepLR(
+            optim, self.lr_step_size, self.lr_gamma
+        )
+
+        return [optim], [scheduler]
+
+    @staticmethod
+    def add_model_specific_args(parent_parser):
+        """
+        Define parameters that only apply to this model
+        """
+        parser = ArgumentParser(parents=[parent_parser], add_help=False)
+        parser = MriModule.add_model_specific_args(parser)
+
+        # network params
+        parser.add_argument(
+            "--num_cascades",
+            default=12,
+            type=int,
+            help="Number of VarNet cascades",
+        )
+        parser.add_argument(
+            "--pools",
+            default=4,
+            type=int,
+            help="Number of U-Net pooling layers in VarNet blocks",
+        )
+        parser.add_argument(
+            "--chans",
+            default=18,
+            type=int,
+            help="Number of channels for U-Net in VarNet blocks",
+        )
+        parser.add_argument(
+            "--sens_pools",
+            default=4,
+            type=int,
+            help=(
+                "Number of pooling layers for sense map estimation U-Net in" " VarNet"
+            ),
+        )
+        parser.add_argument(
+            "--sens_chans",
+            default=8,
+            type=float,
+            help="Number of channels for sense map estimation U-Net in VarNet",
+        )
+        parser.add_argument(
+            "--gno_pools",
+            default=4,
+            type=int,
+            help=("Number of pooling layers for GNO"),
+        )
+        parser.add_argument(
+            "--gno_chans",
+            default=16,
+            type=int,
+            help="Number of channels for GNO",
+        )
+        parser.add_argument(
+            "--gno_radius_cutoff",
+            default=0.02,
+            type=float,
+            required=True,
+            help="GNO module radius_cutoff",
+        )
+        parser.add_argument(
+            "--gno_kernel_shape",
+            default=(6, 7),
+            type=tuple_type,
+            required=True,
+            help="GNO module kernel_shape. Ex: (6, 7)",
+        )
+        parser.add_argument(
+            "--radius_cutoff",
+            default=0.01,
+            type=float,
+            required=True,
+            help="DISCO module radius_cutoff",
+        )
+        parser.add_argument(
+            "--kernel_shape",
+            default=(6, 7),
+            type=tuple_type,
+            required=True,
+            help="DISCO module kernel_shape. Ex: (6, 7)",
+        )
+        parser.add_argument(
+            "--in_shape",
+            default=(640, 320),
+            type=tuple_type,
+            required=True,
+            help="Spatial dimensions of masked_kspace samples. Ex: (640, 320)",
+        )
+        parser.add_argument(
+            "--use_dc_term",
+            default=True,
+            type=bool,
+            help="Whether to use the DC term in the unrolled iterative update step",
+        )
+
+        # training params (opt)
+        parser.add_argument(
+            "--lr", default=0.0003, type=float, help="Adam learning rate"
+        )
+        parser.add_argument(
+            "--lr_step_size",
+            default=40,
+            type=int,
+            help="Epoch at which to decrease step size",
+        )
+        parser.add_argument(
+            "--lr_gamma",
+            default=0.1,
+            type=float,
+            help="Extent to which step size should be decreased",
+        )
+        parser.add_argument(
+            "--weight_decay",
+            default=0.0,
+            type=float,
+            help="Strength of weight decay regularization",
+        )
+        parser.add_argument(
+            "--reduction_method",
+            default="rss",
+            type=str,
+            choices=["rss", "batch"],
+            help="Reduction method used to reduce multi-channel k-space data before inpainting module. Read documentation of GNO for more information.",
+        )
+        parser.add_argument(
+            "--skip_method",
+            default="add_inv",
+            type=str,
+            choices=["add_inv", "add", "concat", "replace"],
+            help="Method for skip connection around inpainting module.",
+        )
+
+        return parser

models/udno.py

@@ -0,0 +1,369 @@
+"""
+U-shaped DISCO Neural Operator
+"""
+
+from typing import List, Tuple
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+from torch_harmonics_local.convolution import (
+    EquidistantDiscreteContinuousConv2d as DISCO2d,
+)
+
+
+class UDNO(nn.Module):
+    """
+    U-shaped DISCO Neural Operator in PyTorch
+    """
+
+    def __init__(
+        self,
+        in_chans: int,
+        out_chans: int,
+        radius_cutoff: float,
+        chans: int = 32,
+        num_pool_layers: int = 4,
+        drop_prob: float = 0.0,
+        in_shape: Tuple[int, int] = (320, 320),
+        kernel_shape: Tuple[int, int] = (3, 4),
+    ):
+        """
+        Parameters
+        ----------
+        in_chans : int
+            Number of channels in the input to the U-Net model.
+        out_chans : int
+            Number of channels in the output to the U-Net model.
+        radius_cutoff : float
+            Control the effective radius of the DISCO kernel. Values are
+            between 0.0 and 1.0. The radius_cutoff is represented as a proportion
+            of the normalized input space, to ensure that kernels are resolution
+            invaraint.
+        chans : int, optional
+            Number of output channels of the first DISCO layer. Default is 32.
+        num_pool_layers : int, optional
+            Number of down-sampling and up-sampling layers. Default is 4.
+        drop_prob : float, optional
+            Dropout probability. Default is 0.0.
+        in_shape : Tuple[int, int]
+            Shape of the input to the UDNO. This is required to dynamically
+            compile DISCO kernels for resolution invariance.
+        kernel_shape : Tuple[int, int], optional
+            Shape of the DISCO kernel. Default is (3, 4). This corresponds to 3
+            rings and 4 anisotropic basis functions. Under the hood, each DISCO
+            kernel has (3 - 1) * 4 + 1 = 9 parameters, equivalent to a standard
+            3x3 convolution kernel.
+
+            Note: This is NOT kernel_size, as under the DISCO framework,
+            kernels are dynamically compiled to support resolution invariance.
+        """
+        super().__init__()
+        assert len(in_shape) == 2, "Input shape must be 2D"
+
+        self.in_chans = in_chans
+        self.out_chans = out_chans
+        self.chans = chans
+        self.num_pool_layers = num_pool_layers
+        self.drop_prob = drop_prob
+        self.in_shape = in_shape
+        self.kernel_shape = kernel_shape
+
+        self.down_sample_layers = nn.ModuleList(
+            [
+                DISCOBlock(
+                    in_chans,
+                    chans,
+                    radius_cutoff,
+                    drop_prob,
+                    in_shape,
+                    kernel_shape,
+                )
+            ]
+        )
+        ch = chans
+        shape = (in_shape[0] // 2, in_shape[1] // 2)
+        radius_cutoff = radius_cutoff * 2
+        for _ in range(num_pool_layers - 1):
+            self.down_sample_layers.append(
+                DISCOBlock(
+                    ch,
+                    ch * 2,
+                    radius_cutoff,
+                    drop_prob,
+                    in_shape=shape,
+                    kernel_shape=kernel_shape,
+                )
+            )
+            ch *= 2
+            shape = (shape[0] // 2, shape[1] // 2)
+            radius_cutoff *= 2
+
+            # test commit
+
+        self.bottleneck = DISCOBlock(
+            ch,
+            ch * 2,
+            radius_cutoff,
+            drop_prob,
+            in_shape=shape,
+            kernel_shape=kernel_shape,
+        )
+
+        self.up = nn.ModuleList()
+        self.up_transpose = nn.ModuleList()
+        for _ in range(num_pool_layers - 1):
+            self.up_transpose.append(
+                TransposeDISCOBlock(
+                    ch * 2,
+                    ch,
+                    radius_cutoff,
+                    in_shape=shape,
+                    kernel_shape=kernel_shape,
+                )
+            )
+            shape = (shape[0] * 2, shape[1] * 2)
+            radius_cutoff /= 2
+            self.up.append(
+                DISCOBlock(
+                    ch * 2,
+                    ch,
+                    radius_cutoff,
+                    drop_prob,
+                    in_shape=shape,
+                    kernel_shape=kernel_shape,
+                )
+            )
+            ch //= 2
+
+        self.up_transpose.append(
+            TransposeDISCOBlock(
+                ch * 2,
+                ch,
+                radius_cutoff,
+                in_shape=shape,
+                kernel_shape=kernel_shape,
+            )
+        )
+        shape = (shape[0] * 2, shape[1] * 2)
+        radius_cutoff /= 2
+        self.up.append(
+            nn.Sequential(
+                DISCOBlock(
+                    ch * 2,
+                    ch,
+                    radius_cutoff,
+                    drop_prob,
+                    in_shape=shape,
+                    kernel_shape=kernel_shape,
+                ),
+                nn.Conv2d(
+                    ch, self.out_chans, kernel_size=1, stride=1
+                ),  # 1x1 conv is always res-invariant (pixel wise channel transformation)
+            )
+        )
+
+    def forward(self, image: torch.Tensor) -> torch.Tensor:
+        """
+        Parameters
+        ----------
+        image : torch.Tensor
+            Input 4D tensor of shape `(N, in_chans, H, W)`.
+
+        Returns
+        -------
+        torch.Tensor
+            Output tensor of shape `(N, out_chans, H, W)`.
+        """
+        stack = []
+        output = image
+
+        # apply down-sampling layers
+        for layer in self.down_sample_layers:
+            output = layer(output)
+            stack.append(output)
+            output = F.avg_pool2d(output, kernel_size=2, stride=2, padding=0)
+
+        output = self.bottleneck(output)
+
+        # apply up-sampling layers
+        for transpose, disco in zip(self.up_transpose, self.up):
+            downsample_layer = stack.pop()
+            output = transpose(output)
+
+            # reflect pad on the right/botton if needed to handle odd input dimensions
+            padding = [0, 0, 0, 0]
+            if output.shape[-1] != downsample_layer.shape[-1]:
+                padding[1] = 1  # padding right
+            if output.shape[-2] != downsample_layer.shape[-2]:
+                padding[3] = 1  # padding bottom
+            if torch.sum(torch.tensor(padding)) != 0:
+                output = F.pad(output, padding, "reflect")
+
+            output = torch.cat([output, downsample_layer], dim=1)
+            output = disco(output)
+
+        return output
+
+
+class DISCOBlock(nn.Module):
+    """
+    A DISCO Block that consists of two DISCO layers each followed by
+    instance normalization, LeakyReLU activation and dropout.
+    """
+
+    def __init__(
+        self,
+        in_chans: int,
+        out_chans: int,
+        radius_cutoff: float,
+        drop_prob: float,
+        in_shape: Tuple[int, int],
+        kernel_shape: Tuple[int, int] = (3, 4),
+    ):
+        """
+        Parameters
+        ----------
+        in_chans : int
+            Number of channels in the input.
+        out_chans : int
+            Number of channels in the output.
+        radius_cutoff : float
+            Control the effective radius of the DISCO kernel. Values are
+            between 0.0 and 1.0. The radius_cutoff is represented as a proportion
+            of the normalized input space, to ensure that kernels are resolution
+            invaraint.
+        in_shape : Tuple[int]
+            Unbatched spatial 2D shape of the input to this block.
+            Rrequired to dynamically compile DISCO kernels for resolution invariance.
+        kernel_shape : Tuple[int, int], optional
+            Shape of the DISCO kernel. Default is (3, 4). This corresponds to 3
+            rings and 4 anisotropic basis functions. Under the hood, each DISCO
+            kernel has (3 - 1) * 4 + 1 = 9 parameters, equivalent to a standard
+            3x3 convolution kernel.
+
+            Note: This is NOT kernel_size, as under the DISCO framework,
+            kernels are dynamically compiled to support resolution invariance.
+        drop_prob : float
+            Dropout probability.
+        """
+        super().__init__()
+
+        self.in_chans = in_chans
+        self.out_chans = out_chans
+        self.drop_prob = drop_prob
+
+        self.layers = nn.Sequential(
+            DISCO2d(
+                in_chans,
+                out_chans,
+                kernel_shape=kernel_shape,
+                in_shape=in_shape,
+                bias=False,
+                radius_cutoff=radius_cutoff,
+                padding_mode="constant",
+            ),
+            nn.InstanceNorm2d(out_chans),
+            nn.LeakyReLU(negative_slope=0.2, inplace=True),
+            nn.Dropout2d(drop_prob),
+            DISCO2d(
+                out_chans,
+                out_chans,
+                kernel_shape=kernel_shape,
+                in_shape=in_shape,
+                bias=False,
+                radius_cutoff=radius_cutoff,
+                padding_mode="constant",
+            ),
+            nn.InstanceNorm2d(out_chans),
+            nn.LeakyReLU(negative_slope=0.2, inplace=True),
+            nn.Dropout2d(drop_prob),
+        )
+
+    def forward(self, image: torch.Tensor) -> torch.Tensor:
+        """
+        Parameters
+        ----------
+        image : ndarray
+            Input 4D tensor of shape `(N, in_chans, H, W)`.
+
+        Returns
+        -------
+        ndarray
+            Output tensor of shape `(N, out_chans, H, W)`.
+        """
+        return self.layers(image)
+
+
+class TransposeDISCOBlock(nn.Module):
+    """
+    A transpose DISCO Block that consists of an up-sampling layer followed by a
+    DISCO layer, instance normalization, and LeakyReLU activation.
+    """
+
+    def __init__(
+        self,
+        in_chans: int,
+        out_chans: int,
+        radius_cutoff: float,
+        in_shape: Tuple[int, int],
+        kernel_shape: Tuple[int, int] = (3, 4),
+    ):
+        """
+        Parameters
+        ----------
+        in_chans : int
+            Number of channels in the input.
+        out_chans : int
+            Number of channels in the output.
+        radius_cutoff : float
+            Control the effective radius of the DISCO kernel. Values are
+            between 0.0 and 1.0. The radius_cutoff is represented as a proportion
+            of the normalized input space, to ensure that kernels are resolution
+            invaraint.
+        in_shape : Tuple[int]
+            Unbatched spatial 2D shape of the input to this block.
+            Rrequired to dynamically compile DISCO kernels for resolution invariance.
+        kernel_shape : Tuple[int, int], optional
+            Shape of the DISCO kernel. Default is (3, 4). This corresponds to 3
+            rings and 4 anisotropic basis functions. Under the hood, each DISCO
+            kernel has (3 - 1) * 4 + 1 = 9 parameters, equivalent to a standard
+            3x3 convolution kernel.
+
+            Note: This is NOT kernel_size, as under the DISCO framework,
+            kernels are dynamically compiled to support resolution invariance
+        """
+        super().__init__()
+
+        self.in_chans = in_chans
+        self.out_chans = out_chans
+
+        self.layers = nn.Sequential(
+            nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True),
+            DISCO2d(
+                in_chans,
+                out_chans,
+                kernel_shape=kernel_shape,
+                in_shape=(2 * in_shape[0], 2 * in_shape[1]),
+                bias=False,
+                radius_cutoff=(radius_cutoff / 2),
+                padding_mode="constant",
+            ),
+            nn.InstanceNorm2d(out_chans),
+            nn.LeakyReLU(negative_slope=0.2, inplace=True),
+        )
+
+    def forward(self, image: torch.Tensor) -> torch.Tensor:
+        """
+        Parameters
+        ----------
+        image : torch.Tensor
+            Input 4D tensor of shape `(N, in_chans, H, W)`.
+
+        Returns
+        -------
+        torch.Tensor
+            Output tensor of shape `(N, out_chans, H*2, W*2)`.
+        """
+        return self.layers(image)

models/unet.py

@@ -0,0 +1,209 @@
+"""
+Copyright (c) Facebook, Inc. and its affiliates.
+
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+from typing import List, Tuple
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+class Unet(nn.Module):
+    """
+    PyTorch implementation of a U-Net model.
+
+    O. Ronneberger, P. Fischer, and Thomas Brox. U-net: Convolutional networks
+    for biomedical image segmentation. In International Conference on Medical
+    image computing and computer-assisted intervention, pages 234–241.
+    Springer, 2015.
+    """
+
+    def __init__(
+        self,
+        in_chans: int,
+        out_chans: int,
+        chans: int = 32,
+        num_pool_layers: int = 4,
+        drop_prob: float = 0.0,
+    ):
+        """
+        Parameters
+        ----------
+        in_chans : int
+            Number of channels in the input to the U-Net model.
+        out_chans : int
+            Number of channels in the output to the U-Net model.
+        chans : int, optional
+            Number of output channels of the first convolution layer. Default is 32.
+        num_pool_layers : int, optional
+            Number of down-sampling and up-sampling layers. Default is 4.
+        drop_prob : float, optional
+            Dropout probability. Default is 0.0.
+        """
+        super().__init__()
+
+        self.in_chans = in_chans
+        self.out_chans = out_chans
+        self.chans = chans
+        self.num_pool_layers = num_pool_layers
+        self.drop_prob = drop_prob
+
+        self.down_sample_layers = nn.ModuleList([ConvBlock(in_chans, chans, drop_prob)])
+        ch = chans
+        for _ in range(num_pool_layers - 1):
+            self.down_sample_layers.append(ConvBlock(ch, ch * 2, drop_prob))
+            ch *= 2
+        self.conv = ConvBlock(ch, ch * 2, drop_prob)
+
+        self.up_conv = nn.ModuleList()
+        self.up_transpose_conv = nn.ModuleList()
+        for _ in range(num_pool_layers - 1):
+            self.up_transpose_conv.append(TransposeConvBlock(ch * 2, ch))
+            self.up_conv.append(ConvBlock(ch * 2, ch, drop_prob))
+            ch //= 2
+
+        self.up_transpose_conv.append(TransposeConvBlock(ch * 2, ch))
+        self.up_conv.append(
+            nn.Sequential(
+                ConvBlock(ch * 2, ch, drop_prob),
+                nn.Conv2d(ch, self.out_chans, kernel_size=1, stride=1),
+            )
+        )
+
+    def forward(self, image: torch.Tensor) -> torch.Tensor:
+        """
+        Parameters
+        ----------
+        image : torch.Tensor
+            Input 4D tensor of shape `(N, in_chans, H, W)`.
+
+        Returns
+        -------
+        torch.Tensor
+            Output tensor of shape `(N, out_chans, H, W)`.
+        """
+        stack = []
+        output = image
+
+        # apply down-sampling layers
+        for layer in self.down_sample_layers:
+            output = layer(output)
+            stack.append(output)
+            output = F.avg_pool2d(output, kernel_size=2, stride=2, padding=0)
+
+        output = self.conv(output)
+
+        # apply up-sampling layers
+        for transpose_conv, conv in zip(self.up_transpose_conv, self.up_conv):
+            downsample_layer = stack.pop()
+            output = transpose_conv(output)
+
+            # reflect pad on the right/botton if needed to handle odd input dimensions
+            padding = [0, 0, 0, 0]
+            if output.shape[-1] != downsample_layer.shape[-1]:
+                padding[1] = 1  # padding right
+            if output.shape[-2] != downsample_layer.shape[-2]:
+                padding[3] = 1  # padding bottom
+            if torch.sum(torch.tensor(padding)) != 0:
+                output = F.pad(output, padding, "reflect")
+
+            output = torch.cat([output, downsample_layer], dim=1)
+            output = conv(output)
+
+        return output
+
+
+class ConvBlock(nn.Module):
+    """
+    A Convolutional Block that consists of two convolution layers each followed by
+    instance normalization, LeakyReLU activation and dropout.
+    """
+
+    def __init__(self, in_chans: int, out_chans: int, drop_prob: float):
+        """
+        Parameters
+        ----------
+        in_chans : int
+            Number of channels in the input.
+        out_chans : int
+            Number of channels in the output.
+        drop_prob : float
+            Dropout probability.
+        """
+        super().__init__()
+
+        self.in_chans = in_chans
+        self.out_chans = out_chans
+        self.drop_prob = drop_prob
+
+        self.layers = nn.Sequential(
+            nn.Conv2d(in_chans, out_chans, kernel_size=3, padding=1, bias=False),
+            nn.InstanceNorm2d(out_chans),
+            nn.LeakyReLU(negative_slope=0.2, inplace=True),
+            nn.Dropout2d(drop_prob),
+            nn.Conv2d(out_chans, out_chans, kernel_size=3, padding=1, bias=False),
+            nn.InstanceNorm2d(out_chans),
+            nn.LeakyReLU(negative_slope=0.2, inplace=True),
+            nn.Dropout2d(drop_prob),
+        )
+
+    def forward(self, image: torch.Tensor) -> torch.Tensor:
+        """
+        Parameters
+        ----------
+        image : ndarray
+            Input 4D tensor of shape `(N, in_chans, H, W)`.
+
+        Returns
+        -------
+        ndarray
+            Output tensor of shape `(N, out_chans, H, W)`.
+        """
+        return self.layers(image)
+
+
+class TransposeConvBlock(nn.Module):
+    """
+    A Transpose Convolutional Block that consists of one convolution transpose
+    layers followed by instance normalization and LeakyReLU activation.
+    """
+
+    def __init__(self, in_chans: int, out_chans: int):
+        """
+        Parameters
+        ----------
+        in_chans : int
+            Number of channels in the input.
+        out_chans : int
+            Number of channels in the output.
+        """
+        super().__init__()
+
+        self.in_chans = in_chans
+        self.out_chans = out_chans
+
+        self.layers = nn.Sequential(
+            nn.ConvTranspose2d(
+                in_chans, out_chans, kernel_size=2, stride=2, bias=False
+            ),
+            nn.InstanceNorm2d(out_chans),
+            nn.LeakyReLU(negative_slope=0.2, inplace=True),
+        )
+
+    def forward(self, image: torch.Tensor) -> torch.Tensor:
+        """
+        Parameters
+        ----------
+        image : torch.Tensor
+            Input 4D tensor of shape `(N, in_chans, H, W)`.
+
+        Returns
+        -------
+        torch.Tensor
+            Output tensor of shape `(N, out_chans, H*2, W*2)`.
+        """
+        return self.layers(image)

models/varnet.py

@@ -0,0 +1,416 @@
+"""
+Copyright (c) Facebook, Inc. and its affiliates.
+
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+import math
+import os
+from typing import List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import fastmri
+from fastmri import transforms
+from models.unet import Unet
+
+
+class NormUnet(nn.Module):
+    """
+    Normalized U-Net model.
+
+    This is the same as a regular U-Net, but with normalization applied to the
+    input before the U-Net. This keeps the values more numerically stable
+    during training.
+    """
+
+    def __init__(
+        self,
+        chans: int,
+        num_pools: int,
+        in_chans: int = 2,
+        out_chans: int = 2,
+        drop_prob: float = 0.0,
+    ):
+        """
+
+        Initialize the VarNet model.
+
+        Parameters
+        ----------
+        chans : int
+            Number of output channels of the first convolution layer.
+        num_pools : int
+            Number of down-sampling and up-sampling layers.
+        in_chans : int, optional
+            Number of channels in the input to the U-Net model. Default is 2.
+        out_chans : int, optional
+            Number of channels in the output to the U-Net model. Default is 2.
+        drop_prob : float, optional
+            Dropout probability. Default is 0.0.
+        """
+        super().__init__()
+
+        self.unet = Unet(
+            in_chans=in_chans,
+            out_chans=out_chans,
+            chans=chans,
+            num_pool_layers=num_pools,
+            drop_prob=drop_prob,
+        )
+
+    def complex_to_chan_dim(self, x: torch.Tensor) -> torch.Tensor:
+        b, c, h, w, two = x.shape
+        assert two == 2
+        return x.permute(0, 4, 1, 2, 3).reshape(b, 2 * c, h, w)
+
+    def chan_complex_to_last_dim(self, x: torch.Tensor) -> torch.Tensor:
+        b, c2, h, w = x.shape
+        assert c2 % 2 == 0
+        c = c2 // 2
+        return x.view(b, 2, c, h, w).permute(0, 2, 3, 4, 1).contiguous()
+
+    def norm(
+        self, x: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        # group norm
+        b, c, h, w = x.shape
+        x = x.view(b, 2, c // 2 * h * w)
+
+        mean = x.mean(dim=2).view(b, 2, 1, 1)
+        std = x.std(dim=2).view(b, 2, 1, 1)
+
+        x = x.view(b, c, h, w)
+
+        return (x - mean) / std, mean, std
+
+    def unnorm(
+        self, x: torch.Tensor, mean: torch.Tensor, std: torch.Tensor
+    ) -> torch.Tensor:
+        return x * std + mean
+
+    def pad(
+        self, x: torch.Tensor
+    ) -> Tuple[torch.Tensor, Tuple[List[int], List[int], int, int]]:
+        _, _, h, w = x.shape
+        w_mult = ((w - 1) | 15) + 1
+        h_mult = ((h - 1) | 15) + 1
+        w_pad = [math.floor((w_mult - w) / 2), math.ceil((w_mult - w) / 2)]
+        h_pad = [math.floor((h_mult - h) / 2), math.ceil((h_mult - h) / 2)]
+        # TODO: fix this type when PyTorch fixes theirs
+        # the documentation lies - this actually takes a list
+        # https://github.com/pytorch/pytorch/blob/master/torch/nn/functional.py#L3457
+        # https://github.com/pytorch/pytorch/pull/16949
+        x = F.pad(x, w_pad + h_pad)
+
+        return x, (h_pad, w_pad, h_mult, w_mult)
+
+    def unpad(
+        self,
+        x: torch.Tensor,
+        h_pad: List[int],
+        w_pad: List[int],
+        h_mult: int,
+        w_mult: int,
+    ) -> torch.Tensor:
+        return x[
+            ..., h_pad[0] : h_mult - h_pad[1], w_pad[0] : w_mult - w_pad[1]
+        ]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if not x.shape[-1] == 2:
+            raise ValueError("Last dimension must be 2 for complex.")
+
+        # get shapes for unet and normalize
+        x = self.complex_to_chan_dim(x)
+        x, mean, std = self.norm(x)
+        x, pad_sizes = self.pad(x)
+
+        x = self.unet(x)
+
+        # get shapes back and unnormalize
+        x = self.unpad(x, *pad_sizes)
+        x = self.unnorm(x, mean, std)
+        x = self.chan_complex_to_last_dim(x)
+
+        return x
+
+
+class SensitivityModel(nn.Module):
+    """
+    Model for learning sensitivity estimation from k-space data.
+
+    This model applies an IFFT to multichannel k-space data and then a U-Net
+    to the coil images to estimate coil sensitivities. It can be used with the
+    end-to-end variational network.
+
+    Input: multi-coil k-space data
+    Output: multi-coil spatial domain sensitivity maps
+    """
+
+    def __init__(
+        self,
+        chans: int,
+        num_pools: int,
+        in_chans: int = 2,
+        out_chans: int = 2,
+        drop_prob: float = 0.0,
+        mask_center: bool = True,
+    ):
+        """
+        Parameters
+        ----------
+        chans : int
+            Number of output channels of the first convolution layer.
+        num_pools : int
+            Number of down-sampling and up-sampling layers.
+        in_chans : int, optional
+            Number of channels in the input to the U-Net model. Default is 2.
+        out_chans : int, optional
+            Number of channels in the output to the U-Net model. Default is 2.
+        drop_prob : float, optional
+            Dropout probability. Default is 0.0.
+        mask_center : bool, optional
+            Whether to mask center of k-space for sensitivity map calculation.
+            Default is True.
+        """
+        super().__init__()
+        self.mask_center = mask_center
+        self.norm_unet = NormUnet(
+            chans,
+            num_pools,
+            in_chans=in_chans,
+            out_chans=out_chans,
+            drop_prob=drop_prob,
+        )
+
+    def chans_to_batch_dim(self, x: torch.Tensor) -> Tuple[torch.Tensor, int]:
+        b, c, h, w, comp = x.shape
+
+        return x.view(b * c, 1, h, w, comp), b
+
+    def batch_chans_to_chan_dim(
+        self,
+        x: torch.Tensor,
+        batch_size: int,
+    ) -> torch.Tensor:
+        bc, _, h, w, comp = x.shape
+        c = bc // batch_size
+
+        return x.view(batch_size, c, h, w, comp)
+
+    def divide_root_sum_of_squares(self, x: torch.Tensor) -> torch.Tensor:
+        return x / fastmri.rss_complex(x, dim=1).unsqueeze(-1).unsqueeze(1)
+
+    def get_pad_and_num_low_freqs(
+        self, mask: torch.Tensor, num_low_frequencies: Optional[int] = None
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if num_low_frequencies is None or any(
+            torch.any(t == 0) for t in num_low_frequencies
+        ):
+            # get low frequency line locations and mask them out
+            squeezed_mask = mask[:, 0, 0, :, 0].to(torch.int8)
+            cent = squeezed_mask.shape[1] // 2
+            # running argmin returns the first non-zero
+            left = torch.argmin(squeezed_mask[:, :cent].flip(1), dim=1)
+            right = torch.argmin(squeezed_mask[:, cent:], dim=1)
+            num_low_frequencies_tensor = torch.max(
+                2 * torch.min(left, right), torch.ones_like(left)
+            )  # force a symmetric center unless 1
+        else:
+            num_low_frequencies_tensor = num_low_frequencies * torch.ones(
+                mask.shape[0], dtype=mask.dtype, device=mask.device
+            )
+
+        pad = (mask.shape[-2] - num_low_frequencies_tensor + 1) // 2
+
+        return pad.type(torch.long), num_low_frequencies_tensor.type(torch.long)
+
+    def forward(
+        self,
+        masked_kspace: torch.Tensor,
+        mask: torch.Tensor,
+        num_low_frequencies: Optional[int] = None,
+    ) -> torch.Tensor:
+        if self.mask_center:
+            pad, num_low_freqs = self.get_pad_and_num_low_freqs(
+                mask, num_low_frequencies
+            )
+            masked_kspace = transforms.batched_mask_center(
+                masked_kspace, pad, pad + num_low_freqs
+            )
+
+        # convert to image space
+        images, batches = self.chans_to_batch_dim(fastmri.ifft2c(masked_kspace))
+
+        # estimate sensitivities
+        return self.divide_root_sum_of_squares(
+            self.batch_chans_to_chan_dim(self.norm_unet(images), batches)
+        )
+
+
+class VarNet(nn.Module):
+    """
+    A full variational network model.
+
+    This model applies a combination of soft data consistency with a U-Net
+    regularizer. To use non-U-Net regularizers, use VarNetBlock.
+
+    Input: multi-channel k-space data
+    Output: single-channel RSS reconstructed image
+    """
+
+    def __init__(
+        self,
+        num_cascades: int = 12,
+        sens_chans: int = 8,
+        sens_pools: int = 4,
+        chans: int = 18,
+        pools: int = 4,
+        mask_center: bool = True,
+    ):
+        """
+        Parameters
+        ----------
+        num_cascades : int
+            Number of cascades (i.e., layers) for variational network.
+        sens_chans : int
+            Number of channels for sensitivity map U-Net.
+        sens_pools : int
+            Number of downsampling and upsampling layers for sensitivity map U-Net.
+        chans : int
+            Number of channels for cascade U-Net.
+        pools : int
+            Number of downsampling and upsampling layers for cascade U-Net.
+        mask_center : bool
+            Whether to mask center of k-space for sensitivity map calculation.
+        """
+
+        super().__init__()
+
+        self.sens_net = SensitivityModel(
+            chans=sens_chans,
+            num_pools=sens_pools,
+            mask_center=mask_center,
+        )
+        self.cascades = nn.ModuleList(
+            [VarNetBlock(NormUnet(chans, pools)) for _ in range(num_cascades)]
+        )
+
+    def forward(
+        self,
+        masked_kspace: torch.Tensor,
+        mask: torch.Tensor,
+        num_low_frequencies: Optional[int] = None,
+    ) -> torch.Tensor:
+        sens_maps = self.sens_net(masked_kspace, mask, num_low_frequencies)
+        kspace_pred = masked_kspace.clone()
+        for cascade in self.cascades:
+            kspace_pred = cascade(kspace_pred, masked_kspace, mask, sens_maps)
+
+        spatial_pred = fastmri.ifft2c(kspace_pred)
+
+        # ---------> FIXME: CHANGE FOR MVUE MODE
+        if self.training and os.getenv("MVUE") in ["yes", "1", "true", "True"]:
+            combined_spatial = fastmri.mvue(spatial_pred, sens_maps, dim=1)
+        else:
+            spatial_pred_abs = fastmri.complex_abs(spatial_pred)
+            combined_spatial = fastmri.rss(spatial_pred_abs, dim=1)
+        return combined_spatial
+
+
+class VarNetBlock(nn.Module):
+    """
+    Model block for end-to-end variational network (refinemnt module)
+
+    This model applies a combination of soft data consistency with the input
+    model as a regularizer. A series of these blocks can be stacked to form
+    the full variational network.
+
+    Input: multi-channel k-space data
+    Output: multi-channel k-space data
+    """
+
+    def __init__(self, model: nn.Module):
+        """
+        Parameters
+        ----------
+        model : nn.Module
+            Module for "regularization" component of variational network.
+        """
+        super().__init__()
+
+        self.model = model
+        self.dc_weight = nn.Parameter(torch.ones(1))
+
+    def sens_expand(
+        self, x: torch.Tensor, sens_maps: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Calculates F (x sens_maps)
+        """
+        return fastmri.fft2c(fastmri.complex_mul(x, sens_maps))
+
+    def sens_reduce(
+        self, x: torch.Tensor, sens_maps: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Calculates F^{-1}(x) \overline{sens_maps}
+        where \overline{sens_maps} is the element-wise applied complex conjugate
+        """
+        return fastmri.complex_mul(
+            fastmri.ifft2c(x), fastmri.complex_conj(sens_maps)
+        ).sum(dim=1, keepdim=True)
+
+    def forward(
+        self,
+        current_kspace: torch.Tensor,
+        ref_kspace: torch.Tensor,
+        mask: torch.Tensor,
+        sens_maps: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Parameters
+        ----------
+        current_kspace : torch.Tensor
+            The current k-space data (frequency domain data) being processed by the network.
+        ref_kspace : torch.Tensor
+            The reference k-space data (measured data) used for data consistency.
+        mask : torch.Tensor
+            A binary mask indicating the locations in k-space where data consistency should be enforced.
+        sens_maps : torch.Tensor
+            Sensitivity maps for the different coils in parallel imaging.
+
+        Returns
+        -------
+        torch.Tensor
+            The output k-space data after applying the variational network block.
+        """
+
+        """
+        Model term:
+            - Reduces the current k-space data using the sensitivity maps (inverse Fourier transform followed by element-wise multiplication and summation).
+            - Applies the neural network model to the reduced data.
+            - Expands the output of the model using the sensitivity maps (element-wise multiplication followed by Fourier transform).
+        """
+
+        model_term = self.sens_expand(
+            self.model(self.sens_reduce(current_kspace, sens_maps)), sens_maps
+        )
+
+        """
+        Soft data consistency term:
+            - Calculates the difference between current k-space and reference k-space where the mask is true.
+            - Multiplies this difference by the data consistency weight.
+        """
+        zero = torch.zeros(1, 1, 1, 1, 1).to(current_kspace)
+        soft_dc = (
+            torch.where(mask, current_kspace - ref_kspace, zero)
+            * self.dc_weight
+        )
+
+        # with data consistency term (removed for single cascade experiments)
+        return current_kspace - soft_dc - model_term

pyproject.toml

@@ -0,0 +1,46 @@
+[build-system]
+requires = ["setuptools>=64.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "no-med"
+version = "0.0.0"
+description = "Neural operators for medical imaging."
+readme = "README.md"
+# readme-content-type = "text/markdown"
+authors = [{ name = "Armeet Singh Jatyani", email = "[email protected]" }]
+license = { text = "MIT" }
+keywords = ["medical imaging", "neural operators", "AI"]
+classifiers = [
+  "Programming Language :: Python :: 3",
+  "License :: OSI Approved :: MIT License",
+  "Operating System :: OS Independent",
+]
+requires-python = ">=3.6"
+
+[project.optional-dependencies]
+dev = ["pytest>=6.0", "black", "flake8", "pdoc3"]
+
+[tool.setuptools]
+include-package-data = true
+packages = { find = {} }    # auto find packages
+
+[tool.black]
+line-length = 80 # Default is 88, but you can set it to 100 or 120 if needed
+target-version = ['py38', 'py39', 'py310'] # Set target Python versions
+include = '\.pyi?$' # Format only .py and .pyi files
+skip-string-normalization = true
+exclude = '''
+/(
+    \.eggs         # Exclude files generated by packaging
+  | \.git          # Exclude version control files
+  | \.mypy_cache   # Exclude mypy caches
+  | \.tox          # Exclude tox environments
+  | \.venv         # Exclude virtual environments
+)/
+'''
+fast = true
+
+[tool.isort]
+profile = "black"
+line_length = 80

pytest.ini

@@ -0,0 +1,4 @@
+[pytest]
+markers =
+    train: marks training tests (long runtime) (deselect with '-m "not train"')
+addopts = -m "not train"

setup_config.py

@@ -0,0 +1,21 @@
+import os
+
+config_filename = "fastmri.yaml"
+default_config_content = """brain_path: brain path here
+knee_path: knee path here
+log_path: log path here
+checkpoint_path: checkpoint path here
+"""
+
+
+def check_and_create_config():
+    if not os.path.exists(config_filename):
+        print(f"{config_filename} not found. Creating with default template...")
+        with open(config_filename, "w") as config_file:
+            config_file.write(default_config_content)
+        print(f"Default configuration file created at {config_filename}.")
+    else:
+        print(f"{config_filename} already exists. No changes made.")
+
+
+check_and_create_config()

torch_harmonics_local/_disco_convolution.py

@@ -0,0 +1,502 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+
+import math
+
+import torch
+
+# triton will only be avaiable on cuda installations of pytorch
+import triton
+import triton.language as tl
+
+BLOCK_SIZE_BATCH = 4
+BLOCK_SIZE_NZ = 8
+BLOCK_SIZE_POUT = 8
+
+
+@triton.jit
+def _disco_s2_contraction_kernel(
+    inz_ptr,
+    vnz_ptr,
+    nnz,
+    inz_stride_ii,
+    inz_stride_nz,
+    vnz_stride,
+    x_ptr,
+    batch_size,
+    nlat_in,
+    nlon_in,
+    x_stride_b,
+    x_stride_t,
+    x_stride_p,
+    y_ptr,
+    kernel_size,
+    nlat_out,
+    nlon_out,
+    y_stride_b,
+    y_stride_f,
+    y_stride_t,
+    y_stride_p,
+    pscale,
+    backward: tl.constexpr,
+    BLOCK_SIZE_BATCH: tl.constexpr,
+    BLOCK_SIZE_NZ: tl.constexpr,
+    BLOCK_SIZE_POUT: tl.constexpr,
+):
+    """
+    Kernel for the sparse-dense contraction for the S2 DISCO convolution.
+    """
+
+    pid_batch = tl.program_id(0)
+    pid_pout = tl.program_id(2)
+
+    # pid_nz should always be 0 as we do not account for larger grids in this dimension
+    pid_nz = tl.program_id(1)  # should be always 0
+    tl.device_assert(pid_nz == 0)
+
+    # create the pointer block for pout
+    pout = pid_pout * BLOCK_SIZE_POUT + tl.arange(0, BLOCK_SIZE_POUT)
+    b = pid_batch * BLOCK_SIZE_BATCH + tl.arange(0, BLOCK_SIZE_BATCH)
+
+    # create pointer blocks for the psi datastructure
+    iinz = tl.arange(0, BLOCK_SIZE_NZ)
+
+    # get the initial pointers
+    fout_ptrs = inz_ptr + iinz * inz_stride_nz
+    tout_ptrs = inz_ptr + iinz * inz_stride_nz + inz_stride_ii
+    tpnz_ptrs = inz_ptr + iinz * inz_stride_nz + 2 * inz_stride_ii
+    vals_ptrs = vnz_ptr + iinz * vnz_stride
+
+    # iterate in a blocked fashion over the non-zero entries
+    for offs_nz in range(0, nnz, BLOCK_SIZE_NZ):
+        # load input output latitude coordinate pairs
+        fout = tl.load(
+            fout_ptrs + offs_nz * inz_stride_nz, mask=(offs_nz + iinz < nnz), other=-1
+        )
+        tout = tl.load(
+            tout_ptrs + offs_nz * inz_stride_nz, mask=(offs_nz + iinz < nnz), other=-1
+        )
+        tpnz = tl.load(
+            tpnz_ptrs + offs_nz * inz_stride_nz, mask=(offs_nz + iinz < nnz), other=-1
+        )
+
+        # load corresponding values
+        vals = tl.load(
+            vals_ptrs + offs_nz * vnz_stride, mask=(offs_nz + iinz < nnz), other=0.0
+        )
+
+        # compute the shifted longitude coordinates p+p' to read in a coalesced fashion
+        tnz = tpnz // nlon_in
+        pnz = tpnz % nlon_in
+
+        # make sure the value is not out of bounds
+        tl.device_assert(fout < kernel_size)
+        tl.device_assert(tout < nlat_out)
+        tl.device_assert(tnz < nlat_in)
+        tl.device_assert(pnz < nlon_in)
+
+        # load corresponding portion of the input array
+        x_ptrs = (
+            x_ptr
+            + tnz[None, :, None] * x_stride_t
+            + ((pnz[None, :, None] + pout[None, None, :] * pscale) % nlon_in)
+            * x_stride_p
+            + b[:, None, None] * x_stride_b
+        )
+        y_ptrs = (
+            y_ptr
+            + fout[None, :, None] * y_stride_f
+            + tout[None, :, None] * y_stride_t
+            + (pout[None, None, :] % nlon_out) * y_stride_p
+            + b[:, None, None] * y_stride_b
+        )
+
+        # precompute the mask
+        mask = (
+            (b[:, None, None] < batch_size) and (offs_nz + iinz[None, :, None] < nnz)
+        ) and (pout[None, None, :] < nlon_out)
+
+        # do the actual computation. Backward is essentially just the same operation with swapped tensors.
+        if not backward:
+            x = tl.load(x_ptrs, mask=mask, other=0.0)
+            y = vals[None, :, None] * x
+
+            # store it to the output array
+            tl.atomic_add(y_ptrs, y, mask=mask)
+        else:
+            y = tl.load(y_ptrs, mask=mask, other=0.0)
+            x = vals[None, :, None] * y
+
+            # store it to the output array
+            tl.atomic_add(x_ptrs, x, mask=mask)
+
+
+def _disco_s2_contraction_fwd(x: torch.Tensor, psi: torch.Tensor, nlon_out: int):
+    """
+    Wrapper function for the triton implementation of the efficient DISCO convolution on the sphere.
+
+    Parameters
+    ----------
+    x: torch.Tensor
+        Input signal on the sphere. Expects a tensor of shape batch_size x channels x nlat_in x nlon_in).
+    psi : torch.Tensor
+        Pre-computed convolution tensor. Expects a sparse tensor of shape kernel_size x nlat_out x (nlat_in * nlon_in).
+    nlon_out: int
+        Number of longitude points the output should have.
+    """
+
+    # check the shapes of all input tensors
+    assert len(psi.shape) == 3
+    assert len(x.shape) == 4
+    assert psi.is_sparse, "Psi must be a sparse COO tensor"
+
+    # TODO: check that Psi is also coalesced
+
+    # get the dimensions of the problem
+    kernel_size, nlat_out, n_in = psi.shape
+    nnz = psi.indices().shape[-1]
+    batch_size, n_chans, nlat_in, nlon_in = x.shape
+    assert nlat_in * nlon_in == n_in
+
+    # TODO: check that Psi index vector is of type long
+
+    # make sure that the grid-points of the output grid fall onto the grid points of the input grid
+    assert nlon_in % nlon_out == 0
+    pscale = nlon_in // nlon_out
+
+    # to simplify things, we merge batch and channel dimensions
+    x = x.reshape(batch_size * n_chans, nlat_in, nlon_in)
+
+    # prepare the output tensor
+    y = torch.zeros(
+        batch_size * n_chans,
+        kernel_size,
+        nlat_out,
+        nlon_out,
+        device=x.device,
+        dtype=x.dtype,
+    )
+
+    # determine the grid for the computation
+    grid = (
+        triton.cdiv(batch_size * n_chans, BLOCK_SIZE_BATCH),
+        1,
+        triton.cdiv(nlon_out, BLOCK_SIZE_POUT),
+    )
+
+    # launch the kernel
+    _disco_s2_contraction_kernel[grid](
+        psi.indices(),
+        psi.values(),
+        nnz,
+        psi.indices().stride(-2),
+        psi.indices().stride(-1),
+        psi.values().stride(-1),
+        x,
+        batch_size * n_chans,
+        nlat_in,
+        nlon_in,
+        x.stride(0),
+        x.stride(-2),
+        x.stride(-1),
+        y,
+        kernel_size,
+        nlat_out,
+        nlon_out,
+        y.stride(0),
+        y.stride(1),
+        y.stride(-2),
+        y.stride(-1),
+        pscale,
+        False,
+        BLOCK_SIZE_BATCH,
+        BLOCK_SIZE_NZ,
+        BLOCK_SIZE_POUT,
+    )
+
+    # reshape y back to expose the correct dimensions
+    y = y.reshape(batch_size, n_chans, kernel_size, nlat_out, nlon_out)
+
+    return y
+
+
+def _disco_s2_contraction_bwd(grad_y: torch.Tensor, psi: torch.Tensor, nlon_in: int):
+    """
+    Backward pass for the triton implementation of the efficient DISCO convolution on the sphere.
+
+    Parameters
+    ----------
+    grad_y: torch.Tensor
+        Input gradient on the sphere. Expects a tensor of shape batch_size x channels x kernel_size x nlat_out x nlon_out.
+    psi : torch.Tensor
+        Pre-computed convolution tensor. Expects a sparse tensor of shape kernel_size x nlat_out x (nlat_in * nlon_in).
+    nlon_in: int
+        Number of longitude points the input used. Is required to infer the correct dimensions
+    """
+
+    # check the shapes of all input tensors
+    assert len(psi.shape) == 3
+    assert len(grad_y.shape) == 5
+    assert psi.is_sparse, "psi must be a sparse COO tensor"
+
+    # TODO: check that Psi is also coalesced
+
+    # get the dimensions of the problem
+    kernel_size, nlat_out, n_in = psi.shape
+    nnz = psi.indices().shape[-1]
+    assert grad_y.shape[-2] == nlat_out
+    assert grad_y.shape[-3] == kernel_size
+    assert n_in % nlon_in == 0
+    nlat_in = n_in // nlon_in
+    batch_size, n_chans, _, _, nlon_out = grad_y.shape
+
+    # make sure that the grid-points of the output grid fall onto the grid points of the input grid
+    assert nlon_in % nlon_out == 0
+    pscale = nlon_in // nlon_out
+
+    # to simplify things, we merge batch and channel dimensions
+    grad_y = grad_y.reshape(batch_size * n_chans, kernel_size, nlat_out, nlon_out)
+
+    # prepare the output tensor
+    grad_x = torch.zeros(
+        batch_size * n_chans, nlat_in, nlon_in, device=grad_y.device, dtype=grad_y.dtype
+    )
+
+    # determine the grid for the computation
+    grid = (
+        triton.cdiv(batch_size * n_chans, BLOCK_SIZE_BATCH),
+        1,
+        triton.cdiv(nlon_out, BLOCK_SIZE_POUT),
+    )
+
+    # launch the kernel
+    _disco_s2_contraction_kernel[grid](
+        psi.indices(),
+        psi.values(),
+        nnz,
+        psi.indices().stride(-2),
+        psi.indices().stride(-1),
+        psi.values().stride(-1),
+        grad_x,
+        batch_size * n_chans,
+        nlat_in,
+        nlon_in,
+        grad_x.stride(0),
+        grad_x.stride(-2),
+        grad_x.stride(-1),
+        grad_y,
+        kernel_size,
+        nlat_out,
+        nlon_out,
+        grad_y.stride(0),
+        grad_y.stride(1),
+        grad_y.stride(-2),
+        grad_y.stride(-1),
+        pscale,
+        True,
+        BLOCK_SIZE_BATCH,
+        BLOCK_SIZE_NZ,
+        BLOCK_SIZE_POUT,
+    )
+
+    # reshape y back to expose the correct dimensions
+    grad_x = grad_x.reshape(batch_size, n_chans, nlat_in, nlon_in)
+
+    return grad_x
+
+
+class _DiscoS2ContractionTriton(torch.autograd.Function):
+    """
+    Helper function to make the triton implementation work with PyTorch autograd functionality
+    """
+
+    @staticmethod
+    def forward(ctx, x: torch.Tensor, psi: torch.Tensor, nlon_out: int):
+        ctx.save_for_backward(psi)
+        ctx.nlon_in = x.shape[-1]
+
+        return _disco_s2_contraction_fwd(x, psi, nlon_out)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        (psi,) = ctx.saved_tensors
+        grad_input = _disco_s2_contraction_bwd(grad_output, psi, ctx.nlon_in)
+        grad_x = grad_psi = None
+
+        return grad_input, None, None
+
+
+class _DiscoS2TransposeContractionTriton(torch.autograd.Function):
+    """
+    Helper function to make the triton implementation work with PyTorch autograd functionality
+    """
+
+    @staticmethod
+    def forward(ctx, x: torch.Tensor, psi: torch.Tensor, nlon_out: int):
+        ctx.save_for_backward(psi)
+        ctx.nlon_in = x.shape[-1]
+
+        return _disco_s2_contraction_bwd(x, psi, nlon_out)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        (psi,) = ctx.saved_tensors
+        grad_input = _disco_s2_contraction_fwd(grad_output, psi, ctx.nlon_in)
+        grad_x = grad_psi = None
+
+        return grad_input, None, None
+
+
+def _disco_s2_contraction_triton(x: torch.Tensor, psi: torch.Tensor, nlon_out: int):
+    return _DiscoS2ContractionTriton.apply(x, psi, nlon_out)
+
+
+def _disco_s2_transpose_contraction_triton(
+    x: torch.Tensor, psi: torch.Tensor, nlon_out: int
+):
+    return _DiscoS2TransposeContractionTriton.apply(x, psi, nlon_out)
+
+
+def _disco_s2_contraction_torch(x: torch.Tensor, psi: torch.Tensor, nlon_out: int):
+    """
+    Reference implementation of the custom contraction as described in [1]. This requires repeated
+    shifting of the input tensor, which can potentially be costly. For an efficient implementation
+    on GPU, make sure to use the custom kernel written in Triton.
+    """
+    assert len(psi.shape) == 3
+    assert len(x.shape) == 4
+    psi = psi.to(x.device)
+
+    batch_size, n_chans, nlat_in, nlon_in = x.shape
+    kernel_size, nlat_out, _ = psi.shape
+
+    assert psi.shape[-1] == nlat_in * nlon_in
+    assert nlon_in % nlon_out == 0
+    assert nlon_in >= nlat_out
+    pscale = nlon_in // nlon_out
+
+    # add a dummy dimension for nkernel and move the batch and channel dims to the end
+    x = x.reshape(1, batch_size * n_chans, nlat_in, nlon_in).permute(0, 2, 3, 1)
+    x = x.expand(kernel_size, -1, -1, -1)
+
+    y = torch.zeros(
+        nlon_out,
+        kernel_size,
+        nlat_out,
+        batch_size * n_chans,
+        device=x.device,
+        dtype=x.dtype,
+    )
+
+    for pout in range(nlon_out):
+        # sparse contraction with psi
+        y[pout] = torch.bmm(psi, x.reshape(kernel_size, nlat_in * nlon_in, -1))
+        # we need to repeatedly roll the input tensor to faciliate the shifted multiplication
+        x = torch.roll(x, -pscale, dims=2)
+
+    # reshape y back to expose the correct dimensions
+    y = y.permute(3, 1, 2, 0).reshape(
+        batch_size, n_chans, kernel_size, nlat_out, nlon_out
+    )
+
+    return y
+
+
+def _disco_s2_transpose_contraction_torch(
+    x: torch.Tensor, psi: torch.Tensor, nlon_out: int
+):
+    """
+    Reference implementation of the custom contraction as described in [1]. This requires repeated
+    shifting of the input tensor, which can potentially be costly. For an efficient implementation
+    on GPU, make sure to use the custom kernel written in Triton.
+    """
+    assert len(psi.shape) == 3
+    assert len(x.shape) == 5
+    psi = psi.to(x.device)
+
+    batch_size, n_chans, kernel_size, nlat_in, nlon_in = x.shape
+    kernel_size, _, n_out = psi.shape
+
+    assert psi.shape[-2] == nlat_in
+    assert n_out % nlon_out == 0
+    nlat_out = n_out // nlon_out
+    assert nlon_out >= nlat_in
+    pscale = nlon_out // nlon_in
+
+    # we do a semi-transposition to faciliate the computation
+    inz = psi.indices()
+    tout = inz[2] // nlon_out
+    pout = inz[2] % nlon_out
+    # flip the axis of longitudes
+    pout = nlon_out - 1 - pout
+    tin = inz[1]
+    inz = torch.stack([inz[0], tout, tin * nlon_out + pout], dim=0)
+    psi_mod = torch.sparse_coo_tensor(
+        inz, psi.values(), size=(kernel_size, nlat_out, nlat_in * nlon_out)
+    )
+
+    # interleave zeros along the longitude dimension to allow for fractional offsets to be considered
+    x_ext = torch.zeros(
+        kernel_size,
+        nlat_in,
+        nlon_out,
+        batch_size * n_chans,
+        device=x.device,
+        dtype=x.dtype,
+    )
+    x_ext[:, :, ::pscale, :] = x.reshape(
+        batch_size * n_chans, kernel_size, nlat_in, nlon_in
+    ).permute(1, 2, 3, 0)
+    # we need to go backwards through the vector, so we flip the axis
+    x_ext = x_ext.contiguous()
+
+    y = torch.zeros(
+        kernel_size,
+        nlon_out,
+        nlat_out,
+        batch_size * n_chans,
+        device=x.device,
+        dtype=x.dtype,
+    )
+
+    for pout in range(nlon_out):
+        # we need to repeatedly roll the input tensor to faciliate the shifted multiplication
+        # TODO: double-check why this has to happen first
+        x_ext = torch.roll(x_ext, -1, dims=2)
+        # sparse contraction with the modified psi
+        y[:, pout, :, :] = torch.bmm(
+            psi_mod, x_ext.reshape(kernel_size, nlat_in * nlon_out, -1)
+        )
+
+    # sum over the kernel dimension and reshape to the correct output size
+    y = y.sum(dim=0).permute(2, 1, 0).reshape(batch_size, n_chans, nlat_out, nlon_out)
+
+    return y

torch_harmonics_local/convolution.py

@@ -0,0 +1,1014 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+
+import abc
+import math
+from functools import partial
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .quadrature import _precompute_grid, _precompute_latitudes
+
+if torch.cuda.is_available():
+    from ._disco_convolution import (
+        _disco_s2_contraction_triton,
+        _disco_s2_transpose_contraction_triton,
+    )
+
+
+def _compute_support_vals_isotropic(
+    r: torch.Tensor, phi: torch.Tensor, nr: int, r_cutoff: float, norm: str = "s2"
+):
+    """
+    Computes the index set that falls into the isotropic kernel's support and returns both indices and values.
+    """
+
+    # compute the support
+    dr = (r_cutoff - 0.0) / nr
+    ikernel = torch.arange(nr).reshape(-1, 1, 1)
+    ir = ikernel * dr
+
+    if norm == "none":
+        norm_factor = 1.0
+    elif norm == "2d":
+        norm_factor = (
+            math.pi * (r_cutoff * nr / (nr + 1)) ** 2
+            + math.pi * r_cutoff**2 * (2 * nr / (nr + 1) + 1) / (nr + 1) / 3
+        )
+    elif norm == "s2":
+        norm_factor = (
+            2
+            * math.pi
+            * (
+                1
+                - math.cos(r_cutoff - dr)
+                + math.cos(r_cutoff - dr)
+                + (math.sin(r_cutoff - dr) - math.sin(r_cutoff)) / dr
+            )
+        )
+    else:
+        raise ValueError(f"Unknown normalization mode {norm}.")
+
+    # find the indices where the rotated position falls into the support of the kernel
+    iidx = torch.argwhere(((r - ir).abs() <= dr) & (r <= r_cutoff))
+    vals = (
+        1 - (r[iidx[:, 1], iidx[:, 2]] - ir[iidx[:, 0], 0, 0]).abs() / dr
+    ) / norm_factor
+    return iidx, vals
+
+
+def _compute_support_vals_anisotropic(
+    r: torch.Tensor,
+    phi: torch.Tensor,
+    nr: int,
+    nphi: int,
+    r_cutoff: float,
+    norm: str = "s2",
+):
+    """
+    Computes the index set that falls into the anisotropic kernel's support and returns both indices and values.
+    """
+
+    # compute the support
+    dr = (r_cutoff - 0.0) / nr
+    dphi = 2.0 * math.pi / nphi
+    kernel_size = (nr - 1) * nphi + 1
+    ikernel = torch.arange(kernel_size).reshape(-1, 1, 1)
+    ir = ((ikernel - 1) // nphi + 1) * dr
+    iphi = ((ikernel - 1) % nphi) * dphi
+
+    if norm == "none":
+        norm_factor = 1.0
+    elif norm == "2d":
+        norm_factor = (
+            math.pi * (r_cutoff * nr / (nr + 1)) ** 2
+            + math.pi * r_cutoff**2 * (2 * nr / (nr + 1) + 1) / (nr + 1) / 3
+        )
+    elif norm == "s2":
+        norm_factor = (
+            2
+            * math.pi
+            * (
+                1
+                - math.cos(r_cutoff - dr)
+                + math.cos(r_cutoff - dr)
+                + (math.sin(r_cutoff - dr) - math.sin(r_cutoff)) / dr
+            )
+        )
+    else:
+        raise ValueError(f"Unknown normalization mode {norm}.")
+
+    # find the indices where the rotated position falls into the support of the kernel
+    cond_r = ((r - ir).abs() <= dr) & (r <= r_cutoff)
+    cond_phi = (
+        (ikernel == 0)
+        | ((phi - iphi).abs() <= dphi)
+        | ((2 * math.pi - (phi - iphi).abs()) <= dphi)
+    )
+    iidx = torch.argwhere(cond_r & cond_phi)
+    vals = (
+        1 - (r[iidx[:, 1], iidx[:, 2]] - ir[iidx[:, 0], 0, 0]).abs() / dr
+    ) / norm_factor
+    vals *= torch.where(
+        iidx[:, 0] > 0,
+        (
+            1
+            - torch.minimum(
+                (phi[iidx[:, 1], iidx[:, 2]] - iphi[iidx[:, 0], 0, 0]).abs(),
+                (
+                    2 * math.pi
+                    - (phi[iidx[:, 1], iidx[:, 2]] - iphi[iidx[:, 0], 0, 0]).abs()
+                ),
+            )
+            / dphi
+        ),
+        1.0,
+    )
+    return iidx, vals
+
+
+def _precompute_convolution_tensor_s2(
+    in_shape,
+    out_shape,
+    kernel_shape,
+    grid_in="equiangular",
+    grid_out="equiangular",
+    theta_cutoff=0.01 * math.pi,
+):
+    """
+    Precomputes the rotated filters at positions $R^{-1}_j \omega_i = R^{-1}_j R_i \nu = Y(-\theta_j)Z(\phi_i - \phi_j)Y(\theta_j)\nu$.
+    Assumes a tensorized grid on the sphere with an equidistant sampling in longitude as described in Ocampo et al.
+    The output tensor has shape kernel_shape x nlat_out x (nlat_in * nlon_in).
+
+    The rotation of the Euler angles uses the YZY convention, which applied to the northpole $(0,0,1)^T$ yields
+    $$
+    Y(\alpha) Z(\beta) Y(\gamma) n =
+        {\begin{bmatrix}
+            \cos(\gamma)\sin(\alpha) + \cos(\alpha)\cos(\beta)\sin(\gamma) \\
+            \sin(\beta)\sin(\gamma) \\
+            \cos(\alpha)\cos(\gamma)-\cos(\beta)\sin(\alpha)\sin(\gamma)
+        \end{bmatrix}}
+    $$
+    """
+
+    assert len(in_shape) == 2
+    assert len(out_shape) == 2
+
+    if len(kernel_shape) == 1:
+        kernel_handle = partial(
+            _compute_support_vals_isotropic,
+            nr=kernel_shape[0],
+            r_cutoff=theta_cutoff,
+            norm="s2",
+        )
+    elif len(kernel_shape) == 2:
+        kernel_handle = partial(
+            _compute_support_vals_anisotropic,
+            nr=kernel_shape[0],
+            nphi=kernel_shape[1],
+            r_cutoff=theta_cutoff,
+            norm="s2",
+        )
+    else:
+        raise ValueError("kernel_shape should be either one- or two-dimensional.")
+
+    nlat_in, nlon_in = in_shape
+    nlat_out, nlon_out = out_shape
+
+    lats_in, _ = _precompute_latitudes(nlat_in, grid=grid_in)
+    lats_in = torch.from_numpy(lats_in).float()
+    lats_out, _ = _precompute_latitudes(nlat_out, grid=grid_out)
+    lats_out = torch.from_numpy(lats_out).float()
+
+    # array for accumulating non-zero indices
+    out_idx = torch.empty([3, 0], dtype=torch.long)
+    out_vals = torch.empty([0], dtype=torch.long)
+
+    # compute the phi differences
+    # It's imporatant to not include the 2 pi point in the longitudes, as it is equivalent to lon=0
+    lons_in = torch.linspace(0, 2 * math.pi, nlon_in + 1)[:-1]
+
+    for t in range(nlat_out):
+        # the last angle has a negative sign as it is a passive rotation, which rotates the filter around the y-axis
+        alpha = -lats_out[t]
+        beta = lons_in
+        gamma = lats_in.reshape(-1, 1)
+
+        # compute cartesian coordinates of the rotated position
+        # This uses the YZY convention of Euler angles, where the last angle (alpha) is a passive rotation,
+        # and therefore applied with a negative sign
+        z = -torch.cos(beta) * torch.sin(alpha) * torch.sin(gamma) + torch.cos(
+            alpha
+        ) * torch.cos(gamma)
+        x = torch.cos(alpha) * torch.cos(beta) * torch.sin(gamma) + torch.cos(
+            gamma
+        ) * torch.sin(alpha)
+        y = torch.sin(beta) * torch.sin(gamma)
+
+        # normalization is emportant to avoid NaNs when arccos and atan are applied
+        # this can otherwise lead to spurious artifacts in the solution
+        norm = torch.sqrt(x * x + y * y + z * z)
+        x = x / norm
+        y = y / norm
+        z = z / norm
+
+        # compute spherical coordinates, where phi needs to fall into the [0, 2pi) range
+        theta = torch.arccos(z)
+        phi = torch.arctan2(y, x) + torch.pi
+
+        # find the indices where the rotated position falls into the support of the kernel
+        iidx, vals = kernel_handle(theta, phi)
+
+        # add the output latitude and reshape such that psi has dimensions kernel_shape x nlat_out x (nlat_in*nlon_in)
+        idx = torch.stack(
+            [
+                iidx[:, 0],
+                t * torch.ones_like(iidx[:, 0]),
+                iidx[:, 1] * nlon_in + iidx[:, 2],
+            ],
+            dim=0,
+        )
+
+        # append indices and values to the COO datastructure
+        out_idx = torch.cat([out_idx, idx], dim=-1)
+        out_vals = torch.cat([out_vals, vals], dim=-1)
+
+    return out_idx, out_vals
+
+
+def _precompute_convolution_tensor_2d(
+    grid_in, grid_out, kernel_shape, radius_cutoff=0.01, periodic=False
+):
+    """
+    Precomputes the translated filters at positions $T^{-1}_j \omega_i = T^{-1}_j T_i \nu$. Similar to the S2 routine,
+    only that it assumes a non-periodic subset of the euclidean plane
+    """
+
+    # check that input arrays are valid point clouds in 2D
+    assert len(grid_in) == 2
+    assert len(grid_out) == 2
+    assert grid_in.shape[0] == 2
+    assert grid_out.shape[0] == 2
+
+    n_in = grid_in.shape[-1]
+    n_out = grid_out.shape[-1]
+
+    if len(kernel_shape) == 1:
+        kernel_handle = partial(
+            _compute_support_vals_isotropic,
+            nr=kernel_shape[0],
+            r_cutoff=radius_cutoff,
+            norm="2d",
+        )
+    elif len(kernel_shape) == 2:
+        kernel_handle = partial(
+            _compute_support_vals_anisotropic,
+            nr=kernel_shape[0],
+            nphi=kernel_shape[1],
+            r_cutoff=radius_cutoff,
+            norm="2d",
+        )
+    else:
+        raise ValueError("kernel_shape should be either one- or two-dimensional.")
+
+    grid_in = grid_in.reshape(2, 1, n_in)
+    grid_out = grid_out.reshape(2, n_out, 1)
+
+    diffs = grid_in - grid_out
+    if periodic:
+        periodic_diffs = torch.where(diffs > 0.0, diffs - 1, diffs + 1)
+        diffs = torch.where(diffs.abs() < periodic_diffs.abs(), diffs, periodic_diffs)
+
+    r = torch.sqrt(diffs[0] ** 2 + diffs[1] ** 2)
+    phi = torch.arctan2(diffs[1], diffs[0]) + torch.pi
+
+    idx, vals = kernel_handle(r, phi)
+    idx = idx.permute(1, 0)
+
+    return idx, vals
+
+
+class DiscreteContinuousConv(nn.Module, abc.ABC):
+    """
+    Abstract base class for DISCO convolutions
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_shape: Union[int, List[int]],
+        groups: Optional[int] = 1,
+        bias: Optional[bool] = True,
+    ):
+        super().__init__()
+
+        if isinstance(kernel_shape, int):
+            self.kernel_shape = [kernel_shape]
+        else:
+            self.kernel_shape = kernel_shape
+
+        if len(self.kernel_shape) == 1:
+            self.kernel_size = self.kernel_shape[0]
+        elif len(self.kernel_shape) == 2:
+            self.kernel_size = (self.kernel_shape[0] - 1) * self.kernel_shape[1] + 1
+        else:
+            raise ValueError("kernel_shape should be either one- or two-dimensional.")
+
+        # groups
+        self.groups = groups
+
+        # weight tensor
+        if in_channels % self.groups != 0:
+            raise ValueError(
+                "Error, the number of input channels has to be an integer multiple of the group size"
+            )
+        if out_channels % self.groups != 0:
+            raise ValueError(
+                "Error, the number of output channels has to be an integer multiple of the group size"
+            )
+        self.groupsize = in_channels // self.groups
+        scale = math.sqrt(1.0 / self.groupsize)
+        self.weight = nn.Parameter(
+            scale * torch.randn(out_channels, self.groupsize, self.kernel_size)
+        )
+
+        if bias:
+            self.bias = nn.Parameter(torch.zeros(out_channels))
+        else:
+            self.bias = None
+
+    @abc.abstractmethod
+    def forward(self, x: torch.Tensor):
+        raise NotImplementedError
+
+
+def _disco_s2_contraction_torch(x: torch.Tensor, psi: torch.Tensor, nlon_out: int):
+    """
+    Reference implementation of the custom contraction as described in [1]. This requires repeated
+    shifting of the input tensor, which can potentially be costly. For an efficient implementation
+    on GPU, make sure to use the custom kernel written in Triton.
+    """
+    assert len(psi.shape) == 3
+    assert len(x.shape) == 4
+    psi = psi.to(x.device)
+
+    batch_size, n_chans, nlat_in, nlon_in = x.shape
+    kernel_size, nlat_out, _ = psi.shape
+
+    assert psi.shape[-1] == nlat_in * nlon_in
+    assert nlon_in % nlon_out == 0
+    assert nlon_in >= nlat_out
+    pscale = nlon_in // nlon_out
+
+    # add a dummy dimension for nkernel and move the batch and channel dims to the end
+    x = x.reshape(1, batch_size * n_chans, nlat_in, nlon_in).permute(0, 2, 3, 1)
+    x = x.expand(kernel_size, -1, -1, -1)
+
+    y = torch.zeros(
+        nlon_out,
+        kernel_size,
+        nlat_out,
+        batch_size * n_chans,
+        device=x.device,
+        dtype=x.dtype,
+    )
+
+    for pout in range(nlon_out):
+        # sparse contraction with psi
+        y[pout] = torch.bmm(psi, x.reshape(kernel_size, nlat_in * nlon_in, -1))
+        # we need to repeatedly roll the input tensor to faciliate the shifted multiplication
+        x = torch.roll(x, -pscale, dims=2)
+
+    # reshape y back to expose the correct dimensions
+    y = y.permute(3, 1, 2, 0).reshape(
+        batch_size, n_chans, kernel_size, nlat_out, nlon_out
+    )
+
+    return y
+
+
+def _disco_s2_transpose_contraction_torch(
+    x: torch.Tensor, psi: torch.Tensor, nlon_out: int
+):
+    """
+    Reference implementation of the custom contraction as described in [1]. This requires repeated
+    shifting of the input tensor, which can potentially be costly. For an efficient implementation
+    on GPU, make sure to use the custom kernel written in Triton.
+    """
+    assert len(psi.shape) == 3
+    assert len(x.shape) == 5
+    psi = psi.to(x.device)
+
+    batch_size, n_chans, kernel_size, nlat_in, nlon_in = x.shape
+    kernel_size, _, n_out = psi.shape
+
+    assert psi.shape[-2] == nlat_in
+    assert n_out % nlon_out == 0
+    nlat_out = n_out // nlon_out
+    assert nlon_out >= nlat_in
+    pscale = nlon_out // nlon_in
+
+    # we do a semi-transposition to faciliate the computation
+    inz = psi.indices()
+    tout = inz[2] // nlon_out
+    pout = inz[2] % nlon_out
+    # flip the axis of longitudes
+    pout = nlon_out - 1 - pout
+    tin = inz[1]
+    inz = torch.stack([inz[0], tout, tin * nlon_out + pout], dim=0)
+    psi_mod = torch.sparse_coo_tensor(
+        inz, psi.values(), size=(kernel_size, nlat_out, nlat_in * nlon_out)
+    )
+
+    # interleave zeros along the longitude dimension to allow for fractional offsets to be considered
+    x_ext = torch.zeros(
+        kernel_size,
+        nlat_in,
+        nlon_out,
+        batch_size * n_chans,
+        device=x.device,
+        dtype=x.dtype,
+    )
+    x_ext[:, :, ::pscale, :] = x.reshape(
+        batch_size * n_chans, kernel_size, nlat_in, nlon_in
+    ).permute(1, 2, 3, 0)
+    # we need to go backwards through the vector, so we flip the axis
+    x_ext = x_ext.contiguous()
+
+    y = torch.zeros(
+        kernel_size,
+        nlon_out,
+        nlat_out,
+        batch_size * n_chans,
+        device=x.device,
+        dtype=x.dtype,
+    )
+
+    for pout in range(nlon_out):
+        # we need to repeatedly roll the input tensor to faciliate the shifted multiplication
+        # TODO: double-check why this has to happen first
+        x_ext = torch.roll(x_ext, -1, dims=2)
+        # sparse contraction with the modified psi
+        y[:, pout, :, :] = torch.bmm(
+            psi_mod, x_ext.reshape(kernel_size, nlat_in * nlon_out, -1)
+        )
+
+    # sum over the kernel dimension and reshape to the correct output size
+    y = y.sum(dim=0).permute(2, 1, 0).reshape(batch_size, n_chans, nlat_out, nlon_out)
+
+    return y
+
+
+class DiscreteContinuousConvS2(DiscreteContinuousConv):
+    """
+    Discrete-continuous convolutions (DISCO) on the 2-Sphere as described in [1].
+
+    [1] Ocampo, Price, McEwen, Scalable and equivariant spherical CNNs by discrete-continuous (DISCO) convolutions, ICLR (2023), arXiv:2209.13603
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        in_shape: Tuple[int],
+        out_shape: Tuple[int],
+        kernel_shape: Union[int, List[int]],
+        groups: Optional[int] = 1,
+        grid_in: Optional[str] = "equiangular",
+        grid_out: Optional[str] = "equiangular",
+        bias: Optional[bool] = True,
+        theta_cutoff: Optional[float] = None,
+    ):
+        super().__init__(in_channels, out_channels, kernel_shape, groups, bias)
+
+        self.nlat_in, self.nlon_in = in_shape
+        self.nlat_out, self.nlon_out = out_shape
+
+        # compute theta cutoff based on the bandlimit of the input field
+        if theta_cutoff is None:
+            theta_cutoff = (
+                (self.kernel_shape[0] + 1) * torch.pi / float(self.nlat_in - 1)
+            )
+
+        if theta_cutoff <= 0.0:
+            raise ValueError("Error, theta_cutoff has to be positive.")
+
+        # integration weights
+        _, wgl = _precompute_latitudes(self.nlat_in, grid=grid_in)
+        quad_weights = (
+            2.0 * torch.pi * torch.from_numpy(wgl).float().reshape(-1, 1) / self.nlon_in
+        )
+        self.register_buffer("quad_weights", quad_weights, persistent=False)
+
+        idx, vals = _precompute_convolution_tensor_s2(
+            in_shape,
+            out_shape,
+            self.kernel_shape,
+            grid_in=grid_in,
+            grid_out=grid_out,
+            theta_cutoff=theta_cutoff,
+        )
+
+        self.register_buffer("psi_idx", idx, persistent=False)
+        self.register_buffer("psi_vals", vals, persistent=False)
+
+    def get_psi(self):
+        psi = torch.sparse_coo_tensor(
+            self.psi_idx,
+            self.psi_vals,
+            size=(self.kernel_size, self.nlat_out, self.nlat_in * self.nlon_in),
+        ).coalesce()
+        return psi
+
+    def forward(self, x: torch.Tensor, use_triton_kernel: bool = True) -> torch.Tensor:
+        # pre-multiply x with the quadrature weights
+        x = self.quad_weights * x
+
+        psi = self.get_psi()
+
+        if x.is_cuda and use_triton_kernel:
+            x = _disco_s2_contraction_triton(x, psi, self.nlon_out)
+        else:
+            x = _disco_s2_contraction_torch(x, psi, self.nlon_out)
+
+        # extract shape
+        B, C, K, H, W = x.shape
+        x = x.reshape(B, self.groups, self.groupsize, K, H, W)
+
+        # do weight multiplication
+        out = torch.einsum(
+            "bgckxy,gock->bgoxy",
+            x,
+            self.weight.reshape(
+                self.groups, -1, self.weight.shape[1], self.weight.shape[2]
+            ),
+        )
+        out = out.reshape(out.shape[0], -1, out.shape[-2], out.shape[-1])
+
+        if self.bias is not None:
+            out = out + self.bias.reshape(1, -1, 1, 1)
+
+        return out
+
+
+class DiscreteContinuousConvTransposeS2(DiscreteContinuousConv):
+    """
+    Discrete-continuous transpose convolutions (DISCO) on the 2-Sphere as described in [1].
+
+    [1] Ocampo, Price, McEwen, Scalable and equivariant spherical CNNs by discrete-continuous (DISCO) convolutions, ICLR (2023), arXiv:2209.13603
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        in_shape: Tuple[int],
+        out_shape: Tuple[int],
+        kernel_shape: Union[int, List[int]],
+        groups: Optional[int] = 1,
+        grid_in: Optional[str] = "equiangular",
+        grid_out: Optional[str] = "equiangular",
+        bias: Optional[bool] = True,
+        theta_cutoff: Optional[float] = None,
+    ):
+        super().__init__(in_channels, out_channels, kernel_shape, groups, bias)
+
+        self.nlat_in, self.nlon_in = in_shape
+        self.nlat_out, self.nlon_out = out_shape
+
+        # bandlimit
+        if theta_cutoff is None:
+            theta_cutoff = (
+                (self.kernel_shape[0] + 1) * torch.pi / float(self.nlat_in - 1)
+            )
+
+        if theta_cutoff <= 0.0:
+            raise ValueError("Error, theta_cutoff has to be positive.")
+
+        # integration weights
+        _, wgl = _precompute_latitudes(self.nlat_in, grid=grid_in)
+        quad_weights = (
+            2.0 * torch.pi * torch.from_numpy(wgl).float().reshape(-1, 1) / self.nlon_in
+        )
+        self.register_buffer("quad_weights", quad_weights, persistent=False)
+
+        # switch in_shape and out_shape since we want transpose conv
+        idx, vals = _precompute_convolution_tensor_s2(
+            out_shape,
+            in_shape,
+            self.kernel_shape,
+            grid_in=grid_out,
+            grid_out=grid_in,
+            theta_cutoff=theta_cutoff,
+        )
+
+        self.register_buffer("psi_idx", idx, persistent=False)
+        self.register_buffer("psi_vals", vals, persistent=False)
+
+    def get_psi(self):
+        psi = torch.sparse_coo_tensor(
+            self.psi_idx,
+            self.psi_vals,
+            size=(self.kernel_size, self.nlat_in, self.nlat_out * self.nlon_out),
+        ).coalesce()
+        return psi
+
+    def forward(self, x: torch.Tensor, use_triton_kernel: bool = True) -> torch.Tensor:
+        # extract shape
+        B, C, H, W = x.shape
+        x = x.reshape(B, self.groups, self.groupsize, H, W)
+
+        # do weight multiplication
+        x = torch.einsum(
+            "bgcxy,gock->bgokxy",
+            x,
+            self.weight.reshape(
+                self.groups, -1, self.weight.shape[1], self.weight.shape[2]
+            ),
+        )
+        x = x.reshape(x.shape[0], -1, x.shape[-3], x.shape[-2], x.shape[-1])
+
+        # pre-multiply x with the quadrature weights
+        x = self.quad_weights * x
+
+        psi = self.get_psi()
+
+        if x.is_cuda and use_triton_kernel:
+            out = _disco_s2_transpose_contraction_triton(x, psi, self.nlon_out)
+        else:
+            out = _disco_s2_transpose_contraction_torch(x, psi, self.nlon_out)
+
+        if self.bias is not None:
+            out = out + self.bias.reshape(1, -1, 1, 1)
+
+        return out
+
+
+class DiscreteContinuousConv2d(DiscreteContinuousConv):
+    """
+    Discrete-continuous convolutions (DISCO) on arbitrary 2d grids.
+
+    [1] Ocampo, Price, McEwen, Scalable and equivariant spherical CNNs by discrete-continuous (DISCO) convolutions, ICLR (2023), arXiv:2209.13603
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        grid_in: torch.Tensor,
+        grid_out: torch.Tensor,
+        kernel_shape: Union[int, List[int]],
+        n_in: Optional[Tuple[int]] = None,
+        n_out: Optional[Tuple[int]] = None,
+        quad_weights: Optional[torch.Tensor] = None,
+        periodic: Optional[bool] = False,
+        groups: Optional[int] = 1,
+        bias: Optional[bool] = True,
+        radius_cutoff: Optional[float] = None,
+    ):
+        super().__init__(in_channels, out_channels, kernel_shape, groups, bias)
+
+        # the instantiator supports convenience constructors for the input and output grids
+        if isinstance(grid_in, torch.Tensor):
+            assert isinstance(quad_weights, torch.Tensor)
+            assert not periodic
+        elif isinstance(grid_in, str):
+            assert n_in is not None
+            assert len(n_in) == 2
+            x, wx = _precompute_grid(n_in[0], grid=grid_in, periodic=periodic)
+            y, wy = _precompute_grid(n_in[1], grid=grid_in, periodic=periodic)
+            x, y = torch.meshgrid(torch.from_numpy(x), torch.from_numpy(y))
+            wx, wy = torch.meshgrid(torch.from_numpy(wx), torch.from_numpy(wy))
+            grid_in = torch.stack([x.reshape(-1), y.reshape(-1)])
+            quad_weights = (wx * wy).reshape(-1)
+        else:
+            raise ValueError(f"Unknown grid input type of type {type(grid_in)}")
+
+        if isinstance(grid_out, torch.Tensor):
+            pass
+        elif isinstance(grid_out, str):
+            assert n_out is not None
+            assert len(n_out) == 2
+            x, wx = _precompute_grid(n_out[0], grid=grid_out, periodic=periodic)
+            y, wy = _precompute_grid(n_out[1], grid=grid_out, periodic=periodic)
+            x, y = torch.meshgrid(torch.from_numpy(x), torch.from_numpy(y))
+            grid_out = torch.stack([x.reshape(-1), y.reshape(-1)])
+        else:
+            raise ValueError(f"Unknown grid output type of type {type(grid_out)}")
+
+        # check that input arrays are valid point clouds in 2D
+        assert len(grid_in.shape) == 2
+        assert len(grid_out.shape) == 2
+        assert len(quad_weights.shape) == 1
+        assert grid_in.shape[0] == 2
+        assert grid_out.shape[0] == 2
+
+        self.n_in = grid_in.shape[-1]
+        self.n_out = grid_out.shape[-1]
+
+        # compute the cutoff radius based on the bandlimit of the input field
+        # TODO: this heuristic is ad-hoc! Verify that we do the right one
+        if radius_cutoff is None:
+            radius_cutoff = (
+                2 * (self.kernel_shape[0] + 1) / float(math.sqrt(self.n_in) - 1)
+            )
+
+        if radius_cutoff <= 0.0:
+            raise ValueError("Error, radius_cutoff has to be positive.")
+
+        # integration weights
+        self.register_buffer("quad_weights", quad_weights, persistent=False)
+
+        idx, vals = _precompute_convolution_tensor_2d(
+            grid_in,
+            grid_out,
+            self.kernel_shape,
+            radius_cutoff=radius_cutoff,
+            periodic=periodic,
+        )
+
+        # to improve performance, we make psi a matrix by merging the first two dimensions
+        # This has to be accounted for in the forward pass
+        idx = torch.stack([idx[0] * self.n_out + idx[1], idx[2]], dim=0)
+
+        self.register_buffer("psi_idx", idx.contiguous(), persistent=False)
+        self.register_buffer("psi_vals", vals.contiguous(), persistent=False)
+
+    def get_psi(self):
+        psi = torch.sparse_coo_tensor(
+            self.psi_idx, self.psi_vals, size=(self.kernel_size * self.n_out, self.n_in)
+        )
+        return psi
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # pre-multiply x with the quadrature weights
+        x = self.quad_weights * x
+
+        psi = self.get_psi()
+
+        # extract shape
+        B, C, _ = x.shape
+
+        # bring into the right shape for the bmm and perform it
+        x = x.reshape(B * C, self.n_in).permute(1, 0).contiguous()
+        x = torch.mm(psi, x)
+        x = x.permute(1, 0).reshape(B, C, self.kernel_size, self.n_out)
+        x = x.reshape(B, self.groups, self.groupsize, self.kernel_size, self.n_out)
+
+        # do weight multiplication
+        out = torch.einsum(
+            "bgckx,gock->bgox",
+            x,
+            self.weight.reshape(
+                self.groups, -1, self.weight.shape[1], self.weight.shape[2]
+            ),
+        )
+        out = out.reshape(out.shape[0], -1, out.shape[-1])
+
+        if self.bias is not None:
+            out = out + self.bias.reshape(1, -1, 1)
+
+        return out
+
+
+class DiscreteContinuousConvTranspose2d(DiscreteContinuousConv):
+    """
+    Discrete-continuous convolutions (DISCO) on arbitrary 2d grids.
+
+    [1] Ocampo, Price, McEwen, Scalable and equivariant spherical CNNs by discrete-continuous (DISCO) convolutions, ICLR (2023), arXiv:2209.13603
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        grid_in: torch.Tensor,
+        grid_out: torch.Tensor,
+        kernel_shape: Union[int, List[int]],
+        n_in: Optional[Tuple[int]] = None,
+        n_out: Optional[Tuple[int]] = None,
+        quad_weights: Optional[torch.Tensor] = None,
+        periodic: Optional[bool] = False,
+        groups: Optional[int] = 1,
+        bias: Optional[bool] = True,
+        radius_cutoff: Optional[float] = None,
+    ):
+        super().__init__(in_channels, out_channels, kernel_shape, groups, bias)
+
+        # the instantiator supports convenience constructors for the input and output grids
+        if isinstance(grid_in, torch.Tensor):
+            assert isinstance(quad_weights, torch.Tensor)
+            assert not periodic
+        elif isinstance(grid_in, str):
+            assert n_in is not None
+            assert len(n_in) == 2
+            x, wx = _precompute_grid(n_in[0], grid=grid_in, periodic=periodic)
+            y, wy = _precompute_grid(n_in[1], grid=grid_in, periodic=periodic)
+            x, y = torch.meshgrid(torch.from_numpy(x), torch.from_numpy(y))
+            wx, wy = torch.meshgrid(torch.from_numpy(wx), torch.from_numpy(wy))
+            grid_in = torch.stack([x.reshape(-1), y.reshape(-1)])
+            quad_weights = (wx * wy).reshape(-1)
+        else:
+            raise ValueError(f"Unknown grid input type of type {type(grid_in)}")
+
+        if isinstance(grid_out, torch.Tensor):
+            pass
+        elif isinstance(grid_out, str):
+            assert n_out is not None
+            assert len(n_out) == 2
+            x, wx = _precompute_grid(n_out[0], grid=grid_out, periodic=periodic)
+            y, wy = _precompute_grid(n_out[1], grid=grid_out, periodic=periodic)
+            x, y = torch.meshgrid(torch.from_numpy(x), torch.from_numpy(y))
+            grid_out = torch.stack([x.reshape(-1), y.reshape(-1)])
+        else:
+            raise ValueError(f"Unknown grid output type of type {type(grid_out)}")
+
+        # check that input arrays are valid point clouds in 2D
+        assert len(grid_in.shape) == 2
+        assert len(grid_out.shape) == 2
+        assert len(quad_weights.shape) == 1
+        assert grid_in.shape[0] == 2
+        assert grid_out.shape[0] == 2
+
+        self.n_in = grid_in.shape[-1]
+        self.n_out = grid_out.shape[-1]
+
+        # compute the cutoff radius based on the bandlimit of the input field
+        # TODO: this heuristic is ad-hoc! Verify that we do the right one
+        if radius_cutoff is None:
+            radius_cutoff = (
+                2 * (self.kernel_shape[0] + 1) / float(math.sqrt(self.n_in) - 1)
+            )
+
+        if radius_cutoff <= 0.0:
+            raise ValueError("Error, radius_cutoff has to be positive.")
+
+        # integration weights
+        self.register_buffer("quad_weights", quad_weights, persistent=False)
+
+        # precompute the transposed tensor
+        idx, vals = _precompute_convolution_tensor_2d(
+            grid_out,
+            grid_in,
+            self.kernel_shape,
+            radius_cutoff=radius_cutoff,
+            periodic=periodic,
+        )
+
+        # to improve performance, we make psi a matrix by merging the first two dimensions
+        # This has to be accounted for in the forward pass
+        idx = torch.stack([idx[0] * self.n_out + idx[2], idx[1]], dim=0)
+
+        self.register_buffer("psi_idx", idx.contiguous(), persistent=False)
+        self.register_buffer("psi_vals", vals.contiguous(), persistent=False)
+
+    def get_psi(self):
+        psi = torch.sparse_coo_tensor(
+            self.psi_idx, self.psi_vals, size=(self.kernel_size * self.n_out, self.n_in)
+        )
+        return psi
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # pre-multiply x with the quadrature weights
+        x = self.quad_weights * x
+
+        psi = self.get_psi()
+
+        # extract shape
+        B, C, _ = x.shape
+
+        # bring into the right shape for the bmm and perform it
+        x = x.reshape(B * C, self.n_in).permute(1, 0).contiguous()
+        x = torch.mm(psi, x)
+        x = x.permute(1, 0).reshape(B, C, self.kernel_size, self.n_out)
+        x = x.reshape(B, self.groups, self.groupsize, self.kernel_size, self.n_out)
+
+        # do weight multiplication
+        out = torch.einsum(
+            "bgckx,gock->bgox",
+            x,
+            self.weight.reshape(
+                self.groups, -1, self.weight.shape[1], self.weight.shape[2]
+            ),
+        )
+        out = out.reshape(out.shape[0], -1, out.shape[-1])
+
+        if self.bias is not None:
+            out = out + self.bias.reshape(1, -1, 1)
+
+        return out
+
+
+class EquidistantDiscreteContinuousConv2d(DiscreteContinuousConv):
+    """
+    Discrete-continuous convolutions (DISCO) on arbitrary 2d grids.
+
+    [1] Ocampo, Price, McEwen, Scalable and equivariant spherical CNNs by discrete-continuous (DISCO) convolutions, ICLR (2023), arXiv:2209.13603
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_shape: Union[int, List[int]],
+        in_shape: Tuple[int],
+        groups: Optional[int] = 1,
+        bias: Optional[bool] = True,
+        radius_cutoff: Optional[float] = None,
+        padding_mode: str = "circular",
+        use_min_dim: bool = True,
+        **kwargs,
+    ):
+        """
+        use_min_dim (bool, optional): Use the minimum dimension of the input
+            shape to compute the cutoff radius. Otherwise use the maximum
+            dimension. Defaults to True.
+        """
+        super().__init__(in_channels, out_channels, kernel_shape, groups, bias)
+
+        self.padding_mode = padding_mode
+
+        # compute the cutoff radius based on the assumption that the grid is [-1, 1]^2
+        # this still assumes a quadratic domain
+        f = min if use_min_dim else max
+        if radius_cutoff is None:
+            radius_cutoff = 2 * (self.kernel_shape[0]) / float(f(*in_shape))
+        # 2 * 0.02 * 7 / 2 + 1 = 1.14
+        self.psi_local_size = math.floor(2 * radius_cutoff * f(*in_shape) / 2) + 1
+
+        # psi_local is essentially the support of the hat functions evaluated locally
+        x = torch.linspace(-radius_cutoff, radius_cutoff, self.psi_local_size)
+        x, y = torch.meshgrid(x, x)
+        grid_in = torch.stack([x.reshape(-1), y.reshape(-1)])
+        grid_out = torch.Tensor([[0.0], [0.0]])
+
+        idx, vals = _precompute_convolution_tensor_2d(
+            grid_in,
+            grid_out,
+            self.kernel_shape,
+            radius_cutoff=radius_cutoff,
+            periodic=False,
+        )
+
+        psi_loc = torch.zeros(
+            self.kernel_size, self.psi_local_size * self.psi_local_size
+        )
+        for ie in range(len(vals)):
+            f = idx[0, ie]
+            j = idx[2, ie]
+            v = vals[ie]
+            psi_loc[f, j] = v
+
+        # compute local version of the filter matrix
+        psi_loc = psi_loc.reshape(
+            self.kernel_size, self.psi_local_size, self.psi_local_size
+        )
+        # normalization by the quadrature weights
+        psi_loc = 4.0 * psi_loc / float(in_shape[0] * in_shape[1])
+
+        self.register_buffer("psi_loc", psi_loc, persistent=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+
+        kernel = torch.einsum("kxy,ogk->ogxy", self.psi_loc, self.weight)
+
+        left_pad = self.psi_local_size // 2
+        right_pad = (self.psi_local_size + 1) // 2 - 1
+        x = F.pad(x, (left_pad, right_pad, left_pad, right_pad), mode=self.padding_mode)
+        out = F.conv2d(
+            x, kernel, self.bias, stride=1, dilation=1, padding=0, groups=self.groups
+        )
+
+        return out

torch_harmonics_local/quadrature.py

@@ -0,0 +1,207 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+
+import numpy as np
+
+
+def _precompute_grid(n, grid="equidistant", a=0.0, b=1.0, periodic=False):
+
+    if (grid != "equidistant") and periodic:
+        raise ValueError(f"Periodic grid is only supported on equidistant grids.")
+
+    # compute coordinates
+    if grid == "equidistant":
+        xlg, wlg = trapezoidal_weights(n, a=a, b=b, periodic=periodic)
+    elif grid == "legendre-gauss":
+        xlg, wlg = legendre_gauss_weights(n, a=a, b=b)
+    elif grid == "lobatto":
+        xlg, wlg = lobatto_weights(n, a=a, b=b)
+    elif grid == "equiangular":
+        xlg, wlg = clenshaw_curtiss_weights(n, a=a, b=b)
+    else:
+        raise ValueError(f"Unknown grid type {grid}")
+
+    return xlg, wlg
+
+
+def _precompute_latitudes(nlat, grid="equiangular"):
+    r"""
+    Convenience routine to precompute latitudes
+    """
+
+    # compute coordinates
+    xlg, wlg = _precompute_grid(nlat, grid=grid, a=-1.0, b=1.0, periodic=False)
+
+    lats = np.flip(np.arccos(xlg)).copy()
+    wlg = np.flip(wlg).copy()
+
+    return lats, wlg
+
+
+def trapezoidal_weights(n, a=-1.0, b=1.0, periodic=False):
+    r"""
+    Helper routine which returns equidistant nodes with trapezoidal weights
+    on the interval [a, b]
+    """
+
+    xlg = np.linspace(a, b, n)
+    wlg = (b - a) / (n - 1) * np.ones(n)
+
+    if not periodic:
+        wlg[0] *= 0.5
+        wlg[-1] *= 0.5
+
+    return xlg, wlg
+
+
+def legendre_gauss_weights(n, a=-1.0, b=1.0):
+    r"""
+    Helper routine which returns the Legendre-Gauss nodes and weights
+    on the interval [a, b]
+    """
+
+    xlg, wlg = np.polynomial.legendre.leggauss(n)
+    xlg = (b - a) * 0.5 * xlg + (b + a) * 0.5
+    wlg = wlg * (b - a) * 0.5
+
+    return xlg, wlg
+
+
+def lobatto_weights(n, a=-1.0, b=1.0, tol=1e-16, maxiter=100):
+    r"""
+    Helper routine which returns the Legendre-Gauss-Lobatto nodes and weights
+    on the interval [a, b]
+    """
+
+    wlg = np.zeros((n,))
+    tlg = np.zeros((n,))
+    tmp = np.zeros((n,))
+
+    # Vandermonde Matrix
+    vdm = np.zeros((n, n))
+
+    # initialize Chebyshev nodes as first guess
+    for i in range(n):
+        tlg[i] = -np.cos(np.pi * i / (n - 1))
+
+    tmp = 2.0
+
+    for i in range(maxiter):
+        tmp = tlg
+
+        vdm[:, 0] = 1.0
+        vdm[:, 1] = tlg
+
+        for k in range(2, n):
+            vdm[:, k] = (
+                (2 * k - 1) * tlg * vdm[:, k - 1] - (k - 1) * vdm[:, k - 2]
+            ) / k
+
+        tlg = tmp - (tlg * vdm[:, n - 1] - vdm[:, n - 2]) / (n * vdm[:, n - 1])
+
+        if max(abs(tlg - tmp).flatten()) < tol:
+            break
+
+    wlg = 2.0 / ((n * (n - 1)) * (vdm[:, n - 1] ** 2))
+
+    # rescale
+    tlg = (b - a) * 0.5 * tlg + (b + a) * 0.5
+    wlg = wlg * (b - a) * 0.5
+
+    return tlg, wlg
+
+
+def clenshaw_curtiss_weights(n, a=-1.0, b=1.0):
+    r"""
+    Computation of the Clenshaw-Curtis quadrature nodes and weights.
+    This implementation follows
+
+    [1] Joerg Waldvogel, Fast Construction of the Fejer and Clenshaw-Curtis Quadrature Rules; BIT Numerical Mathematics, Vol. 43, No. 1, pp. 001–018.
+    """
+
+    assert n > 1
+
+    tcc = np.cos(np.linspace(np.pi, 0, n))
+
+    if n == 2:
+        wcc = np.array([1.0, 1.0])
+    else:
+
+        n1 = n - 1
+        N = np.arange(1, n1, 2)
+        l = len(N)
+        m = n1 - l
+
+        v = np.concatenate([2 / N / (N - 2), 1 / N[-1:], np.zeros(m)])
+        v = 0 - v[:-1] - v[-1:0:-1]
+
+        g0 = -np.ones(n1)
+        g0[l] = g0[l] + n1
+        g0[m] = g0[m] + n1
+        g = g0 / (n1**2 - 1 + (n1 % 2))
+        wcc = np.fft.ifft(v + g).real
+        wcc = np.concatenate((wcc, wcc[:1]))
+
+    # rescale
+    tcc = (b - a) * 0.5 * tcc + (b + a) * 0.5
+    wcc = wcc * (b - a) * 0.5
+
+    return tcc, wcc
+
+
+def fejer2_weights(n, a=-1.0, b=1.0):
+    r"""
+    Computation of the Fejer quadrature nodes and weights.
+    This implementation follows
+
+    [1] Joerg Waldvogel, Fast Construction of the Fejer and Clenshaw-Curtis Quadrature Rules; BIT Numerical Mathematics, Vol. 43, No. 1, pp. 001–018.
+    """
+
+    assert n > 2
+
+    tcc = np.cos(np.linspace(np.pi, 0, n))
+
+    n1 = n - 1
+    N = np.arange(1, n1, 2)
+    l = len(N)
+    m = n1 - l
+
+    v = np.concatenate([2 / N / (N - 2), 1 / N[-1:], np.zeros(m)])
+    v = 0 - v[:-1] - v[-1:0:-1]
+
+    wcc = np.fft.ifft(v).real
+    wcc = np.concatenate((wcc, wcc[:1]))
+
+    # rescale
+    tcc = (b - a) * 0.5 * tcc + (b + a) * 0.5
+    wcc = wcc * (b - a) * 0.5
+
+    return tcc, wcc

type_utils.py

@@ -0,0 +1,4 @@
+def tuple_type(strings):
+    strings = strings.replace("(", "").replace(")", "").replace(" ", "")
+    mapped_int = map(int, strings.split(","))
+    return tuple(mapped_int)