Gradient_clade

September 10, 2021

In [17]:

import baltic as bt
import requests
from io import StringIO as sio

import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from matplotlib.patches import Polygon, PathPatch
from matplotlib.path import Path
from matplotlib.collections import LineCollection
from matplotlib import cm
import matplotlib.patheffects as path_effects

import numpy as np

from datetime import datetime as dt

In [7]:

# From Washginton, Gangavarapu et al. Cell 2021 - https://doi.org/10.1016/j.cell.2021.03.052

nexus_file_path = "https://raw.githubusercontent.com/andersen-lab/paper_2021_early-b117-usa/master/2021-02-13_b117_combined.nexus"
nexus_file = requests.get(nexus_file_path) # Fetch tree

treeFile=sio(nexus_file.text) ## stream from repo copy

ll=bt.loadNexus(treeFile) ## treeFile here can alternatively be a path to a local file
ll.traverse_tree()
ll.treeStats()
ll.sortBranches()

Tree height: 0.978386
Tree length: 40.419198
strictly bifurcating tree
annotations present

Numbers of objects in tree: 1907 (953 nodes and 954 leaves)

In [66]:

loc_color = {
    "California": "#228b22",
    "Florida": "#ff6a6a",
    "USA": "#000000"
}

background_node_color = "#cccccc"

# Visualization inspired by JT McCrone's global SARS-CoV-2 visualization at https://artic.network/hCoV-19-viz/

def get_varying_alpha_line(x1, x2, y, _color):
    """
    Draw line with gradually increasing alpha from x1 to x2 at a given y and _color
    """
    nlines = 100
    coords = np.dstack([np.linspace(x1, x2, nlines), [y] * nlines])[0]
    coords = np.stack([coords[1:], coords[:-1]], axis = 1).reshape(nlines - 1, 2, 2)
    z = np.empty((nlines-1, 4))
    rgb = mcolors.colorConverter.to_rgb(_color)
    z[:,  :3] = rgb
    z[:,-1] = np.append((np.logspace(0, 1, nlines - 2) -1)/10, 1)
    return coords, z

def four_point_bezier(*points):
    """
    Draw a cubic bezier curve with four control points
    https://en.wikipedia.org/wiki/B%C3%A9zier_curve
    """
    p1, p2, p3, p4 = [np.array(i) for i in points]
    t = np.arange(0,1,0.01)
    p = [None, None]
    p[0] = (1- t) ** 3 * p1[0] + 3 * (1-t)**2 *t * p2[0] +3*(1-t) * t ** 2 * p3[0] + t ** 3 * p4[0]
    p[1] = (1- t) ** 3 * p1[1] + 3 * (1-t)**2 *t * p2[1] +3*(1-t) * t ** 2 * p3[1] + t ** 3 * p4[1]
    return p

def get_all_terminals(node):
    """
    Get all terminal nodes that descend from a given node
    """
    terminals = []
    for i in node.children:
        if i.branchType == "leaf":
            terminals.append(i)
        if i.branchType != "leaf":
            terminals.extend(get_all_terminals(i))
    return terminals

def draw_gradient_clade(node, ax, _color, start_x = None, start_y = None, end_x = None, zorder = 200, control_point_fraction = 0.1, offset_y = [0,0]):
    """
    Draw gradient given a node object, node, a matplotlib axis, ax
    start_x: start of the gradient. By default is absoluteTime of node object
    start_y: y axis to start gradient from. By default, it is y attribute of node object
    end_x: x corrdinate to end gradient. By default, this is the maximum absoluteTime of all terminal nodes that descend from node object
    """
    if start_x == None:
        start_x = node.absoluteTime
    if start_y == None:
        start_y = node.y
    terminals = get_all_terminals(node)
    y2 = [min([i.y for i in terminals]) - offset_y[0], max([i.y for i in terminals]) + offset_y[1]]
    x2 = max([i.absoluteTime for i in terminals])
    x1 = start_x
    if end_x == None:
        end_x = x2
    lower_control_points = [[x1,start_y], [x1 + (end_x-x1)* control_point_fraction, start_y], [x1 + (end_x-x1)* control_point_fraction, y2[0]], [end_x,y2[0]]]
    upper_control_points = [[x1,start_y], [x1 + (end_x-x1)* control_point_fraction, start_y], [x1 + (end_x-x1)* control_point_fraction, y2[1]], [end_x,y2[1]]]
    lower_p = four_point_bezier(*lower_control_points)
    ax.plot(lower_p[0], lower_p[1], color=_color, zorder = 200)
    # Lower gradient
    z = np.empty((100, 1, 4))
    rgb = mcolors.colorConverter.to_rgb(_color)
    z[:, :, :3] = rgb
    z[:,:,-1] = np.linspace(0, 1, 100)[:,None]
    im = ax.imshow(np.rot90(z, axes=(1,0)), aspect='auto', extent=[x1, end_x, y2[0], start_y], origin='upper', zorder=zorder)
    xy = np.dstack([lower_p[0], lower_p[1]])[0]
    xy = np.vstack([[x1, start_y], xy, [end_x, y2[0]], [end_x, start_y]])
    clip_path = Polygon(xy, facecolor='none', edgecolor='none', closed=True)
    ax.add_patch(clip_path)
    im.set_clip_path(clip_path)
    codes = [Path.MOVETO]
    codes.extend([Path.LINETO] * (len(xy) - 2))
    bezier_border = PathPatch(Path(xy[:-1], codes), facecolor='none', lw=0.5, zorder = 201, edgecolor="#474747")
    ax.add_patch(bezier_border)
    # Upper gradient
    upper_p = four_point_bezier(*upper_control_points)
    ax.plot(upper_p[0], upper_p[1], color=_color, zorder = 200)
    z = np.empty((100, 1, 4))
    rgb = mcolors.colorConverter.to_rgb(_color)
    z[:, :, :3] = rgb
    z[:,:,-1] = np.linspace(0, 1, 100)[:,None]
    im = ax.imshow(np.rot90(z, axes=(1,0)), aspect='auto', extent=[x1, end_x, start_y, y2[1]], origin='upper', zorder=zorder)
    xy = np.dstack([upper_p[0], upper_p[1]])[0]
    xy = np.vstack([[x1, start_y], xy, [end_x, y2[1]], [end_x, start_y]])
    clip_path = Polygon(xy, facecolor='none', edgecolor='none', closed=True)
    ax.add_patch(clip_path)
    im.set_clip_path(clip_path)
    codes = [Path.MOVETO]
    codes.extend([Path.LINETO] * (len(xy) - 2))
    bezier_border = PathPatch(Path(xy[:-1], codes), facecolor='none', lw=0.5, zorder = 201, edgecolor="#474747")
    ax.add_patch(bezier_border)

x_attr=lambda k: k.absoluteTime


f,ax = plt.subplots(figsize=(15, 15))

# Get LCA of B.1.1.7 lineage
b117_node = next(i for i in ll.Objects if i.branchType == "leaf" and "EPI_ISL_860089" in i.name).parent.children[1]

# Draw gradient starting from LCA of B.1.1.7 lineage
draw_gradient_clade(b117_node, ax, "#cccccc", offset_y = [5, 0])

