Fix and improve relative binning in HeterodynedTransientLikelihoodFD

[thomasckng]

Closes #108.

Summary

This PR corrects three bugs in HeterodynedTransientLikelihoodFD and aligns the implementation with Bilby's relative-binning reference (arXiv:1806.08792).

---

Changes

1. Correct gamma set in _max_phase_diff

The old implementation used gamma = arange(-5, 6) / 3 (11 terms, including zero) with a f_star-based normalisation that is not equivalent to Bilby's. Replaced with Bilby's 5 physically-motivated PN/IMR terms:

$$\gamma \in {-5/3,\ -2/3,\ 1,\ 5/3,\ 7/3}$$

Each term is normalised by d_alpha so that its contribution spans exactly chi * 2π rad across [f_low, f_high]. The array is offset so it starts at zero (cumulative phase from f_low).

2. Correct r0/r1 estimators in _likelihood

The old code estimated the ratio polynomial using a center–left off-centre finite difference:

r0 = waveform_center / waveform_center_ref
r1 = (waveform_low / waveform_low_ref - r0) / (f_low - f_center)

Replaced with Bilby's bin-edge average:

r0 = (r_low + r_high) / 2
r1 = (r_high - r_low) / bin_width

This requires evaluating the reference and proposal waveforms at both bin edges. waveform_high_ref (evaluated at freq_grid_high) replaces waveform_center_ref, and bin_widths is stored as an instance attribute.

3. Adaptive epsilon-based binning

Added an epsilon parameter — the maximum allowed phase change per bin in radians — as the primary way to control bin count, matching Bilby's interface:

n_bins = max(1, int(total_phase / epsilon))

• epsilon=0.5 (the new default when neither argument is given) yields ~60 bins for a typical CBC signal with f=[20, 512] Hz.
• n_bins is still accepted but is now mutually exclusive with epsilon; passing both raises ValueError.
• Non-positive values of either parameter raise ValueError.
• The CLI (CLIHeterodynedConfig) gains a corresponding epsilon field; both default to None.

4. Fix NaN at waveform boundary

The masking condition used freq_grid_center <= f_waveform_max, which could include bins whose upper edge (freq_grid_high) extended beyond the waveform's valid range. At those frequencies the reference waveform is zero, causing r_high = 0 / 0 = NaN. Fixed by using freq_grid_high <= f_waveform_max as the upper-bound condition.

5. Rename helper methods to private

max_phase_diff, make_binning_scheme, and compute_coefficients are internal helpers not intended as public API. Renamed to _max_phase_diff, _make_binning_scheme, and _compute_coefficients. _make_binning_scheme no longer accepts epsilon; the resolution from epsilon to integer n_bins is done in __init__ before calling it.

6. CLI and docs

• CLIHeterodynedConfig gains epsilon: Optional[float] = None alongside n_bins: Optional[int] = None.
• build_likelihood passes epsilon to the constructor and now logs likelihood.n_bins (the actual post-construction bin count) rather than the config value.
• docs/guides/cli.md updated to reflect the new epsilon-based TOML configuration.
• All docstrings updated.

---

Tests

• Unit tests updated: waveform_center_ref → waveform_high_ref; freq_grid_high attribute assertion added. All 13 tests pass.
• Cross-validation tests updated: use epsilon=0.5 directly instead of matching bilby_likelihood.number_of_bins.

Summary by CodeRabbit

Release Notes

• New Features

  • Added alternative binning configuration option for likelihood calculations based on phase tolerance.

• Changes

  • Modified default binning configuration method.

• Documentation

  • Updated CLI guidance to document new binning configuration options.

[thomasckng]

@coderabbitai review

[coderabbitai]

✅ Actions performed

Review triggered.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[coderabbitai]

Warning

Review limit reached

@thomasckng, we couldn't start this review because you've reached your PR review rate limit.

More reviews will be available in 45 minutes and 48 seconds. Learn how PR review limits work.

Your organization has run out of usage credits. Purchase more in the billing tab.

⌛ How to resolve this issue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

We recommend that you space out your commits to avoid hitting the rate limit.

🚦 How do rate limits work?

CodeRabbit enforces hourly rate limits for each developer per organization.

