# sreshwarprasad/hierarchical-gz2 architecture

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

- Interactive diagram: https://gitdiagram.com/sreshwarprasad/hierarchical-gz2
- Repository: https://github.com/sreshwarprasad/hierarchical-gz2

## Overview

#### Purpose and entry points
This is an offline ML research pipeline for probabilistic Galaxy Zoo 2 morphology prediction from galaxy images. It trains image models to emit per-question vote distributions rather than hard labels, then evaluates distribution validity, accuracy, calibration, uncertainty, and attention-based interpretability.

Primary execution surfaces are training in `src/train.py` and `src/train_single.py`, experiment orchestration in `src/ablation.py`, and shell launch integration in `scripts/training/run_model.sh`.

#### Data preparation and dataset boundary
Raw Galaxy Zoo/image assets are converted, labeled, and validated before model training. The data boundary is handled by `scripts/data/convert_zenodo_images.py`, `scripts/data/prepare_labels.py`, and `src/dataset.py`; dataset verification is available in `scripts/data/verify_dataset.py`.

The dataset loader supplies transformed image tensors and hierarchical vote-distribution targets to the training/evaluation stack. External dataset/model hosting is Hugging Face; source imagery is associated with Galaxy Zoo 2.

#### Hierarchical probabilistic model
The core predictor is a Vision Transformer-oriented regression model in `src/model.py`. Its architectural contract is to produce probability distributions over answer branches for Galaxy Zoo questions, respecting each question’s simplex constraint rather than independent unconstrained scalar regression.

Alternative modeling and baseline paths are centralized in `src/baselines.py`, enabling comparison with ResNet, ViT MSE-only, and Dirichlet-style approaches. Experiment variants are selected through YAML configurations such as `configs/base.yaml`, `configs/full_train.yaml`, and `configs/ablation.yaml`.

#### Optimization and experiment execution
Training consumes batches from `src/dataset.py`, forwards them through `src/model.py`, and applies the composite probabilistic objective in `src/loss.py`: KL divergence plus MSE. `src/train.py` is the main multi-stage training workflow, while `src/train_single.py` supports a single configured run.

Ablation control in `src/ablation.py` varies important training choices, notably loss weighting/configuration. Run configuration is externalized in `configs/`, keeping experiment parameters separate from model and loss implementation.

#### Evaluation, validity, and calibration
Full evaluation is orchestrated by `src/evaluate_full.py`, using shared calculations from `src/metrics.py`. This boundary measures MAE/RMSE and probability-quality properties, including calibration and simplex violations—critical because valid per-question distributions are a core model guarantee.

The evaluation path compares proposed and baseline models, with generated manuscript artifacts retained under `outputs/manuscript/03_baseline_comparison/` and `outputs/manuscript/05_calibration/`.

#### Interpretability and uncertainty stages
Attention inspection is a distinct post-training analysis stage in `src/attention_viz.py`, producing ViT attention rollout/entropy artifacts represented in `outputs/manuscript/06_interpretability/`.

Uncertainty analysis is implemented in `src/uncertainty_analysis.py`; its outputs relate predicted uncertainty to prediction error and are published under `outputs/manuscript/07_uncertainty/`. These stages consume trained-model predictions rather than altering the training path.

#### Output collection and publication artifacts
`scripts/training/collect_outputs.py` consolidates experiment results. Analysis scripts such as `scripts/analysis/generate_comparison_tables.py` and `scripts/analysis/plot_training_curves.py` transform collected metrics/checkpoints into tables and figures.

`outputs/manuscript/` is the committed research-result boundary, with `outputs/manuscript/MANIFEST.csv` indexing final artifacts. The repository itself does not contain the described web demo; the deployed demo is an external Vercel service.

Main entry points: [`src/evaluate_full.py`](https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/src/evaluate_full.py) (Full evaluation), [`src/train.py`](https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/src/train.py) (Multi-stage training), [`src/train_single.py`](https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/src/train_single.py) (Single-run training), [`src/dataset.py`](https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/src/dataset.py) (Dataset loader)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_data["Data preparation"]
  node_galaxy_zoo[("Galaxy Zoo 2 imagery<br/>external dataset")]
  node_huggingface[("Hugging Face hosting<br/>external hosting")]
  node_convert_images["Image conversion<br/>data script"]
  node_prepare_labels["Label preparation<br/>data script<br/>[prepare_labels.py]"]
  node_dataset_verify["Dataset verification<br/>data validation<br/>[verify_dataset.py]"]
  node_dataset["Dataset loader<br/>PyTorch dataset<br/>[dataset.py]"]
