import os import sys import gc import re import glob import argparse import numpy as np import xarray as xr import pandas as pd import matplotlib.pyplot as plt from datetime import datetime, timedelta from scipy.interpolate import interp1d from scipy.ndimage import gaussian_filter1d import warnings warnings.filterwarnings('ignore') # Try to import h5py for RAMS and fallback data reading try: import h5py HAS_H5PY = True except ImportError: HAS_H5PY = False print("Warning: h5py not available, some features may be limited") # Try to import netCDF4 for alternative reading try: import netCDF4 HAS_NETCDF4 = True except ImportError: HAS_NETCDF4 = False # ============================================================================ # CONSTANTS # ============================================================================ P0 = 100000.0 # Reference pressure (Pa) RD = 287.05 # Gas constant for dry air (J/kg/K) RV = 461.5 # Gas constant for water vapor (J/kg/K) CP = 1004.0 # Specific heat at constant pressure (J/kg/K) KAPPA = RD / CP # Poisson constant G = 9.81 # Gravitational acceleration (m/s²) EPS = RD / RV # Ratio of gas constants (~0.622) # Target height levels for interpolation (m) TARGET_HEIGHTS = np.array([24, 50, 75, 100, 125, 150, 200, 250, 300, 400, 500, 600, 700, 800, 900, 1000, 1100, 1200, 1300, 1400, 1500, 1700, 1900, 2100, 2300, 2500, 2800, 3100, 3400, 3700, 4000]) # ============================================================================ # PATH CONFIGURATION # ============================================================================ BASE_PATH = '/global/cfs/projectdirs/m4486/Haochen' TRACER_MODEL_PATH = os.path.join(BASE_PATH, 'TRACER_Model') TRACERMIP_PATH = '/global/cfs/projectdirs/m4486/www/TRACERMIP' # ARM Sonde observation data SONDE_FILE = os.path.join(BASE_PATH, 'Domenic_Files/post_processing/ARM_data', 'katia_sonde_data/value_added_ABL_properties_Houston_2022_v1', 'value_added_ABL_properties_Houston_2022_v1.nc') # Output directory OUTPUT_DIR = os.path.join(BASE_PATH, 'Multimodel_Sonde_Compare') # ============================================================================ # MODEL CONFIGURATION # ============================================================================ MODEL_CONFIG = { 'WRF_SBM_ANL': { 'type': 'wrf_anl', 'june17': os.path.join(TRACER_MODEL_PATH, 'WRF_FSBM_ANL/17June_500m/RefCCN'), 'aug7': os.path.join(TRACER_MODEL_PATH, 'WRF_FSBM_ANL/07Aug_500m/RefCCN'), 'color': 'black', 'marker': 's', 'label': 'WRF_SBM_ANL' }, 'WRF_SBM_TAMU': { 'type': 'wrf_tamu', 'june17': os.path.join(TRACER_MODEL_PATH, 'TAMU_WRF_FSBM2/20220617/control'), 'aug7': os.path.join(TRACER_MODEL_PATH, 'TAMU_WRF_FSBM2/20220807/control'), 'color': 'red', 'marker': '^', 'label': 'WRF_SBM_TAMU' }, 'WRF_NU': { 'type': 'wrf_nu', 'june17': os.path.join(TRACER_MODEL_PATH, 'NUWRF/0617_2022/control/10min_output'), 'aug7': os.path.join(TRACER_MODEL_PATH, 'NUWRF/0807_2022/control/10min_output'), 'color': 'blue', 'marker': 'v', 'label': 'WRF_NU' }, 'RAMS': { 'type': 'rams', 'june17': os.path.join(TRACER_MODEL_PATH, 'RAMS/a.jun17.preconv.500m'), 'aug7': os.path.join(TRACER_MODEL_PATH, 'RAMS/a.aug07.preconv.500m'), 'color': 'brown', 'marker': 'D', 'label': 'RAMS' }, 'ICON_KIT': { 'type': 'icon', 'june17': os.path.join(TRACER_MODEL_PATH, 'ICON_KIT/17June_500m/RefCCN'), 'aug7': os.path.join(TRACER_MODEL_PATH, 'ICON_KIT/07Aug_500m/RefCCN'), 'color': 'purple', 'marker': 'o', 'label': 'ICON_KIT' }, 'E3SM_THREAD': { 'type': 'e3sm', 'june17': os.path.join(TRACER_MODEL_PATH, 'E3SM_THREAD/20220617_ctrl'), 'aug7': os.path.join(TRACER_MODEL_PATH, 'E3SM_THREAD/20220807_ctrl'), 'color': 'orange', 'marker': 'p', 'label': 'E3SM_THREAD' }, 'CM1': { 'type': 'cm1', 'june17': os.path.join(TRACER_MODEL_PATH, 'CM1_idealized/June17/control'), 'aug7': os.path.join(TRACER_MODEL_PATH, 'CM1_idealized/Aug7/control'), 'color': 'green', 'marker': 'h', 'label': 'CM1' } } # Case configuration CASE_CONFIG = { 'june17': { 'start_datetime': datetime(2022, 6, 17, 12, 0), # 12:00 UTC 'end_datetime': datetime(2022, 6, 18, 6, 0), # 06:00 UTC next day 'date_label': 'June 17, 2022', 'valid_sites': [0, 1, 2, 5, 8, 11, 12] # Sites with data }, 'aug7': { 'start_datetime': datetime(2022, 8, 7, 12, 0), # 12:00 UTC 'end_datetime': datetime(2022, 8, 8, 6, 0), # 06:00 UTC next day 'date_label': 'August 7, 2022', 'valid_sites': [0, 1, 4, 5, 7, 8] } } # UTC offset for Houston (CDT = UTC-5 in summer) UTC_OFFSET_HOURS = 5 # CM1 domain center for coordinate conversion CM1_CENTER_LON = -95.0792 CM1_CENTER_LAT = 29.4719 # Known TRACER ARM site locations (from Lamer et al. 2024, Scientific Data) # Data file uses 0-based indexing; Paper uses Sites 1-14 # Reference: https://doi.org/10.1038/s41597-024-03477-9 KNOWN_SITE_LOCATIONS = { 0: (29.6710, -95.0595), # Site 1: La Porte Airport 1: (29.3313, -95.7409), # Site 2: Guy 2: (29.5319, -95.2838), # Site 3: Pearland 3: (29.7499, -95.3634), # Site 4: Downtown Houston 4: (30.0703, -95.9380), # Site 5: Waller 5: (29.9011, -95.3262), # Site 6: Aldine 6: (29.5379, -94.9389), # Site 7: Galveston Bay 7: (29.3383, -94.7870), # Site 8: Galveston 8: (29.3861, -95.0421), # Site 9: UH Coastal Center 9: (29.5171, -95.7034), # Site 10: Crabb (George Ranch HS) 10: (29.5267, -95.3831), # Site 11: Manvel 11: (28.9504, -95.3038), # Site 12: Surfside 12: (29.8477, -94.7548), # Site 13: Lake Charlotte 13: (29.6524, -95.6189), # Site 14: Sugarland (Hampton Inn) } # ============================================================================ # THERMODYNAMIC CALCULATIONS # ============================================================================ def calc_saturation_vapor_pressure(T): """ Calculate saturation vapor pressure using Tetens formula. T: Temperature in Kelvin Returns: Saturation vapor pressure in Pa """ T_celsius = T - 273.15 # Use different formula for ice vs water es = np.where(T_celsius >= 0, 611.2 * np.exp(17.67 * T_celsius / (T_celsius + 243.5)), 611.2 * np.exp(21.8745584 * T_celsius / (T_celsius + 265.5))) return es def calc_relative_humidity(qv, T, P): """ Calculate relative humidity from mixing ratio, temperature, and pressure. qv: Water vapor mixing ratio (kg/kg) T: Temperature (K) P: Pressure (Pa) Returns: Relative humidity (%) """ es = calc_saturation_vapor_pressure(T) # Saturation mixing ratio qs = EPS * es / (P - es) # Relative humidity rh = 100.0 * qv / qs return np.clip(rh, 0, 100) def calc_potential_temperature(T, P): """ Calculate potential temperature. T: Temperature (K) P: Pressure (Pa) Returns: Potential temperature (K) """ return T * (P0 / P) ** KAPPA def calc_temperature_from_theta(theta, P): """ Calculate temperature from potential temperature. theta: Potential temperature (K) P: Pressure (Pa) Returns: Temperature (K) """ return theta * (P / P0) ** KAPPA def specific_to_mixing_ratio(qv_specific): """ Convert specific humidity to mixing ratio. qv_specific: Specific humidity (kg/kg) Returns: Mixing ratio (kg/kg) """ return qv_specific / (1.0 - qv_specific) # ============================================================================ # DATA LOADING FUNCTIONS # ============================================================================ def load_sonde_data(case_name): """Load ARM sonde observation data for a specific case.""" print(f"Loading sonde data for {case_name}...") # Check if file exists if not os.path.exists(SONDE_FILE): raise FileNotFoundError(f"Sonde file not found: {SONDE_FILE}") print(f" File exists: {SONDE_FILE}") print(f" File size: {os.path.getsize(SONDE_FILE) / 1e6:.2f} MB") ds = None # Method 1: Try h5py first (often works better with HDF5-based NetCDF4) if HAS_H5PY: print(" Trying h5py directly...") try: with h5py.File(SONDE_FILE, 'r') as f: print(f" Found {len(f.keys())} variables") data_vars = {} coords = {} # Known dimension sizes for this sonde file n_sites = 14 n_times = 1104 n_heights = 133 for key in f.keys(): arr = np.array(f[key]) # Replace fill values with NaN if arr.dtype in [np.float32, np.float64]: arr = np.where(arr == -9999, np.nan, arr) arr = np.where(np.abs(arr) > 1e30, np.nan, arr) # Determine dimensions based on shape shape = arr.shape if arr.ndim == 1: if len(arr) == n_sites: dims = ['Site_dimension'] elif len(arr) == n_heights: dims = ['Height_dimension'] elif len(arr) == n_times: dims = ['Time_dimension'] else: dims = [f'dim_{len(arr)}'] elif arr.ndim == 2: if shape == (n_times, n_heights): dims = ['Time_dimension', 'Height_dimension'] elif shape == (n_sites, n_heights): dims = ['Site_dimension', 'Height_dimension'] elif shape == (n_sites, n_times): dims = ['Site_dimension', 'Time_dimension'] else: dims = [f'dim_{s}' for s in shape] elif arr.ndim == 3: if shape == (n_sites, n_times, n_heights): dims = ['Site_dimension', 'Time_dimension', 'Height_dimension'] else: dims = [f'dim_{s}' for s in shape] else: dims = [f'dim_{i}' for i in range(arr.ndim)] # Check if this is a coordinate variable if key.endswith('_dimension') or key in ['Height_dimension', 'Site_dimension', 'Time_dimension']: coords[key] = arr else: data_vars[key] = (dims, arr) ds = xr.Dataset(data_vars, coords=coords) print(" Successfully read with h5py") except Exception as e: print(f" h5py failed: {str(e)[:100]}") # Method 2: Try xarray with different engines if ds is None: engines_to_try = ['h5netcdf', 'scipy'] for engine in engines_to_try: try: print(f" Trying xarray with engine: {engine}...") ds = xr.open_dataset(SONDE_FILE, engine=engine) print(f" Successfully opened with xarray/{engine}") break except Exception as e: print(f" Failed with {engine}: {str(e)[:80]}") continue if ds is None: raise RuntimeError(f"Could not open sonde file with any method. " f"File may be corrupted or locked.") # Print available variables print(f" Variables available: {list(ds.data_vars)[:10]}...") # Check for height-related variables height_candidates = [v for v in ds.data_vars if 'height' in v.lower() or 'Height' in v or 'altitude' in v.lower()] if height_candidates: print(f" Height-related variables: {height_candidates}") # Get time range for case case_cfg = CASE_CONFIG[case_name] ds_case = ds.copy() # Mask invalid values for var in ['Potential_temperature_mean', 'Relative_humidity_mean']: if var in ds_case: ds_case[var] = ds_case[var].where(ds_case[var] != -9999) return ds_case def find_nearest_point(lat_target, lon_target, lat_2d, lon_2d): """Find nearest grid point to target lat/lon.""" dist = np.sqrt((lat_2d - lat_target)**2 + (lon_2d - lon_target)**2) idx = np.unravel_index(np.argmin(dist), dist.shape) return idx 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)) lon = center_lon + xh_km / km_per_deg_lon lat = center_lat + yh_km / km_per_deg_lat return lon, lat # ============================================================================ # MODEL-SPECIFIC DATA READERS # ============================================================================ def read_wrf_anl_data(data_path, target_time, site_lat, site_lon): """ Read WRF_SBM_ANL data (standard WRF output). Returns: theta, rh, height arrays at target location """ # Find the file closest to target time wrf_files = sorted(glob.glob(os.path.join(data_path, 'wrfout_d02_*'))) if not wrf_files: print(f" No WRF d02 files found in {data_path}") return None, None, None # Find files close to target time and try them in order target_dt = target_time candidate_files = [] for f in wrf_files: basename = os.path.basename(f) try: time_str = basename.replace('wrfout_d02_', '').replace('_', ' ').replace(':00:00', ':00') file_dt = datetime.strptime(time_str[:16], '%Y-%m-%d %H:%M') time_diff = abs((file_dt - target_dt).total_seconds()) candidate_files.append((time_diff, f)) except: continue # Sort by time difference (closest first) candidate_files.sort(key=lambda x: x[0]) # If no candidates found, use middle file if not candidate_files: candidate_files = [(0, wrf_files[len(wrf_files)//2])] last_error = None # Try files in order until one works for time_diff, target_file in candidate_files[:10]: # Try up to 10 files # Try xarray with different engines first for engine in [None, 'h5netcdf', 'scipy']: try: if engine: ds = xr.open_dataset(target_file, engine=engine) else: ds = xr.open_dataset(target_file) # Get coordinates lat = ds['XLAT'].values[0] # Remove time dimension lon = ds['XLONG'].values[0] # Find nearest point j, i = find_nearest_point(site_lat, site_lon, lat, lon) # Calculate height from geopotential ph = ds['PH'].values[0, :, j, i] # Perturbation geopotential phb = ds['PHB'].values[0, :, j, i] # Base geopotential height_stag = (ph + phb) / G height = 0.5 * (height_stag[:-1] + height_stag[1:]) # Destagger # Get potential temperature (T is perturbation, add 300K base) theta = ds['T'].values[0, :, j, i] + 300.0 # Get pressure p = ds['P'].values[0, :, j, i] # Perturbation pb = ds['PB'].values[0, :, j, i] # Base pressure = p + pb # Get water vapor mixing ratio qvapor = ds['QVAPOR'].values[0, :, j, i] # Calculate temperature and RH T = calc_temperature_from_theta(theta, pressure) rh = calc_relative_humidity(qvapor, T, pressure) ds.close() return theta, rh, height except Exception as e: last_error = str(e) continue # Try h5py directly as last resort if HAS_H5PY: try: with h5py.File(target_file, 'r') as f: lat = np.array(f['XLAT'])[0] lon = np.array(f['XLONG'])[0] j, i = find_nearest_point(site_lat, site_lon, lat, lon) theta = np.array(f['T'])[0, :, j, i] + 300.0 ph = np.array(f['PH'])[0, :, j, i] phb = np.array(f['PHB'])[0, :, j, i] height_stag = (ph + phb) / G height = 0.5 * (height_stag[:-1] + height_stag[1:]) p = np.array(f['P'])[0, :, j, i] pb = np.array(f['PB'])[0, :, j, i] pressure = p + pb qvapor = np.array(f['QVAPOR'])[0, :, j, i] T = calc_temperature_from_theta(theta, pressure) rh = calc_relative_humidity(qvapor, T, pressure) return theta, rh, height except Exception as e: last_error = str(e) continue # All files failed print(f" All WRF files failed. Last error: {last_error[:100] if last_error else 'Unknown'}") return None, None, None def read_wrf_tamu_data(data_path, target_time, site_lat, site_lon): """ Read WRF_SBM_TAMU data. May have pre-processed vars (air_potential_temperature, geopotential_height) or standard WRF vars. """ wrf_files = sorted(glob.glob(os.path.join(data_path, 'wrfout_d02_*'))) if not wrf_files: # Try d01 as fallback wrf_files = sorted(glob.glob(os.path.join(data_path, 'wrfout_d01_*'))) if not wrf_files: print(f" No WRF files found in {data_path}") return None, None, None # Find files close to target time and try them in order target_dt = target_time candidate_files = [] for f in wrf_files: basename = os.path.basename(f) try: # Handle both d01 and d02 time_str = basename.replace('wrfout_d02_', '').replace('wrfout_d01_', '') time_str = time_str.replace('_', ' ').replace(':00:00', ':00') file_dt = datetime.strptime(time_str[:16], '%Y-%m-%d %H:%M') time_diff = abs((file_dt - target_dt).total_seconds()) candidate_files.append((time_diff, f)) except: continue # Sort by time difference (closest first) candidate_files.sort(key=lambda x: x[0]) # If no candidates found, use middle file if not candidate_files: candidate_files = [(0, wrf_files[len(wrf_files)//2])] last_error = None # Try files in order until one works for time_diff, target_file in candidate_files[:10]: # Try up to 10 files # Try xarray with different engines first for engine in [None, 'h5netcdf', 'scipy']: try: if engine: ds = xr.open_dataset(target_file, engine=engine) else: ds = xr.open_dataset(target_file) # Get coordinates lat = ds['XLAT'].values lon = ds['XLONG'].values # Handle time dimension if present if lat.ndim == 3: lat = lat[0] lon = lon[0] # Find nearest point j, i = find_nearest_point(site_lat, site_lon, lat, lon) # Check if pre-processed variables exist if 'air_potential_temperature' in ds: # Pre-processed WRF output theta = ds['air_potential_temperature'].values[0, :, j, i] height = ds['geopotential_height'].values[0, :, j, i] pressure = ds['air_pressure'].values[0, :, j, i] qvapor = ds['QVAPOR'].values[0, :, j, i] else: # Standard WRF output - calculate variables theta = ds['T'].values[0, :, j, i] + 300.0 # Perturbation + base # Height from geopotential ph = ds['PH'].values[0, :, j, i] phb = ds['PHB'].values[0, :, j, i] height_stag = (ph + phb) / G height = 0.5 * (height_stag[:-1] + height_stag[1:]) # Pressure p = ds['P'].values[0, :, j, i] pb = ds['PB'].values[0, :, j, i] pressure = p + pb qvapor = ds['QVAPOR'].values[0, :, j, i] # Calculate temperature and RH T = calc_temperature_from_theta(theta, pressure) rh = calc_relative_humidity(qvapor, T, pressure) ds.close() return theta, rh, height except Exception as e: last_error = str(e) continue # Try h5py directly as last resort if HAS_H5PY: try: with h5py.File(target_file, 'r') as f: lat = np.array(f['XLAT']) lon = np.array(f['XLONG']) if lat.ndim == 3: lat = lat[0] lon = lon[0] j, i = find_nearest_point(site_lat, site_lon, lat, lon) # Check if pre-processed or standard WRF if 'air_potential_temperature' in f: theta = np.array(f['air_potential_temperature'])[0, :, j, i] height = np.array(f['geopotential_height'])[0, :, j, i] pressure = np.array(f['air_pressure'])[0, :, j, i] qvapor = np.array(f['QVAPOR'])[0, :, j, i] else: theta = np.array(f['T'])[0, :, j, i] + 300.0 ph = np.array(f['PH'])[0, :, j, i] phb = np.array(f['PHB'])[0, :, j, i] height_stag = (ph + phb) / G height = 0.5 * (height_stag[:-1] + height_stag[1:]) p = np.array(f['P'])[0, :, j, i] pb = np.array(f['PB'])[0, :, j, i] pressure = p + pb qvapor = np.array(f['QVAPOR'])[0, :, j, i] T = calc_temperature_from_theta(theta, pressure) rh = calc_relative_humidity(qvapor, T, pressure) return theta, rh, height except Exception as e: last_error = str(e) continue # All files failed print(f" All WRF_TAMU files failed. Last error: {last_error[:100] if last_error else 'Unknown'}") return None, None, None def read_wrf_nu_data(data_path, target_time, site_lat, site_lon): """ Read WRF_NU data (standard WRF output, same as ANL). """ return read_wrf_anl_data(data_path, target_time, site_lat, site_lon) def read_rams_data(data_path, target_time, site_lat, site_lon): """ Read RAMS data (HDF5 format). Variables: THETA, RV, PI, GLAT, GLON File pattern: a-A-2022-06-17-060000-g2.h5 (YYYY-MM-DD-HHMMSS, g2 = fine grid) """ if not HAS_H5PY: print(" h5py not available, skipping RAMS") return None, None, None # Find HDF5 files - prefer g2 (fine grid) h5_files = sorted(glob.glob(os.path.join(data_path, 'a-A-*-g2.h5'))) if not h5_files: # Fallback to any h5 file h5_files = sorted(glob.glob(os.path.join(data_path, '*.h5'))) if not h5_files: print(f" No RAMS HDF5 files found in {data_path}") return None, None, None # Find file closest to target time # Filename pattern: a-A-2022-06-17-060000-g2.h5 target_file = None best_diff = float('inf') for f in h5_files: basename = os.path.basename(f) try: # Extract timestamp: a-A-YYYY-MM-DD-HHMMSS-gN.h5 parts = basename.split('-') if len(parts) >= 6: # parts: ['a', 'A', '2022', '06', '17', '060000', 'g2.h5'] year = int(parts[2]) month = int(parts[3]) day = int(parts[4]) time_str = parts[5] # '060000' hour = int(time_str[:2]) minute = int(time_str[2:4]) second = int(time_str[4:6]) file_dt = datetime(year, month, day, hour, minute, second) time_diff = abs((file_dt - target_time).total_seconds()) if time_diff < best_diff: best_diff = time_diff target_file = f except: continue # Fallback to middle file if no time match found if target_file is None: target_file = h5_files[min(len(h5_files)//2, len(h5_files)-1)] print(f" RAMS: No time match, using fallback file") # Debug: print selected file timestamp basename = os.path.basename(target_file) parts = basename.split('-') if len(parts) >= 6: print(f"[{parts[2]}-{parts[3]}-{parts[4]}-{parts[5]}]", end=' ') try: with h5py.File(target_file, 'r') as f: # Get coordinates if 'GLAT' in f: lat = f['GLAT'][:] lon = f['GLON'][:] else: print(" GLAT/GLON not found in RAMS file") return None, None, None # Find nearest point j, i = find_nearest_point(site_lat, site_lon, lat, lon) # Get potential temperature if 'THETA' in f: theta = f['THETA'][:, j, i] else: print(" THETA not found in RAMS file") return None, None, None # Get water vapor mixing ratio if 'RV' in f: rv = f['RV'][:, j, i] else: rv = np.zeros_like(theta) # Get Exner function to calculate pressure if 'PI' in f: pi = f['PI'][:, j, i] # PI = Exner * Cp exner = pi / CP pressure = P0 * exner ** (CP / RD) elif 'PI0' in f and 'PP' in f: pi0 = f['PI0'][:, j, i] pp = f['PP'][:, j, i] pi = pi0 + pp exner = pi / CP pressure = P0 * exner ** (CP / RD) else: # Estimate pressure from standard atmosphere print(" Warning: Using estimated pressure for RAMS") pressure = P0 * np.exp(-np.arange(len(theta)) * 100 / 8500) # Calculate height from vertical grid (if available) # RAMS uses terrain-following coordinates if 'TOPT' in f: topo = f['TOPT'][j, i] else: topo = 0 # Estimate height levels (RAMS typically has ~100m spacing near surface) nz = len(theta) height = np.zeros(nz) for k in range(nz): # Simple estimate: stretched grid if k < 20: height[k] = topo + k * 50 # 50m spacing near surface else: height[k] = height[19] + (k - 19) * 100 # 100m spacing above # Calculate temperature and RH T = calc_temperature_from_theta(theta, pressure) rh = calc_relative_humidity(rv, T, pressure) return theta, rh, height except Exception as e: print(f" Error reading RAMS: {e}") return None, None, None def read_icon_data(data_path, target_time, site_lat, site_lon): """ Read ICON_KIT data. Variables: temp, pres, qv (specific humidity), z_mc (height) File pattern: NWP_LAM_DOM01_20220617T060000Z_TRACER_10min_reglatlon.nc """ nc_files = sorted(glob.glob(os.path.join(data_path, '*.nc'))) if not nc_files: print(f" No ICON files found in {data_path}") return None, None, None # Find file closest to target time # Filename pattern: NWP_LAM_DOM01_20220617T060000Z_... target_file = None best_diff = float('inf') for f in nc_files: basename = os.path.basename(f) try: # Extract timestamp from filename: ...20220617T060000Z... match = re.search(r'(\d{8}T\d{6})Z', basename) if match: time_str = match.group(1) file_dt = datetime.strptime(time_str, '%Y%m%dT%H%M%S') time_diff = abs((file_dt - target_time).total_seconds()) if time_diff < best_diff: best_diff = time_diff target_file = f except: continue # Fallback to middle file if no time match found if target_file is None: target_file = nc_files[min(len(nc_files)//2, len(nc_files)-1)] print(f" ICON: No time match, using fallback file") # Debug: print selected file timestamp basename = os.path.basename(target_file) match = re.search(r'(\d{8}T\d{6})Z', basename) if match: print(f"[{match.group(1)}]", end=' ') try: ds = xr.open_dataset(target_file) # Get coordinates - use 1D search to avoid creating large 2D meshgrid lat = ds['lat'].values lon = ds['lon'].values # Find nearest indices directly from 1D arrays i = int(np.argmin(np.abs(lon - site_lon))) j = int(np.argmin(np.abs(lat - site_lat))) # Get variables - only read the column we need temp = ds['temp'].values[0, :, j, i] # Temperature (K) pres = ds['pres'].values[0, :, j, i] # Pressure (Pa) qv_spec = ds['qv'].values[0, :, j, i] # Specific humidity (kg/kg) height = ds['z_mc'].values[:, j, i] # Height (m) ds.close() # ICON levels are top-to-bottom, may need to flip if height[0] > height[-1]: temp = temp[::-1] pres = pres[::-1] qv_spec = qv_spec[::-1] height = height[::-1] # Calculate potential temperature theta = calc_potential_temperature(temp, pres) # Convert specific humidity to mixing ratio qv = specific_to_mixing_ratio(qv_spec) # Calculate RH rh = calc_relative_humidity(qv, temp, pres) return theta, rh, height except Exception as e: print(f" Error reading ICON: {e}") return None, None, None def read_e3sm_data(data_path, target_time, site_lat, site_lon): """ Read E3SM_THREAD data. Variables in separate files: T_mid, p_mid, qv, geopotential_mid File pattern: T_mid_F2010-SCREAMv1.TRACERmipx32subx3v1pg2.001.3D.State.INST.INSTANT.nmins_x2.2022-06-17-00000.nc """ # Find files for T_mid with the correct pattern t_files = sorted(glob.glob(os.path.join(data_path, 'T_mid_F2010-SCREAMv1*.nc'))) if not t_files: # Try alternative pattern t_files = sorted(glob.glob(os.path.join(data_path, 'T_mid_*.nc'))) if not t_files: print(f" No E3SM T_mid files found in {data_path}") return None, None, None # Find file closest to target time (files have timestamps like 2022-06-17-43200 at the end) target_file_t = None best_diff = float('inf') for f in t_files: # Extract timestamp from filename basename = os.path.basename(f) # Pattern: ...nmins_x2.2022-06-17-43200.nc try: # Get the date-time part before .nc date_part = basename.replace('.nc', '').split('.')[-1] # e.g., "2022-06-17-43200" parts = date_part.split('-') if len(parts) >= 4: file_sec = int(parts[-1]) date_str = '-'.join(parts[:3]) file_date = datetime.strptime(date_str, '%Y-%m-%d') file_dt = file_date + timedelta(seconds=file_sec) time_diff = abs((file_dt - target_time).total_seconds()) if time_diff < best_diff: best_diff = time_diff target_file_t = f except Exception as e: # Debug: print parsing errors # print(f" E3SM parse error for {basename}: {e}") continue if target_file_t is None: target_file_t = t_files[min(len(t_files)//2, len(t_files)-1)] print(f" E3SM: No time match, using fallback file") # Debug: print selected file selected_basename = os.path.basename(target_file_t) selected_timestamp = selected_basename.replace('.nc', '').split('.')[-1] print(f"[{selected_timestamp}]", end=' ') # Construct corresponding file paths by replacing variable name base_pattern = os.path.basename(target_file_t) # Replace T_mid with other variable names p_pattern = base_pattern.replace('T_mid_', 'p_mid_') qv_pattern = base_pattern.replace('T_mid_', 'qv_').replace('.3D.State.', '.3D.Water.') z_pattern = base_pattern.replace('T_mid_', 'geopotential_mid_') target_file_p = os.path.join(data_path, p_pattern) target_file_qv = os.path.join(data_path, qv_pattern) target_file_z = os.path.join(data_path, z_pattern) last_error = None # E3SM fill value threshold - values above this are considered invalid FILL_VALUE_THRESHOLD = 1e30 # Try netCDF4 first - most memory efficient for reading single columns if HAS_NETCDF4: try: nc = netCDF4.Dataset(target_file_t, 'r') lat = nc.variables['lat'][:] lon = nc.variables['lon'][:] # Find nearest indices (1D search) i = int(np.argmin(np.abs(lon - site_lon))) j = int(np.argmin(np.abs(lat - site_lat))) # Read only the column we need - time index 0 (files have internal timesteps but we want first) temp = np.asarray(nc.variables['T_mid'][0, :, j, i]) # Check if height variable exists in this file (newer E3SM format) if 'height' in nc.variables: height = np.asarray(nc.variables['height'][0, :, j, i]) else: height = None nc.close() # Read pressure if os.path.exists(target_file_p): nc_p = netCDF4.Dataset(target_file_p, 'r') pres = np.asarray(nc_p.variables['p_mid'][0, :, j, i]) nc_p.close() else: pres = P0 * np.exp(-np.arange(len(temp)) * 250 / 8500) # Read water vapor if os.path.exists(target_file_qv): nc_qv = netCDF4.Dataset(target_file_qv, 'r') qv = np.asarray(nc_qv.variables['qv'][0, :, j, i]) nc_qv.close() else: qv = np.zeros_like(temp) # Read height from geopotential if not already loaded if height is None: if os.path.exists(target_file_z): nc_z = netCDF4.Dataset(target_file_z, 'r') geopotential = np.asarray(nc_z.variables['geopotential_mid'][0, :, j, i]) height = geopotential / G nc_z.close() else: height = -8500 * np.log(pres / P0) # Mask fill values - E3SM uses very large values as fill valid_mask = (np.abs(temp) < FILL_VALUE_THRESHOLD) & \ (np.abs(pres) < FILL_VALUE_THRESHOLD) & \ (np.abs(height) < FILL_VALUE_THRESHOLD) temp = temp[valid_mask] pres = pres[valid_mask] qv = qv[valid_mask] height = height[valid_mask] if len(temp) == 0: print("No valid data after masking fill values") return None, None, None # E3SM levels are top-to-bottom, flip if needed # Check using valid data only if len(height) > 1 and height[0] > height[-1]: temp = temp[::-1] pres = pres[::-1] qv = qv[::-1] height = height[::-1] # Calculate potential temperature theta = calc_potential_temperature(temp, pres) # Calculate RH rh = calc_relative_humidity(qv, temp, pres) return theta, rh, height except Exception as e: last_error = str(e) # Fallback to xarray with different engines for engine in ['h5netcdf', None, 'scipy']: try: # Read temperature if engine: ds_t = xr.open_dataset(target_file_t, engine=engine) else: ds_t = xr.open_dataset(target_file_t) lat = ds_t['lat'].values lon = ds_t['lon'].values # Find nearest indices i = int(np.argmin(np.abs(lon - site_lon))) j = int(np.argmin(np.abs(lat - site_lat))) temp = ds_t['T_mid'].values[0, :, j, i] # Check if height variable exists in this file if 'height' in ds_t: height = ds_t['height'].values[0, :, j, i] else: height = None ds_t.close() # Read pressure if os.path.exists(target_file_p): if engine: ds_p = xr.open_dataset(target_file_p, engine=engine) else: ds_p = xr.open_dataset(target_file_p) pres = ds_p['p_mid'].values[0, :, j, i] ds_p.close() else: pres = P0 * np.exp(-np.arange(len(temp)) * 250 / 8500) # Read water vapor if os.path.exists(target_file_qv): if engine: ds_qv = xr.open_dataset(target_file_qv, engine=engine) else: ds_qv = xr.open_dataset(target_file_qv) qv = ds_qv['qv'].values[0, :, j, i] ds_qv.close() else: qv = np.zeros_like(temp) # Read height from geopotential if not already loaded if height is None: if os.path.exists(target_file_z): if engine: ds_z = xr.open_dataset(target_file_z, engine=engine) else: ds_z = xr.open_dataset(target_file_z) geopotential = ds_z['geopotential_mid'].values[0, :, j, i] height = geopotential / G ds_z.close() else: height = -8500 * np.log(pres / P0) # Mask fill values - E3SM uses very large values as fill valid_mask = (np.abs(temp) < FILL_VALUE_THRESHOLD) & \ (np.abs(pres) < FILL_VALUE_THRESHOLD) & \ (np.abs(height) < FILL_VALUE_THRESHOLD) temp = temp[valid_mask] pres = pres[valid_mask] qv = qv[valid_mask] height = height[valid_mask] if len(temp) == 0: continue # E3SM levels are top-to-bottom, flip if needed if len(height) > 1 and height[0] > height[-1]: temp = temp[::-1] pres = pres[::-1] qv = qv[::-1] height = height[::-1] # Calculate potential temperature theta = calc_potential_temperature(temp, pres) # Calculate RH rh = calc_relative_humidity(qv, temp, pres) return theta, rh, height except Exception as e: last_error = str(e) continue print(f" E3SM read failed: {last_error[:80] if last_error else 'Unknown'}") return None, None, None def read_cm1_data(data_path, target_time, site_lat, site_lon): """ Read CM1 data. Variables: tk (temperature K), prs (pressure Pa), qv, zhval (height m) CM1 uses km coordinates that need conversion to lat/lon. File pattern: cm1out_0600-1800UTC.nc (files contain multiple time steps) """ nc_files = sorted(glob.glob(os.path.join(data_path, 'cm1out_*.nc'))) if not nc_files: nc_files = sorted(glob.glob(os.path.join(data_path, '*.nc'))) if not nc_files: print(f" No CM1 files found in {data_path}") return None, None, None # Get target hour in UTC (0-23, but CM1 uses continuous hours like 30 for 06:00 next day) target_utc_hour = target_time.hour + (target_time.day - 17) * 24 # Relative to June 17 # For June 17 case, target_utc_hour is simply target_time.hour # For June 18 00:00 UTC, it would be 24 # Select the right file based on time range in filename # Filename pattern: cm1out_0600-1800UTC.nc means 06:00 to 18:00 UTC target_file = None for f in nc_files: basename = os.path.basename(f) try: # Extract time range: cm1out_HHMM-HHMM UTC.nc # e.g., cm1out_0600-1800UTC.nc -> start=6, end=18 parts = basename.replace('cm1out_', '').replace('UTC.nc', '').replace('.nc', '') time_parts = parts.split('-') if len(time_parts) == 2: start_hour = int(time_parts[0][:2]) # First 2 digits end_hour = int(time_parts[1][:2]) # First 2 digits # Check if target time falls within this range if start_hour <= target_utc_hour < end_hour or \ (end_hour > 24 and start_hour <= target_utc_hour) or \ (end_hour > 24 and target_utc_hour < end_hour - 24): target_file = f break except: continue # Fallback to first file if no match if target_file is None: target_file = nc_files[0] # Debug: print selected file and time index print(f"[{os.path.basename(target_file)}]", end=' ') last_error = None # Try xarray with different engines for engine in ['h5netcdf', None, 'scipy']: try: if engine: ds = xr.open_dataset(target_file, engine=engine) else: ds = xr.open_dataset(target_file) # Find the correct time index within the file time_idx = 0 # Default to first time step if 'time' in ds.dims or 'time' in ds: n_times = ds.dims.get('time', 1) if n_times > 1: # Try to match time - CM1 time might be in seconds or hours if 'time' in ds and ds['time'].size > 1: time_var = ds['time'].values # Assuming time is in seconds from start of simulation # or hours - we'll pick the closest one based on hour of day target_hour_frac = target_time.hour + target_time.minute / 60.0 # Try to interpret time values if np.max(time_var) > 100: # Likely seconds file_hours = time_var / 3600.0 else: # Likely hours already file_hours = time_var # Find closest time index (accounting for file start time from filename) basename = os.path.basename(target_file) try: start_str = basename.replace('cm1out_', '')[:4] file_start_hour = int(start_str[:2]) except: file_start_hour = 0 # Adjust target hour relative to file start target_rel_hour = target_utc_hour - file_start_hour time_idx = int(np.argmin(np.abs(file_hours - target_rel_hour))) # Get coordinates (CM1 uses xh, yh in km) if 'xh' in ds.coords: xh = ds['xh'].values yh = ds['yh'].values elif 'xh' in ds: xh = ds['xh'].values yh = ds['yh'].values else: xh = ds['x'].values if 'x' in ds else np.arange(ds.dims.get('x', 100)) yh = ds['y'].values if 'y' in ds else np.arange(ds.dims.get('y', 100)) # Vectorized km to lat/lon conversion lon_1d = CM1_CENTER_LON + (xh / (111.32 * np.cos(np.radians(CM1_CENTER_LAT)))) lat_1d = CM1_CENTER_LAT + (yh / 110.574) # Find nearest indices directly i = int(np.argmin(np.abs(lon_1d - site_lon))) j = int(np.argmin(np.abs(lat_1d - site_lat))) # Get variables at the correct time index if 'tk' in ds: temp = ds['tk'].values[time_idx, :, j, i] theta_raw = None elif 'theta' in ds: theta_raw = ds['theta'].values[time_idx, :, j, i] temp = None else: ds.close() continue if 'prs' in ds: pres = ds['prs'].values[time_idx, :, j, i] else: pres = P0 * np.exp(-np.arange(len(temp if temp is not None else theta_raw)) * 100 / 8500) if 'qv' in ds: qv = ds['qv'].values[time_idx, :, j, i] else: qv = np.zeros_like(temp if temp is not None else theta_raw) if 'zhval' in ds: height = ds['zhval'].values[time_idx, :, j, i] elif 'zh' in ds: height = ds['zh'].values elif 'z' in ds: height = ds['z'].values else: height = np.arange(len(temp if temp is not None else theta_raw)) * 100 if temp is not None: theta = calc_potential_temperature(temp, pres) else: theta = theta_raw temp = calc_temperature_from_theta(theta, pres) rh = calc_relative_humidity(qv, temp, pres) ds.close() return theta, rh, height except Exception as e: last_error = str(e) continue # Try h5py as fallback if HAS_H5PY: try: with h5py.File(target_file, 'r') as f: # Find time index time_idx = 0 if 'time' in f: time_var = np.array(f['time']) if len(time_var) > 1: target_hour_frac = target_time.hour + target_time.minute / 60.0 if np.max(time_var) > 100: file_hours = time_var / 3600.0 else: file_hours = time_var basename = os.path.basename(target_file) try: start_str = basename.replace('cm1out_', '')[:4] file_start_hour = int(start_str[:2]) except: file_start_hour = 0 target_rel_hour = target_utc_hour - file_start_hour time_idx = int(np.argmin(np.abs(file_hours - target_rel_hour))) # Get coordinates if 'xh' in f: xh = np.array(f['xh']) yh = np.array(f['yh']) elif 'x' in f: xh = np.array(f['x']) yh = np.array(f['y']) else: print(f" CM1: No coordinate variables found") return None, None, None # Vectorized km to lat/lon conversion lon_1d = CM1_CENTER_LON + (xh / (111.32 * np.cos(np.radians(CM1_CENTER_LAT)))) lat_1d = CM1_CENTER_LAT + (yh / 110.574) i = int(np.argmin(np.abs(lon_1d - site_lon))) j = int(np.argmin(np.abs(lat_1d - site_lat))) # Get variables at the correct time index if 'tk' in f: temp = np.array(f['tk'][time_idx, :, j, i]) theta_raw = None elif 'theta' in f: theta_raw = np.array(f['theta'][time_idx, :, j, i]) temp = None else: return None, None, None if 'prs' in f: pres = np.array(f['prs'][time_idx, :, j, i]) else: pres = P0 * np.exp(-np.arange(len(temp if temp is not None else theta_raw)) * 100 / 8500) if 'qv' in f: qv = np.array(f['qv'][time_idx, :, j, i]) else: qv = np.zeros_like(temp if temp is not None else theta_raw) if 'zhval' in f: height = np.array(f['zhval'][time_idx, :, j, i]) elif 'zh' in f: height = np.array(f['zh']) elif 'z' in f: height = np.array(f['z']) else: height = np.arange(len(temp if temp is not None else theta_raw)) * 100 if temp is not None: theta = calc_potential_temperature(temp, pres) else: theta = theta_raw temp = calc_temperature_from_theta(theta, pres) rh = calc_relative_humidity(qv, temp, pres) return theta, rh, height except Exception as e: last_error = str(e) # All files failed print(f" CM1 read failed: {last_error[:80] if last_error else 'Unknown'}") return None, None, None # ============================================================================ # MAIN DATA READER DISPATCHER # ============================================================================ def read_model_data(model_name, data_path, target_time, site_lat, site_lon): """ Dispatch to appropriate reader based on model type. """ model_type = MODEL_CONFIG[model_name]['type'] if model_type == 'wrf_anl': return read_wrf_anl_data(data_path, target_time, site_lat, site_lon) elif model_type == 'wrf_tamu': return read_wrf_tamu_data(data_path, target_time, site_lat, site_lon) elif model_type == 'wrf_nu': return read_wrf_nu_data(data_path, target_time, site_lat, site_lon) elif model_type == 'rams': return read_rams_data(data_path, target_time, site_lat, site_lon) elif model_type == 'icon': return read_icon_data(data_path, target_time, site_lat, site_lon) elif model_type == 'e3sm': return read_e3sm_data(data_path, target_time, site_lat, site_lon) elif model_type == 'cm1': return read_cm1_data(data_path, target_time, site_lat, site_lon) else: print(f" Unknown model type: {model_type}") return None, None, None def interpolate_to_heights(data, heights, target_heights): """ Interpolate data to target height levels. """ # Remove NaN values valid = ~np.isnan(data) & ~np.isnan(heights) if np.sum(valid) < 2: return np.full(len(target_heights), np.nan) # Create interpolation function f = interp1d(heights[valid], data[valid], bounds_error=False, fill_value=np.nan) return f(target_heights) # ============================================================================ # PLOTTING FUNCTIONS # ============================================================================ def smooth_profile(data, sigma=2): """ Apply Gaussian smoothing to profile data. Handles NaN values by interpolating, smoothing, then restoring NaN positions. Parameters: data: 1D array of profile data sigma: Standard deviation for Gaussian kernel (higher = smoother) Returns: Smoothed data array """ data = np.array(data, dtype=float) # Find valid (non-NaN) indices valid_mask = ~np.isnan(data) if np.sum(valid_mask) < 3: return data # Not enough points to smooth # Get valid data valid_indices = np.where(valid_mask)[0] valid_data = data[valid_mask] # Apply Gaussian smoothing to valid data smoothed_valid = gaussian_filter1d(valid_data, sigma=sigma) # Put smoothed values back result = data.copy() result[valid_mask] = smoothed_valid return result def plot_multimodel_comparison(obs_data, model_data_dict, var_name, site_idx, target_time, case_name, save_path): """ Create comparison plot of observations vs all models. """ fig, ax = plt.subplots(figsize=(10, 8)) # Plot observation obs_height = obs_data['height'] obs_var = obs_data[var_name] # Apply smoothing to RH observations (they tend to be noisy) if var_name == 'rh': obs_var_plot = smooth_profile(obs_var, sigma=2) else: obs_var_plot = obs_var valid = ~np.isnan(obs_var_plot) if np.any(valid): # ARM Sonde: dark grey color, same marker size as models (4) ax.plot(obs_var_plot[valid], obs_height[valid], 'o-', color='dimgray', linewidth=2, markersize=4, label='ARM Sonde', zorder=10) # Plot each model for model_name, model_data in model_data_dict.items(): if model_data is None: continue config = MODEL_CONFIG[model_name] model_var = model_data[var_name] model_height = model_data['height'] valid = ~np.isnan(model_var) if np.any(valid): ax.plot(model_var[valid], model_height[valid], color=config['color'], marker=config['marker'], linestyle='-', linewidth=1.5, markersize=4, label=config['label'], alpha=0.8) # Labels and formatting var_labels = { 'theta': ('Potential Temperature', 'K'), 'rh': ('Relative Humidity', '%') } var_label, var_unit = var_labels.get(var_name, (var_name, '')) ax.set_xlabel(f'{var_label} ({var_unit})', fontsize=12) ax.set_ylabel('Height (m)', fontsize=12) ax.set_ylim(0, 4000) if var_name == 'theta': ax.set_xlim(295, 320) # Updated range for potential temperature legend_loc = 'lower right' # Legend at bottom right for theta elif var_name == 'rh': ax.set_xlim(0, 105) legend_loc = 'lower left' # Legend at bottom left for RH else: legend_loc = 'best' # Format title with date and time in UTC time_str = target_time.strftime("%B %d, %Y %H:%M") ax.set_title(f'{var_label} Profile | Site {site_idx} | {time_str} UTC', fontsize=14) # Legend: bigger font size (11) and position based on variable ax.legend(loc=legend_loc, fontsize=11) ax.grid(True, alpha=0.3) plt.tight_layout() plt.savefig(save_path, dpi=150, bbox_inches='tight') plt.close(fig) plt.close('all') # Ensure all figures are closed print(f" Saved: {os.path.basename(save_path)}") # ============================================================================ # MAIN PROCESSING # ============================================================================ def process_case(case_name, models_to_process=None, plot_type='all'): """ Process a single case, comparing all models against sonde observations. """ print(f"\n{'='*70}") print(f"Processing case: {case_name}") print(f"{'='*70}") # Create output directory case_output_dir = os.path.join(OUTPUT_DIR, case_name) os.makedirs(case_output_dir, exist_ok=True) # Load sonde data ds_sonde = load_sonde_data(case_name) # Get case configuration case_cfg = CASE_CONFIG[case_name] valid_sites = case_cfg['valid_sites'] # Get site locations from sonde data site_lats = ds_sonde['Latitude_mean'].values site_lons = ds_sonde['Longitude_mean'].values # Debug: print shape print(f" Site lat/lon shapes: {site_lats.shape}, {site_lons.shape}") # Determine which models to process if models_to_process is None: models_to_process = list(MODEL_CONFIG.keys()) # Process each valid site for site_idx in valid_sites: print(f"\n--- Site {site_idx} ---") # Handle different possible array shapes for lat/lon # Also properly mask fill values (-9999) if site_lats.ndim == 1: site_lat = float(site_lats[site_idx]) site_lon = float(site_lons[site_idx]) elif site_lats.ndim == 2: # Take mean across time dimension if 2D (Site, Time) # But first mask fill values lat_vals = site_lats[site_idx, :] lon_vals = site_lons[site_idx, :] # Mask fill values (-9999 and other invalid values) lat_vals = np.where(lat_vals <= -9998, np.nan, lat_vals) lon_vals = np.where(lon_vals <= -9998, np.nan, lon_vals) lat_vals = np.where(lat_vals < -900, np.nan, lat_vals) lon_vals = np.where(lon_vals < -900, np.nan, lon_vals) # Also check for physically unreasonable values lat_vals = np.where((lat_vals < 25) | (lat_vals > 35), np.nan, lat_vals) lon_vals = np.where((lon_vals < -100) | (lon_vals > -90), np.nan, lon_vals) site_lat = float(np.nanmean(lat_vals)) site_lon = float(np.nanmean(lon_vals)) else: # For higher dimensions, flatten and take first valid value site_lat = float(np.nanmean(site_lats.flat[site_idx])) site_lon = float(np.nanmean(site_lons.flat[site_idx])) # Check if coordinates are valid (within Houston area: ~28-31°N, ~94-97°W) coords_valid = (not np.isnan(site_lat) and not np.isnan(site_lon) and 28 < site_lat < 31 and -97 < site_lon < -94) # Use known location if file coordinates are invalid if not coords_valid: if site_idx in KNOWN_SITE_LOCATIONS: site_lat, site_lon = KNOWN_SITE_LOCATIONS[site_idx] print(f" Location: ({site_lat:.4f}, {site_lon:.4f}) [from KNOWN_SITE_LOCATIONS]") else: print(f" WARNING: Invalid coordinates and no known location for Site {site_idx}, skipping!") continue else: print(f" Location: ({site_lat:.4f}, {site_lon:.4f}) [from sonde file]") # Find valid time steps for this site # Use isel for integer-based indexing (works with h5py-constructed dataset) if 'Potential_temperature_mean' in ds_sonde: obs_theta_raw = ds_sonde['Potential_temperature_mean'].values # Handle different dimension orders if obs_theta_raw.ndim == 3: # Assume (Site, Time, Height) order obs_theta_site = obs_theta_raw[site_idx, :, :] elif obs_theta_raw.ndim == 2: obs_theta_site = obs_theta_raw else: print(f" Unexpected theta shape: {obs_theta_raw.shape}") continue # Mask invalid values obs_theta_site = np.where(obs_theta_site == -9999, np.nan, obs_theta_site) # Find times with valid data (not all NaN across heights) valid_time_mask = ~np.all(np.isnan(obs_theta_site), axis=1) all_valid_time_indices = np.where(valid_time_mask)[0] else: print(f" Potential_temperature_mean not found in sonde data") continue # Filter time indices to only include times within the case window # Build datetime for each time index and filter # Note: Sonde data uses Hour_LT (Local Time), need to convert to UTC valid_time_indices = [] start_dt = case_cfg['start_datetime'] # Now in UTC end_dt = case_cfg['end_datetime'] # Now in UTC for time_idx in all_valid_time_indices: try: if 'Year' in ds_sonde and 'Month' in ds_sonde and 'Day' in ds_sonde: year = int(ds_sonde['Year'].values[time_idx]) month = int(ds_sonde['Month'].values[time_idx]) day = int(ds_sonde['Day'].values[time_idx]) if 'Hour_LT' in ds_sonde: hour_lt = int(ds_sonde['Hour_LT'].values[time_idx]) else: hour_lt = 12 # Convert Local Time to UTC (add UTC_OFFSET_HOURS) obs_time_lt = datetime(year, month, day, hour_lt, 0) obs_time_utc = obs_time_lt + timedelta(hours=UTC_OFFSET_HOURS) # Check if within case window (now comparing UTC to UTC) if start_dt <= obs_time_utc <= end_dt: valid_time_indices.append(time_idx) except: continue if len(valid_time_indices) == 0: print(f" No valid observation times within case window, skipping site") continue print(f" Found {len(valid_time_indices)} valid time steps within case window") # Process each valid time for time_idx in valid_time_indices: # Get observation height - the variable is Height_AGL (Above Ground Level) obs_height = None # Priority order for height variables - Height_AGL is the correct one for this dataset height_vars = ['Height_AGL', 'Height', 'Height_mean', 'height', 'height_m', 'altitude', 'z'] for hvar in height_vars: if hvar in ds_sonde: obs_height = ds_sonde[hvar].values # If it's 2D or 3D, extract the right slice if obs_height.ndim > 1: # Try to get just the height levels obs_height = obs_height.flatten()[:133] if obs_height.size >= 133 else obs_height.flatten() if time_idx == valid_time_indices[0]: print(f" Using height variable: {hvar}, range: {np.nanmin(obs_height):.1f} - {np.nanmax(obs_height):.1f} m") break # If no height variable found, use fallback if obs_height is None: obs_height = np.arange(133) * 30.0 # Fallback: 30m spacing if time_idx == valid_time_indices[0]: print(f" No height variable found, using fallback: 0 - {np.max(obs_height):.1f} m") obs_theta_profile = obs_theta_site[time_idx, :] # Get RH data if 'Relative_humidity_mean' in ds_sonde: obs_rh_raw = ds_sonde['Relative_humidity_mean'].values if obs_rh_raw.ndim == 3: obs_rh_site = obs_rh_raw[site_idx, :, :] else: obs_rh_site = obs_rh_raw obs_rh_site = np.where(obs_rh_site == -9999, np.nan, obs_rh_site) obs_rh_profile = obs_rh_site[time_idx, :] else: obs_rh_profile = np.full_like(obs_theta_profile, np.nan) # Get time for this observation # Construct datetime from Year, Month, Day, Hour_LT, Minute and convert to UTC try: if 'Year' in ds_sonde and 'Month' in ds_sonde and 'Day' in ds_sonde: year = int(ds_sonde['Year'].values[time_idx]) month = int(ds_sonde['Month'].values[time_idx]) day = int(ds_sonde['Day'].values[time_idx]) if 'Hour_LT' in ds_sonde: hour_lt = int(ds_sonde['Hour_LT'].values[time_idx]) else: hour_lt = 12 # Include minutes if available if 'Minute' in ds_sonde: minute = int(ds_sonde['Minute'].values[time_idx]) else: minute = 0 # Create local time then convert to UTC local_time = datetime(year, month, day, hour_lt, minute) target_time_utc = local_time + timedelta(hours=UTC_OFFSET_HOURS) else: target_time_utc = case_cfg['start_datetime'] + timedelta(hours=time_idx) except: target_time_utc = case_cfg['start_datetime'] + timedelta(hours=time_idx) # target_time_utc is now in UTC - use directly for model file matching time_str = target_time_utc.strftime("%m%d_%H%M") print(f"\n Time: {target_time_utc.strftime('%Y-%m-%d %H:%M')} UTC") # Collect observation data obs_data = { 'height': obs_height, 'theta': obs_theta_profile, 'rh': obs_rh_profile } # Collect model data model_data_dict = {} for model_name in models_to_process: config = MODEL_CONFIG[model_name] data_path = config.get(case_name) if data_path is None or not os.path.exists(data_path): print(f" {model_name}: Path not found") continue print(f" Reading {model_name}...", end=' ') sys.stdout.flush() # Use UTC time for model file matching theta, rh, height = read_model_data( model_name, data_path, target_time_utc, site_lat, site_lon ) if theta is not None: # Interpolate to common height levels theta_interp = interpolate_to_heights(theta, height, TARGET_HEIGHTS) rh_interp = interpolate_to_heights(rh, height, TARGET_HEIGHTS) model_data_dict[model_name] = { 'height': TARGET_HEIGHTS, 'theta': theta_interp, 'rh': rh_interp } print("OK") # Clean up original arrays to free memory del theta, rh, height else: print("Failed") # Force garbage collection after each model gc.collect() # Generate plots if plot_type in ['theta', 'all']: save_path = os.path.join(case_output_dir, f'multimodel_theta_site{site_idx}_{time_str}.png') plot_multimodel_comparison(obs_data, model_data_dict, 'theta', site_idx, target_time_utc, case_name, save_path) if plot_type in ['rh', 'all']: save_path = os.path.join(case_output_dir, f'multimodel_rh_site{site_idx}_{time_str}.png') plot_multimodel_comparison(obs_data, model_data_dict, 'rh', site_idx, target_time_utc, case_name, save_path) # Clean up to free memory before next time step del model_data_dict, obs_data gc.collect() print(f"\nCase {case_name} complete!") print(f"Output saved to: {case_output_dir}") # ============================================================================ # MAIN # ============================================================================ def main(): parser = argparse.ArgumentParser( description='Multi-model comparison with ARM sonde observations' ) parser.add_argument('--case', type=str, default='both', choices=['june17', 'aug7', 'both'], help='Case to process') parser.add_argument('--models', type=str, nargs='+', default=None, help='Models to include (default: all)') parser.add_argument('--skip-models', type=str, nargs='+', default=None, help='Models to skip (e.g., --skip-models E3SM_THREAD CM1)') parser.add_argument('--plot', type=str, default='all', choices=['theta', 'rh', 'all'], help='Variables to plot') args = parser.parse_args() print("="*70) print("Multi-Model Sonde Comparison for TRACER-MIP") print("="*70) print(f"Output directory: {OUTPUT_DIR}") # Create output directory os.makedirs(OUTPUT_DIR, exist_ok=True) # Determine models to process models_to_use = args.models if args.skip_models and models_to_use is None: # Start with all models and remove skipped ones models_to_use = [m for m in MODEL_CONFIG.keys() if m not in args.skip_models] print(f"Skipping models: {args.skip_models}") print(f"Processing models: {models_to_use}") # Determine cases to process if args.case == 'both': cases = ['june17', 'aug7'] else: cases = [args.case] # Process each case for case_name in cases: process_case(case_name, models_to_process=models_to_use, plot_type=args.plot) print("\n" + "="*70) print("All processing complete!") print("="*70) if __name__ == '__main__': main()