feat: autofit[nss] install extra (Phase 4 of nss_first_class_sampler)

[Jammy2211]

Overview

Phase 4 of the nss_first_class_sampler roadmap — make installing af.NSS a single safe command, eliminating the multi-step install saga documented in z_projects/profiling/FINDINGS_v3.md. Closes the install-pain barrier so Phase 5 (workspace tutorial scripts) can recommend af.NSS to users without a 3-hour adventure.

Approach (after audit): Option C from the original prompt — a pyproject.toml extra autofit[nss] listing pinned git+ URLs in dependency order — is the lightest path that works with modern pip. Original prompt assumed pip extras couldn't sequence properly; modern pip 23+ handles URL-direct deps well, and pinning the handley-lab/blackjax fork as the primary blackjax in our extra means pip never installs mainline blackjax in the first place. Fallback to Option D (python -m autofit.install_nss helper) if C fails the fresh-venv smoke.

Plan

• Add [project.optional-dependencies] entry nss = [...] in PyAutoFit/pyproject.toml listing four pinned URL deps in install order: fastprogress<1.1, handley-lab/blackjax fork (git+URL), yallup/nss (git+URL), plus any transitive deps the fork needs.
• Pin specific commit SHAs (not branch heads) for both git+ URLs so installs are reproducible and we control upgrade timing.
• Fresh-venv smoke: python -m venv tmp_venv && tmp_venv/bin/pip install -e PyAutoFit[nss] && tmp_venv/bin/python -c "import autofit as af; af.NSS()" must succeed end-to-end in CI.
• Add a GitHub Actions workflow .github/workflows/nss_install_smoke.yml that runs the fresh-venv smoke on every PR — catches install regressions before users see them.
• If Option C trips on pip resolver issues (mainline blackjax wins despite the URL pin, or transitive dep conflict), fall back to Option D: ship a autofit/install_nss.py helper that subprocess-calls the install commands in order.
• Update af.NSS's ImportError message to reference the new pip install autofit[nss] command (currently points at the manual git+URL install).
• Document the install in PyAutoFit/README.md (or wherever nautilus-sampler install is documented) so users find it via the standard PyAutoFit docs.

Detailed implementation plan

Affected Repositories

• PyAutoFit (primary — library)

Work Classification

Library

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoFit	main	clean

worktree_check_conflict (exit=0): no active task claims PyAutoFit.

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

Implementation Steps

1. Audit upstream commit SHAs. Pull current HEAD from handley-lab/blackjax and yallup/nss. Verify those exact commits work with our pinned JAX 0.4.x. Record SHAs for pinning.

2. Add the nss extra to PyAutoFit/pyproject.toml:

   [project.optional-dependencies]
   nss = [
     "fastprogress<1.1",
     "blackjax @ git+https://github.com/handley-lab/blackjax.git@<COMMIT_SHA>",
     "nss @ git+https://github.com/yallup/nss.git@<COMMIT_SHA>",
   ]

3. Fresh-venv smoke locally to validate before touching CI:

   python -m venv /tmp/nss_install_smoke
   source /tmp/nss_install_smoke/bin/activate
   pip install -e PyAutoFit[nss]
   python -c "import autofit as af; s = af.NSS(); print('NSS import + instantiate OK')"
   deactivate && rm -rf /tmp/nss_install_smoke

   If this succeeds → Option C is viable; proceed to CI. If it fails → diagnose, possibly fall back to Option D.

4. Add a CI workflow at PyAutoFit/.github/workflows/nss_install_smoke.yml:

   • Triggers: pull_request (any path) + weekly cron (catches upstream regressions on unchanged code).
   • Job: fresh ubuntu venv, pip install .[nss], run the same import smoke as step 3.

5. Update af.NSS.__init__'s ImportError message to reference the new install command:

   raise ImportError(
       "af.NSS requires the optional `nss` package. Install via\n"
       "    pip install autofit[nss]\n"
       "(see PyAutoFit's pyproject.toml for the pinned upstream commits)."
   )

6. Document the install in README.md alongside the existing nautilus-sampler install pointer.

