CrateringLesson_v2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction to Craters on Planetary Surfaces\n",
    "\n",
    "Cratering is an important process that has altered the surface topography of the Earth, other terrestrial and icy planets, and both rocky and icy moons, within our solar system and probaby others. That is - cratering is quite a uniform process. However, although we understand the process of cratering quite well, we only have a limited understanding about how a cratered surface influences the role of other surface processes. For example, on heavily cratered surfaces of the Moon, Mars, and Mercury, which have not been altered by plate tectonics like on Earth, how do the craters affect the flow and collection of water? of wind? and of the sediments transported by water and wind? These are funamental questions in our understanding of planetary bodies. Looking at the surfaces of the 4 rocky planets in our solar system, Earth is the odd one out, with topography dominated by mountains and basins created by plate tectonics. Based on this knowledge, what do you expect other planetary surfaces in other solar systems would look like, cratered or Earth-like? \n",
    "\n",
    "\n",
    "#### Learning objectives\n",
    "_Skills_\n",
    "* Create functions that generate crater morphologies using simple known laws\n",
    "* Make changes to input parameters to vary the distribution and size of crater populations\n",
    "* Use of LandLab model grids and components to route any kind of flow over the model surface\n",
    "* Gain hands-on experience with visualizing data in Python\n",
    "\n",
    "_Key concepts_\n",
    "* Explore and understand the process of simulation of craters on a landscape\n",
    "* Explore how an initially cratered landscape is affected by hillslope diffusion and the flow of surface water\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "In this notebook, we also use a Landlab model grid to simulate topography, you can see more about Landlab [here](https://landlab.github.io/#/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from landlab import RasterModelGrid, imshow_grid, NodeStatus, values\n",
    "import random as random\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How do we simulate crater location and size?\n",
    "\n",
    "The size and location of a crater is somewhat random, but to constrain it in numerical models, we \"pick\" numbers randomly from a given probability distribution.   \n",
    "\n",
    "#### Randomly Selecting Crater Locations\n",
    "\n",
    "In this model, we pick crater loctions randomly from a  \"uniform\" distribution, that is, each location has an equally probable chance of having a crater. For example, in the code below we pick an x-coordinate and y-coordinate at random from a uniform distribution using the numpy.random.randint function which generates a random integer from a given minimum to a given maximum:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cratercenter = (np.random.randint(1, 10), np.random.randint(1, 10))\n",
    "print(\"(x, y) = \", cratercenter)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we generate many crater centers and plot the x and y coordinates as a histogram you can see there are _roughly_ equal numbers of each number but this is not a perfect fit as there is some \"noise\" as we randomly selected from this distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cratercenters = [np.random.randint(1, 11) for i in range(100)]\n",
    "\n",
    "plt.hist(cratercenters, bins = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);\n",
    "plt.xlabel(\"Random integer from 1 to 10\", fontsize = 14);\n",
    "plt.ylabel(\"Count\", fontsize = 14);\n",
    "plt.title(\"Histogram of Integers Randomly Selected\\n from a Uniform Distribution\", fontsize = 16);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Randomly Selecting Crater Diameters\n",
    "\n",
    "For crater size, we do not choose from a random distribution because there are actually different probabilities for different crater sizes. \n",
    "* Larger meteorites are less common, and more likely to break up in the atmosphere, so they are less likely to create craters at the surface. \n",
    "* Medium to small meteorites are more common and correspondingly so are their craters, \n",
    "* But, very small meteorites, although very common, are not likely to make it to the surface as they burn up in the atmosphere - even if it is incredibly thin like on Mars. \n",
    "\n",
    "Here we ignore very small craters (<1 km) as in a landscape evolution model we would likely be looking at medium to large scales (>> 1 km), and this simplifies the probability distribution to a **cumulative probability distribtuion** (Figure). Where the number of craters, ND, of diameter, D, that form is related to two scaling factors, Kx and an exponent $ \\delta $:\n",
    "\n",
    "$ N_D = K_x    D^{- \\delta} $  \n",
    "\n",
    "For Mars it has been estimated that: \n",
    "\n",
    "$ K_x $ = 1.0\n",
    "\n",
    "and\n",
    "\n",
    "$ \\delta $ = 2.0 \n",
    "\n",
    "In reality we implement this as shown in the code below, by using the \"number of craters of a given diameter (Nd)\" as a probability weight for that diameter, and we select the crater diameter randomly from a probability distribution of how many of that given diameter's craters there should be: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Kx = 1.0  # Scaling coefficient\n",
    "delta = 2.0  # km, scaling exponent\n",
    "# Ncraters = 2 # Number of craters to add\n",
    "\n",
    "minD = 1 #1 km\n",
    "maxD =  50 # 50 km\n",
    "\n",
    "diameter = list(range(minD, maxD)) ## make a list of all the possible diameters, from minD to maxD\n",
    "\n",
    "NumberDs = [] #This step is required before the loop to instantiate a list which will hold our output\n",
    "for D in range(minD, maxD): #for each diameter from 1 km to 50 km\n",
    "    NumberD = Kx * D ** - delta #calculate the number of craters of that diameter that there should be\n",
    "    NumberDs.append(NumberD) #append the number to a list of all the \"numbers\" of craters for each diameter, from 1 to 50 km\n",
    "    \n",
    "\n",
    "def weighted_choice_sub(weights):\n",
    "    ''' randomly generate a number and see which weight number in the input list it falls under,\n",
    "    return the index of that weight '''\n",
    "    rnd = np.random.random() * sum(weights)\n",
    "    for i, w in enumerate(weights):\n",
    "        rnd -= w\n",
    "        if rnd < 0:\n",
    "            return i #return the index of the weight, which corresponds to the index of the diameter in a \"partner\" list\n",
    "\n",
    "\n",
    "D_index = weighted_choice_sub(NumberDs);\n",
    "print(\"Randomly selected diameter: {} km\".format(diameter[D_index]));\n",
    "## You can run the above cell a few times - what is the most commonly printed diameter or diameter range?\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we generate many crater diameters and plot them as a histogram you can see there are more small craters and the number of craters of larger and larger diameters decreases, but this is not a perfect fit as there is some \"noise\" as we **_randomly_** selected from this distribution:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "craterdiameters = [diameter[weighted_choice_sub(NumberDs)] for i in range(100)]\n",
    "\n",
    "plt.hist(craterdiameters, bins = 10, range = (minD, maxD));\n",
    "plt.xlabel(\"Crater Diameter from {} km to {} km\".format(minD, maxD), fontsize = 14);\n",
    "plt.ylabel(\"Count\", fontsize = 14);\n",
    "plt.title(\"Histogram of Crater Diameters Randomly Selected\\n from a Cumulative Probability Distribution\", fontsize = 16);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How do we simulate crater shape? \n",
    "\n",
    "Good question! \n",
    "\n",
    "Craters are recognized because of the characteristic bowl-like depression that transitions upwards into a crater rim, that is actually above the surrounding terrain and surrounded by an \"ejecta blanket\" of all the material excavated from the crater itself, and with distance away from the crater the rim slopes downwards and the ejecta blanket gets thinner, so their contributions to elevation get smaller and smaller. These changes in elevation can actually be described by mathematical equations such as an exponential function, or in our case, a power law. \n",
    "\n",
    "For a SIMPLE crater (bowl-like depression) with diameter, D (in m), on Mars, the elevation changes are:\n",
    "\n",
    "$ H2 - H1 + H1 \\frac{2r}{D}^{m} $ within the crater (ie. from radial distance = 0 to radial distance = crater radius) (1)\n",
    "\n",
    "$ H2 \\frac{2r}{D}^{-n} $ outside of the crater (ie. where radial distance > crater radius, to describe the ejecta blanket) (2) \n",
    "\n",
    "\n",
    "and where $ r $ is the radial distance from the crater center, and $ H1, H2, $ and the exponents $ m $ and $ n $ are:\n",
    "\n",
    "$ H1 = 2.54 D ^{0.67} $\n",
    "\n",
    "$  H2 = 1.93 D ^{0.52} $\n",
    "  \n",
    "$  m = 0.73 D ^{0.11} $ \n",
    "\n",
    "$ n  = 3 $\n",
    "\n",
    "\n",
    "These equations can be modified slightly to reflect the shape of a Complex crater, although they still do not describe finer crater detail such a central uplift and terracing, which is usually not considered as part of landscape-scale simulations. \n",
    "\n",
    "See [Howard (2007)](https://doi.org/10.1016/j.geomorph.2007.04.017) for more information on simulating craters!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "diameter = 40\n",
    "radius = diameter/2\n",
    "rs = np.arange(start=0, stop=radius*2, step=0.1) #create an array of radial distances from the crater center\n",
    "\n",
    "H1 = 2.54*(diameter*1000)**0.67\n",
    "H2 = 1.93*(diameter*1000)**0.52\n",
    "m = 0.73*(diameter*1000)**0.11 \n",
    "n = 3\n",
    "\n",
    "incrater = rs[rs <= radius]\n",
    "in_idx = np.where(rs <= radius)[0]\n",
    "# equation for inside the crater\n",
    "inDepth = (H2 - H1 + H1*((2*incrater)/(diameter))**m)/1000\n",
    "# mg.at_node['topographic__elevation'][in_idx] += inDepth\n",
    "\n",
    "outcrater = rs[rs >= radius]\n",
    "out_idx = np.where(rs >= radius)[0]\n",
    "# equation for outside the crater (ejecta!)\n",
    "outDepth = (H2*((2*outcrater)/(diameter))**-n)/1000\n",
    "# mg.at_node['topographic__elevation'][out_idx] += outDepth\n",
    "\n",
    "\n",
    "fig, ax = plt.subplots(1, 2, dpi=250, facecolor='w', figsize=(10, 3))\n",
    "# fig.suptitle(\"Representation of Crater Shape\", fontsize = 16);\n",
    "ax[0].plot(incrater/1000, inDepth, 'k', label = 'Inside Crater');\n",
    "ax[0].plot(outcrater/1000, outDepth, 'r', label = 'Outside Crater (ejecta)');\n",
    "ax[0].set_xlabel(\"Radial Distance from Center (km)\", fontsize = 14);\n",
    "ax[0].set_ylabel(\"Elevation (km)\", fontsize =14 );\n",
    "ax[0].set_title(\"Simple Crater\")\n",
    "ax[0].set_ylim(min(inDepth)-0.5, max(inDepth)+0.5)\n",
    "\n",
    "## Lets also compare the shape of a complex crater! (for simplicity; the crater is of the same diameter as the one above)\n",
    "H1 = 12.20*(diameter*1000)**0.49\n",
    "H2 = 0.79*(diameter*1000)**0.6\n",
    "m = 0.64*(diameter*1000)**0.13 \n",
    "n = 3\n",
    "\n",
    "# equation for inside the crater\n",
    "inDepthC = (H2 - H1 + H1*((2*incrater)/(diameter))**m)/1000\n",
    "# mg.at_node['topographic__elevation'][in_idx] += inDepth\n",
    "\n",
    "# equation for outside the crater (ejecta!)\n",
    "outDepthC = (H2*((2*outcrater)/(diameter))**-n)/1000\n",
    "# mg.at_node['topographic__elevation'][out_idx] += outDepth\n",
    "\n",
    "ax[1].plot(incrater/1000, inDepthC, 'k', label = 'Inside Crater');\n",
    "ax[1].plot(outcrater/1000, outDepthC, 'r', label = 'Outside Crater (ejecta)');\n",
    "ax[1].set_xlabel(\"Radial Distance from Center (km)\", fontsize = 14);\n",
    "ax[1].set_ylabel(\"Elevation (km)\", fontsize =14 );\n",
    "ax[1].legend(loc=\"best\", fontsize=14);\n",
    "ax[1].set_title(\"Complex Crater\");\n",
    "ax[1].set_ylim(min(inDepth)-0.5, max(inDepth)+0.5);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Now with all of the cratering theory covered, let's make a cratered surface!\n",
    "\n",
    "This tutorial uses Landlab model grids, a convenient library of modules for building model grids, with many components to then do landscape evolution on the grids! With some tweaks to the code, this cratering module would work just as well on a numpy array of elevation values. \n",
    "\n",
    "First, lets instantiate the Landlab model grid, with a field for elevation values called \"topographic__elevation\":"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "size = 200 #km\n",
    "spacing = 1 #km\n",
    "xy = int(size/spacing) #km\n",
    "rf = 0.1 # how noisy we want the initial topography to be\n",
    "Kx = 1.0  # Scaling coefficient\n",
    "delta = 2.0  # km, scaling exponent\n",
    "\n",
    "mg = RasterModelGrid((xy, xy), xy_spacing = spacing); #initiate surface; see above for variables\n",
    "z = np.zeros(mg.number_of_nodes); # initiate a field which will be our elevation values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can add some random noise to this surface to make it slightly more realistic. The \"seed\" here ensures that the added-randomness is the same everytime, to keep our model grid reproducible! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(30) # Keep this constant (e.g., at 30) so the initial randomness it always the same\n",
    "z += np.random.rand(mg.number_of_nodes)*rf  # make the noise large enough relative to crater\n",
    "mg.add_field(\"topographic__elevation\", z, at=\"node\", units=\"km\", copy=True, clobber=True);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's very easy to visualize model grids with the inbuilt function, imshow_grid:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "imshow_grid(mg, \"topographic__elevation\" , plot_name='Initial, Noisy Surface',\n",
    "        var_name='Elevation', var_units='km', grid_units=('km', 'km')) #imshow(grid, value field)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ADD SOME CRATERS!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we define everything as functions which allows us to easily re-do the cratering over and over without copying code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Everything is packaged as functions: \n",
    "def crater_depth(d, diameter, mg, d_ref=7):\n",
    "    \"\"\"\n",
    "    Define in and out of crater changes to topography.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    d : np.ndarray\n",
    "        Array of distances of nodes to the center of the crater.\n",
    "        From `landlab.grid.raster.RasterModelGrid.calc_distances_of_nodes_to_point`\n",
    "    diameter : int, float\n",
    "        Diameter of the crater\n",
    "    mg : landlab.grid.raster.RasterModelGrid\n",
    "        Landlab raster model grid of the landscape\n",
    "\n",
    "    d_ref: int, float\n",
    "        Diameter at which craters transition from \"simple\" to \"complex\" structural type\n",
    "        Default = 7 km (relevant to Mars)\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    mg : landlab.grid.raster.RasterModelGrid\n",
    "        Landlab raster model grid after crater has modified the topography\n",
    "    \"\"\"\n",
    "    radius = diameter/2\n",
    "    diameter *= 1000\n",
    "    d_ref *= 1000\n",
    "\n",
    "    if diameter <= d_ref:\n",
    "        H1 = 2.54*diameter**0.67\n",
    "        H2 = 1.93*diameter**0.52\n",
    "        m = 0.73*diameter**0.11  # value: 2 to 3\n",
    "\n",
    "    elif diameter > d_ref:\n",
    "        H1 = 12.20*diameter**0.49\n",
    "        H2 = 0.79*diameter**0.6\n",
    "        m = 0.64*diameter**0.13  # value: 2 to 3\n",
    "\n",
    "    # Howard et al, 2007:\n",
    "    # \"The exponent n is constrained such that volume deposited on the rim\n",
    "    # equals the volume excavated from the bowl and ranges from a value of\n",
    "    # about 3 for a 7 km crater to 3.5 for a 250 km crater.\"\"\n",
    "    n = 3\n",
    "    H2H1 = H2 - H1\n",
    "\n",
    "    incrater = d[d <= radius]\n",
    "    in_idx = np.where(d <= radius)[0]\n",
    "    # equation for inside the crater\n",
    "    inDepth = (H2H1 + H1*((2*(incrater*1000))/(diameter))**m)/1000 #In KM (hence /1000)\n",
    "    mg.at_node['topographic__elevation'][in_idx] += inDepth\n",
    "\n",
    "    outcrater = d[d > radius]\n",
    "    out_idx = np.where(d > radius)[0]\n",
    "    # equation for outside the crater (ejecta!)\n",
    "    outDepth = (H2*((2*(outcrater*1000))/(diameter))**-n)/1000 # IN KM (hence /km)\n",
    "    mg.at_node['topographic__elevation'][out_idx] += outDepth\n",
    "    return mg\n",
    "\n",
    "\n",
    "def do_cratering(Ncraters, NDs, minD, maxD, xy, mg, spacing):\n",
    "    \"\"\"\n",
    "    Add craters to some landlab raster model.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    Ncraters : int\n",
    "        Number of craters that impact\n",
    "\n",
    "    NDs : list\n",
    "        List of weights for random sampling\n",
    "\n",
    "    minD : int\n",
    "        Minimum crater diameter\n",
    "\n",
    "    maxD : int\n",
    "        Maximum crater diameter\n",
    "\n",
    "    xy : int\n",
    "        Domain size (# cells) in one direction (domain is square)\n",
    "\n",
    "    mg : landlab.grid.raster.RasterModelGrid\n",
    "        Landlab raster model grid of the landscape\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    mg : landlab.grid.raster.RasterModelGrid\n",
    "        Landlab raster model grid after craters have modified the topography\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    for i in range(Ncraters):  # For N number of craters\n",
    "        a = weighted_choice_sub(NDs)\n",
    "        diameter = list(range(minD, maxD))[a]\n",
    "        cratercenter = (np.random.randint(1, xy*spacing), np.random.randint(1, xy*spacing))\n",
    "        d = mg.calc_distances_of_nodes_to_point(cratercenter)\n",
    "\n",
    "        crater_depth(d, diameter, mg, d_ref=7)\n",
    "\n",
    "    return mg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before running the functions, we need to define the variables we need: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Ncraters = int(input(\"How many craters do you want to add? \")) # see console to add your own number!\n",
    "\n",
    "minD = int(spacing * 3) # want it to be greater than spacing, three times the smallest cell\n",
    "maxD = int((size)/4) #max diameter (m), a quarter of the total domain width\n",
    "print(\"Minimum possible crater size:\", minD, \"km\")\n",
    "print(\"Maximum possible crater size:\", maxD, \"km\")\n",
    "\n",
    "NDs = []\n",
    "for D in range(minD, maxD):\n",
    "    ND = Kx * D ** - delta\n",
    "    NDs.append(ND)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and now, to run the code it's very simple given that we already defined the function \"do_cratering\":"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "random.seed(50) #Chose random seed number 50 (this ensures crater locations are same every time)\n",
    "mg = do_cratering(Ncraters, NDs, minD, maxD, size, mg, spacing)\n",
    "\n",
    "## Display the newly cratered surface!\n",
    "imshow_grid(mg, \"topographic__elevation\", plot_name='Cratered Surface',\n",
    "        var_name='Elevation', var_units='km', grid_units=('km', 'km'), cmap=\"Spectral_r\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that if you rerun the code above, it will keep adding more craters to the surface. To \"reset\" it, rerun all the way from 4 cells above, where we initialize the noisy surface (= mg). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### This is some code that offers better visualisation of the data\n",
    "* use LandLab \"calc_hillshade_at_node\" function to create a hillshade\n",
    "* display the data as a hillshade with semi-transparent colored topography overlain\n",
    "* and display a cross-section of elevation through the middle of the surface"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "## This cell may take a while! don't worry... it's working!\n",
    "size = 200\n",
    "xy = int(size/spacing)\n",
    "spacing = 1\n",
    "rf = 0.1\n",
    "Ncraters = 10 #number of cratesr by timestep\n",
    "\n",
    "mg = RasterModelGrid((xy, xy), xy_spacing = spacing); #initiate surface; see above for variables\n",
    "z = np.zeros(mg.number_of_nodes); # initiate a field which will be our elevation values\n",
    "np.random.seed(30) # Keep this constant (e.g., at 30) so the initial randomness it always the same\n",
    "z += np.random.rand(mg.number_of_nodes)*rf  # make the noise large enough relative to crater\n",
    "mg.add_field(\"topographic__elevation\", z, at=\"node\", units=\"km\", copy=True, clobber=True);\n",
    "\n",
    "# Define where section is taken along y (in km)\n",
    "stk_km = xy/2\n",
    "stk = int(stk_km/spacing) #node location of stk_km\n",
    "# Number of timesteps\n",
    "nsteps = 10\n",
    "\n",
    "cmap1 = \"Greys\"\n",
    "cmap2 = \"Spectral_r\"\n",
    "\n",
    "old_arr = np.zeros((nsteps, xy))\n",
    "\n",
    "for i in range(1, nsteps):\n",
    "    # crater landscape\n",
    "    mg = do_cratering(Ncraters, NDs, minD, maxD, xy, mg, spacing)\n",
    "\n",
    "    # update plot\n",
    "    fig, ax = plt.subplots(1, 2, dpi=250, facecolor='w', figsize=(8, 3))\n",
    "    topo = mg.field_values('node', 'topographic__elevation').reshape((xy, xy)) #create an array of elevation values across the mg domain\n",
    "    hs = mg.calc_hillshade_at_node(elevs='topographic__elevation') #create a hillshade array\n",
    "    hill = np.reshape(hs, (xy, xy)) #reshape the hillshade array to be the same shape as the topo array\n",
    "    img1 = ax[0].imshow(hill, cmap=cmap1, alpha=1, extent = [0,size, 0, size], vmin=0, vmax=1)\n",
    "    img2 = ax[0].imshow(topo, cmap=cmap2, alpha=0.6, extent = [0,size, 0, size])\n",
    "    ax[0].set_title(\"Cratered surface evolution (t = %i)\" %i)\n",
    "    cbar = plt.colorbar(img2, fraction=0.045, ax=ax[0], label = \"Elevation [km]\")\n",
    "    ax[0].set_xlabel('X [km]')\n",
    "    ax[0].set_ylabel('Y [km] ')\n",
    "    #stk = int(stk_km/spacing)\n",
    "    ax[0].plot(np.linspace(0, size), np.ones_like(np.linspace(0, size))*stk_km,\n",
    "               c='k', linestyle='--')\n",
    "    ax[0].set_xlim([0, size])\n",
    "    ax[0].set_ylim([0, size])\n",
    "\n",
    "    ax[1].plot(topo[stk, :], c=[0, 0, 0], zorder=10)\n",
    "    for j in range(old_arr.shape[0]):\n",
    "        if np.sum(old_arr[j, :]) != 0:\n",
    "            ax[1].plot(old_arr[j, :], c=[1-j/nsteps, 1-j/nsteps, 1-j/nsteps])\n",
    "    ax[1].set_title('Topographic section at Y = ' + str(stk_km) + \" km\")\n",
    "    ax[1].set_ylabel('Elevation [km]')\n",
    "    ax[1].set_xlabel('Distance along X')\n",
    "    ax[1].set_ylim() # I think making a dynamic scale is kind of fun, shows how much it changes? Instead of [zmin, zmax]\n",
    "\n",
    "    plt.tight_layout()   \n",
    "\n",
    "    # retain old array section\n",
    "    old_arr[i, :] = topo[stk, :]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How do the craters affect the flow and collection of water? \n",
    "\n",
    "Back to a question posed at the start of this notebook, how does the cratered landscape look with some diffusion? With some 'flowing water'? We can use pre-built Landlab components to explore this!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Some Questions to Think About / Explore:\n",
    "* What does advection (fluvial erosion) do to the cratered landscape? What happens when you differ some of the stream power coefficients? Look at the difference map, where is the most fluvial erosion taking place?\n",
    "* What does diffusion do to the cratered landscape? What happens if you differ the diffusion coefficient? Look at the difference map, where is the most fluvial erosion taking place?\n",
    "* How do the paths of water vary when the cratered landscape is diffused vs. not diffused? "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Each time we initialise a new grid, to avoid using the already cratered grid which we'd just made above\n",
    "size = 200\n",
    "xy = int(size/spacing)\n",
    "spacing = 1\n",
    "rf = 0.1\n",
    "\n",
    "mg = RasterModelGrid((xy, xy), xy_spacing = spacing); #initiate surface; see above for variables\n",
    "z = np.zeros(mg.number_of_nodes); # initiate a field which will be our elevation values\n",
    "np.random.seed(30) # Keep this constant (e.g., at 30) so the initial randomness it always the same\n",
    "z += np.random.rand(mg.number_of_nodes)*rf  # make the noise large enough relative to crater\n",
    "mg.add_field(\"topographic__elevation\", z, at=\"node\", units=\"km\", copy=True, clobber=True);\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## First, we add craters, as before:\n",
    "# Ncraters = int(input(\"How many craters do you want to add per timestep? \")) # see console to add your own number!\n",
    "Ncraters = 100 ## Set the number of craters to add to the surface\n",
    "\n",
    "minD = int(spacing * 3) # want it to be greater than spacing, three times the smallest cell\n",
    "print(\"Minimum possible crater size:\", minD, \"km\")\n",
    "maxD = int((size)/4) #max diameter (m), a quarter of the total domain width\n",
    "print(\"Maximum possible crater size:\", maxD, \"km\")\n",
    "\n",
    "NDs = []\n",
    "for D in range(minD, maxD):\n",
    "    ND = Kx * D ** - delta\n",
    "    NDs.append(ND)\n",
    "\n",
    "random.seed(50) #Chose random seed number 50 (this ensures crater locations are same every time)\n",
    "mg = do_cratering(Ncraters, NDs, minD, maxD, size, mg, spacing) ## This function is defined above\n",
    "\n",
    "fig = plt.figure(dpi=250, facecolor='w', figsize=(8, 3));\n",
    "topo = mg.field_values('node', 'topographic__elevation').reshape((xy, xy)); #create an array of elevation values across the mg domain\n",
    "hs = mg.calc_hillshade_at_node(elevs='topographic__elevation'); #create a hillshade array\n",
    "hill = np.reshape(hs, (xy, xy)); #reshape the hillshade array to be the same shape as the topo array\n",
    "img1 = plt.imshow(hill, cmap=cmap1, alpha=1, extent = [0,size, 0, size], vmin=0, vmax=1);\n",
    "img2 = plt.imshow(topo, cmap=cmap2, alpha=0.4, extent = [0,size, 0, size]);\n",
    "plt.title(\"Cratered Surface\");\n",
    "cbar = plt.colorbar(img2, fraction=0.045, label = \"Elevation [km]\");\n",
    "plt.xlabel('X [km]');\n",
    "plt.ylabel('Y [km] ');\n",
    "plt.xlim([0, size]);\n",
    "plt.ylim([0, size]);\n",
    "plt.tight_layout(); \n",
    "\n",
    "topo_initial = topo.copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## This cell may take some time!\n",
    "########### DIRECT SOME FLOW ACROSS LANDSCAPE ########\n",
    "######################################################\n",
    "\n",
    "## first import the modules needed from LandLab, you can read all about these on ReadTheDocs\n",
    "from landlab.components import FlowDirectorSteepest, FlowAccumulator, FastscapeEroder, DepressionFinderAndRouter\n",
    "from landlab.components import ChannelProfiler\n",
    "\n",
    "## Set a consistent time step for the model runs\n",
    "dt = 1000.0\n",
    "\n",
    "fa = FlowAccumulator(mg, surface='topographic__elevation', flow_director='FlowDirectorD8', \n",
    "                     runoff_rate=None, depression_finder=DepressionFinderAndRouter, routing = 'D8');\n",
    "\n",
    "fa.run_one_step()\n",
    "\n",
    "sp = FastscapeEroder(mg, K_sp=1.0, m_sp=0.5, n_sp=1., threshold_sp=0, \n",
    "                     discharge_field=\"surface_water__discharge\", erode_flooded_nodes = False)\n",
    "\n",
    "sp.run_one_step(dt=dt)\n",
    "\n",
    "fig = plt.figure(dpi=250, facecolor='w', figsize=(8, 3));\n",
    "topo = mg.field_values('node', 'topographic__elevation').reshape((xy, xy)); #create an array of elevation values across the mg domain\n",
    "hs = mg.calc_hillshade_at_node(elevs='topographic__elevation'); #create a hillshade array\n",
    "hill = np.reshape(hs, (xy, xy)); #reshape the hillshade array to be the same shape as the topo array\n",
    "img1 = plt.imshow(hill, cmap=cmap1, alpha=1, extent = [0,size, 0, size], vmin=0, vmax=1);\n",
    "img2 = plt.imshow(topo, cmap=cmap2, alpha=0.4, extent = [0,size, 0, size]);\n",
    "plt.title(\"Evolution of a Cratered surface (+Advection)\");\n",
    "cbar = plt.colorbar(img2, fraction=0.045, label = \"Elevation [km]\");\n",
    "plt.xlabel('X [km]');\n",
    "plt.ylabel('Y [km] ');\n",
    "plt.xlim([0, size]);\n",
    "plt.ylim([0, size]); \n",
    "plt.tight_layout();\n",
    "\n",
    "## Plot a difference map\n",
    "topo_diff = topo_initial - topo\n",
    "fig2 = plt.figure(dpi=250, facecolor='w', figsize=(8, 3));\n",
    "# img1 = plt.imshow(hill, cmap=cmap1, alpha=1, extent = [0,size, 0, size], vmin=0, vmax=1);\n",
    "img2 = plt.imshow(topo_diff, cmap=\"Blues\", alpha=0.4, extent = [0,size, 0, size]);\n",
    "plt.title(\"Surface Difference Map (after advection)\");\n",
    "cbar = plt.colorbar(img2, fraction=0.045, label = \"Elevation Difference [km]\");\n",
    "plt.xlabel('X [km]');\n",
    "plt.ylabel('Y [km] ');\n",
    "plt.xlim([0, size]);\n",
    "plt.ylim([0, size]); \n",
    "plt.tight_layout();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Next: What happens if we diffuse the landscape first?\n",
    "## This cell might take a while to run!\n",
    "#########################\n",
    "\n",
    "## First we should consider the boundary conditions, which boundaries should be open or closed?\n",
    "## Here we have set all the boundaries as open\n",
    "mg.set_closed_boundaries_at_grid_edges(False, False, False, False) #set all as open boundaries\n",
    "\n",
    "## Import the module for Diffusion from Landlab. \n",
    "## There's different landlab modules for this which you can check out!\n",
    "from landlab.components import LinearDiffuser\n",
    "\n",
    "## Set up the linear diffuser model (ld), with a diffusivity coefficient\n",
    "ld = LinearDiffuser(mg, linear_diffusivity=1)\n",
    "\n",
    "#Then run the linear diffused model for one time step (of 100 seconds)\n",
    "for i in range(1):\n",
    "    ld.run_one_step(1.)\n",
    "\n",
    "fig = plt.figure(dpi=250, facecolor='w', figsize=(8, 3));\n",
    "topo = mg.field_values('node', 'topographic__elevation').reshape((xy, xy)); #create an array of elevation values across the mg domain\n",
    "hs = mg.calc_hillshade_at_node(elevs='topographic__elevation'); #create a hillshade array\n",
    "hill = np.reshape(hs, (xy, xy)); #reshape the hillshade array to be the same shape as the topo array\n",
    "img1 = plt.imshow(hill, cmap=cmap1, alpha=1, extent = [0,size, 0, size], vmin=0, vmax=1);\n",
    "img2 = plt.imshow(topo, cmap=cmap2, alpha=0.4, extent = [0,size, 0, size]);\n",
    "plt.title(\"Evolution of a Cratered surface (+Diffusion)\");\n",
    "cbar = plt.colorbar(img2, fraction=0.045, label = \"Elevation [km]\");\n",
    "plt.xlabel('X [km]');\n",
    "plt.ylabel('Y [km] ');\n",
    "plt.xlim([0, size]);\n",
    "plt.ylim([0, size]); \n",
    "plt.tight_layout();\n",
    "\n",
    "## Plot a difference map\n",
    "topo_diff = topo_initial - topo\n",
    "fig2 = plt.figure(dpi=250, facecolor='w', figsize=(8, 3));\n",
    "# img1 = plt.imshow(hill, cmap=cmap1, alpha=1, extent = [0,size, 0, size], vmin=0, vmax=1);\n",
    "img2 = plt.imshow(topo_diff, cmap=\"Blues\", alpha=0.4, extent = [0,size, 0, size]);\n",
    "plt.title(\"Surface Difference Map (after advection)\");\n",
    "cbar = plt.colorbar(img2, fraction=0.045, label = \"Elevation Difference [km]\");\n",
    "plt.xlabel('X [km]');\n",
    "plt.ylabel('Y [km] ');\n",
    "plt.xlim([0, size]);\n",
    "plt.ylim([0, size]); \n",
    "plt.tight_layout();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "########### DIRECT SOME FLOW ACROSS LANDSCAPE ######## (Copied from above: repeat after diffusion of landscape)\n",
    "######################################################\n",
    "\n",
    "## first import the modules needed from LandLab, you can read all about these on ReadTheDocs\n",
    "from landlab.components import FlowDirectorSteepest, FlowAccumulator, FastscapeEroder, DepressionFinderAndRouter\n",
    "from landlab.components import ChannelProfiler\n",
    "\n",
    "## Set a consistent time step for the model runs\n",
    "dt = 1000.0\n",
    "\n",
    "fa = FlowAccumulator(mg, surface='topographic__elevation', flow_director='FlowDirectorD8', \n",
    "                     runoff_rate=None, depression_finder=DepressionFinderAndRouter, routing = 'D8');\n",
    "\n",
    "fa.run_one_step()\n",
    "\n",
    "sp = FastscapeEroder(mg, K_sp=1.0, m_sp=0.5, n_sp=1., threshold_sp=0, \n",
    "                     discharge_field=\"surface_water__discharge\", erode_flooded_nodes = False)\n",
    "\n",
    "sp.run_one_step(dt=dt)\n",
    "\n",
    "fig = plt.figure(dpi=250, facecolor='w', figsize=(8, 3));\n",
    "topo = mg.field_values('node', 'topographic__elevation').reshape((xy, xy)); #create an array of elevation values across the mg domain\n",
    "hs = mg.calc_hillshade_at_node(elevs='topographic__elevation'); #create a hillshade array\n",
    "hill = np.reshape(hs, (xy, xy)); #reshape the hillshade array to be the same shape as the topo array\n",
    "img1 = plt.imshow(hill, cmap=cmap1, alpha=1, extent = [0,size, 0, size], vmin=0, vmax=1);\n",
    "img2 = plt.imshow(topo, cmap=cmap2, alpha=0.4, extent = [0,size, 0, size]);\n",
    "plt.title(\"Evolution of a Cratered surface (+Diffusion & Advection)\");\n",
    "cbar = plt.colorbar(img2, fraction=0.045, label = \"Elevation [km]\");\n",
    "plt.xlabel('X [km]');\n",
    "plt.ylabel('Y [km] ');\n",
    "plt.xlim([0, size]);\n",
    "plt.ylim([0, size]); \n",
    "plt.tight_layout();\n",
    "\n",
    "## Plot a difference map\n",
    "topo_diff = topo_initial - topo\n",
    "fig2 = plt.figure(dpi=250, facecolor='w', figsize=(8, 3));\n",
    "# img1 = plt.imshow(hill, cmap=cmap1, alpha=1, extent = [0,size, 0, size], vmin=0, vmax=1);\n",
    "img2 = plt.imshow(topo_diff, cmap=\"Blues\", alpha=0.4, extent = [0,size, 0, size]);\n",
    "plt.title(\"Surface Difference Map (after diffusion)\");\n",
    "cbar = plt.colorbar(img2, fraction=0.045, label = \"Elevation Difference [km]\");\n",
    "plt.xlabel('X [km]');\n",
    "plt.ylabel('Y [km] ');\n",
    "plt.xlim([0, size]);\n",
    "plt.ylim([0, size]); \n",
    "plt.tight_layout();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.9.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}