Scalar-Tensor theory update draft PR

.gitignore

@@ -152,4 +152,15 @@ dmypy.json
 .pytype/
 
 # Cython debug symbols
-cython_debug/
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+lalsuite/
+
+examples/inference/*/*/outdir/

.pre-commit-config.yaml

@@ -3,18 +3,10 @@ repos:
     rev: 24.1.1
     hooks:
     - id: black
+      exclude: \.ipynb$
   - repo: https://github.com/astral-sh/ruff-pre-commit
     rev: 'v0.1.14'
     hooks:
       - id: ruff
         args: ["--fix"]
-  - repo: https://github.com/nbQA-dev/nbQA
-    rev: 1.7.1
-    hooks:
-    - id: nbqa-black
-      exclude: ^docs/examples/
-      additional_dependencies: [ipython==8.12, black]
-    - id: nbqa-ruff 
-      exclude: ^docs/examples/
-      args: ["--fix", "--ignore=I001"]
-      additional_dependencies: [ipython==8.12, ruff]
+        exclude: \.ipynb$

docs/api/_autosummary/jesterTOV.tov.data_classes.FamilyData.rst

@@ -25,6 +25,7 @@
 
    .. autosummary::
 
+      ~FamilyData.extra
       ~FamilyData.lambdas
       ~FamilyData.log10pcs
       ~FamilyData.masses

docs/api/_autosummary/jesterTOV.tov.data_classes.TOVSolution.rst

@@ -27,6 +27,7 @@
 
       ~TOVSolution.M
       ~TOVSolution.R
+      ~TOVSolution.extra
       ~TOVSolution.k2
 
 

docs/examples/eos_tov/automatic_differentiation.ipynb

@@ -86,13 +86,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "gr_tov_solver = GRTOVSolver()\n",
-    "construct_family_lambda = lambda eos_data: gr_tov_solver.construct_family(eos_data, ndat=100, min_nsat=1.0)"
-   ]
+   "source": "gr_tov_solver = GRTOVSolver()\nconstruct_family_lambda = lambda eos_data: gr_tov_solver.construct_family(eos_data, tov_params={}, ndat=100, min_nsat=1.0)\n"
   },
   {
    "cell_type": "markdown",
@@ -814,4 +811,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}

docs/examples/eos_tov/eos_STtov.ipynb

@@ -641,7 +641,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "jester",
+   "display_name": "jestertov",
    "language": "python",
    "name": "python3"
   },

docs/examples/eos_tov/eos_tov.ipynb

@@ -182,13 +182,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "gr_tov_solver = GRTOVSolver()\n",
-    "sol: FamilyData = gr_tov_solver.construct_family(eos_data=eos_data, ndat = 200, min_nsat = 1.0)"
-   ]
+   "source": "gr_tov_solver = GRTOVSolver()\nsol: FamilyData = gr_tov_solver.construct_family(eos_data=eos_data, tov_params={}, ndat = 200, min_nsat = 1.0)\n"
   },
   {
    "cell_type": "markdown",
@@ -273,4 +270,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}

docs/examples/index.rst

