02_SourceExtract.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run CNMF source extraction on movies\n",
    "Step 2 of the Caiman processing pipeline for dendritic two-photon calcium imaging movies. This part uses mmap files as input. These are created during motion correction with the Caiman toolbox (see `01_Preprocess_MotionCorrect.ipynb`). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Imports & Setup\n",
    "The first cells import the various Python modules required by the notebook. In particular, a number of modules are imported from the Caiman package. In addition, we also setup the environment so that everything works as expected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generic imports\n",
    "# from __future__ import absolute_import, division, print_function\n",
    "# from builtins import *\n",
    "\n",
    "import os, platform, glob, sys, re\n",
    "import json\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.io import savemat\n",
    "\n",
    "from IPython.display import clear_output\n",
    "\n",
    "# Import Bokeh library\n",
    "from bokeh.plotting import Figure, show\n",
    "from bokeh.models import Range1d, CrosshairTool, HoverTool, Legend\n",
    "from bokeh.io import output_notebook\n",
    "from bokeh.io import export_svgs\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# on Linux we have to add the caiman folder to Pythonpath\n",
    "if platform.system() == 'Linux':\n",
    "    sys.path.append(os.path.expanduser('~/caiman'))\n",
    "# environment variables for parallel processing\n",
    "os.environ['MKL_NUM_THREADS']='1'\n",
    "os.environ['OPENBLAS_NUM_THREADS']='1'\n",
    "os.environ['VECLIB_MAXIMUM_THREADS']='1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# CaImAn imports\n",
    "import caiman as cm\n",
    "from caiman.source_extraction.cnmf import cnmf as cnmf\n",
    "from caiman.components_evaluation import estimate_components_quality as estimate_q\n",
    "from caiman.utils.visualization import plot_contours, nb_view_patches, nb_plot_contour\n",
    "from caiman.source_extraction.cnmf import utilities as cnmf_utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This has to be in a separate cell, otherwise it wont work.\n",
    "from bokeh import resources\n",
    "output_notebook(resources=resources.INLINE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Select files and parameters\n",
    "The following need to be specified:\n",
    "- data_folder ... location of the data (relative to ~/Data)\n",
    "- mc_output ... select if output of rigid ('rig') or piece-wise rigid ('els') motion correction should be used (currently only 'rig' is tested and works)\n",
    "- max_files ... maximum number of files to process, e.g. for testing (if 0, all files will be processed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "animal_folder = 'M1_for_processing'\n",
    "date_folder = 'M1_2018-01-31'\n",
    "session_folder = 'S1'\n",
    "group_id = 'G0'\n",
    "\n",
    "mc_output = 'rig'\n",
    "remove_bad_frames = True # remove bad frames specified in Json file\n",
    "\n",
    "n_processes = 0 # number of compute processes (None to select automatically)\n",
    "\n",
    "# create the complete path to the data folder\n",
    "if platform.system() == 'Linux':\n",
    "    data_folder = '/home/ubuntu/Data'\n",
    "elif platform.system() == 'Darwin':\n",
    "    data_folder = '/Users/Henry/polybox/Data_temp/Dendrites_Gwen'\n",
    "data_folder = os.path.join(data_folder, animal_folder, date_folder, session_folder)\n",
    "data_folder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# select the mmap file created during motion correction\n",
    "all_files = os.listdir(data_folder)\n",
    "matches = [x for x in all_files if x.startswith('%s_%s' % (date_folder, session_folder)) \n",
    "           and x.endswith('.mmap') and mc_output in x and group_id in x]\n",
    "\n",
    "if len(matches) == 1: # if only 1 file matches, select it\n",
    "    fname = matches[0]\n",
    "    print('Selected file: %s' % (fname))\n",
    "elif len(matches) > 0: # if several files match, ask which one to choos\n",
    "    print('Found several matching files!')\n",
    "    for ix, match in enumerate(matches):\n",
    "        print('%1.0f: %s' % (ix+1,match))\n",
    "    select = input('Please select your file (1-%1.0f):' % (ix+1))\n",
    "    fname = matches[int(select)-1]\n",
    "    print('Selected file: %s' % (fname))\n",
    "else: # if no file matches, throw an error\n",
    "    raise Exception('Did not find a matching mmap file!')\n",
    "\n",
    "fname = os.path.join(data_folder, fname)\n",
    "\n",
    "# get metadata for corresponding joined file\n",
    "meta_fname = fname[:fname.find('_crop_')+5] + '.json'\n",
    "meta = json.load(open(meta_fname))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup cluster\n",
    "The default backend mode for parallel processing is through the multiprocessing package. This will allow us to use all the cores in the VM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# start the cluster (if a cluster already exists terminate it)\n",
    "if 'dview' in locals():\n",
    "    dview.terminate()\n",
    "c, dview, n_processes = cm.cluster.setup_cluster(\n",
    "    backend='local', n_processes=n_processes, single_thread=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parameters for source extraction\n",
    "Next, we define the important parameters for calcium source extraction. These parameters will have to be iteratively refined for the respective datasets.\n",
    "\n",
    "The parameters are stored in the Python dictionary `params_cnmf`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# parameters for source extraction and deconvolution\n",
    "p = 1                       # order of the autoregressive system (normally 1, 2 for fast indicators / imaging)\n",
    "gnb = 2                     # number of global background components\n",
    "merge_thresh = 0.8          # merging threshold, max correlation allowed\n",
    "rf = None                   # half-size of the patches in pixels. e.g., if rf=25, patches are 50x50 / None: no patches\n",
    "stride_cnmf = 6             # amount of overlap between the patches in pixels\n",
    "K = 20                      # number of components per patch\n",
    "gSig = [4, 4]               # expected half size of neurons\n",
    "init_method = 'sparse_nmf'  # initialization method (if analyzing dendritic data using 'sparse_nmf')\n",
    "is_dendrites = True         # flag for analyzing dendritic data\n",
    "#alpha_snmf = 10e2           # sparsity penalty for dendritic data analysis through sparse NMF\n",
    "alpha_snmf = 1e-6\n",
    "\n",
    "method_deconvolution='oasis'# deconvolution method (oasis or cvxpy)\n",
    "final_rate = 13.1316             # final frame rate in Hz\n",
    "\n",
    "\n",
    "# parameters for component evaluation\n",
    "# min_SNR = 2.5               # signal to noise ratio for accepting a component\n",
    "# rval_thr = 0.8              # space correlation threshold for accepting a component\n",
    "# cnn_thr = 0.8               # threshold for CNN based classifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# store parameters in Python dictionary\n",
    "params = dict()\n",
    "params['p'] = p\n",
    "params['gnb'] = gnb\n",
    "params['merge_thresh'] = merge_thresh\n",
    "params['rf'] = rf\n",
    "params['stride_cnmf'] = stride_cnmf\n",
    "params['K'] = K\n",
    "params['gSig'] = gSig\n",
    "params['init_method'] = init_method\n",
    "params['is_dendrites'] = is_dendrites\n",
    "params['alpha_snmf'] = alpha_snmf\n",
    "\n",
    "params['method_deconvolution'] = method_deconvolution\n",
    "params['final_rate'] = final_rate\n",
    "\n",
    "params"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run CNMF on joined mmap file\n",
    "According to the Caiman tutorials, CNMF source extraction should be run in several stages. First, we run CNMF on patches of the dataset and then evaluate the quality of the extracted components. Next, CNMF is run again but on the full field-of-view. Finally, the extracted components are again evaluated and classified (good and bad)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load mmap file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Yr, dims, T = cm.load_memmap(fname)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# remove bad frames specified in corresponding Json file\n",
    "# Todo: print deleted frames\n",
    "if remove_bad_frames:\n",
    "    bad_frames = json.load(open(fname.replace('.mmap','badFrames.json')))\n",
    "    bad_frames = np.array(bad_frames['frames'])\n",
    "    Yr = np.delete(Yr, bad_frames, axis=1)\n",
    "    T = Yr.shape[1]\n",
    "    images = np.reshape(Yr.T, [T] + list(dims), order='F')\n",
    "    fname_new = cm.save_memmap([images], base_name=os.path.join(data_folder, 'removedFrames'))\n",
    "    Yr, dims, T = cm.load_memmap(fname_new)\n",
    "    print('Deleted %1d frames. Saved to new file %s.' % (len(bad_frames), os.path.basename(fname_new)))\n",
    "    print('Deleted frames:')\n",
    "    print(bad_frames)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "d1, d2 = dims\n",
    "images = np.reshape(Yr.T, [T] + list(dims), order='F')\n",
    "Y = np.reshape(Yr, dims + (T,), order='F')\n",
    "m_images = cm.movie(images)\n",
    "#  checks on movies (might take time if large!)\n",
    "if np.min(images) < 0:\n",
    "    raise Exception('Movie too negative, add_to_movie should be larger')\n",
    "if np.sum(np.isnan(images)) > 0:\n",
    "    raise Exception('Movie contains nan! You did not remove enough borders')\n",
    "# correlation image\n",
    "Cn = cm.local_correlations(Y)\n",
    "Cn[np.isnan(Cn)] = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sanity check the image\n",
    "avg_img = np.mean(images,axis=0)\n",
    "plt.figure(figsize=(10,20))\n",
    "plt.subplot(1,2,1)\n",
    "plt.imshow(avg_img, cmap='gray'), plt.title('Frame average')\n",
    "plt.subplot(1,2,2)\n",
    "plt.imshow(Cn, cmap='jet'), plt.title('Correlation image')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run single CNMF on full FoV, without patches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure CNMF and fit to the data\n",
    "cnmf_single = cnmf.CNMF(n_processes, k=params['K'], gSig=params['gSig'], merge_thresh=params['merge_thresh'], \n",
    "                p=params['p'], dview=dview, Ain=None, gnb=params['gnb'], rf=params['rf'],\n",
    "                memory_fact=1, method_init=params['init_method'], alpha_snmf=params['alpha_snmf'],\n",
    "                method_deconvolution=params['method_deconvolution'])\n",
    "cnmf_single = cnmf_single.fit(images)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Unravel results\n",
    "A, C, b, f, YrA, sn = cnmf_single.A, cnmf_single.C, cnmf_single.b, cnmf_single.f, cnmf_single.YrA, cnmf_single.sn\n",
    "# A   ... n_pixel x n_components sparse matrix (component locations)\n",
    "# C   ... n_component x t np.array (fitted signal)\n",
    "# b   ... ? np.array\n",
    "# f   ... ? np.array (b / f related to global background components)\n",
    "# YrA ... n_component x t np.array (residual)\n",
    "# sn  ... n_pixel np.array (SNR?)\n",
    "\n",
    "print('Detected components: %1.0f' % (C.shape[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evaluate components (discard low quality components)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "r_values_min = .65         # threshold on space consistency\n",
    "fitness_min = -100        # threshold on time variability\n",
    "fitness_delta_min = -100  # threshold on time variability (if nonsparse activity)\n",
    "Npeaks = 100\n",
    "thresh_C = 0.3\n",
    "Athresh = 0.1\n",
    "\n",
    "traces = C + YrA\n",
    "idx_comps, idx_comps_bad, fitness_raw, fitness_delta, r_values = estimate_q(traces, Y, A, C, b, f, final_frate=params['final_rate'], \n",
    "                                      Npeaks=Npeaks, r_values_min=r_values_min, fitness_min=fitness_min, \n",
    "                                      fitness_delta_min=fitness_delta_min, thresh_C=thresh_C, Athresh=Athresh,\n",
    "                                      return_all=True)\n",
    "print('Good / bad components: %1.0f / %1.0f' % (len(idx_comps), len(idx_comps_bad)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot good and bad components\n",
    "plt.figure(figsize=(20,30));\n",
    "plt.subplot(121); crd_good = cm.utils.visualization.plot_contours(A[:,idx_comps], Cn, thr=.8, vmax=0.75)\n",
    "plt.title('Contour plots of accepted components')\n",
    "plt.subplot(122); crd_bad = cm.utils.visualization.plot_contours(A[:,idx_comps_bad], Cn, thr=.8, vmax=0.75)\n",
    "plt.title('Contour plots of rejected components')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot good components on background image and as component map\n",
    "A_dense = A.todense()\n",
    "counter = 1\n",
    "plt.figure(figsize=(20,20));\n",
    "for i_comp in range(len(idx_comps)):\n",
    "    plt.subplot(len(idx_comps),2,counter)\n",
    "    counter += 1\n",
    "    dummy = cm.utils.visualization.plot_contours(A[:,idx_comps[i_comp]], avg_img, cmap='gray', \n",
    "                                                 colors='r', display_numbers=False)\n",
    "    component_img = np.array(np.reshape(A_dense[:,idx_comps[i_comp]], dims, order='F'))\n",
    "    plt.subplot(len(idx_comps),2,counter)\n",
    "    counter += 1\n",
    "    plt.imshow(component_img), plt.title('Good component %1.0f' % (i_comp))\n",
    "    \n",
    "    if i_comp == 0:\n",
    "        component_matrix = component_img\n",
    "    else:\n",
    "        component_matrix = np.dstack((component_matrix, component_img))\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Re-classify some components\n",
    "Exclude some 'good' components (select the index, i.e. 0,1,2 as shown in the plot above). These will be added to the list of bad components"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# before re-classification\n",
    "print('Good components: ')\n",
    "print(idx_comps)\n",
    "print('Bad components: ')\n",
    "print(idx_comps_bad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "comps_to_exclude = [1,2,4,5,6] # should be index of the good components (i.e. 0,1,2 as shown in plot above)\n",
    "\n",
    "# add to bad components\n",
    "idx_comps_bad = np.sort(np.append(idx_comps_bad, idx_comps[comps_to_exclude]))\n",
    "# remove from good components\n",
    "idx_comps = np.delete(idx_comps, comps_to_exclude)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# after re-classification\n",
    "print('Good components: ')\n",
    "print(idx_comps)\n",
    "print('Bad components: ')\n",
    "print(idx_comps_bad)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Save and view results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# saving\n",
    "npz_name = os.path.join(data_folder, '%s_%s_Join_%s_results_CNMF.npz' % (date_folder, session_folder, group_id))\n",
    "np.savez(npz_name, Cn=Cn, A=A.todense(), C=C, b=b, f=f, YrA=YrA, sn=sn, \n",
    "         d1=d1, d2=d2, idx_components=idx_comps, idx_components_bad=idx_comps_bad)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Interactive viewer for traces of accepted and rejected components\n",
    "**Currently does not work remotely!**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# accepted components\n",
    "if len(idx_comps) > 0:\n",
    "    nb_view_patches(Yr, A.tocsc()[:, idx_comps], C[idx_comps], \n",
    "                    b, f, dims[0], dims[1], YrA=YrA[idx_comps], image_neurons = Cn,\n",
    "                    denoised_color = 'red');\n",
    "else:\n",
    "    print(\"No accepted components!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculate DFF and plot traces\n",
    "The CaImAn function `detrend_df_f` uses a sliding window percentile filter to determine the baseline and compute DFF. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "F_dff = cnmf_utils.detrend_df_f(A, b, C, f, YrA = YrA, quantileMin=8, frames_window=100)\n",
    "# select good components\n",
    "F_dff = F_dff[idx_comps,:]\n",
    "\n",
    "t = np.arange(0, F_dff.shape[-1]) / params['final_rate']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For plotting we use the Bokeh interactive visualization library, which offers a number of nice features for plotting and interacting with plots."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define some functions for plotting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def getHover():\n",
    "    \"\"\"Define and return hover tool for a plot\"\"\"\n",
    "    # Define hover tool\n",
    "    hover = HoverTool()\n",
    "    hover.tooltips = [\n",
    "        (\"index\", \"$index\"),\n",
    "        (\"(x,y)\", \"($x, $y)\"),\n",
    "    ]\n",
    "    return hover"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plotTimeseries(p, t, y, legend=None, stack=True, xlabel='', ylabel='', output_backend='canvas'):\n",
    "    \"\"\"\n",
    "    Plot a timeseries in Figure p using the Bokeh library\n",
    "    \n",
    "    Input arguments:\n",
    "    p ... Bokeh figure\n",
    "    t ... 1d time axis vector (numpy array)\n",
    "    y ... 2d data numpy array (number of traces x time)\n",
    "    legend ... list of items to be used as figure legend\n",
    "    stack ... whether to stack traces or nor (True / False)\n",
    "    xlabel ... label for x-axis\n",
    "    ylabel ... label for y-axis\n",
    "    output_backend ... 'canvas' or 'svg'\n",
    "    \"\"\"\n",
    "    \n",
    "    colors_list = ['red', 'green', 'blue', 'yellow', 'cyan', 'orange', 'magenta', 'black', 'gray']\n",
    "    p.add_tools(CrosshairTool(), getHover())\n",
    "    \n",
    "    offset = 0\n",
    "    for i in range(y.shape[0]):\n",
    "        if len(colors_list) < i+1:\n",
    "            colors_list = colors_list + colors_list\n",
    "        \n",
    "        plot_trace = y[i, :]\n",
    "        if stack:\n",
    "            plot_trace = plot_trace - min(plot_trace) + offset\n",
    "            offset = max(plot_trace)\n",
    "        \n",
    "        p.line(t, plot_trace, line_width=2, legend=legend[i], color=colors_list[i])\n",
    "        \n",
    "#     p.legend.location = (0,-30)\n",
    "    p.legend.click_policy=\"hide\"\n",
    "    \n",
    "    # format plot\n",
    "    p.xaxis.axis_label = xlabel\n",
    "    p.yaxis.axis_label = ylabel\n",
    "    \n",
    "    p.x_range = Range1d(np.min(t), np.max(t))\n",
    "    \n",
    "    p.background_fill_color = None\n",
    "    p.border_fill_color = None\n",
    "    \n",
    "    p.output_backend = output_backend\n",
    "\n",
    "    show(p)\n",
    "    \n",
    "    return p"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next cell plots the figure. The function `plotTimeseries` can plot traces on top of each other (`stack=False`) or stacked (`stack=True`). The interactive toolbar on the right of the figures allows panning, zooming, saving etc. One can also hide traces by clicking the corresponding legend item. To save the figure, click the disk icon in the plotting toolbar. With the default `output_backend` ('canvas'), a png file will be saved. To save to svg format, change `output_backend` to 'svg'."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = Figure(plot_width=800, plot_height=300, \n",
    "           title=('%s %s CNMF Results' % (date_folder, session_folder)))    \n",
    "legend_text = ['Component %1d' % (x) for x in range(len(idx_comps))]\n",
    "# this is the call to the plotting function (change args. as required)\n",
    "plotTimeseries(p, t, F_dff, \n",
    "               legend=legend_text, \n",
    "               stack=True, \n",
    "               xlabel='Time [s]', ylabel='DFF',\n",
    "               output_backend='canvas'\n",
    "              )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Split up by trials and save as .mat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# split up F_dff by trials\n",
    "# Todo: correct source_frames: key\n",
    "source_files = meta['source_file']\n",
    "source_frames = np.array(meta['source_frames'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check if our numbers match\n",
    "if not (np.sum(source_frames)-len(bad_frames)) == F_dff.shape[-1]:\n",
    "    raise Exception('Sum of source frames minus number of bad frames must be equal to number of timepoints.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "results_dff = dict()\n",
    "removed_frames = dict()\n",
    "bad_frames_removed = []\n",
    "start_frame = 0\n",
    "for ix, trial_file in enumerate(source_files):\n",
    "    trial_indices = list(range(start_frame, source_frames[ix] + start_frame))\n",
    "#     print(trial_file)\n",
    "#     print('Start / stop: %1.0f / %1.0f' % (trial_indices[0], trial_indices[-1]))\n",
    "    \n",
    "    # check if trial contains bad frames (that have not been removed yet)\n",
    "    bad = [x for x in bad_frames if x in trial_indices and x not in bad_frames_removed]\n",
    "    \n",
    "    # remove indices of bad frames\n",
    "    if len(bad) > 0:\n",
    "#         print('Found %1.0f bad frames' % (len(bad)))\n",
    "        del trial_indices[-len(bad):]\n",
    "#         print('Removed frames (relative to trial start):')\n",
    "        for i_bad in bad:\n",
    "            bad_frames_removed.append(i_bad)\n",
    "        removed_frames_trial = [x-trial_indices[0] for x in bad]\n",
    "    else:\n",
    "        removed_frames_trial = []\n",
    "    \n",
    "#     print('Start / stop: %1.0f / %1.0f\\n' % (trial_indices[0], trial_indices[-1]))\n",
    "    \n",
    "    # create a valid Matlab variable / field name\n",
    "    field_name = str('x' + source_files[ix]).replace('.tif','').replace('-','_')\n",
    "    results_dff[field_name] = F_dff[:,trial_indices]\n",
    "    removed_frames[field_name] = removed_frames_trial\n",
    "    \n",
    "    start_frame = trial_indices[-1] + 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# prepare the dictionary for saving as mat file\n",
    "# the field names will be the variable names in Matlab\n",
    "mdict={'trials': [str(x) for x in source_files], \n",
    "       'dff_trial': results_dff, \n",
    "       'removed_frames': removed_frames,\n",
    "       'mean_image': avg_img,\n",
    "       'spatial_components': component_matrix\n",
    "      }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save the .mat file\n",
    "matfile_name = os.path.join(data_folder, '%s_%s_Join_%s_results_CNMF.mat' % (date_folder, session_folder, group_id))\n",
    "savemat(os.path.join(data_folder, matfile_name), mdict=mdict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "caiman",
   "language": "python",
   "name": "caiman"
  },
  "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.5.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}