Update pressure handling

tyff/_tests/compute/test_prep.py

@@ -87,4 +87,9 @@ def test_pressure_in_gas_configs(dhvap_entry):
 
     for this_config in compute_configs:
         for val in this_config:
+            # both liquid and gas configs should have a pressure key
             assert "pressure" in val
+
+            # liquid configs must have non-null pressure, gas configs must have None
+            assert (val["tag"] == "liquid") == (val["pressure"] is not None)
+            assert (val["tag"] == "gas") == (val["pressure"] is None)

tyff/compute/_equilibrate.py

@@ -92,13 +92,14 @@ def _run_equilibration(
         if topology.getPeriodicBoxVectors() is not None:
             simulation.context.setPeriodicBoxVectors(*topology.getPeriodicBoxVectors())
 
+    # let pressure be None for gas-phase NVT simulations, including reporting as None to TensorReporter later
     pressure = compute_config.get("pressure", None)
 
-    if pressure is None:
-        raise PressureNotDefinedError("Trying to set up NPT simulation but no pressure defined.")
-
     # only set pressure in liquid NPT simulations, not in gas-phase NVT simulations
     if compute_config["tag"] == "liquid":
+        if pressure is None:
+            raise PressureNotDefinedError("Trying to set up NPT simulation in liquid phase but no pressure defined.")
+
         barostat = openmm.MonteCarloBarostat(
             (pressure * openmm.unit.kilopascal).value_in_unit(openmm.unit.bar),  # pressure in bar
             compute_config["temperature"],  # temperature in kelvin
@@ -141,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 * openmm.unit.kilopascal,
+            pressure=pressure if pressure is None else openmm.unit.kilopascal,
         )
     )
 

tyff/compute/_produce.py

@@ -91,19 +91,21 @@ def _run_production(
     logger.info("Reinitializing context (in production step)")
     simulation.context.reinitialize(preserveState=True)
 
+    # let pressure be None for gas-phase NVT simulations, including reporting as None to TensorReporter later
+    pressure = compute_config.get("pressure", None)
+
     if compute_config["tag"] == "liquid":
+        if pressure is None:
+            # this clause should already be triggered in _run_equilibration
+            raise PressureNotDefinedError("Trying to set up NPT simulation but no pressure defined.")
+
         # just double check pressure is loaded back in correctly from state
         barostat = next(
             force for force in simulation.system.getForces() if isinstance(force, openmm.MonteCarloBarostat)
         )
 
         assert barostat.getDefaultPressure() is not None
 
-    pressure = compute_config.get("pressure", None)
-
-    if pressure is None:
-        raise PressureNotDefinedError("Trying to set up NPT simulation but no pressure defined.")
-
     simulation.context.setVelocitiesToTemperature(
         compute_config["temperature"],  # kelvin, but as float
         compute_config["replicate_index"] + 1,
@@ -136,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 * openmm.unit.kilopascal,
+            pressure=pressure if pressure is None else openmm.unit.kilopascal,
         )
     )
 

tyff/compute/prep.py

@@ -164,7 +164,7 @@ def _make_enthalpy_of_vaporization_compute_configs(
                         smiles=data_entry["smiles"],
                         x=data_entry["x"],
                         temperature=data_entry["temperature"],
-                        pressure=data_entry["pressure"],
+                        pressure=None,  # should *always* be None, but not validated in the model
                     ),
                 )
             )

tyff/configs/_compute.py

@@ -18,4 +18,4 @@ class BaseComputeConfig(typing_extensions.TypedDict, extra_items=typing.Any):
 
     temperature: float
 
-    pressure: float
+    pressure: float | None

tyff/configs/gas.py

@@ -4,6 +4,14 @@
 
 
 class VacuumGas(BaseComputeConfig):
+    """
+    General model for vacuum simulations
+
+    Note that `pressure` is inherited from the base class as `float | None`,
+    but should always be `None` in this model.
+
+    """
+
     tag: typing.Literal["gas"]
 
     """