# nasa-impact/surya architecture

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

- Interactive diagram: https://gitdiagram.com/nasa-impact/surya
- Repository: https://github.com/nasa-impact/surya

## Overview

1. **Repository shape**
   - This is a Python 3.11+ PyTorch repository organized around a reusable **core foundation-model library** (`surya/`) plus two thin application layers: **zero-shot inference** (`easy_inference/`) and **task-specific downstream examples** (`downstream_examples/`).

2. **Core model boundary**
   - `surya/models/` contains the actual heliophysics foundation model stack.
   - The main entry points are `helio_spectformer.py`, `spectformer.py`, `transformer_ls.py`, `embedding.py`, and `flow.py`, which together implement the spatiotemporal transformer architecture described in the README.
   - Architecturally, this is the “pretrained backbone” layer that downstream scripts wrap rather than reimplement.

3. **Core data boundary**
   - `surya/datasets/` holds the canonical data handling for model training/inference.
   - `helio.py` is the domain dataset interface, while `transformations.py` centralizes preprocessing/augmentation or normalization logic.
   - This layer is the main boundary between raw solar observations and model-ready tensors.

4. **Shared runtime utilities**
   - `surya/utils/` provides cross-cutting infrastructure: configuration handling, data helpers, distributed execution, logging, miscellaneous helpers, and schemas.
   - These utilities support both standalone inference and distributed fine-tuning, and they keep orchestration concerns out of the model code.

5. **Zero-shot inference application**
   - `easy_inference/` is the simplest end-to-end path for using the foundation model.
   - `run_easy_inference.py` consumes `config_easy.yaml`, prompts for a UTC date range, downloads required `.nc` inputs from the public `nasa-surya-bench` S3 bucket, runs inference, and writes predictions plus metrics under an output directory.
   - The device-selection policy is explicit in config: `cuda -> mps -> cpu`.

6. **Downstream application pattern**
   - Each folder under `downstream_examples/` follows the same mini-architecture: `dataset.py`, `models.py`, `finetune.py`, `infer.py`, `config.yaml`, and `config_infer.yaml`, plus a task README and download script.
   - This indicates a repeated pattern of: load preprocessed data, attach a task head, fine-tune from the pretrained Surya backbone, then run task-specific inference.

7. **Task-specific subsystems**
   - The four downstream tasks are distinct application boundaries on top of the same foundation model:
     - **Solar flare forecasting** (`solar_flare_forcasting/`) with `metrics.py`
     - **Active region segmentation** (`ar_segmentation/`)
     - **Solar wind forecasting** (`solar_wind_forcasting/`)
     - **EUV spectra prediction** (`euv_spectra_prediction/`)
   - These differ mainly in target type and output head, not in overall pipeline structure.

8. **Data preparation and indexing**
   - `downstream_examples/download_sample_train_data.py` and each task’s `download_data.sh` suggest a two-stage data setup: fetch preprocessed assets, then organize them locally for training/inference.
   - `create_ar_csv.py` in active-region segmentation shows that CSV index generation is part of the preprocessing boundary, not the training loop itself.

9. **Model distribution and external services**
   - The pretrained model and datasets are distributed via Hugging Face (`nasa-ibm-ai4science/Surya-1.0` and `nasa-ibm-ai4science/core-sdo`).
   - Easy inference additionally pulls data from a public S3 bucket. These are the main external data/model dependencies that affect runtime architecture.

10. **Validation and CI-style check**
   - `tests/test_surya.py` is an end-to-end smoke test: it downloads weights and test data, runs a forecast, and generates a validation visualization.
   - This test acts as the repository’s integration boundary, confirming that model loading, data access, and inference all work together.

11. **Artifacts and notebooks**
   - `assets/` and per-task `assets/` folders are output/documentation holders, not core runtime code.
   - The notebooks in `easy_inference/` and downstream examples are tutorial surfaces around the same pipelines, useful for understanding flow but not separate subsystems.

12. **How to read the architecture**
   - The cleanest diagram is: **external data/model sources → data prep/indexing → `surya` foundation model → task-specific heads/scripts → predictions/metrics**, with `utils/` and YAML configs orchestrating execution across all entry points.

