test: af.NSS wiring + HST MGE smokes (searches_minimal/nss_first_class{,_gaussian}.py)

searches_minimal/nss_first_class.py

@@ -0,0 +1,137 @@
+"""
+NSS First-Class Search — `af.NSS` on the HST MGE Likelihood
+-----------------------------------------------------------
+
+End-to-end smoke for the Phase 1 `af.NSS` `NonLinearSearch` wrapper
+(PyAutoLabs/PyAutoFit#1271). Drives the same HST MGE problem as the
+bare-bones `nss_jit.py` reference, but through the production
+`search = af.NSS(...).fit(model=model, analysis=analysis)` pipeline so
+the result rides on `Paths`, `Result`, and the aggregator round-trip
+that downstream users actually consume.
+
+Expected outcome (from FINDINGS_v3 `c3_big_delete`):
+- `result.max_log_likelihood_instance.galaxies.lens.mass.einstein_radius`
+  ≈ 1.5996 ± 0.002.
+- `result.samples.log_evidence` finite, similar magnitude to the bare
+  reference (~ -31786).
+- Wall time within 20% of bare `nss_jit.py` on the same machine.
+
+Requirements:
+    pip install git+https://github.com/yallup/nss.git
+
+Run from the workspace root:
+    cd autolens_workspace_developer
+    python searches_minimal/nss_first_class.py
+"""
+
+import time
+from pathlib import Path
+
+import numpy as np
+
+import autofit as af
+
+from searches_minimal._setup import (
+    build_analysis,
+    build_dataset,
+    build_model,
+    format_best_fit,
+)
+
+
+def main():
+    dataset = build_dataset()
+    model = build_model()
+    analysis = build_analysis(dataset, use_jax=True)
+
+    print(f"Model free parameters: {model.total_free_parameters}")
+
+    search = af.NSS(
+        name="nss_first_class",
+        path_prefix=str(Path("searches_minimal") / "output"),
+        n_live=200,
+        num_mcmc_steps=5,
+        num_delete=10,
+        termination=-3.0,
+        seed=42,
+    )
+
+    print(
+        "Running af.NSS via NonLinearSearch.fit (Phase 1 wrapper). JIT compile "
+        "on the first iteration may take 25-30 s.\n"
+    )
+
+    t_start = time.time()
+    result = search.fit(model=model, analysis=analysis)
+    t_elapsed = time.time() - t_start
+
+    best_instance = result.max_log_likelihood_instance
+    samples = result.samples
+    info = samples.samples_info
+
+    max_logl = float(max(samples.log_likelihood_list))
+
+    summary = f"""\
+--- af.NSS (Phase 1 wrapper) Results ---
+Best fit:        {format_best_fit(best_instance)}
+Max log L:       {max_logl:.4f}
+Log evidence:    {info.get("log_evidence"):.4f} +/- {info.get("log_evidence_error"):.4f}
+
+--- Performance ---
+Wall time:           {t_elapsed:.2f} s     (includes JIT compile)
+Sampling time:       {info.get("sampling_time"):.2f} s
+Likelihood evals:    {info.get("total_samples")}
+Time per eval:       {info.get("sampling_time") / max(info.get("total_samples"), 1) * 1e3:.3f} ms
+ESS:                 {info.get("ess")}
+Posterior samples:   {info.get("total_accepted_samples")}
+Sampler config:      n_live={info.get("number_live_points")}, num_mcmc_steps={info.get("num_mcmc_steps")}, num_delete={info.get("num_delete")}, termination={info.get("termination")}
+
+--- Correctness gates ---
+"""
+
+    einstein_radius = float(
+        best_instance.galaxies.lens.mass.einstein_radius
+    )
+    log_evidence = float(info["log_evidence"])
+
+    assertion_lines = []
+
+    er_ok = abs(einstein_radius - 1.5996) < 0.002
+    assertion_lines.append(
+        f"einstein_radius = {einstein_radius:.4f} "
+        f"(target 1.5996 +/- 0.002) -> {'PASS' if er_ok else 'FAIL'}"
+    )
+
+    lz_finite = np.isfinite(log_evidence)
+    assertion_lines.append(
+        f"log_evidence    = {log_evidence:.4f} "
+        f"(must be finite)               -> {'PASS' if lz_finite else 'FAIL'}"
+    )
+
+    summary += "\n".join(assertion_lines) + "\n"
+
+    print()
+    print(summary)
+
+    output_dir = Path(__file__).parent / "output"
+    output_dir.mkdir(parents=True, exist_ok=True)
+    summary_path = output_dir / f"{Path(__file__).stem}_summary.txt"
+    summary_path.write_text(summary)
+    print(f"Summary written to: {summary_path}")
+
+    if not er_ok:
+        raise AssertionError(
+            f"einstein_radius {einstein_radius:.4f} outside the 1.5996 +/- 0.002 "
+            f"correctness window — the af.NSS Phase 1 wrapper has regressed."
+        )
+    if not lz_finite:
+        raise AssertionError(
+            f"log_evidence {log_evidence!r} is not finite — the af.NSS Phase 1 "
+            f"wrapper has not converged or the NSSamples conversion is broken."
+        )
+
+    print("nss_first_class: all correctness gates passed")
+
+
+if __name__ == "__main__":
+    main()

