{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "#imports\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "#define constants\n",
    "\n",
    "g       = 9.81              # gravitational acceleration [m/s²]\n",
    "\n",
    "L       = 100              # length of pipeline [m]\n",
    "rho     = 1000              # density of water [kg/m³]\n",
    "D       = 1                 # pipe diameter \n",
    "Q0      = 2                 # initial flow in whole pipe [m³/s]\n",
    "h       = 20                # water level in upstream reservoir [m]\n",
    "n       = 10                # number of pipe segments in discretization\n",
    "nt      = 150              # number of time steps\n",
    "f_coeff = 0.1              # lambda = 0.01 Friction loss coefficient [m]\n",
    "c       = 400               # propagation velocity of the pressure wave\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "# preparing the discretization and initial conditions\n",
    "\n",
    "dx      = L/n                       # length of each pipe segment\n",
    "dt      = dx/c                      # timestep according to method of characterisitics\n",
    "nn      = n+1                       # number of nodes\n",
    "pl_vec  = np.arange(0,nn*dx,dx)     # pl = pipe-length. position of the nodes on the pipeline\n",
    "t_vec   = np.arange(0,(nt)*dt,dt) # time vector\n",
    "\n",
    "v0 = Q0/(D**2/4*np.pi)\n",
    "p0 = (rho*g*h-v0**2*rho/2)\n",
    "\n",
    "# storage vectors for old parameters\n",
    "v_old = np.full(nn,v0)\n",
    "p_old = p0-(f_coeff*pl_vec/D*rho/2*v0**2)   # ref Wikipedia: Rohrreibungszahls\n",
    "\n",
    "# storage vectors for new parameters\n",
    "v_new = np.zeros_like(v_old)\n",
    "p_new = np.zeros_like(p_old)\n",
    "\n",
    "# storage vector for time evolution of parameters at node nn (at reservoir)\n",
    "p_nn = np.zeros_like(t_vec)\n",
    "v_nn = np.zeros_like(t_vec)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib qt\n",
    "# time loop\n",
    "\n",
    "fig = plt.figure()\n",
    "ax1 = fig.add_subplot(111)\n",
    "lo1, = ax1.plot(pl_vec,np.full_like(pl_vec,p0),marker='.')\n",
    "ax1.set_ylim([-20*p0,20*p0])\n",
    "\n",
    "for it in range(nt):\n",
    "    # set boundary conditions\n",
    "    v_new[-1]   = 0       # in front of the instantaneously closing valve, the velocity is 0\n",
    "    p_new[0]    = p0      # hydrostatic pressure from the reservoir\n",
    "\n",
    "    # calculate the new parameters at first and last node\n",
    "    v_new[0]    = v_old[1]+1/(rho*c)*(p0-p_old[1])-f_coeff*dt/(2*D)*abs(v_old[1])*v_old[1]\n",
    "    p_new[-1]   = p_old[-2]+rho*c*v_old[-2]-rho*c*f_coeff*dt/(2*D) *abs(v_old[-2])*v_old[-2]\n",
    "\n",
    "    # calculate parameters at second to second-to-last nodes \n",
    "    #equation 2-30 plus 2-31 (and refactor for v_i^j+1) in block 2\n",
    "\n",
    "    for i in range(1,nn-1):\n",
    "        v_new[i] = 0.5*(v_old[i-1]+v_old[i+1])+0.5/(rho*c)*(p_old[i-1]-p_old[i+1]) \\\n",
    "            -f_coeff*dt/(4*D)*(abs(v_old[i-1])*v_old[i-1]+abs(v_old[i+1])*v_old[i+1])\n",
    "\n",
    "        p_new[i] = 0.5*rho*c*(v_old[i-1]-v_old[i+1])+0.5*(p_old[i-1]+p_old[i+1]) \\\n",
    "            -rho*c*f_coeff*dt/(4*D)*(abs(v_old[i-1])*v_old[i-1]-abs(v_old[i+1])*v_old[i+1])\n",
    "    \n",
    "    lo1.set_xdata(pl_vec)\n",
    "    lo1.set_ydata(p_new)\n",
    "    ax1.set_title(str(t_vec[it]))\n",
    "    fig.canvas.draw()\n",
    "    plt.pause(0.05)\n",
    "\n",
    "    # store parameters of node 0 (at reservoir)\n",
    "    p_nn[it] = p_old[-1]\n",
    "    v_nn[it] = v_old[-1]\n",
    "\n",
    "    # prepare for next loop\n",
    "    p_old = p_new\n",
    "    v_old = v_new\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x14c45f0bfd0>]"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "plt.plot(v_nn)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.13 ('Georg_DT_Slot3')",
   "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.8.13"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "84fb123bdc47ab647d3782661abcbe80fbb79236dd2f8adf4cef30e8755eb2cd"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}