No negative volumes in rare cases

CHANGELOG.md

@@ -38,6 +38,7 @@ Most recent change on the bottom.
 - Work with `wandb>=0.13.8`
 - Better error for standard deviation with too few data
 - `load_model_state` GPU -> CPU
+- No negative volumes in rare cases 
 
 ### Removed
 - [Breaking] `fixed_fields` machinery (`npz_fixed_field_keys` is still supported, but through a more straightforward implementation)

nequip/nn/_grad_output.py

@@ -330,12 +330,13 @@ def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
             # ^ can only scale by cell volume if we have one...:
             # Rescale stress tensor
             # See https://github.com/atomistic-machine-learning/schnetpack/blob/master/src/schnetpack/atomistic/output_modules.py#L180
+            # See also https://en.wikipedia.org/wiki/Triple_product
+            # See also https://gitlab.com/ase/ase/-/blob/master/ase/cell.py,
+            #          which uses np.abs(np.linalg.det(cell))
             # First dim is batch, second is vec, third is xyz
-            volume = torch.einsum(
-                "zi,zi->z",
-                cell[:, 0, :],
-                torch.cross(cell[:, 1, :], cell[:, 2, :], dim=1),
-            ).unsqueeze(-1)
+            # Note the .abs(), since volume should always be positive
+            # det is equal to a dot (b cross c)
+            volume = torch.linalg.det(cell).abs().unsqueeze(-1)
             stress = virial / volume.view(num_batch, 1, 1)
             data[AtomicDataDict.CELL_KEY] = orig_cell
         else: