PREP01 Prepare Samples.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7805e554-6063-487a-bc76-29d3d6ac138a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import gc\n",
    "import h5py\n",
    "import joblib\n",
    "import json\n",
    "from matplotlib.colors import LogNorm\n",
    "import numpy as np\n",
    "import os\n",
    "import pandas as pd\n",
    "import progressbar\n",
    "import pylab as plt\n",
    "import warnings\n",
    "\n",
    "%matplotlib inline\n",
    "\n",
    "from spacepy import pycdf\n",
    "\n",
    "warnings.filterwarnings(\"ignore\", category=DeprecationWarning) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f57bb5de-1833-467a-a3b1-dc682893bb9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('hyperparams.json') as fh:\n",
    "    hyperparams = json.load(fh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc0e8390-6f38-4736-b0b7-aab3f34c6448",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_train_and_test = pd.read_csv('/mnt/efs/dasilva/compression-cfha/data/test_train_split.csv')\n",
    "df_train_and_test[df_train_and_test.file_path.str.contains('mms1')]\n",
    "df_train = df_train_and_test[df_train_and_test.test_train=='train']\n",
    "df_test = df_train_and_test[df_train_and_test.test_train=='test']\n",
    "df_train.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "981c6102-f27d-4a9c-abd3-bf85bd37e8cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_train.file_path.tolist()[:20]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da5749e2-9543-48a1-b44a-4f9eae1f9f0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_test.file_path.tolist()[:20]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6dcf8f2b-3c40-4f3c-9a19-8f9d3967fc4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "phases = sorted(set(df_train.phase))\n",
    "phases "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a146091-20f7-48d7-87cd-ba448c85644f",
   "metadata": {},
   "outputs": [],
   "source": [
    "for phase in phases:\n",
    "    print(f'{phase.ljust(15)} {df_train[df_train.phase==phase].lengths.sum()} frames')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2bbf5e4-3ebf-4298-b865-c4d30933baae",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_data(approx_count, df, phase):\n",
    "    df_phase = df[df.phase==phase]\n",
    "    \n",
    "    # Collect list of tasks\n",
    "    p = approx_count / df_phase.lengths.sum() \n",
    "    tasks = []\n",
    "    \n",
    "    for _, row in df_phase.iterrows():\n",
    "        bool_mask = np.random.rand(row.lengths) < p\n",
    "        if not bool_mask.any():\n",
    "            continue\n",
    "        tasks.append(joblib.delayed(get_from_file)(row.file_path, row.lengths, bool_mask))\n",
    "    \n",
    "    # Run tasks in parallel and aggregate results\n",
    "    results = joblib.Parallel(n_jobs=50, verbose=10)(tasks)\n",
    "    output = {key: [] for key in results[0].keys()}\n",
    "    \n",
    "    for result in results:\n",
    "        for key in result:\n",
    "            output[key].extend(result[key])\n",
    "\n",
    "    return output\n",
    "\n",
    "\n",
    "def get_from_file(file_path, length, bool_mask):\n",
    "    gc.collect()\n",
    "    warnings.filterwarnings(\"ignore\", category=DeprecationWarning) \n",
    "    \n",
    "    cdf = pycdf.CDF(file_path)\n",
    "    mms_prefix = get_mms_prefix(file_path)\n",
    "\n",
    "    result = {'dist': [], 'counts': [], 'phi': [], 'theta': [], 'E': []}\n",
    "    \n",
    "    energy_table = str(cdf.attrs['Energy_table_name'])\n",
    "    if '12-14' in energy_table:\n",
    "        # Skip solar wind data\n",
    "        return result\n",
    "    \n",
    "    for i in bool_mask.nonzero()[0]:\n",
    "        lossy = bool(cdf[f'{mms_prefix}_dis_compressionloss_brst'][i])\n",
    "        if lossy:\n",
    "            continue # lossy frame-- don't include it\n",
    "        \n",
    "        result['dist'].append(cdf[f'{mms_prefix}_dis_dist_brst'][i, :, :, :])\n",
    "\n",
    "        with warnings.catch_warnings():\n",
    "            warnings.filterwarnings('ignore')\n",
    "            counts = np.square(result['dist'][-1] / cdf[f'{mms_prefix}_dis_disterr_brst'][i, :, :, :])\n",
    "        counts[np.isnan(counts)] = 0\n",
    "        counts = np.rint(counts) # round to nearest int\n",
    "        result['counts'].append(counts)\n",
    "        \n",
    "        result['phi'].append(cdf[f'{mms_prefix}_dis_phi_brst'][i, :])\n",
    "        result['theta'].append(cdf[f'{mms_prefix}_dis_theta_brst'][:])\n",
    "        result['E'].append(cdf[f'{mms_prefix}_dis_energy_brst'][i, :])\n",
    "\n",
    "    cdf.close()\n",
    "\n",
    "    return result\n",
    "\n",
    "\n",
    "def get_mms_prefix(file_path):\n",
    "    for i in range(1, 5):\n",
    "        if f'mms{i}' in file_path:\n",
    "            return f'mms{i}'\n",
    "    return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c876eeda-cfc9-4e22-b9d2-0247f74cb5c6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "tasks = [\n",
    "    ('train', df_train, hyperparams['prep']['sample_sizes']['train']),\n",
    "    ('test', df_test, hyperparams['prep']['sample_sizes']['test'])\n",
    "]\n",
    "\n",
    "for label, df, sample_size in tasks:\n",
    "    for phase in phases:\n",
    "        print(f'{label.ljust(10)} - {phase.ljust(20)}', end='')\n",
    "\n",
    "        output = get_data(sample_size, df, phase)\n",
    "\n",
    "        hdf_path = f'/mnt/efs/dasilva/compression-cfha/data/samples_{label}_n={sample_size}_nosw.hdf'\n",
    "        hdf = h5py.File(hdf_path, 'a')\n",
    "        \n",
    "        try:\n",
    "            group = hdf.create_group(phase)\n",
    "        except ValueError:\n",
    "            del hdf[phase]\n",
    "            group = hdf.create_group(phase)\n",
    "        \n",
    "        for key in output:\n",
    "            group[key] = output[key]\n",
    "        \n",
    "        hdf.close()\n",
    "        print(f'...wrote to {hdf_path}')\n",
    "        \n",
    "        del output\n",
    "        gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d8cb40d-5ae4-45c1-b69e-33d6a8514029",
   "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
}