Providing support for atomistic StructureData

src/aiida/orm/nodes/data/array/kpoints.py

@@ -223,16 +223,26 @@ def set_cell_from_structure(self, structuredata):
 
         :param structuredata: an instance of StructureData
         """
-        from aiida.orm import StructureData
+        from aiida.orm.nodes.data.structure import StructureData, has_atomistic
 
         if not isinstance(structuredata, StructureData):
-            raise ValueError(
-                'An instance of StructureData should be passed to ' 'the KpointsData, found instead {}'.format(
-                    structuredata.__class__
-                )
+            error_message = (
+                'An instance of StructureData or aiida-atomistic StructureData should be passed to '
+                'the KpointsData, found instead {}'.format(structuredata.__class__)
             )
-        cell = structuredata.cell
-        self.set_cell(cell, structuredata.pbc)
+            if has_atomistic():
+                from aiida_atomistic import StructureData as AtomisticStructureData
+
+                if not isinstance(structuredata, AtomisticStructureData):
+                    raise ValueError(error_message)
+                else:
+                    cell = structuredata.cell
+                    self.set_cell(cell, structuredata.pbc)
+            else:
+                raise ValueError(error_message)
+        else:
+            cell = structuredata.cell
+            self.set_cell(cell, structuredata.pbc)
 
     def set_cell(self, cell, pbc=None):
         """Set a cell to be used for symmetry analysis.

src/aiida/orm/nodes/data/structure.py

@@ -102,6 +102,15 @@ def has_pymatgen():
     return True
 
 
+def has_atomistic() -> bool:
+    """:return: True if the StructureData and StructureDataMutable from aiida-atomistic module can be imported, False otherwise."""
+    try:
+        import aiida_atomistic  # noqa: F401
+    except ImportError:
+        return False
+    return True
+
+
 def get_pymatgen_version():
     """:return: string with pymatgen version, None if can not import."""
     if not has_pymatgen():
@@ -1876,6 +1885,32 @@ def _get_object_pymatgen_molecule(self, **kwargs):
         positions = [list(site.position) for site in self.sites]
         return Molecule(species, positions)
 