Our paid plans include higher PR review limits than trial, open-source, and free plans. In all cases, reviews become available again over time. During sustained high-volume PR review activity, CodeRabbit may temporarily slow when the next review becomes available.

Please see our Fair Usage Limits Policy for further information.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: a9573376-3481-4a82-8dcf-8390409db75b

📥 Commits

Reviewing files that changed from the base of the PR and between a59032b and eac5db9.

📒 Files selected for processing (1)

• tests/cross_validation/test_likelihood.py

📝 Walkthrough

Walkthrough

The pull request refactors heterodyned relative binning from a fixed 1000-bin default to support configurable phase-based (epsilon) or fixed-count (n_bins) binning. CLI configuration and likelihood builder integrate the new epsilon parameter; core likelihood implementation rewrites binning scheme computation, reference waveform storage from centre to separate low/high edges, and likelihood evaluation accordingly. Tests validate the new scheme.

Changes

Heterodyned Likelihood Binning Refactor

Layer / File(s)	Summary
CLI Configuration and Likelihood Builder Integration docs/guides/cli.md, src/jimgw/cli/_config.py, src/jimgw/cli/_likelihood.py	Documentation updated to describe epsilon and n_bins as mutually exclusive binning options with epsilon=0.5 default. CLIHeterodynedConfig schema changed to n_bins: Optional[int] = None and added epsilon: Optional[float] = None. Likelihood builder forwards cfg.heterodyne.epsilon to constructor and logs the computed likelihood.n_bins instead of config value.
Constructor Signature and Binning Setup Logic src/jimgw/core/single_event/likelihood.py (lines 537–705)	HeterodynedTransientLikelihoodFD.__init__ signature updated to accept n_bins: Optional[int] and new epsilon: Optional[float] parameter. Initialisation validates mutual exclusivity, defaults epsilon to 0.5 when unset, derives n_bins from maximum phase difference, generates bin edges using private scheme helper, and stores freq_grid_low, freq_grid_high, bin_widths, and computed n_bins.
Phase-Difference and Binning Scheme Helpers src/jimgw/core/single_event/likelihood.py (lines 815–862)	Private helpers _max_phase_diff (with updated PN/IMR term weights) and _make_binning_scheme (with linear phase-space binning and frequency interpolation) replace previous static methods and implement phase-based bin edge computation.
Reference Waveform Storage and Coefficient Computation src/jimgw/core/single_event/likelihood.py (lines 564–744)	Typed attributes changed from centre-frequency reference to separate waveform_low_ref and waveform_high_ref. Frequency overlap region filtered using low/high edge constraints. Per-detector reference waveforms constructed at both edge frequencies; coefficient computation uses renamed private _compute_coefficients helper.
Likelihood Evaluation with Low/High Edge Grids src/jimgw/core/single_event/likelihood.py (lines 768–786)	_likelihood method rewritten to evaluate waveform responses at freq_grid_low and freq_grid_high, compute low/high ratio terms, construct r0/r1 from those ratios, and accumulate match-filter, optimal-SNR, and phase-marginalised contributions using precomputed coefficient arrays.
Cross-Validation and Unit Test Updates tests/cross_validation/test_likelihood.py, tests/unit/core/single_event/test_likelihood.py	TestHeterodynedTransientLikelihoodFD.test_base and test_phase_marginalized updated to specify epsilon=0.5 instead of deriving n_bins from Bilby likelihood. Unit test attributes assertion updated to check waveform_high_ref instead of waveform_center_ref.

🎯 4 (Complex) | ⏱️ ~45 minutes

Suggested labels

enhancement

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately describes the primary change: fixing and improving relative binning in HeterodynedTransientLikelihoodFD, which is the core focus of all code modifications.
Linked Issues check	✅ Passed	All coding objectives from issue #108 are met: differences documented, binning made configurable via epsilon, Jim now aligns with Bilby's implementation through corrected gamma terms and bin-edge estimators.
Out of Scope Changes check	✅ Passed	All changes are directly related to fixing and improving relative binning; no unrelated modifications detected.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch change-RB

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

tests/cross_validation/test_likelihood.py (1)

679-699: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Pin bilby epsilon explicitly and fix the stale comment.

