Add back enthalpy of vaporization and mixing estimation

examples/mvp/run.py

@@ -15,8 +15,8 @@
 from tyff.train import ParameterConfig, Trainable
 
 N_EPOCHS = 10
-N_REPLICATES = 2
-N_MOLECULES = 200
+N_REPLICATES = 1
+N_MOLECULES = 480
 INITIAL_FORCE_FIELD = "openff-2.3.0.offxml"
 
 # these weights should be exposed to the user in some sort of training config
@@ -41,9 +41,6 @@ def load_datasets(root: Path) -> tuple[Dataset, Dataset]:
     thermo_data = Dataset.load_from_disk(root / "data/thermo")
     energy_data = Dataset.load_from_disk(root / "data/energy")
 
-    # temporary - need to go back and add dhmix, dhvap, etc.
-    thermo_data = thermo_data.filter(lambda x: x["tag"] == "density")
-
     logger.info(f"Loaded {len(thermo_data)} thermo entries and {len(energy_data)} energy entries")
 
     return thermo_data, energy_data
@@ -82,17 +79,17 @@ def main(root: Path) -> None:
 
     theta = trainable.to_values().requires_grad_(True)
 
-    # TODO: Possible drop in L-M optimizer
+    # TODO: Possibly drop in L-M optimizer
     optimizer = Adam([theta], lr=0.01)
 
     # TODO: inner/outer loop, or any way of stepping through QM optimization
     #       without running MD jobs each time
     for epoch in range(N_EPOCHS):
