Config
Pydantic configuration models for Jim's samplers.
Each sampler has its own *Config class discriminated by a type literal;
SamplerConfig is the discriminated-union annotation a caller passes to
Jim(..., sampler_config=...).
Attributes:
| Name | Type | Description |
|---|---|---|
SamplerConfig |
Discriminated union of every concrete sampler config. |
SamplerConfig = Annotated[Union[FlowMCConfig, BlackJAXNSAWConfig, BlackJAXNSSConfig, BlackJAXSMCConfig], Discriminator('type')]
module-attribute
¤
Discriminated union of every concrete sampler config.
BaseSamplerConfig
¤
Bases: BaseModel
Fields shared by all sampler configs.
BlackJAXNSAWConfig
¤
Bases: BaseSamplerConfig, _CheckpointMixin
Configuration for the BlackJAX acceptance-walk nested sampler.
Note
This sampler requires the sampling space to be the unit hypercube
[0, 1]^n_dims. When using Jim, this means all
sample_transforms must map the prior support onto the unit cube.
Note
Periodic parameters are not configured here. Pass a
periodic argument to Jim instead.
For NS-AW, bounds are implicit as [0, 1]; just list the
parameter names.
BlackJAXNSSConfig
¤
Bases: BaseSamplerConfig, _CheckpointMixin
Configuration for the BlackJAX nested slice sampler.
Note
Periodic parameters are not configured here. Pass a periodic
argument to Jim instead.
BlackJAXSMCConfig
¤
Bases: BaseSamplerConfig, _CheckpointMixin
Configuration for the BlackJAX SMC sampler.
Parameters¤
batch_size : int, optional
Number of particles to process per sequential batch during the MCMC
update step. When batch_size > 0, the sampler uses jax.lax.map
instead of jax.vmap, which reduces peak GPU memory at the cost
of sequential execution. 0 (default) uses the original full
jax.vmap behaviour.
Note
Periodic parameters are not configured here. Pass a periodic
argument to Jim instead.
FlowMCConfig
¤
Bases: BaseSamplerConfig, _CheckpointMixin
Configuration for FlowMCSampler.
The local_kernel field selects the MCMC kernel used for local proposals:
"MALA"— Metropolis-Adjusted Langevin; default."HMC"— Hamiltonian Monte Carlo."GRW"— Gaussian random walk.
Parallel tempering is off by default. To enable, pass a
ParallelTemperingConfig,
a dict of its fields, or simply True (uses all defaults).
False disables it.
Note
Only the sub-config matching the active local_kernel is used.
Non-default values in inactive sub-configs emit a UserWarning.
Note
Periodic parameters are not configured here. Pass a periodic
argument to Jim instead; Jim resolves
parameter names to dimension indices and passes them to the sampler.
GRWConfig
¤
Bases: BaseModel
Gaussian random-walk local-kernel settings for the flowMC backend.
step_size may be a scalar float (applied uniformly to all
dimensions) or a 1-D np.ndarray of length n_dims for
per-dimension step sizes.
HMCConfig
¤
Bases: BaseModel
HMC local-kernel settings for the flowMC backend.
step_size may be a scalar float (applied uniformly) or a 1-D
np.ndarray of length n_dims for per-dimension sizes.
condition_matrix may also be a scalar or 1-D array and is used
as the mass matrix / preconditioning for HMC leapfrog steps.
MALAConfig
¤
Bases: BaseModel
MALA local-kernel settings for the flowMC backend.
step_size may be a scalar float (applied uniformly to all
dimensions) or a 1-D np.ndarray of length n_dims for
per-dimension step sizes.
ParallelTemperingConfig
¤
Bases: BaseModel
Parallel-tempering settings for the flowMC backend.
Construct directly or pass a plain dict to FlowMCConfig.parallel_tempering.
Use True to enable with all defaults, False / None to disable.