Spaces:
Running
Running
File size: 5,798 Bytes
e76b79a | 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 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 | # Copyright (c) Meta Platforms, Inc. and affiliates.
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from abc import ABC, abstractmethod
from copy import deepcopy
from typing import List, Union
import numpy as np
from biotite.structure import AtomArray
ALL_RESIDUE_TYPES = [
"A",
"R",
"N",
"D",
"C",
"Q",
"E",
"G",
"H",
"I",
"L",
"K",
"M",
"F",
"P",
"S",
"T",
"W",
"Y",
"V",
]
RESIDUE_TYPES_WITHOUT_CYSTEINE = deepcopy(ALL_RESIDUE_TYPES)
RESIDUE_TYPES_WITHOUT_CYSTEINE.remove("C")
RESIDUE_TYPES_1to3 = {
"A": "ALA",
"R": "ARG",
"N": "ASN",
"D": "ASP",
"C": "CYS",
"Q": "GLN",
"E": "GLU",
"G": "GLY",
"H": "HIS",
"I": "ILE",
"L": "LEU",
"K": "LYS",
"M": "MET",
"F": "PHE",
"P": "PRO",
"S": "SER",
"T": "THR",
"W": "TRP",
"Y": "TYR",
"V": "VAL",
}
RESIDUE_TYPES_3to1 = {v: k for k, v in RESIDUE_TYPES_1to3.items()}
class SequenceSegmentFactory(ABC):
def __init__(self) -> None:
pass
@abstractmethod
def get(self) -> str:
pass
@abstractmethod
def mutate(self) -> None:
pass
@abstractmethod
def num_mutation_candidates(self) -> int:
pass
class ConstantSequenceSegment(SequenceSegmentFactory):
def __init__(self, sequence: str) -> None:
super().__init__()
self.sequence = sequence
def get(self) -> str:
return self.sequence
def mutate(self) -> None:
pass
def num_mutation_candidates(self) -> int:
return 0
class FixedLengthSequenceSegment(SequenceSegmentFactory):
def __init__(
self, initial_sequence: Union[str, int], disallow_mutations_to_cysteine=True,
) -> None:
super().__init__()
self.mutation_residue_types = (
RESIDUE_TYPES_WITHOUT_CYSTEINE
if disallow_mutations_to_cysteine
else ALL_RESIDUE_TYPES
)
self.sequence = (
initial_sequence
if type(initial_sequence) == str
else random_sequence(
length=initial_sequence, corpus=self.mutation_residue_types
)
)
def get(self) -> str:
return self.sequence
def mutate(self) -> None:
self.sequence = substitute_one_amino_acid(
self.sequence, self.mutation_residue_types
)
def num_mutation_candidates(self) -> int:
return len(self.sequence)
def substitute_one_amino_acid(sequence: str, corpus: List[str]) -> str:
sequence = list(sequence)
index = np.random.choice(len(sequence))
sequence[index] = np.random.choice(corpus)
return "".join(sequence)
def random_sequence(length: int, corpus: List[str]) -> str:
"Generate a random sequence using amino acids in corpus."
return "".join([np.random.choice(corpus) for _ in range(length)])
def sequence_from_atomarray(atoms: AtomArray) -> str:
return "".join(
[RESIDUE_TYPES_3to1[aa] for aa in atoms[atoms.atom_name == "CA"].res_name]
)
class VariableLengthSequenceSegment(SequenceSegmentFactory):
def __init__(
self,
initial_sequence: Union[str, int],
disallow_mutations_to_cysteine=True,
mutation_operation_probabilities: List[float] = [
3., # Substitution weight.
1., # Deletion weight.
1., # Insertion weight.
],
) -> None:
super().__init__()
self.mutation_residue_types = (
RESIDUE_TYPES_WITHOUT_CYSTEINE
if disallow_mutations_to_cysteine
else ALL_RESIDUE_TYPES
)
self.sequence = (
initial_sequence
if type(initial_sequence) == str
else random_sequence(
length=initial_sequence, corpus=self.mutation_residue_types
)
)
self.mutation_operation_probabilities = np.array(mutation_operation_probabilities)
self.mutation_operation_probabilities /= self.mutation_operation_probabilities.sum()
def get(self) -> str:
return self.sequence
def mutate(self) -> None:
mutation_operation = np.random.choice(
[
self._mutate_substitution,
self._mutate_deletion,
self._mutate_insertion,
],
p=self.mutation_operation_probabilities,
)
mutation_operation()
def _mutate_substitution(self) -> str:
self.sequence = substitute_one_amino_acid(
self.sequence, self.mutation_residue_types
)
def _mutate_deletion(self) -> str:
self.sequence = delete_one_amino_acid(self.sequence)
def _mutate_insertion(self) -> str:
self.sequence = insert_one_amino_acid(
self.sequence, self.mutation_residue_types
)
def num_mutation_candidates(self) -> int:
# NOTE(brianhie): This should be `3*len(self.sequence) + 1`,
# since there are `len(self.sequence)` substitutions and
# deletions, and `len(self.sequence) + 1` insertions.
# However, as this is used to weight sequence segments for
# mutations when combined into a multi-segment program, we
# just weight by `len(self.sequence)` for now.
return len(self.sequence)
def delete_one_amino_acid(sequence: str) -> str:
index = np.random.choice(len(sequence))
return sequence[:index] + sequence[index + 1 :]
def insert_one_amino_acid(sequence: str, corpus: List[str]) -> str:
n = len(sequence)
index = np.random.randint(0, n) if n > 0 else 0
insertion = np.random.choice(corpus)
return sequence[:index] + insertion + sequence[index:]
|