diff --git a/src/aiida_epw/calculations/epw.py b/src/aiida_epw/calculations/epw.py index 3573755..af506ad 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,10 @@ class EpwCalculation(NamelistsCalculation): ("INPUTEPW", "nkf1"), ("INPUTEPW", "nkf2"), ("INPUTEPW", "nkf3"), + ("INPUTEPW", "eliashberg"), + ("INPUTEPW", "scattering"), + ("INPUTEPW", "plrn"), + ("INPUTEPW", "band_plot"), ] _use_kpoints = True @@ -102,6 +126,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, @@ -349,20 +380,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 @@ -386,7 +426,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; " @@ -407,10 +458,27 @@ 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) + 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 +641,37 @@ 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 + elif calc_type == CalculationTypes.WANNIERIZE: + 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 + 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 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..95d9cce --- /dev/null +++ b/src/aiida_epw/common/types.py @@ -0,0 +1,14 @@ +# -*- coding: utf-8 -*- +"""Data types for EPW plugin.""" + +import enum + + +class CalculationTypes(enum.Enum): + """Enumeration of EPW calculation types.""" + + WANNIERIZE = "wannierize" + ELIASHBERG = "eliashberg" + TRANSPORT = "transport" + POLARON = "polaron" + BANDS = "bands" 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" 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)