create_gmm.py

from __future__ import print_function
import IMP
import IMP.em
import IMP.isd
import IMP.isd.gmm_tools
from IMP import ArgumentParser

import os


def parse_args():
    desc = """
    Create a GMM from either density file (.mrc), a pdb file (.pdb)
    Will detect input format from extension.
    Outputs as text and optionally as a density map
    see help(-h)
"""
    p = ArgumentParser(description=desc)

    p.add_argument("-t", "--covar_type", dest="covar_type", default='full',
                   choices=['spherical', 'tied', 'diag', 'full'],
                   help="covariance type for the GMM")
    p.add_argument("-m", "--out_map", dest="out_map", default='',
                   help="write out the gmm to an mrc file")
    p.add_argument("-a", "--apix", dest="apix", default=1.0, type=float,
                   help="if you don't provide a map, set the voxel_size "
                        "here (for sampling)")
    p.add_argument("-n", "--num_samples", dest="num_samples", default=1000000,
                   type=int, help="num samples to draw from the density map")
    p.add_argument("-i", "--num_iter", dest="num_iter", default=100, type=int,
                   help="num iterations of GMM")
    p.add_argument("-s", "--threshold", dest="threshold", default=0.0,
                   type=float, help="threshold for the map before sampling")

    p.add_argument("-f", "--force_radii", dest="force_radii", default=-1.0,
                   type=float,
                   help="force radii to be this value (spherical) "
                        "-1 means deactivated ")
    p.add_argument("-w", "--force_weight", dest="force_weight", default=-1.0,
                   type=float,
                   help="force weight to be this value (spherical) "
                        "-1 means deactivated ")
    p.add_argument("-e", "--force_weight_frac", dest="force_weight_frac",
                   action="store_true", default=False,
                   help="force weight to be 1.0/(num centers). "
                        "Takes precedence over -w")
    p.add_argument("-d", "--use_dirichlet", dest="use_dirichlet",
                   default=False, action="store_true",
                   help="use dirichlet process for fit")

    p.add_argument("-k", "--multiply_by_mass", dest="multiply_by_mass",
                   default=False, action="store_true",
                   help="if set, will multiply all weights by the total mass "
                        "of the particles (PDB ONLY)")
    p.add_argument("-x", "--chain", dest="chain", default=None,
                   help="If you passed a PDB file, read this chain")

    p.add_argument("-z", "--use_cpp", dest="use_cpp", default=False,
                   action="store_true",
                   help="EXPERIMENTAL. Uses the IMP GMM code. "
                        "Requires isd_emxl")
    p.add_argument("data_file", help="data file name")
    p.add_argument("n_centers", type=int, help="number of centers")
    p.add_argument("out_file", help="output file name")
    return p.parse_args()


def run(args):
    data_fn = args.data_file
    ncenters = args.n_centers
    out_txt_fn = args.out_file
    mdl = IMP.Model()

    if not os.path.isfile(data_fn):
        raise Exception("The data file you entered: " + data_fn
                        + " does not exist!")

    # get points for fitting the GMM
    ext = data_fn.split('.')[-1]
    mass_multiplier = 1.0
    if ext == 'pdb':
        mh = IMP.atom.read_pdb(
            data_fn, mdl, IMP.atom.NonWaterNonHydrogenPDBSelector())
        if args.chain:
            mps = IMP.atom.Selection(
                mh, chain=args.chain).get_selected_particles()
        else:
            mps = IMP.core.get_leaves(mh)

        if args.multiply_by_mass:
            mass_multiplier = sum(IMP.atom.Mass(p).get_mass() for p in mps)

        pts = [IMP.core.XYZ(p).get_coordinates() for p in mps]
        bbox = None
    elif ext == 'mrc':
        dmap = IMP.em.read_map(data_fn, IMP.em.MRCReaderWriter())
        bbox = IMP.em.get_bounding_box(dmap)
        dmap.set_was_used(True)
        print('sampling points')
        pts = IMP.isd.sample_points_from_density(
            dmap, args.num_samples, args.threshold)
    else:
        raise ValueError("data_fn extension must be pdb or mrc")

    # Do fitting to points
    if not args.use_cpp:
        density_ps = []
        print('fitting gmm')

        if args.force_weight_frac:
            force_weight = 1.0/ncenters
        else:
            force_weight = args.force_weight
        if force_weight != -1:
            print('weight forced to', force_weight)
        if not args.use_dirichlet:
            IMP.isd.gmm_tools.fit_gmm_to_points(
                pts, ncenters, mdl, density_ps, args.num_iter, args.covar_type,
                force_radii=args.force_radii, force_weight=args.force_weight,
                mass_multiplier=mass_multiplier)
        else:
            IMP.isd.gmm_tools.fit_dirichlet_gmm_to_points(
                pts, ncenters, mdl, density_ps, args.num_iter, args.covar_type,
                mass_multiplier=mass_multiplier)

    else:
        try:
            import isd_emxl  # noqa: F401
        except ImportError:
            print("This option is experimental, only works if you "
                  "have isd_emxl")
        gmm_threshold = 0.01
        density_ps = IMP.isd_emxl.fit_gaussians_to_density(
            mdl, dmap, args.num_samples, ncenters, args.num_iter,
            args.threshold, gmm_threshold)

    # Write to files
    comments = ['Created by create_gmm.py, IMP.isd version %s'
                % IMP.isd.get_module_version()]
    comments.append('data_fn: ' + IMP.get_relative_path(out_txt_fn, data_fn))
    comments.append('ncenters: %d' % ncenters)
    for key in ('covar_type', 'apix', 'num_samples', 'num_iter',
                'threshold', 'force_radii', 'force_weight',
                'force_weight_frac', 'use_dirichlet', 'multiply_by_mass',
                'chain'):
        comments.append('%s: %s' % (key, repr(getattr(args, key))))
    IMP.isd.gmm_tools.write_gmm_to_text(density_ps, out_txt_fn, comments)
    if args.out_map != '':
        IMP.isd.gmm_tools.write_gmm_to_map(density_ps, args.out_map,
                                           args.apix, bbox)


def main():
    args = parse_args()
    run(args)


if __name__ == "__main__":
    main()