clustering of pdb models

import IMP import IMP.base import IMP.core as core import IMP.atom as atom import IMP.em2d as em2d import IMP.em as em import IMP.algebra as alg import IMP.container import random

An Optimizer score to get the values of the statistics after a given set

of evaluations

class WriteStatisticsOptimizerScore(IMP.OptimizerState): count =0 def init(self): IMP.OptimizerState.__init__(self) count = 0 def update(self): if (self.count!=10): self.count += 1 return else: self.count=0 o=self.get_optimizer() m=o.get_model() m.show_restraint_score_statistics() m.show_all_statistics() #for i in range(0,m.get_number_of_restraints()):

r=m.get_restraint(i)

print "restraint",r.get_name(),"value",r.evaluate(False)

def do_show(self, stream): print >> stream, ps

Get model from PDB file

IMP.set_log_level(IMP.TERSE) m = IMP.Model() prot = atom.read_pdb(em2d.get_example_path("1z5s.pdb"),m,atom.ATOMPDBSelector()) atom.add_radii(prot)

get the chains

chains = atom.get_by_type(prot,atom.CHAIN_TYPE) print "there are",len(chains),"chains in 1z5s.pdb"

set the chains as rigid bodies

native_chain_centers = [] rigid_bodies= [] for c in chains: atoms=core.get_leaves(c) rbd=core.RigidBody.setup_particle(c,atoms) rigid_bodies.append(rbd) print "chain has",rbd.get_number_of_members(), \ "atoms","coordinates: ",rbd.get_coordinates() native_chain_centers.append(rbd.get_coordinates())

bb=alg.BoundingBox3D(alg.Vector3D(-25, -40,-60), alg.Vector3D( 25, 40, 60))

rotate and translate the chains

for rbd in rigid_bodies:

rotation

rotation= alg.get_random_rotation_3d() transformation1=alg.get_rotation_about_point(rbd.get_coordinates(),rotation)

translation

transformation2=alg.Transformation3D(alg.get_random_vector_in(bb))

Apply

final_transformation = alg.compose(transformation1,transformation2) core.transform(rbd,final_transformation) print "Writing transformed assembly" atom.write_pdb (prot,"1z5s-transformed.pdb")

set distance restraints measusring some distances between rigid bodies

for the solution.

d01 = alg.get_distance(native_chain_centers[0],native_chain_centers[1]) r01 = core.DistanceRestraint(core.Harmonic(d01,1),chains[0],chains[1]) r01.set_name("distance 0-1") d12 = alg.get_distance(native_chain_centers[1],native_chain_centers[2]) r12 = core.DistanceRestraint(core.Harmonic(d12,1),chains[1],chains[2]) r12.set_name("distance 1-2") d23 = alg.get_distance(native_chain_centers[2],native_chain_centers[3]) r23 = core.DistanceRestraint(core.Harmonic(d23,1),chains[2],chains[3]) r23.set_name("distance 2-3") d30 = alg.get_distance(native_chain_centers[3],native_chain_centers[0]) r30 = core.DistanceRestraint(core.Harmonic(d30,1),chains[3],chains[0]) r30.set_name("distance 3-0") print "Distances in the solution: d01 =", \ d01,"d12 =",d12,"d23 =",d23,"d30 =",d30

set distance restraints

print "adding distance restraints " for r in [r01,r12,r23,r30]: m.add_restraint(r) print "model has ",m.get_number_of_restraints(),"restraints"

set em2D restraint

srw = em2d.SpiderImageReaderWriter() selection_file=em2d.get_example_path("all-1z5s-projections.sel") images_to_read_names=[IMP.base.get_relative_path(selection_file, x) for x in em2d.read_selection_file(selection_file)] em_images =em2d.read_images(images_to_read_names,srw) print len(em_images),"images read"

em2d_restraint = em2d.Em2DRestraint() apix=1.5 # sampling rate of the available EM images

resolution at which you want to generate the projections of the model

In principle you want "perfect" projections, so use the highest resolution

resolution=1

