Skip to content

Jim

Jim ¤

Master class for gravitational-wave parameter estimation.

Wires together a LikelihoodBase, a Prior, optional parameter transforms, and a pluggable JAX Sampler selected via a typed sampler_config object.

Methods:

Name Description
__init__

Initialise Jim and build the internal sampler.
add_name

Convert a flat sampling-space array to a named dict.
evaluate_posterior

Log-posterior in the sampling space.
evaluate_prior

Log-prior in the sampling space (with Jacobian corrections from sample_transforms).
get_diagnostics

Return run-level diagnostics from the most recent sample call.
get_samples

Retrieve posterior samples in prior space, optionally downsampled.
sample

Run the sampler.
sample_initial_positions

Draw n_points initial positions from the prior in sampling space.

__init__(likelihood: LikelihoodBase, prior: Prior, sampler_config: SamplerConfig, *, sample_transforms: Sequence[BijectiveTransform] = (), likelihood_transforms: Sequence[NtoMTransform] = (), periodic: Optional[list[str] | dict[str, tuple[float, float]]] = None, seed: int = 0, verbose: bool = False) -> None ¤

Initialise Jim and build the internal sampler.

Parameters:

Name Type Description Default
likelihood LikelihoodBase

The likelihood to evaluate.

 required
prior Prior

The prior distribution.

 required
sampler_config SamplerConfig

Pydantic config selecting and configuring the sampler backend (e.g. FlowMCConfig).

 required
sample_transforms Sequence[BijectiveTransform]

Bijective transforms applied in the sampling space (reversed when retrieving posterior samples).

 ()
likelihood_transforms Sequence[NtoMTransform]

Transforms applied to reach the likelihood parameter space from the prior parameter space.

 ()
periodic Optional[list[str] | dict[str, tuple[float, float]]]

Periodic sampling-space parameters. For most samplers, pass a dict mapping parameter name to (lo, hi) bounds (e.g. {"phase_c": (0.0, 6.2832)}). For the BlackJAX NS-AW sampler (unit-cube space), pass a list of parameter names (bounds are implicit as [0, 1]).

 None
seed int

Integer random seed. The key for the sampling run is derived from this seed at construction time, so sample is reproducible regardless of any intermediate operations (sanity checks, initial-position draws, etc.).

 0
verbose bool

Enable DEBUG-level logging for all jimgw components. At False (default) INFO-level messages are always shown. Pass True to also see per-step diagnostics and backend-specific progress output (e.g. flowMC training loss).

 False

add_name(x: Float[Array, ' n_dims']) -> dict[str, Float] ¤

Convert a flat sampling-space array to a named dict.

evaluate_posterior(params: Float[Array, ' n_dims']) -> Float ¤

Log-posterior in the sampling space.

evaluate_prior(params: Float[Array, ' n_dims']) -> Float ¤

Log-prior in the sampling space (with Jacobian corrections from sample_transforms).

get_diagnostics() -> dict[str, Any] ¤

Return run-level diagnostics from the most recent sample call.

Returns:

Type Description
dict[str, Any]

Plain dict of backend-specific diagnostics.

get_samples(n_samples: int = 0) -> dict[str, np.ndarray] ¤

Retrieve posterior samples in prior space, optionally downsampled.

Calls Sampler.get_samples on the underlying sampler, which returns equally-weighted posterior samples. Pass n_samples to further downsample.

Parameters:

Name Type Description Default
n_samples int

Target number of samples. If 0 (default) returns all

 0

Returns:

Type Description
dict[str, ndarray]

Dict mapping prior parameter names to 1-D numpy arrays in prior
dict[str, ndarray]

space, plus an extra item containing the log-likelihood values.

sample(initial_position: Optional[Float[Array, 'n_chains n_dims']] = None) -> None ¤

Run the sampler.

The sampling key is pre-reserved at construction time from seed, so results are reproducible regardless of any calls made before this method (e.g. the construction-time posterior verification).

Parameters:

Name Type Description Default
initial_position Optional[Float[Array, 'n_chains n_dims']]

Starting positions in sampling space, or None (default) to draw them from the prior. The expected shape depends on the backend:

  • flowMC: (n_chains, n_dims) or (n_dims,) (broadcast to all chains).
  • BlackJAX NS-AW / NSS: exactly (n_live, n_dims).
  • BlackJAX SMC: exactly (n_particles, n_dims).

The concrete sampler validates the shape and raises ValueError on mismatch.

 None

sample_initial_positions(n_points: int, rng_key: Optional[Key] = None) -> Float[Array, 'n_points n_dims'] ¤

Draw n_points initial positions from the prior in sampling space.

Parameters:

Name Type Description Default
n_points int

Number of positions to draw.

 required
rng_key Optional[Key]

Optional explicit PRNG key. If None, Jim's internal auxiliary key is advanced automatically.

 None

Returns:

Type Description
Float[Array, 'n_points n_dims']

Array of shape (n_points, n_dims) in sampling space.