add in sc2 and rsv handling dais data processing

Author: mandysulli

src/processes/prepare_mira_reports.rs

@@ -112,7 +112,9 @@ fn read_yaml<R: std::io::Read>(reader: R) -> Result<QCConfig, Box<dyn std::error
 }
 
 pub fn prepare_mira_reports_process(args: ReportsArgs) -> Result<(), Box<dyn Error>> {
-    /////////////// Read in all data ///////////////
+    ///////////////////////////////////////////////////////////////////////////
+    /////////////// Read in and process data from IRMA and Dais ///////////////
+    ///////////////////////////////////////////////////////////////////////////
     // Read in samplesheet
     let samplesheet_path = create_reader(args.samplesheet)?;
     let samplesheet = if &args.platform == "illumina" {
@@ -134,13 +136,12 @@ pub fn prepare_mira_reports_process(args: ReportsArgs) -> Result<(), Box<dyn Err
     let qc_yaml_path = create_reader(args.qc_yaml)?;
     let qc_config: QCConfig = read_yaml(qc_yaml_path)?;
 
-    //Read in IRMA data
+    // Read in IRMA data
     let coverage_data = coverage_data_collection(&args.irma_path, &args.platform, &args.runid)?;
     let read_data = reads_data_collection(&args.irma_path, &args.platform, &args.runid)?;
     let vtype_data = create_vtype_data(&read_data);
     let allele_data = allele_data_collection(&args.irma_path)?;
     let indel_data = indels_data_collection(&args.irma_path)?;
-    //TODO: feed in organism from argument
     //let seq_data = amended_consensus_data_collection(&args.irma_path, virus);
     let ref_lengths = match get_reference_lens(&args.irma_path) {
         Ok(data) => data,
@@ -158,10 +159,13 @@ pub fn prepare_mira_reports_process(args: ReportsArgs) -> Result<(), Box<dyn Err
     let dais_ins_data = dais_insertion_data_collection(&args.irma_path)?;
     let dais_del_data = dias_deletion_data_collection(&args.irma_path)?;
     let dais_seq_data = dias_sequence_data_collection(&args.irma_path)?;
-    //TODO: sc2 handling
-    let dais_ref_data = dais_ref_seq_data_collection(&args.workdir_path, &args.virus)?;
-
-    //TODO: remove these at end
+    let mut dais_ref_data: Vec<DaisSeqData> = Vec::new();
+    if args.virus.to_lowercase() == "flu" || args.virus.to_lowercase() == "rsv" {
+        dais_ref_data = dais_ref_seq_data_collection(&args.workdir_path, &args.virus)?;
+    } else if args.virus.to_lowercase() == "sc2-wgs" {
+        dais_ref_data = dais_ref_seq_data_collection(&args.workdir_path, "sc2")?;
+    }
+    //TODO: remove print statements at end
     //println!("{qc_config:?}")
     //println!("cov data: {coverage_data:?}");
     //println!("Allele data: {allele_data:?}");
@@ -171,22 +175,33 @@ pub fn prepare_mira_reports_process(args: ReportsArgs) -> Result<(), Box<dyn Err
     //println!("dais del data: {dais_del_data:#?}");
     //println!("dais seq data: {dais_seq_data:#?}");
     //println!("dais ref data: {dais_ref_data:#?}");
-    println!("ref length data: {ref_lengths:#?}");
+    //println!("ref length data: {ref_lengths:#?}");
 
     //Calculate AA variants for aavars.csv and dais_vars.json
-    //TODO: circle back for the rsv and sc2 situations
-    let dais_vars_data = compute_dais_variants(&dais_ref_data, &dais_seq_data)?;
+    let mut dais_vars_data: Vec<DaisVarsData> = Vec::new();
+
+    if args.virus.to_lowercase() == "flu" {
+        dais_vars_data = compute_dais_variants(&dais_ref_data, &dais_seq_data)?;
+    } else if args.virus.to_lowercase() == "sc2-wgs"
+        || args.virus.to_lowercase() == "sc2"
+        || args.virus.to_lowercase() == "rsv"
+    {
+        println!("sc2");
+        dais_vars_data = compute_cvv_dais_variants(&dais_ref_data, &dais_seq_data)?;
+    }
+
     negative_qc_statement(
         "/home/xpa3/mira-oxide/test/qc_statement.json",
         &read_data,
         &neg_control_list,
     )?;
 
+    // Calculating the % coverage and median coverage for summary
     let melted_reads_df = melt_reads_data(&read_data);
     let mut calculated_cov_df: Vec<ProcessedCoverage> = Vec::new();
     let mut calculated_position_cov_df: Vec<ProcessedCoverage> = Vec::new();
-    // TODO: add in SC2 handling
-    if args.virus.to_lowercase() == "flu" {
+
+    if args.virus.to_lowercase() == "flu" || args.virus.to_lowercase() == "rsv" {
         calculated_cov_df = process_wgs_coverage_data(&coverage_data, &ref_lengths);
     } else if args.virus.to_lowercase() == "sc2-spike" {
         calculated_cov_df =
@@ -199,10 +214,13 @@ pub fn prepare_mira_reports_process(args: ReportsArgs) -> Result<(), Box<dyn Err
 
     //println!("{dais_vars_data:?}");
     //println!("{melted_reads_df:?}");
-    println!("{calculated_cov_df:?}");
-    println!("{calculated_position_cov_df:?}");
-    /*
+    //println!("{calculated_cov_df:?}");
+    //println!("{calculated_position_cov_df:?}");
+
+    /////////////////////////////////////////////////////////////////////////////
     /////////////// Write the structs to JSON files and CSV files ///////////////
+    /////////////////////////////////////////////////////////////////////////////
+
     // Writing out coverage data
     let coverage_struct_values = vec![
         "Sample",
@@ -384,7 +402,7 @@ pub fn prepare_mira_reports_process(args: ReportsArgs) -> Result<(), Box<dyn Err
         &dais_vars_data,
         &aavars_columns,
         &aavars_columns,
-    )?; */
+    )?;
 
     /////////////// Write the structs to parquet files if flag invoked ///////////////
     // Write fields to parq if flag given

src/utils/data_processing.rs

@@ -215,6 +215,7 @@ pub fn return_seg_data(
     (segments, segset, segcolor)
 }
 
+//////////////// Functions used to process the variants found in dais outputs ///////////////
 // Function to calculate the aa variants - this is specifically for flu right now
 pub fn compute_dais_variants(
     ref_seqs_data: &Vec<DaisSeqData>,
@@ -279,7 +280,112 @@ pub fn compute_dais_variants(
     Ok(unique_entries)
 }
 
-//////////////// Functions using to create irma_summary ///////////////
+/// Compute CVV DAIS Variants
+pub fn compute_cvv_dais_variants(
+    ref_seqs_data: &Vec<DaisSeqData>,
+    sample_seqs_data: &Vec<DaisSeqData>,
+) -> Result<Vec<DaisVarsData>, Box<dyn Error>> {
+    let mut merged_data = merge_sequences(ref_seqs_data, sample_seqs_data)?;
+
+    // Compute AA Variants
+    for entry in &mut merged_data {
+        entry.insertion = compute_aa_variants(&entry.aa_aln, &entry.aa_seq);
+    }
+
+    for entry in &mut merged_data {
+        entry.insertions_shift_frame = if entry.insertion.is_empty() {
+            "0".to_string()
+        } else {
+            entry.insertion.split(',').count().to_string()
+        };
+    }
+
+    // Filter and sort the data - keep first
+    merged_data.sort_by(|a, b| {
+        a.protein
+            .cmp(&b.protein)
+            .then(a.sample_id.cmp(&b.sample_id))
+            .then(a.insertions_shift_frame.cmp(&b.insertions_shift_frame))
+    });
+
+    let mut unique_data = HashMap::new();
+    for entry in merged_data {
+        let key = (entry.sample_id.clone(), entry.protein.clone());
+        unique_data.entry(key).or_insert(entry);
+    }
+
+    // Convert DaisSeqData to DaisVarsData and collect into a Vec
+    let result: Vec<DaisVarsData> = unique_data
+        .into_values()
+        .map(|entry| DaisVarsData {
+            sample_id: entry.sample_id,
+            reference_id: entry.reference.clone(),
+            protein: entry.protein.clone(),
+            aa_variant_count: entry.insertions_shift_frame.parse::<i32>().unwrap_or(0),
+            aa_variants: entry.insertion.clone(),
+        })
+        .collect();
+
+    Ok(result)
+}
+
+/// Merge sequences based on Coordspace and Protein - used by compute_cvv_dais_variants fn
+fn merge_sequences(
+    ref_seqs_data: &Vec<DaisSeqData>,
+    sample_seqs_data: &Vec<DaisSeqData>,
+) -> Result<Vec<DaisSeqData>, Box<dyn Error>> {
+    let mut merged_data = Vec::new();
+
+    for sample_entry in sample_seqs_data {
+        for ref_entry in ref_seqs_data {
+            if sample_entry.reference == ref_entry.reference
+                && sample_entry.protein == ref_entry.protein
+            {
+                // Create a new owned DaisSeqData instance - it was this or lifetimes...
+                let merged_entry = DaisSeqData {
+                    sample_id: sample_entry.sample_id.clone(),
+                    ctype: sample_entry.ctype.clone(),
+                    reference: sample_entry.reference.clone(),
+                    protein: sample_entry.protein.clone(),
+                    vh: sample_entry.vh.clone(),
+                    aa_seq: ref_entry.aa_seq.clone(), // Use ref_entry's AA sequence
+                    aa_aln: sample_entry.aa_aln.clone(),
+                    cds_id: sample_entry.cds_id.clone(),
+                    insertion: sample_entry.insertion.clone(),
+                    insertions_shift_frame: sample_entry.insertions_shift_frame.clone(),
+                    cds_sequence: sample_entry.cds_sequence.clone(),
+                    aligned_cds_sequence: sample_entry.aligned_cds_sequence.clone(),
+                    reference_nt_positions: sample_entry.reference_nt_positions.clone(),
+                    sample_nt_positions: sample_entry.sample_nt_positions.clone(),
+                };
+
+                merged_data.push(merged_entry); // Push the owned value
+            }
+        }
+    }
+
+    Ok(merged_data)
+}
+
+/// Compute AA variants - used by compute_cvv_dais_variants fn
+fn compute_aa_variants(aligned_aa_sequence: &str, ref_aa_sequence: &str) -> String {
+    let mut aa_vars = String::new();
+
+    for (index, (sample_aa, ref_aa)) in aligned_aa_sequence
+        .chars()
+        .zip(ref_aa_sequence.chars())
+        .enumerate()
+    {
+        if sample_aa != ref_aa {
+            let pos = index + 1;
+            let variant = format!("{ref_aa}{pos}{sample_aa}");
+            append_with_comma(&mut aa_vars, &variant);
+        }
+    }
+
+    aa_vars
+}
+//////////////// Functions used to create irma_summary ///////////////
 /// Flip orientation of the reads structs
 pub fn melt_reads_data(records: &Vec<ReadsData>) -> Vec<MeltedRecord> {
     let mut result = Vec::new();