NNET02 Interval Visualizer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01ff0b15-3a26-49d2-824b-d637c3ea515a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "import os\n",
    "import random\n",
    "import warnings\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pylab as plt\n",
    "import seaborn as sns\n",
    "from spacepy import pycdf\n",
    "\n",
    "import moms_fast\n",
    "import nnet_evaluate\n",
    "import utils\n",
    "\n",
    "%matplotlib inline\n",
    "\n",
    "sns.set_style('darkgrid')\n",
    "\n",
    "N_EN = 32\n",
    "N_EN_SHELLS = 2\n",
    "N_PHI = 32\n",
    "N_THETA = 16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b06f89ac-f2c6-472c-8639-91adcb83e910",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_data = nnet_evaluate.load_test_data('4D_tail')\n",
    "f1ct = utils.get_f1ct({'4D_tail': test_data}, ['4D_tail'])\n",
    "hidden_layer_size = 50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c17a31b8-0587-4a21-99ea-51069cc24b52",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv('/mnt/efs/dasilva/compression-cfha/data/test_train_split.csv')\n",
    "df = df[df['test_train'] == 'test']\n",
    "df = df[df['phase'] == '4D_tail']\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b4e086c-f85b-47de-8027-741d51a2442a",
   "metadata": {},
   "outputs": [],
   "source": [
    "files = [f for f in list(df.file_path) if 'mms1' in f]\n",
    "cdf_filename = random.choice(files)\n",
    "print(cdf_filename)\n",
    "cdf_filename = '/mnt/efs/dasilva/compression-cfha/data/mms_data/4D_tail/mms1_fpi_brst_l2_dis-dist_20190720043943_v3.3.0.cdf'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98a90ba0-cb82-45c4-84ad-243328415113",
   "metadata": {},
   "outputs": [],
   "source": [
    "cdf = pycdf.CDF(cdf_filename)\n",
    "\n",
    "dist = cdf['mms1_dis_dist_brst'][:]\n",
    "dist_err = cdf['mms1_dis_disterr_brst'][:]\n",
    "epoch = cdf['Epoch'][:]\n",
    "ntime = epoch.size\n",
    "counts = np.zeros((ntime, N_PHI, N_THETA, N_PHI))\n",
    "\n",
    "for i in range(ntime):\n",
    "    with warnings.catch_warnings():\n",
    "        warnings.filterwarnings('ignore')\n",
    "        tmp_counts = np.square(dist[i] / dist_err[i])\n",
    "    tmp_counts[np.isnan(tmp_counts)] = 0\n",
    "    tmp_counts = np.rint(tmp_counts)\n",
    "    counts[i] = tmp_counts\n",
    "    \n",
    "cdf.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "963d8c74-76a0-4283-9e7e-e42ce9700690",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "runs = ['4D_tail']\n",
    "models = {}\n",
    "\n",
    "for run in runs:\n",
    "    for en_index in range(0, N_EN, N_EN_SHELLS):\n",
    "        models[run, en_index] = nnet_evaluate.load_model(\n",
    "            '4D_tail', hidden_layer_size, en_index,\n",
    "            outpath=(f'/mnt/efs/dasilva/compression-cfha/data/nnet_models'\n",
    "                     f'/hidden_layer_exp/{run}/')\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b43b29fa-54b6-4ac0-88d6-95610dd2c548",
   "metadata": {},
   "outputs": [],
   "source": [
    "def reconstruct(counts, models):\n",
    "    counts_recon = np.zeros_like(counts)\n",
    "\n",
    "    for j in range(counts.shape[0]):\n",
    "        # Branch model\n",
    "        num_nonzeros = counts[j].nonzero()[0].size\n",
    "        \n",
    "        for en_index in range(0, N_EN, N_EN_SHELLS):\n",
    "            i, di = en_index, N_EN_SHELLS\n",
    "            \n",
    "            model = models[runs[0], i]\n",
    "            model_input = counts[j, :, :, i:i+di]     \n",
    "            model_output = model([model_input]).numpy()\n",
    "            counts_recon[j, :, :, i:i+di] = model_output[0]\n",
    "    \n",
    "            avg_orig = counts[j, :, :, i:i+di].mean()\n",
    "            avg_recon = counts_recon[j, :, :, i:i+di].mean()\n",
    "    \n",
    "            if avg_orig == 0:\n",
    "                counts_recon[j, :, :, i:i+di] = 0\n",
    "            elif avg_recon > 0:\n",
    "                counts_recon[j, :, :, i:i+di] *=  avg_orig / avg_recon \n",
    "    \n",
    "    return counts_recon "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0508990b-c46c-42b2-9c1d-78f73ef9a4ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "counts_recon = reconstruct(counts, models)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b3848b2-5a18-4007-bbaf-4a7050b5f966",
   "metadata": {},
   "outputs": [],
   "source": [
    "moms_true = [moms_fast.fast_moments(f1ct * c) for c in counts]\n",
    "moms_recon = [moms_fast.fast_moments(f1ct * c) for c in counts_recon]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78af58f4-b868-4b8c-9608-8fdb66c3053c",
   "metadata": {},
   "outputs": [],
   "source": [
    "vars = ['n', 'vx', 'vy', 'vz', 'txx', 'tyy', 'tzz']\n",
    "fig, axes = plt.subplots(len(vars), 1, figsize=(15, 4*len(vars)))\n",
    "\n",
    "for i, var in enumerate(vars):\n",
    "\n",
    "    #axes[i].set_title(var, fontsize=16)\n",
    "    axes[i].plot(epoch, [d[var] for d in moms_true], label=f'{var} True')\n",
    "    axes[i].plot(epoch, [d[var] for d in moms_recon], label=f'{var} Reconstructed')\n",
    "    axes[i].legend()\n",
    "    if var == 'n':\n",
    "        axes[i].set_ylim([0, 1.1 * np.max([d[var] for d in moms_true])])\n",
    "        axes[i].set_ylabel('n ($cm^{-3}$)', fontsize=16)\n",
    "    elif var[0] == 'v':\n",
    "        axes[i].set_ylabel(f'{var} (km/s)', fontsize=16)\n",
    "    elif var[0] == 't':\n",
    "        axes[i].set_ylabel(f'{var} (eV)', fontsize=16)\n",
    "\n",
    "    #axes[i].set_xlim(epoch[0], epoch[epoch.size//4])\n",
    "fig.suptitle(f'Dimensionality Reduction => {100*hidden_layer_size/(32*16*2):.1f}%')\n",
    "fig.tight_layout()\n",
    "os.makedirs('plots', exist_ok=True)\n",
    "fig.savefig('plots/' + os.path.basename(cdf_filename).split('.')[0] + '.png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1e2b77b-07dc-4e57-8cd2-543a8263b481",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ea2e958-477c-47b8-ae7e-c6468bce4e1f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1224e77-8d84-45cf-9520-b0ff3bc324bf",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Compression-CFHA",
   "language": "python",
   "name": "compression-cfha"
  },
  "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.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}