diff --git a/notebooks/searches/nest.ipynb b/notebooks/searches/nest.ipynb index 7d4fc902..59c7ea6c 100644 --- a/notebooks/searches/nest.ipynb +++ b/notebooks/searches/nest.ipynb @@ -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", @@ -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)." ] }, @@ -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." ] }, @@ -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": { @@ -562,4 +439,4 @@ }, "nbformat": 4, "nbformat_minor": 4 -} \ No newline at end of file +} diff --git a/scripts/searches/nest.py b/scripts/searches/nest.py index 7477ca3b..aae2d8e0 100644 --- a/scripts/searches/nest.py +++ b/scripts/searches/nest.py @@ -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__ @@ -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). """ @@ -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}")