Skip to content

docs: remove the NSS section from the searches/nest example #88

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Jammy2211 merged 1 commit into from
Jul 11, 2026
Merged

docs: remove the NSS section from the searches/nest example #88

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
129 changes: 3 additions & 126 deletions notebooks/searches/nest.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,30 +12,13 @@
" - `DynestyStatic`: Dynesty with static nested sampling.\n",
" - `DynestyDynamic`: Dynesty with dynamic nested sampling.\n",
" - `Nautilus`: Nautilus nested sampler.\n",
" - `NSS`: Nested Slice Sampling (JAX-native, optional install).\n",
"\n",
"The first three are boundary-based samplers that work with any Python log-likelihood. `NSS` is a more recent\n",
"JAX-native sampler that runs its inner sampling loop inside `jax.jit` \u2014 when the log-likelihood is itself\n",
"JAX-traceable, the per-evaluation cost drops by roughly an order of magnitude versus the boundary samplers\n",
"(measured on real lensing likelihoods, see the `nss_first_class_sampler` roadmap and FINDINGS_v3 in the\n",
"profiling project for the numbers).\n",
"These are boundary-based samplers that work with any Python log-likelihood.\n",
"\n",
"Relevant links:\n",
"\n",
" - Dynesty: https://dynesty.readthedocs.io/en/latest/\n",
" - Nautilus: https://nautilus-sampler.readthedocs.io/en/stable/\n",
" - NSS (Nested Slice Sampling): https://github.com/yallup/nss\n",
"\n",
"__Install Precondition for NSS__\n",
"\n",
"The `Search: NSS` section at the bottom of this script imports the optional `nss` package. To run that\n",
"section, install the dependencies first:\n",
"\n",
" pip install autofit[nss]\n",
"\n",
"The extra pins the right `handley-lab/blackjax` fork at a known-good commit, so this is a single safe\n",
"command (no `--no-deps` dance, no manual git+ URLs). The other three samplers in this script have no\n",
"additional dependencies and run with the standard `pip install autofit` install.\n",
"\n",
"__Contents__\n",
"\n",
Expand All @@ -46,7 +29,6 @@
"- **Search: DynestyStatic**: Configuring and running the DynestyStatic nested sampler.\n",
"- **Search: DynestyDynamic**: Configuring and running the DynestyDynamic nested sampler.\n",
"- **Search: Nautilus**: Configuring and running the Nautilus nested sampler.\n",
"- **Search: NSS**: Configuring and running the NSS (Nested Slice Sampling) sampler.\n",
"- **Search Internal**: Accessing the internal sampler for advanced use (shown once for DynestyStatic)."
]
},
Expand All @@ -61,7 +43,7 @@
"points the configuration at it. If you are running the notebook elsewhere (e.g. locally via\n",
"your own installation) it does nothing, and you can run it safely.\n",
"\n",
"Colab tip: model-fits run much faster on a GPU \u2014 enable one via \"Runtime\" -> \"Change runtime\n",
"Colab tip: model-fits run much faster on a GPU — enable one via \"Runtime\" -> \"Change runtime\n",
"type\" -> \"Hardware accelerator\" before running the notebook."
]
},
Expand Down Expand Up @@ -433,111 +415,6 @@
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"__Search: NSS__\n",
"\n",
"NSS (Nested Slice Sampling) is a JAX-native nested sampler whose inner sampling loop runs end-to-end inside\n",
"`jax.jit`. The advantage versus the boundary samplers above is **per-evaluation cost**: when your\n",
"log-likelihood is itself JAX-traceable, NSS avoids the Python \u2194 JAX boundary that Nautilus and Dynesty cross\n",
"on every likelihood call. On the production lensing likelihoods that motivated this sampler, the per-eval\n",
"cost is roughly 30 times lower than Nautilus's, and total wall time to convergence drops from tens of\n",
"minutes to a few minutes.\n",
"\n",
"On the trivial 1D Gaussian dataset used by this tutorial the speedup is not visible \u2014 the likelihood is so\n",
"cheap that the per-call cost is dominated by Python overhead rather than the floating-point work, and the\n",
"first NSS run pays a one-off ~25\u201330 second JIT compile that the simpler samplers above skip. The numbers\n",
"below let you confirm NSS converged to roughly the same posterior, not that it ran faster. Try NSS on a\n",
"real autolens or autogalaxy MGE / pixelization likelihood to see the per-eval advantage.\n",
"\n",
"NSS exposes the same `result.samples` interface as the boundary samplers \u2014 swapping `af.Nautilus(...)` for\n",
"`af.NSS(...)` in your existing scripts is a one-line change.\n",
"\n",
"A few NSS-specific kwargs to be aware of:\n",
"\n",
" - `n_live`: live particles maintained throughout the run (production default 200).\n",
" - `num_mcmc_steps`: slice-MCMC inner steps per dead-point batch (production default 5).\n",
" - `num_delete`: particles removed per outer iteration (production default 50; larger values reduce JIT\n",
" overhead per outer iteration at the cost of slightly coarser posterior coverage).\n",
" - `termination`: stopping criterion on `logZ_live - logZ` (default `-3.0`, corresponding to a remaining\n",
" evidence fraction below 1e-3).\n",
" - `checkpoint_interval`: outer iterations between disk-saved checkpoints. NSS writes a resumable state\n",
" file every `checkpoint_interval` iterations, so a SLURM timeout halfway through a long fit is recovered\n",
" automatically the next time you run the same script with the same Paths.\n",
" - `iterations_per_quick_update`: when set non-None, NSS calls `analysis.visualize(...)` with the current\n",
" best live point every N outer iterations \u2014 partial results appear in the image_path directory while the\n",
" run is still in flight.\n",
"\n",
"Settings reference: see `af.NSS.__init__` for the full kwarg list.\n",
"\n",
"__Optional Dependency Guard__\n",
"\n",
"The standard workspace release environment does not install `autofit[nss]`. When the optional stack is\n",
"absent, this script skips only the NSS fit and leaves the Dynesty / Nautilus examples above runnable.\n",
"\n",
"__Analysis Must Be JAX-Traceable__\n",
"\n",
"NSS runs the log-likelihood inside `jax.jit`. The boundary-based samplers above are happy with the default\n",
"NumPy `af.ex.Analysis(data, noise_map)` \u2014 when we hand the same analysis to NSS the JIT trace hits the\n",
"NumPy paths inside the analysis and raises `TracerArrayConversionError`. The fix is to build the analysis\n",
"with `use_jax=True`, which makes its internal arithmetic dispatch through `jax.numpy` instead of `numpy`.\n",
"\n",
"This is the production pattern: for autolens / autogalaxy / autofit analyses that you want to run with NSS,\n",
"construct your `Analysis` with `use_jax=True`. Everything below works identically to the NumPy path \u2014 same\n",
"`log_likelihood_function` API, same `Result` shape \u2014 but the body is now JAX-traceable."
]
},
{
"cell_type": "code",
"metadata": {},
"source": [
"try:\n",
" search = af.NSS(\n",
" path_prefix=path.join(\"searches\"),\n",
" name=\"NSS\",\n",
" n_live=200, # live particles maintained throughout the run\n",
" num_mcmc_steps=5, # slice-MCMC inner steps per dead-point batch\n",
" num_delete=50, # particles removed per outer iteration\n",
" termination=-3.0, # delta-logZ stopping criterion\n",
" seed=42, # JAX PRNG seed for reproducible runs\n",
" checkpoint_interval=100, # SLURM-friendly resume \u2014 see docstring above\n",
" )\n",
"except ImportError as exc:\n",
" print(f\"Skipping NSS example because the optional dependency stack is unavailable:\\n{exc}\")\n",
"else:\n",
" analysis_jax = af.ex.Analysis(data=data, noise_map=noise_map, use_jax=True)\n",
"\n",
" result = search.fit(model=model, analysis=analysis_jax)\n",
"\n",
" model_data = result.max_log_likelihood_instance.model_data_from(\n",
" xvalues=np.arange(data.shape[0])\n",
" )\n",
"\n",
" plt.errorbar(\n",
" x=range(data.shape[0]),\n",
" y=data,\n",
" yerr=noise_map,\n",
" linestyle=\"\",\n",
" color=\"k\",\n",
" ecolor=\"k\",\n",
" elinewidth=1,\n",
" capsize=2,\n",
" )\n",
" plt.plot(range(data.shape[0]), model_data, color=\"r\")\n",
" plt.title(\"NSS model fit to 1D Gaussian dataset.\")\n",
" plt.xlabel(\"x values of profile\")\n",
" plt.ylabel(\"Profile normalization\")\n",
" plt.show()\n",
" plt.close()\n",
"\n",
" print(f\"NSS log evidence: {result.samples.samples_info['log_evidence']:.4f}\")\n",
" print(f\"NSS max log L: {max(result.samples.log_likelihood_list):.4f}\")\n"
],
"outputs": [],
"execution_count": null
}
],
"metadata": {
Expand All @@ -562,4 +439,4 @@
},
"nbformat": 4,
"nbformat_minor": 4
}
}
114 changes: 1 addition & 113 deletions scripts/searches/nest.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,30 +7,13 @@
- `DynestyStatic`: Dynesty with static nested sampling.
- `DynestyDynamic`: Dynesty with dynamic nested sampling.
- `Nautilus`: Nautilus nested sampler.
- `NSS`: Nested Slice Sampling (JAX-native, optional install).

