feat(inference): add production IFU experiment runner with checkpointed orchestration

Author: TobiBu

bench/check_inference_guardrails.py

@@ -0,0 +1,115 @@
+#!/usr/bin/env python
+import argparse
+import json
+from pathlib import Path
+
+import jax.numpy as jnp
+
+from rubix.core.data import Galaxy, GasData, RubixData, StarsData
+from rubix.inference import (
+    OptimizationObjectiveThresholds,
+    RuntimeThresholds,
+    benchmark_ifu_cube_optimization,
+    benchmark_result_to_dict,
+    check_ifu_optimization_guardrails,
+)
+
+
+class _SyntheticPipeline:
+    """Small synthetic pipeline for guardrail smoke checks."""
+
+    def __init__(self, template: jnp.ndarray):
+        self.template = template
+
+    def run_sharded(self, rubixdata: RubixData) -> jnp.ndarray:
+        return rubixdata.stars.age[0] * self.template
+
+
+def _make_data() -> RubixData:
+    return RubixData(
+        galaxy=Galaxy(),
+        stars=StarsData(
+            coords=jnp.zeros((1, 3)),
+            velocity=jnp.zeros((1, 3)),
+            mass=jnp.ones(1),
+            age=jnp.array([0.0]),
+            metallicity=jnp.array([0.01]),
+        ),
+        gas=GasData(
+            coords=jnp.zeros((1, 3)),
+            velocity=jnp.zeros((1, 3)),
+            mass=jnp.ones(1),
+        ),
+    )
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Run optimization benchmark and assert runtime/loss guardrails."
+    )
+    parser.add_argument("--nx", type=int, default=8)
+    parser.add_argument("--ny", type=int, default=8)
+    parser.add_argument("--nw", type=int, default=64)
+    parser.add_argument("--max-steps", type=int, default=120)
+    parser.add_argument("--repeats", type=int, default=2)
+    parser.add_argument("--max-mean-runtime-s", type=float, default=3.0)
+    parser.add_argument("--max-median-runtime-s", type=float, default=3.0)
+    parser.add_argument("--max-final-loss", type=float, default=1e-3)
+    parser.add_argument("--max-best-loss", type=float, default=1e-3)
+    parser.add_argument("--output-json", type=str, default="")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+
+    cube = jnp.ones((args.nx, args.ny, args.nw), dtype=jnp.float32)
+    target = 1.5 * cube
+
+    benchmark_result = benchmark_ifu_cube_optimization(
+        pipeline=_SyntheticPipeline(cube),
+        params_init={"stars": {"age": jnp.array([0.2])}},
+        static_data=_make_data(),
+        target=target,
+        learning_rate=0.1,
+        max_steps=args.max_steps,
+        tol=1e-8,
+        repeats=args.repeats,
+        warmup=True,
+    )
+
+    runtime_thresholds = RuntimeThresholds(
+        max_mean_runtime_s=args.max_mean_runtime_s,
+        max_median_runtime_s=args.max_median_runtime_s,
+    )
+    objective_thresholds = OptimizationObjectiveThresholds(
+        max_final_loss=args.max_final_loss,
+        max_best_loss=args.max_best_loss,
+    )
+
+    check = check_ifu_optimization_guardrails(
+        benchmark_result,
+        runtime_thresholds,
+        objective_thresholds,
+    )
+
+    payload = {
+        "benchmark": benchmark_result_to_dict(benchmark_result),
+        "guardrail": {
+            "passed": check.passed,
+            "message": check.message,
+            "failed_conditions": check.failed_conditions,
+        },
+    }
+
+    if args.output_json:
+        output_path = Path(args.output_json)
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        output_path.write_text(json.dumps(payload, indent=2), encoding="utf-8")
+
+    if not check.passed:
+        raise SystemExit(check.message)
+
+
+if __name__ == "__main__":
+    main()

docs/inference_workflows.rst

@@ -320,6 +320,48 @@ residual metrics) and persist science-ready outputs:
      --vi-steps 200 \
      --num-posterior-draws 16
 
