#!/usr/bin/env python3 import os import glob import warnings import numpy as np import h5py from datetime import datetime, timedelta import re warnings.filterwarnings('ignore', category=FutureWarning) # ============================================================================ # PATH CONFIGURATION # ============================================================================ RAMS_BASE_PATH = "/global/cfs/projectdirs/m4486/Haochen/TRACER_Model/RAMS" OUTPUT_BASE = "/global/cfs/projectdirs/m4486/Haochen/Extracted_Data/Hourly_Precip/RAMS" # Scenario folder names and output names # preconv = CNTL, 3xMore = high_ccn, 3xLess = low_ccn SCENARIO_CONFIG = { 'preconv': 'control', '3xMore': 'high_ccn', '3xLess': 'low_ccn' } CASE_CONFIG = { 'june17': { 'folder_prefix': 'a.jun17', 'folder_suffix': '500m', 'start_date': '2022-06-17', }, 'aug7': { 'folder_prefix': 'a.aug07', 'folder_suffix': '500m', 'start_date': '2022-08-07', } } def parse_rams_filename(filename): """Parse datetime from RAMS filename. Example: a-A-2022-06-17-120000-g2.h5 Returns: datetime(2022, 6, 17, 12, 0, 0) """ basename = os.path.basename(filename) # Extract datetime: a-A-YYYY-MM-DD-HHMMSS-g2.h5 match = re.search(r'a-A-(\d{4})-(\d{2})-(\d{2})-(\d{6})-g2\.h5', basename) if match: year = int(match.group(1)) month = int(match.group(2)) day = int(match.group(3)) time_str = match.group(4) hour = int(time_str[0:2]) minute = int(time_str[2:4]) second = int(time_str[4:6]) return datetime(year, month, day, hour, minute, second) return None def extract_precipitation(case_name, scenario_key, output_experiment): """ Extract ACCPR from RAMS files. Parameters: ----------- case_name : str Case name (june17 or aug7) scenario_key : str RAMS folder key (preconv, 3xMore, 3xLess) output_experiment : str Output experiment name (control, high_ccn, low_ccn) """ config = CASE_CONFIG[case_name] print(f"\n{'='*70}") print(f"Processing: RAMS - {case_name} - {scenario_key} -> {output_experiment}") print(f"{'='*70}") # Build folder path # Pattern: a.jun17.preconv.500m or a.aug07.3xMore.500m folder_name = f"{config['folder_prefix']}.{scenario_key}.{config['folder_suffix']}" case_path = os.path.join(RAMS_BASE_PATH, folder_name) # Output file os.makedirs(OUTPUT_BASE, exist_ok=True) output_file = os.path.join(OUTPUT_BASE, f"RAMS_{case_name}_{output_experiment}.nc") print(f"Input path: {case_path}") print(f"Output file: {output_file}") # Check if input path exists if not os.path.exists(case_path): print(f"ERROR: Input path does not exist: {case_path}") return False # 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 a-A-*-g2.h5 files (10-minute interval analysis files) file_pattern = os.path.join(case_path, "a-A-*-g2.h5") 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)} a-A files") # Extract data from each file accpr_list = [] time_list = [] lat_vals = None lon_vals = None for i, fpath in enumerate(files): if i % 20 == 0 or i == 0: print(f" Reading file {i+1}/{len(files)}: {os.path.basename(fpath)}") try: with h5py.File(fpath, 'r') as f: # Check if ACCPR exists if 'ACCPR' not in f: print(f" Warning: ACCPR not found in {fpath}") continue # Extract accumulated precipitation (2D: ny, nx) accpr = f['ACCPR'][:] accpr_list.append(accpr) # Get time from filename file_dt = parse_rams_filename(fpath) if file_dt: time_list.append(file_dt) # Get coordinates (only once) if lat_vals is None: lat_vals = f['GLAT'][:] lon_vals = f['GLON'][:] print(f" Grid shape: {lat_vals.shape}") print(f" Lat range: {lat_vals.min():.2f} to {lat_vals.max():.2f}") print(f" Lon range: {lon_vals.min():.2f} to {lon_vals.max():.2f}") except Exception as e: print(f" ERROR reading {fpath}: {e}") continue if not accpr_list: print("ERROR: No data extracted") return False # Stack data print(" Combining data...") accpr_arr = np.stack(accpr_list, axis=0).astype(np.float32) print(f" Total timesteps: {len(time_list)}") print(f" Time range: {time_list[0]} to {time_list[-1]}") print(f" Data shape: {accpr_arr.shape}") print(f" ACCPR range: {np.nanmin(accpr_arr):.2f} to {np.nanmax(accpr_arr):.2f} mm") # Compute precipitation rate from accumulated (difference between timesteps) # Rate = delta_accum / delta_time (10 minutes = 600 seconds) print(" Computing precipitation rate...") precip_rate = np.zeros_like(accpr_arr) precip_rate[0, :, :] = 0 # First timestep has no rate precip_rate[1:, :, :] = np.diff(accpr_arr, axis=0) / 600.0 # mm/600s -> mm/s -> kg/m²/s # Ensure non-negative (accumulated should only increase) precip_rate = np.maximum(precip_rate, 0).astype(np.float32) # Create output dataset using xarray import xarray as xr import pandas as pd print(f" Writing to output file...") # Convert datetime to pandas DatetimeIndex for better compatibility time_pd = pd.DatetimeIndex(time_list) # Create xarray dataset ds_out = xr.Dataset({ 'tot_prec': (['time', 'y', 'x'], accpr_arr), 'tot_prec_rate': (['time', 'y', 'x'], precip_rate), 'lat': (['y', 'x'], lat_vals.astype(np.float32)), 'lon': (['y', 'x'], lon_vals.astype(np.float32)), }, coords={ 'time': time_pd, }) # Add variable attributes ds_out['tot_prec'].attrs = { 'long_name': 'Total Accumulated Rain Precipitation', 'units': 'mm', 'description': 'Accumulated rain precipitation (ACCPR) from RAMS', 'original_name': 'ACCPR' } ds_out['tot_prec_rate'].attrs = { 'long_name': 'Total Precipitation Rate', 'units': 'kg m-2 s-1', 'description': 'Precipitation rate computed from accumulated (ACCPR diff / dt)' } ds_out['lat'].attrs = {'long_name': 'Latitude', 'units': 'degrees_north'} ds_out['lon'].attrs = {'long_name': 'Longitude', 'units': 'degrees_east'} # Add global attributes ds_out.attrs = { 'title': 'RAMS Precipitation Data', 'case': case_name, 'scenario': output_experiment, 'description': 'Extracted precipitation from RAMS 10-min output (a-A files)', 'creation_date': datetime.now().strftime('%Y-%m-%d %H:%M:%S'), 'creator': 'Haochen Zhang', 'model': 'RAMS (CSU Regional Atmospheric Modeling System)', 'n_files_processed': str(len(files)), 'source_variable': 'ACCPR (accumulated rain precipitation)', 'grid_type': 'curvilinear (2D lat/lon)' } # Save using scipy engine (NetCDF3 format) to avoid HDF5 issues # Note: scipy doesn't support compression but is more reliable print(f" Writing to output file (NetCDF3 format)...") ds_out.to_netcdf(output_file, engine='scipy', format='NETCDF3_64BIT') ds_out.close() # 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: {lat_vals.shape[0]} x {lat_vals.shape[1]}") print(f" File size: {file_size_mb:.2f} MB") print(f"{'='*70}") return True def main(): """Main function to process all RAMS simulations.""" print("\n" + "="*70) print("RAMS Precipitation Extraction Script") print("="*70) print(f"Start time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}") print("\nScenario mapping:") print(" preconv -> control (CNTL)") print(" 3xMore -> high_ccn") print(" 3xLess -> low_ccn") results = [] total_tasks = len(CASE_CONFIG) * len(SCENARIO_CONFIG) completed = 0 for case_name in CASE_CONFIG.keys(): for scenario_key, output_experiment in SCENARIO_CONFIG.items(): success = extract_precipitation(case_name, scenario_key, 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}: RAMS - {result['case']} - {result['scenario']}") print(f"\nEnd time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}") print("="*70 + "\n") if __name__ == "__main__": main()