mchammers package

Copyright (c) 2024 The Princeton Hammers. All rights reserved.

mchammers: A toy package for sampling posteriors. Our final project for APC 524 at Princeton University.

class mchammers.Sampler(num_step, num_walker, num_dim, prior_bounds, state_init, log_prob_func, args=None, kwargs=None, frac_burn=0.2, seed=None, flatten=True)

Bases: ABC

The Sampler object.

Parameters:

num_step:

The number of steps.

num_walker:

The number of walkers.

num_dim:

The number of dimensions.

prior_bounds:

The prior bounds on the model parameters.

state_init:

The initial state.

log_prob_func:

The log probability function.

args:

The arguments of the log probability function.

kwargs:

The keyword arguments of the log probability function.

frac_burn:

The burn fraction.

seed:

The random number generator seed.

flatten:

Whether to flatten the final sample array.

param num_step:

type num_step:

int

param num_walker:

type num_walker:

int

param num_dim:

type num_dim:

int

param prior_bounds:

type prior_bounds:

Sequence[tuple[Union[int, float], Union[int, float]]]

param state_init:

type state_init:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

param log_prob_func:

type log_prob_func:

Callable[..., ndarray[tuple[int, ...], dtype[Union[int, float]]]]

param args:

type args:

Optional[Sequence[Any]]

param kwargs:

type kwargs:

Optional[Dict[Any, Any]]

param frac_burn:

type frac_burn:

Union[int, float]

param seed:

type seed:

int | None

param flatten:

type flatten:

bool

calc_rate_accept()

Calculate the acceptance rate.

Return type:

None

post_process()

Post-process the samples.

Return type:

None

abstract prob_accept(log_prob_curr, log_prob_prop)

Calculate the acceptance probability.

Return type:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

Parameters:

log_prob_curr:

The log probability of the current state.

log_prob_prop:

The log probability of the proposed state.

Returns:

:

The acceptance probability.

param log_prob_curr:

type log_prob_curr:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

param log_prob_prop:

type log_prob_prop:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

run()

Run the sampler.

Return type:

None

Parameters:

state_init:

The initial state.

abstract sample_prop(idx_group)

Sample the proposal distribution.

Return type:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

Parameters:

idx_group:

The index of the current group.

Returns:

:

An array of samples of the proposal distribution.

param idx_group:

type idx_group:

int

step()

Step the sampler.

Return type:

None

class mchammers.SamplerBasic(num_step, num_walker, num_dim, prior_bounds, state_init, log_prob_func, std_rel_prop, args=None, kwargs=None, frac_burn=0.2, seed=None, flatten=True)

Bases: Sampler

The SamplerBasic object.

Parameters:

std:

The standard deviation of the proposal distribution relative to the prior bounds.

param num_step:

type num_step:

int

param num_walker:

type num_walker:

int

param num_dim:

type num_dim:

int

param prior_bounds:

type prior_bounds:

Sequence[tuple[Union[int, float], Union[int, float]]]

param state_init:

type state_init:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

param log_prob_func:

type log_prob_func:

Callable[..., ndarray[tuple[int, ...], dtype[Union[int, float]]]]

param std_rel_prop:

type std_rel_prop:

Union[int, float]

param args:

type args:

Optional[Sequence[Any]]

param kwargs:

type kwargs:

Optional[Dict[Any, Any]]

param frac_burn:

type frac_burn:

Union[int, float]

param seed:

type seed:

int | None

param flatten:

type flatten:

bool

prob_accept(log_prob_curr, log_prob_prop)

Calculate the acceptance probability.

Return type:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

Parameters:

log_prob_curr:

The log probability of the current state.

log_prob_prop:

The log probability of the proposed state.

Returns:

:

The acceptance probability.

param log_prob_curr:

type log_prob_curr:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

param log_prob_prop:

type log_prob_prop:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

sample_prop(idx_group)

Sample the proposal distribution.

Return type:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

Parameters:

idx_group:

The index of the current group.

Returns:

:

An array of samples of the proposal distribution.

param idx_group:

type idx_group:

int

