Skip to content

Add cache-aware Nested Slice within Gibbs sampler #128

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

Open
yallup wants to merge 2 commits into
base: jim-dev
Choose a base branch
from
Open

Add cache-aware Nested Slice within Gibbs sampler #128

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b65a70c
Add cache-aware Nested SwiG sampler
yallup Jul 16, 2026
d1b5cd9
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jul 16, 2026
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
16 changes: 16 additions & 0 deletions docs/guides/cli.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -300,6 +300,7 @@ The `type` field selects the backend. Each backend has its own set of tuning par
| `flowmc` | Normalizing-flow MCMC | No |
| `blackjax-smc` | Sequential Monte Carlo | Yes |
| `blackjax-nss` | Nested slice sampling | Yes |
| `blackjax-swig` | Nested Slice within Gibbs with waveform caching | Yes |
| `blackjax-ns-aw` | Nested sampling (acceptance-walk) | Yes |

### `type = "flowmc"`
Expand Down Expand Up @@ -356,6 +357,21 @@ n_tempered_steps = 5
| `checkpoint_dir` | `{output.dir}/` | Directory for `checkpoint.pkl`; set by the CLI automatically |
| `checkpoint_interval` | `600.0` | Seconds between checkpoint writes; `0` disables checkpointing |

### `type = "blackjax-swig"`

| Field | Default | Description |
| --- | --- | --- |
| `blocks` | required | Ordered lists of sampling-space parameter names |
| `n_live` | `512` | Number of live points |
| `n_delete_frac` | `0.125` | Fraction of live points replaced per iteration |
| `num_inner_steps_per_dim` | `1` | Slice steps per dimension within each block |
| `num_gibbs_sweeps` | `2` | Complete block sweeps per replacement |
| `max_steps` | `10` | Maximum stepping-out expansions per slice |
| `max_shrinkage` | `100` | Maximum shrinkage evaluations per slice |
| `termination_dlogz` | `exp(-3)` (~`0.0498`) | Stop when the remaining evidence contribution falls below this |
| `checkpoint_dir` | `{output.dir}/` | Directory for `checkpoint.pkl`; set by the CLI automatically |
| `checkpoint_interval` | `600.0` | Seconds between checkpoint writes; `0` disables checkpointing |

### `type = "blackjax-ns-aw"`

