Some fixes for pressure handling

pixi.toml

@@ -7,6 +7,7 @@ platforms = ["linux-64", "osx-arm64"]
 run_mvp = "python examples/mvp/run.py"
 run_mypy = "python -m mypy -p 'tyff'"
 run_tests = { cmd = "python -m pytest -n logical tyff/_tests", env = { OPENMM_DEFAULT_PLATFORM = "CPU" } }
+run_gpu_tests = { cmd = "python -m pytest -n 1 --run-gpu -m gpu tyff/_tests --color=no", env = { OPENMM_DEFAULT_PLATFORM = "GPU" } }
 dev_tests = { cmd = "python -m pytest -s -x --pdb --lf tyff/_tests", env = { OPENMM_DEFAULT_PLATFORM = "CPU" } }
 
 [dependencies]
@@ -83,6 +84,9 @@ cuda-version = "==13.0"
 [feature.py312.dependencies]
 python = "3.12.*"
 
+[feature.test.dependencies]
+pytest = "<9.2"
+
 [feature.test.pypi-dependencies]
 pytest = "*"
 pytest-cov = "*"

tyff/_tests/conftest.py

@@ -24,6 +24,28 @@
 _E = openff.units.unit.elementary_charge
 
 
+def pytest_addoption(parser):
+    parser.addoption(
+        "--run-gpu",
+        action="store_true",
+        default=False,
+        help="run tests marked with @pytest.mark.gpu",
+    )
+
+
+def pytest_configure(config):
+    config.addinivalue_line("markers", "gpu: tests that require a GPU (skipped unless --run-gpu is given)")
+
+
+def pytest_collection_modifyitems(config, items):
+    if config.getoption("--run-gpu"):
+        return
+    skip_gpu = pytest.mark.skip(reason="need --run-gpu option to run")
+    for item in items:
+        if "gpu" in item.keywords:
+            item.add_marker(skip_gpu)
+
+
 @pytest.fixture
 def density_entry() -> DataEntry:
     return {

tyff/_tests/integration_tests/compute/test_workflow.py

@@ -0,0 +1,67 @@
+import json
+
+import pytest
+
+from tyff.compute.configs import local_config
+from tyff.compute.workflow import SimulationWorkflow
+from tyff.targets.thermo import create_dataset
+
+
+@pytest.mark.gpu
+def test_nonstandard_pressure(tmp_path):
+
+    dataset = create_dataset(
+        *[
+            {  # type: ignore
+                "id": "high_pressure_water",
+                "tag": "density",
+                "x": [1.0],
+                "smiles": ["O"],
+                "temperature": 298.15,
+                "pressure": val,
+                "value": 1.0,
+                "std": 0.01,
+                "units": "g/mL",
+                "source": "test_source",
+            }
+            for val in [101.325, 101325]
+        ]
+    )
+
+    with SimulationWorkflow(
+        base_dir=tmp_path,
+        parsl_config=local_config(),
+    ) as workflow:
+        workflow.submit_dataset(
+            dataset,
+            force_field="openff-2.3.0.offxml",
+            n_molecules=480,
+            n_replicates=1,
+        )
+
+        workflow.wait()
+
+        # hack - there needs to be a better way of grabbing results
+        job_dirs = [
+            path
+            for path in (tmp_path).glob("*/")
+            if not path.match("thermo_targets")
+            and not path.match("reference_force_fields")
+            and not path.match("interchanges")
+        ]
+
+        results = dict()
+
+        for job_dir in job_dirs:
+            with open(job_dir / "compute_config.json") as f:
+                compute_config = json.load(f)
+                target_pressure = compute_config["pressure"]
+
+                assert compute_config["temperature"] == 298.15
+
+            with open(job_dir / "ensemble_averages.json") as f:
+                density = json.load(f)["density"]
+
+            results[target_pressure] = density
+
+        assert results[101325.0] > results[101.325]

tyff/compute/_equilibrate.py

@@ -142,7 +142,7 @@ def _run_equilibration(
             output_file=files["msgpack_trajectory"].filepath,
             report_interval=1000,
             beta=1.0 / (openmm.unit.MOLAR_GAS_CONSTANT_R * compute_config["temperature"] * openmm.unit.kelvin),
-            pressure=pressure if pressure is None else openmm.unit.kilopascal,
+            pressure=pressure if pressure is None else pressure * openmm.unit.kilopascal,
         )
     )
 

tyff/compute/_jacobian.py

@@ -78,7 +78,7 @@ def _get_ensemble_average_and_jacobian(
         compute_config = BulkLiquid(**json.load(f))  # type: ignore[typeddict-item]
 
     temperature = compute_config["temperature"]  # kelvin, float
-    pressure = compute_config.get("pressure")  # atmosphere, float | None
+    pressure = compute_config.get("pressure")  # kPa, float | None
 
     interchanges = []
 
@@ -141,7 +141,7 @@ def _get_ensemble_average_and_jacobian(
         worker_force_field,
         frames_path,
         temperature * openmm.unit.kelvin,
-        None if pressure is None else pressure * openmm.unit.atmosphere,
+        None if pressure is None else pressure * openmm.unit.kilopascal,
     )
 
     # originally this was sorted(means),

tyff/compute/_produce.py

@@ -138,7 +138,7 @@ def _run_production(
             output_file=files["msgpack_trajectory"].filepath,
             report_interval=1000,
             beta=1.0 / (openmm.unit.MOLAR_GAS_CONSTANT_R * compute_config["temperature"] * openmm.unit.kelvin),
-            pressure=pressure if pressure is None else openmm.unit.kilopascal,
+            pressure=pressure if pressure is None else pressure * openmm.unit.kilopascal,
         )
     )