Add physical property prediction with PyTorch machinery

examples/mvp/run.py

@@ -2,38 +2,21 @@
 from pathlib import Path
 
 from datasets import Dataset
-from openff.toolkit import ForceField, Molecule
+from openff.toolkit import ForceField
+from rich import print
 
-import tyff
-import tyff.converters
-import tyff.train
 from tyff.compute._logging import _set_up_logger
-from tyff.compute.configs import local_config
-from tyff.compute.workflow import SimulationWorkflow
-from tyff.converters.openff._tensors import convert_tensor_force_field
+from tyff.converters import topologies_for
 from tyff.targets import energy, thermo
+from tyff.targets.new_thermo import predict, simulate
+from tyff.train import ParameterConfig, Trainable
 
 N_EPOCHS = 1
-N_REPLICATES = 1
+N_REPLICATES = 3
 N_MOLECULES = 200
 INITIAL_FORCE_FIELD = "openff-2.3.0.offxml"
 
 
-def topologies_for(
-    smiles_set: set[str],
-    smirnoff_force_field: ForceField,
-) -> tuple[dict[str, tyff.TensorTopology], tyff.TensorForceField]:
-    # this function is maybe meant to be surely there is a more straightforward way to do this?
-    interchanges = [
-        smirnoff_force_field.create_interchange(Molecule.from_mapped_smiles(smiles).to_topology())
-        for smiles in smiles_set
-    ]
-
-    tensor_ff, topologies = tyff.converters.convert_interchange(interchanges)
-
-    return {smiles: topology for smiles, topology in zip(smiles_set, topologies)}, tensor_ff
-
-
 def load_datasets(root: Path) -> tuple[Dataset, Dataset]:
     """Load datasets stored at data/*.tar.gz"""
     if not (root / "data/energy").exists():
@@ -47,79 +30,21 @@ 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
 
 
-def simulate(
-    force_field: tyff.TensorForceField,
-    root_dir: Path,
-    thermo_data: Dataset,
-    n_molecules: int,
-    n_replicates: int,
-) -> None:
-    """
-    Simulate the given force field using the provided thermo data.
-
-    This is meant to be used inside of a training loop.
-
-    Parameters
-    ----------
-    force_field : tyff.TensorForceField
-        The force field to be used in the simulation.
-    root_dir : Path
-        The root directory where simulation results will be stored.
-    thermo_data : Dataset
-        The dataset containing thermodynamic data for the simulation.
-    n_molecules : int
-        The number of molecules to simulate.
-    n_replicates : int
-        The number of replicates for each simulation.
-
-    Returns
-    -------
-    None
-    """
-
-    # write out this epoch's (tensor) force field as a SMIRNOFF force field
-    current_smirnoff_force_field = convert_tensor_force_field(
-        original_force_field=ForceField(INITIAL_FORCE_FIELD),
-        tensor_force_field=force_field,
-    )
-
-    logger.info("Making directory for this epoch's simulations")
-    root_dir.mkdir(exist_ok=True)
-
-    logger.info("Writing out the current SMIRNOFF force field to file")
-    smirnoff_force_field_file = str(root_dir / "current_force_field.offxml")
-    current_smirnoff_force_field.to_file(smirnoff_force_field_file)
-
-    logger.info("Setting up simulation workflow")
-    with SimulationWorkflow(
-        base_dir=str(root_dir / "simulations"),
-        parsl_config=local_config(max_workers=8),
-    ) as workflow:
-        logger.info("Submitting dataset")
-        workflow.submit_dataset(
-            dataset=thermo_data,
-            force_field=smirnoff_force_field_file,
-            n_molecules=n_molecules,
-            n_replicates=n_replicates,
-        )
-
-        workflow.wait()
-
-
 def main(root: Path) -> None:
     global logger
 
     logger = _set_up_logger(f"{root}/mvp.log")
 
     thermo_data, energy_data = load_datasets(root)
 