class mchammers.SamplerBasicMPI(num_step, num_walker, num_dim, prior_bounds, state_init, log_prob_func, comm, mpi_sum, size, rank, std_rel_prop, args=None, kwargs=None, frac_burn=0.2, seed=None, flatten=True)

Bases: SamplerMPI

The SamplerBasic object.

Parameters:

cov:

The covariance of the proposal distribution in each dimension.

param num_step:

type num_step:

int

param num_walker:

type num_walker:

int

param num_dim:

type num_dim:

int

param prior_bounds:

type prior_bounds:

Sequence[tuple[Union[int, float], Union[int, float]]]

param state_init:

type state_init:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

param log_prob_func:

type log_prob_func:

Callable[..., ndarray[tuple[int, ...], dtype[Union[int, float]]]]

param comm:

type comm:

object

param mpi_sum:

type mpi_sum:

object

param size:

type size:

int

param rank:

type rank:

int

param std_rel_prop:

type std_rel_prop:

Union[int, float]

param args:

type args:

Optional[Sequence[Any]]

param kwargs:

type kwargs:

Optional[Dict[Any, Any]]

param frac_burn:

type frac_burn:

Union[int, float]

param seed:

type seed:

int | None

param flatten:

type flatten:

bool

prob_accept(log_prob_curr, log_prob_prop)

Calculate the acceptance probability.

Return type:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

Parameters:

log_prob_curr:

The log probability of the current state.

log_prob_prop:

The log probability of the proposed state.

Returns:

:

The acceptance probability.

param log_prob_curr:

type log_prob_curr:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

param log_prob_prop:

type log_prob_prop:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

sample_prop(idx_group)

Sample the proposal distribution.

Return type:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

Parameters:

idx_group:

The index of the current group.

Returns:

:

An array of samples of the proposal distribution.

param idx_group:

type idx_group:

int

class mchammers.SamplerMPI(num_step, num_walker, num_dim, prior_bounds, state_init, log_prob_func, comm, mpi_sum, size, rank, args=None, kwargs=None, frac_burn=0.2, seed=None, flatten=True)

Bases: Sampler

The SamplerMPI object.

Parameters:

comm:

The MPI communicator.

mpi_sum:

The MPI sum operation.

size:

The number of MPI ranks.

rank:

The MPI rank.

param num_step:

type num_step:

int

param num_walker:

type num_walker:

int

param num_dim:

type num_dim:

int

param prior_bounds:

type prior_bounds:

Sequence[tuple[Union[int, float], Union[int, float]]]

param state_init:

type state_init:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

param log_prob_func:

type log_prob_func:

Callable[..., ndarray[tuple[int, ...], dtype[Union[int, float]]]]

param comm:

type comm:

object

param mpi_sum:

type mpi_sum:

object

param size:

type size:

int

param rank:

type rank:

int

param args:

type args:

Optional[Sequence[Any]]

param kwargs:

type kwargs:

Optional[Dict[Any, Any]]

param frac_burn:

type frac_burn:

Union[int, float]

param seed:

type seed:

int | None

param flatten:

type flatten:

bool

post_process()

Post-process the samples.

Return type:

None

class mchammers.SamplerStretch(num_step, num_walker, num_dim, prior_bounds, state_init, log_prob_func, args=None, kwargs=None, a=2.0, frac_burn=0.2, seed=None, flatten=True)

Bases: Sampler

The SamplerStretch object. Implements the affine-invariant ensemble sampler using the parallel stretch move.

Parameters:

a:

The stretch move parameter, a > 1.

param num_step:

type num_step:

int

param num_walker:

type num_walker:

int

param num_dim:

type num_dim:

int

param prior_bounds:

type prior_bounds:

Sequence[tuple[Union[int, float], Union[int, float]]]

param state_init:

type state_init:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

param log_prob_func:

type log_prob_func:

Callable[..., ndarray[tuple[int, ...], dtype[Union[int, float]]]]

param args:

type args:

Optional[Sequence[Any]]

param kwargs:

type kwargs:

Optional[Dict[Any, Any]]

param a:

type a:

