Code_Part1.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Code_Part1",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/erica-mccormick/widespread-bedrock-water-use/blob/main/Code_Part1.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_a5SRV9N_kSO"
      },
      "source": [
        "[![DOI](https://zenodo.org/badge/357586333.svg)](https://zenodo.org/badge/latestdoi/357586333)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RfRICAHKt6_a"
      },
      "source": [
        "# **PART 1: Code for *Evidence for widespread woody plant use of water stored in bedrock***\n",
        "**Erica McCormick, David Dralle, W. Jesse Hahm, Alison Tune, Logan Schmidt, Dana Chadwick, and Daniella Rempe**\n",
        "\n",
        "---\n",
        "\n",
        "All data products are available in the [Hydroshare repository](https://doi.org/10.4211/hs.a2f0d5fd10f14cd189a3465f72cba6f3) and as GEE assets (Images and ImageCollections, see Part 2).\n",
        "\n",
        "See [Github](https://github.com/erica-mccormick/widespread-bedrock-water-use) or [Zenodo](http://doi.org/10.5281/zenodo.4904036) for Part 2 and more information on data inputs and products. \n",
        "\n",
        "For more information, see [website](https://erica-mccormick.github.io/widespread-bedrock-water-use/).\n",
        "\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "spb3JM2kePfK"
      },
      "source": [
        "#**Calculate $D_{bedrock}$ and $S_{bedrock}$ for CONUS**\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UQCxr-rE2jcT"
      },
      "source": [
        "Press shift+enter on each chunk while highlighted to run code. The output rasters from this script are $S_r$ and $D_{max}$ for CONUS. $S_{bedrock}$ and $D_{bedrock, 2004-2017}$ can be calculated by subtracting $S_{soil}$ (see Methods). The products from this script can be found on [hydroshare](https://www.hydroshare.org/resource/e7ad140edaf54d69ba4f1cf1ec8e7f73/). The protocol for $S_r$ follows that established by [Wang-Erlandsson et al. (2016)](https://doi.org/10.5194/hess-20-1459-2016) and [Dralle et al. (2020)](https://10.5194/hess-25-2861-2021)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IU23BHGj2TIP"
      },
      "source": [
        "### Authenticate Google Earth Engine (GEE) account. First time users sign-up [here](https://earthengine.google.com/new_signup/)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "both",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yxBXo7unYLY1",
        "outputId": "ce4b761a-5485-41e3-b6af-d05f7908bbb1"
      },
      "source": [
        "import ee\n",
        "ee.Authenticate()\n",
        "ee.Initialize()"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "To authorize access needed by Earth Engine, open the following URL in a web browser and follow the instructions. If the web browser does not start automatically, please manually browse the URL below.\n",
            "\n",
            "    https://accounts.google.com/o/oauth2/auth?client_id=517222506229-vsmmajv00ul0bs7p89v5m89qs8eb9359.apps.googleusercontent.com&scope=https%3A%2F%2Fwww.googleapis.com%2Fauth%2Fearthengine+https%3A%2F%2Fwww.googleapis.com%2Fauth%2Fdevstorage.full_control&redirect_uri=urn%3Aietf%3Awg%3Aoauth%3A2.0%3Aoob&response_type=code&code_challenge=oSXBjESbXcvL1R1R9m0G-yfQaWugDfOEKXqVHtYd_Z4&code_challenge_method=S256\n",
            "\n",
            "The authorization workflow will generate a code, which you should paste in the box below. \n",
            "Enter verification code: 4/1AY0e-g4_hoiNN47YIcNmC7-Aie5O70XrOXJOr-x-Jrql-xN0spNWDzuGQCs\n",
            "\n",
            "Successfully saved authorization token.\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yYN4ZkE63Ns3"
      },
      "source": [
        "### GEE assets used in this script:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aA43TR1J2ATp"
      },
      "source": [
        "Publically Available Assets:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pYaSonOI2C28"
      },
      "source": [
        "prism = ee.ImageCollection(\"OREGONSTATE/PRISM/AN81d\")\n",
        "pml = ee.ImageCollection(\"CAS/IGSNRR/PML/V2\")\n",
        "USGS_landcover = ee.ImageCollection(\"USGS/NLCD\")\n",
        "modis_landcover = ee.ImageCollection(\"MODIS/006/MCD12Q1\")\n",
        "snow_cover = ee.ImageCollection(\"MODIS/006/MOD10A1\").select('NDSI_Snow_Cover')\n",
        "bess = ee.Image('users/daviddralle/bessv2')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EL4NxLNQ2DK5"
      },
      "source": [
        "Personal Assets:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LFvuxVJf1YkM"
      },
      "source": [
        "# Regions\n",
        "ca = ee.Feature(ee.FeatureCollection(\"users/daviddralle/ca_et/CA\").first())\n",
        "conus = ee.FeatureCollection('users/ericamccormick/20_RockMoisture/geometries/conus_20m')\n",
        "texas = ee.FeatureCollection('users/ericamccormick/20_RockMoisture/geometries/TX')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3697CljqnUM-"
      },
      "source": [
        "# gNATSGO derived (Soil Survey Staff, 2019):\n",
        "us_soil = ee.Image('users/ericamccormick/20_RockMoisture/products/gNATSGO/Ssoil_500m')\n",
        "densic = ee.Image('users/ericamccormick/20_RockMoisture/products/gNATSGO/densic_weavg_lower_reprojected').reproject(crs='EPSG:4326',scale=90) # Soil Survey Staff, 2020\n",
        "paralithic = ee.Image('users/ericamccormick/20_RockMoisture/products/gNATSGO/paralithic_weavg_lower_reprojected').reproject(crs='EPSG:4326',scale=90) # Soil Survey Staff, 2020\n",
        "lithic = ee.Image(\"users/ericaelmstead/20_RockMoisture/gNATSGO/lithic_weavg_lower_reprojected\").reproject(crs='EPSG:4326',scale=90) # Soil Survey Staff, 2020"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bGHwLp65vCia"
      },
      "source": [
        "### Load packages"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "b_1N3c5DWmW9"
      },
      "source": [
        "import datetime\n",
        "from IPython.display import Image\n",
        "import pandas as pd\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import matplotlib as mpl\n",
        "plt.rcParams['pdf.fonttype'] = 42\n",
        "from google.colab import files\n",
        "from IPython.display import Image\n",
        "%matplotlib inline"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OcMXqg8ncU2I"
      },
      "source": [
        "For running analyses in GoogleColab, download datasets and .tiffs to GoogleDrive and import directly from there. To mount your Drive, find the file name for the drive folder you wish to connect to:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gHFxvSiWcUHj",
        "outputId": "729185df-1e63-4548-cc61-85c6088452fd"
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Mounted at /content/drive\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C6AxmwYzc8ee"
      },
      "source": [
        "#**Create woody vegetation and shallow bedrock masks**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7giCys28u91G",
        "outputId": "4e082404-ef58-44ed-9a1a-b835db18f3a2"
      },
      "source": [
        "# Get the crs of the layer you want to match (here using prism as a test)\n",
        "prism_image = ee.ImageCollection(\"OREGONSTATE/PRISM/AN81d\").first().reproject(crs='EPSG:4326', scale = 500)\n",
        "\n",
        "prism_projection = prism_image.projection();\n",
        "prism_projection.getInfo()"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'crs': 'EPSG:4326',\n",
              " 'transform': [0.004491576420597608, 0, 0, 0, -0.004491576420597608, 0],\n",
              " 'type': 'Projection'}"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6S2E8NajwbvB"
      },
      "source": [
        "## Get binary soil depth > 1.5 and all soil depth pixels = 1 layers"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TjitxEBHwAjL"
      },
      "source": [
        "# Get all of the bedrock pixels = 1 for counting\n",
        "all_bedrock_pixels = paralithic.add(densic).add(lithic)\n",
        "all_bedrock_pixels = all_bedrock_pixels.gt(-1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "b9HIX2Xg2ATL"
      },
      "source": [
        "# mask places with greater than 1.5m soil\n",
        "depth = 150\n",
        "\n",
        "m = paralithic.select('b1').lte(depth)\n",
        "m1 = densic.select('b1').lte(depth)\n",
        "m2 = lithic.select('b1').lte(depth)\n",
        "\n",
        "bedrock_depth_include = m.add(m1).add(m2) #could add.gt(0) here to make sure all values are one, but we are counting so its not a big deal plus theres very little to no overlap between bedrock types\n",
        "\n",
        "included_depth = all_bedrock_pixels.updateMask(bedrock_depth_include)\t"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hLg2Hpw_uWqT"
      },
      "source": [
        "## Coarsen resolution using sum to 500m"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_fWpDWQ3wg-y"
      },
      "source": [
        "\n",
        "count_all = all_bedrock_pixels.reduceResolution(**{\n",
        "  'reducer': ee.Reducer.count()\n",
        "  }).reproject(crs=prism_projection)\n",
        "\n",
        "\n",
        "count_under_threshold = included_depth.reduceResolution(**{\n",
        "  'reducer': ee.Reducer.count()\n",
        "  }).reproject(crs=prism_projection)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BlPgA-91nl-Y"
      },
      "source": [
        "fraction = count_under_threshold.divide(count_all).multiply(100)\n",
        "keep_shallowbedrock = fraction.gte(100)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6uF0ABfYc5Ox"
      },
      "source": [
        "## Export"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "T2iNeEeItVxd"
      },
      "source": [
        "task_config = {\n",
        "        'region': conus.geometry(),\n",
        "        'fileFormat': 'GeoTIFF',\n",
        "        'fileNamePrefix': 'mask_shallowBedrock',\n",
        "        'image': keep_shallowbedrock.toDouble().clip(conus.geometry()),\n",
        "        'description': 'mask_shallowBedrock',\n",
        "        'scale': 500,\n",
        "        'maxPixels': 10000000000000\n",
        "    }\n",
        "\n",
        "task=ee.batch.Export.image.toDrive(**task_config)\n",
        "task.start()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hzcRXDP2dRFG"
      },
      "source": [
        "task_config = {\n",
        "        'region': conus.geometry(),\n",
        "        'image': keep_shallowbedrock,\n",
        "        'description': 'shallowBedrock',\n",
        "        'assetId' : 'users/ericamccormick/mask_shallowBedrock',\n",
        "        'scale': 500,\n",
        "        'maxPixels': 10000000000000\n",
        "    }\n",
        "\n",
        "task=ee.batch.Export.image.toAsset(**task_config)\n",
        "task.start()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "05RV_ivwhYgn"
      },
      "source": [
        "## Repeat for NLCD 30m forest and shrubland classifications"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uUM7CYpvhcZS"
      },
      "source": [
        "usgs_nlcd = ee.ImageCollection(\"USGS/NLCD_RELEASES/2016_REL\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VlbGQNRtjZD_"
      },
      "source": [
        "landsort = usgs_nlcd.sort('system:time_start', False) #sort backcwards to get most recent image\n",
        "landsort = ee.Image(landsort.first()) # make sure to get just first image\n",
        "img1 = landsort.select('landcover').reproject(crs='EPSG:4326',scale=30)\n",
        "\n",
        "## Get shrub at bedrock.lt(depth): \n",
        "shrub = img1.remap([52],[1],0) #eq(52) #mask to get all values = 52\n",
        "#shrub_final = shrub.multiply(ee.Image(2)) # set shrubs as value 2, everything else 0\n",
        "\n",
        "# get classes so forested is 0s with 1s where values 41,42,43. Use function 'reclassify'\n",
        "forested = img1.remap([41,42,43],[1,1,1],0) #forested.remap([from],[to],defaultVal) , defaultval sets everything else to zerofore\n",
        "\n",
        "## add forested_final and shrub_final and return a mask with 1s and 0s:\n",
        "woody_veg = ee.Image(0).add(forested).add(shrub) #.gt(0) #final binary tiff where 1s are places with woody veg"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kF6c24RZoYdH"
      },
      "source": [
        "woody_pixels_only = img1.updateMask(woody_veg)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MK3aXVcioYdJ"
      },
      "source": [
        "count_landcover_all = img1.reduceResolution(**{\n",
        "  'reducer': ee.Reducer.count(),\n",
        "  'maxPixels': 500\n",
        "  }).reproject(crs=prism_projection)\n",
        "\n",
        "\n",
        "count_woody = woody_pixels_only.reduceResolution(**{\n",
        "  'reducer': ee.Reducer.count(),\n",
        "  'maxPixels':500\n",
        "  }).reproject(crs=prism_projection)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6zeb4nsyoqws"
      },
      "source": [
        "## do division and thresholding here\n",
        "fraction_woody = count_woody.divide(count_landcover_all).multiply(100)\n",
        "keep_woodyveg = fraction_woody.gte(75)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-0NTeIVij3J9"
      },
      "source": [
        "task_config = {\n",
        "        'region': conus.geometry(),\n",
        "        'fileFormat': 'GeoTIFF',\n",
        "        'fileNamePrefix': 'mask_woodyVeg',\n",
        "        'image': keep_woodyveg.toDouble().clip(conus.geometry()),\n",
        "        'description': 'mask_woodyVeg',\n",
        "        'scale': 500,\n",
        "        'maxPixels': 10000000000000\n",
        "    }\n",
        "\n",
        "task=ee.batch.Export.image.toDrive(**task_config)\n",
        "task.start()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZePtFwshbCEx"
      },
      "source": [
        "task_config = {\n",
        "        'region': conus.geometry(),\n",
        "        'image': keep_woodyveg,\n",
        "        'description': 'mask_woodyVeg',\n",
        "        'assetId' : 'users/ericamccormick/mask_woodyVeg',\n",
        "        'scale': 500,\n",
        "        'maxPixels': 10000000000000\n",
        "    }\n",
        "\n",
        "task=ee.batch.Export.image.toAsset(**task_config)\n",
        "task.start()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Zqf9tViK2spo"
      },
      "source": [
        "#**Calculate $S_{r}$ for CONUS**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sVPHlzcx-LIX"
      },
      "source": [
        "### Assuming $F_{out}=ET$, calculate $F_{out}$ time series without snow correction (following [Wang-Erlandsson et al (2016)](https://hess.copernicus.org/articles/20/1459/2016/)) and with a snow correction."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ScABVFuLRT8g"
      },
      "source": [
        "#Flux is stored as mm/day, and duration between samples is 8 days\n",
        "multiplier = 8.0 \n",
        "\n",
        "#### USING SNOW CORRECTION METHOD ####\n",
        "# threshold of non-snow-fraction above which we assume that there is \n",
        "# negligible snow in the image\n",
        "threshold = 0.9\n",
        "# mapping function for creating ET time series but set to zero (Fout = 0) during\n",
        "# periods when snow is present\n",
        "def sumpml_snow(image):\n",
        "  #First, get percent snow coverage for nearest datetime\n",
        "  current_datetime = datetime.datetime.now().strftime('%Y-%m-%d')\n",
        "  nearest_snow_image = snow_cover.filterDate(image.get('system:time_start'), \n",
        "                                             current_datetime).first() \n",
        "  #reproject\n",
        "  nearest_snow_image= nearest_snow_image.reproject(crs='EPSG:4326',scale=500)\n",
        "  non_snow_frac = ee.Image(1).subtract(nearest_snow_image.divide(100))\n",
        "  #assign non-snow fraction to 1 if nodata\n",
        "  non_snow_frac=non_snow_frac.unmask(1)\n",
        "\n",
        "  # zero multiplier if there is less than threshold non snow frac\n",
        "  # i.e. if there is significant snow in the image\n",
        "  snow_multiplier = non_snow_frac.gt(threshold)\n",
        "\n",
        "  # get total surface evap and transpiration, zero out if significant snow\n",
        "  temp = image.select('Es').add(image.select('Ec')).multiply(snow_multiplier)\n",
        "\n",
        "  #get the first band of this new temporary image, rename it to 'ET', \n",
        "  #reproject it, and then multiply by 8\n",
        "  temp = temp.select([0], ['ET']).reproject(\n",
        "      crs='EPSG:4326',scale=500\n",
        "      ).multiply(multiplier)\n",
        "  #temp should now be total ET in mm over the time window between images\n",
        "  #assign the datetime stamp and the index from the original image\n",
        "  temp = temp.set('system:time_start', image.get('system:time_start'))\n",
        "  temp = temp.set('system:index', image.get('system:index'))\n",
        "  return temp\n",
        "\n",
        "#now, actually map the pml image collection w/ this function,\n",
        "# in order to make a new combined ET image collection\n",
        "et_snow_corr = ee.ImageCollection(pml.map(sumpml_snow).select('ET'))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "d8x4lRcoI8PL"
      },
      "source": [
        "# function use to calculate A_{t_n} values and running storage deficit\n",
        "def get_AD(et, yearstart, yearend):\n",
        "  et_current = et.filterDate(yearstart, yearend)\n",
        "  prism = ee.ImageCollection('OREGONSTATE/PRISM/AN81d').select('ppt')\n",
        "\n",
        "  # Create starting image for iterator used to calculate A_{t_n}\n",
        "  starter = ee.Image(et_current.first())\n",
        "  # set the time of the starter image to 8 days before first ET image\n",
        "  time0 = starter.get('system:time_start').getInfo() - 691000000\n",
        "  first = ee.List([\n",
        "    ee.Image(0).set('system:time_start', time0).select([0], ['A']).toDouble()\n",
        "  ])\n",
        "\n",
        "  # for each item in the ET (F_out) time series, \n",
        "  # get cumulative precip (F_in) since previous F_out\n",
        "  # Take the difference between F_out and F_in; this is A\n",
        "  def getA(image, thelist):\n",
        "    previous = ee.Image(ee.List(thelist).get(-1))\n",
        "    startdate = previous.get('system:time_start')\n",
        "    enddate = image.get('system:time_start')\n",
        "    windowed_prism = prism.filterDate(startdate,enddate)\n",
        "    prism_total = windowed_prism.reduce(ee.Reducer.sum())\n",
        "    Atemp = image.subtract(prism_total)\n",
        "    Atemp = Atemp.select([0], ['A']).toDouble()\n",
        "    Atemp = Atemp.set({'system:time_start':image.get('system:time_start')})\n",
        "    return ee.List(thelist).add(Atemp)\n",
        "  # iterate using getA() \n",
        "  A = ee.ImageCollection(ee.List(et_current.iterate(getA, first)))\n",
        "\n",
        "  # add band to keep track of running storage deficit calculation\n",
        "  def add_deficit_band(image):\n",
        "    return image.addBands(ee.Image.constant(0).rename('D')).toDouble()\n",
        "  AD_initial = A.map(add_deficit_band)\n",
        "  starter = ee.Image(AD_initial.first()).multiply(0)\n",
        "\n",
        "  ts = AD_initial.first().get('system:time_start')\n",
        "  starter = starter.set('system:time_start', ts)\n",
        "\n",
        "  first = ee.List([starter])\n",
        "  # iterator we will use to calculate running storage deficit\n",
        "  def accumulate(image, thelist):\n",
        "    # get image with previous D values\n",
        "    previous = ee.Image(ee.List(thelist).get(-1))\n",
        "    previousD = previous.select('D')\n",
        "\n",
        "    # get current value of A = F_out - F_in\n",
        "    currentA = image.select('A')\n",
        "\n",
        "    # D is just previous D plus current A\n",
        "    newD = previousD.add(currentA)\n",
        "\n",
        "    # if storage deficit is < 0, set it to zero (can only be positive)\n",
        "    newD = newD.multiply(newD.gt(0))\n",
        "\n",
        "    # add to list over which we are iterating\n",
        "    tempAD = currentA.addBands(newD)\n",
        "    return ee.List(thelist).add(tempAD)\n",
        "\n",
        "  # Get storage deficit time series, take maximum observed storage deficit\n",
        "  AD = ee.ImageCollection(ee.List(AD_initial.iterate(accumulate, first)))\n",
        "  return AD"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JGD9rCRo4H1L"
      },
      "source": [
        "### Use `get_AD()` function to calculate max storage deficits"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "g-h3v-1DmK-g"
      },
      "source": [
        "# get max storage deficit using non-snow-corrected and snow-corrected F_out\n",
        "yearstart = '2003'\n",
        "yearend = '2017'\n",
        "\n",
        "AD_snow_corr = get_AD(et_snow_corr, yearstart, yearend)\n",
        "\n",
        "# Equivelant to Sr\n",
        "maxD_snow_corr = AD_snow_corr.select('D').max().select([0],['S_R_snow_corr'])\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZMW0Xp7AeJFl"
      },
      "source": [
        "### Apply woody veg and shallow bedrock masks and add in additional ET>P mask"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_m8wpoq2eNVe"
      },
      "source": [
        "# ET>P\n",
        "etsum = et_snow_corr.filterDate('2003', '2017').sum().clip(conus.geometry()).reproject(crs='EPSG:4326',scale=500)\n",
        "prism = ee.ImageCollection('OREGONSTATE/PRISM/AN81m').select('ppt')\n",
        "pptsum = prism.filterDate('2003', '2017').sum().clip(conus.geometry()).reproject(crs='EPSG:4326',scale=500)\n",
        "ppt_minus_et = pptsum.subtract(etsum)\n",
        "mask_ET_gt_PPT = ppt_minus_et.gte(0)\n",
        "\n",
        "# Multiply masks together\n",
        "mask = mask_ET_gt_PPT.multiply(keep_woodyveg).multiply(keep_shallowbedrock)\n",
        "\n",
        "mask_reproj = mask.reproject(\n",
        "      crs='EPSG:4326',scale=500)\n",
        "\n",
        "# create multiband raster with S_R values and snowpack band\n",
        "# mask by land cover type\n",
        "#to_save = maxD.addBands(maxD_snow_corr).addBands(mean_winter_snow_cover).addBands(abs_diff)\n",
        "\n",
        "# add difference between snowb\n",
        "#Sr_tosave = maxD_snow_corr.updateMask(mask_reproj)\n",
        "Sr = maxD_snow_corr.updateMask(mask_reproj).clip(conus.geometry())\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pKyubnx6ftqQ"
      },
      "source": [
        "task_config = {\n",
        "        'region': conus.geometry(),\n",
        "        'fileFormat': 'GeoTIFF',\n",
        "        'fileNamePrefix': 'Sr',\n",
        "        'image': Sr.toDouble(),\n",
        "        'description': 'pptsum2',\n",
        "        'scale': 500,\n",
        "        'maxPixels': 10000000000000\n",
        "    }\n",
        "\n",
        "task=ee.batch.Export.image.toDrive(**task_config)\n",
        "task.start()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yg7samGy4otF"
      },
      "source": [
        "## Export\n",
        "\n",
        "Only run this cell if you intend to export the dataset from Earth Engine to your Google Drive. This operation will take several hours to complete. You can check the status of the export on [https://code.earthengine.google.com/](https://code.earthengine.google.com/)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "q33lW8WB4qWi"
      },
      "source": [
        "task_config = {\n",
        "        'region': conus.geometry(),\n",
        "        'fileFormat': 'GeoTIFF',\n",
        "        'fileNamePrefix': 'Sr_' + str(yearstart) + '_' + str(yearend) + '_Mask',\n",
        "        'image': Sr.toDouble(),\n",
        "        'description': 'Sr',\n",
        "        'scale': 500,\n",
        "        'maxPixels': 10000000000000\n",
        "    }\n",
        "\n",
        "task=ee.batch.Export.image.toDrive(**task_config)\n",
        "task.start()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iM09IN02aYcc"
      },
      "source": [
        "task_config = {\n",
        "        'region': conus.geometry(),\n",
        "        'image': Sr,\n",
        "        'description': 'Sr_mask',\n",
        "        'assetId' : 'users/ericamccormick/Sr' + str(yearstart) + '_' + str(yearend) + '_Mask',\n",
        "        'scale': 500,\n",
        "        'maxPixels': 10000000000000\n",
        "    }\n",
        "\n",
        "task=ee.batch.Export.image.toAsset(**task_config)\n",
        "task.start()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OQThYrgGj84D"
      },
      "source": [
        "#**Calculate $D_{max, 2003-2017}$ for CONUS**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jKrHFzZeSnLi"
      },
      "source": [
        "yearstart = '2003'\n",
        "yearend = '2017'\n",
        "\n",
        "for year in range(2003,2017):\n",
        "  yearstart = str(year) + '-10-01'\n",
        "  yearend = str(year+1) + '-10-01'\n",
        "\n",
        "  AD_snow_corr = get_AD(et_snow_corr, yearstart, yearend)\n",
        "\n",
        "  def add_time_band(image):\n",
        "    t = ee.Number(image.get('system:time_start')).toLong() #changed from .toLong()\n",
        "    timeimage = ee.Image(t).select([0],['t'])\n",
        "    image = image.addBands(timeimage)\n",
        "    image = image.toLong()\n",
        "    return image\n",
        "\n",
        "  AD_snow_corr = AD_snow_corr.map(add_time_band, opt_dropNulls=True)\n",
        "\n",
        "  def greater_than_zero(image):\n",
        "    temp = image.select('D').gt(0).select([0],['D_positive']).toLong() \n",
        "    image = image.addBands(temp)\n",
        "    return image\n",
        "\n",
        "  D_positive = AD_snow_corr.select(['t','D']).map(greater_than_zero)\n",
        "  # D_positive has three bands: t, D, and D_positive\n",
        "\n",
        "  # // create a starting image for iterator\n",
        "  starter = ee.Image(D_positive.first())\n",
        "  starter = starter.addBands(ee.Image(0).select([0],['cumulative_time']))\n",
        "  first = ee.List([starter])\n",
        "\n",
        "  def accumulate(image, thelist):\n",
        "    previous = ee.Image(ee.List(thelist).get(-1))\n",
        "    # // accumulate and reset counter where image=0\n",
        "    added = image.select('D_positive').add(previous.select('cumulative_time')).multiply(image.select('D_positive'))\n",
        "    added = added.select([0],['cumulative_time'])\n",
        "    image = image.addBands(added)\n",
        "    # image has 4 bands t, D, D_positive, and cumulative_time\n",
        "    return ee.List(thelist).add(image)\n",
        "\n",
        "  cumulative_list = ee.List(D_positive.iterate(accumulate, first))\n",
        "  cumulative_list = cumulative_list.remove(cumulative_list.get(0))\n",
        "  # get rid of first item in list, keep everything else, convert to imagecollection\n",
        "  cumulative = ee.ImageCollection(cumulative_list)\n",
        "  # argmax = cumulative.qualityMosaic('D').select(['t','D','cumulative_time'])\n",
        "  argmax = cumulative.qualityMosaic('D').select(['t','cumulative_time','D'])\n",
        "\n",
        "  to_save = argmax.updateMask(mask_reproj).clip(conus.geometry())\n",
        "  to_save = to_save.toDouble() # Must comment this out to send to asset.\n",
        "\n",
        "  task_config = {'region': conus.geometry(),\n",
        "          'fileFormat': 'GeoTIFF',\n",
        "          'fileNamePrefix': 'Dmax_' + str(year+1) + '_Mask', \n",
        "          'image': to_save,\n",
        "          'description': 'Dmax_' + str(year+1) + '_Mask',\n",
        "          'scale': 500,\n",
        "          'maxPixels': 10000000000000\n",
        "      }\n",
        "\n",
        "  task=ee.batch.Export.image.toDrive(**task_config)\n",
        "  task.start()\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C9YiTPESs3Gv"
      },
      "source": [
        "Export annual to GEE asset:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eT-RS9TkW5bp"
      },
      "source": [
        "yearstart = '2003'\n",
        "yearend = '2017'\n",
        "\n",
        "for year in range(2003,2017):\n",
        "  yearstart = str(year) + '-10-01'\n",
        "  yearend = str(year+1) + '-10-01'\n",
        "\n",
        "  AD_snow_corr = get_AD(et_snow_corr, yearstart, yearend)\n",
        "\n",
        "  def add_time_band(image):\n",
        "    t = ee.Number(image.get('system:time_start')).toLong() #changed from .toLong()\n",
        "    timeimage = ee.Image(t).select([0],['t'])\n",
        "    image = image.addBands(timeimage)\n",
        "    image = image.toLong()\n",
        "    return image\n",
        "\n",
        "  AD_snow_corr = AD_snow_corr.map(add_time_band, opt_dropNulls=True)\n",
        "\n",
        "  def greater_than_zero(image):\n",
        "    temp = image.select('D').gt(0).select([0],['D_positive']).toLong() \n",
        "    image = image.addBands(temp)\n",
        "    return image\n",
        "\n",
        "  D_positive = AD_snow_corr.select(['t','D']).map(greater_than_zero)\n",
        "  # D_positive has three bands: t, D, and D_positive\n",
        "\n",
        "  # // create a starting image for iterator\n",
        "  starter = ee.Image(D_positive.first())\n",
        "  starter = starter.addBands(ee.Image(0).select([0],['cumulative_time']))\n",
        "  first = ee.List([starter])\n",
        "\n",
        "  def accumulate(image, thelist):\n",
        "    previous = ee.Image(ee.List(thelist).get(-1))\n",
        "    # // accumulate and reset counter where image=0\n",
        "    added = image.select('D_positive').add(previous.select('cumulative_time')).multiply(image.select('D_positive'))\n",
        "    added = added.select([0],['cumulative_time'])\n",
        "    image = image.addBands(added)\n",
        "    # image has 4 bands t, D, D_positive, and cumulative_time\n",
        "    return ee.List(thelist).add(image)\n",
        "\n",
        "  cumulative_list = ee.List(D_positive.iterate(accumulate, first))\n",
        "  cumulative_list = cumulative_list.remove(cumulative_list.get(0))\n",
        "  # get rid of first item in list, keep everything else, convert to imagecollection\n",
        "  cumulative = ee.ImageCollection(cumulative_list)\n",
        "  # argmax = cumulative.qualityMosaic('D').select(['t','D','cumulative_time'])\n",
        "  argmax = cumulative.qualityMosaic('D').select(['t','cumulative_time','D'])\n",
        "\n",
        "  to_save = argmax.updateMask(mask_reproj)\n",
        "\n",
        "  to_save = to_save.clip(conus.geometry()) # added 10/12/20 by Erica from David-Daniella chat\n",
        "\n",
        "  task_config = {\n",
        "          'region': conus.geometry(),\n",
        "          'image': to_save,\n",
        "          'description': 'US_ANNUALargmax_allrock150_' + str(year+1),\n",
        "          'assetId' : 'users/ericaelmstead/US_ANNUALargmax_allrock150_' + str(year+1),\n",
        "          'scale': 1000,\n",
        "          'maxPixels': 10000000000000\n",
        "      }\n",
        "\n",
        "  task=ee.batch.Export.image.toAsset(**task_config)\n",
        "  task.start()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TZ3fmylLi0x_"
      },
      "source": [
        "#**Subtract $S_{soil}$ to calculate $S_{bedrock}$ and $D_{bedrock,Y}$**\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HY-z4uRpkEAW"
      },
      "source": [
        "###Get $S_{soil}$"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5AlDBQnljbxZ"
      },
      "source": [
        "ssoil_mm = ee.Image(\"users/ericaelmstead/20_RockMoisture/gNATSGO/US_AWS_reprojected\").clip(conus.geometry()).multiply(10).reproject(crs = 'EPSG:4326', scale = 90)\n",
        "\n",
        "# Specify for mean soil to be used when scaling to 500m\n",
        "mean_soil = ssoil_mm.reduceResolution(**{\n",
        "  'reducer': ee.Reducer.mean()\n",
        "})\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3U6tD6_Tj621"
      },
      "source": [
        "task_config = {\n",
        "        'region': conus.geometry(),\n",
        "        'fileFormat': 'GeoTIFF',\n",
        "        'fileNamePrefix': 'Ssoil_500m',\n",
        "        'image': mean_soil.toDouble(),\n",
        "        'description': 'Ssoil_500m',\n",
        "        'scale': 500,\n",
        "        'maxPixels': 10000000000000\n",
        "    }\n",
        "\n",
        "task=ee.batch.Export.image.toDrive(**task_config)\n",
        "task.start()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iGwpEMf6j62_"
      },
      "source": [
        "task_config = {\n",
        "        'region': conus.geometry(),\n",
        "        'image': mean_soil,\n",
        "        'description': 'Ssoil_500m',\n",
        "        'assetId' : 'users/ericamccormick/Ssoil_500m',\n",
        "        'scale': 500,\n",
        "        'maxPixels': 10000000000000\n",
        "    }\n",
        "\n",
        "task=ee.batch.Export.image.toAsset(**task_config)\n",
        "task.start()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uai6FYLAkHw8"
      },
      "source": [
        "###Get $S_{bedrock}$"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EJW8MK1kkJ4Q"
      },
      "source": [
        "# Sbedrock = Sr - Ssoil\n",
        "Sbedrock = Sr.subtract(mean_soil).clip(conus.geometry())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CGkOjHLvkJ4Q"
      },
      "source": [
        "task_config = {\n",
        "        'region': conus.geometry(),\n",
        "        'fileFormat': 'GeoTIFF',\n",
        "        'fileNamePrefix': 'Sbedrock',\n",
        "        'image': Sbedrock.toDouble(),\n",
        "        'description': 'Sbedrock',\n",
        "        'scale': 500,\n",
        "        'maxPixels': 10000000000000\n",
        "    }\n",
        "\n",
        "task=ee.batch.Export.image.toDrive(**task_config)\n",
        "task.start()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QubTYt1dkJ4R"
      },
      "source": [
        "task_config = {\n",
        "        'region': conus.geometry(),\n",
        "        'image': Sbedrock,\n",
        "        'description': 'Sbedrock',\n",
        "        'assetId' : 'users/ericamccormick/Sbedrock',\n",
        "        'scale': 500,\n",
        "        'maxPixels': 10000000000000\n",
        "    }\n",
        "\n",
        "task=ee.batch.Export.image.toAsset(**task_config)\n",
        "task.start()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uDRhnFB3lfDt"
      },
      "source": [
        "### Get $D_{bedrock,Y}$"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uNRliHDvlhui"
      },
      "source": [
        "yearstart = '2003'\n",
        "yearend = '2017'\n",
        "\n",
        "for year in range(2003,2017):\n",
        "  yearstart = str(year) + '-10-01'\n",
        "  yearend = str(year+1) + '-10-01'\n",
        "\n",
        "  AD_snow_corr = get_AD(et_snow_corr, yearstart, yearend)\n",
        "\n",
        "  def add_time_band(image):\n",
        "    t = ee.Number(image.get('system:time_start')).toLong() #changed from .toLong()\n",
        "    timeimage = ee.Image(t).select([0],['t'])\n",
        "    image = image.addBands(timeimage)\n",
        "    image = image.toLong()\n",
        "    return image\n",
        "\n",
        "  AD_snow_corr = AD_snow_corr.map(add_time_band, opt_dropNulls=True)\n",
        "\n",
        "  def greater_than_zero(image):\n",
        "    temp = image.select('D').gt(0).select([0],['D_positive']).toLong() \n",
        "    image = image.addBands(temp)\n",
        "    return image\n",
        "\n",
        "  D_positive = AD_snow_corr.select(['t','D']).map(greater_than_zero)\n",
        "  # D_positive has three bands: t, D, and D_positive\n",
        "\n",
        "  # // create a starting image for iterator\n",
        "  starter = ee.Image(D_positive.first())\n",
        "  starter = starter.addBands(ee.Image(0).select([0],['cumulative_time']))\n",
        "  first = ee.List([starter])\n",
        "\n",
        "  def accumulate(image, thelist):\n",
        "    previous = ee.Image(ee.List(thelist).get(-1))\n",
        "    # // accumulate and reset counter where image=0\n",
        "    added = image.select('D_positive').add(previous.select('cumulative_time')).multiply(image.select('D_positive'))\n",
        "    added = added.select([0],['cumulative_time'])\n",
        "    image = image.addBands(added)\n",
        "    # image has 4 bands t, D, D_positive, and cumulative_time\n",
        "    return ee.List(thelist).add(image)\n",
        "\n",
        "  cumulative_list = ee.List(D_positive.iterate(accumulate, first))\n",
        "  cumulative_list = cumulative_list.remove(cumulative_list.get(0))\n",
        "  # get rid of first item in list, keep everything else, convert to imagecollection\n",
        "  cumulative = ee.ImageCollection(cumulative_list)\n",
        "  # argmax = cumulative.qualityMosaic('D').select(['t','D','cumulative_time'])\n",
        "  argmax = cumulative.qualityMosaic('D').select(['t','cumulative_time','D'])\n",
        "\n",
        "  dmax = argmax.select('D').updateMask(mask_reproj)\n",
        "  dbedrock = dmax.subtract(mean_soil).clip(conus.geometry())\n",
        "\n",
        "  task_config = {'region': conus.geometry(),\n",
        "          'fileFormat': 'GeoTIFF',\n",
        "          'fileNamePrefix': 'Dbedrock_' + str(year+1) + '_nomask', \n",
        "          'image': dbedrock.toDouble(),\n",
        "          'description': 'Dbedrock' + str(year+1) + '_nomask',\n",
        "          'scale': 500,\n",
        "          'maxPixels': 10000000000000\n",
        "      }\n",
        "\n",
        "  task=ee.batch.Export.image.toDrive(**task_config)\n",
        "  task.start()\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "stX5PTE20wmF"
      },
      "source": [
        "#**References**\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y6JXSnsk01c-"
      },
      "source": [
        "\n",
        "Daly, C., Halbleib, M., Smith, J. I., Gibson, W. P., Doggett, M. K., Taylor, G. H., Curtis, J., and Pasteris, P. P.: Physiographically Sensitive\n",
        "15 Mapping of Climatological Temperature and Precipitation across the Conterminous United States, International Journal of Climatology,\n",
        "28, 2031–2064, https://doi.org/10.1002/joc.1688, 2008.\n",
        "\n",
        "\n",
        "Dralle, D. N., W. J. Hahm, D. M. Rempe (2020). Dataset for \"Accounting for snow in the estimation of root-zone water storage capacity from precipitation and evapotranspiration fluxes\", HydroShare, http://www.hydroshare.org/resource/ee45c2f5f13042ca85bcb86bbfc9dd37\n",
        "\n",
        "\n",
        "Soil Survey Staff: Gridded National Soil Survey Geographic (gNATSGO) Database for the Conterminous United States, Tech. rep., United\n",
        "States Department of Agriculture, Natural Resources Conservation Service, available online at /https://nrcs.app.box.com/v/soils, 2019.\n",
        "\n",
        "Wang-Erlandsson, L., Bastiaanssen, W. G. M., Gao, H., Jägermeyr, J., Senay, G. B., van Dijk, A. I. J. M., Guerschman, J. P., Keys, P. W.,\n",
        "Gordon, L. J., and Savenije, H. H. G.: Global Root Zone Storage Capacity from Satellite-Based Evaporation, Hydrology and Earth System\n",
        "Sciences, 20, 1459–1481, https://doi.org/10.5194/hess-20-1459-2016, 2016.\n",
        "\n",
        "\n",
        "Yang, L., J. S. D. P. H. C. G. L. C. A. C. C. D. J. F. J. F. M. G. B. R. M. and Xian, G.: A New Generation of the United States National Land\n",
        "20 Cover Database: Requirements, Research Priorities, Design, and Implementation Strategies, pp. 108–123, https://developers.google.com/\n",
        "earth-engine/datasets/catalog/USGS_NLCD#citations, 2018.\n",
        "\n",
        "\n",
        "Zhang, Y., Kong, D., Gan, R., Chiew, F. H. S., McVicar, T. R., Zhang, Q., and Yang, Y.: Coupled Estimation of 500 m and 8-Day\n",
        "Resolution Global Evapotranspiration and Gross Primary Production in 2002–2017, Remote Sensing of Environment, 222, 165–182,\n",
        "305 https://doi.org/10.1016/j.rse.2018.12.031, 2019."
      ]
    }
  ]
}