`module` `adabmDCA.ptt.sampler`

Parallel Trajectory Tempering with snapshot retention and mixing recovery.

All replicas target beta=1. The active ladder can be shortened using a reservoir; optional full sampler mode retains the historical models. The initial profile normalizer is analytic and subsequent ones use bridges.

Global Variables

DEFAULT_PTT_N_CHAINS
DEFAULT_PTT_SAMPLER
IMMOBILE_THRESHOLD
LOCAL_KERNELS

`function` `bridge_increment`

bridge_increment(lower, upper, samples)

Return log(Z_upper/Z_lower) from samples of the lower Hamiltonian.

`class` `PartitionEstimate`

Estimate of the endpoint's log partition function, with its provenance.

Attributes:

log_z: Natural log of the partition function.
model_version: Training update of the endpoint.
ladder_version: Version of the ladder the bridges ran along.
model_id: Hash of the endpoint parameters.
sample_round: Exchange round of the chains used.
method: "ptt_bridge" (forward bridges) or "ptt_bar" (BAR).
status: "estimated".

`method` `init`

__init__(
    log_z: 'float',
    model_version: 'int',
    ladder_version: 'int',
    model_id: 'str',
    sample_round: 'int',
    method: 'str' = 'ptt_bridge',
    status: 'str' = 'estimated'
) → None

`class` `PTTSampler`

A Parallel Trajectory Tempering ladder: models, their chains and its own RNG.

During training (see train_model with ptt=) the ladder runs from an exactly normalized profile model to the model being trained, and snapshots of the training trajectory are kept as intermediate replicas. Each exchange round swaps configurations between adjacent replicas and then runs local Monte Carlo sweeps. The same ladder is saved in the PTT archive and reused to sample the trained model and to estimate its log Z.

Most users only load an archive and sample or measure it:

from_archive, save_archive, fork;
draw, endpoint_samples, endpoint_params;
partition_estimate, entropy, ladder_statistics, ladder_health, measure_renewal.

The remaining public methods (transition_target, update_replica_chain, refresh_reservoir, ...) are the training machinery used by the PTT trainer. Mutating the ladder directly is unsupported. Archives contain arrays and JSON only, never pickled objects.

Example:

sampler = PTTSampler.from_archive("model/ptt.h5", device="cuda")
sampler.prepare_sampling_ladder(2000)
sampler.measure_renewal(local_sweeps=10).status  # 'converged'
sequences = sampler.draw(1000, local_sweeps=10)   # one-hot, (1000, L, q)
sampler.entropy()["entropy"]

`method` `init`

__init__(
    params: 'Mapping[str, Tensor]',
    tokens: 'str',
    n_chains: 'int' = 2048,
    sampler: 'str' = 'metropolized_gibbs',
    config: 'PTTConfig | None' = None,
    seed: 'int' = 0
) → None

Start a two-replica ladder at an uncoupled (profile) model.

Both replicas start at params; training then moves the endpoint. To continue from a saved state use from_archive.

Args:

params: bias of shape (L, q) and an all-zero coupling_matrix of shape (L, q, L, q), float32 or float64, on CPU, CUDA or MPS (float32).
tokens: Ordered alphabet of length q.
n_chains: Chains per replica.
sampler: Local kernel: "metropolized_gibbs", "gibbs" or "metropolis".
config: PTT settings; defaults to PTTConfig().
seed: Random seed of the sampler's own generator.

Raises:

InputValidationError: If the parameters are coupled, invalid or underflow, or a setting is invalid.

`property` n_active

Number of replicas in the active ladder (above the reservoir, if any).

`method` `advance`

advance(
    rounds: 'int' = 1,
    local_sweeps: 'int' = 1,
    is_cancelled: 'Callable[[], bool] | None' = None,
    on_round: 'Callable[[int, int], None] | None' = None,
    until: 'Callable[[], bool] | None' = None,
    return_samples: 'bool' = True
) → Tensor | None

Run exchange rounds at fixed parameters.

Each round swaps configurations between adjacent replicas, permutes chains within replicas, redraws the bottom replica, and runs local sweeps.

Args:

rounds: Largest number of rounds.
local_sweeps: Local sweeps per round.
is_cancelled: Optional callable; returning True stops the run.
on_round: Optional callback on_round(completed, rounds) after each round.
until: Optional callable; returning True stops after the current round.
return_samples: Return one-hot endpoint chains; False skips their construction when only advancing the sampler state.

Returns: A copy of the endpoint chains, one-hot, shape (n_chains, L, q), or None when return_samples=False.

Raises:

OperationCancelledError: If is_cancelled returns True.

`method` `capture_state`

capture_state() → dict[str, Any]

Return a deep copy of the numerical state, for a recovery checkpoint.

Training history and recovery points are excluded; see restore_state.

`method` `draw`

draw(
    n_samples: 'int',
    spacing: 'int' = 1,
    local_sweeps: 'int' = 1,
    is_cancelled: 'Callable[[], bool] | None' = None
) → Tensor

Collect endpoint samples, taking a batch of chains every spacing rounds.

Samples are only as good as the ladder's equilibration: run measure_renewal (or use sample_sequences) first.

Args:

n_samples: Number of samples.
spacing: Exchange rounds between successive batches of n_chains.
local_sweeps: Local sweeps per round.
is_cancelled: Optional callable; returning True stops the run.

Returns: One-hot samples of shape (n_samples, L, q).

`method` `endpoint_params`

endpoint_params() → dict[str, Tensor]

Return a copy of the endpoint (last replica) parameters: the trained model.

`method` `endpoint_samples`

endpoint_samples(copy: 'bool' = True) → Tensor

Return the endpoint chains, one-hot, shape (n_chains, L, q).

Args:

copy: Return a copy; False may return internal storage, which must not be modified.

`method` `entropy`

entropy(estimator: 'str' = 'bar') → dict[str, Any]

Estimate the endpoint entropy as <E> + log Z from the current chains.

Args:

estimator: log Z estimator, "bar" or "forward" (see partition_estimate).

Returns: entropy, mean_energy and free_energy (-log Z), in nats, with the fields of the PartitionEstimate.

`method` `equilibrate`

equilibrate(
    rounds: 'int | None' = None,
    local_sweeps: 'int' = 1,
    is_cancelled: 'Callable[[], bool] | None' = None,
    on_round: 'Callable[[int, int], None] | None' = None
) → Tensor

Run a fixed warmup of advance; this does not certify equilibrium.

Args:

rounds: Rounds to run; defaults to config.equilibration_rounds.
local_sweeps: Local sweeps per round.
is_cancelled: Optional callable; returning True stops the run.
on_round: Optional callback on_round(completed, rounds).

Returns: A copy of the endpoint chains, one-hot.

`method` `equilibrate_for_lag`

equilibrate_for_lag(
    relaxed: 'Callable[[PTTSampler], bool]',
    max_rounds: 'int',
    chunk: 'int' = 10,
    local_sweeps: 'int' = 1,
    is_cancelled: 'Callable[[], bool] | None' = None,
    on_progress: 'Callable[, None] | None' = None
) → tuple[bool, int, int]

Advance the ladder, then evolve at fixed parameters until relaxed(sampler) holds.

The response of the lag optimizer when the endpoint chains trail the model: parameters stay fixed while force_ladder_update gives the endpoint a recent neighbour and the ladder evolves in chunks of chunk rounds, up to max_rounds. Runs on a fork and commits only a healthy result. Returns (healthy, rounds, work).

`method` `estimate_mixing_time`

