xnn.hybrid.models.bamboo#

BAMBOO: a graph equivariant transformer force field (Gong et al. 2024).

A faithful, self-contained re-implementation of the BAMBOO model (ByteDance AI Molecular Simulation Booster, arXiv:2404.07181) built on the xnn abstractions. BAMBOO is a hybrid potential: a graph neural network whose message-passing layers are transformers (the Graph Equivariant Transformer, GET), followed by a physics-based split of the atomic energy into three pieces (Supplementary A.2):

E_i = E_i^NN  +  E_i^elec  +  E_i^disp

  • semi-local E^NN – an MLP on the per-atom GET features;

  • electrostatic E^elec – a charge-equilibrium energy built from predicted partial charges (a per-atom electronegativity/hardness term plus a damped Coulomb sum over all pairs);

  • dispersion E^disp – an optional D3(CSO) correction (off by default, matching the paper, which excludes dispersion from the DFT training data and only adds it during MD).

Architecture (Supplementary A.1): the atom type Z is embedded to the scalar node feature x_i while the vector node feature V_i starts at zero. Each GET layer runs a multi-head QKV attention on the neighbour graph (the shared EdgeMultiheadAttention), scales the neighbour values by a radial edge feature (from the ExpNormalSmearing basis) and the attention weight, and mixes the scalar and vector channels through inner products so both stay rotation-equivariant. Two MLPs read the final scalar features into the per-atom energy and the partial charge.

The model subclasses InteratomicPotential directly (like PhysNet, the other physics-split potential in xnn): it needs no e3nn because equivariance comes from Cartesian vector channels, not spherical harmonics. Forces and stress are added uniformly by ForceStressOutput via autograd – the paper’s separately damped Coulomb force is, on inspection, exactly the gradient of its Coulomb energy (the softplus energy damping differentiates to the sigmoid force damping), so autograd reproduces it.

Given the same weights this matches the original bamboo package (bytedance/bamboo) to machine precision – see tests/test_bamboo.py and examples/fidelity_checks/bamboo_verification.ipynb.

Module Attributes

ELE_FACTOR

Coulomb prefactor k_e * e^2 in kcal/mol * Angstrom / e^2.

DEBYE_EA

Dipole conversion (e * Angstrom -> Debye is the reciprocal).

Classes

BAMBOO([dim, num_rbf, cutoff, n_layers, ...])

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

GETLayer(dim, num_heads, attn_act[, ...])

One Graph Equivariant Transformer layer (BAMBOO Supplementary A.1).
Show Source