Example of how to plot Auspice trees from Nextstrain Groups with BALTIC

2022-03-10-plot-auspice-tree-from-nextstrain-groups-with-baltic.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b4c7cde7-1d1f-4688-9edd-b669249a2aa3",
   "metadata": {},
   "source": [
    "# Plot Auspice tree from Nextstrain Groups with BALTIC"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab57aca3-9c59-417d-a279-b04873afaeab",
   "metadata": {},
   "source": [
    "## Import modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b60eaac-3f64-4286-be1e-f87c0290ef56",
   "metadata": {},
   "outputs": [],
   "source": [
    "import baltic as bt\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "import requests\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "727cb5b9-3259-4936-ae4e-d5e9d2457ae0",
   "metadata": {},
   "source": [
    "## Setup parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d4a3cc6-733e-49a3-9411-83bbda48e4e8",
   "metadata": {},
   "source": [
    "### Matplotlib styles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4512ed1-2d6a-4cb6-af54-4c6541d6d4cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "typeface = 'sans'\n",
    "mpl.rcParams['font.weight'] = 300\n",
    "mpl.rcParams['axes.labelweight'] = 300\n",
    "mpl.rcParams['font.family'] = typeface\n",
    "mpl.rcParams['font.size'] = 22"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cea269a-9560-4c88-ad15-ba24f3f49416",
   "metadata": {},
   "source": [
    "### Input paths\n",
    "\n",
    "Use [Auspice's server API](https://docs.nextstrain.org/projects/auspice/en/latest/server/api.html) to fetch the dataset we want from a Nextstrain Group."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fad76a7-b8b9-46ee-8cba-86ced33de846",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_url = \"https://nextstrain.org/charon/getDataset?prefix=groups/waphl/ncov/wa/4m\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd37a171-30bb-44e0-b697-aa36536c3c6b",
   "metadata": {},
   "source": [
    "## Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e01a8e3a-5671-4822-9bb9-287f7a0a3de0",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_json = requests.get(data_url).json()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88c97841-7973-4627-a4e1-130fb0ce2643",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Allows baltic to find correct attributes in JSON, height and name are required at a minimum.\n",
    "json_translation = {\n",
    "    'absoluteTime': lambda k: k.traits['node_attrs']['num_date']['value'],\n",
    "    'name': 'name',\n",
    "    \"clade_membership\": lambda k: k.traits['node_attrs']['clade_membership']['value'],\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e02776c-fd09-466a-843e-8a99d738c05e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Give loadJSON the name of the tree file, the translation dictionary and (optionally) the meta file.\n",
    "tree, meta = bt.loadJSON(\n",
    "    data_json,\n",
    "    json_translation=json_translation\n",
    ") "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47a51c81-2952-4197-8f3e-d434f2e8495a",
   "metadata": {},
   "source": [
    "## Plot tree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0bd9dd4-3241-4d00-94e3-9dc843e420ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "tree.cmap.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "334ac741-2575-4f1a-9cbd-3eb2ffc44797",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig,ax = plt.subplots(figsize=(20,15),facecolor='w')\n",
    "\n",
    "L=len(list(filter(lambda k: k.branchType=='leaf',tree.Objects)))\n",
    "\n",
    "# x coordinate will be absolute time of node\n",
    "x_attr=lambda k: k.absoluteTime\n",
    "\n",
    "# Colour will be determined by the trait of choice, using the colour map defined by the meta json\n",
    "c_func=lambda k: tree.cmap['region'][k.traits['node_attrs']['region']['value']]\n",
    "\n",
    "# Branch width is determined by how many descendants a node has, otherwise 2 (tips)\n",
    "b_func=lambda k: 2+10.0*len(k.leaves)/float(L) if k.branchType=='node' else 2\n",
    "\n",
    "tree.plotTree(ax,x_attr=x_attr,width=b_func,colour=c_func,zorder=20000) ## plot branches\n",
    "tree.plotPoints(ax,x_attr=x_attr,size=30,colour=c_func,zorder=20002) ## plot coloured points at tips\n",
    "\n",
    "ax.set_ylim(-50, tree.ySpan+50)\n",
    "\n",
    "[ax.spines[loc].set_visible(False) for loc in ['left','right','top']] ## no axes\n",
    "\n",
    "ax.grid(axis='x',ls='-',color='grey')\n",
    "ax.tick_params(axis='y',size=0)\n",
    "ax.tick_params(axis='x',labelsize=26)\n",
    "ax.set_yticklabels([])\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "892c5ef0-9ef8-4ea1-a7a3-5d9c95245418",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}