Add environment variable for output directory

.ci-scripts/build_and_test.sh

@@ -263,7 +263,7 @@ then
     timed_message "Run regression tests"
 
     mpi_exe=$(grep 'MPIEXEC_EXECUTABLE' "${hostconfig_path}" | cut -d'"' -f2 | sed 's/;/ /g')
-    pytest -v -s tests/regression --mpi-exec="${mpi_exe}"
+    pytest -v -s -m regression --mpi-exec="${mpi_exe}"
 
     timed_message "Quandary tests completed"
 fi
@@ -284,7 +284,7 @@ then
     timed_message "Run performance tests"
 
     mpi_exe=$(grep 'MPIEXEC_EXECUTABLE' "${hostconfig_path}" | cut -d'"' -f2 | sed 's/;/ /g')
-    pytest -v -s tests/performance --mpi-exec="${mpi_exe}" --benchmark-json=benchmark_results.json
+    pytest -v -s -m performance --mpi-exec="${mpi_exe}" --benchmark-json=benchmark_results.json
 
     timed_message "Quandary performance tests completed"
 fi

.pytest.ini

@@ -1,3 +1,7 @@
 [pytest]
 addopts = --ignore=blt --ignore=uberenv_libs
 
+markers =
+    regression: mark a test as a regression test
+    performance: mark a test as a performance test
+

README.md

@@ -125,10 +125,19 @@ The `examples/` folder exemplifies the usage of Quandary's Python interface.
 
 # Tests
 
+
 ## Regression tests
-Regression tests are defined in `tests/regression` and can be run with
+Regression tests are defined in `tests/regression` and `tests/python` directories.
+
+You can run all regression tests with:
+```bash
+pytest -m regression
 ```
+
+Or run tests in a specific directory:
+```bash
 pytest tests/regression
+pytest tests/python
 ```
 See `tests/regression/README.md` for more information.
 

quandary.py

