What the Models Produce

The shipped default models can generate both signal-like waveforms and glitch-like waveforms. These samples are standardized strain time series with 512 points.

Install

pip install starccato-jax

For development, clone the repository and use uv:

uv sync
uv run train-vae --help

Train a VAE

The training CLI downloads/loads the requested dataset, standardizes each waveform, trains the VAE, and writes weights plus diagnostic plots into the output directory.

Training runs write model weights and diagnostics under the requested --outdir. A typical reconstruction diagnostic looks like this:

Train the signal VAE:

uv run train-vae \
  --dataset ccsne \
  --outdir runs/ccsne_vae \
  --latent-dim 5 \
  --epochs 1000 \
  --batch-size 64

Train the glitch VAE:

uv run train-vae \
  --dataset blip \
  --outdir runs/blip_vae \
  --latent-dim 5 \
  --epochs 1000 \
  --batch-size 64

Useful knobs:

• --latent-dim controls the dimension of the learned latent space. The default is 5.

• --cycles controls cyclical KL annealing for --no-use-capacity training. It is not used by the capacity objective.

• --use-capacity / --no-use-capacity switches the capacity-controlled KL objective.

• --capacity-end and --capacity-warmup-epochs tune the target KL ramp.

Why the Default Latent Dimension is 5

The VAE training default is latent_dim=5. The intended validation is a latent-dimensionality sweep over small latent sizes, using three families of diagnostics:

• latent geometry: total correlation and maximum off-diagonal correlation;

• cross-model separability: linear discriminant ROC-AUC and Gaussian Jensen-Shannon divergence;

• model fit: train/validation reconstruction loss and final total KL divergence.

The reproducible sweep script trains both waveform families at each requested latent dimension, records the measurements to CSV, and then plots the summary figure from those CSV files:

uv run python docs/studies/latent_dimensionality/run_latent_sweep.py \
  --latent-dims 2,3,4,5,6,8,12,16 \
  --epochs 1000

The script writes:

• docs/studies/latent_dimensionality/out/latent_sweep_metrics.csv with reconstruction, KL, active-dimension, correlation, and total-correlation measurements for each trained model;

• docs/studies/latent_dimensionality/out/latent_sweep_separability.csv with cross-model LDA ROC-AUC and Gaussian Jensen-Shannon divergence;

• docs/assets/latent_dimensionality_sweep.png generated only from those CSV files.

For a quick end-to-end check without claiming publication-quality numbers:

uv run python docs/studies/latent_dimensionality/run_latent_sweep.py \
  --latent-dims 2 \
  --epochs 2 \
  --outdir /tmp/starccato_latent_sweep_smoke

Load and Sample Models

The package ships default weights for both waveform families:

from starccato_jax.waveforms import StarccatoBlip, StarccatoCCSNe

signal_vae = StarccatoCCSNe()
glitch_vae = StarccatoBlip()

signal_waveforms = signal_vae.generate(n=8)
glitch_waveforms = glitch_vae.generate(n=8)

To load a model that you trained yourself, point StarccatoVAE at the output directory:

from starccato_jax import StarccatoVAE

signal_vae = StarccatoVAE("runs/ccsne_vae")
glitch_vae = StarccatoVAE("runs/blip_vae")

The current retraining contract, selected hyperparameters, artifact hashes, and real-data NUTS validation are recorded in docs/studies/vae_retraining/README.md. Run scripts/run_training_for_defaults.sh to reproduce the selected CCSNE and blip candidates.

Compare Signal and Glitch Embeddings

A direct comparison is to generate examples from each VAE, encode both populations through both encoders, and compare the latent norm distributions. This is a compact manifold-overlap check: it reduces each latent vector to its Euclidean norm, then compares the resulting one-dimensional distributions.

This plot is useful for answering: “Does the other waveform family land in the same broad part of this model’s latent space?” It is not, by itself, the best classifier, because two latent clouds can have similar radii while still being separated by direction or local geometry.

Interpretation:

• Low overlap in the ccsne latent space means the signal encoder places blip-like waveforms away from the signal manifold.

• Low overlap in the blip latent space means the glitch encoder places signal-like waveforms away from the glitch manifold.

• The overlap score below is the shared area between two normalized histograms. 0 means separated; 1 means indistinguishable by this one-dimensional norm summary.

• In the current default models, the glitch VAE latent norm is highly discriminating, while the signal VAE latent norm is less discriminating. That asymmetry is useful: it says the glitch model strongly rejects signal-like waveforms by latent radius, but the signal model’s latent radius alone is not enough to reject all glitch-like waveforms.

import numpy as np
import matplotlib.pyplot as plt

from starccato_jax.waveforms import StarccatoBlip, StarccatoCCSNe


def latent_norms(encoder, waveforms):
    z = np.asarray(encoder.encode(waveforms))
    return np.linalg.norm(z, axis=1)


def histogram_overlap(a, b, bins=60):
    values = np.concatenate([a, b])
    lo, hi = np.percentile(values, [1, 99])
    edges = np.geomspace(max(lo, 1e-6), hi, bins)
    ha, _ = np.histogram(a, bins=edges, density=True)
    hb, _ = np.histogram(b, bins=edges, density=True)
    widths = np.diff(edges)
    return float(np.sum(np.minimum(ha, hb) * widths)), edges


n = 5000
signal_vae = StarccatoCCSNe()
glitch_vae = StarccatoBlip()

signal_x = signal_vae.generate(n=n)
glitch_x = glitch_vae.generate(n=n)

signal_in_signal = latent_norms(signal_vae, signal_x)
glitch_in_signal = latent_norms(signal_vae, glitch_x)
glitch_in_glitch = latent_norms(glitch_vae, glitch_x)
signal_in_glitch = latent_norms(glitch_vae, signal_x)

signal_space_overlap, signal_edges = histogram_overlap(
    signal_in_signal, glitch_in_signal
)
glitch_space_overlap, glitch_edges = histogram_overlap(
    glitch_in_glitch, signal_in_glitch
)

print(f"Signal latent-space overlap: {signal_space_overlap:.3f}")
print(f"Glitch latent-space overlap: {glitch_space_overlap:.3f}")

fig, axes = plt.subplots(1, 2, figsize=(10, 4), sharey=True)

axes[0].hist(
    signal_in_signal,
    bins=signal_edges,
    density=True,
    histtype="step",
    lw=2,
    label="signal encoded by signal VAE",
)
axes[0].hist(
    glitch_in_signal,
    bins=signal_edges,
    density=True,
    histtype="step",
    lw=2,
    label="glitch encoded by signal VAE",
)
axes[0].set_title(f"Signal latent space, overlap={signal_space_overlap:.2f}")

axes[1].hist(
    glitch_in_glitch,
    bins=glitch_edges,
    density=True,
    histtype="step",
    lw=2,
    label="glitch encoded by glitch VAE",
)
axes[1].hist(
    signal_in_glitch,
    bins=glitch_edges,
    density=True,
    histtype="step",
    lw=2,
    label="signal encoded by glitch VAE",
)
axes[1].set_title(f"Glitch latent space, overlap={glitch_space_overlap:.2f}")

for ax in axes:
    ax.set_xscale("log")
    ax.set_xlabel("latent vector norm")
    ax.set_ylabel("density")
    ax.legend(frameon=False)

fig.tight_layout()
fig.savefig("vae_embedding_overlap.png", dpi=200)

Class Separation Score

For classification, the strongest simple diagnostic is a cross-reconstruction score. Reconstruct the same waveform with both VAEs and compare the mean squared reconstruction error:

Positive scores favor the signal VAE. Negative scores favor the glitch VAE. On the held-out validation waveforms sampled below, the current default weights separate the two classes cleanly with a threshold at 0.

Why this works:

• Each VAE is trained to reconstruct one waveform family.

• A waveform should reconstruct better under the model whose training distribution it resembles.

