diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..c3abd62 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,96 @@ +# CLAUDE.md + +This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. + +## Overview + +Synference performs simulation-based inference (SBI) SED fitting of galaxy photometry and spectroscopy. It couples [Synthesizer](https://synthesizer-project.github.io) (mock spectra/photometry generation) with [LtU-ILI](https://ltu-ili.readthedocs.io/) (amortised posterior inference, wrapping the `sbi` and `lampe` backends). The core design principle: **mock library generation is decoupled from feature engineering, which is decoupled from model training** — one expensive HDF5 library can feed many trained models with different filter sets, noise models, and normalizations. + +## Commands + +```bash +# Install (editable). ltu-ili and synthesizer cannot come from PyPI — install from git: +pip install -e .[dev] +pip install "ltu_ili@git+https://github.com/maho3/ltu-ili.git" +pip install "cosmos-synthesizer@git+https://github.com/synthesizer-project/synthesizer.git#egg=synthesizer" + +# Test data (required before running tests): +synthesizer-download --test-grids --dust-grid +python -c "from synference.utils import download_test_data; download_test_data()" + +# Tests +pytest # full suite +pytest tests/test_sbi.py # one file +pytest tests/test_sbi.py -k feature_array # one test by keyword + +# Lint / format (ruff, config in pyproject.toml: line-length 100, Google docstrings, D-rules enforced except in examples/) +ruff check --fix . +ruff format . +pre-commit run --all-files # ruff + nb-clean + large-file/merge-conflict checks + +# Docs (executes all notebooks; needs pandoc and .[docs] extras) +cd docs && make html SPHINXOPTS="-j auto" +``` + +`sbi_runner.py` (and therefore `SBI_Fitter`) is imported inside a try/except in `__init__.py` — if `ltu-ili`/`sbi`/`lampe` are missing the package still imports but only library-generation functions are available. If `SBI_Fitter` is mysteriously undefined, check that import error printed at import time. + +## Architecture: the three-stage pipeline + +### Stage 1 — Library generation (`src/synference/library.py`) + +Builds an HDF5 "library" of (parameters → photometry/spectra) pairs by running the Synthesizer pipeline. + +- **`GalaxyBasis`** — defines the model: SPS `Grid`, `EmissionModel`, list of SFH objects, metallicity distributions, redshifts, `Instrument` (filters). Two modes: pass N samples of every parameter (e.g. from `draw_from_hypercube` Latin hypercube draws), or set `build_library=True` to take the outer product of a few basis SFHs/redshifts/ZDists. `GalaxyBasis.create_mock_library()` is the main entry point (multiprocessing via `n_proc`, optional `multi_node` MPI/SLURM operation, batched in `batch_size` chunks). +- **`CombinedBasis`** — what `create_mock_library` delegates to. Galaxies are simulated at a filler mass and renormalized to the requested stellar masses afterwards; this also lets multiple `GalaxyBasis` objects (e.g. different SPS grids) be mass-weighted into composite galaxies via `combination_weights`. +- **Supplementary parameters** — derived quantities (`calculate_muv`, `calculate_sfr`, `calculate_beta`, `calculate_line_ew`, etc.; registry in `SUPP_FUNCTIONS`) are computed per-galaxy at library time and stored alongside the free parameters, so they can later be fitted or used as features. Extra ones are passed as `**extra_analysis_functions`. +- **`GalaxySimulator`** — a self-contained parameters→photometry callable (SFH class + ZDist class + grid + instrument + emission model). Used for online/sequential SBI training, SED recovery from posterior samples, and is serialized into the library so `SBI_Fitter.recreate_simulator_from_library()` can rebuild it. +- **`alt_parametrizations` / `parameter_transforms_to_save`** — mechanism for fitting in a different parametrization than Synthesizer uses (no lambdas; must be picklable). + +**HDF5 library layout** (read by `utils.load_library_from_hdf5`): datasets `Grid/Photometry`, `Grid/Parameters`, `Grid/SupplementaryParameters`, optionally `Grid/Spectra`; root attrs `FilterCodes`, `ParameterNames`, `ParameterUnits`, `PhotometryUnits`, `SupplementaryParameterNames`. Raw photometry is stored noiseless in physical units (typically nJy). + +Default output dir is `library_folder` = `/libraries/` — in this checkout that is a symlink to bulk data storage. + +### Stage 2 — Feature engineering + training (`src/synference/sbi_runner.py`, ~9k lines, the heart of the package) + +**`SBI_Fitter`** — instantiate with `SBI_Fitter.init_from_hdf5(model_name=..., hdf5_path=...)`, then: + +1. **`create_feature_array_from_raw_photometry()`** (or `..._from_raw_spectra`, or plain `create_feature_array()`): converts the raw noiseless photometry grid into the training features. This is where all observational realism is injected: flux normalization (`normalize_method`), unit choice (AB mags, asinh mags, nJy), noise scattering via depths or empirical noise models (`scatter_fluxes=N` produces N noisy realizations per galaxy), errors as extra features, simulated missing bands, extra color features (`extra_features=['F090W - F115W']`, parsed by `FilterArithmeticParser`), and parameter add/drop/transformations. +2. **`run_single_sbi()`**: trains via LtU-ILI. Key knobs: `backend` ('sbi' or 'lampe'), `engine` ('NPE'/'NLE'/'NRE' + sequential variants), `model_type` ('mdn'/'maf'/'nsf'; pass a list for an ensemble), `n_nets`, `learning_type` ('offline' from the library, or 'online' with a `GalaxySimulator`), feature/target scalers (sklearn). Priors are built from the parameter array (`create_priors`). Saves to `/models//` by default: `*_posterior.pkl` (joblib), `*_params.pkl`, `*_summary.json`, plus any noise models as `*_empirical_noise_models.h5`. +3. **Evaluation**: `evaluate_model()`, `plot_coverage()`, `calculate_TARP()`, `calculate_PIT()`, `plot_diagnostics()`, `lc2st()`, `detect_misspecification()`. + +Reload a trained model with `SBI_Fitter.load_saved_model(...)` / `load_model_from_pkl(...)`. + +- **`optimize_sbi()`** runs Optuna hyperparameter searches via **`SBICustomRunner`** (`custom_runner.py`), a subclass of LtU-ILI's `SBIRunner` adding Optuna studies (optionally SQLite/MySQL-backed for multi-node searches), custom training loops, and `CustomUniform` torch priors with bijections. YAML study configs live in `examples/sbi/configs/`. +- **`Simformer_Fitter`** (subclass of `SBI_Fitter`, backed by the pure-PyTorch `src/synference/simformer/` subpackage) — score-diffusion transformer (Gloeckler et al. 2024) trained on the joint `[theta, x]` with per-example condition masks. Key difference from NPE: one trained model can condition on arbitrary subsets of features/parameters (boolean condition masks over nodes `[theta..., x...]`, True = observed), so it natively handles missing bands, acts as likelihood or posterior, and supports interval-constrained sampling via guidance (`sample_posterior_intervals`). The subpackage split: `nn.py` (tokenizer + transformer score net), `sde.py` (VE/VP SDEs; VE default), `masks.py` (condition/attention masks), `sampling.py` (fixed-step integrators + guidance + PF-ODE log-prob), `train.py` (`train_simformer`, config defaults), `model.py` (`SimformerModel` wrapper with `sample`/`sample_batched`/`sample_intervals`/`log_prob`/`save`/`load`). +- **`MissingPhotometryHandler`** — NPE alternative for missing data: KDE-imputes missing bands from nearest library neighbours (chi²-matched), then marginalizes posteriors over imputations. + +### Stage 3 — Inference on real data + +- `sample_posterior(obs_vector)` — posterior samples for one observation. +- `fit_catalogue(observations, columns_to_feature_names=..., flux_units=...)` — batch-fits an astropy Table / pandas DataFrame: builds features from catalogue columns (`create_features_from_observations` handles unit conversion and missing-data flags), runs out-of-distribution checks (pyod outlier ensembles + `test_in_distribution`), samples posteriors with per-row timeouts, returns quantiles appended to the input table, optionally recovers SEDs via the simulator. +- `recover_SED(obs_vector)` — pushes posterior samples back through the `GalaxySimulator` to reconstruct the SED implied by the fit. + +### Noise models (`src/synference/noise_models.py`) + +Class hierarchy under abstract `UncertaintyModel` (all serialize to/from HDF5 groups; `save_unc_model_to_hdf5` / `load_unc_model_from_hdf5`): + +- `DepthUncertaintyModel` — analytic scatter from n-sigma depths. +- `EmpiricalUncertaintyModel` / `AsinhEmpiricalUncertaintyModel` / `GeneralEmpiricalUncertaintyModel` — flux-dependent error distributions binned from real catalogues (e.g. `create_uncertainty_models_from_EPOCHS_cat`); the General variant also models upper limits and SNR-dependent behaviour. +- `SpectralUncertaintyModel` — error kernels for spectra. + +Noise models used in training are saved with the model so `fit_catalogue` applies consistent scattering at inference time. Pre-built models for COSMOS2020/COSMOS2025/JADES live in `priv/` (private, not shipped). + +## Repository layout notes + +- `src/synference/` — the whole package is 6 modules; `library.py` and `sbi_runner.py` contain nearly everything. +- `examples/` — the real documentation of usage patterns: `library_generation/scripts/` (incl. SLURM multi-node library builds), `sbi/` (training scripts, Optuna configs, SLURM), `online/`, `simformer/`, `paper/` (analysis notebooks). +- `models/` — trained model outputs (gitignored artifacts); `libraries/` — symlink to bulk library storage; `priv/` — private data/notebooks not for distribution. +- `tests/` — pytest suite; fixtures in `conftest.py` build small libraries from the Synthesizer test grid and expect the downloaded test data (see Commands). +- Docs are notebook-driven (`docs/source/`, nbsphinx); the docs CI build executes every notebook, so broken notebooks fail CI. + +## Conventions + +- Ruff enforces pydocstyle (Google convention) on `src/` — new public functions/classes need docstrings or CI fails. `examples/` is exempt from D-rules. +- Units are handled with `unyt` throughout; photometry conversions (AB ↔ Jy ↔ asinh) go through helpers in `utils.py` / `UncertaintyModel` staticmethods rather than ad-hoc math. +- Logging goes through the package logger from `utils.setup_mpi_named_logger("synference")` (MPI-rank aware); use `from . import logger`, not `print`. +- Anything serialized into libraries/models (parameter transforms, simulators) must be picklable — no lambdas. diff --git a/README.md b/README.md index acdfc23..7e5059c 100644 --- a/README.md +++ b/README.md @@ -35,7 +35,6 @@ Synference requires Python 3.10 or higher. It has the following dependencies: - [unyt](https://unyt.readthedocs.io/) for unit handling - [matplotlib](https://matplotlib.org/) for plotting and visualization - [tqdm](https://tqdm.github.io/) for progress bars -- [jax](https://jax.readthedocs.io/) for GPU acceleration (optional, for some inference models) These dependencies will be automatically installed when you install Synference using pip. diff --git a/docs/source/advanced_topics/simformer.ipynb b/docs/source/advanced_topics/simformer.ipynb index cdd0364..63928a3 100644 --- a/docs/source/advanced_topics/simformer.ipynb +++ b/docs/source/advanced_topics/simformer.ipynb @@ -2,108 +2,81 @@ "cells": [ { "cell_type": "markdown", - "id": "2291bddd", "metadata": {}, - "source": [ - "### The Simformer\n", - "\n", - "[Gloeckler et al. 2024](https://arxiv.org/abs/2404.09636) introduced the Simformer, for 'all in one simulation based inference'.\n", - "\n", - "They use a novel probablistic diffusion model with a transformer architecture which learns the full joint distribution of parameters and data, allowing for fast, amortized Bayesian inference, without specifying beforehand which parameters are of interest. This makes the simformer approach particularly well suited to missing data, as the use of an attention mechanism allows the model sample from an arbitrary conditional distribution excluding any missing data." - ] + "source": "### The Simformer\n\n[Gloeckler et al. 2024](https://arxiv.org/abs/2404.09636) introduced the Simformer, for 'all in one simulation based inference'.\n\nIt combines a probabilistic diffusion model with a transformer architecture which learns the full joint distribution of parameters and data, allowing for fast, amortized Bayesian inference without specifying beforehand which parameters are of interest. Because the model is trained with per-example *condition masks*, a single trained network provides the posterior, the likelihood, and any other conditional. This makes the Simformer particularly well suited to missing data: at inference time you simply mark missing bands as unobserved in the condition mask.\n\nSynference ships a native PyTorch implementation in `synference.simformer` (no JAX required), exposed through the `Simformer_Fitter` class, which mirrors the standard `SBI_Fitter` workflow." }, { "cell_type": "markdown", - "id": "12ddb0a9", "metadata": {}, - "source": [ - "The Simformer is currently implemented in two ways. The first way, is a seperate class called ```Simformer_Fitter```, which requires the user to install a [fork of the original simformer repo](https://github.com/tHarvey303/simformer/), which requires quite specific versions of CUDA, PyTorch and jax to work. \n", - "\n", - "The second way uses the new simformer implementation in the sbi package, which is currently only available in a pull request, but should be merged into the main branch soon. For now we will deal with this implementation, as it is much easier to install and use. There are examples of using the original approach in the examples/simformer folder of the synference repo.\n", - "\n", - "There are some limitations to the current sbi simformer implementation. Currently it doesn't seem to support serialization (due to the use of lambda functions), so models cannot be saved and loaded like other synference SBI models." - ] + "source": "We start from a photometric library exactly as for any other synference model." }, { "cell_type": "code", + "metadata": {}, "execution_count": null, - "id": "6284d69d", + "outputs": [], + "source": "from synference import Simformer_Fitter, test_data_dir\n\nfitter = Simformer_Fitter.init_from_hdf5(\n model_name=\"simformer_docs_example\",\n hdf5_path=f\"{test_data_dir}/example_model_library.hdf5\",\n)\nfitter.create_feature_array();" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "### Training\n\nTraining is run with `run_single_sbi`, matching the usual synference API. The defaults follow the Simformer paper: a variance-exploding SDE, a 6-layer transformer, and 'structured random' condition-mask sampling (a mixture of joint, posterior, likelihood, and random masks so all conditionals are learned).\n\nThe three override dictionaries control the model, the SDE, and the training loop; unknown keys raise an error rather than being silently ignored. The available keys (and their defaults) are in `synference.simformer.DEFAULT_MODEL_CONFIG`, `DEFAULT_SDE_CONFIG`, and `DEFAULT_TRAIN_CONFIG`. Some useful ones:\n\n1. `sde_config_dict_overrides`: `{\"name\": \"vesde\"}` (default) or `{\"name\": \"vpsde\", \"beta_min\": 0.01, \"beta_max\": 10.0}`.\n2. `model_config_dict_overrides`: `num_layers`, `num_heads`, `token_dim`, `attn_size`, `widening_factor`, `time_embedding_dim`.\n3. `train_config_dict_overrides`: `learning_rate`, `training_batch_size`, `min_number_steps` / `max_number_steps`, `condition_mask_fn`, `validation_fraction`.\n4. `attention_mask_type`: `\"full\"` (dense attention, default) or `\"directed\"` for a structured attention mask where parameters attend only to themselves and data attend to parameters.\n\nHere we use a tiny configuration so the notebook runs quickly \u2014 for real models leave the step counts at their defaults." + }, + { + "cell_type": "code", "metadata": {}, + "execution_count": null, "outputs": [], - "source": [ - "from synference import SBI_Fitter, test_data_dir\n", - "\n", - "fitter = SBI_Fitter.init_from_hdf5(\n", - " model_name=\"test\", hdf5_path=f\"{test_data_dir}/example_model_library.hdf5\"\n", - ")" - ] + "source": "model, stats = fitter.run_single_sbi(\n name_append=\"simformer_docs\",\n load_existing_model=False,\n model_config_dict_overrides={\"num_layers\": 2},\n train_config_dict_overrides={\n \"min_number_steps\": 500,\n \"max_number_steps\": 500,\n \"training_batch_size\": 128,\n },\n num_posterior_draws_per_sample=100,\n evaluate_model=False,\n verbose=False,\n)" }, { "cell_type": "markdown", - "id": "88192510", "metadata": {}, - "source": [ - "We will create our training array as normal using the ```SBI_Fitter``` class." - ] + "source": "### Sampling arbitrary conditionals\n\n`sample_posterior` works as for any other fitter, and returns samples of shape `(n_obs, num_samples, n_params)`. The `condition_mask` argument selects *which* conditional is sampled: it is a boolean array over the nodes `[parameters..., features...]` where `True` marks an observed quantity. The default `\"full\"` conditions on all features (the standard posterior). To handle missing bands, set the corresponding feature entries to `False` and drop those columns from the input." }, { "cell_type": "code", + "metadata": {}, "execution_count": null, - "id": "96c2ff30", + "outputs": [], + "source": "import numpy as np\n\nx_obs = fitter.feature_array[:2]\nsamples = fitter.sample_posterior(x_obs, num_samples=200, num_steps=100)\nprint(samples.shape) # (2, 200, n_params)\n\n# Condition on everything except the first photometric band (e.g. it is missing):\nn_theta = len(fitter.fitted_parameter_names)\nmask = np.array([False] * n_theta + [True] * len(fitter.feature_names))\nmask[n_theta] = False # first feature unobserved\nsamples_missing = fitter.sample_posterior(\n x_obs[:, 1:], num_samples=200, num_steps=100, condition_mask=mask\n)\nprint(samples_missing.shape) # (2, 200, n_params + 1): the missing band is inferred too" + }, + { + "cell_type": "markdown", "metadata": {}, + "source": "### Observation intervals\n\nThe Simformer also supports conditioning on *intervals* rather than exact values, using guidance ([Bansal et al. 2023](https://arxiv.org/abs/2302.07121)): at each step of the reverse diffusion, the clean data estimate is pushed towards the requested box. Here we constrain the first fitted parameter to a narrow range while conditioning on the photometry as usual." + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, "outputs": [], - "source": [ - "fitter.create_feature_array();" - ] + "source": "constraint_mask = np.zeros(n_theta + len(fitter.feature_names), dtype=bool)\nconstraint_mask[0] = True # constrain the first parameter (redshift)\n\nguided = fitter.sample_posterior_intervals(\n fitter.feature_array[0],\n constraint_mask=constraint_mask,\n a=[1.0], # lower bound, original units\n b=[3.0], # upper bound\n num_samples=200,\n num_steps=100,\n)\ninside = ((guided[:, 0] >= 1.0) & (guided[:, 0] <= 3.0)).mean()\nprint(f\"{inside:.0%} of samples inside the requested interval\")" }, { "cell_type": "markdown", - "id": "fe118ab9", "metadata": {}, - "source": [ - "There are a few parameters to be aware of:\n", - "\n", - "1. ```sde_type``` : The type of SDE to use. Options are 've` (variance exploding), 'vp' (variance preserving) or 'subvp' (sub variance preserving). This doesn't do anything for the flow based simformer.\n", - "2. ```simformer_type```: 'score' or 'flow'- whether to use a score based or flow based simformer. \n", - "3 ```learning_rate```: The learning rate to use for training.\n", - "4. ```model_kwargs```: A dictionary of additional keyword arguments to pass to the simformer model. These can include:\n", - " - ```num_layers```: The number of transformer layers to use.\n", - " - ```num_heads```: The number of attention heads to use.\n", - " - ```dim_val```: The dimension of the value vectors in the attention mechanism.\n", - " - ```dim_id```: The dimension of the identity vectors in the attention mechanism.\n", - " - ```mlp_ratio``` : The ratio of the hidden dimension to the input dimension in the MLP layers.\n", - " - ```hidden_features```: The number of hidden features to use in the MLP layers.\n", - " - ```time_embedding_dim```: The dimension of the time embedding.\n", - "\n", - "Like all other synference models we can also set the ```training_batch_size```, ```validation_fraction```, ```stop_after_epochs``` and ```clip_max_norm``` parameters.\n", - "\n", - "```python\n", - "\n", - "fitter.run_single_simformer(\n", - " name_append=\"simformer_test\",\n", - " sde_type=\"ve\",\n", - " simformer_type=\"score\",\n", - " learning_rate=1e-5,\n", - " training_batch_size=64,\n", - " model_kwargs={\n", - " \"hidden_features\": 128,\n", - " \"n_layers\": 6,\n", - " \"dim_val\": 64,\n", - " \"dim_id\": 64,\n", - " \"mlp_ratio\": 4,\n", - " \"time_embedding_dim\": 32,\n", - " \"num_heads\": 4,\n", - " },\n", - " load_existing_model=False,\n", - " validation_fraction=0.1,\n", - " stop_after_epochs=30,\n", - " plot=False, # Currently the LtU-ILI plotting doesn't work with the simformer\n", - ")\n", - "```" - ] + "source": "### Saving, loading and log-probabilities\n\nModels are serialized to a plain `torch` payload alongside the usual synference sidecar files, and can be reloaded in a fresh process with `Simformer_Fitter.load_saved_model(model_name, library_path, model_file)` or directly via `synference.simformer.SimformerModel.load(path)`.\n\nLog probabilities of parameters given observations are computed with the probability-flow ODE:" + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": "log_prob = fitter.log_prob(\n fitter.feature_array[:2], fitter.fitted_parameter_array[:2], num_steps=50\n)\nprint(log_prob)" } ], - "metadata": {}, + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.10" + } + }, "nbformat": 4, "nbformat_minor": 5 -} +} \ No newline at end of file diff --git a/docs/source/getting_started/installation.rst b/docs/source/getting_started/installation.rst index bf18665..4ad728b 100644 --- a/docs/source/getting_started/installation.rst +++ b/docs/source/getting_started/installation.rst @@ -71,7 +71,6 @@ The available groups are: - **Development** (``dev``): Tools to help developing including linting and formatting. - **Testing** (``test``): Frameworks and utilities for running tests. - **Documentation** (``docs``): Packages required to build the project documentation. -- **simformer** (``simformer``): Dependencies for using the Simformer model within Synference. For example, to install with development dependencies, run: diff --git a/examples/sbi/scripts/train_model.py b/examples/sbi/scripts/train_model.py index f083503..afe22af 100644 --- a/examples/sbi/scripts/train_model.py +++ b/examples/sbi/scripts/train_model.py @@ -392,7 +392,6 @@ def log10_floor(x, floor=-6): ) else: args = dict( - backend="jax", num_training_simulations=10_000, train_test_fraction=args.train_test_fraction, random_seed=42, @@ -401,7 +400,6 @@ def log10_floor(x, floor=-6): load_existing_model=True, name_append=args.name_append, save_method="joblib", - task_func=None, ) print("Running SBI training.") diff --git a/examples/simformer/scripts/fit_simformer_model.py b/examples/simformer/scripts/fit_simformer_model.py index 0880e10..b4f7a58 100644 --- a/examples/simformer/scripts/fit_simformer_model.py +++ b/examples/simformer/scripts/fit_simformer_model.py @@ -6,7 +6,7 @@ fitter = Simformer_Fitter.load_saved_model( model_name="simformer_v2_initial", - grid_path="/cosma/apps/dp276/dc-harv3/synference/libraries/grid_BPASS_Chab_DenseBasis_SFH_0.01_z_12_logN_5.0_CF00_v2.hdf5", + library_path="/cosma/apps/dp276/dc-harv3/synference/libraries/grid_BPASS_Chab_DenseBasis_SFH_0.01_z_12_logN_5.0_CF00_v2.hdf5", model_file="/cosma/apps/dp276/dc-harv3/synference/models/simformer_v2", ) diff --git a/examples/simformer/scripts/two_moons_demo.py b/examples/simformer/scripts/two_moons_demo.py new file mode 100644 index 0000000..8022216 --- /dev/null +++ b/examples/simformer/scripts/two_moons_demo.py @@ -0,0 +1,135 @@ +"""Two-moons demo of the native PyTorch Simformer. + +Trains a Simformer on the classic two-moons SBI benchmark and demonstrates: +posterior sampling (bimodal), likelihood and joint conditionals, interval-constrained +sampling via guidance, log probabilities, and save/load round-tripping. + +Run with: python two_moons_demo.py [--quick] +""" + +import argparse +import os + +import corner +import matplotlib.pyplot as plt +import numpy as np +import torch + +from synference.simformer import SimformerModel, train_simformer + + +def two_moons_simulator(theta, rng): + """Two-moons simulator (Lueckmann et al. 2021 parametrization).""" + n = theta.shape[0] + alpha = rng.uniform(-np.pi / 2, np.pi / 2, n) + r = rng.normal(0.1, 0.01, n) + p = np.stack([r * np.cos(alpha) + 0.25, r * np.sin(alpha)], axis=1) + shift = np.stack( + [-np.abs(theta[:, 0] + theta[:, 1]), (-theta[:, 0] + theta[:, 1])], axis=1 + ) / np.sqrt(2) + return p + shift + + +def main(): + """Train and exercise a Simformer on the two-moons task.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--quick", action="store_true", help="Small training budget.") + parser.add_argument("--out-dir", default="two_moons_output", help="Output directory.") + args = parser.parse_args() + os.makedirs(args.out_dir, exist_ok=True) + + rng = np.random.default_rng(42) + n_train = 10_000 if args.quick else 100_000 + steps = 3_000 if args.quick else 30_000 + + theta = rng.uniform(-1, 1, (n_train, 2)) + x = two_moons_simulator(theta, rng) + + model, stats = train_simformer( + theta, + x, + train_config={ + "min_number_steps": steps, + "max_number_steps": steps, + "training_batch_size": 512 if args.quick else 1000, + }, + seed=0, + ) + print(f"Trained for {stats['steps_run']} steps in {stats['training_time_s']:.0f} s.") + + model_path = os.path.join(args.out_dir, "two_moons_simformer.pt") + model.save(model_path) + model = SimformerModel.load(model_path) # prove the round trip + print(f"Model saved to and reloaded from {model_path}.") + + gen = torch.Generator().manual_seed(1) + posterior_mask = np.array([False, False, True, True]) + x_o = [0.0, 0.0] + + # 1. Posterior at the origin: two crescent-shaped modes. + samples = model.sample( + 20_000, x_o=x_o, condition_mask=posterior_mask, num_steps=200, generator=gen + ) + fig = corner.corner( + samples, + labels=[r"$\theta_1$", r"$\theta_2$"], + bins=80, + plot_datapoints=False, + plot_density=True, + ) + fig.suptitle("Two-moons posterior at x = (0, 0)") + fig.savefig(os.path.join(args.out_dir, "posterior_corner.png"), dpi=150) + plt.close(fig) + + # 2. Interval-guided posterior: constrain theta_1 to [0, 0.6]. + constraint = np.array([True, False, False, False]) + guided = model.sample_intervals( + 20_000, + x_o=x_o, + condition_mask=posterior_mask, + constraint_mask=constraint, + a=[0.0], + b=[0.6], + scale_bias=1e-2, + num_steps=200, + generator=gen, + ) + inside = ((guided[:, 0] >= 0) & (guided[:, 0] <= 0.6)).mean() + fig = corner.corner( + guided, + labels=[r"$\theta_1$", r"$\theta_2$"], + bins=80, + plot_datapoints=False, + plot_density=True, + ) + fig.suptitle(rf"Guided posterior, $\theta_1 \in [0, 0.6]$ ({inside:.1%} inside)") + fig.savefig(os.path.join(args.out_dir, "posterior_interval_corner.png"), dpi=150) + plt.close(fig) + + # 3. Likelihood conditional x | theta and the joint. + likelihood_mask = np.array([True, True, False, False]) + x_samples = model.sample( + 5_000, x_o=[0.4, -0.3], condition_mask=likelihood_mask, num_steps=200, generator=gen + ) + x_true = two_moons_simulator(np.tile([0.4, -0.3], (5_000, 1)), rng) + fig, ax = plt.subplots(figsize=(5, 5)) + ax.scatter(*x_true.T, s=2, alpha=0.3, label="Simulator") + ax.scatter(*x_samples.T, s=2, alpha=0.3, label="Simformer likelihood") + ax.legend() + ax.set_xlabel(r"$x_1$") + ax.set_ylabel(r"$x_2$") + fig.savefig(os.path.join(args.out_dir, "likelihood_conditional.png"), dpi=150) + plt.close(fig) + + # 4. Log probabilities of the true parameters for a few test pairs. + test_theta = rng.uniform(-1, 1, (5, 2)) + test_x = two_moons_simulator(test_theta, rng) + for i in range(5): + lp = model.log_prob(test_theta[i], test_x[i], posterior_mask, num_steps=100) + print(f"log p(theta_true | x_{i}) = {float(lp):.2f} (prior: {-np.log(4):.2f})") + + print(f"Plots written to {args.out_dir}/.") + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index cb2bd68..3ad8ad5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,7 +61,6 @@ dependencies = [ "numpy", "scipy", "unyt", - "jax", "simple_parsing", "corner", "plotext", diff --git a/src/synference/__init__.py b/src/synference/__init__.py index 6140915..bbc384e 100644 --- a/src/synference/__init__.py +++ b/src/synference/__init__.py @@ -36,6 +36,8 @@ from .custom_runner import SBICustomRunner +from .simformer import SimformerModel, UniformBoxPrior, train_simformer + try: from .sbi_runner import SBI_Fitter, MissingPhotometryHandler, Simformer_Fitter except ImportError as e: @@ -43,7 +45,6 @@ print('Dependencies for SBI not installed. Only the library generation functions will be available.') -#from .simformer import UncertainityModelTask warnings.filterwarnings('ignore') __all__ = [ @@ -78,6 +79,9 @@ "SBI_Fitter", "MissingPhotometryHandler", "Simformer_Fitter", + "SimformerModel", + "UniformBoxPrior", + "train_simformer", "GalaxySimulator", "generate_random_DB_sfh", "EmpiricalUncertaintyModel", diff --git a/src/synference/sbi_runner.py b/src/synference/sbi_runner.py index 1409c8b..d17da58 100644 --- a/src/synference/sbi_runner.py +++ b/src/synference/sbi_runner.py @@ -17,8 +17,6 @@ from typing import Any, Callable, Dict, List, Optional, Tuple, Union import ili -import jax -import jax.numpy as jnp import matplotlib.pyplot as plt import numpy as np import optuna @@ -64,6 +62,7 @@ load_unc_model_from_hdf5, save_unc_model_to_hdf5, ) +from .simformer import SimformerModel, UniformBoxPrior, train_simformer from .utils import ( FilterArithmeticParser, TimeoutException, @@ -3236,7 +3235,9 @@ def fit_catalogue( num_samples=num_samples, **kwargs, ) - samples_quant = samples + if samples.ndim == 2: + samples = samples[None, ...] + samples_quant = samples.transpose(2, 0, 1) else: samples = self.sample_posterior( @@ -4061,335 +4062,6 @@ def timeout_handler(signum, frame): if old_handler is not None: signal.signal(signal.SIGALRM, old_handler) - def run_single_simformer( - self, - train_test_fraction: float = 0.8, - random_seed: int = None, - train_indices: np.ndarray = None, - test_indices: np.ndarray = None, - save_model: bool = True, - verbose: bool = True, - out_dir: str = f"{code_path}/models/name/", - plot: bool = True, - name_append: str = "timestamp", - save_method: str = "dill", - set_self: bool = True, - override_prior_ranges: dict = {}, - load_existing_model: bool = True, - use_existing_indices: bool = True, - evaluate_model: bool = True, - max_num_epochs: int = 100_000, - sde_type: str = "ve", - simformer_type="score", - model_kwargs: dict = {}, - learning_rate: float = 0.0005, - training_batch_size: int = 200, - validation_fraction: float = 0.1, - stop_after_epochs: int = 20, - clip_max_norm: float = 5.0, - training_args: dict = {}, - ) -> tuple: - r"""Trains a single Simformer model using the SBI implementation. - - May need to be on this branch: - https://github.com/sbi-dev/sbi/pull/1621/ - - Parameters: - train_test_fraction: Fraction of the dataset to be used for training. - random_seed: Random seed for reproducibility. - train_indices: Indices of the training set. - test_indices: Indices of the test set. If None, no test set is used. - save_model: Whether to save the trained model. - verbose: Whether to print verbose output. - out_dir: Directory to save the model. - plot: Whether to plot the results. - name_append: String to append to the model name. - set_self: Whether to set the self attribute. - override_prior_ranges: Dictionary to override prior ranges. - load_existing_model: Whether to load an existing model. - use_existing_indices: Whether to use existing indices. - evaluate_model: Whether to evaluate the model. - max_num_epochs: Maximum number of epochs to train the model. - sde_type: Type of SDE to use ('ve','vp' or 'subvp'). Not used for flow matching. - simformer_type: Type of Simformer to use ('score' or 'flow'). - model_kwargs: Additional keyword arguments to pass to the Simformer builder. - Available kwargs and defaults are: - For both 'score' and 'flow': - - hidden_features: int = 100, - - num_heads: int = 4, - - num_layers: int = 4, - - mlp_ratio: int = 2, - - time_embedding_dim: int = 32, - - embedding_net: nn.Module = nn.Identity(), - - dim_val: int = 64, - - dim_id: int = 32, - - dim_cond: int = 16, - - ada_time: bool = False, - - \**kwargs: Any, - learning_rate: Learning rate for the optimizer. - training_batch_size: Batch size for training. - validation_fraction: Fraction of the training set to use for validation. - stop_after_epochs: Number of epochs without improvement before stopping. - clip_max_norm: Maximum norm for gradient clipping. - training_args: Additional arguments to pass to the training function. - """ - from sbi.inference import FlowMatchingSimformer, Simformer - - assert self.has_features, ( - "Feature array not created. Please create the feature array first." - ) - - if self.fitted_parameter_array is None: - raise ValueError("Parameter grid not created. Please create the parameter grid first.") - - if name_append == "timestamp": - name_append = self._timestamp - - out_dir = os.path.join(os.path.abspath(out_dir), self.name) - - if not os.path.exists(out_dir): - os.makedirs(out_dir) - - run = False - if ( - os.path.exists(f"{out_dir}/{self.name}_{name_append}_params.pkl") - and load_existing_model - ): - logger.info( - f"Loading existing model from {out_dir}/{self.name}_{name_append}_params.pkl" # noqa: E501 - ) - posterior, stats, params = self.load_model_from_pkl( - f"{out_dir}/{self.name}_{name_append}_posterior.pkl", - set_self=set_self, - ) - # return posterior, stats - run = True - - if params is not None: - save_model = False # Don't save the model again if we loaded it. - elif os.path.exists(f"{out_dir}/{self.name}_{name_append}_simformer.pkl"): - logger.info( - "Model with same name already exists. \ - Please change the name of this model or delete the existing one." - ) - return None - - if ( - use_existing_indices - and (self._train_indices is not None) - and (self._test_indices is not None) - ): - train_indices = self._train_indices - test_indices = self._test_indices - logger.info("Using existing train and test indices.") - - if (train_indices is None) or (test_indices is None): - train_indices, test_indices = self.split_dataset( - train_fraction=train_test_fraction, - random_seed=random_seed, - verbose=verbose, - ) - - if set_self: - self._train_indices = train_indices - self._test_indices = test_indices - - X_train = self.feature_array[train_indices] - y_train = self.fitted_parameter_array[train_indices] - X_test = self.feature_array[test_indices] - y_test = self.fitted_parameter_array[test_indices] - - # Construct simformer unflattened features - - X_train = torch.tensor(X_train, dtype=torch.float32, device=self.device) - y_train = torch.tensor(y_train, dtype=torch.float32, device=self.device) - X_test = torch.tensor(X_test, dtype=torch.float32, device=self.device) - y_test = torch.tensor(y_test, dtype=torch.float32, device=self.device) - - inputs = torch.cat([y_train, X_train], dim=1) - - prior = self.create_priors( - override_prior_ranges=override_prior_ranges, - verbose=verbose, - ) - - if not run: - simformer = Simformer if simformer_type == "score" else FlowMatchingSimformer - - from torch.utils.tensorboard import SummaryWriter - - summary_writer = ( - SummaryWriter(log_dir=f"{out_dir}/logs/{self.name}_{name_append}") - if verbose - else None - ) - - # model_kwargs - passed to model_builder. - inference = simformer( - device=self.device, - sde_type=sde_type, - logging_level="INFO", - show_progress_bars=verbose, - summary_writer=summary_writer if verbose else None, - **model_kwargs, - ) - - training_args_default = { - "training_batch_size": training_batch_size, - "learning_rate": learning_rate, - "validation_fraction": validation_fraction, - "stop_after_epochs": stop_after_epochs, - "max_num_epochs": max_num_epochs, - "clip_max_norm": clip_max_norm, - "force_first_round_loss": False, - "discard_prior_samples": False, - "retrain_from_scratch": False, - "show_train_summary": True, - "calibration_kernel": None, - "ema_loss_decay": 0.1, - "validation_times": 10, - "dataloader_kwargs": None, - } - training_args_default.update(training_args) - - inference.append_simulations(inputs, data_device=self.device) - - start_time = time.time() - inference.train(**training_args_default) - end_time = time.time() - # use torch to save score_estimator - - stats = [inference._summary] - - try: - torch.save(inference, f"{out_dir}/{self.name}_{name_append}_simformer.pkl") - except Exception as e: - logger.error( - f"Error saving simformer " - f"to {out_dir}/{self.name}_{name_append}_simformer.pkl: {e}" - ) - - inference.set_condition_indexes( - new_posterior_latent_idx=list(range(len(self.fitted_parameter_names))), - new_posterior_observed_idx=list( - range(len(self.fitted_parameter_names), inputs.shape[1]) - ), - ) - posterior = inference.build_posterior(prior=prior) - - """ This would be nice, but DirectPosterior is not compatible with Simformer yet - it seems (no support for .parameters() method) - posterior = DirectPosterior(posterior_estimator=posterior, prior=prior) - posterior = EnsemblePosterior( - posteriors=[posterior], - weights=torch.tensor([1.0], device=self.device), - theta_transform=posterior.theta_transform, - ) - """ - # Save the posterior in a compatible format. - try: - with open(f"{out_dir}/{self.name}_{name_append}_posterior.pkl", "wb") as f: - if save_method == "dill": - import dill - - dill.dump(posterior, f) - elif save_method == "joblib": - dump(posterior, f, compress=3) - elif save_method == "pickle": - pickle.dump(posterior, f) - else: - torch.save(posterior, f"{out_dir}/{self.name}_{name_append}_posterior.pkl") - except Exception as e: - logger.error( - f"Error saving posterior {out_dir}/{self.name}_{name_append}_posterior.pkl: {e}" - ) - - if set_self: - self.simformer = inference - self.posteriors = posterior - self._prior = prior - self.stats = stats - self._X_train = X_train.cpu().numpy() - self._y_train = y_train.cpu().numpy() - self._X_test = X_test.cpu().numpy() - self._y_test = y_test.cpu().numpy() - - if save_model: - param_dict = { - "train_indices": train_indices, - "test_indices": test_indices, - "sde_type": sde_type, - "simformer_type": simformer_type, - "model_kwargs": model_kwargs, - "random_seed": random_seed, - "training_time": end_time - start_time, - "training_args": training_args_default, - "prior": prior, - "stats": stats, - } - self.save_state( - out_dir=out_dir, - name_append=name_append, - save_method=save_method, - has_grid=True, - **param_dict, - ) - - if plot: - if verbose: - logger.info("Plotting training diagnostics...") - self.plot_diagnostics( - X_train=X_train.cpu().numpy(), - y_train=y_train.cpu().numpy(), - X_test=X_test.cpu().numpy(), - y_test=y_test.cpu().numpy(), - plots_dir=f"{out_dir}/plots/{name_append}/", - stats=None, - sample_method="direct", - posteriors=posterior, - ) - # Evaluate the model - if evaluate_model: - metrics_path = f"{out_dir}/{self.name}_{name_append}_summary.json" - if verbose: - logger.info("Evaluating the model...") - stats = self.evaluate_model( - posteriors=posterior, - X_test=X_test, - y_test=y_test, - ) - - try: - with open(metrics_path, "w") as f: - json.dump(stats, f, indent=4) - except Exception as e: - logger.error(f"Error saving metrics to {metrics_path}: {e}") - - if set_self: - self.stats = stats - - return inference - - # Build conditional just lets you set whichever parameters are missing. - # Build_posterior and build_likelihood are just wrappers around build_conditional. - - # conditional = inference.build_conditional(condition_mask=[False, True]) - # conditional_samples = conditional.sample((10000,), x=x_o) - - # Must set indexes here to indicate which are observed and which are latent. - - # Set condition indexes properly from len(self.fitted_param_names) - # and len(self.feature_names) - - # inference.set_condition_indexes(new_posterior_latent_idx=[0], - # new_posterior_observed_idx=[1]) - - # - # posterior_samples = posterior.sample((10000,), x=x_o) - - # likelihood = inference.build_likelihood() - # likelihood_samples = likelihood.sample((10000,), x=theta_o) - def run_single_sbi( self, train_test_fraction: float = 0.8, @@ -6668,7 +6340,7 @@ def evaluate_model( for metric in metrics: # print(metric, type(metrics[metric])) if ( - isinstance(metrics[metric], (np.ndarray, list, torch.Tensor, jnp.ndarray)) + isinstance(metrics[metric], (np.ndarray, list, torch.Tensor)) and hasattr(metrics[metric], "__len__") and len(metrics[metric]) == len(self.fitted_parameter_names) ): @@ -8039,28 +7711,27 @@ def __init__( class Simformer_Fitter(SBI_Fitter): - """Simformer Fitter for SBI models. + """SBI fitter backed by the native PyTorch Simformer. - This class implements the Simformer architecture for SBI tasks. - - To Do: - - Ensure more inherited methods actually work with Simformer. - - Implement TARP and other metrics which are native to LtU-ILI. - - Ensure methods to act on real observations and recover photometry work. + The Simformer (Gloeckler et al. 2024) is a transformer + score-diffusion model + trained on the joint ``[theta, x]`` with per-example condition masks. A single + trained model provides the posterior, the likelihood, and arbitrary conditionals + (e.g. missing photometric bands), plus interval-constrained sampling via guidance. + The trained model is a :class:`synference.simformer.SimformerModel` stored on + ``self.posteriors``. Node order everywhere is ``[theta..., x...]`` and condition + masks are boolean arrays where True marks an observed (conditioned) node. """ def __init__(self, name: str = "simformer_fitter", **kwargs): """Initialize the Simformer Fitter.""" super().__init__(name=name, **kwargs) - self.simformer_task = None - @classmethod def init_from_hdf5( cls, model_name, hdf5_path: str = None, return_output: bool = False, **kwargs ): - """Initialize the Simformer Fitter.""" + """Initialize the Simformer Fitter from an HDF5 library.""" return super().init_from_hdf5( model_name, hdf5_path, @@ -8070,7 +7741,7 @@ def init_from_hdf5( @classmethod def load_saved_model(cls, model_name: str, library_path: str, model_file: str, **kwargs): - """Load a saved Simformer model from a file.""" + """Load a saved Simformer model given a library and a model directory.""" model = cls.init_from_hdf5( model_name=model_name, hdf5_path=library_path, @@ -8087,7 +7758,6 @@ def load_saved_model(cls, model_name: str, library_path: str, model_file: str, * def run_single_sbi( self, - backend: str = "jax", num_training_simulations: int = 10_000, train_test_fraction: float = 0.9, random_seed: int = 42, @@ -8096,23 +7766,25 @@ def run_single_sbi( load_existing_model: bool = True, name_append: str = "timestamp", save_method: str = "joblib", - task_func: Optional[Callable] = None, model_config_dict_overrides: Optional[Dict[str, Any]] = None, - sde_config_dict: Optional[Dict[str, Any]] = None, - train_config_dict_overrides: Optional[Dict[str, Any]] = None, sde_config_dict_overrides: Optional[Dict[str, Any]] = None, - attention_mask_type: str = "full", + train_config_dict_overrides: Optional[Dict[str, Any]] = None, + attention_mask_type: Union[str, np.ndarray] = "full", evaluate_model: bool = True, + out_dir: str = f"{code_path}/models/", + num_posterior_draws_per_sample: int = 1000, + device: Optional[str] = None, ): - """Train a Simformer model using the provided configurations. + """Train a Simformer model on the feature/parameter arrays. - This method sets up the Simformer task, prepares the data, trains - the model using the specified configurations, and saves the trained - model to a pickle file. + Trains with denoising score matching over the joint ``[theta, x]`` using the + default paper configuration (VE-SDE, structured-random condition masks); see + :data:`synference.simformer.DEFAULT_MODEL_CONFIG`, + :data:`synference.simformer.DEFAULT_SDE_CONFIG`, and + :data:`synference.simformer.DEFAULT_TRAIN_CONFIG` for the available override + keys (unknown keys raise a ``ValueError``). Args: - backend (str, optional): Backend to use for training ('jax' or 'torch'). - Defaults to "jax". num_training_simulations (int, optional): Number of training simulations to generate. Only used if `has_simulator` is True, otherwise the `feature_array` is used. Defaults to 10,000. @@ -8120,266 +7792,178 @@ def run_single_sbi( training. Defaults to 0.9. random_seed (int, optional): Random seed for reproducibility. Defaults to 42. - set_self (bool, optional): If True, sets the trained model and task - to the instance attributes. Defaults to True. + set_self (bool, optional): If True, sets the trained model to the + instance attributes. Defaults to True. verbose (bool, optional): If True, prints progress information during training. Defaults to True. load_existing_model (bool, optional): If True, loads an existing - model from a pickle file if it exists. Defaults to True. + model from disk if it exists. Defaults to True. name_append (str, optional): String to append to the model name in the output file. Defaults to "timestamp". - save_method (str, optional): Method to use for saving the model - (e.g., 'torch', 'pickle'). Defaults to "joblib". - task_func (Callable, optional): Function to create the task. If None, - uses the default `GalaxyPhotometryTask`. Defaults to None. - model_config_dict_overrides (dict, optional): Dictionary to override - the default model configuration. Defaults to None. - sde_config_dict (dict, optional): Dictionary to override configs - for the SDE. Defaults to a pre-defined VPSDE configuration. - train_config_dict_overrides (dict, optional): Dictionary to override the - training configuration. Defaults to a pre-defined configuration. - sde_config_dict_overrides (dict, optional): Dictionary to override - the SDE configuration. Defaults to None. - attention_mask_type (str, optional): Type of attention mask to use - ('full', 'causal', or 'none'). Defaults to "full". - evaluate_model (bool, optional): If True, evaluates the trained - model on the validation set and prints metrics. Defaults to True. - - - - Add support for constraint functions during sampling to allow intervals - (e.g., from `scoresbibm.methods.guidance` import - `generalized_guidance`, `get_constraint_fn`). - """ - from omegaconf import OmegaConf - from scoresbibm.methods.score_transformer import train_transformer_model - - model_config_dict = { - "name": "ScoreTransformer", - "d_model": 128, - "n_heads": 4, - "n_layers": 4, - "d_feedforward": 256, - "dropout": 0.1, - "max_len": 5000, # Adjust based on theta_dim + x_dim - "tokenizer": {"name": "LinearTokenizer", "encoding_dim": 64}, - "use_output_scale_fn": True, - } - if model_config_dict_overrides is not None: - model_config_dict.update(model_config_dict_overrides) - - sde_config_dict = { - "name": "VPSDE", # or "VESDE" - "beta_min": 0.1, - "beta_max": 20.0, - "num_steps": 1000, - "T_min": 1e-05, - "T_max": 1.0, - } - if sde_config_dict_overrides is not None: - sde_config_dict.update(sde_config_dict_overrides) - - train_config_dict = { - "learning_rate": 1e-4, # Initial learning rate for training # used - "min_learning_rate": 1e-6, # Minimum learning rate for training # used - "z_score_data": True, # Whether to z-score the data # used - "total_number_steps_scaling": 5, # Scaling factor for total number of steps - "max_number_steps": 1e8, # Maximum number of steps for training # used - "min_number_steps": 1e4, # Minimum number of steps for training # used - "training_batch_size": 64, # Batch size for training # used - "val_every": 100, # Validate every 100 steps # used - "clip_max_norm": 10.0, # Gradient clipping max norm # used - "condition_mask_fn": { - "name": "joint" - }, # Use the base mask function defined in the task - "edge_mask_fn": {"name": "none"}, - "validation_fraction": 0.1, # Fraction of data to use for validation # used - "val_repeat": 5, # Number of times to repeat validation # used - "stop_early_count": 5, # Number of steps to wait before stopping early # used - "rebalance_loss": False, # Whether to rebalance the loss # used - } - if train_config_dict_overrides is not None: - train_config_dict.update(train_config_dict_overrides) - - if task_func is None: - from .simformer import GalaxyPhotometryTask as task_func + save_method (str, optional): Method used for the sidecar params file + (the model itself is always saved with torch). Defaults to "joblib". + model_config_dict_overrides (dict, optional): Overrides for the model + architecture config. Defaults to None. + sde_config_dict_overrides (dict, optional): Overrides for the SDE config + (e.g. ``{"name": "vpsde", "beta_max": 20.0}``). Defaults to None. + train_config_dict_overrides (dict, optional): Overrides for the training + config. Defaults to None. + attention_mask_type (str or np.ndarray, optional): Base attention mask — + "full" (dense), "directed"/"causal", or a custom boolean array. + Defaults to "full". + evaluate_model (bool, optional): If True, evaluates the trained model on + the test split and saves diagnostic plots/metrics. Defaults to True. + out_dir (str, optional): Root output directory; the model is written to + ``{out_dir}/{self.name}/``. Defaults to ``{code_path}/models/``. + num_posterior_draws_per_sample (int, optional): Posterior draws per test + observation during evaluation. Defaults to 1000. + device (str, optional): Torch device for training. Defaults to + ``self.device``. + Returns: + tuple: The trained :class:`synference.simformer.SimformerModel` and the + training stats dictionary. + """ if name_append == "timestamp": name_append = f"{self._timestamp}" - if len(name_append) > 0 and name_append[0] != "_": name_append = f"_{name_append}" - out_path = f"{code_path}/models/{self.name}/{self.name}{name_append}_posterior.pkl" - if load_existing_model and os.path.exists(out_path): - logger.info(f"Loading existing model from {out_path}") - trained_score_model, meta = self.load_model_from_pkl( - model_dir=f"{code_path}/models/{self.name}/", - model_name=f"{self.name}{name_append}_posterior", - set_self=True, - ) - run = True - else: - run = False + out_dir = os.path.join(os.path.abspath(out_dir), self.name) + model_path = os.path.join(out_dir, f"{self.name}{name_append}_posterior.pkl") + + if device is None: + device = self.device - priors = self.create_priors() + prior = self.create_priors(set_self=True) if self.has_simulator: - logger.info("Using online simulator for training.") - simulator_function = self.simulator - learning_type = "online" + logger.info(f"Generating {num_training_simulations} simulations for training.") + x_all, theta_all = self.generate_pairs_from_simulator(num_training_simulations) + n_train = int(train_test_fraction * len(x_all)) + rng = np.random.default_rng(random_seed) + order = rng.permutation(len(x_all)) + train_idx, test_idx = order[:n_train], order[n_train:] else: - logger.info("Using pre-generated samples for training.") assert self.feature_array is not None, ( "Feature array must be provided for pre-generated samples." ) - simulator_function = None - learning_type = "offline" - - if not self.has_simulator: - # Split the dataset into training and validation sets. - train_indices, test_indices = self.split_dataset( + train_idx, test_idx = self.split_dataset( train_fraction=train_test_fraction, random_seed=random_seed, verbose=verbose, ) + x_all = np.asarray(self.feature_array, dtype=np.float32) + theta_all = np.asarray(self.fitted_parameter_array, dtype=np.float32) - x = jnp.array(self.feature_array, dtype=jnp.float32) - theta = jnp.array(self.fitted_parameter_array, dtype=jnp.float32) + x_train, theta_train = x_all[train_idx], theta_all[train_idx] + x_test, theta_test = x_all[test_idx], theta_all[test_idx] - training_data = { - "theta": theta[train_indices], - "x": x[train_indices], - } - - validation_data = { - "theta": theta[test_indices], - "x": x[test_indices], - } - - task = task_func( - name="galaxy_photometry", - backend=backend, - prior_dict=priors, - param_names_ordered=self.fitted_parameter_names, - run_simulator_fn=simulator_function, - num_filters=len(self.feature_names), - test_X_data=copy.deepcopy(validation_data["x"]), - test_theta_data=copy.deepcopy(validation_data["theta"]), - attention_mask_type=attention_mask_type, - ) - - method_config_dict = { - "device": str(self.device), # Ensure this matches device setup - "sde": sde_config_dict, - "model": model_config_dict, - "train": train_config_dict, + meta = { + "param_names": list(self.fitted_parameter_names), + "feature_names": list(self.feature_names), + "prior_ranges": prior.prior_ranges, } - # Convert the main method_cfg to OmegaConf DictConfig - method_cfg = OmegaConf.create(method_config_dict) - master_rng_key = jax.random.PRNGKey(random_seed) - - if self.has_simulator: - logger.info(f"Generating {num_training_simulations} training simulations...") - training_data = task.get_data(num_samples=num_training_simulations) - - num_validation_simulations = int(num_training_simulations * train_test_fraction) - - validation_data = task.get_data(num_samples=num_validation_simulations) - - if not run: + stats = None + if load_existing_model and os.path.exists(model_path): + logger.info(f"Loading existing model from {model_path}") + trained_model, _ = self.load_model_from_pkl( + model_dir=out_dir, + model_name=f"{self.name}{name_append}_posterior", + set_self=set_self, + ) + else: if verbose: logger.info(f"Starting training at {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}") - trained_score_model = train_transformer_model( - task=task, - data=training_data, - method_cfg=method_cfg, - rng=master_rng_key, + trained_model, stats = train_simformer( + theta_train, + x_train, + model_config=model_config_dict_overrides, + sde_config=sde_config_dict_overrides, + train_config=train_config_dict_overrides, + base_mask=attention_mask_type, + meta=meta, + device=device, + seed=random_seed, + verbose=verbose, ) if verbose: logger.info(f"Finished training at {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}") if set_self: - self.simformer_task = task - self.posteriors = trained_score_model - - if verbose: - logger.info(f"Saving model at {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}") - - self.save_model_to_pkl( - task=task, - posteriors=trained_score_model, - output_folder=f"{code_path}/models/{self.name}/", - name_append=name_append, - method_config_dict=method_config_dict, - save_method=save_method, - extras={ - "model_config_dict": model_config_dict, - "sde_config_dict": sde_config_dict, - "train_config_dict": train_config_dict, - "random_seed": random_seed, - "num_training_simulations": num_training_simulations, - "train_test_fraction": train_test_fraction, - "attention_mask_type": attention_mask_type, - "has_simulator": self.has_simulator, - "learning_type": learning_type, - }, - ) + self.posteriors = trained_model + self.stats = stats - # Evaluate model. - test_x = validation_data["x"] - theta_val = validation_data["theta"] + self.save_model_to_pkl( + posteriors=trained_model, + out_dir=out_dir, + name_append=name_append, + save_method=save_method, + stats=stats, + extras={ + "model_config_dict": model_config_dict_overrides or {}, + "sde_config_dict": sde_config_dict_overrides or {}, + "train_config_dict": train_config_dict_overrides or {}, + "random_seed": random_seed, + "num_training_simulations": num_training_simulations, + "train_test_fraction": train_test_fraction, + "has_simulator": self.has_simulator, + "learning_type": "online" if self.has_simulator else "offline", + }, + ) if set_self: - self._X_test = np.array(test_x) - self._y_test = np.array(theta_val) + self._X_test = np.asarray(x_test) + self._y_test = np.asarray(theta_test) if evaluate_model: if verbose: logger.info(f"Evaluating model at {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}") self.plot_diagnostics( - task=task, - X_test=test_x, - y_test=theta_val, - posteriors=trained_score_model, - num_samples=1000, - num_evaluations=25, + X_test=x_test, + y_test=theta_test, + posteriors=trained_model, + num_samples=num_posterior_draws_per_sample, rng_seed=random_seed, - plots_dir=f"{code_path}/models/{self.name}/plots/{name_append}/", - metric_path=f"{code_path}/models/{self.name}/{self.name}_{name_append}_metrics.json", + plots_dir=os.path.join(out_dir, "plots", name_append.lstrip("_")), + metric_path=os.path.join(out_dir, f"{self.name}{name_append}_metrics.json"), ) + return trained_model, stats + def load_model_from_pkl( self, model_dir: str, model_name: str = "simformer", set_self: bool = True ): - """Load a Simformer model from a pickle file. + """Load a Simformer model and its sidecar parameter file. Parameters: - ----------- - model_dir : str - Directory where the model file is located. - model_name : str, optional - Name of the model file to load. Default is "simformer". - set_self : bool, optional - If True, sets the loaded model and task to the instance attributes. - Default is True. - """ - from .simformer import load_full_model + model_dir: Directory containing the saved model files. + model_name: Name of the model file (with or without the ``_posterior`` + suffix). Default is "simformer". + set_self: If True, sets the loaded model and metadata on the instance. + Returns: + tuple: The loaded :class:`synference.simformer.SimformerModel` and the + metadata dictionary from the sidecar params file. + """ if not model_name.endswith("_posterior"): model_name = f"{model_name}_posterior" - model, meta, task = load_full_model( - model_dir, - model_name, - ) + model_path = os.path.join(model_dir, f"{model_name}.pkl") + model = SimformerModel.load(model_path, map_location=self.device) + + meta_path = os.path.join(model_dir, f"{model_name.replace('_posterior', '_params')}.pkl") + meta = {} + if os.path.exists(meta_path): + meta = load(meta_path) + else: + logger.warning(f"No params file found at {meta_path}.") if set_self: - self.simformer_task = task self.posteriors = model - for item in meta: val = meta[item] if isinstance(val, list): @@ -8397,122 +7981,72 @@ def load_model_from_pkl( setattr(self, item, val) - model.has_features = True + self.has_features = True return model, meta def save_model_to_pkl( self, - task=None, posteriors=None, - name_append="", + name_append: str = "", out_dir: str = f"{code_path}/models/name/", - save_method="joblib", - **extras, + save_method: str = "joblib", + stats=None, + extras: Optional[dict] = None, ): - """Save the Simformer model to a pickle file. + """Save the Simformer model and the fitter state to disk. + + The model itself is written with :meth:`SimformerModel.save` (a plain torch + payload); ``save_method`` only affects the sidecar params file written by + :meth:`SBI_Fitter.save_state`. Parameters: - ************* - task : object, optional - Task object containing the model and data. - If None, uses the simformer_task attribute of the object. - posteriors : object, optional - Posteriors to save. If None, uses the posteriors attribute of the object. - name_append : str, optional - String to append to the model name in the output file. - Default is an empty string. - out_dir : str, optional - Directory to save the model files. - Default is f"{code_path}/models/name/". - save_method : str, optional - Method to use for saving the model. - Options are 'torch', 'joblib', 'pickle', and 'hdf5'. - Default is 'torch'. - extras : dict, optional - Additional parameters to save with the model. - These will be added to the saved dictionary. + posteriors: The :class:`synference.simformer.SimformerModel` to save. + If None, uses ``self.posteriors``. + name_append: String appended to the model name in the output files. + out_dir: Output directory ("name" is replaced by ``self.name``). + save_method: Serialization method for the sidecar params file. + stats: Training stats dictionary (loss traces are summarized). + extras: Additional entries stored in the params file. """ - if task is None: - task = self.simformer_task if posteriors is None: posteriors = self.posteriors - - if task is None or posteriors is None: - raise ValueError("Task and posteriors must be provided.") + if posteriors is None: + raise ValueError("A trained model must be provided.") out_dir = out_dir.replace("name", self.name) - - if not os.path.exists(out_dir): - os.makedirs(out_dir) + os.makedirs(out_dir, exist_ok=True) if len(name_append) > 0 and name_append[0] != "_": name_append = f"_{name_append}" file_name = os.path.join(out_dir, f"{self.name}{name_append}_posterior.pkl") - - if save_method == "torch": - from torch import save - - # Save trained model using PyTorch - save(posteriors, file_name) - elif save_method == "joblib": - from joblib import dump - - # Save trained model using joblib - dump(posteriors, file_name, compress=3) - elif save_method == "pickle": - import pickle - - # Save trained model using pickle - with open(file_name, "wb") as f: - pickle.dump(posteriors, f) - elif save_method == "hdf5": - raise NotImplementedError("HDF5 saving is not implemented yet.") - elif save_method == "dill": - import dill - - with open(file_name, "wb") as f: - dill.dump(posteriors, f) - else: - raise ValueError( - f"Invalid save_method: {save_method}. " - "Choose from 'torch', 'joblib', 'pickle', 'dill', or 'hdf5'." - ) - - from scoresbibm.methods.score_transformer import get_z_score_fn - - # Why is this here? - if ( - posteriors.z_score_params is not None - and posteriors.z_score_params["z_score_fn"] is None - ): - zscore, un_zscore = get_z_score_fn( - posteriors.z_score_params["mean_per_node_id"], - posteriors.z_score_params["std_per_node_id"], - ) - posteriors.z_score_params["z_score_fn"] = zscore - posteriors.z_score_params["un_z_score_fn"] = un_zscore + posteriors.save(file_name) save_dict = { - "_x_dim": task.get_x_dim(), - "_theta_dim": task.get_theta_dim(), - "prior_dict": task.prior_dist.prior_ranges, - "param_names_ordered": task.param_names_ordered, - "backend": task.backend, + "theta_dim": posteriors.theta_dim, + "x_dim": posteriors.x_dim, + "prior_dict": posteriors.meta.get("prior_ranges", {}), + "param_names_ordered": posteriors.meta.get("param_names", []), + "backend": "torch", } + if stats is not None: + summary = { + key: value + for key, value in stats.items() + if isinstance(value, (int, float, bool, str)) or value is None + } + save_dict["stats"] = [summary] + if extras: + save_dict.update(extras) - save_dict.update(extras) - - # Move the posteriors to CPU and convert to numpy if needed. Saves problems recreating - # the arrays later save_dict = make_serializable(save_dict, allowed_types=[np.ndarray]) self.save_state( out_dir=out_dir, name_append=name_append, save_method=save_method, - has_grid=~self.has_simulator, + has_grid=not self.has_simulator, **save_dict, ) @@ -8520,78 +8054,55 @@ def plot_diagnostics( self, X_test=None, y_test=None, - num_samples=1000, - num_evaluations=25, - task=None, + num_samples: int = 1000, posteriors=None, rng_seed: int = 42, plots_dir: str = f"{code_path}/models/name/plots/", metric_path: str = f"{code_path}/models/name/name_metrics.json", overwrite: bool = False, ): - """Plot diagnostics for the Simformer model. + """Evaluate the model on test data and save diagnostic plots and metrics. Args: - X_test (np.ndarray, optional): Test data to evaluate the model. - If None, uses the `X_test` attribute. Defaults to None. - y_test (np.ndarray, optional): True values for the test data. - If None, uses the `y_test` attribute. Defaults to None. - num_samples (int, optional): Number of samples to draw from the - posterior. Defaults to 1000. - num_evaluations (int, optional): Number of evaluations to perform - for coverage. Defaults to 25. - task (object, optional): Task object containing the model and data. - If None, uses the `simformer_task` attribute. Defaults to None. - posteriors (object, optional): Posteriors to use for sampling. - If None, uses the posteriors stored in the object. - Defaults to None. - rng_seed (int, optional): Random seed for reproducibility. - Defaults to 42. - plots_dir (str, optional): Directory to save the plots. - Defaults to f"{code_path}/models/{name}/plots/". - metric_path (str, optional): Path to save the metrics JSON file. - Defaults to f"{code_path}/models/{name}/{name}_metrics.json". - overwrite (bool, optional): If True, overwrites existing plots - and metrics. Defaults to False. + X_test (np.ndarray, optional): Test features. Defaults to ``self._X_test``. + y_test (np.ndarray, optional): True test parameters. Defaults to + ``self._y_test``. + num_samples (int, optional): Posterior draws per observation. + Defaults to 1000. + posteriors (SimformerModel, optional): Model to evaluate. Defaults to + ``self.posteriors``. + rng_seed (int, optional): Random seed. Defaults to 42. + plots_dir (str, optional): Directory for the plots. + metric_path (str, optional): Path for the metrics JSON file. + overwrite (bool, optional): Unused, kept for API parity. Defaults to + False. Returns: - None: - This function saves plots and metrics to disk and does not - return anything. + None: Plots and metrics are written to disk. """ - if task is None: - task = self.simformer_task - if posteriors is None: posteriors = self.posteriors + if X_test is None or y_test is None: + X_test, y_test = self._X_test, self._y_test - """ - eval_inference_task( - task=task, - model=posteriors, - metric_fn=c2st, # Use the c2st metric function - metric_params={"condition_mask_fn": "posterior"}, - rng=master_rng_key, - num_samples=num_samples, - num_evaluations=num_evaluations, - ) - """ samples = self.sample_posterior( X_test=X_test, num_samples=num_samples, posteriors=posteriors, rng_seed=rng_seed, ) + if samples.ndim == 2: + samples = samples[None, ...] metrics = self.evaluate_model( posteriors=posteriors, X_test=X_test, y_test=y_test, samples=samples, + num_samples=num_samples, ) metrics_path = metric_path.replace("name", self.name) - if not os.path.exists(os.path.dirname(metrics_path)): - os.makedirs(os.path.dirname(metrics_path)) + os.makedirs(os.path.dirname(metrics_path), exist_ok=True) try: with open(metrics_path, "w") as f: @@ -8600,10 +8111,8 @@ def plot_diagnostics( logger.error(f"Error saving metrics to {metrics_path}: {e}") self.plot_sample_accuracy( - num_samples=num_samples, X_test=X_test, y_test=y_test, - task=task, posteriors=posteriors, rng_seed=rng_seed, plots_dir=plots_dir, @@ -8611,91 +8120,61 @@ def plot_diagnostics( ) self.plot_coverage( - num_samples=num_samples, - num_evaluations=num_evaluations, - task=task, - posteriors=posteriors, - rng_seed=rng_seed, + y_test=y_test, + samples=samples, plots_dir=plots_dir, - overwrite=overwrite, ) def plot_sample_accuracy( self, - num_samples=1000, + num_samples: int = 1000, X_test=None, y_test=None, - task=None, posteriors=None, rng_seed: int = 42, plots_dir: str = f"{code_path}/models/name/plots/", samples=None, ): - """Plot the accuracy of the sampled posterior distribution. + """Plot recovered posterior quantiles against the true parameter values. Parameters: - ------------ - num_samples : int, optional - Number of samples to draw from the posterior. Default is 1000. - X_test : np.ndarray, optional - Test data to sample from the posterior. If None, uses the - X_test attribute of the object. - y_test : np.ndarray, optional - True values for the test data. If None, uses the y_test - attribute of the object. - task : object, optional - Task object containing the model and data. If None, uses the - simformer_task attribute of the object. - posteriors : object, optional - Posteriors to use for sampling. If None, uses the posteriors - stored in the object. - rng_seed : int, optional - Random seed for reproducibility. Default is 42. - plots_dir : str, optional - Directory to save the plots. Default is - f"{code_path}/models/{name}/plots/". - samples : np.ndarray, optional - Pre-computed samples from the posterior. If None, samples will be - drawn from the posterior using the sample_posterior method. + num_samples: Posterior draws per observation if ``samples`` is None. + X_test: Test features. Defaults to ``self._X_test``. + y_test: True test parameters. Defaults to ``self._y_test``. + posteriors: Model to sample from. Defaults to ``self.posteriors``. + rng_seed: Random seed. + plots_dir: Directory for the plot ("name" replaced by ``self.name``). + samples: Precomputed samples of shape ``(n_obs, num_samples, n_theta)``. """ - if task is None: - task = self.simformer_task - if posteriors is None: posteriors = self.posteriors - - posterior_condition_mask = jnp.array( - [0] * task.get_theta_dim() + [1] * task.get_x_dim(), dtype=jnp.bool_ - ) + if X_test is None or y_test is None: + X_test, y_test = self._X_test, self._y_test if samples is None: - y_test_recovered = self.sample_posterior( + samples = self.sample_posterior( X_test=X_test, - num_samples=1000, + num_samples=num_samples, posteriors=posteriors, rng_seed=rng_seed, - attention_mask=posterior_condition_mask, ) - else: - y_test_recovered = samples - - # Get 16, 50 and 84 percentiles for each parameter for each test sample + if samples.ndim == 2: + samples = samples[None, ...] + y_test = np.asarray(y_test) + param_names = list(self.simple_fitted_parameter_names) fig, ax = plt.subplots( nrows=1, - ncols=len(task.param_names_ordered), - figsize=(len(task.param_names_ordered) * 3, 3), + ncols=len(param_names), + figsize=(len(param_names) * 3, 3), + squeeze=False, ) + ax = ax[0] - for i, param in enumerate(task.param_names_ordered): - # Get the 16th, 50th and 84th percentiles for the parameter - p16, p50, p84 = np.percentile( - y_test_recovered[:, :, i], - [16, 50, 84], - axis=1, - ).squeeze() + for i, param in enumerate(param_names): + p16, p50, p84 = np.percentile(samples[:, :, i], [16, 50, 84], axis=1) ax[i].errorbar( - y_test[:, i], # True values + y_test[:, i], p50, yerr=[p50 - p16, p84 - p50], fmt="o", @@ -8704,285 +8183,304 @@ def plot_sample_accuracy( elinewidth=0.5, alpha=0.75, ) - # Add a 1:1 line ax[i].plot( [y_test[:, i].min(), y_test[:, i].max()], [y_test[:, i].min(), y_test[:, i].max()], "k--", ) - ax[i].set_title(param) ax[i].set_xlabel("True") if i == 0: ax[i].set_ylabel("Predicted") plt.tight_layout() - plots_dir = plots_dir.replace("name", self.name) - - if not os.path.exists(plots_dir): - os.makedirs(plots_dir) - # TODO: Rename to match convention - plt.savefig(os.path.join(plots_dir, "plot_sample_predictions.jpg")) + os.makedirs(plots_dir, exist_ok=True) + fig.savefig(os.path.join(plots_dir, "plot_sample_predictions.jpg")) + plt.close(fig) def plot_coverage( self, - num_samples=1000, - num_evaluations=25, - task=None, + y_test=None, + samples=None, + X_test=None, + num_samples: int = 1000, posteriors=None, rng_seed: int = 42, plots_dir: str = f"{code_path}/models/name/plots/", ): - """Plot the coverage of the posterior distribution. + """Plot TARP coverage and SBC rank histograms for the posterior. Parameters: - ------------ - num_samples : int, optional - Number of samples to draw from the posterior. Default is 1000. - num_evaluations : int, optional - Number of evaluations to perform for coverage. Default is 25. - task : object, optional - Task object containing the model and data. If None, uses the - simformer_task attribute of the object. - posteriors : object, optional - Posteriors to use for sampling. If None, uses the posteriors - stored in the object. - rng_seed : int, optional - Random seed for reproducibility. Default is 42. - plots_dir : str, optional - Directory to save the plots. Default is - f"{code_path}/models/{name}/plots/". + y_test: True test parameters. Defaults to ``self._y_test``. + samples: Precomputed samples of shape ``(n_obs, num_samples, n_theta)``. + Drawn from the posterior if None. + X_test: Test features (used only if ``samples`` is None). + num_samples: Posterior draws per observation if sampling here. + posteriors: Model to sample from. Defaults to ``self.posteriors``. + rng_seed: Random seed. + plots_dir: Directory for the plots ("name" replaced by ``self.name``). """ - master_rng_key = jax.random.PRNGKey(rng_seed) - - from scoresbibm.evaluation.eval_task import eval_coverage - - metric_values, eval_time = eval_coverage( - task=task, - model=posteriors, - metric_params={ - "num_samples": num_samples, - "num_evaluations": num_evaluations, - "condition_mask_fn": "posterior", # posterior, joint, likelihood, - # random or structured random - "num_bins": 20, # Number of bins for histogram - "sample_kwargs": {}, - "log_prob_kwargs": {}, - "batch_size": 64, # Batch size for sampling - }, - rng=master_rng_key, - ) - - plt.plot(metric_values[0], metric_values[1], marker="o", label="Coverage") - plt.plot([0, 1], [0, 1], "k--", label="Ideal Coverage") - plt.xlabel("Predicted Percentile") - plt.ylabel("Empirical Percentile") - plt.legend() - - plt.title(f"Coverage Plot (num_samples={num_samples}, num_evaluations={num_evaluations})") + if y_test is None: + y_test = self._y_test + if samples is None: + if X_test is None: + X_test = self._X_test + samples = self.sample_posterior( + X_test=X_test, + num_samples=num_samples, + posteriors=posteriors, + rng_seed=rng_seed, + ) + if samples.ndim == 2: + samples = samples[None, ...] + y_test = np.asarray(y_test) plots_dir = plots_dir.replace("name", self.name) + os.makedirs(plots_dir, exist_ok=True) + + # TARP expects (num_samples, num_sims, num_dims). + ecp, alpha = tarp.get_tarp_coverage(np.transpose(samples, (1, 0, 2)), y_test, norm=True) + fig, ax = plt.subplots(figsize=(4, 4)) + ax.plot(alpha, ecp, marker="o", markersize=2, label="TARP coverage") + ax.plot([0, 1], [0, 1], "k--", label="Ideal") + ax.set_xlabel("Credibility level") + ax.set_ylabel("Expected coverage") + ax.legend() + fig.tight_layout() + fig.savefig(os.path.join(plots_dir, f"coverage_plot_{self._timestamp}.png")) + plt.close(fig) - if not os.path.exists(plots_dir): - os.makedirs(plots_dir) - - plt.savefig(os.path.join(plots_dir, f"coverage_plot_{self._timestamp}.png")) + # SBC rank histograms: rank of the truth among the posterior samples. + ranks = (samples < y_test[:, None, :]).sum(axis=1) + param_names = list(self.simple_fitted_parameter_names) + fig, ax = plt.subplots( + nrows=1, + ncols=len(param_names), + figsize=(len(param_names) * 3, 3), + squeeze=False, + ) + ax = ax[0] + for i, param in enumerate(param_names): + ax[i].hist(ranks[:, i], bins=20, density=True, alpha=0.8) + ax[i].axhline(1.0 / samples.shape[1] * 20 / 20, color="k", ls="--", lw=1) + ax[i].set_title(param) + ax[i].set_xlabel("Rank") + fig.tight_layout() + fig.savefig(os.path.join(plots_dir, f"sbc_ranks_{self._timestamp}.png")) + plt.close(fig) def plot_posterior(self): - """Plot the posterior distribution.""" + """Plot the posterior distribution (not implemented).""" pass - def log_prob(self, X_test, condition_mask="full", posteriors=None, theta=None): - """Compute the log probability of the data given the model. + def log_prob( + self, + X_test, + theta=None, + posteriors=None, + condition_mask: Union[str, np.ndarray] = "full", + num_steps: int = 100, + num_samples: int = 100, + **kwargs, + ): + """Compute log probabilities of parameters given observations. - Parameters - ---------- - X_test : np.ndarray - Observed data of shape (n_observations, n_features) or (n_features,) - condition_mask : np.ndarray or str - Mask indicating which parts of the data are observed. - If 'full', assumes all features are observed. - posteriors : object, optional - Posteriors to use for computing the log probability. If None, - will use the posteriors stored in the object. - theta : np.ndarray, optional - Parameter samples of shape (n_samples, n_params). If None, - will sample from posterior. + Parameters: + X_test: Observed features of shape ``(n_obs, n_conditioned)`` or + ``(n_conditioned,)``. + theta: Parameter values. If it has one row per observation, a paired + log probability is returned per observation. If ``X_test`` is a + single observation, ``theta`` may hold many rows. If None, + ``num_samples`` posterior samples are drawn and scored per + observation. + posteriors: Model to use. Defaults to ``self.posteriors``. + condition_mask: Boolean node mask (True = observed) or "full" for the + standard posterior mask. + num_steps: Number of probability-flow ODE steps. + num_samples: Posterior draws per observation when ``theta`` is None. + **kwargs: Forwarded to :meth:`SimformerModel.log_prob`. Returns: - ------- - np.ndarray - Log probabilities of shape (n_observations, n_samples) where each - element [i,j] is the log probability of observation i under - posterior sample j. + np.ndarray: Log probabilities — shape ``(n_obs,)`` for paired input, + ``(n_theta_rows,)`` for a single observation, or + ``(n_obs, num_samples)`` when ``theta`` is None. """ if posteriors is None: posteriors = self.posteriors - num_theta = len(self.fitted_parameter_names) - num_x = len(self.feature_names) - - if condition_mask == "full": - condition_mask = jnp.array([0] * num_theta + [1] * num_x, dtype=jnp.bool_) - else: - condition_mask = jnp.array(condition_mask, dtype=jnp.bool_) - - X_test = np.atleast_2d(X_test) - n_observations = X_test.shape[0] + mask = self._resolve_condition_mask(condition_mask) + X_test = np.atleast_2d(np.asarray(X_test, dtype=np.float32)) + n_obs = X_test.shape[0] - # Get posterior samples if not provided if theta is None: - theta = self.sample_posterior( - X_test=X_test, + samples = self.sample_posterior( + X_test, + num_samples=num_samples, posteriors=posteriors, - condition_mask=condition_mask, + condition_mask=mask, ) + if samples.ndim == 2: + samples = samples[None, ...] + log_probs = np.stack( + [ + posteriors.log_prob(samples[i], X_test[i], mask, num_steps=num_steps, **kwargs) + for i in range(n_obs) + ] + ) + return log_probs[0] if n_obs == 1 else log_probs - # Ensure theta is 2D: (n_samples, n_params) - theta = np.atleast_2d(theta) - n_samples = theta.shape[0] - - # Initialize result array - log_probs = np.zeros((n_observations, n_samples)) - - # Compute log probability for each observation and each posterior sample - for i, x_obs in enumerate(X_test): - x_o = jnp.array(x_obs, dtype=jnp.float32) - - for j in range(n_samples): - theta_sample = jnp.array(theta[j], dtype=jnp.float32) + theta = np.atleast_2d(np.asarray(theta, dtype=np.float32)) + if theta.shape[0] == n_obs: + return np.atleast_1d( + posteriors.log_prob(theta, X_test, mask, num_steps=num_steps, **kwargs) + ) + if n_obs == 1: + return np.atleast_1d( + posteriors.log_prob(theta, X_test[0], mask, num_steps=num_steps, **kwargs) + ) + raise ValueError(f"Cannot pair theta of shape {theta.shape} with {n_obs} observations.") - log_prob = posteriors.log_prob( - theta=theta_sample, x_o=x_o, condition_mask=condition_mask - ) - log_probs[i, j] = float(log_prob) - - # Return appropriate shape based on input - if n_observations == 1 and n_samples == 1: - return log_probs[0, 0] # Single scalar - elif n_observations == 1: - return log_probs[0, :] # 1D array of samples for single observation - elif n_samples == 1: - return log_probs[:, 0] # 1D array of observations for single sample - else: - return log_probs # 2D array + def _resolve_condition_mask(self, condition_mask) -> np.ndarray: + """Resolve the "full" sentinel to the standard posterior condition mask.""" + num_theta = len(self.fitted_parameter_names) + num_x = len(self.feature_names) + if isinstance(condition_mask, str): + if condition_mask != "full": + raise ValueError(f"Unknown condition mask '{condition_mask}'. Use 'full'.") + return np.array([False] * num_theta + [True] * num_x) + mask = np.asarray(condition_mask, dtype=bool) + if mask.shape != (num_theta + num_x,): + raise ValueError( + f"condition_mask must have shape ({num_theta + num_x},), got {mask.shape}." + ) + return mask def sample_posterior( self, X_test, num_samples: int = 1000, - posteriors: object = None, + posteriors=None, rng_seed: int = 42, - attention_mask: Union[str, np.ndarray] = "full", + condition_mask: Union[str, np.ndarray] = "full", batch_size: int = 100, + num_steps: Optional[int] = None, + sample_method: str = "sde", **kwargs, ): - """Sample from the posterior distribution. - - Parameters - ---------- - - X_test : np.ndarray - Test data to sample from the posterior. - - num_samples : int, optional - Number of samples to draw from the posterior. Default is 1000. - posteriors : object, optional - Posteriors to use for sampling. If None, will use the posteriors - stored in the object. - attention_mask : Union[str, np.ndarray], optional - Attention mask to use for sampling. Can be 'full' or a numpy array. - Default is 'full'. Full means you have full observations for all bands. - If you have missing bands, you can provide a numpy array with - the shape + """Sample from the posterior (or any conditional) for a set of observations. - TODO: Make this work for multidimensional X_test. + Parameters: + X_test: Observations of shape ``(n_obs, n_conditioned)`` or + ``(n_conditioned,)``. Columns must match the conditioned nodes of + ``condition_mask`` in node order. + num_samples: Samples per observation. Default is 1000. + posteriors: Model to sample from. Defaults to ``self.posteriors``. + rng_seed: Random seed for the sampler. + condition_mask: Boolean node mask (True = observed) or "full" for the + standard posterior mask over ``[theta..., x...]``. + batch_size: Observations integrated simultaneously. + num_steps: Integration steps (default: the model's sampling default). + sample_method: "sde" (stochastic) or "ode" (deterministic). + **kwargs: Ignored extra keyword arguments from generic callers + (a warning is logged). + Returns: + np.ndarray: Samples of shape ``(n_obs, num_samples, n_latent)``, squeezed + to ``(num_samples, n_latent)`` for a single observation. """ if posteriors is None: posteriors = self.posteriors - master_rng_key = jax.random.PRNGKey(rng_seed) - - num_theta = len(self.fitted_parameter_names) - num_x = len(self.feature_names) - - X_test = np.atleast_2d(X_test) - - assert X_test.shape[1] == num_x or attention_mask == "full", ( - "Must provide all features or a manual attention mask. " + if kwargs: + logger.warning(f"sample_posterior ignoring unsupported kwargs: {sorted(kwargs)}") + + mask = self._resolve_condition_mask(condition_mask) + X_test = np.atleast_2d(np.asarray(X_test, dtype=np.float32)) + assert X_test.shape[1] == mask.sum(), ( + f"X_test has {X_test.shape[1]} columns but the condition mask expects " + f"{int(mask.sum())} observed values." ) - if attention_mask == "full": - mask = jnp.array([0] * num_theta + [1] * num_x, dtype=jnp.bool_) - else: - mask = attention_mask.astype(np.bool_) - - all_samples = [] + generator = torch.Generator(device=posteriors.device) + generator.manual_seed(int(rng_seed)) - for x in tqdm(X_test, desc="Sampling from posterior"): - samples = posteriors.sample_batched( - num_samples=num_samples, - x_o=x, - condition_mask=mask, - rng=master_rng_key, - ) - - samples = np.array(samples[0], dtype=np.float32) - all_samples.append(samples) - """ Batched sampling is slow. - nbatches = int(np.ceil(X_test.shape[0] / batch_size)) - for batch_idx in trange(nbatches, desc="Sampling from posterior"): - start_idx = batch_idx * batch_size - end_idx = min((batch_idx + 1) * batch_size, X_test.shape[0]) - x_batch = jnp.array(X_test[start_idx:end_idx], dtype=jnp.float32) - - # Sample from the posterior for the batch - samples_batch = posteriors.sample_batched( - num_samples=num_samples, - x_o=x_batch, - condition_mask=mask, - rng=master_rng_key, - ) - - # Convert to numpy and append to the list - all_samples.extend(np.array(samples_batch, dtype=np.float32)) - """ - - all_samples = np.array(all_samples, dtype=np.float32) + all_samples = posteriors.sample_batched( + num_samples, + X_test, + mask, + batch_size=batch_size, + num_steps=num_steps, + method=sample_method, + generator=generator, + ).astype(np.float32) - # if X_test is 1-dimensional, flatten the output if X_test.shape[0] == 1: - all_samples = all_samples[0] - + return all_samples[0] return all_samples - """theta_samples = samples[:, :theta_dim] - if attention_mask != 'full': - phot_samples = samples[:, theta_dim:] - return theta_samples, phot_samples - else: - return theta_samples""" + def sample_posterior_intervals( + self, + X_test, + constraint_mask, + a=None, + b=None, + num_samples: int = 1000, + posteriors=None, + rng_seed: int = 42, + condition_mask: Union[str, np.ndarray] = "full", + num_steps: Optional[int] = None, + scale_bias: float = 1e-2, + ): + """Sample the posterior with interval constraints on some nodes via guidance. + + Parameters: + X_test: A single observation of shape ``(n_conditioned,)`` (values for + the effective condition mask, i.e. excluding constrained nodes). + constraint_mask: Boolean node mask of interval-constrained nodes. + a: Lower bounds for the constrained nodes (original units), or None. + b: Upper bounds for the constrained nodes (original units), or None. + num_samples: Number of samples. + posteriors: Model to sample from. Defaults to ``self.posteriors``. + rng_seed: Random seed. + condition_mask: Boolean node mask or "full". + num_steps: Integration steps. + scale_bias: Stability bias in the guidance scale (see + :meth:`SimformerModel.sample_intervals`). + + Returns: + np.ndarray: Samples of shape ``(num_samples, n_latent)`` where the + latents include the constrained nodes. + """ + if posteriors is None: + posteriors = self.posteriors + mask = self._resolve_condition_mask(condition_mask) + generator = torch.Generator(device=posteriors.device) + generator.manual_seed(int(rng_seed)) + return posteriors.sample_intervals( + num_samples, + x_o=np.asarray(X_test, dtype=np.float32).reshape(-1), + condition_mask=mask, + constraint_mask=np.asarray(constraint_mask, dtype=bool), + a=a, + b=b, + scale_bias=scale_bias, + num_steps=num_steps, + generator=generator, + ) def create_priors( self, override_prior_ranges: dict = {}, verbose: bool = True, set_self: bool = False ): - """Create priors for the Simformer model. + """Create a uniform box prior over the fitted parameters. Parameters: - ----------- - override_prior_ranges : dict, optional - Dictionary to override the default prior ranges. - Default is an empty dictionary. - verbose : bool, optional - If True, prints information about the prior creation. - Default is True. - set_self : bool, optional - If True, sets the created prior to the instance attribute. + override_prior_ranges: Overrides for the default prior ranges. + verbose: If True, prints information about the prior creation. + set_self: If True, stores the prior on ``self._prior``. + Returns: + UniformBoxPrior: The constructed prior. """ - from .simformer import GalaxyPrior - priors_sbi = ( super() .create_priors( @@ -8998,7 +8496,7 @@ def create_priors( high = float(priors_sbi.high[i]) prior_ranges[name] = (low, high) - prior = GalaxyPrior(prior_ranges, self.fitted_parameter_names) + prior = UniformBoxPrior(prior_ranges, self.fitted_parameter_names) if set_self: self._prior = prior @@ -9006,7 +8504,7 @@ def create_priors( return prior def optimize_sbi(self): - """Optimize the SBI model.""" + """Optimize the SBI model (not implemented for the Simformer).""" raise NotImplementedError("Simformer_Fitter does not implement optimize_sbi method. ") def fit_catalogue( @@ -9025,20 +8523,18 @@ def fit_catalogue( check_out_of_distribution: bool = True, simulator: Optional[GalaxySimulator] = None, rng_seed: int = 42, - attention_mask: Union[str, np.ndarray] = "full", + condition_mask: Union[str, np.ndarray] = "full", batch_size: int = 100, missing_data_mcmc: bool = False, - log_times=False, + log_times: bool = False, ): - """Wrapper for fit_catalogue in parent. - - To Do: Better attention mask. + """Fit a catalogue of observations (wrapper around the base implementation). + See :meth:`SBI_Fitter.fit_catalogue` for the shared parameters. The + Simformer-specific ``condition_mask`` (True = observed node) replaces the + old ``attention_mask`` argument and is forwarded to + :meth:`sample_posterior`. """ - sample_method: str = "direct" - sample_kwargs: dict = {} - timeout_seconds_per_row: int = 5 - return super().fit_catalogue( observations=observations, columns_to_feature_names=columns_to_feature_names, @@ -9053,11 +8549,11 @@ def fit_catalogue( plot_SEDs=plot_SEDs, check_out_of_distribution=check_out_of_distribution, simulator=simulator, - sample_method=sample_method, - sample_kwargs=sample_kwargs, - timeout_seconds_per_row=timeout_seconds_per_row, + sample_method="direct", + sample_kwargs={}, + timeout_seconds_per_row=5, rng_seed=rng_seed, - attention_mask=attention_mask, + condition_mask=condition_mask, batch_size=batch_size, missing_data_mcmc=missing_data_mcmc, log_times=log_times, diff --git a/src/synference/simformer.py b/src/synference/simformer.py deleted file mode 100644 index 4df8e75..0000000 --- a/src/synference/simformer.py +++ /dev/null @@ -1,761 +0,0 @@ -"""Functions for simformer tasks and model saving/loading.""" - -import os -import pickle -from typing import ( - Callable, - Dict, - List, - Tuple, -) # For GalaxySimulator type hints - -import corner -import jax -import jax.numpy as jnp -import joblib -import numpy as np -import torch - -try: - from synthesizer.emission_models import TotalEmission - from synthesizer.emission_models.attenuation import Calzetti2000 - from synthesizer.grid import Grid - from synthesizer.instruments import FilterCollection, Instrument - from synthesizer.parametric import ( - SFH, - ZDist, - ) # Need concrete SFH, ZDist classese -except Exception: - pass - -from unyt import ( - Myr, -) - -try: - from omegaconf import OmegaConf # To create DictConfig-like objects if needed - from scoresbibm.methods.score_transformer import train_transformer_model - from scoresbibm.tasks.base_task import InferenceTask -except Exception as e: - print(e) - - class InferenceTask: - """Dummy InferenceTask class for compatibility.""" - - pass - - -# --- Custom Simformer Task --- -class GalaxyPhotometryTask(InferenceTask): - """A Simformer InferenceTask for the GalaxySimulator.""" - - _theta_dim: int - _x_dim: int - - def __init__( - self, - name: str = "galaxy_photometry_task", - backend: str = "jax", - prior_dict: Dict[str, Tuple[float, float]] = None, - param_names_ordered: List[str] = None, - run_simulator_fn: Callable = None, - num_filters: int = None, - test_X_data: np.ndarray = None, - test_theta_data: np.ndarray = None, - attention_mask_type: str = "full", # or "causal", or a custom jnp.ndarray - ): - """Initializes the GalaxyPhotometryTask. - - If a simulator function is not provided, test_X_data and test_theta_data - should be provided if you want to do any testing or validation. - - Arguments: - name: Name of the task. - backend: Backend to use, e.g., "jax", "torch", or "numpy". - prior_dict: Dictionary defining parameter ranges for the prior. - param_names_ordered: Ordered list of parameter names. - run_simulator_fn: Function that runs the galaxy simulator. - num_filters: Number of filters in the photometry simulation. - test_X_data: Optional test data for x (photometry). - test_theta_data: Optional test data for theta (parameters). - attention_mask_type: Type of attention mask to use, can be "full", "causal", - or a custom jnp.ndarray defining the mask structure. - - """ - super().__init__(name, backend) - - if run_simulator_fn is None and (test_X_data is None or test_theta_data is None): - print("Warning! No simulator function or test data provided. ") - - if prior_dict is None or param_names_ordered is None or num_filters is None: - raise ValueError( - """prior_dict, param_names_ordered, - and num_filters must be provided.""" - ) - - self.param_names_ordered = param_names_ordered - self._theta_dim = len(param_names_ordered) - self._x_dim = num_filters - self.backend = backend - self.test_X_data = test_X_data - self.test_theta_data = test_theta_data - self.attention_mask_type = attention_mask_type - - if isinstance(prior_dict, GalaxyPrior): - self.prior_dist = prior_dict - else: - self.prior_dist = GalaxyPrior( - prior_ranges=prior_dict, param_order=self.param_names_ordered - ) - - if run_simulator_fn is None: - self.fake_simulator = True - - def run_simulator_fn(X, return_type="jax"): - return None - else: - self.fake_simulator = False - - self.run_simulator_fn = run_simulator_fn - - def get_theta_dim(self) -> int: - """Returns the dimension of the theta vector.""" - return self._theta_dim - - def get_x_dim(self) -> int: - """Returns the dimension of the x vector (photometry).""" - return self._x_dim - - def get_prior(self): - """Returns the prior distribution object.""" - return self.prior_dist - - def get_simulator(self): - """Gets a batched simulator function. - - Returns: - A callable that takes a batch of thetas and returns a batch of xs. - """ - - def batched_simulator( - thetas_batch_torch: torch.Tensor, - ) -> torch.Tensor: - xs_list = [] - for i in range(thetas_batch_torch.shape[0]): - theta_sample_torch = thetas_batch_torch[i, :] - # run_simulator_fn expects numpy array if it's not a dict, - # and handles tensor conversion internally. - # It returns a tensor of shape [1, num_filters]. - x_sample_torch = self.run_simulator_fn(theta_sample_torch, return_type="tensor") - xs_list.append(x_sample_torch) - return torch.cat(xs_list, dim=0) # Shape will be (num_samples, num_filters) - - return batched_simulator - - def get_data(self, num_samples: int, **kwargs) -> Dict[str, jnp.ndarray]: - """Returns data for the task. - - Arguments: - num_samples: The number of samples to generate. - **kwargs: Additional keyword arguments for the prior sampling. - - Returns: - A dictionary with keys 'theta' and 'x', containing the sampled parameters - and simulated photometry, respectively. - """ - prior = self.get_prior() - simulator = self.get_simulator() # This is our batched_simulator - - if self.fake_simulator: - if self.test_X_data is None or self.test_theta_data is None: - raise ValueError("No simulator function provided and no test data available.") - # Use test data if available - if num_samples > self.test_X_data.shape[0]: - raise ValueError( - f"Requested {num_samples} samples" - f" but only {self.test_X_data.shape[0]} are available." - ) - thetas_torch = self.test_theta_data[:num_samples] - xs_out = self.test_X_data[:num_samples] - else: - # Sample thetas (parameters) using the prior - # GalaxyPrior.sample returns a PyTorch tensor - thetas_torch = prior.sample((num_samples,), **kwargs) - - # Simulate xs (photometry) using the parameters - # batched_simulator also returns a PyTorch tensor - xs_torch = simulator(thetas_torch) - - if self.backend == "jax": - thetas_out = jnp.array(thetas_torch.cpu().numpy()) - xs_out = jnp.array(xs_torch.cpu().numpy()) - elif self.backend == "numpy": - thetas_out = thetas_torch.cpu().numpy() - xs_out = xs_torch.cpu().numpy() - else: # "torch" or other - thetas_out = thetas_torch - xs_out = xs_torch - - return {"theta": thetas_out, "x": xs_out} - - def get_node_id(self) -> jnp.ndarray: - """Returns an array identifying the nodes (dimensions) of theta and x.""" - dim = self.get_theta_dim() + self.get_x_dim() - if self.backend == "torch": # Should align with SBIBMTask if that's a reference - return torch.arange(dim) - else: # JAX or numpy - return jnp.arange(dim) - - def get_base_mask_fn(self) -> Callable: - """Defines the base attention mask for the transformer.""" - theta_dim = self.get_theta_dim() - x_dim = self.get_x_dim() - - if self.attention_mask_type == "full": - # Block for θ attending to θ (Full self-attention) for parameters - thetas_self_mask = jnp.ones((theta_dim, theta_dim), dtype=jnp.bool_) - # Block for x attending to x (Full self-attention) for data - xs_self_mask = jnp.ones((x_dim, x_dim), dtype=jnp.bool_) - # Block for x attending to θ - data can attend to parameters - # xs_attends_theta_mask = jnp.ones((x_dim, theta_dim), dtype=jnp.bool_) - # Block for θ attending to x (NO attention) - parameters do not attend to data - # thetas_attends_xs_mask = jnp.zeros((theta_dim, x_dim), dtype=jnp.bool_) - elif self.attention_mask_type == "causal": - # Parameters only attend to themselves (or causal if ordered) - thetas_self_mask = jnp.eye(theta_dim, dtype=jnp.bool_) - # Data can attend to previous/current data points (causal within x) - # Or use jnp.ones if full self-attention within x is desired. - xs_self_mask = jnp.tril(jnp.ones((x_dim, x_dim), dtype=jnp.bool_)) - # Data can attend to all parameters - xs_attend_thetas_mask = jnp.ones((x_dim, theta_dim), dtype=jnp.bool_) - # Parameters do not attend to data - thetas_attend_xs_mask = jnp.zeros((theta_dim, x_dim), dtype=jnp.bool_) - elif isinstance(self.attention_mask_type, jnp.ndarray): - # If a custom mask is provided, use it directly - base_mask = self.attention_mask_type - if base_mask.shape != (theta_dim + x_dim, theta_dim + x_dim): - raise ValueError( - "Custom attention mask must be of shape (theta_dim + x_dim, theta_dim + x_dim)." - ) - base_mask = base_mask.astype(jnp.bool_) - - def base_mask_fn(node_ids, node_meta_data): - return base_mask[jnp.ix_(node_ids, node_ids)] - - return base_mask_fn - else: - raise ValueError( - "attention_mask_type must be 'full', 'causal', or a custom jnp.ndarray." - ) - - base_mask = jnp.block( - [ - [thetas_self_mask, thetas_attend_xs_mask], - [xs_attend_thetas_mask, xs_self_mask], - ] - ) - base_mask = base_mask.astype(jnp.bool_) - - def base_mask_fn(node_ids, node_meta_data): - # Handles potential permutation/subsetting of nodes - return base_mask[jnp.ix_(node_ids, node_ids)] - - return base_mask_fn - - -class UncertainityModelTask(InferenceTask): - """Condtional uncertainty model task for galaxy magnitudes and log-uncertainties.""" - - def __init__(self, magnitudes: np.ndarray, log_uncertainties: np.ndarray): - """Initializes the NoiseModelTask with magnitudes and log-uncertainties. - - Args: - magnitudes: Array of shape (num_examples, num_bands) from a catalog. - log_uncertainties: Array of shape (num_examples, num_bands) from a catalog. - """ - super().__init__(name="conditional_noise_model", backend="jax") - - # Ensure data is in the correct format - if magnitudes.shape != log_uncertainties.shape: - raise ValueError("Magnitudes and log_uncertainties must have the same shape.") - - # In this task, 'theta' is the magnitude, 'x' is the log_uncertainty - self._theta_data = jnp.array(magnitudes) - self._x_data = jnp.array(log_uncertainties) - - self._theta_dim = self._theta_data.shape[1] - self._x_dim = self._x_data.shape[1] - - def get_theta_dim(self) -> int: - """Returns the dimension of the theta vector (magnitude).""" - return self._theta_dim - - def get_x_dim(self) -> int: - """Returns the dimension of the x vector (log-uncertainty).""" - return self._x_dim - - def get_data(self, num_samples: int, rng=None) -> dict[str, jnp.ndarray]: - """Returns a random subset of the provided catalog data.""" - if num_samples > self._theta_data.shape[0]: - raise ValueError( - f"Requested {num_samples} samples, but only {self._theta_data.shape[0]} are available." # noqa E501 - ) - - # For simplicity, we sample with replacement. - # A more robust implementation might use a data loader. - indices = np.random.choice(self._theta_data.shape[0], size=num_samples, replace=True) - - return {"theta": self._theta_data[indices], "x": self._x_data[indices]} - - def get_base_mask_fn(self): - """Defines that log-uncertainty 'x' depends on magnitude 'theta'.""" - theta_dim = self.get_theta_dim() - x_dim = self.get_x_dim() - - # Parameters ('theta', magnitudes) only attend to themselves - thetas_self_mask = jnp.eye(theta_dim, dtype=jnp.bool_) - - # Data ('x', log-uncertainties) attend to themselves causally - # (or fully, depending on assumption about correlations between band uncertainties) - xs_self_mask = jnp.tril(jnp.ones((x_dim, x_dim), dtype=jnp.bool_)) - - # Data ('x') can attend to all parameters ('theta') - xs_attend_thetas_mask = jnp.ones((x_dim, theta_dim), dtype=jnp.bool_) - - # Parameters ('theta') do not attend to data ('x') - thetas_attend_xs_mask = jnp.zeros((theta_dim, x_dim), dtype=jnp.bool_) - - base_mask = jnp.block( - [[thetas_self_mask, thetas_attend_xs_mask], [xs_attend_thetas_mask, xs_self_mask]] - ) - base_mask = base_mask.astype(jnp.bool_) - - def base_mask_fn(node_ids, node_meta_data): - return base_mask[jnp.ix_(node_ids, node_ids)] - - return base_mask_fn - - # Methods like get_prior and get_simulator are not needed - # as get_data is implemented directly from a dataset. - - -# --- Helper Class for Prior --- -class GalaxyPrior: - """A prior distribution for galaxy parameters. - - This class uses uniform distributions for each parameter defined in prior_ranges. - It can sample from the prior and compute log probabilities. - """ - - def __init__( - self, - prior_ranges: Dict[str, Tuple[float, float]], - param_order: List[str], - ): - """Initializes the GalaxyPrior with parameter ranges and order. - - Arguments: - prior_ranges: A dictionary mapping parameter names to their (low,high) ranges. - param_order: A list of parameter names in the order they should be sampled. - """ - self.prior_ranges = prior_ranges - self.param_order = param_order - self.theta_dim = len(param_order) - - self.distributions = [] - for param_name in self.param_order: - low, high = self.prior_ranges[param_name] - self.distributions.append( - torch.distributions.Uniform(torch.tensor(float(low)), torch.tensor(float(high))) - ) - - def sample(self, sample_shape: Tuple[int], sample_lhc=False, rng=None) -> torch.Tensor: - """Generates samples from the prior. - - Arguments: - sample_shape: A tuple containing the number of samples, e.g., (num_samples,). - sample_lhc: If True, samples using Latin Hypercube sampling. - rng: Optional random number generator for reproducibility. - - Returns: - A PyTorch tensor of shape (num_samples, theta_dim). - """ - if not sample_lhc: - num_samples = sample_shape[0] - samples_per_param = [dist.sample((num_samples, 1)) for dist in self.distributions] - else: - # Use boundaries, but sample from Latin Hypercube - from scipy.stats.qmc import LatinHypercube - - sampler = LatinHypercube(d=self.theta_dim, rng=rng) - lhc_samples = sampler.random(n=sample_shape[0]) - samples_per_param = [] - for i, dist in enumerate(self.distributions): - low, high = self.prior_ranges[self.param_order[i]] - # Scale LHC samples to the range of the distribution - scaled_samples = low + (high - low) * lhc_samples[:, i : i + 1] - samples_per_param.append(torch.tensor(scaled_samples, dtype=torch.float32)) - - return torch.cat(samples_per_param, dim=1) - - def log_prob(self, theta: torch.Tensor) -> torch.Tensor: - """Calculates the log probability of theta under the prior. - - Arguments: - theta: A PyTorch tensor of shape (num_samples, theta_dim). - - Returns: - A PyTorch tensor of shape (num_samples,) containing log probabilities. - - """ - if theta.ndim == 1: - theta = theta.unsqueeze(0) # Make it (1, theta_dim) - - log_probs_per_param = [] - for i, dist in enumerate(self.distributions): - log_probs_per_param.append(dist.log_prob(theta[:, i])) - - # Sum log_probs for independent parameters - return torch.sum(torch.stack(log_probs_per_param, dim=1), dim=1) - - -def load_full_model(dir_path, model_id, simulator=None): - """Load a full model from the specified directory and model ID.""" - from joblib import load - from scoresbibm.tasks import get_task - from scoresbibm.utils.edge_masks import get_edge_mask_fn - - file_name = f"{dir_path}/{model_id}.pkl" - with open(file_name, "rb") as file: - model = joblib.load(file) - - try: - meta = load(f"{dir_path}/{model_id.replace('posterior', 'params')}.pkl") - model.__dict__.update(meta) - task_name = model.edge_mask_fn_params.get("task") - task = get_task(task_name) - task.__dict__.update(meta) - - if getattr(task, "prior_dict", None) is None: - prior_dict = meta["prior"] - else: - prior_dict = task.prior_dict - task.prior_dist = GalaxyPrior(prior_ranges=prior_dict, param_order=task.param_names_ordered) - - model.edge_mask_fn = get_edge_mask_fn(model.edge_mask_fn_params["name"], task) - - if simulator is not None: - model.simulator = simulator - else: - print("No simulator provided. Please provide a simulator to use with the model.") - except FileNotFoundError as e: - print(f"not found {e}") - meta = {} - task = None - except (EOFError, AttributeError): - print("Pickle corrupted?") - meta = {} - task = None - - if isinstance(model, tuple): - model = model[0] - - return model, meta, task - - -if __name__ == "__main__": - file_path = os.path.dirname(os.path.realpath(__file__)) - grid_folder = os.path.join(os.path.dirname(os.path.dirname(file_path)), "grids") - output_folder = os.path.join(os.path.dirname(os.path.dirname(file_path)), "models") - - # Define sfh and zdist instances or classes as used by GalaxySimulator - sfh_model_class = SFH.LogNormal - zdist_model_class = ZDist.DeltaConstant - - # Example: Define global 'device' if not already defined - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - grid_dir = os.environ["SYNTHESIZER_GRID_DIR"] - - # path for this file - - dir_path = os.path.dirname(os.path.abspath(__file__)) - out_dir = os.path.join(os.path.dirname(os.path.dirname(dir_path)), "grids/") - - grid_name = "bpass-2.2.1-bin_chabrier03-0.1,300.0_cloudy-c23.01-sps.hdf5" - - # --- This part needs functional grid, instrument etc. --- - - grid = Grid(grid_name, grid_dir=grid_dir) - filter_codes = [ - "JWST/NIRCam.F090W", - "JWST/NIRCam.F115W", - "JWST/NIRCam.F150W", - "JWST/NIRCam.F162M", - "JWST/NIRCam.F182M", - "JWST/NIRCam.F200W", - "JWST/NIRCam.F210M", - "JWST/NIRCam.F250M", - "JWST/NIRCam.F277W", - "JWST/NIRCam.F300M", - "JWST/NIRCam.F335M", - "JWST/NIRCam.F356W", - "JWST/NIRCam.F410M", - "JWST/NIRCam.F444W", - ] - filterset = FilterCollection(filter_codes) - instrument = Instrument("JWST", filters=filterset) - emission_model_instance = TotalEmission( - grid=grid, - fesc=0.0, - fesc_ly_alpha=0.1, - dust_curve=Calzetti2000(), - dust_emission_model=None, - ) - emitter_params_dict = {"stellar": ["tau_v"]} - from synference import GalaxySimulator - - galaxy_simulator_instance = GalaxySimulator( - sfh_model=sfh_model_class, # Pass the class - zdist_model=zdist_model_class, # Pass the class - grid=grid, - instrument=instrument, - emission_model=emission_model_instance, - emission_model_key="total", - emitter_params=emitter_params_dict, - param_units={"peak_age": Myr, "max_age": Myr}, # Ensure Myr is defined - normalize_method=None, - output_type="photo_fnu", - out_flux_unit="ABmag", - ) - - inputs_list = [ - "redshift", - "log_mass", - "log10metallicity", - "tau_v", - "peak_age", - "max_age", - "tau", - ] - filter_codes_list = [ - "JWST/NIRCam.F090W", - "JWST/NIRCam.F115W", - "JWST/NIRCam.F150W", - "JWST/NIRCam.F162M", - "JWST/NIRCam.F182M", - "JWST/NIRCam.F200W", - "JWST/NIRCam.F210M", - "JWST/NIRCam.F250M", - "JWST/NIRCam.F277W", - "JWST/NIRCam.F300M", - "JWST/NIRCam.F335M", - "JWST/NIRCam.F356W", - "JWST/NIRCam.F410M", - "JWST/NIRCam.F444W", - ] - - priors_ranges_dict = { - "redshift": (5.0, 10.0), - "log_mass": (7.0, 11.0), - "log10metallicity": (-3.0, -1.3), - "tau_v": (0.0, 2), - "peak_age": ( - 0.0, - 500.0, - ), # Ensure float for peak_age if used with torch.tensor(float(low)) - "max_age": (500.0, 1000.0), - "tau": (0.3, 1.5), - } - - def run_simulator_glob(params, return_type="tensor"): - """Runs the galaxy simulator with given parameters. - - Arguments: - params: A numpy array or dictionary of parameters. - return_type: "tensor" for torch tensor output, "numpy" for numpy array. - - Returns: - A torch tensor or numpy array of simulated photometry. - """ - if isinstance(params, torch.Tensor): - params = params.cpu().numpy() - if isinstance(params, dict): - pass # assumes params are correctly keyed - elif isinstance(params, (list, tuple, np.ndarray)): - params = np.squeeze(params) - params = {inputs_list[i]: params[i] for i in range(len(inputs_list))} - - phot = galaxy_simulator_instance(params) # This line requires galaxy_simulator_instance - - if return_type == "tensor": - return torch.tensor(phot[np.newaxis, :], dtype=torch.float32).to(device) - else: - return phot - - galaxy_task = GalaxyPhotometryTask( - prior_dict=priors_ranges_dict, - param_names_ordered=inputs_list, - run_simulator_fn=run_simulator_glob, - num_filters=len(filter_codes_list), - ) - - # Test data generation - print(f"Theta dim: {galaxy_task.get_theta_dim()}") - print(f"X dim: {galaxy_task.get_x_dim()}") - data_batch = galaxy_task.get_data(num_samples=3) - print("Sampled theta (JAX):", data_batch["theta"]) - print("Shape of theta:", data_batch["theta"].shape) - print("Sampled x (JAX):", data_batch["x"]) - print("Shape of x:", data_batch["x"].shape) - - # Test prior sampling directly - prior_for_test = galaxy_task.get_prior() - theta_samples_torch = prior_for_test.sample((2,)) - print("Direct prior samples (Torch):", theta_samples_torch) - print( - "Log prob of prior samples:", - prior_for_test.log_prob(theta_samples_torch), - ) - - # Test base mask function - mask_fn = galaxy_task.get_base_mask_fn() - node_ids_example = jnp.arange(galaxy_task.get_theta_dim() + galaxy_task.get_x_dim()) - applied_mask = mask_fn(node_ids=node_ids_example, node_meta_data=None) - print("Base mask applied to node_ids:", applied_mask) - - # Model Configuration (from config/method/model/score_transformer.yaml) - # - model_config_dict = { - "name": "ScoreTransformer", - "d_model": 128, - "n_heads": 4, - "n_layers": 4, - "d_feedforward": 256, - "dropout": 0.1, - "max_len": 5000, # Adjust based on theta_dim + x_dim - "tokenizer": {"name": "LinearTokenizer", "encoding_dim": 64}, - "use_output_scale_fn": True, - # Add other model-specific parameters as per the YAML - } - - # SDE Configuration (e.g., from config/method/sde/vpsde.yaml) - # - sde_config_dict = { - "name": "VPSDE", # or "VESDE" - "beta_min": 0.1, - "beta_max": 20.0, - "num_steps": 1000, - "T_min": 1e-05, - "T_max": 1.0, - # "likelihood_weighting": False, - # "schedule_name": "linear", - # Add other SDE-specific parameters - } - - # Training Configuration (from config/method/train/train_score_transformer.yaml) - # - train_config_dict = { - "learning_rate": 1e-4, # Initial learning rate for training # used - "min_learning_rate": 1e-6, # Minimum learning rate for training # used - "z_score_data": True, # Whether to z-score the data # used - "total_number_steps_scaling": 5, # Scaling factor for total number of steps - "max_number_steps": 1e8, # Maximum number of steps for training # used - "min_number_steps": 1e4, # Minimum number of steps for training # used - "training_batch_size": 64, # Batch size for training # used - "val_every": 100, # Validate every 100 steps # used - "clip_max_norm": 10.0, # Gradient clipping max norm # used - "condition_mask_fn": {"name": "joint"}, # Use the base mask function defined in the task - "edge_mask_fn": {"name": "none"}, - "validation_fraction": 0.1, # Fraction of data to use for validation # used - "val_repeat": 5, # Number of times to repeat validation # used - "stop_early_count": 5, # Number of steps to wait before stopping early # used - "rebalance_loss": False, # Whether to rebalance the loss # used - } - - method_config_dict = { - "device": str(device), # Ensure this matches device setup - "sde": sde_config_dict, - "model": model_config_dict, - "train": train_config_dict, - } - - # Convert the main method_cfg to OmegaConf DictConfig - method_cfg = OmegaConf.create(method_config_dict) - - print("Instantiating GalaxyPhotometryTask...") - galaxy_task = GalaxyPhotometryTask( - prior_dict=priors_ranges_dict, - param_names_ordered=inputs_list, - run_simulator_fn=run_simulator_glob, # function to run the simulator - num_filters=len(filter_codes_list), - backend="jax", # Or "torch" - ) - print("Task instantiated.") - - # --- 2. Generate Data --- - num_training_simulations = 5000 # Example number - print(f"Generating {num_training_simulations} training simulations...") - # .get_data() returns a dict with JAX arrays if backend is "jax" - training_data = galaxy_task.get_data(num_samples=num_training_simulations) - theta_train = training_data["theta"] - x_train = training_data["x"] - print(f"Data generated: theta shape {theta_train.shape}, x shape {x_train.shape}") - - # (Optional) Generate validation data if train_config.val_split is 0 - num_validation_simulations = 20 - validation_data = galaxy_task.get_data(num_samples=num_validation_simulations) - theta_val = validation_data["theta"] - x_val = validation_data["x"] - - # --- 3. Set RNG Seed for JAX --- - rng_seed_for_training = 0 - master_rng_key = jax.random.PRNGKey(rng_seed_for_training) - - # --- 4. Train the Model --- - - print("Starting training...") - trained_score_model = train_transformer_model( - task=galaxy_task, - data=training_data, # Expects dict {"theta": ..., "x": ...} with JAX arrays - method_cfg=method_cfg, # The OmegaConf object created above - rng=master_rng_key, - ) - print( - "Training finished. Model returned by train_transformer_model:", - type(trained_score_model), - ) - plot_corner = True - # Take test observation - theta_dim = galaxy_task.get_theta_dim() - x_dim = galaxy_task.get_x_dim() - # Mask for posterior: theta is unknown (0), x is known (1) - posterior_condition_mask = jnp.array([0] * theta_dim + [1] * x_dim, dtype=jnp.bool_) - for i, xobs in enumerate(x_val): - x_val = jnp.array([xobs], dtype=jnp.float32) - samples = trained_score_model.sample_batched( - num_samples=1000, - x_o=x_val, - rng=master_rng_key, - condition_mask=posterior_condition_mask, - ) - if plot_corner: - import corner - import matplotlib.pyplot as plt - - truth = jnp.array(theta_val[i], dtype=jnp.float32) - corner.corner( - samples, - labels=galaxy_task.param_names_ordered, - show_titles=True, - truths=truth, - quantiles=[0.16, 0.5, 0.84], - title_kwargs={"fontsize": 12}, - ) - plt.savefig(f"{output_folder}/simformer/plots/corner_plot_{i}.png") - - import pickle - - with open("trained_galaxy_score_model_params.pkl", "wb") as f: - pickle.dump(trained_score_model.score_model_params, f) - print("Model parameters saved (example).") diff --git a/src/synference/simformer/__init__.py b/src/synference/simformer/__init__.py new file mode 100644 index 0000000..750a163 --- /dev/null +++ b/src/synference/simformer/__init__.py @@ -0,0 +1,69 @@ +"""Native PyTorch implementation of the Simformer (Gloeckler et al. 2024). + +A transformer + score-diffusion all-in-one SBI model: a single network trained on the +joint ``[theta, x]`` with per-example condition masks provides posterior, likelihood, +and arbitrary conditionals, with optional interval-constrained sampling via guidance. +""" + +from .masks import ( + build_base_mask, + get_condition_mask_fn, + joint_condition_mask, + likelihood_condition_mask, + posterior_condition_mask, + structured_random_condition_mask, +) +from .model import SimformerModel, UniformBoxPrior, posterior_mask +from .nn import GaussianFourierEmbedding, ScalarTokenizer, ScoreTransformer, TransformerBlock +from .sampling import ( + euler_maruyama_reverse, + guided_euler_maruyama, + heun_probability_flow, + interval_constraint_score, + probability_flow_log_prob, +) +from .sde import SDE_CLASSES, VESDE, VPSDE, BaseSDE, build_sde +from .train import ( + DEFAULT_MODEL_CONFIG, + DEFAULT_SDE_CONFIG, + DEFAULT_TRAIN_CONFIG, + adaptive_grad_clip_, + denoising_score_matching_loss, + mean_std_per_node, + merge_config, + train_simformer, +) + +__all__ = [ + "SimformerModel", + "UniformBoxPrior", + "posterior_mask", + "ScoreTransformer", + "TransformerBlock", + "ScalarTokenizer", + "GaussianFourierEmbedding", + "BaseSDE", + "VESDE", + "VPSDE", + "SDE_CLASSES", + "build_sde", + "build_base_mask", + "get_condition_mask_fn", + "joint_condition_mask", + "posterior_condition_mask", + "likelihood_condition_mask", + "structured_random_condition_mask", + "train_simformer", + "denoising_score_matching_loss", + "adaptive_grad_clip_", + "mean_std_per_node", + "merge_config", + "DEFAULT_MODEL_CONFIG", + "DEFAULT_SDE_CONFIG", + "DEFAULT_TRAIN_CONFIG", + "euler_maruyama_reverse", + "heun_probability_flow", + "guided_euler_maruyama", + "interval_constraint_score", + "probability_flow_log_prob", +] diff --git a/src/synference/simformer/masks.py b/src/synference/simformer/masks.py new file mode 100644 index 0000000..e1c2657 --- /dev/null +++ b/src/synference/simformer/masks.py @@ -0,0 +1,194 @@ +"""Condition-mask samplers and attention (edge) mask builders. + +Two distinct mask concepts are used by the Simformer — do not confuse them: + +- **Condition mask**: boolean vector over nodes ``[theta..., x...]`` where True marks + an *observed* (conditioned) variable. Sampled per training example; supplied by the + user at inference time. +- **Edge / base (attention) mask**: boolean ``(T, T)`` matrix where ``mask[i, j] = True`` + allows token ``i`` to attend to token ``j``. Encodes dependency structure; ``None`` + means dense attention (the paper's training default). +""" + +from typing import Callable, Optional, Union + +import numpy as np +import torch + + +def _fix_all_true_rows(condition_mask: torch.Tensor) -> torch.Tensor: + """Force rows where every node is conditioned back to fully latent (all False).""" + all_true = condition_mask.all(dim=-1, keepdim=True) + return condition_mask & ~all_true + + +def joint_condition_mask( + num_samples: int, + theta_dim: int, + x_dim: int, + generator: Optional[torch.Generator] = None, +) -> torch.Tensor: + """All-False masks: model the full joint distribution.""" + return torch.zeros(num_samples, theta_dim + x_dim, dtype=torch.bool) + + +def posterior_condition_mask( + num_samples: int, + theta_dim: int, + x_dim: int, + generator: Optional[torch.Generator] = None, +) -> torch.Tensor: + """Condition on all data nodes: standard posterior masks.""" + row = torch.tensor([False] * theta_dim + [True] * x_dim) + return row.expand(num_samples, -1).clone() + + +def likelihood_condition_mask( + num_samples: int, + theta_dim: int, + x_dim: int, + generator: Optional[torch.Generator] = None, +) -> torch.Tensor: + """Condition on all parameter nodes: likelihood masks.""" + row = torch.tensor([True] * theta_dim + [False] * x_dim) + return row.expand(num_samples, -1).clone() + + +def random_condition_mask( + num_samples: int, + theta_dim: int, + x_dim: int, + generator: Optional[torch.Generator] = None, + alpha: float = 1.0, + beta: float = 4.0, +) -> torch.Tensor: + """Bernoulli masks with per-batch probability drawn from ``Beta(alpha, beta)``.""" + beta_dist = torch.distributions.Beta(alpha, beta) + prob = beta_dist.sample() + mask = torch.rand(num_samples, theta_dim + x_dim, generator=generator) < prob + return _fix_all_true_rows(mask) + + +def structured_random_condition_mask( + num_samples: int, + theta_dim: int, + x_dim: int, + generator: Optional[torch.Generator] = None, + p_joint: float = 0.2, + p_posterior: float = 0.2, + p_likelihood: float = 0.2, + p_rnd1: float = 0.2, + p_rnd2: float = 0.2, + rnd1_prob: float = 0.3, + rnd2_prob: float = 0.7, +) -> torch.Tensor: + """Sample per-row masks from {joint, posterior, likelihood, two random masks}. + + Matches ``scoresbibm.utils.condition_masks.sample_strutured_conditional_mask``: + the two Bernoulli candidate masks are drawn once per call and shared across the + rows that select them; rows that end up all-True are forced back to all-False. + + Args: + num_samples: Number of mask rows to draw. + theta_dim: Number of parameter nodes. + x_dim: Number of data nodes. + generator: Optional torch random generator. + p_joint: Probability of an all-latent row. + p_posterior: Probability of a posterior row. + p_likelihood: Probability of a likelihood row. + p_rnd1: Probability of the first random candidate row. + p_rnd2: Probability of the second random candidate row. + rnd1_prob: Bernoulli probability of the first random candidate. + rnd2_prob: Bernoulli probability of the second random candidate. + + Returns: + Boolean condition masks of shape ``(num_samples, theta_dim + x_dim)``. + """ + total = theta_dim + x_dim + candidates = torch.stack( + [ + torch.zeros(total, dtype=torch.bool), + torch.tensor([False] * theta_dim + [True] * x_dim), + torch.tensor([True] * theta_dim + [False] * x_dim), + torch.rand(total, generator=generator) < rnd1_prob, + torch.rand(total, generator=generator) < rnd2_prob, + ] + ) + probs = torch.tensor([p_joint, p_posterior, p_likelihood, p_rnd1, p_rnd2]) + choice = torch.multinomial(probs, num_samples, replacement=True, generator=generator) + return _fix_all_true_rows(candidates[choice]) + + +CONDITION_MASK_FNS = { + "joint": joint_condition_mask, + "posterior": posterior_condition_mask, + "likelihood": likelihood_condition_mask, + "random": random_condition_mask, + "structured_random": structured_random_condition_mask, +} + + +def get_condition_mask_fn(name: str, **kwargs) -> Callable: + """Look up a condition-mask sampler by name. + + Args: + name: One of ``joint``, ``posterior``, ``likelihood``, ``random``, + ``structured_random``. + **kwargs: Fixed keyword arguments bound to the sampler (e.g. ``p_joint``). + + Returns: + A callable ``fn(num_samples, theta_dim, x_dim, generator=None)``. + + Raises: + ValueError: If ``name`` is not a known sampler. + """ + key = name.lower() + if key not in CONDITION_MASK_FNS: + raise ValueError( + f"Unknown condition mask fn '{name}'. Choose from {sorted(CONDITION_MASK_FNS)}." + ) + if kwargs: + from functools import partial + + return partial(CONDITION_MASK_FNS[key], **kwargs) + return CONDITION_MASK_FNS[key] + + +def build_base_mask( + kind: Union[str, np.ndarray, torch.Tensor, None], + theta_dim: int, + x_dim: int, +) -> Optional[torch.Tensor]: + """Build the base (attention) mask over nodes ordered ``[theta..., x...]``. + + Args: + kind: ``"full"``/``None`` for dense attention (returns None); ``"directed"`` + (alias ``"causal"``) for the structured mask where parameters attend only + to themselves, data attend to all parameters and causally within data, and + parameters do not attend to data; or an explicit boolean ``(T, T)`` array. + theta_dim: Number of parameter nodes. + x_dim: Number of data nodes. + + Returns: + A boolean ``(T, T)`` tensor, or None for dense attention. + + Raises: + ValueError: If ``kind`` is an unknown string or a wrongly shaped array. + """ + total = theta_dim + x_dim + if kind is None or (isinstance(kind, str) and kind.lower() == "full"): + return None + if isinstance(kind, str): + if kind.lower() in ("directed", "causal"): + mask = torch.zeros(total, total, dtype=torch.bool) + mask[:theta_dim, :theta_dim] = torch.eye(theta_dim, dtype=torch.bool) + mask[theta_dim:, :theta_dim] = True + mask[theta_dim:, theta_dim:] = torch.tril(torch.ones(x_dim, x_dim, dtype=torch.bool)) + return mask + raise ValueError( + f"Unknown base mask kind '{kind}'. Use 'full', 'directed'/'causal', or an array." + ) + mask = torch.as_tensor(np.asarray(kind), dtype=torch.bool) + if mask.shape != (total, total): + raise ValueError(f"Custom base mask must have shape ({total}, {total}), got {mask.shape}.") + return mask diff --git a/src/synference/simformer/model.py b/src/synference/simformer/model.py new file mode 100644 index 0000000..ac936dd --- /dev/null +++ b/src/synference/simformer/model.py @@ -0,0 +1,596 @@ +"""Trained Simformer model wrapper and box prior. + +:class:`SimformerModel` bundles the trained score network, the SDE, per-node z-score +statistics, and metadata into a single object with a numpy-friendly inference API +(sampling with arbitrary condition masks, interval-constrained sampling via guidance, +and probability-flow log probabilities). It serializes to a plain ``torch.save`` dict +and is reloadable in a fresh process. +""" + +from typing import Dict, Optional, Sequence, Tuple, Union + +import numpy as np +import torch + +from .nn import ScoreTransformer +from .sampling import ( + euler_maruyama_reverse, + guided_euler_maruyama, + heun_probability_flow, + interval_constraint_score, + probability_flow_log_prob, +) +from .sde import BaseSDE, build_sde + +SAVE_FORMAT_VERSION = 1 + + +class UniformBoxPrior: + """Independent uniform prior over a named box, matching the old ``GalaxyPrior`` API. + + Attributes: + prior_ranges: Mapping of parameter name to ``(low, high)``. + param_order: Parameter names in sampling order. + theta_dim: Number of parameters. + """ + + def __init__(self, prior_ranges: Dict[str, Tuple[float, float]], param_order: Sequence[str]): + """Initialize the prior. + + Args: + prior_ranges: Mapping of parameter name to ``(low, high)`` bounds. + param_order: Parameter names defining the dimension order. + """ + self.prior_ranges = {name: tuple(prior_ranges[name]) for name in param_order} + self.param_order = list(param_order) + self.theta_dim = len(self.param_order) + lows = torch.tensor([self.prior_ranges[p][0] for p in self.param_order]) + highs = torch.tensor([self.prior_ranges[p][1] for p in self.param_order]) + self.low = lows.float() + self.high = highs.float() + + def sample( + self, + sample_shape: Union[int, Tuple[int, ...]] = (1,), + generator: Optional[torch.Generator] = None, + ) -> torch.Tensor: + """Draw uniform samples from the box. + + Args: + sample_shape: Number of samples or a shape tuple ``(n,)``. + generator: Optional torch random generator. + + Returns: + Samples of shape ``(n, theta_dim)``. + """ + if isinstance(sample_shape, int): + sample_shape = (sample_shape,) + n = int(np.prod(sample_shape)) + u = torch.rand(n, self.theta_dim, generator=generator) + return self.low + u * (self.high - self.low) + + def sample_n(self, num_samples: int) -> torch.Tensor: + """Alias of :meth:`sample` taking a plain integer (SBI_Fitter compatibility).""" + return self.sample((num_samples,)) + + def log_prob(self, theta: torch.Tensor) -> torch.Tensor: + """Log density of ``theta`` under the box prior. + + Args: + theta: Parameter values of shape ``(..., theta_dim)``. + + Returns: + Log probabilities of shape ``(...,)`` (``-inf`` outside the box). + """ + theta = torch.as_tensor(theta, dtype=torch.float32) + inside = ((theta >= self.low) & (theta <= self.high)).all(dim=-1) + log_volume = torch.log(self.high - self.low).sum() + out = torch.full(theta.shape[:-1], -torch.inf) + out[inside] = -log_volume + return out + + +class SimformerModel: + """A trained all-conditional Simformer score model. + + The node order is ``[theta..., x...]``; condition masks are boolean vectors over + the nodes with True marking observed variables. All public methods accept and + return values in original (un-z-scored) units. + """ + + def __init__( + self, + net: ScoreTransformer, + sde: BaseSDE, + theta_dim: int, + x_dim: int, + node_ids: Optional[np.ndarray] = None, + base_mask: Optional[np.ndarray] = None, + z_score_mean: Optional[np.ndarray] = None, + z_score_std: Optional[np.ndarray] = None, + meta: Optional[dict] = None, + sampling_defaults: Optional[dict] = None, + ): + """Initialize the wrapper. + + Args: + net: Trained score network. + sde: The diffusion SDE (holding per-node data statistics). + theta_dim: Number of parameter nodes. + x_dim: Number of data nodes. + node_ids: Integer node ids, shape ``(theta_dim + x_dim,)``. Defaults to + ``arange``. + base_mask: Optional boolean base attention mask, shape ``(T, T)``. + z_score_mean: Per-node z-score means, shape ``(T,)``, or None. + z_score_std: Per-node z-score stddevs, shape ``(T,)``, or None. + meta: Metadata dict (parameter/feature names, prior ranges, configs). + sampling_defaults: Default sampling settings (``num_steps``, ``method``). + """ + self.net = net + self.sde = sde + self.theta_dim = int(theta_dim) + self.x_dim = int(x_dim) + self.num_nodes = self.theta_dim + self.x_dim + if node_ids is None: + node_ids = np.arange(self.num_nodes) + self.node_ids = np.asarray(node_ids, dtype=np.int64) + self.base_mask = None if base_mask is None else np.asarray(base_mask, dtype=bool) + self.z_score_mean = None if z_score_mean is None else np.asarray(z_score_mean, np.float32) + self.z_score_std = None if z_score_std is None else np.asarray(z_score_std, np.float32) + self.meta = meta or {} + self.sampling_defaults = sampling_defaults or {"num_steps": 500, "method": "sde"} + self.device = torch.device("cpu") + + def to(self, device: Union[str, torch.device]) -> "SimformerModel": + """Move the model to ``device`` and return self.""" + self.device = torch.device(device) + self.net.to(self.device) + self.sde.to(self.device) + return self + + # ------------------------------------------------------------------ helpers + def _as_condition_mask(self, condition_mask) -> torch.Tensor: + mask = torch.as_tensor(np.asarray(condition_mask), dtype=torch.bool) + # make sure mask on the correct device for the score network + mask = mask.to(self.device) + if mask.shape != (self.num_nodes,): + raise ValueError( + f"condition_mask must have shape ({self.num_nodes},), got {tuple(mask.shape)}." + ) + if mask.all(): + raise ValueError("condition_mask cannot condition on every node.") + return mask + + def _z_stats(self) -> Tuple[torch.Tensor, torch.Tensor]: + if self.z_score_mean is None: + mean = torch.zeros(self.num_nodes, device=self.device) + std = torch.ones(self.num_nodes, device=self.device) + else: + mean = torch.as_tensor(self.z_score_mean, device=self.device) + std = torch.as_tensor(self.z_score_std, device=self.device) + return mean, std + + def _resolve_edge_mask(self, edge_mask) -> Optional[torch.Tensor]: + if edge_mask is None: + if self.base_mask is None: + return None + edge_mask = self.base_mask + mask = torch.as_tensor(np.asarray(edge_mask), dtype=torch.bool, device=self.device) + return mask + + def score( + self, + t: torch.Tensor, + x: torch.Tensor, + condition_mask: torch.Tensor, + edge_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Evaluate the score estimate in z-scored space. + + Args: + t: Scalar diffusion time (tensor). + x: State of shape ``(B, T)``. + condition_mask: Boolean mask of shape ``(T,)``. + edge_mask: Optional attention mask; defaults to the stored base mask. + + Returns: + Score of shape ``(B, T)``. + """ + node_ids = torch.as_tensor(self.node_ids, device=x.device) + raw = self.net( + torch.atleast_1d(t), + x[..., None], + node_ids, + condition_mask.to(x.device), + edge_mask=self._resolve_edge_mask(edge_mask), + )[..., 0] + return raw / self.sde.output_scale(t) + + def _score_fn(self, condition_mask: torch.Tensor, edge_mask=None): + def score_fn(t: torch.Tensor, x: torch.Tensor) -> torch.Tensor: + return self.score(t, x, condition_mask, edge_mask=edge_mask) + + return score_fn + + def _init_x_T( + self, + num_samples: int, + x_o_z: torch.Tensor, + condition_mask: torch.Tensor, + generator: Optional[torch.Generator], + ) -> torch.Tensor: + mean_end = self.sde.marginal_mean_end().to(self.device) + std_end = self.sde.marginal_std_end().to(self.device) + noise = torch.randn(num_samples, self.num_nodes, device=self.device, generator=generator) + x_T = mean_end + std_end * noise + x_T[:, condition_mask] = x_o_z + return x_T + + def _prepare_x_o(self, x_o, condition_mask: torch.Tensor) -> torch.Tensor: + """Z-score observed values ordered by the conditioned node indices.""" + x_o = torch.as_tensor(np.asarray(x_o, dtype=np.float32), device=self.device).reshape(-1) + n_cond = int(condition_mask.sum()) + if x_o.numel() != n_cond: + raise ValueError(f"x_o has {x_o.numel()} values but the mask conditions on {n_cond}.") + mean, std = self._z_stats() + return (x_o - mean[condition_mask]) / std[condition_mask] + + # ----------------------------------------------------------------- sampling + @torch.no_grad() + def sample( + self, + num_samples: int, + x_o, + condition_mask, + num_steps: Optional[int] = None, + method: Optional[str] = None, + edge_mask=None, + generator: Optional[torch.Generator] = None, + ) -> np.ndarray: + """Sample the latent (unconditioned) nodes given observed values. + + Args: + num_samples: Number of posterior/conditional samples to draw. + x_o: Observed values (original units) for the conditioned nodes, in node + order; shape ``(n_conditioned,)``. + condition_mask: Boolean mask of shape ``(T,)``; True marks observed nodes. + num_steps: Number of integration steps (default from + ``sampling_defaults``). + method: ``"sde"`` (Euler-Maruyama) or ``"ode"`` (Heun probability flow). + edge_mask: Optional attention mask override. + generator: Optional torch random generator on the model device. + + Returns: + Samples of shape ``(num_samples, n_latent)`` in original units. + + Raises: + ValueError: If ``method`` is unknown. + """ + condition_mask = self._as_condition_mask(condition_mask) + num_steps = num_steps or self.sampling_defaults.get("num_steps", 500) + method = method or self.sampling_defaults.get("method", "sde") + + x_o_z = self._prepare_x_o(x_o, condition_mask) + x_T = self._init_x_T(num_samples, x_o_z, condition_mask, generator) + score_fn = self._score_fn(condition_mask, edge_mask=edge_mask) + + if method == "sde": + x_final = euler_maruyama_reverse( + score_fn, self.sde, x_T, condition_mask, num_steps=num_steps, generator=generator + ) + elif method == "ode": + x_final = heun_probability_flow( + score_fn, self.sde, x_T, condition_mask, num_steps=num_steps + ) + else: + raise ValueError(f"Unknown sampling method '{method}'. Use 'sde' or 'ode'.") + + return self._extract_latents(x_final, condition_mask) + + def _extract_latents(self, x_final: torch.Tensor, condition_mask: torch.Tensor) -> np.ndarray: + mean, std = self._z_stats() + latents = x_final[:, ~condition_mask] + latents = latents * std[~condition_mask] + mean[~condition_mask] + return latents.cpu().numpy() + + @torch.no_grad() + def sample_batched( + self, + num_samples: int, + x_o_batch, + condition_mask, + batch_size: int = 100, + **kwargs, + ) -> np.ndarray: + """Sample conditionals for a batch of observations. + + Args: + num_samples: Samples per observation. + x_o_batch: Observations of shape ``(n_obs, n_conditioned)``. + condition_mask: Boolean mask of shape ``(T,)`` shared by all observations. + batch_size: Number of observations integrated simultaneously (the + effective state batch is ``batch_size * num_samples``). + **kwargs: Forwarded to :meth:`sample` (``num_steps``, ``method``, + ``edge_mask``, ``generator``). + + Returns: + Samples of shape ``(n_obs, num_samples, n_latent)`` in original units. + """ + condition_mask = self._as_condition_mask(condition_mask) + x_o_batch = np.atleast_2d(np.asarray(x_o_batch, dtype=np.float32)) + n_obs = x_o_batch.shape[0] + num_steps = kwargs.pop("num_steps", None) or self.sampling_defaults.get("num_steps", 500) + method = kwargs.pop("method", None) or self.sampling_defaults.get("method", "sde") + edge_mask = kwargs.pop("edge_mask", None) + generator = kwargs.pop("generator", None) + if kwargs: + raise TypeError(f"Unexpected keyword arguments: {sorted(kwargs)}") + + score_fn = self._score_fn(condition_mask, edge_mask=edge_mask) + results = [] + for start in range(0, n_obs, batch_size): + chunk = x_o_batch[start : start + batch_size] + n_chunk = chunk.shape[0] + mean, std = self._z_stats() + x_o_z = (torch.as_tensor(chunk, device=self.device) - mean[condition_mask]) / std[ + condition_mask + ] + x_o_z = x_o_z.repeat_interleave(num_samples, dim=0) + x_T = self._init_x_T(n_chunk * num_samples, x_o_z, condition_mask, generator) + if method == "sde": + x_final = euler_maruyama_reverse( + score_fn, + self.sde, + x_T, + condition_mask, + num_steps=num_steps, + generator=generator, + ) + elif method == "ode": + x_final = heun_probability_flow( + score_fn, self.sde, x_T, condition_mask, num_steps=num_steps + ) + else: + raise ValueError(f"Unknown sampling method '{method}'. Use 'sde' or 'ode'.") + latents = self._extract_latents(x_final, condition_mask) + results.append(latents.reshape(n_chunk, num_samples, -1)) + return np.concatenate(results, axis=0) + + @torch.no_grad() + def sample_intervals( + self, + num_samples: int, + x_o, + condition_mask, + constraint_mask, + a=None, + b=None, + scale_bias: float = 0.0, + num_steps: Optional[int] = None, + edge_mask=None, + generator: Optional[torch.Generator] = None, + ) -> np.ndarray: + """Sample with interval constraints on a subset of nodes via guidance. + + Constrained nodes are removed from the hard condition mask (a node cannot be + both clamped and constrained) and guided towards the box ``[a, b]`` using the + Tweedie generalized-guidance scheme. + + Args: + num_samples: Number of samples to draw. + x_o: Observed values (original units) for nodes that remain in the + effective condition mask (``condition_mask & ~constraint_mask``), + in node order. + condition_mask: Boolean mask of shape ``(T,)`` of observed nodes. + constraint_mask: Boolean mask of shape ``(T,)`` of interval-constrained + nodes. + a: Lower bounds in original units — scalar, array of shape + ``(n_constrained,)``, or None for unbounded below. + b: Upper bounds in original units, like ``a``. + scale_bias: Additive bias in the guidance scale + ``1 / (marginal_var(t) + bias)``; the faithful default is 0, but + values of order 1e-2 improve stability for the VE SDE near ``t=0``. + num_steps: Number of integration steps. + edge_mask: Optional attention mask override. + generator: Optional torch random generator. + + Returns: + Samples of shape ``(num_samples, n_latent)`` in original units, where + latent nodes are all nodes not in the *effective* condition mask + (i.e. constrained nodes are included in the output). + """ + condition_mask = self._as_condition_mask(condition_mask) + constraint_mask = torch.as_tensor(np.asarray(constraint_mask), dtype=torch.bool) + if constraint_mask.shape != (self.num_nodes,): + raise ValueError(f"constraint_mask must have shape ({self.num_nodes},).") + effective_condition = condition_mask & ~constraint_mask + num_steps = num_steps or self.sampling_defaults.get("num_steps", 500) + + mean, std = self._z_stats() + a_full = self._expand_bounds(a, constraint_mask, mean, std) + b_full = self._expand_bounds(b, constraint_mask, mean, std) + + x_o_z = self._prepare_x_o(x_o, effective_condition) + x_T = self._init_x_T(num_samples, x_o_z, effective_condition, generator) + score_fn = self._score_fn(effective_condition, edge_mask=edge_mask) + constraint_mask_dev = constraint_mask.to(self.device) + + def constraint_score_fn(x0_hat: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + scale = 1.0 / (self.sde.transition_var(t) + scale_bias) + return interval_constraint_score(x0_hat, scale, constraint_mask_dev, a=a_full, b=b_full) + + x_final = guided_euler_maruyama( + score_fn, + self.sde, + x_T, + effective_condition, + constraint_score_fn, + num_steps=num_steps, + generator=generator, + ) + return self._extract_latents(x_final, effective_condition) + + def _expand_bounds(self, bound, constraint_mask, mean, std) -> Optional[torch.Tensor]: + """Expand user bounds to a z-scored full-node vector (non-finite = unbounded).""" + if bound is None: + return None + full = torch.full((self.num_nodes,), torch.nan, device=self.device) + values = torch.as_tensor( + np.broadcast_to( + np.asarray(bound, dtype=np.float32), (int(constraint_mask.sum()),) + ).copy(), + device=self.device, + ) + full[constraint_mask] = values + return (full - mean) / std + + # ---------------------------------------------------------------- log prob + def log_prob( + self, + theta, + x_o, + condition_mask, + num_steps: int = 250, + divergence: str = "exact", + hutchinson_probes: int = 8, + edge_mask=None, + generator: Optional[torch.Generator] = None, + ) -> np.ndarray: + """Log probability of latent values via the probability-flow ODE. + + Args: + theta: Latent values (original units), shape ``(n, n_latent)`` or + ``(n_latent,)``, in node order of the unconditioned nodes. + x_o: Observed values (original units) for the conditioned nodes — either a + single vector shared by all rows of ``theta``, or one row per ``theta`` + row (shape ``(n, n_conditioned)``). + condition_mask: Boolean mask of shape ``(T,)``; True marks observed nodes. + num_steps: Number of ODE steps. + divergence: ``"exact"`` or ``"hutchinson"``. + hutchinson_probes: Probes for the Hutchinson estimator. + edge_mask: Optional attention mask override. + generator: Optional torch random generator. + + Returns: + Log probabilities of shape ``(n,)`` (scalar array for 1D input), + in original units (z-score Jacobian included). + """ + condition_mask = self._as_condition_mask(condition_mask) + theta = np.atleast_2d(np.asarray(theta, dtype=np.float32)) + n_latent = int((~condition_mask).sum()) + if theta.shape[1] != n_latent: + raise ValueError(f"theta must have {n_latent} columns, got {theta.shape[1]}.") + + mean, std = self._z_stats() + n_cond = int(condition_mask.sum()) + x_o_arr = np.atleast_2d(np.asarray(x_o, dtype=np.float32)).reshape(-1, n_cond) + if x_o_arr.shape[0] not in (1, theta.shape[0]): + raise ValueError( + f"x_o must be a single vector or one row per theta row, got {x_o_arr.shape}." + ) + x_o_t = torch.as_tensor(x_o_arr, device=self.device) + x_o_z = (x_o_t - mean[condition_mask]) / std[condition_mask] + theta_t = torch.as_tensor(theta, device=self.device) + theta_z = (theta_t - mean[~condition_mask]) / std[~condition_mask] + + x0 = torch.zeros(theta.shape[0], self.num_nodes, device=self.device) + x0[:, condition_mask] = x_o_z + x0[:, ~condition_mask] = theta_z + + score_fn = self._score_fn(condition_mask, edge_mask=edge_mask) + log_p_z = probability_flow_log_prob( + score_fn, + self.sde, + x0, + condition_mask, + num_steps=num_steps, + divergence=divergence, + hutchinson_probes=hutchinson_probes, + generator=generator, + ) + jacobian = torch.log(std[~condition_mask]).sum() + out = (log_p_z - jacobian).cpu().numpy() + return out if out.size > 1 else out.reshape(()) + + # ------------------------------------------------------------ serialization + def save(self, path: str) -> None: + """Serialize the model to ``path`` as a plain ``torch.save`` dict. + + Args: + path: Destination file path. + """ + payload = { + "format_version": SAVE_FORMAT_VERSION, + "state_dict": {k: v.cpu() for k, v in self.net.state_dict().items()}, + "model_config": self.net.config, + "sde_config": self.sde.config, + "sde_x0_mean": self.sde.x0_mean.cpu().numpy(), + "sde_x0_var": self.sde.x0_var.cpu().numpy(), + "z_score_mean": self.z_score_mean, + "z_score_std": self.z_score_std, + "theta_dim": self.theta_dim, + "x_dim": self.x_dim, + "node_ids": self.node_ids, + "base_mask": self.base_mask, + "meta": self.meta, + "sampling_defaults": self.sampling_defaults, + } + torch.save(payload, path) + + @classmethod + def load(cls, path: str, map_location: Union[str, torch.device] = "cpu") -> "SimformerModel": + """Load a model saved with :meth:`save`. + + Args: + path: Path to the saved file. + map_location: Device to load onto. + + Returns: + The reconstructed model. + + Raises: + ValueError: If the file has an unknown format version. + """ + payload = torch.load(path, map_location="cpu", weights_only=False) + version = payload.get("format_version") + if version != SAVE_FORMAT_VERSION: + raise ValueError(f"Unsupported Simformer save format version: {version}.") + net = ScoreTransformer(**payload["model_config"]) + net.load_state_dict(payload["state_dict"]) + net.eval() + sde_config = dict(payload["sde_config"]) + sde = build_sde( + sde_config.pop("name"), + payload["sde_x0_mean"], + payload["sde_x0_var"], + **sde_config, + ) + model = cls( + net=net, + sde=sde, + theta_dim=payload["theta_dim"], + x_dim=payload["x_dim"], + node_ids=payload["node_ids"], + base_mask=payload["base_mask"], + z_score_mean=payload["z_score_mean"], + z_score_std=payload["z_score_std"], + meta=payload["meta"], + sampling_defaults=payload["sampling_defaults"], + ) + return model.to(map_location) + + +def posterior_mask(theta_dim: int, x_dim: int) -> np.ndarray: + """Standard posterior condition mask ``[False]*theta_dim + [True]*x_dim``. + + Args: + theta_dim: Number of parameter nodes. + x_dim: Number of data nodes. + + Returns: + Boolean array of shape ``(theta_dim + x_dim,)``. + """ + return np.array([False] * theta_dim + [True] * x_dim) + + +__all__ = ["SimformerModel", "UniformBoxPrior", "posterior_mask"] diff --git a/src/synference/simformer/nn.py b/src/synference/simformer/nn.py new file mode 100644 index 0000000..b91b555 --- /dev/null +++ b/src/synference/simformer/nn.py @@ -0,0 +1,381 @@ +"""Neural network components for the PyTorch Simformer. + +Faithful port of the score-transformer architecture from Gloeckler et al. (2024), +"All-in-one simulation-based inference" (original JAX/haiku implementation in +``probjax``/``scoresbibm``). Each scalar variable is one token; the network predicts +the (scaled) score of the noised joint distribution at diffusion time ``t``. +""" + +import math +from typing import Optional, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F + +TRUNCATED_NORMAL_STDDEV_FACTOR = 0.87962566103423978 +"""Correction factor so a +/-2 sigma truncated normal has the requested stddev.""" + + +def variance_scaling_init_(tensor: torch.Tensor, scale: float, fan_in: int) -> torch.Tensor: + """Initialize ``tensor`` in-place with haiku-style fan-in variance scaling. + + Draws from a truncated normal (+/- 2 stddev) with stddev ``sqrt(scale / fan_in)``, + matching ``hk.initializers.VarianceScaling(scale)`` defaults. + + Args: + tensor: Tensor to initialize in-place. + scale: Variance scale factor. + fan_in: Number of input units. + + Returns: + The initialized tensor. + """ + stddev = math.sqrt(scale / max(1.0, fan_in)) / TRUNCATED_NORMAL_STDDEV_FACTOR + return nn.init.trunc_normal_(tensor, mean=0.0, std=stddev, a=-2 * stddev, b=2 * stddev) + + +class GaussianFourierEmbedding(nn.Module): + """Gaussian Fourier feature embedding, mostly used to embed diffusion time. + + The random projection matrix ``B`` is drawn once at initialization and frozen + (registered as a buffer so it is saved/restored with the state dict), matching + the stop-gradient behaviour of the original implementation. + """ + + def __init__(self, output_dim: int = 128, input_dim: int = 1): + """Initialize the embedding. + + Args: + output_dim: Output embedding dimension. + input_dim: Dimension of the input (1 for scalar time). + """ + super().__init__() + self.output_dim = output_dim + half_dim = output_dim // 2 + 1 + self.register_buffer("B", torch.randn(half_dim, input_dim)) + + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + """Embed ``inputs`` of shape ``(..., input_dim)`` to ``(..., output_dim)``. + + Args: + inputs: Input tensor with the projected dimension last. + + Returns: + Fourier features ``[cos(2 pi x B^T), sin(2 pi x B^T)]`` truncated to + ``output_dim``. + """ + proj = 2 * math.pi * inputs @ self.B.T + out = torch.cat([torch.cos(proj), torch.sin(proj)], dim=-1) + return out[..., : self.output_dim] + + +class ScalarTokenizer(nn.Module): + """Tokenize scalar variables into embedding vectors. + + Each token is the concatenation of a (frozen) learned node-id embedding and the + scalar value tiled to the remaining width. The node-id embedding is initialized + orthogonally and never trained, replicating the ``stop_gradient`` applied in the + original ``probjax`` tokenizer. + """ + + def __init__(self, output_dim: int, num_nodes: int): + """Initialize the tokenizer. + + Args: + output_dim: Total token dimension (split between id and value parts). + num_nodes: Maximum number of distinct node ids. + """ + super().__init__() + self.output_dim = output_dim + self.id_dim = output_dim // 2 + self.value_dim = output_dim - self.id_dim + self.node_embedding = nn.Embedding(num_nodes, self.id_dim) + nn.init.orthogonal_(self.node_embedding.weight, gain=0.5) + # The original implementation stop-gradients the node embedding, so it stays + # at its random orthogonal initialization. + self.node_embedding.weight.requires_grad_(False) + + def forward(self, node_ids: torch.Tensor, values: torch.Tensor) -> torch.Tensor: + """Tokenize values. + + Args: + node_ids: Integer node ids of shape ``(num_nodes,)`` or ``(B, num_nodes)``. + values: Scalar values of shape ``(B, num_nodes, 1)``. + + Returns: + Tokens of shape ``(B, num_nodes, output_dim)``. + """ + batch, num_nodes, _ = values.shape + node_ids = node_ids.reshape(-1, num_nodes).long() + id_embedding = self.node_embedding(node_ids) # (1 or B, num_nodes, id_dim) + id_embedding = id_embedding.expand(batch, num_nodes, self.id_dim) + value_embedding = values.expand(batch, num_nodes, self.value_dim) + return torch.cat([id_embedding, value_embedding], dim=-1) + + +class MultiHeadAttention(nn.Module): + """Multi-head attention with an optional boolean attention mask. + + ``mask[i, j] = True`` means token ``i`` may attend to token ``j``. Masked logits + are set to ``-1e30`` before the softmax, as in the original implementation. + """ + + def __init__(self, model_size: int, num_heads: int, key_size: int, init_scale: float): + """Initialize projections. + + Args: + model_size: Token embedding width. + num_heads: Number of attention heads. + key_size: Per-head key/query/value size. + init_scale: Variance-scaling factor for weight init. + """ + super().__init__() + self.num_heads = num_heads + self.key_size = key_size + inner = num_heads * key_size + self.query_proj = nn.Linear(model_size, inner) + self.key_proj = nn.Linear(model_size, inner) + self.value_proj = nn.Linear(model_size, inner) + self.out_proj = nn.Linear(inner, model_size) + for layer in (self.query_proj, self.key_proj, self.value_proj): + variance_scaling_init_(layer.weight, init_scale, model_size) + nn.init.zeros_(layer.bias) + variance_scaling_init_(self.out_proj.weight, init_scale, inner) + nn.init.zeros_(self.out_proj.bias) + + def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor: + """Apply masked self-attention. + + Args: + x: Tokens of shape ``(B, T, model_size)``. + mask: Optional boolean mask of shape ``(T, T)`` or ``(B, T, T)``. + + Returns: + Attended tokens of shape ``(B, T, model_size)``. + """ + batch, num_tokens, _ = x.shape + + def split_heads(t: torch.Tensor) -> torch.Tensor: + return t.reshape(batch, num_tokens, self.num_heads, self.key_size).transpose(1, 2) + + q = split_heads(self.query_proj(x)) # (B, H, T, K) + k = split_heads(self.key_proj(x)) + v = split_heads(self.value_proj(x)) + + logits = q @ k.transpose(-2, -1) / math.sqrt(self.key_size) # (B, H, T, T) + if mask is not None: + if mask.ndim == 2: + mask = mask[None, None, :, :] + elif mask.ndim == 3: + mask = mask[:, None, :, :] + else: + raise ValueError(f"Mask must have ndim 2 or 3, got {mask.ndim}.") + logits = logits.masked_fill(~mask, -1e30) + weights = F.softmax(logits, dim=-1) + attn = weights @ v # (B, H, T, K) + attn = attn.transpose(1, 2).reshape(batch, num_tokens, -1) + return self.out_proj(attn) + + +class TransformerBlock(nn.Module): + """Pre-LayerNorm transformer block with time context injected in the MLP. + + Structure (matching ``probjax.nn.transformers.Transformer``):: + + h = LN(h) + h = h + MHA(h, mask) + h = LN(h) + h = h + ( + MLP(h) + + gelu( + Linear(time) + ) + ) + """ + + def __init__( + self, + model_size: int, + num_heads: int, + key_size: int, + widening_factor: int, + num_hidden_layers: int, + time_dim: int, + init_scale: float, + ): + """Initialize the block. + + Args: + model_size: Token embedding width. + num_heads: Number of attention heads. + key_size: Per-head attention size. + widening_factor: MLP hidden width multiplier. + num_hidden_layers: Number of hidden layers in the MLP. + time_dim: Dimension of the time-context embedding. + init_scale: Variance-scaling factor for weight init. + """ + super().__init__() + self.ln_attn = nn.LayerNorm(model_size) + self.attention = MultiHeadAttention(model_size, num_heads, key_size, init_scale) + self.ln_mlp = nn.LayerNorm(model_size) + + hidden = widening_factor * model_size + mlp_layers: list[nn.Module] = [] + in_dim = model_size + for _ in range(num_hidden_layers): + layer = nn.Linear(in_dim, hidden) + variance_scaling_init_(layer.weight, init_scale, in_dim) + nn.init.zeros_(layer.bias) + mlp_layers += [layer, nn.GELU()] + in_dim = hidden + out_layer = nn.Linear(in_dim, model_size) + variance_scaling_init_(out_layer.weight, init_scale, in_dim) + nn.init.zeros_(out_layer.bias) + mlp_layers.append(out_layer) + self.mlp = nn.Sequential(*mlp_layers) + + self.time_proj = nn.Linear(time_dim, model_size) + variance_scaling_init_(self.time_proj.weight, init_scale, time_dim) + nn.init.zeros_(self.time_proj.bias) + + def forward( + self, + h: torch.Tensor, + time_embedding: torch.Tensor, + mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Apply the block. + + Args: + h: Tokens of shape ``(B, T, model_size)``. + time_embedding: Time context of shape ``(B, time_dim)``. + mask: Optional boolean attention mask, ``(T, T)`` or ``(B, T, T)``. + + Returns: + Updated tokens of shape ``(B, T, model_size)``. + """ + h = h + self.attention(self.ln_attn(h), mask=mask) + normed = self.ln_mlp(h) + dense = self.mlp(normed) + context = F.gelu(self.time_proj(time_embedding)) + dense = dense + context[:, None, :] + return h + dense + + +class ScoreTransformer(nn.Module): + """Transformer score network over per-variable tokens. + + The forward pass returns the *unscaled* network output of shape + ``(B, num_nodes, 1)``; dividing by the SDE marginal standard deviation (done in + :class:`synference.simformer.SimformerModel`) turns it into a score estimate. + + Defaults match the paper configuration (``score_transformer_small``). + """ + + def __init__( + self, + num_nodes: int, + token_dim: int = 40, + condition_token_dim: int = 10, + condition_token_init_scale: float = 0.1, + time_embedding_dim: int = 128, + num_heads: int = 4, + num_layers: int = 6, + attn_size: int = 10, + widening_factor: int = 3, + num_hidden_layers: int = 1, + ): + """Initialize the network. + + Args: + num_nodes: Total number of variables (theta dims + x dims). + token_dim: Tokenizer output width (id + value parts). + condition_token_dim: Width of the learned condition token. + condition_token_init_scale: Init stddev of the condition token. + time_embedding_dim: Gaussian Fourier time-embedding width. + num_heads: Attention heads per layer. + num_layers: Number of transformer blocks. + attn_size: Per-head key/query/value size. + widening_factor: MLP hidden width multiplier. + num_hidden_layers: Hidden layers per MLP block. + """ + super().__init__() + self.num_nodes = num_nodes + self.config = { + "num_nodes": num_nodes, + "token_dim": token_dim, + "condition_token_dim": condition_token_dim, + "condition_token_init_scale": condition_token_init_scale, + "time_embedding_dim": time_embedding_dim, + "num_heads": num_heads, + "num_layers": num_layers, + "attn_size": attn_size, + "widening_factor": widening_factor, + "num_hidden_layers": num_hidden_layers, + } + + self.tokenizer = ScalarTokenizer(token_dim, num_nodes) + self.time_embedding = GaussianFourierEmbedding(time_embedding_dim) + self.condition_token = nn.Parameter( + torch.randn(1, 1, condition_token_dim) * condition_token_init_scale + ) + + model_size = token_dim + condition_token_dim + init_scale = 2.0 / num_layers + self.blocks = nn.ModuleList( + TransformerBlock( + model_size=model_size, + num_heads=num_heads, + key_size=attn_size, + widening_factor=widening_factor, + num_hidden_layers=num_hidden_layers, + time_dim=time_embedding_dim, + init_scale=init_scale, + ) + for _ in range(num_layers) + ) + self.final_norm = nn.LayerNorm(model_size) + self.head = nn.Linear(model_size, 1) + variance_scaling_init_(self.head.weight, 1.0, model_size) + nn.init.zeros_(self.head.bias) + + def forward( + self, + t: torch.Tensor, + x: torch.Tensor, + node_ids: torch.Tensor, + condition_mask: torch.Tensor, + edge_mask: Union[torch.Tensor, None] = None, + ) -> torch.Tensor: + """Evaluate the network. + + Args: + t: Diffusion times of shape ``(B,)`` (or scalar, broadcast to the batch). + x: Variable values of shape ``(B, num_nodes, 1)``. + node_ids: Integer node ids of shape ``(num_nodes,)``. + condition_mask: Boolean mask, ``(num_nodes,)`` or ``(B, num_nodes)``; + True marks observed (conditioned) variables. + edge_mask: Optional boolean attention mask, ``(T, T)`` or ``(B, T, T)``. + None means dense attention. + + Returns: + Unscaled network output of shape ``(B, num_nodes, 1)``. + """ + batch, num_nodes, _ = x.shape + t = torch.atleast_1d(t).to(x) + if t.shape[0] == 1 and batch > 1: + t = t.expand(batch) + + tokens = self.tokenizer(node_ids, x) + condition_mask = condition_mask.reshape(-1, num_nodes, 1).to(x) + condition_token = condition_mask * self.condition_token + condition_token = condition_token.expand(batch, num_nodes, -1) + h = torch.cat([tokens, condition_token], dim=-1) + + time_embedding = self.time_embedding(t[..., None]) + + for block in self.blocks: + h = block(h, time_embedding, mask=edge_mask) + h = self.final_norm(h) + return self.head(h) diff --git a/src/synference/simformer/sampling.py b/src/synference/simformer/sampling.py new file mode 100644 index 0000000..c5a25ae --- /dev/null +++ b/src/synference/simformer/sampling.py @@ -0,0 +1,320 @@ +"""Fixed-step SDE/ODE integrators and guidance for the PyTorch Simformer. + +All integrators operate in the model's (z-scored) space on batched states of shape +``(B, T)`` with a boolean condition mask of shape ``(T,)`` (True = observed). Observed +entries are clamped: both drift and diffusion are multiplied by the latent indicator, +matching ``scoresbibm.methods.models.AllConditionalScoreModel``. + +``score_fn(t, x)`` takes a scalar time tensor and the ``(B, T)`` state and returns the +``(B, T)`` score estimate. +""" + +from typing import Callable, Optional, Tuple + +import torch + +from .sde import BaseSDE + + +def euler_maruyama_reverse( + score_fn: Callable, + sde: BaseSDE, + x_T: torch.Tensor, + condition_mask: torch.Tensor, + num_steps: int = 500, + generator: Optional[torch.Generator] = None, +) -> torch.Tensor: + """Integrate the reverse SDE from ``T_max`` to ``T_min`` with Euler-Maruyama. + + Args: + score_fn: Score function ``(t, x) -> score`` on ``(B, T)`` states. + sde: The diffusion SDE. + x_T: Initial state at ``T_max``, shape ``(B, T)``, observed entries already + set to their conditioning values. + condition_mask: Boolean mask of shape ``(T,)``; True marks observed nodes. + num_steps: Number of integration steps. + generator: Optional torch random generator (on the same device as ``x_T``). + + Returns: + The state at ``T_min``, shape ``(B, T)`` (observed entries unchanged). + """ + latent = (~condition_mask).to(x_T) + ts = torch.linspace(0.0, sde.T_max - sde.T_min, num_steps, device=x_T.device) + x = x_T + for n in range(num_steps - 1): + dt = ts[n + 1] - ts[n] + t = sde.T_max - ts[n] + score = score_fn(t, x) + g = sde.diffusion(t) + drift = -(sde.drift(t, x) - g**2 * score) * latent + noise = torch.randn(x.shape, device=x.device, generator=generator) + x = x + drift * dt + g * torch.sqrt(dt) * noise * latent + return x + + +def heun_probability_flow( + score_fn: Callable, + sde: BaseSDE, + x_T: torch.Tensor, + condition_mask: torch.Tensor, + num_steps: int = 500, +) -> torch.Tensor: + """Integrate the reverse probability-flow ODE with Heun's method. + + Args: + score_fn: Score function ``(t, x) -> score`` on ``(B, T)`` states. + sde: The diffusion SDE. + x_T: Initial state at ``T_max``, shape ``(B, T)``. + condition_mask: Boolean mask of shape ``(T,)``; True marks observed nodes. + num_steps: Number of integration steps. + + Returns: + The state at ``T_min``, shape ``(B, T)``. + """ + latent = (~condition_mask).to(x_T) + + def drift_backward(s: torch.Tensor, x: torch.Tensor) -> torch.Tensor: + t = sde.T_max - s + score = score_fn(t, x) + dx = sde.drift(t, x) - 0.5 * sde.diffusion(t) ** 2 * score + return -dx * latent + + ts = torch.linspace(0.0, sde.T_max - sde.T_min, num_steps, device=x_T.device) + x = x_T + for n in range(num_steps - 1): + dt = ts[n + 1] - ts[n] + k1 = drift_backward(ts[n], x) + k2 = drift_backward(ts[n + 1], x + dt * k1) + x = x + 0.5 * dt * (k1 + k2) + return x + + +def interval_constraint_score( + x0_hat: torch.Tensor, + scale: torch.Tensor, + constraint_mask: torch.Tensor, + a: Optional[torch.Tensor] = None, + b: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Gradient of the smoothed box-constraint log-probability at ``x0_hat``. + + Implements the analytic gradient of ``scoresbibm.methods.guidance.log_step_fn``: + ``sum log sigmoid(scale (x - a)) + log sigmoid(scale (b - x))`` over constrained + nodes, differentiated with respect to ``x0_hat``. + + Args: + x0_hat: Tweedie estimate of the clean state, shape ``(B, T)``. + scale: Guidance scale ``s(t)`` (scalar tensor). + constraint_mask: Boolean mask of shape ``(T,)`` marking constrained nodes. + a: Lower bounds, shape ``(T,)``; non-finite entries mean unbounded below. + b: Upper bounds, shape ``(T,)``; non-finite entries mean unbounded above. + + Returns: + Constraint score of shape ``(B, T)`` (zero outside ``constraint_mask``). + """ + grad = torch.zeros_like(x0_hat) + mask = constraint_mask.to(x0_hat) + if a is not None: + finite_a = torch.isfinite(a) + a_safe = torch.where(finite_a, a, torch.zeros_like(a)) + # d/dx log sigmoid(s (x - a)) = s sigmoid(-s (x - a)) + term = scale * torch.sigmoid(-scale * (x0_hat - a_safe)) + grad = grad + term * finite_a.to(x0_hat) * mask + if b is not None: + finite_b = torch.isfinite(b) + b_safe = torch.where(finite_b, b, torch.zeros_like(b)) + # d/dx log sigmoid(s (b - x)) = -s sigmoid(-s (b - x)) + term = -scale * torch.sigmoid(-scale * (b_safe - x0_hat)) + grad = grad + term * finite_b.to(x0_hat) * mask + return grad + + +def guided_euler_maruyama( + score_fn: Callable, + sde: BaseSDE, + x_T: torch.Tensor, + condition_mask: torch.Tensor, + constraint_score_fn: Callable, + num_steps: int = 500, + generator: Optional[torch.Generator] = None, +) -> torch.Tensor: + """Reverse Euler-Maruyama with generalized (Tweedie) guidance. + + At each step the clean state is estimated via Tweedie's formula, + ``x0_hat = (x + std(t)^2 score) / mean_scale(t)``, and + ``constraint_score_fn(x0_hat, t)`` is added to the model score, following + ``scoresbibm.methods.guidance.generalized_guidance``. + + Args: + score_fn: Score function ``(t, x) -> score`` on ``(B, T)`` states. + sde: The diffusion SDE. + x_T: Initial state at ``T_max``, shape ``(B, T)``. + condition_mask: Boolean mask of shape ``(T,)``; True marks observed nodes + (hard-conditioned; constrained nodes must not be in this mask). + constraint_score_fn: Callable ``(x0_hat, t) -> (B, T)`` guidance score. + num_steps: Number of integration steps. + generator: Optional torch random generator. + + Returns: + The state at ``T_min``, shape ``(B, T)``. + """ + latent = (~condition_mask).to(x_T) + ts = torch.linspace(sde.T_min, sde.T_max, num_steps, device=x_T.device) + x = x_T + t1 = ts[-1] + for n in range(num_steps - 2, -1, -1): + t0 = ts[n] + dt = t0 - t1 # negative + score = score_fn(t1, x) + x0_hat = (x + sde.transition_std(t1) ** 2 * score) / sde.mean_scale(t1) + score = score + constraint_score_fn(x0_hat, t1) + g = sde.diffusion(t1) + drift = (sde.drift(t1, x) - g**2 * score) * latent + noise = torch.randn(x.shape, device=x.device, generator=generator) + x = x + drift * dt + g * torch.sqrt(torch.abs(dt)) * noise * latent + t1 = t0 + return x + + +def _divergence_exact( + vector_field: Callable, + t: torch.Tensor, + x: torch.Tensor, + latent_idx: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Vector field and its exact divergence over the latent dimensions. + + Args: + vector_field: Callable ``(t, x) -> (B, T)``. + t: Scalar time tensor. + x: State of shape ``(B, T)``. + latent_idx: Long tensor of latent dimension indices. + + Returns: + Tuple of the vector field value ``(B, T)`` and divergence ``(B,)``. + """ + x = x.detach().requires_grad_(True) + with torch.enable_grad(): + v = vector_field(t, x) + div = torch.zeros(x.shape[0], device=x.device) + for dim in latent_idx.tolist(): + grad = torch.autograd.grad(v[:, dim].sum(), x, create_graph=False, retain_graph=True)[0] + div = div + grad[:, dim] + return v.detach(), div.detach() + + +def _divergence_hutchinson( + vector_field: Callable, + t: torch.Tensor, + x: torch.Tensor, + latent_idx: torch.Tensor, + num_probes: int = 8, + generator: Optional[torch.Generator] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Vector field and Hutchinson (Rademacher) divergence estimate. + + Args: + vector_field: Callable ``(t, x) -> (B, T)``. + t: Scalar time tensor. + x: State of shape ``(B, T)``. + latent_idx: Long tensor of latent dimension indices. + num_probes: Number of Rademacher probe vectors. + generator: Optional torch random generator. + + Returns: + Tuple of the vector field value ``(B, T)`` and divergence estimate ``(B,)``. + """ + latent = torch.zeros(x.shape[-1], device=x.device) + latent[latent_idx] = 1.0 + x = x.detach().requires_grad_(True) + with torch.enable_grad(): + v = vector_field(t, x) + div = torch.zeros(x.shape[0], device=x.device) + for _ in range(num_probes): + probe = torch.randint( + 0, 2, x.shape, device=x.device, generator=generator, dtype=x.dtype + ) + probe = (2 * probe - 1) * latent + vjp = torch.autograd.grad((v * probe).sum(), x, retain_graph=True)[0] + div = div + (vjp * probe).sum(dim=-1) + div = div / num_probes + return v.detach(), div.detach() + + +def probability_flow_log_prob( + score_fn: Callable, + sde: BaseSDE, + x0: torch.Tensor, + condition_mask: torch.Tensor, + num_steps: int = 250, + divergence: str = "exact", + hutchinson_probes: int = 8, + generator: Optional[torch.Generator] = None, +) -> torch.Tensor: + """Log probability of the latent entries of ``x0`` via the probability-flow ODE. + + Integrates the forward (noising) probability-flow ODE from ``T_min`` to ``T_max`` + with Euler steps, accumulating the divergence (instantaneous change of variables), + and evaluates the terminal Gaussian at ``T_max``. All in z-scored model space. + + Args: + score_fn: Score function ``(t, x) -> score`` on ``(B, T)`` states. + sde: The diffusion SDE. + x0: Full state at ``T_min``, shape ``(B, T)``, observed entries set to the + conditioning values, latent entries set to the values being evaluated. + condition_mask: Boolean mask of shape ``(T,)``; True marks observed nodes. + num_steps: Number of integration steps. + divergence: ``"exact"`` (autograd trace over latent dims) or ``"hutchinson"``. + hutchinson_probes: Number of probes for the Hutchinson estimator. + generator: Optional torch random generator (Hutchinson only). + + Returns: + Log probabilities of shape ``(B,)``. + + Raises: + ValueError: If ``divergence`` is not a known estimator. + """ + latent_mask = ~condition_mask + latent_idx = torch.nonzero(latent_mask, as_tuple=False).flatten() + latent = latent_mask.to(x0) + + def vector_field(t: torch.Tensor, x: torch.Tensor) -> torch.Tensor: + score = score_fn(t, x) + return (sde.drift(t, x) - 0.5 * sde.diffusion(t) ** 2 * score) * latent + + if divergence == "exact": + div_fn = lambda t, x: _divergence_exact(vector_field, t, x, latent_idx) # noqa: E731 + elif divergence == "hutchinson": + div_fn = lambda t, x: _divergence_hutchinson( # noqa: E731 + vector_field, t, x, latent_idx, num_probes=hutchinson_probes, generator=generator + ) + else: + raise ValueError(f"Unknown divergence estimator '{divergence}'.") + + ts = torch.linspace(sde.T_min, sde.T_max, num_steps, device=x0.device) + x = x0 + log_det = torch.zeros(x0.shape[0], device=x0.device) + for n in range(num_steps - 1): + dt = ts[n + 1] - ts[n] + v, div = div_fn(ts[n], x) + x = x + v * dt + log_det = log_det + div * dt + + mean_end = sde.marginal_mean_end()[latent_mask] + std_end = sde.marginal_std_end()[latent_mask] + y = x[:, latent_mask] + log_q = ( + -0.5 * ((y - mean_end) / std_end) ** 2 + - torch.log(std_end) + - 0.5 * torch.log(torch.tensor(2 * torch.pi, device=x0.device)) + ).sum(dim=-1) + return log_q + log_det + + +__all__ = [ + "euler_maruyama_reverse", + "heun_probability_flow", + "guided_euler_maruyama", + "interval_constraint_score", + "probability_flow_log_prob", +] diff --git a/src/synference/simformer/sde.py b/src/synference/simformer/sde.py new file mode 100644 index 0000000..31534ab --- /dev/null +++ b/src/synference/simformer/sde.py @@ -0,0 +1,259 @@ +"""Diffusion SDEs for the PyTorch Simformer. + +Ports the variance-exploding (VE) and variance-preserving (VP) SDEs from +``probjax.distributions.sde`` together with the loss weighting and score output +scaling from ``scoresbibm.methods.sde``. + +Conventions (forward SDE ``dx = f(t, x) dt + g(t) dW``, ``t in [T_min, T_max]``): + +- ``mean_scale(t)``: transition-kernel mean scale, ``E[x_t | x_0] = mean_scale(t) x_0``. +- ``transition_var(t)``: transition-kernel variance ``Var[x_t | x_0]``. +- ``weight(t)``: denoising score-matching loss weight. +- ``output_scale(t)``: the raw network output is divided by + ``clamp(sqrt(transition_var(t)), scale_min)`` to produce a score estimate. +- ``marginal_mean_end`` / ``marginal_std_end``: per-node data-marginal statistics at + ``T_max`` (using stored per-node data mean/variance), used to draw ``x_T``. +""" + +import math +from typing import Dict + +import numpy as np +import torch + + +class BaseSDE(torch.nn.Module): + """Base class holding shared time constants and per-node data statistics.""" + + def __init__( + self, + x0_mean: np.ndarray, + x0_var: np.ndarray, + T_min: float = 1e-5, + T_max: float = 1.0, + scale_min: float = 1e-3, + ): + """Initialize the SDE. + + Args: + x0_mean: Per-node mean of the (z-scored) training data, shape ``(T,)``. + x0_var: Per-node variance of the training data, shape ``(T,)``. + T_min: Minimum diffusion time (never evaluate below this). + T_max: Maximum diffusion time. + scale_min: Lower clamp on the marginal stddev in ``output_scale``. + """ + super().__init__() + self.T_min = float(T_min) + self.T_max = float(T_max) + self.scale_min = float(scale_min) + self.register_buffer("x0_mean", torch.as_tensor(x0_mean, dtype=torch.float32)) + self.register_buffer("x0_var", torch.as_tensor(x0_var, dtype=torch.float32)) + + @property + def config(self) -> Dict[str, float]: + """Serializable constructor arguments (excluding data statistics).""" + return {"T_min": self.T_min, "T_max": self.T_max, "scale_min": self.scale_min} + + def drift(self, t: torch.Tensor, x: torch.Tensor) -> torch.Tensor: + """Forward drift ``f(t, x)``.""" + raise NotImplementedError + + def diffusion(self, t: torch.Tensor) -> torch.Tensor: + """Forward diffusion ``g(t)`` (scalar function of time).""" + raise NotImplementedError + + def mean_scale(self, t: torch.Tensor) -> torch.Tensor: + """Transition-kernel mean scale ``E[x_t | x_0] / x_0``.""" + raise NotImplementedError + + def transition_var(self, t: torch.Tensor) -> torch.Tensor: + """Transition-kernel variance ``Var[x_t | x_0]``.""" + raise NotImplementedError + + def transition_std(self, t: torch.Tensor) -> torch.Tensor: + """Transition-kernel standard deviation.""" + return torch.sqrt(self.transition_var(t)) + + def weight(self, t: torch.Tensor) -> torch.Tensor: + """Loss weight for denoising score matching.""" + raise NotImplementedError + + def output_scale(self, t: torch.Tensor) -> torch.Tensor: + """Factor the raw network output is divided by to obtain the score.""" + return torch.clamp(self.transition_std(t), min=self.scale_min) + + def marginal_mean_end(self) -> torch.Tensor: + """Per-node data-marginal mean at ``T_max``, shape ``(T,)``.""" + t = torch.tensor(self.T_max, device=self.x0_mean.device) + return self.mean_scale(t) * self.x0_mean + + def marginal_var_end(self) -> torch.Tensor: + """Per-node data-marginal variance at ``T_max``, shape ``(T,)``.""" + raise NotImplementedError + + def marginal_std_end(self) -> torch.Tensor: + """Per-node data-marginal standard deviation at ``T_max``.""" + return torch.sqrt(self.marginal_var_end()) + + +class VESDE(BaseSDE): + """Variance-exploding SDE (paper default: ``sigma_min=1e-4``, ``sigma_max=15``).""" + + name = "vesde" + + def __init__( + self, + x0_mean: np.ndarray, + x0_var: np.ndarray, + sigma_min: float = 1e-4, + sigma_max: float = 15.0, + T_min: float = 1e-5, + T_max: float = 1.0, + scale_min: float = 1e-3, + ): + """Initialize the VE SDE. + + Args: + x0_mean: Per-node data mean, shape ``(T,)``. + x0_var: Per-node data variance, shape ``(T,)``. + sigma_min: Noise scale at ``t=0``. + sigma_max: Noise scale at ``t=1``. + T_min: Minimum diffusion time. + T_max: Maximum diffusion time. + scale_min: Lower clamp on the output-scale stddev. + """ + super().__init__(x0_mean, x0_var, T_min=T_min, T_max=T_max, scale_min=scale_min) + self.sigma_min = float(sigma_min) + self.sigma_max = float(sigma_max) + self._log_ratio = math.log(self.sigma_max / self.sigma_min) + + @property + def config(self) -> Dict[str, float]: + """Serializable constructor arguments (excluding data statistics).""" + return { + "name": self.name, + "sigma_min": self.sigma_min, + "sigma_max": self.sigma_max, + **super().config, + } + + def drift(self, t: torch.Tensor, x: torch.Tensor) -> torch.Tensor: + """Forward drift is zero for the VE SDE.""" + return torch.zeros_like(x) + + def diffusion(self, t: torch.Tensor) -> torch.Tensor: + """``g(t) = sigma_min (sigma_max/sigma_min)^t sqrt(2 log(sigma_max/sigma_min))``.""" + sigma_t = self.sigma_min * (self.sigma_max / self.sigma_min) ** t + return sigma_t * math.sqrt(2 * self._log_ratio) + + def mean_scale(self, t: torch.Tensor) -> torch.Tensor: + """Mean scale is one for the VE SDE.""" + return torch.ones_like(torch.as_tensor(t, dtype=torch.float32)) + + def transition_var(self, t: torch.Tensor) -> torch.Tensor: + """``sigma_min^2 (sigma_max/sigma_min)^(2t)``.""" + return self.sigma_min**2 * (self.sigma_max / self.sigma_min) ** (2 * t) + + def weight(self, t: torch.Tensor) -> torch.Tensor: + """Standard SDE weighting ``g(t)^2``.""" + return self.diffusion(t) ** 2 + + def marginal_var_end(self) -> torch.Tensor: + """Data variance plus transition variance at ``T_max``.""" + t = torch.tensor(self.T_max, device=self.x0_var.device) + return self.x0_var + self.transition_var(t) + + +class VPSDE(BaseSDE): + """Variance-preserving SDE (``beta_min=0.01``, ``beta_max=10``).""" + + name = "vpsde" + + def __init__( + self, + x0_mean: np.ndarray, + x0_var: np.ndarray, + beta_min: float = 0.01, + beta_max: float = 10.0, + T_min: float = 1e-5, + T_max: float = 1.0, + scale_min: float = 0.0, + ): + """Initialize the VP SDE. + + Args: + x0_mean: Per-node data mean, shape ``(T,)``. + x0_var: Per-node data variance, shape ``(T,)``. + beta_min: Noise schedule value at ``t=0``. + beta_max: Noise schedule value at ``t=1``. + T_min: Minimum diffusion time. + T_max: Maximum diffusion time. + scale_min: Lower clamp on the output-scale stddev. + """ + super().__init__(x0_mean, x0_var, T_min=T_min, T_max=T_max, scale_min=scale_min) + self.beta_min = float(beta_min) + self.beta_max = float(beta_max) + + @property + def config(self) -> Dict[str, float]: + """Serializable constructor arguments (excluding data statistics).""" + return { + "name": self.name, + "beta_min": self.beta_min, + "beta_max": self.beta_max, + **super().config, + } + + def _beta(self, t: torch.Tensor) -> torch.Tensor: + return self.beta_min + t * (self.beta_max - self.beta_min) + + def drift(self, t: torch.Tensor, x: torch.Tensor) -> torch.Tensor: + """``f(t, x) = -0.5 beta(t) x``.""" + return -0.5 * self._beta(t) * x + + def diffusion(self, t: torch.Tensor) -> torch.Tensor: + """``g(t) = sqrt(beta(t))``.""" + return torch.sqrt(self._beta(t)) + + def mean_scale(self, t: torch.Tensor) -> torch.Tensor: + """``exp(-0.25 t^2 (beta_max - beta_min) - 0.5 t beta_min)``.""" + t = torch.as_tensor(t, dtype=torch.float32) + return torch.exp(-0.25 * t**2 * (self.beta_max - self.beta_min) - 0.5 * t * self.beta_min) + + def transition_var(self, t: torch.Tensor) -> torch.Tensor: + """``1 - mean_scale(t)^2``.""" + return 1.0 - self.mean_scale(t) ** 2 + + def weight(self, t: torch.Tensor) -> torch.Tensor: + """Likelihood weighting ``clamp(1 - mean_scale(t)^2, min=1e-4)``.""" + return torch.clamp(self.transition_var(t), min=1e-4) + + def marginal_var_end(self) -> torch.Tensor: + """``1 + mean_scale(T)^2 (var_0 - 1)`` per node.""" + t = torch.tensor(self.T_max, device=self.x0_var.device) + return 1.0 + self.mean_scale(t) ** 2 * (self.x0_var - 1.0) + + +SDE_CLASSES = {"vesde": VESDE, "vpsde": VPSDE} + + +def build_sde(name: str, x0_mean: np.ndarray, x0_var: np.ndarray, **kwargs) -> BaseSDE: + """Build an SDE by name. + + Args: + name: Either ``"vesde"`` or ``"vpsde"`` (case-insensitive). + x0_mean: Per-node data mean, shape ``(T,)``. + x0_var: Per-node data variance, shape ``(T,)``. + **kwargs: Extra keyword arguments passed to the SDE constructor + (e.g. ``sigma_min``, ``beta_max``, ``T_min``). + + Returns: + The constructed SDE. + + Raises: + ValueError: If ``name`` is not a known SDE. + """ + key = name.lower() + if key not in SDE_CLASSES: + raise ValueError(f"Unknown SDE '{name}'. Choose from {sorted(SDE_CLASSES)}.") + return SDE_CLASSES[key](x0_mean, x0_var, **kwargs) diff --git a/src/synference/simformer/train.py b/src/synference/simformer/train.py new file mode 100644 index 0000000..2ca95e0 --- /dev/null +++ b/src/synference/simformer/train.py @@ -0,0 +1,418 @@ +"""Training loop for the PyTorch Simformer. + +Ports ``scoresbibm.methods.score_transformer.train_transformer_model``: denoising +score matching over the joint ``[theta, x]`` with per-example condition masks, +adaptive gradient clipping + Adam with a linear-decay schedule, a Monte-Carlo +validation split with early stopping, and optional per-node z-scoring of the data. +""" + +import copy +import time +from typing import Callable, Dict, Optional, Tuple + +import numpy as np +import torch + +from .. import logger +from .masks import build_base_mask, get_condition_mask_fn +from .model import SimformerModel +from .nn import ScoreTransformer +from .sde import BaseSDE, build_sde + +DEFAULT_MODEL_CONFIG: Dict = { + "token_dim": 40, + "condition_token_dim": 10, + "condition_token_init_scale": 0.1, + "time_embedding_dim": 128, + "num_heads": 4, + "num_layers": 6, + "attn_size": 10, + "widening_factor": 3, + "num_hidden_layers": 1, +} + +DEFAULT_SDE_CONFIG: Dict = { + "name": "vesde", + "sigma_min": 1e-4, + "sigma_max": 15.0, + "T_min": 1e-5, + "T_max": 1.0, + "scale_min": 1e-3, +} + +DEFAULT_TRAIN_CONFIG: Dict = { + "learning_rate": 1e-3, + "min_learning_rate": 1e-6, + "z_score_data": True, + "total_number_steps_scaling": 3, + "max_number_steps": 100_000, + "min_number_steps": 5_000, + "training_batch_size": 1000, + "val_every": 50, + "clip_max_norm": 10.0, + "condition_mask_fn": "structured_random", + "validation_fraction": 0.05, + "val_repeat": 5, + "val_error_ratio": 1.1, + "stop_early_count": 5, + "rebalance_loss": False, + "print_every_fraction": 0.1, +} + + +def merge_config(defaults: Dict, overrides: Optional[Dict], config_name: str) -> Dict: + """Merge override values into a default config, rejecting unknown keys. + + Args: + defaults: The default configuration dictionary. + overrides: User overrides (or None). + config_name: Name used in the error message. + + Returns: + A new merged dictionary. + + Raises: + ValueError: If an override key is not present in the defaults. + """ + config = dict(defaults) + if overrides: + unknown = set(overrides) - set(defaults) + if unknown: + raise ValueError( + f"Unknown {config_name} keys: {sorted(unknown)}. Valid keys: {sorted(defaults)}." + ) + config.update(overrides) + return config + + +def mean_std_per_node(data: np.ndarray, node_ids: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + """Per-node-id mean and clipped standard deviation of the training data. + + Args: + data: Training data of shape ``(N, T)``. + node_ids: Integer node ids of shape ``(T,)`` (repeated ids share statistics). + + Returns: + Tuple of per-node mean and stddev arrays, each of shape ``(T,)`` + (stddev clipped below at 1e-2). + """ + node_ids = np.asarray(node_ids).reshape(-1) + mean_per_id = {} + std_per_id = {} + for node in np.unique(node_ids): + values = data[:, node_ids == node] + mean_per_id[node] = float(np.mean(values)) + std_per_id[node] = max(float(np.std(values)), 1e-2) + mean = np.array([mean_per_id[i] for i in node_ids], dtype=np.float32) + std = np.array([std_per_id[i] for i in node_ids], dtype=np.float32) + return mean, std + + +def denoising_score_matching_loss( + net: ScoreTransformer, + sde: BaseSDE, + data: torch.Tensor, + node_ids: torch.Tensor, + condition_mask: torch.Tensor, + edge_mask: Optional[torch.Tensor] = None, + rebalance_loss: bool = False, + generator: Optional[torch.Generator] = None, +) -> torch.Tensor: + """Masked denoising score-matching loss. + + Times are drawn uniformly in ``[T_min, T_max]``; the state is noised with the + transition kernel except at conditioned nodes, which keep their clean values and + contribute zero loss. The target score is ``-eps / std(t)`` and the summed + per-node error is weighted by ``sde.weight(t)``. + + Args: + net: The score network. + sde: The diffusion SDE. + data: Clean training batch of shape ``(B, T)`` (z-scored space). + node_ids: Integer node ids of shape ``(T,)``. + condition_mask: Boolean masks of shape ``(B, T)``. + edge_mask: Optional attention mask (``(T, T)`` or ``(B, T, T)``). + rebalance_loss: If True, divide each example's loss by its number of + unconditioned nodes. + generator: Optional torch random generator. + + Returns: + Scalar loss tensor. + """ + batch = data.shape[0] + device = data.device + t = torch.rand(batch, device=device, generator=generator) * (sde.T_max - sde.T_min) + sde.T_min + eps = torch.randn(data.shape, device=device, generator=generator) + mean_t = sde.mean_scale(t)[:, None] * data + std_t = sde.transition_std(t)[:, None].expand_as(data) + x_t = mean_t + std_t * eps + x_t = torch.where(condition_mask, data, x_t) + + raw = net(t, x_t[..., None], node_ids, condition_mask, edge_mask=edge_mask)[..., 0] + score_pred = raw / sde.output_scale(t)[:, None] + score_target = -eps / std_t + + sq_err = (score_pred - score_target) ** 2 + sq_err = torch.where(condition_mask, torch.zeros_like(sq_err), sq_err) + per_example = sde.weight(t) * sq_err.sum(dim=-1) + if rebalance_loss: + num_latent = (~condition_mask).sum(dim=-1) + per_example = torch.where( + num_latent > 0, per_example / num_latent.clamp(min=1), torch.zeros_like(per_example) + ) + return per_example.mean() + + +def adaptive_grad_clip_( + parameters, + clipping: float = 10.0, + eps: float = 1e-3, +) -> None: + """In-place adaptive gradient clipping (port of ``optax.adaptive_grad_clip``). + + Each parameter's gradient is rescaled so its unit-wise norm does not exceed + ``clipping`` times the unit-wise parameter norm (floored at ``eps``). + + Args: + parameters: Iterable of parameters with gradients. + clipping: Maximum gradient-to-parameter norm ratio. + eps: Floor on the parameter norm. + """ + for param in parameters: + if param.grad is None: + continue + if param.ndim <= 1: + p_norm = param.norm().clamp(min=eps) + g_norm = param.grad.norm().clamp(min=1e-6) + scale = torch.clamp(clipping * p_norm / g_norm, max=1.0) + param.grad.mul_(scale) + else: + dims = tuple(range(1, param.ndim)) + p_norm = param.norm(dim=dims, keepdim=True).clamp(min=eps) + g_norm = param.grad.norm(dim=dims, keepdim=True).clamp(min=1e-6) + scale = torch.clamp(clipping * p_norm / g_norm, max=1.0) + param.grad.mul_(scale) + + +def train_simformer( + theta: np.ndarray, + x: np.ndarray, + model_config: Optional[Dict] = None, + sde_config: Optional[Dict] = None, + train_config: Optional[Dict] = None, + base_mask=None, + meta: Optional[Dict] = None, + device: str = "cpu", + seed: Optional[int] = None, + verbose: bool = True, + progress_callback: Optional[Callable] = None, +) -> Tuple[SimformerModel, Dict]: + """Train a Simformer on parameter/observation pairs. + + Args: + theta: Parameters of shape ``(N, theta_dim)``. + x: Observations of shape ``(N, x_dim)``. + model_config: Overrides for :data:`DEFAULT_MODEL_CONFIG`. + sde_config: Overrides for :data:`DEFAULT_SDE_CONFIG` (must include ``name`` + only to switch SDE type; unknown keys are rejected per type). + train_config: Overrides for :data:`DEFAULT_TRAIN_CONFIG`. + base_mask: Base attention mask — ``"full"``/None for dense, ``"directed"``, + or an explicit boolean ``(T, T)`` array (see + :func:`synference.simformer.masks.build_base_mask`). + meta: Metadata stored on the returned model (parameter/feature names, prior + ranges, etc.). + device: Torch device for training. + seed: Random seed (None for nondeterministic). + verbose: Log training progress. + progress_callback: Optional callable ``(step, train_loss, val_loss)``. + + Returns: + Tuple of the trained :class:`SimformerModel` (on CPU, eval mode) and a stats + dict (loss traces, steps run, early-stopping info, wall time). + """ + model_config = merge_config(DEFAULT_MODEL_CONFIG, model_config, "model_config") + train_config = merge_config(DEFAULT_TRAIN_CONFIG, train_config, "train_config") + sde_config_full = dict(DEFAULT_SDE_CONFIG) + if sde_config: + sde_name = str(sde_config.get("name", sde_config_full["name"])).lower() + if sde_name != sde_config_full["name"]: + # Switching SDE type: start from that type's own defaults. + sde_config_full = {"name": sde_name} + sde_config_full.update(sde_config) + sde_config_full["name"] = sde_name + + theta = np.asarray(theta, dtype=np.float32) + x = np.asarray(x, dtype=np.float32) + if theta.ndim != 2 or x.ndim != 2 or theta.shape[0] != x.shape[0]: + raise ValueError( + f"theta and x must be 2D with matching first dimension, got {theta.shape}, {x.shape}." + ) + theta_dim, x_dim = theta.shape[1], x.shape[1] + num_nodes = theta_dim + x_dim + data = np.hstack([theta, x]) + node_ids = np.arange(num_nodes) + + generator = torch.Generator(device="cpu") + if seed is not None: + generator.manual_seed(int(seed)) + torch.manual_seed(int(seed)) + + # Per-node z-scoring. + if train_config["z_score_data"]: + z_mean, z_std = mean_std_per_node(data, node_ids) + data = (data - z_mean) / z_std + else: + z_mean = z_std = None + + x0_mean = data.mean(axis=0) + x0_var = data.var(axis=0) + sde_kwargs = {k: v for k, v in sde_config_full.items() if k != "name"} + sde = build_sde(sde_config_full["name"], x0_mean, x0_var, **sde_kwargs) + + net = ScoreTransformer(num_nodes=num_nodes, **model_config) + device = torch.device(device) + net.to(device) + sde.to(device) + + # Validation split (first fraction of rows, as in the original implementation). + n_val = max(int(train_config["validation_fraction"] * data.shape[0]), 0) + perm = torch.randperm(data.shape[0], generator=generator).numpy() + data = data[perm] + data_t = torch.as_tensor(data, device=device) + data_val = data_t[:n_val].repeat(train_config["val_repeat"], 1) if n_val > 0 else None + data_train = data_t[n_val:] + node_ids_t = torch.as_tensor(node_ids, device=device) + + total_steps = int( + np.clip( + data.shape[0] * train_config["total_number_steps_scaling"], + train_config["min_number_steps"], + train_config["max_number_steps"], + ) + ) + batch_size = int(train_config["training_batch_size"]) + val_every = max(total_steps // int(train_config["val_every"]), 1) + print_every = max(int(total_steps * train_config["print_every_fraction"]), 1) + + optimizer = torch.optim.Adam( + [p for p in net.parameters() if p.requires_grad], lr=train_config["learning_rate"] + ) + half = total_steps // 2 + lr_ratio = train_config["min_learning_rate"] / train_config["learning_rate"] + + def lr_lambda(step: int) -> float: + if step < half or half == 0: + return 1.0 + frac = min((step - half) / max(total_steps - half, 1), 1.0) + return 1.0 + frac * (lr_ratio - 1.0) + + scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) + + condition_mask_fn = get_condition_mask_fn(train_config["condition_mask_fn"]) + edge_mask = build_base_mask(base_mask, theta_dim, x_dim) + edge_mask_t = edge_mask.to(device) if edge_mask is not None else None + + def compute_loss(batch_data: torch.Tensor) -> torch.Tensor: + condition_mask = condition_mask_fn( + batch_data.shape[0], theta_dim, x_dim, generator=generator + ).to(device) + return denoising_score_matching_loss( + net, + sde, + batch_data, + node_ids_t, + condition_mask, + edge_mask=edge_mask_t, + rebalance_loss=train_config["rebalance_loss"], + ) + + stats: Dict = { + "train_loss": [], + "val_loss": [], + "val_steps": [], + "total_steps_planned": total_steps, + "early_stopped": False, + } + train_loss_ema = None + min_val_loss = np.inf + best_state = None + early_stopping_counter = 0 + start_time = time.time() + + net.train() + for step in range(total_steps): + idx = torch.randint( + 0, data_train.shape[0], (batch_size,), device="cpu", generator=generator + ) + batch = data_train[idx.to(device)] + + optimizer.zero_grad(set_to_none=True) + loss = compute_loss(batch) + loss.backward() + adaptive_grad_clip_( + [p for p in net.parameters() if p.requires_grad], train_config["clip_max_norm"] + ) + optimizer.step() + scheduler.step() + + loss_value = float(loss.detach()) + train_loss_ema = ( + loss_value if train_loss_ema is None else 0.9 * train_loss_ema + 0.1 * loss_value + ) + stats["train_loss"].append(loss_value) + + if data_val is not None and step > 50 and (step % val_every) == 0: + net.eval() + with torch.no_grad(): + val_loss = float(compute_loss(data_val).detach()) + net.train() + stats["val_loss"].append(val_loss) + stats["val_steps"].append(step) + + if val_loss / train_loss_ema > train_config["val_error_ratio"]: + early_stopping_counter += 1 + else: + early_stopping_counter = 0 + if val_loss < min_val_loss: + min_val_loss = val_loss + best_state = copy.deepcopy({k: v.cpu() for k, v in net.state_dict().items()}) + if early_stopping_counter > train_config["stop_early_count"]: + stats["early_stopped"] = True + if verbose: + logger.info(f"Early stopping at step {step} (val loss {val_loss:.4f}).") + break + + if verbose and (step % print_every) == 0: + message = f"Step {step}/{total_steps}: train loss {train_loss_ema:.4f}" + if stats["val_loss"]: + message += f", val loss {stats['val_loss'][-1]:.4f}" + logger.info(message) + if progress_callback is not None: + progress_callback( + step, train_loss_ema, stats["val_loss"][-1] if stats["val_loss"] else None + ) + + if best_state is not None: + net.load_state_dict(best_state) + + stats["steps_run"] = len(stats["train_loss"]) + stats["best_val_loss"] = None if np.isinf(min_val_loss) else float(min_val_loss) + stats["final_train_loss_ema"] = train_loss_ema + stats["training_time_s"] = time.time() - start_time + + net.eval() + net.to("cpu") + sde.to("cpu") + model = SimformerModel( + net=net, + sde=sde, + theta_dim=theta_dim, + x_dim=x_dim, + node_ids=node_ids, + base_mask=None if edge_mask is None else edge_mask.cpu().numpy(), + z_score_mean=z_mean, + z_score_std=z_std, + meta=meta or {}, + sampling_defaults={"num_steps": 500, "method": "sde"}, + ) + return model, stats diff --git a/tests/test_simformer.py b/tests/test_simformer.py new file mode 100644 index 0000000..a4d1e91 --- /dev/null +++ b/tests/test_simformer.py @@ -0,0 +1,672 @@ +"""Tests for the native PyTorch Simformer implementation.""" + +import numpy as np +import pytest +import torch + +from synference.simformer import ( + DEFAULT_MODEL_CONFIG, + VESDE, + VPSDE, + ScoreTransformer, + SimformerModel, + UniformBoxPrior, + build_base_mask, + denoising_score_matching_loss, + euler_maruyama_reverse, + interval_constraint_score, + merge_config, + structured_random_condition_mask, + train_simformer, +) + +THETA_DIM = 2 +X_DIM = 3 +NUM_NODES = THETA_DIM + X_DIM + + +@pytest.fixture +def small_net(): + """A small randomly initialized score network.""" + torch.manual_seed(0) + return ScoreTransformer(num_nodes=NUM_NODES, num_layers=2, token_dim=16, condition_token_dim=4) + + +@pytest.fixture +def random_model(small_net): + """An untrained SimformerModel wrapping the small network.""" + sde = VESDE(np.zeros(NUM_NODES), np.ones(NUM_NODES)) + return SimformerModel( + net=small_net, + sde=sde, + theta_dim=THETA_DIM, + x_dim=X_DIM, + z_score_mean=np.arange(NUM_NODES, dtype=np.float32), + z_score_std=np.linspace(1.0, 2.0, NUM_NODES).astype(np.float32), + meta={"param_names": ["a", "b"], "feature_names": ["f1", "f2", "f3"]}, + ) + + +class TestScoreTransformer: + """Architecture-level tests.""" + + def test_forward_shapes(self, small_net): + """Output has one scalar per token for various mask shapes.""" + batch = 4 + t = torch.rand(batch) + x = torch.randn(batch, NUM_NODES, 1) + node_ids = torch.arange(NUM_NODES) + cond_1d = torch.zeros(NUM_NODES, dtype=torch.bool) + cond_2d = torch.zeros(batch, NUM_NODES, dtype=torch.bool) + edge_2d = torch.ones(NUM_NODES, NUM_NODES, dtype=torch.bool) + edge_3d = torch.ones(batch, NUM_NODES, NUM_NODES, dtype=torch.bool) + + for cond in (cond_1d, cond_2d): + for edge in (None, edge_2d, edge_3d): + out = small_net(t, x, node_ids, cond, edge_mask=edge) + assert out.shape == (batch, NUM_NODES, 1) + + def test_edge_mask_isolation(self, small_net): + """With an identity attention mask, node outputs only depend on own input.""" + t = torch.full((1,), 0.5) + node_ids = torch.arange(NUM_NODES) + cond = torch.zeros(NUM_NODES, dtype=torch.bool) + eye_mask = torch.eye(NUM_NODES, dtype=torch.bool) + + x = torch.randn(1, NUM_NODES, 1) + x_perturbed = x.clone() + x_perturbed[0, 1, 0] += 10.0 + + out = small_net(t, x, node_ids, cond, edge_mask=eye_mask) + out_perturbed = small_net(t, x_perturbed, node_ids, cond, edge_mask=eye_mask) + # Node 1 changes, all others are unaffected. + assert not torch.allclose(out[0, 1], out_perturbed[0, 1]) + unchanged = [i for i in range(NUM_NODES) if i != 1] + assert torch.allclose(out[0, unchanged], out_perturbed[0, unchanged], atol=1e-6) + + # Dense attention propagates the perturbation everywhere. + out_dense = small_net(t, x, node_ids, cond) + out_dense_pert = small_net(t, x_perturbed, node_ids, cond) + assert not torch.allclose(out_dense[0, 0], out_dense_pert[0, 0]) + + def test_frozen_embeddings(self, small_net): + """Node-id embedding and Fourier projection stay fixed; condition token trains.""" + node_embedding_before = small_net.tokenizer.node_embedding.weight.clone() + fourier_before = small_net.time_embedding.B.clone() + condition_before = small_net.condition_token.clone() + + optimizer = torch.optim.Adam( + [p for p in small_net.parameters() if p.requires_grad], lr=1e-2 + ) + t = torch.rand(8) + x = torch.randn(8, NUM_NODES, 1) + # Condition on some nodes so the (mask-gated) condition token receives gradient. + cond = torch.tensor([False, False, True, True, True]) + loss = small_net(t, x, torch.arange(NUM_NODES), cond).pow(2).mean() + 1.0 + loss.backward() + optimizer.step() + + assert torch.equal(small_net.tokenizer.node_embedding.weight, node_embedding_before) + assert torch.equal(small_net.time_embedding.B, fourier_before) + assert not torch.equal(small_net.condition_token, condition_before) + + +class TestSDE: + """Closed-form checks of the SDE quantities.""" + + def test_vesde_closed_forms(self): + """VE transition variance, diffusion, and weight match the formulas.""" + sde = VESDE(np.zeros(2), np.ones(2), sigma_min=1e-4, sigma_max=15.0) + t = torch.tensor([1e-5, 0.5, 1.0]) + expected_var = 1e-8 * (15.0 / 1e-4) ** (2 * t.numpy()) + assert np.allclose(sde.transition_var(t).numpy(), expected_var, rtol=1e-5) + expected_g = 1e-4 * (15.0 / 1e-4) ** t.numpy() * np.sqrt(2 * np.log(15.0 / 1e-4)) + assert np.allclose(sde.diffusion(t).numpy(), expected_g, rtol=1e-5) + assert np.allclose(sde.weight(t).numpy(), expected_g**2, rtol=1e-5) + assert np.allclose(sde.mean_scale(t).numpy(), 1.0) + # Data-marginal std at T_max: sqrt(var_0 + sigma_max^2). + assert np.allclose(sde.marginal_std_end().numpy(), np.sqrt(1 + 225.0), rtol=1e-5) + + def test_vpsde_closed_forms(self): + """VP mean scale, variance, drift, and weight match the formulas.""" + sde = VPSDE(np.zeros(2), np.ones(2), beta_min=0.01, beta_max=10.0) + t = torch.tensor([1e-5, 0.5, 1.0]) + phi = np.exp(-0.25 * t.numpy() ** 2 * (10.0 - 0.01) - 0.5 * t.numpy() * 0.01) + assert np.allclose(sde.mean_scale(t).numpy(), phi, rtol=1e-5) + assert np.allclose(sde.transition_var(t).numpy(), 1 - phi**2, rtol=1e-4) + assert np.allclose(sde.weight(t).numpy(), np.clip(1 - phi**2, 1e-4, None), rtol=1e-4) + x = torch.ones(2) + beta_half = 0.01 + 0.5 * (10.0 - 0.01) + assert np.allclose(sde.drift(torch.tensor(0.5), x).numpy(), -0.5 * beta_half, rtol=1e-5) + # Unit data variance keeps the marginal variance at one. + assert np.allclose(sde.marginal_var_end().numpy(), 1.0, rtol=1e-5) + + +class TestMasks: + """Condition- and base-mask behaviour.""" + + def test_structured_random_masks(self): + """Structured masks have the right shape, never all-True, and hit all types.""" + gen = torch.Generator().manual_seed(0) + masks = structured_random_condition_mask(2000, THETA_DIM, X_DIM, generator=gen) + assert masks.shape == (2000, NUM_NODES) + assert not masks.all(dim=-1).any() + + posterior_row = torch.tensor([False] * THETA_DIM + [True] * X_DIM) + likelihood_row = torch.tensor([True] * THETA_DIM + [False] * X_DIM) + n_joint = (~masks.any(dim=-1)).sum() + n_posterior = (masks == posterior_row).all(dim=-1).sum() + n_likelihood = (masks == likelihood_row).all(dim=-1).sum() + assert n_joint > 100 + assert n_posterior > 100 + assert n_likelihood > 100 + + def test_build_base_mask(self): + """'full' is dense (None), 'directed' has the expected block structure.""" + assert build_base_mask("full", THETA_DIM, X_DIM) is None + assert build_base_mask(None, THETA_DIM, X_DIM) is None + + mask = build_base_mask("directed", THETA_DIM, X_DIM) + assert mask.shape == (NUM_NODES, NUM_NODES) + # Parameters attend only to themselves. + assert torch.equal(mask[:THETA_DIM, :THETA_DIM], torch.eye(THETA_DIM, dtype=torch.bool)) + # Parameters do not attend to data. + assert not mask[:THETA_DIM, THETA_DIM:].any() + # Data attend to all parameters and causally within data. + assert mask[THETA_DIM:, :THETA_DIM].all() + assert torch.equal( + mask[THETA_DIM:, THETA_DIM:], + torch.tril(torch.ones(X_DIM, X_DIM, dtype=torch.bool)), + ) + + custom = np.eye(NUM_NODES, dtype=bool) + assert torch.equal( + build_base_mask(custom, THETA_DIM, X_DIM), torch.eye(NUM_NODES, dtype=torch.bool) + ) + with pytest.raises(ValueError): + build_base_mask(np.eye(3, dtype=bool), THETA_DIM, X_DIM) + with pytest.raises(ValueError): + build_base_mask("banana", THETA_DIM, X_DIM) + + def test_merge_config_rejects_unknown_keys(self): + """Config overrides with typos raise instead of being silently ignored.""" + with pytest.raises(ValueError, match="d_model"): + merge_config(DEFAULT_MODEL_CONFIG, {"d_model": 128}, "model_config") + + +class TestLoss: + """Denoising score-matching loss semantics.""" + + def test_conditioned_nodes_contribute_zero_loss(self, small_net): + """Corrupting the prediction at conditioned nodes leaves the loss unchanged.""" + sde = VESDE(np.zeros(NUM_NODES), np.ones(NUM_NODES)) + data = torch.randn(16, NUM_NODES) + node_ids = torch.arange(NUM_NODES) + cond = torch.zeros(16, NUM_NODES, dtype=torch.bool) + cond[:, THETA_DIM:] = True # condition on all x nodes + + class CorruptedNet(torch.nn.Module): + """Adds a huge offset to the output at conditioned nodes.""" + + def __init__(self, base): + super().__init__() + self.base = base + + def forward(self, t, x, node_ids, condition_mask, edge_mask=None): + out = self.base(t, x, node_ids, condition_mask, edge_mask=edge_mask) + return out + 1e6 * condition_mask.reshape(out.shape[0], -1, 1).float() + + loss_clean = denoising_score_matching_loss( + small_net, sde, data, node_ids, cond, generator=torch.Generator().manual_seed(3) + ) + loss_corrupted = denoising_score_matching_loss( + CorruptedNet(small_net), + sde, + data, + node_ids, + cond, + generator=torch.Generator().manual_seed(3), + ) + assert torch.allclose(loss_clean, loss_corrupted) + + +class TestSampling: + """Sampling semantics on an untrained model.""" + + def test_conditioned_values_clamped_throughout(self, random_model): + """The integrator never moves observed entries away from x_o.""" + cond = torch.tensor([False, False, True, True, True]) + x_o_z = torch.tensor([0.3, -0.2, 1.5]) + x_T = torch.randn(8, NUM_NODES) + x_T[:, cond] = x_o_z + + seen_states = [] + + def recording_score_fn(t, x): + seen_states.append(x.clone()) + return random_model.score(t, x, cond) + + final = euler_maruyama_reverse( + recording_score_fn, + random_model.sde, + x_T, + cond, + num_steps=10, + generator=torch.Generator().manual_seed(0), + ) + for state in seen_states + [final]: + assert torch.allclose(state[:, cond], x_o_z.expand(8, -1)) + + def test_sample_shapes_and_units(self, random_model): + """sample/sample_batched return original-unit latents of the right shape.""" + cond = np.array([False, False, True, True, True]) + gen = torch.Generator().manual_seed(0) + samples = random_model.sample( + 16, x_o=[1.0, 2.0, 3.0], condition_mask=cond, num_steps=8, generator=gen + ) + assert samples.shape == (16, 2) + assert np.isfinite(samples).all() + + batch = random_model.sample_batched( + 8, np.random.randn(5, 3), condition_mask=cond, num_steps=8, batch_size=2 + ) + assert batch.shape == (5, 8, 2) + assert np.isfinite(batch).all() + + def test_all_conditioned_mask_rejected(self, random_model): + """A mask conditioning on every node is invalid.""" + with pytest.raises(ValueError): + random_model.sample( + 4, x_o=np.zeros(NUM_NODES), condition_mask=np.ones(NUM_NODES, dtype=bool) + ) + + +class TestGuidance: + """Interval-guidance components.""" + + def test_interval_constraint_score_matches_autograd(self): + """The analytic box-constraint gradient matches autograd of log_step_fn.""" + torch.manual_seed(0) + x = torch.randn(6, NUM_NODES, requires_grad=True) + constraint_mask = torch.tensor([True, False, True, False, False]) + a = torch.tensor([-0.5, torch.nan, 0.1, torch.nan, torch.nan]) + b = torch.tensor([0.8, torch.nan, torch.inf, torch.nan, torch.nan]) + scale = torch.tensor(3.0) + + mask_f = constraint_mask.float() + finite_a, finite_b = torch.isfinite(a), torch.isfinite(b) + log_p = ( + torch.nn.functional.logsigmoid(scale * (x - torch.nan_to_num(a))) * (mask_f * finite_a) + ).sum() + ( + torch.nn.functional.logsigmoid(scale * (torch.nan_to_num(b) - x)) * (mask_f * finite_b) + ).sum() + (expected,) = torch.autograd.grad(log_p, x) + + actual = interval_constraint_score(x.detach(), scale, constraint_mask, a=a, b=b) + assert torch.allclose(actual, expected, atol=1e-5) + + def test_sample_intervals_shapes(self, random_model): + """Constrained nodes are returned among the latents.""" + cond = np.array([False, False, True, True, True]) + constraint = np.array([True, False, False, False, False]) + samples = random_model.sample_intervals( + 8, + x_o=[1.0, 2.0, 3.0], + condition_mask=cond, + constraint_mask=constraint, + a=[-1.0], + b=[1.0], + num_steps=8, + scale_bias=1e-2, + generator=torch.Generator().manual_seed(0), + ) + assert samples.shape == (8, 2) + assert np.isfinite(samples).all() + + +class TestLogProb: + """Probability-flow log-probability.""" + + def test_exact_and_hutchinson_agree(self, random_model): + """The two divergence estimators agree within Monte-Carlo tolerance.""" + cond = np.array([False, False, True, True, True]) + theta = np.array([[0.1, 0.2], [0.5, -0.3]], dtype=np.float32) + x_o = [1.0, 2.0, 3.0] + lp_exact = random_model.log_prob(theta, x_o, cond, num_steps=25) + lp_hutch = random_model.log_prob( + theta, + x_o, + cond, + num_steps=25, + divergence="hutchinson", + hutchinson_probes=64, + generator=torch.Generator().manual_seed(0), + ) + assert np.all(np.isfinite(lp_exact)) + assert np.allclose(lp_exact, lp_hutch, rtol=0.15, atol=0.3) + + +class TestSerialization: + """Save/load round-trips.""" + + def test_roundtrip_scores_and_samples(self, random_model, tmp_path): + """A reloaded model reproduces scores exactly and samples with a shared seed.""" + path = str(tmp_path / "model.pt") + random_model.save(path) + reloaded = SimformerModel.load(path) + + cond = torch.tensor([False, False, True, True, True]) + t = torch.tensor(0.5) + x = torch.randn(4, NUM_NODES) + with torch.no_grad(): + s1 = random_model.score(t, x, cond) + s2 = reloaded.score(t, x, cond) + assert torch.equal(s1, s2) + + samples1 = random_model.sample( + 8, + x_o=[1.0, 2.0, 3.0], + condition_mask=cond.numpy(), + num_steps=8, + generator=torch.Generator().manual_seed(7), + ) + samples2 = reloaded.sample( + 8, + x_o=[1.0, 2.0, 3.0], + condition_mask=cond.numpy(), + num_steps=8, + generator=torch.Generator().manual_seed(7), + ) + assert np.array_equal(samples1, samples2) + assert reloaded.meta["param_names"] == ["a", "b"] + assert np.array_equal(reloaded.z_score_mean, random_model.z_score_mean) + + +def two_moons_simulator(theta: np.ndarray, rng: np.random.Generator) -> np.ndarray: + """Two-moons simulator (Lueckmann et al. 2021 benchmark parametrization). + + Args: + theta: Parameters of shape ``(N, 2)`` in ``[-1, 1]^2``. + rng: Numpy random generator. + + Returns: + Observations of shape ``(N, 2)``. + """ + n = theta.shape[0] + alpha = rng.uniform(-np.pi / 2, np.pi / 2, n) + r = rng.normal(0.1, 0.01, n) + p = np.stack([r * np.cos(alpha) + 0.25, r * np.sin(alpha)], axis=1) + shift = np.stack( + [-np.abs(theta[:, 0] + theta[:, 1]), (-theta[:, 0] + theta[:, 1])], axis=1 + ) / np.sqrt(2) + return p + shift + + +@pytest.fixture(scope="module") +def two_moons_model(): + """Train a small Simformer on the two-moons task (shared across tests).""" + rng = np.random.default_rng(42) + n_train = 10_000 + theta = rng.uniform(-1, 1, (n_train, 2)) + x = two_moons_simulator(theta, rng) + + model, stats = train_simformer( + theta, + x, + model_config={"num_layers": 3}, + train_config={ + "min_number_steps": 3000, + "max_number_steps": 3000, + "training_batch_size": 512, + }, + seed=0, + verbose=False, + ) + test_theta = rng.uniform(-1, 1, (100, 2)) + test_x = two_moons_simulator(test_theta, rng) + return model, stats, (theta, x), (test_theta, test_x) + + +class TestTwoMoons: + """End-to-end validation on the two-moons benchmark.""" + + POSTERIOR_MASK = np.array([False, False, True, True]) + LIKELIHOOD_MASK = np.array([True, True, False, False]) + JOINT_MASK = np.array([False, False, False, False]) + + def test_training_ran(self, two_moons_model): + """Training completes and the loss decreases.""" + _, stats, _, _ = two_moons_model + assert stats["steps_run"] > 0 + assert stats["final_train_loss_ema"] < stats["train_loss"][0] + + def test_posterior_is_bimodal(self, two_moons_model): + """The posterior at the origin covers both signs of theta_1 + theta_2.""" + model, _, _, _ = two_moons_model + gen = torch.Generator().manual_seed(0) + samples = model.sample( + 2000, + x_o=[0.0, 0.0], + condition_mask=self.POSTERIOR_MASK, + num_steps=100, + generator=gen, + ) + s = samples[:, 0] + samples[:, 1] + frac_positive = (s > 0).mean() + assert 0.1 < frac_positive < 0.9, f"Posterior lost a mode: {frac_positive:.2f} positive" + + def test_posterior_calibration_tarp(self, two_moons_model): + """TARP expected coverage stays close to nominal at mid-credibility.""" + import tarp + + model, _, _, (test_theta, test_x) = two_moons_model + gen = torch.Generator().manual_seed(1) + samples = model.sample_batched( + 250, + test_x, + condition_mask=self.POSTERIOR_MASK, + num_steps=100, + batch_size=25, + generator=gen, + ) + # tarp expects (num_samples, num_sims, num_dims). + ecp, alpha = tarp.get_tarp_coverage(samples.transpose(1, 0, 2), test_theta, norm=True) + mid = np.argmin(np.abs(alpha - 0.5)) + assert abs(ecp[mid] - 0.5) < 0.15 + + def test_arbitrary_conditionals(self, two_moons_model): + """Likelihood and joint conditionals reproduce simulator statistics.""" + model, _, (theta_train, x_train), _ = two_moons_model + gen = torch.Generator().manual_seed(2) + rng = np.random.default_rng(3) + + theta_o = np.array([0.4, -0.3]) + x_model = model.sample( + 2000, + x_o=theta_o, + condition_mask=self.LIKELIHOOD_MASK, + num_steps=100, + generator=gen, + ) + x_true = two_moons_simulator(np.tile(theta_o, (2000, 1)), rng) + assert np.allclose(x_model.mean(axis=0), x_true.mean(axis=0), atol=0.1) + assert np.allclose(x_model.std(axis=0), x_true.std(axis=0), atol=0.1) + + joint = model.sample( + 2000, x_o=[], condition_mask=self.JOINT_MASK, num_steps=100, generator=gen + ) + train_data = np.hstack([theta_train, x_train]) + assert np.allclose(joint.mean(axis=0), train_data.mean(axis=0), atol=0.15) + assert np.allclose(joint.std(axis=0), train_data.std(axis=0), atol=0.15) + + def test_interval_guidance_respects_box(self, two_moons_model): + """Interval-constrained samples stay inside the box while conditioning on x.""" + model, _, _, _ = two_moons_model + gen = torch.Generator().manual_seed(4) + constraint = np.array([True, False, False, False]) + samples = model.sample_intervals( + 1000, + x_o=[0.0, 0.0], + condition_mask=self.POSTERIOR_MASK, + constraint_mask=constraint, + a=[0.0], + b=[0.6], + scale_bias=1e-2, + num_steps=100, + generator=gen, + ) + inside = ((samples[:, 0] >= 0.0) & (samples[:, 0] <= 0.6)).mean() + assert inside > 0.95 + + def test_log_prob_beats_prior(self, two_moons_model): + """Mean posterior log-prob of the true parameters beats the prior density.""" + model, _, _, (test_theta, test_x) = two_moons_model + log_probs = [] + for i in range(10): + lp = model.log_prob(test_theta[i], test_x[i], self.POSTERIOR_MASK, num_steps=100) + log_probs.append(float(lp)) + prior_log_prob = -np.log(4.0) # U(-1,1)^2 + assert np.isfinite(log_probs).all() + assert np.mean(log_probs) > prior_log_prob + + +@pytest.fixture(scope="module") +def trained_fitter(tmp_path_factory): + """Train a small Simformer_Fitter on the galaxy test library (shared).""" + from synference import Simformer_Fitter + from synference.utils import test_data_dir + + out_dir = tmp_path_factory.mktemp("simformer_models") + fitter = Simformer_Fitter.init_from_hdf5( + model_name="test_simformer", + hdf5_path=f"{test_data_dir}/sbi_test_library.hdf5", + ) + fitter.create_feature_array_from_raw_photometry() + + model, stats = fitter.run_single_sbi( + name_append="pytest", + out_dir=str(out_dir), + random_seed=0, + load_existing_model=False, + model_config_dict_overrides={"num_layers": 2}, + train_config_dict_overrides={ + "min_number_steps": 600, + "max_number_steps": 600, + "training_batch_size": 128, + }, + num_posterior_draws_per_sample=100, + evaluate_model=True, + verbose=False, + ) + return fitter, model, stats, out_dir + + +class TestSimformerFitterLibrary: + """End-to-end Simformer_Fitter run on the galaxy test library.""" + + def test_artifacts_written(self, trained_fitter): + """Training writes the posterior, params, and metrics files plus plots.""" + fitter, model, stats, out_dir = trained_fitter + model_dir = out_dir / "test_simformer" + assert (model_dir / "test_simformer_pytest_posterior.pkl").exists() + assert (model_dir / "test_simformer_pytest_params.pkl").exists() + assert (model_dir / "test_simformer_pytest_metrics.json").exists() + plots = list((model_dir / "plots" / "pytest").glob("*")) + assert len(plots) >= 2 + assert stats["steps_run"] == 600 + + def test_sample_posterior_contract(self, trained_fitter): + """sample_posterior returns finite samples with the documented shapes.""" + fitter, _, _, _ = trained_fitter + n_theta = len(fitter.fitted_parameter_names) + + multi = fitter.sample_posterior(fitter._X_test[:3], num_samples=50, num_steps=50) + assert multi.shape == (3, 50, n_theta) + assert np.isfinite(multi).all() + + single = fitter.sample_posterior(fitter._X_test[0], num_samples=50, num_steps=50) + assert single.shape == (50, n_theta) + + def test_log_prob_paired_contract(self, trained_fitter): + """log_prob(X, y) pairs rows and returns one finite value per observation.""" + fitter, _, _, _ = trained_fitter + lp = fitter.log_prob(fitter._X_test[:3], fitter._y_test[:3], num_steps=25) + assert lp.shape == (3,) + assert np.isfinite(lp).all() + + def test_load_saved_model_roundtrip(self, trained_fitter): + """A model reloaded via load_saved_model reproduces seeded samples.""" + from synference import Simformer_Fitter + from synference.utils import test_data_dir + + fitter, _, _, out_dir = trained_fitter + reloaded = Simformer_Fitter.load_saved_model( + model_name="test_simformer_pytest", + library_path=f"{test_data_dir}/sbi_test_library.hdf5", + model_file=str(out_dir / "test_simformer"), + ) + assert reloaded.posteriors is not None + assert list(reloaded.posteriors.meta["param_names"]) == list(fitter.fitted_parameter_names) + x_obs = fitter._X_test[:1] + s1 = fitter.sample_posterior(x_obs, num_samples=20, num_steps=25, rng_seed=11) + s2 = reloaded.sample_posterior(x_obs, num_samples=20, num_steps=25, rng_seed=11) + assert np.allclose(s1, s2) + + def test_load_in_fresh_process(self, trained_fitter): + """A saved model loads and samples in a brand-new Python interpreter.""" + import subprocess + import sys + + _, _, _, out_dir = trained_fitter + model_path = out_dir / "test_simformer" / "test_simformer_pytest_posterior.pkl" + script = ( + "import numpy as np\n" + "from synference.simformer import SimformerModel\n" + f"model = SimformerModel.load({str(model_path)!r})\n" + "mask = np.array([False] * model.theta_dim + [True] * model.x_dim)\n" + "x_o = np.zeros(model.x_dim) + 25.0\n" + "s = model.sample(10, x_o=x_o, condition_mask=mask, num_steps=10)\n" + "assert s.shape == (10, model.theta_dim)\n" + "assert np.isfinite(s).all()\n" + "print('FRESH_LOAD_OK')\n" + ) + result = subprocess.run( + [sys.executable, "-c", script], capture_output=True, text=True, timeout=300 + ) + assert "FRESH_LOAD_OK" in result.stdout, result.stderr + + def test_fit_catalogue_smoke(self, trained_fitter): + """fit_catalogue produces quantile columns for a small catalogue.""" + from astropy.table import Table + + fitter, _, _, _ = trained_fitter + table = Table() + for i, name in enumerate(fitter.feature_names): + table[name] = np.asarray(fitter._X_test[:3, i]) + + result = fitter.fit_catalogue( + table, + columns_to_feature_names={name: name for name in fitter.feature_names}, + flux_units="AB", + num_samples=50, + check_out_of_distribution=False, + ) + for param in fitter.simple_fitted_parameter_names: + assert f"{param}_50" in result.colnames + assert np.isfinite(result[f"{param}_50"]).all() + + +class TestUniformBoxPrior: + """Box prior behaviour.""" + + def test_sample_and_log_prob(self): + """Samples stay inside the box; log_prob is -log(volume) inside, -inf outside.""" + prior = UniformBoxPrior({"a": (0.0, 2.0), "b": (-1.0, 1.0)}, ["a", "b"]) + samples = prior.sample(100, generator=torch.Generator().manual_seed(0)) + assert samples.shape == (100, 2) + assert (samples[:, 0] >= 0).all() and (samples[:, 0] <= 2).all() + assert (samples[:, 1] >= -1).all() and (samples[:, 1] <= 1).all() + + lp = prior.log_prob(torch.tensor([[1.0, 0.0], [3.0, 0.0]])) + assert torch.isclose(lp[0], torch.tensor(-np.log(4.0).astype(np.float32))) + assert torch.isinf(lp[1]) and lp[1] < 0