OneScience-Group/CodonTransformer

1"""2File: CodonEvaluation.py3---------------------------4Includes functions to calculate various evaluation metrics along with helper5functions.6"""7 8from typing import Dict, List, Tuple9 10import pandas as pd11from CAI import CAI, relative_adaptiveness12from tqdm import tqdm13 14 15def get_CSI_weights(sequences: List[str]) -> Dict[str, float]:16    """17    Calculate the Codon Similarity Index (CSI) weights for a list of DNA sequences.18 19    Args:20        sequences (List[str]): List of DNA sequences.21 22    Returns:23        dict: The CSI weights.24    """25    return relative_adaptiveness(sequences=sequences)26 27 28def get_CSI_value(dna: str, weights: Dict[str, float]) -> float:29    """30    Calculate the Codon Similarity Index (CSI) for a DNA sequence.31 32    Args:33        dna (str): The DNA sequence.34        weights (dict): The CSI weights from get_CSI_weights.35 36    Returns:37        float: The CSI value.38    """39    return CAI(dna, weights)40 41 42def get_organism_to_CSI_weights(43    dataset: pd.DataFrame, organisms: List[str]44) -> Dict[str, dict]:45    """46    Calculate the Codon Similarity Index (CSI) weights for a list of organisms.47 48    Args:49        dataset (pd.DataFrame): Dataset containing organism and DNA sequence info.50        organisms (List[str]): List of organism names.51 52    Returns:53        Dict[str, dict]: A dictionary mapping each organism to its CSI weights.54    """55    organism2weights = {}56 57    # Iterate through each organism to calculate its CSI weights58    for organism in tqdm(organisms, desc="Calculating CSI Weights: ", unit="Organism"):59        organism_data = dataset.loc[dataset["organism"] == organism]60        sequences = organism_data["dna"].to_list()61        weights = get_CSI_weights(sequences)62        organism2weights[organism] = weights63 64    return organism2weights65 66 67def get_GC_content(dna: str, lower: bool = False) -> float:68    """69    Calculate the GC content of a DNA sequence.70 71    Args:72        dna (str): The DNA sequence.73        lower (bool): If True, converts DNA sequence to lowercase before calculation.74 75    Returns:76        float: The GC content as a percentage.77    """78    if lower:79        dna = dna.lower()80    return (dna.count("G") + dna.count("C")) / len(dna) * 10081 82 83def get_cfd(84    dna: str,85    codon_frequencies: Dict[str, Tuple[List[str], List[float]]],86    threshold: float = 0.3,87) -> float:88    """89    Calculate the codon frequency distribution (CFD) metric for a DNA sequence.90 91    Args:92        dna (str): The DNA sequence.93        codon_frequencies (Dict[str, Tuple[List[str], List[float]]]): Codon94            frequency distribution per amino acid.95        threshold (float): Frequency threshold for counting rare codons.96 97    Returns:98        float: The CFD metric as a percentage.99    """100    # Get a dictionary mapping each codon to its normalized frequency101    codon2frequency = {102        codon: freq / max(frequencies)103        for amino, (codons, frequencies) in codon_frequencies.items()104        for codon, freq in zip(codons, frequencies)105    }106 107    cfd = 0108 109    # Iterate through the DNA sequence in steps of 3 to process each codon110    for i in range(0, len(dna), 3):111        codon = dna[i : i + 3]112        codon_frequency = codon2frequency[codon]113 114        if codon_frequency < threshold:115            cfd += 1116 117    return cfd / (len(dna) / 3) * 100118 119 120def get_min_max_percentage(121    dna: str,122    codon_frequencies: Dict[str, Tuple[List[str], List[float]]],123    window_size: int = 18,124) -> List[float]:125    """126    Calculate the %MinMax metric for a DNA sequence.127 128    Args:129        dna (str): The DNA sequence.130        codon_frequencies (Dict[str, Tuple[List[str], List[float]]]): Codon131            frequency distribution per amino acid.132        window_size (int): Size of the window to calculate %MinMax.133 134    Returns:135        List[float]: List of %MinMax values for the sequence.136 137    Credit: https://github.com/chowington/minmax138    """139    # Get a dictionary mapping each codon to its respective amino acid140    codon2amino = {141        codon: amino142        for amino, (codons, frequencies) in codon_frequencies.items()143        for codon in codons144    }145 146    min_max_values = []147    codons = [dna[i : i + 3] for i in range(0, len(dna), 3)]  # Split DNA into codons148 149    # Iterate through the DNA sequence using the specified window size150    for i in range(len(codons) - window_size + 1):151        codon_window = codons[i : i + window_size]  # Codons in the current window152 153        Actual = 0.0  # Average of the actual codon frequencies154        Max = 0.0  # Average of the min codon frequencies155        Min = 0.0  # Average of the max codon frequencies156        Avg = 0.0  # Average of the averages of all frequencies for each amino acid157 158        # Sum the frequencies for codons in the current window159        for codon in codon_window:160            aminoacid = codon2amino[codon]161            frequencies = codon_frequencies[aminoacid][1]162            codon_index = codon_frequencies[aminoacid][0].index(codon)163            codon_frequency = codon_frequencies[aminoacid][1][codon_index]164 165            Actual += codon_frequency166            Max += max(frequencies)167            Min += min(frequencies)168            Avg += sum(frequencies) / len(frequencies)169 170        # Divide by the window size to get the averages171        Actual = Actual / window_size172        Max = Max / window_size173        Min = Min / window_size174        Avg = Avg / window_size175 176        # Calculate %MinMax177        percentMax = ((Actual - Avg) / (Max - Avg)) * 100178        percentMin = ((Avg - Actual) / (Avg - Min)) * 100179 180        # Append the appropriate %MinMax value181        if percentMax >= 0:182            min_max_values.append(percentMax)183        else:184            min_max_values.append(-percentMin)185 186    # Populate the last floor(window_size / 2) entries of min_max_values with None187    for i in range(int(window_size / 2)):188        min_max_values.append(None)189 190    return min_max_values191 192 193def get_sequence_complexity(dna: str) -> float:194    """195    Calculate the sequence complexity score of a DNA sequence.196 197    Args:198        dna (str): The DNA sequence.199 200    Returns:201        float: The sequence complexity score.202    """203 204    def sum_up_to(x):205        """Recursive function to calculate the sum of integers from 1 to x."""206        if x <= 1:207            return 1208        else:209            return x + sum_up_to(x - 1)210 211    def f(x):212        """Returns 4 if x is greater than or equal to 4, else returns x."""213        if x >= 4:214            return 4215        elif x < 4:216            return x217 218    unique_subseq_length = []219 220    # Calculate unique subsequences lengths221    for i in range(1, len(dna) + 1):222        unique_subseq = set()223        for j in range(len(dna) - (i - 1)):224            unique_subseq.add(dna[j : (j + i)])225        unique_subseq_length.append(len(unique_subseq))226 227    # Calculate complexity score228    complexity_score = (229        sum(unique_subseq_length) / (sum_up_to(len(dna) - 1) + f(len(dna)))230    ) * 100231 232    return complexity_score233 234 235def get_sequence_similarity(236    original: str, predicted: str, truncate: bool = True, window_length: int = 1237) -> float:238    """239    Calculate the sequence similarity between two sequences.240 241    Args:242        original (str): The original sequence.243        predicted (str): The predicted sequence.244        truncate (bool): If True, truncate the original sequence to match the length245            of the predicted sequence.246        window_length (int): Length of the window for comparison (1 for amino acids,247            3 for codons).248 249    Returns:250        float: The sequence similarity as a percentage.251 252    Preconditions:253        len(predicted) <= len(original).254    """255    if not truncate and len(original) != len(predicted):256        raise ValueError(257            "Set truncate to True if the length of sequences do not match."258        )259 260    identity = 0.0261    original = original.strip()262    predicted = predicted.strip()263 264    if truncate:265        original = original[: len(predicted)]266 267    if window_length == 1:268        # Simple comparison for amino acid269        for i in range(len(predicted)):270            if original[i] == predicted[i]:271                identity += 1272    else:273        # Comparison for substrings based on window_length274        for i in range(0, len(original) - window_length + 1, window_length):275            if original[i : i + window_length] == predicted[i : i + window_length]:276                identity += 1277 278    return (identity / (len(predicted) / window_length)) * 100279