diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 61ba881a..73dccc7d 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -27,7 +27,7 @@ jobs: run: uv python install ${{ matrix.python-version }} - name: Install dependencies run: | - uv sync --python ${{ matrix.python-version }} --group test --group cross-validation --group nested-sampling + uv sync --python ${{ matrix.python-version }} --group test --group cross-validation if [[ "${{ github.ref }}" == "refs/heads/jim-dev" || "${{ github.base_ref }}" == "jim-dev" ]]; then uv pip install git+https://github.com/GW-JAX-Team/flowMC.git@flowMC-dev uv pip install git+https://github.com/GW-JAX-Team/ripple.git@ripple-dev diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 067ca8f0..436bab83 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -23,4 +23,4 @@ repos: hooks: - id: pyright files: src/ - additional_dependencies: ["astropy>=7.0.1", "beartype>=0.19.0", "evosax>=0.2.0", "flowMC>=0.6.1", "gwpy>=3.0.12", "jax>=0.8.2", "jaxtyping>=0.3.3", pytest, "ripplegw>=0.2.1", typing_extensions, "blackjax@git+https://github.com/blackjax-devs/blackjax.git@main", "pydantic>=2", "anesthetic>=2", "tomli-w>=1.0", "typer>=0.12"] + additional_dependencies: ["astropy>=7.0.1", "beartype>=0.19.0", "evosax>=0.2.0", "flowMC>=0.6.1", "gwpy>=3.0.12", "jax>=0.8.2", "jaxtyping>=0.3.3", pytest, "ripplegw>=0.2.1", typing_extensions, "blackjax>=1.6", "pydantic>=2", "anesthetic>=2", "tomli-w>=1.0", "typer>=0.12"] diff --git a/docs/dev/blackjax_followups.md b/docs/dev/blackjax_followups.md deleted file mode 100644 index 15791fb1..00000000 --- a/docs/dev/blackjax_followups.md +++ /dev/null @@ -1,52 +0,0 @@ -# BlackJAX PyPI follow-ups - -## What we're waiting on - -A new PyPI release of [blackjax-devs/blackjax](https://github.com/blackjax-devs/blackjax) that includes: - -- Nested slice sampler (`blackjax.nss`, `blackjax.ns.{base,adaptive,utils}`). -- `blackjax.ns.utils.finalise` (used by both NS-AW and NSS at run time). - -These modules are already merged into `blackjax-devs/blackjax` main. -Jim's `[tool.uv.sources]` currently pins to that branch directly. - -## What to undo when a PyPI release ships - -### Packaging - -- Remove the `[tool.uv.sources]` block in `pyproject.toml` pointing `blackjax = { git = "https://github.com/blackjax-devs/blackjax.git" }`. -- Drop the `[dependency-groups] nested-sampling` PEP 735 group entirely. -- Bump the `blackjax>=1.6` pin to whatever release first contains all NS features. -- Remove the BlackJAX fork owned by GW JAX Team. - -### Inline imports → module top - -**`src/jimgw/samplers/blackjax/ns_aw.py`** - -| Line | Import | -|------|------------------------------------------| -| 144 | `from blackjax.ns.utils import finalise` | - -**`src/jimgw/samplers/blackjax/nss.py`** - -| Line | Import | -|------|------------------------------------------| -| 107 | `from blackjax.ns.utils import finalise` | - -### Docs - -- `docs/installation.md`: - - Drop the "BlackJAX nested samplers" section entirely. - - `pip install jimgw` is the canonical install instruction. -- `docs/guides/samplers.md`: - - Drop the "BlackJAX samplers" install preamble section. - - Restore `pip install jimgw` as the canonical install instruction. -- `README.md`: any BlackJAX install caveats can be deleted. - -### CI - -- `.github/workflows/CI.yml`: remove `--group nested-sampling` from the test job. The `[dependency-groups] nested-sampling` block in `pyproject.toml` can be deleted at the same time. - -### Tests - -- `tests/unit/samplers/test_blackjax_*.py`: the `pytest.importorskip("blackjax")` lines can stay as defense-in-depth but are no longer load-bearing. diff --git a/docs/guides/cli.md b/docs/guides/cli.md index a58e9391..f3ac1d38 100644 --- a/docs/guides/cli.md +++ b/docs/guides/cli.md @@ -295,12 +295,12 @@ The `type` field selects the backend. Each backend has its own set of tuning par ### Sampler comparison -| `type` | Algorithm | Evidence | Extra install | -| --- | --- | --- | --- | -| `flowmc` | Normalizing-flow MCMC | No | No | -| `blackjax-smc` | Sequential Monte Carlo | Yes | No | -| `blackjax-nss` | Nested slice sampling | Yes | `uv sync --group nested-sampling` | -| `blackjax-ns-aw` | Nested sampling (acceptance-walk) | Yes | `uv sync --group nested-sampling` | +| `type` | Algorithm | Evidence | +| --- | --- | --- | +| `flowmc` | Normalizing-flow MCMC | No | +| `blackjax-smc` | Sequential Monte Carlo | Yes | +| `blackjax-nss` | Nested slice sampling | Yes | +| `blackjax-ns-aw` | Nested sampling (acceptance-walk) | Yes | ### `type = "flowmc"` diff --git a/docs/guides/samplers.md b/docs/guides/samplers.md index 4f48b099..df03d5cb 100644 --- a/docs/guides/samplers.md +++ b/docs/guides/samplers.md @@ -11,12 +11,12 @@ samples = jim.get_samples() # dict[str, np.ndarray] keyed by parameter name ## Sampler overview -| Sampler | Algorithm | Evidence | Extra install | Prior constraint | -| --- | --- | --- | --- | --- | -| [flowMC](#flowmc) | normalizing-flow-enhanced MCMC | No | No | None | -| [NS-AW](#blackjax-ns-aw) | Nested sampling (bilby/dynesty-style acceptance-walk) | Yes | Yes (`nested-sampling`) | Uniform prior; unit-cube sampling space | -| [NSS](#blackjax-nss) | Nested slice sampling | Yes | Yes (`nested-sampling`) | Normalised prior | -| [SMC](#blackjax-smc) | Sequential Monte Carlo | Yes | No | Normalised prior | +| Sampler | Algorithm | Evidence | Prior constraint | +| --- | --- | --- | --- | +| [flowMC](#flowmc) | normalizing-flow-enhanced MCMC | No | None | +| [NS-AW](#blackjax-ns-aw) | Nested sampling (bilby/dynesty-style acceptance-walk) | Yes | Uniform prior; unit-cube sampling space | +| [NSS](#blackjax-nss) | Nested slice sampling | Yes | Normalised prior | +| [SMC](#blackjax-smc) | Sequential Monte Carlo | Yes | Normalised prior | --- @@ -92,21 +92,6 @@ Key parameters: --- -## BlackJAX nested samplers - -The two BlackJAX nested-sampling backends require additional dependencies. -The required modules are merged into upstream BlackJAX but not yet released on PyPI; install from main with: - -```bash -uv sync --group nested-sampling -``` - -This pulls in: - -- **blackjax** — pinned to `blackjax-devs/blackjax` until a PyPI release ships the nested-sampling modules. - ---- - ### BlackJAX NS-AW Nested sampling with a bilby/dynesty-style adaptive differential-evolution acceptance-walk inner kernel. diff --git a/docs/installation.md b/docs/installation.md index 102cfcad..5aeb7536 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -24,28 +24,3 @@ pip install -e . ``` We recommend using [uv](https://docs.astral.sh/uv/) to manage your Python environment. After cloning the repository, run `uv sync` to create a virtual environment with all dependencies installed. - -## BlackJAX nested samplers - -Jim's BlackJAX nested-sampling backends depend on modules not yet released on PyPI. -The features are merged into the upstream [blackjax-devs/blackjax](https://github.com/blackjax-devs/blackjax) main branch but await a versioned release. - -### With uv (recommended) - -Clone the repository if you haven't already, then sync the nested-sampling group: - -```bash -git clone https://github.com/GW-JAX-Team/Jim.git -cd Jim -uv sync --group nested-sampling -``` - -This installs BlackJAX from `blackjax-devs/blackjax` as declared in `[tool.uv.sources]`. - -### From source (pip) - -```bash -pip install "git+https://github.com/blackjax-devs/blackjax.git" -``` - -> **Note:** `pip install jimgw[nested-sampling]` will not work until the required features land in a PyPI release. Use one of the commands above instead. diff --git a/docs/tutorials/gw150914_cli.md b/docs/tutorials/gw150914_cli.md index 59fc9225..14f4b7cf 100644 --- a/docs/tutorials/gw150914_cli.md +++ b/docs/tutorials/gw150914_cli.md @@ -206,7 +206,7 @@ jim-run output/GW150914/config.final.toml ## 5. Next steps -- **Switch to a nested sampler** — change `[sampler]` to `type = "blackjax-ns-aw"` or `type = "blackjax-nss"` (requires `uv sync --group nested-sampling`) to obtain a log-evidence estimate. +- **Switch to a nested sampler** — change `[sampler]` to `type = "blackjax-ns-aw"` or `type = "blackjax-nss"` to obtain a log-evidence estimate. - **Try an injection** — change `[data]` to `type = "injection"` with `injection_parameters` to validate recovery on a known signal. See the [CLI Config Reference](../guides/cli.md) for the full injection syntax. - **Enable heterodyning** — add a `[likelihood.heterodyne]` block for a 10–100× likelihood speedup. See the [CLI Config Reference](../guides/cli.md). - **Full config reference** — [CLI Config Reference](../guides/cli.md) documents every available field and default. diff --git a/pyproject.toml b/pyproject.toml index 37bd3698..e2fc2925 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,7 +24,7 @@ dependencies = [ "anesthetic>=2", "astropy>=7.0.1", "beartype>=0.19.0", - "blackjax>=1.4", + "blackjax>=1.6", "corner>=2.2.3", "equinox>=0.11", "evosax>=0.2.0", @@ -70,10 +70,6 @@ doc = [ "pymdown-extensions", "zensical", ] -nested-sampling = [ - "blackjax", -] - [project.urls] Homepage = "https://github.com/GW-JAX-Team/Jim" Documentation = "https://gw-jax-team.github.io/Jim/" @@ -83,11 +79,6 @@ Changelog = "https://github.com/GW-JAX-Team/Jim/releases" [tool.uv.sources] jimgw = { workspace = true } -# Pinned to upstream main until a PyPI release ships with the NS submodules: -# - blackjax.ns.{base,adaptive,utils} (NS-AW, NSS) -# - blackjax.nss (NSS) -# When a release is cut, drop this pin and bump blackjax>=. -blackjax = { git = "https://github.com/blackjax-devs/blackjax.git" } [build-system] requires = ["hatchling", "hatch-vcs"] diff --git a/src/jimgw/samplers/__init__.py b/src/jimgw/samplers/__init__.py index e996ee99..8931416f 100644 --- a/src/jimgw/samplers/__init__.py +++ b/src/jimgw/samplers/__init__.py @@ -6,10 +6,8 @@ * `SamplerConfig` — discriminated-union annotation of concrete configs. * [`build_sampler`][jimgw.samplers.build_sampler] — factory that dispatches to the right concrete class. -The registry uses lazy loaders so that ``import jimgw.samplers`` does not fail -when an optional backend (e.g. BlackJAX NSS) is not installed; ImportError is -raised only when the caller actually asks for that backend via -`build_sampler`. +The registry uses loaders to defer concrete sampler construction until the +caller requests a backend via `build_sampler`. """ from collections.abc import Callable @@ -47,8 +45,7 @@ def register_sampler(type_str: str, lazy_loader: Callable[[], SamplerBuilder]) - """Register a concrete [`Sampler`][jimgw.samplers.base.Sampler] class under ``type_str``. ``lazy_loader`` is called (with no args) only when `build_sampler` - dispatches to this type — this is how we defer BlackJAX imports until - someone actually asks for a BlackJAX sampler. + dispatches to this type. """ if type_str in _REGISTRY: raise ValueError( @@ -81,8 +78,6 @@ def build_sampler( Raises: KeyError: If no sampler is registered for ``config.type``. - ImportError: If the lazy loader for that type fails (e.g. BlackJAX - missing when requesting a BlackJAX sampler). """ type_str = config.type if type_str not in _REGISTRY: @@ -109,29 +104,10 @@ def build_sampler( register_sampler("blackjax-smc", lambda: BlackJAXSMCSampler) -# --- BlackJAX NS-AW / NSS (require `uv sync --group nested-sampling`) --- +from jimgw.samplers.blackjax.ns_aw import BlackJAXNSAWSampler # noqa: E402 +register_sampler("blackjax-ns-aw", lambda: BlackJAXNSAWSampler) -def _load_ns_aw() -> SamplerBuilder: - import blackjax - from jimgw.samplers.blackjax._imports import require_nested_sampling +from jimgw.samplers.blackjax.nss import BlackJAXNSSSampler # noqa: E402 - require_nested_sampling(blackjax) - from jimgw.samplers.blackjax.ns_aw import BlackJAXNSAWSampler - - return BlackJAXNSAWSampler - - -def _load_nss() -> SamplerBuilder: - import blackjax - from jimgw.samplers.blackjax._imports import require_nested_sampling, require_nss - - require_nested_sampling(blackjax) - require_nss(blackjax) - from jimgw.samplers.blackjax.nss import BlackJAXNSSSampler - - return BlackJAXNSSSampler - - -register_sampler("blackjax-ns-aw", _load_ns_aw) -register_sampler("blackjax-nss", _load_nss) +register_sampler("blackjax-nss", lambda: BlackJAXNSSSampler) diff --git a/src/jimgw/samplers/blackjax/__init__.py b/src/jimgw/samplers/blackjax/__init__.py index 80b70cde..4cb5c9d8 100644 --- a/src/jimgw/samplers/blackjax/__init__.py +++ b/src/jimgw/samplers/blackjax/__init__.py @@ -1,5 +1 @@ -"""BlackJAX-backed samplers for Jim. - -SMC uses ``blackjax>=1.4`` (a core dependency). -NS-AW and NSS additionally need: ``uv sync --group nested-sampling`` -""" +"""BlackJAX-backed samplers for Jim.""" diff --git a/src/jimgw/samplers/blackjax/_imports.py b/src/jimgw/samplers/blackjax/_imports.py deleted file mode 100644 index 43ca0102..00000000 --- a/src/jimgw/samplers/blackjax/_imports.py +++ /dev/null @@ -1,31 +0,0 @@ -"""Lazy import of BlackJAX with feature validation. - -NS-AW and NSS rely on nested-sampling submodules merged into upstream BlackJAX -but not yet released on PyPI. Install them via the ``nested-sampling`` group: - - uv sync --group nested-sampling -""" - -_INSTALL_MSG = ( - "The BlackJAX nested-sampling submodules are required for this sampler " - "but are not available in the installed version. Install them with:\n" - " uv sync --group nested-sampling\n" - "See docs/installation.md for details." -) - - -def require_nested_sampling(bjx) -> None: - """Check that the installed BlackJAX has the ``ns`` nested-sampling submodule.""" - if not hasattr(bjx, "ns"): - raise ImportError( - "Installed BlackJAX is missing the `blackjax.ns` nested-sampling " - "submodule. " + _INSTALL_MSG - ) - - -def require_nss(bjx) -> None: - """Check that the installed BlackJAX exposes top-level ``blackjax.nss``.""" - if not hasattr(bjx, "nss"): - raise ImportError( - "Installed BlackJAX is missing top-level `blackjax.nss`. " + _INSTALL_MSG - ) diff --git a/src/jimgw/samplers/blackjax/ns_aw.py b/src/jimgw/samplers/blackjax/ns_aw.py index 9c17dab3..49b80ae4 100644 --- a/src/jimgw/samplers/blackjax/ns_aw.py +++ b/src/jimgw/samplers/blackjax/ns_aw.py @@ -11,21 +11,17 @@ import numpy as np from jaxtyping import Array, Float, Key from anesthetic.samples import NestedSamples -import blackjax from blackjax.ns.adaptive import AdaptiveNSState from blackjax.ns.base import NSInfo from blackjax.ns.utils import finalise from jimgw.samplers.base import Sampler from jimgw.samplers.blackjax._acceptance_walk_kernel import bilby_adaptive_de_sampler -from jimgw.samplers.blackjax._imports import require_nested_sampling from jimgw.samplers.config import BlackJAXNSAWConfig from jimgw.samplers.periodic import to_unit_cube_stepper logger = logging.getLogger(__name__) -require_nested_sampling(blackjax) - class BlackJAXNSAWSampler(Sampler): """BlackJAX nested sampler using the bilby/dynesty-style adaptive DE acceptance-walk kernel. diff --git a/src/jimgw/samplers/blackjax/nss.py b/src/jimgw/samplers/blackjax/nss.py index 0e28ff94..c6d65e9d 100644 --- a/src/jimgw/samplers/blackjax/nss.py +++ b/src/jimgw/samplers/blackjax/nss.py @@ -19,18 +19,11 @@ from blackjax.ns.nss import live_covariance, sample_direction_from_covariance from blackjax.ns.utils import finalise from jimgw.samplers.base import Sampler -from jimgw.samplers.blackjax._imports import ( - require_nested_sampling, - require_nss, -) from jimgw.samplers.config import BlackJAXNSSConfig from jimgw.samplers.periodic import to_prior_space_proposal logger = logging.getLogger(__name__) -require_nested_sampling(blackjax) -require_nss(blackjax) - class BlackJAXNSSSampler(Sampler): """BlackJAX Nested Slice Sampler (NSS). diff --git a/tests/unit/samplers/blackjax/test_imports.py b/tests/unit/samplers/blackjax/test_imports.py deleted file mode 100644 index 5e8584c7..00000000 --- a/tests/unit/samplers/blackjax/test_imports.py +++ /dev/null @@ -1,43 +0,0 @@ -"""Tests for the BlackJAX import guards in _imports.py.""" - -from __future__ import annotations - -import types - -import pytest - -from jimgw.samplers.blackjax._imports import ( - require_nested_sampling, - require_nss, -) - -# --------------------------------------------------------------------------- -# require_nested_sampling -# --------------------------------------------------------------------------- - - -def test_require_nested_sampling_ok(): - fake = types.SimpleNamespace(ns=object()) - require_nested_sampling(fake) # must not raise - - -def test_require_nested_sampling_missing(): - fake = types.SimpleNamespace() # no `ns` - with pytest.raises(ImportError, match=r"blackjax\.ns"): - require_nested_sampling(fake) - - -# --------------------------------------------------------------------------- -# require_nss -# --------------------------------------------------------------------------- - - -def test_require_nss_ok(): - fake = types.SimpleNamespace(nss=object()) - require_nss(fake) # must not raise - - -def test_require_nss_missing(): - fake = types.SimpleNamespace() # no `nss` - with pytest.raises(ImportError, match=r"blackjax\.nss"): - require_nss(fake) diff --git a/tests/unit/samplers/test_registry.py b/tests/unit/samplers/test_registry.py index 6b2c6ddf..a03c475e 100644 --- a/tests/unit/samplers/test_registry.py +++ b/tests/unit/samplers/test_registry.py @@ -1,13 +1,10 @@ """Tests for the sampler registry and build_sampler factory.""" -import sys - import pytest from jimgw.core.prior import CombinePrior, UniformPrior # type: ignore[attr-defined] from jimgw.samplers import build_sampler from jimgw.samplers.config import ( - BlackJAXNSAWConfig, FlowMCConfig, ) from jimgw.samplers.flowmc import FlowMCSampler # type: ignore[import] @@ -80,23 +77,6 @@ class _FakeConfig(BaseSamplerConfig): ) -def test_build_sampler_blackjax_raises_import_error_when_missing(monkeypatch): - """When blackjax is not installed, requesting a BlackJAX sampler should raise ImportError.""" - monkeypatch.setitem(sys.modules, "blackjax", None) # type: ignore[arg-type] - - prior = _make_prior() - lp, ll, lpost = _make_callables(prior) - cfg = BlackJAXNSAWConfig() - with pytest.raises((ImportError, KeyError)): - build_sampler( - cfg, - n_dims=1, - log_prior_fn=lp, - log_likelihood_fn=ll, - log_posterior_fn=lpost, - ) - - def test_registry_has_all_four_types(): from jimgw.samplers import _REGISTRY diff --git a/uv.lock b/uv.lock index cce1be2c..60f6f198 100644 --- a/uv.lock +++ b/uv.lock @@ -197,8 +197,8 @@ wheels = [ [[package]] name = "blackjax" -version = "1.7.dev10+g2dcea81d9" -source = { git = "https://github.com/blackjax-devs/blackjax.git#2dcea81d9d8cdf768b5e60627ea304f1cd450c84" } +version = "1.6" +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "jax" }, { name = "jaxlib" }, @@ -209,6 +209,10 @@ dependencies = [ { name = "scipy", version = "1.18.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, { name = "typing-extensions" }, ] +sdist = { url = "https://files.pythonhosted.org/packages/42/0f/9da22812a20bd4f7da5b4f67ad79f459324c701fcca434ce583d90bd32e4/blackjax-1.6.tar.gz", hash = "sha256:fc1af2565bcf129974daa5d25067616aea85ea3e22e35d57d492d14bcd74daa5", size = 5174101, upload-time = "2026-07-10T20:30:31.99Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5f/64/3d9dcc0e3fbba2de25c39adf302b6ad2ca3c119299ac087b8002161958b8/blackjax-1.6-py3-none-any.whl", hash = "sha256:d1c135f448789961a4bada1e2eb803a1da55c0622c5f77eddce80bb3f7351fd5", size = 4857402, upload-time = "2026-07-10T20:30:30.492Z" }, +] [[package]] name = "bleach" @@ -1400,9 +1404,6 @@ doc = [ { name = "pymdown-extensions" }, { name = "zensical" }, ] -nested-sampling = [ - { name = "blackjax" }, -] test = [ { name = "coverage" }, { name = "coveralls" }, @@ -1418,7 +1419,7 @@ requires-dist = [ { name = "anesthetic", specifier = ">=2" }, { name = "astropy", specifier = ">=7.0.1" }, { name = "beartype", specifier = ">=0.19.0" }, - { name = "blackjax", git = "https://github.com/blackjax-devs/blackjax.git" }, + { name = "blackjax", specifier = ">=1.6" }, { name = "corner", specifier = ">=2.2.3" }, { name = "equinox", specifier = ">=0.11" }, { name = "evosax", specifier = ">=0.2.0" }, @@ -1449,7 +1450,6 @@ doc = [ { name = "pymdown-extensions" }, { name = "zensical" }, ] -nested-sampling = [{ name = "blackjax", git = "https://github.com/blackjax-devs/blackjax.git" }] test = [ { name = "coverage", specifier = ">=7.7.0" }, { name = "coveralls", specifier = ">=4.0.1" },