| 175 | |
| 176 | |
| 177 | def load_surface(surf_dict): |
| 178 | # Load pickled protein surface |
| 179 | surf_points = torch.tensor(surf_dict["surf_points"]).float() |
| 180 | normals = torch.tensor(surf_dict["surf_normals"]).float() |
| 181 | hks = torch.tensor(surf_dict["surf_hks"]).float() |
| 182 | curvatures = torch.tensor(surf_dict["surf_curvatures"]).float() |
| 183 | |
| 184 | num_surf_points = surf_points.shape[0] |
| 185 | edge_list = torch.zeros((1, 3), dtype=torch.long) |
| 186 | node_feature = torch.cat([hks, curvatures], dim=-1) |
| 187 | |
| 188 | surf_graph = data.Graph(edge_list=edge_list, node_feature=node_feature, bond_feature=None, |
| 189 | num_node=num_surf_points, num_relation=1) |
| 190 | with surf_graph.node(): |
| 191 | surf_graph.normals = normals |
| 192 | surf_graph.node_position = surf_points |
| 193 | |
| 194 | return surf_graph |
| 195 | |
| 196 | |
| 197 | @R.register("datasets.CATH") |