From de02473d540b1eb26c7d90c31345854ab7357311 Mon Sep 17 00:00:00 2001
From: Marnik Bercx <mbercx@gmail.com>
Date: Fri, 15 Dec 2023 21:06:31 +0100
Subject: [PATCH] First round of review

---
 .../workflows/pw/relax.py                     | 30 ++++++++-----------
 1 file changed, 13 insertions(+), 17 deletions(-)

diff --git a/src/aiida_quantumespresso/workflows/pw/relax.py b/src/aiida_quantumespresso/workflows/pw/relax.py
index 80df5e90f..bff8ae16b 100644
--- a/src/aiida_quantumespresso/workflows/pw/relax.py
+++ b/src/aiida_quantumespresso/workflows/pw/relax.py
@@ -23,9 +23,6 @@ def define(cls, spec):
         """Define the process specification."""
         # yapf: disable
         super().define(spec)
-        spec.expose_inputs(PwBaseWorkChain, namespace='base_relax',
-            exclude=('clean_workdir', 'pw.structure', 'pw.parent_folder'),
-            namespace_options={'help': 'Inputs for the `PwBaseWorkChain` for the main relax loop.'})
         spec.expose_inputs(PwBaseWorkChain, namespace='base_init_relax',
             exclude=('clean_workdir', 'pw.structure', 'pw.parent_folder'),
             namespace_options={'required': False, 'populate_defaults': False,
@@ -33,6 +30,9 @@ def define(cls, spec):
                     'Inputs for the `PwBaseWorkChain` that runs an initial geometry optimization, typically with looser'
                     'precision settings to find a geometry that is already closer to the final one quickly.'
                 )})
+        spec.expose_inputs(PwBaseWorkChain, namespace='base_relax',
+            exclude=('clean_workdir', 'pw.structure', 'pw.parent_folder'),
+            namespace_options={'help': 'Inputs for the `PwBaseWorkChain` for the main relax loop.'})
         spec.input('structure', valid_type=orm.StructureData, help='The inputs structure.')
         spec.input('meta_convergence', valid_type=orm.Bool, default=lambda: orm.Bool(True),
             help='If `True` the workchain will perform a meta-convergence on the cell volume.')
@@ -55,10 +55,13 @@ def define(cls, spec):
             ),
             cls.results,
         )
+        spec.exit_code(400, 'ERROR_MAX_ITERATIONS_EXCEEDED',
+            message='The maximum number of meta convergence iterations was exceeded.')
         spec.exit_code(401, 'ERROR_SUB_PROCESS_FAILED_RELAX',
             message='the relax `PwBaseWorkChain` sub process failed')
         spec.exit_code(402, 'ERROR_SUB_PROCESS_FAILED_FINAL_SCF',
-            message='this exit code has been deprecated since the final SCF has been removed.')
+            message='the final scf PwBaseWorkChain sub process failed '
+                    '(deprecated since the final SCF has been removed)')
         spec.exit_code(403, 'ERROR_SUB_PROCESS_FAILED_INIT_RELAX',
             message='the initial relaxation `PwBaseWorkChain` sub process failed')
         spec.expose_outputs(PwBaseWorkChain, exclude=('output_structure',))
@@ -254,7 +257,7 @@ def should_run_relax(self):
         # Stop if the maximum number of meta iterations has been reached
         if self.ctx.iteration == self.inputs.max_meta_convergence_iterations.value:
             self.report('Maximum number of meta convergence iterations reached.')
-            return False
+            return self.exit_codes.ERROR_MAX_ITERATIONS_EXCEEDED
 
         base_relax_workchain = self.ctx.base_relax_workchains[-1]
 
@@ -278,10 +281,11 @@ def should_run_relax(self):
                 }
             }
             new_kpts_mesh, _ = create_kpoints_from_distance(**inputs_create_kpoints).get_kpoints_mesh()
-            if tuple(input_kpts_mesh) != tuple(new_kpts_mesh):
+
+            if not all(k1 <= k2 for k1, k2 in zip(new_kpts_mesh, input_kpts_mesh)):
                 self.report(
-                    'k-points mesh changed for specified `kpoints_distance` due to a change in the unit cell. '
-                    'Running another geometry optimization with new mesh.'
+                    'Density of k-points mesh has increased for the specified `kpoints_distance` due to a change in the'
+                    ' unit cell. Running another geometry optimization with new mesh.'
                 )
                 return True
 
@@ -293,11 +297,6 @@ def run_relax(self):
         self.ctx.iteration += 1
         inputs = self.ctx.relax_inputs
 
-        if 'base_init_relax_workchain' in self.ctx:
-            inputs.pw.structure = self.ctx.base_init_relax_workchain.outputs.output_structure
-        else:
-            inputs.pw.structure = self.ctx.current_structure
-
         # If one of the nested `PwBaseWorkChain`s changed the number of bands, apply it here
         if self.ctx.current_number_of_bands is not None:
             inputs.pw.parameters.setdefault('SYSTEM', {})['nbnd'] = self.ctx.current_number_of_bands
@@ -352,10 +351,7 @@ def results(self):
         if self.ctx.relax_inputs.pw.parameters['CONTROL']['calculation'] != 'scf':
             self.out('output_structure', final_relax_workchain.outputs.output_structure)
 
-        try:
-            self.out_many(self.exposed_outputs(self.ctx.workchain_scf, PwBaseWorkChain))
-        except AttributeError:
-            self.out_many(self.exposed_outputs(final_relax_workchain, PwBaseWorkChain))
+        self.out_many(self.exposed_outputs(final_relax_workchain, PwBaseWorkChain))
 
     def on_terminated(self):
         """Clean the working directories of all child calculations if `clean_workdir=True` in the inputs."""