Add cache-aware Nested Slice within Gibbs sampler#128
Open
yallup wants to merge 2 commits into
Open
Conversation
Contributor
📝 WalkthroughWalkthroughAdds the BlackJAX SwiG sampler with block-based cached waveform evaluation, configuration validation, Jim integration, documentation, a GW170817 example, and unit tests covering configuration, cache reuse, block validation, registry wiring, and sampler behaviour. ChangesBlackJAX SwiG
Estimated code review effort: 4 (Complex) | ~45 minutes Suggested labels: Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant Jim
participant BlackJAXSwiGSampler
participant TransientLikelihoodFD
participant NestedSampler
Jim->>BlackJAXSwiGSampler: build sampler with blocks and cache hooks
BlackJAXSwiGSampler->>NestedSampler: construct nested sampling algorithm
NestedSampler->>BlackJAXSwiGSampler: execute constrained Gibbs slice step
BlackJAXSwiGSampler->>TransientLikelihoodFD: refresh or reuse waveform cache
TransientLikelihoodFD-->>BlackJAXSwiGSampler: return cached log-likelihood
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
yallup
marked this pull request as ready for review
July 20, 2026 10:09
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
src/jimgw/samplers/blackjax/swig.py (1)
99-133: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winBind loop variables into the slice closures to avoid late-binding fragility.
proposal_generator/slice_fnclose overblock_array,covariance, andmust_refreshfrom the enclosingfor block, must_refresh, covariance in zip(...)loop (Ruff B023). It is correct today only becausejax.lax.scantracesone_sliceeagerly within the same iteration; any future change that defers execution would silently reuse the last block's values. Binding them as default arguments removes the hazard and clears the linter.♻️ Bind loop variables explicitly
- def one_slice(carry, key): + def one_slice( + carry, + key, + block_array=block_array, + covariance=covariance, + must_refresh=must_refresh, + ): current, all_accepted, expansions, shrink = carry🤖 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 `@src/jimgw/samplers/blackjax/swig.py` around lines 99 - 133, Bind the enclosing loop variables block_array, covariance, and must_refresh explicitly in the proposal_generator and slice_fn closures within one_slice, using closure/default-argument binding so each iteration retains its own values. Preserve the existing proposal and cache-refresh behavior while eliminating late binding and Ruff B023 warnings.Source: Linters/SAST tools
🤖 Prompt for all review comments with 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.
Inline comments:
In `@examples/GW170817_SwiG.py`:
- Around line 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.
---
Nitpick comments:
In `@src/jimgw/samplers/blackjax/swig.py`:
- Around line 99-133: Bind the enclosing loop variables block_array, covariance,
and must_refresh explicitly in the proposal_generator and slice_fn closures
within one_slice, using closure/default-argument binding so each iteration
retains its own values. Preserve the existing proposal and cache-refresh
behavior while eliminating late binding and Ruff B023 warnings.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: d89b4ef6-b935-4d32-8557-7cd526f6c398
📒 Files selected for processing (13)
docs/guides/cli.mddocs/guides/samplers.mdexamples/GW170817_SwiG.pysrc/jimgw/core/jim.pysrc/jimgw/core/single_event/likelihood.pysrc/jimgw/samplers/__init__.pysrc/jimgw/samplers/blackjax/nss.pysrc/jimgw/samplers/blackjax/swig.pysrc/jimgw/samplers/config.pytests/unit/core/single_event/test_likelihood.pytests/unit/samplers/blackjax/test_swig.pytests/unit/samplers/test_config.pytests/unit/samplers/test_registry.py
Comment on lines
+12
to
+29
| 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) |
Contributor
There was a problem hiding this comment.
🎯 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")
PYRepository: 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
doneRepository: 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'
doneRepository: 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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
BlackJAXSwiGConfigand theblackjax-swigNested Slice within Gibbs backendd_Lupdates reuse its1_zands2_zin one joint slow block and use Ripple's built-inIMRPhenomD_NRTidalv2d_Land marginalises phase and coalescence time over(-0.03, 0.03)Motivation
Long BNS waveforms make repeated waveform generation the dominant cost. Nested Slice within Gibbs separates intrinsic waveform updates from detector projection and distance-amplitude updates, allowing the latter to reuse cached polarizations while retaining the existing BlackJAX nested-sampling runner and evidence calculation.
Validation
208 passedacross the relevant Jim, likelihood, NSS, SwiG, registry, and config unit suitesd_Lreuses the cache under time marginalisationgit diff --checkpassOutstanding validation
The GW170817 example has not yet been fully replicated end-to-end on a GPU. Its block layout, binning, and marginalisation settings may need further tuning before it is treated as a production analysis.
Related to #124. This PR uses the existing time-grid marginalisation but does not add Bilby-style
jitter_time.Summary by CodeRabbit
New Features
Documentation
Bug Fixes