Draft: minimal interface for Quantum ESPRESSO integration

src/autoplex/misc/qe/__init__.py

@@ -0,0 +1,20 @@
+from .jobs import QeStaticMaker
+from .run import run_qe_static
+from .schema import (
+    InputDoc,
+    OutputDoc,
+    QeKpointsSettings,
+    QeRunSettings,
+    TaskDoc,
+)
+from .utils import QeStaticInputGenerator
+
+__all__ = [
+    "QEStaticMaker",
+    "QeInputSet",
+    "QeKpointsSettings",
+    "QeRunResult",
+    "QeRunSettings",
+    "QeStaticInputGenerator",
+    "run_qe_static",
+]

src/autoplex/misc/qe/jobs.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+
+from ase import Atoms
+from jobflow import Flow, Maker
+from pymatgen.core import Structure
+
+from .run import run_qe_static
+from .schema import QeKpointsSettings, QeRunSettings
+from .utils import QeStaticInputGenerator
+
+
+@dataclass
+class QeStaticMaker(Maker):
+    """
+    StaticMaker for Quantum ESPRESSO:
+    - assemble and write one .pwi per structure using InputGenerator;
+    - create a `run_qe_static` job for each input;
+    - assemble flow with all jobs.
+
+    Parameters
+    ----------
+    name : str
+        Name of the Flow.
+    command : str
+        Command to execute QE (e.g. "pw.x" or "mpirun -np 4 pw.x -nk 2").
+    workdir : str | None
+        Directory used to write input/output files. Default: "<cwd>/qe_static".
+    run_settings : QeRunSettings | None
+        Update namelists (&control, &system, &electrons).
+    kpoints : QeKpointsSettings | None
+        Set up for K_POINTS if it is not contained in the template.
+    pseudo : dict[str, str] | None
+        Dictionary of atomic symbols and corresponding pseudopotential files.
+    """
+
+    name: str = "qe_static"
+    command: str = "pw.x"
+    workdir: str | None = None
+    run_settings: QeRunSettings | None = None
+    kpoints: QeKpointsSettings | None = None
+    pseudo: dict[str, str] | None = None
+
+    def make(
+        self,
+        structures: Atoms | list[Atoms] | Structure | list[Structure] | str | list[str],
+    ) -> Flow:
+        """
+        Create a Flow to run static SCF calculations with QE for given structures.
+
+        Parameters
+        ----------
+        structures : Atoms | list[Atoms] | Structure | list[Structure] | str | list[str]
+            Single or list of ASE Atoms, pymatgen Structures, or ASE-readable files.
+
+        Returns
+        -------
+        Flow
+            A jobflow Flow with one `run_qe_static` job per structure.
+        """
+        workdir = self.workdir or os.path.join(os.getcwd(), "qe_static")
+        os.makedirs(workdir, exist_ok=True)
+
+        # Generate one input per structure
+        generator = QeStaticInputGenerator(
+            run_settings=self.run_settings or QeRunSettings(),
+            kpoints=self.kpoints or QeKpointsSettings(),
+            pseudo=self.pseudo or {},
+        )
+        input_sets = generator.generate_for_structures(
+            structures=structures, workdir=workdir, seed_prefix="structure"
+        )
+
+        # If single structure, generate one job
+        if len(input_sets) == 1:
+            job = run_qe_static(input_sets[0], command=self.command)
+            job.name = self.name
+            return job
+
+        # Else create one SCF job per structure and assemble flow
+        jobs = []
+        tasks = []
+        for i, inp in enumerate(input_sets):
+            j = run_qe_static(inp, command=self.command)
+            j.name = f"{self.name}_{i}"
+            jobs.append(j)
+            tasks.append(j.output)
+
+        return Flow(jobs=jobs, output=tasks, name=self.name)

src/autoplex/misc/qe/run.py

