AF3 ProteinDataset类的_patch方法解读

AlphaFold3 protein_dataset模块 ProteinDataset 类  _patch 方法的主要目标是围绕锚点残基（anchor residues）裁剪蛋白质数据，提取一个局部补丁（patch）作为模型输入。
源代码：

1.     def _patch(self, data):
2.         """Cut the data around the anchor residues."""
3.         # adapted from diffab
4.         pos_alpha = data["X"][:, 2]
5.         if self.mask_whole_chains:
6.             mask_ = (data["mask"] * data["masked_res"]).bool()
7.             anchor_points = pos_alpha[mask_].mean(0).unsqueeze(0)
8.             anchor_ind = []
9.         else:
10.             anchor_ind = self.get_anchor_ind(data["masked_res"], data["mask"])
11.             anchor_points = torch.stack([pos_alpha[ind] for ind in anchor_ind], dim=0)
12.         dist_anchor = torch.cdist(pos_alpha, anchor_points, p=2).min(dim=1)[0]  # (L, )
13.         dist_anchor[~data["mask"].bool()] = float("+inf")
14.         initial_patch_idx = torch.topk(
15.             dist_anchor,
16.             k=min(self.initial_patch_size, dist_anchor.size(0)),
17.             largest=False,
18.             sorted=True,
19.         )[
20.             1
21.         ]  # (initial_patch_size, )
22.         patch_mask = data["masked_res"].bool().clone()
23.         patch_mask[[int(x) for x in anchor_ind]] = True
24.         patch_mask[initial_patch_idx] = True
26.         if self.sabdab:
27.             antibody_mask = self._get_antibody_mask(data)
28.             antigen_mask = ~antibody_mask
29.             dist_anchor_antigen = dist_anchor.masked_fill(
30.                

复制代码

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。