Implement method for JobCalculation to create a restart builder

aiida/backends/tests/work/job_processes.py

@@ -19,6 +19,7 @@
 from aiida.orm.calculation.job.simpleplugins.templatereplacer import TemplatereplacerCalculation
 from aiida.work.persistence import ObjectLoader
 from aiida.work.job_processes import JobProcess
+from aiida.work.process_builder import JobProcessBuilder
 from aiida.work import Process
 
 from . import utils
@@ -72,9 +73,20 @@ def tearDown(self):
         self.runner = None
         work.set_runner(None)
 
-    def test_class_loader(self):
-        cl = ObjectLoader()
-        TemplatereplacerProcess = JobProcess.build(TemplatereplacerCalculation)
+    def test_job_calculation_process(self):
+        """Verify that JobCalculation.process returns a sub class of JobProcess with correct calculation class."""
+        process = TemplatereplacerCalculation.process()
+        self.assertTrue(issubclass(process, JobProcess))
+        self.assertEqual(process._calc_class, TemplatereplacerCalculation)
+
+    def test_job_calculation_get_builder(self):
+        """Verify that JobCalculation.get_builder() returns an instance of JobProcessBuilder."""
+        process = TemplatereplacerCalculation.process()
+        builder = TemplatereplacerCalculation.get_builder()
+        self.assertTrue(isinstance(builder, JobProcessBuilder))
+
+        # Class objects are actually different memory instances so can't use assertEqual on simply instances
+        self.assertEqual(builder.process_class.__name__, process.__name__)
 
     def test_job_process_set_label_and_description(self):
         """

aiida/backends/tests/work/test_process_builder.py

@@ -110,6 +110,23 @@ def test_dynamic_getters_doc_string(self):
         self.assertEquals(builder.__class__.c.__doc__, str(TestWorkChain.spec().inputs['c']))
         self.assertEquals(builder.c.__class__.d.__doc__, str(TestWorkChain.spec().inputs['c']['d']))
 
+    def test_job_calculation_get_builder_restart(self):
+        """
+        Test the get_builder_restart method of JobCalculation class
+        """
+        from aiida.orm.calculation.job import JobCalculation
+
+        original = JobCalculation()
+        original.set_option('resources', {'num_machines': 1, 'num_mpiprocs_per_machine': 1})
+        original.set_option('max_wallclock_seconds', 1800)
+        original.set_computer(self.computer)
+        original.label = 'original'
+        original.store()
+
+        builder = original.get_builder_restart()
+
+        self.assertDictEqual(builder.options, original.get_options())
+
     def test_code_get_builder(self):
         """
         Test that the get_builder method of Code returns a builder

aiida/orm/implementation/general/calculation/job/__init__.py

@@ -112,13 +112,63 @@ def get_hash(self, ignore_errors=True, ignored_folder_content=('raw_input',), **
 
     @classmethod
     def process(cls):
+        """
+        Return the JobProcess class constructed based on this JobCalculatin class
+
+        :return: JobProcess class
+        """
         from aiida.work.job_processes import JobProcess
         return JobProcess.build(cls)
 
     @classmethod
     def get_builder(cls):
+        """
+        Return a JobProcessBuilder instance, tailored for this calculation class
+
+        This builder is a mapping of the inputs of the JobCalculation class, supports tab-completion, automatic
+        validation when settings values as well as automated docstrings for each input
+
+        :return: JobProcessBuilder instance
+        """
         return cls.process().get_builder()
 
+    def get_builder_restart(self):
+        """
+        Return a JobProcessBuilder instance, tailored for this calculation instance
+
+        This builder is a mapping of the inputs of the JobCalculation class, supports tab-completion, automatic
+        validation when settings values as well as automated docstrings for each input.
+
+        The fields of the builder will be pre-populated with all the inputs recorded for this instance as well as
+        settings all the options that were explicitly set for this calculation instance.
+
+        This builder can then directly be launched again to effectively run a duplicate calculation. But more useful
+        is that it serves as a starting point to, after changing one or more inputs, launch a similar calculation by
+        using this already completed calculation as a starting point.
+
+        :return: JobProcessBuilder instance
+        """
+        from aiida.work.job_processes import JobProcess
+        from aiida.work.ports import PortNamespace
+
+        inputs = self.get_inputs_dict()
+        options = self.get_options()
+        builder = self.get_builder()
+
+        for port_name, port in self.process().spec().inputs.items():
+            if port_name == JobProcess.OPTIONS_INPUT_LABEL:
+                setattr(builder, port_name, options)
+            elif isinstance(port, PortNamespace):
+                namespace = port_name + '_'
+                sub = {link[len(namespace):]: node for link, node in inputs.items() if link.startswith(namespace)}
+                if sub:
+                    setattr(builder, port_name, sub)
+            else:
+                if port_name in inputs:
+                    setattr(builder, port_name, inputs[port_name])
+
+        return builder
+
     def _init_internal_params(self):
         """
         Define here internal parameters that should be defined right after the __init__