{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "839ca214",
   "metadata": {},
   "outputs": [],
   "source": [
    "import awkward\n",
    "import numpy as np\n",
    "from coffea import nanoevents\n",
    "events = nanoevents\\\n",
    "         .NanoEventsFactory\\\n",
    "         .from_root('../tests/samples/nano_dy.root')\\\n",
    "         .events()\n",
    "\n",
    "def some_event_weight(ones):\n",
    "    return (1.0 + np.array([0.05, -0.05], dtype=np.float32)) * ones[:, None]\n",
    "\n",
    "events.add_systematic(\"RenFactScale\", \"UpDownSystematic\", \"weight\", some_event_weight)\n",
    "events.add_systematic(\"XSectionUncertainty\", \"UpDownSystematic\", \"weight\", some_event_weight)\n",
    "\n",
    "muons = events.Muon\n",
    "jets = events.Jet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b803b91b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "def muon_pt_scale(pt):   \n",
    "    return (1.0 + np.array([0.05, -0.05], dtype=np.float32)) * pt[:, None]\n",
    "\n",
    "def muon_pt_resolution(pt):   \n",
    "    return np.random.normal(pt[:,None], np.array([0.02, 0.01], dtype=np.float32))\n",
    "\n",
    "def muon_eff_weight(ones):\n",
    "    return (1.0 + np.array([0.05, -0.05], dtype=np.float32)) * ones[:, None]\n",
    "\n",
    "muons.add_systematic(\"PtScale\", \"UpDownSystematic\", \"pt\", muon_pt_scale)\n",
    "muons.add_systematic(\"PtResolution\", \"UpDownSystematic\", \"pt\", muon_pt_resolution)\n",
    "muons.add_systematic(\"EfficiencySF\", \"UpDownSystematic\", \"weight\", muon_eff_weight)\n",
    "\n",
    "def jet_pt_scale(pt):   \n",
    "    return (1.0 + np.array([0.10, -0.10], dtype=np.float32)) * pt[:, None]\n",
    "\n",
    "def jet_pt_resolution(pt):   \n",
    "    return np.random.normal(pt[:,None], np.array([0.20, 0.10], dtype=np.float32))\n",
    "\n",
    "jets.add_systematic(\"PtScale\", \"UpDownSystematic\", \"pt\", jet_pt_scale)\n",
    "jets.add_systematic(\"PtResolution\", \"UpDownSystematic\", \"pt\", jet_pt_resolution)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b88cf664",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "awkward.flatten(muons.pt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c89f8fd1",
   "metadata": {},
   "outputs": [],
   "source": [
    "awkward.flatten(muons.systematics.PtScale.up)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f484489",
   "metadata": {},
   "outputs": [],
   "source": [
    "awkward.flatten(muons.systematics.PtScale.up.pt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c806b7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "awkward.flatten(jets.pt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7adf5de",
   "metadata": {},
   "outputs": [],
   "source": [
    "awkward.flatten(jets.systematics.PtScale.up)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f01ca3a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "awkward.flatten(jets.systematics.PtScale.up.pt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89f8dee6",
   "metadata": {},
   "outputs": [],
   "source": [
    "events.systematics.RenFactScale.up.weight_RenFactScale"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3591fbf3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Make it so that syst_muons.Y > X returns boolean values\n",
    "#       for all variations over Y. \n",
    "#       Requires some tracking of (pieces of) \"what\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2184b3bd",
   "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}