Parameter Gradients¶

This example will show how the gradient of the potential energy with respect to force field parameters may be computed.

We start be loading and parameterizing the molecule of interest.

In [1]:

import openff.interchange
import openff.toolkit
import openff.units
import torch

import smee.converters

molecule = openff.toolkit.Molecule.from_smiles("CC(=O)NC1=CC=C(C=C1)O")
molecule.generate_conformers(n_conformers=1)

conformer = torch.tensor(molecule.conformers[0].m_as(openff.units.unit.angstrom))

interchange = openff.interchange.Interchange.from_smirnoff(
    openff.toolkit.ForceField("openff_unconstrained-2.0.0.offxml"),
    molecule.to_topology(),
)
tensor_ff, [tensor_topology] = smee.converters.convert_interchange(interchange)

import openff.interchange import openff.toolkit import openff.units import torch import smee.converters molecule = openff.toolkit.Molecule.from_smiles("CC(=O)NC1=CC=C(C=C1)O") molecule.generate_conformers(n_conformers=1) conformer = torch.tensor(molecule.conformers[0].m_as(openff.units.unit.angstrom)) interchange = openff.interchange.Interchange.from_smirnoff( openff.toolkit.ForceField("openff_unconstrained-2.0.0.offxml"), molecule.to_topology(), ) tensor_ff, [tensor_topology] = smee.converters.convert_interchange(interchange)

We can access the parameters for each SMIRNOFF parameter 'handler' (e.g. vdW, bond, angle, etc.) using the potentials_by_type (or the potentials) attribute of the TensorForceField object.

In [2]:

vdw_potential = tensor_ff.potentials_by_type["vdW"]
vdw_potential.parameters.requires_grad = True

vdw_potential = tensor_ff.potentials_by_type["vdW"] vdw_potential.parameters.requires_grad = True

The gradient of the potential energy with respect to the parameters can then be computed by backpropagating through the energy computation.

In [3]:

import smee

energy = smee.compute_energy(tensor_topology, tensor_ff, conformer)
energy.backward()

for parameter_key, gradient in zip(
    vdw_potential.parameter_keys, vdw_potential.parameters.grad.numpy(), strict=True
):
    parameter_cols = vdw_potential.parameter_cols

    parameter_grads = ", ".join(
        f"dU/d{parameter_col} = {parameter_grad: 8.3f}"
        for parameter_col, parameter_grad in zip(parameter_cols, gradient, strict=True)
    )
    print(f"{parameter_key.id.ljust(15)} - {parameter_grads}")

import smee energy = smee.compute_energy(tensor_topology, tensor_ff, conformer) energy.backward() for parameter_key, gradient in zip( vdw_potential.parameter_keys, vdw_potential.parameters.grad.numpy(), strict=True ): parameter_cols = vdw_potential.parameter_cols parameter_grads = ", ".join( f"dU/d{parameter_col} = {parameter_grad: 8.3f}" for parameter_col, parameter_grad in zip(parameter_cols, gradient, strict=True) ) print(f"{parameter_key.id.ljust(15)} - {parameter_grads}")

[#6X4:1]        - dU/depsilon =   -1.033, dU/dsigma =   -0.139
[#6:1]          - dU/depsilon =   87.490, dU/dsigma =   34.983
[#8:1]          - dU/depsilon =   15.846, dU/dsigma =   17.232
[#7:1]          - dU/depsilon =    0.148, dU/dsigma =    1.187
[#8X2H1+0:1]    - dU/depsilon =   -0.305, dU/dsigma =    0.558
[#1:1]-[#6X4]   - dU/depsilon =    7.630, dU/dsigma =    1.404
[#1:1]-[#7]     - dU/depsilon =   -2.894, dU/dsigma =   -0.074
[#1:1]-[#6X3]   - dU/depsilon =  137.134, dU/dsigma =   12.129
[#1:1]-[#8]     - dU/depsilon =  -22.417, dU/dsigma =   -0.001