AcceptanceEffects.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "#os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"\n",
    "#os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import absolute_import, division, print_function\n",
    "\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from matplotlib.patches import Patch\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "import tensorflow as tf\n",
    "import keras\n",
    "import keras.backend as K\n",
    "from keras.layers import Dense, Input\n",
    "from keras.models import Model\n",
    "from keras.utils import to_categorical\n",
    "from keras.callbacks import EarlyStopping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Global plot settings\n",
    "from matplotlib import rc\n",
    "import matplotlib.font_manager\n",
    "\n",
    "rc('font', family='serif')\n",
    "rc('text', usetex=True)\n",
    "rc('font', size=22)\n",
    "rc('xtick', labelsize=15)\n",
    "rc('ytick', labelsize=15)\n",
    "rc('legend', fontsize=15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define default plot styles\n",
    "plot_style_0 = {\n",
    "    'histtype': 'step',\n",
    "    'color': 'black',\n",
    "    'linewidth': 2,\n",
    "    'linestyle': '--',\n",
    "    'density': True\n",
    "}\n",
    "\n",
    "plot_style_1 = {\n",
    "    'histtype': 'step',\n",
    "    'color': 'black',\n",
    "    'linewidth': 2,\n",
    "    'density': True\n",
    "}\n",
    "\n",
    "plot_style_2 = {'alpha': 0.5, 'density': True}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check Versions\n",
    "print(tf.__version__)  # 1.15.0\n",
    "print(keras.__version__)  # 2.2.4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# OmniFold Gaussian Toy Example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup the Datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mu0 = 0\n",
    "sigma0 = 1\n",
    "\n",
    "eff = 0.1 #fraction of true but not reco\n",
    "fake = 0.1 #fraction of reco but not true"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 10**5\n",
    "# param = (mu, sigma)\n",
    "theta0_param = (mu0, sigma0)  # synthetic sample\n",
    "theta_unknown_param = (1, 1.5)  # this is the data (the target)\n",
    "epsilon = sigma0 / 1.  # Smearing width\n",
    "\n",
    "theta0_G = np.random.normal(theta0_param[0], theta0_param[1],N)  # Generator-level synthetic sample\n",
    "theta0_S = np.array([(x + np.random.normal(0, epsilon)) for x in theta0_G])  # Detector smearing for synthetic sample\n",
    "pass_reco = np.random.binomial(1,1.-eff,len(theta0_G))\n",
    "pass_truth = np.random.binomial(1,1.-fake,len(theta0_G))\n",
    "theta0_S[pass_reco==0] = -10\n",
    "theta0_G[pass_truth==0] = -10\n",
    "\n",
    "theta0 = np.stack([theta0_G, theta0_S], axis=1)\n",
    "labels0 = np.zeros(len(theta0))\n",
    "\n",
    "theta_unknown_G = np.random.normal(theta_unknown_param[0],theta_unknown_param[1], N)\n",
    "theta_unknown_S = np.array([(x + np.random.normal(0, epsilon)) for x in theta_unknown_G]) \n",
    "#pass_reco = np.random.binomial(1,1.-eff,len(theta_unknown_G))\n",
    "#pass_truth = np.random.binomial(1,1.-fake,len(theta_unknown_G))\n",
    "#theta_unknown_S[pass_reco==0] = -10\n",
    "#theta_unknown_G[pass_truth==0] = -10\n",
    "\n",
    "theta_unknown = np.stack([theta_unknown_G, theta_unknown_S], axis=1)\n",
    "labels_unknown = np.ones(len(theta_unknown))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins = np.linspace(-11, 11, 41)\n",
    "\n",
    "fig, ax = plt.subplots(1,\n",
    "                       2,\n",
    "                       figsize=(12, 5),\n",
    "                       constrained_layout=True,\n",
    "                       sharey=True)\n",
    "\n",
    "ax[0].set_xlabel('$x$')\n",
    "ax[0].set_ylabel('Events per bin (normalized)')\n",
    "ax[0].hist(theta0_G[theta0_G!=-10], bins=bins, **plot_style_2, label='Generation')\n",
    "ax[0].hist(theta0_S[theta0_S!=-10], bins=bins, **plot_style_2, label='Simulation')\n",
    "legend = ax[0].legend(\n",
    "    title='Synthetic\\n(w/o MC weights)\\n($\\mu$, $\\sigma$) = {}'.format(theta0_param),\n",
    "    loc='upper left',\n",
    "    frameon=True)\n",
    "plt.setp(legend.get_title(), multialignment='center')\n",
    "\n",
    "ax[1].set_xlabel('$x$')\n",
    "ax[1].hist(theta_unknown_G[theta_unknown_G!=-10], bins=bins, **plot_style_2, label='Truth')\n",
    "ax[1].hist(theta_unknown_S[theta_unknown_S!=-10], bins=bins, **plot_style_2, label='Detector')\n",
    "legend = ax[1].legend(\n",
    "    title='Natural\\n(w/o MC weights)\\n($\\mu$, $\\sigma$) = {}'.format(theta_unknown_param),\n",
    "    loc='upper left',\n",
    "    frameon=True)\n",
    "plt.setp(legend.get_title(), multialignment='center')\n",
    "\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Unfold Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "iterations = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xvals_1 = np.concatenate((theta0_S, theta_unknown_S))\n",
    "yvals_1 = np.concatenate((labels0, labels_unknown))\n",
    "\n",
    "xvals_2 = np.concatenate((theta0_G, theta0_G))\n",
    "yvals_2 = np.concatenate((labels0, labels_unknown))\n",
    "\n",
    "weights = np.empty(shape=(iterations, 2, len(theta0)))\n",
    "# shape = (iteration, step, event)\n",
    "\n",
    "inputs = Input((1, ))\n",
    "hidden_layer_1 = Dense(50, activation='relu')(inputs)\n",
    "hidden_layer_2 = Dense(50, activation='relu')(hidden_layer_1)\n",
    "hidden_layer_3 = Dense(50, activation='relu')(hidden_layer_2)\n",
    "outputs = Dense(1, activation='sigmoid')(hidden_layer_3)\n",
    "\n",
    "model = Model(inputs=inputs, outputs=outputs)\n",
    "\n",
    "earlystopping = EarlyStopping(patience=10,\n",
    "                              verbose=1,\n",
    "                              restore_best_weights=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from NN (DCTR)\n",
    "def reweight(events):\n",
    "    f = model.predict(events, batch_size=10000)\n",
    "    weights = f / (1. - f)\n",
    "    return np.squeeze(np.nan_to_num(weights))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initial iterative weights are ones\n",
    "\n",
    "weights_pull = np.ones(len(theta0_S))\n",
    "weights_push = np.ones(len(theta0_S))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(iterations):\n",
    "    print(\"\\nITERATION: {}\\n\".format(i + 1))\n",
    "\n",
    "    # STEP 1: classify Sim. (which is reweighted by weights_push) to Data\n",
    "    # weights reweighted Sim. --> Data\n",
    "    print(\"STEP 1\\n\")\n",
    "\n",
    "    weights_1 = np.concatenate((weights_push, np.ones(len(theta_unknown_S))))\n",
    "    # actual weights for Sim., ones for Data (not MC weights)\n",
    "\n",
    "    X_train_1, X_test_1, Y_train_1, Y_test_1, w_train_1, w_test_1 = train_test_split(\n",
    "        xvals_1, yvals_1, weights_1)\n",
    "\n",
    "    model.compile(loss='binary_crossentropy',\n",
    "                  optimizer='Adam',\n",
    "                  metrics=['accuracy'])\n",
    "    model.fit(X_train_1[X_train_1!=-10],\n",
    "              Y_train_1[X_train_1!=-10],\n",
    "              sample_weight=w_train_1[X_train_1!=-10],\n",
    "              epochs=200,\n",
    "              batch_size=10000,\n",
    "              validation_data=(X_test_1[X_test_1!=-10], Y_test_1[X_test_1!=-10], w_test_1[X_test_1!=-10]),\n",
    "              callbacks=[earlystopping],\n",
    "              verbose=1)\n",
    "\n",
    "    weights_pull = weights_push * reweight(theta0_S)\n",
    "    weights_pull[theta0_S==-10] = 1. #these are events that don't pass reco; take the prior.\n",
    "    weights[i, :1, :] = weights_pull\n",
    "\n",
    "    # STEP 2: classify Gen. to reweighted Gen. (which is reweighted by weights_pull)\n",
    "    # weights Gen. --> reweighted Gen.\n",
    "    print(\"\\nSTEP 2\\n\")\n",
    "\n",
    "    weights_2 = np.concatenate((np.ones(len(theta0_G)), weights_pull))\n",
    "    # ones for Gen. (not MC weights), actual weights for (reweighted) Gen.\n",
    "\n",
    "    X_train_2, X_test_2, Y_train_2, Y_test_2, w_train_2, w_test_2 = train_test_split(\n",
    "        xvals_2, yvals_2, weights_2)\n",
    "\n",
    "    model.compile(loss='binary_crossentropy',\n",
    "                  optimizer='Adam',\n",
    "                  metrics=['accuracy'])\n",
    "    model.fit(X_train_2,\n",
    "              Y_train_2,\n",
    "              sample_weight=w_train_2,\n",
    "              epochs=200,\n",
    "              batch_size=10000,\n",
    "              validation_data=(X_test_2, Y_test_2, w_test_2),\n",
    "              callbacks=[earlystopping],\n",
    "              verbose=1)\n",
    "\n",
    "    weights_push = reweight(theta0_G)\n",
    "    weights[i, 1:2, :] = weights_push"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins = np.linspace(-8, 8, 41)\n",
    "\n",
    "for i in range(5):\n",
    "    print(\"ITERATION {}:\".format(i + 1))\n",
    "    plt.figure(figsize=(6, 5))\n",
    "    hist0 = plt.hist(theta0_G[theta0_G!=0],\n",
    "                     bins=bins,\n",
    "                     label=r'($\\mu$, $\\sigma$) = {}'.format(theta0_param),\n",
    "                     **plot_style_2)\n",
    "\n",
    "    hist1 = plt.hist(\n",
    "        theta0_G[theta0_G!=0],\n",
    "        bins=bins,\n",
    "        label='($\\mu$, $\\sigma$) = {}\\nDCTR wgt.'.format(theta0_param),\n",
    "        weights=weights[i, 1, :][theta0_G!=0],\n",
    "        **plot_style_1)\n",
    "    hist2 = plt.hist(theta_unknown_G[theta_unknown_G!=0],\n",
    "                     bins=bins,\n",
    "                     label=r'($\\mu$, $\\sigma$) = ?',\n",
    "                     **plot_style_2)\n",
    "    plt.legend(frameon=False)\n",
    "    plt.xlabel(r\"$x_G$\")\n",
    "    plt.ylabel(\"Events per bin (normalized)\")\n",
    "    plt.show()"
   ]
  }
 ],
 "metadata": {
  "gist": {
   "data": {
    "description": "AcceptanceEffects",
    "public": false
   },
   "id": ""
  },
  "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.8"
  },
  "notify_time": "0",
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "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
}