| 196 | |
| 197 | @R.register("datasets.CATH") |
| 198 | class CATH(data.ProteinDataset): |
| 199 | |
| 200 | def __init__(self, path, max_length=None, surf_path=None, transform=None): |
| 201 | path = os.path.expanduser(path) |
| 202 | self.path = path |
| 203 | self.max_length = max_length |
| 204 | |
| 205 | self.pkl_files = sorted([os.path.join(path, f) for f in os.listdir(path) if f.endswith(".pkl")]) |
| 206 | self.transform = transform |
| 207 | if surf_path: |
| 208 | surf_path = os.path.expanduser(surf_path) |
| 209 | self.surf_path = surf_path |
| 210 | |
| 211 | def truncate(self, data_dict, surf_dict=None): |
| 212 | length = data_dict["aatype"].shape[0] |
| 213 | if length <= self.max_length: |
| 214 | return data_dict, surf_dict |
| 215 | start = np.random.randint(length - self.max_length, size=(1,))[0] |
| 216 | end = start + self.max_length |
| 217 | for k in data_dict.keys(): |
| 218 | data_dict[k] = data_dict[k][start:end] |
| 219 | |
| 220 | if surf_dict is not None: |
| 221 | # Remove surfaces of the truncated part |
| 222 | surf_idx = surf_dict["res2surf"][start:end] |
| 223 | surf_mask = np.zeros(surf_dict["surf_points"].shape[0], dtype=bool) |
| 224 | surf_mask[surf_idx.flatten()] = 1 |
| 225 | for k in surf_dict.keys(): |
| 226 | if not k.startswith("surf_"): continue |
| 227 | surf_dict[k] = surf_dict[k][surf_mask] |
| 228 | # Re-index surface graph points |
| 229 | _, surf_dict["res2surf"] = np.unique(surf_idx, return_inverse=True) |
| 230 | surf_dict["res2surf"] = surf_dict["res2surf"].reshape(*surf_idx.shape) |
| 231 | |
| 232 | return data_dict, surf_dict |
| 233 | |
| 234 | def get_item(self, idx): |
| 235 | with open(self.pkl_files[idx], "rb") as fin: |
| 236 | data_dict = pickle.load(fin) |
| 237 | if self.surf_path: |
| 238 | surf_file = os.path.join(self.surf_path, os.path.basename(self.pkl_files[idx])) |
| 239 | with open(surf_file, "rb") as fin: |
| 240 | surf_dict = pickle.load(fin) |
| 241 | else: |
| 242 | surf_dict = None |
| 243 | if self.max_length: |
| 244 | data_dict, surf_dict = self.truncate(data_dict, surf_dict=surf_dict) |
| 245 | protein = load_protein(data_dict) |
| 246 | |
| 247 | item = {"graph": protein} |
| 248 | if surf_dict is not None: |
| 249 | surf_graph = load_surface(surf_dict) |
| 250 | item.update({"surf_graph": surf_graph}) |
| 251 | with protein.residue(): |
| 252 | protein.res2surf = torch.as_tensor(surf_dict["res2surf"]) # Need to transform local index to global index after batching |
| 253 | if self.transform: |
| 254 | item = self.transform(item) |
| 255 | return item |
nothing calls this directly
no outgoing calls
no test coverage detected