import prody as pd  # Note: if you're a Pandas user, it has the same conventional abbreviation "pd", so be careful

"""
Explanation: Lecture 21: Molecular Dynamics
CBIO (CSCI) 4835/6835: Introduction to Computational Biology
Overview and Objectives
In the last lecture, we introduced the idea of computational structural biology and the concept of molecular dynamics simulations to gauge how proteins could move and perform their functions over different timescales. In this lecture, we'll go over some tools you can use (in Python, of course) to look at proteins and analyze MD trajectories. By the end of this lecture, you should be able to:

Download and explore PDB files of proteins using ProDy or PyMol
Understand the basics of performing PCA on covariance matrices of ensembles or trajectories
Use PyMol to visualize protein structures and interface with Python scripts

Part 1: ProDy
<img src="Lecture21/prody-logo.png" />
ProDy stands for Protein Dynamics,

an API that is very suitable for interactive usage 
comes with several command line applications

ProDy is designed for normal mode analysis, but also is

a powerful tool for handling macromolecular structures
useful for analysis of molecular dynamics (MD) trajectories
useful for sequence conservation and coevolution analysis (Evol)

What is "normal mode analysis"?
From Bahar et al 2009, Normal Mode Analysis of Biomolecular Structures: Functional Mechanisms of Membrane Proteins (section 1.1.3):
Normal mode analysis provides information on the equilibrium modes accessible to a system, assuming that the system is stabilized by harmonic potentials. It has been used for several decades in studying classical physical phenomena such as atomic vibrational spectra and transport in the solid state. Its application to proteins dates back to the early 1980s. However, only in the past decade has it become a tool widely used for exploring functional motions. A major reason behind its broader use is the observation that global modes elucidated by NMA bear functional significance. This feature became even more evident with the use of simplified models in coarse-grained (CG) NMA.

One way to look at it: what are the energetically-favorable configurations of a macromolecule?
These folded configurations of the protein(s) have functional significance, so it's very important to understand

what the folded configurations are, and
how the protein reaches those configurations

ProDy basics
It's easy enough to get started: just import the base package.
End of explanation
"""


ubi = pd.parsePDB("1ubi")

"""
Explanation: Most ProDy functions follow a specific naming convention:

an action verb, followed by
some kind of three-letter abbreviation of an object

For example, a function that "does something" would be named in ProDy as doSTH
Just a few examples directly from the ProDy package:

parseEXT(): parse a file in EXT format, e.g. parsePDB, parseDCD
writeEXT(): write a file in EXT format, e.g. writePDB, writeDCD
fetchSTH(): download a file, e.g. fetchPDB, fetchMSA
calcSTH(): calculate something, e.g. calcRMSD, calcGyradius, calcANM
showSTH(): show a plotting of something, e.g. showCrossCorrelations, showProtein
saveSTH(): save a ProDy object instance to disk, e.g. saveAtoms
loadSTH(): save a ProDy object instance to disk, e.g. loadAtoms

We'll touch on a few of these.
Let's dive in with an example, shall we?
One of the coolest things about ProDy--when you request the structure information for a specific macromolecule, it will download that structure directly from the Protein Data Bank (remember PDB?):
End of explanation
"""


print(ubi)

print(ubi.numAtoms())

print(pd.calcGyradius(ubi))  # This function calculates the radius of gyration of the atoms

"""
Explanation: How cool is that?!
PDB
Recall from our last lecture: the Protein Data Bank is a website ( www.rcsb.org/pdb/home/home.do ), but the acronym PDB is the file format used by the website (and pretty much any researcher interested in protein structure) to describe the 3D structures of macromolecules.
<img src="Lecture21/pdb-logo.png" />
Each macromolecule on the PDB is given a 4-digit descriptor; usually the first digit is a number, and the next three are letters related to the name of the protein.
In our code example, we used 1ubi.
<img src="Lecture21/1ubi-landing.png" width="80%" />
Note the wealth of information presented about this chromosomal protein--this even excludes all the literature hits below that cite the use of this macromolecule.
Also note the "Download Files" link in the upper right--this is where you can get the PDB files if you don't already have them.
On the other hand, if you're using ProDy, it will just download them for you!
Back to ProDy! Since it went ahead and downloaded the PDB file for 1ubi for us, let's take a look at what we have.
End of explanation
"""


pd.saveAtoms(ubi)

"""
Explanation: File Handling
If the internet isn't your thing, ProDy has its own formats for interacting with files on your hard drive.
End of explanation
"""


pd.writePDB("ubi.pdb", ubi)  # Save to the file "ubi.pdb"

ubi2 = pd.parsePDB("ubi.pdb")  # Now read from it, just to test that it worked!
print(ubi2)

