21_plot_encoding_examples.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "from scipy import stats\n",
    "\n",
    "import seaborn as sns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "### execute script to load modules here\n",
    "exec(open('setup_aesthetics.py').read())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## create export directory if necessary\n",
    "## foldernames for output plots/lists produced in this notebook\n",
    "import os\n",
    "FIG_DIR = f'./figures/encodings_examples/'\n",
    "os.makedirs(FIG_DIR, exist_ok=True)\n",
    "print(\"All  plots will be stored in: \\n\" + FIG_DIR)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define encodings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### define helper function\n",
    "def logit(x):\n",
    "    return np.log(np.divide(x,1-x))\n",
    "## test\n",
    "logit(0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval_statistic(xf,x0, phi = logit):\n",
    "    return phi(xf) - phi(x0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_statistic(xf = 0.55, x0 = 0.45)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval_statistic_s(xf,x0):\n",
    "    return eval_statistic(xf=xf,x0=x0, phi = logit) \n",
    "\n",
    "def eval_statistic_deltalog(xf,x0):\n",
    "    return eval_statistic(xf=xf,x0=x0, phi =lambda x: np.log(x)) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "palette = sns.color_palette(\"Set2\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### simulate trajectory to fixation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Formulating the ODE problem in python is not obvious, the underlying SCIPY module requires a specific format for the derivative. \n",
    "\n",
    "See an example for ODE solving in Python here: https://pundit.pratt.duke.edu/wiki/Python:Ordinary_Differential_Equations/Examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.integrate import solve_ivp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Logistic Equation of growth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "### define logistic equation\n",
    "def fun(t, y, r = 5):\n",
    "    return r*np.multiply(y,(1-y))*np.heaviside(y,0)\n",
    "\n",
    "## define timewindow of solution\n",
    "tspan = [0,5]\n",
    "\n",
    "### define initial condition\n",
    "y0 = [1e-6]\n",
    "\n",
    "\n",
    "sol = solve_ivp(fun, tspan, y0 = y0, atol = 1e-10, rtol = 1e-8, vectorized = True)\n",
    "yraw = sol.y[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(3,1, figsize = (FIGHEIGHT_TRIPLET, 1*FIGHEIGHT_TRIPLET), sharex = True)\n",
    "\n",
    "ax = axes[0]\n",
    "#ax.set_ylabel('frequency $x$')\n",
    "#twin = ax.twinx()\n",
    "ax.plot(sol.t,yraw, color = palette[0], lw = 3)\n",
    "\n",
    "\n",
    "\n",
    "ax = axes[1]\n",
    "#ax.set_ylabel('$\\log\\; x$')\n",
    "#twin = ax.twinx()\n",
    "ax.plot(sol.t,np.log(yraw),color = palette[1], lw = 3)\n",
    "\n",
    "\n",
    "\n",
    "ax = axes[2]\n",
    "#ax.set_ylabel('$\\mathrm{logit}\\; x$')\n",
    "#twin = ax.twinx()\n",
    "ax.plot(sol.t,logit(yraw), color = palette[2], lw = 3)\n",
    "\n",
    "for ax in axes:\n",
    "    ax.tick_params(left = False, labelleft= False)\n",
    "    ax.set_xlim(sol.t[0],sol.t[-1])\n",
    "    ax.tick_params(left = True, labelleft = True)\n",
    "\n",
    "ax.set_xlabel('time',labelpad =10)\n",
    "ax.tick_params(labelbottom = False)\n",
    "\n",
    "fig.savefig(FIG_DIR + f'example_trajectories_logistic.pdf', DPI = DPI, bbox_inches = 'tight', pad_inches = PAD_INCHES)\n",
    "  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gompertz Equation of growth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "### define gompertz equation\n",
    "def fun(t, y, r = 5):\n",
    "    return r*np.multiply(y,np.log(1/y))*np.heaviside(y,0)\n",
    "\n",
    "## define timewindow of solution\n",
    "tspan = [0,1.5]\n",
    "\n",
    "sol = solve_ivp(fun, tspan, y0 = y0, atol = 1e-10, rtol = 1e-8, vectorized = True)\n",
    "yraw = sol.y[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(3,1, figsize = (FIGHEIGHT_TRIPLET, 1*FIGHEIGHT_TRIPLET), sharex = True)\n",
    "\n",
    "ax = axes[0]\n",
    "#ax.set_ylabel('frequency $x$')\n",
    "#twin = ax.twinx()\n",
    "ax.plot(sol.t,yraw, color = palette[0], lw = 3)\n",
    "\n",
    "\n",
    "\n",
    "ax = axes[1]\n",
    "#ax.set_ylabel('$\\log\\; x$')\n",
    "#twin = ax.twinx()\n",
    "ax.plot(sol.t,np.log(yraw),color = palette[1], lw = 3)\n",
    "\n",
    "\n",
    "\n",
    "ax = axes[2]\n",
    "#ax.set_ylabel('$\\mathrm{logit}\\; x$')\n",
    "#twin = ax.twinx()\n",
    "ax.plot(sol.t,logit(yraw), color = palette[2], lw = 3)\n",
    "\n",
    "for ax in axes:\n",
    "    ax.tick_params(left = False, labelleft= False)\n",
    "    ax.set_xlim(sol.t[0],sol.t[-1])\n",
    "    ax.tick_params(left = True, labelleft = True)\n",
    "\n",
    "ax.set_xlabel('time',labelpad =10)\n",
    "ax.tick_params(labelbottom = False)\n",
    "\n",
    "fig.savefig(FIG_DIR + f'example_trajectories_gompertz.pdf', DPI = DPI, bbox_inches = 'tight', pad_inches = PAD_INCHES)\n",
    "  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot heteroscedasticity for regression on raw variable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats import linregress"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "### define logistic equation\n",
    "def fun(t, y, r = 5):\n",
    "    return r*np.multiply(y,(1-y))*np.heaviside(y,0)\n",
    "\n",
    "## define timewindow of solution\n",
    "tspan = [0,5]\n",
    "\n",
    "### define initial condition\n",
    "y0 = [1e-6]\n",
    "\n",
    "\n",
    "sol = solve_ivp(fun, tspan, y0 = y0, atol = 1e-10, rtol = 1e-8, vectorized = True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## define the true trajectory\n",
    "tvec = sol.t\n",
    "yraw = sol.y[0]\n",
    "\n",
    "def time2freq(t):\n",
    "    return np.interp(t, tvec,yraw)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rng = np.random.default_rng(seed = 304576)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### demonstration: using example from scipy documentation\n",
    "\n",
    "n_trials, p_success = 10, .5  # number of trials, probability of each trial\n",
    "\n",
    "# result of flipping a coin 10 times, tested 1000 times.\n",
    "rng.binomial(n_trials, p_success, 10)/n_trials\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## demonstration: how this sampling works\n",
    "\n",
    "n_trials = int(1e5*0.1) # number of cells that are extracted\n",
    "size = 5 # number of extraction attempts\n",
    "rng.binomial(n_trials, p_success, size = size)/n_trials"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## sample binomial data along a regression\n",
    "n_timepoints = 5\n",
    "### chosse timewindow for sampliing\n",
    "t_start, t_end = 2.2,3.2 # middle window\n",
    "#t_start, t_end = 1.,2. # early window\n",
    "#t_start, t_end = 2.5, 3.5  #late window\n",
    "timepoints = np.linspace(t_start,t_end,num = n_timepoints)\n",
    "\n",
    "### choose parameters for the cell count and frequency estimation\n",
    "n_cells_counted = int(1e2) # number of cells that are extracted\n",
    "n_extractions = 50 # number of extraction attempts\n",
    "\n",
    "### simulate extraction process\n",
    "freq_samples = np.ones((n_timepoints,n_extractions))\n",
    "\n",
    "for i in range(len(timepoints)): \n",
    "    freq_true = time2freq(timepoints[i])\n",
    "    freq_samples[i] =  rng.binomial(n_cells_counted, freq_true, \n",
    "                                    size = n_extractions)/n_cells_counted\n",
    "    \n",
    "\n",
    "\n",
    "## create similar matrix of timepoints\n",
    "timepoint_samples = np.outer(timepoints, np.ones(n_extractions))\n",
    "\n",
    "## flatten both objects\n",
    "freq_samples = freq_samples.flatten()\n",
    "timepoint_samples = timepoint_samples.flatten()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### build regression models for the different encodings\n",
    "\n",
    "## no encoding\n",
    "slope_raw, intercept_raw, _, _, _ = linregress(x=timepoint_samples, y=freq_samples)\n",
    "def time2freq_fit(t):\n",
    "    return intercept_raw + slope_raw*t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "yregression = time2freq_fit(timepoints)\n",
    "freq_fitted = np.outer(yregression,np.ones(n_extractions)).flatten()\n",
    "freq_true = time2freq(timepoint_samples)\n",
    "\n",
    "freq_residuals = freq_samples - freq_fitted\n",
    "\n",
    "freq_residual_mean = np.zeros_like(timepoints)\n",
    "for i in range(len(timepoints)):\n",
    "    t = timepoints[i]\n",
    "    freq_residual_mean[i] = freq_residuals[timepoint_samples ==t].mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "yregression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## log encoding\n",
    "## remove inf values\n",
    "is_finite = freq_samples != 0\n",
    "log_timepoints = timepoint_samples[is_finite]\n",
    "log_samples = np.log(freq_samples[is_finite])\n",
    "\n",
    "slope_log, intercept_log, _, _, _ = linregress(x= log_timepoints, y=log_samples)\n",
    "def time2log_fit(t):\n",
    "    return intercept_log + slope_log*t\n",
    "\n",
    "yregression = time2log_fit(timepoints)\n",
    "log_fitted = np.outer(yregression,np.ones(n_extractions)).flatten()[is_finite]\n",
    "log_true = np.log(freq_true)[is_finite]\n",
    "\n",
    "log_residuals = log_samples - log_fitted\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "log_residual_mean = np.zeros_like(timepoints)\n",
    "for i in range(len(timepoints)):\n",
    "    t = timepoints[i]\n",
    "    log_residual_mean[i] = log_residuals[log_timepoints ==t].mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## logitit encoding\n",
    "## remove inf values\n",
    "is_finite = (freq_samples != 0) & (freq_samples != 1)\n",
    "logit_timepoints = timepoint_samples[is_finite]\n",
    "logit_samples = logit(freq_samples[is_finite])\n",
    "\n",
    "slope_logit, intercept_logit, _, _, _ = linregress(x=logit_timepoints, y=logit_samples)\n",
    "def time2logit_fit(t):\n",
    "    return intercept_logit + slope_logit*t\n",
    "\n",
    "yregression = time2logit_fit(timepoints)\n",
    "logit_fitted = np.outer(yregression,np.ones(n_extractions)).flatten()[is_finite]\n",
    "logit_true = logit(freq_true)[is_finite]\n",
    "\n",
    "logit_residuals = logit_samples - logit_fitted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logit_residual_mean = np.zeros_like(timepoints)\n",
    "for i in range(len(timepoints)):\n",
    "    t = timepoints[i]\n",
    "    logit_residual_mean[i] = logit_residuals[logit_timepoints ==t].mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(3,2, figsize = (2.5*FIGHEIGHT_TRIPLET, 1*FIGHEIGHT_TRIPLET), sharex=True)\n",
    "\n",
    "ax = axes[0,0]\n",
    "ax.set_ylabel('frequency $x$')\n",
    "ax.plot(tvec,yraw, color = palette[0], lw = 3)\n",
    "## plot samples with regression\n",
    "ax.scatter(timepoint_samples, freq_samples, color = 'tab:grey')\n",
    "ax.plot(timepoints,time2freq_fit(timepoints), color = 'tab:red')\n",
    "\n",
    "### plot residuals\n",
    "ax = axes[0,1]\n",
    "ax.scatter(timepoint_samples,freq_residuals, color= 'tab:grey')\n",
    "ax.scatter(timepoints, freq_residual_mean, marker = 'x',color = 'tab:red')\n",
    "ax.plot(timepoints, freq_residual_mean, color = 'tab:red', lw = 3,zorder = 3)\n",
    "ax.axhline(0, ls = '--', color = 'black')\n",
    "\n",
    "\n",
    "ax = axes[1,0]\n",
    "ax.set_ylabel('$\\log\\; x$')\n",
    "ax.plot(tvec,np.log(yraw),color = palette[1], lw = 3)\n",
    "## plot samples with regression\n",
    "\n",
    "ax.scatter(timepoint_samples, np.log(freq_samples), color = 'tab:grey')\n",
    "ax.plot(timepoints,time2log_fit(timepoints), color = 'tab:red')\n",
    "\n",
    "\n",
    "### plot residuals\n",
    "ax = axes[1,1]\n",
    "\n",
    "ax.scatter(log_timepoints,log_residuals, color= 'tab:grey')\n",
    "ax.scatter(timepoints, log_residual_mean, marker = 'x',color = 'tab:red')\n",
    "ax.plot(timepoints, log_residual_mean, color = 'tab:red', lw = 3,zorder = 3)\n",
    "\n",
    "ax.axhline(0, ls = '--', color = 'black')\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "ax = axes[2,0]\n",
    "ax.set_ylabel('$\\mathrm{logit}\\; x$')\n",
    "#twin = ax.twinx()\n",
    "ax.plot(tvec,logit(yraw), color = palette[2], lw = 3)\n",
    "\n",
    "## plot samples with regression\n",
    "ax.scatter(timepoint_samples, logit(freq_samples), color = 'tab:grey')\n",
    "ax.plot(timepoints,time2logit_fit(timepoints), color = 'tab:red')\n",
    "\n",
    "### plot residuals\n",
    "ax = axes[2,1]\n",
    "ax.scatter(logit_timepoints,logit_residuals, color= 'tab:grey')\n",
    "ax.scatter(timepoints, logit_residual_mean, marker = 'x',color = 'tab:red')\n",
    "ax.plot(timepoints, logit_residual_mean, color = 'tab:red', lw = 3, zorder = 3)\n",
    "ax.axhline(0, ls = '--', color = 'black')\n",
    "\n",
    "\n",
    "for ax in axes[:,0]:\n",
    "    ax.tick_params(left = False, labelleft= False)\n",
    "    ax.set_xlim(tvec[0],tvec[-1])\n",
    "    \n",
    "for ax in axes[:,1]:\n",
    "    ymin,ymax = ax.get_ylim()\n",
    "    yabs = np.max(np.abs([ymin,ymax]))*1.1\n",
    "    ax.set_ylim(-yabs,yabs)\n",
    "    \n",
    "    ax.tick_params(left = False, labelleft= False)\n",
    "    ax.set_ylabel('residuals')\n",
    "    \n",
    "\n",
    "\n",
    "axes[2,0].set_xlabel('time')\n",
    "axes[2,1].set_xlabel('time')\n",
    "#ax.tick_params(labelbottom = False)\n",
    "\n",
    "fig.savefig(FIG_DIR + f'heteroscedasticity_comparison.pdf', DPI = DPI, bbox_inches = 'tight', pad_inches = PAD_INCHES)\n",
    "  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### calculate standard deviation from binomial sampling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To calculate the error of sampling, we need to specify\n",
    "- the true frequency `p`\n",
    "- the number of samples drawn from the urn `n`\n",
    "- the size of the population `N` in the urn. \n",
    "\n",
    "The variation on a set of samples `x_i`, where `x_i` are realizations of the bionmial random variable so they are either `x_i = 0` (head) or `x_i=1` tail is given by \n",
    "\n",
    "    mean = sum x_i/n\n",
    "    var = sum (x_i - mean)^2\n",
    "    \n",
    "For the binomial distribution, we know that\n",
    "\n",
    "    var = n*p*(1-p)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What makes this potentially confusing is that in many experiments, the number of samples collected is varied with the true frequency `p`. For example, if the frequency of a focal strain is low, then the experimentalists keeps counting colonies until he has at least `100` colonies counted of the rare type. This effectively means `n*p` is fixed.\n",
    "\n",
    "In a simplification, here we only plot the error from single draw `n=1`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xvec = np.linspace(0.01,0.99,num = 500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "var = np.multiply(xvec,1-xvec)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(3,1, figsize = (1.5*FIGHEIGHT_TRIPLET, 1.33*FIGHEIGHT_TRIPLET), sharex = True)\n",
    "\n",
    "ax = axes[0]\n",
    "ax.set_ylabel('frequency $x$')\n",
    "#twin = ax.twinx()\n",
    "ax.plot(xvec, var, color = palette[0], lw = 3)\n",
    "\n",
    "\n",
    "\n",
    "ax = axes[1]\n",
    "ax.set_ylabel('$\\log\\; x$')\n",
    "#twin = ax.twinx()\n",
    "ax.plot(xvec,np.log(var),color = palette[1], lw = 3)\n",
    "\n",
    "\n",
    "\n",
    "ax = axes[2]\n",
    "ax.set_ylabel('$\\mathrm{logit}\\; x$')\n",
    "#twin = ax.twinx()\n",
    "ax.plot(xvec,logit(var), color = palette[2], lw = 3)\n",
    "\n",
    "for ax in axes:\n",
    "    #ax.tick_params(left = False, labelleft= False)\n",
    "    pass\n",
    "\n",
    "ax.set_xlabel('time')\n",
    "ax.tick_params(labelbottom = False)\n",
    "\n",
    "\n",
    "fig.savefig(FIG_DIR + f'example_error_trajectories.pdf', DPI = DPI, bbox_inches = 'tight', pad_inches = PAD_INCHES)\n",
    "  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Under the encoding we need to compute the variation as follows. \n",
    "\n",
    "First we compute the expected value, then the va\n",
    "\n",
    "    E[f(x)] = f(0)*(1-p) * f(0)* p\n",
    "    \n",
    "The problem: For our transforms `f=log` and `f=logit` the values are infinite in this simple calculation.\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### plot a binomial distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### create a sample of initial frequencies\n",
    "rs = np.random.RandomState(15031998)\n",
    "\n",
    "\n",
    "\n",
    "### size of the random vector\n",
    "### basically, number of replicate experiments\n",
    "size = 100000\n",
    "\n",
    "df_data = pd.DataFrame()\n",
    "\n",
    "n_sampled = 1000 #/xtrue #number of balls drawn from the urn at each replicate experiment\n",
    "n_success  = 100\n",
    "\n",
    "for xtrue, n_sampled in zip([0.99, 0.5, 0.01],[n_success/0.01, n_success/0.5, n_success/0.01]):\n",
    "\n",
    "    \n",
    "\n",
    "    dist = rs.binomial(n=n_sampled,p=xtrue,size = size)/n_sampled\n",
    "\n",
    "    ## sample raw frequencies\n",
    "    df_raw = pd.DataFrame(data=np.vstack([xtrue*np.ones_like(dist),dist]).T, columns = ['true frequency', 'value'])\n",
    "    \n",
    "    \n",
    "    mean,std = df_raw['value'].mean(), df_raw['value'].std()\n",
    "    df_raw['rescaled'] = (df_raw['value'] - mean)/std\n",
    "    df_raw['residual'] = df_raw['value'] - mean\n",
    "    df_raw['type'] = 'no encoding'\n",
    "    df_data = df_data.append(df_raw)\n",
    "    \n",
    "    ### evaluate under logit transform \n",
    "    df_logit = df_raw.copy(deep=True) \n",
    "    df_logit['value'] = np.array([logit(v) for v in df_raw['value']])\n",
    "    df_logit['type'] = 'encoded with logit'\n",
    "    # Replacing infinite with nan\n",
    "    df_logit = df_logit.replace([np.inf, -np.inf], np.nan)\n",
    "    mean,std = df_logit['value'].mean(), df_logit['value'].std()\n",
    "    df_logit['rescaled'] = (df_logit['value'] - mean)/std\n",
    "    df_logit['residual'] = df_logit['value'] - mean\n",
    "    df_data = df_data.append(df_logit)\n",
    "    \n",
    "    ## evaluate under log transform \n",
    "    df_log = df_raw.copy(deep=True) \n",
    "    df_log['value'] = np.array([np.log(v) for v in df_raw['value']])\n",
    "    df_log['type'] = 'encoded with log'\n",
    "    df_log = df_log.replace([np.inf, -np.inf], np.nan)\n",
    "    mean,std = df_log['value'].mean(), df_log['value'].std()\n",
    "    df_log['rescaled'] = (df_log['value'] - mean)/std\n",
    "    df_log['residual'] = df_log['value'] - mean\n",
    "    df_data = df_data.append(df_log)\n",
    "\n",
    "\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "# Replacing infinite with nan\n",
    "df_data = df_data.replace([np.inf, -np.inf], np.nan)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sum(df_data['value'].isna())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "palette = sns.color_palette(\"Set2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## sort\n",
    "df_data = df_data.sort_values(['type', 'true frequency'], ascending = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 3, figsize = (FIGHEIGHT_TRIPLET,FIGHEIGHT_TRIPLET), sharey = True)\n",
    "\n",
    "ax = axes[0]\n",
    "label = 'no encoding'\n",
    "data_to_plot =df_data[df_data['type']== label]\n",
    "sns.violinplot(x='residual',y = 'true frequency', data=data_to_plot, ax =ax,orient = 'h',\n",
    "              label=label,  color = palette[0], scale = 'count', rasterized = True, \n",
    "               inner = None, cut = 0)\n",
    "ax.set_xlabel(label)\n",
    "sns.despine(ax=ax)\n",
    "\n",
    "\n",
    "ax = axes[1]\n",
    "label = 'encoded with log'\n",
    "data_to_plot =df_data[df_data['type']== label]\n",
    "sns.violinplot(x='residual',y = 'true frequency', data=data_to_plot, ax =ax,orient = 'h',\n",
    "              label=label,  color = palette[1], scale = 'count', rasterized = True,\n",
    "              inner = None, cut = 0)\n",
    "sns.despine(ax=ax, left = True)\n",
    "ax.tick_params(left=False)\n",
    "ax.set_ylabel(\"\")\n",
    "ax.set_xlabel('log')\n",
    "\n",
    "ax = axes[2]\n",
    "label = 'encoded with logit'\n",
    "data_to_plot =df_data[df_data['type']== label]\n",
    "sns.violinplot(x='residual',y = 'true frequency', data=data_to_plot, ax =ax,orient = 'h',\n",
    "              label=label,  color = palette[2], scale = 'count', rasterized = True, \n",
    "               inner = None, cut = 0)\n",
    "ax.set_xlabel('logit')\n",
    "ax.set_ylabel(\"\")\n",
    "sns.despine(ax=ax, left = True)\n",
    "ax.tick_params(left=False)\n",
    "\n",
    "\n",
    "for ax in axes: \n",
    "    ### symmetrize\n",
    "    xmin,xmax = ax.get_xlim()\n",
    "    max_abs = np.abs([xmin,xmax]).max()\n",
    "    ax.set_xlim(-max_abs,max_abs)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "fig.savefig(FIG_DIR + f'example_distributions_after_encoding.pdf', DPI = DPI, bbox_inches = 'tight', pad_inches = PAD_INCHES)\n",
    "              "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Standard deviations\")\n",
    "for label in ['no encoding', 'encoded with logit', 'encoded with log']:\n",
    "    print(label)\n",
    "    df_bytype =df_data[df_data['type']== label]\n",
    "    for xtrue in [0.99, 0.5, 0.01]:\n",
    "        data = df_bytype[df_bytype['true frequency'] == xtrue]\n",
    "        print(data['value'].std())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### plot overview of encodings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "palette = sns.color_palette(\"Set2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize = (FIGHEIGHT_TRIPLET,FIGHEIGHT_TRIPLET))\n",
    "\n",
    "x0_vec = np.linspace(0.0001,0.9999, num = 100)\n",
    "\n",
    "ax.plot(x0_vec,x0_vec, color = palette[0], label = '$m=x$', lw = 3)\n",
    "ax.plot(x0_vec, np.log(x0_vec), color = palette[1], ls = '-', label = '$m=\\log(x)$', lw = 3)\n",
    "ax.plot(x0_vec, logit(x0_vec), color = palette[2], ls = '-', label = '$m=\\log(x/1-x)$', lw = 3)\n",
    "\n",
    "ax.axhline(0, color = 'black', ls = 'dotted')\n",
    "ax.set_ylim(-3,3)\n",
    "ax.set_xlim(0,1)\n",
    "\n",
    "ax.set_xlabel('input strain frequency $x$')\n",
    "ax.set_ylabel('output $m$ from encoding function')\n",
    "\n",
    "ax.legend(loc = 'upper left')\n",
    "\n",
    "fig.savefig(FIG_DIR + f'example_encodings.pdf', DPI = DPI, bbox_inches = 'tight', pad_inches = PAD_INCHES)\n",
    "              "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}