#!/usr/bin/env python
# coding: utf-8

# <!--NOTEBOOK_HEADER-->
# *This notebook contains material from [PyRosetta](https://RosettaCommons.github.io/PyRosetta.notebooks);
# content is available [on Github](https://github.com/RosettaCommons/PyRosetta.notebooks.git).*

# <!--NAVIGATION-->
# < [Working With Symmetry](http://nbviewer.jupyter.org/github/RosettaCommons/PyRosetta.notebooks/blob/master/notebooks/10.00-Working-With-Symmetry.ipynb) | [Contents](toc.ipynb) | [Index](index.ipynb) | [Working With Antibodies](http://nbviewer.jupyter.org/github/RosettaCommons/PyRosetta.notebooks/blob/master/notebooks/12.00-Working-With-Antibodies.ipynb) ><p><a href="https://colab.research.google.com/github/RosettaCommons/PyRosetta.notebooks/blob/master/notebooks/11.00-Working-With-Density.ipynb"><img align="left" src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open in Colab" title="Open in Google Colaboratory"></a>

# # Working With Density
# Keywords: crystal, density, refinement, LoadDensityMapMover, map, EM, X-ray
# 
# ## Overview
# 
# Density can be a useful tool in Rosetta and can be used to refine a PDB structure, test Rosetta structure prediction methods, and build de-novo models guided by density of different resolutions. It can also be used as an experimentally-determined guide/constraint for custom methods. 
# 
# Rosetta understands X-ray or electron density through scoring.  In this tutorial, we will first walk through how to create a density file that Rosetta understands, and then load it, score it, and refine a structure using it. We will also cover some common tools used while working with density. 
# 
# More complicated protocols exist through the use of `RosettaScripts` in PyRosetta, and even more protocols through applications built using main Rosetta.  We will not cover these in this introductory tutorial, but please refer to the original references listed below (such as De novo protein structure determination...) for how to run these applications.     
# 
# Note that symmetry and density can be used together!
# 
# ## Documentation
# More information on Density Scoring and relavent applications can be found here:
# 
# - https://www.rosettacommons.org/docs/latest/application_documentation/analysis/density-map-scoring
# - https://www.rosettacommons.org/docs/latest/application_documentation/structure_prediction/mr-protocols
# 
# 
# ## References
# 
# - _Atomic-accuracy models from 4.5-Å cryo-electron microscopy data with density-guided iterative local refinement._
#  - https://www.ncbi.nlm.nih.gov/pubmed/25707030
#  
#  
# - _De novo protein structure determination from near-atomic-resolution cryo-EM maps._
#  - https://www.ncbi.nlm.nih.gov/pubmed/25707029
#  
#  
# - _Tools for Model Building and Optimization into Near-Atomic Resolution Electron Cryo-Microscopy Density Maps._
#  - https://www.ncbi.nlm.nih.gov/pubmed/27572730
#  
#  
# - _Automated structure refinement of macromolecular assemblies from cryo-EM maps using Rosetta._
#  - https://www.ncbi.nlm.nih.gov/pubmed/27669148
#  
#  
# - _Rosetta Structure Prediction as a Tool for Solving Difficult Molecular Replacement Problems._
#  - https://www.ncbi.nlm.nih.gov/pubmed/28573585
#  
#  
# - _RosettaES: a sampling strategy enabling automated interpretation of difficult cryo-EM maps._
#  - https://www.ncbi.nlm.nih.gov/pubmed/28628127
#  
#  
# - _Automatically Fixing Errors in Glycoprotein Structures with Rosetta_
#  - https://www.ncbi.nlm.nih.gov/pubmed/30344107
# 
# 

# In[13]:


get_ipython().system('pip install pyrosettacolabsetup')
import pyrosettacolabsetup; pyrosettacolabsetup.install_pyrosetta()
import pyrosetta; pyrosetta.init()


# ## Creating a Density Map
# 
# We won't actually do this in the tutorial as it requires a third-party application, but the (.ccp4) density map you will be using was created this way. 
# 
# 1) Download and install phenix.maps: https://www.phenix-online.org/documentation/reference/maps.html 
# 
#  - Make sure phenix.maps is in your path.
# 
# 2) Download the CIF file from the Protein Data Bank for the structure you are interested in.
# 
# 3) Use this command to create a map: `phenix.maps pdb_path cif_path` 
# 
#  - Note that depending on the cif file and structure, you may need to rarely edit the defaults to get it work properly, especially for glycan structures.  Please refer to the phenix documentation for this. 

