Improving infra for basic image fitting

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
+  method: map               # MAP is sufficient for noiseless Level 0
+  map_config:
+    num_steps: 200
+    learning_rate: 0.1
   rng_seed: 42
-  map_init:
-    enabled: true
-    num_steps: 1000
-    learning_rate: 0.01

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

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

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
 

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

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`)
 

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
 ```

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

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
 

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`)