Number of projections to use for the initial registration

(coarse registration) to estimate the registration parameters

n_projections=20 params = em2d.Em2DRestraintParameters(apix,resolution,n_projections)

This method (recommended) uses preprocessing of the images and projections

to speed-up the registration

params.coarse_registration_method = IMP.em2d.ALIGN2D_PREPROCESSING

use true if you want to save the projections from the model that best

match the Em images

params.save_match_images = False

score_function=IMP.em2d.EM2DScore() em2d_restraint = IMP.em2d.Em2DRestraint() em2d_restraint.setup(score_function, params) em2d_restraint.set_images(em_images) em2d_restraint.set_name("em2d restraint") container = IMP.container.ListSingletonContainer(core.get_leaves(prot)) em2d_restraint.set_particles(container) em2d_restraints_set=IMP.RestraintSet()

The next two lines are commented, because the optimization of the example

is expensive. To run the full example, uncomment them (It can take a few

hours).

#em2d_restraints_set.add_restraint(em2d_restraint) #em2d_restraints_set.set_weight(1000) # weight for the em2D restraint

print "adding em2d restraint " m.add_restraint(em2d_restraints_set)

Add all restraints to a model

print "model has ",m.get_number_of_restraints(),"restraints"

MONTECARLO OPTIMIZATION

s=core.MonteCarlo(m)

Add movers for the rigid bodies

movers=[] for rbd in rigid_bodies: movers.append(core.RigidBodyMover(rbd,5,2)) s.add_movers(movers) print "MonteCarlo sampler has",s.get_number_of_movers(),"movers"

Add an optimizer state to save intermediate configurations of the hierarchy

o_state=atom.WritePDBOptimizerState(chains,"intermediate-step-%1%.pdb") o_state.set_skip_steps(10) s.add_optimizer_state(o_state)

ostate2 = WriteStatisticsOptimizerScore() s.add_optimizer_state(ostate2)

Perform optimization

m.set_gather_statistics(True) # Writes a lot of information!

temperatures=[200,100,60,40,20,5]

200 steps recommended for accurate optimization; a smaller number is used

here for demonstration purposes

optimization_steps = 10 for T in temperatures: s.set_kt(T) s.optimize(optimization_steps) atom.write_pdb(prot,"solution.pdb")

Check that the optimization achieves distances close to the solution

print "*** End optimization ***" new_centers = [] for rbd in rigid_bodies: print "chain has",rbd.get_number_of_members(), \ "atoms","coordinates: ",rbd.get_coordinates() new_centers.append(rbd.get_coordinates())

d01 = alg.get_distance(new_centers[0],new_centers[1]) d12 = alg.get_distance(new_centers[1],new_centers[2]) d23 = alg.get_distance(new_centers[2],new_centers[3]) d30 = alg.get_distance(new_centers[3],new_centers[0]) print "Distances at the end of the optimization: d01 =", \ d01,"d12 =",d12,"d23 =",d23,"d30 =",d30

import IMP
import IMP.core as core
import IMP.atom as atom
import IMP.em2d as em2d
import os
import sys
import csv
"""
    Clustering of pdb models.
    This script clusters pdb models of an structure, chosen from a
    selection file.
    - It is assumed that all the pdb files belong to the same structure
    and that the order of the atoms in the pdb files is the same in all files
    - After the clustering procedure, a linkage matrix is generated.
"""
if sys.platform == 'win32':
    sys.stderr.write("this example does not currently work on Windows\n")
    sys.exit(0)
def get_columns(fn,cols=[],delimiter=" ",comment="#"):
    """ ge the columns of a file:
        cols - a list of columns to extract. E.g [0,3,5]
               If empty, all the columns are extracted
        lines starting with the comment character are ignored """
    columns=[[] for i in cols]
    # get a reader for the file
    reader=csv.reader(open(fn,"r"),delimiter=delimiter,skipinitialspace=True)
    for row in reader:
        if(row!=[] and row[0][0]!=comment): # not empty or comment row
            if(cols==[]):
                for i in range(0,len(row)):
                    columns[i].append(row[i])
            else:
                for i in range(0,len(cols)):
                    columns[i].append(row[cols[i]])
    return columns
