test_exp_tmin.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import lsqfit\n",
    "from model_avg_paper import *\n",
    "from model_avg_paper.test_tmin import test_vary_tmin_SE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p0_test_ME = {\n",
    "    'A0': 2.0,\n",
    "    'E0': 0.8,\n",
    "    'A1': 10.4,\n",
    "    'E1': 1.16,\n",
    "}\n",
    "Nt = 32\n",
    "noise_params = {\n",
    "    'noise_amp': 0.3,\n",
    "    'noise_samples': 500,\n",
    "    'frac_noise': True,\n",
    "    'cross_val': False,\n",
    "    'cv_frac': 0.1,\n",
    "}\n",
    "obs_name='E0'\n",
    "\n",
    "correlated_data = True\n",
    "rho=0.6\n",
    "\n",
    "# Set seed for consistency of outcome\n",
    "#np.random.seed(10911)  # Fig 3, subfig A; Fig 4\n",
    "#np.random.seed(81890)  # Fig 3, subfig B\n",
    "#np.random.seed(87414)  # Fig 3, subfig C\n",
    "np.random.seed(77700)  # Fig 3, subfig D\n",
    "\n",
    "\n",
    "def ME_model(x,p):\n",
    "    return multi_exp_model(x,p,Nexc=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if correlated_data:\n",
    "    test_data = gen_synth_data_corr(\n",
    "        np.arange(0,Nt), \n",
    "        p0_test_ME, \n",
    "        ME_model,\n",
    "        rho=rho,\n",
    "        **noise_params)\n",
    "else:\n",
    "    test_data = gen_synth_data(\n",
    "        np.arange(0,Nt), \n",
    "        p0_test_ME, \n",
    "        ME_model,\n",
    "        **noise_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_res = test_vary_tmin_SE(test_data, Nt=Nt, max_tmin=26, obs_name=obs_name, IC='AIC', \n",
    "                             cross_val=noise_params['cross_val'])\n",
    "print(test_res['obs_avg'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Figure 3\n",
    "\n",
    "import matplotlib.ticker as ticker\n",
    "\n",
    "gs = plt.GridSpec(2, 1, height_ratios=[3,1])\n",
    "gs.update(hspace=0.06)\n",
    "\n",
    "ax1 = plt.subplot(gs[0])\n",
    "\n",
    "plot_gvcorr([test_res['obs_avg']], x=np.array([1.5]), color='red', markersize=7, marker='s', open_symbol=True, label='Model avg.')\n",
    "plot_gvcorr(test_res['obs'], x=test_res['tmin'], label='Individual fits')\n",
    "\n",
    "ax1.plot(np.array([-1,34]), 0*np.array([0,0])+p0_test_ME[obs_name], linestyle='--', color='k', label='Model truth')\n",
    "#ax1.set_xlabel('$N_p$')\n",
    "ax1.set_ylabel('$E_0$')\n",
    "\n",
    "ax1.legend(loc='center left', bbox_to_anchor=(1,0.5))\n",
    "ax1.set_xlim(0.7,27.3)\n",
    "\n",
    "plt.setp(ax1.get_xticklabels(), visible=False)\n",
    "\n",
    "ax2 = plt.subplot(gs[1])\n",
    "\n",
    "p_norm = test_res['probs'] / np.sum(test_res['probs'])\n",
    "Q_norm = test_res['Qs'] / np.sum(test_res['Qs'])\n",
    "plt.plot(test_res['tmin'], p_norm, color='orange', label='pr$(M|D)$')\n",
    "plt.plot(test_res['tmin'], Q_norm, color='blue', linestyle='-.', label='Fit p-value')  # Note: fit prob != model prob! \n",
    "\n",
    "tick_spacing = 4\n",
    "ax2.xaxis.set_major_locator(ticker.MultipleLocator(tick_spacing))\n",
    "plt.yticks([0,np.max(p_norm)])\n",
    "ax2.yaxis.set_major_formatter(ticker.FuncFormatter(lambda x, pos: '0' if x == 0 else '{:.2f}'.format(x)))\n",
    "ax2.set_xlim(0.7,27.3)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# Put a legend to the right of the current axis\n",
    "ax2.legend(loc='center left', bbox_to_anchor=(1, 0.5))\n",
    "ax2.set_xlabel(r'$t_{\\rm min}$')\n",
    "ax2.set_ylabel('p')\n",
    "\n",
    "# Uncomment to save figure to disk\n",
    "#plt.savefig('plots/exp_avg_4.pdf', bbox_inches = \"tight\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Scaling w/number of samples\n",
    "Nsamp_array = np.array([20, 40, 80, 160, 320, 640, 2040, 4096, 4096*2, 4096*4])\n",
    "Nsamp_max = Nsamp_array[-1]\n",
    "\n",
    "noise_params['noise_samples'] = Nsamp_max\n",
    "if correlated_data:\n",
    "    scale_data = gen_synth_data_corr(\n",
    "        np.arange(0,Nt), \n",
    "        p0_test_ME, \n",
    "        ME_model,\n",
    "        rho=rho,\n",
    "        **noise_params)    \n",
    "else:\n",
    "    scale_data = gen_synth_data(\n",
    "        np.arange(0,Nt), \n",
    "        p0_test_ME, \n",
    "        ME_model,\n",
    "        **noise_params)\n",
    "\n",
    "model_avg_vs_Nsamp = []\n",
    "naive_avg_vs_Nsamp = []\n",
    "fixed_tmin_vs_Nsamp = []\n",
    "fixed_tmin_2_vs_Nsamp = []\n",
    "fw_vs_Nsamp = []\n",
    "\n",
    "fix_tmin = 14\n",
    "fix_tmin_2 = 8\n",
    "\n",
    "\n",
    "for Nsamp in Nsamp_array:\n",
    "    test_data_scale = cut_synth_data_Nsamp(scale_data, Nsamp)\n",
    "    test_res_scale = test_vary_tmin_SE(test_data_scale, Nt=Nt, max_tmin=Nt-4, obs_name=obs_name, IC='AIC')\n",
    "    test_res_scale_naive = test_vary_tmin_SE(test_data_scale, Nt=Nt, max_tmin=Nt-4, obs_name=obs_name, \n",
    "                                             IC='naive')                                             \n",
    "    \n",
    "    model_avg_vs_Nsamp.append(test_res_scale['obs_avg'])\n",
    "    naive_avg_vs_Nsamp.append(test_res_scale_naive['obs_avg'])\n",
    "    fixed_tmin_vs_Nsamp.append(test_res_scale['obs'][fix_tmin])\n",
    "    fixed_tmin_2_vs_Nsamp.append(test_res_scale['obs'][fix_tmin_2])\n",
    "    fw_vs_Nsamp.append(obs_avg_full_width(test_res_scale['obs'], test_res_scale['Qs'], test_res_scale['fits'], bf_i=None))\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Figure 4\n",
    "\n",
    "plot_gvcorr(model_avg_vs_Nsamp, x=np.log(Nsamp_array)+0.1, label='Model avg. (AIC)')\n",
    "plot_gvcorr(fixed_tmin_vs_Nsamp, x=np.log(Nsamp_array)+0.2, color='red', marker='s', markersize=6, label=r'Fixed $t_{\\rm min} = 14$')\n",
    "plot_gvcorr(fw_vs_Nsamp, x=np.log(Nsamp_array)+0.3, marker='X', markersize=8, color='orange', label='Full-width systematic')\n",
    "plot_gvcorr(naive_avg_vs_Nsamp, x=np.log(Nsamp_array)+0.4, color='silver', marker='v', markersize=8, label=r'Model avg. (naive)')\n",
    "\n",
    "plt.plot(np.arange(0,10), 0*np.arange(0,10)+p0_test_ME[obs_name], linestyle='--', color='k', label='Model truth')\n",
    "plt.xlabel(r'$\\log(N_s)$')\n",
    "plt.ylabel(r'$E_0$')\n",
    "plt.xlim(2.7,7.)\n",
    "plt.ylim(0.78,0.82)\n",
    "\n",
    "\n",
    "# Put a legend to the right of the current axis\n",
    "ax = plt.subplot(111)\n",
    "ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))\n",
    "\n",
    "# Uncomment to save figure to disk\n",
    "#plt.savefig('plots/exp_N_scaling.pdf', bbox_inches = \"tight\")\n"
   ]
  },
  {
   "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.10"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}