How nasa-impact/surya works
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: helio_spectformer.py Helio Spectformer · downstream_examples Downstream examples · run_easy_inference.py Inference runner · datasets Data layer
Components 20
External sources
Model and data dependencies outside the repository.
- Hugging Face · Hosts pretrained weights and benchmark datasets for Surya.
- S3 bench · Public bucket used by easy inference to fetch input .nc files.
Surya core
Reusable foundation-model library.
- Model stack · Heliophysics backbone and supporting architecture modules.
surya/models - Helio Spectformer · Main pretrained model entry point for heliophysics forecasting.
surya/models/helio_spectformer.py - Spectformer · Core transformer building block used by the backbone.
surya/models/spectformer.py - Transformer LS · Specialized transformer variant in the model stack.
surya/models/transformer_ls.py - Embeddings · Input embedding and representation utilities for model inputs.
surya/models/embedding.py - Flow head · Output flow or decoding component used by the model stack.
surya/models/flow.py - Data layer · Canonical dataset and preprocessing boundary for solar observations.
surya/datasets - Helio dataset · Domain-specific dataset interface for model training and inference.
surya/datasets/helio.py - Transforms · Normalization, augmentation, and tensor preparation logic.
surya/datasets/transformations.py - Runtime utils · Configuration, data, distributed, logging, and schema helpers.
surya/utils
Application layers
Zero-shot inference and downstream task wrappers.
- Easy inference · End-to-end inference path for prompting, downloading inputs, and writing predictions.
easy_inference - Inference runner · Loads config, resolves device, fetches inputs, runs inference, and saves outputs.
easy_inference/run_easy_inference.py - Downstream examples · Reusable pattern for task-specific fine-tuning and inference examples.
downstream_examples - Flare task · Solar flare forecasting example with metrics and task-specific head.
downstream_examples/solar_flare_forcasting - AR task · Active-region segmentation example including CSV indexing and task head.
downstream_examples/ar_segmentation - Wind task · Solar wind forecasting example built on the shared backbone.
downstream_examples/solar_wind_forcasting - EUV task · 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 · End-to-end validation that exercises model loading, data access, and inference.
tests/test_surya.py
Connections 25
- 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.