欢迎关注我的CSDN：https://spike.blog.csdn.net/
本文地址：https://spike.blog.csdn.net/article/details/134333419
在蛋白质复合物结构预测中，在 TemplatePairEmbedderMultimer 层中 ，构建 Template Pair 特征的源码，即：
- 将特征
template_dgram、pseudo_beta_mask_2d、aatype_one_hot、backbone_mask_2d、unit_vector(x/y/z)特征，通过 linear 层累加到一起。 - 其中，都需要使用
multichain_mask_2d进行固定掩码，选择单链区域。 - 输出维度：
([1, 1102, 1102, 64])，linear层的输出c_out维度是 64。
源码如下：
def forward(
self,
template_dgram: torch.Tensor,
aatype_one_hot: torch.Tensor,
query_embedding: torch.Tensor,
pseudo_beta_mask: torch.Tensor,
backbone_mask: torch.Tensor,
multichain_mask_2d: torch.Tensor,
unit_vector: geometry.Vec3Array,
) -> torch.Tensor:
act = 0.0
pseudo_beta_mask_2d = (
pseudo_beta_mask[..., None] * pseudo_beta_mask[..., None, :]
)
pseudo_beta_mask_2d = pseudo_beta_mask_2d * multichain_mask_2d
template_dgram = template_dgram * pseudo_beta_mask_2d[..., None]
act += self.dgram_linear(template_dgram)
act += self.pseudo_beta_mask_linear(pseudo_beta_mask_2d[..., None])
aatype_one_hot = aatype_one_hot.to(template_dgram.dtype)
act += self.aatype_linear_1(aatype_one_hot[..., None, :, :])
act += self.aatype_linear_2(aatype_one_hot[..., None, :])
backbone_mask_2d = backbone_mask[..., None] * backbone_mask[..., None, :]
backbone_mask_2d = backbone_mask_2d * multichain_mask_2d
x, y, z = [coord * backbone_mask_2d for coord in unit_vector]
act += self.x_linear(x[..., None])
act += self.y_linear(y[..., None])
act += self.z_linear(z[..., None])
act += self.backbone_mask_linear(backbone_mask_2d[..., None])
query_embedding = self.query_embedding_layer_norm(query_embedding)
act += self.query_embedding_linear(query_embedding)
return act
template_dgram 特征：
template_dgram 特征与 multichain_mask_2d：
backbone_mask_2d 特征：
backbone_mask_2d 特征与 multichain_mask_2d：
写入特征，即：
tmp_dict = dict()
tmp_dict["pseudo_beta_mask_2d_prev"] = pseudo_beta_mask_2d.cpu().numpy()
tmp_dict["pseudo_beta_mask_2d_post"] = pseudo_beta_mask_2d.cpu().numpy()
tmp_dict["template_dgram_post"] = template_dgram.cpu().numpy()
tmp_dict["backbone_mask_2d_prev"] = backbone_mask_2d.cpu().numpy()
tmp_dict["backbone_mask_2d_post"] = backbone_mask_2d.cpu().numpy()
import pickle
with open("template_pair_embedder_multimer.pkl", "wb") as f:
pickle.dump(tmp_dict, f)
logger.info(f"[CL] saved template_pair_embedder_multimer!")
读取特征，即：
def load_tensor_dict(input_path):
"""
加载特征文件
['template_dgram', 'z', 'pseudo_beta_mask', 'backbone_mask', 'multichain_mask_2d',
'unit_vector_x', 'unit_vector_y', 'unit_vector_z']
"""
import pickle
with open(input_path, "rb") as f:
obj = pickle.load(f)
print(f"[Info] feat_dict: {
obj.keys()}")
return obj
def process_template_pair_embedder_multimer_dict(feat_dict, output_dir):
print(f"[Info] feat_dict.keys: {
feat_dict.keys()}")
draw_tensor_2d(feat_dict["pseudo_beta_mask_2d_prev"], os.path.join(output_dir, "pseudo_beta_mask_2d_prev.png"))
draw_tensor_2d(feat_dict["pseudo_beta_mask_2d_post"], os.path.join(output_dir, "pseudo_beta_mask_2d_prev.png"))
draw_template_dgram(feat_dict["template_dgram_post"], os.path.join(output_dir, "template_dgram_post.png"))
draw_tensor_2d(feat_dict["backbone_mask_2d_prev"], os.path.join(output_dir, "backbone_mask_2d_prev.png"))
draw_tensor_2d(feat_dict["backbone_mask_2d_post"], os.path.join(output_dir, "backbone_mask_2d_post.png"))
def draw_tensor_2d(feat, output_path):
"""
backbone_mask: torch.Size([1, 1102])
"""
feat = np.squeeze(feat)
f, ax_arr = plt.subplots(1, 1, figsize=(8, 5))
im = ax_arr.imshow(feat)
f.colorbar(im, ax=ax_arr)
plt.savefig(output_path, bbox_inches='tight', format='png')
plt.show()
def draw_template_dgram(feat, output_path):
"""
template_dgram: torch.Size([1, 1102, 1102, 39])
"""
f, ax_arr = plt.subplots(6, 7, figsize=(24, 15))
ax_arr = ax_arr.flatten()
feat = np.squeeze(feat)
print(f"[Info] feat: {
feat.shape}")
for i in range(0, 42):
if i <= 38:
im = ax_arr[i].imshow(feat[:, :, i], interpolation='none')
f.colorbar(im, ax=ax_arr[i])
else:
ax_arr[i].set_axis_off()
plt.savefig(output_path, bbox_inches='tight', format='png')
plt.show()