This is an adapted version of the original Surya source code to work with TerraTorch
☀️ Surya: Foundation Model for Heliophysics ☀️
The first foundation model for heliophysics trained on full-resolution Solar Dynamics Observatory data
📖 Overview
Surya (Sanskrit for "Sun") is a 366M-parameter foundation model for heliophysics, trained on full-resolution multi-instrument SDO observations (AIA & HMI). It learns general-purpose solar representations through spatiotemporal transformers, enabling state-of-the-art performance in solar flare forecasting, active region segmentation, solar wind prediction, and EUV spectra modeling.
Key Features
- Multi-instrument Learning: Trained on 13 channels from SDO's AIA (8 channels) and HMI (5 channels) instruments
- Full Resolution: Native 4096×4096 pixel resolution with 12-minute cadence
- Novel Architecture: Spatiotemporal transformer with spectral gating and long-short range attention
- Zero-shot Capabilities: Forecasts solar dynamics and flare events without additional training
- Versatile Fine-tuning: Parameter-efficient LoRA adaptation for diverse downstream tasks
What Makes Surya Special?
Unlike traditional task-specific models, Surya learns physics-aware representations that generalize across multiple solar phenomena:
- Solar Flare Forecasting
- Active Region Segmentation
- Solar Wind Prediction
- EUV Spectra Modeling
🚀 Quick Start
Prerequisites
- Python 3.11+
- CUDA-capable GPU (recommended)
- uv package manager (recommended)
🛠️ Installation
- Clone the repository
git clone https://github.com/NASA-IMPACT/Surya.git
cd Surya
- Install uv package manager (optional)
curl -LsSf https://astral.sh/uv/install.sh | sh
source ~/.bashrc
- Set up the environment
uv sync
source .venv/bin/activate
🧪 Verify Installation
Run the end-to-end test to ensure everything is working:
python -m pytest -s -o log_cli=true tests/test_surya.py
This will:
- Download the pretrained model and test data
- Generate 2-hour ahead forecasts for 2014-01-07
- Create a validation visualization (
surya_model_validation.png) - Verify model inference
Expected output:
============================= test session starts ==============================
INFO test_surya:test_surya.py:188 GPU detected. Running the test on device 0.
INFO test_surya:test_surya.py:195 Surya FM: 366.19 M total parameters.
INFO test_surya:test_surya.py:199 Loaded weights.
INFO test_surya:test_surya.py:201 Starting inference run.
INFO test_surya:test_surya.py:215 Completed validation run. Local loss 0.31665.
PASSED [100%]
🎯 Downstream Applications
To download the Surya model and a sample dataset for downstream tasks, please follow these steps:
# Step 1: Run pytest to download the model and verify dependencies
python -m pytest -s -o log_cli=true tests/test_surya.py
# Step 2: Navigate to the downstream examples
cd downstream_examples/
# Step 3: Download the sample dataset
python download_data.py
1. Solar Flare Forecasting
Predict M-class and X-class solar flares up to 24 hours in advance.
cd downstream_examples/solar_flare_forcasting
python3 download_data.sh
torchrun --nnodes=1 --nproc_per_node=1 --standalone finetune.py
2. Active Region Segmentation
Segment solar active regions and polarity inversion lines from magnetograms.
cd downstream_examples/ar_segmentation
python3 download_data.sh
torchrun --nnodes=1 --nproc_per_node=1 --standalone finetune.py
3. Solar Wind Forecasting
Predict solar wind speed at L1 point with 4-day lead time.
cd downstream_examples/solar_wind_forcasting
python3 download_data.sh
torchrun --nnodes=1 --nproc_per_node=1 --standalone finetune.py
4. EUV Spectra Modeling
Model extreme ultraviolet irradiance across 1343 spectral bands (5-35 nm).
cd downstream_examples/euv_spectra_prediction
python3 download_data.sh
torchrun --nnodes=1 --nproc_per_node=1 --standalone finetune.py
📥 Data and Model Access
Pretrained Models
The Surya foundation model and datasets are available on HuggingFace 🤗 :
- Model Repository:
nasa-ibm-ai4science/Surya-1.0 - Dataset Repository:
nasa-ibm-ai4science/core-sdo
SDO Data Download
For downstream applications, download the preprocessed SDO data:
cd downstream_examples
python download_data.py
This will:
- Download data from HuggingFace repository
- Extract and organize validation/test datasets
- Generate CSV index files for each downstream task
- Set up data in the expected directory structure
📊 Model Architecture
Surya employs a novel spatiotemporal transformer architecture optimized for solar dynamics:
Core Components
-
Spectral Gating Blocks (2 layers)
- Frequency-domain filtering with learnable complex weights
- Adaptive re-weighting of spectral components
- Noise suppression and feature enhancement
-
Long-Short Attention Blocks (8 layers)
- Local attention: Fine-scale dependencies within spatial windows
- Global attention: Long-range correlations via dynamic projection
- Multi-scale representation learning
-
Decoder Block
- Lightweight projection back to physical domain
- Maintains spatial structure and channel relationships
Training Strategy
- Phase 1: One-step ahead forecasting (160k steps, 128 GPUs)
- Phase 2: Autoregressive rollout tuning (2-5 hour horizons)
- Objective: Mean Squared Error with signum-log normalization
- Data: 2011-2019 SDO observations (~257TB processed)
Data Processing Pipeline
Our preprocessing ensures ML-ready, physics-consistent data:
- Temporal alignment: 12-minute cadence across all instruments
- Spatial registration: Uniform 0.6"/pixel grid, solar north alignment
- Calibration: Instrument degradation correction, exposure normalization
- Quality control: Automated flagging and filtering
🏆 Performance Benchmarks
| Task | Metric | Surya | Baseline | Improvement |
|---|---|---|---|---|
| Solar Flare Forecasting | TSS | 0.436 | 0.358 (AlexNet) | 22% |
| Active Region Segmentation | IoU | 0.768 | 0.688 (UNet) | 12% |
| Solar Wind Prediction | RMSE | 75.92 | 93.76 (ResNet50) | 19% |
| EUV Spectra Modeling | MAPE | 1.48% | 1.68% (AlexNet) | 12% |
📄 Citation
If you use Surya in your research, please cite our paper:
@article{roy2025surya,
title={Surya: Foundation Model for Heliophysics},
author={Roy, Sujit and Schmude, Johannes and Lal, Rohit and Gaur, Vishal and Freitag, Marcus and Kuehnert, Julian and van Kessel, Theodore and Hegde, Dinesha V and Mu{\~n}oz-Jaramillo, Andr{\'e}s and Jakubik, Johannes and others},
journal={arXiv preprint arXiv:2508.14112},
year={2025}
}
📜 License
This project is licensed under the Apache License 2.0. See the LICENSE file for details.
🤝 Contributing
We welcome contributions to the Surya repository! Please see our contribution guidelines and feel free to:
- 🐛 Report bugs and issues
- 💡 Suggest new features or applications
- 🔧 Submit pull requests for improvements
- 📖 Improve documentation and examples
Metadata
Release files for terratorch-surya 0.1.0
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| terratorch_surya-0.1.0.tar.gz | 6.4 MB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| terratorch_surya-0.1.0-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 6.5 MB
Release files / terratorch_surya-0.1.0.tar.gz
| Download URL | terratorch_surya-0.1.0.tar.gz |
|---|---|
| Size | 6.4 MB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
83726ef049a9dc704f2968c3ced3d01acb8f53825719a37a60537f51fd7f9959
|
|
BLAKE2b-256 checksum How to use checksums |
94eeeab12458248e968d17653b7ddd29255a1243c2c8018689ca1b208633c569
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/6.2.0 CPython/3.13.9
|
Release files / terratorch_surya-0.1.0-py3-none-any.whl
| Download URL | terratorch_surya-0.1.0-py3-none-any.whl |
|---|---|
| Size | 154.7 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
07df8f4ca629a8d3a631ded5f795ef2765a6a8669af5b3830bba7b8df7c51c6f
|
|
BLAKE2b-256 checksum How to use checksums |
0f6dcfb204cc2ab969323df5ef8f7221c37a173c900443888550bc34fd9b0b19
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/6.2.0 CPython/3.13.9
|