# s-sayan/shearnet architecture

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

- Interactive diagram: https://gitdiagram.com/s-sayan/shearnet
- Repository: https://github.com/s-sayan/shearnet

## Overview

ShearNet is a JAX-based galaxy-shear estimation product: it simulates or loads labeled galaxy images, trains a model, evaluates observations with ShearNet or ngmix/metacalibration, and merges estimator catalogs. The graph also includes documented pipeline dispatch and the run/recovery layer. Some edges rely on README statements or SOURCE INDEX references; unsampled implementation details are not inferred.

Main entry points: [`shearnet/evaluation/service.py`](https://github.com/s-sayan/shearnet/blob/main/shearnet/evaluation/service.py) (Evaluation service), [`shearnet/training/service.py`](https://github.com/s-sayan/shearnet/blob/main/shearnet/training/service.py) (Training service), [`shearnet/cli/pipeline.py`](https://github.com/s-sayan/shearnet/blob/main/shearnet/cli/pipeline.py) (Pipeline CLI), [`shearnet/cli/train.py`](https://github.com/s-sayan/shearnet/blob/main/shearnet/cli/train.py) (Training CLI)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_configuration["Configuration"]
  node_config["Config validation<br/>[schema.py]"]
  node_configload["Config loading<br/>[config_handler.py]"]
  node_traincli["Training CLI<br/>[train.py]"]
  node_evalcli["Evaluation CLI<br/>[evaluate.py]"]
  node_pipelinecli["Pipeline CLI<br/>[pipeline.py]"]
end

subgraph group_learning["Simulation and training"]
  node_simdata["Galaxy simulation<br/>[dataset.py]"]
  node_jaxdata["Differentiable simulation<br/>[dataset_jax.py]"]
  node_inloop["In-loop generation<br/>[inloop.py]"]
  node_model["Neural models<br/>[models.py]"]
  node_training["Training service<br/>[service.py]"]
end

subgraph group_measurement["Evaluation and catalogs"]
  node_evalservice["Evaluation service<br/>[service.py]"]
  node_render["Observation rendering<br/>[rendering.py]"]
  node_predict["Model prediction<br/>[predictor.py]"]
  node_workers["Measurement workers<br/>[workers.py]"]
  node_ngmix["NGmix metacalibration<br/>[ngmix.py]"]
  node_catalogschema["Catalog schema<br/>[catalog_schema.py]"]
  node_fits[("FITS catalog<br/>[fits_catalog.py]")]
  node_merge["Catalog merge<br/>[merge.py]"]
end

subgraph group_operations["Runs and orchestration"]
  node_runstore[("Run artifacts<br/>[runs.py]")]
  node_checkpoints["Training checkpoints<br/>[checkpoints.py]"]
  node_recovery["Recovery controls<br/>[recovery.py]"]
end

node_user(("Researcher"))
node_slurm["Slurm scheduler"]

node_user -->|"starts training"| node_traincli
node_user -->|"starts evaluation"| node_evalcli
node_user -->|"submits or resumes"| node_pipelinecli
node_traincli -->|"loads config"| node_configload
node_configload -->|"validates settings"| node_config
node_traincli -->|"runs training"| node_training
node_training -->|"renders upfront data"| node_simdata
node_training -->|"selects in-loop"| node_inloop
node_inloop -->|"generates batches"| node_jaxdata
node_training -->|"trains model"| node_model
node_training -->|"writes run"| node_runstore
node_training -->|"saves checkpoints"| node_checkpoints
node_pipelinecli -->|"submits stages"| node_slurm
node_pipelinecli -->|"dispatches training"| node_traincli
node_pipelinecli -->|"dispatches evaluation"| node_evalservice
node_pipelinecli -->|"dispatches merge"| node_merge
node_evalcli -.->|"loads config"| node_configload
node_evalcli -->|"runs estimator"| node_evalservice
node_evalservice -->|"renders observations"| node_render
node_evalservice -->|"uses prediction"| node_predict
node_evalservice -->|"dispatches work"| node_workers
node_workers -->|"measures metacal"| node_ngmix
node_evalservice -->|"builds tables"| node_catalogschema
node_evalservice -->|"writes catalog"| node_fits
node_evalservice -->|"reads run artifacts"| node_runstore
node_evalservice -->|"guards and resumes"| node_recovery
node_user -->|"requests catalog merge"| node_merge
node_merge -->|"reads catalogs"| node_fits
node_merge -->|"checks schema"| node_catalogschema

click node_config "https://github.com/s-sayan/shearnet/blob/main/shearnet/config/schema.py"
click node_configload "https://github.com/s-sayan/shearnet/blob/main/shearnet/config/config_handler.py"
click node_traincli "https://github.com/s-sayan/shearnet/blob/main/shearnet/cli/train.py"
click node_evalcli "https://github.com/s-sayan/shearnet/blob/main/shearnet/cli/evaluate.py"
click node_pipelinecli "https://github.com/s-sayan/shearnet/blob/main/shearnet/cli/pipeline.py"
click node_simdata "https://github.com/s-sayan/shearnet/blob/main/shearnet/core/dataset.py"
click node_jaxdata "https://github.com/s-sayan/shearnet/blob/main/shearnet/core/dataset_jax.py"
click node_inloop "https://github.com/s-sayan/shearnet/blob/main/shearnet/core/inloop.py"
click node_model "https://github.com/s-sayan/shearnet/blob/main/shearnet/core/models.py"
click node_training "https://github.com/s-sayan/shearnet/blob/main/shearnet/training/service.py"
click node_runstore "https://github.com/s-sayan/shearnet/blob/main/shearnet/artifacts/runs.py"
click node_checkpoints "https://github.com/s-sayan/shearnet/blob/main/shearnet/artifacts/checkpoints.py"
click node_recovery "https://github.com/s-sayan/shearnet/blob/main/shearnet/artifacts/recovery.py"
click node_evalservice "https://github.com/s-sayan/shearnet/blob/main/shearnet/evaluation/service.py"
click node_render "https://github.com/s-sayan/shearnet/blob/main/shearnet/evaluation/rendering.py"
click node_predict "https://github.com/s-sayan/shearnet/blob/main/shearnet/evaluation/predictor.py"
click node_workers "https://github.com/s-sayan/shearnet/blob/main/shearnet/evaluation/workers.py"
click node_ngmix "https://github.com/s-sayan/shearnet/blob/main/shearnet/methods/ngmix.py"
click node_catalogschema "https://github.com/s-sayan/shearnet/blob/main/shearnet/io/catalog_schema.py"
click node_fits "https://github.com/s-sayan/shearnet/blob/main/shearnet/io/fits_catalog.py"
click node_merge "https://github.com/s-sayan/shearnet/blob/main/shearnet/io/merge.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_config,node_configload,node_traincli,node_evalcli,node_pipelinecli toneBlue
class node_simdata,node_jaxdata,node_inloop,node_model,node_training toneAmber
class node_evalservice,node_render,node_predict,node_workers,node_ngmix,node_catalogschema,node_fits,node_merge toneMint
class node_runstore,node_checkpoints,node_recovery,node_slurm toneRose
class node_user toneIndigo
```

## Components

### Configuration

- [Config validation](https://github.com/s-sayan/shearnet/blob/main/shearnet/config/schema.py): `shearnet/config/schema.py`
- [Config loading](https://github.com/s-sayan/shearnet/blob/main/shearnet/config/config_handler.py): `shearnet/config/config_handler.py`
- [Training CLI](https://github.com/s-sayan/shearnet/blob/main/shearnet/cli/train.py): `shearnet/cli/train.py`
- [Evaluation CLI](https://github.com/s-sayan/shearnet/blob/main/shearnet/cli/evaluate.py): `shearnet/cli/evaluate.py`
- [Pipeline CLI](https://github.com/s-sayan/shearnet/blob/main/shearnet/cli/pipeline.py): `shearnet/cli/pipeline.py`

### Simulation and training

- [Galaxy simulation](https://github.com/s-sayan/shearnet/blob/main/shearnet/core/dataset.py): `shearnet/core/dataset.py`
- [Differentiable simulation](https://github.com/s-sayan/shearnet/blob/main/shearnet/core/dataset_jax.py): `shearnet/core/dataset_jax.py`
- [In-loop generation](https://github.com/s-sayan/shearnet/blob/main/shearnet/core/inloop.py): `shearnet/core/inloop.py`
- [Neural models](https://github.com/s-sayan/shearnet/blob/main/shearnet/core/models.py): `shearnet/core/models.py`
- [Training service](https://github.com/s-sayan/shearnet/blob/main/shearnet/training/service.py): `shearnet/training/service.py`

### Evaluation and catalogs

- [Evaluation service](https://github.com/s-sayan/shearnet/blob/main/shearnet/evaluation/service.py): `shearnet/evaluation/service.py`
- [Observation rendering](https://github.com/s-sayan/shearnet/blob/main/shearnet/evaluation/rendering.py): `shearnet/evaluation/rendering.py`
- [Model prediction](https://github.com/s-sayan/shearnet/blob/main/shearnet/evaluation/predictor.py): `shearnet/evaluation/predictor.py`
- [Measurement workers](https://github.com/s-sayan/shearnet/blob/main/shearnet/evaluation/workers.py): `shearnet/evaluation/workers.py`
- [NGmix metacalibration](https://github.com/s-sayan/shearnet/blob/main/shearnet/methods/ngmix.py): `shearnet/methods/ngmix.py`
- [Catalog schema](https://github.com/s-sayan/shearnet/blob/main/shearnet/io/catalog_schema.py): `shearnet/io/catalog_schema.py`
- [FITS catalog](https://github.com/s-sayan/shearnet/blob/main/shearnet/io/fits_catalog.py): `shearnet/io/fits_catalog.py`
- [Catalog merge](https://github.com/s-sayan/shearnet/blob/main/shearnet/io/merge.py): `shearnet/io/merge.py`

### Runs and orchestration

- [Run artifacts](https://github.com/s-sayan/shearnet/blob/main/shearnet/artifacts/runs.py): `shearnet/artifacts/runs.py`
- [Training checkpoints](https://github.com/s-sayan/shearnet/blob/main/shearnet/artifacts/checkpoints.py): `shearnet/artifacts/checkpoints.py`
- [Recovery controls](https://github.com/s-sayan/shearnet/blob/main/shearnet/artifacts/recovery.py): `shearnet/artifacts/recovery.py`

### Other components

- **Researcher**
- **Slurm scheduler**

## Connections

- Researcher → Training CLI: starts training (evidence: `README.md`)
- Researcher → Evaluation CLI: starts evaluation (evidence: `README.md`)
- Researcher → Pipeline CLI: submits or resumes (evidence: `README.md`)
- Training CLI → Config loading: loads config (evidence: `shearnet/cli/train.py`)
- Config loading → Config validation: validates settings (evidence: `shearnet/config/config_handler.py`)
- Training CLI → Training service: runs training (evidence: `shearnet/cli/train.py`)
- Training service → Galaxy simulation: renders upfront data (evidence: `shearnet/training/service.py`)
- Training service → In-loop generation: selects in-loop (evidence: `shearnet/training/service.py`)
- In-loop generation → Differentiable simulation: generates batches (evidence: `shearnet/core/inloop.py`)
- Training service → Neural models: trains model (evidence: `shearnet/training/service.py`)
- Training service → Run artifacts: writes run (evidence: `shearnet/training/service.py`)
- Training service → Training checkpoints: saves checkpoints (evidence: `shearnet/training/service.py`)
- Pipeline CLI → Slurm scheduler: submits stages (evidence: `README.md`)
- Pipeline CLI → Training CLI: dispatches training (evidence: `shearnet/cli/pipeline.py`)
- Pipeline CLI → Evaluation service: dispatches evaluation (evidence: `shearnet/cli/pipeline.py`)
- Pipeline CLI → Catalog merge: dispatches merge (evidence: `shearnet/cli/pipeline.py`)
- Evaluation CLI → Config loading: loads config
- Evaluation CLI → Evaluation service: runs estimator (evidence: `README.md`)
- Evaluation service → Observation rendering: renders observations (evidence: `shearnet/evaluation/service.py`)
- Evaluation service → Model prediction: uses prediction (evidence: `shearnet/evaluation/service.py`)
- Evaluation service → Measurement workers: dispatches work (evidence: `shearnet/evaluation/service.py`)
- Measurement workers → NGmix metacalibration: measures metacal (evidence: `shearnet/evaluation/workers.py`)
- Evaluation service → Catalog schema: builds tables (evidence: `shearnet/evaluation/service.py`)
- Evaluation service → FITS catalog: writes catalog (evidence: `shearnet/evaluation/service.py`)
- Evaluation service → Run artifacts: reads run artifacts (evidence: `shearnet/evaluation/service.py`)
- Evaluation service → Recovery controls: guards and resumes (evidence: `shearnet/evaluation/service.py`)
- Researcher → Catalog merge: requests catalog merge (evidence: `README.md`)
- Catalog merge → FITS catalog: reads catalogs (evidence: `shearnet/io/merge.py`)
- Catalog merge → Catalog schema: checks schema (evidence: `shearnet/io/merge.py`)