+
+Production IFU Experiment
+-------------------------
+
+Run a full configuration-driven IFU experiment with optional checkpointing and
+resume support:
+
+.. code-block:: bash
+
+   python scripts/run_ifu_science_experiment.py \
+     --config rubix/config/inference_experiment_template.yml
+
+The template supports:
+
+- deterministic or stochastic mode selection
+- runtime objective selection via ``run.objective``
+- chunked optimization/VI with stage checkpoints
+- posterior predictive output products and masked science metrics
+
+
+Real Data Runbook
+-----------------
+
+1. Prepare cube-side arrays as ``.npy`` or ``.npz`` files with matching
+   ``(nx, ny, nw)`` shapes:
+   ``target``, optional ``mask``, optional ``weights``, optional ``sigma`` or
+   ``inv_variance``.
+2. Copy and edit
+   ``rubix/config/inference_experiment_template.yml``:
+   set ``run.rubix_config_path``, ``data.*_path``, and stage hyperparameters.
+3. Run deterministic fitting first:
+
+   .. code-block:: bash
+
+      python scripts/run_ifu_science_experiment.py \
+        --config rubix/config/inference_experiment_template.yml
+
+4. If interrupted, resume from latest stage checkpoint by setting
+   ``optimization.resume_checkpoint`` or ``variational.resume_checkpoint``.
+5. Inspect outputs in ``run.output_dir``:
+   ``summary.json``, ``predictive_summary.npz``, ``residual_products.npz``.
+
 Benchmarking Full-IFU Optimization
 ----------------------------------
 
@@ -333,3 +375,14 @@ memory diagnostics for full IFU cubes.
      --repeats 3 \
      --max-steps 200 \
      --use-mask --use-weights
+
+For CI-style threshold checks on a representative synthetic IFU size:
+
+.. code-block:: bash
+
+   python bench/check_inference_guardrails.py \
+     --nx 8 --ny 8 --nw 64 \
+     --max-steps 120 \
+     --max-mean-runtime-s 3.0 \
+     --max-final-loss 1e-3 \
+     --output-json outputs/guardrails/opt_guardrail.json

docs/rubix.inference.rst

@@ -20,6 +20,14 @@ rubix.inference.checkpoint module
    :undoc-members:
    :show-inheritance:
 
+rubix.inference.experiment module
+---------------------------------
+
+.. automodule:: rubix.inference.experiment
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
 rubix.inference.losses module
 -----------------------------
 

rubix/config/inference_experiment_template.yml

@@ -0,0 +1,60 @@
+run:
+  rubix_config_path: rubix/config/rubix_config.yml
+  mode: deterministic
+  seed: 0
+  noise_seed: 0
+  output_dir: outputs/ifu_science
+  checkpoint_dir: outputs/ifu_science/checkpoints
+  params_init_overrides:
+    stars:
+      age: [1.0]
+      metallicity: [0.01]
+  objective:
+    kind: combined
+    terms:
+      - kind: gaussian_nll
+        inv_variance_key: inv_variance
+        mask_key: mask
+      - kind: huber
+        delta: 0.2
+        mask_key: mask
+        weight: 0.05
+
+data:
+  target_path: data/target_cube.npy
+  target_key:
+  mask_path: data/mask_cube.npy
+  mask_key:
+  weights_path:
+  weights_key:
+  sigma_path:
+  sigma_key:
+  inv_variance_path: data/inv_variance_cube.npy
+  inv_variance_key:
+
+optimization:
+  enabled: true
+  learning_rate: 0.001
+  max_steps: 500
+  tol: 1.0e-6
+  normalize_loss: true
+  checkpoint_interval_steps: 100
+  resume_checkpoint:
+
+variational:
+  enabled: true
+  learning_rate: 0.005
+  max_steps: 500
+  tol: 1.0e-6
+  num_samples: 4
+  beta_kl: 0.001
+  init_log_std: -2.0
+  normalize_loss: true
+  huber_delta: 0.2
+  huber_weight: 0.0
+  checkpoint_interval_steps: 100
+  resume_checkpoint:
+
+predictive:
+  enabled: true
+  num_draws: 16

rubix/inference/__init__.py

@@ -16,6 +16,11 @@
     resume_variational_from_checkpoint,
     save_checkpoint,
 )
+from .experiment import (
+    normalize_experiment_config,
+    run_ifu_experiment,
+    save_ifu_experiment_outputs,
+)
 from .losses import combine_loss_fns, huber_data_loss, masked_gaussian_nll
 from .modes import get_pipeline_name_for_mode, make_inference_pipeline
 from .objective_config import build_loss_from_config, build_loss_from_user_config
@@ -87,6 +92,7 @@
     "load_checkpoint",
     "make_optimization_checkpoint",
     "make_variational_checkpoint",
+    "normalize_experiment_config",
     "build_age_metallicity_transforms",
     "build_ifu_cube_loss",
     "build_loss_from_config",
@@ -124,6 +130,8 @@
     "summarize_masked_metrics",
     "summarize_predictive_cube_samples",
     "save_checkpoint",
+    "run_ifu_experiment",
+    "save_ifu_experiment_outputs",
     "value_and_grad",
     "run_synthetic_science_recipe",
     "save_science_recipe_outputs",