#!/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-->
# < [Structure Refinement](http://nbviewer.jupyter.org/github/RosettaCommons/PyRosetta.notebooks/blob/master/notebooks/05.00-Structure-Refinement.ipynb) | [Contents](toc.ipynb) | [Index](index.ipynb) | [Refinement Protocol](http://nbviewer.jupyter.org/github/RosettaCommons/PyRosetta.notebooks/blob/master/notebooks/05.02-Refinement-Protocol.ipynb) ><p><a href="https://colab.research.google.com/github/RosettaCommons/PyRosetta.notebooks/blob/master/notebooks/05.01-High-Res-Movers.ipynb"><img align="left" src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open in Colab" title="Open in Google Colaboratory"></a>

# # High-Resolution Movers
# Keywords: keep_history(), MoveMap, SmallMover(), ShearMover(), angle_max(), set_bb(), MinMover(), MonteCarlo(), boltzmann(), TrialMover(), SequenceMover(), RepeatMover()

# In[ ]:


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


# **Make sure you are in the directory with the pdb files:**
# 
# `cd google_drive/MyDrive/student-notebooks/`

# In the last workshop, you encountered the `ClassicFragmentMover`, which inserts a short sequence of backbone torsion angles, and the `SwitchResidueTypeSetMover`, which doesn’t actually change the conformation of the pose but instead swaps out the residue types used.
# 
# 
# In this workshop, we will introduce a variety of other `Movers`, particularly those used in high-resolution refinement (e.g., in Bradley’s 2005 paper).
# 
# Before you start, load the cleaned cetuximab protein 1YY8 that we've worked with previously (but just one Fab fragment, chain A and B) and make a copy of the pose so we can compare later:
# 
# ```
# start = pose_from_pdb("1YY8.clean.pdb")
# test = Pose()
# test.assign(start)
# ```

# In[42]:


### BEGIN SOLUTION
start = pose_from_pdb("inputs/1YY8.clean.pdb")
test = Pose()
test.assign(start)
### END SOLUTION


# OPTIONAL: For convenient viewing in PyMOL, set the names of both poses:
# ```
# start.pdb_info().name("start")
# test.pdb_info().name("test")
# pmm = PyMOLMover()
# ```

# In[43]:


### BEGIN SOLUTION
start.pdb_info().name("start")
test.pdb_info().name("test")
pmm = PyMOLMover()
### END SOLUTION


# We also want to activate the `keep_history` setting so that PyMOL will keep separate frames for each conformation as we modify it (more on this shortly):
# ```
# pmm.keep_history(True)
# pmm.apply(start)
# pmm.apply(test)
# ```

# In[44]:


### BEGIN SOLUTION
pmm.keep_history(True)
pmm.apply(start)
pmm.apply(test)
### END SOLUTION


# ## Small and Shear Moves

# In[1]:


from IPython.core.interactiveshell import InteractiveShell
InteractiveShell.ast_node_interactivity = "all"
from IPython import display


# Small mover (1YY9, residue 277):

# In[2]:


from pathlib import Path
gifPath = Path("./Media/small-mover.gif")
# Display GIF in Jupyter, CoLab, IPython
with open(gifPath,'rb') as f:
    display.Image(data=f.read(), format='png',width='400')


# Shear mover (1YY9, residue 277):

# In[4]:


gifPath = Path("./Media/shear-mover.gif")
# Display GIF in Jupyter, CoLab, IPython
with open(gifPath,'rb') as f:
    display.Image(data=f.read(), format='png',width='400')


# The simplest move types are small moves, which perturb φ or ψ of a random residue by a random small angle, and shear moves, which perturb φ of a random residue by a small angle and ψ of the same residue by the same small angle of opposite sign.
# 
# For convenience, the `SmallMover` and `ShearMover` can do multiple rounds of perturbation. They also check that the new φ/ψ combinations are within an allowable region of the Ramachandran plot by using a Metropolis acceptance criterion based on the rama score component change. (The `rama` score is a statistical score from Simons et al. 1999, parametrized by bins of φ/ψ space.) Because they use the Metropolis criterion, we must also supply $kT$.
# 
# Finally, like most `Movers`, these require a `MoveMap` object to specify which degrees of freedom are fixed and which are free to change. Thus, we can create our `Movers` like this:
# 
# ```
# kT = 1.0
# n_moves = 1
# movemap = MoveMap()
# movemap.set_bb(True)
# small_mover = SmallMover(movemap, kT, n_moves)
# shear_mover = ShearMover(movemap, kT, n_moves)
# ```

# In[45]:


