Skip to content

Energy-Based Models#

Base classes for energy-based models. AbstractEBM defines the interface; AbstractFactorizedEBM adds factor-sum energies. AnnealedEBM implements the standard parallel-tempering path between a reference and a target.

Two properties on AbstractEBM drive how NRPT may temper a model. proper_at_beta_zero says whether the β = 0 member is a proper distribution — False for unbounded state spaces, which is why a ladder starting at exactly β = 0 is rejected for the continuous families. beta_affine says whether interactions interpolate as offset + β·slope rather than scaling linearly with β, which is what AnnealedEBM needs.

hamon.models.AbstractEBM #

Something that has a well-defined energy function (map from a state to a scalar).

energy(state: list[PyTree[Shaped[Array, 'nodes ?*state'], _State]], blocks: BlockSpec | list[Block]) -> Float[Array, ''] #

Evaluate the energy function of the EBM given some state information.

Arguments:

  • state: The state for which to evaluate the energy function. Must be compatible with blocks.
  • blocks: Specifies how the information in state is organized. May be either a pre-built BlockSpec (fast path — avoids rebuilding the spec) or a plain list[Block] for convenience when calling from user code.

Returns:

A scalar representing the energy value associated with state.

hamon.models.AbstractFactorizedEBM #

An EBM that is made up of Factors, i.e., an EBM with an energy function like,

\[\mathcal{E}(x) = \sum_i \mathcal{E}^i(x)\]

where the sum over \(i\) is taken over factors.

Child classes must define a property which returns a list of factors that substantiate the EBM.

Attributes:

  • node_shape_dtypes: the shape/dtypes of the nodes involved in this EBM. Used to generate the BlockSpec that defines the global state that factors receive to compute energy.
__init__(node_shape_dtypes: collections.abc.Mapping[type[hamon.pgm.AbstractNode], PyTree[jax.ShapeDtypeStruct]] = {<class 'hamon.pgm.SpinNode'>: ShapeDtypeStruct(shape=(), dtype=bool), <class 'hamon.pgm.CategoricalNode'>: ShapeDtypeStruct(shape=(), dtype=uint8), <class 'hamon.pgm.GaussianNode'>: ShapeDtypeStruct(shape=(), dtype=float32)}) #

hamon.models.FactorizedEBM #

An EBM that is defined by a concrete list of factors.

Attributes:

  • _factors: the list of factors that defines this EBM.
__init__(factors: list[hamon.models.ebm.EBMFactor], node_shape_dtypes: collections.abc.Mapping[type[hamon.pgm.AbstractNode], PyTree[jax.ShapeDtypeStruct]] = {<class 'hamon.pgm.SpinNode'>: ShapeDtypeStruct(shape=(), dtype=bool), <class 'hamon.pgm.CategoricalNode'>: ShapeDtypeStruct(shape=(), dtype=uint8), <class 'hamon.pgm.GaussianNode'>: ShapeDtypeStruct(shape=(), dtype=float32)}) #

hamon.models.EBMFactor #

A factor that defines an energy function.

__abstractclassvars__ class-attribute #

Build an immutable unordered collection of unique elements.

__abstractmethods__ class-attribute #

Build an immutable unordered collection of unique elements.

__abstractvars__ class-attribute #

Build an immutable unordered collection of unique elements.

__annotations__ class-attribute #

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_fields__ class-attribute #

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_params__ class-attribute #
__doc__ class-attribute #

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__firstlineno__ class-attribute #

int([x]) -> integer int(x, base=10) -> integer

Convert a number or string to an integer, or return 0 if no arguments are given. If x is a number, return x.int(). For floating-point numbers, this truncates towards zero.

If x is not a number or if base is given, then x must be a string, bytes, or bytearray instance representing an integer literal in the given base. The literal can be preceded by '+' or '-' and be surrounded by whitespace. The base defaults to 10. Valid bases are 0 and 2-36. Base 0 means to interpret the base from the string as an integer iteral.

int('0b100', base=0) 4

__match_args__ class-attribute #

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__module__ class-attribute #

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__static_attributes__ class-attribute #

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__init__(node_groups: list[hamon.block_management.Block]) #

Create a batch of Factors.

Practically, this just means writing down some parallel groups of nodes that the batch of Factors acts on. All of the functionality of the Factor is implemented by the method to_interaction_groups.

Arguments:

  • node_groups: The node groups that this batch of factors acts on. A single Factor is defined over node_groups[k][i] for all values of k and a particular batch index i.
to_interaction_groups() -> list[hamon.interaction.InteractionGroup] #

Compile a factor to a set of directed interactions.

energy(global_state: list[Array], block_spec: BlockSpec) -> Float[Array, ''] #

Evaluate the energy function of the factor.

Arguments:

  • global_state: The state information to use to evaluate the energy function. Is a global state of block_spec.
  • block_spec: The BlockSpec used to generate global_state.

Reference annealing#

AnnealedEBM(reference, target, β) implements E_β = (1−β)·E_ref + β·E_target, whose β = 0 member is the reference at full weight. Every rung of the ladder is then proper, so the ladder can cover the full entropic path even when the target alone has no proper β = 0 member. NRPT handles this with an affine interpolation and swap energies Δ = E₁ − E₀ — the shared reference term cancels in every swap ratio.

hamon.models.AnnealedEBM #

The reference-annealing path between two EBMs over the same nodes:

\[\mathcal{E}_\beta(x) = (1-\beta)\,\mathcal{E}_{\text{ref}}(x) + \beta\,\mathcal{E}_{\text{target}}(x) = \mathcal{E}_{\text{ref}} + \beta\,(\mathcal{E}_{\text{target}} - \mathcal{E}_{\text{ref}})\]

— the standard PT path from a reference distribution (β = 0) to the target (β = 1), rather than from the flat/uniform member of the target's own tempered family. Its point: an unbounded-state-space target has no proper β = 0 member (proper_at_beta_zero = False), but annealing from a proper reference — e.g. a diagonal Gaussian — makes every rung of the ladder proper, so β can start at exactly 0 and the ladder covers the full entropic path.

Both EBMs must be defined over the same nodes and be temperature-linear themselves (factors emit β-scaled coefficients, as all hamon models do): the annealed factors are simply the reference's at β' = 1−β plus the target's at β' = β. The sampling program must use a conditional that understands the union of both factor families (e.g. :class:~hamon.models.SliceGibbsConditional handles quartic + quadratic; an annealed pair of Gaussians is handled by :class:~hamon.models.GaussianGibbsConditional).

beta_affine is True: NRPT's template mode interpolates interactions affinely and computes swap energies as Δ = E_target − E_ref (the shared E_ref cancels in every swap ratio).

__init__(reference: AbstractFactorizedEBM, target: AbstractFactorizedEBM, beta) #