02_era5_land_api_dask.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "93744300-0bc8-453e-b801-a4b1ced3b802",
   "metadata": {},
   "source": [
    "# Obtain NetCDFs file from ERA5-Land using the CDSAPI\n",
    "Information here: https://confluence.ecmwf.int/display/CKB/ERA5-Land%3A+data+documentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a5cc283-be10-4cf0-ba8d-d5f3148ddbde",
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import hvplot.xarray\n",
    "import pandas as pd\n",
    "import dask\n",
    "import fsspec\n",
    "import cdsapi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e1b8fdb-c697-4162-9442-e2b4f4cc5c4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "c = cdsapi.Client()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dff5180f-c026-4223-83f4-69698f7d4e15",
   "metadata": {},
   "source": [
    "#### Specify variables, spatial and temporal extents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9c375ea-97c2-4fd9-bfcb-7993ed84249e",
   "metadata": {},
   "outputs": [],
   "source": [
    "var_list = ['snow_depth_water_equivalent', 'soil_temperature_level_1']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68a014cc-566f-428a-8838-f8c74bd08e47",
   "metadata": {},
   "outputs": [],
   "source": [
    "# CONUS\n",
    "north = 49.3457868 \n",
    "west = -124.7844079 \n",
    "east = -66.9513812 \n",
    "south =  24.7433195 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c435078c-9666-4b12-94de-63049ea66de2",
   "metadata": {},
   "outputs": [],
   "source": [
    "fs = fsspec.filesystem('s3', anon=False,  skip_instance_cache=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8912fc6-8673-415c-9b70-05c640e79906",
   "metadata": {},
   "outputs": [],
   "source": [
    "dates = pd.date_range('1980-12-01','2022-01-31', freq='14D')\n",
    "print(dates)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42ee94a5-51cb-4c12-9148-2cbae46b2fd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "start_dates = [date.strftime('%Y-%m-%d') for date in dates[:-1]]\n",
    "stop_dates = [(date+pd.offsets.Day(13)).strftime('%Y-%m-%d') for date in dates[:-1]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08e24413-c808-4c8f-86f5-4fadcdd744b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_files = [f'esip-qhub/usgs/era5_land1/conus_{start_date}.nc' for start_date in start_dates]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "959e8c92-d28a-4f26-9aef-d92e434c4213",
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_files_processed = fs.glob('esip-qhub/usgs/era5_land1/conus_*.nc')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cef589c-d3c2-4daf-ae5f-a0b259e6b420",
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_files_to_create = sorted(list(set(s3_files) - set(s3_files_processed)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73cc0262-b71c-47e9-a5f7-e45019f78d2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(s3_files))\n",
    "print(len(s3_files_to_create))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6b27b20-bde0-4efc-aa5b-4c8eb59a1474",
   "metadata": {},
   "source": [
    "#### test generation of start_date and stop_date from s3file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "932e98db-c8c7-4739-9f7e-51dba0d2f7bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import datetime as dt\n",
    "s3file = s3_files_to_create[0]\n",
    "start_date = s3file.split('_')[-1].split('.')[0]\n",
    "s = start_date.split('-')\n",
    "stop_date = (dt.datetime(int(s[0]),int(s[1]),int(s[2])) + pd.offsets.Day(13)).strftime('%Y-%m-%d')\n",
    "print(start_date)\n",
    "print(stop_date)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0821072-0f7f-4337-a941-6cd44e9da83b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_chunk(s3file, keepbits):\n",
    "    # bitround python routines snipped out of xbitinfo\n",
    "    from numcodecs.bitround import BitRound\n",
    "    import pandas as pd\n",
    "\n",
    "    def _np_bitround(data, keepbits):\n",
    "        \"\"\"Bitround for Arrays.\"\"\"\n",
    "        codec = BitRound(keepbits=keepbits)\n",
    "        data = data.copy()  # otherwise overwrites the input\n",
    "        encoded = codec.encode(data)\n",
    "        return codec.decode(encoded)\n",
    "\n",
    "\n",
    "    def _keepbits_interface(da, keepbits):\n",
    "        \"\"\"Common interface to allowed keepbits types\n",
    "        Parameters\n",
    "        ----------\n",
    "        da : :py:class:`xarray.DataArray`\n",
    "          Input data to bitround\n",
    "        keepbits : int, dict of {str: int}, :py:class:`xarray.DataArray` or :py:class:`xarray.Dataset`\n",
    "          How many bits to keep as int\n",
    "        Returns\n",
    "        -------\n",
    "        keep : int\n",
    "          Number of keepbits for variable given in ``da``\n",
    "        \"\"\"\n",
    "        assert isinstance(da, xr.DataArray)\n",
    "        if isinstance(keepbits, int):\n",
    "            keep = keepbits\n",
    "        elif isinstance(keepbits, dict):\n",
    "            v = da.name\n",
    "            if v in keepbits.keys():\n",
    "                keep = keepbits[v]\n",
    "            else:\n",
    "                raise ValueError(f\"name {v} not for in keepbits: {keepbits.keys()}\")\n",
    "        elif isinstance(keepbits, xr.Dataset):\n",
    "            assert keepbits.coords[\"inflevel\"].shape <= (\n",
    "                1,\n",
    "            ), \"Information content is only allowed for one 'inflevel' here. Please make a selection.\"\n",
    "            if \"dim\" in keepbits.coords:\n",
    "                assert keepbits.coords[\"dim\"].shape <= (\n",
    "                    1,\n",
    "                ), \"Information content is only allowed along one dimension here. Please select one `dim`. To find the maximum keepbits, simply use `keepbits.max(dim='dim')`\"\n",
    "            v = da.name\n",
    "            if v in keepbits.keys():\n",
    "                keep = int(keepbits[v])\n",
    "            else:\n",
    "                raise ValueError(f\"name {v} not for in keepbits: {keepbits.keys()}\")\n",
    "        elif isinstance(keepbits, xr.DataArray):\n",
    "            assert keepbits.coords[\"inflevel\"].shape <= (\n",
    "                1,\n",
    "            ), \"Information content is only allowed for one 'inflevel' here. Please make a selection.\"\n",
    "            assert keepbits.coords[\"dim\"].shape <= (\n",
    "                1,\n",
    "            ), \"Information content is only allowed along one dimension here. Please select one `dim`. To find the maximum keepbits, simply use `keepbits.max(dim='dim')`\"\n",
    "            v = da.name\n",
    "            if v == keepbits.name:\n",
    "                keep = int(keepbits)\n",
    "            else:\n",
    "                raise KeyError(f\"no keepbits found for variable {v}\")\n",
    "        else:\n",
    "            raise TypeError(f\"type {type(keepbits)} is not a valid type for keepbits.\")\n",
    "        return keep\n",
    "    \n",
    "    def xr_bitround(da, keepbits):\n",
    "    \n",
    "        \"\"\"Apply bitrounding based on keepbits from :py:func:`xbitinfo.xbitinfo.get_keepbits` for :py:class:`xarray.Dataset` or :py:class:`xarray.DataArray` wrapping ``numcodecs.bitround``\n",
    "        Parameters\n",
    "        ----------\n",
    "        da : :py:class:`xarray.DataArray` or :py:class:`xarray.Dataset`\n",
    "          Input data to bitround\n",
    "        keepbits : int, dict of {str: int}, :py:class:`xarray.DataArray` or :py:class:`xarray.Dataset`\n",
    "          How many bits to keep as int. Fails if dict or :py:class:`xarray.Dataset` and key or variable not present.\n",
    "        Returns\n",
    "        -------\n",
    "        da_bitrounded : :py:class:`xarray.DataArray` or :py:class:`xarray.Dataset`\n",
    "        Example\n",
    "        -------\n",
    "        >>> ds = xr.tutorial.load_dataset(\"air_temperature\")\n",
    "        >>> info_per_bit = xb.get_bitinformation(ds, dim=\"lon\")\n",
    "        >>> keepbits = xb.get_keepbits(info_per_bit, 0.99)\n",
    "        >>> ds_bitrounded = xb.xr_bitround(ds, keepbits)\n",
    "        \"\"\"\n",
    "        if isinstance(da, xr.Dataset):\n",
    "            da_bitrounded = da.copy()\n",
    "            for v in da.data_vars:\n",
    "                da_bitrounded[v] = xr_bitround(da[v], keepbits)\n",
    "            return da_bitrounded\n",
    "\n",
    "        assert isinstance(da, xr.DataArray)\n",
    "        keep = _keepbits_interface(da, keepbits)\n",
    "\n",
    "        da = xr.apply_ufunc(_np_bitround, da, keep, dask=\"parallelized\", keep_attrs=True)\n",
    "        da.attrs[\"_QuantizeBitRoundNumberOfSignificantDigits\"] = keep\n",
    "        return da\n",
    "\n",
    "    import datetime as dt   \n",
    "    start_date = s3file.split('_')[-1].split('.')[0]\n",
    "    s = start_date.split('-')\n",
    "    stop_date = (dt.datetime(int(s[0]),int(s[1]),int(s[2])) + pd.offsets.Day(13)).strftime('%Y-%m-%d')\n",
    "    local_ncfile = f'era5land_{start_date}.nc'\n",
    "    local_nc4file = f'era5land_{start_date}-nc4.nc'\n",
    "    c.retrieve(\n",
    "        'reanalysis-era5-land',\n",
    "        {\n",
    "            'variable': var_list, \n",
    "            'area'    : f'{north}/{west}/{south}/{east}', \n",
    "            'date'    : f'{start_date}/{stop_date}',\n",
    "            'time': ['00:00', '01:00', '02:00', '03:00', '04:00', '05:00',\n",
    "                     '06:00', '07:00', '08:00', '09:00', '10:00', '11:00',\n",
    "                     '12:00', '13:00', '14:00', '15:00', '16:00', '17:00',\n",
    "                     '18:00', '19:00', '20:00', '21:00', '22:00', '23:00'],\n",
    "            'format':'netcdf'\n",
    "        },\n",
    "        local_ncfile)\n",
    "        \n",
    "    ds = xr.open_dataset(local_ncfile)\n",
    "    ds_bitrounded = xr_bitround(ds, keepbits)\n",
    "    encoding = {}\n",
    "    for data_var in ds.data_vars:\n",
    "        encoding[data_var]=dict(dtype='float32', zlib=True)\n",
    "\n",
    "    encoding['latitude'] = {'_FillValue':None}\n",
    "    encoding['longitude'] = {'_FillValue':None}\n",
    "\n",
    "    ds_bitrounded.to_netcdf(local_nc4file, engine='netcdf4', encoding=encoding, mode='w')  \n",
    "    fs.upload(local_nc4file, s3file)\n",
    "    fs2 = fsspec.filesystem('file')\n",
    "    fs2.rm([local_ncfile, local_nc4file])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1436f294-161f-43b6-afa8-84bb7db25810",
   "metadata": {},
   "source": [
    "### API requests in Parallel using Dask Gateway Cluster"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a38d0e7-c017-4c2c-bfb4-88a999781a9c",
   "metadata": {},
   "source": [
    "#### Spin up Dask Cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ee23cd8-6f94-4551-9e2a-e0487a13b16c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def configure_cluster(resource):\n",
    "    ''' Helper function to configure cluster\n",
    "    '''\n",
    "    if resource == 'denali':\n",
    "        cluster = LocalCluster(threads_per_worker=1)\n",
    "        client = Client(cluster)\n",
    "    \n",
    "    elif resource == 'tallgrass':\n",
    "        from dask_jobqueue import SLURMCluster\n",
    "        cluster = SLURMCluster(queue='cpu', cores=1, interface='ib0',\n",
    "                               job_extra=['--nodes=1', '--ntasks-per-node=1', '--cpus-per-task=1'],\n",
    "                               memory='6GB')\n",
    "        cluster.adapt(maximum_jobs=30)\n",
    "        client = Client(cluster)\n",
    "        \n",
    "    elif resource == 'local':\n",
    "        import os\n",
    "        import warnings\n",
    "        warnings.warn(\"Running locally can result in costly data transfers!\\n\")\n",
    "        n_cores = os.cpu_count() # set to match your machine\n",
    "        cluster = LocalCluster(threads_per_worker=n_cores)\n",
    "        client = Client(cluster)\n",
    "        \n",
    "    elif resource in ['esip-qhub-gateway-v0.4']:   \n",
    "        import sys, os\n",
    "        sys.path.append(os.path.join(os.environ['HOME'],'shared','users','lib'))\n",
    "        import ebdpy as ebd\n",
    "        aws_profile = 'esip-qhub'\n",
    "        aws_region = 'us-west-2'\n",
    "        endpoint = f's3.{aws_region}.amazonaws.com'\n",
    "        ebd.set_credentials(profile=aws_profile, region=aws_region, endpoint=endpoint)\n",
    "        worker_max = 5\n",
    "        client,cluster = ebd.start_dask_cluster(profile=aws_profile, worker_max=worker_max, \n",
    "                                              region=aws_region, use_existing_cluster=True,\n",
    "                                              adaptive_scaling=False, wait_for_cluster=False, \n",
    "                                              worker_profile='Small Worker', propagate_env=True)\n",
    "        \n",
    "    return client, cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "480cec43-bb52-4c7b-ab49-b1ce84cd5d99",
   "metadata": {},
   "outputs": [],
   "source": [
    "resource = 'esip-qhub-gateway-v0.4' #denali, tallgrass, local, esip-qhub-gateway-v0.4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3bb3ba9-11c6-44a3-b87f-813fbd42932c",
   "metadata": {},
   "outputs": [],
   "source": [
    "client, cluster = configure_cluster(resource)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08b688c8-9f77-4bc7-8f8e-11061173fa43",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_workers = 10\n",
    "cluster.scale(n_workers)\n",
    "client.wait_for_workers(n_workers=n_workers)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cff81d9-1c8a-43c3-8fc1-12a009069062",
   "metadata": {},
   "source": [
    "#### Run in Parallel using Dask Delayed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0201d76-2a91-4889-9d96-23172b695899",
   "metadata": {},
   "outputs": [],
   "source": [
    "keepbits = xr.open_dataset('keepbits.nc').load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ef3f61f-2b35-4f66-ba7b-d0d01da440f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "_ = dask.compute(*[dask.delayed(get_chunk)(s3file,keepbits) for s3file in s3_files_to_create[20:40]], retries=10);   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "201bdbdc-b48c-4f79-8455-b204b4e35ed9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Estimated time in hours to complete\n",
    "1070/(20/(6*60+50))/3600"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "users-pangeo",
   "language": "python",
   "name": "conda-env-users-pangeo-py"
  },
  "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}