The first three are boundary-based samplers that work with any Python log-likelihood. `NSS` is a more recent
JAX-native sampler that runs its inner sampling loop inside `jax.jit` — when the log-likelihood is itself
JAX-traceable, the per-evaluation cost drops by roughly an order of magnitude versus the boundary samplers
(measured on real lensing likelihoods, see the `nss_first_class_sampler` roadmap and FINDINGS_v3 in the
profiling project for the numbers).
These are boundary-based samplers that work with any Python log-likelihood.

Relevant links:

- Dynesty: https://dynesty.readthedocs.io/en/latest/
- Nautilus: https://nautilus-sampler.readthedocs.io/en/stable/
- NSS (Nested Slice Sampling): https://github.com/yallup/nss

__Install Precondition for NSS__

The `Search: NSS` section at the bottom of this script imports the optional `nss` package. To run that
section, install the dependencies first:

pip install autofit[nss]

The extra pins the right `handley-lab/blackjax` fork at a known-good commit, so this is a single safe
command (no `--no-deps` dance, no manual git+ URLs). The other three samplers in this script have no
additional dependencies and run with the standard `pip install autofit` install.

__Contents__

Expand All @@ -41,7 +24,6 @@
- **Search: DynestyStatic**: Configuring and running the DynestyStatic nested sampler.
- **Search: DynestyDynamic**: Configuring and running the DynestyDynamic nested sampler.
- **Search: Nautilus**: Configuring and running the Nautilus nested sampler.
- **Search: NSS**: Configuring and running the NSS (Nested Slice Sampling) sampler.
- **Search Internal**: Accessing the internal sampler for advanced use (shown once for DynestyStatic).
"""

Expand Down Expand Up @@ -295,97 +277,3 @@
plt.ylabel("Profile normalization")
plt.show()
plt.close()

"""
__Search: NSS__

