diff --git a/examples/mvp/run.py b/examples/mvp/run.py index 246245d..a402b58 100644 --- a/examples/mvp/run.py +++ b/examples/mvp/run.py @@ -1,21 +1,26 @@ 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.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 -N_EPOCHS = 1 -N_REPLICATES = 3 +N_EPOCHS = 10 +N_REPLICATES = 2 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""" @@ -43,13 +48,19 @@ 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( + # 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), ) @@ -65,9 +76,28 @@ def main(root: Path) -> None: attributes={}, ) 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): - current_ff = trainable.to_force_field(theta) + # 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) # 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( @@ -79,21 +109,58 @@ 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}", ) - 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=) + if MERGE_DATA: + 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() + + logger.info(f"EPOCH {epoch}: loss={loss.item():.6f}") + print(f"EPOCH {epoch}: loss={loss.item():.6f}") + + optimizer.zero_grad() + + loss.backward() + + optimizer.step() + + with no_grad(): + theta.copy_(trainable.clamp(theta)) + + 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(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__": diff --git a/tyff/compute/configs.py b/tyff/compute/configs.py index 24c7d07..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=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