kp_predict_tf.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 435
        },
        "id": "QPecxZR9vilX",
        "outputId": "a2349089-0d5c-4f67-db60-be6e670c882f"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "dsc_fc_summed_spectra_2023_v01.csv\n"
          ]
        },
        {
          "output_type": "error",
          "ename": "FileNotFoundError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mFileNotFoundError\u001b[0m                         Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-8-8460dc05ca6b>\u001b[0m in \u001b[0;36m<cell line: 15>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     17\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     18\u001b[0m         \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfile\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 19\u001b[0;31m         \u001b[0mpd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mread_csv\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'file1.csv'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mon_bad_lines\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'skip'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     20\u001b[0m         \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdata\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0miloc\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     21\u001b[0m         \u001b[0mdata\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"Date\"\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto_datetime\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"Date\"\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/pandas/util/_decorators.py\u001b[0m in \u001b[0;36mwrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    209\u001b[0m                 \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    210\u001b[0m                     \u001b[0mkwargs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mnew_arg_name\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnew_arg_value\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 211\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    212\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    213\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mcast\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mF\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mwrapper\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/pandas/util/_decorators.py\u001b[0m in \u001b[0;36mwrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    329\u001b[0m                     \u001b[0mstacklevel\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mfind_stack_level\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    330\u001b[0m                 )\n\u001b[0;32m--> 331\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    332\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    333\u001b[0m         \u001b[0;31m# error: \"Callable[[VarArg(Any), KwArg(Any)], Any]\" has no\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/pandas/io/parsers/readers.py\u001b[0m in \u001b[0;36mread_csv\u001b[0;34m(filepath_or_buffer, sep, delimiter, header, names, index_col, usecols, squeeze, prefix, mangle_dupe_cols, dtype, engine, converters, true_values, false_values, skipinitialspace, skiprows, skipfooter, nrows, na_values, keep_default_na, na_filter, verbose, skip_blank_lines, parse_dates, infer_datetime_format, keep_date_col, date_parser, dayfirst, cache_dates, iterator, chunksize, compression, thousands, decimal, lineterminator, quotechar, quoting, doublequote, escapechar, comment, encoding, encoding_errors, dialect, error_bad_lines, warn_bad_lines, on_bad_lines, delim_whitespace, low_memory, memory_map, float_precision, storage_options)\u001b[0m\n\u001b[1;32m    948\u001b[0m     \u001b[0mkwds\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mupdate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mkwds_defaults\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    949\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 950\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0m_read\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfilepath_or_buffer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwds\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    951\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    952\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/pandas/io/parsers/readers.py\u001b[0m in \u001b[0;36m_read\u001b[0;34m(filepath_or_buffer, kwds)\u001b[0m\n\u001b[1;32m    603\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    604\u001b[0m     \u001b[0;31m# Create the parser.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 605\u001b[0;31m     \u001b[0mparser\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mTextFileReader\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfilepath_or_buffer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwds\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    606\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    607\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mchunksize\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0miterator\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/pandas/io/parsers/readers.py\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, f, engine, **kwds)\u001b[0m\n\u001b[1;32m   1440\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1441\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mhandles\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mIOHandles\u001b[0m \u001b[0;34m|\u001b[0m \u001b[0;32mNone\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1442\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_engine\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_make_engine\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mengine\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1443\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1444\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mclose\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/pandas/io/parsers/readers.py\u001b[0m in \u001b[0;36m_make_engine\u001b[0;34m(self, f, engine)\u001b[0m\n\u001b[1;32m   1733\u001b[0m                 \u001b[0;32mif\u001b[0m \u001b[0;34m\"b\"\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mmode\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1734\u001b[0m                     \u001b[0mmode\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0;34m\"b\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1735\u001b[0;31m             self.handles = get_handle(\n\u001b[0m\u001b[1;32m   1736\u001b[0m                 \u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1737\u001b[0m                 \u001b[0mmode\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/pandas/io/common.py\u001b[0m in \u001b[0;36mget_handle\u001b[0;34m(path_or_buf, mode, encoding, compression, memory_map, is_text, errors, storage_options)\u001b[0m\n\u001b[1;32m    854\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mioargs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mencoding\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0;34m\"b\"\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mioargs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmode\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    855\u001b[0m             \u001b[0;31m# Encoding\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 856\u001b[0;31m             handle = open(\n\u001b[0m\u001b[1;32m    857\u001b[0m                 \u001b[0mhandle\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    858\u001b[0m                 \u001b[0mioargs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmode\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mFileNotFoundError\u001b[0m: [Errno 2] No such file or directory: 'file1.csv'"
          ]
        }
      ],
      "source": [
        "import os\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import tensorflow as tf\n",
        "from datetime import datetime\n",
        "from tqdm import tqdm\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "# Check if GPU is available\n",
        "device = \"cuda\" if tf.config.list_physical_devices('GPU') else \"cpu\"\n",
        "\n",
        "# Load dscovr_data\n",
        "dscovr_data = {}\n",
        "data_dir = \"/content/sample_data\"\n",
        "for file in os.listdir(data_dir):\n",
        "    if \"dsc\" in file:\n",
        "\n",
        "        print(str(file))\n",
        "        pd.read_csv('file1.csv', on_bad_lines='skip')\n",
        "        data = data.iloc[1:]\n",
        "        data[\"Date\"] = pd.to_datetime(data[\"Date\"])\n",
        "        data[\"DOY\"] = data[\"Date\"].dt.dayofyear\n",
        "        data[\"YEAR\"] = data[\"Date\"].dt.year\n",
        "        data[\"HR\"] = data[\"Time\"].apply(lambda x: int(x.split(\":\")[0]) + int(x.split(\":\")[1]) / 60)\n",
        "        data[\"Bt\"] = np.sqrt(np.sum(data[[\"Bx\", \"By\", \"Bz\"]] ** 2, axis=1))\n",
        "        data = data.drop(columns=[\"Date\", \"Time\"])\n",
        "        key = \"yr_\" + file.split(\"_\")[4]\n",
        "        dscovr_data[key] = data\n",
        "\n",
        "# Load kp_data\n",
        "kp_data = pd.read_csv(\"kP_index_data.csv\")\n",
        "kp_data[\"HR\"] = kp_data[\"HR\"].apply(lambda x: float(x))\n",
        "kp_data[\"kP\"] = kp_data[\"kP\"].apply(lambda x: x / 10)\n",
        "\n",
        "# Define sliding_window function\n",
        "def sliding_window(window_size, x_split, jump, index):\n",
        "    full_window = np.arange(jump * index, window_size + jump * index)\n",
        "    x_idxs = full_window[:x_split]\n",
        "    y_idxs = full_window[x_split:]\n",
        "    return x_idxs, y_idxs\n",
        "\n",
        "# Define ForecastingDataset\n",
        "class ForecastingDataset(tf.data.Dataset):\n",
        "    def __init__(self, data, window_size=360, x_split=240, jump=60):\n",
        "        self.idx_length = (len(data) - window_size) // jump + 1\n",
        "        self.data = data\n",
        "        self.window_size = window_size\n",
        "        self.x_split = x_split\n",
        "        self.jump = jump\n",
        "\n",
        "    def __len__(self):\n",
        "        return self.idx_length\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        x_idxs, y_idxs = sliding_window(self.window_size, self.x_split, self.jump, idx)\n",
        "        return self.data.iloc[x_idxs], self.data.iloc[y_idxs]\n",
        "\n",
        "# Define kPPredictionDataset\n",
        "class kPPredictionDataset(tf.data.Dataset):\n",
        "    def __init__(self, discovr_data, kp_data, window_size=12*60, x_split=12*60-60, jump=180):\n",
        "        self.idx_length = (len(discovr_data) - window_size) // jump + 1\n",
        "        self.discovr_data = discovr_data\n",
        "        self.kp_data = kp_data\n",
        "        self.window_size = window_size\n",
        "        self.x_split = x_split\n",
        "        self.jump = jump\n",
        "\n",
        "    def __len__(self):\n",
        "        return self.idx_length\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        x_idxs, y_idxs = sliding_window(self.window_size, self.x_split, self.jump, idx)\n",
        "        y_kp = []\n",
        "        for idx, row in self.discovr_data.iloc[y_idxs][::-60][::-1].iterrows():\n",
        "            year = row[\"YEAR\"]\n",
        "            doy = row[\"DOY\"]\n",
        "            hour = row[\"HR\"]\n",
        "            y_kp.append(self.kp_data[(self.kp_data[\"HR\"] == hour) & (self.kp_data[\"DOY\"] == doy) & (self.kp_data[\"YEAR\"] == year)].drop(columns=[\"YEAR\", \"DOY\", \"HR\"]).to_numpy())\n",
        "        return tf.convert_to_tensor(self.discovr_data.iloc[x_idxs].drop(columns=[\"YEAR\", \"HR\"])[[\"DOY\", \"Bt\", \"Bx\", \"By\", \"Bz\"]].to_numpy()), tf.convert_to_tensor(np.stack(y_kp))\n",
        "\n",
        "# Create kP datasets\n",
        "kp_ds = kPPredictionDataset(dscovr_data[\"yr_2023\"], kp_data[kp_data[\"YEAR\"]==2023])\n",
        "train_size = int(len(kp_ds) * 0.75)\n",
        "train_set, val_set = tf.data.experimental.random_split(kp_ds, [train_size, len(kp_ds) - train_size])\n",
        "train_dataloader = train_set.batch(16, drop_remainder=True).shuffle(buffer_size=len(train_set))\n",
        "test_dataloader = train_set.batch(16, drop_remainder=True).shuffle(buffer_size=len(train_set))\n",
        "\n",
        "# Define LSTM model\n",
        "class LSTM(tf.keras.Model):\n",
        "    def __init__(self, num_classes, input_size, hidden_size, num_layers):\n",
        "        super(LSTM, self).__init__()\n",
        "        self.num_classes = num_classes\n",
        "        self.num_layers = num_layers\n",
        "        self.input_size = input_size\n",
        "        self.hidden_size = hidden_size\n",
        "        self.lstm = tf.keras.layers.LSTM(\n",
        "            hidden_size, return_sequences=True, return_state=True\n",
        "        )\n",
        "        self.fc = tf.keras.layers.Dense(num_classes)\n",
        "\n",
        "    def call(self, x):\n",
        "        x, h_state, c_state = self.lstm(x)\n",
        "        x = self.fc(x)\n",
        "        return x\n",
        "\n",
        "# Create and compile the LSTM model\n",
        "model = LSTM(num_classes=5, input_size=5, hidden_size=128, num_layers=2)\n",
        "model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.001), loss='mse')\n",
        "\n",
        "# Training loop\n",
        "EPOCHS = 10\n",
        "for epoch in range(EPOCHS):\n",
        "    total_loss = 0\n",
        "    for x, y in train_dataloader:\n",
        "        with tf.GradientTape() as tape:\n",
        "            predictions = model(x)\n",
        "            loss = tf.keras.losses.MeanSquaredError()(y, predictions)\n",
        "        gradients = tape.gradient(loss, model.trainable_variables)\n",
        "        optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
        "        total_loss += loss\n",
        "    print(f\"Epoch {epoch+1}, Loss: {total_loss.numpy() / len(train_dataloader)}\")\n",
        "\n",
        "# Evaluate the model on the test set\n",
        "test_loss = tf.keras.metrics.MeanSquaredError()\n",
        "for x, y in test_dataloader:\n",
        "    predictions = model(x)\n",
        "    test_loss(y, predictions)\n",
        "print(f\"Test Loss: {test_loss.result().numpy()}\")\n",
        "\n",
        "# Plot some predictions\n",
        "n_samples = 5\n",
        "for x, y in test_dataloader.take(n_samples):\n",
        "    predictions = model(x)\n",
        "    plt.figure(figsize=(12, 6))\n",
        "    for i in range(5):\n",
        "        plt.subplot(5, 1, i+1)\n",
        "        plt.plot(y[i, :, 0], label='True')\n",
        "        plt.plot(predictions[i, :, 0], label='Predicted')\n",
        "        plt.legend()\n",
        "    plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "metadata": {
        "id": "7v56idpCAqOD"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}