Ising Models

Ising-specific model, program, and training utilities.

For most problems ising_sample is the entry point: it goes from (biases, edges, weights) to samples in one call, building the coloring, autotuning the ladder, and drawing from the target. IsingEBM is the lower-level construction — the factor graph from a list of nodes, edges, biases, and coupling weights — for when you want to drive the sampler yourself.

Front door

hamon.ising_sample(biases: Shaped[Array, n], edges: Shaped[Array, 'm 2'], weights: Shaped[Array, m], *, key: Key[Array, ''], beta: float | str = 1.0, n_samples: int = 1000, n_warmup: int = 500, steps_per_sample: int = 1, target_acceptance: float = 0.5, max_chains: int = 128, device: str | jaxlib._jax.Device | None = 'auto') -> tuple[Bool[Array, 'n_samples n'], dict]

Sample from an Ising model Boltzmann distribution via fully autotuned NRPT.

A thin Ising-specific front end over :func:hamon.autosample: it builds and colors the graph, then autotunes the full NRPT configuration — chain count, local-exploration count (gibbs_steps_per_round), and schedule — before drawing from the cold chain. Unlike earlier versions, the exploration count is no longer a fixed argument; it is discovered (and device-calibrated) automatically.

A warning is logged if all coupling weights are zero (NRPT is unnecessary) or if all biases are identical (the model has no per-variable preference).

The energy function is

\[\mathcal{E}(s) = -\beta \left( \sum_i b_i s_i + \sum_{(i,j)} J_{ij} s_i s_j \right)\]

Parameters:

Name	Type	Description	Default
biases	Shaped[Array, n]	per-node bias array of shape (n,).	required
edges	Shaped[Array, 'm 2']	integer index pairs of shape (m, 2).	required
weights	Shaped[Array, m]	per-edge coupling of shape (m,).	required
key	Key[Array, '']	JAX PRNG key.	required
beta	float | str	inverse temperature for the target distribution, or "auto" to choose it for ground-state search: the excitation-cost spectrum of the landscape (exact on field-free forests, greedy-descent probe elsewhere) picks the smallest β whose predicted equilibrium excess energy is ≤ 0.1% of the ground-state scale — see :func:ising_estimate_beta. The estimate and its rationale are returned under diagnostics["beta_estimate"].	1.0
n_samples	int	number of samples to return.	1000
n_warmup	int	warmup steps before collecting samples.	500
steps_per_sample	int	Gibbs sweeps between recorded samples.	1
target_acceptance	float	desired per-pair swap acceptance rate for the chain-count search. Default 0.5 — the round-trip-optimal r = 1/2 (N ≈ 2Λ; Syed et al.).	0.5
max_chains	int	ceiling on the discovered chain count.	128
device	str | Device | None	where to run — "auto" (default), "cpu"/"gpu", a concrete jax.Device, or None to leave placement untouched. Resolved once and reused across all autotuning stages; the measured wall time on this device calibrates the chosen gibbs_steps_per_round.	'auto'

Returns:

Type	Description
Bool[Array, 'n_samples n']	A tuple (samples, diagnostics) where samples is a boolean array of
dict	shape (n_samples, n) (True = spin up) and diagnostics is a dict
tuple[Bool[Array, 'n_samples n'], dict]	with keys n_chains, betas, Lambda, gibbs_steps_per_round,
tuple[Bool[Array, 'n_samples n'], dict]	mean_spins (average number of +1 spins per sample), device,
tuple[Bool[Array, 'n_samples n'], dict]	round_trip_diagnostics, and report (the full
tuple[Bool[Array, 'n_samples n'], dict]	class:hamon.AutotuneReport).

Choosing β from the model

beta="auto" reads the coldest useful temperature off the model's own excitation-cost spectrum instead of guessing. Both halves of that estimator are available on their own.

hamon.ising_estimate_beta(biases, edges, weights, *, gap_tol: float = 0.001, n_replicas: int = 64, seed: int = 0)

Estimate the coldest useful β for ground-state search on an Ising model.

A thin front end over :func:hamon.estimate_beta_max: extracts the excitation-cost spectrum (exact on field-free forests, greedy-descent probe elsewhere) and selects the smallest β whose predicted equilibrium excess energy is at most gap_tol of the ground-state scale. Runs on the host in milliseconds — no tuning, no compiles. Returns a :class:hamon.BetaEstimate.

