finetune_FNO4D_DP_LGR.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4cba1655",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Figure size 432x288 with 0 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "import matplotlib.patches as patches\n",
    "import glob\n",
    "import os\n",
    "from utility import *\n",
    "from UnitGaussianNormalizer import *\n",
    "import torch.nn.functional as F\n",
    "import torch\n",
    "from CustomDataset import *\n",
    "import pickle\n",
    "from FNO4D import *\n",
    "import json\n",
    "from visulization_compare import *\n",
    "plt.jet()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6a13b872",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['input', 'output'])\n",
      "dict_keys(['GLOBAL', 'LGR1', 'LGR2', 'LGR3', 'LGR4']) dict_keys(['GLOBAL', 'LGR1', 'LGR2', 'LGR3', 'LGR4'])\n"
     ]
    }
   ],
   "source": [
    "NORMALIZER_DICT = {}\n",
    "d_in, d_out = {}, {}\n",
    "for key in ['GLOBAL', 'LGR1', 'LGR2', 'LGR3', 'LGR4']:\n",
    "    with open(f\"normalizer/input_normalizer_{key}_DP_val.pickle\", 'rb') as f:\n",
    "        input_normalizer = pickle.load(f)\n",
    "    with open(f\"normalizer/output_normalizer_{key}_DP_val.pickle\", 'rb') as f:\n",
    "        output_normalizer = pickle.load(f)\n",
    "    input_normalizer.cuda()\n",
    "    output_normalizer.cuda()\n",
    "    d_in[key] = input_normalizer\n",
    "    d_out[key] = output_normalizer\n",
    "NORMALIZER_DICT['input'] = d_in\n",
    "NORMALIZER_DICT['output'] = d_out\n",
    "\n",
    "print(NORMALIZER_DICT.keys())\n",
    "print(NORMALIZER_DICT['input'].keys(), NORMALIZER_DICT['output'].keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "24678fef",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['GLOBAL', 'LGR1', 'LGR2', 'LGR3', 'LGR4'])\n"
     ]
    }
   ],
   "source": [
    "# Change these models to your trained weights\n",
    "PATH = {}\n",
    "PATH['GLOBAL'] = \"pre_trained_models/FNO4D-GLOBAL-DP.pt\"\n",
    "PATH['LGR1'] = \"pre_trained_models/FNO4D-LGR1-DP.pt\"\n",
    "PATH['LGR2'] = \"pre_trained_models/FNO4D-LGR2-DP.pt\"\n",
    "PATH['LGR3'] = \"pre_trained_models/FNO4D-LGR3-DP.pt\"\n",
    "PATH['LGR4'] = \"pre_trained_models/FNO4D-LGR4-DP.pt\"\n",
    "device = torch.device('cuda')\n",
    "\n",
    "MODEL_DICT = {}\n",
    "for key in ['GLOBAL', 'LGR1', 'LGR2', 'LGR3', 'LGR4']:\n",
    "    model = torch.load(PATH[key])\n",
    "    model.to(device)\n",
    "    model.eval()\n",
    "    MODEL_DICT[key] = model\n",
    "    \n",
    "print(MODEL_DICT.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4d667b14",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "16 2 2\n"
     ]
    }
   ],
   "source": [
    "DATA_LOADER_DICT = torch.load('DATA_LOADER_DICT.pth')\n",
    "train_loader = DATA_LOADER_DICT['GLOBAL']['train']\n",
    "val_loader = DATA_LOADER_DICT['GLOBAL']['val']\n",
    "test_loader = DATA_LOADER_DICT['GLOBAL']['test']\n",
    "n_train = len(train_loader)\n",
    "n_val = len(val_loader)\n",
    "n_test = len(test_loader)\n",
    "print(n_train, n_val, n_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "2334ea1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "it = iter(train_loader)\n",
    "ERR_LIST = []\n",
    "\n",
    "for counter in range(len(train_loader)):\n",
    "    with torch.no_grad():\n",
    "        PRED, TRUE = {}, {}\n",
    "        data = next(it)\n",
    "        x, y, path = data['x'], data['y'], data['path']\n",
    "        x, y = x[None,...].to(device), y[None,...]\n",
    "        x[...,-1:] = NORMALIZER_DICT['input']['GLOBAL'].encode(x.to(device)[...,-1:])\n",
    "        pred = NORMALIZER_DICT['output']['GLOBAL'].decode(MODEL_DICT['GLOBAL'](x)).cpu()\n",
    "        PRED['GLOBAL'] = pred\n",
    "        TRUE['GLOBAL'] = y\n",
    "        slope, idx, well = path[0], path[1], path[2]\n",
    "\n",
    "        meta_data = np.load(f'ECLIPSE/meta_data/{slope}_{idx}.npy', allow_pickle=True).tolist()\n",
    "        WELL_LIST = meta_data[f'case_{idx}']['WELL_LIST']\n",
    "        GRID_IDX_DICT = meta_data[f'case_{idx}']['GRID_IDX_DICT']\n",
    "\n",
    "        for well in WELL_LIST:\n",
    "            lgr_dict, true_dict = {}, {}\n",
    "\n",
    "            data_LGR1 = torch.load(f'dataset/dP_LGR1/{slope}_{idx}_LGR1_{well}_DP.pt')\n",
    "            I1, I2 = GRID_IDX_DICT[well]['LGR1']['I1']-1-15, GRID_IDX_DICT[well]['LGR1']['I2']+15\n",
    "            J1, J2 = GRID_IDX_DICT[well]['LGR1']['J1']-1-15, GRID_IDX_DICT[well]['LGR1']['J2']+15\n",
    "            coarse = np.repeat(PRED['GLOBAL'][0,...][:,I1:I2,J1:J2,:,:],5,axis=-2).permute(-1,1,2,3,0)[...,None]\n",
    "            x_LGR1 = torch.cat((data_LGR1['input'][...,:-1],coarse),axis=-1)\n",
    "            x_LGR1 = x_LGR1.permute(0,4,1,2,3,5).to(device)\n",
    "            x_LGR1[...,-1:] = NORMALIZER_DICT['input']['LGR1'].encode(x_LGR1.to(device)[...,-1:])\n",
    "            pred = NORMALIZER_DICT['output']['LGR1'].decode(MODEL_DICT['LGR1'](x_LGR1)).cpu()\n",
    "            lgr_dict['LGR1'] = pred\n",
    "            y = data_LGR1['output'][...,:1].permute(0,4,1,2,3,5)\n",
    "            true_dict['LGR1'] = y\n",
    "\n",
    "            data_LGR2 = torch.load(f'dataset/dP_LGR2/{slope}_{idx}_LGR2_{well}_DP.pt')\n",
    "            coarse = np.repeat(lgr_dict['LGR1'][0,...],2,axis=-2).permute(-1,1,2,3,0)[...,None]\n",
    "            x_LGR2 = torch.cat((data_LGR2['input'][...,:-1],coarse),axis=-1)\n",
    "            x_LGR2 = x_LGR2.permute(0,4,1,2,3,5).to(device)\n",
    "            x_LGR2[...,-1:] = NORMALIZER_DICT['input']['LGR2'].encode(x_LGR2.to(device)[...,-1:])\n",
    "            pred = NORMALIZER_DICT['output']['LGR2'].decode(MODEL_DICT['LGR2'](x_LGR2)).cpu()\n",
    "            lgr_dict['LGR2'] = pred\n",
    "            y = data_LGR2['output'][...,:1].permute(0,4,1,2,3,5)\n",
    "            true_dict['LGR2'] = y\n",
    "\n",
    "            data_LGR3 = torch.load(f'dataset/dP_LGR3/{slope}_{idx}_LGR3_{well}_DP.pt')\n",
    "            coarse = lgr_dict['LGR2'][0,...].permute(-1,1,2,3,0)[...,None]\n",
    "            x_LGR3 = torch.cat((data_LGR3['input'][...,:-1],coarse),axis=-1)\n",
    "            x_LGR3 = x_LGR3.permute(0,4,1,2,3,5).to(device)\n",
    "            x_LGR3[...,-1:] = NORMALIZER_DICT['input']['LGR3'].encode(x_LGR3.to(device)[...,-1:])\n",
    "            pred = NORMALIZER_DICT['output']['LGR3'].decode(MODEL_DICT['LGR3'](x_LGR3)).cpu()\n",
    "            lgr_dict['LGR3'] = pred\n",
    "            y = data_LGR3['output'][...,:1].permute(0,4,1,2,3,5)\n",
    "            true_dict['LGR3'] = y\n",
    "\n",
    "            data_LGR4 = torch.load(f'dataset/dP_LGR4/{slope}_{idx}_LGR4_{well}_DP.pt')\n",
    "            coarse = lgr_dict['LGR3'][0,...].permute(-1,1,2,3,0)[...,None]\n",
    "            x = data_LGR4['input']\n",
    "            x_LGR4 = torch.cat((x[...,:-1],coarse),axis=-1)\n",
    "            x_LGR4 = x_LGR4.permute(0,4,1,2,3,5).to(device)\n",
    "            x_LGR4[...,-1:] = NORMALIZER_DICT['input']['LGR4'].encode(x_LGR4.to(device)[...,-1:])\n",
    "            pred = NORMALIZER_DICT['output']['LGR4'].decode(MODEL_DICT['LGR4'](x_LGR4)).cpu()\n",
    "            lgr_dict['LGR4'] = pred\n",
    "            y = data_LGR4['output'][...,:1].permute(0,4,1,2,3,5)\n",
    "            true_dict['LGR4'] = y\n",
    "\n",
    "            err = pred - y\n",
    "            ERR_LIST.append(err.numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c198b65",
   "metadata": {},
   "source": [
    "# Save error in matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "27254451",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(43, 24, 40, 40, 50, 1)\n"
     ]
    }
   ],
   "source": [
    "err = np.array(ERR_LIST).reshape(-1, 24, 40, 40, 50, 1)\n",
    "print(err.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dda34ef2",
   "metadata": {},
   "source": [
    "# Reload from trained weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "26e50dde",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "FNO4d(\n",
       "  (fc0): Linear(in_features=9, out_features=28, bias=True)\n",
       "  (conv): Block4d(\n",
       "    (conv0): SpectralConv4d()\n",
       "    (conv1): SpectralConv4d()\n",
       "    (conv2): SpectralConv4d()\n",
       "    (conv3): SpectralConv4d()\n",
       "    (w0): Conv1d(28, 28, kernel_size=(1,), stride=(1,))\n",
       "    (w1): Conv1d(28, 28, kernel_size=(1,), stride=(1,))\n",
       "    (w2): Conv1d(28, 28, kernel_size=(1,), stride=(1,))\n",
       "    (w3): Conv1d(28, 28, kernel_size=(1,), stride=(1,))\n",
       "    (fc1): Linear(in_features=28, out_features=112, bias=True)\n",
       "    (fc2): Linear(in_features=112, out_features=1, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from lploss import *\n",
    "LPloss = LpLoss(size_average=True)\n",
    "\n",
    "from FNO4D import *\n",
    "device = torch.device('cuda')\n",
    "width = 28\n",
    "model = torch.load( f\"pre_trained_models/FNO4D-LGR4-DP.pt\")\n",
    "model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "89f6eaf1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FNO4D-LGR4-DP-1107-1730-finetune16\n"
     ]
    }
   ],
   "source": [
    "from datetime import datetime\n",
    "from datetime import date\n",
    "\n",
    "now = datetime.now()\n",
    "today = date.today()\n",
    "key = 'LGR4' \n",
    "day = today.strftime(\"%m%d\")\n",
    "current_time = now.strftime(\"%H%M\")\n",
    "specs = f'FNO4D-{key}-DP'\n",
    "model_str = f'{day}-{current_time}-finetune{n_train}'\n",
    "print(f'{specs}-{model_str}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "d70d4f15",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(f\"normalizer/input_normalizer_{key}_DP_val.pickle\", 'rb') as f:\n",
    "    input_normalizer = pickle.load(f)\n",
    "    input_normalizer.cuda()\n",
    "    \n",
    "with open(f\"normalizer/output_normalizer_{key}_DP_val.pickle\", 'rb') as f:\n",
    "    output_normalizer = pickle.load(f)\n",
    "    output_normalizer.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "2ed2fc16",
   "metadata": {},
   "outputs": [],
   "source": [
    "from Adam import Adam\n",
    "scheduler_step = 2\n",
    "scheduler_gamma = 0.85\n",
    "learning_rate = 1e-3\n",
    "\n",
    "optimizer = Adam(model.parameters(), lr=learning_rate, weight_decay=0)\n",
    "scheduler = torch.optim.lr_scheduler.StepLR(optimizer, \n",
    "                                        step_size=scheduler_step, \n",
    "                                        gamma=scheduler_gamma)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "f6f92ae1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.tensorboard import SummaryWriter\n",
    "writer = SummaryWriter(f'logs/')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6aabaa3d",
   "metadata": {},
   "source": [
    "# Finetune"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c90b8af",
   "metadata": {},
   "outputs": [],
   "source": [
    "for ep in range(51,60):\n",
    "    model.train()\n",
    "    train_lp = 0, 0\n",
    "    c = 0\n",
    "    \n",
    "    for data in train_loader:\n",
    "        x, y, path = data['x'], data['y'], data['path']\n",
    "        rand_idx = np.random.randint(err.shape[0]) # change this to the size of your saved error\n",
    "        x[...,-1:] += torch.from_numpy(err[rand_idx,...])\n",
    "        x, y = x[None,...].to(device), y[None,...][...,:1].to(device)\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        x[...,-1:] = input_normalizer_GLOBAL.encode(x[...,-1:])\n",
    "        pred = model_global(x)\n",
    "        pred = output_normalizer_GLOBAL.decode(pred)\n",
    "        \n",
    "        loss = LPloss(pred.reshape(1, -1), y.reshape(1, -1))\n",
    "        train_lp += loss.item()\n",
    "        \n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        c += 1\n",
    "        \n",
    "        if c%100 ==0:\n",
    "            writer.add_scalars('dp LPloss', {f'{model_str}_{specs}_train': loss.item()}, ep*n_train+c)\n",
    "            print(f'ep: {ep}, iter: {c}, train lp: {loss.item():.4f}')\n",
    "\n",
    "    scheduler.step()\n",
    "    print('----------------------------------------------------------------------')\n",
    "\n",
    "    torch.save(model_global, f'saved_models/{model_str}-{specs}-ep{ep}.pt')"
   ]
  }
 ],
 "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.9.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}