From eb689fbf907f3167dbd68d6ca9db2d5bffd52109 Mon Sep 17 00:00:00 2001
From: "Matthew W. Thompson" <mattwthompson@protonmail.com>
Date: Wed, 14 Oct 2026 12:30:09 -0500
Subject: [PATCH 01/14] Prototype `predict()` function

---
 prototype.py             | 170 +++++++++++++++++++++++++++++++++++++++
 tyff/compute/workflow.py |  20 ++++-
 2 files changed, 189 insertions(+), 1 deletion(-)
 create mode 100644 prototype.py

diff --git a/prototype.py b/prototype.py
new file mode 100644
index 0000000..01ecec6
--- /dev/null
+++ b/prototype.py
@@ -0,0 +1,170 @@
+import json
+import pathlib
+from collections.abc import Sequence
+
+import torch
+from rich import print
+
+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
+
+    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))
+
+    print(job_id_to_observables)
+
+    return job_id_to_observables
+
+
+def compute_properties(dataset: Sequence[DataEntry], observables):
+    from tyff.compute.prep import _compute_configs_from_data_entry
+
+    means_and_stds = []
+
+    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]
+
+            assert job_ids is not None
+            # 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']}")
+
+        mean = None
+        std = None
+        means_and_stds.append((mean, std.detach()))
+
+    return means_and_stds
+
+
+def predict(
+    dataset: Sequence[DataEntry],
+    force_field: tyff.TensorForceField,
+    job_ids: list[str],
+    base_dir: str,
+):
+    attached_means = load_remote_observables(force_field, base_dir, job_ids)
+
+    means_and_stds = compute_properties(dataset, attached_means)
+
+    return means_and_stds
+    # can either just return means_and_stds or can reformat into
+    # returning (ref_values, ref_stds, predicted_values, predicted_std)
+    # like other targets
+
+
+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",
+    )
diff --git a/tyff/compute/workflow.py b/tyff/compute/workflow.py
index 32f8f72..1da43cb 100644
--- a/tyff/compute/workflow.py
+++ b/tyff/compute/workflow.py
@@ -127,6 +127,20 @@ def submit_target_batch(
 
         from tyff._serialization import dump_tensor_force_field
 
+        # save these targets since they're probably going to be needed later
+        (pathlib.Path(self.base_dir) / "thermo_targets").mkdir(exist_ok=True)
+
+        if (pathlib.Path(self.base_dir) / "thermo_targets" / "targets.json").exists():
+            with open(pathlib.Path(self.base_dir) / "thermo_targets" / "targets.json") as f:
+                targets = json.load(f)
+        else:
+            targets = list()
+
+        targets.extend(target_configs)
+
+        with open(pathlib.Path(self.base_dir) / "thermo_targets" / "targets.json", "w") as f:
+            json.dump(targets, f, indent=4)
+
         unique_smiles = set()
 
         for target in target_configs:
@@ -157,7 +171,11 @@ def submit_target_batch(
 
         # dumps a JSON representation of the tensor representation of the SMIRNOFF force field
         with open(self._reference_force_fields / f"{force_field}.json", "w") as f:
-            json.dump(dump_tensor_force_field(reference_force_field), f)
+            json.dump(
+                dump_tensor_force_field(reference_force_field),
+                f,
+                indent=4,
+            )
 
         return [
             result

From fee02b62ff17ae70c1aa0b80c06963f77ee5eedd Mon Sep 17 00:00:00 2001
From: "Matthew W. Thompson" <mattwthompson@protonmail.com>
Date: Wed, 14 Oct 2026 13:07:42 -0500
Subject: [PATCH 02/14] Stage state

---
 prototype.py | 97 +++++++++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 93 insertions(+), 4 deletions(-)

diff --git a/prototype.py b/prototype.py
index 01ecec6..bc32e54 100644
--- a/prototype.py
+++ b/prototype.py
@@ -101,7 +101,98 @@ def compute_properties(dataset: Sequence[DataEntry], observables):
             # only one job per replicate in density calculatiosn
             job_ids = [make_job_id(replicate[0]) for replicate in compute_configs]
 
-            assert job_ids is not None
+            types_of_observables = [*observables[job_ids[0]].keys()]
+            """
+            ipdb> entry
+            {
+                'id': 74669191270958899039426367776551206167278900382279422138348522811937124722445,
+                'tag': 'density',
+                'smiles': ['COCCO'],
+                'x': [1.0],
+                'temperature': 293.15,
+                'pressure': 101.3,
+                'value': 0.9648800000000002,
+                'std': 5.0000000000000016e-05,
+                'units': 'gram / milliliter',
+                'source': ''
+            }
+            ipdb> job_ids
+            [
+                '5090adac8840d612600a2542c0236d47c2de3cdef6d6517961aa68e84f7ab8da',
+                '2ed40b18f3773e1a0eb756147e2b56e3da4c47149cdbf97c7bf6b85107df5c1c',
+                '86cf94ed4a70093ce9d6d1f480eebaf806c457ca1c7347e5b2f153e59f12d59e'
+            ]
+            ipdb> print([observables[job_id] for job_id in job_ids])
+            [
+                {
+                    'potential_energy': tensor(1835.6803, dtype=torch.float64),
+                    'potential_energy^2': tensor(3371036.2500, dtype=torch.float64),
+                    'volume': tensor(26207.9766, dtype=torch.float64),
+                    'volume^2': tensor(6.8692e+08, dtype=torch.float64),
+                    'density': tensor(0.9644, dtype=torch.float64),
+                    'enthalpy': tensor(3684.1626, dtype=torch.float64),
+                    'enthalpy^2': tensor(13574647., dtype=torch.float64),
+                    'enthalpy_volume': tensor(96557640., dtype=torch.float64)
+                },
+                {
+                    'potential_energy': tensor(1848.0491, dtype=torch.float64),
+                    'potential_energy^2': tensor(3417077., dtype=torch.float64),
+                    'volume': tensor(26305.6504, dtype=torch.float64),
+                    'volume^2': tensor(6.9209e+08, dtype=torch.float64),
+                    'density': tensor(0.9608, dtype=torch.float64),
+                    'enthalpy': tensor(3695.4373, dtype=torch.float64),
+                    'enthalpy^2': tensor(13658783., dtype=torch.float64),
+                    'enthalpy_volume': tensor(97218952., dtype=torch.float64)
+                },
+                {
+                    'potential_energy': tensor(1837.3766, dtype=torch.float64),
+                    'potential_energy^2': tensor(3377881.2500, dtype=torch.float64),
+                    'volume': tensor(26262.2070, dtype=torch.float64),
+                    'volume^2': tensor(6.8976e+08, dtype=torch.float64),
+                    'density': tensor(0.9624, dtype=torch.float64),
+                    'enthalpy': tensor(3681.6003, dtype=torch.float64),
+                    'enthalpy^2': tensor(13557029., dtype=torch.float64),
+                    'enthalpy_volume': tensor(96692128., dtype=torch.float64)
+                }
+            ]
+            """
+            # do we just return the mean and std for each observable this list?
+            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
+            }
+
+            """
+            not sure if these are the right things to return?
+
+            ipdb> print(mean)
+            {
+                'potential_energy': tensor(1840.3687, dtype=torch.float64),
+                'potential_energy^2': tensor(3388664.8333, dtype=torch.float64),
+                'volume': tensor(26258.6113, dtype=torch.float64),
+                'volume^2': tensor(6.8959e+08, dtype=torch.float64),
+                'density': tensor(0.9625, dtype=torch.float64),
+                'enthalpy': tensor(3687.0667, dtype=torch.float64),
+                'enthalpy^2': tensor(13596819.6667, dtype=torch.float64),
+                'enthalpy_volume': tensor(96822906.6667, dtype=torch.float64)
+            }
+            ipdb> print(std)
+            {
+                'potential_energy': tensor(6.7053, dtype=torch.float64),
+                'potential_energy^2': tensor(24842.5425, dtype=torch.float64),
+                'volume': tensor(48.9361, dtype=torch.float64),
+                'volume^2': tensor(2586491.9573, dtype=torch.float64),
+                'density': tensor(0.0018, dtype=torch.float64),
+                'enthalpy': tensor(7.3614, dtype=torch.float64),
+                'enthalpy^2': tensor(54380.0468, dtype=torch.float64),
+                'enthalpy_volume': tensor(349514.9283, dtype=torch.float64)
+            }
+            """
+
             # figure out which job id/s (if multiple replicas) in the observable corresponds to
             # and calculate means and stds
         elif entry["tag"] == "dielectric_constant":
@@ -117,9 +208,7 @@ def compute_properties(dataset: Sequence[DataEntry], observables):
         else:
             raise ValueError(f"I do not know what to do with {entry['tag']}")
 
-        mean = None
-        std = None
-        means_and_stds.append((mean, std.detach()))
+        means_and_stds.append((mean, std))
 
     return means_and_stds
 

From ffa608c509af85dcb3584536ad90907d6cb453c9 Mon Sep 17 00:00:00 2001
From: "Matthew W. Thompson" <mattwthompson@protonmail.com>
Date: Wed, 14 Oct 2026 17:04:12 -0500
Subject: [PATCH 03/14] Updates from review

Co-authored-by: Lily Wang <31115101+lilyminium@users.noreply.github.com>
---
 prototype.py | 101 +++++++--------------------------------------------
 1 file changed, 14 insertions(+), 87 deletions(-)

diff --git a/prototype.py b/prototype.py
index bc32e54..8c67bb0 100644
--- a/prototype.py
+++ b/prototype.py
@@ -3,7 +3,6 @@
 from collections.abc import Sequence
 
 import torch
-from rich import print
 
 import tyff
 import tyff.converters
@@ -33,6 +32,7 @@ def load_remote_observables(
         [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(
@@ -76,15 +76,16 @@ def backward(
 
         job_id_to_observables[job_id] = dict(zip(names, attached_means))
 
-    print(job_id_to_observables)
-
     return job_id_to_observables
 
 
-def compute_properties(dataset: Sequence[DataEntry], 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 = []
+    means_and_stds: list[tuple[dict, dict]] = list()
 
     for entry in dataset:
         if entry["tag"] == "density":
@@ -102,61 +103,9 @@ def compute_properties(dataset: Sequence[DataEntry], observables):
             job_ids = [make_job_id(replicate[0]) for replicate in compute_configs]
 
             types_of_observables = [*observables[job_ids[0]].keys()]
-            """
-            ipdb> entry
-            {
-                'id': 74669191270958899039426367776551206167278900382279422138348522811937124722445,
-                'tag': 'density',
-                'smiles': ['COCCO'],
-                'x': [1.0],
-                'temperature': 293.15,
-                'pressure': 101.3,
-                'value': 0.9648800000000002,
-                'std': 5.0000000000000016e-05,
-                'units': 'gram / milliliter',
-                'source': ''
-            }
-            ipdb> job_ids
-            [
-                '5090adac8840d612600a2542c0236d47c2de3cdef6d6517961aa68e84f7ab8da',
-                '2ed40b18f3773e1a0eb756147e2b56e3da4c47149cdbf97c7bf6b85107df5c1c',
-                '86cf94ed4a70093ce9d6d1f480eebaf806c457ca1c7347e5b2f153e59f12d59e'
-            ]
-            ipdb> print([observables[job_id] for job_id in job_ids])
-            [
-                {
-                    'potential_energy': tensor(1835.6803, dtype=torch.float64),
-                    'potential_energy^2': tensor(3371036.2500, dtype=torch.float64),
-                    'volume': tensor(26207.9766, dtype=torch.float64),
-                    'volume^2': tensor(6.8692e+08, dtype=torch.float64),
-                    'density': tensor(0.9644, dtype=torch.float64),
-                    'enthalpy': tensor(3684.1626, dtype=torch.float64),
-                    'enthalpy^2': tensor(13574647., dtype=torch.float64),
-                    'enthalpy_volume': tensor(96557640., dtype=torch.float64)
-                },
-                {
-                    'potential_energy': tensor(1848.0491, dtype=torch.float64),
-                    'potential_energy^2': tensor(3417077., dtype=torch.float64),
-                    'volume': tensor(26305.6504, dtype=torch.float64),
-                    'volume^2': tensor(6.9209e+08, dtype=torch.float64),
-                    'density': tensor(0.9608, dtype=torch.float64),
-                    'enthalpy': tensor(3695.4373, dtype=torch.float64),
-                    'enthalpy^2': tensor(13658783., dtype=torch.float64),
-                    'enthalpy_volume': tensor(97218952., dtype=torch.float64)
-                },
-                {
-                    'potential_energy': tensor(1837.3766, dtype=torch.float64),
-                    'potential_energy^2': tensor(3377881.2500, dtype=torch.float64),
-                    'volume': tensor(26262.2070, dtype=torch.float64),
-                    'volume^2': tensor(6.8976e+08, dtype=torch.float64),
-                    'density': tensor(0.9624, dtype=torch.float64),
-                    'enthalpy': tensor(3681.6003, dtype=torch.float64),
-                    'enthalpy^2': tensor(13557029., dtype=torch.float64),
-                    'enthalpy_volume': tensor(96692128., dtype=torch.float64)
-                }
-            ]
-            """
-            # do we just return the mean and std for each observable this list?
+
+            # 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
@@ -166,32 +115,7 @@ def compute_properties(dataset: Sequence[DataEntry], observables):
                 for key in types_of_observables
             }
 
-            """
-            not sure if these are the right things to return?
-
-            ipdb> print(mean)
-            {
-                'potential_energy': tensor(1840.3687, dtype=torch.float64),
-                'potential_energy^2': tensor(3388664.8333, dtype=torch.float64),
-                'volume': tensor(26258.6113, dtype=torch.float64),
-                'volume^2': tensor(6.8959e+08, dtype=torch.float64),
-                'density': tensor(0.9625, dtype=torch.float64),
-                'enthalpy': tensor(3687.0667, dtype=torch.float64),
-                'enthalpy^2': tensor(13596819.6667, dtype=torch.float64),
-                'enthalpy_volume': tensor(96822906.6667, dtype=torch.float64)
-            }
-            ipdb> print(std)
-            {
-                'potential_energy': tensor(6.7053, dtype=torch.float64),
-                'potential_energy^2': tensor(24842.5425, dtype=torch.float64),
-                'volume': tensor(48.9361, dtype=torch.float64),
-                'volume^2': tensor(2586491.9573, dtype=torch.float64),
-                'density': tensor(0.0018, dtype=torch.float64),
-                'enthalpy': tensor(7.3614, dtype=torch.float64),
-                'enthalpy^2': tensor(54380.0468, dtype=torch.float64),
-                'enthalpy_volume': tensor(349514.9283, dtype=torch.float64)
-            }
-            """
+            means_and_stds.append((mean, std))
 
             # figure out which job id/s (if multiple replicas) in the observable corresponds to
             # and calculate means and stds
@@ -226,7 +150,7 @@ def predict(
     return means_and_stds
     # can either just return means_and_stds or can reformat into
     # returning (ref_values, ref_stds, predicted_values, predicted_std)
-    # like other targets
+    # like other targets - reference values can be taken straight from DataEntrys
 
 
 if __name__ == "__main__":
@@ -257,3 +181,6 @@ def predict(
         ],
         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

From 0280ddac771726681564e4cc60afb76d31dc1ea7 Mon Sep 17 00:00:00 2001
From: "Matthew W. Thompson" <mattwthompson@protonmail.com>
Date: Wed, 14 Oct 2026 17:16:57 -0500
Subject: [PATCH 04/14] Also return std dev

---
 prototype.py | 24 ++++++++++++++++++++----
 1 file changed, 20 insertions(+), 4 deletions(-)

diff --git a/prototype.py b/prototype.py
index 8c67bb0..11fa553 100644
--- a/prototype.py
+++ b/prototype.py
@@ -115,8 +115,6 @@ def compute_properties(
                 for key in types_of_observables
             }
 
-            means_and_stds.append((mean, std))
-
             # figure out which job id/s (if multiple replicas) in the observable corresponds to
             # and calculate means and stds
         elif entry["tag"] == "dielectric_constant":
@@ -145,9 +143,27 @@ def predict(
 ):
     attached_means = load_remote_observables(force_field, base_dir, job_ids)
 
-    means_and_stds = compute_properties(dataset, attached_means)
+    means_and_stds: list[dict[str, torch.Tensor]] = compute_properties(dataset, attached_means)
 
-    return means_and_stds
+    # return list of 4-long tuples or 4 tuples of list-likes?
+
+    return_values: list[tuple[float, float, float, float]] = list()
+    # putting a lot of trust in ordering here (and I guess everywhere?)
+
+    for reference, predicted in zip(dataset, means_and_stds):
+        return_values.append(
+            (
+                (
+                    reference["value"],
+                    reference["std"],
+                    predicted[0],
+                    predicted[1],
+                )
+            ),
+        )
+
+    # unwrap list of 4-long tuples to 4 separate tuples?
+    return return_values
     # can either just return means_and_stds or can reformat into
     # returning (ref_values, ref_stds, predicted_values, predicted_std)
     # like other targets - reference values can be taken straight from DataEntrys

From 0cc62e1cfd417687d1dca82813e9c8ecd5952f9d Mon Sep 17 00:00:00 2001
From: "Matthew W. Thompson" <mattwthompson@protonmail.com>
Date: Wed, 14 Oct 2026 17:59:27 -0500
Subject: [PATCH 05/14] Reshape return values of `predict`

---
 prototype.py | 33 +++++++++++----------------------
 1 file changed, 11 insertions(+), 22 deletions(-)

diff --git a/prototype.py b/prototype.py
index 11fa553..e9ad198 100644
--- a/prototype.py
+++ b/prototype.py
@@ -140,33 +140,22 @@ def predict(
     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)
 
-    # return list of 4-long tuples or 4 tuples of list-likes?
-
-    return_values: list[tuple[float, float, float, float]] = list()
-    # putting a lot of trust in ordering here (and I guess everywhere?)
-
-    for reference, predicted in zip(dataset, means_and_stds):
-        return_values.append(
-            (
-                (
-                    reference["value"],
-                    reference["std"],
-                    predicted[0],
-                    predicted[1],
-                )
-            ),
-        )
+    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)
 
-    # unwrap list of 4-long tuples to 4 separate tuples?
-    return return_values
-    # can either just return means_and_stds or can reformat into
-    # returning (ref_values, ref_stds, predicted_values, predicted_std)
-    # like other targets - reference values can be taken straight from DataEntrys
+    return ref_values, ref_stds, predicted_values, predicted_stds
 
 
 if __name__ == "__main__":

From 7810f1d733082f5cc9860489a9e84e46cd19f124 Mon Sep 17 00:00:00 2001
From: "Matthew W. Thompson" <mattwthompson@protonmail.com>
Date: Thu, 15 Oct 2026 11:55:04 -0500
Subject: [PATCH 06/14] Move functions from prototype to internal

---
 tyff/mm/_ops.py            |  19 ++++++
 tyff/targets/new_thermo.py | 132 +++++++++++++++++++++++++++++++++++++
 2 files changed, 151 insertions(+)
 create mode 100644 tyff/targets/new_thermo.py

diff --git a/tyff/mm/_ops.py b/tyff/mm/_ops.py
index f652a97..d416020 100644
--- a/tyff/mm/_ops.py
+++ b/tyff/mm/_ops.py
@@ -552,3 +552,22 @@ def reweight_ensemble_averages(
 
     *avg_outputs, columns = _ReweightAverageOp.apply(kwargs, *tensors)
     return {column: avg for avg, column in zip(avg_outputs, columns, strict=True)}
+
+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
diff --git a/tyff/targets/new_thermo.py b/tyff/targets/new_thermo.py
new file mode 100644
index 0000000..1bb2b0a
--- /dev/null
+++ b/tyff/targets/new_thermo.py
@@ -0,0 +1,132 @@
+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
+
+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, RemoteEnsembleAverage
+
+    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
+
+    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

From ecc5cd3a5bab6ffb3db0696bd9ac0840e1b26d2b Mon Sep 17 00:00:00 2001
From: "Matthew W. Thompson" <mattwthompson@protonmail.com>
Date: Thu, 15 Oct 2026 12:27:35 -0500
Subject: [PATCH 07/14] Store target-compute mapping on exit

---
 tyff/_tests/compute/test_workflow.py |  2 ++
 tyff/compute/workflow.py             | 30 +++++++++++++++++++++++-----
 2 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/tyff/_tests/compute/test_workflow.py b/tyff/_tests/compute/test_workflow.py
index b4781c8..52bf83e 100644
--- a/tyff/_tests/compute/test_workflow.py
+++ b/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()
diff --git a/tyff/compute/workflow.py b/tyff/compute/workflow.py
index 1da43cb..d192e81 100644
--- a/tyff/compute/workflow.py
+++ b/tyff/compute/workflow.py
@@ -40,7 +40,7 @@ def __init__(self, base_dir: str, parsl_config: parsl.Config):
 
         # at scale these should become databases or something more robust
         self._targets = dict()
-        self._target_compute_mapping: dict[tuple[int, str, int, int], list[tuple[str, ...]]] = dict()
+        self._target_compute_mapping: dict[tuple[str, str, int, int], list[tuple[str, ...]]] = dict()
 
         # self.logger, maybe?
         logger = _set_up_logger(f"{base_dir}/workflow.log")
@@ -107,11 +107,23 @@ def submit_target(
         )
 
         # track targets (+ other arg) mapped to compute runs (as job_id) so we can look up results later
-        self._target_compute_mapping[(target_config["id"], force_field, n_molecules, n_replicates)] = list()
-        for this_targets_compute_configs in compute_configs:
-            self._target_compute_mapping[(target_config["id"], force_field, n_molecules, n_replicates)].append(
-                tuple([make_job_id(compute_config) for compute_config in this_targets_compute_configs])
+        self._target_compute_mapping[
+            (
+                target_config["id"],
+                force_field,
+                n_molecules,
+                n_replicates,
             )
+        ] = list()
+        for this_targets_compute_configs in compute_configs:
+            self._target_compute_mapping[
+                (
+                    target_config["id"],
+                    force_field,
+                    n_molecules,
+                    n_replicates,
+                )
+            ].append(tuple([make_job_id(compute_config) for compute_config in this_targets_compute_configs]))
 
             # run each compute job - can be >1 compute job per property
             self.run(compute_configs=this_targets_compute_configs)
@@ -336,4 +348,12 @@ def __enter__(self):
         return self
 
     def __exit__(self, *args):
+        with open(pathlib.Path(self.base_dir) / "target_job_map.json", "w") as f:
+            json.dump(
+                # keys are tuples, so just serialize key-val pairs as lists
+                [[list(k), v] for k, v in self._target_compute_mapping.items()],
+                f,
+                indent=4,
+            )
+
         self.shutdown()

From 6162cae82778620588d2a4134aad166c7a5499fe Mon Sep 17 00:00:00 2001
From: "Matthew W. Thompson" <mattwthompson@protonmail.com>
Date: Thu, 15 Oct 2026 12:44:22 -0500
Subject: [PATCH 08/14] Update to use as-submitted `force_field` and
 `n_molecules`

---
 tyff/mm/_ops.py            |  1 +
 tyff/targets/new_thermo.py | 30 ++++++++++++++++--------------
 2 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/tyff/mm/_ops.py b/tyff/mm/_ops.py
index d416020..c25b05a 100644
--- a/tyff/mm/_ops.py
+++ b/tyff/mm/_ops.py
@@ -553,6 +553,7 @@ def reweight_ensemble_averages(
     *avg_outputs, columns = _ReweightAverageOp.apply(kwargs, *tensors)
     return {column: avg for avg, column in zip(avg_outputs, columns, strict=True)}
 
+
 class RemoteEnsembleAverage(torch.autograd.Function):
     @staticmethod
     def forward(
diff --git a/tyff/targets/new_thermo.py b/tyff/targets/new_thermo.py
index 1bb2b0a..f4284fa 100644
--- a/tyff/targets/new_thermo.py
+++ b/tyff/targets/new_thermo.py
@@ -6,18 +6,16 @@
 
 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
 
+
 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, RemoteEnsembleAverage
+    from tyff.mm._ops import RemoteEnsembleAverage, _pack_force_field
 
     tensors, *_ = _pack_force_field(
         force_field_per_epoch
@@ -56,25 +54,25 @@ def load_remote_observables(
 def compute_properties(
     dataset: Sequence[DataEntry],
     observables: dict[str, dict[str, torch.Tensor]],
+    simulation_root: str,
 ) -> list[tuple[dict, dict]]:
-    from tyff.compute.prep import _compute_configs_from_data_entry
 
     means_and_stds: list[tuple[dict, dict]] = list()
 
+    with open(pathlib.Path(simulation_root) / "target_job_map.json") as f:
+        target_compute_mapping = {tuple(k): v for k, v in json.load(f)}
+
     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]
+            compute_configs = list()
+
+            for as_submitted, job_ids in target_compute_mapping.items():
+                if as_submitted[0] == entry["id"]:
+                    compute_configs.append(job_ids)
 
             types_of_observables = [*observables[job_ids[0]].keys()]
 
@@ -117,7 +115,11 @@ def predict(
 ) -> 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)
+    means_and_stds: list[dict[str, torch.Tensor]] = compute_properties(
+        dataset=dataset,
+        observables=attached_means,
+        simulation_root=base_dir,
+    )
 
     ref_values = [entry["value"] for entry in dataset]
     ref_stds = [entry["std"] for entry in dataset]

From dd4da08fdf7fcd1705a8ddd04babb9ba4e44efbe Mon Sep 17 00:00:00 2001
From: "Matthew W. Thompson" <mattwthompson@protonmail.com>
Date: Mon, 19 Oct 2026 12:16:58 -0500
Subject: [PATCH 09/14] Fix typing

---
 tyff/compute/workflow.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tyff/compute/workflow.py b/tyff/compute/workflow.py
index 9cc6baf..ac9c5dc 100644
--- a/tyff/compute/workflow.py
+++ b/tyff/compute/workflow.py
@@ -169,7 +169,7 @@ def submit_dataset(
         else:
             targets = list()
 
-        targets.extend(target_configs)
+        targets.extend(*dataset)
 
         with open(pathlib.Path(self.base_dir) / "thermo_targets" / "targets.json", "w") as f:
             json.dump(targets, f, indent=4)

From 7c277b8b29a8d4944c5e5e32ff8b8e289edd2e4e Mon Sep 17 00:00:00 2001
From: "Matthew W. Thompson" <mattwthompson@protonmail.com>
Date: Mon, 19 Oct 2026 14:22:37 -0500
Subject: [PATCH 10/14] Stage working changes

---
 examples/mvp/run.py      | 161 ++++++++++++++++++++++++++++++++++++++-
 tyff/compute/workflow.py |  21 ++---
 2 files changed, 171 insertions(+), 11 deletions(-)

diff --git a/examples/mvp/run.py b/examples/mvp/run.py
index 90ede52..09310f5 100644
--- a/examples/mvp/run.py
+++ b/examples/mvp/run.py
@@ -1,6 +1,8 @@
+import json
 import tarfile
 from pathlib import Path
 
+import torch
 from datasets import Dataset
 from openff.toolkit import ForceField, Molecule
 
@@ -10,8 +12,11 @@
 from tyff.compute._logging import _set_up_logger
 from tyff.compute.configs import local_config
 from tyff.compute.workflow import SimulationWorkflow
+from tyff.configs.targets.thermo import DataEntry
 from tyff.converters.openff._tensors import convert_tensor_force_field
+from tyff.mm._ops import RemoteEnsembleAverage
 from tyff.targets import energy, thermo
+from tyff.targets.thermo import DataEntry
 
 N_EPOCHS = 1
 N_REPLICATES = 1
@@ -111,6 +116,153 @@ def simulate(
         workflow.wait()
 
 
+def load_remote_observables(
+    force_field_per_epoch: tyff.TensorForceField,
+    base_dir: Path,
+    job_ids: list[str],
+) -> dict[str, dict[str, torch.Tensor]]:
+    """
+    Load observables and attach to graph.
+
+    Returns a mapping between job IDs and associated observables.
+    """
+    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
+
+    job_id_to_observables = {}
+
+    for job_id in job_ids:
+        job_dir = 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: Dataset,
+    observables: dict[str, dict[str, torch.Tensor]],
+    simulation_root: Path,
+) -> list[tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]]:
+
+    means_and_stds: list[tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]] = list()
+
+    with open(simulation_root / "target_job_map.json") as f:
+        target_compute_mapping = {tuple(k): v for k, v in json.load(f)}
+
+    for entry in dataset:  # type: ignore
+        entry: DataEntry
+        if entry["tag"] == "density":
+            compute_configs_for_target = list()
+
+            for as_submitted, job_ids in target_compute_mapping.items():
+                if as_submitted[0] == entry["id"]:
+                    compute_configs_for_target.append(job_ids)
+
+            # this is set with an iteration variable - can we assume that
+            # this will be the same for all job_ids within this target?
+            # all targets will be the same (i.e. always computing dhvap)
+            # but job-to-job maybe not? (i.e. gas vs. liquid jobs)
+            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: dict[str, torch.Tensor] = {
+                key: torch.stack([observables[job_id][key] for job_id in job_ids]).mean()
+                for key in types_of_observables
+            }
+            std: dict[str, torch.Tensor] = {
+                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":
+            import ipdb
+
+            ipdb.set_trace()
+            # 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":
+            import ipdb
+
+            ipdb.set_trace()
+        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: Dataset,
+    force_field: tyff.TensorForceField,
+    base_dir: Path,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    # in order to get all of the (remote) observables for all simulations run from this
+    # dataset, we need the job IDs of these
+
+    with open(base_dir / "target_job_map.json") as f:
+        target_compute_mapping = {tuple(k): v for k, v in json.load(f)}
+
+    this_datasets_ids: set[str] = set(entry["id"] for entry in dataset)  # type: ignore
+
+    job_ids: list[str] = list()
+
+    for (target_id, _, _, _), these_job_ids in target_compute_mapping.items():
+        if target_id in this_datasets_ids:
+            job_ids.extend(these_job_ids)
+
+    attached_observables = load_remote_observables(force_field, base_dir, job_ids)
+
+    means_and_stds: list[tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]] = compute_properties(
+        dataset=dataset,
+        observables=attached_observables,
+        simulation_root=base_dir,
+    )
+
+    ref_values = [entry["value"] for entry in dataset]  # type: ignore
+    ref_stds = [entry["std"] for entry in dataset]  # type: ignore
+    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
+
+
 def main(root: Path) -> None:
     global logger
 
@@ -131,7 +283,7 @@ def main(root: Path) -> None:
     trainable = tyff.train.Trainable(
         tensor_ff,
         parameters={
-            "vdW": tyff.train.ParameterConfig(
+            "vdW": tyff.train.ParameterConfig(  # type: ignore
                 cols=["epsilon", "sigma"],
                 scales={"epsilon": 10.0, "sigma": 1.0},
                 limits={"epsilon": (0.0001, None), "sigma": (0.5, None)},
@@ -152,6 +304,13 @@ def main(root: Path) -> None:
             n_molecules=N_MOLECULES,
         )
 
+        predict(
+            dataset=thermo_data,
+            force_field=current_ff,
+            base_dir=root,
+        )
+        # -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+
     return
 
 
diff --git a/tyff/compute/workflow.py b/tyff/compute/workflow.py
index ac9c5dc..9119c6d 100644
--- a/tyff/compute/workflow.py
+++ b/tyff/compute/workflow.py
@@ -125,7 +125,7 @@ def submit_target(
         # track targets (+ other arg) mapped to compute runs (as job_id) so we can look up results later
         self._target_compute_mapping[
             (
-                target_config["id"],
+                target_config["id"],  # this is unmodified from input data and therefore could look funky
                 force_field,
                 n_molecules,
                 n_replicates,
@@ -169,7 +169,8 @@ def submit_dataset(
         else:
             targets = list()
 
-        targets.extend(*dataset)
+        # list-ify dicts, which should quack like `DataEntry`s, via iterating over datasets
+        targets.extend([*dataset])
 
         with open(pathlib.Path(self.base_dir) / "thermo_targets" / "targets.json", "w") as f:
             json.dump(targets, f, indent=4)
@@ -211,6 +212,14 @@ def submit_dataset(
                 indent=4,
             )
 
+        with open(pathlib.Path(self.base_dir) / "target_job_map.json", "w") as f:
+            json.dump(
+                # keys are tuples, so just serialize key-val pairs as lists
+                [[list(k), v] for k, v in self._target_compute_mapping.items()],
+                f,
+                indent=4,
+            )
+
         return [
             result
             for spec in dataset
@@ -393,13 +402,5 @@ def __enter__(self):
         return self
 
     def __exit__(self, *args):
-        with open(pathlib.Path(self.base_dir) / "target_job_map.json", "w") as f:
-            json.dump(
-                # keys are tuples, so just serialize key-val pairs as lists
-                [[list(k), v] for k, v in self._target_compute_mapping.items()],
-                f,
-                indent=4,
-            )
-
         self.wait()
         self.shutdown()

From b311f853c50a5c5a7ec56bde0cdfe269a90e9ef0 Mon Sep 17 00:00:00 2001
From: "Matthew W. Thompson" <mattwthompson@protonmail.com>
Date: Mon, 19 Oct 2026 16:32:49 -0500
Subject: [PATCH 11/14] More minor updates from live testing

---
 examples/mvp/run.py      | 39 +++++++++++++++++++--------------------
 tyff/compute/workflow.py |  9 ++++++++-
 2 files changed, 27 insertions(+), 21 deletions(-)

diff --git a/examples/mvp/run.py b/examples/mvp/run.py
index 09310f5..7702ba6 100644
--- a/examples/mvp/run.py
+++ b/examples/mvp/run.py
@@ -12,7 +12,6 @@
 from tyff.compute._logging import _set_up_logger
 from tyff.compute.configs import local_config
 from tyff.compute.workflow import SimulationWorkflow
-from tyff.configs.targets.thermo import DataEntry
 from tyff.converters.openff._tensors import convert_tensor_force_field
 from tyff.mm._ops import RemoteEnsembleAverage
 from tyff.targets import energy, thermo
@@ -52,6 +51,8 @@ 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")
 
+    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
@@ -59,7 +60,7 @@ def load_datasets(root: Path) -> tuple[Dataset, Dataset]:
 
 def simulate(
     force_field: tyff.TensorForceField,
-    root_dir: Path,
+    epoch_dir: Path,
     thermo_data: Dataset,
     n_molecules: int,
     n_replicates: int,
@@ -94,16 +95,16 @@ def simulate(
     )
 
     logger.info("Making directory for this epoch's simulations")
-    root_dir.mkdir(exist_ok=True)
+    epoch_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")
+    smirnoff_force_field_file = str(epoch_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),
+        base_dir=str(epoch_dir / "simulations"),
+        parsl_config=local_config(max_workers=6),
     ) as workflow:
         logger.info("Submitting dataset")
         workflow.submit_dataset(
@@ -165,12 +166,13 @@ def load_remote_observables(
 def compute_properties(
     dataset: Dataset,
     observables: dict[str, dict[str, torch.Tensor]],
-    simulation_root: Path,
+    epoch_dir: Path,
 ) -> list[tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]]:
 
     means_and_stds: list[tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]] = list()
 
-    with open(simulation_root / "target_job_map.json") as f:
+    # think about better standardizing some of these paths
+    with open(epoch_dir / "simulations" / "target_job_map.json") as f:
         target_compute_mapping = {tuple(k): v for k, v in json.load(f)}
 
     for entry in dataset:  # type: ignore
@@ -186,6 +188,7 @@ def compute_properties(
             # this will be the same for all job_ids within this target?
             # all targets will be the same (i.e. always computing dhvap)
             # but job-to-job maybe not? (i.e. gas vs. liquid jobs)
+
             types_of_observables = [*observables[job_ids[0]].keys()]
 
             # could be simpler to return these as 1-D tensors - make sure to sort
@@ -205,16 +208,12 @@ def compute_properties(
             raise NotImplementedError()
             # likely identical to density case
         elif entry["tag"] == "enthalpy_of_mixing":
-            import ipdb
-
-            ipdb.set_trace()
+            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":
-            import ipdb
-
-            ipdb.set_trace()
+            raise NotImplementedError()
         else:
             raise ValueError(f"I do not know what to do with {entry['tag']}")
 
@@ -226,12 +225,12 @@ def compute_properties(
 def predict(
     dataset: Dataset,
     force_field: tyff.TensorForceField,
-    base_dir: Path,
+    epoch_dir: Path,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
     # in order to get all of the (remote) observables for all simulations run from this
     # dataset, we need the job IDs of these
 
-    with open(base_dir / "target_job_map.json") as f:
+    with open(epoch_dir / "simulations" / "target_job_map.json") as f:
         target_compute_mapping = {tuple(k): v for k, v in json.load(f)}
 
     this_datasets_ids: set[str] = set(entry["id"] for entry in dataset)  # type: ignore
@@ -242,12 +241,12 @@ def predict(
         if target_id in this_datasets_ids:
             job_ids.extend(these_job_ids)
 
-    attached_observables = load_remote_observables(force_field, base_dir, job_ids)
+    attached_observables = load_remote_observables(force_field, epoch_dir, job_ids)
 
     means_and_stds: list[tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]] = compute_properties(
         dataset=dataset,
         observables=attached_observables,
-        simulation_root=base_dir,
+        epoch_dir=epoch_dir,
     )
 
     ref_values = [entry["value"] for entry in dataset]  # type: ignore
@@ -298,7 +297,7 @@ def main(root: Path) -> None:
 
         simulate(
             force_field=current_ff,
-            root_dir=root / f"EPOCH_{epoch}",
+            epoch_dir=root / f"EPOCH_{epoch}",
             thermo_data=thermo_data,
             n_replicates=N_REPLICATES,
             n_molecules=N_MOLECULES,
@@ -307,7 +306,7 @@ def main(root: Path) -> None:
         predict(
             dataset=thermo_data,
             force_field=current_ff,
-            base_dir=root,
+            epoch_dir=root / f"EPOCH_{epoch}",
         )
         # -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
 
diff --git a/tyff/compute/workflow.py b/tyff/compute/workflow.py
index 9119c6d..ca4ec7b 100644
--- a/tyff/compute/workflow.py
+++ b/tyff/compute/workflow.py
@@ -52,7 +52,14 @@ def __init__(self, base_dir: str, parsl_config: parsl.Config):
 
         # at scale these should become databases or something more robust
         self._targets = dict()
-        self._target_compute_mapping: dict[tuple[str, str, int, int], list[tuple[str, ...]]] = dict()
+
+        if pathlib.Path(pathlib.Path(self.base_dir) / "simulations" / "target_job_map.json").exists():
+            with open(pathlib.Path(self.base_dir) / "simulations" / "target_job_map.json") as f:
+                # don't overwrite this if it already exists
+                self._target_compute_mapping = {tuple(k): v for k, v in json.load(f)}
+        else:
+            # initialize to empty dict otherwise
+            self._target_compute_mapping: dict[tuple[str, str, int, int], list[tuple[str, ...]]] = dict()
 
         # futures submitted via submit_target/submit_target_batch that haven't
         # been waited on yet; drained by wait()

From df0695015e10b0ba807dbaf2fffd4fffd87ce82f Mon Sep 17 00:00:00 2001
From: "Matthew W. Thompson" <mattwthompson@protonmail.com>
Date: Tue, 20 Oct 2026 15:43:05 -0500
Subject: [PATCH 12/14] Several bugfixes

---
 examples/mvp/run.py       | 30 +++++++++++++++----------
 tyff/compute/_jacobian.py | 11 ++++++++--
 tyff/compute/_prepare.py  |  4 ++--
 tyff/compute/workflow.py  | 46 +++++++++++++++++++++++++++------------
 4 files changed, 61 insertions(+), 30 deletions(-)

diff --git a/examples/mvp/run.py b/examples/mvp/run.py
index 7702ba6..4607776 100644
--- a/examples/mvp/run.py
+++ b/examples/mvp/run.py
@@ -18,7 +18,7 @@
 from tyff.targets.thermo import DataEntry
 
 N_EPOCHS = 1
-N_REPLICATES = 1
+N_REPLICATES = 3
 N_MOLECULES = 200
 INITIAL_FORCE_FIELD = "openff-2.3.0.offxml"
 
@@ -139,7 +139,7 @@ def load_remote_observables(
     job_id_to_observables = {}
 
     for job_id in job_ids:
-        job_dir = base_dir / job_id
+        job_dir = base_dir / "simulations" / job_id
 
         assert job_dir.exists()
 
@@ -178,27 +178,26 @@ def compute_properties(
     for entry in dataset:  # type: ignore
         entry: DataEntry
         if entry["tag"] == "density":
-            compute_configs_for_target = list()
-
+            # this lookup can be improved, but INITIAL_FORCE_FIELD
+            # gets mutated by this point so it's not so straightforward
             for as_submitted, job_ids in target_compute_mapping.items():
                 if as_submitted[0] == entry["id"]:
-                    compute_configs_for_target.append(job_ids)
+                    break
 
             # this is set with an iteration variable - can we assume that
             # this will be the same for all job_ids within this target?
             # all targets will be the same (i.e. always computing dhvap)
             # but job-to-job maybe not? (i.e. gas vs. liquid jobs)
-
-            types_of_observables = [*observables[job_ids[0]].keys()]
+            types_of_observables = [*observables[job_ids[0][0]].keys()]
 
             # could be simpler to return these as 1-D tensors - make sure to sort
             # keys to avoid ordering issues
             mean: dict[str, torch.Tensor] = {
-                key: torch.stack([observables[job_id][key] for job_id in job_ids]).mean()
+                key: torch.stack([observables[job_id[0]][key] for job_id in job_ids]).mean()
                 for key in types_of_observables
             }
             std: dict[str, torch.Tensor] = {
-                key: torch.stack([observables[job_id][key] for job_id in job_ids]).std()
+                key: torch.stack([observables[job_id[0]][key] for job_id in job_ids]).std()
                 for key in types_of_observables
             }
 
@@ -235,12 +234,19 @@ def predict(
 
     this_datasets_ids: set[str] = set(entry["id"] for entry in dataset)  # type: ignore
 
+    # this needs to be refactored to pass in _targets_, not jobs.
+    # synonymous with pure density, but for others (dhmix, dhvap, etc.) need to
+    # collect information from multiple jobs (i.e. gas and liquid jobs for dhvap)
+    # and condense down to one value (dhvap) to be reported
     job_ids: list[str] = list()
 
     for (target_id, _, _, _), these_job_ids in target_compute_mapping.items():
         if target_id in this_datasets_ids:
             job_ids.extend(these_job_ids)
 
+    # see comment above
+    job_ids = [x[0] for x in job_ids]
+
     attached_observables = load_remote_observables(force_field, epoch_dir, job_ids)
 
     means_and_stds: list[tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]] = compute_properties(
@@ -295,6 +301,7 @@ 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,
             epoch_dir=root / f"EPOCH_{epoch}",
@@ -303,14 +310,13 @@ def main(root: Path) -> None:
             n_molecules=N_MOLECULES,
         )
 
-        predict(
+        y_ref, y_ref_std, y_pred, y_pred_std = predict(
             dataset=thermo_data,
             force_field=current_ff,
             epoch_dir=root / f"EPOCH_{epoch}",
         )
-        # -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
 
-    return
+        print(y_ref, y_ref_std, y_pred, y_pred_std)
 
 
 if __name__ == "__main__":
diff --git a/tyff/compute/_jacobian.py b/tyff/compute/_jacobian.py
index 0eea661..2745a06 100644
--- a/tyff/compute/_jacobian.py
+++ b/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:
diff --git a/tyff/compute/_prepare.py b/tyff/compute/_prepare.py
index 3c2cc7a..b5fe853 100644
--- a/tyff/compute/_prepare.py
+++ b/tyff/compute/_prepare.py
@@ -43,7 +43,7 @@ def _prepare_openmm_system(
 
     packed_topology: Topology = Topology.from_pdb(
         file_path=packing_files["packed_topology"].filepath,
-        unique_molecules=[Molecule.from_smiles(smiles) for smiles in compute_config["smiles"]],
+        unique_molecules=[Molecule.from_mapped_smiles(smiles) for smiles in compute_config["smiles"]],
     )
 
     force_field = ForceField(compute_config["force_field"])
@@ -55,7 +55,7 @@ def _prepare_openmm_system(
     #
     # ordering is quite fragile, but hope order of smiles list never changes mid-job
     for unique_molecule_index, unique_molecule in enumerate(
-        [Molecule.from_smiles(smiles) for smiles in compute_config["smiles"]]
+        [Molecule.from_mapped_smiles(smiles) for smiles in compute_config["smiles"]]
     ):
         single_molecule_interchange = force_field.create_interchange(unique_molecule.to_topology())
         with open(f"{job_dir}/single_molecule_interchange_{unique_molecule_index}.json", "w") as f:
diff --git a/tyff/compute/workflow.py b/tyff/compute/workflow.py
index ca4ec7b..a8990be 100644
--- a/tyff/compute/workflow.py
+++ b/tyff/compute/workflow.py
@@ -53,10 +53,13 @@ def __init__(self, base_dir: str, parsl_config: parsl.Config):
         # at scale these should become databases or something more robust
         self._targets = dict()
 
-        if pathlib.Path(pathlib.Path(self.base_dir) / "simulations" / "target_job_map.json").exists():
-            with open(pathlib.Path(self.base_dir) / "simulations" / "target_job_map.json") as f:
+        target_job_map_path = pathlib.Path(self.base_dir) / "simulations" / "target_job_map.json"
+        if target_job_map_path.exists():
+            with open(target_job_map_path) as f:
                 # don't overwrite this if it already exists
-                self._target_compute_mapping = {tuple(k): v for k, v in json.load(f)}
+                loaded_mapping = {tuple(k): v for k, v in json.load(f).items()}
+                if len(loaded_mapping) > 0:
+                    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()
@@ -205,28 +208,33 @@ def submit_dataset(
                 f.write(interchange.model_dump_json(indent=4))
 
         self._reference_force_fields = pathlib.Path(self.base_dir) / "reference_force_fields"
-        pathlib.Path(self._reference_force_fields).mkdir(exist_ok=True)
+        self._reference_force_fields.mkdir(exist_ok=True)
 
+        assert self._reference_force_fields.exists()
+        assert self._reference_force_fields.is_dir()
+
+        logger.info(f"Made a path to store reference force fields at {self._reference_force_fields}")
         # this is a "global"/reference force field composed from all unique molecules,
         # don't think we're going to use the topologies here
         reference_force_field, _tensor_topologies = tyff.converters.convert_interchange(interchanges)
 
+        # make sure this is not empty
+        assert len(reference_force_field.potentials) > 0
+
+        logger.info(
+            "Writing reference force field from path (which is probably outside of "
+            f"this workflow's root): {force_field=}"
+            f"to path {self._reference_force_fields / 'reference.json'}"
+        )
+
         # dumps a JSON representation of the tensor representation of the SMIRNOFF force field
-        with open(self._reference_force_fields / f"{force_field}.json", "w") as f:
+        with open(self._reference_force_fields / "reference.json", "w") as f:
             json.dump(
                 dump_tensor_force_field(reference_force_field),
                 f,
                 indent=4,
             )
 
-        with open(pathlib.Path(self.base_dir) / "target_job_map.json", "w") as f:
-            json.dump(
-                # keys are tuples, so just serialize key-val pairs as lists
-                [[list(k), v] for k, v in self._target_compute_mapping.items()],
-                f,
-                indent=4,
-            )
-
         return [
             result
             for spec in dataset
@@ -301,6 +309,14 @@ def wait(self):
         made, so Parsl can actually run them concurrently in the meantime.
         Returns results/errors per job and clears the pending queue.
         """
+        with open(pathlib.Path(self.base_dir) / "target_job_map.json", "w") as f:
+            json.dump(
+                # keys are tuples, so just serialize key-val pairs as lists
+                [[list(k), v] for k, v in self._target_compute_mapping.items()],
+                f,
+                indent=4,
+            )
+
         results = self._drain(self._pending)
         self._pending = []
         return results
@@ -397,7 +413,9 @@ def _common_run(self, compute_config: BaseComputeConfig) -> tuple[dict[str, torc
         jacobian_future = create_jacobian(
             production_future=production_future,
             job_dir=job_dir,
-            reference_force_field_path=self._reference_force_fields / f"{compute_config['force_field']}.json",
+            # pathing here is weird, but there should only be one global reference
+            # force field for this workflow
+            reference_force_field_path=self._reference_force_fields / "reference.json",
         )
 
         return {"job_id": job_id, "future": jacobian_future}

From 6e6c127d4e51f447563cc054187f4c2be83438bc Mon Sep 17 00:00:00 2001
From: "Matthew W. Thompson" <mattwthompson@protonmail.com>
Date: Tue, 20 Oct 2026 16:07:17 -0500
Subject: [PATCH 13/14] Record results from run

---
 examples/mvp/run.py | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/examples/mvp/run.py b/examples/mvp/run.py
index 4607776..c60ac00 100644
--- a/examples/mvp/run.py
+++ b/examples/mvp/run.py
@@ -5,6 +5,7 @@
 import torch
 from datasets import Dataset
 from openff.toolkit import ForceField, Molecule
+from rich import print
 
 import tyff
 import tyff.converters
@@ -262,8 +263,8 @@ def predict(
 
     ref_values = torch.tensor(ref_values)
     ref_stds = torch.tensor(ref_stds)
-    predicted_values = torch.tensor(predicted_values)
-    predicted_stds = torch.tensor(predicted_stds)
+    predicted_values = torch.stack(predicted_values)
+    predicted_stds = torch.stack(predicted_stds)
 
     return ref_values, ref_stds, predicted_values, predicted_stds
 
@@ -316,7 +317,15 @@ def main(root: Path) -> None:
             epoch_dir=root / f"EPOCH_{epoch}",
         )
 
-        print(y_ref, y_ref_std, y_pred, y_pred_std)
+        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__":

From c18e0ce94ae5138af71908664d47776a24976e2e Mon Sep 17 00:00:00 2001
From: "Matthew W. Thompson" <mattwthompson@protonmail.com>
Date: Tue, 20 Oct 2026 16:29:45 -0500
Subject: [PATCH 14/14] Shuffle new functions into library

---
 examples/mvp/run.py                | 248 +----------------------------
 tyff/converters/__init__.py        |   2 +
 tyff/converters/openff/__init__.py |   8 +-
 tyff/converters/openff/_openff.py  |  15 ++
 tyff/targets/new_thermo.py         | 164 ++++++++++++++-----
 5 files changed, 159 insertions(+), 278 deletions(-)

diff --git a/examples/mvp/run.py b/examples/mvp/run.py
index c60ac00..246245d 100644
--- a/examples/mvp/run.py
+++ b/examples/mvp/run.py
@@ -1,22 +1,15 @@
-import json
 import tarfile
 from pathlib import Path
 
-import torch
 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.mm._ops import RemoteEnsembleAverage
+from tyff.converters import topologies_for
 from tyff.targets import energy, thermo
-from tyff.targets.thermo import DataEntry
+from tyff.targets.new_thermo import predict, simulate
+from tyff.train import ParameterConfig, Trainable
 
 N_EPOCHS = 1
 N_REPLICATES = 3
@@ -24,21 +17,6 @@
 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():
@@ -52,6 +30,7 @@ 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")
@@ -59,216 +38,6 @@ def load_datasets(root: Path) -> tuple[Dataset, Dataset]:
     return thermo_data, energy_data
 
 
-def simulate(
-    force_field: tyff.TensorForceField,
-    epoch_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")
-    epoch_dir.mkdir(exist_ok=True)
-
-    logger.info("Writing out the current SMIRNOFF force field to file")
-    smirnoff_force_field_file = str(epoch_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(epoch_dir / "simulations"),
-        parsl_config=local_config(max_workers=6),
-    ) 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 load_remote_observables(
-    force_field_per_epoch: tyff.TensorForceField,
-    base_dir: Path,
-    job_ids: list[str],
-) -> dict[str, dict[str, torch.Tensor]]:
-    """
-    Load observables and attach to graph.
-
-    Returns a mapping between job IDs and associated observables.
-    """
-    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
-
-    job_id_to_observables = {}
-
-    for job_id in job_ids:
-        job_dir = base_dir / "simulations" / 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: Dataset,
-    observables: dict[str, dict[str, torch.Tensor]],
-    epoch_dir: Path,
-) -> list[tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]]:
-
-    means_and_stds: list[tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]] = list()
-
-    # think about better standardizing some of these paths
-    with open(epoch_dir / "simulations" / "target_job_map.json") as f:
-        target_compute_mapping = {tuple(k): v for k, v in json.load(f)}
-
-    for entry in dataset:  # type: ignore
-        entry: DataEntry
-        if entry["tag"] == "density":
-            # this lookup can be improved, but INITIAL_FORCE_FIELD
-            # gets mutated by this point so it's not so straightforward
-            for as_submitted, job_ids in target_compute_mapping.items():
-                if as_submitted[0] == entry["id"]:
-                    break
-
-            # this is set with an iteration variable - can we assume that
-            # this will be the same for all job_ids within this target?
-            # all targets will be the same (i.e. always computing dhvap)
-            # but job-to-job maybe not? (i.e. gas vs. liquid jobs)
-            types_of_observables = [*observables[job_ids[0][0]].keys()]
-
-            # could be simpler to return these as 1-D tensors - make sure to sort
-            # keys to avoid ordering issues
-            mean: dict[str, torch.Tensor] = {
-                key: torch.stack([observables[job_id[0]][key] for job_id in job_ids]).mean()
-                for key in types_of_observables
-            }
-            std: dict[str, torch.Tensor] = {
-                key: torch.stack([observables[job_id[0]][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: Dataset,
-    force_field: tyff.TensorForceField,
-    epoch_dir: Path,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-    # in order to get all of the (remote) observables for all simulations run from this
-    # dataset, we need the job IDs of these
-
-    with open(epoch_dir / "simulations" / "target_job_map.json") as f:
-        target_compute_mapping = {tuple(k): v for k, v in json.load(f)}
-
-    this_datasets_ids: set[str] = set(entry["id"] for entry in dataset)  # type: ignore
-
-    # this needs to be refactored to pass in _targets_, not jobs.
-    # synonymous with pure density, but for others (dhmix, dhvap, etc.) need to
-    # collect information from multiple jobs (i.e. gas and liquid jobs for dhvap)
-    # and condense down to one value (dhvap) to be reported
-    job_ids: list[str] = list()
-
-    for (target_id, _, _, _), these_job_ids in target_compute_mapping.items():
-        if target_id in this_datasets_ids:
-            job_ids.extend(these_job_ids)
-
-    # see comment above
-    job_ids = [x[0] for x in job_ids]
-
-    attached_observables = load_remote_observables(force_field, epoch_dir, job_ids)
-
-    means_and_stds: list[tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]] = compute_properties(
-        dataset=dataset,
-        observables=attached_observables,
-        epoch_dir=epoch_dir,
-    )
-
-    ref_values = [entry["value"] for entry in dataset]  # type: ignore
-    ref_stds = [entry["std"] for entry in dataset]  # type: ignore
-    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.stack(predicted_values)
-    predicted_stds = torch.stack(predicted_stds)
-
-    return ref_values, ref_stds, predicted_values, predicted_stds
-
-
 def main(root: Path) -> None:
     global logger
 
@@ -276,8 +45,6 @@ def main(root: Path) -> None:
 
     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),
@@ -286,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(  # type: ignore
+            "vdW": ParameterConfig(  # type: ignore
                 cols=["epsilon", "sigma"],
                 scales={"epsilon": 10.0, "sigma": 1.0},
                 limits={"epsilon": (0.0001, None), "sigma": (0.5, None)},
@@ -307,6 +74,7 @@ def main(root: Path) -> None:
             force_field=current_ff,
             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,
         )
diff --git a/tyff/converters/__init__.py b/tyff/converters/__init__.py
index 82bfd1a..a24acf4 100644
--- a/tyff/converters/__init__.py
+++ b/tyff/converters/__init__.py
@@ -4,6 +4,7 @@
     convert_handlers,
     convert_interchange,
     smirnoff_parameter_converter,
+    topologies_for,
 )
 from tyff.converters.openff._tensors import convert_tensor_force_field
 from tyff.converters.openmm import (
@@ -24,4 +25,5 @@
     "convert_to_openmm_topology",
     "ffxml_converter",
     "smirnoff_parameter_converter",
+    "topologies_for",
 ]
diff --git a/tyff/converters/openff/__init__.py b/tyff/converters/openff/__init__.py
index ca44cb0..aa53e3e 100644
--- a/tyff/converters/openff/__init__.py
+++ b/tyff/converters/openff/__init__.py
@@ -4,6 +4,12 @@
     convert_handlers,
     convert_interchange,
     smirnoff_parameter_converter,
+    topologies_for,
 )
 
-__all__ = ["convert_handlers", "convert_interchange", "smirnoff_parameter_converter"]
+__all__ = (
+    "convert_handlers",
+    "convert_interchange",
+    "smirnoff_parameter_converter",
+    "topologies_for",
+)
diff --git a/tyff/converters/openff/_openff.py b/tyff/converters/openff/_openff.py
index 08e52a4..f8c9b04 100644
--- a/tyff/converters/openff/_openff.py
+++ b/tyff/converters/openff/_openff.py
@@ -494,3 +494,18 @@ def convert_interchange(
 
     tensor_force_field = tyff.TensorForceField(potentials, v_sites)
     return tensor_force_field, tensor_topologies
+
+
+def topologies_for(
+    smiles_set: set[str],
+    smirnoff_force_field: openff.toolkit.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(openff.toolkit.Molecule.from_mapped_smiles(smiles).to_topology())
+        for smiles in smiles_set
+    ]
+
+    tensor_ff, topologies = convert_interchange(interchanges)
+
+    return {smiles: topology for smiles, topology in zip(smiles_set, topologies)}, tensor_ff
diff --git a/tyff/targets/new_thermo.py b/tyff/targets/new_thermo.py
index f4284fa..853c7bb 100644
--- a/tyff/targets/new_thermo.py
+++ b/tyff/targets/new_thermo.py
@@ -1,21 +1,89 @@
 import json
-import pathlib
-from collections.abc import Sequence
+from pathlib import Path
 
 import torch
+from datasets import Dataset
+from openff.toolkit import ForceField
 
 import tyff
 import tyff.converters
-from tyff.configs.targets.thermo import DataEntry
+import tyff.train
+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.mm._ops import RemoteEnsembleAverage
+from tyff.targets.thermo import DataEntry
+
+
+def simulate(
+    force_field: tyff.TensorForceField,
+    epoch_dir: Path,
+    thermo_data: Dataset,
+    smirnoff_force_field: ForceField,
+    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.
+    epoch_dir : Path
+        The directory for this epoch where simulation results will be stored.
+    thermo_data : Dataset
+        The dataset containing thermodynamic data for the simulation.
+    smirnoff_force_field : ForceField
+        The original SMIRNOFF force field corresponding to the tensor force field.
+    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=smirnoff_force_field,
+        tensor_force_field=force_field,
+    )
+
+    epoch_dir.mkdir(exist_ok=True)
+
+    smirnoff_force_field_file = str(epoch_dir / "current_force_field.offxml")
+    current_smirnoff_force_field.to_file(smirnoff_force_field_file)
+
+    with SimulationWorkflow(
+        base_dir=str(epoch_dir / "simulations"),
+        parsl_config=local_config(max_workers=6),
+    ) as workflow:
+        workflow.submit_dataset(
+            dataset=thermo_data,
+            force_field=smirnoff_force_field_file,
+            n_molecules=n_molecules,
+            n_replicates=n_replicates,
+        )
+
+        workflow.wait()
 
 
 def load_remote_observables(
     force_field_per_epoch: tyff.TensorForceField,
-    base_dir: str,
+    base_dir: Path,
     job_ids: list[str],
-) -> dict[str, dict]:
-    """Load observables and attach to graph."""
-    from tyff.mm._ops import RemoteEnsembleAverage, _pack_force_field
+) -> dict[str, dict[str, torch.Tensor]]:
+    """
+    Load observables and attach to graph.
+
+    Returns a mapping between job IDs and associated observables.
+    """
+    from tyff.mm._ops import _pack_force_field
 
     tensors, *_ = _pack_force_field(
         force_field_per_epoch
@@ -27,7 +95,7 @@ def load_remote_observables(
     job_id_to_observables = {}
 
     for job_id in job_ids:
-        job_dir = pathlib.Path(base_dir) / job_id
+        job_dir = base_dir / "simulations" / job_id
 
         assert job_dir.exists()
 
@@ -52,38 +120,40 @@ def load_remote_observables(
 
 
 def compute_properties(
-    dataset: Sequence[DataEntry],
+    dataset: Dataset,
     observables: dict[str, dict[str, torch.Tensor]],
-    simulation_root: str,
-) -> list[tuple[dict, dict]]:
+    epoch_dir: Path,
+) -> list[tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]]:
 
-    means_and_stds: list[tuple[dict, dict]] = list()
+    means_and_stds: list[tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]] = list()
 
-    with open(pathlib.Path(simulation_root) / "target_job_map.json") as f:
+    # think about better standardizing some of these paths
+    with open(epoch_dir / "simulations" / "target_job_map.json") as f:
         target_compute_mapping = {tuple(k): v for k, v in json.load(f)}
 
-    for entry in dataset:
+    for entry in dataset:  # type: ignore
+        entry: DataEntry
         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 = list()
-
+            # this lookup can be improved, but INITIAL_FORCE_FIELD
+            # gets mutated by this point so it's not so straightforward
             for as_submitted, job_ids in target_compute_mapping.items():
                 if as_submitted[0] == entry["id"]:
-                    compute_configs.append(job_ids)
+                    break
 
-            types_of_observables = [*observables[job_ids[0]].keys()]
+            # this is set with an iteration variable - can we assume that
+            # this will be the same for all job_ids within this target?
+            # all targets will be the same (i.e. always computing dhvap)
+            # but job-to-job maybe not? (i.e. gas vs. liquid jobs)
+            types_of_observables = [*observables[job_ids[0][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()
+            mean: dict[str, torch.Tensor] = {
+                key: torch.stack([observables[job_id[0]][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()
+            std: dict[str, torch.Tensor] = {
+                key: torch.stack([observables[job_id[0]][key] for job_id in job_ids]).std()
                 for key in types_of_observables
             }
 
@@ -108,27 +178,47 @@ def compute_properties(
 
 
 def predict(
-    dataset: Sequence[DataEntry],
+    dataset: Dataset,
     force_field: tyff.TensorForceField,
-    job_ids: list[str],
-    base_dir: str,
+    epoch_dir: Path,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-    attached_means = load_remote_observables(force_field, base_dir, job_ids)
+    # in order to get all of the (remote) observables for all simulations run from this
+    # dataset, we need the job IDs of these
+
+    with open(epoch_dir / "simulations" / "target_job_map.json") as f:
+        target_compute_mapping = {tuple(k): v for k, v in json.load(f)}
+
+    this_datasets_ids: set[str] = set(entry["id"] for entry in dataset)  # type: ignore
+
+    # this needs to be refactored to pass in _targets_, not jobs.
+    # synonymous with pure density, but for others (dhmix, dhvap, etc.) need to
+    # collect information from multiple jobs (i.e. gas and liquid jobs for dhvap)
+    # and condense down to one value (dhvap) to be reported
+    job_ids: list[str] = list()
+
+    for (target_id, _, _, _), these_job_ids in target_compute_mapping.items():
+        if target_id in this_datasets_ids:
+            job_ids.extend(these_job_ids)
+
+    # see comment above
+    job_ids = [x[0] for x in job_ids]
+
+    attached_observables = load_remote_observables(force_field, epoch_dir, job_ids)
 
-    means_and_stds: list[dict[str, torch.Tensor]] = compute_properties(
+    means_and_stds: list[tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]] = compute_properties(
         dataset=dataset,
-        observables=attached_means,
-        simulation_root=base_dir,
+        observables=attached_observables,
+        epoch_dir=epoch_dir,
     )
 
-    ref_values = [entry["value"] for entry in dataset]
-    ref_stds = [entry["std"] for entry in dataset]
+    ref_values = [entry["value"] for entry in dataset]  # type: ignore
+    ref_stds = [entry["std"] for entry in dataset]  # type: ignore
     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)
+    predicted_values = torch.stack(predicted_values)
+    predicted_stds = torch.stack(predicted_stds)
 
     return ref_values, ref_stds, predicted_values, predicted_stds