import random import unittest import warnings import torch from CodonTransformer.CodonData import get_amino_acid_sequence from CodonTransformer.CodonPrediction import ( load_model, load_tokenizer, predict_dna_sequence, ) from CodonTransformer.CodonUtils import ( AMINO_ACIDS, ORGANISM2ID, STOP_SYMBOLS, DNASequencePrediction, ) class TestCodonPrediction(unittest.TestCase): @classmethod def setUpClass(cls): cls.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Suppress warnings about loading from HuggingFace for message in [ "Tokenizer path not provided. Loading from HuggingFace.", "Model path not provided. Loading from HuggingFace.", ]: warnings.filterwarnings("ignore", message=message) cls.model = load_model(device=cls.device) cls.tokenizer = load_tokenizer() def test_predict_dna_sequence_valid_input(self): protein_sequence = "MWWMW" organism = "Escherichia coli general" result = predict_dna_sequence( protein_sequence, organism, device=self.device, tokenizer=self.tokenizer, model=self.model, ) self.assertIsInstance(result.predicted_dna, str) self.assertTrue( all(nucleotide in "ATCG" for nucleotide in result.predicted_dna) ) self.assertEqual(result.predicted_dna, "ATGTGGTGGATGTGGTGA") def test_predict_dna_sequence_non_deterministic(self): protein_sequence = "MFWY" organism = "Escherichia coli general" num_iterations = 100 temperatures = [0.2, 0.5, 0.8] possible_outputs = set() possible_encodings_wo_stop = { "ATGTTTTGGTAT", "ATGTTCTGGTAT", "ATGTTTTGGTAC", "ATGTTCTGGTAC", } for _ in range(num_iterations): for temperature in temperatures: result = predict_dna_sequence( protein=protein_sequence, organism=organism, device=self.device, tokenizer=self.tokenizer, model=self.model, deterministic=False, temperature=temperature, ) possible_outputs.add(result.predicted_dna[:-3]) # Remove stop codon self.assertEqual(possible_outputs, possible_encodings_wo_stop) def test_predict_dna_sequence_invalid_inputs(self): test_cases = [ ("MKTZZFVLLL?", "Escherichia coli general", "invalid protein sequence"), ("MKTFFVLLL", "Alien $%#@!", "invalid organism code"), ("", "Escherichia coli general", "empty protein sequence"), ] for protein_sequence, organism, error_type in test_cases: with self.subTest(error_type=error_type): with self.assertRaises(ValueError): predict_dna_sequence( protein_sequence, organism, device=self.device, tokenizer=self.tokenizer, model=self.model, ) def test_predict_dna_sequence_top_p_effect(self): """Test that changing top_p affects the diversity of outputs.""" protein_sequence = "MFWY" organism = "Escherichia coli general" num_iterations = 50 temperature = 0.5 top_p_values = [0.8, 0.95] outputs_by_top_p = {top_p: set() for top_p in top_p_values} for top_p in top_p_values: for _ in range(num_iterations): result = predict_dna_sequence( protein=protein_sequence, organism=organism, device=self.device, tokenizer=self.tokenizer, model=self.model, deterministic=False, temperature=temperature, top_p=top_p, ) outputs_by_top_p[top_p].add( result.predicted_dna[:-3] ) # Remove stop codon # Assert that higher top_p results in more diverse outputs diversity_lower_top_p = len(outputs_by_top_p[0.8]) diversity_higher_top_p = len(outputs_by_top_p[0.95]) self.assertGreaterEqual( diversity_higher_top_p, diversity_lower_top_p, "Higher top_p should result in more diverse outputs", ) def test_predict_dna_sequence_invalid_temperature_and_top_p(self): """Test that invalid temperature and top_p values raise ValueError.""" protein_sequence = "MWWMW" organism = "Escherichia coli general" invalid_params = [ {"temperature": -0.1, "top_p": 0.95}, {"temperature": 0, "top_p": 0.95}, {"temperature": 0.5, "top_p": -0.1}, {"temperature": 0.5, "top_p": 1.1}, ] for params in invalid_params: with self.subTest(params=params): with self.assertRaises(ValueError): predict_dna_sequence( protein=protein_sequence, organism=organism, device=self.device, tokenizer=self.tokenizer, model=self.model, deterministic=False, temperature=params["temperature"], top_p=params["top_p"], ) def test_predict_dna_sequence_translation_consistency(self): """Test that the predicted DNA translates back to the original protein.""" protein_sequence = "MALWMRLLPLLALLALWGPDPAAAFVNQHLCGSHLVE" organism = "Escherichia coli general" result = predict_dna_sequence( protein=protein_sequence, organism=organism, device=self.device, tokenizer=self.tokenizer, model=self.model, deterministic=True, ) # Translate predicted DNA back to protein translated_protein = get_amino_acid_sequence(result.predicted_dna[:-3]) self.assertEqual( translated_protein, protein_sequence, "Translated protein does not match the original protein sequence", ) def test_predict_dna_sequence_long_protein_sequence(self): """Test the function with a very long protein sequence to check performance and correctness.""" protein_sequence = ( "M" + "MALWMRLLPLLALLALWGPDPAAAFVNQHLCGSHLVEALYLVCGERGFFYTPKTRREAEDLQVGQVELGG" * 20 + STOP_SYMBOLS[0] ) organism = "Escherichia coli general" result = predict_dna_sequence( protein=protein_sequence, organism=organism, device=self.device, tokenizer=self.tokenizer, model=self.model, deterministic=True, ) # Check that the predicted DNA translates back to the original protein dna_sequence = result.predicted_dna[:-3] translated_protein = get_amino_acid_sequence(dna_sequence) self.assertEqual( translated_protein, protein_sequence[:-1], "Translated protein does not match the original long protein sequence", ) def test_predict_dna_sequence_edge_case_organisms(self): """Test the function with organism IDs at the boundaries of the mapping.""" protein_sequence = "MWWMW" # Assuming ORGANISM2ID has IDs starting from 0 to N min_organism_id = min(ORGANISM2ID.values()) max_organism_id = max(ORGANISM2ID.values()) organisms = [min_organism_id, max_organism_id] for organism_id in organisms: with self.subTest(organism_id=organism_id): result = predict_dna_sequence( protein=protein_sequence, organism=organism_id, device=self.device, tokenizer=self.tokenizer, model=self.model, deterministic=True, ) self.assertIsInstance(result.predicted_dna, str) self.assertTrue( all(nucleotide in "ATCG" for nucleotide in result.predicted_dna) ) def test_predict_dna_sequence_concurrent_calls(self): """Test the function's behavior under concurrent execution.""" import threading protein_sequence = "MWWMW" organism = "Escherichia coli general" results = [] def call_predict(): result = predict_dna_sequence( protein=protein_sequence, organism=organism, device=self.device, tokenizer=self.tokenizer, model=self.model, deterministic=True, ) results.append(result.predicted_dna) threads = [threading.Thread(target=call_predict) for _ in range(10)] for thread in threads: thread.start() for thread in threads: thread.join() self.assertEqual(len(results), 10) self.assertTrue(all(dna == results[0] for dna in results)) def test_predict_dna_sequence_random_seed_consistency(self): """Test that setting a random seed results in consistent outputs in non-deterministic mode.""" protein_sequence = "MFWY" organism = "Escherichia coli general" temperature = 0.5 top_p = 0.95 torch.manual_seed(42) result1 = predict_dna_sequence( protein=protein_sequence, organism=organism, device=self.device, tokenizer=self.tokenizer, model=self.model, deterministic=False, temperature=temperature, top_p=top_p, ) torch.manual_seed(42) result2 = predict_dna_sequence( protein=protein_sequence, organism=organism, device=self.device, tokenizer=self.tokenizer, model=self.model, deterministic=False, temperature=temperature, top_p=top_p, ) self.assertEqual( result1.predicted_dna, result2.predicted_dna, "Outputs should be consistent when random seed is set", ) def test_predict_dna_sequence_invalid_tokenizer_and_model(self): """Test that providing invalid tokenizer or model raises appropriate exceptions.""" protein_sequence = "MWWMW" organism = "Escherichia coli general" with self.subTest("Invalid tokenizer"): with self.assertRaises(Exception): predict_dna_sequence( protein=protein_sequence, organism=organism, device=self.device, tokenizer="invalid_tokenizer_path", model=self.model, ) with self.subTest("Invalid model"): with self.assertRaises(Exception): predict_dna_sequence( protein=protein_sequence, organism=organism, device=self.device, tokenizer=self.tokenizer, model="invalid_model_path", ) def test_predict_dna_sequence_stop_codon_handling(self): """Test the function's handling of protein sequences ending with a non '_' or '*' stop symbol.""" protein_sequence = "MWW/" organism = "Escherichia coli general" with self.assertRaises(ValueError): predict_dna_sequence( protein=protein_sequence, organism=organism, device=self.device, tokenizer=self.tokenizer, model=self.model, ) def test_predict_dna_sequence_device_compatibility(self): """Test that the function works correctly on both CPU and GPU devices.""" protein_sequence = "MWWMW" organism = "Escherichia coli general" devices = [torch.device("cpu")] if torch.cuda.is_available(): devices.append(torch.device("cuda")) for device in devices: with self.subTest(device=device): result = predict_dna_sequence( protein=protein_sequence, organism=organism, device=device, tokenizer=self.tokenizer, model=self.model, deterministic=True, ) self.assertIsInstance(result.predicted_dna, str) self.assertTrue( all(nucleotide in "ATCG" for nucleotide in result.predicted_dna) ) def test_predict_dna_sequence_random_proteins(self): """Test random proteins to ensure translated DNA matches the original protein.""" organism = "Escherichia coli general" num_tests = 200 for _ in range(num_tests): # Generate a random protein sequence of random length between 10 and 50 protein_length = random.randint(10, 500) protein_sequence = "M" + "".join( random.choices(AMINO_ACIDS, k=protein_length - 1) ) protein_sequence += random.choice(STOP_SYMBOLS) result = predict_dna_sequence( protein=protein_sequence, organism=organism, device=self.device, tokenizer=self.tokenizer, model=self.model, deterministic=True, ) # Remove stop codon from predicted DNA dna_sequence = result.predicted_dna[:-3] # Translate predicted DNA back to protein translated_protein = get_amino_acid_sequence(dna_sequence) self.assertEqual( translated_protein, protein_sequence[:-1], # Remove stop symbol f"Translated protein does not match the original protein sequence for protein: {protein_sequence}", ) def test_predict_dna_sequence_long_protein_over_max_length(self): """Test that the model handles protein sequences longer than 2048 amino acids.""" # Create a protein sequence longer than 2048 amino acids base_sequence = ( "MALWMRLLPLLALLALWGPDPAAAFVNQHLCGSHLVEALYLVCGERGFFYTPKTRREAEDLQVGQVELGG" ) protein_sequence = base_sequence * 100 # Length > 2048 amino acids organism = "Escherichia coli general" result = predict_dna_sequence( protein=protein_sequence, organism=organism, device=self.device, tokenizer=self.tokenizer, model=self.model, deterministic=True, ) # Remove stop codon from predicted DNA dna_sequence = result.predicted_dna[:-3] translated_protein = get_amino_acid_sequence(dna_sequence) # Due to potential model limitations, compare up to the model's max supported length max_length = len(translated_protein) self.assertEqual( translated_protein[:max_length], protein_sequence[:max_length], "Translated protein does not match the original protein sequence up to the maximum length supported.", ) def test_predict_dna_sequence_multi_output(self): """Test that the function returns multiple sequences when num_sequences > 1.""" protein_sequence = "MFQLLAPWY" organism = "Escherichia coli general" num_sequences = 20 result = predict_dna_sequence( protein=protein_sequence, organism=organism, device=self.device, tokenizer=self.tokenizer, model=self.model, deterministic=False, num_sequences=num_sequences, ) self.assertIsInstance(result, list) self.assertEqual(len(result), num_sequences) for prediction in result: self.assertIsInstance(prediction, DNASequencePrediction) self.assertTrue( all(nucleotide in "ATCG" for nucleotide in prediction.predicted_dna) ) # Check that all predicted DNA sequences translate back to the original protein translated_protein = get_amino_acid_sequence(prediction.predicted_dna[:-3]) self.assertEqual(translated_protein, protein_sequence) def test_predict_dna_sequence_deterministic_multi_raises_error(self): """Test that requesting multiple sequences in deterministic mode raises an error.""" protein_sequence = "MFWY" organism = "Escherichia coli general" with self.assertRaises(ValueError): predict_dna_sequence( protein=protein_sequence, organism=organism, device=self.device, tokenizer=self.tokenizer, model=self.model, deterministic=True, num_sequences=3, ) def test_predict_dna_sequence_multi_diversity(self): """Test that multiple sequences generated are diverse.""" protein_sequence = "MFWYMFWY" organism = "Escherichia coli general" num_sequences = 10 result = predict_dna_sequence( protein=protein_sequence, organism=organism, device=self.device, tokenizer=self.tokenizer, model=self.model, deterministic=False, num_sequences=num_sequences, temperature=0.8, ) unique_sequences = set(prediction.predicted_dna for prediction in result) self.assertGreater( len(unique_sequences), 2, "Multiple sequence generation should produce diverse results", ) # Check that all sequences are valid translations of the input protein for prediction in result: translated_protein = get_amino_acid_sequence(prediction.predicted_dna[:-3]) self.assertEqual(translated_protein, protein_sequence) def test_predict_dna_sequence_match_protein_repetitive(self): """Test that match_protein=True correctly handles highly repetitive and unconventional sequences.""" test_sequences = ( "QQQQQQQQQQQQQQQQ_", "KRKRKRKRKRKRKRKR_", "PGPGPGPGPGPGPGPG_", "DEDEDEDEDEDEDEDEDE_", "M_M_M_M_M_", "MMMMMMMMMM_", "WWWWWWWWWW_", "CCCCCCCCCC_", "MWCHMWCHMWCH_", "Q_QQ_QQQ_QQQQ_", "MWMWMWMWMWMW_", "CCCHHHMMMWWW_", "_", "M_", "MGWC_", ) organism = "Homo sapiens" for protein_sequence in test_sequences: # Generate sequence with match_protein=True result = predict_dna_sequence( protein=protein_sequence, organism=organism, device=self.device, tokenizer=self.tokenizer, model=self.model, deterministic=False, temperature=20, # High temperature to test protein matching match_protein=True, ) dna_sequence = result.predicted_dna translated_protein = get_amino_acid_sequence(dna_sequence) self.assertEqual( translated_protein, protein_sequence, f"Translated protein must match original when match_protein=True. Failed for sequence: {protein_sequence}", ) def test_predict_dna_sequence_match_protein_rare_amino_acids(self): """Test match_protein with rare amino acids that have limited codon options.""" # Methionine (M) and Tryptophan (W) have only one codon each # While Leucine (L) has 6 codons - testing contrast protein_sequence = "MWLLLMWLLL" organism = "Escherichia coli general" # Run multiple predictions results = [] num_iterations = 10 for _ in range(num_iterations): result = predict_dna_sequence( protein=protein_sequence, organism=organism, device=self.device, tokenizer=self.tokenizer, model=self.model, deterministic=False, temperature=20, # High temperature to test protein matching match_protein=True, ) results.append(result.predicted_dna) # Check all sequences for dna_sequence in results: # Verify M always uses ATG m_positions = [0, 5] # Known positions of M in sequence for pos in m_positions: self.assertEqual( dna_sequence[pos * 3 : (pos + 1) * 3], "ATG", "Methionine must use ATG codon.", ) # Verify W always uses TGG w_positions = [1, 6] # Known positions of W in sequence for pos in w_positions: self.assertEqual( dna_sequence[pos * 3 : (pos + 1) * 3], "TGG", "Tryptophan must use TGG codon.", ) # Verify all L codons are valid l_positions = [2, 3, 4, 7, 8, 9] # Known positions of L in sequence l_codons = [dna_sequence[pos * 3 : (pos + 1) * 3] for pos in l_positions] valid_l_codons = {"TTA", "TTG", "CTT", "CTC", "CTA", "CTG"} self.assertTrue( all(codon in valid_l_codons for codon in l_codons), "All Leucine codons must be valid", ) if __name__ == "__main__": unittest.main()