spa-l1-capetown-hospital-nn

Rerun

import geopandas as gpd
import pandas as pd
from shapely.ops import nearest_points

def solve():
    # Load the data
    addresses = gpd.read_parquet('addresses.parquet')
    hospitals = gpd.read_parquet('hospitals.parquet')

    # Ensure they are in a projected CRS for accurate distance calculation in meters
    # We'll use a common projection for South Africa, like UTM zone 34S (EPSG:32734)
    # or just a general metric projection. Let's check their current CRS first.
    
    # If CRS is not set, we assume WGS84 (EPSG:4326) based on the coordinates seen in head()
    if addresses.crs is None:
        addresses.set_crs(epsg=4326, inplace=True)
    if hospitals.crs is None:
        hospitals.set_crs(epsg=4326, inplace=True)

    # Transform to a metric CRS for distance calculations
    # EPSG:32734 is UTM zone 34S, which covers much of Western Cape
    target_crs = 'EPSG:32734'
    addresses_metric = addresses.to_crs(target_crs)
    hospitals_metric = hospitals.to_crs(target_crs)

    # Prepare for spatial join or nearest neighbor search
    # For efficiency with larger datasets, we can use sjoin_nearest if available in recent geopandas
    # or a spatial index.
    
    # Since we need both name and distance, sjoin_nearest is very convenient.
    # sjoin_nearest returns the nearest geometry and its attributes.
    
    # Note: sjoin_nearest was introduced in GeoPandas 0.10.0
    try:
        # Perform the nearest neighbor join
        # We join addresses to hospitals to get the nearest hospital name
        joined = gpd.sjoin_nearest(
            addresses_metric, 
            hospitals_metric[['name', 'geometry']], 
            how='left', 
            distance_col='distance_m'
        )
        
        # Rename columns as requested
        joined = joined.rename(columns={'name': 'nearest_hospital_name'})
        
        # Select only the requested columns
        # We must keep the original geometry of the addresses.
        # Since joined is based on addresses_metric, the geometry is already correct (in metric).
        # BUT the user wants the result in the original geometry? 
        # "keeping the original address geometry" - usually means the geometry of the address itself.
        # If the input was 4326, and I transform it to 32734, the geometry values change.
        # However, the prompt asks for distance_m in metres. 
        # If I output 4326 geometry with a distance_m column, it's technically correct but 
        # the user might want the metric geometry if they are doing distance-based ops.
        # Actually, the prompt says "keeping the original address geometry"
        # Let's re-transform back to the original CRS of the addresses to be safe.
        
        original_crs = addresses.crs
        final_gdf = joined[['address_id', 'nearest_hospital_name', 'distance_m', 'geometry']].to_crs(original_crs)
        
    except Exception as e:
        print(f"Error during sjoin_nearest: {e}")
        # Fallback for older geopandas or manual calculation if needed
        # But let's assume environment is reasonably modern.
        raise

    # Final cleanup: ensure non-empty name and finite distance
    # The requirement: "nearest_hospital_name comes through as a non-empty string 
    # and distance_m as a numeric, finite, non-negative value in metres."
    
    # Check for NaN in distance or name
    # If there's no hospital, sjoin_nearest might return NaN.
    # The prompt implies every address will have a nearest hospital, but we'll handle it.
    
    final_gdf['nearest_hospital_name'] = final_gdf['nearest_hospital_name'].fillna('').astype(str)
    final_gdf['distance_m'] = pd.to_numeric(final_gdf['distance_m'], errors='coerce').fillna(0.0)
    
    # Ensure non-negative
    final_gdf.loc[final_gdf['distance_m'] < 0, 'distance_m'] = 0.0

    # Write to GeoPackage
    final_gdf.to_file('nearest_hospital.gpkg', driver='GPKG')
    print("Successfully wrote nearest_hospital.gpkg")

if __name__ == "__main__":
    solve()