4dn-demo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a747e5eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dataclasses\n",
    "import functools\n",
    "import itertools\n",
    "import pathlib\n",
    "from typing import Iterable, Union\n",
    "import uuid\n",
    "\n",
    "import bioframe as bf\n",
    "import higlass as hg\n",
    "import ipywidgets\n",
    "import jscatter\n",
    "import matplotlib.colors as colors\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "import traitlets\n",
    "from IPython.display import display\n",
    "\n",
    "data_dir = pathlib.Path.cwd() / \"data\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41f192c0",
   "metadata": {},
   "source": [
    "## Linked embeddings of single-locus interaction profiles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1c24b310",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def load_eigs(stages: tuple[str, ...]) -> list[pd.DataFrame]:\n",
    "    eigs = []\n",
    "    for name in stages:\n",
    "        eigs_path = (\n",
    "            data_dir\n",
    "            / f\"eigvecs.proj__ESC-Trypsin-FA-DSG-HindIII.{name}-Trypsin-FA-DSG-HindIII.E1-E10.hg38.50000.pq\"\n",
    "        )\n",
    "        eigs_df = pd.read_parquet(eigs_path).dropna()\n",
    "        eigs.append(eigs_df)\n",
    "\n",
    "    # create shared index for all dataframes\n",
    "    index = functools.reduce(pd.Index.intersection, [df.index for df in eigs])\n",
    "    return [df.loc[index] for df in eigs]\n",
    "\n",
    "def load_track_metadata() -> tuple[pd.DataFrame, dict[str, dict]]:\n",
    "    tracks = pd.read_parquet(data_dir / \"tracks.hg38.50000.pq\")\n",
    "    clusters = pd.read_parquet(\n",
    "        data_dir\n",
    "        / \"clusters.joint.proj__ESC-Trypsin-FA-DSG-HindIII.E1-E10.50000.kmeans_sm.pq\"\n",
    "    ).filter(regex=\"^kmeans_sm\\d+$\")\n",
    "    for col in clusters:\n",
    "        clusters[col] = clusters[col].astype(str)\n",
    "    return tracks.join(clusters), {\n",
    "        \"GC\": dict(norm=colors.Normalize(vmin=0.35, vmax=0.65), map=\"RdYlBu_r\"),\n",
    "        \"centel_abs\": dict(norm=colors.Normalize(vmin=0, vmax=149043529), map=\"Greys\"),\n",
    "        \"kmeans_sm8\": dict(\n",
    "            relabel={ \"7\": \"A1\", \"5\": \"A2\", \"6\": \"AB\", \"4\": \"A3\", \"3\": \"A3\", \"1\": \"B0\", \"0\": \"B0\", \"2\": \"B4\", \"8\": np.nan },\n",
    "            map={'A1': '#e23838', 'A2': '#f78200', 'AB': '#5ebd3e', 'A3': '#ffb900', 'B0': 'cornflowerblue', 'B4': '#973999'}\n",
    "        )\n",
    "    }\n",
    "\n",
    "def init_dropdowns(\n",
    "    x: str,\n",
    "    y: str,\n",
    "    color: str,\n",
    "    scatters: Iterable[jscatter.Scatter],\n",
    "):\n",
    "    tracks, color_config = load_track_metadata()\n",
    "    xy_options = [f\"E{i}\" for i in range(1, 11)]    \n",
    "    color_options = [c for c in tracks.columns if c not in [\"start\", \"end\"]]\n",
    "    color_options.remove(color)\n",
    "    color_options.insert(0, color)\n",
    "\n",
    "    x_dropdown = ipywidgets.Dropdown(options=xy_options, value=x, description=\"x:\")\n",
    "\n",
    "    def on_change_x(change):\n",
    "        for scatter in scatters:\n",
    "            scatter.x(change.new)\n",
    "\n",
    "    y_dropdown = ipywidgets.Dropdown(options=xy_options, value=y, description=\"y:\")\n",
    "\n",
    "    def on_change_y(change):\n",
    "        for scatter in scatters:\n",
    "            scatter.y(change.new)\n",
    "\n",
    "    c_dropdown = ipywidgets.Dropdown(\n",
    "        options=color_options, value=color, description=\"color:\"\n",
    "    )\n",
    "    \n",
    "    def extract_color_series(\n",
    "        track_data: pd.DataFrame,\n",
    "        field: str,\n",
    "        color_kwargs: Union[dict, None],\n",
    "    ):\n",
    "        data = track_data[field]\n",
    "        if color_kwargs is None:\n",
    "            color_kwargs = {}\n",
    "            if data.dtype.name in (\"object\", \"category\"):\n",
    "                data = data.astype(\"category\") # ensure categorical\n",
    "                color_kwargs[\"map\"] = dict(zip(data.cat.categories, jscatter.glasbey_dark))\n",
    "            else:\n",
    "                color_kwargs[\"norm\"] = colors.Normalize(vmin=data.min(), vmax=data.max())\n",
    "                color_kwargs[\"map\"] = \"viridis_r\"\n",
    "            return data, color_kwargs\n",
    "        if color_kwargs and \"relabel\" in color_kwargs:\n",
    "            data = data.map(color_kwargs[\"relabel\"])\n",
    "                        \n",
    "        return data, color_kwargs\n",
    "\n",
    "    def on_change_color(change):\n",
    "        field = change[\"new\"]\n",
    "        for scatter in scatters:\n",
    "            track_data = tracks.loc[scatter._data.index]\n",
    "            data, kwargs = extract_color_series(track_data, field, color_config.get(field))\n",
    "            scatter._data[\"_color\"] = data\n",
    "            scatter.color(by=\"_color\", **kwargs)\n",
    "\n",
    "    x_dropdown.observe(on_change_x, names=[\"value\"])\n",
    "    y_dropdown.observe(on_change_y, names=[\"value\"])\n",
    "    c_dropdown.observe(on_change_color, names=[\"value\"])\n",
    "    on_change_color(dict(new=color))\n",
    "\n",
    "    return x_dropdown, y_dropdown, c_dropdown\n",
    "\n",
    "def init_scatters(\n",
    "    stages: tuple[str, ...] = (\"ESC\", \"DE\", \"HB\", \"iHEP\", \"mHEP\"),\n",
    "    x: str = \"E1\",\n",
    "    y: str = \"E2\",\n",
    "    color: str = \"kmeans_sm8\",\n",
    "):\n",
    "    eigs = load_eigs(stages=stages)\n",
    "    scatters = [jscatter.Scatter(x=x, y=y, data=data, opacity=0.5) for data in eigs]\n",
    "    src = scatters[0]\n",
    "    for target in scatters[1:]:\n",
    "        ipywidgets.jslink((src.widget, \"selection\"), (target.widget, \"selection\"))\n",
    "    \n",
    "    dropdowns = init_dropdowns(x=x, y=y, color=color, scatters=scatters)\n",
    "    \n",
    "    component = ipywidgets.VBox([\n",
    "        ipywidgets.HBox([\n",
    "            ipywidgets.VBox([ipywidgets.Label(name), s.show()])\n",
    "            for name, s in zip(stages, scatters)\n",
    "        ]),\n",
    "        ipywidgets.HBox(dropdowns)\n",
    "    ])\n",
    "    \n",
    "    def extract_coords(ind):\n",
    "        return scatters[0]._data.iloc[ind][[\"chrom\", \"start\", \"end\"]]\n",
    "    \n",
    "    # we expose a single \"selection\" for this component, which the viewer can subscribe to\n",
    "    component.add_traits(coords=traitlets.Any(extract_coords([])))\n",
    "    \n",
    "    ipywidgets.dlink(\n",
    "        source=(scatters[0].widget, \"selection\"),\n",
    "        target=(component, \"coords\"),\n",
    "        transform=extract_coords,\n",
    "    )\n",
    "    \n",
    "    return component"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c5b2f39",
   "metadata": {},
   "source": [
    "## Linked RNA seq violinplot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "bd8eb757",
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_violinplot(scatter):\n",
    "    rna = pd.read_parquet(data_dir / \"tracks.hg38.50000.pq\").filter(regex=\"R1.fwd$\")\n",
    "    fig, ax = plt.figure(), plt.subplot(111)\n",
    "    output = ipywidgets.Output()\n",
    "\n",
    "    ax.set(xlabel=\"stage\", ylabel='log(n + 1)')\n",
    "\n",
    "    def on_selection_change(change):\n",
    "        try:\n",
    "            ax.clear()\n",
    "            data = np.log10(rna.loc[change.new.index] + 1).melt()\n",
    "            sns.violinplot(x=\"variable\", y=\"value\", data=data, ax=ax)\n",
    "            output.clear_output()\n",
    "            with output:\n",
    "                display(fig)\n",
    "        except Exception as e:\n",
    "            with output:\n",
    "                print(e)\n",
    "\n",
    "    scatter.observe(on_selection_change, names=[\"coords\"])\n",
    "    return output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e58b3758",
   "metadata": {},
   "source": [
    "## HiGlass view"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "15231d8f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def _encode(obj):\n",
    "    if isinstance(obj, np.integer):\n",
    "        return int(obj)\n",
    "    elif isinstance(obj, np.floating):\n",
    "        return float(obj)\n",
    "    elif isinstance(obj, np.ndarray):\n",
    "        return obj.tolist()\n",
    "    return obj\n",
    "\n",
    "@dataclasses.dataclass\n",
    "class DynamicTileset:\n",
    "    name: str\n",
    "    chromsizes: pd.Series\n",
    "    uid: str = dataclasses.field(default_factory=lambda: uuid.uuid4().hex)\n",
    "    datatype = \"bedlike\"\n",
    "        \n",
    "    def __post_init__(self):\n",
    "        abslen = self.chromsizes.cumsum()\n",
    "        starts = pd.Series([0] + abslen[:-1].tolist(), index=abslen.index)\n",
    "        self._starts =  dict(starts)\n",
    "        self._tiles = []\n",
    "        \n",
    "    def update(self, df):\n",
    "        starts = self._starts\n",
    "        df = bf.cluster(df).groupby(\"cluster\").agg({\n",
    "            \"chrom\": \"first\",\n",
    "            \"start\": \"min\",\n",
    "            \"end\": \"max\",\n",
    "            \"name\": \"first\",\n",
    "            \"score\": \"first\",\n",
    "            \"strand\": \"first\",\n",
    "            \"thickStart\": \"min\",\n",
    "            \"thickEnd\": \"max\",\n",
    "            \"rgb\": \"first\",\n",
    "        })\n",
    "        self._tiles = [{\n",
    "            \"chrOffset\": int(starts[r[0]]),\n",
    "            \"xStart\": int(starts[r[0]] + r[1]),\n",
    "            \"xEnd\": int(starts[r[0]] + r[2]),\n",
    "            \"importance\": 0,\n",
    "            \"uid\": uuid.uuid4().hex,\n",
    "            \"fields\": tuple(map(_encode, r))\n",
    "        } for r in df.to_records(index=False)]\n",
    "        \n",
    "    def info(self):\n",
    "        genome_length = int(np.sum(self.chromsizes.values))\n",
    "        return {\n",
    "            \"uuid\": self.uid,\n",
    "            \"max_width\": genome_length,\n",
    "            \"min_pos\": [1],\n",
    "            \"max_pos\": [genome_length],\n",
    "            \"max_zoom\": 0,\n",
    "        }\n",
    "     \n",
    "    def tiles(self, _tileids):\n",
    "        return [(f\"{self.uid}.0.0\", self._tiles)]\n",
    "\n",
    "def init_dynamic_tileset():\n",
    "    bins = bf.read_table(\n",
    "        data_dir / \"clusters.proj__all.JOINT.50000.E1-E10.kmeans_sm8.bed\",\n",
    "        schema=\"bed9\",\n",
    "        schema_is_strict=True,\n",
    "    )\n",
    "    tileset = DynamicTileset(name=\"IPG clusters\", chromsizes=bf.fetch_chromsizes('hg38')[:'chrY'])\n",
    "    tileset.update(bins)\n",
    "    return tileset, bins\n",
    "\n",
    "def init_dynamic_track():\n",
    "    tileset, bins = init_dynamic_tileset()\n",
    "    track = hg.server.add(tileset).track(height=30).opts(fillOpacity=1)\n",
    "    def update_tileset(coords):\n",
    "        if len(coords) > 0:\n",
    "            inds = bf.overlap(bins, coords, how=\"inner\", return_index=True, return_input=False)\n",
    "            tileset.update(bins.iloc[inds[\"index\"]])\n",
    "        else:\n",
    "            tileset.update(bins)\n",
    "        \n",
    "    return track, update_tileset\n",
    "\n",
    "def init_viewer(scatter: Union[None, traitlets.HasTraits] = None):\n",
    "    conf = hg.Viewconf.parse_file(\"./viewconf.json\")\n",
    "    track, update_tileset = init_dynamic_track()\n",
    "    view = conf.views[0]\n",
    "    view.tracks.top[1].height= 100\n",
    "    view.tracks.top.append(track)\n",
    "    \n",
    "    hg_viewer = conf.widget()\n",
    "    \n",
    "    if scatter is None:\n",
    "        return hg_viewer\n",
    "    \n",
    "    def on_coords_change(change):\n",
    "        update_tileset(change.new)\n",
    "        hg_viewer.reload(dict(trackId=track.uid, viewId=view.uid))\n",
    "\n",
    "    scatter.observe(on_coords_change, names=[\"coords\"])\n",
    "    return hg_viewer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0bd7697",
   "metadata": {},
   "source": [
    "## Demo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "11c80d45",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "91d4ca726c874e418fc48c62b9c61785",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(HiGlassWidget(), VBox(children=(HBox(children=(VBox(children=(Label(value='ESC'), HBox(children…"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def init():\n",
    "    scatter = init_scatters()\n",
    "    hg_viewer = init_viewer(scatter)\n",
    "    return hg_viewer, scatter\n",
    "\n",
    "ipywidgets.VBox(init())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65253ea3",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2f1f6ae-7883-4627-91ca-f07c32d9bea7",
   "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}