Protocols: consider pbc of StructureData (#907)

Currently none of the `get_builder_from_protocol()` methods consider the periodic boundary conditions (`pbc`) of the input `StructureData`. Here we add support for several cases to the `PwBaseWorkChain` and `PwRelaxWorkChain`:

* 2D structures in the x-y plane with periodic boundary conditions for the first two lattice vectors (`pbc == [True, True, False]`). For this case, the `PwBaseWorkChain` protocol will set `SYSTEM.assume_isolated` to `2D`, and the `PwRelaxWorkChain` protocol will use `CELL.do_free = '2Dxy'` in case the `RelaxType` is `CELL` or `POSITIONS_CELL`.
* 1D structures defined on an arbitrary lattice vector will use the corresponding `CELL.do_free` setting (`'x'`, `'y'`, `'z'`) if  `RelaxType` is `CELL` or `POSITIONS_CELL`.

So far, no validation is done to see if the structure has the correct `cell` specified. This could potentially be added to the `PwCalculation` validation.

Co-authored-by: Marnik Bercx <mbercx@gmail.com>

src/aiida_quantumespresso/workflows/pw/base.py

@@ -180,6 +180,10 @@ def get_builder_from_protocol(
         parameters['SYSTEM']['ecutwfc'] = cutoff_wfc
         parameters['SYSTEM']['ecutrho'] = cutoff_rho
 
+        #If the structure is 2D periodic in the x-y plane, we set assume_isolate to `2D`
+        if structure.pbc == (True, True, False):
+            parameters['SYSTEM']['assume_isolated'] = '2D'
+
         if electronic_type is ElectronicType.INSULATOR:
             parameters['SYSTEM']['occupations'] = 'fixed'
             parameters['SYSTEM'].pop('degauss')

src/aiida_quantumespresso/workflows/pw/relax.py

@@ -143,7 +143,18 @@ def get_builder_from_protocol(
             base.pw.parameters['CELL']['cell_dofree'] = 'shape'
 
         if relax_type in (RelaxType.CELL, RelaxType.POSITIONS_CELL):
-            base.pw.parameters['CELL']['cell_dofree'] = 'all'
+
+            pbc_cell_dofree_map = {
+                (True, True, True): 'all',
+                (True, False, False): 'x',
+                (False, True, False): 'y',
+                (False, False, True): 'z',
+                (True, True, False): '2Dxy',
+            }
+            if structure.pbc in pbc_cell_dofree_map:
+                base.pw.parameters['CELL']['cell_dofree'] = pbc_cell_dofree_map[structure.pbc]
+            else:
+                raise ValueError(f'Structures with periodic boundary conditions `{structure.pbc}` are not supported.')
 
         builder = cls.get_builder()
         builder.base = base

tests/conftest.py

@@ -386,6 +386,23 @@ def _generate_structure(structure_id='silicon'):
             structure = StructureData(cell=cell)
             structure.append_atom(position=(0., 0., 0.), symbols='U', name='U')
             structure.append_atom(position=(param / 4., param / 4., param / 4.), symbols='U', name='U')
+        elif structure_id == '2D-xy-arsenic':
+            cell = [[3.61, 0, 0], [-1.80, 3.13, 0], [0, 0, 21.3]]
+            structure = StructureData(cell=cell, pbc=(True, True, False))
+            structure.append_atom(position=(1.804, 1.042, 11.352), symbols='As', name='As')
+            structure.append_atom(position=(0, 2.083, 9.960), symbols='As', name='As')
+        elif structure_id == '1D-x-carbon':
+            cell = [[4.2, 0, 0], [0, 20, 0], [0, 0, 20]]
+            structure = StructureData(cell=cell, pbc=(True, False, False))
+            structure.append_atom(position=(0, 0, 0), symbols='C', name='C')
+        elif structure_id == '1D-y-carbon':
+            cell = [[20, 0, 0], [0, 4.2, 0], [0, 0, 20]]
+            structure = StructureData(cell=cell, pbc=(False, True, False))
+            structure.append_atom(position=(0, 0, 0), symbols='C', name='C')
+        elif structure_id == '1D-z-carbon':
+            cell = [[20, 0, 0], [0, 20, 0], [0, 0, 4.2]]
+            structure = StructureData(cell=cell, pbc=(False, False, True))
+            structure.append_atom(position=(0, 0, 0), symbols='C', name='C')
         else:
             raise KeyError(f'Unknown structure_id="{structure_id}"')
         return structure

tests/workflows/protocols/pw/test_base.py

@@ -1,4 +1,5 @@
 # -*- coding: utf-8 -*-
+# pylint: disable=no-member,redefined-outer-name
 """Tests for the ``PwBaseWorkChain.get_builder_from_protocol`` method."""
 from aiida.engine import ProcessBuilder
 import pytest
@@ -67,6 +68,24 @@ def test_spin_type(fixture_code, generate_structure):
     assert parameters['SYSTEM']['starting_magnetization'] == {'Si': 0.1}
 
 
+@pytest.mark.parametrize(
+    'struc_name,assume_isolated', (
+        ('silicon', None),
+        ('2D-xy-arsenic', '2D'),
+        ('1D-x-carbon', None),
+        ('1D-y-carbon', None),
+        ('1D-z-carbon', None),
+    )
+)
+def test_pbc_assume_isolated(fixture_code, generate_structure, struc_name, assume_isolated):
+    """Test structures with various ``pbc`` set the correct ``assume_isolated``."""
+    code = fixture_code('quantumespresso.pw')
+    structure = generate_structure(struc_name)
+
+    builder = PwBaseWorkChain.get_builder_from_protocol(code, structure)
+    assert builder.pw.parameters['SYSTEM'].get('assume_isolated', None) == assume_isolated
+
+
 @pytest.mark.parametrize('initial_magnetic_moments', ({}, {'Si1': 1.0, 'Si2': 2.0}))
 def test_initial_magnetic_moments_invalid(fixture_code, generate_structure, initial_magnetic_moments):
     """Test ``PwBaseWorkChain.get_builder_from_protocol`` with invalid ``initial_magnetic_moments`` keyword."""

tests/workflows/protocols/pw/test_relax.py

@@ -126,3 +126,21 @@ def test_options(fixture_code, generate_structure):
         builder.base_final_scf.pw.metadata,
     ):
         assert subspace['options']['queue_name'] == queue_name
+
+
+@pytest.mark.parametrize(
+    'struc_name,cell_dofree', (
+        ('silicon', 'all'),
+        ('2D-xy-arsenic', '2Dxy'),
+        ('1D-x-carbon', 'x'),
+        ('1D-y-carbon', 'y'),
+        ('1D-z-carbon', 'z'),
+    )
+)
+def test_pbc_cell(fixture_code, generate_structure, struc_name, cell_dofree):
+    """Test structures with various ``pbc`` set the correct ``CELL`` parameters."""
+    code = fixture_code('quantumespresso.pw')
+    structure = generate_structure(struc_name)
+
+    builder = PwRelaxWorkChain.get_builder_from_protocol(code, structure)
+    assert builder.base.pw.parameters['CELL'].get('cell_dofree', None) == cell_dofree