Views
No views yet
pip install polymathic-aion1import torch
2from aion import AION
3from aion.codecs import CodecManager
4from aion.modalities import LegacySurveyImage, Z
5
6# Load model and codec manager
7model = AION.from_pretrained('polymathic-ai/aion-base').to('cuda') # or 'aion-large', 'aion-xlarge'
8codec_manager = CodecManager(device='cuda')
9
10# Example: Prepare your astronomical data (e.g., a dummy Legacy Survey image)
11# In a real scenario, 'your_image_tensor' would come from your dataset.
12your_image_tensor = torch.randn(1, 4, 96, 96) # Example: batch_size=1, 4 bands, 96x96 resolution
13image = LegacySurveyImage(
14 flux=your_image_tensor,
15 bands=['DES-G', 'DES-R', 'DES-I', 'DES-Z']
16)
17
18# Encode data to tokens
19tokens = codec_manager.encode(image)
20
21# Option 1: Extract embeddings for downstream tasks
22embeddings = model.encode(tokens, num_encoder_tokens=600)
23print(f"Extracted embeddings shape: {embeddings.shape}")
24
25# Option 2: Generate predictions (e.g., redshift)
26# For this example, we predict redshift (Z) from the image.
27# The target_mask tells the model which modality to generate.
28preds = model(
29 codec_manager.encode(image),
30 target_modality=Z,
31)
32print(f"Predicted redshift logits shape: {preds['tok_z'].shape}")| Category | Description | Token Name(s) |
|---|---|---|
| Imaging (2) | Legacy Survey, HSC Wide | tok_image_ls, tok_image_hsc |
| Catalog (1) | Legacy Survey catalog entries | catalog |
| Spectra (2) | SDSS, DESI | tok_spectrum_sdss, tok_spectrum_desi |
| Gaia (4) | BP/RP spectra, parallax, sky coords | tok_xp_bp, tok_xp_rp, tok_parallax, tok_ra, tok_dec |
| Gaia Photometry (3) | G/BP/RP flux | tok_flux_g_gaia, tok_flux_bp_gaia, tok_flux_rp_gaia |
| Legacy Survey (9) | g,r,i,z bands & WISE W1–W4 flux, E(B–V) | tok_flux_g,…,tok_flux_w4, tok_ebv |
| Legacy Shape (3) | Ellipticity components & effective radius | tok_shape_e1, tok_shape_e2, tok_shape_r |
| HSC Photometry (5) | g,r,i,z,y magnitudes | tok_mag_g,…,tok_mag_y |
| HSC Extinction (5) | g,r,i,z,y extinctions | tok_a_g,…,tok_a_y |
| HSC Shape (3) | Shape components 11,22,12 | tok_shape11, tok_shape22, tok_shape12 |
| Other (1) | Spectroscopic redshift | tok_z |