{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#Simple 1D Richards' Equation Model\n",
    "\n",
    "This is a 1D Richards' equation model, written by Andrew Ireson, 3 November 2015.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Import all of the basic libraries (you will always need these)\n",
    "%matplotlib inline\n",
    "from matplotlib import pyplot as pl\n",
    "import numpy as np\n",
    "\n",
    "# Import a library that contains soil moisture properties and functions\n",
    "import vanGenuchten as vg\n",
    "\n",
    "# Import ODE solvers\n",
    "from scipy.interpolate import interp1d\n",
    "from scipy.integrate import odeint\n",
    "\n",
    "# Select which soil properties to use\n",
    "p=vg.HygieneSandstone()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Richards equation solver\n",
    "# This is a function that calculated the right hand side of Richards' equation. You\n",
    "# will not need to modify this function, unless you are doing something advanced. \n",
    "# This block of code must be executed so that the function can be later called.\n",
    "\n",
    "def RichardsModel(psi,t,dz,n,p,vg,qTop,qBot,psiTop,psiBot):\n",
    "       \n",
    "    # Basic properties:\n",
    "    C=vg.CFun(psi,p)\n",
    "   \n",
    "    # initialize vectors:\n",
    "    q=np.zeros(n+1)\n",
    "    \n",
    "    # Upper boundary\n",
    "    if qTop == []:\n",
    "        KTop=vg.KFun(np.zeros(1)+psiTop,p)\n",
    "        q[n]=-KTop*((psiTop-psi[n-1])/dz*2+1)\n",
    "    else:\n",
    "        q[n]=qTop\n",
    "    \n",
    "    # Lower boundary\n",
    "    if qBot == []:\n",
    "        if psiBot == []:\n",
    "            # Free drainage\n",
    "            KBot=vg.KFun(np.zeros(1)+psi[0],p)\n",
    "            q[0]=-KBot\n",
    "        else:\n",
    "            # Type 1 boundary\n",
    "            KBot=vg.KFun(np.zeros(1)+psiBot,p)\n",
    "            q[0]=-KBot*((psi[0]-psiBot)/dz*2+1.0)    \n",
    "    else:\n",
    "        # Type 2 boundary\n",
    "        q[0]=qBot\n",
    "    \n",
    "    # Internal nodes\n",
    "    i=np.arange(0,n-1)\n",
    "    Knodes=vg.KFun(psi,p)\n",
    "    Kmid=(Knodes[i+1]+Knodes[i])/2.0\n",
    "    \n",
    "    j=np.arange(1,n)\n",
    "    q[j]=-Kmid*((psi[i+1]-psi[i])/dz+1.0)\n",
    "    \n",
    "    \n",
    "    \n",
    "    # Continuity\n",
    "    i=np.arange(0,n)\n",
    "    dpsidt=(-(q[i+1]-q[i])/dz)/C\n",
    "    \n",
    "    return dpsidt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 0, '$\\\\psi$')"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 40524 characters]",
      "text/plain": [
       "<Figure size 500x1000 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "psi = np.linspace(-10,0)\n",
    "theta = vg.thetaFun(psi,p)\n",
    "C=vg.CFun(psi,p)\n",
    "K=vg.KFun(psi,p)\n",
    "\n",
    "pl.rcParams['figure.figsize'] = (5.0, 10.0)\n",
    "pl.subplot(311)\n",
    "pl.plot(psi,theta)\n",
    "pl.ylabel(r'$\\theta$', fontsize=20)\n",
    "pl.subplot(312)\n",
    "pl.plot(psi,C)\n",
    "pl.ylabel(r'$C$',fontsize=20)\n",
    "pl.subplot(313)\n",
    "pl.plot(psi,K)\n",
    "pl.ylabel(r'$K$', fontsize=20)\n",
    "pl.xlabel(r'$\\psi$', fontsize=20)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_3081757/2620229187.py:26: DeprecationWarning: Conversion of an array with ndim > 0 to a scalar is deprecated, and will error in future. Ensure you extract a single element from your array before performing this operation. (Deprecated NumPy 1.25.)\n",
      "  q[0]=-KBot\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model run successfully\n"
     ]
    }
   ],
   "source": [
    "# This block of code sets up and runs the model\n",
    "\n",
    "# Boundary conditions\n",
    "qTop=-0.01\n",
    "qBot=[]\n",
    "psiTop=[]\n",
    "psiBot=[]\n",
    "\n",
    "# Grid in space\n",
    "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",
    "# Initial conditions\n",
    "psi0=-z\n",
    "\n",
    "# Solve\n",
    "psi=odeint(RichardsModel,psi0,t,args=(dz,n,p,vg,qTop,qBot,psiTop,psiBot),mxstep=5000000);\n",
    "\n",
    "print(\"Model run successfully\")    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Post process model output to get useful information\n",
    "\n",
    "# Get water content\n",
    "theta=vg.thetaFun(psi,p)\n",
    "\n",
    "# Get total profile storage\n",
    "S=theta.sum(axis=1)*dz\n",
    "\n",
    "# Get change in storage [dVol]\n",
    "dS=np.zeros(S.size)\n",
    "dS[1:]=np.diff(S)/(t[1]-t[0])\n",
    "\n",
    "# Get infiltration flux\n",
    "if qTop == []:\n",
    "    KTop=vg.KFun(np.zeros(1)+psiTop,p)\n",
    "    qI=-KTop*((psiTop-psi[:,n-1])/dz*2+1)\n",
    "else:\n",
    "    qI=np.zeros(t.size)+qTop\n",
    "    \n",
    "# Get discharge flux\n",
    "if qBot == []:\n",
    "    if psiBot == []:\n",
    "        # Free drainage\n",
    "        KBot=vg.KFun(psi[:,0],p)\n",
    "        qD=-KBot\n",
    "    else:\n",
    "        # Type 1 boundary\n",
    "        KBot=vg.KFun(np.zeros(1)+psiBot,p)\n",
    "        qD=-KBot*((psi[:,0]-psiBot)/dz*2+1.0)\n",
    "else:\n",
    "    qD=np.zeros(t.size)+qBot\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 0, '$\\\\theta$ [-]')"
      ]
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 298108 characters]",
      "text/plain": [
       "<Figure size 1000x1000 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot vertical profiles\n",
    "pl.rcParams['figure.figsize'] = (10.0, 10.0)\n",
    "for i in range(0,t.size-1):\n",
    "    pl.subplot(121)\n",
    "    pl.plot(psi[i,:],z)\n",
    "    pl.subplot(122)\n",
    "    pl.plot(theta[i,:],z)\n",
    "\n",
    "pl.subplot(121)\n",
    "pl.ylabel('Elevation [m]',fontsize=20)\n",
    "pl.xlabel(r'$\\psi$ [m]',fontsize=20)\n",
    "pl.subplot(122)\n",
    "pl.xlabel(r'$\\theta$ [-]',fontsize=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.0, 0.02)"
      ]
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 66260 characters]",
      "text/plain": [
       "<Figure size 1000x1000 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot timeseries\n",
    "dt = t[2]-t[1]\n",
    "pl.plot(t,dS,label='Change in storage')\n",
    "#pl.hold(True)\n",
    "pl.plot(t,-qI,label='Infiltration')\n",
    "pl.plot(t,-qD,label='Discharge')\n",
    "pl.legend(loc=\"upper left\")\n",
    "pl.ylim((0,0.02))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {
    "collapsed": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(101, 50)"
      ]
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "psi.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[-0.05      , -0.15      , -0.25      , ..., -4.75      ,\n",
       "        -4.85      , -4.95      ],\n",
       "       [-1.21320422, -1.21942172, -1.228691  , ..., -3.46875612,\n",
       "        -2.01171495, -1.6235221 ],\n",
       "       [-1.30025791, -1.30650948, -1.31501356, ..., -2.21954275,\n",
       "        -1.70316917, -1.53958328],\n",
       "       ...,\n",
       "       [-1.48483418, -1.4848332 , -1.48483245, ..., -1.4848316 ,\n",
       "        -1.4848316 , -1.4848316 ],\n",
       "       [-1.48483293, -1.48483243, -1.48483204, ..., -1.4848316 ,\n",
       "        -1.4848316 , -1.4848316 ],\n",
       "       [-1.48483229, -1.48483203, -1.48483183, ..., -1.4848316 ,\n",
       "        -1.4848316 , -1.4848316 ]])"
      ]
     },
     "execution_count": 66,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "psi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.io import savemat\n",
    "mdic = {\"num_psi\": psi, \"label\": \"Hygiene Sand Stone\"}\n",
    "savemat('/users/vemuri/Desktop/RichardPINNs/test1112_HSS.mat',mdic)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.save('/users/vemuri/Desktop/RichardPINNs/test1112_HSS.npy', psi, allow_pickle=True, fix_imports=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [
    {
     "ename": "IndexError",
     "evalue": "arrays used as indices must be of integer (or boolean) type",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mIndexError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[68], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mpsi\u001b[49m\u001b[43m[\u001b[49m\u001b[43mt\u001b[49m\u001b[43m]\u001b[49m,\n",
      "\u001b[0;31mIndexError\u001b[0m: arrays used as indices must be of integer (or boolean) type"
     ]
    }
   ],
   "source": [
    "psi[t],"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.12 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "3.10.12"
  },
  "vscode": {
   "interpreter": {
    "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
