rmf/multiresolution.py

## \example rmf/multiresolution.py
# In this example a pdb is converted into a multiresolution model,
# somewhat painfully.

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

IMP.base.setup_from_argv(sys.argv, "Create a multiresolution rmf file")

pdbname = IMP.rmf.get_example_path("big.pdb")

m = IMP.kernel.Model()
h = IMP.atom.read_pdb(pdbname, m)
IMP.atom.add_bonds(h)

chains = IMP.atom.get_by_type(h, IMP.atom.CHAIN_TYPE)

if IMP.base.get_bool_flag("run_quick_test"):
    chains = [chains[0]]


def recursive_approximation(res):
    print("approximating", res)
    lr = len(res)
    if lr <= 1:
        return res
    if lr > 4:
        me = recursive_approximation(res[0:lr // 4])\
            + recursive_approximation(res[lr // 4: lr // 2])\
            + recursive_approximation(res[lr // 2: 3 * lr // 4])\
            + recursive_approximation(res[3 * lr // 4: lr])
    else:
        me = res
    p = IMP.kernel.Particle(m)
    hc = IMP.atom.Hierarchy.setup_particle(p)
    IMP.atom.setup_as_approximation(p, res)
    nm = str(IMP.atom.Residue(res[0]).get_index()) + "-"\
        + str(IMP.atom.Residue(res[-1]).get_index())
    p.set_name(nm)
    for mm in me:
        hc.add_child(mm)
    return [hc]

for c in chains:
    res = IMP.atom.get_by_type(h, IMP.atom.RESIDUE_TYPE)
    for r in res:
        c.remove_child(r)
        lvs = IMP.atom.get_leaves(r)
        IMP.atom.setup_as_approximation(r, lvs)
    lr = len(res)
    me = recursive_approximation(res[0:lr // 4])\
        + recursive_approximation(res[lr // 4: lr // 2])\
        + recursive_approximation(res[lr // 2: 3 * lr // 4])\
        + recursive_approximation(res[3 * lr // 4: lr])
    for mm in me:
        c.add_child(mm)

print("writing")
fn = IMP.base.create_temporary_file_name("multires", ".rmf")
rmf = RMF.create_rmf_file(fn)
print("adding")
IMP.rmf.add_hierarchies(rmf, chains)

print("see file", fn)