xnn.gnn.models.cace.CACE#

class xnn.gnn.models.cace.CACE(species, cutoff=5.5, n_atom_basis=3, n_rbf=8, n_radial_basis=None, max_l=3, max_nu=3, num_message_passing=1, message_types=('M', 'Ar', 'Bchi'), embed_receiver_nodes=False, avg_num_neighbors=10.0, num_polynomial_cutoff=6, trainable_rbf=True, readout_hidden=None, atomic_energies=None)[source]#

Bases: GNNPotential

Faithful CACE (Cheng 2024): Cartesian atomic cluster expansion.

Body-ordered invariant features built entirely in Cartesian coordinates (see the module docstring for the architecture walk-through), read out by a linear + MLP head into per-atom energies. With num_message_passing=0 this is exactly an (optimized-radial-coupling, element-embedded) ACE; each message-passing layer appends one more set of B features.

All architecture options are read from ModelConfig.extra (see from_config()); upstream constructor spellings are translated to the xnn names at config-load time by xnn.common.config.translate.

Parameters:

  • species (list of int) – Atomic numbers of the supported elements, in channel order (upstream zs).

  • cutoff (float, optional) – Radial cutoff in Angstrom, by default 5.5 (the paper’s water model).

  • n_atom_basis (int, optional) – Length of the learnable element embedding theta (paper N_embedding, typically 1-4), by default 3. The number of edge channels is n_atom_basis**2.

  • n_rbf (int, optional) – Number of raw Bessel radial functions, by default 8.

  • n_radial_basis (int or None, optional) – Mixed radial channels after the learned coupling (paper n); None (default) keeps n_rbf.

  • max_l (int, optional) – Maximum total angular momentum of the Cartesian basis, by default 3.

  • max_nu (int, optional) – Maximum body order of the invariant B features (1-4), by default 3.

  • num_message_passing (int, optional) – Number of message-passing layers T (0 = plain Cartesian ACE), by default 1.

  • message_types (sequence of str, optional) – Enabled message mechanisms per layer, subset of ("M", "Ar", "Bchi") (node memory, radial-filter message, recursive edge embedding), by default all three (the upstream default; the paper’s water model uses ("Bchi",)).

  • embed_receiver_nodes (bool, optional) – Use a separate embedding table for receiver atoms in the edge type theta_i (x) theta_j (upstream flag), by default False (sender table shared, as in the upstream constructor default).

  • avg_num_neighbors (float, optional) – Message normalization 1/sqrt(avg_num_neighbors), by default 10.0.

  • num_polynomial_cutoff (int, optional) – Degree p of the polynomial cutoff envelope, by default 6.

  • trainable_rbf (bool, optional) – Learnable Bessel frequencies, by default True (the paper’s “trainable Bessel functions”).

  • readout_hidden (list of int, optional) – Hidden widths of the readout MLP, by default [32, 16] (the upstream example). The readout is the sum of this MLP and a parallel linear layer (paper eq 15).

  • atomic_energies (array-like or None, optional) – Per-species reference energies folded into atom_ref (upstream subtracts them from the training data instead).

node_features_energy(atomic_numbers, edge_index, edge_vec)[source]#

Tensor core: invariant features and per-atom energies.

Parameters:

  • atomic_numbers (Tensor) – Per-atom atomic numbers, shape (N,).

  • edge_index (Tensor) – Edge index (2, E); row 0 is the sender (neighbor), row 1 the receiver (center) – the upstream CACE/MACE convention.

  • edge_vec (Tensor) – Edge vectors pos[receiver] - pos[sender], shape (E, 3).

Returns:

The concatenated invariant B features (N, node_feature_dim) (what the readout consumes, and what LatentEwald maps to latent charges) and the per-atom energies (N,).

Return type:

tuple of Tensor

node_energy(atomic_numbers, edge_index, edge_vec)[source]#

Per-atom energies from raw graph tensors (thin wrapper over node_features_energy()).

Parameters:
Return type:

Tensor

forward(data)[source]#

Predict per-node and total energy for an atomic graph.

Parameters:

data (xnn.common.data.AtomicGraph) – The input atomic graph.

Returns:

"node_energy" (per-atom energies), "energy" (per-structure totals) and "node_features" (the invariant B features, shape (N, node_feature_dim)).

Return type:

dict of str to torch.Tensor

classmethod from_config(cfg)[source]#

Construct a CACE from a core model config.

Reads the CACE hyper-parameters from cfg.extra; upstream spellings are translated by xnn.common.config.translate and value forms coerced by xnn.common.config.coerce. Note that cfg.n_features is not used – CACE’s feature width is set by n_atom_basis / n_radial_basis / max_l / max_nu.

Parameters:

cfg (xnn.common.config.schema.ModelConfig) – The core model config.

Returns:

The instantiated model.

Return type:

CACE