adapting BandsWorkChain for HubbardStructureData

src/aiida_quantumespresso/data/hubbard_structure.py

@@ -5,6 +5,7 @@
 
 from aiida.orm import StructureData
 import numpy as np
+from pymatgen.core import Lattice, PeriodicSite
 
 from aiida_quantumespresso.common.hubbard import Hubbard, HubbardParameters
 
@@ -58,6 +59,7 @@ def hubbard(self) -> Hubbard:
 
         :returns: a :class:`~aiida_quantumespresso.common.hubbard.Hubbard` instance.
         """
+        # pylint: disable=not-context-manager
         with self.base.repository.open(self._hubbard_filename, mode='rb') as handle:
             return Hubbard.model_validate_json(json.load(handle))
 
@@ -109,8 +111,19 @@ def append_hubbard_parameter(
         :param hubbard_type: hubbard type (U, V, J, ...), defaults to 'Ueff'
             (see :class:`~aiida_quantumespresso.common.hubbard.Hubbard` for full allowed values)
         """
-        pymat = self.get_pymatgen_structure()
-        sites = pymat.sites
+        sites = [
+            PeriodicSite(
+                species=site.species,
+                coords=site.coords,
+                lattice=Lattice(self.cell, pbc=self.pbc),
+                coords_are_cartesian=True
+            ) for site in self.get_pymatgen().sites
+        ]
+
+        if any((atom_index > len(sites) - 1, neighbour_index > len(sites) - 1)):
+            raise ValueError(
+                'atom_index and neighbour_index must be within the range of the number of sites in the structure'
+            )
 
         if translation is None:
             _, translation = sites[atom_index].distance_and_image(sites[neighbour_index])

tests/conftest.py

@@ -375,6 +375,10 @@ def _generate_structure(structure_id='silicon'):
             structure = StructureData(cell=cell)
             structure.append_atom(position=(0., 0., 0.), symbols='Si', name=name1)
             structure.append_atom(position=(param / 4., param / 4., param / 4.), symbols='Si', name=name2)
+        elif structure_id == 'cobalt-prim':
+            cell = [[0.0, 2.715, 2.715], [2.715, 0.0, 2.715], [2.715, 2.715, 0.0]]
+            structure = StructureData(cell=cell)
+            structure.append_atom(position=(0.0, 0.0, 0.0), symbols='Co', name='Co')
         elif structure_id == 'water':
             structure = StructureData(cell=[[5.29177209, 0., 0.], [0., 5.29177209, 0.], [0., 0., 5.29177209]])
             structure.append_atom(position=[12.73464656, 16.7741411, 24.35076238], symbols='H', name='H')

tests/data/test_hubbard_structure.py

@@ -62,15 +62,43 @@ def test_from_structure(generate_structure, generate_hubbard):
 
 
 @pytest.mark.usefixtures('aiida_profile')
-def test_append_hubbard_parameters(generate_hubbard_structure):
+@pytest.mark.parametrize('structure_name', ('silicon', '2D-xy-arsenic'))
+@pytest.mark.parametrize(
+    'parameters', (
+        ((0, '1s', 0, '1s', 5.0, (0, 0, 0), 'Ueff'),),
+        ((0, '1s', 1, '1s', 5.0, None, 'V'),),
+        (
+            (0, '1s', 0, '1s', 5.0, (0, 0, 0), 'Ueff'),
+            (0, '1s', 0, '1s', 5.0, (0, 0, 0), 'Ueff'),
+        ),
+        (
+            (0, '1s', 0, '1s', 5.0, (0, 0, 0), 'Ueff'),
+            (0, '1s', 1, '1s', 5.0, (0, 1, 0), 'V'),
+        ),
+    )
+)
+def test_append_hubbard_parameters(data_regression, generate_structure, structure_name, parameters):
     """Test the `append_hubbard_parameters` method."""
