updated algorithm for the data augmentation

Author: dsmutin

scripts/databases/generate_noisy_probes.py

@@ -7,38 +7,6 @@
 from typing import Tuple, List, Dict
 import argparse
 
-def load_formamide_probabilities(file_path: str) -> Tuple[List[float], List[float]]:
-    """Load formamide probabilities from TSV file and normalize them."""
-    # Read the file with explicit float conversion
-    df = pd.read_csv(file_path, sep=r'\s+', header=None, names=['formamide', 'probability'])
-    
-    # Convert to float explicitly and handle any potential NaN values
-    df['probability'] = pd.to_numeric(df['probability'], errors='coerce')
-    df = df.dropna()  # Remove any rows with NaN values
-    
-    if df.empty:
-        raise ValueError("No valid probability values found in the formamide file")
-    
-    # Convert probabilities to numpy array and normalize
-    probs = df['probability'].values.astype(float)
-    
-    # Check for any invalid values
-    if np.any(np.isnan(probs)) or np.any(np.isinf(probs)):
-        raise ValueError("Invalid probability values found (NaN or Inf)")
-    
-    # Normalize probabilities
-    total = np.sum(probs)
-    if total <= 0:
-        raise ValueError("Sum of probabilities must be positive")
-    
-    probs = probs / total
-    
-    # Verify the probabilities sum to 1 (within numerical precision)
-    if not np.isclose(np.sum(probs), 1.0, rtol=1e-5):
-        raise ValueError("Probabilities do not sum to 1 after normalization")
-    
-    return df['formamide'].tolist(), probs.tolist()
-
 def clean_sequence(sequence: str) -> str:
     """Clean sequence by removing non-ACGT characters and converting to uppercase."""
     sequence = sequence.strip().upper()
@@ -52,34 +20,47 @@ def calculate_gc_content(sequence: str) -> int:
     gc_count = sequence.count('G') + sequence.count('C')
     return round((gc_count / len(sequence)) * 100)
 
-def apply_mutations(sequence: str, 
+def apply_mutations_sequence(sequence: str, 
                    insertion_rate: float = 0.01,
                    deletion_rate: float = 0.01,
-                   mutation_rate: float = 0.1) -> str:
+                   mutation_rate: float = 0.1,
+                   values: list=[None]) -> str:
     """Apply random mutations to the sequence."""
     nucleotides = ['A', 'C', 'G', 'T']
     result = []
 
-    for base in sequence:
-        # Apply mutations based on probabilities
-        if random.random() < mutation_rate:
-            # SNP mutation
-            new_base = random.choice([n for n in nucleotides if n != base])
-            result.append(new_base)
-        else:
-            result.append(base)
-            
-        # Insertion after current position
-        if random.random() < insertion_rate:
-            result.append(random.choice(nucleotides))
-            
-        # Deletion of current position
-        if random.random() < deletion_rate:
-            if result:
-                result.pop()
+    while sequence == ''.join(result):
+        for base in sequence:
+            # Apply mutations based on probabilities
+            if random.random() < mutation_rate:
+                # SNP mutation
+                new_base = random.choice([n for n in nucleotides if n != base])
+                result.append(new_base)
+            else:
+                result.append(base)
+                
+            # Insertion after current position
+            if random.random() < insertion_rate:
+                result.append(random.choice(nucleotides))
+                
+            # Deletion of current position
+            if random.random() < deletion_rate:
+                if result:
+                    result.pop()
 
     return ''.join(result)
 
+def apply_mutations_experiment(sequence,
+                               values: list=[None],
+                               insertion_rate: float = 0.01,
+                               deletion_rate: float = 0.01,
+                               mutation_rate: float = 0.1):
+    """Apply random mutations to the experiment parameters."""
+    if random.random() < mutation_rate:
+        return random.choice(values)
+    else:
+        return None
+
 def process_probe_data(df: pd.DataFrame) -> List[Dict]:
     """Process the input DataFrame into a list of probe dictionaries."""
     probes = []
@@ -103,18 +84,37 @@ def process_probe_data(df: pd.DataFrame) -> List[Dict]:
 
 def generate_noisy_probes(input_file: str,
                          output_file: str,
-                         formamide_file: str,
+                         mutation_number: int = 1,
+                         obligate_mutations: list = ["Sequence"],
+                         facultative_mutations: list = ["Formamide [%]","Modified version(s)"],
                          insertion_rate: float = 0.01,
                          deletion_rate: float = 0.01,
                          mutation_rate: float = 0.1,
                          iterations: int = 1) -> None:
-    """Generate noisy probe data with mutations and formamide variations."""
-    
-    # Load formamide probabilities
-    formamide_values, formamide_probs = load_formamide_probabilities(formamide_file)
-    
+    """Generate noisy probe data with sequence mutations and experiment parameter variations.
+    Parameters:
+        input_file (str): Path to the input CSV file containing probe data.
+        output_file (str): Path to the output file where noisy probes will be saved.
+        insertion_rate (float, optional): Probability of inserting a nucleotide at each position. Default is 0.01.
+        deletion_rate (float, optional): Probability of deleting a nucleotide at each position. Default is 0.01.
+        mutation_rate (float, optional): Probability of introducing a SNP at each position. Default is 0.1.
+        iterations (int, optional): Number of noisy probe sets to generate. Default is 1.
+    """
+       
     # Read input probe data
     df = pd.read_csv(input_file, header=None)
+    # Prepare
+    mutation_list = [*obligate_mutations, *facultative_mutations]
+    mutation_functions = {
+        "Sequence":apply_mutations_sequence,
+        "Formamide [%]":apply_mutations_experiment,
+        "Modified version(s)":apply_mutations_experiment
+    }
+    mutation_values = {
+        "Sequence":None,
+        "Formamide [%]":df.loc[df[0] == "Formamide [%]", 1].tolist(),
+        "Modified version(s)":df.loc[df[0] == "Modified version(s)", 1].tolist()
+    }
 
     # Process probes
     probes = process_probe_data(df)
@@ -129,34 +129,43 @@ def generate_noisy_probes(input_file: str,
 
             # Clean sequence
             clean_seq = clean_sequence(probe['Sequence'])
+            probe['Sequence'] = clean_seq
 
             # Apply mutations
-            mutated_seq = apply_mutations(clean_seq, 
-                                        insertion_rate,
-                                        deletion_rate,
-                                        mutation_rate)
+            for OM in obligate_mutations:
+                tmp = mutation_functions[OM]
+                probe[OM] = tmp(sequence = probe[OM],
+                                insertion_rate = insertion_rate,
+                                deletion_rate = deletion_rate,
+                                mutation_rate = mutation_rate,
+                                values = mutation_values[OM])
+                
+            if mutation_number < len(mutation_list):
+                OMs = random.choices(mutation_list, k=mutation_number)
+            else:
+                add = ["Sequence"]*(mutation_number - len(mutation_list))
+                OMs = [*mutation_list, *add]
 
+            for OM in OMs:
+                tmp = mutation_functions[OM]
+                probe[OM] = tmp(sequence = probe[OM],
+                                insertion_rate = insertion_rate,
+                                deletion_rate = deletion_rate,
+                                mutation_rate = mutation_rate,
+                                values = mutation_values[OM])
+                
             # Calculate new properties
-            new_length = len(mutated_seq)
-            new_gc_content = calculate_gc_content(mutated_seq)
-            
-            # Randomly select new formamide value based on probabilities
-            new_formamide = random.choices(formamide_values, 
-                                         weights=formamide_probs, 
-                                         k=1)[0]
-            
-            # Create new probe with updated values
-            new_probe = probe.copy()
-            new_probe['Sequence'] = mutated_seq
-            new_probe['Length [nt]'] = str(new_length)
-            new_probe['G+C content [%]'] = str(new_gc_content)
-            new_probe['Formamide [%]'] = str(new_formamide)
-            
+            new_length = len(probe['Sequence'])
+            new_gc_content = calculate_gc_content(probe['Sequence'])
+            probe['Length [nt]'] = str(new_length)
+            probe['G+C content [%]'] = str(new_gc_content)
+        
             # Generate unique ID for this iteration
-            if 'Accession no.' in new_probe:
-                new_probe['Accession no.'] = f"{new_probe['Accession no.']}_{iteration + 1}"
-            
-            noisy_probes.append(new_probe)
+            if 'Accession no.' in probe:
+                probe['Accession no.'] = f"{probe['Accession no.']}_{iteration + 1}"
+                
+            probe['Mutation number'] = mutation_number
+            noisy_probes.append(probe)
 
         all_noisy_probes.extend(noisy_probes)
 
@@ -174,23 +183,27 @@ def main():
     parser = argparse.ArgumentParser(description='Generate noisy probe data')
     parser.add_argument('--input', required=True, help='Input probeBase CSV file')
     parser.add_argument('--output', type=str, default='data/databases/open/probeBase_false.csv', help='Output noisy probeBase CSV file')
-    parser.add_argument('--formamide', type=str, default='data/databases/open/probeBase_formamide.tsv', help='Formamide probabilities TSV file')
+    parser.add_argument('--mutation-number', type=int, default=5, help='Maximum number of mutations to apply to each probe')
     parser.add_argument('--insertion-rate', type=float, default=0.01, help='Insertion mutation rate')
     parser.add_argument('--deletion-rate', type=float, default=0.01, help='Deletion mutation rate')
     parser.add_argument('--mutation-rate', type=float, default=0.1, help='SNP mutation rate')
     parser.add_argument('--iterations', type=int, default=10, help='Number of iterations to generate noisy data')
 
     args = parser.parse_args()
 
-    generate_noisy_probes(
-        args.input,
-        args.output,
-        args.formamide,
-        args.insertion_rate,
-        args.deletion_rate,
-        args.mutation_rate,
-        args.iterations
-    )
+    for i in range(args.mutation_number):
+        generate_noisy_probes(
+            args.input,
+            args.output,
+            i,
+            insertion_rate=args.insertion_rate,
+            deletion_rate=args.deletion_rate,
+            mutation_rate=args.mutation_rate,
+            iterations=args.iterations
+        )
 
 if __name__ == '__main__':
-    main() 
+    main() 
+    
+# Usage:
+# python scripts/databases/generate_noisy_probes.py --input data/databases/open/probeBase.csv --output data/databases/open/probeBase_mutated.csv