• The log ratio makes the score symmetric and easy to threshold: 0 means both VAEs reconstruct equally well.

• The magnitude is interpretable. A score of +1 means the glitch VAE has about ten times the reconstruction MSE of the signal VAE. A score of -2 means the glitch VAE has about one hundredth of the signal VAE’s reconstruction MSE.

For the current default weights, the held-out validation sample used to make this figure gives:

• CCSNe accuracy at threshold 0: 1.000.

• Blip accuracy at threshold 0: 1.000.

• Balanced accuracy: 1.000.

• Median CCSNe score: +1.156.

• Median blip score: -1.997.

These numbers should be treated as a model diagnostic, not a final detection claim. For an analysis paper, recompute them on the exact held-out data, preprocessing, and trained weights used in the experiment.

import numpy as np

from starccato_jax.data import TrainValData
from starccato_jax.waveforms import StarccatoBlip, StarccatoCCSNe


def reconstruction_mse(model, x):
    xhat = np.asarray(model.reconstruct(x))
    return np.mean((xhat - x) ** 2, axis=1)


signal_vae = StarccatoCCSNe()
glitch_vae = StarccatoBlip()

signal_x = np.asarray(TrainValData.load(source="ccsne").val)
glitch_x = np.asarray(TrainValData.load(source="blip").val)

signal_score = np.log10(
    reconstruction_mse(glitch_vae, signal_x)
    / reconstruction_mse(signal_vae, signal_x)
)
glitch_score = np.log10(
    reconstruction_mse(glitch_vae, glitch_x)
    / reconstruction_mse(signal_vae, glitch_x)
)

signal_accuracy = np.mean(signal_score > 0)
glitch_accuracy = np.mean(glitch_score < 0)
balanced_accuracy = 0.5 * (signal_accuracy + glitch_accuracy)

Latent Projection Study

The norm plot compresses each latent vector to one number. A second useful view is to keep the full latent vectors, combine the signal and glitch encodings in each VAE latent space, and project them to two dimensions with PCA.

This plot is qualitative. It helps show whether separation is visible in the dominant latent directions, but it should not replace a numerical score such as the cross-reconstruction ratio. PCA is used here because it is deterministic, dependency-free, and sufficient for a lightweight documentation diagnostic.

import numpy as np
import matplotlib.pyplot as plt

from starccato_jax.waveforms import StarccatoBlip, StarccatoCCSNe


def encode(encoder, waveforms):
    return np.asarray(encoder.encode(waveforms))


def pca2(z):
    z0 = z - z.mean(axis=0, keepdims=True)
    _, _, vt = np.linalg.svd(z0, full_matrices=False)
    return z0 @ vt[:2].T


n = 2500
signal_vae = StarccatoCCSNe()
glitch_vae = StarccatoBlip()

signal_x = signal_vae.generate(n=n)
glitch_x = glitch_vae.generate(n=n)

signal_space_xy = pca2(
    np.vstack(
        [
            encode(signal_vae, signal_x),
            encode(signal_vae, glitch_x),
        ]
    )
)
glitch_space_xy = pca2(
    np.vstack(
        [
            encode(glitch_vae, glitch_x),
            encode(glitch_vae, signal_x),
        ]
    )
)

Recommended Workflow

Use the diagnostics in this order:

1. Train one ccsne VAE and one blip VAE with matched preprocessing and comparable latent dimensions.

2. Check reconstruction plots from each model to confirm that the VAE learned the basic waveform morphology.

3. Compute cross-reconstruction scores on held-out validation waveforms. This is the main class-separation diagnostic.

4. Use latent norm overlap to summarize whether one model’s latent space broadly accepts or rejects the other waveform family.

5. Use PCA projections as a visual sanity check of the full latent vectors.

If the cross-reconstruction score separates the classes but the latent norm overlap does not, the VAEs are still useful for discrimination: the separation is likely carried by latent direction or decoder reconstruction quality rather than by radius alone.