-    from aiida_quantumespresso.common.hubbard import HubbardParameters
-    hubbard_structure = generate_hubbard_structure()
-    args = (0, '1s', 1, '1s', 5.0, (0, 0, 0), 'U')
-    hubbard_structure.append_hubbard_parameter(*args)
-    params = HubbardParameters.from_tuple(args)
-    assert len(hubbard_structure.hubbard.parameters) == 2
-    assert params == hubbard_structure.hubbard.parameters[1]
+    hubbard_structure = HubbardStructureData.from_structure(generate_structure(structure_name))
+
+    for parameter in parameters:
+        hubbard_structure.append_hubbard_parameter(*parameter)
+
+    data_regression.check(hubbard_structure.hubbard.to_list())
+    assert len(hubbard_structure.hubbard.parameters) == len(set(parameters))
+
+
+@pytest.mark.parametrize('structure_name', ('cobalt-prim', '1D-x-carbon'))
+@pytest.mark.parametrize('parameter', (
+    (0, '1s', 1, '1s', 5.0, None, 'V'),
+    (0, '1s', 1, '1s', 5.0, (0, 0, 0), 'V'),
+))
+def test_append_hubbard_parameters_invalid_index(generate_structure, structure_name, parameter):
+    """Test the `append_hubbard_parameters` method with invalid index."""
+    hubbard_structure = HubbardStructureData.from_structure(generate_structure(structure_name))
+
+    with pytest.raises(ValueError, match='atom_index and neighbour_index must be within the range'):
+        hubbard_structure.append_hubbard_parameter(*parameter)
 
 
 @pytest.mark.usefixtures('aiida_profile')

tests/data/test_hubbard_structure/test_append_hubbard_parameters_parameters0_2D_xy_arsenic_.yml

@@ -0,0 +1,9 @@
+- - 0
+  - 1s
+  - 0
+  - 1s
+  - 5.0
+  - - 0
+    - 0
+    - 0
+  - Ueff

tests/data/test_hubbard_structure/test_append_hubbard_parameters_parameters0_silicon_.yml

@@ -0,0 +1,9 @@
+- - 0
+  - 1s
+  - 0
+  - 1s
+  - 5.0
+  - - 0
+    - 0
+    - 0
+  - Ueff

tests/data/test_hubbard_structure/test_append_hubbard_parameters_parameters1_2D_xy_arsenic_.yml

@@ -0,0 +1,9 @@
+- - 0
+  - 1s
+  - 1
+  - 1s
+  - 5.0
+  - - 0
+    - 0
+    - 0
+  - V

tests/data/test_hubbard_structure/test_append_hubbard_parameters_parameters1_silicon_.yml

@@ -0,0 +1,9 @@
+- - 0
+  - 1s
+  - 1
+  - 1s
+  - 5.0
+  - - -1
+    - 0
+    - 0
+  - V

tests/data/test_hubbard_structure/test_append_hubbard_parameters_parameters2_2D_xy_arsenic_.yml

@@ -0,0 +1,9 @@
+- - 0
+  - 1s
+  - 0
+  - 1s
+  - 5.0
+  - - 0
+    - 0
+    - 0
+  - Ueff

tests/data/test_hubbard_structure/test_append_hubbard_parameters_parameters2_silicon_.yml

@@ -0,0 +1,9 @@
+- - 0
+  - 1s
+  - 0
+  - 1s
+  - 5.0
+  - - 0
+    - 0
+    - 0
+  - Ueff

tests/data/test_hubbard_structure/test_append_hubbard_parameters_parameters3_2D_xy_arsenic_.yml

@@ -0,0 +1,18 @@
+- - 0
+  - 1s
+  - 0
+  - 1s
+  - 5.0
+  - - 0
+    - 0
+    - 0
+  - Ueff
+- - 0
+  - 1s
+  - 1
+  - 1s
+  - 5.0
+  - - 0
+    - 1
+    - 0
+  - V

tests/data/test_hubbard_structure/test_append_hubbard_parameters_parameters3_silicon_.yml

@@ -0,0 +1,18 @@
+- - 0
+  - 1s
+  - 0
+  - 1s
+  - 5.0
+  - - 0
+    - 0
+    - 0
+  - Ueff
+- - 0
+  - 1s
+  - 1
+  - 1s
+  - 5.0
+  - - 0
+    - 1
+    - 0
+  - V