jesterTOV.inference.likelihoods.gw.GWLikelihood

jesterTOV.inference.likelihoods.gw.GWLikelihood#

class GWLikelihood(
event_name,
model_dir,
penalty_value=0.0,
N_masses_evaluation=2000,
N_masses_batch_size=1000,
seed=42,
)[source]#

Bases: LikelihoodBase

Gravitational wave likelihood using pre-sampled masses for deterministic evaluation

This likelihood improves upon GWLikelihoodResampled by pre-sampling mass pairs once at initialization, eliminating the need for the _random_key parameter and providing deterministic likelihood evaluations critical for sampler convergence.

Key improvements over GWLikelihoodResampled: 1. Deterministic: Same EOS parameters → same likelihood value 2. No _random_key hack: Uses fixed seed at initialization 3. Scalable: Can use N=10,000+ samples efficiently on GPU 4. Fair comparison: All EOS evaluated at identical mass points 5. Better convergence: Smooth likelihood surface for MCMC/SMC

The likelihood works by:

  1. Pre-sampling (m1, m2) pairs from the trained flow at initialization

  2. For each EOS evaluation: interpolate Λ1, Λ2 from the candidate EOS at the fixed mass points, evaluate flow log_prob on (m1, m2, Λ1_EOS, Λ2_EOS), apply penalties for masses exceeding Mtov, and average over all pre-sampled mass pairs

Parameters:

  • event_name (str) – Name of the GW event (e.g., “GW170817”)

  • model_dir (str) – Path to directory containing the trained normalizing flow model

  • penalty_value (float, optional) – Penalty value for samples where masses exceed Mtov (default: 0.0, i.e. no penalty)

  • N_masses_evaluation (int, optional) – Number of mass samples to pre-sample (default: 2000) Large values recommended - GPU parallelization makes this cheap!

  • N_masses_batch_size (int, optional) – Batch size for jax.lax.map processing (default: 1000)

  • seed (int, optional) – Random seed for mass pre-sampling (default: 42) Fixed seed ensures reproducibility across runs

Variables:

  • event_name (str) – Name of the GW event

  • model_dir (str) – Path to directory containing the trained normalizing flow model

  • penalty_value (float) – Penalty value for samples where masses exceed Mtov

  • N_masses_evaluation (int) – Number of pre-sampled mass pairs

  • N_masses_batch_size (int) – Batch size for processing

  • seed (int) – Random seed used for pre-sampling

  • flow (Flow) – Normalizing flow model for this GW event

  • fixed_mass_samples (Float[Array, "n_samples 2"]) – Pre-sampled (m1, m2) pairs from the flow, shape [N, 2]

Notes

This class does NOT require _random_key in the parameter dictionary, unlike GWLikelihoodResampled. The seed is only used once at initialization.

GPU parallelization via jax.lax.map means N=10,000 samples costs nearly the same as N=20, so use large N for near-integration accuracy.

Examples

Configure in YAML:

likelihoods:
  - type: "gw"
    enabled: true
    parameters:
      events:
        - name: "GW170817"
      N_masses_evaluation: 2000  # Default value
      N_masses_batch_size: 1000
      seed: 42
__init__(
event_name,
model_dir,
penalty_value=0.0,
N_masses_evaluation=2000,
N_masses_batch_size=1000,
seed=42,
)[source]#

Methods

__init__(event_name, model_dir[, ...])

evaluate(params)

Evaluate log likelihood for given EOS parameters

Attributes

data

The data for the likelihood.

model

The model for the likelihood.

event_name

model_dir

penalty_value

N_masses_evaluation

N_masses_batch_size

seed

flow

fixed_mass_samples

N_masses_batch_size: int#
N_masses_evaluation: int#
evaluate(params)[source]#

Evaluate log likelihood for given EOS parameters

Parameters:

params (dict[str, Float | Array]) –

Must contain: - ‘masses_EOS’: Array of neutron star masses from EOS - ‘Lambdas_EOS’: Array of tidal deformabilities from EOS

Note: Does NOT require ‘_random_key’ (unlike GWLikelihood)

Return type:

Float

Returns:

Float – Log likelihood value for this GW event

event_name: str#
fixed_mass_samples: Float[Array, 'n_samples 2']#
flow: Flow#
model_dir: str#
penalty_value: float#
seed: int#
Contents