# Draw gradient starting from LCA of given set of taxa
# California
ca_main_taxa = ["STM-0000016-H11|2021-01-15|USA|California|None", "STM-0000024-E05|2021-01-14|USA|California|None", "STM-0000024-E11|2021-01-14|USA|California|None", "STM-0000025-F02|2021-01-11|USA|California|None", "STM-0000029-A05|2021-01-18|USA|California|None", "STM-0000029-A11|2021-01-18|USA|California|None", "STM-0000029-D10|2021-01-18|USA|California|None", "STM-0000029-F11|2021-01-18|USA|California|None", "STM-0000031-A09|2021-01-16|USA|California|None", "STM-0000031-B09|2021-01-16|USA|California|None", "STM-0000036-E11|2021-01-19|USA|California|None", "STM-0000041-B02|2021-01-20|USA|California|None", "STM-0000041-B06|2021-01-20|USA|California|None", "STM-0000041-G04|2021-01-20|USA|California|None", "STM-0000081-F04|2021-01-11|USA|California|None", "STM-0000082-D07|2021-01-10|USA|California|None", "STM-0000082-D09|2021-01-10|USA|California|None", "STM-021|2020-12-28|USA|California|None", "STM-064|2020-12-26|USA|California|None", "STM-070|2020-12-26|USA|California|None", "STM-105|2020-12-29|USA|California|None", "STM-127|2020-12-30|USA|California|None", "STM-156|2020-12-31|USA|California|None", "STM-166|2020-12-31|USA|California|None", "STM-A100051|2021-01-04|USA|California|None", "STM-A100119|2021-01-05|USA|California|None", "STM-A100172|2021-01-05|USA|California|None", "STM-A100175|2021-01-05|USA|California|None", "STM-A100197|2021-01-05|USA|California|None", "STM-A100288|2021-01-06|USA|California|None", "STM-A100349|2021-01-06|USA|California|None", "STM-A100478|2021-01-07|USA|California|None", "STM-P019|2020-12-21|USA|California|None", "STM-P025|2020-12-20|USA|California|None", "STM-0000014-F03|2021-01-13|USA|California|None", "STM-0000024-C11|2021-01-14|USA|California|None", "STM-0000024-A12|2021-01-14|USA|California|None", "STM-0000025-B07|2021-01-12|USA|California|None", "STM-0000024-D08|2021-01-14|USA|California|None", "STM-0000024-G10|2021-01-14|USA|California|None", "STM-0000025-C02|2021-01-11|USA|California|None", "STM-0000025-G04|2021-01-11|USA|California|None", "STM-0000031-G10|2021-01-16|USA|California|None", "STM-0000081-G09|2021-01-08|USA|California|None", "STM-A100136|2021-01-05|USA|California|None", "STM-067|2020-12-26|USA|California|None", "STM-076|2020-12-29|USA|California|None", "STM-087|2020-12-29|USA|California|None", "STM-154|2020-12-31|USA|California|None", "STM-165|2020-12-31|USA|California|None", "STM-P017|2020-12-20|USA|California|None", "STM-208|2021-01-02|USA|California|None", "STM-0000025-C04|2021-01-11|USA|California|None", "STM-A100153|2021-01-05|USA|California|None", "STM-069|2020-12-26|USA|California|None", "STM-071|2020-12-26|USA|California|None", "STM-073|2020-12-26|USA|California|None", "STM-A100193|2021-01-05|USA|California|None", "STM-A100148|2021-01-05|USA|California|None", "STM-A100151|2021-01-05|USA|California|None", "STM-A100152|2021-01-05|USA|California|None", "STM-A100173|2021-01-05|USA|California|None", "STM-0000035-H05|2021-01-19|USA|California|None", "STM-022|2020-12-29|USA|California|None", "STM-062|2020-12-27|USA|California|None", "STM-A100413|2021-01-07|USA|California|None", "STM-A100427|2021-01-07|USA|California|None", "STM-0000014-B12|2021-01-13|USA|California|None", "STM-0000024-F07|2021-01-14|USA|California|None", "STM-0000043-G12|2021-01-20|USA|Texas|None", "STM-0000080-C09|2021-01-08|USA|California|None", "STM-0000082-E09|2021-01-10|USA|California|None", "STM-0000016-G03|2021-01-12|USA|California|None", "STM-0000031-G09|2021-01-16|USA|California|None", "STM-037|2020-12-28|USA|California|None", "STM-222|2021-01-02|USA|California|None", "STM-0000025-H09|2021-01-12|USA|California|None", "STM-0000050-B06|2021-01-22|USA|Texas|None", "STM-232|2021-01-02|USA|California|None", "STM-0000031-H08|2021-01-16|USA|California|None", "STM-0000082-E07|2021-01-10|USA|California|None", "STM-130|2020-12-30|USA|California|None", "STM-0000080-E01|2021-01-08|USA|California|None", "STM-0000041-B12|2021-01-20|USA|California|None", "STM-131|2020-12-30|USA|California|None", "STM-A100155|2021-01-05|USA|California|None", "STM-080|2020-12-29|USA|California|None", "STM-126|2020-12-30|USA|California|None", "STM-0000024-B12|2021-01-14|USA|California|None", "STM-A100422|2021-01-07|USA|California|None", "STM-0000031-E11|2021-01-16|USA|California|None", "STM-0000080-C05|2021-01-08|USA|California|None", "STM-0000080-C06|2021-01-08|USA|Georgia|None", "STM-244|2021-01-02|USA|California|None", "STM-310|2021-01-03|USA|California|None", "STM-342|2021-01-03|USA|California|None", "EPI_ISL_751801|2020-12-29|USA|California|San_Diego_County", "STM-0000080-F11|2021-01-08|USA|Georgia|None", "STM-A100147|2021-01-05|USA|California|None", "STM-128|2020-12-30|USA|California|None", "STM-0000082-H09|2021-01-10|USA|California|None", "STM-150|2020-12-31|USA|California|None", "STM-A100432|2021-01-07|USA|California|None", "STM-A100403|2021-01-07|USA|California|None", "STM-149|2020-12-31|USA|California|None", "STM-A100154|2021-01-05|USA|California|None", "STM-155|2020-12-31|USA|California|None", "STM-077|2020-12-29|USA|California|None", "STM-0000082-C07|2021-01-10|USA|California|None", "STM-0000082-F07|2021-01-10|USA|California|None", "STM-A100423|2021-01-07|USA|California|None", "STM-050|2020-12-28|USA|California|None", "STM-153|2020-12-31|USA|Minnesota|None", "STM-0000016-H08|2021-01-12|USA|California|None", "STM-061|2020-12-27|USA|California|None", "STM-0000030-G12|2021-01-17|USA|California|None", "STM-0000031-A10|2021-01-16|USA|California|None", "STM-078|2020-12-29|USA|California|None", "STM-296|2021-01-03|USA|California|None", "STM-303|2021-01-03|USA|California|None", "STM-A100116|2021-01-05|USA|California|None", "STM-A100150|2021-01-05|USA|California|None", "STM-A100162|2021-01-05|USA|California|None", "STM-A100377|2021-01-06|USA|California|None", "STM-A100480|2021-01-07|USA|California|None", "STM-082|2020-12-29|USA|California|None", "STM-0000036-B05|2021-01-19|USA|California|None", "STM-0000041-D04|2021-01-20|USA|California|None", "STM-0000041-G10|2021-01-20|USA|California|None", "STM-0000082-A02|2021-01-11|USA|California|None", "ALSR-6046|2021-01-07|USA|California|San_Diego", "ALSR-6096|2021-01-07|USA|California|San_Diego", "STM-0000016-E08|2021-01-12|USA|California|None", "STM-0000016-F03|2021-01-12|USA|California|None", "STM-0000082-G01|2021-01-11|USA|California|None", "STM-0000014-E11|2021-01-13|USA|California|None", "STM-0000082-G11|2021-01-10|USA|California|None", "STM-A100386|2021-01-07|USA|California|None", "STM-0000014-C07|2021-01-13|USA|California|None", "STM-0000080-A10|2021-01-08|USA|California|None", "STM-0000013-E01|2021-01-15|USA|California|None", "STM-0000036-A03|2021-01-19|USA|California|None", "STM-0000036-G03|2021-01-19|USA|California|None"]
ca_children = [i for i in ll.Objects if i.branchType == "leaf" and i.name in ca_main_taxa]
ca_node = ll.commonAncestor(ca_children)
draw_gradient_clade(ca_node, ax, "#EABF7D", control_point_fraction = 0.05)

