06_LSTM_data_prepare.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import os\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir = '../noaa-runtime/data/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "dst = pd.read_csv(data_dir +\"dst_labels.csv\")\n",
    "dst.timedelta = pd.to_timedelta(dst.timedelta)\n",
    "dst.set_index([\"period\", \"timedelta\"], inplace=True)\n",
    "\n",
    "sunspots = pd.read_csv(data_dir + \"sunspots.csv\")\n",
    "sunspots.timedelta = pd.to_timedelta(sunspots.timedelta)\n",
    "sunspots.set_index([\"period\", \"timedelta\"], inplace=True)\n",
    "\n",
    "solar_wind = pd.read_csv(data_dir + \"solar_wind.csv\")\n",
    "solar_wind.timedelta = pd.to_timedelta(solar_wind.timedelta)\n",
    "solar_wind.set_index([\"period\", \"timedelta\"], inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from numpy.random import seed\n",
    "from tensorflow.random import set_seed\n",
    "\n",
    "seed(2020)\n",
    "set_seed(2021)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Preprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def aggregate_data(feature_df, aggs=[\"mean\", \"std\"], agg_time=\"30min\"):\n",
    "    \"\"\"Aggregates features to the floor of each hour using mean and standard deviation.\n",
    "    e.g. All values from \"11:00:00\" to \"11:30:00\" will be aggregated to \"11:00:00\".\n",
    "    \"\"\"\n",
    "    # group by the floor of each hour use timedelta index\n",
    "    agged = feature_df.groupby(\n",
    "        [\"period\", feature_df.index.get_level_values(1).floor(agg_time)]\n",
    "    ).agg(aggs)\n",
    "    # flatten hierachical column index\n",
    "    agged.columns = [\"_\".join(x) for x in agged.columns]\n",
    "    return agged\n",
    "\n",
    "# interpolate the data\n",
    "def interpolate_data(df):\n",
    "    \"\"\"\n",
    "    # interpolate the data\n",
    "    - `smoothed_ssn`: forward fill\n",
    "    - `solar_wind`: interpolation\n",
    "    \"\"\"\n",
    "    df.smoothed_ssn = df.smoothed_ssn.fillna(method=\"ffill\")\n",
    "    df = df.interpolate()\n",
    "    return df\n",
    "    \n",
    "def preprocess_data(solar_wind, sunspots, scaler=None, subset=None):\n",
    "    \"\"\"\n",
    "    Preprocessing steps:\n",
    "        - Subset the data\n",
    "        - Aggregate hourly\n",
    "        - Join solar wind and sunspot data\n",
    "        - Scale using standard scaler\n",
    "        - Impute missing values\n",
    "    \"\"\"\n",
    "    # select features we want to use\n",
    "    if subset:\n",
    "        solar_wind = solar_wind[subset]\n",
    "    \n",
    "    # aggregate solar wind data and join with sunspots\n",
    "    data_agg = aggregate_data(solar_wind).join(sunspots)\n",
    "\n",
    "    # scaling the data\n",
    "    if scaler is None:\n",
    "        scaler = StandardScaler()\n",
    "        scaler.fit(data_agg)\n",
    "    data_normalized = pd.DataFrame( scaler.transform(data_agg),\n",
    "                                   index=data_agg.index, columns=data_agg.columns)\n",
    "    # Interpolate the missing data\n",
    "    data_interpolated = interpolate_data(data_normalized)\n",
    "    \n",
    "    # Return the scaler object as well to use later during prediction\n",
    "    return data_interpolated, scaler\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SOLAR_WIND_FEATURES = [ \"bt\", \"temperature\", \"bx_gse\", \"by_gse\",\n",
    "                       \"bz_gse\", \"speed\", \"density\"]\n",
    "\n",
    "XCOLS = ( [col + \"_mean\" for col in SOLAR_WIND_FEATURES]\n",
    "         + [col + \"_std\" for col in SOLAR_WIND_FEATURES]\n",
    "         + [\"smoothed_ssn\"] )\n",
    "\n",
    "df_features, scaler = preprocess_data(solar_wind, sunspots, subset=SOLAR_WIND_FEATURES)\n",
    "print(df_features.shape)\n",
    "df_features.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check to make sure missing values are filled\n",
    "assert (df_features.isna().sum() == 0).all()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "YCOLS = [\"t0\", \"t1\"]\n",
    "\n",
    "def process_labels(dst):\n",
    "    y = dst.copy()\n",
    "    y[\"t1\"] = y.groupby(\"period\").dst.shift(-1)\n",
    "    y.columns = YCOLS\n",
    "    return y\n",
    "\n",
    "df_labels = process_labels(dst)\n",
    "df_labels.head(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = df_labels.join(df_features)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_train_test_val(df, test_per_period, val_per_period):\n",
    "    \"\"\"Splits data across periods into train, test, and validation\"\"\"\n",
    "    # assign the last `test_per_period` rows from each period to test\n",
    "    test = df.groupby(\"period\").tail(test_per_period)\n",
    "    interim = df[~df.index.isin(test.index)]\n",
    "    # assign the last `val_per_period` from the remaining rows to validation\n",
    "    val = interim.groupby(\"period\").tail(val_per_period)\n",
    "    # the remaining rows are assigned to train\n",
    "    train = interim[~interim.index.isin(val.index)]\n",
    "    return train, test, val\n",
    "\n",
    "df_train, df_test, df_val = get_train_test_val(df, test_per_period=6_000, val_per_period=3_000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plots():\n",
    "    \n",
    "    ind = [0, 1, 2]\n",
    "    names = [\"train_a\", \"train_b\", \"train_c\"]\n",
    "    width = 0.75\n",
    "    train_cnts = [len(df) for _, df in df_train.groupby(\"period\")]\n",
    "    val_cnts = [len(df) for _, df in df_val.groupby(\"period\")]\n",
    "    test_cnts = [len(df) for _, df in df_test.groupby(\"period\")]\n",
    "\n",
    "    print (train_cnts, val_cnts, test_cnts)\n",
    "    \n",
    "    p1 = plt.barh(ind, train_cnts, width)\n",
    "    p2 = plt.barh(ind, val_cnts, width, left=train_cnts)\n",
    "    p3 = plt.barh(ind, test_cnts, width, left=np.add(val_cnts, train_cnts).tolist())\n",
    "\n",
    "    plt.yticks(ind, names)\n",
    "    plt.ylabel(\"Period\")\n",
    "    plt.xlabel(\"Hourly Timesteps\")\n",
    "    plt.title(\"Train/Validation/Test Splits over Periods\", fontsize=16)\n",
    "    plt.legend([\"Train\", \"Validation\", \"Test\"])\n",
    "    \n",
    "plots()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras.models import Sequential\n",
    "from tensorflow.keras.layers import Dense, LSTM, Dropout\n",
    "from tensorflow.keras import preprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#df[df.index=='train_a'].shape\n",
    "df[df.index.get_level_values(0)=='train_a'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for _, dff in df.groupby(\"period\"):\n",
    "    print (dff[XCOLS].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_config = { \"timesteps\": 32, \"batch_size\": 32}\n",
    "\n",
    "def timeseries_dataset_from_df(df, batch_size):\n",
    "    dataset = None\n",
    "    timesteps = data_config[\"timesteps\"]\n",
    "\n",
    "    # iterate through periods\n",
    "    for _, period_df in df.groupby(\"period\"):\n",
    "        # realign features and labels so that first sequence of 32 is aligned with the 33rd target\n",
    "        inputs = period_df[XCOLS][:-timesteps]\n",
    "        outputs = period_df[YCOLS][timesteps:]\n",
    "\n",
    "        period_ds = preprocessing.timeseries_dataset_from_array(\n",
    "            inputs,\n",
    "            outputs,\n",
    "            timesteps,\n",
    "            batch_size=batch_size,\n",
    "        )\n",
    "\n",
    "        if dataset is None:\n",
    "            dataset = period_ds\n",
    "        else:\n",
    "            dataset = dataset.concatenate(period_ds)\n",
    "\n",
    "    return dataset\n",
    "\n",
    "\n",
    "train_ds = timeseries_dataset_from_df(df_train, data_config[\"batch_size\"])\n",
    "val_ds = timeseries_dataset_from_df(df_val, data_config[\"batch_size\"])\n",
    "\n",
    "print(f\"Number of train batches: {len(train_ds)}\")\n",
    "print(f\"Number of val batches: {len(val_ds)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define our model\n",
    "model_config = {\"n_epochs\": 50, \"n_neurons\": 32, \"dropout\": 32, \"stateful\": False}\n",
    "\n",
    "def make_model1():\n",
    "    model = Sequential()\n",
    "    model.add(\n",
    "        LSTM(\n",
    "            model_config[\"n_neurons\"],\n",
    "            # usually set to (`batch_size`, `sequence_length`, `n_features`)\n",
    "            # setting the batch size to None allows for variable length batches\n",
    "            batch_input_shape=(None, data_config[\"timesteps\"], len(XCOLS)),\n",
    "            stateful=model_config[\"stateful\"],\n",
    "            dropout=model_config[\"dropout\"],\n",
    "        )\n",
    "    )\n",
    "\n",
    "    model.add(Dense(64))\n",
    "    model.add(Dense(32))\n",
    "    \n",
    "    model.add(Dense(len(YCOLS)))\n",
    "    model.compile( loss=\"mean_squared_error\", optimizer=\"adam\")\n",
    "    return model\n",
    "\n",
    "model = make_model1()\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "history = model.fit(\n",
    "    train_ds,\n",
    "    batch_size=data_config[\"batch_size\"],\n",
    "    epochs=model_config[\"n_epochs\"],\n",
    "    verbose=1,\n",
    "    shuffle=False,\n",
    "    validation_data=val_ds,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for name, values in history.history.items():\n",
    "    plt.plot(values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_config = {\"n_epochs\": 10, \"n_neurons\": 64, \"dropout\": 0.4, \"stateful\": False}\n",
    "\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print ('n_epochs=',model_config[\"n_epochs\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "history = model.fit(\n",
    "    train_ds,\n",
    "    batch_size=data_config[\"batch_size\"],\n",
    "    epochs=model_config[\"n_epochs\"],\n",
    "    verbose=1,\n",
    "    shuffle=False,\n",
    "    validation_data=val_ds,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#for name, values in history.history.items():\n",
    "#    plt.plot(values)\n",
    "#loss=\n",
    "loss=history.history['loss']\n",
    "val_loss=history.history['val_loss']\n",
    "plt.plot(range(1, len(loss)+1), loss, label='Loss' )\n",
    "plt.plot(range(1, len(loss)+1), val_loss, label='Val Loss' )\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# define our model\n",
    "model_config = {\"n_epochs\": 50, \"n_neurons\": 32, \"dropout\": 0.4, \"stateful\": False}\n",
    "\n",
    "def make_model2():\n",
    "    model = Sequential()\n",
    "    model.add(\n",
    "        LSTM(\n",
    "            model_config[\"n_neurons\"],\n",
    "            # usually set to (`batch_size`, `sequence_length`, `n_features`)\n",
    "            # setting the batch size to None allows for variable length batches\n",
    "            batch_input_shape=(None, data_config[\"timesteps\"], len(XCOLS)),\n",
    "            stateful=model_config[\"stateful\"],\n",
    "            dropout=model_config[\"dropout\"],\n",
    "        )\n",
    "    )\n",
    "    model.add(Dense(32))\n",
    "    model.add(Dense(32))\n",
    "    model.add(Dense(32))\n",
    "    model.add(Dense(len(YCOLS)))\n",
    "    model.compile(\n",
    "        loss=\"mean_squared_error\",\n",
    "        optimizer=\"adam\",\n",
    "    )\n",
    "    return model\n",
    "\n",
    "model = make_model2()\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# First Run\n",
    "history = model.fit(\n",
    "    train_ds,\n",
    "    batch_size=data_config[\"batch_size\"],\n",
    "    epochs=model_config[\"n_epochs\"],\n",
    "    verbose=1,\n",
    "    shuffle=False,\n",
    "    validation_data=val_ds,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#for name, values in history.history.items():\n",
    "#    plt.plot(values)\n",
    "loss=history.history['loss']\n",
    "val_loss=history.history['val_loss']\n",
    "plt.plot(range(1, len(loss)+1), loss, label='Loss' )\n",
    "plt.plot(range(1, len(loss)+1), val_loss, label='Val Loss' )\n",
    "[plt.axvline(x=i, color='k', lw=.5, ls='--') for i in range(2, 10, 2)]\n",
    "plt.legend()\n",
    "plt.show() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Rerun for a smaller epoch\n",
    "model_config = {\"n_epochs\": 8, \"n_neurons\": 64, \"dropout\": 0.4, \"stateful\": False}\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# First Run\n",
    "history = model.fit(\n",
    "    train_ds,\n",
    "    batch_size=data_config[\"batch_size\"],\n",
    "    epochs=model_config[\"n_epochs\"],\n",
    "    verbose=1,\n",
    "    shuffle=False,\n",
    "    validation_data=val_ds,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss=history.history['loss']\n",
    "val_loss=history.history['val_loss']\n",
    "plt.plot(range(1, len(loss)+1), loss, label='Loss' )\n",
    "plt.plot(range(1, len(loss)+1), val_loss, label='Val Loss' )\n",
    "[plt.axvline(x=i, color='k', lw=.5, ls='--') for i in range(2, 8, 2)]\n",
    "plt.legend()\n",
    "plt.show() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_ds = timeseries_dataset_from_df(test, data_config[\"batch_size\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_ds = timeseries_dataset_from_df(test, data_config[\"batch_size\"])\n",
    "mse = model.evaluate(test_ds)\n",
    "print(f\"Test RMSE: {mse**.5:.2f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import pickle\n",
    "\n",
    "model.save(\"model\")\n",
    "\n",
    "with open(\"scaler.pck\", \"wb\") as f:\n",
    "    pickle.dump(scaler, f)\n",
    "\n",
    "data_config[\"solar_wind_subset\"] = SOLAR_WIND_FEATURES\n",
    "print(data_config)\n",
    "with open(\"config.json\", \"w\") as f:\n",
    "    json.dump(data_config, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# define our model\n",
    "model_config = {\"n_epochs\": 50, \"n_neurons\": 16, \"dropout\": 0.25, \"stateful\": False}\n",
    "\n",
    "def make_model3():\n",
    "    model = Sequential()\n",
    "    model.add(\n",
    "        LSTM(\n",
    "            model_config[\"n_neurons\"],\n",
    "            # usually set to (`batch_size`, `sequence_length`, `n_features`)\n",
    "            # setting the batch size to None allows for variable length batches\n",
    "            batch_input_shape=(None, data_config[\"timesteps\"], len(XCOLS)),\n",
    "            stateful=model_config[\"stateful\"],\n",
    "            dropout=model_config[\"dropout\"],\n",
    "        )\n",
    "    )\n",
    "    model.add(Dense(16))\n",
    "    model.add(Dropout(model_config[\"dropout\"]))\n",
    "    \n",
    "    model.add(Dense(len(YCOLS)))\n",
    "    model.compile(\n",
    "        loss=\"mean_squared_error\",\n",
    "        optimizer=\"adam\",\n",
    "    )\n",
    "    return model\n",
    "\n",
    "model = make_model3()\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# First Run\n",
    "history = model.fit(\n",
    "    train_ds,\n",
    "    batch_size=data_config[\"batch_size\"],\n",
    "    epochs=model_config[\"n_epochs\"],\n",
    "    verbose=1,\n",
    "    shuffle=False,\n",
    "    validation_data=val_ds,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss=history.history['loss']\n",
    "val_loss=history.history['val_loss']\n",
    "plt.plot(range(1, len(loss)+1), loss, label='Loss' )\n",
    "plt.plot(range(1, len(loss)+1), val_loss, label='Val Loss' )\n",
    "[plt.axvline(x=i, color='k', lw=.5, ls='--') for i in range(2, 18, 2)]\n",
    "plt.legend()\n",
    "plt.show() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Rerun for a smaller epoch\n",
    "model_config = {\"n_epochs\": 17, \"n_neurons\": 16, \"dropout\": 0.25, \"stateful\": False}\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# First Run\n",
    "history = model.fit(\n",
    "    train_ds,\n",
    "    batch_size=data_config[\"batch_size\"],\n",
    "    epochs=model_config[\"n_epochs\"],\n",
    "    verbose=1,\n",
    "    shuffle=False,\n",
    "    validation_data=val_ds,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "loss=history.history['loss']\n",
    "val_loss=history.history['val_loss']\n",
    "plt.plot(range(1, len(loss)+1), loss, label='Loss' )\n",
    "plt.plot(range(1, len(loss)+1), val_loss, label='Val Loss' )\n",
    "[plt.axvline(x=i, color='k', lw=.5, ls='--') for i in range(2, 18, 2)]\n",
    "plt.legend()\n",
    "plt.show() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_ds = timeseries_dataset_from_df(test, data_config[\"batch_size\"])\n",
    "mse = model.evaluate(test_ds)\n",
    "print(f\"Test RMSE: {mse**.5:.2f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import pickle\n",
    "\n",
    "model.save(\"model\")\n",
    "\n",
    "with open(\"scaler.pck\", \"wb\") as f:\n",
    "    pickle.dump(scaler, f)\n",
    "\n",
    "data_config[\"solar_wind_subset\"] = SOLAR_WIND_FEATURES\n",
    "print(data_config)\n",
    "with open(\"config.json\", \"w\") as f:\n",
    "    json.dump(data_config, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import pickle\n",
    "from typing import Tuple\n",
    "from tensorflow import keras\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "# Load in serialized model, config, and scaler\n",
    "model = keras.models.load_model(\"model\")\n",
    "\n",
    "with open(\"config.json\", \"r\") as f:\n",
    "    CONFIG = json.load(f)\n",
    "\n",
    "with open(\"scaler.pck\", \"rb\") as f:\n",
    "    scaler = pickle.load(f)\n",
    "\n",
    "# Set global variables    \n",
    "TIMESTEPS = CONFIG[\"timesteps\"]\n",
    "SOLAR_WIND_FEATURES = [\n",
    "    \"bt\",\n",
    "    \"temperature\",\n",
    "    \"bx_gse\",\n",
    "    \"by_gse\",\n",
    "    \"bz_gse\",\n",
    "    \"speed\",\n",
    "    \"density\",\n",
    "]\n",
    "XCOLS = (\n",
    "    [col + \"_mean\" for col in SOLAR_WIND_FEATURES]\n",
    "    + [col + \"_std\" for col in SOLAR_WIND_FEATURES]\n",
    "    + [\"smoothed_ssn\"]\n",
    ")\n",
    "\n",
    "\n",
    "# Define functions for preprocessing\n",
    "def impute_features(feature_df):\n",
    "    \"\"\"Imputes data using the following methods:\n",
    "    - `smoothed_ssn`: forward fill\n",
    "    - `solar_wind`: interpolation\n",
    "    \"\"\"\n",
    "    # forward fill sunspot data for the rest of the month\n",
    "    feature_df.smoothed_ssn = feature_df.smoothed_ssn.fillna(method=\"ffill\")\n",
    "    # interpolate between missing solar wind values\n",
    "    feature_df = feature_df.interpolate()\n",
    "    return feature_df\n",
    "\n",
    "\n",
    "def aggregate_hourly(feature_df, aggs=[\"mean\", \"std\"]):\n",
    "    \"\"\"Aggregates features to the floor of each hour using mean and standard deviation.\n",
    "    e.g. All values from \"11:00:00\" to \"11:59:00\" will be aggregated to \"11:00:00\".\n",
    "    \"\"\"\n",
    "    # group by the floor of each hour use timedelta index\n",
    "    agged = feature_df.groupby([feature_df.index.floor(\"H\")]).agg(aggs)\n",
    "\n",
    "    # flatten hierachical column index\n",
    "    agged.columns = [\"_\".join(x) for x in agged.columns]\n",
    "    return agged\n",
    "\n",
    "\n",
    "def preprocess_features(solar_wind, sunspots, scaler=None, subset=None):\n",
    "    \"\"\"\n",
    "    Preprocessing steps:\n",
    "        - Subset the data\n",
    "        - Aggregate hourly\n",
    "        - Join solar wind and sunspot data\n",
    "        - Scale using standard scaler\n",
    "        - Impute missing values\n",
    "    \"\"\"\n",
    "    # select features we want to use\n",
    "    if subset:\n",
    "        solar_wind = solar_wind[subset]\n",
    "\n",
    "    # aggregate solar wind data and join with sunspots\n",
    "    hourly_features = aggregate_hourly(solar_wind).join(sunspots)\n",
    "\n",
    "    # subtract mean and divide by standard deviation\n",
    "    if scaler is None:\n",
    "        scaler = StandardScaler()\n",
    "        scaler.fit(hourly_features)\n",
    "\n",
    "    normalized = pd.DataFrame(\n",
    "        scaler.transform(hourly_features),\n",
    "        index=hourly_features.index,\n",
    "        columns=hourly_features.columns,\n",
    "    )\n",
    "\n",
    "    # impute missing values\n",
    "    imputed = impute_features(normalized)\n",
    "\n",
    "    # we want to return the scaler object as well to use later during prediction\n",
    "    return imputed, scaler\n",
    "\n",
    "\n",
    "# THIS MUST BE DEFINED FOR YOUR SUBMISSION TO RUN\n",
    "def predict_dst(\n",
    "    solar_wind_7d: pd.DataFrame,\n",
    "    satellite_positions_7d: pd.DataFrame,\n",
    "    latest_sunspot_number: float,\n",
    ") -> Tuple[float, float]:\n",
    "    \"\"\"\n",
    "    Take all of the data up until time t-1, and then make predictions for\n",
    "    times t and t+1.\n",
    "    Parameters\n",
    "    ----------\n",
    "    solar_wind_7d: pd.DataFrame\n",
    "        The last 7 days of satellite data up until (t - 1) minutes [exclusive of t]\n",
    "    satellite_positions_7d: pd.DataFrame\n",
    "        The last 7 days of satellite position data up until the present time [inclusive of t]\n",
    "    latest_sunspot_number: float\n",
    "        The latest monthly sunspot number (SSN) to be available\n",
    "    Returns\n",
    "    -------\n",
    "    predictions : Tuple[float, float]\n",
    "        A tuple of two predictions, for (t and t + 1 hour) respectively; these should\n",
    "        be between -2,000 and 500.\n",
    "    \"\"\"\n",
    "    # Re-format data to fit into our pipeline\n",
    "    sunspots = pd.DataFrame(index=solar_wind_7d.index, columns=[\"smoothed_ssn\"])\n",
    "    sunspots[\"smoothed_ssn\"] = latest_sunspot_number\n",
    "    \n",
    "    # Process our features and grab last 32 (timesteps) hours\n",
    "    features, s = preprocess_features(\n",
    "        solar_wind_7d, sunspots, scaler=scaler, subset=SOLAR_WIND_FEATURES\n",
    "    )\n",
    "    model_input = features[-TIMESTEPS:][XCOLS].values.reshape(\n",
    "        (1, TIMESTEPS, features.shape[1])\n",
    "    )\n",
    "    \n",
    "    # Make a prediction\n",
    "    prediction_at_t0, prediction_at_t1 = model.predict(model_input)[0]\n",
    "\n",
    "    # Optional check for unexpected values\n",
    "    if not np.isfinite(prediction_at_t0):\n",
    "        prediction_at_t0 = -12\n",
    "    if not np.isfinite(prediction_at_t1):\n",
    "        prediction_at_t1 = -12\n",
    "\n",
    "    return prediction_at_t0, prediction_at_t1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}