write_bin_polygon.py

import numpy as np
import json
import geopandas as gpd
from shapely.geometry import mapping
from consts import files_path, binary_data_dir
import pickle
import os
from collections import defaultdict

def build_threshold_rgba(a, C):
    # Bins normalized between 0 and 1
    norm = [(float(i) - min(a)) / (max(a) - min(a)) for i in a]

    # Generate the desired output format
    output = []

    for value, color in zip(a, C):
        rgba_color = (int(color[0]), int(color[1]), int(color[2]), 255)
        output.append({"value": value, "color": rgba_color})

    return output

def get_color_for_value(var_threshold, value):
    """
    Given a value and a list of thresholds, return the corresponding color.
    If the value is non-numeric, return a fully transparent color.
    If the value is below the first threshold, return the first color.
    If the value is above the last threshold, return the last color.
    """
    if not isinstance(value, (int, float)) or value == 0:
        # Fully transparent
        return (0, 0, 0, 0)

    for threshold in var_threshold:
        if value <= threshold["value"]:
            return threshold["color"]
    # If value exceeds all var_threshold, return the last color
    return var_threshold[-1]["color"]

def process_files(files, tracts_gdf):
    """
    Process each variable file to compute average values by census tract and include geometries.
    """

    # Directory to save processed pickle files
    os.makedirs(binary_data_dir, exist_ok=True)  # Create the directory if it doesn't exist

    for file in files:
        var_name = file["var_name"]
        file_path = file["path"]

        print(f"Processing variable: {var_name}")

        # Open and load the GeoJSON file using geopandas
        try:
            points_gdf = gpd.read_file(file_path)
        except FileNotFoundError:
            print(f"File not found: {file_path}. Skipping {var_name}.")
            continue
        except Exception as e:
            print(f"Error loading {file_path}: {e}. Skipping {var_name}.")
            continue

        # Ensure the GeoDataFrame has a geometry column
        if points_gdf.empty or 'geometry' not in points_gdf:
            print(f"No geometry found in {file_path}. Skipping {var_name}.")
            continue

        # Ensure CRS matches between points and tracts
        if points_gdf.crs != tracts_gdf.crs:
            points_gdf = points_gdf.to_crs(tracts_gdf.crs)

        # Extract all years from the properties
        # Assuming year properties are all keys that are purely digits
        sample_properties = points_gdf.iloc[0].drop(labels='geometry').to_dict()
        years = sorted([key for key in sample_properties.keys() if key.isdigit()])

        if not years:
            print(f"No year properties found in {file_path}. Skipping {var_name}.")
            continue

        print(f"Processing variable '{var_name}' for years: {', '.join(years)}")

        for year in years:
            # Select the value column for the current year
            if year not in points_gdf.columns:
                print(f"Year '{year}' not found in {file_path}. Skipping year.")
                continue

            # Create a GeoDataFrame with only necessary columns
            df_year = points_gdf[['geometry', year]].copy()
            df_year = df_year.dropna(subset=[year])

            # Spatial join with tracts to assign each point to a tract
            try:
                joined = gpd.sjoin(df_year, tracts_gdf, how='inner', predicate='within')
            except Exception as e:
                print(f"Error during spatial join for {var_name} {year}: {e}")
                continue

            # Group by tract and compute average
            grouped = joined.groupby('GEOID')[year].mean().reset_index()
            grouped.rename(columns={year: 'average_value'}, inplace=True)

            # Assign colors based on average values
            grouped['color'] = grouped['average_value'].apply(lambda x: get_color_for_value(file["threshold"], x))

            # Merge with tracts_gdf to get geometries
            grouped = grouped.merge(tracts_gdf[['GEOID', 'geometry']], on='GEOID', how='left')

            # Convert geometries to GeoJSON-like dicts
            grouped['geometry'] = grouped['geometry'].apply(lambda geom: mapping(geom))

            # Prepare the binary data
            binary_data = {
                "tracts": grouped[['GEOID', 'average_value', 'color', 'geometry']].to_dict(orient='records')
            }

            # Define the filename
            filename = f"ct_{var_name}_{year}.pickle"
            filepath = os.path.join(binary_data_dir, filename)

            # Save the binary data using pickle
            try:
                with open(filepath, 'wb') as pf:
                    pickle.dump(binary_data, pf)
                print(f"Saved binary data for {var_name} {year} to {filepath}")
            except IOError as e:
                print(f"Failed to save binary data to {filepath}: {e}")

# Define thresholds and colors as in your original code
start, end, n = -50, 50, 38
temp_range = [round(start + (end - start) * i / (n - 1), 1) for i in range(n)]

temp_colors = np.array([
    [145, 0, 63],
    [206, 18, 86],
    [231, 41, 138],
    [223, 101, 176],
    [255, 115, 223],
    [255, 190, 232],
    [255, 255, 255],
    [218, 218, 235],
    [188, 189, 220],
    [158, 154, 200],
    [117, 107, 177],
    [84, 39, 143],
    [13, 0, 125],
    [13, 61, 156],
    [0, 102, 194],
    [41, 158, 255],
    [74, 199, 255],
    [115, 215, 255],
    [173, 255, 255],
    [48, 207, 194],
    [0, 153, 150],
    [18, 87, 87],
    [6, 109, 44],
    [49, 163, 84],
    [116, 196, 118],
    [161, 217, 155],
    [211, 255, 190],
    [255, 255, 179],
    [255, 237, 160],
    [254, 209, 118],
    [254, 174, 42],
    [253, 141, 60],
    [252, 78, 42],
    [227, 26, 28],
    [177, 0, 38],
    [128, 0, 38],
    [89, 0, 66],
    [40, 0, 40]
])

prcp_range_inches = [0, 0.01, 0.1, 0.25, 0.5, 1, 1.5, 2, 3, 4, 6, 8, 10, 15, 20, 30]
prcp_range_mm = [round(value * 25.4, 2) for value in prcp_range_inches]

prcp_colors = np.array([
    [255, 255, 255],
    [199, 233, 192],
    [161, 217, 155],
    [116, 196, 118],
    [49, 163, 83],
    [0, 109, 44],
    [255, 250, 138],
    [255, 204, 79],
    [254, 141, 60],
    [252, 78, 42],
    [214, 26, 28],
    [173, 0, 38],
    [112, 0, 38],
    [59, 0, 48],
    [76, 0, 115],
    [255, 219, 255]
])

temp_threshold = build_threshold_rgba(temp_range, temp_colors)
prcp_threshold = build_threshold_rgba(prcp_range_mm, prcp_colors)

# List of files
geojson_files = [
    {"var_name": "tmin", "path": f"{files_path}/Yearly_tmin_round.json", "threshold": temp_threshold},
    {"var_name": "tmax", "path": f"{files_path}/Yearly_tmax_round.json", "threshold": temp_threshold},
    {"var_name": "prcp", "path": f"{files_path}/Illinois_prcp_risks_round.json", "threshold": prcp_threshold},
]

# Load Census Tracts GeoJSON
tracts_geojson_path = f"{files_path}/tl_2023_17_tract.json"  # Update the path accordingly
try:
    tracts_gdf = gpd.read_file(tracts_geojson_path)
except FileNotFoundError:
    print(f"Census tract file not found: {tracts_geojson_path}. Exiting.")
    exit(1)
except Exception as e:
    print(f"Error loading census tracts GeoJSON: {e}. Exiting.")
    exit(1)

# Ensure there is a unique identifier for each tract
# Replace 'GEOID' with the actual property name in your GeoJSON
if 'GEOID' not in tracts_gdf.columns:
    print("Error: 'GEOID' field not found in census tracts GeoJSON. Please update the code with the correct field name.")
    exit(1)

# Process the files
process_files(geojson_files, tracts_gdf)