estimate_mixing_time(
    local_sweeps: 'int' = 1,
    is_cancelled: 'Callable[[], bool] | None' = None,
    full_ladder: 'bool' = False,
    on_progress: 'Callable[, None] | None' = None,
    bounded: 'bool' = True,
    max_rounds: 'int | None' = None,
    in_place: 'bool' = False,
    reference: 'bool' = False,
    method: 'str | None' = None
) → MixingEstimate

Run reference TRWA on a separate population and return its diagnostics.

Extend the experiment until it contains at least 20 times both estimated correlation times (configurable), or reaches its budget. Training populations and RNG are preserved; diagnostic work counts. full_ladder=True also checks retained history in full sampler mode. bounded=False continues until convergence or cancellation, ignoring the training mixing-round budget. in_place=True measures and equilibrates the current populations, as in reference TRWA. reference=True uses its FFT and fitting rules. The replica index of each uniquely tagged configuration (its lineage) is the observable. method defaults to config.mixing_method. 'renewal' replaces the replica-index experiment with a population-renewal check on the same kind of separate population (see _renewal_mixing_experiment).

Args:

local_sweeps: Local sweeps per exchange round.
is_cancelled: Optional callable; returning True stops the run.
full_ladder: Also check the retained history (full sampler mode).
on_progress: Optional callback on_progress(stage, done, total, **details).
bounded: Stop at config.mixing_max_rounds; False runs until converged.
max_rounds: Round budget overriding the configured one.
in_place: Measure and equilibrate the current populations.
reference: Use the reference TRWA fitting rules.
method: "autocorrelation" or "renewal"; defaults to config.mixing_method.

Returns: A MixingEstimate; converged tells whether the ladder mixes.

`method` `force_ladder_update`

force_ladder_update(
    local_sweeps: 'int' = 1,
    is_cancelled: 'Callable[[], bool] | None' = None,
    on_progress: 'Callable[, None] | None' = None
) → bool

Advance the snapshot procedure now, as if acceptance had just dropped below its target.

With a held snapshot, it replaces the replica below the endpoint (and the reservoir is refreshed when the active ladder grows too long); otherwise the stored snapshot, or a snapshot of the current endpoint when none is stored, is inserted below the endpoint and the endpoint is flagged as a checkpoint. Nothing is inserted twice at the same model version. Returns False when a reservoir refresh fails.

`method` `fork`

fork(seed: 'int | None' = None) → PTTSampler

Return an independent copy of the sampler, without its recovery points.

Args:

seed: New random seed, or None to continue the same random stream.

Returns: A new PTTSampler; changes to it do not affect this one.

`classmethod` `from_archive`

from_archive(
    path: 'str | Path',
    device: 'str' = 'cpu',
    mode: 'str' = 'generate',
    seed: 'int | None' = None
) → PTTSampler

Load a sampler from an HDF5 archive written by training or save_archive.

Every array is validated on load. Generation never writes the source archive. Loading on a different device than the one that saved the archive requires a new seed (and then a warmup): the random stream cannot be continued exactly across devices.

Args:

path: Archive path.
device: "cpu", a CUDA device such as "cuda", or "mps" (float32 archives).
mode: "generate" (sampling; training state dropped), "resume" (continue training exactly) or "inspect" (read only).
seed: New random seed, or None to continue the archived stream.

Returns: The loaded PTTSampler.

Raises:

InputValidationError: If the archive is invalid, or the device changes without a new seed.

`method` `ladder_health`

ladder_health(bootstrap: 'int' = 200, seed: 'int' = 0) → dict[str, Any]

Bidirectional reweighting diagnostics for every adjacent active pair.

See adabmDCA.ptt.health. Per pair: forward, reverse and BAR free energies with bootstrap errors, hysteresis, effective sample sizes and top-1% weight shares in both directions, the Crooks slope (-1 at equilibrium), mean swap acceptance and the fraction of configurations whose average swap probability is below IMMOBILE_THRESHOLD, and quantiles of the per-configuration acceptance in both directions. With birth tracking, the mean age of immobile and mobile upper configurations is included, and replicas gives, per replica, the median and 99th percentile of the configuration ages (rounds since birth) and the flow: the fraction of configurations that reached the top replica since birth. Also returns forward and BAR endpoint log Z.

