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

# # Rosetta Scripts in PyRosetta (Advanced)
# 
# @Author: Jian Huang
# 
# @E-mail: jian.huang@xtalpi.com

# 通过第一节的学习，理论而言，只要我们掌握了RS的各个tag的含义和定义，就可以进行非常复杂且用户定制化的protocols。本节的内容提供关于Minimization和packing的RS高级设置，让我们可以与之前学习过的TaskOperation和MoveMap知识进行联动。并进一步了解RS的编写和用法。
# 
# </br>
# 
# **写在开头：一点回顾**
# 
# *关于packing*
# 
# Rosetta中的packing过程：在一个蛋白pose构象中，由于每个位置的氨基酸都具有其本身的rotamer分布，想要去找到哪种组合的rotamer会让体系能量最低的过程。rotamer其实就是氨基酸侧链的构象分布，在rosetta中使用的rotamer library是在PDB库中采取那些精度特别高（1.8 A）的晶体结构进行统计的结果，虽然rotamer在空间中分布是连续的，但是概率出现最大的几个位置被选出来作为能量优势构象，然后连续的rotamer分布面就被几个分离的概率较大的rotamer作为代表而简化处理了。
# 
# 即便如此，在N个氨基酸组成的蛋白pose中，假设每个氨基酸由M种rotamer，那么搜索的空间就是M的N次方，可以看到随着N的增加，搜索空间将无比巨大。Rosetta的Packer就是使用一种蒙特卡洛的随机搜索算法，以接受能够使构象能量降低的构象变化或以一定概率接受能量升高的构象变化，目的使寻找这些N个氨基酸的哪些rotamer的组合，可以给出能量最优点的构象。从上述描述中可以看到，该搜索算法使随机的，它只能够保证运行时间结束后，得到的构象能量会比之前的小，但却并不一定使全局能量最优构象。（“尽量去搜索最优点，但不保证能够得到最优点”）
# 
# *关于minimizer*
# 
# 第一节使用的MinMover是一个使用最朴实的Minimization的例子，其在底层调用了rosetta中的底层算法——minimizer。
# 
# 之前在学习minimizer的时候，提到所有的minimizer的操作使用了rosetta的**Movemap**来确定构象的自由度范围。
# 
# MoveMap在minimizer中是可以允许用户精准控制能量最小化限制的操作。MoveMap可以允许用户精确定义是否移动backbone（BB）的扭转角度变化（也就是phi角和psi角度）以及侧链扭转角变化（chi角）。其实从这里的描述也可以看到，minimizer与packer十分相似，相比于packer是去寻找哪种组合的rotamer会让能量最低，minimizer也是在空间中搜索哪一些组合的扭转角会让构象的能量最低，只是他们用的搜索算法和优化能量采取的研究对象不同（一个是rotamer，一个是扭转角）罢了。同时也可以看到，如果我们先用不做design的packer优化一下能量，但由于我们采用的rotamer库使用的是分离的侧链构象，近似为理想化的情况，此时做minimizer就很有必要了，因为它可以对构象的扭转角进行一些微调，从而采样到那些并不是理想化rotamer的情况的构象。
# 
# 值得注意的是，尽管我们可以在MoveMap中禁止某些侧链扭转角的变化，也并不代表该残基在空间中不会发生位移。其实这也很好理解，蛋白质链都是依次序连接在一起，如果上游的某个扭转角改变了，势必会影响它接下来的原子在空间中的位置，尽管后续的扭转角都没有变化。在rosetta中，其实是采取了一个叫FoldTree的概念来描述这种连接的关系。（关于FoldTree，请参考本教程相关章节）
# 
# 在使用MoveMap进行minimization的时候，我们必须注意FoldTree。rosetta中默认的FoldTree定义第一条链A的第一个残基作为根残基（root），依次后面称第二个残基是子代残基1，第三个残基是子代残基2，......等等；如果存在多条链，jump会转接到另一条链的第一个残基，而后重复上一步操作，形成完成的FoldTree。任何对于指定残基的自由度的变化都会影响到它在Foldtree中下游子代。
# 
# 在MoveMap中可以自由定义各种自由度。与MinMover相似，packing操作在RS中其实也是通过定义mover来进行，其在底层会调用packer以及控制packer的**TaskOperation**进行精确的packer操作。
# 
# ***
# 
# 我们应该注意，**Movemap**的定义具有不可交换性，某一个残基的Movemap定义以其最后一次出现的定义为准；而**TaskOperation**具有可交换性，交换一个packing过程的各个TaskOperation定义等同。
# 

# ## 1. RS in Pyrosetta：精准定义MoveMap的Minimization

# 在第一节中，我们给出了官网上MinMover的定义范式，在```<MinMover> </MinMover>```标签内含有允许用户定义的亚标签——```<MoveMap></MoveMap>```。
# 
# 我们可以利用该亚标签对MinMover进行细致、精确地定义，以满足我们对某一pose的rosetta采样需求。
# 
# Example：
# 
# ```
# <SCOREFXNS>
#         <ScoreFunction name="ref_2015" weights="ref2015"/>
# </SCOREFXNS>
# <MOVERS>
#     <MinMover name="min_torsion" scorefxn="ref_2015" chi="true" bb="1" cartesian="F" >
#         <MoveMap name="min_torsion_mm">                         
#             <Span begin="1" end="999" chi="false" bb="false" /> 
#             <Span begin="1" end="50" chi="true" bb="true" />    
#             <Span begin="5" end="10" chi="true" bb="false" />   
#         </MoveMap>                                              
#     </MinMover>
# </MOVERS>
# <PROTOCOLS>
#     <Add mover="min_torsion" />
# </PROTOCOLS>
# ```
# 
# 过程：
# 
# 1. 先限制1~999号残基的BB和CHI移动；
# 2. 然后允许1~55号残基的CHI和BB运动；
# 3. 最后限制5~10号残基的BB不运动、CHI可以运动。
# 
# NOTE: 
# 这里先定义了一个很宽的范围1 - 999，注意在Rosetta中即使我们现在正在模拟的蛋白长度没有这么多，也是允许的！只是其他的限制仍然还是在输入构象基础上进行限制的（比如1 - 50，5 - 10）。（假定我们模拟的蛋白长度是76，我们知道MoveMap遵循依次定义限制的规则，一个残基的自由度由最后一次出现它定义的地方指定，比如5 - 10号残基先是禁止所有运动，而后开放了CHI和BB的运动，最后又限制了BB运动，开放CHI运动——最终的结果就是5 - 10号允许CHI运动、不允许BB运动；而1 - 4和11 - 50具有BB和CHI自由度；51 - 76不允许任何移动）

# In[1]:


from pyrosetta.rosetta.protocols.rosetta_scripts import *
from pyrosetta import *
init()


# In[2]:


# 读入初始pose，并且拷贝一份
original_pose = pose_from_pdb("./data/1ubq_clean.pdb")
pose = original_pose.clone()


# In[3]:


# 也可以将rosetta scripts直接在python中用字符串表示
# 省去读入文件一步

RS_mini = \
"""
<ROSETTASCRIPTS>
    <SCOREFXNS>    
        <ScoreFunction name="ref_2015" weights="ref2015" />
    </SCOREFXNS>
    <RESIDUE_SELECTORS>
    </RESIDUE_SELECTORS>
    <TASKOPERATIONS>
    </TASKOPERATIONS>
    <FILTERS>
    </FILTERS>
    <MOVERS>
        <MinMover name="min_torsion" scorefxn="ref_2015" chi="true" bb="1" cartesian="F" >
            <MoveMap name="min_torsion_mm">                         
                <Span begin="1" end="999" chi="false" bb="false" /> 
                <Span begin="1" end="50" chi="true" bb="true" />    
                <Span begin="5" end="10" chi="true" bb="false" />   
            </MoveMap>                                              
        </MinMover>
    </MOVERS>
    <APPLY_TO_POSE>
    </APPLY_TO_POSE>
    <PROTOCOLS>
        <Add mover="min_torsion" />    
    </PROTOCOLS>
    <OUTPUT scorefxn="ref_2015" />
</ROSETTASCRIPTS>
"""


# In[4]:


# it's gonna take a while for Rosetta to parse the RS
xml = XmlObjects.create_from_string(RS_mini)


