{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "daXuA3h2b53x"
   },
   "source": [
    "# Hands-on 05: Time series data and RNNs: Identifying cosmic rays in radio signals\n",
    "Credit: Adapted from <a href=\"http://deeplearningphysics.org\">deeplearningphysics.org</a>\n",
    "\n",
    "Large arrays of radio antennas can be used to measure cosmic rays by recording the electromagnetic radiation generated in the atmosphere.\n",
    "These radio signals are strongly contaminated by galactic noise as well as signals from human origin. \n",
    "Because these signals appear to be similar to the background, the discovery of cosmic-ray events can be challenging.\n",
    "\n",
    "## Identification of signals\n",
    "In this exercise, we design an RNN to classify if the recorded radio signals contain a cosmic-ray event or only noise.\n",
    "\n",
    "The signal-to-noise ratio (SNR) of a measured trace $S(t)$ is defined as follows:\n",
    "\n",
    "$$\\mathrm{SNR}=\\frac{S^{\\mathrm{signal}}(t)_\\mathrm{max}}{\\mathrm{RMS}[S(t)]},$$\n",
    "\n",
    "where $S^{\\mathrm{signal}}(t)_\\mathrm{max}$ denotes the maximum amplitude of the (true) signal.\n",
    "\n",
    "Typical cosmic-ray observatories enable a precise reconstruction at an SNR of roughly 3.\n",
    "\n",
    "We choose a challenging setup in this task and try to identify cosmic-ray events in signal traces with an SNR of 2.  \n",
    "Training RNNs can be computationally demanding, thus, we recommend to use a GPU for this task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "yMkxS2Zrb533",
    "outputId": "0ca1dc6b-3095-46e5-d809-b95bbfaf7525"
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "b7A1BSCHb539"
   },
   "source": [
    "### Load and prepare dataset\n",
    "In this task, we use a simulation of cosmic-ray-induced air showers that are measured by radio antennas.  \n",
    "For more information, see <a href=\"https://arxiv.org/abs/1901.04079\">arXiv:1901.04079</a>.\n",
    "The task is to design an RNN which is able to identify if the measured signal traces (shortened to 500 time steps) contains a signal or not."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ss_iQSeWb54A"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import gdown\n",
    "\n",
    "url = \"https://drive.google.com/u/0/uc?export=download&confirm=HgGH&id=1R-qfxO1jVh88TC9Gnm9JGMomSRg0Zpkx\"\n",
    "output = \"data/radio_data.npz\"\n",
    "\n",
    "if not os.path.exists(output):\n",
    "    gdown.download(url, output, quiet=True)\n",
    "\n",
    "# there are 50k traces total, but let's just look at the first 5k\n",
    "f = np.load(output)\n",
    "n_train = 4000\n",
    "n_test = 1000\n",
    "\n",
    "x_train, x_test = f[\"traces\"][:n_train], f[\"traces\"][n_train : n_train + n_test]  # measured traces (signal + colored noise)\n",
    "signal_train, signal_test = (\n",
    "    f[\"signals\"][:n_train],\n",
    "    f[\"signals\"][n_train : n_train + n_test],\n",
    ")  # signal part (only available for cosmic-ray events)\n",
    "\n",
    "# define training label (1=cosmic event, 0=noise)\n",
    "y_train = (signal_train.std(axis=-1) != 0).astype(float)\n",
    "y_test = (signal_test.std(axis=-1) != 0).astype(float)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hpgUu--lb54B"
   },
   "source": [
    "## Plot example signal traces\n",
    "Left: signal trace containing a cosmic-ray event. The underlying cosmic-ray signal is shown in red, the backgrounds + signal is shown in blue.\n",
    "Right: background noise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 297
    },
    "id": "BiswHp3_b54D",
    "outputId": "9a0bc7a6-3c22-4b9c-b5d3-0728ee629593"
   },
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "\n",
    "fs = 180e6  # Sampling frequency of antenna setup 180 MHz\n",
    "t = np.arange(500) / fs * 1e6\n",
    "\n",
    "plt.figure(1, (12, 4))\n",
    "\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.plot(t, np.real(x_train[y_train.astype(bool)][0]), linewidth=1, color=\"b\", label=\"Measured trace\")\n",
    "plt.plot(t, np.real(signal_train[y_train.astype(bool)][0]), linewidth=1, color=\"r\", label=\"CR signal\")\n",
    "plt.ylabel(\"Amplitude [mV]\")\n",
    "plt.xlabel(\"Time [$\\mu$s]\")\n",
    "plt.legend()\n",
    "plt.title(\"Cosmic-ray event\")\n",
    "plt.grid(True)\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.plot(t, np.real(x_train[~y_train.astype(bool)][0]), linewidth=1, color=\"b\", label=\"Measured trace\")\n",
    "plt.ylabel(\"Amplitude [mV]\")\n",
    "plt.xlabel(\"Time [$\\mu$s]\")\n",
    "plt.legend()\n",
    "plt.title(\"Noise event\")\n",
    "plt.grid(True)\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oRPX-EYOb54L"
   },
   "outputs": [],
   "source": [
    "sigma = x_train.std()\n",
    "x_train /= sigma\n",
    "x_test /= sigma"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lOanYtv_b54G"
   },
   "source": [
    "### Define RNN model\n",
    "In the following, design a cosmic-ray model to identify cosmic-ray events using an RNN-based classifier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "kBzMjinOb54G",
    "outputId": "1e892def-6270-4178-b372-65210f5d8f28"
   },
   "outputs": [],
   "source": [
    "from tensorflow.keras import layers\n",
    "\n",
    "model = keras.models.Sequential(name=\"sequential_1\")\n",
    "model.add(layers.LSTM(32, return_sequences=True, input_shape=(500, 1), name=\"lstm_1\"))\n",
    "model.add(layers.LSTM(64, return_sequences=True, name=\"lstm_2\"))\n",
    "model.add(layers.LSTM(10, return_sequences=True, name=\"lstm_3\"))\n",
    "model.add(layers.Flatten(name=\"flatten_1\"))\n",
    "model.add(layers.Dropout(0.3, name=\"dropout_1\"))\n",
    "model.add(layers.Dense(1, activation=\"sigmoid\", name=\"dense_1\"))\n",
    "\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Bw8tkc1Xb54I"
   },
   "source": [
    "#### RNN training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "fjncDgOcb54N",
    "outputId": "1fd1ec77-5711-401e-9c40-336c3a588f30"
   },
   "outputs": [],
   "source": [
    "lr_schedule = keras.optimizers.schedules.ExponentialDecay(initial_learning_rate=1e-3, decay_steps=100, decay_rate=0.8)\n",
    "\n",
    "model.compile(loss=\"binary_crossentropy\", optimizer=keras.optimizers.Adam(learning_rate=lr_schedule), metrics=[\"accuracy\"])\n",
    "\n",
    "\n",
    "results = model.fit(\n",
    "    x_train[..., np.newaxis],\n",
    "    y_train,\n",
    "    batch_size=256,\n",
    "    epochs=20,\n",
    "    verbose=0,\n",
    "    validation_split=0.1,\n",
    "    callbacks=[\n",
    "        keras.callbacks.EarlyStopping(patience=15, verbose=0, restore_best_weights=True),\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "xzutH-Qkb54P",
    "outputId": "dbf0745e-dab6-4073-9466-c42bf6e10c0e"
   },
   "outputs": [],
   "source": [
    "y_pred = model.predict(x_test[..., np.newaxis], verbose=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "A16SCLCBb54R"
   },
   "source": [
    "### Plot loss and accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from plotting import plot_model_history\n",
    "\n",
    "plot_model_history(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from plotting import make_roc\n",
    "from sklearn.metrics import accuracy_score\n",
    "\n",
    "make_roc(y_test, y_pred, [\"CR signal\"])\n",
    "accuracy_score(y_test, y_pred > 0.5)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercises\n",
    "1. Replace the LSTM layers with bidirectional LSTM layers: `layers.Bidirectional(layers.LSTM(...))` (See https://keras.io/api/layers/recurrent_layers/bidirectional/). Note for the first layer, the `input_shape` argument needs to be an input to the outer `layers.Bidirectional` function instead of the inner `layers.LSTM` function. How does the number of parameters of the model change? How does the performance change?\n",
    "2. Replace the LSTM layers with GRU layers. How does the number of parameters of the model change? How does the performance change?\n",
    "3. Replace the LSTM layers with 1D convolutional layers `layers.Conv1D(...)` with the following parameters. How does the number of parameters of the model change? How does the performance change?\n",
    "   - `filters=32, kernel_size=5, strides=1, padding=\"same\", activation=\"relu\"` \n",
    "   - `filters=64, kernel_size=5, strides=2, padding=\"same\", activation=\"relu\"`\n",
    "   - `filters=10, kernel_size=5, strides=2, padding=\"same\", activation-\"relu\"`\n",
    "4. How can you modify the architecture and data to perform signal trace extraction?\n",
    "5. (BONUS) Implement signal trace extraction."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "name": "Exercise_9_2_solution.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "phys139",
   "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.9.15"
  },
  "vscode": {
   "interpreter": {
    "hash": "d0ea348b636367bcdf67fd2d6d24251712b38670f61fdee14f28eb58fe74f081"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}