-    # should this be convert_interchanges?
-    # is `topologies` list[tyff.TensorTopology] or dict[str, tyff.TensorTopology]?
     _topologies, tensor_ff = topologies_for(
         smiles_set={
             *thermo.extract_smiles(thermo_data),
@@ -128,10 +53,10 @@ def main(root: Path) -> None:
         smirnoff_force_field=ForceField(INITIAL_FORCE_FIELD),
     )
 
-    trainable = tyff.train.Trainable(
+    trainable = Trainable(
         tensor_ff,
         parameters={
-            "vdW": tyff.train.ParameterConfig(
+            "vdW": ParameterConfig(  # type: ignore
                 cols=["epsilon", "sigma"],
                 scales={"epsilon": 10.0, "sigma": 1.0},
                 limits={"epsilon": (0.0001, None), "sigma": (0.5, None)},
@@ -144,15 +69,31 @@ def main(root: Path) -> None:
     for epoch in range(N_EPOCHS):
         current_ff = trainable.to_force_field(theta)
 
+        # no return value?
         simulate(
             force_field=current_ff,
-            root_dir=root / f"EPOCH_{epoch}",
+            epoch_dir=root / f"EPOCH_{epoch}",
             thermo_data=thermo_data,
+            smirnoff_force_field=ForceField(INITIAL_FORCE_FIELD),
             n_replicates=N_REPLICATES,
             n_molecules=N_MOLECULES,
         )
 
-    return
+        y_ref, y_ref_std, y_pred, y_pred_std = predict(
+            dataset=thermo_data,
+            force_field=current_ff,
+            epoch_dir=root / f"EPOCH_{epoch}",
+        )
+
+        print(f"{y_ref=}\n{y_ref_std=}\n{y_pred=}\n{y_pred_std=}\n")
+        # y_ref=tensor([0.8736, 0.8060, 1.1335, 0.8579, 0.8215, 0.8375, 0.8543, 1.0950, 1.1203,
+        #         1.1290])
+        # y_ref_std=tensor([4.0000e-05, 5.7000e-05, 5.0000e-04, 3.3500e-04, 1.1150e-04, 1.0100e-04,
+        #         9.1000e-05, 1.5000e-04, 3.5000e-04, 1.0000e-03])
+        # y_pred=tensor([0.8859, 0.7895, 1.1050, 0.8637, 0.8140, 0.8355, 0.8586, 1.0800, 1.0986,
+        #         1.1028], dtype=torch.float64, grad_fn=<StackBackward0>)
+        # y_pred_std=tensor([0.0037, 0.0043, 0.0068, 0.0021, 0.0024, 0.0025, 0.0065, 0.0065, 0.0092,
+        #         0.0034], dtype=torch.float64, grad_fn=<StackBackward0>)
 
 
 if __name__ == "__main__":

prototype.py

@@ -0,0 +1,191 @@
+import json
+import pathlib
+from collections.abc import Sequence
+
+import torch
+
+import tyff
+import tyff.converters
+from tyff._serialization import load_tensor_force_field
+from tyff.compute.jobs import make_job_id
+from tyff.configs.liquid import BulkLiquid
+from tyff.configs.targets.thermo import DataEntry
+
+# in leiu of better target-job tracking, just define known argumnets as ~globals
+FORCE_FIELD = "openff-2.3.0.offxml"
+N_MOLECULES = 200
+N_REPLICATES = 3
+
+
+def load_remote_observables(
+    force_field_per_epoch: tyff.TensorForceField,
+    base_dir: str,
+    job_ids: list[str],
+) -> dict[str, dict]:
+    """Load observables and attach to graph."""
+    from tyff.mm._ops import _pack_force_field
+
+    tensors, *_ = _pack_force_field(
+        force_field_per_epoch
+    )  # force_field_per_epoch should also be the global reference FF passed to the workers
+    parameters: torch.Tensor = torch.cat(
+        [t.reshape(-1) for t in tensors if t is not None]
+    )  # has to be flat, and not a tuple
+
+    # TODO: Move this class to not be inside of this function
+    class RemoteEnsembleAverage(torch.autograd.Function):
+        @staticmethod
+        def forward(
+            ctx: torch.autograd.function.FunctionCtx,
+            values: torch.Tensor,
+            jacobian: torch.Tensor,
+            parameters: torch.Tensor,
+        ) -> torch.Tensor:
+            ctx.save_for_backward(jacobian)
+            return values.clone()
+
+        @staticmethod
+        def backward(
+            ctx: torch.autograd.function.FunctionCtx,
+            grad_output: torch.Tensor,
+        ) -> tuple[None, None, torch.Tensor]:
+            (jacobian,) = ctx.saved_tensors
+            return None, None, grad_output @ jacobian
+
+    job_id_to_observables = {}
+
+    for job_id in job_ids:
+        job_dir = pathlib.Path(base_dir) / job_id
+
+        assert job_dir.exists()
+
+        with open(job_dir / "ensemble_averages.json") as f:
+            means_dict: dict[str, float] = json.load(f)  # actually str, torch.float64
+
+        jacobians: torch.Tensor = torch.load(job_dir / "jacobian.pt")
+
+        names = list(
+            means_dict
+        )  # get the names of the computed properties in the means_dict so we can convert dict -> tensor -> dict again
+        tensor_means = parameters.new_tensor([means_dict[x] for x in names])
+        attached_means = RemoteEnsembleAverage.apply(
+            tensor_means,
+            jacobians.to("cpu"),  # to(device) ?, get from somewhere?
+            parameters,
+        )
+
+        job_id_to_observables[job_id] = dict(zip(names, attached_means))
+
+    return job_id_to_observables
+
+
+def compute_properties(
+    dataset: Sequence[DataEntry],
+    observables: dict[str, dict[str, torch.Tensor]],
+) -> list[tuple[dict, dict]]:
+    from tyff.compute.prep import _compute_configs_from_data_entry
+
+    means_and_stds: list[tuple[dict, dict]] = list()
+
+    for entry in dataset:
+        if entry["tag"] == "density":
+            # this includes replicates, but only jobs for this physical property target
+            # structure is a list (length n_replicates) of tuples, each tuple is
+            # 1-length and contains a compute config for (a single) liquid job
+            compute_configs: Sequence[tuple[BulkLiquid, ...]] = _compute_configs_from_data_entry(
+                data_entry=entry,
+                force_field=FORCE_FIELD,
+                n_molecules=N_MOLECULES,
+                n_replicates=N_REPLICATES,
+            )
+
+            # only one job per replicate in density calculatiosn
+            job_ids = [make_job_id(replicate[0]) for replicate in compute_configs]
+
+            types_of_observables = [*observables[job_ids[0]].keys()]
+
+            # could be simpler to return these as 1-D tensors - make sure to sort
+            # keys to avoid ordering issues
+            mean = {
+                key: torch.stack([observables[job_id][key] for job_id in job_ids]).mean()
+                for key in types_of_observables
+            }
+            std = {
+                key: torch.stack([observables[job_id][key] for job_id in job_ids]).std()
+                for key in types_of_observables
+            }
+
+            # figure out which job id/s (if multiple replicas) in the observable corresponds to
+            # and calculate means and stds
+        elif entry["tag"] == "dielectric_constant":
+            raise NotImplementedError()
+            # likely identical to density case
+        elif entry["tag"] == "enthalpy_of_mixing":
+            raise NotImplementedError()
+            # figure out the three job IDs necessary from observables, calculate, etc
+        elif entry["tag"] == "excess_molar_volume":
+            raise NotImplementedError()
+        elif entry["tag"] == "enthalpy_of_vaporization":
+            raise NotImplementedError()
+        else:
+            raise ValueError(f"I do not know what to do with {entry['tag']}")
+
+        means_and_stds.append((mean, std))
+
+    return means_and_stds
+
+
+def predict(
+    dataset: Sequence[DataEntry],
+    force_field: tyff.TensorForceField,
+    job_ids: list[str],
+    base_dir: str,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    attached_means = load_remote_observables(force_field, base_dir, job_ids)
+
+    means_and_stds: list[dict[str, torch.Tensor]] = compute_properties(dataset, attached_means)
+
+    ref_values = [entry["value"] for entry in dataset]
+    ref_stds = [entry["std"] for entry in dataset]
+    predicted_values = [mean["density"] for mean, _ in means_and_stds]
+    predicted_stds = [std["density"] for _, std in means_and_stds]
+
+    ref_values = torch.tensor(ref_values)
+    ref_stds = torch.tensor(ref_stds)
+    predicted_values = torch.tensor(predicted_values)
+    predicted_stds = torch.tensor(predicted_stds)
+
+    return ref_values, ref_stds, predicted_values, predicted_stds
+
+
+if __name__ == "__main__":
+    with open("density_example/thermo_targets/targets.json") as f:
+        targets = json.load(f)
+
+    with open(
+        "density_example/reference_force_fields/openff-2.3.0.offxml.json",
+    ) as f:
+        global_force_field = load_tensor_force_field(json.load(f))
+
+    predict(
+        dataset=targets,
+        force_field=global_force_field,
+        job_ids=[
+            "0f6ef5a4fdd601e764c9a70b185a6e27b1b91db6a8b65fb25a8a26922a3fca09",
+            "2ed40b18f3773e1a0eb756147e2b56e3da4c47149cdbf97c7bf6b85107df5c1c",
+            "5090adac8840d612600a2542c0236d47c2de3cdef6d6517961aa68e84f7ab8da",
+            "5a423d141a2e9d2d9ba9d30f988179ca8230bd8b93e7b854cc9c5c61f9ee9093",
+            "8418e9155a226c72f4d382d18e190ba023cbe55c9ad90175333d30637267059a",
+            "86cf94ed4a70093ce9d6d1f480eebaf806c457ca1c7347e5b2f153e59f12d59e",
+            "93e2b533fdec76225cd6af47938c69c40180d90a398a73b99892221471bcebc8",
+            "d721ef4f95b38106e81cc51e1d5638d9f09f372f4245cfe96ce2269d17e6a971",
+            "f08c269b4006de30978674f5ea1615258f471b5c21e045d57641d67ee86c3bd6",
+            "86cf94ed4a70093ce9d6d1f480eebaf806c457ca1c7347e5b2f153e59f12d59e",
+            "5090adac8840d612600a2542c0236d47c2de3cdef6d6517961aa68e84f7ab8da",
+            "2ed40b18f3773e1a0eb756147e2b56e3da4c47149cdbf97c7bf6b85107df5c1c",
+        ],
+        base_dir="density_example",
+    )
+
+    # TODO: Would be nice to dummy-check .backwards, can do this by constructing
+    # a minimal loss function of (predicted - expected) for each density value

tyff/_tests/compute/test_workflow.py

@@ -15,3 +15,5 @@ def test_init(self, tmp_path):
                 assert str(tmp_path) in lines[0]
 
             parsl.clear()
+
+        assert (tmp_path / "target_job_map.json").exists()

tyff/compute/_jacobian.py

@@ -22,10 +22,17 @@ def _gather_from_reference(local: tyff.TensorForceField, reference: tyff.TensorF
 
     for potential in local.potentials:
         ref = reference_by_type[potential.handler_type]  # KeyError: potential type missing from reference
-        assert potential.parameter_cols == ref.parameter_cols, potential.handler_type
+        assert potential.parameter_cols == ref.parameter_cols, (
+            f"parameter_cols doesn't match in potential {potential.handler_type=}"
+        )
 
         ref_idx = {key: i for i, key in enumerate(ref.parameter_keys)}
-        idx = torch.tensor([ref_idx[key] for key in potential.parameter_keys])  # KeyError: unseen parameter
+        try:
+            idx = torch.tensor([ref_idx[key] for key in potential.parameter_keys])  # KeyError: unseen parameter
+        except KeyError as error:
+            raise KeyError(
+                f"Parameter key {error} in local potential {potential.handler_type} not found in reference"
+            ) from error
 
         attributes = potential.attributes
         if potential.attribute_cols is not None: