From 95dbb427312ab8fdb9cd8f0200d98b7c87f0bde3 Mon Sep 17 00:00:00 2001
From: Ryan Lim <rlim@chanzuckerberg.com>
Date: Thu, 9 Feb 2023 17:42:45 +0000
Subject: [PATCH] converts the report json to a csv

---
 workflows/short-read-mngs/postprocess.wdl     | 63 +++++++++++++++++++
 workflows/short-read-mngs/test/local_test.yml |  1 +
 .../short-read-mngs/test/local_test_viral.yml |  2 +-
 3 files changed, 65 insertions(+), 1 deletion(-)

diff --git a/workflows/short-read-mngs/postprocess.wdl b/workflows/short-read-mngs/postprocess.wdl
index ff36ab126..873ccc442 100644
--- a/workflows/short-read-mngs/postprocess.wdl
+++ b/workflows/short-read-mngs/postprocess.wdl
@@ -303,6 +303,57 @@ task CombineTaxonCounts {
   }
 }
 
+task OutputResultsCSV { 
+  input { 
+    String docker_image_id
+    String s3_wd_uri
+    File counts_json_file 
+    File lineage_csv
+  }
+  command <<<
+  set -euxo pipefail
+  gunzip -c "~{lineage_csv}" > "lineage.csv"
+
+  python3 <<CODE
+  import json
+  import pandas as pd
+  from idseq_dag.util.dict import open_file_db_by_extension
+
+  with open("~{counts_json_file}") as f:
+    taxon_counts = json.load(f)["pipeline_output"]["taxon_counts_attributes"]
+  df = pd.DataFrame(taxon_counts)
+  df = df.loc[df["count_type"] != "merged_NT_NR", ~df.columns.isin(['base_count', 'source_count_type'])]
+  lineage = pd.read_csv("lineage.csv", index_col='taxid', usecols = ['taxid','tax_name'], squeeze=True).to_dict()
+  df["family"] = df.family_taxid.map(lambda x: lineage.get(int(x), x))
+  df["genus"] = df.genus_taxid.map(lambda x: lineage.get(int(x), x))
+  df["taxon_name"] = df.tax_id.map(lambda x: lineage.get(int(x), x))
+
+  column_order = [
+    "tax_id",
+    "tax_level",
+    "taxon_name",
+    "genus", 
+    "family",
+    "count_type",
+    "count",
+    "nonunique_count",
+    "dcr",
+    "percent_identity",
+    "alignment_length",
+    "e_value",
+  ]
+  df[column_order].sort_values(by=["count_type", "tax_level", "count"], ascending=False).to_csv("result.csv", index=None)
+
+  CODE
+  >>>
+  output {
+    File result_csv = "result.csv"
+  }
+  runtime {
+    docker: docker_image_id
+  }
+}
+
 task CombineJson {
   input {
     String docker_image_id
@@ -481,6 +532,8 @@ workflow czid_postprocess {
     File lineage_db = "s3://czid-public-references/taxonomy/2021-01-22/taxid-lineages.db"
     File taxon_blacklist = "s3://czid-public-references/taxonomy/2021-01-22/taxon_blacklist.txt"
     File deuterostome_db = "s3://czid-public-references/taxonomy/2021-01-22/deuterostome_taxids.txt"
+    File lineage_csv = "s3://czid-public-references/ncbi-indexes-prod/2021-01-22/index-generation-2/versioned-taxid-lineages.csv.gz"
+    Boolean output_results_csv = false
     Boolean use_deuterostome_filter = true
     Boolean use_taxon_whitelist = false
     Int min_contig_length = 100
@@ -603,6 +656,16 @@ workflow czid_postprocess {
       ]
   }
 
+  if (output_results_csv) { 
+    call OutputResultsCSV { 
+      input:
+      docker_image_id = docker_image_id,
+      s3_wd_uri = s3_wd_uri,
+      counts_json_file = CombineTaxonCounts.assembly_refined_taxon_counts_with_dcr_json,
+      lineage_csv = lineage_csv
+    }
+  }
+
   call CombineJson {
     input:
       docker_image_id = docker_image_id,
diff --git a/workflows/short-read-mngs/test/local_test.yml b/workflows/short-read-mngs/test/local_test.yml
index ce371dad2..05acae638 100644
--- a/workflows/short-read-mngs/test/local_test.yml
+++ b/workflows/short-read-mngs/test/local_test.yml
@@ -15,6 +15,7 @@ host_filter.human_bowtie2_genome: s3://czid-public-references/host_filter/human/
 host_filter.adapter_fasta: https://raw.githubusercontent.com/broadinstitute/viral-pipelines/master/test/input/clipDb.fasta
 host_filter.max_input_fragments: 9000
 host_filter.max_subsample_fragments: 9000
+postprocess.output_results_csv: true
 non_host_rapsearch2_index: s3://czid-public-references/alignment_indexes/2020-04-20/nr.tar.gz
 non_host_gsnap_index: s3://czid-public-references/alignment_indexes/2020-04-20/nt.tar.gz
 non_host_gsnap_genome_name: nt_k16
diff --git a/workflows/short-read-mngs/test/local_test_viral.yml b/workflows/short-read-mngs/test/local_test_viral.yml
index 6772dfe45..f8a4b2d8e 100644
--- a/workflows/short-read-mngs/test/local_test_viral.yml
+++ b/workflows/short-read-mngs/test/local_test_viral.yml
@@ -18,11 +18,11 @@ host_filter.adapter_fasta: https://raw.githubusercontent.com/broadinstitute/vira
 host_filter.max_input_fragments: 9000
 host_filter.max_subsample_fragments: 9000
 non_host_alignment.accession2taxid_db: s3://czid-public-references/mini-database/alignment_indexes/2020-08-20-viral/viral_accessions2taxid.marisa
-non_host_alignment.alignment_scalability: true
 minimap2_local_db_path: s3://czid-public-references/test/viral-alignment-indexes/viral_nt
 diamond_local_db_path: s3://czid-public-references/test/viral-alignment-indexes/viral_nr
 diamond_args: "mid-sensitive" 
 postprocess.nt_db: s3://czid-public-references/test/viral-alignment-indexes/viral_nt
+postprocess.output_results_csv: true
 postprocess.nt_loc_db: s3://czid-public-references/test/viral-alignment-indexes/viral_nt_loc.marisa
 postprocess.nr_db: s3://czid-public-references/test/viral-alignment-indexes/viral_nr
 postprocess.nr_loc_db: s3://czid-public-references/test/viral-alignment-indexes/viral_nr_loc.marisa