From 593dfe9f48374c5eb998ee5378299736a4ce8ee1 Mon Sep 17 00:00:00 2001
From: juacrumar <juacrumar@lairen.eu>
Date: Thu, 18 Apr 2024 11:22:24 +0200
Subject: [PATCH] add the saving of the data as pkl files

---
 .../examples/API_extension_Pineappl.ipynb     | 86 ++++++++++++++++++-
 1 file changed, 82 insertions(+), 4 deletions(-)

diff --git a/validphys2/examples/API_extension_Pineappl.ipynb b/validphys2/examples/API_extension_Pineappl.ipynb
index d2be9e3425..9354ccbce2 100644
--- a/validphys2/examples/API_extension_Pineappl.ipynb
+++ b/validphys2/examples/API_extension_Pineappl.ipynb
@@ -292,13 +292,91 @@
     "figs[0]"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "f8302287-5b04-47b0-99c4-0c621639ef1f",
+   "metadata": {},
+   "source": [
+    "The cell below creates a `.pkl` format with all results with the following structure:\n",
+    "\n",
+    "```\n",
+    "    exp_data: central experimental data\n",
+    "    exp_uncert: experimental uncertainties (+- 1sigma)\n",
+    "    predictions: theoretical predictions for the central scale\n",
+    "    pdf_uncert: pdf uncertainties\n",
+    "    theory_uncert: scale uncertainties\n",
+    "```\n",
+    "\n",
+    "so that it can be reused in other functions."
+   ]
+  },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 6,
    "id": "e3877c80-d3dd-4dc3-b14d-a54c40cff6b2",
    "metadata": {},
-   "outputs": [],
-   "source": []
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Computing predictions for H1_2JET_319GEV_290PB-1_DIF_PTQ2 with NNPDF40_nnlo_as_01180 (tid=600)\n",
+      "Predictions for H1_2JET_319GEV_290PB-1_DIF_PTQ2 written to /tmp/results_H1_2JET_319GEV_290PB-1_DIF_PTQ2_600_NNPDF40_nnlo_as_01180.pkl (pdf='NNPDF40_nnlo_as_01180')\n",
+      "Computing predictions for H1_2JET_319GEV_290PB-1_DIF_PTQ2 with PDF4LHC21_mc (tid=600)\n",
+      "Predictions for H1_2JET_319GEV_290PB-1_DIF_PTQ2 written to /tmp/results_H1_2JET_319GEV_290PB-1_DIF_PTQ2_600_PDF4LHC21_mc.pkl (pdf='PDF4LHC21_mc')\n"
+     ]
+    }
+   ],
+   "source": [
+    "def create_data_pickle(dname, pdf, theoryid, force=False, output_folder=Path(\".\")):\n",
+    "    data_path = Path(output_folder) / f\"results_{dname}_{theoryid}_{pdf}.pkl\"\n",
+    "     \n",
+    "    if data_path.exists() and not force:\n",
+    "        print(f\"Skipping {dname} ({pdf}) since {data_path} already exists. If you want to overwrite it use `force=True`\")\n",
+    "        return\n",
+    "        \n",
+    "    print(f\"Computing predictions for {dname} with {pdf} (tid={theoryid})\")\n",
+    "    generic = {\n",
+    "        \"dataset_input\": {\"dataset\": dname},\n",
+    "        \"theoryid\": theoryid,\n",
+    "        \"pdf\": pdf,\n",
+    "        \"use_cuts\": \"internal\",\n",
+    "    }\n",
+    "\n",
+    "    # These two functions are using the monkey patching from before\n",
+    "    # so you better make sure they are being used in the same notebook!\n",
+    "    # otherwise you will be getting FKTable results and not grids!!!!!!1111\n",
+    "\n",
+    "    res_data, res_theory = API.results(**generic)\n",
+    "    theory_covmat = API.theory_covmat_dataset(\n",
+    "        **generic,\n",
+    "        point_prescription=\"9 point\",\n",
+    "        theoryids={\"from_\": \"scale_variation_theories\"},\n",
+    "        use_theorycovmat=True,\n",
+    "    )\n",
+    "    ret = {\n",
+    "        \"exp_data\": res_data.central_value,\n",
+    "        \"predictions\": res_theory.central_value,\n",
+    "        \"exp_uncert\": res_data.std_error,\n",
+    "        \"pdf_uncert\": res_theory.std_error,\n",
+    "        \"theory_uncert\": np.sqrt(np.diag(theory_covmat)),\n",
+    "    }\n",
+    "    ds = API.dataset(**generic)\n",
+    "    # Read the cuts \n",
+    "    cuts = ds.cuts.load()+1 # kinematics start at 1\n",
+    "    kin_df = ds.commondata.metadata.load_kinematics().loc[cuts]\n",
+    "    results_df = pd.DataFrame(ret, index=kin_df.index)\n",
+    "\n",
+    "    results_df.to_pickle(data_path)\n",
+    "    print(f\"Predictions for {dname} written to {data_path} ({pdf=})\")\n",
+    "\n",
+    "#Example:\n",
+    "all_datasets = [\"H1_2JET_319GEV_290PB-1_DIF_PTQ2\"]\n",
+    "if True:\n",
+    "    for pdf in [\"NNPDF40_nnlo_as_01180\", \"PDF4LHC21_mc\"]:\n",
+    "        for dataset in all_datasets:\n",
+    "            create_data_pickle(dataset, pdf, 600, force=True, output_folder=Path(\"/tmp\"))"
+   ]
   },
   {
    "cell_type": "code",
@@ -325,7 +403,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.2"
+   "version": "3.12.3"
   }
  },
  "nbformat": 4,