From fadb129d50567ec1f59e4ccb535bd592892e96be Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Fri, 5 Jun 2026 19:22:10 -0700 Subject: [PATCH 01/13] feat: patch qe in --- src/py/mat3ra/notebooks_utils/workflow.py | 149 ++++++++++++++++++++++ 1 file changed, 149 insertions(+) create mode 100644 src/py/mat3ra/notebooks_utils/workflow.py diff --git a/src/py/mat3ra/notebooks_utils/workflow.py b/src/py/mat3ra/notebooks_utils/workflow.py new file mode 100644 index 000000000..f77aa11ec --- /dev/null +++ b/src/py/mat3ra/notebooks_utils/workflow.py @@ -0,0 +1,149 @@ +import io +import re +from typing import Any, Dict, List, Mapping, Optional, Tuple + +import f90nml + +JINJA_PATTERN = re.compile(r"\{\{.*?\}\}|\{%-?.*?-?%\}", re.DOTALL) + + +def _normalize_section(section: str) -> str: + return section.lstrip("&").upper() + + +def _protect_jinja(text: str) -> Tuple[str, Dict[str, str]]: + placeholders: Dict[str, str] = {} + + def replacer(match: re.Match[str]) -> str: + key = f"__JINJA_{len(placeholders)}__" + placeholders[key] = match.group(0) + return f"'{key}'" + + return JINJA_PATTERN.sub(replacer, text), placeholders + + +def _restore_jinja(text: str, placeholders: Mapping[str, str]) -> str: + restored = text + for key, original in placeholders.items(): + restored = restored.replace(f"'{key}'", original) + return restored + + +def _extract_namelist_block(content: str, section: str) -> Tuple[str, str, str]: + normalized = _normalize_section(section) + pattern = rf"(?ms)(^&{re.escape(normalized)}\s*\n.*?^/\s*$)" + match = re.search(pattern, content) + if match is None: + raise ValueError(f"Namelist '&{normalized}' not found in input template.") + return content[: match.start()], match.group(0), content[match.end() :] + + +def _restore_namelist_header(block: str, section: str) -> str: + normalized = _normalize_section(section) + return re.sub(r"^&\w+", f"&{normalized}", block, count=1, flags=re.MULTILINE) + + +def _patch_namelist_block(block: str, section: str, parameters: Mapping[str, Any]) -> str: + normalized = _normalize_section(section) + protected, placeholders = _protect_jinja(block) + nml = f90nml.reads(protected) + nml.patch({normalized.lower(): dict(parameters)}) + buffer = io.StringIO() + nml.write(buffer) + patched = _restore_namelist_header(buffer.getvalue(), normalized) + return _restore_jinja(patched, placeholders) + + +def set_content(content: str, section: str, parameters: Mapping[str, Any]) -> str: + before, block, after = _extract_namelist_block(content, section) + patched_block = _patch_namelist_block(block, section, parameters) + return before + patched_block + after + + +def _get_template(input_item): + if isinstance(input_item, dict): + template = input_item.get("template") + if isinstance(template, dict): + return template + return input_item + template = getattr(input_item, "template", None) + if template is not None: + return template + return input_item + + +def _get_input_content(input_item) -> Optional[str]: + template = _get_template(input_item) + if isinstance(template, dict): + return template.get("content") + return getattr(template, "content", None) + + +def _set_input_content(input_item, content: str) -> None: + template = _get_template(input_item) + if isinstance(template, dict): + template["content"] = content + else: + template.content = content + + +def _get_input_name(input_item) -> Optional[str]: + template = _get_template(input_item) + if isinstance(template, dict): + return template.get("name") + return getattr(template, "name", None) + + +def _is_multi_section_parameters(value: Any) -> bool: + if not isinstance(value, Mapping): + return False + return all(isinstance(item, Mapping) for item in value.values()) + + +def _patch_qe_input_single( + unit, + section: str, + parameters: Mapping[str, Any], + input_name: Optional[str] = None, +) -> None: + matching_inputs = 0 + for input_item in getattr(unit, "input", []): + content = _get_input_content(input_item) + if content is None: + continue + if input_name is not None and _get_input_name(input_item) != input_name: + continue + _set_input_content(input_item, set_content(content, section, parameters)) + matching_inputs += 1 + + if matching_inputs == 0: + raise ValueError("No matching input template found for QE patch.") + + +def patch_qe_input( + unit, + section_or_parameters: Any, + parameters: Optional[Mapping[str, Any]] = None, + input_name: Optional[str] = None, +) -> None: + if parameters is None and _is_multi_section_parameters(section_or_parameters): + for section, section_parameters in section_or_parameters.items(): + _patch_qe_input_single(unit, section, section_parameters, input_name) + return + if parameters is None: + raise TypeError("Expected section parameters mapping or a section name with parameters.") + _patch_qe_input_single(unit, section_or_parameters, parameters, input_name) + + +def patch_workflow_qe_input( + workflow, + parameters: Mapping[str, Mapping[str, Any]], + unit_names: List[str], + input_name: Optional[str] = None, +) -> None: + for subworkflow in workflow.subworkflows: + for unit_name in unit_names: + unit = subworkflow.get_unit_by_name(name=unit_name) + if unit: + patch_qe_input(unit, parameters, input_name=input_name) + subworkflow.set_unit(unit) From 5c1eba18396b90719e0a27a617f1f8cf5043242c Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Fri, 5 Jun 2026 19:37:13 -0700 Subject: [PATCH 02/13] update: optimize --- src/py/mat3ra/notebooks_utils/workflow.py | 174 +++++++++------------- 1 file changed, 67 insertions(+), 107 deletions(-) diff --git a/src/py/mat3ra/notebooks_utils/workflow.py b/src/py/mat3ra/notebooks_utils/workflow.py index f77aa11ec..8c739cc0f 100644 --- a/src/py/mat3ra/notebooks_utils/workflow.py +++ b/src/py/mat3ra/notebooks_utils/workflow.py @@ -3,147 +3,107 @@ from typing import Any, Dict, List, Mapping, Optional, Tuple import f90nml - -JINJA_PATTERN = re.compile(r"\{\{.*?\}\}|\{%-?.*?-?%\}", re.DOTALL) - - -def _normalize_section(section: str) -> str: - return section.lstrip("&").upper() +from mat3ra.utils.extra.jinja import JINJA_EXPRESSION_PATTERN def _protect_jinja(text: str) -> Tuple[str, Dict[str, str]]: + """Replace Jinja expressions with f90nml-safe placeholders.""" placeholders: Dict[str, str] = {} + pattern = re.compile(JINJA_EXPRESSION_PATTERN + r"|\{%-?.*?-?%\}", re.DOTALL) - def replacer(match: re.Match[str]) -> str: + def replacer(match): key = f"__JINJA_{len(placeholders)}__" placeholders[key] = match.group(0) return f"'{key}'" - return JINJA_PATTERN.sub(replacer, text), placeholders + return pattern.sub(replacer, text), placeholders def _restore_jinja(text: str, placeholders: Mapping[str, str]) -> str: - restored = text for key, original in placeholders.items(): - restored = restored.replace(f"'{key}'", original) - return restored + text = text.replace(f"'{key}'", original) + return text -def _extract_namelist_block(content: str, section: str) -> Tuple[str, str, str]: - normalized = _normalize_section(section) - pattern = rf"(?ms)(^&{re.escape(normalized)}\s*\n.*?^/\s*$)" - match = re.search(pattern, content) - if match is None: +def set_content(content: str, section: str, parameters: Mapping[str, Any]) -> str: + """Upsert parameters into a QE namelist section, preserving Jinja and case.""" + normalized = section.lstrip("&").upper() + match = re.search(rf"(?ms)(^&{re.escape(normalized)}\s*\n.*?^/\s*$)", content) + if not match: raise ValueError(f"Namelist '&{normalized}' not found in input template.") - return content[: match.start()], match.group(0), content[match.end() :] - - -def _restore_namelist_header(block: str, section: str) -> str: - normalized = _normalize_section(section) - return re.sub(r"^&\w+", f"&{normalized}", block, count=1, flags=re.MULTILINE) - -def _patch_namelist_block(block: str, section: str, parameters: Mapping[str, Any]) -> str: - normalized = _normalize_section(section) + before, block, after = content[: match.start()], match.group(0), content[match.end() :] protected, placeholders = _protect_jinja(block) nml = f90nml.reads(protected) nml.patch({normalized.lower(): dict(parameters)}) buffer = io.StringIO() nml.write(buffer) - patched = _restore_namelist_header(buffer.getvalue(), normalized) - return _restore_jinja(patched, placeholders) + patched = re.sub(r"^&\w+", f"&{normalized}", buffer.getvalue(), count=1, flags=re.MULTILINE) + return before + _restore_jinja(patched, placeholders) + after -def set_content(content: str, section: str, parameters: Mapping[str, Any]) -> str: - before, block, after = _extract_namelist_block(content, section) - patched_block = _patch_namelist_block(block, section, parameters) - return before + patched_block + after - - -def _get_template(input_item): - if isinstance(input_item, dict): - template = input_item.get("template") - if isinstance(template, dict): - return template - return input_item - template = getattr(input_item, "template", None) - if template is not None: - return template - return input_item - - -def _get_input_content(input_item) -> Optional[str]: - template = _get_template(input_item) - if isinstance(template, dict): - return template.get("content") - return getattr(template, "content", None) - - -def _set_input_content(input_item, content: str) -> None: - template = _get_template(input_item) - if isinstance(template, dict): - template["content"] = content +def _get_template_attr(item, attr: str): + """Get attribute from nested template dict/object or flat stub.""" + if isinstance(item, dict): + template = item.get("template", item) + return template.get(attr) if isinstance(template, dict) else None + template = getattr(item, "template", item) + return getattr(template, attr, None) + + +def _set_template_content(item, content: str): + """Set content on nested template dict/object or flat stub.""" + if isinstance(item, dict): + template = item.get("template", item) + (template if isinstance(template, dict) else item)["content"] = content else: - template.content = content - - -def _get_input_name(input_item) -> Optional[str]: - template = _get_template(input_item) - if isinstance(template, dict): - return template.get("name") - return getattr(template, "name", None) - - -def _is_multi_section_parameters(value: Any) -> bool: - if not isinstance(value, Mapping): - return False - return all(isinstance(item, Mapping) for item in value.values()) - - -def _patch_qe_input_single( - unit, - section: str, - parameters: Mapping[str, Any], - input_name: Optional[str] = None, -) -> None: - matching_inputs = 0 - for input_item in getattr(unit, "input", []): - content = _get_input_content(input_item) - if content is None: - continue - if input_name is not None and _get_input_name(input_item) != input_name: - continue - _set_input_content(input_item, set_content(content, section, parameters)) - matching_inputs += 1 - - if matching_inputs == 0: + getattr(item, "template", item).content = content + + +def _patch_unit(unit, section: str, parameters: Mapping[str, Any], input_name: Optional[str]): + """Patch a single namelist section across matching unit inputs.""" + matched = False + for item in getattr(unit, "input", []): + content = _get_template_attr(item, "content") + if content and (not input_name or _get_template_attr(item, "name") == input_name): + _set_template_content(item, set_content(content, section, parameters)) + matched = True + if not matched: raise ValueError("No matching input template found for QE patch.") def patch_qe_input( - unit, - section_or_parameters: Any, - parameters: Optional[Mapping[str, Any]] = None, - input_name: Optional[str] = None, -) -> None: - if parameters is None and _is_multi_section_parameters(section_or_parameters): - for section, section_parameters in section_or_parameters.items(): - _patch_qe_input_single(unit, section, section_parameters, input_name) - return + unit, section_or_params: Any, parameters: Optional[Mapping[str, Any]] = None, input_name: Optional[str] = None +): + """ + Patch QE namelist parameters on a workflow unit. + + Examples: + patch_qe_input(unit, "system", {"vdw_corr": "d3_grimme"}) + patch_qe_input(unit, {"system": {"vdw_corr": "d3_grimme"}}) + """ if parameters is None: - raise TypeError("Expected section parameters mapping or a section name with parameters.") - _patch_qe_input_single(unit, section_or_parameters, parameters, input_name) + if not isinstance(section_or_params, Mapping) or not all( + isinstance(v, Mapping) for v in section_or_params.values() + ): + raise TypeError("Expected section name with parameters or multi-section dict.") + for section, params in section_or_params.items(): + _patch_unit(unit, section, params, input_name) + else: + _patch_unit(unit, section_or_params, parameters, input_name) def patch_workflow_qe_input( - workflow, - parameters: Mapping[str, Mapping[str, Any]], - unit_names: List[str], - input_name: Optional[str] = None, -) -> None: + workflow, parameters: Mapping[str, Mapping[str, Any]], unit_names: List[str], input_name: Optional[str] = None +): + """ + Patch QE inputs across workflow subworkflows for named units. + + Example: + patch_workflow_qe_input(workflow, {"system": {"vdw_corr": "d3_grimme"}}, ["pw_relax"]) + """ for subworkflow in workflow.subworkflows: for unit_name in unit_names: - unit = subworkflow.get_unit_by_name(name=unit_name) - if unit: + if unit := subworkflow.get_unit_by_name(name=unit_name): patch_qe_input(unit, parameters, input_name=input_name) subworkflow.set_unit(unit) From 5ac2d79383b0bd83461dc572cbb3982087250d4d Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Fri, 5 Jun 2026 19:43:48 -0700 Subject: [PATCH 03/13] update: simplify --- src/py/mat3ra/notebooks_utils/workflow.py | 38 ++++++++++++----------- 1 file changed, 20 insertions(+), 18 deletions(-) diff --git a/src/py/mat3ra/notebooks_utils/workflow.py b/src/py/mat3ra/notebooks_utils/workflow.py index 8c739cc0f..d03f48dc3 100644 --- a/src/py/mat3ra/notebooks_utils/workflow.py +++ b/src/py/mat3ra/notebooks_utils/workflow.py @@ -1,6 +1,6 @@ import io import re -from typing import Any, Dict, List, Mapping, Optional, Tuple +from typing import Dict, List, Mapping, Optional, Tuple import f90nml from mat3ra.utils.extra.jinja import JINJA_EXPRESSION_PATTERN @@ -25,7 +25,7 @@ def _restore_jinja(text: str, placeholders: Mapping[str, str]) -> str: return text -def set_content(content: str, section: str, parameters: Mapping[str, Any]) -> str: +def set_content(content: str, section: str, parameters: Mapping[str, object]) -> str: """Upsert parameters into a QE namelist section, preserving Jinja and case.""" normalized = section.lstrip("&").upper() match = re.search(rf"(?ms)(^&{re.escape(normalized)}\s*\n.*?^/\s*$)", content) @@ -60,7 +60,7 @@ def _set_template_content(item, content: str): getattr(item, "template", item).content = content -def _patch_unit(unit, section: str, parameters: Mapping[str, Any], input_name: Optional[str]): +def _patch_unit(unit, section: str, parameters: Mapping[str, object], input_name: Optional[str]): """Patch a single namelist section across matching unit inputs.""" matched = False for item in getattr(unit, "input", []): @@ -73,29 +73,31 @@ def _patch_unit(unit, section: str, parameters: Mapping[str, Any], input_name: O def patch_qe_input( - unit, section_or_params: Any, parameters: Optional[Mapping[str, Any]] = None, input_name: Optional[str] = None -): + unit, + parameters: Mapping[str, Mapping[str, object]], + input_name: Optional[str] = None, +) -> None: """ Patch QE namelist parameters on a workflow unit. - Examples: - patch_qe_input(unit, "system", {"vdw_corr": "d3_grimme"}) + Args: + unit: Execution unit with input templates. + parameters: Namelist parameters as {section: {key: value}}. + input_name: Optional input file name filter. + + Example: patch_qe_input(unit, {"system": {"vdw_corr": "d3_grimme"}}) """ - if parameters is None: - if not isinstance(section_or_params, Mapping) or not all( - isinstance(v, Mapping) for v in section_or_params.values() - ): - raise TypeError("Expected section name with parameters or multi-section dict.") - for section, params in section_or_params.items(): - _patch_unit(unit, section, params, input_name) - else: - _patch_unit(unit, section_or_params, parameters, input_name) + for section, section_parameters in parameters.items(): + _patch_unit(unit, section, section_parameters, input_name) def patch_workflow_qe_input( - workflow, parameters: Mapping[str, Mapping[str, Any]], unit_names: List[str], input_name: Optional[str] = None -): + workflow, + parameters: Mapping[str, Mapping[str, object]], + unit_names: List[str], + input_name: Optional[str] = None, +) -> None: """ Patch QE inputs across workflow subworkflows for named units. From 12d0767d16d332608b3ed22d961f7452d6b5d5a0 Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Fri, 5 Jun 2026 20:19:06 -0700 Subject: [PATCH 04/13] update: simplify --- src/py/mat3ra/notebooks_utils/workflow.py | 105 ++++++++-------------- 1 file changed, 39 insertions(+), 66 deletions(-) diff --git a/src/py/mat3ra/notebooks_utils/workflow.py b/src/py/mat3ra/notebooks_utils/workflow.py index d03f48dc3..b0bc47437 100644 --- a/src/py/mat3ra/notebooks_utils/workflow.py +++ b/src/py/mat3ra/notebooks_utils/workflow.py @@ -1,75 +1,30 @@ -import io import re -from typing import Dict, List, Mapping, Optional, Tuple +from typing import List, Mapping, Optional -import f90nml -from mat3ra.utils.extra.jinja import JINJA_EXPRESSION_PATTERN +def _format_value(value: object) -> str: + """Format Python value as Fortran namelist value.""" + if isinstance(value, bool): + return ".true." if value else ".false." + return f"'{value}'" if isinstance(value, str) else str(value) -def _protect_jinja(text: str) -> Tuple[str, Dict[str, str]]: - """Replace Jinja expressions with f90nml-safe placeholders.""" - placeholders: Dict[str, str] = {} - pattern = re.compile(JINJA_EXPRESSION_PATTERN + r"|\{%-?.*?-?%\}", re.DOTALL) - - def replacer(match): - key = f"__JINJA_{len(placeholders)}__" - placeholders[key] = match.group(0) - return f"'{key}'" - - return pattern.sub(replacer, text), placeholders +def set_content(content: str, section: str, parameters: Mapping[str, object]) -> str: + """Upsert parameters into a QE namelist section.""" + section = section.lstrip("&").upper() + pattern = rf"(?ms)(^&{re.escape(section)}\s*\n)(.*?)(^/\s*$)" + match = re.search(pattern, content) + if not match: + raise ValueError(f"Namelist '&{section}' not found in input template.") -def _restore_jinja(text: str, placeholders: Mapping[str, str]) -> str: - for key, original in placeholders.items(): - text = text.replace(f"'{key}'", original) - return text + before, header, body, footer, after = content[: match.start()], *match.groups(), content[match.end() :] + for param, value in parameters.items(): + line = f" {param} = {_format_value(value)}" + param_pattern = rf"(?m)^\s*{re.escape(param)}\s*=.*$" + body = re.sub(param_pattern, line, body) if re.search(param_pattern, body) else body.rstrip() + f"\n{line}\n" -def set_content(content: str, section: str, parameters: Mapping[str, object]) -> str: - """Upsert parameters into a QE namelist section, preserving Jinja and case.""" - normalized = section.lstrip("&").upper() - match = re.search(rf"(?ms)(^&{re.escape(normalized)}\s*\n.*?^/\s*$)", content) - if not match: - raise ValueError(f"Namelist '&{normalized}' not found in input template.") - - before, block, after = content[: match.start()], match.group(0), content[match.end() :] - protected, placeholders = _protect_jinja(block) - nml = f90nml.reads(protected) - nml.patch({normalized.lower(): dict(parameters)}) - buffer = io.StringIO() - nml.write(buffer) - patched = re.sub(r"^&\w+", f"&{normalized}", buffer.getvalue(), count=1, flags=re.MULTILINE) - return before + _restore_jinja(patched, placeholders) + after - - -def _get_template_attr(item, attr: str): - """Get attribute from nested template dict/object or flat stub.""" - if isinstance(item, dict): - template = item.get("template", item) - return template.get(attr) if isinstance(template, dict) else None - template = getattr(item, "template", item) - return getattr(template, attr, None) - - -def _set_template_content(item, content: str): - """Set content on nested template dict/object or flat stub.""" - if isinstance(item, dict): - template = item.get("template", item) - (template if isinstance(template, dict) else item)["content"] = content - else: - getattr(item, "template", item).content = content - - -def _patch_unit(unit, section: str, parameters: Mapping[str, object], input_name: Optional[str]): - """Patch a single namelist section across matching unit inputs.""" - matched = False - for item in getattr(unit, "input", []): - content = _get_template_attr(item, "content") - if content and (not input_name or _get_template_attr(item, "name") == input_name): - _set_template_content(item, set_content(content, section, parameters)) - matched = True - if not matched: - raise ValueError("No matching input template found for QE patch.") + return before + header + body + footer + after def patch_qe_input( @@ -88,8 +43,20 @@ def patch_qe_input( Example: patch_qe_input(unit, {"system": {"vdw_corr": "d3_grimme"}}) """ - for section, section_parameters in parameters.items(): - _patch_unit(unit, section, section_parameters, input_name) + matched = False + for item in getattr(unit, "input", []): + template = item.template + if input_name and template.name != input_name: + continue + + content = template.content + for section, params in parameters.items(): + content = set_content(content, section, params) + template.set_content(content) + matched = True + + if not matched: + raise ValueError("No matching input template found for QE patch.") def patch_workflow_qe_input( @@ -101,6 +68,12 @@ def patch_workflow_qe_input( """ 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"]) """ From 70de2fac2f73a55018ea720dd8223bbadca1c5eb Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Fri, 5 Jun 2026 20:20:31 -0700 Subject: [PATCH 05/13] update: simplify 2 --- src/py/mat3ra/notebooks_utils/workflow.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/py/mat3ra/notebooks_utils/workflow.py b/src/py/mat3ra/notebooks_utils/workflow.py index b0bc47437..940bd586f 100644 --- a/src/py/mat3ra/notebooks_utils/workflow.py +++ b/src/py/mat3ra/notebooks_utils/workflow.py @@ -2,7 +2,7 @@ from typing import List, Mapping, Optional -def _format_value(value: object) -> str: +def _format_to_f90_value(value: object) -> str: """Format Python value as Fortran namelist value.""" if isinstance(value, bool): return ".true." if value else ".false." @@ -20,7 +20,7 @@ def set_content(content: str, section: str, parameters: Mapping[str, object]) -> before, header, body, footer, after = content[: match.start()], *match.groups(), content[match.end() :] for param, value in parameters.items(): - line = f" {param} = {_format_value(value)}" + line = f" {param} = {_format_to_f90_value(value)}" param_pattern = rf"(?m)^\s*{re.escape(param)}\s*=.*$" body = re.sub(param_pattern, line, body) if re.search(param_pattern, body) else body.rstrip() + f"\n{line}\n" From 282a24b3ca07d79267c5b5dbdfce3f459083e43b Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Fri, 5 Jun 2026 20:33:44 -0700 Subject: [PATCH 06/13] update: use new qe helper --- .../workflows/relaxation.ipynb | 9 +++- .../workflows/valence_band_offset.ipynb | 45 ++++++++++--------- 2 files changed, 31 insertions(+), 23 deletions(-) diff --git a/other/materials_designer/workflows/relaxation.ipynb b/other/materials_designer/workflows/relaxation.ipynb index 225088488..dd1881b2a 100644 --- a/other/materials_designer/workflows/relaxation.ipynb +++ b/other/materials_designer/workflows/relaxation.ipynb @@ -86,6 +86,9 @@ "# Model parameters\n", "MODEL_SUBTYPE = \"gga\" # or \"lda\"\n", "\n", + "# Additional parameters to set in the QE input\n", + "ADDITIONAL_PARAMETERS = None # for example {\"system\": {\"vdw_corr\": \"d3_grimme\"}}\n", + "\n", "# 5. Compute parameters\n", "CLUSTER_NAME = None # specify full or partial name i.e. \"cluster-001\" to select\n", "QUEUE_NAME = QueueName.D\n", @@ -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", @@ -383,7 +387,10 @@ " 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", From 4815fc81351644623ebf3b9cd614c1a68ad4588f Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Fri, 5 Jun 2026 20:34:35 -0700 Subject: [PATCH 07/13] update: use new qe helper 2 --- other/materials_designer/workflows/relaxation.ipynb | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/other/materials_designer/workflows/relaxation.ipynb b/other/materials_designer/workflows/relaxation.ipynb index dd1881b2a..23da7cdb7 100644 --- a/other/materials_designer/workflows/relaxation.ipynb +++ b/other/materials_designer/workflows/relaxation.ipynb @@ -86,9 +86,6 @@ "# Model parameters\n", "MODEL_SUBTYPE = \"gga\" # or \"lda\"\n", "\n", - "# Additional parameters to set in the QE input\n", - "ADDITIONAL_PARAMETERS = None # for example {\"system\": {\"vdw_corr\": \"d3_grimme\"}}\n", - "\n", "# 5. Compute parameters\n", "CLUSTER_NAME = None # specify full or partial name i.e. \"cluster-001\" to select\n", "QUEUE_NAME = QueueName.D\n", @@ -116,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" @@ -381,7 +381,8 @@ " 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", From fd4c557d7bcefc778a90abc2c7daf44bbafb7dc6 Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Fri, 5 Jun 2026 21:01:41 -0700 Subject: [PATCH 08/13] update: ai suggestions --- pyproject.toml | 2 +- src/py/mat3ra/notebooks_utils/workflow.py | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index f1d58e5d1..1698a9b6d 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]" diff --git a/src/py/mat3ra/notebooks_utils/workflow.py b/src/py/mat3ra/notebooks_utils/workflow.py index 940bd586f..165a3d439 100644 --- a/src/py/mat3ra/notebooks_utils/workflow.py +++ b/src/py/mat3ra/notebooks_utils/workflow.py @@ -11,17 +11,17 @@ def _format_to_f90_value(value: object) -> str: def set_content(content: str, section: str, parameters: Mapping[str, object]) -> str: """Upsert parameters into a QE namelist section.""" - section = section.lstrip("&").upper() - pattern = rf"(?ms)(^&{re.escape(section)}\s*\n)(.*?)(^/\s*$)" + section_name = section.lstrip("&") + pattern = rf"(?ims)(^&{re.escape(section_name)}\s*\n)(.*?)(^/\s*$)" match = re.search(pattern, content) if not match: - raise ValueError(f"Namelist '&{section}' not found in input template.") + raise ValueError(f"Namelist '&{section_name.upper()}' not found in input template.") before, header, body, footer, after = content[: match.start()], *match.groups(), content[match.end() :] for param, value in parameters.items(): line = f" {param} = {_format_to_f90_value(value)}" - param_pattern = rf"(?m)^\s*{re.escape(param)}\s*=.*$" + param_pattern = rf"(?im)^\s*{re.escape(param)}\s*=.*$" body = re.sub(param_pattern, line, body) if re.search(param_pattern, body) else body.rstrip() + f"\n{line}\n" return before + header + body + footer + after From 37a34c34dff0713e52e91272f44a42b78234be51 Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Fri, 5 Jun 2026 21:35:52 -0700 Subject: [PATCH 09/13] update: add tests --- tests/py/unit/test_workflow_utils.py | 42 ++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) create mode 100644 tests/py/unit/test_workflow_utils.py diff --git a/tests/py/unit/test_workflow_utils.py b/tests/py/unit/test_workflow_utils.py new file mode 100644 index 000000000..d86471d5c --- /dev/null +++ b/tests/py/unit/test_workflow_utils.py @@ -0,0 +1,42 @@ +import pytest +from mat3ra.notebooks_utils.workflow import set_content + +PW_INPUT = ( + "&system\n VDW_CORR = 'dft-d'\n ibrav = {{ input.IBRAV }}\n/\n" + "&ELECTRONS\n diago_full_acc = .false.\n mixing_beta = 0.3\n/\n" +) + + +@pytest.mark.parametrize( + "content,steps,present,absent,error", + [ + ( + PW_INPUT, + [ + ("system", {"vdw_corr": "d3_grimme"}), + ("electrons", {"mixing_beta": 0.5, "diago_full_acc": True}), + ], + [ + "vdw_corr = 'd3_grimme'", + "ibrav = {{ input.IBRAV }}", + "mixing_beta = 0.5", + "diago_full_acc = .true.", + ], + ["VDW_CORR = 'dft-d'", "mixing_beta = 0.3"], + None, + ), + ("&SYSTEM\n/\n", [("IONS", {"ion_dynamics": "bfgs"})], [], [], "Namelist '&IONS' not found"), + ], +) +def test_set_content(content, steps, present, absent, error): + if error: + with pytest.raises(ValueError, match=error): + set_content(content, *steps[0]) + return + result = content + for section, parameters in steps: + result = set_content(result, section, parameters) + for text in present: + assert text in result + for text in absent: + assert text not in result From a04f609de8f89e03e7899ab28117eec0d220da43 Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Tue, 9 Jun 2026 14:34:14 -0700 Subject: [PATCH 10/13] update: compress logic with LLM --- src/py/mat3ra/notebooks_utils/workflow.py | 90 +++++++---------------- 1 file changed, 27 insertions(+), 63 deletions(-) diff --git a/src/py/mat3ra/notebooks_utils/workflow.py b/src/py/mat3ra/notebooks_utils/workflow.py index 165a3d439..997a8f8a3 100644 --- a/src/py/mat3ra/notebooks_utils/workflow.py +++ b/src/py/mat3ra/notebooks_utils/workflow.py @@ -1,70 +1,15 @@ import re -from typing import List, Mapping, Optional +from typing import Dict, List, Optional - -def _format_to_f90_value(value: object) -> str: - """Format Python value as Fortran namelist value.""" - if isinstance(value, bool): - return ".true." if value else ".false." - return f"'{value}'" if isinstance(value, str) else str(value) - - -def set_content(content: str, section: str, parameters: Mapping[str, object]) -> str: - """Upsert parameters into a QE namelist section.""" - section_name = section.lstrip("&") - pattern = rf"(?ims)(^&{re.escape(section_name)}\s*\n)(.*?)(^/\s*$)" - match = re.search(pattern, content) - if not match: - raise ValueError(f"Namelist '&{section_name.upper()}' not found in input template.") - - before, header, body, footer, after = content[: match.start()], *match.groups(), content[match.end() :] - - for param, value in parameters.items(): - line = f" {param} = {_format_to_f90_value(value)}" - param_pattern = rf"(?im)^\s*{re.escape(param)}\s*=.*$" - body = re.sub(param_pattern, line, body) if re.search(param_pattern, body) else body.rstrip() + f"\n{line}\n" - - return before + header + body + footer + after - - -def patch_qe_input( - unit, - parameters: Mapping[str, Mapping[str, object]], - input_name: Optional[str] = None, -) -> None: - """ - Patch QE namelist parameters on a workflow unit. - - Args: - unit: Execution unit with input templates. - parameters: Namelist parameters as {section: {key: value}}. - input_name: Optional input file name filter. - - Example: - patch_qe_input(unit, {"system": {"vdw_corr": "d3_grimme"}}) - """ - matched = False - for item in getattr(unit, "input", []): - template = item.template - if input_name and template.name != input_name: - continue - - content = template.content - for section, params in parameters.items(): - content = set_content(content, section, params) - template.set_content(content) - matched = True - - if not matched: - raise ValueError("No matching input template found for QE patch.") +from mat3ra.wode import Workflow def patch_workflow_qe_input( - workflow, - parameters: Mapping[str, Mapping[str, object]], + workflow: Workflow, + parameters: Dict[str, Dict[str, object]], unit_names: List[str], input_name: Optional[str] = None, -) -> None: +): """ Patch QE inputs across workflow subworkflows for named units. @@ -77,8 +22,27 @@ def patch_workflow_qe_input( 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 unit := subworkflow.get_unit_by_name(name=unit_name): - patch_qe_input(unit, parameters, input_name=input_name) - subworkflow.set_unit(unit) + 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) From d4f6e4dd5e1b7de53caeab0b0686d45db6df7ef6 Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Tue, 9 Jun 2026 14:35:10 -0700 Subject: [PATCH 11/13] update: test patch qe input --- tests/py/unit/test_workflow_utils.py | 52 ++++++++++++++-------------- 1 file changed, 26 insertions(+), 26 deletions(-) diff --git a/tests/py/unit/test_workflow_utils.py b/tests/py/unit/test_workflow_utils.py index d86471d5c..69aa2f5b1 100644 --- a/tests/py/unit/test_workflow_utils.py +++ b/tests/py/unit/test_workflow_utils.py @@ -1,42 +1,42 @@ import pytest -from mat3ra.notebooks_utils.workflow import set_content +from mat3ra.notebooks_utils.workflow import patch_workflow_qe_input +from mat3ra.standata.workflows import WorkflowStandata +from mat3ra.wode.workflows import Workflow -PW_INPUT = ( - "&system\n VDW_CORR = 'dft-d'\n ibrav = {{ input.IBRAV }}\n/\n" - "&ELECTRONS\n diago_full_acc = .false.\n mixing_beta = 0.3\n/\n" -) +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( - "content,steps,present,absent,error", + "parameters,present,absent,error", [ ( - PW_INPUT, - [ - ("system", {"vdw_corr": "d3_grimme"}), - ("electrons", {"mixing_beta": 0.5, "diago_full_acc": True}), - ], - [ - "vdw_corr = 'd3_grimme'", - "ibrav = {{ input.IBRAV }}", - "mixing_beta = 0.5", - "diago_full_acc = .true.", - ], - ["VDW_CORR = 'dft-d'", "mixing_beta = 0.3"], + {"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, ), - ("&SYSTEM\n/\n", [("IONS", {"ion_dynamics": "bfgs"})], [], [], "Namelist '&IONS' not found"), + ({"FAKESECTION": {"x": 1}}, [], [], "Namelist '&FAKESECTION' not found."), ], ) -def test_set_content(content, steps, present, absent, error): +def test_patch_workflow_qe_input(parameters, present, absent, error): + workflow = _relax_workflow() if error: with pytest.raises(ValueError, match=error): - set_content(content, *steps[0]) + patch_workflow_qe_input(workflow, parameters, unit_names=RELAX_UNIT_NAMES) return - result = content - for section, parameters in steps: - result = set_content(result, section, parameters) + patch_workflow_qe_input(workflow, parameters, unit_names=RELAX_UNIT_NAMES) + content = _pw_relax_content(workflow) for text in present: - assert text in result + assert text in content for text in absent: - assert text not in result + assert text not in content From e14e641e1269289a843073facc2f84c990005203 Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Tue, 9 Jun 2026 14:38:32 -0700 Subject: [PATCH 12/13] chore: wode + standata --- pyproject.toml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 1698a9b6d..c3f38ec22 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -66,6 +66,8 @@ tests = [ "pytest", "pytest-asyncio", "pytest-cov", + "mat3ra-wode", + "mat3ra-standata", ] docs = [ "mkdocs>=1.4.3", From 87d296fc46a8aacfc975f870d0e60cfbe9007de4 Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Tue, 9 Jun 2026 16:07:18 -0700 Subject: [PATCH 13/13] update: return WF --- src/py/mat3ra/notebooks_utils/workflow.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/py/mat3ra/notebooks_utils/workflow.py b/src/py/mat3ra/notebooks_utils/workflow.py index 997a8f8a3..d2ada0db6 100644 --- a/src/py/mat3ra/notebooks_utils/workflow.py +++ b/src/py/mat3ra/notebooks_utils/workflow.py @@ -9,7 +9,7 @@ def patch_workflow_qe_input( 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. @@ -46,3 +46,4 @@ def patch_workflow_qe_input( content = content[: match.start()] + header + body + footer + content[match.end() :] template.set_content(content) subworkflow.set_unit(unit) + return workflow