# edongdongchen/ei architecture

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

- Interactive diagram: https://gitdiagram.com/edongdongchen/ei
- Repository: https://github.com/edongdongchen/ei

## Overview

Equivariant Imaging (EI) is a research framework for reconstructing images from measurements when ground-truth training images are unavailable. Its main workflow loads image batches, applies a task-specific forward model, and trains a neural reconstructor using measurement consistency and equivariance under image transformations; optional adversarial and supervised variants are also present. For evaluation, a measurement is mapped through the pseudoinverse and trained network to produce a reconstruction. The graph reflects the documented workflow and sampled implementations; some closure and CT details are not shown in the excerpts.

Main entry points: [`demo_train.py`](https://github.com/edongdongchen/ei/blob/main/demo_train.py) (Training entry), [`ei/closure/ei.py`](https://github.com/edongdongchen/ei/blob/main/ei/closure/ei.py) (EI objective), [`demo_test_inpainting.py`](https://github.com/edongdongchen/ei/blob/main/demo_test_inpainting.py) (Inpainting evaluation), [`models/unet.py`](https://github.com/edongdongchen/ei/blob/main/models/unet.py) (Reconstruction U-Net)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_data["Image data"]
  node_ct_data[("CT data<br/>[ctdb.py]")]
  node_image_data[("Image datasets<br/>[cvdb.py]")]
end

subgraph group_physics["Imaging operators"]
  node_ct_task["CT operator<br/>[ct.py]"]
  node_inpaint_task["Inpainting operator<br/>[inpainting.py]"]
  node_radon["Radon transforms<br/>[radon.py]"]
  node_rotate["Rotation transforms<br/>[rotate.py]"]
  node_shift["Shift transforms<br/>[shift.py]"]
end

subgraph group_learning["EI learning"]
  node_training["Training entry<br/>[demo_train.py]"]
  node_ei_api["EI training API<br/>[ei.py]"]
  node_ei_closure["EI objective<br/>[ei.py]"]
  node_adv_closure["Adversarial EI<br/>[ei_adv.py]"]
  node_sup_closure["Supervised objectives<br/>[supervised.py]"]
  node_sup_ei_closure["Supervised EI<br/>[supervised_ei.py]"]
  node_generator["Reconstruction U-Net<br/>[unet.py]"]
  node_discriminator["Discriminator<br/>[discriminator.py]"]
end

subgraph group_evaluation["Reconstruction"]
  node_test["Inpainting evaluation"]
  node_metrics["Image metrics<br/>[metric.py]"]
  node_reconstruction["Reconstructed image"]
end

node_user(("Experimenter"))

node_user -->|"configures"| node_training
node_training -->|"loads CT batches"| node_ct_data
node_training -->|"loads image batches"| node_image_data
node_training -->|"selects CT model"| node_ct_task
node_training -->|"selects inpainting model"| node_inpaint_task
node_training -->|"selects rotations"| node_rotate
node_training -->|"selects shifts"| node_shift
node_training -->|"starts training"| node_ei_api
node_ei_api -->|"runs objective"| node_ei_closure
node_ei_closure -->|"reconstructs"| node_generator
node_ei_closure -->|"uses forward model"| node_ct_task
node_ei_closure -->|"uses forward model"| node_inpaint_task
node_ct_task -->|"uses projection"| node_radon
node_ei_closure -->|"transforms estimate"| node_rotate
node_ei_closure -->|"transforms estimate"| node_shift
node_adv_closure -->|"trains generator"| node_generator
node_adv_closure -->|"trains discriminator"| node_discriminator
node_sup_closure -->|"trains reconstructor"| node_generator
node_sup_ei_closure -->|"trains reconstructor"| node_generator
node_user -->|"evaluates checkpoint"| node_test
node_test -->|"loads test images"| node_image_data
node_test -->|"forms measurements"| node_inpaint_task
node_test -->|"applies checkpoint"| node_generator
node_test -->|"scores output"| node_metrics
node_test -->|"produces"| node_reconstruction

click node_training "https://github.com/edongdongchen/ei/blob/main/demo_train.py"
click node_ct_data "https://github.com/edongdongchen/ei/blob/main/dataset/ctdb.py"
click node_image_data "https://github.com/edongdongchen/ei/blob/main/dataset/cvdb.py"
click node_ct_task "https://github.com/edongdongchen/ei/blob/main/physics/ct.py"
click node_inpaint_task "https://github.com/edongdongchen/ei/blob/main/physics/inpainting.py"
click node_radon "https://github.com/edongdongchen/ei/blob/main/physics/radon/radon.py"
click node_rotate "https://github.com/edongdongchen/ei/blob/main/transforms/rotate.py"
click node_shift "https://github.com/edongdongchen/ei/blob/main/transforms/shift.py"
click node_ei_api "https://github.com/edongdongchen/ei/blob/main/ei/ei.py"
click node_ei_closure "https://github.com/edongdongchen/ei/blob/main/ei/closure/ei.py"
click node_adv_closure "https://github.com/edongdongchen/ei/blob/main/ei/closure/ei_adv.py"
click node_sup_closure "https://github.com/edongdongchen/ei/blob/main/ei/closure/supervised.py"
click node_sup_ei_closure "https://github.com/edongdongchen/ei/blob/main/ei/closure/supervised_ei.py"
click node_generator "https://github.com/edongdongchen/ei/blob/main/models/unet.py"
click node_discriminator "https://github.com/edongdongchen/ei/blob/main/models/discriminator.py"
click node_test "https://github.com/edongdongchen/ei/blob/main/demo_test_inpainting.py"
click node_metrics "https://github.com/edongdongchen/ei/blob/main/utils/metric.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_ct_data,node_image_data toneBlue
class node_ct_task,node_inpaint_task,node_radon,node_rotate,node_shift toneAmber
class node_training,node_ei_api,node_ei_closure,node_adv_closure,node_sup_closure,node_sup_ei_closure,node_generator,node_discriminator toneMint
class node_test,node_metrics,node_reconstruction toneRose
class node_user toneIndigo
```

## Components

### Image data

- [CT data](https://github.com/edongdongchen/ei/blob/main/dataset/ctdb.py): `dataset/ctdb.py`
- [Image datasets](https://github.com/edongdongchen/ei/blob/main/dataset/cvdb.py): `dataset/cvdb.py`

### Imaging operators

- [CT operator](https://github.com/edongdongchen/ei/blob/main/physics/ct.py): `physics/ct.py`
- [Inpainting operator](https://github.com/edongdongchen/ei/blob/main/physics/inpainting.py): `physics/inpainting.py`
- [Radon transforms](https://github.com/edongdongchen/ei/blob/main/physics/radon/radon.py): `physics/radon/radon.py`
- [Rotation transforms](https://github.com/edongdongchen/ei/blob/main/transforms/rotate.py): `transforms/rotate.py`
- [Shift transforms](https://github.com/edongdongchen/ei/blob/main/transforms/shift.py): `transforms/shift.py`

### EI learning

- [Training entry](https://github.com/edongdongchen/ei/blob/main/demo_train.py): `demo_train.py`
- [EI training API](https://github.com/edongdongchen/ei/blob/main/ei/ei.py): `ei/ei.py`
- [EI objective](https://github.com/edongdongchen/ei/blob/main/ei/closure/ei.py): `ei/closure/ei.py`
- [Adversarial EI](https://github.com/edongdongchen/ei/blob/main/ei/closure/ei_adv.py): `ei/closure/ei_adv.py`
- [Supervised objectives](https://github.com/edongdongchen/ei/blob/main/ei/closure/supervised.py): `ei/closure/supervised.py`
- [Supervised EI](https://github.com/edongdongchen/ei/blob/main/ei/closure/supervised_ei.py): `ei/closure/supervised_ei.py`
- [Reconstruction U-Net](https://github.com/edongdongchen/ei/blob/main/models/unet.py): `models/unet.py`
- [Discriminator](https://github.com/edongdongchen/ei/blob/main/models/discriminator.py): `models/discriminator.py`

### Reconstruction

- [Inpainting evaluation](https://github.com/edongdongchen/ei/blob/main/demo_test_inpainting.py): `demo_test_inpainting.py`
- [Image metrics](https://github.com/edongdongchen/ei/blob/main/utils/metric.py): `utils/metric.py`
- **Reconstructed image**

### Other components

- **Experimenter**

## Connections

- Experimenter → Training entry: configures
- Training entry → CT data: loads CT batches
- Training entry → Image datasets: loads image batches
- Training entry → CT operator: selects CT model
- Training entry → Inpainting operator: selects inpainting model
- Training entry → Rotation transforms: selects rotations
- Training entry → Shift transforms: selects shifts
- Training entry → EI training API: starts training
- EI training API → EI objective: runs objective
- EI objective → Reconstruction U-Net: reconstructs
- EI objective → CT operator: uses forward model
- EI objective → Inpainting operator: uses forward model
- CT operator → Radon transforms: uses projection
- EI objective → Rotation transforms: transforms estimate
- EI objective → Shift transforms: transforms estimate
- Adversarial EI → Reconstruction U-Net: trains generator
- Adversarial EI → Discriminator: trains discriminator
- Supervised objectives → Reconstruction U-Net: trains reconstructor
- Supervised EI → Reconstruction U-Net: trains reconstructor
- Experimenter → Inpainting evaluation: evaluates checkpoint
- Inpainting evaluation → Image datasets: loads test images
- Inpainting evaluation → Inpainting operator: forms measurements
- Inpainting evaluation → Reconstruction U-Net: applies checkpoint
- Inpainting evaluation → Image metrics: scores output
- Inpainting evaluation → Reconstructed image: produces