NN_model_code_standalone.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "46989dab",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader, TensorDataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "a79c4a38",
   "metadata": {},
   "outputs": [],
   "source": [
    "######### NEURAL NETWORK MODEL ###########\n",
    "\n",
    "def nn_model(t, time_series, hidden_width, epochs, title, save_fig=False):\n",
    "    print(torch.cuda.is_available())\n",
    "    print(torch.cuda.device_count())\n",
    "    print(torch.cuda.current_device())\n",
    "    print(torch.cuda.device(0))\n",
    "    print(torch.cuda.get_device_name(0))\n",
    "    \n",
    "    #Our input dimension is 1, and each entry in the tensor above is a datapoint for training. So, the length of \n",
    "    #the tensor/numpy array should be the batch size:\n",
    "    n_input, n_hidden, n_output = 1, hidden_width, 1\n",
    "    batch_size = 10\n",
    "    \n",
    "    #Turn into tensor\n",
    "    t_tensor = torch.from_numpy(np.reshape(t, (-1, n_input))).to(torch.float32)\n",
    "    y_tensor = torch.from_numpy(np.reshape(time_series, (-1, n_output))).to(torch.float32)\n",
    "\n",
    "    dataset = TensorDataset(t_tensor, y_tensor)\n",
    "    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)\n",
    "    \n",
    "    dataset_size = len(dataloader.dataset)\n",
    "    \n",
    "    learning_rate = 0.001\n",
    "    \n",
    "    plt.ylabel('y')\n",
    "    plt.xlabel('t')\n",
    "    plt.title(\"Given Time Series\")\n",
    "    plt.plot(t_tensor, y_tensor, marker='.', linestyle='none')\n",
    "    plt.show()\n",
    "\n",
    "    model = nn.Sequential(nn.Linear(n_input, n_hidden), nn.ReLU(), nn.Linear(n_hidden, n_output))\n",
    "    print(model)\n",
    "\n",
    "    loss_function = nn.MSELoss(reduction=\"sum\")\n",
    "    print(loss_function)\n",
    "    \n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)\n",
    "\n",
    "    losses = [] # For plotting purposes\n",
    "    print(\"---------------------------------------------------------\")\n",
    "    print(\"START TRAINING\")\n",
    "    print(f\"HIDDEN LAYER WIDTH: {hidden_width}\")\n",
    "    for epoch in range(epochs):\n",
    "        #print(f\"Epoch {epoch}\\n------------------------\")\n",
    "        \n",
    "        losses_epoch = []\n",
    "        for id_batch, (t_batch, y_batch) in enumerate(dataloader):\n",
    "            # Forward pass\n",
    "            y_pred = model(t_batch)\n",
    "            \n",
    "            # Calculate the loss\n",
    "            loss = loss_function(y_pred, y_batch)\n",
    "            losses_epoch.append(loss.item())\n",
    "            \n",
    "            # We set the gradients to zero, as PyTorch accumulates gradients on subsequent backward passes.\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "        losses.append(np.mean(losses_epoch))\n",
    "    print(\"END TRAINING\")\n",
    "    \n",
    "    y_pred_final = model(t_tensor)\n",
    "    rel_error_final = relative_error(y_pred_final.detach().numpy(), y_tensor.detach().numpy())\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        plt.title(f\"{title}\")\n",
    "            \n",
    "        plt.plot(t_tensor, y_tensor, label=\"True\") #, marker='.', linestyle='none')\n",
    "        plt.plot(t_tensor, y_pred_final, label=f\"NN Pred. Width={hidden_width}\") #, marker='.', linestyle='none', color='orange')\n",
    "        plt.legend()\n",
    "        if save_fig:\n",
    "            plt.savefig(savefigs_loc_nn+\"_TRAINING_\"+title+f\"_hw_{hidden_width}__epochs_{epochs}__rel_error_{rel_error_final:.4}_.png\")\n",
    "        plt.show()\n",
    "    \n",
    "    plt.plot(losses)\n",
    "    plt.ylabel('loss')\n",
    "    plt.xlabel('epoch')\n",
    "    plt.title(title+f\"_width={hidden_width}_\"+\"lr = %f\"%(learning_rate))\n",
    "    if save_fig:\n",
    "        plt.savefig(savefigs_loc_nn+\"_LOSS_\"+title+f\"_hw_{hidden_width}__epochs_{epochs}__rel_error_{rel_error_final:.4}_.png\")\n",
    "    plt.show()\n",
    "    \n",
    "    print(f\"Relative error = {rel_error_final:.4}\")\n",
    "    print(\"--------------------------------------\")\n",
    "    return model, rel_error_final, y_pred_final, y_tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "7d11d0e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "#### CONSTRUCT MULTUPLE NEURAL NETWORKS WITH DIFFERENT WIDTHS ####\n",
    "\n",
    "def nn_models(hidden_widths, t, time_series, epochs, title, save_fig=False):\n",
    "    models_dict = {}\n",
    "    rel_err_dict = {}\n",
    "    final_preds_dict = {}\n",
    "    ts_tensor_dict={}\n",
    "    \n",
    "    for hw in hidden_widths:\n",
    "        mod, rel_err, y_pred_final, ts_tensor = nn_model(t, time_series, hw, epochs, title, save_fig)\n",
    "        models_dict[hw]=mod\n",
    "        rel_err_dict[hw]=rel_err\n",
    "        \n",
    "        final_preds_dict[hw] = y_pred_final\n",
    "        ts_tensor_dict[hw] = ts_tensor\n",
    "    \n",
    "    plt.plot(t, ts_tensor_dict[hidden_widths[0]].detach().numpy(), label=\"True\")\n",
    "    for hw in hidden_widths:\n",
    "        plt.plot(t, final_preds_dict[hw].detach().numpy(), label=f\"NN pred. Width={hw}\")\n",
    "    plt.legend()\n",
    "    plt.title(title)\n",
    "    if save_fig:\n",
    "        plt.savefig(savefigs_loc_nn+\"_TRAINING_ALL_HWs__\"+title+f\"__epochs_{epochs}_.png\")\n",
    "    plt.show()\n",
    "    \n",
    "    \n",
    "    x = y = list(range(hidden_widths[-1]))\n",
    "    idx = rel_err_dict.keys()\n",
    "    new_y = [rel_err_dict[i] for i in idx]\n",
    "    plt.plot(range(len(idx)), new_y, 'o-')\n",
    "    plt.xticks(range(len(idx)),idx)\n",
    "    plt.xlabel(\"Hidden layer width\")\n",
    "    plt.ylabel(\"Training, Relative error\")\n",
    "    plt.legend()\n",
    "    if save_fig:\n",
    "        plt.savefig(savefigs_loc_nn+\"_REL_ERRORS_vs_HWs__\"+title+f\"__epochs_{epochs}_.png\")\n",
    "    plt.show()\n",
    "    \n",
    "    best_width = min(rel_err_dict, key=rel_err_dict.get)\n",
    "    print(\"Best width: \", best_width)\n",
    "    return models_dict, best_width"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "3dbf1175",
   "metadata": {},
   "outputs": [],
   "source": [
    "### TEST NEURAL NETWORK ON TEST DATA, ON TRAINING INTERVAL OR FORWARD IN TIME ###\n",
    "\n",
    "def nn_test(model,hidden_width, test_time, test_time_series, title, forward_time=False, save_fig=False):\n",
    "    n_input, n_output = 1, 1\n",
    "    test_time_tensor = torch.from_numpy(np.reshape(test_time, (-1, n_input))).to(torch.float32)\n",
    "    test_time_series_tensor = torch.from_numpy(np.reshape(test_time_series, (-1, n_output))).to(torch.float32)\n",
    "    \n",
    "    pred = model(test_time_tensor)\n",
    "    \n",
    "    plt.title(title)\n",
    "    \n",
    "    if forward_time:\n",
    "        middle_index = test_time.shape[0]//2\n",
    "        forward_rel_err = relative_error(pred.detach().numpy()[middle_index:], test_time_series_tensor.detach().numpy()[middle_index:])\n",
    "          \n",
    "        plt.plot(test_time_tensor.detach().numpy(), test_time_series_tensor.detach().numpy(), label=\"True\")\n",
    "        plt.plot(test_time_tensor.detach().numpy(), pred.detach().numpy(), label=f\"Pred. Width={hidden_width}\")\n",
    "        plt.legend()\n",
    "        if save_fig:\n",
    "            plt.savefig(savefigs_loc_nn+\"_FORWARD_\"+title+f\"_hw_{hidden_width}__forward_rel_error_{forward_rel_err:.4}_.png\")\n",
    "        \n",
    "        print(\"Forward, relative error (on future time points): \", forward_rel_err)\n",
    "    else:\n",
    "        test_rel_err = relative_error(pred.detach().numpy(), test_time_series_tensor.detach().numpy())\n",
    "        \n",
    "        plt.plot(test_time_tensor.detach().numpy(), tesla_2021_data.to_numpy(), color=\"gray\",alpha=0.5, label=\"Original\")\n",
    "        \n",
    "        plt.plot(test_time_tensor.detach().numpy(), test_time_series_tensor.detach().numpy(), label=\"Smoothed\")\n",
    "        plt.plot(test_time_tensor.detach().numpy(), pred.detach().numpy(), label=f\"Pred. Width={hidden_width}\")\n",
    "        plt.legend()\n",
    "        if save_fig:\n",
    "            plt.savefig(savefigs_loc_nn+\"_TESTING_\"+title+f\"_hw_{hidden_width}__forward_rel_error_{test_rel_err:.4}_.png\")\n",
    "\n",
    "        \n",
    "        print(\"Testing (same interval), relative error: \", test_rel_err)\n",
    "    \n",
    "    plt.show()\n",
    "    return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a022d00a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def relative_error(prediction, true):\n",
    "    mse = np.mean(np.power(prediction-true,2))\n",
    "    print(\"mse: \", mse)\n",
    "    true_size = np.mean(np.power(true,2))\n",
    "    print(\"true size: \", true_size)\n",
    "    return mse/true_size"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}