{
 "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 hightly unusualy (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 excercied using the GAN results in other latitude regiems."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Summary\n",
    "\n",
    "This file provides a demonstration of the abcgan project. This includes using the Generator to generate data consistent\n",
    "with the historically observed distribution and using the discriminator to detect whether measurements are consistent\n",
    "with the historically observed distribution."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports\n",
    "\n",
    "Several imports are used. In particular we load `numpy`, which provides the formats for data input output.\n",
    "We also use the numpy library for calculation of statistics and simulating fake data. We load `abcgan` which\n",
    "has both the generator and the discriminator that we will use. Finally, we load `matplotlib.pyplot`, `seaborn`,\n",
    "and `pandas` for display."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np  # numerics\n",
    "import abcgan  # abcgan package\n",
    "import matplotlib.pyplot as plt  # plotting\n",
    "import seaborn as sns  # plotting\n",
    "import pandas as pd  # plotting\n",
    "import h5py  # reading hdf5 files\n",
    "import matplotlib.image as mpimg\n",
    "import abcgan.constants as const\n",
    "import abcgan.transforms as trans\n",
    "from abcgan.interface import estimate_drivers\n",
    "from abcgan.mask import prev_driver_mask\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading Drivers and Background Variables\n",
    "\n",
    "** Replace with or link to drivers notebook? **\n",
    "\n",
    "We need to first get the driving parameters. We first load a set of relevant drivers that are stored in hdf5 format.\n",
    " We print the keys to show the saved list of drivers. We then convert to a standard python dictionary from hdf5 format\n",
    "  then use `abcgan.stack_drivers` to aggregate the drivers for input to the generator. For this demo we will only be\n",
    "  using 500 samples to simplify things. We also gather the unix time stamp for each driver sample, which will\n",
    "   be used later for estimating future driver samples.Some of the samples might be missing a driver feature, which we can mask\n",
    "  out by looking for any nans in a sample."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['Ap', 'F10.7', 'F10.7avg', 'MEI', 'MLT', 'RMM1', 'RMM2', 'SLT', 'SZA', 'TCI', 'ap', 'moon_phase', 'moon_x', 'moon_y', 'moon_z']\n"
     ]
    }
   ],
   "source": [
    "fname = 'tutorial_continuous.h5'\n",
    "start_sample, end_sample = 2000, 2500\n",
    "with h5py.File(fname, 'r') as hf:\n",
    "    driver_dict = hf['Drivers']\n",
    "    print(list(driver_dict.keys()))\n",
    "    driver_dict = {k: v[()] for k, v in driver_dict.items()}\n",
    "    unix_time = hf['UnixTime'][()]\n",
    "    drivers = abcgan.stack_drivers(driver_dict)\n",
    "\n",
    "    valid_samples = ~np.isnan(drivers).any(-1)\n",
    "    unix_time = unix_time[valid_samples][start_sample:end_sample]\n",
    "    drivers = drivers[valid_samples][start_sample:end_sample, :]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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",
    "\n",
    "\n",
    "\n",
    "Each driver is stored as a separate vector in `tutorial_v7.h5`, where each entry in the vector is a separate sample.\n",
    "So the full set of drivers for the first sample consists of the first element in each vector. The stacking routine\n",
    "`abcgan.stack_drivers` collates these into a single data structure which has size nSamples x nDrivers, in the order\n",
    "specified by `abcgan.driver_names`. `abcgan.stack_drivers` can accept either an `h5py.Group` object directly or\n",
    "as shown above a dictionary mapping each driver name to a one dimensional numpy array.\n",
    "\n",
    "See the tutorial on generating and saving drivers in this directory for how to create a driver file."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Background Variables\n",
    "\n",
    "We can load a sample of background variable observations collected at the same time as each of the sampled drivers and will be\n",
    " used later to compare our generated background values with the actual background values in the h5 file. The file contains about 5000 samples,\n",
    "  but for this demo we will only be using the first 500 samples to simplify things. Note that here we directly stack the background\n",
    "   variables using `abcgan.stack_bvs` from the `h5py.Group`. In the hdf5 file each background variable is saved as a single\n",
    "   array of size nSamples x nAltitude, where missing values are padded with nans. Note that the first dimensions (number of samples)\n",
    "    is 500, corresponding to the 500 sets of drivers. The final dimension is 6 corresponding to the number of background variables.\n",
    "    In case there was any invalid driver samples we apply the same valid mask to the background variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(500, 30, 12)\n"
     ]
    }
   ],
   "source": [
    "with h5py.File(fname, 'r') as hf:\n",
    "    bvs = abcgan.stack_bvs(hf['BackgroundValues'])\n",
    "    bvs = bvs[valid_samples][start_sample:end_sample, :, :]\n",
    "    print(bvs.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The background parameters, in the order specified by `abcgan.bv_names`, are:\n",
    "\n",
    "| Parameter | Variable Name | Units |\n",
    "|-----------|---------------|-------|\n",
    "| Average Electron Density | Ne_bac | m$^{-3}$ |\n",
    "| Root Mean Square Variation of Electron Density | Ne_rms | m$^{-3}$ |\n",
    "| Average Electron Temperature | Te_bac | K |\n",
    "| Root Mean Square Variation of Electron Temperature | Te_rms | K |\n",
    "| Average Ion Temperature | Ti_bac | K |\n",
    "| Root Mean Square Variation of Ion Temperature | Ti_rms | K |\n",
    "| Average Ion Velocity East | Ve_bac | m/s |\n",
    "| Root Mean Square Variation of Ion Velocity East | Ve_rms | m/s |\n",
    "| Average Ion Velocity North | Vn_bac | m/s |\n",
    "| Root Mean Square Variation of Ion Velocity North | Vn_rms | m/s |\n",
    "| Average Ion Velocity Up | Vu_bac | m/s |\n",
    "| Root Mean Square Variation of Ion Velocity Up | Vu_rms | m/s |\n",
    "| Average Neutral Density | Nn_bac | m$^{-3}$ |\n",
    "| Average Neutral Temperature | Tn_bac | K |\n",
    "\n",
    "\n",
    "Averages and root mean square values are over the standard smoothing window of 1 hour."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Altitudes\n",
    "The ABC GAN predicts background values on a set, non-uniform altitude grid.  The non-uniform grid was chosen to be higher resolution in the lower atmosphere where important dynamics occur at smaller scales, but lower resolution at high altitudes where the atmosphere is more stable across a larger range and data is sparser.\n",
    "\n",
    "| Altitude Range | Grid Spacing |\n",
    "|----------------|--------------|\n",
    "| 90 - 120 km    | 4.5 km       |\n",
    "| 120 - 150 km   | 9 km         |\n",
    "| 150 - 200 km   | 18 km        |\n",
    "| 200 - 800 km   | 24 km        |\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "altgrid = np.concatenate((np.arange(92.,119.,4.5),np.arange(119.,146.,9.),np.arange(146.,200.,18.),np.arange(200.,800.,24.)))    # altitude grid\n",
    "altgrid = altgrid[:30]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate\n",
    "\n",
    "Given the driving parameters we can now generate a set of measurements, one for each sample.\n",
    "\n",
    "The output is an array of measurements, one for each altitude bin, for each driver. In this case there are different sets of driving parameters. We choose to generate 30 altitude bins worth of data (which is both the default and the maximum as specified in `abcgan.max_alt`). There are six background variables generated for each sample for each altitude bin."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(500, 30, 12)\n",
      "['Ne_bac', 'Ne_rms', 'Te_bac', 'Te_rms', 'Ti_bac', 'Ti_rms', 'Ve_bac', 'Ve_rms', 'Vn_bac', 'Vn_rms', 'Vu_bac', 'Vu_rms']\n"
     ]
    }
   ],
   "source": [
    "meas = abcgan.generate(drivers, n_alt=30)\n",
    "print(meas.shape)\n",
    "print(abcgan.bv_names)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "\n",
    "## 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. Each output will be a separate sample. Here we generate 50 samples using the first set of saved drivers."
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(50, 30, 12)\n"
     ]
    }
   ],
   "source": [
    "meas_sampled = abcgan.generate(drivers[:1, :].repeat(50, 0))\n",
    "print(meas_sampled.shape)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compare Generated Profile\n",
    "As a practical example, look at several real profiles of electron density from the background values data 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<5 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": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": "<Figure size 1080x720 with 2 Axes>",
      "image/png": "[encoded data, 275508 characters]\n"
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "idx = np.argwhere((np.abs(drivers[:,3]-12.)<1.) & (drivers[:,0]<5.) & (drivers[:,1]<80.)).flatten()\n",
    "avg_drivers = np.mean(drivers[idx,:], axis=0)[None,:]\n",
    "\n",
    "gen_profiles = abcgan.generate(avg_drivers.repeat(50, 0))\n",
    "mean_profile = np.mean(gen_profiles, axis=0)\n",
    "std_profile = np.std(gen_profiles, axis=0)\n",
    "\n",
    "fig, ax = plt.subplots(1, 2, figsize=(15, 10))\n",
    "\n",
    "ax[0].plot(bvs[idx,:,0].T, altgrid, zorder=3)\n",
    "ax[0].set_title('Real Ne Profiles from PFISR')\n",
    "\n",
    "ax[1].plot(gen_profiles[:,:,0].T, altgrid[:30], zorder=3)\n",
    "ax[1].set_title('Ne Profiles Generated by GAN')\n",
    "\n",
    "for i in range(2):\n",
    "    ax[i].errorbar(mean_profile[:,0], altgrid[:30], xerr=std_profile[:,0], color='black', 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}$)')\n",
    "    ax[i].set_ylabel('Altitude (km)')\n",
    "    \n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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.\n",
    "\n",
    "The output of the GAN is **statistical**, and each individual generated profile is a sampled profile from the distribution but does not necessarily represent the \"average\" or \"typical\" conditions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Scoring \n",
    "\n",
    "To score the background variables we run the critic providing both the drivers and the measurements. Each altitude\n",
    "is scored separately, corresponding to how well the measurements at that altitude fit the historically observed\n",
    "distribution, conditional upon both the drivers and all lower altitude measurements. To score the entire measurement\n",
    "set overall we sum the contributions from all altitude bins."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(500, 30)\n",
      "(500,)\n"
     ]
    }
   ],
   "source": [
    "real_score = abcgan.discriminate(drivers, bvs)\n",
    "print(real_score.shape)\n",
    "total_score = real_score.sum(-1)\n",
    "print(total_score.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To score only the first ten bins we can just provide the first 10 data points."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(500, 10)\n",
      "(500,)\n"
     ]
    }
   ],
   "source": [
    "bin_score = abcgan.discriminate(drivers, bvs[:, :10, :])\n",
    "print(bin_score.shape)\n",
    "total_score = bin_score.sum(-1)\n",
    "print(total_score.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can directly compare the average scores of the real data versus obviously fake data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Real: 1.576124668121338\n",
      "Fake: 12.636296272277832\n"
     ]
    }
   ],
   "source": [
    "mean_real = real_score.mean()\n",
    "mean_fake = abcgan.discriminate(drivers, np.ones(bvs.shape)).mean()\n",
    "print(f'Real: {mean_real}')\n",
    "print(f'Fake: {mean_fake}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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",
    "- Demonstration showing how to compare the histograms of generated data with historically observed data.\n",
    "- Demonstrate how to use the discriminator score to distinguish real and fake data, including improved calibration."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Select an altitude bin, say the 25th, to compare histograms and filter out any possible outliers in the data using `valid_mask` created from `transforms.scale(bvs)` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "alt_bin = 20\n",
    "_, valid_mask = abcgan.transforms.scale_bv(bvs)\n",
    "real = bvs[valid_mask, alt_bin, :]\n",
    "generated = meas[valid_mask, alt_bin, :]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We can then compare the distributions of the `real` and `generate` values for each of the\n",
    "6 background variables using a kde plot implemented with matplotlib and seaborn. The nested for loop iterates over each of the 6 distubtions plots we want make. One is added to the `real` and `generated` values to avoid log scaling values less than or equal to zero"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": "<Figure size 1080x720 with 6 Axes>",
      "image/png": "[encoded data, 46792 characters]\n"
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots(2, 3, figsize=(15, 10))\n",
    "for i in range(2):\n",
    "    for j in range(3):\n",
    "        var1 = 1 + real[..., 3*i + j].flatten()\n",
    "        var2 = 1 + generated[..., 3*i + j].flatten()\n",
    "        x = np.hstack((var1, var2))\n",
    "        hue = ['Real'] * len(var1) + ['Generated'] * len(var2)\n",
    "        data = pd.DataFrame({'Value': x, 'Type': hue})\n",
    "        sns.histplot(data=data, x='Value', hue='Type', ax=ax[i, j],\n",
    "                     element='step', bins=25, legend=True, stat='density',\n",
    "                     palette='hls', alpha=0.2, common_norm=False,\n",
    "                     hue_order=['Generated', 'Real'], log_scale=True)\n",
    "        ax[i, j].set_title(f'BV {abcgan.constants.bv_names[3*i + j]} Histogram')\n",
    "plt.tight_layout()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Using Previous Drivers To Estimate Current Drivers\n",
    "\n",
    "In the situation that you don't have the current drivers available for predicting the background profile you can use the\n",
    "driver estimation model to predict the current driver features based on previous driver samples. Specifically the driver estimation\n",
    "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\n",
    "in which we only have drivers sample from two hours ago by creating a new driver dataset containing of driver samples that occurred two\n",
    "hours from the actual background variable samples. Since the time difference between each sample isn't always the same in our dataset we\n",
    "have to calculate a mapping vector between current samples and previous samples from two hour ago. The `prev_driver_mask` function\n",
    "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\n",
    "will use when evaluating the driver estimator."
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "outputs": [],
   "source": [
    "prev_dr_map, dr_mask = prev_driver_mask(unix_time)\n",
    "prev_drivers = drivers[prev_dr_map]"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "We can now estimate the current driver samples from the previous by calling the `estimate_drivers` function. As you can see the\n",
    "estimated drivers are very similar to actual driver samples. Note that when using in realtime you can simply call the\n",
    "estimate driver function with two hour old driver samples to predict the current driver samples."
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Original Driver Shape: (500, 15)\n",
      "Estimated Driver Shape: (449, 15)\n"
     ]
    }
   ],
   "source": [
    "est_drivers = estimate_drivers(prev_drivers, model='dr_gan')\n",
    "print(f'Original Driver Shape: {drivers.shape}')\n",
    "print(f'Estimated Driver Shape: {est_drivers.shape}')"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "We can look at how the estimated driver distribution compares to the actual driver samples. As you can see the\n",
    "estimated drivers are very similar to actual driver samples. From here you can pass the estimated driver samples\n",
    "into the background variable generator to estimate the current background profile."
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "outputs": [
    {
     "data": {
      "text/plain": "<Figure size 2880x1440 with 15 Axes>",
      "image/png": "[encoded data, 145076 characters]\n"
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "i_range = 3\n",
    "j_range = 5\n",
    "fig, ax = plt.subplots(i_range, j_range, figsize=(40, 20))\n",
    "for i in range(i_range):\n",
    "    for j in range(j_range):\n",
    "        x1 = drivers[:, j_range * i + j]\n",
    "        x2 = prev_drivers[:, j_range * i + j]\n",
    "        x = np.hstack((x1, x2))\n",
    "        hue = ['Estimated'] * len(x1) + ['Actual'] * len(x2)\n",
    "        data = pd.DataFrame({'Value': x, 'Type': hue})\n",
    "        sns.histplot(data=data, x='Value', hue='Type', ax=ax[i, j],\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, j].set_title(f'{const.driver_names[j_range * i + j]} Distribution')\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "is_executing": true,
     "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\n",
    "generated statistical distribution using a 'log-probability' distance measure."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Background\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\n",
    "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$\n",
    "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\n",
    "log of the exponent to estimate an anomaly score, $A$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": "<matplotlib.image.AxesImage at 0x14a658d00>"
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": "<Figure size 864x864 with 1 Axes>",
      "image/png": "[encoded data, 46680 characters]\n"
     },
     "metadata": {
      "needs_background": "light"
     },
     "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,\n",
    "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": [
    "## Putting it all together\n",
    "\n",
    "For our distance measure, we compute sum the exponential distances between the new sample point and 10,000 GAN-generated samples. The net effect is that we treat each new measurement as a random variable and the anomaly\n",
    "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.\n",
    "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.\n",
    "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": 16,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": "<matplotlib.image.AxesImage at 0x14a6bb340>"
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": "<Figure size 864x864 with 1 Axes>",
      "image/png": "[encoded data, 90664 characters]\n"
     },
     "metadata": {
      "needs_background": "light"
     },
     "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": [
    "## Anomaly Example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Fist we have to Generate fake data.\n",
    "Recall fakes at upper altitudes are conditioned on real data.\n",
    "Here we select a particular record with anomalies in the upper altitudes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "record number selected = [470]\n",
      "data shape = (1, 30, 12)\n",
      "drivers shape = (500, 15)\n",
      "driver names shape = (15,)\n",
      "driver means shape = (18,)\n",
      "driver sigmas shape = (18,)\n",
      "s.shape = (30, 12)\n",
      "std shape = (1, 30)\n",
      "sigma (before repeat) = (1, 30)\n",
      "sigma (after repeat) = (10000, 30)\n",
      "a = (10000, 30)\n"
     ]
    }
   ],
   "source": [
    "bv_type = 'radar '\n",
    "if bv_type == 'lidar':\n",
    "    n_alt = const.max_alt_lidar\n",
    "    n_feat = const.n_lidar_bv_feat\n",
    "else:\n",
    "    n_alt = const.max_alt\n",
    "    n_feat = const.n_bv_feat\n",
    "tic = time.time()\n",
    "\n",
    "# samples to test\n",
    "nr = 1\n",
    "sz = drivers.shape[0]\n",
    "\n",
    "records = np.random.randint(sz, size=[nr])\n",
    "print(\"record number selected = {}\".format(records))\n",
    "data_samples = bvs[records, :, :]\n",
    "driver_samples = drivers[records, :]\n",
    "\n",
    "# collect stats on the model\n",
    "print(\"data shape = {}\".format(data_samples.shape))\n",
    "print(\"drivers shape = {}\".format(drivers.shape))\n",
    "print(\"driver names shape = {}\".format(np.array(const.driver_names).shape))\n",
    "\n",
    "print(\"driver means shape = {}\".format(const.driver_mu.shape))\n",
    "print(\"driver sigmas shape = {}\".format(const.driver_sigma.shape))\n",
    "[fakes, data] = abcgan.gen_stats(drivers=driver_samples, data=data_samples, model='mm_gan_radar')\n",
    "anomalies = abcgan.anomaly_score(drivers=driver_samples, data=data_samples, model='mm_gan_radar')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Finally, we can plot the anomalies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "anomalies = [14.92777941 16.64825084 16.83946098 16.8130951  13.91926708 15.657292\n",
      " 15.65583634 12.15422554 15.18948565 16.39389544 15.39301809 15.31309815\n",
      " 15.64396862 19.10386474 20.12115042 19.15698786 19.41199831 19.08945927\n",
      " 19.16610978 18.53876584 18.68734108 18.77373597 18.68361816 18.25685881\n",
      " 18.35860592 18.37662432 17.97998188 18.28005114 17.71316692 18.07903172]\n"
     ]
    },
    {
     "data": {
      "text/plain": "<Figure size 2160x720 with 6 Axes>",
      "image/png": "[encoded data, 59484 characters]\n"
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "print('anomalies = {}'.format(anomalies))\n",
    "\n",
    "\n",
    "alts = [a[-1] for a in const.alt_bins]\n",
    "fig, ax = plt.subplots(2, 3, figsize=(30, 10))\n",
    "for ii in range(2):\n",
    "    for jj in range(3):\n",
    "        index = 3 * ii + jj\n",
    "        a = alts[index]\n",
    "        mask = fakes[:, a, :].flatten() != float(0)\n",
    "        x = np.hstack((fakes[:, a, :].flatten()[mask],\n",
    "                       data[:, a, :].flatten()[mask]))\n",
    "        hue = ['baseline'] * mask.sum() + ['new_data'] * mask.sum()\n",
    "        d = pd.DataFrame({'Value': x, 'Type': hue})\n",
    "        sns.histplot(data=d, x='Value', hue='Type',\n",
    "                     element='step', bins=100, legend=True, stat='density',\n",
    "                     palette='hls', alpha=0.2, common_norm=False, ax=ax[ii, jj],\n",
    "                     hue_order=['baseline', 'new_data'], log_scale=False)\n",
    "        ax[ii, jj].set_title(f'Alt Bin {const.alt_bins[index][-1]} Anomaly Score: {anomalies[alts[index]]:.2f}')\n",
    "plt.show()"
   ]
  },
  {
   "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
}