@@ -12,6 +12,12 @@ class Quandary:
     """
     This class collects configuration options to run quandary and sets all defaults. Each parameter can be overwritten within the constructor. The number of time-steps required to resolve the time-domain, as well as the resonant carrier wave frequencies are computed within the constructor. If you attempt to change options of the configuration *after* construction by accessing them directly (e.g. myconfig.fieldname = mynewsetting), it is therefore advised to call myconfig.update() afterwards to recompute the number of time-steps and carrier waves.
 
+    Environment Variables:
+    --------------------
+    QUANDARY_BASE_DATADIR  : Base directory for output files. If set, all output will be written to this directory
+                             instead of the current working directory. Relative paths provided to methods will be
+                             considered relative to this base directory.
+
     Parameters
     ----------
     # Quantum system specifications
@@ -295,7 +301,8 @@ def simulate(self, *, pcof0=[], pt0=[], qt0=[], maxcores=-1, datadir="./run_dir"
         pt0           : List of ndarrays for the real part of the control function [MHz] for each oscillator, ndarray size = nsteps+1. Assumes spline_order == 0 and ignores the pcof0 argument. Default: []
         qt0           : Same as pt0, but for the imaginary part.
         maxcores      : Maximum number of processing cores. Default: number of initial conditions
-        datadir       : Data directory for storing output files. Default: "./run_dir"
+        datadir       : Data directory for storing output files. Default: "./run_dir".
+                        If $QUANDARY_BASE_DATADIR is set, this will be relative to that directory, otherwise relative to current working directory
         quandary_exec : Location of Quandary's C++ executable, if not in $PATH
         cygwinbash    : To run on Windows through Cygwin, set the path to Cygwin/bash.exe. Default: None.
         mpi_exec      : String for MPI launcher prefix, e.g. "mpirun -np" or "srun -n". The string should include the flag for core counts, but not the number of cores itself which will be appended automatically
@@ -326,7 +333,8 @@ def optimize(self, *, pcof0=[], pt0=[], qt0=[], maxcores=-1, datadir="./run_dir"
         pcof0          : List of control parameters to start the optimization from. Default: Use initial guess from the Quandary (pcof0, or pcof0_filename, or randomized initial guess)
         maxcores      : Maximum number of processing cores. Default: number of initial conditions
         pt, qt          :  p,q-control pulses [MHz] at each time point for each oscillator (List of list)
-        datadir       : Data directory for storing output files. Default: "./run_dir"
+        datadir       : Data directory for storing output files. Default: "./run_dir".
+                        If $QUANDARY_BASE_DATADIR is set, this will be relative to that directory, otherwise relative to current working directory
         quandary_exec : Location of Quandary's C++ executable, if not in $PATH
         mpi_exec      : String for MPI launcher prefix, e.g. "mpirun -np" or "srun -n". The string should include the flag for core counts, but not the number of cores itself which will be appended automatically
         cygwinbash    : To run on Windows through Cygwin, set the path to Cygwin/bash.exe. Default: None.
@@ -355,7 +363,8 @@ def evalControls(self, *, pcof0=[], points_per_ns=1,datadir="./run_dir", quandar
         --------------------
         pcof0         :  List of control parameters (bspline coefficients) that determine the controls pulse. If not given, the initial guess from Quandary class will be used (pcof0, or filename, or random initial control...)
         points_per_ns :  sample rate of the resulting controls. Default: 1ns 
-        datadir       :  Directory for output files. Default: "./run_dir"
+        datadir       :  Directory for output files. Default: "./run_dir".
+                         If $QUANDARY_BASE_DATADIR is set, this will be relative to that directory, otherwise relative to current working directory
         quandary_exec :  Path to Quandary's C++ executable if not in $PATH
         mpi_exec      : String for MPI launcher prefix, e.g. "mpirun -np" or "srun -n". The string should include the flag for core counts, but not the number of cores itself which will be appended automatically
         cygwinbash    : To run on Windows through Cygwin, set the path to Cygwin/bash.exe. Default: None.
@@ -370,6 +379,8 @@ def evalControls(self, *, pcof0=[], points_per_ns=1,datadir="./run_dir", quandar
         nsteps_org = self.nsteps
         self.nsteps = int(np.floor(self.T * points_per_ns))
 
+        datadir = resolve_datadir(datadir)
+
         # Execute quandary in 'evalcontrols' mode
         datadir_controls = datadir +"_ppns"+str(points_per_ns)
         os.makedirs(datadir_controls, exist_ok=True)
@@ -445,6 +456,8 @@ def __run(self, *, pcof0=[], runtype="optimization", overwrite_popt=False, maxco
           4. Evaluate controls on the input sample rate, if given
         """
 
+        datadir = resolve_datadir(datadir)
+
         # Create quandary data directory and dump configuration file
         os.makedirs(datadir, exist_ok=True)
         config_filename = self.__dump(pcof0=pcof0, runtype=runtype, datadir=datadir)
@@ -715,6 +728,7 @@ def get_results(self, *, datadir="./", ignore_failure=False):
         Parameters:
         -----------
         datadir        (string) : Directory containing Quandary's output files.
+                                  If $QUANDARY_BASE_DATADIR is set, this will be relative to that directory, otherwise relative to current working directory
         ignore_failure (bool)   : Flag to ignore warning when an expected file can't be found
 
         Returns:
