AbstractHypercubeSampler

class pmrf.infer.AbstractHypercubeSampler

Bases: Module

Interface for samplers operating in a unit hypercube (e.g., classical Nested Sampling).

All inputs (u0, init_cube_samples etc.) must be in the unit hypercube, whereas any outputs (e.g. samples in SampleResults) must be in physical space.

abstractmethod run(loglikelihood_fn: Callable[[PyTree, Any], Any], prior_transform_fn: Callable[[PyTree, Any], PyTree], u0: PyTree, args: PyTree[Any], key: Array, init_cube_samples: PyTree | None = None, max_steps: int | None = None, **kwargs) → tuple[SampleResult, Any]

Execute the sampling algorithm.

Parameters:

  • loglikelihood_fn (callable) – A function taking the physical parameters and args as input and returning the log-likelihood.

  • prior_transform_fn (callable) – A function taking the hypercube parameters and args as input and returning the physical parameters.

  • u0 (PyTree) – The initial parameters in the unit hypercube, either for shape reference or as a starting point.

  • args (Any) – Args to pass to fn.

  • key (Array) – A random JAX key.

  • init_cube_samples (PyTree, optional) – An optional batched PyTree the same structure as u0 with initial hypercube samples to warm-start the algorithm.

  • max_steps (int, optional) – The maximum number of sampling steps to take. If None, implies there should be no limit.

  • **kwargs – Runtime arguments forward to the solver backend.

Returns:

A tuple of (pmrf.infer.SampleResult, metrics)`.

Return type:

tuple