def argmin(sequence):
    """ Argmin function: Returns the pair (min_value,min_index),
        where min_index is the index of the minimum value
    """
    min_value = sequence[0]
    min_index = 0
    for i in range(0,len(sequence)):
#        print "argmin - checking ",sequence[i]
        if(sequence[i]<min_value):
            min_value = sequence[i]
            min_index =i
#            print "argmin - selecting ",min_value,min_index
    return min_value,min_index
#***************************
fn_selection = em2d.get_example_path("all-models-1z5s.sel")
fn_em2d_scores = em2d.get_example_path("em2d_scores_for_clustering.data")
# Load models
print "Reading models ..."
model = IMP.Model()
ssel = atom.ATOMPDBSelector()
coords =[]
fn_models = em2d.read_selection_file(fn_selection)
n_models = len(fn_models)
hierarchies=[]
for fn in fn_models:
    fn_model=em2d.get_example_path(fn)
    h=atom.read_pdb(fn_model,model,ssel,True,True);
    hierarchies.append(h)
    xyz=core.XYZs(atom.get_leaves(h))
    coords.append( [x.get_coordinates() for x in xyz])
print "Computing matrix of RMSD ..."
rmsds=[[0.0 for i in range(0,n_models)] for n in range(0,n_models)]
transformations=[[[] for i in range(0,n_models)]  for j in range(0,n_models)]
# fill rmsd and transformations
for i in xrange(0,n_models):
    for j in xrange(i+1,n_models):
        if(i!=j):
            t=IMP.algebra.get_transformation_aligning_first_to_second(
                                              coords[i],
                                              coords[j])
            transformations[i][j]=t
            transformations[j][i]=t.get_inverse()
            temp = [t.get_transformed(v) for v in coords[i]]
            rmsd=IMP.atom.get_rmsd(temp,coords[j])
            rmsds[i][j]=rmsd
            rmsds[j][i]=rmsd
# cluster
print "Clustering (Complete linkage methond)..."
cluster_set = em2d.do_hierarchical_clustering_complete_linkage(rmsds)
mat2=cluster_set.get_linkage_matrix()
print "Complete Linkage Matrix"
for m in mat2:
    print m
# Read scores from the scores file
em2d_scores = get_columns(fn_em2d_scores,[1])
em2d_scores = em2d_scores[0]
# get biggest clusters below a certain rmsd
rmsd_cutoff=1.4
print "clusters below cutoff",rmsd_cutoff,"Angstroms"
clusters=cluster_set.get_clusters_below_cutoff(rmsd_cutoff)
for c in clusters:
    elems=cluster_set.get_cluster_elements(c)
    scores_elements=[]
    for cid in elems:
        scores_elements.append(em2d_scores[cid])
    print "Cluster",c,":",elems,scores_elements,
    # find model with best score
    min_value,min_index= argmin(scores_elements)
    min_elem_id=elems[min_index]
    # The representative element is the one with the minimum em2d score
    print "representative element",min_elem_id,min_value
    for i in elems:
        pdb_name="cluster-%03d-elem-%03d.pdb" % (c,i)
        if(i!=min_elem_id):
            print "Writting element",i,"aligned to ",min_elem_id,":",pdb_name
            T=core.Transform(transformations[i][min_elem_id])
            ps=atom.get_leaves(hierarchies[i])
            for p in ps:
                T.apply(p)
        else:
            print "Writting representative element",min_elem_id,":",pdb_name
        atom.write_pdb(hierarchies[i],pdb_name)

collision cross section

import IMP import IMP.base import IMP.core as core import IMP.atom as atom import IMP.em2d as em2d import IMP.em as em import IMP.algebra as alg import IMP.container import random

An Optimizer score to get the values of the statistics after a given set

of evaluations

r=m.get_restraint(i)

print "restraint",r.get_name(),"value",r.evaluate(False)

