Metadata-Version: 2.4
Name: starccato_jax
Version: 0.2.4
Summary: JAX-VAE stellar core collapse waveform
Author-email: Starccato Team <avi.vajpeyi@gmail.com>
Project-URL: Homepage, https://starccato.github.io/starccato_jax/
Project-URL: Bug Reports, https://giuthub.com/starccato/starccato_jax/issues
Project-URL: Source, https://github.com/starccato/starccato_jax/
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3.8
Requires-Python: >=3.9
Description-Content-Type: text/x-rst
Requires-Dist: jax
Requires-Dist: flax
Requires-Dist: numpy
Requires-Dist: matplotlib
Requires-Dist: pandas
Requires-Dist: h5py
Requires-Dist: pillow
Requires-Dist: natsort
Requires-Dist: tqdm
Requires-Dist: rich
Requires-Dist: uniplot
Requires-Dist: click
Requires-Dist: arviz
Provides-Extra: dev
Requires-Dist: pytest; extra == "dev"
Requires-Dist: pytest-cov; extra == "dev"
Requires-Dist: flake8>=5.0.4; extra == "dev"
Requires-Dist: black>=22.12.0; extra == "dev"
Requires-Dist: black[jupyter]>=22.12.0; extra == "dev"
Requires-Dist: isort; extra == "dev"
Requires-Dist: mypy; extra == "dev"
Requires-Dist: jupyter-book<2; extra == "dev"
Requires-Dist: pre-commit; extra == "dev"

Starccato-JAX
=============

``starccato_jax`` trains and serves compact VAE waveform surrogates for two
classes used in Starccato analyses:

* ``ccsne``: the signal VAE for core-collapse supernova waveforms.
* ``blip``: the glitch VAE for short blip-like detector transients.

The two practical jobs are:

1. Train a VAE for one waveform family.
2. Compare the signal and glitch VAEs in latent space to see how much their
   encoded populations overlap.

The documentation below shows three complementary diagnostics:

* **Latent norm overlap** asks whether the two waveform families occupy
  similar radii after encoding.
* **Cross-reconstruction score** asks which VAE reconstructs a waveform better.
  This is the most useful simple class-separation score.
* **Latent PCA projection** asks whether the full latent vectors separate in a
  low-dimensional view.

The README figures are reproducible. To regenerate them from the shipped
default model artifacts:

.. code-block:: bash

   uv run python docs/studies/make_readme_plots.py

That script writes the figures to ``docs/assets/`` and records the numerical
caption values in ``docs/studies/readme_plot_metrics.csv``.

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.

.. figure:: assets/waveform_samples.png
   :alt: Generated CCSNe and blip waveform samples
   :width: 95%
   :align: center

   Example waveforms generated from the default signal and glitch VAEs.

Install
-------

.. code-block:: bash

   pip install starccato-jax

For development, clone the repository and use ``uv``:

.. code-block:: bash

   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:

.. figure:: model_out/reconstructions.png
   :alt: Example VAE reconstruction diagnostic
   :width: 80%
   :align: center

   Example reconstruction plot produced by the VAE diagnostics.

Train the signal VAE:

.. code-block:: bash

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

Train the glitch VAE:

.. code-block:: bash

   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.
* ``--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:

.. code-block:: bash

   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:

.. code-block:: bash

   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:

.. code-block:: python

   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:

.. code-block:: python

   from starccato_jax import StarccatoVAE

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

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.

.. figure:: assets/embedding_overlap.png
   :alt: Cross-encoded signal and glitch latent norm overlap
   :width: 95%
   :align: center

   Cross-encoding generated waveforms through both VAEs. With the current
   default weights, the signal latent-space norm overlap is about ``0.72`` and
   the glitch latent-space norm overlap is about ``0.005``.

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.

.. code-block:: python

   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:

.. math::

   s(x) = \log_{10}\left(
       \frac{\mathrm{MSE}(\mathrm{glitch\ VAE}, x)}
            {\mathrm{MSE}(\mathrm{signal\ VAE}, x)}
   \right)

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.

.. figure:: assets/cross_reconstruction_score.png
   :alt: Cross-reconstruction score separating held-out CCSNe and blip waveforms
   :width: 95%
   :align: center

   Cross-reconstruction error is a direct class-separation diagnostic. The
   left panel shows the signed decision score; the right panel shows the median
   reconstruction error under each model.

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.

.. code-block:: python

   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.

.. figure:: assets/latent_pca_projection.png
   :alt: PCA projection of cross-encoded signal and glitch latent vectors
   :width: 95%
   :align: center

   Two-dimensional PCA projection of the encoded latent vectors. Each panel is
   fit within one VAE latent space, then both waveform families are plotted in
   that same projected coordinate system.

.. code-block:: python

   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.
