Round-Trip Diagnostics#
Functions for measuring tempering performance: round-trip rates, the communication barrier Λ, chain count recommendations, and the log normalizing constant via thermodynamic integration.
hamon.round_trip_summary(index_state: dict, rejection_rates: jax.Array, betas: jax.Array, n_rounds: int) -> dict
#
Compute full diagnostic summary for NRPT run.
Jitted so the handful of reductions below (Λ, τ̄, the local-barrier profile,
the round-trip rate) fuse into a single compiled kernel instead of ~8 eager
op-by-op dispatches, each of which otherwise pays a first-shape XLA compile
when called once per probe at a new chain count. n_rounds is traced (not
static), so the compile is shared across round counts.
Returns dict with
Lambda: global communication barrier estimate tau_predicted: theoretical optimal round trip rate tau_observed: empirical round trip rate efficiency: tau_observed / tau_predicted (closer to 1 = better) lambda_profile: local barrier at each pair midpoint round_trips_per_chain: per-machine round trip counts restarts_per_chain: per-machine restart counts
hamon.recommend_n_chains(Lambda: float | jax.Array, target_acceptance: float = 0.5) -> int
#
Suggest chain count for a given barrier and target acceptance rate.
For NRPT with equalized rejection rates: Nr ≈ Λ where r = 1 - target_acceptance. Solving: N = Λ / r* = Λ / (1 - target_acceptance).
The default target_acceptance=0.5 (r = 1/2 ⇒ N ≈ 2Λ) is the round-trip- optimal rejection rate from Syed et al., not the 0.77 of reversible PT.
Note: Λ from too few chains is biased low. Use tune_chains
for iterative bootstrapping if recommendations keep increasing.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
Lambda
|
float | Array
|
estimated global communication barrier |
required |
target_acceptance
|
float
|
desired per-pair acceptance rate (default: 0.5 = 50%) |
0.5
|
Returns:
| Type | Description |
|---|---|
int
|
Recommended number of chains (minimum 2). |
hamon.thermodynamic_integration(betas: jax.Array, mean_energies: jax.Array, *, method: str = 'trapezoid') -> jax.Array
#
Log normalizing-constant ratio via thermodynamic integration.
Estimates log Z(β_max) / Z(β_min) = -∫ μ(β) dβ (Syed et al. 2021,
Sec. 5.5), where μ(β) = E_{π^(β)}[V] is the mean base energy and the
integral runs over the supplied ladder [β_min, β_max]. mean_energies
are the per-chain means μ(β_i) — accumulate them with
hamon.NRPTEnergyObserver and divide its (sum_E, count) carry, or
use :func:nrpt_log_normalizing_constant.
The reference chain (β_min, typically 0) has a known normalizer: for a
discrete model with a uniform β=0 reference over M configurations,
log Z(β_min) = log M (e.g. n·log 2 for n spins), so the absolute
log Z(β_max) is this result plus that constant.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
betas
|
Array
|
ascending β ladder, shape |
required |
mean_energies
|
Array
|
per-chain mean base energy μ(β_i), shape |
required |
method
|
str
|
|
'trapezoid'
|
Returns:
| Type | Description |
|---|---|
Array
|
Scalar |
hamon.nrpt_log_normalizing_constant(stats: dict, *, log_z0: float = 0.0, method: str = 'trapezoid') -> jax.Array
#
Log normalizing constant from an NRPT run with an energy observer.
Convenience over :func:thermodynamic_integration: reads the
(sum_E, count) carry left by hamon.NRPTEnergyObserver in
stats["observer_carry"], forms the per-chain mean energies, and
integrates them against stats["betas"].
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
stats
|
dict
|
the stats dict from |
required |
log_z0
|
float
|
|
0.0
|
method
|
str
|
quadrature rule, see :func: |
'trapezoid'
|
Returns:
| Type | Description |
|---|---|
Array
|
Scalar |