def do_show(self, stream): print >> stream, ps

Get model from PDB file

IMP.set_log_level(IMP.TERSE) m = IMP.Model() prot = atom.read_pdb(em2d.get_example_path("1z5s.pdb"),m,atom.ATOMPDBSelector()) atom.add_radii(prot)

get the chains

chains = atom.get_by_type(prot,atom.CHAIN_TYPE) print "there are",len(chains),"chains in 1z5s.pdb"

set the chains as rigid bodies

bb=alg.BoundingBox3D(alg.Vector3D(-25, -40,-60), alg.Vector3D( 25, 40, 60))

rotate and translate the chains

for rbd in rigid_bodies:

rotation

rotation= alg.get_random_rotation_3d() transformation1=alg.get_rotation_about_point(rbd.get_coordinates(),rotation)

translation

transformation2=alg.Transformation3D(alg.get_random_vector_in(bb))

Apply

final_transformation = alg.compose(transformation1,transformation2) core.transform(rbd,final_transformation) print "Writing transformed assembly" atom.write_pdb (prot,"1z5s-transformed.pdb")

set distance restraints measusring some distances between rigid bodies

for the solution.

set distance restraints

print "adding distance restraints " for r in [r01,r12,r23,r30]: m.add_restraint(r) print "model has ",m.get_number_of_restraints(),"restraints"

set em2D restraint

em2d_restraint = em2d.Em2DRestraint() apix=1.5 # sampling rate of the available EM images

resolution at which you want to generate the projections of the model

In principle you want "perfect" projections, so use the highest resolution

resolution=1

Number of projections to use for the initial registration

(coarse registration) to estimate the registration parameters

n_projections=20 params = em2d.Em2DRestraintParameters(apix,resolution,n_projections)

This method (recommended) uses preprocessing of the images and projections

to speed-up the registration

params.coarse_registration_method = IMP.em2d.ALIGN2D_PREPROCESSING

use true if you want to save the projections from the model that best

match the Em images

params.save_match_images = False

The next two lines are commented, because the optimization of the example

is expensive. To run the full example, uncomment them (It can take a few

hours).

#em2d_restraints_set.add_restraint(em2d_restraint) #em2d_restraints_set.set_weight(1000) # weight for the em2D restraint

print "adding em2d restraint " m.add_restraint(em2d_restraints_set)

Add all restraints to a model

print "model has ",m.get_number_of_restraints(),"restraints"

MONTECARLO OPTIMIZATION

s=core.MonteCarlo(m)

Add movers for the rigid bodies

movers=[] for rbd in rigid_bodies: movers.append(core.RigidBodyMover(rbd,5,2)) s.add_movers(movers) print "MonteCarlo sampler has",s.get_number_of_movers(),"movers"

Add an optimizer state to save intermediate configurations of the hierarchy

o_state=atom.WritePDBOptimizerState(chains,"intermediate-step-%1%.pdb") o_state.set_skip_steps(10) s.add_optimizer_state(o_state)

ostate2 = WriteStatisticsOptimizerScore() s.add_optimizer_state(ostate2)

Perform optimization

m.set_gather_statistics(True) # Writes a lot of information!

temperatures=[200,100,60,40,20,5]

200 steps recommended for accurate optimization; a smaller number is used

here for demonstration purposes

optimization_steps = 10 for T in temperatures: s.set_kt(T) s.optimize(optimization_steps) atom.write_pdb(prot,"solution.pdb")

Check that the optimization achieves distances close to the solution

import IMP
import IMP.em2d as em2d
import IMP.atom as atom
"""
Example of how to compute the collision cross section of a molecule
"""
IMP.set_log_level(IMP.TERSE)
m = IMP.Model()
fn = em2d.get_example_path("1z5s.pdb")
prot =  atom.read_pdb(fn, m ,atom.ATOMPDBSelector())
atom.add_radii(prot)
projections = 20
resolution = 1.0
pixel_size = 1.5
img_size = 80
ccs = em2d.CollisionCrossSection(projections, resolution, pixel_size, img_size)
ccs.set_model_particles(IMP.atom.get_leaves(prot))
print "CCS",ccs.get_ccs(),"A**2"

IMP::atom::get_leaves(Hierarchy)

em images conversion

import IMP import IMP.base import IMP.core as core import IMP.atom as atom import IMP.em2d as em2d import IMP.em as em import IMP.algebra as alg import IMP.container import random

An Optimizer score to get the values of the statistics after a given set

of evaluations

r=m.get_restraint(i)

print "restraint",r.get_name(),"value",r.evaluate(False)

def do_show(self, stream): print >> stream, ps

Get model from PDB file

IMP.set_log_level(IMP.TERSE) m = IMP.Model() prot = atom.read_pdb(em2d.get_example_path("1z5s.pdb"),m,atom.ATOMPDBSelector()) atom.add_radii(prot)

get the chains

chains = atom.get_by_type(prot,atom.CHAIN_TYPE) print "there are",len(chains),"chains in 1z5s.pdb"

set the chains as rigid bodies

bb=alg.BoundingBox3D(alg.Vector3D(-25, -40,-60), alg.Vector3D( 25, 40, 60))

rotate and translate the chains

for rbd in rigid_bodies:

rotation

rotation= alg.get_random_rotation_3d() transformation1=alg.get_rotation_about_point(rbd.get_coordinates(),rotation)

translation

transformation2=alg.Transformation3D(alg.get_random_vector_in(bb))

Apply

final_transformation = alg.compose(transformation1,transformation2) core.transform(rbd,final_transformation) print "Writing transformed assembly" atom.write_pdb (prot,"1z5s-transformed.pdb")

set distance restraints measusring some distances between rigid bodies

for the solution.

set distance restraints

print "adding distance restraints " for r in [r01,r12,r23,r30]: m.add_restraint(r) print "model has ",m.get_number_of_restraints(),"restraints"

set em2D restraint

em2d_restraint = em2d.Em2DRestraint() apix=1.5 # sampling rate of the available EM images

resolution at which you want to generate the projections of the model

In principle you want "perfect" projections, so use the highest resolution

resolution=1

Number of projections to use for the initial registration

(coarse registration) to estimate the registration parameters

n_projections=20 params = em2d.Em2DRestraintParameters(apix,resolution,n_projections)

This method (recommended) uses preprocessing of the images and projections

to speed-up the registration

params.coarse_registration_method = IMP.em2d.ALIGN2D_PREPROCESSING

use true if you want to save the projections from the model that best

match the Em images

params.save_match_images = False

The next two lines are commented, because the optimization of the example

is expensive. To run the full example, uncomment them (It can take a few

hours).

#em2d_restraints_set.add_restraint(em2d_restraint) #em2d_restraints_set.set_weight(1000) # weight for the em2D restraint

print "adding em2d restraint " m.add_restraint(em2d_restraints_set)

Add all restraints to a model

print "model has ",m.get_number_of_restraints(),"restraints"

MONTECARLO OPTIMIZATION

s=core.MonteCarlo(m)

Add movers for the rigid bodies

movers=[] for rbd in rigid_bodies: movers.append(core.RigidBodyMover(rbd,5,2)) s.add_movers(movers) print "MonteCarlo sampler has",s.get_number_of_movers(),"movers"

Add an optimizer state to save intermediate configurations of the hierarchy

o_state=atom.WritePDBOptimizerState(chains,"intermediate-step-%1%.pdb") o_state.set_skip_steps(10) s.add_optimizer_state(o_state)

ostate2 = WriteStatisticsOptimizerScore() s.add_optimizer_state(ostate2)

Perform optimization

m.set_gather_statistics(True) # Writes a lot of information!

temperatures=[200,100,60,40,20,5]

200 steps recommended for accurate optimization; a smaller number is used

here for demonstration purposes

optimization_steps = 10 for T in temperatures: s.set_kt(T) s.optimize(optimization_steps) atom.write_pdb(prot,"solution.pdb")

Check that the optimization achieves distances close to the solution

import IMP
import IMP.em2d as em2d
import os
"""
    Conversion of Electron Microscopy Images.
"""
# Read images
fn_selection=em2d.get_example_path("all-1z5s-projections.sel")
srw = em2d.SpiderImageReaderWriter()
trw = em2d.TIFFImageReaderWriter()
fn_images=em2d.read_selection_file(fn_selection)
fn_images = [em2d.get_example_path(x) for x in fn_images]
images = em2d.read_images(fn_images,srw)
# write
fn_saved=em2d.create_filenames(3,"1z5s-projection","tif")
em2d.save_images(images,fn_saved,trw)

optimize em2d with montecarlo

Example of optimizing an EM2D restraint using Monte Carlo.

import IMP
import IMP.base
import IMP.core as core
import IMP.atom as atom
import IMP.em2d as em2d
import IMP.em as em
import IMP.algebra as alg
import IMP.container
import random
# An Optimizer score to get the values of the statistics after a given set
# of evaluations
class WriteStatisticsOptimizerScore(IMP.OptimizerState):
    count =0
    def __init__(self):
        IMP.OptimizerState.__init__(self)
        count = 0
    def update(self):
        if (self.count!=10):
            self.count += 1
            return
        else:
            self.count=0
        o=self.get_optimizer()
        m=o.get_model()
        m.show_restraint_score_statistics()
        m.show_all_statistics()
        #for i in range(0,m.get_number_of_restraints()):
        #    r=m.get_restraint(i)
        #    print "restraint",r.get_name(),"value",r.evaluate(False)
    def do_show(self, stream):
        print >> stream, ps
# Get model from PDB file
IMP.set_log_level(IMP.TERSE)
m = IMP.Model()
prot =  atom.read_pdb(em2d.get_example_path("1z5s.pdb"),m,atom.ATOMPDBSelector())
atom.add_radii(prot)
# get the chains
chains = atom.get_by_type(prot,atom.CHAIN_TYPE)
print "there are",len(chains),"chains in 1z5s.pdb"
# set the chains as rigid bodies
native_chain_centers = []
rigid_bodies= []
for c in chains:
    atoms=core.get_leaves(c)
    rbd=core.RigidBody.setup_particle(c,atoms)
    rigid_bodies.append(rbd)
    print "chain has",rbd.get_number_of_members(), \
                          "atoms","coordinates: ",rbd.get_coordinates()
    native_chain_centers.append(rbd.get_coordinates())
bb=alg.BoundingBox3D(alg.Vector3D(-25, -40,-60),
                         alg.Vector3D( 25,  40, 60))
# rotate and translate the chains
for rbd in rigid_bodies:
    # rotation
    rotation= alg.get_random_rotation_3d()
    transformation1=alg.get_rotation_about_point(rbd.get_coordinates(),rotation)
    # translation
    transformation2=alg.Transformation3D(alg.get_random_vector_in(bb))
    # Apply
    final_transformation = alg.compose(transformation1,transformation2)
    core.transform(rbd,final_transformation)
print "Writing transformed assembly"
atom.write_pdb (prot,"1z5s-transformed.pdb")
# set distance restraints measusring some distances between rigid bodies
# for the solution.
d01 = alg.get_distance(native_chain_centers[0],native_chain_centers[1])
r01 = core.DistanceRestraint(core.Harmonic(d01,1),chains[0],chains[1])
r01.set_name("distance 0-1")
d12 = alg.get_distance(native_chain_centers[1],native_chain_centers[2])
r12 = core.DistanceRestraint(core.Harmonic(d12,1),chains[1],chains[2])
r12.set_name("distance 1-2")
d23 = alg.get_distance(native_chain_centers[2],native_chain_centers[3])
r23 = core.DistanceRestraint(core.Harmonic(d23,1),chains[2],chains[3])
r23.set_name("distance 2-3")
d30 = alg.get_distance(native_chain_centers[3],native_chain_centers[0])
r30 = core.DistanceRestraint(core.Harmonic(d30,1),chains[3],chains[0])
r30.set_name("distance 3-0")
print "Distances in the solution: d01 =", \
    d01,"d12 =",d12,"d23 =",d23,"d30 =",d30
