Use thermo and energy errors in loss function

.pre-commit-config.yaml

@@ -10,17 +10,17 @@ repos:
       exclude: '\.(json)$'
     - id: debug-statements
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.17.2
+  rev: v0.17.5
   hooks:
     - id: ruff-check
       args: ["--fix"]
     - id: ruff-format
 - repo: https://github.com/tox-dev/pyproject-fmt
-  rev: v2.30.3
+  rev: v2.30.6
   hooks:
     - id: pyproject-fmt
 - repo: https://github.com/tombi-toml/tombi-pre-commit
-  rev: v1.9.6
+  rev: v1.11.2
   hooks:
     - id: tombi-format
       args: ["--offline"]

examples/mvp/run.py

@@ -19,7 +19,13 @@
 N_MOLECULES = 200
 INITIAL_FORCE_FIELD = "openff-2.3.0.offxml"
 
-MERGE_DATA = False
+# these weights should be exposed to the user in some sort of training config
+# probably a good phase II task
+COEFFICIENTS = {
+    "energy": 0.1,
+    "force": 0.01,
+    "thermo": 100_000.0,
+}
 
 
 def load_datasets(root: Path) -> tuple[Dataset, Dataset]:
@@ -58,9 +64,7 @@ def main(root: Path) -> None:
         smiles_set={
             *thermo.extract_smiles(thermo_data),
             *energy.extract_smiles(energy_data),
-        }
-        if MERGE_DATA
-        else {*thermo.extract_smiles(thermo_data)},
+        },
         smirnoff_force_field=ForceField(INITIAL_FORCE_FIELD),
     )
 
@@ -75,36 +79,23 @@ def main(root: Path) -> None:
         },
         attributes={},
     )
+
     theta = trainable.to_values().requires_grad_(True)
+
+    # TODO: Possible drop in L-M optimizer
     optimizer = Adam([theta], lr=0.01)
 
     # TODO: inner/outer loop, or any way of stepping through QM optimization
     #       without running MD jobs each time
     for epoch in range(N_EPOCHS):
-        # maybe do this outside of loop for iteration 0, then
-        # at end of loop for [1, ..., N_EPOCHS]?
+        # this is linked under the hood to theta - needed to gradients together
         current_ff = trainable.to_force_field(theta)  # 287
 
-        # import json
-        # from tyff._serialization import load_tensor_force_field
-
-        # sum(t.numel() for t in _pack_force_field(ref)[0] if t is not None)
-        # is probably 208
-
-        # current_ff = load_tensor_force_field(
-        #     json.loads(
-        #         open(
-        #             'examples/mvp/EPOCH_0/simulations/reference_force_fields/reference.json'
-        #         ).read()
-        #     )
-        # )
-
-        # no return value?
         simulate(
             force_field=current_ff,
             epoch_dir=root / f"EPOCH_{epoch}",
             thermo_data=thermo_data,
-            smirnoff_force_field=ForceField(INITIAL_FORCE_FIELD),
+            smirnoff_force_field=INITIAL_FORCE_FIELD,
             n_replicates=N_REPLICATES,
             n_molecules=N_MOLECULES,
         )
@@ -115,27 +106,26 @@ def main(root: Path) -> None:
             epoch_dir=root / f"EPOCH_{epoch}",
         )
 
-        if MERGE_DATA:
-            energy_ref, energy_pred, forces_ref, forces_pred = energy.predict(
-                energy_data,
-                current_ff,
-                topologies,
-                normalize=False,
-            )
+        energy_ref, energy_pred, forces_ref, forces_pred = energy.predict(
+            energy_data,
+            current_ff,
+            topologies,
+            normalize=False,
+        )
 
-            # the weights in front of each category of loss (energy, forces, and thermo)
-            # should be exposed to the user in some sort of training config
-            # probably a good phase II task
-            loss = (
-                (energy_pred - energy_ref).square().mean()
-                + 0.1 * (forces_pred - forces_ref).square().mean()
-                + (y_pred - y_ref).square().mean()
-            )
-        else:
-            loss = (y_pred - y_ref).square().mean()
+        # can construct loss function in many different ways, including modifying coefficients
+        loss = (
+            COEFFICIENTS["energy"] * (energy_pred - energy_ref).square().mean()
+            + COEFFICIENTS["force"] * (forces_pred - forces_ref).square().mean()
+            + COEFFICIENTS["thermo"] * (y_pred - y_ref).square().mean()
+        )
 
         logger.info(f"EPOCH {epoch}: loss={loss.item():.6f}")
         print(f"EPOCH {epoch}: loss={loss.item():.6f}")
+        print(f"\tloss component energy (after applying {COEFFICIENTS=})")
+        print(f"\tenergy {COEFFICIENTS['energy'] * (energy_pred - energy_ref).square().mean()}")
+        print(f"\tforce {COEFFICIENTS['force'] * (forces_pred - forces_ref).square().mean()}")
+        print(f"\tthermo {COEFFICIENTS['thermo'] * (y_pred - y_ref).square().mean()}")
 
         optimizer.zero_grad()
 
@@ -156,7 +146,7 @@ def main(root: Path) -> None:
                 indent=4,
             )
 
-        # write out FF as SMIRNOFF force field
+        # update SMIRNOFF force field with new values from tensor FF
         convert_tensor_force_field(
             original_force_field=ForceField(INITIAL_FORCE_FIELD),
             tensor_force_field=updated_tensor_force_field,