| 80 | return len(self.mutated_sequences) |
| 81 | |
| 82 | def truncate(self, sequence_graph, structure_graph, surface_graph=None): |
| 83 | num_residue = structure_graph.num_residue |
| 84 | start = sequence_graph.start - structure_graph.start |
| 85 | end = sequence_graph.end - structure_graph.start |
| 86 | residue_mask = torch.zeros((num_residue, ), dtype=torch.bool) |
| 87 | residue_mask[start:end] = 1 |
| 88 | structure_graph = structure_graph.subresidue(residue_mask) |
| 89 | |
| 90 | if surface_graph: |
| 91 | surf_idx = structure_graph.res2surf |
| 92 | surf_mask = torch.zeros((surface_graph.num_node, ), dtype=torch.bool) |
| 93 | surf_mask[surf_idx.flatten()] = 1 |
| 94 | surface_graph = surface_graph.subgraph(surf_mask) |
| 95 | |
| 96 | _, res2surf = torch.unique(surf_idx, return_inverse=True) |
| 97 | with structure_graph.residue(): |
| 98 | structure_graph.res2surf = res2surf.view(*surf_idx.shape) |
| 99 | |
| 100 | return sequence_graph, structure_graph, surface_graph |
| 101 | |
| 102 | def assign_structure(self, sequence_graph, structure_graph): |
| 103 | graph = structure_graph.clone() |