| 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: |