From 450290f8ff2b3bb9b91406b98c5dd0d796ccf86b Mon Sep 17 00:00:00 2001
From: Sebastiaan Huber <mail@sphuber.net>
Date: Fri, 18 Nov 2022 16:51:42 +0100
Subject: [PATCH] `PwBaseWorkChain`: Set `restart_mode` in parameters if
 `parent_folder`

Typically, if a user launches a `PwBaseWorkChain` with a `parent_folder`
they want to restart from a previously completed but unconverged run. In
this case the `CONTROL.restart_mode` parameter should be set to
`from_scratch`. This used to be done automatically but was removed in
commit cb32be5361afa05bad617f00f8b187c96eb365ec.

Here we reinstate the behavior by calling `set_restart_type` to
`RestartType.FULL` if the `parent_folder` is defined and the
`CONTROL.restart_mode` hasn't already been set explicitly in the
parameters to `from_scratch`. In this case, it was likely set by the
caller and we don't want to silently override this.
---
 .../workflows/pw/base.py                      |  7 ++++++
 tests/workflows/pw/test_base.py               | 23 +++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/src/aiida_quantumespresso/workflows/pw/base.py b/src/aiida_quantumespresso/workflows/pw/base.py
index ad2a13c47..0a97c44fc 100644
--- a/src/aiida_quantumespresso/workflows/pw/base.py
+++ b/src/aiida_quantumespresso/workflows/pw/base.py
@@ -259,6 +259,13 @@ def setup(self):
 
         self.ctx.inputs.settings = self.ctx.inputs.settings.get_dict() if 'settings' in self.ctx.inputs else {}
 
+        # If a ``parent_folder`` is specified, automatically set the parameters for a ``RestartType.Full`` unless the
+        # ``CONTROL.restart_mode`` has explicitly been set to ``from_scratch``. In that case, the user most likely set
+        # that, and we do not want to override it.
+        restart_mode = self.ctx.inputs.parameters['CONTROL'].get('restart_mode', None)
+        if 'parent_folder' in self.ctx.inputs and restart_mode != 'from_scratch':
+            self.set_restart_type(RestartType.FULL, self.ctx.inputs.parent_folder)
+
     def validate_kpoints(self):
         """Validate the inputs related to k-points.
 
diff --git a/tests/workflows/pw/test_base.py b/tests/workflows/pw/test_base.py
index ec0341a18..3d8cbd3b3 100644
--- a/tests/workflows/pw/test_base.py
+++ b/tests/workflows/pw/test_base.py
@@ -3,6 +3,7 @@
 """Tests for the `PwBaseWorkChain` class."""
 from aiida.common import AttributeDict
 from aiida.engine import ExitCode, ProcessHandlerReport
+from aiida.orm import Dict
 import pytest
 
 from aiida_quantumespresso.calculations.pw import PwCalculation
@@ -230,3 +231,25 @@ def test_set_max_seconds(generate_workchain_pw):
 
     assert 'max_seconds' in process.ctx.inputs['parameters']['CONTROL']
     assert process.ctx.inputs['parameters']['CONTROL']['max_seconds'] == max_seconds
+
+
+@pytest.mark.parametrize('restart_mode, expected', (
+    (None, 'restart'),
+    ('from_scratch', 'from_scratch'),
+))
+def test_parent_folder(generate_workchain_pw, generate_calc_job_node, restart_mode, expected):
+    """Test that ``parameters`` gets automatically updated if ``parent_folder`` in the inputs.
+
+    Specifically, the ``parameters`` should define the ``CONTROL.restart_mode`` unless it was explicitly set to
+    ``from_scratch`` by the caller.
+    """
+    node = generate_calc_job_node('pw', test_name='default')
+
+    inputs = generate_workchain_pw(return_inputs=True)
+    inputs['pw']['parent_folder'] = node.outputs.remote_folder
+    inputs['pw']['parameters'] = Dict({'CONTROL': {'restart_mode': restart_mode}})
+
+    process = generate_workchain_pw(inputs=inputs)
+    process.setup()
+
+    assert process.ctx.inputs['parameters']['CONTROL']['restart_mode'] == expected