Patch for unstable part in branches in postprocessing#128
Conversation
📝 WalkthroughWalkthroughAdded a new helper function to detect non-monotone (unstable) tidal deformability sequences. The plotting functions now segment curves into monotone-decreasing branches before rendering, with logging to track samples exhibiting instability. Test coverage includes the new helper and smoke tests for both plotting functions. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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 |
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@jesterTOV/inference/postprocessing/postprocessing.py`:
- Around line 354-361: Before returning slice ranges, validate that the paired
arrays have equal lengths: check that len(lambdas) == len(masses) at start of
the routine (the code that builds segments using lambdas, masses, segments,
start) and if they differ either raise a ValueError with a clear message or
truncate to the smaller length and use that for the final segment boundary;
ensure the final segments.append uses the validated length (e.g.,
len(lambdas_validated) or min_len) rather than unvalidated len(masses).
- Around line 335-353: Convert the docstring for the function that splits a
mass-Lambda curve (the docstring that documents parameters masses and lambdas in
postprocessing.py) into a raw string (prefix with r""") and replace any
LaTeX/math fragments with Sphinx roles: use :math:`...` for inline expressions
(e.g. variable names like m, Lambda) and use a ``.. math::`` block for any
displayed equations; ensure backslashes are preserved by the raw string and
update examples/phrasing to use :math: roles consistently throughout the
description.
In `@tests/test_inference/test_postprocessing.py`:
- Around line 529-627: Add explicit return type annotations "-> None" to each
new test function definition so they comply with the project's type-hinting
rule: update the defs for test_monotone_decreasing_returns_one_segment,
test_single_lambda_increase_splits_into_two_branches,
test_multiple_lambda_increases_many_branches,
test_constant_lambda_returns_one_segment, test_all_segments_cover_full_range,
test_plot_runs_without_error_when_unstable, and
test_plot_runs_without_error_when_smooth to include "-> None" after the
parameter list.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: e0515b75-e4fb-476b-a530-555ab85b6ed9
📒 Files selected for processing (2)
jesterTOV/inference/postprocessing/postprocessing.pytests/test_inference/test_postprocessing.py
| """Split a mass-Lambda curve into monotone-decreasing segments. | ||
|
|
||
| A branch break is detected wherever Lambda increases as mass increases, | ||
| which signals that the stored curve contains two interleaved stable branches | ||
| (the "third family" / twin-star scenario arising from a fold-back in M(pc)). | ||
|
|
||
| Parameters | ||
| ---------- | ||
| masses : np.ndarray | ||
| Mass array (monotone increasing, uniform grid). | ||
| lambdas : np.ndarray | ||
| Tidal deformability array on the same grid. | ||
|
|
||
| Returns | ||
| ------- | ||
| list of (start, end) index pairs | ||
| Each pair defines a half-open slice ``masses[start:end]`` that is | ||
| monotone decreasing in Lambda. Always contains at least one segment. | ||
| """ |
There was a problem hiding this comment.
Format the math notation in the new docstring.
The helper docstring introduces math notation but does not use Sphinx math roles; if you spell this as LaTeX, make the docstring raw as well.
Proposed fix
-def _split_into_monotone_branches(
- masses: np.ndarray, lambdas: np.ndarray
-) -> list[tuple[int, int]]:
- """Split a mass-Lambda curve into monotone-decreasing segments.
+def _split_into_monotone_branches(
+ masses: np.ndarray, lambdas: np.ndarray
+) -> list[tuple[int, int]]:
+ r"""Split a :math:`M-\Lambda` curve into monotone-decreasing segments.
- A branch break is detected wherever Lambda increases as mass increases,
+ A branch break is detected wherever :math:`\Lambda` increases as mass increases,
which signals that the stored curve contains two interleaved stable branches
- (the "third family" / twin-star scenario arising from a fold-back in M(pc)).
+ (the "third family" / twin-star scenario arising from a fold-back in
+ :math:`M(p_c)`).As per coding guidelines, "All mathematical expressions in docstrings must use Sphinx/reStructuredText formatting with :math: role for inline math and .. math:: directive for display equations" and "Use raw strings (r\"\"\") for docstrings containing LaTeX to avoid Python escape sequence warnings".
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@jesterTOV/inference/postprocessing/postprocessing.py` around lines 335 - 353,
Convert the docstring for the function that splits a mass-Lambda curve (the
docstring that documents parameters masses and lambdas in postprocessing.py)
into a raw string (prefix with r""") and replace any LaTeX/math fragments with
Sphinx roles: use :math:`...` for inline expressions (e.g. variable names like
m, Lambda) and use a ``.. math::`` block for any displayed equations; ensure
backslashes are preserved by the raw string and update examples/phrasing to use
:math: roles consistently throughout the description.
| segments: list[tuple[int, int]] = [] | ||
| start = 0 | ||
| for j in range(1, len(lambdas)): | ||
| if lambdas[j] > lambdas[j - 1]: | ||
| segments.append((start, j)) | ||
| start = j | ||
| segments.append((start, len(masses))) | ||
| return segments |
There was a problem hiding this comment.
Validate the paired array lengths before returning slice ranges.
The function detects breaks over lambdas but ends the last slice at len(masses). If malformed input has mismatched lengths, callers can plot ranges that were not checked for monotonicity.
Proposed fix
segments: list[tuple[int, int]] = []
+ if len(masses) != len(lambdas):
+ raise ValueError(
+ "masses and lambdas must have the same length to split monotone branches"
+ )
+
start = 0🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@jesterTOV/inference/postprocessing/postprocessing.py` around lines 354 - 361,
Before returning slice ranges, validate that the paired arrays have equal
lengths: check that len(lambdas) == len(masses) at start of the routine (the
code that builds segments using lambdas, masses, segments, start) and if they
differ either raise a ValueError with a clear message or truncate to the smaller
length and use that for the final segment boundary; ensure the final
segments.append uses the validated length (e.g., len(lambdas_validated) or
min_len) rather than unvalidated len(masses).
| def test_monotone_decreasing_returns_one_segment(self): | ||
| """Smooth Lambda(M) with no increases → single segment covering all data.""" | ||
| masses = np.array([0.5, 0.8, 1.1, 1.4, 1.7, 2.0]) | ||
| lambdas = np.array([2000.0, 800.0, 300.0, 120.0, 50.0, 20.0]) | ||
| branches = _split_into_monotone_branches(masses, lambdas) | ||
| assert len(branches) == 1 | ||
| assert branches[0] == (0, 6) | ||
|
|
||
| def test_single_lambda_increase_splits_into_two_branches(self): | ||
| """One Lambda increase mid-curve → two branches.""" | ||
| masses = np.array([0.5, 0.8, 1.1, 1.2, 1.5, 1.8]) | ||
| lambdas = np.array([2000.0, 800.0, 300.0, 1500.0, 400.0, 100.0]) | ||
| branches = _split_into_monotone_branches(masses, lambdas) | ||
| assert len(branches) == 2 | ||
| assert branches[0] == (0, 3) | ||
| assert branches[1] == (3, 6) | ||
|
|
||
| def test_multiple_lambda_increases_many_branches(self): | ||
| """Multiple Lambda increases → multiple branches.""" | ||
| masses = np.linspace(0.5, 2.0, 10) | ||
| lambdas = np.array( | ||
| [1000.0, 500.0, 800.0, 200.0, 600.0, 100.0, 300.0, 50.0, 80.0, 20.0] | ||
| ) | ||
| branches = _split_into_monotone_branches(masses, lambdas) | ||
| assert len(branches) > 2 | ||
| assert branches[0][0] == 0 | ||
| assert branches[-1][1] == len(masses) | ||
|
|
||
| def test_constant_lambda_returns_one_segment(self): | ||
| """Flat Lambda → no increases → single segment.""" | ||
| masses = np.linspace(0.5, 2.0, 5) | ||
| lambdas = np.ones(5) * 500.0 | ||
| branches = _split_into_monotone_branches(masses, lambdas) | ||
| assert len(branches) == 1 | ||
|
|
||
| def test_all_segments_cover_full_range(self): | ||
| """Branch segments are contiguous and cover the full array without gaps.""" | ||
| masses = np.linspace(0.5, 2.0, 20) | ||
| rng = np.random.default_rng(42) | ||
| lambdas = np.cumsum(rng.uniform(-200, 100, 20)) + 2000.0 | ||
| branches = _split_into_monotone_branches(masses, lambdas) | ||
| assert branches[0][0] == 0 | ||
| assert branches[-1][1] == len(masses) | ||
| for k in range(len(branches) - 1): | ||
| assert branches[k][1] == branches[k + 1][0] | ||
|
|
||
| def test_plot_runs_without_error_when_unstable(self, tmp_path): | ||
| """make_mass_radius_plot runs cleanly when multi-branch samples are present.""" | ||
| n_eos = 50 | ||
| masses_jagged = np.linspace(0.75, 2.5, n_eos) | ||
| lambdas_jagged = np.concatenate( | ||
| [ | ||
| np.linspace(2000, 400, 25), | ||
| np.linspace(1500, 50, 25), | ||
| ] | ||
| ) | ||
| radii_jagged = np.linspace(12.0, 10.0, n_eos) | ||
|
|
||
| masses_smooth = np.linspace(0.75, 2.5, n_eos) | ||
| lambdas_smooth = np.linspace(2000, 10, n_eos) | ||
| radii_smooth = np.linspace(12.0, 10.0, n_eos) | ||
|
|
||
| data = { | ||
| "masses": np.array([masses_jagged, masses_smooth]), | ||
| "radii": np.array([radii_jagged, radii_smooth]), | ||
| "lambdas": np.array([lambdas_jagged, lambdas_smooth]), | ||
| "densities": np.ones((2, n_eos)), | ||
| "pressures": np.ones((2, n_eos)), | ||
| "cs2": np.ones((2, n_eos)) * 0.3, | ||
| "log_prob": np.array([-10.0, -10.0]), | ||
| "prior_params": {}, | ||
| } | ||
|
|
||
| make_mass_radius_plot(data, prior_data=None, outdir=str(tmp_path)) | ||
| plt.close("all") | ||
| assert (tmp_path / "mass_radius_plot.pdf").exists() | ||
|
|
||
| def test_plot_runs_without_error_when_smooth(self, tmp_path): | ||
| """make_mass_radius_plot runs cleanly when all samples are smooth.""" | ||
| n_eos = 50 | ||
| n_samples = 3 | ||
| masses = np.tile(np.linspace(0.75, 2.5, n_eos), (n_samples, 1)) | ||
| lambdas = np.tile(np.linspace(2000, 10, n_eos), (n_samples, 1)) | ||
| radii = np.tile(np.linspace(12.0, 10.0, n_eos), (n_samples, 1)) | ||
|
|
||
| data = { | ||
| "masses": masses, | ||
| "radii": radii, | ||
| "lambdas": lambdas, | ||
| "densities": np.ones((n_samples, n_eos)), | ||
| "pressures": np.ones((n_samples, n_eos)), | ||
| "cs2": np.ones((n_samples, n_eos)) * 0.3, | ||
| "log_prob": np.full(n_samples, -10.0), | ||
| "prior_params": {}, | ||
| } | ||
|
|
||
| make_mass_radius_plot(data, prior_data=None, outdir=str(tmp_path)) | ||
| plt.close("all") | ||
| assert (tmp_path / "mass_radius_plot.pdf").exists() |
There was a problem hiding this comment.
Add return type annotations to the new tests.
The new test methods should declare -> None to satisfy the project’s type-hinting rule.
Proposed fix
- def test_monotone_decreasing_returns_one_segment(self):
+ def test_monotone_decreasing_returns_one_segment(self) -> None:
@@
- def test_single_lambda_increase_splits_into_two_branches(self):
+ def test_single_lambda_increase_splits_into_two_branches(self) -> None:
@@
- def test_multiple_lambda_increases_many_branches(self):
+ def test_multiple_lambda_increases_many_branches(self) -> None:
@@
- def test_constant_lambda_returns_one_segment(self):
+ def test_constant_lambda_returns_one_segment(self) -> None:
@@
- def test_all_segments_cover_full_range(self):
+ def test_all_segments_cover_full_range(self) -> None:
@@
- def test_plot_runs_without_error_when_unstable(self, tmp_path):
+ def test_plot_runs_without_error_when_unstable(self, tmp_path) -> None:
@@
- def test_plot_runs_without_error_when_smooth(self, tmp_path):
+ def test_plot_runs_without_error_when_smooth(self, tmp_path) -> None:As per coding guidelines, "All new code MUST include comprehensive type hints with standard library types (Python 3.10+ syntax), jaxtyping for JAX arrays, Pydantic for configs, and type aliases for complex types".
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/test_inference/test_postprocessing.py` around lines 529 - 627, Add
explicit return type annotations "-> None" to each new test function definition
so they comply with the project's type-hinting rule: update the defs for
test_monotone_decreasing_returns_one_segment,
test_single_lambda_increase_splits_into_two_branches,
test_multiple_lambda_increases_many_branches,
test_constant_lambda_returns_one_segment, test_all_segments_cover_full_range,
test_plot_runs_without_error_when_unstable, and
test_plot_runs_without_error_when_smooth to include "-> None" after the
parameter list.
Some EOSs might give rise to unstable parts and multiple branches, so that a better treatment to interpolate through the TOV solution is required. A proposal is already given in #106 to improve upon this, but it would take quite some time to fully test it and make sure this implementation is fully correct as also mentioned there in the src code. This PR gives a quick patch by essentially discarding the unstable part in the plotting during postprocessing. Note that this means that the EOS/NS itself did have the unstable part during the likelihood evaluation and is therefore not fully correct. This has negligible impact on the posterior though: for chiEFT, 1.6% of samples are flagged to have been patched up in this way, and it is expected this will reduce to ~0% for cases with eg the radio timing likelihood, which has overall much more well-behaved solutions. Therefore, this is an acceptable fix for the IGWN review timeline, and gives us more time and freedom to explore the best general solution beyond IGWN review.
Summary by CodeRabbit
New Features
Tests