Args:

bootstrap: Bootstrap resamples for the error bars.
seed: Seed of the bootstrap.

Returns: pairs (one dictionary per adjacent pair), replicas (one per replica, empty without birth tracking), log_z_forward, log_z_bar, log_z_bar_error and immobile_threshold.

`method` `ladder_statistics`

ladder_statistics(
    reference: 'Tensor',
    weights: 'Tensor',
    estimator: 'str' = 'bar'
) → list[dict[str, Any]]

Estimate every active model's entropy and weighted data likelihood.

Energies, entropy and logZ are in nats per sequence; likelihood also includes a per-residue column, matching the training convention. logZ uses BAR bridges by default (see partition_estimate); log_z_forward always reports the forward bridges for comparison.

Args:

reference: Data sequences, categorical (M, L) or one-hot (M, L, q).
weights: Weight of each reference sequence.
estimator: "bar" or "forward".

Returns: One dictionary per active replica, with log_z, mean_energy, entropy, log_likelihood, log_likelihood_per_residue, ...

`method` `measure_renewal`

measure_renewal(
    local_sweeps: 'int' = 1,
    max_rounds: 'int' = 20000,
    tolerance: 'float' = 0.01,
    blocks_per_warmup: 'int' = 1,
    stationary: 'bool' = True,
    is_cancelled: 'Callable[[], bool] | None' = None,
    on_progress: 'Callable[, None] | None' = None,
    stages: 'Sequence[str]' = ('renewal_warmup', 'renewal_stationary')
) → RenewalEstimate

Equilibrate the active ladder in place until its population is renewed twice (or once).

Requires birth tracking (prepare_sampling_ladder or start_birth_tracking). A configuration is born when it enters the active ladder: an exact profile draw at rung 0 of a full ladder, or a reservoir emission otherwise. The warmup ends at the first round where at most tolerance * n_chains configurations born before the start remain in the active ladder. The stationary phase then restarts the reference and advances in fixed blocks of the warmup length divided by blocks_per_warmup until the same holds at a block end. Stopping only at block ends keeps the stopping round independent of which configurations are old; shorter blocks waste fewer rounds after the renewal. With stationary=False only the warmup runs, and the measurement converges at the warmup renewal. No autocorrelation fit decides anything; see RenewalEstimate for the reported quantities. max_rounds bounds both phases together; exhausting it reports the old endpoint fraction as trapped_fraction.

Progress events carry the old fraction ladder_old of the whole ladder, the threshold it must reach, and every 10 rounds a renewal_forecast of the phase: decay_rounds, predicted_round and the first decay time of the phase, first_decay_rounds, whose growth signals configurations that leave more and more slowly.

Args:

local_sweeps: Local sweeps per exchange round.
max_rounds: Round budget of both phases together.
tolerance: Largest remaining fraction of old configurations, per replica.
blocks_per_warmup: Stationary phase: blocks per warmup length.
stationary: Run the stationary phase after the warmup.
is_cancelled: Optional callable; returning True stops the run.
on_progress: Optional callback on_progress(stage, done, total, **details).
stages: Stage names reported for the two phases.

Returns: A RenewalEstimate; status is "converged" when renewed in time.

`method` `partition_estimate`

partition_estimate(estimator: 'str' = 'forward') → PartitionEstimate

Estimate log Z of the endpoint from the exact anchor and one bridge per replica pair.

estimator='forward' reweights each lower population to the model above (the historical bridge used during training and stored in archives). 'bar' combines both populations of every pair with Bennett's acceptance ratio, which stays reliable when the forward weights are heavy-tailed.

Args:

estimator: "forward" or "bar".

Returns: A PartitionEstimate with log_z and its provenance.

`method` `prepare_sampling_ladder`

