VAE training and 21-cm kernels

The VAE tools learn a low-dimensional representation of simulated 21-cm power spectra. The trained decoder is consumed by VAEKernTorch.

Train dimensionality-reduction models for the VAE 21-cm kernel.

BetaVAE learns a latent representation of simulated power spectra. VAEFitterPreProc handles preprocessing and training, while FitterResult provides reconstruction and latent-space diagnostics.

class ps_eor.ml_gpr.vae.BetaVAE(in_dim, latent_dim, hidden_dims=None, beta=1, warmup_iters=100, warmup_gamma=0.001, loss_type='H', fc_hidden_dim=None, **kwargs)

Bases: Module

Fully connected beta-VAE for one-dimensional power spectra.

Adapted from PyTorch-VAE.

encode(input)

Return latent means and log variances for a batch of spectra.

decode(z)

Decode latent coordinates into normalized spectra.

reparameterize(mu, logvar)

Draw differentiable latent samples from N(mu, exp(logvar)).

forward(input, **kwargs)

Encode, sample, and reconstruct a batch.

step()

Advance the warm-up schedule by one training epoch.

loss_function(*args, **kwargs)

Return total, reconstruction, and KL losses.

generate(x, **kwargs)

Reconstruct a batch using sampled latent coordinates.

generate_m(x)

Reconstruct a batch from the latent means.

class ps_eor.ml_gpr.vae.PreProcessor

Bases: object

Normalize spectra after removing empty curves and amplitude outliers.

process(ps_data, n_sigma_clip=3)

Filter and normalize a power-spectrum training set.

inv(ps)

Undo preprocessing on decoded spectra.

class ps_eor.ml_gpr.vae.PreProcessorFlatten(k_mean, alpha=1)

Bases: PreProcessor

Flatten a power-law trend before VAE training.

process(ps_data, n_sigma_clip=3)

Normalize spectra after dividing by k_mean**alpha.

inv(ps)

Restore the removed power-law trend.

class ps_eor.ml_gpr.vae.PreProcessorLogScale(k_mean)

Bases: PreProcessor

Map normalized power spectra to a bounded logarithmic scale.

process(ps_data, n_sigma_clip=3)

Filter, normalize, and log-scale a training set.

inv(ps)

Map decoded values back to linear power.

class ps_eor.ml_gpr.vae.AbstractFitter(n_dim, k_mean)

Bases: object

Common interface for saved VAE and PCA fitters.

encode(data)

Map spectra to latent coordinates.

decode(data)

Map latent coordinates to spectra.

reconstruct(data)

Encode and decode spectra.

save(filename)

Save the trained fitter for use as fitter_filename.

static load(filename)

Load a saved fitter, caching it by path and modification time.

class ps_eor.ml_gpr.vae.VAEFitter(model, optimizer, k_mean)

Bases: AbstractFitter

Train a BetaVAE on pre-scaled power spectra.

Prefer VAEFitterPreProc for raw training spectra.

fit(dataloader)

Run one training epoch and return its loss components.

validate(dataloader)

Evaluate loss components on a validation loader.

train(epochs, train_loader, val_loader)

Train for epochs and retain training and validation losses.

ensure_tensor(data)

Convert NumPy input to the model’s tensor type.

ensure_np(data)

Convert tensor output to NumPy.

encode(data)

Sample latent coordinates for input spectra.

decode(latent_data)

Decode latent coordinates to normalized spectra.

reconstruct(data)

Reconstruct spectra from their latent means.

class ps_eor.ml_gpr.vae.PCAFitter(n_cmpt, k_mean)

Bases: AbstractFitter

Linear PCA alternative implementing the fitter interface.

train(train_set)

Fit PCA and record the latent bounds of the training set.

encode(data)

Project spectra into PCA coordinates.

decode(latent_data)

Reconstruct spectra from PCA coordinates.

reconstruct(data)

Project and reconstruct spectra.

class ps_eor.ml_gpr.vae.VAEFitterPreProc(model, optimizer, pre_proc)

Bases: VAEFitter

Train a VAE after preprocessing raw power spectra.

train() preprocesses the library, creates a validation split, and returns a FitterResult.

train(epochs, train_data, frac_validation=0.1, batch_size=128, rng=None)

Preprocess, split, and train on a power-spectrum library.

class ps_eor.ml_gpr.vae.FitterResult(fitter, train_data, test_data)

Bases: AbstractFitter

Diagnostics for a trained fitter.

Use plot_loss(), plot_latent_qq(), and plot_ratio() before using the decoder as a kernel.

plot_loss()

Plot reconstruction and KL losses over training.

plot_latent_qq()

Compare training and validation latent coordinates with a normal law.

get_reco_ratio_train()

Return reconstructed-to-input ratios for the training set.

get_reco_ratio_val()

Return reconstructed-to-input ratios for the validation set.

plot_ratio()

Plot reconstruction ratios for training and validation sets.