Add a list of possible stop symbols.

Adibvafa · Sep 20, 2024 · 9334cdc · 9334cdc
1 parent bbfe03a
commit 9334cdc
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/CodonTransformer/CodonPrediction.py b/CodonTransformer/CodonPrediction.py
@@ -23,6 +23,7 @@
 from CodonTransformer.CodonData import get_merged_seq
 from CodonTransformer.CodonUtils import (
     AMINO_ACIDS,
+    STOP_SYMBOLS,
     INDEX2TOKEN,
     NUM_ORGANISMS,
     ORGANISM2ID,
@@ -148,7 +149,7 @@ def predict_dna_sequence(
         raise ValueError("Protein sequence cannot be empty.")
 
     # Ensure the protein sequence contains only valid amino acids
-    if not all(aminoacid in AMINO_ACIDS for aminoacid in protein):
+    if not all(aminoacid in AMINO_ACIDS + STOP_SYMBOLS for aminoacid in protein):
         raise ValueError("Invalid amino acid found in protein sequence.")
 
     # Validate temperature

diff --git a/CodonTransformer/CodonUtils.py b/CodonTransformer/CodonUtils.py
@@ -38,6 +38,7 @@
     "W",  # Tryptophan
     "Y",  # Tyrosine
 ]
+STOP_SYMBOLS = ["_", "*"]  # Stop codon symbols
 
 # Dictionary ambiguous amino acids to standard amino acids
 AMBIGUOUS_AMINOACID_MAP: Dict[str, str] = {