notebooks/2-Composing-Linking-Scatter-Plots.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fec34920-0fbc-4b8c-b9b4-e52575bd11b9",
   "metadata": {},
   "source": [
    "# Composing and Linking Multiple Scatter Plots\n",
    "\n",
    "In this notebook, we'll learn:\n",
    "1. [About `jscatter`'s API for plotting multiple scatter plots](#API-for-Composing-Multiple-Scatter-Plots)\n",
    "2. [How to synchronize selections using Fashion MNIST embeddings](#Synchronizing-the-Selection-and-Hover)\n",
    "3. [How to synchronize views using LLM-based sentence embeddings](#Synchronizing-Views)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31dfbd5c-fdfb-4a21-bc5b-299c3f69e706",
   "metadata": {},
   "outputs": [],
   "source": [
    "# If you run this notebook in Google Colab, you need to manually install the following packages.\n",
    "# !pip install --quiet jupyter-scatter"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3eec7b4a-b13e-46fd-9200-279154dfc64f",
   "metadata": {},
   "source": [
    "## API for Composing Multiple Scatter Plots\n",
    "\n",
    "We'll start out with a very simple example to get familiar with the API.\n",
    "\n",
    "In the following we'll compose two scatter plots next to each other using `jscatters.compose()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f5991d8-3bfa-4d5f-a184-f0d5bc573035",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jscatter import Scatter, compose\n",
    "from numpy.random import rand\n",
    "\n",
    "a = Scatter(x=rand(500), y=rand(500))\n",
    "b = Scatter(x=rand(5000), y=rand(5000))\n",
    "\n",
    "compose([a, b])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bec7d3f2-3ade-47d5-a450-b4580323ea23",
   "metadata": {},
   "source": [
    "By default, `jscatter` arranges scatter plots into a single row but we can customize this of course."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03db38e9-af0e-4533-b43d-65b984ec3f19",
   "metadata": {},
   "outputs": [],
   "source": [
    "compose([a, b], rows=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f3594f9-ebb4-4efd-8ce5-c443fcde436e",
   "metadata": {},
   "source": [
    "So good so far but the fun part starts when we link/synchronize the scatter plots' views and selections.\n",
    "\n",
    "## Synchronizing the Selection and Hover\n",
    "\n",
    "### Comparing Embedding Methods\n",
    "\n",
    "To demonstrate the usefulness of linked/synchronized selections, let's take a look the [Fashion MNIST dataset](https://github.com/zalandoresearch/fashion-mnist), which we embedded using four different embedding methods:\n",
    "\n",
    "1. [PCA](https://scikit-learn.org/stable/modules/generated/sklearn.decomposition.PCA.html)\n",
    "2. [t-SNE](https://opentsne.readthedocs.io/en/stable/)\n",
    "3. [UMAP](https://umap-learn.readthedocs.io/en/latest/)\n",
    "4. [A convolutional autoencoder](https://blog.keras.io/building-autoencoders-in-keras.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7162c714-8f1e-4f02-b183-dc8e307eee4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p data\n",
    "!curl -L -C - -o data/fashion-mnist-embeddings.pq https://storage.googleapis.com/flekschas/jupyter-scatter-tutorial/fashion-mnist-embeddings.pq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "823ede65-6bf9-4a74-bea1-b11c0b409e8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "fashion_mnist_embeddings = pd.read_parquet('data/fashion-mnist-embeddings.pq')\n",
    "fashion_mnist_embeddings = fashion_mnist_embeddings.replace({\"class\": {0: \"T-shirt/top\", 1: \"Trouser\", 2: \"Pullover\", 3: \"Dress\", 4: \"Coat\", 5: \"Sandal\", 6: \"Shirt\", 7: \"Sneaker\", 8: \"Bag\", 9: \"Ankle boot\"}})\n",
    "fashion_mnist_embeddings['class'] = fashion_mnist_embeddings['class'].astype('category')\n",
    "fashion_mnist_embeddings.head(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d5361bd-baf3-414c-b88e-46004d9a9c80",
   "metadata": {},
   "source": [
    "The dataframe contains pre-embedded x/y locations of each image and the associated class.\n",
    "\n",
    "Since we're going to visualize each embedding using the same visual encoding, we can specify most things upfront:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fddf0dea-ef44-4604-b348-6081f11a6472",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = dict(\n",
    "    background_color='#111111',\n",
    "    color_by='class',\n",
    "    color_map={\n",
    "        \"T-shirt/top\": '#FFFF00',\n",
    "        \"Trouser\": '#1CE6FF',\n",
    "        \"Pullover\": '#FF34FF',\n",
    "        \"Dress\": '#FF4A46',\n",
    "        \"Coat\": '#008941',\n",
    "        \"Sandal\": '#006FA6',\n",
    "        \"Shirt\": '#A30059',\n",
    "        \"Sneaker\": '#FFDBE5',\n",
    "        \"Bag\": '#7A4900',\n",
    "        \"Ankle boot\": '#0000A6'\n",
    "    },\n",
    "    legend=True,\n",
    "    axes=False,\n",
    "    zoom_on_selection=True, # To automatically zoom to selected points\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a6fc25e-e2e1-4c7a-b406-1cdb79b37ff7",
   "metadata": {},
   "source": [
    "Finally, we need to create four `jscatter` instances and compose them in a 2x2 grid. This time however, we're going to link/synchronize the selection and point hovering across all four instances because each scatter plot references the same images from Fashion MNIST."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f44e21c5-fb11-42a9-964d-28f3a8d682d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "pca = Scatter(data=fashion_mnist_embeddings, x='pcaX', y='pcaY', **config)\n",
    "tsne = Scatter(data=fashion_mnist_embeddings, x='tsneX', y='tsneY', **config)\n",
    "umap = Scatter(data=fashion_mnist_embeddings, x='umapX', y='umapY', **config)\n",
    "cae = Scatter(data=fashion_mnist_embeddings, x='caeX', y='caeY', **config)\n",
    "\n",
    "compose(\n",
    "    [pca, tsne, umap, cae],\n",
    "    sync_selection=True,\n",
    "    sync_hover=True,\n",
    "    rows=2,\n",
    "    row_height=240\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f9687a0-5dda-4c41-bca0-302686ce23ff",
   "metadata": {},
   "source": [
    "Because I like to see the selected points within their local neighborhood, I activated `zoom_on_selection`. In combination with the synced selection, this makes all scatter plots automatically zoom to selected points.\n",
    "\n",
    "## Synchronizing Views\n",
    "\n",
    "### For Faceted Exploration or Shared Latent Spaces\n",
    "\n",
    "Beyond synchronizing the selection, `jscatter` also supports view synchronization. However, it does not make much sense to activate this for the above Fashion MNIST example because each scatter plots drew a different embedding space. However, we might want to explore a large dataset where all points share the same latent space. In this case it can be interesting to facet the dataset to for comparison. And since the space is the same, it can be useful to synchronize the view.\n",
    "\n",
    "> 🚨 LLM Alert\n",
    "\n",
    "In the next example we're going to compare news articles from 2012-2022 by their title. For that we're using the fantastic [News Category Dataset](https://www.kaggle.com/datasets/rmisra/news-category-dataset?resource=download) from [Rishabh Misra, 2022](https://arxiv.org/abs/2209.11429). We embedded the titles, abstract, and both using a combination of the pretrained [all-MiniLM-L6-v2 sentence transformer from 🤗](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2) and [UMAP](https://umap-learn.readthedocs.io/en/latest/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37781a24-a5ca-4386-a931-c091dcd225e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p data\n",
    "!curl -L -C - -o data/huffpost-embeddings.pq https://storage.googleapis.com/flekschas/jupyter-scatter-tutorial/huffpost-embeddings.pq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a5ce07a-e827-4636-8e91-17d022c3487e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "huffpost_embeddings = pd.read_parquet('data/huffpost-embeddings.pq')\n",
    "huffpost_embeddings.head(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67a7aa17-2517-4d71-9cd9-0660391740ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jscatter import glasbey_light\n",
    "\n",
    "category_cmap = { cat: glasbey_light[i] for i, cat in enumerate(sorted(huffpost_embeddings.category.unique())) }\n",
    "\n",
    "huffpost_scatter_config = dict(axes=False, background_color='#111111')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69f67600-6635-4514-af35-8631e3b71715",
   "metadata": {},
   "outputs": [],
   "source": [
    "huffpost_scatter = Scatter(\n",
    "    data=huffpost_embeddings,\n",
    "    x='x',\n",
    "    y='y',\n",
    "    color_by='category',\n",
    "    color_map=category_cmap,\n",
    "    height=640,\n",
    "    legend=True,\n",
    "    **huffpost_scatter_config\n",
    ")\n",
    "huffpost_scatter.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8e55325-7747-4277-9a8b-fbc43501d726",
   "metadata": {},
   "source": [
    "As mentioned above, this dataset consists of news articles from 2012 – 2022. An intersting question is, whether the distribution of published articles has changed over the years. To achieve this we're going to facet the data frame by year and plot each year as an individual scatter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4db93e68-345b-4102-9496-9cd0ee66fb4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_annual_scatter(year):\n",
    "    return Scatter(\n",
    "        data=huffpost_embeddings_years[year],\n",
    "        x='x',\n",
    "        y='y',\n",
    "        color_by='category',\n",
    "        color_map=category_cmap,\n",
    "        **huffpost_scatter_config\n",
    "    )\n",
    "\n",
    "years = sorted(huffpost_embeddings.year.unique())\n",
    "\n",
    "huffpost_embeddings_years = {\n",
    "    year: huffpost_embeddings[huffpost_embeddings.year == year] for year in years\n",
    "}\n",
    "\n",
    "huffpost_scatters_years = {\n",
    "    year: create_annual_scatter(year) for year in years\n",
    "}\n",
    "\n",
    "compose([(sc, y) for y, sc in huffpost_scatters_years.items()], sync_view=True, rows=3, cols=4, row_height=320)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c64f89b0-0dae-40ac-9f3b-b5ef2ce7ca1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "huffpost_embeddings_years['2012'].iloc[huffpost_scatters_years['2012'].selection()]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e693978-a65d-433d-83ea-34fa2d4a643d",
   "metadata": {},
   "source": [
    "## Synchronize Everything\n",
    "\n",
    "### Visualizing Multiple Properties of the Same Dataset\n",
    "\n",
    "Sometimes it can be useful to sychronize everything: the view, selection, and hover state. This is typically the case when one wants to simply explore many different properties at the same time.\n",
    "\n",
    "This is common in exploring single-cell data. The purpose of the embedding visualization is not only to inform the biologist about cell type clusters but to also allow them to verify cluster validity by visually collerating clusters with known cell type marker expressions. To demonstrate this use case, let's take another look at the single-cell data from from [Mair et al., 2022](https://www.nature.com/articles/s41586-022-04718-w) that was clustered with [Ozette](https://www.ozette.com/)'s [FAUST method](https://doi.org/10.1016/j.patter.2021.100372) and transformed with [Ozette's Annotation Transformation](https://github.com/flekschas-ozette/ismb-biovis-2022) prior to embedding with [UMAP](https://umap-learn.readthedocs.io/en/latest/).\n",
    "\n",
    "Given the frequent use of this scenario and since I'm lazy, `jscatter` offers a short-hand function called `link` to synchronize everything for you."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a13dc21d-4195-4123-94b1-2c1984bef8cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from itertools import cycle\n",
    "from jscatter import glasbey_light, link\n",
    "\n",
    "mair_2022_tumor_ozette = pd.read_parquet(\"./data/mair-2022-tumor-006-ozette.pq\")\n",
    "\n",
    "mair_2022_colormap = dict(zip(mair_2022_tumor_ozette.faustLabels.unique(), cycle(glasbey_light[1:])))\n",
    "mair_2022_colormap[\"0_0_0_0_0\"] = (0.2, 0.2, 0.2, 1.0)\n",
    "\n",
    "mair_2022_scatter_config = dict(\n",
    "    data=mair_2022_tumor_ozette,\n",
    "    x='umapX',\n",
    "    y='umapY',\n",
    "    background_color=\"#111111\",\n",
    "    axes=False,\n",
    ")\n",
    "\n",
    "link([\n",
    "    (Scatter(color_by='faustLabels', color_map=mair_2022_colormap, **mair_2022_scatter_config), \"Cell Types\"),\n",
    "    (Scatter(color_by='CD4_Windsorized', legend=True, color_labeling=(\"low\", \"high\"), **mair_2022_scatter_config), \"CD4 Expression\"),\n",
    "    (Scatter(color_by='CD8_Windsorized', legend=True, color_labeling=(\"low\", \"high\"), **mair_2022_scatter_config), \"CD8 Expression\"),\n",
    "    (Scatter(color_by='CD19_Windsorized', legend=True, color_labeling=(\"low\", \"high\"), **mair_2022_scatter_config), \"CD19 Expression\")\n",
    "], rows=2, row_height=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad658073-642f-4864-ac11-c94e3ac9a3c1",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Next\n",
    "\n",
    "Next up, we'll how to use everything we've learned so far and build bespoke interfaces for exploring large-scale datasets starting with the LLM-based news article dataset we just explored.\n",
    "\n",
    "➡️ [Building a Bespoke Interface for Exploring LLM-Based Sentence Embeddings.ipynb](3-LLM-Sentence-Embedding.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c750af45-8b47-4970-92cc-dc815bc06762",
   "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}