{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Introduction\n",
    "ABCGAN (atmospheric background characterization 'GAN') is a generative adversarial network that reproduces realistic physical profiles of the upper atmosphere and ionosphere. Unlike traditional empirical models, ABCGAN can produce not only the average profile under particular conditions, but also the expected distribution of parameters. This allows the GAN to estimate the probability that a given profile is part of the historical distribution, enabling it to determine if an event is highly unusual (outside the range of \"normal\" atmospheric variation).\n",
    "\n",
    "The model was trained on historical measurements collected from the Poker Flat Incoherent Scattering Radar (PFISR) from 2007-2020. PFISR measures profiles of the scalar plasma parameters electron density (Ne), electron temperature (Te), and ion temperature (Ti) in the ionosphere. For more information about the training data, visit amisr.com.\n",
    "\n",
    "ABCGAN is valid for predicting average ionospheric parameters over Poker Flat, AK.  The results should also be reasonable for most locations in the auroral zone, but caution should be exercised using the GAN results in other latitude regimes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Tutorial Outline\n",
    "\n",
    "This demo provides a demonstration of the abcgan project. This includes using our BV/HFP GAN models to generate atmospheric\n",
    "background profiles/gravity waves consistent with the historically observed distribution and detecting statistical anomalies\n",
    "the historically observed atmospheric data. The entire outline of the tutorial is listed below.\n",
    "\n",
    "- Import ABCGAN\n",
    "- Import Data\n",
    "    - External Drivers\n",
    "    * Background Variables\n",
    "    * High Frequency Perturbations\n",
    "* Generating Background Profiles & HFP Waves\n",
    "    * Generating BVs and HFPs\n",
    "    * Sampling the Generator\n",
    "* Exploring Model Results\n",
    "    * Evaluating Model Performance\n",
    "    * Comparing Real vs Generated\n",
    "* Driver Estimation\n",
    "* Anomaly Detection\n",
    "    * Background\n",
    "    * N-Dim Estimation\n",
    "    * BV Anomaly Example\n",
    "    * HFP Anomaly Example\n",
    "* Attention Maps\n",
    "    * BV Self-Attention Maps\n",
    "    * HFP Attention Map\n",
    "* TEC Waves\n",
    "    * LSTIDs GAN (Poker Flats)\n",
    "    * MSTIDs GAN (Mt Moses)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Import ABCGAN & Python Modules\n",
    "\n",
    "Several imports are used. In particular, we load `numpy`, which provides the formats for data input output. We also use the numpy library for calculation of statistics and simulating fake data. We load `abcgan` which contains modules for generating atmospheric data, evaluating model performing, estimating anomaly scores, and analysing model behaviors. We also load `matplotlib.pyplot`, `seaborn`, and `pandas` for display."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np  # numerics\n",
    "import abcgan  # abcgan package\n",
    "import matplotlib.pyplot as plt  # plotting\n",
    "import matplotlib\n",
    "import seaborn as sns  # plotting\n",
    "import pandas as pd  # plotting\n",
    "import datetime\n",
    "import matplotlib.image as mpimg\n",
    "import abcgan.constants as const\n",
    "from abcgan.interface import estimate_drivers\n",
    "from abcgan.mask import prev_driver_mask\n",
    "from abcgan.interface import generate_bvs, generate_hfps, generate_multi_bv, generate_multi_hfp\n",
    "from abcgan.anomaly import anomaly_score_bv, anomaly_score_hfp, anomaly_score_wtec\n",
    "from abcgan.interface import load_h5_data, load_wtec_h5, generate_wtec, generate_multi_wtec\n",
    "from abcgan.attention import collect_bv_attn_map, collect_hfp_attn_map\n",
    "from abcgan.interface import hellinger_scores_bv, hellinger_scores_hfp, hellinger_scores_wtec"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Import Data\n",
    "\n",
    "We can load all the data using the `load_h5_data`, which takes care of reading data from the h5 file and filtering out any invalid samples. You can pass in the number of samples to load along with the file name and the function returns numpy arrays of driver parameters, background variables, altitude profile mask for the background variables, HFP waves, HFP wave mask, and the unix time stamp of each sample.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Drivers Shape : (1000, 29)\n",
      "Background Variables Shape : (1000, 30, 12)\n",
      "Alt Mask Shape : (1000, 30)\n",
      "Background Variables Shape : (1000, 1, 8)\n",
      "Alt Mask Shape : (1000, 1)\n",
      "Unix Time Stamp Shape : (1000,)\n"
     ]
    }
   ],
   "source": [
    "fname = 'tutorial_isr.h5'\n",
    "drivers, bvs, alt_mask, hfps, wave_mask, unix_time = load_h5_data(fname, n_samples=1000, random_start=True, load_hfp=True)\n",
    "print(f\"Drivers Shape : {drivers.shape}\")\n",
    "print(f\"Background Variables Shape : {bvs.shape}\")\n",
    "print(f\"Alt Mask Shape : {alt_mask.shape}\")\n",
    "print(f\"Background Variables Shape : {hfps.shape}\")\n",
    "print(f\"Alt Mask Shape : {wave_mask.shape}\")\n",
    "print(f\"Unix Time Stamp Shape : {unix_time.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## External Drivers\n",
    "\n",
    "Our GAN models take the external drivers, which describe the current environmental conditions, and outputs generated background profiles and/or hfp waves. We currently are using 29 external drivers, which are listed and described below\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of Drivers: 29\n",
      "Drivers Names: ['Ap', 'F10.7', 'F10.7avg', 'MEI', 'MLT', 'RMM1', 'RMM2', 'SLT', 'SZA', 'ShadHeight', 'T10', 'T100', 'T2', 'T30', 'T5', 'T70', 'TCI', 'U10', 'U100', 'U2', 'U30', 'U5', 'U70', 'ap', 'dst', 'moon_phase', 'moon_x', 'moon_y', 'moon_z']\n"
     ]
    }
   ],
   "source": [
    "print(f\"Number of Drivers: {drivers.shape[-1]}\")\n",
    "print(f\"Drivers Names: {const.driver_names}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "\n",
    "The list of drivers keys printed above corresponds to the following physical parameters:\n",
    "\n",
    "| Key               | Full Name                      | Description                                                                                               |\n",
    "|-------------------|--------------------------------|-----------------------------------------------------------------------------------------------------------|\n",
    "| SLT               | Solar Local Time               | Local time as determined by the sun's position                                                            |\n",
    "| MLT               | Magnetic Local Time            | Local time calculated used magnetic coordinate system                                                     |\n",
    "| SZA               | Solar Zenith Angle             | Angle between the sun and local zenith                                                                    |\n",
    "| ShadHeight        | Shadow Height                  | Height of the Earth's shadow in the atmosphere                                                            |\n",
    "| moon_x            | Lunar Position - x             | X component of Lunar position in ECEF coordinates                                                         |\n",
    "| moon_y            | Lunar Position - y             | Y component of Lunar position in ECEF coordinates                                                         |\n",
    "| moon_z            | Lunar Position - z             | Z component of Lunar position in ECEF coordinates                                                         |\n",
    "| moon_phase        | Lunar Phase                    | Lunar phase in degrees                                                                                    |\n",
    "| ap                | ap Geomagnetic Index           | Three hour equivilent planetary amplitude                                                                 |\n",
    "| Ap                | Daily Ap Index                 | Daily equivalent planetary aplitude                                                                       |\n",
    "| F10.7             | F10.7 Solar Radio Flux         | Solar rado flux at 10.7 cm (2800 MHz) in s.f.u (10^-22 W m^-2 Hz^-1)                                      |\n",
    "| F10.7a            | Average F10.7 Solar Radio Flux | Average solar rado flux at 10.7 cm (2800 MHz) in s.f.u (10^-22 W m^-2 Hz^-1)                              |\n",
    "| TCI               | Thermosphere Climate Index     | 60-day running average of global cooling power radiated from the thermosphere                             |\n",
    "| MEI               | Multivariate ESNO Index        | Measure of the El Nino/Southern Oscillation (ESNO)/sea surface temperature                                |\n",
    "| RMM1              | First MJO Index                | First component of the Madden-Julian Osscillation Indices                                                 |\n",
    "| RMM2              | Second MJO Index               | Second component of the Madden-Julian Osscillation Indices                                                |\n",
    "| U2/5/10/30/70/100 | Mean zonal winds               | mean zonal winds (eastward) at at altitudes with pressures 2,5,10,30,70, and 100 hPa                      |\n",
    "| T2/5/10/30/70/100 | Mean Stratospheric Temps       | Stratospheric northern hemisphere mean (60 - 90 lat) at altitudes with pressures 2,5,10,30,70, and 100 hPa |\n",
    "| dst               | Distrurbance storm time index  | index of magnetic activity that measures the intensity of the globally symmetrical equatorial electrojet  |\n",
    "\n",
    "\n",
    "\n",
    "Each driver is stored as a separate vector in `tutorial_isr.h5`, where each entry in the vector is a separate sample. So the full set of drivers for the first sample consists of the first element in each vector.The `load_h5_data` function collates all the drivers samples into a single data structure which has size nSamples x nDrivers, in the order specified by `abcgan.driver_names`.\n",
    "\n",
    "See the tutorial on generating and saving drivers in this directory for how to create a driver file."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## High Frequency Perturbations (HFPs)\n",
    "\n",
    "The HFPs data describes atmospheric gravity waves found in the background variable data. Waves are not always present, which is indicated by nans in hfp dataset. The `wave_mask` also indicates whether a wave is present or not at each sample and we can see that only around half the samples have waves. Currently our dataset only allows up to one wave to be present at each sample thus the 2nd dim, which indicates the list of waves, will always be 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "HFP Shape: (1000, 1, 8)\n",
      "1000 Samples\n",
      "8 HFP variables\n",
      "['GW_tau', 'GW_lamh', 'GW_lamv', 'GW_ch', 'GW_cv', 'GW_phi', 'GW_hl', 'GW_hu']\n",
      "\n",
      "Percentage of Samples with Waves: 0.41\n"
     ]
    }
   ],
   "source": [
    "print(f'HFP Shape: {hfps.shape}')\n",
    "print(f'{hfps.shape[0]} Samples')\n",
    "print(f'{hfps.shape[2]} HFP variables')\n",
    "print(const.hfp_names)\n",
    "print(f'\\nPercentage of Samples with Waves: {wave_mask.sum()/len(wave_mask):.2f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "\n",
    "The list of hfps variables printed above corresponds to the following physical parameters:\n",
    "\n",
    "| Key        | Full Name                                               | Description                                                                  |\n",
    "| ---------- | --------------------------------------------------------|----------------------------------------------------------------------------- |\n",
    "| GW_tau     | GW period in minutes                                    | Period using PFISR 5-min integrated nenotr F-region Ne perturbations                              |\n",
    "| GW_lamh    | GW horizontal wavelength in km                          | Horizontal wavelength using PFISR 5-min integrated nenotr F-region Ne perturbations|\n",
    "| GW_lamv    | GW vertical wavelength in km                            | Vertical wavelength using PFISR 5-min integrated nenotr F-region Ne perturbations|\n",
    "| GW_ch      | GW horizontal phase speed in m/s                        | Horizontal phase speed using PFISR 5-min integrated nenotr F-region Ne perturbations|\n",
    "| GW_cv      | GW vertical phase speed in m/s                          | Vertical phase speed using PFISR 5-min integrated nenotr F-region Ne perturbations|\n",
    "| GW_phi     | GW horizontal propagation direction in azimuth degrees  | Horizontal propagation direction using PFISR 5-min integrated nenotr F-region Ne perturbations|\n",
    "| GW_hl      | GW lower altitude boundary in km                        | Lower altitude where the wave was detected|\n",
    "| GW_hu      | GW upper altitude boundary in km                        | Higher altitude where the wave was detected|                 |"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Generating Background Profiles & HFPs\n",
    "\n",
    "Using the BV GAN we can generate background profiles, one for each sample, given an input set of external drivers. The BV GAN generates 12 background variables altitude bin by altitude bin (bottom to top) using measurements generated from subsequent layers to help predict measurements at higher altitudes. To get the generated background profiles we can call the `generate_bvs` function. Below is an example of generating 5 background profiles using the first 5 external driver samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generate BV Profile: 100%|██████████| 1/1 [00:00<00:00,  5.45it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated BV Shape: (5, 30, 12)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "G_bvs = generate_bvs(drivers=drivers[:5])\n",
    "print(f'Generated BV Shape: {G_bvs.shape}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "When generating background profiles we can specify the max altitude bin using the `n_alt` argument. We also have the options to generate upper sections of a background profile conditioned on real measurements at lower sections in the case that you only need to generate a sub section of higher background measurements. For example if I already had the first 10 altitude bins of background measurements and I wanted to generate background profiles that go up to the 20th altitude bin conditioned on the measurements I already have I could call the generate function with the following arguments:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generate BV Profile: 100%|██████████| 1/1 [00:00<00:00, 13.55it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated BV Shape: (5, 20, 12)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "G_bvs = generate_bvs(drivers[:5], bvs[:5, :10, :], n_alt=20)\n",
    "print(f'Generated BV Shape: {G_bvs.shape}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We also have a HFP GAN that generates 8 HFP wave parameters given the same set of external driver parameters. To generate the HFP waves conditioned on the set of real background profiles we can call the same `generate_hfps` and pass in the real background measurements. Since around half of the HFP waves are not actually present at each sample the HFP GAN will also output a probability `G_b`, that the generated wave is present. We can perform a hard decision on whether the generated wave is present by checking to see if the probability `G_b` is less than sample randomly selected from a uniform distribution between 0 and 1. Note that you have to pass a set of bv measurements and if you don't have them you first need to generate them using the `generate_bvs` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generate HFP Batches: 100%|██████████| 1/1 [00:00<00:00, 53.91it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated HFP Shape: (5, 1, 8)\n",
      "Gb Shape: (5,)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "G_hfps, G_b = generate_hfps(drivers=drivers[:5], bv_measurements=bvs[:5])\n",
    "Gb_mask = G_b < np.random.uniform(size=len(G_b))\n",
    "print(f'Generated HFP Shape: {G_hfps.shape}')\n",
    "print(f'Gb Shape: {G_b.shape}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We also have the option to condition HFP GAN using background profile measurements generated from the BV GAN. You would pass in the generated BV measurement as conditioning when you don't have access to the actual BV measurements. This time we'll also generate background profiles and HFP waves each of the 500 external driver samples. Note that the `generate_bvs` and `generate_hfps` function call may take a minute or longer to complete when generating 500 BV/HFP samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generate BV Profile: 100%|██████████| 10/10 [00:27<00:00,  2.80s/it]\n",
      "Generate HFP Batches: 100%|██████████| 10/10 [00:02<00:00,  4.83it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated HFP Shape: (1000, 1, 8)\n",
      "Generated BV Shape: (1000, 30, 12)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "G_bvs = generate_bvs(drivers=drivers)\n",
    "G_hfps, G_b = generate_hfps(drivers=drivers, bv_measurements=G_bvs)\n",
    "Gb_mask = G_b < np.random.uniform(size=len(G_b))\n",
    "print(f'Generated HFP Shape: {G_hfps.shape}')\n",
    "print(f'Generated BV Shape: {G_bvs.shape}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Sampling\n",
    "\n",
    "To generate multiple samples for the same driving parameters (usually in order to compute statistics) pass a repeated version of the drivers. An example of this can be shown below were we are generating 100 background variable altitude profiles samples using the first set drivers. Again for the HFP waves generated we can make a hard decision on whether the wave is present or not by comparing each generated probability `Gb_sampled` to another randomly selected probability taken from a uniform distribution. If `Gb_sampled` is less than the randomly selected sample then we assume the wave generated at that sample is not present.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generate BV Profile: 100%|██████████| 1/1 [00:03<00:00,  3.02s/it]\n",
      "Generate HFP Batches: 100%|██████████| 1/1 [00:00<00:00,  4.65it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sampled Generated BV Shape (100, 30, 12)\n",
      "Sampled Generated HFP Shape (100, 1, 8)\n",
      "Sampled Generated B Shape (100,)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "G_bvs_sampled = generate_bvs(drivers=drivers[:1, :].repeat(100, 0))\n",
    "G_hfps_sampled, Gb_sampled = generate_hfps(drivers=drivers[:1, :].repeat(100, 0), bv_measurements=G_bvs[:100])\n",
    "Gb_mask_sampled = Gb_sampled < np.random.uniform(size=len(Gb_sampled))\n",
    "print(f'Sampled Generated BV Shape {G_bvs_sampled.shape}')\n",
    "print(f'Sampled Generated HFP Shape {G_hfps_sampled.shape}')\n",
    "print(f'Sampled Generated B Shape {Gb_sampled.shape}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Compare Generated BV Profiles\n",
    "As a practical example, look at several real profiles of electron density from the background values with similar driving conditions and compare them to the set of profiles generated by the GAN for these drivers.  First identify points from the sample data set that were close to noon MLT with Ap<10 and F10.7<80.  Then find the average driving conditions for these points and use them to generate sampled profiles with the GAN."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generate BV Profile: 100%|██████████| 1/1 [00:03<00:00,  3.21s/it]\n"
     ]
    },
    {
     "data": {
      "text/plain": "<Figure size 1500x1000 with 2 Axes>",
      "image/png": "[encoded data, 252440 characters]\n"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "idx = np.argwhere((np.abs(drivers[:, const.dr_cols['MLT']]-12.)<1.) & (drivers[:,const.dr_cols['Ap']]<10.) & (drivers[:,const.dr_cols['F10.7']]<80.)).flatten()\n",
    "avg_drivers = np.mean(drivers[idx,:], axis=0)[None,:]\n",
    "\n",
    "gen_profiles = generate_bvs(avg_drivers.repeat(100, 0))\n",
    "mean_profile = np.mean(gen_profiles, axis=0)\n",
    "std_profile = np.std(gen_profiles, axis=0)\n",
    "\n",
    "plt.style.use('default')\n",
    "fig, ax = plt.subplots(1, 2, figsize=(15, 10))\n",
    "pal = sns.color_palette('hls', n_colors=2)\n",
    "plt_idxs = np.random.choice(np.arange(len(gen_profiles)), replace=False, size=5)\n",
    "\n",
    "ax[0].plot(bvs[idx[plt_idxs],:,0].T, const.altgrid[:30], c=pal[1], zorder=3)\n",
    "ax[0].plot(bvs[idx[plt_idxs][-2:-1],:,0].T, const.altgrid[:30], c=pal[1], label='Real Profiles', zorder=3)\n",
    "ax[1].plot(bvs[idx[plt_idxs],:,0].T, const.altgrid[:30], c=pal[1], zorder=3)\n",
    "ax[1].plot(bvs[idx[plt_idxs][-2:-1],:,0].T, const.altgrid[:30], c=pal[1], label='Real Profiles', zorder=3)\n",
    "# ax[0].set_title('Real Ne Profiles from PFISR')\n",
    "ax[0].plot(gen_profiles[plt_idxs,:,0].T, const.altgrid[:30], c=pal[0],zorder=3)\n",
    "ax[0].plot(gen_profiles[plt_idxs,:,0][-2:-1].T, const.altgrid[:30], c=pal[0], label='Generated Profiles',zorder=3)\n",
    "# ax[0].set_title('Real vs Generated (Conditional)', fontsize=25)\n",
    "\n",
    "for i in range(2):\n",
    "    ax[i].errorbar(mean_profile[:,0], const.altgrid[:30], xerr=std_profile[:,0], color='black', label='Generator Error Bar', zorder=5)\n",
    "    ax[i].set_xlim([0.,5.e11])\n",
    "    ax[i].set_ylim([80.,600.])\n",
    "    ax[i].set_xlabel(r'Electron Density (m$^{-3}$)', fontsize=20)\n",
    "    ax[i].set_ylabel('Altitude (km)', fontsize=20)\n",
    "    ax[i].legend(fontsize=20)\n",
    "    ax[i].tick_params(labelsize=20)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The black profile in both plots is the mean and standard deviation of the sampled profiles from the GAN. They agree with the measured profiles, and contain information about the expected spread of the electron density at each altitude. The output of the GAN is **statistical**, and each individual generated profile is a sampled profile from the distribution but doesnot necessarily represent the \"average\" or \"typical\" conditions.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Exploring the results\n",
    "\n",
    "The above utilities are still being tested and developed. We expect to soon have the following included in the tutorial:\n",
    "- Demonstrate how to evaluate BV/HFP GAN using the Hellinger Distance Metric\n",
    "- Demonstration showing how to compare the histograms of generated data with historically observed data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Evaluating Model Performance\n",
    "\n",
    "One way to evaluate the performance of our GAN model is to compare how similar the conditional background profile distributions of the real measurements are to the generated ones. To measure the similarity between two distributions we will use the Hellinger distance metric, which as the following formula:\n",
    "\n",
    "\n",
    "$$H(p, q) = {\\frac{1}{\\sqrt2}}\\sqrt{\\sum \\limits _{i=1}^ {k} (\\sqrt{p_i} - \\sqrt{q_i})^2}$$\n",
    "\n",
    "The hellinger distance metric can score anywhere between 0 (very similar) to 1 (very different). We can perform this calculation by calling the `hellinger_distance_bv` function. This function will compute the Hellinger distance score for each background variable at each altitude using the conditional distributions derived from the inputted set of real and generated background profile samples. Along with inputting the samples we'll pass in the `alt_mask` so we don't end up using any invalid/nan samples in the calculation.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hellinger Distance Shape: (30, 12)\n",
      "Avg Hel Dist: 0.059\n"
     ]
    }
   ],
   "source": [
    "hel_dists = hellinger_scores_bv(bvs, G_bvs, mask=alt_mask)\n",
    "print(f'Hellinger Distance Shape: {hel_dists.shape}')\n",
    "print(f'Avg Hel Dist: {np.mean(hel_dists):.3f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The same process can be done with the HFP data. However, we can only compare samples that have present HFP waves before calling the scoring function we filter out any of the samples don't have a present wave using the `wave_mask`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "HFP Hellinger Distance Shape: (1, 8)\n",
      "Avg Hel Dist: 0.101\n"
     ]
    }
   ],
   "source": [
    "hfp_hel_dists = hellinger_scores_hfp(hfps[wave_mask[:, 0]], G_hfps[wave_mask[:, 0]])\n",
    "print(f'HFP Hellinger Distance Shape: {hfp_hel_dists.shape}')\n",
    "print(f'Avg Hel Dist: {np.mean(hfp_hel_dists):.3f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Comparing Real vs Generated\n",
    "\n",
    "We can also visualize the performance of the BV GAN by plotting and comparing the distributions real and generated measurements for a particular altitude bin. For this example we'll look at the real and generated BV distributions for the 20th altitude bin."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": "<Figure size 2500x1000 with 12 Axes>",
      "image/png": "[encoded data, 205236 characters]\n"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "alt_bin = 20\n",
    "real = bvs[alt_mask[:, alt_bin], alt_bin, :]\n",
    "generated = G_bvs[alt_mask[:, alt_bin], alt_bin, :]\n",
    "\n",
    "plt.style.use('default')\n",
    "fig, ax = plt.subplots(2, 6, figsize=(25, 10))\n",
    "ax = ax.flatten()\n",
    "for i in range(len(ax)):\n",
    "    x = np.hstack((real[... ,i].flatten(), generated[..., i].flatten()))\n",
    "    hue = ['Real'] * len(real[... ,i].flatten()) + ['Generated'] * len(generated[..., i].flatten())\n",
    "    data = pd.DataFrame({f'{const.bv_units[i]}': x, 'Type': hue})\n",
    "    sns.histplot(data=data, x=f'{const.bv_units[i]}', hue='Type', ax=ax[i],\n",
    "                 element='step', bins=25, legend=True, stat='density',\n",
    "                 palette='hls', alpha=0.2, common_norm=False,\n",
    "                 hue_order=['Generated', 'Real'])\n",
    "    ax[i].set_title(f'{abcgan.constants.bv_names[i]} (Hel Dist = {hel_dists[alt_bin, i]:.3f})')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The same process can be done with the HFP data. However, before getting the Hellinger scores we first have to filter out any waves that are not present in the real hfps using the `wave_mask` and also filter out the generated waves that were predicted to not be present using the `Gb_mask` calculated earlier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": "<Figure size 2000x1000 with 8 Axes>",
      "image/png": "[encoded data, 168904 characters]\n"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "real = hfps[wave_mask, ...]\n",
    "generated = G_hfps[Gb_mask, 0, ...]\n",
    "\n",
    "fig, ax = plt.subplots(2, 4, figsize=(20, 10))\n",
    "ax = ax.flatten()\n",
    "for i in range(len(ax)):\n",
    "    x = np.hstack((real[... ,i].flatten(), generated[..., i].flatten()))\n",
    "    hue = ['Real'] * len(real[... ,i].flatten()) + ['Generated'] * len(generated[..., i].flatten())\n",
    "    data = pd.DataFrame({f'{const.hfp_units[i]}': x, 'Type': hue})\n",
    "    sns.histplot(data=data, x=f'{const.hfp_units[i]}', hue='Type', ax=ax[i],\n",
    "                 element='step', bins=25, legend=True, stat='density',\n",
    "                 palette='hls', alpha=0.2, common_norm=False,\n",
    "                 hue_order=['Generated', 'Real'])\n",
    "    ax[i].set_title(f'{abcgan.constants.hfp_names[i]} (Hel Dist: {hfp_hel_dists[0, i]:.3f})')\n",
    "plt.tight_layout()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Using Previous Drivers To Estimate Current Drivers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "In the situation that you don't have the current drivers available for predicting the background profile you can use the driver estimation model to predict the current driver features based on previous driver samples. Specifically the driver estimation model has been trained to predict a driver sample that occur 2 hours from the time of the input sample. We can simulate the situation in which we only have drivers sample from two hours ago by creating a new driver dataset containing of driver samples that occurred two hours from the actual background variable samples. Since the time difference between each sample isn't always the same in our dataset we have to calculate a mapping vector between current samples and previous samples from two hour ago. The `prev_driver_mask` function generates this mapping vector along with a mask for selecting the driver samples that have a previous driver sample from two hours ago, which we will use when evaluating the driver estimator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "curr_drivers, _, _, unix_time = load_h5_data(fname, n_samples=None)\n",
    "prev_dr_map, dr_mask = prev_driver_mask(unix_time)\n",
    "prev_drivers = curr_drivers[prev_dr_map]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We can now estimate the current driver samples from the previous by calling the `estimate_drivers` function. As you can see the estimated drivers are very similar to actual driver samples. Note that when using in realtime you can simply call the estimate driver function with two hour old driver samples to predict the current driver samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Original Driver Shape: (5000, 29)\n",
      "Estimated Driver Shape: (3398, 29)\n"
     ]
    }
   ],
   "source": [
    "est_drivers = estimate_drivers(prev_drivers, model='dr_gan')\n",
    "print(f'Original Driver Shape: {curr_drivers.shape}')\n",
    "print(f'Estimated Driver Shape: {est_drivers.shape}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We can look at how the estimated driver distribution compares to the actual driver samples. As you can see the estimated drivers are very similar to actual driver samples. From here you can pass the estimated driver samples into the background variable generator to estimate the current background profile."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": "<Figure size 4000x2000 with 15 Axes>",
      "image/png": "[encoded data, 384056 characters]\n"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots(3, 5, figsize=(40, 20))\n",
    "ax = ax.flatten()\n",
    "for i in range(len(ax)):\n",
    "    x = np.hstack((curr_drivers[:, i], prev_drivers[:, i]))\n",
    "    hue = ['Estimated'] * len(curr_drivers[:, i]) + ['Actual'] * len(prev_drivers[:, i])\n",
    "    data = pd.DataFrame({'Value': x, 'Type': hue})\n",
    "    sns.histplot(data=data, x='Value', hue='Type', ax=ax[i],\n",
    "                 element='step', bins=100, legend=True, stat='density',\n",
    "                 palette='hls', alpha=0.2,\n",
    "                 hue_order=['Estimated', 'Actual'],\n",
    "                 common_norm=False)\n",
    "    ax[i].set_title(f'{const.driver_names[i]} Distribution')\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Detecting Anomalies with the GAN\n",
    "\n",
    "GANs offer a convenient way to realize an anomaly detector. First, we apply the GAN to regenerate the statistical distribution of our dataset. Then, new data is compared to the generated statistical distribution using a 'log-probability' distance measure.\n",
    "\n",
    "To compute the anomaly score, we apply a distance measure between new samples and a distribution of values generated by the GAN. To motivate our approach, first consider scalar data. For simplicity, assume our model is a random variable, so repeated calls to the GAN will generate a normal distribution centered around a mean value $\\mu$ with standard deviation $\\sigma$, then we may estimate the probability $p$ of a new signal $x$ comes from the distribution as.\n",
    "\n",
    "$$p(x) = \\frac{exp(-\\frac{1}{2}{\\delta}^2)}{\\sigma \\sqrt{2 \\pi}}.$$\n",
    "\n",
    "We define a dimensionless variable $\\delta = (x-\\mu)/\\sigma $. As the distance to the new point $\\delta$ increases, the exponential may underflow, or be too small to evaluate exactly. To fix underflow, we compute the log of the exponent to estimate an anomaly score, $A$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": "<matplotlib.image.AxesImage at 0x139d9e550>"
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": "<Figure size 1200x1200 with 1 Axes>",
      "image/png": "[encoded data, 72544 characters]\n"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "img = mpimg.imread('img/anom-normal.png')\n",
    "plt.figure(figsize = (12,12))\n",
    "plt.imshow(img)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "In plots $a)$ and $b)$, the distance to the distribution is small, so the exponential evaluates to value close to unity. However, as the new sample distance increases as in plots $ c)$ and $ d)$, the exponential evaluates to a very small number and is likely to underflow, meaning the value will be essentially zero."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Multi-Dimensional Estimation\n",
    "\n",
    "For our distance measure, we compute sum the exponential distances between the new sample point and 100 GAN-generated samples. The net effect is that we treat each new measurement as a random variable and the anomaly score tracks whether our new measurement is similar to any of the generated data points. If any of the GAN points overlaps with the new data, then it drives the anomaly score to a positive number. However, if none of the generated points lies near the sample data, then the anomaly score returns a negative number. To capture distances within the joint distribution, we compute the L2 norm between the N-dimensional sample and each generated data point, assuming a standard deviation $\\sigma_N$ computed from the geometric mean of all $N$ features $\\sigma_N = 2* \\sqrt[N]{(\\prod \\sigma_i)}/M$. Note the standard deviation was divided by the number of generated samples $M$ = 10,000. Finally, we compute the anomaly score $A$ for a given sample $x$ with numpy's built-in $logsumexp$ function, which is simply a numerically stable form of\n",
    "\n",
    "$$A(x) = log(\\sum_{i}e^{-1/2 {||\\hat{x}-\\hat{\\mu}||}^2/\\sigma_N^2}).$$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": "<matplotlib.image.AxesImage at 0x13980c190>"
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": "<Figure size 1200x1200 with 1 Axes>",
      "image/png": "[encoded data, 140848 characters]\n"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "img = mpimg.imread('img/marginal_vs_joint.png')\n",
    "plt.figure(figsize = (12,12))\n",
    "plt.imshow(img)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## BV Anomaly Example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Fist we have to Generate sampled fake data, which we can do using the `generate_mutli_bv` function. For this example we are going to generate 500 background profiles for a random driver sample to use in our anomaly estimation. Note that this may take minute to compute because we are generating 500 background profiles. You can lower `n_repeat` to speed up the process, but the fewer samples you generate the more inconsistent the anomaly detection results will be."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating BV Samples: 100%|██████████| 1/1 [00:13<00:00, 13.61s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "G_bvs Shape = (1, 500, 30, 12)\n",
      "N Driver Samples = 1\n",
      "N Background Profiles = 500\n",
      "N alt Bins = 30\n",
      "N Background Variables = 12\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "sample_idx = np.random.choice(np.arange(len(drivers))[wave_mask[:, 0]], size=[1])\n",
    "n_repeat = 500\n",
    "G_sampled_bvs = generate_multi_bv(drivers[sample_idx, ...], n_repeat=n_repeat)\n",
    "\n",
    "print(f\"G_bvs Shape = {G_sampled_bvs.shape}\")\n",
    "print(f\"N Driver Samples = {G_sampled_bvs.shape[0]}\")\n",
    "print(f\"N Background Profiles = {G_sampled_bvs.shape[1]}\")\n",
    "print(f\"N alt Bins = {G_sampled_bvs.shape[2]}\")\n",
    "print(f\"N Background Variables = {G_sampled_bvs.shape[3]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "With our 500 generate background profile samples we can calculate anomaly scores for our randomly selected driver under question using the `anomaly_score_bv` function. This function will produce anomaly scores for every Background feature at each alt bin by using the marginal distributions of a background profile or produce a single anomaly score for each altitude bin by using the joint distributions of a background profile. You can declare which method to use by setting the `method` argument to either \"marginal\" or \"joint\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape of Marginal Method Anomalies: (1, 30, 12), Avg Score: 4.10\n",
      "Shape of Joint Method Anomalies: (1, 30), Avg Score 0.93\n"
     ]
    }
   ],
   "source": [
    "marginal_anomalies = anomaly_score_bv(bvs[sample_idx, ...], G_sampled_bvs, method='marginal', alpha=2.0)\n",
    "print(f'Shape of Marginal Method Anomalies: {marginal_anomalies.shape}, Avg Score: {np.nanmean(marginal_anomalies):.2f}')\n",
    "\n",
    "joint_anomalies = anomaly_score_bv(bvs[sample_idx, ...], G_sampled_bvs, method='joint', alpha=2.0)\n",
    "print(f'Shape of Joint Method Anomalies: {joint_anomalies.shape}, Avg Score {np.nanmean(joint_anomalies):.2f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We can visualize the anomalies scores by plotting the generated background profiles along side the real background profile measurements. The real background profile will be colored based on its anomaly score so we can visualize how the background measurements compare to the calculated anomaly scores across all the altitudes. Remember a high anomaly score means that the sample was found to be normal vs low score means the sample was found to be irregular."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": "<Figure size 3500x1500 with 24 Axes>",
      "image/png": "[encoded data, 1134860 characters]\n"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.style.use('default')\n",
    "fig, ax = plt.subplots(2, 6, figsize=(35, 15))\n",
    "pal = sns.color_palette('hls', n_colors=2)\n",
    "alt_grid = const.altgrid[1:const.max_alt]\n",
    "for i in range(6):\n",
    "    for j in range(2):\n",
    "        bv_index = 2*i + j\n",
    "        for k in range(9):\n",
    "            ax[j, i].plot(G_sampled_bvs[0, k, 1:, bv_index], alt_grid, color=pal[1])\n",
    "        ax[j, i].plot(G_sampled_bvs[0, -1, 1:, bv_index], alt_grid, label='Generated', color=pal[1])\n",
    "        # ax[j, i].plot(bvs[sample_idx[0], :, bv_index], alt_grid, label='Real', color=pal[0],)\n",
    "        alts_interp = np.linspace(const.altgrid[1], const.altgrid[const.max_alt - 1], 10000)\n",
    "        bvs_interp = np.interp(alts_interp, xp=const.altgrid[1:const.max_alt], fp=bvs[sample_idx[0], 1:, bv_index],)\n",
    "        anom_interp = np.interp(alts_interp, xp=const.altgrid[1:const.max_alt], fp=marginal_anomalies[0, 1:, bv_index],)\n",
    "        im = ax[j, i].scatter(bvs_interp, alts_interp, label='Real', c=anom_interp, cmap='plasma', zorder=11,\n",
    "                              vmin=max(np.nanmin(marginal_anomalies[0, 1:, bv_index]), -100), vmax=np.nanmax(marginal_anomalies[0, 1:, bv_index]))\n",
    "        ax[j, i].legend()\n",
    "        ax[j, i].set_xlabel(const.bv_units[bv_index], fontsize=10)\n",
    "        ax[j, i].set_ylabel('Altitude (km)', fontsize=10)\n",
    "        ax[j, i].set_xlabel(const.bv_units[bv_index], fontsize=10)\n",
    "        ax[j, i].set_title(f'{const.bv_full_names[bv_index]} '\n",
    "                           f'(Avg Score = {np.nanmean(marginal_anomalies[0, 1:, bv_index]):.2f})',\n",
    "                           fontsize=15)\n",
    "        fig.colorbar(im, ax=ax[j, i], label='Anomaly Score')\n",
    "\n",
    "plt.suptitle(f'Avg Anomaly Score', fontsize=20)\n",
    "plt.tight_layout(rect=[0, 0, 1, 0.95])\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "\n",
    "## HFP Anomaly Example\n",
    "\n",
    "When looking at detecting anomalies in the HFP data we can compute an anomaly score using generated HFPs that were either conditioned with real background profile measurements or measurements generated from the BV GAN. To get anomaly scores using HFPs conditioned on real background profiles we first call the `multi_generate` function with the real BV measurements passed in as an argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating HFP Samples: 100%|██████████| 1/1 [00:01<00:00,  1.20s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "G_hfps Shape = (1, 500, 1, 8)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "G_sampled_hfps, G_sampled_b = generate_multi_hfp(drivers[sample_idx, ...], bvs=bvs[sample_idx, ...], n_repeat=n_repeat)\n",
    "print(f\"G_hfps Shape = {G_sampled_hfps.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Once we have the sampled HFP waves that were conditioned on real background measurements we pass them into the `anomaly_score_hfp` function. Again we can perform marginal based anomaly estimation, which will return an anomaly score for each wave parameters in a sample, or a joint based anomaly estimation, which will return a single anomaly score for each sample."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape of Anomalies: (1, 1, 8)\n",
      "['GW_tau: 3.08', 'GW_lamh: 5.61', 'GW_lamv: 6.78', 'GW_ch: 5.77', 'GW_cv: 6.12', 'GW_phi: 6.67', 'GW_hl: 6.61', 'GW_hu: -13.87']\n"
     ]
    }
   ],
   "source": [
    "anomalies_hfp = anomaly_score_hfp(hfps[sample_idx, :], G_sampled_hfps, method='marginal', alpha=2.0)\n",
    "print(f'Shape of Anomalies: {anomalies_hfp.shape}')\n",
    "print([f'{const.hfp_names[i]}: {anomalies_hfp[0, 0, i]:.2f}' for i in range(const.n_hfp)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape of N-Dim Anomalies: (1, 1), N-Dim Anomaly Score: 1.83\n"
     ]
    }
   ],
   "source": [
    "nd_anomalies_hfp = anomaly_score_hfp(hfps[sample_idx, :], G_sampled_hfps, method='joint', alpha=2.0)\n",
    "print(f'Shape of N-Dim Anomalies: {nd_anomalies_hfp.shape}, N-Dim Anomaly Score: {nd_anomalies_hfp.item():.2f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "On the other hand we can also compute anomaly scores using HFPs that were not conditioned with real background profile measurements. This will require generating the HFPs again using the `multi_generate` function, but this time we won't pass in any real background profile measurements. Again we can use the `anomaly_score_hfp` function and pass in the non-conditioned generated HFPs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating HFP Samples: 100%|██████████| 1/1 [00:01<00:00,  1.23s/it]\n"
     ]
    }
   ],
   "source": [
    "G_sampled_hfps, G_sampled_b = generate_multi_hfp(drivers[sample_idx, ...], bvs=G_bvs[sample_idx, ...], n_repeat=n_repeat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Non-Conditioned Scores\n",
      "['GW_tau: -57.65', 'GW_lamh: -5.20', 'GW_lamv: 4.52', 'GW_ch: 5.64', 'GW_cv: 2.43', 'GW_phi: 5.51', 'GW_hl: 7.17', 'GW_hu: -1.00']\n"
     ]
    }
   ],
   "source": [
    "anomalies_hfp = anomaly_score_hfp(hfps[sample_idx, :], G_sampled_hfps, method='marginal', alpha=1.0)\n",
    "print(f'Non-Conditioned Scores')\n",
    "print([f'{const.hfp_names[i]}: {anomalies_hfp[0, 0, i]:.2f}' for i in range(const.n_hfp)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Non-Conditioned N-Dim Anomaly Score: 0.58\n"
     ]
    }
   ],
   "source": [
    "nd_anomalies_hfp = anomaly_score_hfp(hfps[sample_idx, :], G_sampled_hfps, method='joint', alpha=1.0)\n",
    "print(f'Non-Conditioned N-Dim Anomaly Score: {nd_anomalies_hfp.item():.2f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Comparing Real Data Vs Fake Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We can also use the anomaly score as a metric of how well a set of background profile or waves measurements fits the historically observed distributions that our model has learned to characterize. A simple experiment we can perform to verify this is to compare the real background profiles from our dataset to obviously fake background profile and wave measurements. For this example the obviously fake background and wave measurements are randomly selected from uniform distributions between known upper and lower bounds for each variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "fake_hfp_data = np.zeros_like(G_sampled_hfps)\n",
    "fake_bv_data = np.zeros_like(G_sampled_bvs)\n",
    "for i in range(const.n_hfp_feat):\n",
    "    fake_hfp_data[:, :, :, i] = np.random.uniform(low=const.hfp_thresholds[i, 0],\n",
    "                                                  high=const.hfp_thresholds[i, 1], size=(1, n_repeat, 1))\n",
    "for i in range(const.n_bv_feat):\n",
    "    fake_bv_data[:, :, :, i] = np.random.uniform(low=const.bv_thresholds[i, 0],\n",
    "                                                 high=const.bv_thresholds[i, 1], size=(1, n_repeat, const.max_alt))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We will compute the anomaly scores of the fake data we just randomly generated and compare them to the anomaly scores generated in the previous example. We expect that the anomaly scores of the fake data will be significantly lower than the anomaly scores we generated using background profiles generated from our model, which is exactly what we see when we compare the scores below. Note we are using `np.nanmean()` to get an overall score because some samples have missing wave and/or background profile measurements which would produce nans in those positions as expected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Real Background Measurements Anomaly Score: -4.068878002833192\n",
      "Fake Background Measurements Anomaly Score: -135.8449236056502\n"
     ]
    }
   ],
   "source": [
    "real_anomaly_scores_bv = anomaly_score_bv(bvs[sample_idx, ...], G_sampled_bvs, alpha=1.0, method='joint')\n",
    "fake_anomaly_scores_bv = anomaly_score_bv(bvs[sample_idx, ...], fake_bv_data, alpha=1.0, method='joint')\n",
    "\n",
    "print(f'Real Background Measurements Anomaly Score: {np.nanmean(real_anomaly_scores_bv)}')\n",
    "print(f'Fake Background Measurements Anomaly Score: {np.nanmean(fake_anomaly_scores_bv)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We can do the same for the HFPs Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Real HFP Measurements Anomaly Score: 0.5806393550076\n",
      "Fake HFP Measurements Anomaly Score: -36.216158096763955\n"
     ]
    }
   ],
   "source": [
    "real_anomaly_scores_hfp = anomaly_score_hfp(hfps[sample_idx, ...], G_sampled_hfps, alpha=1.0, method='joint')\n",
    "fake_anomaly_scores_hfp = anomaly_score_hfp(hfps[sample_idx, ...], fake_hfp_data, alpha=1.0, method='joint')\n",
    "\n",
    "print(f'Real HFP Measurements Anomaly Score: {np.nanmean(real_anomaly_scores_hfp)}')\n",
    "print(f'Fake HFP Measurements Anomaly Score: {np.nanmean(fake_anomaly_scores_hfp)}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Attention Map Analysis\n",
    "\n",
    "Both BV and HFP GANs are conditioned using a transformer model, which makes use an attention mechanism to weigh the importance of background measurements from subsequent altitude bins. The transformer is split into an encoder, which takes in background variables and produces the conditional inputs for the BV Generator and a decoder which takes in the encoder outputs and produces conditional inputs for the HPF Generator. The encoder has a self-attention mechanism that looks at previous altitude measurements to help make predictions for higher altitude bins. The self-attention mechanism learns to weigh subsequent measurements differently allowing it to pay more attention to inputs that help the encoder produce better conditional inputs for the BV Generator. Like the encoder the decoder also has an attention mechanism that learns to weigh the importance of outputs from the encoder at specific altitude bins.\n",
    "\n",
    "These attention weights shows what our models learned to pay attention to in order background profiles and waves under specific conditions and may display patterns that show some underlying physics the has model learned in order to generate BV profiles and HFP waves. We can view attention maps for the encoder and decoder given a set of inputs (Drivers/BVs) using the `collect_bv_attn_map` and  `collect_hfp_attn_map` functions. The BV self-attention map is collected from the encoder and has the output shape of (n_alt x n_alt) and the HFP attention map is collected from the decoder and has th output shape of (1 x n_alt).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Attention Map Shape: (1000, 30, 30)\n",
      "(1000, 30, 30) Samples\n",
      "30 x 30 Self-Attention Weights Per Altitude\n",
      "\n",
      "HFP Attention Map Shape: (1000, 1, 30)\n",
      "(1000, 1, 30) Samples\n",
      "1 x 30 Attention Weights Per Altitude\n"
     ]
    }
   ],
   "source": [
    "attn_map_bv = collect_bv_attn_map(drivers, bvs)\n",
    "attn_map_hfp = collect_hfp_attn_map(drivers, bvs)\n",
    "\n",
    "print(f'Attention Map Shape: {attn_map_bv.shape}')\n",
    "print(f'{attn_map_bv.shape} Samples')\n",
    "print(f'{attn_map_bv.shape[1]} x {attn_map_bv.shape[2]} Self-Attention Weights Per Altitude\\n')\n",
    "print(f'HFP Attention Map Shape: {attn_map_hfp.shape}')\n",
    "print(f'{attn_map_hfp.shape} Samples')\n",
    "print(f'{attn_map_hfp.shape[1]} x {attn_map_hfp.shape[2]} Attention Weights Per Altitude')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## BV Self-Attention Map\n",
    "\n",
    "We can visualize some general attention patterns learned by the model by plotting the maps averaged across all the sample. Since we generate the BV profile layer by layer from bottom to top we have to mask the upper values to ensure that the encoder is only allowed to use values from subsequent lower altitudes to predict values of the next altitude bin, which explains why the upper diagonal of the attention map has no attention. If we look at the BV's self-attention map from the encoder can still see that there are few distinct patterns. Some obvious ones include the consistently high attention along 0th altitude bin and the sudden decrease of attention that higher altitude layers place on measurements below the 13th altitude bin.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": "<Figure size 1600x1400 with 2 Axes>",
      "image/png": "[encoded data, 187860 characters]\n"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.style.use('dark_background')\n",
    "fig, ax = plt.subplots(figsize=(16, 14))\n",
    "sns.heatmap(np.mean(attn_map_bv, axis=0),\n",
    "            ax=ax, vmax=0.2, vmin=.01, cmap='inferno',\n",
    "            yticklabels=[f'({const.altgrid[i].astype(int)} km) {i}' for i in range(const.max_alt)],)\n",
    "cbar = ax.collections[0].colorbar\n",
    "vmax, vmin = cbar.get_ticks()[0], cbar.get_ticks()[-1]\n",
    "cbar.set_ticks(np.hstack((np.array([vmin]),np.array([vmax]))))\n",
    "cbar.set_ticklabels(['High Attention'] + ['Low Attention'])\n",
    "cbar.ax.tick_params(labelsize=20)\n",
    "ax.tick_params(axis='x', labelsize=15, rotation=45)\n",
    "ax.tick_params(axis='y', labelsize=15)\n",
    "ax.set_ylabel(\"Altitude Bins\", fontsize=20)\n",
    "ax.set_xlabel(\"Altitude Bins\", fontsize=20)\n",
    "ax.set_title(\"Background Self-Attention Map\", fontsize=25)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## HFP Attention Map\n",
    "\n",
    "We can also visualize the general attention pattern that decoder has on the inputs from the encoder. Since there is only a max of one wave per sample the attention map is really just an array, which we can see below. As you can see on average most of the attention is directed toward the values between altitude bin 11 - 18, which is makes sense because most of the HFP waves occur in that altitude region as shown by the histogram of upper and lower hfp bounds in the bottom figure.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": "<Figure size 2000x500 with 2 Axes>",
      "image/png": "[encoded data, 98596 characters]\n"
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "<Figure size 1500x500 with 1 Axes>",
      "image/png": "[encoded data, 44844 characters]\n"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.style.use('dark_background')\n",
    "fig, ax = plt.subplots(figsize=(20, 5))\n",
    "sns.heatmap(np.mean(attn_map_hfp, axis=0),\n",
    "            ax=ax, cmap='inferno',\n",
    "            xticklabels=[f'({const.altgrid[i].astype(int)} km) {i}' for i in range(const.max_alt)],)\n",
    "cbar = ax.collections[0].colorbar\n",
    "vmax, vmin = cbar.get_ticks()[0], cbar.get_ticks()[-1]\n",
    "cbar.set_ticks(np.hstack((np.array([vmin]),np.array([vmax]))))\n",
    "cbar.set_ticklabels(['High Attention'] + ['Low Attention'])\n",
    "cbar.ax.tick_params(labelsize=20)\n",
    "ax.set_yticks([])\n",
    "ax.tick_params(axis='x', labelsize=15, rotation=45)\n",
    "ax.set_xlabel(\"Altitude Bins\", fontsize=20)\n",
    "ax.set_title(\"HFP Attention Map\", fontsize=25)\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "plt.style.use('default')\n",
    "fig, ax = plt.subplots(figsize=(15, 5))\n",
    "x = np.hstack((hfps[wave_mask, 6], hfps[wave_mask, 7]))\n",
    "hue = [\"Lower Bound\"] * len(hfps[wave_mask, 6]) + [\"Upper Bound\"] * len(hfps[wave_mask, 7])\n",
    "sns.histplot(x=np.hstack((hfps[wave_mask, 6], hfps[wave_mask, 7])),\n",
    "             hue=[\"Lower Bound\"] * len(hfps[wave_mask, 6]) + [\"Upper Bound\"] * len(hfps[wave_mask, 7]),\n",
    "             ax=ax, stat='density', bins=25, common_norm=False)\n",
    "ax.set_ylabel('Denity')\n",
    "ax.set_ylim(0, 0.05)\n",
    "ax.set_xlabel('Altitude (km)')\n",
    "ax.set_title(\"HFP Upper and Lower Bound Histograms\", fontsize=25)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "# TEC Waves\n",
    "\n",
    "ABC machine learning methods can be applied to characterize the atmosphere using alternative more globally available data sources. Specifically, this next section will show how we can use the same machine learning techniques to characterize atmospheric waves derived from TEC measurements collected from various GNSS stations near pokers. In theory, these TEC waves can be derived for any location which has nearby GNSS stations and later in the tutorial we will be showing how we trained our GAN to characterize the atmospheric over Mt Moses in Nevada using TEC wave data. Using the more globally available TEC data let us generalize our models to potentially work on a global scale, i.e. not just at Poker Flats.\n",
    "\n",
    "The data team for ABC created two types of wTEC datasets. One dataset contains large scale transient ionospheric disturbances (LSTIDs), which contain waves that have longer periods between 3600-7200 seconds, and the other dataset contains medium scale transient ionospheric disturbances (MSTIDs), which contain waves with periods between 600-3600 seconds. We can load each dataset using the `load_wtec_h5` function. The wTEC data has the structure `(n_samples, n_wave_paramters)`.  We will first start off by loading the LSTIDs, which can be seen bellow. Like the previous data loader for the ISR the wTEC data loader will return the aligned drivers and unix timestamps. The data files used in this tutorial only contain a small subset of 5000 samples, which range from 2013-2014 shown in the plot below. The full dataset contains over 60,0000 wave samples and ranges from 2010-2020.\n",
    "\n",
    "\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Drivers Shape: (5000, 29)\n",
      "wTEC Shape: (5000, 18)\n",
      "Unix Timestamp Shape: (5000,)\n",
      "wtEC Wave Parameter Names: ['Amp_min', 'Amp_avg', 'Amp_max', 'Tau_min', 'Tau_avg', 'Tau_max', 'Lambda_min', 'Lambda_avg', 'Lambda_max', 'C1_min', 'C1_avg', 'C1_max', 'C2_min', 'C2_avg', 'C2_max', 'Phi_min', 'Phi_avg', 'Phi_max']\n"
     ]
    },
    {
     "data": {
      "text/plain": "<Figure size 800x400 with 1 Axes>",
      "image/png": "[encoded data, 21680 characters]\n"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dataset_name = 'LSTIDs_Poker'\n",
    "wtec_model = f'wtec_gan_{dataset_name}'\n",
    "drivers, wtec_lstids, unix_time = load_wtec_h5(f'wtec_{dataset_name}.h5')\n",
    "print(f'Drivers Shape: {drivers.shape}')\n",
    "print(f'wTEC Shape: {wtec_lstids.shape}')\n",
    "print(f'Unix Timestamp Shape: {unix_time.shape}')\n",
    "\n",
    "print(f'wtEC Wave Parameter Names: {const.wtec_names}')\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(8, 4))\n",
    "ax.hist([datetime.datetime.utcfromtimestamp(t) for t in unix_time], bins=40)\n",
    "ax.set_ylabel(\"Count\")\n",
    "ax.set_xlabel('UTC Time')\n",
    "ax.set_title('LSTIDs UTC Distribution')\n",
    "plt.show()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## wTEC GAN (LSTIDs Poker Flats)\n",
    "Like the HFP GAN, we created a similar more simplified wTEC GAN, which aims to generate historically accurate transient ionospheric disturbances, i.e. atmospheric waves, using a dataset we just loaded above. Unlike the HFP GAN the wTEC GAN is only conditioned on the external drivers when generating wave parameters and thus doesn't require an altitude based transformer for conditioning. To generate LSTID wTECs given a set of external driver you can call the `generate_wtec` function as shown below."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating TEC Waves: 100%|██████████| 50/50 [00:00<00:00, 288.37it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated wTEC LSTIDs Shape: (5000, 18)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "G_wtec_lstids = generate_wtec(drivers=drivers, dataset_name=dataset_name, wtec_model=wtec_model)\n",
    "print(f'Generated wTEC LSTIDs Shape: {G_wtec_lstids.shape}')"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "We can evaluate the wTEC using the same Hellinger distance metric by calling the `hellinger_scores_wtec` function and passing in the real and generated samples as shown below. You can also see in the plot below how our generated distributions match the real distributions for each the wave parameters."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['Amp_min Dist: 0.063', 'Amp_avg Dist: 0.065', 'Amp_max Dist: 0.065', 'Tau_min Dist: 0.089', 'Tau_avg Dist: 0.084', 'Tau_max Dist: 0.120', 'Lambda_min Dist: 0.077', 'Lambda_avg Dist: 0.089', 'Lambda_max Dist: 0.113', 'C1_min Dist: 0.058', 'C1_avg Dist: 0.068', 'C1_max Dist: 0.064', 'C2_min Dist: 0.083', 'C2_avg Dist: 0.084', 'C2_max Dist: 0.082', 'Phi_min Dist: 0.107', 'Phi_avg Dist: 0.109', 'Phi_max Dist: 0.108']\n"
     ]
    },
    {
     "data": {
      "text/plain": "<Figure size 2500x4000 with 18 Axes>",
      "image/png": "[encoded data, 359548 characters]\n"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dists = hellinger_scores_wtec(wtec_lstids, G_wtec_lstids, dataset_name=dataset_name, filter_length=5)\n",
    "print([f'{n} Dist: {d:.3f}' for (n, d) in zip(const.wtec_names, dists)])\n",
    "\n",
    "plt.style.use('default')\n",
    "fig, ax = plt.subplots(len(const.wtec_vars), 3, figsize=(25, 40))\n",
    "ax = ax.flatten()\n",
    "for i in range(len(ax)):\n",
    "    x = np.hstack((wtec_lstids[... ,i].flatten(), G_wtec_lstids[..., i].flatten()))\n",
    "    hue = ['Real'] * len(wtec_lstids[... ,i].flatten()) + ['Generated'] * len(G_wtec_lstids[..., i].flatten())\n",
    "    data = pd.DataFrame({f'{const.wtec_units[i]}': x, 'Type': hue})\n",
    "    sns.histplot(data=data, x=f'{const.wtec_units[i]}', hue='Type', ax=ax[i],\n",
    "                 element='step', bins=25, legend=True, stat='density',\n",
    "                 palette='hls', alpha=0.2, common_norm=False,\n",
    "                 hue_order=['Generated', 'Real'])\n",
    "    ax[i].set_title(f'{const.wtec_names[i]} (Hel Dist = {dists[i]:.3f})')\n",
    "plt.show()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## TEC Wave anomaly scores\n",
    "\n",
    "Again we can use the same techniques previously used on the ISR data to estimate anomaly scores for the observations in the TEC wave dataset. We first generate 1,000 samples for a particular external driver and then use our generated distribution to estimate anomaly scores. To generate the samples for each driver we call the `generate_multi_wtec` function passing in the drivers as conditioning. The multi generate function will produce  (n_samples, n_repeat, n_wtec). This is shown below for the first 500 samples to speed  things up.\n",
    "\n",
    "You can then pass the sampled wave parameters into the `anomaly_score_wtec` function, which will perform the same joint anomaly score estimation that is performed with the HFP data. The `anomaly_score_wtec` function will estimate a anomaly score for each input sample based on the generated distribution. The plot below show the distribution of estimate anomaly scores produced from the wTEC GAN. The lower the score the more anomaly the observation is. Most of the scores sit around a score of 1, but you can see a couple samples in the left tail of the distribution with highly anomalous scores around -15."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating Samples: 100%|██████████| 500/500 [00:18<00:00, 27.11it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sampled wTECs shape: (500, 1000, 18)\n",
      "wTEC anomaly Score Shape: (500,)\n",
      "Average wTEC Score: 1.207\n"
     ]
    },
    {
     "data": {
      "text/plain": "<Figure size 800x400 with 1 Axes>",
      "image/png": "[encoded data, 20276 characters]\n"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "G_wtec_lstids = generate_multi_wtec(drivers=drivers[:500], n_repeat=1000, dataset_name=dataset_name, wtec_model=wtec_model, verbose=1)\n",
    "print(f'Sampled wTECs shape: {G_wtec_lstids.shape}')\n",
    "\n",
    "wtec_anomaly_scores = anomaly_score_wtec(wtec_lstids[:500], G_wtec_lstids, dataset_name=dataset_name)\n",
    "print(f'wTEC anomaly Score Shape: {wtec_anomaly_scores.shape}')\n",
    "print(f'Average wTEC Score: {np.mean(wtec_anomaly_scores):.3f}')\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(8, 4))\n",
    "ax.hist(wtec_anomaly_scores, bins=40)\n",
    "ax.set_xlabel('Joint Anomaly Score')\n",
    "ax.set_ylabel('Count')\n",
    "ax.set_title('wTEC Anomaly Score Distribution')\n",
    "plt.show()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Mt Moses (MSTID Waves)\n",
    "\n",
    "We also have an identical wTEC GAN trained to characterize medium transient ionospheric disturbances (MSTIDs) over Mt Moses in Nevada, which can use the name generation and anomaly score functions as we did for the Poker Flat's LSTIDs waves. Just make sure to pass the correct dataset name/model into the generator and other functions. To load the MSTIDs call the `load_wtec_h5` function and pass in the MSTID wave dataset file. The MSTIDs file for this tutorial contains observations during the year of 2016 as shown below."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Drivers Shape: (5000, 29)\n",
      "wTEC Shape: (5000, 18)\n",
      "Unix Timestamp Shape: (5000,)\n",
      "wtEC Wave Parameter Names: ['Amp_min', 'Amp_avg', 'Amp_max', 'Tau_min', 'Tau_avg', 'Tau_max', 'Lambda_min', 'Lambda_avg', 'Lambda_max', 'C1_min', 'C1_avg', 'C1_max', 'C2_min', 'C2_avg', 'C2_max', 'Phi_min', 'Phi_avg', 'Phi_max']\n"
     ]
    },
    {
     "data": {
      "text/plain": "<Figure size 800x400 with 1 Axes>",
      "image/png": "[encoded data, 22456 characters]\n"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dataset_name = 'MSTIDs_MtMoses'\n",
    "wtec_model = f'wtec_gan_{dataset_name}'\n",
    "drivers, wtec_mstids, unix_time = load_wtec_h5(f'wtec_{dataset_name}.h5')\n",
    "print(f'Drivers Shape: {drivers.shape}')\n",
    "print(f'wTEC Shape: {wtec_mstids.shape}')\n",
    "print(f'Unix Timestamp Shape: {unix_time.shape}')\n",
    "\n",
    "print(f'wtEC Wave Parameter Names: {const.wtec_names}')\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(8, 4))\n",
    "ax.hist([datetime.datetime.utcfromtimestamp(t) for t in unix_time], bins=40)\n",
    "ax.set_ylabel(\"Count\")\n",
    "ax.set_xlabel('UTC Time')\n",
    "ax.set_title(f'{dataset_name} UTC Distribution')\n",
    "plt.show()\n",
    "\n",
    "# G_wtec_mstids = generate_wtec(drivers=drivers, dataset_name=dataset_name, wtec_model=wtec_model)\n",
    "# print(f'Generated wTEC {dataset_name} Shape: {G_wtec_mstids.shape}')\n",
    "#\n",
    "# dists = hellinger_scores_wtec(wtec_mstids, G_wtec_mstids, dataset_name=dataset_name, filter_length=2)\n",
    "# print([f'{n} Dist: {d:.3f}' for (n, d) in zip(const.wtec_names, dists)])\n",
    "#\n",
    "#\n",
    "# plt.style.use('default')\n",
    "# fig, ax = plt.subplots(len(const.wtec_vars), 3, figsize=(25, 40))\n",
    "# ax = ax.flatten()\n",
    "# for i in range(len(ax)):\n",
    "#     x = np.hstack((wtec_mstids[... ,i].flatten(), G_wtec_mstids[..., i].flatten()))\n",
    "#     hue = ['Real'] * len(wtec_mstids[... ,i].flatten()) + ['Generated'] * len(G_wtec_mstids[..., i].flatten())\n",
    "#     data = pd.DataFrame({f'{const.wtec_units[i]}': x, 'Type': hue})\n",
    "#     sns.histplot(data=data, x=f'{const.wtec_units[i]}', hue='Type', ax=ax[i],\n",
    "#                  element='step', bins=25, legend=True, stat='density',\n",
    "#                  palette='hls', alpha=0.2, common_norm=False,\n",
    "#                  hue_order=['Generated', 'Real'])\n",
    "#     ax[i].set_title(f'{const.wtec_names[i]} (Hel Dist = {dists[i]:.3f})')\n",
    "# plt.show()\n",
    "#\n",
    "#\n",
    "# G_wtec_mstids = generate_multi_wtec(drivers=drivers[:500], n_repeat=1000, dataset_name=dataset_name, wtec_model=wtec_model, verbose=1)\n",
    "# print(f'Sampled wTECs shape: {G_wtec_mstids.shape}')\n",
    "#\n",
    "# wtec_anomaly_scores = anomaly_score_wtec(wtec_mstids[:500], G_wtec_mstids, dataset_name=dataset_name)\n",
    "# print(f'wTEC anomaly Score Shape: {wtec_anomaly_scores.shape}')\n",
    "# print(f'Average wTEC Score: {np.mean(wtec_anomaly_scores):.3f}')\n",
    "#\n",
    "# fig, ax = plt.subplots(figsize=(8, 4))\n",
    "# ax.hist(wtec_anomaly_scores, bins=40)\n",
    "# ax.set_xlabel('Joint Anomaly Score')\n",
    "# ax.set_ylabel('Count')\n",
    "# ax.set_title('wTEC Anomaly Score Distribution')\n",
    "# plt.show()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating TEC Waves: 100%|██████████| 50/50 [00:06<00:00,  7.53it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated wTEC MSTIDs_MtMoses Shape: (5000, 18)\n",
      "['Amp_min Dist: 0.097', 'Amp_avg Dist: 0.092', 'Amp_max Dist: 0.093', 'Tau_min Dist: 0.142', 'Tau_avg Dist: 0.089', 'Tau_max Dist: 0.184', 'Lambda_min Dist: 0.113', 'Lambda_avg Dist: 0.094', 'Lambda_max Dist: 0.106', 'C1_min Dist: 0.112', 'C1_avg Dist: 0.097', 'C1_max Dist: 0.104', 'C2_min Dist: 0.111', 'C2_avg Dist: 0.111', 'C2_max Dist: 0.113', 'Phi_min Dist: 0.103', 'Phi_avg Dist: 0.103', 'Phi_max Dist: 0.100']\n"
     ]
    },
    {
     "data": {
      "text/plain": "<Figure size 2500x4000 with 18 Axes>",
      "image/png": "[encoded data, 376436 characters]\n"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "G_wtec_mstids = generate_wtec(drivers=drivers, dataset_name=dataset_name, wtec_model=wtec_model)\n",
    "print(f'Generated wTEC {dataset_name} Shape: {G_wtec_mstids.shape}')\n",
    "\n",
    "dists = hellinger_scores_wtec(wtec_mstids, G_wtec_mstids, dataset_name=dataset_name, filter_length=2)\n",
    "print([f'{n} Dist: {d:.3f}' for (n, d) in zip(const.wtec_names, dists)])\n",
    "\n",
    "\n",
    "plt.style.use('default')\n",
    "fig, ax = plt.subplots(len(const.wtec_vars), 3, figsize=(25, 40))\n",
    "ax = ax.flatten()\n",
    "for i in range(len(ax)):\n",
    "    x = np.hstack((wtec_mstids[... ,i].flatten(), G_wtec_mstids[..., i].flatten()))\n",
    "    hue = ['Real'] * len(wtec_mstids[... ,i].flatten()) + ['Generated'] * len(G_wtec_mstids[..., i].flatten())\n",
    "    data = pd.DataFrame({f'{const.wtec_units[i]}': x, 'Type': hue})\n",
    "    sns.histplot(data=data, x=f'{const.wtec_units[i]}', hue='Type', ax=ax[i],\n",
    "                 element='step', bins=25, legend=True, stat='density',\n",
    "                 palette='hls', alpha=0.2, common_norm=False,\n",
    "                 hue_order=['Generated', 'Real'])\n",
    "    ax[i].set_title(f'{const.wtec_names[i]} (Hel Dist = {dists[i]:.3f})')\n",
    "plt.suptitle(f'{dataset_name} GAN Results', fontsize=30)\n",
    "plt.show()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Use the Mt Moses MSTID GAN to generate 500 sets of wave parameters for 200 drivers."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating Samples: 100%|██████████| 200/200 [02:23<00:00,  1.39it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sampled wTECs shape: (200, 500, 18)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "n_samples = 200\n",
    "G_wtec_mstids = generate_multi_wtec(drivers=drivers[:n_samples], n_repeat=500, dataset_name=dataset_name, wtec_model=wtec_model, verbose=1)\n",
    "print(f'Sampled wTECs shape: {G_wtec_mstids.shape}')"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Call the same `anomaly_score_wtec` function with the appropriate dataset name and generated PDFs to estimate anomaly scores for the medium scale TEC waves over Mt Moses."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "wTEC anomaly Score Shape: (200,)\n",
      "Average wTEC Score: -3.487\n"
     ]
    },
    {
     "data": {
      "text/plain": "<Figure size 800x400 with 1 Axes>",
      "image/png": "[encoded data, 38236 characters]\n"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "wtec_anomaly_scores = anomaly_score_wtec(wtec_mstids[:n_samples], G_wtec_mstids, dataset_name=dataset_name)\n",
    "print(f'wTEC anomaly Score Shape: {wtec_anomaly_scores.shape}')\n",
    "print(f'Average wTEC Score: {np.mean(wtec_anomaly_scores):.3f}')\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(8, 4))\n",
    "ax.hist(wtec_anomaly_scores, bins=20)\n",
    "ax.set_xlabel('Joint Anomaly Score')\n",
    "ax.set_ylabel('Count')\n",
    "ax.set_title(f'{dataset_name} Anomaly Score Distribution', fontsize=30)\n",
    "plt.show()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "file_extension": ".py",
  "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.9.6"
  },
  "mimetype": "text/x-python",
  "name": "python",
  "npconvert_exporter": "python",
  "pygments_lexer": "ipython3",
  "version": 3
 },
 "nbformat": 4,
 "nbformat_minor": 4
}