From e569d41d0689b65089611e49442ae155a530f0a1 Mon Sep 17 00:00:00 2001 From: EiffL Date: Sun, 8 Feb 2026 11:14:50 +0100 Subject: [PATCH 1/5] updating to support several inference methods --- configs/validation/level0_base.yaml | 37 +++-- docs/api/config.md | 4 +- docs/api/inference.md | 7 +- docs/api/validation/cli.md | 2 +- docs/api/validation/extraction.md | 6 +- docs/architecture.md | 10 +- docs/configuration.md | 118 +++++++++++--- docs/getting-started.md | 20 +-- docs/validation/batched.md | 8 +- docs/validation/index.md | 10 +- docs/validation/level0.md | 31 ++-- examples/level0_batched_inference.ipynb | 191 +++++++++++++++++++++++ shine/config.py | 66 +++++++- shine/inference.py | 129 ++++++++++++--- shine/main.py | 2 +- shine/validation/cli.py | 4 +- shine/validation/extraction.py | 74 +++++++-- shine/validation/plots.py | 132 +++++++++++++++- tests/test_config.py | 120 ++++++++++++-- tests/test_validation/test_extraction.py | 129 ++++++++++++++- 20 files changed, 959 insertions(+), 141 deletions(-) create mode 100644 examples/level0_batched_inference.ipynb diff --git a/configs/validation/level0_base.yaml b/configs/validation/level0_base.yaml index 42bfd20..3450035 100644 --- a/configs/validation/level0_base.yaml +++ b/configs/validation/level0_base.yaml @@ -1,33 +1,40 @@ -# Level 0 base SHINE config — low-noise sanity check -# Noise sigma is small so the posterior is tight but not degenerate. +# Level 0 base SHINE config — noiseless sanity check +# Parameters matched to esheldon/ngmix metacal example: +# https://github.com/esheldon/ngmix/blob/master/examples/metacal/metacal.py +# +# Differences from metacal example: +# - PSF ellipticity (g1=0.02, g2=-0.01 on PSF) not supported yet; using round PSF +# - Galaxy position fixed at center (metacal uses random subpixel offsets) +# - Flux=1000 in SHINE prior; data generated with flux=1 matching metacal default image: - pixel_scale: 0.1 # arcsec/pixel + pixel_scale: 0.263 # arcsec/pixel (LSST-like, matches metacal) size_x: 48 size_y: 48 n_objects: 1 fft_size: 128 noise: type: Gaussian - sigma: 0.1 # Low noise for Level 0 (avoids degenerate posteriors) + sigma: 1.0e-6 # Effectively noiseless (matches metacal default) psf: - type: Gaussian - sigma: 0.1 # arcsec + type: Moffat + sigma: 0.9 # FWHM in arcsec (matches metacal psf_fwhm=0.9) + beta: 2.5 # Moffat beta (matches metacal) gal: type: Exponential - flux: 1000.0 - half_light_radius: 0.5 # arcsec + flux: 1.0 + half_light_radius: 0.5 # arcsec (matches metacal gal_hlr=0.5) shear: type: G1G2 g1: type: Normal - mean: 0.02 + mean: 0.0 # Prior center; truth (0.01) set via bias run config sigma: 0.05 g2: type: Normal - mean: -0.01 + mean: 0.0 # Prior center; truth (0.00) set via bias run config sigma: 0.05 position: type: Uniform @@ -37,12 +44,8 @@ gal: y_max: 24.5 inference: - warmup: 500 - samples: 1000 - chains: 2 - dense_mass: false - rng_seed: 42 - map_init: - enabled: true + method: map # MAP is sufficient for noiseless Level 0 + map_config: num_steps: 1000 learning_rate: 0.01 + rng_seed: 42 diff --git a/docs/api/config.md b/docs/api/config.md index 8674ff4..a65e69f 100644 --- a/docs/api/config.md +++ b/docs/api/config.md @@ -4,6 +4,8 @@ Configuration handling with Pydantic models. Parses YAML configuration files and validates all parameters. Distribution parameters (Normal, LogNormal, Uniform) are automatically treated as latent -variables for Bayesian inference. +variables for Bayesian inference. The `InferenceConfig` supports three +inference methods (NUTS, MAP, VI) with method-specific config blocks +(`NUTSConfig`, `MAPConfig`, `VIConfig`). ::: shine.config diff --git a/docs/api/inference.md b/docs/api/inference.md index 7a8434d..fd24009 100644 --- a/docs/api/inference.md +++ b/docs/api/inference.md @@ -1,8 +1,9 @@ # shine.inference -Bayesian inference engine with optional MAP initialization. +Bayesian inference engine supporting NUTS/MCMC, MAP, and Variational Inference. -Wraps NumPyro's NUTS sampler with support for MAP pre-initialization -using Adam optimization to improve MCMC convergence. +Dispatches on `InferenceConfig.method` to run one of three inference paths. +All methods return ArviZ `InferenceData` so the downstream pipeline works +uniformly. ::: shine.inference diff --git a/docs/api/validation/cli.md b/docs/api/validation/cli.md index e904717..af8d501 100644 --- a/docs/api/validation/cli.md +++ b/docs/api/validation/cli.md @@ -2,7 +2,7 @@ CLI entry points for the three-stage bias measurement pipeline. -- **Stage 1** (`shine-bias-run`): Generate data + run MCMC +- **Stage 1** (`shine-bias-run`): Generate data + run inference (NUTS, MAP, or VI) - **Stage 2** (`shine-bias-extract`): Load posteriors, extract diagnostics, write CSV - **Stage 3** (`shine-bias-stats`): Read CSV, compute bias, check acceptance, plot diff --git a/docs/api/validation/extraction.md b/docs/api/validation/extraction.md index 1f08b1c..e7e3436 100644 --- a/docs/api/validation/extraction.md +++ b/docs/api/validation/extraction.md @@ -2,7 +2,9 @@ Extract structured results from ArviZ InferenceData. -Provides convergence diagnostics (R-hat, ESS, divergences, BFMI) and -shear summary statistics from posterior samples. +Provides method-aware convergence diagnostics (R-hat, ESS, divergences, BFMI) +and shear summary statistics from posterior samples. Automatically adapts +to the inference method (NUTS, MAP, or VI) via the `inference_method` +attribute on the posterior. ::: shine.validation.extraction diff --git a/docs/architecture.md b/docs/architecture.md index 20a2eee..e566c10 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -60,12 +60,14 @@ The forward model is fully differentiable, enabling gradient-based samplers. ### Inference Engine (`shine.inference`) -Runs Bayesian inference with optional MAP initialization: +Runs Bayesian inference using one of three methods, configured via YAML: -1. **MAP phase** (optional): Adam optimizer finds a good starting point -2. **MCMC phase**: NUTS sampler explores the posterior +- **NUTS**: MCMC sampling with the No-U-Turn Sampler, optionally preceded by MAP initialization. +- **MAP**: Maximum a posteriori point estimation (fast, no posterior samples). +- **VI**: Variational Inference with an AutoNormal guide (approximate posterior). -Results are returned as ArviZ `InferenceData` objects with full diagnostics. +All three methods return ArviZ `InferenceData` objects, so the downstream +pipeline (extraction, diagnostics, plots) works uniformly. ### Data Loading (`shine.data`) diff --git a/docs/configuration.md b/docs/configuration.md index 5e59581..8637a15 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -155,31 +155,95 @@ gal: ## Inference Section -Controls the MCMC sampler and optional MAP initialization. +Controls the inference method and its settings. SHINE supports three methods: + +- **`nuts`** (default): NUTS/MCMC sampling, optionally preceded by MAP initialization. +- **`map`**: MAP point estimation only (fast, no posterior samples). +- **`vi`**: Variational Inference with an AutoNormal guide (approximate posterior). + +All three methods return ArviZ `InferenceData`, so the downstream pipeline +(extraction, diagnostics, plots) works uniformly. + +### Method selection ```yaml inference: - warmup: 500 # NUTS warmup steps - samples: 1000 # posterior samples - chains: 2 # number of parallel chains - dense_mass: false # use dense mass matrix - rng_seed: 42 # reproducibility seed - map_init: - enabled: true # run MAP before MCMC - num_steps: 1000 # Adam optimization steps - learning_rate: 0.01 + method: nuts # "nuts", "map", or "vi" + rng_seed: 42 # JAX PRNG seed (shared across all methods) ``` +Each method reads its own config block; the others are ignored. When a +method's config block is omitted, defaults are used. + +### NUTS config + +```yaml +inference: + method: nuts + nuts_config: + warmup: 500 # NUTS warmup steps + samples: 1000 # posterior samples per chain + chains: 2 # number of parallel chains + dense_mass: false # use dense mass matrix + map_init: # optional MAP pre-initialization + enabled: true + num_steps: 1000 + learning_rate: 0.01 + rng_seed: 42 +``` + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `nuts_config.warmup` | int > 0 | `500` | NUTS warmup iterations | +| `nuts_config.samples` | int > 0 | `1000` | Number of posterior samples | +| `nuts_config.chains` | int > 0 | `1` | Number of MCMC chains | +| `nuts_config.dense_mass` | bool | `false` | Dense mass matrix for correlated parameters | +| `nuts_config.map_init.enabled` | bool | `false` | Enable MAP pre-initialization | +| `nuts_config.map_init.num_steps` | int > 0 | `1000` | Optimization steps for MAP | +| `nuts_config.map_init.learning_rate` | float > 0 | `0.01` | Adam learning rate for MAP | + +### MAP config + +```yaml +inference: + method: map + map_config: + num_steps: 2000 + learning_rate: 0.005 + rng_seed: 42 +``` + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `map_config.num_steps` | int > 0 | `1000` | Adam optimization steps | +| `map_config.learning_rate` | float > 0 | `0.01` | Adam learning rate | + +MAP returns a single point estimate (1 chain, 1 draw in the InferenceData). + +### VI config + +```yaml +inference: + method: vi + vi_config: + num_steps: 5000 + learning_rate: 0.001 + num_samples: 2000 + rng_seed: 42 +``` + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `vi_config.num_steps` | int > 0 | `5000` | SVI optimization steps | +| `vi_config.learning_rate` | float > 0 | `0.001` | Adam learning rate | +| `vi_config.num_samples` | int > 0 | `1000` | Posterior samples drawn from fitted guide | + +### Common parameters + | Parameter | Type | Default | Description | |-----------|------|---------|-------------| -| `warmup` | int > 0 | -- | NUTS warmup iterations | -| `samples` | int > 0 | -- | Number of posterior samples | -| `chains` | int > 0 | -- | Number of MCMC chains | -| `dense_mass` | bool | `false` | Dense mass matrix for correlated parameters | -| `rng_seed` | int >= 0 | -- | JAX PRNG seed | -| `map_init.enabled` | bool | `false` | Enable MAP pre-initialization | -| `map_init.num_steps` | int > 0 | -- | Optimization steps for MAP | -| `map_init.learning_rate` | float > 0 | -- | Adam learning rate for MAP | +| `method` | `"nuts"` / `"map"` / `"vi"` | `"nuts"` | Inference method | +| `rng_seed` | int >= 0 | `0` | JAX PRNG seed | ## Complete Example @@ -220,13 +284,15 @@ gal: y_max: 24.5 inference: - warmup: 500 - samples: 1000 - chains: 2 - dense_mass: false + method: nuts + nuts_config: + warmup: 500 + samples: 1000 + chains: 2 + dense_mass: false + map_init: + enabled: true + num_steps: 1000 + learning_rate: 0.01 rng_seed: 42 - map_init: - enabled: true - num_steps: 1000 - learning_rate: 0.01 ``` diff --git a/docs/getting-started.md b/docs/getting-started.md index bcbde2e..e180d59 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -70,15 +70,17 @@ gal: sigma: 0.05 inference: - warmup: 200 - samples: 500 - chains: 1 - dense_mass: false + method: nuts # "nuts", "map", or "vi" + nuts_config: + warmup: 200 + samples: 500 + chains: 1 + dense_mass: false + map_init: + enabled: true + num_steps: 500 + learning_rate: 0.01 rng_seed: 42 - map_init: - enabled: true - num_steps: 500 - learning_rate: 0.01 ``` Here, `flux` and `half_light_radius` are fixed values. The shear components @@ -95,7 +97,7 @@ This will: 1. Generate synthetic data from the config (since no `data_path` is specified) 2. Build the NumPyro probabilistic model -3. Run MAP initialization followed by NUTS MCMC +3. Run inference using the configured method (NUTS with MAP init in this example) 4. Save the posterior as `results/posterior.nc` (ArviZ NetCDF format) Override the output directory with `--output`: diff --git a/docs/validation/batched.md b/docs/validation/batched.md index 7fd3f26..ac46f0e 100644 --- a/docs/validation/batched.md +++ b/docs/validation/batched.md @@ -6,15 +6,17 @@ and makes better use of GPU parallelism. ## How It Works -Instead of running N separate MCMC jobs, batched inference: +Instead of running N separate inference jobs, batched inference: 1. Generates N synthetic observations and stacks them into a single array 2. Builds a batched NumPyro model that `vmap`s over the batch dimension -3. Runs one MCMC chain that samples all N shear posteriors simultaneously +3. Runs one inference pass that samples all N shear posteriors simultaneously 4. Splits the combined posterior back into per-realization outputs Each realization gets its own shear latent variables (`g1_0`, `g1_1`, ...) so -they are independent despite sharing the same MCMC chain. +they are independent despite sharing the same inference run. The inference +method (NUTS, MAP, or VI) is determined by the YAML config's `inference.method` +field. ## Usage diff --git a/docs/validation/index.md b/docs/validation/index.md index f428f16..a2a56a7 100644 --- a/docs/validation/index.md +++ b/docs/validation/index.md @@ -26,19 +26,21 @@ Stage 1 (Run) Stage 2 (Extract) Stage 3 (Stats) Config + shear → posterior.nc → summary.csv → bias_results.json ↓ ↓ ↓ Generate data Extract diagnostics Compute m, c -Run MCMC Check convergence Check acceptance +Run inference Check convergence Check acceptance Save posterior Write CSV Generate plots ``` ### Stage 1: Run (`shine-bias-run`) -Generates synthetic data with an explicit shear override and runs MCMC inference. +Generates synthetic data with an explicit shear override and runs inference. +The inference method (NUTS, MAP, or VI) is determined by the `inference.method` +field in the SHINE config YAML. **Outputs** (per realization): -- `posterior.nc` -- ArviZ InferenceData with posterior samples +- `posterior.nc` -- ArviZ InferenceData (posterior samples, or point estimate for MAP) - `truth.json` -- ground truth shear values and seed -- `convergence.json` -- R-hat, ESS, divergences, BFMI +- `convergence.json` -- convergence diagnostics (method-aware) ### Stage 2: Extract (`shine-bias-extract`) diff --git a/docs/validation/level0.md b/docs/validation/level0.md index cc5541f..3348036 100644 --- a/docs/validation/level0.md +++ b/docs/validation/level0.md @@ -1,8 +1,9 @@ # Level 0 Walkthrough Level 0 is a noiseless sanity check: a single galaxy with fixed morphology, -very low noise, and a known shear. The posterior should collapse tightly -around the true shear values. +very low noise, and a known shear. Since there is effectively no noise, MAP +estimation is sufficient -- the point estimate should land directly on the +true shear values. ## Prerequisites @@ -12,7 +13,9 @@ pip install -e . ## Step 1: Run a single realization -Use `shine-bias-run` to generate data with a known shear and run MCMC: +Use `shine-bias-run` to generate data with a known shear and run MAP inference. +The default Level 0 config (`configs/validation/level0_base.yaml`) uses +`method: map`: ```bash shine-bias-run \ @@ -28,9 +31,9 @@ This produces: ``` results/validation/level0/r0001/ -├── posterior.nc # ArviZ InferenceData +├── posterior.nc # ArviZ InferenceData (MAP point estimate) ├── truth.json # {"g1": 0.02, "g2": -0.01} -└── convergence.json # R-hat, ESS, divergences +└── convergence.json # Convergence diagnostics (sentinels for MAP) ``` ## Step 2: Extract results @@ -61,7 +64,7 @@ shine-bias-stats \ This produces: - `stats/bias_results.json` -- bias values and overall pass/fail -- `stats/plots/` -- diagnostic plots (trace, marginals, pair plot) +- `stats/plots/` -- diagnostic plots (MAP estimate vs truth) ## Acceptance Criteria @@ -69,8 +72,8 @@ Level 0 checks three conditions: | Criterion | Threshold | Meaning | |-----------|-----------|---------| -| Posterior width | $\sigma < 0.01$ | Posterior should be tight | -| Offset from truth | $< 1\sigma$ | Mean should be near truth | +| Posterior width | $\sigma < 0.01$ | Posterior should be tight (0 for MAP) | +| Offset from truth | $< 1\sigma$ | Estimate should be near truth | | Multiplicative bias | $\|m\| < 0.01$ | Less than 1% bias | ## Inspecting Results @@ -82,14 +85,12 @@ import arviz as az idata = az.from_netcdf("results/validation/level0/r0001/posterior.nc") -# Summary table -print(az.summary(idata, var_names=["g1", "g2"])) +# Check the inference method +print(idata.posterior.attrs.get("inference_method")) # "map" -# Trace plot -az.plot_trace(idata, var_names=["g1", "g2"]) - -# Pair plot -az.plot_pair(idata, var_names=["g1", "g2"], kind="kde") +# Point estimates +print(f"g1 = {float(idata.posterior.g1.values.flatten()[0]):.6f}") +print(f"g2 = {float(idata.posterior.g2.values.flatten()[0]):.6f}") ``` ## Running Multiple Realizations diff --git a/examples/level0_batched_inference.ipynb b/examples/level0_batched_inference.ipynb new file mode 100644 index 0000000..8d8a74f --- /dev/null +++ b/examples/level0_batched_inference.ipynb @@ -0,0 +1,191 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": "# Level 0: Batched MAP Shear Inference\n\nThis notebook demonstrates SHINE's **Level 0 sanity check** on a batch of 10 galaxies\nsharing the same true shear. Level 0 is the noiseless self-consistency test: when data\nis generated by the exact same forward model with effectively zero noise, the MAP\nestimate should recover the truth exactly.\n\nSince there is no noise, MAP (point estimation) is the natural choice -- full MCMC\nis unnecessary and much slower.\n\nSimulation parameters are matched to the\n[ngmix metacal example](https://github.com/esheldon/ngmix/blob/master/examples/metacal/metacal.py):\n\n| Parameter | Value | Source |\n|-----------|-------|--------|\n| Galaxy | Exponential, hlr=0.5\" | metacal `gal_hlr=0.5` |\n| PSF | Moffat, $\\beta$=2.5, FWHM=0.9\" | metacal `psf_fwhm=0.9` |\n| Pixel scale | 0.263\"/px | metacal `scale=0.263` |\n| Noise | $\\sigma = 10^{-6}$ | metacal `noise=1e-6` |\n| Shear | $g_1=0.01$, $g_2=0.00$ | metacal `shear_true=[0.01, 0.00]` |\n\n**What we do:**\n1. Generate 10 synthetic observations with the same shear\n2. Run MAP inference on each realization independently\n3. Check that MAP estimates match truth with negligible bias" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "import jax\nimport jax.numpy as jnp\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport arviz as az\n\nfrom shine.config import (\n ShineConfig,\n ImageConfig,\n NoiseConfig,\n PSFConfig,\n GalaxyConfig,\n ShearConfig,\n EllipticityConfig,\n PositionConfig,\n InferenceConfig,\n MAPConfig,\n DistributionConfig,\n)\nfrom shine.scene import SceneBuilder\nfrom shine.inference import Inference\nfrom shine.validation.simulation import generate_biased_observation\nfrom shine.validation.extraction import (\n extract_convergence_diagnostics,\n extract_shear_estimates,\n check_convergence,\n)\nfrom shine.validation.bias_config import ConvergenceThresholds\nfrom shine.validation.statistics import compute_bias_single_point\n\n# Use 64-bit precision for accurate shear recovery\njax.config.update(\"jax_enable_x64\", True)\n\nprint(f\"JAX devices: {jax.devices()}\")" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Configuration\n", + "\n", + "Parameters matched to the [ngmix metacal example](https://github.com/esheldon/ngmix/blob/master/examples/metacal/metacal.py):\n", + "- **Galaxy**: Exponential, hlr=0.5\" (metacal `gal_hlr=0.5`)\n", + "- **PSF**: Moffat, $\\beta = 2.5$, FWHM=0.9\" (metacal `psf_fwhm=0.9, beta=2.5`)\n", + "- **Pixel scale**: 0.263\"/px (metacal `scale=0.263`)\n", + "- **Noise**: $\\sigma = 10^{-6}$ (metacal default `noise=1e-6`)\n", + "- **Shear**: $g_1 = 0.01$, $g_2 = 0.00$ (metacal `shear_true=[0.01, 0.00]`)\n", + "- **Position**: Fixed at center (metacal uses random subpixel offsets)\n", + "\n", + "**Note:** The metacal PSF has intrinsic ellipticity ($g_1^{\\rm PSF}=0.02$, $g_2^{\\rm PSF}=-0.01$).\n", + "SHINE does not yet support PSF ellipticity, so we use a round Moffat PSF here." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "# Ground truth shear (matches metacal shear_true=[0.01, 0.00])\nG1_TRUE = 0.01\nG2_TRUE = 0.00\nN_BATCH = 10\n\nconfig = ShineConfig(\n image=ImageConfig(\n pixel_scale=0.263, # arcsec/pixel (metacal scale=0.263)\n size_x=48,\n size_y=48,\n n_objects=1,\n fft_size=128,\n noise=NoiseConfig(type=\"Gaussian\", sigma=1e-6), # metacal noise=1e-6\n ),\n psf=PSFConfig(\n type=\"Moffat\",\n sigma=0.9, # FWHM in arcsec (metacal psf_fwhm=0.9)\n beta=2.5, # metacal beta=2.5\n ),\n gal=GalaxyConfig(\n type=\"Exponential\", # metacal galsim.Exponential\n flux=1.0, # metacal default flux=1\n half_light_radius=0.5, # arcsec (metacal gal_hlr=0.5)\n ellipticity=EllipticityConfig(type=\"E1E2\", e1=0.0, e2=0.0),\n shear=ShearConfig(\n type=\"G1G2\",\n g1=DistributionConfig(type=\"Normal\", mean=0.0, sigma=0.05),\n g2=DistributionConfig(type=\"Normal\", mean=0.0, sigma=0.05),\n ),\n position=PositionConfig(\n type=\"Uniform\",\n x_min=23.5, x_max=24.5,\n y_min=23.5, y_max=24.5,\n ),\n ),\n inference=InferenceConfig(\n method=\"map\",\n map_config=MAPConfig(num_steps=1000, learning_rate=0.01),\n rng_seed=42,\n ),\n)\n\nprint(f\"Image: {config.image.size_x}x{config.image.size_y} px, \"\n f\"scale={config.image.pixel_scale}\\\"/px\")\nprint(f\"Galaxy: {config.gal.type}, flux={config.gal.flux}, \"\n f\"hlr={config.gal.half_light_radius}\\\"\")\nprint(f\"PSF: {config.psf.type}, FWHM={config.psf.sigma}\\\", beta={config.psf.beta}\")\nprint(f\"Noise sigma: {config.image.noise.sigma}\")\nprint(f\"True shear: g1={G1_TRUE}, g2={G2_TRUE}\")\nprint(f\"Batch size: {N_BATCH}\")\nprint(f\"Inference method: {config.inference.method}\")" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## 2. Generate Observations\n\nAll 10 galaxies share the same true shear but have independent (effectively zero) noise\nrealizations. We use `generate_biased_observation()` per realization." + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "# Generate N_BATCH independent observations with the same true shear\nseeds = list(range(100, 100 + N_BATCH))\nsim_results = []\nfor seed in seeds:\n sim = generate_biased_observation(config, G1_TRUE, G2_TRUE, seed)\n sim_results.append(sim)\n\nprint(f\"Generated {N_BATCH} observations\")\nprint(f\"Image shape: {sim_results[0].observation.image.shape}\")\nprint(f\"PSF type: {type(sim_results[0].observation.psf_model).__name__}\")" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "# Visualize a few of the generated images\nfig, axes = plt.subplots(2, 5, figsize=(15, 6))\nfor i, ax in enumerate(axes.flat):\n im = ax.imshow(sim_results[i].observation.image, origin=\"lower\", cmap=\"gray_r\")\n ax.set_title(f\"Galaxy {i}\", fontsize=10)\n ax.set_xticks([])\n ax.set_yticks([])\nfig.suptitle(\n f\"Level 0: 10 Exponential galaxies, g1={G1_TRUE}, g2={G2_TRUE}, \"\n f\"noise={config.image.noise.sigma}\",\n fontsize=12,\n)\nplt.tight_layout()\nplt.show()" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## 3. Run MAP Inference\n\nFor each realization, we build a model and run MAP estimation. Since Level 0\nis noiseless, MAP finds the maximum a posteriori point estimate which should\nmatch the truth exactly." + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "# Run MAP inference for each realization\nscene = SceneBuilder(config)\nmodel_fn = scene.build_model()\n\nmap_cfg = config.inference.map_config\nprint(f\"Inference method: {config.inference.method}\")\nprint(f\"MAP: {map_cfg.num_steps} steps, lr={map_cfg.learning_rate}\")\n\nidata_list = []\nfor i, sim in enumerate(sim_results):\n rng_key = jax.random.PRNGKey(config.inference.rng_seed + i)\n engine = Inference(model=model_fn, config=config.inference)\n idata = engine.run(\n rng_key=rng_key,\n observed_data=sim.observation.image,\n extra_args={\"psf\": sim.observation.psf_model},\n )\n idata_list.append(idata)\n g1_val = float(idata.posterior.g1.values.flatten()[0])\n g2_val = float(idata.posterior.g2.values.flatten()[0])\n print(f\" Realization {i}: g1={g1_val:+.6f}, g2={g2_val:+.6f}\")\n\nprint(f\"\\nCompleted {N_BATCH} MAP estimations\")" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## 4. Extract Estimates\n\nFor each realization, extract the MAP point estimates and check that they\nmatch the truth." + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "run_ids = [f\"level0_{i:04d}\" for i in range(N_BATCH)]\n\nprint(f\"Extracted {N_BATCH} MAP estimates\")\nprint(f\"Example: {run_ids[0]}, posterior vars: {list(idata_list[0].posterior.data_vars)}\")\nprint(f\" inference_method: {idata_list[0].posterior.attrs.get('inference_method')}\")" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## 5. Extract Diagnostics\n\nFor each realization, extract:\n- **Shear estimates**: MAP point estimate (mean = median = value, std = 0)\n- **Convergence diagnostics**: sentinel values for MAP (rhat=1, ess=1)\n\nLevel 0 acceptance criterion for MAP: the estimate should be very close to truth." + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "# Extract results for all realizations\nresults = []\nfor run_id, single_idata in zip(run_ids, idata_list):\n g1_est = extract_shear_estimates(single_idata, \"g1\")\n g2_est = extract_shear_estimates(single_idata, \"g2\")\n diag = extract_convergence_diagnostics(single_idata)\n method = single_idata.posterior.attrs.get(\"inference_method\", \"nuts\")\n passed = check_convergence(diag, ConvergenceThresholds(), method=method)\n results.append({\n \"run_id\": run_id,\n \"g1_est\": g1_est,\n \"g2_est\": g2_est,\n \"diagnostics\": diag,\n \"passed\": passed,\n })\n\n# Summary table\nprint(f\"{'Run ID':<14} {'g1 estimate':>14} {'g2 estimate':>14} {'Pass':>5}\")\nprint(\"-\" * 55)\nfor r in results:\n print(\n f\"{r['run_id']:<14} \"\n f\"{r['g1_est'].mean:>14.6f} \"\n f\"{r['g2_est'].mean:>14.6f} \"\n f\"{'OK' if r['passed'] else 'FAIL':>5}\"\n )" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## 6. Acceptance Criteria Check\n\nFor Level 0 MAP, each realization must satisfy:\n1. MAP estimate close to truth (absolute offset $< 10^{-3}$)\n2. Convergence always passes for MAP (no sampling diagnostics)" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "MAX_ABS_OFFSET = 1e-3\n\nall_passed = True\nprint(\"Level 0 Acceptance Criteria (MAP)\")\nprint(\"=\" * 80)\n\nfor r in results:\n run_id = r[\"run_id\"]\n g1, g2 = r[\"g1_est\"], r[\"g2_est\"]\n\n # For MAP: check absolute offset from truth\n g1_offset = abs(g1.mean - G1_TRUE)\n g2_offset = abs(g2.mean - G2_TRUE)\n offset_ok = g1_offset < MAX_ABS_OFFSET and g2_offset < MAX_ABS_OFFSET\n\n passed = offset_ok and r[\"passed\"]\n all_passed = all_passed and passed\n\n status = \"PASS\" if passed else \"FAIL\"\n print(f\"\\n{run_id} [{status}]\")\n print(f\" g1: truth={G1_TRUE:+.4f} MAP={g1.mean:+.6f} \"\n f\"|offset|={g1_offset:.2e} {'ok' if g1_offset < MAX_ABS_OFFSET else 'FAIL'}\")\n print(f\" g2: truth={G2_TRUE:+.4f} MAP={g2.mean:+.6f} \"\n f\"|offset|={g2_offset:.2e} {'ok' if g2_offset < MAX_ABS_OFFSET else 'FAIL'}\")\n\nprint(\"\\n\" + \"=\" * 80)\nprint(f\"Overall Level 0 result: {'ALL PASSED' if all_passed else 'SOME FAILED'}\")\nprint(f\" {sum(1 for r in results if r['passed'])}/{len(results)} \"\n f\"realizations passed\")" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 7. Multiplicative Bias\n", + "\n", + "For each realization, compute the multiplicative bias:\n", + "$$m = \\frac{\\bar{g}_{\\rm est}}{g_{\\rm true}} - 1$$\n", + "\n", + "At Level 0 (noiseless, self-consistent model), we expect $m \\approx 0$.\n", + "\n", + "Since $g_2^{\\rm true} = 0$ (matching metacal), we cannot compute $m$ for $g_2$\n", + "(division by zero). Instead we report the additive residual $c_2 = \\bar{g}_2 - 0$." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "bias_g1_list = []\n\nprint(f\"{'Run ID':<14} {'m(g1)':>12} {'c(g2)':>12}\")\nprint(\"-\" * 42)\n\nfor r in results:\n b1 = compute_bias_single_point(G1_TRUE, r[\"g1_est\"].mean, r[\"g1_est\"].std, \"g1\")\n bias_g1_list.append(b1)\n # g2_true=0 so we report additive residual instead of multiplicative bias\n c2 = r[\"g2_est\"].mean - G2_TRUE\n print(f\"{r['run_id']:<14} {b1.m:>12.6f} {c2:>12.2e}\")\n\n# Ensemble average\nm_g1_vals = np.array([b.m for b in bias_g1_list])\nc_g2_vals = np.array([r[\"g2_est\"].mean - G2_TRUE for r in results])\nprint(\"-\" * 42)\nprint(f\"{'Ensemble mean':<14} {m_g1_vals.mean():>12.6f} {c_g2_vals.mean():>12.2e}\")" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 8. Diagnostic Plots" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "# Collect all g1, g2 MAP estimates\ng1_means = np.array([r[\"g1_est\"].mean for r in results])\ng2_means = np.array([r[\"g2_est\"].mean for r in results])\nindices = np.arange(N_BATCH)\n\nfig, axes = plt.subplots(1, 2, figsize=(14, 5))\n\n# g1 estimates\naxes[0].scatter(indices, g1_means, marker=\"o\", s=60, color=\"steelblue\",\n zorder=5, label=\"MAP estimate\")\naxes[0].axhline(G1_TRUE, color=\"red\", ls=\"--\", lw=2, label=f\"Truth = {G1_TRUE}\")\naxes[0].fill_between([-0.5, N_BATCH - 0.5],\n G1_TRUE - MAX_ABS_OFFSET,\n G1_TRUE + MAX_ABS_OFFSET,\n color=\"red\", alpha=0.1, label=f\"$\\\\pm${MAX_ABS_OFFSET}\")\naxes[0].set_xlabel(\"Realization\")\naxes[0].set_ylabel(\"$g_1$\")\naxes[0].set_title(\"$g_1$ Recovery (MAP)\")\naxes[0].legend(fontsize=9)\naxes[0].set_xticks(indices)\n\n# g2 estimates\naxes[1].scatter(indices, g2_means, marker=\"o\", s=60, color=\"coral\",\n zorder=5, label=\"MAP estimate\")\naxes[1].axhline(G2_TRUE, color=\"red\", ls=\"--\", lw=2, label=f\"Truth = {G2_TRUE}\")\naxes[1].fill_between([-0.5, N_BATCH - 0.5],\n G2_TRUE - MAX_ABS_OFFSET,\n G2_TRUE + MAX_ABS_OFFSET,\n color=\"red\", alpha=0.1, label=f\"$\\\\pm${MAX_ABS_OFFSET}\")\naxes[1].set_xlabel(\"Realization\")\naxes[1].set_ylabel(\"$g_2$\")\naxes[1].set_title(\"$g_2$ Recovery (MAP)\")\naxes[1].legend(fontsize=9)\naxes[1].set_xticks(indices)\n\nfig.suptitle(\"Level 0: MAP Shear Recovery Across 10 Realizations\", fontsize=13)\nplt.tight_layout()\nplt.show()" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "# Bias plots: m(g1) and c(g2)\nfig, axes = plt.subplots(1, 2, figsize=(14, 5))\n\n# g1: multiplicative bias\naxes[0].scatter(indices, m_g1_vals, marker=\"s\", s=60, color=\"steelblue\")\naxes[0].axhline(0, color=\"red\", ls=\"--\", lw=2, label=\"$m = 0$ (no bias)\")\naxes[0].axhline(m_g1_vals.mean(), color=\"green\", ls=\":\", lw=1.5,\n label=f\"Mean $m$ = {m_g1_vals.mean():.2e}\")\naxes[0].set_xlabel(\"Realization\")\naxes[0].set_ylabel(\"$m_{g_1}$\")\naxes[0].set_title(\"Multiplicative Bias $g_1$\")\naxes[0].legend(fontsize=9)\naxes[0].set_xticks(indices)\n\n# g2: additive bias (g2_true = 0)\naxes[1].scatter(indices, c_g2_vals, marker=\"s\", s=60, color=\"coral\")\naxes[1].axhline(0, color=\"red\", ls=\"--\", lw=2, label=\"$c = 0$ (no bias)\")\naxes[1].axhline(c_g2_vals.mean(), color=\"green\", ls=\":\", lw=1.5,\n label=f\"Mean $c$ = {c_g2_vals.mean():.2e}\")\naxes[1].set_xlabel(\"Realization\")\naxes[1].set_ylabel(\"$c_{g_2}$\")\naxes[1].set_title(\"Additive Bias $g_2$ ($g_2^{\\\\rm true} = 0$)\")\naxes[1].legend(fontsize=9)\naxes[1].set_xticks(indices)\n\nfig.suptitle(\"Level 0: Bias per Realization (MAP)\", fontsize=13)\nplt.tight_layout()\nplt.show()" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "# g1 vs g2 MAP estimates across all realizations\nfig, ax = plt.subplots(1, 1, figsize=(6, 6))\nax.scatter(g1_means, g2_means, s=60, color=\"steelblue\", zorder=5, label=\"MAP estimates\")\nax.scatter([G1_TRUE], [G2_TRUE], c=\"red\", s=150, marker=\"*\",\n zorder=10, label=\"Truth\")\nax.set_xlabel(\"$g_1$\")\nax.set_ylabel(\"$g_2$\")\nax.set_title(\"MAP Estimates: $g_1$ vs $g_2$\")\nax.legend()\nfig.tight_layout()\nplt.show()" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## 9. Summary\n\nThis Level 0 test validates that SHINE's forward model is **self-consistent**: when\nthe data is generated from the same model with no noise, the MAP estimate recovers\nthe true shear values with negligible bias.\n\nSince Level 0 is noiseless, MAP is the natural and fastest inference method --\nfull MCMC is unnecessary. For higher validation levels (Level 1+) with realistic\nnoise, NUTS or VI should be used instead." + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "# Final summary\nn_passed = sum(1 for r in results if r[\"passed\"])\n\nprint(\"Level 0 MAP Inference Summary\")\nprint(\"=\" * 50)\nprint(f\" Batch size: {N_BATCH}\")\nprint(f\" True shear: g1={G1_TRUE}, g2={G2_TRUE}\")\nprint(f\" Inference method: MAP\")\nprint(f\" All passed: {n_passed}/{N_BATCH}\")\nprint(f\" Mean m(g1): {m_g1_vals.mean():.2e}\")\nprint(f\" Mean c(g2): {c_g2_vals.mean():.2e}\")\nprint(f\" Max |g1 offset|: {max(abs(r['g1_est'].mean - G1_TRUE) for r in results):.2e}\")\nprint(f\" Max |g2 offset|: {max(abs(r['g2_est'].mean - G2_TRUE) for r in results):.2e}\")\nprint(\"=\" * 50)" + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv (3.14.2)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.14.2" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/shine/config.py b/shine/config.py index 9f658f5..faa517f 100644 --- a/shine/config.py +++ b/shine/config.py @@ -1,5 +1,5 @@ from pathlib import Path -from typing import Optional, Union +from typing import Literal, Optional, Union import yaml from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator @@ -297,16 +297,15 @@ def validate_learning_rate_positive(cls, v: float) -> float: return v -class InferenceConfig(BaseModel): - """Configuration for Bayesian inference settings. +class NUTSConfig(BaseModel): + """Configuration for NUTS/MCMC inference. Attributes: warmup: Number of warmup/burn-in steps for MCMC (default 500). samples: Number of posterior samples to draw (default 1000). chains: Number of independent MCMC chains to run (default 1). dense_mass: Whether to use dense mass matrix in NUTS (default False). - map_init: Optional MAP initialization configuration. - rng_seed: Random number generator seed for reproducibility (default 0). + map_init: Optional MAP initialization before NUTS. """ warmup: int = 500 @@ -314,7 +313,6 @@ class InferenceConfig(BaseModel): chains: int = 1 dense_mass: bool = False map_init: Optional[MAPConfig] = None - rng_seed: int = 0 @field_validator("warmup", "samples") @classmethod @@ -332,6 +330,62 @@ def validate_chains_positive(cls, v: int) -> int: raise ValueError(f"Number of chains must be positive, got {v}") return v + +class VIConfig(BaseModel): + """Configuration for Variational Inference. + + Attributes: + num_steps: Number of SVI optimization steps (default 5000). + learning_rate: Learning rate for SVI optimizer (default 1e-3). + num_samples: Number of posterior samples to draw from fitted guide (default 1000). + """ + + num_steps: int = 5000 + learning_rate: float = 1e-3 + num_samples: int = 1000 + + @field_validator("num_steps", "num_samples") + @classmethod + def validate_positive_integers(cls, v: int, info) -> int: + """Validate that num_steps and num_samples are positive.""" + if v <= 0: + raise ValueError(f"{info.field_name} must be positive, got {v}") + return v + + @field_validator("learning_rate") + @classmethod + def validate_learning_rate_positive(cls, v: float) -> float: + """Validate that learning rate is positive.""" + if v <= 0: + raise ValueError(f"Learning rate must be positive, got {v}") + return v + + +class InferenceConfig(BaseModel): + """Configuration for Bayesian inference settings. + + Supports three inference methods: + - "nuts": NUTS/MCMC sampling (default), optionally with MAP initialization. + - "map": MAP point estimation only. + - "vi": Variational Inference with AutoNormal guide. + + Each method reads its own config block; the others are ignored. + When a method's config block is None, defaults are applied. + + Attributes: + method: Inference method ("nuts", "map", or "vi"). + nuts_config: Configuration for NUTS/MCMC (used when method="nuts"). + map_config: Configuration for MAP estimation (used when method="map"). + vi_config: Configuration for Variational Inference (used when method="vi"). + rng_seed: Random number generator seed for reproducibility (default 0). + """ + + method: Literal["nuts", "map", "vi"] = "nuts" + nuts_config: Optional[NUTSConfig] = None + map_config: Optional[MAPConfig] = None + vi_config: Optional[VIConfig] = None + rng_seed: int = 0 + @field_validator("rng_seed") @classmethod def validate_rng_seed_non_negative(cls, v: int) -> int: diff --git a/shine/inference.py b/shine/inference.py index 8040970..33b19dc 100644 --- a/shine/inference.py +++ b/shine/inference.py @@ -4,20 +4,21 @@ import arviz as az import jax import jax.numpy as jnp +import numpy as np import numpyro from numpyro.infer import MCMC, NUTS, SVI, Trace_ELBO -from numpyro.infer.autoguide import AutoDelta +from numpyro.infer.autoguide import AutoDelta, AutoNormal -from shine.config import InferenceConfig, MAPConfig +from shine.config import InferenceConfig, MAPConfig, NUTSConfig, VIConfig logger = logging.getLogger(__name__) class Inference: - """Inference engine supporting HMC/MCMC and optional MAP initialization. + """Inference engine supporting NUTS/MCMC, MAP, and Variational Inference. - MAP estimation can be used as an initialization step before running - MCMC chains to improve convergence. + All three methods return az.InferenceData so the downstream pipeline + (extraction, diagnostics, plots) works uniformly. """ def __init__(self, model: Callable, config: InferenceConfig) -> None: @@ -66,6 +67,81 @@ def run_map( logger.info("MAP estimation complete.") return map_estimates + @staticmethod + def _map_estimates_to_idata(map_estimates: Dict[str, Any]) -> az.InferenceData: + """Wrap MAP point estimates as InferenceData (1 chain, 1 draw). + + Args: + map_estimates: Dictionary of MAP parameter estimates. + + Returns: + ArviZ InferenceData with posterior group containing point estimates. + """ + posterior_dict = {} + for name, value in map_estimates.items(): + arr = jnp.atleast_1d(jnp.asarray(value)) + posterior_dict[name] = np.array(arr)[None, None, ...] # (1, 1, ...) + idata = az.from_dict(posterior=posterior_dict) + idata.posterior.attrs["inference_method"] = "map" + return idata + + def run_vi( + self, + rng_key: jax.random.PRNGKey, + observed_data: jnp.ndarray, + extra_args: Optional[Dict[str, Any]] = None, + ) -> az.InferenceData: + """Run Variational Inference with AutoNormal guide. + + Args: + rng_key: JAX random key. + observed_data: Observed image data. + extra_args: Extra keyword arguments passed to the model (e.g., psf). + + Returns: + ArviZ InferenceData with posterior samples from the fitted guide. + """ + if extra_args is None: + extra_args = {} + + vi_config = self.config.vi_config or VIConfig() + guide = AutoNormal(self.model) + optimizer = numpyro.optim.Adam(step_size=vi_config.learning_rate) + svi = SVI(self.model, guide, optimizer, loss=Trace_ELBO()) + + logger.info( + f"Running VI: {vi_config.num_steps} steps, " + f"lr={vi_config.learning_rate}..." + ) + svi_result = svi.run( + rng_key, vi_config.num_steps, observed_data=observed_data, **extra_args + ) + + # Draw posterior samples from fitted guide + sample_key, _ = jax.random.split(rng_key) + predictive = numpyro.infer.Predictive( + guide, params=svi_result.params, num_samples=vi_config.num_samples + ) + vi_samples = predictive( + sample_key, observed_data=observed_data, **extra_args + ) + + # Wrap as InferenceData (1 chain, N draws), filtering out "obs" + posterior_dict = { + k: np.array(v)[None, ...] + for k, v in vi_samples.items() + if k != "obs" + } + idata = az.from_dict(posterior=posterior_dict) + idata.posterior.attrs["inference_method"] = "vi" + idata.posterior.attrs["vi_final_loss"] = float(svi_result.losses[-1]) + + logger.info( + f"VI complete. Final ELBO loss: {svi_result.losses[-1]:.4f}, " + f"{vi_config.num_samples} posterior samples drawn." + ) + return idata + def run_mcmc( self, rng_key: jax.random.PRNGKey, @@ -87,6 +163,8 @@ def run_mcmc( if extra_args is None: extra_args = {} + nuts_cfg = self.config.nuts_config or NUTSConfig() + # init_to_uniform is robust for unbounded distributions where init_to_median may fail if init_params is not None: init_strategy = numpyro.infer.init_to_value(values=init_params) @@ -95,19 +173,19 @@ def run_mcmc( kernel = NUTS( self.model, - dense_mass=self.config.dense_mass, + dense_mass=nuts_cfg.dense_mass, init_strategy=init_strategy, ) mcmc = MCMC( kernel, - num_warmup=self.config.warmup, - num_samples=self.config.samples, - num_chains=self.config.chains, + num_warmup=nuts_cfg.warmup, + num_samples=nuts_cfg.samples, + num_chains=nuts_cfg.chains, ) logger.info( - f"Running MCMC: {self.config.warmup} warmup, " - f"{self.config.samples} samples, {self.config.chains} chain(s)..." + f"Running MCMC: {nuts_cfg.warmup} warmup, " + f"{nuts_cfg.samples} samples, {nuts_cfg.chains} chain(s)..." ) mcmc.run(rng_key, observed_data=observed_data, **extra_args) mcmc.print_summary() @@ -120,10 +198,7 @@ def run( observed_data: jnp.ndarray, extra_args: Optional[Dict[str, Any]] = None, ) -> az.InferenceData: - """Run full inference pipeline with optional MAP initialization. - - If MAP initialization is enabled in config, runs MAP first to find - good starting points, then runs MCMC. + """Run inference pipeline, dispatching on the configured method. Args: rng_key: JAX random key. @@ -131,17 +206,29 @@ def run( extra_args: Extra keyword arguments passed to the model (e.g., psf). Returns: - ArviZ InferenceData object with posterior samples. + ArviZ InferenceData object with posterior samples/estimates. """ - init_params = None + method = self.config.method + + if method == "map": + map_cfg = self.config.map_config or MAPConfig() + estimates = self.run_map(rng_key, observed_data, extra_args, map_cfg) + return self._map_estimates_to_idata(estimates) - map_init = self.config.map_init - if map_init is not None and map_init.enabled: + if method == "vi": + return self.run_vi(rng_key, observed_data, extra_args) + + # NUTS: optional MAP init then MCMC + nuts_cfg = self.config.nuts_config or NUTSConfig() + init_params = None + if nuts_cfg.map_init is not None and nuts_cfg.map_init.enabled: map_key, rng_key = jax.random.split(rng_key) init_params = self.run_map( - map_key, observed_data, extra_args, map_init + map_key, observed_data, extra_args, nuts_cfg.map_init ) else: logger.info("Skipping MAP initialization.") - return self.run_mcmc(rng_key, observed_data, extra_args, init_params) + idata = self.run_mcmc(rng_key, observed_data, extra_args, init_params) + idata.posterior.attrs["inference_method"] = "nuts" + return idata diff --git a/shine/main.py b/shine/main.py index 2b37a7d..49f04a5 100644 --- a/shine/main.py +++ b/shine/main.py @@ -78,7 +78,7 @@ def main() -> None: model_fn = scene_builder.build_model() # 6. Run Bayesian inference - logger.info("Starting Bayesian inference pipeline...") + logger.info(f"Starting {config.inference.method.upper()} inference pipeline...") try: rng_key = jax.random.PRNGKey(config.inference.rng_seed) engine = Inference(model=model_fn, config=config.inference) diff --git a/shine/validation/cli.py b/shine/validation/cli.py index 164d627..f71f1fb 100644 --- a/shine/validation/cli.py +++ b/shine/validation/cli.py @@ -155,7 +155,7 @@ def run_bias_realization() -> None: rng_key = jax.random.PRNGKey(shine_config.inference.rng_seed) engine = Inference(model=model_fn, config=shine_config.inference) - logger.info("Running MCMC inference...") + logger.info(f"Running {shine_config.inference.method.upper()} inference...") idata = engine.run( rng_key=rng_key, observed_data=sim_result.observation.image, @@ -273,7 +273,7 @@ def _run_batched(args: argparse.Namespace) -> None: rng_key = jax.random.PRNGKey(shine_config.inference.rng_seed) engine = Inference(model=model_fn, config=shine_config.inference) - logger.info(f"Running batched MCMC ({n_batch} realizations)...") + logger.info(f"Running batched {shine_config.inference.method.upper()} ({n_batch} realizations)...") idata = engine.run( rng_key=rng_key, observed_data=batch_result.images, diff --git a/shine/validation/extraction.py b/shine/validation/extraction.py index 7ec86cf..45afdd5 100644 --- a/shine/validation/extraction.py +++ b/shine/validation/extraction.py @@ -83,6 +83,13 @@ def extract_convergence_diagnostics( ) -> ConvergenceDiagnostics: """Extract convergence diagnostics from an InferenceData object. + Method-aware: reads ``inference_method`` from ``idata.posterior.attrs`` + to determine which diagnostics are applicable. + + - MAP (1 chain, 1 draw): returns sentinel values (rhat=1, ess=1). + - VI (1 chain, N draws): computes ESS, sets rhat=1, no MCMC stats. + - NUTS (default): full MCMC diagnostics. + Args: idata: ArviZ InferenceData with posterior and sample_stats groups. params: Parameter names to compute diagnostics for (default: ["g1", "g2"]). @@ -93,6 +100,37 @@ def extract_convergence_diagnostics( if params is None: params = ["g1", "g2"] + method = idata.posterior.attrs.get("inference_method", "nuts") + posterior = idata.posterior + n_chains = posterior.sizes.get("chain", 1) + n_samples_per_chain = posterior.sizes.get("draw", 0) + n_samples = n_chains * n_samples_per_chain + + if method == "map": + return ConvergenceDiagnostics( + rhat={p: 1.0 for p in params}, + ess={p: 1.0 for p in params}, + divergences=0, + divergence_frac=0.0, + bfmi=[], + n_samples=n_samples, + n_chains=n_chains, + ) + + if method == "vi": + # ESS is meaningful for VI samples; rhat is not (single chain) + ess_data = az.ess(idata, var_names=params) + return ConvergenceDiagnostics( + rhat={p: 1.0 for p in params}, + ess={p: float(ess_data[p].values) for p in params}, + divergences=0, + divergence_frac=0.0, + bfmi=[], + n_samples=n_samples, + n_chains=n_chains, + ) + + # NUTS: full MCMC diagnostics # R-hat rhat_data = az.rhat(idata, var_names=params) rhat = {p: float(rhat_data[p].values) for p in params} @@ -110,7 +148,6 @@ def extract_convergence_diagnostics( else: divergences = 0 divergence_frac = 0.0 - total_samples = 0 # BFMI try: @@ -120,12 +157,6 @@ def extract_convergence_diagnostics( logger.debug(f"Could not compute BFMI: {exc}") bfmi = [] - # Chain/sample counts - posterior = idata.posterior - n_chains = posterior.sizes.get("chain", 1) - n_samples_per_chain = posterior.sizes.get("draw", 0) - n_samples = n_chains * n_samples_per_chain - return ConvergenceDiagnostics( rhat=rhat, ess=ess, @@ -170,16 +201,38 @@ def extract_shear_estimates( def check_convergence( diagnostics: ConvergenceDiagnostics, thresholds: ConvergenceThresholds, + method: str = "nuts", ) -> bool: - """Check if MCMC convergence diagnostics meet thresholds. + """Check if convergence diagnostics meet thresholds. + + Method-aware: + - MAP: always returns True (point estimate, no convergence to check). + - VI: only checks ESS. + - NUTS: all four checks (rhat, ESS, divergences, BFMI). Args: diagnostics: Computed convergence diagnostics. thresholds: Threshold criteria to check against. + method: Inference method ("nuts", "map", or "vi"). Returns: - True if all diagnostics pass, False otherwise. + True if all applicable diagnostics pass, False otherwise. """ + if method == "map": + return True + + if method == "vi": + # Only check ESS for VI samples + for param, ess_val in diagnostics.ess.items(): + if ess_val < thresholds.ess_min: + logger.warning( + f"ESS for {param} = {ess_val:.0f} below " + f"threshold {thresholds.ess_min}" + ) + return False + return True + + # NUTS: full checks # Check R-hat (NaN/inf from degenerate posteriors are treated as failures) for param, rhat_val in diagnostics.rhat.items(): if np.isnan(rhat_val) or np.isinf(rhat_val): @@ -249,7 +302,8 @@ def extract_realization( diagnostics = extract_convergence_diagnostics(idata) g1_est = extract_shear_estimates(idata, "g1") g2_est = extract_shear_estimates(idata, "g2") - passed = check_convergence(diagnostics, thresholds) + method = idata.posterior.attrs.get("inference_method", "nuts") + passed = check_convergence(diagnostics, thresholds, method=method) return RealizationResult( run_id=run_id, diff --git a/shine/validation/plots.py b/shine/validation/plots.py index dafa492..ab08c49 100644 --- a/shine/validation/plots.py +++ b/shine/validation/plots.py @@ -69,12 +69,13 @@ def plot_level0_diagnostics( g2_true: float, output_dir: str, ) -> List[Path]: - """Generate Level 0 diagnostic plots. + """Generate Level 0 diagnostic plots, dispatching on inference method. - Produces: - - Trace plots for g1 and g2 - - Marginal posterior histograms with truth lines - - Pair plot (g1 vs g2) + Reads ``inference_method`` from ``idata.posterior.attrs`` to select + the appropriate plotting style: + - NUTS: trace plots + histograms + pair plot + - VI: histograms + pair plot (no trace since no chains) + - MAP: bar at point estimate + truth line (2 panels: g1, g2) Args: idata: ArviZ InferenceData with posterior samples. @@ -85,6 +86,22 @@ def plot_level0_diagnostics( Returns: List of saved plot file paths. """ + method = idata.posterior.attrs.get("inference_method", "nuts") + if method == "map": + return _plot_map_diagnostics(idata, g1_true, g2_true, output_dir) + elif method == "vi": + return _plot_vi_diagnostics(idata, g1_true, g2_true, output_dir) + else: + return _plot_nuts_diagnostics(idata, g1_true, g2_true, output_dir) + + +def _plot_nuts_diagnostics( + idata: az.InferenceData, + g1_true: float, + g2_true: float, + output_dir: str, +) -> List[Path]: + """Generate NUTS diagnostic plots: trace + histograms + pair plot.""" output_path = Path(output_dir) output_path.mkdir(parents=True, exist_ok=True) saved = [] @@ -157,6 +174,111 @@ def plot_level0_diagnostics( return saved +def _plot_vi_diagnostics( + idata: az.InferenceData, + g1_true: float, + g2_true: float, + output_dir: str, +) -> List[Path]: + """Generate VI diagnostic plots: histograms + pair plot (no trace).""" + output_path = Path(output_dir) + output_path.mkdir(parents=True, exist_ok=True) + saved = [] + + g1_samples = idata.posterior.g1.values.flatten() + g2_samples = idata.posterior.g2.values.flatten() + + # --- Marginal posteriors --- + fig, axes = plt.subplots(1, 2, figsize=(12, 5)) + + _safe_hist(axes[0], g1_samples, "steelblue", "black") + axes[0].axvline(g1_true, color="red", ls="--", lw=2, label="Truth") + axes[0].axvline(g1_samples.mean(), color="green", lw=2, label="Mean") + axes[0].set( + xlabel="g1", + title=f"VI: g1 = {g1_samples.mean():.4f} ± {g1_samples.std():.4f}", + ) + axes[0].legend() + + _safe_hist(axes[1], g2_samples, "coral", "black") + axes[1].axvline(g2_true, color="red", ls="--", lw=2, label="Truth") + axes[1].axvline(g2_samples.mean(), color="green", lw=2, label="Mean") + axes[1].set( + xlabel="g2", + title=f"VI: g2 = {g2_samples.mean():.4f} ± {g2_samples.std():.4f}", + ) + axes[1].legend() + + fig.tight_layout() + hist_path = output_path / "vi_posterior.png" + fig.savefig(hist_path, dpi=150) + plt.close(fig) + saved.append(hist_path) + logger.info(f"Saved VI posterior plot to {hist_path}") + + # --- Pair plot --- + fig2, ax2 = plt.subplots(1, 1, figsize=(6, 6)) + ax2.scatter(g1_samples, g2_samples, alpha=0.1, s=2, color="steelblue") + ax2.axvline(g1_true, color="red", ls="--", lw=1.5, label="g1 truth") + ax2.axhline(g2_true, color="red", ls="--", lw=1.5, label="g2 truth") + ax2.plot(g1_true, g2_true, "r*", ms=15, label="Truth") + ax2.plot(g1_samples.mean(), g2_samples.mean(), "g*", ms=15, label="Mean") + ax2.set(xlabel="g1", ylabel="g2", title="VI: g1 vs g2 posterior") + ax2.legend() + fig2.tight_layout() + pair_path = output_path / "pair_plot.png" + fig2.savefig(pair_path, dpi=150) + plt.close(fig2) + saved.append(pair_path) + logger.info(f"Saved VI pair plot to {pair_path}") + + return saved + + +def _plot_map_diagnostics( + idata: az.InferenceData, + g1_true: float, + g2_true: float, + output_dir: str, +) -> List[Path]: + """Generate MAP diagnostic plots: bar at point estimate + truth line.""" + output_path = Path(output_dir) + output_path.mkdir(parents=True, exist_ok=True) + saved = [] + + g1_val = float(idata.posterior.g1.values.flatten()[0]) + g2_val = float(idata.posterior.g2.values.flatten()[0]) + + fig, axes = plt.subplots(1, 2, figsize=(10, 4)) + + # g1 MAP estimate + bar_width = max(abs(g1_true) * 0.1, 1e-4) + axes[0].bar(g1_val, 1.0, width=bar_width, alpha=0.7, + color="steelblue", edgecolor="black", label="MAP") + axes[0].axvline(g1_true, color="red", ls="--", lw=2, label="Truth") + axes[0].set(xlabel="g1", title=f"MAP: g1 = {g1_val:.6f} (truth = {g1_true:.4f})") + axes[0].set_ylabel("(point estimate)") + axes[0].legend() + + # g2 MAP estimate + bar_width = max(abs(g2_true) * 0.1, 1e-4) + axes[1].bar(g2_val, 1.0, width=bar_width, alpha=0.7, + color="coral", edgecolor="black", label="MAP") + axes[1].axvline(g2_true, color="red", ls="--", lw=2, label="Truth") + axes[1].set(xlabel="g2", title=f"MAP: g2 = {g2_val:.6f} (truth = {g2_true:.4f})") + axes[1].set_ylabel("(point estimate)") + axes[1].legend() + + fig.tight_layout() + map_path = output_path / "map_estimate.png" + fig.savefig(map_path, dpi=150) + plt.close(fig) + saved.append(map_path) + logger.info(f"Saved MAP estimate plot to {map_path}") + + return saved + + def plot_bias_vs_shear( g_true_values: np.ndarray, g_est_means: np.ndarray, diff --git a/tests/test_config.py b/tests/test_config.py index 31fb913..8cbf5ef 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -15,9 +15,11 @@ InferenceConfig, MAPConfig, NoiseConfig, + NUTSConfig, PSFConfig, ShearConfig, ShineConfig, + VIConfig, ) @@ -178,25 +180,122 @@ def test_map_config_custom_values(self): assert config.learning_rate == 1e-3 -class TestInferenceConfig: - """Test InferenceConfig.""" +class TestNUTSConfig: + """Test NUTSConfig.""" - def test_inference_config_defaults(self): - """Test inference configuration with defaults.""" - config = InferenceConfig() + def test_nuts_config_defaults(self): + """Test NUTS configuration with defaults.""" + config = NUTSConfig() assert config.warmup == 500 assert config.samples == 1000 assert config.chains == 1 assert config.dense_mass is False assert config.map_init is None - def test_inference_config_with_map(self): - """Test inference configuration with MAP initialization.""" + def test_nuts_config_with_map_init(self): + """Test NUTS configuration with MAP initialization.""" map_config = MAPConfig(enabled=True) - config = InferenceConfig(map_init=map_config) + config = NUTSConfig(map_init=map_config) assert config.map_init is not None assert config.map_init.enabled is True + def test_nuts_config_custom_values(self): + """Test NUTS configuration with custom values.""" + config = NUTSConfig(warmup=200, samples=500, chains=4, dense_mass=True) + assert config.warmup == 200 + assert config.samples == 500 + assert config.chains == 4 + assert config.dense_mass is True + + def test_nuts_config_invalid_warmup(self): + """Test that non-positive warmup raises error.""" + with pytest.raises(ValueError, match="warmup must be positive"): + NUTSConfig(warmup=0) + + def test_nuts_config_invalid_chains(self): + """Test that non-positive chains raises error.""" + with pytest.raises(ValueError, match="Number of chains must be positive"): + NUTSConfig(chains=0) + + +class TestVIConfig: + """Test VIConfig.""" + + def test_vi_config_defaults(self): + """Test VI configuration with defaults.""" + config = VIConfig() + assert config.num_steps == 5000 + assert config.learning_rate == 1e-3 + assert config.num_samples == 1000 + + def test_vi_config_custom_values(self): + """Test VI configuration with custom values.""" + config = VIConfig(num_steps=10000, learning_rate=5e-4, num_samples=2000) + assert config.num_steps == 10000 + assert config.learning_rate == 5e-4 + assert config.num_samples == 2000 + + def test_vi_config_invalid_num_steps(self): + """Test that non-positive num_steps raises error.""" + with pytest.raises(ValueError, match="num_steps must be positive"): + VIConfig(num_steps=0) + + def test_vi_config_invalid_learning_rate(self): + """Test that non-positive learning rate raises error.""" + with pytest.raises(ValueError, match="Learning rate must be positive"): + VIConfig(learning_rate=-0.01) + + def test_vi_config_invalid_num_samples(self): + """Test that non-positive num_samples raises error.""" + with pytest.raises(ValueError, match="num_samples must be positive"): + VIConfig(num_samples=0) + + +class TestInferenceConfig: + """Test InferenceConfig.""" + + def test_inference_config_defaults(self): + """Test inference configuration with defaults.""" + config = InferenceConfig() + assert config.method == "nuts" + assert config.nuts_config is None + assert config.map_config is None + assert config.vi_config is None + assert config.rng_seed == 0 + + def test_inference_config_method_nuts(self): + """Test inference configuration with NUTS method.""" + nuts = NUTSConfig(warmup=200, samples=500) + config = InferenceConfig(method="nuts", nuts_config=nuts) + assert config.method == "nuts" + assert config.nuts_config.warmup == 200 + assert config.nuts_config.samples == 500 + + def test_inference_config_method_map(self): + """Test inference configuration with MAP method.""" + map_cfg = MAPConfig(enabled=True, num_steps=2000) + config = InferenceConfig(method="map", map_config=map_cfg) + assert config.method == "map" + assert config.map_config.num_steps == 2000 + + def test_inference_config_method_vi(self): + """Test inference configuration with VI method.""" + vi = VIConfig(num_steps=5000, learning_rate=0.001, num_samples=2000) + config = InferenceConfig(method="vi", vi_config=vi) + assert config.method == "vi" + assert config.vi_config.num_steps == 5000 + assert config.vi_config.num_samples == 2000 + + def test_inference_config_invalid_method(self): + """Test that invalid method raises validation error.""" + with pytest.raises(ValueError): + InferenceConfig(method="invalid") + + def test_inference_config_no_method_defaults_to_nuts(self): + """Test backward compatibility: no method field defaults to nuts.""" + config = InferenceConfig(rng_seed=42) + assert config.method == "nuts" + class TestShineConfig: """Test full ShineConfig.""" @@ -248,7 +347,10 @@ def test_load_valid_config(self): "half_light_radius": 1.0, "shear": {"type": "G1G2", "g1": 0.01, "g2": -0.02}, }, - "inference": {"warmup": 500, "samples": 1000}, + "inference": { + "method": "nuts", + "nuts_config": {"warmup": 500, "samples": 1000}, + }, "output_path": "results", } diff --git a/tests/test_validation/test_extraction.py b/tests/test_validation/test_extraction.py index 702e25d..9bd80a3 100644 --- a/tests/test_validation/test_extraction.py +++ b/tests/test_validation/test_extraction.py @@ -24,8 +24,13 @@ def _make_mock_idata( n_chains=2, n_samples=500, n_divergences=0, + inference_method=None, ): - """Create a mock InferenceData object for testing.""" + """Create a mock InferenceData object for testing. + + Args: + inference_method: If provided, sets the inference_method attr on posterior. + """ rng = np.random.default_rng(42) posterior = { "g1": rng.normal(g1_mean, g1_std, size=(n_chains, n_samples)), @@ -44,7 +49,10 @@ def _make_mock_idata( sample_stats = {"diverging": diverging, "energy": energy} - return az.from_dict(posterior=posterior, sample_stats=sample_stats) + idata = az.from_dict(posterior=posterior, sample_stats=sample_stats) + if inference_method is not None: + idata.posterior.attrs["inference_method"] = inference_method + return idata class TestExtractConvergenceDiagnostics: @@ -197,3 +205,120 @@ def test_full_extraction(self): assert isinstance(result.diagnostics, ConvergenceDiagnostics) assert isinstance(result.passed_convergence, bool) assert result.seed == 42 + + +class TestMethodAwareDiagnostics: + """Tests for method-aware convergence diagnostics.""" + + def test_map_diagnostics_sentinels(self): + """MAP idata (1 chain, 1 draw) returns sentinel diagnostics.""" + idata = _make_mock_idata( + n_chains=1, n_samples=1, inference_method="map" + ) + diag = extract_convergence_diagnostics(idata) + assert diag.rhat == {"g1": 1.0, "g2": 1.0} + assert diag.ess == {"g1": 1.0, "g2": 1.0} + assert diag.divergences == 0 + assert diag.bfmi == [] + assert diag.n_chains == 1 + assert diag.n_samples == 1 + + def test_map_check_convergence_always_true(self): + """MAP method always passes convergence.""" + diag = ConvergenceDiagnostics( + rhat={"g1": 1.0, "g2": 1.0}, + ess={"g1": 1.0, "g2": 1.0}, + divergences=0, + divergence_frac=0.0, + bfmi=[], + n_samples=1, + n_chains=1, + ) + thresholds = ConvergenceThresholds(ess_min=100) + # MAP always returns True regardless of ESS + assert check_convergence(diag, thresholds, method="map") is True + + def test_vi_diagnostics_ess_computed(self): + """VI idata (1 chain, N draws) computes ESS, rhat=1.0.""" + idata = _make_mock_idata( + n_chains=1, n_samples=1000, inference_method="vi" + ) + diag = extract_convergence_diagnostics(idata) + assert diag.rhat == {"g1": 1.0, "g2": 1.0} + assert diag.ess["g1"] > 0 + assert diag.ess["g2"] > 0 + assert diag.divergences == 0 + assert diag.bfmi == [] + assert diag.n_chains == 1 + + def test_vi_check_convergence_only_checks_ess(self): + """VI convergence only checks ESS, not rhat/divergences/bfmi.""" + diag = ConvergenceDiagnostics( + rhat={"g1": 2.0, "g2": 2.0}, # Would fail for NUTS + ess={"g1": 500, "g2": 500}, + divergences=100, # Would fail for NUTS + divergence_frac=0.1, + bfmi=[0.01], # Would fail for NUTS + n_samples=1000, + n_chains=1, + ) + thresholds = ConvergenceThresholds() + # VI only checks ESS, so this should pass + assert check_convergence(diag, thresholds, method="vi") is True + + def test_vi_check_convergence_fails_low_ess(self): + """VI convergence fails when ESS is too low.""" + diag = ConvergenceDiagnostics( + rhat={"g1": 1.0, "g2": 1.0}, + ess={"g1": 10, "g2": 500}, + divergences=0, + divergence_frac=0.0, + bfmi=[], + n_samples=1000, + n_chains=1, + ) + thresholds = ConvergenceThresholds(ess_min=100) + assert check_convergence(diag, thresholds, method="vi") is False + + def test_nuts_default_method(self): + """Without inference_method attr, defaults to NUTS behavior.""" + idata = _make_mock_idata(n_chains=2, n_samples=500) + diag = extract_convergence_diagnostics(idata) + # Should compute full MCMC diagnostics + assert diag.rhat["g1"] < 1.1 + assert diag.ess["g1"] > 0 + assert isinstance(diag.bfmi, list) + assert len(diag.bfmi) == 2 + + def test_extract_shear_estimates_map(self): + """extract_shear_estimates works for MAP (1 chain, 1 draw).""" + idata = _make_mock_idata( + g1_mean=0.01, n_chains=1, n_samples=1, inference_method="map" + ) + est = extract_shear_estimates(idata, "g1") + assert isinstance(est, ShearEstimates) + # For a single sample, mean == median == the value + assert est.mean == est.median + assert est.std == 0.0 + + def test_extract_shear_estimates_vi(self): + """extract_shear_estimates works for VI (1 chain, N draws).""" + idata = _make_mock_idata( + g1_mean=0.02, n_chains=1, n_samples=1000, inference_method="vi" + ) + est = extract_shear_estimates(idata, "g1") + assert isinstance(est, ShearEstimates) + assert est.mean == pytest.approx(0.02, abs=0.005) + + def test_extract_realization_map(self): + """extract_realization works end-to-end for MAP.""" + idata = _make_mock_idata( + g1_mean=0.01, g2_mean=0.0, + n_chains=1, n_samples=1, inference_method="map" + ) + thresholds = ConvergenceThresholds() + result = extract_realization( + idata, g1_true=0.01, g2_true=0.0, + run_id="map_test", seed=0, thresholds=thresholds, + ) + assert result.passed_convergence is True From 9115cc05279aca7cacda8d3daa2592ef95d2bd0a Mon Sep 17 00:00:00 2001 From: EiffL Date: Sun, 8 Feb 2026 11:31:33 +0100 Subject: [PATCH 2/5] adjusting setting --- configs/validation/level0_base.yaml | 4 ++-- examples/level0_batched_inference.ipynb | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/configs/validation/level0_base.yaml b/configs/validation/level0_base.yaml index 3450035..8fe8b1a 100644 --- a/configs/validation/level0_base.yaml +++ b/configs/validation/level0_base.yaml @@ -46,6 +46,6 @@ gal: inference: method: map # MAP is sufficient for noiseless Level 0 map_config: - num_steps: 1000 - learning_rate: 0.01 + num_steps: 3000 + learning_rate: 0.005 rng_seed: 42 diff --git a/examples/level0_batched_inference.ipynb b/examples/level0_batched_inference.ipynb index 8d8a74f..00c828b 100644 --- a/examples/level0_batched_inference.ipynb +++ b/examples/level0_batched_inference.ipynb @@ -35,7 +35,7 @@ "execution_count": null, "metadata": {}, "outputs": [], - "source": "# Ground truth shear (matches metacal shear_true=[0.01, 0.00])\nG1_TRUE = 0.01\nG2_TRUE = 0.00\nN_BATCH = 10\n\nconfig = ShineConfig(\n image=ImageConfig(\n pixel_scale=0.263, # arcsec/pixel (metacal scale=0.263)\n size_x=48,\n size_y=48,\n n_objects=1,\n fft_size=128,\n noise=NoiseConfig(type=\"Gaussian\", sigma=1e-6), # metacal noise=1e-6\n ),\n psf=PSFConfig(\n type=\"Moffat\",\n sigma=0.9, # FWHM in arcsec (metacal psf_fwhm=0.9)\n beta=2.5, # metacal beta=2.5\n ),\n gal=GalaxyConfig(\n type=\"Exponential\", # metacal galsim.Exponential\n flux=1.0, # metacal default flux=1\n half_light_radius=0.5, # arcsec (metacal gal_hlr=0.5)\n ellipticity=EllipticityConfig(type=\"E1E2\", e1=0.0, e2=0.0),\n shear=ShearConfig(\n type=\"G1G2\",\n g1=DistributionConfig(type=\"Normal\", mean=0.0, sigma=0.05),\n g2=DistributionConfig(type=\"Normal\", mean=0.0, sigma=0.05),\n ),\n position=PositionConfig(\n type=\"Uniform\",\n x_min=23.5, x_max=24.5,\n y_min=23.5, y_max=24.5,\n ),\n ),\n inference=InferenceConfig(\n method=\"map\",\n map_config=MAPConfig(num_steps=1000, learning_rate=0.01),\n rng_seed=42,\n ),\n)\n\nprint(f\"Image: {config.image.size_x}x{config.image.size_y} px, \"\n f\"scale={config.image.pixel_scale}\\\"/px\")\nprint(f\"Galaxy: {config.gal.type}, flux={config.gal.flux}, \"\n f\"hlr={config.gal.half_light_radius}\\\"\")\nprint(f\"PSF: {config.psf.type}, FWHM={config.psf.sigma}\\\", beta={config.psf.beta}\")\nprint(f\"Noise sigma: {config.image.noise.sigma}\")\nprint(f\"True shear: g1={G1_TRUE}, g2={G2_TRUE}\")\nprint(f\"Batch size: {N_BATCH}\")\nprint(f\"Inference method: {config.inference.method}\")" + "source": "# Ground truth shear (matches metacal shear_true=[0.01, 0.00])\nG1_TRUE = 0.01\nG2_TRUE = 0.00\nN_BATCH = 10\n\nconfig = ShineConfig(\n image=ImageConfig(\n pixel_scale=0.263, # arcsec/pixel (metacal scale=0.263)\n size_x=48,\n size_y=48,\n n_objects=1,\n fft_size=128,\n noise=NoiseConfig(type=\"Gaussian\", sigma=1e-6), # metacal noise=1e-6\n ),\n psf=PSFConfig(\n type=\"Moffat\",\n sigma=0.9, # FWHM in arcsec (metacal psf_fwhm=0.9)\n beta=2.5, # metacal beta=2.5\n ),\n gal=GalaxyConfig(\n type=\"Exponential\", # metacal galsim.Exponential\n flux=1.0, # metacal default flux=1\n half_light_radius=0.5, # arcsec (metacal gal_hlr=0.5)\n ellipticity=EllipticityConfig(type=\"E1E2\", e1=0.0, e2=0.0),\n shear=ShearConfig(\n type=\"G1G2\",\n g1=DistributionConfig(type=\"Normal\", mean=0.0, sigma=0.05),\n g2=DistributionConfig(type=\"Normal\", mean=0.0, sigma=0.05),\n ),\n position=PositionConfig(\n type=\"Uniform\",\n x_min=23.5, x_max=24.5,\n y_min=23.5, y_max=24.5,\n ),\n ),\n inference=InferenceConfig(\n method=\"map\",\n map_config=MAPConfig(num_steps=3000, learning_rate=0.005),\n rng_seed=42,\n ),\n)\n\nprint(f\"Image: {config.image.size_x}x{config.image.size_y} px, \"\n f\"scale={config.image.pixel_scale}\\\"/px\")\nprint(f\"Galaxy: {config.gal.type}, flux={config.gal.flux}, \"\n f\"hlr={config.gal.half_light_radius}\\\"\")\nprint(f\"PSF: {config.psf.type}, FWHM={config.psf.sigma}\\\", beta={config.psf.beta}\")\nprint(f\"Noise sigma: {config.image.noise.sigma}\")\nprint(f\"True shear: g1={G1_TRUE}, g2={G2_TRUE}\")\nprint(f\"Batch size: {N_BATCH}\")\nprint(f\"Inference method: {config.inference.method}\")" }, { "cell_type": "markdown", From 870e5a37bac8defe75277113928817de29dd0b57 Mon Sep 17 00:00:00 2001 From: EiffL Date: Sun, 8 Feb 2026 12:14:46 +0100 Subject: [PATCH 3/5] fixing inference in batch --- configs/validation/level0_base.yaml | 4 ++-- shine/validation/extraction.py | 2 ++ 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/configs/validation/level0_base.yaml b/configs/validation/level0_base.yaml index 8fe8b1a..c973200 100644 --- a/configs/validation/level0_base.yaml +++ b/configs/validation/level0_base.yaml @@ -46,6 +46,6 @@ gal: inference: method: map # MAP is sufficient for noiseless Level 0 map_config: - num_steps: 3000 - learning_rate: 0.005 + num_steps: 100 + learning_rate: 0.1 rng_seed: 42 diff --git a/shine/validation/extraction.py b/shine/validation/extraction.py index 45afdd5..43e3c02 100644 --- a/shine/validation/extraction.py +++ b/shine/validation/extraction.py @@ -379,6 +379,8 @@ def split_batched_idata( posterior=post_dict, sample_stats=stats_dict, ) + # Preserve posterior attributes (e.g., inference_method) + single_idata.posterior.attrs.update(posterior.attrs) results.append((run_ids[i], single_idata)) logger.info(f"Split batched InferenceData into {n_batch} per-realization objects") From 28a5436547e8053037817d2900d321f707bd1336 Mon Sep 17 00:00:00 2001 From: EiffL Date: Sun, 8 Feb 2026 12:41:13 +0100 Subject: [PATCH 4/5] adding example notebook --- examples/level0_batched_inference.ipynb | 420 ++++++++++++++++++++++-- 1 file changed, 399 insertions(+), 21 deletions(-) diff --git a/examples/level0_batched_inference.ipynb b/examples/level0_batched_inference.ipynb index 00c828b..40cbb3b 100644 --- a/examples/level0_batched_inference.ipynb +++ b/examples/level0_batched_inference.ipynb @@ -3,14 +3,67 @@ { "cell_type": "markdown", "metadata": {}, - "source": "# Level 0: Batched MAP Shear Inference\n\nThis notebook demonstrates SHINE's **Level 0 sanity check** on a batch of 10 galaxies\nsharing the same true shear. Level 0 is the noiseless self-consistency test: when data\nis generated by the exact same forward model with effectively zero noise, the MAP\nestimate should recover the truth exactly.\n\nSince there is no noise, MAP (point estimation) is the natural choice -- full MCMC\nis unnecessary and much slower.\n\nSimulation parameters are matched to the\n[ngmix metacal example](https://github.com/esheldon/ngmix/blob/master/examples/metacal/metacal.py):\n\n| Parameter | Value | Source |\n|-----------|-------|--------|\n| Galaxy | Exponential, hlr=0.5\" | metacal `gal_hlr=0.5` |\n| PSF | Moffat, $\\beta$=2.5, FWHM=0.9\" | metacal `psf_fwhm=0.9` |\n| Pixel scale | 0.263\"/px | metacal `scale=0.263` |\n| Noise | $\\sigma = 10^{-6}$ | metacal `noise=1e-6` |\n| Shear | $g_1=0.01$, $g_2=0.00$ | metacal `shear_true=[0.01, 0.00]` |\n\n**What we do:**\n1. Generate 10 synthetic observations with the same shear\n2. Run MAP inference on each realization independently\n3. Check that MAP estimates match truth with negligible bias" + "source": [ + "# Level 0: Batched MAP Shear Inference\n", + "\n", + "This notebook demonstrates SHINE's **Level 0 sanity check** on a batch of 10 galaxies\n", + "sharing the same true shear. Level 0 is the noiseless self-consistency test: when data\n", + "is generated by the exact same forward model with effectively zero noise, the MAP\n", + "estimate should recover the truth exactly.\n", + "\n", + "Since there is no noise, MAP (point estimation) is the natural choice -- full MCMC\n", + "is unnecessary and much slower.\n", + "\n", + "**Configuration:** Exponential galaxy (hlr=0.5\"), Moffat PSF ($\\beta$=2.5, FWHM=0.9\"),\n", + "pixel scale 0.263\"/px, noise $\\sigma = 10^{-6}$, true shear $g_1=0.01$, $g_2=0.00$.\n", + "\n", + "**What we do:**\n", + "1. Generate 10 synthetic observations with the same shear\n", + "2. Run MAP inference on each realization independently\n", + "3. Check that MAP estimates match truth with negligible bias" + ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], - "source": "import jax\nimport jax.numpy as jnp\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport arviz as az\n\nfrom shine.config import (\n ShineConfig,\n ImageConfig,\n NoiseConfig,\n PSFConfig,\n GalaxyConfig,\n ShearConfig,\n EllipticityConfig,\n PositionConfig,\n InferenceConfig,\n MAPConfig,\n DistributionConfig,\n)\nfrom shine.scene import SceneBuilder\nfrom shine.inference import Inference\nfrom shine.validation.simulation import generate_biased_observation\nfrom shine.validation.extraction import (\n extract_convergence_diagnostics,\n extract_shear_estimates,\n check_convergence,\n)\nfrom shine.validation.bias_config import ConvergenceThresholds\nfrom shine.validation.statistics import compute_bias_single_point\n\n# Use 64-bit precision for accurate shear recovery\njax.config.update(\"jax_enable_x64\", True)\n\nprint(f\"JAX devices: {jax.devices()}\")" + "source": [ + "import jax\n", + "import jax.numpy as jnp\n", + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "import arviz as az\n", + "\n", + "from shine.config import (\n", + " ShineConfig,\n", + " ImageConfig,\n", + " NoiseConfig,\n", + " PSFConfig,\n", + " GalaxyConfig,\n", + " ShearConfig,\n", + " EllipticityConfig,\n", + " PositionConfig,\n", + " InferenceConfig,\n", + " MAPConfig,\n", + " DistributionConfig,\n", + ")\n", + "from shine.scene import SceneBuilder\n", + "from shine.inference import Inference\n", + "from shine.validation.simulation import generate_biased_observation\n", + "from shine.validation.extraction import (\n", + " extract_convergence_diagnostics,\n", + " extract_shear_estimates,\n", + " check_convergence,\n", + ")\n", + "from shine.validation.bias_config import ConvergenceThresholds\n", + "from shine.validation.statistics import compute_bias_single_point\n", + "\n", + "# Use 64-bit precision for accurate shear recovery\n", + "jax.config.update(\"jax_enable_x64\", True)\n", + "\n", + "print(f\"JAX devices: {jax.devices()}\")" + ] }, { "cell_type": "markdown", @@ -35,74 +88,269 @@ "execution_count": null, "metadata": {}, "outputs": [], - "source": "# Ground truth shear (matches metacal shear_true=[0.01, 0.00])\nG1_TRUE = 0.01\nG2_TRUE = 0.00\nN_BATCH = 10\n\nconfig = ShineConfig(\n image=ImageConfig(\n pixel_scale=0.263, # arcsec/pixel (metacal scale=0.263)\n size_x=48,\n size_y=48,\n n_objects=1,\n fft_size=128,\n noise=NoiseConfig(type=\"Gaussian\", sigma=1e-6), # metacal noise=1e-6\n ),\n psf=PSFConfig(\n type=\"Moffat\",\n sigma=0.9, # FWHM in arcsec (metacal psf_fwhm=0.9)\n beta=2.5, # metacal beta=2.5\n ),\n gal=GalaxyConfig(\n type=\"Exponential\", # metacal galsim.Exponential\n flux=1.0, # metacal default flux=1\n half_light_radius=0.5, # arcsec (metacal gal_hlr=0.5)\n ellipticity=EllipticityConfig(type=\"E1E2\", e1=0.0, e2=0.0),\n shear=ShearConfig(\n type=\"G1G2\",\n g1=DistributionConfig(type=\"Normal\", mean=0.0, sigma=0.05),\n g2=DistributionConfig(type=\"Normal\", mean=0.0, sigma=0.05),\n ),\n position=PositionConfig(\n type=\"Uniform\",\n x_min=23.5, x_max=24.5,\n y_min=23.5, y_max=24.5,\n ),\n ),\n inference=InferenceConfig(\n method=\"map\",\n map_config=MAPConfig(num_steps=3000, learning_rate=0.005),\n rng_seed=42,\n ),\n)\n\nprint(f\"Image: {config.image.size_x}x{config.image.size_y} px, \"\n f\"scale={config.image.pixel_scale}\\\"/px\")\nprint(f\"Galaxy: {config.gal.type}, flux={config.gal.flux}, \"\n f\"hlr={config.gal.half_light_radius}\\\"\")\nprint(f\"PSF: {config.psf.type}, FWHM={config.psf.sigma}\\\", beta={config.psf.beta}\")\nprint(f\"Noise sigma: {config.image.noise.sigma}\")\nprint(f\"True shear: g1={G1_TRUE}, g2={G2_TRUE}\")\nprint(f\"Batch size: {N_BATCH}\")\nprint(f\"Inference method: {config.inference.method}\")" + "source": [ + "# Ground truth shear (matches metacal shear_true=[0.01, 0.00])\n", + "G1_TRUE = 0.01\n", + "G2_TRUE = 0.00\n", + "N_BATCH = 10\n", + "\n", + "config = ShineConfig(\n", + " image=ImageConfig(\n", + " pixel_scale=0.263, # arcsec/pixel (metacal scale=0.263)\n", + " size_x=48,\n", + " size_y=48,\n", + " n_objects=1,\n", + " fft_size=128,\n", + " noise=NoiseConfig(type=\"Gaussian\", sigma=1e-6), # metacal noise=1e-6\n", + " ),\n", + " psf=PSFConfig(\n", + " type=\"Moffat\",\n", + " sigma=0.9, # FWHM in arcsec (metacal psf_fwhm=0.9)\n", + " beta=2.5, # metacal beta=2.5\n", + " ),\n", + " gal=GalaxyConfig(\n", + " type=\"Exponential\", # metacal galsim.Exponential\n", + " flux=1.0, # metacal default flux=1\n", + " half_light_radius=0.5, # arcsec (metacal gal_hlr=0.5)\n", + " ellipticity=EllipticityConfig(type=\"E1E2\", e1=0.0, e2=0.0),\n", + " shear=ShearConfig(\n", + " type=\"G1G2\",\n", + " g1=DistributionConfig(type=\"Normal\", mean=0.0, sigma=0.05),\n", + " g2=DistributionConfig(type=\"Normal\", mean=0.0, sigma=0.05),\n", + " ),\n", + " position=PositionConfig(\n", + " type=\"Uniform\",\n", + " x_min=23.5, x_max=24.5,\n", + " y_min=23.5, y_max=24.5,\n", + " ),\n", + " ),\n", + " inference=InferenceConfig(\n", + " method=\"map\",\n", + " map_config=MAPConfig(num_steps=200, learning_rate=0.1),\n", + " rng_seed=42,\n", + " ),\n", + ")\n", + "\n", + "print(f\"Image: {config.image.size_x}x{config.image.size_y} px, \"\n", + " f\"scale={config.image.pixel_scale}\\\"/px\")\n", + "print(f\"Galaxy: {config.gal.type}, flux={config.gal.flux}, \"\n", + " f\"hlr={config.gal.half_light_radius}\\\"\")\n", + "print(f\"PSF: {config.psf.type}, FWHM={config.psf.sigma}\\\", beta={config.psf.beta}\")\n", + "print(f\"Noise sigma: {config.image.noise.sigma}\")\n", + "print(f\"True shear: g1={G1_TRUE}, g2={G2_TRUE}\")\n", + "print(f\"Batch size: {N_BATCH}\")\n", + "print(f\"Inference method: {config.inference.method}\")" + ] }, { "cell_type": "markdown", "metadata": {}, - "source": "## 2. Generate Observations\n\nAll 10 galaxies share the same true shear but have independent (effectively zero) noise\nrealizations. We use `generate_biased_observation()` per realization." + "source": [ + "## 2. Generate Observations\n", + "\n", + "All 10 galaxies share the same true shear but have independent (effectively zero) noise\n", + "realizations. We use `generate_biased_observation()` per realization." + ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], - "source": "# Generate N_BATCH independent observations with the same true shear\nseeds = list(range(100, 100 + N_BATCH))\nsim_results = []\nfor seed in seeds:\n sim = generate_biased_observation(config, G1_TRUE, G2_TRUE, seed)\n sim_results.append(sim)\n\nprint(f\"Generated {N_BATCH} observations\")\nprint(f\"Image shape: {sim_results[0].observation.image.shape}\")\nprint(f\"PSF type: {type(sim_results[0].observation.psf_model).__name__}\")" + "source": [ + "# Generate N_BATCH independent observations with the same true shear\n", + "seeds = list(range(100, 100 + N_BATCH))\n", + "sim_results = []\n", + "for seed in seeds:\n", + " sim = generate_biased_observation(config, G1_TRUE, G2_TRUE, seed)\n", + " sim_results.append(sim)\n", + "\n", + "print(f\"Generated {N_BATCH} observations\")\n", + "print(f\"Image shape: {sim_results[0].observation.image.shape}\")\n", + "print(f\"PSF type: {type(sim_results[0].observation.psf_model).__name__}\")" + ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], - "source": "# Visualize a few of the generated images\nfig, axes = plt.subplots(2, 5, figsize=(15, 6))\nfor i, ax in enumerate(axes.flat):\n im = ax.imshow(sim_results[i].observation.image, origin=\"lower\", cmap=\"gray_r\")\n ax.set_title(f\"Galaxy {i}\", fontsize=10)\n ax.set_xticks([])\n ax.set_yticks([])\nfig.suptitle(\n f\"Level 0: 10 Exponential galaxies, g1={G1_TRUE}, g2={G2_TRUE}, \"\n f\"noise={config.image.noise.sigma}\",\n fontsize=12,\n)\nplt.tight_layout()\nplt.show()" + "source": [ + "# Visualize a few of the generated images\n", + "fig, axes = plt.subplots(2, 5, figsize=(15, 6))\n", + "for i, ax in enumerate(axes.flat):\n", + " im = ax.imshow(sim_results[i].observation.image, origin=\"lower\", cmap=\"gray_r\")\n", + " ax.set_title(f\"Galaxy {i}\", fontsize=10)\n", + " ax.set_xticks([])\n", + " ax.set_yticks([])\n", + "fig.suptitle(\n", + " f\"Level 0: 10 Exponential galaxies, g1={G1_TRUE}, g2={G2_TRUE}, \"\n", + " f\"noise={config.image.noise.sigma}\",\n", + " fontsize=12,\n", + ")\n", + "plt.tight_layout()\n", + "plt.show()" + ] }, { "cell_type": "markdown", "metadata": {}, - "source": "## 3. Run MAP Inference\n\nFor each realization, we build a model and run MAP estimation. Since Level 0\nis noiseless, MAP finds the maximum a posteriori point estimate which should\nmatch the truth exactly." + "source": [ + "## 3. Run MAP Inference\n", + "\n", + "For each realization, we build a model and run MAP estimation. Since Level 0\n", + "is noiseless, MAP finds the maximum a posteriori point estimate which should\n", + "match the truth exactly." + ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], - "source": "# Run MAP inference for each realization\nscene = SceneBuilder(config)\nmodel_fn = scene.build_model()\n\nmap_cfg = config.inference.map_config\nprint(f\"Inference method: {config.inference.method}\")\nprint(f\"MAP: {map_cfg.num_steps} steps, lr={map_cfg.learning_rate}\")\n\nidata_list = []\nfor i, sim in enumerate(sim_results):\n rng_key = jax.random.PRNGKey(config.inference.rng_seed + i)\n engine = Inference(model=model_fn, config=config.inference)\n idata = engine.run(\n rng_key=rng_key,\n observed_data=sim.observation.image,\n extra_args={\"psf\": sim.observation.psf_model},\n )\n idata_list.append(idata)\n g1_val = float(idata.posterior.g1.values.flatten()[0])\n g2_val = float(idata.posterior.g2.values.flatten()[0])\n print(f\" Realization {i}: g1={g1_val:+.6f}, g2={g2_val:+.6f}\")\n\nprint(f\"\\nCompleted {N_BATCH} MAP estimations\")" + "source": [ + "# Run MAP inference for each realization\n", + "scene = SceneBuilder(config)\n", + "model_fn = scene.build_model()\n", + "\n", + "map_cfg = config.inference.map_config\n", + "print(f\"Inference method: {config.inference.method}\")\n", + "print(f\"MAP: {map_cfg.num_steps} steps, lr={map_cfg.learning_rate}\")\n", + "\n", + "idata_list = []\n", + "for i, sim in enumerate(sim_results):\n", + " rng_key = jax.random.PRNGKey(config.inference.rng_seed + i)\n", + " engine = Inference(model=model_fn, config=config.inference)\n", + " idata = engine.run(\n", + " rng_key=rng_key,\n", + " observed_data=sim.observation.image,\n", + " extra_args={\"psf\": sim.observation.psf_model},\n", + " )\n", + " idata_list.append(idata)\n", + " g1_val = float(idata.posterior.g1.values.flatten()[0])\n", + " g2_val = float(idata.posterior.g2.values.flatten()[0])\n", + " print(f\" Realization {i}: g1={g1_val:+.6f}, g2={g2_val:+.6f}\")\n", + "\n", + "print(f\"\\nCompleted {N_BATCH} MAP estimations\")" + ] }, { "cell_type": "markdown", "metadata": {}, - "source": "## 4. Extract Estimates\n\nFor each realization, extract the MAP point estimates and check that they\nmatch the truth." + "source": [ + "## 4. Extract Estimates\n", + "\n", + "For each realization, extract the MAP point estimates and check that they\n", + "match the truth." + ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], - "source": "run_ids = [f\"level0_{i:04d}\" for i in range(N_BATCH)]\n\nprint(f\"Extracted {N_BATCH} MAP estimates\")\nprint(f\"Example: {run_ids[0]}, posterior vars: {list(idata_list[0].posterior.data_vars)}\")\nprint(f\" inference_method: {idata_list[0].posterior.attrs.get('inference_method')}\")" + "source": [ + "run_ids = [f\"level0_{i:04d}\" for i in range(N_BATCH)]\n", + "\n", + "print(f\"Extracted {N_BATCH} MAP estimates\")\n", + "print(f\"Example: {run_ids[0]}, posterior vars: {list(idata_list[0].posterior.data_vars)}\")\n", + "print(f\" inference_method: {idata_list[0].posterior.attrs.get('inference_method')}\")" + ] }, { "cell_type": "markdown", "metadata": {}, - "source": "## 5. Extract Diagnostics\n\nFor each realization, extract:\n- **Shear estimates**: MAP point estimate (mean = median = value, std = 0)\n- **Convergence diagnostics**: sentinel values for MAP (rhat=1, ess=1)\n\nLevel 0 acceptance criterion for MAP: the estimate should be very close to truth." + "source": [ + "## 5. Extract Diagnostics\n", + "\n", + "For each realization, extract:\n", + "- **Shear estimates**: MAP point estimate (mean = median = value, std = 0)\n", + "- **Convergence diagnostics**: sentinel values for MAP (rhat=1, ess=1)\n", + "\n", + "Level 0 acceptance criterion for MAP: the estimate should be very close to truth." + ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], - "source": "# Extract results for all realizations\nresults = []\nfor run_id, single_idata in zip(run_ids, idata_list):\n g1_est = extract_shear_estimates(single_idata, \"g1\")\n g2_est = extract_shear_estimates(single_idata, \"g2\")\n diag = extract_convergence_diagnostics(single_idata)\n method = single_idata.posterior.attrs.get(\"inference_method\", \"nuts\")\n passed = check_convergence(diag, ConvergenceThresholds(), method=method)\n results.append({\n \"run_id\": run_id,\n \"g1_est\": g1_est,\n \"g2_est\": g2_est,\n \"diagnostics\": diag,\n \"passed\": passed,\n })\n\n# Summary table\nprint(f\"{'Run ID':<14} {'g1 estimate':>14} {'g2 estimate':>14} {'Pass':>5}\")\nprint(\"-\" * 55)\nfor r in results:\n print(\n f\"{r['run_id']:<14} \"\n f\"{r['g1_est'].mean:>14.6f} \"\n f\"{r['g2_est'].mean:>14.6f} \"\n f\"{'OK' if r['passed'] else 'FAIL':>5}\"\n )" + "source": [ + "# Extract results for all realizations\n", + "results = []\n", + "for run_id, single_idata in zip(run_ids, idata_list):\n", + " g1_est = extract_shear_estimates(single_idata, \"g1\")\n", + " g2_est = extract_shear_estimates(single_idata, \"g2\")\n", + " diag = extract_convergence_diagnostics(single_idata)\n", + " method = single_idata.posterior.attrs.get(\"inference_method\", \"nuts\")\n", + " passed = check_convergence(diag, ConvergenceThresholds(), method=method)\n", + " results.append({\n", + " \"run_id\": run_id,\n", + " \"g1_est\": g1_est,\n", + " \"g2_est\": g2_est,\n", + " \"diagnostics\": diag,\n", + " \"passed\": passed,\n", + " })\n", + "\n", + "# Summary table\n", + "print(f\"{'Run ID':<14} {'g1 estimate':>14} {'g2 estimate':>14} {'Pass':>5}\")\n", + "print(\"-\" * 55)\n", + "for r in results:\n", + " print(\n", + " f\"{r['run_id']:<14} \"\n", + " f\"{r['g1_est'].mean:>14.6f} \"\n", + " f\"{r['g2_est'].mean:>14.6f} \"\n", + " f\"{'OK' if r['passed'] else 'FAIL':>5}\"\n", + " )" + ] }, { "cell_type": "markdown", "metadata": {}, - "source": "## 6. Acceptance Criteria Check\n\nFor Level 0 MAP, each realization must satisfy:\n1. MAP estimate close to truth (absolute offset $< 10^{-3}$)\n2. Convergence always passes for MAP (no sampling diagnostics)" + "source": [ + "## 6. Acceptance Criteria Check\n", + "\n", + "For Level 0 MAP, each realization must satisfy:\n", + "1. MAP estimate close to truth (absolute offset $< 10^{-3}$)\n", + "2. Convergence always passes for MAP (no sampling diagnostics)" + ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], - "source": "MAX_ABS_OFFSET = 1e-3\n\nall_passed = True\nprint(\"Level 0 Acceptance Criteria (MAP)\")\nprint(\"=\" * 80)\n\nfor r in results:\n run_id = r[\"run_id\"]\n g1, g2 = r[\"g1_est\"], r[\"g2_est\"]\n\n # For MAP: check absolute offset from truth\n g1_offset = abs(g1.mean - G1_TRUE)\n g2_offset = abs(g2.mean - G2_TRUE)\n offset_ok = g1_offset < MAX_ABS_OFFSET and g2_offset < MAX_ABS_OFFSET\n\n passed = offset_ok and r[\"passed\"]\n all_passed = all_passed and passed\n\n status = \"PASS\" if passed else \"FAIL\"\n print(f\"\\n{run_id} [{status}]\")\n print(f\" g1: truth={G1_TRUE:+.4f} MAP={g1.mean:+.6f} \"\n f\"|offset|={g1_offset:.2e} {'ok' if g1_offset < MAX_ABS_OFFSET else 'FAIL'}\")\n print(f\" g2: truth={G2_TRUE:+.4f} MAP={g2.mean:+.6f} \"\n f\"|offset|={g2_offset:.2e} {'ok' if g2_offset < MAX_ABS_OFFSET else 'FAIL'}\")\n\nprint(\"\\n\" + \"=\" * 80)\nprint(f\"Overall Level 0 result: {'ALL PASSED' if all_passed else 'SOME FAILED'}\")\nprint(f\" {sum(1 for r in results if r['passed'])}/{len(results)} \"\n f\"realizations passed\")" + "source": [ + "MAX_ABS_OFFSET = 1e-3\n", + "\n", + "all_passed = True\n", + "print(\"Level 0 Acceptance Criteria (MAP)\")\n", + "print(\"=\" * 80)\n", + "\n", + "for r in results:\n", + " run_id = r[\"run_id\"]\n", + " g1, g2 = r[\"g1_est\"], r[\"g2_est\"]\n", + "\n", + " # For MAP: check absolute offset from truth\n", + " g1_offset = abs(g1.mean - G1_TRUE)\n", + " g2_offset = abs(g2.mean - G2_TRUE)\n", + " offset_ok = g1_offset < MAX_ABS_OFFSET and g2_offset < MAX_ABS_OFFSET\n", + "\n", + " passed = offset_ok and r[\"passed\"]\n", + " all_passed = all_passed and passed\n", + "\n", + " status = \"PASS\" if passed else \"FAIL\"\n", + " print(f\"\\n{run_id} [{status}]\")\n", + " print(f\" g1: truth={G1_TRUE:+.4f} MAP={g1.mean:+.6f} \"\n", + " f\"|offset|={g1_offset:.2e} {'ok' if g1_offset < MAX_ABS_OFFSET else 'FAIL'}\")\n", + " print(f\" g2: truth={G2_TRUE:+.4f} MAP={g2.mean:+.6f} \"\n", + " f\"|offset|={g2_offset:.2e} {'ok' if g2_offset < MAX_ABS_OFFSET else 'FAIL'}\")\n", + "\n", + "print(\"\\n\" + \"=\" * 80)\n", + "print(f\"Overall Level 0 result: {'ALL PASSED' if all_passed else 'SOME FAILED'}\")\n", + "print(f\" {sum(1 for r in results if r['passed'])}/{len(results)} \"\n", + " f\"realizations passed\")" + ] }, { "cell_type": "markdown", @@ -124,7 +372,25 @@ "execution_count": null, "metadata": {}, "outputs": [], - "source": "bias_g1_list = []\n\nprint(f\"{'Run ID':<14} {'m(g1)':>12} {'c(g2)':>12}\")\nprint(\"-\" * 42)\n\nfor r in results:\n b1 = compute_bias_single_point(G1_TRUE, r[\"g1_est\"].mean, r[\"g1_est\"].std, \"g1\")\n bias_g1_list.append(b1)\n # g2_true=0 so we report additive residual instead of multiplicative bias\n c2 = r[\"g2_est\"].mean - G2_TRUE\n print(f\"{r['run_id']:<14} {b1.m:>12.6f} {c2:>12.2e}\")\n\n# Ensemble average\nm_g1_vals = np.array([b.m for b in bias_g1_list])\nc_g2_vals = np.array([r[\"g2_est\"].mean - G2_TRUE for r in results])\nprint(\"-\" * 42)\nprint(f\"{'Ensemble mean':<14} {m_g1_vals.mean():>12.6f} {c_g2_vals.mean():>12.2e}\")" + "source": [ + "bias_g1_list = []\n", + "\n", + "print(f\"{'Run ID':<14} {'m(g1)':>12} {'c(g2)':>12}\")\n", + "print(\"-\" * 42)\n", + "\n", + "for r in results:\n", + " b1 = compute_bias_single_point(G1_TRUE, r[\"g1_est\"].mean, r[\"g1_est\"].std, \"g1\")\n", + " bias_g1_list.append(b1)\n", + " # g2_true=0 so we report additive residual instead of multiplicative bias\n", + " c2 = r[\"g2_est\"].mean - G2_TRUE\n", + " print(f\"{r['run_id']:<14} {b1.m:>12.6f} {c2:>12.2e}\")\n", + "\n", + "# Ensemble average\n", + "m_g1_vals = np.array([b.m for b in bias_g1_list])\n", + "c_g2_vals = np.array([r[\"g2_est\"].mean - G2_TRUE for r in results])\n", + "print(\"-\" * 42)\n", + "print(f\"{'Ensemble mean':<14} {m_g1_vals.mean():>12.6f} {c_g2_vals.mean():>12.2e}\")" + ] }, { "cell_type": "markdown", @@ -138,33 +404,145 @@ "execution_count": null, "metadata": {}, "outputs": [], - "source": "# Collect all g1, g2 MAP estimates\ng1_means = np.array([r[\"g1_est\"].mean for r in results])\ng2_means = np.array([r[\"g2_est\"].mean for r in results])\nindices = np.arange(N_BATCH)\n\nfig, axes = plt.subplots(1, 2, figsize=(14, 5))\n\n# g1 estimates\naxes[0].scatter(indices, g1_means, marker=\"o\", s=60, color=\"steelblue\",\n zorder=5, label=\"MAP estimate\")\naxes[0].axhline(G1_TRUE, color=\"red\", ls=\"--\", lw=2, label=f\"Truth = {G1_TRUE}\")\naxes[0].fill_between([-0.5, N_BATCH - 0.5],\n G1_TRUE - MAX_ABS_OFFSET,\n G1_TRUE + MAX_ABS_OFFSET,\n color=\"red\", alpha=0.1, label=f\"$\\\\pm${MAX_ABS_OFFSET}\")\naxes[0].set_xlabel(\"Realization\")\naxes[0].set_ylabel(\"$g_1$\")\naxes[0].set_title(\"$g_1$ Recovery (MAP)\")\naxes[0].legend(fontsize=9)\naxes[0].set_xticks(indices)\n\n# g2 estimates\naxes[1].scatter(indices, g2_means, marker=\"o\", s=60, color=\"coral\",\n zorder=5, label=\"MAP estimate\")\naxes[1].axhline(G2_TRUE, color=\"red\", ls=\"--\", lw=2, label=f\"Truth = {G2_TRUE}\")\naxes[1].fill_between([-0.5, N_BATCH - 0.5],\n G2_TRUE - MAX_ABS_OFFSET,\n G2_TRUE + MAX_ABS_OFFSET,\n color=\"red\", alpha=0.1, label=f\"$\\\\pm${MAX_ABS_OFFSET}\")\naxes[1].set_xlabel(\"Realization\")\naxes[1].set_ylabel(\"$g_2$\")\naxes[1].set_title(\"$g_2$ Recovery (MAP)\")\naxes[1].legend(fontsize=9)\naxes[1].set_xticks(indices)\n\nfig.suptitle(\"Level 0: MAP Shear Recovery Across 10 Realizations\", fontsize=13)\nplt.tight_layout()\nplt.show()" + "source": [ + "# Collect all g1, g2 MAP estimates\n", + "g1_means = np.array([r[\"g1_est\"].mean for r in results])\n", + "g2_means = np.array([r[\"g2_est\"].mean for r in results])\n", + "indices = np.arange(N_BATCH)\n", + "\n", + "fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n", + "\n", + "# g1 estimates\n", + "axes[0].scatter(indices, g1_means, marker=\"o\", s=60, color=\"steelblue\",\n", + " zorder=5, label=\"MAP estimate\")\n", + "axes[0].axhline(G1_TRUE, color=\"red\", ls=\"--\", lw=2, label=f\"Truth = {G1_TRUE}\")\n", + "axes[0].fill_between([-0.5, N_BATCH - 0.5],\n", + " G1_TRUE - MAX_ABS_OFFSET,\n", + " G1_TRUE + MAX_ABS_OFFSET,\n", + " color=\"red\", alpha=0.1, label=f\"$\\\\pm${MAX_ABS_OFFSET}\")\n", + "axes[0].set_xlabel(\"Realization\")\n", + "axes[0].set_ylabel(\"$g_1$\")\n", + "axes[0].set_title(\"$g_1$ Recovery (MAP)\")\n", + "axes[0].legend(fontsize=9)\n", + "axes[0].set_xticks(indices)\n", + "\n", + "# g2 estimates\n", + "axes[1].scatter(indices, g2_means, marker=\"o\", s=60, color=\"coral\",\n", + " zorder=5, label=\"MAP estimate\")\n", + "axes[1].axhline(G2_TRUE, color=\"red\", ls=\"--\", lw=2, label=f\"Truth = {G2_TRUE}\")\n", + "axes[1].fill_between([-0.5, N_BATCH - 0.5],\n", + " G2_TRUE - MAX_ABS_OFFSET,\n", + " G2_TRUE + MAX_ABS_OFFSET,\n", + " color=\"red\", alpha=0.1, label=f\"$\\\\pm${MAX_ABS_OFFSET}\")\n", + "axes[1].set_xlabel(\"Realization\")\n", + "axes[1].set_ylabel(\"$g_2$\")\n", + "axes[1].set_title(\"$g_2$ Recovery (MAP)\")\n", + "axes[1].legend(fontsize=9)\n", + "axes[1].set_xticks(indices)\n", + "\n", + "fig.suptitle(\"Level 0: MAP Shear Recovery Across 10 Realizations\", fontsize=13)\n", + "plt.tight_layout()\n", + "plt.show()" + ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], - "source": "# Bias plots: m(g1) and c(g2)\nfig, axes = plt.subplots(1, 2, figsize=(14, 5))\n\n# g1: multiplicative bias\naxes[0].scatter(indices, m_g1_vals, marker=\"s\", s=60, color=\"steelblue\")\naxes[0].axhline(0, color=\"red\", ls=\"--\", lw=2, label=\"$m = 0$ (no bias)\")\naxes[0].axhline(m_g1_vals.mean(), color=\"green\", ls=\":\", lw=1.5,\n label=f\"Mean $m$ = {m_g1_vals.mean():.2e}\")\naxes[0].set_xlabel(\"Realization\")\naxes[0].set_ylabel(\"$m_{g_1}$\")\naxes[0].set_title(\"Multiplicative Bias $g_1$\")\naxes[0].legend(fontsize=9)\naxes[0].set_xticks(indices)\n\n# g2: additive bias (g2_true = 0)\naxes[1].scatter(indices, c_g2_vals, marker=\"s\", s=60, color=\"coral\")\naxes[1].axhline(0, color=\"red\", ls=\"--\", lw=2, label=\"$c = 0$ (no bias)\")\naxes[1].axhline(c_g2_vals.mean(), color=\"green\", ls=\":\", lw=1.5,\n label=f\"Mean $c$ = {c_g2_vals.mean():.2e}\")\naxes[1].set_xlabel(\"Realization\")\naxes[1].set_ylabel(\"$c_{g_2}$\")\naxes[1].set_title(\"Additive Bias $g_2$ ($g_2^{\\\\rm true} = 0$)\")\naxes[1].legend(fontsize=9)\naxes[1].set_xticks(indices)\n\nfig.suptitle(\"Level 0: Bias per Realization (MAP)\", fontsize=13)\nplt.tight_layout()\nplt.show()" + "source": [ + "# Bias plots: m(g1) and c(g2)\n", + "fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n", + "\n", + "# g1: multiplicative bias\n", + "axes[0].scatter(indices, m_g1_vals, marker=\"s\", s=60, color=\"steelblue\")\n", + "axes[0].axhline(0, color=\"red\", ls=\"--\", lw=2, label=\"$m = 0$ (no bias)\")\n", + "axes[0].axhline(m_g1_vals.mean(), color=\"green\", ls=\":\", lw=1.5,\n", + " label=f\"Mean $m$ = {m_g1_vals.mean():.2e}\")\n", + "axes[0].set_xlabel(\"Realization\")\n", + "axes[0].set_ylabel(\"$m_{g_1}$\")\n", + "axes[0].set_title(\"Multiplicative Bias $g_1$\")\n", + "axes[0].legend(fontsize=9)\n", + "axes[0].set_xticks(indices)\n", + "\n", + "# g2: additive bias (g2_true = 0)\n", + "axes[1].scatter(indices, c_g2_vals, marker=\"s\", s=60, color=\"coral\")\n", + "axes[1].axhline(0, color=\"red\", ls=\"--\", lw=2, label=\"$c = 0$ (no bias)\")\n", + "axes[1].axhline(c_g2_vals.mean(), color=\"green\", ls=\":\", lw=1.5,\n", + " label=f\"Mean $c$ = {c_g2_vals.mean():.2e}\")\n", + "axes[1].set_xlabel(\"Realization\")\n", + "axes[1].set_ylabel(\"$c_{g_2}$\")\n", + "axes[1].set_title(\"Additive Bias $g_2$ ($g_2^{\\\\rm true} = 0$)\")\n", + "axes[1].legend(fontsize=9)\n", + "axes[1].set_xticks(indices)\n", + "\n", + "fig.suptitle(\"Level 0: Bias per Realization (MAP)\", fontsize=13)\n", + "plt.tight_layout()\n", + "plt.show()" + ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], - "source": "# g1 vs g2 MAP estimates across all realizations\nfig, ax = plt.subplots(1, 1, figsize=(6, 6))\nax.scatter(g1_means, g2_means, s=60, color=\"steelblue\", zorder=5, label=\"MAP estimates\")\nax.scatter([G1_TRUE], [G2_TRUE], c=\"red\", s=150, marker=\"*\",\n zorder=10, label=\"Truth\")\nax.set_xlabel(\"$g_1$\")\nax.set_ylabel(\"$g_2$\")\nax.set_title(\"MAP Estimates: $g_1$ vs $g_2$\")\nax.legend()\nfig.tight_layout()\nplt.show()" + "source": [ + "# g1 vs g2 MAP estimates across all realizations\n", + "fig, ax = plt.subplots(1, 1, figsize=(6, 6))\n", + "ax.scatter(g1_means, g2_means, s=60, color=\"steelblue\", zorder=5, label=\"MAP estimates\")\n", + "ax.scatter([G1_TRUE], [G2_TRUE], c=\"red\", s=150, marker=\"*\",\n", + " zorder=10, label=\"Truth\")\n", + "ax.set_xlabel(\"$g_1$\")\n", + "ax.set_ylabel(\"$g_2$\")\n", + "ax.set_title(\"MAP Estimates: $g_1$ vs $g_2$\")\n", + "ax.legend()\n", + "fig.tight_layout()\n", + "plt.show()" + ] }, { "cell_type": "markdown", "metadata": {}, - "source": "## 9. Summary\n\nThis Level 0 test validates that SHINE's forward model is **self-consistent**: when\nthe data is generated from the same model with no noise, the MAP estimate recovers\nthe true shear values with negligible bias.\n\nSince Level 0 is noiseless, MAP is the natural and fastest inference method --\nfull MCMC is unnecessary. For higher validation levels (Level 1+) with realistic\nnoise, NUTS or VI should be used instead." + "source": [ + "## 9. Summary\n", + "\n", + "This Level 0 test validates that SHINE's forward model is **self-consistent**: when\n", + "the data is generated from the same model with no noise, the MAP estimate recovers\n", + "the true shear values with negligible bias.\n", + "\n", + "Since Level 0 is noiseless, MAP is the natural and fastest inference method --\n", + "full MCMC is unnecessary. For higher validation levels (Level 1+) with realistic\n", + "noise, NUTS or VI should be used instead." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Final summary\n", + "n_passed = sum(1 for r in results if r[\"passed\"])\n", + "\n", + "print(\"Level 0 MAP Inference Summary\")\n", + "print(\"=\" * 50)\n", + "print(f\" Batch size: {N_BATCH}\")\n", + "print(f\" True shear: g1={G1_TRUE}, g2={G2_TRUE}\")\n", + "print(f\" Inference method: MAP\")\n", + "print(f\" All passed: {n_passed}/{N_BATCH}\")\n", + "print(f\" Mean m(g1): {m_g1_vals.mean():.2e}\")\n", + "print(f\" Mean c(g2): {c_g2_vals.mean():.2e}\")\n", + "print(f\" Max |g1 offset|: {max(abs(r['g1_est'].mean - G1_TRUE) for r in results):.2e}\")\n", + "print(f\" Max |g2 offset|: {max(abs(r['g2_est'].mean - G2_TRUE) for r in results):.2e}\")\n", + "print(\"=\" * 50)" + ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], - "source": "# Final summary\nn_passed = sum(1 for r in results if r[\"passed\"])\n\nprint(\"Level 0 MAP Inference Summary\")\nprint(\"=\" * 50)\nprint(f\" Batch size: {N_BATCH}\")\nprint(f\" True shear: g1={G1_TRUE}, g2={G2_TRUE}\")\nprint(f\" Inference method: MAP\")\nprint(f\" All passed: {n_passed}/{N_BATCH}\")\nprint(f\" Mean m(g1): {m_g1_vals.mean():.2e}\")\nprint(f\" Mean c(g2): {c_g2_vals.mean():.2e}\")\nprint(f\" Max |g1 offset|: {max(abs(r['g1_est'].mean - G1_TRUE) for r in results):.2e}\")\nprint(f\" Max |g2 offset|: {max(abs(r['g2_est'].mean - G2_TRUE) for r in results):.2e}\")\nprint(\"=\" * 50)" + "source": [] } ], "metadata": { @@ -188,4 +566,4 @@ }, "nbformat": 4, "nbformat_minor": 4 -} \ No newline at end of file +} From a44d0163ddad347ad884ff225e6686da4ed2cb28 Mon Sep 17 00:00:00 2001 From: EiffL Date: Sun, 8 Feb 2026 12:41:43 +0100 Subject: [PATCH 5/5] increased number of steps --- configs/validation/level0_base.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/configs/validation/level0_base.yaml b/configs/validation/level0_base.yaml index c973200..148c74e 100644 --- a/configs/validation/level0_base.yaml +++ b/configs/validation/level0_base.yaml @@ -46,6 +46,6 @@ gal: inference: method: map # MAP is sufficient for noiseless Level 0 map_config: - num_steps: 100 + num_steps: 200 learning_rate: 0.1 rng_seed: 42