| 69 | |
| 70 | |
| 71 | class MutantDataset(torch_data.Dataset): |
| 72 | |
| 73 | def __init__(self, mutated_sequences, wild_type, surf_graph=None, transform=None): |
| 74 | self.mutated_sequences = mutated_sequences |
| 75 | self.wild_type = wild_type |
| 76 | self.surf_graph = surf_graph |
| 77 | self.transform = transform |
| 78 | |
| 79 | def __len__(self): |
| 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() |
| 104 | # Assume the backbone structure won't change |
| 105 | assert graph.num_residue == sequence_graph.num_residue |
| 106 | with graph.residue(): |
| 107 | graph.residue_type = sequence_graph.residue_type |
| 108 | return graph |
| 109 | |
| 110 | def __getitem__(self, index): |
| 111 | sequence_graph = self.mutated_sequences[index] |
| 112 | structure_graph = self.wild_type |
| 113 | surface_graph = self.surf_graph |
| 114 | |
| 115 | # we need to truncate structure if structure is longer than sequence |
| 116 | if structure_graph.start <= sequence_graph.start and structure_graph.end >= sequence_graph.end: |
| 117 | sequence_graph, structure_graph, surface_graph = self.truncate(sequence_graph, structure_graph, surface_graph) |
| 118 | elif structure_graph.start != sequence_graph.start or structure_graph.end != sequence_graph.end: |
| 119 | raise ValueError("the structure range (%d, %d) doesn't match the sequence range (%d, %d)" % |
| 120 | (structure_graph.start, structure_graph.end, sequence_graph.start, sequence_graph.end)) |
| 121 | |
| 122 | graph = self.assign_structure(sequence_graph, structure_graph) |
| 123 | item = {"graph": graph} |
| 124 | if surface_graph: |
| 125 | item["surf_graph"] = surface_graph |
| 126 | if self.transform: |
| 127 | item = self.transform(item) |
| 128 | return item |
nothing calls this directly
no outgoing calls
no test coverage detected