Adds a property string to ase_to_tensormap

src/equisolve/utils/convert.py

@@ -18,7 +18,8 @@
 def ase_to_tensormap(
     frames: List[ase.Atoms], energy: str = None, forces: str = None, stress: str = None
 ) -> TensorMap:
-    """Store informations from :class:`ase.Atoms` in a :class:`equistore.TensorMap`.
+    """Store informations from :class:`ase.Atoms`
+    in a :class:`equistore.TensorMap`.
 
     :param frames:
         ase.Atoms or list of ase.Atoms
@@ -35,17 +36,27 @@ def ase_to_tensormap(
     if not isinstance(frames, list):
         frames = [frames]
 
-    values = [f.info[energy] for f in frames]
+    if energy is not None:
+        values = [f.info[energy] for f in frames]
+    else:
+        energy = "energy"
+        values = [f.get_potential_energy() for f in frames]
 
     if forces is not None:
         positions_gradients = [-f.arrays[forces] for f in frames]
     else:
-        positions_gradients = None
+        try:
+            positions_gradients = [-f.get_forces() for f in frames]
+        except ase.ase.calculators.calculator.PropertyNotImplementedError:
+            positions_gradients = None
 
     if stress is not None:
         cell_gradients = [-f.info[stress] for f in frames]
     else:
-        cell_gradients = None
+        try:
+            cell_gradients = [-f.get_stress(voigt=False) for f in frames]
+        except ase.ase.calculators.calculator.PropertyNotImplementedError:
+            cell_gradients = None
 
     return properties_to_tensormap(
         values, positions_gradients, cell_gradients, property_name=energy

tests/equisolve_tests/utils/convert.py

@@ -9,6 +9,8 @@
 import ase
 import numpy as np
 import pytest
+from ase.calculators.calculator import Calculator, all_changes
+from ase.stress import voigt_6_to_full_3x3_stress
 from numpy.testing import assert_equal
 
 from equisolve.utils import ase_to_tensormap, properties_to_tensormap
@@ -35,7 +37,9 @@ def forces(self):
 
     @pytest.fixture
     def stress(self):
-        return [i for i in self.rng.random([self.n_strucs, 3, 3])]
+        return [
+            voigt_6_to_full_3x3_stress(i) for i in self.rng.random([self.n_strucs, 6])
+        ]
 
     def test_ase_to_tensormap(self, energies, forces, stress):
         frames = []
@@ -61,6 +65,67 @@ def test_ase_to_tensormap(self, energies, forces, stress):
             block.gradient("cell").values, -np.array(stress).reshape(-1, 3, 3, 1)
         )
 
+    def test_ase_to_tensormap_w_calculator(self, energies, forces, stress):
+        class CustomCalculator(Calculator):
+            implemented_properties = ("energy", "forces", "stress")
+
+            def __init__(self, energy, forces, stress, **kwargs):
+                Calculator.__init__(self, **kwargs)
+                self.energy = energy  # Predefined potential energy
+                self.forces = forces  # Predefined forces
+                self.stress = stress  # Predefined stress
+
+            def calculate(
+                self, atoms=None, properties=("energy"), system_changes=all_changes
+            ):
+                super().calculate(atoms, properties, system_changes)
+
+                self.results["energy"] = self.energy
+                self.results["forces"] = self.forces
+                self.results["stress"] = self.stress
+
+        frames = []
+        for i in range(len(energies)):
+            frame = ase.Atoms(self.n_atoms * "H")
+            frame.calc = CustomCalculator(energies[i], forces[i], stress[i])
+            frame.info["energy"] = energies[i]
+            frame.arrays["forces"] = forces[i]
+            frame.info["stress"] = stress[i]
+            frames.append(frame)
+
+        property_tm = ase_to_tensormap(frames, "energy", "forces", "stress")
+
+        # Use `[0]` function without parameters to check that TensorMap
+        # only has one block.
+        block = property_tm[0]
+
+        assert_equal(block.values, np.array(energies).reshape(-1, 1))
+        assert_equal(
+            block.gradient("positions").values,
+            -np.concatenate(forces, axis=0).reshape(-1, 3, 1),
+        )
+        assert_equal(
+            block.gradient("cell").values, -np.array(stress).reshape(-1, 3, 3, 1)
+        )
+
+        property_tm = ase_to_tensormap(frames)
+        block = property_tm[0]
+
+        assert_equal(block.values, np.array(energies).reshape(-1, 1))
+        assert_equal(
+            block.gradient("positions").values,
+            -np.concatenate(forces, axis=0).reshape(-1, 3, 1),
+        )
+
+        print(block.gradient("cell").values.shape)
+        print(np.array(stress).reshape(-1, 3, 3, 1).shape)
+
+        print(block.gradient("cell").values + np.array(stress).reshape(-1, 3, 3, 1))
+
+        assert_equal(
+            block.gradient("cell").values, -np.array(stress).reshape(-1, 3, 3, 1)
+        )
+
     def test_properties_to_tensormap(self, energies, forces, stress):
         property_tm = properties_to_tensormap(energies, forces, stress)
         block = property_tm[0]