#!/usr/bin/env python3 import os import glob import numpy as np import xarray as xr from datetime import datetime, timedelta import re # ============================================================================ # PATH CONFIGURATION # ============================================================================ ICONKIT_BASE_PATH = "/global/cfs/projectdirs/m4486/Haochen/TRACER_Model/ICON_KIT" OUTPUT_BASE = "/global/cfs/projectdirs/m4486/Haochen/Extracted_Data/Hourly_Precip/ICON_KIT" # Scenario folder names and output names SCENARIO_CONFIG = { 'RefCCN': 'control', 'HighCCN': 'high_ccn', 'LowCCN': 'low_ccn' } CASE_CONFIG = { 'june17': { 'case_folder': '17June_500m', 'base_datetime': datetime(2022, 6, 17, 6, 0), 'file_pattern': '*_10min_reglatlon.nc', }, 'aug7': { 'case_folder': '07Aug_500m', 'base_datetime': datetime(2022, 8, 7, 6, 0), 'file_pattern': '*_10min_reglatlon.nc', } } def parse_iconkit_filename(filename): """Parse datetime from ICON_KIT filename. Example: NWP_LAM_DOM01_20220617T060000Z_TRACER_10min_reglatlon.nc """ basename = os.path.basename(filename) # Extract datetime string: YYYYMMDDTHHMMSS match = re.search(r'(\d{8}T\d{6})Z', basename) if match: dt_str = match.group(1) return datetime.strptime(dt_str, '%Y%m%dT%H%M%S') return None def extract_precipitation(case_name, scenario_folder, output_experiment): """ Extract tot_prec and tot_prec_rate from ICON_KIT files. Parameters: ----------- case_name : str Case name (june17 or aug7) scenario_folder : str ICON_KIT folder name (RefCCN, HighCCN, LowCCN) output_experiment : str Output experiment name (control, high_ccn, low_ccn) """ config = CASE_CONFIG[case_name] print(f"\n{'='*70}") print(f"Processing: ICON_KIT - {case_name} - {scenario_folder} -> {output_experiment}") print(f"{'='*70}") # Build file path case_path = os.path.join(ICONKIT_BASE_PATH, config['case_folder'], scenario_folder) file_pattern = os.path.join(case_path, config['file_pattern']) # Output file os.makedirs(OUTPUT_BASE, exist_ok=True) output_file = os.path.join(OUTPUT_BASE, f"ICON_KIT_{case_name}_{output_experiment}.nc") print(f"Input path: {case_path}") print(f"Output file: {output_file}") # Check if output exists if os.path.exists(output_file): print(f"Output file already exists. Skipping... (delete to re-extract)") return True # Find all 10-min files files = sorted(glob.glob(file_pattern)) if not files: print(f"ERROR: No files found matching {file_pattern}") return False print(f"Found {len(files)} files") # Extract data from each file tot_prec_list = [] tot_prec_rate_list = [] time_list = [] lat_vals = None lon_vals = None for i, fpath in enumerate(files): if i % 20 == 0: print(f" Reading file {i+1}/{len(files)}: {os.path.basename(fpath)}") try: ds = xr.open_dataset(fpath) # Get precipitation variables if 'tot_prec' not in ds or 'tot_prec_rate' not in ds: print(f" Warning: Missing variables in {fpath}") ds.close() continue # Extract data (squeeze time dimension if single timestep) tot_prec = ds['tot_prec'].values tot_prec_rate = ds['tot_prec_rate'].values # Handle dimensions - squeeze if needed if tot_prec.ndim == 3 and tot_prec.shape[0] == 1: tot_prec = tot_prec[0, :, :] if tot_prec_rate.ndim == 3 and tot_prec_rate.shape[0] == 1: tot_prec_rate = tot_prec_rate[0, :, :] tot_prec_list.append(tot_prec) tot_prec_rate_list.append(tot_prec_rate) # Get time from filename file_dt = parse_iconkit_filename(fpath) if file_dt: time_list.append(file_dt) # Get coordinates (only once) if lat_vals is None: lat_vals = ds['lat'].values lon_vals = ds['lon'].values ds.close() except Exception as e: print(f" ERROR reading {fpath}: {e}") continue if not tot_prec_list: print("ERROR: No data extracted") return False # Stack data print(" Combining data...") tot_prec_arr = np.stack(tot_prec_list, axis=0) tot_prec_rate_arr = np.stack(tot_prec_rate_list, axis=0) time_arr = np.array(time_list, dtype='datetime64[ns]') print(f" Total timesteps: {len(time_list)}") print(f" Time range: {time_list[0]} to {time_list[-1]}") print(f" Data shape: {tot_prec_arr.shape}") # Create output dataset ds_out = xr.Dataset({ 'tot_prec': (['time', 'lat', 'lon'], tot_prec_arr), 'tot_prec_rate': (['time', 'lat', 'lon'], tot_prec_rate_arr), }, coords={ 'time': time_arr, 'lat': lat_vals, 'lon': lon_vals }) # Add metadata ds_out.attrs = { 'title': 'ICON_KIT Precipitation Data', 'case': case_name, 'scenario': output_experiment, 'description': 'Extracted precipitation from ICON_KIT 10-min output', 'creation_date': datetime.now().strftime('%Y-%m-%d %H:%M:%S'), 'creator': 'Haochen Zhang', 'model': 'ICON_KIT', 'n_files_processed': len(files) } ds_out['tot_prec'].attrs = { 'long_name': 'Total Accumulated Precipitation', 'units': 'mm', 'description': 'Time-integrated total precipitation (kg/m² = mm)' } ds_out['tot_prec_rate'].attrs = { 'long_name': 'Total Precipitation Rate', 'units': 'kg m-2 s-1', 'description': 'Instantaneous total precipitation rate' } ds_out['lat'].attrs = {'long_name': 'Latitude', 'units': 'degrees_north'} ds_out['lon'].attrs = {'long_name': 'Longitude', 'units': 'degrees_east'} # Save to NetCDF with compression print(f" Writing to output file...") encoding = { 'tot_prec': {'zlib': True, 'complevel': 4, 'dtype': 'float32'}, 'tot_prec_rate': {'zlib': True, 'complevel': 4, 'dtype': 'float32'}, 'lat': {'zlib': True, 'complevel': 4, 'dtype': 'float32'}, 'lon': {'zlib': True, 'complevel': 4, 'dtype': 'float32'} } ds_out.to_netcdf(output_file, encoding=encoding) # Get file size file_size_mb = os.path.getsize(output_file) / (1024**2) print(f"\n{'='*70}") print(f"SUCCESS: File saved!") print(f" Output file: {output_file}") print(f" Time steps: {len(time_list)}") print(f" Grid: {len(lat_vals)} x {len(lon_vals)}") print(f" File size: {file_size_mb:.2f} MB") print(f"{'='*70}") ds_out.close() return True def main(): """Main function to process all ICON_KIT simulations.""" print("\n" + "="*70) print("ICON_KIT Precipitation Extraction Script") print("="*70) print(f"Start time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}") results = [] total_tasks = len(CASE_CONFIG) * len(SCENARIO_CONFIG) completed = 0 for case_name in CASE_CONFIG.keys(): for scenario_folder, output_experiment in SCENARIO_CONFIG.items(): success = extract_precipitation(case_name, scenario_folder, output_experiment) completed += 1 results.append({ 'case': case_name, 'scenario': output_experiment, 'success': success }) print(f"\nProgress: {completed}/{total_tasks} tasks completed\n") # Print final summary print("\n" + "="*70) print("FINAL SUMMARY") print("="*70) print(f"Total tasks: {total_tasks}") print(f"Successful: {sum(1 for r in results if r['success'])}") print(f"Failed: {sum(1 for r in results if not r['success'])}") print("\nDetailed results:") for result in results: status = "✓ SUCCESS" if result['success'] else "✗ FAILED" print(f" {status}: ICON_KIT - {result['case']} - {result['scenario']}") print(f"\nEnd time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}") print("="*70 + "\n") if __name__ == "__main__": main()