PolarMACE periodic stress may omit explicit strain dependence of the long-range electrostatic energy

[Alexandrina-Chen]

Describe the bug

Hi, I may be overlooking an intended derivative path, but after tracing the current main-branch implementation of PolarMACE, I believe its analytical stress is incomplete for periodic systems.

For a homogeneous strain tensor $\boldsymbol{\epsilon}$, the stress should be obtained from the complete strain derivative of the total energy:

$$\boldsymbol{\sigma} = \frac{1}{V} \frac{\partial E_{\mathrm{tot}}}{\partial \boldsymbol{\epsilon}}.$$

The Cartesian positions and cell should transform consistently:

$$\mathbf R(\boldsymbol{\epsilon}) = \mathbf R(\mathbf I+\boldsymbol{\epsilon}), \qquad \mathbf H(\boldsymbol{\epsilon}) = \mathbf H(\mathbf I+\boldsymbol{\epsilon}).$$

For a reciprocal-space electrostatic energy, the strain dependence also enters through the reciprocal lattice and cell volume:

$$\mathbf G(\boldsymbol{\epsilon}) = 2\pi\mathbf H(\boldsymbol{\epsilon})^{-T}, \qquad V(\boldsymbol{\epsilon}) = \det\mathbf H(\boldsymbol{\epsilon}).$$

Therefore, the long-range contribution contains explicit geometric derivatives such as

$$\frac{\partial E_{\mathrm{LR}}}{\partial \boldsymbol{\epsilon}} = \frac{\partial E_{\mathrm{LR}}}{\partial \mathbf R} : \frac{\partial \mathbf R}{\partial \boldsymbol{\epsilon}} + \frac{\partial E_{\mathrm{LR}}}{\partial \mathbf G} : \frac{\partial \mathbf G}{\partial \boldsymbol{\epsilon}} + \frac{\partial E_{\mathrm{LR}}}{\partial V} \frac{\partial V}{\partial \boldsymbol{\epsilon}} +\cdots.$$

In the current implementation, the positions, cell, reciprocal cell, and volume used by the PolarMACE long-range calculation do not appear to be connected to the symbolic strain variable. The returned stress therefore seems to omit the explicit affine deformation of the long-range electrostatic geometry.

This would affect periodic stress, stress-supervised training, variable-cell relaxation, NPT dynamics, equations of state, and elastic properties. Energies and forces evaluated at a fixed cell are not necessarily affected in the same way.

To Reproduce / code path inspected

1. In mace/modules/utils.py, prepare_graph first stores the original positions and cell:

positions = data["positions"]
cell = data["cell"]

When stress or virials are requested, get_symmetric_displacement constructs strain-dependent positions and periodic shifts:

p, s, displacement = get_symmetric_displacement(
    positions=positions,
    unit_shifts=data["unit_shifts"],
    cell=cell,
    edge_index=data["edge_index"],
    num_graphs=num_graphs,
    batch=data["batch"],
    displacement=data.get("displacement"),
)

data["positions"], data["shifts"] = p, s

However, the returned GraphContext still contains the original local variables:

return GraphContext(
    ...
    positions=positions,
    ...
    cell=cell,
    ...
)

Thus:

• data["positions"] is strain-dependent;
• data["shifts"] is strain-dependent;
• the edge vectors used by the local MACE backbone are strain-dependent;
• ctx.positions remains the original position tensor;
• ctx.cell remains the original cell tensor.

2. In PolarMACE.forward in mace/modules/extensions.py, the long-range calculation uses the tensors from GraphContext:

positions = ctx.positions
cell = ctx.cell

The reciprocal-space vectors are then constructed using the original cell and the precomputed reciprocal cell:

compute_k_vectors_flat(
    self.kspace_cutoff,
    cell.view(-1, 3, 3),
    data["rcell"].view(-1, 3, 3),
)

The electrostatic feature and energy routines also receive the original positions and the precomputed volume:

node_positions=positions
volume=data["volume"]

3. In mace/data/atomic_data.py, volume and rcell are calculated once when the graph data are created:

volume = torch.linalg.det(cell)

if torch.abs(volume) > 0:
    rcell = 2 * torch.pi * torch.linalg.inv(cell.mT)
else:
    rcell = torch.zeros(3, 3, dtype=torch.get_default_dtype())

These tensors are not subsequently recomputed from a strain-dependent cell inside PolarMACE.forward.

4. At the end of PolarMACE.forward, the complete local plus long-range energy is passed to the usual MACE derivative routine:

forces, virials, stress, ... = get_outputs(
    energy=total_energy,
    positions=positions,
    displacement=displacement,
    cell=cell,
    ...
)

The local MACE features depend on displacement through the strained edge vectors. Therefore, the learned charge-density coefficients may retain an indirect strain dependence through the local representation.

However, the geometric quantities passed to the long-range calculation appear fixed with respect to displacement:

$$\mathbf R_{\mathrm{LR}}=\mathbf R, \qquad \mathbf H_{\mathrm{LR}}=\mathbf H, \qquad \mathbf G_{\mathrm{LR}}=2\pi\mathbf H^{-T}, \qquad V_{\mathrm{LR}}=\det\mathbf H.$$

The current derivative therefore appears to include indirect changes in the learned charge density or multipoles, but not the explicit affine deformation of the positions, reciprocal lattice, and volume used by the long-range electrostatic calculation.

Expected behavior

All geometric quantities used by the long-range energy should be constructed from the same symbolic strain tensor used by the local model. Schematically:

symmetric_displacement = 0.5 * (
    displacement + displacement.transpose(-1, -2)
)

positions_strained = positions + torch.einsum(
    "be,bec->bc",
    positions,
    symmetric_displacement[batch],
)

cell_strained = cell.view(-1, 3, 3)
cell_strained = cell_strained + torch.matmul(
    cell_strained,
    symmetric_displacement,
)

volume_strained = torch.linalg.det(cell_strained)

rcell_strained = (
    2
    * torch.pi
    * torch.linalg.inv(cell_strained.transpose(-1, -2))
)

The PolarMACE electrostatic feature construction, Coulomb energy, dipole calculation, and field-dependent nonlocal energy should then use:

positions_strained
cell_strained
rcell_strained
volume_strained

before the complete total energy is differentiated with respect to displacement.

The returned stress should satisfy

$$\boldsymbol{\sigma}_{\mathrm{PolarMACE}} = \frac{1}{V} \frac{\partial}{\partial\boldsymbol{\epsilon}} \left( E_{\mathrm{local}} + E_{\mathrm{electrostatic}} + E_{\mathrm{nonlocal}} \right),$$

with all position- and cell-dependent quantities transformed consistently.

Here, the final prefactor $1/V$ can use the current unstrained cell volume at $\boldsymbol{\epsilon}=0$. The separate volume argument that enters the long-range energy itself must vary with the symbolic strain.

For the reciprocal-space derivative, the integer reciprocal-mode indices should normally remain fixed under an infinitesimal strain, while the physical reciprocal vectors associated with those indices vary with the strained cell.

Why this matters

The omitted terms may be small for an isolated molecule placed in a very large periodic box. They can be much more important for dense periodic systems.

Could the maintainers confirm whether there is another cell-derivative path that I have overlooked? If not, I believe the current periodic PolarMACE stress is incomplete until the long-range positions, cell, reciprocal cell, and volume are connected to the symbolic strain tensor.