Python mdtraj.rmsd函数代码示例

def pairwise_distances(X, Y=None, index=None, metric="euclidean"):
    '''
    Compute the distance matrix from a vector array X and optional Y.
    This method takes either a vector array or a distance matrix,
    and returns a distance matrix. If the input is a vector array,
    the distances are computed. If the input is a distances matrix,
    it is returned instead.
    This method provides a safe way to take a distance matrix as input,
    while preserving compatibility with many other algorithms that take
    a vector array.

    :param X:  array [n_samples_a, n_samples_a]
        Array of pairwise distances between samples, or a feature array.
    :param Y:   array [n_samples_b, n_features]
        A second feature array only if X has shape [n_samples_a, n_features].
    :param index:  int, the index of element in X array
    :param metric: The metric to use when calculating distance between instances in a feature array.
        If metric ='rmsd', it should be computed by MDTraj
    :return: The distances
    '''
    if metric == "rmsd":
        if Y is None:
            distances_ = md.rmsd(X, X, index, parallel=True, precentered=True)
        else:
            #distances_ = np.empty((len(X), len(Y)), dtype=np.float32)
           # for i in xrange(len(Y)):
            distances_ = md.rmsd(X, Y, index, parallel=True, precentered=True)
        return distances_
    else:
        if Y is None:
            print "if Y is None"
            return sp.pairwise_distances(X, X[index], metric=metric)
        if index is None:
            print "if index is None, pairwise XX"
            return sp.pairwise_distances(X, X, metric=metric)

def main():
    parser = argparse.ArgumentParser(description='custom featurization of clc fah trjs')
    parser.add_argument('--ref', type=str, help='homology model pdb file')
    parser.add_argument('--trj', type=str, help='trajectory file')
    parser.add_argument('--mol2', type=str, help='homology model mol2 file (charges needed for dipole calc)')
    args = parser.parse_args()
     
    # load system data
    trj = mdtraj.load(args.trj, top=args.ref)
    hmodel = mdtraj.load(args.ref)
    
    ### feature 0: protein RMSD from hmodel ###
    pi_noh = [atom.index for atom in trj.top.atoms if ((atom.residue.is_protein) and (atom.element.symbol != 'H'))]
    p_rmsd = mdtraj.rmsd(trj, hmodel, atom_indices=pi_noh)

    ### feature 1: GLU128 RMSD from hmodel ###
    e128 = res_ndxs(hmodel, vs_ri['glu128'])
    e128_rmsd = mdtraj.rmsd(trj, hmodel, atom_indices=e128)
    
    ### feature 2: LYS317 and GLU318 RMSD from hmodel ###
    tl = np.concatenate((res_ndxs(hmodel, vs_ri['lys317']), res_ndxs(hmodel, vs_ri['glu318'])))
    tl_rmsd = mdtraj.rmsd(trj, hmodel, atom_indices=tl)

    ### feature 2: distance between ASP32 and LYS127 ###
    a32 = ele_ndxs(hmodel, vs_ri['asp32'], ['OD1', 'OD2'])
    l127 = ele_ndxs(hmodel, vs_ri['lys127'], ['NZ'])
    al_pairs = cartesian([a32, l127])
    # i think the asp oxygens are degenerate, so i'll look at the min here
    al_dist = np.amin(al_pairs, axis=1)