@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+import logging
+import os
+import re
+import subprocess
+
+from ase.io import read
+from ase.units import GPa
+from jobflow import job
+from pymatgen.io.ase import AseAtomsAdaptor
+
+from .schema import InputDoc, OutputDoc, TaskDoc
+
+logger = logging.getLogger(__name__)
+
+
+_ENERGY_RE = re.compile(r"!\s+total energy\s+=\s+([-\d\.Ee+]+)\s+Ry")
+
+
+def _parse_total_energy_ev(pwo_path: str) -> float | None:
+    """
+    Extract and return total energy (eV) if found in QE output (.pwo)
+
+    Parameters
+    ----------
+    pwo_path : str
+        Path to QE output file (.pwo)
+
+    Returns
+    -------
+    float | None
+        Total energy in eV if found, else None
+    """
+    if not os.path.exists(pwo_path):
+        return None
+
+    try:
+        with open(pwo_path, errors="ignore") as fh:
+            for line in fh:
+                m = _ENERGY_RE.search(line)
+                if m:
+                    # convert energy in eV
+                    ry = float(m.group(1))
+                    return ry * 13.605693009
+    except Exception:
+        return None
+    return None
+
+
+@job
+def run_qe_static(input: InputDoc, command: str) -> TaskDoc:
+    """
+    Execute single QE SCF static calculation from .pwi file.
+    Parse output .pwo file to extract total energy, forces, stress, and final structure.
+
+    Parameters
+    ----------
+    input : InputDoc
+        Input document containing paths and settings for the QE run.
+    command : str
+        Command to execute QE (e.g. "pw.x" or "mpirun -np 4 pw.x -nk 2").
+
+    Returns
+    -------
+    TaskDoc
+        Document containing input, output, and metadata of the QE run.
+    """
+    pwi_path = input.pwi_path
+    pwo_path = pwi_path.replace(".pwi", ".pwo")
+    # Assemble pwscf run command e.g. "pw.x < input.pwi >> input.pwo"
+    run_cmd = f"{command} < {pwi_path} >> {pwo_path}"
+
+    success = False
+    try:
+        subprocess.run(run_cmd, shell=True, check=True, executable="/bin/bash")
+    except subprocess.CalledProcessError as exc:
+        logger.error("QE failed for %s: %s", pwi_path, exc)
+
+    # # Manual parse of total energy in eV from .pwo
+    # energy_ev = _parse_total_energy_ev(pwo_path)
+
+    # Parse with ASE
+    atoms = read(pwo_path)
+    energy_ev = atoms.get_total_energy()
+    forces_evA = atoms.get_forces()
+    stress_kbar = atoms.get_stress(voigt=False) * (-10 / GPa)
+    final_structure = AseAtomsAdaptor().get_structure(atoms)
+
+    output = OutputDoc(
+        energy=energy_ev,
+        forces=forces_evA.tolist(),
+        stress=stress_kbar.tolist(),
+        energy_per_atom=energy_ev / len(atoms) if energy_ev is not None else None,
+    )
+
+    return TaskDoc(
+        structure=final_structure,
+        dir_name=os.path.dirname(pwi_path),
+        task_label="qe_scf",
+        input=input,
+        output=output,
+    )

src/autoplex/misc/qe/schema.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+from emmet.core.math import Matrix3D, Vector3D
+from emmet.core.structure import StructureMetadata
+from pydantic import BaseModel, Field, field_validator
+from pymatgen.core import Structure
+
+
+class QeRunSettings(BaseModel):
+    """
+    Set QE namelists for static calculation SCF
+    Standard namelists: CONTROL, SYSTEM, ELECTRONS
+    """
+
+    control: dict[str, object] = Field(default_factory=dict)
+    system: dict[str, object] = Field(default_factory=dict)
+    electrons: dict[str, object] = Field(default_factory=dict)
+
+    @field_validator("control", "system", "electrons")
+    @classmethod
+    def _lowercase_keys(cls, v: dict[str, object]) -> dict[str, object]:
+        # default lowercase keywords
+        return {str(k).lower(): v[k] for k in v}
+
+
+class QeKpointsSettings(BaseModel):
+    """
+    K-points: use k-space resoultion with automatic Monkhorst-Pack grid
+    k-points offset as in Quantum ESPRESSO manual: 0: False, 1: True
+    """
+
+    kspace_resolution: float | None = None  # angstrom^-1
+    koffset: list[bool] = Field(default_factory=lambda: [False, False, False])
+
+    @field_validator("koffset")
+    @classmethod
+    def _len3(cls, v: list[bool]) -> list[bool]:
+        if len(v) != 3:
+            raise ValueError("koffset must be a list of 3 booleans.")
+        return v
+
+
+class InputDoc(BaseModel):
+    """
+    Inputs and contexts used to run the static SCF job
+    """
+
+    workdir: str
+    pwi_path: str
+    seed: str
+    run_settings: QeRunSettings = Field(
+        None, description="QE namelist section with: &control, &system, &electrons"
+    )
+    pseudo: dict[str, str] = Field(
+        None,
+        description="Dictionary of atomic symbols and corresponding pseudopotential files.",
+    )
+    kpoints: QeKpointsSettings = Field(None, description="QE K_POINTS settings")
+
+
+class OutputDoc(BaseModel):
+    """
+    The outputs of this jobs
+    """
+
+    energy: float | None = Field(None, description="Total energy in units of eV.")
+
+    energy_per_atom: float | None = Field(
+        None,
+        description="Energy per atom of the final molecule or structure "
+        "in units of eV/atom.",
+    )
+
+    forces: list[Vector3D] | None = Field(
+        None,
+        description=(
+            "The force on each atom in units of eV/A for the final molecule "
+            "or structure."
+        ),
+    )
+
+    # NOTE: units for stresses were converted to kbar (* -10 from standard output)
+    #       to comply with MP convention
+    stress: Matrix3D | None = Field(
+        None, description="The stress on the cell in units of kbar."
+    )
+
+
+class TaskDoc(StructureMetadata):
+    """Document containing information on structure manipulation using Quantum ESPRESSO."""
+
+    structure: Structure = Field(
+        None, description="Final output structure from the task"
+    )
+
+    input: InputDoc = Field(
+        None, description="The input information used to run this job."
+    )
+
+    output: OutputDoc = Field(None, description="The output information from this job.")
+
+    task_label: str = Field(
+        None,
+        description="Description of the QE task (e.g., static, relax)",
+    )
+
+    dir_name: str | None = Field(
+        None, description="Directory where the QE calculations are performed."
+    )