Python mdtraj.compute_dihedrals函数代码示例

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('assignments', default='Macro4/MacroAssignments.h5', help='Path to an assignments file. (default=Macro4/MacroAssignments.h5)')
    parser.add_argument('--project', default='ProjectInfo.yaml', help='Path to ProjectInfo.yaml file. (default=ProjectInfo.yaml)')
    args = parser.parse_args()

    project = Project.load_from(args.project)
    t = reduce(operator.add, (project.load_traj(i) for i in range(project.n_trajs)))

    phi_angles = md.compute_dihedrals(t, [PHI_INDICES]) * 180.0 / np.pi
    psi_angles = md.compute_dihedrals(t, [PSI_INDICES]) * 180.0 / np.pi
    state_index = np.hstack(io.loadh(args.assignments)['arr_0'])

    for i in np.unique(state_index):
        pp.plot(phi_angles[np.where(state_index == i)],
                psi_angles[np.where(state_index == i)],
                'x', label='State %d' % i)


    pp.title("Alanine Dipeptide Macrostates")
    pp.xlabel(r"$\phi$")
    pp.ylabel(r"$\psi$")
    annotate()

    pp.legend(loc=1, labelspacing=0.075, prop={'size': 8.0}, scatterpoints=1,
              markerscale=0.5, numpoints=1)
    pp.xlim([-180, 180])
    pp.ylim([-180, 180])
    pp.show()

def test_generator():
    N_FRAMES = 2
    N_ATOMS = 5
    xyz = np.asarray(np.random.randn(N_FRAMES, N_ATOMS, 3), dtype=np.float32)
    ptraj = md.Trajectory(xyz=xyz, topology=None)

    quartets = np.array(list(itertools.combinations(range(N_ATOMS), 4)), dtype=np.int32)
    quartets2 = itertools.combinations(range(N_ATOMS), 4)
    a = md.compute_dihedrals(ptraj, quartets)
    b = md.compute_dihedrals(ptraj, quartets2)
    eq(a, b)

    def partial_transform(self, traj):
        """Featurize an MD trajectory into a vector space via calculation
        of dihedral (torsion) angles of alpha carbon backbone

        Parameters
        ----------
        traj : mdtraj.Trajectory
            A molecular dynamics trajectory to featurize.

        Returns
        -------
        features : np.ndarray, dtype=float, shape=(n_samples, n_features)
            A featurized trajectory is a 2D array of shape
            `(length_of_trajectory x n_features)` where each `features[i]`
            vector is computed by applying the featurization function
            to the `i`th snapshot of the input trajectory.

        """

        ca = [a.index for a in traj.top.atoms if a.name == 'CA']
        if len(ca) < 4:
            return np.zeros((len(traj), 0), dtype=np.float32)

        alpha_indices = np.array(
            [(ca[i - 1], ca[i], ca[i+1], ca[i + 2]) for i in range(1, len(ca) - 2)])
        result = md.compute_dihedrals(traj, alpha_indices)

        x = []
        if self.atom_indices is None:
            self.atom_indices = np.vstack(alpha_indices)
        if self.sincos:
            x.extend([np.cos(result), np.sin(result)])
        else:
            x.append(result)
        return np.hstack(x)

def test_no_indices():
    for fn in ['2EQQ.pdb', '1bpi.pdb']:
        for opt in [True, False]:
            t = md.load(get_fn(fn))
            assert md.compute_distances(t, np.zeros((0,2), dtype=int), opt=opt).shape == (t.n_frames, 0)
            assert md.compute_angles(t, np.zeros((0,3), dtype=int), opt=opt).shape == (t.n_frames, 0)
            assert md.compute_dihedrals(t, np.zeros((0,4), dtype=int), opt=opt).shape == (t.n_frames, 0)

