Cage optimisation workflow using OpenMM/OpenFF

Here we implement a cage optimisation workflow similar to that found here but using the open-source OpenMM and OpenFF infrastructure.

This is shown in an example script that we run through below.

Warning

OpenMM/OpenFF workflows are more efficient when using GPUs, which is possible with stko if you have a GPU and if you install the cuda option, like so: pip install 'stko[cuda]'.

First we build a cage, the classic CC3 porous organic cage. But note, we are not handling the detailed stereochemistry of this system here.

from pathlib import Path

import openmm
import stk
from openff.toolkit import ForceField

import stko

# Construct a cage.
bb1 = stk.BuildingBlock(
    smiles="C1CCC(C(C1)N)N",
    functional_groups=[stk.PrimaryAminoFactory()],
)
bb2 = stk.BuildingBlock(
    smiles="C1=C(C=C(C=C1C=O)C=O)C=O",
    functional_groups=[stk.AldehydeFactory()],
)
cage = stk.ConstructedMolecule(
    topology_graph=stk.cage.FourPlusSix((bb1, bb2)),
)

First we define some settings and a function for generating a new integrator from these settings.

def integrator(
    *,
    temperature: float,
    friction: float,
    time_step: float,
) -> openmm.LangevinIntegrator:
    """Define integrator."""
    integrator = openmm.LangevinIntegrator(temperature, friction, time_step)
    integrator.setRandomNumberSeed(127)
    return integrator

# Settings.
force_field = ForceField("openff_unconstrained-2.1.0.offxml")
partial_charges = "mmff94"
temperature = 700 * openmm.unit.kelvin
friction = 10 / openmm.unit.picoseconds
time_step = 1 * openmm.unit.femtoseconds

We can then run an stk.OptimizerSequence built from OpenMM classes to get the structure below in a few minutes!

# Define sequence.
optimisation_sequence = stko.OptimizerSequence(
    # Restricted true to optimised the constructed bonds.
    stko.OpenMMForceField(
        force_field=force_field,
        restricted=True,
        partial_charges_method=partial_charges,
    ),
    # Unrestricted optimisation.
    stko.OpenMMForceField(
        # Load the openff-2.1.0 force field appropriate for
        # vacuum calculations (without constraints)
        force_field=force_field,
        restricted=False,
        partial_charges_method=partial_charges,
    ),
    # Molecular dynamics, short for equilibration.
    stko.OpenMMMD(
        force_field=force_field,
        output_directory=output_directory / "md_optimisation",
        integrator=integrator(
            temperature=temperature,
            friction=friction,
            time_step=time_step,
        ),
        random_seed=275,
        partial_charges_method=partial_charges,
        # Frequency here is not related to the num confs tested.
        reporting_freq=100,
        trajectory_freq=100,
        # 10 ps
        num_steps=10_000,
        num_conformers=10,
        platform="CUDA",
        conformer_optimiser=stko.OpenMMForceField(
            force_field=force_field,
            restricted=False,
            partial_charges_method=partial_charges,
        ),
    ),
    # Long MD, for collecting lowest energy conformers.
    stko.OpenMMMD(
        force_field=force_field,
        output_directory=output_directory / "md_optimisation",
        integrator=integrator(
            temperature=temperature,
            friction=friction,
            time_step=time_step,
        ),
        random_seed=275,
        partial_charges_method=partial_charges,
        # Frequency here is not related to the num confs tested.
        reporting_freq=100,
        trajectory_freq=100,
        # 0.2 ns
        num_steps=200_000,
        # 1 every 4 ps
        num_conformers=50,
        platform="CUDA",
        conformer_optimiser=stko.OpenMMForceField(
            force_field=force_field,
            restricted=False,
            partial_charges_method=partial_charges,
        ),
    ),
)

optimised_cage = optimisation_sequence.optimize(cage)