DL_LEM_Numerical_Solver.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff1fb37e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from landlab.components import FlowAccumulator\n",
    "from       IPython.display import clear_output\n",
    "from         landlab    import RasterModelGrid\n",
    "import                 matplotlib.pylab as plt\n",
    "from                     scipy import optimize\n",
    "from                           mpmath import *\n",
    "import                             numpy as np\n",
    "import                                   scipy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa2d4064",
   "metadata": {},
   "outputs": [],
   "source": [
    "def DinfEroder(mg, dt, K_sp, m_sp, n_sp):\n",
    "    # This function computes the erosion term implicitly using the D_infinity flow direction method\n",
    "    '''\n",
    "    Function arguments:\n",
    "    mg  :                    Raster grid used for the simulation\n",
    "    dt  :                                  Time-step to be taken\n",
    "    K_sp:                                    Erosion coefficient\n",
    "    m_sp: Exponent of specific drainage area in the erosion term\n",
    "    n_sp:            Exponent of local slope in the erosion term\n",
    "    '''\n",
    "    \n",
    "    # Initializing local arrays for the function block\n",
    "    nodes      =                                                      np.arange(mg.number_of_nodes)\n",
    "    receiver   =                                                  mg.at_node['flow__receiver_node']\n",
    "    A          =                                               np.copy(mg.at_node['drainage_area'])\n",
    "    proportion =                                           mg.at_node['flow__receiver_proportions']\n",
    "    ele        =                                      np.copy(mg.at_node['topographic__elevation'])\n",
    "    neighbor   = np.append(mg.diagonal_adjacent_nodes_at_node, mg.adjacent_nodes_at_node, axis = 1)\n",
    "    donor      =                                                            np.zeros_like(neighbor)\n",
    "    \n",
    "    # Listing the receivers of each node's neighbors\n",
    "    receiver_of_neighbor        =                                         receiver[neighbor]\n",
    "    receiver_of_neighbor[:,:,0] = np.where(neighbor == -1 , -1, receiver_of_neighbor[:,:,0])\n",
    "    receiver_of_neighbor[:,:,1] = np.where(neighbor == -1 , -1, receiver_of_neighbor[:,:,1])\n",
    "    \n",
    "    # Determining the donors for each node in the domain\n",
    "    for nei in range(8):\n",
    "        donor[:, nei] = np.where((receiver_of_neighbor[:, nei, 0] == nodes) |\n",
    "                                 (receiver_of_neighbor[:, nei, 1] == nodes) , neighbor[:, nei], -1)\n",
    "\n",
    "    # Initializing the queue by adding sink nodes\n",
    "    queue           = list(np.arange(mg.number_of_nodes)[mg.at_node['flow__sink_flag'] == 1])\n",
    "    \n",
    "    # Marking the sink nodes as processed\n",
    "    processed_nodes        =   np.zeros_like(nodes, dtype=bool)\n",
    "    processed_nodes[queue] =                               True\n",
    "    \n",
    "    # Traversing the remaining network\n",
    "    while len(queue) > 0:\n",
    "        \n",
    "        # Popping the first element of the queue\n",
    "        node = queue.pop(0)\n",
    "        \n",
    "        # Visiting its donors\n",
    "        for nei in donor[node, donor[node] >= 0]:\n",
    "            \n",
    "            # Processing the donor if all the receivers of that donor are processesd\n",
    "            if processed_nodes[nei] == False and np.all(processed_nodes[receiver[nei][receiver[nei] >= 0]]):\n",
    "                processed_nodes[nei] = True\n",
    "                \n",
    "                # Finding the slope using the elevation of neighbors in cardinal and diagonal directions\n",
    "                h_0 , h_1 = None, None\n",
    "                for res, pro in zip(receiver[nei], proportion[nei]):\n",
    "                    if res != -1:\n",
    "                        if res // ncols == nei // ncols or res % ncols == nei % ncols:\n",
    "                            h_0 = ele[res]\n",
    "                        else:\n",
    "                            h_1 = ele[res]\n",
    "                h = ele[nei]\n",
    "                width = dx\n",
    "                if h_0 is not None and h_1 is not None:\n",
    "                    f = lambda h_next: h_next - h + K_sp * (A[nei] / width)**m_sp *\\\n",
    "                                ((h_0 - h_1)**2 + (h_next - h_0)**2)**(n_sp / 2.) * dt / dx**n_sp\n",
    "                    min_h = min(h_0, h_1)\n",
    "\n",
    "                elif h_0 is not None:\n",
    "                    f = lambda h_next: h_next - h + K_sp * (A[nei] / width)**m_sp *\\\n",
    "                                ((h_next - h_0)**2)**(n_sp / 2.) * dt / dx**n_sp\n",
    "                    min_h = h_0\n",
    "                elif h_1 is not None:\n",
    "                    f = lambda h_next: h_next - h + K_sp * (A[nei] / width)**m_sp *\\\n",
    "                                (h_next - h_1)**n_sp * dt / (2**.5 * dx)**n_sp\n",
    "\n",
    "                    min_h = h_1\n",
    "                    \n",
    "                # Updating the elevation implicitly using Brent’s method in Python\n",
    "                ele[nei] = optimize.brenth(f, h, min_h)\n",
    "                \n",
    "                # Appending the processed node to the queue\n",
    "                queue.append(nei)\n",
    "    return ele\n",
    "\n",
    "def implicit_diffusion(zgiven, dt, D):    \n",
    "    # This function implicitly solves the diffusion term with the domain boundary nodes\n",
    "    # at zero fixed elevation value. For any other boundary conditions, the diagonals of the \n",
    "    # linear (matrix) system should be modified accordingly.\n",
    "    '''\n",
    "    Function arguments:\n",
    "    dt     :          Time-step to be taken\n",
    "    zgiven :                Elevation field\n",
    "    D      :         Soil creep coefficient\n",
    "    '''   \n",
    "    # dx is grid resolution, dt is the time step, D is the coefficient\n",
    "    kd   =                       D*dt/(dx**2)\n",
    "    d1m1 = np.repeat(-kd, len(mg.core_nodes))\n",
    "    d1m1[ncols-3:d1m1.size:ncols-2]       = 0\n",
    "    A    = scipy.sparse.diags([-kd, d1m1, 1 + 4*kd, d1m1, -kd], [-ncols+2, -1, 0, 1, ncols-2], shape=(len(mg.core_nodes), len(mg.core_nodes)))\n",
    "    return scipy.sparse.linalg.lgmres(A, zgiven, atol=0.000000000001)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a4b9903",
   "metadata": {},
   "outputs": [],
   "source": [
    "#          Initializing the coefficients, domain-size, and model parameters\n",
    "K_sp  = 0.0025 #                                Fluvial erosion coefficient\n",
    "D     =  0.005 #                                     Soil creep coefficient\n",
    "U     =  0.001 #                                                Uplift rate\n",
    "m_sp  =    0.5 # Exponent of the specific drainage area in the erosion term\n",
    "n_sp  =    1.0 #            Exponent of the local slope in the erosion term\n",
    "dx    =    1.0 #                  Raster grid resolution for the simulation\n",
    "nrows =    500 #                                             Number of rows\n",
    "ncols =    100 #                                          Number of columns\n",
    "\n",
    "CI    = (K_sp*((ncols)*dx)**(m_sp + n_sp))/(D**n_sp * U**(1-n_sp))\n",
    "print(\"The channelization index for this run is\",       round(CI))\n",
    "print(\"This index tells the tendency of the system to form channels. More information is in the manuscript of the code.\")\n",
    "\n",
    "# Providing the initial condition for the elevation field (z)\n",
    "z     = np.load('Zinitial_500_100.npy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8b8d0cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initializing the raster grid\n",
    "mg    =                           RasterModelGrid((nrows,ncols), dx)\n",
    "_     =               mg.add_zeros('node', 'topographic__elevation')\n",
    "\n",
    "# Setting the boundary conditions for the numerical model\n",
    "for edge in (mg.nodes_at_left_edge, mg.nodes_at_right_edge):\n",
    "    mg.status_at_node[edge] = RasterModelGrid.BC_NODE_IS_FIXED_VALUE\n",
    "for edge in (mg.nodes_at_bottom_edge,mg.nodes_at_top_edge):\n",
    "    mg.status_at_node[edge] = RasterModelGrid.BC_NODE_IS_FIXED_VALUE\n",
    "\n",
    "# Initializing the flow-direction algorithm for the given topography    \n",
    "mg.at_node['topographic__elevation']   =     z.reshape(nrows * ncols)\n",
    "fc = FlowAccumulator(mg, flow_director =          'FlowDirectorDINF')\n",
    "fc.run_one_step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4ab3fbf",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Lists, variables to keep track of time-steps, temporal changes in elevation, etc., during the simulation run.\n",
    "count        =        0\n",
    "count_max    =      500\n",
    "min_try      =      500\n",
    "i            =       -1\n",
    "t            =        0\n",
    "diff_list    =       []\n",
    "steady_state =    False\n",
    "\n",
    "# Beginning of the simulation run\n",
    "while steady_state is False:\n",
    "    i    += 1\n",
    "    alpha = min((i + 1) * 0.01, 1.)\n",
    "    dt    = alpha * min(dx / (K_sp * np.max((mg.at_node['drainage_area'] / dx) ** m_sp)) , 20. * dx ** 2 / (2 * D))\n",
    "    dt    = min(dt, 500.)\n",
    "\n",
    "    ele_1      = np.copy(mg.at_node['topographic__elevation'])\n",
    "    erode_done =                                         False\n",
    "    \n",
    "    while erode_done is False:\n",
    "        try:\n",
    "            fc = FlowAccumulator(mg, flow_director = 'FlowDirectorDINF')\n",
    "            fc.run_one_step()\n",
    "            \n",
    "            # This function implements the implicit algorithm for updating the erosion term in Anand et al.(2020).\n",
    "            mg.at_node['topographic__elevation'] = DinfEroder(mg, dt, K_sp, m_sp, n_sp)\n",
    "            \n",
    "            # This function updates the diffusion term implicitly for the fixed elevation boundary conditions.\n",
    "            mg.at_node['topographic__elevation'][mg.core_nodes] = implicit_diffusion(mg.at_node['topographic__elevation'][mg.core_nodes] + U * dt, dt, D)\n",
    "            erode_done = True\n",
    "\n",
    "        except ValueError:\n",
    "            print('Adjusting dt')\n",
    "            mg.at_node['topographic__elevation'] = ele_1\n",
    "            dt = dt / 2.\n",
    "\n",
    "    ele_2             = np.copy(mg.at_node['topographic__elevation'])\n",
    "    t                          +=                                  dt\n",
    "    ele_diff                    =         np.abs(ele_1 - ele_2).max()\n",
    "    max_diff_location           =      np.abs(ele_1 - ele_2).argmax()\n",
    "    ele_diff_mean               = np.abs(ele_1.mean() - ele_2.mean())\n",
    "    diff_list.append([t, ele_diff, ele_diff_mean, max_diff_location])\n",
    "\n",
    "    if ele_diff_mean < 0.00005 * U * dt:\n",
    "        count = count + 1\n",
    "    else:\n",
    "        count = 0\n",
    "\n",
    "    if  i >= min_try and ele_diff < U * 0.001 * dt and ele_diff_mean < 0.0001 * U * dt:\n",
    "        steady_state = True\n",
    "\n",
    "    if count == count_max:\n",
    "        steady_state = True\n",
    "    \n",
    "    # Visualizing the transient elevation and specific drainage area solutions. If one is only\n",
    "    # interested in steady-state solutions, consider commenting the generated plots for speed-up.\n",
    "    if i%50 == 0:\n",
    "        clear_output(wait=True)\n",
    "        plt.figure(figsize=(8, 4.5))\n",
    "        plt.subplot(211)\n",
    "        plt.title('Specific Drainage Area [m]')\n",
    "        plt.imshow(np.transpose((mg.at_node['drainage_area'].reshape(nrows,ncols)/dx)**0.5), cmap ='seismic')\n",
    "        plt.colorbar()\n",
    "        plt.xlabel('X [m]')\n",
    "        plt.ylabel('Y [m]')\n",
    "        plt.subplot(212)\n",
    "        plt.title('Elevation [m]')\n",
    "        plt.imshow(    np.transpose(mg.at_node['topographic__elevation'].reshape(nrows,ncols)), cmap ='BrBG')\n",
    "        plt.xlabel('X [m]')\n",
    "        plt.ylabel('Y [m]')\n",
    "        plt.colorbar()\n",
    "        plt.tight_layout()\n",
    "        plt.show()\n",
    "\n",
    "    print(i, t, round(ele_diff/ (U * 0.001 * dt), 1), round(ele_diff_mean/ (U * 0.0001 * dt), 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d32be2f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "#                                Saving the steady-state solution arrays\n",
    "z           = str(round(chi)) + '_steady_ele_at_' + str(int(t)) + '.npy'\n",
    "a           = str(round(chi)) + '_steady_acc_at_' + str(int(t)) + '.npy'\n",
    "np.save(z,    mg.at_node['topographic__elevation'].reshape(nrows,ncols))\n",
    "np.save(a,          mg.at_node['drainage_area'].reshape(nrows,ncols)/dx)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "lemv2 [~/.conda/envs/lemv2/]",
   "language": "python",
   "name": "conda_lemv2"
  },
  "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}