# Florida
fl_main_taxa = ["STM-0000013-F04|2021-01-14|USA|Florida|None", "STM-0000013-F12|2021-01-12|USA|Florida|None", "STM-0000014-A06|2021-01-12|USA|Florida|None", "STM-0000014-B11|2021-01-12|USA|Florida|None", "STM-0000025-B04|2021-01-10|USA|Florida|None", "STM-0000025-C11|2021-01-14|USA|Florida|None", "STM-0000029-C11|2021-01-17|USA|Florida|None", "STM-0000029-G12|2021-01-17|USA|Florida|None", "STM-0000030-E07|2021-01-16|USA|Florida|None", "STM-0000035-E11|2021-01-17|USA|Florida|None", "STM-0000035-F08|2021-01-17|USA|Florida|None", "STM-0000035-H07|2021-01-18|USA|Florida|None", "STM-0000035-H11|2021-01-17|USA|Florida|None", "STM-0000036-B06|2021-01-18|USA|Florida|None", "STM-0000036-F01|2021-01-18|USA|Florida|None", "STM-0000036-G10|2021-01-18|USA|Florida|None", "STM-0000043-B04|2021-01-20|USA|Florida|None", "STM-0000053-C08|2021-01-23|USA|Florida|None", "STM-0000053-D03|2021-01-23|USA|Florida|None", "STM-0000053-D06|2021-01-23|USA|Florida|None", "STM-0000054-A02|2021-01-26|USA|North_Carolina|None", "STM-0000054-H02|2021-01-26|USA|Florida|None", "STM-0000056-H10|2021-01-27|USA|Florida|None", "STM-0000058-A08|2021-01-27|USA|Florida|None", "STM-0000058-D05|2021-01-28|USA|Florida|None", "STM-0000059-C11|2021-01-28|USA|Florida|None", "STM-0000059-D11|2021-01-28|USA|Florida|None", "STM-0000059-E09|2021-01-28|USA|Florida|None", "STM-0000059-G10|2021-01-28|USA|Florida|None", "STM-0000072-C10|2021-01-29|USA|Florida|None", "STM-0000072-E07|2021-01-29|USA|Florida|None", "STM-0000080-D03|2021-01-08|USA|Florida|None", "STM-0000081-G05|2021-01-10|USA|Florida|None", "STM-024|2020-12-28|USA|Florida|None", "STM-085|2020-12-28|USA|Florida|None", "STM-100|2020-12-28|USA|Florida|None", "STM-115|2020-12-29|USA|Florida|None", "STM-161|2020-12-30|USA|Florida|None", "STM-A015|2021-01-02|USA|Florida|None", "STM-A100002|2021-01-04|USA|Florida|None", "STM-A100266|2021-01-02|USA|Florida|None", "STM-A100330|2021-01-05|USA|Florida|None", "STM-0000036-G09|2021-01-18|USA|Florida|None", "STM-0000082-H04|2021-01-09|USA|Florida|None", "STM-0000053-F12|2021-01-26|USA|Florida|None", "STM-0000058-A05|2021-01-28|USA|Florida|None", "STM-0000059-A08|2021-01-28|USA|Florida|None", "STM-0000059-G07|2021-01-28|USA|Florida|None", "STM-A100199|2021-01-03|USA|Florida|None", "STM-A100227|2021-01-05|USA|Florida|None", "STM-0000013-G09|2021-01-11|USA|Florida|None", "STM-0000041-H05|2021-01-19|USA|Florida|None", "STM-0000058-A11|2021-01-28|USA|Florida|None", "STM-0000016-B04|2021-01-11|USA|Florida|None", "STM-0000025-F03|2021-01-10|USA|Florida|None", "STM-0000058-D03|2021-01-27|USA|Florida|None", "STM-0000041-C10|2021-01-19|USA|Florida|None", "STM-0000042-B07|2021-01-21|USA|Florida|None", "STM-0000043-F06|2021-01-20|USA|Florida|None", "STM-0000054-B02|2021-01-26|USA|Florida|None", "STM-0000053-C06|2021-01-23|USA|Florida|None", "STM-0000053-H05|2021-01-23|USA|Florida|None", "STM-0000082-G08|2021-01-09|USA|Florida|None", "STM-088|2020-12-28|USA|Florida|None", "STM-A100428|2021-01-06|USA|Florida|None", "STM-0000013-G11|2021-01-11|USA|Florida|None", "STM-0000041-A06|2021-01-20|USA|Florida|None", "STM-0000058-F07|2021-01-27|USA|Florida|None", "STM-0000080-C12|2021-01-07|USA|Florida|None", "STM-0000053-E01|2021-01-23|USA|Florida|None", "STM-0000058-C03|2021-01-27|USA|Florida|None", "STM-0000059-A09|2021-01-28|USA|Florida|None", "STM-0000059-G12|2021-01-28|USA|Florida|None", "STM-0000059-H10|2021-01-28|USA|Florida|None", "STM-0000013-C11|2021-01-11|USA|Florida|None", "STM-0000013-F05|2021-01-14|USA|Florida|None", "STM-0000054-G10|2021-01-27|USA|Florida|None", "STM-0000058-B02|2021-01-27|USA|Florida|None", "STM-0000082-H05|2021-01-09|USA|Florida|None", "STM-194|2020-12-31|USA|Florida|None", "STM-0000014-B05|2021-01-12|USA|Florida|None", "STM-0000041-A08|2021-01-19|USA|Florida|None", "STM-0000041-C08|2021-01-19|USA|Florida|None", "STM-0000035-A08|2021-01-17|USA|Florida|None", "STM-0000050-D01|2021-01-22|USA|Florida|None", "STM-0000056-G12|2021-01-27|USA|Florida|None", "STM-0000056-D10|2021-01-27|USA|Florida|None", "STM-0000056-H11|2021-01-27|USA|Florida|None", "STM-0000013-E10|2021-01-11|USA|Michigan|None", "STM-0000029-C02|2021-01-16|USA|Florida|None", "STM-0000030-A02|2021-01-15|USA|Pennsylvania|None", "STM-0000035-E02|2021-01-18|USA|Florida|None", "STM-0000050-C06|2021-01-22|USA|Florida|None", "STM-0000054-E12|2021-01-27|USA|Pennsylvania|None", "STM-0000054-G06|2021-01-26|USA|Florida|None", "STM-0000058-C05|2021-01-28|USA|Florida|None", "STM-0000059-G09|2021-01-28|USA|Florida|None", "STM-0000030-D10|2021-01-16|USA|Florida|None", "STM-0000041-A03|2021-01-19|USA|Florida|None", "STM-0000053-A07|2021-01-23|USA|Florida|None", "STM-055|2020-12-26|USA|Florida|None", "STM-184|2020-12-31|USA|Florida|None", "STM-0000053-F06|2021-01-23|USA|Florida|None", "STM-145|2020-12-30|USA|Florida|None", "STM-0000053-E08|2021-01-23|USA|Florida|None", "STM-0000081-F07|2021-01-07|USA|Florida|None", "STM-0000058-B05|2021-01-28|USA|Florida|None", "STM-056|2020-12-26|USA|Florida|None", "STM-0000058-A02|2021-01-27|USA|Florida|None", "STM-0000013-B12|2021-01-11|USA|Florida|None", "STM-0000053-H06|2021-01-23|USA|Florida|None", "STM-0000013-D07|2021-01-14|USA|Florida|None", "STM-0000013-E07|2021-01-14|USA|Florida|None", "STM-0000025-F01|2021-01-10|USA|Florida|None", "STM-0000035-D01|2021-01-18|USA|Florida|None", "STM-0000036-B10|2021-01-18|USA|Florida|None", "STM-0000041-G09|2021-01-19|USA|Florida|None", "STM-0000082-E04|2021-01-09|USA|Florida|None", "STM-0000054-E03|2021-01-26|USA|Florida|None", "STM-0000058-C06|2021-01-27|USA|Florida|None", "STM-0000030-F01|2021-01-15|USA|Florida|None", "STM-0000036-B03|2021-01-18|USA|Florida|None", "STM-0000059-F07|2021-01-28|USA|Florida|None", "STM-0000013-H11|2021-01-11|USA|Florida|None", "STM-0000082-E10|2021-01-09|USA|Georgia|None", "STM-032|2020-12-27|USA|Florida|None", "STM-044|2020-12-26|USA|Florida|None", "STM-049|2020-12-26|USA|Florida|None", "STM-189|2020-12-31|USA|Florida|None", "STM-A100353|2021-01-05|USA|Florida|None", "STM-0000035-H04|2021-01-17|USA|Florida|None", "STM-0000041-B05|2021-01-19|USA|Florida|None", "STM-0000054-F04|2021-01-26|USA|Florida|None", "STM-0000059-G06|2021-01-28|USA|Florida|None", "STM-0000080-H09|2021-01-07|USA|Florida|None", "STM-292|2021-01-02|USA|Florida|None", "STM-033|2020-12-27|USA|Florida|None", "STM-0000053-A08|2021-01-23|USA|Florida|None", "STM-0000050-A09|2021-01-23|USA|North_Carolina|None", "STM-0000080-B08|2021-01-08|USA|Florida|None", "STM-0000042-H03|2021-01-21|USA|Indiana|None", "STM-0000035-H09|2021-01-17|USA|Florida|None", "STM-083|2020-12-28|USA|Florida|None", "STM-0000053-F09|2021-01-23|USA|Florida|None", "STM-0000058-E06|2021-01-27|USA|Florida|None", "STM-0000053-F07|2021-01-23|USA|Florida|None", "STM-147|2020-12-30|USA|Florida|None", "STM-042|2020-12-26|USA|Florida|None", "STM-0000080-H05|2021-01-07|USA|Florida|None", "STM-0000035-C01|2021-01-18|USA|Indiana|None", "STM-0000056-F02|2021-01-26|USA|Florida|None", "STM-0000081-B05|2021-01-10|USA|Florida|None", "STM-0000058-E05|2021-01-27|USA|Florida|None", "STM-0000058-F06|2021-01-27|USA|Florida|None", "STM-0000058-A07|2021-01-27|USA|Florida|None", "STM-0000059-H02|2021-01-28|USA|Texas|None", "STM-P012|2020-12-19|USA|Florida|None", "STM-0000013-E08|2021-01-14|USA|Florida|None", "STM-0000024-H01|2021-01-12|USA|Florida|None", "STM-0000013-F08|2021-01-14|USA|Florida|None", "STM-0000014-H02|2021-01-12|USA|Florida|None", "STM-0000058-D06|2021-01-27|USA|Florida|None", "STM-0000059-G08|2021-01-28|USA|Florida|None", "STM-0000082-A04|2021-01-09|USA|Florida|None", "STM-151|2020-12-30|USA|Florida|None", "STM-158|2020-12-30|USA|Florida|None", "STM-0000030-E10|2021-01-16|USA|Florida|None", "STM-124|2020-12-29|USA|Florida|None", "STM-0000054-E02|2021-01-26|USA|Florida|None", "STM-157|2020-12-30|USA|Florida|None", "STM-0000016-B12|2021-01-14|USA|Georgia|None", "STM-0000024-E01|2021-01-12|USA|Georgia|None", "STM-0000030-D01|2021-01-15|USA|Florida|None", "STM-0000036-B08|2021-01-18|USA|Florida|None", "STM-0000036-C10|2021-01-19|USA|Florida|None", "STM-0000041-E06|2021-01-19|USA|Georgia|None", "STM-0000053-C02|2021-01-23|USA|Georgia|None", "STM-0000054-D02|2021-01-26|USA|Florida|None", "STM-0000056-G09|2021-01-27|USA|Georgia|None", "STM-0000059-C08|2021-01-28|USA|Florida|None", "STM-0000059-F08|2021-01-28|USA|Florida|None", "STM-171|2020-12-31|USA|Florida|None", "STM-191|2020-12-31|USA|Florida|None", "STM-315|2021-01-02|USA|Florida|None", "STM-0000053-D04|2021-01-23|USA|Florida|None", "STM-0000080-G02|2021-01-08|USA|Georgia|None", "STM-135|2020-12-29|USA|Florida|None", "STM-0000013-A11|2021-01-11|USA|Florida|None", "STM-0000080-E08|2021-01-08|USA|Georgia|None", "STM-0000013-F02|2021-01-14|USA|Georgia|None", "STM-0000013-H01|2021-01-14|USA|Georgia|None", "STM-0000030-B04|2021-01-15|USA|Georgia|None", "STM-199|2020-12-31|USA|Florida|None", "STM-0000042-A10|2021-01-21|USA|Florida|None", "STM-A100364|2021-01-05|USA|Georgia|None", "STM-0000014-H10|2021-01-12|USA|Florida|None", "STM-0000030-D09|2021-01-16|USA|Florida|None", "STM-0000031-A12|2021-01-15|USA|Florida|None", "STM-0000035-A02|2021-01-18|USA|Florida|None", "STM-0000035-C02|2021-01-17|USA|Florida|None", "STM-0000035-D12|2021-01-17|USA|Florida|None", "STM-0000035-G10|2021-01-17|USA|Florida|None", "STM-0000041-B08|2021-01-19|USA|Florida|None", "STM-0000043-C06|2021-01-20|USA|Florida|None", "STM-0000054-B03|2021-01-26|USA|Florida|None", "STM-0000054-H10|2021-01-27|USA|Florida|None", "STM-0000056-F03|2021-01-26|USA|Florida|None", "STM-0000058-B08|2021-01-27|USA|Florida|None", "STM-0000058-G01|2021-01-27|USA|Florida|None", "STM-0000059-C07|2021-01-28|USA|Florida|None", "STM-0000059-E07|2021-01-28|USA|Florida|None", "STM-0000059-G05|2021-01-28|USA|Florida|None", "STM-0000081-F12|2021-01-08|USA|Florida|None", "STM-A034|2021-01-02|USA|Florida|None", "STM-0000013-G06|2021-01-14|USA|Florida|None", "STM-0000030-A10|2021-01-16|USA|Florida|None", "STM-0000013-H10|2021-01-11|USA|Florida|None", "STM-0000056-G02|2021-01-26|USA|Florida|None", "STM-0000024-B06|2021-01-13|USA|Florida|None", "STM-0000080-B10|2021-01-07|USA|Florida|None", "STM-0000080-B03|2021-01-08|USA|Florida|None", "STM-A100316|2021-01-05|USA|Florida|None", "STM-181|2020-12-31|USA|Florida|None", "STM-0000013-H07|2021-01-14|USA|Florida|None", "STM-0000080-A07|2021-01-07|USA|Florida|None", "STM-0000054-B01|2021-01-26|USA|Florida|None", "STM-0000054-G03|2021-01-26|USA|Florida|None", "STM-0000056-F01|2021-01-26|USA|Florida|None", "STM-0000056-A07|2021-01-27|USA|Florida|None", "STM-0000072-D06|2021-01-29|USA|Florida|None", "STM-0000013-A07|2021-01-14|USA|North_Carolina|None", "STM-0000029-C04|2021-01-16|USA|Florida|None", "STM-0000029-F07|2021-01-17|USA|North_Carolina|None", "STM-0000036-E10|2021-01-19|USA|Florida|None", "STM-0000041-H04|2021-01-20|USA|Florida|None", "STM-0000053-A02|2021-01-23|USA|Illinois|None", "STM-0000058-E07|2021-01-27|USA|Florida|None", "STM-0000025-E01|2021-01-10|USA|Florida|None", "STM-0000041-G11|2021-01-19|USA|Florida|None", "STM-0000059-G04|2021-01-28|USA|Florida|None", "STM-0000056-F12|2021-01-27|USA|Florida|None", "STM-0000058-F04|2021-01-27|USA|Florida|None", "STM-0000029-B12|2021-01-17|USA|Florida|None", "STM-0000053-G12|2021-01-25|USA|Florida|None", "STM-0000054-B07|2021-01-26|USA|Florida|None", "STM-0000059-B08|2021-01-28|USA|Florida|None", "STM-0000054-C01|2021-01-26|USA|Florida|None", "STM-0000042-B05|2021-01-21|USA|Georgia|None", "STM-0000054-D05|2021-01-26|USA|Florida|None", "STM-0000016-D12|2021-01-14|USA|Florida|None", "STM-0000024-C06|2021-01-13|USA|Florida|None", "STM-0000035-E05|2021-01-17|USA|Florida|None", "STM-0000043-A07|2021-01-20|USA|Florida|None", "STM-0000053-B07|2021-01-23|USA|Florida|None", "STM-0000053-E06|2021-01-23|USA|Florida|None", "STM-0000059-H04|2021-01-28|USA|Florida|None", "STM-0000043-F09|2021-01-20|USA|Florida|None", "STM-0000053-A05|2021-01-23|USA|Florida|None", "STM-0000059-H09|2021-01-28|USA|Florida|None", "STM-0000059-D10|2021-01-28|USA|Florida|None", "STM-0000029-E03|2021-01-16|USA|Florida|None", "STM-0000035-H08|2021-01-17|USA|Florida|None", "STM-0000013-G12|2021-01-15|USA|Florida|None", "STM-0000058-C07|2021-01-27|USA|Florida|None", "STM-0000050-G09|2021-01-23|USA|Florida|None", "STM-0000053-B05|2021-01-23|USA|Florida|None", "STM-0000053-G01|2021-01-23|USA|Florida|None", "STM-0000054-C03|2021-01-26|USA|Florida|None", "STM-0000054-F01|2021-01-26|USA|Florida|None", "STM-0000056-G07|2021-01-27|USA|Florida|None", "STM-0000059-D12|2021-01-28|USA|Florida|None", "STM-0000072-D08|2021-01-29|USA|Florida|None", "STM-0000056-G03|2021-01-27|USA|Florida|None", "STM-0000042-D03|2021-01-21|USA|Florida|None", "STM-0000072-E06|2021-01-29|USA|Florida|None", "STM-0000058-H07|2021-01-27|USA|Florida|None", "STM-0000059-C10|2021-01-28|USA|Florida|None", "STM-0000072-F08|2021-01-29|USA|Florida|None", "STM-0000054-F12|2021-01-27|USA|Florida|None", "STM-0000059-F11|2021-01-28|USA|Florida|None"]
fl_children = [i for i in ll.Objects if i.branchType == "leaf" and i.name in fl_main_taxa]
fl_node = ll.commonAncestor(fl_children)
draw_gradient_clade(fl_node, ax, "#A5C982", control_point_fraction = 0.05)


