From 4aefe3de10fe0edfc839fbeb11fce9919927acf6 Mon Sep 17 00:00:00 2001
From: bastonero <bastonero.lorenzo@gmail.com>
Date: Thu, 23 Nov 2023 17:10:14 +0000
Subject: [PATCH 1/2] Test/Improve `PhBaseWorkChain` overrides/protocol

The `PhBaseWorkChain.get_builder_from_protocol` was untested.
Tests are added, and the `electronic_type` input is added to allow
for the computation of the dielectric and effective charge tensors
when an insulator is in input. This is useful in other workchains, so
that the `electronic_type` input can be passed seamlessly as a
kwargs input at an higher level.
---
 .../workflows/ph/base.py                      | 21 +++-
 tests/conftest.py                             |  2 +-
 tests/workflows/protocols/ph/__init__.py      |  0
 tests/workflows/protocols/ph/test_base.py     | 98 +++++++++++++++++++
 .../protocols/ph/test_base/test_default.yml   | 20 ++++
 5 files changed, 139 insertions(+), 2 deletions(-)
 create mode 100644 tests/workflows/protocols/ph/__init__.py
 create mode 100644 tests/workflows/protocols/ph/test_base.py
 create mode 100644 tests/workflows/protocols/ph/test_base/test_default.yml

diff --git a/src/aiida_quantumespresso/workflows/ph/base.py b/src/aiida_quantumespresso/workflows/ph/base.py
index e673ccb52..81f80b79f 100644
--- a/src/aiida_quantumespresso/workflows/ph/base.py
+++ b/src/aiida_quantumespresso/workflows/ph/base.py
@@ -9,6 +9,7 @@
 from aiida.plugins import CalculationFactory
 
 from aiida_quantumespresso.calculations.functions.merge_ph_outputs import merge_ph_outputs
+from aiida_quantumespresso.common.types import ElectronicType
 from aiida_quantumespresso.workflows.protocols.utils import ProtocolMixin
 
 PhCalculation = CalculationFactory('quantumespresso.ph')
@@ -64,7 +65,16 @@ def get_protocol_filepath(cls):
         return files(ph_protocols) / 'base.yaml'
 
     @classmethod
-    def get_builder_from_protocol(cls, code, parent_folder=None, protocol=None, overrides=None, options=None, **_):
+    def get_builder_from_protocol(
+        cls,
+        code,
+        parent_folder=None,
+        protocol=None,
+        overrides=None,
+        electronic_type=ElectronicType.METAL,
+        options=None,
+        **_
+    ):
         """Return a builder prepopulated with inputs selected according to the chosen protocol.
 
         :param code: the ``Code`` instance configured for the ``quantumespresso.ph`` plugin.
@@ -73,6 +83,7 @@ def get_builder_from_protocol(cls, code, parent_folder=None, protocol=None, over
         :param overrides: optional dictionary of inputs to override the defaults of the protocol.
         :param options: A dictionary of options that will be recursively set for the ``metadata.options`` input of all
             the ``CalcJobs`` that are nested in this work chain.
+        :param electronic_type: indicate the electronic character of the system through ``ElectronicType`` instance.
         :return: a process builder instance with all inputs defined ready for launch.
         """
         from aiida_quantumespresso.workflows.protocols.utils import recursive_merge
