README.md

---
license: mit
library_name: pytorch
pipeline_tag: image-to-image
tags:
  - solar-physics
  - solar-forecasting
  - swin-transformer
  - pytorch
  - aia
  - patch-size-4
datasets:
  - hrrsmjd/AIA_12hour_512x512
---

# Solaris Small Patch 4

This repository contains a Solaris-Small checkpoint trained for 12-hour multi-wavelength solar forecasting, following the Solaris pretraining setup from [Solaris: A Foundation Model of the Sun](https://arxiv.org/abs/2411.16339).

This run uses patch size 4. The earlier patch-size-8 checkpoint is published separately as `hrrsmjd/solaris_small_patch8`.

The checkpoint was trained on `hrrsmjd/AIA_12hour_512x512` for 7750 optimizer steps using two history frames (`t-12h`, `t`) to predict all eight pretraining wavelengths at `t+12h`.

## Files

- `solaris_small_patch4_model_state_dict.pt`: reusable PyTorch checkpoint containing `model_state_dict`, learned normalization coefficients, wavelengths, scale factors, patch size, seed, training step, and final training loss.
- `config.json`: lightweight metadata for reconstructing the model and normalization.
- `assets/solaris_small_patch4_test0_prediction.png`: example qualitative test prediction plot.
- `eval/solaris_pretrain_paperloss_p4_ema_seed42_test_mse_subset_0352.md`: full test-split raw-scale MSE/RMSE/MAE report.

## Example Plot

The plot below shows one test sample with rows for `input t-12h`, `input t`, `target t+12h`, `prediction t+12h`, and `prediction - target` across all eight wavelengths.

![Solaris Small Patch 4 test prediction](assets/solaris_small_patch4_test0_prediction.png)

## Model Details

- Architecture: `SolarisSmall`
- Patch size: `4`
- Embedding dimension: `256`
- Encoder depths: `(2, 6, 2)`
- Decoder depths: `(2, 6, 2)`
- Output wavelengths: `0094`, `0131`, `0171`, `0193`, `0211`, `0304`, `0335`, `1600`
- Training dataset: `hrrsmjd/AIA_12hour_512x512`
- Training target: 12-hour forecast
- Training budget: 7750 optimizer steps, batch size 8, gradient accumulation 4
- Seed: `42`

## Loading

```python
import torch

from solaris.model.solaris import SolarisSmall

checkpoint = torch.load("solaris_small_patch4_model_state_dict.pt", map_location="cpu", weights_only=False)

model = SolarisSmall(
    out_levels=len(checkpoint["wavelengths"]),
    patch_size=checkpoint["patch_size"],
)
model.load_state_dict(checkpoint["model_state_dict"])
model.eval()

scale_factors = torch.tensor(checkpoint["scale_factors"], dtype=torch.float32)
norm_coeff_1 = checkpoint["norm_coeff_1"]
norm_coeff_2 = checkpoint["norm_coeff_2"]
```

Inputs should be normalized with the Solaris transform used during training. Model outputs are normalized intensities; multiply by the per-wavelength scale factors before comparing to raw-intensity targets.

## Test Metrics

Metrics below use all 352 test samples and are computed on the raw intensity scale. Regular final weights are recommended; EMA weights from the training checkpoint were worse on the full test split and are not included in this model-state checkpoint.

| Wavelength (A) | MSE | RMSE | MAE |
|---:|---:|---:|---:|
| 0094 | 8.87581 | 2.97923 | 0.240353 |
| 0131 | 243.534 | 15.6056 | 1.446 |
| 0171 | 9067.35 | 95.2226 | 37.4301 |
| 0193 | 18337.8 | 135.417 | 56.6422 |
| 0211 | 3811.31 | 61.7358 | 23.776 |
| 0304 | 1089.31 | 33.0047 | 12.0236 |
| 0335 | 31.7769 | 5.6371 | 1.63412 |
| 1600 | 54.0763 | 7.35366 | 3.33375 |
| **Mean** | **4080.5** | **44.6195** | **17.0658** |

## Training Notes

Scale factors were computed as half the average per-image maximum over unique train-split timestamps:

```text
[58.224720422037755, 216.21549287451052, 1616.446579054541, 2551.0149615718674, 1190.0182024885178, 887.1800787601859, 112.33733897339224, 266.61844876445224]
```

The final logged training-batch metrics at step 7750 were:

```text
weighted MAE: 0.009133
mean raw RMSE: 25.278
per-wavelength raw RMSE: [2.535, 4.554, 62.817, 67.771, 28.201, 22.793, 2.572, 10.982]
```