From eda62b38bd2ae55278535bff7d2dc063a3c2cdc7 Mon Sep 17 00:00:00 2001
From: tyler-a-cox <tyler.a.cox@berkeley.edu>
Date: Sun, 28 Jul 2024 11:44:14 -0600
Subject: [PATCH] update lststack notebook

---
 .../notebooks/lststack.ipynb                  | 37 +++++++++++--------
 1 file changed, 22 insertions(+), 15 deletions(-)

diff --git a/hera_notebook_templates/notebooks/lststack.ipynb b/hera_notebook_templates/notebooks/lststack.ipynb
index 7aa829e..ca1c7e8 100644
--- a/hera_notebook_templates/notebooks/lststack.ipynb
+++ b/hera_notebook_templates/notebooks/lststack.ipynb
@@ -66,6 +66,7 @@
    "source": [
     "import os\n",
     "import sys\n",
+    "import pickle\n",
     "from pathlib import Path\n",
     "from functools import partial\n",
     "from datetime import datetime\n",
@@ -759,6 +760,7 @@
    "source": [
     "%%time\n",
     "if do_lstcal and n_bl_chunks==1: # can't fit all the baselines in memory if not redavg'd, and need all of them at once to do lstcal\n",
+    "    all_calibration_parameters = {}\n",
     "    for i, (stack, lstavg_model, auto_model) in enumerate(zip(cross_stacks, cross_lstavg, autos_lstavg)):\n",
     "        calibration_parameters, gains = lstbin_absolute_calibration(\n",
     "            stack=stack, \n",
@@ -768,8 +770,9 @@
     "            auto_stack=auto_stacks[i],\n",
     "            auto_model=auto_model['data'],\n",
     "            calibrate_inplace=True, # calibrate inplace\n",
-    "            run_phase_cal=True, # run phase calibration\n",
     "            run_amplitude_cal=True, # run amplitude calibration\n",
+    "            run_phase_cal=True, # run phase calibration\n",
+    "            run_cross_pol_phase_cal=True\n",
     "        )\n",
     "            \n",
     "        # Write out calibration parameters and metadata\n",
@@ -779,18 +782,22 @@
     "        calibration_parameters[\"antpairs\"] = stack.antpairs\n",
     "        calibration_parameters[\"lst\"] = stackconf.lst_grid[i]\n",
     "        \n",
-    "        # Get the calibration filename\n",
-    "        cal_fname = lstbin.io.format_outfile_name(\n",
-    "            lst=stackconf.lst_grid_edges[i],\n",
-    "            pols=stackconf.pols,\n",
-    "            fname_format=fname_format.replace(\"sum.uvh5\", \"lstcal.npz\"),\n",
-    "            inpaint_mode=True,\n",
-    "            lst_branch_cut=stackconf.properties[\"lst_branch_cut\"],\n",
-    "            kind=\"LST\"\n",
-    "        )\n",
-    "        outfile = outdir / cal_fname\n",
+    "        all_calibration_parameters[stackconf.lst_grid[i]] = calibration_parameters\n",
+    "        \n",
+    "    # Get the calibration filename\n",
+    "    cal_fname = lstbin.io.format_outfile_name(\n",
+    "        lst=stackconf.lst_grid_edges[0],\n",
+    "        pols=stackconf.pols,\n",
+    "        fname_format=fname_format.replace(\"sum.uvh5\", \"lstcal.pkl\"),\n",
+    "        inpaint_mode=True,\n",
+    "        lst_branch_cut=stackconf.properties[\"lst_branch_cut\"],\n",
+    "        kind=\"LST\"\n",
+    "    )\n",
+    "    outfile = outdir / cal_fname\n",
     "        \n",
-    "        np.savez(outfile, **calibration_parameters)\n",
+    "    # Write out calibration parameters\n",
+    "    with open(outfile, 'wb') as handle:\n",
+    "        pickle.dump(all_calibration_parameters, handle, protocol=pickle.HIGHEST_PROTOCOL)\n",
     "            \n",
     "        \n",
     "    # Recompute auto and cross averages after calibration - needed for STD files \n",
@@ -2754,9 +2761,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "h6c",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
-   "name": "h6c"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
@@ -2768,7 +2775,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.6"
+   "version": "3.11.3"
   },
   "toc": {
    "base_numbering": 1,