feature/SOF 7924

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"
    ]
   },
   {

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",

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",

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

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