From 834e75d177368d8c520ce1c879d64eee2ab2521b Mon Sep 17 00:00:00 2001 From: "Matthew W. Thompson" Date: Tue, 20 Oct 2026 17:14:45 -0500 Subject: [PATCH 1/5] Add energy and force predictions, loss optimization --- examples/mvp/run.py | 39 +++++++++++++++++++++++++++++++++++++-- 1 file changed, 37 insertions(+), 2 deletions(-) diff --git a/examples/mvp/run.py b/examples/mvp/run.py index 246245d..648babc 100644 --- a/examples/mvp/run.py +++ b/examples/mvp/run.py @@ -1,17 +1,21 @@ import tarfile +from json import dump from pathlib import Path from datasets import Dataset from openff.toolkit import ForceField from rich import print +from torch import no_grad +from torch.optim import Adam +from tyff._serialization import dump_tensor_force_field from tyff.compute._logging import _set_up_logger from tyff.converters import topologies_for from tyff.targets import energy, thermo from tyff.targets.new_thermo import predict, simulate from tyff.train import ParameterConfig, Trainable -N_EPOCHS = 1 +N_EPOCHS = 10 N_REPLICATES = 3 N_MOLECULES = 200 INITIAL_FORCE_FIELD = "openff-2.3.0.offxml" @@ -43,9 +47,11 @@ def main(root: Path) -> None: logger = _set_up_logger(f"{root}/mvp.log") + (root / "force_fields").mkdir(exist_ok=True) + thermo_data, energy_data = load_datasets(root) - _topologies, tensor_ff = topologies_for( + topologies, tensor_ff = topologies_for( smiles_set={ *thermo.extract_smiles(thermo_data), *energy.extract_smiles(energy_data), @@ -65,6 +71,7 @@ def main(root: Path) -> None: attributes={}, ) theta = trainable.to_values().requires_grad_(True) + optimizer = Adam([theta], lr=0.01) for epoch in range(N_EPOCHS): current_ff = trainable.to_force_field(theta) @@ -95,6 +102,34 @@ def main(root: Path) -> None: # 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=) + energy_ref, energy_pred, forces_ref, forces_pred = energy.predict( + energy_data, + current_ff, + topologies, + normalize=False, + ) + + loss = ( + (energy_pred - energy_ref).square().mean() + + 0.1 * (forces_pred - forces_ref).square().mean() + + (y_pred - y_ref).square().mean() + ) + loss.backward() + optimizer.step() + optimizer.zero_grad() + with no_grad(): + theta.copy_(trainable.clamp(theta)) + + # probably write out FF and end of each epoch? + with open(root / "force_fields" / f"after_EPOCH_{epoch}.ff.json", "w") as f: + dump( + dump_tensor_force_field( + trainable.to_force_field(theta), + ), + f, + indent=4, + ) + if __name__ == "__main__": main(root=Path(__file__).resolve().parent) From 7991891bb87261c78d24c3df65efc18a22dd43b6 Mon Sep 17 00:00:00 2001 From: "Matthew W. Thompson" Date: Tue, 20 Oct 2026 18:05:41 -0500 Subject: [PATCH 2/5] Minor cleanup, write force fields at end of each epoch --- examples/mvp/run.py | 33 ++++++++++++++++----------------- 1 file changed, 16 insertions(+), 17 deletions(-) diff --git a/examples/mvp/run.py b/examples/mvp/run.py index 648babc..029c99f 100644 --- a/examples/mvp/run.py +++ b/examples/mvp/run.py @@ -4,13 +4,12 @@ from datasets import Dataset from openff.toolkit import ForceField -from rich import print from torch import no_grad from torch.optim import Adam from tyff._serialization import dump_tensor_force_field from tyff.compute._logging import _set_up_logger -from tyff.converters import topologies_for +from tyff.converters import convert_tensor_force_field, topologies_for from tyff.targets import energy, thermo from tyff.targets.new_thermo import predict, simulate from tyff.train import ParameterConfig, Trainable @@ -73,7 +72,11 @@ def main(root: Path) -> None: theta = trainable.to_values().requires_grad_(True) 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]? current_ff = trainable.to_force_field(theta) # no return value? @@ -92,16 +95,6 @@ def main(root: Path) -> None: epoch_dir=root / f"EPOCH_{epoch}", ) - print(f"{y_ref=}\n{y_ref_std=}\n{y_pred=}\n{y_pred_std=}\n") - # y_ref=tensor([0.8736, 0.8060, 1.1335, 0.8579, 0.8215, 0.8375, 0.8543, 1.0950, 1.1203, - # 1.1290]) - # y_ref_std=tensor([4.0000e-05, 5.7000e-05, 5.0000e-04, 3.3500e-04, 1.1150e-04, 1.0100e-04, - # 9.1000e-05, 1.5000e-04, 3.5000e-04, 1.0000e-03]) - # y_pred=tensor([0.8859, 0.7895, 1.1050, 0.8637, 0.8140, 0.8355, 0.8586, 1.0800, 1.0986, - # 1.1028], dtype=torch.float64, grad_fn=) - # 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=) - energy_ref, energy_pred, forces_ref, forces_pred = energy.predict( energy_data, current_ff, @@ -114,22 +107,28 @@ def main(root: Path) -> None: + 0.1 * (forces_pred - forces_ref).square().mean() + (y_pred - y_ref).square().mean() ) - loss.backward() + optimizer.step() optimizer.zero_grad() with no_grad(): theta.copy_(trainable.clamp(theta)) - # probably write out FF and end of each epoch? + updated_tensor_force_field = trainable.to_force_field(theta) + + # write out FF as tensor with open(root / "force_fields" / f"after_EPOCH_{epoch}.ff.json", "w") as f: dump( - dump_tensor_force_field( - trainable.to_force_field(theta), - ), + dump_tensor_force_field(updated_tensor_force_field), f, indent=4, ) + # write out FF as SMIRNOFF force field + convert_tensor_force_field( + original_force_field=ForceField(INITIAL_FORCE_FIELD), + tensor_force_field=updated_tensor_force_field, + ).to_file(str(root / "force_fields" / f"after_EPOCH_{epoch}.offxml")) + if __name__ == "__main__": main(root=Path(__file__).resolve().parent) From ea4d1a02f31f628e91338f72c876b57e48ed0250 Mon Sep 17 00:00:00 2001 From: "Matthew W. Thompson" Date: Mon, 26 Oct 2026 15:28:44 -0500 Subject: [PATCH 3/5] Stage error --- 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 029c99f..c578bd2 100644 --- a/examples/mvp/run.py +++ b/examples/mvp/run.py @@ -15,7 +15,7 @@ from tyff.train import ParameterConfig, Trainable N_EPOCHS = 10 -N_REPLICATES = 3 +N_REPLICATES = 2 N_MOLECULES = 200 INITIAL_FORCE_FIELD = "openff-2.3.0.offxml" @@ -89,7 +89,7 @@ def main(root: Path) -> None: n_molecules=N_MOLECULES, ) - y_ref, y_ref_std, y_pred, y_pred_std = 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}", @@ -108,8 +108,17 @@ def main(root: Path) -> None: + (y_pred - y_ref).square().mean() ) - optimizer.step() + logger.info(f"EPOCH {epoch}: loss={loss.item():.6f}") + optimizer.zero_grad() + + loss.backward() + # ipdb> loss.backward() + # *** RuntimeError: Function RemoteEnsembleAverageBackward returned an invalid gradient at + # index 2 - got [208] but expected shape compatible with [287] + + optimizer.step() + with no_grad(): theta.copy_(trainable.clamp(theta)) From b0e83dc495116e4cf9d2c24ee16cfa43b3fc621e Mon Sep 17 00:00:00 2001 From: "Matthew W. Thompson" Date: Tue, 27 Oct 2026 10:48:09 -0500 Subject: [PATCH 4/5] Run fitting loop with only thermo data --- examples/mvp/run.py | 56 +++++++++++++++++++++++++++++------------ tyff/compute/configs.py | 2 +- tyff/mm/_ops.py | 1 + 3 files changed, 42 insertions(+), 17 deletions(-) diff --git a/examples/mvp/run.py b/examples/mvp/run.py index c578bd2..a402b58 100644 --- a/examples/mvp/run.py +++ b/examples/mvp/run.py @@ -19,6 +19,8 @@ N_MOLECULES = 200 INITIAL_FORCE_FIELD = "openff-2.3.0.offxml" +MERGE_DATA = False + def load_datasets(root: Path) -> tuple[Dataset, Dataset]: """Load datasets stored at data/*.tar.gz""" @@ -50,11 +52,15 @@ def main(root: Path) -> None: thermo_data, energy_data = load_datasets(root) + # tensor_ff is a subset of all parameters, capturing the parameters + # used in the datasets topologies, tensor_ff = topologies_for( 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), ) @@ -77,7 +83,21 @@ def main(root: Path) -> None: for epoch in range(N_EPOCHS): # maybe do this outside of loop for iteration 0, then # at end of loop for [1, ..., N_EPOCHS]? - current_ff = trainable.to_force_field(theta) + 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( @@ -95,27 +115,31 @@ def main(root: Path) -> None: epoch_dir=root / f"EPOCH_{epoch}", ) - energy_ref, energy_pred, forces_ref, forces_pred = energy.predict( - energy_data, - current_ff, - topologies, - normalize=False, - ) + if MERGE_DATA: + energy_ref, energy_pred, forces_ref, forces_pred = energy.predict( + energy_data, + current_ff, + topologies, + normalize=False, + ) - loss = ( - (energy_pred - energy_ref).square().mean() - + 0.1 * (forces_pred - forces_ref).square().mean() - + (y_pred - y_ref).square().mean() - ) + # 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() logger.info(f"EPOCH {epoch}: loss={loss.item():.6f}") + print(f"EPOCH {epoch}: loss={loss.item():.6f}") optimizer.zero_grad() loss.backward() - # ipdb> loss.backward() - # *** RuntimeError: Function RemoteEnsembleAverageBackward returned an invalid gradient at - # index 2 - got [208] but expected shape compatible with [287] optimizer.step() diff --git a/tyff/compute/configs.py b/tyff/compute/configs.py index 24c7d07..938344a 100644 --- a/tyff/compute/configs.py +++ b/tyff/compute/configs.py @@ -30,7 +30,7 @@ def local_config(max_workers: int = 10): cores_per_worker=1, ) ], - strategy=None, # Disable dynamic scaling for simpler local execution + strategy='htex_auto_scale', monitoring=MonitoringHub( resource_monitoring_interval=5, # seconds; 0 disables per-task resource sampling ), diff --git a/tyff/mm/_ops.py b/tyff/mm/_ops.py index c25b05a..a3e429a 100644 --- a/tyff/mm/_ops.py +++ b/tyff/mm/_ops.py @@ -571,4 +571,5 @@ def backward( grad_output: torch.Tensor, ) -> tuple[None, None, torch.Tensor]: (jacobian,) = ctx.saved_tensors + return None, None, grad_output @ jacobian From ac81b28990188fa7539136344b0f932240727ae1 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 27 Oct 2026 15:48:31 +0000 Subject: [PATCH 5/5] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tyff/compute/configs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tyff/compute/configs.py b/tyff/compute/configs.py index 938344a..303d1df 100644 --- a/tyff/compute/configs.py +++ b/tyff/compute/configs.py @@ -30,7 +30,7 @@ def local_config(max_workers: int = 10): cores_per_worker=1, ) ], - strategy='htex_auto_scale', + strategy="htex_auto_scale", monitoring=MonitoringHub( resource_monitoring_interval=5, # seconds; 0 disables per-task resource sampling ),