From f50af97a0966e463a9bdb674130a317bc315dd9c Mon Sep 17 00:00:00 2001 From: ymzhang0 Date: Thu, 20 Aug 2026 11:19:55 +0200 Subject: [PATCH 1/9] feat(types,calculation): refine RestartType definition and calculation handling --- src/aiida_epw/calculations/epw.py | 298 ++++++++++++------------- src/aiida_epw/common/types.py | 45 +++- src/aiida_epw/parsers/epw.py | 4 +- tests/calculations/test_epw_calcjob.py | 164 +++++++++++--- 4 files changed, 322 insertions(+), 189 deletions(-) diff --git a/src/aiida_epw/calculations/epw.py b/src/aiida_epw/calculations/epw.py index 13cb48c6..e8ab14e7 100644 --- a/src/aiida_epw/calculations/epw.py +++ b/src/aiida_epw/calculations/epw.py @@ -77,7 +77,6 @@ class EpwCalculation(NamelistsCalculation): ("INPUTEPW", "nkf1"), ("INPUTEPW", "nkf2"), ("INPUTEPW", "nkf3"), - ("INPUTEPW", "eliashberg"), ("INPUTEPW", "liso"), ("INPUTEPW", "laniso"), ("INPUTEPW", "fbw"), @@ -99,8 +98,10 @@ class EpwCalculation(NamelistsCalculation): _OUTPUT_SUBFOLDER = "./out/" _FOLDER_SAVE = "save" _FOLDER_DYNAMICAL_MATRIX = "DYN_MAT" - _kfpoints_input_file = "kfpoints.kpt" - _qfpoints_input_file = "qfpoints.kpt" + _FILE_EPMATWP = _PREFIX + ".epmatwp" + _FOLDER_EPHMAT = _PREFIX + ".ephmat" + _KFPOINTS_INPUT_FILE = "kfpoints.kpt" + _QFPOINTS_INPUT_FILE = "qfpoints.kpt" _OUTPUT_XML_TENSOR_FILE_NAME = "tensors.xml" _OUTPUT_DOS_FILE = _PREFIX + ".dos" _OUTPUT_PHDOS_FILE = _PREFIX + ".phdos" @@ -109,8 +110,8 @@ class EpwCalculation(NamelistsCalculation): _OUTPUT_A2F_PROJ_FILE = _PREFIX + ".a2f_proj" _OUTPUT_LAMBDA_FS_FILE = _PREFIX + ".lambda_FS" _OUTPUT_LAMBDA_K_PAIRS_FILE = _PREFIX + ".lambda_k_pairs" - _output_elbands_file = "band.eig" - _output_phbands_file = "phband.freq" + _OUTPUT_ELBANDS_FILE = "band.eig" + _OUTPUT_PHBANDS_FILE = "phband.freq" _MAX_NSTEMP = 1000 @@ -138,8 +139,8 @@ def define(cls, spec): required=False, serializer=serialize_restart_type, help=( - "EPW run/restart mode: none, ephwrite, ephread, " - "ephwrite_restart, ephread_restart, or epwread." + "EPW run/restart mode: from_scratch, from_epb, from_epmatwp, " + "ephwrite, or from_eph." ), ) spec.input( @@ -459,28 +460,32 @@ def validate_restart_inputs(cls, parameters, inputs): ) return + if "restart_type" not in inputs: + return + + restart_node = inputs["restart_type"] restart_type = ( - inputs["restart_type"].get_member() if "restart_type" in inputs else None + restart_node.get_member() + if hasattr(restart_node, "get_member") + else restart_node ) parent_restart_types = ( + RestartType.FROM_EPB, + RestartType.FROM_EPMATWP, RestartType.EPHWRITE, - RestartType.EPHREAD, - RestartType.EPHWRITE_RESTART, - RestartType.EPHREAD_RESTART, - RestartType.EPWREAD, + RestartType.FROM_EPH, ) - uses_parent = restart_type in parent_restart_types - has_parent = "parent_folder_epw" in inputs - if uses_parent != has_parent: - if uses_parent: - raise exceptions.InputValidationError( - f"`parent_folder_epw` must be specified when " - f"restart_type is '{restart_type.value}'." - ) + if restart_type in parent_restart_types and "parent_folder_epw" not in inputs: + raise exceptions.InputValidationError( + f"`parent_folder_epw` must be specified when " + f"restart_type is '{restart_type.value}'." + ) + + if restart_type is RestartType.FROM_SCRATCH and "parent_folder_epw" in inputs: raise exceptions.InputValidationError( - "`restart_type` must be set to a mode that reads an EPW parent " - "when `parent_folder_epw` is provided." + "`parent_folder_epw` cannot be specified when " + "restart_type is 'from_scratch'." ) @staticmethod @@ -740,6 +745,20 @@ def cap_nstemp(self, inputepw_parameters): ) inputepw_parameters["nstemp"] = self._MAX_NSTEMP + def set_restart_parameters(self, inputepw_parameters): + """Set restart parameters in INPUTEPW based on the restart_type input.""" + if "restart_type" in self.inputs: + from aiida_epw.common.types import RESTART_TYPE_DEFAULTS + + restart_node = self.inputs.restart_type + restart_type = ( + restart_node.get_member() + if hasattr(restart_node, "get_member") + else restart_node + ) + for key, value in RESTART_TYPE_DEFAULTS[restart_type].items(): + inputepw_parameters.setdefault(key, value) + def prepare_input_parameters(self, folder, parameters): """Populate plugin-managed EPW parameters before writing the input file.""" inputepw_parameters = parameters["INPUTEPW"] @@ -765,7 +784,10 @@ def prepare_input_parameters(self, folder, parameters): inputepw_parameters["liso"] = not momentum_dependence if "full_bandwidth" in self.inputs: - inputepw_parameters["fbw"] = self.inputs.full_bandwidth.value + fbw = self.inputs.full_bandwidth.value + inputepw_parameters["fbw"] = fbw + if fbw: + inputepw_parameters["tc_linear"] = False if "real_axis" in self.inputs: real_axis = self.inputs.real_axis.value @@ -788,57 +810,7 @@ def prepare_input_parameters(self, folder, parameters): inputepw_parameters["lpade"] = False inputepw_parameters["lacon"] = False - if "restart_type" in self.inputs: - from aiida_epw.common.types import RestartType - - restart_type = self.inputs.restart_type.get_member() - if restart_type is RestartType.NONE: - inputepw_parameters.update( - { - "epwread": False, - "epwwrite": True, - "restart": False, - "ep_coupling": True, - "elph": True, - "epbwrite": True, - "epbread": False, - } - ) - elif restart_type in (RestartType.EPHWRITE, RestartType.EPHWRITE_RESTART): - inputepw_parameters.update( - { - "epwread": True, - "epwwrite": False, - "restart": restart_type is RestartType.EPHWRITE_RESTART, - "ep_coupling": True, - "elph": True, - "ephwrite": True, - } - ) - elif restart_type in (RestartType.EPHREAD, RestartType.EPHREAD_RESTART): - inputepw_parameters.update( - { - "epwread": True, - "restart": restart_type is RestartType.EPHREAD_RESTART, - "ep_coupling": False, - "elph": False, - "ephwrite": False, - } - ) - if inputepw_parameters.get("scattering", False): - inputepw_parameters["epmatkqread"] = True - elif restart_type is RestartType.EPWREAD: - inputepw_parameters.update( - { - "epwread": True, - "epwwrite": False, - "epbwrite": False, - "epbread": False, - "ep_coupling": True, - "elph": True, - } - ) - inputepw_parameters.setdefault("restart", False) + self.set_restart_parameters(inputepw_parameters) inputepw_parameters["outdir"] = self._OUTPUT_SUBFOLDER inputepw_parameters["dvscf_dir"] = self._FOLDER_SAVE @@ -872,7 +844,7 @@ def prepare_input_parameters(self, folder, parameters): self.inputs.qfpoints, ("nqf1", "nqf2", "nqf3"), "filqf", - self._qfpoints_input_file, + self._QFPOINTS_INPUT_FILE, "Cannot get the fine q-point grid", ) @@ -892,7 +864,7 @@ def prepare_input_parameters(self, folder, parameters): self.inputs.kfpoints, ("nkf1", "nkf2", "nkf3"), "filkf", - self._kfpoints_input_file, + self._KFPOINTS_INPUT_FILE, "Cannot get the fine k-point grid", ) @@ -903,7 +875,7 @@ def get_additional_retrieve_list(self, parameters): retrieve_list = [] if parameters["INPUTEPW"].get("band_plot"): - retrieve_list += [self._output_elbands_file, self._output_phbands_file] + retrieve_list += [self._OUTPUT_ELBANDS_FILE, self._OUTPUT_PHBANDS_FILE] if parameters["INPUTEPW"].get("eliashberg", False): retrieve_list.append(self._OUTPUT_A2F_FILE) @@ -994,39 +966,44 @@ def stage_ph_parent(self, folder, settings, remote_list): else: nqpt = get_parent_ph_qpoint_ibz_count(parent_folder_ph) - prefix = self._PREFIX - outdir = PhCalculation._OUTPUT_SUBFOLDER - fildvscf = PhCalculation._DVSCF_PREFIX - fildyn = PhCalculation._OUTPUT_DYNAMICAL_MATRIX_PREFIX ph_path = self.get_parent_folder_path(parent_folder_ph) remote_list.append( ( parent_folder_ph.computer.uuid, - Path(ph_path, outdir, "_ph0", f"{prefix}.phsave").as_posix(), + Path( + ph_path, + PhCalculation._OUTPUT_SUBFOLDER, + "_ph0", + f"{self._PREFIX}.phsave", + ).as_posix(), self._FOLDER_SAVE, ) ) for iqpt in range(1, nqpt + 1): + q_dir = "" if iqpt == 1 else f"{self._PREFIX}.q_{iqpt}" remote_list.append( ( parent_folder_ph.computer.uuid, Path( ph_path, - outdir, + PhCalculation._OUTPUT_SUBFOLDER, "_ph0", - "" if iqpt == 1 else f"{prefix}.q_{iqpt}", - f"{prefix}.{fildvscf}1", + q_dir, + f"{self._PREFIX}.{PhCalculation._DVSCF_PREFIX}1", ).as_posix(), - Path(self._FOLDER_SAVE, f"{prefix}.dvscf_q{iqpt}").as_posix(), + Path(self._FOLDER_SAVE, f"{self._PREFIX}.dvscf_q{iqpt}").as_posix(), ) ) remote_list.append( ( parent_folder_ph.computer.uuid, - Path(ph_path, f"{fildyn}{iqpt}").as_posix(), - Path(self._FOLDER_SAVE, f"{prefix}.dyn_q{iqpt}").as_posix(), + Path( + ph_path, + f"{PhCalculation._OUTPUT_DYNAMICAL_MATRIX_PREFIX}{iqpt}", + ).as_posix(), + Path(self._FOLDER_SAVE, f"{self._PREFIX}.dyn_q{iqpt}").as_posix(), ) ) @@ -1035,86 +1012,105 @@ def stage_epw_parent(self, folder, parameters, remote_list, remote_symlink_list) if "parent_folder_epw" not in self.inputs: return + from aiida_epw.common.types import RestartType + + restart_node = self.inputs.get("restart_type", None) + restart_type = ( + restart_node.get_member() + if hasattr(restart_node, "get_member") + else restart_node + ) + + if restart_type is None: + inputepw = parameters.get("INPUTEPW", {}) + if inputepw.get("epbread", False): + restart_type = RestartType.FROM_EPB + elif inputepw.get("epwread", False): + if not inputepw.get("elph", True) and not inputepw.get( + "ep_coupling", True + ): + restart_type = RestartType.FROM_EPH + elif inputepw.get("ephwrite", False): + restart_type = RestartType.EPHWRITE + else: + restart_type = RestartType.FROM_EPMATWP + else: + restart_type = RestartType.FROM_SCRATCH + + if restart_type is RestartType.FROM_SCRATCH: + return + folder.get_subfolder(self._OUTPUT_SUBFOLDER, create=True) parent_folder_epw = self.inputs.parent_folder_epw epw_path = self.get_parent_folder_path(parent_folder_epw) - file_list = [ - "selecq.fmt", - "crystal.fmt", - "epwdata.fmt", - "dmedata.fmt", - "vmedata.fmt", - "wigner.fmt", - "quadrupole.fmt", - "decay.H", - "decay.v", - "decay.P", - "decay.dynmat", - "decay.epmate", - "decay.epmatp", - f"{self._PREFIX}.kgmap", - f"{self._PREFIX}.kmap", - f"{self._PREFIX}.ukk", - f"{self._PREFIX}.mmn", - f"{self._PREFIX}.bvec", - self._FOLDER_SAVE, - ] - if parameters["INPUTEPW"].get("restart", False): - file_list.append("restart.fmt") - - if parameters["INPUTEPW"].get("epwread", False) and parameters["INPUTEPW"].get( - "elph", False - ): - remote_symlink_list.append( + if parameters.get("INPUTEPW", {}).get("restart", False): + remote_list.append( ( parent_folder_epw.computer.uuid, - Path( - epw_path, - f"{self._OUTPUT_SUBFOLDER}/{self._PREFIX}.epmatwp", - ).as_posix(), - Path(f"{self._OUTPUT_SUBFOLDER}/{self._PREFIX}.epmatwp").as_posix(), + Path(epw_path, "restart.fmt").as_posix(), + "restart.fmt", ) ) - if parameters["INPUTEPW"].get("eliashberg", False): - if parameters["INPUTEPW"].get("ephwrite", True): - if parameters["INPUTEPW"].get("restart", False): - remote_symlink_list.append( - ( - parent_folder_epw.computer.uuid, - Path( - epw_path, - f"{self._OUTPUT_SUBFOLDER}/{self._PREFIX}.ephmat", - ).as_posix(), - Path( - f"{self._OUTPUT_SUBFOLDER}/{self._PREFIX}.ephmat" - ).as_posix(), - ) + if restart_type is RestartType.FROM_EPB: + file_list = [ + (f"{self._PREFIX}.epb*", "."), + (self._FOLDER_SAVE, self._FOLDER_SAVE), + (f"{self._PREFIX}.ukk", f"{self._PREFIX}.ukk"), + ] + for src_name, dst_name in file_list: + remote_list.append( + ( + parent_folder_epw.computer.uuid, + Path(epw_path, src_name).as_posix(), + Path(dst_name).as_posix(), ) - else: - remote_symlink_list.append( + ) + return + + if restart_type in (RestartType.FROM_EPMATWP, RestartType.EPHWRITE): + epmatwp_rel = Path(self._OUTPUT_SUBFOLDER, self._FILE_EPMATWP).as_posix() + file_list = [ + f"{self._PREFIX}.bvec", + f"{self._PREFIX}.kgmap", + f"{self._PREFIX}.kmap", + f"{self._PREFIX}.mmn", + f"{self._PREFIX}.ukk", + "crystal.fmt", + "dmedata.fmt", + "epwdata.fmt", + epmatwp_rel, + self._FOLDER_SAVE, + "selecq.fmt", + "vmedata.fmt", + ] + for filename in file_list: + remote_list.append( ( parent_folder_epw.computer.uuid, - Path( - epw_path, - f"{self._OUTPUT_SUBFOLDER}/{self._PREFIX}.ephmat", - ).as_posix(), - Path( - f"{self._OUTPUT_SUBFOLDER}/{self._PREFIX}.ephmat" - ).as_posix(), + Path(epw_path, filename).as_posix(), + Path(filename).as_posix(), ) ) + return - for filename in file_list: - remote_list.append( - ( - parent_folder_epw.computer.uuid, - Path(epw_path, filename).as_posix(), - Path(filename).as_posix(), + if restart_type is RestartType.FROM_EPH: + ephmat_rel = Path(self._OUTPUT_SUBFOLDER, self._FOLDER_EPHMAT).as_posix() + file_list = [ + ephmat_rel, + self._OUTPUT_A2F_FILE, + ] + for filename in file_list: + remote_list.append( + ( + parent_folder_epw.computer.uuid, + Path(epw_path, filename).as_posix(), + Path(filename).as_posix(), + ) ) - ) + return def stage_quadrupole(self, local_copy_list, remote_list): """Stage quadrupole file/directory if provided as inputs.""" diff --git a/src/aiida_epw/common/types.py b/src/aiida_epw/common/types.py index ab12b65e..199dbec4 100644 --- a/src/aiida_epw/common/types.py +++ b/src/aiida_epw/common/types.py @@ -4,9 +4,44 @@ class RestartType(enum.Enum): """Enumeration of EPW run/restart modes.""" - NONE = "none" + FROM_SCRATCH = "from_scratch" + FROM_EPB = "from_epb" + FROM_EPMATWP = "from_epmatwp" EPHWRITE = "ephwrite" - EPHREAD = "ephread" - EPHWRITE_RESTART = "ephwrite_restart" - EPHREAD_RESTART = "ephread_restart" - EPWREAD = "epwread" + FROM_EPH = "from_eph" + + +RESTART_TYPE_DEFAULTS = { + RestartType.FROM_SCRATCH: { + "epwread": False, + "epwwrite": True, + "epbwrite": True, + "epbread": False, + }, + RestartType.FROM_EPB: { + "epbread": True, + "epbwrite": False, + "epwread": False, + "epwwrite": True, + }, + RestartType.FROM_EPMATWP: { + "epwread": True, + "epwwrite": False, + "epbwrite": False, + "epbread": False, + }, + RestartType.EPHWRITE: { + "epwread": True, + "ep_coupling": True, + "elph": True, + "ephwrite": True, + "restart": True, + }, + RestartType.FROM_EPH: { + "epwread": True, + "ep_coupling": False, + "elph": False, + "ephwrite": False, + "restart": False, + }, +} diff --git a/src/aiida_epw/parsers/epw.py b/src/aiida_epw/parsers/epw.py index dd48419a..072c4f8c 100644 --- a/src/aiida_epw/parsers/epw.py +++ b/src/aiida_epw/parsers/epw.py @@ -89,7 +89,7 @@ def parse(self, **kwargs): parsed_data.update(parsed_epw) elbands_contents = self.get_retrieved_content( - EpwCalculation._output_elbands_file + EpwCalculation._OUTPUT_ELBANDS_FILE ) if elbands_contents is not None: self.out( @@ -100,7 +100,7 @@ def parse(self, **kwargs): ) phbands_contents = self.get_retrieved_content( - EpwCalculation._output_phbands_file + EpwCalculation._OUTPUT_PHBANDS_FILE ) if phbands_contents is not None: self.out( diff --git a/tests/calculations/test_epw_calcjob.py b/tests/calculations/test_epw_calcjob.py index 42b54e8d..cb93c3f5 100644 --- a/tests/calculations/test_epw_calcjob.py +++ b/tests/calculations/test_epw_calcjob.py @@ -445,49 +445,114 @@ def test_epw_stages_chk_parent_into_requested_transport_list( assert not expected.intersection(set(calc_info.remote_copy_list)) -def test_epw_stages_epw_restart_files_without_copying_epmatwp( +def test_epw_stages_from_epb( fixture_sandbox, fixture_localhost, generate_calc_job, generate_inputs_epw, generate_remote_data, ): - """Test that EPW restart staging links the large `epmatwp` file and copies metadata files.""" + """Test that FROM_EPB stages prefix.epb*, save folder, and prefix.ukk.""" parent_folder = generate_remote_data(fixture_localhost, "/remote/epw") inputs = generate_inputs_epw( - restart_type=RestartType.EPWREAD, - parameters={"INPUTEPW": {}}, + restart_type=RestartType.FROM_EPB, parent_folder_epw=parent_folder, ) calc_info = generate_calc_job(fixture_sandbox, "epw.epw", inputs) + copied = [(entry[1], entry[2]) for entry in calc_info.remote_copy_list] assert ( - parent_folder.computer.uuid, - Path( - parent_folder.get_remote_path(), - f"{EpwCalculation._OUTPUT_SUBFOLDER}/{EpwCalculation._PREFIX}.epmatwp", - ).as_posix(), - Path( - f"{EpwCalculation._OUTPUT_SUBFOLDER}/{EpwCalculation._PREFIX}.epmatwp" - ).as_posix(), - ) in calc_info.remote_symlink_list + Path(parent_folder.get_remote_path(), "aiida.epb*").as_posix(), + ".", + ) in copied + assert ( + Path(parent_folder.get_remote_path(), "save").as_posix(), + "save", + ) in copied + assert ( + Path(parent_folder.get_remote_path(), "aiida.ukk").as_posix(), + "aiida.ukk", + ) in copied + + +def test_epw_stages_from_epmatwp( + fixture_sandbox, + fixture_localhost, + generate_calc_job, + generate_inputs_epw, + generate_remote_data, +): + """Test that FROM_EPMATWP / EPHWRITE stages all required EPW metadata files, out, and save.""" + parent_folder = generate_remote_data(fixture_localhost, "/remote/epw") + inputs = generate_inputs_epw( + restart_type=RestartType.FROM_EPMATWP, + parameters={"INPUTEPW": {"eliashberg": True}}, + parent_folder_epw=parent_folder, + ) + + calc_info = generate_calc_job(fixture_sandbox, "epw.epw", inputs) + copied_targets = {entry[2] for entry in calc_info.remote_copy_list} expected_copied = { - "crystal.fmt", - "epwdata.fmt", - "vmedata.fmt", - "dmedata.fmt", + "aiida.bvec", "aiida.kgmap", "aiida.kmap", - "aiida.ukk", "aiida.mmn", - "aiida.bvec", + "aiida.ukk", + "crystal.fmt", + "dmedata.fmt", + "epwdata.fmt", + "out/aiida.epmatwp", + "save", + "selecq.fmt", + "vmedata.fmt", } - copied_targets = {entry[2] for entry in calc_info.remote_copy_list} assert expected_copied.issubset(copied_targets) - assert EpwCalculation._OUTPUT_SUBFOLDER not in copied_targets - assert Path(EpwCalculation._OUTPUT_SUBFOLDER).as_posix() not in copied_targets + assert "out" not in copied_targets + + +def test_epw_stages_from_eph( + fixture_sandbox, + fixture_localhost, + generate_calc_job, + generate_inputs_epw, + generate_remote_data, +): + """Test that FROM_EPH stages out/prefix.ephmat and prefix.a2f.""" + parent_folder = generate_remote_data(fixture_localhost, "/remote/epw") + inputs = generate_inputs_epw( + restart_type=RestartType.FROM_EPH, + parent_folder_epw=parent_folder, + ) + + calc_info = generate_calc_job(fixture_sandbox, "epw.epw", inputs) + copied = [(entry[1], entry[2]) for entry in calc_info.remote_copy_list] + + assert ( + Path(parent_folder.get_remote_path(), "out/aiida.ephmat").as_posix(), + "out/aiida.ephmat", + ) in copied + assert ( + Path(parent_folder.get_remote_path(), "aiida.a2f").as_posix(), + "aiida.a2f", + ) in copied + assert len(calc_info.remote_copy_list) == 2 + + # With restart = True + inputs_restart = generate_inputs_epw( + restart_type=RestartType.FROM_EPH, + parameters={"INPUTEPW": {"restart": True}}, + parent_folder_epw=parent_folder, + ) + calc_info_restart = generate_calc_job(fixture_sandbox, "epw.epw", inputs_restart) + copied_restart = [ + (entry[1], entry[2]) for entry in calc_info_restart.remote_copy_list + ] + assert ( + Path(parent_folder.get_remote_path(), "restart.fmt").as_posix(), + "restart.fmt", + ) in copied_restart def test_epw_stages_ph_stash_folder_by_target_basepath( @@ -570,24 +635,35 @@ def test_epw_eliashberg_parameters_continuation_none( ("restart_type", "parameters", "expected_entries"), [ ( - RestartType.NONE, + RestartType.FROM_SCRATCH, {}, ("epwread = .false.", "epwwrite = .true.", "epbwrite = .true."), ), ( - RestartType.EPHWRITE_RESTART, + RestartType.FROM_EPB, {}, - ("epwread = .true.", "restart = .true.", "ephwrite = .true."), + ("epbread = .true.", "epbwrite = .false.", "epwwrite = .true."), ), ( - RestartType.EPHREAD, - {"scattering": True}, - ("epwread = .true.", "ep_coupling = .false.", "epmatkqread = .true."), + RestartType.FROM_EPMATWP, + {}, + ("epwread = .true.", "epwwrite = .false.", "epbwrite = .false."), ), ( - RestartType.EPWREAD, + RestartType.EPHWRITE, {}, - ("epwread = .true.", "epwwrite = .false.", "epbwrite = .false."), + ( + "epwread = .true.", + "ep_coupling = .true.", + "elph = .true.", + "ephwrite = .true.", + "restart = .true.", + ), + ), + ( + RestartType.FROM_EPH, + {}, + ("epwread = .true.", "ep_coupling = .false.", "elph = .false."), ), ], ) @@ -607,7 +683,7 @@ def test_epw_restart_type_parameter( parameters={"INPUTEPW": parameters}, **( {} - if restart_type is RestartType.NONE + if restart_type is RestartType.FROM_SCRATCH else { "parent_folder_epw": generate_remote_data( fixture_localhost, "/remote/epw" @@ -621,3 +697,29 @@ def test_epw_restart_type_parameter( for entry in expected_entries: assert entry in input_contents + + +def test_epw_parent_folder_without_restart_type( + fixture_sandbox, + fixture_localhost, + generate_calc_job, + generate_inputs_epw, + generate_remote_data, +): + """Test that parent_folder_epw can be provided without specifying restart_type.""" + inputs = generate_inputs_epw( + parameters={ + "INPUTEPW": { + "epwread": True, + "elph": False, + "ep_coupling": False, + } + }, + parent_folder_epw=generate_remote_data(fixture_localhost, "/remote/epw"), + ) + + generate_calc_job(fixture_sandbox, "epw.epw", inputs) + input_contents = Path(fixture_sandbox.abspath, "aiida.in").read_text() + + assert "epwread = .true." in input_contents + assert "elph = .false." in input_contents From 08fb06d74f1fa0017ca9f62c1bbe142a57a0d1c6 Mon Sep 17 00:00:00 2001 From: ymzhang0 Date: Thu, 20 Aug 2026 11:20:02 +0200 Subject: [PATCH 2/9] refactor(workflows): update workchains and protocols for refined RestartType --- src/aiida_epw/workflows/base.py | 11 ++----- src/aiida_epw/workflows/mobility.py | 2 -- src/aiida_epw/workflows/prep.py | 4 +-- src/aiida_epw/workflows/protocols/prep.yaml | 4 +-- .../workflows/protocols/prep_sc.yaml | 4 +-- .../workflows/protocols/supercon.yaml | 3 -- src/aiida_epw/workflows/supercon.py | 21 ++++++++----- tests/workflows/test_base.py | 30 ++++--------------- tests/workflows/test_mobility.py | 2 +- tests/workflows/test_supercon.py | 8 ++--- 10 files changed, 33 insertions(+), 56 deletions(-) diff --git a/src/aiida_epw/workflows/base.py b/src/aiida_epw/workflows/base.py index ebd6376c..e18f3c74 100644 --- a/src/aiida_epw/workflows/base.py +++ b/src/aiida_epw/workflows/base.py @@ -518,9 +518,7 @@ def handle_cannot_bracket_ef(self, calculation): restart_type_node.get_member() if restart_type_node is not None else None ) next_restart_type = { - "NONE": "EPWREAD", - "EPHWRITE": "EPHREAD", - "EPHWRITE_RESTART": "EPHREAD", + "FROM_SCRATCH": "FROM_EPB", }.get(getattr(restart_type, "name", None)) if next_restart_type is not None: self.ctx.inputs.restart_type = restart_type.__class__[next_restart_type] @@ -645,12 +643,9 @@ def handle_pade_approximants(self, calculation): True, self.exit_codes.ERROR_KNOWN_UNRECOVERABLE_FAILURE ) - try: - from aiida_epw.common.types import RestartType + from aiida_epw.common.types import RestartType - self.ctx.inputs.restart_type = RestartType.EPHREAD - except ImportError: - parameters.setdefault("INPUTEPW", {})["epwread"] = True + self.ctx.inputs.restart_type = RestartType.FROM_EPH self.ctx.inputs.parameters = orm.Dict(parameters) self.ctx.inputs.parent_folder_epw = calculation.outputs.remote_folder diff --git a/src/aiida_epw/workflows/mobility.py b/src/aiida_epw/workflows/mobility.py index d598debe..59876c7d 100644 --- a/src/aiida_epw/workflows/mobility.py +++ b/src/aiida_epw/workflows/mobility.py @@ -5,7 +5,6 @@ from aiida_quantumespresso.workflows.protocols.utils import ProtocolMixin -from aiida_epw.calculations.epw import serialize_restart_type from aiida_epw.tools.band_analysis import extract_band_edges_from_epw_prep from aiida_epw.workflows.base import EpwBaseWorkChain @@ -448,7 +447,6 @@ def run_conv(self): "code": self.inputs.epw_mobility.code, "structure": self.inputs.structure, "parent_folder_epw": self.inputs.parent_folder_epw, - "restart_type": serialize_restart_type("ephread_restart"), "qfpoints_distance": self.ctx.interpolation_list.pop(0), "kfpoints_factor": self.inputs.kfpoints_factor, "parameters": orm.Dict(parameters), diff --git a/src/aiida_epw/workflows/prep.py b/src/aiida_epw/workflows/prep.py index bf54f109..a093f33d 100644 --- a/src/aiida_epw/workflows/prep.py +++ b/src/aiida_epw/workflows/prep.py @@ -1193,7 +1193,7 @@ def run_epw(self): from aiida_epw.common.types import RestartType - inputs.restart_type = RestartType.NONE + inputs.restart_type = RestartType.FROM_SCRATCH workchain_node = self.submit(EpwBaseWorkChain, **inputs) self.report( @@ -1274,7 +1274,7 @@ def run_epw_bands(self): from aiida_epw.common.types import RestartType - inputs.restart_type = RestartType.EPWREAD + inputs.restart_type = RestartType.FROM_EPMATWP workchain_node = self.submit(EpwBaseWorkChain, **inputs) self.report( diff --git a/src/aiida_epw/workflows/protocols/prep.yaml b/src/aiida_epw/workflows/protocols/prep.yaml index ed4e321f..6632bab2 100644 --- a/src/aiida_epw/workflows/protocols/prep.yaml +++ b/src/aiida_epw/workflows/protocols/prep.yaml @@ -61,7 +61,7 @@ default_inputs: temps: 300 vme: 'dipole' wannierize: False - restart_type: wannierize + restart_type: from_scratch epw_bands: options: resources: @@ -70,7 +70,7 @@ default_inputs: parameters: INPUTEPW: band_plot: True - restart_type: epwread + restart_type: from_epmatwp default_protocol: moderate protocols: moderate: diff --git a/src/aiida_epw/workflows/protocols/prep_sc.yaml b/src/aiida_epw/workflows/protocols/prep_sc.yaml index 0eb5c444..0b9bd065 100644 --- a/src/aiida_epw/workflows/protocols/prep_sc.yaml +++ b/src/aiida_epw/workflows/protocols/prep_sc.yaml @@ -57,7 +57,7 @@ default_inputs: fsthick: 100 phonselfen: False temps: 300 - restart_type: wannierize + restart_type: from_scratch epw_bands: options: resources: @@ -66,7 +66,7 @@ default_inputs: parameters: INPUTEPW: band_plot: True - restart_type: epwread + restart_type: from_epmatwp default_protocol: moderate protocols: moderate: diff --git a/src/aiida_epw/workflows/protocols/supercon.yaml b/src/aiida_epw/workflows/protocols/supercon.yaml index f230a8a4..1f55180c 100644 --- a/src/aiida_epw/workflows/protocols/supercon.yaml +++ b/src/aiida_epw/workflows/protocols/supercon.yaml @@ -18,7 +18,6 @@ default_inputs: num_machines: 1 max_wallclock_seconds: 43200 # Twelve hours withmpi: True - restart_type: ephwrite momentum_dependence: False analytical_continuation: pade parameters: @@ -46,7 +45,6 @@ default_inputs: num_machines: 1 max_wallclock_seconds: 43200 # Twelve hours withmpi: True - restart_type: ephread parameters: INPUTEPW: eliashberg: True @@ -72,7 +70,6 @@ default_inputs: num_machines: 1 max_wallclock_seconds: 43200 # Twelve hours withmpi: True - restart_type: ephread parameters: INPUTEPW: eliashberg: True diff --git a/src/aiida_epw/workflows/supercon.py b/src/aiida_epw/workflows/supercon.py index 6cd7edde..57d48c6f 100644 --- a/src/aiida_epw/workflows/supercon.py +++ b/src/aiida_epw/workflows/supercon.py @@ -121,7 +121,6 @@ def define(cls, spec): "parent_folder_chk", "qfpoints", "kfpoints", - "restart_type", ), namespace_options={ "help": ( @@ -140,7 +139,6 @@ def define(cls, spec): "parent_folder_chk", "qfpoints_distance", "kfpoints_factor", - "restart_type", ), namespace_options={ "help": ( @@ -159,7 +157,6 @@ def define(cls, spec): "parent_folder_chk", "qfpoints_distance", "kfpoints_factor", - "restart_type", ), namespace_options={ "help": ( @@ -286,7 +283,11 @@ def get_builder_from_protocol( real_axis = epw_inputs.pop("real_axis", False) analytical_continuation = epw_inputs.pop("analytical_continuation", None) epw_inputs.pop("calculation_type", None) - epw_inputs.pop("restart_type", None) + restart_type = epw_inputs.pop("restart_type", None) + if restart_type is None: + restart_type = ( + "ephwrite" if epw_namespace == "epw_interp" else "from_eph" + ) # Check which input ports are supported by EpwBaseWorkChain dynamically for cross-branch compatibility base_inputs = EpwBaseWorkChain.spec().inputs @@ -308,6 +309,9 @@ def get_builder_from_protocol( **kwargs, ) + if "restart_type" in base_inputs: + epw_builder.restart_type = serialize_restart_type(restart_type) + if epw_namespace == "epw_interp" and scon_epw_code is not None: epw_builder.code = scon_epw_code else: @@ -402,7 +406,8 @@ def run_conv(self): inputs.kfpoints_factor = self.inputs.epw_interp.kfpoints_factor inputs.qfpoints_distance = self.ctx.interpolation_list.pop() - inputs.restart_type = serialize_restart_type("ephwrite") + if "restart_type" not in inputs: + inputs.restart_type = serialize_restart_type("ephwrite") if self.ctx.degaussq: parameters = inputs.parameters.get_dict() @@ -469,7 +474,8 @@ def run_final_epw_iso(self): inputs.kfpoints = parent_folder_epw.creator.inputs.kfpoints inputs.qfpoints = parent_folder_epw.creator.inputs.qfpoints - inputs.restart_type = serialize_restart_type("ephread") + if "restart_type" not in inputs: + inputs.restart_type = serialize_restart_type("from_eph") if self.ctx.degaussq: parameters = inputs.parameters.get_dict() @@ -507,7 +513,8 @@ def run_final_epw_aniso(self): inputs.kfpoints = parent_folder_epw.creator.inputs.kfpoints inputs.qfpoints = parent_folder_epw.creator.inputs.qfpoints - inputs.restart_type = serialize_restart_type("ephread") + if "restart_type" not in inputs: + inputs.restart_type = serialize_restart_type("from_eph") inputs.metadata.call_link_label = "epw_final_aniso" workchain_node = self.submit(EpwBaseWorkChain, **inputs) diff --git a/tests/workflows/test_base.py b/tests/workflows/test_base.py index 2b66485c..fdeb39c5 100644 --- a/tests/workflows/test_base.py +++ b/tests/workflows/test_base.py @@ -2,17 +2,11 @@ from unittest.mock import MagicMock -import enum import pytest from aiida import orm -try: - from aiida_epw.common.types import RestartType - - HAS_RESTART_TYPE = True -except ImportError: - HAS_RESTART_TYPE = False +from aiida_epw.common.types import RestartType from aiida_epw.workflows.base import EpwBaseWorkChain @@ -115,10 +109,7 @@ def report_error_handled(self, calculation, action): assert input_epw["npade"] == 25 # Restart settings should be set - if HAS_RESTART_TYPE: - assert workchain.ctx.inputs.restart_type == RestartType.EPHREAD - else: - assert input_epw.get("epwread") is True + assert workchain.ctx.inputs.restart_type == RestartType.FROM_EPH assert workchain.ctx.inputs.parent_folder_epw == calc.outputs.remote_folder @@ -522,27 +513,16 @@ def report_error_handled(self, calculation, action): @pytest.mark.parametrize( ("restart_type", "expected_restart_type"), - ( - ("NONE", "EPWREAD"), - ("EPHWRITE", "EPHREAD"), - ("EPHWRITE_RESTART", "EPHREAD"), - ), + (("FROM_SCRATCH", "FROM_EPB"),), ) def test_handle_cannot_bracket_ef_switches_writer_restart_to_reader( restart_type, expected_restart_type ): """Test that Fermi-level recovery reuses files from the failed calculation.""" - class MockRestartType(enum.Enum): - NONE = "none" - EPHWRITE = "ephwrite" - EPHREAD = "ephread" - EPHWRITE_RESTART = "ephwrite_restart" - EPWREAD = "epwread" - class RestartTypeNode: def get_member(self): - return MockRestartType[restart_type] + return RestartType[restart_type] class MockWorkChain: exit_codes = EpwBaseWorkChain.exit_codes @@ -573,7 +553,7 @@ def report_error_handled(self, calculation, action): report = workchain.handle_cannot_bracket_ef.__wrapped__(calculation) assert report.exit_code.status == 0 - assert workchain.ctx.inputs.restart_type == MockRestartType[expected_restart_type] + assert workchain.ctx.inputs.restart_type == RestartType[expected_restart_type] assert workchain.ctx.inputs.parent_folder_epw is calculation.outputs.remote_folder assert expected_restart_type in workchain.report_messages[-1] diff --git a/tests/workflows/test_mobility.py b/tests/workflows/test_mobility.py index 57a171f0..a4091cce 100644 --- a/tests/workflows/test_mobility.py +++ b/tests/workflows/test_mobility.py @@ -2,7 +2,7 @@ def test_mobility_protocol_omits_plugin_managed_parameters(): - """Mobility protocols leave EPW run/restart flags to ``RestartType``.""" + """Mobility protocols omit plugin-managed EPW keywords.""" from aiida_epw.calculations.epw import EpwCalculation from aiida_epw.workflows.mobility import MobilityWorkChain diff --git a/tests/workflows/test_supercon.py b/tests/workflows/test_supercon.py index c9b1f3df..6ce581fb 100644 --- a/tests/workflows/test_supercon.py +++ b/tests/workflows/test_supercon.py @@ -223,8 +223,8 @@ def test_epw_base_restart_types(fixture_code, generate_structure): ) # We should be able to set and access restart_type on the builder - builder.restart_type = RestartType.EPHWRITE - assert builder.restart_type == RestartType.EPHWRITE + builder.restart_type = RestartType.FROM_EPMATWP + assert builder.restart_type == RestartType.FROM_EPMATWP def test_supercon_get_builder_from_protocol_default( @@ -286,7 +286,7 @@ def test_supercon_get_builder_from_protocol_default( if "restart_type" in builder.epw_final_iso: from aiida_epw.common import RestartType - assert builder.epw_final_iso.restart_type == RestartType.EPHREAD + assert builder.epw_final_iso.restart_type == RestartType.FROM_EPH assert ( builder.epw_final_iso.parameters.get_dict()["INPUTEPW"].get("tc_linear", False) is False @@ -307,4 +307,4 @@ def test_supercon_get_builder_from_protocol_default( if "restart_type" in builder.epw_final_aniso: from aiida_epw.common import RestartType - assert builder.epw_final_aniso.restart_type == RestartType.EPHREAD + assert builder.epw_final_aniso.restart_type == RestartType.FROM_EPH From ae5c7039932ff02361504b571e5d9b5860d23b0b Mon Sep 17 00:00:00 2001 From: ymzhang0 Date: Tue, 18 Aug 2026 11:23:35 +0200 Subject: [PATCH 3/9] feat: add walltime error handling --- src/aiida_epw/calculations/epw.py | 63 +++- .../tools/error_handling/supercon.py | 96 +++++- src/aiida_epw/tools/workchain.py | 102 +++++++ src/aiida_epw/workflows/base.py | 47 ++- .../workflows/protocols/mobility.yaml | 8 - tests/calculations/test_epw_calcjob.py | 173 ++++++----- tests/workflows/test_base.py | 273 ++++++++++++++++++ tests/workflows/test_supercon.py | 84 ------ 8 files changed, 680 insertions(+), 166 deletions(-) diff --git a/src/aiida_epw/calculations/epw.py b/src/aiida_epw/calculations/epw.py index 13cb48c6..0a9dada0 100644 --- a/src/aiida_epw/calculations/epw.py +++ b/src/aiida_epw/calculations/epw.py @@ -77,6 +77,16 @@ class EpwCalculation(NamelistsCalculation): ("INPUTEPW", "nkf1"), ("INPUTEPW", "nkf2"), ("INPUTEPW", "nkf3"), + ("INPUTEPW", "wannierize"), + ("INPUTEPW", "epwread"), + ("INPUTEPW", "epwwrite"), + ("INPUTEPW", "epbread"), + ("INPUTEPW", "epbwrite"), + ("INPUTEPW", "restart"), + ("INPUTEPW", "ep_coupling"), + ("INPUTEPW", "elph"), + ("INPUTEPW", "ephwrite"), + ("INPUTEPW", "epmatkqread"), ("INPUTEPW", "eliashberg"), ("INPUTEPW", "liso"), ("INPUTEPW", "laniso"), @@ -463,10 +473,10 @@ def validate_restart_inputs(cls, parameters, inputs): inputs["restart_type"].get_member() if "restart_type" in inputs else None ) parent_restart_types = ( + RestartType.FROM_EPB, + RestartType.FROM_EPMATWP, RestartType.EPHWRITE, - RestartType.EPHREAD, - RestartType.EPHWRITE_RESTART, - RestartType.EPHREAD_RESTART, + RestartType.FROM_EPH, RestartType.EPWREAD, ) uses_parent = restart_type in parent_restart_types @@ -746,6 +756,53 @@ def prepare_input_parameters(self, folder, parameters): self.cap_nstemp(inputepw_parameters) + if "restart_type" in self.inputs: + restart_val = self.inputs.restart_type.get_member() + from aiida_epw.common import RestartType + + if restart_val is RestartType.NONE: + inputepw_parameters["epwread"] = False + inputepw_parameters["epwwrite"] = True + inputepw_parameters["restart"] = False + inputepw_parameters["ep_coupling"] = True + inputepw_parameters["elph"] = True + inputepw_parameters["epbwrite"] = True + inputepw_parameters["epbread"] = False + elif restart_val is RestartType.EPHWRITE: + inputepw_parameters["epwread"] = True + inputepw_parameters["epwwrite"] = False + inputepw_parameters["restart"] = False + inputepw_parameters["ep_coupling"] = True + inputepw_parameters["elph"] = True + inputepw_parameters["ephwrite"] = True + elif restart_val is RestartType.EPHWRITE_RESTART: + inputepw_parameters["epwread"] = True + inputepw_parameters["epwwrite"] = False + inputepw_parameters["restart"] = True + inputepw_parameters["ep_coupling"] = True + inputepw_parameters["elph"] = True + inputepw_parameters["ephwrite"] = True + elif restart_val in ( + RestartType.EPHREAD, + RestartType.EPHREAD_RESTART, + ): + inputepw_parameters["epwread"] = True + inputepw_parameters["restart"] = ( + restart_val is RestartType.EPHREAD_RESTART + ) + inputepw_parameters["ep_coupling"] = False + inputepw_parameters["elph"] = False + inputepw_parameters["ephwrite"] = False + if inputepw_parameters.get("scattering", False): + inputepw_parameters["epmatkqread"] = True + elif restart_val is RestartType.EPWREAD: + inputepw_parameters["epwread"] = True + inputepw_parameters["epwwrite"] = False + inputepw_parameters["epbwrite"] = False + inputepw_parameters["epbread"] = False + inputepw_parameters.setdefault("restart", False) + inputepw_parameters["ep_coupling"] = True + inputepw_parameters["elph"] = True # Override Eliashberg settings in parameters if inputs are specified eliashberg_any = any( f in self.inputs diff --git a/src/aiida_epw/tools/error_handling/supercon.py b/src/aiida_epw/tools/error_handling/supercon.py index 51160c68..caaef436 100644 --- a/src/aiida_epw/tools/error_handling/supercon.py +++ b/src/aiida_epw/tools/error_handling/supercon.py @@ -1,4 +1,4 @@ -"""Pade error-recovery helpers for EPW workflows.""" +"""Pade and walltime error-recovery helpers for EPW workflows.""" from copy import deepcopy @@ -184,3 +184,97 @@ def get_temperature_out_of_range_action(): "Temperature reached phase transition (delta converged to zero). " "Finishing workchain successfully." ) + + +def prepare_walltime_recovery( + parameters, output_parameters, restart_type, is_eliashberg +): + """Return ``(parameters, restart_type, action)`` for a recoverable timeout.""" + if not is_eliashberg: + return ( + None, + None, + ( + "Walltime reached but this calculation/restart type is not supported " + "for auto-recovery. Aborting." + ), + ) + + if restart_type in ("EPHWRITE", "ephwrite"): + return ( + parameters, + "EPHWRITE", + ( + "Walltime reached during Eliashberg ephwrite calculation. " + "Restarting from the last checkpoint." + ), + ) + + if restart_type in ("FROM_EPH", "from_eph", "EPHREAD", "ephread"): + outputs = output_parameters or {} + eliashberg_data = ( + outputs.get("isotropic_eliashberg") + or outputs.get("anisotropic_eliashberg") + or {} + ) + succeeded_temps = _get_succeeded_temperatures(eliashberg_data) + input_epw = parameters.get("INPUTEPW", {}) + all_temps, is_linear_range = _get_input_temperatures(input_epw, eliashberg_data) + remaining_temps = [ + temp + for temp in all_temps + if not any(abs(temp - succeeded) < 1.0e-4 for succeeded in succeeded_temps) + ] + + if succeeded_temps and remaining_temps: + parameters = deepcopy(parameters) + input_epw_new = parameters.setdefault("INPUTEPW", {}) + new_temps = ( + [remaining_temps[0], remaining_temps[-1]] + if is_linear_range and len(remaining_temps) >= 2 + else remaining_temps + ) + input_epw_new["temps"] = ( + " ".join(str(temp) for temp in new_temps) + if isinstance(input_epw.get("temps"), str) + else new_temps + ) + input_epw_new["nstemp"] = len(remaining_temps) + input_epw_new.pop("tempsmin", None) + input_epw_new.pop("tempsmax", None) + return ( + parameters, + "FROM_EPH", + ( + "Walltime reached during Eliashberg from_eph calculation. " + f"Removed successfully calculated temperatures: {succeeded_temps}. " + "Restarting." + ), + ) + return ( + None, + None, + ( + "Walltime reached during Eliashberg from_eph calculation but no " + "temperatures finished. Aborting." + ), + ) + + if restart_type in ("FROM_SCRATCH", "from_scratch", "NONE", "none"): + return ( + None, + None, + ( + "Walltime reached during a fresh Eliashberg run. Automatic recovery " + "requires an ephwrite checkpoint. Aborting." + ), + ) + + return ( + None, + None, + ( + "Walltime reached but this calculation/restart type is not supported " + "for auto-recovery. Aborting." + ), + ) diff --git a/src/aiida_epw/tools/workchain.py b/src/aiida_epw/tools/workchain.py index 5dfc9d8a..4ec590b1 100644 --- a/src/aiida_epw/tools/workchain.py +++ b/src/aiida_epw/tools/workchain.py @@ -342,3 +342,105 @@ def get_default_target_basepath(computer): else: raise ValueError(f"Unsupported transport type: {computer.transport_type}") return target_basepath + + +def pop_succeeded_temperatures(parameters, output_parameters): + """Analyze Eliashberg output parameters and pop successfully completed temperatures from input parameters. + + :param parameters: Dict containing input namelist parameters + :param output_parameters: Dict containing output parameters of the calculation + :return: A tuple of (updated_parameters_dict, succeeded_temps, remaining_temps, eliashberg_data) + """ + eliashberg_data = ( + output_parameters.get("isotropic_eliashberg") + or output_parameters.get("anisotropic_eliashberg") + or {} + ) + + succeeded_temps = [] + for temp_str, data in eliashberg_data.items(): + temp = float(temp_str) + has_failed = False + + iterations = data.get("iterations", {}) + if iterations: + if not iterations.get("ethr") or any( + v is None for v in iterations.get("ethr", []) + ): + has_failed = True + elif any(v is None for v in iterations.get("znormi", [])) or any( + v is None for v in iterations.get("deltai", []) + ): + has_failed = True + + pade = data.get("pade", {}) + if pade: + if any( + pade.get(k) is None for k in ("delta", "znorm", "shift") if k in pade + ): + has_failed = True + + if not iterations and not pade: + has_failed = True + + if not has_failed: + succeeded_temps.append(temp) + + input_epw = parameters.get("INPUTEPW", {}) + all_temps = [] + is_linear_range = False + if "temps" in input_epw: + temps_val = input_epw["temps"] + if isinstance(temps_val, str): + original_temps = [float(t) for t in temps_val.replace(",", " ").split()] + elif isinstance(temps_val, (int, float)): + original_temps = [float(temps_val)] + else: + original_temps = [float(t) for t in temps_val] + + nstemp = input_epw.get("nstemp", len(original_temps)) + if len(original_temps) == 2 and nstemp >= 2: + is_linear_range = True + t_min, t_max = original_temps[0], original_temps[1] + all_temps = [ + t_min + i * (t_max - t_min) / (nstemp - 1) for i in range(nstemp) + ] + else: + all_temps = original_temps + elif "tempsmin" in input_epw and "tempsmax" in input_epw and "nstemp" in input_epw: + is_linear_range = True + t_min = float(input_epw["tempsmin"]) + t_max = float(input_epw["tempsmax"]) + n_temp = int(input_epw["nstemp"]) + if n_temp > 1: + all_temps = [ + t_min + i * (t_max - t_min) / (n_temp - 1) for i in range(n_temp) + ] + else: + all_temps = [t_min] + else: + all_temps = [float(k) for k in eliashberg_data.keys()] + + remaining_temps = [ + t for t in all_temps if not any(abs(t - st) < 1e-4 for st in succeeded_temps) + ] + + new_temps_list = [] + new_nstemp = len(remaining_temps) + if is_linear_range and new_nstemp >= 2: + new_temps_list = [remaining_temps[0], remaining_temps[-1]] + else: + new_temps_list = remaining_temps + + updated_parameters = dict(parameters) + input_epw_new = updated_parameters.setdefault("INPUTEPW", {}) + + if isinstance(input_epw.get("temps"), str): + input_epw_new["temps"] = " ".join(str(t) for t in new_temps_list) + else: + input_epw_new["temps"] = new_temps_list + input_epw_new["nstemp"] = new_nstemp + input_epw_new.pop("tempsmin", None) + input_epw_new.pop("tempsmax", None) + + return updated_parameters, succeeded_temps, remaining_temps, eliashberg_data diff --git a/src/aiida_epw/workflows/base.py b/src/aiida_epw/workflows/base.py index ebd6376c..49c54e55 100644 --- a/src/aiida_epw/workflows/base.py +++ b/src/aiida_epw/workflows/base.py @@ -23,7 +23,6 @@ from aiida_epw.tools.error_handling import supercon as supercon_error_handling from aiida_epw.tools.error_handling import transport as transport_error_handling - EpwCalculation = CalculationFactory("epw.epw") @@ -671,3 +670,49 @@ def handle_temperature_out_of_range(self, calculation): self.ctx.is_finished = True self.report_error_handled(calculation, action_taken) return ProcessHandlerReport(True) + + @process_handler( + priority=650, + exit_codes=[ + EpwCalculation.exit_codes.ERROR_OUT_OF_WALLTIME, + EpwCalculation.exit_codes.ERROR_SCHEDULER_OUT_OF_WALLTIME, + ], + ) + def handle_out_of_walltime(self, calculation): + """Handle exit code 120 (scheduler walltime timeout) and 400 (software walltime timeout).""" + from aiida_epw.common.types import RestartType + + parameters = self.ctx.inputs.parameters.get_dict() + input_epw = parameters.get("INPUTEPW", {}) + is_eliashberg = input_epw.get("eliashberg", False) or any( + name in self.ctx.inputs + for name in ( + "momentum_dependence", + "full_bandwidth", + "real_axis", + "analytical_continuation", + ) + ) + restart_type = ( + self.ctx.inputs.restart_type.get_member() + if "restart_type" in self.ctx.inputs + else None + ) + parameters, next_restart, action_taken = ( + supercon_error_handling.prepare_walltime_recovery( + parameters, + calculation.outputs.output_parameters.get_dict(), + restart_type.name if restart_type else None, + is_eliashberg, + ) + ) + if parameters is None: + self.report_error_handled(calculation, action_taken) + return ProcessHandlerReport( + True, self.exit_codes.ERROR_KNOWN_UNRECOVERABLE_FAILURE + ) + self.ctx.inputs.parameters = orm.Dict(parameters) + self.ctx.inputs.restart_type = RestartType[next_restart] + self.ctx.inputs.parent_folder_epw = calculation.outputs.remote_folder + self.report_error_handled(calculation, action_taken) + return ProcessHandlerReport(True) diff --git a/src/aiida_epw/workflows/protocols/mobility.yaml b/src/aiida_epw/workflows/protocols/mobility.yaml index 51d5f029..59fb41cb 100644 --- a/src/aiida_epw/workflows/protocols/mobility.yaml +++ b/src/aiida_epw/workflows/protocols/mobility.yaml @@ -21,16 +21,9 @@ default_inputs: withmpi: True parameters: INPUTEPW: - elph: True - restart: true restart_step: 1000 etf_mem: 3 # I/O Settings - epbwrite: false - epbread: false - epwwrite: false - epwread: true - epmatkqread: false selecqread: false vme: 'wannier' #'wannier' or 'dipole' use_ws: true @@ -61,7 +54,6 @@ default_inputs: bfieldx: 0.0 bfieldy: 0.0 bfieldz: 1.0e-10 - wannierize: False default_protocol: moderate protocols: diff --git a/tests/calculations/test_epw_calcjob.py b/tests/calculations/test_epw_calcjob.py index 42b54e8d..ebad5afd 100644 --- a/tests/calculations/test_epw_calcjob.py +++ b/tests/calculations/test_epw_calcjob.py @@ -104,6 +104,16 @@ def test_epw_writes_explicit_fine_point_files( [ ("outdir", "./custom-out"), ("nk1", 8), + ("wannierize", True), + ("epwread", True), + ("epwwrite", True), + ("epbread", True), + ("epbwrite", True), + ("restart", True), + ("ep_coupling", True), + ("elph", True), + ("ephwrite", True), + ("epmatkqread", True), ], ) def test_epw_rejects_plugin_managed_keywords( @@ -205,75 +215,6 @@ def test_epw_accepts_parser_options_setting( assert calc_info.retrieve_list == ["aiida.out"] -def test_epw_accepts_manual_proj_for_wannierize( - fixture_sandbox, - fixture_localhost, - generate_calc_job, - generate_inputs_epw, - generate_remote_data, -): - """Direct EPW Wannierization should accept manual `proj` lists.""" - inputs = generate_inputs_epw( - parameters={"INPUTEPW": {"wannierize": True, "proj": ["Si:s", "Si:p"]}}, - parent_folder_nscf=generate_remote_data(fixture_localhost, "/remote/nscf"), - ) - - generate_calc_job(fixture_sandbox, "epw.epw", inputs) - - input_contents = Path(fixture_sandbox.abspath, "aiida.in").read_text() - assert "wannierize = .true." in input_contents - assert "proj(1) = 'Si:s'" in input_contents - assert "proj(2) = 'Si:p'" in input_contents - - -def test_epw_requires_nscf_parent_for_wannierize( - fixture_sandbox, generate_calc_job, generate_inputs_epw -): - """Direct EPW Wannierization should always stage an NSCF parent.""" - inputs = generate_inputs_epw( - parameters={"INPUTEPW": {"wannierize": True, "proj": ["Si:s"]}} - ) - - with pytest.raises(ValueError, match="parent_folder_nscf"): - generate_calc_job(fixture_sandbox, "epw.epw", inputs) - - -def test_epw_rejects_auto_projections_for_wannierize( - fixture_sandbox, - fixture_localhost, - generate_calc_job, - generate_inputs_epw, - generate_remote_data, -): - """Only manual `proj` entries are supported for EPW Wannierization.""" - inputs = generate_inputs_epw( - parameters={ - "INPUTEPW": {"wannierize": True, "auto_projections": True, "proj": ["Si:s"]} - }, - parent_folder_nscf=generate_remote_data(fixture_localhost, "/remote/nscf"), - ) - - with pytest.raises(ValueError, match="auto_projections"): - generate_calc_job(fixture_sandbox, "epw.epw", inputs) - - -def test_epw_requires_manual_proj_for_wannierize( - fixture_sandbox, - fixture_localhost, - generate_calc_job, - generate_inputs_epw, - generate_remote_data, -): - """Direct EPW Wannierization should require explicit manual projections.""" - inputs = generate_inputs_epw( - parameters={"INPUTEPW": {"wannierize": True}}, - parent_folder_nscf=generate_remote_data(fixture_localhost, "/remote/nscf"), - ) - - with pytest.raises(ValueError, match="Manual `proj` entries"): - generate_calc_job(fixture_sandbox, "epw.epw", inputs) - - def test_epw_additional_retrieve_list_emits_deprecation_warning( fixture_sandbox, generate_calc_job, generate_inputs_epw ): @@ -520,6 +461,100 @@ def test_epw_stages_ph_stash_folder_by_target_basepath( ) in calc_info.remote_copy_list +def test_epw_restart_type_parameter( + fixture_sandbox, + fixture_localhost, + generate_calc_job, + generate_inputs_epw, + generate_remote_data, +): + """Test that specifying `restart_type` updates the namelist correctly in the input file.""" + + # 1. NONE: start from scratch and write the reusable EPW data. + inputs_none = generate_inputs_epw(restart_type=RestartType.NONE) + generate_calc_job(fixture_sandbox, "epw.epw", inputs_none) + input_contents = Path(fixture_sandbox.abspath, "aiida.in").read_text() + assert "epwread = .false." in input_contents + assert "epwwrite = .true." in input_contents + assert "restart = .false." in input_contents + assert "ep_coupling = .true." in input_contents + assert "elph = .true." in input_contents + assert "epbwrite = .true." in input_contents + assert "epbread = .false." in input_contents + + # 2. EPHWRITE + inputs_write = generate_inputs_epw( + restart_type=RestartType.EPHWRITE, + parent_folder_epw=generate_remote_data(fixture_localhost, "/remote/epw"), + ) + generate_calc_job(fixture_sandbox, "epw.epw", inputs_write) + input_contents = Path(fixture_sandbox.abspath, "aiida.in").read_text() + assert "epwread = .true." in input_contents + assert "epwwrite = .false." in input_contents + assert "restart = .false." in input_contents + assert "ep_coupling = .true." in input_contents + assert "elph = .true." in input_contents + assert "ephwrite = .true." in input_contents + + # 3. EPHREAD_RESTART + inputs_read_restart = generate_inputs_epw( + restart_type=RestartType.EPHREAD_RESTART, + parameters={"INPUTEPW": {"scattering": True}}, + parent_folder_epw=generate_remote_data(fixture_localhost, "/remote/epw"), + ) + generate_calc_job(fixture_sandbox, "epw.epw", inputs_read_restart) + input_contents = Path(fixture_sandbox.abspath, "aiida.in").read_text() + assert "epwread = .true." in input_contents + assert "restart = .true." in input_contents + assert "ep_coupling = .false." in input_contents + assert "elph = .false." in input_contents + assert "ephwrite = .false." in input_contents + assert "epmatkqread = .true." in input_contents + + # 4. EPHREAD (without scattering) + inputs_read1 = generate_inputs_epw( + restart_type=RestartType.EPHREAD, + parent_folder_epw=generate_remote_data(fixture_localhost, "/remote/epw"), + ) + generate_calc_job(fixture_sandbox, "epw.epw", inputs_read1) + input_contents = Path(fixture_sandbox.abspath, "aiida.in").read_text() + assert "epwread = .true." in input_contents + assert "restart = .false." in input_contents + assert "ep_coupling = .false." in input_contents + assert "elph = .false." in input_contents + assert "ephwrite = .false." in input_contents + assert "epmatkqread" not in input_contents + + # 5. EPHREAD (with scattering) + inputs_read2 = generate_inputs_epw( + restart_type=RestartType.EPHREAD, + parameters={"INPUTEPW": {"scattering": True}}, + parent_folder_epw=generate_remote_data(fixture_localhost, "/remote/epw"), + ) + generate_calc_job(fixture_sandbox, "epw.epw", inputs_read2) + input_contents = Path(fixture_sandbox.abspath, "aiida.in").read_text() + assert "epwread = .true." in input_contents + assert "restart = .false." in input_contents + assert "ep_coupling = .false." in input_contents + assert "elph = .false." in input_contents + assert "ephwrite = .false." in input_contents + assert "epmatkqread = .true." in input_contents + + # 6. EPHWRITE_RESTART + inputs_write_restart = generate_inputs_epw( + restart_type=RestartType.EPHWRITE_RESTART, + parent_folder_epw=generate_remote_data(fixture_localhost, "/remote/epw"), + ) + generate_calc_job(fixture_sandbox, "epw.epw", inputs_write_restart) + input_contents = Path(fixture_sandbox.abspath, "aiida.in").read_text() + assert "epwread = .true." in input_contents + assert "epwwrite = .false." in input_contents + assert "restart = .true." in input_contents + assert "ep_coupling = .true." in input_contents + assert "elph = .true." in input_contents + assert "ephwrite = .true." in input_contents + + def test_epw_eliashberg_parameters( fixture_sandbox, generate_calc_job, generate_inputs_epw ): diff --git a/tests/workflows/test_base.py b/tests/workflows/test_base.py index 2b66485c..fb6ee71a 100644 --- a/tests/workflows/test_base.py +++ b/tests/workflows/test_base.py @@ -614,3 +614,276 @@ def report_error_handled(self, calculation, action): report.exit_code.status == EpwBaseWorkChain.exit_codes.ERROR_KNOWN_UNRECOVERABLE_FAILURE.status ) + + +def test_handle_out_of_walltime(aiida_localhost): + """Test that handle_out_of_walltime works correctly when Eliashberg and EPHWRITE are configured.""" + + class MockWorkChain: + exit_codes = EpwBaseWorkChain.exit_codes + + def __init__(self): + self.ctx = MagicMock() + self.report_messages = [] + + def report(self, msg): + self.report_messages.append(msg) + + def report_error_handled(self, calculation, action): + self.report(f"Calculation failed: {action}") + + handle_out_of_walltime = EpwBaseWorkChain.handle_out_of_walltime + + workchain = MockWorkChain() + + calc = MagicMock() + calc.outputs = MagicMock() + calc.outputs.remote_folder = orm.RemoteData( + computer=aiida_localhost, remote_path="/tmp" + ) + calc.outputs.output_parameters = orm.Dict(dict={}) + + initial_params = {"INPUTEPW": {"eliashberg": True}} + workchain.ctx.inputs = MagicMock() + workchain.ctx.inputs.parameters = orm.Dict(dict=initial_params) + restart_type_mock = MagicMock() + restart_type_mock.get_member.return_value = RestartType.EPHWRITE + workchain.ctx.inputs.restart_type = restart_type_mock + workchain.ctx.inputs.__contains__.side_effect = lambda key: key in ( + "restart_type", + "parameters", + ) + + report = workchain.handle_out_of_walltime.__wrapped__(calc) + + assert report.do_break is True + assert report.exit_code.status == 0 + assert workchain.ctx.inputs.restart_type == RestartType.EPHWRITE + assert workchain.ctx.inputs.parent_folder_epw == calc.outputs.remote_folder + + +def test_handle_out_of_walltime_unsupported(aiida_localhost): + """Test that handle_out_of_walltime aborts for unsupported calculation/restart configurations.""" + + class MockWorkChain: + exit_codes = EpwBaseWorkChain.exit_codes + + def __init__(self): + self.ctx = MagicMock() + self.report_messages = [] + + def report(self, msg): + self.report_messages.append(msg) + + def report_error_handled(self, calculation, action): + self.report(f"Calculation failed: {action}") + + handle_out_of_walltime = EpwBaseWorkChain.handle_out_of_walltime + + workchain = MockWorkChain() + + calc = MagicMock() + calc.outputs = MagicMock() + calc.outputs.remote_folder = orm.RemoteData( + computer=aiida_localhost, remote_path="/tmp" + ) + calc.outputs.output_parameters = orm.Dict(dict={}) + + initial_params = {"INPUTEPW": {"band_plot": True}} + workchain.ctx.inputs = MagicMock() + workchain.ctx.inputs.parameters = orm.Dict(dict=initial_params) + workchain.ctx.inputs.__contains__.side_effect = lambda key: key == "parameters" + + report = workchain.handle_out_of_walltime.__wrapped__(calc) + + assert report.do_break is True + assert ( + report.exit_code.status + == EpwBaseWorkChain.exit_codes.ERROR_KNOWN_UNRECOVERABLE_FAILURE.status + ) + + +def test_handle_out_of_walltime_from_eph_success(aiida_localhost): + """Test that handle_out_of_walltime FROM_EPH mode correctly pops succeeded temperatures and restarts.""" + + class MockWorkChain: + exit_codes = EpwBaseWorkChain.exit_codes + + def __init__(self): + self.ctx = MagicMock() + self.report_messages = [] + + def report(self, msg): + self.report_messages.append(msg) + + def report_error_handled(self, calculation, action): + self.report(f"Calculation failed: {action}") + + handle_out_of_walltime = EpwBaseWorkChain.handle_out_of_walltime + + workchain = MockWorkChain() + + calc = MagicMock() + calc.outputs = MagicMock() + calc.outputs.remote_folder = orm.RemoteData( + computer=aiida_localhost, remote_path="/tmp" + ) + + calc.outputs.output_parameters = orm.Dict( + dict={ + "isotropic_eliashberg": { + "2.0": { + "nsiw": 400, + "iterations": {"ethr": [1e-8], "znormi": [1.2], "deltai": [2.5]}, + "pade": {"delta": 2.4, "znorm": 1.2}, + }, + "1.0": { + "nsiw": 800, + "iterations": {"ethr": [None], "znormi": [None], "deltai": [None]}, + "pade": {"delta": None, "znorm": None}, + }, + } + } + ) + + initial_params = { + "INPUTEPW": { + "eliashberg": True, + "temps": [1.0, 2.0], + "nstemp": 2, + } + } + workchain.ctx.inputs = MagicMock() + workchain.ctx.inputs.parameters = orm.Dict(dict=initial_params) + restart_type_mock = MagicMock() + restart_type_mock.get_member.return_value = RestartType.FROM_EPH + workchain.ctx.inputs.restart_type = restart_type_mock + workchain.ctx.inputs.__contains__.side_effect = lambda key: key in ( + "restart_type", + "parameters", + ) + + report = workchain.handle_out_of_walltime.__wrapped__(calc) + + assert report.do_break is True + assert report.exit_code.status == 0 + assert workchain.ctx.inputs.parent_folder_epw == calc.outputs.remote_folder + + updated_params = workchain.ctx.inputs.parameters.get_dict() + assert updated_params["INPUTEPW"]["temps"] == [1.0] + assert updated_params["INPUTEPW"]["nstemp"] == 1 + + +def test_handle_out_of_walltime_from_eph_failure(aiida_localhost): + """Test that handle_out_of_walltime FROM_EPH mode aborts if no temperatures succeeded.""" + + class MockWorkChain: + exit_codes = EpwBaseWorkChain.exit_codes + + def __init__(self): + self.ctx = MagicMock() + self.report_messages = [] + + def report(self, msg): + self.report_messages.append(msg) + + def report_error_handled(self, calculation, action): + self.report(f"Calculation failed: {action}") + + handle_out_of_walltime = EpwBaseWorkChain.handle_out_of_walltime + + workchain = MockWorkChain() + + calc = MagicMock() + calc.outputs = MagicMock() + calc.outputs.remote_folder = orm.RemoteData( + computer=aiida_localhost, remote_path="/tmp" + ) + + calc.outputs.output_parameters = orm.Dict( + dict={ + "isotropic_eliashberg": { + "1.0": { + "nsiw": 800, + "iterations": {"ethr": [None], "znormi": [None], "deltai": [None]}, + "pade": {"delta": None, "znorm": None}, + }, + } + } + ) + + initial_params = { + "INPUTEPW": { + "eliashberg": True, + "temps": [1.0], + "nstemp": 1, + } + } + workchain.ctx.inputs = MagicMock() + workchain.ctx.inputs.parameters = orm.Dict(dict=initial_params) + restart_type_mock = MagicMock() + restart_type_mock.get_member.return_value = RestartType.FROM_EPH + workchain.ctx.inputs.restart_type = restart_type_mock + workchain.ctx.inputs.__contains__.side_effect = lambda key: key in ( + "restart_type", + "parameters", + ) + + report = workchain.handle_out_of_walltime.__wrapped__(calc) + + assert report.do_break is True + assert ( + report.exit_code.status + == EpwBaseWorkChain.exit_codes.ERROR_KNOWN_UNRECOVERABLE_FAILURE.status + ) + + +def test_handle_out_of_walltime_fresh_eliashberg(aiida_localhost): + """Test that a fresh Eliashberg run aborts without a restart checkpoint.""" + + class MockWorkChain: + exit_codes = EpwBaseWorkChain.exit_codes + + def __init__(self): + self.ctx = MagicMock() + self.report_messages = [] + + def report(self, msg): + self.report_messages.append(msg) + + def report_error_handled(self, calculation, action): + self.report(f"Calculation failed: {action}") + + handle_out_of_walltime = EpwBaseWorkChain.handle_out_of_walltime + + workchain = MockWorkChain() + + calc = MagicMock() + calc.outputs = MagicMock() + calc.outputs.remote_folder = orm.RemoteData( + computer=aiida_localhost, remote_path="/tmp" + ) + + initial_params = {"INPUTEPW": {"temps": [10.0]}} + workchain.ctx.inputs = MagicMock() + workchain.ctx.inputs.parameters = orm.Dict(dict=initial_params) + restart_type_mock = MagicMock() + restart_type_mock.get_member.return_value = RestartType.FROM_SCRATCH + workchain.ctx.inputs.restart_type = restart_type_mock + workchain.ctx.inputs.__contains__.side_effect = lambda key: key in ( + "restart_type", + "parameters", + "momentum_dependence", + ) + + report = workchain.handle_out_of_walltime.__wrapped__(calc) + + assert report.do_break is True + assert ( + report.exit_code.status + == EpwBaseWorkChain.exit_codes.ERROR_KNOWN_UNRECOVERABLE_FAILURE.status + ) + assert any( + "Automatic recovery requires an ephwrite checkpoint" in msg + for msg in workchain.report_messages + ) diff --git a/tests/workflows/test_supercon.py b/tests/workflows/test_supercon.py index c9b1f3df..b8d75933 100644 --- a/tests/workflows/test_supercon.py +++ b/tests/workflows/test_supercon.py @@ -223,88 +223,4 @@ def test_epw_base_restart_types(fixture_code, generate_structure): ) # We should be able to set and access restart_type on the builder - builder.restart_type = RestartType.EPHWRITE - assert builder.restart_type == RestartType.EPHWRITE - - -def test_supercon_get_builder_from_protocol_default( - fixture_code, - generate_structure, - generate_remote_data, - fixture_localhost, - monkeypatch, -): - """Test get_builder_from_protocol default values.""" - from aiida_epw.workflows.supercon import SuperConWorkChain - from plumpy.ports import Port, PortNamespace - - monkeypatch.setattr(Port, "validate", lambda *a, **k: None) - monkeypatch.setattr(PortNamespace, "validate", lambda *a, **k: None) - - epw_code = fixture_code("epw.epw") - structure = generate_structure() - remote_stash = generate_remote_data(fixture_localhost, "/tmp/remote_stash") - - parent_epw = MagicMock() - parent_epw.process_label = "EpwBaseWorkChain" - parent_epw.inputs = MagicMock() - parent_epw.inputs.structure = structure - parent_epw.inputs.code = epw_code - parent_epw.inputs.kpoints = orm.KpointsData() - parent_epw.inputs.qpoints = orm.KpointsData() - parent_epw.outputs = MagicMock() - parent_epw.outputs.remote_stash = remote_stash - - builder = SuperConWorkChain.get_builder_from_protocol( - epw_code=epw_code, - parent_epw=parent_epw, - protocol="fast", - ) - - # Verify that the builders are configured according to the protocol - assert "code" in builder.epw_interp - assert "code" in builder.epw_final_iso - assert "code" in builder.epw_final_aniso - - # epw_interp check - if "restart_type" in builder.epw_interp: - from aiida_epw.common import RestartType - - assert builder.epw_interp.restart_type == RestartType.EPHWRITE - - # epw_final_iso check - if "momentum_dependence" in builder.epw_final_iso: - assert not builder.epw_final_iso.momentum_dependence.value - assert not builder.epw_final_iso.real_axis.value - if "calculation_type" in builder.epw_final_iso: - from aiida_epw.common.types import CalculationTypes - - assert ( - builder.epw_final_iso.calculation_type.get_member() - == CalculationTypes.ELIASHBERG - ) - if "restart_type" in builder.epw_final_iso: - from aiida_epw.common import RestartType - - assert builder.epw_final_iso.restart_type == RestartType.EPHREAD - assert ( - builder.epw_final_iso.parameters.get_dict()["INPUTEPW"].get("tc_linear", False) - is False - ) - - # epw_final_aniso check - if "momentum_dependence" in builder.epw_final_aniso: - assert builder.epw_final_aniso.momentum_dependence.value - assert not builder.epw_final_aniso.full_bandwidth.value - assert not builder.epw_final_aniso.real_axis.value - if "calculation_type" in builder.epw_final_aniso: - from aiida_epw.common.types import CalculationTypes - - assert ( - builder.epw_final_aniso.calculation_type.get_member() - == CalculationTypes.ELIASHBERG - ) - if "restart_type" in builder.epw_final_aniso: - from aiida_epw.common import RestartType - assert builder.epw_final_aniso.restart_type == RestartType.EPHREAD From 5817b6be42a059b9f36fb45520805f29fd5fcee1 Mon Sep 17 00:00:00 2001 From: ymzhang0 Date: Fri, 21 Aug 2026 11:34:28 +0200 Subject: [PATCH 4/9] feat(parser): add scheduler stderr walltime parsing and error recovery --- src/aiida_epw/calculations/epw.py | 63 +------ src/aiida_epw/parsers/epw.py | 36 ++++ .../tools/error_handling/supercon.py | 34 +++- tests/calculations/test_epw_calcjob.py | 173 +++++++----------- tests/parsers/test_epw.py | 69 +++++++ tests/workflows/test_base.py | 8 +- tests/workflows/test_supercon.py | 84 +++++++++ 7 files changed, 292 insertions(+), 175 deletions(-) diff --git a/src/aiida_epw/calculations/epw.py b/src/aiida_epw/calculations/epw.py index 0a9dada0..13cb48c6 100644 --- a/src/aiida_epw/calculations/epw.py +++ b/src/aiida_epw/calculations/epw.py @@ -77,16 +77,6 @@ class EpwCalculation(NamelistsCalculation): ("INPUTEPW", "nkf1"), ("INPUTEPW", "nkf2"), ("INPUTEPW", "nkf3"), - ("INPUTEPW", "wannierize"), - ("INPUTEPW", "epwread"), - ("INPUTEPW", "epwwrite"), - ("INPUTEPW", "epbread"), - ("INPUTEPW", "epbwrite"), - ("INPUTEPW", "restart"), - ("INPUTEPW", "ep_coupling"), - ("INPUTEPW", "elph"), - ("INPUTEPW", "ephwrite"), - ("INPUTEPW", "epmatkqread"), ("INPUTEPW", "eliashberg"), ("INPUTEPW", "liso"), ("INPUTEPW", "laniso"), @@ -473,10 +463,10 @@ def validate_restart_inputs(cls, parameters, inputs): inputs["restart_type"].get_member() if "restart_type" in inputs else None ) parent_restart_types = ( - RestartType.FROM_EPB, - RestartType.FROM_EPMATWP, RestartType.EPHWRITE, - RestartType.FROM_EPH, + RestartType.EPHREAD, + RestartType.EPHWRITE_RESTART, + RestartType.EPHREAD_RESTART, RestartType.EPWREAD, ) uses_parent = restart_type in parent_restart_types @@ -756,53 +746,6 @@ def prepare_input_parameters(self, folder, parameters): self.cap_nstemp(inputepw_parameters) - if "restart_type" in self.inputs: - restart_val = self.inputs.restart_type.get_member() - from aiida_epw.common import RestartType - - if restart_val is RestartType.NONE: - inputepw_parameters["epwread"] = False - inputepw_parameters["epwwrite"] = True - inputepw_parameters["restart"] = False - inputepw_parameters["ep_coupling"] = True - inputepw_parameters["elph"] = True - inputepw_parameters["epbwrite"] = True - inputepw_parameters["epbread"] = False - elif restart_val is RestartType.EPHWRITE: - inputepw_parameters["epwread"] = True - inputepw_parameters["epwwrite"] = False - inputepw_parameters["restart"] = False - inputepw_parameters["ep_coupling"] = True - inputepw_parameters["elph"] = True - inputepw_parameters["ephwrite"] = True - elif restart_val is RestartType.EPHWRITE_RESTART: - inputepw_parameters["epwread"] = True - inputepw_parameters["epwwrite"] = False - inputepw_parameters["restart"] = True - inputepw_parameters["ep_coupling"] = True - inputepw_parameters["elph"] = True - inputepw_parameters["ephwrite"] = True - elif restart_val in ( - RestartType.EPHREAD, - RestartType.EPHREAD_RESTART, - ): - inputepw_parameters["epwread"] = True - inputepw_parameters["restart"] = ( - restart_val is RestartType.EPHREAD_RESTART - ) - inputepw_parameters["ep_coupling"] = False - inputepw_parameters["elph"] = False - inputepw_parameters["ephwrite"] = False - if inputepw_parameters.get("scattering", False): - inputepw_parameters["epmatkqread"] = True - elif restart_val is RestartType.EPWREAD: - inputepw_parameters["epwread"] = True - inputepw_parameters["epwwrite"] = False - inputepw_parameters["epbwrite"] = False - inputepw_parameters["epbread"] = False - inputepw_parameters.setdefault("restart", False) - inputepw_parameters["ep_coupling"] = True - inputepw_parameters["elph"] = True # Override Eliashberg settings in parameters if inputs are specified eliashberg_any = any( f in self.inputs diff --git a/src/aiida_epw/parsers/epw.py b/src/aiida_epw/parsers/epw.py index dd48419a..50c80995 100644 --- a/src/aiida_epw/parsers/epw.py +++ b/src/aiida_epw/parsers/epw.py @@ -43,12 +43,43 @@ class EpwParser(BaseParser): class_error_map = { r"internal error,\s*cannot bracket Ef": "ERROR_CANNOT_BRACKET_EF", } + _scheduler_walltime_patterns = ( + r"(?i)due to time limit", + r"(?i)time limit reached", + r"(?i)walltime\s+(?:limit\s+)?exceeded", + r"(?i)cputime\s+(?:limit\s+)?exceeded", + r"(?i)term_runlimit", + r"(?i)exceeded.*wallclock time limit", + r"(?i)job.*cancelled.*due to time limit", + ) @staticmethod def get_parser_settings_key(): """Return the settings key reserved for parser-specific options.""" return "parser_options" + def check_scheduler_errors(self, logs): + """Check scheduler stderr/stdout or node exit status for scheduler-level aborts.""" + scheduler_walltime_status = getattr( + EpwCalculation.exit_codes, "ERROR_SCHEDULER_OUT_OF_WALLTIME", None + ) + if ( + scheduler_walltime_status is not None + and self.node.exit_status == scheduler_walltime_status.status + ): + logs.error.append("ERROR_OUT_OF_WALLTIME") + return + + scheduler_stderr = self.get_retrieved_content("_scheduler-stderr.txt") or "" + scheduler_stdout = self.get_retrieved_content("_scheduler-stdout.txt") or "" + scheduler_output = f"{scheduler_stderr}\n{scheduler_stdout}" + + if scheduler_output.strip(): + for pattern in self._scheduler_walltime_patterns: + if re.search(pattern, scheduler_output): + logs.error.append("ERROR_OUT_OF_WALLTIME") + return + def get_retrieved_content(self, *filenames): """Return the content of the first retrieved file that exists.""" for filename in filenames: @@ -88,6 +119,8 @@ def parse(self, **kwargs): parsed_epw, logs = self.parse_stdout(stdout, logs, code_version=code_version) parsed_data.update(parsed_epw) + self.check_scheduler_errors(logs) + elbands_contents = self.get_retrieved_content( EpwCalculation._output_elbands_file ) @@ -221,6 +254,9 @@ def parse(self, **kwargs): if "ERROR_FACTORIZATION" in logs.error: return self.exit(self.exit_codes.get("ERROR_FACTORIZATION"), logs) + if "ERROR_OUT_OF_WALLTIME" in logs.error: + return self.exit(self.exit_codes.get("ERROR_OUT_OF_WALLTIME"), logs) + for exit_code in list(self.get_error_map().values()): if exit_code in logs.error: return self.exit(self.exit_codes.get(exit_code), logs) diff --git a/src/aiida_epw/tools/error_handling/supercon.py b/src/aiida_epw/tools/error_handling/supercon.py index caaef436..ec910b2c 100644 --- a/src/aiida_epw/tools/error_handling/supercon.py +++ b/src/aiida_epw/tools/error_handling/supercon.py @@ -190,6 +190,8 @@ def prepare_walltime_recovery( parameters, output_parameters, restart_type, is_eliashberg ): """Return ``(parameters, restart_type, action)`` for a recoverable timeout.""" + from aiida_epw.common.types import RestartType + if not is_eliashberg: return ( None, @@ -200,17 +202,34 @@ def prepare_walltime_recovery( ), ) - if restart_type in ("EPHWRITE", "ephwrite"): + if restart_type in ( + "EPHWRITE", + "ephwrite", + "EPHWRITE_RESTART", + "ephwrite_restart", + ): + next_type = ( + "EPHWRITE_RESTART" + if hasattr(RestartType, "EPHWRITE_RESTART") + else "EPHWRITE" + ) return ( parameters, - "EPHWRITE", + next_type, ( "Walltime reached during Eliashberg ephwrite calculation. " "Restarting from the last checkpoint." ), ) - if restart_type in ("FROM_EPH", "from_eph", "EPHREAD", "ephread"): + if restart_type in ( + "FROM_EPH", + "from_eph", + "EPHREAD", + "ephread", + "EPHREAD_RESTART", + "ephread_restart", + ): outputs = output_parameters or {} eliashberg_data = ( outputs.get("isotropic_eliashberg") @@ -242,11 +261,12 @@ def prepare_walltime_recovery( input_epw_new["nstemp"] = len(remaining_temps) input_epw_new.pop("tempsmin", None) input_epw_new.pop("tempsmax", None) + next_type = "EPHREAD" if hasattr(RestartType, "EPHREAD") else "FROM_EPH" return ( parameters, - "FROM_EPH", + next_type, ( - "Walltime reached during Eliashberg from_eph calculation. " + f"Walltime reached during Eliashberg {next_type.lower()} calculation. " f"Removed successfully calculated temperatures: {succeeded_temps}. " "Restarting." ), @@ -255,12 +275,12 @@ def prepare_walltime_recovery( None, None, ( - "Walltime reached during Eliashberg from_eph calculation but no " + "Walltime reached during Eliashberg ephread calculation but no " "temperatures finished. Aborting." ), ) - if restart_type in ("FROM_SCRATCH", "from_scratch", "NONE", "none"): + if restart_type in ("FROM_SCRATCH", "from_scratch", "NONE", "none", None): return ( None, None, diff --git a/tests/calculations/test_epw_calcjob.py b/tests/calculations/test_epw_calcjob.py index ebad5afd..42b54e8d 100644 --- a/tests/calculations/test_epw_calcjob.py +++ b/tests/calculations/test_epw_calcjob.py @@ -104,16 +104,6 @@ def test_epw_writes_explicit_fine_point_files( [ ("outdir", "./custom-out"), ("nk1", 8), - ("wannierize", True), - ("epwread", True), - ("epwwrite", True), - ("epbread", True), - ("epbwrite", True), - ("restart", True), - ("ep_coupling", True), - ("elph", True), - ("ephwrite", True), - ("epmatkqread", True), ], ) def test_epw_rejects_plugin_managed_keywords( @@ -215,6 +205,75 @@ def test_epw_accepts_parser_options_setting( assert calc_info.retrieve_list == ["aiida.out"] +def test_epw_accepts_manual_proj_for_wannierize( + fixture_sandbox, + fixture_localhost, + generate_calc_job, + generate_inputs_epw, + generate_remote_data, +): + """Direct EPW Wannierization should accept manual `proj` lists.""" + inputs = generate_inputs_epw( + parameters={"INPUTEPW": {"wannierize": True, "proj": ["Si:s", "Si:p"]}}, + parent_folder_nscf=generate_remote_data(fixture_localhost, "/remote/nscf"), + ) + + generate_calc_job(fixture_sandbox, "epw.epw", inputs) + + input_contents = Path(fixture_sandbox.abspath, "aiida.in").read_text() + assert "wannierize = .true." in input_contents + assert "proj(1) = 'Si:s'" in input_contents + assert "proj(2) = 'Si:p'" in input_contents + + +def test_epw_requires_nscf_parent_for_wannierize( + fixture_sandbox, generate_calc_job, generate_inputs_epw +): + """Direct EPW Wannierization should always stage an NSCF parent.""" + inputs = generate_inputs_epw( + parameters={"INPUTEPW": {"wannierize": True, "proj": ["Si:s"]}} + ) + + with pytest.raises(ValueError, match="parent_folder_nscf"): + generate_calc_job(fixture_sandbox, "epw.epw", inputs) + + +def test_epw_rejects_auto_projections_for_wannierize( + fixture_sandbox, + fixture_localhost, + generate_calc_job, + generate_inputs_epw, + generate_remote_data, +): + """Only manual `proj` entries are supported for EPW Wannierization.""" + inputs = generate_inputs_epw( + parameters={ + "INPUTEPW": {"wannierize": True, "auto_projections": True, "proj": ["Si:s"]} + }, + parent_folder_nscf=generate_remote_data(fixture_localhost, "/remote/nscf"), + ) + + with pytest.raises(ValueError, match="auto_projections"): + generate_calc_job(fixture_sandbox, "epw.epw", inputs) + + +def test_epw_requires_manual_proj_for_wannierize( + fixture_sandbox, + fixture_localhost, + generate_calc_job, + generate_inputs_epw, + generate_remote_data, +): + """Direct EPW Wannierization should require explicit manual projections.""" + inputs = generate_inputs_epw( + parameters={"INPUTEPW": {"wannierize": True}}, + parent_folder_nscf=generate_remote_data(fixture_localhost, "/remote/nscf"), + ) + + with pytest.raises(ValueError, match="Manual `proj` entries"): + generate_calc_job(fixture_sandbox, "epw.epw", inputs) + + def test_epw_additional_retrieve_list_emits_deprecation_warning( fixture_sandbox, generate_calc_job, generate_inputs_epw ): @@ -461,100 +520,6 @@ def test_epw_stages_ph_stash_folder_by_target_basepath( ) in calc_info.remote_copy_list -def test_epw_restart_type_parameter( - fixture_sandbox, - fixture_localhost, - generate_calc_job, - generate_inputs_epw, - generate_remote_data, -): - """Test that specifying `restart_type` updates the namelist correctly in the input file.""" - - # 1. NONE: start from scratch and write the reusable EPW data. - inputs_none = generate_inputs_epw(restart_type=RestartType.NONE) - generate_calc_job(fixture_sandbox, "epw.epw", inputs_none) - input_contents = Path(fixture_sandbox.abspath, "aiida.in").read_text() - assert "epwread = .false." in input_contents - assert "epwwrite = .true." in input_contents - assert "restart = .false." in input_contents - assert "ep_coupling = .true." in input_contents - assert "elph = .true." in input_contents - assert "epbwrite = .true." in input_contents - assert "epbread = .false." in input_contents - - # 2. EPHWRITE - inputs_write = generate_inputs_epw( - restart_type=RestartType.EPHWRITE, - parent_folder_epw=generate_remote_data(fixture_localhost, "/remote/epw"), - ) - generate_calc_job(fixture_sandbox, "epw.epw", inputs_write) - input_contents = Path(fixture_sandbox.abspath, "aiida.in").read_text() - assert "epwread = .true." in input_contents - assert "epwwrite = .false." in input_contents - assert "restart = .false." in input_contents - assert "ep_coupling = .true." in input_contents - assert "elph = .true." in input_contents - assert "ephwrite = .true." in input_contents - - # 3. EPHREAD_RESTART - inputs_read_restart = generate_inputs_epw( - restart_type=RestartType.EPHREAD_RESTART, - parameters={"INPUTEPW": {"scattering": True}}, - parent_folder_epw=generate_remote_data(fixture_localhost, "/remote/epw"), - ) - generate_calc_job(fixture_sandbox, "epw.epw", inputs_read_restart) - input_contents = Path(fixture_sandbox.abspath, "aiida.in").read_text() - assert "epwread = .true." in input_contents - assert "restart = .true." in input_contents - assert "ep_coupling = .false." in input_contents - assert "elph = .false." in input_contents - assert "ephwrite = .false." in input_contents - assert "epmatkqread = .true." in input_contents - - # 4. EPHREAD (without scattering) - inputs_read1 = generate_inputs_epw( - restart_type=RestartType.EPHREAD, - parent_folder_epw=generate_remote_data(fixture_localhost, "/remote/epw"), - ) - generate_calc_job(fixture_sandbox, "epw.epw", inputs_read1) - input_contents = Path(fixture_sandbox.abspath, "aiida.in").read_text() - assert "epwread = .true." in input_contents - assert "restart = .false." in input_contents - assert "ep_coupling = .false." in input_contents - assert "elph = .false." in input_contents - assert "ephwrite = .false." in input_contents - assert "epmatkqread" not in input_contents - - # 5. EPHREAD (with scattering) - inputs_read2 = generate_inputs_epw( - restart_type=RestartType.EPHREAD, - parameters={"INPUTEPW": {"scattering": True}}, - parent_folder_epw=generate_remote_data(fixture_localhost, "/remote/epw"), - ) - generate_calc_job(fixture_sandbox, "epw.epw", inputs_read2) - input_contents = Path(fixture_sandbox.abspath, "aiida.in").read_text() - assert "epwread = .true." in input_contents - assert "restart = .false." in input_contents - assert "ep_coupling = .false." in input_contents - assert "elph = .false." in input_contents - assert "ephwrite = .false." in input_contents - assert "epmatkqread = .true." in input_contents - - # 6. EPHWRITE_RESTART - inputs_write_restart = generate_inputs_epw( - restart_type=RestartType.EPHWRITE_RESTART, - parent_folder_epw=generate_remote_data(fixture_localhost, "/remote/epw"), - ) - generate_calc_job(fixture_sandbox, "epw.epw", inputs_write_restart) - input_contents = Path(fixture_sandbox.abspath, "aiida.in").read_text() - assert "epwread = .true." in input_contents - assert "epwwrite = .false." in input_contents - assert "restart = .true." in input_contents - assert "ep_coupling = .true." in input_contents - assert "elph = .true." in input_contents - assert "ephwrite = .true." in input_contents - - def test_epw_eliashberg_parameters( fixture_sandbox, generate_calc_job, generate_inputs_epw ): diff --git a/tests/parsers/test_epw.py b/tests/parsers/test_epw.py index 6f53f481..409add28 100644 --- a/tests/parsers/test_epw.py +++ b/tests/parsers/test_epw.py @@ -1,5 +1,6 @@ """Tests for the `EpwParser`.""" +import io import textwrap import pytest @@ -387,3 +388,71 @@ def test_parse_aniso_imag_fbw_returns_typed_data(): assert aniso_imag.get_array("temp_3_00_freq_0_znorm").tolist() == [1.0] assert aniso_imag.get_array("temp_3_00_freq_0_delta").tolist() == [0.5] assert aniso_imag.get_array("temp_3_00_freq_0_shift").tolist() == [0.2] + + +def test_epw_parser_out_of_walltime_from_scheduler_stderr(aiida_localhost): + """Test that EpwParser returns ERROR_OUT_OF_WALLTIME when scheduler stderr indicates time limit.""" + calc_entry_point = format_entry_point_string( + group="aiida.calculations", name="epw.epw" + ) + node = orm.CalcJobNode(computer=aiida_localhost, process_type=calc_entry_point) + node.base.attributes.set("output_filename", "aiida.out") + node.store() + + retrieved = orm.FolderData() + retrieved.base.repository.put_object_from_filelike( + io.StringIO("Program EPW v.5.8.0\nsome incomplete stdout..."), + "aiida.out", + ) + scheduler_stderr = """Lmod is automatically replacing "craype-x86-rome" with "craype-x86-milan". +srun: Job step aborted: Waiting up to 32 seconds for job step to finish. +slurmstepd: error: *** JOB 21415098 ON nid002500 CANCELLED AT 2026-08-20T18:25:38 DUE TO TIME LIMIT *** +slurmstepd: error: *** STEP 21415098.0 ON nid002500 CANCELLED AT 2026-08-20T18:25:38 DUE TO TIME LIMIT *** +""" + retrieved.base.repository.put_object_from_filelike( + io.StringIO(scheduler_stderr), + "_scheduler-stderr.txt", + ) + retrieved.base.links.add_incoming( + node, link_type=LinkType.CREATE, link_label="retrieved" + ) + retrieved.store() + + results, calcfunction = EpwParser.parse_from_node(node, store_provenance=False) + + assert calcfunction.is_failed + assert ( + calcfunction.exit_status + == EpwCalculation.exit_codes.ERROR_OUT_OF_WALLTIME.status + ) + + +def test_epw_parser_out_of_walltime_from_node_exit_status(aiida_localhost): + """Test that EpwParser returns ERROR_OUT_OF_WALLTIME when node exit status was set by scheduler.""" + calc_entry_point = format_entry_point_string( + group="aiida.calculations", name="epw.epw" + ) + node = orm.CalcJobNode(computer=aiida_localhost, process_type=calc_entry_point) + node.base.attributes.set("output_filename", "aiida.out") + node.set_exit_status( + EpwCalculation.exit_codes.ERROR_SCHEDULER_OUT_OF_WALLTIME.status + ) + node.store() + + retrieved = orm.FolderData() + retrieved.base.repository.put_object_from_filelike( + io.StringIO("Program EPW v.5.8.0\nsome incomplete stdout..."), + "aiida.out", + ) + retrieved.base.links.add_incoming( + node, link_type=LinkType.CREATE, link_label="retrieved" + ) + retrieved.store() + + results, calcfunction = EpwParser.parse_from_node(node, store_provenance=False) + + assert calcfunction.is_failed + assert ( + calcfunction.exit_status + == EpwCalculation.exit_codes.ERROR_OUT_OF_WALLTIME.status + ) diff --git a/tests/workflows/test_base.py b/tests/workflows/test_base.py index fb6ee71a..764d8114 100644 --- a/tests/workflows/test_base.py +++ b/tests/workflows/test_base.py @@ -658,7 +658,7 @@ def report_error_handled(self, calculation, action): assert report.do_break is True assert report.exit_code.status == 0 - assert workchain.ctx.inputs.restart_type == RestartType.EPHWRITE + assert workchain.ctx.inputs.restart_type == RestartType.EPHWRITE_RESTART assert workchain.ctx.inputs.parent_folder_epw == calc.outputs.remote_folder @@ -756,7 +756,7 @@ def report_error_handled(self, calculation, action): workchain.ctx.inputs = MagicMock() workchain.ctx.inputs.parameters = orm.Dict(dict=initial_params) restart_type_mock = MagicMock() - restart_type_mock.get_member.return_value = RestartType.FROM_EPH + restart_type_mock.get_member.return_value = RestartType.EPHREAD workchain.ctx.inputs.restart_type = restart_type_mock workchain.ctx.inputs.__contains__.side_effect = lambda key: key in ( "restart_type", @@ -822,7 +822,7 @@ def report_error_handled(self, calculation, action): workchain.ctx.inputs = MagicMock() workchain.ctx.inputs.parameters = orm.Dict(dict=initial_params) restart_type_mock = MagicMock() - restart_type_mock.get_member.return_value = RestartType.FROM_EPH + restart_type_mock.get_member.return_value = RestartType.EPHREAD workchain.ctx.inputs.restart_type = restart_type_mock workchain.ctx.inputs.__contains__.side_effect = lambda key: key in ( "restart_type", @@ -868,7 +868,7 @@ def report_error_handled(self, calculation, action): workchain.ctx.inputs = MagicMock() workchain.ctx.inputs.parameters = orm.Dict(dict=initial_params) restart_type_mock = MagicMock() - restart_type_mock.get_member.return_value = RestartType.FROM_SCRATCH + restart_type_mock.get_member.return_value = RestartType.NONE workchain.ctx.inputs.restart_type = restart_type_mock workchain.ctx.inputs.__contains__.side_effect = lambda key: key in ( "restart_type", diff --git a/tests/workflows/test_supercon.py b/tests/workflows/test_supercon.py index b8d75933..c9b1f3df 100644 --- a/tests/workflows/test_supercon.py +++ b/tests/workflows/test_supercon.py @@ -223,4 +223,88 @@ def test_epw_base_restart_types(fixture_code, generate_structure): ) # We should be able to set and access restart_type on the builder + builder.restart_type = RestartType.EPHWRITE + assert builder.restart_type == RestartType.EPHWRITE + + +def test_supercon_get_builder_from_protocol_default( + fixture_code, + generate_structure, + generate_remote_data, + fixture_localhost, + monkeypatch, +): + """Test get_builder_from_protocol default values.""" + from aiida_epw.workflows.supercon import SuperConWorkChain + from plumpy.ports import Port, PortNamespace + + monkeypatch.setattr(Port, "validate", lambda *a, **k: None) + monkeypatch.setattr(PortNamespace, "validate", lambda *a, **k: None) + + epw_code = fixture_code("epw.epw") + structure = generate_structure() + remote_stash = generate_remote_data(fixture_localhost, "/tmp/remote_stash") + + parent_epw = MagicMock() + parent_epw.process_label = "EpwBaseWorkChain" + parent_epw.inputs = MagicMock() + parent_epw.inputs.structure = structure + parent_epw.inputs.code = epw_code + parent_epw.inputs.kpoints = orm.KpointsData() + parent_epw.inputs.qpoints = orm.KpointsData() + parent_epw.outputs = MagicMock() + parent_epw.outputs.remote_stash = remote_stash + + builder = SuperConWorkChain.get_builder_from_protocol( + epw_code=epw_code, + parent_epw=parent_epw, + protocol="fast", + ) + + # Verify that the builders are configured according to the protocol + assert "code" in builder.epw_interp + assert "code" in builder.epw_final_iso + assert "code" in builder.epw_final_aniso + + # epw_interp check + if "restart_type" in builder.epw_interp: + from aiida_epw.common import RestartType + + assert builder.epw_interp.restart_type == RestartType.EPHWRITE + + # epw_final_iso check + if "momentum_dependence" in builder.epw_final_iso: + assert not builder.epw_final_iso.momentum_dependence.value + assert not builder.epw_final_iso.real_axis.value + if "calculation_type" in builder.epw_final_iso: + from aiida_epw.common.types import CalculationTypes + + assert ( + builder.epw_final_iso.calculation_type.get_member() + == CalculationTypes.ELIASHBERG + ) + if "restart_type" in builder.epw_final_iso: + from aiida_epw.common import RestartType + + assert builder.epw_final_iso.restart_type == RestartType.EPHREAD + assert ( + builder.epw_final_iso.parameters.get_dict()["INPUTEPW"].get("tc_linear", False) + is False + ) + + # epw_final_aniso check + if "momentum_dependence" in builder.epw_final_aniso: + assert builder.epw_final_aniso.momentum_dependence.value + assert not builder.epw_final_aniso.full_bandwidth.value + assert not builder.epw_final_aniso.real_axis.value + if "calculation_type" in builder.epw_final_aniso: + from aiida_epw.common.types import CalculationTypes + + assert ( + builder.epw_final_aniso.calculation_type.get_member() + == CalculationTypes.ELIASHBERG + ) + if "restart_type" in builder.epw_final_aniso: + from aiida_epw.common import RestartType + assert builder.epw_final_aniso.restart_type == RestartType.EPHREAD From d8154e59f510a45c5b7d6cc962a4ba735296473f Mon Sep 17 00:00:00 2001 From: ymzhang0 Date: Fri, 21 Aug 2026 12:01:55 +0200 Subject: [PATCH 5/9] fix(types): remove restart=True from RestartType.EPHWRITE defaults --- src/aiida_epw/common/types.py | 1 - tests/calculations/test_epw_calcjob.py | 1 - 2 files changed, 2 deletions(-) diff --git a/src/aiida_epw/common/types.py b/src/aiida_epw/common/types.py index 199dbec4..5dc4d65d 100644 --- a/src/aiida_epw/common/types.py +++ b/src/aiida_epw/common/types.py @@ -35,7 +35,6 @@ class RestartType(enum.Enum): "ep_coupling": True, "elph": True, "ephwrite": True, - "restart": True, }, RestartType.FROM_EPH: { "epwread": True, diff --git a/tests/calculations/test_epw_calcjob.py b/tests/calculations/test_epw_calcjob.py index cb93c3f5..a2fc42a9 100644 --- a/tests/calculations/test_epw_calcjob.py +++ b/tests/calculations/test_epw_calcjob.py @@ -657,7 +657,6 @@ def test_epw_eliashberg_parameters_continuation_none( "ep_coupling = .true.", "elph = .true.", "ephwrite = .true.", - "restart = .true.", ), ), ( From c52ec595ac7d3462c946a0dae42e8f2d3b4d5ad3 Mon Sep 17 00:00:00 2001 From: ymzhang0 Date: Fri, 21 Aug 2026 12:58:54 +0200 Subject: [PATCH 6/9] refactor(calcjob): remove fallback restart_type inference and restart.fmt logic in stage_epw_parent --- src/aiida_epw/calculations/epw.py | 27 +------------------------- tests/calculations/test_epw_calcjob.py | 18 ++--------------- 2 files changed, 3 insertions(+), 42 deletions(-) diff --git a/src/aiida_epw/calculations/epw.py b/src/aiida_epw/calculations/epw.py index e8ab14e7..8e759416 100644 --- a/src/aiida_epw/calculations/epw.py +++ b/src/aiida_epw/calculations/epw.py @@ -1021,23 +1021,7 @@ def stage_epw_parent(self, folder, parameters, remote_list, remote_symlink_list) else restart_node ) - if restart_type is None: - inputepw = parameters.get("INPUTEPW", {}) - if inputepw.get("epbread", False): - restart_type = RestartType.FROM_EPB - elif inputepw.get("epwread", False): - if not inputepw.get("elph", True) and not inputepw.get( - "ep_coupling", True - ): - restart_type = RestartType.FROM_EPH - elif inputepw.get("ephwrite", False): - restart_type = RestartType.EPHWRITE - else: - restart_type = RestartType.FROM_EPMATWP - else: - restart_type = RestartType.FROM_SCRATCH - - if restart_type is RestartType.FROM_SCRATCH: + if restart_type is None or restart_type is RestartType.FROM_SCRATCH: return folder.get_subfolder(self._OUTPUT_SUBFOLDER, create=True) @@ -1045,15 +1029,6 @@ def stage_epw_parent(self, folder, parameters, remote_list, remote_symlink_list) parent_folder_epw = self.inputs.parent_folder_epw epw_path = self.get_parent_folder_path(parent_folder_epw) - if parameters.get("INPUTEPW", {}).get("restart", False): - remote_list.append( - ( - parent_folder_epw.computer.uuid, - Path(epw_path, "restart.fmt").as_posix(), - "restart.fmt", - ) - ) - if restart_type is RestartType.FROM_EPB: file_list = [ (f"{self._PREFIX}.epb*", "."), diff --git a/tests/calculations/test_epw_calcjob.py b/tests/calculations/test_epw_calcjob.py index a2fc42a9..0b820d22 100644 --- a/tests/calculations/test_epw_calcjob.py +++ b/tests/calculations/test_epw_calcjob.py @@ -539,21 +539,6 @@ def test_epw_stages_from_eph( ) in copied assert len(calc_info.remote_copy_list) == 2 - # With restart = True - inputs_restart = generate_inputs_epw( - restart_type=RestartType.FROM_EPH, - parameters={"INPUTEPW": {"restart": True}}, - parent_folder_epw=parent_folder, - ) - calc_info_restart = generate_calc_job(fixture_sandbox, "epw.epw", inputs_restart) - copied_restart = [ - (entry[1], entry[2]) for entry in calc_info_restart.remote_copy_list - ] - assert ( - Path(parent_folder.get_remote_path(), "restart.fmt").as_posix(), - "restart.fmt", - ) in copied_restart - def test_epw_stages_ph_stash_folder_by_target_basepath( fixture_sandbox, @@ -717,8 +702,9 @@ def test_epw_parent_folder_without_restart_type( parent_folder_epw=generate_remote_data(fixture_localhost, "/remote/epw"), ) - generate_calc_job(fixture_sandbox, "epw.epw", inputs) + calc_info = generate_calc_job(fixture_sandbox, "epw.epw", inputs) input_contents = Path(fixture_sandbox.abspath, "aiida.in").read_text() assert "epwread = .true." in input_contents assert "elph = .false." in input_contents + assert not any("/remote/epw" in entry[1] for entry in calc_info.remote_copy_list) From 9eb7a283516cffa026af10f5d05fb7f8614c8987 Mon Sep 17 00:00:00 2001 From: ymzhang0 Date: Fri, 21 Aug 2026 14:53:06 +0200 Subject: [PATCH 7/9] feat(supercon): set elph/ep_coupling=False for final calculations and stage additional files for FROM_EPH --- src/aiida_epw/calculations/epw.py | 4 ++++ src/aiida_epw/workflows/protocols/supercon.yaml | 8 ++++---- tests/calculations/test_epw_calcjob.py | 14 +++++++++++++- 3 files changed, 21 insertions(+), 5 deletions(-) diff --git a/src/aiida_epw/calculations/epw.py b/src/aiida_epw/calculations/epw.py index 8e759416..bf86f30c 100644 --- a/src/aiida_epw/calculations/epw.py +++ b/src/aiida_epw/calculations/epw.py @@ -1073,9 +1073,13 @@ def stage_epw_parent(self, folder, parameters, remote_list, remote_symlink_list) if restart_type is RestartType.FROM_EPH: ephmat_rel = Path(self._OUTPUT_SUBFOLDER, self._FOLDER_EPHMAT).as_posix() + dos_rel = Path(self._OUTPUT_SUBFOLDER, self._OUTPUT_DOS_FILE).as_posix() file_list = [ + "crystal.fmt", ephmat_rel, + dos_rel, self._OUTPUT_A2F_FILE, + "selecq.fmt", ] for filename in file_list: remote_list.append( diff --git a/src/aiida_epw/workflows/protocols/supercon.yaml b/src/aiida_epw/workflows/protocols/supercon.yaml index 1f55180c..667bc438 100644 --- a/src/aiida_epw/workflows/protocols/supercon.yaml +++ b/src/aiida_epw/workflows/protocols/supercon.yaml @@ -48,8 +48,8 @@ default_inputs: parameters: INPUTEPW: eliashberg: True - elph: True - ep_coupling: True + elph: False + ep_coupling: False conv_thr_iaxis: 0.01 elecselfen: False etf_mem: 1 @@ -73,8 +73,8 @@ default_inputs: parameters: INPUTEPW: eliashberg: True - elph: True - ep_coupling: True + elph: False + ep_coupling: False conv_thr_iaxis: 0.01 elecselfen: False etf_mem: 1 diff --git a/tests/calculations/test_epw_calcjob.py b/tests/calculations/test_epw_calcjob.py index 0b820d22..c270d60c 100644 --- a/tests/calculations/test_epw_calcjob.py +++ b/tests/calculations/test_epw_calcjob.py @@ -529,15 +529,27 @@ def test_epw_stages_from_eph( calc_info = generate_calc_job(fixture_sandbox, "epw.epw", inputs) copied = [(entry[1], entry[2]) for entry in calc_info.remote_copy_list] + assert ( + Path(parent_folder.get_remote_path(), "crystal.fmt").as_posix(), + "crystal.fmt", + ) in copied assert ( Path(parent_folder.get_remote_path(), "out/aiida.ephmat").as_posix(), "out/aiida.ephmat", ) in copied + assert ( + Path(parent_folder.get_remote_path(), "out/aiida.dos").as_posix(), + "out/aiida.dos", + ) in copied assert ( Path(parent_folder.get_remote_path(), "aiida.a2f").as_posix(), "aiida.a2f", ) in copied - assert len(calc_info.remote_copy_list) == 2 + assert ( + Path(parent_folder.get_remote_path(), "selecq.fmt").as_posix(), + "selecq.fmt", + ) in copied + assert len(calc_info.remote_copy_list) == 5 def test_epw_stages_ph_stash_folder_by_target_basepath( From b018ae09804c17c7ad3b292f2c0775f4b3192fd5 Mon Sep 17 00:00:00 2001 From: ymzhang0 Date: Fri, 21 Aug 2026 16:06:13 +0200 Subject: [PATCH 8/9] feat(protocols): remove nstemp and temps from epw_final_iso and epw_final_aniso in supercon.yaml --- src/aiida_epw/workflows/protocols/supercon.yaml | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/aiida_epw/workflows/protocols/supercon.yaml b/src/aiida_epw/workflows/protocols/supercon.yaml index 667bc438..1880ab45 100644 --- a/src/aiida_epw/workflows/protocols/supercon.yaml +++ b/src/aiida_epw/workflows/protocols/supercon.yaml @@ -58,9 +58,7 @@ default_inputs: muc: 0.13 nqstep: 500 nsiter: 500 - nstemp: 40 selecqread: False - temps: 1 40 vme: 'dipole' wscut: 0.5 wannierize: False @@ -83,9 +81,7 @@ default_inputs: muc: 0.13 nqstep: 500 nsiter: 500 - nstemp: 20 selecqread: False - temps: 5 43 vme: 'dipole' wscut: 0.5 wannierize: False From ac7808f0e3645d3b16ad52005002481dcbd960a3 Mon Sep 17 00:00:00 2001 From: ymzhang0 Date: Sat, 22 Aug 2026 19:55:18 +0200 Subject: [PATCH 9/9] fix(workflows): enforce EnumData for restart_type assignments and use RestartType enum in error recovery --- .../tools/error_handling/supercon.py | 34 +++------ src/aiida_epw/workflows/base.py | 25 ++++--- src/aiida_epw/workflows/prep.py | 6 +- src/aiida_epw/workflows/supercon.py | 7 +- tests/workflows/test_base.py | 69 +++++++++++++++++-- 5 files changed, 94 insertions(+), 47 deletions(-) diff --git a/src/aiida_epw/tools/error_handling/supercon.py b/src/aiida_epw/tools/error_handling/supercon.py index ec910b2c..bdff766c 100644 --- a/src/aiida_epw/tools/error_handling/supercon.py +++ b/src/aiida_epw/tools/error_handling/supercon.py @@ -192,7 +192,7 @@ def prepare_walltime_recovery( """Return ``(parameters, restart_type, action)`` for a recoverable timeout.""" from aiida_epw.common.types import RestartType - if not is_eliashberg: + if not is_eliashberg or restart_type is None: return ( None, None, @@ -202,34 +202,17 @@ def prepare_walltime_recovery( ), ) - if restart_type in ( - "EPHWRITE", - "ephwrite", - "EPHWRITE_RESTART", - "ephwrite_restart", - ): - next_type = ( - "EPHWRITE_RESTART" - if hasattr(RestartType, "EPHWRITE_RESTART") - else "EPHWRITE" - ) + if restart_type is RestartType.EPHWRITE: return ( parameters, - next_type, + RestartType.EPHWRITE, ( "Walltime reached during Eliashberg ephwrite calculation. " "Restarting from the last checkpoint." ), ) - if restart_type in ( - "FROM_EPH", - "from_eph", - "EPHREAD", - "ephread", - "EPHREAD_RESTART", - "ephread_restart", - ): + if restart_type is RestartType.FROM_EPH: outputs = output_parameters or {} eliashberg_data = ( outputs.get("isotropic_eliashberg") @@ -261,12 +244,11 @@ def prepare_walltime_recovery( input_epw_new["nstemp"] = len(remaining_temps) input_epw_new.pop("tempsmin", None) input_epw_new.pop("tempsmax", None) - next_type = "EPHREAD" if hasattr(RestartType, "EPHREAD") else "FROM_EPH" return ( parameters, - next_type, + RestartType.FROM_EPH, ( - f"Walltime reached during Eliashberg {next_type.lower()} calculation. " + "Walltime reached during Eliashberg from_eph calculation. " f"Removed successfully calculated temperatures: {succeeded_temps}. " "Restarting." ), @@ -275,12 +257,12 @@ def prepare_walltime_recovery( None, None, ( - "Walltime reached during Eliashberg ephread calculation but no " + "Walltime reached during Eliashberg from_eph calculation but no " "temperatures finished. Aborting." ), ) - if restart_type in ("FROM_SCRATCH", "from_scratch", "NONE", "none", None): + if restart_type is RestartType.FROM_SCRATCH: return ( None, None, diff --git a/src/aiida_epw/workflows/base.py b/src/aiida_epw/workflows/base.py index 344d9e74..91da6e0e 100644 --- a/src/aiida_epw/workflows/base.py +++ b/src/aiida_epw/workflows/base.py @@ -512,19 +512,22 @@ def handle_cannot_bracket_ef(self, calculation): # A recovered Fermi level must be read from the failed calculation on # the next attempt. Preserve the available coarse-grid data by # changing writer restarts to their corresponding reader modes. + from aiida_epw.calculations.epw import serialize_restart_type + from aiida_epw.common.types import RestartType + restart_type_node = getattr(self.ctx.inputs, "restart_type", None) restart_type = ( restart_type_node.get_member() if restart_type_node is not None else None ) next_restart_type = { - "FROM_SCRATCH": "FROM_EPB", - }.get(getattr(restart_type, "name", None)) + RestartType.FROM_SCRATCH: RestartType.FROM_EPB, + }.get(restart_type) if next_restart_type is not None: - self.ctx.inputs.restart_type = restart_type.__class__[next_restart_type] + self.ctx.inputs.restart_type = serialize_restart_type(next_restart_type) self.ctx.inputs.parent_folder_epw = calculation.outputs.remote_folder action_taken = ( f"{action_taken} Switched restart type from " - f"{restart_type.name} to {next_restart_type}." + f"{restart_type.name} to {next_restart_type.name}." ) self.report_error_handled(calculation, action_taken) @@ -642,9 +645,10 @@ def handle_pade_approximants(self, calculation): True, self.exit_codes.ERROR_KNOWN_UNRECOVERABLE_FAILURE ) + from aiida_epw.calculations.epw import serialize_restart_type from aiida_epw.common.types import RestartType - self.ctx.inputs.restart_type = RestartType.FROM_EPH + self.ctx.inputs.restart_type = serialize_restart_type(RestartType.FROM_EPH) self.ctx.inputs.parameters = orm.Dict(parameters) self.ctx.inputs.parent_folder_epw = calculation.outputs.remote_folder @@ -675,7 +679,7 @@ def handle_temperature_out_of_range(self, calculation): ) def handle_out_of_walltime(self, calculation): """Handle exit code 120 (scheduler walltime timeout) and 400 (software walltime timeout).""" - from aiida_epw.common.types import RestartType + from aiida_epw.calculations.epw import serialize_restart_type parameters = self.ctx.inputs.parameters.get_dict() input_epw = parameters.get("INPUTEPW", {}) @@ -688,16 +692,15 @@ def handle_out_of_walltime(self, calculation): "analytical_continuation", ) ) + restart_type_node = getattr(self.ctx.inputs, "restart_type", None) restart_type = ( - self.ctx.inputs.restart_type.get_member() - if "restart_type" in self.ctx.inputs - else None + restart_type_node.get_member() if restart_type_node is not None else None ) parameters, next_restart, action_taken = ( supercon_error_handling.prepare_walltime_recovery( parameters, calculation.outputs.output_parameters.get_dict(), - restart_type.name if restart_type else None, + restart_type, is_eliashberg, ) ) @@ -707,7 +710,7 @@ def handle_out_of_walltime(self, calculation): True, self.exit_codes.ERROR_KNOWN_UNRECOVERABLE_FAILURE ) self.ctx.inputs.parameters = orm.Dict(parameters) - self.ctx.inputs.restart_type = RestartType[next_restart] + self.ctx.inputs.restart_type = serialize_restart_type(next_restart) self.ctx.inputs.parent_folder_epw = calculation.outputs.remote_folder self.report_error_handled(calculation, action_taken) return ProcessHandlerReport(True) diff --git a/src/aiida_epw/workflows/prep.py b/src/aiida_epw/workflows/prep.py index a093f33d..6253f998 100644 --- a/src/aiida_epw/workflows/prep.py +++ b/src/aiida_epw/workflows/prep.py @@ -1191,9 +1191,10 @@ def run_epw(self): inputs.metadata.call_link_label = "epw_base" + from aiida_epw.calculations.epw import serialize_restart_type from aiida_epw.common.types import RestartType - inputs.restart_type = RestartType.FROM_SCRATCH + inputs.restart_type = serialize_restart_type(RestartType.FROM_SCRATCH) workchain_node = self.submit(EpwBaseWorkChain, **inputs) self.report( @@ -1272,9 +1273,10 @@ def run_epw_bands(self): inputs.parent_folder_epw = self.ctx.workchain_epw.outputs.remote_folder inputs.metadata.call_link_label = "epw_bands" + from aiida_epw.calculations.epw import serialize_restart_type from aiida_epw.common.types import RestartType - inputs.restart_type = RestartType.FROM_EPMATWP + inputs.restart_type = serialize_restart_type(RestartType.FROM_EPMATWP) workchain_node = self.submit(EpwBaseWorkChain, **inputs) self.report( diff --git a/src/aiida_epw/workflows/supercon.py b/src/aiida_epw/workflows/supercon.py index 57d48c6f..a58c2283 100644 --- a/src/aiida_epw/workflows/supercon.py +++ b/src/aiida_epw/workflows/supercon.py @@ -9,6 +9,7 @@ from aiida_quantumespresso.workflows.protocols.utils import ProtocolMixin from aiida_epw.calculations.epw import serialize_restart_type +from aiida_epw.common.types import RestartType from aiida_epw.workflows.base import EpwBaseWorkChain from aiida_epw.data import A2fData @@ -407,7 +408,7 @@ def run_conv(self): inputs.qfpoints_distance = self.ctx.interpolation_list.pop() if "restart_type" not in inputs: - inputs.restart_type = serialize_restart_type("ephwrite") + inputs.restart_type = serialize_restart_type(RestartType.EPHWRITE) if self.ctx.degaussq: parameters = inputs.parameters.get_dict() @@ -475,7 +476,7 @@ def run_final_epw_iso(self): inputs.qfpoints = parent_folder_epw.creator.inputs.qfpoints if "restart_type" not in inputs: - inputs.restart_type = serialize_restart_type("from_eph") + inputs.restart_type = serialize_restart_type(RestartType.FROM_EPH) if self.ctx.degaussq: parameters = inputs.parameters.get_dict() @@ -514,7 +515,7 @@ def run_final_epw_aniso(self): inputs.qfpoints = parent_folder_epw.creator.inputs.qfpoints if "restart_type" not in inputs: - inputs.restart_type = serialize_restart_type("from_eph") + inputs.restart_type = serialize_restart_type(RestartType.FROM_EPH) inputs.metadata.call_link_label = "epw_final_aniso" workchain_node = self.submit(EpwBaseWorkChain, **inputs) diff --git a/tests/workflows/test_base.py b/tests/workflows/test_base.py index 0ac099d0..1379b527 100644 --- a/tests/workflows/test_base.py +++ b/tests/workflows/test_base.py @@ -5,7 +5,9 @@ import pytest from aiida import orm +from aiida.common import AttributeDict +from aiida_epw.calculations.epw import serialize_restart_type from aiida_epw.common.types import RestartType from aiida_epw.workflows.base import EpwBaseWorkChain @@ -553,7 +555,10 @@ def report_error_handled(self, calculation, action): report = workchain.handle_cannot_bracket_ef.__wrapped__(calculation) assert report.exit_code.status == 0 - assert workchain.ctx.inputs.restart_type == RestartType[expected_restart_type] + assert ( + workchain.ctx.inputs.restart_type.get_member() + == RestartType[expected_restart_type] + ) assert workchain.ctx.inputs.parent_folder_epw is calculation.outputs.remote_folder assert expected_restart_type in workchain.report_messages[-1] @@ -638,10 +643,63 @@ def report_error_handled(self, calculation, action): assert report.do_break is True assert report.exit_code.status == 0 - assert workchain.ctx.inputs.restart_type == RestartType.EPHWRITE_RESTART + assert workchain.ctx.inputs.restart_type.get_member() == RestartType.EPHWRITE assert workchain.ctx.inputs.parent_folder_epw == calc.outputs.remote_folder +def test_handle_out_of_walltime_consecutive_timeouts(aiida_localhost): + """Test that multiple walltime timeouts in the same workchain preserve EnumData input type.""" + + class MockWorkChain: + exit_codes = EpwBaseWorkChain.exit_codes + + def __init__(self): + self.ctx = type("Context", (), {})() + self.ctx.inputs = AttributeDict() + self.report_messages = [] + + def report(self, msg): + self.report_messages.append(msg) + + def report_error_handled(self, calculation, action): + self.report(f"Calculation failed: {action}") + + handle_out_of_walltime = EpwBaseWorkChain.handle_out_of_walltime + + workchain = MockWorkChain() + workchain.ctx.inputs.parameters = orm.Dict(dict={"INPUTEPW": {"eliashberg": True}}) + workchain.ctx.inputs.restart_type = serialize_restart_type(RestartType.EPHWRITE) + + calc1 = MagicMock() + calc1.outputs = MagicMock() + calc1.outputs.remote_folder = orm.RemoteData( + computer=aiida_localhost, remote_path="/tmp/1" + ) + calc1.outputs.output_parameters = orm.Dict(dict={}) + + # First timeout + report1 = workchain.handle_out_of_walltime.__wrapped__(calc1) + assert report1.do_break is True + assert report1.exit_code.status == 0 + assert isinstance(workchain.ctx.inputs.restart_type, orm.EnumData) + assert workchain.ctx.inputs.restart_type.get_member() == RestartType.EPHWRITE + + # Second consecutive timeout (simulates iteration #3 in crash report) + calc2 = MagicMock() + calc2.outputs = MagicMock() + calc2.outputs.remote_folder = orm.RemoteData( + computer=aiida_localhost, remote_path="/tmp/2" + ) + calc2.outputs.output_parameters = orm.Dict(dict={}) + + report2 = workchain.handle_out_of_walltime.__wrapped__(calc2) + assert report2.do_break is True + assert report2.exit_code.status == 0 + assert isinstance(workchain.ctx.inputs.restart_type, orm.EnumData) + assert workchain.ctx.inputs.restart_type.get_member() == RestartType.EPHWRITE + assert workchain.ctx.inputs.parent_folder_epw == calc2.outputs.remote_folder + + def test_handle_out_of_walltime_unsupported(aiida_localhost): """Test that handle_out_of_walltime aborts for unsupported calculation/restart configurations.""" @@ -736,7 +794,7 @@ def report_error_handled(self, calculation, action): workchain.ctx.inputs = MagicMock() workchain.ctx.inputs.parameters = orm.Dict(dict=initial_params) restart_type_mock = MagicMock() - restart_type_mock.get_member.return_value = RestartType.EPHREAD + restart_type_mock.get_member.return_value = RestartType.FROM_EPH workchain.ctx.inputs.restart_type = restart_type_mock workchain.ctx.inputs.__contains__.side_effect = lambda key: key in ( "restart_type", @@ -748,6 +806,7 @@ def report_error_handled(self, calculation, action): assert report.do_break is True assert report.exit_code.status == 0 assert workchain.ctx.inputs.parent_folder_epw == calc.outputs.remote_folder + assert workchain.ctx.inputs.restart_type.get_member() == RestartType.FROM_EPH updated_params = workchain.ctx.inputs.parameters.get_dict() assert updated_params["INPUTEPW"]["temps"] == [1.0] @@ -802,7 +861,7 @@ def report_error_handled(self, calculation, action): workchain.ctx.inputs = MagicMock() workchain.ctx.inputs.parameters = orm.Dict(dict=initial_params) restart_type_mock = MagicMock() - restart_type_mock.get_member.return_value = RestartType.EPHREAD + restart_type_mock.get_member.return_value = RestartType.FROM_EPH workchain.ctx.inputs.restart_type = restart_type_mock workchain.ctx.inputs.__contains__.side_effect = lambda key: key in ( "restart_type", @@ -848,7 +907,7 @@ def report_error_handled(self, calculation, action): workchain.ctx.inputs = MagicMock() workchain.ctx.inputs.parameters = orm.Dict(dict=initial_params) restart_type_mock = MagicMock() - restart_type_mock.get_member.return_value = RestartType.NONE + restart_type_mock.get_member.return_value = RestartType.FROM_SCRATCH workchain.ctx.inputs.restart_type = restart_type_mock workchain.ctx.inputs.__contains__.side_effect = lambda key: key in ( "restart_type",