# ## Setup for Density Scoring
# 
# Here, we will make a scorefunction and set `fast_elec_dens` to a proper value and use the `LoadDensityMapMover` to actually load the density map.  Note that (currently) the density map is GLOBAL data within Rosetta - so only a single structure can be modeled/refined at a time. 
# 
# ### From the docs:
# 
# Several scoring functions have been added to Rosetta which describe how well a structure agrees to experimental density data. Density map data is read in CCP4/MRC format (the density has to minimally cover the asymmetric unit).  The various scoring functions trade off speed versus accuracy, and their use should be primarily determined by the resolution of the density map data:
# 
# -   **elec\_dens\_fast** is recommended for most cases.
# 
# Additionally, a slower but more precise scoring function is available.  This is only recommended if **elec\_dens\_fast** performs poorly (for example, if map quality varies significantly throughout the map):
# -   **elec\_dens\_window** - Uses the sum of correlations of a sliding window of residues versus the experimental data; structure density only uses all heavy atoms.
# 
# The weights to use vary depending on resolution of data but the following give reasonable ranges:
# 
# - **elec\_dens\_fast**: ~25 is generally good, higher for high-resolution (<3Å) and lower for low-resolution (>6Å)
# - **elec\_dens\_window**: ~0.2 is generally good, higher for high-resolution (<3Å) and lower for low-resolution (>6Å)
# 
# 
# 

# In[14]:


from pyrosetta import *
from pyrosetta.rosetta import *
from pyrosetta.teaching import *
import os

init('-ignore_unrecognized_res -load_PDB_components false -ignore_zero_occupancy false @inputs/glycan_flags')


# In[15]:


from rosetta.protocols.cryst import *
from rosetta.protocols.rosetta_scripts import *

p = pose_from_pdb('inputs/1jnd.pdb')

original = p.clone()

#The LoadDensityMapMover unfortunately does not have getters and setters yet.  
# This has been updated in the Rosetta C++ code, but for now, we have to use the XML interface. 

#setup_dens = LoadDensityMapMover("inputs/1jnd_2mFo-DFc_map.ccp4")
setup_dens = XmlObjects.static_get_mover('<LoadDensityMap name="loaddens" mapfile="inputs/1jnd_2mFo-DFc_map.ccp4"/>')
setup_dens.apply(p)


# Now, we need to set the pose up for density scoring.  The `SetupForDensityScoring` mover sets a specific foldtree to the pose to allow scoring properly. We will then load a scorefunction with our density scoreterm, and load a pre-refined pose that was refined into the density using the pareto-optimal protocol with density. 

# In[16]:


setup_dens_pose = rosetta.protocols.electron_density.SetupForDensityScoringMover()

ref = pose_from_pdb('inputs/1jnd_refined.pdb.gz')

setup_dens_pose.apply(p)
setup_dens_pose.apply(ref)

score = get_score_function()
score_dens = get_score_function()
score_dens.set_weight(rosetta.core.scoring.elec_dens_fast, 25)

print("crystal", score_dens(p))
print("refined", score_dens(ref))


# Now lets minimize our pose. 

# In[17]:


minmover = MinMover()
mm = MoveMap()
mm.set_bb(True)
mm.set_chi(True)
minmover.set_movemap(mm)

if not os.getenv("DEBUG"):
    for i in range(1, 5):
        minmover.apply(p)
    
print(score_dens.score(p))


# Why has the score gotten worse here??  Because we are minimizing in dihedral space instead of cartesian space - so we make certain energies better, but crystal refinement works best in cartesian space.  Lets now minimize in cartesian and see what happens. 

# In[18]:


p = original.clone()
setup_dens_pose.apply(p)

score_dens_cart = create_score_function("ref2015_cart")
score_dens_cart.set_weight(rosetta.core.scoring.elec_dens_fast, 25)

#Set Bondlengths and angles to true. This is easier and more straightforward to do if using a MoveMapFactory.
mm.set(rosetta.core.id.THETA, True)

minmover.cartesian(True)
minmover.score_function(score_dens_cart)
minmover.set_movemap(mm)

if not os.getenv("DEBUG"):
    for i in range(1, 5):
        minmover.apply(p)

(print(score_dens.score(p)))


# Is this closer to the pre-refined model (which was pre-refined in Relax)?  Do they fit the density better?  Are they much more within the density?  How does the glycan density compare to the full protein density?

# ## Conclusion
# 
# That should get you started using density with typical Rosetta modeling tasks! See the references for more complex protocols. 

# In[ ]:





# **Chapter contributors:**
# 
# - Jared Adolf-Bryfogle (Scripps; Institute for Protein Innovation)

# <!--NAVIGATION-->
# < [Working With Symmetry](http://nbviewer.jupyter.org/github/RosettaCommons/PyRosetta.notebooks/blob/master/notebooks/10.00-Working-With-Symmetry.ipynb) | [Contents](toc.ipynb) | [Index](index.ipynb) | [Working With Antibodies](http://nbviewer.jupyter.org/github/RosettaCommons/PyRosetta.notebooks/blob/master/notebooks/12.00-Working-With-Antibodies.ipynb) ><p><a href="https://colab.research.google.com/github/RosettaCommons/PyRosetta.notebooks/blob/master/notebooks/11.00-Working-With-Density.ipynb"><img align="left" src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open in Colab" title="Open in Google Colaboratory"></a>