joint-demo/simulate.py at main · omsf/joint-demo · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
import mdtraj
import nglview
import openmm
import openmm.app
import openmm.unit
from openff.interchange import Interchange
from openff.toolkit import ForceField, Molecule
from rdkit import Chem


def simulate_and_visualize(
    rdmol: Chem.Mol,
    force_field: ForceField,
) -> nglview.NGLWidget:

    interchange = force_field.create_interchange(
        Molecule.from_rdkit(rdmol).to_topology(),
    )

    integrator = openmm.LangevinMiddleIntegrator(
        300 * openmm.unit.kelvin,
        1 / openmm.unit.picosecond,
        0.002 * openmm.unit.picoseconds,
    )

    simulation = interchange.to_openmm_simulation(integrator)

    # OpenMM setup boilerplate
    simulation.minimizeEnergy()
    dcd_reporter = openmm.app.DCDReporter(file="trajectory.dcd", reportInterval=20)
    simulation.reporters.append(dcd_reporter)
    simulation.step(1000)

    # Visualize the trajectory
    trajectory: mdtraj.Trajectory = mdtraj.load(
        "trajectory.dcd",
        top=mdtraj.Topology.from_openmm(interchange.topology.to_openmm()),
    )
    view = nglview.show_mdtraj(trajectory)
    view.add_representation("line", selection="protein")
    view.add_line(selection="water")

    return view

def run_openmm_half_minute(interchange: Interchange, trj_file):
    integrator = openmm.LangevinMiddleIntegrator(
        300 * openmm.unit.kelvin,
        1 / openmm.unit.picosecond,
        0.002 * openmm.unit.picoseconds,
    )

    simulation = interchange.to_openmm_simulation(integrator)

    simulation.minimizeEnergy(tolerance=100)

    dcd_reporter = openmm.app.DCDReporter(
        file=trj_file,
        reportInterval=100,
    )
    simulation.reporters.append(dcd_reporter)

    # run for (literally) half a minute
    simulation.runForClockTime(0.5 * openmm.unit.minute)