SARS2-spike-predictor-phenos.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "93f86296-548c-4d78-b295-d041b371c92e",
   "metadata": {},
   "source": [
    "# Tabulate phenotypes of SARS-CoV-2 spike that might have predictive power in understanding evolution for different mutations and clades\n",
    "\n",
    "Import Python modules:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4356fd9-5204-4b7f-a587-f2bc90bb719f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import datetime\n",
    "import gzip\n",
    "import json\n",
    "import os\n",
    "import urllib.request\n",
    "\n",
    "import altair as alt\n",
    "\n",
    "import numpy\n",
    "\n",
    "import pandas as pd\n",
    "\n",
    "import ruamel.yaml as yaml\n",
    "\n",
    "_ = alt.data_transformers.disable_max_rows()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e88e8e74-6172-43f3-a4e9-47d533125250",
   "metadata": {},
   "source": [
    "Read the configuration YAML:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85d21482-675b-40fa-b897-f7606b70cfb9",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "# This cell is tagged as parameters. So if you run the notebook using papermill\n",
    "# (https://papermill.readthedocs.io/en/latest/usage-cli.html) with \n",
    "# `-p config_yaml <config_yaml_file>` then the file defined here will be replaced\n",
    "# with the one you pass\n",
    "config_yaml = \"config.yaml\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b33ae279-2351-413f-8a5f-f39fd7be684c",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(config_yaml) as f:\n",
    "    config = yaml.YAML().load(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bddf27e-86c9-475d-90e9-5b9a2be2b1df",
   "metadata": {},
   "source": [
    "## Mutation effects on phenotypes\n",
    "Read all the specified mutation-level phenotypic effects into a Data Frame.\n",
    "Write that data frame to a file.\n",
    "Also created data frame and file in which the mutation-level phenotypic data are randomized among mutations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb07d139-b566-414d-84df-26d917192c91",
   "metadata": {},
   "outputs": [],
   "source": [
    "mutation_phenotypes = []\n",
    "req_cols = [\"site\", \"wildtype\", \"mutant\"]\n",
    "ref_clades = []\n",
    "for ref_clade, ref_clade_d in config[\"mutation_phenotype_csvs\"].items():\n",
    "    ref_clades.append(ref_clade)\n",
    "    for phenotype_category, category_d in ref_clade_d.items():\n",
    "        pheno_cols = category_d[\"columns\"]\n",
    "        print(\n",
    "            f\"Reading {phenotype_category} phenotypes from {category_d['csv']}\\n  \"\n",
    "            + \"\\n  \".join(pheno_cols)\n",
    "        )\n",
    "        if set(req_cols).intersection(pheno_cols):\n",
    "            raise ValueError(f\"{phenotype_category} columns {pheno_cols} include {req_cols}\")\n",
    "        df = pd.read_csv(category_d[\"csv\"])\n",
    "        if not set(req_cols + pheno_cols).issubset(df.columns):\n",
    "            raise ValueError(f\"cannot find expected columns for {phenotype_category}\")\n",
    "        mutation_phenotypes.append(\n",
    "            df\n",
    "            .melt(\n",
    "                id_vars=req_cols,\n",
    "                value_vars=pheno_cols,\n",
    "                var_name=\"phenotype\",\n",
    "                value_name=\"mutation_effect\",\n",
    "            )\n",
    "            .assign(\n",
    "                ref_clade=ref_clade,\n",
    "                phenotype=lambda x: (\n",
    "                    phenotype_category\n",
    "                    if len(pheno_cols) == 1 and pheno_cols[0] == phenotype_category\n",
    "                    else phenotype_category + \" \" + x[\"phenotype\"]\n",
    "                ),\n",
    "            )\n",
    "            [[\"ref_clade\", \"phenotype\", *req_cols, \"mutation_effect\"]]\n",
    "            .query(\"mutation_effect.notnull()\")\n",
    "            .query(\"wildtype != mutant\")  # drop any wildtype to wildtype mutations\n",
    "        )\n",
    "\n",
    "mutation_phenotypes = pd.concat(mutation_phenotypes, ignore_index=True)\n",
    "\n",
    "mutation_phenotypes_csv = config[\"mutation_phenotypes_csv\"]\n",
    "os.makedirs(os.path.dirname(mutation_phenotypes_csv), exist_ok=True)\n",
    "print(f\"\\nWriting mutation phenotypes to {mutation_phenotypes_csv}\")\n",
    "mutation_phenotypes.to_csv(mutation_phenotypes_csv, index=False, float_format=\"%.4g\")\n",
    "\n",
    "n_randomizations = config[\"n_randomizations\"]\n",
    "print(f\"\\nNow randomizing each mutation phenotype {n_randomizations} times\")\n",
    "mutation_phenotypes_randomized = []\n",
    "for _, df in mutation_phenotypes.groupby([\"ref_clade\", \"phenotype\"]):\n",
    "    for random_seed in range(n_randomizations):\n",
    "        numpy.random.seed(random_seed)\n",
    "        mutation_phenotypes_randomized.append(\n",
    "            df\n",
    "            .assign(\n",
    "                mutation_effect=lambda x: numpy.random.permutation(x[\"mutation_effect\"]),\n",
    "                random_seed=random_seed\n",
    "            )\n",
    "            [[\"random_seed\"] + df.columns.tolist()]\n",
    "        )\n",
    "mutation_phenotypes_randomized = pd.concat(mutation_phenotypes_randomized, ignore_index=True)\n",
    "mutation_phenotypes_randomized_csv = config[\"mutation_phenotypes_randomized_csv\"]\n",
    "os.makedirs(os.path.dirname(mutation_phenotypes_randomized_csv), exist_ok=True)\n",
    "print(f\"Writing randomized mutation phenotypes to {mutation_phenotypes_randomized_csv}\")\n",
    "mutation_phenotypes_randomized.to_csv(mutation_phenotypes_randomized_csv, index=False, float_format=\"%.4g\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "457037e6-00c9-46a9-8e57-aad4804e303a",
   "metadata": {},
   "source": [
    "## Read the Pango clades\n",
    "Read Pango clade definitions come from Cornelius Roemer's GitHub repo ([https://github.com/corneliusroemer/pango-sequences](https://github.com/corneliusroemer/pango-sequences)):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a52b188-af5f-4bb8-ae0b-44655517d7a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "pango_json = config[\"pango_json\"]\n",
    "print(f\"Reading Pango clade definitions from {pango_json}\")\n",
    "with urllib.request.urlopen(pango_json) as url:\n",
    "    pango_clades = json.load(url)\n",
    "print(f\"Read definitions for {len(pango_clades)} clades\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da113cfa-0a6a-40fe-92db-b962994f504f",
   "metadata": {},
   "source": [
    "Read Pango clade growth rates from Bedford lab repo ([https://github.com/nextstrain/forecasts-ncov/](https://github.com/nextstrain/forecasts-ncov/)):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef9a53a9-3ddf-432e-8ea4-e65a04180790",
   "metadata": {},
   "outputs": [],
   "source": [
    "pango_growth_json = config[\"pango_growth_json\"]\n",
    "print(f\"Reading Pango clade growth rates from {pango_growth_json}\")\n",
    "with urllib.request.urlopen(pango_growth_json) as url:\n",
    "    pango_growth_dict = json.loads(gzip.decompress(url.read()))\n",
    "\n",
    "# convert to data frame\n",
    "pango_growth = (\n",
    "    pd.DataFrame(pango_growth_dict[\"data\"])\n",
    "    .query(\"location == 'hierarchical'\")\n",
    "    .query(\"site == 'ga'\")\n",
    "    .pivot_table(index=\"variant\", values=\"value\", columns=\"ps\")\n",
    "    .reset_index(names=\"clade\")\n",
    "    .assign(\n",
    "        clade_growth_HDI_95=lambda x: x.apply(\n",
    "            lambda row: f\"{row['HDI_95_lower']:.3g} to {row['HDI_95_upper']:.3g}\",\n",
    "            axis=1,\n",
    "        ),\n",
    "    )\n",
    "    .rename(columns={\"median\": \"clade growth\", \"clade_growth_HDI_95\": \"clade growth HDI 95\"})\n",
    "    [[\"clade\", \"clade growth\", \"clade growth HDI 95\"]]\n",
    "    .query(\"clade != 'other'\")\n",
    ")\n",
    "assert len(pango_growth) == pango_growth[\"clade\"].nunique()\n",
    "extra_growth_clades = set(pango_growth[\"clade\"]) - set(pango_clades)\n",
    "if extra_growth_clades:\n",
    "    raise ValueError(f\"some clades have growth data but are not defined {extra_growth_clades}\")\n",
    "\n",
    "print(f\"Read growth data for {len(pango_growth)} clades\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0206fdf5-21f3-4bc4-9942-5bc502bd4330",
   "metadata": {},
   "source": [
    "Get a data frame of all Pango clades along with other relevant information:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "119c1bfd-10ec-41bb-a8e8-69345341da40",
   "metadata": {},
   "outputs": [],
   "source": [
    "def is_descendant_of(clade, ancestor):\n",
    "    \"\"\"Returns True iff `clade` is a descendant of `ancestor`.\"\"\"\n",
    "    if pango_clades[clade][\"parent\"] == ancestor:\n",
    "        return True\n",
    "    elif pango_clades[clade][\"parent\"]:\n",
    "        return is_descendant_of(pango_clades[clade][\"parent\"], ancestor)\n",
    "    else:\n",
    "        return False\n",
    "\n",
    "def relative_mutations(clade_muts, reference_muts):\n",
    "    \"\"\"Get mutation in `clade_muts` relative `reference_muts`.\"\"\"\n",
    "    shared_muts = set(clade_muts).intersection(reference_muts)\n",
    "    clade_sites = {\n",
    "        r: (wt, m) for (wt, r, m) in [tup for tup in clade_muts if tup not in shared_muts]\n",
    "    }\n",
    "    reference_sites = {\n",
    "        r: (wt, m) for (wt, r, m) in [tup for tup in reference_muts if tup not in shared_muts]\n",
    "    }\n",
    "    muts = []\n",
    "    for r, (wt, m) in clade_sites.items():\n",
    "        if r in reference_sites:\n",
    "            assert wt == reference_sites[r][0]\n",
    "            muts.append((r, reference_sites[r][1], m))\n",
    "        else:\n",
    "            muts.append((r, wt, m))\n",
    "    for r, (wt, m) in reference_sites.items():\n",
    "        if r in clade_sites:\n",
    "            assert wt == clade_sites[r][0]\n",
    "            pass  # already counted\n",
    "        else:\n",
    "            muts.append((r, m, wt))\n",
    "    return [(wt, r, m) for (r, wt, m) in sorted(muts)]\n",
    "\n",
    "def parse_spike_muts(clade_d):\n",
    "    \"\"\"Parse spike mutations from dict for a clade.\"\"\"\n",
    "    return [\n",
    "        (mut.split(\":\")[1][0], int(mut.split(\":\")[1][1: -1]), mut.split(\":\")[1][-1])\n",
    "        for mut in clade_d[\"aaSubstitutions\"] + clade_d[\"aaDeletions\"]\n",
    "        if mut and mut.startswith(\"S:\")\n",
    "    ]\n",
    "\n",
    "clades_df = {}\n",
    "\n",
    "ref_clade_spike_muts = {\n",
    "    ref_clade: parse_spike_muts(pango_clades[ref_clade])\n",
    "    for ref_clade in ref_clades\n",
    "}\n",
    "\n",
    "for clade, clade_d in pango_clades.items():\n",
    "    spike_muts = parse_spike_muts(clade_d)\n",
    "    clades_df[clade] = {\n",
    "        \"date\": clade_d[\"designationDate\"] if clade_d[\"designationDate\"] else pd.NA,\n",
    "        \"parent\": clade_d[\"parent\"] if clade_d[\"parent\"] else pd.NA,\n",
    "        \"spike muts from Wuhan-Hu-1\": \" \".join(f\"{wt}{r}{m}\" for (wt, r, m) in spike_muts),\n",
    "        \"number spike muts from Wuhan-Hu-1\": len(spike_muts),\n",
    "        **{\n",
    "            f\"spike muts from {ref_clade}\":\n",
    "                relative_mutations(spike_muts, ref_clade_spike_muts[ref_clade])\n",
    "            for ref_clade in ref_clades\n",
    "        },\n",
    "        **{\n",
    "            f\"descendant of {ancestor}\": is_descendant_of(clade, ancestor)\n",
    "            for ancestor in config[\"classify_descendants_of\"]\n",
    "        },\n",
    "        \n",
    "    }\n",
    "    \n",
    "clades_df = pd.DataFrame.from_dict(clades_df, orient=\"index\").reset_index(names=\"clade\")\n",
    "\n",
    "# add growth information\n",
    "clades_df = clades_df.merge(pango_growth, on=\"clade\", validate=\"one_to_one\", how=\"left\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae904577-1dd4-42c6-b937-c474fed25c8c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-02-23T22:58:02.715536Z",
     "iopub.status.busy": "2024-02-23T22:58:02.715236Z",
     "iopub.status.idle": "2024-02-23T22:58:02.719701Z",
     "shell.execute_reply": "2024-02-23T22:58:02.719092Z",
     "shell.execute_reply.started": "2024-02-23T22:58:02.715516Z"
    }
   },
   "source": [
    "## Assign phenotypes to clades\n",
    "Phenotypes are just sum of mutation effects of all mutations with respect to the `spike_muts_relative_to` clade.\n",
    "Mutations without measured phenotype effects are ignored (given values of zero)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9e6ca47-7520-4cbe-9a3f-97e202e4ad8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PhenotypeAssigner:\n",
    "    \"\"\"Assign phenotypes to sets of mutations.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    mutation_phenotypes_df : pandas.DataFrame\n",
    "        Should have columns `site`, `wildtype`, `mutant`, `mutation_effect`.\n",
    "        \n",
    "    \"\"\"\n",
    "    def __init__(self, mutation_phenotypes_df):\n",
    "        assert len(mutation_phenotypes_df) == len(\n",
    "            mutation_phenotypes_df[[\"site\", \"mutant\"]].drop_duplicates()\n",
    "        )\n",
    "        self.sites = sorted(set(mutation_phenotypes_df[\"site\"]))\n",
    "        assert len(self.sites) == len(\n",
    "            mutation_phenotypes_df[[\"site\", \"wildtype\"]].drop_duplicates()\n",
    "        )\n",
    "        self.wts = mutation_phenotypes_df.set_index(\"site\")[\"wildtype\"].to_dict()\n",
    "        self.effects = {\n",
    "            site: site_df.set_index(\"mutant\")[\"mutation_effect\"].to_dict()\n",
    "            for site, site_df in mutation_phenotypes_df.groupby(\"site\")\n",
    "        }\n",
    "        for site, wt in self.wts.items():\n",
    "            assert wt not in self.effects[site]\n",
    "            self.effects[site][wt] = 0.0\n",
    "\n",
    "    def phenotype(self, muts):\n",
    "        \"\"\"Returns phenotype for list of `muts` as `(wildtype, site, mutant)`.\"\"\"\n",
    "        pheno = 0.0\n",
    "        for wt, site, m in muts:\n",
    "            if (site in self.effects) and (wt in self.effects[site]) and (m in self.effects[site]):\n",
    "                pheno += self.effects[site][m] - self.effects[site][wt]\n",
    "        return pheno\n",
    "\n",
    "# assign clade phenotypes, Hamming distance from reference, and mutation lists as strings\n",
    "clade_phenos = clades_df.copy()\n",
    "new_clade_pheno_cols = []\n",
    "for (ref_clade, phenotype), df in mutation_phenotypes.groupby(\n",
    "    [\"ref_clade\", \"phenotype\"],\n",
    "    sort=False,\n",
    "):\n",
    "    phenos = PhenotypeAssigner(df)\n",
    "    new_clade_pheno_cols.append(\n",
    "        clade_phenos[f\"spike muts from {ref_clade}\"].map(phenos.phenotype).rename(\n",
    "            f\"{phenotype} relative to {ref_clade}\"\n",
    "        )\n",
    "    )\n",
    "for ref_clade in ref_clades:\n",
    "    col = f\"spike muts from {ref_clade}\"\n",
    "    new_clade_pheno_cols.append(\n",
    "        clade_phenos[col].map(len).rename(f\"Hamming distance relative to {ref_clade}\")\n",
    "    )\n",
    "    clade_phenos[col] = clade_phenos[col].map(\n",
    "        lambda mutlist: \" \".join(f\"{wt}{r}{m}\" for (wt, r, m) in mutlist)\n",
    "    )\n",
    "clade_phenos = pd.concat([clade_phenos] + new_clade_pheno_cols, axis=1)\n",
    "clade_phenotypes_csv = config[\"clade_phenotypes_csv\"]\n",
    "os.makedirs(os.path.dirname(clade_phenotypes_csv), exist_ok=True)\n",
    "print(f\"Writing clade phenotypes to {clade_phenotypes_csv}\")\n",
    "clade_phenos.to_csv(clade_phenotypes_csv, index=False, float_format=\"%.4g\")\n",
    "\n",
    "# assign clade phenotypes from randomized data\n",
    "print(f\"\\nComputing clade phenotypes from {n_randomizations} randomizations of mutation phenotypes\")\n",
    "clade_phenos_randomized = clades_df.copy()\n",
    "new_clade_pheno_randomized_cols = []\n",
    "for (random_seed, ref_clade, phenotype), df in mutation_phenotypes_randomized.groupby(\n",
    "    [\"random_seed\", \"ref_clade\", \"phenotype\"]\n",
    "):\n",
    "    phenos = PhenotypeAssigner(df)\n",
    "    new_clade_pheno_randomized_cols.append(\n",
    "        clade_phenos_randomized[f\"spike muts from {ref_clade}\"].map(phenos.phenotype).rename(\n",
    "            f\"random_{random_seed} {phenotype} relative to {ref_clade}\"\n",
    "        )\n",
    "    )\n",
    "for ref_clade in ref_clades:\n",
    "    col = f\"spike muts from {ref_clade}\"\n",
    "    for random_seed in mutation_phenotypes_randomized[\"random_seed\"].unique():\n",
    "        new_clade_pheno_randomized_cols.append(\n",
    "            clade_phenos_randomized[col].map(len).rename(\n",
    "                f\"random_{random_seed} Hamming distance relative to {ref_clade}\"\n",
    "            )\n",
    "        )\n",
    "    clade_phenos_randomized[col] = clade_phenos_randomized[col].map(\n",
    "        lambda mutlist: \" \".join(f\"{wt}{r}{m}\" for (wt, r, m) in mutlist)\n",
    "    )\n",
    "clade_phenos_randomized = pd.concat([clade_phenos_randomized] + new_clade_pheno_randomized_cols, axis=1)\n",
    "clade_phenotypes_randomized_csv = config[\"clade_phenotypes_randomized_csv\"]\n",
    "os.makedirs(os.path.dirname(clade_phenotypes_randomized_csv), exist_ok=True)\n",
    "print(f\"Writing randomized clade phenotypes to {clade_phenotypes_randomized_csv}\")\n",
    "clade_phenos_randomized.to_csv(clade_phenotypes_randomized_csv, index=False, float_format=\"%.4g\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a112e2df-7a7d-4bd7-a769-5532eb7bec87",
   "metadata": {},
   "source": [
    "## Make some summary plots\n",
    "First get data frame appropriate for plotting with Altair, which includes dropping clades without a date (these are early clades for the most part).\n",
    "We also drop the columns with the spike mutations as those are very large."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ece54087-d8c2-4696-b99f-25e6101fa67b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def dt_to_dec(dt):\n",
    "    \"\"\"Convert a datetime to decimal year (https://stackoverflow.com/a/29852163)\"\"\"\n",
    "    year_start = datetime.datetime(dt.year, 1, 1)\n",
    "    year_end = year_start.replace(year=dt.year+1)\n",
    "    return dt.year + ((dt - year_start).total_seconds() /  # seconds so far\n",
    "        float((year_end - year_start).total_seconds()))  # seconds in year\n",
    "\n",
    "# for Altair plotting, remove \".\" in column names as these cause problems,\n",
    "# and make date column into decimal year dates\n",
    "clade_phenos_alt_df = (\n",
    "    clade_phenos\n",
    "    .drop(columns=[c for c in clade_phenos if c.startswith(\"spike muts from\")])\n",
    "    .rename(columns={c: c.replace(\".\", \"_\") for c in clade_phenos.columns})\n",
    "    .query(\"date.notnull()\")\n",
    "    .assign(date=lambda x: pd.to_datetime(x[\"date\"]).map(dt_to_dec))\n",
    ")\n",
    "\n",
    "pheno_cols = clade_phenos_alt_df.columns.tolist()[\n",
    "    clade_phenos_alt_df.columns.tolist().index(\"clade growth HDI 95\") + 1:\n",
    "]\n",
    "assert all(col.split()[-3: -1] == [\"relative\", \"to\"] for col in pheno_cols)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41171988-5f23-4521-85c4-516ebc40fbfb",
   "metadata": {},
   "source": [
    "Now make the interactive chart:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "246b4cf7-04e7-40f2-8634-b25c5ad6678b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://stackoverflow.com/questions/70937066/make-dropdown-selection-responsive-for-y-axis-altair-python\n",
    "xaxis_param = alt.param(\n",
    "    value=pheno_cols[0],\n",
    "    bind=alt.binding_select(\n",
    "        options=pheno_cols,\n",
    "        name=\"x-axis variable\",\n",
    "    ),\n",
    ")\n",
    "yaxis_param = alt.param(\n",
    "    value=pheno_cols[1],\n",
    "    bind=alt.binding_select(\n",
    "        options=pheno_cols,\n",
    "        name=\"y-axis variable\",\n",
    "    ),\n",
    ")\n",
    "\n",
    "no_growth_point_size = alt.param(\n",
    "    value=50,\n",
    "    bind=alt.binding_range(\n",
    "        min=0,\n",
    "        max=100,\n",
    "        name=\"size of points for clades without growth estimate\",\n",
    "    ),\n",
    ")\n",
    "\n",
    "earliest_date = alt.param(\n",
    "    # initial value shows clades designated in last 1.5 year\n",
    "    value=round(dt_to_dec(pd.to_datetime(datetime.date.today())), 2) - 1.5,\n",
    "    bind=alt.binding_range(\n",
    "        min=round(clade_phenos_alt_df[\"date\"].min(), 2) - 0.01,\n",
    "        max=round(clade_phenos_alt_df[\"date\"].max(), 2) + 0.01,\n",
    "        name=\"only show clades designated after this date\",\n",
    "    ),\n",
    ")\n",
    "\n",
    "# make chart to select descendants\n",
    "classify_descendant_cols = [\"descendant of any\"] + [\n",
    "    col for col in clade_phenos_alt_df.columns if col.startswith(\"descendant of\")\n",
    "]\n",
    "assert len(classify_descendant_cols) == len(config[\"classify_descendants_of\"]) + 1\n",
    "\n",
    "select_descendants_of = alt.selection_point(\n",
    "    fields=[\"descendant_of\"],\n",
    "    value=classify_descendant_cols[1],\n",
    "    toggle=\"true\",\n",
    ")\n",
    "\n",
    "descendants_of_chart = (\n",
    "    alt.Chart(pd.DataFrame({\"descendant_of\": classify_descendant_cols}))\n",
    "    .add_params(select_descendants_of)\n",
    "    .encode(\n",
    "        alt.Y(\n",
    "            \"descendant_of\",\n",
    "            sort=classify_descendant_cols,\n",
    "            axis=alt.Axis(labelExpr=\"split(datum.label, ' ')[2]\", labelFontSize=11),\n",
    "            title=\"click clade(s) to show descendants\",\n",
    "        ),\n",
    "        fill=alt.condition(select_descendants_of, alt.value(\"gray\"), alt.value(\"lightgray\")),\n",
    "    )\n",
    "    .mark_rect(stroke=\"black\", strokeWidth=2)\n",
    ")\n",
    "\n",
    "select_clade = alt.selection_point(on=\"mouseover\", fields=[\"clade\"], empty=False)\n",
    "\n",
    "scatter_chart = (\n",
    "    alt.Chart(clade_phenos_alt_df.assign(**{\"descendant of any\": True}))\n",
    "    .add_params(xaxis_param, yaxis_param, earliest_date, no_growth_point_size, select_clade)\n",
    "    .transform_filter(alt.expr.isValid(alt.datum[\"clade growth\"]) | (no_growth_point_size > 0))\n",
    "    .transform_filter(alt.datum[\"date\"] >= earliest_date)\n",
    "    .transform_fold(classify_descendant_cols, [\"descendant_of\", \"is_descendant\"])\n",
    "    .transform_filter(select_descendants_of)\n",
    "    .transform_aggregate(\n",
    "        is_descendant=\"sum(is_descendant)\",\n",
    "        groupby=[c for c in clade_phenos_alt_df.columns if c not in classify_descendant_cols],\n",
    "    )\n",
    "    .transform_filter(alt.datum[\"is_descendant\"] > 0)\n",
    "    .transform_calculate(\n",
    "        x=f\"datum[{xaxis_param.name}]\",\n",
    "        y=f\"datum[{yaxis_param.name}]\",\n",
    "        point_size=alt.expr.if_(\n",
    "            alt.expr.isValid(alt.datum[\"clade growth\"]),\n",
    "            100,\n",
    "            no_growth_point_size,\n",
    "        ),\n",
    "    )\n",
    "    .encode(\n",
    "        alt.X(\n",
    "            \"x:Q\",\n",
    "            title=\"x-axis variable (select with dropdown below)\",\n",
    "            scale=alt.Scale(nice=False, zero=False, padding=9),\n",
    "        ),\n",
    "        alt.Y(\n",
    "            \"y:Q\",\n",
    "            title=\"y-axis variable (select with dropdown below)\",\n",
    "            scale=alt.Scale(nice=False, zero=False, padding=9),\n",
    "        ),\n",
    "        # show clades with / without growth estimates differently:\n",
    "        # https://stackoverflow.com/a/73063556\n",
    "        fill=alt.condition(\n",
    "            alt.expr.isValid(alt.datum[\"clade growth\"]),\n",
    "            alt.Fill(\n",
    "                \"clade growth:Q\",\n",
    "                scale=alt.Scale(scheme=\"viridis\"),\n",
    "                title=\"clade growth (gray means no estimate)\",\n",
    "                legend=alt.Legend(\n",
    "                    orient=\"bottom\",\n",
    "                    titleLimit=500,\n",
    "                    titleFontSize=13,\n",
    "                    gradientStrokeColor=\"black\",\n",
    "                    gradientStrokeWidth=1,\n",
    "                ),\n",
    "            ),\n",
    "            alt.value(\"gray\"),\n",
    "        ),\n",
    "        size=alt.Size(\"point_size:Q\", legend=None, scale=alt.Scale(range=[0, 200])),\n",
    "        fillOpacity=alt.Opacity(\"point_size:Q\", legend=None, scale=alt.Scale(range=[0.4, 0.85])),\n",
    "        strokeOpacity=alt.condition(\n",
    "            select_clade,\n",
    "            alt.value(1),\n",
    "            alt.Opacity(\"point_size:Q\", legend=None, scale=alt.Scale(range=[0.4, 1])),\n",
    "        ),\n",
    "        stroke=alt.condition(select_clade, alt.value(\"red\"), alt.value(\"black\")),\n",
    "        strokeWidth=alt.condition(select_clade, alt.value(2), alt.value(0.75)),\n",
    "        tooltip=[\n",
    "            \"clade\",\n",
    "            alt.Tooltip(\"date\", format=\".2f\"),\n",
    "            \"clade growth HDI 95\",\n",
    "        ],\n",
    "    )\n",
    "    .mark_circle(size=50, invalid=None)\n",
    "    .interactive()\n",
    "    .properties(width=350, height=350)\n",
    ")\n",
    "\n",
    "chart = (\n",
    "    alt.hconcat(scatter_chart, descendants_of_chart, center=True)\n",
    "    .configure_axis(grid=False, titleFontSize=13)\n",
    "    .properties(\n",
    "        title=alt.TitleParams(\n",
    "            \"Spike phenotypes and growth of SARS-CoV-2 clades\",\n",
    "            subtitle=[\n",
    "                \"This chart is interactive: dropdown menus at bottom to adjust which phenotypes to show,\",\n",
    "                \"filter dates of clades shown, and otherwise adjust appearance. Click on the boxes at right\",\n",
    "                \"to only show clades that are descendants of key clades, and mouse over points for details.\",\n",
    "                \"\",\n",
    "                \"See github.com/jbloomlab/SARS2-spike-predictor-phenos for code and details.\",\n",
    "                \"\",\n",
    "                f\"This plot generated {datetime.date.today().strftime('%b-%d-%Y')}.\",\n",
    "            ],\n",
    "            fontSize=15,\n",
    "            dy=-10\n",
    "        )\n",
    "    )\n",
    ")\n",
    "\n",
    "chart"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e6f65ad-678a-481d-abf8-35c3ba7843b7",
   "metadata": {},
   "source": [
    "Save chart:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "951af52c-40ce-4e47-a881-5ef1759aa99b",
   "metadata": {},
   "outputs": [],
   "source": [
    "chart_file = config[\"clade_phenotype_chart\"]\n",
    "os.makedirs(os.path.dirname(chart_file), exist_ok=True)\n",
    "print(f\"Saving chart to {chart_file}\")\n",
    "chart.save(chart_file)"
   ]
  }
 ],
 "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.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}