demos/notebooks/demo_realtime_cnmfE.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pipeline for real-time processing of microendoscopic data with CaImAn\n",
    "This demo presents 3 approaches for processing microendoscopic data in real time using CaImAn. \n",
    "1. Sufficiently long initialization phase to identify all ROIs  followed by tracking\n",
    "2. Short initalization phase followed by online processing using OnACID-E \n",
    "3. Short initalization phase followed by online processing using Ring-CNN+OnACID\n",
    "\n",
    "All approached include:\n",
    "- Motion Correction using the NoRMCorre algorithm\n",
    "- Source extraction using a variant of the CNMF algorithm\n",
    "- Deconvolution using the OASIS algorithm\n",
    "\n",
    "OnACID-E and Ring-CNN further include\n",
    "- Detection of new components\n",
    "- Refinement of neural spatial footprints (and background in the case of OnACID-E)\n",
    "\n",
    "@author: Johannes Friedrich @j-friedrich. Special thanks to Bernardo Sabatini and his lab at Harvard Medical School for sharing the data used in this demo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "get_ipython().run_line_magic('load_ext' , 'autoreload')\n",
    "get_ipython().run_line_magic('autoreload', 2)\n",
    "\n",
    "import os\n",
    "os.environ['MKL_NUM_THREADS'] = '1'\n",
    "os.environ['OPENBLAS_NUM_THREADS'] = '1'\n",
    "\n",
    "import bokeh.plotting as bpl\n",
    "import cv2\n",
    "import holoviews as hv\n",
    "import logging\n",
    "import matplotlib.pyplot as plt\n",
    "from multiprocessing import Process, Queue, set_start_method\n",
    "import numpy as np\n",
    "from scipy.sparse import csr_matrix\n",
    "from sys import platform\n",
    "import tensorflow as tf\n",
    "from time import time, sleep\n",
    "\n",
    "logging.basicConfig(format=\n",
    "                          \"%(relativeCreated)12d [%(filename)s:%(funcName)10s():%(lineno)s] [%(process)d] %(message)s\",\n",
    "                    # filename=\"/tmp/caiman.log\",\n",
    "                    level=logging.WARNING)\n",
    "\n",
    "import caiman as cm\n",
    "from caiman.source_extraction import cnmf\n",
    "from caiman.source_extraction.cnmf.online_cnmf import demix1p\n",
    "from caiman.utils.nn_models import (fit_NL_model, create_LN_model, quantile_loss, rate_scheduler)\n",
    "from caiman.utils.utils import download_demo\n",
    "\n",
    "bpl.output_notebook()\n",
    "hv.notebook_extension('bokeh')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if tf.__version__[0] == '1':\n",
    "    tf.enable_eager_execution()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below, we mention the code that can be used to read directly from the camera.<br/>\n",
    "In this demo we simulate it by reading frame by frame from a file instead, so that everybody can execute the demo.\n",
    "The `download_demo` function will download the file (if not already present) and store it inside your caiman_data/example_movies folder. We create an `iterator` over the file that returns the next imaged frame when calling `next(iterator)`.<br/>\n",
    "Here we define a function that creates and returns such an iterator. It works for both, a file or an imaging device."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_iterator(device=0, fr=None):\n",
    "    \"\"\"\n",
    "    device: device number (int) or filename (string) for reading from camera or file respectively\n",
    "    fr: frame rate\n",
    "    \"\"\"\n",
    "    if isinstance(device, int):  # capture from camera\n",
    "        def capture_iter(device=device, fr=fr):\n",
    "            cap = cv2.VideoCapture(device)\n",
    "            if fr is not None:  # set frame rate\n",
    "                cap.set(cv2.CAP_PROP_FPS, fr)\n",
    "            while True:\n",
    "                yield cv2.cvtColor(cap.read()[1], cv2.COLOR_BGR2GRAY)\n",
    "        iterator = capture_iter(device, fr)\n",
    "    else:  # read frame by frame from file\n",
    "        iterator = cm.base.movies.load_iter(device, var_name_hdf5='Y')\n",
    "    return iterator "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Sufficiently long initialization phase to identify all ROIs  followed by tracking"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Record for few minutes\n",
    "The `download_demo` function will download the file (if not already present) and store it inside your caiman_data/example_movies folder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "init_batch = 3000  # number of frames to use for initialization\n",
    "T = 6000           # total number of frames \n",
    "fr = 10            # frame rate (Hz)\n",
    "\n",
    "iterator = get_iterator(download_demo('blood_vessel_10Hz.mat'))\n",
    "\n",
    "m = cm.movie(np.array([next(iterator) for t in range(init_batch)], dtype='float32'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Take a break from imaging to process recorded data\n",
    "Taking a break to keep this demo simple. One could in parallel continue to save the otherwise \"lost\" frames to disk if one was not only intersted in the real-time experiment but post-analysis of the entire session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fname_init = m.save('init.mmap', order='C')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params_dict = {'fnames': fname_init,            # filename(s) to be processed\n",
    "               'fr': fr,                        # frame rate (Hz)\n",
    "               'method_init': 'corr_pnr',       # use corr_pnr for initialization of 1p data\n",
    "               'K': None,                       # upper bound on number of components per patch, in general None\n",
    "               'gSig': (3, 3),                  # gaussian width of a 2D gaussian kernel, which approximates a neuron\n",
    "               'gSiz': (13, 13),                # average diameter of a neuron, in general 4*gSig+1\n",
    "               'merge_thr': .65,                # merging threshold, max correlation allowed\n",
    "               'p': 1,                          # order of the autoregressive system\n",
    "               'tsub': 1,                       # downsampling factor in time for initialization\n",
    "               'ssub': 1,                       # downsampling factor in space for initialization\n",
    "               'only_init': True,               # set it to True to run CNMF-E\n",
    "               'nb': 0,                         # number of background components (rank) if positive,\n",
    "                                                # set it to 0 to use the exact ring model and return background as b0 and W\n",
    "               'min_corr': .7,                  # min peak value from correlation image\n",
    "               'min_pnr': 7,                    # min peak to noise ration from PNR image\n",
    "               'normalize_init': False,         # whether to equalize the movies during initialization\n",
    "               'ring_size_factor': 1.4,         # radius of ring is gSiz*ring_size_factor\n",
    "               'center_psf': True,              # whether to center the Gaussian convolution kernel, set it to True for 1 photon          \n",
    "               'init_iter': 1,                  # number of iterations during corr_pnr (1p) initialization\n",
    "               's_min': -10,                    # minimum spike threshold\n",
    "               'init_batch': init_batch,        # length of mini batch for initialization\n",
    "               'init_method': 'cnmf',           # initialization method for initial batch\n",
    "               'use_dense': False,              # flag for representation and storing of A and b\n",
    "               'motion_correct': True,          # flag for performing motion correction\n",
    "               'gSig_filt': (3, 3)}             # size of high pass spatial filtering, used in 1p data\n",
    "opts = cnmf.params.CNMFParams(params_dict=params_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cnm1 = cnmf.online_cnmf.OnACID(dview=None, params=opts)\n",
    "cnm1.initialize_online(T=T)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Alternatively we could initialize using any other pipeline that models the movie data as $Y=A C + B + N$ with spatial components $A$, temporal components $C$, background $B$, and noise $N$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# def init_from_other_pipeline(cnm, movie, A, C, T):\n",
    "#     Yr = movie.reshape(len(m), -1).T\n",
    "#     e = cnm.estimates\n",
    "#     e.A = A\n",
    "#     e.C = np.zeros_like(C)\n",
    "#     e.bl = np.zeros(len(C))\n",
    "#     e.c1 = np.zeros(len(C))\n",
    "#     e.g = np.zeros((len(C), 1))\n",
    "#     e.neurons_sn = np.zeros(len(C))\n",
    "#     e.S = np.zeros_like(C)\n",
    "#     e.lam = np.zeros(len(C))\n",
    "#     for i,c in enumerate(C):\n",
    "#         e.C[i], e.bl[i], e.c1[i], e.g[i], e.neurons_sn[i], e.S[i], e.lam[i] = (\n",
    "#         cm.source_extraction.cnmf.deconvolution.constrained_foopsi(\n",
    "#         c, p=1, bas_nonneg=False, noise_method='mean', fudge_factor=.97, optimize_g=5))\n",
    "#     e.YrA = C-e.C\n",
    "#     e.W, e.b = cm.source_extraction.cnmf.initialization.compute_W(\n",
    "#         Yr, e.A, e.C, movie.shape[1:],\n",
    "#         cnm.params.get('init', 'gSiz')[0] * cnm.params.get('init', 'ring_size_factor'),\n",
    "#         ssub=cnm.params.get('init', 'ssub_B'))\n",
    "#     cnm._prepare_object(Yr, T)\n",
    "#     return cnm\n",
    "#\n",
    "# replace A,C with the results of your pipeline, and m with your initialization movie \n",
    "# A, C = cnm.estimates.Ab, cnm.estimates.noisyC[:cnm.N]\n",
    "# cnm = cnmf.online_cnmf.OnACID(dview=None, params=opts)\n",
    "# cnm = init_from_other_pipeline(cnm, m, A, C, T)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Start real-time processing "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cnm1.t_read = []\n",
    "cnm1.t_demix = []\n",
    "cnm1.t_deconvolve = []\n",
    "est = cnm1.estimates\n",
    "ssub_B = cnm1.params.get('init', 'ssub_B') * cnm1.params.get('init', 'ssub')\n",
    "for t in range(init_batch, T):\n",
    "    # read frame\n",
    "    t0 = time()\n",
    "    frame = next(iterator)\n",
    "    cnm1.t_read.append(time()-t0)\n",
    "    # motion correct\n",
    "    t0 = time()\n",
    "    frame = cnm1.mc_next(t, frame.astype(np.float32))\n",
    "    cnm1.t_motion.append(time()-t0)\n",
    "    # get noisy fluorescence value via NNLS (project data on shapes & demix)\n",
    "    t0 = time()\n",
    "    C_in = est.noisyC[:cnm1.M, t - 1].copy()\n",
    "    est.C_on[:cnm1.M, t], est.noisyC[:cnm1.M, t] = demix1p(\n",
    "        frame.ravel(order='F'), est.Ab, C_in, est.AtA,\n",
    "        Atb=est.Atb, AtW=est.AtW, AtWA=est.AtWA,\n",
    "        iters=3, groups=est.groups, ssub_B=ssub_B, \n",
    "        downscale_matrix=est.downscale_matrix if ssub_B > 1 else None)\n",
    "    cnm1.t_demix.append(time()-t0)\n",
    "    # denoise & deconvolve\n",
    "    t0 = time()\n",
    "    if cnm1.params.get('preprocess', 'p'):\n",
    "        for i, o in enumerate(est.OASISinstances):\n",
    "            o.fit_next(est.noisyC[i, t])\n",
    "            est.C_on[i, t - o.get_l_of_last_pool() +\n",
    "                      1: t + 1] = o.get_c_of_last_pool()\n",
    "    cnm1.t_deconvolve.append(time()-t0)\n",
    "    # add code to display whatever you want in order to guide the closed-loop experiment below\n",
    "    # e.g. print the indices of neurons that just spiked\n",
    "#     print('\\r', ' '*200, end=\"\\r\")\n",
    "#     for i, o in enumerate(est.OASISinstances):\n",
    "#         if o.get_l_of_last_pool() == 1:\n",
    "#             print(i, end=' ')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "del iterator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot results "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate time one would have to wait for next frame to arrive if reading directly from camera\n",
    "t_all = np.cumsum(cnm1.t_read) + np.cumsum(cnm1.t_motion) + np.cumsum(cnm1.t_demix) + np.cumsum(cnm1.t_deconvolve)\n",
    "\n",
    "t_wait=[1]\n",
    "t_wait_total=0\n",
    "for i in range(1, T-init_batch):\n",
    "    t_wait.append(max(i/fr - t_all[i-1]-t_wait_total, 0))\n",
    "    t_wait_total += t_wait[-1]\n",
    "realtime = np.array(t_wait)>0\n",
    "print('%g%s processed in real time. %g/%g frames' % \n",
    "      (100 * realtime.sum() / (T-init_batch), '%', realtime.sum(), T-init_batch))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(12,4))\n",
    "for i, f in enumerate((lambda a: 1000*np.array(a), np.cumsum)):\n",
    "    plt.subplot(1,2,1+i)\n",
    "    plt.stackplot(np.arange(T-init_batch), f(cnm1.t_read), f(cnm1.t_motion),\n",
    "                  f(cnm1.t_demix), f(cnm1.t_deconvolve))\n",
    "    plt.gca().add_artist(plt.legend(labels=['read', 'motion', 'demix', 'deconvolve'], loc=2))\n",
    "    plt.title('Processing time allocation')\n",
    "    plt.xlabel('Frame #')\n",
    "    plt.ylabel(('Processing time per frame [ms]', 'Cumulative processing time [s]')[i])\n",
    "    if i==0:\n",
    "        plt.fill_between(range(T-init_batch),[0]*(T-init_batch),\n",
    "                 [plt.ylim()[1]]*(T-init_batch), where=realtime, \n",
    "                 color='y', alpha=.1, edgecolor='y', zorder=-11, label='real time')\n",
    "        plt.gca().add_artist(plt.legend())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cnm1.estimates.A = cnm1.estimates.Ab\n",
    "cnm1.estimates.C = cnm1.estimates.C_on\n",
    "cnm1.estimates.YrA = cnm1.estimates.noisyC-cnm1.estimates.C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cn, pnr = cm.summary_images.correlation_pnr(cm.load(download_demo('blood_vessel_10Hz.mat'), var_name_hdf5='Y'), gSig=3, swap_dim=False) # change swap dim if output looks weird, it is a problem with tiffile\n",
    "cnm1.estimates.coordinates = None\n",
    "cnm1.estimates.plot_contours_nb(img=cn, thr=.6);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cnm1.estimates.nb_view_components(img=cn, denoised_color='red');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Short initalization phase followed by online processing using OnACID-E "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Record for some seconds\n",
    "The `download_demo` function will download the file (if not already present) and store it inside your caiman_data/example_movies folder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "init_batch = 500  # number of frames to use for initialization\n",
    "T = 6000          # total number of frames \n",
    "fr = 10           # frame rate (Hz)\n",
    "\n",
    "iterator = get_iterator(download_demo('blood_vessel_10Hz.mat'))\n",
    "\n",
    "m = cm.movie(np.array([next(iterator) for t in range(init_batch)], dtype='float32'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Take a break from imaging to process recorded data\n",
    "Taking a break to keep this demo simple. One could in parallel continue to save the otherwise \"lost\" frames to disk if one was not only intersted in the real-time experiment but post-analysis of the entire session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fname_init = m.save('init.mmap', order='C')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params_dict = {'fnames': fname_init,            # filename(s) to be processed\n",
    "               'fr': fr,                        # frame rate (Hz)\n",
    "               'method_init': 'corr_pnr',       # use corr_pnr for initialization of 1p data\n",
    "               'K': None,                       # upper bound on number of components per patch, in general None\n",
    "               'gSig': (3, 3),                  # gaussian width of a 2D gaussian kernel, which approximates a neuron\n",
    "               'gSiz': (13, 13),                # average diameter of a neuron, in general 4*gSig+1\n",
    "               'merge_thr': .65,                # merging threshold, max correlation allowed\n",
    "               'p': 1,                          # order of the autoregressive system\n",
    "               'tsub': 1,                       # downsampling factor in time for initialization\n",
    "               'ssub': 1,                       # downsampling factor in space for initialization\n",
    "               'only_init': True,               # set it to True to run CNMF-E\n",
    "               'nb': 0,                         # number of background components (rank) if positive,\n",
    "                                                # set it to 0 to use the exact ring model and return background as b0 and W\n",
    "               'min_corr': .65,                 # min peak value from correlation image\n",
    "               'min_pnr': 6,                    # min peak to noise ration from PNR image\n",
    "               'normalize_init': False,         # whether to equalize the movies during initialization\n",
    "               'ring_size_factor': 1.4,         # radius of ring is gSiz*ring_size_factor\n",
    "               'center_psf': True,              # whether to center the Gaussian convolution kernel, set it to True for 1 photon          \n",
    "               'init_iter': 1,                  # number of iterations during corr_pnr (1p) initialization\n",
    "               's_min': -10,                    # minimum spike threshold\n",
    "               'init_batch': init_batch,        # length of mini batch for initialization\n",
    "               'init_method': 'cnmf',           # initialization method for initial batch\n",
    "               'batch_update_suff_stat': True,  # flag for updating sufficient statistics (used for updating shapes) in batch\n",
    "               'expected_comps': 600,           # number of expected components\n",
    "               'rval_thr': .55,                 # space correlation threshold\n",
    "               'thresh_fitness_raw': -130,      # threshold for trace SNR \n",
    "               'thresh_fitness_delta': -20,     # threshold for trace SNR \n",
    "               'use_dense': False,              # flag for representation and storing of A and b\n",
    "               'motion_correct': True,          # flag for performing motion correction\n",
    "               'gSig_filt': (3, 3),             # size of high pass spatial filtering, used in 1p data\n",
    "               'use_cnn': False}\n",
    "opts = cnmf.params.CNMFParams(params_dict=params_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cnm2 = cnmf.online_cnmf.OnACID(dview=None, params=opts)\n",
    "cnm2.initialize_online(T=T)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Start real-time processing \n",
    "The spatial footprints and background parameters are updated every 200 frames. This results in a long processing time for the individual frames for which these updates occur. Thus instead of the usual waiting for the camera to provide the next image, there are few frames for which the images are acquired faster than processed. We use a separate `Process` to acquire and add the images to a FIFO `Queue` at regular time intervals. The main process reads the next image from this `Queue` and waits for it if the `Queue` is empty."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if platform in ('linux', 'darwin'):\n",
    "    \"\"\" This code assumes you are running on a Unix-based computing system (like Linux or macOS)\n",
    "        that allows to start a Process uing a fork() call, NOT Windows. For the latter see the else call below\n",
    "        https://docs.python.org/3.10/library/multiprocessing.html#the-spawn-and-forkserver-start-methods\"\"\"\n",
    "    set_start_method('fork')\n",
    "    q = Queue()\n",
    "    cnm2.t_wait = []\n",
    "    cnm2.t_motion = []\n",
    "    cnm2.t_fit = []\n",
    "    realtime = []\n",
    "\n",
    "    def append_to_queue(q, init_batch, T, fr):\n",
    "        t_start = time()\n",
    "        for t in range(init_batch, T):\n",
    "            # read frame and append to queue\n",
    "            frame = next(iterator)\n",
    "            q.put(frame)\n",
    "            sleep(max(0, (t+1-init_batch)/fr - time()+t_start))\n",
    "\n",
    "    producer = Process(target=append_to_queue, args=(q,init_batch,T,fr))\n",
    "    producer.start()\n",
    "\n",
    "    # process first frame\n",
    "    t = init_batch\n",
    "    # read form queue (wait for next frame if empty)\n",
    "    t0=time()\n",
    "    frame = q.get()\n",
    "    cnm2.t_wait.append(time()-t0)\n",
    "    t_start = time() + (1-init_batch)/fr\n",
    "    # motion correct\n",
    "    t0 = time()\n",
    "    frame = cnm2.mc_next(t, frame)\n",
    "    cnm2.t_motion.append(time()-t0)\n",
    "    # fit\n",
    "    t0 = time()\n",
    "    cnm2.fit_next(t, frame.ravel(order='F'))\n",
    "    cnm2.t_fit.append(time()-t0)\n",
    "    rt = time() <= t_start + t/fr\n",
    "    realtime.append(rt)\n",
    "\n",
    "    # process remaining frames\n",
    "    for t in range(init_batch+1, T):\n",
    "        # read form queue (wait for next frame if empty)\n",
    "        t0=time()\n",
    "        frame = q.get()\n",
    "        cnm2.t_wait.append(time()-t0)\n",
    "        # motion correct\n",
    "        t0 = time()\n",
    "        frame = cnm2.mc_next(t, frame)\n",
    "        cnm2.t_motion.append(time()-t0)\n",
    "        # fit\n",
    "        t0 = time()\n",
    "        cnm2.fit_next(t, frame.ravel(order='F'))\n",
    "        cnm2.t_fit.append(time()-t0)\n",
    "        rt = time() <= t_start + t/fr\n",
    "        realtime.append(rt)\n",
    "        # add code to display whatever you want in order to guide the closed-loop experiment below\n",
    "        print('Realtime: ' + (\"\\x1b[32mTrue\\x1b[0m\" if rt else \"\\x1b[31mFalse\\x1b[0m\"), end=\"  \\r\", flush=True)\n",
    "else:\n",
    "\n",
    "    \"\"\"Windows requires as workaround to move some code to the external file 'win.py'.\n",
    "       The new Process is being spawned not forked. Expect performance to be somewhat impeded.\"\"\"\n",
    "    print('Windows lacks os.fork(). Expect performance to be somewhat impeded.')\n",
    "    import win\n",
    "\n",
    "    if __name__ == '__main__':\n",
    "        q = Queue()\n",
    "        cnm2.t_wait = []\n",
    "        cnm2.t_motion = []\n",
    "        cnm2.t_fit = []\n",
    "        realtime = []\n",
    "        producer = Process(target=win.append_to_queue, args=(q,init_batch,T,fr))\n",
    "        producer.start()\n",
    "\n",
    "        # process first frame\n",
    "        t = init_batch\n",
    "        # read form queue (wait for next frame if empty)\n",
    "        t0=time()\n",
    "        frame = q.get()\n",
    "        cnm2.t_wait.append(time()-t0)\n",
    "        t_start = time() + (1-init_batch)/fr\n",
    "        # motion correct\n",
    "        t0 = time()\n",
    "        frame = cnm2.mc_next(t, frame)\n",
    "        cnm2.t_motion.append(time()-t0)\n",
    "        # fit\n",
    "        t0 = time()\n",
    "        cnm2.fit_next(t, frame.ravel(order='F'))\n",
    "        cnm2.t_fit.append(time()-t0)\n",
    "        rt = time() <= t_start + t/fr\n",
    "        realtime.append(rt)\n",
    "\n",
    "        #process remaining frames\n",
    "        for t in range(init_batch+1, T):\n",
    "            # read form queue (wait for next frame if empty)\n",
    "            t0=time()\n",
    "            frame = q.get()\n",
    "            cnm2.t_wait.append(time()-t0)\n",
    "            # motion correct\n",
    "            t0 = time()\n",
    "            frame = cnm2.mc_next(t, frame)\n",
    "            cnm2.t_motion.append(time()-t0)\n",
    "            # fit\n",
    "            t0 = time()\n",
    "            cnm2.fit_next(t, frame.ravel(order='F'))\n",
    "            cnm2.t_fit.append(time()-t0)\n",
    "            rt = time() <= t_start + t/fr\n",
    "            realtime.append(rt)\n",
    "            # add code to display whatever you want in order to guide the closed-loop experiment below\n",
    "            print('Realtime: ' + (\"\\x1b[32mTrue\\x1b[0m\" if rt else \"\\x1b[31mFalse\\x1b[0m\"), end=\"  \\r\", flush=True)\n",
    "        producer.join()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "del iterator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot results "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('%g%s processed in real time. %g/%g frames' % \n",
    "      (100 * np.sum(realtime) / (T-init_batch), '%', np.sum(realtime), T-init_batch))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(12,4))\n",
    "cnm2.t_read = np.zeros(T-init_batch)\n",
    "for i, f in enumerate((lambda a: 1000*np.array(a), np.cumsum)):\n",
    "    plt.subplot(1,2,1+i)\n",
    "    plt.stackplot(np.arange(len(cnm2.t_fit)), f(cnm2.t_read), f(cnm2.t_motion),\n",
    "                  f(np.array(cnm2.t_fit) - np.array([cnm2.t_detect, cnm2.t_shapes, cnm2.t_stat]).sum(0)),\n",
    "                  f(cnm2.t_detect), f(cnm2.t_shapes)+f(cnm2.t_stat))\n",
    "    plt.gca().add_artist(plt.legend(labels=['read', 'motion', 'process', 'detect', 'shapes'], loc=2))\n",
    "    plt.title('Processing time allocation')\n",
    "    plt.xlabel('Frame #')\n",
    "    plt.ylabel(('Processing time per frame [ms]', 'Cumulative processing time [s]')[i])\n",
    "    if i==0:\n",
    "        plt.ylim(0, 100)\n",
    "        plt.fill_between(range(T-init_batch),[0]*(T-init_batch),\n",
    "                 [100]*(T-init_batch), where=realtime, \n",
    "                 color='y', alpha=.1, edgecolor='y', zorder=-11, label='real time')\n",
    "        plt.gca().add_artist(plt.legend())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cnm2.estimates.A = cnm2.estimates.Ab\n",
    "cnm2.estimates.C = cnm2.estimates.C_on[:cnm2.N]\n",
    "cnm2.estimates.YrA = cnm2.estimates.noisyC[:cnm2.N]-cnm2.estimates.C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cn, pnr = cm.summary_images.correlation_pnr(cm.load(download_demo('blood_vessel_10Hz.mat'), var_name_hdf5='Y'), gSig=3, swap_dim=False) # change swap dim if output looks weird, it is a problem with tiffile\n",
    "cnm2.estimates.coordinates = None\n",
    "cnm2.estimates.plot_contours_nb(img=cn, thr=.6);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cnm2.estimates.nb_view_components(img=cn, denoised_color='red');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Short initalization phase followed by online processing using Ring-CNN+OnACID"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Record for some seconds\n",
    "The `download_demo` function will download the file (if not already present) and store it inside your caiman_data/example_movies folder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "init_batch = 500  # number of frames to use for initialization\n",
    "T = 6000          # total number of frames \n",
    "fr = 10           # frame rate (Hz)\n",
    "\n",
    "iterator = get_iterator(download_demo('blood_vessel_10Hz.mat'))\n",
    "\n",
    "m = cm.movie(np.array([next(iterator) for t in range(init_batch)], dtype='float32'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Take a break from imaging to process recorded data\n",
    "Taking a break to keep this demo simple. One could in parallel continue to save the otherwise \"lost\" frames to disk if one was not only intersted in the real-time experiment but post-analysis of the entire session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fname_init = m.save('init.mmap', order='C')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reuse_model = False                                       # set to True to re-use an existing ring model\n",
    "path_to_model = None                                      # specify a pre-trained model here if needed \n",
    "gSig = (7, 7)                                             # expected half size of neurons\n",
    "gnb = 2                                                   # number of background components for OnACID\n",
    "\n",
    "params_dict = {'fnames': fname_init,                      # filename(s) to be processed\n",
    "               'var_name_hdf5': 'Y',                      # name of variable inside mat file where the data is stored\n",
    "               'fr': fr,                                  # frame rate (Hz)\n",
    "               'decay_time': 0.5,                         # approximate length of transient event in seconds\n",
    "               'gSig': gSig,                              # gaussian width of a 2D gaussian kernel, which approximates a neuron\n",
    "               'p': 1,                                    # order of AR indicator dynamics\n",
    "               'ring_CNN': True,                          # SET TO TRUE TO USE RING CNN \n",
    "               'min_SNR': 2.65,                           # minimum SNR for accepting new components\n",
    "               'SNR_lowest': 0.75,                        # reject components with SNR below this value\n",
    "               'use_cnn': False,                          # do not use CNN based test for components\n",
    "               'use_ecc': True,                           # test eccentricity\n",
    "               'max_ecc': 2.625,                          # reject components with eccentricity above this value\n",
    "               'rval_thr': 0.70,                          # correlation threshold for new component inclusion\n",
    "               'rval_lowest': 0.25,                       # reject components with corr below that value\n",
    "               'ds_factor': 1,                            # spatial downsampling factor (increases speed but may lose some fine structure)\n",
    "               'nb': gnb,                                 # number of background components (rank)\n",
    "               'motion_correct': True,                    # Flag for motion correction\n",
    "               'init_batch': init_batch,                  # number of frames for initialization (presumably from the first file)\n",
    "               'init_method': 'bare',                     # initialization method\n",
    "               'normalize': False,                        # Whether to normalize each frame prior to online processing\n",
    "               'expected_comps': 700,                     # maximum number of expected components used for memory pre-allocation (exaggerate here)\n",
    "               'sniper_mode': False,                      # flag using a CNN to detect new neurons (o/w space correlation is used)\n",
    "               'dist_shape_update' : True,                # flag for updating shapes in a distributed way\n",
    "               'min_num_trial': 5,                        # number of candidate components per frame\n",
    "               'epochs': 1,                               # number of total passes over the data\n",
    "               'stop_detection': False,                   # Run a last epoch without detecting new neurons  \n",
    "               'K': 50,                                   # initial number of components\n",
    "               'lr': 6e-4,                                # (initial) learning rate\n",
    "               'lr_scheduler': [0.9, 6000, 10000],        # learning rate scheduler\n",
    "               'pct': 0.01,                               # quantile of the quantile loss function\n",
    "               'path_to_model': path_to_model,            # where the ring CNN model is saved/loaded\n",
    "               'reuse_model': reuse_model,                # flag for re-using a ring CNN model\n",
    "              }\n",
    "opts = cnmf.params.CNMFParams(params_dict=params_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Train or load Ring CNN "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cnm3 = cnmf.online_cnmf.OnACID(params=opts)\n",
    "\n",
    "if cnm3.params.get('ring_CNN', 'loss_fn') == 'pct':\n",
    "    loss_fn = quantile_loss(cnm3.params.get('ring_CNN', 'pct'))\n",
    "else:\n",
    "    loss_fn = cnm3.params.get('ring_CNN', 'loss_fn')\n",
    "if cnm3.params.get('ring_CNN', 'lr_scheduler') is None:\n",
    "    sch = None\n",
    "else:\n",
    "    sch = rate_scheduler(*cnm3.params.get('ring_CNN', 'lr_scheduler'))\n",
    "model_LN = create_LN_model(m, shape=m.shape[1:] + (1,),\n",
    "                           n_channels=cnm3.params.get('ring_CNN', 'n_channels'),\n",
    "                           lr=cnm3.params.get('ring_CNN', 'lr'),\n",
    "                           gSig=cnm3.params.get('init', 'gSig')[0],\n",
    "                           loss=loss_fn, width=cnm3.params.get('ring_CNN', 'width'),\n",
    "                           use_add=cnm3.params.get('ring_CNN', 'use_add'),\n",
    "                           use_bias=cnm3.params.get('ring_CNN', 'use_bias'))\n",
    "if cnm3.params.get('ring_CNN', 'reuse_model'):\n",
    "    model_LN.load_weights(cnm3.params.get('ring_CNN', 'path_to_model'))\n",
    "else:\n",
    "    model_LN, history, path_to_model = fit_NL_model(\n",
    "        model_LN, m, epochs=cnm3.params.get('ring_CNN', 'max_epochs'),\n",
    "        patience=cnm3.params.get('ring_CNN', 'patience'), schedule=sch)\n",
    "    cnm3.params.set('ring_CNN', {'path_to_model': path_to_model})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Initialize OnACID "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cnm3.initialize_online(T=T, model_LN=model_LN)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# if no GPU is available prediction is faster when the CNN is converted into a sparse matrix \n",
    "if tf.test.is_gpu_available():\n",
    "    predict = model_LN.predict\n",
    "    # dummy prediction to initialize CNN-ring model model_LN\n",
    "    predict(m[:1, ..., None])\n",
    "else:\n",
    "    dims = m[0].shape\n",
    "    kernel_size = model_LN.layers[1].weights[0].shape[0]\n",
    "    kernel_half = (kernel_size-1)//2\n",
    "    w_1 = model_LN.layers[1].weights[0].numpy().squeeze().transpose(2,0,1)\n",
    "    w_1 = w_1.reshape(len(w_1), -1)\n",
    "    w_2 = model_LN.layers[2].weights[0].numpy()\n",
    "    data = []\n",
    "    indices = []\n",
    "    indptr = [0]\n",
    "    for i in range(dims[0]):\n",
    "        for j in range(dims[1]):\n",
    "            tmp = np.zeros((dims[0]+kernel_size, dims[1]+kernel_size), dtype=np.float32)\n",
    "            tmp[i:i+kernel_size, j:j+kernel_size] = w_2[i, j].dot(w_1).reshape(kernel_size, kernel_size)\n",
    "            tmp = tmp[kernel_half:kernel_half+dims[0], kernel_half:kernel_half+dims[1]].ravel()\n",
    "            ind = np.where(tmp!=0)[0]\n",
    "            newdata = list(tmp[ind])\n",
    "            data += newdata\n",
    "            indices += list(ind)\n",
    "            indptr += [len(newdata) + indptr[-1]]\n",
    "    w = csr_matrix((data, indices, indptr)).todia()\n",
    "\n",
    "    def predict(frame):\n",
    "        return w.dot(frame.ravel()).reshape(dims)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Start real-time processing "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cnm3.t_read = []\n",
    "cnm3.t_bkgrd = []\n",
    "cnm3.t_motion = []\n",
    "cnm3.t_fit = []\n",
    "for t in range(init_batch, T):\n",
    "    # read frame\n",
    "    t0 = time()\n",
    "    frame = next(iterator)\n",
    "    cnm3.t_read.append(time()-t0)\n",
    "    # remove background\n",
    "    t0 = time()\n",
    "    frame = np.maximum(frame - np.squeeze(predict(\n",
    "        frame.astype(np.float32)[None,...,None])), 0)\n",
    "    cnm3.t_bkgrd.append(time()-t0)\n",
    "    # motion correct\n",
    "    t0 = time()\n",
    "    frame = cnm3.mc_next(t, frame.astype(np.float32))\n",
    "    cnm3.t_motion.append(time()-t0)\n",
    "    # fit\n",
    "    t0 = time()\n",
    "    cnm3.fit_next(t, frame.ravel(order='F'))\n",
    "    cnm3.t_fit.append(time()-t0)\n",
    "    # add code to display whatever you want in order to guide the closed-loop experiment below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "del iterator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot results "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate time one would have to wait for next frame to arrive if reading directly from camera\n",
    "t_all = np.cumsum(cnm3.t_read) + np.cumsum(cnm3.t_bkgrd) + np.cumsum(cnm3.t_motion) + np.cumsum(cnm3.t_fit)\n",
    "\n",
    "t_wait=[1]\n",
    "t_wait_total=0\n",
    "for i in range(1, T-init_batch):\n",
    "    t_wait.append(max(i/fr - t_all[i-1]-t_wait_total, 0))\n",
    "    t_wait_total += t_wait[-1]\n",
    "realtime = np.array(t_wait)>0\n",
    "print('%g%s processed in real time. %g/%g frames' % \n",
    "      (100 * realtime.sum() / (T-init_batch), '%', realtime.sum(), T-init_batch))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(12,4))\n",
    "for i, f in enumerate((lambda a: 1000*np.array(a), np.cumsum)):\n",
    "    plt.subplot(1,2,1+i)\n",
    "    plt.stackplot(np.arange(len(cnm3.t_fit)), f(cnm3.t_read), f(cnm3.t_bkgrd), f(cnm3.t_motion),\n",
    "                  f(np.array(cnm3.t_fit) - np.array([cnm3.t_detect, cnm3.t_shapes, cnm3.t_stat]).sum(0)),\n",
    "                  f(cnm3.t_detect), f(cnm3.t_shapes)+f(cnm3.t_stat))\n",
    "    plt.gca().add_artist(plt.legend(labels=['read', 'background', 'motion', 'process', 'detect', 'shapes'], loc=2))\n",
    "    plt.title('Processing time allocation')\n",
    "    plt.xlabel('Frame #')\n",
    "    plt.ylabel(('Processing time per frame [ms]', 'Cumulative processing time [s]')[i])\n",
    "    if i==0:\n",
    "        plt.fill_between(range(T-init_batch),[0]*(T-init_batch),\n",
    "                 [plt.ylim()[1]]*(T-init_batch), where=realtime, \n",
    "                 color='y', alpha=.1, edgecolor='y', zorder=-11, label='real time')\n",
    "        plt.gca().add_artist(plt.legend())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cnm3.estimates.A = cnm3.estimates.Ab[:,gnb:]\n",
    "cnm3.estimates.C = cnm3.estimates.C_on[gnb:cnm3.M]\n",
    "cnm3.estimates.YrA = cnm3.estimates.noisyC[gnb:cnm3.M]-cnm3.estimates.C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cn, pnr = cm.summary_images.correlation_pnr(cm.load(download_demo('blood_vessel_10Hz.mat'), var_name_hdf5='Y'), gSig=3, swap_dim=False) # change swap dim if output looks weird, it is a problem with tiffile\n",
    "cnm3.estimates.coordinates = None\n",
    "cnm3.estimates.plot_contours_nb(img=cn, thr=.6);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cnm3.estimates.nb_view_components(img=cn,denoised_color='red');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}