Three working approaches to predicting crystal energy and forces with neural networks — a pretrained equivariant foundation model (MACE-MP-0), an invariant network (SchNet), and a custom GNN built from scratch in PyTorch — with full runnable code.
MACE-MP-0 — pretrained, recommended
SchNet — invariant GNN
Custom GNN from scratch
F = −dE/dR via autodiff
This post moves from theory to running code. All three approaches answer the same question — given atomic positions, predict total energy, then obtain forces via automatic differentiation — but differ in how much directional geometric information they preserve internally, exactly as covered in Post 17.
1. Build a graph from atomic positions (nodes = atoms, edges = nearby pairs).
2. Predict a single scalar: total energy E.
3. Call torch.autograd.grad(E, positions) to get forces —
no separate force-prediction code is ever written by hand.
1. Setup — Shared Across All Three Approaches
import torch
import numpy as np
from ase import Atoms
from ase.build import bulk
# Example: rock-salt NiO — a correlated transition-metal oxide
atoms = bulk('NiO', crystalstructure='rocksalt', a=4.17)
positions = torch.tensor(atoms.get_positions(), dtype=torch.float32,
requires_grad=True) # ← critical: enables autodiff
atomic_numbers = torch.tensor(atoms.get_atomic_numbers())
Setting requires_grad=True on the position tensor is what makes
force computation possible at all — without it, PyTorch never builds the
computational graph needed to differentiate energy with respect to position,
and torch.autograd.grad will fail.
2. Approach 1 — MACE-MP-0 (Pretrained, Equivariant, Recommended)
MACE-MP-0 is a foundation model trained on the full Materials Project relaxation trajectory database, as introduced in Post 16. It can be applied zero-shot — no training required — to essentially any inorganic crystal, including transition metal compounds.
from mace.calculators import mace_mp
from ase.optimize import BFGS
# Load the pretrained foundation model — downloads weights on first use
calc = mace_mp(model="medium", dispersion=False, default_dtype="float64")
atoms.calc = calc
# Energy and forces — both come from ONE forward pass + autodiff internally
energy = atoms.get_potential_energy() # eV
forces = atoms.get_forces() # eV/Å, shape (N_atoms, 3)
print(f"NiO total energy: {energy:.4f} eV")
print(f"Max force component: {np.abs(forces).max():.4f} eV/Å")
# Relax the structure — ASE handles the optimisation loop
opt = BFGS(atoms)
opt.run(fmax=0.01) # converge until max force < 0.01 eV/Å
print(f"Relaxed energy: {atoms.get_potential_energy():.4f} eV")
No DFT data from your own calculations is required to run this code — MACE-MP-0 has already seen tens of thousands of Materials Project structures, including many transition-metal oxides and chalcogenides similar to NiO. This is the recommended starting point in Post 21's roadmap.
3. Approach 2 — SchNet (Invariant GNN)
Using SchNetPack gives a transparent, invariant baseline — useful for understanding the architecture from Post 18, and as a comparison point against MACE's equivariant features.
import schnetpack as spk
import schnetpack.transform as trn
# Build a small SchNet model (untrained — for architecture demonstration)
cutoff = 5.0 # Å — neighbour cutoff radius
n_atom_basis = 64
pairwise_distance = spk.atomistic.PairwiseDistances()
representation = spk.representation.SchNet(
n_atom_basis=n_atom_basis, n_interactions=3,
radial_basis=spk.nn.GaussianRBF(n_rbf=20, cutoff=cutoff),
cutoff_fn=spk.nn.CosineCutoff(cutoff)
)
pred_energy = spk.atomistic.Atomwise(n_in=n_atom_basis, output_key="energy")
nnpot = spk.model.NeuralNetworkPotential(
representation=representation,
input_modules=[pairwise_distance],
output_modules=[pred_energy],
)
# Forward pass: predicted energy (untrained, illustrative only)
inputs = {"_positions": positions, "_atomic_numbers": atomic_numbers}
output = nnpot(inputs)
energy_pred = output["energy"]
# Forces via autodiff — exactly the same pattern as MACE
forces_pred = -torch.autograd.grad(energy_pred.sum(), positions,
create_graph=True)[0]
4. Approach 3 — A Custom GNN Built From Scratch
Building a minimal GNN by hand makes the Behler-Parrinello "sum of per-atom contributions" pattern from Post 18 completely explicit, with no library abstracting it away.
import torch.nn as nn
class SimpleGNN(nn.Module):
"""Minimal Behler-Parrinello-style potential: per-atom energy MLP,
summed over atoms. Invariant (distance-based), like SchNet."""
def __init__(self, hidden=32, cutoff=5.0):
super().__init__()
self.cutoff = cutoff
# Per-atom MLP — same weights applied to every atom
self.mlp = nn.Sequential(
nn.Linear(1, hidden), nn.SiLU(), # input: 1 distance per neighbour pair
nn.Linear(hidden, hidden), nn.SiLU(),
nn.Linear(hidden, 1)
)
def forward(self, positions):
n = positions.shape[0]
# All pairwise distances
diff = positions.unsqueeze(1) - positions.unsqueeze(0) # (n,n,3)
dist = torch.norm(diff, dim=-1) # (n,n) — invariant by construction
# Per-pair energy contribution, masked beyond cutoff and self-pairs
mask = (dist > 1e-6) & (dist < self.cutoff)
pair_energy = self.mlp(dist.unsqueeze(-1)).squeeze(-1)
pair_energy = pair_energy * mask
# Sum per-atom (Behler-Parrinello pattern): each atom collects
# half the energy of every pair it participates in, then sum over atoms
total_energy = pair_energy.sum() / 2.0
return total_energy
# Usage — identical pattern to the previous two approaches
model = SimpleGNN()
energy = model(positions)
forces = -torch.autograd.grad(energy, positions, create_graph=True)[0]
print(f"Custom GNN energy: {energy.item():.4f} | Force shape: {forces.shape}")
5. Comparing the Three Approaches
| Approach | Training needed? | Invariant/Equivariant | Accuracy | When to use |
|---|---|---|---|---|
| MACE-MP-0 | None — zero-shot | Equivariant | Near-DFT for most inorganics | Always start here (Post 21) |
| SchNet | Required (not pretrained above) | Invariant | Good after training, needs more data | Architecture study, baseline comparison |
| Custom GNN | Required, very minimal model | Invariant | Low (toy model, 2 features) | Pedagogical — see the mechanics with no library abstraction |
Every approach above predicts a single scalar energy, then calls
torch.autograd.grad on it — never once is a force written
down by hand. This is the practical payoff of the F = −∂E/∂r relation
introduced in Post 16: build a good enough energy model, and accurate
forces follow automatically.
Quick Check
1. Why must requires_grad=True be set on the positions tensor before computing forces?
2. What is the main practical advantage of starting with MACE-MP-0 rather than training SchNet or a custom GNN from scratch?
3. In the custom GNN code, why is the summed pairwise energy divided by 2.0 before being returned as the total energy?
My name is Dr Abderrahmane Reggad. I'm a 54 year old and I am a researcher in materials' sciences. 