Hidden-Markury Notebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " <font size=\"+20\">Hidden-Markury</font>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hidden-Markov-Model Fitting for FRET-Traces\n",
    "The Hidden-Markury Notebook is made for the analysis of single-pair FRET-traces of immobilized molecules. It is divided into 10 sections to  load and preprocess the traces, generate and fit the Hidden-Markov-Model (HMM), and evaluate the prediction and fitting results:\n",
    "\n",
    "0. [Check Requirements and Import Packages](#section_0)\n",
    "1. [Load and Preprocess Data](#section_1)\n",
    "2. [Trace Viewer](#section_2)  <font color=\"gray\"><i>(optional)</i></font>\n",
    "3. [Select Fit Model](#section_3)\n",
    "4. [Generate Guess and Fit States](#section_4)\n",
    "5. [Specify HMM and Create Initials](#section_5)\n",
    "6. [Generate HMM and Fit Data](#section_6)\n",
    "7. [Plot Prediction Results](#section_7)   <font color=\"gray\"><i>(optional)</i></font>\n",
    "8. [Lifetime Analysis](#section_8)   <font color=\"gray\"><i>(optional)</i></font>\n",
    "9. [Vary Fitting (Subset of Traces) and Estimate Error](#section_9)   <font color=\"gray\"><i>(optional)</i></font>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\">\n",
    "<b>Tip:</b> You should follow the order of all sections, however, some sections can be skipped (optional).\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "***\n",
    "<a id=\"section_0\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Check Requirements and Import Packages\n",
    "\n",
    "This line checks if all libraries are installed and loads all required packages at the beginning. It returns feedback on which packages to install."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import load_packages\n",
    "from load_packages import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"section_1\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Load and Preprocess Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When loading the data, you have to adapt the code below according to your data and your file format. There is a built-in `load_trace` function, which can import traces with time-stamps in the first column/row and either the photon data (donor channel/acceptor channel) or FRET efficiencies in one of the other columns/rorws.\n",
    "\n",
    "Here is an example of an data-format and the required specifications below to load the data properly.\n",
    "\n",
    "The files are of the form `trace_*.txt` with tab-delimiter in one folder `traces`:\n",
    "```\n",
    "%t (s)\tIdd (a.u.)\tIda (a.u.)\tIaa (a.u.)\tFRET E\n",
    "0.000\t+2.113e+03\t+5.966e+03\t+6.828e+03\t+0.738\n",
    "0.200\t+5.413e+03\t+1.918e+03\t+6.286e+03\t+0.262\n",
    "0.400\t+5.120e+03\t+1.690e+03\t+7.507e+03\t+0.248\n",
    "0.600\t+6.331e+03\t+1.856e+03\t+8.474e+03\t+0.227\n",
    "0.800\t+6.492e+03\t+6.661e+02\t+9.783e+03\t+0.093\n",
    "```\n",
    "\n",
    "\n",
    "```python\n",
    "#files\n",
    "base_folder = r'C:\\Users\\Musterfrau\\examples\\sim_190628_100337_level1'\n",
    "delimiter = '\\t'\n",
    "file_tag = \"trace_\"\n",
    "# data columns/rows\n",
    "DD_col = 1\n",
    "DA_col = 2\n",
    "E_col  = 4\n",
    "```\n",
    "If you do not have the data of the donor (DD) / acceptor channel (DA) or the FRET efficiency (E), set the corresponding variable `<X>_col = -1`.\n",
    "\n",
    "If you have a file format, which is not covered by this imput function, you need to modify the import."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\">\n",
    "<b>Tip:</b> Check out the example folder first, before modifying the import to your file format.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# === DATA PARAMETER ===\n",
    "base_folder = r'ABSOLUT_PATH_TO_FODLER'\n",
    "delimiter = '\\t'\n",
    "file_tag = \"trace_\"\n",
    "\n",
    "# data columns/rows\n",
    "DD_col = 1\n",
    "DA_col = 2\n",
    "E_col  = -1\n",
    "\n",
    "scale_photon_counts = 0.001 # scale photon counts, e.g. change from cps to kps\n",
    "# === DATA PARAMETER ===\n",
    "\n",
    "# load traces (modify this part if needed)\n",
    "print(\"Folder: \", base_folder)\n",
    "files = oh.list_selected_files(base_folder,tag=file_tag)\n",
    "print(\"-- %d traces found --\" % len(files))\n",
    "traces_DD = []\n",
    "traces_DA = []\n",
    "traces_E = []\n",
    "dT = -1\n",
    "len_tot = 0\n",
    "for trace_file in files:\n",
    "    T, DD, DA, E, _ = hh.load_trace(trace_file,DD_col,DA_col,E_col,delimiter=delimiter)\n",
    "    if dT==-1:\n",
    "        dT = np.mean(np.array(T[1:])-np.array(T[:-1]))\n",
    "    else:\n",
    "        pass\n",
    "    if len(DD)>0:\n",
    "        traces_DD.append(np.array([T,DD*scale_photon_counts]))\n",
    "    if len(DA)>0:\n",
    "        traces_DA.append(np.array([T,DA*scale_photon_counts]))\n",
    "    if len(E)>0:\n",
    "        traces_E.append(np.array([T,E]))\n",
    "    len_tot += len(T)\n",
    "print(\"Total length:  \"+str(len_tot))\n",
    "print(\"Time interval: \"+str(dT))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At this point, you should have loaded the traces for donor and acceptor channel + efficiencies.\n",
    "The variables `traces_DD` + `traces_DA`, and/or `traces_E` should be a list of np.array pairs of the form\n",
    "\n",
    "`traces_<X> = [trace1, trace2, trace3, ...]`\n",
    " \n",
    " with `trace<N>` of the form\n",
    " \n",
    " `np.array([np.array([t1, t2, t3, ...]),np.array([value1, value2, value3, ...])])`\n",
    " \n",
    " where `t<N>` are the timestamps of the trace and `value<N>` the corresponding measurement values."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"section_2\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Trace Viewer\n",
    "\n",
    "Run cell and click `Next` or `Previous` to iterate through traces."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "%load_ext disable_autoscroll\n",
    "   \n",
    "# === PLOTTING ===\n",
    "N_traces = 3   #number of shown traces (at a time)\n",
    "trace_idx = 0  #start trace\n",
    "\n",
    "out=widgets.Output()\n",
    "info = widgets.Output()\n",
    "\n",
    "# Plot first traces\n",
    "with out:\n",
    "    hh.plot_traces(N_traces,trace_idx,traces_DD,traces_DA,traces_E)\n",
    "with info:\n",
    "    print(\"Traces {0:d} - {1:d}\".format(trace_idx,trace_idx+N_traces-1))    \n",
    "\n",
    "# Arrange output (buttons + text + plot)\n",
    "b_prev=widgets.Button(description='Previous')\n",
    "b_next=widgets.Button(description='Next')\n",
    "hbox = widgets.HBox(children=(b_prev,b_next,info))\n",
    "vbox=widgets.VBox(children=(hbox,out))\n",
    "display(vbox)\n",
    "\n",
    "# Define button events (recalculate trace plots)\n",
    "def click(direction):\n",
    "    global trace_idx\n",
    "    if direction=='next':\n",
    "        trace_idx += N_traces\n",
    "    elif direction=='back':\n",
    "        trace_idx -= N_traces\n",
    "    else:\n",
    "        raise Exception(\"Command not defined\")\n",
    "    with out:\n",
    "        clear_output()\n",
    "        hh.plot_traces(N_traces,trace_idx,traces_DD,traces_DA,traces_E)\n",
    "    with info:\n",
    "        clear_output(wait=True)\n",
    "        print(\"Traces {0:d} - {1:d}\".format(trace_idx,trace_idx+N_traces-1))\n",
    "def click_p(b):\n",
    "    click('back')\n",
    "def click_n(b):\n",
    "    click('next')    \n",
    "b_prev.on_click(click_p)\n",
    "b_next.on_click(click_n)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"section_3\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Select Fit Model\n",
    "Select model to analyse your data (number of states and analysis of E-trace or DD/DA-traces). Set `mode = '2D'` for combined donor & acceptor trace analysis or `mode = '1D'` for FRET efficiency trace analysis.\n",
    "\n",
    "Select peak positions of the states by clicking on the peaks in the iterative plot:\n",
    "\n",
    "<img src=\"images\\2D_FitSelection.png\" width=\"240\" align=\"left\">\n",
    "\n",
    "#### `2D` for combined donor & acceptor trace analysis \n",
    "Click on the centers of the populations to give a first estimate for fitting the states. The values are stored in variable `pos`. In this plot for example:\n",
    "\n",
    "`pos = [[0.81, 6.6, 303], [3.8, 4.0, 228], [7.2, 0.5, 105]]`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"images\\1D_FitSelection.png\" width=\"240\" align=\"left\">\n",
    "\n",
    "#### `1D` for FRET efficiency trace analysis\n",
    "Click on the peaks of the states to give a first estimate for fitting the states. The values are stored in variable `pos`. In this plot for example:\n",
    "\n",
    "`pos = [[0.11, 2152], [0.55, 4158], [0.85, 5263]]`\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For correcting the estimate, click `Clear`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-warning\">\n",
    "<b>Warning:</b> The switch from inline to notebook backend fails sometimes unpredictable. If no plot appears, run cell again until interactive plot appears.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# === SELECT MODE ===\n",
    "\n",
    "mode = '2D' \n",
    "\n",
    "# '2D' for combined donor & acceptor trace analysis\n",
    "# '1D' for FRET efficiency trace analysis\n",
    "# === SELECT MODE ===\n",
    "\n",
    "#Switching from inline to notebook does not always work (see link)\n",
    "#https://stackoverflow.com/questions/43545050/using-matplotlib-notebook-after-matplotlib-inline-in-jupyter-notebook-doesnt\n",
    "for i in range(5):\n",
    "    %matplotlib notebook\n",
    "    time.sleep(0.2)\n",
    "\n",
    "#Validate selection\n",
    "if mode == '2D':\n",
    "    assert len(traces_DD)>0 and len(traces_DA)==len(traces_DD) , \"Mode does not work with type of imported data - no photon count trace for donor and acceptor loaded\"\n",
    "elif mode == '1D':\n",
    "    assert len(traces_E)>0 , \"Mode does not work with type of imported data - no FRET efficiency trace loaded\"\n",
    "else:\n",
    "    assert -1>0 , \"Mode '\" + mode + \"' not defined\"\n",
    "\n",
    "# === PLOTTING ===\n",
    "#%load_ext disable_autoscroll\n",
    "fig = plt.figure(figsize=(5,5))\n",
    "ax = fig.add_subplot(111)\n",
    "if mode=='2D':\n",
    "    scatter_DD = traces_DD[0][1]\n",
    "    scatter_DA = traces_DA[0][1]\n",
    "    for trace_DD, trace_DA in zip(traces_DD[1:],traces_DA[1:]):\n",
    "        scatter_DD = np.append(scatter_DD, trace_DD[1])\n",
    "        scatter_DA = np.append(scatter_DA, trace_DA[1])\n",
    "    hist_data, xbins, ybins = np.histogram2d(scatter_DD, scatter_DA, bins=50)\n",
    "    DD_min = min([min(t[:,1]) for t in traces_DD])\n",
    "    DD_max = max([max(t[:,1]) for t in traces_DD])\n",
    "    DA_min = min([min(t[:,1]) for t in traces_DA])\n",
    "    DA_max = max([max(t[:,1]) for t in traces_DA])\n",
    "    \n",
    "    ax.hist2d(scatter_DD, scatter_DA, bins=(50, 50), cmap=plt.cm.jet)\n",
    "    try:\n",
    "        idx = np.log10(scale_photon_counts)/3\n",
    "        print\n",
    "        if idx.is_integer():\n",
    "            cts = \"[\" + [\"\",\"m\",\"mu\",\"M\",\"k\"][int(idx)] + \"cps]\"\n",
    "        else:\n",
    "            cts = \"[a.u.]\"\n",
    "    except:\n",
    "        cts = \"[a.u.]\"\n",
    "    ax.set_xlabel(\"DD \"+cts)\n",
    "    ax.set_ylabel(\"DA \"+cts)\n",
    "else:\n",
    "    long_trace_E = traces_E[0][1]\n",
    "    for trace_E  in traces_E[1:]:\n",
    "        long_trace_E = np.append(long_trace_E, trace_E[1])\n",
    "    bins=np.linspace(-0.2,1.2,num=50)\n",
    "    print(min(long_trace_E))\n",
    "    print(max(long_trace_E))\n",
    "    ax.hist(long_trace_E, bins=bins)\n",
    "    hist_data, bins_h1D =np.histogram(long_trace_E, bins=bins)\n",
    "    max_hist = max(hist_data)\n",
    "    ax.set_xlabel(\"E\")\n",
    "    ax.set_ylabel(\"Events\")\n",
    "    ax.set_xlim(xmin=-0.2, xmax = 1.2)\n",
    "   \n",
    "pos = []\n",
    "def onclick(event):\n",
    "    if mode=='2D':\n",
    "        x_pos = np.abs(xbins - event.xdata).argmin()\n",
    "        y_pos = np.abs(ybins - event.ydata).argmin()\n",
    "        z = hist_data[x_pos,y_pos]\n",
    "        pos.append([event.xdata,event.ydata,z])\n",
    "    else:\n",
    "        pos.append([event.xdata,event.ydata])\n",
    "    with pos_out:\n",
    "        clear_output(wait=True)\n",
    "        print(\"Selected positions: \",pos)   \n",
    "cid = fig.canvas.mpl_connect('button_press_event', onclick)\n",
    "out=widgets.Output()\n",
    "pos_out = widgets.Output()\n",
    "button=widgets.Button(description='Clear')\n",
    "vbox=widgets.VBox(children=(out,pos_out,button))\n",
    "display(vbox)\n",
    "\n",
    "with pos_out:\n",
    "    print(\"Selected positions: \",pos)\n",
    "    \n",
    "# button to clear selection\n",
    "def click(b):\n",
    "    global pos\n",
    "    pos = []\n",
    "    with pos_out:\n",
    "        clear_output(wait=True)\n",
    "        print(\"Selected positions: \",pos)\n",
    "        #text.set_text(\"\")\n",
    "button.on_click(click)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"section_4\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Generate Guess and Fit States\n",
    "Funchtion `guess` generates a guess for fitting based on the selection of the previous cell. You can override the variable `guess`, if the fit in the following cell fails.\n",
    "\n",
    "For `2D`-mode, `guess`is of the form\n",
    "```python\n",
    "guess = [A1, mu_DD1, sigma_DD1, mu_DA1, sigma_DA1, rho1, A2, mu_DD2, sigma_DD2, mu_DA2, sigma_DA2, rho2, ...]\n",
    "```\n",
    "For `1D`-mode, `guess`is of the form\n",
    "```python\n",
    "guess = [A1, mu_E1, sigma_E1, A2, mu_E2, sigma_E2, ...]\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "guess = []\n",
    "if mode=='2D':\n",
    "    hist_data, xbins, ybins = np.histogram2d(scatter_DD, scatter_DA, bins=50)\n",
    "    hist_data_ravel = hist_data.T.ravel()\n",
    "    x_fit, y_fit = np.meshgrid(xbins[:-1]+(xbins[1]-xbins[0])/2, ybins[:-1]+(ybins[1]-ybins[0])/2)\n",
    "    for p in pos:\n",
    "        guess += [p[2], p[0], (xbins[-1]-xbins[0])/8., p[1], (ybins[-1]-ybins[0])/10.,0]\n",
    "    GAUSS_N = 6\n",
    "else:\n",
    "    bins=np.linspace(-0.2,1.2,num=50)\n",
    "    hist_data, bins_h1D =np.histogram(long_trace_E, bins=bins)\n",
    "    x_fit = np.linspace(bins_h1D[0]+(bins_h1D[1]-bins_h1D[0])/2,bins_h1D[-1]-(bins_h1D[1]-bins_h1D[0])/2,num=len(bins)-1)\n",
    "    for p in pos:\n",
    "        guess += [p[1], p[0], 0.1]\n",
    "    GAUSS_N = 3\n",
    "print(\"guess = \"+str(guess))\n",
    "hh.print_gaussians(guess, mode)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "# === FITTING ===\n",
    "if mode=='2D':\n",
    "    popt_init, pcov = curve_fit(hh.gauss_2D, (x_fit.ravel(),y_fit.ravel()), hist_data_ravel, p0=guess)\n",
    "else:\n",
    "    popt_init, pcov = curve_fit(hh.gauss_1D, x_fit, hist_data, p0=guess)\n",
    "popt_init = [(val if (idx%5!=2 and idx%5!=4) else abs(val)) for idx, val in enumerate(popt_init)]\n",
    "print(\"popt_init = \",str(popt_init))\n",
    "hh.print_gaussians(popt_init,mode)\n",
    "\n",
    "# === PLOTTING ===\n",
    "if mode=='2D':\n",
    "    fig, axes = plt.subplots(1, 3, figsize=(15,5))\n",
    "    axes[0].hist2d(scatter_DD, scatter_DA, bins=(50, 50), cmap=plt.cm.jet)\n",
    "\n",
    "    data_fitted = hh.gauss_2D((x_fit.ravel(), y_fit.ravel()), *popt_init)\n",
    "    max_hist = np.max(hist_data)\n",
    "    axes[0].contour(xbins[:-1]+(xbins[1]-xbins[0])/2,ybins[:-1]+(ybins[1]-ybins[0])/2, data_fitted.reshape(50, 50), levels=np.array([0.05,0.1,0.2,0.4,0.7])*max_hist, colors='w')\n",
    "\n",
    "    x_ax = np.linspace(xbins[0], xbins[-1], num=500)\n",
    "    y_ax = np.linspace(ybins[0], ybins[-1], num=500)\n",
    "    para_x = list(np.ravel([[sigma2*A, mu, sigma] for A,mu,sigma,mu2,sigma2,rho in hh.grouped_list(popt_init, GAUSS_N)]))\n",
    "    x_hist = hh.gauss_1D(x_ax,*para_x)\n",
    "    para_y = list(np.ravel([[sigma*A, mu2, sigma2] for A,mu,sigma,mu2,sigma2,rho in hh.grouped_list(popt_init, GAUSS_N)]))\n",
    "    y_hist = hh.gauss_1D(y_ax,*para_y)\n",
    "    axes[0].set_xlabel(\"DD [kcps]\")\n",
    "    axes[0].set_xlabel(\"DA [kcps]\")\n",
    "    axes[1].hist(scatter_DD, bins=50)\n",
    "    axes[1].plot(x_ax, 4*np.pi*x_hist)\n",
    "    axes[1].set_xlabel(\"DD [kcps]\")\n",
    "    #axes[1].set_xlabel(\"DA [kcps]\")\n",
    "    axes[2].hist(scatter_DA, bins=50)\n",
    "    axes[2].plot(y_ax, 4*np.pi*y_hist)\n",
    "    axes[2].set_xlabel(\"DA [kcps]\")\n",
    "    #axes[2].set_xlabel(\"DA [kcps]\")\n",
    "else:\n",
    "    fig, axes = plt.subplots(1, 1, figsize=(5,5))\n",
    "    bins=np.linspace(-0.2,1.2,num=50)\n",
    "    axes.hist(long_trace_E, bins=bins)\n",
    "    E_fit = hh.gauss_1D(np.linspace(-0.2,1.2,1000),*popt_init)\n",
    "    axes.plot(np.linspace(-0.2,1.2,1000),E_fit)\n",
    "    axes.set_xlabel(\"E\")\n",
    "    axes.set_ylabel(\"Events\")\n",
    "    axes.set_xlim(xmin=-0.2, xmax = 1.2)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"section_5\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Specify Hidden-Markov-Model (HMM) and Create Initials\n",
    "Here, you have to select some general parameter of the HMM, before initializing the model and fit the trace data:\n",
    "1. Set `degenerated_states = True` if you expect degenerated states \n",
    "2. If there are degenerated states, specify degeneration for each state: e.g. `degeneration = [2,1,1]` means that the first state is degenerated, the others not\n",
    "3. If ther is no degenerated state, specify if you want to optimize the states durch the HMM-fitting or not (for degenerated states, the states have to be fixed)\n",
    "\n",
    "The function `generate_hmm_initials` pre-calculates initial values for the states (mean & covariance), the transition matrix and the start probability of each state. You can override the values, if the fit fails or you have a better guess."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# === SPECIFY HMM ===\n",
    "degenerated_states = False\n",
    "degeneration = [1,1]\n",
    "fix_states = True #for degenerated states, states are always fixed (assignment is ignored)\n",
    "# == SPECIFY HMM ===\n",
    "\n",
    "if degenerated_states: \n",
    "    assert len(degeneration) < np.sum(degeneration), \"Degenerated states does not fit total number of states\"\n",
    "    nbr_states = np.sum(degeneration)\n",
    "    fix_states = True\n",
    "else:\n",
    "    nbr_states = int(len(popt_init)/GAUSS_N)\n",
    "    degeneration = [1 for i in range(nbr_states)]\n",
    "                \n",
    "means, covs, trans_mat, starts = hh.generate_hmm_initials(popt_init, degeneration, mode)\n",
    "\n",
    "means, covs, trans_mat, starts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\">\n",
    "<b>Tip:</b> For complex systemes (e.g. many states), you might want to specify the initial transition matrix manually and set some transitions zero.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-warning\">\n",
    "<b>Info:</b> Adapt the matrix dimensions/entries and run this cell only if the initial guess does not converge to a adequate solution!.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# OVERRIDE INITIAL GUESSES AND SET \"MANUEL=True\" ONLY IF NEEDED\n",
    "MANUAL = False\n",
    "if MANUAL:\n",
    "    #trans_mat = np.array([[0.96, 0.02,0.01,0.01],\n",
    "    #                      [0.05,0.7,0.24,0.01],\n",
    "    #                      [0.01,0.25,0.73,0.01],\n",
    "    #                      [0.02,0.01,0.01,0.96]])\n",
    "    trans_mat = np.array([[0.90, 0.10],\n",
    "                          [0.08,0.92]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"section_6\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Generate HMM and Fit Data\n",
    "This cell contains the important steps:\n",
    "1. The HMM is created as `model` and initialised with the settings from above\n",
    "2. The data are prepared for fitting (one long trace, length specification of individual traces)\n",
    "3. The model is fit\n",
    "4. The states are predicted based on the fitted model\n",
    "5. The prediction results are added to the traces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Several different starting points are recommended (gradient based approach)\n",
    "\n",
    "\"\"\"\n",
    "# === CREATE MODEL ===\n",
    "if fix_states or degenerated_states:\n",
    "    model = hmm.GaussianHMM(n_components=nbr_states, covariance_type=\"full\", init_params=\"\",params=\"ts\", n_iter=100)\n",
    "else:\n",
    "    model = hmm.GaussianHMM(n_components=nbr_states, covariance_type=\"full\", init_params=\"\",params=\"mcts\", n_iter=100)\n",
    "model.startprob_ = np.array(starts)\n",
    "model.transmat_ = np.array(trans_mat)\n",
    "model.means_ = np.array(means)\n",
    "model.covars_ = np.array(covs)\n",
    "\n",
    "# === FIT MODEL ===\n",
    "#prepare data\n",
    "if mode=='2D':\n",
    "    traces_y = [np.array([trace_DD[1], trace_DA[1]]).T for trace_DD, trace_DA in zip(traces_DD,traces_DA)]\n",
    "else:\n",
    "     traces_y = [np.array([trace_E[1]]).T for trace_E in traces_E]\n",
    "traces_all = np.concatenate(traces_y)\n",
    "lengths = [len(t) for t in traces_y]\n",
    "\n",
    "#fit model and calculate rates from transition matrix\n",
    "fit_result = model.fit(traces_all, lengths)\n",
    "rates = np.array([(1-model.transmat_[i,i])/dT for i in range(len(model.transmat_))])\n",
    "\n",
    "#predict values\n",
    "predictions = model.predict(traces_all, lengths)\n",
    "\n",
    "#append fitted values to traces\n",
    "len_cum = 0\n",
    "traces_DD_fit = []\n",
    "traces_DA_fit = []\n",
    "traces_E_fit = []\n",
    "prediction_traces = []\n",
    "for idx, length in enumerate(lengths):\n",
    "    prediction = predictions[len_cum:len_cum+length]\n",
    "    prediction_traces.append(np.array(prediction))\n",
    "    if mode=='1D':\n",
    "        E_fit = np.array([means[s][0] for s in  prediction])\n",
    "        traces_E_fit.append(np.array([traces_E[idx][0],traces_E[idx][1],E_fit]))\n",
    "    if mode=='2D':\n",
    "        DD_fit = np.array([means[s][0] for s in  prediction])\n",
    "        DA_fit = np.array([means[s][1] for s in  prediction])\n",
    "        traces_DD_fit.append(np.array([traces_DD[idx][0],traces_DD[idx][1],DD_fit]))\n",
    "        traces_DA_fit.append(np.array([traces_DA[idx][0],traces_DA[idx][1],DA_fit]))\n",
    "        E_fit = np.divide(DA_fit,(DA_fit+DD_fit))\n",
    "        if len(traces_E)>0:\n",
    "            traces_E_fit.append(np.array([traces_E[idx][0],traces_E[idx][1],E_fit]))        \n",
    "    len_cum += length\n",
    "    \n",
    "model.monitor_.converged, fit_result, model.monitor_, model.means_, model.transmat_, rates"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"section_7\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Plot Prediction Results\n",
    "\n",
    "Run cell and click `Next` or `Previous` to iterate through traces."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "%load_ext disable_autoscroll\n",
    "reload(hh)\n",
    "# === PLOTTING ===\n",
    "N_traces = 3   #number of shown traces (at a time)\n",
    "trace_idx = 0  #start trace\n",
    "\n",
    "out=widgets.Output()\n",
    "info = widgets.Output()\n",
    "\n",
    "# Plot first traces\n",
    "with out:\n",
    "    hh.plot_traces(N_traces,start=trace_idx, traces_DD=traces_DD_fit, traces_DA=traces_DA_fit,traces_E=traces_E_fit,prediction=True)\n",
    "with info:\n",
    "    print(\"Traces {0:d} - {1:d}\".format(trace_idx,trace_idx+N_traces-1))    \n",
    "\n",
    "# Arrange output (buttons + text + plot)\n",
    "b_prev=widgets.Button(description='Previous')\n",
    "b_next=widgets.Button(description='Next')\n",
    "hbox = widgets.HBox(children=(b_prev,b_next,info))\n",
    "vbox=widgets.VBox(children=(hbox,out))\n",
    "display(vbox)\n",
    "\n",
    "# Define button events (recalculate trace plots)\n",
    "def click(direction):\n",
    "    global trace_idx\n",
    "    if direction=='next':\n",
    "        trace_idx += N_traces\n",
    "    elif direction=='back':\n",
    "        trace_idx -= N_traces\n",
    "    else:\n",
    "        raise Exception(\"Command not defined\")\n",
    "    with out:\n",
    "        clear_output()\n",
    "        hh.plot_traces(N_traces,start=trace_idx, traces_DD=traces_DD_fit, traces_DA=traces_DA_fit,traces_E=traces_E_fit,prediction=True)\n",
    "    with info:\n",
    "        clear_output(wait=True)\n",
    "        print(\"Traces {0:d} - {1:d}\".format(trace_idx,trace_idx+N_traces-1))\n",
    "def click_p(b):\n",
    "    click('back')\n",
    "def click_n(b):\n",
    "    click('next')    \n",
    "b_prev.on_click(click_p)\n",
    "b_next.on_click(click_n)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"section_8\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Lifetime Analysis\n",
    "\n",
    "Dwell times are extracted from prediction results, where the last state is discarded due to bleaching. Bin widths are automatically generated based on averaged dwell times. Exponential fit is performed with 4 slightly different bin sizes and averaged."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "#estract dwelltimes from traces\n",
    "lifetime_traces = [hh.split_trace(pred, nbr_max=(nbr_states-1),skip_last=True) for pred in prediction_traces]\n",
    "lifetimes_all = hh.cumulate_lifetimes(lifetime_traces)\n",
    "lifetime_states = [np.array(l)*dT for l in lifetimes_all]\n",
    "\n",
    "np.seterr(over='ignore')\n",
    "fits_all = []\n",
    "for rnd, bin_div in enumerate([1.2,1.5,2,2.5]): #do lifetime analysis for different bin sizes\n",
    "    #define individual bin-width and range for diffent states\n",
    "    dTbin = []\n",
    "    s_T = 0\n",
    "    e_T = []\n",
    "    for i in range(nbr_states):\n",
    "        mean = np.mean(lifetime_states[i])\n",
    "        dTbin.append(mean/bin_div)\n",
    "        e_T.append(20*dTbin[i])\n",
    "    print(\"-- bin widths (states {0:d}-{1:d}) --\".format(1,nbr_states))\n",
    "    print(dTbin)\n",
    "    xs = [np.linspace(s_T+dTbin[i]/2, e_T[i]-dTbin[i]/2, num=int(np.round((e_T[i]-s_T)/dTbin[i]))) for i in range(nbr_states)]\n",
    "    xs_fit =[np.linspace(s_T+dTbin[i]/2, e_T[i]-dTbin[i]/2, num=1000) for i in range(nbr_states)]\n",
    "    bins = [np.linspace(s_T, e_T[i], num=int(np.round((e_T[i]-s_T)/dTbin[i])+1)) for i in range(nbr_states)]\n",
    "\n",
    "    #fit exponential and plot\n",
    "    hist_T = []\n",
    "    fits = []\n",
    "    fit_res = []\n",
    "    if rnd==1:\n",
    "        fig, axes = plt.subplots(1,nbr_states,figsize=(15,10))\n",
    "        fig.text(0.5, 0.04, 'Time, t [s]', ha='center')\n",
    "        fig.text(0.04, 0.5, 'Events', va='center', rotation='vertical')\n",
    "    for i in range(nbr_states):\n",
    "        y,_ = np.histogram(lifetime_states[i], bins=bins[i])\n",
    "        hist_T.append(y)\n",
    "        guess=[max(y),0.5]\n",
    "        popt, pcov = curve_fit(hh.func_exp, xs[i], y, p0=guess)\n",
    "        fit = hh.func_exp(xs_fit[i], *popt)\n",
    "        fits.append(fit)\n",
    "        fit_res.append(popt)\n",
    "        if rnd==1:\n",
    "            axes[i].hist(lifetime_states[i], bins=bins[i])\n",
    "            axes[i].plot(xs_fit[i], fit)\n",
    "    fits_all.append(fit_res)\n",
    "\n",
    "for i, taus in enumerate(np.array(fits_all).T[1]):    \n",
    "    axes[i].annotate(r'$\\tau={0:.2f} \\pm {1:.2f}$'.format(np.mean(1/taus),np.std(1/taus)), xy=(0.45, 0.8), xycoords='axes fraction', fontsize=12,\n",
    "                horizontalalignment='left', verticalalignment='bottom')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"section_9\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. Vary Fitting (Subset of Traces) and Estimate Error\n",
    "The fit quality of the model is estimated by variation of the imput data. The variable `rounds = 20` sets the number of variations and `selection = 0.6` specifies the size of the subset relative to the total number of traces.\n",
    "\n",
    "Note that multiple occurence of one trace in the subset is possible (sampling with replacement)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# === VARIATION PARAMETER ===\n",
    "#prepare shuffle\n",
    "rounds = 20      # number of variation rounds\n",
    "selection = 0.4  # ratio of selection size compared to total number of traces\n",
    "# === VARIATION PARAMETER ===\n",
    "\n",
    "# === CREATE MODEL ===\n",
    "if fix_states or degenerated_states:\n",
    "    model = hmm.GaussianHMM(n_components=nbr_states, covariance_type=\"full\", init_params=\"\",params=\"ts\", n_iter=100)\n",
    "else:\n",
    "    model = hmm.GaussianHMM(n_components=nbr_states, covariance_type=\"full\", init_params=\"\",params=\"mcts\", n_iter=100)\n",
    "model.startprob_ = np.array(starts)\n",
    "model.transmat_ = np.array(trans_mat)\n",
    "model.means_ = np.array(means)\n",
    "model.covars_ = np.array(covs)\n",
    "\n",
    "# === FIT MODEL ===\n",
    "tag_DA_E = \"E\" if mode==\"1D\" else \"DA\"\n",
    "tag_fix = \"fix\" if fix_states else \"free\"\n",
    "tag_sel = str(selection)\n",
    "timestr = time.strftime(\"%Y-%m-%d-%H.%M.%S\")\n",
    "folder_name = \"model_fit_\"  + tag_DA_E + \"_\" + tag_fix + \"_\" + tag_sel + \"_\" + timestr\n",
    "folder_model_fit = os.path.join(base_folder,folder_name)\n",
    "if not os.path.exists(folder_model_fit):\n",
    "    os.makedirs(folder_model_fit)\n",
    "\n",
    "for rnd in range(rounds):\n",
    "\n",
    "    randTraceIdx = np.random.randint(len(traces_DD), size=int(selection*len(traces_DD)))\n",
    "    randTraces_DD = [traces_DD[i] for i in randTraceIdx]\n",
    "    randTraces_DA = [traces_DA[i] for i in randTraceIdx]\n",
    "    randTraces_E = [traces_E[i] for i in randTraceIdx]\n",
    "\n",
    "\n",
    "    #prepare data\n",
    "    if mode=='2D':\n",
    "        traces_y = [np.array([trace_DD[1], trace_DA[1]]).T for trace_DD, trace_DA in zip(randTraces_DD,randTraces_DA)]\n",
    "    else:\n",
    "         traces_y = [np.array([trace_E[1]]).T for trace_E in randTraces_E]\n",
    "    traces_all = np.concatenate(traces_y)\n",
    "    lengths = [len(t) for t in traces_y]\n",
    "\n",
    "    #fit model and calculate rates from transition matrix\n",
    "    fit_result = model.fit(traces_all, lengths)\n",
    "    \n",
    "\n",
    "    name_model = \"model_\" + tag_DA_E + \"_\" + tag_fix + \"_\" + tag_sel + \"_\" + str(rnd) + \".txt\"\n",
    "    path_model = os.path.join(folder_model_fit,name_model)\n",
    "    hh.save_model(model,path_model)\n",
    "    name_tm = \"transition_matrix_\" + tag_DA_E + \"_\" + tag_fix + \"_\" + tag_sel + \"_\" + str(rnd) + \".tsv\"\n",
    "    path_tm = os.path.join(folder_model_fit,name_tm)\n",
    "    np.savetxt(path_tm, model.transmat_)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate Mean and Standard-Deviation of fit runs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load matrices\n",
    "print(\"Folder: \", folder_model_fit)\n",
    "files = oh.list_selected_files(folder_model_fit,tag=\"transition_matrix\")\n",
    "tms = []\n",
    "for file in files:\n",
    "    tms.append(np.loadtxt(file)/dT)\n",
    "tms = np.array(tms)\n",
    "\n",
    "tm_mean = np.mean(tms,axis=0)\n",
    "path = os.path.join(folder_model_fit,\"transition_matix_mean.tsv\")\n",
    "np.savetxt(path, tm_mean)\n",
    "print(\"-- mean transition matrix --\")\n",
    "print(tm_mean)\n",
    "tm_stdDev = np.std(tms,axis=0)\n",
    "path = os.path.join(folder_model_fit,\"transition_matix_stdDev.tsv\")\n",
    "np.savetxt(path, tm_stdDev)\n",
    "print(\"-- standard deviation transition matrix --\")\n",
    "print(tm_stdDev)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}