prepare_sampling_ladder(n_chains: 'int') → None

Rebuild the full ladder for generation, with fresh chains drawn from the profile.

The generation ladder is the training ladder: the exact profile, every update flagged ptt during training, and the final endpoint. This switches the sampler to generation mode and starts birth tracking, as measure_renewal requires.

Args:

n_chains: Chains per replica.

Raises:

InputValidationError: If the archive lacks the flagged models.

`method` `prepare_steering`

prepare_steering(
    steering: 'Steering',
    strength: 'float',
    proposal_steps: 'int | None' = None,
    target_acceptance: 'float' = 0.3,
    local_sweeps: 'int' = 1,
    trial_rounds: 'int' = 20,
    max_rungs: 'int' = 64,
    is_cancelled: 'Callable[[], bool] | None' = None,
    on_progress: 'Callable[, None] | None' = None
) → list[float]

Extend the generation ladder with steered rungs up to strength.

Every new rung has the endpoint's DCA parameters plus the steering potential at a strength between 0 and strength; the last one is at strength and becomes the sampled endpoint. Rungs are placed one at a time: a candidate strength is kept when, after trial_rounds rounds of warmup and trial_rounds of measurement, the swap acceptance with the rung below is at least target_acceptance; otherwise the step is halved. The bottom of the ladder, the exact profile, is unchanged, so population renewal still certifies mixing. Proposal blocks of the steered rungs adapt during placement and are then frozen.

Call after prepare_sampling_ladder. A steered ladder cannot be saved with save_archive.

Args:

steering: The validated potential.
strength: Final steering strength.
proposal_steps: Site updates per steered proposal, or None to adapt.
target_acceptance: Smallest swap acceptance between adjacent steered rungs.
local_sweeps: Local sweeps per exchange round.
trial_rounds: Rounds of warmup, then of measurement, per candidate rung.
max_rungs: Largest number of steered rungs.
is_cancelled: Optional callable; returning True stops the run.
on_progress: Optional callback on_progress(stage, done, total, **details).

Returns: The strengths of the steered rungs, increasing to strength.

Raises:

InputValidationError: If the ladder is not a generation ladder, the potential is invalid, or strength is 0.
ConvergenceError: If max_rungs rungs do not reach strength.

`method` `refresh_reservoir`

refresh_reservoir(
    local_sweeps: 'int' = 1,
    is_cancelled: 'Callable[[], bool] | None' = None,
    on_progress: 'Callable[, None] | None' = None
) → bool

Training: collect a new reservoir from the ladder and shorten the active ladder.

Returns: False if the ladder failed the mixing check needed to collect it.

`method` `restore_state`

restore_state(state: 'Mapping[str, Any]') → None

Restore a state returned by capture_state, clearing the training state.

`method` `save_archive`

save_archive(path: 'str | Path') → Path

Write the complete sampler to an HDF5 archive, replacing path atomically.

The archive is validated before it replaces an existing file, so an interrupted save never leaves a broken archive.

Args:

path: Archive path, usually ptt.h5.

Returns: The written path.

`method` `set_coupling_dtype`

set_coupling_dtype(dtype: 'dtype | None' = None) → None

Select the precision of the couplings read by the local kernels.

torch.bfloat16 makes the local sweeps read a BF16 copy of each replica's couplings, refreshed whenever the couplings change; biases, fields, acceptance, exchanges and every stored state stay in the parameter precision. This reduces memory traffic, which speeds up large protein models. The setting is not archived. None restores full precision.

Args:

dtype: torch.bfloat16 or None.

Raises:

InputValidationError: If BF16 is requested without float32 parameters on an Ampere-or-newer CUDA GPU with Triton.

`method` `set_generation_kernel`

set_generation_kernel(name: 'str | None' = None) → None

Select the local kernel used during generation; None restores the archived one.

Every kernel leaves each replica's distribution invariant, so a model trained with one kernel can be sampled with another. Archives keep the training kernel as local_sampler.

