display/displaying_ensembles.py

The script shows a couple experiments with trying to visualize an ensemble of structures. The ensemble is fairly tight on the assembly scale, but there is significant variation between the location and orientation of the individual proteins (which were modeled as rigid bodies). To save space, the models have had their sidechain atoms removed.

## \example display/displaying_ensembles.py
# The script shows a couple experiments with trying to visualize an
# ensemble of structures. The ensemble is fairly tight on the assembly
# scale, but there is significant variation between the location and
# orientation of the individual proteins (which were modeled as rigid
# bodies). To save space, the models have had their sidechain atoms
# removed.

import IMP.display
import IMP.atom
import sys

IMP.setup_from_argv(
    sys.argv, "Experiments with trying to visualize an ensemble of structures")

Segment = IMP.algebra.Segment3D
Cylinder = IMP.algebra.Cylinder3D

# turn off internal checks to speed things up
IMP.set_check_level(IMP.USAGE)


def read(m, beyond_file):
    print("reading")
    hs = []
    for i in range(0, beyond_file):
        # create a simplified version for each chain to speed up computations
        name = IMP.display.get_example_path(
            "ensemble/aligned-" + str(i) + ".pdb")
        h = IMP.atom.read_pdb(name, m, IMP.atom.CAlphaPDBSelector())
        hr = IMP.atom.Hierarchy.setup_particle(IMP.Particle(m))
        hs.append(hr)
        for c in IMP.atom.get_by_type(h, IMP.atom.CHAIN_TYPE):
            simp = IMP.atom.create_simplified_along_backbone(
                IMP.atom.Chain(c), 4)
            hr.add_child(simp)
        IMP.atom.destroy(h)
        print(" ", i)
    return hs


def add_markers(h, c, w):
    """Add markers to a the passed conformation. The marker locations are chosen
    pretty thoughtlessly and don't really illustrate the technique well."""
    def add_marker(s, name):
        g = IMP.core.XYZRGeometry(s.get_selected_particles()[0])
        g.set_name(name)
        g.set_color(c)
        w.add_geometry(g)
    s = IMP.atom.Selection(h, chain='B', residue_index=317)
    add_marker(s, "m0")
    s = IMP.atom.Selection(h, chain='G', residue_index=212)
    add_marker(s, "m1")
    s = IMP.atom.Selection(h, chain='I', residue_index=237)
    add_marker(s, "m2")
    s = IMP.atom.Selection(h, chain='F', residue_index=101)
    add_marker(s, "m3")


def get_nice_name(h):
    nm = h.get_name()
    return nm[nm.find('-') + 1:nm.rfind('.')]


def add_axis(h, c, w, chain_colors):
    """Add a coordinate axis to show the relative orientation of the protein"""
    for hc in IMP.atom.get_by_type(h, IMP.atom.CHAIN_TYPE):
        rb = IMP.core.RigidMember(hc).get_rigid_body()
        g = IMP.display.ReferenceFrameGeometry(rb.get_reference_frame())
        g.set_name(get_nice_name(h) + "_orient")
        if c:
            g.set_color(c)
        else:
            g.set_color(chain_colors[IMP.atom.Chain(hc).get_id()])
        w.add_geometry(g)


def add_skeleton(h, c, r, w, chain_colors):
    """Show the connectivity skeleton of the conformation to give an idea of
    how things are laid out"""
    for hc0 in IMP.atom.get_by_type(h, IMP.atom.CHAIN_TYPE):
        for hc1 in IMP.atom.get_by_type(h, IMP.atom.CHAIN_TYPE):
            if hc1 <= hc0:
                continue
            d = ps.evaluate_index(h.get_model(),
                                  (hc0.get_particle_index(),
                                   hc1.get_particle_index()), None)
            if d < 1:
                d0 = IMP.core.XYZ(hc0)
                d1 = IMP.core.XYZ(hc1)
                mp = .5 * (d0.get_coordinates() + d1.get_coordinates())
                g = IMP.display.CylinderGeometry(
                    Cylinder(Segment(d0.get_coordinates(), mp), r))
                if c:
                    g.set_color(c)
                else:
                    g.set_color(chain_colors[IMP.atom.Chain(d0).get_id()])
                g.set_name(get_nice_name(h) + "_skel")
                w.add_geometry(g)
                g = IMP.display.CylinderGeometry(
                    Cylinder(Segment(d1.get_coordinates(), mp), r))
                if c:
                    g.set_color(c)
                else:
                    g.set_color(chain_colors[IMP.atom.Chain(d1).get_id()])
                g.set_name(get_nice_name(h) + "_skel")
                w.add_geometry(g)


IMP.set_log_level(IMP.TERSE)
m = IMP.Model()

# change to 46 to display all of them
hs = read(m, 3)

# used to test of two molecules are touching one another
ps = IMP.core.KClosePairsPairScore(
    IMP.core.SphereDistancePairScore(IMP.core.HarmonicUpperBound(10, 1)),
    IMP.core.LeavesRefiner(IMP.atom.Hierarchy.get_traits()))
ps.set_log_level(IMP.SILENT)


print("creating rigid bodies")
base_chains = {}
for hc in IMP.atom.get_by_type(hs[0], IMP.atom.CHAIN_TYPE):
    c = IMP.atom.Chain(hc)
    base_chains[c.get_id()] = c

for i, h in enumerate(hs):
    for hc in IMP.atom.get_by_type(h, IMP.atom.CHAIN_TYPE):
        c = IMP.atom.Chain(hc)
        if h == hs[0]:
            crb = IMP.atom.create_rigid_body(hc)
        else:
            # Make sure the rigid bodies have equivalent defining reference
            # frames. If we just used IMP.atom.create_rigid_body, globular
            # proteins are likely to have different axes computed when
            # starting in different orientations
            crb = IMP.atom.create_compatible_rigid_body(
                hc, base_chains[c.get_id()])
    print(" ", i)

chains = IMP.atom.get_by_type(hs[0], IMP.atom.CHAIN_TYPE)
chains.sort(key=lambda x: IMP.core.XYZ(x).get_x() + IMP.core.XYZ(x).get_y())
chain_colors = {}
for i, c in enumerate(chains):
    id = IMP.atom.Chain(c).get_id()
    color = IMP.display.get_display_color(i)
    # IMP.display.get_jet_color(f)
    chain_colors[id] = color

w = IMP.display.PymolWriter("markers.pym")
add_markers(hs[0], IMP.display.Color(1, 1, 1), w)
hso = hs[1:]


# sort them spatially so the colors are nicely arranged and allow one to
# visually connect the position of one end with that of the other
hso.sort(key=lambda h: IMP.core.XYZ(IMP.atom.Selection(
    h, chain='I', residue_index=237).get_selected_particles()[0]).get_z())
print("adding markers", end=' ')
for i, h in enumerate(hso):
    c = IMP.display.get_interpolated_rgb(
        IMP.display.Color(1, 0, 0), IMP.display.Color(0, 0, 1), i / 50.)
    add_markers(h, c, w)
    print(" ", i)
w = IMP.display.PymolWriter("axis.pym")
print("adding axis", end=' ')
add_axis(hs[0], IMP.display.Color(1, 1, 1), w, chain_colors)
for i, h in enumerate(hs[1:]):
    add_axis(h, None, w, chain_colors)
    print(i, end=' ')

w = IMP.display.PymolWriter("skeletons.pym")
add_skeleton(hs[0], IMP.display.Color(1, 1, 1), 5, w, chain_colors)
print("adding skeleton", end=' ')
for i, h in enumerate(hs[1:]):
    add_skeleton(h, None, 1, w, chain_colors)
    print(" ", i)