-        # this is linked under the hood to theta - needed to gradients together
-        current_ff = trainable.to_force_field(theta)  # 287
+        # this is linked under the hood to theta - needed to hook gradients together
+        current_ff = trainable.to_force_field(theta)
 
         simulate(
-            force_field=current_ff,
+            tensor_force_field=current_ff,
             epoch_dir=root / f"EPOCH_{epoch}",
             thermo_data=thermo_data,
             smirnoff_force_field=INITIAL_FORCE_FIELD,
@@ -102,7 +99,7 @@ def main(root: Path) -> None:
 
         y_ref, _y_ref_std, y_pred, _y_pred_std = predict(
             dataset=thermo_data,
-            force_field=current_ff,
+            tensor_force_field=current_ff,
             epoch_dir=root / f"EPOCH_{epoch}",
         )
 

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,110 @@
+from pathlib import Path
+
+import openff.toolkit
+import pytest
+import torch
+
+from tyff.converters.openff import topologies_for
+from tyff.targets.new_thermo import predict, simulate
+from tyff.targets.thermo import create_dataset
+
+# this variable is magically detected by pytest, important that it is defined at module level
+pytestmark = pytest.mark.gpu
+
+# turn on to store simulation results in local path for debugging/iteration, otherwise (as in CI)
+# let pytest start everything fresh in a temporary directory
+USE_LOCAL_FILES = False
+
+
+class TestEnthalpyOfVaporization:
+    def test_enthalpy_of_vaporization(self, tmp_path):
+
+        dataset = create_dataset(
+            *[
+                {  # type: ignore
+                    "id": "methanol",
+                    "tag": "enthalpy_of_vaporization",
+                    "x": [1.0],
+                    "smiles": ["CO"],
+                    "temperature": 300.0,
+                    "pressure": 101.325,
+                    "value": 38.3,
+                    "std": 0.01,  # arbitrary/unused
+                    "units": "kJ/mol",
+                    "source": "test_source",
+                }
+            ]
+        )
+
+        # I could be forgetting a cleaner way of going from SMIRNOFF + set[SMILES] -> TensorForceField
+        _topologies, tensor_force_field = topologies_for(
+            smiles_set={smiles for smiles_list in dataset["smiles"] for smiles in smiles_list},
+            smirnoff_force_field=openff.toolkit.ForceField("openff-2.3.0.offxml"),
+        )
+
+        simulate(
+            tensor_force_field=tensor_force_field,
+            epoch_dir=Path("vaporization") if USE_LOCAL_FILES else tmp_path,
+            thermo_data=dataset,
+            smirnoff_force_field="openff-2.3.0.offxml",
+            n_replicates=2,
+            n_molecules=240,
+        )
+
+        y_ref, y_ref_std, y_pred, y_pred_std = predict(
+            dataset=dataset,
+            tensor_force_field=tensor_force_field,
+            epoch_dir=Path("vaporization") if USE_LOCAL_FILES else tmp_path,
+        )
+
+        assert not torch.isnan(torch.stack([y_ref, y_ref_std, y_pred, y_pred_std])).any()
+
+        # experimental value is around 38 kJ/mol, could get away with much smaller range
+        assert 30 < y_pred < 50, f"{y_pred=} {y_ref=}, {y_pred_std=}, {y_ref_std=}"
+
+
+class TestEnthalpyOfMixing:
+    def test_enthalpy_of_mixing(self, tmp_path):
+
+        dataset = create_dataset(
+            *[
+                {  # type: ignore
+                    "id": "foo",
+                    "tag": "enthalpy_of_mixing",
+                    "x": [0.8, 0.2],
+                    "smiles": ["O", "CCO"],
+                    "temperature": 298.15,
+                    "pressure": 101.325,
+                    "value": -0.765,
+                    "std": 0.01,  # arbitrary/unused
+                    "units": "kJ/mol",
+                    "source": "test_source",
+                }
+            ]
+        )
+
+        # I could be forgetting a cleaner way of going from SMIRNOFF + set[SMILES] -> TensorForceField
+        _topologies, tensor_force_field = topologies_for(
+            smiles_set={smiles for smiles_list in dataset["smiles"] for smiles in smiles_list},
+            smirnoff_force_field=openff.toolkit.ForceField("openff-2.3.0.offxml"),
+        )
+
+        simulate(
+            tensor_force_field=tensor_force_field,
+            epoch_dir=Path("mixing") if USE_LOCAL_FILES else tmp_path,
+            thermo_data=dataset,
+            smirnoff_force_field="openff-2.3.0.offxml",
+            n_replicates=3,
+            n_molecules=480,
+        )
+
+        y_ref, y_ref_std, y_pred, y_pred_std = predict(
+            dataset=dataset,
+            tensor_force_field=tensor_force_field,
+            epoch_dir=Path("mixing") if USE_LOCAL_FILES else tmp_path,
+        )
+
+        assert not torch.isnan(torch.stack([y_ref, y_ref_std, y_pred, y_pred_std])).any()
+
+        # experimental value is around -0.7 kJ/mol, but it should be negative (for all concentrations)
+        assert -1 < y_pred < 0, f"{y_pred=} {y_ref=}, {y_pred_std=}, {y_ref_std=}"

tyff/compute/workflow.py

@@ -62,7 +62,19 @@ def __init__(self, base_dir: str, parsl_config: parsl.Config):
                     self._target_compute_mapping = loaded_mapping
         else:
             # initialize to empty dict otherwise
-            self._target_compute_mapping: dict[tuple[str, str, int, int], list[tuple[str, ...]]] = dict()
+            self._target_compute_mapping: dict[
+                tuple[
+                    str,  # target id
+                    str,  # target tag
+                    str,  # force field path
+                    int,  # n_molecules
+                    int,  # n_replicates
+                ],
+                list[  # len n_replicates
+                    # each job_dir in this replicate, i.e. 1 for density, 2 for dhvap, 3 for dhmix, etc.
+                    tuple[str, ...]
+                ],
+            ] = dict()
 
         # futures submitted via submit_target/submit_target_batch that haven't
         # been waited on yet; drained by wait()
@@ -136,6 +148,7 @@ def submit_target(
         self._target_compute_mapping[
             (
                 target_config["id"],  # this is unmodified from input data and therefore could look funky
+                target_config["tag"],  # this is unmodified from input data and therefore could look funky
                 force_field,
                 n_molecules,
                 n_replicates,
@@ -145,6 +158,7 @@ def submit_target(
             self._target_compute_mapping[
                 (
                     target_config["id"],
+                    target_config["tag"],
                     force_field,
                     n_molecules,
                     n_replicates,
@@ -237,7 +251,15 @@ def estimate_target(
         if target_config["tag"] == "density":
             # pull job ids for all replicates of this target/force field/n_molecules combination
             # TODO: other querying into jobs i.e. given a target and force field but any n_molecules
-            job_ids = self._target_compute_mapping[(target_config["id"], force_field, n_molecules, n_replicates)]
+            job_ids = self._target_compute_mapping[
+                (
+                    target_config["id"],
+                    target_config["tag"],
+                    force_field,
+                    n_molecules,
+                    n_replicates,
+                )
+            ]
 
             density_futures = [
                 run_density_analysis(
@@ -258,7 +280,15 @@ def estimate_target(
         elif target_config["tag"] == "enthalpy_of_vaporization":
             # pull job ids for all replicates of this target/force field/n_molecules combination
             # TODO: other querying into jobs i.e. given a target and force field but any n_molecules
-            job_ids = self._target_compute_mapping[(target_config["id"], force_field, n_molecules, n_replicates)]
+            job_ids = self._target_compute_mapping[
+                (
+                    target_config["id"],
+                    target_config["tag"],
+                    force_field,
+                    n_molecules,
+                    n_replicates,
+                )
+            ]
 
             dhvap_futures = [
                 # this function should separate the two job ids into gas and liquid,
@@ -281,6 +311,16 @@ def estimate_target(
         else:
             raise ValueError(f"Unknown target tag {target_config['tag']}")
 
+    def estimate_dataset(
+        self,
+        dataset: datasets.Dataset,  # should be something that unpacks to DataEntry on iteration
+        tensor_force_field: tyff.TensorForceField,
+        force_field: str,
+        n_molecules: int,
+        n_replicates: int = 3,
+    ):
+        raise NotImplementedError()
+
     def run(self, compute_configs: Sequence[BaseComputeConfig]):
         """Submit a batch and block until all complete."""
         pending = self._submit_compute_batch(compute_configs)