#!/usr/bin/env python3 """ TRACER-MIP: Extract HOURLY Precipitation (E3SM_THREAD / SCREAMv1) ================================================================= E3SM SCREAMv1 / THREAD specifics -------------------------------- - File FORMAT: NETCDF3_64BIT_DATA (CDF5), NOT netCDF4-on-HDF5. Must use the netCDF4 reader; h5py cannot read CDF5. - File LAYOUT: one variable per file, hourly filenames, 30 timesteps each at 2-min cadence: {var}_F2010-SCREAMv1.TRACERmipx32subx3v1pg2.001.2D.INST.INSTANT.nmins_x2.YYYY-MM-DD-SSSSS.nc SSSSS = seconds-of-day at start of hour (00000, 03600, ...) 48 files per variable per scenario (2 days x 24 hours). - DIMENSIONS: (time=30, lat=502, lon=583) for the 2D fields. time encoded as "days since 2022-{06,08}-16 00:00:00", noleap. - SCENARIO MAPPING: ctrl -> RefCCN 3x -> HighCCN 0.3x -> LowCCN - FILL VALUES: SCREAM 2D fields may carry _FillValue ~3.4e+33. netCDF4 auto-masks, but np.asarray() loses the mask, so we apply an explicit |x| > 1e30 -> NaN guard after every read. Precipitation variable ----------------------- precip_total_surf_mass_flux : instantaneous surface precipitation rate, units m s-1 (water-equivalent). Hourly accumulation method -------------------------- The native output is a 2-min INSTANTANEOUS rate (30 samples per clock hour). For each clock hour we average the available 2-min rate samples that fall in [HH:00, HH+1:00), then convert to an hourly depth: precip_hour [mm] = mean(rate over the hour) [m s-1] * 3600 s * 1000 mm/m This mirrors the rate-integration convention used by the other rate-based models (CM1, RAMS) in this project. The hourly mean of 2-min instantaneous samples is the trapezoid/midpoint estimate of the hourly accumulation; the (small) sampling error is recorded in the output file attributes as a known trade-off. Output schema: matches the other TRACER-MIP extraction outputs (h5py file, unlimited Time axis in "hours since 2022-01-01 00:00:00", 1D lat/lon promoted to 2D XLAT/XLONG). Output directory: /global/cfs/projectdirs/m4486/Haochen/Precipitation/E3SM_THREAD/{June17,Aug7}/ Output naming: E3SM_THREAD_{scenario}_PRECIP_HOURLY.nc Usage ----- python Extract_Precip_E3SM.py # both cases, all scenarios python Extract_Precip_E3SM.py --case june17 # single case python Extract_Precip_E3SM.py --scenario RefCCN # single scenario python Extract_Precip_E3SM.py --case aug7 --scenario HighCCN python Extract_Precip_E3SM.py --dryrun # list files, no write python Extract_Precip_E3SM.py --force # overwrite existing """ import os os.environ.setdefault("HDF5_USE_FILE_LOCKING", "FALSE") import sys import re import gc import glob import argparse from datetime import datetime, timedelta import numpy as np import h5py from netCDF4 import Dataset as NC4Dataset import warnings warnings.filterwarnings("ignore") # ============================================================================ # Configuration # ============================================================================ MODEL_NAME = "E3SM_THREAD" BASE_PATH = "/global/cfs/projectdirs/m4486/Haochen/TRACER_Model/E3SM_THREAD" OUTPUT_BASE = "/global/cfs/projectdirs/m4486/Haochen/Precipitation/E3SM_THREAD" # Precipitation variable in the E3SM 2D files (water-equivalent, m s-1) PRECIP_VAR = "precip_total_surf_mass_flux" PRECIP_FILE_PREFIX = "precip_total_surf_mass_flux" # Disk scenario folder -> project-standard scenario name SCENARIO_CONFIG = { "ctrl": "RefCCN", "3x": "HighCCN", "0.3x": "LowCCN", } # Case windows. Other TRACER-MIP extractions use the 12Z->06Z analysis # window; the precip rate is integrated only over hours inside the window. CASE_CONFIG = { "june17": { "date_strs": ["2022-06-17", "2022-06-18"], "folder_suffix": "20220617", "start_dt": datetime(2022, 6, 17, 12, 0), "end_dt": datetime(2022, 6, 18, 6, 0), "output_subdir": "June17", }, "aug7": { "date_strs": ["2022-08-07", "2022-08-08"], "folder_suffix": "20220807", "start_dt": datetime(2022, 8, 7, 12, 0), "end_dt": datetime(2022, 8, 8, 6, 0), "output_subdir": "Aug7", }, } DT_INTERNAL = 120 # native 2-min cadence (seconds) SECONDS_PER_HR = 3600 M_TO_MM = 1000.0 FILL_THRESHOLD = 1e30 # |x| beyond this is treated as fill -> NaN REF_TIME = datetime(2022, 1, 1, 0, 0) # Time-axis origin # {var}_..._YYYY-MM-DD-SSSSS.nc -> capture date + seconds-of-day E3SM_FILE_RE = re.compile(r"\.(\d{4})-(\d{2})-(\d{2})-(\d+)\.nc$") # ============================================================================ # Helpers # ============================================================================ def get_scenario_dir(case, scenario_folder): """Disk path for a (case, disk-scenario-folder) pair.""" suffix = CASE_CONFIG[case]["folder_suffix"] return os.path.join(BASE_PATH, f"{suffix}_{scenario_folder}") def parse_e3sm_filename(fp): """Return datetime of the file's first timestep (start-of-hour).""" m = E3SM_FILE_RE.search(os.path.basename(fp)) if not m: return None yyyy, mm, dd, secs = m.groups() base = datetime(int(yyyy), int(mm), int(dd)) return base + timedelta(seconds=int(secs)) def find_precip_files(scenario_dir, date_strs): """Sorted list of hourly precip 2D files covering the requested dates.""" files = [] for ds in date_strs: pat = os.path.join( scenario_dir, f"{PRECIP_FILE_PREFIX}_*2D*{ds}*.nc") files.extend(glob.glob(pat)) # Fallback if the 2D token is absent from the name if not files: for ds in date_strs: pat = os.path.join(scenario_dir, f"{PRECIP_FILE_PREFIX}_*{ds}*.nc") files.extend(glob.glob(pat)) files = sorted(set(files)) dated = [] for f in files: t0 = parse_e3sm_filename(f) if t0 is not None: dated.append((t0, f)) dated.sort(key=lambda x: x[0]) return dated def read_rate_masked(nc, varname): """ Read the full (time, lat, lon) rate array with fill-value masking. netCDF4 auto-masks _FillValue, but converting to a plain ndarray drops the mask, so we additionally NaN out |x| > FILL_THRESHOLD and any non-finite values. """ if varname not in nc.variables: return None d = nc.variables[varname][:] if hasattr(d, "filled"): arr = d.filled(np.nan).astype(np.float32) else: arr = np.asarray(d, dtype=np.float32) arr = np.where(np.isfinite(arr), arr, np.nan) arr = np.where(np.abs(arr) > FILL_THRESHOLD, np.nan, arr) return arr def read_grid(nc): """Return 1D lat, lon arrays from an E3SM file.""" lat = np.array(nc.variables["lat"][:], dtype=np.float32) lon = np.array(nc.variables["lon"][:], dtype=np.float32) return lat, lon def file_time_axis(nc, fallback_start): """ Datetimes for each of the (up to 30) internal timesteps. E3SM 'time' is "days since 2022-MM-16 00:00:00" (noleap). We parse the units string when present and fall back to fallback_start + k*120 s. """ n = nc.variables[PRECIP_VAR].shape[0] try: tvar = nc.variables["time"] units = getattr(tvar, "units", "") m = re.search(r"days since (\d{4})-(\d{2})-(\d{2})", units) if m: origin = datetime(int(m.group(1)), int(m.group(2)), int(m.group(3))) vals = np.array(tvar[:], dtype="float64") return [origin + timedelta(days=float(v)) for v in vals] except Exception: pass return [fallback_start + timedelta(seconds=k * DT_INTERNAL) for k in range(n)] # ============================================================================ # Output writer (h5py, matches other TRACER-MIP outputs) # ============================================================================ def make_output_path(case, scenario): sub = CASE_CONFIG[case]["output_subdir"] out_dir = os.path.join(OUTPUT_BASE, sub) os.makedirs(out_dir, exist_ok=True) return os.path.join(out_dir, f"{MODEL_NAME}_{scenario}_PRECIP_HOURLY.nc") def open_output(out_path, ny, nx, lat_2d, lon_2d, case, scenario): if os.path.exists(out_path): bak = out_path + ".bak" if os.path.exists(bak): os.remove(bak) os.rename(out_path, bak) print(f" Existing file backed up to {os.path.basename(bak)}") fh = h5py.File(out_path, "w") fh.create_dataset( "PRECIP_HOURLY", shape=(0, ny, nx), maxshape=(None, ny, nx), dtype="float32", chunks=(1, ny, nx), compression="gzip", compression_opts=4, fillvalue=np.nan) fh["PRECIP_HOURLY"].attrs["units"] = "mm" fh["PRECIP_HOURLY"].attrs["long_name"] = ( "Hourly accumulated surface precipitation (water equivalent)") fh.create_dataset("Time", shape=(0,), maxshape=(None,), dtype="float64") fh["Time"].attrs["units"] = "hours since 2022-01-01 00:00:00" fh["Time"].attrs["calendar"] = "standard" fh["Time"].attrs["long_name"] = ( "Valid time at the END of each accumulation hour") fh.create_dataset("XLAT", data=lat_2d.astype("float32")) fh["XLAT"].attrs["units"] = "degrees_north" fh.create_dataset("XLONG", data=lon_2d.astype("float32")) fh["XLONG"].attrs["units"] = "degrees_east" fh.attrs["model"] = MODEL_NAME fh.attrs["case"] = case fh.attrs["scenario"] = scenario fh.attrs["variable"] = "PRECIP_HOURLY" fh.attrs["source_var"] = PRECIP_VAR fh.attrs["created"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S") fh.attrs["method"] = ( "Hourly mean of 2-min instantaneous surface precip rate " "(m/s) * 3600 s * 1000 mm/m -> mm/hr accumulation") fh.attrs["accumulation_note"] = ( "Native field is a 2-min INSTANTANEOUS rate; the hourly value is " "the mean of the in-hour samples integrated over 3600 s. Small " "sampling error vs. a true model-integrated bucket is a known, " "accepted trade-off, consistent with the CM1/RAMS rate handling.") fh.attrs["time_convention"] = ( "Time = END of the accumulation hour (e.g. value at 13:00 covers " "12:00:00 - 13:00:00).") fh.attrs["fill_mask"] = f"|x| > {FILL_THRESHOLD:g} set to NaN" return fh def append_timestep(fh, depth_2d, t_hours): cur = fh["PRECIP_HOURLY"].shape[0] fh["PRECIP_HOURLY"].resize(cur + 1, axis=0) fh["Time"].resize(cur + 1, axis=0) fh["PRECIP_HOURLY"][cur] = depth_2d fh["Time"][cur] = t_hours # ============================================================================ # Per (case, scenario) extraction # ============================================================================ def extract_one(case, scenario_folder, force=False, dryrun=False): scenario = SCENARIO_CONFIG[scenario_folder] cfg = CASE_CONFIG[case] sdir = get_scenario_dir(case, scenario_folder) start_dt = cfg["start_dt"] end_dt = cfg["end_dt"] print(f"\n{'-'*70}") print(f" {MODEL_NAME} / {case} / {scenario_folder} -> {scenario}") print(f" Path: {sdir}") print(f" Window: {start_dt} -> {end_dt}") print(f"{'-'*70}") if not os.path.isdir(sdir): print(f" ERROR: scenario directory not found: {sdir}") return False out_path = make_output_path(case, scenario) if os.path.exists(out_path) and not force and not dryrun: print(f" Output exists, skipping (use --force to overwrite):") print(f" {out_path}") return True dated_files = find_precip_files(sdir, cfg["date_strs"]) if not dated_files: print(f" ERROR: no '{PRECIP_FILE_PREFIX}' files found in {sdir}") return False print(f" Found {len(dated_files)} hourly precip files") if dryrun: for t0, f in dated_files[:5]: print(f" {t0} {os.path.basename(f)}") if len(dated_files) > 5: print(f" ... ({len(dated_files) - 5} more)") return True # ---------------------------------------------------------------- # Accumulate 2-min instantaneous samples into clock-hour buckets. # bucket key = datetime at the START of the clock hour. # ---------------------------------------------------------------- sums = {} # hour_start -> running sum of rate (m/s) counts = {} # hour_start -> number of samples lat_1d = lon_1d = None for fi, (t0, fpath) in enumerate(dated_files): try: nc = NC4Dataset(fpath, "r") except Exception as e: print(f" SKIP (cannot open) {os.path.basename(fpath)}: {e}") continue try: if lat_1d is None: lat_1d, lon_1d = read_grid(nc) rate = read_rate_masked(nc, PRECIP_VAR) # (nt, ny, nx) m/s if rate is None: print(f" SKIP (no {PRECIP_VAR}) " f"{os.path.basename(fpath)}") nc.close() continue times = file_time_axis(nc, t0) nt = min(len(times), rate.shape[0]) for k in range(nt): tk = times[k] # Keep only samples inside the analysis window if tk < start_dt or tk >= end_dt: continue hour_start = tk.replace(minute=0, second=0, microsecond=0) slc = rate[k] if hour_start not in sums: sums[hour_start] = np.zeros_like(slc, dtype=np.float64) counts[hour_start] = np.zeros(slc.shape, dtype=np.int32) valid = np.isfinite(slc) sums[hour_start][valid] += slc[valid] counts[hour_start][valid] += 1 nc.close() del rate except Exception as e: print(f" ERROR on {os.path.basename(fpath)}: {e}") try: nc.close() except Exception: pass continue if fi % 6 == 0: gc.collect() if not sums: print(f" WARNING: no in-window samples accumulated; nothing written") return False # ---------------------------------------------------------------- # Convert each hour's mean rate to an hourly depth and write. # Output Time is the END of the accumulation hour. # ---------------------------------------------------------------- lon_2d, lat_2d = np.meshgrid(lon_1d, lat_1d) ny, nx = lat_2d.shape fh = open_output(out_path, ny, nx, lat_2d, lon_2d, case, scenario) n_written = 0 for hour_start in sorted(sums.keys()): cnt = counts[hour_start] with np.errstate(invalid="ignore", divide="ignore"): mean_rate = np.where(cnt > 0, sums[hour_start] / np.maximum(cnt, 1), np.nan) depth_mm = (mean_rate * SECONDS_PER_HR * M_TO_MM).astype(np.float32) hour_end = hour_start + timedelta(hours=1) t_hours = (hour_end - REF_TIME).total_seconds() / 3600.0 append_timestep(fh, depth_mm, t_hours) n_written += 1 mx = np.nanmax(depth_mm) if np.isfinite(depth_mm).any() else 0.0 print(f" {hour_start:%Y-%m-%d %H}:00 -> " f"{hour_end:%H}:00 max={mx:6.2f} mm " f"(samples/cell up to {int(cnt.max()) if cnt.size else 0})") fh.attrs["n_hours"] = n_written fh.close() size_mb = os.path.getsize(out_path) / (1024.0 ** 2) print(f" DONE: wrote {n_written} hourly steps -> " f"{out_path} ({size_mb:.2f} MB)") gc.collect() return True # ============================================================================ # CLI # ============================================================================ def main(): p = argparse.ArgumentParser( description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) p.add_argument("--case", choices=["june17", "aug7", "both"], default="both") p.add_argument("--scenario", choices=["RefCCN", "HighCCN", "LowCCN", "all"], default="all") p.add_argument("--force", action="store_true", help="overwrite existing output files") p.add_argument("--dryrun", action="store_true", help="list input files only, do not write") args = p.parse_args() cases = ["june17", "aug7"] if args.case == "both" else [args.case] # Map requested scenario name back to disk folder(s) name_to_folder = {v: k for k, v in SCENARIO_CONFIG.items()} if args.scenario == "all": scen_folders = list(SCENARIO_CONFIG.keys()) else: scen_folders = [name_to_folder[args.scenario]] print("=" * 70) print(f"{MODEL_NAME} Hourly Precipitation Extraction") print(f" Source var: {PRECIP_VAR} (m/s, instantaneous)") print(f" Cases: {cases}") print(f" Scenarios: " f"{[SCENARIO_CONFIG[s] for s in scen_folders]}") print(f" Output base: {OUTPUT_BASE}") print(f" Start time: {datetime.now():%Y-%m-%d %H:%M:%S}") print("=" * 70) n_ok = n_fail = 0 for case in cases: for sf in scen_folders: try: ok = extract_one(case, sf, force=args.force, dryrun=args.dryrun) except Exception as e: import traceback traceback.print_exc() ok = False if ok: n_ok += 1 else: n_fail += 1 gc.collect() print("\n" + "=" * 70) print(f"FINISHED successful={n_ok} failed={n_fail}") print(f" End time: {datetime.now():%Y-%m-%d %H:%M:%S}") print("=" * 70) sys.exit(0 if n_fail == 0 else 1) if __name__ == "__main__": main()