7. Smoke test on a fresh venv end-to-end — pip install autofit[nss] from scratch, then run autolens_workspace_developer/searches_minimal/nss_first_class_gaussian.py (Phase 1's Gaussian smoke) to confirm the full path works without any manual install steps.

Fallback path (Option D)

If step 3 fails, add PyAutoFit/autofit/install_nss.py:

"""Helper to install `af.NSS`'s dependencies in the right order."""
import subprocess
import sys

_COMMANDS = [
    ["pip", "install", "fastprogress<1.1"],
    ["pip", "install", "git+https://github.com/handley-lab/blackjax.git@<SHA>"],
    ["pip", "install", "git+https://github.com/yallup/nss.git@<SHA>", "--no-deps"],
]

def main():
    for cmd in _COMMANDS:
        print(f"Running: {' '.join(cmd)}")
        subprocess.check_call([sys.executable, "-m"] + cmd)
    print("af.NSS dependencies installed.")

if __name__ == "__main__":
    main()

User runs pip install autofit && python -m autofit.install_nss. Same smoke gate applies.

Key Files

• PyAutoFit/pyproject.toml — add [project.optional-dependencies] nss = [...]
• PyAutoFit/.github/workflows/nss_install_smoke.yml — new CI workflow
• PyAutoFit/autofit/non_linear/search/nest/nss/search.py — update ImportError message
• PyAutoFit/README.md — install documentation
• PyAutoFit/autofit/install_nss.py — fallback only, ship only if Option C fails

Out of scope

• Vendoring nss / blackjax (Option A) — heavy maintenance commitment; only revisit if Option C and D both fail
• Upstream coordination with yallup (Option B) — pursue informally in parallel; not gating this PR
• nss_grad / HMC variant install — separate prompt when that lands
• conda / mamba install instructions — pip-first
• pocomc / numpyro install pain — those are separate samplers, not in scope here

Risks / open questions

1. Pip resolver may install mainline blackjax anyway despite our URL-direct pin. Verify in step 3's fresh-venv smoke. If it does, the workaround is --no-deps on the nss install line, which works as a pip install flag but I need to confirm it works inside a pyproject.toml extra. If pyproject.toml can't express --no-deps, that's the Option D fallback trigger.

2. Commit SHA staleness. Pinning specific commits insulates from upstream churn but means we miss bug fixes. Mitigation: periodically (e.g. quarterly) bump the pins, run the smoke, and ship a CHANGELOG note.

3. License audit. Verify yallup/nss (BSD-3-Clause per latest check) and handley-lab/blackjax (Apache-2.0 from the file headers) licenses are compatible with PyAutoFit before pinning. They are, but document it in a top-of-file comment in the extra.

4. CI cost. The fresh-venv install pulls JAX, JAX deps, fastprogress, optax, chex, equinox, jaxtyping, etc. — likely 5-10 min CI step. Run on a single Python version (3.12) in this workflow to keep cost bounded.

Original Prompt

Click to expand starting prompt

Make installing af.NSS a single safe command for a science user
installing PyAutoFit fresh — eliminate the multi-hour multi-step
install saga that the profiling project survived in session-of-2026-05-11.

This is Phase 4 of z_features/nss_first_class_sampler.md.
Depends on Phase 1 (so we know what API surface to vendor or
pin against), and informs Phase 5 (workspace tutorials can only
recommend af.NSS to users if installation isn't a 3-hour adventure).

[... full prompt as authored — see PyAutoPrompt/issued/nss_install_simplification.md for the verbatim source.]

[Jammy2211]

Shipped — 2026-05-16

PR

• feat: autofit[nss] install extra (Phase 4 of nss_first_class_sampler) #1277 (merged)

Summary

`pip install autofit[nss]` is now the single command that installs `af.NSS` and all its dependencies. Replaces the multi-step install saga documented in FINDINGS_v3.md.

```bash
pip install autofit[nss]
python -c "import autofit as af; af.NSS()" # just works
```

Approach (Option C from the prompt)

The Phase 4 prompt recommended Option A (full vendoring of ~1300 LOC). Investigation revealed modern pip 23+ handles URL-direct deps cleanly, so the lighter Option C path (pyproject.toml extra with pinned git+ URLs) works without ongoing re-vendor maintenance burden. Verified empirically on a fresh venv.

Critical pin: `ef45acd2f` for handley-lab/blackjax

First attempt with HEAD (`5dbf89a9`) failed: HEAD adds `numpy>=1.25` which conflicts with PyAutoFit's `anesthetic==2.8.14` (numpy<2.0). Pinned to `ef45acd2f` instead — the May 2026 "Merge PR #60 — double_compile" commit, validated through Phases 1-3. Bump only when the anesthetic numpy cap moves (likely when Python 3.13 anesthetic>=2.9 takes over).

Validation

• `pytest test_autofit` — 1258 passed, 1 skipped (no regression).
• Fresh `python -m venv /tmp/v && pip install -e PyAutoFit[nss]` on Python 3.12 — completes in ~3 min, no `--no-deps` flag, no resolver conflict.
• Fresh-venv `af.NSS()` instantiate + `blackjax.ns.adaptive.init` attribute check (canonical fork marker) — pass.
• New CI workflow runs the fresh-venv install + import smoke on every PR + a Sunday 03:00 UTC cron (catches upstream drift past the pinned SHAs).

nss_first_class_sampler roadmap status

Phase	Status
0 — JAX-native priors	shipped (#1262 + #1266)
1 — af.NSS wrapper	shipped (#1271 → #1272)
2 — checkpointing/resume	shipped (#1273 → #1274)
3 — on-the-fly viz	shipped (#1273 → #1274)
4 — install simplification	shipped (#1276 → #1277)
5 — workspace tutorial scripts	not yet — autolens_workspace/searches/nss.py + autofit/autogalaxy cookbooks

Only Phase 5 remains. With install now `pip install autofit[nss]`, the workspace tutorials can recommend af.NSS without any disclaimers.