#!/usr/bin/env python3 """ Weather ETL for Wuhan Respiratory Disease Risk Prediction Platform. Processes weather CSV files into daily Wuhan parquet. """ import argparse from pathlib import Path import pandas as pd # Wuhan station metadata (from station list CSV) WUHAN_STATIONS = {} # Coordinate bounding box for Wuhan area WUHAN_LAT_MIN, WUHAN_LAT_MAX = 29.9, 31.5 WUHAN_LON_MIN, WUHAN_LON_MAX = 113.7, 115.2 # Pollutant type mapping to output schema POLLUTANT_MAP = { 'AQI': 'AQI', 'PM2.5': 'PM25', 'PM2.5_24h': 'PM25_24h', 'PM10': 'PM10', 'PM10_24h': 'PM10_24h', 'SO2': 'SO2', 'SO2_24h': 'SO2_24h', 'NO2': 'NO2', 'NO2_24h': 'NO2_24h', 'O3': 'O3', 'O3_24h': 'O3_24h', 'O3_8h': 'O3_8h', 'O3_8h_24h': 'O3_8h_24h', 'CO': 'CO', 'CO_24h': 'CO_24h', 'NOx': 'NOX', 'primary_pollutant': 'PRIMARY_POLLUTANT', 'air_quality_level': 'AIR_QUALITY_LEVEL', } # Core pollutants for output (7 pollutants as per plan) OUTPUT_POLLUTANTS = ['AQI', 'PM25', 'PM10', 'SO2', 'NO2', 'O3', 'CO'] def load_station_list(station_file: str) -> dict: """Load station metadata from station list CSV.""" global WUHAN_STATIONS stations = {} df = pd.read_csv(station_file, encoding='utf-8') for _, row in df.iterrows(): station_id = str(row['监测点编码']).strip() city = str(row['城市']).strip() if pd.notna(row['城市']) else '' lat_val = row['纬度'] lon_val = row['经度'] try: lat = float(lat_val) if pd.notna(lat_val) and lat_val != '-' else 0 except (ValueError, TypeError): lat = 0 try: lon = float(lon_val) if pd.notna(lon_val) and lon_val != '-' else 0 except (ValueError, TypeError): lon = 0 district = str(row['监测点名称']).strip() if pd.notna(row['监测点名称']) else '' if (city == '武汉' or (WUHAN_LAT_MIN <= lat <= WUHAN_LAT_MAX and WUHAN_LON_MIN <= lon <= WUHAN_LON_MAX)): stations[station_id] = { 'name': district, 'lat': lat, 'lon': lon, 'district': district, } WUHAN_STATIONS = stations return stations def process_daily_csv(csv_path: str, wuhan_stations: list[str]) -> pd.DataFrame: """Process a single daily CSV file. Args: csv_path: Path to china_sites_YYYYMMDD.csv wuhan_stations: List of Wuhan station IDs. If empty, process ALL stations. Returns: DataFrame with columns: datetime, station_id, pollutant, value """ df = pd.read_csv(csv_path) if wuhan_stations: wuhan_cols = ['date', 'hour', 'type'] + wuhan_stations available_cols = [c for c in wuhan_cols if c in df.columns] else: available_cols = df.columns.tolist() df = df[available_cols] id_vars = ['date', 'hour', 'type'] value_vars = [c for c in available_cols if c not in id_vars] if not value_vars: return pd.DataFrame(columns=['datetime', 'station_id', 'pollutant', 'value']) df_long = df.melt( id_vars=id_vars, value_vars=value_vars, var_name='station_id', value_name='value', ) df_long['datetime'] = pd.to_datetime( df_long['date'].astype(str) + df_long['hour'].astype(str).str.zfill(2), format='%Y%m%d%H' ) df_long['pollutant'] = df_long['type'].map(POLLUTANT_MAP) return df_long[['datetime', 'station_id', 'pollutant', 'value']] def aggregate_to_daily(df_long: pd.DataFrame, wuhan_metadata: dict) -> pd.DataFrame: """Aggregate hourly data to daily level per station. Uses mean for all pollutants. """ # Filter to output pollutants only df_pollutants = df_long[df_long['pollutant'].isin(OUTPUT_POLLUTANTS)].copy() # Extract date (without time) from datetime df_pollutants['date'] = df_pollutants['datetime'].dt.date # First aggregate by (date, station_id, pollutant) to get daily mean df_daily_pollutant = df_pollutants.groupby( ['date', 'station_id', 'pollutant'], as_index=False )['value'].mean() # Pivot to wide format: one column per pollutant df_pivot = df_daily_pollutant.pivot_table( index=['date', 'station_id'], columns='pollutant', values='value', aggfunc='mean' ).reset_index() # Flatten column names df_pivot.columns.name = None # Add metadata df_pivot['district'] = df_pivot['station_id'].map( lambda x: wuhan_metadata.get(x, {}).get('district', '') ) df_pivot['lat'] = df_pivot['station_id'].map( lambda x: wuhan_metadata.get(x, {}).get('lat', 0) ) df_pivot['lon'] = df_pivot['station_id'].map( lambda x: wuhan_metadata.get(x, {}).get('lon', 0) ) # Ensure output schema columns exist for col in OUTPUT_POLLUTANTS: if col not in df_pivot.columns: df_pivot[col] = None # Reorder columns output_cols = ['date', 'station_id', 'district', 'lat', 'lon'] + OUTPUT_POLLUTANTS df_pivot = df_pivot[[c for c in output_cols if c in df_pivot.columns]] return df_pivot def process_year(input_dir: str, output_dir: str, year: int, station_file: str = None) -> None: """Process all CSV files for a given year.""" from pathlib import Path import glob as glob_module import os output_path = Path(output_dir) output_path.mkdir(parents=True, exist_ok=True) # Find station list file dynamically if not provided if station_file is None or not os.path.exists(station_file): base_dir = '/home/akiba/CA/Datas' station_files = glob_module.glob(os.path.join(base_dir, '*空气*', '*列表*.csv')) if station_files: station_file = station_files[0] print(f'Found station list: {station_file}') else: print(f'ERROR: No station list file found') return if station_file and os.path.exists(station_file): print(f'Loading station list from {station_file}...') wuhan_stations = load_station_list(station_file) station_ids = list(wuhan_stations.keys()) print(f' Found {len(station_ids)} Wuhan stations: {station_ids}') else: print(f'ERROR: Station file not found: {station_file}') return # Find year directory dynamically base_dir = '/home/akiba/CA/Datas' year_dirs = glob_module.glob(os.path.join(base_dir, '*空气*', f'站点_{year}*')) # Filter out .zip files and Zone.Identifier year_dirs = [d for d in year_dirs if os.path.isdir(d)] if not year_dirs: print(f'No directory found for year {year}') return year_dir = year_dirs[0] print(f'Using year directory: {year_dir}') csv_pattern = os.path.join(year_dir, f'china_sites_{year}*.csv') csv_files = sorted(glob_module.glob(csv_pattern)) if not csv_files: print(f'No CSV files found for year {year}') return print(f'Processing {len(csv_files)} files for year {year}...') all_data = [] for i, csv_file in enumerate(csv_files): try: df = process_daily_csv(csv_file, station_ids) all_data.append(df) if i == 0: print(f' First file processed: {df.shape}') if (i + 1) % 50 == 0: print(f' Processed {i + 1}/{len(csv_files)} files...') except Exception as e: print(f'Error processing {csv_file}: {e}') if not all_data: print('No data processed successfully.') return df_combined = pd.concat(all_data, ignore_index=True) print(f'Combined data shape: {df_combined.shape}') df_daily = aggregate_to_daily(df_combined, wuhan_stations) df_daily = df_daily.sort_values(['date', 'station_id']).reset_index(drop=True) output_file = output_path / f'daily_wuhan_{year}.parquet' df_daily.to_parquet(output_file, index=False) print(f'Output: {output_file}') print(f'Shape: {df_daily.shape}') print(f'Columns: {list(df_daily.columns)}') print(f'Date range: {df_daily["date"].min()} to {df_daily["date"].max()}') print(f'Stations: {df_daily["station_id"].nunique()}') def main(): parser = argparse.ArgumentParser(description='Process weather data for Wuhan') parser.add_argument('--year', type=int, required=True, help='Year to process (e.g., 2022)') parser.add_argument( '--output-dir', type=str, default='processed/weather', help='Output directory for parquet files' ) parser.add_argument( '--station-file', type=str, default=None, help='Station list CSV file (auto-detected if not provided)' ) args = parser.parse_args() process_year(None, args.output_dir, args.year, args.station_file) if __name__ == '__main__': main()