rmf/simulation.py

This example shows writing a brownian dynamics simulation to a rmf file. It includes a bounding box, restraints and a hierarchy of particles.

## \example rmf/simulation.py
# This example shows writing a brownian dynamics simulation to a rmf file.
# It includes a bounding box, restraints and a hierarchy of particles.

from __future__ import print_function, division
import IMP.atom
import IMP.rmf
import RMF
import IMP.container
import IMP.display
import sys

IMP.setup_from_argv(sys.argv, "simulation")

k = 10.0
np = 10
frames = 10

m = IMP.Model()
bb = IMP.algebra.BoundingBox3D(IMP.algebra.Vector3D(0, 0, 0),
                               IMP.algebra.Vector3D(10, 10, 10))

ps = [IMP.core.XYZR.setup_particle(IMP.Particle(m))
      for i in range(0, np)]
h = IMP.atom.Hierarchy.setup_particle(IMP.Particle(m, "root"))
for p in ps:
    p.set_radius(1)
    p.set_coordinates_are_optimized(True)
    IMP.atom.Mass.setup_particle(p, 30)
    IMP.atom.Diffusion.setup_particle(p)
    p.set_coordinates(IMP.algebra.get_random_vector_in(bb))
rs = []
link = IMP.core.HarmonicUpperBoundSphereDistancePairScore(0, k)
pl = []
for i in range(0, len(ps) // 2):
    pp = (ps[2 * i], ps[2 * i + 1])
    pl.append(pp)
    # create OK staring position
    pp[1].set_coordinates(IMP.algebra.get_random_vector_on(pp[0].get_sphere()))
    IMP.display.Colored.setup_particle(pp[0], IMP.display.get_display_color(i))
    IMP.display.Colored.setup_particle(pp[1], IMP.display.get_display_color(i))
    hr = IMP.atom.Hierarchy.setup_particle(
        IMP.Particle(m, "molecule " + str(i)))
    hr.add_child(IMP.atom.Hierarchy.setup_particle(pp[0]))
    hr.add_child(IMP.atom.Hierarchy.setup_particle(pp[1]))
    h.add_child(hr)
r = IMP.container.PairsRestraint(link, pl, "Bonds")
rs.append(r)

ev = IMP.core.ExcludedVolumeRestraint(ps, k, 1, "EV")
rs.append(ev)

bbss = IMP.core.BoundingBox3DSingletonScore(
    IMP.core.HarmonicUpperBound(0, k), bb)
bbr = IMP.container.SingletonsRestraint(bbss, ps)
rs.append(bbr)

sf = IMP.core.RestraintsScoringFunction(rs, "SF")

bd = IMP.atom.BrownianDynamics(m)
bd.set_log_level(IMP.SILENT)
bd.set_scoring_function(sf)

bd.set_maximum_time_step(10)

rmf = RMF.create_rmf_file("brownian.rmf")
rmf.set_description("Brownian dynamics trajectory with 10fs timestep.\n")
IMP.rmf.add_hierarchy(rmf, h)
IMP.rmf.add_restraints(rmf, rs)
IMP.rmf.add_geometry(rmf, IMP.display.BoundingBoxGeometry(bb))

os = IMP.rmf.SaveOptimizerState(m, rmf)
os.update_always("initial conformation")
os.set_log_level(IMP.SILENT)
os.set_simulator(bd)
bd.add_optimizer_state(os)

bd.optimize(frames)

print("file is", rmf.get_name())