Restructure valid waveform mask in Heterodyne likelihood

[SSL32081]

This PR refactors much of the code of the masking applied to the waveform and frequencies during initialisation of the Heterodyned Likelihood. This mask aimed to mask out non-physical/invalid and trivial waveform pieces, which was introduced in kazewong/jim@93023114b.

The mask is defined as:

h_amp = jnp.sum(
    jnp.array([jnp.abs(h_sky[pol]) for pol in h_sky.keys()]), axis=0
)
f_valid = frequency_original[jnp.where(h_amp > 0)[0]]

mask_heterodyne_center = jnp.where(
    (self.freq_grid_center <= jnp.max(f_valid))
    & (self.freq_grid_center >= jnp.min(f_valid))
)[0]

At the time, it was necessary to mask out the NaNs from the waveform, since the self.frequencies attribute at the time begins at zero (see e.g. here), where the waveform evaluation is invalid. But this is no longer needed with the introduction of sliced_frequencies in kazewong/jim@b745e04.

That being said, the mask is still needed to remove the waveform where it is trivial (= 0), which occurs soon after the merger frequency. For otherwise it would lead to NaN (from 0/0) later when taking the ratio between waveform pieces.

Summary by CodeRabbit [edited]

• Improvements
  • Enhanced consistency and stability of frequency binning for heterodyned transient likelihood calculations.
  • Improved likelihood coefficient normalisation for more robust evaluation across detector duration.
• Refactor
  • Updated bin-boundary handling to simplify and align coefficient preparation while preserving expected results.
• Tests
  • Strengthened likelihood initialisation checks (including finiteness validation) and updated test expectations for low-frequency reference cutoff behaviour.

[coderabbitai]

Important

Review skipped

Review was skipped due to path filters

⛔ Files ignored due to path filters (1)

• uv.lock is excluded by !**/*.lock

CodeRabbit blocks several paths by default. You can override this behavior by explicitly including those paths in the path filters. For example, including **/dist/** will override the default block on the dist directory, by removing the pattern from both the lists.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 6069a113-df94-497f-a867-39901af7fadb

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

HeterodynedTransientLikelihoodFD now uses masked bin edges as the primary binning representation, derives bin centres internally, and normalises coefficients using detector duration. Tests validate finite stored arrays and the updated coefficient interface.

Changes

Heterodyne bin-edge representation

Layer / File(s)	Summary
Bin-edge binning initialisation src/jimgw/core/single_event/likelihood.py	Likelihood annotations and initialisation now derive valid frequency bounds, widths, reference responses, and coefficients from masked bin edges.
Binning helper and coefficient contracts src/jimgw/core/single_event/likelihood.py	_make_binning_scheme returns edges only; _mask_and_set_frequency_arrays establishes valid frequency metadata; _compute_coefficients derives centres internally and uses detector duration for normalisation.
Likelihood validation updates tests/unit/core/single_event/test_likelihood.py	Tests check finite stored detector arrays and adapt the coefficient test double to receive bin edges.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Possibly related PRs

• GW-JAX-Team/Jim#116: Reworks related relative-binning and heterodyne logic in HeterodynedTransientLikelihoodFD.

Suggested reviewers: narolaharsh

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly matches the main change: restructuring the valid waveform mask in the Heterodyne likelihood.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch rm-obsolete-mask

---

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.}

[coderabbitai]

Actionable comments posted: 1

🤖 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 697-699: In HeterodynedTransientLikelihoodFD.__init__, reorder the
assignments so self.freq_grid_low and self.freq_grid_high are set from freq_grid
before calculating self.bin_widths; then derive bin_widths from those
initialized attributes.

🪄 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: 9b56adbe-eb3a-4fa0-a369-4536edb125be

📥 Commits

Reviewing files that changed from the base of the PR and between 4d7f611 and 4cf526c.

📒 Files selected for processing (1)

• src/jimgw/core/single_event/likelihood.py

[coderabbitai]

Actionable comments posted: 3

🧹 Nitpick comments (1)

src/jimgw/core/single_event/likelihood.py (1)

556-561: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Remove the spaces from these jaxtyping shape strings — " n_bins + 1" / " n_valid+1" are parsed as separate axes, so these annotations need compact symbolic expressions. Ruff already ignores F722, so no config change is needed.

🤖 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/core/single_event/likelihood.py` around lines 556 - 561, The
jaxtyping shape annotations in the affected likelihood definitions use spaces
around symbolic expressions, causing them to parse as separate axes. Update the
relevant annotations, including `freq_grid_low`, `freq_grid_high`, `bin_widths`,
`waveform_low_ref`, `waveform_high_ref`, and `summary_data`, to use compact
shape strings such as `"n_bins+1"` and `"n_valid+1"` without spaces; leave Ruff
configuration unchanged.

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 `@src/jimgw/core/single_event/likelihood.py`:
- Around line 849-875: Remove the self references from _compute_coefficients,
using the supplied f_bins to derive the bin edges and centers needed for the
mask and frequency-shift calculations. Replace self.freq_grid_low and
self.freq_grid_high with values computed from f_bins while preserving the
final-bin inclusive equality handling.
- Line 817: In the relevant likelihood class initialization logic, replace the
assignment to self.nbins with self.n_bins so the valid-bin count is stored in
the declared attribute expected by downstream callers and tests.

In `@tests/unit/core/single_event/test_likelihood.py`:
- Around line 1174-1234: Restore and activate
test_low_frequency_reference_cutoff_does_not_reindex_summary_data, updating its
assertions and setup to match the final HeterodynedTransientLikelihoodFD
contract; retain coverage for masked bin-edge reindexing, reduced n_bins,
frequency bounds, and per-detector coefficient dimensions, and remove the
commented-out test block.

---

Nitpick comments:
In `@src/jimgw/core/single_event/likelihood.py`:
- Around line 556-561: The jaxtyping shape annotations in the affected
likelihood definitions use spaces around symbolic expressions, causing them to
parse as separate axes. Update the relevant annotations, including
`freq_grid_low`, `freq_grid_high`, `bin_widths`, `waveform_low_ref`,
`waveform_high_ref`, and `summary_data`, to use compact shape strings such as
`"n_bins+1"` and `"n_valid+1"` without spaces; leave Ruff configuration
unchanged.

🪄 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: ee9ecd0d-aa26-4e37-800c-91668c5e69b5

📥 Commits

Reviewing files that changed from the base of the PR and between 4cf526c and 2d34459.

📒 Files selected for processing (2)

• src/jimgw/core/single_event/likelihood.py
• tests/unit/core/single_event/test_likelihood.py