Run fit with Adams optimizer, across SLURM distributed compute, co-optimised with energies and forces #148

import json
import pathlib

import datasets
import openff.toolkit
import torch
from dimsim.compute.configs import slurm_config
from dimsim.compute.workflow import SimulationWorkflow
from parsl import Config, python_app
from torch.autograd.function import FunctionCtx

import fitlib
import fitlib.mm
import fitlib.train
from fitlib.mm._ops import _pack_force_field, _unpack_force_field
from fitlib.targets import energy, thermo

N_EPOCHS = 2
N_REPLICATES = 3
N_MOLECULES = 256
PROPERTY_SCALES = {"density": 1 / 0.02, "hmix": 1 / 0.05, "hvap": 1 / 0.5}


# Proposed fitlib.targets.thermo.simulate body.
def simulate(
    data: datasets.Dataset,
    force_field: fitlib.TensorForceField,
    output_dir: pathlib.Path,
    compute: SimulationWorkflow,
    n_replicates: int,
    force_field_template: openff.toolkit.ForceField,
    n_molecules: int,
) -> tuple[
    dict[str, pathlib.Path], list[list[tuple[str, ...]]], pathlib.Path, pathlib.Path
]:
    # see issue about exporting to OpenFF FF
    offxml = write_offxml(
        force_field_template,
        force_field,
        output_dir / "force-field.offxml",
    )
    targets = ...
    results = compute.submit_target_batch(
        targets,
        force_field=str(offxml),
        n_molecules=n_molecules,
        n_replicates=n_replicates,
    )
    for result in results:
        if "error" in result:
            raise RuntimeError(f"{result['job_id']}: {result['error']}")

    return (
        {
            result["job_id"]: pathlib.Path(
                result["result"]["simulation_files"]["msgpack_trajectory"].filepath
            )
            for result in results
        },
        [
            compute._target_compute_mapping[
                (target["id"], str(offxml), n_molecules, n_replicates)
            ]
            for target in targets
        ],
        pathlib.Path(compute.base_dir),
        offxml,
    )


# Proposed fitlib.mm worker.
@python_app
def ensemble_average_jacobian(
    system: fitlib.TensorSystem,
    frames_path: pathlib.Path,
    temperature: float,
    pressure: float | None,
    force_field: fitlib.TensorForceField,
) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
    import openmm.unit
    import torch

    import fitlib.mm
    from fitlib.mm._ops import _pack_force_field, _unpack_force_field

    # Use existing fitlib packing order, including attributes and optional v-sites.
    tensors, parameter_lookup, attribute_lookup, has_v_sites = _pack_force_field(
        force_field
    )
    parameters = torch.cat(
        [t.detach().reshape(-1) for t in tensors if t is not None]
    ).requires_grad_(True)
    pieces = iter(parameters.split([t.numel() for t in tensors if t is not None]))
    tensors = tuple(
        None if t is None else next(pieces).reshape(t.shape) for t in tensors
    )
    worker_ff = _unpack_force_field(
        tensors,
        parameter_lookup,
        attribute_lookup,
        has_v_sites,
        force_field,
    )
    means, _ = fitlib.mm.compute_ensemble_averages(
        system,
        worker_ff,
        frames_path,
        temperature * openmm.unit.kelvin,
        None if pressure is None else pressure * openmm.unit.atmosphere,
    )
    # Jacobian rows follow sorted observable names; columns follow packed parameters.
    jacobian = torch.stack(
        [
            torch.autograd.grad(means[name], parameters, retain_graph=True)[0]
            for name in sorted(means)
        ]
    )
    return {name: value.detach() for name, value in means.items()}, jacobian.detach()


# Proposed fitlib.mm.RemoteEnsembleAverage; executed on the driver.
class RemoteEnsembleAverage(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx: FunctionCtx,
        values: torch.Tensor,
        jacobian: torch.Tensor,
        parameters: torch.Tensor,
    ) -> torch.Tensor:
        ctx.save_for_backward(jacobian)
        return values.clone()

    @staticmethod
    def backward(
        ctx: FunctionCtx,
        grad_output: torch.Tensor,
    ) -> tuple[None, None, torch.Tensor]:
        (jacobian,) = ctx.saved_tensors
        return None, None, grad_output @ jacobian


# Proposed fitlib.targets.thermo.predict body -- just passing msgpack trajectories in
def predict(
    data: datasets.Dataset,
    force_field: fitlib.TensorForceField,
    topologies: dict[str, fitlib.TensorTopology],
    simulations: tuple[
        dict[str, pathlib.Path], list[list[tuple[str, ...]]], pathlib.Path, pathlib.Path
    ],
    per_type_scales: dict[str, float],
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    ... # returns predicted values, predicted stds, reference values, reference errors

def main(root: pathlib.Path, parsl_config: Config) -> None:
    thermo_data = thermo.create_from_evaluator(root / "properties.json")
    energy_data = datasets.Dataset.load_from_disk(str(root / "qm"))
    starting_ff = openff.toolkit.ForceField("openff-2.3.0.offxml")
    topologies, tensor_ff = fitlib.topologies_for(
        {*thermo.extract_smiles(thermo_data), *energy.extract_smiles(energy_data)},
        starting_ff,
    )
    trainable = fitlib.train.Trainable(
        tensor_ff,
        parameters={
            "vdW": fitlib.train.ParameterConfig(
                cols=["epsilon", "sigma"],
                scales={"epsilon": 10.0, "sigma": 1.0},
                limits={"epsilon": (0.0001, None), "sigma": (0.5, None)},
            )
        },
        attributes={},
    )
    theta = trainable.to_values().requires_grad_(True)
    optimizer = torch.optim.Adam([theta], lr=0.01)

    with SimulationWorkflow(root / "simulations", parsl_config) as workflow:
        for epoch in range(N_EPOCHS):
            current_ff = trainable.to_force_field(theta)

            # 1. simulate with dimsim; return its results and replicate mappings.
            simulations = simulate(
                thermo_data,
                current_ff,
                output_dir=root / "epochs" / f"epoch-{epoch:03d}",
                compute=workflow,
                n_replicates=N_REPLICATES,
                force_field_template=starting_ff,
                n_molecules=N_MOLECULES,
            )

            # 2. evalute saved frames, return Jacobians, apply RemoteEnsembleAverage.
            y_ref, y_ref_std, y_pred, y_pred_std = predict(
                thermo_data,
                current_ff,
                topologies,
                simulations,
                PROPERTY_SCALES,
            )
            # This example evaluates at the sampling parameters. Changed parameters
            # require explicit reweighting or new simulation.
            energy_ref, energy_pred, forces_ref, forces_pred = energy.predict(
                energy_data,
                current_ff,
                topologies,
                normalize=False,
            )

            # 3. fit -- both targets reach the same optimizer theta.
            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 torch.no_grad():
                theta.copy_(trainable.clamp(theta))