"""
Explanation: You can also save to and load from more standard file formats, like PDB:
End of explanation
"""


type(ubi)

"""
Explanation: The key point: if the argument you specify to parsePDB doesn't exist on your computer, then it'll connect to PDB directly to try and download it.
Note the type of the variable that comes back from a call to parsePDB:
End of explanation
"""


print(ubi.numAtoms())

ag = pd.parsePDB('1vrt')
print(ag.numAtoms())

"""
Explanation: Why AtomGroup?

not Molecule, because structures are usually made up from multiple molecules
not Structure, because PDB format is sometimes used for storing small-molecules

AtomGroup made sense for handling bunch of atoms, and is used by some other packages too.
Some AtomGroup methods
As we've seen, we can check on how many atoms there are:
End of explanation
"""


names = ag.getNames()
print(names)

"""
Explanation: We can also ask for specific properties of the macromolecule.
End of explanation
"""


len(names)

type(names)

"""
Explanation: What do you think these are?
End of explanation
"""


coords = ag.getCoords()
print(coords)

print(coords.shape)

"""
Explanation: Oh hey, we recognize that!
We could ask for more detail on the macromolecule, such as its location in space:
End of explanation
"""


a0 = ag[0]
print(a0)

print(a0.getName())

"""
Explanation: (Would you be able to compute the generalized coordinates of this macromolecule?)
Atom instances
You can get the names of all the atoms in the macromolecule via the getNames method, but you can also index the macromolecule directly as though it were an array:
End of explanation
"""


every_other_atom = ag[::2]
print(every_other_atom)

"""
Explanation: Taking that same thinking further, we can even slice out subgroups of atoms from the macromolecule:
End of explanation
"""


print(len(every_other_atom))
print(len(ag))

"""
Explanation: The type is a Selection object, but we can see that we get what we'd expect:
End of explanation
"""


print(ag.numChains())

print(ag.numResidues())

"""
Explanation: ProDy Hierarchy
Atoms, structures, residues... all terms we understand from a biological perspective, but how do they play into ProDy?
ProDy arranges these concepts into a hierarchy within a macromolecule. The hierarchy looks something like this:

Atom: lowest level of the hierarchy
Residue: an amino acid, nucleotide, small molecule, or ion
Chain: a polypeptide or nucleic acid chain
Segment: used by simulation programs and comprise multiple chains
End of explanation
"""


# Printing out each chain and the number of residues each has.
for chain in ag.iterChains():
    print(chain, chain.numResidues())

# Here, we'll print out each chain and their first 10 residues.
for chain in ag.iterChains():
    print(chain)
    residues = 0
    for residue in chain:  # We can also loop through residues on a chain!
        print(' | - ', residue)
        residues += 1
        if residues >= 10: break
    print("...")

"""
Explanation: We can set up a loop to iterate through the chains, using the iterChains method:
End of explanation
"""


# Select all the alpha Carbon atoms
sel = ag.select("protein and name CA")
print(sel)
print(sel.numAtoms())

# Shorthand 
sel2 = ag.select("calpha")
sel3 = ag.select("ca")
sel2 == sel3

"""
Explanation: Other methods for looping over structures in a macromolecule:

iterAtoms
iterBonds
iterCoordsets
iterFragments
iterSegments

Two time-saving asides:

In Jupyter, you can TAB-complete with partial function names to see the list of all the functions available to you.
Also in Jupyter, you can type out a function name, but at the end, put a question mark ? and hit ENTER. This will bring up the documentation for how to use that function.

Selection Grammar
This is a very complicated, but very powerful, interface to searching for specific properties of your molecule. We won't spend a lot of time here, but this basically allows you to search for specific atoms, residues, or chains using plain English:
End of explanation
"""


import numpy as np
origin = np.zeros(3)

sel = ag.select("within 5 of origin", origin = origin)
print(sel)
print(pd.calcDistance(sel, origin))

sel = ag.select("within 5 of center", center = pd.calcCenter(ag))
print(sel)

# You can even use dot-selection shorthand, instead of the "select" method!
ag.calpha

ag.name_CA_and_resname_ALA

"""
Explanation: You can also select atoms or residues by proximity:
End of explanation
"""


p38 = pd.parsePDB("5uoj")
bound = pd.parsePDB("1zz2")

matches = pd.matchChains(p38, bound)
print(len(matches))
p38_ch, bnd_ch, seqid, seqov = matches[0]

print(bnd_ch)
print(seqid)
print(seqov)

"""
Explanation: See the full documentation on selection grammar here: http://csb.pitt.edu/ProDy/reference/atomic/select.html
Chain Matching and RMSD
ProDy will even try to find matching portions of chains in macromolecules, and align one macromolecule to another.
End of explanation
"""


print(pd.calcRMSD(p38_ch, bnd_ch))

