{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Core tracking API by examples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook illustrates the use of core APIs with a simple example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    import google.colab\n",
    "\n",
    "    %pip install -q --upgrade laptrack matplotlib spacy flask pandas\n",
    "    # upgrade packages to avoid pip warnings if the notebook is run in colab\n",
    "except:\n",
    "    %pip install -q --upgrade laptrack matplotlib pandas"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note**: Restart the runtime when you are on Google Colab and you see an error in matplotlib."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Importing packages\n",
    "\n",
    "`laptrack.LapTrack` is the core object for tracking."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from os import path\n",
    "import pandas as pd\n",
    "from IPython.display import display\n",
    "from matplotlib import pyplot as plt\n",
    "from laptrack import LapTrack\n",
    "from laptrack import datasets\n",
    "\n",
    "plt.rcParams[\"font.family\"] = \"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the example point coordinates from a CSV file.\n",
    "\n",
    "`spots_df` has columns `[\"frame\", \"position_x\", \"position_y\"]` \n",
    "(can be arbitrary, and the column names for tracking will be specified later)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spots_df = datasets.simple_tracks()\n",
    "display(spots_df.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tracking"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initializing LapTrack object"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, initialize the `LapTrack` object with parameters. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_distance = 15\n",
    "lt = LapTrack(\n",
    "    metric=\"sqeuclidean\",  # The similarity metric for particles. See `scipy.spatial.distance.cdist` for allowed values.\n",
    "    splitting_metric=\"sqeuclidean\",\n",
    "    merging_metric=\"sqeuclidean\",\n",
    "    # the square of the cutoff distance for the \"sqeuclidean\" metric\n",
    "    cutoff=max_distance**2,\n",
    "    splitting_cutoff=max_distance**2,  # or False for non-splitting case\n",
    "    merging_cutoff=max_distance**2,  # or False for non-merging case\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example 1: using `predict_dataframe` to track pandas DataFrame coordinates"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`predict_dataframe` is the easiest option when you have the coordinate data in pandas DataFrame.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "track_df, split_df, merge_df = lt.predict_dataframe(\n",
    "    spots_df,\n",
    "    coordinate_cols=[\n",
    "        \"position_x\",\n",
    "        \"position_y\",\n",
    "    ],  # the column names for the coordinates\n",
    "    frame_col=\"frame\",  # the column name for the frame (default \"frame\")\n",
    "    only_coordinate_cols=False,  # if False, returned track_df includes columns not in coordinate_cols.\n",
    "    # False will be the default in the major release.\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`track_df` is the original dataframe with additional columns \"track_id\" and \"tree_id\".\n",
    "\n",
    "The track_id is a unique id for each track segments without branches. A new id is assigned when a splitting and merging occured. \n",
    "\n",
    "The tree_id is a unique id for each \"clonal\" tracks sharing the same ancestor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "keys = [\"position_x\", \"position_y\", \"track_id\", \"tree_id\"]\n",
    "display(track_df[keys].head())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`split_df` is a dataframe for splitting events with the following columns:\n",
    "- \"parent_track_id\" : the track id of the parent\n",
    "- \"child_track_id\" : the track id of the parent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(split_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`merge_df` is a dataframe for merging events with the following columns:\n",
    "- \"parent_track_id\" : the track id of the parent\n",
    "- \"child_track_id\" : the track id of the parent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(merge_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can display tracks in `napari` as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Does not work in Colab\n",
    "#\n",
    "# import napari\n",
    "# v = napari.Viewer()\n",
    "# v.add_points(spots_df[[\"frame\", \"position_x\", \"position_y\"]])\n",
    "# track_df2 = track_df.reset_index()\n",
    "# v.add_tracks(track_df2[[\"track_id\", \"frame\", \"position_x\", \"position_y\"]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plotting tracks in `matplotlib`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbsphinx-thumbnail": {
     "tooltip": "Simple example tracking results."
    }
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(3, 3))\n",
    "frames = track_df.index.get_level_values(\"frame\")\n",
    "frame_range = [frames.min(), frames.max()]\n",
    "k1, k2 = \"position_y\", \"position_x\"\n",
    "keys = [k1, k2]\n",
    "\n",
    "\n",
    "def get_track_end(track_id, first=True):\n",
    "    df = track_df[track_df[\"track_id\"] == track_id].sort_index(level=\"frame\")\n",
    "    return df.iloc[0 if first else -1][keys]\n",
    "\n",
    "\n",
    "for track_id, grp in track_df.groupby(\"track_id\"):\n",
    "    df = grp.reset_index().sort_values(\"frame\")\n",
    "    plt.scatter(df[k1], df[k2], c=df[\"frame\"], vmin=frame_range[0], vmax=frame_range[1])\n",
    "    for i in range(len(df) - 1):\n",
    "        pos1 = df.iloc[i][keys]\n",
    "        pos2 = df.iloc[i + 1][keys]\n",
    "        plt.plot([pos1[0], pos2[0]], [pos1[1], pos2[1]], \"-k\")\n",
    "    for _, row in list(split_df.iterrows()) + list(merge_df.iterrows()):\n",
    "        pos1 = get_track_end(row[\"parent_track_id\"], first=False)\n",
    "        pos2 = get_track_end(row[\"child_track_id\"], first=True)\n",
    "        plt.plot([pos1[0], pos2[0]], [pos1[1], pos2[1]], \"-k\")\n",
    "\n",
    "\n",
    "plt.xticks([])\n",
    "plt.yticks([])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example 2: using `predict` to track frame-wise-organized coordinates to make networkx tree"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`predict_dataframe` is a thin wrapper of the `predict` function, the core tracking function of the `LapTrack` object. \n",
    "\n",
    "One can directly use this function with the input of the frame-wise coordinate list and the output of networkx `DiGraph` object representing the lineage tree."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "frame_max = spots_df[\"frame\"].max()\n",
    "coords = []\n",
    "for i in range(frame_max):\n",
    "    df = spots_df[spots_df[\"frame\"] == i]\n",
    "    coords.append(df[[\"position_x\", \"position_y\"]].values)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The input variable `coords` should be organized as the frame-wise list of the point coordinates.\n",
    "\n",
    "The coordinate dimension is `(particle, dimension)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(coords[:5])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`predict` function generates a networkx `DiGraph` object from the coordinates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "track_tree = lt.predict(\n",
    "    coords,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The returned `track_tree` is the connection between spots, represented as ((frame1,index1), (frame2,index2)) ...\n",
    "\n",
    "For example, `((0, 0), (1, 1))` means the connection between `[178.41257464, 185.18866074]` and `[185.1758993 , 185.18866074]` in this example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for edge in list(track_tree.edges())[:5]:\n",
    "    print(edge)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.5 64-bit ('image_analysis')",
   "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.10.11"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "112c4654da1b8403c1d0efe0680f4afdb1a0816fe1d810c3856d3bd0e24d8e17"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}