diff --git a/src/equisolve/utils/convert.py b/src/equisolve/utils/convert.py index fcd09c6..288e3dd 100644 --- a/src/equisolve/utils/convert.py +++ b/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,19 +36,34 @@ 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 + if (frames[0].get_calculator() is not None) and ("forces" in frames[0].get_calculator().implemented_properties): + positions_gradients = [-f.get_forces() if (f.get_calculator() is not None) and ("forces" in f.get_calculator().implemented_properties) + else raise AttributeError(f"First frame has calculator with forces, so assumed all frames have forces but frame {i} does not have forces. Please add or remove forces as property everywhere.") + for i, f in enumerate(frames)] + else: + positions_gradients = None + if stress is not None: - cell_gradients = [-f.arrays[stress] for f in frames] + 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) + return properties_to_tensormap( + values, positions_gradients, cell_gradients, property_name=energy + ) def properties_to_tensormap( @@ -55,6 +71,7 @@ def properties_to_tensormap( positions_gradients: List[np.ndarray] = None, cell_gradients: List[np.ndarray] = None, is_structure_property: bool = True, + property_name: str = "property", ) -> TensorMap: """Create a :class:`equistore.TensorMap` from array like properties. @@ -94,7 +111,7 @@ def properties_to_tensormap( values=np.asarray(values).reshape(-1, 1), samples=Labels(["structure"], np.arange(n_structures).reshape(-1, 1)), components=[], - properties=Labels(["property"], np.array([(0,)])), + properties=Labels([property_name], np.array([(0,)])), ) if positions_gradients is not None: diff --git a/tests/equisolve_tests/utils/convert.py b/tests/equisolve_tests/utils/convert.py index f0552dd..b0992e5 100644 --- a/tests/equisolve_tests/utils/convert.py +++ b/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 = [] @@ -43,7 +47,7 @@ def test_ase_to_tensormap(self, energies, forces, stress): frame = ase.Atoms(self.n_atoms * "H") frame.info["energy"] = energies[i] frame.arrays["forces"] = forces[i] - frame.arrays["stress"] = stress[i] + frame.info["stress"] = stress[i] frames.append(frame) property_tm = ase_to_tensormap(frames, "energy", "forces", "stress") @@ -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]