Add a QboxContext helper class

Author: p-zubieta

pysages/backends/contexts.py

@@ -6,12 +6,26 @@
 class to hold the simulation data.
 """
 
+import weakref
+from dataclasses import dataclass
 from importlib import import_module
-
-from pysages.typing import Any, Callable, JaxArray, NamedTuple, Optional
-from pysages.utils import is_file
+from pathlib import Path
+from xml.etree import ElementTree as et
+
+from pysages.typing import (
+    Any,
+    Callable,
+    Iterable,
+    JaxArray,
+    NamedTuple,
+    Optional,
+    Union,
+)
+from pysages.utils import dispatch, is_file, splitlines
 
 JaxMDState = Any
+QboxInstance = Any
+XMLElement = et.Element
 
 
 class JaxMDContextState(NamedTuple):
@@ -63,44 +77,118 @@ class JaxMDContext(NamedTuple):
     dt: float
 
 
+@dataclass(frozen=True)
 class QboxContextGenerator:
     """
     Provides an interface for setting up Qbox-backed simulations.
 
     Arguments
     ---------
+
     launch_command: str
         Specifies the command that will be used to run Qbox in interactive mode,
         e.g. `qb` or `mpirun -n 4 qb`.
 
-    input_script: str
-        Path to the Qbox input script.
+    script: str
+        File or multile string with the Qbox input script.
+
+    nitscf: Optional[int]
+        Same as Qbox's `run` command parameter. The maximum number of self-consistent
+        iterations.
 
-    output_filename: Union[Path, str]
+    nite: Optional[int]
+        Same as Qbox's `run` command parameter. The number of electronic iterations
+        performed between updates of the charge density.
+
+    logfile: Union[Path, str]
         Name for the output file. It must not exist on the working directory.
         Defaults to `qb.r`.
     """
 
-    def __init__(self, launch_command, input_script, output_filename="qb.r"):
-        self.cmd = launch_command
-        self.script = input_script
-        self.logfile = output_filename
+    # NOTE: we leave `niter` as non-configurable for now.
+    # niter: int
+    #     Same as Qbox's `run` command parameter. The number of steps during which atomic
+    #     positions are updated. Defaults to 1.
 
-    def __call__(self, **kwargs):
-        if not is_file(self.script):
-            raise FileNotFoundError(f"Unable to find or open {self.script}")
+    launch_command: str
+    script: str
+    nitscf: Optional[int] = None
+    nite: Optional[int] = None
+    logfile: Union[Path, str] = Path("qb.r")
 
+    def __call__(self, **kwargs):
         if is_file(self.logfile):
-            msg = f"Delete {self.logfile} or choose a different output file name"
+            msg = f"Rename or delete {self.logfile}, or choose a different log file name"
             raise FileExistsError(msg)
 
-        pexpect = import_module("pexpect")
-
-        qb = pexpect.spawn(self.cmd)
-        qb.logfile_read = open(self.logfile, "wb")
-        qb.expect(r"\[qbox\] ")
-
-        qb.sendline(self.script)
-        qb.expect(r"\[qbox\] ")
-
-        return qb
+        return QboxContext(
+            self.launch_command, self.script, self.logfile, 1, self.nitscf, self.nite
+        )
+
+
+@dataclass(frozen=True)
+class QboxContext:
+    instance: QboxInstance
+    niter: int
+    nitscf: Optional[int]
+    nite: Optional[int]
+    species_masses: dict
+    initial_state: XMLElement
+    state: XMLElement
+
+    @dispatch
+    def __init__(
+        self, launch_command: str, script: str, logfile: Union[Path, str], niter, nitscf, nite
+    ):
+        pexpect = import_module("pexpect.popen_spawn")
+
+        def finalize(qb):
+            if not qb.flag_eof:
+                qb.sendline("quit")
+                qb.expect(pexpect.EOF)
+
+        qb = pexpect.PopenSpawn(launch_command)
+        weakref.finalize(qb, lambda: finalize(qb))
+        qb.logfile_read = open(logfile, "wb")
+        i = qb.expect([r"\[qbox\] ", pexpect.EOF])
+
+        if i == 1:  # EOF was written to the log file
+            preamble = (
+                "The command:\n\n  "
+                f"{launch_command}\n\n"
+                "for running Qbox failed, it returned the following:\n\n"
+            )
+            raise ChildProcessError(preamble + qb.before.decode())
+
+        super().__setattr__("instance", qb)
+        super().__setattr__("niter", niter)
+        super().__setattr__("nitscf", "" if nitscf is None else nitscf)
+        super().__setattr__("nite", "" if nite is None else nite)
+
+        initial_state = qb.before
+        state = self.process_input(script)  # sets `self.state`
+
+        if self.state.find("error") is not None:
+            try:
+                qb.expect(pexpect.EOF, timeout=3)
+            finally:
+                raise ChildProcessError("Qbox encountered the following error:\n" + state.decode())
+
+        initial_state += state + b"\n</fpmd:simulation>"
+        super().__setattr__("initial_state", et.fromstring(initial_state))
+
+        k = 1822.888486  # to convert amu to atomic units
+        species = self.initial_state.iter("species")
+        species_masses = {s.attrib["name"]: k * float(s.find("mass").text) for s in species}
+        super().__setattr__("species_masses", species_masses)
+
+    def process_input(self, entries: Union[str, Iterable[str]], target=r"\[qbox\] ", timeout=None):
+        qb = self.instance
+        state = b""
+        for entry in splitlines(entries):
+            qb.sendline(entry)
+            qb.expect(target, timeout=timeout)
+            state += qb.before
+        # We add tags to ensure that the state corresponds to a valid xml section
+        super().__setattr__("state", et.fromstring(b"<root>\n" + state + b"\n</root>"))
+        return state