pipeline_clean.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setting Up"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## load modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import itertools as itt\n",
    "import os\n",
    "import sys\n",
    "\n",
    "import holoviews as hv\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "from dask.distributed import Client, LocalCluster\n",
    "from holoviews.operation.datashader import datashade, regrid\n",
    "from holoviews.util import Dynamic\n",
    "from IPython.core.display import display"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## set path and parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Set up Initial Basic Parameters#\n",
    "minian_path = \".\"\n",
    "dpath = \"data/wf-RSC/Miniscope/\"\n",
    "minian_ds_path = os.path.join(dpath, \"minian\")\n",
    "intpath = \"./minian_intermediate\"\n",
    "subset = dict(frame=slice(0, None))\n",
    "subset_mc = None\n",
    "interactive = True\n",
    "output_size = 100\n",
    "n_workers = int(os.getenv(\"MINIAN_NWORKERS\", 4))\n",
    "param_save_minian = {\n",
    "    \"dpath\": minian_ds_path,\n",
    "    \"meta_dict\": dict(session=-1, animal=-2),\n",
    "    \"overwrite\": True,\n",
    "}\n",
    "\n",
    "# Pre-processing Parameters#\n",
    "param_load_videos = {\n",
    "    \"pattern\": \"[0-9]+\\.avi$\",\n",
    "    \"dtype\": np.uint8,\n",
    "    \"downsample\": dict(frame=1, height=1, width=1),\n",
    "    \"downsample_strategy\": \"subset\",\n",
    "}\n",
    "param_denoise = {\"method\": \"median\", \"ksize\": 7}\n",
    "param_background_removal = {\"method\": \"tophat\", \"wnd\": 15}\n",
    "\n",
    "# Motion Correction Parameters#\n",
    "subset_mc = None\n",
    "param_estimate_motion = {\"dim\": \"frame\"}\n",
    "\n",
    "# Initialization Parameters#\n",
    "param_seeds_init = {\n",
    "    \"wnd_size\": 1000,\n",
    "    \"method\": \"rolling\",\n",
    "    \"stp_size\": 500,\n",
    "    \"max_wnd\": 15,\n",
    "    \"diff_thres\": 3,\n",
    "}\n",
    "param_pnr_refine = {\"noise_freq\": 0.06, \"thres\": 1}\n",
    "param_ks_refine = {\"sig\": 0.05}\n",
    "param_seeds_merge = {\"thres_dist\": 10, \"thres_corr\": 0.8, \"noise_freq\": 0.06}\n",
    "param_initialize = {\"thres_corr\": 0.8, \"wnd\": 10, \"noise_freq\": 0.06}\n",
    "param_init_merge = {\"thres_corr\": 0.8}\n",
    "\n",
    "# CNMF Parameters#\n",
    "param_get_noise = {\"noise_range\": (0.06, 0.5)}\n",
    "param_first_spatial = {\n",
    "    \"dl_wnd\": 10,\n",
    "    \"sparse_penal\": 0.01,\n",
    "    \"size_thres\": (25, None),\n",
    "}\n",
    "param_first_temporal = {\n",
    "    \"noise_freq\": 0.06,\n",
    "    \"sparse_penal\": 1,\n",
    "    \"p\": 1,\n",
    "    \"add_lag\": 20,\n",
    "    \"jac_thres\": 0.2,\n",
    "}\n",
    "param_first_merge = {\"thres_corr\": 0.8}\n",
    "param_second_spatial = {\n",
    "    \"dl_wnd\": 10,\n",
    "    \"sparse_penal\": 0.01,\n",
    "    \"size_thres\": (25, None),\n",
    "}\n",
    "param_second_temporal = {\n",
    "    \"noise_freq\": 0.06,\n",
    "    \"sparse_penal\": 1,\n",
    "    \"p\": 1,\n",
    "    \"add_lag\": 20,\n",
    "    \"jac_thres\": 0.4,\n",
    "}\n",
    "\n",
    "os.environ[\"OMP_NUM_THREADS\"] = \"1\"\n",
    "os.environ[\"MKL_NUM_THREADS\"] = \"1\"\n",
    "os.environ[\"OPENBLAS_NUM_THREADS\"] = \"1\"\n",
    "os.environ[\"MINIAN_INTERMEDIATE\"] = intpath"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## import minian"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.tags": "worksheet-0",
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "sys.path.append(minian_path)\n",
    "from minian.cnmf import (\n",
    "    compute_AtC,\n",
    "    compute_trace,\n",
    "    get_noise_fft,\n",
    "    smooth_sig,\n",
    "    unit_merge,\n",
    "    update_spatial,\n",
    "    update_temporal,\n",
    "    update_background,\n",
    ")\n",
    "from minian.initialization import (\n",
    "    gmm_refine,\n",
    "    initA,\n",
    "    initC,\n",
    "    intensity_refine,\n",
    "    ks_refine,\n",
    "    pnr_refine,\n",
    "    seeds_init,\n",
    "    seeds_merge,\n",
    ")\n",
    "from minian.motion_correction import apply_transform, estimate_motion\n",
    "from minian.preprocessing import denoise, remove_background\n",
    "from minian.utilities import (\n",
    "    TaskAnnotation,\n",
    "    get_optimal_chk,\n",
    "    load_videos,\n",
    "    open_minian,\n",
    "    save_minian,\n",
    ")\n",
    "from minian.visualization import (\n",
    "    CNMFViewer,\n",
    "    VArrayViewer,\n",
    "    generate_videos,\n",
    "    visualize_gmm_fit,\n",
    "    visualize_motion,\n",
    "    visualize_preprocess,\n",
    "    visualize_seeds,\n",
    "    visualize_spatial_update,\n",
    "    visualize_temporal_update,\n",
    "    write_video,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## module initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.tags": "worksheet-0",
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "dpath = os.path.abspath(dpath)\n",
    "hv.notebook_extension(\"bokeh\", width=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## start cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster = LocalCluster(\n",
    "    n_workers=n_workers,\n",
    "    memory_limit=\"2GB\",\n",
    "    resources={\"MEM\": 1},\n",
    "    threads_per_worker=2,\n",
    "    dashboard_address=\":8787\",\n",
    ")\n",
    "annt_plugin = TaskAnnotation()\n",
    "cluster.scheduler.add_plugin(annt_plugin)\n",
    "client = Client(cluster)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pre-processing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## loading videos and visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.tags": "worksheet-0",
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "varr = load_videos(dpath, **param_load_videos)\n",
    "chk, _ = get_optimal_chk(varr, dtype=float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "varr = save_minian(\n",
    "    varr.chunk({\"frame\": chk[\"frame\"], \"height\": -1, \"width\": -1}).rename(\"varr\"),\n",
    "    intpath,\n",
    "    overwrite=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## visualize raw data and optionally set roi for motion correction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.output(size=output_size)\n",
    "if interactive:\n",
    "    vaviewer = VArrayViewer(varr, framerate=5, summary=[\"mean\", \"max\"])\n",
    "    display(vaviewer.show())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if interactive:\n",
    "    try:\n",
    "        subset_mc = list(vaviewer.mask.values())[0]\n",
    "    except IndexError:\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subset = {'frame': slice(10, None)}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## subset part of video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "varr_ref = varr.sel(subset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## glow removal and visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "varr_min = varr_ref.min(\"frame\").compute()\n",
    "varr_ref = varr_ref - varr_min"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.Image(varr_ref.max('frame').compute())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.Curve(varr_ref.max(['height', 'width']).compute()).opts(frame_width=1300, tools=['hover'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.Image(varr_ref.sel(frame=12430))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.output(size=int(output_size * 0.7))\n",
    "if interactive:\n",
    "    vaviewer = VArrayViewer(\n",
    "        [varr.rename(\"original\"), varr_ref.rename(\"glow_removed\")],\n",
    "        framerate=5,\n",
    "        summary=None,\n",
    "        layout=True,\n",
    "    )\n",
    "    display(vaviewer.show())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## denoise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.output(size=int(output_size * 0.6))\n",
    "if interactive:\n",
    "    display(\n",
    "        visualize_preprocess(\n",
    "            varr_ref.isel(frame=0).compute(),\n",
    "            denoise,\n",
    "            method=[\"median\"],\n",
    "            ksize=[5, 7, 9],\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following cell would carry out denoise step.\n",
    "Be sure to [change the parameters](https://minian.readthedocs.io/page/start_guide/faq.html#i-don-t-know-python-can-i-still-use-the-pipeline) based on visualization results before running the following cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "varr_ref = denoise(varr_ref, **param_denoise)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## background removal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.output(size=int(output_size * 0.6))\n",
    "if interactive:\n",
    "    display(\n",
    "        visualize_preprocess(\n",
    "            varr_ref.isel(frame=0).compute(),\n",
    "            remove_background,\n",
    "            method=[\"tophat\"],\n",
    "            wnd=[10, 15, 20],\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following cell would carry out background removal step.\n",
    "Be sure to [change the parameters](https://minian.readthedocs.io/page/start_guide/faq.html#i-don-t-know-python-can-i-still-use-the-pipeline) based on visualization results before running the following cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "varr_ref = remove_background(varr_ref, **param_background_removal)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## save result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "varr_ref = save_minian(varr_ref.rename(\"varr_ref\"), dpath=intpath, overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Motion Correction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## estimate motion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "motion = estimate_motion(varr_ref.sel(subset_mc), **param_estimate_motion)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## save motion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "motion = save_minian(\n",
    "    motion.rename(\"motion\").chunk({\"frame\": chk[\"frame\"]}), **param_save_minian\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## visualization of motion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.output(size=output_size)\n",
    "visualize_motion(motion)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## apply transform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "Y = apply_transform(varr_ref, motion, fill=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## save result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "Y_fm_chk = save_minian(Y.astype(float).rename(\"Y_fm_chk\"), intpath, overwrite=True)\n",
    "Y_hw_chk = save_minian(\n",
    "    Y_fm_chk.rename(\"Y_hw_chk\"),\n",
    "    intpath,\n",
    "    overwrite=True,\n",
    "    chunks={\"frame\": -1, \"height\": chk[\"height\"], \"width\": chk[\"width\"]},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## visualization of motion-correction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.output(size=int(output_size * 0.7))\n",
    "if interactive:\n",
    "    vaviewer = VArrayViewer(\n",
    "        [varr_ref.rename(\"before_mc\"), Y_fm_chk.rename(\"after_mc\")],\n",
    "        framerate=5,\n",
    "        summary=None,\n",
    "        layout=True,\n",
    "    )\n",
    "    display(vaviewer.show())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "im_opts = dict(\n",
    "    frame_width=500,\n",
    "    aspect=varr_ref.sizes[\"width\"] / varr_ref.sizes[\"height\"],\n",
    "    cmap=\"Viridis\",\n",
    "    colorbar=True,\n",
    ")\n",
    "(\n",
    "    regrid(\n",
    "        hv.Image(\n",
    "            varr_ref.max(\"frame\").compute().astype(np.float32),\n",
    "            [\"width\", \"height\"],\n",
    "            label=\"before_mc\",\n",
    "        ).opts(**im_opts)\n",
    "    )\n",
    "    + regrid(\n",
    "        hv.Image(\n",
    "            Y_hw_chk.max(\"frame\").compute().astype(np.float32),\n",
    "            [\"width\", \"height\"],\n",
    "            label=\"after_mc\",\n",
    "        ).opts(**im_opts)\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## generate video for motion-correction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "vid_arr = xr.concat([varr_ref, Y_fm_chk], \"width\").chunk({\"width\": -1})\n",
    "write_video(vid_arr, \"minian_mc.mp4\", dpath)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Initialization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## compute max projection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_proj = save_minian(\n",
    "    Y_fm_chk.max(\"frame\").rename(\"max_proj\"), **param_save_minian\n",
    ").compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## generating over-complete set of seeds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "seeds = seeds_init(Y_fm_chk, **param_seeds_init)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.output(size=output_size)\n",
    "visualize_seeds(max_proj, seeds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## peak-noise-ratio refine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "if interactive:\n",
    "    noise_freq_list = [0.005, 0.01, 0.02, 0.06, 0.1, 0.2, 0.3, 0.45, 0.6, 0.8]\n",
    "    example_seeds = seeds.sample(6, axis=\"rows\")\n",
    "    example_trace = Y_hw_chk.sel(\n",
    "        height=example_seeds[\"height\"].to_xarray(),\n",
    "        width=example_seeds[\"width\"].to_xarray(),\n",
    "    ).rename(**{\"index\": \"seed\"})\n",
    "    smooth_dict = dict()\n",
    "    for freq in noise_freq_list:\n",
    "        trace_smth_low = smooth_sig(example_trace, freq)\n",
    "        trace_smth_high = smooth_sig(example_trace, freq, btype=\"high\")\n",
    "        trace_smth_low = trace_smth_low.compute()\n",
    "        trace_smth_high = trace_smth_high.compute()\n",
    "        hv_trace = hv.HoloMap(\n",
    "            {\n",
    "                \"signal\": (\n",
    "                    hv.Dataset(trace_smth_low)\n",
    "                    .to(hv.Curve, kdims=[\"frame\"])\n",
    "                    .opts(frame_width=300, aspect=2, ylabel=\"Signal (A.U.)\")\n",
    "                ),\n",
    "                \"noise\": (\n",
    "                    hv.Dataset(trace_smth_high)\n",
    "                    .to(hv.Curve, kdims=[\"frame\"])\n",
    "                    .opts(frame_width=300, aspect=2, ylabel=\"Signal (A.U.)\")\n",
    "                ),\n",
    "            },\n",
    "            kdims=\"trace\",\n",
    "        ).collate()\n",
    "        smooth_dict[freq] = hv_trace"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.output(size=int(output_size * 0.7))\n",
    "if interactive:\n",
    "    hv_res = (\n",
    "        hv.HoloMap(smooth_dict, kdims=[\"noise_freq\"])\n",
    "        .collate()\n",
    "        .opts(aspect=2)\n",
    "        .overlay(\"trace\")\n",
    "        .layout(\"seed\")\n",
    "        .cols(3)\n",
    "    )\n",
    "    display(hv_res)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following cell would carry out pnr refinement step.\n",
    "Be sure to [change](https://minian.readthedocs.io/page/start_guide/faq.html#i-don-t-know-python-can-i-still-use-the-pipeline) the `noise_freq` based on visualization results before running the following cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "seeds, pnr, gmm = pnr_refine(Y_hw_chk, seeds, **param_pnr_refine)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if gmm:\n",
    "    display(visualize_gmm_fit(pnr, gmm, 100))\n",
    "else:\n",
    "    print(\"nothing to show\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.output(size=output_size)\n",
    "visualize_seeds(max_proj, seeds, \"mask_pnr\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ks refine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "seeds = ks_refine(Y_hw_chk, seeds, **param_ks_refine)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.output(size=output_size)\n",
    "visualize_seeds(max_proj, seeds, \"mask_ks\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## merge seeds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "seeds_final = seeds[seeds[\"mask_ks\"] & seeds[\"mask_pnr\"]].reset_index(drop=True)\n",
    "seeds_final = seeds_merge(Y_hw_chk, max_proj, seeds_final, **param_seeds_merge)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.output(size=output_size)\n",
    "visualize_seeds(max_proj, seeds_final, \"mask_mrg\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## initialize spatial matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "A_init = initA(Y_hw_chk, seeds_final[seeds_final[\"mask_mrg\"]], **param_initialize)\n",
    "A_init = save_minian(A_init.rename(\"A_init\"), intpath, overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## initialize temporal matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "C_init = initC(Y_fm_chk, A_init)\n",
    "C_init = save_minian(\n",
    "    C_init.rename(\"C_init\"), intpath, overwrite=True, chunks={\"unit_id\": 1, \"frame\": -1}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## merge units"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "A, C = unit_merge(A_init, C_init, **param_init_merge)\n",
    "A = save_minian(A.rename(\"A\"), intpath, overwrite=True)\n",
    "C = save_minian(C.rename(\"C\"), intpath, overwrite=True)\n",
    "C_chk = save_minian(\n",
    "    C.rename(\"C_chk\"),\n",
    "    intpath,\n",
    "    overwrite=True,\n",
    "    chunks={\"unit_id\": -1, \"frame\": chk[\"frame\"]},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## initialize background terms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "b, f = update_background(Y_fm_chk, A, C_chk)\n",
    "f = save_minian(f.rename(\"f\"), intpath, overwrite=True)\n",
    "b = save_minian(b.rename(\"b\"), intpath, overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## visualization of initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.output(size=int(output_size * 0.55))\n",
    "im_opts = dict(\n",
    "    frame_width=500,\n",
    "    aspect=A.sizes[\"width\"] / A.sizes[\"height\"],\n",
    "    cmap=\"Viridis\",\n",
    "    colorbar=True,\n",
    ")\n",
    "cr_opts = dict(frame_width=750, aspect=1.5 * A.sizes[\"width\"] / A.sizes[\"height\"])\n",
    "(\n",
    "    regrid(\n",
    "        hv.Image(\n",
    "            A.max(\"unit_id\").rename(\"A\").compute().astype(np.float32),\n",
    "            kdims=[\"width\", \"height\"],\n",
    "        ).opts(**im_opts)\n",
    "    ).relabel(\"Initial Spatial Footprints\")\n",
    "    + regrid(\n",
    "        hv.Image(\n",
    "            C.rename(\"C\").compute().astype(np.float32), kdims=[\"frame\", \"unit_id\"]\n",
    "        ).opts(cmap=\"viridis\", colorbar=True, **cr_opts)\n",
    "    ).relabel(\"Initial Temporal Components\")\n",
    "    + regrid(\n",
    "        hv.Image(\n",
    "            b.rename(\"b\").compute().astype(np.float32), kdims=[\"width\", \"height\"]\n",
    "        ).opts(**im_opts)\n",
    "    ).relabel(\"Initial Background Sptial\")\n",
    "    + datashade(hv.Curve(f.rename(\"f\").compute(), kdims=[\"frame\"]), min_alpha=200)\n",
    "    .opts(**cr_opts)\n",
    "    .relabel(\"Initial Background Temporal\")\n",
    ").cols(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# CNMF"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## estimate spatial noise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "sn_spatial = get_noise_fft(Y_hw_chk, **param_get_noise)\n",
    "sn_spatial = save_minian(sn_spatial.rename(\"sn_spatial\"), intpath, overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## first spatial update"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### parameter exploration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if interactive:\n",
    "    units = np.random.choice(A.coords[\"unit_id\"], 10, replace=False)\n",
    "    units.sort()\n",
    "    A_sub = A.sel(unit_id=units).persist()\n",
    "    C_sub = C.sel(unit_id=units).persist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "if interactive:\n",
    "    sprs_ls = [0.005, 0.01, 0.05]\n",
    "    A_dict = dict()\n",
    "    C_dict = dict()\n",
    "    for cur_sprs in sprs_ls:\n",
    "        cur_A, cur_mask, cur_norm = update_spatial(\n",
    "            Y_hw_chk,\n",
    "            A_sub,\n",
    "            C_sub,\n",
    "            sn_spatial,\n",
    "            in_memory=True,\n",
    "            dl_wnd=param_first_spatial[\"dl_wnd\"],\n",
    "            sparse_penal=cur_sprs,\n",
    "        )\n",
    "        if cur_A.sizes[\"unit_id\"]:\n",
    "            A_dict[cur_sprs] = cur_A.compute()\n",
    "            C_dict[cur_sprs] = C_sub.sel(unit_id=cur_mask).compute()\n",
    "    hv_res = visualize_spatial_update(A_dict, C_dict, kdims=[\"sparse penalty\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.output(size=int(output_size * 0.6))\n",
    "if interactive:\n",
    "    display(hv_res)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### spatial update"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following cell would carry out spatial update.\n",
    "Be sure to [change the parameters](https://minian.readthedocs.io/page/start_guide/faq.html#i-don-t-know-python-can-i-still-use-the-pipeline) based on visualization results before running the following cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "A_new, mask, norm_fac = update_spatial(\n",
    "    Y_hw_chk, A, C, sn_spatial, **param_first_spatial\n",
    ")\n",
    "C_new = save_minian(\n",
    "    (C.sel(unit_id=mask) * norm_fac).rename(\"C_new\"), intpath, overwrite=True\n",
    ")\n",
    "C_chk_new = save_minian(\n",
    "    (C_chk.sel(unit_id=mask) * norm_fac).rename(\"C_chk_new\"), intpath, overwrite=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After obtaining a better estimation of spatial footprints, we update the background terms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "b_new, f_new = update_background(Y_fm_chk, A_new, C_chk_new)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### visualization of spatial footprints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.output(size=int(output_size * 0.6))\n",
    "opts = dict(\n",
    "    plot=dict(height=A.sizes[\"height\"], width=A.sizes[\"width\"], colorbar=True),\n",
    "    style=dict(cmap=\"Viridis\"),\n",
    ")\n",
    "(\n",
    "    regrid(\n",
    "        hv.Image(\n",
    "            A.max(\"unit_id\").compute().astype(np.float32).rename(\"A\"),\n",
    "            kdims=[\"width\", \"height\"],\n",
    "        ).opts(**opts)\n",
    "    ).relabel(\"Spatial Footprints Initial\")\n",
    "    + regrid(\n",
    "        hv.Image(\n",
    "            (A.fillna(0) > 0).sum(\"unit_id\").compute().astype(np.uint8).rename(\"A\"),\n",
    "            kdims=[\"width\", \"height\"],\n",
    "        ).opts(**opts)\n",
    "    ).relabel(\"Binary Spatial Footprints Initial\")\n",
    "    + regrid(\n",
    "        hv.Image(\n",
    "            A_new.max(\"unit_id\").compute().astype(np.float32).rename(\"A\"),\n",
    "            kdims=[\"width\", \"height\"],\n",
    "        ).opts(**opts)\n",
    "    ).relabel(\"Spatial Footprints First Update\")\n",
    "    + regrid(\n",
    "        hv.Image(\n",
    "            (A_new > 0).sum(\"unit_id\").compute().astype(np.uint8).rename(\"A\"),\n",
    "            kdims=[\"width\", \"height\"],\n",
    "        ).opts(**opts)\n",
    "    ).relabel(\"Binary Spatial Footprints First Update\")\n",
    ").cols(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### visualization of background"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.output(size=int(output_size * 0.55))\n",
    "opts_im = dict(\n",
    "    plot=dict(height=b.sizes[\"height\"], width=b.sizes[\"width\"], colorbar=True),\n",
    "    style=dict(cmap=\"Viridis\"),\n",
    ")\n",
    "opts_cr = dict(plot=dict(height=b.sizes[\"height\"], width=b.sizes[\"height\"] * 2))\n",
    "(\n",
    "    regrid(\n",
    "        hv.Image(b.compute().astype(np.float32), kdims=[\"width\", \"height\"]).opts(\n",
    "            **opts_im\n",
    "        )\n",
    "    ).relabel(\"Background Spatial Initial\")\n",
    "    + hv.Curve(f.compute().rename(\"f\").astype(np.float16), kdims=[\"frame\"])\n",
    "    .opts(**opts_cr)\n",
    "    .relabel(\"Background Temporal Initial\")\n",
    "    + regrid(\n",
    "        hv.Image(b_new.compute().astype(np.float32), kdims=[\"width\", \"height\"]).opts(\n",
    "            **opts_im\n",
    "        )\n",
    "    ).relabel(\"Background Spatial First Update\")\n",
    "    + hv.Curve(f_new.compute().rename(\"f\").astype(np.float16), kdims=[\"frame\"])\n",
    "    .opts(**opts_cr)\n",
    "    .relabel(\"Background Temporal First Update\")\n",
    ").cols(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### save results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "A = save_minian(\n",
    "    A_new.rename(\"A\"),\n",
    "    intpath,\n",
    "    overwrite=True,\n",
    "    chunks={\"unit_id\": 1, \"height\": -1, \"width\": -1},\n",
    ")\n",
    "b = save_minian(b_new.rename(\"b\"), intpath, overwrite=True)\n",
    "f = save_minian(\n",
    "    f_new.chunk({\"frame\": chk[\"frame\"]}).rename(\"f\"), intpath, overwrite=True\n",
    ")\n",
    "C = save_minian(C_new.rename(\"C\"), intpath, overwrite=True)\n",
    "C_chk = save_minian(C_chk_new.rename(\"C_chk\"), intpath, overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## first temporal update"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### parameter exploration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if interactive:\n",
    "    units = np.random.choice(A.coords[\"unit_id\"], 10, replace=False)\n",
    "    units.sort()\n",
    "    A_sub = A.sel(unit_id=units).persist()\n",
    "    C_sub = C_chk.sel(unit_id=units).persist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "if interactive:\n",
    "    p_ls = [1]\n",
    "    sprs_ls = [0.1, 0.5, 1, 2]\n",
    "    add_ls = [20]\n",
    "    noise_ls = [0.06]\n",
    "    YA_dict, C_dict, S_dict, g_dict, sig_dict, A_dict = [dict() for _ in range(6)]\n",
    "    YrA = (\n",
    "        compute_trace(Y_fm_chk, A_sub, b, C_sub, f)\n",
    "        .persist()\n",
    "        .chunk({\"unit_id\": 1, \"frame\": -1})\n",
    "    )\n",
    "    for cur_p, cur_sprs, cur_add, cur_noise in itt.product(\n",
    "        p_ls, sprs_ls, add_ls, noise_ls\n",
    "    ):\n",
    "        ks = (cur_p, cur_sprs, cur_add, cur_noise)\n",
    "        print(\n",
    "            \"p:{}, sparse penalty:{}, additional lag:{}, noise frequency:{}\".format(\n",
    "                cur_p, cur_sprs, cur_add, cur_noise\n",
    "            )\n",
    "        )\n",
    "        cur_C, cur_S, cur_b0, cur_c0, cur_g, cur_mask = update_temporal(\n",
    "            A_sub,\n",
    "            C_sub,\n",
    "            YrA=YrA,\n",
    "            sparse_penal=cur_sprs,\n",
    "            p=cur_p,\n",
    "            use_smooth=True,\n",
    "            add_lag=cur_add,\n",
    "            noise_freq=cur_noise,\n",
    "        )\n",
    "        YA_dict[ks], C_dict[ks], S_dict[ks], g_dict[ks], sig_dict[ks], A_dict[ks] = (\n",
    "            YrA.compute(),\n",
    "            cur_C.compute(),\n",
    "            cur_S.compute(),\n",
    "            cur_g.compute(),\n",
    "            (cur_C + cur_b0 + cur_c0).compute(),\n",
    "            A_sub.compute(),\n",
    "        )\n",
    "    hv_res = visualize_temporal_update(\n",
    "        YA_dict,\n",
    "        C_dict,\n",
    "        S_dict,\n",
    "        g_dict,\n",
    "        sig_dict,\n",
    "        A_dict,\n",
    "        kdims=[\"p\", \"sparse penalty\", \"additional lag\", \"noise frequency\"],\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.output(size=int(output_size * 0.6))\n",
    "if interactive:\n",
    "    display(hv_res)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### temporal update"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "YrA = save_minian(\n",
    "    compute_trace(Y_fm_chk, A, b, C_chk, f).rename(\"YrA\"),\n",
    "    intpath,\n",
    "    overwrite=True,\n",
    "    chunks={\"unit_id\": 1, \"frame\": -1},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following cell would carry out temporal update.\n",
    "Be sure to [change the parameters](https://minian.readthedocs.io/page/start_guide/faq.html#i-don-t-know-python-can-i-still-use-the-pipeline) based on visualization results before running the following cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "C_new, S_new, b0_new, c0_new, g, mask = update_temporal(\n",
    "    A, C, YrA=YrA, **param_first_temporal\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### visualization of temporal components"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.output(size=int(output_size * 0.6))\n",
    "opts_im = dict(frame_width=500, aspect=2, colorbar=True, cmap=\"Viridis\")\n",
    "(\n",
    "    regrid(\n",
    "        hv.Image(\n",
    "            C.compute().astype(np.float32).rename(\"ci\"), kdims=[\"frame\", \"unit_id\"]\n",
    "        ).opts(**opts_im)\n",
    "    ).relabel(\"Temporal Trace Initial\")\n",
    "    + hv.Div(\"\")\n",
    "    + regrid(\n",
    "        hv.Image(\n",
    "            C_new.compute().astype(np.float32).rename(\"c1\"), kdims=[\"frame\", \"unit_id\"]\n",
    "        ).opts(**opts_im)\n",
    "    ).relabel(\"Temporal Trace First Update\")\n",
    "    + regrid(\n",
    "        hv.Image(\n",
    "            S_new.compute().astype(np.float32).rename(\"s1\"), kdims=[\"frame\", \"unit_id\"]\n",
    "        ).opts(**opts_im)\n",
    "    ).relabel(\"Spikes First Update\")\n",
    ").cols(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### visualization of dropped units"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.output(size=int(output_size * 0.6))\n",
    "if interactive:\n",
    "    h, w = A.sizes[\"height\"], A.sizes[\"width\"]\n",
    "    im_opts = dict(aspect=w / h, frame_width=500, cmap=\"Viridis\")\n",
    "    cr_opts = dict(aspect=3, frame_width=1000)\n",
    "    bad_units = mask.where(mask == False, drop=True).coords[\"unit_id\"].values\n",
    "    if len(bad_units) > 0:\n",
    "        hv_res = (\n",
    "            hv.NdLayout(\n",
    "                {\n",
    "                    \"Spatial Footprint\": Dynamic(\n",
    "                        hv.Dataset(A.sel(unit_id=bad_units).compute().rename(\"A\"))\n",
    "                        .to(hv.Image, kdims=[\"width\", \"height\"])\n",
    "                        .opts(**im_opts)\n",
    "                    ),\n",
    "                    \"Spatial Footprints of Accepted Units\": Dynamic(\n",
    "                        hv.Image(\n",
    "                            A.sel(unit_id=mask).sum(\"unit_id\").compute().rename(\"A\"),\n",
    "                            kdims=[\"width\", \"height\"],\n",
    "                        ).opts(**im_opts)\n",
    "                    ),\n",
    "                }\n",
    "            )\n",
    "            + datashade(\n",
    "                hv.Dataset(YrA.sel(unit_id=bad_units).rename(\"raw\")).to(\n",
    "                    hv.Curve, kdims=[\"frame\"]\n",
    "                )\n",
    "            )\n",
    "            .opts(**cr_opts)\n",
    "            .relabel(\"Temporal Trace\")\n",
    "        ).cols(1)\n",
    "        display(hv_res)\n",
    "    else:\n",
    "        print(\"No rejected units to display\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### visualization of accepted units"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.output(size=int(output_size * 0.6))\n",
    "if interactive:\n",
    "    sig = C_new + b0_new + c0_new\n",
    "    display(\n",
    "        visualize_temporal_update(\n",
    "            YrA.sel(unit_id=mask),\n",
    "            C_new,\n",
    "            S_new,\n",
    "            g,\n",
    "            sig,\n",
    "            A.sel(unit_id=mask),\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### save results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "C = save_minian(\n",
    "    C_new.rename(\"C\").chunk({\"unit_id\": 1, \"frame\": -1}), intpath, overwrite=True\n",
    ")\n",
    "C_chk = save_minian(\n",
    "    C.rename(\"C_chk\"),\n",
    "    intpath,\n",
    "    overwrite=True,\n",
    "    chunks={\"unit_id\": -1, \"frame\": chk[\"frame\"]},\n",
    ")\n",
    "S = save_minian(\n",
    "    S_new.rename(\"S\").chunk({\"unit_id\": 1, \"frame\": -1}), intpath, overwrite=True\n",
    ")\n",
    "b0 = save_minian(\n",
    "    b0_new.rename(\"b0\").chunk({\"unit_id\": 1, \"frame\": -1}), intpath, overwrite=True\n",
    ")\n",
    "c0 = save_minian(\n",
    "    c0_new.rename(\"c0\").chunk({\"unit_id\": 1, \"frame\": -1}), intpath, overwrite=True\n",
    ")\n",
    "A = A.sel(unit_id=C.coords[\"unit_id\"].values)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## merge units"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "A_mrg, C_mrg, [sig_mrg] = unit_merge(A, C, [C + b0 + c0], **param_first_merge)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.output(size=int(output_size * 0.6))\n",
    "opts_im = dict(frame_width=500, aspect=2, colorbar=True, cmap=\"Viridis\")\n",
    "(\n",
    "    regrid(\n",
    "        hv.Image(\n",
    "            C.compute().astype(np.float32).rename(\"c1\"), kdims=[\"frame\", \"unit_id\"]\n",
    "        )\n",
    "        .relabel(\"Temporal Signals Before Merge\")\n",
    "        .opts(**opts_im)\n",
    "    )\n",
    "    + regrid(\n",
    "        hv.Image(\n",
    "            C_mrg.compute().astype(np.float32).rename(\"c2\"), kdims=[\"frame\", \"unit_id\"]\n",
    "        )\n",
    "        .relabel(\"Temporal Signals After Merge\")\n",
    "        .opts(**opts_im)\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "A = save_minian(A_mrg.rename(\"A_mrg\"), intpath, overwrite=True)\n",
    "C = save_minian(C_mrg.rename(\"C_mrg\"), intpath, overwrite=True)\n",
    "C_chk = save_minian(\n",
    "    C.rename(\"C_mrg_chk\"),\n",
    "    intpath,\n",
    "    overwrite=True,\n",
    "    chunks={\"unit_id\": -1, \"frame\": chk[\"frame\"]},\n",
    ")\n",
    "sig = save_minian(sig_mrg.rename(\"sig_mrg\"), intpath, overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## second spatial update"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### parameter exploration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if interactive:\n",
    "    units = np.random.choice(A.coords[\"unit_id\"], 10, replace=False)\n",
    "    units.sort()\n",
    "    A_sub = A.sel(unit_id=units).persist()\n",
    "    C_sub = sig.sel(unit_id=units).persist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "if interactive:\n",
    "    sprs_ls = [5e-3, 1e-2, 5e-2]\n",
    "    A_dict = dict()\n",
    "    C_dict = dict()\n",
    "    for cur_sprs in sprs_ls:\n",
    "        cur_A, cur_mask, cur_norm = update_spatial(\n",
    "            Y_hw_chk,\n",
    "            A_sub,\n",
    "            C_sub,\n",
    "            sn_spatial,\n",
    "            in_memory=True,\n",
    "            dl_wnd=param_second_spatial[\"dl_wnd\"],\n",
    "            sparse_penal=cur_sprs,\n",
    "        )\n",
    "        if cur_A.sizes[\"unit_id\"]:\n",
    "            A_dict[cur_sprs] = cur_A.compute()\n",
    "            C_dict[cur_sprs] = C_sub.sel(unit_id=cur_mask).compute()\n",
    "    hv_res = visualize_spatial_update(A_dict, C_dict, kdims=[\"sparse penalty\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.output(size=int(output_size * 0.6))\n",
    "if interactive:\n",
    "    display(hv_res)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### spatial update"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "A_new, mask, norm_fac = update_spatial(\n",
    "    Y_hw_chk, A, C, sn_spatial, **param_second_spatial\n",
    ")\n",
    "C_new = save_minian(\n",
    "    (C.sel(unit_id=mask) * norm_fac).rename(\"C_new\"), intpath, overwrite=True\n",
    ")\n",
    "C_chk_new = save_minian(\n",
    "    (C_chk.sel(unit_id=mask) * norm_fac).rename(\"C_chk_new\"), intpath, overwrite=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "b_new, f_new = update_background(Y_fm_chk, A_new, C_chk_new)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### visualization of spatial footprints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.output(size=int(output_size * 0.6))\n",
    "opts = dict(\n",
    "    plot=dict(height=A.sizes[\"height\"], width=A.sizes[\"width\"], colorbar=True),\n",
    "    style=dict(cmap=\"Viridis\"),\n",
    ")\n",
    "(\n",
    "    regrid(\n",
    "        hv.Image(\n",
    "            A.max(\"unit_id\").compute().astype(np.float32).rename(\"A\"),\n",
    "            kdims=[\"width\", \"height\"],\n",
    "        ).opts(**opts)\n",
    "    ).relabel(\"Spatial Footprints Last\")\n",
    "    + regrid(\n",
    "        hv.Image(\n",
    "            (A.fillna(0) > 0).sum(\"unit_id\").compute().astype(np.uint8).rename(\"A\"),\n",
    "            kdims=[\"width\", \"height\"],\n",
    "        ).opts(**opts)\n",
    "    ).relabel(\"Binary Spatial Footprints Last\")\n",
    "    + regrid(\n",
    "        hv.Image(\n",
    "            A_new.max(\"unit_id\").compute().astype(np.float32).rename(\"A\"),\n",
    "            kdims=[\"width\", \"height\"],\n",
    "        ).opts(**opts)\n",
    "    ).relabel(\"Spatial Footprints New\")\n",
    "    + regrid(\n",
    "        hv.Image(\n",
    "            (A_new > 0).sum(\"unit_id\").compute().astype(np.uint8).rename(\"A\"),\n",
    "            kdims=[\"width\", \"height\"],\n",
    "        ).opts(**opts)\n",
    "    ).relabel(\"Binary Spatial Footprints New\")\n",
    ").cols(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### visualization of background"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.output(size=int(output_size * 0.55))\n",
    "opts_im = dict(\n",
    "    plot=dict(height=b.sizes[\"height\"], width=b.sizes[\"width\"], colorbar=True),\n",
    "    style=dict(cmap=\"Viridis\"),\n",
    ")\n",
    "opts_cr = dict(plot=dict(height=b.sizes[\"height\"], width=b.sizes[\"height\"] * 2))\n",
    "(\n",
    "    regrid(\n",
    "        hv.Image(b.compute().astype(np.float32), kdims=[\"width\", \"height\"]).opts(\n",
    "            **opts_im\n",
    "        )\n",
    "    ).relabel(\"Background Spatial Last\")\n",
    "    + hv.Curve(f.compute().rename(\"f\").astype(np.float16), kdims=[\"frame\"])\n",
    "    .opts(**opts_cr)\n",
    "    .relabel(\"Background Temporal Last\")\n",
    "    + regrid(\n",
    "        hv.Image(b_new.compute().astype(np.float32), kdims=[\"width\", \"height\"]).opts(\n",
    "            **opts_im\n",
    "        )\n",
    "    ).relabel(\"Background Spatial New\")\n",
    "    + hv.Curve(f_new.compute().rename(\"f\").astype(np.float16), kdims=[\"frame\"])\n",
    "    .opts(**opts_cr)\n",
    "    .relabel(\"Background Temporal New\")\n",
    ").cols(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### save results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "A = save_minian(\n",
    "    A_new.rename(\"A\"),\n",
    "    intpath,\n",
    "    overwrite=True,\n",
    "    chunks={\"unit_id\": 1, \"height\": -1, \"width\": -1},\n",
    ")\n",
    "b = save_minian(b_new.rename(\"b\"), intpath, overwrite=True)\n",
    "f = save_minian(\n",
    "    f_new.chunk({\"frame\": chk[\"frame\"]}).rename(\"f\"), intpath, overwrite=True\n",
    ")\n",
    "C = save_minian(C_new.rename(\"C\"), intpath, overwrite=True)\n",
    "C_chk = save_minian(C_chk_new.rename(\"C_chk\"), intpath, overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## second temporal update"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### parameter exploration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if interactive:\n",
    "    units = np.random.choice(A.coords[\"unit_id\"], 10, replace=False)\n",
    "    units.sort()\n",
    "    A_sub = A.sel(unit_id=units).persist()\n",
    "    C_sub = C_chk.sel(unit_id=units).persist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "if interactive:\n",
    "    p_ls = [1]\n",
    "    sprs_ls = [0.1, 0.5, 1, 2]\n",
    "    add_ls = [20]\n",
    "    noise_ls = [0.06]\n",
    "    YA_dict, C_dict, S_dict, g_dict, sig_dict, A_dict = [dict() for _ in range(6)]\n",
    "    YrA = (\n",
    "        compute_trace(Y_fm_chk, A_sub, b, C_sub, f)\n",
    "        .persist()\n",
    "        .chunk({\"unit_id\": 1, \"frame\": -1})\n",
    "    )\n",
    "    for cur_p, cur_sprs, cur_add, cur_noise in itt.product(\n",
    "        p_ls, sprs_ls, add_ls, noise_ls\n",
    "    ):\n",
    "        ks = (cur_p, cur_sprs, cur_add, cur_noise)\n",
    "        print(\n",
    "            \"p:{}, sparse penalty:{}, additional lag:{}, noise frequency:{}\".format(\n",
    "                cur_p, cur_sprs, cur_add, cur_noise\n",
    "            )\n",
    "        )\n",
    "        cur_C, cur_S, cur_b0, cur_c0, cur_g, cur_mask = update_temporal(\n",
    "            A_sub,\n",
    "            C_sub,\n",
    "            YrA=YrA,\n",
    "            sparse_penal=cur_sprs,\n",
    "            p=cur_p,\n",
    "            use_smooth=True,\n",
    "            add_lag=cur_add,\n",
    "            noise_freq=cur_noise,\n",
    "        )\n",
    "        YA_dict[ks], C_dict[ks], S_dict[ks], g_dict[ks], sig_dict[ks], A_dict[ks] = (\n",
    "            YrA.compute(),\n",
    "            cur_C.compute(),\n",
    "            cur_S.compute(),\n",
    "            cur_g.compute(),\n",
    "            (cur_C + cur_b0 + cur_c0).compute(),\n",
    "            A_sub.compute(),\n",
    "        )\n",
    "    hv_res = visualize_temporal_update(\n",
    "        YA_dict,\n",
    "        C_dict,\n",
    "        S_dict,\n",
    "        g_dict,\n",
    "        sig_dict,\n",
    "        A_dict,\n",
    "        kdims=[\"p\", \"sparse penalty\", \"additional lag\", \"noise frequency\"],\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.output(size=int(output_size * 0.6))\n",
    "if interactive:\n",
    "    display(hv_res)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### temporal update"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "YrA = save_minian(\n",
    "    compute_trace(Y_fm_chk, A, b, C_chk, f).rename(\"YrA\"),\n",
    "    intpath,\n",
    "    overwrite=True,\n",
    "    chunks={\"unit_id\": 1, \"frame\": -1},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "C_new, S_new, b0_new, c0_new, g, mask = update_temporal(\n",
    "    A, C, YrA=YrA, **param_second_temporal\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### visualization of temporal components"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "You might notice \"Temporal Trace Last\" has different number of units than \"Spikes Last\".\n",
    "This is because `C` was merged in previous steps while `S` was not.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.output(size=int(output_size * 0.6))\n",
    "opts_im = dict(frame_width=500, aspect=2, colorbar=True, cmap=\"Viridis\")\n",
    "(\n",
    "    regrid(\n",
    "        hv.Image(\n",
    "            C.compute().astype(np.float32).rename(\"c1\"), kdims=[\"frame\", \"unit_id\"]\n",
    "        ).opts(**opts_im)\n",
    "    ).relabel(\"Temporal Trace Last\")\n",
    "    + regrid(\n",
    "        hv.Image(\n",
    "            S.compute().astype(np.float32).rename(\"s1\"), kdims=[\"frame\", \"unit_id\"]\n",
    "        ).opts(**opts_im)\n",
    "    ).relabel(\"Spikes Last\")\n",
    "    + regrid(\n",
    "        hv.Image(\n",
    "            C_new.compute().astype(np.float32).rename(\"c2\"), kdims=[\"frame\", \"unit_id\"]\n",
    "        ).opts(**opts_im)\n",
    "    ).relabel(\"Temporal Trace New\")\n",
    "    + regrid(\n",
    "        hv.Image(\n",
    "            S_new.compute().astype(np.float32).rename(\"s2\"), kdims=[\"frame\", \"unit_id\"]\n",
    "        ).opts(**opts_im)\n",
    "    ).relabel(\"Spikes New\")\n",
    ").cols(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### visualization of dropped units"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.output(size=int(output_size * 0.6))\n",
    "if interactive:\n",
    "    h, w = A.sizes[\"height\"], A.sizes[\"width\"]\n",
    "    im_opts = dict(aspect=w / h, frame_width=500, cmap=\"Viridis\")\n",
    "    cr_opts = dict(aspect=3, frame_width=1000)\n",
    "    bad_units = mask.where(mask == False, drop=True).coords[\"unit_id\"].values\n",
    "    if len(bad_units) > 0:\n",
    "        hv_res = (\n",
    "            hv.NdLayout(\n",
    "                {\n",
    "                    \"Spatial Footprint\": Dynamic(\n",
    "                        hv.Dataset(A.sel(unit_id=bad_units).compute().rename(\"A\"))\n",
    "                        .to(hv.Image, kdims=[\"width\", \"height\"])\n",
    "                        .opts(**im_opts)\n",
    "                    ),\n",
    "                    \"Spatial Footprints of Accepted Units\": Dynamic(\n",
    "                        hv.Image(\n",
    "                            A.sel(unit_id=mask).sum(\"unit_id\").compute().rename(\"A\"),\n",
    "                            kdims=[\"width\", \"height\"],\n",
    "                        ).opts(**im_opts)\n",
    "                    ),\n",
    "                }\n",
    "            )\n",
    "            + datashade(\n",
    "                hv.Dataset(YrA.sel(unit_id=bad_units).rename(\"raw\")).to(\n",
    "                    hv.Curve, kdims=[\"frame\"]\n",
    "                )\n",
    "            )\n",
    "            .opts(**cr_opts)\n",
    "            .relabel(\"Temporal Trace\")\n",
    "        ).cols(1)\n",
    "        display(hv_res)\n",
    "    else:\n",
    "        print(\"No rejected units to display\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### visualization of accepted units"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.output(size=int(output_size * 0.6))\n",
    "if interactive:\n",
    "    sig = C_new + b0_new + c0_new\n",
    "    display(\n",
    "        visualize_temporal_update(\n",
    "            YrA.sel(unit_id=mask),\n",
    "            C_new,\n",
    "            S_new,\n",
    "            g,\n",
    "            sig,\n",
    "            A.sel(unit_id=mask),\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### save results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "C = save_minian(\n",
    "    C_new.rename(\"C\").chunk({\"unit_id\": 1, \"frame\": -1}), intpath, overwrite=True\n",
    ")\n",
    "C_chk = save_minian(\n",
    "    C.rename(\"C_chk\"),\n",
    "    intpath,\n",
    "    overwrite=True,\n",
    "    chunks={\"unit_id\": -1, \"frame\": chk[\"frame\"]},\n",
    ")\n",
    "S = save_minian(\n",
    "    S_new.rename(\"S\").chunk({\"unit_id\": 1, \"frame\": -1}), intpath, overwrite=True\n",
    ")\n",
    "b0 = save_minian(\n",
    "    b0_new.rename(\"b0\").chunk({\"unit_id\": 1, \"frame\": -1}), intpath, overwrite=True\n",
    ")\n",
    "c0 = save_minian(\n",
    "    c0_new.rename(\"c0\").chunk({\"unit_id\": 1, \"frame\": -1}), intpath, overwrite=True\n",
    ")\n",
    "A = A.sel(unit_id=C.coords[\"unit_id\"].values)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "generate_videos(varr.sel(subset), Y_fm_chk, A=A, C=C_chk, vpath=dpath)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "if interactive:\n",
    "    cnmfviewer = CNMFViewer(A=A, C=C, S=S, org=Y_fm_chk)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.output(size=int(output_size * 0.35))\n",
    "if interactive:\n",
    "    display(cnmfviewer.show())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## save unit labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if interactive:\n",
    "    A = A.assign_coords(unit_labels=(\"unit_id\", cnmfviewer.unit_labels))\n",
    "    C = C.assign_coords(unit_labels=(\"unit_id\", cnmfviewer.unit_labels))\n",
    "    S = S.assign_coords(unit_labels=(\"unit_id\", cnmfviewer.unit_labels))\n",
    "    c0 = c0.assign_coords(unit_labels=(\"unit_id\", cnmfviewer.unit_labels))\n",
    "    b0 = b0.assign_coords(unit_labels=(\"unit_id\", cnmfviewer.unit_labels))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## save final results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "A = save_minian(A.rename(\"A\"), **param_save_minian)\n",
    "C = save_minian(C.rename(\"C\"), **param_save_minian)\n",
    "S = save_minian(S.rename(\"S\"), **param_save_minian)\n",
    "c0 = save_minian(c0.rename(\"c0\"), **param_save_minian)\n",
    "b0 = save_minian(b0.rename(\"b0\"), **param_save_minian)\n",
    "b = save_minian(b.rename(\"b\"), **param_save_minian)\n",
    "f = save_minian(f.rename(\"f\"), **param_save_minian)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## close cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.close()\n",
    "cluster.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.15"
  },
  "name": "pipeline.ipynb",
  "toc-autonumbering": false,
  "toc-showcode": false,
  "toc-showmarkdowntxt": false,
  "toc-showtags": true
 },
 "nbformat": 4,
 "nbformat_minor": 4
}