hecras_curve_number_calculation.txt

import os
import geopandas as gpd
import pandas as pd

# Placeholder strings for InfilGPT to update variables set equal to "<UPDATE>"
SOIL_SHAPEFILE_PATH = r"<UPDATE>"
LAND_USE_SHAPEFILE_PATH = r"<UPDATE>"
HSG_FIELD_NAME = "<UPDATE>"
LAND_USE_FIELD_NAME = "<UPDATE>"
CN_A_FIELD_NAME = "<UPDATE>"
CN_B_FIELD_NAME = "<UPDATE>"
CN_C_FIELD_NAME = "<UPDATE>"
CN_D_FIELD_NAME = "<UPDATE>"
OUTPUT_SHAPEFILE_PATH = r"<UPDATE>"
OUTPUT_SPREADSHEET_PATH = r"<UPDATE>"

def load_shapefiles(soil_shapefile_path, land_use_shapefile_path, subbasins_shapefile_path=None):
    soil_df = gpd.read_file(soil_shapefile_path)
    land_use_df = gpd.read_file(land_use_shapefile_path)
    if soil_df.crs != land_use_df.crs:
        land_use_df = land_use_df.to_crs(soil_df.crs)
    return soil_df, land_use_df

def intersect_geometries_custom(df1, df2):
    intersections = []
    for index1, row1 in df1.iterrows():
        for index2, row2 in df2.iterrows():
            if row1['geometry'].intersects(row2['geometry']):
                intersected_geom = row1['geometry'].intersection(row2['geometry'])
                combined_attributes = row1.to_dict()
                combined_attributes.update(row2.to_dict())
                combined_attributes['geometry'] = intersected_geom
                intersections.append(combined_attributes)
    return gpd.GeoDataFrame(intersections)

def intersect_shapefiles(soil_df, land_use_df):
    intersected_df = intersect_geometries_custom(soil_df, land_use_df)
    return intersected_df

def assign_curve_number(intersected_df, hsg_field_name, cn_a_field_name, cn_b_field_name, cn_c_field_name, cn_d_field_name):
    curve_number_mapping = {
        'A': cn_a_field_name,
        'B': cn_b_field_name,
        'C': cn_c_field_name,
        'D': cn_d_field_name
    }

    existing_hsgs = intersected_df[hsg_field_name].unique()
    curve_number_mapping = {hsg: curve_number_mapping[hsg] for hsg in existing_hsgs if hsg in curve_number_mapping}

    def get_curve_number(row):
        hsg = row[hsg_field_name]
        cn_field = curve_number_mapping.get(hsg)
        return row[cn_field] if cn_field and cn_field in row else None

    intersected_df['Final_CN'] = intersected_df.apply(get_curve_number, axis=1)
    return intersected_df

def create_new_shapefile(intersected_df, output_path):
    if not isinstance(intersected_df, gpd.GeoDataFrame):
        intersected_df = gpd.GeoDataFrame(intersected_df, geometry='geometry')
    
    export_columns = ['LandUseHSG', 'Final_CN', 'geometry']
    intersected_df[export_columns].to_file(output_path, driver='ESRI Shapefile')

def export_spreadsheet(intersected_df, output_spreadsheet_path):
    export_columns = ['LandUseHSG', 'Final_CN']
    unique_df = intersected_df[export_columns].drop_duplicates().sort_values('LandUseHSG')
    unique_df.to_csv(os.path.join(output_spreadsheet_path), index=False)

soil_df, land_use_df = load_shapefiles(SOIL_SHAPEFILE_PATH, LAND_USE_SHAPEFILE_PATH)
intersected_df = intersect_shapefiles(soil_df, land_use_df)
intersected_df = assign_curve_number(intersected_df, HSG_FIELD_NAME, CN_A_FIELD_NAME, CN_B_FIELD_NAME, CN_C_FIELD_NAME, CN_D_FIELD_NAME)
    
intersected_df['LandUseHSG'] = intersected_df[LAND_USE_FIELD_NAME] + ": " + intersected_df[HSG_FIELD_NAME]
    
export_spreadsheet(intersected_df, OUTPUT_SPREADSHEET_PATH)
create_new_shapefile(intersected_df, OUTPUT_SHAPEFILE_PATH)