AbstractSplitSampler

class pmrf.infer.AbstractSplitSampler

Bases: Module

Interface for samplers needing separate likelihood and prior densities (e.g., modern Nested Sampling).

abstractmethod run(loglikelihood_fn: Callable[[PyTree, Any], Any], logprior_fn: Callable[[PyTree, Any], Shaped[jaxlib._jax.Array, '']], y0: PyTree, args: PyTree[Any], key: Array, init_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 parameters and args as input and returning the log-likelihood.

  • logprior_fn (callable) – A function taking the parameters and args as input and returning the log prior probability.

  • y0 (PyTree) – The initial parameters, either for shape reference or as a starting point.

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

  • key (Array) – A random JAX key.

  • init_samples (PyTree, optional) – An optional batched PyTree the same structure as y0 with initial 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