s_func=lambda k: 50-30*k.height/ll.treeHeight # Size of tips
z_func=lambda k: 100 # Z order
cu_func=lambda k: 'k' # For plotting a black outline of tip circles
su_func=lambda k: 2*(50-30*k.height/ll.treeHeight) # Black outline is typically twice as big as tip circle


# Iterate over objects in tree to only plot branches leading upto the clade visualzied using the gradient
gradient_nodes = [b117_node, ca_node, fl_node]

nodes_to_plot = []
leaves_to_plot = []
for i in ll.Objects:
    path_to_root = []
    j = i
    while j != ll.root:
        path_to_root.append(j)
        j = j.parent
    if all([node not in path_to_root for node in gradient_nodes]):
        nodes_to_plot.append(i)
    if any([node in path_to_root for node in gradient_nodes]) and i.branchType == "leaf":
        leaves_to_plot.append(i)

target_func = lambda x: x in leaves_to_plot
tree_target_func = lambda x: x in nodes_to_plot
ll.plotTree(ax,x_attr=x_attr,colour="#000000",width=0.5,target=tree_target_func) # Only plot branches leading upto CA and FL clades
# For testing draw all branches
#ll.plotTree(ax,x_attr=x_attr,colour="#000000",width=0.5) 


state_color_scheme = {
    "California": "#EABF7D",
    "Florida": "#A5C982",
}

