#!/usr/bin/env python3 import argparse import os import glob import warnings from datetime import datetime import matplotlib matplotlib.use('Agg') warnings.filterwarnings('ignore', category=FutureWarning) import numpy as np import pandas as pd import xarray as xr import matplotlib.pyplot as plt import matplotlib.colors as mcolors import cartopy.crs as ccrs import cartopy.feature as cfeature # ============================================================================ # PATH CONFIGURATION # ============================================================================ EXTRACTED_DATA_BASE = "/global/cfs/projectdirs/m4486/Haochen/Extracted_Data/Hourly_Precip" OUTPUT_DIR = "/global/cfs/projectdirs/m4486/Haochen/Precip/output_figures" # E3SM_THREAD paths - different for June vs August E3SM_THREAD_DIR_JUNE = os.path.join(EXTRACTED_DATA_BASE, 'E3SM_THREAD') # Original path for June E3SM_THREAD_DIR_AUG = os.path.join(EXTRACTED_DATA_BASE, 'E3SM_THREAD_New') # New path for August (extracted from ctrl_rainfrac2) # Model order for plotting (9 panels: MRMS + 8 models, WRFs grouped together) MODEL_ORDER = ['MRMS', 'WRF_SBM_ANL', 'WRF_SBM_TAMU', 'WRF_NU', 'WRF_DRI', 'RAMS', 'ICON_KIT', 'E3SM_THREAD', 'CM1'] # Model configurations (E3SM_THREAD will be handled specially based on case) MODEL_CONFIG = { 'MRMS': { 'extracted_dir': os.path.join(EXTRACTED_DATA_BASE, 'MRMS'), 'file_prefix': 'MRMS', 'label': 'MRMS', 'type': 'mrms' }, 'WRF_SBM_ANL': { 'extracted_dir': os.path.join(EXTRACTED_DATA_BASE, 'ANL_WRF'), 'file_prefix': 'ANL_WRF', 'label': 'WRF_SBM_ANL', 'type': 'wrf' }, 'WRF_SBM_TAMU': { 'extracted_dir': os.path.join(EXTRACTED_DATA_BASE, 'TAMU_WRF'), 'file_prefix': 'TAMU_WRF', 'label': 'WRF_SBM_TAMU', 'type': 'wrf' }, 'WRF_NU': { 'extracted_dir': os.path.join(EXTRACTED_DATA_BASE, 'NU_WRF'), 'file_prefix': 'NU_WRF', 'label': 'WRF_NU', 'type': 'wrf' }, 'WRF_DRI': { 'extracted_dir': os.path.join(EXTRACTED_DATA_BASE, 'WRF_DRI'), 'file_prefix': 'WRF_DRI', 'label': 'WRF_DRI', 'type': 'wrf' }, 'RAMS': { 'extracted_dir': os.path.join(EXTRACTED_DATA_BASE, 'RAMS'), 'file_prefix': 'RAMS', 'label': 'RAMS', 'type': 'rams' }, 'ICON_KIT': { 'extracted_dir': os.path.join(EXTRACTED_DATA_BASE, 'ICON_KIT'), 'file_prefix': 'ICON_KIT', 'label': 'ICON_KIT', 'type': 'icon' }, 'E3SM_THREAD': { # extracted_dir will be set dynamically based on case 'extracted_dir': None, 'file_prefix': 'E3SM_THREAD', 'label': 'E3SM_THREAD', 'type': 'e3sm' }, 'CM1': { 'extracted_dir': os.path.join(EXTRACTED_DATA_BASE, 'CM1'), 'file_prefix': 'CM1', 'label': 'CM1', 'type': 'cm1' } } # CM1 domain center for coordinate conversion CM1_CENTER_LON = -95.0792 CM1_CENTER_LAT = 29.4719 CASE_CONFIG = { 'june17': { 'start_datetime': datetime(2022, 6, 17, 12, 0), 'end_datetime': datetime(2022, 6, 18, 6, 0), 'analysis_box': {'lon_min': -96.25, 'lon_max': -95.35, 'lat_min': 29.6, 'lat_max': 30.2}, 'date_label': 'June 17, 2022', 'e3sm_dir': E3SM_THREAD_DIR_JUNE # Use original E3SM path }, 'aug7': { 'start_datetime': datetime(2022, 8, 7, 12, 0), 'end_datetime': datetime(2022, 8, 8, 6, 0), 'analysis_box': {'lon_min': -95.92, 'lon_max': -95.1, 'lat_min': 29.5, 'lat_max': 30.17}, 'date_label': 'August 7, 2022', 'e3sm_dir': E3SM_THREAD_DIR_AUG # Use new E3SM path } } # ============================================================================ # COLORMAP # ============================================================================ def create_precip_colormap(): p_clevs = [0.25, 1, 2, 3, 4, 6, 8, 10, 12, 14, 16, 20, 24] custom_colors = ['#2b0057', '#3b0f7e', '#0000cd', '#0066ff', '#00b3b3', '#00cc66', '#66ff66', '#ccff33', '#ffff00', '#ffcc00', '#ff6600', '#ff0000'] p_cmap = mcolors.ListedColormap(custom_colors) p_norm = mcolors.BoundaryNorm(p_clevs, len(custom_colors)) p_cmap.set_over('#ff66ff') p_cmap.set_under('white') return p_cmap, p_norm, p_clevs # ============================================================================ # COORDINATE CONVERSION # ============================================================================ def km_to_latlon(xh_km, yh_km, center_lon=CM1_CENTER_LON, center_lat=CM1_CENTER_LAT): """Convert CM1 km coordinates to lat/lon.""" km_per_deg_lat = 111.0 km_per_deg_lon = 111.0 * np.cos(np.radians(center_lat)) x_center_km = (xh_km.min() + xh_km.max()) / 2.0 y_center_km = (yh_km.min() + yh_km.max()) / 2.0 lon_1d = center_lon + (xh_km - x_center_km) / km_per_deg_lon lat_1d = center_lat + (yh_km - y_center_km) / km_per_deg_lat lon_2d, lat_2d = np.meshgrid(lon_1d, lat_1d) return lon_2d, lat_2d # ============================================================================ # DATA READING FUNCTIONS # ============================================================================ def get_e3sm_dir(case_name): """Get the appropriate E3SM_THREAD directory based on case.""" return CASE_CONFIG[case_name]['e3sm_dir'] def read_mrms_data(case_name): """Read extracted MRMS precipitation data.""" config = CASE_CONFIG[case_name] # Use extracted MRMS file mrms_file = os.path.join(EXTRACTED_DATA_BASE, 'MRMS', f"MRMS_{case_name}_10min.nc") if not os.path.exists(mrms_file): print(f"Warning: MRMS file not found: {mrms_file}") return None print(f"Reading: {os.path.basename(mrms_file)}") ds = xr.open_dataset(mrms_file) # Filter by time range start_dt = np.datetime64(config['start_datetime']) end_dt = np.datetime64(config['end_datetime']) # Check for time dimension name if 'time' in ds.dims: ds = ds.sel(time=slice(start_dt, end_dt)) elif 'Time' in ds.dims: ds = ds.sel(Time=slice(start_dt, end_dt)) return ds def read_model_data(model_name, case_name, scenario='control'): """Read model data based on model type.""" config = CASE_CONFIG[case_name] model_config = MODEL_CONFIG[model_name].copy() # Special handling for MRMS if model_config['type'] == 'mrms': return read_mrms_data(case_name) # Special handling for E3SM_THREAD - use case-specific directory if model_name == 'E3SM_THREAD': model_config['extracted_dir'] = get_e3sm_dir(case_name) print(f" Using E3SM_THREAD directory: {model_config['extracted_dir']}") # Construct file path extracted_file = os.path.join( model_config['extracted_dir'], f"{model_config['file_prefix']}_{case_name}_{scenario}.nc" ) if not os.path.exists(extracted_file): print(f"Warning: File not found: {extracted_file}") return None print(f"Reading: {os.path.basename(extracted_file)}") ds = xr.open_dataset(extracted_file) # Filter by time range start_dt = np.datetime64(config['start_datetime']) end_dt = np.datetime64(config['end_datetime']) model_type = model_config['type'] if model_type == 'wrf': ds = ds.sel(Time=slice(start_dt, end_dt)) # Already has XLAT, XLONG, RAINNC elif model_type == 'rams': ds = ds.sel(time=slice(start_dt, end_dt)) # RAMS has 2D lat/lon and tot_prec if 'y' in ds.dims and 'x' in ds.dims: ds['XLONG'] = ds['lon'] ds['XLAT'] = ds['lat'] elif 'lat' in ds.dims and 'lon' in ds.dims: if ds['lat'].ndim == 1: lon_2d, lat_2d = np.meshgrid(ds['lon'].values, ds['lat'].values) ds['XLONG'] = xr.DataArray(lon_2d, dims=['lat', 'lon']) ds['XLAT'] = xr.DataArray(lat_2d, dims=['lat', 'lon']) else: ds['XLONG'] = ds['lon'] ds['XLAT'] = ds['lat'] ds['RAINNC'] = ds['tot_prec'] elif model_type == 'icon': ds = ds.sel(time=slice(start_dt, end_dt)) lon_2d, lat_2d = np.meshgrid(ds['lon'].values, ds['lat'].values) ds['XLONG'] = xr.DataArray(lon_2d, dims=['lat', 'lon']) ds['XLAT'] = xr.DataArray(lat_2d, dims=['lat', 'lon']) ds['RAINNC'] = ds['tot_prec'] ds = ds.rename({'time': 'Time'}) elif model_type == 'e3sm': ds = ds.sel(time=slice(start_dt, end_dt)) lon_2d, lat_2d = np.meshgrid(ds['lon'].values, ds['lat'].values) ds['XLONG'] = xr.DataArray(lon_2d, dims=['lat', 'lon']) ds['XLAT'] = xr.DataArray(lat_2d, dims=['lat', 'lon']) ds['RAINNC'] = ds['precip_accum'] ds = ds.rename({'time': 'Time'}) elif model_type == 'cm1': ds = ds.sel(time=slice(start_dt, end_dt)) lon_2d, lat_2d = km_to_latlon(ds['xh'].values, ds['yh'].values) ds['XLONG'] = xr.DataArray(lon_2d, dims=['yh', 'xh']) ds['XLAT'] = xr.DataArray(lat_2d, dims=['yh', 'xh']) ds['RAINNC'] = ds['rain'] ds = ds.rename({'time': 'Time'}) return ds def compute_accumulated_precip(ds, model_type): """Compute accumulated precipitation from dataset.""" if model_type == 'mrms': # Extracted MRMS data: precip_rate in mm/hr at 10-min intervals # Accumulation = sum(rate) * time_interval_in_hours # 10 min = 10/60 = 1/6 hour precip_rate = ds['precip_rate'] precip_rate = precip_rate.where(precip_rate > 0, 0) # Set negative/nan to 0 time_dim = 'time' if 'time' in ds.dims else 'Time' # Sum rates and multiply by time interval (10 min = 1/6 hour) p_acc = (precip_rate.sum(dim=time_dim, skipna=True) * (10.0 / 60.0)).values # Apply threshold p_acc = np.where(p_acc > 0.25, p_acc, 0) # Get coordinates lons = ds['lon'].values lats = ds['lat'].values if lons.ndim == 1: lons, lats = np.meshgrid(lons, lats) return p_acc, lons, lats else: # Model: difference between last and first time step rainnc = ds['RAINNC'].values p_acc = rainnc[-1, :, :] - rainnc[0, :, :] p_acc = np.where(p_acc > 0.25, p_acc, 0) lons = ds['XLONG'].values lats = ds['XLAT'].values if lons.ndim == 1: lons, lats = np.meshgrid(lons, lats) return p_acc, lons, lats # ============================================================================ # PLOTTING FUNCTION # ============================================================================ def plot_multipanel_control(case_name): """Create 8-panel figure comparing MRMS with all model controls.""" config = CASE_CONFIG[case_name] print("\n" + "="*60) print(f"Processing: {config['date_label']} - MRMS vs All Models Control") print(f"E3SM_THREAD directory: {config['e3sm_dir']}") print("="*60) # Create colormap p_cmap, p_norm, p_clevs = create_precip_colormap() # Load Harris County boundary gdf = None shapefile_path = "/pscratch/sd/d/dbrooks/tracer_houston_case/HCAD_Harris_County_Boundary/HCAD_Harris_County_Boundary.shp" try: import geopandas as gpd if os.path.exists(shapefile_path): gdf = gpd.read_file(shapefile_path) gdf = gdf.to_crs(epsg=4326) print("Harris County boundary loaded.") except: pass # Create figure: 3 rows x 3 columns (MRMS + 8 models = 9 panels) fig, axes = plt.subplots(3, 3, figsize=(20, 16), subplot_kw={'projection': ccrs.PlateCarree()}) # Adjust spacing: minimal horizontal gaps; extra room above for suptitle + panel titles plt.subplots_adjust(wspace=0.02, hspace=0.10, left=0.06, right=0.98, top=0.91, bottom=0.08) axes = axes.flatten() # Get reference extent from first available model ref_extent = None # Load all data first to determine common extent all_data = {} for model_name in MODEL_ORDER: model_config = MODEL_CONFIG[model_name] ds = read_model_data(model_name, case_name, 'control') if ds is not None: p_acc, lons, lats = compute_accumulated_precip(ds, model_config['type']) all_data[model_name] = {'p_acc': p_acc, 'lons': lons, 'lats': lats, 'ds': ds} if ref_extent is None and model_name != 'MRMS': ref_extent = [lons.min(), lons.max(), lats.min(), lats.max()] # Plot each panel for idx, model_name in enumerate(MODEL_ORDER): ax = axes[idx] model_config = MODEL_CONFIG[model_name] if model_name not in all_data: ax.set_visible(False) continue data = all_data[model_name] p_acc, lons, lats = data['p_acc'], data['lons'], data['lats'] # Set extent (use reference extent for MRMS, own extent for models) if model_name == 'MRMS' and ref_extent is not None: ax.set_extent(ref_extent) # Subset MRMS to reference extent ds = data['ds'] lon_mask = (ds['lon'] >= ref_extent[0]) & (ds['lon'] <= ref_extent[1]) lat_mask = (ds['lat'] >= ref_extent[2]) & (ds['lat'] <= ref_extent[3]) ds_sub = ds.where(lon_mask & lat_mask, drop=True) # precip_rate in mm/hr, 10-min intervals precip_rate = ds_sub['precip_rate'].where(ds_sub['precip_rate'] > 0, 0) p_acc = (precip_rate.sum(dim='time', skipna=True) * (10.0 / 60.0)).values p_acc = np.where(p_acc > 0.25, p_acc, 0) lons = ds_sub['lon'].values lats = ds_sub['lat'].values if lons.ndim == 1: lons, lats = np.meshgrid(lons, lats) else: ax.set_extent([lons.min(), lons.max(), lats.min(), lats.max()]) # Plot precipitation pb = ax.pcolormesh(lons, lats, p_acc, cmap=p_cmap, norm=p_norm, transform=ccrs.PlateCarree()) # Add features ax.add_feature(cfeature.STATES, edgecolor="black", linewidths=0.5, alpha=0.6) ax.add_feature(cfeature.COASTLINE, edgecolor="black", linewidths=0.5, alpha=0.6) if gdf is not None: gdf.plot(ax=ax, facecolor='none', edgecolor='black', linestyle='-', linewidth=1, zorder=3) # Gridlines - control label visibility gl = ax.gridlines(draw_labels=True, linewidth=0.5, color='gray', alpha=0.25, linestyle='--') gl.top_labels = False gl.right_labels = False # Only show y-labels on leftmost column (idx 0, 3, 6) for 3x3 grid if idx not in [0, 3, 6]: gl.left_labels = False # Only show x-labels on bottom row (idx 6, 7, 8) for 3x3 grid if idx < 6: gl.bottom_labels = False gl.xlabel_style = {'size': 16} gl.ylabel_style = {'size': 16} # Title ax.set_title(f'{model_config["label"]}', fontsize=20, weight='bold') # Add common colorbar cbar_ax = fig.add_axes([0.15, 0.03, 0.7, 0.02]) cbar = fig.colorbar(plt.cm.ScalarMappable(norm=p_norm, cmap=p_cmap), cax=cbar_ax, orientation='horizontal', extend='both') cbar.set_label('Accumulated Precipitation (mm)', fontsize=18) cbar.set_ticks(p_clevs) cbar.set_ticklabels(['0.25' if t == 0.25 else str(int(t)) for t in p_clevs]) cbar.ax.tick_params(labelsize=16) # Add main title fig.suptitle(f'Total Accumulated Precipitation: {config["date_label"]} (12Z - 06Z +1 day)', fontsize=22, weight='bold', y=0.96) # Save figure case_output_dir = os.path.join(OUTPUT_DIR, case_name) os.makedirs(case_output_dir, exist_ok=True) save_path = os.path.join(case_output_dir, 'MRMS_vs_AllModels_Control.png') plt.savefig(save_path, dpi=150, bbox_inches='tight') print(f"\nSaved: {save_path}") plt.close() # ============================================================================ # MAIN # ============================================================================ if __name__ == '__main__': parser = argparse.ArgumentParser(description='Plot MRMS vs All Models Control') parser.add_argument('--case', choices=['june17', 'aug7', 'both'], default='both', help='Case to process (default: both)') args = parser.parse_args() cases = ['june17', 'aug7'] if args.case == 'both' else [args.case] for case in cases: plot_multipanel_control(case) print("\n" + "="*60) print("All processing complete!") print("="*60)