Extract_River_Segments_from_HydroRIVERS.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "31fd200c",
   "metadata": {},
   "source": [
    "# Extract river segements from HydroRIVERS\n",
    "\n",
    "The following code takes a continental-scale rivers shapefile from the HydroATLAS \"HydroRIVERS\" database, which is a large polyline file containing segments of all large rivers in the world. The user should first open the shpaefile in a GIS and find the desired river to extract. The user should then choose the most upstream segment they want for the extraction, and record the 'HYRIV_ID' for that segment. 'HYRIV_ID' is a unique numeric for each segment in the dataset. The user then enters the desired dataset name, path, the 'HYRIV_ID', a river name, and an output path (optional). Executing the code extracts all river segments downstream of and including the given 'HYRIV_ID' segment, terminating at the ocean or depositional basin. This is written as a shapefile to the desired output folder and named after the river.\n",
    "\n",
    "Author: James (Huck) Rees; PhD Student, UCSB Geography\n",
    "\n",
    "Date: June 25th, 2024"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0346ff4d",
   "metadata": {},
   "source": [
    "## Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b498dad0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import requests\n",
    "import zipfile\n",
    "from urllib.parse import urlparse\n",
    "import geopandas as gpd\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0b842ae",
   "metadata": {},
   "source": [
    "## Initialize functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e365e85e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def download_and_extract_shapefile(url, extract_to='shapefile_data'):\n",
    "    # Step 1: Download the zip file\n",
    "    local_zip = 'shapefile.zip'\n",
    "    with requests.get(url, stream=True) as r:\n",
    "        r.raise_for_status()\n",
    "        with open(local_zip, 'wb') as f:\n",
    "            for chunk in r.iter_content(chunk_size=8192):\n",
    "                f.write(chunk)\n",
    "\n",
    "    # Step 2: Extract the zip file\n",
    "    with zipfile.ZipFile(local_zip, 'r') as zip_ref:\n",
    "        zip_ref.extractall(extract_to)\n",
    "    \n",
    "    # Step 3: Extract the subfolder name from the URL\n",
    "    parsed_url = urlparse(url)\n",
    "    subfolder_name = os.path.basename(parsed_url.path).replace('.zip', '')\n",
    "\n",
    "    # Construct the full path to the subfolder\n",
    "    subfolder_path = os.path.join(extract_to, subfolder_name)\n",
    "\n",
    "    # Ensure the subfolder exists\n",
    "    if not os.path.exists(subfolder_path):\n",
    "        raise FileNotFoundError(f\"Expected subfolder {subfolder_name} not found in {extract_to}.\")\n",
    "\n",
    "    # Step 4: Find the shapefile (.shp) within the subfolder\n",
    "    shapefile_path = None\n",
    "    for root, dirs, files in os.walk(subfolder_path):\n",
    "        for file in files:\n",
    "            if file.endswith('.shp'):\n",
    "                shapefile_path = os.path.join(root, file)\n",
    "                break\n",
    "        if shapefile_path:\n",
    "            break\n",
    "\n",
    "    # Step 5: Load the shapefile into a GeoDataFrame\n",
    "    if shapefile_path:\n",
    "        gdf = gpd.read_file(shapefile_path)\n",
    "        return gdf\n",
    "    else:\n",
    "        raise FileNotFoundError(f\"No .shp file found in the subfolder {subfolder_name}.\")\n",
    "        \n",
    "def extract_downstream_river(gdf, start_hyriv_id, method=\"specific_segment\", terminal_hyriv_id=0, threshold_length_km=None):\n",
    "    \"\"\"\n",
    "    Extracts downstream segments of a river starting from a given HYRIV_ID.\n",
    "    \n",
    "    Parameters:\n",
    "        gdf (GeoDataFrame): The GeoDataFrame containing river segments.\n",
    "        start_hyriv_id (int): The HYRIV_ID to start the extraction from.\n",
    "        method (str): The method to determine when to stop extraction. \n",
    "                      Options are \"specific_segment\" (default) or \"length\".\n",
    "        terminal_hyriv_id (int): The HYRIV_ID to stop at when using \"specific_segment\" method. Default is 0 (ocean).\n",
    "        threshold_length_km (float): The cumulative distance in km to stop at when using \"length\" method.\n",
    "    \n",
    "    Returns:\n",
    "        river (GeoDataFrame): A GeoDataFrame containing the extracted downstream river segments.\n",
    "    \"\"\"\n",
    "    # Prepare to store the segments\n",
    "    downstream_segments = []\n",
    "    total_length_km = 0\n",
    "    \n",
    "    # Create a dictionary for quick access to segments by HYRIV_ID\n",
    "    segment_dict = gdf.set_index('HYRIV_ID').to_dict(orient='index')\n",
    "    \n",
    "    # Track visited segments to avoid loops\n",
    "    visited = set()\n",
    "    \n",
    "    # Start with the given HYRIV_ID\n",
    "    current_id = start_hyriv_id\n",
    "    \n",
    "    while current_id != 0 and current_id not in visited:\n",
    "        # Mark the current segment as visited\n",
    "        visited.add(current_id)\n",
    "        \n",
    "        # Get the current segment from the dictionary\n",
    "        current_segment = segment_dict.get(current_id)\n",
    "        \n",
    "        if current_segment is not None:\n",
    "            downstream_segments.append(current_segment)\n",
    "            total_length_km += current_segment['LENGTH_KM']\n",
    "            \n",
    "            # Check if we've reached the stopping condition\n",
    "            if method == \"specific_segment\" and current_id == terminal_hyriv_id:\n",
    "                break\n",
    "            if method == \"length\" and threshold_length_km is not None and total_length_km >= threshold_length_km:\n",
    "                break\n",
    "            \n",
    "            # Move to the next downstream segment\n",
    "            current_id = current_segment['NEXT_DOWN']\n",
    "        else:\n",
    "            break\n",
    "    \n",
    "    # Convert the list of segments back into a GeoDataFrame\n",
    "    if downstream_segments:\n",
    "        # Ensure that the geometry column is properly set\n",
    "        river = gpd.GeoDataFrame(downstream_segments, geometry=[seg['geometry'] for seg in downstream_segments], crs=gdf.crs)\n",
    "    else:\n",
    "        river = gpd.GeoDataFrame(columns=gdf.columns, crs=gdf.crs)\n",
    "    \n",
    "    return river\n",
    "\n",
    "def extract_and_save_segments_from_csv(gdf, df):\n",
    "    \"\"\"\n",
    "    Extract segments from a GeoDataFrame (gdf) based on criteria provided in each row of a DataFrame (df)\n",
    "    and save each extracted river to a shapefile.\n",
    "\n",
    "    Parameters:\n",
    "        gdf (GeoDataFrame): The GeoDataFrame containing river segments.\n",
    "        df (DataFrame): The DataFrame containing criteria for extraction.\n",
    "    \"\"\"\n",
    "    for idx, row in df.iterrows():\n",
    "        # Extracting the necessary parameters from the current row\n",
    "        river_name = row['river_name']\n",
    "        start_hyriv_id = row['us_hyriv_id']  # Column name for starting HYRIV_ID\n",
    "        end_condition = row['end_condition']  # Segment or length condition\n",
    "        max_length_km = row['max_length_km']  # Maximum length in km\n",
    "        terminal_hyriv_id = row['ds_hyriv_id']  # Terminal HYRIV_ID if applicable\n",
    "        local_shapefile_directory = row['local_shapefile_directory']\n",
    "\n",
    "        # Construct the output directory path\n",
    "        output_directory = os.path.join(local_shapefile_directory, \"Extracted_Rivers\", river_name)\n",
    "        os.makedirs(output_directory, exist_ok=True)  # Create the directory if it doesn't exist\n",
    "\n",
    "        # Define the full path for the output shapefile\n",
    "        output_shapefile_path = os.path.join(output_directory, f\"{river_name}.shp\")\n",
    "\n",
    "        # Call the extract_downstream_river function\n",
    "        extracted_river = extract_downstream_river(\n",
    "            gdf, \n",
    "            start_hyriv_id=start_hyriv_id, \n",
    "            method=end_condition, \n",
    "            terminal_hyriv_id=terminal_hyriv_id, \n",
    "            threshold_length_km=max_length_km\n",
    "        )\n",
    "        \n",
    "        # Save the extracted river to a shapefile\n",
    "        extracted_river.to_file(output_shapefile_path)\n",
    "        \n",
    "def main(csv_path):\n",
    "    \"\"\"\n",
    "    Main function to process the CSV, download and extract shapefiles, and extract and save river segments.\n",
    "    \n",
    "    Parameters:\n",
    "        csv_path (str): The path to the input CSV file.\n",
    "    \"\"\"\n",
    "    # Step 1: Read the CSV into a DataFrame\n",
    "    df = pd.read_csv(csv_path)\n",
    "    \n",
    "    # Step 2: Find unique URLs in the 'online_url' column\n",
    "    unique_urls = df['online_url'].unique()\n",
    "    \n",
    "    for url in unique_urls:\n",
    "        # Step 3: Download and extract the shapefile, resulting in a GeoDataFrame\n",
    "        gdf = download_and_extract_shapefile(url)\n",
    "        \n",
    "        # Step 4: Create a sub-DataFrame containing only the records with the current URL\n",
    "        sub_df = df[df['online_url'] == url]\n",
    "        \n",
    "        # Step 5: Run the extraction and saving process with the sub-DataFrame and associated gdf\n",
    "        extract_and_save_segments_from_csv(gdf, sub_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf0a3967",
   "metadata": {},
   "source": [
    "## Input variables and run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "cf50e090",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\huckr\\miniconda3\\Lib\\site-packages\\geopandas\\io\\file.py:608: UserWarning: You are attempting to write an empty DataFrame to file. For some drivers, this operation may fail.\n",
      "  _to_file_fiona(df, filename, driver, schema, crs, mode, **kwargs)\n",
      "C:\\Users\\huckr\\miniconda3\\Lib\\site-packages\\geopandas\\io\\file.py:608: UserWarning: You are attempting to write an empty DataFrame to file. For some drivers, this operation may fail.\n",
      "  _to_file_fiona(df, filename, driver, schema, crs, mode, **kwargs)\n",
      "C:\\Users\\huckr\\miniconda3\\Lib\\site-packages\\geopandas\\io\\file.py:608: UserWarning: You are attempting to write an empty DataFrame to file. For some drivers, this operation may fail.\n",
      "  _to_file_fiona(df, filename, driver, schema, crs, mode, **kwargs)\n",
      "C:\\Users\\huckr\\miniconda3\\Lib\\site-packages\\geopandas\\io\\file.py:608: UserWarning: You are attempting to write an empty DataFrame to file. For some drivers, this operation may fail.\n",
      "  _to_file_fiona(df, filename, driver, schema, crs, mode, **kwargs)\n",
      "C:\\Users\\huckr\\miniconda3\\Lib\\site-packages\\geopandas\\io\\file.py:608: UserWarning: You are attempting to write an empty DataFrame to file. For some drivers, this operation may fail.\n",
      "  _to_file_fiona(df, filename, driver, schema, crs, mode, **kwargs)\n",
      "C:\\Users\\huckr\\miniconda3\\Lib\\site-packages\\geopandas\\io\\file.py:608: UserWarning: You are attempting to write an empty DataFrame to file. For some drivers, this operation may fail.\n",
      "  _to_file_fiona(df, filename, driver, schema, crs, mode, **kwargs)\n"
     ]
    }
   ],
   "source": [
    "csv_path = r\"C:\\Users\\huckr\\Desktop\\UCSB\\Dissertation\\Data\\RiverMapping\\Combo_river_datasheet.csv\" # Replace with the actual path to your CSV file\n",
    "main(csv_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "650acfc0",
   "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}