README.md

---
license: apache-2.0
language:
- en
tags:
- OneScience
- Earth Science
- Heliophysics
- Sun-Earth Space Forecasting
- Spatiotemporal Foundation Model
- Spectral Gating
- Long-Short Attention
frameworks: PyTorch
datasets:
- SDO/AIA
- SDO/HMI
---

<p align="center">
  <strong><span style="font-size: 30px;">Surya</span></strong>
</p>

# Model Introduction

Surya is a spatiotemporal foundation model for heliophysics and space weather research. It learns solar dynamics from multi-instrument SDO observations. The model combines spectral gating with long-short attention to predict future solar images from two historical time steps and supports autoregressive multi-step forecasting.

Paper: Surya: Foundation Model for Heliophysics  
https://ntrs.nasa.gov/citations/20250008498

# Model Description

Surya was proposed by institutions including NASA, IBM Research, and the University of Alabama in Huntsville. The model uses eight AIA channels and five HMI products from SDO as training data. It is suitable for one-step prediction and autoregressive multi-step forecasting of solar images.

# Use Cases

| Scenario | Description |
| :---: | :--- |
| Solar dynamics forecasting | Predict future multichannel solar images from two historical time steps. |
| Space weather research | Analyze solar activity and errors at different forecast lead times. |
| SDO multi-instrument modeling | Jointly process multichannel observations from AIA and HMI. |
| Solar flare forecasting | Fine-tune model representations to predict M-class and X-class solar flares within future time periods. |
| Solar active region segmentation | Fine-tune on magnetograms to segment solar active regions and polarity inversion lines. |
| Local quick validation | Use synthetic data to validate training, inference, evaluation, and visualization. |
| Multi-GPU training | Launch distributed data-parallel training with `torchrun`. |

# Usage Guide

## 1. OneCode Usage

Experience intelligent one-click AI4S programming through the OneCode online environment:

[Click to Experience Intelligent One-Click AI4S Programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)

## 2. Manual Installation and Usage

**Hardware Requirements**

- A GPU or DCU is recommended.
- CPU can be used for pipeline validation with the current default small configuration.
- Training the original paper configuration with `4096x4096` resolution and 366 million parameters requires large-scale multi-GPU computing resources.

### Download the Model Package

```bash
hf download OneScience-Group/Surya --local-dir ./Surya
cd Surya
```

### Install the Runtime Environment

**DCU Environment**

```bash
conda create -n onescience311 python=3.11 -y
conda activate onescience311
pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
```

**GPU Environment**

```bash
conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
conda activate onescience311
pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
```

### Training Data Introduction

The paper uses SDO data from 2010 to 2024, including AIA 94, 131, 171, 193, 211, 304, 335, and 1600 Å, as well as HMI magnetic-field and Doppler-velocity products. The data are standardized to a 12-minute interval and spatially registered. The original training data are approximately 257 TB.

This model package uses synthetic data by default:

```bash
python scripts/fake_data.py
```

Each NPZ file contains at least `inputs`, `targets`, and `activity`. `inputs` has shape `[N, 2, 13, H, W]`, `targets` has shape `[N, S, 13, H, W]`, and `activity` stores the solar activity intensity at each forecast time step. When using real data, prepare files with the same fields and modify `conf/config.yaml`.

When using real data, do not run `fake_data.py`. Save temporally aligned, spatially registered, channel-organized, and signum-log normalized data to `data/`, replacing the synthetic data files:

```text
data/train.npz
data/test.npz
```

Each NPZ file contains at least:

```text
inputs:   float32 [N,2,13,H,W]
targets:  float32 [N,S,13,H,W]
activity: float32 [N,S]
```

The 13 channels correspond in order to eight AIA channels and five HMI products, and `S` is consistent with `forecast_steps` in `conf/config.yaml`. Modify `conf/config.yaml` according to the image size, number of forecast steps, channel statistics, and data paths of the real data. After data preparation, continue to use the unified training, inference, and evaluation commands below; use script arguments to override the default paths if other file locations are required.

### Training

Single GPU:

```bash
python scripts/train.py
```

Multiple GPUs:

```bash
torchrun --nproc_per_node=8 --nnodes=1 --rdzv_id=1000 --rdzv_backend=c10d --max_restarts=0 --master_addr="localhost" --master_port=29500 scripts/train.py
```

Training outputs:

```text
result/checkpoints/surya.pt
result/training/metrics.json
```

The training outputs include a model checkpoint for subsequent multi-step forecasting and training metrics that reflect different training stages, overall loss, and learning-rate changes, facilitating training-state preservation and model optimization analysis. Results on synthetic data are only for code-flow validation and do not represent paper results.

### Training Weights

This repository provides weights trained on SDO/AIA and SDO/HMI solar observations in the `weight/` folder. The weight files will be uploaded soon and are expected to be available in the near future.

### Inference

```bash
python scripts/inference.py
```

Inference results are output to:

```text
result/output/forecast.npz
```

### Evaluation and Visualization

```bash
python scripts/result.py
```

Evaluation and visualization outputs are written to:

```text
result/evaluation/metrics.json
result/evaluation/rollout_forecast_skill.png
result/evaluation/solar_dynamics_forecast.png
result/evaluation/solar_activity_evolution.png
result/evaluation/sdo_channel_error.png
result/evaluation/persistence_skill.png
```

The evaluation results comprehensively measure multi-step forecast errors, forecast skill relative to a persistence baseline, performance across AIA/HMI channels, and changes in solar activity intensity. Forecast images and lead-time curves show the evolution of solar dynamics and the long-term forecast stability of the model. The results follow the autoregressive forecasting protocol from the paper; results on synthetic data are only for validating the Surya workflow and do not represent performance on the paper's downstream tasks.

# Official OneScience Resources

| Platform | OneScience Main Repository | Skills Repository |
| --- | --- | --- |
| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |

# Citation and License

This repository is a reproduction of the original Surya paper.