processing.py

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

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

# This module contains several functions that can be used to load and clean datasets 
# for the SAWS project.

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

import requests
import zipfile
import shutil
import io
import pandas as pd
import geopandas as gpd
import numpy as np
import rasterio
from rasterio.windows import from_bounds
from rasterio.features import geometry_mask
from shapely.geometry import Point
import os
import sys
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

from project_utils import saws_crs, saws_dir, scale_to_geo_unit_cols_dict, load_saws_gdf, saws_scales

##########################################
### LOADING DATA
##########################################

def load_data_from_url(url, zipped=False, filepath=None, **kwargs):
    """
    Extracts and returns raw data from a URL.

    Parameters:
    url : str
        The URL of the file to be extracted.
    zipped : bool, optional
        If True, the file is expected to be inside a zip file. Default is False.
    filepath : str, optional
        If zipped is True, this is the path to the file inside the archive.
    **kwargs :
        Additional arguments passed to the appropriate reader function 
        (e.g., pd.read_csv, pd.read_excel, gpd.read_file, rasterio.open)

    Returns:
    data : DataFrame, GeoDataFrame, or rasterio.DatasetReader
    """
    def detect_file_type(path):
        ext = os.path.splitext(path.lower())[1]
        if ext in ['.csv']: return 'csv'
        if ext in ['.txt']: return 'txt'
        elif ext in ['.xls']: return 'xls'
        elif ext in ['.xlsx']: return 'xlsx'
        elif ext in ['.shp', '.geojson', '.gpkg', '.json', '.kml']: return 'vector'
        elif ext in ['.tif', '.tiff']: return 'raster'
        else: return 'unknown'

    response = requests.get(url, headers={'User-Agent': 'Mozilla/5.0'})
    if response.status_code != 200:
        response.raise_for_status()

    if not zipped:
        file_type = detect_file_type(url)
        buffer = io.BytesIO(response.content)

        if file_type == 'csv' or file_type == 'txt':
            data = pd.read_csv(buffer, low_memory=False, **kwargs)
        elif file_type == 'xls':
            data = pd.read_excel(buffer, engine='xlrd', **kwargs)
        elif file_type == 'xlsx':
            data = pd.read_excel(buffer, engine='openpyxl', **kwargs)
        elif file_type == 'vector':
            data = gpd.read_file(buffer, **kwargs).to_crs(saws_crs)
        elif file_type == 'raster':
            data = rasterio.open(buffer, **kwargs)
        else: # fallback on csv
            try:
                data = pd.read_csv(buffer, low_memory=False, **kwargs) 
            except:
                raise ValueError(f"Could not download data from this URL. Please check URL and try again.")

    else:
        with zipfile.ZipFile(io.BytesIO(response.content)) as zip_file:
            if filepath is None:
                raise ValueError("Must specify 'filepath' within ZIP archive.")
            file_type = detect_file_type(filepath)

            try:
                with zip_file.open(filepath) as file:
                    if file_type == 'csv':
                        data = pd.read_csv(file, low_memory=False, **kwargs)
                    elif file_type == 'xls':
                        data = pd.read_excel(file, engine='xlrd', **kwargs)
                    elif file_type == 'xlsx':
                        data = pd.read_excel(file, engine='openpyxl', **kwargs)
                    elif file_type in ['vector', 'raster']:
                        raise NotImplementedError("Shapefiles and rasters require extraction.")
            except:
                zip_file.extractall('extracted_data')
                full_path = os.path.abspath(os.path.join('extracted_data', filepath))
                if file_type == 'csv':
                    data = pd.read_csv(full_path, low_memory=False, **kwargs)
                elif file_type == 'xls':
                    data = pd.read_excel(full_path, engine='xlrd', **kwargs)
                elif file_type == 'xlsx':
                    data = pd.read_excel(full_path, engine='openpyxl', **kwargs)
                elif file_type == 'vector':
                    data = gpd.read_file(full_path, **kwargs).to_crs(saws_crs)
                elif file_type == 'raster':
                    data = rasterio.open(full_path, **kwargs)
                shutil.rmtree('extracted_data')

    return data

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

