SSP, Gas, And AGN Compression

Purpose

The compression implemented in this repository is runtime compression, not disk compression. The goal is to reduce the spectral tensors that are resident in JAX/GPU memory so larger galaxy batches fit on the same GPU.

HDF5 gzip or shuffle compression can make files smaller on disk, but it does not help if load_context expands the file into a dense float32 tensor before JAX sees it. The compressed configs therefore load compact arrays directly:

basis[k, wave]
coeff[physical_axes..., k]
scale[physical_axes...]

The public FS2 config loads compact arrays directly and is the recommended path for high-throughput FS2 MAP fits.

Compressed Assets

The active FS2 config uses:

runtime: gpu

fit:
  trace_mode: optimizer
  trace_interval: 20

model:
  ssp_model: compressed_basis
  compressed_ssp_path: Data/popcosmos_chabrier_stellar_ssp_basis_k64_coeff16.h5
  nebular_model: compressed_gas_grid
  compressed_gas_grid_path: Data/popcosmos_chabrier_gas_grid_basis_k64_mixed16.h5
  agn_model: compressed_fsps_component_grid
  compressed_agn_component_grid_path: Data/popcosmos_chabrier_agn_component_basis_k12_fagnlinear_coeff16.h5

These assets are generated from the dense Chabrier FSPS assets. They keep the same wavelength, age, metallicity, gas, and AGN axes, but replace the repeated dense wavelength axis by a low-rank basis.

How The SVD Compression Works

For the base SSP cube, the dense tensor is:

ssp_flux[Zstar, age, wave]

The compressor flattens Zstar and age into one curve axis:

F[curve, wave]

Each curve is normalized before SVD:

scale[curve] = norm(F[curve, :])
F_norm[curve, wave] = F[curve, wave] / scale[curve]

SVD then finds a wavelength basis ordered by decreasing singular value:

F_norm ~= U_k S_k V_k.T

The HDF5 asset stores:

basis = V_k.T
coeff = U_k S_k
scale = scale

At runtime DSPS reconstructs only the metallicity-specific slice needed by the current forward pass:

coeff_z = interpolate_Zstar(ssp_coeff, lgmet_abs)      # [age, k]
scale_z = interpolate_Zstar(ssp_scale, lgmet_abs)      # [age]
weighted_coeff = coeff_z * scale_z[:, None]
ssp_flux_z = jnp.einsum("ak,kw->aw", weighted_coeff, ssp_basis)

This gives ssp_flux_z[age, wave] without loading the full ssp_flux[Zstar, age, wave] tensor into JAX.

Gas Compression

The dense gas grid is much larger:

ssp_flux[Zgas, Ugas, Zstar, age, wave]

where:

  • Zgas is log10_gas_metallicity = log10(Zgas / Zsun);

  • Ugas is log10_gas_ionization = log10 U;

  • Zstar is the stellar metallicity axis;

  • age is the SSP age axis;

  • wave is the wavelength axis.

The compressed gas asset stores:

gas_basis[k, wave]
gas_coeff[Zgas, Ugas, Zstar, age, k]
gas_scale[Zgas, Ugas, Zstar, age]

DSPS interpolates in gas metallicity, gas ionization, and stellar metallicity in coefficient space, then expands only the requested [age, wave] slice. This avoids materializing the dense gas grid or four dense interpolation corners.

Gas spectra contain narrow emission lines, so the gas basis is harder to compress than a smooth stellar continuum. The current production compromise is k=64 with mixed precision storage.

AGN Compression

The dense FSPS-native AGN component grid is:

agn_lnu_per_mformed[fagn, tauagn, Zstar, age, wave]

The fitted AGN parameters are:

ln_fagn    -> fagn = exp(ln_fagn)
ln_tauagn  -> tauagn = exp(ln_tauagn)

The AGN grid is effectively linear in fagn for the active asset family. The compressed format therefore removes the fagn axis from storage and applies fagn as a runtime multiplier:

agn_basis[k, wave]
agn_coeff[tauagn, Zstar, age, k]
agn_scale[tauagn, Zstar, age]

Runtime reconstruction:

agn_flux[age, wave]
  = fagn * agn_scale[tauagn, Zstar, age]
         * agn_coeff[tauagn, Zstar, age, k] @ agn_basis[k, wave]

This is why AGN compression is especially effective: the dense fagn axis does not need to be resident on GPU.

Why SVD, Not Wavelets

A sparse Haar-wavelet benchmark was added in scripts/benchmark_wavelet_compression.py. It compares the current SVD-style compressed assets to an oracle sparse Haar representation that stores the top wavelet coefficients and their indices per spectrum.

The result was:

  • AGN: current fagn-factored SVD is much smaller and much more accurate.

  • Gas: current SVD k64 mixed precision is smaller and more accurate than the tested sparse Haar candidates.

  • Stellar SSP: Haar is competitive on a sampled median spectral metric, but its worst-tail error is worse and it has no validated JAX sparse runtime path.

Decision: keep SVD/factored-SVD as the production compression method. Wavelets remain a research option only after photometric closure and JAX sparse-layout benchmarks.

Build Compressed Assets

The dense input assets must already exist. Then build:

python scripts/build_compressed_ssp_grid.py \
  --input Data/fsps_v0.4.7_mist_c3k_a_chabrier_wNE_logGasU-2.0_logGasZ0.0.h5 \
  --output Data/popcosmos_chabrier_stellar_ssp_basis_k64_coeff16.h5 \
  --k 64 --basis-dtype float32 --coeff-dtype float16 --overwrite

python scripts/build_compressed_gas_grid.py \
  --input Data/popcosmos_chabrier_gas_ssp_grid.h5 \
  --output Data/popcosmos_chabrier_gas_grid_basis_k64_mixed16.h5 \
  --k 64 --basis-dtype float16 --coeff-dtype float16 --overwrite

python scripts/build_compressed_agn_component_grid.py \
  --input Data/popcosmos_chabrier_agn_component_ssp_grid.h5 \
  --output Data/popcosmos_chabrier_agn_component_basis_k12_fagnlinear_coeff16.h5 \
  --k 12 --factor-fagn --basis-dtype float32 --coeff-dtype float16 --overwrite

Validate compressed assets:

python scripts/validate_compressed_spectral_asset.py \
  Data/popcosmos_chabrier_stellar_ssp_basis_k64_coeff16.h5

python scripts/validate_compressed_spectral_asset.py \
  Data/popcosmos_chabrier_gas_grid_basis_k64_mixed16.h5

python scripts/validate_compressed_spectral_asset.py \
  Data/popcosmos_chabrier_agn_component_basis_k12_fagnlinear_coeff16.h5

Production Fit Command

Use the FS2 GPU config for large MAP batches:

conda activate shine

python -m euclid_dsps.cli \
  --config configs/fs2_gpu.yaml \
  fit \
  --limit 1000 \
  --batch-size 128 \
  --fit-maxiter 200 \
  --out outputs/runs/fs2_gpu_map_n1000_bs128 \
  --sed-samples 0 \
  --reporting-level light

If the target GPU is smaller, reduce --batch-size first. The compressed config already requests the GPU runtime preset and TF_GPU_ALLOCATOR setting used in the local successful runs.

Checks

Run the lightweight checks after changing compression code or configs:

python -m compileall euclid_dsps scripts
pytest tests/test_config.py
pytest tests/test_model.py

For scientific closure, run dense-vs-compressed and FSPS/Prospector benchmarks before using a new compressed asset family for final science.