Observers

Observers collect statistics during sampling. StateObserver records raw states; MomentAccumulatorObserver computes running means and variances without storing every sample.

hamon.AbstractObserver

Interface for objects that inspect the sampling program while it is running.

A concrete Observer is called once per block-sampling iteration and can maintain an arbitrary "carry" state across calls (e.g. running averages, histogram buffers, log-probs, etc.).

init() -> PyTree

Initialize the memory for the observer. Defaults to None.

hamon.StateObserver

Observer which logs the raw state of some set of nodes.

This observer is stateless: its carry is always None and iteration is ignored.

Attributes:

• blocks_to_sample: the list of Blocks which the states are logged for

__init__(blocks_to_sample: list[hamon.block_management.Block]) -> None

Initialize self. See help(type(self)) for accurate signature.

hamon.MomentAccumulatorObserver

Observer that accumulates and updates the provided moments.

It doesn't log any samples, and will only accumulate moments. Note that this observer does not scale the accumulated values by the number of times it was called. It simply records a running sum of a product of some state variables,

\[\sum_i f(x_1^i) f(x_2^i) \dots f(x_N^i)\]

Attributes:

• blocks_to_sample: the blocks to accumulate the moments over. These are for constructing the final state, and aren't truly "blocks" in the algorithmic sense (they can be connected to each other). There is one block per node type.
• flat_nodes_list: a list of all of the nodes in the moments (each occurring only once, so len(set(x)) = len(x)).
• flat_to_type_slices_list: a list over node types in which each element is an array of indices of the flat_node_list which that type corresponds to
• flat_to_full_moment_slices: a list over moment types in which each element is a 2D array, which matches the shape of the moment_spec[i] and of which each element is the index in the flat_node_list.
• f_transform: the element-wise transformation \(f\) to apply to sample values before accumulation.
• _flat_scatter_index: precomputed concatenation of all flat_to_type_slices_list arrays, used to build flat_state in a single scatter call.
• _flat_scatter_sizes: number of entries contributed by each node type, used to split the concatenated sampled state before scattering.
• _flat_value_order: precomputed argsort(_flat_scatter_index); used in __call__ to permute the concatenated sampled values into flat-node order without allocating a zeros array.
• _accumulate_dtype: dtype for the accumulator, fixed at construction time.

__init__(moment_spec: collections.abc.Sequence[collections.abc.Sequence[collections.abc.Sequence[hamon.pgm.AbstractNode]]], f_transform: Callable = _f_identity, dtype: dtype = float32)

Create a MomentAccumulatorObserver.

Arguments:

• moment_spec: A 3 depth sequence. The first is a sequence over different moment types. A given moment type should have the same number of nodes in each moment. Then for each moment type, there is a sequence over moments. Each given moment is defined by a certain set of nodes.

  For example, to get the first and second moments on a simple o-o graph:

  [ [(node1,), (node2,)], [(node1, node2)] ]

• f_transform: A function that takes in (state, blocks) and returns something with the same structure as state. Defines a transformation \(y=f(x)\) so accumulated moments are \(\langle f(x_1) f(x_2) \rangle\).

• dtype: Accumulator dtype, fixed at construction. Defaults to jnp.float32. Use jnp.float64 for double-precision models. Fixing this here avoids a per-step cast inside the scan body.

NRPT observers

Observers called once per NRPT round (after the Gibbs sweeps and swaps). NRPTStateObserver records chain states; NRPTEnergyObserver accumulates the per-chain mean energy used for the log normalizing constant (thermodynamic_integration).

hamon.AbstractNRPTObserver

Observer for NRPT rounds, called once per round after Gibbs sweeps and swaps.

Concrete subclasses must implement __call__ and may override init to provide a non-trivial carry. The observation returned by __call__ is stacked by lax.scan into a pytree with a leading axis of size n_rounds. Return None as the observation for accumulate-only observers that do not need per-round storage.

init() -> PyTree

Initialize the observer carry. Defaults to None.

hamon.NRPTStateObserver

Collect raw chain states at specified chain indices each round.

This observer is stateless (carry is always None). The returned observation is a list of arrays — one per free block — each of shape (len(chain_indices), ...). After lax.scan stacking the leading axis becomes n_rounds.

Attributes:

• chain_indices: Tuple of chain positions to record. Use (-1,) to collect only the cold chain (the default).

__init__(chain_indices: tuple[int, ...] = (-1,))

hamon.NRPTEnergyObserver

Accumulate per-chain mean base energy μ(β_i) = E[V^(β_i)] for thermodynamic integration of the log normalizing constant.

Each round the NRPT loop hands every observer the post-swap base_energies (shape (n_chains,), aligned to chain/β positions). Under stationarity these are samples of V^(β_i), so a running mean estimates μ(β_i), which integrates to log Z via :func:hamon.round_trips.thermodynamic_integration.

Accumulate-only: it returns None as the per-round observation, so it adds no per-round output stack — only a tiny carry (sum_E, count). Read the mean energies after a run from stats["observer_carry"]::

obs = NRPTEnergyObserver(n_chains)
states, stats = tune_schedule(..., observer=obs)
sum_E, count = stats["observer_carry"]
mean_energies = sum_E / count

or use the one-call :func:hamon.round_trips.nrpt_log_normalizing_constant.

.. note:: Attaching any observer (this one included) switches the NRPT round loop from the dynamic-trip-count lax.fori_loop fast path to lax.scan, which compiles once per distinct n_rounds. The default no-observer path is unaffected. In tune_schedule the observer is attached only to the production run, so accumulation is naturally post-tuning (no burn-in from the tuning phases is included).

Attributes:

• n_chains: number of chains in the ladder (sets the carry shape).
• _dtype: accumulator dtype, fixed at construction.

__init__(n_chains: int, dtype: dtype = float32)

Create an energy observer.

Arguments:

• n_chains: the number of chains (len(betas)).
• dtype: accumulator dtype, fixed at construction. Defaults to jnp.float32; use jnp.float64 for double-precision models.

hamon.nrpt_node_samples(observations: list[Array], program: BlockSamplingProgram, nodes: collections.abc.Sequence[hamon.pgm.AbstractNode], chain_index: int = 0) -> Array

Reorder NRPT observer output into node order.

stats["observations"] from hamon.nrpt with an hamon.NRPTStateObserver is a list of per-free-block arrays of shape (n_rounds, len(chain_indices), block_len, ...) — block-local layout, in free-block order. Assembling per-node samples from that requires concatenating blocks and inverting the block→node permutation, which is easy to get silently wrong (a forgotten inversion produces plausible-looking but scrambled samples). This helper does it once, correctly:

obs = NRPTStateObserver(chain_indices=(-1,))
states, stats = nrpt(..., observer=obs)
samples = nrpt_node_samples(stats["observations"], program, nodes)
# samples[r, i] is the state of nodes[i] at round r — guaranteed.

Arguments:

• observations: stats["observations"] — one array per free block.
• program: The sampling program the observations came from (any of the per-chain programs, or the template program in temperature-linear mode; only its gibbs_spec is used).
• nodes: The nodes, in the order you want the output columns. All must share one node type and belong to free (not clamped) blocks.
• chain_index: Which entry of the observer's chain_indices tuple to extract (default 0). Note this indexes the recorded chains, not the temperature ladder: for NRPTStateObserver(chain_indices=(0, -1)), the cold chain is chain_index=1.

Returns:

Array of shape (n_rounds, len(nodes), ...) with column i holding the state of nodes[i].