Merge pull request #790 from yimuchen/correctionlib_wrapper_docs

docs: Added example code for correctionlib_wrapper
CoffeaTeam · Jun 4, 2023 · 492e38d · 492e38d
2 parents 3aaddc5 + ca59c7c
commit 492e38d
Showing 1 changed file with 82 additions and 2 deletions.
diff --git a/binder/applying_corrections.ipynb b/binder/applying_corrections.ipynb
@@ -623,10 +623,11 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "At the time of writing, correctionlib does not support jagged arrays. But we can work around this using awkward utilities `flatten` and `unflatten`, as shown below"
+    "At the time of writing, `correctionlib` does not support jagged arrays. A `correctionlib_wrapper` provided in `coffea.lookup_tools` allows for the processing of jagged array inputs."
    ]
   },
   {
@@ -636,15 +637,94 @@
    "outputs": [
     {
      "data": {
+      "text/html": [
+       "<pre>[[0.994, 0.313, 1.08, 0.904, 0.904],\n",
+       " [0.142, 0.352, 0.895, 0.73, 0.904, 0.73, 1.16, 0.904],\n",
+       " [1, 0.522, 0.696, 0.582, 0.65],\n",
+       " [0.663, 0.576, 0.73],\n",
+       " [0.284, 0.124, 1.08, 0.73, 0.73],\n",
+       " [0.844, 0.352, 0.84, 1.12, 0.827, 0.73, 0.904, 1.22],\n",
+       " [0.994, 0.521, 0.576, 0.864],\n",
+       " [0.576, 0.864, 0.827, 0.73],\n",
+       " [0.827],\n",
+       " [0.249, 0.68, 0.352, 0.84, 0.864, 1.08, 0.904, 1.16, 0.904],\n",
+       " ...,\n",
+       " [0.929, 1.12],\n",
+       " [0.827, 0.904, 0.904, 0.904],\n",
+       " [0.313, 0.54, 0.864, 1.01, 1.22, 0.73, 1.16, 1.16, 0.904],\n",
+       " [0.76, 1.08, 1.12],\n",
+       " [1.01, 1.22],\n",
+       " [0.895, 0.65, 1.22],\n",
+       " [1.07],\n",
+       " [0.351, 0.643, 0.65, 0.827, 0.65, 1.22],\n",
+       " [0.73, 1.16]]\n",
+       "-------------------------------------------------------------\n",
+       "type: 40 * var * float64</pre>"
+      ],
       "text/plain": [
-       "<Array [[0.498, 0.211, 0.86, ... [1.15, 1.17]] type='40 * var * float64'>"
+       "<Array [[0.994, 0.313, 1.08, 0.904, 0.904], ...] type='40 * var * float64'>"
       ]
      },
      "execution_count": 15,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
+   "source": [
+    "from coffea.lookup_tools.correctionlib_wrapper import correctionlib_wrapper\n",
+    "\n",
+    "wrap_c = correctionlib_wrapper(ceval[\"gen2_to_gen1\"])\n",
+    "wrap_c(events.Jet.pt, events.Jet.eta)"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Alternatively, we can use the awkward utilities `flatten` and `unflatten` to convert awkward arrays into numpy arrays for evaluation."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<pre>[[0.994, 0.313, 1.08, 0.904, 0.904],\n",
+       " [0.142, 0.352, 0.895, 0.73, 0.904, 0.73, 1.16, 0.904],\n",
+       " [1, 0.522, 0.696, 0.582, 0.65],\n",
+       " [0.663, 0.576, 0.73],\n",
+       " [0.284, 0.124, 1.08, 0.73, 0.73],\n",
+       " [0.844, 0.352, 0.84, 1.12, 0.827, 0.73, 0.904, 1.22],\n",
+       " [0.994, 0.521, 0.576, 0.864],\n",
+       " [0.576, 0.864, 0.827, 0.73],\n",
+       " [0.827],\n",
+       " [0.249, 0.68, 0.352, 0.84, 0.864, 1.08, 0.904, 1.16, 0.904],\n",
+       " ...,\n",
+       " [0.929, 1.12],\n",
+       " [0.827, 0.904, 0.904, 0.904],\n",
+       " [0.313, 0.54, 0.864, 1.01, 1.22, 0.73, 1.16, 1.16, 0.904],\n",
+       " [0.76, 1.08, 1.12],\n",
+       " [1.01, 1.22],\n",
+       " [0.895, 0.65, 1.22],\n",
+       " [1.07],\n",
+       " [0.351, 0.643, 0.65, 0.827, 0.65, 1.22],\n",
+       " [0.73, 1.16]]\n",
+       "-------------------------------------------------------------\n",
+       "type: 40 * var * float64</pre>"
+      ],
+      "text/plain": [
+       "<Array [[0.994, 0.313, 1.08, 0.904, 0.904], ...] type='40 * var * float64'>"
+      ]
+     },
+     "execution_count": 16,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "source": [
     "def myJetSF(jets):\n",
     "    j, nj = ak.flatten(jets), ak.num(jets)\n",