@@ -726,6 +740,8 @@ def get_results(self, *, datadir="./", ignore_failure=False):
         population      :  Evolution of the population of each oscillator, of each initial condition. (expectedEnergy[oscillator][initialcondition])
         """
 
+        datadir = resolve_datadir(datadir)
+
         # Get control parameters
         filename = os.path.join(datadir, "params.dat")
         try:
@@ -1024,6 +1040,36 @@ def map_to_oscillators(id, Ne, Ng):
 
     return localIDs 
 
+def resolve_datadir(datadir: str) -> str:
+    """Helper function to resolve the output directory using environment variable
+
+    Parameters:
+    ----------
+    datadir : Output directory
+              If $QUANDARY_BASE_DATADIR is set, relative paths will be resolved against it
+
+    Returns:
+    -------
+    Resolved absolute or relative path for the data directory
+
+    Raises:
+    ------
+    ValueError: If QUANDARY_BASE_DATADIR is set but doesn't exist or isn't a directory
+    """
+    if os.path.isabs(datadir):
+        return datadir
+
+    base_dir = os.environ.get("QUANDARY_BASE_DATADIR")
+    if base_dir:
+        if not os.path.exists(base_dir):
+            raise ValueError(f"Environment variable QUANDARY_BASE_DATADIR points to non-existent path: {base_dir}")
+        if not os.path.isdir(base_dir):
+            raise ValueError(f"Environment variable QUANDARY_BASE_DATADIR is not a directory: {base_dir}")
+
+        datadir = os.path.join(base_dir, datadir)
+
+    return os.path.normpath(datadir)
+
 
 def hamiltonians(*, N, freq01, selfkerr, crosskerr=[], Jkl = [], rotfreq=[], verbose=True):
     """  

tests/performance/performance_test.py

@@ -7,6 +7,9 @@
 from pydantic import BaseModel, TypeAdapter
 from tests.utils.common import build_mpi_command
 
+# Mark all tests in this file as performance tests
+pytestmark = pytest.mark.performance
+
 TEST_PATH = os.path.dirname(os.path.realpath(__file__))
 TEST_CASES_PATH = os.path.join(TEST_PATH, "test_cases.json")
 TEST_CONFIG_PATH = os.path.join(TEST_PATH, "configs")

tests/python/test_env_variable.py

