%matplotlib inline
import matplotlib.pyplot as plt
import mdtraj as md
traj = md.load("5550217/kras.xtc", top="5550217/kras.pdb")
topology = traj.topology

from contact_map import ContactFrequency, ContactDifference

"""
Explanation: Overview
Contact Map Explorer is a tool for studying contacts in biomolecules. It builds on tools in MDTraj. These notebooks teach how to use Contact Map Explorer by showing examples.
As our initial example system, we'll use part of a trajectory of the KRas protein bound to GTP, which was provided by Sander Roet. KRas is a protein that plays a role in many cancers. For simplicity, the waters were removed from the trajectory (although ions are still included). To run this notebook, download the example files from https://figshare.com/s/453b1b215cf2f9270769 (total download size about 1.2 MB). Download all files, and extract in the same directory that you started Jupyter from (so that you have a directory called 5550217 in your current working directory).
End of explanation
"""


%%time
frame_contacts = ContactFrequency(traj[0])

"""
Explanation: Contact map for a single frame
First we make the contact map for the 0th frame. For default parameters (and how to change them) see the notebook called "Changing defaults."
End of explanation
"""


%%time
fig, ax = frame_contacts.residue_contacts.plot()
plt.xlabel("Residue")
_ = plt.ylabel("Residue")
# setting the return of plt.ylabel to the dummy var _ is a way to swallow output

"""
Explanation: Plotting requires matplotlib, and the default plotting parameters use the 'seismic' colormap. For details on customizing the appearance of the contact map plots, see the notebook on "Customizing contact map plots."
End of explanation
"""


%%time
traj_contacts = ContactFrequency(traj)

%%time
fig, ax = traj_contacts.residue_contacts.plot()
plt.xlabel("Residue")
_ = plt.ylabel("Residue")

"""
Explanation: Contact map for a trajectory
If you provide a single frame, the values in the contact map are 0 or 1: either the contact was present or it was not. But if you provide a longer trajectory, the values in the contact map are the fraction of frames during which the contact was present.
Other than using the entire trajectory, the process is exactly the same as for the single-frame example.
End of explanation
"""


fig.savefig("traj_contacts.pdf")

"""
Explanation: Note that some of the reds are a little lighter now -- these contacts were not present 100% of the time!
If we want to save the figure, we can use matplotlib's savefig functionality:
End of explanation
"""


traj_contacts.save_to_file("traj_contacts.p")

# to reload it:
ContactFrequency.from_file("traj_contacts.p")

"""
Explanation: If we want to save the contact data to a file for later usage, we can also do that:
End of explanation
"""


%%time
diff = traj_contacts - frame_contacts
# equivalently:  diff = ContactDifference(traj_contacts, frame_contacts)

%%time
fig, ax = diff.residue_contacts.plot()

"""
Explanation: For many users, this is all that Contact Map Explorer is needed for. You can make publication-quality figures in only a few lines of Python. However, Contact Map Explorer is also intended to help you analyze your data and help you understand it better. The rest of this notebook, and other notebooks in the documentation, will show you some of the possibilities.
Differences between contact maps
One question you might have is, how do two contact maps differ? For this purpose, we can calculate contact map differences.
End of explanation
"""


# show the 10 with the highest value
diff.residue_contacts.most_common()[:10]

"""
Explanation: The contact map gives a lot of information, and frequently we want to ask follow-up questions. We'll looks at the results from the difference above, but all the same techniques apply to any contact results.
For example, after looking at that difference, we might want to know which were much more likely to be in the full trajectory than in the single frame:
End of explanation
"""


list(reversed(diff.residue_contacts.most_common()))[:10]

"""
Explanation: If you wanted to see the least common (that is, the ones from the single frame that were least likely to be in the rest of the trajectory), you can use standard Python tricks for reversing a list:
End of explanation
"""


gtp201 = topology.residue(166)
print(gtp201)  # it's good to print this to double-check you got the right one!

"""
Explanation: What if you wanted to know which residues were most commonly in contact with a specific residue, say the GTP in KRAS? First, we need to get the Residue object for the GTP from the MDTraj topology:
End of explanation
"""


for contact in traj_contacts.residue_contacts.most_common(gtp201):
    if contact[1] > 0.1:
        print(contact)

"""
Explanation: We extended the standard .most_common() to take an optional argument. When the argument is given, it will filter the output to only include the ones where that argument is part of the contact. So to get the most common residue-residue contacts involving GTP201:
End of explanation
"""


traj_contacts.most_common_atoms_for_residue(gtp201)[:10]

"""
Explanation: Finding atoms associated with residue contacts
One of the particularly useful features of Contact Map Explorer is that it preserves information about the atom-atom contacts, as well as the residue-residue contacts.
Above, we looked at which residues were more commonly in contact with GTP201. But what if we wanted to know which atom-atom contacts involved GTP201? We can do that with contacts.most_common_atoms_for_residue.
End of explanation
"""


val81 = topology.residue(80)
asn116 = topology.residue(115)
print(val81, asn116)

traj_contacts.most_common_atoms_for_contact([val81, asn116])

"""
Explanation: On the other hand, what if you wanted to know the most common atoms for a specific contact pair? Let's do that for the contact pair Val81-Asn116. First we'll get the residue objects, then we'll use the most_common_atoms_for_contact method.
End of explanation
"""


%%time
traj_contacts.atom_contacts.plot();

"""
Explanation: Atom-atom contact maps
In addition to residue-residue contacts, Contact Map Explorer also keeps information about atom-atom contacts. Everything done above with contacts.residue_contacts can also be done with contacts.atom_contacts. 
Plotting atom-atom contacts
We can plot atom-atom contacts, but because the matrix is so large, there are often problems displaying it because there aren't enough pixels on the screen.
End of explanation
"""


traj_contacts.atom_contacts.most_common()[:10]

asn116cb = topology.atom(1809)
print(asn116cb)

traj_contacts.atom_contacts.most_common(asn116cb)

"""
Explanation: You'll notice that you don't see many points here. As the warning states: That is because the points are smaller than a single pixel at this resolution. To fix that, output or save this figure as a vector image (like PDF), increase the figure's size (by using the figsize parameter) or resolution (by using the dpi parameter), or both. 
To get Jupyter to output PDF by default instead of PNG:
python
from IPython.display import set_matplotlib_formats
set_matplotlib_formats('pdf')
To increase the figure size:
python
traj_contacts.atom_contacts.plot(figsize=(38, 38))
To increase DPI (and make the figure square):
python
traj_contacts.atom_contacts.plot(figsize=(6, 6), dpi=454)
Using most_common
The most_common method works the same on atom contacts as on residue contacts.
End of explanation
"""