### BEGIN SOLUTION
kT = 1.0
n_moves = 1
movemap = MoveMap()
movemap.set_bb(True)
small_mover = SmallMover(movemap, kT, n_moves)
shear_mover = ShearMover(movemap, kT, n_moves)
### END SOLUTION


# We can also adjust the maximum magnitude of the perturbations and get the information back from the `SmallMover` by printing it:
# 
# ```
# small_mover.angle_max("H", 25)
# small_mover.angle_max("E", 25)
# small_mover.angle_max("L", 25)
# print(SmallMover)
# ```

# In[46]:


### BEGIN SOLUTION
small_mover.angle_max("H", 25)
small_mover.angle_max("E", 25)
small_mover.angle_max("L", 25)
print(SmallMover)
### END SOLUTION


# Here, *"H"*, *"E"*, and *"L"* refer to helical, sheet, and loop residues — as they did in the fragment library file — and the magnitude is in degrees. We will set all the maximum angles to 25° to make the changes easy to visualize. (The default values in Rosetta are 0°, 5°, and 6°, respectively.)

# ### Test your mover by applying it to your pose
# ```
# small_mover.apply(test)
# ```

# In[47]:


### BEGIN SOLUTION
small_mover.apply(test)
### END SOLUTION


# Confirm that the change has occurred by comparing the start and test poses in PyMOL.
# 
# ```
# pmm.apply(test)
# ```

# In[48]:


### BEGIN SOLUTION
pmm.apply(test)
### END SOLUTION


# Second, try the PyMOL animation controls on the bottom right corner of the Viewer window. There should be a play button (►) as well as frame-forward, rewind, etc. Play the movie to watch PyMOL shuffle your pose move back and forth.
# 
# __Question:__ Can you identify which torsion angles changed? By how much? If it is hard to view on the screen, it may help to use your programs from previous workshops to compare torsion angles or coordinates.

# In[ ]:





# ### Comparing small and shear movers

# Reset the test pose by re-assigning it the conformation from `start`, and create and view a second test pose (`test2`) in the same manner. Reset the existing `MoveMap` object to only allow the backbone angles of residue 50 to move. (Hint: Set all residues to `False`, then set just residues 50 and 51 to `True`). Note that the `SmallMover` contains a pointer to your `MoveMap`, and so it will automatically know you have made these changes and use the modified `MoveMap` in future moves.
# 
# ```
# test2 = Pose()
# test2.assign(start)
# test2.pdb_info().name("test2")
# pmm.apply(test2)
# 
# movemap.set_bb(False)
# movemap.set_bb(50, True)
# movemap.set_bb(51, True)
# print(movemap)
# ```

# In[49]:


### BEGIN SOLUTION
test2 = Pose()
test2.assign(start)
test2.pdb_info().name("test2")
pmm.apply(test2)

movemap.set_bb(False)
movemap.set_bb(50, True)
movemap.set_bb(51, True)
print(movemap)
### END SOLUTION


# Make one small move on one of your test poses and one shear move on the other test pose. Output both poses to PyMOL using the `PyMOLMover`. Be sure to set the name of each pose so they are distinguishable in PyMOL. Show only backbone atoms and view as lines or sticks. Identify the torsion angle changes that occurred.
# 
# ```
# small_mover.apply(test)
# shear_mover.apply(test2)
# pmm.apply(test)
# pmm.apply(test2)
# ```

# In[50]:


### BEGIN SOLUTION
small_mover.apply(test)
shear_mover.apply(test2)
pmm.apply(test)
pmm.apply(test2)
### END SOLUTION


# __Question:__ What was the magnitude of the change in the sheared pose? How does the displacement of residue 60 compare between the small- and shear-perturbed poses?

# In[ ]:





# ## Minimization Moves

# In[5]:


from IPython.display import Image
Image('./Media/minmover.png',width='300')


# The `MinMover` carries out a gradient-based minimization to find the nearest local minimum in the energy function, such as that used in one step of the Monte-Carlo-plus-Minimization algorithm of Li & Scheraga.
# 
# ```
# min_mover = MinMover()
# ```

# In[51]:


### BEGIN SOLUTION
min_mover = MinMover()
### END SOLUTION


# 3. The minimization mover needs at least a `MoveMap` and a `ScoreFunction`. You can also specify different minimization algorithms and a tolerance. (See Appendix A). For now, set up a new `MoveMap` that is flexible from residues 40 to 60, inclusive, using:
# 
# ```
# mm4060 = MoveMap()
# mm4060.set_bb_true_range(40, 60)
# ```

# In[52]:


### BEGIN SOLUTION
mm4060 = MoveMap()
mm4060.set_bb_true_range(40, 60)
### END SOLUTION


