From 0ebfd5d4262075b20a531fde733c13dfc30568ce Mon Sep 17 00:00:00 2001
From: Lorenzo <79980269+bastonero@users.noreply.github.com>
Date: Wed, 6 Dec 2023 12:09:58 +0100
Subject: [PATCH] Feat: clean the workdirs at each iteration (#47)

Modify the `SelfConsistentHubbardWorkChain`
to clean the work directories at each iteration instead
of cleaning everything at the end, eventually saving
lots of storage space. This may be crucial for large scale
calculations or for users with low disk availability.
---
 .../workflows/hubbard.py                      | 21 ++++++++++++-------
 1 file changed, 13 insertions(+), 8 deletions(-)

diff --git a/src/aiida_quantumespresso_hp/workflows/hubbard.py b/src/aiida_quantumespresso_hp/workflows/hubbard.py
index 7ea2506..529c982 100644
--- a/src/aiida_quantumespresso_hp/workflows/hubbard.py
+++ b/src/aiida_quantumespresso_hp/workflows/hubbard.py
@@ -172,6 +172,7 @@ def define(cls, spec):
         spec.inputs.validator = validate_inputs
         spec.inputs['hubbard']['hp'].validator = None
 
+        # yapf: disable
         spec.outline(
             cls.setup,
             while_(cls.should_run_iteration)(
@@ -188,10 +189,16 @@ def define(cls, spec):
                 ),
                 cls.run_hp,
                 cls.inspect_hp,
-                if_(cls.should_check_convergence)(cls.check_convergence,),
+                if_(cls.should_check_convergence)(
+                    cls.check_convergence,
+                ),
+                if_(cls.should_clean_workdir)(
+                    cls.clean_iteration,
+                )
             ),
             cls.run_results,
         )
+        # yapf: enable
 
         spec.output(
             'hubbard_structure',
@@ -673,14 +680,12 @@ def run_results(self):
         self.report(f'Hubbard parameters self-consistently converged in {self.ctx.iteration} iterations')
         self.out('hubbard_structure', self.ctx.current_hubbard_structure)
 
-    def on_terminated(self):
-        """Clean the working directories of all child calculations if `clean_workdir=True` in the inputs."""
-        super().on_terminated()
-
-        if self.inputs.clean_workdir.value is False:
-            self.report('remote folders will not be cleaned')
-            return
+    def should_clean_workdir(self):
+        """Whether to clean the work directories at each iteration."""
+        return self.inputs.clean_workdir.value
 
+    def clean_iteration(self):
+        """Clean all work directiories of the current iteration."""
         cleaned_calcs = []
 
         for called_descendant in self.node.called_descendants: