Systematic energy shift between simulation and rerun

[jacopoventurin]

Description

When recomputing energies from saved simulation coordinates using the specialized model checkpoint, a systematic energy shift is observed in some energy components. Interestingly, with SchNet models, shifts in different components can partially compensate each other, making the total energy shift smaller than individual component shifts.
The problem becomes more evident with RANGESchNet models where the shifts do not compensate.

To reproduce

1. Run simulation saving energy components, total energy and coordinates
2. Load specialized model checkpoint and recompute energies on saved coordinates
3. Compare energy components between simulation and rerun

Observed behavior

Mean energy distribution shift $\Delta E = E_{simulation}-E_{rerun}$

Component	Shift (kcal/mol)
Model (SchNet)	~-0.1
Model (SchNet/RANGESchNet)	~2.5
non_bonded	~0.15
bonds	~0
angles	~0
dihedrals	~0
Forces	~0

Notably there is no energy shift between different rerun $E_{rerun1}-E_{rerun2} = 0$

Key observations:

• Shifts in model energy and non_bonded partially compensate each other in schnet
• Forces match perfectly (centered at 0)
• Rerun vs rerun gives 0 shift (model is deterministic)
• Shift is present from frame 0 and does not accumulate over time
• Dummy model returning sum of positions shows no shift

Environment

• mlcg version: 0.1.3
• Device: cuda
• Number of parallel simulations: 5
• Models: Transferable SchNet and RANGESchNet

Images reporting bug

Code for rerun

Adapted from @kbno.

import torch 
import numpy as np
import pandas as pd
from tqdm import tqdm
import os
import mdtraj as md
from copy import deepcopy
from glob import glob
import argparse
from torch_geometric.data.collate import collate
from mlcg.datasets.utils import chunker
import sys 
from pathlib import Path

def parse_cli():
    dir_path = os.path.dirname(os.path.realpath(__file__))
    parser = argparse.ArgumentParser(
        description="""
    Script for computing energies of bound and unbound frames on mutations for ddG computation.
    """
    )
    parser.add_argument(
        "-n",
        "--sim_name",
        type=str,
        help=(
            "name of the simulation"
        ),
    )
    parser.add_argument(
        "-d",
        "--sim_dir",
        type=str,
        help=(
            "directory of the simulation"
        ),
    )
    parser.add_argument(
        "-m",
        "--model_name",
        type=str,
        default="SchNet",
        help=(
            "model name: `SchNet` or `RANGESchNet"
        ),
    )
    return parser.parse_args()

if __name__ == "__main__":
    args = parse_cli()

    device = "cuda"

    model_path = os.path.join(args.sim_dir, f"{args.sim_name}_specialized_model_and_config.pt")
    model, config = torch.load(model_path, weights_only=False)

    data, _, _ = collate(
                            config[0].__class__,
                            data_list=config,
                            increment=True,
                            add_batch=True,
                        )

    model = model.eval().to(device=torch.device(device), dtype=torch.float32)
    for param in model.parameters():
        param.requires_grad_(False)
    data = data.to(torch.device(device))

    coords_paths = sorted(glob(os.path.join(args.sim_dir, f"{args.sim_name}_coords_*.npy")))
    for path in tqdm(coords_paths, desc=f"Computing energies for {args.sim_name}"):
        coords = np.load(path)

        # Initialize energy arrays
        forces = np.zeros_like(coords)
        energies = np.zeros((coords.shape[0], coords.shape[1]), dtype=coords.dtype)
        model_energies = np.zeros((coords.shape[0], coords.shape[1]))
        bonds_energies = np.zeros((coords.shape[0], coords.shape[1]))
        angles_energies = np.zeros((coords.shape[0], coords.shape[1]))
        dihedrals_energies = np.zeros((coords.shape[0], coords.shape[1]))
        repulsion_energies = np.zeros((coords.shape[0], coords.shape[1]))

        # Compute energies
        for fr_idx in range(coords.shape[1]):
            fr = coords[:, fr_idx]
            assert data.pos.shape[0] == fr.shape[0]*fr.shape[1]
            data.pos = torch.as_tensor(np.concatenate(fr), device=torch.device(device))
            data.out = {}
            data = model(data)
            forces[:, fr_idx, :, :] = data.out["forces"].detach().cpu().numpy().reshape(coords.shape[0], coords.shape[2], 3)
            energies[:, fr_idx] = data.out["energy"].detach().cpu().numpy()
            model_energies[:, fr_idx] = data.out[args.model_name]["energy"].detach().cpu().numpy()
            bonds_energies[:, fr_idx] = data.out["bonds"]["energy"].detach().cpu().numpy()
            angles_energies[:, fr_idx] = data.out["angles"]["energy"].detach().cpu().numpy()
            dihedrals_energies[:, fr_idx] = data.out["dihedrals"]["energy"].detach().cpu().numpy()
            repulsion_energies[:, fr_idx] = data.out["non_bonded"]["energy"].detach().cpu().numpy()
        
        # Save energies
        p = Path(path)
        name_forces = p.stem.replace("_coords_", "_forces_rerun_")
        name = p.stem.replace("_coords_", "_potential_rerun_")
        name_components = p.stem.replace("_coords_", f"_energy_components_rerun_")

        np.save(os.path.join(p.parent, name_forces), forces)
        np.save(os.path.join(p.parent, name), energies)
        np.savez(
            os.path.join(p.parent, name_components),
            model=model_energies,
            bonds=bonds_energies,
            angles=angles_energies,
            dihedrals=dihedrals_energies,
            non_bonded=repulsion_energies
        )

[sayeg84]

TL; DR: there are some simulation numerical setups that can affect a naive recalculation

Culprits:

After a long morning of debugging, @jacopoventurin and I found the two culprits responsible for this:

1. The matmul precision set to high in line 17 of the mlcg-nvt_langevin.py simulation script
2. The compilation of the model.

Matmul precision

The matmul precision is set by default as highest by torch, whereas we set it to high in the simulation as there is some speedup in ineference when doing this. This matmul change affects the outputs of the network as the priors don't really do a lot of matrix multiplication. The effect is more widely seen on range because its a more complicated architecture.

Compilation

The compilation of the model changes the spread in the non_bonded energies. We are not sure what is the reason for this but I think it could be related to the fact that torch_scatter can be prone to numerical issues. If we don't compile the network the problem disappears

Reruns taking into care matmul and compilation:

If one is consistent with the way the energies are evaluated between the rerun script and the values saved in simulation, this stops being a problem:

• Simulation and rerun with compiled network and highest matmul shows only symmetric spread in non-bonded energies, consistent with numerical errors.

• Simulation and rerun with non compiled network and highest matmul shows the same.

• Simulation and rerun with compiled network and matmul high shows symmetric spread in both schnet and non-bonded energies.

• Simulation and rerun with not compiled network and matmul high shows symmetric spread in both schnet and non-bonded energies.

Conclusions

• You need to use the same matmul percision and compilation status for doing energy recalculations.

• We analyzed this problem in two different hardware's: an RTX 3090 and RTX 6000 and the results were consistent in both.

• Note that for the forces there is also a more significant error between the compiled and not compiled version.

• if you did energy recalculations without taking recompilation or matmul precision this might have affected slightly your values

Things to do

We can improve the documentation to make this very clear. Perhaps it would also be a good idea to make the matmul precision a parameter of the simulation class

Acknowledgements

Thanks also to @EdoardoRolando who was the one that suggested that the matmul precision could be a problem.