👌 PwBaseWorkChain: Improve handling of SCF convergence issues

src/aiida_quantumespresso/workflows/protocols/pw/base.yaml

@@ -24,7 +24,7 @@ default_inputs:
                 smearing: cold
                 degauss: 0.02
             ELECTRONS:
-                electron_maxstep: 80
+                electron_maxstep: 50
                 mixing_beta: 0.4
 default_protocol: balanced
 protocols:

src/aiida_quantumespresso/workflows/pw/base.py

@@ -29,11 +29,15 @@ class PwBaseWorkChain(ProtocolMixin, BaseRestartWorkChain):
         'qe': qe_defaults,
         'delta_threshold_degauss': 30,
         'delta_factor_degauss': 0.1,
-        'delta_factor_mixing_beta': 0.8,
+        'delta_factor_mixing_beta': 0.5,
+        'delta_addition_mixing_ndim': 4,
         'delta_factor_max_seconds': 0.95,
         'delta_factor_nbnd': 0.05,
         'delta_minimum_nbnd': 4,
         'delta_factor_trust_radius_min': 0.1,
+        'conv_slope_threshold': -0.1,
+        'conv_slope_range': 30,
+        'low_symmetry_threshold': 6
     })
 
     @classmethod
@@ -635,15 +639,87 @@ def handle_electronic_convergence_not_reached(self, calculation):
 
         Decrease the mixing beta and fully restart from the previous calculation.
         """
-        factor = self.defaults.delta_factor_mixing_beta
-        mixing_beta = self.ctx.inputs.parameters.get('ELECTRONS', {}).get('mixing_beta', self.defaults.qe.mixing_beta)
-        mixing_beta_new = mixing_beta * factor
+        import numpy
 
-        self.ctx.inputs.parameters['ELECTRONS']['mixing_beta'] = mixing_beta_new
-        action = f'reduced beta mixing from {mixing_beta} to {mixing_beta_new} and restarting from the last calculation'
+        scf_accuracy = calculation.tools.get_scf_accuracy(-1)[-self.defaults.conv_slope_range:]
+        scf_accuracy_slope = numpy.polyfit(numpy.arange(0, len(scf_accuracy)), numpy.log(scf_accuracy), 1)[0]
+
+        mixing_beta = self.ctx.inputs.parameters['ELECTRONS'].get('mixing_beta', self.defaults.qe.mixing_beta)
+        mixing_ndim = self.ctx.inputs.parameters['ELECTRONS'].get('mixing_ndim', self.defaults.qe.mixing_ndim)
+        mixing_mode = self.ctx.inputs.parameters['ELECTRONS'].get('mixing_mode', self.defaults.qe.mixing_mode)
+        low_symmetry_structure = (
+            len(calculation.outputs.output_parameters.get_dict()['symmetries']) < self.defaults.low_symmetry_threshold
+        )
+        low_dim_structure = calculation.inputs.structure.pbc != (True, True, True)
+        nbnd_cur = calculation.outputs.output_parameters.get_dict()['number_of_bands']
+        diagonalization = self.ctx.inputs.parameters['ELECTRONS'].get('diagonalization', 'david')
+
+        self.report(f'number of bands: {nbnd_cur}')
+        self.report(f'scf accuracy slope: {scf_accuracy_slope:.2f}')
+        self.report(f'mixing beta: {mixing_beta:.2f}')
+        self.report(f'mixing ndim: {mixing_ndim}')
+        self.report(f'mixing mode: {mixing_mode}')
+        self.report(f"structure symmetries: {len(calculation.outputs.output_parameters.get_dict()['symmetries'])}")
+        self.report(f'low symmetry structure: {low_symmetry_structure}')
+        self.report(f'low dimension structure: {low_dim_structure}')
+
+        if 'scf_failed_once' not in self.ctx:
+            nbnd_new = nbnd_cur + max(int(nbnd_cur * self.defaults.delta_factor_nbnd), self.defaults.delta_minimum_nbnd)
+            self.ctx.inputs.parameters['SYSTEM']['nbnd'] = nbnd_new
+            self.report(
+                f'First SCF failure encountered: increasing number of bands to {nbnd_new}'
+            )
+            self.ctx.scf_failed_once = True
+
+        if scf_accuracy_slope < self.defaults.conv_slope_threshold:
+            action = (
+                f'electronic convergence not reached but the scf accuracy slope ({scf_accuracy_slope:.2f}) is smaller '
+                f'than the threshold ({self.defaults.conv_slope_threshold:.2f}): restart from the last calculation.'
+            )
+            self.set_restart_type(RestartType.FROM_CHARGE_DENSITY, calculation.outputs.remote_folder)
+            self.report_error_handled(calculation, action)
+            return ProcessHandlerReport(True)
+
+        if mixing_mode == 'plain' and low_dim_structure:
+
+            self.ctx.inputs.parameters['ELECTRONS'].setdefault('mixing_mode', 'local-TF')
+            action = (
+                'electronic convergence not reached and structure is low dimensional: switch to local-TF mixing and '
+                'restart from the last calculation.'
+            )
+            self.set_restart_type(RestartType.FROM_CHARGE_DENSITY, calculation.outputs.remote_folder)
+            self.report_error_handled(calculation, action)
+            return ProcessHandlerReport(True)
+
+        if 'diagonalizations' not in self.ctx:
+            # Initialize a list to track diagonalisations that haven't been tried in reverse order or preference
+            self.ctx.diagonalizations = [value for value in ['cg', 'paro', 'ppcg', 'rmm-paro', 'david'] if value != diagonalization.lower()]
+
+        try:
+            new = self.ctx.diagonalizations.pop()
+            self.ctx.inputs.parameters['ELECTRONS']['diagonalization'] = new
+            action = f'electronic convergence not reached: switching to `{new}` diagonalization.'
+            self.set_restart_type(RestartType.FROM_CHARGE_DENSITY, calculation.outputs.remote_folder)
+            self.report_error_handled(calculation, action)
+            return ProcessHandlerReport(True)
+        except IndexError:
+            pass
+
+        if mixing_beta > 0.1:
+
+            mixing_beta_new = mixing_beta * self.defaults.delta_factor_mixing_beta
+            mixing_ndim_new = mixing_ndim + self.defaults.delta_addition_mixing_ndim
+
+            self.ctx.inputs.parameters['ELECTRONS']['mixing_beta'] = mixing_beta_new
+            self.ctx.inputs.parameters['ELECTRONS']['mixing_ndim'] = mixing_ndim_new
+            action = (
+                f'reduced beta mixing from {mixing_beta} to {mixing_beta_new}, increased `mixing_ndim` from '
+                f'{mixing_ndim} to {mixing_ndim_new} and restarting from the last calculation.'
+            )
+            self.set_restart_type(RestartType.FROM_CHARGE_DENSITY, calculation.outputs.remote_folder)
+            self.report_error_handled(calculation, action)
+            return ProcessHandlerReport(True)
 
-        self.set_restart_type(RestartType.FULL, calculation.outputs.remote_folder)
-        self.report_error_handled(calculation, action)
         return ProcessHandlerReport(True)
 
     @process_handler(priority=420, exit_codes=[

tests/conftest.py

@@ -245,6 +245,7 @@ def _generate_calc_job_node(
             filename = os.path.join(entry_point_name[len('quantumespresso.'):], test_name)
             filepath_folder = os.path.join(basepath, 'parsers', 'fixtures', filename)
             filepath_input = os.path.join(filepath_folder, 'aiida.in')
+            print(basepath)
 
         entry_point = format_entry_point_string('aiida.calculations', entry_point_name)
 
@@ -748,7 +749,7 @@ def _generate_inputs_pw():
         structure = generate_structure()
         inputs = {
             'code': fixture_code('quantumespresso.pw'),
-            'structure': generate_structure(),
+            'structure': structure,
             'kpoints': generate_kpoints_mesh(2),
             'parameters': parameters,
             'pseudos': {kind: generate_upf_data(kind) for kind in structure.get_kind_names()},

tests/workflows/protocols/pw/test_bands/test_default.yml

@@ -21,7 +21,7 @@ bands:
         conv_thr: 4.0e-10
         diago_full_acc: true
         diagonalization: paro
-        electron_maxstep: 80
+        electron_maxstep: 50
         mixing_beta: 0.4
         startingpot: file
       SYSTEM:
@@ -98,7 +98,7 @@ relax:
           tstress: true
         ELECTRONS:
           conv_thr: 4.0e-10
-          electron_maxstep: 80
+          electron_maxstep: 50
           mixing_beta: 0.4
         SYSTEM:
           degauss: 0.02
@@ -133,7 +133,7 @@ scf:
         tstress: true
       ELECTRONS:
         conv_thr: 4.0e-10
-        electron_maxstep: 80
+        electron_maxstep: 50
         mixing_beta: 0.4
       SYSTEM:
         degauss: 0.02

tests/workflows/protocols/pw/test_base/test_default.yml

@@ -19,7 +19,7 @@ pw:
       tstress: true
     ELECTRONS:
       conv_thr: 4.0e-10
-      electron_maxstep: 80
+      electron_maxstep: 50
       mixing_beta: 0.4
     SYSTEM:
       degauss: 0.02

tests/workflows/protocols/pw/test_relax/test_default.yml

@@ -59,7 +59,40 @@ base_relax:
         tstress: true
       ELECTRONS:
         conv_thr: 4.0e-10
-        electron_maxstep: 80
+        electron_maxstep: 50
+        mixing_beta: 0.4
+      SYSTEM:
+        degauss: 0.02
+        ecutrho: 240.0
+        ecutwfc: 30.0
+        nosym: false
+        occupations: smearing
+        smearing: cold
+    pseudos:
+      Si: Si<md5=57fa15d98af99972c7b7aa5c179b0bb8>
+base_final_scf:
+  kpoints_distance: 0.15
+  kpoints_force_parity: false
+  max_iterations: 5
+  pw:
+    code: test.quantumespresso.pw@localhost
+    metadata:
+      options:
+        max_wallclock_seconds: 43200
+        resources:
+          num_machines: 1
+          num_mpiprocs_per_machine: 1
+        withmpi: true
+    parameters:
+      CONTROL:
+        calculation: scf
+        etot_conv_thr: 2.0e-05
+        forc_conv_thr: 0.0001
+        tprnfor: true
+        tstress: true
+      ELECTRONS:
+        conv_thr: 4.0e-10
+        electron_maxstep: 50
         mixing_beta: 0.4
       SYSTEM:
         degauss: 0.02

tests/workflows/protocols/test_pdos/test_default.yml

@@ -33,7 +33,7 @@ nscf:
         tstress: true
       ELECTRONS:
         conv_thr: 4.0e-10
-        electron_maxstep: 80
+        electron_maxstep: 50
         mixing_beta: 0.4
       SYSTEM:
         ecutrho: 240.0
@@ -76,7 +76,7 @@ scf:
         tstress: true
       ELECTRONS:
         conv_thr: 4.0e-10
-        electron_maxstep: 80
+        electron_maxstep: 50
         mixing_beta: 0.4
       SYSTEM:
         degauss: 0.02

tests/workflows/pw/test_base.py

@@ -52,19 +52,29 @@ def test_handle_out_of_walltime(
         assert process.ctx.inputs.structure == output_structure
 
 
-def test_handle_electronic_convergence_not_reached(generate_workchain_pw, fixture_localhost, generate_remote_data):
+def test_handle_electronic_convergence_not_reached(
+        generate_workchain_pw, fixture_localhost, generate_remote_data, generate_structure,
+        generate_calc_job_node
+    ):
     """Test `PwBaseWorkChain.handle_electronic_convergence_not_reached`."""
     remote_data = generate_remote_data(computer=fixture_localhost, remote_path='/path/to/remote')
 
     process = generate_workchain_pw(
         exit_code=PwCalculation.exit_codes.ERROR_ELECTRONIC_CONVERGENCE_NOT_REACHED,
-        pw_outputs={'remote_folder': remote_data}
+        pw_outputs={'remote_folder': remote_data},
     )
     process.setup()
 
+    cj_node = generate_calc_job_node('quantumespresso.pw', fixture_localhost, 'test')
+
+    print(cj_node.outputs.output_parameters.get_dict()) 
+
+    process.ctx.structure = generate_structure('2D-xy-arsenic')
+
     process.ctx.inputs.parameters['ELECTRONS']['mixing_beta'] = 0.5
 
     result = process.handle_electronic_convergence_not_reached(process.ctx.children[-1])
+
     assert isinstance(result, ProcessHandlerReport)
     assert process.ctx.inputs.parameters['ELECTRONS']['mixing_beta'] == \
         process.defaults.delta_factor_mixing_beta * 0.5