File size: 4,083 Bytes
4d20b62 | 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 | # Copyright 2025 ByteDance and/or its affiliates.
#
# 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.
import random
import biotite.structure as struc
import torch
from protenix.data import ccd
from pxdesign.data.constants import (
DNA_STD_RESIDUES,
PROT_STD_RESIDUES_ONE_TO_THREE,
RNA_STD_RESIDUES,
STD_RESIDUES_WITH_GAP,
)
def encoder(encode_def_list: list[str], input_list: list[str]) -> torch.Tensor:
"""
Encode a list of input values into a binary format using a specified encoding definition list.
Args:
encode_def_list (list): A list of encoding definitions.
input_list (list): A list of input values to be encoded.
Returns:
torch.Tensor: A tensor representing the binary encoding of the input values.
"""
onehot_dict = {}
num_keys = len(encode_def_list)
for index, key in enumerate(encode_def_list):
onehot = [0] * num_keys
onehot[index] = 1
onehot_dict[key] = onehot
onehot_encoded_data = [onehot_dict[item] for item in input_list]
onehot_tensor = torch.Tensor(onehot_encoded_data)
return onehot_tensor
def restype_onehot_encoded(restype_list: list[str]) -> torch.Tensor:
return encoder(list(STD_RESIDUES_WITH_GAP.keys()), restype_list)
def cano_seq_resname_with_mask(atom_array):
"""
Assign to each atom the three-letter residue name (resname)
corresponding to its place in the canonical sequences.
Non-standard residues are mapped to standard ones.
Residues that cannot be mapped to standard residues and ligands are all labeled as "UNK".
Note: Some CCD Codes in the canonical sequence are mapped to three letters. It is labeled as one "UNK".
Args:
atom_array (AtomArray): Biotite AtomArray object
Returns:
AtomArray: Biotite AtomArray object with "cano_seq_resname" annotation added.
"""
cano_seq_resname = []
starts = struc.get_residue_starts(atom_array, add_exclusive_stop=True)
for start, stop in zip(starts[:-1], starts[1:]):
res_atom_nums = stop - start
mol_type = atom_array.mol_type[start]
resname = atom_array.res_name[start]
if resname == "xpb":
one_letter_code = "j"
else:
one_letter_code = ccd.get_one_letter_code(resname)
if one_letter_code is None or len(one_letter_code) != 1:
# Some non-standard residues cannot be mapped back to one standard residue.
one_letter_code = "X" if mol_type == "protein" else "N"
if mol_type == "protein":
res_name_in_cano_seq = PROT_STD_RESIDUES_ONE_TO_THREE.get(
one_letter_code, "UNK"
)
elif mol_type == "dna":
res_name_in_cano_seq = "D" + one_letter_code
if res_name_in_cano_seq not in DNA_STD_RESIDUES:
res_name_in_cano_seq = "DN"
elif mol_type == "rna":
res_name_in_cano_seq = one_letter_code
if res_name_in_cano_seq not in RNA_STD_RESIDUES:
res_name_in_cano_seq = "N"
else:
# some molecules attached to a polymer like ATP-RNA. e.g.
res_name_in_cano_seq = "UNK"
cano_seq_resname.extend([res_name_in_cano_seq] * res_atom_nums)
return cano_seq_resname
def make_random_mask(ref_tensor, mask_ratio=None):
if mask_ratio == None:
mask_ratio = random.random()
mask = torch.zeros_like(ref_tensor)
random_tensor = torch.rand_like(ref_tensor.to(torch.float32))
mask[random_tensor >= mask_ratio] = 1
return mask.to(torch.int64)
|