@@ -0,0 +1,169 @@
+import os
+import pytest
+from quandary import Quandary
+
+# Mark all tests in this file as regression tests
+pytestmark = pytest.mark.regression
+
+BASE_DATADIR = "QUANDARY_BASE_DATADIR"
+
+
+def quandary_simulate(datadir, mpi_exec):
+    return Quandary(
+        Ne=[2],
+        Ng=[0],
+        freq01=[4.0],
+        selfkerr=[0.2],
+        T=1.0,
+        nsteps=10,
+        maxiter=1,
+        spline_order=0,
+    ).simulate(
+        datadir=datadir,
+        mpi_exec=mpi_exec,
+        maxcores=2
+    )
+
+
+def quandary_optimize(datadir, mpi_exec):
+    return Quandary(
+        Ne=[2],
+        targetstate=[0.0, 1.0],
+        initialcondition="basis",
+        tol_infidelity=1e-2,
+        nsteps=1,
+        maxiter=1,
+        spline_order=0
+    ).optimize(
+        datadir=datadir,
+        mpi_exec=mpi_exec,
+        maxcores=2
+    )
+
+
+test_cases = [
+    quandary_optimize,
+    quandary_simulate,
+]
+
+
+@pytest.mark.parametrize("quandary", test_cases)
+def test_relative_output_path_without_env_var(quandary, request, cd_tmp_path, clean_env_var, mpi_exec):
+    datadir_name = request.node.name
+    datadir_path = os.path.join(os.getcwd(), datadir_name)
+
+    quandary(datadir=datadir_name, mpi_exec=mpi_exec)
+
+    assert_output_files(datadir_path)
+
+
+@pytest.mark.parametrize("quandary", test_cases)
+def test_absolute_output_path_without_env_var(quandary, request, tmp_path, clean_env_var, mpi_exec):
+    datadir_name = request.node.name
+    datadir_path = os.path.join(tmp_path, datadir_name)
+
+    quandary(datadir=datadir_path, mpi_exec=mpi_exec)
+
+    assert_output_files(datadir_path)
+
+
+@pytest.mark.parametrize("quandary", test_cases)
+def test_relative_output_path_with_env_var(quandary, request, tmp_path, clean_env_var, mpi_exec):
+    base_dir = str(tmp_path)
+    os.environ[BASE_DATADIR] = base_dir
+    datadir_name = request.node.name
+    datadir_path = os.path.join(base_dir, datadir_name)
+
+    quandary(datadir=datadir_name, mpi_exec=mpi_exec)
+
+    assert_output_files(datadir_path)
+
+
+@pytest.mark.parametrize("quandary", test_cases)
+def test_absolute_output_path_with_env_var(quandary, request, tmp_path, clean_env_var, mpi_exec):
+    os.environ[BASE_DATADIR] = "should_not_use_this/path"
+    datadir_name = request.node.name
+    datadir_path = os.path.join(tmp_path, datadir_name)
+
+    quandary(datadir=datadir_path, mpi_exec=mpi_exec)
+
+    assert_output_files(datadir_path)
+    assert not os.path.exists(os.environ[BASE_DATADIR])
+
+
+@pytest.mark.parametrize("quandary", test_cases)
+def test_nonexistent_base_directory(quandary, request, tmp_path, clean_env_var, mpi_exec):
+    nonexistent_path = os.path.join(tmp_path, "nonexistent_directory")
+    os.environ[BASE_DATADIR] = nonexistent_path
+    datadir_name = "some_output_dir"
+
+    with pytest.raises(ValueError) as excinfo:
+        quandary(datadir=datadir_name, mpi_exec=mpi_exec)
+
+    assert "non-existent path" in str(excinfo.value)
+    assert nonexistent_path in str(excinfo.value)
+
+
+@pytest.mark.parametrize("quandary", test_cases)
+def test_file_as_base_directory(quandary, request, tmp_path, clean_env_var, mpi_exec):
+    file_path = os.path.join(tmp_path, "this_is_a_file.txt")
+    with open(file_path, 'w') as f:
+        f.write("This is a file, not a directory")
+
+    os.environ[BASE_DATADIR] = file_path
+    datadir_name = "some_output_dir"
+
+    with pytest.raises(ValueError) as excinfo:
+        quandary(datadir=datadir_name, mpi_exec=mpi_exec)
+
+    assert "not a directory" in str(excinfo.value)
+    assert file_path in str(excinfo.value)
+
+
+@pytest.fixture
+def mpi_exec(request):
+    """Get MPI executor from pytest option."""
+    executor = request.config.getoption("--mpi-exec")
+    if executor != "mpirun":
+        return f"{executor} -n "
+    return "mpirun -np "
+
+
+@pytest.fixture
+def cd_tmp_path(tmp_path):
+    """Change to a temporary directory for the test and return afterward."""
+    original_cwd = os.getcwd()
+    os.chdir(tmp_path)
+    try:
+        yield tmp_path
+    finally:
+        os.chdir(original_cwd)
+
+
+@pytest.fixture
+def clean_env_var():
+    """Fixture to ensure env var is restored to previous state after the tests"""
+    orig_value = os.environ.get(BASE_DATADIR)
+
+    if BASE_DATADIR in os.environ:
+        del os.environ[BASE_DATADIR]
+
+    yield
+
+    if orig_value is not None:
+        os.environ[BASE_DATADIR] = orig_value
+    elif BASE_DATADIR in os.environ:
+        del os.environ[BASE_DATADIR]
+
+
+def assert_output_files(datadir):
+    expected_output_files = [
+        "config.cfg",
+        "optim_history.dat",
+        "params.dat",
+        "control0.dat"
+    ]
+
+    assert os.path.exists(datadir), f"directory {datadir} does not exist"
+    for file in expected_output_files:
+        assert os.path.exists(os.path.join(datadir, file)), f"file {file} does not exist"

tests/regression/regression_test.py

@@ -8,6 +8,9 @@
 
 from tests.utils.common import build_mpi_command
 
+# Mark all tests in this file as regression tests
+pytestmark = pytest.mark.regression
+
 REL_TOL = 1.0e-7
 ABS_TOL = 1.0e-15