vaexede 0.1.0

Trained models for EDE and LCDM VAEs

VAExEDE

A repository to host the trained models from https://arxiv.org/abs/2502.09810, where we trained a variational autoencoder (VAE) on CMB temperature power spectra for early dark energy (EDE) models, to discover novel, data-driven parametrizations.

Installation

To use the trained models, follow these steps:

1. (optional) conda create -n vaexede python=3.11 jupyter (create a custom conda environment with python 3.11)

2. (optional) conda activate vaexede (activate it)

3. Install the package:

    pip install vaexede
    python -c 'from vaexede.load_models import load_model'
   

   or alternatively, clone the repository and install it:

    git clone https://github.com/dpiras/VAExEDE.git
    cd VAExEDE
    pip install . 
    python -c 'from vaexede.load_models import load_model'
   

The latter option will also give you access to a Jupyter notebook with a quick walkthrough on how to use the models.

Usage

A simple way to load and use the trained models looks like this:

model = 'lcdm' # either 'lcdm' or 'ede'
lcdm_model = load_model(model) # load the trained network

# add your unitless D_ell temperature spectrum here, in the ell range [30, 2500]
# should also support batches of data, has not been tested though
# D_ell = C_ell * ell * (ell+1) / 2pi
input_spectrum = # your spectrum here

# and preprocess it as described in the paper
input_spectrum_preprocess = preprocess(input_spectrum, model)

# here the preprocessed spectrum gets encoded, samples from the latent space are obtained
# then the spectrum is decoded and unpreprocessed
mean, logvar = lcdm_model.encode(input_spectrum_preprocess.reshape(1, -1))
z = lcdm_model.reparameterize(mean, logvar) # here we sample from the latent distribution
decoded_spectrum = lcdm_model.decode(z)
output_spectrum_lcdm = unpreprocess(decoded_spectrum[0, :, 0], model)

Disclaimer

The repository contains only some of the material needed to reproduce the paper. If you need more or would like to add a feature, feel free to fork this repository to work on it; otherwise, please raise an issue or contact Davide Piras.

Contributors

Laura Herold and Luisa Lucie-Smith contributed to this code.

Citation

If you use this code, please cite the corresponding paper:

 @article{Piras:2025eip,
 author = "Piras, Davide and Herold, Laura and Lucie-Smith, Luisa and Komatsu, Eiichiro",
 title = "{$\Lambda$CDM and early dark energy in latent space: a data-driven parametrization of the CMB temperature power spectrum}",
 eprint = "2502.09810",
 archivePrefix = "arXiv",
 primaryClass = "astro-ph.CO",
 month = "2",
 year = "2025"
 }

License

This code is released under the GPL-3 license - see LICENSE-, subject to the non-commercial use condition - see LICENSE_EXT.

 VAExEDE
 Copyright (C) 2025 Davide Piras & contributors

 This program is released under the GPL-3 license (see LICENSE.txt), 
 subject to a non-commercial use condition (see LICENSE_EXT.txt).

 This program is distributed in the hope that it will be useful,
 but WITHOUT ANY WARRANTY; without even the implied warranty of
 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.