tree_nine.wdl

version 1.0

import "https://raw.githubusercontent.com/aofarrel/SRANWRP/v1.2.1/tasks/processing_tasks.wdl" as processing
import "https://raw.githubusercontent.com/aofarrel/dropkick/1.1.0/dropkick.wdl" as dropkick
import "https://raw.githubusercontent.com/aofarrel/microreact_WDLs/1.0.0/share_projects_with_team_via_file.wdl"
import "./matutils_and_friends.wdl" as matWDLlib
import "./split_cluster_tasks.wdl" as clusterlib

# User notes:
# * If user doesn't define input_tree, a rudimentary 7K sample tree will serve as the base tree. This base tree
#   doesn't represent genetic diversity of MTBC well and should not be used for anything besides quick testing.
# * Should be run with --copy-input-files on miniwdl (required if clustering, may work w/o it if not clustering)

# Dev notes:
# * Anything marked !ForwardReference is using a bogus fallback value with select_first() to coerce to not-optional

workflow Tree_Nine {
	input {
		Array[File] diffs
		File? input_tree
		File? existing_diffs
		File? existing_samples
		String? listener_bucket
		
		# matUtils/UShER options
		Boolean detailed_clades          = false
		Float?  max_low_coverage_sites
		File?   matutils_clade_annotations
		Boolean optimize = true
		String? reroot_to_this_node
		Boolean summarize_tree_before_placing_samples   = false 
		Boolean summarize_tree_after_placing_samples    = false

		# related to clustering/distance matrix
		Boolean identify_clusters              = false
		Boolean cluster_entire_tree            = false
		File? special_samples
		File? persistent_denylist

		# metadata file; expected to be pulled via the FISS API but not strictly required
		File? sample_metadata_tsv
		Boolean strictly_check_metadata = true

		Boolean adhoc = false

		# if you are running with pre-existing clusters, all three of these must be filled in
		# if you are identifying clusters ad-hoc, all three of these must be undefined
		File? persistent_cluster_meta    # vital for persistent clusters
		File? persistent_cluster_ids     # vital for persistent clusters
		File? previous_run_cluster_json  # only used to generate a change report
		
		# related to putting clusters on Microreact
		Boolean upload_clusters_to_microreact  = false
		File? microreact_blank_template_json
		File? microreact_decimated_template_json
		File? microreact_key
		File? microreact_update_template_json

		# non-exclusive ways of sharing your microreact projects
		String? microreact_share_email
		String? microreact_share_team
		
		# rarely used files (see parameter_meta)
		Array[File]? coverage_reports
		File? ref_genome               # do not define this if you're using H37Rv!
		
		# output file names and prefixes, extension not included
		String? comment
		Array[String]? rename_samples
		Boolean datestamp_outs         = true
		String out_prefix              = "bigtree"
		String out_diffs               = "_combined"
		
		# testing functions
		Boolean DEBUG_concat_files_then_exit = false
		File?   DEBUG_override_latest_samples
		File?   DEBUG_override_latest_clusters
		Boolean DEBUG_generate_debug_mr_jsons = false
	}

	parameter_meta {
		comment: "String that gets copied directly to output (useful for Terra data tables)"
		
		diffs: "Array of single-sample MAPLE-formatted diff files from myco"

		input_tree: "The base MAT tree (.pb) that samples will be placed upon; will fall back to a test tree on ~7K TB samples from SRA if not defined (test tree should ONLY be used for quick debugging; it is not representative of MTBC diversity nor proper sample QC)"

		existing_diffs: "A pre-concatenated multi-sample .diff file which will be concatenated with the single-sample diffs input (requires existing_samples)"
		
		existing_samples: "A file listing all samples within the multiple-sample existing_diffs file (requires existing_diffs)"
		
		matutils_clade_annotations: "Two column TSV for clade annotation via matUtils"
		
		cluster_entire_tree: "If true, matrix and cluster all samples on tree; if false, only matrix and cluster special_samples (if defined) or newly added samples."
		
		cluster_max_distance: "Soft-maximum SNP distance between two samples for them to be in the same cluster. NOTE if cluster_max_distance=10, A:B=10, B:C=5, and A:C=15, then all three will still be in a cluster even though A:C is above cluster_max_distance, since both are within 10 of another sample in that cluster."
		
		coverage_reports: "NOT USUALLY NEEDED - Single line text files generated by Lily's vcf to diff script, used to filter samples with low overall coverage. By default, myco's vcf_to_diff.py filters sites per site coverage, so this isn't typically needed here in Tree Nine."
		
		detailed_clades: "usher_sampled_diff -D"
		
		max_low_coverage_sites: "Maximum percentage of low coverage sites a sample can have before throwing it out (requires coverage_reports, does not apply to backmasked diffs)"
		
		special_samples: "Provide an override file containing names of the only samples to consider for matrix and clustering. If this isn't defined, matrixing and clustering is done on either entire tree (if cluster_entire_tree) or all samples with a diff file (if not cluster_entire_tree)."
		
		ref_genome: "Reference genome, equivalent to UShER's ref argument, default is H37Rv (M tuberculosis)"
		
		rename_samples: "For file at index i in diffs[i], rename it to the corresponding string at rename_samples[i]."
		
		reroot_to_this_node: "matUtils extract -y (Reroot the output tree relative to this node, leave blank to not reroot)"
		
		out_prefix: "Prefix for all output files"
		
		upload_clusters_to_microreact: "If you know, you know"
	}

	# Apparent Cromwell bug in Array[Pair[String, String]] means some stuff is temporarily hardcoded.
	#
	# These were previously user-accessible workflow-level outputs:
	# Array[String]? microreact_metadata_columns = ["Epi_Duplication","Year_Collected","Patient_County","State","Country","20_Cluster_Date","10_Cluster_Date","5_Cluster_Date","Lineage_TBProf","Resistance_TBProf","Submitter_Facility","Submitter_Facility_Sample_ID","Sequencing_Facility","Latitude","Longitude"]
	# Array[Pair[String, String]] microreact_metadata_column_renames = [("tbd_strain_per_tbprof", "Lineage_TBProf"), ("tbd_resistance", "Resistance_TBProf")] as a user
	# We now hardcode the metadata columns and do not attempt any column renames.
	Array[String]? microreact_metadata_columns = ["Epi_Duplication","Year_Collected","Patient_County","State","Country","20_Cluster_Date","10_Cluster_Date","5_Cluster_Date","tbd_strain_per_tbprof","tbd_resistance","Submitter_Facility","Submitter_Facility_Sample_ID","Sequencing_Facility","Latitude","Longitude"]
	#
	# This section builds Array[Pair[String, String]] "dictionaries" for TBProfiler lineage replacements per CDPH request, but due to https://github.com/broadinstitute/cromwell/issues/7883
	# I cannot actually pass these into process_metadata without the pipeline crashing. This doesn't happen on miniwdl so I'm reasonably confident this a Cromwell bug. For the time
	# being I'm going turn to leave this here (because for some reason the type checker is fine with unless it actually becomes part of a task) and disable column too, as I'd prefer
	# not to rewrite the process_metadata task.
	# Update: I am now handling column renames within process_CDPH_clusters using two hardcoded renames, and skipping value replacements.
	# Implementation note: if in microreact_metadata_column_renames, use the post-rename name
	String replace_values_in_this_column = "Lineage_TBProf"
	Array[Pair[String, String]] value_replacements_1 = [("La1", "M. bovis (La1)"), ("La1.1", "M. bovis (La1.1)")]
	Array[Pair[String, String]] value_replacements_2 = [("La1.2", "M. bovis not-BCG-but-BCG-like (La1.2; note TBProfiler can call BCG specifically as La1.2.BCG but did not)"), ("La1.2.BCG", "M. bovis BCG (La1.2.BCG)")]
	Array[Pair[String, String]] value_replacements_3 = [("La1.3", "M. bovis (La1.3)"), ("La1.4", "M. bovis (La1.4)"), ("La1.5", "M. bovis (La1.5)"), ("La1.6", "M. bovis (La1.6)")]
	Array[Pair[String, String]] value_replacements_4 = [("La1.7", "M. bovis (La1.7)"), ("La1.7.1", "M. bovis (La1.7.1)"), ("La1.8.1", "M. bovis (La1.8.1)"), ("La1.8.2", "M. bovis (La1.8.2)"), ("La2", "M. caprae (La2)"), ("La3", "M. orygis (La3)")]
	Array[Pair[String, String]] value_replacements = flatten([value_replacements_1, value_replacements_2, value_replacements_3, value_replacements_4])

	call matWDLlib.validate_treenine_inputs as validate_inputs {
		input:
			input_tree = input_tree,
			existing_diffs = existing_diffs,
			existing_samples = existing_samples,
			persistent_cluster_meta = persistent_cluster_meta,
			persistent_cluster_ids = persistent_cluster_ids,
			previous_run_cluster_json = previous_run_cluster_json,
			microreact_blank_template_json = microreact_blank_template_json,
			microreact_decimated_template_json = microreact_decimated_template_json,
			microreact_key = microreact_key,
			microreact_update_template_json = microreact_update_template_json,
			upload_clusters_to_microreact = upload_clusters_to_microreact,
			DEBUG_generate_debug_mr_jsons = DEBUG_generate_debug_mr_jsons,
			ref_genome = ref_genome,
			adhoc = adhoc
	}

	if (defined(sample_metadata_tsv)) {

		# Surely if you reference an optional value within a defined() block, you can input it as a non-optional, right?
		# Not so! You will see here (and elsewhere) bogus select_first() fallbacks because of this quirk of WDL.

		call processing.process_metadata_table as process_metadata {
			input:
				table = select_first([sample_metadata_tsv, diffs[0]]),
				desired_columns = microreact_metadata_columns,
				strict = strictly_check_metadata
		}
	}

	call processing.cat_files as cat_diff_files {
		input:
			new_files_to_concat = diffs,
			out_concat_file = out_prefix + out_diffs,
			keep_only_unique_lines = false,
			keep_only_unique_files = true,                             # STRICTLY NECESSARY UNLESS YOUR DATA *AND* SAMPLE IDS ARE DEDUPLICATED
			new_files_quality_reports = coverage_reports,
			quality_report_removal_threshold = max_low_coverage_sites,
			out_sample_names = "samples_added",
			new_files_override_sample_names = rename_samples,
			king_file = existing_diffs,
			king_file_sample_names = existing_samples,
			and_then_exit_1 = DEBUG_concat_files_then_exit,
			datestamp_main_files = true,  # does not datestamp diffs
			out_concat_extension = ".diff"
	}

	File samples_considered_for_clustering = select_first([special_samples, cat_diff_files.first_lines, usher_sampled_diff.usher_tree]) #!ForwardReference

	# Tree Nine attempts to use a clear naming scheme to make its large number of output files unambigious, but you might have a better
	# system than I do, so I'm going to define all remaining major outfile-controlling variables here so you can edit it easily.

	String empty_string = ""
	if(!(datestamp_outs)) { String no_datestamp = "" }
	String optional_datestamp = select_first([no_datestamp, cat_diff_files.today])
	String presumed_input_mat_basename  = basename(select_first([input_tree, "default_debug-only_basetree"]))
	
	String outfile_annotated_input_tree = "input_" + presumed_input_mat_basename + optional_datestamp + ".pb"
	String outfile_usher_tree_raw       = out_prefix + optional_datestamp + "_raw.pb"
	String outfile_usher_tree_optimized = basename(outfile_usher_tree_raw, "_raw.pb") + "_optimized.pb"
	String outfile_usher_tree_annotated = basename(outfile_usher_tree_raw, "_raw.pb") + "_annotated.pb"
	String outfile_usher_tree_rerooted  = basename(outfile_usher_tree_raw, "_raw.pb") + "_reroot_to_" + select_first([reroot_to_this_node, empty_string]) + ".pb"
	String outfile_taxonium_tree        = basename(outfile_usher_tree_raw, "_raw.pb") + "_taxonium.jsonl.gz"
	String outfile_nextstrain_tree      = basename(outfile_usher_tree_raw, "_raw.pb") + ".json"
	#String outfile_nwk_matutils_tree    = basename(outfile_usher_tree_raw, "_raw.pb") + ".nwk"
	# There is also a nwk tree generated by the clustering script, that one is always called a000000.nwk and should be identical to the matutils one
	
	String outfile_input_tree_summaries               = "input_" + presumed_input_mat_basename + "_"
	String outfile_usher_tree_summaries_before_reroot = out_prefix + optional_datestamp + "_before_reroot_"
	String outfile_usher_tree_summaries_final         = out_prefix + optional_datestamp + "_final_"

	
	if(summarize_tree_before_placing_samples) {
		if (defined(input_tree)) {

			# iff there is a metadata tsv, annotate input tree with it before summarizing
			if (defined(matutils_clade_annotations)) {

				call matWDLlib.annotate as annotate_input_tree {
					input:
						input_mat = select_first([input_tree, usher_sampled_diff.usher_tree]), #!ForwardReference
						metadata_tsv = select_first([matutils_clade_annotations, usher_sampled_diff.usher_tree]), #!ForwardReference
						outfile_mat = outfile_annotated_input_tree
				}
			}

			File possibly_annotated_input_tree = select_first([annotate_input_tree.annotated_tree, input_tree])

			call matWDLlib.summarize as summarize_input_tree {
				input:
					input_mat = possibly_annotated_input_tree,
					prefix_outs = outfile_input_tree_summaries
			}
		}
	}

	call matWDLlib.usher_sampled_diff as usher_sampled_diff {
		input:
			detailed_clades = detailed_clades,
			diff = cat_diff_files.outfile,
			input_mat = input_tree,
			output_mat = outfile_usher_tree_raw,
			ref_genome = ref_genome,
			noop_boolean = validate_inputs.didnt_crash
	}

	if (optimize) {
		call matWDLlib.matOptimize as matOptimize_usher {
			input:
				input_mat = usher_sampled_diff.usher_tree,
				output_mat = outfile_usher_tree_optimized
		}
	}

	File optimized_or_raw_tree = select_first([matOptimize_usher.optimized_tree, usher_sampled_diff.usher_tree])
	

	if (defined(matutils_clade_annotations)) {
		call matWDLlib.annotate as annotate_usher {
			input:
				input_mat = optimized_or_raw_tree,
				metadata_tsv = select_first([matutils_clade_annotations, usher_sampled_diff.usher_tree]), # bogus fallback
				outfile_mat = outfile_usher_tree_annotated
		}
	}

	File possibly_annotated_maximal_output_tree = select_first([annotate_usher.annotated_tree, optimized_or_raw_tree])

	if(defined(reroot_to_this_node)) {

		if(summarize_tree_after_placing_samples) {
			call matWDLlib.summarize as summarize_before_reroot {
				input:
					input_mat = possibly_annotated_maximal_output_tree,
					prefix_outs = outfile_usher_tree_summaries_before_reroot
			}
		}

		call matWDLlib.reroot as reroot_usher {
			input:
				input_mat = possibly_annotated_maximal_output_tree,
				reroot_to_this_node = select_first([reroot_to_this_node, ""]),
				output_mat = outfile_usher_tree_rerooted
		}
	}

	File final_maximal_output_tree = select_first([reroot_usher.rerooted_tree, possibly_annotated_maximal_output_tree])

	# defined(matutils_clade_annotations)   defined(reroot_to_this_node)          final_maximal_output_tree
	# ----------------------------------------------------------------------------------------------------------------
	#       true                       true                    annotated and rerooted
	#       true                      false                    annotated
	#       false                      true                    rerooted
	#       false                     false                    neither, just the output matOptimize_usher.optimized_tree

	call matWDLlib.convert_to_taxonium as to_taxonium {
		input:
			input_mat = final_maximal_output_tree,
			outfile_taxonium = outfile_taxonium_tree
	}

	if (identify_clusters) {
		if (!defined(DEBUG_override_latest_samples)) {
			call clusterlib.find_CDPH_clusters as find_clusters {
				input:
					input_mat_with_new_samples = final_maximal_output_tree,
					special_samples = samples_considered_for_clustering,
					only_matrix_special_samples = !(cluster_entire_tree),
					datestamp = cat_diff_files.today
			}
		}
		
		call clusterlib.process_CDPH_clusters as process_clusters {
			input:
				shareemail = microreact_share_email,
				input_mat_with_new_samples = final_maximal_output_tree,
				special_samples = samples_considered_for_clustering,
				combined_diff_file = cat_diff_files.outfile,
				only_matrix_special_samples = !(cluster_entire_tree),
				persistent_ids = persistent_cluster_ids,
				persistent_cluster_meta = persistent_cluster_meta,
				previous_run_cluster_json = previous_run_cluster_json,
				microreact_key = microreact_key,
				microreact_update_template_json = microreact_update_template_json,
				microreact_blank_template_json = microreact_blank_template_json,
				microreact_decimated_template_json = microreact_decimated_template_json,
				persistent_denylist = persistent_denylist,
				upload_clusters_to_microreact = upload_clusters_to_microreact,
				datestamp = cat_diff_files.today,
				sample_metadata_tsv = process_metadata.processed_metadata_table,
				microreact_metadata_columns = microreact_metadata_columns,
				latest_samples_tsv = select_first([find_clusters.latest_samples_tsv, DEBUG_override_latest_samples]),
				latest_clusters_tsv = select_first([find_clusters.latest_clusters_tsv, DEBUG_override_latest_clusters]),
				cluster_matrices_randomIDs_tarball = find_clusters.cluster_matrices_randomIDs,
				cluster_subtrees_randomIDs_tarball = find_clusters.cluster_subtrees_randomIDs,
				DEBUG_generate_debug_mr_jsons = DEBUG_generate_debug_mr_jsons
		}

		# This is some trickery to prevent Cromwell from complaining about us putting an "optional" output
		# into a non-optional task input. We can do this because the "optional" output actually gets created
		# in non-error cases (at least, this is the case with how we call the task here in Tree Nine)
		# so it will never actually fall back on optimized_or_raw_tree -- and if it does error, the entire
		# pipeline crashes so what happens here is moot.
		File coerced_unclustered_txt = select_first([find_clusters.unclustered_samples, optimized_or_raw_tree])

		if (defined(listener_bucket)) {
			# More coercion workarounds here, this one is even sillier because we're in a defined() block, alas
			# this is required.
			String coerced_destination_bucket = select_first([listener_bucket, "nonsense fallback value"])
			call dropkick.Dropkick_Curl as upload_cluster_json {
				input:
					destination_bucket = coerced_destination_bucket,
					files_to_upload = [process_clusters.final_cluster_information_json]
			}

			call dropkick.Dropkick_Curl as upload_unclustered_txt {
				input:
					destination_bucket = coerced_destination_bucket,
					files_to_upload = [coerced_unclustered_txt]
			}
		}

		if (defined(microreact_share_team)) {
			if (defined(microreact_key)) {
				if (defined(process_clusters.updated_mr_URIs_file)) { # must explictly check as it could be undefined if nothing got updated 
					File coerced_microeract_key = select_first([microreact_key, optimized_or_raw_tree])
					String coerced_microreact_share_team = select_first([microreact_share_team, "nonsense fallback value"])
					File coerced_updated_mr_URIs_file = select_first([process_clusters.updated_mr_URIs_file, optimized_or_raw_tree])
					call share_projects_with_team_via_file.Microreact_Share_Projects_With_Team {
						input:
							token = coerced_microeract_key,
							team_uri = coerced_microreact_share_team,
							project_uris = coerced_updated_mr_URIs_file
					}
				}
			}
		}

		call matWDLlib.convert_to_nextstrain_single_terra_compatiable as to_nextstrain_cluster {
			input:
				input_mat = final_maximal_output_tree,
				outfile_nextstrain = outfile_nextstrain_tree,
				one_metadata_file = process_clusters.samp_cluster_ten
		}
	}

	if (!(identify_clusters)) {
		call matWDLlib.convert_to_nextstrain_single_terra_compatiable as to_nextstrain {
				input:
					input_mat = final_maximal_output_tree,
					outfile_nextstrain = outfile_nextstrain_tree
			}
	}
	

	if(summarize_tree_after_placing_samples) {
		call matWDLlib.summarize as summarize_final {
			input:
				input_mat = final_maximal_output_tree,
				prefix_outs = outfile_usher_tree_summaries_final
		}
	}
	
	output {
		String? out_comment = comment
		File?   unclusted_samples = find_clusters.unclustered_samples

		# big trees - protobuff
		#
		# note that tree_usher_rerooted is annotated if defined(matutils_clade_annotations), but tree_usher_annotated is NOT rerooted
		# even if defined(reroot_to_this_node) -- this was done on purpose so people can get two annotated trees if they
		# want to easily compare the tree before and after rerooting
		#
		File  BIG_tree_usher = optimized_or_raw_tree
		File  BIG_tree_usher_raw_dont_use = usher_sampled_diff.usher_tree
		File? BIG_tree_reroot = reroot_usher.rerooted_tree
		File? BIG_tree_ushanno = annotate_usher.annotated_tree

		# big trees - other formats
		#
		# iff defined(reroot_to_this_node), these are based on usher_tree_rerooted
		# else, these are based on usher_tree_raw (and usher_tree_rerooted doesn't exist)
		#
		File?  BIG_tree_nwk_raw = find_clusters.bigtree_raw
		File?  BIG_tree_nwk_gen = find_clusters.bigtree_gen
		File   BIG_tree_taxonium = to_taxonium.taxonium_tree
		File?  BIG_tree_json_noanno = to_nextstrain.nextstrain_singular_tree
		File?  BIG_tree_json_clusteranno = to_nextstrain_cluster.nextstrain_singular_tree

		# cluster subtrees/matrices -- the ones from process_clusters have the expected persistent IDs, the ones
		# from find_clusters DO NOT, so we'll only include the process_clusters tarballs here. the reason why
		# these are tarballs now is because Terra has a hard limit on the number of workflow outputs.
		File? cluster_subtrees = process_clusters.cluster_trees_persisIDs                          # formerly Array[File] CLUSTER_trees_nwk
		File? cluster_matrices = process_clusters.cluster_matrices_persisIDs                       # formerly Array[File] CLUSTER_dmatrices
		File? cluster_matrices_backmasked= process_clusters.cluster_matrices_persisIDs             # formerly Array[File] BM_CLUSTER_dmatrices
		File? cluster_subtrees_backmasked = process_clusters.cluster_matrices_persisIDs_backmasked # formerly Array[File] BM_CLUSTER_trees_nwk

		# if you want to do persistent clustering in the future, you need all five of these files
		File updated_diff_file = cat_diff_files.outfile
		File updated_diff_contents = samples_considered_for_clustering
		File? updated_persistent_ids = process_clusters.new_persistent_ids
		File? updated_persistent_meta = process_clusters.new_persistent_meta
		File? updated_cluster_information_json = process_clusters.final_cluster_information_json

		#### "stats for the nerds" section, most users don't need these but they're good context ####

		# cat_diff_files
		Int   n_new_samps_input = cat_diff_files.files_input
		Int   n_new_samps_skipped = cat_diff_files.files_removed
		Array[String] samples_dropped = cat_diff_files.removed_files

		# cluster-related
		File? BIG_matrix_nb = find_clusters.bigtree_matrix    # nb as in "not backmasked" although there is no backmasked version
		File? all_samples_nearest_relatives = find_clusters.all_nearest_relatives
		File? all_samples_that_clustered = process_clusters.all_samples_cluster_information
		File? new_samples_that_clustered = process_clusters.new_samples_cluster_information
		Int?  n_20SNP_clusters = find_clusters.n_big_clusters
		Int?  n_samps_unclustered = find_clusters.n_unclustered
		Int?  n_samps_clustered = find_clusters.n_samples_in_clusters
		Int?  n_samps_processed = find_clusters.n_samples_processed
		File? unclustered_subtrees_and_info = find_clusters.unclustered_subtrees_etc
		File? mr_uris_updated = process_clusters.updated_mr_URIs_file  # awkward name because not required for subsequent runs
		
		# tree summary tasks
		File? in_summary = summarize_input_tree.summary
		File? nb_summary_preroot = summarize_before_reroot.summary      # iff defined(reroot_to_this_node)
		File? nb_summary_final = summarize_final.summary
		File? in_list_samples = summarize_input_tree.samples
		File? nb_list_samples_preroot = summarize_before_reroot.samples # iff defined(reroot_to_this_node)
		File? nb_list_samples_final = summarize_final.samples
		

	}
}