Skip to content

Commit

Permalink
Providing support for atomistic StructureData
Browse files Browse the repository at this point in the history
- orm.StructureData `to_atomistic` method and `has_atomistic` function
- adapted `set_cell_from_structure` method for KpointsData
- added tests for both structure and kpoints.
  • Loading branch information
mikibonacci committed Nov 22, 2024
1 parent 779cc29 commit 1c24082
Show file tree
Hide file tree
Showing 4 changed files with 150 additions and 8 deletions.
24 changes: 16 additions & 8 deletions src/aiida/orm/nodes/data/array/kpoints.py
Original file line number Diff line number Diff line change
Expand Up @@ -223,16 +223,24 @@ 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__
)
)
cell = structuredata.cell
self.set_cell(cell, structuredata.pbc)
error_message = 'An instance of StructureData or aiida-atomistic StructureData should be passed to ' 'the KpointsData, found instead {}'.format(
structuredata.__class__
)
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.
Expand Down
34 changes: 34 additions & 0 deletions src/aiida/orm/nodes/data/structure.py
Original file line number Diff line number Diff line change
Expand Up @@ -101,6 +101,14 @@ def has_pymatgen():
return False
return True

def has_atomistic():
""":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."""
Expand Down Expand Up @@ -1876,6 +1884,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 as AtomisticStructureData
from aiida_atomistic import StructureDataMutable as AtomisticStructureDataMutable

atomistic = AtomisticStructureDataMutable()
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 AtomisticStructureData.from_mutable(atomistic)

class Kind:
"""This class contains the information about the species (kinds) of the system.
Expand Down
77 changes: 77 additions & 0 deletions tests/orm/nodes/data/test_kpoints.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +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:
"""Test for the `Kpointsdata` class."""
Expand Down Expand Up @@ -81,3 +83,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 StructureDataMutable
from aiida_atomistic import StructureData as AtomisticStructureData

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 = AtomisticStructureData.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
23 changes: 23 additions & 0 deletions tests/test_dataclasses.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
has_ase,
has_pymatgen,
has_spglib,
has_atomistic,
)


Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -1850,6 +1852,27 @@ def test_clone(self):
for i in range(3):
assert round(abs(c.sites[1].position[i] - 1.0), 7) == 0

@skip_atomistic
class TestAtomisticStructureData:
"""Tests the creation of StructureData objects (cell and pbc), and its conversion to the new atomistic StructureData."""

def test_to_atomistic(self):
"""Test the conversion 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."""
Expand Down

0 comments on commit 1c24082

Please sign in to comment.