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=default
CLI 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=37
python scripts/main.py --config-name=pretrain_tiny
python scripts/main.py --config-name=dimming_tiny
To 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>.yaml
Once 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 --tags
This 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://github.com/spaceml-org/SDO-FM},
year = {2024}
}
This work is the research product of the SDO-FM: A Multi-Modal Foundation Model POC for SDO. This has been funded and supported by NASA under Grant award No 80NSSC24K0701. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views of the National Aeronautics and Space Administration (NASA). The research and its outputs have been designed, managed and delivered by Trillium Technologies Inc (https://trillium.tech). Trillium is a research and development company with a focus on intelligent systems and collaborative communities for planetary stewardship, space exploration and human health. Trillium aspires to ensure that the latest tools and techniques in Artificial Intelligence (AI) and Machine Learning (ML) are applied to developing open science for all Humankind.