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

# # 构象的模型表示（Representation of conformation）

# @Author: Jian Huang
# 
# @email：jian.huang@xtalpi.com
# 
# @pictures: partly from weikun.wu@xtalpi.com
# 
# 经过第一章的学习，我们知道在pyrosetta中通过Pose储存构象的所有信息。
# 
# 对于一个大分子构象而言，例如一个简单的含有一个蛋白Model的PDB文件，Rosetta对其中构象的处理主要分为**全原子描述**（Full atom representation）和**质心描述**（centriod representation）。
# 
# **全原子描述**，即对蛋白的所有原子都采用精确（x, y, z）坐标进行描述的方法。相比之下，**质心描述**保留骨架部分的精确全原子描述，但简化了氨基酸残基的侧链，通过对残基整体性质（坐标、原子质量、体积等）计算构建出一个假原子（pseudo atom，又称CEN原子），以该假原子作为原本侧链剩余部分的描述。侧链质心的位置基于侧链整体的质心决定，而假原子的大小由侧链的平均大小决定。可见，这种处理方式简化了蛋白构象，或者说丢失了部分信息，是一种“**粗粒化**”的描述。
# 
# ***
# 
# ***问题一：为什么要费事进行“粗粒化”，而不全部采取精确全原子描述？***
# 
# 理想的情况而言，当我们拥有无尽的时间和无限的计算能力，我们当然是希望所有的体系优化的时候都能够使用全原子表示进行。但往往这是不切实际的，这是出现了简化的质心表示的第一个原因。其次，选择质心表示也可以让能量面比之于全原子表示更为平坦，更有效率地在空间中进行大规模采样，也更容易通过蒙特卡洛标准。（全原子表示的能量面非常崎岖，在执行蒙特卡洛搜索的时候，很多的搜索都会被蒙特卡洛规则拒绝掉而大概率陷入局域最小值）
# 
# ***问题二：质心模式的缺点是什么？***
# 
# 前文提过，质心描述以牺牲分辨率、丢失信息为代价而获得较快的构象空间搜索效率。这种信息的丢失，主要是侧链，例如侧链的氢键作用、侧链范式作用等细节，就没有被显式捕获到。相比之下，rosetta中会通过成对的统计学势能（pair-wise statistical potenhtial）和范式作用的范围型近似(VdW sphere approximation)去隐式捕获这些作用，但是这种处理是近似的、不精确的。
# ***
# 
# 
# **总结**
# 
# 这种两种原子模型在rosetta大部分protocal中都会涉及到。基本思路：首先通过质心粗颗粒描述快速搜索大量的构象空间，称为“low-resolution/Coarse-grain phase”，这样可以快速在能量面上找到能量较低的范围；然后通过全原子描述在该范围内进一步精确、优化地寻找低能构象的细节，称为“high-resolution refinement phase”。

# <center><img src="./img/lowRes-HighRes.jpg" width = "1000" height = "200" align=center /></center>
# 
# (图片来源: google搜索拼接)

# ### 一、单个氨基酸全原子和质心描述的直观感受

# 在Rosetta中，PDB的读入默认采取全原子的方式进行表示。

# In[1]:


from pyrosetta import init, pose_from_sequence
init()


# In[2]:


# 按照rosetta标准库模板，创建一个Trp氨基酸的pose
pose = pose_from_sequence('W')

# 判断一个pose是否为全原子描述
print("Is the pose a full-atom description?\n", pose.is_fullatom())

# 同样地，pose.is_centroid()可以判断是否为质心描述

# 我们先将该构象储存为PDB文件
pose.dump_pdb("./data/trp_fullatom.pdb")


# 当Pose的表示形式需要转换时，可以使用相关的函数来实现(SwitchResidueTypeSetMover)

# In[3]:


# TODO:转化全原子描述为质心描述
from pyrosetta.rosetta.protocols.simple_moves import SwitchResidueTypeSetMover

# 初始化一个SwitchResidueTypeSetMover对象用来操作我们的pose
to_centroid = SwitchResidueTypeSetMover('centroid')
# print(type(to_centroid))

# 对pose进行操作
to_centroid.apply(pose)

# 储存质心描述的pose为PDB文件
pose.dump_pdb("./data/trp_centroid.pdb")


