project_utils.py

#######################################################################################

# PROJECT: SAWS
# FILE NAME: projcect_utils.py
# AUTHOR: Caroline Adkins
# LAST UPDATED: June 2025

# This module contains utility functions and constants for the SAWS project.

#######################################################################################
#######################################################################################

import requests
import pandas as pd
import geopandas as gpd
import zipfile
import shutil
import io
import os

##########################################
### PROJECT PARAMETERS
##########################################
saws_crs = 'ESRI:102003' # coordinate reference system for project
saws_scales = ['county', 'state', 'huc8', 'huc6', 'huc4', 'huc2'] # resolutions supported by SAWS

##########################################
### LOCAL DIRECTORIES
##########################################
saws_dir = os.path.abspath(os.path.join(os.path.dirname(__file__))) # main saws directory
shapefiles_dir = os.path.abspath(os.path.join(saws_dir, 'shapefiles')) # directory for shapefiles
data_dir = os.path.abspath(os.path.join(saws_dir, 'data')) # directory for data files

##########################################
### SCALE DEFINITIONS AND CONVERSIONS
##########################################
scale_to_scale_name_dict = { # dictionary of scale codes to scale names
    'county' : 'County',
    'state'  : 'State',
    'huc8'   : 'HUC-8',
    'huc6'   : 'HUC-6',
    'huc4'   : 'HUC-4',
    'huc2'   : 'HUC-2'
}

scale_to_boundary_type_dict = { # dictionary of type of boundary for each scale
    'county' : {'political'},
    'state'  : {'political', 'hydrological'}, # HUC shapefiles are split by state so we include hydrological
    'huc8'   : {'hydrological'},
    'huc6'   : {'hydrological'},
    'huc4'   : {'hydrological'},
    'huc2'   : {'hydrological'}
}

scale_to_geo_unit_cols_dict = { # dictionary of scale codes to geo unit columns
    'county' : 'geoid',
    'state'  : 'state_alpha',
    'huc8'   : 'huc8id',
    'huc6'   : 'huc6id',
    'huc4'   : 'huc4id',
    'huc2'   : 'huc2id',
    'nation' : 'nation_name'
}

##########################################
### HELPER FUNCTIONS
##########################################

def dict_from_df(df, key_col, val_col):
    """
    This function converts a DataFrame into a dictionary where the values of one column
    are the keys and the values of another column are the values.

    Arguments:
    df : DataFrame
        The DataFrame to be converted.
    key_col : str
        The name of the column to be used as keys in the dictionary.
    val_col : str
        The name of the column to be used as values in the dictionary.

    Returns:
    new_dict : dict
        The dictionary created from the DataFrame.

    """
    if key_col == 'index':
        new_dict = dict(zip(df.index, df[val_col]))
    else:
        new_dict = dict(zip(df[key_col], df[val_col]))

    return new_dict

##########################################

def load_saws_gdf(*scales, full_res = False):
    """
    This function loads all or a subset of available SAWS shapefiles either in full or
    simplified resolution.

    Arguments:
    *scales: sequence of str
        Sequence of scale codes for which to load shapefiles. Can refer to
        individual shapefiles or 'all'.
    full_res : bool, optional
        If True, loads full resolution shapefiles. If False, loads simplified 
        resolution shapefiles. Default is False.
    
    Returns:
    gdf : GeoDataFrame
        If only one scale is requested, returns a GeoDataFrame for that scale.
    gdf_dict : dictionary
        If multiple scales are requested, returns a dictionary where keys are scale 
        codes and values are GeoDataFrames.    
    
    """
    # determine which set of shapefiles to load
    scales_to_load = []
    for scale in scales:
        if scale == 'all':
            scales_to_load = saws_scales
        else:
            scales_to_load += [scale]
    scales_to_load = set(scales_to_load)

    # laod shapefiles and store in dictionary
    gdf_dict = {}
    for scale in scales_to_load:
        # identify file name based on resolution
        if full_res==True:
            filename =f'conus_{scale}_boundaries_full_res.geojson'
            filepath = os.path.abspath(os.path.join(shapefiles_dir, 'full_res', filename))
        else:
            filename = f'conus_{scale}_boundaries.geojson'
            filepath = os.path.abspath(os.path.join(shapefiles_dir, filename))
        
        # make sure file exists
        if not os.path.exists(filepath):
            raise FileNotFoundError(f'{filename} not found in shapefiles directory.')
        try:
            gdf = gpd.read_file(filepath).set_crs(saws_crs, allow_override=True)
            gdf['geometry'] = gdf['geometry'].make_valid()
            gdf = gdf.sort_values(by=scale_to_geo_unit_cols_dict[scale]).reset_index(drop=True)
            gdf_dict[scale] = gdf.copy()
        except:
            gdf_dict[scale] = None

        # if loading county dataframe, add a column for state alpha
        if scale == 'county':
            statefp_to_alpha_dict = dict_from_df(load_saws_gdf('state'), 'statefp', 'state_alpha')
            gdf_dict[scale].insert(2, 'state_alpha', gdf_dict[scale]['statefp'].map(statefp_to_alpha_dict))
    
    if len(scales_to_load) == 1:
        return gdf_dict[scales_to_load.pop()]
    else:
        return gdf_dict