end

subgraph group_training["Model experiments"]
  node_configs["Experiment configs<br/>YAML configuration<br/>[base.yaml]"]
  node_model["Hierarchical ViT predictor<br/>probabilistic model<br/>[model.py]"]
  node_baselines["Baseline models<br/>model variants<br/>[baselines.py]"]
  node_loss["Composite objective<br/>training loss<br/>[loss.py]"]
  node_train["Multi-stage training<br/>training workflow<br/>[train.py]"]
  node_train_single["Single-run training<br/>training workflow<br/>[train_single.py]"]
  node_ablation["Ablation orchestrator<br/>experiment control<br/>[ablation.py]"]
  node_run_model["Shell launcher<br/>launch integration<br/>[run_model.sh]"]
end

subgraph group_analysis["Evaluation and publication"]
  node_evaluate["Full evaluation<br/>evaluation workflow<br/>[evaluate_full.py]"]
  node_metrics["Probability metrics<br/>metrics library<br/>[metrics.py]"]
  node_attention["Attention analysis<br/>interpretability stage<br/>[attention_viz.py]"]
  node_uncertainty["Uncertainty analysis<br/>post-training analysis"]
  node_collect["Result collection<br/>output consolidation<br/>[collect_outputs.py]"]
  node_publication["Manuscript artifacts<br/>published results<br/>[MANIFEST.csv]"]
end

node_galaxy_zoo -->|"source images"| node_convert_images
node_convert_images -->|"prepared assets"| node_prepare_labels
node_prepare_labels -->|"labeled dataset"| node_dataset_verify
node_dataset_verify -->|"validated data"| node_dataset
node_huggingface -.->|"hosted assets"| node_dataset
node_configs -->|"parameters"| node_train
node_configs -->|"run configuration"| node_train_single
node_configs -->|"variant definitions"| node_ablation
node_run_model -->|"launches"| node_train
node_ablation -->|"configured runs"| node_train_single
node_dataset -->|"image batches and targets"| node_train
node_model -->|"forward model"| node_train
node_loss -->|"optimization objective"| node_train
node_dataset -->|"image batches and targets"| node_train_single
node_model -->|"forward model"| node_train_single
node_loss -->|"optimization objective"| node_train_single
node_train -->|"trained predictions"| node_evaluate
node_train_single -->|"trained predictions"| node_evaluate
node_baselines -->|"comparison predictions"| node_evaluate
node_metrics -->|"shared calculations"| node_evaluate
node_evaluate -->|"trained model outputs"| node_attention
node_evaluate -->|"predictions and errors"| node_uncertainty
node_evaluate -->|"metrics and checkpoints"| node_collect
node_attention -->|"interpretability artifacts"| node_collect
node_uncertainty -->|"uncertainty artifacts"| node_collect
node_collect -->|"tables and figures"| node_publication

