使用numpy计算分子内坐标
技术背景
当我们打开一个用于表示分子构象的xyz文件或者pdb文件,很容易可以理解这种基于笛卡尔坐标的空间表征方法。但是除了笛卡尔坐标表示方法之外,其实也有很多其他的方法用于粗粒化或者其他目的的表征方法,比如前一篇文章中所介绍的在AlphaFold2中所使用的残基的刚体表示方法。而这种刚体坐标,在本质上来说也是一种特殊的分子内坐标表示方法,因为对于每一个残基而言只有旋转和平移的自由度,而残基内部是保持互相之间相对静止的。换句话说,每一个残基的内坐标是保持不变的,本文主要介绍分子的内坐标表示方法。
具体表示方法
在笛卡尔坐标系中,我们使用绝对坐标来表示每一个原子的空间位置,虽然也可以用于计算分子之间的相对位置,但是如果每一次更新之后都需要重新计算一遍这个相对位置的话,在演化效率上来说会比较低。因此,我们考虑有没有一种方法,可以直接对分子的“相对位置”进行演化,直到演化结束之后,再转化回笛卡尔坐标进行可视化。其实,市面上已经有一些软件可以直接可视化这种基于“相对位置”的内坐标了,这里我们主要探讨从绝对坐标到相对坐标的算法。
绝对坐标描述的是B个体系,每个体系A个原子的三维空间坐标riri,在python编程中我们可以用张量的维度来描述这个三维坐标(B, A, D)。
相对坐标的第一个参量是原子之间的距离li,i+1li,i+1,其张量维度为(B, A-1, 1)。
相对坐标的第二个参量是原子之间的夹角ai,i+1,i+2ai,i+1,i+2,其张量维度为(B, A-2, 1)。
相对坐标的第三个参量是原子之间的二面角di,i+1,i+2,i+3di,i+1,i+2,i+3,其张量维度为(B, A-3, 1)。
最后在计算得到所有的内坐标参量之后,我们可以用concat的方法把它们按照最后一个维度进行拼接,并且在原子数维度进行扩展,最终得到一个(B, A, 3)的张量,也就是我们所需要的最终的内坐标。
代码实现
其实这个算法逻辑是很简单的,我们更多的注重一个原生算子的使用以及代码的复用。以下是几个相关的关注点:
在计算距离、角度和二面角的过程中,我们都会使用到序列原子之间的相对矢量(B, A-1, D),那么在计算过一次之后我们应该保存下来以供几个不同的函数使用。
在numpy或者是一些常用的深度学习框架中,我们最好在代码实现阶段就去避免x0x0这种情况的出现,一般在遇到除法、反三角函数或者对数函数的时候,我们可以在对应的位置加一个小量ϵ=1e−08ϵ=1e−08以避免出现nan的情况。
在计算相对矢量的时候我们一般使用的是错位相减,比如可以使用crd[1:]-crd[:-1],但是这里我们在计算过程中使用的是numpy.roll对数组进行滚动之后做减法,最后再去掉一个结果。事实上用前面的这种算法会更加简单高效一些。
# inner_crd.pyimport numpy as npnp.random.seed(1)EPSILON = 1e-08def get_vec(crd):""" Get the vector of the sequential coordinate. """# (B, A, D)crd_ = np.roll(crd, -1, axis=-2)vec = crd_ - crd# (B, A-1, D)return vec[:, :-1, :]def get_dis(crd):""" Get the distance of the sequential coordinate. """# (B, A-1, D)vec = get_vec(crd)# (B, A-1, 1)dis = np.linalg.norm(vec, axis=-1, keepdims=True)return dis, vecdef get_angle(crd):""" Get the bond angle of the sequential coordinate. """# (B, A-1, 1), (B, A-1, D)dis, vec = get_dis(crd)vec_ = np.roll(vec, -1, axis=-2)dis_ = np.roll(dis, -1, axis=-2)# (B, A-1, 1)angle = np.einsum('ijk,ijk->ij', vec, vec_)[..., None] / (dis * dis_ + EPSILON)# (B, A-2, 1), (B, A-1, 1), (B, A-1, D)return np.arccos(angle[:, :-1, :]), dis, vecdef get_dihedral(crd):""" Get the dihedrals of the sequential coordinate. """# (B, A-2, 1), (B, A-1, 1), (B, A-1, D)angle, dis, vec_0 = get_angle(crd)# (B, A-1, D)vec_1 = np.roll(vec_0, -1, axis=-2)vec_2 = np.roll(vec_1, -1, axis=-2)vec_01 = np.cross(vec_0, vec_1)vec_12 = np.cross(vec_1, vec_2)vec_01 /= np.linalg.norm(vec_01, axis=-1, keepdims=True) + EPSILONvec_12 /= np.linalg.norm(vec_12, axis=-1, keepdims=True) + EPSILON# (B, A-1, 1)dihedral = np.einsum('ijk,ijk->ij', vec_01, vec_12)[..., None]# (B, A-3, 1), (B, A-2, 1), (B, A-1, 1)return np.arccos(dihedral[:, :-2, :]), angle, disdef get_inner_crd(crd):""" Concat the distance, angles and dihedrals to get the inner coordinate. """# (B, A-3, 1), (B, A-2, 1), (B, A-1, 1)dihedral, angle, dis = get_dihedral(crd)# (B, A, 1)dihedral_ = np.pad(dihedral, ((0, 0), (3, 0), (0, 0)), mode='constant', constant_values=0)angle_ = np.pad(angle, ((0, 0), (2, 0), (0, 0)), mode='constant', constant_values=0)dis_ = np.pad(dis, ((0, 0), (1, 0), (0, 0)), mode='constant', constant_values=0)# (B, A, 3)inner_crd = np.concatenate((dis_, angle_, dihedral_), axis=-1)return inner_crdif __name__ == '__main__':B = 1A = 6D = 3# (B, A, D)origin_crd = np.random.random((B, A, D))# (B, A, 3)icrd = get_inner_crd(origin_crd)print (icrd)
上述代码执行的输出结果如下所示:
[[[0. 0. 0. ][0.59214856 0. 0. ][0.38167145 1.89801242 0. ][0.46143538 1.2138982 1.46589893][0.86899521 2.32255675 1.61009033][0.84368274 2.92999231 1.97853456]]]
这个结果就是我们所需要的分子内坐标。
总结概要
本文主要介绍了在numpy的框架下实现的分子内坐标的计算,类似的方法可以应用于MindSpore和Pytorch、Jax等深度学习相关的框架中。分子的内坐标,可以更加直观的描述分子内的相对运动,通过键长键角和二面角这三个参数。