@@ -18,9 +18,15 @@ Basic Examples
 
    :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`
+
 .. toctree::
    :hidden:
 
    eos_tov/eos_tov
    eos_tov/automatic_differentiation
    eos_tov/eos_STtov
+   eos_tov/EoS_beyondGR

examples/inference/st_smc_random_walk/chiEFT/config.yaml

@@ -0,0 +1,37 @@
+seed: 44
+dry_run: false
+validate_only: false
+eos:
+  type: metamodel_cse
+  ndat_metamodel: 100
+  nmax_nsat: 25.0
+  nb_CSE: 8
+  crust_name: DH
+tov:
+  type: scalar_tensor
+  min_nsat_TOV: 0.75
+  ndat_TOV: 100
+  nb_masses: 100
+prior:
+  specification_file: prior.prior
+likelihoods:
+- type: constraints_eos
+  enabled: true
+- 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
+sampler:
+  type: smc-rw
+  n_particles: 2000
+  n_mcmc_steps: 10
+  target_ess: 0.9
+  random_walk_sigma: 0.1
+  output_dir: ./outdir/
+postprocessing:
+  enabled: true
+  make_cornerplot: false
+  make_massradius: true
+  make_pressuredensity: true
+debug_nans: true

examples/inference/st_smc_random_walk/chiEFT/prior.prior

@@ -0,0 +1,13 @@
+E_sat = UniformPrior(-16.1, -15.9, parameter_names=["E_sat"])
+K_sat = UniformPrior(150.0, 300.0, parameter_names=["K_sat"])
+Q_sat = UniformPrior(-500.0, 1100.0, parameter_names=["Q_sat"])
+Z_sat = UniformPrior(-2500.0, 1500.0, parameter_names=["Z_sat"])
+E_sym = UniformPrior(28.0, 45.0, parameter_names=["E_sym"])
+L_sym = UniformPrior(10.0, 200.0, parameter_names=["L_sym"])
+K_sym = UniformPrior(-400.0, 200.0, parameter_names=["K_sym"])
+Q_sym = UniformPrior(-1000.0, 1500.0, parameter_names=["Q_sym"])
+Z_sym = UniformPrior(-2000.0, 1500.0, parameter_names=["Z_sym"])
+nbreak = UniformPrior(0.16, 0.32, parameter_names=["nbreak"])
+beta_ST = UniformPrior(-6, 0, parameter_names=["beta_ST"])
+phi_c = UniformPrior(0.999, 1, parameter_names=["phi_c"])
+phi_inf_tgt = UniformPrior(0.999e-3, 1e-3, parameter_names=["phi_inf_tgt"])

examples/inference/st_smc_random_walk/chiEFT/submit.sh

@@ -0,0 +1,28 @@
+#!/bin/bash -l
+#Set job requirements
+#SBATCH -N 1
+#SBATCH -n 1
+#SBATCH -p gpu_h100
+#SBATCH -t 10:00:00
+#SBATCH --gpus-per-node=1
+#SBATCH --cpus-per-gpu=1
+#SBATCH --mem-per-gpu=40G
+#SBATCH --output="./outdir/log_smc_rw.out"
+#SBATCH --job-name="SMC-RW-chiEFT"
+
+now=$(date)
+echo "$now"
+
+# Loading modules
+# module load 2024
+# module load Python/3.10.4-GCCcore-11.3.0
+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

jesterTOV/inference/config/schema.py

@@ -39,6 +39,7 @@
 from .schemas.tov import (
     BaseTOVConfig,
     GRTOVConfig,
+    ScalarTensorTOVConfig,
     AnisotropyTOVConfig,
     TOVConfig,
 )
@@ -233,6 +234,7 @@ def validate_seed(cls, v: int) -> int:
     # TOV
     "BaseTOVConfig",
     "GRTOVConfig",
+    "ScalarTensorTOVConfig",
     "AnisotropyTOVConfig",
     "TOVConfig",
     # Likelihoods

jesterTOV/inference/config/schemas/tov.py

@@ -56,6 +56,32 @@ class GRTOVConfig(BaseTOVConfig):
     type: Literal["gr"] = "gr"  # type: ignore[override]  # Literal["gr"] ⊂ str
 
 
+class ScalarTensorTOVConfig(BaseTOVConfig):
+    """Configuration for the Scalar-Tensor TOV solver.
+
+    Attributes
+    ----------
+    type : Literal["scalar_tensor"]
+        TOV solver type identifier
+    beta_ST : float
+        Scalar-tensor coupling parameter beta_ST (default: 0.0)
+    phi_inf_tgt : float
+        Target asymptotic scalar field at infinity (default: 1e-3)
+    phi_c : float
+        Central scalar field value (default: 1.0)
+    """
+
+    type: Literal["scalar_tensor"] = "scalar_tensor"  # type: ignore[override]
+
+    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")
+
+
 class AnisotropyTOVConfig(BaseTOVConfig):
     """Configuration for the post-TOV solver with beyond-GR corrections.
 
@@ -72,7 +98,8 @@ class AnisotropyTOVConfig(BaseTOVConfig):
     type: Literal["anisotropy"] = "anisotropy"  # type: ignore[override]  # Literal["anisotropy"] ⊂ str
 
 
+# TOVConfig is a discriminated union of all available TOV solver configs.
 TOVConfig = Annotated[
-    Union[GRTOVConfig, AnisotropyTOVConfig],
+    Union[GRTOVConfig, AnisotropyTOVConfig, ScalarTensorTOVConfig],
     Discriminator("type"),
 ]

jesterTOV/inference/transforms/transform.py

@@ -21,6 +21,7 @@
     SpectralEOSConfig,
     BaseTOVConfig,
     GRTOVConfig,
+    ScalarTensorTOVConfig,
     AnisotropyTOVConfig,
 )
 from jesterTOV.inference.likelihoods.constraints import check_all_constraints
@@ -283,12 +284,14 @@ def _create_tov_solver(config: BaseTOVConfig) -> TOVSolverBase:
         """
         if isinstance(config, GRTOVConfig):
             return GRTOVSolver()
+        elif isinstance(config, ScalarTensorTOVConfig):
+            from jesterTOV.tov.scalar_tensor import ScalarTensorTOVSolver
 
+            return ScalarTensorTOVSolver()
         elif isinstance(config, AnisotropyTOVConfig):
             from jesterTOV.tov.anisotropy import AnisotropyTOVSolver
 
             return AnisotropyTOVSolver()
-
         else:
             raise ValueError(f"Unknown TOV solver type: {type(config).__name__}")
 

jesterTOV/tov/base.py

@@ -6,6 +6,7 @@
 """
 
 from abc import ABC, abstractmethod
+from typing import Optional
 import jax
 import jax.numpy as jnp
 from jaxtyping import Float, Array
@@ -103,6 +104,20 @@ def construct_family(
         Returns:
             FamilyData: Mass-radius-tidal curves in physical units
         """
+        # Evaluate EOS validity once for the entire family
+        is_valid_eos = (
+            jnp.all(eos_data.ps > 0)
+            & jnp.all(eos_data.hs > 0)
+            & jnp.all(eos_data.es > 0)
+            & jnp.all(eos_data.cs2 >= 0)
+            & jnp.all(eos_data.cs2 <= 1)
+        )
+
+        # # Poison the EOS data with NaNs if unphysical
+        # eos_data = jax.tree_util.tree_map(
+        #     lambda x: jnp.where(is_valid_eos, x, jnp.nan), eos_data
+        # )
+
         # Create central pressure grid
         pc_min = self._get_pc_min(eos_data, min_nsat)
         pc_max = self._get_pc_max(eos_data)
@@ -120,8 +135,13 @@ def solve_single_pc(pc):
         radii: Float[Array, "ndat"] = solutions.R  # type: ignore[assignment]
         k2s: Float[Array, "ndat"] = solutions.k2  # type: ignore[assignment]
 
+        # Extract optional extra fields from solver-specific quantities
+        extra_fields: Optional[dict[str, Float[Array, "ndat"]]] = solutions.extra  # type: ignore[assignment]
+
         # Convert to physical units and compute tidal deformability
-        return self._create_family_data(pcs, masses, radii, k2s, ndat)
+        return self._create_family_data(
+            pcs, masses, radii, k2s, ndat, extra=extra_fields
+        )
 
     def _get_pc_min(self, eos_data: EOSData, min_nsat: float) -> Float[Array, ""]:
         """
@@ -170,6 +190,7 @@ def _create_family_data(
         radii: Float[Array, "ndat"],
         k2s: Float[Array, "ndat"],
         ndat: int,
+        extra: Optional[dict[str, Float[Array, "ndat"]]] = None,
     ) -> FamilyData:
         """
         Shared post-processing: unit conversion, compactness limits, interpolation.
@@ -180,6 +201,7 @@ def _create_family_data(
             radii: Radii [geometric units]
             k2s: Love numbers [dimensionless]
             ndat: Number of points for output grid
+            extra: Optional dictionary of solver-specific quantities (e.g., from ScalarTensorTOVSolver)
 
         Returns:
             FamilyData: Processed family curves in physical units
@@ -194,21 +216,30 @@ def _create_family_data(
         # Calculate tidal deformability
         lambdas = 2.0 / 3.0 * k2s * jnp.power(compactness, -5.0)
 
-        # Limit masses to be below MTOV (removes unstable branch)
-        pcs_lim, masses_lim, radii_lim, lambdas_lim = utils.limit_by_MTOV(
-            pcs, masses_solar, radii_km, lambdas
+        # Limit masses to be below MTOV and interpolate onto a uniform log(pc) grid on stable segments
+        pcs_lim, masses_lim, radii_lim, lambdas_lim = (
+            utils.limit_by_MTOV_and_interpolate(
+                pcs, masses_solar, radii_km, lambdas, ndat
+            )
         )
-
-        # Get a mass grid and interpolate, since we might have some duplicate points
-        mass_grid = jnp.linspace(jnp.min(masses_lim), jnp.max(masses_lim), ndat)
-        radii_interp = jnp.interp(mass_grid, masses_lim, radii_lim)
-        lambdas_interp = jnp.interp(mass_grid, masses_lim, lambdas_lim)
-        pcs_interp = jnp.interp(mass_grid, masses_lim, pcs_lim)
-        log10pcs = jnp.log10(pcs_interp)
+        # @TODO: merge extras interpolation in a single call in above.
+        # Process extra solver-specific fields if provided
+        extra_processed: Optional[dict[str, Float[Array, "ndat"]]] = None
+        if extra is not None:
+            # Use tree_map to efficiently apply interpolation to all extra fields at once
+            extra_processed = jax.tree_util.tree_map(
+                lambda val: utils.limit_by_MTOV_and_interpolate(
+                    pcs, masses_solar, radii_km, val, ndat
+                )[
+                    3
+                ],  # Take only the interpolated fourth element (the 'lambda'-like array)
+                extra,
+            )
 
         return FamilyData(
-            log10pcs=log10pcs,
-            masses=mass_grid,
-            radii=radii_interp,
-            lambdas=lambdas_interp,
+            log10pcs=jnp.log10(pcs_lim),
+            masses=masses_lim,
+            radii=radii_lim,
+            lambdas=lambdas_lim,
+            extra=extra_processed,
         )