from prody import *
from pylab import *
%matplotlib inline

"""
Explanation: Essential Dynamics Analysis
This example shows how to perform essential dynamics analysis of molecular dynamics (MD) trajectories. A EDA instance that stores covariance matrix and principal modes that describes the essential dynamics of the system observed in the simulation will be built. EDA and principal modes (Mode) can be used as input to functions in dynamics module for further analysis.
User needs to provide trajectory in DCD file format and PDB file of the system.
You can download the trajectory from the following link:
http://prody.csb.pitt.edu/tutorials/trajectory_analysis/trajectory_analysis_files.zip
Setting up environment
End of explanation
"""


structure = parsePDB('mdm2.pdb')

structure

"""
Explanation: Parse reference structure
The PDB file provided with this example contains and X-ray structure which will be useful in a number of places, so let’s start with parsing this file first:
End of explanation
"""


ensemble = parseDCD('mdm2.dcd')

ensemble.setCoords(structure)

ensemble.setAtoms(structure.calpha)

ensemble

ensemble.superpose()

eda_ensemble = EDA('MDM2 Ensemble')

eda_ensemble.buildCovariance( ensemble )

eda_ensemble.calcModes()

eda_ensemble

"""
Explanation: EDA calculations
Depending on the size of your trajectory, you suggested to use different classes. If you are analyzing a small trajectory, you can use an Ensemble instance obtained by parsing the trajectory at once using parseDCD():
End of explanation
"""


dcd = DCDFile('mdm2.dcd')

dcd.link(structure)

dcd.setAtoms(structure.calpha)

dcd

eda_trajectory = EDA('MDM2 Trajectory')

eda_trajectory.buildCovariance( dcd )

eda_trajectory.calcModes()

eda_trajectory

"""
Explanation: If you are analyzing a large trajectory, you can pass the trajectory instance to the PCA.buildCovariance() method as follows:
End of explanation
"""


printOverlapTable(eda_ensemble[:3], eda_trajectory[:3])

"""
Explanation: Comparison
End of explanation
"""


trajectory = Trajectory('mdm2.dcd')

trajectory.addFile('mdm2sim2.dcd')

trajectory

trajectory.link(structure)

trajectory.setCoords(structure)

trajectory.setAtoms(structure.calpha)

trajectory

eda = EDA('mdm2')

eda.buildCovariance( trajectory )

eda.calcModes()

eda

"""
Explanation: Multiple files
It is also possible to analyze multiple trajectory files without concatenating them. In this case we will use data from two independent simulations
End of explanation
"""


for mode in eda_trajectory[:4]:
    print(calcFractVariance(mode).round(2))

"""
Explanation: Analysis
Let’s print fraction of variance for top ranking 4 essential modes:
End of explanation
"""


mdm2ca_sim1 = trajectory[:500]

mdm2ca_sim1.superpose()

mdm2ca_sim2 = trajectory[500:]

mdm2ca_sim2.superpose()

showProjection(mdm2ca_sim1, eda[:3], color='red', marker='.');
showProjection(mdm2ca_sim2, eda[:3], color='blue', marker='.');
showProjection(mdm2ca_sim1[0], eda[:3], color='red', marker='o', ms=12);
showProjection(mdm2ca_sim2[0], eda[:3], color='blue', marker='o', ms=12);
showProjection(mdm2ca_sim1[-1], eda[:3], color='red', marker='s', ms=12);
showProjection(mdm2ca_sim2[-1], eda[:3], color='blue', marker='s', ms=12);

"""
Explanation: Plotting
Now, let’s project the trajectories onto top three essential modes:
End of explanation
"""