def c_func(n):
    k = n.name
    if k == "NC_045512.2" or "USA" not in k.split("|")[2]:
        return "#ECECEC"
    if k.split("|")[3].replace("_", " ") in state_color_scheme:
        state_name = k.split("|")[3].replace("_", " ")
        return state_color_scheme[state_name]
    return loc_color["USA"]

node_target_func = lambda x: x in leaves_to_plot
ll.plotPoints(ax,x_attr=x_attr,size=s_func,colour=c_func,zorder=202, target = node_target_func) ## Only plot leaves if branch within 

# Draw lines with increasing alpha emerging from LCA of each clae visualized using a gradient
branch_lines = []
line_colors = []
for node in gradient_nodes:
    for i in get_all_terminals(node):
        coords, colors = get_varying_alpha_line((node.absoluteTime + i.absoluteTime)/2, i.absoluteTime, i.y, "#000000")
        branch_lines.extend(coords)
        line_colors.extend(colors)

ax.add_collection(LineCollection(branch_lines, color = line_colors, lw = 1))

# Draw branches from parent to gradient nodes
for i in gradient_nodes:
    ax.plot([i.parent.absoluteTime, i.parent.absoluteTime], [i.parent.y, i.y], color = "#000000", lw = 0.5)
    ax.plot([i.parent.absoluteTime, i.absoluteTime], [i.y, i.y], color = "#000000", lw = 0.5)

for i in ["top", "left", "right"]:
    ax.spines[i].set_visible(False)

# Label nodes with posterios >= 0.9
#for i in ll.Objects:
#    if "posterior" not in i.traits:
#        continue
#    if i.traits["posterior"] >= 0.9:
#        ax.text(i.absoluteTime, i.y, round(i.traits["posterior"], 2))

dates = ["2020-02-01"]
dates.extend(["2020-{}-01".format(i) for i in range(2,13)])
dates.extend(["2021-01-01", "2020-02-01"])
xticks = [bt.decimalDate(i) for i in dates]
xticklabels = [dt.strptime(i, "%Y-%m-%d").strftime("%b %Y") for i in dates]
ax.set_ylim(-10,ll.ySpan+5)
ax.set_xlim(bt.decimalDate("2020-02-01"), bt.decimalDate("2021-02-01"))
ax.set_yticks([])

for i in xticks:
    ax.axvline(i, ls="--", color="#ECECEC", zorder = -1)

ax.set_xticks(xticks)
ax.set_xticklabels(xticklabels)
plt.tight_layout()
plt.show()
plt.close()

In [ ]: