Compute and save ensemble averages and Jacobians for each job

.pre-commit-config.yaml

@@ -10,7 +10,7 @@ repos:
       exclude: '\.(json)$'
     - id: debug-statements
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.16.7
+  rev: v0.16.9
   hooks:
     - id: ruff-check
       args: ["--fix"]

examples/density/run.py

@@ -1,3 +1,4 @@
+import json
 import pathlib
 import socket
 
@@ -7,6 +8,8 @@
 
 from tyff.compute.configs import hpc3_config, local_config
 from tyff.compute.fetch import fetch_trajectory_paths_from_target
+from tyff.compute.jobs import make_job_id
+from tyff.compute.prep import _compute_configs_from_data_entry
 from tyff.compute.workflow import SimulationWorkflow
 from tyff.targets.thermo import DataEntry
 
@@ -73,6 +76,7 @@
             n_replicates=5,
         )
 
+        # naive property estimate, no tensor math
         for density_target in density_targets:
             workflow.estimate_target(
                 density_target,
@@ -94,6 +98,7 @@
             n_replicates=3,
         )
 
+        # naive property estimate, no tensor math
         for density_target in density_targets:
             workflow.estimate_target(
                 density_target,
@@ -104,6 +109,29 @@
 
 # TODO: Show how to check status while running
 
+# this generally happens under the hood
+compute_configs = [
+    _compute_configs_from_data_entry(
+        target,
+        force_field="openff-2.3.0.offxml",
+        n_molecules=200,
+    )
+    for target in density_targets
+]
+
+print(f"Number of targets for which compute configs were generated: {len(compute_configs)}")
+print(f"Number of replicates for each target: {[len(replicates) for replicates in compute_configs]}")
+
+# this is a list of job directories of each individual job (total of 9, 3x replicates of 3 targets)
+job_ids = [make_job_id(liquid_job) for replicate in compute_configs for row in replicate for liquid_job in row]
+
+for job_id in job_ids:
+    print(job_id)
+    print(json.loads(open(f"density_example/{job_id}/compute_config.json").read()))
+    print(json.loads(open(f"density_example/{job_id}/ensemble_averages.json").read()))
+
+# now for the old-school statistics, etc.
+
 # get trajectory paths from root job directory and per-target info (without knowing internal compute configs)
 trajectory_paths = [
     fetch_trajectory_paths_from_target(

pixi.toml

@@ -33,6 +33,12 @@ rdkit-stubs = "*"
 parsl = "*"
 rich = "*"
 tyff = { path = ".", editable = true }
+# probably optional-ish below here
+jupyter = "*"
+nbconvert = "*"
+tensorboard = "*"
+tensorboardx = "*"
+uncertainties = "*"
 
 [dependencies]
 openmm = "8.4.*" # need to help Pixi, since conda solve constraints PyPI solve
@@ -42,15 +48,12 @@ openff-nagl-base = "*"
 openff-nagl-models = "*"
 openmmtools = "*"
 pydantic-units = "*"
+sqlalchemy = ">=2.0.14,<2.1.0"
 pytorch = "*"
+NNPOps = "*"
 datasets = "*"
 bokeh = "3.9.*"
 # probably optional-ish below here
-jupyter = "*"
-nbconvert = "*"
-tensorboard = "*"
-tensorboardx = "*"
-uncertainties = "*"
 openff-packmol = "0.1.*"
 openff-evaluator-base = "*"
 # shims

tyff/_tests/compute/test_jacobians.py

@@ -0,0 +1,47 @@
+import shutil
+from importlib.resources import files
+
+import pytest
+from parsl import File
+
+from tyff.compute._files import ProductionFiles
+from tyff.compute._jacobian import _get_ensemble_average_and_jacobian
+
+
+@pytest.fixture
+def production_future() -> dict[str, ProductionFiles]:
+    return {
+        "simulation_files": ProductionFiles(
+            topology=File(files("tyff") / "_tests/data/app_files/sample_density/production_topology.pdb"),
+            dcd_trajectory=File(files("tyff") / "_tests/data/app_files/sample_density/production_trajectory.dcd"),
+            msgpack_trajectory=File(
+                files("tyff") / "_tests/data/app_files/sample_density/production_trajectory.msgpack"
+            ),
+            log=File(files("tyff") / "_tests/data/app_files/sample_density/produce.log"),
+            state_data=File(files("tyff") / "_tests/data/app_files/sample_density/production.csv"),
+            system=File(files("tyff") / "_tests/data/app_files/sample_density/production_system.xml"),
+            integrator=File(files("tyff") / "_tests/data/app_files/sample_density/production_integrator.xml"),
+            checkpoint=File(files("tyff") / "_tests/data/app_files/sample_density/production_checkpoint.chk"),
+        )
+    }
+
+
+def test_jacobian(production_future, tmp_path):
+    for file in [
+        "compute_config.json",
+        "production_trajectory.msgpack",
+        "single_molecule_interchange_0.json",
+        "single_molecule_interchange_1.json",
+    ]:
+        shutil.copy(
+            str(files("tyff") / f"_tests/data/app_files/sample_density/{file}"),
+            str(tmp_path / file),
+        )
+
+    averages, jacobians = _get_ensemble_average_and_jacobian(
+        production_future,
+        tmp_path,
+        reference_force_field_path=File(files("tyff") / "_tests/data/app_files/sample_density/ref.ff.json"),
+    )
+
+    assert len(averages) == jacobians.shape[0]

tyff/_tests/data/app_files/make_sample_density_files.py

@@ -2,8 +2,11 @@
 import pathlib
 import random
 
+import openff.toolkit
 from rich.pretty import pprint as print
 
+import tyff.converters
+from tyff._serialization import dump_tensor_force_field
 from tyff.compute._analyze import _run_density_analysis
 from tyff.compute._equilibrate import _run_equilibration
 from tyff.compute._minimize import _minimize_energy
@@ -15,9 +18,11 @@
 
 pathlib.Path("sample_density/").mkdir(exist_ok=True)
 
+FORCE_FIELD = "openff-2.3.0.offxml"
+
 target = DataEntry(
     **{
-        "id": random.randint(10**15, 10**16 - 1),
+        "id": random.randint(10**15, 10**16 - 1),  # TODO: Make this the real ID
         "tag": "density",
         "x": [0.5, 0.5],
         "smiles": [
@@ -33,12 +38,28 @@
     }
 )
 
-with open("sample_density/target_config.json", "w") as target_config:
-    json.dump(target, target_config, indent=4)
+with open("sample_density/target_config.json", "w") as target_config_file:
+    json.dump(target, target_config_file, indent=4)
+
+smirnoff_force_field = openff.toolkit.ForceField(FORCE_FIELD)
+
+interchanges = [
+    smirnoff_force_field.create_interchange(openff.toolkit.Molecule.from_smiles(smiles_).to_topology())
+    for smiles_ in target["smiles"]
+]
+
+for index, interchange in enumerate(interchanges):
+    with open(f"sample_density/single_molecule_interchange_{index}.json", "w") as f:
+        f.write(interchange.model_dump_json())
+
+reference_force_field, _ = tyff.converters.convert_interchange(interchanges)
+
+with open("sample_density/ref.ff.json", "w") as reference_force_field_file:
+    json.dump(dump_tensor_force_field(reference_force_field), reference_force_field_file, indent=4)
 
 compute = _make_liquid_density_compute_configs(
     data_entry=target,
-    force_field="openff-2.3.0.offxml",
+    force_field=FORCE_FIELD,
     n_molecules=200,
 )[0][0]