tutorials/T01-DataPreparation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "84d8a02e-05df-4638-a426-3a0445407442",
   "metadata": {},
   "source": [
    "# Tutorial 01 - Data Preparation\n",
    "\n",
    "This tutorial guides you towards the creation of the training data used in this study."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "849d957b-f833-4ff0-b20d-7b905729ed22",
   "metadata": {},
   "source": [
    "### Downloading inputs\n",
    "\n",
    "The original inputs are available from Zenodo for various samples:\n",
    "\n",
    "- QCD background from official LHCO dataset: https://zenodo.org/records/4536377/files/events_anomalydetection_v2.features.h5\n",
    "- Extra QCD background : https://zenodo.org/records/8370758/files/events_anomalydetection_qcd_extra_inneronly_features.h5\n",
    "- Parametric W->X(qq)Y(qq) signal : https://zenodo.org/records/11188685/files/events_anomalydetection_Z_XY_qq_parametric.h5\n",
    "- Parametric W->X(qq)Y(qqq) signal : https://zenodo.org/records/11188685/files/events_anomalydetection_Z_XY_qqq_parametric.h5\n",
    "\n",
    "You can download the inputs manually and put them according to the file structure `<datadir>/original/<input>.h5`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b309c74-ba49-4bb8-abb9-c449b9de981e",
   "metadata": {},
   "source": [
    "A more straight forward way to download the inputs and make sure it has the correct file structure is to use the paws Command Line Interface (CLI)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b86e47a0-6ee8-45c9-bce9-5862b3803b90",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Usage: paws download_data [OPTIONS]\n",
      "\n",
      "  Download datasets used in this study.\n",
      "\n",
      "Options:\n",
      "  -s, --samples [QCD|extra_QCD|W_qq|W_qqq]\n",
      "                                  List of data samples to download (separated\n",
      "                                  by commas). Available samples are: ['QCD',\n",
      "                                  'extra_QCD', 'W_qq', 'W_qqq']. By default,\n",
      "                                  all samples will be downloaded.\n",
      "  -d, --datadir TEXT              Base directory for storing datasets. The\n",
      "                                  downloaded data will be stored in\n",
      "                                  <datadir>/raw.  [default: datasets]\n",
      "  --help                          Show this message and exit.\n"
     ]
    }
   ],
   "source": [
    "!paws download_data --help"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d1a619d8-eb63-4424-82c7-959651cd7fc3",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] Downloading sample \"QCD\" from https://zenodo.org/records/4536377/files/events_anomalydetection_v2.features.h5\n",
      "--2024-06-01 19:36:24--  https://zenodo.org/records/4536377/files/events_anomalydetection_v2.features.h5\n",
      "Resolving zenodo.org (zenodo.org)... 2001:1458:d00:3b::100:200, 2001:1458:d00:9::100:195, 2001:1458:d00:3a::100:33a, ...\n",
      "Connecting to zenodo.org (zenodo.org)|2001:1458:d00:3b::100:200|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 74315238 (71M) [application/octet-stream]\n",
      "Saving to: ‘datasets/original/events_anomalydetection_v2.features.h5’\n",
      "\n",
      "datasets/original/e 100%[===================>]  70.87M  23.1MB/s    in 3.9s    \n",
      "\n",
      "2024-06-01 19:36:29 (18.2 MB/s) - ‘datasets/original/events_anomalydetection_v2.features.h5’ saved [74315238/74315238]\n",
      "\n",
      "[INFO] File downloaded to datasets/original/events_anomalydetection_v2.features.h5\n",
      "[INFO] Downloading sample \"extra_QCD\" from https://zenodo.org/records/8370758/files/events_anomalydetection_qcd_extra_inneronly_features.h5\n",
      "--2024-06-01 19:36:29--  https://zenodo.org/records/8370758/files/events_anomalydetection_qcd_extra_inneronly_features.h5\n",
      "Resolving zenodo.org (zenodo.org)... 2001:1458:d00:3b::100:200, 2001:1458:d00:9::100:195, 2001:1458:d00:3a::100:33a, ...\n",
      "Connecting to zenodo.org (zenodo.org)|2001:1458:d00:3b::100:200|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 38554784 (37M) [application/octet-stream]\n",
      "Saving to: ‘datasets/original/events_anomalydetection_qcd_extra_inneronly_features.h5’\n",
      "\n",
      "datasets/original/e 100%[===================>]  36.77M  15.0MB/s    in 2.5s    \n",
      "\n",
      "2024-06-01 19:36:32 (15.0 MB/s) - ‘datasets/original/events_anomalydetection_qcd_extra_inneronly_features.h5’ saved [38554784/38554784]\n",
      "\n",
      "[INFO] File downloaded to datasets/original/events_anomalydetection_qcd_extra_inneronly_features.h5\n",
      "[INFO] Downloading sample \"W_qq\" from https://zenodo.org/records/11188685/files/events_anomalydetection_Z_XY_qq_parametric.h5\n",
      "--2024-06-01 19:36:32--  https://zenodo.org/records/11188685/files/events_anomalydetection_Z_XY_qq_parametric.h5\n",
      "Resolving zenodo.org (zenodo.org)... 2001:1458:d00:3b::100:200, 2001:1458:d00:9::100:195, 2001:1458:d00:3a::100:33a, ...\n",
      "Connecting to zenodo.org (zenodo.org)|2001:1458:d00:3b::100:200|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 2435840888 (2.3G) [application/octet-stream]\n",
      "Saving to: ‘datasets/original/events_anomalydetection_Z_XY_qq_parametric.h5’\n",
      "\n",
      "datasets/original/e 100%[===================>]   2.27G  24.5MB/s    in 1m 44s  \n",
      "\n",
      "2024-06-01 19:38:17 (22.4 MB/s) - ‘datasets/original/events_anomalydetection_Z_XY_qq_parametric.h5’ saved [2435840888/2435840888]\n",
      "\n",
      "[INFO] File downloaded to datasets/original/events_anomalydetection_Z_XY_qq_parametric.h5\n",
      "[INFO] Downloading sample \"W_qqq\" from https://zenodo.org/records/11188685/files/events_anomalydetection_Z_XY_qqq_parametric.h5\n",
      "--2024-06-01 19:38:17--  https://zenodo.org/records/11188685/files/events_anomalydetection_Z_XY_qqq_parametric.h5\n",
      "Resolving zenodo.org (zenodo.org)... 2001:1458:d00:3b::100:200, 2001:1458:d00:9::100:195, 2001:1458:d00:3a::100:33a, ...\n",
      "Connecting to zenodo.org (zenodo.org)|2001:1458:d00:3b::100:200|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 2424698488 (2.3G) [application/octet-stream]\n",
      "Saving to: ‘datasets/original/events_anomalydetection_Z_XY_qqq_parametric.h5’\n",
      "\n",
      "datasets/original/e 100%[===================>]   2.26G  25.3MB/s    in 1m 49s  \n",
      "\n",
      "2024-06-01 19:40:07 (21.1 MB/s) - ‘datasets/original/events_anomalydetection_Z_XY_qqq_parametric.h5’ saved [2424698488/2424698488]\n",
      "\n",
      "[INFO] File downloaded to datasets/original/events_anomalydetection_Z_XY_qqq_parametric.h5\n"
     ]
    }
   ],
   "source": [
    "# to download all input samples with datadir = \"datasets\"\n",
    "!paws download_data -d datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9b162c2-5f6f-449e-85bf-fbe2809d711f",
   "metadata": {},
   "source": [
    "Alternatively, you may use the paws API:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7c3e0cc4-52ef-4f02-8ae3-25245ab10370",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "datadir = datasets\n",
      "samples = ['QCD', 'extra_QCD', 'W_qq', 'W_qqq']\n"
     ]
    }
   ],
   "source": [
    "from paws import PathManager\n",
    "from paws.data_preparation import download_file\n",
    "from paws.settings import Sample, SampleURLs\n",
    "\n",
    "# modify your base data directory here\n",
    "datadir = \"datasets\"\n",
    "\n",
    "# choose the samples to download, you will most likely download all of them\n",
    "samples = ['QCD', 'extra_QCD', 'W_qq', 'W_qqq']\n",
    "\n",
    "print(f\"datadir = {datadir}\")\n",
    "print(f\"samples = {samples}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "66ed45c4-8229-4eb3-b89a-af98af70de67",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{<Sample.QCD: 0>: 'https://zenodo.org/records/4536377/files/events_anomalydetection_v2.features.h5',\n",
       " <Sample.EXTRA_QCD: 1>: 'https://zenodo.org/records/8370758/files/events_anomalydetection_qcd_extra_inneronly_features.h5',\n",
       " <Sample.W_QQ: 2>: 'https://zenodo.org/records/11188685/files/events_anomalydetection_Z_XY_qq_parametric.h5',\n",
       " <Sample.W_QQQ: 3>: 'https://zenodo.org/records/11188685/files/events_anomalydetection_Z_XY_qqq_parametric.h5'}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# check the urls for each sample\n",
    "SampleURLs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "10697dbb-48d6-4964-9102-2e93643544a0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "outdir = datasets/original\n"
     ]
    }
   ],
   "source": [
    "# the actual output directory is handled by PathManager\n",
    "path_manager = PathManager(directories={\"dataset\": datadir})\n",
    "outdir = path_manager.get_directory('original_dataset')\n",
    "print(f\"outdir = {outdir}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8357d30e-aacc-4dbd-be56-cbe12e8b301f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# now download the inputs\n",
    "for sample in samples:\n",
    "    url = SampleURLs[Sample.parse(sample)]\n",
    "    print(f'Downloading sample \"{sample}\" from {url}')\n",
    "    download_file(url, outdir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95808aeb-acf1-47a4-8aea-9af1928b7a99",
   "metadata": {},
   "source": [
    "### Prepare dedicated datasets\n",
    "\n",
    "Dedicated datasets are the processed datasets in tfrecord format per sample per mass point, i.e. (mX, mY). For the background, the dataset is repeated for each mass point.\n",
    "\n",
    "These datasets are used for training of the datacated supervised and the weakly supervised models which we train over signal of a specific mass point.\n",
    "\n",
    "To use the CLI:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a20bf74e-6cef-41bf-b72f-9861b074b408",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Usage: paws create_dedicated_datasets [OPTIONS]\n",
      "\n",
      "  Create dedicated datasets for model training.\n",
      "\n",
      "Options:\n",
      "  -s, --samples [QCD|extra_QCD|W_qq|W_qqq]\n",
      "                                  List of data samples to process (separated\n",
      "                                  by commas). Available samples are: ['QCD',\n",
      "                                  'extra_QCD', 'W_qq', 'W_qqq']. By default,\n",
      "                                  all samples will be processed.\n",
      "  -d, --datadir TEXT              Base directory for storing datasets.\n",
      "                                  [default: datasets]\n",
      "  --cache / --no-cache            Whether to cache existing results.\n",
      "                                  [default: cache]\n",
      "  --parallel INTEGER              Parallelize job across the N workers\n",
      "                                  Case  0: Jobs are run sequentially (for debugging)\n",
      "                                  Case -1: Jobs are run across N_CPU workers.  [default: -1]\n",
      "  -v, --verbosity [DEBUG|INFO|WARNING|ERROR]\n",
      "                                  Verbosity level.  [default: INFO]\n",
      "  --help                          Show this message and exit.\n"
     ]
    }
   ],
   "source": [
    "!paws create_dedicated_datasets --help"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd584426-c670-423c-b9d2-4a9d469981b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# note that each command will consume all the CPUs by default, please limit the CPU usage by setting --parallel N_CPU\n",
    "!paws create_dedicated_datasets --s QCD -d datasets\n",
    "!paws create_dedicated_datasets --s extra_QCD -d datasets\n",
    "!paws create_dedicated_datasets --s W_qq -d datasets\n",
    "!paws create_dedicated_datasets --s W_qqq -d datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a053b2b0-4742-4d9f-bfd0-c450e105b908",
   "metadata": {},
   "source": [
    "To use the API:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1dad0683-f5c1-43ab-846b-bab3fa67b098",
   "metadata": {},
   "outputs": [],
   "source": [
    "from paws.data_preparation import create_high_level_dedicated_datasets\n",
    "\n",
    "datadir = \"datasets\"\n",
    "samples = ['QCD', 'extra_QCD', 'W_qq', 'W_qqq']\n",
    "\n",
    "for sample in samples:\n",
    "    create_high_level_dedicated_datasets(sample, datadir=datadir, parallel=-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f8e6ba1-f780-47ea-a8e6-e3738e3b077b",
   "metadata": {},
   "source": [
    "### Prepare dedicated datasets\n",
    "\n",
    "Dedicated datasets are the combined dedicated dataset for each decay mode scenario. The datasets are pre-shuffled with a given seed.\n",
    "\n",
    "Notice that for each sample and mass point, the dataset is split into N = 100 shards. We prepare the training, validation and test dataset splits based on the set of shard indices that go into each split to ensure their orthogonality and reproducibility. In the weakly supervised training, the validation and test datasets will be turned into the training datasets instead (and the training split into the validation and test) to avoid bias from the supervised training.\n",
    "\n",
    "These datasets are used for training of the parameterised supervised models which we train over a parametric set of signal mass points.\n",
    "\n",
    "To use the CLI:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "15f5cc6d-8bf3-402d-a898-639ce73ba7a9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Usage: paws create_param_datasets [OPTIONS]\n",
      "\n",
      "  Create parameterised datasets for model training.\n",
      "\n",
      "Options:\n",
      "  -s, --samples [QCD|extra_QCD|W_qq|W_qqq]\n",
      "                                  List of data samples to include (separated\n",
      "                                  by commas). Available samples are: ['QCD',\n",
      "                                  'extra_QCD', 'W_qq', 'W_qqq']. By default,\n",
      "                                  all samples will be included. Note that for\n",
      "                                  two-prong / three-prong training, only the\n",
      "                                  two-prong / three-prong signals should be\n",
      "                                  included.\n",
      "  -d, --datadir TEXT              Base directory for storing datasets.\n",
      "                                  [default: datasets]\n",
      "  --shards TEXT                   Process datasets with the specific shard\n",
      "                                  indices (separated by commas). By default,\n",
      "                                  all shards will be processed. Use\n",
      "                                  \"start_index:end_index\" to indicate a slice\n",
      "                                  of shard indices.\n",
      "  --seed INTEGER                  Random seed used in dataset shuffling.\n",
      "                                  [default: 2023]\n",
      "  --cache / --no-cache            Whether to cache existing results.\n",
      "                                  [default: cache]\n",
      "  --parallel INTEGER              Parallelize job across the N workers\n",
      "                                  Case  0: Jobs are run sequentially (for debugging)\n",
      "                                  Case -1: Jobs are run across N_CPU workers.  [default: -1]\n",
      "  -v, --verbosity [DEBUG|INFO|WARNING|ERROR]\n",
      "                                  Verbosity level.  [default: INFO]\n",
      "  --help                          Show this message and exit.\n"
     ]
    }
   ],
   "source": [
    "!paws create_param_datasets --help"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53936f0a-2227-421c-baee-4551416c4699",
   "metadata": {},
   "outputs": [],
   "source": [
    "# two-prong dataset\n",
    "!paws create_param_datasets -s QCD,extra_QCD,W_qq -d datasets\n",
    "# three-prong dataset\n",
    "!paws create_param_datasets -s QCD,extra_QCD,W_qqq -d datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51ca93d3-55cd-4a82-9ab7-c9f70ea51c8d",
   "metadata": {},
   "source": [
    "To use the API:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41c77ae7-6239-4892-a71d-360fb19ea9a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "    \n",
    "from paws.data_preparation import create_parameterised_datasets\n",
    "from paws.settings import NUM_SHARDS\n",
    "\n",
    "datadir = \"datasets\n",
    "\n",
    "two_prong_samples = [\"W_qq\", \"QCD\", \"extra_QCD\"]\n",
    "three_prong_samples = [\"W_qqq\", \"QCD\", \"extra_QCD\"]\n",
    "\n",
    "# you can process a subset of shards in each job to potentially speed up the procedure\n",
    "# process all shards in one go for now\n",
    "shard_indices = np.arange(NUM_SHARDS)\n",
    "\n",
    "# you might run into memory issue if you process too many shards at the same time, try to reduce the number of parallel workers accordingly\n",
    "\n",
    "# two-prong dataset\n",
    "create_parameterised_datasets(shard_indices, sample=two_prong_samples, datadir=datadir, parallel=16)\n",
    "# three-prong dataset\n",
    "create_parameterised_datasets(shard_indices, sample=three_prong_samples, datadir=datadir, parallel=16)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}