# neuraloperator/mri architecture

> Architecture diagram of the GitHub repository neuraloperator/mri, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-09-28). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/neuraloperator/mri
- Repository: https://github.com/neuraloperator/mri

## Overview

NoMRI is a compressed-sensing MRI reconstruction research system. Its main workflow builds undersampling masks and slice datasets, trains or validates a reconstruction model, and evaluates reconstructed images. The README also documents conversion of raw fastMRI samples to LMDB, a conventional VarNet baseline, and pattern-independent NO reconstruction. The graph reflects documented capabilities where source was not sampled; those uncertain implementation links are omitted.

Main entry points: [`main.py`](https://github.com/neuraloperator/mri/blob/main/main.py) (Train / validate), [`fastmri/datasets.py`](https://github.com/neuraloperator/mri/blob/main/fastmri/datasets.py) (Slice datasets), [`models/lightning/no_varnet_module.py`](https://github.com/neuraloperator/mri/blob/main/models/lightning/no_varnet_module.py) (NO Lightning module), [`models/lightning/varnet_module.py`](https://github.com/neuraloperator/mri/blob/main/models/lightning/varnet_module.py) (VarNet Lightning module)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_workflow["Experiment workflow"]
  node_main["Train / validate<br/>[main.py]"]
  node_converter["LMDB conversion"]
end

subgraph group_data["MRI data and sampling"]
  node_rawdata[("fastMRI HDF5")]
  node_lmdb[("LMDB slices")]
  node_datasets["Slice datasets<br/>[datasets.py]"]
  node_sampling["Mask patterns<br/>[subsample.py]"]
  node_transforms["MRI transforms<br/>[transforms.py]"]
  node_fastmri_ops["FFT and coil ops<br/>[fftc.py]"]
end

subgraph group_reconstruction["Reconstruction models"]
  node_no_varnet["NO-VarNet<br/>[no_varnet.py]"]
  node_udno["U-shaped DISCO operator<br/>[udno.py]"]
  node_varnet["VarNet baseline<br/>[varnet.py]"]
  node_unet["U-Net regularizer<br/>[unet.py]"]
end

subgraph group_operator["Operator implementation"]
  node_convolution["DISCO convolution<br/>[convolution.py]"]
  node_disco_kernel["DISCO kernels"]
end

subgraph group_evaluation["Training and evaluation"]
  node_no_module["NO Lightning module"]
  node_varnet_module["VarNet Lightning module<br/>[varnet_module.py]"]
  node_metrics["Validation metrics<br/>[mri_module.py]"]
end

node_user(("Researcher"))
node_wandb(("Weights &amp; Biases"))
node_reconstruction_result["MRI reconstructions"]

node_user -->|"configures run"| node_main
node_user -->|"requests conversion"| node_converter
node_rawdata -->|"provides samples"| node_converter
node_converter -->|"writes slices"| node_lmdb
node_main -->|"creates masks"| node_sampling
node_main -->|"builds loaders"| node_datasets
node_datasets -->|"reads slices"| node_lmdb
node_datasets -->|"applies masks"| node_sampling
node_datasets -->|"prepares samples"| node_transforms
node_main -->|"configures model"| node_no_module
node_main -.->|"configures model"| node_varnet_module
node_no_module -->|"runs reconstruction"| node_no_varnet
node_varnet_module -.->|"runs reconstruction"| node_varnet
node_no_varnet -->|"uses operator"| node_udno
node_varnet -->|"uses regularizer"| node_unet
node_udno -->|"uses convolution"| node_convolution
node_convolution -->|"dispatches kernels"| node_disco_kernel
node_no_varnet -->|"uses MRI operations"| node_fastmri_ops
node_varnet -->|"uses MRI operations"| node_fastmri_ops
node_no_module -->|"logs validation"| node_metrics
node_varnet_module -.->|"logs validation"| node_metrics
node_no_module -->|"produces images"| node_reconstruction_result
node_varnet_module -.->|"produces images"| node_reconstruction_result
node_metrics -.->|"logs metrics"| node_wandb

click node_main "https://github.com/neuraloperator/mri/blob/main/main.py"
click node_converter "https://github.com/neuraloperator/mri/blob/main/scripts/gen_lmdb_dataset.py"
click node_datasets "https://github.com/neuraloperator/mri/blob/main/fastmri/datasets.py"
click node_sampling "https://github.com/neuraloperator/mri/blob/main/fastmri/subsample.py"
click node_transforms "https://github.com/neuraloperator/mri/blob/main/fastmri/transforms.py"
click node_no_module "https://github.com/neuraloperator/mri/blob/main/models/lightning/no_varnet_module.py"
click node_varnet_module "https://github.com/neuraloperator/mri/blob/main/models/lightning/varnet_module.py"
click node_metrics "https://github.com/neuraloperator/mri/blob/main/models/lightning/mri_module.py"
click node_no_varnet "https://github.com/neuraloperator/mri/blob/main/models/no_varnet.py"
click node_udno "https://github.com/neuraloperator/mri/blob/main/models/udno.py"
click node_varnet "https://github.com/neuraloperator/mri/blob/main/models/varnet.py"
click node_unet "https://github.com/neuraloperator/mri/blob/main/models/unet.py"
click node_convolution "https://github.com/neuraloperator/mri/blob/main/torch_harmonics/convolution.py"
click node_disco_kernel "https://github.com/neuraloperator/mri/blob/main/torch_harmonics/_disco_convolution.py"
click node_fastmri_ops "https://github.com/neuraloperator/mri/blob/main/fastmri/fftc.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_main,node_converter toneBlue
class node_rawdata,node_lmdb,node_datasets,node_sampling,node_transforms,node_fastmri_ops toneAmber
class node_no_varnet,node_udno,node_varnet,node_unet toneMint
class node_convolution,node_disco_kernel toneRose
class node_no_module,node_varnet_module,node_metrics,node_user,node_wandb,node_reconstruction_result toneIndigo
```

## Components

### Experiment workflow

- [Train / validate](https://github.com/neuraloperator/mri/blob/main/main.py): `main.py`
- [LMDB conversion](https://github.com/neuraloperator/mri/blob/main/scripts/gen_lmdb_dataset.py): `scripts/gen_lmdb_dataset.py`

### MRI data and sampling

- **fastMRI HDF5**
- **LMDB slices**
- [Slice datasets](https://github.com/neuraloperator/mri/blob/main/fastmri/datasets.py): `fastmri/datasets.py`
- [Mask patterns](https://github.com/neuraloperator/mri/blob/main/fastmri/subsample.py): `fastmri/subsample.py`
- [MRI transforms](https://github.com/neuraloperator/mri/blob/main/fastmri/transforms.py): `fastmri/transforms.py`
- [FFT and coil ops](https://github.com/neuraloperator/mri/blob/main/fastmri/fftc.py): `fastmri/fftc.py`

### Reconstruction models

- [NO-VarNet](https://github.com/neuraloperator/mri/blob/main/models/no_varnet.py): `models/no_varnet.py`
- [U-shaped DISCO operator](https://github.com/neuraloperator/mri/blob/main/models/udno.py): `models/udno.py`
- [VarNet baseline](https://github.com/neuraloperator/mri/blob/main/models/varnet.py): `models/varnet.py`
- [U-Net regularizer](https://github.com/neuraloperator/mri/blob/main/models/unet.py): `models/unet.py`

### Operator implementation

- [DISCO convolution](https://github.com/neuraloperator/mri/blob/main/torch_harmonics/convolution.py): `torch_harmonics/convolution.py`
- [DISCO kernels](https://github.com/neuraloperator/mri/blob/main/torch_harmonics/_disco_convolution.py): `torch_harmonics/_disco_convolution.py`

### Training and evaluation

- [NO Lightning module](https://github.com/neuraloperator/mri/blob/main/models/lightning/no_varnet_module.py): `models/lightning/no_varnet_module.py`
- [VarNet Lightning module](https://github.com/neuraloperator/mri/blob/main/models/lightning/varnet_module.py): `models/lightning/varnet_module.py`
- [Validation metrics](https://github.com/neuraloperator/mri/blob/main/models/lightning/mri_module.py): `models/lightning/mri_module.py`

### Other components

- **Researcher**
- **Weights & Biases**
- **MRI reconstructions**

## Connections

- Researcher → Train / validate: configures run
- Researcher → LMDB conversion: requests conversion
- fastMRI HDF5 → LMDB conversion: provides samples
- LMDB conversion → LMDB slices: writes slices
- Train / validate → Mask patterns: creates masks
- Train / validate → Slice datasets: builds loaders
- Slice datasets → LMDB slices: reads slices
- Slice datasets → Mask patterns: applies masks
- Slice datasets → MRI transforms: prepares samples
- Train / validate → NO Lightning module: configures model
- Train / validate → VarNet Lightning module: configures model
- NO Lightning module → NO-VarNet: runs reconstruction
- VarNet Lightning module → VarNet baseline: runs reconstruction
- NO-VarNet → U-shaped DISCO operator: uses operator
- VarNet baseline → U-Net regularizer: uses regularizer
- U-shaped DISCO operator → DISCO convolution: uses convolution
- DISCO convolution → DISCO kernels: dispatches kernels
- NO-VarNet → FFT and coil ops: uses MRI operations
- VarNet baseline → FFT and coil ops: uses MRI operations
- NO Lightning module → Validation metrics: logs validation
- VarNet Lightning module → Validation metrics: logs validation
- NO Lightning module → MRI reconstructions: produces images
- VarNet Lightning module → MRI reconstructions: produces images
- Validation metrics → Weights & Biases: logs metrics