@@ -81,9 +92,16 @@ def get_builder_from_protocol(cls, code, parent_folder=None, protocol=None, over
             code = orm.load_code(code)
 
         type_check(code, orm.AbstractCode)
+        type_check(electronic_type, ElectronicType)
+
+        if electronic_type not in [ElectronicType.METAL, ElectronicType.INSULATOR]:
+            raise NotImplementedError(f'electronic type `{electronic_type}` is not supported.')
 
         inputs = cls.get_protocol_inputs(protocol, overrides)
 
+        if electronic_type is ElectronicType.INSULATOR:
+            inputs['ph']['parameters']['INPUTPH']['epsil'] = True
+
         qpoints_mesh = inputs['ph'].pop('qpoints')
         qpoints = orm.KpointsData()
         qpoints.set_kpoints_mesh(qpoints_mesh)
@@ -104,6 +122,7 @@ def get_builder_from_protocol(cls, code, parent_folder=None, protocol=None, over
             builder.ph['settings'] = orm.Dict(inputs['ph']['settings'])
         builder.clean_workdir = orm.Bool(inputs['clean_workdir'])
         builder.ph['qpoints'] = qpoints
+        builder.max_iterations = orm.Int(inputs['max_iterations'])
         # pylint: enable=no-member
 
         return builder
diff --git a/tests/conftest.py b/tests/conftest.py
index a42e5c0ca..38c5e42db 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -470,7 +470,7 @@ def generate_remote_data():
     """Return a `RemoteData` node."""
 
     def _generate_remote_data(computer, remote_path, entry_point_name=None):
-        """Return a `KpointsData` with a mesh of npoints in each direction."""
+        """Return a `RemoteData` node."""
         from aiida.common.links import LinkType
         from aiida.orm import CalcJobNode, RemoteData
         from aiida.plugins.entry_point import format_entry_point_string
diff --git a/tests/workflows/protocols/ph/__init__.py b/tests/workflows/protocols/ph/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/workflows/protocols/ph/test_base.py b/tests/workflows/protocols/ph/test_base.py
new file mode 100644
index 000000000..74bb1b710
--- /dev/null
+++ b/tests/workflows/protocols/ph/test_base.py
@@ -0,0 +1,98 @@
+# -*- coding: utf-8 -*-
+# pylint: disable=no-member,redefined-outer-name
+"""Tests for the ``PhBaseWorkChain.get_builder_from_protocol`` method."""
+from aiida.engine import ProcessBuilder
+import pytest
+
+from aiida_quantumespresso.common.types import ElectronicType
+from aiida_quantumespresso.workflows.ph.base import PhBaseWorkChain
+
+
+def test_get_available_protocols():
+    """Test ``PhBaseWorkChain.get_available_protocols``."""
+    protocols = PhBaseWorkChain.get_available_protocols()
+    assert sorted(protocols.keys()) == ['fast', 'moderate', 'precise']
+    assert all('description' in protocol for protocol in protocols.values())
+
+
+def test_get_default_protocol():
+    """Test ``PhBaseWorkChain.get_default_protocol``."""
+    assert PhBaseWorkChain.get_default_protocol() == 'moderate'
+
+
+def test_default(fixture_code, data_regression, serialize_builder):
+    """Test ``PhBaseWorkChain.get_builder_from_protocol`` for the default protocol."""
+    code = fixture_code('quantumespresso.ph')
+    builder = PhBaseWorkChain.get_builder_from_protocol(code)
+
+    assert isinstance(builder, ProcessBuilder)
+    data_regression.check(serialize_builder(builder))
+
+
+def test_electronic_type(fixture_code):
+    """Test ``PhBaseWorkChain.get_builder_from_protocol`` with ``electronic_type`` keyword."""
+    code = fixture_code('quantumespresso.ph')
+
+    with pytest.raises(NotImplementedError):
+        for electronic_type in [ElectronicType.AUTOMATIC]:
+            PhBaseWorkChain.get_builder_from_protocol(code, electronic_type=electronic_type)
+
+    builder = PhBaseWorkChain.get_builder_from_protocol(code, electronic_type=ElectronicType.INSULATOR)
+    parameters = builder.ph.parameters.get_dict()  # pylint: disable=no-member
+
+    assert parameters['INPUTPH']['epsil']
+
+
+def test_parameter_overrides(fixture_code):
+    """Test specifying parameter ``overrides`` for the ``get_builder_from_protocol()`` method."""
+    code = fixture_code('quantumespresso.ph')
+
+    overrides = {'ph': {'parameters': {'INPUTHP': {'nmix_ph': 20}}}}
+    builder = PhBaseWorkChain.get_builder_from_protocol(code, overrides=overrides)
+    assert builder.ph.parameters['INPUTHP']['nmix_ph'] == 20  # pylint: disable=no-member
+
+
+def test_settings_overrides(fixture_code):
+    """Test specifying settings ``overrides`` for the ``get_builder_from_protocol()`` method."""
+    code = fixture_code('quantumespresso.ph')
+
+    overrides = {'ph': {'settings': {'cmdline': ['--kickass-mode']}}}
+    builder = PhBaseWorkChain.get_builder_from_protocol(code, overrides=overrides)
+    assert builder.ph.settings['cmdline'] == ['--kickass-mode']  # pylint: disable=no-member
+
+
+def test_metadata_overrides(fixture_code):
+    """Test specifying metadata ``overrides`` for the ``get_builder_from_protocol()`` method."""
+    code = fixture_code('quantumespresso.ph')
+
+    overrides = {'ph': {'metadata': {'options': {'resources': {'num_machines': 1e90}, 'max_wallclock_seconds': 1}}}}
+    builder = PhBaseWorkChain.get_builder_from_protocol(code, overrides=overrides)
+    metadata = builder.ph.metadata  # pylint: disable=no-member
+
+    assert metadata['options']['resources']['num_machines'] == 1e90
+    assert metadata['options']['max_wallclock_seconds'] == 1
+
+
+def test_options(fixture_code):
+    """Test specifying ``options`` for the ``get_builder_from_protocol()`` method."""
+    code = fixture_code('quantumespresso.ph')
+
+    queue_name = 'super-fast'
+    withmpi = False  # The protocol default is ``True``
+
+    options = {'queue_name': queue_name, 'withmpi': withmpi}
+    builder = PhBaseWorkChain.get_builder_from_protocol(code, options=options)
+    metadata = builder.ph.metadata  # pylint: disable=no-member
+
+    assert metadata['options']['queue_name'] == queue_name
+    assert metadata['options']['withmpi'] == withmpi
+
+
+def test_parent_folder(fixture_code, generate_remote_data, fixture_localhost, fixture_sandbox):
+    """Test specifying ``options`` for the ``get_builder_from_protocol()`` method."""
+    code = fixture_code('quantumespresso.ph')
+    remote_folder = generate_remote_data(fixture_localhost, fixture_sandbox.abspath, 'quantumespresso.pw')
+
+    builder = PhBaseWorkChain.get_builder_from_protocol(code, parent_folder=remote_folder)
+
+    assert builder.ph.parent_folder == remote_folder  # pylint: disable=no-member
diff --git a/tests/workflows/protocols/ph/test_base/test_default.yml b/tests/workflows/protocols/ph/test_base/test_default.yml
new file mode 100644
index 000000000..845fe3211
--- /dev/null
+++ b/tests/workflows/protocols/ph/test_base/test_default.yml
@@ -0,0 +1,20 @@
+clean_workdir: false
+max_iterations: 5
+ph:
+  code: test.quantumespresso.ph@localhost
+  metadata:
+    options:
+      max_wallclock_seconds: 43200
+      resources:
+        num_machines: 1
+      withmpi: true
+  parameters:
+    INPUTPH:
+      tr2_ph: 1.0e-18
+  qpoints:
+  - - 3
+    - 3
+    - 3
+  - - 0.0
+    - 0.0
+    - 0.0

From 1ff03a5159b48ab80a589894ee4dbf05537a9bd0 Mon Sep 17 00:00:00 2001
From: bastonero <bastonero.lorenzo@gmail.com>
Date: Thu, 23 Nov 2023 17:12:35 +0000
Subject: [PATCH 2/2] Fix missing `max_iterations` in overrides Fixes #966

All the base workchains missed to override the `max_iterations`
input. This is now fixed for all of them.
---
 src/aiida_quantumespresso/workflows/protocols/ph/base.yaml     | 1 +
 src/aiida_quantumespresso/workflows/protocols/pw/base.yaml     | 1 +
 .../workflows/protocols/xspectra/base.yaml                     | 1 +
 src/aiida_quantumespresso/workflows/pw/base.py                 | 1 +
 src/aiida_quantumespresso/workflows/xspectra/base.py           | 1 +
 tests/workflows/protocols/pw/test_bands/test_default.yml       | 3 +++
 tests/workflows/protocols/pw/test_base/test_default.yml        | 1 +
 tests/workflows/protocols/pw/test_relax/test_default.yml       | 2 ++
 tests/workflows/protocols/test_pdos/test_default.yml           | 2 ++
 tests/workflows/protocols/xspectra/test_base/test_default.yml  | 1 +
 tests/workflows/protocols/xspectra/test_core/test_default.yml  | 1 +
 .../workflows/protocols/xspectra/test_crystal/test_default.yml | 2 ++
 12 files changed, 17 insertions(+)

diff --git a/src/aiida_quantumespresso/workflows/protocols/ph/base.yaml b/src/aiida_quantumespresso/workflows/protocols/ph/base.yaml
index 503750d45..126c6b836 100644
--- a/src/aiida_quantumespresso/workflows/protocols/ph/base.yaml
+++ b/src/aiida_quantumespresso/workflows/protocols/ph/base.yaml
@@ -1,5 +1,6 @@
 default_inputs:
     clean_workdir: False
+    max_iterations: 5
     ph:
         metadata:
             options:
diff --git a/src/aiida_quantumespresso/workflows/protocols/pw/base.yaml b/src/aiida_quantumespresso/workflows/protocols/pw/base.yaml
index 200edd9da..b8f93db1c 100644
--- a/src/aiida_quantumespresso/workflows/protocols/pw/base.yaml
+++ b/src/aiida_quantumespresso/workflows/protocols/pw/base.yaml
@@ -2,6 +2,7 @@ default_inputs:
     clean_workdir: False
     kpoints_distance: 0.15
     kpoints_force_parity: False
+    max_iterations: 5
     meta_parameters:
         conv_thr_per_atom: 0.2e-9
         etot_conv_thr_per_atom: 1.e-5
diff --git a/src/aiida_quantumespresso/workflows/protocols/xspectra/base.yaml b/src/aiida_quantumespresso/workflows/protocols/xspectra/base.yaml
index b96eb9fe7..7cc855fe8 100644
--- a/src/aiida_quantumespresso/workflows/protocols/xspectra/base.yaml
+++ b/src/aiida_quantumespresso/workflows/protocols/xspectra/base.yaml
@@ -2,6 +2,7 @@ default_inputs:
     clean_workdir: False
     kpoints_distance: 0.15
     kpoints_force_parity: False
+    max_iterations: 5
     xspectra:
         parameters:
             INPUT_XSPECTRA:
diff --git a/src/aiida_quantumespresso/workflows/pw/base.py b/src/aiida_quantumespresso/workflows/pw/base.py
index cfea4d4a9..fccac5114 100644
--- a/src/aiida_quantumespresso/workflows/pw/base.py
+++ b/src/aiida_quantumespresso/workflows/pw/base.py
@@ -227,6 +227,7 @@ def get_builder_from_protocol(
         else:
             builder.kpoints_distance = orm.Float(inputs['kpoints_distance'])
         builder.kpoints_force_parity = orm.Bool(inputs['kpoints_force_parity'])
+        builder.max_iterations = orm.Int(inputs['max_iterations'])
         # pylint: enable=no-member
 
         return builder
diff --git a/src/aiida_quantumespresso/workflows/xspectra/base.py b/src/aiida_quantumespresso/workflows/xspectra/base.py
index 7ee8e6d35..52413850e 100644
--- a/src/aiida_quantumespresso/workflows/xspectra/base.py
+++ b/src/aiida_quantumespresso/workflows/xspectra/base.py
@@ -137,6 +137,7 @@ def get_builder_from_protocol(
         if 'settings' in inputs['xspectra']:
             builder.xspectra['settings'] = orm.Dict(inputs['xspectra']['settings'])
         builder.clean_workdir = orm.Bool(inputs['clean_workdir'])
+        builder.max_iterations = orm.Int(inputs['max_iterations'])
         # pylint: enable=no-member
 
         return builder
diff --git a/tests/workflows/protocols/pw/test_bands/test_default.yml b/tests/workflows/protocols/pw/test_bands/test_default.yml
index f3ab01bde..f1b73428b 100644
--- a/tests/workflows/protocols/pw/test_bands/test_default.yml
+++ b/tests/workflows/protocols/pw/test_bands/test_default.yml
@@ -1,4 +1,5 @@
 bands:
+  max_iterations: 5
   pw:
     code: test.quantumespresso.pw@localhost
     metadata:
@@ -39,6 +40,7 @@ relax:
   base:
     kpoints_distance: 0.15
     kpoints_force_parity: false
+    max_iterations: 5
     pw:
       code: test.quantumespresso.pw@localhost
       metadata:
@@ -77,6 +79,7 @@ relax:
 scf:
   kpoints_distance: 0.15
   kpoints_force_parity: false
+  max_iterations: 5
   pw:
     code: test.quantumespresso.pw@localhost
     metadata:
diff --git a/tests/workflows/protocols/pw/test_base/test_default.yml b/tests/workflows/protocols/pw/test_base/test_default.yml
index 3e08ba059..0e7179be1 100644
--- a/tests/workflows/protocols/pw/test_base/test_default.yml
+++ b/tests/workflows/protocols/pw/test_base/test_default.yml
@@ -1,6 +1,7 @@
 clean_workdir: false
 kpoints_distance: 0.15
 kpoints_force_parity: false
+max_iterations: 5
 pw:
   code: test.quantumespresso.pw@localhost
   metadata:
diff --git a/tests/workflows/protocols/pw/test_relax/test_default.yml b/tests/workflows/protocols/pw/test_relax/test_default.yml
index 10ad26638..56684a53a 100644
--- a/tests/workflows/protocols/pw/test_relax/test_default.yml
+++ b/tests/workflows/protocols/pw/test_relax/test_default.yml
@@ -1,6 +1,7 @@
 base:
   kpoints_distance: 0.15
   kpoints_force_parity: false
+  max_iterations: 5
   pw:
     code: test.quantumespresso.pw@localhost
     metadata:
@@ -36,6 +37,7 @@ base:
 base_final_scf:
   kpoints_distance: 0.15
   kpoints_force_parity: false
+  max_iterations: 5
   pw:
     code: test.quantumespresso.pw@localhost
     metadata:
diff --git a/tests/workflows/protocols/test_pdos/test_default.yml b/tests/workflows/protocols/test_pdos/test_default.yml
index ac1c27ad6..991b87fec 100644
--- a/tests/workflows/protocols/test_pdos/test_default.yml
+++ b/tests/workflows/protocols/test_pdos/test_default.yml
@@ -13,6 +13,7 @@ dos:
 nscf:
   kpoints_distance: 0.1
   kpoints_force_parity: false
+  max_iterations: 5
   pw:
     code: test.quantumespresso.pw@localhost
     metadata:
@@ -55,6 +56,7 @@ projwfc:
 scf:
   kpoints_distance: 0.15
   kpoints_force_parity: false
+  max_iterations: 5
   pw:
     code: test.quantumespresso.pw@localhost
     metadata:
diff --git a/tests/workflows/protocols/xspectra/test_base/test_default.yml b/tests/workflows/protocols/xspectra/test_base/test_default.yml
index ffefc8c2f..b41746ee8 100644
--- a/tests/workflows/protocols/xspectra/test_base/test_default.yml
+++ b/tests/workflows/protocols/xspectra/test_base/test_default.yml
@@ -6,6 +6,7 @@ kpoints:
 - - 0.0
   - 0.0
   - 0.0
+max_iterations: 5
 xspectra:
   code: test.quantumespresso.xspectra@localhost
   core_wfc_data: '# number of core states 3 =  1 0;  2 0;
diff --git a/tests/workflows/protocols/xspectra/test_core/test_default.yml b/tests/workflows/protocols/xspectra/test_core/test_default.yml
index 982fac60f..5b329bde6 100644
--- a/tests/workflows/protocols/xspectra/test_core/test_default.yml
+++ b/tests/workflows/protocols/xspectra/test_core/test_default.yml
@@ -19,6 +19,7 @@ get_powder_spectrum: false
 scf:
   kpoints_distance: 0.15
   kpoints_force_parity: false
+  max_iterations: 5
   pw:
     code: test.quantumespresso.pw@localhost
     metadata:
diff --git a/tests/workflows/protocols/xspectra/test_crystal/test_default.yml b/tests/workflows/protocols/xspectra/test_crystal/test_default.yml
index ee0c5bc9b..e2a310b3c 100644
--- a/tests/workflows/protocols/xspectra/test_crystal/test_default.yml
+++ b/tests/workflows/protocols/xspectra/test_crystal/test_default.yml
@@ -7,6 +7,7 @@ core:
     clean_workdir: false
     kpoints_distance: 0.15
     kpoints_force_parity: false
+    max_iterations: 5
     pw:
       code: test.quantumespresso.pw@localhost
       metadata:
@@ -79,6 +80,7 @@ relax:
   base:
     kpoints_distance: 0.15
     kpoints_force_parity: false
+    max_iterations: 5
     pw:
       code: test.quantumespresso.pw@localhost
       metadata: