fix(models): symmetrize prepare_strain displacement (#111)

Apply only the symmetric part of the displacement tensor when prepare_strain deforms positions and cells.
Document that the returned displacement remains an unconstrained autograd leaf, and add tests for symmetric stress from an asymmetric toy energy plus preservation of a symmetric raw stress response.

Signed-off-by: Roman Zubatyuk <rzubatiuk@nvidia.com>

Author: zubatyuk

nvalchemi/models/_utils.py

@@ -91,9 +91,10 @@ def prepare_strain(
     """Set up the affine strain trick for autograd stress computation.
 
     Creates a per-system 3x3 displacement tensor with
-    ``requires_grad=True``, scales positions and cell through it, and
-    returns all three tensors.  After running the model on the scaled
-    positions/cell, compute stresses with standard PyTorch autograd::
+    ``requires_grad=True``, scales positions and cell through the symmetric
+    part of it, and returns all three tensors.  After running the model on
+    the scaled positions/cell, compute stresses with standard PyTorch
+    autograd::
 
         scaled_pos, scaled_cell, displacement = prepare_strain(
             positions, cell, batch_idx
@@ -131,6 +132,8 @@ def prepare_strain(
     tuple[torch.Tensor, torch.Tensor, torch.Tensor]
         ``(scaled_positions, scaled_cell, displacement)`` where
         ``displacement`` is ``[B, 3, 3]`` with ``requires_grad=True``.
+        The returned tensor is an unconstrained autograd leaf; only its
+        symmetric part is applied as strain.
     """
     n_systems = cell.shape[0]
     displacement = torch.zeros(
@@ -141,15 +144,17 @@ def prepare_strain(
         device=positions.device,
     )
     displacement.requires_grad_(True)
-    symmetric = (
-        torch.eye(3, dtype=positions.dtype, device=positions.device) + displacement
+    symmetric_displacement = 0.5 * (displacement + displacement.mT)
+    deformation = (
+        torch.eye(3, dtype=positions.dtype, device=positions.device)
+        + symmetric_displacement
     )
-    # Scale positions: pos'[n] = pos[n] @ symmetric[system_of_atom[n]]
-    # Index into symmetric per-atom, then batch-matmul each atom's row.
-    per_atom_symmetric = symmetric[batch_idx]  # [N, 3, 3]
-    scaled_positions = torch.einsum("ni,nij->nj", positions, per_atom_symmetric)
-    # Scale cell: cell'[b] = cell[b] @ symmetric[b]
-    scaled_cell = torch.einsum("bij,bjk->bik", cell, symmetric)
+    # Scale positions: pos'[n] = pos[n] @ deformation[system_of_atom[n]]
+    # Index into deformation per-atom, then batch-matmul each atom's row.
+    per_atom_deformation = deformation[batch_idx]  # [N, 3, 3]
+    scaled_positions = torch.einsum("ni,nij->nj", positions, per_atom_deformation)
+    # Scale cell: cell'[b] = cell[b] @ deformation[b]
+    scaled_cell = torch.einsum("bij,bjk->bik", cell, deformation)
     return scaled_positions, scaled_cell, displacement
 
 

test/models/test_pipeline.py

@@ -1314,6 +1314,77 @@ def test_gradient_flows_through_displacement(self):
         assert grad is not None
         assert grad.shape == (1, 3, 3)
 
+    def test_stress_from_asymmetric_energy_is_symmetric(self):
+        """prepare_strain projects displacement gradients onto symmetric strain."""
+        from nvalchemi.models._utils import autograd_stresses, prepare_strain
+
+        positions = torch.tensor(
+            [[0.3, -0.7, 1.1], [1.2, 0.4, -0.2], [-0.5, 0.8, 0.6]],
+            dtype=torch.float64,
+        )
+        cell = torch.eye(3, dtype=torch.float64).unsqueeze(0) * 5.0
+        batch_idx = torch.zeros(positions.shape[0], dtype=torch.long)
+
+        scaled_pos, scaled_cell, displacement = prepare_strain(
+            positions,
+            cell,
+            batch_idx,
+        )
+        x, y = scaled_pos[:, 0], scaled_pos[:, 1]
+        energy = (x * y**2).sum() + scaled_cell[0, 0, 0] * scaled_cell[0, 1, 1] ** 2
+        stress = autograd_stresses(energy, displacement, cell, num_graphs=1)
+
+        torch.testing.assert_close(
+            stress,
+            stress.transpose(-1, -2),
+            atol=1e-12,
+            rtol=0,
+        )
+
+    def test_symmetric_response_matches_raw_displacement_reference(self):
+        """Symmetric strain preserves models with symmetric raw stress response."""
+        from nvalchemi.models._utils import autograd_stresses, prepare_strain
+
+        positions = torch.tensor(
+            [[0.3, -0.7, 1.1], [1.2, 0.4, -0.2], [-0.5, 0.8, 0.6]],
+            dtype=torch.float64,
+        )
+        cell = torch.eye(3, dtype=torch.float64).unsqueeze(0) * 5.0
+        batch_idx = torch.zeros(positions.shape[0], dtype=torch.long)
+
+        scaled_pos, scaled_cell, displacement = prepare_strain(
+            positions,
+            cell,
+            batch_idx,
+        )
+        energy = (scaled_pos**2).sum() + (scaled_cell**2).sum()
+        stress = autograd_stresses(energy, displacement, cell, num_graphs=1)
+
+        raw_displacement = torch.zeros(
+            1,
+            3,
+            3,
+            dtype=positions.dtype,
+            requires_grad=True,
+        )
+        raw_deformation = torch.eye(3, dtype=positions.dtype).unsqueeze(0)
+        raw_deformation = raw_deformation + raw_displacement
+        raw_scaled_pos = torch.einsum(
+            "ni,nij->nj",
+            positions,
+            raw_deformation[batch_idx],
+        )
+        raw_scaled_cell = torch.einsum("bij,bjk->bik", cell, raw_deformation)
+        raw_energy = (raw_scaled_pos**2).sum() + (raw_scaled_cell**2).sum()
+        raw_stress = autograd_stresses(
+            raw_energy,
+            raw_displacement,
+            cell,
+            num_graphs=1,
+        )
+
+        torch.testing.assert_close(stress, raw_stress, atol=1e-12, rtol=0)
+
     def test_multi_system_batches(self):
         """Each system gets its own displacement."""
         from nvalchemi.models._utils import prepare_strain