xnn.hybrid.models.bamboo.BAMBOO#

class xnn.hybrid.models.bamboo.BAMBOO(dim=64, num_rbf=32, cutoff=5.0, n_layers=3, num_heads=16, charge_ub=2.0, charge_mlp_layers=2, energy_mlp_layers=2, n_elements=87, act_fn='silu', attn_act_fn='gelu', coul_damping_beta=18.7, coul_damping_r0=2.2, ele_factor=332.06349451357806, use_electrostatics=True, use_dispersion=False, disp_cutoff=10.0, d3_references='2010', species=None)[source]#

Bases: InteratomicPotential

Graph equivariant transformer force field (Gong et al. 2024).

See the module docstring for the architecture. All hyper-parameters have the upstream defaults; the shared fields (cutoff, n_features, n_rbf, n_interactions) come from the core ModelConfig and the rest from its extra dict (see from_config()).

Parameters:

  • dim (int, optional) – Node scalar feature width (n_features), by default 64. Must be divisible by num_heads.

  • num_rbf (int, optional) – Number of exponential-normal radial basis functions (n_rbf), by default 32.

  • cutoff (float, optional) – Semi-local radial cutoff r_cut in Angstrom, by default 5.0.

  • n_layers (int, optional) – Number of GET layers (n_interactions); must be >= 2, by default 3.

  • num_heads (int, optional) – Number of attention heads, by default 16.

  • charge_ub (float, optional) – Upper bound of the per-atom partial charge (a tanh squashes the raw charge into [-charge_ub, charge_ub]), by default 2.0.

  • charge_mlp_layers (int, optional) – Hidden-layer counts of the charge and energy read-out MLPs, by default 2 each.

  • energy_mlp_layers (int, optional) – Hidden-layer counts of the charge and energy read-out MLPs, by default 2 each.

  • n_elements (int, optional) – Size of the atom-type embedding table (indexed directly by atomic number Z), by default 87 (H..Rn, upstream default).

  • act_fn (str or torch.nn.Module, optional) – Activation of the read-out MLPs (upstream SiLU), by default "silu".

  • attn_act_fn (str or torch.nn.Module, optional) – Activation of the attention and the radial-projection (upstream GELU), by default "gelu".

  • coul_damping_beta (float, optional) – Softplus sharpness of the short-range Coulomb damping, by default 18.7.

  • coul_damping_r0 (float, optional) – Onset distance of the Coulomb damping in Angstrom, by default 2.2.

  • ele_factor (float, optional) – Coulomb prefactor k_e e^2 in kcal/mol*Angstrom, by default ELE_FACTOR.

  • use_electrostatics (bool, optional) – Include the charge-equilibrium electrostatic energy, by default True.

  • use_dispersion (bool, optional) – Add the D3(CSO) dispersion energy, by default False (matching the paper’s training set).

  • disp_cutoff (float, optional) – Dispersion cutoff in Angstrom (used only when use_dispersion), by default 10.0.

  • d3_references (str, optional) – D3 reference systems of the D3(CSO) term, "2010" (default, upstream’s tables) or "2024" (current simple-dftd3 references); see D3CSODispersion.

  • species (list of int, optional) – Supported atomic numbers, recorded for bookkeeping/config round-trips (the embedding is indexed by Z regardless). Default None.

Variables:

node_feature_dim (int) – Width of the invariant per-atom features exposed as node_features (equal to dim), consumed e.g. by LatentEwald.

Raises:

ValueError – If n_layers < 2 or dim is not divisible by num_heads.

Notes

The built-in electrostatics is the molecular / gas-phase-cluster charge-equilibrium form: a damped Coulomb summed over all intra-structure pairs from the raw positions (as in the paper’s cluster training and the upstream predict path). It is therefore not minimum-imaged, so for genuinely periodic long-range electrostatics use the LAMMPS Ewald route (as the original does) or wrap the model with LatentEwald (BAMBOO exposes the required node_features); set use_electrostatics=False to drop the built-in term in that case.

coulomb_energy(charges, edge_vec, row, col)[source]#

Per-pair damped Coulomb energy (upstream get_coulomb, no Ewald).

The short-range damping softens the 1/r singularity smoothly:

E_ij = ele_factor * q_i q_j / r *  (r / r0) / (1 + softplus((r-r0)/r0))

Its gradient with respect to r equals the paper’s separately-damped Coulomb force (softplus differentiates to sigmoid), so autograd forces are the intended ones.

Parameters:

  • charges (Tensor) – Per-atom partial charges (N,).

  • edge_vec (Tensor) – Displacement vectors of the (all-pairs) list (P, 3).

  • row (Tensor) – The two atoms of each pair, shape (P,) each.

  • col (Tensor) – The two atoms of each pair, shape (P,) each.

Returns:

Per-pair Coulomb energy (P,).

Return type:

Tensor

forward(data)[source]#

Predict energies, partial charges and the dipole for a graph.

Parameters:

data (xnn.common.data.AtomicGraph) – The input atomic graph. If it carries a total_charge attribute (per-structure tensor) the partial charges are conserved to it; otherwise neutral structures are assumed.

Returns:

"node_energy" (N,) and "energy" (B,) (as every xnn model), plus "node_features" (N, dim) (invariant features), "charges" (N,) (conserved partial charges), "dipole" (B, 3) in Debye, and the component energies "energy_nn" / "energy_elec" (B,).

Return type:

dict of str to torch.Tensor

classmethod from_config(cfg)[source]#

Construct a BAMBOO from a core model config.

Core fields map as cfg.cutoff -> cutoff, cfg.n_features -> dim, cfg.n_rbf -> num_rbf and cfg.n_interactions -> n_layers; every other option is read from cfg.extra (upstream spellings are translated by xnn.common.config.translate).

Parameters:

cfg (xnn.common.config.schema.ModelConfig) – The core model config, whose extra dict carries the BAMBOO architecture options.

Returns:

The instantiated model.

Return type:

BAMBOO