From a802d4b2bd3cd1f7e9e389ac5fb858845aa8c0c5 Mon Sep 17 00:00:00 2001
From: Niema Moshiri <niemamoshiri@gmail.com>
Date: Fri, 8 Mar 2024 16:29:06 -0800
Subject: [PATCH] Added support for single introduction

---
 .gitignore                                      |  4 +++-
 global.json                                     |  7 ++++++-
 plugins/viral_phylogeny_seeds/__init__.py       |  1 +
 plugins/viral_phylogeny_seeds/common_treesap.py | 11 +++++++++--
 4 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/.gitignore b/.gitignore
index ec3eae5..dd291dc 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,5 +1,7 @@
-# test file
+# test files
+config.json
 test.json
+tmp
 
 # Byte-compiled / optimized / DLL files
 __pycache__/
diff --git a/global.json b/global.json
index 0428c83..f3740f0 100644
--- a/global.json
+++ b/global.json
@@ -1,5 +1,5 @@
 {
-    "VERSION": "0.0.4",
+    "VERSION": "0.0.5",
     "CONFIG_KEYS": ["Contact Network", "Transmission Network", "Sample Times", "Viral Phylogeny (Transmissions)", "Viral Phylogeny (Seeds)", "Mutation Rates", "Ancestral Sequence", "Sequence Evolution"],
     "TYPES": ["positive integer"],
     "DESC": {
@@ -1138,6 +1138,11 @@
                 },
                 "DESC": "No phylogeny is simulated. Must select \"None\" for all downstream steps."
             },
+            "Single Introduction": {
+                "PARAM": {},
+                "REQS": {},
+                "DESC": "Assumes that there is only 1 introduction (and thus 1 transmission chain) in the transmission model, so no seed phylogeny is simulated."
+            },
             "Yule": {
                 "PARAM": {
                     "height": {
diff --git a/plugins/viral_phylogeny_seeds/__init__.py b/plugins/viral_phylogeny_seeds/__init__.py
index d32afa0..feb7e4f 100644
--- a/plugins/viral_phylogeny_seeds/__init__.py
+++ b/plugins/viral_phylogeny_seeds/__init__.py
@@ -6,5 +6,6 @@
     "Dual-Birth": common_treesap.treesap_dualbirth,
     "Non-Homogeneous Yule": common_treesap.treesap_nonhom_yule,
     "None": DUMMY_PLUGIN_FUNC,
+    "Single Introduction": common_treesap.treesap_single_intro,
     "Yule": common_treesap.treesap_yule,
 }
diff --git a/plugins/viral_phylogeny_seeds/common_treesap.py b/plugins/viral_phylogeny_seeds/common_treesap.py
index f329e22..82becf0 100644
--- a/plugins/viral_phylogeny_seeds/common_treesap.py
+++ b/plugins/viral_phylogeny_seeds/common_treesap.py
@@ -7,7 +7,7 @@
 except:
     error("Unable to import treesap. Install with: pip install treesap")
 try:
-    from treeswift import read_tree_newick
+    from treeswift import read_tree_newick, Tree
 except:
     error("Unable to import treeswift. Install with: pip install treeswift")
 
@@ -20,11 +20,16 @@ def treesap_seed(model, params, out_fn, config, GLOBAL, verbose=True):
         tree = dualbirth_tree(params['r'], 1., end_num_leaves=len(chain_trees))
     elif model == "Non-Homogeneous Yule":
         tree = nonhomogeneous_yule_tree(lambda t: eval(params['rate_func']), end_num_leaves=len(chain_trees))
+    elif model == "Single Introduction": # doesn't actually use TreeSAP
+        if len(chain_trees) != 1:
+            raise ValueError("Using 'Single Introduction', but there are %d transmission chain phylogenies" % len(chain_trees))
+        tree = Tree()
     elif model == "Yule":
         tree = yule_tree(1, end_num_leaves=len(chain_trees))
     else:
         error("Invalid TreeSAP model: %s" % model)
-    tree.scale_edges(params['height']/tree.height())
+    if model != "Single Introduction":
+        tree.scale_edges(params['height']/tree.height())
     for i, node in enumerate(tree.traverse_leaves()):
         node.label = str(i)
     tree.write_tree_newick(out_fn['viral_phylogeny_seed_time'])
@@ -40,6 +45,8 @@ def treesap_coalescent_const_pop(params, out_fn, config, GLOBAL, verbose=True):
     treesap_seed("Coalescent (Neutral)", params, out_fn, config, GLOBAL, verbose=verbose)
 def treesap_dualbirth(params, out_fn, config, GLOBAL, verbose=True):
     treesap_seed("Dual-Birth", params, out_fn, config, GLOBAL, verbose=verbose)
+def treesap_single_intro(params, out_fn, config, GLOBAL, verbose=True):
+    treesap_seed("Single Introduction", params, out_fn, config, GLOBAL, verbose=verbose)
 def treesap_nonhom_yule(params, out_fn, config, GLOBAL, verbose=True):
     treesap_seed("Non-Homogeneous Yule", params, out_fn, config, GLOBAL, verbose=verbose)
 def treesap_yule(params, out_fn, config, GLOBAL, verbose=True):