# **使用pymol查看两种描述的PDB文件**
# 
# <center><img src="./img/Trp_centroid_fullatom.png" width="500" height="400" align="center"/></center>
# 
# 左：全原子描述；                          右：质心描述
# 
# 这里顺带强调一下PDB文件，我们用文本编辑器打开trp_centroid.pdb文件，如下：
# 
# ```
# ATOM      1  N   TRP A   1       0.000   0.000   0.000  1.00  0.00           N  
# ATOM      2  CA  TRP A   1       1.458   0.000   0.000  1.00  0.00           C  
# ATOM      3  C   TRP A   1       2.009   1.420   0.000  1.00  0.00           C  
# ATOM      4  O   TRP A   1       2.058   2.045   1.023  1.00  0.00           O  
# ATOM      5  CB  TRP A   1       1.991  -0.769  -1.210  1.00  0.00           C  
# ATOM      6  **CEN** TRP A   1       2.937  -2.180  -1.648  1.00  0.00           X  
# ATOM      7  H   TRP A   1      -0.500  -0.433  -0.750  1.00  0.00           H  
# TER
# ```
# 
# 可以看到Trp侧链的假原子名字为**CEN**。（请打开trp_fullatom.pdb查看其中的ATOM描述，观察差别）

# ### 二、一个更“有趣”的例子
# 接下来的案例，我们将展示Rosetta对PDB->Pose过程中的处理逻辑，实现过程:
# 1. 从PDB库中下载的1ubq.pdb文件，并进行简单处理
# 2. 转化为质心描述，保存质心描述的构象为PDB
# 3. 将该质心描述的构象再转化为全原子描述（称为new_fullatom）
# 4. 比较从PDB库下载后简单处理的全原子描述的构象与再次转化后的全原子描述构象

# In[4]:


from pyrosetta.toolbox.rcsb import pose_from_rcsb
from pyrosetta import pose_from_pdb
from pyrosetta.toolbox import cleanATOM

origin_pose = pose_from_rcsb('1ubq')
origin_pose.dump_pdb("./data/1ubq.pdb")

cleanATOM("./data/1ubq.pdb", out_file="./data/1ubq_clean.pdb")


# In[5]:


# 查看pose的信息
print(origin_pose)
print(origin_pose.is_fullatom())

# 简要描述
# print("Brief description about the origin_pose:", origin_pose.pdb_info().short_desc())

# 该输出打印了非常详细的信息，总的残基数、序列、Foldtree。


# In[6]:


# 初始化两个SwitchResidueTypeSetMover对象
to_centroid = SwitchResidueTypeSetMover('centroid')
to_fullatom = SwitchResidueTypeSetMover('full_atom')

# 读取pose，转化为质心描述
pose_1ubq = pose_from_pdb("./data/1ubq_clean.pdb")
to_centroid.apply(pose_1ubq)


# In[7]:


# 是否是质心描述呢？
print(pose_1ubq.is_centroid())
pose_1ubq.dump_pdb("./data/1ubq_clean_centroid.pdb")


# In[8]:


# 此时进行第三步，将质心描述转化为全原子描述
# 可以直接从已经变成centroid描述的pose_1ubq开始操作
to_fullatom.apply(pose_1ubq)

# 也可以从./data/1ubq_clean_centroid.pdb重新读入
pose_1ubq_centroid = pose_from_pdb("./data/1ubq_clean_centroid.pdb")
print(pose_1ubq_centroid.is_fullatom())


# 在这里注意pose_from_pdb函数将会直接将原本为质心描述的构象转化为全原子描述！
# 
# 查看输出日志可以发现 [ WARNING ] discarding xx atoms at position xx in file
# 
# 这种信息经常出现，这就是Rosetta的运行逻辑决定的，该函数会检查侧链残基，发现CEN就全部丢弃掉，然后按照理想的氨基酸侧链补全了缺失的原子，并进行了一定的侧链的能量优化。
# 
# **因此如果从PDB文件读入并生成Pose时，使用init('-use_input_sc')可以保留原始的结构侧链信息。**

# In[9]:


# 输出再次pack上侧链的全原子构象
pose_1ubq.dump_pdb("./data/1ubq_clean_newfullatom.pdb")

# 输出从质心PDB文件读取、恢复的全原子构象
pose_1ubq_centroid.dump_pdb("./data/1ubq_clean_newfullatom2.pdb")


# ### 思考

# 1. 1ubq_clean_newfullatom.pdb文件和1ubq_clean.pdb文件的构象有什么区别？为什么会有这种区别？(难度指数：*)
# 
# ![title](./img/1ubq_fullatom.png)

# 2. 1ubq_clean_newfullatom.1ubq_clean_newfullatom2.pdb文件的构象有什么区别？为什么会有这种区别？（难度指数：***）

# In[ ]: