DMLAB-GSU/Surya-AR-Segmentation

Surya AR Segmentation Checkpoints

This repository contains the trained model checkpoints for the Active Region (AR) Segmentation downstream application built upon the Surya Solar Foundation Model.

For more details on the foundation model, please refer to the Surya paper and the main GitHub repository.

Project Motivation

Active Regions (ARs) with polarity inversion lines are detected using conventional computer vision techniques. For more details, please check the ARPILs repository. The binary AR masks generated by these techniques are used as ground truth. We map two consecutive 13-channel 4096x4096 inputs (4D) to 2D 4096x4096 maps to evaluate whether the foundation model can successfully detect ARs with polarity inversion lines.

Model Description

The checkpoints in this repository contain the fine-tuned weights for the custom head architectures built on top of Surya embeddings. The models were trained to detect Active Regions (ARs) given full-disk SDO observations.

• Checkpoints: The .pth files correspond to different training configurations and architectures (e.g., U-Net, MobileNet, LoRA variants, etc.) evaluated for AR segmentation.
• Input Data: The input consists of 4K (4096x4096) 13-channel SDO images.

Data Preparation

• The full AR Segmentation dataset requires ~1.8TB of disk space, including AR masks (target) and input 4K SDO images.
• You can download the AR masks and extract them using scripts from the main GitHub repository.
• SDO images (.nc files) are hosted on an S3 bucket: s3://nasa-surya-bench/.

Usage & Inference Demo

To use these checkpoints for inference or fine-tuning, please use the codebase provided in the Surya AR Segmentation & Flare Forecasting GitHub Repository.

Prerequisites

• Python 3.11+
• A CUDA-capable GPU (An H100 80GB is highly recommended for 4096x4096 inputs).

Running Inference

# Clone the repository
git clone https://github.com/DMLAB-GSU/surya_arseg_flare.git
cd surya_arseg_flare

# Create and activate conda environment
conda env create -f environment.yml
conda activate surya_arseg_flare
pip install -e .

# Run the inference script (make sure to place the downloaded checkpoints in the assets/ folder)
python scripts/arseg/infer.py experiment=arseg/sample_inference

Support & Troubleshooting

If you encounter any issues or have general questions about the codebase, please reach out to us:

• Jinsu Hong: jhong36@gsu.edu
• Kang Yang: kyang30@gsu.edu
• Berkay Aydin: baydin2@gsu.edu

Acknowledgements

This material is based upon work supported by NASA under Award No. 80NSSC24PC377. Computational resources were provided by NSF ACCESS, GSU ARCTIC, and NASA Advanced Supercomputing (NAS). The authors acknowledge the NASA Solar Dynamics Observatory (SDO) team and NOAA/GOES for providing the data used in this work.