aspire.samplers.smc.blackjax

Attributes

logger

Classes

BlackJAXSMC

BlackJAX SMC sampler.

Module Contents

aspire.samplers.smc.blackjax.logger

class aspire.samplers.smc.blackjax.BlackJAXSMC(log_likelihood, log_prior, dims, prior_flow, xp, dtype=None, parameters=None, preconditioning_transform=None, rng=None)

Bases: aspire.samplers.smc.base.SMCSampler

BlackJAX SMC sampler.

Parameters:

rng (numpy.random.Generator | None)

key = None

rng

log_prob(x, beta=None)

Log probability function compatible with BlackJAX.

sample(n_samples, n_steps=None, adaptive=True, target_efficiency=0.5, target_efficiency_rate=1.0, n_final_samples=None, sampler_kwargs=None, rng_key=None, checkpoint_callback=None, checkpoint_every=None, checkpoint_file_path=None, resume_from=None)

Sample using BlackJAX SMC.

Parameters:

• n_samples (int) – Number of samples to draw.

• n_steps (int) – Number of SMC steps.

• adaptive (bool) – Whether to use adaptive tempering.

• target_efficiency (float) – Target efficiency for adaptive tempering.

• n_final_samples (int | None) – Number of final samples to return.

• sampler_kwargs (dict | None) – Additional arguments for the BlackJAX sampler. - algorithm: str, one of “nuts”, “hmc”, “rwmh”, “random_walk” - n_steps: int, number of MCMC steps per mutation - step_size: float, step size for HMC/NUTS - inverse_mass_matrix: array, inverse mass matrix - sigma: float or array, proposal covariance for random walk MH - num_integration_steps: int, integration steps for HMC

• rng_key (jax.random.key| None) – JAX random key for reproducibility.

• target_efficiency_rate (float)

• checkpoint_every (int | None)

• checkpoint_file_path (str | None)

• resume_from (str | bytes | dict | None)

mutate(particles, beta, n_steps=None)

Mutate particles using BlackJAX MCMC.