def load_raw_eia_data(params):
    """
    This function loads EIA AEO data into a pandas dataframe using the provided
    query params. 

    Arguments:
    params : dict
        A dictionary of query parameters to filter the data.
    
    Returns:
    df : pandas.DataFrame
        A dataframe containing the EIA AEO data.

    """
    # base URL
    base_url = 'https://api.eia.gov/v2/aeo/2025/data/'

    # send GET request
    response = requests.get(base_url, params=params)

    # check if it worked
    if response.status_code != 200:
        response.raise_for_status()
    df = pd.DataFrame(response.json()['response']['data'])

    return df

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

def load_roots(scale, *root_codes, version='v1'):
    """
    This function loads the root datasets for a given process at the specified scale.

    Arguments:
    scale : str
        The resolution of data to load. ('county', 'state', 'huc2', 'huc4', 'huc6', 'huc8')
    *root_codes : str
        The dataset codes to load. These should be the codes used in the SAWS project.
    version : str, optional
        The version of the SAWS database from which the datasets come. Default is 'v1'.

    Returns:
    df_dict : dict
        A dictionary where keys are dataset codes and values are DataFrames
        containing the loaded datasets.
    """
    # check if scale is valid
    if scale not in saws_scales + ['overlay']:
        raise ValueError(f"Invalid scale '{scale}'. Must be one of {saws_scales}.")
    
    # build filepath for version directory
    version_dir = os.path.abspath(os.path.join(saws_dir, version))

    # initialize a dictionary to store the datasets
    df_dict = {}

    # iterate through dataset codes and load each dataset
    for dataset_code in root_codes:
        # if scale is standard, load as a DataFrame
        if scale in saws_scales:
            dataset_path = os.path.join(version_dir, dataset_code, 'data', f'{scale}-{dataset_code}.csv')
            # check if the file exists
            if not os.path.exists(dataset_path):
                raise FileNotFoundError(f"Dataset '{dataset_code}' at scale '{scale}' not found in version '{version}'.")
            df = pd.read_csv(dataset_path)
        # if scale is overlay, load as a GeoDataFrame
        elif scale == 'overlay':
            dataset_path = os.path.join(version_dir, dataset_code, 'data', f'{scale}-{dataset_code}.geojson')
            # check if the file exists
            if not os.path.exists(dataset_path):
                raise FileNotFoundError(f"Dataset '{dataset_code}' at scale '{scale}' not found in version '{version}'.")
            df = gpd.read_file(dataset_path).to_crs(saws_crs)

        # if scale is county, format geoid col as 5 char
        if scale == 'county':
            df['geoid'] = df['geoid'].astype(str).str.zfill(5)

        # add to dictionary
        df_dict[dataset_code] = df

    return df_dict

##########################################
### CLEANING DATA
##########################################
def configure_columns(df, col_config_dict):
    """
    This function configures a dataframe's columns to desired specifications.
    Columns will be reconfigured in the order of keys in col_config_dict and renamed
    according to their values or dropped if value is None (or key isn't in dictionary).

    Arguments:
    df : DataFrame or GeoDataFrame
        The dataframe to be configured.
    col_config_dict : dict
        A dictionary where keys are the current column names to be kept and values are the
        desired new column names.
    verbose : bool, optional
        If True, prints status updates during the loading process. Default is True.

    Returns:
    config_df : DataFrame or GeoDataFrame
        The dataframe with columns reconfigured according to col_config_dict.

    """
    # suppress pandas performance warnings
    from warnings import simplefilter
    simplefilter(action="ignore", category=pd.errors.PerformanceWarning)

    # create a new dataframe to store reconfigured data
    config_df = pd.DataFrame()

    # iterate through col_config_dict and populate config_df
    for old_col, new_col in col_config_dict.items():
        if old_col not in df.columns: # check to make sure column is in dataframe
            continue
        if col_config_dict[old_col] is not None: # if column is to be retained
            if isinstance(df[old_col], gpd.GeoSeries): # if column is geometry convert to gdf
                config_df = gpd.GeoDataFrame(config_df, geometry=df[old_col])
            else:
                config_df[new_col] = df[old_col].copy() # copy column to new dataframe

    return config_df

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

def filter_to_conus(df, scale):
    """
    This function filters a dataframe to only include rows that are within the
    contiguous United States (CONUS). 

    Arguments:
    df : DataFrame or GeoDataFrame
        The dataframe to be filtered.
    scale : str
        The resolution of data in df. ('county', 'state', 'huc2', 'huc4', 'huc6', 'huc8')

    """
    geo_units_in_conus = list(load_saws_gdf(scale)[scale_to_geo_unit_cols_dict[scale]])
    df = df[df[scale_to_geo_unit_cols_dict[scale]].isin(geo_units_in_conus)]

    return df

##########################################
### SPECIAL RESOLUTION PROCESSING 
##########################################

def lat_long_to_point(df, lat_col, long_col):
    
    """
    This function takes a DataFrame with lat/long columns stored as floats
    and converts it into a GeoDataFrame with a point geometry.

    Arguments:
    df : DataFrame
        The DataFrame to be converted.
    lat_col : str
        The name of the column containing latitude values.
    long_col : str
        The name of the column containing longitude values.

    Returns:
    df : GeoDataFrame
        The modified DataFrame as a GeoDataFrame with point geometry.
    """

    df_geo = [Point(lon, lat) for lon, lat in zip(df[long_col], df[lat_col])]
    df = gpd.GeoDataFrame(df, geometry=df_geo, crs = 'EPSG:4326')
    df = df.to_crs(saws_crs)
    
    return df

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

def raster_to_scale(raster_path, polygon_gdf, geo_unit_col, aggregation_method='mean'):
    """
    This function converts the provided raster file to the resolution of the 
    provided polygon_gdf by performing the specified aggregation method.

    Arguments:
    raster_path: str
        Local path to raster file.
    polygon_gdf: geopandas GeoDataFrame
        GeoDataFrame of polygons to which the raster will be converted
    geo_unit_col : str
        The column of polygon_gdf that should be used to index results.
    aggregation_method: str, default 'mean'
        Method to apply when aggregating raster values to polygon

    Returns:
    results : dict
        Dictionary of results from the raster to polygon conversion (keys are
        polygon IDs and values are the aggregated raster values)

    """

    # convert polygons to raster crs
    geo_units = polygon_gdf.to_crs(rasterio.open(raster_path).crs)

    # pull data from raster and aggregate by polygon
    with rasterio.open(raster_path) as raster:

        # initialize a dictionary to store the results:
        results_dict = {}

        # iterate over each polygon to compute zonal statistics
        for _, unit in geo_units.iterrows():
            try:
                # get the bounding box of the polygon
                minx, miny, maxx, maxy = unit.geometry.bounds
                window = from_bounds(minx, miny, maxx, maxy, raster.transform)

                # check if bounding box is valid
                if window.width <= 0.5 or window.height <= 0.5:
                    # skip this geometry if the window size is invalid
                    results_dict.update({
                        unit[geo_unit_col]: None
                    })
                    continue

                # read window from raster
                data = raster.read(1, window=window)

                # calculate transform for window
                window_transform = raster.window_transform(window)
            
                # mask raster to unit geometry
                unit_mask = geometry_mask(
                    [unit.geometry],
                    transform = window_transform,
                    invert=True,
                    out_shape = data.shape
                )
                masked_data = data[unit_mask]

                # calculate the aggregated value in geographic unit (ignore nodata values)
                valid_data = masked_data[masked_data != raster.nodata]

                if valid_data.size == 0:
                    value = None  # return None when there is no valid data
                else:
                    if aggregation_method == 'mean':
                        value = np.nanmean(valid_data) 
                    elif aggregation_method == 'sum':
                        value = np.nansum(valid_data)
                    else:
                        raise ValueError('Invalid input for argument aggregation_method.')
            except:
                value = None
            
            results_dict.update({
                unit[geo_unit_col]: value
            })

    return results_dict