From 7f92cbc803aff96b75c3c28fd68806d686ca788f Mon Sep 17 00:00:00 2001 From: ymzhang0 Date: Wed, 24 Jun 2026 12:27:47 +0200 Subject: [PATCH 1/6] feat(calculation-types): add CalculationTypes enum and integrate with EpwCalculation --- src/aiida_epw/calculations/epw.py | 66 ++++++++++++++++++++++++++ src/aiida_epw/common/__init__.py | 6 +++ src/aiida_epw/common/types.py | 12 +++++ tests/calculations/test_epw_calcjob.py | 58 ++++++++++++++++++++++ 4 files changed, 142 insertions(+) create mode 100644 src/aiida_epw/common/__init__.py create mode 100644 src/aiida_epw/common/types.py diff --git a/src/aiida_epw/calculations/epw.py b/src/aiida_epw/calculations/epw.py index 3573755..1d79fcc 100644 --- a/src/aiida_epw/calculations/epw.py +++ b/src/aiida_epw/calculations/epw.py @@ -35,6 +35,26 @@ def _lowercase_dict(dictionary, dict_name): return _case_transform_dict(dictionary, dict_name, "_lowercase_dict", str.lower) +def serialize_calculation_type(value): + """Serialize input parameter into an AiiDA EnumData for CalculationTypes.""" + from aiida.orm import EnumData + from aiida_epw.common.types import CalculationTypes + + if isinstance(value, EnumData): + return value + if isinstance(value, CalculationTypes): + return EnumData(value) + if isinstance(value, str): + try: + return EnumData(CalculationTypes(value.lower())) + except ValueError as exc: + raise ValueError( + f"Invalid calculation type '{value}'. Supported values: " + f"{[member.value for member in CalculationTypes]}" + ) from exc + raise TypeError(f"Cannot serialize {value} to EnumData of CalculationTypes") + + class EpwCalculation(NamelistsCalculation): """`CalcJob` implementation for the epw.x code of Quantum ESPRESSO.""" @@ -55,6 +75,9 @@ class EpwCalculation(NamelistsCalculation): ("INPUTEPW", "nkf1"), ("INPUTEPW", "nkf2"), ("INPUTEPW", "nkf3"), + ("INPUTEPW", "eliashberg"), + ("INPUTEPW", "scattering"), + ("INPUTEPW", "plrn"), ] _use_kpoints = True @@ -102,6 +125,13 @@ def define(cls, spec): valid_type=orm.Dict, help="Parameters for the `epw.x` input file.", ) + spec.input( + "calculation_type", + valid_type=orm.EnumData, + required=False, + serializer=serialize_calculation_type, + help="EPW calculation type: Eliashberg, transport, or polaron.", + ) spec.input( "kpoints", valid_type=orm.KpointsData, @@ -411,6 +441,24 @@ def validate_parameters_inputs(cls, parameters, inputs): "`parameters.INPUTEPW.wannierize` is true." ) + calculation_type = inputs.get("calculation_type", None) + if calculation_type is not None: + calc_type = calculation_type.get_member() + from aiida_epw.common.types import CalculationTypes + + if calc_type != CalculationTypes.ELIASHBERG: + for f in ( + "momentum_dependence", + "full_bandwidth", + "real_axis", + "analytical_continuation", + ): + if f in inputs: + raise exceptions.InputValidationError( + f"Eliashberg parameter '{f}' cannot be specified when " + f"calculation_type is '{calc_type.value}'." + ) + @classmethod def set_blocked_keywords(cls, parameters): """Validate plugin-managed keywords without mutating the parameter dictionary.""" @@ -573,6 +621,24 @@ def prepare_input_parameters(self, folder, parameters): self.cap_nstemp(inputepw_parameters) + # Override calculation type settings in parameters if calculation_type is specified + if "calculation_type" in self.inputs: + calc_type = self.inputs.calculation_type.get_member() + from aiida_epw.common.types import CalculationTypes + + if calc_type == CalculationTypes.ELIASHBERG: + inputepw_parameters["eliashberg"] = True + inputepw_parameters["scattering"] = False + inputepw_parameters["plrn"] = False + elif calc_type == CalculationTypes.TRANSPORT: + inputepw_parameters["eliashberg"] = False + inputepw_parameters["scattering"] = True + inputepw_parameters["plrn"] = False + elif calc_type == CalculationTypes.POLARON: + inputepw_parameters["eliashberg"] = False + inputepw_parameters["scattering"] = False + inputepw_parameters["plrn"] = True + inputepw_parameters["outdir"] = self._OUTPUT_SUBFOLDER inputepw_parameters["dvscf_dir"] = self._FOLDER_SAVE inputepw_parameters["prefix"] = self._PREFIX diff --git a/src/aiida_epw/common/__init__.py b/src/aiida_epw/common/__init__.py new file mode 100644 index 0000000..bf6d178 --- /dev/null +++ b/src/aiida_epw/common/__init__.py @@ -0,0 +1,6 @@ +# -*- coding: utf-8 -*- +"""Common utilities and types for EPW plugin.""" + +from .types import CalculationTypes + +__all__ = ("CalculationTypes",) diff --git a/src/aiida_epw/common/types.py b/src/aiida_epw/common/types.py new file mode 100644 index 0000000..a01ecef --- /dev/null +++ b/src/aiida_epw/common/types.py @@ -0,0 +1,12 @@ +# -*- coding: utf-8 -*- +"""Data types for EPW plugin.""" + +import enum + + +class CalculationTypes(enum.Enum): + """Enumeration of EPW calculation types.""" + + ELIASHBERG = "eliashberg" + TRANSPORT = "transport" + POLARON = "polaron" diff --git a/tests/calculations/test_epw_calcjob.py b/tests/calculations/test_epw_calcjob.py index fce3111..f3234f4 100644 --- a/tests/calculations/test_epw_calcjob.py +++ b/tests/calculations/test_epw_calcjob.py @@ -497,3 +497,61 @@ def test_epw_stages_ph_stash_folder_by_target_basepath( ).as_posix(), "save", ) in calc_info.remote_copy_list + + +def test_epw_calculation_type_parameter( + fixture_sandbox, generate_calc_job, generate_inputs_epw +): + """Test that calculation_type correctly overrides the namelist parameters.""" + from aiida_epw.common.types import CalculationTypes + from aiida_epw.calculations.epw import EpwCalculation + import pytest + + # 1. Test Eliashberg mode + inputs = generate_inputs_epw( + calculation_type=orm.EnumData(CalculationTypes.ELIASHBERG), + ) + generate_calc_job(fixture_sandbox, "epw.epw", inputs) + input_contents = Path(fixture_sandbox.abspath, "aiida.in").read_text() + assert "eliashberg = .true." in input_contents + assert "scattering = .false." in input_contents + assert "plrn = .false." in input_contents + + # 2. Test Transport mode + inputs = generate_inputs_epw( + calculation_type=orm.EnumData(CalculationTypes.TRANSPORT), + ) + generate_calc_job(fixture_sandbox, "epw.epw", inputs) + input_contents = Path(fixture_sandbox.abspath, "aiida.in").read_text() + assert "eliashberg = .false." in input_contents + assert "scattering = .true." in input_contents + assert "plrn = .false." in input_contents + + # 3. Test Polaron mode + inputs = generate_inputs_epw( + calculation_type=orm.EnumData(CalculationTypes.POLARON), + ) + generate_calc_job(fixture_sandbox, "epw.epw", inputs) + input_contents = Path(fixture_sandbox.abspath, "aiida.in").read_text() + assert "eliashberg = .false." in input_contents + assert "scattering = .false." in input_contents + assert "plrn = .true." in input_contents + + # 4. Test validation of Eliashberg inputs in non-Eliashberg modes + if "real_axis" in EpwCalculation.spec().inputs: + with pytest.raises( + ValueError, + match="Eliashberg parameter 'real_axis' cannot be specified", + ): + inputs_invalid = generate_inputs_epw( + calculation_type=orm.EnumData(CalculationTypes.TRANSPORT), + real_axis=orm.Bool(True), + ) + generate_calc_job(fixture_sandbox, "epw.epw", inputs_invalid) + + # 5. Test blocked parameter validation + with pytest.raises(ValueError, match="parameters.INPUTEPW.scattering"): + inputs_blocked = generate_inputs_epw( + parameters={"INPUTEPW": {"scattering": True}} + ) + generate_calc_job(fixture_sandbox, "epw.epw", inputs_blocked) From 8f2042110573ed89e04819c89f1fd457c933a7b5 Mon Sep 17 00:00:00 2001 From: ymzhang0 Date: Tue, 7 Jul 2026 16:34:21 +0200 Subject: [PATCH 2/6] feat(types): move WANNIERIZE to CalculationTypes enum and update epw.py validation/overrides --- src/aiida_epw/calculations/epw.py | 44 +++++++++++++++++++++++++------ src/aiida_epw/common/types.py | 1 + 2 files changed, 37 insertions(+), 8 deletions(-) diff --git a/src/aiida_epw/calculations/epw.py b/src/aiida_epw/calculations/epw.py index 1d79fcc..9fb67e7 100644 --- a/src/aiida_epw/calculations/epw.py +++ b/src/aiida_epw/calculations/epw.py @@ -379,20 +379,29 @@ def validate_restart_inputs(cls, parameters, inputs): """Validate restart-related input combinations against the EPW parameters.""" inputepw = parameters["INPUTEPW"] - if not inputepw.get("wannierize", False): + is_wannierize = False + calculation_type = inputs.get("calculation_type", None) + if calculation_type is not None: + calc_type = calculation_type.get_member() + from aiida_epw.common.types import CalculationTypes + + if calc_type is CalculationTypes.WANNIERIZE: + is_wannierize = True + if not is_wannierize: + is_wannierize = inputepw.get("wannierize", False) + + if not is_wannierize: return for input_name in ("parent_folder_epw", "parent_folder_chk"): if input_name in inputs: raise exceptions.InputValidationError( - f"`{input_name}` cannot be specified when " - "`parameters.INPUTEPW.wannierize` is true." + f"`{input_name}` cannot be specified when wannierize is enabled." ) if "parent_folder_nscf" not in inputs: raise exceptions.InputValidationError( - "`parent_folder_nscf` must be specified when " - "`parameters.INPUTEPW.wannierize` is true." + "`parent_folder_nscf` must be specified when wannierize is enabled." ) @staticmethod @@ -416,7 +425,18 @@ def validate_parameters_inputs(cls, parameters, inputs): cls.validate_restart_inputs(parameters, inputs) inputepw = parameters["INPUTEPW"] - if inputepw.get("wannierize", False): + is_wannierize = False + calculation_type = inputs.get("calculation_type", None) + if calculation_type is not None: + calc_type = calculation_type.get_member() + from aiida_epw.common.types import CalculationTypes + + if calc_type is CalculationTypes.WANNIERIZE: + is_wannierize = True + if not is_wannierize: + is_wannierize = inputepw.get("wannierize", False) + + if is_wannierize: if inputepw.get("auto_projections", False): raise exceptions.InputValidationError( "`parameters.INPUTEPW.auto_projections` is not supported; " @@ -437,8 +457,7 @@ def validate_parameters_inputs(cls, parameters, inputs): if not cls.has_manual_projections(inputepw): raise exceptions.InputValidationError( - "Manual `proj` entries must be provided when " - "`parameters.INPUTEPW.wannierize` is true." + "Manual `proj` entries must be provided when wannierize is enabled." ) calculation_type = inputs.get("calculation_type", None) @@ -638,6 +657,15 @@ def prepare_input_parameters(self, folder, parameters): inputepw_parameters["eliashberg"] = False inputepw_parameters["scattering"] = False inputepw_parameters["plrn"] = True + elif calc_type == CalculationTypes.WANNIERIZE: + inputepw_parameters["wannierize"] = True + inputepw_parameters["epwread"] = False + inputepw_parameters["epwwrite"] = True + inputepw_parameters.setdefault("restart", False) + inputepw_parameters["ep_coupling"] = True + inputepw_parameters["elph"] = True + inputepw_parameters["epbwrite"] = True + inputepw_parameters["epbread"] = False inputepw_parameters["outdir"] = self._OUTPUT_SUBFOLDER inputepw_parameters["dvscf_dir"] = self._FOLDER_SAVE diff --git a/src/aiida_epw/common/types.py b/src/aiida_epw/common/types.py index a01ecef..13f49e1 100644 --- a/src/aiida_epw/common/types.py +++ b/src/aiida_epw/common/types.py @@ -7,6 +7,7 @@ class CalculationTypes(enum.Enum): """Enumeration of EPW calculation types.""" + WANNIERIZE = "wannierize" ELIASHBERG = "eliashberg" TRANSPORT = "transport" POLARON = "polaron" From 7cbeb7ea3eda2471d0ebe4c090fb73ec007fc331 Mon Sep 17 00:00:00 2001 From: ymzhang0 Date: Tue, 7 Jul 2026 17:08:53 +0200 Subject: [PATCH 3/6] feat(types): remove wannierize parameter from WANNIERIZE calc_type --- src/aiida_epw/calculations/epw.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/aiida_epw/calculations/epw.py b/src/aiida_epw/calculations/epw.py index 9fb67e7..bcec90b 100644 --- a/src/aiida_epw/calculations/epw.py +++ b/src/aiida_epw/calculations/epw.py @@ -658,7 +658,6 @@ def prepare_input_parameters(self, folder, parameters): inputepw_parameters["scattering"] = False inputepw_parameters["plrn"] = True elif calc_type == CalculationTypes.WANNIERIZE: - inputepw_parameters["wannierize"] = True inputepw_parameters["epwread"] = False inputepw_parameters["epwwrite"] = True inputepw_parameters.setdefault("restart", False) From a8b9e17318989fb2ac9eb2a8de3222218a819cff Mon Sep 17 00:00:00 2001 From: ymzhang0 Date: Wed, 8 Jul 2026 10:58:32 +0200 Subject: [PATCH 4/6] feat(types): add BANDS calculation type and support in EpwCalculation --- src/aiida_epw/calculations/epw.py | 5 +++++ src/aiida_epw/common/types.py | 1 + 2 files changed, 6 insertions(+) diff --git a/src/aiida_epw/calculations/epw.py b/src/aiida_epw/calculations/epw.py index bcec90b..4c31cf1 100644 --- a/src/aiida_epw/calculations/epw.py +++ b/src/aiida_epw/calculations/epw.py @@ -665,6 +665,11 @@ def prepare_input_parameters(self, folder, parameters): inputepw_parameters["elph"] = True inputepw_parameters["epbwrite"] = True inputepw_parameters["epbread"] = False + elif calc_type == CalculationTypes.BANDS: + inputepw_parameters["band_plot"] = True + inputepw_parameters["eliashberg"] = False + inputepw_parameters["scattering"] = False + inputepw_parameters["plrn"] = False inputepw_parameters["outdir"] = self._OUTPUT_SUBFOLDER inputepw_parameters["dvscf_dir"] = self._FOLDER_SAVE diff --git a/src/aiida_epw/common/types.py b/src/aiida_epw/common/types.py index 13f49e1..95d9cce 100644 --- a/src/aiida_epw/common/types.py +++ b/src/aiida_epw/common/types.py @@ -11,3 +11,4 @@ class CalculationTypes(enum.Enum): ELIASHBERG = "eliashberg" TRANSPORT = "transport" POLARON = "polaron" + BANDS = "bands" From 48a34ac109feb8fb1e7471c642c6ad5e681a2f9f Mon Sep 17 00:00:00 2001 From: ymzhang0 Date: Thu, 9 Jul 2026 11:17:42 +0200 Subject: [PATCH 5/6] feat(prep): exclude calculation_type from prep workchain inputs and assign programmatically --- src/aiida_epw/workflows/prep.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/src/aiida_epw/workflows/prep.py b/src/aiida_epw/workflows/prep.py index 2416588..082fb0b 100644 --- a/src/aiida_epw/workflows/prep.py +++ b/src/aiida_epw/workflows/prep.py @@ -276,6 +276,7 @@ def define(cls, spec): "parent_folder_nscf", "parent_folder_epw", "parent_folder_chk", + "calculation_type", ), namespace_options={"help": "Inputs for the `EpwBaseWorkChain`."}, ) @@ -292,6 +293,7 @@ def define(cls, spec): "qfpoints_distance", "kfpoints_factor", "parent_folder_epw", + "calculation_type", ), namespace_options={ "help": "Inputs namespace for `EpwBaseWorkChain` that runs the `epw.x` calculation in interpolation mode, i.e. the interpolated electron and phonon band structures." @@ -1189,6 +1191,10 @@ def run_epw(self): inputs.metadata.call_link_label = "epw_base" + from aiida_epw.common.types import CalculationTypes + + inputs.calculation_type = CalculationTypes.WANNIERIZE + workchain_node = self.submit(EpwBaseWorkChain, **inputs) self.report( f"launching EpwBaseWorkChain<{workchain_node.pk}> in transformation mode" @@ -1265,6 +1271,11 @@ def run_epw_bands(self): inputs.kfpoints = bands_kpoints inputs.parent_folder_epw = self.ctx.workchain_epw.outputs.remote_folder inputs.metadata.call_link_label = "epw_bands" + + from aiida_epw.common.types import CalculationTypes + + inputs.calculation_type = CalculationTypes.BANDS + workchain_node = self.submit(EpwBaseWorkChain, **inputs) self.report( f"launching EpwBaseWorkChain<{workchain_node.pk}> in bands interpolation mode" From f61229fbd79a9c867d763277bfb72b9ad4cf892a Mon Sep 17 00:00:00 2001 From: ymzhang0 Date: Thu, 9 Jul 2026 11:22:44 +0200 Subject: [PATCH 6/6] feat(types): block band_plot keyword in EpwCalculation --- src/aiida_epw/calculations/epw.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/aiida_epw/calculations/epw.py b/src/aiida_epw/calculations/epw.py index 4c31cf1..af506ad 100644 --- a/src/aiida_epw/calculations/epw.py +++ b/src/aiida_epw/calculations/epw.py @@ -78,6 +78,7 @@ class EpwCalculation(NamelistsCalculation): ("INPUTEPW", "eliashberg"), ("INPUTEPW", "scattering"), ("INPUTEPW", "plrn"), + ("INPUTEPW", "band_plot"), ] _use_kpoints = True