{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "87b7f26c-c5f3-416b-a9b9-b331724fe929",
   "metadata": {},
   "source": [
    "# Processing many files using SLURM jobs\n",
    "_Author: Philipp Schreiner_<br>\n",
    "_Created: Aug. 29, 2025_<br>\n",
    "_Last updated: Jul. 2, 2026 by Philipp Schreiner_\n",
    "\n",
    "---\n",
    "\n",
    "Defining the elements of an analysis on a single file and scaling it up to many files can be done in many different ways (mostly depending on personal preference).\n",
    "\n",
    "Here, we show a version where you start from a notebook, send out the tasks to the cluster, and collect the results in the same notebook again, for you to continue your analysis. This approach easily scales and you can add more files by just extending a list of file names.\n",
    "\n",
    "We assume that at this point, you already know how to create SEVs and OFs, apply quality cuts and calculate spectra, and that you just want to trigger many files and perform those processing steps on them. Of course, the principle can be applied to all sorts of different tasks, so feel free to adapt the scripts to your needs.\n",
    "\n",
    "```{tip}\n",
    "Even though the procedure might look intimidating at first, we **recommend** to start following it as early as possible because it automatically results in a clean notebook/script structure even for many analysis files.\n",
    "```\n",
    "\n",
    "To get started, make sure to copy the [Files](#files) below to your working directory, together with the analysis notebook **`my_analysis.ipynb`** that we will be working in:\n",
    "\n",
    ">```\n",
    ">- analysis_info.py\n",
    ">- my_analysis.ipynb\n",
    ">- submit.py\n",
    ">- workflow.py\n",
    ">```\n",
    "\n",
    "We also recommend to create thre folders `files`, `job_files`, and `hdf5`, to keep things organized.\n",
    "\n",
    "Check the files that you copied for user specific variables (e.g. paths of streams, user names, container names, etc.) and update them.\n",
    "\n",
    "The most important one is `analysis_info.py`. Here, all analysis information like file paths, channel names, record length, trigger configurations, etc. are saved in one place.\n",
    "\n",
    "```{warning}\n",
    "This tutorial serves as an illustration. You cannot just run it like the previous ones.\n",
    "```\n",
    "\n",
    "```{tip}\n",
    "Here, we work in a single notebook, `my_analysis.ipynb`. What the author of this tutorial prefers is to have two notebooks. One for template/filter creation (e.g. `0_templates_and_filters.ipynb`) and one for the main analysis (e.g. `1_analysis.ipynb`). They share and use *the same* `analysis_info.py` file, which grows as the analysis develops. See below.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8ba74b9-c137-43da-90c0-5be2174eca77",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "84da9e8b-e7f2-4fa9-9bc2-b16582a6e742",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import scipy as sp\n",
    "\n",
    "import cait as ai\n",
    "import cait.versatile as vai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "609362bd-daa9-46cd-a3c2-65809003c708",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The analysis_info file stores things like the record length that we want to \n",
    "# use, and the SEVs/NPSs/OFs if we created them already. We load them here.\n",
    "from analysis_info import fdir, fdirh5, ch_name, record_length, OF, SEV, NPS\n",
    "\n",
    "# Have a look at analysis_info.py! You'll find a list of file names for the\n",
    "# template creation and the analysis. There, you can easily extend the list\n",
    "# even after you started this notebook to scale your analysis.\n",
    "from analysis_info import fnames_analysis as fnames\n",
    "\n",
    "n_channels = 2\n",
    "\n",
    "# This is just to keep things organised.\n",
    "# I usually format my HDF5 file names in that way.\n",
    "# For the appendix, I use things like '_templates' for when I'm creating SEVs, etc.\n",
    "# Feel free to adjust it.\n",
    "detector_name = \"myDet\"\n",
    "run = \"042\"\n",
    "hdf5_appendix = \"\"\n",
    "hdf5_name = f\"{detector_name}_run{run}{hdf5_appendix}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5152a934-8361-4435-a39d-15075164e3ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here we create a dictionary of all the stream files to organize them naicely.\n",
    "# The exact way you do it depends on the stream hardware. For VDAQ2 files, the one\n",
    "# below works quite well\n",
    "streams = {fname.split('.')[0]: vai.Stream(\"vdaq2\", os.path.join(fdir, fname)) for fname in fnames}\n",
    "streams"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9be754e0-b96e-4040-8094-1b6661fe62cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We also create a dictionary of DataHandlers where each entry corresponds to one stream file\n",
    "dhs = {\n",
    "    fname: ai.DataHandler(\n",
    "        record_length=record_length, \n",
    "        nmbr_channels=len(channel_names)+len(passive_names), \n",
    "        sample_frequency=s.sample_frequency,\n",
    "    )\n",
    "    for fname, stream in streams.items()\n",
    "}\n",
    "\n",
    "# Initialize the empty DataHandlers\n",
    "for k, v in dhs.items():\n",
    "    v.set_filepath(path_h5=fdirh5, fname=f\"{k}_{detector_name}{hdf5_appendix}\", appendix=False)\n",
    "    v.init_empty()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3002436b-7749-4ea3-ad1c-9bd0b26ab43b",
   "metadata": {},
   "source": [
    "### Preprocessing\n",
    "We now have `(Stream, DataHandler)`-pairs and we want to perform a number of operations on each pair (and save the results in the `DataHandler`s). What operations to perform is defined in the `STAGES_CONFIG` dictionary in the `analysis_info.py` file (have a look below!). We see triggering stages, one for main parameter calculation (CMP) and optimum filter/template fit amplitude reconstruction.\n",
    "\n",
    "#### Submitting to the cluster\n",
    "Let's say we want to run stages `[\"OF_TRIGGER\", \"CMP\", \"REAPPLY_OF\", \"TEMPLATE_FIT\"]`. For that, we can use the `define_job` function from the `workflow.py` script. After calling it, a `.json` file containing information for submitting SLURM jobs is created.\n",
    "\n",
    "The `fname` will be the name of the `.json` file. We recommend to store it in the `job_files` folder. The stages must be valid keys in the `STAGES_CONFIG` dictionary in the `analysis_info.py` file. If you want to run the job only on a subset of `Stream`s and `DataHandler`s, you can also filter them here. E.g. if you have `bck` and `ncal` files, you can define a job for `ncal` files only by doing something like `streams=[s for name, s in streams.items() if 'ncal' in name]`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be439d83-3128-4047-a688-e1290020f079",
   "metadata": {},
   "outputs": [],
   "source": [
    "from workflow import define_job\n",
    "\n",
    "define_job(\n",
    "    fname=\"job_files/of_trigger_cmp_of_tf\", \n",
    "    stages=[\"OF_TRIGGER\", \"CMP\", \"REAPPLY_OF\", \"TEMPLATE_FIT\"],\n",
    "    streams=list(streams.values()),\n",
    "    dhs=list(dhs.values())\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a1e3334-6bec-427f-afbd-22da87eac926",
   "metadata": {},
   "source": [
    "You can then log into your computing cluster, go to"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ca93a53-7950-4db1-98ff-248fcbc22358",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pwd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e612469c-08a8-48fa-a032-9eea3190ffa4",
   "metadata": {},
   "source": [
    "From here, we run\n",
    "```bash\n",
    "python3 submit.py job_files/of_trigger_cmp_of_tf.json\n",
    "```\n",
    "\n",
    "We can check if the jobs are running using `squeue -u $USER` or `slurm q` and what they are doing using `sattach <jobid>.0`.\n",
    "\n",
    "Additionally, job logs will be written to `job_files/logs` and you **must check them for problems**. If all jobs completed (i.e. they disappear from `slurm q`) and they didn't stop with an error (as seen from the logs), you can move on to the next step.\n",
    "\n",
    "```{seealso}\n",
    "Saving the ``DataHandler`` and ``Stream`` in ``.json`` files is done through a process called [serialization](https://en.wikipedia.org/wiki/Serialization). ``cait`` implements its own serialization mechanism which you can read more about in the documentation of [**`cait.serialize`**](cait.serialize).\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01c628fd-f554-46a9-a2ff-ac6eb3a93794",
   "metadata": {},
   "source": [
    "#### Running in notebook (not on cluster):\n",
    "\n",
    "Alternatively (if you don't want to use SLURM jobs), you can just run all calculations on the files sequentially in the notebook (otherwise this loop would just be run in parallel on the cluster, the result is identical):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61d0a8ae-da75-4db5-8bb3-d943d690683e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from workflow import job\n",
    "\n",
    "for dh, s in zip(dhs.values(), streams.values()):\n",
    "    print(f\"Working on file '{dh.fname}' ...\")\n",
    "    job(stages, s, dh)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a0dfd53-ee67-485c-9aeb-92b378d1b3c3",
   "metadata": {},
   "source": [
    "#### Saving file information\n",
    "Sometimes, you have different 'kinds of files'. E.g., background data files, 57Co calibration files, neutron calibration files, etc. You can either create separate `DataHandler`s for them, or (as preferred by the author of this tutorial) you can just have a single `DataHandler` and save flags that let you select the respective files easily."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1d89488-d091-48b7-92fb-aaf3d75af09d",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, (name, dh) in enumerate(dhs.items()):\n",
    "    for g in [\n",
    "        \"events\", \n",
    "        \"testpulses\", \n",
    "        \"controlpulses\",\n",
    "        \"noise\",\n",
    "    ]:\n",
    "        bck_flag = np.zeros(dh[f'{g}/hours'].shape[-1], dtype=bool)\n",
    "        ncal_flag = np.zeros(dh[f'{g}/hours'].shape[-1], dtype=bool)\n",
    "        file_index = i*np.ones(dh[f'{g}/hours'].shape[-1], dtype=np.int32)\n",
    "        \n",
    "        if name.startswith(\"bck\"):\n",
    "            bck_flag[:] = True\n",
    "        elif name.startswith(\"ncal\"):\n",
    "            ncal_flag[:] = True\n",
    "        else:\n",
    "            raise NotImplementedError\n",
    "\n",
    "        dh.set(g, flag_bck=bck_flag, flag_ncal=ncal_flag, dtype=bool, overwrite_existing=True)\n",
    "        dh.set(g, file_index=file_index, dtype=np.int32, overwrite_existing=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4337941f-0f94-4445-879e-383ccf136938",
   "metadata": {},
   "source": [
    "#### Combining the processed HDF5 files\n",
    "Running the job produced an HDF5 file for every (`Stream`, `DataHandler`)-pair. We now combine them to do analysis on a single `DataHandler`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb16e459-158c-4344-a483-bdaf8fe0b3de",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Combine HDF5 files into a single one for analysis\n",
    "ai.data.combine_h5(\n",
    "    fname=hdf5_name, \n",
    "    files=[dh.fname for dh in dhs.values()], \n",
    "    src_dir=fdirh5, \n",
    "    out_dir=fdirh5, \n",
    "    groups_combine=[\n",
    "        \"events\", \n",
    "        \"testpulses\", \n",
    "        \"noise\", \n",
    "        \"controlpulses\",\n",
    "    ],\n",
    ") "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1187d866-7a4f-4978-afc1-618bbd738eab",
   "metadata": {},
   "source": [
    "### Start Analysis\n",
    "Now that you have a single HDF5 file, you can continue your analysis as usual by creating a `DataHandler` instance:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f50aedd-581f-48f2-9d45-f42ef2593f62",
   "metadata": {},
   "outputs": [],
   "source": [
    "dh = ai.DataHandler(\n",
    "    record_length=record_length, \n",
    "    nmbr_channels=len(channel_names) + len(passive_names), \n",
    "    sample_frequency=list(streams.values())[0].sample_frequency,\n",
    ")\n",
    "dh.set_filepath(path_h5=fdirh5, fname=hdf5_name, appendix=False)\n",
    "print(dh)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fb3782e-ee7c-47e0-8764-810e3504f9ee",
   "metadata": {},
   "source": [
    "From there, you can define cuts, plot spectra, etc.\n",
    "\n",
    "```{note}\n",
    "We recommend to create two such notebooks. One for template/filter creation (e.g. `0_templates_and_filters.ipynb`) and one for the main analysis (e.g. `1_analysis.ipynb`). In the first one, the goal is just to save SEV, NPS, OF to `.xy` files and to estimate the baseline resolution, respectively threshold to be used for the of-trigger. Then, one starts with the of-trigger in the second notebook for the main analysis.\n",
    "```\n",
    "\n",
    "The working directory would then look something like this:\n",
    "\n",
    "```\n",
    "files\n",
    "    - OF.xy\n",
    "    - NPS.xy\n",
    "    - SEV.xy\n",
    "hdf5\n",
    "    - filter_creation_file0.h5\n",
    "    - filter_creation_file1.h5\n",
    "    - filter_creation.h5\n",
    "    - main_analysis_file0.h5\n",
    "    - main_analysis_file1.h5\n",
    "    - main_analysis_file2.h5\n",
    "    - main_analysis_file3.h5\n",
    "    - main_analysis.h5\n",
    "job_files\n",
    "    - trigger_templates.json\n",
    "    - trigger_of.json\n",
    "    - template_fit_and_reapply_of.json\n",
    "\n",
    "0_templates_and_filters.ipynb\n",
    "1_analysis.ipynb\n",
    "analysis_info.py\n",
    "submit.py\n",
    "workflow.py\n",
    "```\n",
    "\n",
    "```{note}\n",
    "Note that you don't have to define everything in the `analysis_info.py` script from the beginning. Usually, one starts with the z-score trigger and only runs `stages=[\"Z-SCORE_TRIGGER\", \"CMP\"]`. Then one creates SEV, NPS, OF and fills in the of-trigger information, then runs it again for `stages=[\"OF_TRIGGER\", \"CMP\"]`. After validating which parameters to use for `dh.apply_ofilter` and `dh.apply_template_fit` (using the `preview` feature), one fills it into `analysis_info.py` and runs `stages=[\"REAPPLY_OF\", \"TEMPLATE_FIT\"]`, and so on. In the end, all the relevant information is nicely collected in `analysis_info.py`. And if someone asks you which trigger threshold you used in your analysis, you know exactly where to find it.\n",
    "```\n",
    "\n",
    "```{tip}\n",
    "Often, the analysis on a few files is already done while more files are being written. The approach presented above is very convenient for that situation because you just have to extend the `fnames` list in `analysis_info.py` and run all stages again. Since the stages check whether the information has been computed before, only the new files are processed. Once the new files are processed, you combine the HDF5 files again and work with the same `DataHandler` as before.\n",
    "\n",
    "For this to be as convenient as possible, we recommend to **do all costly calculations on the single `DataHandler`s individually**, not on the combined `DataHandler` (since those results would be lost once you combine HDF5 files again). An example that loops through all single files is given in [Files](#saving-file-information).\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07d90949-2fbc-4114-be9a-4160abb3ede0",
   "metadata": {},
   "source": [
    "## Files\n",
    "Below, you find the script files that collect the analyis configuration and the helper functions for cluster job submission."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32175cd8-1e1e-4166-a387-8fc2d0413a9f",
   "metadata": {
    "tags": []
   },
   "source": [
    "### **`analysis_info.py`:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00319111-372e-4735-ba60-611de0240360",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import cait.versatile as vai\n",
    "\n",
    "# BASIC FILE INFO (added immediately)\n",
    "fdir = \"/where/the/stream/files/are/located/\"\n",
    "fdirh5 = \"/your/working/directory/\"\n",
    "\n",
    "# Files used for template/filter creation\n",
    "fnames_creation = [\n",
    "    \"bck_001\",\n",
    "    \"bck_002\",\n",
    "]\n",
    "\n",
    "# Files that are included in the analysis (in temporally increasing order)\n",
    "fnames_analysis = [\n",
    "    \"bck_001\",\n",
    "    \"bck_002\",\n",
    "    \"bck_003\",\n",
    "    \"ncal_001\",\n",
    "    \"ncal_002\",\n",
    "    \"ncal_003\",\n",
    "]\n",
    "\n",
    "# BASIC TRIGGER INFO (added immediately)\n",
    "ch_name = {\n",
    "    \"Phonon\": \"ADC1\",\n",
    "    \"Phonon_TP\": \"DAC1\",\n",
    "    \"Light\": \"ADC2\",\n",
    "    \"Light_TP\": \"DAC2\",\n",
    "}\n",
    "# Record length and random noise sampling frequency for triggering\n",
    "record_length = 2**14\n",
    "f_noise = 1000 # per hour\n",
    "\n",
    "# CONTROLPULSE INFO (added after inspecting the available TPAs)\n",
    "CP_above = {\n",
    "    \"Phonon\": 10, # V\n",
    "    \"Light\": 10, # V\n",
    "}\n",
    "\n",
    "# ENERGY RESOLUTIONS (added after calculating them in the template creation notebook)\n",
    "bl_res = {\n",
    "    \"Phonon\": 0.123e-3, # V\n",
    "    \"Light\": 1.234e-3, # V\n",
    "}\n",
    "\n",
    "# TRUNCATION LIMITS \n",
    "# (added after determining it in the template creation notebook)\n",
    "truncation_limit = {\n",
    "    \"Phonon\": 0.8, # V\n",
    "    \"Light\": 1.5, # V\n",
    "}\n",
    "\n",
    "# SEVs (added after building and saving them in the template creation notebook)\n",
    "SEV = {\n",
    "    \"Phonon\": vai.SEV.from_file(\"files/SEV_phonon\"),\n",
    "    \"Phonon_TP\": vai.SEV.from_file(\"files/SEV_phonon_TP\"),\n",
    "    \"Light\": vai.SEV.from_file(\"files/SEV_light\"),\n",
    "    \"Light_TP\": vai.SEV.from_file(\"files/SEV_light_TP\"),\n",
    "}\n",
    "\n",
    "\n",
    "# NPSs (added after building and saving them in the template creation notebook)\n",
    "NPS = {\n",
    "    \"Phonon\": vai.NPS.from_file(\"files/NPS_phonon\"),\n",
    "    \"Light\": vai.NPS.from_file(\"files/NPS_light\"),\n",
    "}\n",
    "\n",
    "# OFs (added after building and saving them in the template creation notebook)\n",
    "OF = {\n",
    "    \"Phonon\": vai.OF.from_file(\"files/OF_phonon\"),\n",
    "    \"Phonon_TP\": vai.OF.from_file(\"files/OF_phonon_TP\"),\n",
    "    \"Light\": vai.OF.from_file(\"files/OF_light\"),\n",
    "    \"Light_TP\": vai.OF.from_file(\"files/OF_light_TP\"),\n",
    "}\n",
    "\n",
    "# JOB STAGES (added and extended as the analysis progresses)\n",
    "STAGES_CONFIG = {\n",
    "    # ALWAYS DONE AFTER TRIGGERING (added immediately)\n",
    "    # Needs arguments for dh.cmp()\n",
    "    \"CMP\": [\n",
    "        { \"group\": \"events\" },\n",
    "        { \"group\": \"testpulses\" },\n",
    "        { \"group\": \"noise\" },\n",
    "        { \"group\": \"controlpulses\" },\n",
    "    ],\n",
    "    # FIRST ROUND OF TRIGGERING \n",
    "    # (added after 'controlpulses_above' is known)\n",
    "    # Needs arguments for dh.trigger_zscore()\n",
    "    \"Z-SCORE_TRIGGER\": {\n",
    "        \"trigger_channels\": [ch_name[\"Phonon\"]],\n",
    "        \"passive_channels\": [ch_name[\"Light\"]],\n",
    "        \"thresholds\": [5], # sigma\n",
    "        \"testpulse_channels\": [\n",
    "            ch_name[\"Phonon_TP\"], \n",
    "            ch_name[\"Light_TP\"],\n",
    "        ],\n",
    "        \"controlpulses_above\": [\n",
    "            CP_above[\"Phonon\"], \n",
    "            CP_above[\"Light\"],\n",
    "        ],\n",
    "        \"copy_events\": False,\n",
    "        \"f_noise\": f_noise,\n",
    "        \"reuse_triggers\": False,\n",
    "    },\n",
    "    # SECOND ROUND OF TRIGGERING \n",
    "    # (added after building templates/filters in the template creation notebook)\n",
    "    # Needs arguments for dh.trigger_of()\n",
    "    \"OF_TRIGGER\": {\n",
    "        \"trigger_channels\": [ch_name[\"Phonon\"]],\n",
    "        \"passive_channels\": [ch_name[\"Light\"]],\n",
    "        \"of\": OF[\"Phonon\"], \n",
    "        \"thresholds\": [bl_res[\"Phonon\"]*5],\n",
    "        \"testpulse_channels\": [\n",
    "            ch_name[\"Phonon_TP\"], \n",
    "            ch_name[\"Light_TP\"],\n",
    "        ],\n",
    "        \"controlpulses_above\": [\n",
    "            CP_above[\"Phonon\"], \n",
    "            CP_above[\"Light\"],\n",
    "        ],\n",
    "        \"copy_events\": False,\n",
    "        \"f_noise\": f_noise,\n",
    "        \"reuse_triggers\": False\n",
    "    },\n",
    "    # REAPPLYING THE FILTER\n",
    "    # (added after building templates/filters in the template creation notebook)\n",
    "    # Needs arguments for dh.apply_ofilter()\n",
    "    \"REAPPLY_OF\": [\n",
    "        {\n",
    "            \"group\": \"events\",\n",
    "            \"of\": OF[\"Phonon\"],\n",
    "            \"sev\": SEV[\"Light\"],\n",
    "            \"on_stream\": True,\n",
    "            \"tag\": \"uncorrelated\",\n",
    "        },\n",
    "        {\n",
    "            \"group\": \"events\",\n",
    "            \"of\": [OF[\"Phonon\"], OF[\"Light\"]],\n",
    "            \"sev\": [SEV[\"Phonon\"], SEV[\"Light\"]],\n",
    "            \"max_search\": [(0.2, 0.4), (-38, -28)],\n",
    "            \"relative_to\": [None, 0],\n",
    "            \"on_stream\": True,\n",
    "            \"tag\": \"correlated\",\n",
    "        },\n",
    "        # ... add more configurations as needed\n",
    "    ],\n",
    "    # TEMPLATE FIT\n",
    "    # (added after building templates/filters in the template creation notebook)\n",
    "    # Needs arguments for dh.apply_template_fit()\n",
    "    \"TEMPLATE_FIT\": [\n",
    "        {\n",
    "            \"group\": \"events\",\n",
    "            \"sev\": [SEV[\"Phonon\"], SEV[\"Light\"]],\n",
    "            \"bl_poly_order\": [3, 1], \n",
    "            \"truncation_limit\": [truncation_limit[\"Phonon\"], truncation_limit[\"Light\"]],\n",
    "            \"correlated\": True,\n",
    "            \"fit_onset\": [True, False],\n",
    "        },\n",
    "        # ... add more configurations as needed\n",
    "    ]\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ab17bcc-550e-4874-a5b6-683d1a039846",
   "metadata": {
    "tags": []
   },
   "source": [
    "### **`workflow.py`:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e6a4324-217d-4933-a136-288fe0788227",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import json\n",
    "\n",
    "import cait as ai\n",
    "import cait.versatile as vai\n",
    "\n",
    "from analysis_info import STAGES_CONFIG\n",
    "\n",
    "# Function that checks if all requested stages are present in the STAGES_CONFIG\n",
    "def _check_stages(stages: List[str]):\n",
    "    if not all([stage in STAGES_CONFIG.keys() for stage in stages]):\n",
    "        raise KeyError(f\"All selected stages must be present in the configured STAGES dictionary. Available: {list(STAGES.keys())}. Got: {stages}.\")\n",
    "\n",
    "def define_job(\n",
    "    fname: str, \n",
    "    stages: List[str], \n",
    "    streams: List[vai.stream.streambase.StreamBaseClass], \n",
    "    dhs: List[ai.DataHandler],\n",
    "):\n",
    "    \"\"\"\n",
    "    Save a json-file with 'fname' (omit .json extension) that defines a job consisting of stage names 'stages' and is applied to each (Stream, DataHandler)-pair in 'streams'/'dhs'.\n",
    "    \n",
    "    Jobs can be submitted via 'python3 submit.py fname.json'\n",
    "    \"\"\"\n",
    "    _check_stages(stages)\n",
    "    d = {\n",
    "        \"stages\": stages,\n",
    "        \"streams\": [s.to_dict() for s in streams],\n",
    "        \"dhs\": [d.to_dict() for d in dhs]\n",
    "    }\n",
    "    \n",
    "    with open(f\"{fname}.json\", \"w\") as f:\n",
    "        json.dump(d, f, indent=4)\n",
    "\n",
    "def job(\n",
    "    stages: List[str], \n",
    "    stream: vai.stream.streambase.StreamBaseClass, \n",
    "    dh: ai.DataHandler,\n",
    "):\n",
    "    \"\"\"Perform each analysis step defined in 'stages' to the (Stream, DataHandler)-pair.\"\"\"\n",
    "    _check_stages(stages)\n",
    "        \n",
    "    # In each stage, we check, whether the dataset is available already\n",
    "    # and only calculate it, if it's not\n",
    "    for stage in stages:\n",
    "        print(ai.styles.txt_fmt(f\"Entering stage {stage} ...\", color=\"yellow\", style=\"bold\"))\n",
    "        \n",
    "        if stage.startswith(\"Z-SCORE_TRIGGER\"):\n",
    "            if not dh.exists(f\"event_building-z-score\"):\n",
    "                dh.trigger_zscore(stream=stream, **STAGES_CONFIG[stage])\n",
    "                \n",
    "        if stage.startswith(\"OF_TRIGGER\"):\n",
    "            if not dh.exists(f\"event_building-of\"):\n",
    "                dh.trigger_of(stream=stream, **STAGES_CONFIG[stage])\n",
    "                \n",
    "        if stage.startswith(\"CMP\"):\n",
    "            for substage in STAGES_CONFIG[stage]:\n",
    "                has_events = True\n",
    "                try: \n",
    "                    # Note that the check 'dh.exists(\"group/events\")' is unreliable as the events\n",
    "                    # could be virtually stored. The easiest check is therefore to see if\n",
    "                    # get_event_iterator() fails\n",
    "                    dh.get_event_iterator(substage['group'])\n",
    "                except: \n",
    "                    has_events = False\n",
    "                if has_events:\n",
    "                    if not dh.exists(f\"{substage['group']}/pulse_height\"):\n",
    "                        dh.cmp(**substage)\n",
    "                    \n",
    "        if stage.startswith(\"REAPPLY_OF\"):\n",
    "            for substage in STAGES_CONFIG[stage]:\n",
    "                appendix = substage.get(\"tag\", \"\")\n",
    "                appendix = \"-\" + appendix if appendix else \"\"\n",
    "                print(f\"{substage['group']=}, {appendix=}\")\n",
    "                if dh.exists(f\"{substage['group']}\") and not dh.exists(f\"{substage['group']}/of_ph{appendix}\"):\n",
    "                    dh.apply_ofilter(**substage)\n",
    "                    \n",
    "        if stage.startswith(\"TEMPLATE_FIT\"):\n",
    "            for substage in STAGES_CONFIG[stage]:\n",
    "                appendix = substage.get(\"tag\", \"\")\n",
    "                appendix = \"-\" + appendix if appendix else \"\"\n",
    "                print(f\"{substage['group']=}, {appendix=}\")\n",
    "                if dh.exists(f\"{substage['group']}\") and not dh.exists(f\"{substage['group']}/templatefit_pars{appendix}\"):\n",
    "                    dh.apply_template_fit(**substage)\n",
    "                    \n",
    "        # YOU CAN ADD CUSTOM STAGE HANDLING BELOW. \n",
    "        # Add the respective entry in the STAGES_CONFIG dictionary in analysis_info.py and handle\n",
    "        # it as demonstrated above for 'CMP', 'REAPPLY_OF', ...\n",
    "        \n",
    "        # ...\n",
    "        \n",
    "        print(ai.styles.txt_fmt(f\"DONE with stage {stage} ..\", color=\"green\", style=\"bold\"))\n",
    "        \n",
    "if __name__ == \"__main__\":\n",
    "    import os\n",
    "    import sys\n",
    "    import json\n",
    "    \n",
    "    print(f\"Received script arguments: {sys.argv[1:]}\")\n",
    "    \n",
    "    json_path = sys.argv[1]\n",
    "    \n",
    "    try:\n",
    "        with open(json_path, \"r\") as f:\n",
    "            job_dict = json.load(f)\n",
    "\n",
    "        print(f\"Running job defined by dict: {job_dict}\")\n",
    "\n",
    "        job(\n",
    "            job_dict[\"stages\"], \n",
    "            ai.serialize.load(job_dict[\"stream\"]),\n",
    "            ai.serialize.load(job_dict[\"dh\"]),\n",
    "        )\n",
    "        \n",
    "    finally:\n",
    "        print(\"Cleaning up temporary file ...\")\n",
    "        try:\n",
    "            os.remove(json_path)\n",
    "            print(f\"Deleted temp file {json_path}\")\n",
    "        except Exception as e:\n",
    "            print(f\"Could not delete temp file {json_path}: {e}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "205a8ed0-327f-4636-8e71-45ffc11f249d",
   "metadata": {
    "tags": []
   },
   "source": [
    "### **`submit.py`:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e081e3be-875f-4fc7-b66d-da77f267f6f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import subprocess\n",
    "\n",
    "# The analysis directory (where the workflow.py, analysis_info.py and submit.py files are).\n",
    "DIR = os.path.dirname(os.path.abspath(__file__))\n",
    "\n",
    "# We will write temporary files and job logs into this directory\n",
    "JOB_FILES_DIR = f\"{DIR}/job_files\"\n",
    "\n",
    "# Path of the container\n",
    "CONTAINER = \"/path/to/cait_develop-slim.sif\"\n",
    "\n",
    "# Script to be executed in the jobs\n",
    "SCRIPT = \"workflow.py\"\n",
    "\n",
    "# Job name in the 'squeue -u $USER' list\n",
    "JOB_NAME = \"cait-analysis\"\n",
    "\n",
    "# RESOURCES\n",
    "TIME, CORES, MEMORY = \"02:00:00\", 4, \"16G\"\n",
    "\n",
    "# Validation function for a dictionary that defines a job\n",
    "def _check_valid_json(json_dict: dict):\n",
    "    for k in [\"stages\", \"streams\", \"dhs\"]:\n",
    "        if k not in json_dict.keys():\n",
    "            raise KeyError(f\"JSON files defining a valid job need to have a '{k}' key.\")\n",
    "\n",
    "        if not isinstance(json_dict[k], list):\n",
    "            raise TypeError(f\"JSON files defining a valid job need to have a '{k}' key which contains a list.\")\n",
    "            \n",
    "    if len(json_dict[\"streams\"]) != len(json_dict[\"dhs\"]):\n",
    "        raise ValueError(f\"Lists of streams and DataHandlers have to have same length. Got {len(json_dict['streams'])} and {len(json_dict['dhs'])}.\")\n",
    "        \n",
    "# Function that submits the job to the cluster's queue (for each (Stream, DataHandler)-pair separately; temp_json_files are created below)\n",
    "def submit_job(job_tag: str, temp_json_file: str):\n",
    "    command = f\"srun singularity run {CONTAINER} python {SCRIPT} {temp_json_file}\"\n",
    "    \n",
    "    sbatch_command = [\n",
    "        \"sbatch\",\n",
    "        f\"--job-name={JOB_NAME}-{job_tag}\",\n",
    "        f\"--chdir={DIR}\",\n",
    "        f\"--output={JOB_FILES_DIR}/logs/{job_tag}.out\",\n",
    "        f\"--error={JOB_FILES_DIR}/logs/{job_tag}.err\",\n",
    "        f\"--time={TIME}\",\n",
    "        f\"--mem={MEMORY}\",\n",
    "        f\"--cpus-per-task={CORES}\",\n",
    "        \"--partition=c\",\n",
    "        \"--qos=c_medium\",\n",
    "        \"--parsable\",\n",
    "        \"--wrap\", command\n",
    "    ]\n",
    "    \n",
    "    print(f\"Job command: {sbatch_command}\")\n",
    "    \n",
    "    subprocess.run(sbatch_command, check=True)\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    import sys\n",
    "    import json\n",
    "    import tempfile\n",
    "    \n",
    "    # Load json information\n",
    "    json_path = sys.argv[1]\n",
    "    # Name of json file is used to make the job more\n",
    "    # recognizable (job name and log folder)\n",
    "    job_tag = os.path.basename(json_path).split(\".\")[0]\n",
    "    \n",
    "    os.makedirs(JOB_FILES_DIR, exist_ok=True)\n",
    "    \n",
    "    # Load the .json file that defines all job stages for all (Stream, DataHandler)-pairs\n",
    "    with open(json_path, \"r\") as f:\n",
    "        job_dict = json.load(f)\n",
    "    \n",
    "    # Validate it\n",
    "    _check_valid_json(job_dict)\n",
    "    stages = job_dict[\"stages\"]\n",
    "    \n",
    "    # Submit a job for each (Stream, DataHandler)-pair in the json file.\n",
    "    # For that purpose, a smaller json file (per job) is created to\n",
    "    # transfer them to the job script more easily.\n",
    "    for i, (s, dh) in enumerate(zip(job_dict[\"streams\"], job_dict[\"dhs\"])):\n",
    "        with tempfile.NamedTemporaryFile(mode=\"w\", delete=False, suffix=\".json\", dir=JOB_FILES_DIR) as tmpfile:\n",
    "            json.dump({\"stages\": stages, \"stream\": s, \"dh\": dh}, tmpfile)\n",
    "            tmpfile_path = tmpfile.name\n",
    "            \n",
    "            # Ensure the file is readable (some clusters might restrict tmp file access)\n",
    "            os.chmod(tmpfile_path, 0o644)\n",
    " \n",
    "        submit_job(f\"{job_tag}-{i}\", tmpfile_path)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv_cait",
   "language": "python",
   "name": "venv_cait"
  },
  "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.11.9"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}