# In[5]:


protocol = xml.get_mover("ParsedProtocol")
protocol.apply(pose)


# ## 2. RS in Pyrosetta：精准定义TaskOperation的packing

# </br>
# 
# 与Minimizer控制二面角自由度类似，packing过程通过TaskOperation控制packer的氨基酸design和packing的自由度。
# 
# RS中对应经典的packer的Mover名为：PackRotamersMover

# </br>
# 
# 同样地，我们可以从官方网页上找到定义PackRotamersMover的定义范式：
# （https://new.rosettacommons.org/docs/latest/scripting_documentation/RosettaScripts/Movers/movers_pages/PackRotamersMover）
# 
# ```
# <TASKOPERATIONS>
#     # Some TaskOperation definition here
# </TASKOPERATIONS>
# <PackRotamersMover name="(&string;)" nloop="(1 &non_negative_integer;)"
#         scorefxn="(&string;)"
#         task_operations="(&task_operation_comma_separated_list;)"
#         packer_palette="(&named_packer_palette;)" />
#         
# ```
# >nloop: Equivalent to "-ndruns".Number of complete packing runs before an output (best score) is produced.
# >
# >scorefxn: Name of score function to use
# >
# >task_operations: A comma-separated list of TaskOperations to use.
# >
# >packer_palette: A previously-defined PackerPalette to use, which specifies the set of residue types with which to design (to be pruned with TaskOperations).
# 
# 
# **注意**
# 
# 1. RS中有单独的```<TASKOPERATIONS></TASKOPERATIONS>```标签定义全局的TaskOperation内容。然后通过在PackRotamersMover的选项：task_operation控制需要使用的TaskOperations。
# 
# 2. PackRotamersMover中的task_operation选项可以直接删掉（定义为默认值），但是默认的PackRotamersMover会允许每一个位置使用20种的氨基酸的所有rotamer进行packing！
# 
# ***
# 
# 故此，我们还需要了解如何定义```<TASKOPERATIONS></TASKOPERATIONS>```。官方例子如下：
# 
# ```
# ...
# <TASKOPERATIONS>
#   <ReadResfile name="rrf"/>
#   <ReadResfile name="rrf2" filename="resfile2"/>
#   <PreventRepacking name="NotBeingUsedHereButPresenceOkay"/>
#   <RestrictResidueToRepacking name="restrict_Y100" resnum="100"/>
#   <RestrictToRepacking name="rtrp"/>
#   <OperateOnResidueSubset name="NoPackHelix">
#       <SecondaryStructure ss="H" />
#       <PreventRepackingRLT/>
#   </OperateOnResidueSubset>
# </TASKOPERATIONS>
# ...
# <MOVERS>
#   <PackRotamersMover name="packrot" scorefxn="sf" task_operations="rrf,NoPackHelix,rtrp,restrict_Y100"/>
# </MOVERS>
# ...
# ```
# 
# 
# Rosetta为了设定TaskOperation的便利，内置定义了多个独特的TASKOPERATIONS，例如PreventRepacking、RestrictResidueToRepacking等等。

# 除此之外，TaskOperations还会控制侧链如何被采样的具体过程，rosetta中默认是基于rotamer的采样，但常常可能希望加入一些额外的rotamer的采样。例如，添加一些偏移标准rotamer库的因子，ExtraRotamersGeneric的TaskOperations允许我们控制rotamer采样的级别，一般来说在CHI1和CHI2中加入一些额外的采样比较有效，虽然可能会花费更长的时间进行packing。此外InitializeFromCommandline也允许使用-ex1 -ex2的选项进行rotamer的采样控制。
# 
# ```
# 
# <ExtraRotamersGeneric name="(&string)"
# ex1="(0 &boolean)" ex2="(0 &boolean)" ex3="(0 &boolean)" ex4="(0 &boolean)"
# ex1aro="(0 &boolean)" ex2aro="(0 &boolean)" ex1aro_exposed="(0 &boolean)" ex2aro_exposed="(0 &boolean)"
# ex1_sample_level="(7 &Size)" ex2_sample_level="(7 &Size)" ex3_sample_level="(7 &Size)" ex4_sample_level="(7 &Size)"
# ex1aro_sample_level="(7 &Size)" ex2aro_sample_level="(7 &Size)" ex1aro_exposed_sample_level="(7 &Size)" ex2aro_exposed_sample_level="(7 &Size)"
# exdna_sample_level="(7 &Size)"
# extrachi_cutoff="(18 &Size)"/>
# 
# ```

# In[6]:


# 定义本例中的RS
RS_pack = \
"""
<ROSETTASCRIPTS>
        <SCOREFXNS>
                <ScoreFunction name="r15" weights="ref2015" />
        </SCOREFXNS>
        <RESIDUE_SELECTORS>
        </RESIDUE_SELECTORS>
        <TASKOPERATIONS>
                <RestrictToRepacking name="no_design" />    
                <ExtraRotamersGeneric name="extrachi" ex1="1" ex2="1" ex1_sample_level="1" ex2_sample_level="1" />
        </TASKOPERATIONS>
        <FILTERS>
        </FILTERS>
        <MOVERS>
                <PackRotamersMover name="pack1" scorefxn="r15" task_operations="no_design,extrachi" />
        </MOVERS>
        <APPLY_TO_POSE>
        </APPLY_TO_POSE>
        <PROTOCOLS>
                <Add mover="pack1" />
        </PROTOCOLS>
        <OUTPUT scorefxn="r15" />
</ROSETTASCRIPTS>
"""

# 第一个TASKOPERATION：RestrictToRepacking 会禁止design；
# 第二个TASKOPERATION：ExtraRotamersGeneric会允许额外rotamer的采样


# In[7]:


pose_2 = original_pose.clone()

xml_2 = XmlObjects.create_from_string(RS_pack)

protocol_2 = xml_2.get_mover("ParsedProtocol")

protocol_2.apply(pose_2)


# ## 3. RS in Pyrosetta VS python scripts

# </br>
# 
# 下面我们以pyrosetta的python脚本和RS分别实现同样的过程，过程中感受一下两种方式的优缺点。
# 
# 这里我们结合以上两个小节：精确定义MoveMap的Minimization（MinMover）和精确定义TaskOperation的Packing（PackRotamersMover）结合起来构成一个新的**RS protocol**，要求：
# 
# 1. 全部过程使用ref2015能量函数；
# 2. 先使用定义好的pack1（与第二小节保持一致）进行packing；
# 3. 再使用定义好的min_torsion(与第一小节保持一致)进行minimization；
# 4. 计算初始pose和结束pose的能量值
# 
# 其次，我们再使用**pyrostta的python脚本版本**进行这样的操作。比较两者。

# **使用RS**

# In[8]:


# RS
RS_pack_min = \
"""
<ROSETTASCRIPTS>
    <SCOREFXNS>    
    
        <ScoreFunction name="ref_2015" weights="ref2015" />
        
    </SCOREFXNS>
    <RESIDUE_SELECTORS>
    </RESIDUE_SELECTORS>
    <TASKOPERATIONS>
    
        <RestrictToRepacking name="no_design" />  
        
        <ExtraRotamersGeneric name="extrachi" ex1="1" ex2="1" ex1_sample_level="1" ex2_sample_level="1" />
        
    </TASKOPERATIONS>
    <FILTERS>
    </FILTERS>
    <MOVERS>
    
        <MinMover name="min_torsion" scorefxn="ref_2015" chi="true" bb="1" cartesian="F" >
            <MoveMap name="min_torsion_mm">                         
                <Span begin="1" end="999" chi="false" bb="false" /> 
                <Span begin="1" end="50" chi="true" bb="true" />    
                <Span begin="5" end="10" chi="true" bb="false" />   
            </MoveMap>                                              
        </MinMover>
        
        <PackRotamersMover name="pack1" scorefxn="ref_2015" task_operations="no_design,extrachi" />
        
    </MOVERS>
    <APPLY_TO_POSE>
    </APPLY_TO_POSE>
    <PROTOCOLS>
        <Add mover="pack1" /> 
        <Add mover="min_torsion" />
    </PROTOCOLS>
    <OUTPUT scorefxn="ref_2015" />
</ROSETTASCRIPTS>
"""

pose_3 = original_pose.clone()

