| 17 | |
| 18 | |
| 19 | def process(input_dir, output_dir, pkl_file): |
| 20 | with open(os.path.join(input_dir, pkl_file), "rb") as fin: |
| 21 | data_dict = pickle.load(fin) |
| 22 | output_fname = os.path.join(output_dir, pkl_file) |
| 23 | |
| 24 | num_residue = len(data_dict["aatype"]) |
| 25 | atom_position = torch.as_tensor(data_dict["atom_positions"][:, :3]).float() # take backbone atoms |
| 26 | atom_type = torch.tensor([3, 0, 0], dtype=torch.long)[None, :].repeat(num_residue, 1) # (N, CA, C) |
| 27 | atom_position = atom_position.flatten(0, 1).cuda() |
| 28 | atom_type = F.one_hot(atom_type.flatten(0, 1), num_classes=6).cuda() |
| 29 | num_atom = len(atom_position) |
| 30 | batch = torch.zeros((num_atom,), dtype=torch.long).cuda() |
| 31 | |
| 32 | surf_points, surf_normals, _ = surface.atoms_to_points_normals(atom_position, batch, atomtypes=atom_type) |
| 33 | num_surf_points = len(surf_points) |
| 34 | |
| 35 | # Surface -> residue graph correspondence (num_residue, 21) each element is a surface point index |
| 36 | res2surf, _ = surface.knn_atoms(atom_position, surf_points, k=20) |
| 37 | res2surf = res2surf.view(num_residue, 3, -1) |
| 38 | |
| 39 | batch_surf = torch.zeros((num_surf_points,), dtype=torch.long) #.cuda() |
| 40 | surf_curvatures = surface.compute_curvatures(surf_points, surf_normals, batch=batch_surf, curvature_scales=[1.0, 2.0, 3.0, 5.0, 10.0]) |
| 41 | |
| 42 | surf_points = surf_points.cpu().detach().numpy() |
| 43 | eigs_ratio = 0.01 if num_surf_points > 20000 else 0.06 |
| 44 | surf_eig_vals, surf_eig_vecs, surf_eig_vecs_inv = surface.compute_eigens(num_surf_points, surf_points, min_n_eigs=50, eigs_ratio=eigs_ratio) |
| 45 | |
| 46 | surf_hks = surface.compute_HKS(surf_eig_vecs, surf_eig_vals, num_t=32, t_min=0.1, t_max=1000, scale=1000) |
| 47 | |
| 48 | surf_data_dict = { |
| 49 | "surf_points": surf_points.astype(np.float32), |
| 50 | "surf_normals": surf_normals.cpu().detach().numpy().astype(np.float32), |
| 51 | "surf_hks": surf_hks.astype(np.float32), |
| 52 | "surf_curvatures": surf_curvatures.cpu().detach().numpy().astype(np.float32), |
| 53 | "res2surf": res2surf.cpu().detach().numpy(), |
| 54 | } |
| 55 | |
| 56 | with open(output_fname, "wb") as fout: |
| 57 | pickle.dump(surf_data_dict, fout) |
| 58 | |
| 59 | |
| 60 | if __name__ == "__main__": |