Run dftbplus_thermo in an optional isolated working directory

ThermoScreening/thermo/api.py

@@ -1,4 +1,6 @@
 import logging
+import os
+from contextlib import contextmanager
 from pathlib import Path
 
 import numpy as np
@@ -494,11 +496,36 @@ def execute(input_file: str) -> Thermo:
         engine=engine,
     )
 
+@contextmanager
+def _run_in_directory(directory):
+    """
+    Run the enclosed block in ``directory`` (created if needed), restoring the
+    previous working directory afterwards. A no-op when ``directory`` is None.
+
+    The DFTB+ steps read and write fixed filenames (geo_opt.gen, hessian.out,
+    vibrations.tag, ...) in the current directory, so giving each job its own
+    directory keeps batch runs from clobbering each other.
+    """
+    if directory is None:
+        yield
+        return
+
+    previous = Path.cwd()
+    directory = Path(directory)
+    directory.mkdir(parents=True, exist_ok=True)
+    os.chdir(directory)
+    try:
+        yield
+    finally:
+        os.chdir(previous)
+
+
 def dftbplus_thermo(
-        atoms, 
+        atoms,
         temperature=298.15,
         pressure=101325,
         charge=0.0,
+        directory=None,
         **kwargs
     ):
     """
@@ -514,6 +541,10 @@ def dftbplus_thermo(
         The pressure in Pa. Default is 101325.
     charge : float
         The system charge. Default is 0.0.
+    directory : str, optional
+        Working directory to run the DFTB+ steps in (created if needed). The
+        geometry/Hessian/modes files are written here, so separate jobs can run
+        without clobbering each other. Defaults to the current directory.
 
     Other Parameters
     ----------------
@@ -525,29 +556,30 @@ def dftbplus_thermo(
     Thermo
         The thermo calculation object.
     """
-
-    # run geometry optimization
-    geoopt = Geoopt(atoms=atoms, charge=charge, **kwargs)
-    potential_energy = geoopt.potential_energy()
-    optimized_atoms = geoopt.read()
-
-    # run hessian calculation
-    Hessian(atoms=optimized_atoms, charge=charge, **kwargs)
-
-    # run normal mode calculation
-    modes = Modes()
-    frequencies = modes.wave_numbers
-
-    # run thermo calculation on the optimized geometry directly (DFTB+ writes
-    # geo_opt.gen, not geo_opt.xyz, so avoid the disk round-trip entirely)
-    thermo = run_thermo(
-        frequencies,
-        atoms=optimized_atoms,
-        temperature=temperature,
-        pressure=pressure,
-        energy=potential_energy,
-        engine='dftb+',
-        charge=charge,
-    )
+
+    with _run_in_directory(directory):
+        # run geometry optimization
+        geoopt = Geoopt(atoms=atoms, charge=charge, **kwargs)
+        potential_energy = geoopt.potential_energy()
+        optimized_atoms = geoopt.read()
+
+        # run hessian calculation
+        Hessian(atoms=optimized_atoms, charge=charge, **kwargs)
+
+        # run normal mode calculation
+        modes = Modes()
+        frequencies = modes.wave_numbers
+
+        # run thermo calculation on the optimized geometry directly (DFTB+ writes
+        # geo_opt.gen, not geo_opt.xyz, so avoid the disk round-trip entirely)
+        thermo = run_thermo(
+            frequencies,
+            atoms=optimized_atoms,
+            temperature=temperature,
+            pressure=pressure,
+            energy=potential_energy,
+            engine='dftb+',
+            charge=charge,
+        )
 
     return thermo

tests/thermo/test_api.py

@@ -299,5 +299,86 @@ def test_run_thermo_ase_atoms_matches_file_path(self):
             from_file.total_gibbs_free_energy("H")
         )
 
+
+def test_run_in_directory_isolates_and_restores(tmp_path):
+    from ThermoScreening.thermo.api import _run_in_directory
+
+    start = Path.cwd()
+    job = tmp_path / "job1"
+
+    with _run_in_directory(str(job)):
+        assert Path.cwd().resolve() == job.resolve()
+        Path("artifact.txt").write_text("x", encoding="utf-8")
+
+    assert Path.cwd() == start
+    assert (job / "artifact.txt").exists()
+
+
+def test_run_in_directory_restores_on_error(tmp_path):
+    from ThermoScreening.thermo.api import _run_in_directory
+
+    start = Path.cwd()
+
+    with pytest.raises(RuntimeError):
+        with _run_in_directory(str(tmp_path / "job2")):
+            raise RuntimeError("boom")
+
+    assert Path.cwd() == start
+
+
+def test_run_in_directory_none_is_noop():
+    from ThermoScreening.thermo.api import _run_in_directory
+
+    start = Path.cwd()
+    with _run_in_directory(None):
+        assert Path.cwd() == start
+    assert Path.cwd() == start
+
+
+def test_dftbplus_thermo_runs_pipeline_in_directory(monkeypatch, tmp_path):
+    # Drive dftbplus_thermo with the DFTB+ steps mocked, so the directory wiring
+    # is exercised without the binaries.
+    import ThermoScreening.thermo.api as api
+
+    seen = {}
+
+    class FakeGeoopt:
+        def __init__(self, atoms, charge, **kwargs):
+            seen["cwd_during"] = Path.cwd().resolve()
+
+        def potential_energy(self):
+            return -1.0
+
+        def read(self):
+            return "optimized-atoms"
+
+    class FakeHessian:
+        def __init__(self, atoms, charge, **kwargs):
+            seen["hessian_atoms"] = atoms
+
+    class FakeModes:
+        def __init__(self):
+            self.wave_numbers = np.array([1.0, 2.0, 3.0])
+
+    def fake_run_thermo(frequencies, atoms=None, **kwargs):
+        seen["run_thermo_atoms"] = atoms
+        return "thermo-result"
+
+    monkeypatch.setattr(api, "Geoopt", FakeGeoopt)
+    monkeypatch.setattr(api, "Hessian", FakeHessian)
+    monkeypatch.setattr(api, "Modes", FakeModes)
+    monkeypatch.setattr(api, "run_thermo", fake_run_thermo)
+
+    job = tmp_path / "job"
+    start = Path.cwd()
+    result = api.dftbplus_thermo("initial-atoms", directory=str(job))
+
+    assert result == "thermo-result"
+    assert seen["run_thermo_atoms"] == "optimized-atoms"
+    assert seen["hessian_atoms"] == "optimized-atoms"
+    assert seen["cwd_during"] == job.resolve()  # pipeline ran inside the job dir
+    assert Path.cwd() == start  # working directory restored afterwards
+
+
 if __name__ == "__main__":
     unittest.main()