Merge pull request #21 from Oisin-M/feat/gfn

Feat/gfn

Author: fpichi

tutorials/pooling/gfn.ipynb

@@ -0,0 +1,211 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "\n",
+    "# GFN-ROM is a resolution-invariant method for MOR suitable for multifidelity applications.\n",
+    "\n",
+    "# For further details, see the [GFN repo](https://github.com/Oisin-M/GFN) and [GFN paper](https://arxiv.org/abs/2406.03569)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "CtOVuaD1oJGf"
+   },
+   "outputs": [],
+   "source": [
+    "# Install PyTorch\n",
+    "try:\n",
+    "  import torch\n",
+    "  from torch import nn\n",
+    "except ImportError:\n",
+    "  !pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118\n",
+    "  import torch\n",
+    "  from torch import nn"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "rXVVqRjT0CHr"
+   },
+   "outputs": [],
+   "source": [
+    "# Clone and import gfn-rom\n",
+    "import sys\n",
+    "try:\n",
+    "    from gfn_rom import pde, defaults, preprocessing, initialisation, gfn_rom, train, test, plotting\n",
+    "except ImportError:\n",
+    "    try:\n",
+    "        sys.path.append('GFN')\n",
+    "        from gfn_rom import pde, defaults, preprocessing, initialisation, gfn_rom, train, test, plotting\n",
+    "    except ImportError:\n",
+    "        !git clone https://github.com/Oisin-M/GFN.git\n",
+    "        from gfn_rom import pde, defaults, preprocessing, initialisation, gfn_rom, train, test, plotting\n",
+    "\n",
+    "import numpy as np\n",
+    "import scipy"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pname = 'advection'\n",
+    "\n",
+    "# training and test fidelities\n",
+    "train_fidelities = ['3967']\n",
+    "test_fidelities = ['3967']\n",
+    "\n",
+    "# Naming convention for saving the model\n",
+    "save_name = ''.join(train_fidelities)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "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",
+    "\n",
+    "xs=scipy.io.loadmat(\"../../dataset/advection_unstructured.mat\")['xx'][:,0]\n",
+    "ys=scipy.io.loadmat(\"../../dataset/advection_unstructured.mat\")['yy'][:,0]\n",
+    "meshes_train = [np.vstack([xs,ys]).T]\n",
+    "meshes_test = meshes_train\n",
+    "def get_scaled_data(fname=\"../../dataset/advection_unstructured.mat\"):\n",
+    "    U = scipy.io.loadmat(fname)['U']\n",
+    "    U_orig = torch.tensor(U)\n",
+    "    scale, U_sc = preprocessing.scaling(U_orig)\n",
+    "    print('reconstruction error', ((U_orig - preprocessing.undo_scaling(U_sc, scale))**2).sum())\n",
+    "    return scale, U_sc\n",
+    "sols_train = [get_scaled_data()[1]]\n",
+    "sols_test = [get_scaled_data()]\n",
+    "\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": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "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": null,
+   "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": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "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": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "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": {
+  "accelerator": "GPU",
+  "colab": {
+   "authorship_tag": "ABX9TyPxps/Yo6EhPBLUacBJicyu",
+   "gpuType": "T4",
+   "include_colab_link": true,
+   "provenance": []
+  },
+  "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": 4
+}