{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Take three outputs of neural network. psi,K and Theta. And see how they behave. \n",
    "\n",
    "Basically output psi, theta and K and also add loss terms relating the Outputs of NNs.\n",
    "\n"
   ]
  },
  {
   "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 Gardner(psi, theta_r, theta_s, alpha, K_s):\n",
    "    K = K_s * torch.exp(alpha * psi)\n",
    "    theta = theta_r + (theta_s - theta_r) * torch.exp(alpha * psi)\n",
    "    return theta, K\n",
    "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) # /cm\n",
    "K_s = dde.Variable(0.1) # cm/hour\n",
    "theta_s = dde.Variable(0.1)\n",
    "theta_r = dde.Variable(0.06)\n",
    "L_star = dde.Variable(10.0) # cm\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot3D(x,t,y,name=\"F\"):\n",
    "  x_plot =x.squeeze(1) \n",
    "  t_plot =t.squeeze(1)\n",
    "  X,T= torch.meshgrid(x_plot,t_plot)\n",
    "  F_xt = y\n",
    "  fig,ax=plt.subplots(1,1)\n",
    "  cp = ax.contourf(T,X, F_xt,20,cmap=\"terrain\")\n",
    "  fig.colorbar(cp) # Add a colorbar to a plot\n",
    "  ax.set_title(name)\n",
    "  ax.set_xlabel('t')\n",
    "  ax.set_ylabel('x')\n",
    "  plt.show()\n",
    "  ax = plt.axes(projection='3d')\n",
    "  ax.plot_surface(T.numpy(), X.numpy(), F_xt.numpy(),cmap=\"terrain\")\n",
    "  ax.set_xlabel('t')\n",
    "  ax.set_ylabel('x')\n",
    "  ax.set_zlabel(name)\n",
    "  plt.show()\n",
    "     \n",
    "def plot3D_Matrix(x,t,y,name=\"F\"):\n",
    "  X,T= x,t\n",
    "  F_xt = y\n",
    "  fig,ax=plt.subplots(1,1)\n",
    "  cp = ax.contourf(T,X, F_xt,20,cmap=\"rainbow\")\n",
    "  fig.colorbar(cp) # Add a colorbar to a plot\n",
    "  ax.set_title(name)\n",
    "  ax.set_xlabel('t')\n",
    "  ax.set_ylabel('x')\n",
    "  plt.show()\n",
    "  ax = plt.axes(projection='3d')\n",
    "  ax.plot_surface(T.numpy(), X.numpy(), F_xt.numpy(),cmap=\"rainbow\")\n",
    "  ax.set_xlabel('t')\n",
    "  ax.set_ylabel('x')\n",
    "  ax.set_zlabel(name)\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "\n",
    "    Gardner_K = K_s * torch.exp(alpha * psi)\n",
    "    residual_1 = K - Gardner_K\n",
    "\n",
    "    Gardner_theta = theta_r + (theta_s - theta_r) * torch.exp(alpha * psi)\n",
    "\n",
    "    residual_2 = theta - Gardner_theta\n",
    "    \n",
    "    return [residual,residual_1,residual_2]  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "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": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/vemuri/anaconda3/envs/newtorch/lib/python3.9/site-packages/torch/functional.py:504: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at /opt/conda/conda-bld/pytorch_1682343964576/work/aten/src/ATen/native/TensorShape.cpp:3483.)\n",
      "  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 22436 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 97876 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "x = np.linspace(0,10,101).reshape((101,1))\n",
    "t = np.linspace(0,10,101).reshape((101,1))\n",
    "usol = np.load(\"Srivastava_psi_homogeneous.npy\")\n",
    "plot3D(torch.from_numpy(x),torch.from_numpy(t),torch.from_numpy(usol),name=\"psi\") \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 22436 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 97876 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 30516 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 111132 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 26852 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 103520 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def Gardner_np(psi, theta_r, theta_s, alpha, K_s):\n",
    "    K = K_s * np.exp(alpha * psi)\n",
    "    theta = theta_r + (theta_s - theta_r) * np.exp(alpha * psi)\n",
    "    return theta, K\n",
    "theta_True,K_True = Gardner_np(usol,0.06,0.4,1.0,1.0)\n",
    "theta_pred,K_pred = Gardner_np(usol,0.06,0.4,1.0,1.0)\n",
    "\n",
    "plot3D(torch.from_numpy(x),torch.from_numpy(t),torch.from_numpy(usol),name=\"psi\") \n",
    "\n",
    "\n",
    "plot3D(torch.from_numpy(x),torch.from_numpy(t),torch.from_numpy(theta_pred),name=\"Theta_True_SandyLoam\") \n",
    "\n",
    "plot3D(torch.from_numpy(x),torch.from_numpy(t),torch.from_numpy(K_pred),name=\"K_True_SandyLoam\") \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.cuda.is_available()\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "X, T = np.meshgrid(x,t)   \n",
    "X_true = np.hstack((X.flatten()[:,None], T.flatten()[:,None]))\n",
    "lb = X_true[0]  # [-1. 0.]\n",
    "ub = X_true[-1] # [1.  0.99]\n",
    "U_true = usol.flatten('F')[:,None] #Fortran style (Column Major)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "We have 10201 points. We will select 10000 points to train our model.\n"
     ]
    }
   ],
   "source": [
    "# Obtain random points for interior\n",
    "id_f = np.random.choice(101*101, 10000, replace=False)# Randomly chosen points for Interior\n",
    "X_train = X_true[id_f]\n",
    "U_train= U_true[id_f]\n",
    "\n",
    "print(\"We have\",101*101,\"points. We will select\",X_train.shape[0],\"points to train our model.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "geom = dde.geometry.Interval(0, 10)\n",
    "timedomain = dde.geometry.TimeDomain(0, 10)\n",
    "geomtime = dde.geometry.GeometryXTime(geom, timedomain)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "observe_y = dde.icbc.PointSetBC(X_train, U_train,component=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "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": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "layer_size = [2] + [50] * 6 + [3]\n",
    "activation = \"swish\"\n",
    "initializer = \"Glorot uniform\"\n",
    "net = dde.nn.FNN(layer_size, activation, initializer)\n",
    "model = dde.Model(data, net)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Compiling model...\n",
      "'compile' took 0.000753 s\n",
      "\n",
      "Training model...\n",
      "\n",
      "Step      Train loss                                  Test loss                                   Test metric\n",
      "0         [4.29e-04, 8.84e-01, 1.26e-02, 3.58e-01]    [4.29e-04, 8.84e-01, 1.26e-02, 3.58e-01]    []  \n",
      "1000      [4.44e-04, 2.09e-04, 2.17e-04, 5.52e-04]    [4.44e-04, 2.09e-04, 2.17e-04, 5.52e-04]    []  \n",
      "2000      [4.31e-04, 6.77e-05, 2.16e-04, 2.70e-04]    [4.31e-04, 6.77e-05, 2.16e-04, 2.70e-04]    []  \n",
      "3000      [4.13e-04, 5.01e-05, 2.29e-04, 1.46e-04]    [4.13e-04, 5.01e-05, 2.29e-04, 1.46e-04]    []  \n",
      "4000      [4.07e-04, 4.71e-05, 2.13e-04, 8.95e-05]    [4.07e-04, 4.71e-05, 2.13e-04, 8.95e-05]    []  \n",
      "5000      [3.95e-04, 7.57e-05, 2.08e-04, 6.83e-05]    [3.95e-04, 7.57e-05, 2.08e-04, 6.83e-05]    []  \n",
      "6000      [4.08e-04, 4.46e-05, 2.08e-04, 6.32e-05]    [4.08e-04, 4.46e-05, 2.08e-04, 6.32e-05]    []  \n",
      "7000      [4.04e-04, 4.64e-05, 2.06e-04, 5.79e-05]    [4.04e-04, 4.64e-05, 2.06e-04, 5.79e-05]    []  \n",
      "8000      [4.04e-04, 5.67e-05, 2.00e-04, 6.16e-05]    [4.04e-04, 5.67e-05, 2.00e-04, 6.16e-05]    []  \n",
      "9000      [3.91e-04, 4.76e-05, 2.12e-04, 6.01e-05]    [3.91e-04, 4.76e-05, 2.12e-04, 6.01e-05]    []  \n",
      "10000     [4.01e-04, 5.47e-05, 2.01e-04, 5.54e-05]    [4.01e-04, 5.47e-05, 2.01e-04, 5.54e-05]    []  \n",
      "11000     [4.51e-04, 5.02e-05, 1.86e-04, 9.44e-05]    [4.51e-04, 5.02e-05, 1.86e-04, 9.44e-05]    []  \n",
      "12000     [4.03e-04, 4.60e-05, 2.02e-04, 5.43e-05]    [4.03e-04, 4.60e-05, 2.02e-04, 5.43e-05]    []  \n",
      "13000     [4.02e-04, 4.67e-05, 2.01e-04, 5.38e-05]    [4.02e-04, 4.67e-05, 2.01e-04, 5.38e-05]    []  \n",
      "14000     [3.96e-04, 4.74e-05, 2.06e-04, 5.35e-05]    [3.96e-04, 4.74e-05, 2.06e-04, 5.35e-05]    []  \n",
      "15000     [3.90e-04, 4.94e-05, 2.10e-04, 5.87e-05]    [3.90e-04, 4.94e-05, 2.10e-04, 5.87e-05]    []  \n",
      "16000     [3.67e-04, 7.46e-05, 2.77e-04, 3.57e-04]    [3.67e-04, 7.46e-05, 2.77e-04, 3.57e-04]    []  \n",
      "17000     [3.95e-04, 4.72e-05, 2.07e-04, 5.23e-05]    [3.95e-04, 4.72e-05, 2.07e-04, 5.23e-05]    []  \n",
      "18000     [3.96e-04, 4.63e-05, 2.13e-04, 7.34e-05]    [3.96e-04, 4.63e-05, 2.13e-04, 7.34e-05]    []  \n",
      "19000     [3.96e-04, 4.59e-05, 2.09e-04, 5.56e-05]    [3.96e-04, 4.59e-05, 2.09e-04, 5.56e-05]    []  \n",
      "20000     [3.95e-04, 4.67e-05, 2.08e-04, 5.31e-05]    [3.95e-04, 4.67e-05, 2.08e-04, 5.31e-05]    []  \n",
      "\n",
      "Best model at step 17000:\n",
      "  train loss: 7.02e-04\n",
      "  test loss: 7.02e-04\n",
      "  test metric: []\n",
      "\n",
      "'train' took 729.714602 s\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# callbacks for storing results\n",
    "fnamevar = \"variables.dat\"\n",
    "variable = dde.callbacks.VariableValue([alpha], period=100, filename=fnamevar)\n",
    "\n",
    "# Compile, train and save model\n",
    "model.compile(\"adam\", lr=1e-3)\n",
    "loss_history, train_state = model.train(\n",
    "    iterations=20000 ,display_every=1000, disregard_previous_best=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "Preds = model.predict(X_true)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[-5.96188009e-03, -9.58561003e-02, -1.84753314e-01, ...,\n",
       "        -2.24285746e+00, -2.20015240e+00, -2.14840364e+00],\n",
       "       [-3.55165638e-03, -9.31254104e-02, -1.81688577e-01, ...,\n",
       "        -9.73647237e-01, -8.80911887e-01, -7.92154908e-01],\n",
       "       [-1.82352215e-03, -9.11483020e-02, -1.79455280e-01, ...,\n",
       "        -6.52529955e-01, -6.03442550e-01, -5.57623684e-01],\n",
       "       ...,\n",
       "       [-2.07765214e-03, -1.13452598e-02, -1.98374484e-02, ...,\n",
       "        -1.04790568e-01, -1.04730904e-01, -1.04670718e-01],\n",
       "       [-2.36764736e-03, -1.16149550e-02, -2.00875811e-02, ...,\n",
       "        -1.04784712e-01, -1.04724139e-01, -1.04662813e-01],\n",
       "       [-2.67046690e-03, -1.18972370e-02, -2.03502886e-02, ...,\n",
       "        -1.04779176e-01, -1.04717717e-01, -1.04655728e-01]], dtype=float32)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "psi_preds.reshape((101,101))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 139732 characters]",
      "text/plain": [
       "<Figure size 2000x500 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "psi_preds= Preds[:,0]\n",
    "theta_preds= Preds[:,1]\n",
    "K_preds = Preds[:,2]\n",
    "\n",
    "\n",
    "#psi_reshaped_preds = psi_preds.reshape((100,100))\n",
    "\n",
    "fig, ax = plt.subplots(1,3,figsize=(20,5))\n",
    "\n",
    "im0=ax[0].contourf(T,X,psi_preds.reshape((101,101)),vmin=-2.4,vmax=0.0)\n",
    "fig.colorbar(im0)\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",
    "im1 = ax[1].contourf(T,X,U_true.reshape((101,101)),vmin=-2.4,vmax=0.0)\n",
    "fig.colorbar(im1)\n",
    "\n",
    "im2 = ax[2].scatter(X_true[:,1],X_true[:,0],c =np.abs((U_true[:,0]-psi_preds)),marker=\".\",s=None)\n",
    "fig.colorbar(im2)\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",
    "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",
    "fig.suptitle(\"PINN prediction for Hydraulic Pressure\",fontsize=15)\n",
    "fig.savefig(\"GardnerPsi.pdf\",dpi=500)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 151052 characters]",
      "text/plain": [
       "<Figure size 2000x500 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import numpy as np\n",
    "def Gardner(psi, theta_r=0.06, theta_s=0.4, alpha=1.00, K_s=1.0):\n",
    "    K = K_s * np.exp(alpha * psi)\n",
    "    theta = theta_r + (theta_s - theta_r) * np.exp(alpha * psi)\n",
    "    return theta, K\n",
    "\n",
    "theta, K = Gardner(U_true)\n",
    "theta_pred,K_pred =Gardner (psi_preds) \n",
    "fig, ax = plt.subplots(1,3,figsize=(20,5))\n",
    "\n",
    "im0=ax[0].contourf(T,X,theta_pred.reshape((101,101)),vmin=0,vmax=0.4)\n",
    "fig.colorbar(im0)\n",
    "\n",
    "\n",
    "plt.xticks(fontsize=15)\n",
    "plt.yticks(fontsize=15)\n",
    "im1 = ax[1].contourf(T,X,theta.reshape((101,101)),vmin=0,vmax=0.4)\n",
    "fig.colorbar(im1)\n",
    "\n",
    "im2 = ax[2].scatter(X_true[:,1],X_true[:,0],c =np.abs((theta[:,0]-0.99*theta_pred)),marker=\".\",s=None)\n",
    "fig.colorbar(im2)\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",
    "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",
    "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",
    "fig.suptitle(\"PINN prediction for Volumetric Water Content (Theta)\",fontsize=15)\n",
    "fig.savefig(\"GardnerTheta.pdf\",dpi=500)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 161092 characters]",
      "text/plain": [
       "<Figure size 2000x500 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots(1,3,figsize=(20,5))\n",
    "\n",
    "im0=ax[0].contourf(T,X,K_preds.reshape((101,101)),vmin=0,vmax=1.0)\n",
    "fig.colorbar(im0)\n",
    "\n",
    "\n",
    "im1 = ax[1].contourf(T,X,K.reshape((101,101)),vmin=0,vmax=1.0)\n",
    "fig.colorbar(im1)\n",
    "\n",
    "im2 = ax[2].scatter(X_true[:,1],X_true[:,0],c =np.abs((K[:,0]-K_preds)),marker=\".\",s=None)\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",
    "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(\"GardnerK.pdf\",dpi=500)"
   ]
  },
  {
   "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": 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"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "813e009bea45919a754463d31638190cf3c033fdb5fd5c2966a42653051a1217"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