def calc_obs(traj):
    arg_cz_id = 2442
    glu_cd_id = 862
    lys_nz_id = 634
    tyr_oh_id = 2019
    inactive = mdt.load("./topologies/inactive.pdb")
    active = mdt.load("./topologies/active.pdb")

    aloop_atoms_list = [i.index for residue in np.arange(147, 168) for i in inactive.topology.residue(residue).atoms]
    all_heavy = [i.index for i in inactive.topology.atoms if i.residue.is_protein and i.element.name != "hydrogen"]
    print("Processing %s" % traj)
    # load the trajectory
    trj = mdt.load(traj, atom_indices=np.arange(inactive.n_atoms))

    inactive_rms = mdt.rmsd(trj, inactive, atom_indices=all_heavy)
    active_rms = mdt.rmsd(trj, active, atom_indices=all_heavy)
    aloop_rms = mdt.rmsd(trj, inactive, frame=0, atom_indices=aloop_atoms_list)
    distances = mdt.compute_distances(trj, np.vstack(([arg_cz_id, glu_cd_id], [lys_nz_id, glu_cd_id])))
    return dict(
        fname=os.path.basename(traj),
        inactive_rmsd=inactive_rms,
        active_rmsd=active_rms,
        aloop_inactive_rmsd=aloop_rms,
        glu_arg=distances[:, 0],
        gly_lys=distances[:, 1],
    )

    def test_ComparetoMDtraj(self):
        import mdtraj as md
        traj = pt.load(filename="./data/Tc5b.x",
                       top="./data/Tc5b.top")
        m_top = md.load_prmtop("./data/Tc5b.top")
        m_traj = md.load_mdcrd("./data/Tc5b.x", m_top)
        m_traj.xyz = m_traj.xyz * 10  # convert `nm` to `Angstrom` unit

        arr0 = pt.rmsd(traj, ref=0)
        arr1 = pt.rmsd(traj, ref=0)
        arr2 = pt.rmsd(traj, )
        a_md0 = md.rmsd(m_traj, m_traj, 0)
        aa_eq(arr0, arr1)
        aa_eq(arr0, arr2)
        aa_eq(arr0, a_md0)

        arr0 = pt.rmsd(traj, ref=-1)
        arr1 = pt.rmsd(traj, ref=-1)
        a_md = md.rmsd(m_traj, m_traj, -1)
        aa_eq(arr0, arr1)
        aa_eq(arr0, a_md)

        mask = ":[email protected],C"
        atm = traj.top(mask)
        arr0 = pt.rmsd(traj, ref=-1, mask=mask)
        arr1 = pt.rmsd(traj, mask=atm.indices, ref=-1)
        arr2 = pt.rmsd(traj, mask=list(atm.indices), ref=-1)
        arr3 = pt.rmsd(traj, mask=tuple(atm.indices), ref=-1)
        a_md = md.rmsd(m_traj, m_traj, -1, atm.indices)
        aa_eq(arr0, a_md)
        aa_eq(arr1, a_md)
        aa_eq(arr2, a_md)
        aa_eq(arr3, a_md)

        fa = Trajectory(traj)
        arr0 = pt.rmsd(fa, ref=-1, mask=mask)
        arr1 = pt.rmsd(fa, mask=atm.indices, ref=-1)
        arr2 = pt.rmsd(fa, mask=list(atm.indices), ref=-1)
        arr3 = pt.rmsd(fa, mask=tuple(atm.indices), ref=-1)
        a_md = md.rmsd(m_traj, m_traj, -1, atm.indices)
        aa_eq(arr0, a_md)
        aa_eq(arr1, a_md)
        aa_eq(arr2, a_md)
        aa_eq(arr3, a_md)

        fa = Trajectory(traj)
        mask = "[email protected]="
        atm = fa.top(mask)
        arr0 = pt.rmsd(fa, ref=4, mask=mask)
        a_md = md.rmsd(m_traj, m_traj, 4, atm.indices)

        # exclude 0-th frame for ref
        aa_eq(arr0, a_md)

def test_rmsd_atom_indices():
    native = md.load(get_fn('native.pdb'))
    t1 = md.load(get_fn('traj.h5'))

    atom_indices = np.arange(10)
    dist1 = md.rmsd(t1, native, atom_indices=atom_indices)

    t2 = md.load(get_fn('traj.h5'))
    t2.restrict_atoms(atom_indices)
    native.restrict_atoms(atom_indices)
    dist2 = md.rmsd(t2, native)
    
    eq(dist1, dist2)

