pmi/dof/__init__.py

"""@namespace IMP.pmi.dof
   Create movers and set up constraints for PMI objects.
   See the documentation of the DegreesOfFreedom class for more information.
"""

import IMP
import IMP.atom
import IMP.algebra
import IMP.pmi
import IMP.pmi.topology
import IMP.pmi.samplers
import IMP.pmi.tools
import IMP.pmi.samplers
import warnings


def create_rigid_body_movers(dof, maxtrans, maxrot):
    mvs = []
    for rb in dof.rigid_bodies:
        mvs.append(IMP.core.RigidBodyMover(rb.get_model(), rb,
                                           maxtrans, maxrot))
    return mvs


class DegreesOfFreedom:
    """Simplify creation of constraints and movers for an IMP Hierarchy.

       * The various "create X" functions make movers for system components
         as well as set up necessary constraints. For each of these functions,
         you can generally pass PMI objects like
         [Molecule](@ref IMP::pmi::topology::Molecule) or slices thereof.
       * DegreesOfFreedom.create_rigid_body() lets you rigidify a molecule
         (but allows you to also pass "nonrigid" components which move with
         the body and also independently).
       * DegreesOfFreedom.create_super_rigid_body() sets up a special
         "Super Rigid Body" which moves rigidly but is not always constrained
         to be rigid (so you can later move the parts separately). This is
         good for speeding up sampling.
       * DegreesOfFreedom.create_flexible_beads() sets up particles to move
         individually.
       * DegreesOfFreedom.setup_md() sets up particles to move with molecular
         dynamics. Note that this is not (yet) compatible with rigid bodies,
         and only works with some restraints.
       * DegreesOfFreedom.constrain_symmetry() makes a symmetry constraint so
         that clones automatically move with their references. If instead you
         want a softer restraint, check out the
         [SymmetryRestraint](@ref IMP::pmi::restraints::stereochemistry::SymmetryRestraint).
       * When you are done you can access all movers with
         DegreesOfFreedom.get_movers(). If you have set up rigid, super rigid,
         or flexible beads, pass the movers to the `monte_carlo_sample_objects`
         argument of
         [ReplicaExchange](@ref IMP::pmi::macros::ReplicaExchange).
       * If you are running MD, you have to separately pass the particles
         (also returned from DegreesOfFreedom.setup_md()) to the
         `molecular_dynamics_sample_objects` argument of
         [ReplicaExchange](@ref IMP::pmi::macros::ReplicaExchange). Check
         out an [MD example here](pmi_2atomistic_8py-example.html).
    """  # noqa: E501
    def __init__(self, model):
        self.model = model
        self.movers = []
        self.fb_movers = []  # stores movers corresponding to floppy parts
        self.rigid_bodies = []  # stores rigid body objects
        # stores all beads including nonrigid members of rigid bodies
        self.flexible_beads = []
        self.nuisances = []
        # Keys are the RigidBody objects, values are list of movers
        self._rb2mov = {}
        # the following is needed to keep track of disabled movers
        self.movers_particles_map = IMP.pmi.tools.OrderedDict()
        self.movers_rb_map = {}
        self.movers_xyz_map = {}
        self.disabled_movers = []

    def _get_nonrigid_hiers(self, nonrigid_parts, rigid_hiers, resolution):
        """Get Hierarchy objects for nonrigid parts. Make sure that they are
           a subset of the rigid body Hierarchies."""
        if not nonrigid_parts:
            return
        nr_hiers = IMP.pmi.tools.input_adaptor(nonrigid_parts, resolution,
                                               flatten=True)
        if nr_hiers:
            rb_idxs = set(h.get_particle_index() for h in rigid_hiers)
            for h in nr_hiers:
                p = h.get_particle()
                if p.get_index() not in rb_idxs:
                    raise ValueError(
                                "You tried to create nonrigid members from "
                                "particles that aren't in the RigidBody!")
        return nr_hiers

    def create_rigid_body(self, rigid_parts, nonrigid_parts=None,
                          max_trans=4.0, max_rot=0.5, nonrigid_max_trans=4.0,
                          resolution='all', name=None):
        """Create rigid body constraint and mover
        @param rigid_parts Can be one of the following inputs:
               IMP Hierarchy, PMI System/State/Molecule/TempResidue, a list/set
               (of list/set) of them or a RigidBody object.
               Must be uniform input, however. No mixing object types.
        @param nonrigid_parts Same input format as rigid_parts.
               Must be a subset of rigid_parts particles.
        @param max_trans Maximum rigid body translation
        @param max_rot Maximum rigid body rotation
        @param nonrigid_max_trans Maximum step for the nonrigid (bead)
               particles
        @param resolution Only used if you pass PMI objects. Probably you
               want 'all'.
        @param name Rigid body name (if None, use IMP default)
        @return (rb_movers,rb_object)
        @note If you want all resolutions, pass PMI objects because this
              function will get them all. Alternatively you can do your
              selection elsewhere and just pass hierarchies.
        """

        rb_movers = []

        # ADD CHECK: these particles are not already part of some RB or SRB

        # First, is this already a rigid body?
        if type(rigid_parts) is IMP.core.RigidBody:
            warnings.warn("Rigid Body Already Setup", IMP.pmi.ParameterWarning)
            rb = rigid_parts
            model = rb.get_model()
            if name is None:
                name = rb.get_name()
            hiers = [IMP.atom.get_leaves(IMP.atom.Hierarchy(
                model.get_particle(i)))[0]
                for i in rb.get_member_particle_indexes()]
        else:
            # Otherwise, setup RB
            hiers = IMP.pmi.tools.input_adaptor(rigid_parts,
                                                resolution,
                                                flatten=True)

            if not hiers:
                warnings.warn(
                        "No hierarchies were passed to create_rigid_body()",
                        IMP.pmi.ParameterWarning)
                return []

        # Need to do this before rigid body is set up so that we leave the
        # system in a consistent state if a sanity check fails
        nr_hiers = self._get_nonrigid_hiers(nonrigid_parts, hiers, resolution)

        if type(rigid_parts) is not IMP.core.RigidBody:
            model = hiers[0].get_model()

            # we need to use the following constructor because the
            # IMP.core.create_rigid_body seems to construct an arbitrary
            # reference frame, which will be different for all copies.
            # therefore, symmetry won't work all the time

            com = IMP.atom.CenterOfMass.setup_particle(IMP.Particle(model),
                                                       hiers)
            comcoor = IMP.core.XYZ(com).get_coordinates()

            # Don't needlessly update the COM of this particle for the
            # entire sampling run; we are done with both the ScoreState
            # and the Particle.
            IMP.atom.CenterOfMass.teardown_particle(com)
            model.remove_particle(com)

            tr = IMP.algebra.Transformation3D(
                IMP.algebra.get_identity_rotation_3d(), comcoor)
            rf = IMP.algebra.ReferenceFrame3D(tr)
            rbp = IMP.Particle(model)
            rb = IMP.core.RigidBody.setup_particle(rbp, rf)
            for h in hiers:
                rb.add_member(h.get_particle())

        self.rigid_bodies.append(rb)
        rb.set_coordinates_are_optimized(True)
        rb_mover = IMP.core.RigidBodyMover(rb.get_model(), rb, max_trans,
                                           max_rot)
        if name is not None:
            rb.set_name(name)
            rb_mover.set_name(name)
        rb_movers.append(rb_mover)
        self.movers_particles_map[rb_mover] = []
        self.movers_rb_map[rb_mover] = [rb]
        rb_mover.set_was_used(True)
        for h in hiers:
            self.movers_particles_map[rb_mover] += IMP.atom.get_leaves(h)
        # setup nonrigid parts
        if nr_hiers:
            floatkeys = IMP.core.RigidBodyMember.get_internal_coordinate_keys()
            for h in nr_hiers:
                self.flexible_beads.append(h)
                p = h.get_particle()
                rb.set_is_rigid_member(p.get_index(), False)
                for fk in floatkeys:
                    p.set_is_optimized(fk, True)
                fbmv = IMP.core.BallMover(p.get_model(), p,
                                          IMP.FloatKeys(floatkeys),
                                          nonrigid_max_trans)
                self.fb_movers.append(fbmv)
                self.movers_particles_map[fbmv] = IMP.atom.get_leaves(h)
                self.movers_xyz_map[fbmv] = IMP.atom.get_leaves(h)
                fbmv.set_was_used(True)
                rb_movers.append(fbmv)

        self.movers += rb_movers  # probably need to store more info
        self._rb2mov[rb] = rb_movers  # dictionary relating rb to movers

        return rb_movers, rb

    def create_main_chain_mover(self, molecule, resolution=10,
                                lengths=[5, 10]):
        """Create crankshaft moves from a set of SUPER rigid body mover
           from one molecule.
        See http://scfbm.biomedcentral.com/articles/10.1186/1751-0473-3-12
        """
        hiers = IMP.pmi.tools.input_adaptor(molecule, resolution, flatten=True)

        for length in lengths:
            for n in range(len(hiers)-length):
                hs = hiers[n+1:n+length]
                self.create_super_rigid_body(
                    hs, max_trans=0.0, max_rot=0.05,
                    axis=(hiers[n].get_particle(),
                          hiers[n+length].get_particle()))

    def create_super_rigid_body(self, srb_parts, max_trans=1.0, max_rot=0.1,
                                chain_min_length=None, chain_max_length=None,
                                resolution='all', name=None, axis=None):
        """Create SUPER rigid body mover from one or more hierarchies.

           Can also create chain of SRBs. If you don't pass chain min/max,
           it'll treat everything you pass as ONE rigid body.
           If you DO pass chain min/max, it'll expect srb_parts is a list
           and break it into bits.
           @param srb_parts Can be one of the following inputs:
                  IMP Hierarchy, PMI System/State/Molecule/TempResidue,
                  or a list/set (of list/set) of them.
                  Must be uniform input, however. No mixing object types.
           @param max_trans Maximum super rigid body translation
           @param max_rot Maximum super rigid body rotation
           @param chain_min_length Create a CHAIN of super rigid bodies -
                  must provide list; this parameter is the minimum
                  chain length.
           @param chain_max_length Maximum chain length
           @param resolution Only used if you pass PMI objects. Probably you
                  want 'all'.
           @param name The name of the SRB (hard to assign a good one
                  automatically)
           @param axis A tuple containing two particles which are used to
                  compute the rotation axis of the SRB. The default is None,
                  meaning that the rotation axis is random.

           @note If you set the chain parameters, will NOT create an SRB from
                 all of them together, but rather in groups made from the
                 outermost list.
        """

        srb_movers = []

        # organize hierarchies based on chains
        if chain_min_length is None and chain_max_length is None:
            # the "chain" is just everything together
            h = IMP.pmi.tools.input_adaptor(srb_parts, resolution,
                                            flatten=True)
            if len(h) == 0:
                warnings.warn(
                    'No hierarchies were passed to create_super_rigid_body()',
                    IMP.pmi.ParameterWarning)
                return srb_movers
            srb_groups = [h]
        else:
            if not hasattr(srb_parts, '__iter__'):
                raise Exception("You tried to make a chain without a list!")
            srb_groups = [IMP.pmi.tools.input_adaptor(
                h, resolution, flatten=True, warn_about_slices=False)
                for h in srb_parts]

        # create SRBs either from all hierarchies or chain
        if chain_min_length is None and chain_max_length is None:
            mv = self._setup_srb(srb_groups, max_trans, max_rot, axis)
            if mv:
                mv.set_was_used(True)
                srb_movers.append(mv)
        elif chain_min_length is not None and chain_max_length is not None:
            for hs in IMP.pmi.tools.sublist_iterator(
                    srb_groups, chain_min_length, chain_max_length):
                mv = self._setup_srb(hs, max_trans, max_rot, axis)
                if mv:
                    mv.set_was_used(True)
                    srb_movers.append(mv)
        else:
            raise Exception(
                "DegreesOfFreedom: SetupSuperRigidBody: if you want "
                "chain, specify min AND max")
        self.movers += srb_movers
        if name is not None:
            if len(srb_movers) > 1:
                for n, mv in enumerate(srb_movers):
                    mv.set_name(name + '_' + str(n))
            else:
                srb_movers[0].set_name(name)
        return srb_movers

    def _setup_srb(self, hiers, max_trans, max_rot, axis):
        if axis is None:
            srbm = IMP.pmi.TransformMover(
                hiers[0][0].get_model(), max_trans, max_rot)
        else:
            srbm = IMP.pmi.TransformMover(
                hiers[0][0].get_model(), axis[0], axis[1], max_trans, max_rot)
        srbm.set_name("Super rigid body transform mover")
        srbm.set_was_used(True)
        super_rigid_rbs, super_rigid_xyzs \
            = IMP.pmi.tools.get_rbs_and_beads(hiers)
        ct = 0
        self.movers_particles_map[srbm] = []
        for h in hiers:
            self.movers_particles_map[srbm] += IMP.atom.get_leaves(h)
        for xyz in super_rigid_xyzs:
            srbm.add_xyz_particle(xyz)
            ct += 1
        for rb in super_rigid_rbs:
            srbm.add_rigid_body_particle(rb)
            ct += 1
        if ct > 1:
            return srbm
        else:
            return 0

    def create_flexible_beads(self, flex_parts, max_trans=3.0,
                              resolution='all'):
        """Create a chain of flexible beads

           @param flex_parts Can be one of the following inputs:
                  IMP Hierarchy, PMI System/State/Molecule/TempResidue,
                  or a list/set (of list/set) of them.
                  Must be uniform input, however. No mixing object types.
           @param max_trans Maximum flexible bead translation
           @param resolution Only used if you pass PMI objects. Probably
                  you want 'all'.
        """

        fb_movers = []
        hiers = IMP.pmi.tools.input_adaptor(flex_parts,
                                            resolution,
                                            flatten=True)
        if not hiers or len(hiers) == 0:
            warnings.warn(
                'No hierarchies were passed to create_flexible_beads()',
                IMP.pmi.ParameterWarning)
            return fb_movers
        for h in hiers:
            p = h.get_particle()
            IMP.core.XYZ(p).set_coordinates_are_optimized(True)
            if IMP.core.RigidMember.get_is_setup(h) \
                    or IMP.core.NonRigidMember.get_is_setup(h):
                raise Exception(
                    "Cannot create flexible beads from members of rigid body")
            self.flexible_beads.append(h)
            fbmv = IMP.core.BallMover(p.get_model(), p, max_trans)
            fb_movers.append(fbmv)
            fbmv.set_was_used(True)
            self.fb_movers.append(fbmv)
            self.movers_particles_map[fbmv] = IMP.atom.get_leaves(h)
            self.movers_xyz_map[fbmv] = IMP.atom.get_leaves(h)
        self.movers += fb_movers
        return fb_movers

    def create_nuisance_mover(self, nuisance_p, step_size, name=None):
        """Create MC normal mover for nuisance particles.
        We will add an easier interface to add all of them from a PMI restraint
        @param nuisance_p The Nuisance particle (an ISD::Scale)
        @param step_size The maximum step size for Monte Carlo
        @param name The name of the mover, useful for better output reading.
        """
        mv = IMP.core.NormalMover([nuisance_p],
                                  IMP.FloatKeys([IMP.FloatKey("nuisance")]),
                                  step_size)
        if name is not None:
            mv.set_name(name)
        mv.set_was_used(True)
        self.nuisances.append(nuisance_p)
        self.movers.append(mv)
        return [mv]

    def setup_md(self, hspec):
        """Setup particles for MD simulation. Returns all particles, just
        pass this to molecular_dynamics_sample_objects in ReplicaExchange.
        @param hspec Can be one of the following inputs:
               IMP Hierarchy, PMI System/State/Molecule/TempResidue,
               or a list/set (of list/set) of them.
               Must be uniform input, however. No mixing object types.
        """
        hiers = IMP.pmi.tools.input_adaptor(hspec, flatten=True)
        model = hiers[0].get_model()
        all_ps = []
        for hl in hiers:
            for h in IMP.core.get_leaves(hl):
                p = h.get_particle()
                pxyz = IMP.core.XYZ(model, p.get_index())
                pxyz.set_coordinates_are_optimized(True)
                IMP.atom.LinearVelocity.setup_particle(p, [0., 0., 0.])
                all_ps.append(p)
        return all_ps

    def constrain_symmetry(self, references, clones, transform,
                           resolution='all', type="AXIAL"):
        """Create a symmetry constraint. Checks:
        same number of particles
        disable ANY movers involving symmetry copies
         (later may support moving them WITH references,
         but requires code to propagate constraint)
        @param references Can be one of the following inputs:
               IMP Hierarchy, PMI System/State/Molecule/TempResidue,
               or a list/set (of list/set) of them
        @param clones Same format as references
        @param transform The transform that moves a clone onto a reference
               IMP.algebra.Transformation3D
        @param resolution Only used if you pass PMI objects.
               If you have a multires system, assuming each are rigid
               bodies you probably only need one resolution.
        @param type of symmetry. Implemented = AXIAL, RIGID_BODY
        """

        # get all RBs and particles
        href = IMP.pmi.tools.input_adaptor(references, resolution,
                                           flatten=True)
        hclones = IMP.pmi.tools.input_adaptor(clones, resolution, flatten=True)

        ref_rbs, ref_beads = IMP.pmi.tools.get_rbs_and_beads(href)
        clones_rbs, clones_beads = IMP.pmi.tools.get_rbs_and_beads(hclones)

        for ref, clone in zip(ref_rbs, clones_rbs):
            IMP.core.Reference.setup_particle(clone, ref)

        for ref, clone in zip(ref_beads, clones_beads):
            IMP.core.Reference.setup_particle(clone, ref)

        # removing movers involved in clones
        self.disable_movers(hclones)

        if type == "AXIAL":
            sm = IMP.core.TransformationSymmetry(transform)

        if type == "RIGID_BODY":
            p = IMP.Particle(self.model)
            p.set_name("RigidBody_Symmetry")
            rb = IMP.core.RigidBody.setup_particle(
                p, IMP.algebra.ReferenceFrame3D(transform))
            for cp in [(10, 0, 0), (0, 10, 0), (0, 0, 10)]:
                p = IMP.Particle(self.model)
                IMP.core.XYZ.setup_particle(p, cp)
                rb.add_member(p)
            sm = IMP.core.TransformationSymmetry(rb.get_particle_index())
            self.rigid_bodies.append(rb)
            rb.set_coordinates_are_optimized(True)
            rb_mover_tr = IMP.core.RigidBodyMover(
                rb.get_model(), rb.get_particle_index(), 0.0, 1.0)
            rb_mover_rt = IMP.core.RigidBodyMover(
                rb.get_model(), rb.get_particle_index(), 10.0, 0.0)

            rb_mover_tr.set_name("RigidBody_Symmetry_Mover_Translate")
            rb_mover_rt.set_name("RigidBody_Symmetry_Mover_Rotate")
            print('Created rigid body symmetry restraint')
            self.movers_particles_map[rb_mover_tr] = []
            self.movers_particles_map[rb_mover_rt] = []
            self.movers_rb_map[rb_mover_tr] = [rb]
            self.movers_rb_map[rb_mover_rt] = [rb]
            for h in hclones:
                self.movers_particles_map[rb_mover_tr] \
                    += IMP.atom.get_leaves(h)
                self.movers_particles_map[rb_mover_rt] \
                    += IMP.atom.get_leaves(h)
            self.movers.append(rb_mover_tr)  # probably need to store more info
            self.movers.append(rb_mover_rt)  # probably need to store more info
            # dictionary relating rb to movers
            self._rb2mov[rb] = [rb_mover_tr, rb_mover_rt]

        lsc = IMP.container.ListSingletonContainer(
            self.model,
            [bead.get_particle().get_index()
             for bead in clones_rbs+clones_beads])
        c = IMP.container.SingletonsConstraint(sm, None, lsc)
        self.model.add_score_state(c)
        print('Created symmetry restraint for', len(ref_rbs),
              'rigid bodies and', len(ref_beads), 'flexible beads')
        self.model.update()

    def __repr__(self):
        # would like something fancy like this:
        # - super-rigid "SRB1"
        #  - rigid "Mol1" (8 rigid, 3 nonrigid)
        #  - rigid "Mol2" (8 rigid, 3 nonrigid)
        #  - rigid "Mol3" (8 rigid, 3 nonrigid)
        return 'DegreesOfFreedom: ' + \
                "\n".join(repr(m) for m in self.movers)

    def optimize_flexible_beads(self, nsteps, temperature=1.0):
        '''Set up MC run with just flexible beads.
        Optimization works much better when restraints
        are already set up.'''
        mvs = list(_get_floppy_body_movers(self.get_flexible_beads(), 1.0))
        if mvs:
            mc = IMP.pmi.samplers.MonteCarlo(self.model, mvs, temperature)
            print("optimize_flexible_beads: optimizing %i flexible beads"
                  % len(self.get_flexible_beads()))
            mc.optimize(nsteps)
        else:
            print("optimize_flexible_beads: no particle to optimize")

    def get_movers(self):
        """Returns Enabled movers"""
        if self.disabled_movers:
            filtered_mover_list = []
            for mv in self.movers:
                if mv not in self.disabled_movers:
                    filtered_mover_list.append(mv)
            return filtered_mover_list
        else:
            return self.movers

    def get_floppy_body_movers(self):
        """Return all movers corresponding to individual beads"""
        return self.fb_movers

    def get_rigid_bodies(self):
        """Return list of rigid body objects"""
        return self.rigid_bodies

    def get_flexible_beads(self):
        "Return all flexible beads, including nonrigid members of rigid bodies"
        return self.flexible_beads

    def disable_movers(self, objects, mover_types=None):
        """Fix the position of the particles by disabling the corresponding
           movers.
        @param objects Can be one of the following inputs:
            IMP Hierarchy, PMI System/State/Molecule/TempResidue, or a
            list/set (of list/set) of them.
            Must be uniform input, however. No mixing object types.
        @param mover_types further filter the mover type that will be
            disabled; it can be a list of IMP.core.RigidBodyMover,
            IMP.core.BallMover etc etc if one wants to fix the corresponding
            rigid body, or the floppy bodies.
            An empty mover_types list is interpreted as all possible movers.
        It returns the list of fixed xyz particles (ie, floppy bodies/beads)
        and rigid bodies whose movers were disabled
        """
        hierarchies = IMP.pmi.tools.input_adaptor(objects,
                                                  pmi_resolution='all',
                                                  flatten=True)
        tmp_set = set()
        fixed_rb = set()
        fixed_xyz = set()
        if mover_types is None:
            mover_types = []

        inv_map = {}
        for mv, ps in self.movers_particles_map.items():
            for p in ps:
                if p in inv_map:
                    inv_map[p].append(mv)
                else:
                    inv_map[p] = [mv]

        for h in hierarchies:
            if h in inv_map:
                for mv in inv_map[h]:
                    if (type(mv) in mover_types or not mover_types):
                        tmp_set.add(mv)
                        if mv in self.movers_rb_map:
                            fixed_rb |= set(self.movers_rb_map[mv])
                        if mv in self.movers_xyz_map:
                            fixed_xyz |= set(self.movers_xyz_map[mv])
        print("Fixing %s movers" % (str(len(list(tmp_set)))))
        self.disabled_movers += list(tmp_set)
        return list(fixed_xyz), list(fixed_rb)

    def enable_all_movers(self):
        """Re-enable all movers: previously fixed particles will be released"""
        self.disabled_movers = []

    def get_nuisances_from_restraint(self, r):
        """Extract the nuisances from get_particles_to_sample()"""
        try:
            pslist = r.get_particles_to_sample()
        except AttributeError:
            raise Exception(
                "dof.get_nuisances_from_restraint(): the passed object "
                "does not have a get_particles_to_sample() function")
        for name in pslist:
            is_sampled = True
            if len(pslist[name]) == 3:
                ps, maxtrans, is_sampled = pslist[name]
            else:
                ps, maxtrans = pslist[name]
            if is_sampled:
                self.create_nuisance_mover(ps[0], maxtrans, name)


def _get_floppy_body_movers(fbs, maxtrans):
    for fb in fbs:
        # check is that is a rigid body member:
        if IMP.core.NonRigidMember.get_is_setup(fb):
            # if so force the particles to move anyway
            floatkeys = \
                IMP.core.RigidBodyMember.get_internal_coordinate_keys()
            for fk in floatkeys:
                fb.set_is_optimized(fk, True)
            yield IMP.core.BallMover(fb.get_model(), fb,
                                     IMP.FloatKeys(floatkeys),
                                     maxtrans)
        else:
            # otherwise use the normal ball mover
            yield IMP.core.BallMover(fb.get_model(), fb, maxtrans)