Assembling a network from new layers

The twelve cookbook layers compose into a complete two-layer E3GNN that predicts Hamiltonians for any chemistry. Nothing in the model is inherited from the default network except spherical harmonics and the operator decoder.

The recipe

# A complete E3GNN assembled exclusively from the cookbook layers in layers.py.
#
# Every stage reads named feature slots and writes new ones. Interpolations
# such as ${species.order} and ${layout.coupled_irreps} are resolved against
# the orbital vocabulary of the training data before the model is built.
schema_version: 1
species: {ordering: atomic_number}
symmetrize_hamiltonian_targets: false
symmetrize_validation_output: false
apply_cutoff_to_targets: false
revert_on_spike: false
checkpoint_monitor: val/loss
parameters:
  cutoff: 6.0
  angular_irreps: 1x0e+1x1o+1x2e+1x3o+1x4e
  edge_irreps: 16x0e+8x1o+8x2e+4x3o+4x4e
  node_irreps: 32x0e+16x1o+16x1e+8x2e+8x2o+8x3o+8x3e+4x4e
  pair_irreps: 16x0e+8x1o+8x1e+8x2e+8x2o+4x3o+4x3e+4x4e
data:
  representation: packed_shells
  basis: openmx
  graph_source: radius
  graph_radius: 6.0
  validate_hermitian: false
  validate_support: false
  canonicalize_edges: false
training:
  seed: 7
  dtype: float32
  batch_size: 1
  max_epochs: 200
objective: {kind: mse, epoch_weight: elements}
optimization:
  optimizer: {kind: adam, lr: 0.003}
pipeline:
  stages:
    # -- geometry embeddings ------------------------------------------------
    - id: radial
      transform: bessel_basis
      inputs: [edge_lengths]
      outputs: [radial]
      options: {cutoff: '${parameters.cutoff}', n_basis: 8}
    - id: angular
      transform: spherical_harmonics
      inputs: [edge_vectors]
      outputs: [angular]
      options: {irreps: '${parameters.angular_irreps}'}
    # -- encoders -----------------------------------------------------------
    - id: node_encoder
      role: encoder
      transform: periodic_table_embedding
      inputs: [node_type_idx]
      outputs: [node_0]
      options: {elements: '${species.order}', width: 32}
    - id: edge_encoder
      role: encoder
      transform: radial_angular_encoder
      inputs: [radial, angular]
      outputs: [edge_0]
      options: {irreps_out: '${parameters.edge_irreps}'}
    # -- message passing layer 1 -------------------------------------------
    - id: convolution_1
      role: message
      transform: depthwise_convolution
      inputs: [node_0, angular, radial, edge_index]
      outputs: [messages_1]
      options: {lmax: 4}
    - id: aggregation_1
      role: message
      transform: tempered_softmax_aggregation
      inputs: [messages_1, radial, edge_index, node_0]
      outputs: [aggregated_1]
    - id: mlp_1
      transform: equivariant_glu
      inputs: [aggregated_1]
      outputs: [mixed_1]
      options: {irreps_out: '${parameters.node_irreps}'}
    - id: norm_1
      transform: degree_rms_norm
      inputs: [mixed_1]
      outputs: [normalized_1]
    - id: activation_1
      transform: squash_activation
      inputs: [normalized_1]
      outputs: [node_1]
    # -- message passing layer 2 -------------------------------------------
    - id: convolution_2
      role: message
      transform: depthwise_convolution
      inputs: [node_1, angular, radial, edge_index]
      outputs: [messages_2]
      options: {lmax: 4}
    - id: aggregation_2
      role: message
      transform: tempered_softmax_aggregation
      inputs: [messages_2, radial, edge_index, node_1]
      outputs: [aggregated_2]
    - id: mlp_2
      transform: equivariant_glu
      inputs: [aggregated_2]
      outputs: [mixed_2]
      options: {irreps_out: '${parameters.node_irreps}'}
    - id: activation_2
      transform: squash_activation
      inputs: [mixed_2]
      outputs: [update_2]
    - id: residual_2
      transform: rezero_residual
      inputs: [node_1, update_2]
      outputs: [node_2]
    # -- pair features and operator readout --------------------------------
    - id: edge_update
      role: message
      transform: endpoint_bilinear_edge_update
      inputs: [node_2, edge_0, edge_index]
      outputs: [edge_1]
      options: {irreps_out: '${parameters.pair_irreps}'}
    - id: readout
      role: readout
      transform: pair_mixture_readout
      inputs: [edge_1, edge_type_idx]
      outputs: [coupled]
      options: {irreps_out: '${layout.coupled_irreps}', num_pairs: '${species.pair_count}', rank: 4}
    - id: envelope
      role: readout
      transform: learned_pair_envelope
      inputs: [coupled, edge_lengths, edge_type_idx]
      outputs: [coupled_scaled]
      options: {num_pairs: '${species.pair_count}'}
    - id: decode
      role: readout
      transform: shell_tensor_decode
      inputs: [coupled_scaled]
      outputs: [hamiltonian_native]
      options: {shell_pairs: '${layout.shell_pairs}', spinful: '${layout.spinful}'}
  outputs: {hamiltonian: hamiltonian_native}

