data_processing.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "import h5py as h5\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "import pickle\n",
    "import sys\n",
    "sys.path.append(os.path.expanduser(\"~/TMPredictor/survival_tm/auton-survival\"))\n",
    "from auton_survival.preprocessing import Scaler\n",
    "import optuna\n",
    "from sklearn.model_selection import ParameterGrid\n",
    "#sys.path.append('/projects/EKOLEMEN/survival_tm/train_models/auton-survival')\n",
    "sys.path.append(os.path.expanduser(\"~/TMPredictor/survival_tm/auton-survival\"))\n",
    "from sklearn.model_selection import ParameterGrid\n",
    "from auton_survival.estimators import SurvivalModel\n",
    "from auton_survival.metrics import survival_regression_metric\n",
    "from auton_survival.models.dsm import DeepSurvivalMachines\n",
    "from sksurv.metrics import concordance_index_ipcw, brier_score, cumulative_dynamic_auc\n",
    "from sklearn.decomposition import PCA"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "FROM TM LABELS AND DATABASE, CREATE X, T, E, SHOTS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['pinj', 'tinj', 'betan_EFITRT2', 'qmin_EFITRT2', 'ech_pwr_total', 'ip', 'bt', 'li_EFITRT2', 'aminor_EFITRT2', 'rmaxis_EFITRT2', 'tribot_EFITRT2', 'tritop_EFITRT2', 'kappa_EFITRT2', 'volume_EFITRT2', 't_ip_flat_sql', 'ip_flat_duration_sql', 'thomson_temp_mtanh_1d', 'cer_temp_csaps_1d', 'thomson_density_mtanh_1d', 'cer_rot_csaps_1d', 'qpsi_EFITRT2', 'pres_EFITRT2'])\n"
     ]
    }
   ],
   "source": [
    "\n",
    "#database = h5.File('/projects/EKOLEMEN/profile_predictor/joe_hiro_models/diiid_data.h5', 'r')\n",
    "tm_data = h5.File('/projects/EKOLEMEN/survival_tm/tm_labels.h5', 'r')\n",
    "\n",
    "signals = ['pinj', 'tinj', 'betan_EFITRT2', 'qmin_EFITRT2', 'ech_pwr_total', 'ip', 'bt', 'li_EFITRT2', 'aminor_EFITRT2', \n",
    "          'rmaxis_EFITRT2', 'tribot_EFITRT2', 'tritop_EFITRT2', 'kappa_EFITRT2', 'volume_EFITRT2']\n",
    "\n",
    "prof_signals = ['thomson_temp_mtanh_1d', 'cer_temp_csaps_1d', 'thomson_density_mtanh_1d', 'cer_rot_csaps_1d', 'qpsi_EFITRT2', 'pres_EFITRT2']\n",
    "\n",
    "with open('data_dict-2.pkl', 'rb') as f:\n",
    "    data = pickle.load(f)\n",
    "# this shot has messed up datastructure\n",
    "del data['192209']\n",
    "shots = list(data.keys())\n",
    "print(data[shots[0]].keys())\n",
    "def check_all_signals(shot_data, prof_signals, signals, ind):\n",
    "    # Check scalars for nan values\n",
    "    for sig in signals:\n",
    "        if np.isnan(shot_data[sig][ind]):\n",
    "            return False\n",
    "    \n",
    "    # Check profiles for nan values\n",
    "    for sig in prof_signals:\n",
    "        if np.sum(np.isnan(shot_data[sig][ind,:])) > 0:\n",
    "            return False\n",
    "       \n",
    "    return True\n",
    "\n",
    "def gather_feature(shot_data, prof_signals, signals, ind):\n",
    "    feature = np.zeros(len(prof_signals)*33+len(signals))\n",
    "    \n",
    "    # Load scalar signals\n",
    "    for i, sig in enumerate(signals):\n",
    "        feature[i] = shot_data[sig][ind]\n",
    "    \n",
    "    # Load profile signals\n",
    "    for i, sig in enumerate(prof_signals):\n",
    "        feature[33*i+len(signals):33*(i+1)+len(signals)] = shot_data[sig][ind,:]\n",
    "    \n",
    "    return feature\n",
    "\n",
    "x = []\n",
    "t = []\n",
    "e = []\n",
    "shots_rt = []\n",
    "times = np.arange(0, 6000, 20)[::-1]\n",
    "# I should check the timesteps are the same for all scalars\n",
    "for shot in shots:\n",
    "    if shot in list(tm_data.keys()):\n",
    "        # First figure out time of TM, or time of end of shot\n",
    "        label = tm_data[shot]['label'][:]\n",
    "        time = tm_data[shot]['time'][:]\n",
    "        tm_ind = np.argmax(label>0)\n",
    "        tm_time = time[tm_ind]\n",
    "        # Find start and stop of flattop\n",
    "        start = np.array(data[shot]['t_ip_flat_sql'])\n",
    "        duration = np.array(data[shot]['ip_flat_duration_sql'])\n",
    "        stop = start+duration\n",
    "        # Get start and stop inds\n",
    "        start_ind = np.argmin(np.abs(time - start))\n",
    "        end_ind = np.argmin(np.abs(time - stop))\n",
    "        \n",
    "        # End shot when TM occurs\n",
    "        if end_ind > tm_ind and tm_ind > 0:\n",
    "            end_ind = tm_ind\n",
    "        \n",
    "        # Go through valid indices and gather dataset\n",
    "        for ind in range(start_ind, end_ind):\n",
    "            # Check that we have all signals at this time\n",
    "            if check_all_signals(data[shot], prof_signals, signals, ind):\n",
    "                feature = gather_feature(data[shot], prof_signals, signals, ind)\n",
    "                x.append(feature)\n",
    "                if tm_time > 0:\n",
    "                    # Time to TM\n",
    "                    t.append(tm_time - time[ind])\n",
    "                    shots_rt.append(shot)\n",
    "                    e.append(1)\n",
    "                else:\n",
    "                    # Time to end of shot\n",
    "                    t.append(time[end_ind] - time[ind])\n",
    "                    shots_rt.append(shot)\n",
    "                    e.append(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 168,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''with open('data/rt_x.pkl', 'wb') as f:\n",
    "    pickle.dump(np.array(x), f)\n",
    "with open('data/rt_t.pkl', 'wb') as f:\n",
    "    pickle.dump(t, f)\n",
    "with open('data/rt_e.pkl', 'wb') as f:\n",
    "    pickle.dump(e, f)\n",
    "with open('data/rt_shots.pkl', 'wb') as f:\n",
    "    pickle.dump(shots_rt, f)'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 188,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_3615927/4178029613.py:12: RuntimeWarning: divide by zero encountered in divide\n",
      "  data[:, 14+33*4:14+33*5] = 1/data[:, 14+33*4:14+33*5]\n"
     ]
    }
   ],
   "source": [
    "data = x\n",
    "mean = np.mean(data, axis=0)\n",
    "std_dev = np.std(data, axis=0)\n",
    "mean[2] = 0\n",
    "std_dev[2] = 1\n",
    "mean[7] = 0\n",
    "std_dev[7] = 1\n",
    "\n",
    "# Define the threshold for filtering\n",
    "threshold = 8*std_dev\n",
    "data[:, 14+33*4:14+33*5] = 1/data[:, 14+33*4:14+33*5]\n",
    "# Identify rows with any values that are more than 10 standard deviations away from the mean for any column\n",
    "rows_with_anomalies = np.any(np.abs(data - mean) > threshold, axis=1)\n",
    "data[:, 14+33*4:14+33*5] = 1/data[:, 14+33*4:14+33*5]\n",
    "# Filter out the rows with anomalies\n",
    "filtered_data = data[~rows_with_anomalies]\n",
    "\n",
    "'''with open('data/rt_filtered_x.pkl', 'wb') as f:\n",
    "    pickle.dump(filtered_data, f)\n",
    "with open('data/rt_filtered_t.pkl', 'wb') as f:\n",
    "    pickle.dump(np.array(t)[~rows_with_anomalies], f)\n",
    "with open('data/rt_filtered_e.pkl', 'wb') as f:\n",
    "    pickle.dump(np.array(e)[~rows_with_anomalies], f)\n",
    "with open('data/rt_filtered_shots.pkl', 'wb') as f:\n",
    "    pickle.dump(np.array(shots_rt)[~rows_with_anomalies], f)'''"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "CREATE MAIN PREPROCESSED DATABASE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "shots = np.load('/projects/EKOLEMEN/survival_tm/shots.npy')\n",
    "tm_shots = np.load('/projects/EKOLEMEN/survival_tm/tm_shots.npy')\n",
    "st_shots = np.load('/projects/EKOLEMEN/survival_tm/st_shots.npy')\n",
    "\n",
    "def load_data(data_type):\n",
    "    with open(f'/projects/EKOLEMEN/survival_tm/formatted_labels/{data_type}.pkl', 'rb') as f:\n",
    "        data = pickle.load(f)\n",
    "    \n",
    "    return data['x'], data['t'], data['e']\n",
    "\n",
    "# Don't want whole dataset right now so 0.5 factor to reduce\n",
    "n = len(shots)\n",
    "\n",
    "tr_size = int(n*0.80)\n",
    "vl_size = int(n*0.10)\n",
    "te_size = int(n*0.10)\n",
    "\n",
    "train_shots = shots[:tr_size]\n",
    "test_shots = shots[-te_size:]\n",
    "valid_shots = shots[tr_size:tr_size+vl_size]\n",
    "\n",
    "x_train, t_train, e_train = load_data('train')\n",
    "x_test,  t_test,  e_test  = load_data('test')\n",
    "x_valid, t_valid, e_valid = load_data('valid')\n",
    "\n",
    "# Get inds for time <600ms\n",
    "# is this needed? Not sure.\n",
    "'''inds = np.where(t_train < 600)[0]\n",
    "\n",
    "x_train = x_train[inds]\n",
    "t_train = t_train[inds]\n",
    "e_train = e_train[inds]'''\n",
    "\n",
    "tm_inds = np.where(e_train == 1)[0]\n",
    "st_inds = np.where(e_train == 0)[0]\n",
    "new_st_inds = np.random.choice(st_inds, size=len(tm_inds), replace=False)\n",
    "combined_inds = np.concatenate((tm_inds, new_st_inds))\n",
    "sorted_combined_inds = np.sort(combined_inds)\n",
    "\n",
    "x_train = x_train[sorted_combined_inds]\n",
    "t_train = t_train[sorted_combined_inds]\n",
    "e_train = e_train[sorted_combined_inds]\n",
    "\n",
    "'''x_train = np.concatenate((x_train[tm_inds], x_train[new_st_inds]), axis=0)\n",
    "t_train = np.concatenate((t_train[tm_inds], t_train[new_st_inds]), axis=0)\n",
    "e_train = np.concatenate((e_train[tm_inds], e_train[new_st_inds]), axis=0)\n",
    "plt.plot(t_train[10000:11000])\n",
    "# Shuffle arrays because currently all 1s followed by all 0s\n",
    "p = np.random.permutation(len(t_train))\n",
    "x_train = x_train[p,:]\n",
    "t_train = t_train[p]\n",
    "e_train = e_train[p]'''\n",
    "\n",
    "x_train_df = pd.DataFrame(x_train)\n",
    "t_train_df = pd.DataFrame(t_train)\n",
    "e_train_df = pd.DataFrame(e_train)\n",
    "\n",
    "x_valid_df = pd.DataFrame(x_valid)\n",
    "t_valid_df = pd.DataFrame(t_valid)\n",
    "e_valid_df = pd.DataFrame(e_valid)\n",
    "\n",
    "x_test_df = pd.DataFrame(x_test)\n",
    "t_test_df = pd.DataFrame(t_test)\n",
    "e_test_df = pd.DataFrame(e_test)\n",
    "\n",
    "outcomes_valid_df = pd.DataFrame({'time': t_valid, 'event': e_valid})\n",
    "\n",
    "outcomes_train_df = pd.DataFrame({'time': t_train, 'event': e_train})\n",
    "outcomes_test_df = pd.DataFrame({'time': t_test, 'event': e_test})\n",
    "#normalize\n",
    "'''scaler = Scaler()\n",
    "transformer = scaler.fit(x_train_df)\n",
    "x_train_df_normed = transformer.transform(x_train_df)\n",
    "x_valid_df_normed = transformer.transform(x_valid_df)\n",
    "x_test_df_normed = transformer.transform(x_test_df)'''\n",
    "'''\n",
    "with open('data/x_train_df1.pkl', 'wb') as f:\n",
    "    pickle.dump(x_train_df, f)\n",
    "with open('data/x_valid_df1.pkl', 'wb') as f:\n",
    "    pickle.dump(x_valid_df, f)\n",
    "with open('data/x_test_df1.pkl', 'wb') as f:\n",
    "    pickle.dump(x_test_df, f)\n",
    "with open('data/outcomes_valid_df1.pkl', 'wb') as f:\n",
    "    pickle.dump(outcomes_valid_df, f)\n",
    "with open('data/outcomes_train_df1.pkl', 'wb') as f:\n",
    "    pickle.dump(outcomes_train_df, f)\n",
    "with open('data/outcomes_test_df1.pkl', 'wb') as f:\n",
    "    pickle.dump(outcomes_test_df, f)'''"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "FIND NORMALIZATION (UNUSED, WE NOW NORMALIZE AFTER PCA)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# q will be normalized by inverting rather than normalizing\n",
    "# TRY NORMALIZING PROFILE BY PROFILE NOT EACH ELEMENT\n",
    "q_start_index = 146\n",
    "train_means = np.mean(x_train[:, :q_start_index], axis=0)\n",
    "train_std = np.std(x_train[:, :q_start_index], axis=0)\n",
    "valid_means = np.mean(x_valid[:, :q_start_index], axis=0)\n",
    "valid_std = np.std(x_valid[:, :q_start_index], axis=0)\n",
    "train_iotas = 1/x_train[:, q_start_index:]\n",
    "valid_iotas = 1/x_valid[:, q_start_index:]\n",
    "x_train_norm = (x_train[:, :q_start_index] - train_means)/train_std\n",
    "x_valid_norm = (x_valid[:, :q_start_index] - valid_means)/valid_std\n",
    "\n",
    "x_train_norm = np.concatenate((x_train_norm, train_iotas), axis=1)\n",
    "x_valid_norm = np.concatenate((x_valid_norm, valid_iotas), axis=1)\n",
    "train_mask = np.all(x_train_norm <= 10, axis=1)\n",
    "valid_mask = np.all(x_valid_norm <= 10, axis=1)\n",
    "# Filter the array to keep only rows where all values are <= 10\n",
    "x_train_norm = pd.DataFrame(x_train_norm[train_mask])\n",
    "x_valid_norm = pd.DataFrame(x_valid_norm[valid_mask])\n",
    "t_train = t_train[train_mask]\n",
    "e_train = e_train[train_mask]\n",
    "t_valid = t_valid[valid_mask]\n",
    "e_valid = e_valid[valid_mask]\n",
    "outcomes_valid_df = pd.DataFrame({'time': t_valid, 'event': e_valid})\n",
    "outcomes_train_df = pd.DataFrame({'time': t_train, 'event': e_train})\n",
    "\n",
    "'''with open('data/x_train_manual_norm.pkl', 'wb') as f:\n",
    "    pickle.dump(x_train_norm, f)\n",
    "with open('data/x_valid_manual_norm.pkl', 'wb') as f:\n",
    "    pickle.dump(x_valid_norm, f)\n",
    "\n",
    "with open('data/outcomes_valid_manual_norm.pkl', 'wb') as f:\n",
    "    pickle.dump(outcomes_valid_df, f)\n",
    "with open('data/outcomes_train_manual_norm.pkl', 'wb') as f:\n",
    "    pickle.dump(outcomes_train_df, f)'''"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "CREATE PCA DATABASE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 190,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''with open('data/x_train_df1.pkl', 'rb') as f:\n",
    "    x_train_df = pickle.load(f)\n",
    "with open('data/x_valid_df1.pkl', 'rb') as f:\n",
    "    x_valid_df = pickle.load(f)\n",
    "with open('data/outcomes_train_df1.pkl', 'rb') as f:\n",
    "    outcomes_train_df = pickle.load(f)\n",
    "with open('data/outcomes_valid_df1.pkl', 'rb') as f:\n",
    "    outcomes_valid_df = pickle.load(f)'''\n",
    "\n",
    "with open('data/rt_filtered_x.pkl', 'rb') as f:\n",
    "    x_train_df = pickle.load(f)\n",
    "\n",
    "# separate each normalized profile\n",
    "scalars = np.array(x_train_df)[:, 0:14]\n",
    "te = np.array(x_train_df)[:, 14:33+14]\n",
    "ti = np.array(x_train_df)[:, 33+14:66+14]\n",
    "ne = np.array(x_train_df)[:, 66+14:99+14]\n",
    "rot = np.array(x_train_df)[:, 99+14:132+14]\n",
    "qpsi = np.array(x_train_df)[:, 132+14:165+14]\n",
    "qpsi = 1/qpsi\n",
    "qpsi = np.where(qpsi == np.inf, 1, qpsi)\n",
    "pres = np.array(x_train_df)[:, 165+14:198+14]\n",
    "\n",
    "# upsample to PCS dimensions\n",
    "x_old = np.linspace(0, 1, 33)\n",
    "x_long = np.linspace(0, 1, 100)\n",
    "x_for_q = np.linspace(0, 1, 65)\n",
    "te = np.array([np.interp(x_long, x_old, te[i]) for i in range(len(te))])\n",
    "ti = np.array([np.interp(x_long, x_old, ti[i]) for i in range(len(ti))])\n",
    "ne = np.array([np.interp(x_long, x_old, ne[i]) for i in range(len(ne))])\n",
    "rot = np.array([np.interp(x_long, x_old, rot[i]) for i in range(len(rot))])\n",
    "qpsi = np.array([np.interp(x_for_q, x_old, qpsi[i]) for i in range(len(qpsi))])\n",
    "pres = np.array([np.interp(x_for_q, x_old, pres[i]) for i in range(len(pres))])\n",
    "\n",
    "pca_components = []\n",
    "# Fit PCA on the data\n",
    "pca = PCA(n_components=4)\n",
    "pca.fit(te)\n",
    "te_pca = pca.transform(te)\n",
    "pca_components.append(pca.components_)\n",
    "pca_components.append(pca.mean_)\n",
    "\n",
    "pca = PCA(n_components=4)\n",
    "pca.fit(ti)\n",
    "ti_pca = pca.transform(ti)\n",
    "pca_components.append(pca.components_)\n",
    "pca_components.append(pca.mean_)\n",
    "\n",
    "pca = PCA(n_components=4)\n",
    "pca.fit(ne)\n",
    "ne_pca = pca.transform(ne)\n",
    "pca_components.append(pca.components_)\n",
    "pca_components.append(pca.mean_)\n",
    "\n",
    "pca = PCA(n_components=4)\n",
    "pca.fit(rot)\n",
    "rot_pca = pca.transform(rot)\n",
    "pca_components.append(pca.components_)\n",
    "pca_components.append(pca.mean_)\n",
    "\n",
    "pca = PCA(n_components=4)\n",
    "pca.fit(qpsi)\n",
    "qpsi_pca = pca.transform(qpsi)\n",
    "pca_components.append(pca.components_)\n",
    "pca_components.append(pca.mean_)\n",
    "\n",
    "pca = PCA(n_components=4)\n",
    "pca.fit(pres)\n",
    "pres_pca = pca.transform(pres)\n",
    "pca_components.append(pca.components_)\n",
    "pca_components.append(pca.mean_)\n",
    "# normalize PCA components\n",
    "\n",
    "scalars_mean = np.mean(scalars, axis=0)\n",
    "scalars_std = np.std(scalars, axis=0)\n",
    "scalars_normed = (scalars - scalars_mean)/scalars_std\n",
    "\n",
    "te_pca_mean = np.mean(te_pca)\n",
    "te_pca_std = np.std(te_pca)\n",
    "te_pca_normed = (te_pca - te_pca_mean)/te_pca_std\n",
    "\n",
    "ti_pca_mean = np.mean(ti_pca)\n",
    "ti_pca_std = np.std(ti_pca)\n",
    "ti_pca_normed = (ti_pca - ti_pca_mean)/ti_pca_std\n",
    "\n",
    "ne_pca_mean = np.mean(ne_pca)\n",
    "ne_pca_std = np.std(ne_pca)\n",
    "ne_pca_normed = (ne_pca - ne_pca_mean)/ne_pca_std\n",
    "\n",
    "rot_pca_mean = np.mean(rot_pca)\n",
    "rot_pca_std = np.std(rot_pca)\n",
    "rot_pca_normed = (rot_pca - rot_pca_mean)/rot_pca_std\n",
    "\n",
    "qpsi_pca_mean = np.mean(qpsi_pca)\n",
    "qpsi_pca_std = np.std(qpsi_pca)\n",
    "qpsi_pca_normed = (qpsi_pca - qpsi_pca_mean)/qpsi_pca_std\n",
    "\n",
    "pres_pca_mean = np.mean(pres_pca)\n",
    "pres_pca_std = np.std(pres_pca)\n",
    "pres_pca_normed = (pres_pca - pres_pca_mean)/pres_pca_std\n",
    "\n",
    "normalizations = [(scalars_mean, scalars_std), (te_pca_mean, te_pca_std), (ti_pca_mean, ti_pca_std), (ne_pca_mean, ne_pca_std), (rot_pca_mean, rot_pca_std), (qpsi_pca_mean, qpsi_pca_std), (pres_pca_mean, pres_pca_std)]\n",
    "\n",
    "x_train_pca = pd.DataFrame(np.concatenate((scalars_normed, te_pca_normed, ti_pca_normed, ne_pca_normed, rot_pca_normed, qpsi_pca_normed, pres_pca_normed), axis=1))\n",
    "'''\n",
    "with open('data/rt_x_pca.pkl', 'wb') as f:\n",
    "    pickle.dump(x_train_pca, f)\n",
    "\n",
    "with open('data/rt_normalizations.pkl', 'wb') as f:\n",
    "    pickle.dump(normalizations, f)\n",
    "\n",
    "with open('data/rt_pca_components.pkl', 'wb') as f:\n",
    "    pickle.dump(pca_components, f)'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_peaks_in_data(data):\n",
    "    peaks = []\n",
    "    for i in range(1, len(data) - 1):\n",
    "        if data[i-1] < data[i] > data[i+1]:\n",
    "            peaks.append(i)\n",
    "    return peaks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "CREATE MEMORY DATABASE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"with open('data/x_valid_memory_normed.pkl', 'wb') as f:\\n    pickle.dump(x_valid_memory_normed, f)\\nwith open('data/x_test_memory_normed.pkl', 'wb') as f:\\n    pickle.dump(x_test_memory_normed, f)\\nwith open('data/x_train_memory_normed.pkl', 'wb') as f:\\n    pickle.dump(x_train_memory_normed, f)\\nwith open('data/outcomes_valid_memory.pkl', 'wb') as f:\\n    pickle.dump(outcomes_valid_memory, f)\\nwith open('data/outcomes_train_memory.pkl', 'wb') as f:\\n    pickle.dump(outcomes_train_memory, f)\\nwith open('data/outcomes_test_memory.pkl', 'wb') as f:\\n    pickle.dump(outcomes_test_memory, f)\""
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# we give each timestep all the info from a few timesteps in the past. \n",
    "\n",
    "x_train_np = np.array(x_train)\n",
    "t_train_np = np.array(t_train)[4:]\n",
    "e_train_np = np.array(e_train)[4:]\n",
    "\n",
    "x_valid_np = np.array(x_valid)\n",
    "t_valid_np = np.array(t_valid)[4:]\n",
    "e_valid_np = np.array(e_valid)[4:]\n",
    "\n",
    "x_test_np = np.array(x_test)\n",
    "t_test_np = np.array(t_test)[4:]\n",
    "e_test_np = np.array(e_test)[4:]\n",
    "\n",
    "outcomes_valid_memory = pd.DataFrame({'time': t_valid_np, 'event': e_valid_np})\n",
    "outcomes_test_memory = pd.DataFrame({'time': t_test_np, 'event': e_test_np})\n",
    "outcomes_train_memory = pd.DataFrame({'time': t_train_np, 'event': e_train_np})\n",
    "\n",
    "# make the state of each timestep contain 5 timesteps of data\n",
    "x_valid_memory = np.array([x_valid_np[i-4:i+1].flatten() for i in range(4, len(x_valid_np))])\n",
    "x_test_memory = np.array([x_test_np[i-4:i+1].flatten() for i in range(4, len(x_test_np))])\n",
    "x_train_memory = np.array([x_train_np[i-4:i+1].flatten() for i in range(4, len(x_train_np))])\n",
    "\n",
    "# remove the first 4 timesteps of each shot from the dataset. This is so that data from the previous shot isn't saved in the current shot\n",
    "train_peaks = np.array(find_peaks_in_data(np.array(outcomes_train_memory['time'])))\n",
    "valid_peaks = np.array(find_peaks_in_data(np.array(outcomes_valid_memory['time'])))\n",
    "test_peaks = np.array(find_peaks_in_data(np.array(outcomes_test_memory['time'])))\n",
    "new_train_peaks = np.sort(np.concatenate((train_peaks, train_peaks+1, train_peaks+2, train_peaks+3)))\n",
    "new_valid_peaks = np.sort(np.concatenate((valid_peaks, valid_peaks+1, valid_peaks+2, valid_peaks+3)))\n",
    "new_test_peaks = np.sort(np.concatenate((test_peaks, test_peaks+1, test_peaks+2, test_peaks+3)))\n",
    "x_valid_memory = pd.DataFrame(np.delete(x_valid_memory, new_valid_peaks, axis=0))\n",
    "x_test_memory = pd.DataFrame(np.delete(x_test_memory, new_test_peaks, axis=0))\n",
    "x_train_memory = pd.DataFrame(np.delete(x_train_memory, new_train_peaks, axis=0))\n",
    "\n",
    "# normalize\n",
    "scaler = Scaler()\n",
    "transformer = scaler.fit(x_train_memory)\n",
    "x_train_memory_normed = transformer.transform(x_train_memory)\n",
    "x_valid_memory_normed = transformer.transform(x_valid_memory)\n",
    "x_test_memory_normed = transformer.transform(x_test_memory)\n",
    "\n",
    "'''with open('data/x_valid_memory_normed.pkl', 'wb') as f:\n",
    "    pickle.dump(x_valid_memory_normed, f)\n",
    "with open('data/x_test_memory_normed.pkl', 'wb') as f:\n",
    "    pickle.dump(x_test_memory_normed, f)\n",
    "with open('data/x_train_memory_normed.pkl', 'wb') as f:\n",
    "    pickle.dump(x_train_memory_normed, f)\n",
    "with open('data/outcomes_valid_memory.pkl', 'wb') as f:\n",
    "    pickle.dump(outcomes_valid_memory, f)\n",
    "with open('data/outcomes_train_memory.pkl', 'wb') as f:\n",
    "    pickle.dump(outcomes_train_memory, f)\n",
    "with open('data/outcomes_test_memory.pkl', 'wb') as f:\n",
    "    pickle.dump(outcomes_test_memory, f)'''"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "CREATE FUTURE ACTUATORS DATABASE. SAME LOGIC AS ABOVE, BUT ONLY SAVING SCALARS DATA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"with open('data/x_valid_future_normed.pkl', 'wb') as f:\\n    pickle.dump(x_valid_future_normed, f)\\nwith open('data/x_test_future_normed.pkl', 'wb') as f:\\n    pickle.dump(x_test_future_normed, f)\\nwith open('data/x_train_future_normed.pkl', 'wb') as f:\\n    pickle.dump(x_train_future_normed, f)\\nwith open('data/outcomes_valid_future.pkl', 'wb') as f:\\n    pickle.dump(outcomes_valid_future, f)\\nwith open('data/outcomes_train_future.pkl', 'wb') as f:\\n    pickle.dump(outcomes_train_future, f)\\nwith open('data/outcomes_test_future.pkl', 'wb') as f:\\n    pickle.dump(outcomes_test_future, f)\""
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "signals = ['pinj', 'tinj', 'betan_EFIT01', 'qmin_EFIT01', 'ech_pwr_total', 'ip', 'bt', 'li_EFIT01', 'aminor_EFIT01', \n",
    "          'rmaxis_EFIT01', 'tribot_EFIT01', 'tritop_EFIT01', 'kappa_EFIT01', 'volume_EFIT01']\n",
    "\n",
    "# we give each timestep all the info from a few timesteps in the past. \n",
    "\n",
    "x_train_np = np.array(x_train)\n",
    "t_train_np = np.array(t_train)[:-4]\n",
    "e_train_np = np.array(e_train)[:-4]\n",
    "\n",
    "x_valid_np = np.array(x_valid)\n",
    "t_valid_np = np.array(t_valid)[:-4]\n",
    "e_valid_np = np.array(e_valid)[:-4]\n",
    "\n",
    "x_test_np = np.array(x_test)\n",
    "t_test_np = np.array(t_test)[:-4]\n",
    "e_test_np = np.array(e_test)[:-4]\n",
    "\n",
    "outcomes_valid_future = pd.DataFrame({'time': t_valid_np, 'event': e_valid_np})\n",
    "outcomes_test_future = pd.DataFrame({'time': t_test_np, 'event': e_test_np})\n",
    "outcomes_train_future = pd.DataFrame({'time': t_train_np, 'event': e_train_np})\n",
    "\n",
    "# we keep the present and 4 in the future\n",
    "x_valid_future = pd.DataFrame([np.concatenate((x_valid_np[i:i+5, 0:len(signals)].flatten(), x_valid_np[i, len(signals):])) for i in range(len(x_valid_np)-4)])\n",
    "x_test_future = pd.DataFrame([np.concatenate((x_test_np[i:i+5, 0:len(signals)].flatten(), x_test_np[i, len(signals):])) for i in range(len(x_test_np)-4)])\n",
    "x_train_future = pd.DataFrame([np.concatenate((x_train_np[i:i+5, 0:len(signals)].flatten(), x_train_np[i, len(signals):])) for i in range(len(x_train_np)-4)])\n",
    "\n",
    "train_peaks = np.array(find_peaks_in_data(np.array(outcomes_train_future['time'])))\n",
    "valid_peaks = np.array(find_peaks_in_data(np.array(outcomes_valid_future['time'])))\n",
    "test_peaks = np.array(find_peaks_in_data(np.array(outcomes_test_future['time'])))\n",
    "\n",
    "new_train_peaks = np.sort(np.concatenate((train_peaks-1, train_peaks-2, train_peaks-3, train_peaks-4)))\n",
    "new_valid_peaks = np.sort(np.concatenate((valid_peaks-1, valid_peaks-2, valid_peaks-3, valid_peaks-4)))\n",
    "new_test_peaks = np.sort(np.concatenate((test_peaks-1, test_peaks-2, test_peaks-3, test_peaks-4)))\n",
    "\n",
    "x_valid_future = pd.DataFrame(np.delete(x_valid_future, new_valid_peaks, axis=0))\n",
    "x_test_future = pd.DataFrame(np.delete(x_test_future, new_test_peaks, axis=0))\n",
    "x_train_future = pd.DataFrame(np.delete(x_train_future, new_train_peaks, axis=0))\n",
    "\n",
    "scaler = Scaler()\n",
    "transformer = scaler.fit(x_train_future)\n",
    "x_train_future_normed = transformer.transform(x_train_future)\n",
    "x_valid_future_normed = transformer.transform(x_valid_future)\n",
    "x_test_future_normed = transformer.transform(x_test_future)\n",
    "\n",
    "'''with open('data/x_valid_future_normed.pkl', 'wb') as f:\n",
    "    pickle.dump(x_valid_future_normed, f)\n",
    "with open('data/x_test_future_normed.pkl', 'wb') as f:\n",
    "    pickle.dump(x_test_future_normed, f)\n",
    "with open('data/x_train_future_normed.pkl', 'wb') as f:\n",
    "    pickle.dump(x_train_future_normed, f)\n",
    "with open('data/outcomes_valid_future.pkl', 'wb') as f:\n",
    "    pickle.dump(outcomes_valid_future, f)\n",
    "with open('data/outcomes_train_future.pkl', 'wb') as f:\n",
    "    pickle.dump(outcomes_train_future, f)\n",
    "with open('data/outcomes_test_future.pkl', 'wb') as f:\n",
    "    pickle.dump(outcomes_test_future, f)'''"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.11.7"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}