Development Notes

Development Notes#

Module layout#

The model wrappers keep DeePMD-facing APIs in the top-level model modules and move reusable glue code into small helpers:

deepmd_gnn.mace

The DeePMD BaseModel implementation for MACE. It owns parameter serialization, statistics, lower-interface forwarding, force/virial gradients, export hooks, and MPI feature communication.

deepmd_gnn.mace_network

Construction and conversion of the underlying MACE ScaleShiftMACE module. Keep MACE package imports, cuEquivariance configuration, scripted model creation, and cueq-to-e3nn weight transfer here instead of mixing them into the DeePMD wrapper.

deepmd_gnn.export

Shared helpers for PyTorch exportable backend traces. This includes neighbor list padding for symbolic shapes and one-shot cleanup of guards generated by torch.export.

deepmd_gnn.deepmd_ops

Compatibility hooks for DeePMD custom PyTorch ops. The placeholder border_op exists only so TorchScript can compile wrappers before the real DeePMD extension is loaded.

deepmd_gnn.stat_compat

Version-tolerant access to DeePMD observed-type statistic helpers used by both MACE and NequIP wrappers.

MACE forward path#

MaceModel.forward_lower_common is intentionally kept in mace.py because it is the main adapter between DeePMD tensors and MACE graph tensors. When changing it, keep these stages separate:

  1. Normalize DeePMD lower-interface tensors and build an edge index.

  2. Convert extended ghost atoms back to mapped local atoms when no MPI communication dictionary is available.

  3. Build the MACE input graph and optional displacement tensor used for virial gradients.

  4. Select the energy path: eager MACE, exportable module-by-module execution, or layer-wise MPI communication.

  5. Differentiate the reduced energy to produce DeePMD-compatible force and virial outputs.

Avoid mixing export-only logic into the normal eager path unless the same behavior is required by training, inference, and freezing. The export path has different constraints because it must be traceable with symbolic atom counts.