2021 alquraishi laboratory 2021 deepmind technologies limited licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license construct denser atom positions 14 dimensions instead of 37 restypeatom14toatom37list restypeatom37toatom14list restypeatom14masklist for rt in rc restypes atomnames rc restypenametoatom14namesrc restype1to3rt restypeatom14toatom37list appendrc atomordername if name else 0 for name in atomnames atomnametoidx14 name i for i name in enumerateatomnames restypeatom37toatom14list append atomnametoidx14name if name in atomnametoidx14 else 0 for name in rc atomtypes restypeatom14masklist append1 0 if name else 0 0 for name in atomnames add dummy mapping for restype unk restypeatom14toatom37list append0 14 restypeatom37toatom14list append0 37 restypeatom14masklist append0 0 14 restypeatom14toatom37 torch tensor restypeatom14toatom37list dtypetorch int32 deviceproteinaatype device restypeatom37toatom14 torch tensor restypeatom37toatom14list dtypetorch int32 deviceproteinaatype device restypeatom14mask torch tensor restypeatom14masklist dtypetorch float32 deviceproteinaatype device proteinaatype proteinaatype totorch long create the mapping for residx atom14 atom37 i e an array with shape numres 14 containing the atom37 indices for this protein residxatom14toatom37 restypeatom14toatom37proteinaatype residxatom14mask restypeatom14maskproteinaatype proteinatom14atomexists residxatom14mask proteinresidxatom14toatom37 residxatom14toatom37 long create the gather indices for mapping back residxatom37toatom14 restypeatom37toatom14proteinaatype proteinresidxatom37toatom14 residxatom37toatom14 long create the corresponding mask restypeatom37mask torch zeros21 37 dtypetorch float32 deviceproteinaatype device for restype restypeletter in enumeraterc restypes restypename rc restype1to3restypeletter atomnames rc residueatomsrestypename for atomname in atomnames atomtype rc atomorderatomname restypeatom37maskrestype atomtype 1 residxatom37mask restypeatom37maskproteinaatype proteinatom37atomexists residxatom37mask return protein def makeatom14masksnpbatch dictstr torch tensor dictstr np ndarray batch treemaplambda n torch tensorn devicebatchaatype device batch np ndarray out tensortreemaplambda t np arrayt makeatom14masksbatch return out 2021 alquraishi laboratory 2021 deepmind technologies limited licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license construct denser atom positions 14 dimensions instead of 37 add dummy mapping for restype unk create the mapping for residx atom14 atom37 i e an array with shape num_res 14 containing the atom37 indices for this protein create the gather indices for mapping back create the corresponding mask
from typing import Dict

import numpy as np
import torch

from . import residue_constants as rc
from .tensor_utils import tensor_tree_map, tree_map


def make_atom14_masks(protein: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
    
    restype_atom14_to_atom37_list = []
    restype_atom37_to_atom14_list = []
    restype_atom14_mask_list = []

    for rt in rc.restypes:
        atom_names = rc.restype_name_to_atom14_names[rc.restype_1to3[rt]]
        restype_atom14_to_atom37_list.append([(rc.atom_order[name] if name else 0) for name in atom_names])
        atom_name_to_idx14 = {name: i for i, name in enumerate(atom_names)}
        restype_atom37_to_atom14_list.append(
            [(atom_name_to_idx14[name] if name in atom_name_to_idx14 else 0) for name in rc.atom_types]
        )

        restype_atom14_mask_list.append([(1.0 if name else 0.0) for name in atom_names])

    
    restype_atom14_to_atom37_list.append([0] * 14)
    restype_atom37_to_atom14_list.append([0] * 37)
    restype_atom14_mask_list.append([0.0] * 14)

    restype_atom14_to_atom37 = torch.tensor(
        restype_atom14_to_atom37_list,
        dtype=torch.int32,
        device=protein["aatype"].device,
    )
    restype_atom37_to_atom14 = torch.tensor(
        restype_atom37_to_atom14_list,
        dtype=torch.int32,
        device=protein["aatype"].device,
    )
    restype_atom14_mask = torch.tensor(
        restype_atom14_mask_list,
        dtype=torch.float32,
        device=protein["aatype"].device,
    )
    protein_aatype = protein["aatype"].to(torch.long)

    
    
    residx_atom14_to_atom37 = restype_atom14_to_atom37[protein_aatype]
    residx_atom14_mask = restype_atom14_mask[protein_aatype]

    protein["atom14_atom_exists"] = residx_atom14_mask
    protein["residx_atom14_to_atom37"] = residx_atom14_to_atom37.long()

    
    residx_atom37_to_atom14 = restype_atom37_to_atom14[protein_aatype]
    protein["residx_atom37_to_atom14"] = residx_atom37_to_atom14.long()

    
    restype_atom37_mask = torch.zeros([21, 37], dtype=torch.float32, device=protein["aatype"].device)
    for restype, restype_letter in enumerate(rc.restypes):
        restype_name = rc.restype_1to3[restype_letter]
        atom_names = rc.residue_atoms[restype_name]
        for atom_name in atom_names:
            atom_type = rc.atom_order[atom_name]
            restype_atom37_mask[restype, atom_type] = 1

    residx_atom37_mask = restype_atom37_mask[protein_aatype]
    protein["atom37_atom_exists"] = residx_atom37_mask

    return protein


def make_atom14_masks_np(batch: Dict[str, torch.Tensor]) -> Dict[str, np.ndarray]:
    batch = tree_map(lambda n: torch.tensor(n, device=batch["aatype"].device), batch, np.ndarray)
    out = tensor_tree_map(lambda t: np.array(t), make_atom14_masks(batch))
    return out