ngen_forcing/VERYROUGH_RTI_Forcing_example.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4b19730-d455-418f-a3a3-15efec3a8387",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install config\n",
    "!pip install geoparquet\n",
    "import geoparquet as gpq\n",
    "import pickle\n",
    "\n",
    "# import config\n",
    "import os\n",
    "import gc\n",
    "\n",
    "import geopandas as gpd\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "import rasterio as rio\n",
    "\n",
    "from rasterio.io import MemoryFile\n",
    "\n",
    "from rasterio.features import rasterize\n",
    "\n",
    "# from utils import parquet_to_gdf\n",
    "# from Grid_to_parquet import get_dataset\n",
    "\n",
    "TEMPLATE_BLOB_NAME = (\n",
    "    \"nwm.20221001/forcing_medium_range/nwm.t00z.medium_range.forcing.f001.conus.nc\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9159453-76c6-4bbe-b08f-a746a54f48f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "NWM_BUCKET = \"national-water-model\"\n",
    "\n",
    "# WKT strings extracted from NWM grids\n",
    "CONUS_NWM_WKT = 'PROJCS[\"Lambert_Conformal_Conic\",GEOGCS[\"GCS_Sphere\",DATUM[\"D_Sphere\",SPHEROID[\"Sphere\",6370000.0,0.0]], \\\n",
    "PRIMEM[\"Greenwich\",0.0],UNIT[\"Degree\",0.0174532925199433]],PROJECTION[\"Lambert_Conformal_Conic_2SP\"],PARAMETER[\"false_easting\",0.0],\\\n",
    "PARAMETER[\"false_northing\",0.0],PARAMETER[\"central_meridian\",-97.0],PARAMETER[\"standard_parallel_1\",30.0],\\\n",
    "PARAMETER[\"standard_parallel_2\",60.0],PARAMETER[\"latitude_of_origin\",40.0],UNIT[\"Meter\",1.0]]'\n",
    "\n",
    "HI_NWM_WKT = 'PROJCS[\"Lambert_Conformal_Conic\",GEOGCS[\"GCS_Sphere\",DATUM[\"D_Sphere\",SPHEROID[\"Sphere\",6370000.0,0.0]],\\\n",
    "PRIMEM[\"Greenwich\",0.0],UNIT[\"Degree\",0.0174532925199433]],PROJECTION[\"Lambert_Conformal_Conic_2SP\"],PARAMETER[\"false_easting\",0.0],\\\n",
    "PARAMETER[\"false_northing\",0.0],PARAMETER[\"central_meridian\",-157.42],PARAMETER[\"standard_parallel_1\",10.0],\\\n",
    "PARAMETER[\"standard_parallel_2\",30.0],PARAMETER[\"latitude_of_origin\",20.6],UNIT[\"Meter\",1.0]]'\n",
    "\n",
    "PR_NWM_WKT = 'PROJCS[\"Sphere_Lambert_Conformal_Conic\",GEOGCS[\"GCS_Sphere\",DATUM[\"D_Sphere\",SPHEROID[\"Sphere\",6370000.0,0.0]],\\\n",
    "PRIMEM[\"Greenwich\",0.0],UNIT[\"Degree\",0.0174532925199433]],PROJECTION[\"Lambert_Conformal_Conic_2SP\"],PARAMETER[\"false_easting\",0.0],\\\n",
    "PARAMETER[\"false_northing\",0.0],PARAMETER[\"central_meridian\",-65.91],PARAMETER[\"standard_parallel_1\",18.1],\\\n",
    "PARAMETER[\"standard_parallel_2\",18.1],PARAMETER[\"latitude_of_origin\",18.1],UNIT[\"Meter\",1.0]]'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eaba8b05-2b98-4aea-974f-d26de67f729b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from pathlib import Path\n",
    "\n",
    "\n",
    "CACHE_DIR = Path(\".\", \"data\", \"cache\")\n",
    "NWM_CACHE_DIR = os.path.join(CACHE_DIR, \"nwm\")\n",
    "USGS_CACHE_DIR = os.path.join(CACHE_DIR, \"usgs\")\n",
    "GEO_CACHE_DIR = os.path.join(CACHE_DIR, \"geo\")\n",
    "\n",
    "NWM_CACHE_H5 = os.path.join(NWM_CACHE_DIR, \"gcp_client.h5\")\n",
    "\n",
    "PARQUET_CACHE_DIR = os.path.join(CACHE_DIR, \"parquet\")\n",
    "MEDIUM_RANGE_FORCING_PARQUET = os.path.join(PARQUET_CACHE_DIR, \"forcing_medium_range\")\n",
    "FORCING_ANALYSIS_ASSIM_PARQUET = os.path.join(\n",
    "    PARQUET_CACHE_DIR, \"forcing_analysis_assim\"\n",
    ")\n",
    "MEDIUM_RANGE_PARQUET = os.path.join(PARQUET_CACHE_DIR, \"medium_range\")\n",
    "USGS_PARQUET = os.path.join(PARQUET_CACHE_DIR, \"usgs\")\n",
    "\n",
    "HUC10_SHP_FILEPATH = os.path.join(GEO_CACHE_DIR, \"wbdhu10_conus.shp\")\n",
    "HUC10_PARQUET_FILEPATH = os.path.join(GEO_CACHE_DIR, \"wbdhu10_conus.parquet\")\n",
    "HUC10_MEDIUM_RANGE_WEIGHTS_FILEPATH = os.path.join(\n",
    "    GEO_CACHE_DIR, \"wbdhu10_medium_range_weights.pkl\"\n",
    ")\n",
    "\n",
    "ROUTE_LINK_FILE = os.path.join(NWM_CACHE_DIR, \"RouteLink_CONUS.nc\")\n",
    "ROUTE_LINK_PARQUET = os.path.join(NWM_CACHE_DIR, \"route_link_conus.parquet\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82731a4e-70c8-4cc8-a6e1-2de5ff14a947",
   "metadata": {},
   "outputs": [],
   "source": [
    "def parquet_to_gdf(parquet_filepath: str) -> gpd.GeoDataFrame:\n",
    "    gdf = gpd.read_parquet(parquet_filepath)\n",
    "    return gdf\n",
    "\n",
    "\n",
    "def get_cache_dir(create: bool = True):\n",
    "    if not os.path.exists(NWM_CACHE_DIR) and create:\n",
    "        os.mkdir(NWM_CACHE_DIR)\n",
    "    if not os.path.exists(NWM_CACHE_DIR):\n",
    "        raise NotADirectoryError\n",
    "    return NWM_CACHE_DIR\n",
    "\n",
    "\n",
    "def make_parent_dir(filepath):\n",
    "    Path(filepath).parent.mkdir(parents=True, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74895a20-a42e-48ff-bff8-313a0ae26d14",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_dataset(blob_name: str, use_cache: bool = True) -> xr.Dataset:\n",
    "    \"\"\"Retrieve a blob from the data service as xarray.Dataset.\n",
    "    Based largely on OWP HydroTools.\n",
    "    Parameters\n",
    "    ----------\n",
    "    blob_name: str, required\n",
    "        Name of blob to retrieve.\n",
    "    use_cache: bool, default True\n",
    "        If cache should be used.\n",
    "        If True, checks to see if file is in cache, and\n",
    "        If fetched from remote, will save to cache.\n",
    "    Returns\n",
    "    -------\n",
    "    ds : xarray.Dataset\n",
    "        The data stored in the blob.\n",
    "    \"\"\"\n",
    "    # TODO: Check to see if this does any better than kerchunk\n",
    "    # the caching should help, but probably needs to be managed to function asynchronously.\n",
    "    # Perhaps if the files is not cached, we can create the dataset from\n",
    "    # kerchunk with a remote path and then asynchronously do a download to cache it\n",
    "    # for next time. The hypothesis would be that the download speed will not be any slower than\n",
    "    # just accessing the file remotely.\n",
    "    nc_filepath = os.path.join(get_cache_dir(), blob_name)\n",
    "    make_parent_dir(nc_filepath)\n",
    "\n",
    "    # If the file exists and use_cache = True\n",
    "    if os.path.exists(nc_filepath) and use_cache:\n",
    "        # Get dataset from cache\n",
    "        ds = xr.load_dataset(\n",
    "            nc_filepath,\n",
    "            engine=\"h5netcdf\",\n",
    "        )\n",
    "        return ds\n",
    "    else:\n",
    "        # Get raw bytes\n",
    "        raw_bytes = get_blob(blob_name)\n",
    "        # Create Dataset\n",
    "        ds = xr.load_dataset(\n",
    "            MemoryFile(raw_bytes),\n",
    "            engine=\"h5netcdf\",\n",
    "        )\n",
    "        if use_cache:\n",
    "            # Subset and cache\n",
    "            ds[\"RAINRATE\"].to_netcdf(\n",
    "                nc_filepath,\n",
    "                engine=\"h5netcdf\",\n",
    "            )\n",
    "        return ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b199b72a-0220-4583-ac3a-9f27c1586fe1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_weights_file(\n",
    "    gdf: gpd.GeoDataFrame,\n",
    "    src: xr.DataArray,\n",
    "    weights_filepath: str,\n",
    "    crosswalk_dict_key: str,\n",
    "):\n",
    "    \"\"\"Generate a weights file.\"\"\"\n",
    "\n",
    "    gdf_proj = gdf.to_crs(CONUS_NWM_WKT)\n",
    "\n",
    "    crosswalk_dict = {}\n",
    "\n",
    "    # This is a probably a really poor performing way to do this\n",
    "    # TODO: Consider vectorizing -- would require digging into the\n",
    "    # other end of these where we unpack the weights...\n",
    "    for index, row in gdf_proj.iterrows():\n",
    "        geom_rasterize = rasterize(\n",
    "            [(row[\"geometry\"], 1)],\n",
    "            out_shape=src.rio.shape,\n",
    "            transform=src.rio.transform(),\n",
    "            all_touched=True,\n",
    "            fill=0,  # IS FILL 0\n",
    "            dtype=\"uint8\",\n",
    "        )\n",
    "        if crosswalk_dict_key:\n",
    "            crosswalk_dict[row[crosswalk_dict_key]] = np.where(geom_rasterize == 1)\n",
    "        else:\n",
    "            crosswalk_dict[index] = np.where(geom_rasterize == 1)\n",
    "\n",
    "    with open(weights_filepath, \"wb\") as f:\n",
    "        # TODO: This is a dict of ndarrays, which could be easily stored as a set of parquet files for safekeeping.\n",
    "        pickle.dump(crosswalk_dict, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dec7dd6a-483b-4401-9690-4fb213d2fc85",
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_zonalstats_to_gdf_weights(\n",
    "    gdf: gpd.GeoDataFrame,\n",
    "    src: xr.DataArray,\n",
    "    weights_filepath: str,\n",
    ") -> gpd.GeoDataFrame:\n",
    "    \"\"\"Calculates zonal stats and adds to GeoDataFrame\"\"\"\n",
    "\n",
    "    df = calc_zonal_stats_weights(src, weights_filepath)\n",
    "    gdf_map = gdf.merge(df, left_on=\"huc10\", right_on=\"catchment_id\")\n",
    "\n",
    "    return gdf_map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cad942dd-29ec-4ed8-be6d-e5cc13dde0ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "from google.cloud import storage\n",
    "\n",
    "\n",
    "def get_blob(blob_name: str, bucket: str = NWM_BUCKET) -> bytes:\n",
    "    \"\"\"Retrieve a blob from the data service as bytes.\n",
    "    Based largely on OWP HydroTools.\n",
    "    Parameters\n",
    "    ----------\n",
    "    blob_name : str, required\n",
    "        Name of blob to retrieve.\n",
    "    Returns\n",
    "    -------\n",
    "    data : bytes\n",
    "        The data stored in the blob.\n",
    "    \"\"\"\n",
    "    # Setup anonymous client and retrieve blob data\n",
    "    client = storage.Client.create_anonymous_client()\n",
    "    bucket = client.bucket(bucket)\n",
    "    return bucket.blob(blob_name).download_as_bytes(timeout=120)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb25a27c-fd98-478c-97ef-c53dbd73a6f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "### MAIN\n",
    "ds = get_dataset(TEMPLATE_BLOB_NAME, use_cache=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8e7f87d-ae19-454e-b958-2668e15d5c10",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "polygonfile = gpd.read_file(\"16/nextgen_16.gpkg\", layer=\"divides\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95f6ad0a-ec08-42aa-a9a9-4694dd7ad1e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "polygonfile.to_parquet(\"16/ng_16.parquet\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d540ca78-43d5-4387-bd77-c6d8848ca576",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "gdf = parquet_to_gdf(\"16/ng_16.parquet\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb1dc54e-89ff-4d91-b73e-350da31e3ebb",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "src = ds[\"RAINRATE\"]\n",
    "generate_weights_file(gdf, src, \"data/weights.pkl\", crosswalk_dict_key=\"id\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ffb4bf4-8e97-47f6-9bf4-d42be5c6b4f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"data/weights.pkl\", \"rb\") as f:\n",
    "    a = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "135a4496-e9f0-4640-b321-b3f06ac8bf88",
   "metadata": {},
   "outputs": [],
   "source": [
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c64c68d-bc00-4a20-b28b-1d9d3d9b3099",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "with open(\"data/weights.json\", \"w\") as f:\n",
    "    json.dump(a, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2101a68-2365-41db-a226-e70f661ed57a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install ipdb\n",
    "def calc_zonal_stats_weights(\n",
    "    src: xr.DataArray,\n",
    "    weights_filepath: str,\n",
    ") -> pd.DataFrame:\n",
    "    \"\"\"Calculates zonal stats\"\"\"\n",
    "\n",
    "    # Open weights dict from pickle\n",
    "    # This could probably be done once and passed as a reference.\n",
    "    with open(weights_filepath, \"rb\") as f:\n",
    "        crosswalk_dict = pickle.load(f)\n",
    "\n",
    "    r_array = src.values[0]\n",
    "    r_array[r_array == src.rio.nodata] = np.nan\n",
    "\n",
    "    mean_dict = {}\n",
    "    for key, value in crosswalk_dict.items():\n",
    "        mean_dict[key] = np.nanmean(r_array[value])\n",
    "\n",
    "    df = pd.DataFrame.from_dict(mean_dict, orient=\"index\", columns=[\"value\"])\n",
    "\n",
    "    df.reset_index(inplace=True, names=\"catchment_id\")\n",
    "\n",
    "    # This should not be needed, but without memory usage grows\n",
    "    del crosswalk_dict\n",
    "    del f\n",
    "    gc.collect()\n",
    "\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc543b89-9335-4a30-bdae-83c30293d6dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "calc_zonal_stats_weights(src, \"data/weights.pkl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa8c71d3-2fd1-46eb-8fee-d3f0ddb00929",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_forcing_dict_RTIway(\n",
    "    pickle_file,  # This would be a Feature list for parallel calling --\n",
    "    # if there is a stored weights file, we use it\n",
    "    # (checking for an optional flag to force re-creation of the weights...)\n",
    "    folder_prefix,\n",
    "    file_list,\n",
    "):\n",
    "\n",
    "    var = \"RAINRATE\"\n",
    "    reng = \"rasterio\"\n",
    "    filehandles = [\n",
    "        xr.open_dataset(folder_prefix / f, engine=reng)[var] for f in file_list\n",
    "    ]\n",
    "    # filehandles = [get_dataset(\"data/\" + f, use_cache=True) for f in file_list]\n",
    "    stats = []\n",
    "\n",
    "    for _i, f in enumerate(filehandles):\n",
    "        print(f\"{_i}, {round(_i/len(file_list), 2)*100}\".ljust(40), end=\"\\r\")\n",
    "        stats.append(calc_zonal_stats_weights(f, pickle_file))\n",
    "\n",
    "    [f.close() for f in filehandles]\n",
    "    return stats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "129e1681-d89c-4636-a487-0432fe7e283c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_forcing_dict_RTIway2(\n",
    "    pickle_file,  # This would be a Feature list for parallel calling --\n",
    "    # if there is a stored weights file, we use it\n",
    "    # (checking for an optional flag to force re-creation of the weights...)\n",
    "    gpkg_divides,\n",
    "    folder_prefix,\n",
    "    filelist,\n",
    "    var_list,\n",
    "):\n",
    "    reng = \"rasterio\"\n",
    "    pick_val = \"value\"\n",
    "\n",
    "    df_dict = {}\n",
    "    dl_dict = {}\n",
    "    for _v in var_list:\n",
    "        df_dict[_v] = pd.DataFrame(index=gpkg_divides.index)\n",
    "        dl_dict[_v] = []\n",
    "\n",
    "    # ds_list = []\n",
    "    for _i, _nc_file in enumerate(filelist):\n",
    "        # _nc_file = (\"nwm.t00z.medium_range.forcing.f001.conus.nc\")\n",
    "        _full_nc_file = folder_prefix.joinpath(_nc_file)\n",
    "\n",
    "        try:\n",
    "            # with xr.open_dataset(_full_nc_file, engine=reng) as _xds:\n",
    "            with xr.open_dataset(_full_nc_file) as _xds:\n",
    "                # _xds = ds_list[_i]\n",
    "                # _xds.rio.write_crs(rasterio.crs.CRS.from_wkt(CONUS_NWM_WKT), inplace=True)\n",
    "                print(f\"{_i}, {round(_i/len(filelist), 5)*100}\".ljust(40), end=\"\\r\")\n",
    "                for _v in var_list:\n",
    "                    _src = _xds[_v]\n",
    "                    _df_zonal_stats = calc_zonal_stats_weights(_src, pickle_file)\n",
    "                    # if adding statistics back to original GeoDataFrame\n",
    "                    # gdf3 = pd.concat([gpkg_divides, _df_zonal_stats], axis=1)\n",
    "                    _df = pd.DataFrame(index=gpkg_divides.index)\n",
    "                    _df[_xds.time.values[0]] = _df_zonal_stats[pick_val]\n",
    "                    # TODO: This same line could add the new values directly\n",
    "                    # to the same dictionary. But after adding about 100 of them,\n",
    "                    # pandas starts to complain about degraded performance due to\n",
    "                    # fragmentation of the dataframe. We tried it this was as a\n",
    "                    # workaround, with the loop below to accomplish the concatenation.\n",
    "                    dl_dict[_v].append(_df)\n",
    "        except:\n",
    "            print(f\"No such file: {_full_nc_file}\")\n",
    "\n",
    "    for _v in var_list:\n",
    "        df_dict[_v] = pd.concat(dl_dict[_v], axis=1)\n",
    "\n",
    "    # [_xds.close() for _xds in ds_list]\n",
    "\n",
    "    return df_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59249a13-8986-4312-aaf8-a0f60202f6f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "a = pd.DataFrame(index=gpkg_divides.index)\n",
    "a[fd2[\"U2D\"][0][0]] = fd2[\"U2D\"][0][1]\n",
    "b = pd.DataFrame(index=gpkg_divides.index)\n",
    "b[fd2[\"U2D\"][1][0]] = fd2[\"U2D\"][1][1]\n",
    "pd.concat([a, b], axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ddeca13-2dac-4f7a-aa1c-1c38ab496333",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "reng = \"rasterio\"\n",
    "folder_prefix = Path(\"./data/d1980\")\n",
    "_nc_file = \"198001010100.LDASIN_DOMAIN1\"\n",
    "_full_nc_file = folder_prefix.joinpath(_nc_file)\n",
    "_v = \"U2D\"\n",
    "_xds = xr.open_dataset(_full_nc_file)\n",
    "# _src = _xds[_v]\n",
    "# _df_zonal_stats = (calc_zonal_stats_weights(_src, pickle_file))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddb52923-e1ac-4c05-b44f-d8f1c288720e",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "_xds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d36db950-aae4-4c5f-9f27-44baabb3a19e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import rasterio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d626a03-27ff-4ad1-8664-0efe224ad547",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "_xds.rio.write_crs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc1afb13-df7a-4a00-b00b-8cc0b888b0c3",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "_xds.rio.write_crs(rasterio.crs.CRS.from_wkt(CONUS_NWM_WKT), inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82ece2a3-2a6e-4621-be08-317e21769507",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "_xds.rio.crs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c391ab6-3362-40e2-af06-ac26c3e8fc77",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "list_of_files = [\n",
    "    f\"1980{_m:02}{_d:02}{_h:02}00.LDASIN_DOMAIN1\"\n",
    "    for _m in range(1, 13)\n",
    "    for _d in range(1, 32)\n",
    "    for _h in range(0, 24)\n",
    "]\n",
    "# list_of_files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5633984-4043-4c6c-9ca0-45c78c766063",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "import time\n",
    "\n",
    "folder_prefix = Path(\"data/d1980\")\n",
    "list_of_files = [\n",
    "    f\"nwm.t12z.medium_range.forcing.f{_r:03}.conus.nc\" for _r in range(1, 241)\n",
    "]\n",
    "list_of_files = [\n",
    "    f\"1980{_m:02}{_d:02}{_h:02}00.LDASIN_DOMAIN1\"\n",
    "    for _m in range(1, 13)\n",
    "    for _d in range(1, 32)\n",
    "    for _h in range(0, 24)\n",
    "]\n",
    "\n",
    "f_16 = \"16/nextgen_16.gpkg\"\n",
    "gpkg_divides = gpd.read_file(f_16, layer=\"divides\")\n",
    "var_list = [\n",
    "    \"U2D\",\n",
    "    \"V2D\",\n",
    "    \"LWDOWN\",\n",
    "    \"RAINRATE\",\n",
    "    \"T2D\",\n",
    "    \"Q2D\",\n",
    "    \"PSFC\",\n",
    "    \"SWDOWN\",\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d4fa89d-f86c-4978-a536-b95f480c1906",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# file_list = list_of_files[0:18]\n",
    "file_list = list_of_files[0:3]\n",
    "file_list = list_of_files[0:]\n",
    "\n",
    "pickle_file = \"data/weights.pkl\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ba70fab-e2ab-4bf7-a1a1-febb19ce18c4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "var_list\n",
    "start_time = time.time()\n",
    "print(f\"Working on the new way\")\n",
    "fd2 = get_forcing_dict_RTIway2(\n",
    "    pickle_file,\n",
    "    gpkg_divides,\n",
    "    folder_prefix,\n",
    "    file_list,\n",
    "    var_list,\n",
    ")\n",
    "print(time.time() - start_time)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d37d2590-1de1-4e9c-85f2-9c7a1172f07c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# %%time\n",
    "# if 1 == 1:\n",
    "#     _f = folder_prefix / list_of_files[0]\n",
    "#     src = xr.open_dataset(_f)[\"RAINRATE\"]\n",
    "#     generate_weights_file(gdf, src, \"data_3/weights_alt.pkl\", crosswalk_dict_key=\"id\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea5887b5-b637-4642-bd47-ade6118e6532",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "len(fd2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61e4d101-7973-4e64-8619-4ca6ab9a0a18",
   "metadata": {},
   "outputs": [],
   "source": [
    "fd2[\"U2D\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b47e38e-26d1-44bb-be4c-68636035467c",
   "metadata": {},
   "outputs": [],
   "source": [
    "pcp_var.transpose()[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3885ed0-8b83-46af-afc9-75639ff8b7c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "# pcp_var = pd.read_csv(\"RAINRATE.csv\", index_col=0)#.rename(\"APCP_surface\")\n",
    "# lw_var = pd.read_csv(\"LWDOWN.csv\", index_col=0)#.rename(\"DLWRF_surface\")\n",
    "# sw_var = pd.read_csv(\"SWDOWN.csv\", index_col=0)#.rename(\"DSWRF_surface\")\n",
    "# sp_var = pd.read_csv(\"PSFC.csv\", index_col=0)#.rename(\"SPFH_2maboveground\")\n",
    "# tmp_var = pd.read_csv(\"T2D.csv\", index_col=0)#.rename(\"TMP_2maboveground\")\n",
    "# u2d_var = pd.read_csv(\"U2D.csv\", index_col=0)#.rename(\"UGRD_10maboveground\")\n",
    "# v2d_var = pd.read_csv(\"V2D.csv\", index_col=0)#.rename(\"VGRD_10maboveground\")\n",
    "# pcp_var2 = pd.read_csv(\"RAINRATE.csv\", index_col=0)#.rename(\"precip_rate\") ##BROKEN!!\n",
    "pcp_var = fd2[\"RAINRATE\"]\n",
    "lw_var = fd2[\"LWDOWN\"]\n",
    "sw_var = fd2[\"SWDOWN\"]\n",
    "sp_var = fd2[\"PSFC\"]\n",
    "tmp_var = fd2[\"T2D\"]\n",
    "u2d_var = fd2[\"U2D\"]\n",
    "v2d_var = fd2[\"V2D\"]\n",
    "pcp_var2 = fd2[\"RAINRATE\"]\n",
    "\n",
    "for _i in range(0, 40000):\n",
    "    # _i = 0\n",
    "    try:\n",
    "        pcp_var_0 = pcp_var.transpose()[_i].rename(\"APCP_surface\")\n",
    "        lw_var_0 = lw_var.transpose()[_i].rename(\"DLWRF_surface\")\n",
    "        sw_var_0 = sw_var.transpose()[_i].rename(\"DSWRF_surface\")\n",
    "        sp_var_0 = sp_var.transpose()[_i].rename(\"SPFH_2maboveground\")\n",
    "        tmp_var_0 = tmp_var.transpose()[_i].rename(\"TMP_2maboveground\")\n",
    "        u2d_var_0 = u2d_var.transpose()[_i].rename(\"UGRD_10maboveground\")\n",
    "        v2d_var_0 = v2d_var.transpose()[_i].rename(\"VGRD_10maboveground\")\n",
    "        pcp_var2_0 = pcp_var2.transpose()[_i].rename(\"precip_rate\")  ##BROKEN!!\n",
    "\n",
    "        d = pd.concat(\n",
    "            [\n",
    "                pcp_var_0,\n",
    "                lw_var_0,\n",
    "                sw_var_0,\n",
    "                sp_var_0,\n",
    "                tmp_var_0,\n",
    "                u2d_var_0,\n",
    "                v2d_var_0,\n",
    "                pcp_var2_0,\n",
    "            ],\n",
    "            axis=1,\n",
    "        )\n",
    "        d.index.name = \"time\"\n",
    "\n",
    "        d.to_csv(f\"input_data/cat16_{_i:07}.csv\")\n",
    "    except:\n",
    "        print(f\"no data for watershed {_i}\", end=\"\\t\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "358d9a99-2016-4c51-b65c-e9cd92fce358",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Make a shell script string to rename the csvs...\n",
    "gpkg_divides[\"id\"]\n",
    "for _i, cat_id in enumerate(gpkg_divides[\"id\"]):\n",
    "    print(f\"mv cat16_{_i:07}.csv cat16_{cat_id}.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52383127-cc94-49cc-a0ba-3d3e2057fc08",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}