+    def to_atomistic(self):
+        """
+        Returns the atomistic StructureData version of the orm.StructureData one.
+        """
+        if not has_atomistic():
+            raise ImportError(
+                'aiida-atomistic plugin is not installed, \
+                please install it to have full support for atomistic structures'
+            )
+        else:
+            from aiida_atomistic import StructureData, StructureDataMutable
+
+        atomistic = StructureDataMutable()
+        atomistic.set_pbc(self.pbc)
+        atomistic.set_cell(self.cell)
+
+        for site in self.sites:
+            atomistic.add_atom(
+                symbols=self.get_kind(site.kind_name).symbol,
+                masses=self.get_kind(site.kind_name).mass,
+                positions=site.position,
+                kinds=site.kind_name,
+            )
+
+        return StructureData.from_mutable(atomistic)
+
 
 class Kind:
     """This class contains the information about the species (kinds) of the system.

tests/orm/nodes/data/test_kpoints.py

@@ -11,6 +11,9 @@
 import numpy as np
 import pytest
 from aiida.orm import KpointsData, StructureData, load_node
+from aiida.orm.nodes.data.structure import has_atomistic
+
+skip_atomistic = pytest.mark.skipif(not has_atomistic(), reason='Unable to import aiida-atomistic')
 
 
 class TestKpoints:
@@ -81,3 +84,78 @@ def test_get_kpoints(self):
         kpt2 = load_node(kpt.pk)
         assert np.abs(kpt2.get_kpoints() - np.array(kpoints)).sum() == 0.0
         assert np.abs(kpt2.get_kpoints(cartesian=True) - np.array(cartesian_kpoints)).sum() == 0.0
+
+
+@skip_atomistic
+class TestKpointsAtomisticStructureData:
+    """Test for the `Kpointsdata` class using the new atomistic StructureData."""
+
+    @pytest.fixture(autouse=True)
+    def init_profile(self):
+        """Initialize the profile."""
+
+        from aiida_atomistic import StructureData, StructureDataMutable
+
+        alat = 5.430  # angstrom
+        cell = [
+            [
+                0.5 * alat,
+                0.5 * alat,
+                0.0,
+            ],
+            [
+                0.0,
+                0.5 * alat,
+                0.5 * alat,
+            ],
+            [0.5 * alat, 0.0, 0.5 * alat],
+        ]
+        self.alat = alat
+        mutable = StructureDataMutable()
+        mutable.set_cell(cell)
+        mutable.add_atom(positions=(0.000 * alat, 0.000 * alat, 0.000 * alat), symbols='Si')
+        mutable.add_atom(positions=(0.250 * alat, 0.250 * alat, 0.250 * alat), symbols='Si')
+        self.structure = StructureData.from_mutable(mutable)
+        # Define the expected reciprocal cell
+        val = 2.0 * np.pi / alat
+        self.expected_reciprocal_cell = np.array([[val, val, -val], [-val, val, val], [val, -val, val]])
+
+    def test_reciprocal_cell(self):
+        """Test the `reciprocal_cell` method.
+
+        This is a regression test for #2749.
+        """
+        kpt = KpointsData()
+        kpt.set_cell_from_structure(self.structure)
+
+        assert np.abs(kpt.reciprocal_cell - self.expected_reciprocal_cell).sum() == 0.0
+
+        # Check also after storing
+        kpt.store()
+        kpt2 = load_node(kpt.pk)
+        assert np.abs(kpt2.reciprocal_cell - self.expected_reciprocal_cell).sum() == 0.0
+
+    def test_get_kpoints(self):
+        """Test the `get_kpoints` method."""
+        kpt = KpointsData()
+        kpt.set_cell_from_structure(self.structure)
+
+        kpoints = [
+            [0.0, 0.0, 0.0],
+            [0.5, 0.5, 0.5],
+        ]
+
+        cartesian_kpoints = [
+            [0.0, 0.0, 0.0],
+            [np.pi / self.alat, np.pi / self.alat, np.pi / self.alat],
+        ]
+
+        kpt.set_kpoints(kpoints)
+        assert np.abs(kpt.get_kpoints() - np.array(kpoints)).sum() == 0.0
+        assert np.abs(kpt.get_kpoints(cartesian=True) - np.array(cartesian_kpoints)).sum() == 0.0
+
+        # Check also after storing
+        kpt.store()
+        kpt2 = load_node(kpt.pk)
+        assert np.abs(kpt2.get_kpoints() - np.array(kpoints)).sum() == 0.0
+        assert np.abs(kpt2.get_kpoints(cartesian=True) - np.array(cartesian_kpoints)).sum() == 0.0

tests/test_dataclasses.py

@@ -26,6 +26,7 @@
     get_formula,
     get_pymatgen_version,
     has_ase,
+    has_atomistic,
     has_pymatgen,
     has_spglib,
 )
@@ -67,6 +68,7 @@ def simplify(string):
 skip_spglib = pytest.mark.skipif(not has_spglib(), reason='Unable to import spglib')
 skip_pycifrw = pytest.mark.skipif(not has_pycifrw(), reason='Unable to import PyCifRW')
 skip_pymatgen = pytest.mark.skipif(not has_pymatgen(), reason='Unable to import pymatgen')
+skip_atomistic = pytest.mark.skipif(not has_atomistic(), reason='Unable to import aiida-atomistic')
 
 
 @skip_pymatgen
@@ -1851,6 +1853,26 @@ def test_clone(self):
             assert round(abs(c.sites[1].position[i] - 1.0), 7) == 0
 
 
+@skip_atomistic
+def test_to_atomistic(self):
+    """Test the conversion from orm.StructureData to the atomistic structure."""
+
+    # Create a structure with a single atom
+    from aiida_atomistic import StructureData as AtomisticStructureData
+
+    legacy = StructureData(cell=((1.0, 0.0, 0.0), (0.0, 2.0, 0.0), (0.0, 0.0, 3.0)))
+    legacy.append_atom(position=(0.0, 0.0, 0.0), symbols=['Ba'], name='Ba1')
+
+    # Convert to atomistic structure
+    structure = legacy.to_atomistic()
+
+    # Check that the structure is as expected
+    assert isinstance(structure, AtomisticStructureData)
+    assert structure.properties.sites[0].kinds == legacy.sites[0].kind_name
+    assert structure.properties.sites[0].positions == list(legacy.sites[0].position)
+    assert structure.properties.cell == legacy.cell
+
+
 class TestStructureDataFromAse:
     """Tests the creation of Sites from/to a ASE object."""