Layer Stack#

Enerzyme builds models as ordered layers rather than monolithic classes. The architecture name selects a Core; pre- and post-core layers handle geometry, embeddings, physics, and output reductions.

Registered layers#

From enerzyme/models/layers/:

Geometry

DistanceLayer, RangeSeparationLayer, RadiusGraphLayer

Radial basis

GaussianSmearing, ExponentialGaussianRBFLayer, ExponentialBernsteinRBFLayer, BesselRBFLayer, BernsteinRBFLayer, SincRBFLayer

Embeddings

RandomAtomEmbedding, NuclearEmbedding, ElectronicEmbedding, ScalarDenseEmbedding, GatherAtomEmbedding

Core

Architecture-specific message passing (Core with architecture in Modelhub)

Physics / post-processing

AtomicAffine, ChargeConservation, ElectrostaticEnergy, AtomicCharge2Dipole, GrimmeD3Energy, GrimmeD4Energy, ZBLRepulsionEnergy

Output

EnergyReduce, Force, ShallowEnsembleReduce

Typical charge-aware stack#

layers:
  - name: RangeSeparation
  - name: ExponentialBernsteinRBF
  - name: NuclearEmbedding
  - name: ElectronicEmbedding
    params:
        attribute: charge
  - name: Core
    params:
        num_modules: 6
        shallow_ensemble_size: 10
  - name: AtomicAffine
  - name: ChargeConservation
  - name: ElectrostaticEnergy
    params:
        flavor: SpookyNet
        dielectric_constant: 10.0
  - name: AtomicCharge2Dipole
  - name: EnergyReduce
  - name: ShallowEnsembleReduce
    params:
        var: [E]
        train_only: true
  - name: Force
  - name: ShallowEnsembleReduce
    params:
        var: [E, Fa]
        eval_only: true

Layer ordering matters#

  1. Build geometric features (range separation, RBFs)

  2. Embed atoms and optional scalar features

  3. Message passing (Core)

  4. Normalize atomic outputs (AtomicAffine)

  5. Enforce physics (charge conservation, electrostatics, dispersion)

  6. Reduce to molecular properties (EnergyReduce)

  7. Optional ensemble statistics

  8. Analytic forces via autograd (Force)

Shared build_params#

Common keys in build_params:

  • cutoff_sr, cutoff_lr, cutoff_fn

  • dim_embedding, num_rbf, max_Za

  • Hartree_in_E, Bohr_in_R

Layers inherit these unless params overrides them.

Monitoring energy terms#

Optional Trainer.Monitor lists terms such as E_ele (electrostatic), E_disp (D3/D4), E_zbl for debugging layer contributions during training.