技术背景

在前面的这一篇博客中，比较全面的介绍了组成蛋白质的各种氨基酸的三维结构。由于每个氨基酸大小不一，在传统的蛋白质折叠预测的方案中，一般会考虑全原子方案或者是粗粒化方案。对于全原子方案而言，即时去除了氢原子，也包含了极大的原子数，对于计算量来说是一个非常大的考验。而将一个氨基酸近似为一个点的方案，因为往往忽略了太多的信息，比如氨基酸之间的二面角等，因此无法达到很好的预测效果。在AlphaFold中，将每一个氨基酸在主链上的位置，用一个三角形刚体来表示。这个三角形的三个顶点分别是C原子、N原子和\(\alpha\)位的C原子。由于一个三角形就可以确定一个平面，因此每一个氨基酸可以通过一个三角形的内坐标和一个顶点的笛卡尔坐标来完整的表示出来。

算法实现

这里我们使用MindSpore复现一下AlphaFold2中的算法21，也就是刚体表示算法。由于是一个刚体，我们可以将其中除了用于定位顶点的笛卡尔坐标以外的内坐标做一个归一化处理。该算法的基本流程为：

输入三角形\(\Delta CC_{\alpha}N\)的笛卡尔坐标；
计算\(\vec{C_{\alpha}C}\)的单位向量，表示为\(\vec{e_1}\)；
计算\(\vec{C_{\alpha}C}\)到\(N\)的垂向量的单位向量，表示为\(\vec{e_2}\)；
计算\(\vec{e_1}\)和\(\vec{e_2}\)的叉乘，得到一个跟三角形\(\Delta CC_{\alpha}N\)平面相垂直的单位向量\(\vec{e_3}\)；
用\(\vec{e_1}\)、\(\vec{e_2}\)和\(\vec{e_3}\)表示当前氨基酸的空间位置，保留\(C_{\alpha}\)的坐标用于表示当前氨基酸在空间中的绝对位置。

具体算法实现如下所示：

algorithm21.py

关闭MindSpore的Warning信息

import os
os.environ['GLOG_v'] = '4'

import numpy as np
import mindspore as ms
from mindspore import Tensor, ops
from mindspore import numpy as msnp

def rigidFrom3Points(crd):
""" Transform the coordinates formulation. """
v1 = crd[:, 2] - crd[:, 1]
v2 = crd[:, 0] - crd[:, 1]
e1 = v1 / v1.norm(axis=-1, keep_dims=True)
u2 = v2 - e1 * ops.Einsum('ij,ij->i')((e1, v2))[:, None]
e2 = u2 / u2.norm(axis=-1, keep_dims=True)
e3 = msnp.cross(e1, e2, axisc=-1)
R = ops.Concat(axis=-2)((e1[:, None], e2[:, None], e3[:, None]))
t = crd[:, 1][:, None]
new_crd = ops.Concat(axis=-2)((R, t))
return new_crd

if name == 'main':
np.random.seed(1)
BATCHES = 2
DIMENSIONS = 3
ATOMS = 3
# 定义原始坐标
origin_crd = Tensor(np.random.random((BATCHES, ATOMS, DIMENSIONS)), ms.float32)
print ('The original coordinates is: \n{}'.format(origin_crd))
trans_crd = rigidFrom3Points(origin_crd)
print ('The transformed coordinates is: \n{}'.format(trans_crd))

该算法的输出结果为：

The original coordinates is:
[[[4.17021990e-01 7.20324516e-01 1.14374816e-04]
  [3.02332580e-01 1.46755889e-01 9.23385918e-02]
  [1.86260208e-01 3.45560730e-01 3.96767467e-01]]

 [[5.38816750e-01 4.19194520e-01 6.85219526e-01]
  [2.04452246e-01 8.78117442e-01 2.73875929e-02]
  [6.70467496e-01 4.17304814e-01 5.58689833e-01]]]
The transformed coordinates is:
[[[-0.30411592  0.5208795   0.79762024]
  [ 0.30775824  0.8461014  -0.43519795]
  [-0.90155333  0.11312358 -0.41761786]
  [ 0.30233258  0.14675589  0.09233859]]

 [[ 0.55235845 -0.5461919   0.6297417 ]
  [-0.73928094  0.02810205  0.67281044]
  [-0.38518062 -0.8371886  -0.38826686]
  [ 0.20445225  0.87811744  0.02738759]]]

如此一来，我们就可以通过调整给定蛋白质坐标的\(C_{\alpha}\)位置和旋转用于表示氨基酸内坐标的三个向量\(\vec{e_1}\)、\(\vec{e_2}\)和\(\vec{e_3}\)，以得到一个使得整体模型的Loss最低的采样器，就可以用于预测生成一系列给定氨基酸序列的蛋白质构象。