NSS (Nested Slice Sampling) is a JAX-native nested sampler whose inner sampling loop runs end-to-end inside
`jax.jit`. The advantage versus the boundary samplers above is **per-evaluation cost**: when your
log-likelihood is itself JAX-traceable, NSS avoids the Python ↔ JAX boundary that Nautilus and Dynesty cross
on every likelihood call. On the production lensing likelihoods that motivated this sampler, the per-eval
cost is roughly 30 times lower than Nautilus's, and total wall time to convergence drops from tens of
minutes to a few minutes.

On the trivial 1D Gaussian dataset used by this tutorial the speedup is not visible — the likelihood is so
cheap that the per-call cost is dominated by Python overhead rather than the floating-point work, and the
first NSS run pays a one-off ~25–30 second JIT compile that the simpler samplers above skip. The numbers
below let you confirm NSS converged to roughly the same posterior, not that it ran faster. Try NSS on a
real autolens or autogalaxy MGE / pixelization likelihood to see the per-eval advantage.

NSS exposes the same `result.samples` interface as the boundary samplers — swapping `af.Nautilus(...)` for
`af.NSS(...)` in your existing scripts is a one-line change.

A few NSS-specific kwargs to be aware of:

- `n_live`: live particles maintained throughout the run (production default 200).
- `num_mcmc_steps`: slice-MCMC inner steps per dead-point batch (production default 5).
- `num_delete`: particles removed per outer iteration (production default 50; larger values reduce JIT
overhead per outer iteration at the cost of slightly coarser posterior coverage).
- `termination`: stopping criterion on `logZ_live - logZ` (default `-3.0`, corresponding to a remaining
evidence fraction below 1e-3).
- `checkpoint_interval`: outer iterations between disk-saved checkpoints. NSS writes a resumable state
file every `checkpoint_interval` iterations, so a SLURM timeout halfway through a long fit is recovered
automatically the next time you run the same script with the same Paths.
- `iterations_per_quick_update`: when set non-None, NSS calls `analysis.visualize(...)` with the current
best live point every N outer iterations — partial results appear in the image_path directory while the
run is still in flight.

Settings reference: see `af.NSS.__init__` for the full kwarg list.

__Optional Dependency Guard__

The standard workspace release environment does not install `autofit[nss]`. When the optional stack is
absent, this script skips only the NSS fit and leaves the Dynesty / Nautilus examples above runnable.

__Analysis Must Be JAX-Traceable__

NSS runs the log-likelihood inside `jax.jit`. The boundary-based samplers above are happy with the default
NumPy `af.ex.Analysis(data, noise_map)` — when we hand the same analysis to NSS the JIT trace hits the
NumPy paths inside the analysis and raises `TracerArrayConversionError`. The fix is to build the analysis
with `use_jax=True`, which makes its internal arithmetic dispatch through `jax.numpy` instead of `numpy`.

This is the production pattern: for autolens / autogalaxy / autofit analyses that you want to run with NSS,
construct your `Analysis` with `use_jax=True`. Everything below works identically to the NumPy path — same
`log_likelihood_function` API, same `Result` shape — but the body is now JAX-traceable.
"""
try:
search = af.NSS(
path_prefix=path.join("searches"),
name="NSS",
n_live=200, # live particles maintained throughout the run
num_mcmc_steps=5, # slice-MCMC inner steps per dead-point batch
num_delete=50, # particles removed per outer iteration
termination=-3.0, # delta-logZ stopping criterion
seed=42, # JAX PRNG seed for reproducible runs
checkpoint_interval=100, # SLURM-friendly resume — see docstring above
)
except ImportError as exc:
print(f"Skipping NSS example because the optional dependency stack is unavailable:\n{exc}")
else:
analysis_jax = af.ex.Analysis(data=data, noise_map=noise_map, use_jax=True)

result = search.fit(model=model, analysis=analysis_jax)

model_data = result.max_log_likelihood_instance.model_data_from(
xvalues=np.arange(data.shape[0])
)

plt.errorbar(
x=range(data.shape[0]),
y=data,
yerr=noise_map,
linestyle="",
color="k",
ecolor="k",
elinewidth=1,
capsize=2,
)
plt.plot(range(data.shape[0]), model_data, color="r")
plt.title("NSS model fit to 1D Gaussian dataset.")
plt.xlabel("x values of profile")
plt.ylabel("Profile normalization")
plt.show()
plt.close()

print(f"NSS log evidence: {result.samples.samples_info['log_evidence']:.4f}")
print(f"NSS max log L: {max(result.samples.log_likelihood_list):.4f}")
Loading