you've got clones!

dolly_the_ligand.py

from openmmtools.testsystems import HostGuestVacuum
from simtk import openmm as mm

testsystem = HostGuestVacuum(constraints=None)
ligand_residue = list(testsystem.topology.residues())[1]
atoms_in_ligand = [atom.index for atom in ligand_residue.atoms()]


def filter_from_subset(atom_subset=None):
    if type(atom_subset) == type(None):
        atom_filter = lambda i: True
    else:
        atom_filter = lambda i: (i in atom_subset)
    return atom_filter


def clone_bond_force(bond_force, atom_subset=None):
    atom_filter = filter_from_subset(atom_subset)

    new_force = mm.HarmonicBondForce()
    for bond in range(bond_force.getNumBonds()):
        (p1, p2, length, k) = bond_force.getBondParameters(bond)
        if atom_filter(p1) and atom_filter(p2):
            new_force.addBond(p1, p2, length, k)

    return new_force


def clone_angle_force(angle_force, atom_subset=None):
    atom_filter = filter_from_subset(atom_subset)

    new_force = mm.HarmonicAngleForce()
    for angle in range(angle_force.getNumAngles()):
        (p1, p2, p3, theta0, k) = angle_force.getAngleParameters(angle)
        if atom_filter(p1) and atom_filter(p2) and atom_filter(p3):
            new_force.addAngle(p1, p2, p3, theta0, k)

    return new_force


def clone_torsion_force(torsion_force, atom_subset=None):
    atom_filter = filter_from_subset(atom_subset)

    new_force = mm.PeriodicTorsionForce()
    for torsion in range(torsion_force.getNumTorsions()):
        (p1, p2, p3, p4, periodicity, phase, k) = torsion_force.getTorsionParameters(torsion)
        if atom_filter(p1) and atom_filter(p2) and atom_filter(p3) and atom_filter(
                p4):  # and ((k / unit.kilojoule_per_mole) != 0.0):
            new_force.addTorsion(p1, p2, p3, p4, periodicity, phase, k)

    return new_force


def clone_nonbonded_force(nonbonded_force, atom_subset=None):
    if type(atom_subset) == type(None):
        atom_subset = list(range(nonbonded_force.getNumParticles()))

    new_force = mm.NonbondedForce()
    atom_mapping = {}
    for i in range(nonbonded_force.getNumParticles()):
        if i in atom_subset:
            charge, sigma, epsilon = nonbonded_force.getParticleParameters(i)
            atom_mapping[i] = new_force.addParticle(charge, sigma, epsilon)

    # get all the exceptions
    for exception in range(nonbonded_force.getNumExceptions()):
        (i, j, charge_prod, sig, eps) = nonbonded_force.getExceptionParameters(exception)
        if (i in atom_subset) and (j in atom_subset):
            new_force.addException(atom_mapping[i], atom_mapping[j], charge_prod, sig, eps)

    # clone all the other things...
    new_force.setNonbondedMethod(nonbonded_force.getNonbondedMethod())
    new_force.setCutoffDistance(nonbonded_force.getCutoffDistance())
    new_force.setEwaldErrorTolerance(nonbonded_force.getEwaldErrorTolerance())
    new_force.setPMEParameters(*nonbonded_force.getPMEParameters())
    new_force.setLJPMEParameters(*nonbonded_force.getLJPMEParameters())
    new_force.setUseSwitchingFunction(nonbonded_force.getUseSwitchingFunction())
    new_force.setUseDispersionCorrection(nonbonded_force.getUseDispersionCorrection())
    new_force.setSwitchingDistance(nonbonded_force.getSwitchingDistance())
    new_force.setReactionFieldDielectric(nonbonded_force.getReactionFieldDielectric())

    # Assert that we're not missing any of the things we aren't cloning at this time
    # TODO: clone these too
    assert (nonbonded_force.getNumExceptionParameterOffsets() == 0)
    assert (nonbonded_force.getNumParticleParameterOffsets() == 0)
    assert (nonbonded_force.getNumGlobalParameters() == 0)

    return new_force


def clone_system(system, atom_subset=None):
    if type(atom_subset) == type(None):
        atom_subset = list(range(system.getNumParticles()))

    # TODO: don't assume order here...
    bond_force, angle_force, torsion_force, nonbonded_force = system.getForces()[:4]

    new_system = mm.System()
    for i in atom_subset:
        new_system.addParticle(system.getParticleMass(i))

    new_system.addForce(clone_bond_force(bond_force, atom_subset))
    new_system.addForce(clone_angle_force(angle_force, atom_subset))
    new_system.addForce(clone_torsion_force(torsion_force, atom_subset))
    new_system.addForce(clone_nonbonded_force(nonbonded_force, atom_subset))

    return new_system