advection/gfn_rom_tiny.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d18472d4",
   "metadata": {},
   "source": [
    "# GFN-ROM\n",
    "\n",
    "This notebook trains the GFN-ROM model and evaluates the performance."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bda34993",
   "metadata": {},
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "aaca74e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "import numpy as np\n",
    "\n",
    "from torchinfo import summary\n",
    "\n",
    "import sys\n",
    "sys.path.append(\"..\")\n",
    "from gfn_rom import pde, defaults, preprocessing, initialisation, gfn_rom, train, test, plotting"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d9055c8",
   "metadata": {},
   "source": [
    "# Hyperparameter and Problem Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "bfff03e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "pname = 'advection'\n",
    "\n",
    "# training and test fidelities\n",
    "train_fidelities = ['tiny']\n",
    "test_fidelities = ['large', 'medium', 'small', 'tiny']\n",
    "\n",
    "# Naming convention for saving the model\n",
    "save_name = ''.join(train_fidelities)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "41bc5a3e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "reconstruction error tensor(1.2396e-28)\n",
      "reconstruction error tensor(3.2023e-27)\n",
      "reconstruction error tensor(1.0696e-27)\n",
      "reconstruction error tensor(3.4364e-28)\n",
      "reconstruction error tensor(1.2396e-28)\n"
     ]
    }
   ],
   "source": [
    "dev = initialisation.set_device()\n",
    "initialisation.set_precision(defaults.precision)\n",
    "initialisation.create_directories()\n",
    "params = torch.tensor(pde.params(pname)).to(dev)\n",
    "np.random.seed(defaults.split_seed)\n",
    "train_trajs, test_trajs = preprocessing.train_test_split(params, len(train_fidelities), defaults.rate)\n",
    "meshes_train, sols_train, meshes_test, sols_test = preprocessing.load_and_process_datasets(train_fidelities, test_fidelities)\n",
    "sols_train = [x.to(dev) for x in sols_train]\n",
    "initialisation.set_seed(defaults.seed)\n",
    "start_mesh = sorted(meshes_train, key=lambda x: x.shape[0])[-1]\n",
    "update_master = defaults.mode == 'adapt'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8c5ba89",
   "metadata": {},
   "source": [
    "# Model Initialisation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "23fa4daa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(326, 2)\n",
      "(326, 2)\n"
     ]
    }
   ],
   "source": [
    "model = gfn_rom.GFN_ROM(start_mesh, defaults.N_basis_factor, params.shape[1], defaults.act, defaults.ae_sizes, defaults.mapper_sizes).to(dev)\n",
    "print(model.GFN.mesh_m.shape)\n",
    "\n",
    "# We do all of the possible expansions apriori in the preadaptive case\n",
    "# This is a preprocessing step so we don't do any speedup steps here\n",
    "if defaults.mode=='preadapt':\n",
    "    count = np.inf\n",
    "    while count!=0:\n",
    "        count = 0\n",
    "        for mesh_n in meshes_train:\n",
    "            n_exp, n_agg = model.GFN.reshape_weights(mesh_n, update_master=True)\n",
    "            count += n_exp\n",
    "    print(model.GFN.mesh_m.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "242dd65e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "=================================================================\n",
       "Layer (type:depth-idx)                   Param #\n",
       "=================================================================\n",
       "GFN_ROM                                  --\n",
       "├─GFN_AE: 1-1                            130,926\n",
       "├─Tanh: 1-2                              --\n",
       "├─Sequential: 1-3                        --\n",
       "│    └─Linear: 2-1                       603\n",
       "│    └─Tanh: 2-2                         --\n",
       "├─Sequential: 1-4                        --\n",
       "│    └─Linear: 2-3                       800\n",
       "│    └─Tanh: 2-4                         --\n",
       "├─Sequential: 1-5                        --\n",
       "│    └─Linear: 2-5                       150\n",
       "│    └─Tanh: 2-6                         --\n",
       "│    └─Linear: 2-7                       2,550\n",
       "│    └─Tanh: 2-8                         --\n",
       "│    └─Linear: 2-9                       2,550\n",
       "│    └─Tanh: 2-10                        --\n",
       "│    └─Linear: 2-11                      2,550\n",
       "│    └─Tanh: 2-12                        --\n",
       "│    └─Linear: 2-13                      153\n",
       "=================================================================\n",
       "Total params: 140,282\n",
       "Trainable params: 140,282\n",
       "Non-trainable params: 0\n",
       "================================================================="
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "summary(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "834c95b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not update_master:\n",
    "    opt = torch.optim.Adam(model.parameters(), lr=defaults.lr, weight_decay=defaults.lambda_)\n",
    "else:\n",
    "    # Cannot update GFN parameters using Adam anymore since we use adaptive method\n",
    "    # and weights can change shape at each iteration\n",
    "    # Similarly, cannot use momentum\n",
    "    opt = torch.optim.SGD(model.parameters(), lr=defaults.lr, weight_decay=defaults.lambda_)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d40bd139",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a04e245d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading saved network\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    model.load_state_dict(torch.load(\"models/best_model_\"+save_name+\".pt\"))\n",
    "    print(\"Loading saved network\")\n",
    "except FileNotFoundError:\n",
    "    print(\"Training network\")\n",
    "    train_losses, test_losses = train.train(model, opt, meshes_train, sols_train, params, train_trajs, test_trajs, update_master, defaults.epochs, defaults.mapper_weight, save_name)\n",
    "    model.load_state_dict(torch.load(\"models/best_model_\"+save_name+\".pt\"))\n",
    "    plotting.plot_losses(train_losses, test_losses, save_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc201167",
   "metadata": {},
   "source": [
    "# Testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d7c70583",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------\n",
      "TEST MESH: large\n",
      "\n",
      "Maximum relative error for latent  =  0.646364046095673\n",
      "Mean relative error for latent =  0.15373685846837878\n",
      "Minimum relative error for latent =  0.015265202061717536\n",
      "\n",
      "Maximum absolute error for field  =  39.15898584957899\n",
      "Mean absolute error for field  =  9.324528137823002\n",
      "Minimum absolute error for field  =  0.3853870100707307\n",
      "\n",
      "Maximum relative error for field  =  0.17666213232425698\n",
      "Mean relative error for field  =  0.12348877055433212\n",
      "Minimum relative error for field  =  0.07226659067257475\n",
      "\n",
      "----------------------------------------\n",
      "TEST MESH: medium\n",
      "\n",
      "Maximum relative error for latent  =  0.6634368946484427\n",
      "Mean relative error for latent =  0.15836886375024012\n",
      "Minimum relative error for latent =  0.022433965599711694\n",
      "\n",
      "Maximum absolute error for field  =  21.838688428480637\n",
      "Mean absolute error for field  =  5.149770603414611\n",
      "Minimum absolute error for field  =  0.2129012637151258\n",
      "\n",
      "Maximum relative error for field  =  0.17664841717858548\n",
      "Mean relative error for field  =  0.12304023054461821\n",
      "Minimum relative error for field  =  0.0738239514788489\n",
      "\n",
      "----------------------------------------\n",
      "TEST MESH: small\n",
      "\n",
      "Maximum relative error for latent  =  0.6381040932910931\n",
      "Mean relative error for latent =  0.14722939017564676\n",
      "Minimum relative error for latent =  0.03668449179135777\n",
      "\n",
      "Maximum absolute error for field  =  12.446952137456055\n",
      "Mean absolute error for field  =  2.902570062171734\n",
      "Minimum absolute error for field  =  0.11817120632754638\n",
      "\n",
      "Maximum relative error for field  =  0.17534850662544033\n",
      "Mean relative error for field  =  0.11949690580397902\n",
      "Minimum relative error for field  =  0.07448794055540987\n",
      "\n",
      "----------------------------------------\n",
      "TEST MESH: tiny\n",
      "\n",
      "Maximum relative error for latent  =  0.474721769722547\n",
      "Mean relative error for latent =  0.08408625772440592\n",
      "Minimum relative error for latent =  0.012517647972424804\n",
      "\n",
      "Maximum absolute error for field  =  4.731296504211749\n",
      "Mean absolute error for field  =  0.6864411501636155\n",
      "Minimum absolute error for field  =  0.006076201480747502\n",
      "\n",
      "Maximum relative error for field  =  0.14431249589000278\n",
      "Mean relative error for field  =  0.04697245577520337\n",
      "Minimum relative error for field  =  0.008521214239651399\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for i in range(len(test_fidelities)):\n",
    "    \n",
    "    print('-'*40)\n",
    "    print(f'TEST MESH: {test_fidelities[i]}')\n",
    "    \n",
    "    scale, U = sols_test[i]\n",
    "    U = U.to('cpu')\n",
    "    mesh = meshes_test[i]\n",
    "\n",
    "    model.eval()\n",
    "    model.to('cpu')\n",
    "    \n",
    "    Z, Z_net, x_enc, x_map = test.evaluate_results(model, mesh, U, scale, params.to('cpu'))\n",
    "    error = abs(Z - Z_net)\n",
    "    error, rel_error = test.print_results(Z, Z_net, x_enc, x_map)\n",
    "\n",
    "    np.savetxt('errors/relative_errors_train'+save_name+'_test'+test_fidelities[i]+'.txt', [max(rel_error), sum(rel_error)/len(rel_error), min(rel_error)])\n",
    "    print()"
   ]
  }
 ],
 "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}