Untitled

import numpy as np
import mdtraj as md
from sklearn.decomposition import PCA

# right! Okay you're going to need to combine all your trajectories in a single xtc file or w/e MD engine you use
# these need to be aligned!! - basically whatever you fed into the PCA before with ALL frames from ALL simulations! 
traj = md.load('/Asal's/nicely/aligned/trajectory.xtc', top='/frame/from/sim.pdb')

# selecting CA atoms
ca_atoms = traj.topology.select('name CA')
ca_traj = traj.xyz[:, ca_atoms, :]
# reshape to stack xyz coordinates in a 'single column'
ca_traj_reshaped = ca_traj.reshape(ca_traj.shape[0], ca_traj.shape[1] * 3)

# PCA 
no_components = 3
pca = PCA(n_components=no_components)
pca_traj = pca.fit_transform(ca_traj_reshaped)

# this gives us the component vectors - that we can use to use to generate the PC motion
pc_vectors = pca.components_.reshape(no_components, ca_traj.shape[1], 3)


def save_pc_motion(pc_vector, savename, n_frames=20, scale_factor=10, reverse_to_loop=True):

    # this is the average structure from all simulation frames (super important to be aligned!)
    mean_structure = np.mean(ca_traj, axis=0)

    # We're going to linearly interpolate between the mean structure and apply the pc componet vector to the idententified PC movement
    frames = []
    ca_topology = traj.top.subset(ca_atoms)

    #forward motion frames
    for i in range(n_frames):
        displacement = scale_factor * ((i / (n_frames - 1)) - 0.5) * pc_vector
        interpolated_structure = mean_structure + displacement
        frames.append(interpolated_structure)

    # If we want to loop the motion - think Instagram boomerang haha - Otherwise (reverse_to_loop = False)
    if reverse_to_loop:
        #backward motion frames 
        for i in range(n_frames - 1, -1, -1):  
            frames.append(frames[i])

    # save the characteristic motion to a PDB file
    motion_traj = md.Trajectory(np.array(frames), ca_topology)
    motion_traj.save(f'{savename}_motion.pdb')


save_pc_motion(pc_vector=pc_vectors[0], savename='PC1', n_frames=20, scale_factor=20, reverse_to_loop=True)
save_pc_motion(pc_vector=pc_vectors[1], savename='PC2', n_frames=20, scale_factor=20, reverse_to_loop=False)

# This isnt nessarily going to create physically possible movement, its characteristic of the motion captured by the PC
# Mess around with the scale factor, increase it to get more exaggerated motion, decrease if it looks crazy
# more frames, more ... s m o o t h 

# If I was you, I'd use the tube representation when you load it into VMD - remember its only CA positons saved
# Message me if you need any help with getting nice figures or anything else! :)
# > Josh