Merge pull request ACEsuit#296 from ACEsuit/develop

Change stress input + update version

Author: ilyes319

mace/__version__.py

@@ -1 +1 @@
-__version__ = "0.3.3"
+__version__ = "0.3.4"

mace/data/atomic_data.py

@@ -175,7 +175,9 @@ def from_config(
             else None
         )
         virials = (
-            torch.tensor(config.virials, dtype=torch.get_default_dtype()).unsqueeze(0)
+            voigt_to_matrix(
+                torch.tensor(config.virials, dtype=torch.get_default_dtype())
+            ).unsqueeze(0)
             if config.virials is not None
             else None
         )

mace/data/utils.py

@@ -17,8 +17,8 @@
 Vector = np.ndarray  # [3,]
 Positions = np.ndarray  # [..., 3]
 Forces = np.ndarray  # [..., 3]
-Stress = np.ndarray  # [6, ]
-Virials = np.ndarray  # [3,3]
+Stress = np.ndarray  # [6, ], [3,3], [9, ]
+Virials = np.ndarray  # [6, ], [3,3], [9, ]
 Charges = np.ndarray  # [..., 1]
 Cell = np.ndarray  # [3,3]
 Pbc = tuple  # (3,)

mace/tools/torch_tools.py

@@ -107,7 +107,7 @@ def cartesian_to_spherical(t: torch.Tensor):
 def voigt_to_matrix(t: torch.Tensor):
     """
     Convert voigt notation to matrix notation
-    :param t: (6,) tensor or (3, 3) tensor
+    :param t: (6,) tensor or (3, 3) tensor or (9,) tensor
     :return: (3, 3) tensor
     """
     if t.shape == (3, 3):
@@ -121,9 +121,11 @@ def voigt_to_matrix(t: torch.Tensor):
             ],
             dtype=t.dtype,
         )
+    if t.shape == (9,):
+        return t.view(3, 3)
 
     raise ValueError(
-        f"Stress tensor must be of shape (6,) or (3, 3), but has shape {t.shape}"
+        f"Stress tensor must be of shape (6,) or (3, 3), or (9,) but has shape {t.shape}"
     )