models.py

#!/usr/bin/env python

from __future__ import print_function
import IMP.multifit
from IMP import OptionParser

__doc__ = "Write output models as PDB files."

# analyse the ensemble, first we will do the rmsd stuff

def parse_args():
    usage = """%prog [options] <asmb.input> <proteomics.input>
           <mapping.input> <combinations> <model prefix>

Write output models.
"""
    parser = OptionParser(usage)
    parser.add_option("-m", "--max", type="int", dest="max", default=None,
                      help="maximum number of models to write")
    (options, args) = parser.parse_args()
    if len(args) != 5:
        parser.error("incorrect number of arguments")
    return options, args


def run(asmb_fn, proteomics_fn, mapping_fn, combs_fn, model_output, max_comb):
    # get rmsd for subunits
    mdl = IMP.kernel.Model()
    combs = IMP.multifit.read_paths(combs_fn)
    sd = IMP.multifit.read_settings(asmb_fn)
    sd.set_was_used(True)
    prot_data = IMP.multifit.read_proteomics_data(proteomics_fn)
    mapping_data = IMP.multifit.read_protein_anchors_mapping(prot_data,
                                                             mapping_fn)
    ensmb = IMP.multifit.load_ensemble(sd, mdl, mapping_data)
    mhs = ensmb.get_molecules()
    print("number of combinations:", len(combs), max_comb)
    for i, comb in enumerate(combs[:max_comb]):
        if i % 500 == 0:
            print(i)
        ensmb.load_combination(comb)
        print(model_output + "." + str(i) + ".pdb")
        IMP.atom.write_pdb(mhs, model_output + "." + str(i) + ".pdb")
        ensmb.unload_combination(comb)


def main():
    options, args = parse_args()
    run(args[0], args[1], args[2], args[3], args[4], options.max)

if __name__ == "__main__":
    main()