def test_dihedral_1():
    pymol = find_executable('pymol')
    if pymol is None:
        raise SkipTest("pymol executable not found")
    
    xyz = '''MODEL        0
ATOM      1    A ACE     1       4.300  13.100   8.600  1.00  0.00
ATOM      2    B ACE     1       5.200  13.600   8.800  1.00  0.00
ATOM      3    C ACE     1       4.900  14.300   9.600  1.00  0.00
ATOM      4    D ACE     1       5.600  14.200   7.900  1.00  0.00
    '''
    script = '''
with open('output.txt', 'w') as f:
    f.write('%f' % cmd.get_dihedral('1/A', '1/B', '1/C', '1/D'))
'''
    
    with enter_temp_directory():
        with open('xyz.pdb', 'w') as f:
            f.write(xyz)
        with open('pymolscript.py', 'w') as f:
            f.write(script)

        os.system('%s %s -cr %s' % (pymol, 'xyz.pdb', 'pymolscript.py'))
        with open('output.txt') as f:
            pymol_value = np.deg2rad(float(f.read()))
        t = md.load('xyz.pdb')

    mdtraj_value = md.compute_dihedrals(t, [[0,1,2,3]])[0,0]
    
    np.testing.assert_array_almost_equal(pymol_value, mdtraj_value)

 def get_phipsi(self, trajs, phi, psi):
     #phi = [6, 8, 14, 16]
     #psi = [4, 6, 8, 14]
     PHI_INDICES = []
     PSI_INDICES = []
     for i in xrange(len(phi)):
         PHI_INDICES.append(self.atom_indices.index(phi[i]))
         PSI_INDICES.append(self.atom_indices.index(psi[i]))
     len_trajs = len(trajs)
     print "PSI:", PSI_INDICES
     print "PHI:", PHI_INDICES
     phi_angles = md.compute_dihedrals(trajs, [PHI_INDICES]) * 180.0 / np.pi
     psi_angles = md.compute_dihedrals(trajs, [PSI_INDICES]) * 180.0 / np.pi
     #phi_psi=np.column_stack((phi_angles, psi_angles))
     #return phi_psi
     return phi_angles, psi_angles

 def chi1_feat(traj, res):
     chi1 = traj.topology.select('resid %i and (name C or name CA or name CB or name CG or name SG or name CG1 or name OG or name OG1)' %res)
     if chi1.shape[0] != 4:
         return None
     chi1 = chi1.reshape([1,4])
     traj_chi1 = md.compute_dihedrals(traj, chi1)
     return traj_chi1

def sidechain_example(yaml_file):
    # Parse a YAML configuration, return as Dict
    cfg = Settings(yaml_file).asDict()
    structure = cfg['Structure']

    #Select move type
    sidechain = SideChainMove(structure, [1])
    #Iniitialize object that selects movestep
    sidechain_mover = MoveEngine(sidechain)

    #Generate the openmm.Systems outside SimulationFactory to allow modifications
    systems = SystemFactory(structure, sidechain.atom_indices, cfg['system'])

    #Generate the OpenMM Simulations
    simulations = SimulationFactory(systems, sidechain_mover, cfg['simulation'], cfg['md_reporters'],
                                    cfg['ncmc_reporters'])

    # Run BLUES Simulation
    blues = BLUESSimulation(simulations, cfg['simulation'])
    blues.run()

    #Analysis
    import mdtraj as md
    import numpy as np

    traj = md.load_netcdf('vacDivaline-test/vacDivaline.nc', top='tests/data/vacDivaline.prmtop')
    indicies = np.array([[0, 4, 6, 8]])
    dihedraldata = md.compute_dihedrals(traj, indicies)
    with open("vacDivaline-test/dihedrals.txt", 'w') as output:
        for value in dihedraldata:
            output.write("%s\n" % str(value)[1:-1])

