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