searches_minimal/nss_first_class_gaussian.py

@@ -0,0 +1,149 @@
+"""
+af.NSS Wiring Smoke — Fast 2D Gaussian
+--------------------------------------
+
+End-to-end wiring smoke for the Phase 1 ``af.NSS`` ``NonLinearSearch``
+wrapper (PyAutoLabs/PyAutoFit#1271) on a trivial 2-parameter
+JAX-traceable Gaussian likelihood with ``GaussianPrior`` priors. Runs in
+~30 seconds on CPU and validates the full path:
+
+1. ``af.NSS().fit(model, analysis)`` builds JAX closures and calls
+   ``nss.ns.run_nested_sampling``.
+2. The returned ``final_state`` + ``results`` are repackaged into the
+   ``_NSSInternal`` holder via ``_fit``.
+3. ``samples_via_internal_from`` builds ``NSSamples``.
+4. ``AbstractNest.perform_update`` writes ``samples.csv`` /
+   ``samples_summary.json`` / etc. through ``paths``.
+5. ``Result.max_log_likelihood_instance`` returns the expected type and
+   recovers the prior mean (since the likelihood is constant and the
+   prior is Gaussian).
+
+This complements the heavier ``nss_first_class.py`` HST MGE numerical
+smoke — that one validates the science (``einstein_radius=1.5996``),
+this one validates the wrapper plumbing in 1/100 of the wall time.
+
+Run from the workspace root:
+
+    python searches_minimal/nss_first_class_gaussian.py
+"""
+
+import time
+from pathlib import Path
+
+import numpy as np
+
+import autofit as af
+
+
+class FlatLikelihoodModel:
+    """Trivial 2-param model: ``x``, ``y``."""
+
+    def __init__(self, x: float = 0.0, y: float = 0.0):
+        self.x = x
+        self.y = y
+
+
+class FlatLikelihoodAnalysis(af.Analysis):
+    """Flat log-likelihood — the posterior is the prior itself."""
+
+    def log_likelihood_function(self, instance):
+        return 0.0
+
+
+def main():
+    model = af.Model(FlatLikelihoodModel)
+    model.x = af.GaussianPrior(mean=2.5, sigma=1.0)
+    model.y = af.GaussianPrior(mean=-1.0, sigma=0.5)
+
+    analysis = FlatLikelihoodAnalysis()
+
+    search = af.NSS(
+        name="nss_first_class_gaussian",
+        path_prefix=str(Path("searches_minimal") / "output"),
+        n_live=40,
+        num_mcmc_steps=2,
+        num_delete=5,
+        termination=-1.0,
+        seed=0,
+    )
+
+    print(
+        "Running af.NSS on 2D Gaussian flat-likelihood smoke. "
+        "JIT compile on first iteration may take 10-20 s.\n"
+    )
+
+    t_start = time.time()
+    result = search.fit(model=model, analysis=analysis)
+    t_elapsed = time.time() - t_start
+
+    samples = result.samples
+    info = samples.samples_info
+    best_instance = result.max_log_likelihood_instance
+
+    parameter_lists = samples.parameter_lists
+    x_samples = np.array([p[0] for p in parameter_lists])
+    y_samples = np.array([p[1] for p in parameter_lists])
+    weights = np.array(samples.weight_list)
+
+    weight_total = float(weights.sum())
+    x_weighted_mean = float((x_samples * weights).sum() / weight_total)
+    y_weighted_mean = float((y_samples * weights).sum() / weight_total)
+
+    summary = f"""\
+--- af.NSS Wiring Smoke (2D Gaussian, flat L) Results ---
+Best instance:         x={best_instance.x:.4f}, y={best_instance.y:.4f}
+Weighted posterior mean: x={x_weighted_mean:.4f} (target 2.5), y={y_weighted_mean:.4f} (target -1.0)
+Log evidence:          {info["log_evidence"]:.4f} +/- {info["log_evidence_error"]:.4f}
+Wall time:             {t_elapsed:.2f} s
+Posterior samples:     {info["total_accepted_samples"]}
+ESS:                   {info["ess"]}
+Likelihood evals:      {info["total_samples"]}
+"""
+    print()
+    print(summary)
+
+    output_dir = Path(__file__).parent / "output"
+    output_dir.mkdir(parents=True, exist_ok=True)
+    (output_dir / f"{Path(__file__).stem}_summary.txt").write_text(summary)
+
+    assert hasattr(best_instance, "x") and hasattr(best_instance, "y"), (
+        "Result.max_log_likelihood_instance is missing model attributes; "
+        "NSSamples conversion or Result construction is broken."
+    )
+    assert isinstance(best_instance.x, float), (
+        f"best_instance.x has type {type(best_instance.x)}, expected float"
+    )
+    assert np.isfinite(info["log_evidence"]), (
+        f"log_evidence = {info['log_evidence']!r} is not finite"
+    )
+    assert np.isfinite(info["log_evidence_error"]), (
+        f"log_evidence_error = {info['log_evidence_error']!r} is not finite"
+    )
+    assert info["total_accepted_samples"] > 0, (
+        "No accepted samples — NSS run produced an empty posterior."
+    )
+    assert weights.min() >= 0.0, "Negative posterior weight encountered"
+    assert abs(weight_total - 1.0) < 1e-6, (
+        f"Posterior weights do not sum to 1: total = {weight_total!r}"
+    )
+    # Loose recovery check — termination=-1 doesn't fully converge but the
+    # posterior should at least be in the right ballpark of the prior mean.
+    assert abs(x_weighted_mean - 2.5) < 1.5, (
+        f"x weighted mean {x_weighted_mean:.4f} too far from prior mean 2.5"
+    )
+    assert abs(y_weighted_mean - (-1.0)) < 1.0, (
+        f"y weighted mean {y_weighted_mean:.4f} too far from prior mean -1.0"
+    )
+
+    # Verify samples.csv was written through the Paths pipeline.
+    samples_csv = Path(result.paths._files_path) / "samples.csv"
+    assert samples_csv.exists(), (
+        f"samples.csv not written to {samples_csv} — AbstractNest.perform_update "
+        f"path appears broken for af.NSS."
+    )
+
+    print("nss_first_class_gaussian: all wiring assertions passed")
+
+
+if __name__ == "__main__":
+    main()