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

# # PyMOL Visualization
# PyMOL作为最流行的分子结构可视化的工具之一，在PyRosetta中内置了相关的服务可以直接和PyMOL进行工作联动。这使得我们可以实时观察结构和打分的变化。这大大有助于理解采样和计算过程中发了什么。

# ### 一、PyMOLRosettaServer的配置
# PyRosetta的一大特点在于可以使用PyMOL作为可视化的窗口，实时观测采样的过程。 首先安装PYMOL, 通过conda可以轻松获取。
# 
# 运行命令:
# ```
# conda install -c schrodinger pymol
# ```

# 启动PyMOl命令:
# ```
# pymol
# ```
# 
# 图形界面菜单栏-点击-File>Edit pymolrc. 在其中添加:（路径根据自己PyRosetta安装的地址进行填写。）
# ```
# run /opt/miniconda3/envs/pyrosetta/lib/python3.6/site-packages/pyrosetta/PyMOLRosettaServer.py
# ```
# 
# **重启PyMOL后如有以下字样说明配置成功**:
# 
# PyMOL>run /opt/miniconda3/envs/pyrosetta/lib/python3.7/site-packages/pyrosetta/PyMOLRosettaServer.py
# 
# PyMOL <---> PyRosetta link started!
# 
# at 127.0.0.1 port 65000
# 
# <img src="./img/pymolrosettaserver.png" width = "800" height = "200" align=center />
# (图片来源: 知乎Rosetta研习社)

# ### 二、使用PyMOL观察建模结果

# #### 2.1 发送Pose至PyMOL
# 在前面章节中，我们已经学会如何得到Pose的4种方法，那就先来生成Pose，并发送到PyMOL中进行观察。

# In[1]:


# 读入Pose;
from pyrosetta import pose_from_pdb,init,create_score_function
init()
pdb_pose = pose_from_pdb('./data/pose_demo.pdb')


# PyRosetta和PyMOL的交互主要通过PyMOLMover这个模块进行实现。

# In[2]:


# 展示如何发送pose到pymol中
from pyrosetta.rosetta.protocols.moves import PyMOLMover
pymover = PyMOLMover()
pymover.update_energy(True)
pymover.apply(pdb_pose)


# 此处截图可见，Pose已经发送至PyMOL窗口。
# 
# <img src="./img/sendpose.png" width = "900" height = "200" align=center />
# (图片来源: 知乎Rosetta研习社)

# #### 2.2 调用PyMOLMover进行能量分析
# PyMOLMover还有一些跟能量或性质进行着色管理，PyMOLMover能量展示相关的api函数:
# - send_hbonds 
# - send_energy/send_RAW_Energies
# - send_polars(注意:不兼容新版本的pymol)

# In[3]:


# 对pose先进行能量计算（必须执行）
scoring = create_score_function('ref2015')
scoring(pdb_pose)


# In[4]:


# 显示所有氢键的情况:
pymover.send_hbonds(pdb_pose)


# <img src="./img/pymol_hbnet.png" width = "600" height = "200" align=center />
# (图片来源: 知乎Rosetta研习社)

# In[5]:


# 发送pose的能量着色到pymol中
pymover.send_energy(pdb_pose)


# 蓝色代表这个残基的能量越低(favorable energy)，残基着色越红代表能量越高(unfavorable energy)
# <img src="./img/pymol_energy.png" width = "600" height = "200" align=center />
# (图片来源: 知乎Rosetta研习社)

# 除了整体着色以外，也可以支持仅对某个能量项进行着色，比如，我目前最关心分子内部的溶剂化能（看看哪些区域有很大暴露的疏水氨基酸，这会导致蛋白质折叠不稳定）。运行结果表示，β折叠片上的氨基酸侧链上的疏水原子有比较好的packing，暴露的疏水面积较α-螺旋上的少。

# In[6]:


# 发送某个能量打分项到pymol中
pymover.send_energy(pdb_pose,'fa_sol')


# <img src="./img/pymover_fasol.png" width = "700" height = "200" align=center />

# #### 2.3 调用PyMOLObserver进行实时轨迹观察

# 除了静态的分析能量，在PyMOLMover中可以调用PyMOL Observer实时观察并记录构象变化的过程。
# 
# 其中两个关键的模块:
# - AddPyMOLObserver_to_conformation，每当Pose发生构象变化时，用观察器记录构象并输出到PyMOL窗口。
# - AddPyMOLObserver_to_energies，仅当能量发生变化时，用观察器记录构象并输出到PyMOL窗口。

# In[7]:


# loading a pose;
from pyrosetta import pose_from_pdb,init,create_score_function
init()
peptide_pose = pose_from_pdb('./data/4jfx_peptide.pdb')

# 对pose先进行能量计算（必须执行）
scoring = create_score_function('ref2015')
scoring(peptide_pose)


# In[8]:


#AddPyMOLObserver(pose, True)
from pyrosetta.rosetta.protocols.moves import AddPyMOLObserver_to_conformation, AddPyMOLObserver_to_energies
from pyrosetta.rosetta.protocols.minimization_packing import MinMover
from pyrosetta.rosetta.core.kinematics import MoveMap


# In[9]:


# 创建pymol构象变化观察器;
AddPyMOLObserver_to_conformation(peptide_pose, True)

# 创建pymol能量变化观察器;（只能选其一）
# AddPyMOLObserver_to_energies(peptide_pose, True)

# 对多肽进行能量最小化;
mp = MoveMap()
mp.set_bb(True)
mp.set_chi(True)
min_mover = MinMover()
min_mover.movemap(mp)
min_mover.apply(peptide_pose)


# <center><img src="./img/PyMOLObserver_conformers.gif" width = "500" height = "200" align=center /><center/>
# (图片来源: 知乎Rosetta研习社)

# 在后面的一些章节中，还可以是用一些特殊的残基选择器得到可以再PyMOL中进行选择的命令语句，配合使用也可以进行相应的可视化功能。

# In[ ]: