Merge pull request #73 from broadinstitute/dp-beast

BEAST WDL workflows

pipes/WDL/tasks/tasks_interhost.wdl

@@ -39,6 +39,7 @@ task multi_align_mafft_ref {
     docker: "${docker}"
     memory: select_first([machine_mem_gb, 28]) + " GB"
     cpu: 8
+    disks: "local-disk 200 HDD"
     dx_instance_type: "mem2_ssd1_v2_x8"
   }
 }
@@ -80,10 +81,54 @@ task multi_align_mafft {
     docker: "${docker}"
     memory: select_first([machine_mem_gb, 7]) + " GB"
     cpu: 8
+    disks: "local-disk 200 HDD"
     dx_instance_type: "mem1_ssd1_v2_x8"
   }
 }
 
+task beast {
+  input {
+    File     beauti_xml
+
+    String   docker="quay.io/broadinstitute/beast-beagle-cuda"
+  }
+
+  # TO DO: parameterize gpuType and gpuCount
+
+  command {
+    set -e
+    beast -beagle_info
+    nvidia-smi
+    bash -c "sleep 60; nvidia-smi" &
+    beast \
+      -beagle_multipartition off \
+      -beagle_GPU -beagle_cuda -beagle_SSE \
+      -beagle_double -beagle_scaling always \
+      -beagle_order 1,2,3,4 \
+      ${beauti_xml}
+  }
+
+  output {
+    File        beast_log    = glob("*.log")[0]
+    Array[File] trees        = glob("*.trees")
+    File        beast_stdout = stdout()
+  }
+
+  runtime {
+    docker: "${docker}"
+    memory: "7 GB"
+    cpu:    4
+    disks: "local-disk 300 HDD"
+    bootDiskSizeGb: 50
+    gpu:              true                # dxWDL
+    acceleratorType:  "nvidia-tesla-k80"  # GCP PAPIv2
+    acceleratorCount: 4                   # GCP PAPIv2
+    gpuType:          "nvidia-tesla-k80"  # Terra
+    gpuCount:         4                   # Terra
+    nvidiaDriverVersion: "396.37"
+  }
+}
+
 task index_ref {
   input {
     File     referenceGenome
@@ -111,7 +156,9 @@ task index_ref {
   }
   runtime {
     docker: "${docker}"
+    cpu: 2
     memory: select_first([machine_mem_gb, 4]) + " GB"
+    disks: "local-disk 100 HDD"
   }
 }
 
@@ -135,6 +182,8 @@ task trimal_clean_msa {
   runtime {
     docker: "${docker}"
     memory: select_first([machine_mem_gb, 7]) + " GB"
+    cpu: 4
+    disks: "local-disk 100 HDD"
     dx_instance_type: "mem1_ssd1_v2_x8"
   }
 }

pipes/WDL/tasks/tasks_nextstrain.wdl

@@ -399,39 +399,75 @@ task translate_augur_tree {
     }
 }
 
+task augur_import_beast {
+    meta {
+        description: "Import BEAST tree into files ready for augur export, including a Newick-formatted tree and node data in json format. See both https://nextstrain-augur.readthedocs.io/en/stable/faq/import-beast.html and https://nextstrain-augur.readthedocs.io/en/stable/usage/cli/import.html for details."
+    }
+    input {
+        File    beast_mcc_tree
+
+        Float?  most_recent_tip_date
+        String? tip_date_regex
+        String? tip_date_format
+        String? tip_date_delimiter
+
+        Int?    machine_mem_gb
+        String  docker = "nextstrain/base"
+    }
+    String tree_basename = basename(beast_mcc_tree, ".tree")
+    command {
+        augur import beast \
+            --mcc "~{beast_mcc_tree}" \
+            --output-tree "~{tree_basename}.nwk" \
+            --output-node-data "~{tree_basename}.json" \
+            ~{"--most-recent-tip-date " + most_recent_tip_date} \
+            ~{"--tip-date-regex " + tip_date_regex} \
+            ~{"--tip-date-format " + tip_date_format} \
+            ~{"--tip-date-delimeter " + tip_date_delimiter}
+    }
+    runtime {
+        docker: docker
+        memory: select_first([machine_mem_gb, 3]) + " GB"
+        cpu :   2
+        disks:  "local-disk 50 HDD"
+        dx_instance_type: "mem1_ssd1_v2_x2"
+        preemptible: 2
+    }
+    output {
+        File tree_newick    = "~{tree_basename}.nwk"
+        File node_data_json = "~{tree_basename}.json"
+    }
+}
+
 task export_auspice_json {
     meta {
-        description: "export augur files to json suitable for auspice visualization. See https://nextstrain-augur.readthedocs.io/en/stable/usage/cli/export.html"
+        description: "export augur files to json suitable for auspice visualization. The metadata tsv input is generally required unless the node_data_jsons comprehensively capture all of it. See https://nextstrain-augur.readthedocs.io/en/stable/usage/cli/export.html"
     }
     input {
-        File   auspice_config
-        File   metadata
-        File   refined_tree
-        File?  branch_lengths
-        File?  traits
-        File?  nt_muts
-        File?  aa_muts
-        File?  lat_longs_tsv
-        File?  colors_tsv
-        String basename
+        File        auspice_config
+        File?       metadata
+        File        tree
+        Array[File] node_data_jsons
+
+        File?       lat_longs_tsv
+        File?       colors_tsv
 
         Int?   machine_mem_gb
         String docker = "nextstrain/base"
     }
+    String out_basename = basename(basename(tree, ".nwk"), "_refined_tree")
     command {
-
         NODE_DATA_FLAG=""
-        if [[ -n "~{branch_lengths}" || -n "~{traits}" || -n "~{nt_muts}" || -n "~{aa_muts}" ]]; then
+        if [ -n "~{sep=' ' node_data_jsons}" ]; then
           NODE_DATA_FLAG="--node-data "
         fi
-
-        augur export v2 --tree ~{refined_tree} \
-            --metadata ~{metadata} \
-            $NODE_DATA_FLAG ~{sep=' ' select_all([branch_lengths,traits,nt_muts,aa_muts])}\
+        augur export v2 --tree ~{tree} \
+            ~{"--metadata " + metadata} \
+            $NODE_DATA_FLAG ~{sep=' ' node_data_jsons}\
             --auspice-config ~{auspice_config} \
             ~{"--lat-longs " + lat_longs_tsv} \
             ~{"--colors " + colors_tsv} \
-            --output ~{basename}_auspice.json
+            --output ~{out_basename}_auspice.json
     }
     runtime {
         docker: docker
@@ -442,6 +478,6 @@ task export_auspice_json {
         preemptible: 2
     }
     output {
-        File virus_json = "~{basename}_auspice.json"
+        File virus_json = "~{out_basename}_auspice.json"
     }
 }

pipes/WDL/workflows/beast_gpu.wdl

@@ -0,0 +1,12 @@
+version 1.0
+
+import "../tasks/tasks_interhost.wdl" as interhost
+
+workflow beast_gpu {
+    call interhost.beast
+    output {
+        File        beast_log    = beast.beast_log
+        Array[File] trees        = beast.trees
+        String      beast_stdout = beast.beast_stdout
+    }
+}

pipes/WDL/workflows/beast_to_auspice.wdl

@@ -0,0 +1,36 @@
+version 1.0
+
+import "../tasks/tasks_nextstrain.wdl" as nextstrain
+
+workflow beast_to_auspice {
+    meta {
+        description: "Visualize BEAST output with Nextstrain. This workflow converts a BEAST MCC tree (.tree file) into an Auspice v2 json file. See https://nextstrain-augur.readthedocs.io/en/stable/faq/import-beast.html for details."
+    }
+
+    input {
+        File    beast_mcc_tree
+    }
+
+    parameter_meta {
+        beast_mcc_tree: {
+          description: "A maximum clade credibility (MCC) tree (.tree file) that is output from a BEAST run.",
+          patterns: ["*.tree"]
+        }
+    }
+
+    call nextstrain.augur_import_beast {
+        input:
+            beast_mcc_tree = beast_mcc_tree
+    }
+    call nextstrain.export_auspice_json {
+        input:
+            tree            = augur_import_beast.tree_newick,
+            node_data_jsons = [augur_import_beast.node_data_json]
+    }
+
+    output {
+        File  beast_mcc_tree_newick      = augur_import_beast.tree_newick
+        File  node_data_json             = augur_import_beast.node_data_json
+        File  auspice_input_json         = export_auspice_json.virus_json
+    }
+}

pipes/WDL/workflows/build_augur_tree.wdl

@@ -88,13 +88,13 @@ workflow build_augur_tree {
     }
     call nextstrain.export_auspice_json {
         input:
-            refined_tree   = refine_augur_tree.tree_refined,
-            metadata       = sample_metadata,
-            branch_lengths = refine_augur_tree.branch_lengths,
-            traits         = ancestral_traits.node_data_json,
-            nt_muts        = ancestral_tree.nt_muts_json,
-            aa_muts        = translate_augur_tree.aa_muts_json,
-            basename       = virus
+            tree            = refine_augur_tree.tree_refined,
+            metadata        = sample_metadata,
+            node_data_jsons = select_all([
+                                refine_augur_tree.branch_lengths,
+                                ancestral_traits.node_data_json,
+                                ancestral_tree.nt_muts_json,
+                                translate_augur_tree.aa_muts_json])
     }
 
     output {

requirements-modules.txt

@@ -2,5 +2,6 @@ broadinstitute/viral-core=v2.0.21
 broadinstitute/viral-assemble=v2.0.21.0
 broadinstitute/viral-classify=v2.0.21.3
 broadinstitute/viral-phylo=v2.0.21.0
+broadinstitute/beast-beagle-cuda=1.10.5
 nextstrain/base=build-20200506T095107Z
 andersenlabapps/ivar=1.2.1

test/input/WDL/test_inputs-beast_gpu-local_nonauto.json

@@ -0,0 +1,4 @@
+{
+  "beast_gpu.beast.beauti_xml": "test/input/beast-simple-ebov-20k.xml"
+}
+

test/input/WDL/test_inputs-beast_to_auspice.wdl-local.json

@@ -0,0 +1,5 @@
+{
+  "beast_to_auspice.beast_mcc_tree": "test/input/beast-mers-mcc.tree",
+  "beast_to_auspice.export_auspice_json.auspice_config": "test/input/beast-mers-auspice_config.json"
+}
+