hamon.ising_excitation_costs(biases, edges, weights, *, n_replicas: int = 64, seed: int = 0) -> tuple[numpy.ndarray, float, str]

Elementary excitation-cost spectrum of an Ising landscape.

Returns (costs, energy_scale, method). Field-free forests are exact: bond defects are independent with cost 2|J| and |E_GS| = Σ|J|. Anything else uses the greedy-descent probe (2|local field| per-site minima across replicas, energy scale = best minimum found). On highly regular graphs the probe minima can have a zero local field at every site (a genuine degeneracy), leaving no positive costs; there we fall back to the coupling-magnitude spectrum 2|J|.

Model construction

hamon.models.IsingEBM

An EBM with the energy function,

\[\mathcal{E}(s) = -\beta \left( \sum_{i \in S_1} b_i s_i + \sum_{(i, j) \in S_2} J_{ij} s_i s_j \right)\]

where \(S_1\) and \(S_2\) are the sets of biases and weights that make up the model, respectively. \(b_i\) represents the bias associated with the spin \(s_i\) and \(J_{ij}\) is a weight that couples \(s_i\) and \(s_j\). \(\beta\) is the usual temperature parameter.

Attributes:

• nodes: the nodes that have an associated bias (i.e \(S_1\))
• biases: the bias associated with each node in nodes.
• edges: the edges that have an associated weight (i.e \(S_2\))
• weights: the weight associated with each pair of nodes in edges.
• beta: the scalar temperature parameter for the model.

nodes and edges are stored as immutable, identity-hashed sequences (a single pytree leaf each): the EBM is passed to jitted functions (hinton_init, the NRPT round loop), and flattening plain lists would visit and hash every node and edge endpoint — O(|graph|) host work — on every call. They still index, iterate, and len() like lists; with_beta passes the same objects through, which is what keeps the jit cache hitting.

__init__(nodes: collections.abc.Sequence[hamon.pgm.AbstractNode], edges: collections.abc.Sequence[tuple[hamon.pgm.AbstractNode, hamon.pgm.AbstractNode]], biases: Array, weights: Array, beta: Array)

hamon.models.IsingSamplingProgram

Thin wrapper specializing :class:ModelSamplingProgram to an Ising model.

__init__(ebm: IsingEBM, free_blocks: list[tuple[hamon.block_management.Block, ...] | hamon.block_management.Block], clamped_blocks: list[hamon.block_management.Block], *, _gibbs_spec: hamon.block_sampling.BlockGibbsSpec | None = None)

hamon.models.hinton_init(key: Key[Array, ''], model: IsingEBM, blocks: list[hamon.block_management.Block[hamon.pgm.AbstractNode]], batch_shape: tuple[int, ...]) -> list[Bool[Array, 'batch_size block_size']]

Initialize the blocks according to the marginal bias.

Each binary unit \(i\) in a block is sampled independently as

\[\mathbb{P}(S_i = 1) = \sigma(\beta h_i) = \frac{1}{1 + e^{-\beta h_i}}\]

where \(h_i\) is the bias of unit i and \(\beta\) is the inverse-temperature scaling factor. See Hinton (2012) for a discussion of this initialization heuristic.

Units are drawn independently across all blocks at once; blocks may have different sizes.

Parameters:

Name	Type	Description	Default
key	Key[Array, '']	the JAX PRNG key to use	required
model	IsingEBM	the Ising model to initialize for	required
blocks	list[Block[AbstractNode]]	the blocks that are to be initialized	required
batch_shape	tuple[int, ...]	the pre-pended batch dimension	required

Returns:

Type	Description
list[Bool[Array, 'batch_size block_size']]	the initialized blocks as a list of bool arrays, one per block

Training

estimate_kl_grad computes the contrastive-divergence gradient of the KL objective — the positive phase clamped to data, the negative phase free — for an IsingTrainingSpec that pairs the model with its two sampling programs. Pass return_negative_state=True and feed the returned state into the next step for persistent contrastive divergence.

hamon.models.IsingTrainingSpec

Contains a complete specification of an Ising EBM that can be trained using sampling-based gradients.

Defines sampling programs and schedules that allow for collection of the positive and negative phase samples required for Monte Carlo estimation of the gradient of the KL-divergence between the model and a data distribution.

