SDO-FM is a prototype foundation model that integrates data from SDO’s HMI and AMI instruments to encapsulate the Sun's dynamics within an embedding space. Encapsulating the Sun's complex physical interactions in a multi-modal model unlocks the potential for many downstream investigations, lowering the complexity and costs for machine learning in solar physics. Building a small fine-tuned adapter can be done for a fraction of the cost of building an AI model with classical methods.
SDO-FM consists of four key components: an ingestion pipeline to create machine learning ready datasets, model architecture and training, embeddings and fine-tunable models, and downstream fine-tuned applications. SDO-FM utilizes existing datasets (SDOMLv2) and well defined downstream tasks.
The engineering focused on adapting two model architectures to the SDO dataset, and building a benchmarking harness for scientific validation. Training was engineered for a TPU-distributed approach, designed to be compute-agnostic and aligned with industry best practices.
SDO-FM was built using a science-led approach and active involvement of the scientific community. Validation tasks for the prototype as well as future proposed validation tasks were chosen through a series of workshops and feedback from heliophysics researchers, teams building other foundation models in weather prediction, climate science, and earth sciences, and SDO instrument experts. This collaborative approach will help us optimize the model for useful downstream applications and adapt to emerging tools and methodologies.
The model is validated by adapting the embeddings and comparing outcomes against published known results using classical machine learning methods. The four validation tasks used were predicting the Earth’s thermospheric density (F10.7), reconstructing missing channels, autocalibration of the AIA instrument, and the virtualization of the broken MEGS-A instrument (Virtual EVE).
The project will also investigate if the foundation model can either replicate or leverage SDOML (the data product developed in FDL.AI).
├── assets # images for this readme
├── experiments # configuration files for different trials
├── notebooks # visualisation/testing ipynb
├── scripts # entrypoint and highest level executors
├── sdofm # python package
│ ├── ablation # models without backbone integration
│ ├── benchmarks # metrics for comparison
│ ├── datasets # dataloaders/modules
│ ├── finetuning # modules for finetuning
│ ├── models # model components
│ ├── pretraining # modules for pretraining
└── └── visualisation # various graphing utilities| Name | Description | Granularity & Source |
|---|---|---|
| NASA’s Solar Dynamics Observatory (SDO) Pesnell et al. 2012 | Three instruments:
|
4096x4096 12 second cadence: AIA - Lemen et al. 2012. HMI - Hoeksema et al. 2014. 1024 x 2048: EUV - Woods et al. 2012. Downsampled 512x512/0.6, 512x512/0.5 arcsec 6 (AIA) 12 (HMI) minute cadence for machine learning: Galvez et al. 2019 via sdoml.org. |
| Name | Paper |
|---|---|
| Masked Autoencoders Are Scalable Vision Learners | He, Kaiming, et al. "Masked autoencoders are scalable vision learners." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022 (link) |
| Foundation Models for Generalist Geospatial Artificial Intelligence (Prithvi) | Jakubik, Johannes, et al. "Foundation models for generalist geospatial artificial intelligence." arXiv preprint arXiv:2310.18660 (2023) (link) |
| NVAE: A Deep Hierarchical Variational Autoencoder | Vahdat, Arash, and Jan Kautz. "NVAE: A deep hierarchical variational autoencoder." Advances in neural information processing systems 33 (2020): 19667-19679 (link) |
| Name | Paper |
|---|---|
| Multichannel autocalibration for the Atmospheric Imaging Assembly using machine learning | Dos Santos, Luiz FG, et al. "Multichannel autocalibration for the Atmospheric Imaging Assembly using machine learning." Astronomy & Astrophysics 648 (2021): A53 (link) |
| Virtual EVE: a Deep Learning Model for Solar Irradiance Prediction | Indaco, Manuel, et al. "Virtual EVE: a Deep Learning Model for Solar Irradiance Prediction." (2023) (link) |
SDO-FM can be installed locally by directly installing the package in this repository.
pip install -e .To run any task we assume execution inside a container with the image described in the Dockerfile and Hydra configurations, these are kept in the experiments directory. The entry point is main.py and args will select a configuration:
python scripts/main.py --config-name=defaultCLI overrides are still possible with this selection but be aware of some shells not escaping quotes or sqaure brackets:
python scripts/main.py --config-name=default experiment.seed=37python scripts/main.py --config-name=pretrain_tinypython scripts/main.py --config-name=dimming_tinyTo train a model in the cloud start by branching off main and prepare your model and experiment configuration.
git branch -b <NEW_EXPERIMENT>
cd SDO-FM
cp experiments/default.yaml <NEW_EXPERIMENT_CONFIG>.yamlOnce ready commit your changes and tag it with a version number of the format v*.*.*, e.g. v0.1.2
git add .
git commit -m "Added <NEW EXPERIMENT>"
git tag v0.1.2
git push -u origin <NEW_EXPERIMENT>
git push --tagsThis will trigger an action to build this repository state in Google Cloud Build, it takes around 10 mins 🍵. Once completed it will be available in W&B as a job here. To run it, select "Launch" and define your overrides, e.g.
{
"args": [
"--config-name=<NEW_EXPERIMENT_CONFIG>",
"experiment_name=<NEW_EXPERIMENT_RUN>",
],
"run_config": {},
"entry_point": ["/src/scripts/main.py"]
}Set your compute resources abiding by this table, e.g.
MACHINE_TYPE: a2-ultragpu-8g
ACCELERATOR_TYPE: NVIDIA_A100_80GB
ACCELERATOR_COUNT: 8@software{SDOFM_2024,
title = {{Solar Dynamics Observatory Foundation Model}},
repository-code = {https:/spaceml-org/SDO-FM},
year = {2024}
}