File size: 4,182 Bytes
35cdf53 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 |
"""Process Structure Data."""
from flax_model.alphafold3.constants import atom_types
from flax_model.alphafold3.constants import residue_names
from flax_model.alphafold3.constants import side_chains
import numpy as np
NUM_DENSE = atom_types.DENSE_ATOM_NUM
NUM_AA = len(residue_names.PROTEIN_TYPES)
NUM_AA_WITH_UNK_AND_GAP = len(
residue_names.PROTEIN_TYPES_ONE_LETTER_WITH_UNKNOWN_AND_GAP
)
NUM_RESTYPES_WITH_UNK_AND_GAP = (
residue_names.POLYMER_TYPES_NUM_WITH_UNKNOWN_AND_GAP
)
def _make_restype_rigidgroup_dense_atom_idx():
"""Create Mapping from rigid_groups to dense_atom indices."""
# Create an array with the atom names.
# shape (num_restypes, num_rigidgroups, 3_atoms):
# (31, 8, 3)
base_atom_indices = np.zeros(
(NUM_RESTYPES_WITH_UNK_AND_GAP, 8, 3), dtype=np.int32
)
# 4,5,6,7: 'chi1,2,3,4-group'
for restype, restype_letter in enumerate(
residue_names.PROTEIN_TYPES_ONE_LETTER
):
resname = residue_names.PROTEIN_COMMON_ONE_TO_THREE[restype_letter]
dense_atom_names = atom_types.ATOM14[resname]
# 0: backbone frame
base_atom_indices[restype, 0, :] = [
dense_atom_names.index(atom) for atom in ['C', 'CA', 'N']
]
# 3: 'psi-group'
base_atom_indices[restype, 3, :] = [
dense_atom_names.index(atom) for atom in ['CA', 'C', 'O']
]
for chi_idx in range(4):
if side_chains.CHI_ANGLES_MASK[restype][chi_idx]:
atom_names = side_chains.CHI_ANGLES_ATOMS[resname][chi_idx]
base_atom_indices[restype, chi_idx + 4, :] = [
dense_atom_names.index(atom) for atom in atom_names[1:]
]
dense_atom_names = atom_types.DENSE_ATOM['A']
nucleic_rigid_atoms = [
dense_atom_names.index(atom) for atom in ["C1'", "C3'", "C4'"]
]
for nanum, _ in enumerate(residue_names.NUCLEIC_TYPES):
# 0: backbone frame only.
# we have aa + unk + gap, so we want to start after those
resnum = nanum + NUM_AA_WITH_UNK_AND_GAP
base_atom_indices[resnum, 0, :] = nucleic_rigid_atoms
return base_atom_indices
RESTYPE_RIGIDGROUP_DENSE_ATOM_IDX = _make_restype_rigidgroup_dense_atom_idx()
def _make_restype_pseudobeta_idx():
"""Returns indices of residue's pseudo-beta."""
restype_pseudobeta_index = np.zeros(
(NUM_RESTYPES_WITH_UNK_AND_GAP,), dtype=np.int32
)
for restype, restype_letter in enumerate(
residue_names.PROTEIN_TYPES_ONE_LETTER
):
restype_name = residue_names.PROTEIN_COMMON_ONE_TO_THREE[restype_letter]
atom_names = list(atom_types.ATOM14[restype_name])
if restype_name in {'GLY'}:
restype_pseudobeta_index[restype] = atom_names.index('CA')
else:
restype_pseudobeta_index[restype] = atom_names.index('CB')
for nanum, resname in enumerate(residue_names.NUCLEIC_TYPES):
atom_names = list(atom_types.DENSE_ATOM[resname])
# 0: backbone frame only.
# we have aa + unk , so we want to start after those
restype = nanum + NUM_AA_WITH_UNK_AND_GAP
if resname in {'A', 'G', 'DA', 'DG'}:
restype_pseudobeta_index[restype] = atom_names.index('C4')
else:
restype_pseudobeta_index[restype] = atom_names.index('C2')
return restype_pseudobeta_index
RESTYPE_PSEUDOBETA_INDEX = _make_restype_pseudobeta_idx()
def _make_aatype_dense_atom_to_atom37():
"""Map from dense_atom to atom37 per residue type."""
restype_dense_atom_to_atom37 = [] # mapping (restype, dense_atom) --> atom37
for rt in residue_names.PROTEIN_TYPES_ONE_LETTER:
atom_names = list(
atom_types.ATOM14_PADDED[residue_names.PROTEIN_COMMON_ONE_TO_THREE[rt]]
)
atom_names.extend([''] * (NUM_DENSE - len(atom_names)))
restype_dense_atom_to_atom37.append(
[(atom_types.ATOM37_ORDER[name] if name else 0) for name in atom_names]
)
# Add dummy mapping for restype 'UNK', '-' (gap), and nucleics [but not DN].
for _ in range(2 + len(residue_names.NUCLEIC_TYPES_WITH_UNKNOWN)):
restype_dense_atom_to_atom37.append([0] * NUM_DENSE)
restype_dense_atom_to_atom37 = np.array(
restype_dense_atom_to_atom37, dtype=np.int32
)
return restype_dense_atom_to_atom37
PROTEIN_AATYPE_DENSE_ATOM_TO_ATOM37 = _make_aatype_dense_atom_to_atom37()
|