From 5577c30fac7bf0e79a58e7729a48f5c6cd96c651 Mon Sep 17 00:00:00 2001
From: Logan Ward <ward.logan.t@gmail.com>
Date: Thu, 7 Dec 2023 09:49:46 -0500
Subject: [PATCH] Run ASE inside a temporary directory

Avoid calcualtions from intefering with each other
---
 jitterbug/parsl.py  | 22 +++++++++++++++++-----
 tests/test_parsl.py |  7 +++++++
 2 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/jitterbug/parsl.py b/jitterbug/parsl.py
index 2567fb5..7f308e0 100644
--- a/jitterbug/parsl.py
+++ b/jitterbug/parsl.py
@@ -1,12 +1,15 @@
 """Wrappers for functions compatible with the Parsl workflow engine"""
+import os
+from tempfile import TemporaryDirectory
 from typing import Optional
+from pathlib import Path
 
 import ase
 
 from jitterbug.utils import make_calculator, write_to_string
 
 
-def get_energy(atoms: ase.Atoms, method: str, basis: Optional[str], **kwargs) -> str:
+def get_energy(atoms: ase.Atoms, method: str, basis: Optional[str], scratch_dir: Optional[str] = None, **kwargs) -> str:
     """Compute the energy of an atomic structure
 
     Keyword arguments are passed to :meth:`make_calculator`.
@@ -15,11 +18,20 @@ def get_energy(atoms: ase.Atoms, method: str, basis: Optional[str], **kwargs) ->
         atoms: Structure to evaluate
         method: Name of the method to use (e.g., B3LYP)
         basis: Basis set to use (e.g., cc-PVTZ)
+        scratch_dir: Path to the scratch directory.
     Returns:
         Atoms record serialized with the energy and any other data produced by the calculator
     """
 
-    calc = make_calculator(method, basis, **kwargs)
-    atoms.calc = calc
-    atoms.get_potential_energy()
-    return write_to_string(atoms, 'json')
+    # Make a temporary directory
+    start_dir = Path.cwd()
+    tmp = TemporaryDirectory(dir=scratch_dir, prefix='jitterbug_')
+    try:
+        os.chdir(tmp.name)
+        calc = make_calculator(method, basis, directory=tmp.name, **kwargs)
+        atoms.calc = calc
+        atoms.get_potential_energy()
+        return write_to_string(atoms, 'json')
+    finally:
+        os.chdir(start_dir)
+        tmp.cleanup()
diff --git a/tests/test_parsl.py b/tests/test_parsl.py
index c9931dd..598a7f9 100644
--- a/tests/test_parsl.py
+++ b/tests/test_parsl.py
@@ -1,3 +1,5 @@
+from pathlib import Path
+
 from ase.io import read
 
 from jitterbug.parsl import get_energy
@@ -5,7 +7,12 @@
 
 
 def test_energy(xyz_path):
+    mopac_out = Path('mopac.out')
+    mopac_out.unlink(missing_ok=True)
+
     atoms = read(xyz_path)
     atoms_msg = get_energy(atoms, 'pm7', None)
     new_atoms = read_from_string(atoms_msg, 'json')
     assert 'energy' in new_atoms.calc.results
+
+    assert not mopac_out.exists()