def test_dihedral_1(pymol, tmpdir):
    xyz = '''MODEL        0
ATOM      1    A ACE     1       4.300  13.100   8.600  1.00  0.00
ATOM      2    B ACE     1       5.200  13.600   8.800  1.00  0.00
ATOM      3    C ACE     1       4.900  14.300   9.600  1.00  0.00
ATOM      4    D ACE     1       5.600  14.200   7.900  1.00  0.00
    '''
    script = '''
from pymol import cmd
with open('output.txt', 'w') as f:
    f.write('%f' % cmd.get_dihedral('1/A', '1/B', '1/C', '1/D'))
'''
    prevdir = os.path.abspath('.')
    try:
        os.chdir(tmpdir)
        with open('xyz.pdb', 'w') as f:
            f.write(xyz)
        with open('pymolscript.py', 'w') as f:
            f.write(script)

        os.system('%s %s -cr %s' % (pymol, 'xyz.pdb', 'pymolscript.py'))
        with open('output.txt') as f:
            pymol_value = np.deg2rad(float(f.read()))
        t = md.load('xyz.pdb')
    finally:
        os.chdir(prevdir)

    mdtraj_value = md.compute_dihedrals(t, [[0, 1, 2, 3]])[0, 0]

    np.testing.assert_array_almost_equal(pymol_value, mdtraj_value)

    def calc_dihedral_energy(self, traj, improper=False, sum=True):
        """Energy for dihedral interactions

        Parameters
        ----------
        traj : mdtraj.Trajectory

        sum : bool (opt.)
            If sum=True return the total energy.
        """
        phi = md.compute_dihedrals(traj, self._dihedral_idxs)
        #if improper:
        #    phi = np.pi + md.compute_dihedrals(traj, self._dihedral_idxs) # ?
        #else:
        #    phi = -temp_phi.copy()
        #    phi[temp_phi > 0] = 2.*np.pi - temp_phi[temp_phi > 0]

        if sum:
            Edihedral = np.zeros(traj.n_frames, float)
        else:
            Edihedral = np.zeros((traj.n_frames, self.n_dihedrals), float)

        for i in range(self.n_dihedrals):
            if sum:
                Edihedral += self._dihedrals[i].V(phi[:,i])
            else:
                Edihedral[:,i] = self._dihedrals[i].V(phi[:,i])
        return Edihedral

def DFG_dihedral_states(trajectories,def_DFG):

    dihedral = []

    for traj in trajectories:
        dihedral.append(md.compute_dihedrals(traj,[def_DFG]))

    return [dihedral]

 def load_data(self):
     load_time_start = time.time()
     data = []
     for tfn in self.filenames:
         kwargs = {} if tfn.endswith('h5') else {'top': self.top}
         for t in md.iterload(tfn, chunk=self.args.split, **kwargs):
             item = np.asarray(md.compute_dihedrals(t, self.indices), np.double)
             data.append(item)
     return data

def calculate_dihedrals():
    print("Calculating dihedrals...")
    traj_files = sorted(glob.glob("traj*xtc"))
    traj = [ md.load(filename, top='structure.gro') for filename in traj_files ]
    indices = list(traj[0].topology.select('backbone'))
    dihedral_quartets = np.array([indices[i:i+4] for i in range(len(indices)-4)])
    for i in range(len(traj)):
        thetas = md.compute_dihedrals(traj[i], dihedral_quartets)
        cos_sin_dihedrals = np.hstack([np.cos(thetas), np.sin(thetas)])
        np.save('out_' + str(i) + '.npy', thetas)

    def transform(self, traj):

        # compute dihedral energy
        phi = md.compute_dihedrals(traj, self.dihedrals)
        Edih = np.array(map(lambda x,y: x(y), self.Vdih, phi.T)).T

        # compute pair energy
        r = md.compute_distances(traj, self.pairs)
        Epair = np.array(map(lambda x,y: x(y), self.Vpair, r.T)).T

        return np.hstack((Edih, Epair))

def DFG_dihedral(trajectories, def_DFG):

    dihedral = []

    for traj in trajectories:

        dihedral.append(md.compute_dihedrals(traj, [def_DFG]))

    flattened_dihedral = np.asarray([val for sublist in dihedral for val in sublist])

    return [flattened_dihedral]