def test_rmsd_ref_ainds():
    native = md.load(get_fn('native.pdb'))
    t1 = md.load(get_fn('traj.h5'))

    atom_indices = np.arange(10)
    dist1 = md.rmsd(t1, native, atom_indices=atom_indices,
                    ref_atom_indices=atom_indices)

    bad_atom_indices = np.arange(10, 20)
    t2 = md.load(get_fn('traj.h5'))
    dist2 = md.rmsd(t2, native, atom_indices=atom_indices,
                    ref_atom_indices=bad_atom_indices)

    assert np.all(dist2 > dist1)

    def one_to_many(self, prepared_traj1, prepared_traj2, index1, indices2):
        """Calculate a vector of distances from one frame of the first trajectory
        to many frames of the second trajectory

        The distances calculated are from the `index1`th frame of `prepared_traj1`
        to the frames in `prepared_traj2` with indices `indices2`

        Parameters
        ----------
        prepared_traj1 : rmsd.TheoData
            First prepared trajectory
        prepared_traj2 : rmsd.TheoData
            Second prepared trajectory
        index1 : int
            index in `prepared_trajectory`
        indices2 : ndarray
            list of indices in `prepared_traj2` to calculate the distances to

        Returns
        -------
        Vector of distances of length len(indices2)

        Notes
        -----
        If the omp_parallel optional argument is True, we use shared-memory
        parallelization in C to do this faster. Using omp_parallel = False is
        advised if indices2 is a short list and you are paralellizing your
        algorithm (say via mpi) at a different
        level.
        """
        return md.rmsd(prepared_traj1, prepared_traj2, index1, parallel=self.omp_parallel, precentered=True)[indices2]

    def one_to_all(self, prepared_traj1, prepared_traj2, index1):
        """Calculate a vector of distances from one frame of the first trajectory
        to all of the frames in the second trajectory

        The distances calculated are from the `index1`th frame of `prepared_traj1`
        to the frames in `prepared_traj2`

        Parameters
        ----------
        prepared_traj1 : rmsd.TheoData
            First prepared trajectory
        prepared_traj2 : rmsd.TheoData
            Second prepared trajectory
        index1 : int
            index in `prepared_trajectory`

        Returns
        -------
        Vector of distances of length len(prepared_traj2)

        Notes
        -----
        If the omp_parallel optional argument is True, we use shared-memory
        parallelization in C to do this faster.
        """
        return md.rmsd(prepared_traj2, prepared_traj1, index1, parallel=self.omp_parallel, precentered=True)

def main(opts):
    print 'Loading atom indices file for trajectories', opts.ndx
    ndx = np.loadtxt(opts.ndx, dtype=np.int)

    print 'Loading cells from', opts.cells
    cells = mdtraj.load(opts.topol, atom_indices=ndx)
    cells.xyz = load_cells_gps(opts.cells)

    print 'Loading trajectories', ' '.join(opts.trajs)
    traj = mdtraj.load(opts.trajs, top=opts.topol, atom_indices=ndx)

    print 'Assigning to {} cells'.format(len(cells))
    rmsds = -np.ones((len(cells), len(traj)))
    for i in xrange(len(cells)):
        rmsds[i] = mdtraj.rmsd(traj, cells, frame=i)
    rmsds = rmsds.T
    A = -np.ones((len(traj),), dtype=np.int)
    for f in xrange(len(traj)):
        A[f] = rmsds[f].argmin()

    np.savetxt(opts.assignments, A, fmt='%d')

    print 'Computing populations'
    P = np.bincount(A)
    np.savetxt(opts.populations, P, fmt='%d')

def plot_rmsd_distribution(cells, topol, atom_indices, bins=50):
    assert type(topol) is mdtraj.Trajectory, 'Expected Trajectory but got {}'.format(type(topotl))

    trajs = []
    for state in cells.L:
        t = copy.deepcopy(topol)
        t.xyz = state.x
        trajs.append(t)

    traj = trajs[0]
    traj = traj.join(trajs[1:])


    rmsds = []
    for frame in xrange(len(traj)):
        r = mdtraj.rmsd(traj, traj, frame=frame, atom_indices=atom_indices)
        rmsds.append(r)
    rmsds = np.vstack(rmsds)

    triu = np.triu_indices(len(rmsds))
    rmsds[triu] = -1
    np.fill_diagonal(rmsds, -1)
    rmsds = rmsds[np.where(rmsds >= 0)]

    plt.hist(rmsds, bins=bins)

