Scalar-Tensor theory update draft PR#106
Open
diodaberjalan wants to merge 22 commits into
Open
Conversation
Unused for now
Stt inference
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds scalar-tensor TOV solver support: new Pydantic config, transform dispatch, JIT-compiled scalar-tensor solver, solver-specific "extra" fields propagated through family construction, new interpolation utilities, example inference workflow, tests, and minor tooling/docs changes. Changes
Sequence Diagram(s)sequenceDiagram
participant CLI as Inference CLI
participant Transform as JesterTransform
participant Solver as ScalarTensorTOVSolver
participant Utils as utils (interpolate)
participant Data as FamilyData
rect rgba(200,220,255,0.5)
CLI->>Transform: load config (ScalarTensorTOVConfig)
Transform->>Solver: _create_tov_solver(config)
end
rect rgba(200,255,200,0.5)
Transform->>Solver: solve(pc, tov_params...)
Solver->>Solver: _compiled_tov_solve(...) (JIT, lax.while_loop)
Solver-->>Transform: TOVSolution(M,R,k2, extra)
end
rect rgba(255,230,200,0.5)
Transform->>Data: construct_family(..., extra=solutions.extra)
Data->>Utils: limit_by_MTOV_and_interpolate(..., ndat)
Utils-->>Data: interpolated pc, m, r, l (and extra arrays)
Data-->>CLI: FamilyData (includes extra)
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 10
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
docs/examples/index.rst (1)
16-24:⚠️ Potential issue | 🔴 CriticalRemove duplicate "Scalar-tensor theory TOV solver" sections.
There are two identical sections (lines 16-19 and 21-24) with the same title and description, but pointing to different notebooks (
eos_tov/eos_STtovandeos_tov/EoS_beyondGR). This creates documentation confusion.Please either:
- Remove one section if the notebooks cover the same content
- Give them distinct titles and descriptions if they demonstrate different aspects of scalar-tensor theory
- Replace the old section with the new one if
EoS_beyondGRsupersedeseos_STtov📝 Proposed fix assuming EoS_beyondGR replaces eos_STtov
**Scalar-tensor theory TOV solver** Exploration of modified gravity theories using JESTER's scalar-tensor TOV solver to study deviations from general relativity. - :doc:`eos_tov/eos_STtov` - -**Scalar-tensor theory TOV solver** - Exploration of modified gravity theories using JESTER's scalar-tensor TOV solver to study deviations from general relativity. - :doc:`eos_tov/EoS_beyondGR`Then update the toctree accordingly:
eos_tov/eos_tov eos_tov/automatic_differentiation - eos_tov/eos_STtov eos_tov/EoS_beyondGR🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@docs/examples/index.rst` around lines 16 - 24, The two identical "Scalar-tensor theory TOV solver" entries in docs/examples/index.rst point to eos_tov/eos_STtov and eos_tov/EoS_beyondGR and must be de-duplicated: remove the older eos_tov/eos_STtov entry and keep (or replace it with) eos_tov/EoS_beyondGR, updating the section title/description if needed to reflect any behavioral differences and ensure the toctree references only the chosen notebook (update the :doc:`...` link to eos_tov/EoS_beyondGR and adjust the surrounding title/description accordingly).docs/examples/eos_tov/eos_STtov.ipynb (2)
203-221:⚠️ Potential issue | 🟡 MinorFilter non-positive
Lambdasbefore switching the axis to log scale.The saved output already shows
plt.yscale("log")warning because someLambdasare zero/negative after the mass cut. That leaves the example with a noisy warning and an invalid plot.♻️ Proposed fix
-# Limit masses to be above certain mass to make plot prettier -m_min = 0.5 -mask = masses > m_min +# Limit masses to be above certain mass and keep only positive Lambdas for log plots +m_min = 0.5 +mask = (masses > m_min) & (Lambdas > 0) masses = masses[mask] radii = radii[mask] Lambdas = Lambdas[mask]🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@docs/examples/eos_tov/eos_STtov.ipynb` around lines 203 - 221, The Lambda(R) subplot uses plt.yscale("log") but Lambdas contains zero/negative values after the m_min mask, causing warnings and invalid plotting; before plotting in the Lambda plot (after computing mask and slicing masses, radii, Lambdas) further filter for Lambdas > 0 (e.g., compute a positive_lambda_mask = Lambdas > 0 and use it to index Lambdas and the corresponding masses) so only positive Lambda values are passed to plt.plot(masses, Lambdas) prior to plt.yscale("log").
64-67:⚠️ Potential issue | 🟠 MajorReplace the nonexistent
jesterTOV.eos.familiesimport with the correct scalar-tensor solver class.The notebook imports
construct_family_STandconstruct_family_ST_solfrom a modulejesterTOV.eos.familiesthat does not exist in the repository, causing the committed output to showModuleNotFoundError: No module named 'jesterTOV.eos.families'. Scalar-tensor support is provided by theScalarTensorTOVSolverclass injesterTOV.tov.scalar_tensor. Use:from jesterTOV.tov.scalar_tensor import ScalarTensorTOVSolver solver = ScalarTensorTOVSolver() family_data = solver.construct_family(eos_data, ndat=50, min_nsat=2.0, tov_params=params)instead of calling nonexistent standalone functions. Re-execute the notebook after fixing the imports to ensure the example runs cleanly.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@docs/examples/eos_tov/eos_STtov.ipynb` around lines 64 - 67, The notebook imports nonexistent construct_family_ST functions from jesterTOV.eos.families; replace that import with the ScalarTensorTOVSolver class from jesterTOV.tov.scalar_tensor, instantiate ScalarTensorTOVSolver (e.g., solver = ScalarTensorTOVSolver()) and call its construct_family method (solver.construct_family(...)) with the existing eos_data, ndat, min_nsat and tov_params instead of the removed standalone functions; then re-run the notebook cells to ensure the example executes without ModuleNotFoundError.
🧹 Nitpick comments (5)
examples/inference/st_smc_random_walk/chiEFT/config.yaml (1)
20-23: Prefer the built-in ChiEFT data defaults here.
ChiEFTLikelihood.__init__already falls back to the packagedlow.datandhigh.datfiles when these keys are omitted. Keeping repo-relative paths here makes the example more brittle and unnecessarily couples it to a source checkout.Suggested change
- type: chieft enabled: true - low_filename: ../../../../jesterTOV/inference/data/chiEFT/2402.04172/low.dat - high_filename: ../../../../jesterTOV/inference/data/chiEFT/2402.04172/high.dat nb_n: 100🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/inference/st_smc_random_walk/chiEFT/config.yaml` around lines 20 - 23, Remove the explicit repo-relative low_filename and high_filename entries so the example uses the built-in packaged defaults provided by ChiEFTLikelihood.__init__; keep the block with type: chieft and enabled: true, and delete the low_filename and high_filename keys (they are optional and will default to the bundled low.dat/high.dat inside ChiEFTLikelihood).tests/test_integration.py (1)
114-114: Derivenbreakfromnsatinstead of hard-coding0.32.The comment says “2*nsat”, but the test encodes that relationship twice. If
nsatchanges later, this silently becomes wrong and the test starts exercising a different regime than intended.♻️ Proposed fix
- "nbreak": 0.32, # Break density in fm^-3 (2*nsat) + "nbreak": 2.0 * model.nsat, # Break density in fm^-3 (2*nsat)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/test_integration.py` at line 114, The test currently hard-codes "nbreak = 0.32" while the comment states it should be 2*nsat; change the test to compute nbreak from the existing nsat variable (set nbreak = 2 * nsat) instead of using the magic number so the relationship stays correct if nsat changes; update any inline comment to reflect that nbreak is derived from nsat and run the test to ensure no other places expect the literal 0.32.jesterTOV/utils.py (1)
282-326: Use Sphinx math markup consistently in this new docstring.This helper mixes
.. math::blocks with inline$...$expressions, so part of the rendered documentation will show raw TeX instead of math. Please convert the inline expressions to:math:as well.As per coding guidelines "All mathematical expressions in docstrings must use Sphinx/reStructuredText formatting: use
:math:role for inline math,.. math::directive for display equations, always use raw strings (r""") for docstrings containing LaTeX".🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@jesterTOV/utils.py` around lines 282 - 326, The docstring for the stable-branch interpolation block mixes inline TeX ($...$) with ``.. math::`` and must be converted to consistent Sphinx math roles; update the raw docstring (starts with r""" shown) to replace all inline dollar-delimited expressions like `$M_\odot$`, `$km$`, `$S = \sum_i s_i$`, and `$s_i = \log p_{c,\mathrm{end}}^{(i)} - \log p_{c,\mathrm{start}}^{(i)}$` with the :math: role (e.g. :math:`M_\odot`) and similarly convert any other inline dollar math to :math:`...`, leaving display equations in ``.. math::`` blocks unchanged; keep references to function variables (pc, m, r, l, ndat) and preserve the raw-string r""" marker so JAX doc rendering remains correct.jesterTOV/tov/base.py (1)
224-234: Consider batching extra field interpolation to avoid redundant computation.Each call to
limit_by_MTOV_and_interpolaterecomputes stable segment detection from the samepcsandmasses_solararrays. The TODO acknowledges this. A more efficient approach would extract the interpolation grid computation once and apply it to all fields.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@jesterTOV/tov/base.py` around lines 224 - 234, The current mapping calls utils.limit_by_MTOV_and_interpolate for each extra field, recomputing the same stability/interpolation grid from pcs and masses_solar; instead, call utils.limit_by_MTOV_and_interpolate once with pcs, masses_solar, radii_km and ndat to capture the shared outputs (e.g., the stable mask, interpolation grid, and the fourth returned array), then reuse those precomputed components when building extra_processed by applying a lightweight interpolation-only function over each value in extra (via jax.tree_util.tree_map or a loop) so limit_by_MTOV_and_interpolate is not invoked per-field; reference symbols: extra, extra_processed, utils.limit_by_MTOV_and_interpolate, pcs, masses_solar, radii_km, ndat.jesterTOV/tov/scalar_tensor.py (1)
298-301: Add type hints to_compiled_tov_solvefunction.Per coding guidelines, all new code must include comprehensive type hints. This function lacks type annotations for parameters and return type.
Proposed signature with type hints
`@functools.partial`(jax.jit, static_argnames=["max_iterations"]) def _compiled_tov_solve( - pc, beta_ST, phi_inf_target, phi0, ps, hs, es, dloge_dlogps, max_iterations=100 -): + pc: float, + beta_ST: float, + phi_inf_target: float, + phi0: float, + ps: Float[Array, "n_points"], + hs: Float[Array, "n_points"], + es: Float[Array, "n_points"], + dloge_dlogps: Float[Array, "n_points"], + max_iterations: int = 100, +) -> tuple[float, float, float, float, float, float, float]:🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@jesterTOV/tov/scalar_tensor.py` around lines 298 - 301, Add full type annotations to the _compiled_tov_solve signature: annotate array-like inputs (pc, phi_inf_target, phi0, ps, hs, es, dloge_dlogps) as jax.Array (or jnp.ndarray), beta_ST and phi0 scalars as float if appropriate, max_iterations as int, and declare an explicit return type (e.g. Tuple[jax.Array, ...] matching the function's returned values). Also add the needed typing imports (e.g. from typing import Tuple) and ensure the `@functools.partial`(jax.jit, static_argnames=["max_iterations"]) decorator still applies to the annotated function.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@examples/inference/st_smc_random_walk/chiEFT/submit.sh`:
- Around line 19-28: The script currently sources an absolute venv and runs
run_jester_inference assuming the CWD; change it to resolve the script directory
and cd into it, make the virtualenv path configurable via an environment
variable (e.g., VENV with a sensible default) and source "$VENV/bin/activate",
and invoke run_jester_inference with the config and data paths prefixed by the
resolved script directory (so references to config.yaml, prior.prior and ChiEFT
data paths remain valid regardless of the initial working directory). Ensure you
update the activate path and all file arguments to use the script-dir-relative
paths and preserve the nvidia-smi and echo steps.
- Line 10: The SBATCH output path uses ./outdir which may not exist; before
submitting, ensure the directory exists by creating it (e.g., run a
directory-creation step prior to SBATCH in submit.sh) or change the --output
value to a guaranteed existing path; update the submit.sh script so the
directory for the --output (referenced by the SBATCH
--output="./outdir/log_smc_rw.out" directive) is created with mkdir -p or
replaced with an existing folder.
In `@jesterTOV/inference/transforms/transform.py`:
- Around line 287-289: The ScalarTensorTOVConfig fields (beta_ST, phi_inf_tgt,
phi_c) are not propagated to the solver because the branch returning
ScalarTensorTOVSolver() discards config; update the instantiation in the branch
handling ScalarTensorTOVConfig to pass the config values into the solver (or set
them on the solver instance) so that construct_eos_and_solve_tov and ultimately
ScalarTensorTOVSolver.solve() receive them; specifically, when config is an
instance of ScalarTensorTOVConfig, extract beta_ST/phi_inf_tgt/phi_c and supply
them to ScalarTensorTOVSolver (or ensure fetch_params reads from the config) so
the solver uses the configured scalar-tensor parameters.
In `@jesterTOV/tov/scalar_tensor.py`:
- Around line 630-632: The default beta_ST=0.0 in tov_params.get("beta_ST", 0.0)
leads to downstream division-by-zero in the Jordan-frame conversion code; remove
the unsafe default so beta_ST must be provided (e.g., replace
tov_params.get("beta_ST", 0.0) with a direct lookup like tov_params["beta_ST"]
or tov_params.get("beta_ST") without a fallback), leaving phi_inf_target and
phi0 handling unchanged; ensure this relies on existing
get_required_parameters()/fetch_params()/solve() validation to enforce presence
of beta_ST.
- Around line 578-602: Lambda_S_J, Lambda_ST1_J and Lambda_ST2_J perform
divisions by beta_ST and phi_inf_target and can produce inf/nan when beta_ST==0
or phi_inf_target==0; add a guard at the start of the conversion block that
checks for small |beta_ST| or |phi_inf_target| (e.g. abs(x) < eps) and either:
(a) raise a clear ValueError suggesting GR limit users should use GRTOVSolver,
or (b) use a safe fallback (set the problematic Lambda_*_J expressions to their
GR-consistent limits or zero) depending on intended behavior; implement the
check around the Lambda_*_J computations so Lambda_S_J, Lambda_ST1_J,
Lambda_ST2_J are never computed with a zero divisor and include beta_ST and
phi_inf_target in the error message if raising.
In `@jesterTOV/utils.py`:
- Around line 354-364: The code hard-codes four segments by using jnp.where(...,
size=4, ...) and unpacking into s1_start..s4_end which drops any extra stable
branches; change to compute starts = jnp.nonzero(diff == 1)[0] and ends =
jnp.nonzero(diff == -1)[0] (or jnp.where without a fixed size) to get
variable-length index arrays, then replace the fixed unpacking of
s1_start..s4_end with code that iterates over the computed starts/ends (or
returns them as arrays) so all detected segments are preserved (refer to the
variables diff, starts, ends, and the s1_start..s4_end handling in this
function).
- Around line 469-480: The centered differences (central_diff_numer,
central_diff_denom, interior_m1_m) are computed across the entire xp array and
can leak slope info across stitched branch boundaries; update the slope
computation to operate per contiguous stable segment: detect segment boundaries
in xp (e.g., where consecutive index or xp-step indicates a gap), then for each
segment compute forward_m0 / backward_m_end at its ends and central differences
only for its interior points, and finally concatenate per-segment dydx pieces to
form the full dydx array so no centered difference uses xp values from a
different segment.
In `@tests/test_inference/test_anisotropy_tov_inference.py`:
- Around line 279-290: The test currently interpolates R1.4 over the full mass
curves (masses_EOS_gr/radii_EOS_gr and masses_EOS_mg/radii_EOS_mg) which can be
non-monotonic past the maximum-mass point; first find the index of the maximum
mass (e.g., using jnp.argmax on masses_EOS_gr and masses_EOS_mg), slice each
mass and radius array to the stable branch up through that index (inclusive) so
x is monotonic, then perform the jnp.interp on those sliced arrays to compute
R1_4_gr and R1_4_mg before comparing. Ensure you reference masses_EOS_gr,
radii_EOS_gr, masses_EOS_mg, radii_EOS_mg and update how R1_4_gr / R1_4_mg are
computed accordingly.
In `@tests/test_inference/test_config.py`:
- Around line 160-175: The test test_valid_scalar_tensor_tov_config currently
only asserts beta_ST; update it to also assert that the validated config is a
schema.ScalarTensorTOVConfig and that the scalar-tensor-specific fields
phi_inf_tgt and phi_c are correctly parsed (e.g., exact values and/or types)
after using TypeAdapter(schema.TOVConfig). In short, inside
test_valid_scalar_tensor_tov_config add assertions that config.phi_inf_tgt ==
1e-3 and config.phi_c == 1.0 (and optional type checks) so regressions in
parsing those keys are caught.
---
Outside diff comments:
In `@docs/examples/eos_tov/eos_STtov.ipynb`:
- Around line 203-221: The Lambda(R) subplot uses plt.yscale("log") but Lambdas
contains zero/negative values after the m_min mask, causing warnings and invalid
plotting; before plotting in the Lambda plot (after computing mask and slicing
masses, radii, Lambdas) further filter for Lambdas > 0 (e.g., compute a
positive_lambda_mask = Lambdas > 0 and use it to index Lambdas and the
corresponding masses) so only positive Lambda values are passed to
plt.plot(masses, Lambdas) prior to plt.yscale("log").
- Around line 64-67: The notebook imports nonexistent construct_family_ST
functions from jesterTOV.eos.families; replace that import with the
ScalarTensorTOVSolver class from jesterTOV.tov.scalar_tensor, instantiate
ScalarTensorTOVSolver (e.g., solver = ScalarTensorTOVSolver()) and call its
construct_family method (solver.construct_family(...)) with the existing
eos_data, ndat, min_nsat and tov_params instead of the removed standalone
functions; then re-run the notebook cells to ensure the example executes without
ModuleNotFoundError.
In `@docs/examples/index.rst`:
- Around line 16-24: The two identical "Scalar-tensor theory TOV solver" entries
in docs/examples/index.rst point to eos_tov/eos_STtov and eos_tov/EoS_beyondGR
and must be de-duplicated: remove the older eos_tov/eos_STtov entry and keep (or
replace it with) eos_tov/EoS_beyondGR, updating the section title/description if
needed to reflect any behavioral differences and ensure the toctree references
only the chosen notebook (update the :doc:`...` link to eos_tov/EoS_beyondGR and
adjust the surrounding title/description accordingly).
---
Nitpick comments:
In `@examples/inference/st_smc_random_walk/chiEFT/config.yaml`:
- Around line 20-23: Remove the explicit repo-relative low_filename and
high_filename entries so the example uses the built-in packaged defaults
provided by ChiEFTLikelihood.__init__; keep the block with type: chieft and
enabled: true, and delete the low_filename and high_filename keys (they are
optional and will default to the bundled low.dat/high.dat inside
ChiEFTLikelihood).
In `@jesterTOV/tov/base.py`:
- Around line 224-234: The current mapping calls
utils.limit_by_MTOV_and_interpolate for each extra field, recomputing the same
stability/interpolation grid from pcs and masses_solar; instead, call
utils.limit_by_MTOV_and_interpolate once with pcs, masses_solar, radii_km and
ndat to capture the shared outputs (e.g., the stable mask, interpolation grid,
and the fourth returned array), then reuse those precomputed components when
building extra_processed by applying a lightweight interpolation-only function
over each value in extra (via jax.tree_util.tree_map or a loop) so
limit_by_MTOV_and_interpolate is not invoked per-field; reference symbols:
extra, extra_processed, utils.limit_by_MTOV_and_interpolate, pcs, masses_solar,
radii_km, ndat.
In `@jesterTOV/tov/scalar_tensor.py`:
- Around line 298-301: Add full type annotations to the _compiled_tov_solve
signature: annotate array-like inputs (pc, phi_inf_target, phi0, ps, hs, es,
dloge_dlogps) as jax.Array (or jnp.ndarray), beta_ST and phi0 scalars as float
if appropriate, max_iterations as int, and declare an explicit return type (e.g.
Tuple[jax.Array, ...] matching the function's returned values). Also add the
needed typing imports (e.g. from typing import Tuple) and ensure the
`@functools.partial`(jax.jit, static_argnames=["max_iterations"]) decorator still
applies to the annotated function.
In `@jesterTOV/utils.py`:
- Around line 282-326: The docstring for the stable-branch interpolation block
mixes inline TeX ($...$) with ``.. math::`` and must be converted to consistent
Sphinx math roles; update the raw docstring (starts with r""" shown) to replace
all inline dollar-delimited expressions like `$M_\odot$`, `$km$`, `$S = \sum_i
s_i$`, and `$s_i = \log p_{c,\mathrm{end}}^{(i)} - \log
p_{c,\mathrm{start}}^{(i)}$` with the :math: role (e.g. :math:`M_\odot`) and
similarly convert any other inline dollar math to :math:`...`, leaving display
equations in ``.. math::`` blocks unchanged; keep references to function
variables (pc, m, r, l, ndat) and preserve the raw-string r""" marker so JAX doc
rendering remains correct.
In `@tests/test_integration.py`:
- Line 114: The test currently hard-codes "nbreak = 0.32" while the comment
states it should be 2*nsat; change the test to compute nbreak from the existing
nsat variable (set nbreak = 2 * nsat) instead of using the magic number so the
relationship stays correct if nsat changes; update any inline comment to reflect
that nbreak is derived from nsat and run the test to ensure no other places
expect the literal 0.32.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 092a53d4-31c2-465f-a0e9-c8a1e036198a
📒 Files selected for processing (24)
.gitignoredocs/api/_autosummary/jesterTOV.tov.data_classes.FamilyData.rstdocs/api/_autosummary/jesterTOV.tov.data_classes.TOVSolution.rstdocs/examples/eos_tov/EoS_beyondGR.ipynbdocs/examples/eos_tov/automatic_differentiation.ipynbdocs/examples/eos_tov/eos_STtov.ipynbdocs/examples/eos_tov/eos_tov.ipynbdocs/examples/index.rstexamples/inference/st_smc_random_walk/chiEFT/config.yamlexamples/inference/st_smc_random_walk/chiEFT/prior.priorexamples/inference/st_smc_random_walk/chiEFT/submit.shjesterTOV/inference/config/schema.pyjesterTOV/inference/config/schemas/tov.pyjesterTOV/inference/transforms/transform.pyjesterTOV/tov/base.pyjesterTOV/tov/data_classes.pyjesterTOV/tov/scalar_tensor.pyjesterTOV/tov/scalar_tensor_utils.pyjesterTOV/utils.pytests/test_inference/test_anisotropy_tov_inference.pytests/test_inference/test_config.pytests/test_inference/test_transforms.pytests/test_integration.pytests/test_scalar_tensor.py
| #SBATCH --gpus-per-node=1 | ||
| #SBATCH --cpus-per-gpu=1 | ||
| #SBATCH --mem-per-gpu=40G | ||
| #SBATCH --output="./outdir/log_smc_rw.out" |
There was a problem hiding this comment.
🧩 Analysis chain
🌐 Web query:
Does Slurm automatically create the parent directory for a #SBATCH --output path, or must that directory already exist before job submission?
💡 Result:
Slurm does not create missing parent directories for #SBATCH --output=... (or --error). The target directory must already exist and be writable at submit/run time; otherwise Slurm can’t open the file and the job will typically fail (often with an error like “No such file or directory”). (hpc-wiki.info)
Citations:
Create the output directory or use a path that is guaranteed to exist.
Slurm does not automatically create missing parent directories for #SBATCH --output. If ./outdir does not exist before job submission, Slurm will fail to open the log file and the job will fail. Either create the directory before submission or place the log file in a location that is guaranteed to exist.
Suggested change
-#SBATCH --output="./outdir/log_smc_rw.out"
+#SBATCH --output="log_smc_rw.out"📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| #SBATCH --output="./outdir/log_smc_rw.out" | |
| `#SBATCH` --output="log_smc_rw.out" |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@examples/inference/st_smc_random_walk/chiEFT/submit.sh` at line 10, The
SBATCH output path uses ./outdir which may not exist; before submitting, ensure
the directory exists by creating it (e.g., run a directory-creation step prior
to SBATCH in submit.sh) or change the --output value to a guaranteed existing
path; update the submit.sh script so the directory for the --output (referenced
by the SBATCH --output="./outdir/log_smc_rw.out" directive) is created with
mkdir -p or replaced with an existing folder.
Comment on lines
+19
to
+28
| source /home/twouters2/projects/jester_review/jester/.venv/bin/activate | ||
|
|
||
| # Display GPU name | ||
| nvidia-smi --query-gpu=name --format=csv,noheader | ||
|
|
||
| echo "==========================================" | ||
| echo "=== Running jester inference (SMC-RW) ===" | ||
| echo "==========================================" | ||
|
|
||
| run_jester_inference config.yaml |
There was a problem hiding this comment.
Make the example independent of one workstation and one current working directory.
The absolute virtualenv path on Line 19 only exists on one machine, and run_jester_inference config.yaml on Line 28 only works when the job starts in this directory. Resolve the script directory first, cd there, and make the venv path configurable so config.yaml, prior.prior, and the ChiEFT data paths stay valid.
Suggested change
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd -- "$SCRIPT_DIR/../../../../" && pwd)"
+cd "$SCRIPT_DIR"
+
- source /home/twouters2/projects/jester_review/jester/.venv/bin/activate
+JESTER_VENV="${JESTER_VENV:-$REPO_ROOT/.venv}"
+if [[ ! -f "$JESTER_VENV/bin/activate" ]]; then
+ echo "Virtualenv not found: $JESTER_VENV" >&2
+ exit 1
+fi
+source "$JESTER_VENV/bin/activate"
@@
-run_jester_inference config.yaml
+run_jester_inference config.yaml📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| source /home/twouters2/projects/jester_review/jester/.venv/bin/activate | |
| # Display GPU name | |
| nvidia-smi --query-gpu=name --format=csv,noheader | |
| echo "==========================================" | |
| echo "=== Running jester inference (SMC-RW) ===" | |
| echo "==========================================" | |
| run_jester_inference config.yaml | |
| SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" | |
| REPO_ROOT="$(cd -- "$SCRIPT_DIR/../../../../" && pwd)" | |
| cd "$SCRIPT_DIR" | |
| JESTER_VENV="${JESTER_VENV:-$REPO_ROOT/.venv}" | |
| if [[ ! -f "$JESTER_VENV/bin/activate" ]]; then | |
| echo "Virtualenv not found: $JESTER_VENV" >&2 | |
| exit 1 | |
| fi | |
| source "$JESTER_VENV/bin/activate" | |
| # Display GPU name | |
| nvidia-smi --query-gpu=name --format=csv,noheader | |
| echo "==========================================" | |
| echo "=== Running jester inference (SMC-RW) ===" | |
| echo "==========================================" | |
| run_jester_inference config.yaml |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@examples/inference/st_smc_random_walk/chiEFT/submit.sh` around lines 19 - 28,
The script currently sources an absolute venv and runs run_jester_inference
assuming the CWD; change it to resolve the script directory and cd into it, make
the virtualenv path configurable via an environment variable (e.g., VENV with a
sensible default) and source "$VENV/bin/activate", and invoke
run_jester_inference with the config and data paths prefixed by the resolved
script directory (so references to config.yaml, prior.prior and ChiEFT data
paths remain valid regardless of the initial working directory). Ensure you
update the activate path and all file arguments to use the script-dir-relative
paths and preserve the nvidia-smi and echo steps.
Comment on lines
+287
to
+289
| elif isinstance(config, ScalarTensorTOVConfig): | ||
| from jesterTOV.tov.scalar_tensor import ScalarTensorTOVSolver | ||
| return ScalarTensorTOVSolver() |
There was a problem hiding this comment.
Scalar-tensor config values are dropped in this path.
ScalarTensorTOVConfig defines beta_ST, phi_inf_tgt, and phi_c, but this branch only returns a bare solver. Later, construct_eos_and_solve_tov() sources TOV inputs exclusively from fetch_params(params), so values set in the config never reach ScalarTensorTOVSolver.solve(). The new schema is effectively ignored unless the same keys are injected separately.
One way to wire the config through
def from_config(
cls,
eos_config: BaseEOSConfig,
tov_config: BaseTOVConfig,
keep_names: list[str] | None = None,
max_nbreak_nsat: float | None = None,
fixed_params: dict[str, float] | None = None,
) -> "JesterTransform":
@@
- # Create transform
+ effective_fixed_params = dict(fixed_params or {})
+ if isinstance(tov_config, ScalarTensorTOVConfig):
+ effective_fixed_params.setdefault("beta_ST", tov_config.beta_ST)
+ effective_fixed_params.setdefault("phi_inf_tgt", tov_config.phi_inf_tgt)
+ effective_fixed_params.setdefault("phi_c", tov_config.phi_c)
+
+ # Create transform
return cls(
eos=eos,
tov_solver=tov_solver,
keep_names=keep_names,
ndat_TOV=tov_config.ndat_TOV,
min_nsat_TOV=tov_config.min_nsat_TOV,
- fixed_params=fixed_params,
+ fixed_params=effective_fixed_params,
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| elif isinstance(config, ScalarTensorTOVConfig): | |
| from jesterTOV.tov.scalar_tensor import ScalarTensorTOVSolver | |
| return ScalarTensorTOVSolver() | |
| `@classmethod` | |
| def from_config( | |
| cls, | |
| eos_config: BaseEOSConfig, | |
| tov_config: BaseTOVConfig, | |
| keep_names: list[str] | None = None, | |
| max_nbreak_nsat: float | None = None, | |
| fixed_params: dict[str, float] | None = None, | |
| ) -> "JesterTransform": | |
| effective_fixed_params = dict(fixed_params or {}) | |
| if isinstance(tov_config, ScalarTensorTOVConfig): | |
| effective_fixed_params.setdefault("beta_ST", tov_config.beta_ST) | |
| effective_fixed_params.setdefault("phi_inf_tgt", tov_config.phi_inf_tgt) | |
| effective_fixed_params.setdefault("phi_c", tov_config.phi_c) | |
| # Create transform | |
| return cls( | |
| eos=eos, | |
| tov_solver=tov_solver, | |
| keep_names=keep_names, | |
| ndat_TOV=tov_config.ndat_TOV, | |
| min_nsat_TOV=tov_config.min_nsat_TOV, | |
| fixed_params=effective_fixed_params, | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@jesterTOV/inference/transforms/transform.py` around lines 287 - 289, The
ScalarTensorTOVConfig fields (beta_ST, phi_inf_tgt, phi_c) are not propagated to
the solver because the branch returning ScalarTensorTOVSolver() discards config;
update the instantiation in the branch handling ScalarTensorTOVConfig to pass
the config values into the solver (or set them on the solver instance) so that
construct_eos_and_solve_tov and ultimately ScalarTensorTOVSolver.solve() receive
them; specifically, when config is an instance of ScalarTensorTOVConfig, extract
beta_ST/phi_inf_tgt/phi_c and supply them to ScalarTensorTOVSolver (or ensure
fetch_params reads from the config) so the solver uses the configured
scalar-tensor parameters.
Comment on lines
+578
to
+602
| Lambda_T_J = lambda_T * jnp.power(M_inf, -5.0) | ||
| Lambda_S_J = ( | ||
| ( | ||
| R, | ||
| M_inf, | ||
| nu_s, | ||
| phi_inf, | ||
| psi_s, | ||
| phi_s, | ||
| M_s, | ||
| H0_surface_1, | ||
| H0_prime_surface_1, | ||
| delta_phi_surface_1, | ||
| delta_phi_prime_surface_1, | ||
| H0_surface_2, | ||
| H0_prime_surface_2, | ||
| delta_phi_surface_2, | ||
| delta_phi_prime_surface_2, | ||
| ) = run_iteration(phi0) | ||
| # define scalar charge q (Eq. 4.18 & 4.19) | ||
| nu_s_prime = 2 * M_s / (R * (R - 2 * M_s)) + R * psi_s * psi_s | ||
| q = 2 * psi_s / nu_s_prime | ||
|
|
||
| exterior_basis_matrix = build_exterior_basis(M_inf, q, R) | ||
| exterior_basis_matrix_prime = build_exterior_basis_autodiff(M_inf, q, R) | ||
|
|
||
| # The idea: we have 6 coefficients with 4 equations. To reduce coefficients, we set two particular cases | ||
| # Case 1: H0 = 0 so cET = 0 | ||
| # Case 2: dphi = 0 so cES = 0 | ||
| # And then we normalize with one of interior solution coefficient (here case 2) | ||
| # the ratios between coefficients are then used to calculate tidal deformability | ||
| # therefore, normalization has to be consistent for all matrix. | ||
|
|
||
| # CASE 1 (Scalar deformability) | ||
| # Set cET = 0, so replace second lhs matrix column with -H01 or -dphi01 | ||
| interior_sol = ( | ||
| H0_surface_2, | ||
| H0_prime_surface_2, | ||
| delta_phi_surface_2, | ||
| delta_phi_prime_surface_2, | ||
| ) | ||
| exterior_basis_matrix_1 = exterior_basis_matrix | ||
| exterior_basis_matrix_prime_1 = exterior_basis_matrix_prime | ||
|
|
||
| mat1_p0 = jnp.array(exterior_basis_matrix_1[0]) | ||
| mat1_p1 = jnp.array(exterior_basis_matrix_1[1]) | ||
| mat1_prime_p0 = jnp.array(exterior_basis_matrix_prime_1[0]) | ||
| mat1_prime_p1 = jnp.array(exterior_basis_matrix_prime_1[1]) | ||
| mat1_p0 = mat1_p0.at[1].set(-H0_surface_1) | ||
| mat1_p1 = mat1_p1.at[1].set(-delta_phi_surface_1) | ||
| mat1_prime_p0 = mat1_prime_p0.at[1].set(-H0_prime_surface_1) | ||
| mat1_prime_p1 = mat1_prime_p1.at[1].set(-delta_phi_prime_surface_1) | ||
| exterior_basis_matrix_1 = (mat1_p0, mat1_p1) | ||
| exterior_basis_matrix_prime_1 = (mat1_prime_p0, mat1_prime_p1) | ||
|
|
||
| coeffs_1 = coeff_solver( | ||
| interior_sol, exterior_basis_matrix_1, exterior_basis_matrix_prime_1 | ||
| jnp.exp(2 * beta_ST * jnp.power(phi_inf_target, 2)) | ||
| / (4 * beta_ST * beta_ST * phi_inf_target * phi_inf_target) | ||
| ) | ||
| cQT1, c2, cQS1, cES = coeffs_1 | ||
|
|
||
| # CASE 2 (tensor deformability) | ||
| # Setting cES = 0 | ||
| # so change the coeffs into cQT, cET, cQS, c2, where c2 is particular soluion coeff | ||
| # and replace equation relating to ES to be -H01 or -dphi01 | ||
| interior_sol = ( | ||
| H0_surface_2, | ||
| H0_prime_surface_2, | ||
| delta_phi_surface_2, | ||
| delta_phi_prime_surface_2, | ||
| * lambda_S | ||
| * jnp.power(M_inf, -5.0) | ||
| ) | ||
| Lambda_ST1_J = ( | ||
| ( | ||
| -jnp.exp(beta_ST * jnp.power(phi_inf_target, 2)) | ||
| / (2 * beta_ST * phi_inf_target) | ||
| ) | ||
| exterior_basis_matrix_2 = exterior_basis_matrix | ||
| exterior_basis_matrix_prime_2 = exterior_basis_matrix_prime | ||
|
|
||
| mat2_part0 = jnp.array(exterior_basis_matrix_2[0]) | ||
| mat2_part1 = jnp.array(exterior_basis_matrix_2[1]) | ||
| mat2_prime_part0 = jnp.array(exterior_basis_matrix_prime_2[0]) | ||
| mat2_prime_part1 = jnp.array(exterior_basis_matrix_prime_2[1]) | ||
| mat2_part0 = mat2_part0.at[3].set(-H0_surface_1) | ||
| mat2_part1 = mat2_part1.at[3].set(-delta_phi_surface_1) | ||
| mat2_prime_part0 = mat2_prime_part0.at[3].set(-H0_prime_surface_1) | ||
| mat2_prime_part1 = mat2_prime_part1.at[3].set(-delta_phi_prime_surface_1) | ||
|
|
||
| exterior_basis_matrix_2 = (mat2_part0, mat2_part1) | ||
| exterior_basis_matrix_prime_2 = (mat2_prime_part0, mat2_prime_part1) | ||
|
|
||
| coeffs_2 = coeff_solver( | ||
| interior_sol, exterior_basis_matrix_2, exterior_basis_matrix_prime_2 | ||
| * lambda_ST1 | ||
| * jnp.power(M_inf, -5.0) | ||
| ) | ||
| Lambda_ST2_J = ( | ||
| ( | ||
| -jnp.exp(beta_ST * jnp.power(phi_inf_target, 2)) | ||
| / (2 * beta_ST * phi_inf_target) | ||
| ) | ||
| cQT2, cET, cQS2, c2 = coeffs_2 | ||
| * lambda_ST2 | ||
| * jnp.power(M_inf, -5.0) | ||
| ) |
There was a problem hiding this comment.
Potential division by zero in GR limit when beta_ST → 0 or phi_inf_target → 0.
The Jordan frame tidal deformability conversions divide by beta_ST and phi_inf_target. While the default phi_inf_target=1e-3 (line 631) is safe, if beta_ST=0 (GR limit) is passed, these computations will produce inf or nan.
Consider documenting this limitation or adding a guard. If GR limit testing is not a use case for this solver, document that users should use GRTOVSolver instead.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@jesterTOV/tov/scalar_tensor.py` around lines 578 - 602, Lambda_S_J,
Lambda_ST1_J and Lambda_ST2_J perform divisions by beta_ST and phi_inf_target
and can produce inf/nan when beta_ST==0 or phi_inf_target==0; add a guard at the
start of the conversion block that checks for small |beta_ST| or
|phi_inf_target| (e.g. abs(x) < eps) and either: (a) raise a clear ValueError
suggesting GR limit users should use GRTOVSolver, or (b) use a safe fallback
(set the problematic Lambda_*_J expressions to their GR-consistent limits or
zero) depending on intended behavior; implement the check around the Lambda_*_J
computations so Lambda_S_J, Lambda_ST1_J, Lambda_ST2_J are never computed with a
zero divisor and include beta_ST and phi_inf_target in the error message if
raising.
Comment on lines
+630
to
+632
| beta_ST = tov_params.get("beta_ST", 0.0) | ||
| phi_inf_target = tov_params.get("phi_inf_tgt", 1e-3) | ||
| phi0 = tov_params.get("phi_c", 1.0) |
There was a problem hiding this comment.
Default beta_ST=0.0 conflicts with downstream division.
The default value of 0.0 for beta_ST will cause division-by-zero errors in the Jordan frame conversions (lines 582, 590, 598). Either remove the default to require explicit parameter passing, or use a small non-zero default that represents a near-GR regime.
Proposed fix: remove defaults to enforce explicit parameters
- beta_ST = tov_params.get("beta_ST", 0.0)
- phi_inf_target = tov_params.get("phi_inf_tgt", 1e-3)
- phi_c = tov_params.get("phi_c", 1.0)
+ beta_ST = tov_params["beta_ST"]
+ phi_inf_target = tov_params["phi_inf_tgt"]
+ phi_c = tov_params["phi_c"]This is safe because get_required_parameters() lists all three parameters, and fetch_params() in the base class validates their presence before solve() is called.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@jesterTOV/tov/scalar_tensor.py` around lines 630 - 632, The default
beta_ST=0.0 in tov_params.get("beta_ST", 0.0) leads to downstream
division-by-zero in the Jordan-frame conversion code; remove the unsafe default
so beta_ST must be provided (e.g., replace tov_params.get("beta_ST", 0.0) with a
direct lookup like tov_params["beta_ST"] or tov_params.get("beta_ST") without a
fallback), leaving phi_inf_target and phi0 handling unchanged; ensure this
relies on existing get_required_parameters()/fetch_params()/solve() validation
to enforce presence of beta_ST.
Comment on lines
+354
to
+364
| # Extract indices. Assuming a max of 4 segments for static functions (for jax compatibility). | ||
| # We use size=4 and fill with a safe dummy index (len(m)-1) | ||
| # If a segment is dummy, start=end, so length=0, and it is ignored. | ||
| starts = jnp.where(diff == 1, size=4, fill_value=len(m)-1)[0] | ||
| ends = jnp.where(diff == -1, size=4, fill_value=len(m)-1)[0] | ||
|
|
||
| # Extract segment bounds for manual processing | ||
| s1_start, s1_end = starts[0], ends[0] | ||
| s2_start, s2_end = starts[1], ends[1] | ||
| s3_start, s3_end = starts[2], ends[2] | ||
| s4_start, s4_end = starts[3], ends[3] |
There was a problem hiding this comment.
Don't hard-code the stable-branch count to four.
jnp.where(..., size=4, ...) plus the manual s1-s4 unpacking silently drops every stable segment after the fourth. A multi-transition EOS would therefore lose physically valid branches with no signal to the caller, which is exactly the regime this helper is trying to preserve.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@jesterTOV/utils.py` around lines 354 - 364, The code hard-codes four segments
by using jnp.where(..., size=4, ...) and unpacking into s1_start..s4_end which
drops any extra stable branches; change to compute starts = jnp.nonzero(diff ==
1)[0] and ends = jnp.nonzero(diff == -1)[0] (or jnp.where without a fixed size)
to get variable-length index arrays, then replace the fixed unpacking of
s1_start..s4_end with code that iterates over the computed starts/ends (or
returns them as arrays) so all detected segments are preserved (refer to the
variables diff, starts, ends, and the s1_start..s4_end handling in this
function).
Comment on lines
+469
to
+480
| # Finite difference | ||
| central_diff_numer = fp[2:] - fp[:-2] | ||
| central_diff_denom = xp[2:] - xp[:-2] | ||
| forward_m0 = (fp[1] - fp[0]) / (xp[1] - xp[0]) # for starting index | ||
| interior_m1_m = central_diff_numer / central_diff_denom # for middle indices | ||
| backward_m_end = (fp[-1] - fp[-2]) / (xp[-1] - xp[-2]) # for last index | ||
|
|
||
| # Concatenate them to form the final array of derivatives dydx (length m) | ||
| dydx = jnp.concatenate(( | ||
| forward_m0[None], # [None] adds a dimension to match the others | ||
| interior_m1_m, | ||
| backward_m_end[None] |
There was a problem hiding this comment.
Make the Hermite slopes respect segment boundaries.
These centered differences always use xp[i-1] and xp[i+1] for interior nodes. When limit_by_MTOV_and_interpolate() stitches disjoint stable branches, the first/last node of each retained branch therefore picks up slope information from the discarded unstable region next to it, so the spline can overshoot right at the branch entrances and exits. Compute slopes per stable segment, or fall back to one-sided slopes at every segment edge.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@jesterTOV/utils.py` around lines 469 - 480, The centered differences
(central_diff_numer, central_diff_denom, interior_m1_m) are computed across the
entire xp array and can leak slope info across stitched branch boundaries;
update the slope computation to operate per contiguous stable segment: detect
segment boundaries in xp (e.g., where consecutive index or xp-step indicates a
gap), then for each segment compute forward_m0 / backward_m_end at its ends and
central differences only for its interior points, and finally concatenate
per-segment dydx pieces to form the full dydx array so no centered difference
uses xp values from a different segment.
Comment on lines
+279
to
+290
| # Get masses and radii from both solutions | ||
| masses_EOS_gr = result_gr["masses_EOS"] | ||
| radii_EOS_gr = result_gr["radii_EOS"] | ||
|
|
||
| masses_EOS_mg = result_mg["masses_EOS"] | ||
| radii_EOS_mg = result_mg["radii_EOS"] | ||
|
|
||
| # Compute R1.4 | ||
| R1_4_gr = jnp.interp(1.4, masses_EOS_gr, radii_EOS_gr) | ||
| R1_4_mg = jnp.interp(1.4, masses_EOS_mg, radii_EOS_mg) | ||
| diff = float(abs(R1_4_gr - R1_4_mg) / R1_4_gr) | ||
| assert diff > 1e-6, "gamma should shift the M-R family" |
There was a problem hiding this comment.
🧩 Analysis chain
🌐 Web query:
Does jax.numpy.interp require the xp array to be increasing, and what happens if xp is non-monotonic?
💡 Result:
Yes. jax.numpy.interp expects xp to be a 1D sorted (increasing) array. (docs.jax.dev)
If xp is non-monotonic, the function generally does not raise an error, but the interpolation it returns is not well-defined / may be wrong (NumPy’s reference behavior is: “results are meaningless” when xp is non-increasing, and this is not explicitly enforced). (numpy.org)
Citations:
- 1: https://docs.jax.dev/en/latest/_autosummary/jax.numpy.interp.html
- 2: https://numpy.org/doc/stable/reference/generated/numpy.interp.html
🏁 Script executed:
# Check the file exists and view lines around 279-290
wc -l tests/test_inference/test_anisotropy_tov_inference.pyRepository: nuclear-multimessenger-astronomy/jester
Length of output: 140
🏁 Script executed:
# View the relevant section
sed -n '270,295p' tests/test_inference/test_anisotropy_tov_inference.pyRepository: nuclear-multimessenger-astronomy/jester
Length of output: 878
Avoid interpolating on the full mass curve.
masses_EOS_* are not guaranteed to remain monotonic after the maximum-mass point. jnp.interp requires sorted x-values and will silently produce incorrect results on non-monotonic input, making this test unreliable. Slice to the stable branch first before interpolating.
Suggested fix
+ gr_max_idx = jnp.argmax(masses_EOS_gr)
+ mg_max_idx = jnp.argmax(masses_EOS_mg)
+
- R1_4_gr = jnp.interp(1.4, masses_EOS_gr, radii_EOS_gr)
- R1_4_mg = jnp.interp(1.4, masses_EOS_mg, radii_EOS_mg)
+ R1_4_gr = jnp.interp(
+ 1.4, masses_EOS_gr[: gr_max_idx + 1], radii_EOS_gr[: gr_max_idx + 1]
+ )
+ R1_4_mg = jnp.interp(
+ 1.4, masses_EOS_mg[: mg_max_idx + 1], radii_EOS_mg[: mg_max_idx + 1]
+ )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/test_inference/test_anisotropy_tov_inference.py` around lines 279 -
290, The test currently interpolates R1.4 over the full mass curves
(masses_EOS_gr/radii_EOS_gr and masses_EOS_mg/radii_EOS_mg) which can be
non-monotonic past the maximum-mass point; first find the index of the maximum
mass (e.g., using jnp.argmax on masses_EOS_gr and masses_EOS_mg), slice each
mass and radius array to the stable branch up through that index (inclusive) so
x is monotonic, then perform the jnp.interp on those sliced arrays to compute
R1_4_gr and R1_4_mg before comparing. Ensure you reference masses_EOS_gr,
radii_EOS_gr, masses_EOS_mg, radii_EOS_mg and update how R1_4_gr / R1_4_mg are
computed accordingly.
Comment on lines
+160
to
+175
| def test_valid_scalar_tensor_tov_config(self): | ||
| """Test valid Scalar-Tensor TOV configuration.""" | ||
| from pydantic import TypeAdapter | ||
|
|
||
| adapter = TypeAdapter(schema.TOVConfig) | ||
| config_dict = { | ||
| "type": "scalar_tensor", | ||
| "beta_ST": -1.0, | ||
| "phi_inf_tgt": 1e-3, | ||
| "phi_c": 1.0, | ||
| } | ||
| config = adapter.validate_python(config_dict) | ||
| assert isinstance(config, schema.ScalarTensorTOVConfig) | ||
| assert config.type == "scalar_tensor" | ||
| assert config.beta_ST == -1.0 | ||
|
|
There was a problem hiding this comment.
Assert the other scalar-tensor fields here too.
This only verifies beta_ST, so a regression in phi_inf_tgt or phi_c parsing would still pass even though the solver path depends on those exact keys.
Suggested test addition
config = adapter.validate_python(config_dict)
assert isinstance(config, schema.ScalarTensorTOVConfig)
assert config.type == "scalar_tensor"
assert config.beta_ST == -1.0
+ assert config.phi_inf_tgt == 1e-3
+ assert config.phi_c == 1.0📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| def test_valid_scalar_tensor_tov_config(self): | |
| """Test valid Scalar-Tensor TOV configuration.""" | |
| from pydantic import TypeAdapter | |
| adapter = TypeAdapter(schema.TOVConfig) | |
| config_dict = { | |
| "type": "scalar_tensor", | |
| "beta_ST": -1.0, | |
| "phi_inf_tgt": 1e-3, | |
| "phi_c": 1.0, | |
| } | |
| config = adapter.validate_python(config_dict) | |
| assert isinstance(config, schema.ScalarTensorTOVConfig) | |
| assert config.type == "scalar_tensor" | |
| assert config.beta_ST == -1.0 | |
| def test_valid_scalar_tensor_tov_config(self): | |
| """Test valid Scalar-Tensor TOV configuration.""" | |
| from pydantic import TypeAdapter | |
| adapter = TypeAdapter(schema.TOVConfig) | |
| config_dict = { | |
| "type": "scalar_tensor", | |
| "beta_ST": -1.0, | |
| "phi_inf_tgt": 1e-3, | |
| "phi_c": 1.0, | |
| } | |
| config = adapter.validate_python(config_dict) | |
| assert isinstance(config, schema.ScalarTensorTOVConfig) | |
| assert config.type == "scalar_tensor" | |
| assert config.beta_ST == -1.0 | |
| assert config.phi_inf_tgt == 1e-3 | |
| assert config.phi_c == 1.0 |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/test_inference/test_config.py` around lines 160 - 175, The test
test_valid_scalar_tensor_tov_config currently only asserts beta_ST; update it to
also assert that the validated config is a schema.ScalarTensorTOVConfig and that
the scalar-tensor-specific fields phi_inf_tgt and phi_c are correctly parsed
(e.g., exact values and/or types) after using TypeAdapter(schema.TOVConfig). In
short, inside test_valid_scalar_tensor_tov_config add assertions that
config.phi_inf_tgt == 1e-3 and config.phi_c == 1.0 (and optional type checks) so
regressions in parsing those keys are caught.
Comment on lines
+80
to
+100
| def test_from_config_scalar_tensor(self): | ||
| """Test creating Scalar-Tensor transform via from_config.""" | ||
| eos_config = MetamodelEOSConfig( | ||
| type="metamodel", | ||
| ndat_metamodel=100, | ||
| nmax_nsat=2.0, | ||
| nb_CSE=0, | ||
| crust_name="DH", | ||
| ) | ||
| tov_config = ScalarTensorTOVConfig( | ||
| type="scalar_tensor", | ||
| beta_ST=-1.0, | ||
| phi_inf_tgt=1e-3, | ||
| phi_c=1.0, | ||
| ) | ||
|
|
||
| transform = JesterTransform.from_config(eos_config, tov_config) | ||
|
|
||
| assert transform is not None | ||
| assert "MetaModel_EOS_model" in transform.get_eos_type() | ||
| assert "ScalarTensorTOVSolver" in str(type(transform.tov_solver)) |
There was a problem hiding this comment.
This only tests solver selection, not scalar-tensor parameter plumbing.
Because the assertions stop at the solver type, this still passes if from_config() drops beta_ST, phi_inf_tgt, and phi_c. Please add one assertion that those config values actually land in the transform path.
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
jesterTOV/tov/base.py (1)
107-144:⚠️ Potential issue | 🟠 MajorDon’t leave the EOS validity check disconnected.
Line 107 computes
is_valid_eos, but nothing ever consumes it, so this method still buildspcsand vmapssolve()over invalid EOS tables. Non-positive EOS entries can therefore still reachinterp_in_logspace()and generate NaNs here. Please apply the invalid-EOS masking or early-return path before the pressure grid is built.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@jesterTOV/tov/base.py` around lines 107 - 144, The computed is_valid_eos is ignored—ensure invalid EOS tables are excluded before building the pressure grid by either masking eos_data or early-returning: use is_valid_eos to either (A) apply jax.tree_util.tree_map to replace fields in eos_data with jnp.nan when invalid (so downstream _get_pc_min/_get_pc_max/solve see NaNs), or (B) detect invalid entries and return an appropriately shaped empty/NaN family via _create_family_data immediately; update the code around is_valid_eos, pcs, _get_pc_min, _get_pc_max, solve, jax.vmap, and the final _create_family_data call so no invalid eos_data reaches interp_in_logspace.
♻️ Duplicate comments (4)
jesterTOV/inference/transforms/transform.py (1)
287-290:⚠️ Potential issue | 🔴 CriticalScalar-tensor config values are still not propagated to solver inputs.
At Lines 287–290, the new branch selects
ScalarTensorTOVSolver, butbeta_ST,phi_inf_tgt, andphi_cfromScalarTensorTOVConfigare not injected anywhere. Since Line 354 sources TOV params fromfetch_params(params), config-set values are ignored unless duplicated in priors orfixed_params.Suggested fix (wire config defaults into fixed_params in
from_config())@@ def from_config( cls, eos_config: BaseEOSConfig, tov_config: BaseTOVConfig, keep_names: list[str] | None = None, max_nbreak_nsat: float | None = None, fixed_params: dict[str, float] | None = None, ) -> "JesterTransform": @@ - # Create transform + effective_fixed_params = dict(fixed_params or {}) + if isinstance(tov_config, ScalarTensorTOVConfig): + effective_fixed_params.setdefault("beta_ST", tov_config.beta_ST) + effective_fixed_params.setdefault("phi_inf_tgt", tov_config.phi_inf_tgt) + effective_fixed_params.setdefault("phi_c", tov_config.phi_c) + + # Create transform return cls( eos=eos, tov_solver=tov_solver, keep_names=keep_names, ndat_TOV=tov_config.ndat_TOV, min_nsat_TOV=tov_config.min_nsat_TOV, - fixed_params=fixed_params, + fixed_params=effective_fixed_params, )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@jesterTOV/inference/transforms/transform.py` around lines 287 - 290, The ScalarTensorTOVConfig values (beta_ST, phi_inf_tgt, phi_c) are created but never propagated into solver inputs; update from_config() so when it returns a ScalarTensorTOVSolver it also injects these config defaults into the parameter source used by fetch_params(params) (e.g., merge them into fixed_params or the params dict) so the solver receives beta_ST, phi_inf_tgt, and phi_c; locate the branch that instantiates ScalarTensorTOVSolver (the isinstance(config, ScalarTensorTOVConfig) block), the from_config() factory, and fetch_params(params) and ensure the config fields are copied into fixed_params before the solver is returned or before fetch_params is called.jesterTOV/tov/scalar_tensor.py (1)
626-631:⚠️ Potential issue | 🟠 MajorDon’t default required ST parameters to GR-like values.
fetch_params()/get_required_parameters()already make these inputs mandatory. Falling back tobeta_ST=0.0andphi_inf_tgt=1e-3silently changes the model when a key is missing, and zero-valuedbeta_ST/phi_inf_tgtlater feed zero divisors into the Jordan-frameLambda_S/Lambda_ST*conversions. Use direct lookups here and either guard the GR limit explicitly or route it toGRTOVSolver.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@jesterTOV/tov/scalar_tensor.py` around lines 626 - 631, The solve method currently uses tov_params.get with defaults for beta_ST and phi_inf_tgt which silently changes behavior and can introduce zero divisors in the Jordan-frame conversions; change the lookups in solve to use direct key access for "beta_ST" and "phi_inf_tgt" (and "phi_c" if required) so missing parameters raise immediately, and then explicitly handle the GR limit: either detect beta_ST (or phi_inf_tgt) values that indicate GR and dispatch to GRTOVSolver, or raise a clear error if GR is unsupported; ensure the code paths that compute Lambda_S / Lambda_ST* validate nonzero denominators before conversion.jesterTOV/utils.py (2)
354-364:⚠️ Potential issue | 🟠 MajorDon’t truncate stable branches after the fourth segment.
jnp.where(..., size=4)plus the manuals1–s4unpacking still drops every later stable branch. On multi-transition EOS tables, that silently removes physically valid M–R–Λ segments and can miss the actual MTOV branch entirely. Please drive the interpolation grid from the fullstarts/endsset instead of a fixed four-segment budget.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@jesterTOV/utils.py` around lines 354 - 364, The current logic truncates segments by using jnp.where(..., size=4) and unpacking s1_start..s4_end, which drops any stable branches beyond four; change this to compute full starts and ends without the fixed size limit (remove size=4 and fill_value use) and iterate over the resulting starts and ends arrays to build the interpolation grid dynamically instead of using s1_start, s2_start, s3_start, s4_start (and corresponding ends). Ensure you treat dummy/invalid segments by checking start==end (or other existing sentinel logic) when constructing segments so all valid M–R–Λ branches are preserved.
474-487:⚠️ Potential issue | 🟠 MajorCompute Hermite slopes per retained segment.
The derivatives on Lines 474-478 still use
xp[i-1]andxp[i+1]across the whole stitched curve. For the first and last node of each retained stable branch, that pulls slope information from the discarded unstable region and can overshoot right at the branch entrance or exit. Please computedydxon each stable slice separately, or use one-sided slopes at every segment boundary.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@jesterTOV/utils.py` around lines 474 - 487, The current slope construction uses central differences across the entire stitched curve (central_diff_numer/central_diff_denom -> interior_m1_m) and single forward/backward endpoints (forward_m0, backward_m_end) which lets discarded unstable nodes leak into slopes; update dydx to compute slopes per retained stable slice instead: iterate over each retained segment boundary (or partition xp/fp by the retained indices) and compute central differences only within that slice, using one-sided differences at that slice's local first and last points, then concatenate those per-slice slope arrays to form dydx so that forward_m0/backward_m_end are local to each retained branch rather than global.
🧹 Nitpick comments (2)
.pre-commit-config.yaml (1)
6-12: Notebooks are excluded fromblackandruffwithout alternate linting coverage.Lines 6 and 12 exclude
.ipynbfiles, but there are nonbqa-*hooks in the pre-commit configuration to lint/format notebooks. Since no notebooks are modified in this PR, this doesn't directly impact the current changes. However, if notebooks indocs/examples/eos_tov/should maintain code quality standards, consider addingnbqahooks or documenting why notebook formatting is intentionally excluded.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In @.pre-commit-config.yaml around lines 6 - 12, The pre-commit config currently excludes .ipynb via the ruff hook (id: ruff, args: ["--fix"], exclude: \.ipynb$) but provides no nbqa hooks to lint/format notebooks; add nbqa hooks (e.g., ids like nbqa-black and nbqa-ruff from the nbQA repository) to run black and ruff against notebooks in place of direct exclusion, or alternatively add a short comment/doc entry explaining why notebooks under docs/examples/eos_tov/ are intentionally exempt from notebook linting/formatting to ensure code quality expectations are clear.jesterTOV/tov/scalar_tensor.py (1)
299-302: Add full type annotations to_compiled_tov_solve().This new helper drives the entire ST solve, but it currently has an untyped parameter list and return tuple. Please annotate the scalar inputs, EOS arrays, and returned scalars explicitly, then rerun
uv run pyrighton the changed files. As per coding guidelines, 'All new code MUST include comprehensive type hints using standard Python 3.10+ syntax (e.g., list[float], dict[str, float]), jaxtyping for JAX arrays (e.g., Float[Array, "n_points"]), and Pydantic for configs' and 'Run pyright type checking withuv run pyrightbefore considering changes ready for human review'.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@jesterTOV/tov/scalar_tensor.py` around lines 299 - 302, The function _compiled_tov_solve currently has no type annotations; add full PEP 484/3.10-style type hints: annotate scalar inputs (pc, beta_ST, phi_inf_target, phi0, max_iterations) as float|int where appropriate (use float for continuous scalars), annotate EOS-related JAX arrays ps, hs, es, dloge_dlogps with jaxtyping types like Float[Array, "n_points"] (or appropriate shape names), and annotate the return as a tuple of floats (e.g., tuple[float, float, ...] matching the actual returned scalars). Keep the `@functools.partial`(jax.jit, static_argnames=["max_iterations"]) decorator and ensure imports for jaxtyping Array/Float types are present, then rerun the pyright check (uv run pyright) to fix any remaining type issues.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Outside diff comments:
In `@jesterTOV/tov/base.py`:
- Around line 107-144: The computed is_valid_eos is ignored—ensure invalid EOS
tables are excluded before building the pressure grid by either masking eos_data
or early-returning: use is_valid_eos to either (A) apply jax.tree_util.tree_map
to replace fields in eos_data with jnp.nan when invalid (so downstream
_get_pc_min/_get_pc_max/solve see NaNs), or (B) detect invalid entries and
return an appropriately shaped empty/NaN family via _create_family_data
immediately; update the code around is_valid_eos, pcs, _get_pc_min, _get_pc_max,
solve, jax.vmap, and the final _create_family_data call so no invalid eos_data
reaches interp_in_logspace.
---
Duplicate comments:
In `@jesterTOV/inference/transforms/transform.py`:
- Around line 287-290: The ScalarTensorTOVConfig values (beta_ST, phi_inf_tgt,
phi_c) are created but never propagated into solver inputs; update from_config()
so when it returns a ScalarTensorTOVSolver it also injects these config defaults
into the parameter source used by fetch_params(params) (e.g., merge them into
fixed_params or the params dict) so the solver receives beta_ST, phi_inf_tgt,
and phi_c; locate the branch that instantiates ScalarTensorTOVSolver (the
isinstance(config, ScalarTensorTOVConfig) block), the from_config() factory, and
fetch_params(params) and ensure the config fields are copied into fixed_params
before the solver is returned or before fetch_params is called.
In `@jesterTOV/tov/scalar_tensor.py`:
- Around line 626-631: The solve method currently uses tov_params.get with
defaults for beta_ST and phi_inf_tgt which silently changes behavior and can
introduce zero divisors in the Jordan-frame conversions; change the lookups in
solve to use direct key access for "beta_ST" and "phi_inf_tgt" (and "phi_c" if
required) so missing parameters raise immediately, and then explicitly handle
the GR limit: either detect beta_ST (or phi_inf_tgt) values that indicate GR and
dispatch to GRTOVSolver, or raise a clear error if GR is unsupported; ensure the
code paths that compute Lambda_S / Lambda_ST* validate nonzero denominators
before conversion.
In `@jesterTOV/utils.py`:
- Around line 354-364: The current logic truncates segments by using
jnp.where(..., size=4) and unpacking s1_start..s4_end, which drops any stable
branches beyond four; change this to compute full starts and ends without the
fixed size limit (remove size=4 and fill_value use) and iterate over the
resulting starts and ends arrays to build the interpolation grid dynamically
instead of using s1_start, s2_start, s3_start, s4_start (and corresponding
ends). Ensure you treat dummy/invalid segments by checking start==end (or other
existing sentinel logic) when constructing segments so all valid M–R–Λ branches
are preserved.
- Around line 474-487: The current slope construction uses central differences
across the entire stitched curve (central_diff_numer/central_diff_denom ->
interior_m1_m) and single forward/backward endpoints (forward_m0,
backward_m_end) which lets discarded unstable nodes leak into slopes; update
dydx to compute slopes per retained stable slice instead: iterate over each
retained segment boundary (or partition xp/fp by the retained indices) and
compute central differences only within that slice, using one-sided differences
at that slice's local first and last points, then concatenate those per-slice
slope arrays to form dydx so that forward_m0/backward_m_end are local to each
retained branch rather than global.
---
Nitpick comments:
In @.pre-commit-config.yaml:
- Around line 6-12: The pre-commit config currently excludes .ipynb via the ruff
hook (id: ruff, args: ["--fix"], exclude: \.ipynb$) but provides no nbqa hooks
to lint/format notebooks; add nbqa hooks (e.g., ids like nbqa-black and
nbqa-ruff from the nbQA repository) to run black and ruff against notebooks in
place of direct exclusion, or alternatively add a short comment/doc entry
explaining why notebooks under docs/examples/eos_tov/ are intentionally exempt
from notebook linting/formatting to ensure code quality expectations are clear.
In `@jesterTOV/tov/scalar_tensor.py`:
- Around line 299-302: The function _compiled_tov_solve currently has no type
annotations; add full PEP 484/3.10-style type hints: annotate scalar inputs (pc,
beta_ST, phi_inf_target, phi0, max_iterations) as float|int where appropriate
(use float for continuous scalars), annotate EOS-related JAX arrays ps, hs, es,
dloge_dlogps with jaxtyping types like Float[Array, "n_points"] (or appropriate
shape names), and annotate the return as a tuple of floats (e.g., tuple[float,
float, ...] matching the actual returned scalars). Keep the
`@functools.partial`(jax.jit, static_argnames=["max_iterations"]) decorator and
ensure imports for jaxtyping Array/Float types are present, then rerun the
pyright check (uv run pyright) to fix any remaining type issues.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 0e55c4ae-1fb2-47fb-8351-d4074b5bbfd2
📒 Files selected for processing (6)
.pre-commit-config.yamljesterTOV/inference/config/schemas/tov.pyjesterTOV/inference/transforms/transform.pyjesterTOV/tov/base.pyjesterTOV/tov/scalar_tensor.pyjesterTOV/utils.py
🚧 Files skipped from review as they are similar to previous changes (1)
- jesterTOV/inference/config/schemas/tov.py
…ius curve. 2. Add config toggle to turn off tidal deformability calculation in scalar tensor, to save resource when not needed.
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (5)
jesterTOV/utils.py (1)
632-641: Add docstring and type hints forget_MR_split_index.This function lacks documentation and proper type annotations. Also, the magic number
1e-12for numerical stability should be documented.📝 Proposed documentation and type hints
-def get_MR_split_index(x, y, jump_threshold=20.0): - N = x.shape[0] - ds2 = jnp.diff(x)**2 + jnp.diff(y)**2 - - # Find the largest gap - max_idx = jnp.argmax(ds2) - median_ds2 = jnp.median(ds2) + 1e-12 - is_jump = ((ds2[max_idx] > (jump_threshold * median_ds2)) & (max_idx > 0)).astype(jnp.int32) - - split_idx = is_jump * (max_idx + 1) + (1 - is_jump) * N - return split_idx +def get_MR_split_index( + x: Float[Array, "n"], y: Float[Array, "n"], jump_threshold: float = 20.0 +) -> Int[Array, ""]: + r""" + Find the index where an M-R curve should be split due to a discontinuity. + + Detects the largest gap in the (x, y) curve by computing squared distances + between consecutive points. If the maximum gap exceeds ``jump_threshold`` + times the median gap, returns the split index; otherwise returns N (no split). + + Parameters + ---------- + x : Float[Array, "n"] + x-coordinates (e.g., radius) + y : Float[Array, "n"] + y-coordinates (e.g., mass) + jump_threshold : float + Multiplier for median gap to identify a discontinuity (default: 20.0) + + Returns + ------- + Int[Array, ""] + Split index, or N if no significant gap is found + """ + N = x.shape[0] + ds2 = jnp.diff(x) ** 2 + jnp.diff(y) ** 2 + + # Find the largest gap + max_idx = jnp.argmax(ds2) + # Add small epsilon for numerical stability when median is near zero + median_ds2 = jnp.median(ds2) + 1e-12 + is_jump = ((ds2[max_idx] > (jump_threshold * median_ds2)) & (max_idx > 0)).astype( + jnp.int32 + ) + + split_idx = is_jump * (max_idx + 1) + (1 - is_jump) * N + return split_idx🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@jesterTOV/utils.py` around lines 632 - 641, Add a clear docstring and Python type hints to the get_MR_split_index function: annotate parameters x and y as jnp.ndarray (or Array) and jump_threshold as float, and the return type as int; in the docstring describe purpose, parameters, return value, and behavior (how split index is computed, meaning of returned N when no jump); document the numeric stabilizer 1e-12 and why it's added to median_ds2; keep existing logic (use jnp.diff, jnp.argmax, jnp.median) and mention that split_idx is 1-based offset (max_idx + 1) when a jump is detected and N otherwise so callers understand semantics.jesterTOV/inference/likelihoods/nicer.py (4)
707-709: Hardcoded mass grid range may not suit all applications.The mass grid is hardcoded to
[0.1, 3.5]solar masses. Consider making this configurable via__init__parameters if different mass ranges are needed for different mock scenarios.♻️ Optional enhancement
def __init__( self, csv_file: str, penalty_value: float = -99999.0, - N_masses_evaluation: int = 200 + N_masses_evaluation: int = 200, + mass_grid_min: float = 0.1, + mass_grid_max: float = 3.5, ) -> None: ... + self.mass_grid_min = mass_grid_min + self.mass_grid_max = mass_grid_maxThen in
evaluate:- m_grid = jnp.linspace(0.1, 3.5, self.N_masses_evaluation) + m_grid = jnp.linspace(self.mass_grid_min, self.mass_grid_max, self.N_masses_evaluation)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@jesterTOV/inference/likelihoods/nicer.py` around lines 707 - 709, The mass grid range in nicer.NicerLikelihood (m_grid and dm created in the evaluation code) is hardcoded to [0.1, 3.5]; make the lower and upper bounds configurable by adding new __init__ parameters (e.g., mass_min and mass_max) and storing them on self, defaulting to 0.1 and 3.5, then replace the hardcoded literals in the block that constructs m_grid and dm with self.mass_min and self.mass_max so different mock scenarios can pass alternate ranges (also ensure any docs/tests that reference N_masses_evaluation still work and consider using the new attributes in evaluate where appropriate).
294-307: Clarify the combination logic comment.The comment at lines 294-295 is unclear. The likelihood combination works correctly because:
- When no phase transition exists, segment 2 terms are
-infand contributeexp(-inf) = 0to the logsumexp- The
-log(2)normalizes across the 2 analysis groups (Amsterdam/Maryland), not the 4 termsConsider updating the comment to explain this more precisely:
- # Average over all samples for each group - # In this case, even if data is accidentally splitted somewhere, log_likelihood will be exactly the same. + # Average over all samples for each group (log-mean = logsumexp - log(N)) + # Empty segment 2 contributes -inf terms, which yield exp(-inf) = 0 in logsumexp, + # so the result is equivalent to evaluating only the valid segment.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@jesterTOV/inference/likelihoods/nicer.py` around lines 294 - 307, Update the unclear comment above the likelihood combination to explicitly state that when there is no phase transition the segment-2 terms (amsterdam_logprobs_2, maryland_logprobs_2) become -inf and therefore contribute zero after exp/logsumexp, and that the final subtraction of jnp.log(2.0) is normalizing across the two analysis groups (Amsterdam and Maryland) rather than the four intermediate terms; reference the computed intermediates logL_amsterdam_1, logL_maryland_1, logL_amsterdam_2, logL_maryland_2 and the use of logsumexp so readers understand why seg2 can be ignored and why -jnp.log(2.0) is applied.
217-244: Consider extracting segment-splitting logic into a helper function.The segment-splitting logic (lines 217-244) is duplicated almost identically in
NICERKDELikelihood.evaluate(lines 448-474) andMockMRLikelihood.evaluate(lines 685-705). This violates DRY and increases maintenance burden.♻️ Proposed helper function
Add a helper at module level or in
utils.py:def split_mr_segments( masses_EOS: Float[Array, " n"], radii_EOS: Float[Array, " n"], ) -> tuple[ Float[Array, " n"], Float[Array, " n"], Float, Float, # seg1: m, r, min, max Float[Array, " n"], Float[Array, " n"], Float, Float, # seg2: m, r, min, max ]: """Split M-R arrays into two segments based on phase transition detection.""" split_idx = utils.get_MR_split_index(masses_EOS, radii_EOS) idx = jnp.arange(masses_EOS.shape[0]) mask1, mask2 = idx < split_idx, idx >= split_idx # Segment 1 m_eos_1 = jnp.where(mask1, masses_EOS, jnp.inf) r_eos_1 = jnp.where(mask1, radii_EOS, 0.0) sort_1 = jnp.argsort(m_eos_1) m_eos_1, r_eos_1 = m_eos_1[sort_1], r_eos_1[sort_1] seg1_min = m_eos_1[0] seg1_max = jnp.max(jnp.where(m_eos_1 == jnp.inf, -jnp.inf, m_eos_1)) # Segment 2 m_eos_2 = jnp.where(mask2, masses_EOS, jnp.inf) r_eos_2 = jnp.where(mask2, radii_EOS, 0.0) sort_2 = jnp.argsort(m_eos_2) m_eos_2, r_eos_2 = m_eos_2[sort_2], r_eos_2[sort_2] seg2_min = m_eos_2[0] seg2_max = jnp.max(jnp.where(m_eos_2 == jnp.inf, -jnp.inf, m_eos_2)) return m_eos_1, r_eos_1, seg1_min, seg1_max, m_eos_2, r_eos_2, seg2_min, seg2_max🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@jesterTOV/inference/likelihoods/nicer.py` around lines 217 - 244, Extract the duplicated M-R segment splitting into a single helper (e.g., split_mr_segments) placed at module level or in utils.py; the helper should take masses_EOS and radii_EOS, call utils.get_MR_split_index, build masks, produce sorted m_eos_1/r_eos_1 and m_eos_2/r_eos_2 plus seg1_min/seg1_max and seg2_min/seg2_max, and return them; then replace the duplicated blocks in NICERKDELikelihood.evaluate and MockMRLikelihood.evaluate with a call to split_mr_segments(...) and use the returned variables (m_eos_1, r_eos_1, seg1_min, seg1_max, m_eos_2, r_eos_2, seg2_min, seg2_max).
627-642: Enhance class docstring with parameters and attributes documentation.The docstring is minimal compared to
NICERLikelihoodandNICERKDELikelihood. Consider adding Parameters and Attributes sections for consistency and developer experience.📝 Suggested docstring enhancement
class MockMRLikelihood(LikelihoodBase): """ Mock MR Likelihood evaluating deterministic skewed correlated posteriors. Integrates probability density over a mass grid and normalizes by K samples correctly in log-space. Used for testing and development purposes. Parameters ---------- csv_file : str Path to CSV file containing skewed Gaussian parameters. Required columns: Mass_Center_Noise, Radius_Center_Noise, Std_Mass, Std_Radius, Covariance, Skew_Mass, Skew_Radius. penalty_value : float, optional Penalty value for samples where mass exceeds Mtov (default: -99999.0) N_masses_evaluation : int, optional Number of mass grid points for numerical integration (default: 200) Attributes ---------- K : int Number of mock posterior samples loaded from CSV centers : Float[Array, "K 2"] Center positions (mass, radius) for each mock posterior covs : Float[Array, "K 2 2"] Covariance matrices for each mock posterior """🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@jesterTOV/inference/likelihoods/nicer.py` around lines 627 - 642, Add a detailed docstring to MockMRLikelihood mirroring NICERLikelihood/NICERKDELikelihood: expand the top-level description and add Parameters and Attributes sections that document constructor inputs and key fields (e.g., csv_file, penalty_value, N_masses_evaluation) and the class attributes K, centers, covs, inv_covs, log_det_covs, skews, omegas, alpha_primes; ensure the wording explains types and units/shape (e.g., centers: Float[Array, "K 2"], covs: Float[Array, "K 2 2"]) and default behavior for penalty_value and mass grid resolution so future readers can find usage in MockMRLikelihood and its base LikelihoodBase-derived behavior.
🤖 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/config/schemas/likelihoods.py`:
- Around line 749-757: Add an example config under examples/inference/... that
shows using MockMRLikelihood with the default fields penalty_value and
N_masses_evaluation (set to -99999.0 and 200 respectively) so users can see the
expected YAML structure and field names; then add a unit test in
tests/test_inference/test_config.py that imports MockMRLikelihoodConfig (or the
Pydantic model backing MockMRLikelihood) and asserts its default values for
penalty_value == -99999.0 and N_masses_evaluation == 200 to ensure defaults
remain synchronized with MockMRLikelihood.__init__ and the schema.
In `@jesterTOV/inference/likelihoods/nicer.py`:
- Around line 676-678: The computation of self.alpha_primes = self.skews /
self.omegas can divide by zero when any diagonal of self.covs is zero; in the
precompute skew modifiers block guard against zero omegas by either validating
inputs (in the class __init__ check that the source standard deviations/std_m or
std_r are > 0 and raise ValueError with a clear message) or replace zeros with a
small positive epsilon before division (e.g., compute a safe_omegas =
jnp.where(self.omegas <= 0, epsilon, self.omegas) and use that for
self.alpha_primes), ensuring you reference self.covs, self.omegas and self.skews
while making the change.
- Around line 680-683: The evaluate method's params annotation is too broad;
change the signature of evaluate in class containing evaluate (and any sibling
implementations like NICERLikelihood) from params: dict to params: dict[str,
Float | Array] to match the base class and sister classes; update the type
import/aliases if necessary so Float and Array are available in this module and
ensure consistency with the base-class signature used elsewhere.
In `@tests/test_scalar_tensor.py`:
- Around line 94-104: The parametrized test test_parameter_sensitivity passes
beta_ST=0.0 which causes division-by-zero in Jordan frame conversions; update
the parameter set in test_parameter_sensitivity to exclude 0.0 (e.g., use [-5.0,
-4.5] only) or add a guard to skip/xfail when beta_ST == 0.0 before calling
ScalarTensorTOVSolver.solve so the division-by-zero is avoided; make sure to
reference the params dict ("params = {'beta_ST': beta_ST, 'phi_inf_tgt': 1e-3,
'phi_c': 1.0}") and adjust the pytest.mark.parametrize values accordingly.
- Around line 66-83: The test test_gr_limit causes division-by-zero in the
ScalarTensorTOVSolver's Jordan-frame conversions when beta_ST == 0.0; fix the
test by setting params = {"beta_ST": 1e-10, "phi_inf_tgt": 1e-3, "phi_c": 1.0}
in test_gr_limit so the solver approximates the GR limit without triggering
division-by-zero in the Jordan-frame conversion logic (the conversions that
compute Lambda_S_J, Lambda_ST1_J, Lambda_ST2_J inside ScalarTensorTOVSolver).
---
Nitpick comments:
In `@jesterTOV/inference/likelihoods/nicer.py`:
- Around line 707-709: The mass grid range in nicer.NicerLikelihood (m_grid and
dm created in the evaluation code) is hardcoded to [0.1, 3.5]; make the lower
and upper bounds configurable by adding new __init__ parameters (e.g., mass_min
and mass_max) and storing them on self, defaulting to 0.1 and 3.5, then replace
the hardcoded literals in the block that constructs m_grid and dm with
self.mass_min and self.mass_max so different mock scenarios can pass alternate
ranges (also ensure any docs/tests that reference N_masses_evaluation still work
and consider using the new attributes in evaluate where appropriate).
- Around line 294-307: Update the unclear comment above the likelihood
combination to explicitly state that when there is no phase transition the
segment-2 terms (amsterdam_logprobs_2, maryland_logprobs_2) become -inf and
therefore contribute zero after exp/logsumexp, and that the final subtraction of
jnp.log(2.0) is normalizing across the two analysis groups (Amsterdam and
Maryland) rather than the four intermediate terms; reference the computed
intermediates logL_amsterdam_1, logL_maryland_1, logL_amsterdam_2,
logL_maryland_2 and the use of logsumexp so readers understand why seg2 can be
ignored and why -jnp.log(2.0) is applied.
- Around line 217-244: Extract the duplicated M-R segment splitting into a
single helper (e.g., split_mr_segments) placed at module level or in utils.py;
the helper should take masses_EOS and radii_EOS, call utils.get_MR_split_index,
build masks, produce sorted m_eos_1/r_eos_1 and m_eos_2/r_eos_2 plus
seg1_min/seg1_max and seg2_min/seg2_max, and return them; then replace the
duplicated blocks in NICERKDELikelihood.evaluate and MockMRLikelihood.evaluate
with a call to split_mr_segments(...) and use the returned variables (m_eos_1,
r_eos_1, seg1_min, seg1_max, m_eos_2, r_eos_2, seg2_min, seg2_max).
- Around line 627-642: Add a detailed docstring to MockMRLikelihood mirroring
NICERLikelihood/NICERKDELikelihood: expand the top-level description and add
Parameters and Attributes sections that document constructor inputs and key
fields (e.g., csv_file, penalty_value, N_masses_evaluation) and the class
attributes K, centers, covs, inv_covs, log_det_covs, skews, omegas,
alpha_primes; ensure the wording explains types and units/shape (e.g., centers:
Float[Array, "K 2"], covs: Float[Array, "K 2 2"]) and default behavior for
penalty_value and mass grid resolution so future readers can find usage in
MockMRLikelihood and its base LikelihoodBase-derived behavior.
In `@jesterTOV/utils.py`:
- Around line 632-641: Add a clear docstring and Python type hints to the
get_MR_split_index function: annotate parameters x and y as jnp.ndarray (or
Array) and jump_threshold as float, and the return type as int; in the docstring
describe purpose, parameters, return value, and behavior (how split index is
computed, meaning of returned N when no jump); document the numeric stabilizer
1e-12 and why it's added to median_ds2; keep existing logic (use jnp.diff,
jnp.argmax, jnp.median) and mention that split_idx is 1-based offset (max_idx +
1) when a jump is detected and N otherwise so callers understand semantics.
🪄 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: fac9f49a-4cfa-4567-8580-3856cf8ed87f
📒 Files selected for processing (13)
docs/examples/eos_tov/EoS_beyondGR.ipynbdocs/inference/yaml_reference.mdjesterTOV/inference/config/generate_yaml_reference.pyjesterTOV/inference/config/schema.pyjesterTOV/inference/config/schemas/likelihoods.pyjesterTOV/inference/config/schemas/tov.pyjesterTOV/inference/likelihoods/__init__.pyjesterTOV/inference/likelihoods/factory.pyjesterTOV/inference/likelihoods/nicer.pyjesterTOV/inference/transforms/transform.pyjesterTOV/tov/scalar_tensor.pyjesterTOV/utils.pytests/test_scalar_tensor.py
🚧 Files skipped from review as they are similar to previous changes (3)
- jesterTOV/inference/config/schema.py
- jesterTOV/inference/transforms/transform.py
- jesterTOV/inference/config/schemas/tov.py
Comment on lines
+749
to
+757
| penalty_value: float = Field( | ||
| default=-99999.0, | ||
| description="Log-likelihood penalty for mass exceeding M_TOV", | ||
| ) | ||
| N_masses_evaluation: int = Field( | ||
| default=200, | ||
| gt=0, | ||
| description="Number of mass points for numerical integration", | ||
| ) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Verify default values are consistent across locations
echo "=== Checking penalty_value defaults ==="
rg -n "penalty_value.*-99999" --type py
echo ""
echo "=== Checking N_masses_evaluation defaults (200) ==="
rg -n "N_masses_evaluation.*200" --type py
echo ""
echo "=== Checking MockMRLikelihood __init__ signature ==="
ast-grep --pattern 'def __init__(self, $$$, penalty_value=$_, $$$)'Repository: nuclear-multimessenger-astronomy/jester
Length of output: 1019
🏁 Script executed:
#!/bin/bash
# Check generate_yaml_reference.py for hardcoded defaults
echo "=== Checking generate_yaml_reference.py ==="
rg -n "penalty_value|N_masses_evaluation" jesterTOV/inference/config/generate_yaml_reference.py -A 2 -B 2
echo ""
echo "=== Checking yaml_reference.md for these parameters ==="
rg -n "penalty_value|N_masses_evaluation" docs/inference/yaml_reference.md -A 1 -B 1 | head -50
echo ""
echo "=== Checking example configs ==="
find jesterTOV/inference/examples -name "*.yaml" -o -name "*.yml" | head -5
fd "config\.ya?ml$" examples/inference/ --exec grep -l "penalty_value\|N_masses_evaluation" {} \;
echo ""
echo "=== Checking test_config.py ==="
rg -n "penalty_value|N_masses_evaluation" tests/test_inference/test_config.py -A 1 -B 1Repository: nuclear-multimessenger-astronomy/jester
Length of output: 10213
🏁 Script executed:
#!/bin/bash
# Check example configs for MockMRLikelihood usage
echo "=== Checking example configs for MockMRLikelihood ==="
grep -r "MockMRLikelihood\|mock_mr\|mock.*mr" examples/inference/ --include="*.yaml" -l 2>/dev/null
echo ""
echo "=== Checking if MockMRLikelihood is tested in test_config.py ==="
rg -n "MockMRLikelihood" tests/test_inference/test_config.py -A 10 -B 2
echo ""
echo "=== Searching all test files for MockMRLikelihood ==="
fd "test.*\.py$" tests/ --exec grep -l "MockMRLikelihood" {} \;
echo ""
echo "=== Checking full likelihoods.py for MockMRLikelihood class definition ==="
rg -n "class.*MockMR" jesterTOV/inference/config/schemas/likelihoods.py -A 20Repository: nuclear-multimessenger-astronomy/jester
Length of output: 1088
🏁 Script executed:
#!/bin/bash
# Check if MockMRLikelihood is used anywhere in the codebase
echo "=== All references to MockMRLikelihood ==="
rg "MockMRLikelihood" --type py -l
echo ""
echo "=== Check if MockMRLikelihood has corresponding implementation in likelihoods ==="
grep -n "class MockMRLikelihood" jesterTOV/inference/likelihoods/*.py
echo ""
echo "=== Check if mock_mr is referenced in examples ==="
rg "mock_mr\|MockMR" examples/ -l 2>/dev/null | head -10
echo ""
echo "=== Summary of what we found ==="
echo "MockMRLikelihoodConfig in schemas at:"
rg -n "class MockMRLikelihoodConfig" jesterTOV/inference/config/schemas/likelihoods.py
echo ""
echo "Default values in MockMRLikelihoodConfig docstring:"
rg -n "penalty_value.*-99999\|N_masses_evaluation.*200" jesterTOV/inference/config/schemas/likelihoods.py | grep -E "739|740"Repository: nuclear-multimessenger-astronomy/jester
Length of output: 741
🏁 Script executed:
#!/bin/bash
# Check MockMRLikelihood implementation in nicer.py
echo "=== MockMRLikelihood __init__ signature in nicer.py ==="
rg -n "class MockMRLikelihood" jesterTOV/inference/likelihoods/nicer.py -A 30 | head -50
echo ""
echo "=== Check if MockMRLikelihood defaults match ==="
sed -n '627,660p' jesterTOV/inference/likelihoods/nicer.py
echo ""
echo "=== Check MockMRLikelihoodConfig fields in schemas ==="
sed -n '725,760p' jesterTOV/inference/config/schemas/likelihoods.pyRepository: nuclear-multimessenger-astronomy/jester
Length of output: 3756
🏁 Script executed:
#!/bin/bash
# Final check: Are there tests for MockMRLikelihood anywhere?
echo "=== Searching all test files for mock_mr or MockMRLikelihood ==="
rg "mock_mr\|MockMRLikelihood" tests/ -l 2>/dev/null
echo ""
echo "=== Check factory.py to see if MockMRLikelihood is properly registered ==="
rg -n "MockMRLikelihood\|mock_mr" jesterTOV/inference/likelihoods/factory.py -B 2 -A 2
echo ""
echo "=== Check schema.py discriminated union includes MockMRLikelihoodConfig ==="
rg -n "MockMRLikelihoodConfig" jesterTOV/inference/config/schema.py -B 2 -A 2Repository: nuclear-multimessenger-astronomy/jester
Length of output: 636
Confirm default values are synchronized; add missing example config and test.
Default values for penalty_value (-99999.0) and N_masses_evaluation (200) are correctly synchronized across:
- ✅
config/schemas/likelihoods.py - ✅
likelihoods/nicer.py(MockMRLikelihood.init) - ✅
config/generate_yaml_reference.py - ✅
docs/inference/yaml_reference.md
However, per coding guidelines, two required locations are missing:
- ❌ No example config in
examples/inference/**/config.yamldemonstrates MockMRLikelihood usage - ❌ No test in
tests/test_inference/test_config.pyvalidates MockMRLikelihoodConfig defaults
Add these to complete the implementation.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@jesterTOV/inference/config/schemas/likelihoods.py` around lines 749 - 757,
Add an example config under examples/inference/... that shows using
MockMRLikelihood with the default fields penalty_value and N_masses_evaluation
(set to -99999.0 and 200 respectively) so users can see the expected YAML
structure and field names; then add a unit test in
tests/test_inference/test_config.py that imports MockMRLikelihoodConfig (or the
Pydantic model backing MockMRLikelihood) and asserts its default values for
penalty_value == -99999.0 and N_masses_evaluation == 200 to ensure defaults
remain synchronized with MockMRLikelihood.__init__ and the schema.
Comment on lines
+676
to
+678
| # Precompute skew modifiers | ||
| self.omegas = jnp.sqrt(jnp.diagonal(self.covs, axis1=1, axis2=2)) | ||
| self.alpha_primes = self.skews / self.omegas |
There was a problem hiding this comment.
Potential division by zero when computing alpha_primes.
If any row in the CSV has Std_Mass or Std_Radius equal to zero, self.omegas will contain zeros, causing division by zero in self.alpha_primes = self.skews / self.omegas.
🛡️ Proposed defensive fix
# Precompute skew modifiers
self.omegas = jnp.sqrt(jnp.diagonal(self.covs, axis1=1, axis2=2))
- self.alpha_primes = self.skews / self.omegas
+ # Add small epsilon to avoid division by zero for degenerate covariances
+ self.alpha_primes = self.skews / (self.omegas + 1e-12)Alternatively, validate input data in __init__:
if np.any(std_m <= 0) or np.any(std_r <= 0):
raise ValueError("Standard deviations must be positive")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@jesterTOV/inference/likelihoods/nicer.py` around lines 676 - 678, The
computation of self.alpha_primes = self.skews / self.omegas can divide by zero
when any diagonal of self.covs is zero; in the precompute skew modifiers block
guard against zero omegas by either validating inputs (in the class __init__
check that the source standard deviations/std_m or std_r are > 0 and raise
ValueError with a clear message) or replace zeros with a small positive epsilon
before division (e.g., compute a safe_omegas = jnp.where(self.omegas <= 0,
epsilon, self.omegas) and use that for self.alpha_primes), ensuring you
reference self.covs, self.omegas and self.skews while making the change.
Comment on lines
+94
to
+104
| @pytest.mark.parametrize("beta_ST", [-5.0, -4.5, 0.0]) | ||
| def test_parameter_sensitivity(self, sample_eos_data, beta_ST): | ||
| """Test sensitivity to beta_ST parameter.""" | ||
| pc = float(sample_eos_data.ps[25]) | ||
| solver = ScalarTensorTOVSolver() | ||
|
|
||
| params = {"beta_ST": beta_ST, "phi_inf_tgt": 1e-3, "phi_c": 1.0} | ||
|
|
||
| solution = solver.solve(sample_eos_data, pc, tov_params=params) | ||
| assert jnp.isfinite(solution.M) | ||
| assert solution.M > 0 |
There was a problem hiding this comment.
Parametrized test includes beta_ST=0.0 which has division-by-zero issues.
Similar to the GR limit test, the parametrized test_parameter_sensitivity includes beta_ST=0.0 in the test cases. This will trigger the same division-by-zero issues in the Jordan frame conversions.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/test_scalar_tensor.py` around lines 94 - 104, The parametrized test
test_parameter_sensitivity passes beta_ST=0.0 which causes division-by-zero in
Jordan frame conversions; update the parameter set in test_parameter_sensitivity
to exclude 0.0 (e.g., use [-5.0, -4.5] only) or add a guard to skip/xfail when
beta_ST == 0.0 before calling ScalarTensorTOVSolver.solve so the
division-by-zero is avoided; make sure to reference the params dict ("params =
{'beta_ST': beta_ST, 'phi_inf_tgt': 1e-3, 'phi_c': 1.0}") and adjust the
pytest.mark.parametrize values accordingly.
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (3)
jesterTOV/utils.py (2)
354-364:⚠️ Potential issue | 🟠 MajorDon't silently cap the stable branch count at four.
jnp.where(..., size=4)plus the manuals1-s4unpacking drops every stable segment after the fourth. Any EOS with more than four disjoint stable regions will therefore lose physically valid M-R-Λ data with no signal to the caller.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@jesterTOV/utils.py` around lines 354 - 364, The current code silently limits stable-segment detection to four by calling jnp.where(..., size=4) and unpacking into s1_start..s4_end, which drops any additional segments; change the logic in the segment-extraction area (the diff, m, starts, ends and s1_start..s4_start unpack pattern) to handle an arbitrary number of start/end indices: compute starts = jnp.where(diff == 1)[0] and ends = jnp.where(diff == -1)[0] (or equivalent jax-safe operations that do not truncate), pair them into start/end segment tuples dynamically (ensuring equal counts or handling edge cases by padding a single dummy index only when necessary), and replace the fixed s1..s4 unpacking with processing over the resulting list of segments so no stable region is silently dropped.
473-487:⚠️ Potential issue | 🟠 MajorMake the Hermite slopes respect retained-segment boundaries.
These centered differences always use
xp[i-1]andxp[i+1]. Afterlimit_by_MTOV_and_interpolate()stitches disjoint stable branches, the first and last point of each kept branch therefore pick up slope information from the removed unstable gap, which can overshoot at branch entrances and exits.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@jesterTOV/utils.py` around lines 473 - 487, The Hermite slope computation (central_diff_numer/interior_m1_m, forward_m0, backward_m_end assembled into dydx) incorrectly uses neighbors across removed gaps created by limit_by_MTOV_and_interpolate(); change it to compute centered differences only within each retained contiguous segment and compute slopes at segment boundaries using one-sided differences (or copy adjacent one-sided slope) so no derivative uses xp from a different stitched segment; identify segments from the output of limit_by_MTOV_and_interpolate() (or by detecting non-consecutive xp indices), iterate per segment to compute interior central differences and set the first/last slope using forward/backward difference respectively, then concatenate per-segment dydx into the final dydx array.jesterTOV/tov/scalar_tensor.py (1)
577-601:⚠️ Potential issue | 🔴 CriticalGuard the Jordan-frame scalar/mixed lambda conversions against the GR-limit path.
Line 581, Line 589, and Line 597 divide by
beta_STandphi_inf_target, and Line 629 still makesbeta_ST=0.0the default. That makes the newextraoutputs reachable asinf/nanon the default path, andjesterTOV/tov/base.pynow propagates those values intoFamilyData.extraas well. Please branch before these conversions—either reject exact GR-limit inputs here or route them to the GR solver.Also applies to: 629-631
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@jesterTOV/tov/scalar_tensor.py` around lines 577 - 601, The Jordan-frame conversions for Lambda_S_J, Lambda_ST1_J, and Lambda_ST2_J currently divide by beta_ST and phi_inf_target (symbols: Lambda_S_J, Lambda_ST1_J, Lambda_ST2_J, beta_ST, phi_inf_target) and the code also leaves beta_ST defaulting to 0.0; guard these conversions by checking for the GR-limit (beta_ST == 0.0 or phi_inf_target == 0.0) before performing the math: if in the GR-limit, either raise/return an explicit error or short-circuit to the GR solver path instead of computing the divisions, otherwise compute the existing formulas using lambda_T, lambda_S, lambda_ST1, lambda_ST2 and M_inf; also remove or change the default beta_ST = 0.0 initialization so the GR-limit is handled explicitly rather than producing inf/nan in FamilyData.extra.
🧹 Nitpick comments (1)
tests/test_scalar_tensor.py (1)
9-10: Mark this suite as integration-level.These tests drive the real JAX/Diffrax solver and family-construction path, so they belong in the integration bucket rather than the default unit-test bucket. Please add the repo's
integrationmarker here, and addslowselectively if runtime is nontrivial.Example marker placement
+@pytest.mark.integration class TestScalarTensorTOVSolver:As per coding guidelines, "Test markers: use
@pytest.mark.slowfor slow tests (skipped in regular CI),@pytest.mark.integrationfor integration tests,@pytest.mark.e2efor end-to-end tests,@pytest.mark.blackjaxfor BlackJAX-dependent tests."🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/test_scalar_tensor.py` around lines 9 - 10, This test class TestScalarTensorTOVSolver should be marked as integration (and slow if runtime is nontrivial): add the pytest marker decorator(s) `@pytest.mark.integration` (and optionally `@pytest.mark.slow`) immediately above the class definition for TestScalarTensorTOVSolver so the suite is run under the integration test bucket (and skipped in regular CI if marked slow).
🤖 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/config/schemas/tov.py`:
- Around line 76-82: Scalar-tensor fields (beta_ST, phi_inf_tgt, phi_c) are
declared on ScalarTensorTOVConfig but never merged into tov_params, so the
solver falls back to hardcoded defaults; either remove these fields and document
them as required prior parameters, or implement merging of the config defaults
into params (the same way fixed_params are merged) so _create_tov_solver /
transform.py (tov_params = self.tov_solver.fetch_params(params)) passes them
through and solve() uses the config defaults rather than .get() fallbacks;
additionally add a validation in run_inference.py to assert the prior supplies
beta_ST, phi_inf_tgt, and phi_c (or to allow them to be provided via config if
you keep the fields) to prevent silent defaults.
---
Duplicate comments:
In `@jesterTOV/tov/scalar_tensor.py`:
- Around line 577-601: The Jordan-frame conversions for Lambda_S_J,
Lambda_ST1_J, and Lambda_ST2_J currently divide by beta_ST and phi_inf_target
(symbols: Lambda_S_J, Lambda_ST1_J, Lambda_ST2_J, beta_ST, phi_inf_target) and
the code also leaves beta_ST defaulting to 0.0; guard these conversions by
checking for the GR-limit (beta_ST == 0.0 or phi_inf_target == 0.0) before
performing the math: if in the GR-limit, either raise/return an explicit error
or short-circuit to the GR solver path instead of computing the divisions,
otherwise compute the existing formulas using lambda_T, lambda_S, lambda_ST1,
lambda_ST2 and M_inf; also remove or change the default beta_ST = 0.0
initialization so the GR-limit is handled explicitly rather than producing
inf/nan in FamilyData.extra.
In `@jesterTOV/utils.py`:
- Around line 354-364: The current code silently limits stable-segment detection
to four by calling jnp.where(..., size=4) and unpacking into s1_start..s4_end,
which drops any additional segments; change the logic in the segment-extraction
area (the diff, m, starts, ends and s1_start..s4_start unpack pattern) to handle
an arbitrary number of start/end indices: compute starts = jnp.where(diff ==
1)[0] and ends = jnp.where(diff == -1)[0] (or equivalent jax-safe operations
that do not truncate), pair them into start/end segment tuples dynamically
(ensuring equal counts or handling edge cases by padding a single dummy index
only when necessary), and replace the fixed s1..s4 unpacking with processing
over the resulting list of segments so no stable region is silently dropped.
- Around line 473-487: The Hermite slope computation
(central_diff_numer/interior_m1_m, forward_m0, backward_m_end assembled into
dydx) incorrectly uses neighbors across removed gaps created by
limit_by_MTOV_and_interpolate(); change it to compute centered differences only
within each retained contiguous segment and compute slopes at segment boundaries
using one-sided differences (or copy adjacent one-sided slope) so no derivative
uses xp from a different stitched segment; identify segments from the output of
limit_by_MTOV_and_interpolate() (or by detecting non-consecutive xp indices),
iterate per segment to compute interior central differences and set the
first/last slope using forward/backward difference respectively, then
concatenate per-segment dydx into the final dydx array.
---
Nitpick comments:
In `@tests/test_scalar_tensor.py`:
- Around line 9-10: This test class TestScalarTensorTOVSolver should be marked
as integration (and slow if runtime is nontrivial): add the pytest marker
decorator(s) `@pytest.mark.integration` (and optionally `@pytest.mark.slow`)
immediately above the class definition for TestScalarTensorTOVSolver so the
suite is run under the integration test bucket (and skipped in regular CI if
marked slow).
🪄 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: e0d18fb1-97be-4285-8d34-eb2a3bc87424
📒 Files selected for processing (8)
docs/examples/eos_tov/EoS_beyondGR.ipynbdocs/examples/eos_tov/generate_mock.ipynbjesterTOV/inference/config/schema.pyjesterTOV/inference/config/schemas/tov.pyjesterTOV/inference/transforms/transform.pyjesterTOV/tov/scalar_tensor.pyjesterTOV/utils.pytests/test_scalar_tensor.py
✅ Files skipped from review due to trivial changes (1)
- jesterTOV/inference/config/schema.py
🚧 Files skipped from review as they are similar to previous changes (1)
- jesterTOV/inference/transforms/transform.py
Comment on lines
+76
to
+82
| beta_ST: float = Field( | ||
| default=0.0, description="Scalar-tensor coupling parameter beta_ST" | ||
| ) | ||
| phi_inf_tgt: float = Field( | ||
| default=1e-3, description="Target asymptotic scalar field at infinity" | ||
| ) | ||
| phi_c: float = Field(default=1.0, description="Central scalar field value") |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "### fetch_params definitions"
rg -n -A20 -B5 'def fetch_params\(' jesterTOV
echo
echo "### transform-side TOV parameter assembly"
sed -n '345,370p' jesterTOV/inference/transforms/transform.py
echo
echo "### scalar-tensor config schema"
sed -n '59,105p' jesterTOV/inference/config/schemas/tov.py
echo
echo "### scalar-tensor solver parameter consumption"
sed -n '626,672p' jesterTOV/tov/scalar_tensor.py
echo
echo "### inference-side references to the scalar-tensor parameter names"
rg -n -C3 '\b(beta_ST|phi_inf_tgt|phi_c)\b' jesterTOV/inferenceRepository: nuclear-multimessenger-astronomy/jester
Length of output: 8364
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "### ScalarTensorTOVSolver.get_required_parameters() implementation"
rg -n -A10 'class ScalarTensorTOVSolver' jesterTOV/tov/scalar_tensor.py | head -50
echo
echo "### Full get_required_parameters for ScalarTensorTOVSolver"
sed -n '673,690p' jesterTOV/tov/scalar_tensor.py
echo
echo "### How run_inference.py uses the config and passes params to the solver"
rg -n -B5 -A15 'self.tov_solver' jesterTOV/inference/run_inference.py | head -80
echo
echo "### Config instantiation and schema usage in transforms"
rg -n -B3 -A10 'TOVConfig\|ScalarTensorTOVConfig' jesterTOV/inference/transforms/transform.py | head -60
echo
echo "### Parameter prior handling and merging"
rg -n -B5 -A10 'def setup_transforms' jesterTOV/inference/run_inference.py | head -60Repository: nuclear-multimessenger-astronomy/jester
Length of output: 816
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "### Check if TOV solver is instantiated from config in run_inference"
rg -n 'tov_solver\|TOVSolver' jesterTOV/inference/run_inference.py | head -40
echo
echo "### Look for where solver config values would be merged into params"
rg -n -B5 -A10 'class.*Transform' jesterTOV/inference/transforms/transform.py | head -80
echo
echo "### Check setup_transforms or __init__ in Transform"
sed -n '1,150p' jesterTOV/inference/transforms/transform.py
echo
echo "### Check how params dict is built from prior"
rg -n -B3 -A8 'params.*prior\|prior.*params' jesterTOV/inference/run_inference.pyRepository: nuclear-multimessenger-astronomy/jester
Length of output: 151
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "### Examine run_inference.py structure"
wc -l jesterTOV/inference/run_inference.py
head -100 jesterTOV/inference/run_inference.py
echo
echo "### Look for Transform instantiation"
grep -n 'Transform' jesterTOV/inference/run_inference.py
echo
echo "### Examine Transform __init__"
sed -n '1,250p' jesterTOV/inference/transforms/transform.pyRepository: nuclear-multimessenger-astronomy/jester
Length of output: 12999
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "### _create_tov_solver implementation"
sed -n '250,450p' jesterTOV/inference/transforms/transform.py
echo
echo "### forward method and how params are used"
sed -n '300,380p' jesterTOV/inference/transforms/transform.py
echo
echo "### Check if config values are extracted and merged into params anywhere"
rg -n -B3 -A8 'tov_config\|ScalarTensorTOVConfig' jesterTOV/inference/transforms/transform.pyRepository: nuclear-multimessenger-astronomy/jester
Length of output: 9984
🏁 Script executed:
#!/bin/bash
set -euo pipefixo pipefail
echo "### Check how fixed_params flow through to forward()"
sed -n '180,250p' jesterTOV/inference/base.py
echo
echo "### Check if tov_config values are ever used after _create_tov_solver"
rg -n 'tov_config\|tov_solver' jesterTOV/inference/transforms/transform.py | tail -40
echo
echo "### Verify the gap: config passed but not merged"
sed -n '235,260p' jesterTOV/inference/transforms/transform.pyRepository: nuclear-multimessenger-astronomy/jester
Length of output: 136
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "### Check how fixed_params flow through to forward()"
sed -n '180,250p' jesterTOV/inference/base.py
echo
echo "### Check if tov_config values are ever used after _create_tov_solver"
rg -n 'tov_config\|tov_solver' jesterTOV/inference/transforms/transform.py | tail -40
echo
echo "### Verify the gap: config passed but not merged"
sed -n '235,260p' jesterTOV/inference/transforms/transform.pyRepository: nuclear-multimessenger-astronomy/jester
Length of output: 223
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "### Find base.py location"
fd 'base\.py' jesterTOV/inference
echo
echo "### Check from_config method in transform.py more carefully"
sed -n '200,260p' jesterTOV/inference/transforms/transform.py
echo
echo "### Critical: After _create_tov_solver returns, are tov_config values used?"
sed -n '219,235p' jesterTOV/inference/transforms/transform.py
echo
echo "### Look for where params dict gets augmented with config defaults"
rg -n 'fixed_params\|tov_config' jesterTOV/inference/transforms/transform.pyRepository: nuclear-multimessenger-astronomy/jester
Length of output: 2857
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "### Trace: _create_tov_solver receives tov_config but what happens to its field values?"
sed -n '264,290p' jesterTOV/inference/transforms/transform.py
echo
echo "### The critical question: are tov_config field values EVER stored or merged anywhere?"
grep -n 'beta_ST\|phi_inf_tgt\|phi_c' jesterTOV/inference/transforms/transform.py
echo
echo "### How does run_inference.py call from_config?"
rg -n -B5 -A10 'from_config' jesterTOV/inference/run_inference.py | head -60
echo
echo "### What does the prior file provide? Are beta_ST etc. expected to be sampled?"
rg -n 'beta_ST\|phi_inf_tgt\|phi_c' jesterTOV/inferenceRepository: nuclear-multimessenger-astronomy/jester
Length of output: 1037
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "### Check AnisotropyTOVConfig to see the pattern"
sed -n '85,110p' jesterTOV/inference/config/schemas/tov.py
echo
echo "### How are AnisotropyTOVSolver params handled in _create_tov_solver?"
sed -n '264,295p' jesterTOV/inference/transforms/transform.py
echo
echo "### CRITICAL: Check if config solver-specific params should be passed to solver __init__"
rg -n 'class.*TOVSolver' jesterTOV/tov/
echo
echo "### Check AnisotropyTOVSolver __init__ for reference"
rg -n -A15 'class AnisotropyTOVSolver' jesterTOV/tov/anisotropy.py | head -40
echo
echo "### Check ScalarTensorTOVSolver __init__"
rg -n -A15 'class ScalarTensorTOVSolver' jesterTOV/tov/scalar_tensor.py | head -40Repository: nuclear-multimessenger-astronomy/jester
Length of output: 3707
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "### Final verification: Does AnisotropyTOVSolver.__init__ exist and take no params?"
rg -n -A20 'class AnisotropyTOVSolver' jesterTOV/tov/anisotropy.py | head -30
echo
echo "### Does ScalarTensorTOVSolver have __init__?"
sed -n '614,680p' jesterTOV/tov/scalar_tensor.py | grep -A10 '__init__\|def solve'
echo
echo "### Full AnisotropyTOVConfig docstring - note it says params come from PRIOR"
sed -n '85,102p' jesterTOV/inference/config/schemas/tov.py
echo
echo "### Compare: ScalarTensorTOVConfig docstring says fields are here with defaults"
sed -n '59,83p' jesterTOV/inference/config/schemas/tov.pyRepository: nuclear-multimessenger-astronomy/jester
Length of output: 3047
These scalar-tensor config fields are defined but not wired into tov_params—they will silently default via .get() rather than enforcing their schema values.
ScalarTensorTOVConfig defines beta_ST, phi_inf_tgt, and phi_c with Pydantic defaults, but _create_tov_solver() only uses the config to select the solver class; it never extracts or passes these field values to the solver or merges them into params.
In transform.py, tov_params = self.tov_solver.fetch_params(params) extracts these keys from the sampled prior dict. The solver then does tov_params.get("beta_ST", 0.0) etc., which silently falls back to the hardcoded defaults in solve() if the prior is missing these parameters—not the config defaults.
This differs from AnisotropyTOVConfig, whose docstring correctly states: "The six theory parameters are specified in the prior file, not here." The scalar-tensor config should either:
- Remove the field definitions and document these as required prior parameters, or
- Merge config defaults into
paramswhen keys are missing (similar tofixed_paramshandling)
Either way, add validation in run_inference.py (per prior validation guidelines) to ensure these parameters are present in the prior file.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@jesterTOV/inference/config/schemas/tov.py` around lines 76 - 82,
Scalar-tensor fields (beta_ST, phi_inf_tgt, phi_c) are declared on
ScalarTensorTOVConfig but never merged into tov_params, so the solver falls back
to hardcoded defaults; either remove these fields and document them as required
prior parameters, or implement merging of the config defaults into params (the
same way fixed_params are merged) so _create_tov_solver / transform.py
(tov_params = self.tov_solver.fetch_params(params)) passes them through and
solve() uses the config defaults rather than .get() fallbacks; additionally add
a validation in run_inference.py to assert the prior supplies beta_ST,
phi_inf_tgt, and phi_c (or to allow them to be provided via config if you keep
the fields) to prevent silent defaults.
…ius curve. 2. Add config toggle to turn off tidal deformability calculation in scalar tensor, to save resource when not needed.
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
2 participants
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.
This PR is continuation of PR #103 which was accidentally closed automatically due to history mismatch.
Some concerns pointed out by reviewer is resolved in this PR, except in:
Changes:
docs/examples/eos_tov/EoS_beyondGR.ipynbnanto ensure physically invalid EoS not passed into solver and introduce bias. Moved from TOV solver to ensure these filters done beforevmapoperation.ndat) for each M-R-Lambda curve.interpax.Summary by CodeRabbit
New Features
Documentation
Tests
Chores