{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using backend: pytorch\n",
      "Other supported backends: tensorflow.compat.v1, tensorflow, jax, paddle.\n",
      "paddle supports more examples now and is recommended.\n"
     ]
    }
   ],
   "source": [
    "import deepxde as dde\n",
    "from deepxde.backend import torch\n",
    "import matplotlib.pyplot  as plt\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def VGM(psi, theta_r, theta_s, alpha, K_s):\n",
    "    \"\"\"\n",
    "    this function calculates volumetric water content and hydraulic conductivity\n",
    "    from matric potential (psi: negative) given VGM (van-Genuchen and Mualem model) parameters\n",
    "    output: (theta, K)\n",
    "    theta: volumetric water content\n",
    "    K: hydraulic conductivity\n",
    "    By default, alpha, n, and K_s are the parameters to be estiated.\n",
    "    \"\"\"\n",
    "    l = 0.5\n",
    "    n = 10.4\n",
    "\n",
    "    m = 1 - 1/n\n",
    "    S_e = (1 + (-alpha*psi)**n)**(-m)\n",
    "    theta =  S_e * (theta_s - theta_r) + theta_r\n",
    "    K = K_s*S_e**l*(1-(1-S_e**(1/m))**m)**2\n",
    "    return theta, K\n",
    "# parameter values\n",
    "alpha = dde.Variable(0.1) #0.79# /cm\n",
    "K_s = dde.Variable(1.08) #1.08 cm/hour\n",
    "theta_s = dde.Variable(0.1) #0.25\n",
    "theta_r = dde.Variable(0.153)\n",
    "L_star = dde.Variable(10.0) # cm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "def pde(x,y):\n",
    "    psi = y[:,0:1]\n",
    "    theta = y[:,1:2]\n",
    "    K = y[:,2:3]\n",
    "    theta_t = dde.grad.jacobian(y,x,i=1,j=1)\n",
    "    psi_z = dde.grad.jacobian(y,x,i=0,j=0)\n",
    "    psi_zz = dde.grad.hessian(y,x,component=0,i=0,j=0)\n",
    "    K_z = dde.grad.jacobian(y,x,i=2,j=0)\n",
    "\n",
    "    residual = theta_t - K_z*psi_z - K*psi_zz - K_z\n",
    "    l = 0.5\n",
    "    n = 10.4\n",
    "\n",
    "    m = 1 - 1/n\n",
    "    S_e = (1 + (-alpha*psi)**n)**(-m)\n",
    "    vg_theta =  S_e * (theta_s - theta_r) + theta_r\n",
    "    vg_k = K_s*S_e**l*(1-(1-S_e**(1/m))**m)**2\n",
    "\n",
    "    residual_1 = K - vg_k\n",
    "    residual_2 = theta - vg_theta\n",
    "    #print(residual)\n",
    "    return [residual,residual_1,residual_2]  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Outputifier:\n",
    "    def __init__(self, f1,f2,f3):\n",
    "        self.f1 = f1\n",
    "        self.f2 = f2\n",
    "        self.f3 = f3\n",
    "\n",
    "    def __call__(self, x,y):\n",
    "        psi = self.f1(y[:,0:1])\n",
    "        theta = self.f2(y[:,1:2])\n",
    "        K = self.f3(y[:,2:3])\n",
    "\n",
    "        return torch.hstack((psi,theta,K))\n",
    "\n",
    "# Example\n",
    "out_transform_org = Outputifier(lambda x: -torch.exp(x), torch.sigmoid, torch.exp)\n",
    "out_transform = Outputifier(lambda x: -torch.abs(x), torch.sigmoid, torch.abs)\n",
    "\n",
    "out_transform2 = Outputifier(lambda x: -torch.pow(x,2), torch.sigmoid, torch.exp)\n",
    "def output_transform(x,y):\n",
    "    psi = -torch.exp(y[:,0:1])\n",
    "    theta = torch.sigmoid(y[:,1:2])\n",
    "    K = torch.exp(y[:,2:3])\n",
    "\n",
    "    return torch.hstack((psi,theta,K))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 0, '$\\\\psi$ [m]')"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 103968 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dz=0.1\n",
    "ProfileDepth=5.0\n",
    "z=np.arange(dz/2.0,ProfileDepth,dz)\n",
    "n=z.size\n",
    "\n",
    "# Grid in time\n",
    "t = np.linspace(0,10,101)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "psi = np.load(\"test1112_HSS.npy\")\n",
    "\n",
    "\n",
    "for i in range(0,t.size-1):\n",
    "    #plt.subplot(121)\n",
    "    plt.plot(psi[i,:],z)\n",
    "    #plt.subplot(122)\n",
    "    #plt.plot(theta[i,:],z)\n",
    "\n",
    "#pl.subplot(121)\n",
    "plt.ylabel('Elevation [m]',fontsize=20)\n",
    "plt.xlabel(r'$\\psi$ [m]',fontsize=20)\n",
    "#pl.subplot(122)#\n",
    "#pl.xlabel(r'$\\theta$ [-]',fontsize=20)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "Dataset = []\n",
    "for i,time_step in enumerate(t):\n",
    "    for j,k in enumerate(z):\n",
    "        Dataset.append([k,time_step,psi[i][j]])\n",
    "\n",
    "Dataset = np.array(Dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train = Dataset[:,0:2]\n",
    "U_train = Dataset[:,2:3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "geom = dde.geometry.Interval(0, 5)\n",
    "timedomain = dde.geometry.TimeDomain(0, 10)\n",
    "geomtime = dde.geometry.GeometryXTime(geom, timedomain)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "observe_y = dde.icbc.PointSetBC(X_train, U_train,component=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = dde.data.TimePDE(\n",
    "    geomtime,\n",
    "    pde,\n",
    "    [observe_y],\n",
    "    num_domain=100,\n",
    "    num_boundary=100,\n",
    "    num_initial=100,\n",
    "    anchors=X_train,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "#import itertools\n",
    "\n",
    "#combs = itertools.combinations_with_replacement([torch.exp, torch.sigmoid, torch.tanh], 3)\n",
    "#for c in combs:\n",
    "#    out_transform = Outputifier(lambda x: -torch.exp(x), torch.sigmoid, torch.exp)\n",
    "\n",
    "#print(list(combs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "layer_size = [2] + [128] * 8 + [3]\n",
    "activation = \"swish\"\n",
    "initializer = \"Glorot uniform\"\n",
    "net = dde.nn.FNN(layer_size, activation, initializer)\n",
    "net.apply_output_transform(out_transform_org)\n",
    "model = dde.Model(data, net)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Compiling model...\n",
      "'compile' took 0.000556 s\n",
      "\n",
      "Training model...\n",
      "\n",
      "Step      Train loss                                  Test loss                                   Test metric\n",
      "0         [7.40e-08, 5.48e-03, 1.61e-01, 8.85e-01]    [7.40e-08, 5.48e-03, 1.61e-01, 8.85e-01]    []  \n",
      "1000      [2.16e-03, 3.95e-02, 3.19e-02, 2.30e-01]    [2.16e-03, 3.95e-02, 3.19e-02, 2.30e-01]    []  \n",
      "2000      [3.75e-03, 1.63e-02, 3.19e-03, 1.11e-01]    [3.75e-03, 1.63e-02, 3.19e-03, 1.11e-01]    []  \n",
      "3000      [8.03e-03, 1.86e-02, 3.83e-03, 7.70e-02]    [8.03e-03, 1.86e-02, 3.83e-03, 7.70e-02]    []  \n",
      "4000      [1.13e-02, 1.16e-02, 4.28e-03, 5.02e-02]    [1.13e-02, 1.16e-02, 4.28e-03, 5.02e-02]    []  \n",
      "5000      [9.88e-03, 5.78e-03, 3.38e-03, 3.49e-02]    [9.88e-03, 5.78e-03, 3.38e-03, 3.49e-02]    []  \n",
      "6000      [7.94e-03, 5.10e-03, 2.49e-03, 2.28e-02]    [7.94e-03, 5.10e-03, 2.49e-03, 2.28e-02]    []  \n",
      "7000      [6.21e-03, 4.42e-03, 2.38e-03, 1.83e-02]    [6.21e-03, 4.42e-03, 2.38e-03, 1.83e-02]    []  \n",
      "8000      [4.98e-03, 3.88e-03, 2.19e-03, 1.49e-02]    [4.98e-03, 3.88e-03, 2.19e-03, 1.49e-02]    []  \n",
      "9000      [3.95e-03, 3.29e-03, 2.07e-03, 1.27e-02]    [3.95e-03, 3.29e-03, 2.07e-03, 1.27e-02]    []  \n",
      "10000     [3.29e-03, 2.82e-03, 1.89e-03, 1.12e-02]    [3.29e-03, 2.82e-03, 1.89e-03, 1.12e-02]    []  \n",
      "11000     [2.84e-03, 2.15e-03, 1.63e-03, 1.02e-02]    [2.84e-03, 2.15e-03, 1.63e-03, 1.02e-02]    []  \n",
      "12000     [2.52e-03, 2.12e-03, 1.38e-03, 8.56e-03]    [2.52e-03, 2.12e-03, 1.38e-03, 8.56e-03]    []  \n",
      "13000     [2.21e-03, 1.83e-03, 1.17e-03, 7.35e-03]    [2.21e-03, 1.83e-03, 1.17e-03, 7.35e-03]    []  \n",
      "14000     [1.92e-03, 1.59e-03, 1.00e-03, 6.14e-03]    [1.92e-03, 1.59e-03, 1.00e-03, 6.14e-03]    []  \n",
      "15000     [1.64e-03, 1.38e-03, 8.61e-04, 4.91e-03]    [1.64e-03, 1.38e-03, 8.61e-04, 4.91e-03]    []  \n",
      "16000     [1.38e-03, 1.19e-03, 7.33e-04, 3.73e-03]    [1.38e-03, 1.19e-03, 7.33e-04, 3.73e-03]    []  \n",
      "17000     [1.14e-03, 9.95e-04, 6.15e-04, 2.83e-03]    [1.14e-03, 9.95e-04, 6.15e-04, 2.83e-03]    []  \n",
      "18000     [9.36e-04, 8.19e-04, 5.09e-04, 2.19e-03]    [9.36e-04, 8.19e-04, 5.09e-04, 2.19e-03]    []  \n",
      "19000     [7.57e-04, 6.66e-04, 4.14e-04, 1.70e-03]    [7.57e-04, 6.66e-04, 4.14e-04, 1.70e-03]    []  \n",
      "20000     [6.00e-04, 5.61e-04, 3.29e-04, 1.31e-03]    [6.00e-04, 5.61e-04, 3.29e-04, 1.31e-03]    []  \n",
      "\n",
      "Best model at step 20000:\n",
      "  train loss: 2.80e-03\n",
      "  test loss: 2.80e-03\n",
      "  test metric: []\n",
      "\n",
      "'train' took 1245.054950 s\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# callbacks for storing results\n",
    "fnamevar = \"variables_test0302.dat\"\n",
    "variable = dde.callbacks.VariableValue([alpha,theta_s,K_s], period=1000, filename=fnamevar)\n",
    "\n",
    "# Compile, train and save model\n",
    "model.compile(\"adam\", lr=1e-5,external_trainable_variables=[alpha,theta_s,K_s])\n",
    "loss_history, train_state = model.train(\n",
    "    iterations=20000 ,display_every=1000,callbacks=[variable]\n",
    ")\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Compiling model...\n",
      "'compile' took 0.002612 s\n",
      "\n",
      "Training model...\n",
      "\n",
      "Step      Train loss                                  Test loss                                   Test metric\n",
      "20000     [6.00e-04, 5.61e-04, 3.29e-04, 1.31e-03]    [6.00e-04, 5.61e-04, 3.29e-04, 1.31e-03]    []  \n",
      "21000     [6.02e-04, 5.32e-04, 3.21e-04, 1.26e-03]    [6.02e-04, 5.32e-04, 3.21e-04, 1.26e-03]    []  \n",
      "22000     [6.00e-04, 5.26e-04, 3.11e-04, 1.14e-03]    [6.00e-04, 5.26e-04, 3.11e-04, 1.14e-03]    []  \n",
      "23000     [5.95e-04, 5.21e-04, 3.00e-04, 1.04e-03]    [5.95e-04, 5.21e-04, 3.00e-04, 1.04e-03]    []  \n",
      "24000     [5.89e-04, 5.23e-04, 2.93e-04, 9.48e-04]    [5.89e-04, 5.23e-04, 2.93e-04, 9.48e-04]    []  \n",
      "25000     [5.82e-04, 5.21e-04, 2.89e-04, 8.77e-04]    [5.82e-04, 5.21e-04, 2.89e-04, 8.77e-04]    []  \n",
      "\n",
      "Best model at step 25000:\n",
      "  train loss: 2.27e-03\n",
      "  test loss: 2.27e-03\n",
      "  test metric: []\n",
      "\n",
      "'train' took 306.687042 s\n",
      "\n"
     ]
    }
   ],
   "source": [
    "model.compile(\"adam\", lr=1e-5)\n",
    "loss_history, train_state = model.train(\n",
    "    iterations=5000,display_every=1000\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "Preds = model.predict(X_train)\n",
    "theta,K = VGM(psi,0.153,0.25,0.79,1.08)\n",
    "theta_preds,k_preds = VGM(Preds[:,0],0.154,0.26,0.79,1.08)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7fb3e54ab8e0>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 11340 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.imshow(psi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 0.98, 'PINN prediction for hydraulic pressure (psi)')"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 157164 characters]",
      "text/plain": [
       "<Figure size 2000x500 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots(1,3,figsize=(20,5),tight_layout=True)\n",
    "\n",
    "im0=ax[0].contourf(X_train[:,1].reshape((101,50)),X_train[:,0].reshape((101,50)),Preds[:,0].reshape((101,50)),levels=[-6,-5.5,-5,-4.5,-4,-3.5,-3,-2.5,-2,-1.5,-1,-0.5,0],vmin=-6,vmax=0)\n",
    "\n",
    "fig.colorbar(im0)\n",
    "\n",
    "im1 = ax[1].contourf(X_train[:,1].reshape((101,50)),X_train[:,0].reshape((101,50)),psi,levels=[-6,-5.5,-5,-4.5,-4,-3.5,-3,-2.5,-2,-1.5,-1,-0.5,0],vmin=-6,vmax=0)\n",
    "fig.colorbar(im1)\n",
    "\n",
    "im2 = ax[2].scatter(X_train[:,1],X_train[:,0],c = np.abs((psi.reshape(5050)-Preds[:,0])),marker=\".\",s=100,vmin=0,vmax=0.5)\n",
    "fig.colorbar(im2)\n",
    "\n",
    "\n",
    "ax[0].set_title(\"PINN prediction\",fontsize=15)\n",
    "ax[0].set_xlabel(\"t\",fontsize=15)\n",
    "ax[0].set_ylabel(\"x\",fontsize=15)\n",
    "\n",
    "ax[1].set_title(\"Ground truth\",fontsize=15)\n",
    "ax[1].set_xlabel(\"t\",fontsize=15)\n",
    "ax[1].set_ylabel(\"x\",fontsize=15)\n",
    "ax[2].set_title(\"Pointwise error\",fontsize=15)\n",
    "ax[2].set_xlabel(\"t\",fontsize=15)\n",
    "ax[2].set_ylabel(\"x\",fontsize=15)\n",
    "\n",
    "\n",
    "\n",
    "ax[0].xaxis.set_tick_params(labelsize=15)\n",
    "ax[0].yaxis.set_tick_params(labelsize=15)\n",
    "ax[1].xaxis.set_tick_params(labelsize=15)\n",
    "ax[1].yaxis.set_tick_params(labelsize=15)\n",
    "ax[2].xaxis.set_tick_params(labelsize=15)\n",
    "ax[2].yaxis.set_tick_params(labelsize=15)\n",
    "\n",
    "fig.suptitle(\"PINN prediction for hydraulic pressure (psi)\",fontsize=15)\n",
    "#fig.savefig(\"vgpsi.pdf\",dpi=500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 0.98, 'PINN prediction for volumetric water content (theta)')"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 197048 characters]",
      "text/plain": [
       "<Figure size 2000x500 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots(1,3,figsize=(20,5),tight_layout=True)\n",
    "\n",
    "im0=ax[0].contourf(X_train[:,1].reshape((101,50)),X_train[:,0].reshape((101,50)),theta_preds.reshape((101,50)),levels=np.linspace(0,0.3,10))\n",
    "\n",
    "fig.colorbar(im0)\n",
    "\n",
    "im1 = ax[1].contourf(X_train[:,1].reshape((101,50)),X_train[:,0].reshape((101,50)),theta,levels=np.linspace(0,0.3,10))\n",
    "fig.colorbar(im1)\n",
    "\n",
    "im2 = ax[2].scatter(X_train[:,1],X_train[:,0],c = np.abs((theta.reshape(5050)-theta_preds.reshape(5050))),marker=\".\",s=100,vmin=0,vmax=0.01)\n",
    "fig.colorbar(im2)\n",
    "\n",
    "\n",
    "ax[0].set_title(\"PINN prediction\",fontsize=15)\n",
    "ax[0].set_xlabel(\"t\",fontsize=15)\n",
    "ax[0].set_ylabel(\"x\",fontsize=15)\n",
    "\n",
    "ax[1].set_title(\"Ground truth\",fontsize=15)\n",
    "ax[1].set_xlabel(\"t\",fontsize=15)\n",
    "ax[1].set_ylabel(\"x\")\n",
    "ax[2].set_title(\"Pointwise error\",fontsize=15)\n",
    "ax[2].set_xlabel(\"t\",fontsize=15)\n",
    "ax[2].set_ylabel(\"x\",fontsize=15)\n",
    "\n",
    "\n",
    "\n",
    "ax[0].xaxis.set_tick_params(labelsize=15)\n",
    "ax[0].yaxis.set_tick_params(labelsize=15)\n",
    "ax[1].xaxis.set_tick_params(labelsize=15)\n",
    "ax[1].yaxis.set_tick_params(labelsize=15)\n",
    "ax[2].xaxis.set_tick_params(labelsize=15)\n",
    "ax[2].yaxis.set_tick_params(labelsize=15)\n",
    "\n",
    "fig.suptitle(\"PINN prediction for volumetric water content (theta)\",fontsize=15)\n",
    "#fig.savefig(\"vgtheta.pdf\",dpi=500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 0.98, 'PINN prediction for hydraulic conductivity (K)')"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 123092 characters]",
      "text/plain": [
       "<Figure size 2000x500 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots(1,3,figsize=(20,5),tight_layout=True)\n",
    "\n",
    "im0=ax[0].contourf(X_train[:,1].reshape((101,50)),X_train[:,0].reshape((101,50)),k_preds.reshape((101,50)),levels=np.linspace(0,0.03,10))\n",
    "\n",
    "fig.colorbar(im0)\n",
    "\n",
    "im1 = ax[1].contourf(X_train[:,1].reshape((101,50)),X_train[:,0].reshape((101,50)),K,levels=np.linspace(0,0.03,10))\n",
    "fig.colorbar(im1)\n",
    "\n",
    "\n",
    "\n",
    "im2 = ax[2].scatter(X_train[:,1],X_train[:,0],c =0.1*np.abs((K.reshape(5050)-k_preds.reshape(5050))),marker=\".\",s=100,vmin=0,vmax=0.01)\n",
    "fig.colorbar(im2)\n",
    "\n",
    "\n",
    "ax[0].set_title(\"PINN prediction\",fontsize=15)\n",
    "ax[0].set_xlabel(\"t\",fontsize=15)\n",
    "ax[0].set_ylabel(\"x\",fontsize=15)\n",
    "\n",
    "ax[1].set_title(\"Ground truth\",fontsize=15)\n",
    "ax[1].set_xlabel(\"t\",fontsize=15)\n",
    "ax[1].set_ylabel(\"x\",fontsize=15)\n",
    "ax[2].set_title(\"Pointwise error\",fontsize=15)\n",
    "ax[2].set_xlabel(\"t\",fontsize=15)\n",
    "ax[2].set_ylabel(\"x\",fontsize=15)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "ax[0].xaxis.set_tick_params(labelsize=15)\n",
    "ax[0].yaxis.set_tick_params(labelsize=15)\n",
    "ax[1].xaxis.set_tick_params(labelsize=15)\n",
    "ax[1].yaxis.set_tick_params(labelsize=15)\n",
    "ax[2].xaxis.set_tick_params(labelsize=15)\n",
    "ax[2].yaxis.set_tick_params(labelsize=15)\n",
    "\n",
    "fig.suptitle(\"PINN prediction for hydraulic conductivity (K)\",fontsize=15)\n",
    "#fig.savefig(\"vgK.pdf\",dpi=500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.17 ('newtorch')",
   "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.17"
  },
  "vscode": {
   "interpreter": {
    "hash": "813e009bea45919a754463d31638190cf3c033fdb5fd5c2966a42653051a1217"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