Main entry points: [`surya/models/helio_spectformer.py`](https://github.com/nasa-impact/surya/blob/main/surya/models/helio_spectformer.py) (Helio Spectformer), [`downstream_examples`](https://github.com/nasa-impact/surya/tree/main/downstream_examples) (Downstream examples), [`easy_inference/run_easy_inference.py`](https://github.com/nasa-impact/surya/blob/main/easy_inference/run_easy_inference.py) (Inference runner), [`surya/datasets`](https://github.com/nasa-impact/surya/tree/main/surya/datasets) (Data layer)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_external["External sources"]
  node_hf_model[("Hugging Face<br/>model/data hub")]
  node_s3_bucket[("S3 bench<br/>object store")]
end

subgraph group_core["Surya core"]
  node_surya_models["Model stack<br/>foundation models"]
  node_helio_model["Helio Spectformer<br/>backbone entry"]
  node_spectformer["Spectformer<br/>transformer core<br/>[spectformer.py]"]
  node_transformer_ls["Transformer LS<br/>sequence module<br/>[transformer_ls.py]"]
  node_embedding["Embeddings<br/>feature encoder<br/>[embedding.py]"]
  node_flow["Flow head<br/>prediction head<br/>[flow.py]"]
  node_surya_data["Data layer<br/>datasets"]
  node_helio_ds["Helio dataset<br/>domain dataset<br/>[helio.py]"]
  node_transformations["Transforms<br/>preprocessing<br/>[transformations.py]"]
  node_surya_utils["Runtime utils<br/>shared infra"]
end

subgraph group_apps["Application layers"]
  node_easy_inference["Easy inference<br/>zero-shot app"]
  node_run_easy["Inference runner<br/>CLI pipeline"]
  node_downstream["Downstream examples<br/>task apps"]
  node_flare_task["Flare task<br/>classification app"]
  node_ar_task["AR task<br/>segmentation app"]
  node_wind_task["Wind task<br/>forecasting app"]
  node_euv_task["EUV task<br/>spectra app"]
end

subgraph group_qa["Validation"]
  node_tests["Smoke test<br/>integration check<br/>[test_surya.py]"]
end

node_hf_model -->|"weights"| node_helio_model
node_hf_model -->|"datasets"| node_downstream
node_s3_bucket -->|"inputs"| node_run_easy
node_surya_models -->|"contains"| node_helio_model
node_helio_model -->|"uses"| node_spectformer
node_helio_model -->|"uses"| node_transformer_ls
node_helio_model -->|"encodes"| node_embedding
node_helio_model -->|"predicts"| node_flow
node_surya_data -->|"contains"| node_helio_ds
node_surya_data -->|"contains"| node_transformations
node_transformations -->|"prepares"| node_helio_ds
node_surya_utils -->|"orchestrates"| node_run_easy
node_surya_utils -->|"orchestrates"| node_downstream
node_easy_inference -->|"entry point"| node_run_easy
node_run_easy -->|"loads"| node_helio_model
node_downstream -->|"includes"| node_flare_task
node_downstream -->|"includes"| node_ar_task
node_downstream -->|"includes"| node_wind_task
node_downstream -->|"includes"| node_euv_task
node_flare_task -->|"fine-tunes"| node_helio_model
node_ar_task -->|"fine-tunes"| node_helio_model
node_wind_task -->|"fine-tunes"| node_helio_model
node_euv_task -->|"fine-tunes"| node_helio_model
node_tests -->|"exercises"| node_run_easy
node_tests -->|"validates"| node_helio_model

click node_surya_models "https://github.com/nasa-impact/surya/tree/main/surya/models"
click node_helio_model "https://github.com/nasa-impact/surya/blob/main/surya/models/helio_spectformer.py"
click node_spectformer "https://github.com/nasa-impact/surya/blob/main/surya/models/spectformer.py"
click node_transformer_ls "https://github.com/nasa-impact/surya/blob/main/surya/models/transformer_ls.py"
click node_embedding "https://github.com/nasa-impact/surya/blob/main/surya/models/embedding.py"
click node_flow "https://github.com/nasa-impact/surya/blob/main/surya/models/flow.py"
click node_surya_data "https://github.com/nasa-impact/surya/tree/main/surya/datasets"
click node_helio_ds "https://github.com/nasa-impact/surya/blob/main/surya/datasets/helio.py"
click node_transformations "https://github.com/nasa-impact/surya/blob/main/surya/datasets/transformations.py"
click node_surya_utils "https://github.com/nasa-impact/surya/tree/main/surya/utils"
click node_easy_inference "https://github.com/nasa-impact/surya/tree/main/easy_inference"
click node_run_easy "https://github.com/nasa-impact/surya/blob/main/easy_inference/run_easy_inference.py"
click node_downstream "https://github.com/nasa-impact/surya/tree/main/downstream_examples"
click node_flare_task "https://github.com/nasa-impact/surya/tree/main/downstream_examples/solar_flare_forcasting"
click node_ar_task "https://github.com/nasa-impact/surya/tree/main/downstream_examples/ar_segmentation"
click node_wind_task "https://github.com/nasa-impact/surya/tree/main/downstream_examples/solar_wind_forcasting"
click node_euv_task "https://github.com/nasa-impact/surya/tree/main/downstream_examples/euv_spectra_prediction"
click node_tests "https://github.com/nasa-impact/surya/blob/main/tests/test_surya.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_hf_model,node_s3_bucket toneBlue
class node_surya_models,node_helio_model,node_spectformer,node_transformer_ls,node_embedding,node_flow,node_surya_data,node_helio_ds,node_transformations,node_surya_utils toneAmber
class node_easy_inference,node_run_easy,node_downstream,node_flare_task,node_ar_task,node_wind_task,node_euv_task toneMint
class node_tests toneRose
```

## Components

### External sources

Model and data dependencies outside the repository.

- **Hugging Face**: model/data hub · Hosts pretrained weights and benchmark datasets for Surya.
- **S3 bench**: object store · Public bucket used by easy inference to fetch input .nc files.

### Surya core

Reusable foundation-model library.

- [Model stack](https://github.com/nasa-impact/surya/tree/main/surya/models): foundation models · Heliophysics backbone and supporting architecture modules. · `surya/models`
- [Helio Spectformer](https://github.com/nasa-impact/surya/blob/main/surya/models/helio_spectformer.py): backbone entry · Main pretrained model entry point for heliophysics forecasting. · `surya/models/helio_spectformer.py`
- [Spectformer](https://github.com/nasa-impact/surya/blob/main/surya/models/spectformer.py): transformer core · Core transformer building block used by the backbone. · `surya/models/spectformer.py`
- [Transformer LS](https://github.com/nasa-impact/surya/blob/main/surya/models/transformer_ls.py): sequence module · Specialized transformer variant in the model stack. · `surya/models/transformer_ls.py`
- [Embeddings](https://github.com/nasa-impact/surya/blob/main/surya/models/embedding.py): feature encoder · Input embedding and representation utilities for model inputs. · `surya/models/embedding.py`
- [Flow head](https://github.com/nasa-impact/surya/blob/main/surya/models/flow.py): prediction head · Output flow or decoding component used by the model stack. · `surya/models/flow.py`
- [Data layer](https://github.com/nasa-impact/surya/tree/main/surya/datasets): datasets · Canonical dataset and preprocessing boundary for solar observations. · `surya/datasets`
- [Helio dataset](https://github.com/nasa-impact/surya/blob/main/surya/datasets/helio.py): domain dataset · Domain-specific dataset interface for model training and inference. · `surya/datasets/helio.py`
- [Transforms](https://github.com/nasa-impact/surya/blob/main/surya/datasets/transformations.py): preprocessing · Normalization, augmentation, and tensor preparation logic. · `surya/datasets/transformations.py`
- [Runtime utils](https://github.com/nasa-impact/surya/tree/main/surya/utils): shared infra · Configuration, data, distributed, logging, and schema helpers. · `surya/utils`

### Application layers

Zero-shot inference and downstream task wrappers.

- [Easy inference](https://github.com/nasa-impact/surya/tree/main/easy_inference): zero-shot app · End-to-end inference path for prompting, downloading inputs, and writing predictions. · `easy_inference`
- [Inference runner](https://github.com/nasa-impact/surya/blob/main/easy_inference/run_easy_inference.py): CLI pipeline · Loads config, resolves device, fetches inputs, runs inference, and saves outputs. · `easy_inference/run_easy_inference.py`
- [Downstream examples](https://github.com/nasa-impact/surya/tree/main/downstream_examples): task apps · Reusable pattern for task-specific fine-tuning and inference examples. · `downstream_examples`
- [Flare task](https://github.com/nasa-impact/surya/tree/main/downstream_examples/solar_flare_forcasting): classification app · Solar flare forecasting example with metrics and task-specific head. · `downstream_examples/solar_flare_forcasting`
- [AR task](https://github.com/nasa-impact/surya/tree/main/downstream_examples/ar_segmentation): segmentation app · Active-region segmentation example including CSV indexing and task head. · `downstream_examples/ar_segmentation`
- [Wind task](https://github.com/nasa-impact/surya/tree/main/downstream_examples/solar_wind_forcasting): forecasting app · Solar wind forecasting example built on the shared backbone. · `downstream_examples/solar_wind_forcasting`
- [EUV task](https://github.com/nasa-impact/surya/tree/main/downstream_examples/euv_spectra_prediction): spectra app · EUV spectra prediction example with its own model head and pipeline. · `downstream_examples/euv_spectra_prediction`

### Validation

Repository-level smoke and integration checks.

- [Smoke test](https://github.com/nasa-impact/surya/blob/main/tests/test_surya.py): integration check · End-to-end validation that exercises model loading, data access, and inference. · `tests/test_surya.py`

## Connections

- Hugging Face → Helio Spectformer: weights · Pretrained foundation model weights are sourced externally.
- Hugging Face → Downstream examples: datasets · Benchmark and training assets are distributed via Hugging Face.
- S3 bench → Inference runner: inputs · Easy inference downloads required .nc files from public storage.
- Model stack → Helio Spectformer: contains · Model stack is organized around the heliophysics backbone entry point.
- Helio Spectformer → Spectformer: uses · Backbone composition relies on the core transformer module.
- Helio Spectformer → Transformer LS: uses · Backbone includes specialized sequence modeling components.
- Helio Spectformer → Embeddings: encodes · Input representations are prepared through shared embedding logic.
- Helio Spectformer → Flow head: predicts · Prediction or decoding is handled by a flow-style output component.
- Data layer → Helio dataset: contains · Domain dataset interface is the main entry into the data layer.
- Data layer → Transforms: contains · Preprocessing and augmentation are centralized in transformations.
- Transforms → Helio dataset: prepares · Transforms convert raw observations into model-ready tensors.
- Runtime utils → Inference runner: orchestrates · Shared runtime utilities support configuration, logging, and execution.
- Runtime utils → Downstream examples: orchestrates · Downstream tasks reuse shared config, distributed, and data helpers.
- Easy inference → Inference runner: entry point · Package contains the zero-shot inference CLI and config surface.
- Inference runner → Helio Spectformer: loads · Inference runner instantiates the pretrained backbone for prediction.
- Downstream examples → Flare task: includes · One of the downstream task implementations.
- Downstream examples → AR task: includes · One of the downstream task implementations.
- Downstream examples → Wind task: includes · One of the downstream task implementations.
- Downstream examples → EUV task: includes · One of the downstream task implementations.
- Flare task → Helio Spectformer: fine-tunes · Task-specific heads wrap the shared Surya backbone.
- AR task → Helio Spectformer: fine-tunes · Task-specific heads wrap the shared Surya backbone.
- Wind task → Helio Spectformer: fine-tunes · Task-specific heads wrap the shared Surya backbone.
- EUV task → Helio Spectformer: fine-tunes · Task-specific heads wrap the shared Surya backbone.
- Smoke test → Inference runner: exercises · Smoke test validates the inference path end to end.
- Smoke test → Helio Spectformer: validates · Smoke test confirms backbone loading and prediction work together.