rliterman@0: #!/usr/bin/env python3
rliterman@0: 
rliterman@0: import sys
rliterman@0: import os
rliterman@0: import glob
rliterman@0: import pandas as pd
rliterman@0: from itertools import chain
rliterman@0: import scipy.stats
rliterman@0: import numpy as np
rliterman@0: import datetime
rliterman@0: import time
rliterman@0: import argparse
rliterman@0: 
rliterman@0: def getWarnings(df):
rliterman@0: 
rliterman@0:     df_measures = list(set(df['Measure']))
rliterman@0:     warn_list = []
rliterman@0:     if 'Preserved_Diff' in df_measures:
rliterman@0:         for index,row in df.iterrows():
rliterman@0:             if pd.isna(row['Zscore']):
rliterman@0:                 warn_list.append(np.nan)
rliterman@0:             elif 2.5 <= row['Zscore'] < 3:
rliterman@0:                     warn_list.append("Warning")
rliterman@0:             elif row['Zscore'] >=3:
rliterman@0:                 warn_list.append("Failure")
rliterman@0:             else:
rliterman@0:                 warn_list.append("Pass")
rliterman@0:     elif 'Contig_Count' in df_measures:
rliterman@0:         for index,row in df.iterrows():
rliterman@0:             if pd.isna(row['Zscore']):
rliterman@0:                 warn_list.append(np.nan)
rliterman@0:             elif row['Measure'] in ["Contig_Count","L50","L90"]:
rliterman@0:                 if 2.5 <= row['Zscore'] < 3:
rliterman@0:                     warn_list.append("Warning")
rliterman@0:                 elif row['Zscore'] >=3:
rliterman@0:                     warn_list.append("Failure")
rliterman@0:                 else:
rliterman@0:                     warn_list.append("Pass")
rliterman@0:             elif row['Measure'] == "Assembly_Bases":
rliterman@0:                 if 2.5 <= abs(row['Zscore']) < 3:
rliterman@0:                     warn_list.append("Warning")
rliterman@0:                 elif abs(row['Zscore']) >=3:
rliterman@0:                     warn_list.append("Failure")
rliterman@0:                 else:
rliterman@0:                     warn_list.append("Pass")     
rliterman@0:             elif row['Measure'] in ["N50","N90"]:
rliterman@0:                 if -3 < row['Zscore'] <= -2.5:
rliterman@0:                     warn_list.append("Warning")
rliterman@0:                 elif row['Zscore'] <= -3:
rliterman@0:                     warn_list.append("Failure")
rliterman@0:                 else:
rliterman@0:                     warn_list.append("Pass")
rliterman@0:             else:
rliterman@0:                 sys.exit(f"{row['Measure']}")
rliterman@0:     elif ('Raw_Distance_StdDev' in df_measures) | ('Preserved_Distance_StdDev' in df_measures):
rliterman@0:         for index,row in df.iterrows():
rliterman@0:             if pd.isna(row['Zscore']):
rliterman@0:                 warn_list.append(np.nan)
rliterman@0:             elif 2.5 <= row['Zscore'] < 3:
rliterman@0:                 warn_list.append("Warning")
rliterman@0:             elif row['Zscore'] >=3:
rliterman@0:                 warn_list.append("Failure")
rliterman@0:             else:
rliterman@0:                 warn_list.append("Pass")  
rliterman@0:     elif 'Unique_Kmers' in df_measures:
rliterman@0:         for index,row in df.iterrows():
rliterman@0:             if pd.isna(row['Zscore']):
rliterman@0:                 warn_list.append(np.nan)
rliterman@0:             elif row['Measure'] in ["Align_Percent_Diff","Unique_Kmers","gIndels","Missing_Kmers"]:
rliterman@0:                 if 2.5 <= row['Zscore'] < 3:
rliterman@0:                     warn_list.append("Warning")
rliterman@0:                 elif row['Zscore'] >=3:
rliterman@0:                     warn_list.append("Failure")
rliterman@0:                 else:
rliterman@0:                     warn_list.append("Pass")   
rliterman@0:             elif row['Measure'] in ["Compare_Aligned","Kmer_Similarity","Self_Aligned","Median_Alignment_Length"]:
rliterman@0:                 if -3 < row['Zscore'] <= -2.5:
rliterman@0:                     warn_list.append("Warning")
rliterman@0:                 elif row['Zscore'] <= -3:
rliterman@0:                     warn_list.append("Failure")
rliterman@0:                 else:
rliterman@0:                     warn_list.append("Pass")
rliterman@0:             else:
rliterman@0:                 sys.exit(f"{row['Measure']}")
rliterman@0:     
rliterman@0:     elif ('SNPs_Cocalled' in df_measures) | ('Raw_SNPs_Cocalled' in df_measures) | ('Preserved_SNPs_Cocalled' in df_measures):
rliterman@0:         for index,row in df.iterrows():
rliterman@0:             if pd.isna(row['Zscore']):
rliterman@0:                 warn_list.append(np.nan)
rliterman@0:             elif -3 < row['Zscore'] <= -2.5:
rliterman@0:                     warn_list.append("Warning")
rliterman@0:             elif row['Zscore'] <= -3:
rliterman@0:                     warn_list.append("Failure")
rliterman@0:             else:
rliterman@0:                 warn_list.append("Pass")
rliterman@0:     else:
rliterman@0:         sys.exit(f"{df_measures}")
rliterman@0:     
rliterman@0:     return warn_list
rliterman@0: 
rliterman@0: start_time = time.time()
rliterman@0: 
rliterman@0: # Get args
rliterman@0: parser = argparse.ArgumentParser(description='CSP2 SNP Pipeline Compiler')
rliterman@0: parser.add_argument('--snp_dirs_file', type=str, help='Path to the file containing SNP directories')
rliterman@0: parser.add_argument('--output_directory', type=str, help='Path to the output directory')
rliterman@0: parser.add_argument('--isolate_data_file', type=str, help='Path to the isolate data file')
rliterman@0: parser.add_argument('--mummer_data_file', type=str, help='Path to the MUMmer data file')
rliterman@0: args = parser.parse_args()
rliterman@0: 
rliterman@0: snp_dirs = [line.strip() for line in open(args.snp_dirs_file, 'r')]
rliterman@0: raw_snp_distance_files = list(chain.from_iterable([glob.glob(snp_dir + '/snp_distance_pairwise.tsv') for snp_dir in snp_dirs]))
rliterman@0: screening_files = list(chain.from_iterable([glob.glob(snp_dir + '/Reference_Screening.tsv') for snp_dir in snp_dirs]))
rliterman@0: 
rliterman@0: # Set paths
rliterman@0: output_directory = args.output_directory
rliterman@0: log_file = f"{output_directory}/Compilation.log"
rliterman@0: mean_isolate_file = f"{output_directory}/Mean_Assembly_Stats.tsv"
rliterman@0: isolate_assembly_stats_file = f"{output_directory}/Isolate_Assembly_Stats.tsv"
rliterman@0: align_stats_file =  f"{output_directory}/Isolate_Alignment_Stats.tsv"
rliterman@0: ref_mean_summary_file = f"{output_directory}/Align_Summary_by_Reference.tsv"
rliterman@0: snp_comparison_file = f"{output_directory}/SNP_Distance_Summary.tsv"
rliterman@0: qc_file = f"{output_directory}/QC_Warnings_Failures.tsv"
rliterman@0: 
rliterman@0: with open(log_file,"w+") as log:
rliterman@0:     log.write("CSP2 SNP Pipeline Compiler\n")
rliterman@0:     log.write(str(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"))+"\n")
rliterman@0:     log.write("-------------------------------------------------------\n\n")
rliterman@0:     if len(raw_snp_distance_files) == 0:
rliterman@0:         log.write("\t- CSP2 SNP Pipeline Compiler cannot detected any SNP pipeline output files\n")
rliterman@0:         log.write("\t- Compiler stopping...\n")
rliterman@0:         sys.exit(0)
rliterman@0:     
rliterman@0: isolate_data = pd.read_csv(args.isolate_data_file, sep="\t")
rliterman@0: raw_isolate_count = isolate_data.shape[0]
rliterman@0: 
rliterman@0: mummer_data = pd.read_csv(args.mummer_data_file, sep="\t")
rliterman@0: 
rliterman@0: # Get reference IDs
rliterman@0: reference_ids = list(set(isolate_data[isolate_data['Isolate_Type'] == "Reference"]['Isolate_ID']))
rliterman@0: raw_ref_count = len(reference_ids)
rliterman@0: 
rliterman@0: raw_snp_distance_df = pd.concat([pd.read_csv(file, sep='\t').assign(Reference_ID=os.path.basename(os.path.dirname(file))) for file in raw_snp_distance_files])
rliterman@0: raw_snp_distance_df['Comparison'] = raw_snp_distance_df.apply(lambda row: ';'.join(sorted([str(row['Query_1']), str(row['Query_2'])])), axis=1)
rliterman@0: 
rliterman@0: # Check for preserved data
rliterman@0: preserved_snp_distance_files = list(chain.from_iterable([glob.glob(snp_dir + '/snp_distance_pairwise_preserved.tsv') for snp_dir in snp_dirs]))
rliterman@0: if len(preserved_snp_distance_files) == 0:
rliterman@0:     has_preserved = False
rliterman@0: else:
rliterman@0:     has_preserved = True
rliterman@0:     preserved_snp_distance_df = pd.concat([pd.read_csv(file, sep='\t').assign(Reference_ID=os.path.basename(os.path.dirname(file))) for file in preserved_snp_distance_files])
rliterman@0:     preserved_snp_distance_df['Comparison'] = preserved_snp_distance_df.apply(lambda row: ';'.join(sorted([str(row['Query_1']), str(row['Query_2'])])), axis=1)
rliterman@0:     
rliterman@0: screening_df = pd.concat([pd.read_csv(file, sep='\t').assign(Reference_ID=os.path.basename(os.path.dirname(file))) for file in screening_files])
rliterman@0: 
rliterman@0: snp_isolates = list(set(raw_snp_distance_df['Query_1'].tolist() + raw_snp_distance_df['Query_2'].tolist()))
rliterman@0: 
rliterman@0: with open(log_file,"a+") as log:
rliterman@0:     log.write(f"- Detected SNP distance data for {len(snp_isolates)} isolates out of {raw_isolate_count} total isolates analyzed\n")
rliterman@0:     if len(snp_isolates) <= 2:
rliterman@0:         log.write("\t- CSP2 SNP Pipeline Compiler cannot do much with 2 or fewer isolates\n")
rliterman@0:         log.write("\t- Compiler stopping...\n")
rliterman@0:         sys.exit(0)
rliterman@0:     else:
rliterman@0:         failed_combinations = screening_df.loc[screening_df['Screen_Category'] != "Pass"]
rliterman@0:         failed_comparisons = []
rliterman@0: 
rliterman@0:         if failed_combinations.shape[0] > 0:
rliterman@0:             reference_query_dict = failed_combinations.groupby('Reference_ID')['Query_ID'].apply(list).to_dict()
rliterman@0:             for index,row in failed_combinations.iterrows():
rliterman@0:                 failed_comparisons.append(";".join(sorted([row['Query_ID'], row['Reference_ID']]))) 
rliterman@0:             log.write("\n- The following query-reference alignments did not pass QC\n")
rliterman@0:             for key, value in reference_query_dict.items():
rliterman@0:                 log.write(f"\nReference {key}\n{', '.join(map(str, value))}\n")
rliterman@0: 
rliterman@0: # Prune isolate data
rliterman@0: isolate_data = isolate_data.loc[isolate_data['Isolate_ID'].isin(snp_isolates)]
rliterman@0: reference_ids = [x for x in reference_ids if x in snp_isolates]
rliterman@0: ref_count = len(reference_ids)
rliterman@0: with open(log_file,"a+") as log:
rliterman@0:     log.write(f"- Detected SNP distance data for {ref_count} reference isolates out of {raw_ref_count} total reference isolates analyzed\n")
rliterman@0:     for ref in reference_ids:
rliterman@0:         log.write(f"\t- {ref}\n")
rliterman@0: 
rliterman@0: # Prune MUMmer data
rliterman@0: mummer_data['Comparison'] = mummer_data.apply(lambda row: ';'.join(sorted([str(row['Query_ID']), str(row['Reference_ID'])])), axis=1)
rliterman@0: mummer_data = mummer_data.loc[~mummer_data['Comparison'].isin(failed_comparisons), ['SNPDiffs_File','Query_ID','Reference_ID','Comparison','Reference_Percent_Aligned','Query_Percent_Aligned','Median_Alignment_Length','Kmer_Similarity','Reference_Unique_Kmers','Query_Unique_Kmers','gIndels']]
rliterman@0: max_align_values = np.maximum(mummer_data['Reference_Percent_Aligned'], mummer_data['Query_Percent_Aligned'])
rliterman@0: min_align_values = np.minimum(mummer_data['Reference_Percent_Aligned'], mummer_data['Query_Percent_Aligned'])
rliterman@0: mummer_data['Align_Percent_Diff'] = 100*((max_align_values - min_align_values)/min_align_values)
rliterman@0: 
rliterman@0: # Run basic assembly stats
rliterman@0: isolate_stats = isolate_data.melt(id_vars=['Isolate_ID', 'Isolate_Type'], value_vars = ['Contig_Count','Assembly_Bases','N50','N90','L50','L90'],value_name='Value',var_name = "Measure")
rliterman@0: isolate_stats['Zscore'] = isolate_stats.groupby('Measure')['Value'].transform(scipy.stats.zscore).astype('float').round(3)
rliterman@0: isolate_stats['QC'] = getWarnings(isolate_stats)
rliterman@0: isolate_stats['Value'] = isolate_stats['Value'].astype('int')
rliterman@0: 
rliterman@0: # Reformat for final TSV
rliterman@0: isolate_stats['Min'] = np.nan
rliterman@0: isolate_stats['Max'] = np.nan
rliterman@0: isolate_stats['StdDev'] = np.nan
rliterman@0: isolate_stats = isolate_stats[['Isolate_ID','Isolate_Type','Measure','Min','Value','Max','StdDev','Zscore','QC']].rename(columns = {'Value':'Mean'})
rliterman@0: isolate_stats['Count'] = 1
rliterman@0: 
rliterman@0: # Get mean values
rliterman@0: isolate_mean_df = isolate_stats.groupby(by=['Measure'])['Mean'].agg(Min = 'min',Mean = "mean",Max = 'max',StdDev = 'std',Count = 'count')
rliterman@0: isolate_mean_df['Mean'] = isolate_mean_df['Mean'].astype("int")
rliterman@0: isolate_mean_df['StdDev'] = isolate_mean_df['StdDev'].astype("float").round(3)
rliterman@0: with open(log_file,"a+") as log:
rliterman@0:     log.write("- Read in and processed isolate data\n\n")
rliterman@0:     for index, row in isolate_mean_df.iterrows():
rliterman@0:         log.write(f"{index}:\tMin: {row['Min']}\tMean: {row['Mean']}\tMax: {row['Max']}\tStdDev: {row['StdDev']}\n")
rliterman@0:     log.write("\n")
rliterman@0: 
rliterman@0: # Run basic alignment stats
rliterman@0: isolate_mummer = pd.DataFrame(columns=['Isolate_ID', 'Compare_ID', 'Self_Aligned', 'Compare_Aligned','Align_Percent_Diff','Median_Alignment_Length', 'Kmer_Similarity', 'Unique_Kmers', 'Missing_Kmers', 'gIndels', 'SNPDiffs_File'])
rliterman@0: 
rliterman@0: for isolate in snp_isolates:
rliterman@0:     
rliterman@0:     temp_mummer = mummer_data[(mummer_data['Query_ID'] == isolate) | (mummer_data['Reference_ID'] == isolate)].drop_duplicates(subset=['Comparison']).assign(Isolate_ID = isolate)
rliterman@0:     
rliterman@0:     for index, row in temp_mummer.iterrows():
rliterman@0:         if row['Query_ID'] == isolate:
rliterman@0:             temp_isolate_mummer = row[['Isolate_ID', 'Reference_ID', 'Query_Percent_Aligned', 'Reference_Percent_Aligned','Align_Percent_Diff', 'Median_Alignment_Length', 'Kmer_Similarity', 'Query_Unique_Kmers', 'Reference_Unique_Kmers', 'gIndels', 'SNPDiffs_File']].to_frame().T
rliterman@0:             temp_isolate_mummer.columns = ['Isolate_ID', 'Compare_ID', 'Self_Aligned', 'Compare_Aligned', 'Align_Percent_Diff','Median_Alignment_Length', 'Kmer_Similarity', 'Unique_Kmers', 'Missing_Kmers', 'gIndels', 'SNPDiffs_File']
rliterman@0:             isolate_mummer = pd.concat([isolate_mummer,temp_isolate_mummer])
rliterman@0: 
rliterman@0:         elif row['Reference_ID'] == isolate:
rliterman@0:             temp_isolate_mummer = row[['Isolate_ID', 'Query_ID', 'Reference_Percent_Aligned', 'Query_Percent_Aligned', 'Align_Percent_Diff', 'Median_Alignment_Length', 'Kmer_Similarity', 'Reference_Unique_Kmers', 'Query_Unique_Kmers', 'gIndels', 'SNPDiffs_File']].to_frame().T
rliterman@0:             temp_isolate_mummer.columns = ['Isolate_ID', 'Compare_ID', 'Self_Aligned', 'Compare_Aligned', 'Align_Percent_Diff', 'Median_Alignment_Length', 'Kmer_Similarity', 'Unique_Kmers', 'Missing_Kmers', 'gIndels', 'SNPDiffs_File']
rliterman@0:             isolate_mummer = pd.concat([isolate_mummer,temp_isolate_mummer])
rliterman@0: 
rliterman@0: isolate_mummer['Isolate_Type'] = isolate_mummer['Isolate_ID'].apply(lambda x: 'Reference' if x in reference_ids else 'Query')
rliterman@0: isolate_mummer_df = isolate_mummer.melt(id_vars=['Isolate_ID','Isolate_Type'], value_vars = ['Self_Aligned','Compare_Aligned', 'Align_Percent_Diff','Median_Alignment_Length','Kmer_Similarity','Unique_Kmers','Missing_Kmers','gIndels'],value_name='Value',var_name = "Measure")
rliterman@0: isolate_mummer_df['Value'] = isolate_mummer_df['Value'].astype("float")
rliterman@0: isolate_mummer_df = isolate_mummer_df.groupby(by=['Isolate_ID','Isolate_Type','Measure'])['Value'].agg(Count = "count",Min = "min",Value = "mean",Max = "max",StdDev = 'std').reset_index()
rliterman@0: isolate_mummer_df['Value'] = isolate_mummer_df['Value'].astype("float").round(2)
rliterman@0: 
rliterman@0: # Get Zscores
rliterman@0: isolate_mummer_df['Zscore'] =  isolate_mummer_df.groupby('Measure')['Value'].transform(scipy.stats.zscore).astype('float').round(3)
rliterman@0: isolate_mummer_df['QC'] = getWarnings(isolate_mummer_df)
rliterman@0: 
rliterman@0: # Reformat for final TSV
rliterman@0: align_stats = isolate_mummer_df[['Isolate_ID','Isolate_Type','Measure','Min','Value','Max','StdDev','Zscore','QC','Count']].copy().rename(columns = {"Value":"Mean"})
rliterman@0: align_stats['StdDev'] = align_stats['StdDev'].astype('float').round(3)
rliterman@0: 
rliterman@0: with open(log_file,"a+") as log:
rliterman@0:     log.write("- Read in and processed alignment data\n")
rliterman@0:     
rliterman@0: # Process cocalled data
rliterman@0: raw_cocalled_df = raw_snp_distance_df[['Comparison','Query_1','Query_2','Reference_ID','SNPs_Cocalled']]
rliterman@0: isolate_cocalled_df = pd.DataFrame(columns = ['Isolate_ID','Count','Min','Mean','Max','StdDev'])
rliterman@0: 
rliterman@0: for isolate in snp_isolates:
rliterman@0:     temp_cocalled = raw_cocalled_df[(raw_cocalled_df['Query_1'] == isolate) | (raw_cocalled_df['Query_2'] == isolate)].drop_duplicates(subset=['Comparison','Reference_ID']).assign(Isolate_ID = isolate)
rliterman@0:     temp_cocalled = temp_cocalled.groupby(['Isolate_ID'])['SNPs_Cocalled'].agg(Count = "count", Min = "min", Value = "mean", Max = "max",StdDev = 'std').reset_index()
rliterman@0:     isolate_cocalled_df = pd.concat([isolate_cocalled_df,temp_cocalled])
rliterman@0: 
rliterman@0: isolate_cocalled_df['Measure'] = 'Raw_SNPs_Cocalled'
rliterman@0: isolate_cocalled_df['Value'] = isolate_cocalled_df['Value'].astype('int')
rliterman@0: isolate_cocalled_df['Zscore'] = isolate_cocalled_df['Value'].transform(scipy.stats.zscore).astype('float').round(3)
rliterman@0: isolate_cocalled_df['QC'] = getWarnings(isolate_cocalled_df)
rliterman@0: 
rliterman@0: # Format for final TSV
rliterman@0: isolate_cocalled_df['Isolate_Type'] = isolate_cocalled_df['Isolate_ID'].apply(lambda x: 'Reference' if x in reference_ids else 'Query')
rliterman@0: isolate_cocalled_stats = isolate_cocalled_df[['Isolate_ID','Isolate_Type','Measure','Min','Value','Max','StdDev','Zscore','QC','Count']].copy().rename(columns={'Value':'Mean'})
rliterman@0: 
rliterman@0: if has_preserved:
rliterman@0:     preserved_cocalled_df = preserved_snp_distance_df[['Comparison','Query_1','Query_2','Reference_ID','SNPs_Cocalled']]
rliterman@0:     isolate_preserved_cocalled_df = pd.DataFrame(columns = ['Isolate_ID','Count','Min','Mean','Max','StdDev'])
rliterman@0: 
rliterman@0:     for isolate in snp_isolates:
rliterman@0:         temp_cocalled = preserved_cocalled_df[(preserved_cocalled_df['Query_1'] == isolate) | (preserved_cocalled_df['Query_2'] == isolate)].drop_duplicates(subset=['Comparison','Reference_ID']).assign(Isolate_ID = isolate)
rliterman@0:         temp_cocalled = temp_cocalled.groupby(['Isolate_ID'])['SNPs_Cocalled'].agg(Count = "count", Min = "min", Value = "mean", Max = "max",StdDev = 'std').reset_index()
rliterman@0:         isolate_preserved_cocalled_df = pd.concat([isolate_preserved_cocalled_df,temp_cocalled])
rliterman@0: 
rliterman@0:     isolate_preserved_cocalled_df['Measure'] = 'Preserved_SNPs_Cocalled'
rliterman@0:     isolate_preserved_cocalled_df['Value'] = isolate_preserved_cocalled_df['Value'].astype('int')
rliterman@0:     isolate_preserved_cocalled_df['Zscore'] = isolate_preserved_cocalled_df['Value'].transform(scipy.stats.zscore).astype('float').round(3)
rliterman@0:     isolate_preserved_cocalled_df['QC'] = getWarnings(isolate_preserved_cocalled_df)
rliterman@0: 
rliterman@0:     # Format for final TSV
rliterman@0:     isolate_preserved_cocalled_df['Isolate_Type'] = isolate_preserved_cocalled_df['Isolate_ID'].apply(lambda x: 'Reference' if x in reference_ids else 'Query')
rliterman@0:     isolate_preserved_cocalled_df['StdDev'] = isolate_preserved_cocalled_df['StdDev'].astype('float').round(3)
rliterman@0:     isolate_cocalled_stats = pd.concat([isolate_cocalled_stats,isolate_preserved_cocalled_df[['Isolate_ID','Isolate_Type','Measure','Min','Value','Max','StdDev','Zscore','QC','Count']].copy().rename(columns={'Value':'Mean'})])
rliterman@0:     
rliterman@0: with open(log_file,"a+") as log:
rliterman@0:     log.write("- Processed cocalled SNP data\n")
rliterman@0:     
rliterman@0: if has_preserved:
rliterman@0:     raw_snp_df = raw_snp_distance_df[['Comparison','Query_1','Query_2','Reference_ID','SNP_Distance']].rename(columns = {'SNP_Distance':'Raw_SNP_Distance'})
rliterman@0:     preserved_snp_df = preserved_snp_distance_df[['Comparison','Query_1','Query_2','Reference_ID','SNP_Distance']].rename(columns = {'SNP_Distance':'Preserved_SNP_Distance'})
rliterman@0:     snp_df = pd.merge(raw_snp_df,preserved_snp_df,how="left",on=['Comparison','Query_1','Query_2','Reference_ID'])
rliterman@0:     snp_df['Preserved_Diff'] = abs(snp_df['Preserved_SNP_Distance'] - snp_df['Raw_SNP_Distance'])
rliterman@0:         
rliterman@0:     isolate_snp_df = pd.DataFrame(columns = ['Isolate_ID','Count','Min','Value','Max','StdDev'])
rliterman@0: 
rliterman@0:     for isolate in snp_isolates:
rliterman@0:         temp_snp = snp_df[(snp_df['Query_1'] == isolate) | (snp_df['Query_2'] == isolate)].drop_duplicates(subset=['Comparison','Reference_ID']).assign(Isolate_ID = isolate)
rliterman@0:         temp_snp = temp_snp.groupby(['Isolate_ID'])['Preserved_Diff'].agg(Count = "count", Min = "min", Value = "mean", Max = "max",StdDev = 'std').reset_index()
rliterman@0:         isolate_snp_df = pd.concat([isolate_snp_df,temp_snp])
rliterman@0: 
rliterman@0:     isolate_snp_df['Measure'] = 'Preserved_Diff'
rliterman@0:     isolate_snp_df['Value'] = isolate_snp_df['Value'].astype("float")
rliterman@0:     isolate_snp_df['Zscore'] = isolate_snp_df['Value'].transform(scipy.stats.zscore).astype('float').round(3)
rliterman@0:     isolate_snp_df['Value'] = isolate_snp_df['Value'].astype('float').round(3)
rliterman@0:     isolate_snp_df['QC'] = getWarnings(isolate_snp_df)
rliterman@0: 
rliterman@0:     # Format for final TSV
rliterman@0:     isolate_snp_df['Isolate_Type'] = isolate_snp_df['Isolate_ID'].apply(lambda x: 'Reference' if x in reference_ids else 'Query')
rliterman@0:     isolate_snp_df['StdDev'] = isolate_snp_df['StdDev'].astype('float').round(3)
rliterman@0:     isolate_snp_stats = isolate_snp_df[['Isolate_ID','Isolate_Type','Measure','Min','Value','Max','StdDev','Zscore','QC','Count']].copy().rename(columns={'Value':'Mean'})
rliterman@0:     with open(log_file,"a+") as log:
rliterman@0:         log.write("- Processed preserved SNP data\n")
rliterman@0: else:
rliterman@0:     with open(log_file,"a+") as log:
rliterman@0:         log.write("- No preserved SNP data to process\n")
rliterman@0:     isolate_snp_stats = pd.DataFrame(columns = ['Isolate_ID','Isolate_Type','Measure','Min','Value','Max','StdDev','Zscore','QC','Count'])
rliterman@0: 
rliterman@0: # Compare SNPs across refs
rliterman@0: if len(reference_ids) == 1:
rliterman@0:     isolate_stdev_stats = pd.DataFrame(columns =['Isolate_ID','Isolate_Type','Measure','Min','Mean','Max','StdDev','Zscore','QC','Count'])
rliterman@0:     with open(log_file,"a+") as log:
rliterman@0:         log.write("- 1 reference provided, SNP distances have no comparisons\n")
rliterman@0: else:
rliterman@0:     # Get comparison stats
rliterman@0:     raw_comparison_df = raw_snp_distance_df.groupby(by=['Comparison'])['SNP_Distance'].agg(Count = 'count', Min = 'min', Mean = 'mean', Max = 'max', StdDev = 'std').reset_index()
rliterman@0:     raw_comparison_df['StdDev'] = raw_comparison_df['StdDev'].astype('float').round(3)
rliterman@0:     raw_comparison_df['Mean'] = raw_comparison_df['Mean'].astype('int')
rliterman@0:     raw_comparison_df[['Query_1', 'Query_2']] = raw_comparison_df['Comparison'].str.split(';', expand=True)
rliterman@0:     raw_comparison_df['SNP_Spread'] = abs(raw_comparison_df['Max'] - raw_comparison_df['Min'])
rliterman@0: 
rliterman@0:     comparison_df = raw_comparison_df[['Comparison','Query_1','Query_2','Mean','StdDev','Min','Max','SNP_Spread','Count']].copy()
rliterman@0:     
rliterman@0:     # Get isolate stats
rliterman@0:     isolate_stdev_df = pd.DataFrame(columns = ['Isolate_ID','Count','Min','Value','Max','StdDev'])
rliterman@0: 
rliterman@0:     for isolate in snp_isolates:
rliterman@0:         temp_compare = raw_comparison_df[(raw_comparison_df['Query_1'] == isolate) | (raw_comparison_df['Query_2'] == isolate)].drop_duplicates(subset=['Comparison']).assign(Isolate_ID = isolate)
rliterman@0:         temp_compare = temp_compare.groupby(by=['Isolate_ID'])['StdDev'].agg(Count = 'count', Min = 'min', Value = 'mean', Max = 'max', StdDev = 'std').reset_index()
rliterman@0:         isolate_stdev_df = pd.concat([isolate_stdev_df,temp_compare])
rliterman@0:     
rliterman@0:     isolate_stdev_df['Measure'] = "Raw_Distance_StdDev"
rliterman@0:     isolate_stdev_df['Value'] = isolate_stdev_df['Value'].astype("float")
rliterman@0:     isolate_stdev_df['Zscore'] = isolate_stdev_df['Value'].transform(scipy.stats.zscore).astype('float').round(3)
rliterman@0:     isolate_stdev_df['Value'] = isolate_stdev_df['Value'].astype('float').round(3)
rliterman@0:     isolate_stdev_df['Isolate_Type'] = isolate_stdev_df['Isolate_ID'].apply(lambda x: 'Reference' if x in reference_ids else 'Query')
rliterman@0:     isolate_stdev_df['QC'] = getWarnings(isolate_stdev_df)
rliterman@0:     
rliterman@0:     isolate_stdev_stats = isolate_stdev_df[['Isolate_ID','Isolate_Type','Measure','Min','Value','Max','StdDev','Zscore','QC','Count']].copy().rename(columns={'Value':'Mean'})
rliterman@0:         
rliterman@0:     if has_preserved:
rliterman@0:         comparison_df.columns = ['Comparison','Query_1','Query_2','Raw_Mean','Raw_StdDev','Raw_Min','Raw_Max','Raw_SNP_Spread','Raw_Count']
rliterman@0: 
rliterman@0:         preserved_comparison_df = preserved_snp_distance_df.groupby(by=['Comparison'])['SNP_Distance'].agg(Preserved_Count = 'count', Preserved_Min = 'min', Preserved_Mean = 'mean', Preserved_Max = 'max', Preserved_StdDev = 'std').reset_index()
rliterman@0:         preserved_comparison_df['Preserved_StdDev'] = preserved_comparison_df['Preserved_StdDev'].astype('float').round(3)
rliterman@0:         preserved_comparison_df['Preserved_Mean'] = preserved_comparison_df['Preserved_Mean'].astype('int')
rliterman@0:         preserved_comparison_df[['Query_1', 'Query_2']] = preserved_comparison_df['Comparison'].str.split(';', expand=True)
rliterman@0:         preserved_comparison_df['Preserved_SNP_Spread'] = abs(preserved_comparison_df['Preserved_Max'] - preserved_comparison_df['Preserved_Min'])
rliterman@0:         
rliterman@0:         comparison_df = comparison_df.merge(preserved_comparison_df,how = "left", on=['Comparison','Query_1','Query_2'])
rliterman@0:         comparison_df['Mean_Preserved_Diff'] = abs(comparison_df['Preserved_Mean'] - comparison_df['Raw_Mean'])
rliterman@0:         comparison_df = comparison_df[['Query_1','Query_2','Raw_Mean','Preserved_Mean','Mean_Preserved_Diff','Raw_StdDev','Preserved_StdDev','Raw_SNP_Spread','Preserved_SNP_Spread','Raw_Min','Raw_Max','Preserved_Min','Preserved_Max','Raw_Count','Preserved_Count']]
rliterman@0:         
rliterman@0:         isolate_stdev_df = pd.DataFrame(columns = ['Isolate_ID','Count','Min','Value','Max','StdDev'])
rliterman@0: 
rliterman@0:         for isolate in snp_isolates:
rliterman@0:             temp_compare = preserved_comparison_df[(preserved_comparison_df['Query_1'] == isolate) | (preserved_comparison_df['Query_2'] == isolate)].drop_duplicates(subset=['Comparison']).assign(Isolate_ID = isolate)
rliterman@0:             temp_compare = temp_compare.groupby(by=['Isolate_ID'])['Preserved_StdDev'].agg(Count = 'count', Min = 'min', Value = 'mean', Max = 'max', StdDev = 'std').reset_index()
rliterman@0:             isolate_stdev_df = pd.concat([isolate_stdev_df,temp_compare])
rliterman@0:     
rliterman@0:         isolate_stdev_df['Measure'] = "Preserved_Distance_StdDev"
rliterman@0:         isolate_stdev_df['Value'] = isolate_stdev_df['Value'].astype('float')
rliterman@0:         isolate_stdev_df['Zscore'] = isolate_stdev_df['Value'].transform(scipy.stats.zscore).astype('float').round(3)
rliterman@0:         isolate_stdev_df['Value'] = isolate_stdev_df['Value'].astype('float').round(3)
rliterman@0:         isolate_stdev_df['Isolate_Type'] = isolate_stdev_df['Isolate_ID'].apply(lambda x: 'Reference' if x in reference_ids else 'Query')
rliterman@0:         isolate_stdev_df['QC'] = getWarnings(isolate_stdev_df)
rliterman@0:         isolate_stdev_stats = pd.concat([isolate_stdev_stats,isolate_stdev_df[['Isolate_ID','Isolate_Type','Measure','Min','Value','Max','StdDev','Zscore','QC','Count']].copy().rename(columns={'Value':'Mean'})])
rliterman@0:         with open(log_file,"a+") as log:
rliterman@0:             log.write("- Compared results across references\n")
rliterman@0:     else:
rliterman@0:         with open(log_file,"a+") as log:
rliterman@0:             log.write("- Compared results across references\n")
rliterman@0: # Group by ref
rliterman@0: 
rliterman@0: #### Isolate ####
rliterman@0: ref_isolate_df = isolate_stats.loc[isolate_stats['Isolate_Type'] == "Reference"][['Isolate_ID','Measure','Mean','StdDev','Min','Max','Zscore','QC','Count']].rename(columns = {'Isolate_ID':'Reference_ID'})
rliterman@0: 
rliterman@0: #### StdDev ####
rliterman@0: ref_stdev_df = isolate_stdev_stats.loc[isolate_stdev_stats['Isolate_Type'] == "Reference"][['Isolate_ID','Measure','Mean','StdDev','Min','Max','Zscore','QC','Count']].rename(columns = {'Isolate_ID':'Reference_ID'})
rliterman@0: 
rliterman@0: #### MUMmer ####
rliterman@0: ref_mummer_df = pd.DataFrame(columns = ['Reference_ID','Measure','Mean','StdDev','Min','Max','Count'])
rliterman@0: for ref in reference_ids:
rliterman@0:     ref_mummer =  isolate_mummer[(isolate_mummer['Isolate_ID'] == ref) | (isolate_mummer['Compare_ID'] == ref)].assign(Focal_Reference = ref)
rliterman@0:     ref_mummer['Comparison'] = ref_mummer.apply(lambda row: ';'.join(sorted([str(row['Isolate_ID']), str(row['Compare_ID'])])), axis=1)
rliterman@0:     ref_mummer = ref_mummer.drop_duplicates(subset=['Comparison'])
rliterman@0: 
rliterman@0:     ref_mummer = ref_mummer.melt(id_vars=['Focal_Reference','Isolate_ID','Compare_ID'], value_vars = ['Align_Percent_Diff','Median_Alignment_Length','Kmer_Similarity','gIndels'],value_name='Value',var_name = "Measure")
rliterman@0:     ref_mummer['Value'] = ref_mummer['Value'].astype("float")
rliterman@0:     ref_mummer = ref_mummer.groupby(by=['Measure'])['Value'].agg(Count = "count",Min = "min",Mean = "mean",Max = "max",StdDev = 'std').reset_index().assign(Reference_ID = ref)
rliterman@0:     ref_mummer = ref_mummer[['Reference_ID','Measure','Mean','StdDev','Min','Max','Count']]
rliterman@0:     ref_mummer_df = pd.concat([ref_mummer_df,ref_mummer])
rliterman@0: ref_mummer_df['QC'] = np.nan
rliterman@0: ref_mummer_df['Zscore'] = np.nan
rliterman@0: 
rliterman@0: ref_mummer_summary_df = pd.concat([ref_mummer_df[['Reference_ID','Measure','Mean','StdDev','Min','Max','Zscore','QC','Count']],align_stats.loc[(align_stats['Isolate_Type'] == "Reference") & (align_stats['Measure'].isin(['Self_Aligned','Compare_Aligned','Unique_Kmers','Missing_Kmers']))][['Isolate_ID','Measure','Mean','StdDev','Min','Max','Zscore','QC','Count']].rename(columns = {'Isolate_ID':'Reference_ID'})])
rliterman@0: 
rliterman@0: #### Cocalled ####
rliterman@0: ref_cocalled_summary_df = raw_cocalled_df.groupby(by=['Reference_ID'])['SNPs_Cocalled'].agg(Mean = "mean",StdDev = 'std',Min = "min",Max = "max",Count = 'count').reset_index()
rliterman@0: ref_cocalled_summary_df['Measure'] = "Raw_SNPs_Cocalled"
rliterman@0: ref_cocalled_summary_df['QC'] = np.nan
rliterman@0: ref_cocalled_summary_df['Zscore'] = np.nan
rliterman@0: ref_cocalled_summary_df = ref_cocalled_summary_df[['Reference_ID','Measure','Mean','StdDev','Min','Max','Zscore','QC','Count']]
rliterman@0: 
rliterman@0: if has_preserved:
rliterman@0:     preserved_cocalled_summary = preserved_cocalled_df.groupby(by=['Reference_ID'])['SNPs_Cocalled'].agg(Mean = "mean",StdDev = 'std',Min = "min",Max = "max",Count = 'count').reset_index()
rliterman@0:     preserved_cocalled_summary['Measure'] = "Preserved_SNPs_Cocalled"
rliterman@0:     preserved_cocalled_summary['QC'] = np.nan
rliterman@0:     preserved_cocalled_summary['Zscore'] = np.nan
rliterman@0:     ref_cocalled_summary_df = pd.concat([ref_cocalled_summary_df,preserved_cocalled_summary[['Reference_ID','Measure','Mean','StdDev','Min','Max','Zscore','QC','Count']]])
rliterman@0: 
rliterman@0: #### Preserved Diff ####
rliterman@0: if has_preserved:
rliterman@0:     ref_summary_preserved_df = snp_df.groupby(by=['Reference_ID'])['Preserved_Diff'].agg(Mean = "mean",StdDev = 'std',Min = "min",Max = "max",Count = 'count').reset_index()
rliterman@0:     ref_summary_preserved_df['Measure'] = "Preserved_Diff"
rliterman@0:     ref_summary_preserved_df['QC'] = np.nan
rliterman@0:     ref_summary_preserved_df['Zscore'] = np.nan
rliterman@0:     ref_summary_preserved_df = ref_summary_preserved_df[['Reference_ID','Measure','Mean','StdDev','Min','Max','Count','Zscore','QC']].copy()
rliterman@0:     
rliterman@0: #### Compile ####
rliterman@0: ref_summary_df = pd.concat([ref_mummer_summary_df,
rliterman@0:                             ref_cocalled_summary_df,
rliterman@0:                             ref_isolate_df,ref_stdev_df]).sort_values(by=['Measure'])
rliterman@0: 
rliterman@0: ref_summary_df['Mean'] = ref_summary_df['Mean'].astype("float").round(3)
rliterman@0: ref_summary_df['Min'] = ref_summary_df['Min'].astype("float").round(3)
rliterman@0: ref_summary_df['Max'] = ref_summary_df['Max'].astype("float").round(3)
rliterman@0: ref_summary_df['StdDev'] = ref_summary_df['StdDev'].astype("float").round(3)
rliterman@0: 
rliterman@0: # Catch warnings and failures
rliterman@0: all_isolate_stats = pd.concat([isolate_stats,align_stats,isolate_stdev_stats,isolate_cocalled_stats,isolate_snp_stats]).sort_values(by=['Zscore'])
rliterman@0: 
rliterman@0: warn_fail_df = all_isolate_stats.loc[all_isolate_stats['QC'].isin(['Failure','Warning'])].copy()
rliterman@0: warn_fail_df['abs_Zscore'] = warn_fail_df['Zscore'].abs()
rliterman@0: warn_fail_df = warn_fail_df.sort_values(by='abs_Zscore',ascending=False).drop('abs_Zscore',axis=1)
rliterman@0: 
rliterman@0: warn_fail_isolates = list(set(warn_fail_df['Isolate_ID']))
rliterman@0: if len(warn_fail_isolates) > 0:
rliterman@0:     with open(log_file,"a+") as log:
rliterman@0:         log.write("\n- The following samples had QC warnings or failures:\n")
rliterman@0:         for isolate in warn_fail_isolates:
rliterman@0:             isolate_warn_fail = warn_fail_df.loc[warn_fail_df['Isolate_ID'] == isolate]
rliterman@0:             
rliterman@0:             if isolate in reference_ids:
rliterman@0:                 log.write(f"\n{isolate} (Reference):\n")
rliterman@0:             else:
rliterman@0:                 log.write(f"\n{isolate} (Query):\n")
rliterman@0:             
rliterman@0:             for index,row in isolate_warn_fail.iterrows():
rliterman@0:                 log.write(f"\t- {row['Measure']} - Mean: {row['Mean']}; Zscore: {row['Zscore']}; QC: {row['QC']}\n")
rliterman@0: else:
rliterman@0:     with open(log_file,"a+") as log:
rliterman@0:         log.write("-There were no QC warnings or failures\n")
rliterman@0: 
rliterman@0: # Output data
rliterman@0: 
rliterman@0: # Mean assembly stats
rliterman@0: isolate_mean_df.reset_index().to_csv(mean_isolate_file,sep='\t',index=False)
rliterman@0: 
rliterman@0: # Isolate assembly stats
rliterman@0: isolate_assembly_stats = isolate_stats.loc[isolate_stats['Measure'].isin(['Contig_Count','Assembly_Bases','L50','L90','N50','N90'])].drop(['Min','Max','StdDev','Count'],axis=1).rename(columns = {'Mean':'Value'})
rliterman@0: isolate_assembly_stats.to_csv(isolate_assembly_stats_file,sep='\t',index=False)
rliterman@0: 
rliterman@0: # Isolate alignment stats
rliterman@0: isolate_align_stats = pd.concat([align_stats,isolate_cocalled_stats,isolate_snp_stats,isolate_stdev_stats]).reset_index(drop=True)
rliterman@0: for col in ['Min', 'Mean', 'Max', 'StdDev', 'Zscore']:
rliterman@0:     isolate_align_stats[col] = isolate_align_stats[col].astype("float").round(3)
rliterman@0: isolate_align_stats.to_csv(align_stats_file,sep='\t',index=False)
rliterman@0: 
rliterman@0: # Reference Assembly Stats
rliterman@0: ref_align_summary_df = ref_summary_df.loc[(~ref_summary_df['Measure'].isin(['Contig_Count','Assembly_Bases','L50','L90','N50','N90'])) & (~pd.isna(ref_summary_df['Zscore']))]
rliterman@0: ref_mean_summary_df = ref_summary_df.loc[(~ref_summary_df['Measure'].isin(['Contig_Count','Assembly_Bases','L50','L90','N50','N90'])) & (pd.isna(ref_summary_df['Zscore']))].drop(['Zscore','QC'],axis =1)
rliterman@0: ref_mean_summary_df['Zscore'] = np.nan
rliterman@0: ref_mean_summary_df['QC'] = np.nan
rliterman@0: 
rliterman@0: # Add alignment stats
rliterman@0: if has_preserved:
rliterman@0:     ref_mean_summary_df = pd.concat([ref_mean_summary_df,ref_summary_preserved_df])
rliterman@0: 
rliterman@0: ref_isolate_align_stats = align_stats.loc[(align_stats['Isolate_Type'] == "Reference") & (align_stats['Measure'].isin(['Self_Aligned','Compare_Aligned']))].drop(['Isolate_Type'],axis=1).rename(columns = {'Isolate_ID':'Reference_ID'})[['Reference_ID','Measure','Mean','StdDev','Min','Max','Count','Zscore','QC']]
rliterman@0: 
rliterman@0: ref_mean_summary_stats = pd.concat([ref_mean_summary_df,ref_isolate_align_stats])
rliterman@0: ref_mean_summary_stats.to_csv(ref_mean_summary_file,sep='\t',index=False)
rliterman@0: 
rliterman@0: end_time = time.time()
rliterman@0: 
rliterman@0: with open(log_file,"a+") as log:
rliterman@0:     log.write(f"\n- Completed compilation in {end_time - start_time:.2f} seconds\n")
rliterman@0:     log.write(f"\t- Saved mean isolate assembly data to {mean_isolate_file}\n")
rliterman@0:     log.write(f"\t- Saved raw isolate assembly data to {isolate_assembly_stats_file}\n")
rliterman@0:     log.write(f"\t- Saved isolate alignment data to {align_stats_file}\n")
rliterman@0:     log.write(f"\t- Saved reference summary data to {ref_mean_summary_file}\n")
rliterman@0:     
rliterman@0:     # Comparisons if multiple refs
rliterman@0:     if len(reference_ids) > 1:
rliterman@0:         comparison_df.to_csv(snp_comparison_file,sep="\t",index = False)
rliterman@0:         log.write(f"\t- Saved SNP distance comparisons across references to {snp_comparison_file}\n")
rliterman@0:     
rliterman@0:     # Failures/warnings
rliterman@0:     if warn_fail_df.shape[0] > 0:
rliterman@0:         warn_fail_df.to_csv(qc_file,sep="\t",index=False)
rliterman@0:         log.write(f"\t- Saved QC warnings/failures to {qc_file}\n")