xml_3 = XmlObjects.create_from_string(RS_pack_min)

protocol_3 = xml_3.get_mover("ParsedProtocol")

protocol_3.apply(pose_3)


# In[9]:


scorefxn_RS = xml_3.get_score_function("ref_2015")
before = scorefxn_RS.score(original_pose)
after = scorefxn_RS.score(pose_3)
print(before, after)


# ***
# **使用Python Script**

# In[10]:


#Python
from pyrosetta import *
from pyrosetta.rosetta import *

#Core Includes
from rosetta.core.kinematics import MoveMap
from rosetta.core.kinematics import FoldTree
from rosetta.core.pack.task import TaskFactory
from rosetta.core.pack.task import operation
from rosetta.core.simple_metrics import metrics
from rosetta.core.select import residue_selector as selections
from rosetta.core import select
from rosetta.core.select.movemap import *
from rosetta.protocols import minimization_packing as pack_min


# In[11]:


init()


# In[12]:


pose_4 = original_pose.clone()
# 查看pose基本信息
print(pose)


# In[13]:


# Use TaskFactory to control TaskOperations -- packing
tf = TaskFactory()
# InitializeFromCommandline会调用初始init（）中的选项
tf.push_back(operation.InitializeFromCommandline())
# RestrictToRepacking可以限制不允许进行design
tf.push_back(operation.RestrictToRepacking())
# ExtraRotamersGeneric允许额外rotamer采样
extra_rotamer = operation.ExtraRotamersGeneric()
extra_rotamer.ex1(True)
extra_rotamer.ex2(True)
extra_rotamer.ex1_sample_level(pyrosetta.rosetta.core.pack.task.ExtraRotSample.EX_ONE_STDDEV)
extra_rotamer.ex2_sample_level(pyrosetta.rosetta.core.pack.task.ExtraRotSample.EX_ONE_STDDEV)
tf.push_back(extra_rotamer)


# In[14]:


selection_1_4 = selections.ResidueIndexSelector("1-4")
selection_5_10 = selections.ResidueIndexSelector("5-10")
selection_11_50 = selections.ResidueIndexSelector("11-50")


# In[15]:


print("residue 1 - 4", selection_1_4.apply(pose))
print("residue 5 - 10", selection_5_10.apply(pose))
print("residue 11 - 50", selection_11_50.apply(pose))


# In[16]:


# Use MoveMapFactory to control MoveMap -- Minimization
# 默认为自由度全部关闭状态！按照需求打开自由度
mmf = MoveMapFactory()
# 开放1 - 4 残基骨架和侧链的自由度
mmf.add_chi_action(mm_enable, selection_1_4)
mmf.add_bb_action(mm_enable, selection_1_4)
# 仅开放5 - 10号允许CHI运动
mmf.add_chi_action(mm_enable, selection_5_10)
# 开放11 - 50 残基骨架和侧链的自由度
mmf.add_chi_action(mm_enable, selection_11_50)
mmf.add_bb_action(mm_enable, selection_11_50)
# 关闭51 - 76任何自由度：无需操作，默认关闭
mm  = mmf.create_movemap_from_pose(pose)
# 查看MoveMap，观察是否与设想的一致
print(mm)


# In[17]:


# run protocol as follows:
scorefxn = create_score_function( "ref2015" )

# 首先进行packing -- 对应于pack1
packer = pack_min.PackRotamersMover()
packer.score_function(scorefxn)
packer.task_factory(tf)
packer.apply(pose_4)

# 再进行minimization -- 对应于min_torsion
minimizer = pack_min.MinMover()
minimizer.score_function(scorefxn)
minimizer.cartesian(False)
minimizer.set_movemap(mm)
minimizer.apply(pose_4)

# calculate scores
before = scorefxn.score(original_pose)
after = scorefxn.score(pose_4)
print(before, after)


# ***

# ## 小结与练习
# 
# 在RS中进行的定义可以与Pyrosetta的python版本一一对应。
# 
# 1. 对比两种方式，各有什么优缺点？你会选择什么方式定制自己的protocol呢？
# 
# 2. 请将./data/Example2-MinMover.xml改为python script，并尝试在两个MinMover中设置不同的MoveMap。