examples/heterogeneous_data_analysis.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using DiffusionMLE, Plots"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a minimal example on how to apply our maximum likelihood estimator to a set of heterogeneous single-particle tracking data.  Here, \"heterogeneous\" refers to the fact that the data originate from subpopulations with differing diffusion coefficients.  \n",
    "\n",
    "The code relevant for the analysis of heterogeneous data exploits threading, so it is recommended to run the command <code>export JULIA_NUM_THREADS=n</code>, with <code>n</code> being the number of available (physical) cores, before launching Julia.  This speeds up the numerics significantly.  \n",
    "\n",
    "For more details on the theoretical framework, please refer to the associated preprint:\n",
    "> J. T. Bullerjahn and G. Hummer, \"Maximum likelihood estimates of diffusion coefficients from single-molecule tracking experiments\", https://arxiv.org/abs/2011.09955"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate trajectories"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each trajectory can be seen as a $d$-dimensional array (<code>Array{Float64,2}</code>), so the data set should be of the type <code>Array{Array{Float64,2},1}</code>.  \n",
    "\n",
    "Here, we generate mock data, made up of $M$ $d$-dimensional trajectories of different lengths $N = \\{N_{1}, N_{2}, \\dots, N_{M}\\}$.  The $N_{i}$ are distributed uniformly on the interval $[3,100]$.  The data form a $3:4:3$-mixture of trajectories generated using three distinct diffusive dynamics.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "include(\"smeared_trajectory_integrator.jl\")\n",
    "\n",
    "const M = 1000 # Number of trajectories\n",
    "const d = 2 # Dimension of trajectories\n",
    "\n",
    "const N_sub = 100 # Number of substeps over which the trajectory is smeared out\n",
    "\n",
    "N = [ rand(3:100) for i = 1 : M ] # Array of trajectory lengths\n",
    "\n",
    "# Consider three subpopulations, characterized by the following parameters:\n",
    "const a2_1 = 0.5\n",
    "const a2_2 = 2.0\n",
    "const a2_3 = 1.0\n",
    "const σ2_1 = 0.1\n",
    "const σ2_2 = 1.0\n",
    "const σ2_3 = 10.0\n",
    "\n",
    "B = [1/6 for m = 1 : M] # Array of blurring coefficients, where we have assumed a uniform illumination profile\n",
    "data = vcat([make_2D_data(N[1:300],N_sub,a2_1,σ2_1), \n",
    "        make_2D_data(N[301:700],N_sub,a2_2,σ2_2), \n",
    "        make_2D_data(N[701:1000],N_sub,a2_3,σ2_3)]...); # Mock data set"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Analyzing the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function print_results(estimates,uncertainties)\n",
    "    K = size(estimates,2)\n",
    "    for k = 1 : K\n",
    "        println(string(\"a2_\", k, \" = \", estimates[1,k], \" ± \", uncertainties[1,k]))\n",
    "    end\n",
    "    for k = 1 : K\n",
    "        println(string(\"σ2_\", k, \" = \", estimates[2,k], \" ± \", uncertainties[2,k]))\n",
    "    end\n",
    "    for k = 1 : K\n",
    "        println(string(\"P_\", k, \" = \", estimates[3,k]))\n",
    "    end\n",
    "end\n",
    "\n",
    "const N_local = 500 # Max number of expectation-maximization cycles\n",
    "const N_global = 100 # Number of iterations with different initial parameters\n",
    "\n",
    "# Ranges from which the initial values for the parameters are drawn:\n",
    "a2_range = [ 0.01, 20. ]\n",
    "σ2_range = [ 0.01, 20. ];"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Available cores for threading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Base.Threads\n",
    "println(string(\"Number of available cores for threading: \", nthreads()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Assuming a single population"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parameters = MLE_estimator(B,data)\n",
    "parameter_matrix = reshape(vcat([parameters,[1.0]]...), 3, 1)\n",
    "P1_estimates, P1_L, P1_T = local_EM_estimator!(d,M,1,N_local,parameter_matrix,B,data)\n",
    "P1_uncertainties = MLE_errors(B,data,parameters)\n",
    "\n",
    "println(\"Estimates:\")\n",
    "print_results(P1_estimates,P1_uncertainties)\n",
    "\n",
    "Q_sub = subpopulation_analysis(P1_T,P1_estimates,B,data)\n",
    "Q = vcat(Q_sub...)\n",
    "\n",
    "println()\n",
    "\n",
    "println(\"Kuiper statistic:\")\n",
    "println(\"κ = \",Kuiper_statistic!(Q))\n",
    "\n",
    "histogram(Q, normed=true, legend=false, \n",
    "    xlabel=\"quality factor Q\", \n",
    "    ylabel=\"probability distribution function p(Q)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The distribution of quality factors, which should be uniform, and the associated Kuiper statistic reveal that a single-population model does not suffice.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Assuming two subpopulations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "P2_estimates, P2_L, P2_T = global_EM_estimator(2,N_local,N_global,a2_range,σ2_range,B,data);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "B_sub, X_sub = sort_trajectories(2,P2_T,B,data)\n",
    "P2_uncertainties = hcat([ MLE_errors(B_sub[k], X_sub[k], P2_estimates[1:2,k]) for k = 1 : 2 ]...)\n",
    "\n",
    "println(\"Estimates:\")\n",
    "print_results(P2_estimates,P2_uncertainties)\n",
    "\n",
    "Q_sub = subpopulation_analysis(P2_T,P2_estimates,B,data)\n",
    "Q = vcat(Q_sub...)\n",
    "\n",
    "println()\n",
    "\n",
    "println(\"Kuiper statistic:\")\n",
    "println(\"K=\",Kuiper_statistic!(Q))\n",
    "\n",
    "histogram(Q, normed=true, legend=false, bins=10,\n",
    "    xlabel=\"quality factor Q\", \n",
    "    ylabel=\"probability distribution function p(Q)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Kuiper statistic is slightly better, but overall a two-population model cannot explain the data.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Assuming three subpopulations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "P3_estimates, P3_L, P3_T = global_EM_estimator(3,N_local,N_global,a2_range,σ2_range,B,data);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "B_sub, X_sub = sort_trajectories(3,P3_T,B,data)\n",
    "P3_uncertainties = hcat([ MLE_errors(B_sub[k], X_sub[k], P3_estimates[1:2,k]) for k = 1 : 3 ]...)\n",
    "\n",
    "println(\"Estimates:\")\n",
    "print_results(P3_estimates,P3_uncertainties)\n",
    "\n",
    "println()\n",
    "\n",
    "println(\"Ground truth (up to permutations):\")\n",
    "println(string(\"a2_1 = \", a2_1))\n",
    "println(string(\"a2_2 = \", a2_2))\n",
    "println(string(\"a2_3 = \", a2_3))\n",
    "println(string(\"σ2_1 = \", σ2_1))\n",
    "println(string(\"σ2_2 = \", σ2_2))\n",
    "println(string(\"σ2_3 = \", σ2_3))\n",
    "\n",
    "Q_sub = subpopulation_analysis(P3_T,P3_estimates,B,data)\n",
    "Q = vcat(Q_sub...)\n",
    "\n",
    "println()\n",
    "\n",
    "println(\"Kuiper statistic:\")\n",
    "println(\"K=\",Kuiper_statistic!(Q))\n",
    "\n",
    "histogram(Q, normed=true, legend=false,\n",
    "    xlabel=\"quality factor Q\", \n",
    "    ylabel=\"probability distribution function p(Q)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The quality factors are uniformly distributed, and the Kuiper statistic evaluates close to unity for the three-population model.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Diffusion coefficients"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Diffusion coefficients can be extracted from the $\\sigma^2$-values, irrespective of the dimension $d$, as follows:\n",
    "\\begin{equation*}\n",
    "D = \\frac{\\sigma^2}{2 \\Delta t} \\, .  \n",
    "\\end{equation*}\n",
    "Here, $\\sigma$ has the same dimension as the data, i.e., if the trajectories are recorded on the nanometer scale then $[\\sigma] = \\textrm{nm}$, and $\\Delta t$ denotes the time step between two observations.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.5.0",
   "language": "julia",
   "name": "julia-1.5"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.5.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}