#!/usr/bin/env python3 import argparse import os import glob import warnings from datetime import datetime, timedelta import matplotlib matplotlib.use('Agg') warnings.filterwarnings('ignore', category=FutureWarning) warnings.filterwarnings('ignore', category=RuntimeWarning) import numpy as np import pandas as pd import xarray as xr import h5py import matplotlib.pyplot as plt import matplotlib.colors as mcolors import cartopy.crs as ccrs import cartopy.feature as cfeature # ============================================================================ # PATH CONFIGURATION # ============================================================================ BASE_PATH = "/global/cfs/projectdirs/m4486/Haochen" REFLECTIVITY_BASE = os.path.join(BASE_PATH, "Reflectivity") DOMENIC_BASE = os.path.join(BASE_PATH, "Domenic_Files/post_processing") OUTPUT_DIR = os.path.join(REFLECTIVITY_BASE, "output_figures") # Model order for plotting (9 panels: NEXRAD + 8 models, WRFs grouped together) MODEL_ORDER = ['NEXRAD', 'WRF_SBM_ANL', 'WRF_SBM_TAMU', 'WRF_NU', 'WRF_DRI', 'RAMS', 'ICON_KIT', 'E3SM_THREAD', 'CM1'] # Case configuration CASE_CONFIG = { 'june17': { 'start_datetime': datetime(2022, 6, 17, 12, 0), 'end_datetime': datetime(2022, 6, 18, 6, 0), 'analysis_box': {'lon_min': -96.10, 'lon_max': -94.0, 'lat_min': 28.5, 'lat_max': 30.6}, 'date_label': 'June 17, 2022' }, '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' } } # Model-specific configurations MODEL_CONFIG = { 'NEXRAD': { 'label': 'NEXRAD', 'type': 'nexrad', 'june17': { 'data_dir': os.path.join(DOMENIC_BASE, 'june17/nexrad_data/20220617-18_KHGX_3D_data'), }, 'aug7': { 'data_dir': os.path.join(DOMENIC_BASE, 'aug7/nexrad_data/20220807-08_KHGX_3D_data'), } }, 'WRF_SBM_ANL': { 'label': 'WRF_SBM_ANL', 'type': 'wrf_anl', 'june17': { 'data_file': os.path.join(DOMENIC_BASE, 'june17/reflectivity_data/cntl_aero.nc'), }, 'aug7': { 'data_file': os.path.join(DOMENIC_BASE, 'aug7/reflectivity_data/cntl_aero.nc'), } }, 'WRF_SBM_TAMU': { 'label': 'WRF_SBM_TAMU', 'type': 'composite_nc', 'june17': { 'data_file': os.path.join(REFLECTIVITY_BASE, 'TAMU_WRF_Ref/june17/tamu_wrf_composite_june17_control.nc'), }, 'aug7': { 'data_file': os.path.join(REFLECTIVITY_BASE, 'TAMU_WRF_Ref/aug7/tamu_wrf_composite_aug7_control.nc'), } }, 'WRF_NU': { 'label': 'WRF_NU', 'type': 'composite_nc', 'june17': { 'data_file': os.path.join(REFLECTIVITY_BASE, 'NU_WRF_Ref/june17/nu_wrf_composite_june17_control.nc'), }, 'aug7': { 'data_file': os.path.join(REFLECTIVITY_BASE, 'NU_WRF_Ref/aug7/nu_wrf_composite_aug7_control.nc'), } }, 'WRF_DRI': { 'label': 'WRF_DRI', 'type': 'composite_nc', 'june17': { 'data_file': os.path.join(REFLECTIVITY_BASE, 'WRF_DRI_Ref/june17/wrf_dri_composite_june17_control.nc'), }, 'aug7': { 'data_file': os.path.join(REFLECTIVITY_BASE, 'WRF_DRI_Ref/aug7/wrf_dri_composite_aug7_control.nc'), } }, 'RAMS': { 'label': 'RAMS', 'type': 'rams_h5', 'june17': { 'data_file': os.path.join(REFLECTIVITY_BASE, 'RAMS_Ref/june17/rams_composite_june17_control.h5'), }, 'aug7': { 'data_file': os.path.join(REFLECTIVITY_BASE, 'RAMS_Ref/aug7/rams_composite_aug7_control.h5'), } }, 'ICON_KIT': { 'label': 'ICON_KIT', 'type': 'composite_nc', 'june17': { 'data_file': os.path.join(REFLECTIVITY_BASE, 'ICON_KIT_Ref/june17/iconkit_composite_june17_control.nc'), }, 'aug7': { 'data_file': os.path.join(REFLECTIVITY_BASE, 'ICON_KIT_Ref/aug7/iconkit_composite_aug7_control.nc'), } }, 'E3SM_THREAD': { 'label': 'E3SM_THREAD', 'type': 'composite_nc', 'june17': { 'data_file': os.path.join(REFLECTIVITY_BASE, 'E3SM_THREAD_Ref/june17/e3sm_composite_june17_control.nc'), }, 'aug7': { 'data_file': os.path.join(REFLECTIVITY_BASE, 'E3SM_THREAD_Ref/aug7/e3sm_composite_aug7_control.nc'), } }, 'CM1': { 'label': 'CM1', 'type': 'composite_nc', 'june17': { 'data_file': os.path.join(REFLECTIVITY_BASE, 'CM1_Ref/june17/cm1_composite_june17_control.nc'), }, 'aug7': { 'data_file': os.path.join(REFLECTIVITY_BASE, 'CM1_Ref/aug7/cm1_composite_aug7_control.nc'), } } } # CM1 domain center for coordinate conversion CM1_CENTER_LON = -95.0792 CM1_CENTER_LAT = 29.4719 # ============================================================================ # COLORMAP FOR REFLECTIVITY # ============================================================================ def create_reflectivity_colormap(): """Create NWS-style reflectivity colormap.""" ref_clevs = [5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55] custom_colors = [ '#9C9C9C', # 5 - light gray '#767676', # 10 - gray '#00FFFF', # 15 - cyan '#00BFFF', # 20 - deep sky blue '#0000FF', # 25 - blue '#00FF00', # 30 - green '#00C800', # 35 - darker green '#FFFF00', # 40 - yellow '#FFA500', # 45 - orange '#FF0000', # 50 - red ] ref_cmap = mcolors.ListedColormap(custom_colors) ref_norm = mcolors.BoundaryNorm(ref_clevs, len(custom_colors)) ref_cmap.set_over('#FF00FF') # magenta for >55 ref_cmap.set_under('white') # white for <5 return ref_cmap, ref_norm, ref_clevs # ============================================================================ # HELPER FUNCTIONS # ============================================================================ def open_netcdf_with_h5py(filepath): """ Read NetCDF-4 file using h5py (since NetCDF-4 is HDF5 format). Returns a dictionary with variables as numpy arrays. """ try: hf = h5py.File(filepath, 'r') print(f" (opened with h5py)") return hf except Exception as e: print(f" Failed to open with h5py: {e}") return None def get_h5_variable(hf, var_names): """Get variable from h5py file, trying multiple possible names.""" if isinstance(var_names, str): var_names = [var_names] for name in var_names: if name in hf: return hf[name][:] return None def get_h5_attr_or_default(hf, var_name, attr_name, default=None): """Get attribute from h5py variable, with default fallback.""" try: if var_name in hf and attr_name in hf[var_name].attrs: return hf[var_name].attrs[attr_name] except: pass return default def select_region_2d(data, lat, lon, box_config): """Select data within a lat/lon box.""" lon_min = box_config['lon_min'] lon_max = box_config['lon_max'] lat_min = box_config['lat_min'] lat_max = box_config['lat_max'] if lat.ndim == 2: mask = (lat >= lat_min) & (lat <= lat_max) & (lon >= lon_min) & (lon <= lon_max) return np.where(mask, data, np.nan) else: # 1D lat/lon arrays - create 2D mask lon_2d, lat_2d = np.meshgrid(lon, lat) mask = (lat_2d >= lat_min) & (lat_2d <= lat_max) & (lon_2d >= lon_min) & (lon_2d <= lon_max) return np.where(mask, data, np.nan) def hours_to_datetime(hours, ref_datetime=datetime(2022, 1, 1, 0, 0, 0)): """Convert hours since reference to datetime.""" return ref_datetime + timedelta(hours=float(hours)) 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 # ============================================================================ # NEXRAD DATA FUNCTIONS # ============================================================================ def parse_nexrad_filename(filepath): """Extract datetime from NEXRAD filename.""" basename = os.path.basename(filepath) try: date_str = basename[4:12] time_str = basename[13:19] dt = datetime.strptime(date_str + time_str, "%Y%m%d%H%M%S") return dt except: return None def read_nexrad_with_coords(filepath): """Read NEXRAD file and get composite reflectivity using h5py.""" try: hf = h5py.File(filepath, 'r') # Get REF variable # REF is stored as ubyte with scale_factor=0.5, add_offset=-32.5 ref_raw = hf['REF'][:].astype(float) ref_raw[ref_raw == 255] = np.nan # Fill value # Apply scale and offset manually (h5py doesn't auto-apply these) # Check if scale_factor/add_offset exist as attributes if 'scale_factor' in hf['REF'].attrs: scale = hf['REF'].attrs['scale_factor'] offset = hf['REF'].attrs['add_offset'] ref_3d = ref_raw * scale + offset else: # Default WRF NEXRAD scaling ref_3d = ref_raw * 0.5 + (-32.5) # Composite: max over z dimension (axis=0) ref = np.nanmax(ref_3d, axis=0) # Get radar location if 'latitude' in hf: radar_lat = float(np.array(hf['latitude']).flatten()[0]) else: radar_lat = 29.4719 # KHGX default if 'longitude' in hf: radar_lon = float(np.array(hf['longitude']).flatten()[0]) else: radar_lon = -95.0792 # KHGX default hf.close() # Get grid dimensions ny, nx = ref.shape # The NEXRAD data is gridded to MAAS domain # 500 points, 0.5 km spacing = 250 km total domain dx_km = 0.5 # Create coordinate arrays (km from radar center) x_km = (np.arange(nx) - nx/2) * dx_km y_km = (np.arange(ny) - ny/2) * dx_km # Convert to lat/lon lat_1d = radar_lat + y_km / 111.0 lon_1d = radar_lon + x_km / (111.0 * np.cos(np.radians(radar_lat))) # Create 2D meshgrid lon, lat = np.meshgrid(lon_1d, lat_1d) return ref, lat, lon except Exception as e: print(f" Error reading NEXRAD file: {e}") return None, None, None def read_nexrad_data(case_name, config): """Read NEXRAD data and find max reflectivity time.""" model_config = MODEL_CONFIG['NEXRAD'][case_name] nexrad_path = os.path.join(model_config['data_dir'], 'KHGX*.nc') file_list = sorted(glob.glob(nexrad_path)) if not file_list: print(f" NEXRAD: No files found") return None print(f" NEXRAD: Found {len(file_list)} files") # Filter files within analysis time window case_config = CASE_CONFIG[case_name] valid_files = [] for f in file_list: dt = parse_nexrad_filename(f) if dt and case_config['start_datetime'] <= dt <= case_config['end_datetime']: valid_files.append((f, dt)) if not valid_files: print(f" NEXRAD: No files in time window ({case_config['start_datetime']} to {case_config['end_datetime']})") return None print(f" NEXRAD: {len(valid_files)} files in time window") # Find max reflectivity time box_config = case_config['analysis_box'] max_val = -999 max_file = None max_time = None files_read = 0 files_failed = 0 for filepath, dt in valid_files: result = read_nexrad_with_coords(filepath) if result[0] is None: files_failed += 1 continue ref, lat, lon = result files_read += 1 ref_in_box = select_region_2d(ref, lat, lon, box_config) valid_data = ref_in_box[~np.isnan(ref_in_box)] if len(valid_data) > 0: current_max = float(np.max(valid_data)) if current_max > max_val: max_val = current_max max_file = filepath max_time = dt print(f" NEXRAD: Read {files_read} files, {files_failed} failed") if max_file: ref, lat, lon = read_nexrad_with_coords(max_file) print(f" NEXRAD: Max {max_val:.1f} dBZ at {max_time.strftime('%H:%M')}") return {'data': ref, 'lat': lat, 'lon': lon, 'time': max_time} else: print(f" NEXRAD: No valid data found") return None # ============================================================================ # WRF ANL DATA FUNCTIONS (uses mdbz variable directly) # ============================================================================ def read_wrf_anl_data(case_name, config): """Read WRF ANL data (mdbz variable) and find max reflectivity time.""" model_config = MODEL_CONFIG['WRF_SBM_ANL'][case_name] filepath = model_config['data_file'] if not os.path.exists(filepath): print(f" WRF_SBM_ANL: File not found: {filepath}") return None print(f" WRF_SBM_ANL: Reading {os.path.basename(filepath)}") try: hf = open_netcdf_with_h5py(filepath) if hf is None: return None # Get mdbz variable (max dBZ - already composite) if 'mdbz' not in hf: print(f" WRF_SBM_ANL: mdbz variable not found") print(f" Available variables: {list(hf.keys())}") hf.close() return None mdbz = hf['mdbz'][:] # Shape: (Time, south_north, west_east) # Get coordinates xlat = get_h5_variable(hf, ['XLAT', 'lat']) xlong = get_h5_variable(hf, ['XLONG', 'lon']) if xlat is None or xlong is None: print(f" WRF_SBM_ANL: Coordinates not found") hf.close() return None # Handle 3D coordinates (Time, south_north, west_east) if xlat.ndim == 3: lat_arr = xlat[0, :, :] lon_arr = xlong[0, :, :] else: lat_arr = xlat lon_arr = xlong # Get XTIME for timestamps xtime = get_h5_variable(hf, ['XTIME']) # Find max reflectivity time case_config = CASE_CONFIG[case_name] box_config = case_config['analysis_box'] max_val = -999 max_idx = 0 n_times = mdbz.shape[0] for t in range(n_times): ref_t = mdbz[t, :, :] ref_t = np.where(ref_t > 1e10, np.nan, ref_t) ref_in_box = select_region_2d(ref_t, lat_arr, lon_arr, box_config) valid_data = ref_in_box[~np.isnan(ref_in_box)] if len(valid_data) > 0: current_max = float(np.max(valid_data)) if current_max > max_val: max_val = current_max max_idx = t # Get data at max time ref_data = mdbz[max_idx, :, :] ref_data = np.where(ref_data > 1e10, np.nan, ref_data) # Get timestamp from XTIME if xtime is not None: xtime_val = xtime[max_idx] if xtime.ndim == 1 else xtime[max_idx, 0] max_time = case_config['start_datetime'] + pd.Timedelta(minutes=float(xtime_val)) else: max_time = case_config['start_datetime'] print(f" WRF_SBM_ANL: Max {max_val:.1f} dBZ at {max_time.strftime('%H:%M')}") hf.close() return {'data': ref_data, 'lat': lat_arr, 'lon': lon_arr, 'time': max_time} except Exception as e: print(f" WRF_SBM_ANL: Error - {e}") import traceback traceback.print_exc() return None # ============================================================================ # COMPOSITE NC DATA FUNCTIONS (for TAMU, NU, ICON_KIT, E3SM, CM1) # ============================================================================ def read_composite_nc_data(model_name, case_name, config): """ Read pre-extracted composite reflectivity NetCDF file. NOTE on engine choice: We use xarray.open_dataset, which honors _FillValue / missing_value attributes automatically (unlike raw h5py, which leaks fill values like -999 or 9.97e+36 into data and coords and crashes pcolormesh on non-finite XLAT/XLONG). On Perlmutter, the bundled netCDF4 library has occasionally rejected these files with `[Errno -101] NetCDF: HDF error` due to library/filter version mismatches. To be robust across environments, we try the netCDF4 engine first and fall back to h5netcdf (which reads via h5py but still gives us xarray's CF decoding and masking). """ import xarray as xr def _to_nan(da): """Get plain float64 numpy array from a DataArray, NaN at fills.""" arr = da.values if isinstance(arr, np.ma.MaskedArray): return arr.filled(np.nan).astype(np.float64) return np.asarray(arr, dtype=np.float64) def _open(filepath): """Try netCDF4 engine; if it fails, fall back to h5netcdf.""" last_err = None for engine in ('netcdf4', 'h5netcdf'): try: ds = xr.open_dataset( filepath, engine=engine, decode_times=False, # avoid CF time-decode surprises mask_and_scale=True, # honor _FillValue / scale_factor ) print(f" (opened with xarray + {engine})") return ds except Exception as e: last_err = e continue raise last_err model_config = MODEL_CONFIG[model_name][case_name] filepath = model_config['data_file'] if not os.path.exists(filepath): print(f" {model_name}: File not found: {filepath}") return None print(f" {model_name}: Reading {os.path.basename(filepath)}") try: ds = _open(filepath) # --- composite reflectivity ----------------------------------------- ref_da = None for vname in ['COMPOSITE_REFL_10CM', 'COMPOSITE_REFL']: if vname in ds.variables: ref_da = ds[vname] break if ref_da is None: print(f" {model_name}: No composite reflectivity variable found") print(f" Available variables: {list(ds.variables.keys())}") ds.close() return None composite_ref = _to_nan(ref_da) # --- coordinates ----------------------------------------------------- xlat_da = None xlong_da = None for vname in ['XLAT', 'lat']: if vname in ds.variables: xlat_da = ds[vname] break for vname in ['XLONG', 'lon']: if vname in ds.variables: xlong_da = ds[vname] break if xlat_da is not None and xlong_da is not None: xlat = _to_nan(xlat_da) xlong = _to_nan(xlong_da) if xlat.ndim == 1 and xlong.ndim == 1: xlong_2d, xlat_2d = np.meshgrid(xlong, xlat) else: xlat_2d = xlat xlong_2d = xlong elif 'xh' in ds.variables and 'yh' in ds.variables: # CM1 uses km coordinates xh_km = _to_nan(ds['xh']) yh_km = _to_nan(ds['yh']) xlong_2d, xlat_2d = km_to_latlon(xh_km, yh_km) else: print(f" {model_name}: No coordinates found") ds.close() return None # If lat/lon arrays still contain NaN (corrupt extraction), # try a best-effort recovery so pcolormesh doesn't crash. if not (np.all(np.isfinite(xlat_2d)) and np.all(np.isfinite(xlong_2d))): n_bad = int(np.sum(~np.isfinite(xlat_2d)) + np.sum(~np.isfinite(xlong_2d))) print(f" {model_name}: WARNING — coordinate arrays contain {n_bad} non-finite " f"values; trying to fill from a clean row/column.") for r in range(xlat_2d.shape[0]): if np.all(np.isfinite(xlat_2d[r, :])) and np.all(np.isfinite(xlong_2d[r, :])): xlat_2d = np.broadcast_to(xlat_2d[r:r+1, :], xlat_2d.shape).copy() break for c in range(xlong_2d.shape[1]): if np.all(np.isfinite(xlong_2d[:, c])) and np.all(np.isfinite(xlat_2d[:, c])): xlong_2d = np.broadcast_to(xlong_2d[:, c:c+1], xlong_2d.shape).copy() break # --- time ------------------------------------------------------------ time_data = None for vname in ['time', 'Time', 'XTIME']: if vname in ds.variables: time_data = ds[vname].values break times_str = ds['Times'].values if 'Times' in ds.variables else None # --- find max reflectivity time ------------------------------------- case_config = CASE_CONFIG[case_name] box_config = case_config['analysis_box'] max_val = -999.0 max_idx = 0 n_times = composite_ref.shape[0] for t in range(n_times): ref_t = composite_ref[t, :, :] ref_in_box = select_region_2d(ref_t, xlat_2d, xlong_2d, box_config) valid_data = ref_in_box[~np.isnan(ref_in_box)] if len(valid_data) > 0: current_max = float(np.max(valid_data)) if current_max > max_val: max_val = current_max max_idx = t ref_data = composite_ref[max_idx, :, :] # --- timestamp ------------------------------------------------------- max_time = case_config['start_datetime'] # default if times_str is not None: try: time_bytes = times_str[max_idx] if isinstance(time_bytes, np.ndarray): time_str = time_bytes.tobytes().decode('utf-8').strip().replace('\x00', '') else: time_str = str(time_bytes).strip() max_time = datetime.strptime(time_str, '%Y-%m-%d_%H:%M:%S') except Exception: pass elif time_data is not None: try: time_val = time_data[max_idx] if isinstance(time_val, np.datetime64): ts = (time_val - np.datetime64('1970-01-01T00:00:00')) / np.timedelta64(1, 's') max_time = datetime.utcfromtimestamp(ts) elif isinstance(time_val, (float, np.floating, int, np.integer)): max_time = hours_to_datetime(float(time_val)) except Exception: pass print(f" {model_name}: Max {max_val:.1f} dBZ at {max_time.strftime('%H:%M')}") ds.close() return {'data': ref_data, 'lat': xlat_2d, 'lon': xlong_2d, 'time': max_time} except Exception as e: print(f" {model_name}: Error - {e}") import traceback traceback.print_exc() return None # ============================================================================ # RAMS DATA FUNCTIONS (HDF5 format) # ============================================================================ def read_rams_data(case_name, config): """Read RAMS composite data (HDF5 format) and find max reflectivity time.""" model_config = MODEL_CONFIG['RAMS'][case_name] filepath = model_config['data_file'] if not os.path.exists(filepath): print(f" RAMS: File not found: {filepath}") return None print(f" RAMS: Reading {os.path.basename(filepath)}") try: hf = h5py.File(filepath, 'r') composite_ref = hf['COMPOSITE_REFL'][:] time_data = hf['time'][:] xlat = hf['XLAT'][:] xlong = hf['XLONG'][:] # Find max reflectivity time case_config = CASE_CONFIG[case_name] box_config = case_config['analysis_box'] max_val = -999 max_idx = 0 for t in range(composite_ref.shape[0]): ref_t = composite_ref[t, :, :] ref_in_box = select_region_2d(ref_t, xlat, xlong, box_config) valid_data = ref_in_box[~np.isnan(ref_in_box)] if len(valid_data) > 0: current_max = float(np.max(valid_data)) if current_max > max_val: max_val = current_max max_idx = t ref_data = composite_ref[max_idx, :, :] max_time = hours_to_datetime(time_data[max_idx]) print(f" RAMS: Max {max_val:.1f} dBZ at {max_time.strftime('%H:%M')}") hf.close() return {'data': ref_data, 'lat': xlat, 'lon': xlong, 'time': max_time} except Exception as e: print(f" RAMS: Error - {e}") import traceback traceback.print_exc() return None # ============================================================================ # DATA LOADING DISPATCHER # ============================================================================ def load_model_data(model_name, case_name, config): """Load data for a specific model.""" model_type = MODEL_CONFIG[model_name]['type'] if model_type == 'nexrad': return read_nexrad_data(case_name, config) elif model_type == 'wrf_anl': return read_wrf_anl_data(case_name, config) elif model_type == 'rams_h5': return read_rams_data(case_name, config) elif model_type == 'composite_nc': return read_composite_nc_data(model_name, case_name, config) else: print(f" {model_name}: Unknown type {model_type}") return None # ============================================================================ # MULTI-PANEL PLOTTING # ============================================================================ def plot_multipanel(case_name, model_data, ref_extent, gdf=None): """Create 3x3 multi-panel figure of composite reflectivity (NEXRAD + 8 models).""" config = CASE_CONFIG[case_name] ref_cmap, ref_norm, ref_clevs = create_reflectivity_colormap() # Create figure: 3 rows x 3 columns (NEXRAD + 8 models = 9 panels) fig, axes = plt.subplots(3, 3, figsize=(20, 16), subplot_kw={'projection': ccrs.PlateCarree()}) axes = axes.flatten() # Plot domain from reference extent [lon_min, lon_max, lat_min, lat_max] plot_lon_min = ref_extent[0] plot_lon_max = ref_extent[1] plot_lat_min = ref_extent[2] plot_lat_max = ref_extent[3] # Store a valid pcolormesh for colorbar pb = None for idx, model_name in enumerate(MODEL_ORDER): ax = axes[idx] ax.set_extent([plot_lon_min, plot_lon_max, plot_lat_min, plot_lat_max]) data = model_data.get(model_name) if data is not None: # Filter data for plotting data_plot = np.where(data['data'] >= ref_clevs[0], data['data'], np.nan) # Plot reflectivity pb = ax.pcolormesh(data['lon'], data['lat'], data_plot, cmap=ref_cmap, norm=ref_norm, transform=ccrs.PlateCarree(), shading='auto') # Time label if data['time']: time_str = data['time'].strftime("%H:%Mz") else: time_str = "" else: # No data - show empty panel time_str = "No Data" # Add geographic features ax.add_feature(cfeature.STATES, edgecolor="black", linewidths=0.65, alpha=0.6) ax.add_feature(cfeature.COASTLINE, edgecolor="black", linewidths=0.65, alpha=0.6) # Add county boundary if available if gdf is not None: gdf.plot(ax=ax, facecolor='none', edgecolor='black', linestyle='-', linewidth=1.5, zorder=3) # Add gridlines gl = ax.gridlines(draw_labels=True, linewidth=0.5, color='gray', alpha=0.25, linestyle='--') gl.top_labels = False gl.right_labels = False gl.xlabel_style = {'size': 16} gl.ylabel_style = {'size': 16} # 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 # Title with model name and time model_label = MODEL_CONFIG[model_name]['label'] if time_str: ax.set_title(f'{model_label}\n{time_str}', fontsize=20, weight='bold') else: ax.set_title(f'{model_label}', fontsize=20, weight='bold') # Add common colorbar at bottom cbar_ax = fig.add_axes([0.15, 0.04, 0.7, 0.02]) cbar = fig.colorbar(plt.cm.ScalarMappable(norm=ref_norm, cmap=ref_cmap), cax=cbar_ax, orientation='horizontal', extend='both') cbar.set_label('Composite Reflectivity (dBZ)', fontsize=18) cbar.set_ticks(ref_clevs) cbar.ax.tick_params(labelsize=16) # Main title - position chosen to leave breathing room above panel titles fig.suptitle(f'Composite Reflectivity at Maximum Intensity: {config["date_label"]}', fontsize=22, weight='bold', y=0.96) plt.subplots_adjust(left=0.06, right=0.98, top=0.88, bottom=0.08, wspace=0.02, hspace=0.16) # 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, 'NEXRAD_vs_AllModels_Control.png') plt.savefig(save_path, dpi=150, bbox_inches='tight') print(f"\nSaved: {save_path}") plt.close() # ============================================================================ # MAIN PROCESSING # ============================================================================ def process_case(case_name): """Process a single case.""" config = CASE_CONFIG[case_name] print("\n" + "="*70) print(f"Processing: {config['date_label']}") print("="*70) # Load Harris County boundary shapefile 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 ImportError: print("Note: geopandas not available, skipping county boundary") except Exception as e: print(f"Note: Could not load shapefile: {e}") # Load data for all models print("\nLoading model data...") model_data = {} ref_extent = None for model_name in MODEL_ORDER: model_data[model_name] = load_model_data(model_name, case_name, config) # Get reference extent from first available non-NEXRAD model if model_data[model_name] is not None and ref_extent is None and model_name != 'NEXRAD': lons = model_data[model_name]['lon'] lats = model_data[model_name]['lat'] ref_extent = [lons.min(), lons.max(), lats.min(), lats.max()] print(f" Reference extent from {model_name}: lon [{ref_extent[0]:.2f}, {ref_extent[1]:.2f}], lat [{ref_extent[2]:.2f}, {ref_extent[3]:.2f}]") # Fallback if no model data available if ref_extent is None: ref_extent = [-96.10, -94.0, 28.5, 30.6] print(" Using default extent (no model data available)") # Create multi-panel plot print("\nCreating multi-panel figure...") plot_multipanel(case_name, model_data, ref_extent, gdf) print(f"\nCase {case_name} complete!") # ============================================================================ # ENTRY POINT # ============================================================================ if __name__ == '__main__': parser = argparse.ArgumentParser( description='Plot NEXRAD vs All Models composite reflectivity comparison (2x4 panels)' ) 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: process_case(case) print("\n" + "="*70) print("All processing complete!") print("="*70)