"""
Explanation: matchChains takes 2 arguments: two AtomGroup objects to compare.
It returns however many matches it finds (in our example, only 1).
Each returned match contains 4 values:

the matching chain from the first argument
the matching chain from the second argument
percent identity of the match
percent sequence overlap of the match

We can then use the matching chains from the two proteins to perform an alignment: finding a pose of one of the chains with respect to the other one.
Right now, even though these chains match, they don't align very well:
End of explanation
"""


bnd_ch, transformation = pd.superpose(bnd_ch, p38_ch)
print(pd.calcRMSD(bnd_ch, p38_ch))

"""
Explanation: Recall our discussion of RMSD (Root Mean Squared Deviation) from the previous lecture--it's basically the Euclidean distance between corresponding points in 3D space.
We have an alignment of these chains from the matchChains function, but they differ considerably in terms of their physical, spatial poses. Here, we want to ask if their structure--while similar, not identical--allows them to overlap even in space.
We can use the function superpose to create a superposition of these chains.
 - The first argument is considered the mobile chain
 - The second argument is considered the target, or fixed, chain
End of explanation
"""


ubi = pd.parsePDB('2lz3')

"""
Explanation: Much better!
Dynamics Analysis
ProDy can even perform some analysis of molecular dynamics.
PDB has some more complicated macromolecules that include several conformers of the same protein, called an ensemble. This is basically a fancy term for "set of molecules that are the same but in different spatial poses"--as in, what you'd get from the output of an MD simulation.
End of explanation
"""


ubi_ensemble = pd.Ensemble(ubi.calpha)  # Why calpha?
ubi_ensemble

"""
Explanation: WARNING: There is a bug in ProDy. Running this command on a PDB structure with more than one component (i.e., an ensemble of macromolecule models in a single file) will result in an error unless you use the version directly available through GitHub (or fix it yourself).
There are 21 conformers in this single file (PDB predicts there are 200 total for this molecule!). ProDy will recognize the ensemble nature, though, if we run the following:
End of explanation
"""


print(ubi_ensemble.getRMSDs())

"""
Explanation: The next step is to minimize the differences between each conformer.
End of explanation
"""


ubi_ensemble.iterpose()  # This performs an iterative alignment.

print(ubi_ensemble.getRMSDs())  # Did that improve things any?

"""
Explanation: Initially, the conformers aren't aligned; their respective RMSDs to some reference (by default, the first one; hence its RMSD is 0) is decently high.
We can fix this with the interpose function:
End of explanation
"""


pca = pd.PCA()

pca.buildCovariance(ubi_ensemble)
cov = pca.getCovariance()
print(cov.shape)

%matplotlib inline
import seaborn as sns
import matplotlib.pyplot as plt

plt.imshow(cov, cmap = 'Blues')
plt.colorbar()

"""
Explanation: Once we've aligned the conformers, we can do some analysis. Remember PCA?
End of explanation
"""


pca.calcModes()  # Performs the actual PCA analysis.

"""
Explanation: Initial results don't tell us a whole lot, except that there seem to be some parts of the ensemble that are positively correlated, and some negatively correlated. Let's dig in a bit more.
End of explanation
"""


plt.plot(pca.getEigvals())

"""
Explanation: Now that we've calculated the modes--or principal components!--we can see everything that we discussed in the previous lecture.
For starters, let's look at the eigenvalues.
End of explanation
"""


for mode in pca:
    print(pd.calcFractVariance(mode).round(2))

"""
Explanation: Remember what we said about how PCA works: it rotates the data such that each dimension / axis of the data contains a certain amount of the original variance.
PCA then returns to us the new axes of the data (as the eigenvectors), and the relative contributions (or importance) of each axis to the data (as the eigenvalues).
Therefore, rather than retain all the data, we can just keep the top handful of dimensions, as specified by the eigenvalues--which quantify precisely how much each dimension "counts".
End of explanation
"""


pd.showSqFlucts(pca)

"""
Explanation: The first dimension alone contains over half the variance in the original data.
Include the second dimension, and you already have nearly 80% of the original variance.
With just the first four dimensions, you have nearly 95% of the original variance with but a tiny fraction of the original quantity of data.

Put another way, these first four modes explain 95% of the motion observed in the molecular ensemble. Unless you're interested in really super-high frequency movement, this will probably be sufficient for future analyses.
We can also use the eigenvectors that were computed to look at the fluctuations in each atom.
End of explanation
"""


pd.showProjection(ubi_ensemble, pca[:3])  # The first 3 principal components

"""
Explanation: Of the nearly-60 atoms in the chains, the atoms around certain indices move quite a bit more than the others.
ProDy even comes with a neat projection visualization to view RMSDs for each mode.
End of explanation
"""