def _deprecated_models_regular_spatial_clustering(templateids, traj, atom_indices=None, cutoff=0.06):
    """
    Superseded by models_regular_spatial_clustering
    """
    mdtraj_rmsd_args = {}
    if atom_indices:
        mdtraj_rmsd_args['atom_indices'] = atom_indices

    unique_templateids = []
    min_rmsd = []
    # Iterate through models
    for (t, templateid) in enumerate(templateids):
        # Add the first templateid to the list of uniques
        if t==0:
            unique_templateids.append(templateid)
            continue

        # Calculate rmsds of models up to t against the model t.
        rmsds = mdtraj.rmsd(traj[0:t], traj[t], parallel=False, **mdtraj_rmsd_args)
        min_rmsd.append(min(rmsds))

        # If any rmsd is less than cutoff, discard; otherwise add to list of uniques
        if min_rmsd[-1] < cutoff:
            continue
        else:
            unique_templateids.append(templateid)

    return unique_templateids

def test_lprmsd_5(get_fn):
    t = md.load(get_fn('frame0.h5'))
    t1 = md.load(get_fn('frame0.h5'))

    r = md.rmsd(t, t1, 0)
    a = md.lprmsd(t, t1, 0, permute_groups=[[]], superpose=True)
    eq(a, r, decimal=3)

def rmsd_connector(traj, inactive, residues_map = None):
	residues = [121, 282]
	if residues_map is not None:
		residues = map_residues(residues_map, residues)

	nonsymmetric = ["CG2", "CG1", "CD1", "CD2", "CE1", "CE2"]	
	connector_atoms = [(a.index, str(a)) for a in traj.topology.atoms if a.residue.resSeq in [121, 282] and "hydrogen" not in a.element and not any(substring in str(a) for substring in nonsymmetric)]
	
 	#print(connector_atom_names)
	#print connector_atoms
	connector_atoms = sorted(connector_atoms, key=operator.itemgetter(1), reverse = True)
	#print(connector_atoms)
	connector_atoms = [a[0] for a in connector_atoms]
	traj_stripped = traj.atom_slice(connector_atoms)


	connector_atoms_target = [(a.index,str(a)) for a in inactive.topology.atoms if a.residue.resSeq in [121, 282] and "hydrogen" not in a.element and not any(substring in str(a) for substring in nonsymmetric)]
	
	#connector_atom_names = [(a, a.element, a.index, a.residue) for a in inactive.topology.atoms if a.residue.resSeq in [121, 282] and "hydrogen" not in a.element]
	#print(connector_atom_names)
	#print connector_atoms_target
	connector_atoms_target = sorted(connector_atoms_target, key=operator.itemgetter(1), reverse = True)
	#print(connector_atoms_target)
	connector_atoms_target = [a[0] for a in connector_atoms_target]
	inactive_stripped = inactive.atom_slice(connector_atoms_target)

	traj_stripped_aligned = traj_stripped.superpose(inactive_stripped)
	rmsds = md.rmsd(traj_stripped, inactive_stripped) * 10.0
	return rmsds

def shukla_coords(trajectories,KER,Aloop,SRC2):

    difference = []
    rmsd = []

    for traj in trajectories:

        # append difference
        k295e310 = md.compute_contacts(traj, [KER[0]])
        e310r409 = md.compute_contacts(traj, [KER[1]])
        difference.append(10*(e310r409[0] - k295e310[0])) # 10x because mdtraj is naturally in nm

        # append rmsd
        Activation_Loop_SRC2 = SRC2.top.select("backbone and (resid %s to %s)" %(140,160))
        Activation_Loop_kinase = traj.top.select("backbone and (resid %s to %s)" %(Aloop[0],Aloop[1]))

        SRC2_cut = SRC2.atom_slice(Activation_Loop_SRC2)
        traj_cut = traj.atom_slice(Activation_Loop_kinase)

        rmsd.append(10*(md.rmsd(traj_cut,SRC2_cut,frame=0))) # 10x because mdtraj is naturaly in nm

    # flatten list of arrays
    flattened_difference = np.asarray([val for sublist in difference for val in sublist])
    flattened_rmsd = np.asarray([val for sublist in rmsd for val in sublist])

    return [flattened_rmsd, flattened_difference]

def calculate_rmsd(trajectory, topology, reference):
    import mdtraj
    traj = mdtraj.load(trajectory, top=topology)
    ref = mdtraj.load(reference)
    rmsd = mdtraj.rmsd(traj, ref)
    data = {"step": str(traj.n_frames), "rmsd": str(rmsd[-1])}
    return data