# Create a standard, full-atom `ScoreFunction`, attach these objects to the default `MinMover`, and print out the information in the `MinMover` with the following methods and check that everything looks right:
# 
# ```
# scorefxn = #get the full-atom score function
# min_mover.movemap(mm4060)
# min_mover.score_function(scorefxn)
# print(min_mover)
# ```

# In[53]:


### BEGIN SOLUTION
scorefxn = get_fa_scorefxn()
min_mover.movemap(mm4060)
min_mover.score_function(scorefxn)
print(min_mover)
### END SOLUTION


# Finally, attach an “observer”. The observer is configured to execute a `PyMOLMover.apply()` every time a change is observed in the pose coordinates. The `True` is a flag to ensure that PyMOL keeps a history of the moves.
# 
# ```
# observer = pyrosetta.rosetta.protocols.moves.AddPyMOLObserver(test2, True)
# ```

# In[54]:


### BEGIN SOLUTION
observer = pyrosetta.rosetta.protocols.moves.AddPyMOLObserver(test2, True)
### END SOLUTION


# 4. Apply the `MinMover` to your sheared pose. Observe the output in PyMOL. (This may take a couple minutes — the Observer can slow down the minimization significantly).
# 
# ```
# min_mover.apply(test2)
# ```

# In[55]:


### BEGIN SOLUTION
min_mover.apply(test2)
### END SOLUTION


# __Question:__ How much motion do you see, relative to the original shear move? How many coordinate updates does the *MinMover* try? How does the magnitude of the motion change as the minimization continues? At the end, how far has the Cα atom of residue 60 moved?

# In[ ]:





# ## Monte Carlo Object
# 
# 
# PyRosetta has several object classes for convenience for building more complex algorithms. One example is the `MonteCarlo` object. This object performs all the bookkeeping you need for creating a Monte Carlo search. That is, it can decide whether to accept or reject a trial conformation, and it keeps track of the lowest-energy conformation and other statistics about the search. Having the Monte Carlo operations packaged together is convenient, especially if we want multiple Monte Carlo loops to nest within each other or to operate on different parts of the protein.
# 
# To create the object, you need an initial test pose, a score function, and a temperature factor:
# 
# ```
# mc = MonteCarlo(test, scorefxn, kT)
# ```

# In[56]:


### BEGIN SOLUTION
mc = MonteCarlo(test, scorefxn, kT)
### END SOLUTION


# After the pose is modified by a mover, we tell the `MonteCarlo` object to automatically accept or reject the new conformation and update a set of internal counters by calling:
# 
# ```
# mc.boltzmann(test)
# ```

# In[57]:


### BEGIN SOLUTION
mc.boltzmann(test)
### END SOLUTION


# 5. Test out a `MonteCarlo` object. Before doing so, you may need to adjust your small and shear moves to smaller maximum angles (3–5°) so they are more likely to be accepted. Apply several small or shear moves on your `test` pose, output the score using `print(scorefxn(test))`, then call the `mc.boltzmann(test)` method of the `MonteCarlo` object. A response of `True` indicates the move is accepted, and `False` indicates that the move is rejected. If the move is rejected, the pose is automatically reverted for you to its last accepted state. Manually iterate a few times between moves and calls to `mc.boltzmann()`. Call `pmm.apply(test)` every time you get a `True` back from the `mc.boltzmann(test)` method. Do enough cycles to observe at least two `True` and two `False` outputs. Do the acceptances match what you expect given the scores you obtain? After doing a few cycles, use `mc.show_scores()` to find the score of the last accepted state and the lowest energy state. What energies do you find? Is the last accepted energy equal to the lowest energy?
# 
# ```
# # adjust the SmallMover
# small_mover.angle_max("H", 3)
# small_mover.angle_max("E", 5)
# small_mover.angle_max("L", 6)
# # and the ShearMover
# shear_mover.angle_max("H", 3)
# shear_mover.angle_max("E", 5)
# shear_mover.angle_max("L", 6)
# ```
# 
# Then write your MonteCarlo loop below:

# In[60]:


# adjust the SmallMover
### BEGIN SOLUTION
small_mover.angle_max("H", 3)
small_mover.angle_max("E", 5)
small_mover.angle_max("L", 6)
### END SOLUTION


# In[61]:


# and the ShearMover
### BEGIN SOLUTION
shear_mover.angle_max("H", 3)
shear_mover.angle_max("E", 5)
shear_mover.angle_max("L", 6)
### END SOLUTION


# 6. See what information is stored in the Monte Carlo object using:
# ```
# mc.show_scores()
# mc.show_counters()
# mc.show_state()
# ```
# 
# __Question:__ What information do you get from each of these?

# In[ ]:





# ## Trial Mover

# In[6]:


Image('./Media/trialmover.png',width='250')


