Spaces:
Sleeping
Sleeping
| """ | |
| https://github.com/ProteinDesignLab/protpardelle | |
| License: MIT | |
| Author: Alex Chu | |
| Various utils for handling protein data. | |
| """ | |
| import os | |
| import shlex | |
| import subprocess | |
| import sys | |
| import torch | |
| import yaml | |
| import argparse | |
| from einops import rearrange, repeat | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| import Bio | |
| from Bio.PDB.DSSP import DSSP | |
| from core import protein | |
| from core import protein_mpnn | |
| from core import residue_constants | |
| PATH_TO_TMALIGN = "/home/alexechu/essentials_kit/ml_utils/align/TMalign/TMalign" | |
| ################ STRUCTURE/FORMAT UTILS ############################# | |
| def aatype_to_seq(aatype, seq_mask=None): | |
| if seq_mask is None: | |
| seq_mask = torch.ones_like(aatype) | |
| mapping = residue_constants.restypes_with_x | |
| mapping = mapping + ["<mask>"] | |
| unbatched = False | |
| if len(aatype.shape) == 1: | |
| unbatched = True | |
| aatype = [aatype] | |
| seq_mask = [seq_mask] | |
| seqs = [] | |
| for i, ai in enumerate(aatype): | |
| seq = [] | |
| for j, aa in enumerate(ai): | |
| if seq_mask[i][j] == 1: | |
| try: | |
| seq.append(mapping[aa]) | |
| except IndexError: | |
| print(aatype[i]) | |
| raise Exception(f"Error in mapping {aa} at {i},{j}") | |
| seqs.append("".join(seq)) | |
| if unbatched: | |
| seqs = seqs[0] | |
| return seqs | |
| def seq_to_aatype(seq, num_tokens=21): | |
| if num_tokens == 20: | |
| mapping = residue_constants.restype_order | |
| if num_tokens == 21: | |
| mapping = residue_constants.restype_order_with_x | |
| if num_tokens == 22: | |
| mapping = residue_constants.restype_order_with_x | |
| mapping["<mask>"] = 21 | |
| return torch.Tensor([mapping[aa] for aa in seq]).long() | |
| def batched_seq_to_aatype_and_mask(seqs, max_len=None): | |
| if max_len is None: | |
| max_len = max([len(s) for s in seqs]) | |
| aatypes = [] | |
| seq_mask = [] | |
| for s in seqs: | |
| pad_size = max_len - len(s) | |
| aatype = seq_to_aatype(s) | |
| aatypes.append(F.pad(aatype, (0, pad_size))) | |
| mask = torch.ones_like(aatype).float() | |
| seq_mask.append(F.pad(mask, (0, pad_size))) | |
| return torch.stack(aatypes), torch.stack(seq_mask) | |
| def atom37_mask_from_aatype(aatype, seq_mask=None): | |
| # source_mask is (21,37) originally | |
| source_mask = torch.Tensor(residue_constants.restype_atom37_mask).to(aatype.device) | |
| bb_atoms = source_mask[residue_constants.restype_order["G"]][None] | |
| # Use only the first 20 plus bb atoms for X, mask | |
| source_mask = torch.cat([source_mask[:-1], bb_atoms, bb_atoms], 0) | |
| atom_mask = source_mask[aatype] | |
| if seq_mask is not None: | |
| atom_mask *= seq_mask[..., None] | |
| return atom_mask | |
| def atom37_coords_from_atom14(atom14_coords, aatype, return_mask=False): | |
| # Unbatched | |
| device = atom14_coords.device | |
| atom37_coords = torch.zeros((atom14_coords.shape[0], 37, 3)).to(device) | |
| for i in range(atom14_coords.shape[0]): # per residue | |
| aa = aatype[i] | |
| aa_3name = residue_constants.restype_1to3[residue_constants.restypes[aa]] | |
| atom14_atoms = residue_constants.restype_name_to_atom14_names[aa_3name] | |
| for j in range(14): | |
| atom_name = atom14_atoms[j] | |
| if atom_name != "": | |
| atom37_idx = residue_constants.atom_order[atom_name] | |
| atom37_coords[i, atom37_idx, :] = atom14_coords[i, j, :] | |
| if return_mask: | |
| atom37_mask = atom37_mask_from_aatype(aatype) | |
| return atom37_coords, atom37_mask | |
| return atom37_coords | |
| def atom73_mask_from_aatype(aatype, seq_mask=None): | |
| source_mask = torch.Tensor(residue_constants.restype_atom73_mask).to(aatype.device) | |
| atom_mask = source_mask[aatype] | |
| if seq_mask is not None: | |
| atom_mask *= seq_mask[..., None] | |
| return atom_mask | |
| def atom37_to_atom73(atom37, aatype, return_mask=False): | |
| # Unbatched | |
| atom73 = torch.zeros((atom37.shape[0], 73, 3)).to(atom37) | |
| for i in range(atom37.shape[0]): | |
| aa = aatype[i] | |
| aa1 = residue_constants.restypes[aa] | |
| for j, atom37_name in enumerate(residue_constants.atom_types): | |
| atom73_name = atom37_name | |
| if atom37_name not in ["N", "CA", "C", "O", "CB"]: | |
| atom73_name = aa1 + atom73_name | |
| if atom73_name in residue_constants.atom73_names_to_idx: | |
| atom73_idx = residue_constants.atom73_names_to_idx[atom73_name] | |
| atom73[i, atom73_idx, :] = atom37[i, j, :] | |
| if return_mask: | |
| atom73_mask = atom73_mask_from_aatype(aatype) | |
| return atom73, atom73_mask | |
| return atom73 | |
| def atom73_to_atom37(atom73, aatype, return_mask=False): | |
| # Unbatched | |
| atom37_coords = torch.zeros((atom73.shape[0], 37, 3)).to(atom73) | |
| for i in range(atom73.shape[0]): # per residue | |
| aa = aatype[i] | |
| aa1 = residue_constants.restypes[aa] | |
| for j, atom_type in enumerate(residue_constants.atom_types): | |
| atom73_name = atom_type | |
| if atom73_name not in ["N", "CA", "C", "O", "CB"]: | |
| atom73_name = aa1 + atom73_name | |
| if atom73_name in residue_constants.atom73_names_to_idx: | |
| atom73_idx = residue_constants.atom73_names_to_idx[atom73_name] | |
| atom37_coords[i, j, :] = atom73[i, atom73_idx, :] | |
| if return_mask: | |
| atom37_mask = atom37_mask_from_aatype(aatype) | |
| return atom37_coords, atom37_mask | |
| return atom37_coords | |
| def get_dmap(pdb, atoms=["N", "CA", "C", "O"], batched=True, out="torch", device=None): | |
| def _dmap_from_coords(coords): | |
| coords = coords.contiguous() | |
| dmaps = torch.cdist(coords, coords).unsqueeze(1) | |
| if out == "numpy": | |
| return dmaps.detach().cpu().numpy() | |
| elif out == "torch": | |
| if device is not None: | |
| return dmaps.to(device) | |
| else: | |
| return dmaps | |
| if isinstance(pdb, str): # input is pdb file | |
| coords = load_coords_from_pdb(pdb, atoms=atoms).view(1, -1, 3) | |
| return _dmap_from_coords(coords) | |
| elif len(pdb.shape) == 2: # single set of coords | |
| if isinstance(pdb, np.ndarray): | |
| pdb = torch.Tensor(pdb) | |
| return _dmap_from_coords(pdb.unsqueeze(0)) | |
| elif len(pdb.shape) == 3 and batched: | |
| return _dmap_from_coords(pdb) | |
| elif len(pdb.shape) == 3 and not batched: | |
| return _dmap_from_coords(pdb.view(1, -1, 3)) | |
| elif len(pdb.shape) == 4: | |
| return _dmap_from_coords(pdb.view(pdb.size(0), -1, 3)) | |
| def get_channeled_dmap(coords): | |
| # coords is b, nres, natom, 3 | |
| coords = coords.permute(0, 2, 1, 3) | |
| dvecs = coords[..., None, :] - coords[..., None, :, :] # b, natom, nres, nres, 3 | |
| dists = torch.sqrt(dvecs.pow(2).sum(-1) + 1e-8) | |
| return dists | |
| def fill_in_cbeta_for_atom37(coords): | |
| b = coords[..., 1, :] - coords[..., 0, :] | |
| c = coords[..., 2, :] - coords[..., 1, :] | |
| a = torch.cross(b, c, dim=-1) | |
| cbeta = -0.58273431 * a + 0.56802827 * b - 0.54067466 * c + coords[..., 1, :] | |
| new_coords = torch.clone(coords) | |
| new_coords[..., 3, :] = cbeta | |
| return new_coords | |
| def get_distogram(coords, n_bins=20, start=2, return_onehot=True, seq_mask=None): | |
| # coords is b, nres, natom, 3 | |
| # distogram for cb atom (assume 3rd atom) | |
| coords_with_cb = fill_in_cbeta_for_atom37(coords) | |
| dists = get_channeled_dmap(coords_with_cb[:, :, 3:4]).squeeze(1) | |
| bins = torch.arange(start, start + n_bins - 1).to(dists.device) | |
| dgram = torch.bucketize(dists, bins) | |
| dgram_oh = F.one_hot(dgram, n_bins) | |
| if seq_mask is not None: | |
| mask_2d = seq_mask[:, :, None] * seq_mask[:, None, :] | |
| dgram = dgram * mask_2d | |
| dgram_oh = dgram_oh * mask_2d[..., None] | |
| if return_onehot: | |
| return dgram_oh | |
| return dgram | |
| def get_contacts(coords=None, distogram=None, seq_mask=None): | |
| if distogram is None: | |
| distogram = get_distogram(coords) | |
| contacts = (distogram.argmax(-1) < 6).float() | |
| if seq_mask is not None: | |
| contacts *= seq_mask[..., None] * seq_mask[..., None, :] | |
| return contacts | |
| def dihedral(a, b, c, d): | |
| # inputs can be (1,3), (n,3), or (bs,n,3) | |
| b1 = a - b | |
| b2 = b - c | |
| b3 = c - d | |
| n1 = F.normalize(torch.cross(b1, b2), dim=-1) | |
| n2 = F.normalize(torch.cross(b2, b3), dim=-1) | |
| m1 = torch.cross(n1, b2 / b2.norm(dim=-1).unsqueeze(-1)) | |
| y = (m1 * n2).sum(dim=-1) | |
| x = (n1 * n2).sum(dim=-1) | |
| return torch.atan2(y, x) | |
| def get_torsions_from_coords( | |
| coords, atoms=["N", "CA", "C", "O"], batched=True, out="torch", device=None | |
| ): | |
| """ | |
| Returns a n-dim array of shape (bs, nres, ntors), where ntors is the | |
| number of torsion angles (e.g. 2 if using phi and psi), with units of radians. | |
| """ | |
| if isinstance(coords, np.ndarray): | |
| coords = torch.Tensor(coords) | |
| if len(coords.shape) == 2: | |
| coords = coords.unsqueeze(0) | |
| if len(coords.shape) == 4: | |
| coords = coords.view(coords.size(0), -1, 3) | |
| if len(coords.shape) == 3 and not batched: | |
| coords = coords.view(1, -1, 3) | |
| if len(coords.shape) == 3: | |
| bs = coords.size(0) | |
| if "O" in atoms: | |
| idxs = [ | |
| i for i in range(coords.size(1)) if i % 4 != 3 | |
| ] # deselect O atoms for N-Ca-C-O coords | |
| coords = coords[:, idxs, :] | |
| a, b, c, d = ( | |
| coords[:, :-3, :], | |
| coords[:, 1:-2, :], | |
| coords[:, 2:-1, :], | |
| coords[:, 3:, :], | |
| ) | |
| torsions = dihedral( | |
| a, b, c, d | |
| ) # output order is psi-omega-phi, reorganize to (bs, nres, 3) | |
| torsions = torsions.view(bs, torsions.size(1) // 3, 3) | |
| omegaphi = torch.cat( | |
| (torch.zeros(bs, 1, 2).to(coords.device), torsions[:, :, 1:]), 1 | |
| ) | |
| psi = torch.cat((torsions[:, :, 0], torch.zeros(bs, 1).to(coords.device)), 1) | |
| torsions = torch.cat( | |
| ( | |
| omegaphi[:, :, 1].unsqueeze(-1), | |
| psi.unsqueeze(-1), | |
| omegaphi[:, :, 0].unsqueeze(-1), | |
| ), | |
| -1, | |
| ) | |
| else: | |
| raise Exception("input coords not of correct dims") | |
| if out == "numpy": | |
| return torsions.detach().cpu().numpy() | |
| elif out == "torch": | |
| if device is not None: | |
| return torsions.to(device) | |
| else: | |
| return torsions | |
| def get_trig_from_torsions(torsions, out="torch", device=None): | |
| """ | |
| Calculate unit circle projections from coords input. | |
| Returns a n-dim array of shape (bs, nres, ntors, 2), where ntors is the | |
| number of torsion angles (e.g. 2 if using phi and psi), and the last | |
| dimension is the xy unit-circle coordinates of the corresponding angle. | |
| """ | |
| if isinstance(torsions, np.ndarray): | |
| torsions = torch.Tensor(torsions) | |
| x = torsions.cos() | |
| y = torsions.sin() | |
| trig = torch.cat((x.unsqueeze(-1), y.unsqueeze(-1)), -1) | |
| if out == "numpy": | |
| return trig.detach().cpu().numpy() | |
| elif out == "torch": | |
| if device is not None: | |
| return trig.to(device) | |
| else: | |
| return trig | |
| def get_abego_string_from_torsions(torsions): | |
| A_bin = (-75, 50) | |
| G_bin = (-100, 100) | |
| torsions = torsions * 180.0 / np.pi | |
| phi, psi = torsions[:, :, 0], torsions[:, :, 1] | |
| abego_vec = np.zeros((torsions.size(0), torsions.size(1))).astype(str) | |
| A = (phi <= 0) & (psi <= A_bin[1]) & (psi > A_bin[0]) | |
| B = (phi <= 0) & ((psi > A_bin[1]) | (psi <= A_bin[0])) | |
| G = (phi > 0) & (psi <= G_bin[1]) & (psi > G_bin[0]) | |
| E = (phi > 0) & ((psi > G_bin[1]) | (psi <= G_bin[0])) | |
| abego_vec[A] = "A" | |
| abego_vec[B] = "B" | |
| abego_vec[G] = "G" | |
| abego_vec[E] = "E" | |
| abego_strs = ["".join(v) for v in abego_vec] | |
| return abego_strs | |
| def get_bond_lengths_from_coords(coords, batched=True, out="torch", device=None): | |
| """ | |
| Returns array of shape (bs, n_res, 4), where final dim is bond lengths | |
| in order of N-Ca, Ca-C, C-O, C-N (none for last residue) | |
| """ | |
| if isinstance(coords, np.ndarray): | |
| coords = torch.Tensor(coords) | |
| if len(coords.shape) == 2: | |
| coords = coords.unsqueeze(0) | |
| if len(coords.shape) == 3 and not batched: | |
| coords = coords.view(1, -1, 3) | |
| if len(coords.shape) == 4: | |
| coords = coords.view(coords.size(0), -1, 3) | |
| N = coords[:, ::4, :] | |
| Ca = coords[:, 1::4, :] | |
| C = coords[:, 2::4, :] | |
| O = coords[:, 3::4, :] | |
| NCa = (Ca - N).norm(dim=-1).unsqueeze(-1) | |
| CaC = (C - Ca).norm(dim=-1).unsqueeze(-1) | |
| CO = (O - C).norm(dim=-1).unsqueeze(-1) | |
| CN = (N[:, 1:] - C[:, :-1]).norm(dim=-1) | |
| CN = torch.cat([CN, torch.zeros(CN.size(0), 1).to(CN.device)], 1).unsqueeze(-1) | |
| blengths = torch.cat((NCa, CaC, CO, CN), -1) | |
| if out == "numpy": | |
| return blengths.detach().cpu().numpy() | |
| elif out == "torch": | |
| if device is not None: | |
| return blengths.to(device) | |
| else: | |
| return blengths | |
| def get_bond_angles_from_coords(coords, batched=True, out="torch", device=None): | |
| """ | |
| Returns array of shape (bs, n_res, 5), where final dim is bond angles | |
| in order of N-Ca-C, Ca-C-O, Ca-C-N, O-C-N, C-N-Ca (none for last residue) | |
| """ | |
| def _angle(v1, v2): | |
| cos = (v1 * v2).sum(-1) / (v1.norm(dim=-1) * v2.norm(dim=-1)) | |
| return cos.acos() | |
| if isinstance(coords, np.ndarray): | |
| coords = torch.Tensor(coords) | |
| if len(coords.shape) == 2: | |
| coords = coords.unsqueeze(0) | |
| if len(coords.shape) == 3 and not batched: | |
| coords = coords.view(1, -1, 3) | |
| if len(coords.shape) == 4: | |
| coords = coords.view(coords.size(0), -1, 3) | |
| N = coords[:, ::4, :] | |
| Nnext = coords[:, 4::4, :] | |
| Ca = coords[:, 1::4, :] | |
| Canext = coords[:, 5::4, :] | |
| C = coords[:, 2::4, :] | |
| O = coords[:, 3::4, :] | |
| CaN = N - Ca | |
| CaC = C - Ca | |
| CCa = Ca - C | |
| CO = O - C | |
| CNnext = Nnext - C[:, :-1, :] | |
| NnextC = -1 * CNnext | |
| NnextCanext = Canext - Nnext | |
| NCaC = _angle(CaN, CaC).unsqueeze(-1) | |
| CaCO = _angle(CCa, CO).unsqueeze(-1) | |
| CaCN = _angle(CCa[:, :-1], CNnext).unsqueeze(-1) | |
| CaCN = _extend(CaCN) | |
| OCN = _angle(CO[:, :-1], CNnext).unsqueeze(-1) | |
| OCN = _extend(OCN) | |
| CNCa = _angle(NnextC, NnextCanext).unsqueeze(-1) | |
| # CNCa = torch.cat([CNCa, torch.zeros(CNCa.size(0), 1).to(CNCa.device)], 1).unsqueeze(-1) | |
| CNCa = _extend(CNCa) | |
| bangles = torch.cat((NCaC, CaCO, CaCN, OCN, CNCa), -1) | |
| if out == "numpy": | |
| return bangles.detach().cpu().numpy() | |
| elif out == "torch": | |
| if device is not None: | |
| return bangles.to(device) | |
| else: | |
| return bangles | |
| def get_buried_positions_mask(coords, seq_mask=None, threshold=6.0): | |
| ca_idx = residue_constants.atom_order["CA"] # typically 1 | |
| cb_idx = residue_constants.atom_order["CB"] # typically 3 | |
| if seq_mask is None: | |
| seq_mask = torch.ones_like(coords)[..., 0, 0] | |
| coords = fill_in_cbeta_for_atom37(coords) | |
| # get 8 closest neighbors by CB | |
| neighbor_coords = coords[:, :, cb_idx] | |
| ca_neighbor_dists, edge_index = protein_mpnn.get_closest_neighbors( | |
| neighbor_coords, seq_mask, 9 | |
| ) | |
| edge_index = edge_index[..., 1:].contiguous() | |
| # compute avg CB distance | |
| cb_coords = coords[:, :, cb_idx] | |
| neighbor_cb = protein_mpnn.gather_nodes(cb_coords, edge_index) | |
| avg_cb_dist = (neighbor_cb - cb_coords[..., None, :]).pow(2).sum(-1).sqrt().mean(-1) | |
| buried_positions_mask = (avg_cb_dist < threshold).float() * seq_mask | |
| return buried_positions_mask | |
| def get_fullatom_bond_lengths_from_coords( | |
| coords, aatype, atom_mask=None, return_format="per_aa" | |
| ): | |
| # Also return sidechain bond angles. All unbatched. return list of dicts | |
| def dist(xyz1, xyz2): | |
| return (xyz1 - xyz2).pow(2).sum().sqrt().detach().cpu().item() | |
| assert aatype.max() <= 19 | |
| seq = aatype_to_seq(aatype) | |
| # residue-wise list of dicts [{'N-CA': a, 'CA-C': b}, {'N-CA': a, 'CA-C': b}] | |
| all_bond_lens_by_pos = [] | |
| # aa-wise dict of dicts of lists {'A': {'N-CA': [a, b, c], 'CA-C': [a, b, c]}} | |
| all_bond_lens_by_aa = {aa: {} for aa in residue_constants.restypes} | |
| for i, res in enumerate(coords): | |
| aa3 = residue_constants.restype_1to3[seq[i]] | |
| res_bond_lens = {} | |
| for j, atom1 in enumerate(residue_constants.atom_types): | |
| for k, atom2 in enumerate(residue_constants.atom_types): | |
| if j < k and protein.are_atoms_bonded(aa3, atom1, atom2): | |
| if atom_mask is None or ( | |
| atom_mask[i, j] > 0.5 and atom_mask[i, k] > 0.5 | |
| ): | |
| bond_name = f"{atom1}-{atom2}" | |
| bond_len = dist(res[j], res[k]) | |
| res_bond_lens[bond_name] = bond_len | |
| all_bond_lens_by_pos.append(res_bond_lens) | |
| for key, val in res_bond_lens.items(): | |
| all_bond_lens_by_aa[seq[i]].setdefault(key, []).append(val) | |
| if return_format == "per_aa": | |
| return all_bond_lens_by_aa | |
| elif return_format == "per_position": | |
| return all_bond_lens_by_pos | |
| def batched_fullatom_bond_lengths_from_coords( | |
| coords, aatype, atom_mask=None, return_format="per_aa" | |
| ): | |
| # Expects trimmed coords (no mask) | |
| if return_format == "per_position": | |
| batched_bond_lens = [] | |
| elif return_format == "per_aa": | |
| batched_bond_lens = {aa: {} for aa in residue_constants.restypes} | |
| for i, c in enumerate(coords): | |
| atom_mask_i = None if atom_mask is None else atom_mask[i] | |
| bond_lens = get_fullatom_bond_lengths_from_coords( | |
| c, aatype[i], atom_mask=atom_mask_i, return_format=return_format | |
| ) | |
| if return_format == "per_position": | |
| batched_bond_lens.extend(bond_lens) | |
| elif return_format == "per_aa": | |
| for aa, d in bond_lens.items(): | |
| for bond, lengths in d.items(): | |
| batched_bond_lens[aa].setdefault(bond, []).extend(lengths) | |
| return batched_bond_lens | |
| def batched_fullatom_bond_angles_from_coords(coords, aatype, return_format="per_aa"): | |
| # Expects trimmed coords (no mask) | |
| if return_format == "per_position": | |
| batched_bond_angles = [] | |
| elif return_format == "per_aa": | |
| batched_bond_angles = {aa: {} for aa in residue_constants.restypes} | |
| for i, c in enumerate(coords): | |
| bond_angles = get_fullatom_bond_angles_from_coords( | |
| c, aatype[i], return_format=return_format | |
| ) | |
| if return_format == "per_position": | |
| batched_bond_angles.extend(bond_angles) | |
| elif return_format == "per_aa": | |
| for aa, d in bond_angles.items(): | |
| for bond, lengths in d.items(): | |
| batched_bond_angles[aa].setdefault(bond, []).extend(lengths) | |
| return batched_bond_angles | |
| def get_chi_angles(coords, aatype, atom_mask=None, seq_mask=None): | |
| # unbatched | |
| # return (n, 4) chis in degrees and mask | |
| chis = [] | |
| chi_mask = [] | |
| atom_order = residue_constants.atom_order | |
| seq = aatype_to_seq(aatype, seq_mask=seq_mask) | |
| for i, aa1 in enumerate(seq): # per residue | |
| if seq_mask is not None and seq_mask[i] == 0: | |
| chis.append([0, 0, 0, 0]) | |
| chi_mask.append([0, 0, 0, 0]) | |
| else: | |
| chi = [] | |
| mask = [] | |
| chi_atoms = residue_constants.chi_angles_atoms[ | |
| residue_constants.restype_1to3[aa1] | |
| ] | |
| for j in range(4): # per chi angle | |
| if j > len(chi_atoms) - 1: | |
| chi.append(0) | |
| mask.append(0) | |
| elif atom_mask is not None and any( | |
| [atom_mask[i, atom_order[a]] < 0.5 for a in chi_atoms[j]] | |
| ): | |
| chi.append(0) | |
| mask.append(0) | |
| else: | |
| # Four atoms per dihedral | |
| xyz4 = [coords[i, atom_order[a]] for a in chi_atoms[j]] | |
| angle = dihedral(*xyz4) * 180 / np.pi | |
| chi.append(angle) | |
| mask.append(1) | |
| chis.append(chi) | |
| chi_mask.append(mask) | |
| chis = torch.Tensor(chis) | |
| chi_mask = torch.Tensor(chi_mask) | |
| return chis, chi_mask | |
| def fill_Os_from_NCaC_coords( | |
| coords: torch.Tensor, out: str = "torch", device: str = None | |
| ): | |
| """Given NCaC coords, add O atom coordinates in. | |
| (bs, 3n, 3) -> (bs, 4n, 3) | |
| """ | |
| CO_LEN = 1.231 | |
| if len(coords.shape) == 2: | |
| coords = coords.unsqueeze(0) | |
| Cs = coords[:, 2:-1:3, :] # all but last C | |
| CCa_norm = F.normalize(coords[:, 1:-2:3, :] - Cs, dim=-1) # all but last Ca | |
| CN_norm = F.normalize(coords[:, 3::3, :] - Cs, dim=-1) # all but first N | |
| Os = F.normalize(CCa_norm + CN_norm, dim=-1) * -CO_LEN | |
| Os += Cs | |
| # TODO place C-term O atom properly | |
| Os = torch.cat([Os, coords[:, -1, :].view(-1, 1, 3) + 1], 1) | |
| coords_out = [] | |
| for i in range(Os.size(1)): | |
| coords_out.append(coords[:, i * 3 : (i + 1) * 3, :]) | |
| coords_out.append(Os[:, i, :].view(-1, 1, 3)) | |
| coords_out = torch.cat(coords_out, 1) | |
| if out == "numpy": | |
| return coords_out.detach().cpu().numpy() | |
| elif out == "torch": | |
| if device is not None: | |
| return coords_out.to(device) | |
| else: | |
| return coords_out | |
| def _extend(x, axis=1, n=1, prepend=False): | |
| # Add an extra zeros 'residue' to the end (or beginning, prepend=True) of a Tensor | |
| # Used to extend torsions when there is no 'psi' for last residue | |
| shape = list(x.shape) | |
| shape[axis] = n | |
| if prepend: | |
| return torch.cat([torch.zeros(shape).to(x.device), x], axis) | |
| else: | |
| return torch.cat([x, torch.zeros(shape).to(x.device)], axis) | |
| def trim_coords(coords, n_res, batched=True): | |
| if batched: # Return list of tensors | |
| front = (coords.shape[1] - n_res) // 2 | |
| return [ | |
| coords[i, front[i] : front[i] + n_res[i]] for i in range(coords.shape[0]) | |
| ] | |
| else: | |
| if isinstance(n_res, torch.Tensor): | |
| n_res = n_res.int() | |
| front_pad = (coords.shape[0] - n_res) // 2 | |
| return coords[front_pad : front_pad + n_res] | |
| def batch_align_on_calpha(x, y): | |
| aligned_x = [] | |
| for i, xi in enumerate(x): | |
| xi_calpha = xi[:, 1, :] | |
| _, (R, t) = kabsch_align(xi_calpha, y[i, :, 1, :]) | |
| xi_ctr = xi - xi_calpha.mean(0, keepdim=True) | |
| xi_aligned = xi_ctr @ R.t() + t | |
| aligned_x.append(xi_aligned) | |
| return torch.stack(aligned_x) | |
| def kabsch_align(p, q): | |
| if len(p.shape) > 2: | |
| p = p.reshape(-1, 3) | |
| if len(q.shape) > 2: | |
| q = q.reshape(-1, 3) | |
| p_ctr = p - p.mean(0, keepdim=True) | |
| t = q.mean(0, keepdim=True) | |
| q_ctr = q - t | |
| H = p_ctr.t() @ q_ctr | |
| U, S, V = torch.svd(H) | |
| R = V @ U.t() | |
| I_ = torch.eye(3).to(p) | |
| I_[-1, -1] = R.det().sign() | |
| R = V @ I_ @ U.t() | |
| p_aligned = p_ctr @ R.t() + t | |
| return p_aligned, (R, t) | |
| def get_dssp_string(pdb): | |
| try: | |
| structure = Bio.PDB.PDBParser(QUIET=True).get_structure(pdb[:-3], pdb) | |
| dssp = DSSP(structure[0], pdb, dssp="mkdssp") | |
| dssp_string = "".join([dssp[k][2] for k in dssp.keys()]) | |
| return dssp_string | |
| except Exception as e: | |
| print(e) | |
| return None | |
| def pool_dssp_symbols(dssp_string, newchar=None, chars=["-", "T", "S", "C", " "]): | |
| """Replaces all instances of chars with newchar. DSSP chars are helix=GHI, strand=EB, loop=- TSC""" | |
| if newchar is None: | |
| newchar = chars[0] | |
| string_out = dssp_string | |
| for c in chars: | |
| string_out = string_out.replace(c, newchar) | |
| return string_out | |
| def get_3state_dssp(pdb=None, coords=None): | |
| if coords is not None: | |
| pdb = "temp_dssp.pdb" | |
| write_coords_to_pdb(coords, pdb, batched=False) | |
| dssp_string = get_dssp_string(pdb) | |
| if dssp_string is not None: | |
| dssp_string = pool_dssp_symbols(dssp_string, newchar="L") | |
| dssp_string = pool_dssp_symbols(dssp_string, chars=["H", "G", "I"]) | |
| dssp_string = pool_dssp_symbols(dssp_string, chars=["E", "B"]) | |
| if coords is not None: | |
| subprocess.run(shlex.split(f"rm {pdb}")) | |
| return dssp_string | |
| ############## SAVE/LOAD UTILS ################################# | |
| def load_feats_from_pdb( | |
| pdb, bb_atoms=["N", "CA", "C", "O"], load_atom73=False, **kwargs | |
| ): | |
| feats = {} | |
| with open(pdb, "r") as f: | |
| pdb_str = f.read() | |
| protein_obj = protein.from_pdb_string(pdb_str, **kwargs) | |
| bb_idxs = [residue_constants.atom_order[a] for a in bb_atoms] | |
| bb_coords = torch.from_numpy(protein_obj.atom_positions[:, bb_idxs]) | |
| feats["bb_coords"] = bb_coords.float() | |
| for k, v in vars(protein_obj).items(): | |
| feats[k] = torch.Tensor(v) | |
| feats["aatype"] = feats["aatype"].long() | |
| if load_atom73: | |
| feats["atom73_coords"], feats["atom73_mask"] = atom37_to_atom73( | |
| feats["atom_positions"], feats["aatype"], return_mask=True | |
| ) | |
| return feats | |
| def load_coords_from_pdb( | |
| pdb, | |
| atoms=["N", "CA", "C", "O"], | |
| method="raw", | |
| also_bfactors=False, | |
| normalize_bfactors=True, | |
| ): | |
| """Returns array of shape (1, n_res, len(atoms), 3)""" | |
| coords = [] | |
| bfactors = [] | |
| if method == "raw": # Raw numpy implementation, faster than biopdb | |
| # Indexing into PDB format, allowing XXXX.XXX | |
| coords_in_pdb = [slice(30, 38), slice(38, 46), slice(46, 54)] | |
| # Indexing into PDB format, allowing XXX.XX | |
| bfactor_in_pdb = slice(60, 66) | |
| with open(pdb, "r") as f: | |
| resi_prev = 1 | |
| counter = 0 | |
| for l in f: | |
| l_split = l.rstrip("\n").split() | |
| if len(l_split) > 0 and l_split[0] == "ATOM" and l_split[2] in atoms: | |
| resi = l_split[5] | |
| if resi == resi_prev: | |
| counter += 1 | |
| else: | |
| counter = 0 | |
| if counter < len(atoms): | |
| xyz = [ | |
| np.array(l[s].strip()).astype(float) for s in coords_in_pdb | |
| ] | |
| coords.append(xyz) | |
| if also_bfactors: | |
| bfactor = np.array(l[bfactor_in_pdb].strip()).astype(float) | |
| bfactors.append(bfactor) | |
| resi_prev = resi | |
| coords = torch.Tensor(np.array(coords)).view(1, -1, len(atoms), 3) | |
| if also_bfactors: | |
| bfactors = torch.Tensor(np.array(bfactors)).view(1, -1, len(atoms)) | |
| elif method == "biopdb": | |
| structure = Bio.PDB.PDBParser(QUIET=True).get_structure(pdb[:-3], pdb) | |
| for model in structure: | |
| for chain in model: | |
| for res in chain: | |
| for atom in atoms: | |
| try: | |
| coords.append(np.asarray(res[atom].get_coord())) | |
| if also_bfactors: | |
| bfactors.append(np.asarray(res[atom].get_bfactor())) | |
| except: | |
| continue | |
| else: | |
| raise NotImplementedError(f"Invalid method for reading coords: {method}") | |
| if also_bfactors: | |
| if normalize_bfactors: # Normalize over Calphas | |
| mean_b = bfactors[..., 1].mean() | |
| std_b = bfactors[..., 1].var().sqrt() | |
| bfactors = (bfactors - mean_b) / (std_b + 1e-6) | |
| return coords, bfactors | |
| return coords | |
| def feats_to_pdb_str( | |
| atom_positions, | |
| aatype=None, | |
| atom_mask=None, | |
| residue_index=None, | |
| chain_index=None, | |
| b_factors=None, | |
| atom_lines_only=True, | |
| conect=False, | |
| **kwargs, | |
| ): | |
| # Expects unbatched, cropped inputs. needs at least one of atom_mask, aatype | |
| # Uses all-GLY aatype if aatype not given: does not infer from atom_mask | |
| assert aatype is not None or atom_mask is not None | |
| if atom_mask is None: | |
| aatype = aatype.cpu() | |
| atom_mask = atom37_mask_from_aatype(aatype, torch.ones_like(aatype)) | |
| if aatype is None: | |
| seq_mask = atom_mask[:, residue_constants.atom_order["CA"]].cpu() | |
| aatype = seq_mask * residue_constants.restype_order["G"] | |
| if residue_index is None: | |
| residue_index = torch.arange(aatype.shape[-1]) | |
| if chain_index is None: | |
| chain_index = torch.ones_like(aatype) | |
| if b_factors is None: | |
| b_factors = torch.ones_like(atom_mask) | |
| cast = lambda x: np.array(x.detach().cpu()) if isinstance(x, torch.Tensor) else x | |
| prot = protein.Protein( | |
| atom_positions=cast(atom_positions), | |
| atom_mask=cast(atom_mask), | |
| aatype=cast(aatype), | |
| residue_index=cast(residue_index), | |
| chain_index=cast(chain_index), | |
| b_factors=cast(b_factors), | |
| ) | |
| pdb_str = protein.to_pdb(prot, conect=conect) | |
| if conect: | |
| pdb_str, conect_str = pdb_str | |
| if atom_lines_only: | |
| pdb_lines = pdb_str.split("\n") | |
| atom_lines = [ | |
| l for l in pdb_lines if len(l.split()) > 1 and l.split()[0] == "ATOM" | |
| ] | |
| pdb_str = "\n".join(atom_lines) + "\n" | |
| if conect: | |
| pdb_str = pdb_str + conect_str | |
| return pdb_str | |
| def bb_coords_to_pdb_str(coords, atoms=["N", "CA", "C", "O"]): | |
| def _bb_pdb_line(atom, atomnum, resnum, coords, elem, res="GLY"): | |
| atm = "ATOM".ljust(6) | |
| atomnum = str(atomnum).rjust(5) | |
| atomname = atom.center(4) | |
| resname = res.ljust(3) | |
| chain = "A".rjust(1) | |
| resnum = str(resnum).rjust(4) | |
| x = str("%8.3f" % (float(coords[0]))).rjust(8) | |
| y = str("%8.3f" % (float(coords[1]))).rjust(8) | |
| z = str("%8.3f" % (float(coords[2]))).rjust(8) | |
| occ = str("%6.2f" % (float(1))).rjust(6) | |
| temp = str("%6.2f" % (float(20))).ljust(6) | |
| elname = elem.rjust(12) | |
| return "%s%s %s %s %s%s %s%s%s%s%s%s\n" % ( | |
| atm, | |
| atomnum, | |
| atomname, | |
| resname, | |
| chain, | |
| resnum, | |
| x, | |
| y, | |
| z, | |
| occ, | |
| temp, | |
| elname, | |
| ) | |
| n = coords.shape[0] | |
| na = len(atoms) | |
| pdb_str = "" | |
| for j in range(0, n, na): | |
| for idx, atom in enumerate(atoms): | |
| pdb_str += _bb_pdb_line( | |
| atom, | |
| j + idx + 1, | |
| (j + na) // na, | |
| coords[j + idx], | |
| atom[0], | |
| ) | |
| return pdb_str | |
| def write_coords_to_pdb( | |
| coords_in, | |
| filename, | |
| batched=True, | |
| write_to_frames=False, | |
| conect=False, | |
| **all_atom_feats, | |
| ): | |
| def _write_pdb_string(pdb_str, filename, append=False): | |
| write_mode = "a" if append else "w" | |
| with open(filename, write_mode) as f: | |
| if write_to_frames: | |
| f.write("MODEL\n") | |
| f.write(pdb_str) | |
| if write_to_frames: | |
| f.write("ENDMDL\n") | |
| if not (batched or write_to_frames): | |
| coords_in = [coords_in] | |
| filename = [filename] | |
| all_atom_feats = {k: [v] for k, v in all_atom_feats.items()} | |
| n_atoms_in = coords_in[0].shape[-2] | |
| is_bb_or_ca_pdb = n_atoms_in <= 4 | |
| for i, c in enumerate(coords_in): | |
| n_res = c.shape[0] | |
| if isinstance(filename, list): | |
| fname = filename[i] | |
| elif write_to_frames or len(coords_in) == 1: | |
| fname = filename | |
| else: | |
| fname = f"{filename[:-4]}_{i}.pdb" | |
| if is_bb_or_ca_pdb: | |
| c_flat = rearrange(c, "n a c -> (n a) c") | |
| if n_atoms_in == 1: | |
| atoms = ["CA"] | |
| if n_atoms_in == 3: | |
| atoms = ["N", "CA", "C"] | |
| if n_atoms_in == 4: | |
| atoms = ["N", "CA", "C", "O"] | |
| pdb_str = bb_coords_to_pdb_str(c_flat, atoms) | |
| else: | |
| feats_i = {k: v[i][:n_res] for k, v in all_atom_feats.items()} | |
| pdb_str = feats_to_pdb_str(c, conect=conect, **feats_i) | |
| _write_pdb_string(pdb_str, fname, append=write_to_frames and i > 0) | |
| ###################### LOSSES ################################### | |
| def masked_cross_entropy(logprobs, target, loss_mask): | |
| # target is onehot | |
| cel = -(target * logprobs) | |
| cel = cel * loss_mask[..., None] | |
| cel = cel.sum((-1, -2)) / loss_mask.sum(-1).clamp(min=1e-6) | |
| return cel | |
| def masked_mse(x, y, mask, weight=None): | |
| data_dims = tuple(range(1, len(x.shape))) | |
| mse = (x - y).pow(2) * mask | |
| if weight is not None: | |
| mse = mse * expand(weight, mse) | |
| mse = mse.sum(data_dims) / mask.sum(data_dims).clamp(min=1e-6) | |
| return mse | |
| ###################### ALIGN ################################### | |
| def quick_tmalign( | |
| p, p_sele, q_sele, tmscore_type="avg", differentiable_rmsd=False, rmsd_type="ca" | |
| ): | |
| # sota 210712 | |
| write_coords_to_pdb(p_sele[:, 1:2], "temp_p.pdb", atoms=["CA"], batched=False) | |
| write_coords_to_pdb(q_sele[:, 1:2], "temp_q.pdb", atoms=["CA"], batched=False) | |
| cmd = f"{PATH_TO_TMALIGN} temp_p.pdb temp_q.pdb -m temp_matrix.txt" | |
| outputs = subprocess.run(shlex.split(cmd), capture_output=True, text=True) | |
| # Get RMSD and TM scores | |
| tmout = outputs.stdout.split("\n") | |
| rmsd = float(tmout[16].split()[4][:-1]) | |
| tmscore1 = float(tmout[17].split()[1]) | |
| tmscore2 = float(tmout[18].split()[1]) | |
| if tmscore_type == "avg": | |
| tmscore = (tmscore1 + tmscore2) / 2 | |
| elif tmscore_type == "1" or tmscore_type == "query": | |
| tmscore = tmscore1 | |
| elif tmscore_type == "2": | |
| tmscore = tmscore2 | |
| elif tmscore_type == "both": | |
| tmscore = (tmscore1, tmscore2) | |
| # Get R, t and transform p coords | |
| m = open("temp_matrix.txt", "r").readlines()[2:5] | |
| m = [l.strip()[1:].strip() for l in m] | |
| m = torch.Tensor([[float(i) for i in l.split()] for l in m]).to(p_sele.device) | |
| R = m[:, 1:].t() | |
| t = m[:, 0] | |
| aligned_psele = p_sele @ R + t | |
| aligned = p @ R + t | |
| # Option 2 for rms - MSE of aligned against target coords using TMalign seq alignment. Differentiable | |
| if differentiable_rmsd: | |
| pi, qi = 0, 0 | |
| p_idxs, q_idxs = [], [] | |
| for i, c in enumerate(tmout[23]): | |
| if c in [":", "."]: | |
| p_idxs.append(pi) | |
| q_idxs.append(qi) | |
| if tmout[22][i] != "-": | |
| pi += 1 | |
| if tmout[24][i] != "-": | |
| qi += 1 | |
| tmalign_seq_p = p_sele[p_idxs] | |
| tmalign_seq_q = q_sele[q_idxs] | |
| if rmsd_type == "ca": | |
| tmalign_seq_p = tmalign_seq_p[:, 1] | |
| tmalign_seq_q = tmalign_seq_q[:, 1] | |
| elif rmsd_type == "bb": | |
| pass | |
| rmsd = (tmalign_seq_p - tmalign_seq_q).pow(2).sum(-1).sqrt().mean() | |
| # Delete temp files: p.pdb, q.pdb, matrix.txt, tmalign.out | |
| subprocess.run(shlex.split("rm temp_p.pdb")) | |
| subprocess.run(shlex.split("rm temp_q.pdb")) | |
| subprocess.run(shlex.split("rm temp_matrix.txt")) | |
| return {"aligned": aligned, "rmsd": rmsd, "tm_score": tmscore, "R": R, "t": t} | |
| ###################### OTHER ################################### | |
| def expand(x, tgt=None, dim=1): | |
| if tgt is None: | |
| for _ in range(dim): | |
| x = x[..., None] | |
| else: | |
| while len(x.shape) < len(tgt.shape): | |
| x = x[..., None] | |
| return x | |
| def hookfn(name, verbose=False): | |
| def f(grad): | |
| if check_nan_inf(grad) > 0: | |
| print(name, "grad nan/infs", grad.shape, check_nan_inf(grad), grad) | |
| if verbose: | |
| print(name, "grad shape", grad.shape, "norm", grad.norm()) | |
| return f | |
| def trigger_nan_check(name, x): | |
| if check_nan_inf(x) > 0: | |
| print(name, check_nan_inf(x)) | |
| raise Exception | |
| def check_nan_inf(x): | |
| return torch.isinf(x).sum() + torch.isnan(x).sum() | |
| def directory_find(atom, root="."): | |
| for path, dirs, files in os.walk(root): | |
| if atom in dirs: | |
| return os.path.join(path, atom) | |
| def dict2namespace(config): | |
| namespace = argparse.Namespace() | |
| for key, value in config.items(): | |
| if isinstance(value, dict): | |
| new_value = dict2namespace(value) | |
| else: | |
| new_value = value | |
| setattr(namespace, key, new_value) | |
| return namespace | |
| def load_config(path, return_dict=False): | |
| with open(path, "r") as f: | |
| config_dict = yaml.safe_load(f) | |
| config = dict2namespace(config_dict) | |
| if return_dict: | |
| return config, config_dict | |
| else: | |
| return config | |