Requires all sampling-space parameters to lie in $[0, 1]$. The CLI enforces this automatically.
Expand Down
56 changes: 55 additions & 1 deletion docs/guides/samplers.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ samples = jim.get_samples() # dict[str, np.ndarray] keyed by parameter name
| [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 |
| [SwiG](#blackjax-swig) | Nested Slice within Gibbs with waveform caching | Yes | Normalised prior |
| [SMC](#blackjax-smc) | Sequential Monte Carlo | Yes | Normalised prior |

---
Expand Down Expand Up @@ -168,6 +169,58 @@ Key parameters:

---

### BlackJAX SwiG

Nested Slice within Gibbs (SwiG) applies covariance-shaped slice updates to
named parameter blocks. With `TransientLikelihoodFD` or
`HeterodynedTransientLikelihoodFD`, Jim caches the generated waveform
polarizations: blocks that affect waveform inputs regenerate the cache, while
detector-projection-only blocks reuse it. Luminosity distance is also factored
out analytically, so a `d_L` block reuses the cache. Jim determines the other
dependencies after sample transforms, likelihood transforms, fixed parameters,
and analytic marginalisation have been applied.

```python
import math

from jimgw.samplers.config import BlackJAXSwiGConfig

jim = Jim(
likelihood,
prior,
sampler_config=BlackJAXSwiGConfig(
blocks=[
["M_c", "q", "lambda_1", "lambda_2"],
["s1_z", "s2_z"],
["iota"],
["ra", "dec"],
["psi"],
["t_c"],
],
n_live=512,
n_delete_frac=0.125,
num_inner_steps_per_dim=1,
num_gibbs_sweeps=2,
termination_dlogz=math.exp(-3.0),
),
likelihood_transforms=likelihood_transforms,
)
```

Blocks must form an exact partition of the sampling parameters. A block that
contains any waveform dependency is treated as cache-refreshing; mixed blocks
are therefore safe but may be less efficient. Parameters that are analytically
marginalised must not appear in the blocks.

Key parameters:

- `blocks` — ordered sampling-space parameter blocks for each Gibbs sweep.
- `n_live` / `n_delete_frac` — live-set size and replacement fraction, matching NSS.
- `num_inner_steps_per_dim` — random-direction slice steps per block dimension.
- `num_gibbs_sweeps` — complete block sweeps per particle replacement.

---

## Checkpointing and resuming

All samplers support checkpoint/resume so long-running jobs can survive interruptions.
Expand Down Expand Up @@ -205,7 +258,8 @@ jim = Jim(
jim.sample() # resumes from ./my_run/checkpoint.pkl
```

The same fields work identically for `FlowMCConfig`, `BlackJAXNSAWConfig`, and `BlackJAXNSSConfig`.
The same fields work identically for `FlowMCConfig`, `BlackJAXNSAWConfig`,
`BlackJAXNSSConfig`, and `BlackJAXSwiGConfig`.

| Field | Default | Notes |
| --- | --- | --- |
Expand Down
102 changes: 102 additions & 0 deletions examples/GW170817_SwiG.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,102 @@
"""GW170817-style 128 s BNS analysis with BlackJAX Nested SwiG.

This example deliberately uses Ripple's built-in aligned-spin
IMRPhenomD_NRTidalv2 model. Coalescence time and phase are analytically
marginalized; the remaining parameters are partitioned into waveform and
projection blocks.
"""

import time
import math

import jax
import jax.numpy as jnp
from jimgw.core.jim import Jim
from jimgw.core.prior import (
CombinePrior,
CosinePrior,
PowerLawPrior,
SinePrior,
UniformPrior,
)
from jimgw.core.single_event.data import Data
from jimgw.core.single_event.detector import get_H1, get_L1, get_V1
from jimgw.core.single_event.likelihood import TransientLikelihoodFD
from jimgw.core.single_event.transforms import MassRatioToSymmetricMassRatioTransform
from jimgw.core.single_event.waveform import RippleIMRPhenomD_NRTidalv2
from jimgw.samplers.config import BlackJAXSwiGConfig

jax.config.update("jax_enable_x64", True)
Comment on lines +12 to +29

@coderabbitai coderabbitai Bot Jul 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '\n== File: examples/GW170817_SwiG.py ==\n'
cat -n examples/GW170817_SwiG.py | sed -n '1,80p'

printf '\n== Search for jax_enable_x64 in repo ==\n'
rg -n "jax_enable_x64|enable_x64" .

printf '\n== Locate imported modules ==\n'
python3 - <<'PY'
from pathlib import Path
paths = [
    "jimgw/core/jim.py",
    "jimgw/core/prior.py",
    "jimgw/core/single_event/data.py",
    "jimgw/core/single_event/detector.py",
    "jimgw/core/single_event/likelihood.py",
    "jimgw/core/single_event/transforms.py",
    "jimgw/core/single_event/waveform.py",
    "jimgw/samplers/config.py",
]
for p in paths:
    fp = Path(p)
    print(f"\n== {p} ==")
    if fp.exists():
        print(fp.resolve())
        print("lines:", sum(1 for _ in fp.open()))
    else:
        print("MISSING")
PY

Repository: GW-JAX-Team/Jim

Length of output: 5196

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '\n== Relevant source files ==\n'
git ls-files 'src/jimgw/**/*.py' 'src/jimgw/*.py' | sed -n '1,200p'

printf '\n== Top of imported modules ==\n'
for f in \
  src/jimgw/core/jim.py \
  src/jimgw/core/prior.py \
  src/jimgw/core/single_event/data.py \
  src/jimgw/core/single_event/detector.py \
  src/jimgw/core/single_event/likelihood.py \
  src/jimgw/core/single_event/transforms.py \
  src/jimgw/core/single_event/waveform.py \
  src/jimgw/samplers/config.py
do
  if [ -f "$f" ]; then
    echo
    echo "== $f =="
    wc -l "$f"
    sed -n '1,220p' "$f" | cat -n | sed -n '1,120p'
  else
    echo
    echo "== $f =="
    echo "MISSING"
  fi
done

Repository: GW-JAX-Team/Jim

Length of output: 38525

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '\n== time_utils top ==\n'
cat -n src/jimgw/core/single_event/time_utils.py | sed -n '1,120p'

printf '\n== detector x64 checks ==\n'
sed -n '180,220p' src/jimgw/core/single_event/detector.py | cat -n

printf '\n== example import order samples ==\n'
for f in \
  examples/GW170817_multiband.py \
  examples/GW170817_heterodyne.py \
  examples/GW150914_NS_AW.py \
  examples/injection.py \
  examples/GW150914_flowMC.py \
  examples/GW150914_SMC.py \
  examples/GW150914_NSS.py
do
  echo
  echo "== $f =="
  cat -n "$f" | sed -n '1,25p'
done

Repository: GW-JAX-Team/Jim

Length of output: 12633

Move jax_enable_x64 above the jimgw imports.

src/jimgw/core/single_event/time_utils.py requires 64-bit mode at import time, so this example can fail before it reaches line 29. Put the config update immediately after the JAX imports, before any jimgw imports.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@examples/GW170817_SwiG.py` around lines 12 - 29, Move the jax.config.update
call enabling x64 immediately after the jax and jax.numpy imports, before all
jimgw imports in the example. Keep the existing configuration unchanged and
preserve the current import structure otherwise.



gps = 1187008882.43
duration = 128.0
start = gps + 2.0 - duration
end = start + duration
psd_start = start - 2048.0

ifos = [get_H1(), get_L1(), get_V1()]
for ifo in ifos:
strain = Data.from_gwosc(ifo.name, start, end)
ifo.set_data(strain)
psd_data = Data.from_gwosc(ifo.name, psd_start, start)
ifo.set_psd(
psd_data.to_psd(nperseg=int(strain.duration * strain.sampling_frequency))
)

distance_prior = PowerLawPrior(30.0, 150.0, 2.0, parameter_names=["d_L"])
prior = CombinePrior(
[
UniformPrior(1.18, 1.21, parameter_names=["M_c"]),
UniformPrior(0.5, 1.0, parameter_names=["q"]),
UniformPrior(-0.05, 0.05, parameter_names=["s1_z"]),
UniformPrior(-0.05, 0.05, parameter_names=["s2_z"]),
UniformPrior(0.0, 5000.0, parameter_names=["lambda_1"]),
UniformPrior(0.0, 5000.0, parameter_names=["lambda_2"]),
SinePrior(parameter_names=["iota"]),
distance_prior,
UniformPrior(0.0, 2.0 * jnp.pi, parameter_names=["ra"]),
CosinePrior(parameter_names=["dec"]),
UniformPrior(0.0, jnp.pi, parameter_names=["psi"]),
]
)

likelihood = TransientLikelihoodFD(
ifos,
waveform=RippleIMRPhenomD_NRTidalv2(f_ref=20.0),
trigger_time=gps,
f_min=20.0,
f_max=2048.0,
phase_marginalization=True,
time_marginalization={"tc_range": (-0.03, 0.03)},
)

jim = Jim(
likelihood,
prior,
likelihood_transforms=[MassRatioToSymmetricMassRatioTransform],
periodic={
"ra": (0.0, 2.0 * float(jnp.pi)),
"psi": (0.0, float(jnp.pi)),
},
sampler_config=BlackJAXSwiGConfig(
blocks=[
["M_c", "q", "lambda_1", "lambda_2"],
["s1_z", "s2_z"],
["iota"],
["d_L"],
["ra", "dec"],
["psi"],
],
n_live=512,
n_delete_frac=0.125,
num_inner_steps_per_dim=1,
num_gibbs_sweeps=2,
termination_dlogz=math.exp(-3.0),
),
)

start_time = time.time()
jim.sample()
print(f"Sampling took {(time.time() - start_time) / 60.0:.2f} minutes")
print(jim.get_diagnostics())
Loading
Loading