# set distance restraints
print "adding distance restraints "
for r in [r01,r12,r23,r30]:
    m.add_restraint(r)
print "model has ",m.get_number_of_restraints(),"restraints"
# set em2D restraint
srw = em2d.SpiderImageReaderWriter()
selection_file=em2d.get_example_path("all-1z5s-projections.sel")
images_to_read_names=[IMP.base.get_relative_path(selection_file, x) for x in
        em2d.read_selection_file(selection_file)]
em_images =em2d.read_images(images_to_read_names,srw)
print len(em_images),"images read"
em2d_restraint = em2d.Em2DRestraint()
apix=1.5 # sampling rate of the available EM images
# resolution at which you want to generate the projections of the model
# In principle you want "perfect" projections, so use the highest resolution
resolution=1
# Number of projections to use for the initial registration
#  (coarse registration) to estimate the registration parameters
n_projections=20
params = em2d.Em2DRestraintParameters(apix,resolution,n_projections)
# This method (recommended) uses preprocessing of the images and projections
# to speed-up the registration
params.coarse_registration_method = IMP.em2d.ALIGN2D_PREPROCESSING
# use true if you want to save the projections from the model that best
# match the Em images
params.save_match_images = False
score_function=IMP.em2d.EM2DScore()
em2d_restraint = IMP.em2d.Em2DRestraint()
em2d_restraint.setup(score_function, params)
em2d_restraint.set_images(em_images)
em2d_restraint.set_name("em2d restraint")
container = IMP.container.ListSingletonContainer(core.get_leaves(prot))
em2d_restraint.set_particles(container)
em2d_restraints_set=IMP.RestraintSet()
# The next two lines are commented, because the optimization of the example
# is expensive. To run the full example, uncomment them (It can take a few
# hours).
#em2d_restraints_set.add_restraint(em2d_restraint)
#em2d_restraints_set.set_weight(1000) # weight for the em2D restraint
print "adding em2d restraint "
m.add_restraint(em2d_restraints_set)
# Add all restraints to a model
print "model has ",m.get_number_of_restraints(),"restraints"
# MONTECARLO OPTIMIZATION
s=core.MonteCarlo(m)
# Add movers for the rigid bodies
movers=[]
for rbd in rigid_bodies:
    movers.append(core.RigidBodyMover(rbd,5,2))
s.add_movers(movers)
print "MonteCarlo sampler has",s.get_number_of_movers(),"movers"
# Add an optimizer state to save intermediate configurations of the hierarchy
o_state=atom.WritePDBOptimizerState(chains,"intermediate-step-%1%.pdb")
o_state.set_skip_steps(10)
s.add_optimizer_state(o_state)
ostate2 = WriteStatisticsOptimizerScore()
s.add_optimizer_state(ostate2)
# Perform optimization
# m.set_gather_statistics(True) # Writes a lot of information!
temperatures=[200,100,60,40,20,5]
# 200 steps recommended for accurate optimization; a smaller number is used
# here for demonstration purposes
optimization_steps = 10
for T in temperatures:
    s.set_kt(T)
    s.optimize(optimization_steps)
atom.write_pdb(prot,"solution.pdb")
# Check that the optimization achieves distances close to the solution
print "*** End optimization ***"
new_centers = []
for rbd in rigid_bodies:
    print "chain has",rbd.get_number_of_members(), \
        "atoms","coordinates: ",rbd.get_coordinates()
    new_centers.append(rbd.get_coordinates())
d01 = alg.get_distance(new_centers[0],new_centers[1])
d12 = alg.get_distance(new_centers[1],new_centers[2])
d23 = alg.get_distance(new_centers[2],new_centers[3])
d30 = alg.get_distance(new_centers[3],new_centers[0])
print "Distances at the end of the optimization: d01 =", \
    d01,"d12 =",d12,"d23 =",d23,"d30 =",d30