Points worth noting:

• elements: ${species.order}, ${species.pair_count} and ${layout.coupled_irreps} make the recipe chemistry-agnostic. The same file builds a silicon model (Si: s2p2d1) and the H/O/Si model of the verification script.

• node_irreps contains both parities of every degree up to 3 and 4e, and the bilinear edge update produces pair_irreps; together they cover every coupled irrep of s, p and d blocks (Sparse operators and their irreps).

• The residual of layer 2 reads node_1 long after node_1 was written: slots are write-once, so skip connections are plain references.

Verification

examples/layer_cookbook/verify.py resolves this recipe for the species H (s), O (s, p) and Si (s, p, d), builds the model in double precision, and checks the operator transformation law on a periodic structure with a skewed cell:

def check_model():
    print("Model covariance")
    model = build_model().eval()
    atoms = ("H", "O", "Si", "O")
    positions = torch.tensor(
        [[0.1, 0.2, 0.3], [1.3, 0.4, 0.6], [-0.2, 1.5, 0.8], [0.9, 1.1, -0.7]]
    )
    box = torch.tensor([[4.6, 0.2, 0.3], [0.1, 4.9, 0.2], [0.3, 0.1, 5.2]])
    rotation = OrbitalFrameRotation(dtype=torch.float64)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    for phase in ("initial", "trained"):
        base = predict(model, atoms, positions, box)
        for label, g in group_elements().items():
            observed = predict(model, atoms, positions @ g.T, box @ g.T)
            worst = 0.0
            for edge, (key, row) in base.lookup.items():
                left, right = (SPECIES_SHELLS[s] for s in key.split("-"))
                expected = rotation.rotate_blocks(
                    base.pair_blocks[key][row], g.T, left, right
                )
                worst = max(worst, (observed[edge] - expected).abs().max().item())
            if worst > 1e-8:
                raise AssertionError(f"{phase} model, {label}: error {worst:.2e}")
            print(f"  {phase:<8s} {label:<18s} max |dH| = {worst:.1e}")
        if phase == "initial":
            optimizer.zero_grad()
            loss = sum(block.square().mean() for block in base.pair_blocks.values())
            loss.backward()
            missing = [
                name
                for name, p in model.named_parameters()
                if p.grad is None or not torch.isfinite(p.grad).all()
            ]
            if missing:
                raise AssertionError(f"Parameters without finite gradients: {missing}")
            optimizer.step()
    print(f"  parameters: {sum(p.numel() for p in model.parameters()):,}")

rotation.rotate_blocks applies \(\mathcal D_i(Q)\,H_{ij}\,\mathcal D_j(Q)^\top\) in the native orbital order of the packed representation. The test is run before and after an optimizer step, because a symmetry violation can hide in parameters that are zero at initialization. The observed deviations are at the level of floating-point round-off (\(\sim10^{-10}\) Hartree for blocks of order \(10^{-3}\)), and zero for the inversion, which the network reproduces bit for bit.

Training

The standard trainer accepts the recipe once the layers are registered:

PYTHONPATH=src:. python examples/layer_cookbook/train.py \
    --config examples/layer_cookbook/cookbook.yaml \
    --set 'data.splits.train=[data/small/Si]' \
    --set 'data.splits.val=[data/small/Si_perturbed_scale1_078]' \
    --set training.max_epochs=150 \
    --output runs/cookbook_silicon --accelerator cpu

Trained on a single eight-atom silicon structure and validated on a perturbed one, the packed-shell MSE (Hartree²) evolves as follows on a laptop CPU (about one second per epoch once the dataset is cached):

Epoch	Training loss	Validation loss
1	0.377	0.243
25	0.271	0.194
50	0.154	0.142
100	0.0728	0.0904
150	0.0305	0.0651

The run directory contains run_manifest.json (resolved recipe, input hashes, environment), split_manifest.json, latest_checkpoint.pt for continuation and best_model.pt for evaluation. Restoring a checkpoint of a model with custom layers requires the same registration import; the evaluation example accepts it as --plugin:

PYTHONPATH=src:. python examples/quickstart/evaluate.py \
    --checkpoint runs/cookbook_silicon/best_model.pt \
    --structure data/small/Si_perturbed_scale1_078 \
    --plugin examples/layer_cookbook/layers.py

From here, any single layer can be swapped for its built-in counterpart — gaussian_basis for bessel_basis, projected_attention for tempered_softmax_aggregation, irrep_activation for squash_activation — by editing one stage, which makes controlled ablations of individual layer kinds a matter of configuration.