analyse_condrealnvp_cww.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os,pickle\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import matplotlib\n",
    "matplotlib.use('Agg')\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from process_data import preprocess_conditional_flow_data_cww\n",
    "from Model.ConditionalRealNVP import ConditionalRealNVP\n",
    "from Utils.ObjDict import ObjDict\n",
    "from Utils.mkdir_p import mkdir_p"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'os' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-2-5e31f6aa301a>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0minput_csv_path\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m\"data/train_cww.npy\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      6\u001b[0m \u001b[0msaved_model_path\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m\"output/train_condrealnvp_cww_210208_v1/saved_model_700.h5\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 7\u001b[0;31m \u001b[0moutput_dir\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdirname\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msaved_model_path\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      8\u001b[0m \u001b[0mevent_size\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m4000\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      9\u001b[0m \u001b[0mndim\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m3\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'os' is not defined"
     ]
    }
   ],
   "source": [
    "# __________________________________________________________________ ||\n",
    "# Basic configurables\n",
    "# __________________________________________________________________ ||\n",
    "\n",
    "input_csv_path = \"data/train_cww.npy\"\n",
    "saved_model_path = \"output/train_condrealnvp_cww_210208_v1/saved_model_700.h5\"\n",
    "output_dir = os.path.dirname(saved_model_path)\n",
    "event_size = 4000\n",
    "ndim = 3\n",
    "ncond = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# __________________________________________________________________ ||\n",
    "# Load models\n",
    "# __________________________________________________________________ ||\n",
    "\n",
    "event_size = 1\n",
    "nf_model = ConditionalRealNVP(num_coupling_layers=10,ndim=ndim,ncond=ncond)\n",
    "samples = nf_model.distribution.sample(event_size)\n",
    "condition = 1.0 * np.ones((event_size,1))\n",
    "_,_ = nf_model.predict([samples,condition,])\n",
    "nf_model.load_weights(saved_model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr = np.load(input_csv_path)\n",
    "arr_list,_ = preprocess_conditional_flow_data_cww(arr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/lucien/.pyenv/versions/3.7.6/lib/python3.7/site-packages/ipykernel/__main__.py:13: RuntimeWarning: More than 20 figures have been opened. Figures created through the pyplot interface (`matplotlib.pyplot.figure`) are retained until explicitly closed and may consume too much memory. (To control this warning, see the rcParam `figure.max_open_warning`).\n"
     ]
    }
   ],
   "source": [
    "# __________________________________________________________________ ||\n",
    "# Make plots for different conditions\n",
    "# __________________________________________________________________ ||\n",
    "\n",
    "n_dim = 3\n",
    "param_grid = [arr_list[idx_param] for idx_param in np.random.randint(0,len(arr_list),n_dim*n_dim)]\n",
    "param_grid.sort(key=lambda x: x.condition[0])\n",
    "figsize = (50,50)\n",
    "event_size = 5000\n",
    "\n",
    "samples = nf_model.distribution.sample(event_size)\n",
    "fig_m4l,ax_m4l = plt.subplots(n_dim,n_dim,figsize=figsize)\n",
    "fig_mz1,ax_mz1 = plt.subplots(n_dim,n_dim,figsize=figsize)\n",
    "fig_mz2,ax_mz2 = plt.subplots(n_dim,n_dim,figsize=figsize)\n",
    "\n",
    "for i,m in enumerate(param_grid):\n",
    "    \n",
    "    ix = int(i / n_dim)\n",
    "    iy = i % n_dim\n",
    "    \n",
    "    condition_str = str(m.condition[0])\n",
    "    condition = np.ones((event_size,1)) * m.condition[0]\n",
    "\n",
    "    idx_batch = np.random.randint(0,m.x.shape[0],event_size)\n",
    "    \n",
    "    x_true = m.x[idx_batch]\n",
    "    x_gen,_ = nf_model.predict([samples,condition,])\n",
    "\n",
    "    ax_m4l[ix,iy].hist(x_true[:,0],bins=100,density=1.,histtype='step',range=[-10.,10.],label='True '+condition_str)\n",
    "    ax_m4l[ix,iy].hist(x_gen[:,0],bins=100,density=1.,histtype='step',range=[-10.,10.],label='Flow '+condition_str)\n",
    "    ax_m4l[ix,iy].legend(loc='best')\n",
    "    ax_m4l[ix,iy].set_title(condition_str)\n",
    "    \n",
    "    ax_mz1[ix,iy].hist(x_true[:,1],bins=100,density=1.,histtype='step',range=[-10.,10.],label='True '+condition_str)\n",
    "    ax_mz1[ix,iy].hist(x_gen[:,1],bins=100,density=1.,histtype='step',range=[-10.,10.],label='Flow '+condition_str)\n",
    "    ax_mz1[ix,iy].legend(loc='best')\n",
    "    ax_mz1[ix,iy].set_title(condition_str)\n",
    "    \n",
    "    ax_mz2[ix,iy].hist(x_true[:,2],bins=100,density=1.,histtype='step',range=[-10.,10.],label='True '+condition_str)\n",
    "    ax_mz2[ix,iy].hist(x_gen[:,2],bins=100,density=1.,histtype='step',range=[-5.,5.],label='Flow '+condition_str)\n",
    "    ax_mz2[ix,iy].legend(loc='best')\n",
    "    ax_mz2[ix,iy].set_title(condition_str)\n",
    "    \n",
    "fig_m4l.savefig(os.path.join(output_dir,'m4l.png'))\n",
    "fig_mz1.savefig(os.path.join(output_dir,'mZ1.png'))\n",
    "fig_mz2.savefig(os.path.join(output_dir,'mZ2.png'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'arr_list' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-3-c2d22916dd31>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      7\u001b[0m \u001b[0mn_dim\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m5\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      8\u001b[0m \u001b[0mselect_param\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;36m5\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m30\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m45\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m60\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m80\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 9\u001b[0;31m \u001b[0mparam_grid\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0marr_list\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0midx_param\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0midx_param\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mselect_param\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     10\u001b[0m \u001b[0mparam_grid\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msort\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mkey\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mlambda\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcondition\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     11\u001b[0m \u001b[0mfigsize\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;36m25\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m10\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-3-c2d22916dd31>\u001b[0m in \u001b[0;36m<listcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m      7\u001b[0m \u001b[0mn_dim\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m5\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      8\u001b[0m \u001b[0mselect_param\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;36m5\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m30\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m45\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m60\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m80\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 9\u001b[0;31m \u001b[0mparam_grid\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0marr_list\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0midx_param\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0midx_param\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mselect_param\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     10\u001b[0m \u001b[0mparam_grid\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msort\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mkey\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mlambda\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcondition\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     11\u001b[0m \u001b[0mfigsize\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;36m25\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m10\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'arr_list' is not defined"
     ]
    }
   ],
   "source": [
    "# __________________________________________________________________ ||\n",
    "# Make plots for likelihood\n",
    "# __________________________________________________________________ ||\n",
    "\n",
    "import time\n",
    "\n",
    "n_dim = 5\n",
    "select_param = [5,30,45,60,80]\n",
    "param_grid = [arr_list[idx_param] for idx_param in select_param]\n",
    "param_grid.sort(key=lambda x: x.condition[0])\n",
    "figsize = (25,10)\n",
    "\n",
    "plot_low = 0.0\n",
    "plot_high = 0.2\n",
    "n_grid = 100\n",
    "x_grid = [plot_low+(plot_high-plot_low)/n_grid*i for i in range(n_grid+1)]\n",
    "\n",
    "fig, ax = plt.subplots(len(select_param),n_dim,figsize=figsize,constrained_layout=True)\n",
    "\n",
    "def make_array(event_size,x_grid,p,nf_model):\n",
    "    log_prob = np.zeros(n_grid+1)\n",
    "    log_prob_sm = np.zeros(n_grid+1)\n",
    "    \n",
    "    idx_batch = np.random.randint(0,p.x.shape[0],event_size)\n",
    "    condition_concat = np.concatenate([np.ones((event_size,1)) * x for ix,x in enumerate(x_grid)])\n",
    "    x_data_concat = np.concatenate([p.x[idx_batch] for ix,x in enumerate(x_grid)])\n",
    "    \n",
    "    z_concat = nf_model.batch_log_loss([x_data_concat,condition_concat])\n",
    "\n",
    "    for ig,x in enumerate(x_grid):\n",
    "        arg = tf.expand_dims(z_concat[ig*idx_batch.shape[0]:(ig+1)*idx_batch.shape[0]],axis=1)\n",
    "        log_prob[ig] = tf.reduce_sum(nf_model.distribution.log_prob(arg))\n",
    "    return log_prob.astype(np.float64),idx_batch\n",
    "\n",
    "for i,p in enumerate(param_grid):\n",
    "    \n",
    "    print(\"-\"*100)\n",
    "    print(\"Drawing plot \",i,\" with param \",p.condition[0])\n",
    "    \n",
    "    #ix = int(i / n_dim)\n",
    "    #iy = i % n_dim\n",
    "    ix = i\n",
    "    \n",
    "    start_time = time.time()\n",
    "    \n",
    "    condition_str = str(p)\n",
    "    \n",
    "    for event_size in [5,50,1000]:\n",
    "        nf_model.direction = -1\n",
    "        log_prob,idx_batch = make_array(event_size,x_grid,p,nf_model)\n",
    "        \n",
    "        y_grid = tf.nn.softmax(log_prob/event_size) * n_grid\n",
    "        ax[ix,0].plot(x_grid,y_grid,label=str(event_size)+' events',)\n",
    "        ax[ix,0].set_title(\"cww: \"+\"%4.4f\"%p.param)\n",
    "        ax[ix,0].legend(loc='best')\n",
    "        ax[ix,0].grid(True)\n",
    "        ax[ix,0].set_ylim(0.,40.)\n",
    "        ylims = ax[ix,0].get_ylim()\n",
    "        ax[ix,0].arrow(p.condition[0], ylims[1], 0., ylims[0]-ylims[1],)\n",
    "        \n",
    "        ax[ix,1].plot(x_grid,-log_prob-np.min(-log_prob),label=str(event_size)+' events',)\n",
    "        ax[ix,1].set_title(\"cww: \"+\"%4.4f\"%p.param)\n",
    "        ax[ix,1].legend(loc='best')\n",
    "        ax[ix,1].grid(True)\n",
    "        ax[ix,1].set_ylim(0.,10.)\n",
    "        ylims = ax[ix,1].get_ylim()\n",
    "        ax[ix,1].arrow(p.condition[0], ylims[1], 0., ylims[0]-ylims[1],)\n",
    "        \n",
    "        if event_size == 1000:\n",
    "            nf_model.direction = 1\n",
    "            samples = nf_model.distribution.sample(idx_batch.shape[0])\n",
    "            x_pred,_ = nf_model.predict([samples,np.ones((idx_batch.shape[0],1))*p.param])\n",
    "            ax[ix,2].hist(p.x[idx_batch,0],bins=50,histtype='step',range=[-1.,1.],label='True',)\n",
    "            ax[ix,2].hist(x_pred[:,0],bins=50,histtype='step',range=[-1.,1.],label='Flow',)\n",
    "            ax[ix,2].legend(loc='best')\n",
    "            \n",
    "            ax[ix,3].hist(p.x[idx_batch,1],bins=20,histtype='step',range=[-2.,2.],label='True',)\n",
    "            ax[ix,3].hist (x_pred[:,1],bins=20,histtype='step',range=[-2.,2.],label='Flow',)\n",
    "            ax[ix,3].legend(loc='best')\n",
    "            \n",
    "            ax[ix,4].hist(p.x[idx_batch,2],bins=20,histtype='step',range=[0.,5.],label='True',)\n",
    "            ax[ix,4].hist(x_pred[:,2],bins=20,histtype='step',range=[0.,5.],label='Flow',)\n",
    "            ax[ix,4].legend(loc='best')\n",
    "    \n",
    "    elapsed_time = time.time() - start_time\n",
    "    print(\"Time used: \"+str(elapsed_time)+\"s\")\n",
    "    \n",
    "fig.savefig(os.path.join(output_dir,'log_loss.png'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py3.7.6",
   "language": "python",
   "name": "py3.7.6"
  },
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}