diff --git a/searches_minimal/nss_first_class.py b/searches_minimal/nss_first_class.py new file mode 100644 index 0000000..638d454 --- /dev/null +++ b/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() diff --git a/searches_minimal/nss_first_class_gaussian.py b/searches_minimal/nss_first_class_gaussian.py new file mode 100644 index 0000000..bd4ed58 --- /dev/null +++ b/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()