diff --git a/src/aiida_epw/workgraphs/__init__.py b/src/aiida_epw/workgraphs/__init__.py new file mode 100644 index 00000000..c813636a --- /dev/null +++ b/src/aiida_epw/workgraphs/__init__.py @@ -0,0 +1,11 @@ +"""WorkGraph builders for aiida-epw.""" + +from .prep import build_task_inputs, prep, prep_from_inputs +from .supercon import supercon + +__all__ = ( + "build_task_inputs", + "prep", + "prep_from_inputs", + "supercon", +) diff --git a/src/aiida_epw/workgraphs/prep.py b/src/aiida_epw/workgraphs/prep.py new file mode 100644 index 00000000..62cc4b85 --- /dev/null +++ b/src/aiida_epw/workgraphs/prep.py @@ -0,0 +1,1011 @@ +"""WorkGraph implementation of the EPW preparation workflow.""" + +from __future__ import annotations + +from typing import Any + +from aiida import orm +from aiida_workgraph import If, WorkGraph, spec, task +from aiida_quantumespresso.workflows.protocols.utils import recursive_merge +from aiida_quantumespresso.workflows.ph.base import PhBaseWorkChain +from aiida_quantumespresso.workflows.pw.base import PwBaseWorkChain +from aiida_workgraph.utils import get_dict_from_builder +from aiida_wannier90_workflows.workflows import ( + Wannier90BandsWorkChain, + Wannier90OptimizeWorkChain, +) +from aiida_wannier90_workflows.common.types import WannierProjectionType + +from aiida_epw.tools.workchain import ( + get_target_basepath, + validate_parent_ph_inputs, +) +from aiida_epw.workflows.base import EpwBaseWorkChain + +PwBaseTask = task(PwBaseWorkChain) +Wannier90OptimizeTask = task(Wannier90OptimizeWorkChain) +Wannier90BandsTask = task(Wannier90BandsWorkChain) +PhBaseTask = task(PhBaseWorkChain) +EpwBaseTask = task(EpwBaseWorkChain) + +__all__ = ( + "prep", + "prep_from_inputs", + "build_task_inputs", +) + + +def _copy_nested_containers(value: Any) -> Any: + """Recursively copy Python containers while preserving AiiDA ORM nodes.""" + if isinstance(value, dict): + return {key: _copy_nested_containers(item) for key, item in value.items()} + if isinstance(value, list): + return [_copy_nested_containers(item) for item in value] + if isinstance(value, tuple): + return tuple(_copy_nested_containers(item) for item in value) + return value + + +def _normalize_nested_python_values(value: Any) -> Any: + """Normalize nested Python containers to avoid non-JSON-serializable values.""" + if isinstance(value, orm.Dict): + return orm.Dict(dict=_normalize_nested_python_values(value.get_dict())) + if isinstance(value, dict): + return { + key: _normalize_nested_python_values(item) for key, item in value.items() + } + if isinstance(value, list): + return [_normalize_nested_python_values(item) for item in value] + if isinstance(value, tuple): + return tuple(_normalize_nested_python_values(item) for item in value) + if isinstance(value, range): + return list(value) + return value + + +def _get_nested_value(mapping: dict[str, Any], path: str) -> Any: + """Return a nested value from a mapping or ``None`` if the path is missing.""" + current: Any = mapping + for key in path.split("."): + if not isinstance(current, dict) or key not in current: + return None + current = current[key] + return current + + +def _set_nested_value(mapping: dict[str, Any], path: str, value: Any) -> None: + """Set a nested value inside a mapping, creating intermediate dictionaries.""" + keys = path.split(".") + current = mapping + for key in keys[:-1]: + current = current.setdefault(key, {}) + current[keys[-1]] = value + + +def _restore_nested_fields( + target: dict[str, Any], + source: dict[str, Any], + paths: tuple[str, ...], +) -> dict[str, Any]: + """Restore selected nested fields from a source mapping into a target mapping.""" + restored = _copy_nested_containers(target) + for path in paths: + value = _get_nested_value(source, path) + if value is not None: + _set_nested_value(restored, path, _copy_nested_containers(value)) + return restored + + +def _set_socket_value(namespace: Any, path: str, value: Any) -> None: + """Set a nested socket value using a dotted path.""" + current = namespace + keys = path.split(".") + for key in keys[:-1]: + current = current[key] + current[keys[-1]] = value + + +def _apply_socket_overrides( + namespace: Any, + source: dict[str, Any], + paths: tuple[str, ...], +) -> None: + """Apply selected values from a source mapping onto a socket namespace.""" + for path in paths: + value = _get_nested_value(source, path) + if value is not None: + _set_socket_value(namespace, path, value) + + +def validate_inputs( # pylint: disable=unused-argument,inconsistent-return-statements + inputs, ctx=None +): + """Validate the inputs of the `EpwPrepWorkChain`.""" + has_w90_bands = "w90_bands" in inputs + use_epw_wannierize = should_epw_wannierize(inputs) + + if has_w90_bands and use_epw_wannierize: + return ( + "`w90_bands` inputs and `epw_base.parameters.INPUTEPW.wannierize = True` " + "are mutually exclusive." + ) + + if not has_w90_bands and not use_epw_wannierize: + return ( + "Either provide `w90_bands` inputs or set " + "`epw_base.parameters.INPUTEPW.wannierize = True`." + ) + + if use_epw_wannierize: + missing = [ + namespace for namespace in ("scf", "nscf") if namespace not in inputs + ] + if missing: + return ( + "`scf` and `nscf` inputs are required when " + "`epw_base.parameters.INPUTEPW.wannierize = True`." + ) + + if "parent_folder_ph" not in inputs and "ph_base" not in inputs: + return "Either provide `ph_base` inputs or set `parent_folder_ph`." + + +def should_epw_wannierize(inputs) -> bool: + """Return whether the EPW namespace is configured to run Wannierization directly.""" + epw_base = _as_dict(inputs.get("epw_base")) + parameters = _as_dict(epw_base.get("parameters")) + inputepw = _as_dict(parameters.get("INPUTEPW", parameters.get("inputepw"))) + return bool(inputepw.get("wannierize", False)) + + +def _as_bool(value: Any) -> bool: + """Convert Python and AiiDA booleans into a plain boolean.""" + if isinstance(value, orm.Bool): + return value.value + return bool(value) + + +def _as_dict(value: Any) -> dict[str, Any]: + """Convert AiiDA and plain mapping inputs into a plain dictionary.""" + if value is None: + return {} + if isinstance(value, orm.Dict): + return value.get_dict() + if isinstance(value, dict): + return value + return dict(value) + + +def _create_kpoints_from_distance_node( + structure: orm.StructureData, + distance: orm.Float, + force_parity: orm.Bool, +) -> orm.KpointsData: + """Create a k-point mesh with the same symmetry handling as the QE helper.""" + from numpy import linalg + + epsilon = 1e-5 + + kpoints = orm.KpointsData() + kpoints.set_cell_from_structure(structure) + kpoints.set_kpoints_mesh_from_density( + distance.value, + force_parity=force_parity.value, + ) + + lengths_vector = [linalg.norm(vector) for vector in structure.cell] + lengths_kpoint = kpoints.get_kpoints_mesh()[0] + + is_symmetric_cell = all( + abs(length - lengths_vector[0]) < epsilon for length in lengths_vector + ) + is_symmetric_mesh = all(length == lengths_kpoint[0] for length in lengths_kpoint) + + if is_symmetric_cell and not is_symmetric_mesh: + nkpoints = max(lengths_kpoint) + kpoints.set_kpoints_mesh([nkpoints if pbc else 1 for pbc in structure.pbc]) + + return kpoints + + +def get_protocol_inputs( + protocol: str | None = None, + overrides: dict | None = None, +) -> dict: + """Return the inputs for the EPW preparation workflow based on a protocol.""" + from importlib_resources import files + from aiida_epw.workflows import protocols + from aiida_quantumespresso.workflows.protocols.utils import ProtocolMixin + + # 1. Load protocol file from workflows side + filepath = files(protocols) / "prep.yaml" + + AdHocProtocol = type( + "AdHocProtocol", + (ProtocolMixin,), + { + "get_protocol_filepath": classmethod(lambda cls: filepath), + "_validate_override_keys": classmethod(lambda cls, overrides: None), + }, + ) + + return AdHocProtocol.get_protocol_inputs(protocol, overrides) + + +def _add_metadata_stash_target_base(inputs: dict[str, Any], code) -> None: + """Populate a ``metadata.options.stash`` target base when missing.""" + if code is None: + return + + stash = ( + inputs.setdefault("metadata", {}) + .setdefault("options", {}) + .setdefault("stash", {}) + ) + if "target_base" not in stash: + stash["target_base"] = get_target_basepath(code.computer) + stash["stash_mode"] = stash.get("stash_mode", "copy") + + +def _add_options_stash_target_base(inputs: dict[str, Any], code) -> None: + """Populate an ``options.stash`` target base when missing.""" + if code is None: + return + + stash = inputs.setdefault("options", {}).setdefault("stash", {}) + if "target_base" not in stash: + stash["target_base"] = get_target_basepath(code.computer) + stash["stash_mode"] = stash.get("stash_mode", "copy") + + +def _set_call_link_label(inputs: dict[str, Any], label: str) -> dict[str, Any]: + """Return inputs with a top-level metadata call link label.""" + updated = _copy_nested_containers(inputs) + metadata = updated.setdefault("metadata", {}) + metadata["call_link_label"] = label + return updated + + +def _validate_reference_bands_projection_type( + reference_bands: orm.BandsData | None, + wannier_projection_type: WannierProjectionType, +) -> None: + """Reject unsupported optimization modes for the selected projection type.""" + if ( + reference_bands is not None + and wannier_projection_type == WannierProjectionType.ANALYTIC + ): + raise ValueError( + "`reference_bands` with `WannierProjectionType.ANALYTIC` is not " + "supported in `EpwPrepWorkChain`: the optimize branch uses " + "`Wannier90OptimizeWorkChain`, which optimizes `dis_proj_min/max`. " + "Use `Wannier90BandsWorkChain` without `reference_bands` and tune " + "`dis_win_*`/`dis_froz_*` manually." + ) + + +def _validate_parent_folder_ph(parent_folder_ph, structure) -> None: + """Validate that a provided phonon parent folder can be reused safely.""" + if parent_folder_ph is None: + return + + validate_parent_ph_inputs(parent_folder_ph, structure) + + +def _build_wannier90_inputs( + *, + codes: dict[str, Any], + structure: orm.StructureData, + protocol_inputs: dict[str, Any], + pseudo_family: Any, + wannier_projection_type, + reference_bands: orm.BandsData | None, + bands_kpoints: orm.KpointsData | None, +) -> dict[str, Any]: + """Build the static inputs for the Wannier90 task.""" + _validate_reference_bands_projection_type(reference_bands, wannier_projection_type) + w90_overrides = _copy_nested_containers(protocol_inputs.get("w90_bands", {})) + if reference_bands is not None: + w90_builder = Wannier90OptimizeWorkChain.get_builder_from_protocol( + structure=structure, + codes=codes, + pseudo_family=pseudo_family, + overrides=protocol_inputs.get("w90_bands", {}), + projection_type=wannier_projection_type, + reference_bands=reference_bands, + bands_kpoints=bands_kpoints, + ) + w90_builder.separate_plotting = False + else: + w90_builder = Wannier90BandsWorkChain.get_builder_from_protocol( + structure=structure, + codes=codes, + pseudo_family=pseudo_family, + overrides=protocol_inputs.get("w90_bands", {}), + projection_type=wannier_projection_type, + bands_kpoints=bands_kpoints, + ) + + w90_bands = get_dict_from_builder(w90_builder) + w90_bands = _normalize_nested_python_values(w90_bands) + w90_bands = _restore_nested_fields( + w90_bands, + w90_overrides, + ( + "scf.pw.metadata", + "nscf.pw.metadata", + "pw2wannier90.pw2wannier90.metadata", + "wannier90.wannier90.metadata", + ), + ) + if wannier_projection_type == WannierProjectionType.ATOMIC_PROJECTORS_QE: + w90_bands.pop("projwfc", None) + + w90_bands.pop("structure", None) + w90_bands.pop("open_grid", None) + return w90_bands + + +def _build_ph_inputs( + *, + codes: dict[str, Any], + protocol: str | None, + protocol_inputs: dict[str, Any], + kwargs: dict[str, Any], +) -> dict[str, Any]: + """Build the static inputs for the phonon task.""" + if codes.get("ph") is None: + return {} + + ph_base_inputs = protocol_inputs.get("ph_base", {}) + _add_metadata_stash_target_base( + ph_base_inputs.setdefault("ph", {}), codes.get("ph") + ) + + ph_base_builder = PhBaseWorkChain.get_builder_from_protocol( + codes["ph"], None, protocol, overrides=ph_base_inputs, **kwargs + ) + ph_base = get_dict_from_builder(ph_base_builder) + ph_base = _normalize_nested_python_values(ph_base) + ph_base = _restore_nested_fields(ph_base, ph_base_inputs, ("ph.metadata",)) + ph_base.pop("clean_workdir", None) + ph_base.pop("qpoints_distance", None) + return ph_base + + +def _build_pw_inputs( + *, + code: Any, + structure: orm.StructureData, + protocol: str | None, + pseudo_family: Any, + namespace_inputs: dict[str, Any], + kwargs: dict[str, Any], +) -> dict[str, Any]: + """Build static inputs for a direct PW task.""" + pw_inputs = recursive_merge( + {"pseudo_family": pseudo_family}, + namespace_inputs, + ) + pw_builder = PwBaseWorkChain.get_builder_from_protocol( + code=code, + structure=structure, + protocol=protocol, + overrides=pw_inputs, + **kwargs, + ) + inputs = get_dict_from_builder(pw_builder) + inputs = _normalize_nested_python_values(inputs) + inputs = _restore_nested_fields(inputs, namespace_inputs, ("pw.metadata",)) + inputs.pop("clean_workdir", None) + inputs.pop("kpoints_distance", None) + return inputs + + +def _build_epw_inputs( + *, + codes: dict[str, Any], + structure: orm.StructureData, + protocol: str | None, + protocol_inputs: dict[str, Any], + kwargs: dict[str, Any], +) -> tuple[dict[str, Any], dict[str, Any]]: + """Build the static inputs for the transformation and bands EPW tasks.""" + epw_base: dict[str, Any] = {} + epw_bands: dict[str, Any] = {} + epw_code = codes.get("epw") + + for namespace in ("epw_base", "epw_bands"): + epw_inputs = protocol_inputs.get(namespace, {}) + if namespace == "epw_base": + _add_options_stash_target_base(epw_inputs, epw_code) + + epw_builder = EpwBaseWorkChain.get_builder_from_protocol( + code=codes["epw"], + structure=structure, + protocol=protocol, + overrides=epw_inputs, + **kwargs, + ) + + if "settings" in epw_inputs: + epw_builder.settings = orm.Dict(epw_inputs["settings"]) + if "parallelization" in epw_inputs: + epw_builder.parallelization = orm.Dict(epw_inputs["parallelization"]) + + if namespace == "epw_base": + epw_base = get_dict_from_builder(epw_builder) + epw_base = _normalize_nested_python_values(epw_base) + epw_base = _restore_nested_fields(epw_base, epw_inputs, ("options",)) + else: + epw_bands = get_dict_from_builder(epw_builder) + epw_bands = _normalize_nested_python_values(epw_bands) + epw_bands = _restore_nested_fields(epw_bands, epw_inputs, ("options",)) + + return epw_base, epw_bands + + +def _drop_epw_mesh_generation_inputs(inputs: dict[str, Any]) -> dict[str, Any]: + """Drop builder-side mesh generation inputs when explicit meshes are wired in the graph.""" + cleaned = _copy_nested_containers(inputs) + cleaned.pop("kfpoints_factor", None) + cleaned.pop("qfpoints_distance", None) + return cleaned + + +def _build_wannier90_task_inputs( + w90_bands: dict[str, Any], + structure: orm.StructureData, + kpoints_scf: orm.KpointsData, + kpoints_nscf: orm.KpointsData, + wannier90_kpoints: orm.KpointsData, + parameters: Any, +) -> dict[str, Any]: + """Build inputs for a Wannier90 workchain task.""" + inputs = _copy_nested_containers(w90_bands) + inputs["structure"] = structure + inputs["scf"]["kpoints"] = kpoints_scf + inputs["nscf"]["kpoints"] = kpoints_nscf + inputs["wannier90"]["wannier90"]["kpoints"] = wannier90_kpoints + inputs["wannier90"]["wannier90"]["parameters"] = parameters + return inputs + + +def build_task_inputs( + *, + codes: dict, + structure: orm.StructureData, + protocol: str | None = None, + overrides: dict | None = None, + wannier_projection_type=None, + reference_bands: orm.BandsData | None = None, + bands_kpoints: orm.KpointsData | None = None, + parent_folder_ph: orm.RemoteData | orm.RemoteStashFolderData | None = None, + **kwargs, +) -> dict[str, Any]: + """Build and validate the static inputs required by the prep WorkGraph.""" + if hasattr(codes, "items"): + codes = dict(codes.items()) + if overrides and hasattr(overrides, "items"): + overrides = dict(overrides.items()) + + protocol_inputs = get_protocol_inputs(protocol, overrides) + if parent_folder_ph is not None: + _validate_parent_folder_ph(parent_folder_ph, structure) + use_epw_wannierize = should_epw_wannierize(protocol_inputs) + protocol_inputs = _copy_nested_containers(protocol_inputs) + if parent_folder_ph is not None: + protocol_inputs["parent_folder_ph"] = parent_folder_ph + if use_epw_wannierize: + protocol_inputs.pop("w90_bands", None) + else: + protocol_inputs.pop("scf", None) + protocol_inputs.pop("nscf", None) + + validation_error = validate_inputs(protocol_inputs) + if validation_error is not None: + raise ValueError(validation_error) + + if wannier_projection_type is None: + wannier_projection_type = WannierProjectionType.ATOMIC_PROJECTORS_QE + + pseudo_family = protocol_inputs.pop("pseudo_family", None) + w90_bands: dict[str, Any] = {} + scf: dict[str, Any] = {} + nscf: dict[str, Any] = {} + + if use_epw_wannierize: + scf = _build_pw_inputs( + code=codes["pw"], + structure=structure, + protocol=protocol, + pseudo_family=pseudo_family, + namespace_inputs=protocol_inputs.get("scf", {}), + kwargs=kwargs, + ) + nscf = _build_pw_inputs( + code=codes["pw"], + structure=structure, + protocol=protocol, + pseudo_family=pseudo_family, + namespace_inputs=protocol_inputs.get("nscf", {}), + kwargs=kwargs, + ) + else: + w90_bands = _build_wannier90_inputs( + codes=codes, + structure=structure, + protocol_inputs=protocol_inputs, + pseudo_family=pseudo_family, + wannier_projection_type=wannier_projection_type, + reference_bands=reference_bands, + bands_kpoints=bands_kpoints, + ) + ph_base = {} + if parent_folder_ph is None: + ph_base = _build_ph_inputs( + codes=codes, + protocol=protocol, + protocol_inputs=protocol_inputs, + kwargs=kwargs, + ) + epw_base, epw_bands = _build_epw_inputs( + codes=codes, + structure=structure, + protocol=protocol, + protocol_inputs=protocol_inputs, + kwargs=kwargs, + ) + + return { + "w90_bands": w90_bands, + "scf": scf, + "nscf": nscf, + "ph_base": ph_base, + "epw_base": epw_base, + "epw_bands": epw_bands, + "qpoints_distance": orm.Float(protocol_inputs["qpoints_distance"]), + "kpoints_distance_scf": orm.Float(protocol_inputs["kpoints_distance_scf"]), + "kpoints_factor_nscf": orm.Int(protocol_inputs["kpoints_factor_nscf"]), + "kpoints_force_parity": orm.Bool( + protocol_inputs.get("kpoints_force_parity", False) + ), + "parent_folder_ph": parent_folder_ph, + } + + +@task.calcfunction( + outputs=spec.namespace( + kpoints_scf=Any, + qpoints=Any, + kpoints_nscf=Any, + ) +) +def generate_reciprocal_points( + structure, + force_parity, + kpoints_distance_scf, + qpoints_distance, + kpoints_factor_nscf, + parent_folder_ph=None, +): + """Generate the SCF k-point mesh, Q-point mesh, and NSCF k-point mesh for the EPW parameters.""" + kpoints_scf = _create_kpoints_from_distance_node( + structure=structure, + distance=kpoints_distance_scf, + force_parity=force_parity, + ) + if parent_folder_ph is not None: + qpoints = validate_parent_ph_inputs(parent_folder_ph, structure) + else: + qpoints = _create_kpoints_from_distance_node( + structure=structure, + distance=qpoints_distance, + force_parity=force_parity, + ) + qpoints_mesh = qpoints.get_kpoints_mesh()[0] + kpoints_nscf = orm.KpointsData() + kpoints_nscf.set_kpoints_mesh( + [value * kpoints_factor_nscf.value for value in qpoints_mesh] + ) + + return { + "kpoints_scf": kpoints_scf, + "qpoints": qpoints, + "kpoints_nscf": kpoints_nscf, + } + + +@task() +def should_run_wannier90(w90_parameters) -> bool: + """Mirror the outline guard for the Wannier90 branch.""" + bands_plot = False + if w90_parameters is not None: + bands_plot = _as_dict(w90_parameters).get("bands_plot", False) + return orm.Bool(bands_plot) + + +@task.calcfunction( + outputs=spec.namespace( + nscf_kpoints=Any, + wannier90_kpoints=Any, + parameters=Any, + ) +) +def prepare_wannier90_runtime_inputs(parameters, kpoints_nscf): + """Build runtime NSCF and Wannier90 kpoints from the NSCF mesh.""" + from aiida_wannier90_workflows.utils.kpoints import get_explicit_kpoints + + explicit_kpoints = get_explicit_kpoints(kpoints_nscf) + points, weights = explicit_kpoints.get_kpoints(also_weights=True) + labels = explicit_kpoints.labels if explicit_kpoints.labels else None + + nscf_kpoints = orm.KpointsData() + nscf_kpoints.set_kpoints(points, weights=weights, labels=labels) + + wannier90_kpoints = orm.KpointsData() + wannier90_kpoints.set_kpoints(points, weights=weights, labels=labels) + + updated = parameters.get_dict() + updated["mp_grid"] = kpoints_nscf.get_kpoints_mesh()[0] + return { + "nscf_kpoints": nscf_kpoints, + "wannier90_kpoints": wannier90_kpoints, + "parameters": orm.Dict(updated), + } + + +@task.calcfunction() +def create_kpoints_gamma(): + """Create the gamma-only mesh used by the transformation EPW run.""" + gamma = orm.KpointsData() + gamma.set_kpoints_mesh([1, 1, 1]) + return gamma + + +@task.calcfunction() +def prepare_nscf_runtime_kpoints(kpoints_nscf): + """Build explicit NSCF k-points from the coarse NSCF mesh.""" + from aiida_wannier90_workflows.utils.kpoints import get_explicit_kpoints + + return get_explicit_kpoints(kpoints_nscf) + + +@task.workfunction() +def get_seekpath_explicit_kpoints(reference_output): + """Trace the explicit k-point path back to the internal seekpath calcfunction.""" + from aiida.common.links import LinkType + from aiida.orm import ProcessNode + + workchain = None + for link in reference_output.base.links.get_incoming( + link_type=LinkType.RETURN + ).all(): + if not isinstance(link.node, ProcessNode): + continue + outgoing = link.node.base.links.get_outgoing(link_type=LinkType.CALL_CALC).all() + if any(out.link_label == "seekpath_structure_analysis" for out in outgoing): + workchain = link.node + break + + if workchain is None: + raise ValueError( + "Could not find a parent workchain with a `seekpath_structure_analysis` call." + ) + + for link in workchain.base.links.get_outgoing(link_type=LinkType.CALL_CALC).all(): + if link.link_label == "seekpath_structure_analysis": + return link.node.outputs.explicit_kpoints + + raise ValueError( + "The parent workchain does not expose an `explicit_kpoints` output from seekpath." + ) + + +@task.calcfunction() +def extract_kpoints_path(band_structure): + """Convert a ``BandsData`` output into a standalone ``KpointsData`` node.""" + kpoints = orm.KpointsData() + kpoints.set_cell(band_structure.cell, band_structure.pbc) + + points, weights = band_structure.get_kpoints(also_weights=True) + labels = band_structure.labels if band_structure.labels else None + kpoints.set_kpoints(points, weights=weights, labels=labels) + + return kpoints + + +@task.workfunction() +def generate_seekpath_explicit_kpoints(structure): + """Generate an explicit band path directly from the input structure.""" + from aiida_quantumespresso.calculations.functions.seekpath_structure_analysis import ( + seekpath_structure_analysis, + ) + + result = seekpath_structure_analysis( + structure=structure, + metadata={"call_link_label": "seekpath_structure_analysis"}, + ) + return result["explicit_kpoints"] + + +@task( + outputs=spec.namespace( + retrieved=Any, + epw_folder=Any, + epw_stash=Any, + ) +) +def results(retrieved, epw_folder, epw_stash=None, epw_kpoints=None, epw_qpoints=None): + """Compatibility task retained so previously saved prep graphs can still be loaded.""" + return { + "retrieved": retrieved, + "epw_folder": epw_folder, + "epw_stash": epw_stash, + "epw_kpoints": epw_kpoints, + "epw_qpoints": epw_qpoints, + } + + +def prep_from_inputs( + *, + structure: orm.StructureData, + inputs: dict[str, Any], + use_wannier90_optimize: bool, + w90_chk_to_ukk_script: orm.RemoteData | None = None, +): + """Create the prep WorkGraph from pre-built static task inputs.""" + use_epw_wannierize = should_epw_wannierize(inputs) + w90_bands = inputs.get("w90_bands", {}) + scf = inputs.get("scf", {}) + nscf = inputs.get("nscf", {}) + ph_base = inputs.get("ph_base", {}) + epw_base = inputs["epw_base"] + epw_bands = inputs["epw_bands"] + + with WorkGraph( + name="prep", + outputs=spec.namespace( + retrieved=Any, + epw_folder=Any, + epw_stash=Any, + ), + ) as wg: + parent_folder_ph = inputs.get("parent_folder_ph") + reciprocal_points = generate_reciprocal_points( + structure=structure, + force_parity=inputs["kpoints_force_parity"], + kpoints_distance_scf=inputs["kpoints_distance_scf"], + qpoints_distance=inputs["qpoints_distance"], + kpoints_factor_nscf=inputs["kpoints_factor_nscf"], + parent_folder_ph=parent_folder_ph, + ) + + if use_epw_wannierize: + scf_inputs = recursive_merge( + _copy_nested_containers(scf), + { + "pw": {"structure": structure}, + "kpoints": reciprocal_points.kpoints_scf, + }, + ) + scf_run = PwBaseTask(**_set_call_link_label(scf_inputs, "scf")) + _apply_socket_overrides( + scf_run._task.inputs, + scf, + ("pw.code", "pw.metadata"), + ) + + nscf_kpoints = prepare_nscf_runtime_kpoints( + kpoints_nscf=reciprocal_points.kpoints_nscf, + ) + nscf_inputs = recursive_merge( + _copy_nested_containers(nscf), + { + "pw": { + "structure": structure, + "parent_folder": scf_run.remote_folder, + }, + "kpoints": nscf_kpoints.result, + }, + ) + nscf_run = PwBaseTask(**_set_call_link_label(nscf_inputs, "nscf")) + _apply_socket_overrides( + nscf_run._task.inputs, + nscf, + ("pw.code", "pw.metadata"), + ) + + scf_remote = scf_run.remote_folder + nscf_remote = nscf_run.remote_folder + chk_folder = None + else: + should_run_w90 = should_run_wannier90( + w90_parameters=w90_bands.get("wannier90", {}) + .get("wannier90", {}) + .get("parameters") + ) + with If(should_run_w90.result): + wannier90_runtime_inputs = prepare_wannier90_runtime_inputs( + parameters=w90_bands["wannier90"]["wannier90"]["parameters"], + kpoints_nscf=reciprocal_points.kpoints_nscf, + ) + if use_wannier90_optimize: + wannier90_inputs = _build_wannier90_task_inputs( + w90_bands=w90_bands, + structure=structure, + kpoints_scf=reciprocal_points.kpoints_scf, + kpoints_nscf=wannier90_runtime_inputs.nscf_kpoints, + wannier90_kpoints=wannier90_runtime_inputs.wannier90_kpoints, + parameters=wannier90_runtime_inputs.parameters, + ) + wannier90_run_proxy = Wannier90OptimizeTask( + **_set_call_link_label(wannier90_inputs, "w90_bands"), + ) + _apply_socket_overrides( + wannier90_run_proxy._task.inputs, + w90_bands, + ( + "scf.pw.code", + "scf.pw.metadata", + "nscf.pw.code", + "nscf.pw.metadata", + "pw2wannier90.pw2wannier90.code", + "pw2wannier90.pw2wannier90.metadata", + "wannier90.wannier90.code", + "wannier90.wannier90.metadata", + ), + ) + + scf_remote = wannier90_run_proxy.scf.remote_folder + nscf_remote = wannier90_run_proxy.nscf.remote_folder + + if _as_bool(w90_bands.get("optimize_disproj", False)): + chk_folder = wannier90_run_proxy.wannier90_optimal.remote_folder + else: + chk_folder = wannier90_run_proxy.wannier90.remote_folder + else: + wannier90_inputs = _build_wannier90_task_inputs( + w90_bands=w90_bands, + structure=structure, + kpoints_scf=reciprocal_points.kpoints_scf, + kpoints_nscf=wannier90_runtime_inputs.nscf_kpoints, + wannier90_kpoints=wannier90_runtime_inputs.wannier90_kpoints, + parameters=wannier90_runtime_inputs.parameters, + ) + wannier90_run_proxy = Wannier90BandsTask( + **_set_call_link_label(wannier90_inputs, "w90_bands"), + ) + _apply_socket_overrides( + wannier90_run_proxy._task.inputs, + w90_bands, + ( + "scf.pw.code", + "scf.pw.metadata", + "nscf.pw.code", + "nscf.pw.metadata", + "pw2wannier90.pw2wannier90.code", + "pw2wannier90.pw2wannier90.metadata", + "wannier90.wannier90.code", + "wannier90.wannier90.metadata", + ), + ) + + scf_remote = wannier90_run_proxy.scf.remote_folder + nscf_remote = wannier90_run_proxy.nscf.remote_folder + chk_folder = wannier90_run_proxy.wannier90.remote_folder + + if parent_folder_ph is None: + phonons_inputs = recursive_merge( + ph_base, + { + "qpoints": reciprocal_points.qpoints, + "ph": {"parent_folder": scf_remote}, + }, + ) + phonons_run = PhBaseTask(**_set_call_link_label(phonons_inputs, "ph_base")) + _apply_socket_overrides( + phonons_run._task.inputs, + ph_base, + ("ph.code", "ph.metadata"), + ) + parent_ph_for_epw = phonons_run.remote_folder + else: + parent_ph_for_epw = parent_folder_ph + + kfpoints = create_kpoints_gamma() + epw_inputs = recursive_merge( + _drop_epw_mesh_generation_inputs(epw_base), + { + "structure": structure, + "parent_folder_ph": parent_ph_for_epw, + "parent_folder_nscf": nscf_remote, + "kpoints": reciprocal_points.kpoints_nscf, + "kfpoints": kfpoints.result, + "qpoints": reciprocal_points.qpoints, + "qfpoints": kfpoints.result, + }, + ) + if chk_folder is not None: + epw_inputs["parent_folder_chk"] = chk_folder + if w90_chk_to_ukk_script is not None: + epw_inputs["w90_chk_to_ukk_script"] = w90_chk_to_ukk_script + + epw_run = EpwBaseTask(**_set_call_link_label(epw_inputs, "epw_base")) + _apply_socket_overrides( + epw_run._task.inputs, + epw_base, + ("code", "options"), + ) + + if epw_bands: + epw_bands_inputs = recursive_merge( + _drop_epw_mesh_generation_inputs(epw_bands), + { + "structure": structure, + "parent_folder_epw": epw_run.remote_stash, + "kpoints": reciprocal_points.kpoints_nscf, + "qpoints": reciprocal_points.qpoints, + }, + ) + + if "bands_kpoints" in w90_bands: + bands_kpoints_source = w90_bands["bands_kpoints"] + elif w90_bands: + bands_kpoints_source = get_seekpath_explicit_kpoints( + reference_output=nscf_remote + ).result + else: + bands_kpoints_source = generate_seekpath_explicit_kpoints( + structure=structure + ).result + + epw_bands_inputs["qfpoints"] = bands_kpoints_source + epw_bands_inputs["kfpoints"] = bands_kpoints_source + + epw_bands_run = EpwBaseTask( + **_set_call_link_label(epw_bands_inputs, "epw_bands"), + ) + _apply_socket_overrides( + epw_bands_run._task.inputs, + epw_bands, + ("code", "options"), + ) + + wg.outputs.retrieved = epw_run.retrieved + wg.outputs.epw_folder = epw_run.remote_folder + wg.outputs.epw_stash = epw_run.remote_stash + + return wg + + +def prep( + codes: dict, + structure: orm.StructureData, + protocol: str | None = None, + overrides: dict | None = None, + wannier_projection_type=None, + reference_bands: orm.BandsData | None = None, + bands_kpoints: orm.KpointsData | None = None, + parent_folder_ph: orm.RemoteData | orm.RemoteStashFolderData | None = None, + w90_chk_to_ukk_script: orm.RemoteData | None = None, + **kwargs, +): + """Build the EPW preparation WorkGraph.""" + inputs = build_task_inputs( + codes=codes, + structure=structure, + protocol=protocol, + overrides=overrides, + wannier_projection_type=wannier_projection_type, + reference_bands=reference_bands, + bands_kpoints=bands_kpoints, + parent_folder_ph=parent_folder_ph, + **kwargs, + ) + return prep_from_inputs( + structure=structure, + inputs=inputs, + use_wannier90_optimize=reference_bands is not None, + w90_chk_to_ukk_script=w90_chk_to_ukk_script, + ) diff --git a/src/aiida_epw/workgraphs/supercon.py b/src/aiida_epw/workgraphs/supercon.py new file mode 100644 index 00000000..ccf157ea --- /dev/null +++ b/src/aiida_epw/workgraphs/supercon.py @@ -0,0 +1,392 @@ +"""WorkGraph implementation of the superconductivity workflow.""" + +from __future__ import annotations + +from typing import Any, Union + +from aiida import orm +from aiida.common import AttributeDict +from aiida.engine import ProcessBuilder +from aiida.engine.processes.builder import ProcessBuilderNamespace +from aiida_workgraph import If, spec, task +from aiida_workgraph.utils import get_dict_from_builder + +from aiida_epw.tools.workchain import find_related_calculation +from aiida_epw.workflows.base import EpwBaseWorkChain + +BuilderNamespace = Union[ProcessBuilderNamespace, ProcessBuilder] + +__all__ = ("supercon",) + + +def get_protocol_inputs( + protocol: str | None = None, + overrides: dict | None = None, +) -> dict: + """Return the inputs for the EPW preparation workflow based on a protocol.""" + from importlib_resources import files + from aiida_epw.workflows import protocols + from aiida_quantumespresso.workflows.protocols.utils import ProtocolMixin + + # 1. Load protocol file from workflows side + filepath = files(protocols) / "supercon.yaml" + + AdHocProtocol = type( + "AdHocProtocol", + (ProtocolMixin,), + { + "get_protocol_filepath": classmethod(lambda cls: filepath), + "_validate_override_keys": classmethod(lambda cls, overrides: None), + }, + ) + + return AdHocProtocol.get_protocol_inputs(protocol, overrides) + + +def _namespace_to_dict(namespace: BuilderNamespace) -> dict[str, Any]: + """Convert a builder namespace into a plain nested mapping.""" + return get_dict_from_builder(namespace) + + +def _sorted_interpolation_distances( + interpolation_distance: orm.Float | orm.List, +) -> list[orm.Float]: + """Normalize the interpolation distance input into an ascending list of floats.""" + if isinstance(interpolation_distance, orm.Float): + values = [interpolation_distance.value] + else: + values = [float(value) for value in interpolation_distance.get_list()] + + if not values: + raise ValueError("At least one interpolation distance must be provided.") + + return [orm.Float(value) for value in sorted(values)] + + +@task.calcfunction() +def update_inputepw_degaussq(parameters, degaussq): + """Return a new parameters node with an updated ``INPUTEPW.degaussq`` value.""" + updated = parameters.get_dict() + updated.setdefault("INPUTEPW", {})["degaussq"] = degaussq.value + return orm.Dict(updated) + + +@task.calcfunction() +def derive_degaussq_from_a2f(a2f): + """Derive the smearing used by the original workflow from the highest A2F frequency.""" + frequency = a2f.get_frequency() + return orm.Float(frequency[-1] / 100) + + +@task.calcfunction() +def extract_allen_dynes_tc(output_parameters): + """Extract the Allen-Dynes critical temperature from parsed output parameters.""" + return orm.Float(output_parameters.get_dict()["Allen_Dynes_Tc"]) + + +@task.calcfunction() +def has_converged(previous_tc, current_tc, threshold, total_runs): + """Evaluate the same convergence condition as the original workchain.""" + if total_runs.value < 3: + return orm.Bool(False) + + relative_change = abs(previous_tc.value - current_tc.value) / current_tc.value + return orm.Bool(relative_change < threshold.value) + + +@task.calcfunction() +def should_run_final_epw(is_converged, always_run_final): + """Decide whether the final isotropic and anisotropic runs should be launched.""" + return orm.Bool(is_converged.value or always_run_final.value) + + +@task.calcfunction(outputs=spec.namespace(kfpoints=Any, qfpoints=Any)) +def extract_restart_meshes(parent_folder_epw): + """Extract the fine k/q meshes from the EPW calculation used as restart parent.""" + restart_calculation = find_related_calculation(parent_folder_epw) + return { + "kfpoints": restart_calculation.inputs.kfpoints, + "qfpoints": restart_calculation.inputs.qfpoints, + } + + +def _get_prep_reciprocal_points(parent_prep): + """Return the reciprocal meshes produced by the prep workgraph.""" + reciprocal_points = ( + parent_prep.base.links.get_outgoing( + link_label_filter="generate_reciprocal_points" + ) + .first() + .node + ) + return AttributeDict( + { + "kpoints": reciprocal_points.outputs.kpoints_nscf, + "qpoints": reciprocal_points.outputs.qpoints, + } + ) + + +def _get_prep_structure(parent_prep): + """Return the structure used by the prep workgraph.""" + reciprocal_points = ( + parent_prep.base.links.get_outgoing( + link_label_filter="generate_reciprocal_points" + ) + .first() + .node + ) + return reciprocal_points.inputs.structure + + +def _get_prep_epw_base(parent_prep): + """Return the internal EPW base task launched by the prep workgraph.""" + link = parent_prep.base.links.get_outgoing(link_label_filter="epw_base").first() + if link is None: + raise ValueError( + "Could not find the `epw_base` subprocess in the prep workgraph." + ) + return link.node + + +def _get_prep_restart_parent_folder(parent_prep): + """Return the restart parent folder exposed by the prep workgraph.""" + try: + return parent_prep.outputs.epw_stash + except AttributeError: + pass + + try: + epw_folder = parent_prep.outputs.epw_folder + except AttributeError as exception: + raise ValueError( + "Could not determine a restart folder from the prep workgraph outputs." + ) from exception + + epw_base = _get_prep_epw_base(parent_prep) + clean_workdir = getattr(epw_base.inputs, "clean_workdir", orm.Bool(False)) + if clean_workdir.value: + raise ValueError( + "The prep workgraph does not expose `epw_stash`, and its internal `epw_base` " + "was configured with `clean_workdir=True`, so `epw_folder` is not a safe restart source." + ) + + return epw_folder + + +@task( + outputs=spec.namespace( + converged=Any, + epw_final_a2f_output_parameters=Any, + epw_final_a2f_a2f=Any, + ) +) +def results(converged, output_parameters, a2f): + """Expose the final outputs following the original ``results`` step.""" + return { + "converged": converged, + "epw_final_a2f_output_parameters": output_parameters, + "epw_final_a2f_a2f": a2f, + } + + +@task.graph( + outputs=["converged", "epw_final_a2f_output_parameters", "epw_final_a2f_a2f"] +) +def supercon( + code, + parent_epw, + protocol=None, + overrides=None, + parent_folder_epw=None, + **kwargs, +): + """Superconductivity WorkGraph.""" + if overrides and hasattr(overrides, "items"): + overrides = dict(overrides.items()) + + inputs = get_protocol_inputs(protocol, overrides) + + if parent_epw.process_label == "WorkGraph": + reciprocal_points = _get_prep_reciprocal_points(parent_epw) + structure = _get_prep_structure(parent_epw) + epw_source = AttributeDict( + { + "inputs": AttributeDict( + { + "structure": structure, + "kpoints": reciprocal_points.kpoints, + "qpoints": reciprocal_points.qpoints, + } + ) + } + ) + elif parent_epw.process_label == "EpwPrepWorkChain": + epw_source = ( + parent_epw.base.links.get_outgoing(link_label_filter="epw_base") + .first() + .node + ) + elif parent_epw.process_label == "EpwBaseWorkChain": + epw_source = parent_epw + else: + raise ValueError(f"Invalid parent_epw process: {parent_epw.process_label}") + + if parent_folder_epw is None: + if parent_epw.process_label == "WorkGraph": + parent_folder_epw = _get_prep_restart_parent_folder(parent_epw) + else: + from aiida_epw.workflows.supercon import get_restart_parent_folder + + parent_folder_epw = get_restart_parent_folder(parent_epw) + + sub_inputs = {} + for epw_namespace in ( + "epw_interp", + "epw_final_iso", + "epw_final_aniso", + ): + epw_inputs = inputs.get(epw_namespace, {}) + epw_builder = EpwBaseWorkChain.get_builder_from_protocol( + code=code, + structure=epw_source.inputs.structure, + protocol=protocol, + overrides=epw_inputs, + **kwargs, + ) + epw_builder.kpoints = epw_source.inputs.kpoints + epw_builder.qpoints = epw_source.inputs.qpoints + if "settings" in epw_inputs: + epw_builder.settings = orm.Dict(epw_inputs["settings"]) + + sub_inputs[epw_namespace] = get_dict_from_builder(epw_builder) + + interp_inputs = sub_inputs["epw_interp"] + final_iso_inputs = sub_inputs["epw_final_iso"] + final_aniso_inputs = sub_inputs["epw_final_aniso"] + + distance_input = inputs["interpolation_distance"] + if isinstance(distance_input, float): + distance_node = orm.Float(distance_input) + elif isinstance(distance_input, list): + distance_node = orm.List(distance_input) + else: + distance_node = distance_input + + interpolation_distances = _sorted_interpolation_distances(distance_node) + + convergence_threshold = inputs.get("convergence_threshold") + convergence_threshold_node = ( + orm.Float(convergence_threshold) if convergence_threshold is not None else None + ) + always_run_final = orm.Bool(inputs.get("always_run_final", False)) + structure = epw_source.inputs.structure + kfpoints_factor = orm.Int(inputs.get("kfpoints_factor", 1)) + + interpolation_tasks = [] + allen_dynes_tasks = [] + updated_interp_parameters = None + degaussq_task = None + previous_interpolation_task = None + + for index, distance in enumerate(interpolation_distances, start=1): + EpwInterpTask = task(identifier=f"epw_interp_{index:02d}")(EpwBaseWorkChain) + + inputs_merged = dict(interp_inputs) + inputs_merged.update( + { + "structure": structure, + "parent_folder_epw": parent_folder_epw, + "kfpoints_factor": kfpoints_factor, + "qfpoints_distance": distance, + } + ) + if updated_interp_parameters is not None: + inputs_merged["parameters"] = updated_interp_parameters.result + + interpolation_task = EpwInterpTask(**inputs_merged) + + if previous_interpolation_task is not None: + interpolation_task._task.waiting_on.add(previous_interpolation_task._task) + + interpolation_tasks.append(interpolation_task) + previous_interpolation_task = interpolation_task + + allen_dynes_task = extract_allen_dynes_tc( + output_parameters=interpolation_task.output_parameters, + ) + allen_dynes_tasks.append(allen_dynes_task) + + if index == 1: + degaussq_task = derive_degaussq_from_a2f( + a2f=interpolation_task.a2f, + ) + updated_interp_parameters = update_inputepw_degaussq( + parameters=interp_inputs["parameters"], + degaussq=degaussq_task.result, + ) + + if convergence_threshold_node is not None and len(allen_dynes_tasks) >= 2: + converged_source = has_converged( + previous_tc=allen_dynes_tasks[-2].result, + current_tc=allen_dynes_tasks[-1].result, + threshold=convergence_threshold_node, + total_runs=orm.Int(len(interpolation_tasks)), + ).result + elif convergence_threshold_node is not None: + converged_source = orm.Bool(False) + else: + converged_source = orm.Bool(True) + + run_final_task = should_run_final_epw( + is_converged=converged_source, + always_run_final=always_run_final, + ) + + last_interpolation_task = interpolation_tasks[-1] + + with If(run_final_task.result): + restart_meshes = extract_restart_meshes( + parent_folder_epw=last_interpolation_task.remote_folder, + ) + + final_iso_parameters = None + if degaussq_task is not None: + final_iso_parameters = update_inputepw_degaussq( + parameters=final_iso_inputs["parameters"], + degaussq=degaussq_task.result, + ) + + EpwFinalIsoTask = task(identifier="epw_final_iso")(EpwBaseWorkChain) + final_iso_inputs_merged = dict(final_iso_inputs) + final_iso_inputs_merged.update( + { + "structure": structure, + "parent_folder_epw": last_interpolation_task.remote_folder, + "kfpoints": restart_meshes.kfpoints, + "qfpoints": restart_meshes.qfpoints, + } + ) + if final_iso_parameters is not None: + final_iso_inputs_merged["parameters"] = final_iso_parameters.result + + final_iso_task = EpwFinalIsoTask(**final_iso_inputs_merged) + + EpwFinalAnisoTask = task(identifier="epw_final_aniso")(EpwBaseWorkChain) + final_aniso_inputs_merged = dict(final_aniso_inputs) + final_aniso_inputs_merged.update( + { + "structure": structure, + "parent_folder_epw": last_interpolation_task.remote_folder, + "kfpoints": restart_meshes.kfpoints, + "qfpoints": restart_meshes.qfpoints, + } + ) + final_aniso_task = EpwFinalAnisoTask(**final_aniso_inputs_merged) + final_aniso_task._task.waiting_on.add(final_iso_task._task) + + return results( + converged=converged_source, + output_parameters=last_interpolation_task.output_parameters, + a2f=last_interpolation_task.a2f, + ) diff --git a/tests/workgraphs/test_prep.py b/tests/workgraphs/test_prep.py new file mode 100644 index 00000000..037025b3 --- /dev/null +++ b/tests/workgraphs/test_prep.py @@ -0,0 +1,1490 @@ +"""Tests for ``aiida_epw.workgraphs.prep``.""" + +import json +from importlib import import_module + +import numpy as np +import pytest +from aiida import orm +from aiida.common import AttributeDict +from aiida.common.links import LinkType +from aiida.calculations.arithmetic.add import ArithmeticAddCalculation +from aiida.engine import ToContext, WorkChain, calcfunction +from aiida_workgraph import WorkGraph +from aiida_workgraph import task +from aiida_workgraph.engine.task_manager import TaskManager +from aiida_workgraph.utils import restore_workgraph_data_from_raw_inputs +from aiida_quantumespresso.utils.mapping import prepare_process_inputs +from aiida_wannier90_workflows.common.types import WannierProjectionType + +prep_module = import_module("aiida_epw.workgraphs.prep") + + +class NestedMetadataWorkChain(WorkChain): + """Minimal nested workchain used to probe runtime process-task submission.""" + + @classmethod + def define(cls, spec): + super().define(spec) + spec.expose_inputs(ArithmeticAddCalculation, namespace="scf") + spec.outline(cls.noop) + + def noop(self): + """Do nothing.""" + + +class InnerPwWorkChain(WorkChain): + """Intermediate workchain exposing a nested calcjob namespace.""" + + @classmethod + def define(cls, spec): + super().define(spec) + spec.expose_inputs(ArithmeticAddCalculation, namespace="pw") + spec.outline(cls.noop) + + def noop(self): + """Do nothing.""" + + +class DoubleNestedMetadataWorkChain(WorkChain): + """Minimal double-nested workchain mirroring ``scf.pw.metadata`` style inputs.""" + + @classmethod + def define(cls, spec): + super().define(spec) + spec.expose_inputs(InnerPwWorkChain, namespace="scf") + spec.outline(cls.noop) + + def noop(self): + """Do nothing.""" + + +class InnerSubmitArithmeticWorkChain(WorkChain): + """Intermediate workchain that forwards nested metadata to a CalcJob.""" + + @classmethod + def define(cls, spec): + super().define(spec) + spec.expose_inputs(ArithmeticAddCalculation, namespace="pw") + spec.outline(cls.run_pw, cls.inspect_pw) + + def run_pw(self): + inputs = AttributeDict( + self.exposed_inputs(ArithmeticAddCalculation, namespace="pw") + ) + inputs.metadata.call_link_label = "pw" + inputs = prepare_process_inputs(ArithmeticAddCalculation, inputs) + return ToContext(workchain_pw=self.submit(ArithmeticAddCalculation, **inputs)) + + def inspect_pw(self): + """Do nothing.""" + + +class OuterSubmitInnerWorkChain(WorkChain): + """Outer workchain that forwards ``scf.pw.metadata`` style inputs to an inner workchain.""" + + @classmethod + def define(cls, spec): + super().define(spec) + spec.expose_inputs(InnerSubmitArithmeticWorkChain, namespace="scf") + spec.outline(cls.run_scf, cls.inspect_scf) + + def run_scf(self): + inputs = AttributeDict( + self.exposed_inputs(InnerSubmitArithmeticWorkChain, namespace="scf") + ) + inputs.metadata.call_link_label = "scf" + return ToContext( + workchain_scf=self.submit(InnerSubmitArithmeticWorkChain, **inputs) + ) + + def inspect_scf(self): + """Do nothing.""" + + +class FakeWannier90BandsRuntimeWorkChain(WorkChain): + """Lightweight substitute for the high-level Wannier bands task.""" + + @classmethod + def define(cls, spec): + super().define(spec) + spec.input("structure", valid_type=orm.StructureData) + spec.input("clean_workdir", valid_type=orm.Bool, required=False) + spec.input_namespace("scf", dynamic=True) + spec.input_namespace("nscf", dynamic=True) + spec.input_namespace("pw2wannier90", dynamic=True) + spec.input_namespace("wannier90", dynamic=True) + spec.outline(cls.generate_outputs) + spec.output("scf.remote_folder", valid_type=orm.RemoteData) + spec.output("nscf.remote_folder", valid_type=orm.RemoteData) + spec.output("wannier90.remote_folder", valid_type=orm.RemoteData) + spec.output("primitive_structure", valid_type=orm.StructureData) + spec.output("band_structure", valid_type=orm.BandsData) + + def generate_outputs(self): + """Emit the outputs that the prep graph expects downstream.""" + code = self.inputs.scf["pw"]["code"] + seekpath = _fake_seekpath_structure_analysis( + structure=self.inputs.structure, + metadata={"call_link_label": "seekpath_structure_analysis"}, + ) + self.out( + "scf.remote_folder", + _make_remote_data(code, orm.Str(f"{self.node.uuid}-scf")), + ) + self.out( + "nscf.remote_folder", + _make_remote_data(code, orm.Str(f"{self.node.uuid}-nscf")), + ) + self.out( + "wannier90.remote_folder", + _make_remote_data(code, orm.Str(f"{self.node.uuid}-wannier90")), + ) + self.out("primitive_structure", seekpath["primitive_structure"]) + self.out("band_structure", _make_bands_data(self.inputs.structure)) + + +class FakePhBaseRuntimeWorkChain(WorkChain): + """Lightweight substitute for the high-level phonon task.""" + + @classmethod + def define(cls, spec): + super().define(spec) + spec.input_namespace("ph", dynamic=True) + spec.input("qpoints", valid_type=orm.KpointsData) + spec.outline(cls.generate_outputs) + spec.output("remote_folder", valid_type=orm.RemoteData) + + def generate_outputs(self): + """Emit the remote folder expected by the EPW task.""" + code = self.inputs.ph["code"] + self.out( + "remote_folder", _make_remote_data(code, orm.Str(f"{self.node.uuid}-ph")) + ) + + +class FakePwBaseRuntimeWorkChain(WorkChain): + """Lightweight substitute for the high-level PW task.""" + + @classmethod + def define(cls, spec): + super().define(spec) + spec.input_namespace("pw", dynamic=True) + spec.input("kpoints", valid_type=orm.KpointsData) + spec.outline(cls.generate_outputs) + spec.output("remote_folder", valid_type=orm.RemoteData) + + def generate_outputs(self): + """Emit the remote folder expected downstream.""" + code = self.inputs.pw["code"] + self.out( + "remote_folder", _make_remote_data(code, orm.Str(f"{self.node.uuid}-pw")) + ) + + +class FakeEpwBaseRuntimeWorkChain(WorkChain): + """Lightweight substitute for the high-level EPW task.""" + + @classmethod + def define(cls, spec): + super().define(spec) + spec.input("code", valid_type=orm.AbstractCode) + spec.input("structure", valid_type=orm.StructureData) + spec.input("parameters", valid_type=orm.Dict, required=False) + spec.input("options", valid_type=orm.Dict, required=False) + spec.input("settings", valid_type=orm.Dict, required=False) + spec.input("parallelization", valid_type=orm.Dict, required=False) + spec.input("clean_workdir", valid_type=orm.Bool, required=False) + spec.input("max_iterations", valid_type=orm.Int, required=False) + spec.input("qfpoints_distance", valid_type=orm.Float, required=False) + spec.input("kfpoints_factor", valid_type=orm.Int, required=False) + spec.input("parent_folder_ph", valid_type=orm.RemoteData, required=False) + spec.input("parent_folder_nscf", valid_type=orm.RemoteData, required=False) + spec.input("parent_folder_chk", valid_type=orm.RemoteData, required=False) + spec.input("parent_folder_epw", valid_type=orm.RemoteData, required=False) + spec.input("w90_chk_to_ukk_script", valid_type=orm.RemoteData, required=False) + spec.input("kpoints", valid_type=orm.KpointsData, required=False) + spec.input("kfpoints", valid_type=orm.KpointsData, required=False) + spec.input("qpoints", valid_type=orm.KpointsData, required=False) + spec.input("qfpoints", valid_type=orm.KpointsData, required=False) + spec.outline(cls.generate_outputs) + spec.output("retrieved", valid_type=orm.FolderData) + spec.output("remote_folder", valid_type=orm.RemoteData) + spec.output("remote_stash", valid_type=orm.RemoteData) + + def generate_outputs(self): + """Emit the outputs that the prep graph expects downstream.""" + code = self.inputs.code + self.out("retrieved", _make_folder_data()) + self.out( + "remote_folder", _make_remote_data(code, orm.Str(f"{self.node.uuid}-epw")) + ) + self.out( + "remote_stash", _make_remote_data(code, orm.Str(f"{self.node.uuid}-stash")) + ) + + +def _metadata_dict( + *, + account: str = "elph", + queue_name: str = "debug", + num_mpiprocs_per_machine: int = 128, + max_wallclock_seconds: int = 3600, +) -> dict: + """Return scheduler metadata used in the workgraph tests.""" + return { + "options": { + "resources": { + "num_machines": 1, + "num_mpiprocs_per_machine": num_mpiprocs_per_machine, + }, + "max_wallclock_seconds": max_wallclock_seconds, + "withmpi": True, + "account": account, + "queue_name": queue_name, + } + } + + +def _fake_wannier90_builder(codes): + """Return a minimal Wannier90 builder-like namespace.""" + builder = AttributeDict() + builder.structure = orm.StructureData( + cell=[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]] + ) + builder.open_grid = {} + builder.clean_workdir = orm.Bool(False) + builder.scf = AttributeDict( + { + "max_iterations": orm.Int(2), + "kpoints_distance": orm.Float(0.15), + "kpoints_force_parity": orm.Bool(False), + "pw": AttributeDict( + { + "code": codes["pw"], + "metadata": _metadata_dict(), + "parameters": orm.Dict({"SYSTEM": {"occupations": "smearing"}}), + "pseudos": {}, + } + ), + } + ) + builder.nscf = AttributeDict( + { + "max_iterations": orm.Int(2), + "kpoints_force_parity": orm.Bool(False), + "pw": AttributeDict( + { + "code": codes["pw"], + "metadata": _metadata_dict(), + "parameters": orm.Dict({"SYSTEM": {"occupations": "smearing"}}), + "pseudos": {}, + } + ), + } + ) + builder.pw2wannier90 = AttributeDict( + { + "max_iterations": orm.Int(2), + "scdm_sigma_factor": orm.Float(3.0), + "pw2wannier90": AttributeDict( + { + "code": codes["pw2wannier90"], + "metadata": _metadata_dict(num_mpiprocs_per_machine=8), + "parameters": orm.Dict({"INPUTPP": {}}), + } + ), + } + ) + builder.wannier90 = AttributeDict( + { + "max_iterations": orm.Int(2), + "shift_energy_windows": orm.Bool(True), + "auto_energy_windows": orm.Bool(False), + "auto_energy_windows_threshold": orm.Float(0.5), + "wannier90": AttributeDict( + { + "code": codes["wannier90"], + "metadata": _metadata_dict(num_mpiprocs_per_machine=1), + "parameters": orm.Dict( + { + "bands_plot": True, + "num_bands": 8, + "num_wann": 4, + "exclude_bands": range(1, 3), + } + ), + "settings": orm.Dict({}), + } + ), + } + ) + return builder + + +def _fake_ph_builder(code, *_args, **_kwargs): + """Return a minimal phonon builder-like namespace.""" + builder = AttributeDict() + builder.clean_workdir = orm.Bool(False) + builder.qpoints_distance = orm.Float(0.4) + builder.ph = AttributeDict( + { + "code": code, + "metadata": _metadata_dict(), + "parameters": orm.Dict({"INPUTPH": {}}), + } + ) + return builder + + +def _fake_pw_builder(code, structure, *_args, **_kwargs): + """Return a minimal PW builder-like namespace.""" + builder = AttributeDict() + builder.clean_workdir = orm.Bool(False) + builder.kpoints_distance = orm.Float(0.2) + builder.pw = AttributeDict( + { + "code": code, + "metadata": _metadata_dict(), + "parameters": orm.Dict({"SYSTEM": {"occupations": "smearing"}}), + "pseudos": {}, + "structure": structure, + } + ) + return builder + + +def _fake_epw_builder(code, structure, *_args, **kwargs): + """Return a minimal EPW builder-like namespace.""" + overrides = kwargs.get("overrides", {}) + builder = AttributeDict() + builder.code = code + builder.structure = structure + builder.parameters = orm.Dict( + overrides.get("parameters", {"INPUTEPW": {"band_plot": True}}) + ) + builder.options = orm.Dict(overrides.get("options", _metadata_dict()["options"])) + builder.settings = orm.Dict(overrides.get("settings", {})) + builder.parallelization = orm.Dict(overrides.get("parallelization", {})) + builder.clean_workdir = orm.Bool(False) + builder.max_iterations = orm.Int(2) + builder.qfpoints_distance = orm.Float(0.1) + builder.kfpoints_factor = orm.Int(2) + return builder + + +def _create_explicit_kpoints(mesh): + """Create an explicit KpointsData list from a mesh.""" + kpoints = orm.KpointsData() + points = [] + for i in range(mesh[0]): + for j in range(mesh[1]): + for k in range(mesh[2]): + points.append( + [ + i / mesh[0], + j / mesh[1], + k / mesh[2], + ] + ) + kpoints.set_kpoints(points) + return kpoints + + +def _mesh_kpoints(mesh): + """Create a mesh-style KpointsData node.""" + kpoints = orm.KpointsData() + kpoints.set_kpoints_mesh(mesh) + return kpoints + + +@calcfunction +def _make_remote_data(code, label): + """Create a lightweight remote data node for tests.""" + return orm.RemoteData( + computer=code.computer, + remote_path=f"/tmp/{label.value}", + ) + + +@calcfunction +def _make_bands_data(structure): + """Create a simple band structure node for tests.""" + bands_kpoints = orm.KpointsData() + bands_kpoints.set_cell(structure.cell, structure.pbc) + bands_kpoints.set_kpoints([[0.0, 0.0, 0.0], [0.5, 0.0, 0.0]]) + + band_structure = orm.BandsData() + band_structure.set_kpointsdata(bands_kpoints) + band_structure.set_bands(np.array([[0.0, 1.0], [0.5, 1.5]])) + return band_structure + + +@calcfunction +def _make_folder_data(): + """Create a folder data node for tests.""" + return orm.FolderData() + + +@calcfunction +def _fake_seekpath_structure_analysis(structure): + """Create lightweight seekpath-like outputs with provenance for tests.""" + primitive_structure = orm.StructureData(ase=structure.get_ase()) + return { + "primitive_structure": primitive_structure, + "explicit_kpoints": _create_explicit_kpoints([2, 2, 2]), + } + + +def test_prepare_wannier90_runtime_inputs_returns_distinct_nodes(): + """The runtime helper should not reuse the same node for two calcfunction outputs.""" + kpoints_nscf = orm.KpointsData() + kpoints_nscf.set_kpoints_mesh([2, 2, 2]) + + result = prep_module.prepare_wannier90_runtime_inputs._callable( + parameters=orm.Dict({"bands_plot": True}), + kpoints_nscf=kpoints_nscf, + ) + + nscf_points, nscf_weights = result["nscf_kpoints"].get_kpoints(also_weights=True) + wannier_points, wannier_weights = result["wannier90_kpoints"].get_kpoints( + also_weights=True + ) + + assert result["nscf_kpoints"].uuid != result["wannier90_kpoints"].uuid + assert np.array_equal(nscf_points, wannier_points) + assert np.array_equal(nscf_weights, wannier_weights) + assert result["parameters"].get_dict()["mp_grid"] == [2, 2, 2] + + +def test_drop_epw_mesh_generation_inputs_removes_builder_side_mesh_hints(): + """Explicit EPW meshes in the graph should override builder-side mesh-generation hints.""" + original = { + "options": {"account": "elph"}, + "kfpoints_factor": orm.Int(2), + "qfpoints_distance": orm.Float(0.1), + } + + cleaned = prep_module._drop_epw_mesh_generation_inputs(original) + + assert "kfpoints_factor" not in cleaned + assert "qfpoints_distance" not in cleaned + assert cleaned["options"]["account"] == "elph" + assert "kfpoints_factor" in original + assert "qfpoints_distance" in original + + +def test_validate_inputs_accepts_direct_epw_wannierize_inputs(): + """The workgraph should accept the direct EPW Wannierization entry point.""" + message = prep_module.validate_inputs( + { + "scf": {}, + "nscf": {}, + "ph_base": {}, + "epw_base": { + "parameters": {"INPUTEPW": {"wannierize": True, "proj": ["Si:s"]}} + }, + } + ) + + assert message is None + + +def test_validate_inputs_accepts_direct_epw_wannierize_with_epw_bands(): + """Direct EPW Wannierization can still request the EPW bands branch.""" + message = prep_module.validate_inputs( + { + "scf": {}, + "nscf": {}, + "ph_base": {}, + "epw_base": { + "parameters": {"INPUTEPW": {"wannierize": True, "proj": ["Si:s"]}} + }, + "epw_bands": {}, + } + ) + + assert message is None + + +def test_validate_inputs_accepts_parent_folder_ph_without_ph_base(): + """A reusable phonon parent folder should make the phonon namespace optional.""" + message = prep_module.validate_inputs( + { + "w90_bands": {}, + "parent_folder_ph": object(), + "epw_base": {}, + } + ) + + assert message is None + + +def test_generate_reciprocal_points_prefers_restart_qpoints_from_parent_ph( + fixture_localhost, + generate_remote_data, + monkeypatch, +): + """Restarting from an existing PhCalculation should make its q-points authoritative.""" + restart_qpoints = orm.KpointsData() + restart_qpoints.set_kpoints_mesh([5, 5, 5]) + restart_parent = generate_remote_data(fixture_localhost, "/remote/ph-restart") + + monkeypatch.setattr( + prep_module, + "_create_kpoints_from_distance_node", + lambda *args, **kwargs: _mesh_kpoints([4, 4, 4]), + ) + monkeypatch.setattr( + prep_module, + "validate_parent_ph_inputs", + lambda _folder, _structure: restart_qpoints, + ) + + result = prep_module.generate_reciprocal_points._callable( + structure=orm.StructureData(cell=[[1, 0, 0], [0, 1, 0], [0, 0, 1]]), + force_parity=orm.Bool(False), + kpoints_distance_scf=orm.Float(0.15), + qpoints_distance=orm.Float(0.3), + kpoints_factor_nscf=orm.Int(2), + parent_folder_ph=restart_parent, + ) + + assert result["qpoints"] == restart_qpoints + assert result["kpoints_scf"].get_kpoints_mesh()[0] == [4, 4, 4] + assert result["kpoints_nscf"].get_kpoints_mesh()[0] == [10, 10, 10] + + +def test_prep_passes_w90_chk_to_ukk_script_only_to_epw_base( + fixture_code, + generate_structure, + monkeypatch, +): + """The optional chk-to-ukk script should only be wired to the transformation EPW step.""" + codes = { + "pw": fixture_code("quantumespresso.pw"), + "pw2wannier90": fixture_code("quantumespresso.pw2wannier90"), + "wannier90": fixture_code("wannier90.wannier90"), + "ph": fixture_code("quantumespresso.ph"), + "epw": fixture_code("epw.epw"), + } + structure = generate_structure() + script = orm.RemoteData( + computer=codes["epw"].computer, remote_path="/tmp/w90_chk_to_ukk.jl" + ) + + monkeypatch.setattr( + prep_module, + "get_protocol_inputs", + lambda protocol=None, overrides=None: { + "pseudo_family": "PseudoDojo/0.5/PBE/SR/standard/upf", + "qpoints_distance": 0.5, + "kpoints_distance_scf": 0.15, + "kpoints_factor_nscf": 2, + "kpoints_force_parity": False, + "w90_bands": { + "scf": {"pw": {"metadata": _metadata_dict()}}, + "nscf": {"pw": {"metadata": _metadata_dict()}}, + "pw2wannier90": { + "pw2wannier90": { + "metadata": _metadata_dict(num_mpiprocs_per_machine=8) + } + }, + "wannier90": { + "wannier90": { + "metadata": _metadata_dict(num_mpiprocs_per_machine=1) + } + }, + }, + "ph_base": {"ph": {"metadata": _metadata_dict()}}, + "epw_base": {"options": _metadata_dict()["options"]}, + "epw_bands": {"options": _metadata_dict()["options"]}, + }, + ) + monkeypatch.setattr( + prep_module.Wannier90BandsWorkChain, + "get_builder_from_protocol", + lambda **kwargs: _fake_wannier90_builder(kwargs["codes"]), + ) + monkeypatch.setattr( + prep_module.Wannier90OptimizeWorkChain, + "get_builder_from_protocol", + lambda **kwargs: _fake_wannier90_builder(kwargs["codes"]), + ) + monkeypatch.setattr( + prep_module.PhBaseWorkChain, "get_builder_from_protocol", _fake_ph_builder + ) + monkeypatch.setattr( + prep_module.EpwBaseWorkChain, "get_builder_from_protocol", _fake_epw_builder + ) + + wg = prep_module.prep( + codes=codes, + structure=structure, + protocol="fast", + overrides={}, + w90_chk_to_ukk_script=script, + ) + engine_inputs = wg.to_engine_inputs(metadata=None) + + assert ( + engine_inputs["tasks"]["epw_base"]["w90_chk_to_ukk_script"].uuid == script.uuid + ) + assert "w90_chk_to_ukk_script" not in engine_inputs["tasks"]["epw_bands"] + + +def test_prep_workgraph_runs_with_fake_high_level_workchains( + fixture_code, + generate_structure, + monkeypatch, +): + """Run the real prep graph with lightweight stand-ins for the heavy workchains.""" + codes = { + "pw": fixture_code("quantumespresso.pw"), + "pw2wannier90": fixture_code("quantumespresso.pw2wannier90"), + "wannier90": fixture_code("wannier90.wannier90"), + "ph": fixture_code("quantumespresso.ph"), + "epw": fixture_code("epw.epw"), + } + structure = generate_structure() + + monkeypatch.setattr( + prep_module, + "get_protocol_inputs", + lambda protocol=None, overrides=None: { + "pseudo_family": "PseudoDojo/0.5/PBE/SR/standard/upf", + "qpoints_distance": 0.5, + "kpoints_distance_scf": 0.15, + "kpoints_factor_nscf": 2, + "kpoints_force_parity": False, + "w90_bands": { + "scf": {"pw": {"metadata": _metadata_dict()}}, + "nscf": {"pw": {"metadata": _metadata_dict()}}, + "pw2wannier90": { + "pw2wannier90": { + "metadata": _metadata_dict(num_mpiprocs_per_machine=8) + } + }, + "wannier90": { + "wannier90": { + "metadata": _metadata_dict(num_mpiprocs_per_machine=1) + } + }, + }, + "ph_base": {"ph": {"metadata": _metadata_dict()}}, + "epw_base": {"options": _metadata_dict()["options"]}, + "epw_bands": {"options": _metadata_dict()["options"]}, + }, + ) + monkeypatch.setattr( + prep_module.Wannier90BandsWorkChain, + "get_builder_from_protocol", + lambda **kwargs: _fake_wannier90_builder(kwargs["codes"]), + ) + monkeypatch.setattr( + prep_module.PhBaseWorkChain, + "get_builder_from_protocol", + _fake_ph_builder, + ) + monkeypatch.setattr( + prep_module.EpwBaseWorkChain, + "get_builder_from_protocol", + _fake_epw_builder, + ) + monkeypatch.setattr( + prep_module, "_apply_socket_overrides", lambda *args, **kwargs: None + ) + monkeypatch.setattr( + prep_module, + "Wannier90BandsTask", + task(FakeWannier90BandsRuntimeWorkChain), + ) + monkeypatch.setattr( + prep_module, + "PhBaseTask", + task(FakePhBaseRuntimeWorkChain), + ) + monkeypatch.setattr( + prep_module, + "EpwBaseTask", + task(FakeEpwBaseRuntimeWorkChain), + ) + + wg = prep_module.prep( + codes=codes, + structure=structure, + protocol="fast", + overrides={}, + ) + wg.run() + + assert wg.process is not None + assert wg.process.is_finished_ok + assert ( + wg.process.outputs.epw_stash.uuid + == wg.tasks["epw_base"].process.outputs.remote_stash.uuid + ) + + wannier_node = wg.tasks["w90_bands"].process + phonon_node = wg.tasks["ph_base"].process + epw_node = wg.tasks["epw_base"].process + epw_bands_node = wg.tasks["epw_bands"].process + assert wannier_node is not None + assert phonon_node is not None + assert epw_node is not None + assert epw_bands_node is not None + + nscf_kpoints = wannier_node.inputs.nscf.kpoints + wannier_kpoints = wannier_node.inputs.wannier90.wannier90.kpoints + with pytest.raises(AttributeError): + nscf_kpoints.get_kpoints_mesh() + with pytest.raises(AttributeError): + wannier_kpoints.get_kpoints_mesh() + assert len(nscf_kpoints.get_kpoints()) == len(wannier_kpoints.get_kpoints()) + + assert ( + phonon_node.inputs.ph.parent_folder.uuid + == wannier_node.outputs.scf.remote_folder.uuid + ) + assert ( + epw_node.inputs.parent_folder_ph.uuid == phonon_node.outputs.remote_folder.uuid + ) + assert ( + epw_node.inputs.parent_folder_nscf.uuid + == wannier_node.outputs.nscf.remote_folder.uuid + ) + assert ( + epw_node.inputs.parent_folder_chk.uuid + == wannier_node.outputs.wannier90.remote_folder.uuid + ) + reciprocal_points_node = wg.tasks["generate_reciprocal_points"].process + assert reciprocal_points_node is not None + assert ( + reciprocal_points_node.outputs.kpoints_nscf.get_kpoints_mesh()[0] + == epw_node.inputs.kpoints.get_kpoints_mesh()[0] + ) + assert ( + reciprocal_points_node.outputs.qpoints.get_kpoints_mesh()[0] + == epw_node.inputs.qpoints.get_kpoints_mesh()[0] + ) + assert epw_node.inputs.kfpoints.get_kpoints_mesh()[0] == [1, 1, 1] + assert epw_node.inputs.qfpoints.get_kpoints_mesh()[0] == [1, 1, 1] + with pytest.raises(AttributeError): + epw_bands_node.inputs.kfpoints.get_kpoints_mesh() + with pytest.raises(AttributeError): + epw_bands_node.inputs.qfpoints.get_kpoints_mesh() + epw_bands_kf = epw_bands_node.inputs.kfpoints.get_kpoints() + epw_bands_qf = epw_bands_node.inputs.qfpoints.get_kpoints() + assert len(epw_bands_kf) > 0 + assert np.array_equal(epw_bands_kf, epw_bands_qf) + seekpath_node = next( + link.node + for link in wannier_node.base.links.get_outgoing( + link_type=LinkType.CALL_CALC + ).all() + if link.link_label == "seekpath_structure_analysis" + ) + assert ( + epw_bands_node.inputs.kfpoints.uuid + == seekpath_node.outputs.explicit_kpoints.uuid + ) + + +def test_build_task_inputs_uses_direct_pw_namespaces_for_epw_wannierize( + fixture_code, + generate_structure, + monkeypatch, +): + """Direct EPW Wannierization should build standalone SCF/NSCF namespaces instead of `w90_bands`.""" + codes = { + "pw": fixture_code("quantumespresso.pw"), + "epw": fixture_code("epw.epw"), + } + structure = generate_structure() + + monkeypatch.setattr( + prep_module, + "get_protocol_inputs", + lambda protocol=None, overrides=None: { + "pseudo_family": "PseudoDojo/0.5/PBE/SR/standard/upf", + "qpoints_distance": 0.5, + "kpoints_distance_scf": 0.15, + "kpoints_factor_nscf": 2, + "kpoints_force_parity": False, + "scf": {"pw": {"metadata": _metadata_dict()}}, + "nscf": {"pw": {"metadata": _metadata_dict()}}, + "ph_base": {"ph": {"metadata": _metadata_dict()}}, + "epw_base": { + "options": _metadata_dict()["options"], + "parameters": {"INPUTEPW": {"wannierize": True, "proj": ["Si:s"]}}, + }, + }, + ) + monkeypatch.setattr( + prep_module.PwBaseWorkChain, "get_builder_from_protocol", _fake_pw_builder + ) + monkeypatch.setattr( + prep_module.PhBaseWorkChain, "get_builder_from_protocol", _fake_ph_builder + ) + monkeypatch.setattr( + prep_module.EpwBaseWorkChain, "get_builder_from_protocol", _fake_epw_builder + ) + + prepared_inputs = prep_module.build_task_inputs( + codes=codes, + structure=structure, + protocol="fast", + overrides={}, + ) + + assert prepared_inputs["w90_bands"] == {} + assert "pw" in prepared_inputs["scf"] + assert "pw" in prepared_inputs["nscf"] + assert "code" in prepared_inputs["epw_bands"] + + +def test_build_task_inputs_supports_analytic_wannier_projections( + fixture_code, + generate_structure, + monkeypatch, +): + """The workgraph helper should preserve analytic Wannier90 projections.""" + codes = { + "pw": fixture_code("quantumespresso.pw"), + "pw2wannier90": fixture_code("quantumespresso.pw2wannier90"), + "wannier90": fixture_code("wannier90.wannier90"), + "ph": fixture_code("quantumespresso.ph"), + "epw": fixture_code("epw.epw"), + } + structure = generate_structure() + captured = {} + + monkeypatch.setattr( + prep_module, + "get_protocol_inputs", + lambda protocol=None, overrides=None: { + "pseudo_family": "PseudoDojo/0.5/PBE/SR/standard/upf", + "qpoints_distance": 0.5, + "kpoints_distance_scf": 0.15, + "kpoints_factor_nscf": 2, + "kpoints_force_parity": False, + "w90_bands": { + "scf": {"pw": {"metadata": _metadata_dict()}}, + "nscf": {"pw": {"metadata": _metadata_dict()}}, + "pw2wannier90": {"pw2wannier90": {"metadata": _metadata_dict()}}, + "wannier90": {"wannier90": {"metadata": _metadata_dict()}}, + }, + "ph_base": {"ph": {"metadata": _metadata_dict()}}, + "epw_base": {"options": _metadata_dict()["options"]}, + }, + ) + + def fake_w90_builder(**kwargs): + captured["projection_type"] = kwargs["projection_type"] + builder = _fake_wannier90_builder(kwargs["codes"]) + builder.wannier90["wannier90"]["projections"] = orm.List(list=["Si:s", "Si:p"]) + return builder + + monkeypatch.setattr( + prep_module.Wannier90BandsWorkChain, + "get_builder_from_protocol", + fake_w90_builder, + ) + monkeypatch.setattr( + prep_module.Wannier90OptimizeWorkChain, + "get_builder_from_protocol", + fake_w90_builder, + ) + monkeypatch.setattr( + prep_module.PhBaseWorkChain, + "get_builder_from_protocol", + _fake_ph_builder, + ) + monkeypatch.setattr( + prep_module.EpwBaseWorkChain, + "get_builder_from_protocol", + _fake_epw_builder, + ) + + prepared_inputs = prep_module.build_task_inputs( + codes=codes, + structure=structure, + protocol="fast", + overrides={}, + wannier_projection_type=WannierProjectionType.ANALYTIC, + ) + + assert captured["projection_type"] == WannierProjectionType.ANALYTIC + assert prepared_inputs["w90_bands"]["wannier90"]["wannier90"]["projections"] == [ + "Si:s", + "Si:p", + ] + + +def test_build_task_inputs_rejects_reference_bands_for_analytic_projections( + fixture_code, + generate_structure, + monkeypatch, +): + """Analytic projections should not enter the optimization builder path.""" + codes = { + "pw": fixture_code("quantumespresso.pw"), + "pw2wannier90": fixture_code("quantumespresso.pw2wannier90"), + "wannier90": fixture_code("wannier90.wannier90"), + "ph": fixture_code("quantumespresso.ph"), + "epw": fixture_code("epw.epw"), + } + + monkeypatch.setattr( + prep_module, + "get_protocol_inputs", + lambda protocol=None, overrides=None: { + "pseudo_family": "PseudoDojo/0.5/PBE/SR/standard/upf", + "qpoints_distance": 0.5, + "kpoints_distance_scf": 0.15, + "kpoints_factor_nscf": 2, + "kpoints_force_parity": False, + "w90_bands": { + "scf": {"pw": {"metadata": _metadata_dict()}}, + "nscf": {"pw": {"metadata": _metadata_dict()}}, + "pw2wannier90": {"pw2wannier90": {"metadata": _metadata_dict()}}, + "wannier90": {"wannier90": {"metadata": _metadata_dict()}}, + }, + "ph_base": {"ph": {"metadata": _metadata_dict()}}, + "epw_base": {"options": _metadata_dict()["options"]}, + }, + ) + + with pytest.raises(ValueError, match="Wannier90OptimizeWorkChain"): + prep_module.build_task_inputs( + codes=codes, + structure=generate_structure(), + protocol="fast", + overrides={}, + wannier_projection_type=WannierProjectionType.ANALYTIC, + reference_bands=orm.BandsData(), + ) + + +def test_prep_workgraph_runs_direct_epw_wannierize_with_fake_high_level_workchains( + fixture_code, + fixture_localhost, + generate_remote_data, + generate_structure, + monkeypatch, +): + """The workgraph should mirror the classical direct EPW Wannierization path.""" + codes = { + "pw": fixture_code("quantumespresso.pw"), + "ph": fixture_code("quantumespresso.ph"), + "epw": fixture_code("epw.epw"), + } + structure = generate_structure() + restart_qpoints = _mesh_kpoints([3, 3, 3]) + parent_folder_ph = generate_remote_data(fixture_localhost, "/remote/ph-restart") + + monkeypatch.setattr( + prep_module, + "get_protocol_inputs", + lambda protocol=None, overrides=None: { + "pseudo_family": "PseudoDojo/0.5/PBE/SR/standard/upf", + "qpoints_distance": 0.5, + "kpoints_distance_scf": 0.15, + "kpoints_factor_nscf": 2, + "kpoints_force_parity": False, + "scf": {"pw": {"metadata": _metadata_dict()}}, + "nscf": {"pw": {"metadata": _metadata_dict()}}, + "ph_base": {"ph": {"metadata": _metadata_dict()}}, + "epw_base": { + "options": _metadata_dict()["options"], + "parameters": {"INPUTEPW": {"wannierize": True, "proj": ["Si:s"]}}, + }, + }, + ) + monkeypatch.setattr( + prep_module.PwBaseWorkChain, "get_builder_from_protocol", _fake_pw_builder + ) + monkeypatch.setattr( + prep_module.EpwBaseWorkChain, "get_builder_from_protocol", _fake_epw_builder + ) + monkeypatch.setattr( + prep_module, + "validate_parent_ph_inputs", + lambda _folder, _structure: restart_qpoints, + ) + monkeypatch.setattr( + prep_module, "_apply_socket_overrides", lambda *args, **kwargs: None + ) + monkeypatch.setattr(prep_module, "PwBaseTask", task(FakePwBaseRuntimeWorkChain)) + monkeypatch.setattr(prep_module, "EpwBaseTask", task(FakeEpwBaseRuntimeWorkChain)) + + wg = prep_module.prep( + codes=codes, + structure=structure, + protocol="fast", + overrides={}, + parent_folder_ph=parent_folder_ph, + ) + wg.run() + + assert wg.process is not None + assert wg.process.is_finished_ok + assert "w90_bands" not in {task.name for task in wg.tasks} + + scf_node = wg.tasks["scf"].process + nscf_node = wg.tasks["nscf"].process + epw_node = wg.tasks["epw_base"].process + epw_bands_node = wg.tasks["epw_bands"].process + reciprocal_points_node = wg.tasks["generate_reciprocal_points"].process + + assert scf_node is not None + assert nscf_node is not None + assert epw_node is not None + assert epw_bands_node is not None + assert reciprocal_points_node is not None + + assert "ph_base" not in wg.tasks + assert epw_node.inputs.parent_folder_ph.uuid == parent_folder_ph.uuid + assert ( + epw_node.inputs.parent_folder_nscf.uuid == nscf_node.outputs.remote_folder.uuid + ) + assert "parent_folder_chk" not in epw_node.inputs + assert reciprocal_points_node.outputs.qpoints.uuid == restart_qpoints.uuid + assert epw_node.inputs.qpoints.uuid == restart_qpoints.uuid + assert ( + epw_node.inputs.kpoints.get_kpoints_mesh()[0] + == reciprocal_points_node.outputs.kpoints_nscf.get_kpoints_mesh()[0] + ) + assert ( + epw_bands_node.inputs.parent_folder_epw.uuid + == epw_node.outputs.remote_stash.uuid + ) + assert epw_bands_node.inputs.kfpoints.uuid == epw_bands_node.inputs.qfpoints.uuid + + +def test_prep_workgraph_runs_direct_epw_bands_without_w90( + fixture_code, + generate_structure, + monkeypatch, +): + """Direct EPW Wannierization should still build the bands branch from seekpath.""" + codes = { + "pw": fixture_code("quantumespresso.pw"), + "ph": fixture_code("quantumespresso.ph"), + "epw": fixture_code("epw.epw"), + } + structure = generate_structure() + monkeypatch.setattr( + prep_module, + "get_protocol_inputs", + lambda protocol=None, overrides=None: { + "pseudo_family": "PseudoDojo/0.5/PBE/SR/standard/upf", + "qpoints_distance": 0.5, + "kpoints_distance_scf": 0.15, + "kpoints_factor_nscf": 2, + "kpoints_force_parity": False, + "scf": {"pw": {"metadata": _metadata_dict()}}, + "nscf": {"pw": {"metadata": _metadata_dict()}}, + "ph_base": {"ph": {"metadata": _metadata_dict()}}, + "epw_base": { + "options": _metadata_dict()["options"], + "parameters": {"INPUTEPW": {"wannierize": True, "proj": ["Si:s"]}}, + }, + "epw_bands": {"options": _metadata_dict()["options"]}, + }, + ) + monkeypatch.setattr( + prep_module.PwBaseWorkChain, "get_builder_from_protocol", _fake_pw_builder + ) + monkeypatch.setattr( + prep_module.PhBaseWorkChain, "get_builder_from_protocol", _fake_ph_builder + ) + monkeypatch.setattr( + prep_module.EpwBaseWorkChain, "get_builder_from_protocol", _fake_epw_builder + ) + monkeypatch.setattr( + prep_module, "_apply_socket_overrides", lambda *args, **kwargs: None + ) + monkeypatch.setattr(prep_module, "PwBaseTask", task(FakePwBaseRuntimeWorkChain)) + monkeypatch.setattr(prep_module, "PhBaseTask", task(FakePhBaseRuntimeWorkChain)) + monkeypatch.setattr(prep_module, "EpwBaseTask", task(FakeEpwBaseRuntimeWorkChain)) + + wg = prep_module.prep( + codes=codes, + structure=structure, + protocol="fast", + overrides={}, + ) + wg.run() + + epw_bands_node = wg.tasks["epw_bands"].process + + assert epw_bands_node is not None + assert epw_bands_node.inputs.kfpoints.uuid == epw_bands_node.inputs.qfpoints.uuid + + +def test_prep_workgraph_preserves_nested_scheduler_options_and_codes( + fixture_code, + generate_structure, + monkeypatch, +): + """The prep workgraph should keep nested scheduler options on high-level tasks.""" + codes = { + "pw": fixture_code("quantumespresso.pw"), + "pw2wannier90": fixture_code("quantumespresso.pw2wannier90"), + "wannier90": fixture_code("wannier90.wannier90"), + "ph": fixture_code("quantumespresso.ph"), + "epw": fixture_code("epw.epw"), + } + structure = generate_structure() + + monkeypatch.setattr( + prep_module, + "get_protocol_inputs", + lambda protocol=None, overrides=None: { + "pseudo_family": "PseudoDojo/0.5/PBE/SR/standard/upf", + "qpoints_distance": 0.5, + "kpoints_distance_scf": 0.15, + "kpoints_factor_nscf": 2, + "kpoints_force_parity": False, + "w90_bands": { + "scf": {"pw": {"metadata": _metadata_dict()}}, + "nscf": {"pw": {"metadata": _metadata_dict()}}, + "pw2wannier90": { + "pw2wannier90": { + "metadata": _metadata_dict(num_mpiprocs_per_machine=8) + } + }, + "wannier90": { + "wannier90": { + "metadata": _metadata_dict(num_mpiprocs_per_machine=1) + } + }, + }, + "ph_base": {"ph": {"metadata": _metadata_dict()}}, + "epw_base": {"options": _metadata_dict()["options"]}, + }, + ) + monkeypatch.setattr( + prep_module.Wannier90BandsWorkChain, + "get_builder_from_protocol", + lambda **kwargs: _fake_wannier90_builder(kwargs["codes"]), + ) + monkeypatch.setattr( + prep_module.Wannier90OptimizeWorkChain, + "get_builder_from_protocol", + lambda **kwargs: _fake_wannier90_builder(kwargs["codes"]), + ) + monkeypatch.setattr( + prep_module.PhBaseWorkChain, + "get_builder_from_protocol", + _fake_ph_builder, + ) + monkeypatch.setattr( + prep_module.EpwBaseWorkChain, + "get_builder_from_protocol", + _fake_epw_builder, + ) + + prepared_inputs = prep_module.build_task_inputs( + codes=codes, + structure=structure, + protocol="fast", + overrides={}, + ) + assert prepared_inputs["w90_bands"]["wannier90"]["wannier90"]["parameters"][ + "exclude_bands" + ] == [1, 2] + + wg = prep_module.prep( + codes=codes, structure=structure, protocol="fast", overrides={} + ) + engine_inputs = wg.to_engine_inputs(metadata=None) + restored = WorkGraph.from_dict( + restore_workgraph_data_from_raw_inputs(engine_inputs) + ) + + wannier_task = next(task for task in wg.tasks if task.name == "w90_bands") + phonon_task = next(task for task in wg.tasks if task.name == "ph_base") + epw_task = next(task for task in wg.tasks if task.name == "epw_base") + restored_wannier_task = next( + task for task in restored.tasks if task.name == "w90_bands" + ) + restored_phonon_task = next( + task for task in restored.tasks if task.name == "ph_base" + ) + restored_epw_task = next(task for task in restored.tasks if task.name == "epw_base") + + assert wannier_task.inputs["scf"]["pw"]["code"].value == codes["pw"] + assert ( + wannier_task.inputs["scf"]["pw"]["metadata"]["options"]["account"].value + == "elph" + ) + assert ( + wannier_task.inputs["scf"]["pw"]["metadata"]["options"]["queue_name"].value + == "debug" + ) + assert ( + wannier_task.inputs["nscf"]["pw"]["metadata"]["options"]["account"].value + == "elph" + ) + assert ( + wannier_task.inputs["pw2wannier90"]["pw2wannier90"]["metadata"]["options"][ + "account" + ].value + == "elph" + ) + assert ( + wannier_task.inputs["wannier90"]["wannier90"]["metadata"]["options"][ + "account" + ].value + == "elph" + ) + assert phonon_task.inputs["ph"]["metadata"]["options"]["account"].value == "elph" + assert epw_task.inputs["options"].value["account"] == "elph" + assert epw_task.inputs["options"].value["queue_name"] == "debug" + assert ( + restored_wannier_task.inputs["scf"]["pw"]["metadata"]["options"][ + "account" + ].value + == "elph" + ) + assert ( + restored_wannier_task.inputs["nscf"]["pw"]["metadata"]["options"][ + "account" + ].value + == "elph" + ) + assert ( + restored_wannier_task.inputs["pw2wannier90"]["pw2wannier90"]["metadata"][ + "options" + ]["account"].value + == "elph" + ) + assert ( + restored_wannier_task.inputs["wannier90"]["wannier90"]["metadata"]["options"][ + "account" + ].value + == "elph" + ) + assert ( + restored_phonon_task.inputs["ph"]["metadata"]["options"]["account"].value + == "elph" + ) + assert restored_epw_task.inputs["options"].value["account"] == "elph" + + wannier_engine = engine_inputs["tasks"]["w90_bands"] + phonon_engine = engine_inputs["tasks"]["ph_base"] + epw_engine = engine_inputs["tasks"]["epw_base"] + + assert wannier_engine["scf"]["pw"]["metadata"]["options"]["account"] == "elph" + assert wannier_engine["nscf"]["pw"]["metadata"]["options"]["account"] == "elph" + assert ( + wannier_engine["pw2wannier90"]["pw2wannier90"]["metadata"]["options"]["account"] + == "elph" + ) + assert ( + wannier_engine["wannier90"]["wannier90"]["metadata"]["options"]["account"] + == "elph" + ) + assert phonon_engine["ph"]["metadata"]["options"]["account"] == "elph" + assert epw_engine["options"]["account"] == "elph" + + manager = TaskManager.__new__(TaskManager) + manager.process = AttributeDict({"wg": restored}) + manager.ctx = AttributeDict( + { + "_task_results": { + "generate_reciprocal_points": { + "kpoints_scf": orm.KpointsData(), + "qpoints": orm.KpointsData(), + "kpoints_nscf": orm.KpointsData(), + }, + "prepare_wannier90_runtime_inputs": { + "nscf_kpoints": _create_explicit_kpoints([2, 2, 2]), + "wannier90_kpoints": _create_explicit_kpoints([2, 2, 2]), + "parameters": orm.Dict({"bands_plot": True, "mp_grid": [2, 2, 2]}), + }, + } + } + ) + manager.ctx._task_results["generate_reciprocal_points"][ + "kpoints_scf" + ].set_kpoints_mesh([2, 2, 2]) + manager.ctx._task_results["generate_reciprocal_points"]["qpoints"].set_kpoints_mesh( + [1, 1, 1] + ) + manager.ctx._task_results["generate_reciprocal_points"][ + "kpoints_nscf" + ].set_kpoints_mesh([2, 2, 2]) + + runtime_inputs = manager.get_inputs("w90_bands")["kwargs"] + assert runtime_inputs["scf"]["pw"]["metadata"]["options"]["account"] == "elph" + assert runtime_inputs["nscf"]["pw"]["metadata"]["options"]["account"] == "elph" + assert ( + runtime_inputs["pw2wannier90"]["pw2wannier90"]["metadata"]["options"]["account"] + == "elph" + ) + assert ( + runtime_inputs["wannier90"]["wannier90"]["metadata"]["options"]["account"] + == "elph" + ) + assert len(runtime_inputs["nscf"]["kpoints"].get_kpoints()) == 8 + assert len(runtime_inputs["wannier90"]["wannier90"]["kpoints"].get_kpoints()) == 8 + json.dumps(wg.to_widget_value()) + + +def test_high_level_workgraph_task_keeps_nested_metadata_at_runtime(fixture_code): + """A high-level workchain task should preserve nested scheduler metadata when submitted.""" + from aiida.engine.utils import instantiate_process + from aiida.manage.manager import get_manager + + nested_task = task(NestedMetadataWorkChain) + code = fixture_code("arithmetic.add") + scf_inputs = { + "code": code, + "x": orm.Int(1), + "y": orm.Int(2), + "metadata": { + "options": { + "resources": {"num_machines": 1, "num_mpiprocs_per_machine": 1}, + "max_wallclock_seconds": 3600, + "withmpi": False, + "account": "elph", + "queue_name": "debug", + } + }, + } + + direct_process = instantiate_process( + get_manager().get_runner(), NestedMetadataWorkChain, scf=scf_inputs + ) + assert direct_process.inputs.scf.metadata.options.account == "elph" + assert direct_process.inputs.scf.metadata.options.queue_name == "debug" + assert ( + direct_process.node.base.attributes.get("metadata_inputs")["scf"]["metadata"][ + "options" + ]["account"] + == "elph" + ) + assert ( + direct_process.node.base.attributes.get("metadata_inputs")["scf"]["metadata"][ + "options" + ]["queue_name"] + == "debug" + ) + direct_process.close() + + with WorkGraph(name="nested_runtime_metadata") as wg: + task_node = nested_task( + metadata={"store_provenance": True}, + scf=scf_inputs, + ) + + wg.run() + + process_node = wg.tasks[task_node._task.name].process + assert process_node is not None + metadata_inputs = process_node.base.attributes.get("metadata_inputs") + assert metadata_inputs["scf"]["metadata"]["options"]["account"] == "elph" + assert metadata_inputs["scf"]["metadata"]["options"]["queue_name"] == "debug" + + +def test_high_level_workgraph_task_keeps_double_nested_metadata_at_runtime( + fixture_code, +): + """A high-level workchain task should preserve ``scf.pw.metadata`` style inputs.""" + from aiida.engine.utils import instantiate_process + from aiida.manage.manager import get_manager + + nested_task = task(DoubleNestedMetadataWorkChain) + code = fixture_code("arithmetic.add") + scf_inputs = { + "pw": { + "code": code, + "x": orm.Int(1), + "y": orm.Int(2), + "metadata": { + "options": { + "resources": {"num_machines": 1, "num_mpiprocs_per_machine": 1}, + "max_wallclock_seconds": 3600, + "withmpi": False, + "account": "elph", + "queue_name": "debug", + } + }, + } + } + + direct_process = instantiate_process( + get_manager().get_runner(), + DoubleNestedMetadataWorkChain, + scf=scf_inputs, + ) + direct_meta = direct_process.node.base.attributes.get("metadata_inputs") + assert direct_meta["scf"]["pw"]["metadata"]["options"]["account"] == "elph" + assert direct_meta["scf"]["pw"]["metadata"]["options"]["queue_name"] == "debug" + direct_process.close() + + with WorkGraph(name="double_nested_runtime_metadata") as wg: + task_node = nested_task( + metadata={"store_provenance": True}, + scf=scf_inputs, + ) + + wg.run() + + process_node = wg.tasks[task_node._task.name].process + assert process_node is not None + metadata_inputs = process_node.base.attributes.get("metadata_inputs") + assert metadata_inputs["scf"]["pw"]["metadata"]["options"]["account"] == "elph" + assert metadata_inputs["scf"]["pw"]["metadata"]["options"]["queue_name"] == "debug" + + +def test_workgraph_to_nested_workchains_preserves_account_to_calcjob(fixture_code): + """Nested workchains should keep ``scf.pw.metadata.options.account`` all the way to the CalcJob.""" + nested_task = task(OuterSubmitInnerWorkChain) + code = fixture_code("arithmetic.add") + scf_inputs = { + "pw": { + "code": code, + "x": orm.Int(1), + "y": orm.Int(2), + "metadata": { + "options": { + "resources": {"num_machines": 1, "num_mpiprocs_per_machine": 1}, + "max_wallclock_seconds": 3600, + "withmpi": False, + "account": "elph", + "queue_name": "debug", + } + }, + } + } + + with WorkGraph(name="nested_submit_chain_metadata") as wg: + task_node = nested_task( + metadata={"store_provenance": True}, + scf=scf_inputs, + ) + + wg.run() + + outer_node = wg.tasks[task_node._task.name].process + assert outer_node is not None + outer_meta = outer_node.base.attributes.get("metadata_inputs") + assert outer_meta["scf"]["pw"]["metadata"]["options"]["account"] == "elph" + assert outer_meta["scf"]["pw"]["metadata"]["options"]["queue_name"] == "debug" + + inner_links = outer_node.base.links.get_outgoing(link_type=LinkType.CALL_WORK).all() + assert len(inner_links) == 1 + inner_node = inner_links[0].node + inner_meta = inner_node.base.attributes.get("metadata_inputs") + assert inner_meta["pw"]["metadata"]["options"]["account"] == "elph" + assert inner_meta["pw"]["metadata"]["options"]["queue_name"] == "debug" + + calc_links = inner_node.base.links.get_outgoing(link_type=LinkType.CALL_CALC).all() + assert len(calc_links) == 1 + calc_node = calc_links[0].node + assert calc_node.get_option("account") == "elph" + assert calc_node.get_option("queue_name") == "debug" diff --git a/tests/workgraphs/test_supercon.py b/tests/workgraphs/test_supercon.py new file mode 100644 index 00000000..0ed1bbf6 --- /dev/null +++ b/tests/workgraphs/test_supercon.py @@ -0,0 +1,279 @@ +"""Tests for ``aiida_epw.workgraphs.supercon``.""" + +from types import SimpleNamespace +from importlib import import_module + +import pytest +from aiida import orm +from aiida.common import AttributeDict + +supercon_module = import_module("aiida_epw.workgraphs.supercon") + + +def _metadata_dict( + *, + account: str = "elph", + queue_name: str = "debug", + num_mpiprocs_per_machine: int = 32, + max_wallclock_seconds: int = 1800, +) -> dict: + """Return scheduler metadata used in the workgraph tests.""" + return { + "options": { + "resources": { + "num_machines": 1, + "num_mpiprocs_per_machine": num_mpiprocs_per_machine, + }, + "max_wallclock_seconds": max_wallclock_seconds, + "withmpi": True, + "account": account, + "queue_name": queue_name, + } + } + + +def _fake_epw_builder(code): + """Return a minimal EPW builder-like namespace.""" + return AttributeDict( + { + "code": code, + "parameters": orm.Dict({"INPUTEPW": {"eliashberg": True}}), + "options": orm.Dict(_metadata_dict()["options"]), + "settings": orm.Dict({"CMDLINE": ["-nk", "2"]}), + "parallelization": orm.Dict({"npool": 2}), + "clean_workdir": orm.Bool(False), + "max_iterations": orm.Int(2), + } + ) + + +def test_supercon_workgraph_builds_with_fake_builders( + fixture_code, + generate_kpoints_mesh, + generate_remote_data, + generate_structure, + fixture_localhost, + monkeypatch, +): + """The supercon graph should build and wire the protocol-derived EPW inputs.""" + code = fixture_code("epw.epw") + parent_folder_epw = generate_remote_data(fixture_localhost, "/remote/epw") + structure = generate_structure() + kpoints = generate_kpoints_mesh([6, 6, 6]) + qpoints = generate_kpoints_mesh([3, 3, 3]) + parent_epw = SimpleNamespace( + process_label="EpwBaseWorkChain", + inputs=SimpleNamespace( + structure=structure, + kpoints=kpoints, + qpoints=qpoints, + ), + ) + + monkeypatch.setattr( + supercon_module, + "get_protocol_inputs", + lambda protocol=None, overrides=None: { + "interpolation_distance": [0.2, 0.1], + "kfpoints_factor": 3, + "convergence_threshold": 0.05, + "always_run_final": True, + "epw_interp": { + "settings": {"CMDLINE": ["-nk", "2"]}, + }, + "epw_final_iso": { + "settings": {"CMDLINE": ["-nk", "2"]}, + }, + "epw_final_aniso": { + "settings": {"CMDLINE": ["-nk", "2"]}, + }, + }, + ) + monkeypatch.setattr( + supercon_module.EpwBaseWorkChain, + "get_builder_from_protocol", + lambda code, **kwargs: _fake_epw_builder(code), + ) + + wg = supercon_module.supercon.build( + code=code, + parent_epw=parent_epw, + parent_folder_epw=parent_folder_epw, + protocol="fast", + overrides={}, + ) + + epw_tasks = [task for task in wg.tasks if task.name.startswith("EpwBaseWorkChain")] + assert len(epw_tasks) == 4 + + interp_01 = epw_tasks[0] + interp_02 = epw_tasks[1] + final_iso = epw_tasks[2] + final_aniso = epw_tasks[3] + + assert interp_01.inputs.code.value.uuid == code.uuid + assert interp_02.inputs.code.value.uuid == code.uuid + assert final_iso.inputs.code.value.uuid == code.uuid + assert final_aniso.inputs.code.value.uuid == code.uuid + assert interp_01.inputs.structure.value.uuid == structure.uuid + assert interp_01.inputs.parent_folder_epw.value.uuid == parent_folder_epw.uuid + assert interp_01.inputs.kpoints.value.get_kpoints_mesh()[0] == [6, 6, 6] + assert interp_01.inputs.qpoints.value.get_kpoints_mesh()[0] == [3, 3, 3] + assert interp_01.inputs.kfpoints_factor.value == 3 + assert interp_01.inputs.qfpoints_distance.value == 0.1 + assert interp_02.inputs.qfpoints_distance.value == 0.2 + + +def test_supercon_workgraph_uses_prep_graph_stash_output_as_restart_parent( + fixture_code, + generate_kpoints_mesh, + generate_remote_data, + generate_structure, + fixture_localhost, + monkeypatch, +): + """A prep workgraph parent should contribute `epw_stash` as the EPW restart folder.""" + code = fixture_code("epw.epw") + epw_stash = generate_remote_data(fixture_localhost, "/remote/epw-stash") + structure = generate_structure() + kpoints = generate_kpoints_mesh([6, 6, 6]) + qpoints = generate_kpoints_mesh([3, 3, 3]) + parent_prep_graph = SimpleNamespace( + process_label="WorkGraph", + inputs=SimpleNamespace(structure=structure), + outputs=SimpleNamespace( + epw_stash=epw_stash, + ), + base=SimpleNamespace(), + ) + + monkeypatch.setattr( + supercon_module, + "get_protocol_inputs", + lambda protocol=None, overrides=None: { + "interpolation_distance": [0.1], + "kfpoints_factor": 2, + "always_run_final": False, + "epw_interp": {}, + "epw_final_iso": {}, + "epw_final_aniso": {}, + }, + ) + monkeypatch.setattr( + supercon_module.EpwBaseWorkChain, + "get_builder_from_protocol", + lambda code, **kwargs: _fake_epw_builder(code), + ) + monkeypatch.setattr( + supercon_module, + "_get_prep_reciprocal_points", + lambda _parent: AttributeDict({"kpoints": kpoints, "qpoints": qpoints}), + ) + monkeypatch.setattr( + supercon_module, + "_get_prep_structure", + lambda _parent: structure, + ) + + wg = supercon_module.supercon.build( + code=code, + parent_epw=parent_prep_graph, + protocol="fast", + overrides={}, + ) + + epw_tasks = [task for task in wg.tasks if task.name.startswith("EpwBaseWorkChain")] + assert len(epw_tasks) >= 1 + assert epw_tasks[0].inputs.parent_folder_epw.value.uuid == epw_stash.uuid + assert epw_tasks[0].inputs.kpoints.value.get_kpoints_mesh()[0] == [6, 6, 6] + assert epw_tasks[0].inputs.qpoints.value.get_kpoints_mesh()[0] == [3, 3, 3] + + +def test_supercon_workgraph_falls_back_to_prep_epw_folder_when_stash_missing( + fixture_code, + generate_kpoints_mesh, + generate_remote_data, + generate_structure, + fixture_localhost, + monkeypatch, +): + """A prep workgraph parent should fall back to `epw_folder` when `epw_stash` is unavailable.""" + code = fixture_code("epw.epw") + epw_folder = generate_remote_data(fixture_localhost, "/remote/epw-folder") + structure = generate_structure() + kpoints = generate_kpoints_mesh([6, 6, 6]) + qpoints = generate_kpoints_mesh([3, 3, 3]) + parent_prep_graph = SimpleNamespace( + process_label="WorkGraph", + outputs=SimpleNamespace(epw_folder=epw_folder), + base=SimpleNamespace(), + ) + + monkeypatch.setattr( + supercon_module, + "get_protocol_inputs", + lambda protocol=None, overrides=None: { + "interpolation_distance": [0.1], + "kfpoints_factor": 2, + "always_run_final": False, + "epw_interp": {}, + "epw_final_iso": {}, + "epw_final_aniso": {}, + }, + ) + monkeypatch.setattr( + supercon_module.EpwBaseWorkChain, + "get_builder_from_protocol", + lambda code, **kwargs: _fake_epw_builder(code), + ) + monkeypatch.setattr( + supercon_module, + "_get_prep_reciprocal_points", + lambda _parent: AttributeDict({"kpoints": kpoints, "qpoints": qpoints}), + ) + monkeypatch.setattr( + supercon_module, + "_get_prep_structure", + lambda _parent: structure, + ) + monkeypatch.setattr( + supercon_module, + "_get_prep_epw_base", + lambda _parent: SimpleNamespace( + inputs=SimpleNamespace(clean_workdir=orm.Bool(False)) + ), + ) + + wg = supercon_module.supercon.build( + code=code, + parent_epw=parent_prep_graph, + protocol="fast", + overrides={}, + ) + + epw_tasks = [task for task in wg.tasks if task.name.startswith("EpwBaseWorkChain")] + assert len(epw_tasks) >= 1 + assert epw_tasks[0].inputs.parent_folder_epw.value.uuid == epw_folder.uuid + + +def test_prep_epw_folder_fallback_rejects_cleaned_remote_folder( + generate_remote_data, fixture_localhost +): + """The fallback to `epw_folder` should fail if the prep EPW task cleaned its remote folder.""" + epw_folder = generate_remote_data(fixture_localhost, "/remote/epw-folder") + parent_prep_graph = SimpleNamespace( + outputs=SimpleNamespace(epw_folder=epw_folder), + base=SimpleNamespace(), + ) + + clean_epw_base = SimpleNamespace( + inputs=SimpleNamespace(clean_workdir=orm.Bool(True)) + ) + + original_getter = supercon_module._get_prep_epw_base + try: + supercon_module._get_prep_epw_base = lambda _parent: clean_epw_base + with pytest.raises(ValueError, match="clean_workdir=True"): + supercon_module._get_prep_restart_parent_folder(parent_prep_graph) + finally: + supercon_module._get_prep_epw_base = original_getter