# AION-1 : AstronomIcal Omnimodal Network

AION is a family of foundation models specifically designed for astronomical observations, representing a paradigm shift in how we approach machine learning for astronomy. Built on the Multimodal Universe (MMU) dataset, AION can process and generate predictions across 39 different astronomical modalities.

## 🌟 Overview

AION addresses key limitations of traditional deep learning in astronomy:
- **Limited supervised training data** for rare objects
- **Poor model reusability** across different tasks
- **Difficulty in cross-survey analysis**

By leveraging foundation model principles and training on 100TB+ of astronomical data, AION provides:
- 🔄 **Any-to-any multimodal capabilities**
- 🌐 **Cross-survey understanding** 
- 📊 **Linear probing for downstream tasks**
- 🚀 **Zero/few-shot learning abilities**

## 📡 Supported Data Sources

AION-1 is trained on data from major astronomical surveys:

- **DES** (Dark Energy Survey) - _griz_ bands
- **HSC** (Hyper Suprime-Cam) - _grizy_ bands  
- **DESI** (Dark Energy Spectroscopic Instrument) - Spectra
- **Gaia** - Astrometry and photometry
- **SDSS** (Sloan Digital Sky Survey)

## 🎯 Key Capabilities

### Cross-Survey Translation
```python
# Translate DES observations to HSC
hsc_pred = model.forward(
    input_dict={'des': des_images},
    target_mask={'hsc': hsc_mask}
)
```

### Cross-Modal Prediction  
```python
# Predict DESI spectra from HSC images
spectra = model.forward(
    input_dict={'hsc': hsc_images},
    target_mask={'desi': spectrum_mask}
)
```

### Physical Parameter Estimation
```python
# Estimate redshift from multiple inputs
z_pred = model.forward(
    input_dict={'des': des_images, 'gaia': gaia_features},
    target_mask={'redshift': z_mask}
)
```

### Spectrum Super-Resolution
```python
# Enhance low-resolution spectra
hires_spec = model.forward(
    input_dict={'gaia': gaia_spectrum},
    target_mask={'desi': desi_mask}
)
```

## 🏗️ Model Architecture

AION extends the 4M (Massively Multimodal Masked Modeling) architecture with:
- **Modality-specific tokenizers** using FSQ (Finite Scalar Quantization)
- **Transformer encoder-decoder** with cross-attention
- **Generative masked modeling** pre-training objective

### Model Sizes

| Model | Parameters | Training Hardware | Training Time |
|-------|------------|------------------|---------------|
| AION-1 Base | 300M | 64 H100s | 1.5 days |
| AION-1 Large | 800M | 100 H100s | 2.5 days |
| AION-1 XLarge | 3B | 288 H100s | 3.5 days |


## Installation

Setting up your environment
```bash
python -m venv venv
source venv/bin/activate
pip install --upgrade pip
pip install torch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0 --index-url https://download.pytorch.org/whl/cu124
pip install -e .
```

```bash
from aion import AION

model = AION.from_pretrained('path/to/model/weights')
```