__init__(ebm: IsingEBM, data_blocks: list[hamon.block_management.Block], conditioning_blocks: list[hamon.block_management.Block], positive_sampling_blocks: list[tuple[hamon.block_management.Block, ...] | hamon.block_management.Block], negative_sampling_blocks: list[tuple[hamon.block_management.Block, ...] | hamon.block_management.Block], schedule_positive: SamplingSchedule, schedule_negative: SamplingSchedule)

hamon.models.estimate_moments(key: Key[Array, ''], first_moment_nodes: list[hamon.pgm.AbstractNode], second_moment_edges: list[tuple[hamon.pgm.AbstractNode, hamon.pgm.AbstractNode]], program: BlockSamplingProgram, schedule: SamplingSchedule, init_state: list[Array], clamped_data: list[Array], *, return_state: bool = False, device: str | jaxlib._jax.Device | None = 'auto')

Estimates the first and second moments of an Ising model Boltzmann distribution via sampling.

Parameters:

Name	Type	Description	Default
key	Key[Array, '']	the jax PRNG key	required
first_moment_nodes	list[AbstractNode]	the nodes that represent the variables we want to estimate the first moments of	required
second_moment_edges	list[tuple[AbstractNode, AbstractNode]]	the edges that connect the variables we want to estimate the second moments of	required
program	BlockSamplingProgram	the BlockSamplingProgram to be used for sampling	required
schedule	SamplingSchedule	the schedule to use for sampling	required
init_state	list[Array]	the variable values to use to initialize the sampling	required
clamped_data	list[Array]	the variable values to assign to the clamped nodes	required
return_state	bool	when True, also return the final free-block chain state (the state at the last recorded sample), so callers can continue the chain — e.g. persistent-chain (PCD) training.	False

Returns: the first and second moment data, plus the final chain state when return_state is set.

hamon.models.estimate_kl_grad(key: Key[Array, ''], training_spec: IsingTrainingSpec, bias_nodes: list[hamon.pgm.AbstractNode], weight_edges: list[tuple[hamon.pgm.AbstractNode, hamon.pgm.AbstractNode]], data: list[Array], conditioning_values: list[Array], init_state_positive: list[Array], init_state_negative: list[Array], *, return_negative_state: bool = False, device: str | jaxlib._jax.Device | None = 'auto') -> tuple

Estimate the KL-gradients of an Ising model with respect to its weights and biases.

Uses the standard two-term Monte Carlo estimator of the gradient of the KL-divergence between an Ising model and a data distribution.

The gradients are:

\[\Delta W = -\beta (\langle s_i s_j \rangle_{+} - \langle s_i s_j \rangle_{-})\]

\[\Delta b = -\beta (\langle s_i \rangle_{+} - \langle s_i \rangle_{-})\]

Here, \(\langle\cdot\rangle_{+}\) denotes an expectation under the positive phase (data-clamped Boltzmann distribution) and \(\langle\cdot\rangle_{-}\) under the negative phase (model distribution).

Parameters:

Name	Type	Description	Default
key	Key[Array, '']	the JAX PRNG key	required
training_spec	IsingTrainingSpec	the Ising EBM for which to estimate the gradients	required
bias_nodes	list[AbstractNode]	the nodes for which to estimate the bias gradients	required
weight_edges	list[tuple[AbstractNode, AbstractNode]]	the edges for which to estimate the weight gradients	required
data	list[Array]	The data values to use for the positive phase of the gradient estimate. Each array has shape [batch nodes]	required
conditioning_values	list[Array]	values to assign to the nodes that the model is conditioned on. Each array has shape [nodes]	required
init_state_positive	list[Array]	initial state for the positive sampling chain. Each array has shape [n_chains_pos batch nodes]	required
init_state_negative	list[Array]	initial state for the negative sampling chain. Each array has shape [n_chains_neg nodes]	required
return_negative_state	bool	when True, append the negative chains' final states (same structure as init_state_negative) to the returned tuple. Feeding them back as the next step's init_state_negative gives persistent-chain (PCD) training: the chains track the slowly moving model distribution instead of re-warming from scratch every gradient step, so the negative schedule's n_warmup can drop to ~0. (The positive phase is clamped to per-batch data, so persisting it across batches is not meaningful and it is not returned.)	False

Returns: the weight gradients and the bias gradients (plus the final negative chain state when return_negative_state is set)