click node_convert_images "https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/scripts/data/convert_zenodo_images.py"
click node_prepare_labels "https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/scripts/data/prepare_labels.py"
click node_dataset_verify "https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/scripts/data/verify_dataset.py"
click node_dataset "https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/src/dataset.py"
click node_configs "https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/configs/base.yaml"
click node_model "https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/src/model.py"
click node_baselines "https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/src/baselines.py"
click node_loss "https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/src/loss.py"
click node_train "https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/src/train.py"
click node_train_single "https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/src/train_single.py"
click node_ablation "https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/src/ablation.py"
click node_run_model "https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/scripts/training/run_model.sh"
click node_evaluate "https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/src/evaluate_full.py"
click node_metrics "https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/src/metrics.py"
click node_attention "https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/src/attention_viz.py"
click node_uncertainty "https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/src/uncertainty_analysis.py"
click node_collect "https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/scripts/training/collect_outputs.py"
click node_publication "https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/outputs/manuscript/MANIFEST.csv"

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_galaxy_zoo,node_huggingface,node_convert_images,node_prepare_labels,node_dataset_verify,node_dataset toneBlue
class node_configs,node_model,node_baselines,node_loss,node_train,node_train_single,node_ablation,node_run_model toneAmber
class node_evaluate,node_metrics,node_attention,node_uncertainty,node_collect,node_publication toneMint
```

## Components

### Data preparation

- **Galaxy Zoo 2 imagery**: external dataset
- **Hugging Face hosting**: external hosting · Hosts dataset and model assets.
- [Image conversion](https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/scripts/data/convert_zenodo_images.py): data script · `scripts/data/convert_zenodo_images.py`
- [Label preparation](https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/scripts/data/prepare_labels.py): data script · Builds hierarchical vote-distribution targets. · `scripts/data/prepare_labels.py`
- [Dataset verification](https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/scripts/data/verify_dataset.py): data validation · `scripts/data/verify_dataset.py`
- [Dataset loader](https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/src/dataset.py): PyTorch dataset · Supplies transformed images and hierarchical targets. · `src/dataset.py`

### Model experiments

- [Experiment configs](https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/configs/base.yaml): YAML configuration · Selects training and ablation variants. · `configs/base.yaml`
- [Hierarchical ViT predictor](https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/src/model.py): probabilistic model · Emits simplex-constrained answer distributions per question. · `src/model.py`
- [Baseline models](https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/src/baselines.py): model variants · ResNet, MSE-only ViT, and Dirichlet-style comparisons. · `src/baselines.py`
- [Composite objective](https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/src/loss.py): training loss · Combines KL divergence and MSE. · `src/loss.py`
- [Multi-stage training](https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/src/train.py): training workflow · `src/train.py`
- [Single-run training](https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/src/train_single.py): training workflow · `src/train_single.py`
- [Ablation orchestrator](https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/src/ablation.py): experiment control · Varies loss weighting and training choices. · `src/ablation.py`
- [Shell launcher](https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/scripts/training/run_model.sh): launch integration · `scripts/training/run_model.sh`

### Evaluation and publication

- [Full evaluation](https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/src/evaluate_full.py): evaluation workflow · Evaluates validity, accuracy, calibration, and model comparisons. · `src/evaluate_full.py`
- [Probability metrics](https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/src/metrics.py): metrics library · Calculates errors, calibration, and simplex violations. · `src/metrics.py`
- [Attention analysis](https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/src/attention_viz.py): interpretability stage · Produces ViT attention rollout and entropy artifacts. · `src/attention_viz.py`
- [Uncertainty analysis](https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/src/uncertainty_analysis.py): post-training analysis · Relates predicted uncertainty to prediction error. · `src/uncertainty_analysis.py`
- [Result collection](https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/scripts/training/collect_outputs.py): output consolidation · `scripts/training/collect_outputs.py`
- [Manuscript artifacts](https://github.com/sreshwarprasad/hierarchical-gz2/blob/main/outputs/manuscript/MANIFEST.csv): published results · Manifest-indexed committed research tables and figures. · `outputs/manuscript/MANIFEST.csv`

## Connections

- Galaxy Zoo 2 imagery → Image conversion: source images
- Image conversion → Label preparation: prepared assets
- Label preparation → Dataset verification: labeled dataset
- Dataset verification → Dataset loader: validated data
- Hugging Face hosting → Dataset loader: hosted assets
- Experiment configs → Multi-stage training: parameters
- Experiment configs → Single-run training: run configuration
- Experiment configs → Ablation orchestrator: variant definitions
- Shell launcher → Multi-stage training: launches
- Ablation orchestrator → Single-run training: configured runs
- Dataset loader → Multi-stage training: image batches and targets
- Hierarchical ViT predictor → Multi-stage training: forward model
- Composite objective → Multi-stage training: optimization objective
- Dataset loader → Single-run training: image batches and targets
- Hierarchical ViT predictor → Single-run training: forward model
- Composite objective → Single-run training: optimization objective
- Multi-stage training → Full evaluation: trained predictions
- Single-run training → Full evaluation: trained predictions
- Baseline models → Full evaluation: comparison predictions
- Probability metrics → Full evaluation: shared calculations
- Full evaluation → Attention analysis: trained model outputs
- Full evaluation → Uncertainty analysis: predictions and errors
- Full evaluation → Result collection: metrics and checkpoints
- Attention analysis → Result collection: interpretability artifacts
- Uncertainty analysis → Result collection: uncertainty artifacts
- Result collection → Manuscript artifacts: tables and figures