Skip to content

feat: af.NSS NonLinearSearch wrapper (Phase 1 of nss_first_class_sampler) #1271

Description

@Jammy2211

Overview

Add af.NSS — a first-class NonLinearSearch for Nested Slice Sampling (JAX-native, yallup/nss) — so users can drop search = af.NSS(...) into any production autolens/autogalaxy/autofit script in place of af.Nautilus(...). This is Phase 1 of the nss_first_class_sampler roadmap, unblocked by Phase 0 (#1262, the JAX-native priors). z_projects/profiling/ measured nss_jit at 0.52 ms/eval and 7 min wall on the HST MGE lensing problem (vs Nautilus's 17.9 ms/eval and 25 min wall) — turning the bare nss.ns.run_nested_sampling call into a production sampler is what unlocks that for real science fits.

Plan

  • Build af.NSS mirror of af.Nautilus — new module PyAutoFit/autofit/non_linear/search/nest/nss/ with search.py, samples.py, __init__.py. Same lifecycle / Paths integration / Result round-trip as Nautilus.
  • JAX-traceable likelihood + prior closures: _fit builds log_likelihood(unit) and log_prior(unit) through Phase 0's xp=jnp plumbing, with the standard jnp.where(jnp.isfinite, raw, -1e30) NaN guard.
  • Call nss.ns.run_nested_sampling(rng_key, loglikelihood_fn, prior_logprob, num_mcmc_steps, initial_samples, num_delete, termination) with n_live initial samples drawn from unit cube → physical via jax.random.PRNGKey(self.seed).
  • NSSamples subclass of SamplesNest converts final_state.particles + results.logZs to PyAutoFit's Sample list; log_evidence = logZs.mean(), log_evidence_error = logZs.std().
  • Write the standard Nautilus-style output set (samples.csv, samples_summary.json, samples_info.json, search.json, settings.json, covariance.csv), lifting nautilus/samples.py:save_results patterns verbatim where they apply.
  • Guard the nss import so import autofit keeps working when nss isn't installed; raise an informative ImportError from NSS.__init__ pointing users at Phase 4's install simplification.
  • Stub Phase 2 (checkpointing) and Phase 3 (on-the-fly visualization) — accept the kwargs and log warnings; full wiring is deferred to those phases.
  • Validate end-to-end via autolens_workspace_developer/searches_minimal/nss_first_class.py against the bare-script einstein_radius=1.5996 ± 0.0002 baseline.
Detailed implementation plan

Affected Repositories

  • PyAutoFit (primary — library)
  • autolens_workspace_developer (secondary — dev-workspace verification script)

Work Classification

Library (with a dev-workspace smoke script ridealong on the same branch)

Branch Survey

Repository Current Branch Dirty?
./PyAutoFit main clean
./autolens_workspace_developer main clean

worktree_check_conflict: no active task claims either repo. Clear.

Suggested branch: feature/nss-search-wrapper
Worktree root: ~/Code/PyAutoLabs-wt/nss-search-wrapper/ (created by /start_library)

Implementation Steps

  1. New nss/ package under non_linear/search/nest/ — directory PyAutoFit/autofit/non_linear/search/nest/nss/ with __init__.py re-exporting NSS, search.py holding the class, and samples.py holding NSSamples.

  2. NSS.__init__ signature (mirrors Nautilus shape, prefers yallup/nss parameter names):

    class NSS(AbstractNest):
        def __init__(
            self,
            name: str = "",
            path_prefix: str = "",
            unique_tag: Optional[str] = None,
            # nss-specific knobs (production defaults from FINDINGS_v3 c3_big_delete)
            n_live: int = 200,
            num_mcmc_steps: int = 5,
            num_delete: int = 50,
            termination: float = -3.0,
            # standard nest knobs
            iterations_per_quick_update: Optional[int] = None,
            number_of_cores: int = 1,           # API parity only, ignored
            seed: int = 42,
            session=None,
            **kwargs,
        ):

    number_of_cores is API parity only — nss runs on whatever device JAX is configured for; document explicitly in the docstring and log a warning if user sets >1.

  3. Optional-import guard at module top:

    try:
        from nss.ns import run_nested_sampling
        _HAS_NSS = True
    except ImportError:
        _HAS_NSS = False

    NSS.__init__ raises ImportError("af.NSS requires nss. Install via 'pip install autofit[nss]' (Phase 4) or follow PyAutoPrompt/autofit/nss_install_simplification.md.") if _HAS_NSS is False. This makes import autofit succeed even without nss.

  4. _fit(model, analysis) body:

    • Build a JAX-traceable loglikelihood_fn(unit_vector):
      physical = model.vector_from_unit_vector(unit_vector, xp=jnp)
      instance = model.instance_from_vector(physical, xp=jnp)
      raw = analysis.log_likelihood_function(instance=instance)
      return jnp.where(jnp.isfinite(raw), raw, -1e30)
    • Build a JAX-traceable prior_logprob(unit_vector):
      physical = model.vector_from_unit_vector(unit_vector, xp=jnp)
      return sum(model.log_prior_list_from_vector(vector=physical, xp=jnp))
      This is what Phase 0 + the sign-convention fix (fix: log_prior_from_value sign-convention bug across Prior subclasses #1266) make possible — log_prior_list_from_vector now returns proper density-form log-priors.
    • Initial samples: draw n_live unit-cube points via jax.random.uniform(key, (n_live, ndim)) then map to physical via vmap(vector_from_unit_vector, xp=jnp).
    • Call nss.ns.run_nested_sampling(rng_key, loglikelihood_fn, prior_logprob, num_mcmc_steps, initial_samples, num_delete, termination).
    • Time the call and store wall-time in samples_info.
  5. NSSamples(SamplesNest) override of samples_via_internal_from(model, search_internal):

    • parameter_lists from final_state.particles (each particle is already a physical vector under the prompt's physical-space-sampling decision).
    • log_likelihood_list from final_state.loglikelihoods.
    • weight_list from results.logZs increments per dead-point batch (specific computation: lift from nss.ns internals if exposed, otherwise compute from per-iteration log-evidence deltas).
    • log_prior_list recomputed from parameter vectors via model.log_prior_list_from_vector(vector=parameters) (NumPy path; matches how Nautilus's samples_via_internal_from does it at nautilus/search.py:550-553).
    • log_evidence = float(results.logZs.mean()).
  6. samples_info_from(search_internal) mirrors Nautilus's shape:

    return {
        "log_evidence": float(results.logZs.mean()),
        "log_evidence_error": float(results.logZs.std()),
        "total_samples": int(results.evals),
        "total_accepted_samples": int(len(final_state.particles)),
        "time": self.timer.time if self.timer else None,
        "number_live_points": int(self.n_live),
        "num_mcmc_steps": self.num_mcmc_steps,
        "num_delete": self.num_delete,
        "termination": self.termination,
    }
  7. Output writing — leverage the existing AbstractNest.perform_update plumbing. samples.csv / samples_summary.json / samples_info.json / search.json / settings.json / covariance.csv are produced automatically once samples_via_internal_from returns a populated SamplesNest. Verify by manual inspection of the smoke run's output directory.

  8. Stub Phase 2/3 hooks_fit checks Path(paths.search_internal_path / "state.json").exists() and logs a warning that resume is not yet supported (then proceeds with a fresh fit). iterations_per_quick_update, if non-trivially set, logs a warning that quick-update visualization is not yet wired.

  9. Top-level autofit namespace — add NSS to PyAutoFit/autofit/__init__.py (mirror the Nautilus import line).

  10. Unit tests — new directory test_autofit/non_linear/search/nest/nss/:

    • test_search.py::test_init_accepts_documented_kwargs
    • test_search.py::test_init_raises_without_nss (monkeypatch _HAS_NSS=False)
    • test_search.py::test_init_warns_when_number_of_cores_gt_one
    • test_search.py::test_samples_round_trip — synthetic samples.csv + samples_info.json → load via aggregator → re-extract log_evidence identically
    • test_search.py::test_aggregator_round_trip — serialise → reload via af.Aggregator → check agg.values("samples").log_evidence matches
    • No JAX in unit tests (library policy). Anything requiring a real nss run goes to autolens_workspace_developer instead.
  11. Verification smoke scriptautolens_workspace_developer/searches_minimal/nss_first_class.py runs search = af.NSS(...) end-to-end on the HST MGE dataset (the same setup as nss_jit.py already in that folder) and asserts:

    • result.max_log_likelihood_instance.einstein_radius ≈ 1.5996 ± 0.0002
    • result.samples.log_evidence finite
    • Total wall time within 20% of the bare nss_jit.py reference

Key Files

  • PyAutoFit/autofit/non_linear/search/nest/nss/search.py — new, NSS class
  • PyAutoFit/autofit/non_linear/search/nest/nss/samples.py — new, NSSamples class
  • PyAutoFit/autofit/non_linear/search/nest/nss/__init__.py — new, package init
  • PyAutoFit/autofit/__init__.py — add NSS to top-level re-exports
  • PyAutoFit/test_autofit/non_linear/search/nest/nss/test_search.py — new, unit tests
  • autolens_workspace_developer/searches_minimal/nss_first_class.py — new, integration smoke

Reference Files

  • PyAutoFit/autofit/non_linear/search/nest/nautilus/search.py — class structure template
  • PyAutoFit/autofit/non_linear/search/nest/nautilus/samples.pysamples_via_internal_from / save_results template
  • PyAutoFit/autofit/non_linear/search/nest/abstract_nest.pyAbstractNest interface
  • z_projects/profiling/scripts/nss_jit_one_config.py — validated bare-bones nss_jit reference (HPC)
  • autolens_workspace_developer/searches_minimal/nss_jit.py — bare-bones nss_jit in the searches_minimal style

Out of scope

  • Resumption / checkpointing — Phase 2 (needs an upstream yallup/nss PR)
  • iterations_per_quick_update visualization wiring — Phase 3
  • pip install autofit[nss] extra and install simplification — Phase 4
  • Workspace tutorial scripts — Phase 5
  • HMC variant (nss_grad) — gated on upstream NaN gradient fix
  • JIT persistent cache (paths.search_internal/jax_cache/) — Phase 1 ships without; cold compile is 25-30 s per fit. Investigate as a follow-up.

Decisions baked into Phase 1

  1. Sampling space: physical (status quo for nss). Slice walks have the natural metric of the problem, no per-step remapping cost. Documented in the af.NSS docstring.
  2. log_evidence_error: logZs.std() — the stochastic batch error of the NS estimator across the live ensemble.
  3. number_of_cores: API parity only — ignored by nss. Documented.
  4. Fitness layer: unchanged. af.NSS bypasses Fitness._call and goes directly from the model + analysis to JAX closures, because the whole point is to run inside jax.jit.

Original Prompt

Click to expand starting prompt

Add af.NSS — a first-class NonLinearSearch for Nested Slice
Sampling (JAX-native, yallup/nss) — so users can drop
search = af.NSS(...) into any production autolens/autogalaxy/autofit
script in place of af.Nautilus(...).

This is Phase 1 of z_features/nss_first_class_sampler.md.
Blocked on Phase 0 (autofit/priors_jax_native.md) — without
JAX-native priors the cube → physical mapping has to use
jax.pure_callback, costing ~18 ms/eval and erasing the 30× per-eval
advantage that motivates this whole feature (see
z_projects/profiling/FINDINGS_v3.md).

Why this matters

The user-facing goal from the z_feature:

"Be able to put search = af.NSS(name=..., n_live=200, num_mcmc_steps=5, num_delete=50) into any production
autolens_workspace / autogalaxy_workspace script and have it
work end-to-end — checkpointing, on-the-fly visualization,
samples.csv / model.results / aggregator round-trip — just like
af.Nautilus."

z_projects/profiling/ already proved the numerical headline
(nss_jit c3_big_delete: 7 min wall, 0.57 ms/eval, recovered
einstein_radius=1.5996 consistently across 4 of 6 sweep configs).
This prompt turns the bare nss.ns.run_nested_sampling call into a
production sampler.

Reference implementation pattern

Model the new class on af.Nautilus:

  • @PyAutoFit/autofit/non_linear/search/nest/nautilus/search.py
    class structure, __init__ kwargs, _fit method, paths integration.
  • @PyAutoFit/autofit/non_linear/search/nest/nautilus/samples.py
    posterior extraction, log_Z, ESS.
  • @PyAutoFit/autofit/non_linear/search/nest/abstract.py — the
    abstract NonLinearSearch interface every sampler inherits from.

The working bare-bones reference for the nss_jit numerical pipeline:
@z_projects/profiling/scripts/nss_jit_one_config.py. Lift the
build_dataset, model setup, log_likelihood, log_prior, and
run_nested_sampling call structure verbatim — those are already
validated on the HPC.

What to build

New directory:
@PyAutoFit/autofit/non_linear/search/nest/nss/ with __init__.py,
search.py, samples.py, mirroring the Nautilus layout.

af.NSS __init__ signature (mirror Nautilus, prefer parameter names from yallup/nss)

class NSS(AbstractNest):
    def __init__(
        self,
        name: str = "",
        path_prefix: str = "",
        unique_tag: str | None = None,
        # nss-specific knobs (production defaults from FINDINGS_v3 c3)
        n_live: int = 200,
        num_mcmc_steps: int = 5,
        num_delete: int = 50,
        termination: float = -3.0,
        # standard nest knobs
        iterations_per_quick_update: int = 10_000,
        number_of_cores: int = 1,           # nss is JAX; this is here for API parity, ignored
        seed: int = 42,
        session=None,                       # SQL session
        **kwargs,
    ):

number_of_cores is API parity only — nss_jit runs on whatever
device JAX is configured for. Document that explicitly in the
docstring.

_fit(model, analysis) body

Outline:

  1. Build the JAX-traceable log_likelihood(unit_vector) closure:
    instance = model.instance_from_vector(model.vector_from_unit_vector(unit_vector, xp=jnp), xp=jnp)
    then call analysis.log_likelihood_function(instance=instance).
    Wrap with the NaN guard
    jnp.where(jnp.isfinite(raw), raw, -1e30).

  2. Build the JAX-traceable log_prior(unit_vector) closure: sum of
    model.log_prior_list_from_vector(physical, xp=jnp) mapped from
    the unit cube. Phase 0 makes this possible. Today's
    @z_projects/profiling/scripts/nss_jit_one_config.py uses a
    hard-bound box — replace that with the proper log-prior once
    Phase 0 lands.

  3. Initial samples: n_live unit-cube draws mapped to physical via
    model.vector_from_unit_vector(u, xp=jnp). Seed from
    jax.random.PRNGKey(self.seed).

  4. Call nss.ns.run_nested_sampling(rng_key, loglikelihood_fn=..., prior_logprob=..., num_mcmc_steps=..., initial_samples=..., num_delete=..., termination=...).

  5. After return, build an NSSamples object (new class in
    samples.py) from final_state.particles + results.logZs.

  6. Write the standard PyAutoFit output schema (see "Output writing"
    below).

Output writing

Match the file set Nautilus produces in paths.samples_path:

  • samples.csv — every accepted sample's (parameter_0, ..., parameter_N, log_likelihood, log_prior, weight). For nss the
    weights come from logZs increments per dead-point.
  • samples_summary.json — max log L, log Z + error, max ESS, best fit.
  • samples_info.json — sampler metadata (n_live, num_mcmc_steps,
    num_delete, termination, evals, sampling_time).
  • search.json — sampler config (n_live, etc.) serialised so
    Aggregator can reproduce.
  • settings.jsoniterations_per_quick_update, seed, etc.
  • covariance.csv — posterior covariance over the dead points.

Most of this is mechanical and the Nautilus implementation in
@PyAutoFit/autofit/non_linear/search/nest/nautilus/samples.py:save_results
is the template.

NSSamples class

Inherit from SamplesNest (or whatever the existing nested-sampler
samples base class is). Override:

  • parameter_lists — list of physical parameter vectors per dead point.
  • log_likelihood_list, weight_list, log_prior_list.
  • log_evidenceresults.logZs.mean().
  • log_evidence_errresults.logZs.std() (nss returns an array
    of logZ estimates across the live ensemble).
  • total_samplesresults.evals.
  • total_accepted_sampleslen(final_state.particles).

Resumption + visualization

Stub — these are Phases 2 and 3. The Phase 1 implementation
should:

  • Accept iterations_per_quick_update in __init__ but log a
    warning that it's not yet wired up.
  • Detect paths.search_internal/state.json existing and warn that
    resume isn't supported yet (then proceed with a fresh fit).

The Phase 2 prompt (TBD, depends on an upstream nss PR) wires both
properly.

Plotting hooks

AbstractNest has standard plot_results methods. nss's posterior
is final_state.particles + per-point log-likelihood, mappable to
the corner.py interface the existing nested-sampling plot machinery
uses. Lift the Nautilus plot_results body verbatim if the API
matches.

What to verify

  1. Numerical parity smoke. Run the searches_minimal/nss_jit.py
    case (or z_projects/profiling/scripts/nss_jit_one_config.py
    c3_big_delete) through the new af.NSS(...) and confirm
    einstein_radius=1.5996 ± 0.0002, max log L within 1 nat of
    the bare-script value (~31786). Wall time within 20% of the bare
    reference.

  2. Unit tests. Test files under
    @PyAutoFit/test_autofit/non_linear/search/nest/nss/. Cover:

    • NSS.__init__ accepts the documented kwargs and rejects unknown
      ones cleanly.
    • NSS.fit returns a Result with result.samples non-empty.
    • samples_from(paths) round-trip — write → load → re-extract
      max log L identically.
    • Aggregator round-trip: serialise a fit, load via af.Aggregator,
      check agg.values("samples").log_evidence returns the same number.
  3. autolens_workspace_developer smoke. Add an opt-in script
    searches_minimal/nss_first_class.py that uses
    search = af.NSS(...) end-to-end on the HST MGE dataset. Compare
    the resulting result.max_log_likelihood_instance to the
    nautilus_jax baseline — should agree to <1 nat.

  4. Existing samplers unaffected. Run /smoke_test on each PyAuto
    workspace_test. The change is additive (new module under
    non_linear/search/nest/nss/) so should not regress anything;
    verify anyway.

  5. import autofit does not require nss. Phase 4 will ship a
    pyautofit[nss] extra; the Phase 1 implementation must guard the
    nss import so import autofit keeps working when nss is not
    installed:

    try:
        from nss.ns import run_nested_sampling
        _HAS_NSS = True
    except ImportError:
        _HAS_NSS = False
    
    class NSS(AbstractNest):
        def __init__(self, ...):
            if not _HAS_NSS:
                raise ImportError(
                    "af.NSS requires nss. Install via 'pip install autofit[nss]' "
                    "or follow PyAutoPrompt/autofit/nss_install_simplification.md."
                )
            ...

Out of scope

  • Resumption / checkpointing — Phase 2.
  • iterations_per_quick_update visualization — Phase 3.
  • Install simplification — Phase 4 (autofit/nss_install_simplification.md).
  • Workspace tutorial scripts — Phase 5.
  • HMC variant (nss_grad) — gated on the upstream NaN gradient fix
    from session-1's probe_grad.py.

Risks / open questions

  1. Sampling space. nss currently samples in physical parameter
    space (initial samples are model.vector_from_unit_vector(u)
    per particle). The Nautilus and PocoMC convention is to sample in
    unit-cube space and apply the prior transform inside the
    likelihood. There's an argument either way:

    • Physical-space sampling (status quo): the slice walks have the
      natural metric of the problem, no remapping cost per step.
    • Unit-cube sampling: matches Nautilus/PocoMC mental model,
      simplifies log_prior to jnp.where(in_unit_cube, 0, -inf).

    Decision needed before writing the wrapper. Recommendation:
    keep physical-space sampling (it's what nss is designed for) but
    document the convention in the af.NSS docstring so users know
    why the model.priors_ordered_by_id log-priors are summed inside
    the likelihood.

  2. logZs interpretation. nss returns results.logZs as an
    array of log-evidence estimates per dead-point batch (the
    stochastic batch error of the NS estimator). PyAutoFit's
    Samples.log_evidence expects a single number. Use
    logZs.mean() and expose logZs.std() as
    log_evidence_error. Verify this matches how Nautilus reports
    it.

  3. JIT cache. Every fit recompiles
    run_nested_sampling's while_loop (~25–30 s cold). Production
    users running batched fits should not pay this each time —
    investigate jax.persistent_cache directory baked into PyAutoFit's
    Paths (probably paths.search_internal/jax_cache/). Worth a
    one-line jax.config.update(...) call in _fit's setup; verify
    the cache directory survives across python invocations.

Reference

  • @PyAutoFit/autofit/non_linear/search/nest/nautilus/search.py
    reference structure
  • @PyAutoFit/autofit/non_linear/search/nest/abstract.py — abstract
    base
  • @z_projects/profiling/scripts/nss_jit_one_config.py — validated
    bare-bones reference for the JAX likelihood + nss call
  • @z_projects/profiling/FINDINGS_v3.md — wall-time and ms/eval
    numbers
  • @PyAutoPrompt/z_features/nss_first_class_sampler.md — Phase 1 in
    the sequenced roadmap
  • @PyAutoPrompt/autofit/priors_jax_native.md — Phase 0 blocker

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions