diff --git a/other/materials_designer/workflows/relaxation.ipynb b/other/materials_designer/workflows/relaxation.ipynb index 225088488..23da7cdb7 100644 --- a/other/materials_designer/workflows/relaxation.ipynb +++ b/other/materials_designer/workflows/relaxation.ipynb @@ -113,11 +113,14 @@ "source": [ "# Method parameters\n", "PSEUDOPOTENTIAL_TYPE = \"us\" # \"us\" (ultrasoft), \"nc\" (norm-conserving), \"paw\"\n", - "FUNCTIONAL = \"pbe\" # for gga: \"pbe\", \"pbesol\"; for lda: \"pz\"\n", + "FUNCTIONAL = \"pbe\" # for gga: \"pbe\", \"pbesol\"; for lda: \"pz\"\n", "\n", "# K-grid for the relax step (if not set, KPPRA default is used)\n", "KGRID = None # e.g. [4, 4, 4]\n", "\n", + "# Additional parameters to set in the QE input\n", + "ADDITIONAL_PARAMETERS = None # for example {\"system\": {\"vdw_corr\": \"d3_grimme\"}}\n", + "\n", "# Energy cutoffs\n", "ECUTWFC = 40\n", "ECUTRHO = 200" @@ -365,6 +368,7 @@ "outputs": [], "source": [ "from mat3ra.wode.context.providers import PlanewaveCutoffsContextProvider, PointsGridDataProvider\n", + "from mat3ra.notebooks_utils.workflow import patch_workflow_qe_input\n", "\n", "relax_unit_names = [\"pw_relax\", \"pw_vc-relax\"]\n", "\n", @@ -377,13 +381,17 @@ " swf.set_unit(unit)\n", "\n", "if ECUTWFC is not None:\n", - " cutoffs_context = PlanewaveCutoffsContextProvider(wavefunction=ECUTWFC, density=ECUTRHO, isEdited=True).get_context_item_data()\n", + " cutoffs_context = PlanewaveCutoffsContextProvider(wavefunction=ECUTWFC, density=ECUTRHO,\n", + " isEdited=True).get_context_item_data()\n", " for swf in workflow.subworkflows:\n", " for unit_name in relax_unit_names:\n", " unit = swf.get_unit_by_name(name=unit_name)\n", " if unit:\n", " unit.add_context(cutoffs_context)\n", - " swf.set_unit(unit)" + " swf.set_unit(unit)\n", + "\n", + "if ADDITIONAL_PARAMETERS:\n", + " patch_workflow_qe_input(workflow, ADDITIONAL_PARAMETERS, unit_names=relax_unit_names)\n" ] }, { diff --git a/other/materials_designer/workflows/valence_band_offset.ipynb b/other/materials_designer/workflows/valence_band_offset.ipynb index 9ae7314f1..cdfa1ab7f 100644 --- a/other/materials_designer/workflows/valence_band_offset.ipynb +++ b/other/materials_designer/workflows/valence_band_offset.ipynb @@ -129,9 +129,13 @@ "SCF_KGRID = None # e.g. [8, 8, 1]\n", "KPATH = None # e.g. [{\"point\": \"G\", \"steps\": 20}, {\"point\": \"M\", \"steps\": 20}]\n", "\n", - "# SCF diagonalization and mixing\n", - "DIAGONALIZATION = \"david\" # \"david\" or \"cg\"\n", - "MIXING_BETA = 0.3\n", + "# Set SCF diagonalization and mixing in QE input\n", + "ADDITIONAL_PARAMETERS = {\"electrons\":\n", + " {\n", + " \"diagonalization\": \"david\", # \"david\" or \"cg\"\n", + " \"mixing_beta\": 0.3\n", + " }\n", + "}\n", "\n", "# Energy cutoffs\n", "ECUTWFC = 40\n", @@ -428,16 +432,9 @@ " PointsGridDataProvider,\n", " PointsPathDataProvider,\n", ")\n", + "from mat3ra.notebooks_utils.workflow import patch_workflow_qe_input\n", "\n", - "\n", - "def set_pw_electrons_parameters(unit, diagonalization, mixing_beta):\n", - " unit.replace_in_input_content(r\"diagonalization\\s*=\\s*'[^']*'\", f\"diagonalization = '{diagonalization}'\")\n", - " unit.replace_in_input_content(r\"mixing_beta\\s*=\\s*[-+0-9.eE]+\", f\"mixing_beta = {mixing_beta}\")\n", - " for input in unit.input:\n", - " if isinstance(input, dict) and \"content\" in input:\n", - " input[\"rendered\"] = input[\"content\"]\n", - " return unit\n", - "\n", + "scf_unit_names = [\"pw_scf\", \"pw_bands\"]\n", "\n", "for subworkflow in workflow.subworkflows:\n", " if subworkflow.application.name != APPLICATION_NAME:\n", @@ -455,16 +452,19 @@ " unit.add_context(PointsPathDataProvider(path=KPATH, isEdited=True).get_context_item_data())\n", " subworkflow.set_unit(unit)\n", "\n", - " cutoffs_context = PlanewaveCutoffsContextProvider(\n", - " wavefunction=ECUTWFC, density=ECUTRHO, isEdited=True\n", - " ).get_context_item_data()\n", - " for unit_name in [\"pw_scf\", \"pw_bands\"]:\n", - " if unit_name not in unit_names:\n", - " continue\n", - " unit = subworkflow.get_unit_by_name(name=unit_name)\n", - " unit.add_context(cutoffs_context)\n", - " unit = set_pw_electrons_parameters(unit, DIAGONALIZATION, MIXING_BETA)\n", - " subworkflow.set_unit(unit)\n" + " if ECUTWFC is not None:\n", + " cutoffs_context = PlanewaveCutoffsContextProvider(\n", + " wavefunction=ECUTWFC, density=ECUTRHO, isEdited=True\n", + " ).get_context_item_data()\n", + " for unit_name in scf_unit_names:\n", + " if unit_name not in unit_names:\n", + " continue\n", + " unit = subworkflow.get_unit_by_name(name=unit_name)\n", + " unit.add_context(cutoffs_context)\n", + " subworkflow.set_unit(unit)\n", + "\n", + "if ADDITIONAL_PARAMETERS:\n", + " patch_workflow_qe_input(workflow, ADDITIONAL_PARAMETERS, unit_names=scf_unit_names)\n" ] }, { @@ -683,6 +683,7 @@ "outputs": [], "source": [ "from mat3ra.notebooks_utils.plot import configure_matplotlib_renderer\n", + "\n", "configure_matplotlib_renderer()\n", "import matplotlib.pyplot as plt\n", "import numpy as np\n", diff --git a/pyproject.toml b/pyproject.toml index f1d58e5d1..c3f38ec22 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,7 +43,7 @@ workflows = [ "mat3ra-ade>=2026.5.29.post0", "mat3ra-prode", "mat3ra-ide", - "mat3ra-notebooks-utils[api]", + "mat3ra-notebooks-utils[api]" ] all = [ "mat3ra-notebooks-utils[workflows]" @@ -66,6 +66,8 @@ tests = [ "pytest", "pytest-asyncio", "pytest-cov", + "mat3ra-wode", + "mat3ra-standata", ] docs = [ "mkdocs>=1.4.3", diff --git a/src/py/mat3ra/notebooks_utils/workflow.py b/src/py/mat3ra/notebooks_utils/workflow.py new file mode 100644 index 000000000..d2ada0db6 --- /dev/null +++ b/src/py/mat3ra/notebooks_utils/workflow.py @@ -0,0 +1,49 @@ +import re +from typing import Dict, List, Optional + +from mat3ra.wode import Workflow + + +def patch_workflow_qe_input( + workflow: Workflow, + parameters: Dict[str, Dict[str, object]], + unit_names: List[str], + input_name: Optional[str] = None, +) -> Workflow: + """ + Patch QE inputs across workflow subworkflows for named units. + + Args: + workflow: Workflow with subworkflows. + parameters: Multi-section parameters {section: {key: val}}. + unit_names: List of unit names to patch. + input_name: Optional input file name filter. + + Example: + patch_workflow_qe_input(workflow, {"system": {"vdw_corr": "d3_grimme"}}, ["pw_relax"]) + """ + f90 = lambda value: ( # noqa: E731 + f".{str(value).lower()}." if isinstance(value, bool) else repr(value) if isinstance(value, str) else str(value) + ) + for subworkflow in workflow.subworkflows: + for unit_name in unit_names: + if not (unit := subworkflow.get_unit_by_name(name=unit_name)): + continue + for input_item in getattr(unit, "input", []): + template = input_item.template + if input_name not in (None, template.name): + continue + content = template.content + for section, updates in parameters.items(): + name = section.lstrip("&") + match = re.search(rf"(?ims)(^&{re.escape(name)}\s*\n)(.*?)(^/\s*$)", content) + if not match: + raise ValueError(f"Namelist '&{name.upper()}' not found.") + header, body, footer = match.groups() + for key, value in updates.items(): + line, pattern = f" {key} = {f90(value)}", rf"(?im)^\s*{re.escape(key)}\s*=.*$" + body = re.sub(pattern, line, body) if re.search(pattern, body) else f"{body.rstrip()}\n{line}\n" + content = content[: match.start()] + header + body + footer + content[match.end() :] + template.set_content(content) + subworkflow.set_unit(unit) + return workflow diff --git a/tests/py/unit/test_workflow_utils.py b/tests/py/unit/test_workflow_utils.py new file mode 100644 index 000000000..69aa2f5b1 --- /dev/null +++ b/tests/py/unit/test_workflow_utils.py @@ -0,0 +1,42 @@ +import pytest +from mat3ra.notebooks_utils.workflow import patch_workflow_qe_input +from mat3ra.standata.workflows import WorkflowStandata +from mat3ra.wode.workflows import Workflow + +FIXED_CELL_RELAXATION = "fixed_cell_relaxation.json" +RELAX_UNIT_NAMES = ["pw_relax"] + + +def _relax_workflow(): + config = WorkflowStandata.filter_by_application("espresso").get_by_name_first_match(FIXED_CELL_RELAXATION) + return Workflow.create(config) + + +def _pw_relax_content(workflow): + return workflow.subworkflows[0].get_unit_by_name(name="pw_relax").input[0].template.content + + +@pytest.mark.parametrize( + "parameters,present,absent,error", + [ + ( + {"system": {"vdw_corr": "d3_grimme"}, "electrons": {"mixing_beta": 0.5, "diago_full_acc": True}}, + ["vdw_corr = 'd3_grimme'", "{{ input.IBRAV }}", "mixing_beta = 0.5", "diago_full_acc = .true."], + ["mixing_beta = 0.3"], + None, + ), + ({"FAKESECTION": {"x": 1}}, [], [], "Namelist '&FAKESECTION' not found."), + ], +) +def test_patch_workflow_qe_input(parameters, present, absent, error): + workflow = _relax_workflow() + if error: + with pytest.raises(ValueError, match=error): + patch_workflow_qe_input(workflow, parameters, unit_names=RELAX_UNIT_NAMES) + return + patch_workflow_qe_input(workflow, parameters, unit_names=RELAX_UNIT_NAMES) + content = _pw_relax_content(workflow) + for text in present: + assert text in content + for text in absent: + assert text not in content