• tests/cross_validation/test_likelihood.py sets Jim with epsilon=0.5 in both cross-validation segments, but both bilby RelativeBinningGravitationalWaveTransient(...) initialisations omit epsilon (relying on bilby’s default).
• With this repo’s pinned bilby==2.8.0, the default epsilon is already 0.5 (so today the two sides match), but passing epsilon=0.5 prevents future drift if bilby’s default changes.
• The nearby comments still describe “reading back the number of bins” / an n_bins matching flow, which the current code doesn’t appear to use.

Proposed patch

-        # Initialize bilby first with its default epsilon so we can read back
-        # the number of bins it chose, then initialize Jim with the same count
-        # for an apples-to-apples comparison.
+        # Use explicit epsilon on both sides to avoid default-driven drift
+        # across bilby versions.
         bilby_likelihood = (
             bilby.gw.likelihood.RelativeBinningGravitationalWaveTransient(
                 interferometers=setup["bilby_ifos"],
                 waveform_generator=build_bilby_waveform_generator(
                     setup["duration"], setup["sampling_frequency"]
                 ),
                 fiducial_parameters=bilby_ref_params,
+                epsilon=0.5,
             )
         )
@@
-        # Initialize bilby first with its default epsilon so we can read back
-        # the number of bins it chose, then initialize Jim with the same count
-        # for an apples-to-apples comparison.
+        # Use explicit epsilon on both sides to avoid default-driven drift
+        # across bilby versions.
         bilby_likelihood = (
             bilby.gw.likelihood.RelativeBinningGravitationalWaveTransient(
                 interferometers=setup["bilby_ifos"],
                 waveform_generator=build_bilby_waveform_generator(
                     setup["duration"], setup["sampling_frequency"]
                 ),
                 fiducial_parameters=bilby_ref_params,
                 phase_marginalization=True,
                 priors=priors,
+                epsilon=0.5,
             )
         )

🤖 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 `@tests/cross_validation/test_likelihood.py` around lines 679 - 699, The test
currently relies on bilby's default epsilon and has a stale comment about
reading back number of bins; explicitly pass epsilon=0.5 to the bilby
RelativeBinningGravitationalWaveTransient construction (the same value used when
creating HeterodynedTransientLikelihoodFD) to avoid future drift, and update the
surrounding comment near bilby_likelihood /
RelativeBinningGravitationalWaveTransient to remove or replace the "reading back
the number of bins / n_bins matching" wording so it accurately reflects that
epsilon is being pinned for parity.

🤖 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 `@src/jimgw/core/single_event/likelihood.py`:
- Around line 567-568: Normalize all jaxtyping shape strings in forward
annotations in src/jimgw/core/single_event/likelihood.py: remove
leading/trailing spaces inside the quoted dimension names, replace empty-string
dims with meaningful identifiers, and convert any arithmetic or expression-like
strings into single valid identifiers (e.g. change " n_bin" → "n_bin", "" →
"n_bin" or "n_freq" as appropriate, "n_bins + 1" → "n_bins_plus_1", "n_bands+1
2" → "n_bands_plus_1"). Apply these fixes to the annotated variables and
signatures such as waveform_low_ref and waveform_high_ref and the other forward
annotations flagged (around the original comment locations) so that every
Float[Array, "..."] uses a parsable identifier-only dimension name.

---

Outside diff comments:
In `@tests/cross_validation/test_likelihood.py`:
- Around line 679-699: The test currently relies on bilby's default epsilon and
has a stale comment about reading back number of bins; explicitly pass
epsilon=0.5 to the bilby RelativeBinningGravitationalWaveTransient construction
(the same value used when creating HeterodynedTransientLikelihoodFD) to avoid
future drift, and update the surrounding comment near bilby_likelihood /
RelativeBinningGravitationalWaveTransient to remove or replace the "reading back
the number of bins / n_bins matching" wording so it accurately reflects that
epsilon is being pinned for parity.

🪄 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: 04f5ad8f-9fa8-46f0-922e-8202300a1be1

📥 Commits

Reviewing files that changed from the base of the PR and between ae9d317 and a59032b.

📒 Files selected for processing (6)

• docs/guides/cli.md
• src/jimgw/cli/_config.py
• src/jimgw/cli/_likelihood.py
• src/jimgw/core/single_event/likelihood.py
• tests/cross_validation/test_likelihood.py
• tests/unit/core/single_event/test_likelihood.py