# A `TrialMover` combines a specified `Mover` with a `MonteCarlo` object. Each time a `TrialMover` is called, it performs a trial move and tests that move’s acceptance with the `MonteCarlo` object. You can create a `TrialMover` from any other type of `Mover`. You might imagine that, as we start nesting these together, we can build some complex algorithms!
# 
# ```
# trial_mover = TrialMover(small_mover, mc)
# trial_mover.apply(test)
# ```

# In[62]:


### BEGIN SOLUTION
trial_mover = TrialMover(small_mover, mc)
trial_mover.apply(test)
### END SOLUTION


# 7. Apply the `TrialMover` above ten times. Using `trial_mover.num_accepts()` and `trial_mover.acceptance_rate()`, what do you find?

# In[63]:


### BEGIN SOLUTION
for i in range(10):
    trial_mover.apply(test)
    
print(trial_mover.num_accepts())
print(trial_mover.acceptance_rate())
### END SOLUTION


# 8. The `TrialMover` also communicates information to the `MonteCarlo` object about the type of moves being tried. Create a second `TrialMover` (`trial_mover2`) using a `ShearMover` and the same `MonteCarlo` object, and apply this second `TrialMover` ten times like above. After, look at the `MonteCarlo` object state (`mc.show_state()`).
# 
# __Question:__ Using information from `mc.show_state()`, what are the acceptance rates of each mover (`SmallMover` and `ShearMover`)? Which mover is accepted most often, and which has the largest energy drop per trial? What are the average energy drops?

# In[ ]:





# ## Sequence and Repeat Movers
# 
# 
# A `SequenceMover` is another combination `Mover` and applies several `Movers` in succession. It is useful for building up complex routines and is constructed (and confirmed with a print statement) as follows:
# 
# ```
# seq_mover = SequenceMover()
# seq_mover.add_mover(small_mover)
# seq_mover.add_mover(shear_mover)
# seq_mover.add_mover(min_mover)
# print(seq_mover)
# ```

# In[64]:


### BEGIN SOLUTION
seq_mover = SequenceMover()
seq_mover.add_mover(small_mover)
seq_mover.add_mover(shear_mover)
seq_mover.add_mover(min_mover)
print(seq_mover)
### END SOLUTION


# The above example mover will apply first the small, then the shear, and finally the minimization movers.
# 
# 9. Create and print a `TrialMover` using the `SequenceMover` above, and apply it five times to the test pose. How is the sequence mover shown by `mc.show_state()`?

# In[ ]:





# A `RepeatMover` will apply its input Mover `n_repeats` times each time it is applied:
# 
# ```
# n_repeats = 3
# repeat_mover = RepeatMover(trial_mover, n_repeats)
# print(repeat_mover)
# ```

# In[65]:


### BEGIN SOLUTION
n_repeats = 3
repeat_mover = RepeatMover(trial_mover, n_repeats)
print(repeat_mover)
### END SOLUTION


# 10. Use these tools to build up your own *ab initio* protocol. Create `TrialMovers` for 9-mer and 3-mer fragment insertions. First, create `RepeatMovers` for each and then create the `TrialMovers` using the same `MonteCarlo` object for each. Create a `SequenceMover` to do the 9-mer trials and then the 3-mer trials, and iterate the sequence 10 times.
# 
# __Problem:__ Use a pen and paper to write out a flowchartalgorithm:

# In[ ]:





# 11. *Hierarchical search*. Construct a `TrialMover` that tries to insert a 9-mer fragment and then refines the protein with 100 alternating small and shear trials before the next 9-mer fragment trial. The interesting part is this: you will use one `MonteCarlo` object for the small and shear trials, inside the whole 9-mer combination mover. But use a separate `MonteCarlo` object for the 9-mer trials. In this way, if a 9-mer fragment insertion is evaluated after the optimization by small and shear moves and is rejected, the pose goes all the way back to before the 9-mer fragment insertion.

# In[ ]:





# <!--NAVIGATION-->
# < [Structure Refinement](http://nbviewer.jupyter.org/github/RosettaCommons/PyRosetta.notebooks/blob/master/notebooks/05.00-Structure-Refinement.ipynb) | [Contents](toc.ipynb) | [Index](index.ipynb) | [Refinement Protocol](http://nbviewer.jupyter.org/github/RosettaCommons/PyRosetta.notebooks/blob/master/notebooks/05.02-Refinement-Protocol.ipynb) ><p><a href="https://colab.research.google.com/github/RosettaCommons/PyRosetta.notebooks/blob/master/notebooks/05.01-High-Res-Movers.ipynb"><img align="left" src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open in Colab" title="Open in Google Colaboratory"></a>