Updating gen_mut_model to output and read proper trinucleotide count files

Author: joshfactorial

neat/gen_mut_model/runner.py

@@ -85,18 +85,23 @@ def runner(reference_index,
 
     # Pre-parsing to find all the matching chromosomes between ref and vcf
     _LOG.info('Processing VCF file...')
-    matching_variants, matching_chromosomes = read_and_filter_variants(vcf_to_process,
-                                                                       reference_index,
-                                                                       ignore)
+    matching_variants, matching_chromosomes = read_and_filter_variants(
+        vcf_to_process,
+        reference_index,
+        ignore
+    )
 
     if not matching_variants or not matching_chromosomes:
         _LOG.error("No valid variants detected. Check names in vcf versus reference and/or bed.")
         sys.exit(1)
 
-    trinuc_ref_count, bed_track_length = count_trinucleotides(reference_index,
-                                                              bed,
-                                                              outcounts_file,
-                                                              matching_chromosomes)
+    trinuc_ref_count, bed_track_length = count_trinucleotides(
+        reference_index,
+        bed,
+        outcounts_file,
+        matching_chromosomes,
+        save_trinuc
+    )
 
     if not trinuc_ref_count:
         _LOG.error("No valid trinucleotides detected in reference.")
@@ -296,47 +301,6 @@ def runner(reference_index,
         for k in sorted(snp_trans_freq):
             _LOG.info(f'p({k[0]} --> {k[1]} | SNP occurs) = {snp_trans_freq[k]}')
 
-    # Save counts, if requested
-    if save_trinuc:
-        trinuc_output_data = {
-            'p(snp)': average_snp_freq,
-            'p(insertion)': average_insertion_frequency,
-            'p(deletion)': average_deletion_frequency,
-            'overall average mut rate': avg_mut_rate,
-            'total variants processed': total_var,
-            'homozygous probability': homozygous_frequency
-        }
-
-        for k in sorted(trinuc_mut_prob):
-            trinuc_output_data[f'p({k} mutates)'] = trinuc_mut_prob[k]
-
-        for k in sorted(trinuc_trans_probs):
-            trinuc_output_data[f'p({k[0]} --> {k[1]} | {k[0]} mutates)'] = trinuc_trans_probs[k]
-
-        for k in sorted(deletion_counts):
-            trinuc_output_data[f'p(del length = {abs(k)} | deletion occurs)'] = deletion_frequency[k]
-
-        for k in sorted(insertion_counts):
-            trinuc_output_data[f'p(insert length = {abs(k)} | insertion occurs)'] = insertion_freqency[k]
-
-        for k in sorted(snp_trans_freq):
-            trinuc_output_data[f'p({k[0]} --> {k[1]} | SNP occurs)'] = snp_trans_freq[k]
-
-        trinuc_output_path = Path(output)
-        trinuc_output_path = trinuc_output_path.with_suffix('')
-        trinuc_output_path = trinuc_output_path.with_suffix('.trinuc.pickle.gz')
-        _LOG.info(f'Saving trinucleotide counts to: {trinuc_output_path}')
-
-        with open_output(trinuc_output_path, 'w+') as trinuc_outfile:
-
-            pickle.dump(trinuc_output_data, trinuc_outfile)
-
-            # Human-readable content of file below
-            trinuc_outfile.write('\n')
-
-            for key, value in trinuc_output_data.items():
-                trinuc_outfile.write(f'{key}: {value}\n')
-
     _LOG.info(f'p(snp)   = {average_snp_freq}')
     _LOG.info(f'p(insertion) = {average_insertion_frequency}')
     _LOG.info(f'p(deletion) = {average_deletion_frequency}')
@@ -405,6 +369,8 @@ def compute_mut_runner(reference,
     if outcounts:
         validate_input_path(outcounts)
         outcounts = Path(outcounts)
+    elif save_trinuc:
+        outcounts = Path(output + '.trinuc.pickle.gz')
 
     print('Processing reference...')
     reference_index = SeqIO.index(reference, 'fasta')
@@ -477,8 +443,18 @@ def compute_mut_runner(reference,
     output = Path(output + '.pickle.gz')
     validate_output_path(output, overwrite=overwrite_output)
 
-    runner(reference_index, vcf_to_process, vcf_columns, outcounts, show_trinuc, save_trinuc,
-           output, bed, human_sample, skip_common)
+    runner(
+        reference_index,
+        vcf_to_process,
+        vcf_columns,
+        outcounts,
+        show_trinuc,
+        save_trinuc,
+        output,
+        bed,
+        human_sample,
+        skip_common
+    )
 
     if os.path.exists('temp.vcf'):
         os.remove('temp.vcf')

neat/gen_mut_model/utils.py

@@ -4,6 +4,8 @@
 
 import json
 import sys
+import pickle
+import gzip
 
 import numpy as np
 
@@ -165,14 +167,16 @@ def cluster_list(list_to_cluster, delta):
 def count_trinucleotides(reference_index,
                          bed,
                          trinuc_counts,
-                         matching_chroms):
+                         matching_chroms,
+                         save_trinuc_file):
     """
     Counts the frequency of the various trinucleotide combinations in the dataset
 
     :param dict reference_index: The index of the fasta, from SeqIO
     :param Path bed: Full path to bed file, if using
     :param Path trinuc_counts: Full path to existing trinculeotide counts file, if input, or the output path
     :param list matching_chroms: List of matching chromosomes
+    :param save_trinuc_file: Boolean determines whether to save trinucleotide counts file.
     :return dict, int: A dictionary of trinculeotides and counts, the number of bases spanned
     """
     # Data format: TRINUC: COUNT (e.g., "AAA": 12), where COUNT is the number of times trinuc was observed
@@ -183,11 +187,11 @@ def count_trinucleotides(reference_index,
     track_len = 0
 
     trinuc_counts_exists = False
-    save_trinuc_file = False
     if trinuc_counts:
         if trinuc_counts.is_file():
+            _LOG.info("Trinucleotide counts file exists, skipping save to avoid overwriting data.")
             trinuc_counts_exists = True
-        save_trinuc_file = True
+            save_trinuc_file = False
 
     if bed:
         _LOG.info("Counting trinucleotide combinations in bed regions")
@@ -207,7 +211,6 @@ def count_trinucleotides(reference_index,
                         trinuc_ref_count = count_trinuc(trinuc, trinuc_ref_count)
 
     # Solution to attribute error (needs to be checked)
-    # TODO remove ref_name from this dict
     elif not trinuc_counts_exists:
         _LOG.info("Counting trinucleotide combinations in reference.")
         for ref_name in matching_chroms:
@@ -220,13 +223,12 @@ def count_trinucleotides(reference_index,
             with open_output(trinuc_counts, 'w') as countfile:
                 _LOG.info('Saving trinuc counts to file...')
                 # Convert all values to writable formats
-                trinuc_ref_count = {str(x): int(y) for x, y in trinuc_ref_count.items()}
-                countfile.write(json.dumps(trinuc_ref_count))
+                pickle.dump(trinuc_ref_count, countfile)
 
     else:
         _LOG.info(f'Loading file: {trinuc_counts}.')
-        with open_input(trinuc_counts) as counts:
-            trinuc_ref_count = json.load(counts)
+        with gzip.open(trinuc_counts, 'rb') as counts:
+            trinuc_ref_count = pickle.load(counts)
         if save_trinuc_file:
             _LOG.warning('Existing trinucelotide file will not be changed.')