{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "source": [
    "# 2D Transient simulation -- checkpoint your work\n",
    "\n",
    "## Learning Goals\n",
    "\n",
    "- Write numerical state out to an npz file for later use\n",
    "- Write function and class definitions out to a python module for later import\n",
    "\n",
    "## Context\n",
    "\n",
    "The 2D transient simulation is approach a level of complexity that makes keeping it all\n",
    "in a single notebook problematic.  Here is a demonstration of how we can modify the\n",
    "original notebook to save our work for later use.  We will export all of our functions\n",
    "and classes to a python module called build_2D_matrix.py, and export all of our numeric\n",
    "data (boundary conditions, initial concentration, etc.) to a numpy.npz file called \n",
    "savestate.npz\n",
    "\n",
    "We will read these back into the `2D_assignment_restart.ipynb` notebook\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "import matplotlib.cm as cmap\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from mpl_toolkits.axes_grid1 import AxesGrid\n",
    "from numpy.testing import assert_allclose\n",
    "from scipy import special"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Our functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this function deals with harmonic averaging when diffusion is not the same everywhere.\n",
    "# It doesn't change anything when diffusion is homogeneous but you can try to see how it affects the behavior.\n",
    "\n",
    "\n",
    "def avg(Di, Dj):\n",
    "    \"\"\"\n",
    "    Computes the harmonic average between two values Di and Dj\n",
    "    Returns 0 if either of them is zero\n",
    "    \"\"\"\n",
    "    if (Di * Dj) == 0:\n",
    "        return 0\n",
    "    else:\n",
    "        return 2 / (1 / Di + 1 / Dj)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def index_to_row_col(ind, nrows, ncol):\n",
    "    \"\"\"\n",
    "    in a 2D array, returns the row and column value\n",
    "    associated with a 1D index\n",
    "    Bottom left is index is zero (0-th row, 0-th column)\n",
    "    while index one is the 0-th row and 1st column\n",
    "    \"\"\"\n",
    "    if ind > nrows * ncol - 1:\n",
    "        return 0\n",
    "\n",
    "    row = int(np.floor(ind / ncol))\n",
    "    col = int(ind - row * ncol)\n",
    "    return row, col"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_2D_matrix(bc_dict, problem, D_matrix, Qsource):\n",
    "    \"\"\"\n",
    "    Constructs a coefficient matrix A and an array b corresponding to the system Ac = b\n",
    "    This system corresponds either to a 1D or 2D problem\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    bc_dict: dict\n",
    "       dictionary with Boundary_Def objects defining the boundary conditions\n",
    "    D_matrix: (float vector if 1D or matrix if 2D)\n",
    "        values of the diffusion coefficient at each grid point(dm^2/day)\n",
    "        if 2D, dimension is [problem.ny, problem.nx]\n",
    "    width_x: (float)\n",
    "        x-extent of the domain (dm)\n",
    "    width_y: (float)\n",
    "        y-extent of the domain (dm)\n",
    "    poro (float)\n",
    "        porosity value\n",
    "    Qsource: (float array)\n",
    "      volumetric source term (mg/L/day)\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "\n",
    "    the_matrix, rhs: tuple\n",
    "       where the_matrix=A and rhs =b\n",
    "       in the discretized diffusion problem\n",
    "       Ax=b\n",
    "    \"\"\"\n",
    "    number_of_rows = problem.ny\n",
    "    number_of_cols = problem.nx\n",
    "    n = problem.nx * problem.ny\n",
    "    is1D = False\n",
    "    if number_of_rows == 1 or number_of_cols == 1:\n",
    "        is1D = True\n",
    "        number_of_cols = n\n",
    "    the_matrix = np.zeros((n, n))\n",
    "    rhs = np.zeros(n)\n",
    "\n",
    "    if is1D:\n",
    "        dx = max(problem.wx, problem.wy) / (max(problem.ny, problem.nx) - 1)\n",
    "        coef_x = problem.poro / dx / dx\n",
    "    else:\n",
    "        dx = problem.wx / (problem.ny - 1)\n",
    "        dy = problem.wy / (problem.nx - 1)\n",
    "        coef_x = problem.poro / dx / dx\n",
    "        coef_y = problem.poro / dy / dy\n",
    "\n",
    "    for ind in range(n):\n",
    "        if is1D:\n",
    "            j = ind\n",
    "            i = -1\n",
    "        else:\n",
    "            i, j = index_to_row_col(ind, number_of_rows, number_of_cols)\n",
    "        if j == 0:  # WEST BOUNDARY\n",
    "            if bc_dict[\"west\"].btype == \"const\":\n",
    "                rhs[ind] = bc_dict[\"west\"].val\n",
    "                the_matrix[ind, ind] = 1\n",
    "            else:  # flux boundary condition\n",
    "                the_matrix[ind, ind] = 1\n",
    "                the_matrix[ind, ind + 1] = -1\n",
    "                rhs[ind] = bc_dict[\"west\"].val / dx\n",
    "\n",
    "        elif j == number_of_cols - 1:  # EAST BOUNDARY\n",
    "            if bc_dict[\"east\"].btype == \"const\":\n",
    "                rhs[ind] = bc_dict[\"east\"].val\n",
    "                the_matrix[ind, ind] = 1\n",
    "            else:  # flux boundary condition\n",
    "                the_matrix[ind, ind] = 1\n",
    "                the_matrix[ind, ind - 1] = -1\n",
    "                rhs[ind] = bc_dict[\"east\"].val / dx\n",
    "        elif i == 0 and problem.ny > 1:  # SOUTH BOUNDARY\n",
    "            if bc_dict[\"south\"].btype == \"const\":\n",
    "                rhs[ind] = bc_dict[\"south\"].val\n",
    "                the_matrix[ind, ind] = 1\n",
    "            else:  # flux boundary condition\n",
    "                the_matrix[ind, ind] = 1\n",
    "                the_matrix[ind, ind + number_of_cols] = -1\n",
    "                rhs[ind] = bc_dict[\"south\"].val / dy\n",
    "\n",
    "        elif i == number_of_rows - 1 and problem.ny > 1:  # NORTH BOUNDARY\n",
    "            if bc_dict[\"north\"].btype == \"const\":\n",
    "                rhs[ind] = bc_dict[\"west\"].val\n",
    "                the_matrix[ind, ind] = 1\n",
    "            else:  # flux boundary condition\n",
    "                the_matrix[ind, ind] = 1\n",
    "                the_matrix[ind, ind - number_of_cols] = -1\n",
    "                rhs[ind] = bc_dict[\"north\"].val / dy\n",
    "        else:\n",
    "            if is1D:\n",
    "                north = 0\n",
    "                south = 0\n",
    "                rhs[ind] = Qsource[ind]\n",
    "                east = coef_x * avg(D_matrix[ind + 1], D_matrix[ind])\n",
    "                west = coef_x * avg(D_matrix[ind - 1], D_matrix[ind])\n",
    "            else:\n",
    "                north = coef_y * avg(D_matrix[i, j], D_matrix[i + 1, j])\n",
    "                south = coef_y * avg(D_matrix[i, j], D_matrix[i - 1, j])\n",
    "                east = coef_x * avg(D_matrix[i, j], D_matrix[i, j + 1])\n",
    "                west = coef_x * avg(D_matrix[i, j], D_matrix[i, j - 1])\n",
    "                the_matrix[ind, ind + number_of_cols] = -north\n",
    "                the_matrix[ind, ind - number_of_cols] = -south\n",
    "                rhs[ind] = Qsource[i, j]\n",
    "\n",
    "            the_matrix[ind, ind] = east + west + north + south\n",
    "            the_matrix[ind, ind + 1] = -east\n",
    "            the_matrix[ind, ind - 1] = -west\n",
    "\n",
    "    return the_matrix, rhs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The boundary/problem classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Boundary_Def:\n",
    "    \"\"\"\n",
    "    this class holds the boundary type btype ('flux' or 'const')\n",
    "    and the value of the boundary condition (derivitive of the concentration if 'flux'\n",
    "    value of the concentration if 'const')\n",
    "    \"\"\"\n",
    "\n",
    "    btype: str\n",
    "    val: float\n",
    "\n",
    "    def __init__(self, btype, val):\n",
    "        self.btype = btype\n",
    "        self.val = val"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Problem_Def:\n",
    "    \"\"\"\n",
    "    this class holds the specifcation for the domain,\n",
    "    including the value of the porosity\n",
    "    \"\"\"\n",
    "\n",
    "    nx: int\n",
    "    ny: int\n",
    "    poro: float\n",
    "    wx: float\n",
    "    wy: float\n",
    "\n",
    "    def __init__(self, nx, ny, poro, wx, wy):\n",
    "        self.nx = nx\n",
    "        self.ny = ny\n",
    "        self.poro = poro\n",
    "        self.wx = wx\n",
    "        self.wy = wy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create the boundary conditions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here we create 4 boundaries, west has a constant concentration at c0, east has a constant boundary at 0;\n",
    "c0 = 1  # mg/L\n",
    "west = Boundary_Def(\"const\", val=c0)\n",
    "east = Boundary_Def(\"const\", val=0)\n",
    "\n",
    "# For 1D problem, the used boundaries are west and east.\n",
    "\n",
    "# The other south and north boundaries have a zero flux (impermeable)\n",
    "\n",
    "north = Boundary_Def(\"flux\", val=0)\n",
    "south = Boundary_Def(\"flux\", val=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "bc_dict = {\"west\": west, \"north\": north, \"east\": east, \"south\": south}\n",
    "# The latter array bc_dict will be sent to the different functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build the 2D matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "n_x = 51\n",
    "n_y = 1 # 1D -- x only\n",
    "Diff = 2e-9 * 100 * 24 * 3600  # dm²/day\n",
    "width_x = 10  # dm\n",
    "width_y = 0\n",
    "n = n_x * n_y\n",
    "x = np.linspace(0, width_x, n_x)\n",
    "c_init = np.zeros(n_x)\n",
    "c_init[0] = c0\n",
    "D_matrix = Diff * np.ones(n)\n",
    "poro = 0.4\n",
    "prob = Problem_Def(n_x, n_y, poro, width_x, width_y)\n",
    "Qsource = np.zeros(n)\n",
    "A, b = build_2D_matrix(bc_dict, prob, D_matrix, Qsource)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Make a checkpoint\n",
    "\n",
    "Save everything you need to reconstruct A and B outside of this notebook\n",
    "\n",
    "use https://docs.scipy.org/doc/numpy/reference/generated/numpy.savez.html to\n",
    "dump the state to a file named \"savestate.npz\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Save our work in an npz file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "writing the savez file savestate.npz\n"
     ]
    }
   ],
   "source": [
    "out_name= 'savestate.npz'\n",
    "np.savez(out_name,n_x=n_x,n_y=n_y,poro=poro,width_x=width_x,\n",
    "        width_y=width_y,Qsource=Qsource, c0=c0, c_init=c_init,Diff=Diff,\n",
    "        save_date=\"March 3, 2019\",\n",
    "        case_name=\"1D simulation\",\n",
    "        comment=\"simple x-only domain\",\n",
    "        history=\"written by 2D_Assignment_Transient_Error_checkpoint.ipynb\",\n",
    "        units=[\"Diff: dm²/day\",\"length: dm\",\"concentration: mg/L\"])\n",
    "print(f\"writing the savez file {out_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Write functions and classes to a module file\n",
    "\n",
    "The cell below saves the function and class definitions into\n",
    "a new module call `build_2D_matrix.py`.  It uses the special IPython \"magic\" function\n",
    "`%%writefile build_2D_matrix.py`\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting build_2D_matrix.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile build_2D_matrix.py\n",
    "\"\"\"\n",
    "functions and classes for 2D matrix construction\n",
    "\"\"\"\n",
    "def index_to_row_col(ind, nrows, ncol):\n",
    "    \"\"\"\n",
    "    in a 2D array, returns the row and column value\n",
    "    associated with a 1D index\n",
    "    Bottom left is index is zero (0-th row, 0-th column)\n",
    "    while index one is the 0-th row and 1st column\n",
    "    \"\"\"\n",
    "    if ind > nrows * ncol - 1:\n",
    "        return 0\n",
    "    row = int(np.floor(ind / ncol))\n",
    "    col = int(ind - row * ncol)\n",
    "    return row, col\n",
    "\n",
    "def avg(Di, Dj):\n",
    "    \"\"\"\n",
    "    Computes the harmonic average between two values Di and Dj\n",
    "    Returns 0 if either of them is zero\n",
    "    \"\"\"\n",
    "    if (Di * Dj) == 0:\n",
    "        return 0\n",
    "    else:\n",
    "        return 2 / (1 / Di + 1 / Dj)\n",
    "    \n",
    "class Boundary_Def:\n",
    "    \"\"\"\n",
    "    this class holds the boundary type btype ('flux' or 'const')\n",
    "    and the value of the boundary condition (derivitive of the concentration if 'flux'\n",
    "    value of the concentration if 'const')\n",
    "    \"\"\"\n",
    "\n",
    "    btype: str\n",
    "    val: float\n",
    "\n",
    "    def __init__(self, btype, val):\n",
    "        self.btype = btype\n",
    "        self.val = val\n",
    "        \n",
    "class Problem_Def:\n",
    "    \"\"\"\n",
    "    this class holds the specifcation for the domain,\n",
    "    including the value of the porosity\n",
    "    \"\"\"\n",
    "\n",
    "    nx: int\n",
    "    ny: int\n",
    "    poro: float\n",
    "    wx: float\n",
    "    wy: float\n",
    "\n",
    "    def __init__(self, nx, ny, poro, wx, wy):\n",
    "        self.nx = nx\n",
    "        self.ny = ny\n",
    "        self.poro = poro\n",
    "        self.wx = wx\n",
    "        self.wy = wy        \n",
    "\n",
    "def build_2D_matrix(bc_dict, problem, D_matrix, Qsource):\n",
    "    \"\"\"\n",
    "    Constructs a coefficient matrix A and an array b corresponding to the system Ac = b\n",
    "    This system corresponds either to a 1D or 2D problem\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    bc_dict: dict\n",
    "       dictionary with Boundary_Def objects defining the boundary conditions\n",
    "    D_matrix: (float vector if 1D or matrix if 2D)\n",
    "        values of the diffusion coefficient at each grid point(dm^2/day)\n",
    "        if 2D, dimension is [problem.ny, problem.nx]\n",
    "    width_x: (float)\n",
    "        x-extent of the domain (dm)\n",
    "    width_y: (float)\n",
    "        y-extent of the domain (dm)\n",
    "    poro (float)\n",
    "        porosity value\n",
    "    Qsource: (float array)\n",
    "      volumetric source term (mg/L/day)\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "\n",
    "    the_matrix, rhs: tuple\n",
    "       where the_matrix=A and rhs =b\n",
    "       in the discretized diffusion problem\n",
    "       Ax=b\n",
    "    \"\"\"\n",
    "    import numpy as np\n",
    "    number_of_rows = problem.ny\n",
    "    number_of_cols = problem.nx\n",
    "    n = problem.nx * problem.ny\n",
    "    is1D = False\n",
    "    if number_of_rows == 1 or number_of_cols == 1:\n",
    "        is1D = True\n",
    "        number_of_cols = n\n",
    "    the_matrix = np.zeros((n, n))\n",
    "    rhs = np.zeros(n)\n",
    "\n",
    "    if is1D:\n",
    "        dx = max(problem.wx, problem.wy) / (max(problem.ny, problem.nx) - 1)\n",
    "        coef_x = problem.poro / dx / dx\n",
    "    else:\n",
    "        dx = problem.wx / (problem.ny - 1)\n",
    "        dy = problem.wy / (problem.nx - 1)\n",
    "        coef_x = problem.poro / dx / dx\n",
    "        coef_y = problem.poro / dy / dy\n",
    "\n",
    "    for ind in range(n):\n",
    "        if is1D:\n",
    "            j = ind\n",
    "            i = -1\n",
    "        else:\n",
    "            i, j = index_to_row_col(ind, number_of_rows, number_of_cols)\n",
    "        if j == 0:  # WEST BOUNDARY\n",
    "            if bc_dict[\"west\"].btype == \"const\":\n",
    "                rhs[ind] = bc_dict[\"west\"].val\n",
    "                the_matrix[ind, ind] = 1\n",
    "            else:  # flux boundary condition\n",
    "                the_matrix[ind, ind] = 1\n",
    "                the_matrix[ind, ind + 1] = -1\n",
    "                rhs[ind] = bc_dict[\"west\"].val / dx\n",
    "\n",
    "        elif j == number_of_cols - 1:  # EAST BOUNDARY\n",
    "            if bc_dict[\"east\"].btype == \"const\":\n",
    "                rhs[ind] = bc_dict[\"east\"].val\n",
    "                the_matrix[ind, ind] = 1\n",
    "            else:  # flux boundary condition\n",
    "                the_matrix[ind, ind] = 1\n",
    "                the_matrix[ind, ind - 1] = -1\n",
    "                rhs[ind] = bc_dict[\"east\"].val / dx\n",
    "        elif i == 0 and problem.ny > 1:  # SOUTH BOUNDARY\n",
    "            if bc_dict[\"south\"].btype == \"const\":\n",
    "                rhs[ind] = bc_dict[\"south\"].val\n",
    "                the_matrix[ind, ind] = 1\n",
    "            else:  # flux boundary condition\n",
    "                the_matrix[ind, ind] = 1\n",
    "                the_matrix[ind, ind + number_of_cols] = -1\n",
    "                rhs[ind] = bc_dict[\"south\"].val / dy\n",
    "\n",
    "        elif i == number_of_rows - 1 and problem.ny > 1:  # NORTH BOUNDARY\n",
    "            if bc_dict[\"north\"].btype == \"const\":\n",
    "                rhs[ind] = bc_dict[\"west\"].val\n",
    "                the_matrix[ind, ind] = 1\n",
    "            else:  # flux boundary condition\n",
    "                the_matrix[ind, ind] = 1\n",
    "                the_matrix[ind, ind - number_of_cols] = -1\n",
    "                rhs[ind] = bc_dict[\"north\"].val / dy\n",
    "        else:\n",
    "            if is1D:\n",
    "                north = 0\n",
    "                south = 0\n",
    "                rhs[ind] = Qsource[ind]\n",
    "                east = coef_x * avg(D_matrix[ind + 1], D_matrix[ind])\n",
    "                west = coef_x * avg(D_matrix[ind - 1], D_matrix[ind])\n",
    "            else:\n",
    "                north = coef_y * avg(D_matrix[i, j], D_matrix[i + 1, j])\n",
    "                south = coef_y * avg(D_matrix[i, j], D_matrix[i - 1, j])\n",
    "                east = coef_x * avg(D_matrix[i, j], D_matrix[i, j + 1])\n",
    "                west = coef_x * avg(D_matrix[i, j], D_matrix[i, j - 1])\n",
    "                the_matrix[ind, ind + number_of_cols] = -north\n",
    "                the_matrix[ind, ind - number_of_cols] = -south\n",
    "                rhs[ind] = Qsource[i, j]\n",
    "\n",
    "            the_matrix[ind, ind] = east + west + north + south\n",
    "            the_matrix[ind, ind + 1] = -east\n",
    "            the_matrix[ind, ind - 1] = -west\n",
    "\n",
    "    return the_matrix, rhs        "
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "all",
   "formats": "ipynb,python//py:percent",
   "notebook_metadata_filter": "all"
  },
  "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.6.7"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": true,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}