Union[int, float]

param frac_burn:

type frac_burn:

Union[int, float]

param seed:

type seed:

int | None

param flatten:

type flatten:

bool

prob_accept(log_prob_curr, log_prob_prop)

Calculate the acceptance probability.

Return type:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

Parameters:

log_prob_curr:

The log probability of the current state.

log_prob_prop:

The log probability of the proposed state.

Returns:

:

The acceptance probability.

param log_prob_curr:

type log_prob_curr:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

param log_prob_prop:

type log_prob_prop:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

sample_Z(size)

Sample Z from g(Z) ∝ 1/sqrt(Z) within [1/a, a] using the inverse transform sampling

Parameters:

size (int)

Return type:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

sample_prop(idx_group)

Sample the proposal distribution using the stretch move.

Return type:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

Parameters:

group:

Indices of the walkers in the current group.

comp_group:

Indices of the walkers in the complementary group.

Returns:

:

An array of proposed samples.

param idx_group:

type idx_group:

int

Submodules

mchammers.hammer module

class mchammers.hammer.Sampler(num_step, num_walker, num_dim, prior_bounds, state_init, log_prob_func, args=None, kwargs=None, frac_burn=0.2, seed=None, flatten=True)

Bases: ABC

The Sampler object.

Parameters:

num_step:

The number of steps.

num_walker:

The number of walkers.

num_dim:

The number of dimensions.

prior_bounds:

The prior bounds on the model parameters.

state_init:

The initial state.

log_prob_func:

The log probability function.

args:

The arguments of the log probability function.

kwargs:

The keyword arguments of the log probability function.

frac_burn:

The burn fraction.

seed:

The random number generator seed.

flatten:

Whether to flatten the final sample array.

param num_step:

type num_step:

int

param num_walker:

type num_walker:

int

param num_dim:

type num_dim:

int

param prior_bounds:

type prior_bounds:

Sequence[tuple[Union[int, float], Union[int, float]]]

param state_init:

type state_init:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

param log_prob_func:

type log_prob_func:

Callable[..., ndarray[tuple[int, ...], dtype[Union[int, float]]]]

param args:

type args:

Optional[Sequence[Any]]

param kwargs:

type kwargs:

Optional[Dict[Any, Any]]

param frac_burn:

type frac_burn:

Union[int, float]

param seed:

type seed:

int | None

param flatten:

type flatten:

bool

calc_rate_accept()

Calculate the acceptance rate.

Return type:

None

post_process()

Post-process the samples.

Return type:

None

abstract prob_accept(log_prob_curr, log_prob_prop)

Calculate the acceptance probability.

Return type:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

Parameters:

log_prob_curr:

The log probability of the current state.

log_prob_prop:

The log probability of the proposed state.

Returns:

:

The acceptance probability.

param log_prob_curr:

type log_prob_curr:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

param log_prob_prop:

type log_prob_prop:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

run()

Run the sampler.

Return type:

None

Parameters:

state_init:

The initial state.

abstract sample_prop(idx_group)

Sample the proposal distribution.

Return type:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

Parameters:

idx_group:

The index of the current group.

Returns:

:

An array of samples of the proposal distribution.

param idx_group:

type idx_group:

int

step()

Step the sampler.

Return type:

None

class mchammers.hammer.SamplerBasic(num_step, num_walker, num_dim, prior_bounds, state_init, log_prob_func, std_rel_prop, args=None, kwargs=None, frac_burn=0.2, seed=None, flatten=True)

Bases: Sampler

The SamplerBasic object.

Parameters:

std:

The standard deviation of the proposal distribution relative to the prior bounds.

param num_step:

type num_step:

int

param num_walker:

type num_walker:

int

param num_dim:

type num_dim:

int

param prior_bounds:

type prior_bounds:

Sequence[tuple[Union[int, float], Union[int, float]]]

param state_init:

type state_init:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

param log_prob_func:

type log_prob_func:

Callable[..., ndarray[tuple[int, ...], dtype[Union[int, float]]]]

param std_rel_prop:

type std_rel_prop:

Union[int, float]

param args:

type args:

Optional[Sequence[Any]]

param kwargs:

type kwargs:

Optional[Dict[Any, Any]]

param frac_burn:

type frac_burn:

Union[int, float]

param seed:

type seed:

int | None

param flatten:

type flatten:

bool

prob_accept(log_prob_curr, log_prob_prop)

Calculate the acceptance probability.

Return type:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

Parameters:

log_prob_curr:

The log probability of the current state.

log_prob_prop:

The log probability of the proposed state.

Returns:

:

The acceptance probability.

param log_prob_curr:

type log_prob_curr:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

param log_prob_prop:

type log_prob_prop:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

sample_prop(idx_group)

Sample the proposal distribution.

Return type:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

Parameters:

idx_group:

The index of the current group.

Returns:

:

An array of samples of the proposal distribution.

param idx_group:

type idx_group:

int

class mchammers.hammer.SamplerBasicMPI(num_step, num_walker, num_dim, prior_bounds, state_init, log_prob_func, comm, mpi_sum, size, rank, std_rel_prop, args=None, kwargs=None, frac_burn=0.2, seed=None, flatten=True)

Bases: SamplerMPI

The SamplerBasic object.

Parameters:

cov:

The covariance of the proposal distribution in each dimension.

param num_step:

type num_step:

int

param num_walker:

type num_walker:

int

param num_dim:

type num_dim:

int

param prior_bounds:

type prior_bounds:

Sequence[tuple[Union[int, float], Union[int, float]]]

param state_init:

type state_init:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

param log_prob_func:

type log_prob_func:

Callable[..., ndarray[tuple[int, ...], dtype[Union[int, float]]]]

param comm:

type comm:

object

param mpi_sum:

type mpi_sum:

object

param size:

type size:

int

param rank:

type rank:

int

param std_rel_prop:

type std_rel_prop:

Union[int, float]

param args:

type args:

Optional[Sequence[Any]]

param kwargs:

type kwargs:

Optional[Dict[Any, Any]]

param frac_burn:

type frac_burn:

Union[int, float]

param seed:

type seed:

int | None

param flatten:

type flatten:

bool

prob_accept(log_prob_curr, log_prob_prop)

Calculate the acceptance probability.

Return type:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

Parameters:

log_prob_curr:

The log probability of the current state.

log_prob_prop:

The log probability of the proposed state.

Returns:

:

The acceptance probability.

param log_prob_curr:

type log_prob_curr:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

param log_prob_prop:

type log_prob_prop:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

sample_prop(idx_group)

Sample the proposal distribution.

Return type:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

Parameters:

idx_group:

The index of the current group.

Returns:

:

An array of samples of the proposal distribution.

param idx_group:

type idx_group:

int

class mchammers.hammer.SamplerMPI(num_step, num_walker, num_dim, prior_bounds, state_init, log_prob_func, comm, mpi_sum, size, rank, args=None, kwargs=None, frac_burn=0.2, seed=None, flatten=True)

Bases: Sampler

The SamplerMPI object.

Parameters:

comm:

The MPI communicator.

mpi_sum:

The MPI sum operation.

size:

The number of MPI ranks.

rank:

The MPI rank.

param num_step:

type num_step:

int

param num_walker:

type num_walker:

int

param num_dim:

type num_dim:

int

param prior_bounds:

type prior_bounds:

Sequence[tuple[Union[int, float], Union[int, float]]]

param state_init:

type state_init:

ndarray[tuple[int, ...], dtype[Union[int, float]]]

param log_prob_func:

type log_prob_func:

Callable[..., ndarray[tuple[int, ...], dtype[Union[int, float]]]]

param comm:

type comm:

object

param mpi_sum:

type mpi_sum:

object

param size:

type size:

int

param rank:

type rank:

int

param args:

type args:

Optional[Sequence[Any]]

param kwargs:

type kwargs:

Optional[Dict[Any, Any]]

param frac_burn:

type frac_burn:

Union[int, float]

param seed:

type seed:

int | None

param flatten:

type flatten:

bool

post_process()

Post-process the samples.

Return type:

None