Refinement MVP

src/aiida_dftk/calculations.py

@@ -63,6 +63,7 @@ def define(cls, spec):
         spec.exit_code(102, 'ERROR_MISSING_FORCES_FILE', message='The output file containing forces is missing.')
         spec.exit_code(103, 'ERROR_MISSING_STRESSES_FILE', message='The output file containing stresses is missing.')
         spec.exit_code(104, 'ERROR_MISSING_BANDS_FILE',message='The output file containing bands is missing.')
+        spec.exit_code(105, 'ERROR_MISSING_REFINEMENT_FILE', message='The output file containing refinement data is missing.')
         spec.exit_code(500, 'ERROR_SCF_CONVERGENCE_NOT_REACHED', message='The SCF minimization cycle did not converge, and the POSTSCF functions were not executed.')
         spec.exit_code(501, 'ERROR_SCF_OUT_OF_WALLTIME',message='The SCF was interuptted due to out of walltime. Non-recovarable error.')
         spec.exit_code(502, 'ERROR_POSTSCF_OUT_OF_WALLTIME',message='The POSTSCF was interuptted due to out of walltime.')
@@ -81,6 +82,7 @@ def define(cls, spec):
         spec.output('output_forces', valid_type=orm.ArrayData, required=False, help='forces array')
         spec.output('output_stresses', valid_type=orm.ArrayData, required=False, help='stresses array')
         spec.output('output_bands', valid_type=orm.BandsData, required=False, help='bandstructure')
+        spec.output('output_refinement', valid_type=orm.Dict, required=False, help='refinement')
 
         # TODO: bands and DOS implementation required on DFTK side
         # spec.output('output_bands', valid_type=orm.BandsData, required=False,
@@ -202,6 +204,8 @@ def _generate_retrieve_list(self, parameters: orm.Dict) -> list:
             f"{item['$function']}.json" if item['$function'] == 'compute_bands' else f"{item['$function']}.hdf5"
             for item in parameters['postscf']
         ]
+        if 'refinement' in parameters:
+            retrieve_list.append('refinement.json')
         retrieve_list.append(self.LOGFILE)
         retrieve_list.append('timings.json')
         retrieve_list.append(f'{self.SCFRES_SUMMARY_NAME}')

src/aiida_dftk/parsers.py

@@ -76,6 +76,12 @@ def parse(self, **kwargs):
                 self.exit_codes.ERROR_MISSING_BANDS_FILE,
                 self._parse_output_bands,
             )
+
+            self._parse_optional_result(
+                'refinement.json',
+                self.exit_codes.ERROR_MISSING_REFINEMENT_FILE,
+                self._parse_output_refinement,
+            )
         except ParsingFailedException as e:
             return e.exitcode
 
@@ -165,6 +171,12 @@ def _parse_output_bands(self, file_path):
 
         return None
 
+    def _parse_output_refinement(self, file_path):
+        with open(file_path, 'r') as f:
+            refinement_dict = json.load(f)
+        self.out('output_refinement', Dict(refinement_dict))
+        return None
+
     @staticmethod
     def _hdf5_to_dict(hdf5_file):
         """Convert an HDF5 file to a Python dictionary.

tests/test_workflows.py

@@ -39,7 +39,15 @@ def test_silicon_workflow(get_dftk_code, generate_structure, generate_kpoints_me
             {
                 "$function": "compute_stresses_cart"
             },
-        ]
+        ],
+        "refinement": {
+            "basis_kwargs": {
+                "Ecut": 40,
+            },
+            "$kwargs": {
+                "tol": 1e-12,
+            },
+        },
     })
 
     # Disable MPI for now. If we ever enable MPI,
@@ -66,3 +74,6 @@ def test_silicon_workflow(get_dftk_code, generate_structure, generate_kpoints_me
     assert_allclose(output_parameters["energies"]["total"], _REFERENCE_ENERGY, rtol=1e-2)
     assert_allclose(result.outputs.output_forces.get_array(), _REFERENCE_FORCES, rtol=1e-2)
     assert_allclose(result.outputs.output_stresses.get_array(), _REFERENCE_STRESSES, rtol=1e-2)
+
+    refinement_output = result.outputs.output_refinement.get_dict()
+    assert_allclose(refinement_output["E_refined_new"][0], -8.43893821)