Args:

name: "metropolized_gibbs", "gibbs", "metropolis" or None.

`method` `start_birth_tracking`

start_birth_tracking() → None

Mark every current configuration as born before the next round.

From then on advance stamps configurations entering the active ladder with the current round: exact profile draws at rung 0 in a full ladder, or configurations emitted by the reservoir otherwise. Exchanges and permutations carry birth rounds with configurations.

`method` `steering_log_z_ratio`

steering_log_z_ratio(estimator: 'str' = 'bar') → float

Estimate log Z_s - log Z_0: steered top rung against the unsteered endpoint.

-log_z_ratio is the free-energy cost of the steering, and exp(log_z_ratio) equals <exp(-V(x, s))> under the unsteered model.

Args:

estimator: "bar" or "forward" bridges.

Returns: The log ratio of partition functions.

`method` `transition_target`

transition_target(
    params: 'Mapping[str, Tensor]',
    local_sweeps: 'int' = 1,
    is_cancelled: 'Callable[[], bool] | None' = None,
    on_progress: 'Callable[, None] | None' = None
) → tuple[bool, int]

Training: move the endpoint to params, then run the snapshot procedure.

The transition runs on a fork and is committed only if the ladder stays healthy; otherwise last_failure describes the mixing failure.

Args:

params: New endpoint parameters, same shape, precision and device.
local_sweeps: Local sweeps per exchange round.
is_cancelled: Optional callable; returning True stops the run.
on_progress: Optional progress callback.

Returns: A pair (accepted, work), whether the transition was committed and the local sweeps it used.

`method` `update_replica_chain`

update_replica_chain(
    local_sweeps: 'int' = 1,
    is_cancelled: 'Callable[[], bool] | None' = None,
    on_progress: 'Callable[, None] | None' = None
) → bool

Training: run the snapshot procedure after an endpoint update.

When the endpoint's swap acceptance falls below 2 * target_acceptance a snapshot of it is stored; below target_acceptance the stored snapshot is inserted into the ladder (or a held one replaces the replica below).

Returns: False if a required reservoir refresh failed its mixing check.

This file was automatically generated via lazydocs.

module adabmDCA.ptt.sampler

Global Variables

function bridge_increment

class PartitionEstimate

method __init__

class PTTSampler

method __init__

property n_active

method advance

method capture_state

method draw

method endpoint_params

method endpoint_samples

method entropy

method equilibrate

method equilibrate_for_lag

method estimate_mixing_time

method force_ladder_update

method fork

classmethod from_archive

method ladder_health

method ladder_statistics

method measure_renewal

method partition_estimate

method prepare_sampling_ladder

method prepare_steering

method refresh_reservoir

method restore_state

method save_archive

method set_coupling_dtype

method set_generation_kernel

method start_birth_tracking

method steering_log_z_ratio

method transition_target

method update_replica_chain

`module` `adabmDCA.ptt.sampler`

`function` `bridge_increment`

`class` `PartitionEstimate`

`method` `init`

`class` `PTTSampler`

`method` `init`

`property` n_active

`method` `advance`

`method` `capture_state`

`method` `draw`

`method` `endpoint_params`

`method` `endpoint_samples`

`method` `entropy`

`method` `equilibrate`

`method` `equilibrate_for_lag`

`method` `estimate_mixing_time`

`method` `force_ladder_update`

`method` `fork`

`classmethod` `from_archive`

`method` `ladder_health`

`method` `ladder_statistics`

`method` `measure_renewal`

`method` `partition_estimate`

`method` `prepare_sampling_ladder`

`method` `prepare_steering`

`method` `refresh_reservoir`

`method` `restore_state`

`method` `save_archive`

`method` `set_coupling_dtype`

`method` `set_generation_kernel`

`method` `start_birth_tracking`

`method` `steering_log_z_ratio`

`method` `transition_target`

`method` `update_replica_chain`