diff --git a/da_methods/.gitignore b/da_methods/.gitignore new file mode 100644 index 00000000..e90dda7d --- /dev/null +++ b/da_methods/.gitignore @@ -0,0 +1,10 @@ +*.png +!figures/**/*.png +!**/figures/**/*.png +*.pdf +*.pptx +__pycache__/ +~$* +poster_kriging_troute.md +poster_kriging_troute_v2.md +poster_f1_vs_f2.md diff --git a/da_methods/1_distributed_cfe/calibrate_catchment_cfe_da_v2.py b/da_methods/1_distributed_cfe/calibrate_catchment_cfe_da_v2.py new file mode 100644 index 00000000..f1268200 --- /dev/null +++ b/da_methods/1_distributed_cfe/calibrate_catchment_cfe_da_v2.py @@ -0,0 +1,866 @@ +""" +Per-catchment CFE test run with optional EnKF data assimilation. + +This script runs the test period (Oct 2023 – Oct 2024, includes Hurricane +Helene) on top of pre-calibrated CFE parameters, with optional state +assimilation against per-catchment kriging streamflow observations. + + The CFE calibration step (DDS against the 2020–2022 retro period) is NOT + part of this script. It is performed once, separately, via the calibrate-cfe + workflow at NWC-CUAHSI-Summer-Institute/calibrate-cfe (held-out gauges branch). + This script loads the resulting _best_params.json from --out-dir and + uses those parameters fixed through the entire test run. + +================================================================================ +ENSEMBLE KALMAN FILTER (true stochastic EnKF — Burgers/van Leeuwen/Evensen 1998) +================================================================================ +Each hour during the test period: + - N CFE members run in parallel with perturbed precip + PET + - forecast variance Pyy = ensemble variance of Q (no heuristic) + - Kalman gain VECTOR: one K per state, from cross-covariance Pxy(state, Q) + (no hardcoded state-split weights — the data decides each hour) + - observation perturbed N times: obs_i = obs + sqrt(R) × N(0,1) + - each member updated with its own obs_i and its own innovation + - per-state per-timestep process noise prevents ensemble collapse + - output time series is the ensemble mean + +Measured improvement on USGS gauge 03463300 (21 sub-catchments, Helene year): + mean Test KGE 0.692 (no DA) → 0.817 (with DA). + +================================================================================ +STATES UPDATED (4 states, applied per ensemble member with mass cascade) +================================================================================ + 1. soil_reservoir["storage_m"] via BMI: SOIL_CONCEPTUAL_STORAGE + 2. gw_reservoir["storage_m"] direct attribute (not exposed via BMI) + 3. nash_storage[0] direct attribute (upstream Nash bucket) + 4. nash_storage[1] direct attribute (feeds streamflow) + +Mass-conserving cascade — when a state hits its bound, the excess/deficit is +redirected along CFE's physical flow direction instead of being silently clipped: + Positive corrections (water added): + soil full → excess pushed to Nash[0] + GW full → excess pushed to Nash[1] + Negative corrections (water removed): + Nash[1] < 0 → deficit absorbed from Nash[0] + Nash[0] < 0 → deficit absorbed from soil + soil < 0 → deficit absorbed from GW + GW < 0 → accept loss (logged in mm via total_overflow_lost_mm) + +================================================================================ +INPUTS +================================================================================ +Forcing (training): per-catchment NWM retro CSV + columns: time, APCP_surface, DSWRF_surface, TMP_2maboveground, ... +Forcing (testing): per-catchment NWM operational CSVs (two dirs concatenated) + APCP_surface in kg/m²/s → converted to mm/h inside +Obs (kriging): per-catchment CSV. Three column layouts auto-handled: + a) datetime, qkrig (format A) + b) datetime, qkrig, variance (format A with-variance) + c) time, qkrig_mm_hr (format B standard) + d) time, qkrig_mm_hr, qkrig_variance (format B with-variance) + If a per-hour variance column is present it is used as R; + otherwise R = --enkf-obs-error-std^2 is used as fallback. + +Time splits: + Spinup (cal): 2019-01-01 00:00 → 2019-12-31 23:00 + Calibration: 2020-01-01 00:00 → 2022-12-31 23:00 + Spinup (test): 2023-02-01 00:00 → 2023-09-30 23:00 + Test (Helene): 2023-10-01 00:00 → 2024-10-31 23:00 + +================================================================================ +KEY CLI FLAGS +================================================================================ + --enkf-enabled Turn DA on + --enkf-members 20 N ensemble members (>=2 required) + --enkf-obs-error-std 0.05 Fallback obs std (mm/h) when no variance column + and --no-vrugt-r is set + +OBSERVATION ERROR VARIANCE R + Production default: Vrugt et al. 2005 (SODA paper) heteroscedastic R + scaled by the kriging variance: + R(t) = (alpha * y_obs(t))^2 + scale * sigma^2_krig(t) + Flow-magnitude term keeps R small at low flow (so DA fires) and larger + at peaks (where obs is also more uncertain). Kriging variance brings + per-hour catchment-specific info in, scaled so it does not dominate. + Defaults: alpha=0.10 (10% relative error per Vrugt), scale=0.001. + + Pass --no-vrugt-r to revert to the raw kriging variance behavior (or + --enkf-obs-error-std^2 fallback when no variance column). Not + recommended for this basin — raw kriging variance is 8-10x the obs + magnitude, which collapses the Kalman gain. + +Test-loop perturbation defaults (see EnKFAssimilator.__init__): + precip: lognormal multiplier with sigma=0.15, mean=1 + (literature standard; guarantees non-negative) + PET: ×N(1, 0.10), clipped at 0 + initial state: ×N(1, 0.05) on soil/GW; small additive jitter on Nash + observation: additive noise with std sqrt(R) (Burgers/Evensen) + process noise per state (per-hour multiplicative; prevents ensemble + collapse — required by textbook EnKF): + soil : 0.002 (0.2%) slow-evolving, small noise sufficient + GW : 0.0015 (0.15%) slowest state + Nash : 0.005 (0.5%) + additive floor; needs more because nash + spread does not develop naturally + Magnitudes match NWC-CUAHSI/data_assimilation_with_bmi reference impl. + +================================================================================ +USAGE +================================================================================ +Production run (true EnKF, N=20, Vrugt R on by default): + python3 calibrate_catchment_cfe_da_v2.py \\ + --cat-id cat-1016300 \\ + --forcing-dir /mnt/disk2/suma_helen_poster/nwm_retro_catchment_forcings \\ + --obs-dir /mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance \\ + --cfe-dir /mnt/disk2/suma_helen_poster/cfe_py \\ + --config-file /mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json \\ + --param-bounds /mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json \\ + --out-dir /mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt \\ + --test-forcing-dir1 /mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings \\ + --test-forcing-dir2 /mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings \\ + --enkf-enabled --enkf-members 20 + + Use the with-variance obs dir so the per-hour kriging variance enters R(t) + via the Vrugt + kriging-scaling formula. Pass --no-vrugt-r to disable the + Vrugt R and fall back to raw kriging variance (not recommended). + +NOTE: pre-stage the calibrated _best_params.json into the out-dir +before running: + mkdir -p / + cp //_best_params.json // + +To run all 21 catchments, launch in batches of ~5 (each catchment uses N CFE +instances simultaneously; 21 × 20 = 420 processes if run all at once). +""" + +import argparse +import os +import sys +import json +import numpy as np +import pandas as pd +import matplotlib +matplotlib.use('Agg') +import matplotlib.pyplot as plt +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +TIME_SPLIT = { + "spinup-for-calibration": { + "start": "2019-01-01 00:00:00", + "end": "2019-12-31 23:00:00", + }, + "calibration": { + "start": "2020-01-01 00:00:00", + "end": "2022-12-31 23:00:00", + }, + "spinup-for-testing": { + "start": "2023-02-01 00:00:00", + "end": "2023-09-30 23:00:00", + }, + "testing": { + "start": "2023-10-01 00:00:00", + "end": "2024-10-31 23:00:00", + }, +} + +# Set at runtime +CAT_ID = None +FORCING_FILE = None +TEST_FORCING_FILE = None +OBS_FILE = None +CFE_CONFIG_FILE = None +PARAM_BOUNDS_FILE = None +OUT_DIR = None +bmi_cfe = None +ENKF_ENABLED = False +ENKF_CONFIG = {} + + +class EnKFAssimilator: + """ + True stochastic Ensemble Kalman Filter (Burgers / van Leeuwen / Evensen 1998) + for CFE state updates from per-catchment kriging observations. + + `update_states(models, ...)` runs the N-member ensemble update: + Kalman gain per state derived from cross-covariance Pxy / (Pyy + R), + observations perturbed per member, each member updated independently. + + Updates 4 CFE states per member (soil, GW, Nash[0], Nash[1]) with a + mass-conserving overflow/underflow cascade so corrections are redirected + along CFE's flow path instead of being silently clipped. + """ + + def __init__(self, n_members=20, obs_error_std=0.05, obs_file=None, + precip_perturb_frac=0.15, pet_perturb_frac=0.10, + init_state_perturb_frac=0.05, + soil_process_noise_frac=0.002, + gw_process_noise_frac=0.0015, + nash_process_noise_frac=0.005, + use_vrugt_r=True, vrugt_alpha=0.10, vrugt_scale=0.001, + rng_seed=None): + """ + Initialize EnKF assimilator. + + Args: + n_members (int): Number of ensemble members (>=2 for true EnKF) + obs_error_std (float): Default obs std dev when CSV lacks 'variance' + obs_file (str): Path to kriging observations CSV + precip_perturb_frac (float): Lognormal sigma for precip noise (mean=1) + pet_perturb_frac (float): Multiplicative Gaussian noise std for PET + init_state_perturb_frac (float): Multiplicative noise std on init states + soil_process_noise_frac (float): Per-hour multiplicative process noise std + for soil reservoir. Default 0.002 (0.2%) — soil is slow-evolving so a + small noise is sufficient and avoids cumulative-random-walk drift. + Matches NWC-CUAHSI reference implementation. + gw_process_noise_frac (float): Per-hour multiplicative process noise std + for GW reservoir. Default 0.0015 (0.15%) — slowest state, smallest noise. + nash_process_noise_frac (float): Per-hour multiplicative process noise std + for Nash[0] and Nash[1]. Default 0.005 (0.5%) — needed larger than + soil/GW because Nash spread does not develop naturally (buckets start + at 0 in CFE init). Combined with a small additive floor so noise + survives when the bucket is empty. + use_vrugt_r (bool): If True (default, production setting), compute + observation error variance R as a Vrugt 2005 (SODA) heteroscedastic + function of flow magnitude, scaled by the kriging variance: + R(t) = (alpha * y_obs)^2 + scale * sigma^2_krig. Fixes the case + where raw kriging variance is too large to allow DA to fire, and + gives the most uniform performance across catchments. Set to False + to revert to the raw kriging variance / fallback std behavior. + vrugt_alpha (float): Relative-error fraction in the Vrugt R. Default 0.10 + (10% of flow). Standard hydrology DA value (Vrugt et al. 2005). + vrugt_scale (float): Scaling on the kriging-variance term. Default 0.001. + Picked so the kriging contribution is the same order of magnitude as + the flow-magnitude term at typical flows for this basin. + rng_seed (int or None): Seed for the ensemble random generator + """ + self.n_members = n_members + self.obs_error_std = obs_error_std + self.obs_error_var = obs_error_std ** 2 + self.precip_perturb_frac = precip_perturb_frac + self.pet_perturb_frac = pet_perturb_frac + self.init_state_perturb_frac = init_state_perturb_frac + self.soil_process_noise_frac = soil_process_noise_frac + self.gw_process_noise_frac = gw_process_noise_frac + self.nash_process_noise_frac = nash_process_noise_frac + self.use_vrugt_r = use_vrugt_r + self.vrugt_alpha = vrugt_alpha + self.vrugt_scale = vrugt_scale + self.rng = np.random.default_rng(rng_seed) + + # Load kriging observations + # Handles three obs-CSV layouts: + # Format A: columns = 'datetime', 'qkrig', optional 'variance' + # Format B (no variance): columns = 'time', 'qkrig_mm_hr' + # Format B with-variance: columns = 'time', 'qkrig_mm_hr', 'qkrig_variance' + self.obs_dict = {} + self.obs_var_dict = {} + if obs_file and os.path.exists(obs_file): + obs_df = pd.read_csv(obs_file) + if 'time' in obs_df.columns and 'qkrig_mm_hr' in obs_df.columns: + obs_df = obs_df.rename(columns={'time': 'datetime', 'qkrig_mm_hr': 'qkrig'}) + if 'qkrig_variance' in obs_df.columns: + obs_df = obs_df.rename(columns={'qkrig_variance': 'variance'}) + obs_df['date'] = pd.to_datetime(obs_df['datetime']).dt.strftime('%Y-%m-%d %H:%M:%S') + has_variance = 'variance' in obs_df.columns + for _, row in obs_df.iterrows(): + y_obs = row['qkrig'] + krig_var = row['variance'] if has_variance else self.obs_error_var + self.obs_dict[row['date']] = y_obs + # Vrugt et al. 2005 (SODA) heteroscedastic R, scaled by kriging + # Vrugt 2005 heteroscedastic R scaled by kriging variance: + # R(t) = (alpha * y_obs(t))^2 + scale * sigma^2_krig(t) + # Flow-magnitude term makes R small at low flow (so DA fires), + # larger at peaks (where obs is also more uncertain). Kriging + # variance brings per-hour catchment-specific info in, scaled + # down so it doesn't dominate. + if self.use_vrugt_r and not pd.isna(y_obs): + self.obs_var_dict[row['date']] = ( + (self.vrugt_alpha * y_obs) ** 2 + + self.vrugt_scale * krig_var + ) + else: + self.obs_var_dict[row['date']] = krig_var + + self.n_updates = 0 + self.total_increment = 0.0 + self.total_overflow_lost_mm = 0.0 # water lost at GW underflow (mass-conservation breach) + + # True-EnKF diagnostics (test path) + self.avg_pyy = 0.0 + self.avg_K_sm = 0.0 + self.avg_K_gw = 0.0 + self.avg_K_n0 = 0.0 + self.avg_K_n1 = 0.0 + + # ------------------------------------------------------------------ + # Ensemble helpers (used by the test loop) + # ------------------------------------------------------------------ + def perturb_forcing(self, precip_mm_h, pet_mm_h): + """Return arrays of length n_members with perturbed (P, PET). + + Precip uses **lognormal** noise (literature standard for precip in + hydrology DA — Clark 2008, Renard 2010). Guarantees non-negative + precip and matches the heavy-tailed empirical distribution of + precip errors. Lognormal is parameterized so the multiplier has + E[noise] = 1 (mu = -sigma²/2). + + PET uses multiplicative Gaussian (more symmetric distribution). + """ + N = self.n_members + sigma_p = self.precip_perturb_frac + mu_p = -0.5 * sigma_p ** 2 # makes E[exp(mu + sigma·Z)] = 1 + p = precip_mm_h * self.rng.lognormal(mu_p, sigma_p, N) + # Lognormal × non-negative is non-negative; no clipping needed. + + e = pet_mm_h * (1.0 + self.pet_perturb_frac * self.rng.standard_normal(N)) + return p, np.maximum(e, 0.0) + + def perturb_initial_state(self, value, upper=None, floor=1e-6): + """Multiplicatively perturb a single initial-state scalar across N members.""" + v = value * (1.0 + self.init_state_perturb_frac * self.rng.standard_normal(self.n_members)) + v = np.maximum(v, floor) + if upper is not None: + v = np.minimum(v, upper) + return v + + def add_process_noise(self, models): + """Inject per-timestep state perturbations on every member. + + Process noise (the Q matrix in classical Kalman / "model error" term in + EnKF) is required by textbook EnKF to prevent ensemble collapse over + long runs. Without it, the analysis pulls all members toward similar + states, spread shrinks, Pyy → 0, and the filter goes blind. + + Each state gets its own σ (calibrated by physical response time): + soil : 0.2%/hr — slow-evolving; small noise prevents random-walk drift + GW : 0.15%/hr — slowest state + Nash : 0.5%/hr + small additive floor — needs more because Nash + spread does not develop naturally (buckets start at 0). + Magnitudes match NWC-CUAHSI/data_assimilation_with_bmi reference impl. + """ + sigma_soil = self.soil_process_noise_frac + sigma_gw = self.gw_process_noise_frac + sigma_nash = self.nash_process_noise_frac + if sigma_soil <= 0 and sigma_gw <= 0 and sigma_nash <= 0: + return + for m in models: + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + # Additive floor for nash so noise survives at zero. + nash_floor = max(sm_max * 1e-4, 1e-7) + if sigma_soil > 0: + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1.0 + sigma_soil * self.rng.standard_normal()), + 0.0, sm_max)) + if sigma_gw > 0: + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1.0 + sigma_gw * self.rng.standard_normal()), + 0.0, gw_max)) + if sigma_nash > 0: + m.nash_storage[0] = max( + n00 * (1.0 + sigma_nash * self.rng.standard_normal()) + + nash_floor * self.rng.standard_normal(), 0.0) + m.nash_storage[1] = max( + n10 * (1.0 + sigma_nash * self.rng.standard_normal()) + + nash_floor * self.rng.standard_normal(), 0.0) + + # ------------------------------------------------------------------ + # True stochastic EnKF (Burgers / van Leeuwen / Evensen 1998) + # Used by run_testing_period — replaces 30/15/20/35 + 0.3-heuristic with + # ensemble-derived Kalman gains. + # ------------------------------------------------------------------ + def update_states(self, models, current_date, forecast_runoffs_mm_h): + """ + Perform a true stochastic EnKF update across N ensemble members. + + Args: + models: list of N CFE BMI model instances (one per ensemble member) + current_date (str): 'YYYY-MM-DD HH:MM:SS' + forecast_runoffs_mm_h: array-like of N forecast runoffs (mm/h) + + Returns: + dict with update statistics (mean innovation, Pyy, K per state). + """ + stats = {"updated": False, "innovation_mean": 0.0, "pyy": 0.0, + "K_sm": 0.0, "K_gw": 0.0, "K_n0": 0.0, "K_n1": 0.0} + + if current_date not in self.obs_dict: + return stats + obs = self.obs_dict[current_date] + obs_var = self.obs_var_dict[current_date] + if pd.isna(obs): + return stats + + q = np.asarray(forecast_runoffs_mm_h, dtype=float) + if np.any(np.isnan(q)): + return stats + + N = self.n_members + if N < 2: + return stats # need ≥2 members for a meaningful covariance + + # ---- Collect ensemble of 4 states (meters) ---- + sm = np.array([m.get_value('SOIL_CONCEPTUAL_STORAGE') for m in models], dtype=float) + gw = np.array([m.gw_reservoir["storage_m"] for m in models], dtype=float) + n0 = np.array([float(m.nash_storage[0]) for m in models], dtype=float) + n1 = np.array([float(m.nash_storage[1]) for m in models], dtype=float) + + # ---- Ensemble means and anomalies ---- + sm_a = sm - sm.mean() + gw_a = gw - gw.mean() + n0_a = n0 - n0.mean() + n1_a = n1 - n1.mean() + q_mean = q.mean() + q_a = q - q_mean + + # Pyy = ensemble variance of Q (mm/h)^2 ; degenerate if all members agree + Pyy = float((q_a * q_a).sum() / (N - 1)) + denom = Pyy + obs_var + if denom < 1e-12: + return stats + + # Pxy = state/Q cross-covariance (units: m × mm/h) + Pxy_sm = float((sm_a * q_a).sum() / (N - 1)) + Pxy_gw = float((gw_a * q_a).sum() / (N - 1)) + Pxy_n0 = float((n0_a * q_a).sum() / (N - 1)) + Pxy_n1 = float((n1_a * q_a).sum() / (N - 1)) + + K_sm = Pxy_sm / denom # units: m / (mm/h) + K_gw = Pxy_gw / denom + K_n0 = Pxy_n0 / denom + K_n1 = Pxy_n1 / denom + + # Perturb the observation N times (Burgers/Evensen) + obs_pert = obs + np.sqrt(max(obs_var, 0.0)) * self.rng.standard_normal(N) + + # Per-member innovation (mm/h); state increments (m) + innov = obs_pert - q + sm_delta = K_sm * innov + gw_delta = K_gw * innov + n0_delta = K_n0 * innov + n1_delta = K_n1 * innov + + # ---- Apply per member with overflow-aware cascade ---- + for i, m in enumerate(models): + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + + sm_raw = sm[i] + sm_delta[i] + gw_raw = gw[i] + gw_delta[i] + n0_raw = n0[i] + n0_delta[i] + n1_raw = n1[i] + n1_delta[i] + + # Positive overflow cascade: soil→Nash[0], GW→Nash[1] + so = max(sm_raw - sm_max, 0.0); sm_raw -= so; n0_raw += so + go = max(gw_raw - gw_max, 0.0); gw_raw -= go; n1_raw += go + + # Negative underflow cascade: Nash[1]→Nash[0]→soil→GW→loss + n1u = max(-n1_raw, 0.0); n1_raw += n1u; n0_raw -= n1u + n0u = max(-n0_raw, 0.0); n0_raw += n0u; sm_raw -= n0u + smu = max(-sm_raw, 0.0); sm_raw += smu; gw_raw -= smu + gwu = max(-gw_raw, 0.0); gw_raw += gwu + self.total_overflow_lost_mm += gwu * 1000.0 + + if (np.isnan(sm_raw) or np.isnan(gw_raw) + or np.isnan(n0_raw) or np.isnan(n1_raw)): + continue + + sm_new = float(np.clip(sm_raw, 0.0, sm_max)) + gw_new = float(np.clip(gw_raw, 0.0, gw_max)) + n0_new = float(np.clip(n0_raw, 0.0, None)) + n1_new = float(np.clip(n1_raw, 0.0, None)) + + m.set_value('SOIL_CONCEPTUAL_STORAGE', sm_new) + m.gw_reservoir["storage_m"] = gw_new + m.nash_storage[0] = n0_new + m.nash_storage[1] = n1_new + + # ---- Diagnostics ---- + self.n_updates += 1 + innov_mean = float(innov.mean()) + self.total_increment += abs(innov_mean) # repurpose for ensemble: mean innov magnitude + # Running averages of Pyy and Kalman gains for end-of-run reporting + a = self.n_updates + self.avg_pyy = ((a - 1) * self.avg_pyy + Pyy) / a + self.avg_K_sm = ((a - 1) * self.avg_K_sm + K_sm) / a + self.avg_K_gw = ((a - 1) * self.avg_K_gw + K_gw) / a + self.avg_K_n0 = ((a - 1) * self.avg_K_n0 + K_n0) / a + self.avg_K_n1 = ((a - 1) * self.avg_K_n1 + K_n1) / a + + stats.update(updated=True, innovation_mean=innov_mean, pyy=Pyy, + K_sm=K_sm, K_gw=K_gw, K_n0=K_n0, K_n1=K_n1) + + # Log sparingly — every hour is noisy. Print at storms (large Pyy). + if Pyy > 1e-4 or self.n_updates % 200 == 0: + print(f" [EnKF N={N}] {current_date} | obs={obs:.3f} | q_mean={q_mean:.3f} | " + f"innov_mean={innov_mean:+.4f} | Pyy={Pyy:.6f} | " + f"K_sm={K_sm:+.3e} K_gw={K_gw:+.3e} K_n0={K_n0:+.3e} K_n1={K_n1:+.3e}") + + return stats + + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + """PT-PET from shortwave radiation + temperature. Returns mm/h.""" + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_forcing(path=None): + """Load NWM forcing CSV (retro or operational) and add PET column. Returns DataFrame.""" + path = path or FORCING_FILE + df = pd.read_csv(path) + df = df.rename(columns={"time": "date", "APCP_surface": "total_precipitation"}) + df['date'] = pd.to_datetime(df['date']).dt.strftime('%Y-%m-%d %H:%M:%S') + df['potential_evaporation'] = priestley_taylor_pet( + df['DSWRF_surface'].values, + df['TMP_2maboveground'].values + ) + return df + + +def load_test_forcing(): + """Load NWM operational forcing CSV. APCP_surface is in kg/m²/s → convert to mm/h.""" + df = pd.read_csv(TEST_FORCING_FILE) + df['date'] = pd.to_datetime(df['time']).dt.strftime('%Y-%m-%d %H:%M:%S') + df['total_precipitation'] = df['APCP_surface'] * 3600.0 # kg/m²/s → mm/h + df['potential_evaporation'] = priestley_taylor_pet( + df['DSWRF_surface'].values, + df['TMP_2maboveground'].values + ) + return df + + +def _kge_score(obs, sim): + """Kling-Gupta Efficiency on the overlap of obs and sim (NaNs skipped).""" + m = ~(np.isnan(obs) | np.isnan(sim)) + o, s = obs[m], sim[m] + if len(o) < 2: + return float('nan') + denom = float(np.sqrt(((o - o.mean()) ** 2).sum() * ((s - s.mean()) ** 2).sum())) + if denom == 0: + return float('nan') + r = float(((o - o.mean()) * (s - s.mean())).sum() / denom) + alpha = float(s.std() / o.std()) if o.std() != 0 else float('nan') + beta = float(s.mean() / o.mean()) if o.mean() != 0 else float('nan') + return 1.0 - float(np.sqrt((r - 1) ** 2 + (alpha - 1) ** 2 + (beta - 1) ** 2)) + + +def _nse_score(obs, sim): + """Nash-Sutcliffe Efficiency on the overlap of obs and sim (NaNs skipped).""" + m = ~(np.isnan(obs) | np.isnan(sim)) + o, s = obs[m], sim[m] + if len(o) < 2: + return float('nan') + denom = ((o - o.mean()) ** 2).sum() + if denom == 0: + return float('nan') + return 1.0 - float(((o - s) ** 2).sum() / denom) + + +def run_testing_period(best_param_dict): + """Run CFE over test period using NWM operational forcing with optional EnKF-DA.""" + + def custom_load_forcing(self_cfe): + df = load_test_forcing() + self_cfe.forcing_data = df.rename(columns={"date": "time"}) + + with open(CFE_CONFIG_FILE) as f: + cfg = json.load(f) + + cfg['forcing_file'] = TEST_FORCING_FILE + cfg['soil_params']['bb'] = best_param_dict['bb'] + cfg['soil_params']['smcmax'] = best_param_dict['smcmax'] + cfg['soil_params']['satdk'] = best_param_dict['satdk'] + cfg['slop'] = best_param_dict['slop'] + cfg['max_gw_storage'] = best_param_dict['max_gw_storage'] + cfg['expon'] = best_param_dict['expon'] + cfg['Cgw'] = best_param_dict['Cgw'] + cfg['K_lf'] = best_param_dict['K_lf'] + cfg['K_nash'] = best_param_dict['K_nash'] + cfg['partition_scheme'] = "Schaake" if best_param_dict['scheme'] <= 0.5 else "Xinanjiang" + + tmp_cfg = str(OUT_DIR / f'{CAT_ID}_bmi_config_temp_test.json') + with open(tmp_cfg, 'w') as f: + json.dump(cfg, f) + + # Initialize EnKF assimilator first so its RNG/perturbation knobs are available + enkf_test = None + N_members = 1 + if ENKF_ENABLED: + enkf_test = EnKFAssimilator( + n_members=ENKF_CONFIG.get('n_members', 20), + obs_error_std=ENKF_CONFIG.get('obs_error_std', 0.05), + obs_file=OBS_FILE, + use_vrugt_r=ENKF_CONFIG.get('use_vrugt_r', True), + vrugt_alpha=ENKF_CONFIG.get('vrugt_alpha', 0.10), + vrugt_scale=ENKF_CONFIG.get('vrugt_scale', 0.001), + ) + N_members = enkf_test.n_members + if N_members < 2: + raise ValueError(f"--enkf-members must be >=2 for true EnKF (got {N_members})") + print(f"[EnKF] Initializing {N_members} ensemble members for {CAT_ID}") + + # Build N CFE model instances. With EnKF, each member starts from a slightly + # perturbed initial state so the ensemble has spread from t=0. + models = [] + for i in range(N_members): + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.load_forcing_file = custom_load_forcing.__get__(m) + m.initialize() + if ENKF_ENABLED and i > 0: + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + r = enkf_test.init_state_perturb_frac + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1 + r * enkf_test.rng.standard_normal()), 1e-6, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1 + r * enkf_test.rng.standard_normal()), 1e-6, gw_max)) + # Nash buckets often start at 0 → use small additive jitter instead of multiplicative + jitter = max(sm_max * 1e-6, 1e-9) + m.nash_storage[0] = max(n00 + jitter * enkf_test.rng.standard_normal(), 0.0) + m.nash_storage[1] = max(n10 + jitter * enkf_test.rng.standard_normal(), 0.0) + models.append(m) + + df = load_test_forcing() + + # ----- Spinup (with perturbed forcing per member if ensemble) ----- + sp_mask = (df['date'] >= TIME_SPLIT['spinup-for-testing']['start']) & \ + (df['date'] <= TIME_SPLIT['spinup-for-testing']['end']) + df_sp = df[sp_mask] + for p, e in zip(df_sp['total_precipitation'], df_sp['potential_evaporation']): + if ENKF_ENABLED: + p_arr, e_arr = enkf_test.perturb_forcing(p, e) + for i, m in enumerate(models): + p_i = float(p_arr[i]) if ENKF_ENABLED else p + e_i = float(e_arr[i]) if ENKF_ENABLED else e + m.set_value('atmosphere_water__time_integral_of_precipitation_mass_flux', p_i / 1000) + m.set_value('water_potential_evaporation_flux', e_i / 1000 / 3600) + m.update() + # Process noise during spinup keeps spread alive before DA starts + if ENKF_ENABLED: + enkf_test.add_process_noise(models) + + # ----- Test period ----- + t_mask = (df['date'] >= TIME_SPLIT['testing']['start']) & \ + (df['date'] <= TIME_SPLIT['testing']['end']) + df_test = df[t_mask] + outputs = models[0].get_output_var_names() + out_lists = {o: [] for o in outputs} + + for p, e, current_date in zip(df_test['total_precipitation'], + df_test['potential_evaporation'], + df_test['date']): + if ENKF_ENABLED: + p_arr, e_arr = enkf_test.perturb_forcing(p, e) + + # Run each member forward one hour, collect ensemble Q_sim + ensemble_q_mm_h = np.empty(N_members, dtype=float) + for i, m in enumerate(models): + p_i = float(p_arr[i]) if ENKF_ENABLED else p + e_i = float(e_arr[i]) if ENKF_ENABLED else e + m.set_value('atmosphere_water__time_integral_of_precipitation_mass_flux', p_i / 1000) + m.set_value('water_potential_evaporation_flux', e_i / 1000 / 3600) + m.update() + ensemble_q_mm_h[i] = m.get_value('land_surface_water__runoff_depth') * 1000 # m/h → mm/h + + # True ensemble EnKF state update + if ENKF_ENABLED and enkf_test is not None: + enkf_test.update_states(models, current_date, ensemble_q_mm_h) + + # Process noise on analyzed states — keeps ensemble spread alive across + # the long test run so the filter doesn't collapse (Pyy → 0). + if ENKF_ENABLED and enkf_test is not None: + enkf_test.add_process_noise(models) + + # Recorded Q_sim at time t is the FORECAST (pre-analysis): it was computed + # by model.update() above, before this hour's DA touched the states. DA + # affects the recorded series only from t+1 onward, via the analyzed states + # carried into the next iteration. This is standard sequential-filter + # forecast verification — comparing Q_sim(t) to obs(t) is not circular. + for o in outputs: + vals = [m.get_value(o) for m in models] + out_lists[o].append(float(np.mean(vals))) + + for m in models: + m.finalize() + + sim_test = np.array(out_lists['land_surface_water__runoff_depth']) * 1000 # m/h → mm/h + test_dates = pd.to_datetime(df_test['date'].values) + + # Load kriging obs for test period + obs_raw = pd.read_csv(OBS_FILE) + if 'time' in obs_raw.columns and 'qkrig_mm_hr' in obs_raw.columns: + obs_raw = obs_raw.rename(columns={'time': 'datetime', 'qkrig_mm_hr': 'qkrig'}) + obs_raw['date'] = pd.to_datetime(obs_raw['datetime']).dt.strftime('%Y-%m-%d %H:%M:%S') + obs_raw = obs_raw.rename(columns={"qkrig": "obs_mm_h"}) + test_dates_df = pd.DataFrame({'date': test_dates.strftime('%Y-%m-%d %H:%M:%S')}) + merged = test_dates_df.merge(obs_raw[['date', 'obs_mm_h']], on='date', how='left') + obs_test = merged['obs_mm_h'].values + + kge_test = _kge_score(obs_test, sim_test) + nse_test = _nse_score(obs_test, sim_test) + + da_suffix = " | EnKF-DA" if ENKF_ENABLED else "" + print(f"Test KGE: {kge_test:.4f} | Test NSE: {nse_test:.4f}{da_suffix}") + + df_out = pd.DataFrame({ + 'date': test_dates.strftime('%Y-%m-%d %H:%M:%S'), + 'sim_mm_h': sim_test, + 'obs_mm_h': obs_test, + 'precip_mm_h': df_test['total_precipitation'].values, + }) + df_out.to_csv(OUT_DIR / f'{CAT_ID}_test_results.csv', index=False) + + # Plot test period + fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(20, 10)) + + ax1.plot(test_dates, sim_test, 'tomato', lw=1.5, label='simulated') + ax1.plot(test_dates, obs_test, 'k', lw=1, label='observed (kriging)') + da_label = " (EnKF-DA)" if ENKF_ENABLED else "" + ax1.set_ylabel('Discharge (mm/h)') + ax1.set_title(f'{CAT_ID} | Test (Oct 2023–Oct 2024){da_label} | KGE={kge_test:.4f} | NSE={nse_test:.4f}') + ax1.legend() + ax1_twin = ax1.twinx() + ax1_twin.plot(test_dates, df_test['total_precipitation'].values, 'steelblue', lw=0.8, alpha=0.5) + ax1_twin.set_ylim([80, 0]) + ax1_twin.set_ylabel('Precip (mm/h)') + + helene = (test_dates >= pd.Timestamp('2024-09-20')) & (test_dates <= pd.Timestamp('2024-10-05')) + ax2.plot(test_dates[helene], sim_test[helene], 'tomato', lw=2, label='simulated') + ax2.plot(test_dates[helene], obs_test[helene], 'k', lw=1.5, label='observed') + ax2.set_ylabel('Discharge (mm/h)') + ax2.set_title('Helene Zoom: Sep 20 – Oct 5, 2024') + ax2.legend() + ax2_twin = ax2.twinx() + ax2_twin.bar(test_dates[helene], df_test['total_precipitation'].values[helene], + color='steelblue', alpha=0.4, width=0.04) + ax2_twin.set_ylim([50, 0]) + ax2_twin.set_ylabel('Precip (mm/h)') + + plt.tight_layout() + plt.savefig(OUT_DIR / f'{CAT_ID}_test_plot.png', dpi=150, bbox_inches='tight') + plt.close() + + if ENKF_ENABLED and enkf_test is not None: + print(f"[EnKF Stats - Test] N={enkf_test.n_members} | updates: {enkf_test.n_updates} | " + f"mean |innov|: {enkf_test.total_increment/max(enkf_test.n_updates,1):.4f} mm/h | " + f"avg Pyy: {enkf_test.avg_pyy:.6f} | " + f"avg K_sm={enkf_test.avg_K_sm:+.3e} K_gw={enkf_test.avg_K_gw:+.3e} " + f"K_n0={enkf_test.avg_K_n0:+.3e} K_n1={enkf_test.avg_K_n1:+.3e} | " + f"Mass lost at GW underflow: {enkf_test.total_overflow_lost_mm:.4f} mm") + + return kge_test, nse_test + + +def main(): + global CAT_ID, FORCING_FILE, TEST_FORCING_FILE, OBS_FILE + global CFE_CONFIG_FILE, PARAM_BOUNDS_FILE, OUT_DIR, bmi_cfe + global ENKF_ENABLED, ENKF_CONFIG + + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', required=True, help='Catchment ID, e.g. cat-1016300') + parser.add_argument('--forcing-dir', required=True, help='Dir with per-catchment NWM retro CSVs') + parser.add_argument('--obs-dir', required=True, help='Dir with per-catchment kriging obs CSVs') + parser.add_argument('--cfe-dir', required=True, help='Path to cfe_py directory') + parser.add_argument('--config-file', required=True, help='Base CFE BMI config JSON') + parser.add_argument('--param-bounds', required=True, help='Parameter bounds JSON') + parser.add_argument('--out-dir', required=True, help='Output directory') + parser.add_argument('--test-forcing-dir1', default=None, help='NWM operational forcings dir 1 (2023-Feb2024)') + parser.add_argument('--test-forcing-dir2', default=None, help='NWM operational forcings dir 2 (Feb2024-2025)') + parser.add_argument('--enkf-enabled', action='store_true', help='Enable EnKF-based Data Assimilation') + parser.add_argument('--enkf-members', type=int, default=20, help='Number of ensemble members (default: 20)') + parser.add_argument('--enkf-obs-error-std', type=float, default=0.05, help='Observation error std dev in mm/h (default: 0.05)') + # Vrugt R is ON by default in the production config. Pass --no-vrugt-r to disable. + parser.add_argument('--no-vrugt-r', action='store_true', + help='Disable Vrugt 2005 heteroscedastic R (revert to raw kriging variance)') + parser.add_argument('--vrugt-alpha', type=float, default=0.10, + help='Relative-error fraction in Vrugt R (default: 0.10)') + parser.add_argument('--vrugt-scale', type=float, default=0.001, + help='Scaling on kriging-variance term in Vrugt R (default: 0.001)') + args = parser.parse_args() + + CAT_ID = args.cat_id + FORCING_FILE = os.path.join(args.forcing_dir, f'{CAT_ID}.csv') + OBS_FILE = os.path.join(args.obs_dir, f'{CAT_ID}.csv') + CFE_CONFIG_FILE = args.config_file + PARAM_BOUNDS_FILE = args.param_bounds + OUT_DIR = Path(args.out_dir) / CAT_ID + OUT_DIR.mkdir(parents=True, exist_ok=True) + + # Set EnKF configuration + ENKF_ENABLED = args.enkf_enabled + ENKF_CONFIG = { + 'n_members': args.enkf_members, + 'obs_error_std': args.enkf_obs_error_std, + 'use_vrugt_r': not args.no_vrugt_r, # default ON; --no-vrugt-r disables + 'vrugt_alpha': args.vrugt_alpha, + 'vrugt_scale': args.vrugt_scale, + } + + if ENKF_ENABLED: + print(f"\n{'='*80}") + print(f"EnKF Data Assimilation ENABLED") + print(f" Ensemble members: {ENKF_CONFIG['n_members']}") + print(f" Observation error std dev: {ENKF_CONFIG['obs_error_std']} mm/h") + if ENKF_CONFIG['use_vrugt_r']: + print(f" Vrugt R: alpha={ENKF_CONFIG['vrugt_alpha']}, scale={ENKF_CONFIG['vrugt_scale']}") + else: + print(f" Vrugt R: disabled (using raw kriging variance / fallback)") + print(f"{'='*80}\n") + + # Build combined test forcing file if both dirs provided + if args.test_forcing_dir1 and args.test_forcing_dir2: + f1 = os.path.join(args.test_forcing_dir1, f'{CAT_ID}.csv') + f2 = os.path.join(args.test_forcing_dir2, f'{CAT_ID}.csv') + if os.path.exists(f1) and os.path.exists(f2): + df1 = pd.read_csv(f1) + df2 = pd.read_csv(f2) + combined = pd.concat([df1, df2], ignore_index=True) + combined = combined.drop_duplicates(subset='time').sort_values('time') + TEST_FORCING_FILE = str(OUT_DIR / f'{CAT_ID}_nwm_operational_combined.csv') + combined.to_csv(TEST_FORCING_FILE, index=False) + else: + print(f"Warning: test forcing files not found for {CAT_ID}, skipping test period") + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + bmi_cfe = _bmi_cfe + + best_params_file = OUT_DIR / f'{CAT_ID}_best_params.json' + + if not TEST_FORCING_FILE or not os.path.exists(TEST_FORCING_FILE): + print("No test forcing available; nothing to do. Pre-stage the operational " + "forcing dirs via --test-forcing-dir1/--test-forcing-dir2.") + return + + if not best_params_file.exists(): + print(f"No best params file at {best_params_file}. Pre-stage the Run 3 " + f"best_params.json into the out-dir before running.") + return + + with open(best_params_file) as f: + saved = json.load(f) + best_param_dict = saved['best_parameters'] + print(f"\nLoaded calibrated parameters from {best_params_file}") + print(f"Running test period (Oct 2023 – Oct 2024)...") + run_testing_period(best_param_dict) + + +if __name__ == '__main__': + main() diff --git a/da_methods/2_troute_routing/run_route.py b/da_methods/2_troute_routing/run_route.py new file mode 100644 index 00000000..acbdb151 --- /dev/null +++ b/da_methods/2_troute_routing/run_route.py @@ -0,0 +1,350 @@ +#!/usr/bin/env python3 +""" +run_route.py +Route a single deterministic DA trajectory through T-route Muskingum-Cunge +to gauge 03463300 (South Toe River Near Celo, NC). + +Reads per-catchment *_test_results.csv files (sim_mm_h column) from a DA +results directory, routes the full timeseries through the channel network, +and writes routed_Q_test.csv at the gauge outlet. + +Works for any DA run that produces the standard _test_results.csv layout: + columns: date, sim_mm_h, obs_mm_h, precip_mm_h + +Usage: + /home/svyas/miniconda3/envs/troute/bin/python run_route.py \\ + --gpkg /mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg \\ + --da-dir /mnt/disk2/1400_sites_helene/da_results_dynamic_vrugt_seeded \\ + --out-dir /mnt/disk2/suma_helen_poster/da_results/dynamic_vrugt_seeded_routed \\ + --usgs-csv /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv \\ + --kv-dir /mnt/disk2/1400_sites_helene/catchment_ts_03463300_spliced_dyn_helene + + # For the no-Vrugt (raw variance) seeded run: + /home/svyas/miniconda3/envs/troute/bin/python run_route.py \\ + --gpkg /mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg \\ + --da-dir /mnt/disk2/1400_sites_helene/da_results_dynamic_novrugt_seeded \\ + --out-dir /mnt/disk2/suma_helen_poster/da_results/dynamic_novrugt_seeded_routed \\ + --usgs-csv /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv \\ + --kv-dir /mnt/disk2/1400_sites_helene/catchment_ts_03463300_spliced_dyn_helene +""" + +import argparse +import os +import sys +import sqlite3 +import types as _types +from functools import partial + +import numpy as np +import pandas as pd + +# Coverage stub — required for numba/troute on Python 3.10 +_stub = _types.ModuleType('coverage.types') +for _cls in ['Tracer', 'TTraceData', 'TShouldTraceFn', 'TFileDisposition', + 'TShouldStartContextFn', 'TWarnFn', 'TTraceFn']: + setattr(_stub, _cls, type(_cls, (), {})) +sys.modules['coverage.types'] = _stub + +import troute.nhd_network as nhd_network +from troute.routing.fast_reach.mc_reach import compute_network_structured + +TERMINAL_INT = 1016283 # wb-1016283 → gauge 03463300 +DT = 3600.0 # 1-hour timestep (seconds) +QTS_SUBDIVISIONS = 1 +WATERSHED_AREA_KM2 = 113.18 + +CATS = [ + 'cat-1016279', 'cat-1016280', 'cat-1016281', 'cat-1016282', 'cat-1016283', + 'cat-1016300', 'cat-1016301', 'cat-1016302', 'cat-1016303', 'cat-1016304', + 'cat-1016305', 'cat-1016306', 'cat-1016307', 'cat-1016308', 'cat-1016309', + 'cat-1016310', 'cat-1016311', 'cat-1016312', 'cat-1016313', 'cat-1016314', + 'cat-1016315', +] + + +def seg_int(wb_id): + return int(wb_id[3:]) + + +def read_network(gpkg): + con = sqlite3.connect(gpkg) + fp_attr = pd.read_sql( + 'SELECT link, "to", BtmWdth, TopWdth, TopWdthCC, n, nCC, ChSlp, So, Length_m ' + 'FROM "flowpath-attributes"', con) + fp = pd.read_sql('SELECT divide_id, areasqkm FROM flowpaths', con) + con.close() + return fp_attr, fp + + +def build_connections(fp_attr): + link_set = {seg_int(r) for r in fp_attr['link']} + connections = {} + for _, row in fp_attr.iterrows(): + us = seg_int(row['link']) + ds = int(row['to'][4:]) + connections[us] = [ds] if ds in link_set else [] + return connections + + +def build_param_df(fp_attr): + rows = [{ + 'seg_id': seg_int(r['link']), + 'dt': float(DT), + 'bw': float(r['BtmWdth']), 'tw': float(r['TopWdth']), + 'twcc':float(r['TopWdthCC']), 'dx': float(r['Length_m']), + 'n': float(r['n']), 'ncc': float(r['nCC']), + 'cs': float(r['ChSlp']), 's0': float(r['So']), + 'alt': 0.0, + } for _, r in fp_attr.iterrows()] + df = pd.DataFrame(rows).set_index('seg_id').sort_index() + return df.astype('float32') + + +def route_timeseries(reaches_wTypes, upstreams, param_df, q0_df, + qlat_arr, nts, terminal_pos): + """Route qlat_arr (n_segs × nts, m³/s) and return Q at terminal (nts,).""" + e1i = np.zeros(0, dtype='int32') + e1f = np.zeros(0, dtype='float32') + e2f = np.zeros((0, nts), dtype='float32') + e00f32 = np.zeros((0, 0), dtype='float32') + e00f64 = np.zeros((0, 0), dtype='float64') + e00i32 = np.zeros((0, 0), dtype='int32') + + results = compute_network_structured( + nts, DT, QTS_SUBDIVISIONS, + reaches_wTypes, upstreams, + param_df.index.values.astype('int64'), + param_df.columns.values, + param_df.values, + q0_df.values.astype('float32'), + qlat_arr.astype('float32'), + [], e00f64, {}, e00i32, False, + '2023-10-01_00:00:00', + e2f, e1i, e1i, e1i, e1f, e1f, 0.0, + e2f, e1i, e1f, e1f, e1f, e1f, e1f, + e2f, e1i, e1f, e1f, e1f, e1f, e1f, + e2f, e1i, e1i, [], e1i, e1i, e1f, e1i, e1i, + e1i, e1i, e1f, e1i, e1f, e1i, e1i, e00f32, + ) + seg_ids = np.asarray(results[0]) + fvd = np.asarray(results[1]) + Q_var = fvd[terminal_pos, :nts] + Q_time = fvd[terminal_pos, 0::3] + return Q_var if Q_var.max() > Q_time.max() else Q_time + + +def load_usgs(usgs_csv): + if not usgs_csv or not os.path.exists(usgs_csv): + return None + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ('datetime', 'date', 'time', 'timestamp')) + q_col = next(c for c in df.columns + if any(k in c.lower() for k in ('q', 'flow', 'discharge'))) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if 'mm' in q_col.lower(): + series = series * (WATERSHED_AREA_KM2 * 1000.0 / 3600.0) + return series + + +def load_krig(kv_dir, dates): + """Load Qkrig (mm/h) from the outlet catchment obs CSV and convert to m³/s.""" + for fname in ('cat-1016300.csv',): + p = os.path.join(kv_dir, fname) + if not os.path.exists(p): + print(f' WARNING: Qkrig file not found: {p}') + return None + df = pd.read_csv(p) + df.columns = [c.strip() for c in df.columns] + # Accept any date-like first column + date_col = next((c for c in df.columns + if c.lower() in ('datetime', 'date', 'time', 'timestamp')), + df.columns[0]) + df[date_col] = pd.to_datetime(df[date_col]) + # Accept qkrig, qkrig_mm_hr, or any qkrig column without 'var' + q_col = next((c for c in df.columns + if 'qkrig' in c.lower() and 'var' not in c.lower()), None) + if q_col is None: + print(f' WARNING: no qkrig column in {p}. Columns: {list(df.columns)}') + return None + df = df.set_index(date_col).sort_index() + area_m2 = 113.18 * 1e6 # full watershed — Qkrig at outlet catchment + krig_m3s = df[q_col].astype(float) / 1000.0 / 3600.0 * area_m2 + return krig_m3s.reindex(dates) + return None + + +def kge(obs, sim): + mask = np.isfinite(obs) & np.isfinite(sim) + o, s = obs[mask], sim[mask] + if len(o) < 2 or np.std(o) == 0: + return np.nan + r = np.corrcoef(o, s)[0, 1] + return 1.0 - np.sqrt((r-1)**2 + (np.std(s)/np.std(o)-1)**2 + + (np.mean(s)/np.mean(o)-1)**2) + + +def nse(obs, sim): + mask = np.isfinite(obs) & np.isfinite(sim) + o, s = obs[mask], sim[mask] + if len(o) < 2: + return np.nan + denom = np.sum((o - o.mean())**2) + return 1.0 - np.sum((o - s)**2) / denom if denom > 0 else np.nan + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--gpkg', required=True, + help='GeoPackage with flowpath-attributes and flowpaths tables') + parser.add_argument('--da-dir', required=True, + help='Dir containing /_test_results.csv (sim_mm_h column)') + parser.add_argument('--out-dir', required=True, + help='Output directory for routed_Q_test.csv') + parser.add_argument('--usgs-csv', default=None, + help='USGS obs CSV for KGE/NSE summary (optional)') + parser.add_argument('--kv-dir', default=None, + help='Obs dir with per-catchment Qkrig CSVs (optional)') + args = parser.parse_args() + os.makedirs(args.out_dir, exist_ok=True) + + # ── Network ──────────────────────────────────────────────────────────── + print('Reading network from GPKG...') + fp_attr, fp = read_network(args.gpkg) + connections = build_connections(fp_attr) + rconn = nhd_network.reverse_network(connections) + path_func = partial(nhd_network.split_at_junction, rconn) + reach_list = nhd_network.dfs_decomposition(rconn, path_func) + reaches_wTypes = [(r, 0) for r in reach_list] + upstreams = dict(rconn) + param_df = build_param_df(fp_attr) + n_segs = len(param_df) + + area_map = {int(r['divide_id'][4:]): r['areasqkm'] * 1e6 + for _, r in fp.iterrows() + if r['divide_id'] and str(r['divide_id']).startswith('cat-')} + print(f' {n_segs} segments | {len(area_map)} catchment areas loaded') + + seg_ids_sorted = param_df.index.values + terminal_pos = int(np.where(seg_ids_sorted == TERMINAL_INT)[0][0]) + print(f' Terminal wb-{TERMINAL_INT} at index {terminal_pos}') + + q0_df = pd.DataFrame( + np.zeros((n_segs, 3), dtype='float32'), + index=param_df.index, + columns=['qu0', 'qd0', 'h0']) + + # ── Load per-catchment sim_mm_h ──────────────────────────────────────── + print('Loading DA outputs...') + cat_series = {} + missing = [] + for cat in CATS: + sid = int(cat[4:]) + p = os.path.join(args.da_dir, cat, f'{cat}_test_results.csv') + if os.path.exists(p): + df = pd.read_csv(p, parse_dates=['date']) + df = df.set_index('date').sort_index() + cat_series[sid] = df['sim_mm_h'].astype(float) + continue + # Fallback 1: use ensemble median from production per-member CSV + p_prod = os.path.join(args.da_dir, cat, f'{cat}_production_per_member.csv') + if os.path.exists(p_prod): + df = pd.read_csv(p_prod, parse_dates=['date']) + df = df.set_index('date').sort_index() + mem_cols = [c for c in df.columns if c.startswith('member_')] + cat_series[sid] = df[mem_cols].median(axis=1).astype(float) + continue + # Fallback 2: reconstruct nowcast from forcing arm lead_hour=1 + # valid_time = issue_time + 1h; covers only the Helene assimilation window + p_arm = os.path.join(args.da_dir, cat, f'{cat}_da_forcing_arm.csv') + if os.path.exists(p_arm): + df = pd.read_csv(p_arm, parse_dates=['issue_time']) + now = df[df['lead_hour'] == 1].copy() + now['valid_time'] = now['issue_time'] + pd.Timedelta(hours=1) + now = now.set_index('valid_time').sort_index() + mem_cols = [c for c in now.columns if c.startswith('member_')] + cat_series[sid] = now[mem_cols].median(axis=1).astype(float) + continue + missing.append(cat) + + if missing: + print(f' WARNING: missing outputs for {missing} — zero inflow assumed') + print(f' Loaded {len(cat_series)}/21 catchments') + + # Common time index from first available catchment + ref_dates = next(iter(cat_series.values())).index + nts = len(ref_dates) + print(f' Timesteps: {nts} ({ref_dates[0]} → {ref_dates[-1]})') + + # ── Build lateral inflow array (n_segs × nts, m³/s) ─────────────────── + print('Building lateral inflow array...') + qlat = np.zeros((n_segs, nts), dtype='float32') + seg_pos = {sid: i for i, sid in enumerate(seg_ids_sorted)} + + for sid, q_mm_h in cat_series.items(): + if sid not in seg_pos: + continue + area_m2 = area_map.get(sid, 0.0) + if area_m2 == 0.0: + print(f' WARNING: no area for cat-{sid}') + continue + q_m3s = q_mm_h.reindex(ref_dates).fillna(0.0).values / 1000.0 / 3600.0 * area_m2 + qlat[seg_pos[sid], :] = q_m3s.astype('float32') + + # ── Route ────────────────────────────────────────────────────────────── + print(f'Routing {nts} timesteps through T-route...') + Q_routed = route_timeseries( + reaches_wTypes, upstreams, param_df, q0_df, qlat, nts, terminal_pos) + print(f' Done. Peak routed Q = {Q_routed.max():.2f} m³/s') + + # ── Build output DataFrame ───────────────────────────────────────────── + out_df = pd.DataFrame({'date': ref_dates, 'Q_routed_m3s': Q_routed}) + out_df = out_df.set_index('date') + + # Attach USGS obs + usgs = load_usgs(args.usgs_csv) + if usgs is not None: + out_df['Q_usgs_m3s'] = usgs.reindex(ref_dates) + + # Attach Qkrig + if args.kv_dir: + krig = load_krig(args.kv_dir, ref_dates) + if krig is not None: + out_df['Q_krig_m3s'] = krig.values + + out_path = os.path.join(args.out_dir, 'routed_Q_test.csv') + out_df.to_csv(out_path) + print(f'Saved: {out_path}') + + # ── KGE / NSE summary ───────────────────────────────────────────────── + if 'Q_usgs_m3s' in out_df.columns: + obs = out_df['Q_usgs_m3s'].values + sim = out_df['Q_routed_m3s'].values + + helene = ((out_df.index >= '2024-09-24') & + (out_df.index <= '2024-09-29 23:00:00')) + + print('\n── Routed vs USGS ─────────────────────────────────────────') + print(f' Full period : KGE={kge(obs, sim):+.3f} NSE={nse(obs, sim):+.3f}' + f' peak_sim={sim.max():.1f} peak_obs={np.nanmax(obs):.1f} m³/s') + if helene.sum() > 0: + oh, sh = obs[helene], sim[helene] + print(f' Helene window: KGE={kge(oh, sh):+.3f} NSE={nse(oh, sh):+.3f}' + f' peak_sim={sh.max():.1f} peak_obs={np.nanmax(oh):.1f} m³/s' + f' ({sh.max()/np.nanmax(oh)*100:.0f}% of USGS)') + + if 'Q_krig_m3s' in out_df.columns: + krig_v = out_df['Q_krig_m3s'].values + print(f'\n── Qkrig-routed vs USGS ───────────────────────────────────') + print(f' Full period : KGE={kge(obs, krig_v):+.3f} ' + f'NSE={nse(obs, krig_v):+.3f}') + if helene.sum() > 0: + ok, sk = obs[helene], krig_v[helene] + print(f' Helene window: KGE={kge(ok, sk):+.3f} ' + f'NSE={nse(ok, sk):+.3f}') + print() + + +if __name__ == '__main__': + main() diff --git a/da_methods/2_troute_routing/run_route_crossed_ensemble.py b/da_methods/2_troute_routing/run_route_crossed_ensemble.py new file mode 100644 index 00000000..22fa34e8 --- /dev/null +++ b/da_methods/2_troute_routing/run_route_crossed_ensemble.py @@ -0,0 +1,254 @@ +#!/usr/bin/env python3 +""" +run_route_crossed_ensemble.py +Route the 600-member crossed ensemble through t-route Muskingum-Cunge +to gauge 03463300 (South Toe River Near Celo, NC). + +Reads per-catchment crossed_ensemble.parquet files (all 21 catchments), +routes each member's 18-lead forecast through the channel network, +and writes a single routed parquet at the gauge outlet. + +Input (one per catchment, from run_crossed_ensemble.py): + //_crossed_ensemble.parquet + Columns: issue_time, lead_hour, member_0000..member_0599 (q in mm/h) + +Output: + /routed_crossed_ensemble.parquet + Columns: issue_time, lead_hour, member_0000..member_0599 (Q in m3/s) + +Usage: + python3 run_route_crossed_ensemble.py \\ + --gpkg /mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg \\ + --ensemble-dir /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble \\ + --out-dir /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble_routed \\ + --usgs-csv /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv +""" + +import argparse +import sqlite3 +import os +import sys +import types as _types +from functools import partial + +import numpy as np +import pandas as pd + +# Patch coverage.types for numba/troute on Python 3.10 +_stub = _types.ModuleType('coverage.types') +for _cls in ['Tracer','TTraceData','TShouldTraceFn','TFileDisposition', + 'TShouldStartContextFn','TWarnFn','TTraceFn']: + setattr(_stub, _cls, type(_cls, (), {})) +sys.modules['coverage.types'] = _stub + +import troute.nhd_network as nhd_network +from troute.routing.fast_reach.mc_reach import compute_network_structured + +TERMINAL_INT = 1016283 # wb-1016283 = gauge 03463300 +DT = 3600.0 # 1-hour timestep (seconds) +QTS_SUBDIVISIONS = 1 + +CATS = [ + 'cat-1016279','cat-1016280','cat-1016281','cat-1016282','cat-1016283', + 'cat-1016300','cat-1016301','cat-1016302','cat-1016303','cat-1016304', + 'cat-1016305','cat-1016306','cat-1016307','cat-1016308','cat-1016309', + 'cat-1016310','cat-1016311','cat-1016312','cat-1016313','cat-1016314', + 'cat-1016315', +] + + +def seg_int(wb_id): + return int(wb_id[3:]) + + +def read_network(gpkg): + con = sqlite3.connect(gpkg) + fp_attr = pd.read_sql( + 'SELECT link, "to", BtmWdth, TopWdth, TopWdthCC, n, nCC, ChSlp, So, Length_m ' + 'FROM "flowpath-attributes"', con) + fp = pd.read_sql('SELECT divide_id, areasqkm FROM flowpaths', con) + con.close() + return fp_attr, fp + + +def build_connections(fp_attr): + link_set = {seg_int(r) for r in fp_attr['link']} + connections = {} + for _, row in fp_attr.iterrows(): + us = seg_int(row['link']) + ds = int(row['to'][4:]) + connections[us] = [ds] if ds in link_set else [] + return connections + + +def build_param_df(fp_attr): + rows = [{ + 'seg_id': seg_int(r['link']), + 'dt': float(DT), + 'bw': float(r['BtmWdth']), 'tw': float(r['TopWdth']), + 'twcc': float(r['TopWdthCC']), 'dx': float(r['Length_m']), + 'n': float(r['n']), 'ncc': float(r['nCC']), + 'cs': float(r['ChSlp']), 's0': float(r['So']), + 'alt': 0.0, + } for _, r in fp_attr.iterrows()] + df = pd.DataFrame(rows).set_index('seg_id').sort_index() + return df.astype('float32') + + +def route_one(reaches_wTypes, upstreams, param_df, q0_df, + qlat_arr, nts, n_segs, terminal_pos): + """Route a single member's qlat through t-route; return Q at terminal.""" + e1i = np.zeros(0, dtype='int32') + e1f = np.zeros(0, dtype='float32') + e2f = np.zeros((0, nts), dtype='float32') + e00f32 = np.zeros((0, 0), dtype='float32') + e00f64 = np.zeros((0, 0), dtype='float64') + e00i32 = np.zeros((0, 0), dtype='int32') + + results = compute_network_structured( + nts, DT, QTS_SUBDIVISIONS, + reaches_wTypes, upstreams, + param_df.index.values.astype('int64'), + param_df.columns.values, + param_df.values, + q0_df.values.astype('float32'), + qlat_arr.astype('float32'), + [], e00f64, {}, e00i32, False, + '2024-09-24_00:00:00', + e2f, e1i, e1i, e1i, e1f, e1f, 0.0, + e2f, e1i, e1f, e1f, e1f, e1f, e1f, + e2f, e1i, e1f, e1f, e1f, e1f, e1f, + e2f, e1i, e1i, [], e1i, e1i, e1f, e1i, e1i, + e1i, e1i, e1f, e1i, e1f, e1i, e1i, e00f32, + ) + seg_ids = np.asarray(results[0]) + fvd = np.asarray(results[1]) + Q_var = fvd[terminal_pos, :nts] + Q_time = fvd[terminal_pos, 0::3] + return Q_var if Q_var.max() > Q_time.max() else Q_time + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--gpkg', required=True) + parser.add_argument('--ensemble-dir', required=True, + help='Dir containing /_crossed_ensemble.parquet') + parser.add_argument('--out-dir', required=True) + parser.add_argument('--usgs-csv', default=None, + help='Optional: USGS obs CSV for KGE/NSE summary') + args = parser.parse_args() + os.makedirs(args.out_dir, exist_ok=True) + + # -- Network ------------------------------------------------------- + print('Reading network from GPKG...') + fp_attr, fp = read_network(args.gpkg) + connections = build_connections(fp_attr) + rconn = nhd_network.reverse_network(connections) + path_func = partial(nhd_network.split_at_junction, rconn) + reach_list = nhd_network.dfs_decomposition(rconn, path_func) + reaches_wTypes = [(r, 0) for r in reach_list] + upstreams = dict(rconn) + param_df = build_param_df(fp_attr) + n_segs = len(param_df) + + area_map = {int(r['divide_id'][4:]): r['areasqkm'] * 1e6 + for _, r in fp.iterrows() + if r['divide_id'] and str(r['divide_id']).startswith('cat-')} + print(f' {len(fp_attr)} segments | total area {sum(area_map.values())/1e6:.2f} km2') + + # Terminal position in param_df + seg_ids_sorted = param_df.index.values + terminal_pos = int(np.where(seg_ids_sorted == TERMINAL_INT)[0][0]) + print(f' Terminal: wb-{TERMINAL_INT} at param_df index {terminal_pos}') + + q0_df = pd.DataFrame( + np.zeros((n_segs, 3), dtype='float32'), + index=param_df.index, + columns=['qu0', 'qd0', 'h0']) + + # -- Load per-catchment ensemble parquets -------------------------- + print('Loading per-catchment ensemble parquets...') + cat_data = {} # cat_seg_int -> DataFrame (issue_time, lead_hour, member_*) + missing = [] + for cat in CATS: + sid = int(cat[4:]) + pq = os.path.join(args.ensemble_dir, cat, f'{cat}_crossed_ensemble.parquet') + if os.path.exists(pq): + df = pd.read_parquet(pq) + df['issue_time'] = pd.to_datetime(df['issue_time']) + cat_data[sid] = df + else: + missing.append(cat) + if missing: + print(f' WARNING: missing parquets for: {missing}') + print(' These catchments will have zero lateral inflow.') + print(f' Loaded {len(cat_data)}/21 catchments') + + # Identify shared issue_times and member columns + ref_df = next(iter(cat_data.values())) + issue_times = sorted(ref_df['issue_time'].unique()) + member_cols = sorted([c for c in ref_df.columns if c.startswith('member_')]) + n_members = len(member_cols) + n_leads = int(ref_df['lead_hour'].max()) + print(f' {len(issue_times)} issue times | {n_leads} leads | {n_members} members') + + # -- Route --------------------------------------------------------- + # For each issue_time, route all 600 members through t-route. + # qlat shape for one member: (n_segs, n_leads) in m3/s + all_records = [] + + for t_idx, t0 in enumerate(issue_times): + t0_str = t0.strftime('%Y-%m-%d %H:%M:%S') + + # Build qlat cube: (n_members, n_segs, n_leads) in m3/s + qlat_cube = np.zeros((n_members, n_segs, n_leads), dtype='float32') + + for sid, df in cat_data.items(): + seg_pos = int(np.where(seg_ids_sorted == sid)[0]) if sid in seg_ids_sorted else -1 + if seg_pos < 0: + continue + area = area_map.get(sid, 0.0) + + sub = df[df['issue_time'] == t0].sort_values('lead_hour') + if sub.empty: + continue + + # mem_vals shape: (n_leads, n_members), in mm/h -> m3/s + mem_vals = sub[member_cols].to_numpy(dtype='float32') # (n_leads, n_mem) + mem_vals = mem_vals / 1000.0 / 3600.0 * area # mm/h -> m3/s + # qlat_cube[:, seg_pos, :] = mem_vals.T (n_mem, n_leads) + qlat_cube[:, seg_pos, :] = mem_vals.T + + # Route each member + Q_members = np.full((n_leads, n_members), np.nan, dtype='float32') + for m in range(n_members): + qlat_m = qlat_cube[m] # (n_segs, n_leads) + qlat_df_m = pd.DataFrame( + qlat_m, index=param_df.index, + columns=range(n_leads)) + qlat_df_m = qlat_df_m.reindex(param_df.index, fill_value=0.0) + Q_t = route_one(reaches_wTypes, upstreams, param_df, q0_df, + qlat_df_m.values.astype('float32'), + n_leads, n_segs, terminal_pos) + Q_members[:, m] = Q_t[:n_leads] + + # Pack into rows (one per lead_hour) + for lead in range(1, n_leads + 1): + row = {'issue_time': t0_str, 'lead_hour': lead} + for m, col in enumerate(member_cols): + row[col] = float(Q_members[lead - 1, m]) + all_records.append(row) + + if (t_idx + 1) % 10 == 0 or (t_idx + 1) == len(issue_times): + print(f' Routed {t_idx + 1}/{len(issue_times)} issue times') + + # -- Save ---------------------------------------------------------- + df_out = pd.DataFrame(all_records) + out_path = os.path.join(args.out_dir, 'routed_crossed_ensemble.parquet') + df_out.to_parquet(out_path, index=False) + print(f'Saved: {out_path}') + print(f' Shape: {df_out.shape}') + + +if __name__ == '__main__': + main() diff --git a/da_methods/2_troute_routing/run_route_leadtime_forecasts.py b/da_methods/2_troute_routing/run_route_leadtime_forecasts.py new file mode 100644 index 00000000..afbc3f68 --- /dev/null +++ b/da_methods/2_troute_routing/run_route_leadtime_forecasts.py @@ -0,0 +1,257 @@ +#!/usr/bin/env python3 +""" +run_route_leadtime_forecasts.py +Route per-catchment lead-time forecast CSVs through T-route to gauge 03463300. + +Reads the DA and open-loop lead-time forecast CSVs produced by +run_lead_time_forecast_sweep.py (one per catchment), routes all 20 members +through Muskingum-Cunge, and writes two output parquets at the gauge outlet. + +Input (per catchment, from run_lead_time_forecast_sweep.py): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv + Columns: issue_time, lead_hour, valid_time, member_00..member_19 (q in mm/h) + +Output: + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + Columns: issue_time, lead_hour, member_00..member_19 (Q in m³/s at gauge) + +Usage: + python3 run_route_leadtime_forecasts.py \\ + --gpkg /mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg \\ + --forecast-dir /mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt \\ + --out-dir /mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt_leadtime_routed +""" + +import argparse +import sqlite3 +import os +import sys +import types as _types +from functools import partial + +import numpy as np +import pandas as pd + +# Patch coverage.types for numba/troute on Python 3.10 +_stub = _types.ModuleType('coverage.types') +for _cls in ['Tracer','TTraceData','TShouldTraceFn','TFileDisposition', + 'TShouldStartContextFn','TWarnFn','TTraceFn']: + setattr(_stub, _cls, type(_cls, (), {})) +sys.modules['coverage.types'] = _stub + +import troute.nhd_network as nhd_network +from troute.routing.fast_reach.mc_reach import compute_network_structured + +TERMINAL_INT = 1016283 # wb-1016283 = gauge 03463300 +DT = 3600.0 # 1-hour timestep (seconds) +QTS_SUBDIVISIONS = 1 + +CATS = [ + 'cat-1016279','cat-1016280','cat-1016281','cat-1016282','cat-1016283', + 'cat-1016300','cat-1016301','cat-1016302','cat-1016303','cat-1016304', + 'cat-1016305','cat-1016306','cat-1016307','cat-1016308','cat-1016309', + 'cat-1016310','cat-1016311','cat-1016312','cat-1016313','cat-1016314', + 'cat-1016315', +] + + +def seg_int(wb_id): + return int(wb_id[3:]) + + +def read_network(gpkg): + con = sqlite3.connect(gpkg) + fp_attr = pd.read_sql( + 'SELECT link, "to", BtmWdth, TopWdth, TopWdthCC, n, nCC, ChSlp, So, Length_m ' + 'FROM "flowpath-attributes"', con) + fp = pd.read_sql('SELECT divide_id, areasqkm FROM flowpaths', con) + con.close() + return fp_attr, fp + + +def build_connections(fp_attr): + link_set = {seg_int(r) for r in fp_attr['link']} + connections = {} + for _, row in fp_attr.iterrows(): + us = seg_int(row['link']) + ds = int(row['to'][4:]) + connections[us] = [ds] if ds in link_set else [] + return connections + + +def build_param_df(fp_attr): + rows = [{ + 'seg_id': seg_int(r['link']), + 'dt': float(DT), + 'bw': float(r['BtmWdth']), 'tw': float(r['TopWdth']), + 'twcc': float(r['TopWdthCC']), 'dx': float(r['Length_m']), + 'n': float(r['n']), 'ncc': float(r['nCC']), + 'cs': float(r['ChSlp']), 's0': float(r['So']), + 'alt': 0.0, + } for _, r in fp_attr.iterrows()] + df = pd.DataFrame(rows).set_index('seg_id').sort_index() + return df.astype('float32') + + +def route_one(reaches_wTypes, upstreams, param_df, q0_df, + qlat_arr, nts, n_segs, terminal_pos): + """Route a single member's qlat through t-route; return Q at terminal.""" + e1i = np.zeros(0, dtype='int32') + e1f = np.zeros(0, dtype='float32') + e2f = np.zeros((0, nts), dtype='float32') + e00f32 = np.zeros((0, 0), dtype='float32') + e00f64 = np.zeros((0, 0), dtype='float64') + e00i32 = np.zeros((0, 0), dtype='int32') + + results = compute_network_structured( + nts, DT, QTS_SUBDIVISIONS, + reaches_wTypes, upstreams, + param_df.index.values.astype('int64'), + param_df.columns.values, + param_df.values, + q0_df.values.astype('float32'), + qlat_arr.astype('float32'), + [], e00f64, {}, e00i32, False, + '2024-09-24_00:00:00', + e2f, e1i, e1i, e1i, e1f, e1f, 0.0, + e2f, e1i, e1f, e1f, e1f, e1f, e1f, + e2f, e1i, e1f, e1f, e1f, e1f, e1f, + e2f, e1i, e1i, [], e1i, e1i, e1f, e1i, e1i, + e1i, e1i, e1f, e1i, e1f, e1i, e1i, e00f32, + ) + seg_ids = np.asarray(results[0]) + fvd = np.asarray(results[1]) + Q_var = fvd[terminal_pos, :nts] + Q_time = fvd[terminal_pos, 0::3] + return Q_var if Q_var.max() > Q_time.max() else Q_time + + +def route_scenario(label, cat_data, seg_ids_sorted, area_map, + reaches_wTypes, upstreams, param_df, q0_df, + member_cols, issue_times, n_leads, n_members, n_segs, + terminal_pos, out_path): + """Route one scenario (da or openloop) and save output parquet.""" + all_records = [] + + for t_idx, t0 in enumerate(issue_times): + t0_str = t0.strftime('%Y-%m-%d %H:%M:%S') + + qlat_cube = np.zeros((n_members, n_segs, n_leads), dtype='float32') + + for sid, df in cat_data.items(): + seg_pos = int(np.where(seg_ids_sorted == sid)[0][0]) if sid in seg_ids_sorted else -1 + if seg_pos < 0: + continue + area = area_map.get(sid, 0.0) + + sub = df[df['issue_time'] == t0].sort_values('lead_hour') + if sub.empty: + continue + + mem_vals = sub[member_cols].to_numpy(dtype='float32') # (n_leads, n_members) + mem_vals = mem_vals / 1000.0 / 3600.0 * area # mm/h -> m³/s + qlat_cube[:, seg_pos, :mem_vals.shape[0]] = mem_vals.T + + Q_members = np.full((n_leads, n_members), np.nan, dtype='float32') + for m in range(n_members): + qlat_m = qlat_cube[m] # (n_segs, n_leads) + qlat_df_m = pd.DataFrame( + qlat_m, index=param_df.index, + columns=range(n_leads)) + qlat_df_m = qlat_df_m.reindex(param_df.index, fill_value=0.0) + Q_t = route_one(reaches_wTypes, upstreams, param_df, q0_df, + qlat_df_m.values.astype('float32'), + n_leads, n_segs, terminal_pos) + Q_members[:, m] = Q_t[:n_leads] + + for lead in range(1, n_leads + 1): + row = {'issue_time': t0_str, 'lead_hour': lead} + for m, col in enumerate(member_cols): + row[col] = float(Q_members[lead - 1, m]) + all_records.append(row) + + if (t_idx + 1) % 20 == 0 or (t_idx + 1) == len(issue_times): + print(f' [{label}] Routed {t_idx + 1}/{len(issue_times)} issue times') + + df_out = pd.DataFrame(all_records) + df_out.to_parquet(out_path, index=False) + print(f'Saved: {out_path} shape={df_out.shape}') + return df_out + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--gpkg', required=True, + help='GPKG hydrofabric (gage-03463300_subset.gpkg)') + parser.add_argument('--forecast-dir', required=True, + help='Dir with /_lead_time_forecasts_da.csv etc.') + parser.add_argument('--out-dir', required=True, + help='Output dir for routed parquets') + args = parser.parse_args() + os.makedirs(args.out_dir, exist_ok=True) + + print('Reading network from GPKG...') + fp_attr, fp = read_network(args.gpkg) + connections = build_connections(fp_attr) + rconn = nhd_network.reverse_network(connections) + path_func = partial(nhd_network.split_at_junction, rconn) + reach_list = nhd_network.dfs_decomposition(rconn, path_func) + reaches_wTypes = [(r, 0) for r in reach_list] + upstreams = dict(rconn) + param_df = build_param_df(fp_attr) + n_segs = len(param_df) + + area_map = {int(r['divide_id'][4:]): r['areasqkm'] * 1e6 + for _, r in fp.iterrows() + if r['divide_id'] and str(r['divide_id']).startswith('cat-')} + print(f' {len(fp_attr)} segments | total area {sum(area_map.values())/1e6:.2f} km2') + + seg_ids_sorted = param_df.index.values + terminal_pos = int(np.where(seg_ids_sorted == TERMINAL_INT)[0][0]) + print(f' Terminal: wb-{TERMINAL_INT} at param_df index {terminal_pos}') + + q0_df = pd.DataFrame( + np.zeros((n_segs, 3), dtype='float32'), + index=param_df.index, + columns=['qu0', 'qd0', 'h0']) + + for scenario in ('da', 'openloop'): + print(f'\nLoading per-catchment lead-time CSVs ({scenario})...') + cat_data = {} + missing = [] + for cat in CATS: + sid = int(cat[4:]) + csv_path = os.path.join(args.forecast_dir, cat, + f'{cat}_lead_time_forecasts_{scenario}.csv') + if os.path.exists(csv_path): + df = pd.read_csv(csv_path) + df['issue_time'] = pd.to_datetime(df['issue_time']) + cat_data[sid] = df + else: + missing.append(cat) + if missing: + print(f' WARNING: missing CSVs for: {missing}') + print(' These catchments will have zero lateral inflow.') + print(f' Loaded {len(cat_data)}/21 catchments') + + ref_df = next(iter(cat_data.values())) + issue_times = sorted(ref_df['issue_time'].unique()) + member_cols = sorted([c for c in ref_df.columns if c.startswith('member_')]) + n_members = len(member_cols) + n_leads = int(ref_df['lead_hour'].max()) + print(f' {len(issue_times)} issue times | {n_leads} leads | {n_members} members') + + out_path = os.path.join(args.out_dir, f'routed_leadtime_{scenario}_full.parquet') + route_scenario( + scenario, cat_data, seg_ids_sorted, area_map, + reaches_wTypes, upstreams, param_df, q0_df, + member_cols, issue_times, n_leads, n_members, n_segs, + terminal_pos, out_path) + + print('\nDone.') + + +if __name__ == '__main__': + main() diff --git a/da_methods/README.md b/da_methods/README.md new file mode 100644 index 00000000..30177a39 --- /dev/null +++ b/da_methods/README.md @@ -0,0 +1,391 @@ +# CFE + EnKF Data Assimilation — Hurricane Helene Forecasting + +**Gauge:** USGS 03463300 — South Toe River Near Celo, NC +**Watershed:** 21 NWM catchments · 113.18 km² · outlet reach wb-1016283 +**Event:** Hurricane Helene, September 24–29, 2024 (USGS peak: 1885.7 m³/s) +**Test period:** October 2023 – October 2024 (13 months) +**Model:** CFE (Conceptual Functional Equivalent) hydrological model via BMI interface +**DA method:** Ensemble Kalman Filter (EnKF) assimilating Qkrig streamflow observations + +--- + +## What Makes This Setup Different + +### Distributed CFE — not lumped + +CFE is run as **21 independent instances — one per catchment** in the South Toe River +watershed. Each instance has its own calibrated parameters (`best_params.json`), its own +NWM per-catchment forcings, and its own runoff output. + +This is different from the lumped CFE configuration used in earlier work (e.g. summer +institute projects), where a single CFE instance represents the entire watershed with +one parameter set and one forcing input. The distributed setup captures spatial +variability in soil, land cover, and catchment response that a lumped model cannot. + +### Custom T-route wrapper — not ngen-troute + +Two T-route use cases exist in this repo, both using a **custom wrapper** that calls +T-route's Muskingum-Cunge routing function (`compute_network_structured`) directly: + +| Approach | What initializes T-route | When used | +|---|---|---| +| **Qkrig → T-route** | Kriged streamflow pseudo-observations | Evaluate kriging directly as routing input | +| **Distributed CFE → T-route** | Per-catchment CFE runoff (mm/h) | Evaluate calibrated + DA-corrected CFE | + +Both differ from **ngen-troute** — the standard T-route integration inside the full +NextGen framework — which expects NWM-formatted YAML configs and full hydrofabric +network inputs. The custom wrapper was built because the CFE-BMI DA pipeline outputs +per-catchment CSVs that are incompatible with the ngen-troute CLI interface. + +### DA builds on top + +The EnKF data assimilation layer (this repo's primary contribution) sits on top of +distributed CFE and the custom T-route wrapper: + +``` +Layer 1 — Distributed CFE per-catchment runoff (21 instances) +Layer 2 — Custom T-route wrapper routes CFE or Qkrig output to outlet +Layer 3 — EnKF DA corrects CFE states using Qkrig observations +``` + +--- + +## Pipeline Design + +### 1. Calibrate CFE with Qkrig + +CFE parameters are calibrated per catchment using Qkrig kriged streamflow as +pseudo-observations. Calibration is shared across all four R-formula experiments — +the R formula is **not involved at calibration time**. + +Script: `1_calibrate/calibrate_catchment_cfe_da_v2.py` +Best params saved as: `{cat-id}_best_params.json` + +| Flag | Default | Purpose | +|---|---|---| +| `--enkf-enabled` | off | Turn DA on during calibration run | +| `--enkf-members 20` | 20 | Ensemble size | +| `--no-vrugt-r` | off | Use raw kriging variance instead of Vrugt formula | +| `--vrugt-alpha 0.10` | 0.10 | α in Vrugt R formula | +| `--vrugt-scale 0.001` | 0.001 | Kriging variance scaling factor | + +--- + +### 2. Assimilation Script for CFE + +Script: `2_assimilation/run_perturbation_da_on.py` + +#### 2a — Update Met Forcings (Forcing Arm, 30 members) + +Stochastic perturbation of meteorological forcing to quantify uncertainty from +precipitation and PET inputs, with DA-corrected initial states: + +- Precip: lognormal multiplier σ=0.15 +- PET: Gaussian multiplier σ=0.10, clipped at 0 + +Produces `{cat-id}_da_forcing_arm.csv` — 30-member spaghetti over Helene window. +**Figure 2a:** ensemble spread from met forcing uncertainty shows how sensitive +18-hr forecasts are to precipitation perturbations alone. + +#### 2b — Update Hydro States using Kriging Error Variance in Vrugt R (Hydro-state Arm, 20 members) + +Stochastic perturbation of initial hydrologic states (soil moisture, groundwater +storage) with deterministic forcing. DA uses Qkrig observations with Vrugt R: + +``` +R(t) = (α · y_obs(t))² + scale · σ²_krig(t) + α = 0.10, scale = 0.001 +``` + +Produces `{cat-id}_da_hydro_arm.csv` — 20-member spaghetti over Helene window. +**Figure 2b:** ensemble spread from initial state uncertainty shows how sensitive +18-hr forecasts are to hydrologic state perturbations alone. + +#### 2c — 18-hour Forecast Cycle (save and restart) + +At each initialization time t0: +1. EnKF analyses the ensemble state using Qkrig observation +2. Analyzed state snapshotted to `_da_snapshots.parquet` — enables restart +3. DA switched off — ensemble free-runs for 18 hours +4. At t0+1, DA resumes; a new 18-hour fork begins from the next snapshot + +Script: `2_assimilation/run_lead_time_forecast_sweep.py` + +#### 2d — Crossed Ensemble Design (600 members) + +The forcing arm (2a) and hydro-state arm (2b) are crossed to form a 600-member ensemble +that captures both sources of uncertainty simultaneously: + +``` +member_k → forcing draw k // 20 (selects 1 of 30 met perturbations) + → hydro draw k % 20 (selects 1 of 20 initial-state draws) +``` + +Script: `2_assimilation/run_crossed_ensemble.py` + +--- + +### 3. run_route.py with DA + +Per-catchment runoff (mm/h) from each ensemble member is routed through the +channel network at every hour in the 18-hr forecast cycle for each initialization +timestep via T-route Muskingum-Cunge (`compute_network_structured`) to the +terminal reach wb-1016283 at USGS gauge 03463300. + +Scripts: `3_routing/run_route.py`, `run_route_crossed_ensemble.py`, +`route_lead_time_forecasts.py` + +GPKG: `/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg` + +--- + +### 4. Evaluate at the Gauge + +#### 4a — Forecast Error Decay: Convergence to Open Loop + +How forecast skill degrades from lead hour 1 → 18. DA-initialized forecasts are +compared against the open loop (no DA) at each lead hour. Shows whether the +benefit of DA-corrected initial states persists or converges to open loop skill +by the end of the 18-hour window. + +Scripts: `4_evaluation/4a_error_decay/` + +#### 4b — Compare 600-member Ensemble vs Observations during Helene + +Full 600-member crossed ensemble spread plotted against USGS obs at gauge 03463300 +during Hurricane Helene. Shows whether the ensemble brackets the observed peak and +how the DA-corrected ensemble performs vs open loop. + +Scripts: `4_evaluation/4b_ensemble_vs_obs/` + +#### 4c — Ensemble Mean with Spread from All 18-hour Forecast Initializations + +Reconstructed timeseries: for each valid_time during Helene, all forecasts that +land on that moment (across all 151 issue times × all lead hours × all members) +are pooled. Median + 5th/95th percentile envelope plotted against USGS obs. +Shows overall forecast skill as one continuous picture across the storm. + +Scripts: `4_evaluation/4c_timeseries/` + +--- + +## What This Repo Tests + +Four experiments comparing how the **observation error variance R** is formulated +in the EnKF update. R controls how strongly the DA pulls model states toward the +kriging-interpolated observation (Qkrig) at each hourly timestep: + +| Folder | R Formula | Key idea | +|---|---|---| +| [folder1_variance_scaled_vrugt/](folder1_variance_scaled_vrugt/) | `(0.10·y)² + 0.001·σ²_krig` | Flow-scaled R — uncertainty grows with flow magnitude (Vrugt 2005) | +| [folder2_fixed_r_007/](folder2_fixed_r_007/) | `0.07` (constant) | Simple fixed R — high constant Kalman gain regardless of flow | +| [folder3_dynamic_vrugt_seeded/](folder3_dynamic_vrugt_seeded/) | `(0.10·y)² + 0.001·σ²_krig` | Same Vrugt formula with dynamic kriging variance and fixed RNG seed | +| [folder4_dynamic_variance_direct/](folder4_dynamic_variance_direct/) | `σ²_krig` directly | Raw kriging variance as R — no flow-magnitude term | + +--- + +## Key Results + +| Folder | Full KGE | Full NSE | Helene KGE | Helene NSE | Helene peak | +|---|---|---|---|---|---| +| F1 — Vrugt | +0.277 | **+0.555** | +0.213 | **+0.459** | 645.7 m³/s (34%) | +| F2 — Fixed R=0.07 | **+0.503** | +0.163 | **+0.439** | -0.009 | **1227.1 m³/s (65%)** | +| F3 — Vrugt seeded | +0.261 | +0.485 | +0.189 | +0.372 | 693.8 m³/s (37%) | +| F4 — Direct σ²_krig | +0.200 | +0.428 | +0.132 | +0.303 | 667.9 m³/s (35%) | + +**USGS Helene peak: 1885.7 m³/s** + +**Finding:** Fixed R=0.07 (F2) gives the best KGE and peak capture. Constant small R +keeps Kalman gain K = P/(P+R) near 1 throughout the flood, so the model tracks +observations aggressively at every timestep. The Vrugt formula inflates R at high flows +(R ≈ 35,000 at the Helene peak), collapsing the gain exactly when the DA update matters +most. F1 achieves the best NSE because it fits the overall hydrograph shape better — +F2's high constant gain overshoots at low flows, hurting NSE, but captures the peak +far better. + +--- + +## Which Comparisons Are Controlled + +**Not all four folders can be directly compared to each other.** Two design differences +exist between F1/F2 and F3/F4: + +### 1 — Different RNG seeds + +| Folder | RNG seed | +|---|---| +| F1, F2 | `hash(catchment_id) & 0x7fffffff` — unique per catchment | +| F3, F4 | `42` — fixed, same for all catchments | + +Different seeds → different perturbation realizations → different analysis trajectories, +even with the same R formula. The gap between F1 (+0.277) and F3 (+0.261) is explained +by seed difference alone, not methodology. + +### 2 — Different kriging variance during Helene + +F1/F2 use static σ²_krig ≈ 4.17 throughout. +F3/F4 use dynamic σ²_krig that drops to ~1.6 at the Helene peak. + +The **Qkrig flow values are identical** in both datasets; only the variance changes. +For F3 (Vrugt), the 0.001 weight on σ²_krig makes the variance difference negligible. +For F4 (direct σ²_krig), the lower variance at Helene peak means stronger DA updates — +but this is offset by the absence of the flow-magnitude term entirely. + +### Valid controlled comparisons + +| Pair | Seed | Obs variance | R formula | Valid? | +|---|---|---|---|---| +| F1 vs F2 | both hash | both static ~4.17 | Vrugt vs Fixed 0.07 | ✅ **clean** | +| F3 vs F4 | both seed=42 | both dynamic | Vrugt vs Direct σ² | ✅ **clean** | +| F1 vs F3 | different | small difference | same (Vrugt) | ⚠️ seed confound | +| F2 vs F4 | different | larger difference | different | ⚠️ seed + obs confound | + +To fully compare all four R formulas in a controlled way, F3 and F4 would need to +be re-run with the hash-based seed and the same obs dataset as F1/F2. See +[EXPERIMENTS.md](EXPERIMENTS.md) for the full completeness matrix. + +--- + +## Scientific Approach + +### EnKF update + +At each hourly timestep t: + +``` +K(t) = P(t) / (P(t) + R(t)) Kalman gain +x̂(t) = x(t) + K(t) · (y_obs(t) − H·x(t)) state update +``` + +where P is ensemble state variance, y_obs is Qkrig, and R is the observation error +variance. Each experiment differs only in how R is computed. + +### Vrugt 2005 heteroscedastic R (F1, F3) + +``` +R(t) = (α · y_obs(t))² + scale · σ²_krig(t) + α = 0.10 (10% relative error on the observation) + scale = 0.001 +``` + +R scales with flow magnitude: small at low flow (gain near 1, aggressive DA) and +large at peak flow. At Helene peak (~3.2 mm/h per catchment): R ≈ 0.1, K collapses. + +### Fixed R (F2) + +``` +R = 0.07 (constant) +``` + +Kalman gain K = P/(P+0.07) stays near 1 throughout the simulation including the peak. +The DA update is equally aggressive at low and high flows. + +### Direct kriging variance (F4) + +``` +R(t) = σ²_krig(t) +``` + +σ²_krig is the spatial kriging interpolation uncertainty — typically ~4.17 (static obs) +or ~1.6–2.5 at Helene peak (dynamic obs). This gives weaker DA than Fixed R=0.07 at +low flows and stronger DA at peak when σ²_krig drops. + +--- + +## Folder Structure + +Each experiment is self-contained with identical pipeline structure: + +``` +da_methods/ +├── README.md ← this file +├── EXPERIMENTS.md ← results table, completeness matrix, server paths +├── compare_all_folders.py ← cross-experiment comparison plots +├── build_pptx.py ← builds helene_da_results.pptx +├── pptx_figures/ ← all presentation figures (tracked) +│ +├── folder1_variance_scaled_vrugt/ R = (0.10·y)² + 0.001·σ²_krig, hash seed +├── folder2_fixed_r_007/ R = 0.07 constant, hash seed +├── folder3_dynamic_vrugt_seeded/ R = (0.10·y)² + 0.001·σ²_krig, seed=42 +└── folder4_dynamic_variance_direct/ R = σ²_krig directly, seed=42 + +Each folder/: +├── 1_calibrate/ calibrate_catchment_cfe_da_v2.py (shared params, R not involved) +├── 2_assimilation/ run_perturbation_da_on.py, run_crossed_ensemble.py, +│ run_lead_time_forecast_sweep.py, batch scripts +├── 3_routing/ run_route.py, run_route_crossed_ensemble.py, +│ route_lead_time_forecasts.py +└── 4_evaluation/ + ├── 4a_error_decay/ forecast skill decay vs lead time + ├── 4b_ensemble_vs_obs/ 600-member ensemble vs USGS + └── 4c_timeseries/ reconstructed timeseries + spaghetti plots +``` + +--- + +## Running an Experiment (F3 as example) + +```bash +# Stage A — perturbation arms (2a/2b) for all 21 catchments +bash folder3_dynamic_vrugt_seeded/2_assimilation/batch_run_all_f3.sh arms + +# Stage B — 18hr forecast cycles (2c) +bash folder3_dynamic_vrugt_seeded/2_assimilation/batch_run_all_f3.sh forecast + +# Stage C — 600-member crossed ensemble (2d) +bash folder3_dynamic_vrugt_seeded/2_assimilation/batch_run_all_f3.sh ensemble + +# Route analysis trajectory + ensemble + forecast cycles (3_routing/) +# See folder README for exact run_route.py invocations + +# Evaluate (4_evaluation/4c_timeseries/) +bash folder3_dynamic_vrugt_seeded/4_evaluation/4c_timeseries/run_4c_f3.sh +``` + +F4 works identically — replace `f3` with `f4` and use `batch_run_all_f4.sh`. + +--- + +## Ensemble Traceability + +Any of the 600 members can be fully reconstructed: + +```python +member_k → forcing_draw_i = k // 20 # which of 30 met perturbation draws + → hydro_draw_j = k % 20 # which of 20 initial-state draws +``` + +Four provenance files are saved per catchment: + +| File | Contents | +|---|---| +| `{cat-id}_da_snapshots.parquet` | DA-analyzed state at every issue time — enables restart | +| `{cat-id}_member_manifest.csv` | Decoder: member_k → (forcing_draw_i, hydro_draw_j) | +| `{cat-id}_hydro_draw_states.parquet` | Exact perturbed initial states for each of 20 hydro draws | +| `{cat-id}_forcing_draw_sequences.parquet` | Exact P/PET scale factors for each of 30 forcing draws | + +--- + +## Server Paths (dualearth1) + +All paths under `/mnt/disk2/` unless noted. + +| Resource | Path | +|---|---| +| Retro forcing | `suma_helen_poster/nwm_retro_catchment_forcings/` | +| Test forcing 1 | `/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings/` | +| Test forcing 2 | `/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings/` | +| Obs (F1/F2, static σ²) | `1400_sites_helene/catchment_ts_03463300_with_variance/` | +| Obs (F3/F4, dynamic σ²) | `1400_sites_helene/catchment_ts_03463300_spliced_dyn_helene/` | +| BMI config | `suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json` | +| Param bounds | `suma_helen_poster/run_gpu/CFE_parameter_bounds.json` | +| CFE source | `suma_helen_poster/cfe_py/` | +| GPKG | `/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg` | +| USGS hourly obs | `suma_helen_poster/03463300_usgs_hourly_2018_2024.csv` | +| F1 analysis | `suma_helen_poster/da_results/v2_true_enkf_vrugt/` | +| F2 analysis | `suma_helen_poster/da_results/v2_fixed_r007_analysis/` | +| F3 analysis | `1400_sites_helene/da_results_dynamic_vrugt_seeded/` | +| F4 analysis | `1400_sites_helene/da_results_dynamic_novrugt_seeded/` | + +See [EXPERIMENTS.md](EXPERIMENTS.md) for full server paths including forecast cycles, +ensemble directories, and routed output paths for all 4 folders. diff --git a/da_methods/helene_da_results.pptx b/da_methods/helene_da_results.pptx new file mode 100644 index 00000000..66a30706 Binary files /dev/null and b/da_methods/helene_da_results.pptx differ diff --git a/da_methods/pf.json b/da_methods/pf.json deleted file mode 100644 index afcdb214..00000000 --- a/da_methods/pf.json +++ /dev/null @@ -1 +0,0 @@ -{"n": 1, "smcmax":0.4300243541374813, "m":1, "N": 1000, "D":2} \ No newline at end of file diff --git a/da_methods/pf_method.py b/da_methods/pf_method.py deleted file mode 100644 index 34fb331d..00000000 --- a/da_methods/pf_method.py +++ /dev/null @@ -1,43 +0,0 @@ -import numpy as np -from scipy.stats import norm - - - -class particle_filter: - def __init__(self, N, n, m, smcmax, D): - self.N = N # number of ensembles - self.n = n # number of states - self.current_step = 0 - self.smcmax = smcmax - self.D = D - self.m = m - self.storage_max_m = self.smcmax * self.D - self.storage_init = self.storage_max_m * 0.667 - self.outputs = np.zeros((self.m, self.N)) - # self.state_estimates = np.zeros((self.n)) - self.likelihoods = np.zeros(self.N) - # self.observation = np.zeros((self.n)) - - def add_noise (self, simulations): - perturbation_factor_sim = np.random.randn() - error_factor_sim = 0.2 - for i in range(self.N): - self.outputs[:, i] = simulations[i] - self.outputs[:, i] += (perturbation_factor_sim * simulations[i] * error_factor_sim) - return self.outputs - - def calculate_likelihood(self): - # Estimate the parameters of the Gaussian distribution - mu, sigma = norm.fit(observation) - # Calculate the PDF for each number - self.likelihoods = [norm.pdf(number, loc=mu, scale=sigma) for number in self.outputs] - return self.likelihoods - - - def execute_pf (self): - weights = self.likelihoods / self.likelihoods.sum() - combined = list(zip(self.state_variables, self.outputs, weights)) - sorted_combined = sorted(combined, key=lambda x: x[2]) - sorted_states = [x[0] for x in sorted_combined] - self.state_estimates = np.random.choice(sorted_states, size=1, replace=True, p=weights) - return self.state_estimates \ No newline at end of file diff --git a/da_methods/run_pf_synthetic.ipynb b/da_methods/run_pf_synthetic.ipynb deleted file mode 100644 index 6a8567e8..00000000 --- a/da_methods/run_pf_synthetic.ipynb +++ /dev/null @@ -1,315 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": 1, - "id": "8956f120-6cc1-474c-88f7-2bd868c1c726", - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "#! pip install bmipy\n", - "#! pip install eto\n", - "import sys \n", - "import time\n", - "import numpy as np\n", - "import pandas as pd\n", - "import json\n", - "import matplotlib.pyplot as plt\n", - "import bmi_cfe_perturb_ens\n", - "import bmi_pf\n", - "import bmi_cfe\n", - "import bmi_PET\n", - "from fao_pet import ETRCalculator" - ] - }, - { - "cell_type": "code", - "execution_count": 2, - "id": "8bf7c5f2-a1a7-46ac-a569-ca1661202411", - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "pet_instance = bmi_PET.BMI_pet_model()\n", - "pet_instance.initialize(cfg_file='./pet.json', current_time_step=1)\n", - "pet_instance.update()\n", - "pet_instance.finalize()\n", - "pet = pet_instance.get_value('etr').values" - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "id": "978812da-d61b-42c7-9fca-d3c5c6a26d21", - "metadata": { - "tags": [] - }, - "outputs": [ - { - "ename": "AttributeError", - "evalue": "'particle_filter_bmi' object has no attribute 'simulations'", - "output_type": "error", - "traceback": [ - "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", - "\u001b[0;31mAttributeError\u001b[0m Traceback (most recent call last)", - "Cell \u001b[0;32mIn[3], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m enkf_instance \u001b[38;5;241m=\u001b[39m \u001b[43mbmi_pf\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mparticle_filter_bmi\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43m./pf.json\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m 2\u001b[0m cfe_purturb_instance \u001b[38;5;241m=\u001b[39m bmi_cfe_perturb_ens\u001b[38;5;241m.\u001b[39mBMI_CFE(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m./new_cfe_perturb.json\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m 3\u001b[0m cfe_instance \u001b[38;5;241m=\u001b[39m bmi_cfe\u001b[38;5;241m.\u001b[39mBMI_CFE(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m./new_cfe.json\u001b[39m\u001b[38;5;124m'\u001b[39m)\n", - "File \u001b[0;32m~/home/jovyan/data_assimilation/SI2023/assimilators/synthetic/bmi_pf.py:20\u001b[0m, in \u001b[0;36mparticle_filter_bmi.__init__\u001b[0;34m(self, json_file)\u001b[0m\n\u001b[1;32m 18\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39moutputs \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mzeros((\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mn))\n\u001b[1;32m 19\u001b[0m \u001b[38;5;66;03m# self.simulations = np.zeros((self.m, self.N))\u001b[39;00m\n\u001b[0;32m---> 20\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpf \u001b[38;5;241m=\u001b[39m particle_filter(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mN, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mn, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mm, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39msmcmax, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mD, \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msimulations\u001b[49m)\n\u001b[1;32m 21\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstate_estimates \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mzeros((\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mn))\n\u001b[1;32m 22\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mobservation \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mzeros((\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mn))\n", - "\u001b[0;31mAttributeError\u001b[0m: 'particle_filter_bmi' object has no attribute 'simulations'" - ] - } - ], - "source": [ - "enkf_instance = bmi_pf.particle_filter_bmi('./pf.json')\n", - "cfe_purturb_instance = bmi_cfe_perturb_ens.BMI_CFE('./new_cfe_perturb.json')\n", - "cfe_instance = bmi_cfe.BMI_CFE('./new_cfe.json')" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "a9e47055-8e4d-4a3d-809f-16e01d568011", - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "#initialization\n", - "cfe_purturb_instance.initialize()\n", - "cfe_instance.initialize()\n", - "enkf_instance.initialize()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "27f403fd-9791-43d0-8e53-4e2214f83a2a", - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "cfe_purturb_inputs = cfe_purturb_instance.get_input_var_names()\n", - "input_lists_cfe_purturb = {inputs:{ens:[] for ens in range(0)} for inputs in cfe_purturb_inputs}\n", - "print(\"CFE Perturbed code inputs are:\", input_lists_cfe_purturb)\n", - "\n", - "cfe_inputs = cfe_instance.get_input_var_names()\n", - "input_lists_cfe = {inputs:{ens:[] for ens in range(0)} for inputs in cfe_inputs}\n", - "print(\"CFE code inputs are:\", input_lists_cfe)\n", - "\n", - "enkf_inputs = enkf_instance.get_input_var_names()\n", - "input_lists_enkf = {inputs:{ens:[] for ens in range(0)} for inputs in enkf_inputs}\n", - "print(\"EnKF code inputs are:\", input_lists_enkf)\n", - "\n", - "\n", - "cfe_purturb_outputs = cfe_purturb_instance.get_output_var_names()\n", - "output_lists_cfe_purturb = {output:{ens:[] for ens in range(0)} for output in cfe_purturb_outputs}\n", - "print(\"CFE Perturbed code outputs are:\", output_lists_cfe_purturb)\n", - "\n", - "\n", - "cfe_outputs = cfe_instance.get_output_var_names()\n", - "output_lists_cfe = {output:{ens:[] for ens in range(0)} for output in cfe_outputs}\n", - "print(\"CFE code outputs are:\", output_lists_cfe)\n", - "\n", - "enkf_outputs = enkf_instance.get_output_var_names()\n", - "output_lists_enkf = {output:{ens:[] for ens in range(0)} for output in enkf_outputs}\n", - "print(\"EnKF code outputs are:\", output_lists_enkf)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "6cedde09-76d3-44eb-af5c-1d5e3bebe753", - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "cfe_purturb_outputs = cfe_purturb_instance.get_output_var_names()\n", - "output_lists_cfe_purturb = {output:{ens:[] for ens in range(cfe_purturb_instance.n_cfe_ensembles)} for output in cfe_outputs}\n", - "\n", - "cfe_outputs = cfe_instance.get_output_var_names()\n", - "output_lists_cfe = {output:{ens:[] for ens in range(cfe_purturb_instance.n_cfe_ensembles)} for output in cfe_outputs}\n", - "\n", - "enkf_outputs = enkf_instance.get_output_var_names()\n", - "output_lists_enkf = {output:{ens:[] for ens in range(cfe_purturb_instance.n_cfe_ensembles)} for output in enkf_outputs}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "dc030590-e5c8-49a6-9806-83e3eaa00087", - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "### Observation data\n", - "\n", - "observation_no_noise = []\n", - "\n", - "\n", - "with open(cfe_instance.forcing_file, 'r') as f:\n", - " df_forcing = pd.read_csv(f)\n", - " \n", - "for t, precip in enumerate(df_forcing['APCP_surface']):\n", - " cfe_instance.set_value('atmosphere_water__time_integral_of_precipitation_mass_flux', precip)\n", - " cfe_instance.set_value('water_potential_evaporation_flux', pet[t])\n", - " cfe_instance.update()\n", - " obs = cfe_instance.get_value(\"land_surface_water__runoff_volume_flux\")\n", - " observation_no_noise.append(obs)\n", - " \n", - "cfe_instance.finalize()\n", - "\n", - "perturbation_factor_obs = np.random.standard_normal(size=None)\n", - "error_factor_obs = 0.15\n", - "observation_no_noise = np.array(observation_no_noise) \n", - "observation_noisy = (perturbation_factor_obs * observation_no_noise * error_factor_obs) + observation_no_noise\n", - "\n", - "\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "4871177c-87e2-43d2-8e0f-291e55777ae5", - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "######Synthetic loop\n", - "\n", - "cfe_purturb_instance.initialize()\n", - "cfe_instance.initialize()\n", - "\n", - "\n", - "with open(cfe_purturb_instance.forcing_file, 'r') as f:\n", - " df_forcing = pd.read_csv(f)\n", - "\n", - "\n", - "da_simulations_all = []\n", - "da_simulations_all_mean = []\n", - "state_variables_all = []\n", - "state_variables_all_mean = []\n", - "simulations_all = []\n", - "observations_all = []\n", - "updated_state_all = []\n", - "updated_simulations_all = []\n", - "gw_state_all = []\n", - "gw_state_all_mean = []\n", - "updated_gw_state_all = []\n", - "\n", - "\n", - "for t, precip in enumerate(df_forcing['APCP_surface']):\n", - " \n", - " cfe_purturb_instance.set_value('atmosphere_water__time_integral_of_precipitation_mass_flux', precip)\n", - " cfe_purturb_instance.set_value('water_potential_evaporation_flux', pet[t])\n", - " cfe_purturb_instance.update()\n", - " simulations = cfe_purturb_instance.get_value(\"land_surface_water__runoff_volume_flux\")\n", - " simulations = np.array(list(simulations.values()))\n", - " da_simulations_all.append(simulations)\n", - " mean_simulations = np.mean(simulations)\n", - " da_simulations_all_mean.append(mean_simulations)\n", - " \n", - " state_variables = cfe_purturb_instance.get_value(\"SOIL_CONCEPTUAL_STORAGE\")\n", - " state_variablesssss = np.array(list(state_variables))\n", - " state_variables_all.append(state_variables)\n", - " state_variables_mean = np.mean(state_variablesssss)\n", - " state_variables_all_mean.append(state_variables_mean)\n", - " \n", - " gw_state = cfe_purturb_instance.get_value(\"DEEP_GW_TO_CHANNEL_FLUX\")\n", - " gw_statesssss = np.array(list(gw_state))\n", - " gw_state_all.append(gw_state)\n", - " gw_state_mean = np.mean(gw_statesssss)\n", - " gw_state_all_mean.append(gw_state_mean)\n", - " \n", - "\n", - " Flow = observation_noisy[t]\n", - " observations_all.append(Flow)\n", - " \n", - " enkf_instance.set_value('Observations', Flow) \n", - " enkf_instance.set_value('Simulations', simulations)\n", - " enkf_instance.set_value('State Variables', state_variables)\n", - " enkf_instance.update()\n", - " updated_state_sm = enkf_instance.get_value(\"state_estimates\")\n", - " updated_state_all.append(updated_state_sm)\n", - " updated_state_sm = updated_state_sm[0]\n", - " \n", - " enkf_instance.finalize()\n", - " cfe_instance.initialize()\n", - " print(simulations)\n", - "\n", - " enkf_instance.set_value('Observations', Flow) \n", - " enkf_instance.set_value('Simulations', simulations)\n", - " enkf_instance.set_value('State Variables', gw_state)\n", - " enkf_instance.update()\n", - " updated_state_gw = enkf_instance.get_value(\"state_estimates\")\n", - " updated_gw_state_all.append(updated_state_gw)\n", - " updated_state_gw = updated_state_gw[0]\n", - " \n", - " enkf_instance.finalize()\n", - " cfe_instance.initialize()\n", - "\n", - " cfe_instance.set_value('atmosphere_water__time_integral_of_precipitation_mass_flux', precip)\n", - " cfe_instance.set_value('water_potential_evaporation_flux', pet[t])\n", - " cfe_instance.set_value('SOIL_CONCEPTUAL_STORAGE', updated_state_sm)\n", - " cfe_instance.set_value('DEEP_GW_TO_CHANNEL_FLUX', updated_state_gw)\n", - " cfe_instance.update()\n", - " simulation_updated = cfe_instance.get_value(\"land_surface_water__runoff_volume_flux\")\n", - " updated_simulations_all.append(simulation_updated)\n", - " \n", - " list_of_states_sm = [updated_state_sm] * 100\n", - " cfe_purturb_instance.set_value('SOIL_CONCEPTUAL_STORAGE', list_of_states_sm)\n", - " \n", - " \n", - " list_of_states_gw = [updated_state_gw] * 100\n", - " cfe_purturb_instance.set_value('DEEP_GW_TO_CHANNEL_FLUX', list_of_states_gw)\n", - " \n", - "cfe_instance.finalize()\n", - "cfe_purturb_instance.finalize()\n", - "enkf_instance.finalize()\n", - "\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "09a592a0-bd71-4eee-855d-656383b88eb0", - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "920b7542-4ddd-49fe-9f2c-9027e9f181a5", - "metadata": {}, - "outputs": [], - "source": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.9.16" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/da_methods/test_1_heldout_gauge/1_calibration/batch_calibrate_all_cats.sh b/da_methods/test_1_heldout_gauge/1_calibration/batch_calibrate_all_cats.sh new file mode 100644 index 00000000..143bccc1 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/1_calibration/batch_calibrate_all_cats.sh @@ -0,0 +1,57 @@ +#!/usr/bin/env bash +# Batch DDS calibration for all 21 catchments — 1 held-out gauge Qkrig obs. +# Runs one background job per catchment; logs go to $OUT_DIR/_cal.log. +# After all jobs finish, check logs for any non-zero exit. +# +# Usage (from server home dir, troute env active): +# bash batch_calibrate_all_cats.sh + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +PYTHON=/home/svyas/miniconda3/envs/troute/bin/python3 +CALIBRATE=/mnt/disk2/suma_helen_poster/calibrate_catchment_nwm.py +FORCING_DIR=/mnt/disk2/suma_helen_poster/nwm_retro_catchment_forcings +OBS_DIR=$HOME/catchment_ts_no_03463300_gapfilled +CFE_DIR=/mnt/disk2/suma_helen_poster/cfe_py +CONFIG_FILE=/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json +PARAM_BOUNDS=/mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json +OUT_DIR=/mnt/disk2/suma_helen_poster/catchment_results_1gauge_heldout +TEST_FORCING_DIR1=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings +TEST_FORCING_DIR2=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings +N=1000 + +mkdir -p "$OUT_DIR" +echo "Starting calibration for ${#CATS[@]} catchments" +echo " obs-dir : $OBS_DIR" +echo " out-dir : $OUT_DIR" +echo "" + +for CAT in "${CATS[@]}"; do + LOG="$OUT_DIR/${CAT}_cal.log" + echo " Launching $CAT → $LOG" + "$PYTHON" "$CALIBRATE" \ + --cat-id "$CAT" \ + --forcing-dir "$FORCING_DIR" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --param-bounds "$PARAM_BOUNDS" \ + --out-dir "$OUT_DIR" \ + --test-forcing-dir1 "$TEST_FORCING_DIR1" \ + --test-forcing-dir2 "$TEST_FORCING_DIR2" \ + --N "$N" \ + > "$LOG" 2>&1 & +done + +echo "" +echo "All ${#CATS[@]} jobs launched. Waiting for completion..." +wait +echo "" +echo "All catchments done. Check for failures:" +grep -l "Error\|Traceback\|error" "$OUT_DIR"/cat-*_cal.log 2>/dev/null || echo " No errors found in logs." diff --git a/da_methods/test_1_heldout_gauge/README.md b/da_methods/test_1_heldout_gauge/README.md new file mode 100644 index 00000000..570ee91b --- /dev/null +++ b/da_methods/test_1_heldout_gauge/README.md @@ -0,0 +1,79 @@ +# Test: 1 Held-Out Gauge (03463300) + +## What this experiment tests + +Same four DA R-formula experiments as `test_20pct_heldout_gauges/`, but with a +fundamentally better Qkrig observation field: + +| | 20pct holdout | **This experiment** | +|---|---|---| +| Gauges withheld from kriging | ~20% of all South Toe gauges | Only gauge 03463300 (the outlet we evaluate) | +| Qkrig quality near outlet | Degraded — many local gauges missing | Much better — only 1 gauge missing | +| Expected DA performance | Baseline | **Should improve** — EnKF gets better observations | + +The hypothesis: with only one gauge held out, the kriging field is far more +accurate near the South Toe outlet. The EnKF should therefore correct CFE states +more reliably, improving routed streamflow at 03463300. + +--- + +## Steps to run (in order) + +### Step 1 — Get new Qkrig observations (PENDING DATA) +New per-catchment Qkrig time series with only gauge 03463300 held out. + +``` +# TODO: update when data is available +NEW_OBS_DIR = /mnt/disk2/???/catchment_ts_1gauge_heldout_with_variance/ +``` + +### Step 2 — Re-calibrate CFE (uses shared script) +Run `1_distributed_cfe/calibrate_catchment_cfe_da_v2.py` for all 21 catchments +using the new obs dir above. + +```bash +python ../../1_distributed_cfe/calibrate_catchment_cfe_da_v2.py \ + --cat-id cat-XXXXXXX \ + --obs-dir \ # <-- changes from 20pct version + --forcing-dir \ + --cfe-dir /mnt/disk2/suma_helen_poster/cfe_py \ + --config-file \ # <-- new per-catchment config from re-cal + --param-bounds /mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json \ + --N 1000 +``` + +Outputs per catchment: `best_params.json`, `cat-XXXXXXX_cal_results.csv` + +### Step 3 — T-route routing (uses shared script) +Same as 20pct version. Run `2_troute_routing/run_route.py` using the new +calibration outputs from Step 2. + +### Step 4 — Run DA experiments (F1–F4) +Scripts are copied from `test_20pct_heldout_gauges/`. Before running, update +these paths in each folder's `2_assimilation/run_perturbation_da_on.py`: + +```python +# Lines to update in run_perturbation_da_on.py for each of F1–F4: +--obs-dir # 1-gauge holdout Qkrig +--config-file # from re-calibration +--out-dir # separate from 20pct results +``` + +Static and dynamic variance paths for F1/F3/F4 will also need updating +when the new calibration results are available. + +### Step 5 — Evaluation +Same evaluation scripts as `test_20pct_heldout_gauges/`. Compare results +against 20pct to quantify the improvement from using better Qkrig observations. + +--- + +## PENDING — waiting on + +- [ ] New Qkrig obs dir path (1 gauge held out) +- [ ] New per-catchment config files from re-calibration +- [ ] New output directory paths on server +- [ ] Static/dynamic variance paths for F3/F4 + +Once these are provided, update paths in each `2_assimilation/run_perturbation_da_on.py` +and the corresponding batch shell scripts. diff --git a/da_methods/test_1_heldout_gauge/batch_plot_f5_lead_time_decay.sh b/da_methods/test_1_heldout_gauge/batch_plot_f5_lead_time_decay.sh new file mode 100644 index 00000000..82fbb2cb --- /dev/null +++ b/da_methods/test_1_heldout_gauge/batch_plot_f5_lead_time_decay.sh @@ -0,0 +1,41 @@ +#!/usr/bin/env bash +# Plot lead-time error decay curves for F5 — all 21 catchments. +# Run AFTER batch_run_f5_lead_time_sweep.sh has completed. +# +# Usage: +# bash batch_plot_f5_lead_time_decay.sh + +set -euo pipefail + +PYTHON=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4a_error_decay/plot_lead_time_decay.py +F5_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder5_rekrig_variance_direct + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +SKIP=0; DONE=0; FAIL=0 + +for CAT in "${CATS[@]}"; do + DA_CSV="$F5_DIR/$CAT/${CAT}_lead_time_forecasts_da.csv" + if [ ! -f "$DA_CSV" ]; then + echo " [$CAT] No lead-time CSV — skipping (run sweep first)" + SKIP=$((SKIP + 1)); continue + fi + + echo " [$CAT] plotting..." + "$PYTHON" "$SCRIPT" \ + --cat-id "$CAT" \ + --leadtime-dir "$F5_DIR" \ + --da-dir "$F5_DIR" \ + && DONE=$((DONE + 1)) \ + || { echo " FAILED: $CAT"; FAIL=$((FAIL + 1)); } +done + +echo "" +echo "=== F5 lead-time decay plots done: done=$DONE skipped=$SKIP failed=$FAIL ===" diff --git a/da_methods/test_1_heldout_gauge/batch_run_f5_lead_time_sweep.sh b/da_methods/test_1_heldout_gauge/batch_run_f5_lead_time_sweep.sh new file mode 100644 index 00000000..8ce7c5a5 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/batch_run_f5_lead_time_sweep.sh @@ -0,0 +1,77 @@ +#!/usr/bin/env bash +# Lead-time forecast sweep for F5 (re-kriged variance direct) — all 21 catchments. +# Runs run_lead_time_forecast_sweep.py with --direct-variance and --no-vrugt-r, +# generating _lead_time_forecasts_da.csv and _lead_time_forecasts_openloop.csv +# per catchment. Plot with batch_plot_f5_lead_time_decay.sh after this completes. +# +# Runtime: ~30-60 min per catchment (8760-hour test period × 2 trajectories). +# +# Usage: +# nohup bash batch_run_f5_lead_time_sweep.sh > ~/logs/f5_lead_time_sweep.log 2>&1 & + +set -euo pipefail + +PYTHON=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/run_lead_time_forecast_sweep.py +PROD_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/1_distributed_cfe/calibrate_catchment_cfe_da_v2.py + +OBS_DIR=/mnt/disk2/1400_sites_helene/catchment_ts_no_03463300_dynamic_variance_rekrig +CFE_DIR=/mnt/disk2/suma_helen_poster/cfe_py +CONFIG_FILE=/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json +PARAM_BOUNDS=/mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder5_rekrig_variance_direct +FORCING_DIR=/mnt/disk2/suma_helen_poster/nwm_retro_catchment_forcings +FORCING_DIR1=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings +FORCING_DIR2=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +echo "F5 lead-time sweep — out: $OUT_DIR" +echo "" + +SKIP=0; DONE=0; FAIL=0 + +for CAT in "${CATS[@]}"; do + echo "===============================" + echo "=== $CAT ===" + echo "===============================" + + DA_CSV="$OUT_DIR/$CAT/${CAT}_lead_time_forecasts_da.csv" + OL_CSV="$OUT_DIR/$CAT/${CAT}_lead_time_forecasts_openloop.csv" + + if [ -f "$DA_CSV" ] && [ -f "$OL_CSV" ]; then + echo " Already exists — skipping" + SKIP=$((SKIP + 1)); continue + fi + + if [ ! -f "$OUT_DIR/$CAT/${CAT}_best_params.json" ]; then + echo " WARNING: No best_params for $CAT — skipping" + SKIP=$((SKIP + 1)); continue + fi + + "$PYTHON" "$SCRIPT" \ + --cat-id "$CAT" \ + --forcing-dir "$FORCING_DIR" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --param-bounds "$PARAM_BOUNDS" \ + --out-dir "$OUT_DIR" \ + --test-forcing-dir1 "$FORCING_DIR1" \ + --test-forcing-dir2 "$FORCING_DIR2" \ + --enkf-members 20 \ + --no-vrugt-r \ + --direct-variance \ + --prod-script "$PROD_SCRIPT" \ + && DONE=$((DONE + 1)) \ + || { echo " FAILED: $CAT"; FAIL=$((FAIL + 1)); } +done + +echo "" +echo "=== F5 lead-time sweep done: done=$DONE skipped=$SKIP failed=$FAIL ===" diff --git a/da_methods/test_1_heldout_gauge/batch_run_openloop.sh b/da_methods/test_1_heldout_gauge/batch_run_openloop.sh new file mode 100644 index 00000000..fda15128 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/batch_run_openloop.sh @@ -0,0 +1,70 @@ +#!/usr/bin/env bash +# Open-loop CFE run (no DA) — 1 gauge holdout, all 21 catchments. +# Runs calibrate_catchment_cfe_da_v2.py WITHOUT --enkf-enabled. +# Uses best_params pre-staged from F4 (calibration is shared across experiments). +# Output: routed_Q_test.csv via run_route.py gives the no-DA baseline KGE. +# +# Usage: +# bash batch_run_openloop.sh > ~/logs/openloop.log 2>&1 + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PYTHON=/home/svyas/miniconda3/envs/troute/bin/python3 +PROD_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/1_distributed_cfe/calibrate_catchment_cfe_da_v2.py + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +OBS_DIR=/mnt/disk2/1400_sites_helene/catchment_ts_no_03463300_dynamic_variance +CFE_DIR=/mnt/disk2/suma_helen_poster/cfe_py +CONFIG_FILE=/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json +PARAM_BOUNDS=/mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/openloop +FORCING_DIR=/mnt/disk2/suma_helen_poster/nwm_retro_catchment_forcings +FORCING_DIR1=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings +FORCING_DIR2=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings + +echo "Open-loop (no DA) run — out: $OUT_DIR" +echo "" + +SKIP=0; DONE=0; FAIL=0 + +for CAT in "${CATS[@]}"; do + echo "===============================" + echo "=== $CAT ===" + echo "===============================" + + CAT_OUT="$OUT_DIR/$CAT" + OUT_FILE="$CAT_OUT/${CAT}_test_results.csv" + + if [ ! -f "$CAT_OUT/${CAT}_best_params.json" ]; then + echo " WARNING: No best_params for $CAT — skipping" + SKIP=$((SKIP + 1)); continue + fi + + if [ -f "$OUT_FILE" ]; then + echo " Already exists — skipping" + SKIP=$((SKIP + 1)); continue + fi + + "$PYTHON" "$PROD_SCRIPT" \ + --cat-id "$CAT" \ + --forcing-dir "$FORCING_DIR" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --param-bounds "$PARAM_BOUNDS" \ + --out-dir "$OUT_DIR" \ + --test-forcing-dir1 "$FORCING_DIR1" \ + --test-forcing-dir2 "$FORCING_DIR2" \ + --enkf-members 1 \ + && DONE=$((DONE + 1)) \ + || { echo " FAILED: $CAT"; FAIL=$((FAIL + 1)); } +done + +echo "" +echo "=== Open-loop done: done=$DONE skipped=$SKIP failed=$FAIL ===" diff --git a/da_methods/test_1_heldout_gauge/compare_all_folders.py b/da_methods/test_1_heldout_gauge/compare_all_folders.py new file mode 100644 index 00000000..7eab2dd8 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/compare_all_folders.py @@ -0,0 +1,215 @@ +""" +compare_all_folders.py + +Four-way comparison of DA analysis trajectories routed to gauge 03463300 +(South Toe River near Celo, NC), one curve per R-formula experiment: + + Open loop — No DA + Folder 1 — Variance scaled Vrugt R(t)=(0.10*y)^2 + 0.001*sigma2_krig + Folder 2 — Fixed R=0.07 R=0.07 (constant) + Folder 4 — Dynamic variance direct R(t)=sigma2_krig + Folder 5 — Re-kriged variance R(t)=sigma2_krig (re-kriged network) + +Reads routed_Q_test.csv (columns: date, Q_routed_m3s, Q_usgs_m3s) from each +folder. Skips any folder where the CSV does not exist. + +Outputs: + /compare_all_folders_full.png + /compare_all_folders_helene.png + /compare_all_folders_kge_table.csv + +Usage: + python3 compare_all_folders.py + python3 compare_all_folders.py --out-dir /path/to/out +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +HELENE_START = pd.Timestamp("2024-09-24") +HELENE_END = pd.Timestamp("2024-09-29 23:00:00") + +BASE = "/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout" + +FOLDERS = [ + { + "label": "Open loop\n(no DA)", + "short": "Open loop", + "color": "#7f7f7f", + "lw": 1.4, + "linestyle": "--", + "csv": f"{BASE}/openloop/routed_Q_test.csv", + }, + { + "label": "F1: Variance scaled Vrugt\nR(t)=(0.10·y)²+0.001·σ²_krig", + "short": "F1 Vrugt", + "color": "#1f77b4", + "lw": 2.0, + "linestyle": "-", + "csv": f"{BASE}/folder1_vrugt/routed_Q_test.csv", + }, + { + "label": "F2: Fixed R=0.07", + "short": "F2 R=0.07", + "color": "#ff7f0e", + "lw": 1.6, + "linestyle": "-", + "csv": f"{BASE}/folder2_fixed_r007/routed_Q_test.csv", + }, + { + "label": "F4: Dynamic variance direct\nR(t)=σ²_krig", + "short": "F4 Direct σ²", + "color": "#d62728", + "lw": 1.4, + "linestyle": "-", + "csv": f"{BASE}/folder4_dynamic_variance_direct/routed_Q_test.csv", + }, + { + "label": "F5: Re-kriged variance\nR(t)=σ²_krig (re-kriged network)", + "short": "F5 Rekrig σ²", + "color": "#9467bd", + "lw": 2.2, + "linestyle": "-", + "csv": f"{BASE}/folder5_rekrig_variance_direct/routed_Q_test.csv", + }, +] + + +def kge(obs, sim): + mask = np.isfinite(obs) & np.isfinite(sim) + o, s = obs[mask], sim[mask] + if len(o) < 2 or np.std(o) == 0: + return np.nan + r = np.corrcoef(o, s)[0, 1] + return 1.0 - np.sqrt((r - 1) ** 2 + (np.std(s) / np.std(o) - 1) ** 2 + + (np.mean(s) / np.mean(o) - 1) ** 2) + + +def nse(obs, sim): + mask = np.isfinite(obs) & np.isfinite(sim) + o, s = obs[mask], sim[mask] + if len(o) < 2: + return np.nan + denom = np.sum((o - o.mean()) ** 2) + return 1.0 - np.sum((o - s) ** 2) / denom if denom > 0 else np.nan + + +def load_folder(cfg): + p = cfg["csv"] + if not os.path.exists(p): + print(f" MISSING: {p}") + return None + df = pd.read_csv(p, parse_dates=["date"]).set_index("date").sort_index() + return df + + +def plot_comparison(dfs, obs, dates, helene_mask, out_path, helene_only=False): + fig, ax = plt.subplots(figsize=(15, 6)) + + if helene_only: + plot_dates = dates[helene_mask] + obs_plot = obs[helene_mask] + else: + plot_dates = dates + obs_plot = obs + + ax.plot(plot_dates, obs_plot, + color="black", lw=2.2, zorder=6, label="USGS obs") + + for cfg, df in zip(FOLDERS, dfs): + if df is None: + continue + sim = df["Q_routed_m3s"].reindex(dates).values + sim_plot = sim[helene_mask] if helene_only else sim + obs_sel = obs[helene_mask] if helene_only else obs + kg = kge(obs_sel, sim_plot) + ns = nse(obs_sel, sim_plot) + ax.plot(plot_dates, sim_plot, + color=cfg["color"], lw=cfg["lw"], + linestyle=cfg.get("linestyle", "-"), zorder=4, + label=f"{cfg['short']} KGE={kg:+.3f} NSE={ns:+.3f}") + + if not helene_only: + ax.axvspan(HELENE_START, HELENE_END, color="gold", alpha=0.12, zorder=1) + + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_xlabel("Date (UTC)", fontsize=11) + title = ("Helene window — " if helene_only else "Full test period — ") + ax.set_title(title + "DA R-formula comparison at gauge 03463300\n" + "(1-gauge holdout, test Oct 2023 – Oct 2024)", fontsize=12) + ax.legend(fontsize=8.5, loc="upper left", framealpha=0.9) + ax.grid(True, alpha=0.22) + + if helene_only: + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=1)) + peak_usgs = np.nanmax(obs[helene_mask]) + ax.axhline(peak_usgs, color="black", lw=0.7, linestyle=":", alpha=0.5) + ax.text(HELENE_END - pd.Timedelta(hours=6), peak_usgs * 1.02, + f"USGS peak {peak_usgs:.0f} m³/s", + fontsize=8.5, ha="right", color="black") + else: + ax.xaxis.set_major_formatter(mdates.DateFormatter("%Y-%m")) + ax.xaxis.set_major_locator(mdates.MonthLocator(interval=2)) + + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right", fontsize=9) + plt.tight_layout() + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--out-dir", default=f"{BASE}/comparison_plots") + args = parser.parse_args() + os.makedirs(args.out_dir, exist_ok=True) + + print("Loading routed CSVs...") + dfs = [load_folder(cfg) for cfg in FOLDERS] + + ref = next((df for df in dfs if df is not None), None) + if ref is None: + raise RuntimeError("No routed CSVs found — run routing scripts first.") + dates = ref.index + obs = ref["Q_usgs_m3s"].values if "Q_usgs_m3s" in ref.columns else np.full(len(dates), np.nan) + helene = (dates >= HELENE_START) & (dates <= HELENE_END) + + plot_comparison(dfs, obs, dates, helene, + os.path.join(args.out_dir, "compare_all_folders_full.png"), + helene_only=False) + plot_comparison(dfs, obs, dates, helene, + os.path.join(args.out_dir, "compare_all_folders_helene.png"), + helene_only=True) + + rows = [] + for cfg, df in zip(FOLDERS, dfs): + if df is None: + rows.append({"folder": cfg["short"], "full_kge": np.nan, + "full_nse": np.nan, "helene_kge": np.nan, + "helene_nse": np.nan, "helene_peak_m3s": np.nan}) + continue + sim = df["Q_routed_m3s"].reindex(dates).values + rows.append({ + "folder": cfg["short"], + "full_kge": round(kge(obs, sim), 3), + "full_nse": round(nse(obs, sim), 3), + "helene_kge": round(kge(obs[helene], sim[helene]), 3), + "helene_nse": round(nse(obs[helene], sim[helene]), 3), + "helene_peak_m3s": round(float(np.nanmax(sim[helene])), 1), + }) + + tbl = pd.DataFrame(rows) + tbl_path = os.path.join(args.out_dir, "compare_all_folders_kge_table.csv") + tbl.to_csv(tbl_path, index=False) + print(f"Saved: {tbl_path}") + print("\n" + tbl.to_string(index=False)) + print(f"\nUSGS Helene peak: {np.nanmax(obs[helene]):.1f} m³/s") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/cat-1016300_2a_forcing_arm_helene.png b/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/cat-1016300_2a_forcing_arm_helene.png new file mode 100644 index 00000000..32d0b2f3 Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/cat-1016300_2a_forcing_arm_helene.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/cat-1016300_2ab_arms_comparison.png b/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/cat-1016300_2ab_arms_comparison.png new file mode 100644 index 00000000..48c02d86 Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/cat-1016300_2ab_arms_comparison.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/cat-1016300_2b_hydro_arm_helene.png b/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/cat-1016300_2b_hydro_arm_helene.png new file mode 100644 index 00000000..d79555e8 Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/cat-1016300_2b_hydro_arm_helene.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/cat-1016300_4b_crossed_ensemble_helene.png b/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/cat-1016300_4b_crossed_ensemble_helene.png new file mode 100644 index 00000000..42caace9 Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/cat-1016300_4b_crossed_ensemble_helene.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/cat-1016300_4b_crossed_ensemble_peak.png b/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/cat-1016300_4b_crossed_ensemble_peak.png new file mode 100644 index 00000000..0119d41e Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/cat-1016300_4b_crossed_ensemble_peak.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/helene_routed_ensemble_vs_usgs_full.png b/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/helene_routed_ensemble_vs_usgs_full.png new file mode 100644 index 00000000..add14ef0 Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/helene_routed_ensemble_vs_usgs_full.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/helene_routed_ensemble_vs_usgs_peak.png b/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/helene_routed_ensemble_vs_usgs_peak.png new file mode 100644 index 00000000..78fe3146 Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/helene_routed_ensemble_vs_usgs_peak.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/helene_sensitivity_spaghetti_appendix.png b/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/helene_sensitivity_spaghetti_appendix.png new file mode 100644 index 00000000..de9a25e9 Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/helene_sensitivity_spaghetti_appendix.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/helene_sensitivity_spaghetti_main.png b/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/helene_sensitivity_spaghetti_main.png new file mode 100644 index 00000000..ea42e9a2 Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/helene_sensitivity_spaghetti_main.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/helene_sensitivity_spread_appendix.png b/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/helene_sensitivity_spread_appendix.png new file mode 100644 index 00000000..623b95fa Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/helene_sensitivity_spread_appendix.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/helene_sensitivity_spread_main.png b/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/helene_sensitivity_spread_main.png new file mode 100644 index 00000000..73345516 Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f1_variance_scaled_vrugt/helene_sensitivity_spread_main.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f2_fixed_r_007/cat-1016300_2a_forcing_arm_helene.png b/da_methods/test_1_heldout_gauge/figures/f2_fixed_r_007/cat-1016300_2a_forcing_arm_helene.png new file mode 100644 index 00000000..be7244ba Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f2_fixed_r_007/cat-1016300_2a_forcing_arm_helene.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f2_fixed_r_007/cat-1016300_2b_hydro_arm_helene.png b/da_methods/test_1_heldout_gauge/figures/f2_fixed_r_007/cat-1016300_2b_hydro_arm_helene.png new file mode 100644 index 00000000..acce6927 Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f2_fixed_r_007/cat-1016300_2b_hydro_arm_helene.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f2_fixed_r_007/cat-1016300_4b_crossed_ensemble_helene.png b/da_methods/test_1_heldout_gauge/figures/f2_fixed_r_007/cat-1016300_4b_crossed_ensemble_helene.png new file mode 100644 index 00000000..dd8566a3 Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f2_fixed_r_007/cat-1016300_4b_crossed_ensemble_helene.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f2_fixed_r_007/cat-1016300_4b_crossed_ensemble_peak.png b/da_methods/test_1_heldout_gauge/figures/f2_fixed_r_007/cat-1016300_4b_crossed_ensemble_peak.png new file mode 100644 index 00000000..988ff27c Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f2_fixed_r_007/cat-1016300_4b_crossed_ensemble_peak.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f2_fixed_r_007/helene_ensemble_twopanel.png b/da_methods/test_1_heldout_gauge/figures/f2_fixed_r_007/helene_ensemble_twopanel.png new file mode 100644 index 00000000..3e9c2881 Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f2_fixed_r_007/helene_ensemble_twopanel.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f2_fixed_r_007/helene_ensemble_vs_usgs.png b/da_methods/test_1_heldout_gauge/figures/f2_fixed_r_007/helene_ensemble_vs_usgs.png new file mode 100644 index 00000000..1f212e90 Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f2_fixed_r_007/helene_ensemble_vs_usgs.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f2_fixed_r_007/helene_routed_ensemble_vs_usgs_full.png b/da_methods/test_1_heldout_gauge/figures/f2_fixed_r_007/helene_routed_ensemble_vs_usgs_full.png new file mode 100644 index 00000000..41c952a5 Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f2_fixed_r_007/helene_routed_ensemble_vs_usgs_full.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f2_fixed_r_007/helene_routed_ensemble_vs_usgs_peak.png b/da_methods/test_1_heldout_gauge/figures/f2_fixed_r_007/helene_routed_ensemble_vs_usgs_peak.png new file mode 100644 index 00000000..dd6b4ff7 Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f2_fixed_r_007/helene_routed_ensemble_vs_usgs_peak.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f2_fixed_r_007/routed_ensemble_vs_usgs_full.png b/da_methods/test_1_heldout_gauge/figures/f2_fixed_r_007/routed_ensemble_vs_usgs_full.png new file mode 100644 index 00000000..41c952a5 Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f2_fixed_r_007/routed_ensemble_vs_usgs_full.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f2_fixed_r_007/routed_ensemble_vs_usgs_peak.png b/da_methods/test_1_heldout_gauge/figures/f2_fixed_r_007/routed_ensemble_vs_usgs_peak.png new file mode 100644 index 00000000..dd6b4ff7 Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f2_fixed_r_007/routed_ensemble_vs_usgs_peak.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f4_dynamic_variance_direct/cat-1016300_2a_forcing_arm_helene.png b/da_methods/test_1_heldout_gauge/figures/f4_dynamic_variance_direct/cat-1016300_2a_forcing_arm_helene.png new file mode 100644 index 00000000..6a858d2d Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f4_dynamic_variance_direct/cat-1016300_2a_forcing_arm_helene.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f4_dynamic_variance_direct/cat-1016300_2b_hydro_arm_helene.png b/da_methods/test_1_heldout_gauge/figures/f4_dynamic_variance_direct/cat-1016300_2b_hydro_arm_helene.png new file mode 100644 index 00000000..1cfbf528 Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f4_dynamic_variance_direct/cat-1016300_2b_hydro_arm_helene.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f4_dynamic_variance_direct/cat-1016300_4b_crossed_ensemble_helene.png b/da_methods/test_1_heldout_gauge/figures/f4_dynamic_variance_direct/cat-1016300_4b_crossed_ensemble_helene.png new file mode 100644 index 00000000..b7792014 Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f4_dynamic_variance_direct/cat-1016300_4b_crossed_ensemble_helene.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f4_dynamic_variance_direct/cat-1016300_4b_crossed_ensemble_peak.png b/da_methods/test_1_heldout_gauge/figures/f4_dynamic_variance_direct/cat-1016300_4b_crossed_ensemble_peak.png new file mode 100644 index 00000000..573fa972 Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f4_dynamic_variance_direct/cat-1016300_4b_crossed_ensemble_peak.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f4_dynamic_variance_direct/helene_ensemble_twopanel.png b/da_methods/test_1_heldout_gauge/figures/f4_dynamic_variance_direct/helene_ensemble_twopanel.png new file mode 100644 index 00000000..3d99493c Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f4_dynamic_variance_direct/helene_ensemble_twopanel.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f4_dynamic_variance_direct/helene_ensemble_vs_usgs.png b/da_methods/test_1_heldout_gauge/figures/f4_dynamic_variance_direct/helene_ensemble_vs_usgs.png new file mode 100644 index 00000000..996bd798 Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f4_dynamic_variance_direct/helene_ensemble_vs_usgs.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f4_dynamic_variance_direct/routed_ensemble_vs_usgs_full.png b/da_methods/test_1_heldout_gauge/figures/f4_dynamic_variance_direct/routed_ensemble_vs_usgs_full.png new file mode 100644 index 00000000..7d1e27af Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f4_dynamic_variance_direct/routed_ensemble_vs_usgs_full.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f4_dynamic_variance_direct/routed_ensemble_vs_usgs_peak.png b/da_methods/test_1_heldout_gauge/figures/f4_dynamic_variance_direct/routed_ensemble_vs_usgs_peak.png new file mode 100644 index 00000000..88e6af84 Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f4_dynamic_variance_direct/routed_ensemble_vs_usgs_peak.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f5_rekrig_variance_direct/cat-1016300_2a_forcing_arm_helene.png b/da_methods/test_1_heldout_gauge/figures/f5_rekrig_variance_direct/cat-1016300_2a_forcing_arm_helene.png new file mode 100644 index 00000000..6ea3c623 Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f5_rekrig_variance_direct/cat-1016300_2a_forcing_arm_helene.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f5_rekrig_variance_direct/cat-1016300_2b_hydro_arm_helene.png b/da_methods/test_1_heldout_gauge/figures/f5_rekrig_variance_direct/cat-1016300_2b_hydro_arm_helene.png new file mode 100644 index 00000000..7f12c565 Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f5_rekrig_variance_direct/cat-1016300_2b_hydro_arm_helene.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f5_rekrig_variance_direct/cat-1016300_4b_crossed_ensemble_helene.png b/da_methods/test_1_heldout_gauge/figures/f5_rekrig_variance_direct/cat-1016300_4b_crossed_ensemble_helene.png new file mode 100644 index 00000000..ef34b317 Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f5_rekrig_variance_direct/cat-1016300_4b_crossed_ensemble_helene.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f5_rekrig_variance_direct/cat-1016300_4b_crossed_ensemble_peak.png b/da_methods/test_1_heldout_gauge/figures/f5_rekrig_variance_direct/cat-1016300_4b_crossed_ensemble_peak.png new file mode 100644 index 00000000..73824993 Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f5_rekrig_variance_direct/cat-1016300_4b_crossed_ensemble_peak.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f5_rekrig_variance_direct/helene_ensemble_twopanel.png b/da_methods/test_1_heldout_gauge/figures/f5_rekrig_variance_direct/helene_ensemble_twopanel.png new file mode 100644 index 00000000..a7a00deb Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f5_rekrig_variance_direct/helene_ensemble_twopanel.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f5_rekrig_variance_direct/helene_ensemble_vs_usgs.png b/da_methods/test_1_heldout_gauge/figures/f5_rekrig_variance_direct/helene_ensemble_vs_usgs.png new file mode 100644 index 00000000..5ba3eaa2 Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f5_rekrig_variance_direct/helene_ensemble_vs_usgs.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f5_rekrig_variance_direct/routed_ensemble_vs_usgs_full.png b/da_methods/test_1_heldout_gauge/figures/f5_rekrig_variance_direct/routed_ensemble_vs_usgs_full.png new file mode 100644 index 00000000..af06f9f1 Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f5_rekrig_variance_direct/routed_ensemble_vs_usgs_full.png differ diff --git a/da_methods/test_1_heldout_gauge/figures/f5_rekrig_variance_direct/routed_ensemble_vs_usgs_peak.png b/da_methods/test_1_heldout_gauge/figures/f5_rekrig_variance_direct/routed_ensemble_vs_usgs_peak.png new file mode 100644 index 00000000..5249967b Binary files /dev/null and b/da_methods/test_1_heldout_gauge/figures/f5_rekrig_variance_direct/routed_ensemble_vs_usgs_peak.png differ diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/batch_run_crossed_f1.sh b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/batch_run_crossed_f1.sh new file mode 100644 index 00000000..4ac812e5 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/batch_run_crossed_f1.sh @@ -0,0 +1,64 @@ +#!/usr/bin/env bash +# Batch crossed ensemble for all 21 catchments — F1 Vrugt R formula. +# Reads best_params from the F1 DA out dir (already staged). +# Outputs _crossed_ensemble.parquet per catchment to the same dir. +# +# Usage: +# nohup bash ~/da_1gauge_f1/2_assimilation/batch_run_crossed_f1.sh > ~/crossed_f1.log 2>&1 & + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PYTHON=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT="$SCRIPT_DIR/run_crossed_ensemble.py" + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +OBS_DIR=/home/svyas/catchment_ts_no_03463300_gapfilled +CFE_DIR=/mnt/disk2/suma_helen_poster/cfe_py +CONFIG_FILE=/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt +FORCING_DIR1=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings +FORCING_DIR2=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings + +echo "F1 Vrugt crossed ensemble — out: $OUT_DIR" +echo "" + +SKIP=0; DONE=0; FAIL=0 + +for CAT in "${CATS[@]}"; do + echo "===============================" + echo "=== $CAT ===" + echo "===============================" + + CAT_OUT="$OUT_DIR/$CAT" + + if [ ! -f "$CAT_OUT/${CAT}_best_params.json" ]; then + echo " WARNING: No best_params for $CAT — skipping" + SKIP=$((SKIP + 1)); continue + fi + + if [ -f "$CAT_OUT/${CAT}_crossed_ensemble.parquet" ]; then + echo " Crossed ensemble already exists — skipping" + SKIP=$((SKIP + 1)); continue + fi + + "$PYTHON" "$SCRIPT" \ + --cat-id "$CAT" \ + --forcing-dir "$FORCING_DIR1" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --out-dir "$OUT_DIR" \ + --test-forcing-dir1 "$FORCING_DIR1" \ + --test-forcing-dir2 "$FORCING_DIR2" \ + && DONE=$((DONE + 1)) \ + || { echo " FAILED: $CAT"; FAIL=$((FAIL + 1)); } +done + +echo "" +echo "=== F1 crossed done: done=$DONE skipped=$SKIP failed=$FAIL ===" diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/batch_run_da_on_all_cats.sh b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/batch_run_da_on_all_cats.sh new file mode 100644 index 00000000..864eeba6 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/batch_run_da_on_all_cats.sh @@ -0,0 +1,81 @@ +#!/usr/bin/env bash +# Batch DA for all 21 catchments — F1 Vrugt R formula. +# R = (0.10 * Q)^2 + 0.001 * krig_var (no hardcoded-r) +# Stages best_params.json from calibration results before running. +# +# Usage (from server, after SCPing this folder): +# bash batch_run_da_on_all_cats.sh + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PYTHON=/home/svyas/miniconda3/envs/troute/bin/python3 +DA_SCRIPT="$SCRIPT_DIR/run_perturbation_da_on.py" + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +OBS_DIR=/home/svyas/catchment_ts_no_03463300_gapfilled +BEST_PARAMS_SRC=/mnt/disk2/suma_helen_poster/catchment_results_1gauge_heldout +CFE_DIR=/mnt/disk2/suma_helen_poster/cfe_py +CONFIG_FILE=/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json +PARAM_BOUNDS=/mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt +FORCING_DIR1=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings +FORCING_DIR2=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings + +mkdir -p "$OUT_DIR" +echo "F1 Vrugt DA — obs: $OBS_DIR" +echo " out: $OUT_DIR" +echo "" + +SKIP=0; DONE=0; FAIL=0 + +for CAT in "${CATS[@]}"; do + echo "===============================" + echo "=== $CAT ===" + echo "===============================" + + CAT_OUT="$OUT_DIR/$CAT" + mkdir -p "$CAT_OUT" + + # Stage best_params from calibration results + if [ ! -f "$CAT_OUT/${CAT}_best_params.json" ]; then + SRC="$BEST_PARAMS_SRC/$CAT/${CAT}_best_params.json" + if [ -f "$SRC" ]; then + cp "$SRC" "$CAT_OUT/" + echo " Staged best_params" + else + echo " WARNING: No best_params for $CAT — skipping" + SKIP=$((SKIP + 1)); continue + fi + else + echo " best_params already staged" + fi + + # Skip if both arm CSVs already exist + if [ -f "$CAT_OUT/${CAT}_da_forcing_arm.csv" ] && \ + [ -f "$CAT_OUT/${CAT}_da_hydro_arm.csv" ]; then + echo " Arms already complete — skipping" + SKIP=$((SKIP + 1)); continue + fi + + "$PYTHON" "$DA_SCRIPT" \ + --cat-id "$CAT" \ + --forcing-dir "$FORCING_DIR1" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --param-bounds "$PARAM_BOUNDS" \ + --out-dir "$OUT_DIR" \ + --test-forcing-dir1 "$FORCING_DIR1" \ + --test-forcing-dir2 "$FORCING_DIR2" \ + && DONE=$((DONE + 1)) \ + || { echo " FAILED: $CAT"; FAIL=$((FAIL + 1)); } +done + +echo "" +echo "=== F1 done: done=$DONE skipped=$SKIP failed=$FAIL ===" diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/batch_run_leadtime_f1.sh b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/batch_run_leadtime_f1.sh new file mode 100644 index 00000000..071636fc --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/batch_run_leadtime_f1.sh @@ -0,0 +1,76 @@ +#!/usr/bin/env bash +# Batch lead-time forecast sweep — F1 Vrugt, 1 gauge holdout. +# Runs run_lead_time_forecast_sweep.py for all 21 catchments. +# Uses Vrugt dynamic R formula (no --hardcoded-r flag). +# Skips if both output CSVs already exist. +# +# After this completes, run route_lead_time_forecasts.py to route the +# per-catchment CSVs through T-route and produce routed_leadtime_*.parquet. +# +# Usage: +# nohup bash ~/da_1gauge_f1/2_assimilation/batch_run_leadtime_f1.sh \ +# > ~/leadtime_f1.log 2>&1 & + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PYTHON=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT="$SCRIPT_DIR/run_lead_time_forecast_sweep.py" + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +OBS_DIR=/home/svyas/catchment_ts_no_03463300_gapfilled +CFE_DIR=/mnt/disk2/suma_helen_poster/cfe_py +CONFIG_FILE=/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json +PARAM_BOUNDS=/mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt +FORCING_DIR1=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings +FORCING_DIR2=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings + +echo "F1 lead-time sweep (Vrugt R) — out: $OUT_DIR" +echo "" + +SKIP=0; DONE=0; FAIL=0 + +for CAT in "${CATS[@]}"; do + echo "===============================" + echo "=== $CAT ===" + echo "===============================" + + CAT_OUT="$OUT_DIR/$CAT" + DA_CSV="$CAT_OUT/${CAT}_lead_time_forecasts_da.csv" + OL_CSV="$CAT_OUT/${CAT}_lead_time_forecasts_openloop.csv" + + if [ ! -f "$CAT_OUT/${CAT}_best_params.json" ]; then + echo " WARNING: No best_params for $CAT — skipping" + SKIP=$((SKIP + 1)); continue + fi + + if [ -f "$DA_CSV" ] && [ -f "$OL_CSV" ]; then + echo " Lead-time CSVs already exist — skipping" + SKIP=$((SKIP + 1)); continue + fi + + "$PYTHON" "$SCRIPT" \ + --cat-id "$CAT" \ + --forcing-dir "$FORCING_DIR1" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --param-bounds "$PARAM_BOUNDS" \ + --out-dir "$OUT_DIR" \ + --test-forcing-dir1 "$FORCING_DIR1" \ + --test-forcing-dir2 "$FORCING_DIR2" \ + --enkf-members 20 \ + --base-step-h 6 \ + --prod-script "$SCRIPT_DIR/calibrate_catchment_cfe_da_v2.py" \ + && DONE=$((DONE + 1)) \ + || { echo " FAILED: $CAT"; FAIL=$((FAIL + 1)); } +done + +echo "" +echo "=== F1 lead-time sweep done: done=$DONE skipped=$SKIP failed=$FAIL ===" diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/batch_run_production_f1.sh b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/batch_run_production_f1.sh new file mode 100644 index 00000000..d5fa8dde --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/batch_run_production_f1.sh @@ -0,0 +1,70 @@ +#!/usr/bin/env bash +# Batch production per-member run — F1 Vrugt, 1 gauge holdout. +# Runs run_production_per_member.py for all 21 catchments. +# Skips if output CSV already exists. +# +# Usage: +# nohup bash ~/da_1gauge_f1/2_assimilation/batch_run_production_f1.sh \ +# > ~/production_f1.log 2>&1 & + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PYTHON=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT="$SCRIPT_DIR/run_production_per_member.py" + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +OBS_DIR=/home/svyas/catchment_ts_no_03463300_gapfilled +CFE_DIR=/mnt/disk2/suma_helen_poster/cfe_py +CONFIG_FILE=/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json +PARAM_BOUNDS=/mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt +FORCING_DIR1=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings +FORCING_DIR2=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings + +echo "F1 production per-member — out: $OUT_DIR" +echo "" + +SKIP=0; DONE=0; FAIL=0 + +for CAT in "${CATS[@]}"; do + echo "===============================" + echo "=== $CAT ===" + echo "===============================" + + CAT_OUT="$OUT_DIR/$CAT" + OUT_FILE="$CAT_OUT/${CAT}_production_per_member.csv" + + if [ ! -f "$CAT_OUT/${CAT}_best_params.json" ]; then + echo " WARNING: No best_params for $CAT — skipping" + SKIP=$((SKIP + 1)); continue + fi + + if [ -f "$OUT_FILE" ]; then + echo " Production per-member CSV already exists — skipping" + SKIP=$((SKIP + 1)); continue + fi + + "$PYTHON" "$SCRIPT" \ + --cat-id "$CAT" \ + --forcing-dir "$FORCING_DIR1" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --param-bounds "$PARAM_BOUNDS" \ + --out-dir "$OUT_DIR" \ + --test-forcing-dir1 "$FORCING_DIR1" \ + --test-forcing-dir2 "$FORCING_DIR2" \ + --enkf-members 20 \ + --prod-script "$SCRIPT_DIR/calibrate_catchment_cfe_da_v2.py" \ + && DONE=$((DONE + 1)) \ + || { echo " FAILED: $CAT"; FAIL=$((FAIL + 1)); } +done + +echo "" +echo "=== F1 production per-member done: done=$DONE skipped=$SKIP failed=$FAIL ===" diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/batch_run_sensitivity_f1.sh b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/batch_run_sensitivity_f1.sh new file mode 100644 index 00000000..e4cee157 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/batch_run_sensitivity_f1.sh @@ -0,0 +1,68 @@ +#!/usr/bin/env bash +# Batch perturbation sensitivity analysis — F1 Vrugt, 1 gauge holdout. +# Runs run_perturbation_sensitivity.py for all 21 catchments × 3 sources +# (init, forcing, process). Skips if output CSV already exists. +# +# Usage: +# nohup bash ~/da_1gauge_f1/2_assimilation/batch_run_sensitivity_f1.sh \ +# > ~/sensitivity_f1.log 2>&1 & + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PYTHON=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT="$SCRIPT_DIR/run_perturbation_sensitivity.py" + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +SOURCES=(init forcing process) + +OBS_DIR=/home/svyas/catchment_ts_no_03463300_gapfilled +CFE_DIR=/mnt/disk2/suma_helen_poster/cfe_py +CONFIG_FILE=/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json +PARAM_BOUNDS=/mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt +FORCING_DIR1=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings +FORCING_DIR2=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings + +echo "F1 sensitivity analysis — out: $OUT_DIR" +echo "" + +SKIP=0; DONE=0; FAIL=0 + +for CAT in "${CATS[@]}"; do + for SRC in "${SOURCES[@]}"; do + OUT_FILE="$OUT_DIR/$CAT/${CAT}_sensitivity_${SRC}.csv" + + if [ -f "$OUT_FILE" ]; then + echo " [$CAT/$SRC] already exists — skipping" + SKIP=$((SKIP + 1)) + continue + fi + + echo "===============================" + echo "=== $CAT source=$SRC ===" + echo "===============================" + + "$PYTHON" "$SCRIPT" \ + --cat-id "$CAT" \ + --source "$SRC" \ + --forcing-dir "$FORCING_DIR1" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --param-bounds "$PARAM_BOUNDS" \ + --out-dir "$OUT_DIR" \ + --test-forcing-dir1 "$FORCING_DIR1" \ + --test-forcing-dir2 "$FORCING_DIR2" \ + && DONE=$((DONE + 1)) \ + || { echo " FAILED: $CAT/$SRC"; FAIL=$((FAIL + 1)); } + done +done + +echo "" +echo "=== F1 sensitivity done: done=$DONE skipped=$SKIP failed=$FAIL ===" diff --git a/da_methods/input_enkf_new.json b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/input_enkf_new.json similarity index 100% rename from da_methods/input_enkf_new.json rename to da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/input_enkf_new.json diff --git a/da_methods/new_EnKF.py b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/new_EnKF.py similarity index 100% rename from da_methods/new_EnKF.py rename to da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/new_EnKF.py diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/run_crossed_ensemble.py b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/run_crossed_ensemble.py new file mode 100644 index 00000000..325a64e2 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/run_crossed_ensemble.py @@ -0,0 +1,452 @@ +""" +run_crossed_ensemble.py -- 4b: 600-member crossed ensemble for Helene. + +Crosses the two DA-on perturbation arms: + Forcing arm : 30 met draws (from run_perturbation_da_on.py Phase 2A) + Hydro arm : 20 state draws (from run_perturbation_da_on.py Phase 2B) + +For each Helene issue time t0: + Run 600 members = 30 forcing draws x 20 hydro-state draws. + Each member (i, j) gets: + - Initial state = perturbed DA analysis state (draw j of hydro arm) + - Met forcing = perturbed precip/PET sequence (draw i of forcing arm) + +This gives the full uncertainty envelope combining both sources — +the operational probabilistic forecast picture. + +Output (per catchment): + //_crossed_ensemble.parquet + Columns: issue_time, lead_hour, member_0000..member_0599 + +Usage: + python3 run_crossed_ensemble.py \\ + --cat-id cat-1016300 \\ + --forcing-dir \\ + --obs-dir /mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance \\ + --cfe-dir /mnt/disk2/suma_helen_poster/cfe_py \\ + --config-file /mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json \\ + --out-dir /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble \\ + --test-forcing-dir1 \\ + --test-forcing-dir2 + +Strategy: + Re-runs DA analysis (same as Phase 1 of run_perturbation_da_on.py) to get + the per-issue-time analysis state. Then generates N_FORCING sets of perturbed + forcing sequences and N_HYDRO perturbed states. Runs all 600 combinations. + + To avoid re-running Phase 1, if the Phase 1 snapshot file exists from a + prior run_perturbation_da_on.py run, it is loaded directly. + +NOTE: All special characters kept ASCII to avoid encoding issues on GPU. +""" + +import argparse +import os +import sys +import json +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +N_FORCING = 30 +N_HYDRO = 20 +N_MEMBERS = N_FORCING * N_HYDRO # 600 + +TEST_START = "2024-08-24 00:00:00" # 1 month spinup before Helene window +TEST_END = "2024-10-31 23:00:00" + +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-30 06:00:00") +FORECAST_HOURS = 18 + +PRECIP_SIGMA = 0.15 +PET_SIGMA = 0.10 +STATE_FRAC = 0.05 + +CAT_ID = None +OBS_FILE = None +CFE_CONFIG_FILE = None +TEST_FORCING_FILE = None +OUT_DIR = None +HARDCODED_R = None +bmi_cfe = None + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(): + df = pd.read_csv(TEST_FORCING_FILE) + df["date"] = pd.to_datetime(df["time"]).dt.strftime("%Y-%m-%d %H:%M:%S") + df["total_precipitation"] = df["APCP_surface"] * 3600.0 + df["potential_evaporation"] = priestley_taylor_pet( + df["DSWRF_surface"].values, df["TMP_2maboveground"].values) + return df + + +def load_obs(): + df = pd.read_csv(OBS_FILE) + date_col = next((c for c in df.columns + if c.lower() in ("date", "time", "datetime", "timestamp")), None) + q_col = next((c for c in df.columns + if "qkrig" in c.lower() and "var" not in c.lower()), None) + var_col = next((c for c in df.columns if "var" in c.lower()), None) + df[date_col] = pd.to_datetime(df[date_col]) + df = df.sort_values(date_col).reset_index(drop=True) + date_strs = df[date_col].dt.strftime("%Y-%m-%d %H:%M:%S") + obs_dict = dict(zip(date_strs, df[q_col])) + var_vals = df[var_col] if var_col else pd.Series([0.07] * len(df)) + var_dict = dict(zip(date_strs, var_vals)) + return obs_dict, var_dict + + +def write_model_config(best_params): + """Write the BMI config file once; reused for every model instantiation.""" + with open(CFE_CONFIG_FILE) as f: + cfg = json.load(f) + cfg["forcing_file"] = TEST_FORCING_FILE + cfg["soil_params"]["bb"] = best_params["bb"] + cfg["soil_params"]["smcmax"] = best_params["smcmax"] + cfg["soil_params"]["satdk"] = best_params["satdk"] + cfg["slop"] = best_params["slop"] + cfg["max_gw_storage"] = best_params["max_gw_storage"] + cfg["expon"] = best_params["expon"] + cfg["Cgw"] = best_params["Cgw"] + cfg["K_lf"] = best_params["K_lf"] + cfg["K_nash"] = best_params["K_nash"] + cfg["partition_scheme"] = ("Schaake" if best_params["scheme"] <= 0.5 + else "Xinanjiang") + tmp_cfg = str(OUT_DIR / f"{CAT_ID}_bmi_config_crossed_stable.json") + with open(tmp_cfg, "w") as f: + json.dump(cfg, f) + return tmp_cfg + + +def make_model(cfg_path): + """Instantiate and initialize a CFE model from a pre-written config path.""" + m = bmi_cfe.BMI_CFE(cfg_file=cfg_path) + m.initialize() + return m + + +def get_state(m): + return { + "soil_m": m.soil_reservoir["storage_m"], + "gw_m": m.gw_reservoir["storage_m"], + "nash0": float(m.nash_storage[0]), + "nash1": float(m.nash_storage[1]), + } + + +def set_state(m, state): + m.soil_reservoir["storage_m"] = state["soil_m"] + m.gw_reservoir["storage_m"] = state["gw_m"] + m.nash_storage[0] = state["nash0"] + m.nash_storage[1] = state["nash1"] + + +def perturb_state(state, rng, frac=STATE_FRAC): + return { + "soil_m": max(state["soil_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "nash0": max(state["nash0"] * (1.0 + frac * rng.standard_normal()), 0.0), + "nash1": max(state["nash1"] * (1.0 + frac * rng.standard_normal()), 0.0), + } + + +def run_model_step(m, precip_mmh, pet_mmh): + m.set_value("atmosphere_water__time_integral_of_precipitation_mass_flux", + precip_mmh / 1000.0) + m.set_value("water_potential_evaporation_flux", + pet_mmh / 1000.0 / 3600.0) + m.update() + return m.get_value("land_surface_water__runoff_depth") * 1000.0 + + +def enkf_update(state, q_sim, y_obs, krig_var, rng): + R = HARDCODED_R if HARDCODED_R is not None else max( + (0.10 * max(y_obs, 0.0)) ** 2 + 0.001 * krig_var, 1e-6) + R = max(R, 1e-6) + P_yy = max(q_sim * 0.01, 1e-6) + K = P_yy / (P_yy + R) + scale = K * (y_obs - q_sim) / max(abs(q_sim), 1e-6) + return { + "soil_m": max(state["soil_m"] * (1.0 + scale), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + scale), 1e-6), + "nash0": max(state["nash0"] + state["nash0"] * scale, 0.0), + "nash1": max(state["nash1"] + state["nash1"] * scale, 0.0), + } + + +def run_phase1_da(cfg_path, df_test, obs_dict, var_dict, rng): + """Run DA through test period; return snapshot dict {timestamp -> state}.""" + print(f" Phase 1: DA analysis through test period...") + model = make_model(cfg_path) + snapshots = {} + for _, row in df_test.iterrows(): + date_str = row["date"] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + q_sim = run_model_step(model, P, E) + y_obs = obs_dict.get(date_str, np.nan) + if not np.isnan(y_obs): + krig_var = var_dict.get(date_str, 0.07) + state = get_state(model) + updated = enkf_update(state, q_sim, y_obs, krig_var, rng) + set_state(model, updated) + date_ts = pd.Timestamp(date_str) + if HELENE_START <= date_ts <= HELENE_END: + snapshots[date_ts] = get_state(model) + model.finalize() + print(f" Phase 1 done: {len(snapshots)} issue-time snapshots") + return snapshots + + +def main(): + global CAT_ID, OBS_FILE, CFE_CONFIG_FILE, TEST_FORCING_FILE + global OUT_DIR, HARDCODED_R, bmi_cfe + + parser = argparse.ArgumentParser() + parser.add_argument("--cat-id", required=True) + parser.add_argument("--forcing-dir", required=True) + parser.add_argument("--obs-dir", required=True) + parser.add_argument("--cfe-dir", required=True) + parser.add_argument("--config-file", required=True) + parser.add_argument("--out-dir", required=True) + parser.add_argument("--test-forcing-dir1", default=None) + parser.add_argument("--test-forcing-dir2", default=None) + parser.add_argument("--hardcoded-r", type=float, default=None, + help="Fixed R value for EnKF (e.g. 0.07). " + "Omit to use dynamic Vrugt R: R=(0.10*y)^2 + 0.001*krig_var") + parser.add_argument("--n-forcing", type=int, default=N_FORCING, + help="Number of met forcing draws (default 30)") + parser.add_argument("--n-hydro", type=int, default=N_HYDRO, + help="Number of hydro-state draws (default 20)") + args = parser.parse_args() + + CAT_ID = args.cat_id + OBS_FILE = os.path.join(args.obs_dir, f"{CAT_ID}.csv") + CFE_CONFIG_FILE = args.config_file + HARDCODED_R = args.hardcoded_r + OUT_DIR = Path(args.out_dir) / CAT_ID + OUT_DIR.mkdir(parents=True, exist_ok=True) + + n_fa = args.n_forcing + n_ha = args.n_hydro + n_total = n_fa * n_ha + print(f"[{CAT_ID}] Crossed ensemble: {n_fa} forcing x {n_ha} hydro = {n_total} members") + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + bmi_cfe = _bmi_cfe + + # Reuse combined forcing CSV from run_perturbation_da_on.py if it exists + # (avoids writing a duplicate large file to disk) + da_on_combined = (Path(args.out_dir).parent / "v2_perturbation_da_on" / + CAT_ID / f"{CAT_ID}_nwm_operational_combined.csv") + crossed_combined = OUT_DIR / f"{CAT_ID}_nwm_operational_combined.csv" + + if da_on_combined.exists(): + TEST_FORCING_FILE = str(da_on_combined) + print(f" Reusing combined forcing: {TEST_FORCING_FILE}") + elif crossed_combined.exists(): + TEST_FORCING_FILE = str(crossed_combined) + print(f" Reusing combined forcing: {TEST_FORCING_FILE}") + elif args.test_forcing_dir1 and args.test_forcing_dir2: + f1 = os.path.join(args.test_forcing_dir1, f"{CAT_ID}.csv") + f2 = os.path.join(args.test_forcing_dir2, f"{CAT_ID}.csv") + df1, df2 = pd.read_csv(f1), pd.read_csv(f2) + combined = (pd.concat([df1, df2], ignore_index=True) + .drop_duplicates(subset="time").sort_values("time")) + TEST_FORCING_FILE = str(crossed_combined) + combined.to_csv(TEST_FORCING_FILE, index=False) + print(f" Wrote combined forcing: {TEST_FORCING_FILE}") + else: + print("Need --test-forcing-dir1 and --test-forcing-dir2"); return + + best_params_file = OUT_DIR / f"{CAT_ID}_best_params.json" + if not best_params_file.exists(): + # Also check the DA results dir from run_perturbation_da_on.py + alt = Path(args.out_dir).parent / "v2_perturbation_da_on" / CAT_ID / f"{CAT_ID}_best_params.json" + if alt.exists(): + best_params_file = alt + else: + print(f"Missing best params at {best_params_file}"); return + with open(best_params_file) as f: + best_params = json.load(f)["best_parameters"] + + obs_dict, var_dict = load_obs() + df_forcing = load_test_forcing() + t_mask = (df_forcing["date"] >= TEST_START) & (df_forcing["date"] <= TEST_END) + df_test = df_forcing[t_mask].reset_index(drop=True) + + # Write BMI config once — reused for all 90,600 model instantiations + cfg_path = write_model_config(best_params) + print(f" BMI config written once: {cfg_path}") + + rng = np.random.default_rng(hash(CAT_ID) & 0x7fffffff) + + # Phase 1: get DA analysis snapshots + # Check if a prior run already saved them (skip re-running DA if so) + snap_cache = OUT_DIR / f"{CAT_ID}_da_snapshots.parquet" + if snap_cache.exists(): + print(f" Loading cached DA snapshots from {snap_cache}") + df_snaps = pd.read_parquet(snap_cache) + snapshots = { + pd.Timestamp(row["issue_time"]): { + "soil_m": row["soil_m"], "gw_m": row["gw_m"], + "nash0": row["nash0"], "nash1": row["nash1"], + } + for _, row in df_snaps.iterrows() + } + else: + snapshots = run_phase1_da(cfg_path, df_test, obs_dict, var_dict, rng) + + issue_times = sorted(snapshots.keys()) + t0_to_idx = {t: i for i, t in enumerate(df_test["date"].apply(pd.Timestamp))} + + # ------------------------------------------------------------------ # + # Crossed ensemble: for each issue_time, build n_fa forcing draws and # + # n_ha state draws, then run all n_fa x n_ha combinations. # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Running crossed ensemble ({n_total} members x " + f"{len(issue_times)} issue times)...") + + all_records = [] + all_hydro_draw_records = [] # (issue_time, hydro_draw_j, states) + all_forcing_draw_records = [] # (issue_time, forcing_draw_i, lead, P_pert, E_pert, scales) + + for t_idx, t0 in enumerate(issue_times): + snap = snapshots[t0] + t0_str = t0.strftime("%Y-%m-%d %H:%M:%S") + if t0 not in t0_to_idx: + continue + start_idx = t0_to_idx[t0] + + # Pre-draw n_fa forcing perturbation sequences + forcing_seqs = [] + for i in range(n_fa): + seq = [] + for lead in range(1, FORECAST_HOURS + 1): + fi = start_idx + lead + if fi >= len(df_test): + seq.append((0.0, 0.0)) + all_forcing_draw_records.append({ + "issue_time": t0_str, "forcing_draw_i": i, + "lead_hour": lead, + "P_pert_mmh": 0.0, "E_pert_mmh": 0.0, + "precip_scale": np.nan, "pet_scale": np.nan, + }) + continue + row = df_test.iloc[fi] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + mu_p = -0.5 * PRECIP_SIGMA ** 2 + P_pert = float(P * rng.lognormal(mu_p, PRECIP_SIGMA)) if P > 0 else 0.0 + E_pert = max(float(E * (1.0 + PET_SIGMA * rng.standard_normal())), 0.0) + seq.append((P_pert, E_pert)) + all_forcing_draw_records.append({ + "issue_time": t0_str, "forcing_draw_i": i, + "lead_hour": lead, + "P_pert_mmh": P_pert, "E_pert_mmh": E_pert, + "precip_scale": P_pert / P if P > 0 else np.nan, + "pet_scale": E_pert / E if E > 0 else np.nan, + }) + forcing_seqs.append(seq) + + # Pre-draw n_ha hydro-state perturbations + state_draws = [snap if j == 0 else perturb_state(snap, rng) + for j in range(n_ha)] + + # Record each hydro draw's actual initial state + for j, st in enumerate(state_draws): + all_hydro_draw_records.append({ + "issue_time": t0_str, "hydro_draw_j": j, + "soil_m": st["soil_m"], "gw_m": st["gw_m"], + "nash0": st["nash0"], "nash1": st["nash1"], + }) + + # q_matrix shape: (FORECAST_HOURS, n_fa, n_ha) + q_matrix = np.full((FORECAST_HOURS, n_fa, n_ha), np.nan) + + for i, f_seq in enumerate(forcing_seqs): + for j, state_j in enumerate(state_draws): + m = make_model(cfg_path) + set_state(m, state_j) + for lead in range(1, FORECAST_HOURS + 1): + P_pert, E_pert = f_seq[lead - 1] + q = run_model_step(m, P_pert, E_pert) + q_matrix[lead - 1, i, j] = q + m.finalize() + + # Flatten 600-member matrix into rows + for lead in range(1, FORECAST_HOURS + 1): + row_data = {"issue_time": t0_str, "lead_hour": lead} + mem_idx = 0 + for i in range(n_fa): + for j in range(n_ha): + row_data[f"member_{mem_idx:04d}"] = q_matrix[lead - 1, i, j] + mem_idx += 1 + all_records.append(row_data) + + if (t_idx + 1) % 10 == 0 or (t_idx + 1) == len(issue_times): + print(f" {t_idx + 1}/{len(issue_times)} issue times complete") + + df_out = pd.DataFrame(all_records) + out_path = OUT_DIR / f"{CAT_ID}_crossed_ensemble.parquet" + df_out.to_parquet(out_path, index=False) + print(f"Saved: {out_path}") + print(f" Shape: {df_out.shape} " + f"({df_out['issue_time'].nunique()} issue times x " + f"{FORECAST_HOURS} leads x {n_total} members)") + + # ------------------------------------------------------------------ # + # Provenance files — enable full ensemble traceability and restart # + # ------------------------------------------------------------------ # + + # 1. DA analysis snapshots — unperturbed state at each issue_time. + # Load any row and set_state() to restart from that init time. + snap_records = [ + {"issue_time": t.strftime("%Y-%m-%d %H:%M:%S"), **s} + for t, s in snapshots.items() + ] + snap_path = OUT_DIR / f"{CAT_ID}_da_snapshots.parquet" + pd.DataFrame(snap_records).to_parquet(snap_path, index=False) + print(f"Saved: {snap_path} ({len(snap_records)} snapshots)") + + # 2. Member manifest — decoder ring: member_col -> (forcing_draw_i, hydro_draw_j). + # member_k uses forcing draw k//n_ha and hydro draw k%n_ha. + manifest = [ + {"member": f"member_{k:04d}", "member_idx": k, + "forcing_draw_i": k // n_ha, "hydro_draw_j": k % n_ha} + for k in range(n_total) + ] + manifest_path = OUT_DIR / f"{CAT_ID}_member_manifest.csv" + pd.DataFrame(manifest).to_csv(manifest_path, index=False) + print(f"Saved: {manifest_path} ({n_total} members)") + + # 3. Hydro draw states — actual perturbed initial state for each of + # n_ha draws at every issue_time. Links member -> initial conditions. + hydro_path = OUT_DIR / f"{CAT_ID}_hydro_draw_states.parquet" + pd.DataFrame(all_hydro_draw_records).to_parquet(hydro_path, index=False) + print(f"Saved: {hydro_path} ({len(all_hydro_draw_records)} rows)") + + # 4. Forcing draw sequences — actual P/PET perturbation values for each + # of n_fa draws at every (issue_time, lead_hour). Links member -> forcing. + forcing_path = OUT_DIR / f"{CAT_ID}_forcing_draw_sequences.parquet" + pd.DataFrame(all_forcing_draw_records).to_parquet(forcing_path, index=False) + print(f"Saved: {forcing_path} ({len(all_forcing_draw_records)} rows)") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/run_lead_time_forecast_sweep.py b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/run_lead_time_forecast_sweep.py new file mode 100644 index 00000000..19f3b4d5 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/run_lead_time_forecast_sweep.py @@ -0,0 +1,508 @@ +""" +Forecast lead-time evaluation. + +For each issue-time t0 in a sampling schedule across the test period, fork the +ensemble at t0 and free-run 18 hours of forecast. Two parallel trajectories are +maintained throughout the test period: + + da: Production setup. DA on, forcing perturbed, process noise on. + openloop: No DA from t=0. Forcing perturbed, process noise on. + +At each issue time t0, the state of each scenario's ensemble is copied into a +forecast ensemble that free-runs 18 hours with: + - DA off (no observations consumed during the forecast window) + - process noise off (no obs to collapse toward, so no anti-collapse needed) + - forcing perturbed (lognormal precip, Gaussian PET — proxy for forecast + precip uncertainty in lieu of actual HEFS forecasts) + +Issue-time schedule: + - Base cadence: every --base-step-h hours across the full test period + (default 6h → ~1600 issue times for the year) + - Densified to hourly across --dense-start..--dense-end if provided + (default: 2024-09-24 → 2024-09-28, the Helene window) + +Outputs (stacked across all issue times): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv + Columns: issue_time, lead_hour, valid_time, member_00 .. member_19 + (units: mm/h) + +Post-processing pipeline (separate scripts, run after this): + route_lead_time_forecasts.py — pushes each (t0, lead_hour) forecast through + T-route to the gauge + plot_lead_time_decay.py — error-vs-lead-time curve, DA vs open-loop + +Imports EnKFAssimilator from calibrate_catchment_cfe_da_v2.py so the DA math +matches production exactly. +""" + +import argparse +import os +import sys +import json +import importlib.util +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +SPINUP_START = "2023-02-01 00:00:00" +SPINUP_END = "2023-09-30 23:00:00" +TEST_START = "2023-10-01 00:00:00" +TEST_END = "2024-10-31 23:00:00" + +# Forecast lead time (hours after each issue time) +FORECAST_LEAD_HOURS = 18 + +# Default densification window — Hurricane Helene +DEFAULT_DENSE_START = "2024-09-24 00:00:00" +DEFAULT_DENSE_END = "2024-09-28 23:00:00" + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(test_forcing_file): + df = pd.read_csv(test_forcing_file) + df['date'] = pd.to_datetime(df['time']).dt.strftime('%Y-%m-%d %H:%M:%S') + df['total_precipitation'] = df['APCP_surface'] * 3600.0 + df['potential_evaporation'] = priestley_taylor_pet( + df['DSWRF_surface'].values, df['TMP_2maboveground'].values + ) + return df + + +def import_enkf_class(script_path): + """Load EnKFAssimilator from the production script without triggering its main().""" + spec = importlib.util.spec_from_file_location("prod_v2", script_path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod.EnKFAssimilator + + +# ---------------- State snapshot / restore ----------------------------------- +# +# The EnKF only ever touches these 4 states, so they're the only states that +# diverge between members. Copy these from src → dst to fork the ensemble. + +def snapshot_states(models): + return [ + { + "soil_m": float(m.soil_reservoir["storage_m"]), + "gw_m": float(m.gw_reservoir["storage_m"]), + "nash0_m": float(m.nash_storage[0]), + "nash1_m": float(m.nash_storage[1]), + } + for m in models + ] + + +def restore_states(models, snapshot): + for m, s in zip(models, snapshot): + m.soil_reservoir["storage_m"] = s["soil_m"] + m.gw_reservoir["storage_m"] = s["gw_m"] + m.nash_storage[0] = s["nash0_m"] + m.nash_storage[1] = s["nash1_m"] + + +# ---------------- Issue-time schedule ---------------------------------------- + +def build_issue_time_schedule(dates_list, base_step_h, dense_start, dense_end): + """Return a sorted list of issue-time strings sampled from dates_list. + + Every base_step_h hours across the full test period, plus every hour inside + [dense_start, dense_end] if those are provided. + """ + dates_dt = pd.to_datetime(dates_list) + selected = set() + + # Base cadence (every base_step_h hours from the first date) + base_mask = (np.arange(len(dates_dt)) % base_step_h == 0) + for d in dates_dt[base_mask]: + selected.add(d.strftime('%Y-%m-%d %H:%M:%S')) + + # Densified hourly window + if dense_start is not None and dense_end is not None: + ds = pd.Timestamp(dense_start) + de = pd.Timestamp(dense_end) + dense_mask = (dates_dt >= ds) & (dates_dt <= de) + for d in dates_dt[dense_mask]: + selected.add(d.strftime('%Y-%m-%d %H:%M:%S')) + + return sorted(selected) + + +# ---------------- CFE helpers ------------------------------------------------ + +def build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, enkf, + apply_init_perturbation=True, member_zero_clean=True): + """Build N CFE BMI instances with optional init-state perturbation.""" + def custom_load_forcing(self_cfe): + df = load_test_forcing(test_forcing_file) + self_cfe.forcing_data = df.rename(columns={"date": "time"}) + + models = [] + for i in range(N): + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.load_forcing_file = custom_load_forcing.__get__(m) + m.initialize() + if apply_init_perturbation and not (member_zero_clean and i == 0): + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + r = enkf.init_state_perturb_frac + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1 + r * enkf.rng.standard_normal()), 1e-6, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1 + r * enkf.rng.standard_normal()), 1e-6, gw_max)) + jitter = max(sm_max * 1e-6, 1e-9) + m.nash_storage[0] = max(n00 + jitter * enkf.rng.standard_normal(), 0.0) + m.nash_storage[1] = max(n10 + jitter * enkf.rng.standard_normal(), 0.0) + models.append(m) + return models + + +def step_ensemble(models, p_arr, e_arr): + """Advance each member one hour with member-specific (P, PET). Returns Q (mm/h).""" + q = np.empty(len(models), dtype=float) + for i, m in enumerate(models): + m.set_value('atmosphere_water__time_integral_of_precipitation_mass_flux', + float(p_arr[i]) / 1000.0) + m.set_value('water_potential_evaporation_flux', + float(e_arr[i]) / 1000.0 / 3600.0) + m.update() + q[i] = m.get_value('land_surface_water__runoff_depth') * 1000.0 # m/h → mm/h + return q + + +def do_sanity_check(prod_models, fcst_models, dates_list, forcing_by_date, h, N): + """One-shot verification that the 4-state snapshot is sufficient to reproduce + next-hour forecast Q. + + At main-loop hour h (after prod's hour-h step has completed including DA + + process noise), snapshot prod, restore into fcst, then step BOTH ensembles + one hour using identical unperturbed forcing. If snapshot/restore captures + everything that matters, prod's and fcst's hour-(h+1) Q should match member- + by-member to machine precision. + + Note: this mutates prod's state (h+1 step with non-perturbed forcing and no + DA), so the caller should exit the script right after. + """ + if h + 1 >= len(dates_list): + print("[sanity-check] not enough horizon for h+1 step — skipping.") + return False + snap = snapshot_states(prod_models) + restore_states(fcst_models, snap) + + next_date = dates_list[h + 1] + p_next, e_next = forcing_by_date[next_date] + p_uniform = np.full(N, p_next, dtype=float) + e_uniform = np.full(N, e_next, dtype=float) + + q_prod = step_ensemble(prod_models, p_uniform, e_uniform) + q_fcst = step_ensemble(fcst_models, p_uniform, e_uniform) + + diff = q_prod - q_fcst + max_abs = float(np.max(np.abs(diff))) + print("[sanity-check] forecast Q at h+1 with identical unperturbed forcing:") + print(f" prod first 5 members: {q_prod[:5]}") + print(f" fcst first 5 members: {q_fcst[:5]}") + print(f" max |prod - fcst|: {max_abs:.3e} mm/h") + if max_abs < 1e-9: + print(" OK: snapshot/restore preserves state to machine precision.") + return True + print(" WARN: nonzero divergence — likely a state not in the 4-state snapshot.") + return False + + +def run_forecast(fcst_models, enkf_fcst, dates_list, forcing_by_date, + t0_idx, n_lead): + """Run an n_lead-hour free-forecast starting from the current state of fcst_models. + + Forcing is perturbed per member (lognormal precip + Gaussian PET, same as + production). No DA. No process noise. + + Returns: (lead_hours, valid_times, q_matrix shape (n_lead, N)). + """ + N = len(fcst_models) + q_matrix = np.full((n_lead, N), np.nan, dtype=float) + lead_hours = [] + valid_times = [] + for lead in range(1, n_lead + 1): + idx = t0_idx + lead + if idx >= len(dates_list): + break + valid_date = dates_list[idx] + p, e = forcing_by_date[valid_date] + p_arr, e_arr = enkf_fcst.perturb_forcing(p, e) + q = step_ensemble(fcst_models, p_arr, e_arr) + q_matrix[lead - 1, :] = q + lead_hours.append(lead) + valid_times.append(valid_date) + return lead_hours, valid_times, q_matrix + + +# ---------------- Main run --------------------------------------------------- + +def run(args, bmi_cfe, EnKFAssimilator): + cat_id = args.cat_id + out_dir = Path(args.out_dir) / cat_id + out_dir.mkdir(parents=True, exist_ok=True) + + # Combined test forcing (same idiom as the per-member script) + f1 = os.path.join(args.test_forcing_dir1, f'{cat_id}.csv') + f2 = os.path.join(args.test_forcing_dir2, f'{cat_id}.csv') + df1 = pd.read_csv(f1) + df2 = pd.read_csv(f2) + combined = pd.concat([df1, df2], ignore_index=True) + combined = combined.drop_duplicates(subset='time').sort_values('time') + test_forcing_file = str(out_dir / f'{cat_id}_nwm_operational_combined.csv') + combined.to_csv(test_forcing_file, index=False) + + # Pre-staged calibrated params + best_params_file = out_dir / f'{cat_id}_best_params.json' + if not best_params_file.exists(): + print(f"Missing {best_params_file}. Pre-stage from Run 3.") + return + with open(best_params_file) as f: + best = json.load(f)['best_parameters'] + + # Build temp CFE config with calibrated params + with open(args.config_file) as f: + cfg = json.load(f) + cfg['forcing_file'] = test_forcing_file + cfg['soil_params']['bb'] = best['bb'] + cfg['soil_params']['smcmax'] = best['smcmax'] + cfg['soil_params']['satdk'] = best['satdk'] + cfg['slop'] = best['slop'] + cfg['max_gw_storage'] = best['max_gw_storage'] + cfg['expon'] = best['expon'] + cfg['Cgw'] = best['Cgw'] + cfg['K_lf'] = best['K_lf'] + cfg['K_nash'] = best['K_nash'] + cfg['partition_scheme'] = "Schaake" if best['scheme'] <= 0.5 else "Xinanjiang" + tmp_cfg = str(out_dir / f'{cat_id}_bmi_config_temp_leadtime.json') + with open(tmp_cfg, 'w') as f: + json.dump(cfg, f) + + # Three EnKF instances: + # enkf_da — runs the DA-on trajectory; consumes obs each hour + # enkf_ol — runs the open-loop trajectory; never calls update_states + # enkf_fcst — used only inside the 18-hour free-runs (perturb_forcing only) + # All three share defaults with production. Distinct RNG seeds keep them + # independent so the openloop and forecast spreads are not coupled to DA. + obs_file = os.path.join(args.obs_dir, f'{cat_id}.csv') + seed_base = args.rng_seed if args.rng_seed is not None else 0 + enkf_da = EnKFAssimilator( + n_members=args.enkf_members, obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, use_vrugt_r=(not args.no_vrugt_r), + vrugt_alpha=args.vrugt_alpha, vrugt_scale=args.vrugt_scale, + rng_seed=seed_base if args.rng_seed is not None else None, + ) + enkf_ol = EnKFAssimilator( + n_members=args.enkf_members, obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, use_vrugt_r=(not args.no_vrugt_r), + vrugt_alpha=args.vrugt_alpha, vrugt_scale=args.vrugt_scale, + rng_seed=(seed_base + 1) if args.rng_seed is not None else None, + ) + enkf_fcst = EnKFAssimilator( + n_members=args.enkf_members, obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, use_vrugt_r=(not args.no_vrugt_r), + vrugt_alpha=args.vrugt_alpha, vrugt_scale=args.vrugt_scale, + rng_seed=(seed_base + 2) if args.rng_seed is not None else None, + ) + # Optional R override — replace every per-hour obs variance with a constant. + # Matches DualEarth/new_EnKF.py (R=0.07). Set via --hardcoded-r on the CLI. + # Effect: at storm peak ≈ existing Vrugt R, at low flow ≈ 15-17× larger → + # near-zero gain at low flow where kriging obs is noisy. + if args.hardcoded_r is not None: + for enkf in (enkf_da, enkf_ol, enkf_fcst): + for date in enkf.obs_var_dict: + enkf.obs_var_dict[date] = args.hardcoded_r + print(f"[lead-time] R hardcoded to {args.hardcoded_r} mm^2/h^2 " + f"(overrides Vrugt + kriging variance formula)") + + N = enkf_da.n_members + print(f"[lead-time] {cat_id} | N={N} | DA on/off both run | " + f"forecast lead = {FORECAST_LEAD_HOURS}h") + + # Build three ensembles (init-perturbed except the forecast ensemble, whose + # state will be overwritten at each issue time before each free-run) + prod_models = build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, enkf_da) + openloop_models = build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, enkf_ol) + fcst_models = build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, + enkf_fcst, apply_init_perturbation=False) + + # ----- Spinup with perturbed forcing on both trajectories ----- + df = load_test_forcing(test_forcing_file) + sp_mask = (df['date'] >= SPINUP_START) & (df['date'] <= SPINUP_END) + df_sp = df[sp_mask] + print(f"[lead-time] spinup: {len(df_sp)} hours") + for p, e in zip(df_sp['total_precipitation'], df_sp['potential_evaporation']): + p_da, e_da = enkf_da.perturb_forcing(p, e) + step_ensemble(prod_models, p_da, e_da) + enkf_da.add_process_noise(prod_models) + p_ol, e_ol = enkf_ol.perturb_forcing(p, e) + step_ensemble(openloop_models, p_ol, e_ol) + enkf_ol.add_process_noise(openloop_models) + + # ----- Test period: step both, fork forecasts at issue times ----- + t_mask = (df['date'] >= TEST_START) & (df['date'] <= TEST_END) + df_test = df[t_mask].reset_index(drop=True) + dates_list = list(df_test['date'].values) + forcing_by_date = dict(zip( + df_test['date'].values, + zip(df_test['total_precipitation'].values, + df_test['potential_evaporation'].values))) + + issue_times = build_issue_time_schedule( + dates_list, args.base_step_h, + args.dense_start, args.dense_end) + issue_set = set(issue_times) + print(f"[lead-time] test period: {len(dates_list)} hours | " + f"{len(issue_times)} issue times " + f"(base={args.base_step_h}h, dense={args.dense_start or 'none'}" + f"..{args.dense_end or 'none'})") + + # Pre-build the date→index map so forecast windows are fast to look up + date_to_idx = {d: i for i, d in enumerate(dates_list)} + + da_rows = [] # (issue_time, lead_hour, valid_time, *member_values) + ol_rows = [] + sanity_done = False + + for h, current_date in enumerate(dates_list): + p, e = forcing_by_date[current_date] + + # --- DA trajectory step --- + p_da, e_da = enkf_da.perturb_forcing(p, e) + q_da = step_ensemble(prod_models, p_da, e_da) + enkf_da.update_states(prod_models, current_date, q_da) + enkf_da.add_process_noise(prod_models) + + # --- Open-loop trajectory step (no DA) --- + p_ol, e_ol = enkf_ol.perturb_forcing(p, e) + q_ol = step_ensemble(openloop_models, p_ol, e_ol) + enkf_ol.add_process_noise(openloop_models) + + # --- One-shot sanity check at the first scheduled issue time --- + if args.sanity_check and not sanity_done and current_date in issue_set: + ok = do_sanity_check( + prod_models, fcst_models, dates_list, forcing_by_date, h, N) + sanity_done = True + print(f"[sanity-check] exiting (re-run without --sanity-check for " + f"the full sweep). Result: {'OK' if ok else 'FAIL'}") + for m in prod_models + openloop_models + fcst_models: + m.finalize() + sys.exit(0 if ok else 1) + + # --- Fork forecasts at scheduled issue times --- + if current_date in issue_set: + for scenario_label, src_models, rows_acc in [ + ('da', prod_models, da_rows), + ('openloop', openloop_models, ol_rows), + ]: + snap = snapshot_states(src_models) + restore_states(fcst_models, snap) + leads, valids, qm = run_forecast( + fcst_models, enkf_fcst, + dates_list, forcing_by_date, + t0_idx=h, n_lead=FORECAST_LEAD_HOURS) + for k, (lead, vt) in enumerate(zip(leads, valids)): + rows_acc.append((current_date, lead, vt, *qm[k, :])) + + for m in prod_models + openloop_models + fcst_models: + m.finalize() + + # ----- Save forecast CSVs ----- + cols = ['issue_time', 'lead_hour', 'valid_time'] + [f'member_{i:02d}' for i in range(N)] + da_path = out_dir / f'{cat_id}_lead_time_forecasts_da.csv' + ol_path = out_dir / f'{cat_id}_lead_time_forecasts_openloop.csv' + pd.DataFrame(da_rows, columns=cols).to_csv(da_path, index=False) + pd.DataFrame(ol_rows, columns=cols).to_csv(ol_path, index=False) + print(f"[lead-time] saved {da_path}") + print(f"[lead-time] saved {ol_path}") + + # Issue-time schedule (small, useful for downstream scripts) + sched_path = out_dir / f'{cat_id}_lead_time_issue_times.csv' + pd.DataFrame({'issue_time': issue_times}).to_csv(sched_path, index=False) + print(f"[lead-time] saved {sched_path}") + + print(f"[lead-time] DA assimilator: updates={enkf_da.n_updates} | " + f"avg Pyy={enkf_da.avg_pyy:.6f} | " + f"mass lost={enkf_da.total_overflow_lost_mm:.3f} mm") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', required=True) + parser.add_argument('--forcing-dir', required=True) + parser.add_argument('--obs-dir', required=True) + parser.add_argument('--cfe-dir', required=True) + parser.add_argument('--config-file', required=True) + parser.add_argument('--param-bounds', required=True, + help='Kept for symmetry; not read in this script') + parser.add_argument('--out-dir', required=True) + parser.add_argument('--test-forcing-dir1', required=True) + parser.add_argument('--test-forcing-dir2', required=True) + parser.add_argument('--enkf-members', type=int, default=20) + parser.add_argument('--enkf-obs-error-std', type=float, default=0.05) + parser.add_argument('--no-vrugt-r', action='store_true') + parser.add_argument('--vrugt-alpha', type=float, default=0.10) + parser.add_argument('--vrugt-scale', type=float, default=0.001) + parser.add_argument('--rng-seed', type=int, default=None) + parser.add_argument('--base-step-h', type=int, default=6, + help='Issue-time cadence in hours across the full test ' + 'period (default 6h)') + parser.add_argument('--dense-start', type=str, default=DEFAULT_DENSE_START, + help='Start of hourly-densification window ' + '(default 2024-09-24 00:00:00 — Helene)') + parser.add_argument('--dense-end', type=str, default=DEFAULT_DENSE_END, + help='End of hourly-densification window ' + '(default 2024-09-28 23:00:00 — Helene)') + parser.add_argument('--sanity-check', action='store_true', + help='At the first issue time, fork fcst_models and step ' + 'both prod and fcst one hour with identical unperturbed ' + 'forcing. Confirms the 4-state snapshot is sufficient ' + 'to reproduce next-hour Q to machine precision, then ' + 'exits. Use before kicking off the full sweep.') + parser.add_argument('--hardcoded-r', type=float, default=None, + help='If set, override the Vrugt + kriging-variance R formula ' + 'and use this constant value (mm^2/h^2) at every hour. ' + 'Matches the DualEarth/new_EnKF.py reference (R=0.07). ' + 'Suppresses DA at low flow where kriging is noisy while ' + 'preserving DA strength at storm peaks. Applied to all ' + 'three EnKF instances (da, openloop, fcst) for consistency.') + parser.add_argument('--prod-script', default=None, + help='Path to calibrate_catchment_cfe_da_v2.py for ' + 'importing EnKFAssimilator. Defaults to next-to-this-file.') + args = parser.parse_args() + + here = os.path.dirname(os.path.abspath(__file__)) + prod_script = args.prod_script or os.path.join(here, 'calibrate_catchment_cfe_da_v2.py') + if not os.path.exists(prod_script): + raise FileNotFoundError( + f"Could not find production script at {prod_script}. " + f"Pass --prod-script to override.") + EnKFAssimilator = import_enkf_class(prod_script) + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + + run(args, _bmi_cfe, EnKFAssimilator) + + +if __name__ == '__main__': + main() diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/run_perturbation_da_on.py b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/run_perturbation_da_on.py new file mode 100644 index 00000000..1b9d9546 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/run_perturbation_da_on.py @@ -0,0 +1,392 @@ +""" +run_perturbation_da_on.py — 2a/2b: ensemble spread with DA ON. + +Runs two separated perturbation arms for the Helene window (Sep 24-30), +with DA actively running during the analysis period: + + ARM A — forcing only (30 members): + DA analyzes states using Qkrig at every hour (DA ON). + At each issue_time: fork 30 members with ONLY met forcing perturbed + (different precip/PET draw per member). Hydro states = DA analysis + mean (same for all 30). Free-run 18 hours, no further DA. + + ARM B — hydro states only (20 members): + Same DA analysis. At each issue_time: fork 20 members with ONLY + hydro states perturbed (multiplicative noise on soil/GW/Nash from + the DA analysis ensemble). Forcing = deterministic. Free-run 18 hours. + +These two arms isolate "where does forecast spread come from WITH DA?" +— figures for the assimilation script section. + +Output (per catchment, per arm): + //_da_forcing_arm.csv + //_da_hydro_arm.csv + Columns: issue_time, lead_hour, member_00..member_29 (forcing arm) + issue_time, lead_hour, member_00..member_19 (hydro arm) + +Usage (run for each catchment, both arms): + python3 run_perturbation_da_on.py \\ + --cat-id cat-1016300 \\ + --forcing-dir \\ + --obs-dir /mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance \\ + --cfe-dir /mnt/disk2/suma_helen_poster/cfe_py \\ + --config-file /mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json \\ + --param-bounds /mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json \\ + --out-dir /mnt/disk2/suma_helen_poster/da_results/v2_perturbation_da_on \\ + --test-forcing-dir1 \\ + --test-forcing-dir2 \\ + --hardcoded-r 0.07 +""" + +import argparse +import os +import sys +import json +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +N_FORCING = 30 # members in forcing arm +N_HYDRO = 20 # members in hydro-state arm + +# Analysis period — run DA through entire test period +TEST_START = "2024-08-24 00:00:00" # 1 month spinup before Helene window +TEST_END = "2024-10-31 23:00:00" + +# Issue times for the Helene window (hourly, Sep 24-30) +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-30 06:00:00") +FORECAST_HOURS = 18 + +# Perturbation magnitudes (same as production) +PRECIP_SIGMA = 0.15 +PET_SIGMA = 0.10 +STATE_FRAC = 0.05 # hydro-state perturbation fraction + +# Set at runtime +CAT_ID = None +OBS_FILE = None +CFE_CONFIG_FILE = None +TEST_FORCING_FILE = None +OUT_DIR = None +HARDCODED_R = None +bmi_cfe = None + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(): + df = pd.read_csv(TEST_FORCING_FILE) + df["date"] = pd.to_datetime(df["time"]).dt.strftime("%Y-%m-%d %H:%M:%S") + df["total_precipitation"] = df["APCP_surface"] * 3600.0 + df["potential_evaporation"] = priestley_taylor_pet( + df["DSWRF_surface"].values, df["TMP_2maboveground"].values) + return df + + +def load_obs(): + df = pd.read_csv(OBS_FILE) + date_col = next((c for c in df.columns + if c.lower() in ("date", "time", "datetime", "timestamp")), None) + if date_col is None: + raise ValueError(f"No date column in {OBS_FILE}. Columns: {df.columns.tolist()}") + q_col = next((c for c in df.columns if "qkrig" in c.lower() and "var" not in c.lower()), None) + if q_col is None: + raise ValueError(f"No qkrig column in {OBS_FILE}. Columns: {df.columns.tolist()}") + var_col = next((c for c in df.columns if "var" in c.lower()), None) + df[date_col] = pd.to_datetime(df[date_col]) + df = df.sort_values(date_col).reset_index(drop=True) + date_strs = df[date_col].dt.strftime("%Y-%m-%d %H:%M:%S") + obs_dict = dict(zip(date_strs, df[q_col])) + var_vals = df[var_col] if var_col else pd.Series([0.07] * len(df)) + var_dict = dict(zip(date_strs, var_vals)) + print(f" Obs: date='{date_col}', q='{q_col}', var='{var_col}'") + return obs_dict, var_dict + + +def make_model(best_params): + with open(CFE_CONFIG_FILE) as f: + cfg = json.load(f) + cfg["forcing_file"] = TEST_FORCING_FILE + cfg["soil_params"]["bb"] = best_params["bb"] + cfg["soil_params"]["smcmax"] = best_params["smcmax"] + cfg["soil_params"]["satdk"] = best_params["satdk"] + cfg["slop"] = best_params["slop"] + cfg["max_gw_storage"] = best_params["max_gw_storage"] + cfg["expon"] = best_params["expon"] + cfg["Cgw"] = best_params["Cgw"] + cfg["K_lf"] = best_params["K_lf"] + cfg["K_nash"] = best_params["K_nash"] + cfg["partition_scheme"] = ("Schaake" if best_params["scheme"] <= 0.5 + else "Xinanjiang") + tmp_cfg = str(OUT_DIR / f"{CAT_ID}_bmi_config_temp_da_on.json") + with open(tmp_cfg, "w") as f: + json.dump(cfg, f) + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.initialize() + return m + + +def get_state(m): + return { + "soil_m": m.soil_reservoir["storage_m"], + "gw_m": m.gw_reservoir["storage_m"], + "nash0": float(m.nash_storage[0]), + "nash1": float(m.nash_storage[1]), + } + + +def set_state(m, state): + m.soil_reservoir["storage_m"] = state["soil_m"] + m.gw_reservoir["storage_m"] = state["gw_m"] + m.nash_storage[0] = state["nash0"] + m.nash_storage[1] = state["nash1"] + + +def perturb_state(state, rng, frac=STATE_FRAC): + sm_max = 1.0 # placeholder — will be clipped in set_state; actual max enforced by model + out = { + "soil_m": max(state["soil_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "nash0": max(state["nash0"] * (1.0 + frac * rng.standard_normal()), 0.0), + "nash1": max(state["nash1"] * (1.0 + frac * rng.standard_normal()), 0.0), + } + return out + + +def run_model_step(m, precip_mmh, pet_mmh): + m.set_value("atmosphere_water__time_integral_of_precipitation_mass_flux", + precip_mmh / 1000.0) + m.set_value("water_potential_evaporation_flux", + pet_mmh / 1000.0 / 3600.0) + m.update() + return m.get_value("land_surface_water__runoff_depth") * 1000.0 # m/h -> mm/h + + +def enkf_update(state, q_sim, y_obs, krig_var, rng, n_members=1): + """Simple scalar EnKF update on the mean analysis state.""" + if HARDCODED_R is not None: + R = HARDCODED_R + else: + alpha = 0.10 + R = (alpha * max(y_obs, 0.0)) ** 2 + 0.001 * krig_var + R = max(R, 1e-6) + P_yy = max(q_sim * 0.01, 1e-6) # rough prior variance on Q + K = P_yy / (P_yy + R) + innovation = y_obs - q_sim + # Apply correction to each state proportionally (simple scalar gain) + scale = K * innovation / max(abs(q_sim), 1e-6) + updated = { + "soil_m": max(state["soil_m"] * (1.0 + scale), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + scale), 1e-6), + "nash0": max(state["nash0"] + state["nash0"] * scale, 0.0), + "nash1": max(state["nash1"] + state["nash1"] * scale, 0.0), + } + return updated + + +def main(): + global CAT_ID, OBS_FILE, CFE_CONFIG_FILE, TEST_FORCING_FILE + global OUT_DIR, HARDCODED_R, bmi_cfe + + parser = argparse.ArgumentParser() + parser.add_argument("--cat-id", required=True) + parser.add_argument("--forcing-dir", required=True) + parser.add_argument("--obs-dir", required=True) + parser.add_argument("--cfe-dir", required=True) + parser.add_argument("--config-file", required=True) + parser.add_argument("--param-bounds", required=True) + parser.add_argument("--out-dir", required=True) + parser.add_argument("--test-forcing-dir1", default=None) + parser.add_argument("--test-forcing-dir2", default=None) + parser.add_argument("--hardcoded-r", type=float, default=None) + args = parser.parse_args() + + CAT_ID = args.cat_id + OBS_FILE = os.path.join(args.obs_dir, f"{CAT_ID}.csv") + CFE_CONFIG_FILE = args.config_file + HARDCODED_R = args.hardcoded_r + OUT_DIR = Path(args.out_dir) / CAT_ID + OUT_DIR.mkdir(parents=True, exist_ok=True) + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + bmi_cfe = _bmi_cfe + + # Build combined test forcing + if args.test_forcing_dir1 and args.test_forcing_dir2: + f1 = os.path.join(args.test_forcing_dir1, f"{CAT_ID}.csv") + f2 = os.path.join(args.test_forcing_dir2, f"{CAT_ID}.csv") + df1, df2 = pd.read_csv(f1), pd.read_csv(f2) + combined = (pd.concat([df1, df2], ignore_index=True) + .drop_duplicates(subset="time").sort_values("time")) + TEST_FORCING_FILE = str(OUT_DIR / f"{CAT_ID}_nwm_operational_combined.csv") + combined.to_csv(TEST_FORCING_FILE, index=False) + else: + print("Need --test-forcing-dir1 and --test-forcing-dir2"); return + + best_params_file = OUT_DIR / f"{CAT_ID}_best_params.json" + if not best_params_file.exists(): + print(f"Missing best params at {best_params_file}"); return + with open(best_params_file) as f: + best_params = json.load(f)["best_parameters"] + + obs_dict, var_dict = load_obs() + df_forcing = load_test_forcing() + t_mask = (df_forcing["date"] >= TEST_START) & (df_forcing["date"] <= TEST_END) + df_test = df_forcing[t_mask].reset_index(drop=True) + + rng = np.random.default_rng(hash(CAT_ID) & 0x7fffffff) + + # ------------------------------------------------------------------ # + # Phase 1: Run DA through test period, snapshot analysis state at # + # each Helene issue time # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Phase 1: DA analysis through test period...") + model = make_model(best_params) + + helene_snapshots = {} # issue_time -> analysis state dict + + for h, row in df_test.iterrows(): + date_str = row["date"] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + q_sim = run_model_step(model, P, E) + + date_ts = pd.Timestamp(date_str) + + # DA update if obs available + y_obs = obs_dict.get(date_str, np.nan) + if not np.isnan(y_obs): + krig_var = var_dict.get(date_str, 0.07) + state = get_state(model) + updated = enkf_update(state, q_sim, y_obs, krig_var, rng) + set_state(model, updated) + + # Snapshot if this is a Helene issue time + if HELENE_START <= date_ts <= HELENE_END: + helene_snapshots[date_ts] = get_state(model) + + model.finalize() + print(f" Snapshots saved: {len(helene_snapshots)} issue times") + + # ------------------------------------------------------------------ # + # Phase 2A: Forcing arm — 30 members, fixed analyzed state, varying # + # met forcing # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Phase 2A: Forcing arm ({N_FORCING} members)...") + forcing_records = [] + + issue_times = sorted(helene_snapshots.keys()) + t0_to_idx = {t: i for i, t in enumerate(df_test["date"].apply(pd.Timestamp))} + + for t0 in issue_times: + snap = helene_snapshots[t0] + t0_str = t0.strftime("%Y-%m-%d %H:%M:%S") + + # Find index in df_test + if t0 not in t0_to_idx: + continue + start_idx = t0_to_idx[t0] + + q_matrix = np.full((FORECAST_HOURS, N_FORCING), np.nan) + + for mem in range(N_FORCING): + m = make_model(best_params) + set_state(m, snap) # all members start from same DA analysis state + + for lead in range(1, FORECAST_HOURS + 1): + fi = start_idx + lead + if fi >= len(df_test): + break + row = df_test.iloc[fi] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + # Perturb forcing only + mu_p = -0.5 * PRECIP_SIGMA ** 2 + P_pert = float(P * rng.lognormal(mu_p, PRECIP_SIGMA)) if P > 0 else 0.0 + E_pert = max(float(E * (1.0 + PET_SIGMA * rng.standard_normal())), 0.0) + q = run_model_step(m, P_pert, E_pert) + q_matrix[lead - 1, mem] = q + + m.finalize() + + row_data = {"issue_time": t0_str} + for lead in range(1, FORECAST_HOURS + 1): + row_data["lead_hour"] = lead + for mem in range(N_FORCING): + row_data[f"member_{mem:02d}"] = q_matrix[lead - 1, mem] + forcing_records.append(dict(row_data)) + + df_forcing_arm = pd.DataFrame(forcing_records) + out_a = OUT_DIR / f"{CAT_ID}_da_forcing_arm.csv" + df_forcing_arm.to_csv(out_a, index=False) + print(f" Saved: {out_a}") + + # ------------------------------------------------------------------ # + # Phase 2B: Hydro-state arm — 20 members, deterministic forcing, # + # perturbed analysis states # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Phase 2B: Hydro-state arm ({N_HYDRO} members)...") + hydro_records = [] + + for t0 in issue_times: + snap = helene_snapshots[t0] + t0_str = t0.strftime("%Y-%m-%d %H:%M:%S") + + if t0 not in t0_to_idx: + continue + start_idx = t0_to_idx[t0] + + q_matrix = np.full((FORECAST_HOURS, N_HYDRO), np.nan) + + for mem in range(N_HYDRO): + # Perturb the DA analysis state for this member + if mem == 0: + state_m = snap # member 0 = unperturbed analysis mean + else: + state_m = perturb_state(snap, rng) + + m = make_model(best_params) + set_state(m, state_m) + + for lead in range(1, FORECAST_HOURS + 1): + fi = start_idx + lead + if fi >= len(df_test): + break + row = df_test.iloc[fi] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + # Deterministic forcing + q = run_model_step(m, P, E) + q_matrix[lead - 1, mem] = q + + m.finalize() + + row_data = {"issue_time": t0_str} + for lead in range(1, FORECAST_HOURS + 1): + row_data["lead_hour"] = lead + for mem in range(N_HYDRO): + row_data[f"member_{mem:02d}"] = q_matrix[lead - 1, mem] + hydro_records.append(dict(row_data)) + + df_hydro_arm = pd.DataFrame(hydro_records) + out_b = OUT_DIR / f"{CAT_ID}_da_hydro_arm.csv" + df_hydro_arm.to_csv(out_b, index=False) + print(f" Saved: {out_b}") + print(f"[{CAT_ID}] Done.") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/run_perturbation_sensitivity.py b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/run_perturbation_sensitivity.py new file mode 100644 index 00000000..57b8d185 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/run_perturbation_sensitivity.py @@ -0,0 +1,293 @@ +""" +Perturbation-source sensitivity analysis for the multi-catchment CFE DA experiment. + +Ensemble-spread attribution: run 20 CFE members through the test +period with ONLY ONE perturbation source active at a time, no DA, no spinup. + +Three sub-experiments (selected via --source): + init : perturb initial states once at t=0 (multiplicative ~5%) + forcing : perturb hourly precip (lognormal sigma=0.15) and PET (Gaussian sigma=0.10) + process : apply per-hour process noise on states (soil 0.2%, GW 0.15%, Nash 0.5%) + +Each sub-experiment runs N=20 CFE members. DA is always OFF. Spinup is skipped — each +member starts from CFE's default initial state (modified by the perturbation if --source +init). This is the cleanest setup for attributing test-period spread to its source. + +Output: + //_sensitivity_.csv + Columns: date, member_00, member_01, ..., member_19 + +Each member column is that member's hourly Q_sim (mm/h) over the test period. + +Reads the same calibrated best_params.json as the production script. Does not need +the kriging obs file (no DA), but takes the same --obs-dir argument so it can also +log Q_obs alongside for plotting reference. + +Usage: + python3 run_perturbation_sensitivity.py \ + --cat-id cat-1016300 \ + --source forcing \ + --forcing-dir \ + --obs-dir \ + --cfe-dir \ + --config-file \ + --param-bounds \ + --out-dir \ + --test-forcing-dir1 \ + --test-forcing-dir2 +""" + +import argparse +import os +import sys +import json +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 +N_MEMBERS = 20 + +# We only run the test window for this sensitivity analysis. +TEST_START = "2023-10-01 00:00:00" +TEST_END = "2024-10-31 23:00:00" + +# Perturbation magnitudes — same as production defaults +INIT_STATE_FRAC = 0.05 +PRECIP_PERTURB = 0.15 # lognormal sigma +PET_PERTURB = 0.10 # Gaussian sigma +SOIL_PROC_NOISE = 0.002 +GW_PROC_NOISE = 0.0015 +NASH_PROC_NOISE = 0.005 + +# Set at runtime +CAT_ID = None +FORCING_FILE = None +TEST_FORCING_FILE = None +OBS_FILE = None +CFE_CONFIG_FILE = None +OUT_DIR = None +bmi_cfe = None + + +# ---------- Forcing helpers (same as production) ---------- +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(): + df = pd.read_csv(TEST_FORCING_FILE) + df['date'] = pd.to_datetime(df['time']).dt.strftime('%Y-%m-%d %H:%M:%S') + df['total_precipitation'] = df['APCP_surface'] * 3600.0 # kg/m^2/s -> mm/h + df['potential_evaporation'] = priestley_taylor_pet( + df['DSWRF_surface'].values, + df['TMP_2maboveground'].values, + ) + return df + + +# ---------- Perturbation helpers ---------- +def perturb_precip_lognormal(P, rng, sigma=PRECIP_PERTURB): + """Lognormal multiplier with mean 1; preserves zero precip.""" + mu = -0.5 * sigma ** 2 + return float(P * rng.lognormal(mu, sigma)) + + +def perturb_pet_gaussian(E, rng, sigma=PET_PERTURB): + """Gaussian multiplicative noise on PET, clipped at 0.""" + return max(float(E * (1.0 + sigma * rng.standard_normal())), 0.0) + + +def apply_init_state_perturb(m, rng, frac=INIT_STATE_FRAC): + """One-time multiplicative perturbation on the four states (called per member at t=0).""" + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1.0 + frac * rng.standard_normal()), 1e-6, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1.0 + frac * rng.standard_normal()), 1e-6, gw_max)) + # Nash buckets often start at 0; small additive jitter so they're not identical. + jitter = max(sm_max * 1e-6, 1e-9) + m.nash_storage[0] = max(n00 + jitter * rng.standard_normal(), 0.0) + m.nash_storage[1] = max(n10 + jitter * rng.standard_normal(), 0.0) + + +def apply_process_noise(m, rng): + """Per-timestep process noise on all 4 states for one member.""" + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + nash_floor = max(sm_max * 1e-4, 1e-7) + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1.0 + SOIL_PROC_NOISE * rng.standard_normal()), 0.0, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1.0 + GW_PROC_NOISE * rng.standard_normal()), 0.0, gw_max)) + m.nash_storage[0] = max( + n00 * (1.0 + NASH_PROC_NOISE * rng.standard_normal()) + + nash_floor * rng.standard_normal(), 0.0) + m.nash_storage[1] = max( + n10 * (1.0 + NASH_PROC_NOISE * rng.standard_normal()) + + nash_floor * rng.standard_normal(), 0.0) + + +# ---------- Main per-source run ---------- +def run_sensitivity(source, best_param_dict): + """Run N members through the test period with only `source` perturbation active.""" + + def custom_load_forcing(self_cfe): + df = load_test_forcing() + self_cfe.forcing_data = df.rename(columns={"date": "time"}) + + # Build a per-catchment temporary CFE config with the calibrated params + with open(CFE_CONFIG_FILE) as f: + cfg = json.load(f) + cfg['forcing_file'] = TEST_FORCING_FILE + cfg['soil_params']['bb'] = best_param_dict['bb'] + cfg['soil_params']['smcmax'] = best_param_dict['smcmax'] + cfg['soil_params']['satdk'] = best_param_dict['satdk'] + cfg['slop'] = best_param_dict['slop'] + cfg['max_gw_storage'] = best_param_dict['max_gw_storage'] + cfg['expon'] = best_param_dict['expon'] + cfg['Cgw'] = best_param_dict['Cgw'] + cfg['K_lf'] = best_param_dict['K_lf'] + cfg['K_nash'] = best_param_dict['K_nash'] + cfg['partition_scheme'] = "Schaake" if best_param_dict['scheme'] <= 0.5 else "Xinanjiang" + + tmp_cfg = str(OUT_DIR / f'{CAT_ID}_bmi_config_temp_sensitivity_{source}.json') + with open(tmp_cfg, 'w') as f: + json.dump(cfg, f) + + # Deterministic seed per (catchment, source) so each run is reproducible + seed = hash((CAT_ID, source)) & 0x7fffffff + rng = np.random.default_rng(seed) + + print(f"[sensitivity] {CAT_ID} | source={source} | N={N_MEMBERS} | seed={seed}") + + # Build N members. Initial-state perturbation only happens when source == 'init'. + models = [] + for i in range(N_MEMBERS): + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.load_forcing_file = custom_load_forcing.__get__(m) + m.initialize() + if source == 'init' and i > 0: + apply_init_state_perturb(m, rng) + models.append(m) + + # Skip spinup. Go directly to the test period. + df = load_test_forcing() + t_mask = (df['date'] >= TEST_START) & (df['date'] <= TEST_END) + df_test = df[t_mask] + + # Pre-allocate per-member Q arrays + n_hours = len(df_test) + q_matrix = np.full((n_hours, N_MEMBERS), np.nan, dtype=float) + dates_out = df_test['date'].values + + for h, (p, e) in enumerate(zip(df_test['total_precipitation'], + df_test['potential_evaporation'])): + for i, m in enumerate(models): + if source == 'forcing': + p_i = perturb_precip_lognormal(p, rng) + e_i = perturb_pet_gaussian(e, rng) + else: + p_i = float(p) + e_i = float(e) + m.set_value('atmosphere_water__time_integral_of_precipitation_mass_flux', p_i / 1000) + m.set_value('water_potential_evaporation_flux', e_i / 1000 / 3600) + m.update() + q_matrix[h, i] = m.get_value('land_surface_water__runoff_depth') * 1000 # m/h -> mm/h + + # Process noise applied after each hour, but only when source == 'process' + if source == 'process': + for m in models: + apply_process_noise(m, rng) + + for m in models: + m.finalize() + + # Write per-member CSV: date + 20 columns + out = {'date': dates_out} + for i in range(N_MEMBERS): + out[f'member_{i:02d}'] = q_matrix[:, i] + df_out = pd.DataFrame(out) + out_path = OUT_DIR / f'{CAT_ID}_sensitivity_{source}.csv' + df_out.to_csv(out_path, index=False) + print(f"[sensitivity] saved {out_path}") + print(f"[sensitivity] mean ensemble spread (std across members) over test period: " + f"{float(q_matrix.std(axis=1).mean()):.5f} mm/h") + + +def main(): + global CAT_ID, FORCING_FILE, TEST_FORCING_FILE, OBS_FILE, CFE_CONFIG_FILE + global OUT_DIR, bmi_cfe + + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', required=True) + parser.add_argument('--source', required=True, + choices=['init', 'forcing', 'process'], + help="Which perturbation source to isolate") + parser.add_argument('--forcing-dir', required=True) + parser.add_argument('--obs-dir', required=True, + help='Kept for symmetry with production script; not used since DA is off') + parser.add_argument('--cfe-dir', required=True) + parser.add_argument('--config-file', required=True) + parser.add_argument('--param-bounds', required=True, + help='Kept for symmetry; not read in this script') + parser.add_argument('--out-dir', required=True) + parser.add_argument('--test-forcing-dir1', default=None) + parser.add_argument('--test-forcing-dir2', default=None) + args = parser.parse_args() + + CAT_ID = args.cat_id + FORCING_FILE = os.path.join(args.forcing_dir, f'{CAT_ID}.csv') + OBS_FILE = os.path.join(args.obs_dir, f'{CAT_ID}.csv') + CFE_CONFIG_FILE = args.config_file + OUT_DIR = Path(args.out_dir) / CAT_ID + OUT_DIR.mkdir(parents=True, exist_ok=True) + + # Build combined test forcing CSV (same idiom as production) + if args.test_forcing_dir1 and args.test_forcing_dir2: + f1 = os.path.join(args.test_forcing_dir1, f'{CAT_ID}.csv') + f2 = os.path.join(args.test_forcing_dir2, f'{CAT_ID}.csv') + if os.path.exists(f1) and os.path.exists(f2): + df1 = pd.read_csv(f1) + df2 = pd.read_csv(f2) + combined = pd.concat([df1, df2], ignore_index=True) + combined = combined.drop_duplicates(subset='time').sort_values('time') + TEST_FORCING_FILE = str(OUT_DIR / f'{CAT_ID}_nwm_operational_combined.csv') + combined.to_csv(TEST_FORCING_FILE, index=False) + else: + print(f"Warning: test forcing files not found for {CAT_ID}, exiting") + return + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + bmi_cfe = _bmi_cfe + + best_params_file = OUT_DIR / f'{CAT_ID}_best_params.json' + if not best_params_file.exists(): + print(f"No best params file at {best_params_file}. Pre-stage from Run 3.") + return + with open(best_params_file) as f: + saved = json.load(f) + best_param_dict = saved['best_parameters'] + + run_sensitivity(args.source, best_param_dict) + + +if __name__ == '__main__': + main() diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/run_production_per_member.py b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/run_production_per_member.py new file mode 100644 index 00000000..d9c1c4ac --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/2_assimilation/run_production_per_member.py @@ -0,0 +1,314 @@ +""" +Run the production DA pipeline for ONE catchment and save per-member streamflow. + +Identical to calibrate_catchment_cfe_da_v2.py in algorithm (init perturbation + +forcing perturbation + process noise + true EnKF with Vrugt R), but writes one +column per ensemble member instead of only the ensemble mean. This is for +the per-member factor-decomposition plot. + +To keep the comparison apples-to-apples with the sensitivity runs (which skip +spinup), this script also SKIPS the spinup loop. Each member starts directly +from CFE's default initial state, plus the t=0 init perturbation. + +Output: + //_production_per_member.csv + Columns: date, member_00, member_01, ..., member_19 + +Imports EnKFAssimilator from calibrate_catchment_cfe_da_v2 to guarantee the +DA math matches production exactly. + +Usage: + python3 run_production_per_member.py \ + --cat-id cat-1016300 \ + --forcing-dir /mnt/.../nwm_retro_catchment_forcings \ + --obs-dir /mnt/.../catchment_ts_03463300_with_variance \ + --cfe-dir /mnt/.../cfe_py \ + --config-file /mnt/.../cat_03463300_bmi_config_cfe.json \ + --param-bounds /mnt/.../CFE_parameter_bounds.json \ + --out-dir /mnt/.../v2_production_per_member \ + --test-forcing-dir1 ... \ + --test-forcing-dir2 ... \ + --enkf-members 20 +""" + +import argparse +import os +import sys +import json +import importlib.util +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +TEST_START = "2023-10-01 00:00:00" +TEST_END = "2024-10-31 23:00:00" + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(test_forcing_file): + df = pd.read_csv(test_forcing_file) + df['date'] = pd.to_datetime(df['time']).dt.strftime('%Y-%m-%d %H:%M:%S') + df['total_precipitation'] = df['APCP_surface'] * 3600.0 + df['potential_evaporation'] = priestley_taylor_pet( + df['DSWRF_surface'].values, df['TMP_2maboveground'].values + ) + return df + + +def import_enkf_class(script_path): + """Load EnKFAssimilator from the production script without triggering its main().""" + spec = importlib.util.spec_from_file_location("prod_v2", script_path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod.EnKFAssimilator + + +def run(args, bmi_cfe, EnKFAssimilator): + cat_id = args.cat_id + out_dir = Path(args.out_dir) / cat_id + out_dir.mkdir(parents=True, exist_ok=True) + + # Combined test forcing + f1 = os.path.join(args.test_forcing_dir1, f'{cat_id}.csv') + f2 = os.path.join(args.test_forcing_dir2, f'{cat_id}.csv') + df1 = pd.read_csv(f1) + df2 = pd.read_csv(f2) + combined = pd.concat([df1, df2], ignore_index=True) + combined = combined.drop_duplicates(subset='time').sort_values('time') + test_forcing_file = str(out_dir / f'{cat_id}_nwm_operational_combined.csv') + combined.to_csv(test_forcing_file, index=False) + + # Load best_params + best_params_file = out_dir / f'{cat_id}_best_params.json' + if not best_params_file.exists(): + print(f"Missing {best_params_file}. Pre-stage from Run 3.") + return + with open(best_params_file) as f: + saved = json.load(f) + best = saved['best_parameters'] + + # Build temp CFE config with calibrated params + with open(args.config_file) as f: + cfg = json.load(f) + cfg['forcing_file'] = test_forcing_file + cfg['soil_params']['bb'] = best['bb'] + cfg['soil_params']['smcmax'] = best['smcmax'] + cfg['soil_params']['satdk'] = best['satdk'] + cfg['slop'] = best['slop'] + cfg['max_gw_storage'] = best['max_gw_storage'] + cfg['expon'] = best['expon'] + cfg['Cgw'] = best['Cgw'] + cfg['K_lf'] = best['K_lf'] + cfg['K_nash'] = best['K_nash'] + cfg['partition_scheme'] = "Schaake" if best['scheme'] <= 0.5 else "Xinanjiang" + tmp_cfg = str(out_dir / f'{cat_id}_bmi_config_temp_per_member.json') + with open(tmp_cfg, 'w') as f: + json.dump(cfg, f) + + # Build the EnKF assimilator (same defaults as production) + obs_file = os.path.join(args.obs_dir, f'{cat_id}.csv') + enkf = EnKFAssimilator( + n_members=args.enkf_members, + obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, + use_vrugt_r=(not args.no_vrugt_r) and (not getattr(args, 'direct_variance', False)), + vrugt_alpha=args.vrugt_alpha, + vrugt_scale=args.vrugt_scale, + rng_seed=args.rng_seed, + ) + N = enkf.n_members + print(f"[per-member] {cat_id} | N={N} | use_vrugt_r={enkf.use_vrugt_r}") + if getattr(args, 'hardcoded_r', None) is not None: + enkf.obs_var_dict = {k: args.hardcoded_r for k in enkf.obs_var_dict} + print(f"[per-member] obs_var_dict overridden: R = {args.hardcoded_r} (fixed)") + if getattr(args, 'direct_variance', False): + obs_df = pd.read_csv(obs_file) + t_col = next(c for c in obs_df.columns if c.lower() in ('datetime', 'date', 'time', 'timestamp')) + v_col = next(c for c in obs_df.columns if 'var' in c.lower()) + obs_df[t_col] = pd.to_datetime(obs_df[t_col]) + obs_df = obs_df.set_index(t_col).sort_index() + enkf.obs_var_dict = { + str(ts): float(v) + for ts, v in obs_df[v_col].items() + if pd.notna(v) and float(v) > 0 + } + print(f"[per-member] obs_var_dict overridden: R = σ² direct ({len(enkf.obs_var_dict)} timesteps)") + + # Custom forcing loader for CFE BMI + def custom_load_forcing(self_cfe): + df = load_test_forcing(test_forcing_file) + self_cfe.forcing_data = df.rename(columns={"date": "time"}) + + # Build N CFE instances; perturb initial states (member 0 left clean) + models = [] + init_states_per_member = [] # capture initial-state snapshot after perturbation + for i in range(N): + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.load_forcing_file = custom_load_forcing.__get__(m) + m.initialize() + if i > 0: + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + r = enkf.init_state_perturb_frac + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1 + r * enkf.rng.standard_normal()), 1e-6, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1 + r * enkf.rng.standard_normal()), 1e-6, gw_max)) + jitter = max(sm_max * 1e-6, 1e-9) + m.nash_storage[0] = max(n00 + jitter * enkf.rng.standard_normal(), 0.0) + m.nash_storage[1] = max(n10 + jitter * enkf.rng.standard_normal(), 0.0) + # Snapshot the (possibly-perturbed) initial state for this member + init_states_per_member.append({ + "soil_m": float(m.soil_reservoir["storage_m"]), + "gw_m": float(m.gw_reservoir["storage_m"]), + "nash0_m": float(m.nash_storage[0]), + "nash1_m": float(m.nash_storage[1]), + "soil_max_m": float(m.soil_reservoir["storage_max_m"]), + "gw_max_m": float(m.gw_reservoir["storage_max_m"]), + }) + models.append(m) + + # SKIP SPINUP — go directly to the test period (matches sensitivity-script setup) + df = load_test_forcing(test_forcing_file) + t_mask = (df['date'] >= TEST_START) & (df['date'] <= TEST_END) + df_test = df[t_mask] + n_hours = len(df_test) + q_matrix = np.full((n_hours, N), np.nan, dtype=float) + precip_matrix = np.full((n_hours, N), np.nan, dtype=float) # mm/h per member + pet_matrix = np.full((n_hours, N), np.nan, dtype=float) # mm/h per member + dates_out = df_test['date'].values + + for h, (p, e, current_date) in enumerate(zip( + df_test['total_precipitation'], + df_test['potential_evaporation'], + df_test['date'])): + + # Perturb forcing per member + p_arr, e_arr = enkf.perturb_forcing(p, e) + # Capture the per-member perturbed forcing for later plotting/diagnostics + precip_matrix[h, :] = p_arr + pet_matrix[h, :] = e_arr + + # Advance each member one hour + ensemble_q = np.empty(N, dtype=float) + for i, m in enumerate(models): + m.set_value('atmosphere_water__time_integral_of_precipitation_mass_flux', + float(p_arr[i]) / 1000) + m.set_value('water_potential_evaporation_flux', + float(e_arr[i]) / 1000 / 3600) + m.update() + ensemble_q[i] = m.get_value('land_surface_water__runoff_depth') * 1000 # mm/h + + # Record pre-analysis forecast per member (this is the "actual production" value) + q_matrix[h, :] = ensemble_q + + # DA step (uses production EnKFAssimilator) + enkf.update_states(models, current_date, ensemble_q) + + # Process noise after DA + enkf.add_process_noise(models) + + for m in models: + m.finalize() + + # ----- Save per-member CSVs ----- + def _save_matrix(matrix, suffix): + out = {'date': dates_out} + for i in range(N): + out[f'member_{i:02d}'] = matrix[:, i] + df_out = pd.DataFrame(out) + out_path = out_dir / f'{cat_id}_production_per_member_{suffix}.csv' + df_out.to_csv(out_path, index=False) + return out_path + + # Streamflow outputs + q_path = out_dir / f'{cat_id}_production_per_member.csv' + out = {'date': dates_out} + for i in range(N): + out[f'member_{i:02d}'] = q_matrix[:, i] + pd.DataFrame(out).to_csv(q_path, index=False) + print(f"[per-member] saved {q_path}") + + # Perturbed forcings (so each member's actual inputs are recoverable) + precip_path = _save_matrix(precip_matrix, "precip") + pet_path = _save_matrix(pet_matrix, "pet") + print(f"[per-member] saved {precip_path}") + print(f"[per-member] saved {pet_path}") + + # Initial states (small, JSON is fine) + init_states_path = out_dir / f'{cat_id}_production_per_member_initial_states.json' + with open(init_states_path, 'w') as f: + json.dump({ + "catchment_id": cat_id, + "n_members": N, + "members": {f"member_{i:02d}": init_states_per_member[i] for i in range(N)}, + }, f, indent=2) + print(f"[per-member] saved {init_states_path}") + + print(f"[per-member] ensemble mean spread over test period: " + f"{float(q_matrix.std(axis=1).mean()):.5f} mm/h") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', required=True) + parser.add_argument('--forcing-dir', required=True) + parser.add_argument('--obs-dir', required=True) + parser.add_argument('--cfe-dir', required=True) + parser.add_argument('--config-file', required=True) + parser.add_argument('--param-bounds', required=True, + help='Kept for symmetry; not read in this script') + parser.add_argument('--out-dir', required=True) + parser.add_argument('--test-forcing-dir1', required=True) + parser.add_argument('--test-forcing-dir2', required=True) + parser.add_argument('--enkf-members', type=int, default=20) + parser.add_argument('--enkf-obs-error-std', type=float, default=0.05) + parser.add_argument('--no-vrugt-r', action='store_true') + parser.add_argument('--hardcoded-r', type=float, default=None, + help='Fix R to this constant value for all timesteps (overrides Vrugt formula)') + parser.add_argument('--direct-variance', action='store_true', default=False, + help='Use kriging variance column as R per timestep (R=σ²; overrides Vrugt formula)') + parser.add_argument('--vrugt-alpha', type=float, default=0.10) + parser.add_argument('--vrugt-scale', type=float, default=0.001) + parser.add_argument('--rng-seed', type=int, default=None) + parser.add_argument('--prod-script', default=None, + help='Path to calibrate_catchment_cfe_da_v2.py for importing ' + 'EnKFAssimilator. Defaults to the script located next to this file.') + args = parser.parse_args() + + # Locate the production script (default: same dir as this script) + if args.prod_script is None: + here = os.path.dirname(os.path.abspath(__file__)) + prod_script = os.path.join(here, 'calibrate_catchment_cfe_da_v2.py') + else: + prod_script = args.prod_script + if not os.path.exists(prod_script): + raise FileNotFoundError(f"Could not find production script at {prod_script}. " + f"Pass --prod-script to override.") + EnKFAssimilator = import_enkf_class(prod_script) + + # Dynamic CFE import (same idiom as production script) + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + + run(args, _bmi_cfe, EnKFAssimilator) + + +if __name__ == '__main__': + main() diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/3_routing/route_det_f1.sh b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/3_routing/route_det_f1.sh new file mode 100644 index 00000000..374743b6 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/3_routing/route_det_f1.sh @@ -0,0 +1,28 @@ +#!/usr/bin/env bash +# F1 variance-scaled Vrugt (1 gauge holdout) — deterministic T-route routing. +# Routes per-catchment _test_results.csv (sim_mm_h column) through Muskingum-Cunge +# to gauge 03463300 and writes routed_Q_test.csv + KGE/NSE summary. +# +# Usage: +# bash route_det_f1.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +ROUTE_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/2_troute_routing/run_route.py + +F1_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt +GPKG=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +echo "[F1] Deterministic T-route routing..." +echo " da-dir : $F1_DIR" +echo " out-dir: $F1_DIR" + +$TROUTE "$ROUTE_SCRIPT" \ + --gpkg "$GPKG" \ + --da-dir "$F1_DIR" \ + --out-dir "$F1_DIR" \ + --usgs-csv "$USGS_CSV" + +echo "[F1] Deterministic routing done. Output: $F1_DIR/routed_Q_test.csv" diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/3_routing/route_ensemble_f1.sh b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/3_routing/route_ensemble_f1.sh new file mode 100644 index 00000000..539b9d07 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/3_routing/route_ensemble_f1.sh @@ -0,0 +1,30 @@ +#!/usr/bin/env bash +# F1 Vrugt (1 gauge holdout) — crossed-ensemble T-route routing. +# Routes all 600 crossed-ensemble members through Muskingum-Cunge +# to gauge 03463300 and writes routed_crossed_ensemble.parquet. +# +# Must be run AFTER crossed ensemble parquets exist in F1_DIR. +# +# Usage: +# bash route_ensemble_f1.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +ROUTE_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/2_troute_routing/run_route_crossed_ensemble.py + +F1_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt +GPKG=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +echo "[F1] Crossed-ensemble T-route routing..." +echo " ensemble-dir: $F1_DIR" +echo " out-dir : $F1_DIR" + +$TROUTE "$ROUTE_SCRIPT" \ + --gpkg "$GPKG" \ + --ensemble-dir "$F1_DIR" \ + --out-dir "$F1_DIR" \ + --usgs-csv "$USGS_CSV" + +echo "[F1] Ensemble routing done. Output: $F1_DIR/routed_crossed_ensemble.parquet" diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/3_routing/route_leadtime_f1.sh b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/3_routing/route_leadtime_f1.sh new file mode 100644 index 00000000..f65c2f96 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/3_routing/route_leadtime_f1.sh @@ -0,0 +1,31 @@ +#!/usr/bin/env bash +# F1 variance-scaled Vrugt (1 gauge holdout) — route lead-time forecast CSVs. +# Routes all 21-catchment lead-time forecast CSVs through Muskingum-Cunge +# to gauge 03463300. Writes routed_leadtime_da_full.parquet and +# routed_leadtime_openloop_full.parquet for gauge-level error decay analysis. +# +# Run AFTER batch_run_leadtime_f1.sh finishes. +# +# Usage: +# bash route_leadtime_f1.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +ROUTE_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/2_troute_routing/run_route_leadtime_forecasts.py + +F1_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt +GPKG=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt_leadtime_routed + +echo "[F1] Routing lead-time forecasts through T-route..." +echo " forecast-dir: $F1_DIR" +echo " out-dir : $OUT_DIR" + +$TROUTE "$ROUTE_SCRIPT" \ + --gpkg "$GPKG" \ + --forecast-dir "$F1_DIR" \ + --out-dir "$OUT_DIR" + +echo "[F1] Lead-time routing done. Outputs in: $OUT_DIR" +ls "$OUT_DIR"/*.parquet 2>/dev/null || echo " (no parquets found)" diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4a_error_decay/plot_forecast_error_fixed_target.py b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4a_error_decay/plot_forecast_error_fixed_target.py new file mode 100644 index 00000000..fc4d155a --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4a_error_decay/plot_forecast_error_fixed_target.py @@ -0,0 +1,205 @@ +""" +plot_forecast_error_fixed_target.py — 4a: error decay, fixed-target-time view. +F1 variance-scaled Vrugt (1 gauge holdout). + +For each target verification time T in the Helene peak window: + Collect all forecasts that verify AT T: + issue_time = T - lead_hour*1h, for lead in 1..18 + error[lead] = ensemble_mean(q at T, initialized T-lead) - USGS_obs(T) + +This gives the correct operational picture: + - lead 1 = initialized 1 hr before T (DA just ran -> small error) + - lead 18 = initialized 18 hr before T (DA long ago -> error ~ open loop) + +Two panels: + Top : signed error (m³/s) vs lead hour, one curve per target time + Bot : same for open-loop +Plus a summary panel: mean across all target times, DA vs OL. + +Outputs: + /error_fixed_target_helene.png (per-target spaghetti, DA vs OL) + /error_fixed_target_mean.png (mean across targets, DA vs OL) +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.cm as cm + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt_leadtime_routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +# Target verification times: hourly through the Helene peak window +TARGET_START = pd.Timestamp("2024-09-26 18:00:00") +TARGET_END = pd.Timestamp("2024-09-28 06:00:00") + + +def load_parquet(path): + df = pd.read_parquet(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + df["ens_mean"] = df[member_cols].mean(axis=1) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + return df[["issue_time", "lead_hour", "valid_time", "ens_mean"]] + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() + or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def build_fixed_target_errors(df, obs_series, target_times): + df_idx = df.set_index(["issue_time", "lead_hour"])["ens_mean"] + results = {} + for T in target_times: + obs_val = obs_series.get(T, np.nan) + if np.isnan(obs_val): + continue + curve = {} + for lead in range(1, 19): + t0 = T - pd.Timedelta(hours=lead) + try: + q_fc = df_idx.loc[(t0, lead)] + curve[lead] = float(q_fc) - obs_val + except KeyError: + curve[lead] = np.nan + results[T] = curve + return results + + +def plot_spaghetti(ax, error_dict, color_da, label_prefix, linestyle="-", lw=0.9, alpha=0.35): + leads = list(range(1, 19)) + all_curves = [] + target_times = sorted(error_dict.keys()) + cmap = cm.get_cmap("plasma", len(target_times)) + + for i, T in enumerate(target_times): + curve = [error_dict[T].get(l, np.nan) for l in leads] + ax.plot(leads, curve, + color=cmap(i), lw=lw, alpha=alpha, + linestyle=linestyle, zorder=2) + all_curves.append(curve) + + if all_curves: + mean_curve = np.nanmean(np.array(all_curves), axis=0) + ax.plot(leads, mean_curve, + color=color_da, lw=2.8, alpha=0.95, + linestyle=linestyle, zorder=5, + label=f"{label_prefix} — mean across targets") + return all_curves + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--route-dir", default=DEFAULT_ROUTE_DIR) + parser.add_argument("--out-dir", default=None) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--da-name", default="routed_leadtime_da_full.parquet") + parser.add_argument("--ol-name", default="routed_leadtime_openloop_full.parquet") + parser.add_argument("--target-start", default=str(TARGET_START)) + parser.add_argument("--target-end", default=str(TARGET_END)) + args = parser.parse_args() + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + print("Loading parquets...") + da = load_parquet(os.path.join(args.route_dir, args.da_name)) + ol = load_parquet(os.path.join(args.route_dir, args.ol_name)) + obs = load_usgs(args.usgs_csv) + + target_times = pd.date_range(args.target_start, args.target_end, freq="1h") + print(f" Target verification times: {len(target_times)} " + f"({target_times[0]} → {target_times[-1]})") + + da_errors = build_fixed_target_errors(da, obs, target_times) + ol_errors = build_fixed_target_errors(ol, obs, target_times) + print(f" Targets with obs: DA={len(da_errors)} OL={len(ol_errors)}") + + leads = list(range(1, 19)) + + fig, (ax_da, ax_ol) = plt.subplots(2, 1, figsize=(13, 10), sharex=True, sharey=True) + + ax_da.axhline(0, color="black", lw=0.8, linestyle="--", alpha=0.4) + plot_spaghetti(ax_da, da_errors, "tab:blue", "DA", linestyle="-") + ax_da.set_ylabel("Error: forecast − USGS obs (m³/s)", fontsize=10) + ax_da.set_title("DA — error at each lead for fixed target times (Helene peak window)", fontsize=11) + ax_da.grid(True, alpha=0.2) + ax_da.legend(fontsize=9) + + ax_ol.axhline(0, color="black", lw=0.8, linestyle="--", alpha=0.4) + plot_spaghetti(ax_ol, ol_errors, "tab:gray", "Open-loop", linestyle="--") + ax_ol.set_ylabel("Error: forecast − USGS obs (m³/s)", fontsize=10) + ax_ol.set_xlabel("Forecast lead hour (hours before target)", fontsize=11) + ax_ol.set_title("Open-loop — error at each lead for fixed target times", fontsize=11) + ax_ol.set_xticks(leads) + ax_ol.grid(True, alpha=0.2) + ax_ol.legend(fontsize=9) + + fig.suptitle( + "Forecast error vs lead time — fixed verification time, Helene peak window\n" + f"F1 (variance-scaled Vrugt) | Target: {TARGET_START.strftime('%b %d %H UTC')} → " + f"{TARGET_END.strftime('%b %d %H UTC')} | USGS 03463300", + fontsize=12, y=1.01, + ) + plt.tight_layout() + out1 = os.path.join(out_dir, "error_fixed_target_helene.png") + plt.savefig(out1, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + fig, ax = plt.subplots(figsize=(13, 6)) + ax.axhline(0, color="black", lw=0.8, linestyle="--", alpha=0.4) + + def mean_curve(error_dict): + arr = np.array([ + [error_dict[T].get(l, np.nan) for l in leads] + for T in sorted(error_dict.keys()) + ]) + return np.nanmean(arr, axis=0), np.nanstd(arr, axis=0) + + da_mean, da_std = mean_curve(da_errors) + ol_mean, ol_std = mean_curve(ol_errors) + + ax.fill_between(leads, da_mean - da_std, da_mean + da_std, + color="tab:blue", alpha=0.15, zorder=2) + ax.fill_between(leads, ol_mean - ol_std, ol_mean + ol_std, + color="tab:gray", alpha=0.15, zorder=2) + ax.plot(leads, da_mean, color="tab:blue", lw=2.6, marker="o", + zorder=5, label="DA — mean error (±1 std shaded)") + ax.plot(leads, ol_mean, color="tab:gray", lw=2.6, marker="s", + linestyle="--", zorder=5, label="Open-loop — mean error (±1 std shaded)") + + ax.set_xlabel("Forecast lead hour (hours before target verification time)", fontsize=11) + ax.set_ylabel("Mean error: forecast − USGS obs (m³/s)", fontsize=11) + ax.set_xticks(leads) + ax.set_title( + "Mean forecast error vs lead time — fixed verification time, Helene peak window\n" + f"F1 (variance-scaled Vrugt) | Target: {TARGET_START.strftime('%b %d %H UTC')} → " + f"{TARGET_END.strftime('%b %d %H UTC')} | Lead 1 = init 1 hr before target", + fontsize=11, + ) + ax.legend(fontsize=10, loc="lower right", frameon=True, framealpha=0.92) + ax.grid(True, alpha=0.25) + plt.tight_layout() + out2 = os.path.join(out_dir, "error_fixed_target_mean.png") + plt.savefig(out2, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4a_error_decay/plot_forecast_error_per_init.py b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4a_error_decay/plot_forecast_error_per_init.py new file mode 100644 index 00000000..4377c866 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4a_error_decay/plot_forecast_error_per_init.py @@ -0,0 +1,197 @@ +""" +plot_forecast_error_per_init.py — F1 variance-scaled Vrugt (1 gauge holdout). + +Error decay by initialization time for the Helene window (Sep 24-28 2024). + +For each initialization time t0 in the Helene window: + error[lead] = ensemble_mean(q_gauge_m3s at t0+lead) - USGS_obs(t0+lead) + +Plotted as: + DA : thin colored lines (one per init time, colored by date) + thick mean + OL : thin gray dashed lines + thick gray dashed mean + +x-axis: forecast lead hour (1 -> 18) +y-axis: signed error (m³/s), positive = forecast too high + +The expected signal: DA error is small at lead 1 (just assimilated), grows +and converges toward the OL error curve by lead 18. + +Inputs: + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + +Obs: + /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +Outputs: + /forecast_error_per_init_helene.png (signed error) + /forecast_mae_per_lead_helene.png (mean absolute error per lead) +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.patches as mpatches + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt_leadtime_routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +HELENE_INIT_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_INIT_END = pd.Timestamp("2024-09-28 23:00:00") + +DATE_COLORS = { + "2024-09-24": "#1f77b4", + "2024-09-25": "#ff7f0e", + "2024-09-26": "#2ca02c", + "2024-09-27": "#d62728", + "2024-09-28": "#9467bd", +} + + +def load_parquet(path): + df = pd.read_parquet(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + df["ens_mean"] = df[member_cols].mean(axis=1) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + return df[["issue_time", "lead_hour", "valid_time", "ens_mean"]] + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() + or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def build_error_table(df, obs_series, init_start, init_end): + df = df[(df["issue_time"] >= init_start) & (df["issue_time"] <= init_end)].copy() + df["obs"] = df["valid_time"].map(obs_series) + df["error"] = df["ens_mean"] - df["obs"] + return df.dropna(subset=["obs", "error"]) + + +def plot_error(ax, err_df, color, alpha_thin, lw_thin, linestyle, label_prefix): + leads = sorted(err_df["lead_hour"].unique()) + all_curves = [] + + for t0, grp in err_df.groupby("issue_time"): + date_str = str(pd.Timestamp(t0).date()) + c = DATE_COLORS.get(date_str, color) + grp_sorted = grp.sort_values("lead_hour") + curve = grp_sorted.set_index("lead_hour")["error"].reindex(leads).values + ax.plot(leads, curve, + color=c, lw=lw_thin, alpha=alpha_thin, + linestyle=linestyle, zorder=2) + all_curves.append(curve) + + if all_curves: + mean_curve = np.nanmean(np.array(all_curves), axis=0) + ax.plot(leads, mean_curve, + color=color, lw=2.8, alpha=0.95, + linestyle=linestyle, zorder=5, + label=f"{label_prefix} — mean across all init times") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--route-dir", default=DEFAULT_ROUTE_DIR) + parser.add_argument("--out-dir", default=None) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--da-name", default="routed_leadtime_da_full.parquet") + parser.add_argument("--ol-name", default="routed_leadtime_openloop_full.parquet") + args = parser.parse_args() + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + print("Loading parquets...") + da = load_parquet(os.path.join(args.route_dir, args.da_name)) + ol = load_parquet(os.path.join(args.route_dir, args.ol_name)) + obs = load_usgs(args.usgs_csv) + print(f" DA issue_times: {da['issue_time'].nunique()} " + f"OL issue_times: {ol['issue_time'].nunique()}") + + da_err = build_error_table(da, obs, HELENE_INIT_START, HELENE_INIT_END) + ol_err = build_error_table(ol, obs, HELENE_INIT_START, HELENE_INIT_END) + print(f" DA init times in Helene window: {da_err['issue_time'].nunique()}") + + leads = sorted(da_err["lead_hour"].unique()) + + fig, ax = plt.subplots(figsize=(12, 6)) + ax.axhline(0, color="black", lw=0.8, linestyle="--", alpha=0.5, zorder=1) + + plot_error(ax, ol_err, color="tab:gray", alpha_thin=0.12, lw_thin=0.7, + linestyle="--", label_prefix="Open-loop") + plot_error(ax, da_err, color="tab:blue", alpha_thin=0.18, lw_thin=0.8, + linestyle="-", label_prefix="DA") + + patches = [mpatches.Patch(color=c, label=f"Init {d}") + for d, c in DATE_COLORS.items()] + handles, labels = ax.get_legend_handles_labels() + ax.legend(handles=handles + patches, + fontsize=8, loc="upper left", frameon=True, framealpha=0.9, ncol=2) + + ax.set_xlabel("Forecast lead hour", fontsize=11) + ax.set_ylabel("Error: forecast − USGS obs (m³/s)", fontsize=11) + ax.set_xticks(np.arange(1, 19)) + ax.set_title( + "Forecast error vs lead time — per initialization time, Helene window\n" + "F1 (variance-scaled Vrugt) | DA (blue solid) vs Open-loop (gray dashed) | " + "Sep 24–28 2024 | USGS 03463300", + fontsize=11, + ) + ax.grid(True, alpha=0.25, lw=0.4) + plt.tight_layout() + out1 = os.path.join(out_dir, "forecast_error_per_init_helene.png") + plt.savefig(out1, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + fig, ax = plt.subplots(figsize=(12, 6)) + + def mean_abs_error_by_lead(err_df): + return err_df.groupby("lead_hour")["error"].apply( + lambda x: float(np.nanmean(np.abs(x))) + ) + + da_mae = mean_abs_error_by_lead(da_err) + ol_mae = mean_abs_error_by_lead(ol_err) + + ax.plot(da_mae.index, da_mae.values, + color="tab:blue", lw=2.4, marker="o", label="DA — mean |error|") + ax.plot(ol_mae.index, ol_mae.values, + color="tab:gray", lw=2.4, marker="s", linestyle="--", + label="Open-loop — mean |error|") + + ax.set_xlabel("Forecast lead hour", fontsize=11) + ax.set_ylabel("Mean |error| vs USGS obs (m³/s)", fontsize=11) + ax.set_xticks(np.arange(1, 19)) + ax.set_title( + "Mean absolute forecast error vs lead time — Helene window\n" + "F1 (variance-scaled Vrugt) | DA (blue) vs Open-loop (gray) | " + "Sep 24–28 2024 | USGS 03463300", + fontsize=11, + ) + ax.legend(fontsize=10, loc="upper left", frameon=True, framealpha=0.9) + ax.grid(True, alpha=0.25, lw=0.4) + plt.tight_layout() + out2 = os.path.join(out_dir, "forecast_mae_per_lead_helene.png") + plt.savefig(out2, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4a_error_decay/plot_lead_time_decay.py b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4a_error_decay/plot_lead_time_decay.py new file mode 100644 index 00000000..7da0b64c --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4a_error_decay/plot_lead_time_decay.py @@ -0,0 +1,174 @@ +""" +Catchment-level error-vs-lead-time decay curve + ensemble spread by lead. + +Two-panel figure for ONE catchment: + TOP — RMSE of ensemble-mean forecast at each lead hour (1..18) vs the + catchment's kriging observation. DA solid, open-loop dashed, with + a shaded band showing the min/max of per-member RMSE. + BOTTOM — Mean ensemble spread (std-dev across 20 members, averaged across + all issue times) at each lead hour. Tells you whether forcing + perturbation alone keeps the forecast ensemble diverse during the + 18-hour free-run — useful for inspecting individual members + without needing the full spaghetti view. + +This is the catchment-level analog of the gauge-level decay curve. +Routing to the gauge is a separate post-step (route_lead_time_forecasts.py ++ a gauge-level decay script); this script lets us look at the catchment-level +signal without T-route in the loop. + +Inputs (from run_lead_time_forecast_sweep.py): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv +Obs (from production EnKF run): + //_test_results.csv + +Output: + //_lead_time_decay.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" + +DEFAULT_LEADTIME_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast" +DEFAULT_DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" + +# Module-level globals set by main() so the helpers see them +CAT = "cat-1016300" +LEADTIME_DIR = DEFAULT_LEADTIME_DIR +DA_DIR = DEFAULT_DA_DIR +OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay.png") + + +def load_forecasts(path): + """Return (df, member_cols).""" + if not os.path.exists(path): + raise FileNotFoundError(path) + df = pd.read_csv(path, parse_dates=['issue_time', 'valid_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + return df, member_cols + + +OBS_DIR = None # set by main() when --obs-dir is supplied + + +def load_obs(): + if OBS_DIR is not None: + obs_path = os.path.join(OBS_DIR, f"{CAT}.csv") + df = pd.read_csv(obs_path) + time_col = 'datetime' if 'datetime' in df.columns else 'date' + df[time_col] = pd.to_datetime(df[time_col]) + col = 'qkrig' if 'qkrig' in df.columns else 'obs_mm_h' + return df.set_index(time_col)[col].rename('obs_mm_h') + obs_path = os.path.join(DA_DIR, CAT, f"{CAT}_test_results.csv") + df = pd.read_csv(obs_path, parse_dates=['date']) + return df.set_index('date')['obs_mm_h'] + + +def metrics_by_lead(df, member_cols, obs_series): + """For each lead hour, return (rmse_mean, rmse_min_member, rmse_max_member, + mean_ensemble_std).""" + df = df.copy() + df['obs'] = df['valid_time'].map(obs_series) + df = df.dropna(subset=['obs']) + df['ens_mean'] = df[member_cols].mean(axis=1) + df['ens_std'] = df[member_cols].std(axis=1) + + leads = sorted(df['lead_hour'].unique()) + rmse_mean = np.full(len(leads), np.nan) + rmse_min_member = np.full(len(leads), np.nan) + rmse_max_member = np.full(len(leads), np.nan) + mean_ens_std = np.full(len(leads), np.nan) + for i, L in enumerate(leads): + sub = df[df['lead_hour'] == L] + if len(sub) == 0: + continue + err = sub['ens_mean'].values - sub['obs'].values + rmse_mean[i] = float(np.sqrt(np.mean(err ** 2))) + per_member_rmse = np.array([ + np.sqrt(np.mean((sub[c].values - sub['obs'].values) ** 2)) + for c in member_cols + ]) + rmse_min_member[i] = float(np.nanmin(per_member_rmse)) + rmse_max_member[i] = float(np.nanmax(per_member_rmse)) + mean_ens_std[i] = float(np.nanmean(sub['ens_std'].values)) + return (np.asarray(leads), rmse_mean, rmse_min_member, + rmse_max_member, mean_ens_std) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', default="cat-1016300") + parser.add_argument('--leadtime-dir', default=DEFAULT_LEADTIME_DIR, + help='Dir holding /_lead_time_forecasts_{da,openloop}.csv') + parser.add_argument('--da-dir', default=DEFAULT_DA_DIR, + help='Dir holding /_test_results.csv (for obs)') + parser.add_argument('--obs-dir', default=None, + help='Kriging obs dir holding .csv with qkrig column; ' + 'overrides --da-dir for obs loading') + args = parser.parse_args() + + global CAT, LEADTIME_DIR, DA_DIR, OBS_DIR, OUT_PNG + CAT = args.cat_id + LEADTIME_DIR = args.leadtime_dir + DA_DIR = args.da_dir + OBS_DIR = args.obs_dir + OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay.png") + + da_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_da.csv") + ol_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_openloop.csv") + df_da, m_cols = load_forecasts(da_path) + df_ol, _ = load_forecasts(ol_path) + obs_series = load_obs() + + leads_da, rmse_da, rmse_da_lo, rmse_da_hi, std_da = metrics_by_lead( + df_da, m_cols, obs_series) + leads_ol, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = metrics_by_lead( + df_ol, m_cols, obs_series) + + fig, (ax_top, ax_bot) = plt.subplots(2, 1, figsize=(10, 10), sharex=True, + gridspec_kw={"height_ratios": [1.4, 1.0]}) + + # ----- TOP: RMSE decay curve ----- + ax_top.fill_between(leads_da, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2, + label="DA — member spread (min/max RMSE)") + ax_top.fill_between(leads_ol, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2, + label="Open-loop — member spread") + ax_top.plot(leads_da, rmse_da, color=DA_COLOR, lw=2.4, marker='o', + zorder=4, label="DA on (ensemble-mean RMSE)") + ax_top.plot(leads_ol, rmse_ol, color=OL_COLOR, lw=2.4, marker='s', + linestyle='--', zorder=4, label="Open-loop (ensemble-mean RMSE)") + ax_top.set_ylabel("RMSE vs Qkrig obs (mm/h)", fontsize=11) + ax_top.set_title( + f"Forecast lead-time error decay — {CAT}\n" + f"Issue times pooled across test period (Oct 2023 – Oct 2024)", + fontsize=12, + ) + ax_top.grid(True, alpha=0.25) + ax_top.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + # ----- BOTTOM: ensemble spread by lead hour ----- + ax_bot.plot(leads_da, std_da, color=DA_COLOR, lw=2.2, marker='o', + label="DA — mean ensemble std") + ax_bot.plot(leads_ol, std_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', label="Open-loop — mean ensemble std") + ax_bot.set_xlabel("Forecast lead time (hours)", fontsize=11) + ax_bot.set_ylabel("Mean ensemble std (mm/h)", fontsize=11) + ax_bot.set_xticks(np.arange(1, 19)) + ax_bot.grid(True, alpha=0.25) + ax_bot.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + plt.tight_layout() + plt.savefig(OUT_PNG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4a_error_decay/plot_lead_time_decay_by_regime.py b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4a_error_decay/plot_lead_time_decay_by_regime.py new file mode 100644 index 00000000..92b680b0 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4a_error_decay/plot_lead_time_decay_by_regime.py @@ -0,0 +1,214 @@ +""" +Lead-time decay curve, split by flow regime at issue time. + +The pooled lead-time curve (plot_lead_time_decay.py) showed DA losing to +open-loop across most lead hours, but the test period is ~99% low-flow. +This script splits the same forecast CSVs by the flow regime at the issue +time t0, so we can see whether DA helps when it matters (storms / Helene) +and hurts when it doesn't (low flow). + +Three regimes are partitioned on obs(t0) (kriging obs at the issue time): + + helene issue_time ∈ [2024-09-24, 2024-09-28] (the 5-day Helene window) + storm obs(t0) > 0.5 mm/h (high-flow issue times) + low_flow obs(t0) < 0.1 mm/h (low-flow issue times) + +Same 2-panel layout per regime: RMSE vs lead (top), ensemble spread vs lead +(bottom). Three regime columns × 2 metric rows in one figure. + +Inputs (from run_lead_time_forecast_sweep.py — no re-run needed): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv +Obs (from production EnKF): + //_test_results.csv + +Output: + //_lead_time_decay_by_regime.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" + +DEFAULT_LEADTIME_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast" +DEFAULT_DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" + +# Module-level globals set by main() so helpers see them +CAT = "cat-1016300" +LEADTIME_DIR = DEFAULT_LEADTIME_DIR +DA_DIR = DEFAULT_DA_DIR +OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay_by_regime.png") + +# Regime definitions +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-28 23:00:00") +STORM_OBS_THRESHOLD = 0.5 # mm/h +LOWFLOW_OBS_THRESHOLD = 0.1 # mm/h + + +def load_forecasts(path): + if not os.path.exists(path): + raise FileNotFoundError(path) + df = pd.read_csv(path, parse_dates=['issue_time', 'valid_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + return df, member_cols + + +OBS_DIR = None # set by main() when --obs-dir is supplied + + +def load_obs(): + if OBS_DIR is not None: + obs_path = os.path.join(OBS_DIR, f"{CAT}.csv") + df = pd.read_csv(obs_path) + time_col = 'datetime' if 'datetime' in df.columns else 'date' + df[time_col] = pd.to_datetime(df[time_col]) + col = 'qkrig' if 'qkrig' in df.columns else 'obs_mm_h' + return df.set_index(time_col)[col].rename('obs_mm_h') + obs_path = os.path.join(DA_DIR, CAT, f"{CAT}_test_results.csv") + df = pd.read_csv(obs_path, parse_dates=['date']) + return df.set_index('date')['obs_mm_h'] + + +def regime_mask(issue_times, obs_at_issue, regime): + """Boolean mask over issue_times for the named regime.""" + if regime == "helene": + return (issue_times >= HELENE_START) & (issue_times <= HELENE_END) + if regime == "storm": + return obs_at_issue > STORM_OBS_THRESHOLD + if regime == "low_flow": + return obs_at_issue < LOWFLOW_OBS_THRESHOLD + raise ValueError(regime) + + +def metrics_by_lead(df, member_cols, obs_series, issue_mask): + """Compute (leads, rmse_mean, rmse_min, rmse_max, mean_std) restricted + to issue times where issue_mask is True.""" + df = df.copy() + df = df[df['issue_time'].isin(issue_mask)] + if len(df) == 0: + return None + df['obs'] = df['valid_time'].map(obs_series) + df = df.dropna(subset=['obs']) + if len(df) == 0: + return None + df['ens_mean'] = df[member_cols].mean(axis=1) + df['ens_std'] = df[member_cols].std(axis=1) + + leads = sorted(df['lead_hour'].unique()) + rmse_mean = np.full(len(leads), np.nan) + rmse_lo = np.full(len(leads), np.nan) + rmse_hi = np.full(len(leads), np.nan) + std_mean = np.full(len(leads), np.nan) + for i, L in enumerate(leads): + sub = df[df['lead_hour'] == L] + if len(sub) == 0: + continue + err = sub['ens_mean'].values - sub['obs'].values + rmse_mean[i] = float(np.sqrt(np.mean(err ** 2))) + per_member = np.array([ + np.sqrt(np.mean((sub[c].values - sub['obs'].values) ** 2)) + for c in member_cols + ]) + rmse_lo[i] = float(np.nanmin(per_member)) + rmse_hi[i] = float(np.nanmax(per_member)) + std_mean[i] = float(np.nanmean(sub['ens_std'].values)) + return np.asarray(leads), rmse_mean, rmse_lo, rmse_hi, std_mean + + +def plot_pair(ax_rmse, ax_std, da_metrics, ol_metrics, + regime_label, n_issue_times): + if da_metrics is None or ol_metrics is None: + ax_rmse.set_title(f"{regime_label}\n(no issue times match)", fontsize=11) + ax_rmse.axis('off'); ax_std.axis('off') + return + leads, rmse_da, rmse_da_lo, rmse_da_hi, std_da = da_metrics + _, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = ol_metrics + + # RMSE panel + ax_rmse.fill_between(leads, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2) + ax_rmse.fill_between(leads, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2) + ax_rmse.plot(leads, rmse_da, color=DA_COLOR, lw=2.2, marker='o', + zorder=4, label="DA on") + ax_rmse.plot(leads, rmse_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', zorder=4, label="Open-loop") + ax_rmse.set_title(f"{regime_label}\n({n_issue_times} issue times)", fontsize=11) + ax_rmse.set_ylabel("RMSE (mm/h)", fontsize=10) + ax_rmse.grid(True, alpha=0.25) + ax_rmse.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + # Spread panel + ax_std.plot(leads, std_da, color=DA_COLOR, lw=2.0, marker='o', label="DA") + ax_std.plot(leads, std_ol, color=OL_COLOR, lw=2.0, marker='s', + linestyle='--', label="Open-loop") + ax_std.set_xlabel("Forecast lead time (hours)", fontsize=10) + ax_std.set_ylabel("Mean ensemble std (mm/h)", fontsize=10) + ax_std.set_xticks(np.arange(1, 19)) + ax_std.grid(True, alpha=0.25) + ax_std.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', default="cat-1016300") + parser.add_argument('--leadtime-dir', default=DEFAULT_LEADTIME_DIR, + help='Dir holding /_lead_time_forecasts_{da,openloop}.csv') + parser.add_argument('--da-dir', default=DEFAULT_DA_DIR, + help='Dir holding /_test_results.csv (for obs)') + parser.add_argument('--obs-dir', default=None, + help='Kriging obs dir holding .csv with qkrig column; ' + 'overrides --da-dir for obs loading') + args = parser.parse_args() + + global CAT, LEADTIME_DIR, DA_DIR, OBS_DIR, OUT_PNG + CAT = args.cat_id + LEADTIME_DIR = args.leadtime_dir + DA_DIR = args.da_dir + OBS_DIR = args.obs_dir + OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay_by_regime.png") + + da_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_da.csv") + ol_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_openloop.csv") + df_da, m_cols = load_forecasts(da_path) + df_ol, _ = load_forecasts(ol_path) + obs_series = load_obs() + + # Unique issue times across the run, with obs at each + issue_times = pd.to_datetime(sorted(df_da['issue_time'].unique())) + obs_at_issue = pd.Series(issue_times, index=issue_times).map(obs_series) + + regimes = [ + ("helene", "Helene window (Sep 24–28 2024)"), + ("storm", f"Storm hours (obs(t0) > {STORM_OBS_THRESHOLD} mm/h)"), + ("low_flow", f"Low flow (obs(t0) < {LOWFLOW_OBS_THRESHOLD} mm/h)"), + ] + + fig, axes = plt.subplots(2, 3, figsize=(18, 10), + gridspec_kw={"height_ratios": [1.4, 1.0]}) + for col, (key, label) in enumerate(regimes): + mask = regime_mask(issue_times, obs_at_issue, key) + kept = issue_times[mask] + kept_set = set(pd.to_datetime(kept)) + da_m = metrics_by_lead(df_da, m_cols, obs_series, kept_set) + ol_m = metrics_by_lead(df_ol, m_cols, obs_series, kept_set) + plot_pair(axes[0, col], axes[1, col], da_m, ol_m, label, len(kept)) + + fig.suptitle( + f"Lead-time error decay by flow regime — {CAT}\n" + f"Same forecast CSVs as the pooled view, partitioned on obs(t0).", + fontsize=13, y=0.995, + ) + plt.tight_layout(rect=(0, 0, 1, 0.97)) + plt.savefig(OUT_PNG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4a_error_decay/plot_lead_time_decay_gauge.py b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4a_error_decay/plot_lead_time_decay_gauge.py new file mode 100644 index 00000000..71ce54fc --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4a_error_decay/plot_lead_time_decay_gauge.py @@ -0,0 +1,311 @@ +""" +Gauge-level lead-time forecast decay curve. + +Routes 21-catchment forecast CSVs through T-route to USGS gauge 03463300 and +plots RMSE vs lead time (DA solid, open-loop dashed) against USGS observed Q. +Same 2-panel layout as the catchment-level plot, plus a regime split for the +Helene window. + +The routing step is done separately by route_lead_time_forecasts.py; +this script just reads the resulting parquets and produces the plot. + +Inputs: + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + Each: one row per (issue_time, lead_hour); member columns hold q_gauge_m3s. + (Long-format with explicit `member` and `q_gauge_m3s` columns also supported.) +Obs: + /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + USGS hourly Q in m³/s at gauge 03463300 (South Toe River near Celo, NC). + +Output: + /lead_time_decay_gauge_pooled.png + /lead_time_decay_gauge_by_regime.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" +OBS_COLOR = "black" + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/leadtime_troute_routing" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +# Watershed area at gauge 03463300 (sum of 21 catchments per GPKG hydrofabric) +WATERSHED_AREA_KM2 = 113.18 +# 1 mm/h depth × 113.18 km² = 113.18e3 m³/h = 31.439 m³/s +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-28 23:00:00") +STORM_OBS_THRESHOLD_M3S = 20.0 # gauge-level storm threshold (was 50 — too high + # for this 113 km² basin; only 2 issue times qualified) +LOWFLOW_OBS_THRESHOLD_M3S = 5.0 # gauge-level low-flow threshold + +USGS_HELENE_PEAK_M3S = 1886.0 # reference Sep 27 14:00 + + +def load_parquet_long(path): + """Load a routed parquet and normalize to long format: + columns = issue_time, lead_hour, member, q_gauge_m3s. + + Handles both wide format (member_00..member_19 columns) and long format + (explicit `member` + `q_gauge_m3s` columns). Auto-detects. + """ + if not os.path.exists(path): + raise FileNotFoundError(path) + df = pd.read_parquet(path) + + # Long format detection + if 'q_gauge_m3s' in df.columns and 'member' in df.columns: + df['issue_time'] = pd.to_datetime(df['issue_time']) + return df[['issue_time', 'lead_hour', 'member', 'q_gauge_m3s']] + + # Some variants might call it 'q_m3s' or similar + q_col_candidates = [c for c in df.columns if c.lower() in + ('q_gauge_m3s', 'q_m3s', 'q_gauge', 'q')] + if 'member' in df.columns and q_col_candidates: + qc = q_col_candidates[0] + df['issue_time'] = pd.to_datetime(df['issue_time']) + out = df[['issue_time', 'lead_hour', 'member', qc]].copy() + out = out.rename(columns={qc: 'q_gauge_m3s'}) + return out + + # Wide format: member_00..member_19 columns + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + if not member_cols: + raise ValueError(f"Couldn't identify member columns in {path}. " + f"Columns present: {df.columns.tolist()}") + df['issue_time'] = pd.to_datetime(df['issue_time']) + keep = ['issue_time', 'lead_hour'] + if 'valid_time' in df.columns: + keep.append('valid_time') + long = df[keep + member_cols].melt( + id_vars=keep, + value_vars=member_cols, + var_name='member', + value_name='q_gauge_m3s', + ) + return long + + +def load_usgs_obs(usgs_csv): + """Load USGS gauge obs and ensure m³/s. + + The CSV at gauge 03463300 has the column `QObs(mm/h)` (catchment-averaged + depth, not gauge discharge in m³/s). If the column name contains 'mm', we + convert mm/h → m³/s by multiplying by the watershed area factor: + m³/s = mm/h × WATERSHED_AREA_KM2 × 1000 / 3600 + Verified: 59.978 mm/h × 113.18 km² = 1886 m³/s (USGS Helene peak). + """ + df = pd.read_csv(usgs_csv) + date_candidates = [c for c in df.columns if c.lower() in + ('datetime', 'date', 'time', 'timestamp')] + q_candidates = [c for c in df.columns if 'q' in c.lower() + or 'flow' in c.lower() or 'discharge' in c.lower()] + if not date_candidates or not q_candidates: + raise ValueError(f"Could not identify date/Q columns in {usgs_csv}. " + f"Columns: {df.columns.tolist()}") + dc, qc = date_candidates[0], q_candidates[0] + df[dc] = pd.to_datetime(df[dc]) + series = df.set_index(dc)[qc].astype(float) + if 'mm' in qc.lower(): + series = series * MM_H_TO_M3_S + print(f" Converted obs from mm/h (column '{qc}') to m³/s " + f"using area = {WATERSHED_AREA_KM2} km² (× {MM_H_TO_M3_S:.4f})") + else: + print(f" Loaded obs as-is from column '{qc}' (assumed m³/s)") + return series + + +def metrics_by_lead(df_long, obs_series, issue_mask=None): + """Compute per-lead-hour: ensemble-mean RMSE, per-member min/max RMSE, + mean ensemble std. df_long must have issue_time, lead_hour, member, q_gauge_m3s. + """ + df = df_long.copy() + if issue_mask is not None: + df = df[df['issue_time'].isin(issue_mask)] + if len(df) == 0: + return None + df['valid_time'] = df['issue_time'] + pd.to_timedelta(df['lead_hour'], unit='h') + df['obs'] = df['valid_time'].map(obs_series) + df = df.dropna(subset=['obs']) + if len(df) == 0: + return None + + # Ensemble mean per (issue_time, lead_hour) + grouped = df.groupby(['issue_time', 'lead_hour']) + ens_mean = grouped['q_gauge_m3s'].mean().reset_index(name='ens_mean') + ens_std = grouped['q_gauge_m3s'].std().reset_index(name='ens_std') + obs_per = grouped['obs'].first().reset_index(name='obs') + panel = ens_mean.merge(ens_std, on=['issue_time', 'lead_hour']).merge( + obs_per, on=['issue_time', 'lead_hour']) + + leads = sorted(panel['lead_hour'].unique()) + rmse_mean = np.full(len(leads), np.nan) + rmse_lo = np.full(len(leads), np.nan) + rmse_hi = np.full(len(leads), np.nan) + std_mean = np.full(len(leads), np.nan) + for i, L in enumerate(leads): + psub = panel[panel['lead_hour'] == L] + if len(psub) == 0: + continue + err = psub['ens_mean'].values - psub['obs'].values + rmse_mean[i] = float(np.sqrt(np.mean(err ** 2))) + # Per-member RMSE — recompute from the underlying long frame + dsub = df[df['lead_hour'] == L] + member_rmses = [] + for m, g in dsub.groupby('member'): + err_m = g['q_gauge_m3s'].values - g['obs'].values + if len(err_m) >= 1: + member_rmses.append(float(np.sqrt(np.mean(err_m ** 2)))) + if member_rmses: + rmse_lo[i] = float(np.nanmin(member_rmses)) + rmse_hi[i] = float(np.nanmax(member_rmses)) + std_mean[i] = float(np.nanmean(psub['ens_std'].values)) + return np.asarray(leads), rmse_mean, rmse_lo, rmse_hi, std_mean + + +def plot_two_panel(out_path, leads_da, da_metrics, leads_ol, ol_metrics, title): + fig, (ax_top, ax_bot) = plt.subplots(2, 1, figsize=(10, 10), sharex=True, + gridspec_kw={"height_ratios": [1.4, 1.0]}) + leads, rmse_da, rmse_da_lo, rmse_da_hi, std_da = da_metrics + _, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = ol_metrics + + ax_top.fill_between(leads, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2, + label="DA — member spread (min/max RMSE)") + ax_top.fill_between(leads, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2, + label="Open-loop — member spread") + ax_top.plot(leads, rmse_da, color=DA_COLOR, lw=2.4, marker='o', + zorder=4, label="DA on (ensemble-mean RMSE)") + ax_top.plot(leads, rmse_ol, color=OL_COLOR, lw=2.4, marker='s', + linestyle='--', zorder=4, label="Open-loop (ensemble-mean RMSE)") + ax_top.set_ylabel("RMSE vs USGS obs (m³/s)", fontsize=11) + ax_top.set_title(title, fontsize=12) + ax_top.grid(True, alpha=0.25) + ax_top.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + ax_bot.plot(leads, std_da, color=DA_COLOR, lw=2.2, marker='o', + label="DA — mean ensemble std") + ax_bot.plot(leads, std_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', label="Open-loop — mean ensemble std") + ax_bot.set_xlabel("Forecast lead time (hours)", fontsize=11) + ax_bot.set_ylabel("Mean ensemble std (m³/s)", fontsize=11) + ax_bot.set_xticks(np.arange(1, 19)) + ax_bot.grid(True, alpha=0.25) + ax_bot.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + plt.tight_layout() + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + +def plot_three_panel_regime(out_path, regime_metrics): + fig, axes = plt.subplots(2, 3, figsize=(18, 10), + gridspec_kw={"height_ratios": [1.4, 1.0]}) + for col, (label, n_issue, da_m, ol_m) in enumerate(regime_metrics): + ax_rmse = axes[0, col] + ax_std = axes[1, col] + if da_m is None or ol_m is None: + ax_rmse.set_title(f"{label}\n(no issue times match)", fontsize=11) + ax_rmse.axis('off'); ax_std.axis('off') + continue + leads, rmse_da, rmse_da_lo, rmse_da_hi, std_da = da_m + _, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = ol_m + ax_rmse.fill_between(leads, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2) + ax_rmse.fill_between(leads, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2) + ax_rmse.plot(leads, rmse_da, color=DA_COLOR, lw=2.2, marker='o', + zorder=4, label="DA on") + ax_rmse.plot(leads, rmse_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', zorder=4, label="Open-loop") + ax_rmse.set_title(f"{label}\n({n_issue} issue times)", fontsize=11) + ax_rmse.set_ylabel("RMSE (m³/s)", fontsize=10) + ax_rmse.grid(True, alpha=0.25) + ax_rmse.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + ax_std.plot(leads, std_da, color=DA_COLOR, lw=2.0, marker='o', label="DA") + ax_std.plot(leads, std_ol, color=OL_COLOR, lw=2.0, marker='s', + linestyle='--', label="Open-loop") + ax_std.set_xlabel("Forecast lead time (hours)", fontsize=10) + ax_std.set_ylabel("Mean ensemble std (m³/s)", fontsize=10) + ax_std.set_xticks(np.arange(1, 19)) + ax_std.grid(True, alpha=0.25) + ax_std.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + fig.suptitle("Gauge-level lead-time error decay by regime — USGS 03463300", + fontsize=13, y=0.995) + plt.tight_layout(rect=(0, 0, 1, 0.97)) + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--route-dir', default=DEFAULT_ROUTE_DIR, + help='Dir holding routed_leadtime_{da,openloop}_full.parquet') + parser.add_argument('--out-dir', default=None, + help='Where to write the output PNGs. Defaults to --route-dir ' + '(which may not be writable if owned by another user — ' + 'pass an explicit path then).') + parser.add_argument('--usgs-csv', default=DEFAULT_USGS_CSV) + parser.add_argument('--da-name', default='routed_leadtime_da_full.parquet') + parser.add_argument('--ol-name', default='routed_leadtime_openloop_full.parquet') + args = parser.parse_args() + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + da_path = os.path.join(args.route_dir, args.da_name) + ol_path = os.path.join(args.route_dir, args.ol_name) + df_da = load_parquet_long(da_path) + df_ol = load_parquet_long(ol_path) + obs_series = load_usgs_obs(args.usgs_csv) + + print(f"DA rows: {len(df_da):,} | unique issue_times: {df_da['issue_time'].nunique()}") + print(f"OL rows: {len(df_ol):,} | unique issue_times: {df_ol['issue_time'].nunique()}") + print(f"USGS obs range: {obs_series.index.min()} .. {obs_series.index.max()} " + f"({len(obs_series):,} hours)") + print(f"USGS Helene peak (reference): {USGS_HELENE_PEAK_M3S:.0f} m³/s") + + # ----- Pooled (all issue times) ----- + da_pooled = metrics_by_lead(df_da, obs_series) + ol_pooled = metrics_by_lead(df_ol, obs_series) + out_pooled = os.path.join(out_dir, "lead_time_decay_gauge_pooled.png") + plot_two_panel(out_pooled, da_pooled[0], da_pooled, ol_pooled[0], ol_pooled, + title="Gauge-level lead-time decay — USGS 03463300\n" + "All issue times pooled (Oct 2023 – Oct 2024)") + + # ----- Regime split: Helene, storm, low-flow at gauge ----- + issue_times = pd.to_datetime(sorted(df_da['issue_time'].unique())) + obs_at_issue = pd.Series(issue_times, index=issue_times).map(obs_series) + + helene_mask = (issue_times >= HELENE_START) & (issue_times <= HELENE_END) + storm_mask = obs_at_issue > STORM_OBS_THRESHOLD_M3S + lowflow_mask = obs_at_issue < LOWFLOW_OBS_THRESHOLD_M3S + + regime_metrics = [] + for label, mask in [ + (f"Helene window (Sep 24–28 2024)", helene_mask), + (f"Storm hours (obs(t0) > {STORM_OBS_THRESHOLD_M3S:.0f} m³/s)", storm_mask), + (f"Low flow (obs(t0) < {LOWFLOW_OBS_THRESHOLD_M3S:.0f} m³/s)", lowflow_mask), + ]: + kept = set(pd.to_datetime(issue_times[mask])) + da_m = metrics_by_lead(df_da, obs_series, kept) + ol_m = metrics_by_lead(df_ol, obs_series, kept) + regime_metrics.append((label, len(kept), da_m, ol_m)) + + out_regime = os.path.join(out_dir, "lead_time_decay_gauge_by_regime.png") + plot_three_panel_regime(out_regime, regime_metrics) + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4a_error_decay/run_4a_f1.sh b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4a_error_decay/run_4a_f1.sh new file mode 100644 index 00000000..a75ef20e --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4a_error_decay/run_4a_f1.sh @@ -0,0 +1,78 @@ +#!/usr/bin/env bash +# F1 Vrugt (1 gauge holdout) — 4a lead-time decay plots. +# +# Runs 3 plot scripts per catchment: +# plot_lead_time_decay.py — pooled RMSE vs lead (catchment-level) +# plot_lead_time_decay_by_regime.py — same split by flow regime +# plot_lead_time_decay_gauge.py — gauge-level (requires routed parquets) +# +# plot_lead_time_decay_gauge.py is skipped if the routed parquets don't exist. +# +# Usage: +# bash run_4a_f1.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +LEADTIME_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt +DA_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt +OBS_DIR=/home/svyas/catchment_ts_no_03463300_gapfilled +ROUTE_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt_leadtime_routed +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +echo "[F1-4a] Lead-time decay plots — leadtime dir: $LEADTIME_DIR" + +for CAT in "${CATS[@]}"; do + DA_CSV="$LEADTIME_DIR/$CAT/${CAT}_lead_time_forecasts_da.csv" + if [ ! -f "$DA_CSV" ]; then + echo " [$CAT] No lead-time CSV — skipping" + continue + fi + + echo " [$CAT] Plotting pooled decay..." + $TROUTE "$SCRIPT_DIR/plot_lead_time_decay.py" \ + --cat-id "$CAT" \ + --leadtime-dir "$LEADTIME_DIR" \ + --da-dir "$DA_DIR" \ + --obs-dir "$OBS_DIR" + + echo " [$CAT] Plotting decay by regime..." + $TROUTE "$SCRIPT_DIR/plot_lead_time_decay_by_regime.py" \ + --cat-id "$CAT" \ + --leadtime-dir "$LEADTIME_DIR" \ + --da-dir "$DA_DIR" \ + --obs-dir "$OBS_DIR" +done + +# Gauge-level decay (needs routed parquets — run after route_lead_time_forecasts.py) +if [ -f "$ROUTE_DIR/routed_leadtime_da_full.parquet" ]; then + echo "[F1-4a] Gauge-level decay (routed)..." + $TROUTE "$SCRIPT_DIR/plot_lead_time_decay_gauge.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" + + echo "[F1-4a] Fixed-target error plots..." + $TROUTE "$SCRIPT_DIR/plot_forecast_error_fixed_target.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" + + echo "[F1-4a] Per-init error plots..." + $TROUTE "$SCRIPT_DIR/plot_forecast_error_per_init.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" +else + echo "[F1-4a] Routed parquets not found — skipping gauge-level and error plots." + echo " Run route_leadtime_f1.sh first, then re-run this script." +fi + +echo "[F1-4a] Done." diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_crossed_ensemble.py b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_crossed_ensemble.py new file mode 100644 index 00000000..b498bb53 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_crossed_ensemble.py @@ -0,0 +1,229 @@ +""" +plot_crossed_ensemble.py -- 4b: 600-member crossed ensemble vs USGS obs. + +Loads the crossed ensemble parquet produced by run_crossed_ensemble.py: + //_crossed_ensemble.parquet + Columns: issue_time, lead_hour, member_0000..member_0599 + +Projects all members to valid_time = issue_time + lead_hour hours. +Aggregates at each valid_time across ALL members from ALL issue times +to build a probabilistic envelope (5th/25th/50th/75th/95th percentiles). + +Two output figures: + 4b_crossed_ensemble_helene.png + -- Full Helene window (Sep 24-30) probabilistic envelope + USGS obs. + Percentile shading in two layers (5-95 outer, 25-75 inner). + 4b_crossed_ensemble_peak.png + -- Zoomed to Helene peak (Sep 26 12 UTC -> Sep 28 06 UTC), + same shading + USGS obs + individual spaghetti from Sep 27 inits. + +Usage: + python3 plot_crossed_ensemble.py \\ + --ensemble-dir /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble \\ + --cat-id cat-1016300 \\ + --usgs-csv /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv \\ + --out-dir /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble/cat-1016300 +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import matplotlib.patches as mpatches + +DEFAULT_ENSEMBLE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" +DEFAULT_CAT_ID = "cat-1016300" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +PLOT_START = pd.Timestamp("2024-09-24 00:00:00") +PLOT_END = pd.Timestamp("2024-09-28 18:00:00") # cut trailing spike +PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +PEAK_END = pd.Timestamp("2024-09-28 06:00:00") +HELENE_PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +HELENE_PEAK_END = pd.Timestamp("2024-09-28 00:00:00") + + +def load_ensemble(path): + df = pd.read_parquet(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + return df, member_cols + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def build_percentile_envelope(df, member_cols, t_start, t_end): + """ + For each valid_time in [t_start, t_end], stack all member values + from all issue times and compute percentiles. + Returns DataFrame indexed by valid_time with columns p05, p25, p50, p75, p95. + """ + mask = (df["valid_time"] >= t_start) & (df["valid_time"] <= t_end) + sub = df[mask].copy() + sub["valid_time"] = sub["valid_time"].dt.floor("1h") + + records = [] + for vt, grp in sub.groupby("valid_time"): + vals = grp[member_cols].to_numpy(dtype=float).flatten() + vals = vals[~np.isnan(vals)] * MM_H_TO_M3_S + if len(vals) == 0: + continue + records.append({ + "valid_time": vt, + "p05": np.percentile(vals, 5), + "p25": np.percentile(vals, 25), + "p50": np.percentile(vals, 50), + "p75": np.percentile(vals, 75), + "p95": np.percentile(vals, 95), + "n": len(vals), + }) + return pd.DataFrame(records).set_index("valid_time").sort_index() + + +def _add_helene_band(ax, label=True): + ax.axvspan(HELENE_PEAK_START, HELENE_PEAK_END, + color="salmon", alpha=0.12, zorder=0) + if label: + ax.text(HELENE_PEAK_START + pd.Timedelta(hours=6), 0.97, "Helene peak", + transform=ax.get_xaxis_transform(), + fontsize=9, color="firebrick", ha="left", va="top") + + +def _format_xaxis(ax, major_interval=1, minor_hours=6): + ax.xaxis.set_major_locator(mdates.DayLocator(interval=major_interval)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.xaxis.set_minor_locator(mdates.HourLocator(interval=minor_hours)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.2, lw=0.4) + + +def plot_envelope(ax, env, obs_series, t_start, t_end, + ens_color="tab:blue", obs_color="black", helene_label=True): + _add_helene_band(ax, label=helene_label) + + # Outer band: 5-95 + ax.fill_between(env.index, env["p05"], env["p95"], + color=ens_color, alpha=0.15, zorder=2, + label="5th-95th percentile") + # Inner band: 25-75 + ax.fill_between(env.index, env["p25"], env["p75"], + color=ens_color, alpha=0.30, zorder=3, + label="25th-75th percentile") + # Median + ax.plot(env.index, env["p50"], + color=ens_color, lw=2.2, zorder=5, + label="Ensemble median (600 members)") + + # USGS obs + obs_w = obs_series.loc[t_start:t_end] + ax.plot(obs_w.index, obs_w.values, + color=obs_color, lw=1.8, zorder=6, label="USGS obs") + + ax.set_xlim(t_start, t_end) + ax.set_ylabel("Discharge (m3/s)", fontsize=11) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--ensemble-dir", default=DEFAULT_ENSEMBLE_DIR) + parser.add_argument("--cat-id", default=DEFAULT_CAT_ID) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--out-dir", default=None) + args = parser.parse_args() + + cat_dir = os.path.join(args.ensemble_dir, args.cat_id) + out_dir = args.out_dir or cat_dir + os.makedirs(out_dir, exist_ok=True) + + parquet_path = os.path.join(cat_dir, f"{args.cat_id}_crossed_ensemble.parquet") + print(f"Loading crossed ensemble: {parquet_path}") + df, member_cols = load_ensemble(parquet_path) + obs = load_usgs(args.usgs_csv) + print(f" {df['issue_time'].nunique()} issue times, " + f"{len(member_cols)} members per row") + + # ------------------------------------------------------------------ # + # Figure 4b-full: full Helene window probabilistic envelope # + # ------------------------------------------------------------------ # + print("Building percentile envelope (full window)...") + env_full = build_percentile_envelope(df, member_cols, PLOT_START, PLOT_END) + print(f" Valid times with data: {len(env_full)}, " + f"median members per timestep: {env_full['n'].median():.0f}") + + fig, ax = plt.subplots(figsize=(17, 7)) + plot_envelope(ax, env_full, obs, PLOT_START, PLOT_END) + ax.set_xlabel("Date", fontsize=11) + ax.set_title( + f"4b -- 600-member crossed ensemble (30 met x 20 hydro-state) | {args.cat_id}\n" + "Shaded: 5-95th (light) and 25-75th (dark) percentile across all members " + "and issue times | Sep 24-30 2024 | USGS 03463300", + fontsize=11, + ) + ax.legend(fontsize=10, loc="upper left", frameon=True, framealpha=0.9) + _format_xaxis(ax) + plt.tight_layout() + out1 = os.path.join(out_dir, f"{args.cat_id}_4b_crossed_ensemble_helene.png") + plt.savefig(out1, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + # ------------------------------------------------------------------ # + # Figure 4b-peak: zoomed to Helene peak window # + # ------------------------------------------------------------------ # + print("Building percentile envelope (peak window)...") + env_peak = build_percentile_envelope(df, member_cols, PEAK_START, PEAK_END) + + fig, (ax_main, ax_zoom) = plt.subplots( + 2, 1, figsize=(14, 11), + gridspec_kw={"height_ratios": [1, 1.4]}, sharex=False) + + # Top panel: full window context (smaller) + plot_envelope(ax_main, env_full, obs, PLOT_START, PLOT_END, helene_label=False) + ax_main.set_title("Full Helene window context (Sep 24 - Sep 28 18 UTC)", fontsize=10) + _format_xaxis(ax_main) + # Shade the zoom region on context panel + ax_main.axvspan(PEAK_START, PEAK_END, color="gold", alpha=0.18, zorder=1) + ax_main.text(PEAK_START + pd.Timedelta(hours=1), 0.97, "zoom", + transform=ax_main.get_xaxis_transform(), + fontsize=8, color="goldenrod", ha="left", va="top") + + # Bottom panel: peak zoom + plot_envelope(ax_zoom, env_peak, obs, PEAK_START, PEAK_END, helene_label=True) + ax_zoom.set_xlabel("Date", fontsize=11) + ax_zoom.set_title( + "Helene peak zoom (Sep 26 12 UTC - Sep 28 06 UTC)", fontsize=10) + _format_xaxis(ax_zoom, major_interval=1, minor_hours=3) + ax_zoom.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %HUTC")) + ax_zoom.legend(fontsize=10, loc="lower right", frameon=True, framealpha=0.9) + + fig.suptitle( + f"4b -- 600-member crossed ensemble vs USGS 03463300 | {args.cat_id}\n" + "30 met forcing draws x 20 DA analysis state draws | R = (0.10·Q)² + 0.001·σ²_krig (Vrugt)", + fontsize=12, y=1.01, + ) + plt.tight_layout() + out2 = os.path.join(out_dir, f"{args.cat_id}_4b_crossed_ensemble_peak.png") + plt.savefig(out2, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_per_member_factor_decomposition.py b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_per_member_factor_decomposition.py new file mode 100644 index 00000000..2a1eff34 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_per_member_factor_decomposition.py @@ -0,0 +1,169 @@ +""" +Per-member factor-decomposition plot for ONE catchment. + +For each of the 20 ensemble members, show the member's actual forecast +(production: init + forcing + process + DA all active) alongside the member's +trajectory under each ISOLATED perturbation source (init only / forcing only / +process noise only). The viewer reads each panel as 'why did THIS member give +THIS forecast — which perturbation source pushed it where?' + +Note: 'same member_i' across the four CSVs is a column-name correspondence +only; the underlying random draws are independent across the four runs. So +each panel is an illustrative comparison, not a rigorous matched-seed Shapley +decomposition. The story still reads correctly: each colored line shows what +ONE realization of that perturbation source produces, and the thick line +shows what the full production setup produces. + +Inputs (per catchment): + //_production_per_member.csv (20 cols) + //_sensitivity_init.csv (20 cols) + //_sensitivity_forcing.csv (20 cols) + //_sensitivity_process.csv (20 cols) + //_test_results.csv (for Qkrig obs) + +Output: + //_per_member_factor_decomp.png +""" +import argparse +import os +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +# ---------- Configuration ---------- +CAT = "cat-1016300" # change here to do a different catchment + +PROD_DIR = "/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt" +SEN_DIR = "/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt" +DA_DIR = "/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt" +OBS_DIR = "/home/svyas/catchment_ts_no_03463300_gapfilled" +OUT_PNG = os.path.join(PROD_DIR, CAT, f"{CAT}_per_member_factor_decomp.png") + +ZOOM_START = pd.Timestamp("2024-09-24") +ZOOM_END = pd.Timestamp("2024-09-28") + +PROD_COLOR = "tab:purple" +INIT_COLOR = "tab:red" +FORCING_COLOR = "tab:blue" +PROC_COLOR = "tab:green" +OBS_COLOR = "black" + + +def load_member_csv(path): + if not os.path.exists(path): + return None, None + df = pd.read_csv(path, parse_dates=["date"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + return df["date"].values, df[member_cols] + + +def load_obs(): + obs_p = os.path.join(OBS_DIR, f"{CAT}.csv") + if os.path.exists(obs_p): + df = pd.read_csv(obs_p) + time_col = 'datetime' if 'datetime' in df.columns else 'date' + df[time_col] = pd.to_datetime(df[time_col]) + return df[time_col].values, df["qkrig"].values + p = os.path.join(DA_DIR, CAT, f"{CAT}_test_results.csv") + if not os.path.exists(p): + return None, None + df = pd.read_csv(p, parse_dates=["date"]) + return df["date"].values, df["obs_mm_h"].values + + +def main(): + global CAT, OUT_PNG + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', default=CAT) + args = parser.parse_args() + CAT = args.cat_id + OUT_PNG = os.path.join(PROD_DIR, CAT, f"{CAT}_per_member_factor_decomp.png") + + prod_dates, prod = load_member_csv(os.path.join(PROD_DIR, CAT, f"{CAT}_production_per_member.csv")) + init_dates, init = load_member_csv(os.path.join(SEN_DIR, CAT, f"{CAT}_sensitivity_init.csv")) + forc_dates, forc = load_member_csv(os.path.join(SEN_DIR, CAT, f"{CAT}_sensitivity_forcing.csv")) + proc_dates, proc = load_member_csv(os.path.join(SEN_DIR, CAT, f"{CAT}_sensitivity_process.csv")) + obs_dates, obs = load_obs() + + if prod is None: + raise FileNotFoundError( + f"Missing {PROD_DIR}/{CAT}/{CAT}_production_per_member.csv — " + "run run_production_per_member.py first.") + if init is None or forc is None or proc is None: + raise FileNotFoundError( + f"Missing one of the sensitivity CSVs under {SEN_DIR}/{CAT}/. " + "Run run_perturbation_sensitivity.py for all three sources first.") + + def to_mask(dates): + d = pd.to_datetime(dates) + return d, (d >= ZOOM_START) & (d <= ZOOM_END) + + pd_dates, pd_mask = to_mask(prod_dates) + id_dates, id_mask = to_mask(init_dates) + fd_dates, fd_mask = to_mask(forc_dates) + cd_dates, cd_mask = to_mask(proc_dates) + if obs_dates is not None: + od, om = to_mask(obs_dates) + else: + od, om = None, None + + N = prod.shape[1] + # Lay out 4 cols x 5 rows = 20 panels + n_cols = 4 + n_rows = (N + n_cols - 1) // n_cols + fig, axes = plt.subplots(n_rows, n_cols, figsize=(22, 4 * n_rows), sharex=True) + axes = axes.flatten() + + for i in range(N): + ax = axes[i] + col = f"member_{i:02d}" + + # Each member's INIT-only trajectory (thin red) + if col in init.columns: + ax.plot(id_dates[id_mask], init[col].values[id_mask], + color=INIT_COLOR, lw=0.9, alpha=0.9, label="Init only") + # FORCING-only (thin blue) + if col in forc.columns: + ax.plot(fd_dates[fd_mask], forc[col].values[fd_mask], + color=FORCING_COLOR,lw=0.9, alpha=0.9, label="Forcing only") + # PROCESS-only (thin green) + if col in proc.columns: + ax.plot(cd_dates[cd_mask], proc[col].values[cd_mask], + color=PROC_COLOR, lw=0.9, alpha=0.9, label="Process noise only") + # PRODUCTION (thick purple) — actual forecast with DA on, all perturbations + if col in prod.columns: + ax.plot(pd_dates[pd_mask], prod[col].values[pd_mask], + color=PROD_COLOR, lw=2.2, alpha=0.95, label="Production (DA on)") + # Qkrig observation (thick black dashed) + if od is not None: + ax.plot(od[om], obs[om], + color=OBS_COLOR, lw=1.4, linestyle="--", alpha=0.9, label="Qkrig (obs)") + + ax.set_title(f"{col}", fontsize=10) + ax.tick_params(axis="x", rotation=45, labelsize=7) + ax.tick_params(axis="y", labelsize=8) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=1)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + if i == 0: + ax.legend(fontsize=7, loc="upper left") + + # Hide any extra blank axes + for j in range(N, len(axes)): + axes[j].axis("off") + + fig.suptitle( + f"Per-member factor decomposition - {CAT} - Hurricane Helene peak (Sep 24-28, 2024)\n" + "Purple thick = production member (DA on, all perturbations) | " + "Red = init only | Blue = forcing only | Green = process noise only | " + "Black dashed = Qkrig obs", + fontsize=12, y=1.00, + ) + fig.text(0.005, 0.5, "Discharge (mm/h)", va="center", rotation="vertical", fontsize=12) + plt.tight_layout() + plt.savefig(OUT_PNG, dpi=140, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_per_member_input_output_diagnostic.py b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_per_member_input_output_diagnostic.py new file mode 100644 index 00000000..f6be63d8 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_per_member_input_output_diagnostic.py @@ -0,0 +1,177 @@ +""" +Per-member input/output diagnostic for ONE catchment. + +For each ensemble member, show a single combined panel (hydrograph-style) that +traces the inputs the member actually received and the output it produced: + + Each member panel uses an internal 2-row stack: + TOP sub-panel: perturbed precip (bars, left axis) + + perturbed PET (line, right twinx axis) + BOTTOM sub-panel: simulated streamflow Q (purple) + + Qkrig observation (black dashed) + Shared x-axis between the two sub-panels. + + Initial states for that member are shown in the title line. + +Layout: 5 rows x 4 columns = 20 member panels. + +Inputs (all from run_production_per_member.py): + //_production_per_member.csv (Q per member) + //_production_per_member_precip.csv (precip per member) + //_production_per_member_pet.csv (PET per member) + //_production_per_member_initial_states.json + //_test_results.csv (Qkrig obs) + +Output: + //_per_member_io_diagnostic.png +""" +import os +import json +import pandas as pd +import numpy as np +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import matplotlib.gridspec as gridspec + +# ----- Configuration ----- +CAT = "cat-1016300" + +PROD_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_production_per_member" +DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" +OUT_PNG = os.path.join(PROD_DIR, CAT, f"{CAT}_per_member_io_diagnostic.png") + +# Plot window (Helene 4-day zoom) +ZOOM_START = pd.Timestamp("2024-09-24") +ZOOM_END = pd.Timestamp("2024-09-28") + +PRECIP_COLOR = "tab:blue" +PET_COLOR = "tab:orange" +Q_COLOR = "tab:purple" +OBS_COLOR = "black" + +N_ROWS = 5 # grid rows of member panels +N_COLS = 4 # grid cols of member panels + + +def load_member_csv(path): + if not os.path.exists(path): + return None, None + df = pd.read_csv(path, parse_dates=["date"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + return df["date"].values, df[member_cols] + + +def main(): + q_dates, q_df = load_member_csv(os.path.join(PROD_DIR, CAT, f"{CAT}_production_per_member.csv")) + precip_dates, precip_df = load_member_csv(os.path.join(PROD_DIR, CAT, f"{CAT}_production_per_member_precip.csv")) + pet_dates, pet_df = load_member_csv(os.path.join(PROD_DIR, CAT, f"{CAT}_production_per_member_pet.csv")) + + init_states_path = os.path.join(PROD_DIR, CAT, f"{CAT}_production_per_member_initial_states.json") + if os.path.exists(init_states_path): + with open(init_states_path) as f: + init_states = json.load(f)["members"] + else: + init_states = {} + + obs_path = os.path.join(DA_DIR, CAT, f"{CAT}_test_results.csv") + if os.path.exists(obs_path): + obs_df = pd.read_csv(obs_path, parse_dates=["date"]) + obs_dates = pd.to_datetime(obs_df["date"].values) + obs_vals = obs_df["obs_mm_h"].values + else: + obs_dates = None + obs_vals = None + + if q_df is None or precip_df is None or pet_df is None: + raise FileNotFoundError( + "Missing one of the per-member CSVs. Run run_production_per_member.py " + "with the updated script that saves precip + PET + initial states.") + + member_cols = list(q_df.columns) + N = len(member_cols) + assert N <= N_ROWS * N_COLS, f"Grid too small for {N} members" + + qd = pd.to_datetime(q_dates) + pd_ = pd.to_datetime(precip_dates) + ed = pd.to_datetime(pet_dates) + qmask = (qd >= ZOOM_START) & (qd <= ZOOM_END) + pmask = (pd_ >= ZOOM_START) & (pd_ <= ZOOM_END) + emask = (ed >= ZOOM_START) & (ed <= ZOOM_END) + if obs_dates is not None: + omask = (obs_dates >= ZOOM_START) & (obs_dates <= ZOOM_END) + + # Compute global y-axis ranges so all member panels are directly comparable + p_ymax = float(np.nanmax(precip_df.values[pmask, :])) * 1.10 + e_ymax = float(np.nanmax(pet_df.values[emask, :])) * 1.10 + q_data_max = float(np.nanmax(q_df.values[qmask, :])) + if obs_dates is not None: + q_data_max = max(q_data_max, float(np.nanmax(obs_vals[omask]))) + q_ymax = q_data_max * 1.10 + + fig = plt.figure(figsize=(22, 19)) + outer = gridspec.GridSpec(N_ROWS, N_COLS, hspace=0.75, wspace=0.30) + + for i in range(N): + col = member_cols[i] + row, c = divmod(i, N_COLS) + # Each member cell is a 2-row inner grid (precip+PET on top, Q on bottom) + inner = gridspec.GridSpecFromSubplotSpec( + 2, 1, subplot_spec=outer[row, c], hspace=0.20, height_ratios=[1.0, 1.6]) + + # --- Top sub-panel: precip bars + PET line on twinx --- + ax_top = fig.add_subplot(inner[0]) + ax_top.bar(pd_[pmask], precip_df[col].values[pmask], + width=0.04, color=PRECIP_COLOR, alpha=0.85, label="P") + ax_top.set_ylim(0, p_ymax) + ax_top.set_ylabel("P (mm/h)", fontsize=8, color=PRECIP_COLOR) + ax_top.tick_params(axis="y", labelsize=7, labelcolor=PRECIP_COLOR) + ax_top.tick_params(axis="x", which="both", bottom=False, labelbottom=False) + ax_top.grid(True, alpha=0.15) + # PET on twinx (different scale) + ax_pet = ax_top.twinx() + ax_pet.plot(ed[emask], pet_df[col].values[emask], + color=PET_COLOR, lw=1.0, label="PET") + ax_pet.set_ylim(0, e_ymax) + ax_pet.set_ylabel("PET (mm/h)", fontsize=8, color=PET_COLOR) + ax_pet.tick_params(axis="y", labelsize=7, labelcolor=PET_COLOR) + # Title: member name on line 1, initial states on line 2 (avoids overlap with neighbors) + s = init_states.get(col, {}) + title = (f"{col}\n" + f"soil={s.get('soil_m', float('nan')):.3f} m | " + f"GW={s.get('gw_m', float('nan')):.4f} m | " + f"Nash[0]={s.get('nash0_m', float('nan')):.1e} " + f"Nash[1]={s.get('nash1_m', float('nan')):.1e}") + ax_top.set_title(title, fontsize=8, loc="left") + + # --- Bottom sub-panel: Q + obs (shares x-axis with top) --- + ax_bot = fig.add_subplot(inner[1], sharex=ax_top) + ax_bot.plot(qd[qmask], q_df[col].values[qmask], + color=Q_COLOR, lw=1.6, label=col) + if obs_dates is not None: + ax_bot.plot(obs_dates[omask], obs_vals[omask], + color=OBS_COLOR, lw=1.2, linestyle="--", label="Qkrig (obs)") + ax_bot.set_ylim(0, q_ymax) + ax_bot.set_ylabel("Q (mm/h)", fontsize=8) + ax_bot.tick_params(labelsize=7) + ax_bot.xaxis.set_major_locator(mdates.DayLocator(interval=1)) + ax_bot.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax_bot.grid(True, alpha=0.15) + if i == 0: + ax_bot.legend(fontsize=7, loc="upper left", ncol=2, + frameon=True, framealpha=0.85) + + fig.suptitle( + f"Per-member input/output diagnostic - {CAT} - " + f"Hurricane Helene peak (Sep 24-28, 2024)\n" + "Each member panel: TOP = perturbed precip (blue bars) + perturbed PET (orange line) | " + "BOTTOM = simulated Q (purple) vs Qkrig obs (black dashed). " + "Title shows initial states at t=0.", + fontsize=12, y=0.995, + ) + plt.savefig(OUT_PNG, dpi=140, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_perturbation_category_shaded.py b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_perturbation_category_shaded.py new file mode 100644 index 00000000..69c45ec3 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_perturbation_category_shaded.py @@ -0,0 +1,176 @@ +""" +Per-catchment shaded ensemble-band plot, organized by perturbation category. + +Three categories: + 1. Initial states (red) + 2. Meteorological forcings (blue) + 3. Hydrological states (green) + +For each category, all 20 ensemble members are shown as a shaded band +(min-max envelope fill) plus a thicker median line in the same color. +Qkrig observation overlaid in black. Hurricane Helene peak window shaded +in pink. Styled after a standard ensemble-forecast figure layout. + +Inputs (existing per-source sensitivity CSVs from run_perturbation_sensitivity.py): + //_sensitivity_init.csv (20 members, init only) + //_sensitivity_forcing.csv (20 members, forcing only) + //_sensitivity_process.csv (20 members, process noise only) + //_test_results.csv (Qkrig obs) + +Outputs: + //_perturbation_categories_linear.png + //_perturbation_categories_log.png +""" +import argparse +import os +import pandas as pd +import numpy as np +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import matplotlib.patches as mpatches + +CAT = "cat-1016300" + +SEN_DIR = "/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt" +OBS_DIR = "/home/svyas/catchment_ts_no_03463300_gapfilled" + +OUT_LINEAR = os.path.join(SEN_DIR, CAT, f"{CAT}_perturbation_categories_linear.png") +OUT_LOG = os.path.join(SEN_DIR, CAT, f"{CAT}_perturbation_categories_log.png") + +# Plot window — wider context, similar to the paper's Sep 10 - Oct 08 +PLOT_START = pd.Timestamp("2024-09-20") +PLOT_END = pd.Timestamp("2024-10-05") + +# Helene peak band +HELENE_START = pd.Timestamp("2024-09-26 12:00:00") +HELENE_END = pd.Timestamp("2024-09-28 00:00:00") + +# Category configuration: file suffix, display label, color +CATEGORIES = [ + ("init", "Initial states only", "tab:red"), + ("forcing", "Meteorological forcings only", "tab:blue"), + ("process", "Hydrological states only", "tab:green"), +] + + +def load_members(source): + path = os.path.join(SEN_DIR, CAT, f"{CAT}_sensitivity_{source}.csv") + if not os.path.exists(path): + return None, None + df = pd.read_csv(path, parse_dates=["date"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + return df["date"].values, df[member_cols].to_numpy(dtype=float) + + +def load_obs(): + p = os.path.join(OBS_DIR, f"{CAT}.csv") + if not os.path.exists(p): + return None, None + df = pd.read_csv(p) + time_col = 'datetime' if 'datetime' in df.columns else 'date' + df[time_col] = pd.to_datetime(df[time_col]) + return df[time_col].values, df["qkrig"].values + + +def plot_panel(ax, obs_dates, obs_vals, log_y=False): + handles_labels = [] # for the legend + + # Plot each category as a shaded band + median line + for source, label, color in CATEGORIES: + dates, q = load_members(source) + if dates is None: + continue + d = pd.to_datetime(dates) + mask = (d >= PLOT_START) & (d <= PLOT_END) + if mask.sum() == 0: + continue + # Min-max envelope across all 20 members per timestep (widest possible band). + # Bands are visually narrow even with min/max because perturbations are + # tuned for production EnKF stability, not for max visible spread. + q_window = q[mask, :] + qmin = np.nanmin(q_window, axis=1) + qmax = np.nanmax(q_window, axis=1) + median = np.nanmedian(q_window, axis=1) + + ax.fill_between(d[mask], qmin, qmax, + color=color, alpha=0.30, zorder=2, + edgecolor="none") + line, = ax.plot(d[mask], median, + color=color, lw=1.7, alpha=0.95, zorder=3, + label=f"{label} (N=20)") + handles_labels.append((line, label)) + + # Helene peak shaded band (vertical) + ax.axvspan(HELENE_START, HELENE_END, + color="salmon", alpha=0.15, zorder=1) + ax.text((HELENE_START + (HELENE_END - HELENE_START) / 2), + ax.get_ylim()[1] if not log_y else 1.0, + "Helene peak", + fontsize=9, color="salmon", + ha="center", va="bottom", zorder=3) + + # Observation + if obs_dates is not None: + od = pd.to_datetime(obs_dates) + om = (od >= PLOT_START) & (od <= PLOT_END) + ax.plot(od[om], obs_vals[om], + color="black", lw=1.4, label="Qkrig (obs)", zorder=4) + + ax.set_xlabel("Date", fontsize=10) + if log_y: + ax.set_yscale("log") + ax.set_ylabel(r"q (mm hr$^{-1}$) [log scale]", fontsize=10) + ax.set_ylim(0.01, None) + else: + ax.set_ylabel("Discharge (mm/h)", fontsize=10) + ax.tick_params(axis="x", rotation=30, labelsize=9) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=2)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.grid(True, alpha=0.2) + ax.legend(loc="upper left", fontsize=9, frameon=True, framealpha=0.9) + + +def main(): + global CAT, OUT_LINEAR, OUT_LOG + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', default=CAT) + args = parser.parse_args() + CAT = args.cat_id + OUT_LINEAR = os.path.join(SEN_DIR, CAT, f"{CAT}_perturbation_categories_linear.png") + OUT_LOG = os.path.join(SEN_DIR, CAT, f"{CAT}_perturbation_categories_log.png") + + obs_dates, obs_vals = load_obs() + + # ----- Linear-y ----- + fig, ax = plt.subplots(1, 1, figsize=(20, 7)) + plot_panel(ax, obs_dates, obs_vals, log_y=False) + fig.suptitle( + f"Ensemble forecast by perturbation category - {CAT} - " + f"Sep 20 - Oct 5, 2024 (Hurricane Helene window)\n" + "Shaded bands = min-max envelope across 20 members. " + "Lines = ensemble median. DA off in all three sub-experiments.", + fontsize=11, y=0.995, + ) + plt.tight_layout() + plt.savefig(OUT_LINEAR, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_LINEAR}") + + # ----- Log-y (paper style) ----- + fig, ax = plt.subplots(1, 1, figsize=(20, 7)) + plot_panel(ax, obs_dates, obs_vals, log_y=True) + fig.suptitle( + f"Ensemble forecast by perturbation category - {CAT} - log-scale q - " + f"Sep 20 - Oct 5, 2024 (Hurricane Helene window)\n" + "Shaded bands = min-max envelope across 20 members. " + "Lines = ensemble median. DA off in all three sub-experiments.", + fontsize=11, y=0.995, + ) + plt.tight_layout() + plt.savefig(OUT_LOG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_LOG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_production_ensemble_forecast.py b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_production_ensemble_forecast.py new file mode 100644 index 00000000..1999cad3 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_production_ensemble_forecast.py @@ -0,0 +1,179 @@ +""" +Single-panel ensemble forecast plot for one catchment, paper-style. + +Plots all 20 production-per-member streamflow trajectories overlaid on a +single time-series panel, with the Hurricane Helene window highlighted by a +pink shaded vertical band and the Qkrig observation drawn on top in black. + +Ensemble forecast figure style (50/100/200 km +variogram-range ensemble panels): every member as a thin colored line, storm +window shaded, observation as a thick dark series, optional log-scaled y-axis. + +Inputs: + //_production_per_member.csv (date + 20 members) + //_test_results.csv (Qkrig obs) + +Outputs: + //_production_ensemble_forecast_linear.png + //_production_ensemble_forecast_log.png +""" + +import argparse +import os +import pandas as pd +import numpy as np +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import matplotlib.patches as mpatches +from matplotlib import cm + +# ----- Configuration ----- +CAT = "cat-1016300" + +PROD_DIR = "/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt" +DA_DIR = "/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt" +OBS_DIR = "/home/svyas/catchment_ts_no_03463300_gapfilled" + +# Plot window: a few weeks around Helene so the storm is in context +PLOT_START = pd.Timestamp("2024-09-20") +PLOT_END = pd.Timestamp("2024-10-05") + +# Helene peak window highlighted in pink +HELENE_START = pd.Timestamp("2024-09-26 12:00:00") +HELENE_END = pd.Timestamp("2024-09-28 00:00:00") + +OUT_LINEAR = os.path.join(PROD_DIR, CAT, f"{CAT}_production_ensemble_forecast_linear.png") +OUT_LOG = os.path.join(PROD_DIR, CAT, f"{CAT}_production_ensemble_forecast_log.png") + + +def kge_score(obs, sim): + m = ~(np.isnan(obs) | np.isnan(sim)) + o, s = obs[m], sim[m] + if len(o) < 2: + return float("nan") + denom = float(np.sqrt(((o - o.mean()) ** 2).sum() * ((s - s.mean()) ** 2).sum())) + if denom == 0: + return float("nan") + r = float(((o - o.mean()) * (s - s.mean())).sum() / denom) + alpha = float(s.std() / o.std()) if o.std() != 0 else float("nan") + beta = float(s.mean() / o.mean()) if o.mean() != 0 else float("nan") + return 1.0 - float(np.sqrt((r - 1) ** 2 + (alpha - 1) ** 2 + (beta - 1) ** 2)) + + +def main(): + global CAT, OUT_LINEAR, OUT_LOG + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', default=CAT) + args = parser.parse_args() + CAT = args.cat_id + OUT_LINEAR = os.path.join(PROD_DIR, CAT, f"{CAT}_production_ensemble_forecast_linear.png") + OUT_LOG = os.path.join(PROD_DIR, CAT, f"{CAT}_production_ensemble_forecast_log.png") + + prod_path = os.path.join(PROD_DIR, CAT, f"{CAT}_production_per_member.csv") + if not os.path.exists(prod_path): + raise FileNotFoundError( + f"Missing {prod_path}. Run run_production_per_member.py for {CAT} first.") + + prod_df = pd.read_csv(prod_path, parse_dates=["date"]) + member_cols = sorted([c for c in prod_df.columns if c.startswith("member_")]) + member_arr = prod_df[member_cols].to_numpy(dtype=float) + dates = pd.to_datetime(prod_df["date"].values) + + obs_dates = None + obs_vals = None + obs_krig = os.path.join(OBS_DIR, f"{CAT}.csv") + obs_tr = os.path.join(DA_DIR, CAT, f"{CAT}_test_results.csv") + if os.path.exists(obs_krig): + df = pd.read_csv(obs_krig) + time_col = 'datetime' if 'datetime' in df.columns else 'date' + df[time_col] = pd.to_datetime(df[time_col]) + obs_dates = df[time_col].values + obs_vals = df["qkrig"].values + elif os.path.exists(obs_tr): + obs_df = pd.read_csv(obs_tr, parse_dates=["date"]) + obs_dates = pd.to_datetime(obs_df["date"].values) + obs_vals = obs_df["obs_mm_h"].values + else: + print(f"Note: no obs found for {CAT}; plotting without observation overlay.") + + # Compute per-member peak (within plot window) for the legend label + pw_mask = (dates >= PLOT_START) & (dates <= PLOT_END) + peaks = member_arr[pw_mask, :].max(axis=0) + + # Optional: compute per-member KGE vs obs across the plot window + member_kges = np.full(len(member_cols), np.nan) + if obs_dates is not None: + obs_series = pd.Series(obs_vals, index=obs_dates) + for i, _ in enumerate(member_cols): + member_series = pd.Series(member_arr[:, i], index=dates) + joined = pd.concat([obs_series, member_series], axis=1, join="inner").dropna() + if len(joined) >= 2: + member_kges[i] = kge_score(joined.iloc[:, 0].values, joined.iloc[:, 1].values) + + def plot_panel(ax, log_y=False): + colormap = cm.get_cmap("turbo", len(member_cols)) + # Thin colored lines per member + for i, col in enumerate(member_cols): + label = f"{col} | peak={peaks[i]:.1f} mm/h" + ax.plot(dates[pw_mask], member_arr[pw_mask, i], + color=colormap(i), lw=0.8, alpha=0.85, label=label, zorder=2) + + # Helene peak shaded band + ax.axvspan(HELENE_START, HELENE_END, color="salmon", alpha=0.18, zorder=1) + ax.text(HELENE_START + (HELENE_END - HELENE_START) / 2, + ax.get_ylim()[1] if not log_y else 1.0, + "Helene peak", fontsize=8, color="salmon", + ha="center", va="bottom", zorder=3) + + # Observation as thick black series + if obs_dates is not None: + obs_mask = (obs_dates >= PLOT_START) & (obs_dates <= PLOT_END) + ax.plot(obs_dates[obs_mask], obs_vals[obs_mask], + color="black", lw=1.6, label="Qkrig (obs)", zorder=4) + + ax.set_xlabel("Date", fontsize=10) + if log_y: + ax.set_yscale("log") + ax.set_ylabel(r"q (mm hr$^{-1}$) [log scale]", fontsize=10) + ax.set_ylim(0.01, max(member_arr.max(), (obs_vals.max() if obs_vals is not None else 0)) * 1.2) + else: + ax.set_ylabel("Discharge (mm/h)", fontsize=10) + + ax.tick_params(axis="x", rotation=30, labelsize=9) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=2)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.grid(True, alpha=0.2) + + # ---- LINEAR Y figure ---- + fig, ax = plt.subplots(1, 1, figsize=(20, 7)) + plot_panel(ax, log_y=False) + ax.legend(fontsize=6.5, loc="upper left", ncol=2, + frameon=True, framealpha=0.85, markerfirst=False) + fig.suptitle( + f"Production ensemble forecast (N=20) at {CAT} - " + f"Sep 20 - Oct 5, 2024 (Hurricane Helene window)", + fontsize=12, y=0.99, + ) + plt.tight_layout() + plt.savefig(OUT_LINEAR, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_LINEAR}") + + # ---- LOG Y figure (matches paper style) ---- + fig, ax = plt.subplots(1, 1, figsize=(20, 7)) + plot_panel(ax, log_y=True) + ax.legend(fontsize=6.5, loc="upper left", ncol=2, + frameon=True, framealpha=0.85, markerfirst=False) + fig.suptitle( + f"Production ensemble forecast (N=20) at {CAT} - log-scale q - " + f"Sep 20 - Oct 5, 2024 (Hurricane Helene window)", + fontsize=12, y=0.99, + ) + plt.tight_layout() + plt.savefig(OUT_LOG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_LOG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_combined.py b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_combined.py new file mode 100644 index 00000000..c6d531a0 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_combined.py @@ -0,0 +1,235 @@ +#!/usr/bin/env python3 +""" +plot_routed_ensemble_combined.py +Combined poster plot: 600-member routed ensemble envelope + +deterministic DA lines (Vrugt/no-Vrugt) vs USGS gauge obs at 03463300. + +Inputs: + --vrugt-csv : routed_Q_test.csv from vrugt_dynamic_routed/ + --novrugt-csv : routed_Q_test.csv from novrugt_dynamic_routed/ + --ensemble-pq : routed_crossed_ensemble.parquet from run_route_crossed_ensemble.py + --out-dir : output directory + +The ensemble parquet columns: issue_time, lead_hour, member_0000..member_0599 (Q in m3/s) + +Usage: + python3 plot_routed_ensemble_combined.py \ + --vrugt-csv /mnt/disk2/suma_helen_poster/da_results/vrugt_dynamic_routed/routed_Q_test.csv \ + --novrugt-csv /mnt/disk2/suma_helen_poster/da_results/novrugt_dynamic_routed/routed_Q_test.csv \ + --ensemble-pq /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble_routed/routed_crossed_ensemble.parquet \ + --out-dir /mnt/disk2/suma_helen_poster/da_results/comparison_plots +""" + +import argparse +import os + +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +HELENE_START = pd.Timestamp("2024-09-25 00:00:00") +HELENE_END = pd.Timestamp("2024-09-29 00:00:00") +ZOOM_START = pd.Timestamp("2024-09-24 00:00:00") +ZOOM_END = pd.Timestamp("2024-09-28 18:00:00") # cut trailing ensemble bump + +COLOR_OBS = "black" +COLOR_VRUGT = "#1f77b4" # blue +COLOR_NOVRUGT = "#d62728" # red +COLOR_ENS = "#1f77b4" # same blue family as Vrugt + + +def load_det_csv(path): + df = pd.read_csv(path, parse_dates=["date"]) + return df.set_index("date").sort_index() + + +def build_envelope(pq_path, zoom_start, zoom_end): + """ + Load routed ensemble parquet and build percentile envelope by valid_time. + Returns DataFrame indexed by valid_time with columns p05, p25, p50, p75, p95. + """ + df = pd.read_parquet(pq_path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = [c for c in df.columns if c.startswith("member_")] + + records = [] + for _, row in df.iterrows(): + t0 = row["issue_time"] + lead = int(row["lead_hour"]) + vt = t0 + pd.Timedelta(hours=lead) + vals = row[member_cols].values.astype(float) + records.append({"valid_time": vt, "vals": vals}) + + # Group by valid_time (floor to hour), stack all member values across issue_times + from collections import defaultdict + groups = defaultdict(list) + for rec in records: + vt = rec["valid_time"].floor("1h") + groups[vt].extend(rec["vals"].tolist()) + + rows = [] + for vt in sorted(groups): + if vt < zoom_start or vt > zoom_end: + continue + v = np.array(groups[vt]) + v = v[~np.isnan(v)] + if len(v) == 0: + continue + rows.append({ + "valid_time": vt, + "p05": np.percentile(v, 5), + "p25": np.percentile(v, 25), + "p50": np.percentile(v, 50), + "p75": np.percentile(v, 75), + "p95": np.percentile(v, 95), + }) + + return pd.DataFrame(rows).set_index("valid_time") + + +def compute_kge(obs, sim): + mask = ~(np.isnan(obs) | np.isnan(sim)) + o, s = obs[mask], sim[mask] + if len(o) < 2 or np.std(o) == 0: + return np.nan + r = np.corrcoef(o, s)[0, 1] + return 1.0 - np.sqrt((r-1)**2 + (np.std(s)/np.std(o)-1)**2 + (np.mean(s)/np.mean(o)-1)**2) + + +def _format_xaxis(ax, locator, fmt, rotation=30): + ax.xaxis.set_major_locator(locator) + ax.xaxis.set_major_formatter(mdates.DateFormatter(fmt)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=rotation, ha="right", fontsize=9) + + +def _add_helene_band(ax): + ax.axvspan(HELENE_START, HELENE_END, color="gold", alpha=0.15, zorder=0) + + +def plot_zoom(ax, dates, obs, sim_v, sim_nv, kge_v, kge_nv, env=None): + mask = (dates >= ZOOM_START) & (dates <= ZOOM_END) + dz = dates[mask] + + # Ensemble envelope (bottom layer) + if env is not None: + ax.fill_between(env.index, env["p05"], env["p95"], + color=COLOR_ENS, alpha=0.18, label="Ensemble 5-95th pct") + ax.fill_between(env.index, env["p25"], env["p75"], + color=COLOR_ENS, alpha=0.30, label="Ensemble 25-75th pct") + ax.plot(env.index, env["p50"], + color=COLOR_ENS, lw=1.2, linestyle="--", alpha=0.7, label="Ensemble median") + + # Deterministic DA lines + ax.plot(dz, sim_nv[mask], color=COLOR_NOVRUGT, lw=1.4, linestyle="--", + label=f"DA -- constant R KGE={kge_nv:.3f}", zorder=3) + ax.plot(dz, sim_v[mask], color=COLOR_VRUGT, lw=1.6, + label=f"DA -- dynamic R (Vrugt) KGE={kge_v:.3f}", zorder=4) + + # USGS obs on top + ax.plot(dz, obs[mask], color=COLOR_OBS, lw=1.6, label="USGS obs (gauge 03463300)", zorder=5) + + peak_usgs = np.nanmax(obs[mask]) + ax.axhline(peak_usgs, color=COLOR_OBS, lw=0.7, linestyle=":", alpha=0.5) + # Place annotation to the right of the peak to avoid legend overlap + ax.text(pd.Timestamp("2024-09-28 06:00:00"), peak_usgs * 1.02, + f"USGS peak {peak_usgs:.0f} m³/s", fontsize=8.5, color="black", + alpha=0.8, ha="right") + + _add_helene_band(ax) + _format_xaxis(ax, mdates.DayLocator(interval=1), "%b %d") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--vrugt-csv", required=True) + parser.add_argument("--novrugt-csv", required=True) + parser.add_argument("--ensemble-pq", required=True, + help="routed_crossed_ensemble.parquet from run_route_crossed_ensemble.py") + parser.add_argument("--out-dir", required=True) + args = parser.parse_args() + os.makedirs(args.out_dir, exist_ok=True) + + vrugt = load_det_csv(args.vrugt_csv) + novrugt = load_det_csv(args.novrugt_csv) + + obs = vrugt["Q_usgs_m3s"].values + dates = vrugt.index + sim_v = vrugt["Q_routed_m3s"].values + sim_nv = novrugt["Q_routed_m3s"].reindex(dates).values + + kge_v = compute_kge(obs, sim_v) + kge_nv = compute_kge(obs, sim_nv) + + print("Building ensemble envelope...") + env = build_envelope(args.ensemble_pq, ZOOM_START, ZOOM_END) + print(f" Envelope: {len(env)} timesteps, " + f"peak p95={env['p95'].max():.1f} m3/s, peak p50={env['p50'].max():.1f} m3/s") + + peak_usgs = np.nanmax(obs[(dates >= ZOOM_START) & (dates <= ZOOM_END)]) + in_env = env["p95"].max() >= peak_usgs + print(f" USGS peak {peak_usgs:.1f} m3/s {'IS' if in_env else 'IS NOT'} within p95 envelope") + + # ------------------------------------------------------------------ # + # Figure 1: Helene zoom only -- main poster panel + # ------------------------------------------------------------------ # + fig, ax = plt.subplots(figsize=(11, 5.5)) + plot_zoom(ax, dates, obs, sim_v, sim_nv, kge_v, kge_nv, env) + ax.set_xlabel("Date (UTC)", fontsize=11) + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_title( + "Hurricane Helene -- Probabilistic discharge forecast at gauge 03463300 (South Toe River)\n" + "600-member CFE ensemble + DA (Muskingum routing)", + fontsize=11) + ax.legend(fontsize=9, loc="lower right") + ax.grid(True, alpha=0.25) + plt.tight_layout() + out1 = os.path.join(args.out_dir, "helene_ensemble_vs_usgs.png") + plt.savefig(out1, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + # ------------------------------------------------------------------ # + # Figure 2: two-panel (full period + Helene zoom) + # ------------------------------------------------------------------ # + fig, (ax_top, ax_bot) = plt.subplots( + 2, 1, figsize=(14, 9), + gridspec_kw={"height_ratios": [1, 1.5]}) + + # Top: full test period (deterministic only, no ensemble for readability) + ax_top.plot(dates, obs, color=COLOR_OBS, lw=0.8, label="USGS obs", zorder=3) + ax_top.plot(dates, sim_v, color=COLOR_VRUGT, lw=0.9, + label=f"DA dynamic R (Vrugt) KGE={kge_v:.3f}", zorder=2) + ax_top.plot(dates, sim_nv, color=COLOR_NOVRUGT, lw=0.9, linestyle="--", + label=f"DA constant R KGE={kge_nv:.3f}", zorder=2) + ax_top.axvspan(ZOOM_START, ZOOM_END, color="gray", alpha=0.12, zorder=0) + ax_top.axvspan(HELENE_START, HELENE_END, color="gold", alpha=0.15, zorder=0) + ax_top.text(HELENE_START + pd.Timedelta(days=0.5), + np.nanmax(obs) * 0.88, "Helene", + fontsize=9, color="goldenrod", fontweight="bold") + ax_top.set_ylabel("Discharge (m³/s)", fontsize=10) + ax_top.set_title( + "Routed discharge at gauge 03463300 (South Toe River) -- Test period\n" + "DA with dynamic Vrugt R vs. constant R -- Muskingum routing", + fontsize=11) + ax_top.legend(fontsize=9, loc="upper left") + ax_top.grid(True, alpha=0.25) + _format_xaxis(ax_top, mdates.MonthLocator(interval=2), "%Y-%m") + + # Bottom: Helene zoom with ensemble + plot_zoom(ax_bot, dates, obs, sim_v, sim_nv, kge_v, kge_nv, env) + ax_bot.set_xlabel("Date (UTC)", fontsize=10) + ax_bot.set_ylabel("Discharge (m³/s)", fontsize=10) + ax_bot.set_title("Hurricane Helene window -- 600-member ensemble envelope", fontsize=10) + ax_bot.legend(fontsize=9, loc="lower right") + ax_bot.grid(True, alpha=0.25) + + plt.tight_layout() + out2 = os.path.join(args.out_dir, "helene_ensemble_twopanel.png") + plt.savefig(out2, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_vs_usgs.py b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_vs_usgs.py new file mode 100644 index 00000000..55c52d85 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_vs_usgs.py @@ -0,0 +1,176 @@ +#!/usr/bin/env python3 +""" +plot_routed_ensemble_vs_usgs.py +Plot T-route routed 600-member ensemble envelope vs USGS at gauge 03463300. + +Input: routed_crossed_ensemble.parquet from run_route_crossed_ensemble.py + Columns: issue_time, lead_hour, member_0000..member_0599 (Q in m3/s) + +Usage: + python3 plot_routed_ensemble_vs_usgs.py \ + --routed-pq /mnt/disk2/.../folder1_vrugt_routed/routed_crossed_ensemble.parquet \ + --usgs-csv /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv \ + --label "F1 Vrugt — 1 gauge holdout" \ + --out-dir ~/plots_f1_1gauge +""" + +import argparse +from collections import defaultdict +from pathlib import Path + +import matplotlib.dates as mdates +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd + +WATERSHED_AREA_KM2 = 113.18 # cat-1016300 drainage area +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +PLOT_START = pd.Timestamp("2024-09-24 00:00:00") +PLOT_END = pd.Timestamp("2024-09-28 18:00:00") +PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +PEAK_END = pd.Timestamp("2024-09-28 06:00:00") +HELENE_PEAK = pd.Timestamp("2024-09-27 12:00:00") + +COLOR_ENS = "#1f77b4" +COLOR_OBS = "black" + + +def build_envelope(pq_path, t_start, t_end): + df = pd.read_parquet(pq_path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = [c for c in df.columns if c.startswith("member_")] + + groups = defaultdict(list) + for _, row in df.iterrows(): + vt = (row["issue_time"] + pd.Timedelta(hours=int(row["lead_hour"]))).floor("1h") + if t_start <= vt <= t_end: + groups[vt].extend(row[member_cols].values.astype(float).tolist()) + + rows = [] + for vt in sorted(groups): + v = np.array(groups[vt]) + v = v[~np.isnan(v)] + if len(v) == 0: + continue + rows.append({"valid_time": vt, + "p05": np.percentile(v, 5), "p25": np.percentile(v, 25), + "p50": np.percentile(v, 50), "p75": np.percentile(v, 75), + "p95": np.percentile(v, 95)}) + return pd.DataFrame(rows).set_index("valid_time") + + +def load_usgs(path): + df = pd.read_csv(path) + print(f" USGS CSV columns: {list(df.columns)}") + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() or "discharge" in c.lower()) + print(f" Using date_col={date_col!r}, q_col={q_col!r}") + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series.sort_index() + + +def compute_kge(obs, sim): + mask = ~(np.isnan(obs) | np.isnan(sim)) + o, s = obs[mask], sim[mask] + if len(o) < 2 or np.std(o) == 0: + return np.nan + r = np.corrcoef(o, s)[0, 1] + return 1 - np.sqrt((r-1)**2 + (np.std(s)/np.std(o)-1)**2 + (np.mean(s)/np.mean(o)-1)**2) + + +def _shade(ax, env): + ax.fill_between(env.index, env["p05"], env["p95"], + color=COLOR_ENS, alpha=0.18, label="5th–95th percentile") + ax.fill_between(env.index, env["p25"], env["p75"], + color=COLOR_ENS, alpha=0.32, label="25th–75th percentile") + ax.plot(env.index, env["p50"], color=COLOR_ENS, lw=1.5, label="Ensemble median (600 members)") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--routed-pq", required=True) + parser.add_argument("--usgs-csv", required=True) + parser.add_argument("--label", default="F1 Vrugt — 1 gauge holdout") + parser.add_argument("--out-dir", required=True) + args = parser.parse_args() + + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + print("Building full-window envelope...") + env_full = build_envelope(args.routed_pq, PLOT_START, PLOT_END) + print("Building peak-window envelope...") + env_peak = build_envelope(args.routed_pq, PEAK_START, PEAK_END) + + print("Loading USGS obs...") + try: + obs_all = load_usgs(args.usgs_csv) + obs_full = obs_all.loc[PLOT_START:PLOT_END] + obs_peak = obs_all.loc[PEAK_START:PEAK_END] + print(f" Loaded {len(obs_all)} USGS obs rows; " + f"full window: {len(obs_full)}, peak window: {len(obs_peak)}") + except Exception as e: + print(f" Warning: USGS load failed ({e}) — plotting without obs") + obs_full = obs_peak = None + + kge = np.nan + if obs_full is not None and len(env_full) > 0: + med = env_full["p50"] + obs_hourly = obs_full.copy() + obs_hourly.index = obs_hourly.index.floor("1h") + obs_hourly = obs_hourly[~obs_hourly.index.duplicated(keep="first")] + aligned = obs_hourly.reindex(med.index) + print(f" KGE alignment: {aligned.notna().sum()} matched out of {len(med)} timesteps") + kge = compute_kge(aligned.values, med.values) + + # ── Figure 1: Full Helene window ────────────────────────────────────────── + fig, ax = plt.subplots(figsize=(12, 4.5)) + _shade(ax, env_full) + if obs_full is not None: + ax.plot(obs_full.index, obs_full.values, color=COLOR_OBS, lw=1.6, + label="USGS obs (gauge 03463300)") + ax.axvspan(PEAK_START, PEAK_END, color="salmon", alpha=0.12, zorder=0, label="Helene peak window") + ax.set_ylabel("Discharge (m³/s)") + ax.set_xlabel("Date") + ax.set_title(f"Routed 600-member ensemble vs USGS | {args.label}\n" + f"Ensemble median KGE = {kge:.3f} | Sep 24–28 2024") + ax.legend(fontsize=9, loc="upper left") + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + plt.tight_layout() + out = out_dir / "routed_ensemble_vs_usgs_full.png" + fig.savefig(out, dpi=150) + print(f"Saved: {out}") + plt.close(fig) + + # ── Figure 2: Peak zoom ─────────────────────────────────────────────────── + fig, ax = plt.subplots(figsize=(10, 4.5)) + _shade(ax, env_peak) + if obs_peak is not None: + ax.plot(obs_peak.index, obs_peak.values, color=COLOR_OBS, lw=1.6, + label="USGS obs (gauge 03463300)") + peak_val = float(obs_peak.max()) + ax.axhline(peak_val, color=COLOR_OBS, lw=0.7, ls=":", alpha=0.5) + ax.annotate(f"USGS peak\n{peak_val:.0f} m³/s", + xy=(HELENE_PEAK, peak_val), xytext=(10, -40), + textcoords="offset points", fontsize=8, color=COLOR_OBS) + ax.set_ylabel("Discharge (m³/s)") + ax.set_xlabel("Date") + ax.set_title(f"Helene peak zoom | {args.label}") + ax.legend(fontsize=9) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %HUTC")) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + plt.tight_layout() + out = out_dir / "routed_ensemble_vs_usgs_peak.png" + fig.savefig(out, dpi=150) + print(f"Saved: {out}") + plt.close(fig) + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/run_4b_crossed_f1.sh b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/run_4b_crossed_f1.sh new file mode 100644 index 00000000..af536b6d --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/run_4b_crossed_f1.sh @@ -0,0 +1,40 @@ +#!/usr/bin/env bash +# F1 Vrugt (1 gauge holdout) — 4b crossed ensemble plots. +# Runs plot_crossed_ensemble.py for all 21 catchments. +# Reads //_crossed_ensemble.parquet (already generated). +# +# Usage: +# bash run_4b_crossed_f1.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +DATA_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +echo "[F1-4b-crossed] Crossed ensemble plots — data: $DATA_DIR" + +for CAT in "${CATS[@]}"; do + PQ="$DATA_DIR/$CAT/${CAT}_crossed_ensemble.parquet" + if [ ! -f "$PQ" ]; then + echo " [$CAT] No crossed ensemble parquet — skipping" + continue + fi + echo " [$CAT] crossed ensemble..." + $TROUTE "$SCRIPT_DIR/plot_crossed_ensemble.py" \ + --ensemble-dir "$DATA_DIR" \ + --cat-id "$CAT" \ + --usgs-csv "$USGS_CSV" +done + +echo "[F1-4b-crossed] Done." diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/run_4b_f1.sh b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/run_4b_f1.sh new file mode 100644 index 00000000..51976cb0 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/run_4b_f1.sh @@ -0,0 +1,52 @@ +#!/usr/bin/env bash +# F1 Vrugt (1 gauge holdout) — 4b ensemble vs obs plots. +# +# Per-catchment: perturbation_category_shaded, per_member_factor_decomp, +# production_ensemble_forecast. +# Multi-catchment (runs once): sensitivity_spaghetti, sensitivity_spread. +# +# Usage: +# bash run_4b_f1.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +EVAL_DIR="$(cd "$SCRIPT_DIR/.." && pwd)" + +DATA_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +echo "[F1-4b] Per-catchment ensemble plots — data: $DATA_DIR" + +for CAT in "${CATS[@]}"; do + PROD_CSV="$DATA_DIR/$CAT/${CAT}_production_per_member.csv" + SEN_CSV="$DATA_DIR/$CAT/${CAT}_sensitivity_init.csv" + + if [ ! -f "$PROD_CSV" ] || [ ! -f "$SEN_CSV" ]; then + echo " [$CAT] Missing prod or sensitivity CSV — skipping" + continue + fi + + echo " [$CAT] perturbation category shaded..." + $TROUTE "$SCRIPT_DIR/plot_perturbation_category_shaded.py" --cat-id "$CAT" + + echo " [$CAT] per-member factor decomposition..." + $TROUTE "$SCRIPT_DIR/plot_per_member_factor_decomposition.py" --cat-id "$CAT" + + echo " [$CAT] production ensemble forecast..." + $TROUTE "$SCRIPT_DIR/plot_production_ensemble_forecast.py" --cat-id "$CAT" +done + +echo "[F1-4b] Multi-catchment sensitivity plots..." +$TROUTE "$EVAL_DIR/plot_perturbation_sensitivity_spaghetti.py" +$TROUTE "$EVAL_DIR/plot_perturbation_sensitivity_spread.py" + +echo "[F1-4b] Done." diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4c_timeseries/plot_forecast_spaghetti.py b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4c_timeseries/plot_forecast_spaghetti.py new file mode 100644 index 00000000..655edaef --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4c_timeseries/plot_forecast_spaghetti.py @@ -0,0 +1,176 @@ +""" +plot_forecast_spaghetti.py + +"All forecasts ending at this hour" — verification-time-fixed spaghetti view. + +Window: Sep 24 18 UTC -> Sep 30 06 UTC. +For every initialization time whose 18-hour forecast overlaps this window: + DA : shaded min/max band across 20 members + ensemble mean line + OL : ensemble mean line only (no shading), dashed + +Lines are colored by initialization DATE (7 colors, Sep 24-30) so the +temporal progression is readable. USGS obs overlaid in black. + +Inputs: + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + Wide format: issue_time, lead_hour, member_00..member_19 (q in m³/s) + +Obs: + /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +Outputs: + /forecast_spaghetti_helene.png +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import matplotlib.patches as mpatches + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast_hardcoded_r_routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +# Verification window — only trajectories whose valid_time falls here are shown +PLOT_START = pd.Timestamp("2024-09-24 18:00:00") +PLOT_END = pd.Timestamp("2024-09-30 06:00:00") + +HELENE_PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +HELENE_PEAK_END = pd.Timestamp("2024-09-28 00:00:00") + +# One color per init date (Sep 24-30) +DATE_COLORS = { + "2024-09-24": "#1f77b4", + "2024-09-25": "#ff7f0e", + "2024-09-26": "#2ca02c", + "2024-09-27": "#d62728", + "2024-09-28": "#9467bd", + "2024-09-29": "#8c564b", + "2024-09-30": "#e377c2", +} + + +def load_parquet(path): + df = pd.read_parquet(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + df["ens_mean"] = df[member_cols].mean(axis=1) + df["ens_min"] = df[member_cols].min(axis=1) + df["ens_max"] = df[member_cols].max(axis=1) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + return df + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() + or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--route-dir", default=DEFAULT_ROUTE_DIR) + parser.add_argument("--out-dir", default=None) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--da-name", default="routed_leadtime_da_full.parquet") + parser.add_argument("--ol-name", default="routed_leadtime_openloop_full.parquet") + args = parser.parse_args() + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + print("Loading parquets...") + da = load_parquet(os.path.join(args.route_dir, args.da_name)) + ol = load_parquet(os.path.join(args.route_dir, args.ol_name)) + obs = load_usgs(args.usgs_csv) + + # Keep only valid_times inside the plot window + da = da[(da["valid_time"] >= PLOT_START) & (da["valid_time"] <= PLOT_END)] + ol = ol[(ol["valid_time"] >= PLOT_START) & (ol["valid_time"] <= PLOT_END)] + + issue_times = sorted(da["issue_time"].unique()) + print(f" {len(issue_times)} initialization times contribute to this window") + + fig, ax = plt.subplots(figsize=(18, 6)) + + # Helene peak shading + ax.axvspan(HELENE_PEAK_START, HELENE_PEAK_END, + color="salmon", alpha=0.12, zorder=0) + ax.text(HELENE_PEAK_START + pd.Timedelta(hours=6), 1.0, "Helene peak", + transform=ax.get_xaxis_transform(), + fontsize=9, color="firebrick", ha="left", va="top") + + # Plot each initialization time's trajectory + for t0 in issue_times: + date_str = str(pd.Timestamp(t0).date()) + color = DATE_COLORS.get(date_str, "gray") + + da_t = da[da["issue_time"] == t0].sort_values("valid_time") + ol_t = ol[ol["issue_time"] == t0].sort_values("valid_time") + + if da_t.empty: + continue + + # DA: shaded band (min-max across 20 members) + ensemble mean + ax.fill_between(da_t["valid_time"], da_t["ens_min"], da_t["ens_max"], + color=color, alpha=0.06, zorder=2) + ax.plot(da_t["valid_time"], da_t["ens_mean"], + color=color, lw=0.8, alpha=0.55, zorder=3) + + # OL: ensemble mean only, dashed + if not ol_t.empty: + ax.plot(ol_t["valid_time"], ol_t["ens_mean"], + color=color, lw=0.6, alpha=0.30, linestyle="--", zorder=2) + + # USGS obs + obs_w = obs.loc[PLOT_START:PLOT_END] + ax.plot(obs_w.index, obs_w.values, + color="black", lw=1.8, zorder=6, label="USGS obs") + + # Legend: one patch per init date + obs + DA/OL style + patches = [mpatches.Patch(color=c, label=f"Init {d}") + for d, c in DATE_COLORS.items()] + patches.append(plt.Line2D([0], [0], color="black", lw=1.8, label="USGS obs")) + patches.append(plt.Line2D([0], [0], color="gray", lw=1.2, alpha=0.7, + label="DA — ens mean (solid) + spread (shaded)")) + patches.append(plt.Line2D([0], [0], color="gray", lw=1.0, linestyle="--", + alpha=0.5, label="OL — ens mean (dashed)")) + ax.legend(handles=patches, fontsize=8, loc="upper left", + frameon=True, framealpha=0.9, ncol=2) + + ax.set_xlim(PLOT_START, PLOT_END) + ax.set_xlabel("Date", fontsize=11) + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_title( + "All 18-hour forecast trajectories ending in this window — USGS 03463300\n" + "Sep 24 18 UTC → Sep 30 06 UTC | DA: shaded band + mean | OL: mean dashed", + fontsize=11, + ) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=1)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.xaxis.set_minor_locator(mdates.HourLocator(interval=6)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.2, lw=0.4) + plt.tight_layout() + + out_path = os.path.join(out_dir, "forecast_spaghetti_helene.png") + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4c_timeseries/plot_helene_issue_time_hydrograph.py b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4c_timeseries/plot_helene_issue_time_hydrograph.py new file mode 100644 index 00000000..c6bc5dd8 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4c_timeseries/plot_helene_issue_time_hydrograph.py @@ -0,0 +1,212 @@ +""" +Per-issue-time forecast hydrograph diagnostic. + +For two contrasting issue times — one storm-peak (Helene) and one low-flow — +show what the DA and open-loop ensembles actually produce over the 18-hour +forecast window, alongside the kriging observation. Designed to answer in one +picture: is DA over-shooting, collapsing, or oscillating compared to open-loop? + +Each panel shows: + - All 20 DA members as faint purple lines + median (thick purple) + - All 20 open-loop members as faint gray lines + median (thick gray, dashed) + - Kriging obs as a thick black line (context: t0-6 through t0+18) + - Vertical line at t0 + annotation with obs(t0) and lead-1 ensemble means + +Inputs (from run_lead_time_forecast_sweep.py): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv +Obs (from production EnKF): + //_test_results.csv + +Output: + //_issue_time_hydrograph_helene_vs_lowflow.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +DEFAULT_LEADTIME_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast" +DEFAULT_DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" + +# Module-level globals set by main() +CAT = "cat-1016300" +LEADTIME_DIR = DEFAULT_LEADTIME_DIR +DA_DIR = DEFAULT_DA_DIR + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" +OBS_COLOR = "black" + +# Two issue times — pick a Helene-peak one and a typical low-flow one +HELENE_T0 = pd.Timestamp("2024-09-26 12:00:00") +LOWFLOW_T0 = pd.Timestamp("2024-03-15 12:00:00") + +CONTEXT_HOURS_BEFORE = 6 # hours of obs context shown before t0 +LEAD_HOURS_AFTER = 18 # forecast horizon + +OUT_PNG = os.path.join( + LEADTIME_DIR, CAT, + f"{CAT}_issue_time_hydrograph_helene_vs_lowflow.png", +) + + +def load_forecasts(path): + df = pd.read_csv(path, parse_dates=['issue_time', 'valid_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + return df, member_cols + + +OBS_DIR = None # set by main() when --obs-dir is supplied + + +def load_obs(): + if OBS_DIR is not None: + obs_path = os.path.join(OBS_DIR, f"{CAT}.csv") + df = pd.read_csv(obs_path) + time_col = 'datetime' if 'datetime' in df.columns else 'date' + df[time_col] = pd.to_datetime(df[time_col]) + col = 'qkrig' if 'qkrig' in df.columns else 'obs_mm_h' + return df.set_index(time_col)[col].rename('obs_mm_h') + obs_path = os.path.join(DA_DIR, CAT, f"{CAT}_test_results.csv") + df = pd.read_csv(obs_path, parse_dates=['date']) + return df.set_index('date')['obs_mm_h'] + + +def slice_forecast(df, member_cols, t0): + """Return (valid_times, member_array shape (lead, N)) for a single issue time.""" + sub = df[df['issue_time'] == t0].sort_values('lead_hour') + if len(sub) == 0: + raise ValueError(f"No forecast rows for issue_time={t0}") + valid_times = pd.to_datetime(sub['valid_time'].values) + member_arr = sub[member_cols].to_numpy(dtype=float) + return valid_times, member_arr + + +def plot_panel(ax, t0, df_da, df_ol, member_cols, obs_series, title_extra=""): + # Forecast trajectories + vt_da, q_da = slice_forecast(df_da, member_cols, t0) + vt_ol, q_ol = slice_forecast(df_ol, member_cols, t0) + + # Obs context window + ctx_start = t0 - pd.Timedelta(hours=CONTEXT_HOURS_BEFORE) + ctx_end = t0 + pd.Timedelta(hours=LEAD_HOURS_AFTER) + obs_window = obs_series.loc[ctx_start:ctx_end] + + # Open-loop members (drawn first so DA paints on top) + for j in range(q_ol.shape[1]): + ax.plot(vt_ol, q_ol[:, j], color=OL_COLOR, lw=0.6, alpha=0.45, zorder=2) + # DA members + for j in range(q_da.shape[1]): + ax.plot(vt_da, q_da[:, j], color=DA_COLOR, lw=0.6, alpha=0.45, zorder=3) + # Medians + ax.plot(vt_ol, np.nanmedian(q_ol, axis=1), color=OL_COLOR, lw=2.6, + linestyle="--", zorder=4, label="Open-loop median") + ax.plot(vt_da, np.nanmedian(q_da, axis=1), color=DA_COLOR, lw=2.6, + zorder=5, label="DA median") + # Obs + ax.plot(obs_window.index, obs_window.values, color=OBS_COLOR, lw=1.8, + zorder=6, label="Qkrig (obs)") + + # t0 marker + annotation + ax.axvline(t0, color="black", lw=1.0, linestyle=":", alpha=0.6, zorder=1) + obs_t0 = float(obs_series.get(t0, np.nan)) + obs_lead1 = float(obs_series.get(t0 + pd.Timedelta(hours=1), np.nan)) + da_lead1_mean = float(np.nanmean(q_da[0, :])) + ol_lead1_mean = float(np.nanmean(q_ol[0, :])) + + info_text = ( + f"t0 = {t0.strftime('%Y-%m-%d %H:%M')}\n" + f"obs(t0) = {obs_t0:.3f} mm/h\n" + f"obs(t0+1) = {obs_lead1:.3f} mm/h\n" + f"DA mean(t0+1) = {da_lead1_mean:.3f} mm/h\n" + f"OL mean(t0+1) = {ol_lead1_mean:.3f} mm/h" + ) + ax.text(0.02, 0.97, info_text, + transform=ax.transAxes, + fontsize=9, va='top', ha='left', family='monospace', + bbox=dict(boxstyle='round,pad=0.4', facecolor='white', + alpha=0.85, edgecolor='gray')) + + ax.set_title(title_extra, fontsize=12) + ax.set_ylabel("Discharge (mm/h)", fontsize=10) + ax.set_xlabel("Date", fontsize=10) + ax.xaxis.set_major_locator(mdates.HourLocator(interval=3)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %H:%M")) + ax.tick_params(axis='x', rotation=30, labelsize=8) + ax.grid(True, alpha=0.2) + ax.legend(loc='upper right', fontsize=9, frameon=True, framealpha=0.9) + + +def snap_to_nearest(t0, available_times): + """Return the available issue_time closest to t0.""" + times = pd.DatetimeIndex(available_times) + idx = np.abs((times - t0).total_seconds()).argmin() + snapped = times[idx] + if snapped != t0: + print(f" Snapped {t0} -> {snapped} (nearest available issue time)") + return snapped + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', default="cat-1016300") + parser.add_argument('--leadtime-dir', default=DEFAULT_LEADTIME_DIR, + help='Dir holding /_lead_time_forecasts_{da,openloop}.csv') + parser.add_argument('--da-dir', default=DEFAULT_DA_DIR, + help='Dir holding /_test_results.csv (for obs)') + parser.add_argument('--obs-dir', default=None, + help='Kriging obs dir holding .csv with qkrig column; ' + 'overrides --da-dir for obs loading') + parser.add_argument('--helene-t0', default=None, + help='Issue time for the Helene panel (default: 2024-09-26 12:00:00, ' + 'snapped to nearest available)') + parser.add_argument('--lowflow-t0', default=None, + help='Issue time for the low-flow panel (default: 2024-03-15 00:00:00, ' + 'snapped to nearest available)') + args = parser.parse_args() + + global CAT, LEADTIME_DIR, DA_DIR, OBS_DIR, OUT_PNG + CAT = args.cat_id + LEADTIME_DIR = args.leadtime_dir + DA_DIR = args.da_dir + OBS_DIR = args.obs_dir + OUT_PNG = os.path.join(LEADTIME_DIR, CAT, + f"{CAT}_issue_time_hydrograph_helene_vs_lowflow.png") + + da_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_da.csv") + ol_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_openloop.csv") + df_da, m_cols = load_forecasts(da_path) + df_ol, _ = load_forecasts(ol_path) + obs_series = load_obs() + + available = df_da['issue_time'].unique() + helene_t0 = snap_to_nearest( + pd.Timestamp(args.helene_t0) if args.helene_t0 else HELENE_T0, available) + lowflow_t0 = snap_to_nearest( + pd.Timestamp(args.lowflow_t0) if args.lowflow_t0 else pd.Timestamp("2024-03-15 00:00:00"), + available) + + fig, (ax_h, ax_l) = plt.subplots(1, 2, figsize=(20, 7)) + + plot_panel(ax_h, helene_t0, df_da, df_ol, m_cols, obs_series, + title_extra=f"Helene peak issue time — {helene_t0.strftime('%Y-%m-%d %H:%M')}") + plot_panel(ax_l, lowflow_t0, df_da, df_ol, m_cols, obs_series, + title_extra=f"Low-flow issue time — {lowflow_t0.strftime('%Y-%m-%d %H:%M')}") + + fig.suptitle( + f"Per-issue-time forecast hydrograph — {CAT}\n" + "Faint lines = 20 ensemble members. Thick lines = ensemble medians. " + "Vertical dotted line = t0 (DA stops here for both scenarios).", + fontsize=12, y=1.00, + ) + plt.tight_layout() + plt.savefig(OUT_PNG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4c_timeseries/plot_lead_time_reconstructed_timeseries.py b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4c_timeseries/plot_lead_time_reconstructed_timeseries.py new file mode 100644 index 00000000..a6d03b03 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4c_timeseries/plot_lead_time_reconstructed_timeseries.py @@ -0,0 +1,263 @@ +""" +Reconstructed time series at gauge 03463300, from our routed +lead-time forecasts. Mirrors the methodology of plot_timeseries_hourly_ensemble.py +(qSpatialAR) but feeds from the EnKF + T-route pipeline instead of the CNN. + +For each target hour t, pool ALL forecasts that land on t (across many +issue_times × many lead_hours × 20 members). Compute median + 5th/95th +percentile envelope. Plot DA and open-loop reconstructions overlaid on USGS +obs, with the Helene window shaded. + +Input parquets (from run_route_troute_leadtime.py): + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + Wide format: issue_time, lead_hour, member_00..member_19 (q in m³/s) + +Obs: + /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +Output: + /lead_time_reconstructed_timeseries.png + Default plot window: 2024-09-10 → 2024-10-10. + +Each scenario's curve is annotated with NSE vs USGS obs over the plot window. +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" +OBS_COLOR = "black" + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast_hardcoded_r_routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +# Watershed area at gauge 03463300 (sum of 21 catchments per GPKG hydrofabric) +WATERSHED_AREA_KM2 = 113.18 +# mm/h depth → m³/s: × 113.18 km² × 1000 / 3600 = 31.439 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +PLOT_START_DEFAULT = pd.Timestamp("2024-09-10 00:00:00") +PLOT_END_DEFAULT = pd.Timestamp("2024-10-10 23:00:00") + +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-30 23:00:00") + + +def load_parquet_long(path): + df = pd.read_parquet(path) + df['issue_time'] = pd.to_datetime(df['issue_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + if not member_cols: + raise ValueError(f"No member columns in {path}") + long = df.melt( + id_vars=['issue_time', 'lead_hour'], + value_vars=member_cols, + var_name='member', + value_name='q_m3s', + ) + long['valid_time'] = (long['issue_time'] + + pd.to_timedelta(long['lead_hour'], unit='h')) + return long + + +def reconstruct(long_df): + """For each valid_time, pool all (issue_time × lead × member) forecasts and + return median, p05, p95, count.""" + grouped = long_df.groupby('valid_time')['q_m3s'].agg( + median='median', + p05=lambda s: s.quantile(0.05), + p95=lambda s: s.quantile(0.95), + count='count', + ).reset_index() + return grouped + + +def load_usgs(usgs_csv): + """Load USGS gauge obs and ensure m³/s. + + The CSV at gauge 03463300 has the column `QObs(mm/h)` — catchment-averaged + depth, not gauge discharge in m³/s. Auto-convert when column name contains + 'mm': m³/s = mm/h × WATERSHED_AREA_KM2 × 1000 / 3600. + """ + df = pd.read_csv(usgs_csv) + date_col = next((c for c in df.columns + if c.lower() in ('datetime', 'date', 'time', 'timestamp')), None) + q_col = next((c for c in df.columns + if 'q' in c.lower() or 'flow' in c.lower() + or 'discharge' in c.lower()), None) + if not date_col or not q_col: + raise ValueError(f"Couldn't id date/Q columns in {usgs_csv}: {df.columns.tolist()}") + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if 'mm' in q_col.lower(): + series = series * MM_H_TO_M3_S + print(f" Converted obs from mm/h (column '{q_col}') to m³/s " + f"using area = {WATERSHED_AREA_KM2} km² (× {MM_H_TO_M3_S:.4f})") + else: + print(f" Loaded obs as-is from column '{q_col}' (assumed m³/s)") + return series + + +def nse(obs, sim): + m = np.isfinite(obs) & np.isfinite(sim) + if m.sum() < 2: + return np.nan + o, s = obs[m], sim[m] + denom = ((o - o.mean()) ** 2).sum() + if denom < 1e-10: + return np.nan + return 1.0 - float(((o - s) ** 2).sum() / denom) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--route-dir', default=DEFAULT_ROUTE_DIR) + parser.add_argument('--out-dir', default=None, + help='Default: --route-dir') + parser.add_argument('--usgs-csv', default=DEFAULT_USGS_CSV) + parser.add_argument('--plot-start', default=str(PLOT_START_DEFAULT)) + parser.add_argument('--plot-end', default=str(PLOT_END_DEFAULT)) + parser.add_argument('--da-name', default='routed_leadtime_da_full.parquet') + parser.add_argument('--ol-name', default='routed_leadtime_openloop_full.parquet') + args = parser.parse_args() + + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + plot_start = pd.Timestamp(args.plot_start) + plot_end = pd.Timestamp(args.plot_end) + + print(f"Loading DA parquet: {args.da_name}") + da_long = load_parquet_long(os.path.join(args.route_dir, args.da_name)) + print(f"Loading OL parquet: {args.ol_name}") + ol_long = load_parquet_long(os.path.join(args.route_dir, args.ol_name)) + print(f"Loading USGS obs: {args.usgs_csv}") + obs = load_usgs(args.usgs_csv) + + print("Reconstructing DA time series (overlapping-leads pool)...") + da_rec = reconstruct(da_long) + print(f" {len(da_rec)} unique target hours") + print("Reconstructing OL time series...") + ol_rec = reconstruct(ol_long) + + # Clip to plot window and join obs + da_rec = da_rec[(da_rec['valid_time'] >= plot_start) & (da_rec['valid_time'] <= plot_end)].copy() + ol_rec = ol_rec[(ol_rec['valid_time'] >= plot_start) & (ol_rec['valid_time'] <= plot_end)].copy() + da_rec['obs'] = da_rec['valid_time'].map(obs) + ol_rec['obs'] = ol_rec['valid_time'].map(obs) + + nse_da = nse(da_rec['obs'].values, da_rec['median'].values) + nse_ol = nse(ol_rec['obs'].values, ol_rec['median'].values) + print(f"NSE (plot window) — DA: {nse_da:.3f} | OL: {nse_ol:.3f}") + + # ---- Plot ---- + fig, ax = plt.subplots(figsize=(16, 5.5)) + fig.patch.set_facecolor("white") + + # Helene shaded + ax.axvspan(HELENE_START, HELENE_END, + color="firebrick", alpha=0.08, zorder=0) + ax.text(HELENE_START + pd.Timedelta(hours=12), 1.0, "Helene", + transform=ax.get_xaxis_transform(), + fontsize=9, color="firebrick", ha="left", va="top", + fontweight="bold") + + # Open-loop band + median + ax.fill_between(ol_rec['valid_time'], ol_rec['p05'], ol_rec['p95'], + color=OL_COLOR, alpha=0.18, zorder=2, + label="Open-loop [5th–95th pctile envelope]") + ax.plot(ol_rec['valid_time'], ol_rec['median'], + color=OL_COLOR, lw=1.2, linestyle='--', zorder=3, + label=f"Open-loop median (NSE={nse_ol:.3f})") + + # DA band + median + ax.fill_between(da_rec['valid_time'], da_rec['p05'], da_rec['p95'], + color=DA_COLOR, alpha=0.22, zorder=4, + label="DA [5th–95th pctile envelope]") + ax.plot(da_rec['valid_time'], da_rec['median'], + color=DA_COLOR, lw=1.4, zorder=5, + label=f"DA median (NSE={nse_da:.3f})") + + # Observed as dots + obs_window = obs.loc[plot_start:plot_end] + finite = np.isfinite(obs_window.values) + ax.scatter(obs_window.index[finite], obs_window.values[finite], + s=6, color=OBS_COLOR, linewidths=0, zorder=6, + label="USGS obs") + + ax.set_xlabel("Date", fontsize=11) + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_title( + f"Reconstructed forecast time series at USGS 03463300\n" + f"Overlapping-leads pool from EnKF forecast ensemble " + f"({plot_start.date()} – {plot_end.date()})", + fontsize=12, + ) + ax.set_xlim(plot_start, plot_end) + ax.xaxis.set_major_locator(mdates.WeekdayLocator(interval=1)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.25, lw=0.4) + ax.legend(loc='upper left', fontsize=9, frameon=True, framealpha=0.92) + + plt.tight_layout() + out_path = os.path.join(out_dir, "lead_time_reconstructed_timeseries.png") + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + # Also produce a Helene-zoom version + helene_zoom_start = pd.Timestamp("2024-09-24 00:00:00") + helene_zoom_end = pd.Timestamp("2024-09-29 23:00:00") + da_z = da_rec[(da_rec['valid_time'] >= helene_zoom_start) + & (da_rec['valid_time'] <= helene_zoom_end)] + ol_z = ol_rec[(ol_rec['valid_time'] >= helene_zoom_start) + & (ol_rec['valid_time'] <= helene_zoom_end)] + nse_da_z = nse(da_z['obs'].values, da_z['median'].values) + nse_ol_z = nse(ol_z['obs'].values, ol_z['median'].values) + print(f"NSE (Helene zoom) — DA: {nse_da_z:.3f} | OL: {nse_ol_z:.3f}") + + fig, ax = plt.subplots(figsize=(14, 5.5)) + fig.patch.set_facecolor("white") + ax.fill_between(ol_z['valid_time'], ol_z['p05'], ol_z['p95'], + color=OL_COLOR, alpha=0.18, zorder=2, + label="Open-loop [5th–95th]") + ax.plot(ol_z['valid_time'], ol_z['median'], color=OL_COLOR, + lw=1.4, linestyle='--', zorder=3, + label=f"Open-loop median (NSE={nse_ol_z:.3f})") + ax.fill_between(da_z['valid_time'], da_z['p05'], da_z['p95'], + color=DA_COLOR, alpha=0.22, zorder=4, + label="DA [5th–95th]") + ax.plot(da_z['valid_time'], da_z['median'], color=DA_COLOR, + lw=1.6, zorder=5, + label=f"DA median (NSE={nse_da_z:.3f})") + obs_z = obs.loc[helene_zoom_start:helene_zoom_end] + finite = np.isfinite(obs_z.values) + ax.scatter(obs_z.index[finite], obs_z.values[finite], + s=12, color=OBS_COLOR, linewidths=0, zorder=6, + label="USGS obs") + ax.set_xlabel("Date", fontsize=11) + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_title( + "Reconstructed forecast time series at USGS 03463300 — Helene zoom\n" + "Sep 24 – 29, 2024 (overlapping-leads pool from EnKF forecast ensemble)", + fontsize=12, + ) + ax.xaxis.set_major_locator(mdates.HourLocator(interval=12)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %H:%M")) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.25, lw=0.4) + ax.legend(loc='upper left', fontsize=9, frameon=True, framealpha=0.92) + plt.tight_layout() + out_zoom = os.path.join(out_dir, "lead_time_reconstructed_timeseries_helene.png") + plt.savefig(out_zoom, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_zoom}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4c_timeseries/run_4c_f1.sh b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4c_timeseries/run_4c_f1.sh new file mode 100644 index 00000000..0c56bc5f --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/4c_timeseries/run_4c_f1.sh @@ -0,0 +1,51 @@ +#!/usr/bin/env bash +# F1 Vrugt (1 gauge holdout) — 4c reconstructed timeseries plots. +# +# Reads routed_leadtime_da_full.parquet + routed_leadtime_openloop_full.parquet +# from the F1 leadtime route dir and writes output PNGs there. +# Must run AFTER route_lead_time_forecasts.py. +# +# Usage: +# bash run_4c_f1.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +LEADTIME_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt +ROUTE_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt_leadtime_routed +OBS_DIR=/home/svyas/catchment_ts_no_03463300_gapfilled +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +echo "[F1-4c] Reconstructed timeseries from: $ROUTE_DIR" +$TROUTE "$SCRIPT_DIR/plot_lead_time_reconstructed_timeseries.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" + +echo "[F1-4c] Forecast spaghetti..." +$TROUTE "$SCRIPT_DIR/plot_forecast_spaghetti.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" + +echo "[F1-4c] Helene vs low-flow issue-time hydrograph (per catchment)..." +for CAT in "${CATS[@]}"; do + DA_CSV="$LEADTIME_DIR/$CAT/${CAT}_lead_time_forecasts_da.csv" + [ -f "$DA_CSV" ] || { echo " [$CAT] no leadtime CSV — skipping"; continue; } + echo " [$CAT] helene hydrograph..." + $TROUTE "$SCRIPT_DIR/plot_helene_issue_time_hydrograph.py" \ + --cat-id "$CAT" \ + --leadtime-dir "$LEADTIME_DIR" \ + --obs-dir "$OBS_DIR" +done + +echo "[F1-4c] Done. Outputs in: $ROUTE_DIR and $LEADTIME_DIR" +ls "$ROUTE_DIR"/*.png 2>/dev/null || echo " (no PNGs in route dir)" diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/compare_da_vs_qkrig_vs_usgs.py b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/compare_da_vs_qkrig_vs_usgs.py new file mode 100644 index 00000000..ae7bc03d --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/compare_da_vs_qkrig_vs_usgs.py @@ -0,0 +1,162 @@ +""" +compare_da_vs_qkrig_vs_usgs.py + +Three-way comparison at gauge 03463300: + DA-routed Q vs routed Qkrig vs USGS obs + +Reads the routed_Q_test.csv produced by run_route.py (which already contains +Q_routed_m3s, Q_usgs_m3s, and Q_krig_m3s columns) and prints KGE/NSE/peak +for the full test period and the Helene window separately. + +Also saves a two-panel comparison figure. + +Usage: + python3 compare_da_vs_qkrig_vs_usgs.py \ + --vrugt-csv /mnt/disk2/suma_helen_poster/da_results/vrugt_dynamic_routed/routed_Q_test.csv \ + --out-dir /mnt/disk2/suma_helen_poster/da_results/comparison_plots +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-29 23:00:00") + + +def kge(obs, sim): + mask = np.isfinite(obs) & np.isfinite(sim) + o, s = obs[mask], sim[mask] + if len(o) < 2 or np.std(o) == 0: + return np.nan + r = np.corrcoef(o, s)[0, 1] + return 1.0 - np.sqrt((r-1)**2 + (np.std(s)/np.std(o)-1)**2 + (np.mean(s)/np.mean(o)-1)**2) + + +def nse(obs, sim): + mask = np.isfinite(obs) & np.isfinite(sim) + o, s = obs[mask], sim[mask] + if len(o) < 2: + return np.nan + denom = np.sum((o - o.mean())**2) + return 1.0 - np.sum((o - s)**2) / denom if denom > 0 else np.nan + + +def peak_ratio(obs, sim): + return np.nanmax(sim) / np.nanmax(obs) + + +def print_stats(label, obs, sim, dates, window_name="full period"): + mask = np.isfinite(obs) & np.isfinite(sim) + print(f" {label:35s} KGE={kge(obs,sim):+.3f} NSE={nse(obs,sim):+.3f} " + f"peak_sim={np.nanmax(sim):.1f} peak_obs={np.nanmax(obs):.1f} " + f"ratio={peak_ratio(obs,sim):.2f}x [{window_name}]") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--vrugt-csv", required=True, + help="routed_Q_test.csv from vrugt_dynamic_routed/") + parser.add_argument("--out-dir", required=True) + args = parser.parse_args() + os.makedirs(args.out_dir, exist_ok=True) + + df = pd.read_csv(args.vrugt_csv, parse_dates=["date"]).set_index("date").sort_index() + + if "Q_usgs_m3s" not in df.columns: + raise ValueError("CSV missing Q_usgs_m3s — run run_route.py with --obs-csv") + if "Q_krig_m3s" not in df.columns: + raise ValueError("CSV missing Q_krig_m3s — need run_route.py with --kv-dir or --obs-csv") + + obs = df["Q_usgs_m3s"].values + da = df["Q_routed_m3s"].values + krig = df["Q_krig_m3s"].values + dates = df.index + + helene = (dates >= HELENE_START) & (dates <= HELENE_END) + + print("\n" + "="*85) + print("THREE-WAY COMPARISON — gauge 03463300 (South Toe River)") + print("="*85) + + print("\n[FULL TEST PERIOD]") + print_stats("DA-routed vs USGS", obs, da, dates, "full") + print_stats("Qkrig-routed vs USGS", obs, krig, dates, "full") + print_stats("DA-routed vs Qkrig-routed", krig, da, dates, "full") + + print("\n[HELENE WINDOW Sep 24–29]") + print_stats("DA-routed vs USGS", obs[helene], da[helene], dates[helene], "Helene") + print_stats("Qkrig-routed vs USGS", obs[helene], krig[helene], dates[helene], "Helene") + print_stats("DA-routed vs Qkrig-routed", krig[helene], da[helene], dates[helene], "Helene") + + print("\n[PEAK VALUES (Helene window)]") + print(f" USGS peak : {np.nanmax(obs[helene]):.1f} m³/s") + print(f" DA-routed peak : {np.nanmax(da[helene]):.1f} m³/s " + f"({np.nanmax(da[helene])/np.nanmax(obs[helene])*100:.0f}% of USGS)") + print(f" Qkrig-routed peak: {np.nanmax(krig[helene]):.1f} m³/s " + f"({np.nanmax(krig[helene])/np.nanmax(obs[helene])*100:.0f}% of USGS)") + print(f" DA vs Qkrig peak : DA is " + f"{'higher' if np.nanmax(da[helene]) > np.nanmax(krig[helene]) else 'lower'} " + f"by {abs(np.nanmax(da[helene])-np.nanmax(krig[helene])):.1f} m³/s") + print("="*85 + "\n") + + # ---- Plot ---- + fig, axes = plt.subplots(2, 1, figsize=(15, 9), + gridspec_kw={"height_ratios": [1, 1.6]}) + + # Top: full period + ax = axes[0] + ax.plot(dates, obs, color="black", lw=0.8, label="USGS obs", zorder=4) + ax.plot(dates, krig, color="#e6820e", lw=0.8, linestyle="--", + label=f"Qkrig-routed KGE={kge(obs,da):.3f}", zorder=2) + ax.plot(dates, da, color="#1f77b4", lw=0.9, + label=f"DA-routed KGE={kge(obs,da):.3f}", zorder=3) + ax.axvspan(HELENE_START, HELENE_END, color="gold", alpha=0.15, zorder=0) + ax.set_ylabel("Discharge (m³/s)", fontsize=10) + ax.set_title("Full test period — DA-routed vs Qkrig-routed vs USGS obs", fontsize=11) + ax.legend(fontsize=9, loc="upper left") + ax.grid(True, alpha=0.25) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%Y-%m")) + ax.xaxis.set_major_locator(mdates.MonthLocator(interval=2)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right", fontsize=8) + + # Bottom: Helene zoom + ax = axes[1] + dh = dates[helene] + kge_da_h = kge(obs[helene], da[helene]) + kge_krig_h = kge(obs[helene], krig[helene]) + nse_da_h = nse(obs[helene], da[helene]) + nse_krig_h = nse(obs[helene], krig[helene]) + + ax.plot(dh, obs[helene], color="black", lw=1.8, zorder=4, label="USGS obs") + ax.plot(dh, krig[helene], color="#e6820e", lw=1.4, linestyle="--", zorder=2, + label=f"Qkrig-routed KGE={kge_krig_h:.3f} NSE={nse_krig_h:.3f}") + ax.plot(dh, da[helene], color="#1f77b4", lw=1.6, zorder=3, + label=f"DA-routed KGE={kge_da_h:.3f} NSE={nse_da_h:.3f}") + + peak_usgs = np.nanmax(obs[helene]) + ax.axhline(peak_usgs, color="black", lw=0.6, linestyle=":", alpha=0.5) + ax.text(HELENE_END - pd.Timedelta(hours=6), peak_usgs * 1.02, + f"USGS peak {peak_usgs:.0f} m³/s", fontsize=8.5, ha="right", color="black") + + ax.set_ylabel("Discharge (m³/s)", fontsize=10) + ax.set_xlabel("Date (UTC)", fontsize=10) + ax.set_title("Helene window — does DA add value beyond Qkrig?", fontsize=11) + ax.legend(fontsize=9, loc="upper left") + ax.grid(True, alpha=0.25) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=1)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right", fontsize=9) + + plt.tight_layout() + out_path = os.path.join(args.out_dir, "da_vs_qkrig_vs_usgs.png") + plt.savefig(out_path, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/plot_da_perturbation_arms.py b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/plot_da_perturbation_arms.py new file mode 100644 index 00000000..6b462aeb --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/plot_da_perturbation_arms.py @@ -0,0 +1,450 @@ +""" +plot_da_perturbation_arms.py — 2a/2b: ensemble spread WITH DA on. + +Loads the two arm CSVs produced by run_perturbation_da_on.py: + //_da_forcing_arm.csv (30 members) + //_da_hydro_arm.csv (20 members) + +Plots: + Panel 2a — Forcing arm: 30-member spaghetti over Helene window + (spread = met forcing uncertainty with DA-corrected initial states) + Panel 2b — Hydro-state arm: 20-member spaghetti over Helene window + (spread = initial state uncertainty with deterministic forcing) + Optional: --ol-csv overlays the open-loop grand median as a + thick dashed gray line for comparison. + Panel 2c — Comparison: median ± spread envelope, both arms + USGS obs + +All trajectories projected to valid_time = issue_time + lead_hour hours. +Colored by initialization date (Sep 24-30). USGS obs in black. + +Outputs: + //_2a_forcing_arm_helene.png + //_2b_hydro_arm_helene.png + //_2ab_arms_comparison.png +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import matplotlib.patches as mpatches + +DEFAULT_ARM_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_perturbation_da_on" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" +DEFAULT_CAT_ID = "cat-1016300" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +PLOT_START = pd.Timestamp("2024-09-24 00:00:00") +PLOT_END = pd.Timestamp("2024-09-30 06:00:00") +HELENE_PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +HELENE_PEAK_END = pd.Timestamp("2024-09-28 00:00:00") + +DATE_COLORS = { + "2024-09-24": "#1f77b4", + "2024-09-25": "#ff7f0e", + "2024-09-26": "#2ca02c", + "2024-09-27": "#d62728", + "2024-09-28": "#9467bd", + "2024-09-29": "#8c564b", + "2024-09-30": "#e377c2", +} + + +def load_arm(path): + """Load arm CSV, compute valid_time, member columns.""" + df = pd.read_csv(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + return df, member_cols + + +def load_openloop_ts(path): + """Load a simple timeseries CSV and return a pd.Series (datetime → m³/s) + clipped to the plot window. + + Handles two formats: + - routed_Q_test.csv : date index + Q_routed_m3s column (already m³/s) + - *_test_results.csv : date column + sim_mm_h column (converted mm/h→m³/s) + """ + df = pd.read_csv(path) + time_col = next( + (c for c in df.columns + if c.lower() in ("time", "datetime", "date", "timestamp")), + None, + ) + if time_col is None: + raise ValueError(f"No time column found in {path}") + df[time_col] = pd.to_datetime(df[time_col]) + + # Prefer routed m³/s columns; fall back to sim_mm_h + sim_col = next( + (c for c in df.columns + if any(k in c.lower() for k in ("q_routed", "routed", "q_cms", "q_m3s"))), + None, + ) + if sim_col is None: + sim_col = next( + (c for c in df.columns + if "sim" in c.lower() and ("mm" in c.lower() or "q" in c.lower())), + None, + ) + if sim_col is None: + sim_col = next((c for c in df.columns if "sim" in c.lower()), None) + if sim_col is None: + raise ValueError(f"No discharge column found in {path}") + + series = df.set_index(time_col)[sim_col].astype(float) + if "mm" in sim_col.lower(): + series = series * MM_H_TO_M3_S + return series.loc[PLOT_START:PLOT_END] + + +def load_openloop(path): + """Load open-loop lead-time file (CSV or Parquet) and return grand median + indexed by valid_time. + + Accepts: + - routed_leadtime_openloop_full.parquet (T-route output, m³/s, recommended) + - cat-*_lead_time_forecasts_openloop.csv (unrouted CFE mm/h, single catchment) + + Returns a pd.Series (valid_time → median q in m³/s) clipped to plot window. + """ + if path.endswith(".parquet"): + df = pd.read_parquet(path) + else: + df = pd.read_csv(path) + + df["issue_time"] = pd.to_datetime(df["issue_time"]) + if "valid_time" in df.columns: + df["valid_time"] = pd.to_datetime(df["valid_time"]) + elif "lead_hour" in df.columns: + df["valid_time"] = (df["issue_time"] + + pd.to_timedelta(df["lead_hour"], unit="h")) + else: + raise ValueError("Open-loop file must have valid_time or lead_hour column") + + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + mask = (df["valid_time"] >= PLOT_START) & (df["valid_time"] <= PLOT_END) + df = df[mask].copy() + if df.empty: + return pd.Series(dtype=float) + + vals = df[member_cols].to_numpy(dtype=float) + # Only convert if values are clearly in mm/h (unrouted CFE output). + # Routed parquet is already in m³/s — do not convert. + if path.endswith(".csv") and np.nanmedian(vals[vals > 0]) < 5: + vals = vals * MM_H_TO_M3_S + + df["q_grand_median"] = np.nanmedian(vals, axis=1) + series = (df.groupby("valid_time")["q_grand_median"] + .median() + .sort_index()) + return series + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def _add_helene_band(ax): + ax.axvspan(HELENE_PEAK_START, HELENE_PEAK_END, + color="salmon", alpha=0.12, zorder=0) + ax.text(HELENE_PEAK_START + pd.Timedelta(hours=6), 1.0, "Helene peak", + transform=ax.get_xaxis_transform(), + fontsize=9, color="firebrick", ha="left", va="top") + + +def _format_xaxis(ax): + ax.set_xlim(PLOT_START, PLOT_END) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=1)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.xaxis.set_minor_locator(mdates.HourLocator(interval=6)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.2, lw=0.4) + + +def plot_spaghetti_arm(ax, df, member_cols, arm_color, obs_series, + title, label_stem, lw_thin=0.7, alpha_thin=0.35): + """ + Plot one thin trajectory per (issue_time, member) pair. + Each trajectory's x = valid_time values for that issue_time. + """ + _add_helene_band(ax) + + issue_times = sorted(df["issue_time"].unique()) + all_means = [] + + for t0 in issue_times: + date_str = str(pd.Timestamp(t0).date()) + color = DATE_COLORS.get(date_str, arm_color) + sub = df[df["issue_time"] == t0].sort_values("valid_time") + vt = sub["valid_time"].values + # Filter to plot window + mask = (sub["valid_time"] >= PLOT_START) & (sub["valid_time"] <= PLOT_END) + sub_w = sub[mask] + if sub_w.empty: + continue + vt_w = sub_w["valid_time"].values + + mem_vals = sub_w[member_cols].to_numpy(dtype=float) # shape (T, N_mem) + if "mm" not in label_stem.lower(): + mem_vals = mem_vals * MM_H_TO_M3_S # mm/h -> m³/s + + # Min-max envelope + median line + ax.fill_between(vt_w, + np.nanmin(mem_vals, axis=1), + np.nanmax(mem_vals, axis=1), + color=color, alpha=0.06, zorder=2) + ax.plot(vt_w, np.nanmedian(mem_vals, axis=1), + color=color, lw=lw_thin, alpha=alpha_thin + 0.1, zorder=3) + + all_means.append( + pd.Series(np.nanmedian(mem_vals, axis=1), index=vt_w)) + + # Overall mean across all issue times (interpolated to common grid) + if all_means: + full_idx = pd.date_range(PLOT_START, PLOT_END, freq="1h") + stacked = pd.concat(all_means, axis=1).reindex(full_idx) + grand_mean = stacked.mean(axis=1) + ax.plot(grand_mean.index, grand_mean.values, + color=arm_color, lw=2.4, alpha=0.95, zorder=5, + label=f"{label_stem} — mean across all forecasts") + + # USGS obs + obs_w = obs_series.loc[PLOT_START:PLOT_END] + ax.plot(obs_w.index, obs_w.values, + color="black", lw=1.8, zorder=6, label="USGS obs") + + ax.set_ylabel("Discharge (m³/s)", fontsize=10) + ax.set_title(title, fontsize=11) + _format_xaxis(ax) + + +def compute_envelope(df, member_cols): + """ + Aggregate all members across all issue_times by valid_time. + Returns DataFrame indexed by valid_time with columns: q_min, q_med, q_max. + """ + records = [] + for t0, grp in df.groupby("issue_time"): + mask = (grp["valid_time"] >= PLOT_START) & (grp["valid_time"] <= PLOT_END) + sub = grp[mask] + if sub.empty: + continue + vals = sub[member_cols].to_numpy(dtype=float) * MM_H_TO_M3_S + for i, row in enumerate(sub.itertuples()): + records.append({ + "valid_time": row.valid_time, + "q_min": np.nanmin(vals[i]), + "q_med": np.nanmedian(vals[i]), + "q_max": np.nanmax(vals[i]), + }) + if not records: + return pd.DataFrame(columns=["valid_time", "q_min", "q_med", "q_max"]) + + env_df = pd.DataFrame(records) + env_df = (env_df.groupby("valid_time") + .agg(q_min=("q_min", "min"), + q_med=("q_med", "mean"), + q_max=("q_max", "max")) + .reset_index() + .sort_values("valid_time")) + return env_df + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--arm-dir", default=DEFAULT_ARM_DIR) + parser.add_argument("--cat-id", default=DEFAULT_CAT_ID) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--out-dir", default=None) + parser.add_argument( + "--ol-csv", default=None, + help=( + "Path to open-loop lead-time forecast CSV " + "(e.g. cat-1016300_lead_time_forecasts_openloop.csv). " + "When provided, the grand median is overlaid on panel 2b " + "as a thick dashed gray line." + ), + ) + parser.add_argument( + "--ol-ts-csv", default=None, + help=( + "Path to a simple open-loop timeseries CSV with 'time' and " + "'sim_mm_h' columns (e.g. cat-*_test_results.csv from the " + "openloop run). Overlaid on panel 2b as a thick dashed black line." + ), + ) + parser.add_argument( + "--ymax", type=float, default=None, + help=( + "Hard y-axis upper limit (m³/s) applied to panels 2a, 2b, and 2ab. " + "Use to clip explosive outlier members (e.g. --ymax 2200 for F2)." + ), + ) + args = parser.parse_args() + + cat_dir = os.path.join(args.arm_dir, args.cat_id) + out_dir = args.out_dir or cat_dir + os.makedirs(out_dir, exist_ok=True) + + forcing_path = os.path.join(cat_dir, f"{args.cat_id}_da_forcing_arm.csv") + hydro_path = os.path.join(cat_dir, f"{args.cat_id}_da_hydro_arm.csv") + + print(f"Loading arm CSVs for {args.cat_id}...") + df_fa, fa_cols = load_arm(forcing_path) + df_ha, ha_cols = load_arm(hydro_path) + obs = load_usgs(args.usgs_csv) + print(f" Forcing arm: {df_fa['issue_time'].nunique()} issue times, " + f"{len(fa_cols)} members") + print(f" Hydro arm: {df_ha['issue_time'].nunique()} issue times, " + f"{len(ha_cols)} members") + + # ------------------------------------------------------------------ # + # Figure 2a: Forcing arm spaghetti # + # ------------------------------------------------------------------ # + fig, ax = plt.subplots(figsize=(17, 6)) + plot_spaghetti_arm( + ax, df_fa, fa_cols, + arm_color="tab:blue", obs_series=obs, + title=(f"2a — Forcing arm (DA on, {len(fa_cols)} members): " + "met forcing uncertainty | Sep 24-30 2024\n" + f"{args.cat_id} | Shaded = member min-max | " + "Line = median per init time | Thick = grand mean"), + label_stem="Forcing arm", + ) + # Date-color legend + patches = [mpatches.Patch(color=c, label=f"Init {d}") + for d, c in DATE_COLORS.items()] + handles, labels = ax.get_legend_handles_labels() + ax.legend(handles=handles + patches, + fontsize=8, loc="upper left", frameon=True, framealpha=0.9, ncol=2) + if args.ymax is not None: + ax.set_ylim(bottom=0, top=args.ymax) + plt.tight_layout() + out_a = os.path.join(out_dir, f"{args.cat_id}_2a_forcing_arm_helene.png") + plt.savefig(out_a, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_a}") + + # ------------------------------------------------------------------ # + # Figure 2b: Hydro-state arm spaghetti # + # ------------------------------------------------------------------ # + fig, ax = plt.subplots(figsize=(17, 6)) + plot_spaghetti_arm( + ax, df_ha, ha_cols, + arm_color="tab:green", obs_series=obs, + title=(f"2b — Hydro-state arm (DA on, {len(ha_cols)} members): " + "initial state uncertainty | Sep 24-30 2024\n" + f"{args.cat_id} | Shaded = member min-max | " + "Line = median per init time | Thick = grand mean"), + label_stem="Hydro-state arm", + ) + + # Optional open-loop overlay — lead-time forecast format + if args.ol_csv: + print(f"Loading open-loop CSV: {args.ol_csv}") + ol_series = load_openloop(args.ol_csv) + if not ol_series.empty: + ax.plot( + ol_series.index, ol_series.values, + color="black", lw=2.5, ls="--", alpha=0.95, zorder=7, + label="Open loop (no DA) — grand median", + ) + else: + print(" Warning: open-loop CSV yielded no data in plot window.") + + # Optional open-loop overlay — simple timeseries format (*_test_results.csv) + if args.ol_ts_csv: + print(f"Loading open-loop timeseries: {args.ol_ts_csv}") + ol_ts = load_openloop_ts(args.ol_ts_csv) + if not ol_ts.empty: + ax.plot( + ol_ts.index, ol_ts.values, + color="black", lw=2.5, ls="--", alpha=0.95, zorder=7, + label="Open loop (no DA) — grand median", + ) + else: + print(" Warning: open-loop timeseries CSV yielded no data in plot window.") + + patches = [mpatches.Patch(color=c, label=f"Init {d}") + for d, c in DATE_COLORS.items()] + handles, labels = ax.get_legend_handles_labels() + ax.legend(handles=handles + patches, + fontsize=8, loc="upper left", frameon=True, framealpha=0.9, ncol=2) + if args.ymax is not None: + ax.set_ylim(bottom=0, top=args.ymax) + plt.tight_layout() + out_b = os.path.join(out_dir, f"{args.cat_id}_2b_hydro_arm_helene.png") + plt.savefig(out_b, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_b}") + + # ------------------------------------------------------------------ # + # Figure 2ab: Comparison — envelopes overlaid, both arms # + # ------------------------------------------------------------------ # + print("Computing aggregated envelopes for comparison plot...") + env_fa = compute_envelope(df_fa, fa_cols) + env_ha = compute_envelope(df_ha, ha_cols) + + fig, ax = plt.subplots(figsize=(17, 7)) + _add_helene_band(ax) + + if not env_fa.empty: + ax.fill_between(env_fa["valid_time"], + env_fa["q_min"], env_fa["q_max"], + color="tab:blue", alpha=0.18, zorder=2, + label=f"Forcing arm spread (N={len(fa_cols)} members)") + ax.plot(env_fa["valid_time"], env_fa["q_med"], + color="tab:blue", lw=2.0, zorder=4, + label="Forcing arm — grand median") + + if not env_ha.empty: + ax.fill_between(env_ha["valid_time"], + env_ha["q_min"], env_ha["q_max"], + color="tab:green", alpha=0.18, zorder=2, + label=f"Hydro-state arm spread (N={len(ha_cols)} members)") + ax.plot(env_ha["valid_time"], env_ha["q_med"], + color="tab:green", lw=2.0, zorder=4, + label="Hydro-state arm — grand median") + + obs_w = obs.loc[PLOT_START:PLOT_END] + ax.plot(obs_w.index, obs_w.values, + color="black", lw=1.8, zorder=6, label="USGS obs") + + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_xlabel("Date", fontsize=11) + ax.set_title( + f"2a vs 2b — Forcing arm (blue) vs Hydro-state arm (green) | " + f"DA on | {args.cat_id}\n" + "Shaded = full member spread (min-max aggregated across all issue times). " + "Lines = grand median.", + fontsize=11, + ) + ax.legend(fontsize=9, loc="upper left", frameon=True, framealpha=0.9) + _format_xaxis(ax) + if args.ymax is not None: + ax.set_ylim(bottom=0, top=args.ymax) + plt.tight_layout() + out_c = os.path.join(out_dir, f"{args.cat_id}_2ab_arms_comparison.png") + plt.savefig(out_c, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_c}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/plot_da_v2_vrugt_helene_grid.py b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/plot_da_v2_vrugt_helene_grid.py new file mode 100644 index 00000000..f7857bd3 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/plot_da_v2_vrugt_helene_grid.py @@ -0,0 +1,98 @@ +""" +Helene comparison grid: Run 3 (no DA) vs DA v2 (true EnKF + Vrugt R) vs Qkrig obs. +3 x 7 grid of all 21 catchments, zoomed to Sep 20 - Oct 5, 2024. + +Produces: helene_da_v2_vrugt_vs_run3_grid.png + +Expects per-catchment *_test_results.csv (columns: date, sim_mm_h, obs_mm_h) +in RUN3_DIR and DA_DIR. Output files are created by +calibrate_catchment_cfe_da_v2.py run_testing_period(). + +Run inside the same conda env that has pandas + matplotlib (e.g. troute). +""" +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import os + +CATS = [ + "cat-1016279", "cat-1016280", "cat-1016281", "cat-1016282", "cat-1016283", + "cat-1016300", "cat-1016301", "cat-1016302", "cat-1016303", "cat-1016304", + "cat-1016305", "cat-1016306", "cat-1016307", "cat-1016308", "cat-1016309", + "cat-1016310", "cat-1016311", "cat-1016312", "cat-1016313", "cat-1016314", + "cat-1016315", +] + +RUN3_DIR = "/mnt/disk2/suma_helen_poster/catchment_results_range100_run3" +DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" +OUT_PNG = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt/helene_da_v2_vrugt_vs_run3_grid.png" +DA_LABEL = "DA v2 (true EnKF + Vrugt R, N=20)" + +HELENE_START = pd.Timestamp("2024-09-20") +HELENE_END = pd.Timestamp("2024-10-05") + + +def kge(obs, sim): + """Kling-Gupta Efficiency on full overlap of obs and sim.""" + m = ~(pd.isna(obs) | pd.isna(sim)) + o, s = obs[m], sim[m] + if len(o) < 2: + return float('nan') + denom = (((o - o.mean()) ** 2).sum() * ((s - s.mean()) ** 2).sum()) ** 0.5 + if denom == 0: + return float('nan') + r = ((o - o.mean()) * (s - s.mean())).sum() / denom + alpha = s.std() / o.std() + beta = s.mean() / o.mean() if o.mean() != 0 else float('nan') + return 1 - ((r - 1) ** 2 + (alpha - 1) ** 2 + (beta - 1) ** 2) ** 0.5 + + +def main(): + fig, axes = plt.subplots(3, 7, figsize=(28, 12), sharex=True) + axes = axes.flatten() + + for i, cat in enumerate(CATS): + ax = axes[i] + run3_csv = os.path.join(RUN3_DIR, cat, f"{cat}_test_results.csv") + da_csv = os.path.join(DA_DIR, cat, f"{cat}_test_results.csv") + + if not (os.path.exists(run3_csv) and os.path.exists(da_csv)): + ax.set_title(f"{cat} (missing)") + continue + + df_run3 = pd.read_csv(run3_csv, parse_dates=["date"]) + df_da = pd.read_csv(da_csv, parse_dates=["date"]) + + mask3 = (df_run3["date"] >= HELENE_START) & (df_run3["date"] <= HELENE_END) + maskd = (df_da["date"] >= HELENE_START) & (df_da["date"] <= HELENE_END) + + kge_run3 = kge(df_run3["obs_mm_h"], df_run3["sim_mm_h"]) + kge_da = kge(df_da["obs_mm_h"], df_da["sim_mm_h"]) + + ax.plot(df_da.loc[maskd, "date"], df_da.loc[maskd, "obs_mm_h"], + color="black", lw=1.8, label="Qkrig (obs)", zorder=1) + ax.plot(df_run3.loc[mask3, "date"], df_run3.loc[mask3, "sim_mm_h"], + color="steelblue", lw=4.5, alpha=0.45, label=f"Run 3 KGE={kge_run3:.2f}", zorder=2) + ax.plot(df_da.loc[maskd, "date"], df_da.loc[maskd, "sim_mm_h"], + color="tomato", lw=1.3, label=f"DA KGE={kge_da:.2f}", zorder=3) + + ax.set_title(f"{cat}", fontsize=10) + ax.tick_params(axis="x", rotation=45, labelsize=7) + ax.tick_params(axis="y", labelsize=8) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=4)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.legend(fontsize=6, loc="upper right") + + fig.suptitle( + f"Hurricane Helene (Sep 20 - Oct 5, 2024) - {DA_LABEL} vs Run 3 baseline", + fontsize=14, y=1.00, + ) + fig.text(0.005, 0.5, "Discharge (mm/h)", va="center", rotation="vertical", fontsize=12) + plt.tight_layout() + plt.savefig(OUT_PNG, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/plot_da_v2_vrugt_helene_zoomed_grid.py b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/plot_da_v2_vrugt_helene_zoomed_grid.py new file mode 100644 index 00000000..1a3a94b6 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/plot_da_v2_vrugt_helene_zoomed_grid.py @@ -0,0 +1,100 @@ +""" +Hurricane Helene PEAK zoom: Run 3 (no DA) vs DA v2 (true EnKF + Vrugt R) vs Qkrig obs. +3 x 7 grid of all 21 catchments, zoomed to the 4-day peak window Sep 24 - Sep 28, 2024. + +Same data as plot_da_v2_vrugt_helene_grid.py, just a tighter x-axis to see +the peak detail. + +Produces: helene_da_v2_vrugt_vs_run3_grid_zoomed.png + +Run inside the same conda env that has pandas + matplotlib (e.g. troute). +""" +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import os + +CATS = [ + "cat-1016279", "cat-1016280", "cat-1016281", "cat-1016282", "cat-1016283", + "cat-1016300", "cat-1016301", "cat-1016302", "cat-1016303", "cat-1016304", + "cat-1016305", "cat-1016306", "cat-1016307", "cat-1016308", "cat-1016309", + "cat-1016310", "cat-1016311", "cat-1016312", "cat-1016313", "cat-1016314", + "cat-1016315", +] + +RUN3_DIR = "/mnt/disk2/suma_helen_poster/catchment_results_range100_run3" +DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" +OUT_PNG = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt/helene_da_v2_vrugt_vs_run3_grid_zoomed.png" +DA_LABEL = "DA v2 (true EnKF + Vrugt R, N=20)" + +# Tighter 4-day peak window +ZOOM_START = pd.Timestamp("2024-09-24") +ZOOM_END = pd.Timestamp("2024-09-28") + + +def kge(obs, sim): + """Kling-Gupta Efficiency on full overlap of obs and sim.""" + m = ~(pd.isna(obs) | pd.isna(sim)) + o, s = obs[m], sim[m] + if len(o) < 2: + return float('nan') + denom = (((o - o.mean()) ** 2).sum() * ((s - s.mean()) ** 2).sum()) ** 0.5 + if denom == 0: + return float('nan') + r = ((o - o.mean()) * (s - s.mean())).sum() / denom + alpha = s.std() / o.std() + beta = s.mean() / o.mean() if o.mean() != 0 else float('nan') + return 1 - ((r - 1) ** 2 + (alpha - 1) ** 2 + (beta - 1) ** 2) ** 0.5 + + +def main(): + fig, axes = plt.subplots(3, 7, figsize=(28, 12), sharex=True) + axes = axes.flatten() + + for i, cat in enumerate(CATS): + ax = axes[i] + run3_csv = os.path.join(RUN3_DIR, cat, f"{cat}_test_results.csv") + da_csv = os.path.join(DA_DIR, cat, f"{cat}_test_results.csv") + + if not (os.path.exists(run3_csv) and os.path.exists(da_csv)): + ax.set_title(f"{cat} (missing)") + continue + + df_run3 = pd.read_csv(run3_csv, parse_dates=["date"]) + df_da = pd.read_csv(da_csv, parse_dates=["date"]) + + mask3 = (df_run3["date"] >= ZOOM_START) & (df_run3["date"] <= ZOOM_END) + maskd = (df_da["date"] >= ZOOM_START) & (df_da["date"] <= ZOOM_END) + + # KGE on full test period (matches log-line value); displayed in legend + kge_run3 = kge(df_run3["obs_mm_h"], df_run3["sim_mm_h"]) + kge_da = kge(df_da["obs_mm_h"], df_da["sim_mm_h"]) + + ax.plot(df_da.loc[maskd, "date"], df_da.loc[maskd, "obs_mm_h"], + color="black", lw=1.8, label="Qkrig (obs)", zorder=1) + ax.plot(df_run3.loc[mask3, "date"], df_run3.loc[mask3, "sim_mm_h"], + color="steelblue", lw=4.0, alpha=0.45, label=f"Run 3 KGE={kge_run3:.2f}", zorder=2) + ax.plot(df_da.loc[maskd, "date"], df_da.loc[maskd, "sim_mm_h"], + color="tomato", lw=1.4, label=f"DA KGE={kge_da:.2f}", zorder=3) + + ax.set_title(f"{cat}", fontsize=10) + ax.tick_params(axis="x", rotation=45, labelsize=7) + ax.tick_params(axis="y", labelsize=8) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=1)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.xaxis.set_minor_locator(mdates.HourLocator(byhour=[6, 12, 18])) + ax.legend(fontsize=6, loc="upper right") + + fig.suptitle( + f"Hurricane Helene peak zoom (Sep 24 - Sep 28, 2024) - {DA_LABEL} vs Run 3 baseline", + fontsize=14, y=1.00, + ) + fig.text(0.005, 0.5, "Discharge (mm/h)", va="center", rotation="vertical", fontsize=12) + plt.tight_layout() + plt.savefig(OUT_PNG, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/plot_perturbation_sensitivity_spaghetti.py b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/plot_perturbation_sensitivity_spaghetti.py new file mode 100644 index 00000000..40e6acfa --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/plot_perturbation_sensitivity_spaghetti.py @@ -0,0 +1,163 @@ +""" +Plot perturbation-source sensitivity spaghetti plots. + +Reads the three per-source sensitivity CSVs per catchment and plots all 20 +member streamflow trajectories as colored bundles overlaid on one panel, +zoomed to the Hurricane Helene peak window (Sep 24 - Sep 28, 2024). + +Color legend: + red = initial state perturbation only + blue = forcing perturbation only (precip + PET) + green = process noise on hydrologic states only + +Produces: + helene_sensitivity_spaghetti_main.png (3-catchment subset for the main figure) + helene_sensitivity_spaghetti_appendix.png (full 3 x 7 grid of all 21 catchments) + +Requires the sensitivity runs to have completed first (see +run_perturbation_sensitivity.py). +""" +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import os + +# ---------- Configuration ---------- +ALL_CATS = [ + "cat-1016279", "cat-1016280", "cat-1016281", "cat-1016282", "cat-1016283", + "cat-1016300", "cat-1016301", "cat-1016302", "cat-1016303", "cat-1016304", + "cat-1016305", "cat-1016306", "cat-1016307", "cat-1016308", "cat-1016309", + "cat-1016310", "cat-1016311", "cat-1016312", "cat-1016313", "cat-1016314", + "cat-1016315", +] + +# Main figure: 3 catchments spanning the Run 3 KGE distribution. +# cat-1016311 = worst Run 3 (0.50); cat-1016300 = median (0.74); cat-1016302 = best (0.83) +MAIN_CATS = ["cat-1016311", "cat-1016300", "cat-1016302"] + +DA_DIR = "/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt" +SEN_DIR = "/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt" +OBS_DIR = "/home/svyas/catchment_ts_no_03463300_gapfilled" +MAIN_PNG = "/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt/helene_sensitivity_spaghetti_main.png" +APPENDIX_PNG = "/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt/helene_sensitivity_spaghetti_appendix.png" + +# Helene 4-day peak zoom +ZOOM_START = pd.Timestamp("2024-09-24") +ZOOM_END = pd.Timestamp("2024-09-28") + +SOURCE_COLORS = { + "init": "tab:red", + "forcing": "tab:blue", + "process": "tab:green", +} +SOURCE_LABELS = { + "init": "Initial state only", + "forcing": "Forcing only (P, PET)", + "process": "Process noise on states only", +} + + +def load_sensitivity_csv(cat, source): + """Returns (dates, q_matrix) for the test period (member columns).""" + path = os.path.join(SEN_DIR, cat, f"{cat}_sensitivity_{source}.csv") + if not os.path.exists(path): + return None, None + df = pd.read_csv(path, parse_dates=["date"]) + member_cols = [c for c in df.columns if c.startswith("member_")] + return df["date"].values, df[member_cols].values + + +def load_qkrig_obs(cat): + obs_p = os.path.join(OBS_DIR, f"{cat}.csv") + if os.path.exists(obs_p): + df = pd.read_csv(obs_p) + time_col = 'datetime' if 'datetime' in df.columns else 'date' + df[time_col] = pd.to_datetime(df[time_col]) + return df[time_col].values, df["qkrig"].values + path = os.path.join(DA_DIR, cat, f"{cat}_test_results.csv") + if not os.path.exists(path): + return None, None + df = pd.read_csv(path, parse_dates=["date"]) + return df["date"].values, df["obs_mm_h"].values + + +def plot_one_panel(ax, cat): + """Render the three colored bundles plus the obs line for one catchment.""" + # Each source: 20 thin colored lines + for source, color in SOURCE_COLORS.items(): + dates, q = load_sensitivity_csv(cat, source) + if dates is None: + continue + df_dates = pd.to_datetime(dates) + mask = (df_dates >= ZOOM_START) & (df_dates <= ZOOM_END) + if mask.sum() == 0: + continue + # Plot all 20 members as thin transparent lines, plus a thicker mean line + for i in range(q.shape[1]): + ax.plot(df_dates[mask], q[mask, i], + color=color, lw=0.5, alpha=0.35, zorder=1) + ax.plot(df_dates[mask], q[mask].mean(axis=1), + color=color, lw=1.6, alpha=0.95, zorder=2, + label=SOURCE_LABELS[source]) + + # Black observation overlay + obs_dates, obs = load_qkrig_obs(cat) + if obs_dates is not None: + od = pd.to_datetime(obs_dates) + mask = (od >= ZOOM_START) & (od <= ZOOM_END) + ax.plot(od[mask], obs[mask], + color="black", lw=1.4, label="Qkrig (obs)", zorder=3) + + ax.set_title(cat, fontsize=10) + ax.tick_params(axis="x", rotation=45, labelsize=7) + ax.tick_params(axis="y", labelsize=8) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=1)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + + +def make_main_figure(): + """3-catchment side-by-side panel for the main paper / poster figure.""" + fig, axes = plt.subplots(1, 3, figsize=(18, 5), sharex=True) + for ax, cat in zip(axes, MAIN_CATS): + plot_one_panel(ax, cat) + # Single legend on the rightmost panel (avoid clutter on the others) + axes[-1].legend(fontsize=8, loc="upper right") + axes[0].set_ylabel("Discharge (mm/h)", fontsize=11) + fig.suptitle( + "Ensemble spread by perturbation source — Hurricane Helene peak (Sep 24-28, 2024)\n" + "20 members per source. DA off. Worst / median / best Run 3 catchment.", + fontsize=12, y=1.02, + ) + plt.tight_layout() + plt.savefig(MAIN_PNG, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {MAIN_PNG}") + + +def make_appendix_figure(): + """Full 3 x 7 grid of all 21 catchments.""" + fig, axes = plt.subplots(3, 7, figsize=(28, 12), sharex=True) + axes = axes.flatten() + for i, cat in enumerate(ALL_CATS): + plot_one_panel(axes[i], cat) + # Single legend on the first panel + axes[0].legend(fontsize=7, loc="upper right") + fig.text(0.005, 0.5, "Discharge (mm/h)", va="center", rotation="vertical", fontsize=12) + fig.suptitle( + "Ensemble spread by perturbation source - Hurricane Helene peak (Sep 24-28, 2024) - all 21 sub-catchments\n" + "Red = initial state only | Blue = forcing only | Green = process noise only | Black = Qkrig obs", + fontsize=14, y=1.00, + ) + plt.tight_layout() + plt.savefig(APPENDIX_PNG, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {APPENDIX_PNG}") + + +def main(): + make_main_figure() + make_appendix_figure() + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/plot_perturbation_sensitivity_spread.py b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/plot_perturbation_sensitivity_spread.py new file mode 100644 index 00000000..434566f3 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/plot_perturbation_sensitivity_spread.py @@ -0,0 +1,168 @@ +""" +Plot ensemble-spread time series, one curve per perturbation source. + +For each catchment and each of the three sensitivity sub-experiments, compute the +hourly std-dev of streamflow across the 20 members. Plot all three as colored +lines on the same axes so the relative contribution of each source is directly +comparable hour by hour. + +This is the cleanest answer to "which perturbation source contributes most to +ensemble spread" — the spaghetti version is hard to read when the bundles +overlap. A single std-dev curve per source removes the visual clutter. + +Color legend: + red = std-dev across initial-state-only ensemble + blue = std-dev across forcing-only ensemble + green = std-dev across process-noise-only ensemble + +Produces: + helene_sensitivity_spread_main.png (3-catchment subset for the main figure) + helene_sensitivity_spread_appendix.png (full 3 x 7 grid of all 21 catchments) + +Consumes the same per-source CSVs produced by run_perturbation_sensitivity.py. +""" +import pandas as pd +import numpy as np +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import os + +ALL_CATS = [ + "cat-1016279", "cat-1016280", "cat-1016281", "cat-1016282", "cat-1016283", + "cat-1016300", "cat-1016301", "cat-1016302", "cat-1016303", "cat-1016304", + "cat-1016305", "cat-1016306", "cat-1016307", "cat-1016308", "cat-1016309", + "cat-1016310", "cat-1016311", "cat-1016312", "cat-1016313", "cat-1016314", + "cat-1016315", +] + +MAIN_CATS = ["cat-1016311", "cat-1016300", "cat-1016302"] + +DA_DIR = "/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt" +SEN_DIR = "/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt" +OBS_DIR = "/home/svyas/catchment_ts_no_03463300_gapfilled" +MAIN_PNG = "/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt/helene_sensitivity_spread_main.png" +APPENDIX_PNG = "/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt/helene_sensitivity_spread_appendix.png" + +ZOOM_START = pd.Timestamp("2024-09-24") +ZOOM_END = pd.Timestamp("2024-09-28") + +SOURCE_COLORS = { + "init": "tab:red", + "forcing": "tab:blue", + "process": "tab:green", +} +SOURCE_LABELS = { + "init": "Initial state perturbation", + "forcing": "Forcing perturbation (P, PET)", + "process": "Process noise on states", +} + + +def load_member_spread(cat, source): + """Returns (dates, std_per_hour) for the test period. + + std_per_hour is the std-dev of the 20 member streamflow values at each hour. + """ + path = os.path.join(SEN_DIR, cat, f"{cat}_sensitivity_{source}.csv") + if not os.path.exists(path): + return None, None + df = pd.read_csv(path, parse_dates=["date"]) + member_cols = [c for c in df.columns if c.startswith("member_")] + member_arr = df[member_cols].to_numpy(dtype=float) + # Hourly std-dev across members (sample std, ddof=1 to match Pyy convention) + std = member_arr.std(axis=1, ddof=1) + return df["date"].values, std + + +def load_qkrig_obs(cat): + obs_p = os.path.join(OBS_DIR, f"{cat}.csv") + if os.path.exists(obs_p): + df = pd.read_csv(obs_p) + time_col = 'datetime' if 'datetime' in df.columns else 'date' + df[time_col] = pd.to_datetime(df[time_col]) + return df[time_col].values, df["qkrig"].values + path = os.path.join(DA_DIR, cat, f"{cat}_test_results.csv") + if not os.path.exists(path): + return None, None + df = pd.read_csv(path, parse_dates=["date"]) + return df["date"].values, df["obs_mm_h"].values + + +def plot_one_panel(ax, cat, show_obs=True): + """Render std-dev time series for the three sources plus optional obs reference.""" + for source, color in SOURCE_COLORS.items(): + dates, std = load_member_spread(cat, source) + if dates is None: + continue + d = pd.to_datetime(dates) + mask = (d >= ZOOM_START) & (d <= ZOOM_END) + if mask.sum() == 0: + continue + ax.plot(d[mask], std[mask], color=color, lw=1.6, + label=SOURCE_LABELS[source], zorder=2) + + if show_obs: + # Plot the observation on a secondary y-axis for context (storm timing) + obs_dates, obs = load_qkrig_obs(cat) + if obs_dates is not None: + od = pd.to_datetime(obs_dates) + mask = (od >= ZOOM_START) & (od <= ZOOM_END) + ax2 = ax.twinx() + ax2.plot(od[mask], obs[mask], + color="black", lw=1.0, alpha=0.35, zorder=1, label="Qkrig (obs, ref)") + ax2.set_ylabel("Qkrig (mm/h)", fontsize=8, color="0.4") + ax2.tick_params(axis="y", labelsize=7, colors="0.4") + ax2.spines["right"].set_color("0.7") + + ax.set_title(cat, fontsize=10) + ax.tick_params(axis="x", rotation=45, labelsize=7) + ax.tick_params(axis="y", labelsize=8) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=1)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.set_ylabel("Ensemble std-dev (mm/h)", fontsize=8) + + +def make_main_figure(): + fig, axes = plt.subplots(1, 3, figsize=(18, 5), sharex=True) + for ax, cat in zip(axes, MAIN_CATS): + plot_one_panel(ax, cat, show_obs=True) + axes[-1].legend(fontsize=8, loc="upper right") + fig.suptitle( + "Ensemble spread (std-dev across 20 members) by perturbation source\n" + "Hurricane Helene peak (Sep 24-28, 2024) | DA off | " + "worst / median / best Run 3 catchment\n" + "Grey line on right axis = Qkrig observation (for storm-timing context only)", + fontsize=11, y=1.06, + ) + plt.tight_layout() + plt.savefig(MAIN_PNG, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {MAIN_PNG}") + + +def make_appendix_figure(): + fig, axes = plt.subplots(3, 7, figsize=(28, 12), sharex=True) + axes = axes.flatten() + for i, cat in enumerate(ALL_CATS): + plot_one_panel(axes[i], cat, show_obs=True) + axes[0].legend(fontsize=7, loc="upper right") + fig.suptitle( + "Ensemble spread (hourly std-dev across 20 members) by perturbation source - " + "Hurricane Helene peak (Sep 24-28, 2024) - all 21 sub-catchments\n" + "Red = init state | Blue = forcing | Green = process noise | " + "Grey on right axis = Qkrig obs (storm-timing context only)", + fontsize=13, y=1.00, + ) + plt.tight_layout() + plt.savefig(APPENDIX_PNG, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {APPENDIX_PNG}") + + +def main(): + make_main_figure() + make_appendix_figure() + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/plot_vrugt_comparison.py b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/plot_vrugt_comparison.py new file mode 100644 index 00000000..d241bc04 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/plot_vrugt_comparison.py @@ -0,0 +1,188 @@ +#!/usr/bin/env python3 +""" +plot_vrugt_comparison.py +Compare routed discharge: dynamic Vrugt R vs. constant R (no-Vrugt) +against USGS gauge obs at 03463300 (South Toe River Near Celo, NC). + +Reads routed_Q_test.csv from both run directories (produced by run_route.py +with --obs-csv). Produces: + - Two-panel figure: full test period + Helene zoom + - Single-panel Helene zoom only + +Usage: + python3 plot_vrugt_comparison.py \ + --vrugt-csv /mnt/disk2/suma_helen_poster/da_results/vrugt_dynamic_routed/routed_Q_test.csv \ + --novrugt-csv /mnt/disk2/suma_helen_poster/da_results/novrugt_dynamic_routed/routed_Q_test.csv \ + --out-dir /mnt/disk2/suma_helen_poster/da_results/comparison_plots +""" + +import argparse +import os + +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import matplotlib.patches as mpatches + +HELENE_START = pd.Timestamp("2024-09-25 00:00:00") +HELENE_END = pd.Timestamp("2024-09-29 00:00:00") +ZOOM_START = pd.Timestamp("2024-09-24 00:00:00") +ZOOM_END = pd.Timestamp("2024-09-30 00:00:00") + +COLOR_OBS = "black" +COLOR_VRUGT = "#1f77b4" # blue +COLOR_NOVRUGT = "#d62728" # red + + +def load_csv(path): + df = pd.read_csv(path, parse_dates=["date"]) + df = df.set_index("date").sort_index() + return df + + +def compute_kge(obs, sim): + mask = ~(np.isnan(obs) | np.isnan(sim)) + o, s = obs[mask], sim[mask] + if len(o) < 2 or np.std(o) == 0: + return np.nan + r = np.corrcoef(o, s)[0, 1] + return 1.0 - np.sqrt((r - 1)**2 + (np.std(s) / np.std(o) - 1)**2 + (np.mean(s) / np.mean(o) - 1)**2) + + +def compute_nse(obs, sim): + mask = ~(np.isnan(obs) | np.isnan(sim)) + o, s = obs[mask], sim[mask] + if len(o) < 2 or np.sum((o - np.mean(o))**2) == 0: + return np.nan + return 1.0 - np.sum((o - s)**2) / np.sum((o - np.mean(o))**2) + + +def _add_helene_band(ax): + ax.axvspan(HELENE_START, HELENE_END, + color="gold", alpha=0.18, zorder=0, label="_nolegend_") + + +def _format_xaxis(ax, locator, fmt, rotation=30): + ax.xaxis.set_major_locator(locator) + ax.xaxis.set_major_formatter(mdates.DateFormatter(fmt)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=rotation, ha="right", fontsize=9) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--vrugt-csv", required=True) + parser.add_argument("--novrugt-csv", required=True) + parser.add_argument("--out-dir", required=True) + args = parser.parse_args() + os.makedirs(args.out_dir, exist_ok=True) + + vrugt = load_csv(args.vrugt_csv) + novrugt = load_csv(args.novrugt_csv) + + # USGS obs from vrugt CSV (same timestamps, both runs used same --obs-csv) + obs = vrugt["Q_usgs_m3s"].values + dates = vrugt.index + + sim_v = vrugt["Q_routed_m3s"].values + sim_nv = novrugt["Q_routed_m3s"].reindex(dates).values + + kge_v = compute_kge(obs, sim_v) + nse_v = compute_nse(obs, sim_v) + kge_nv = compute_kge(obs, sim_nv) + nse_nv = compute_nse(obs, sim_nv) + + peak_usgs = np.nanmax(obs) + peak_v = np.nanmax(sim_v) + peak_nv = np.nanmax(sim_nv) + + print(f"USGS peak: {peak_usgs:.1f} m3/s") + print(f"Vrugt peak: {peak_v:.1f} m3/s KGE={kge_v:.3f} NSE={nse_v:.3f}") + print(f"No-Vrugt peak: {peak_nv:.1f} m3/s KGE={kge_nv:.3f} NSE={nse_nv:.3f}") + + # ------------------------------------------------------------------ # + # Figure 1: two-panel (full period + Helene zoom) + # ------------------------------------------------------------------ # + fig, (ax_full, ax_zoom) = plt.subplots( + 2, 1, figsize=(14, 9), + gridspec_kw={"height_ratios": [1, 1.4]}) + + # -- Top: full period -- + ax_full.plot(dates, obs, color=COLOR_OBS, lw=0.8, label="USGS obs (gauge 03463300)", zorder=3) + ax_full.plot(dates, sim_v, color=COLOR_VRUGT, lw=0.9, label=f"Vrugt (dynamic R) KGE={kge_v:.3f}", zorder=2) + ax_full.plot(dates, sim_nv, color=COLOR_NOVRUGT, lw=0.9, linestyle="--", + label=f"No-Vrugt (const R) KGE={kge_nv:.3f}", zorder=2) + _add_helene_band(ax_full) + ax_full.text(HELENE_START + pd.Timedelta(hours=12), ax_full.get_ylim()[1] * 0.85, + "Helene", fontsize=9, color="goldenrod", fontweight="bold") + ax_full.set_ylabel("Discharge (m³/s)", fontsize=10) + ax_full.set_title( + "Routed discharge at gauge 03463300 (South Toe River)\n" + "DA with dynamic Vrugt R vs. constant R -- Muskingum routing", + fontsize=11) + ax_full.legend(fontsize=9, loc="upper left") + ax_full.grid(True, alpha=0.25) + _format_xaxis(ax_full, mdates.MonthLocator(interval=2), "%Y-%m") + + # Add zoom indicator + ax_full.axvspan(ZOOM_START, ZOOM_END, color="gray", alpha=0.10, zorder=0) + ax_full.annotate("", xy=(ZOOM_END, ax_full.get_ylim()[1] * 0.5), + xytext=(ZOOM_START, ax_full.get_ylim()[1] * 0.5), + arrowprops=dict(arrowstyle="<->", color="gray", lw=1.0)) + + # -- Bottom: Helene zoom -- + mask_zoom = (dates >= ZOOM_START) & (dates <= ZOOM_END) + dz = dates[mask_zoom] + ax_zoom.plot(dz, obs[mask_zoom], color=COLOR_OBS, lw=1.2, label="USGS obs (gauge 03463300)", zorder=3) + ax_zoom.plot(dz, sim_v[mask_zoom], color=COLOR_VRUGT, lw=1.4, + label=f"Vrugt (dynamic R) KGE={kge_v:.3f} NSE={nse_v:.3f} peak={peak_v:.0f} m³/s", zorder=2) + ax_zoom.plot(dz, sim_nv[mask_zoom], color=COLOR_NOVRUGT, lw=1.4, linestyle="--", + label=f"No-Vrugt (const R) KGE={kge_nv:.3f} NSE={nse_nv:.3f} peak={peak_nv:.0f} m³/s", zorder=2) + _add_helene_band(ax_zoom) + ax_zoom.axhline(peak_usgs, color=COLOR_OBS, lw=0.7, linestyle=":", alpha=0.6) + ax_zoom.text(ZOOM_START + pd.Timedelta(hours=3), peak_usgs * 1.01, + f"USGS peak {peak_usgs:.0f} m³/s", fontsize=8, color="black", alpha=0.8) + ax_zoom.set_xlabel("Date (UTC)", fontsize=10) + ax_zoom.set_ylabel("Discharge (m³/s)", fontsize=10) + ax_zoom.set_title("Hurricane Helene window (Sep 24-30, 2024)", fontsize=10) + ax_zoom.legend(fontsize=9, loc="lower right") + ax_zoom.grid(True, alpha=0.25) + _format_xaxis(ax_zoom, mdates.DayLocator(interval=1), "%b %d") + + plt.tight_layout() + out1 = os.path.join(args.out_dir, "vrugt_vs_novrugt_twopanel.png") + plt.savefig(out1, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + # ------------------------------------------------------------------ # + # Figure 2: Helene zoom only (poster-ready single panel) + # ------------------------------------------------------------------ # + fig, ax = plt.subplots(figsize=(10, 5)) + ax.plot(dz, obs[mask_zoom], color=COLOR_OBS, lw=1.4, label="USGS obs (gauge 03463300)", zorder=3) + ax.plot(dz, sim_v[mask_zoom], color=COLOR_VRUGT, lw=1.6, + label=f"DA -- dynamic R (Vrugt) KGE={kge_v:.3f}", zorder=2) + ax.plot(dz, sim_nv[mask_zoom], color=COLOR_NOVRUGT, lw=1.6, linestyle="--", + label=f"DA -- constant R KGE={kge_nv:.3f}", zorder=2) + _add_helene_band(ax) + ax.axhline(peak_usgs, color=COLOR_OBS, lw=0.8, linestyle=":", alpha=0.5) + ax.text(ZOOM_START + pd.Timedelta(hours=3), peak_usgs * 1.02, + f"USGS peak {peak_usgs:.0f} m³/s", fontsize=9, color="black", alpha=0.75) + ax.set_xlabel("Date (UTC)", fontsize=11) + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_title( + "Hurricane Helene -- Routed discharge at gauge 03463300 (South Toe River)\n" + "CFE + DA (Muskingum routing): dynamic vs. constant observation error variance", + fontsize=11) + ax.legend(fontsize=10, loc="lower right") + ax.grid(True, alpha=0.25) + _format_xaxis(ax, mdates.DayLocator(interval=1), "%b %d") + plt.tight_layout() + out2 = os.path.join(args.out_dir, "vrugt_vs_novrugt_helene_zoom.png") + plt.savefig(out2, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016279_best_params.json b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016279_best_params.json new file mode 100644 index 00000000..2b05b093 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016279_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016279", + "best_kge": 0.8206116996389099, + "best_parameters": { + "bb": 2.623111719150318, + "smcmax": 0.2, + "satdk": 9.781515618242848e-05, + "slop": 0.49040228658878326, + "max_gw_storage": 0.22502382651758357, + "expon": 3.9507553183626896, + "Cgw": 0.00028265811089655624, + "K_lf": 0.33708866739487964, + "K_nash": 0.7457712895070667, + "scheme": 0.024063276949108925 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016280_best_params.json b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016280_best_params.json new file mode 100644 index 00000000..577b7703 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016280_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016280", + "best_kge": 0.7845371451180185, + "best_parameters": { + "bb": 3.6853461244917503, + "smcmax": 0.2, + "satdk": 5.2992817070504554e-05, + "slop": 0.19547785108903476, + "max_gw_storage": 0.2491156800177733, + "expon": 2.254619375138279, + "Cgw": 0.000560475352848284, + "K_lf": 1.0, + "K_nash": 0.3230987974065112, + "scheme": 0.34480885239395503 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016281_best_params.json b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016281_best_params.json new file mode 100644 index 00000000..17346f8b --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016281_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016281", + "best_kge": 0.7653169268425993, + "best_parameters": { + "bb": 3.6065627281526345, + "smcmax": 0.2, + "satdk": 6.865558286144331e-05, + "slop": 0.16111639333094452, + "max_gw_storage": 0.17693006455477892, + "expon": 2.0040076166999006, + "Cgw": 0.00038006296817321816, + "K_lf": 0.7105167430808037, + "K_nash": 9.127e-08, + "scheme": 0.14298989888102898 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016282_best_params.json b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016282_best_params.json new file mode 100644 index 00000000..0e53037b --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016282_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016282", + "best_kge": 0.7868853451455085, + "best_parameters": { + "bb": 2.627220084315586, + "smcmax": 0.2, + "satdk": 7.889407420323658e-05, + "slop": 0.4136394925802219, + "max_gw_storage": 0.23896390723786864, + "expon": 7.187637073442723, + "Cgw": 9.498445725943469e-06, + "K_lf": 0.44532012273911126, + "K_nash": 0.30386284986679846, + "scheme": 0.4286174650536479 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016283_best_params.json b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016283_best_params.json new file mode 100644 index 00000000..57a06e77 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016283_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016283", + "best_kge": 0.7456494190014495, + "best_parameters": { + "bb": 2.500182877645966, + "smcmax": 0.2, + "satdk": 7.708052707732115e-05, + "slop": 0.43914051816633154, + "max_gw_storage": 0.17844520971245803, + "expon": 1.636624858961519, + "Cgw": 0.0007186031814004434, + "K_lf": 0.7604006750247336, + "K_nash": 0.935851019129656, + "scheme": 0.4295615557232241 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016300_best_params.json b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016300_best_params.json new file mode 100644 index 00000000..99b56bcf --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016300_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016300", + "best_kge": 0.7852863461640692, + "best_parameters": { + "bb": 2.1364476151747556, + "smcmax": 0.22502190033753489, + "satdk": 8.571865512107468e-05, + "slop": 3.06e-05, + "max_gw_storage": 0.25, + "expon": 5.11973472709136, + "Cgw": 0.0001071612486323186, + "K_lf": 0.25307693130285724, + "K_nash": 0.3553147233842257, + "scheme": 0.37275121762887575 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016301_best_params.json b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016301_best_params.json new file mode 100644 index 00000000..a4fc0214 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016301_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016301", + "best_kge": 0.7852546470254989, + "best_parameters": { + "bb": 2.4254415208275426, + "smcmax": 0.2, + "satdk": 7.587190338000397e-05, + "slop": 0.7895909975757629, + "max_gw_storage": 0.2175504021153558, + "expon": 2.748699844202327, + "Cgw": 0.0004359795527131938, + "K_lf": 0.7766373897901322, + "K_nash": 0.3599347814597565, + "scheme": 0.024075462516893846 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016302_best_params.json b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016302_best_params.json new file mode 100644 index 00000000..e9d65dff --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016302_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016302", + "best_kge": 0.7438367117368438, + "best_parameters": { + "bb": 2.0, + "smcmax": 0.2597182603107496, + "satdk": 0.00010927248622434935, + "slop": 0.74928833629841, + "max_gw_storage": 0.20590055515471217, + "expon": 2.6302920260523215, + "Cgw": 0.00036611733203253385, + "K_lf": 0.8653682881680377, + "K_nash": 0.023172715332685445, + "scheme": 1.621e-05 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016303_best_params.json b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016303_best_params.json new file mode 100644 index 00000000..6516c488 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016303_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016303", + "best_kge": 0.7717901438015875, + "best_parameters": { + "bb": 2.0011493088041257, + "smcmax": 0.25267740195882527, + "satdk": 0.00010267632057588901, + "slop": 0.7436120787827551, + "max_gw_storage": 0.25, + "expon": 3.40584201489453, + "Cgw": 0.0002855314649977274, + "K_lf": 2.22e-06, + "K_nash": 0.03462430327931909, + "scheme": 0.4006007187079474 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016304_best_params.json b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016304_best_params.json new file mode 100644 index 00000000..aeef2972 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016304_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016304", + "best_kge": 0.785816115160746, + "best_parameters": { + "bb": 2.627220084315586, + "smcmax": 0.2, + "satdk": 6.623684502570866e-05, + "slop": 0.7710190313472335, + "max_gw_storage": 0.23151061983908472, + "expon": 2.852593661529338, + "Cgw": 0.0003041527852735811, + "K_lf": 0.5192505545426351, + "K_nash": 0.77484040937936, + "scheme": 0.17801104705809254 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016305_best_params.json b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016305_best_params.json new file mode 100644 index 00000000..6bae6748 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016305_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016305", + "best_kge": 0.7744236016300623, + "best_parameters": { + "bb": 2.9787357329850916, + "smcmax": 0.2, + "satdk": 2.01708703759236e-05, + "slop": 0.2671853140329312, + "max_gw_storage": 0.20533656076576476, + "expon": 1.0, + "Cgw": 0.0012738648070631134, + "K_lf": 7.201e-07, + "K_nash": 0.20780337905391905, + "scheme": 0.14434010116001775 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016306_best_params.json b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016306_best_params.json new file mode 100644 index 00000000..5cf772c1 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016306_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016306", + "best_kge": 0.7210357839982497, + "best_parameters": { + "bb": 2.0, + "smcmax": 0.27768197452250365, + "satdk": 0.00011181892180957246, + "slop": 0.17080432119242306, + "max_gw_storage": 0.18472414564626138, + "expon": 1.0299420125608787, + "Cgw": 0.001262070070347035, + "K_lf": 1.049e-05, + "K_nash": 0.10556306992792686, + "scheme": 0.23455263978683152 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016307_best_params.json b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016307_best_params.json new file mode 100644 index 00000000..92a410bf --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016307_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016307", + "best_kge": 0.801808946394892, + "best_parameters": { + "bb": 3.4434053536999176, + "smcmax": 0.2, + "satdk": 5.6465464765815605e-05, + "slop": 0.4263217361110973, + "max_gw_storage": 0.25, + "expon": 4.337261095363485, + "Cgw": 0.0001245938975937281, + "K_lf": 0.845893106799097, + "K_nash": 0.2386693544648092, + "scheme": 0.13456997008570973 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016308_best_params.json b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016308_best_params.json new file mode 100644 index 00000000..b644d485 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016308_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016308", + "best_kge": 0.8061757655923723, + "best_parameters": { + "bb": 3.6702131149274404, + "smcmax": 0.2, + "satdk": 6.003079869425282e-05, + "slop": 0.41011011063560954, + "max_gw_storage": 0.23346373674800242, + "expon": 3.9123522328874523, + "Cgw": 0.000183853837621649, + "K_lf": 0.04420155753170629, + "K_nash": 0.4655759687155024, + "scheme": 0.2840640813933577 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016309_best_params.json b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016309_best_params.json new file mode 100644 index 00000000..c03c2c55 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016309_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016309", + "best_kge": 0.7407364262156276, + "best_parameters": { + "bb": 3.833847789992432, + "smcmax": 0.2, + "satdk": 5.364111110266205e-05, + "slop": 0.23910717165219175, + "max_gw_storage": 0.23812574692579214, + "expon": 3.010993385187965, + "Cgw": 0.0002958484512966826, + "K_lf": 0.3302542952661117, + "K_nash": 0.2737476576326264, + "scheme": 0.4849311470054365 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016310_best_params.json b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016310_best_params.json new file mode 100644 index 00000000..a193e2ad --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016310_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016310", + "best_kge": 0.7414257687945072, + "best_parameters": { + "bb": 4.77839967410216, + "smcmax": 0.2, + "satdk": 0.0005122349525110497, + "slop": 0.2343084566046406, + "max_gw_storage": 0.1972951075076314, + "expon": 2.689926753507914, + "Cgw": 0.00025277943884201336, + "K_lf": 0.10451650547499411, + "K_nash": 0.4766491757683344, + "scheme": 0.14703543348764395 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016311_best_params.json b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016311_best_params.json new file mode 100644 index 00000000..88e348ad --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016311_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016311", + "best_kge": 0.7370239364407825, + "best_parameters": { + "bb": 3.7224279985565554, + "smcmax": 0.2, + "satdk": 5.36848284289468e-05, + "slop": 0.2747382803914872, + "max_gw_storage": 0.20236610051190926, + "expon": 1.2298400851415858, + "Cgw": 0.0008686565105330384, + "K_lf": 0.5491364936227011, + "K_nash": 0.5637455013613035, + "scheme": 0.3172192584604878 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016312_best_params.json b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016312_best_params.json new file mode 100644 index 00000000..12e8f931 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016312_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016312", + "best_kge": 0.7651323332103412, + "best_parameters": { + "bb": 3.4163907102604436, + "smcmax": 0.2, + "satdk": 5.8900848701685016e-05, + "slop": 0.952690690782515, + "max_gw_storage": 0.20339558487059237, + "expon": 2.0393414268931, + "Cgw": 0.0005321846942954315, + "K_lf": 0.005616844228055001, + "K_nash": 0.10556197423792017, + "scheme": 0.16715606346334788 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016313_best_params.json b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016313_best_params.json new file mode 100644 index 00000000..b08b0bc6 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016313_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016313", + "best_kge": 0.7541333608718466, + "best_parameters": { + "bb": 2.0, + "smcmax": 0.3236473885603026, + "satdk": 0.0006067307257057383, + "slop": 1.0, + "max_gw_storage": 0.23803639343037397, + "expon": 4.052333909196705, + "Cgw": 0.0002604427783083463, + "K_lf": 1.0, + "K_nash": 0.013987580178250053, + "scheme": 0.8686112513355525 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016314_best_params.json b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016314_best_params.json new file mode 100644 index 00000000..26e34c1b --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016314_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016314", + "best_kge": 0.8237475701476421, + "best_parameters": { + "bb": 2.781255453729351, + "smcmax": 0.20819967102864187, + "satdk": 7.660284170122033e-05, + "slop": 0.41683958183300995, + "max_gw_storage": 0.23944464809927757, + "expon": 5.8409631177306265, + "Cgw": 5.492004041138577e-05, + "K_lf": 0.846968216875224, + "K_nash": 0.46122458817838097, + "scheme": 0.3779302228964909 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016315_best_params.json b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016315_best_params.json new file mode 100644 index 00000000..e3dbfa20 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/best_params/cat-1016315_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016315", + "best_kge": 0.8232409477622364, + "best_parameters": { + "bb": 2.6967758837666995, + "smcmax": 0.2, + "satdk": 7.062588287975823e-05, + "slop": 5.173e-07, + "max_gw_storage": 0.24187360050863188, + "expon": 4.042512273338185, + "Cgw": 0.0003098824042018927, + "K_lf": 0.1631469310671928, + "K_nash": 0.07301338855959635, + "scheme": 0.431430924337884 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/batch_run_crossed_f2.sh b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/batch_run_crossed_f2.sh new file mode 100644 index 00000000..be9d6953 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/batch_run_crossed_f2.sh @@ -0,0 +1,65 @@ +#!/usr/bin/env bash +# Batch crossed ensemble for all 21 catchments — F2 fixed R=0.07. +# Reads best_params from the F2 DA out dir (already staged). +# Outputs _crossed_ensemble.parquet per catchment to the same dir. +# +# Usage: +# nohup bash ~/da_1gauge_f2/2_assimilation/batch_run_crossed_f2.sh > ~/crossed_f2.log 2>&1 & + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PYTHON=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT="$SCRIPT_DIR/run_crossed_ensemble.py" + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +OBS_DIR=/home/svyas/catchment_ts_no_03463300_gapfilled +CFE_DIR=/mnt/disk2/suma_helen_poster/cfe_py +CONFIG_FILE=/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder2_fixed_r007 +FORCING_DIR1=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings +FORCING_DIR2=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings + +echo "F2 Fixed R=0.07 crossed ensemble — out: $OUT_DIR" +echo "" + +SKIP=0; DONE=0; FAIL=0 + +for CAT in "${CATS[@]}"; do + echo "===============================" + echo "=== $CAT ===" + echo "===============================" + + CAT_OUT="$OUT_DIR/$CAT" + + if [ ! -f "$CAT_OUT/${CAT}_best_params.json" ]; then + echo " WARNING: No best_params for $CAT — skipping" + SKIP=$((SKIP + 1)); continue + fi + + if [ -f "$CAT_OUT/${CAT}_crossed_ensemble.parquet" ]; then + echo " Crossed ensemble already exists — skipping" + SKIP=$((SKIP + 1)); continue + fi + + "$PYTHON" "$SCRIPT" \ + --cat-id "$CAT" \ + --forcing-dir "$FORCING_DIR1" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --out-dir "$OUT_DIR" \ + --test-forcing-dir1 "$FORCING_DIR1" \ + --test-forcing-dir2 "$FORCING_DIR2" \ + --hardcoded-r 0.07 \ + && DONE=$((DONE + 1)) \ + || { echo " FAILED: $CAT"; FAIL=$((FAIL + 1)); } +done + +echo "" +echo "=== F2 crossed done: done=$DONE skipped=$SKIP failed=$FAIL ===" diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/batch_run_da_on_all_cats.sh b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/batch_run_da_on_all_cats.sh new file mode 100644 index 00000000..211dc573 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/batch_run_da_on_all_cats.sh @@ -0,0 +1,82 @@ +#!/usr/bin/env bash +# Batch DA for all 21 catchments — F2 fixed R=0.07. +# R = 0.07 (hardcoded, constant observation error variance) +# Stages best_params.json from calibration results before running. +# +# Usage (from server, after SCPing this folder): +# bash batch_run_da_on_all_cats.sh + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PYTHON=/home/svyas/miniconda3/envs/troute/bin/python3 +DA_SCRIPT="$SCRIPT_DIR/run_perturbation_da_on.py" + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +OBS_DIR=/home/svyas/catchment_ts_no_03463300_gapfilled +BEST_PARAMS_SRC=/mnt/disk2/suma_helen_poster/catchment_results_1gauge_heldout +CFE_DIR=/mnt/disk2/suma_helen_poster/cfe_py +CONFIG_FILE=/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json +PARAM_BOUNDS=/mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder2_fixed_r007 +FORCING_DIR1=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings +FORCING_DIR2=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings + +mkdir -p "$OUT_DIR" +echo "F2 Fixed R=0.07 DA — obs: $OBS_DIR" +echo " out: $OUT_DIR" +echo "" + +SKIP=0; DONE=0; FAIL=0 + +for CAT in "${CATS[@]}"; do + echo "===============================" + echo "=== $CAT ===" + echo "===============================" + + CAT_OUT="$OUT_DIR/$CAT" + mkdir -p "$CAT_OUT" + + # Stage best_params from calibration results + if [ ! -f "$CAT_OUT/${CAT}_best_params.json" ]; then + SRC="$BEST_PARAMS_SRC/$CAT/${CAT}_best_params.json" + if [ -f "$SRC" ]; then + cp "$SRC" "$CAT_OUT/" + echo " Staged best_params" + else + echo " WARNING: No best_params for $CAT — skipping" + SKIP=$((SKIP + 1)); continue + fi + else + echo " best_params already staged" + fi + + # Skip if both arm CSVs already exist + if [ -f "$CAT_OUT/${CAT}_da_forcing_arm.csv" ] && \ + [ -f "$CAT_OUT/${CAT}_da_hydro_arm.csv" ]; then + echo " Arms already complete — skipping" + SKIP=$((SKIP + 1)); continue + fi + + "$PYTHON" "$DA_SCRIPT" \ + --cat-id "$CAT" \ + --forcing-dir "$FORCING_DIR1" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --param-bounds "$PARAM_BOUNDS" \ + --out-dir "$OUT_DIR" \ + --test-forcing-dir1 "$FORCING_DIR1" \ + --test-forcing-dir2 "$FORCING_DIR2" \ + --hardcoded-r 0.07 \ + && DONE=$((DONE + 1)) \ + || { echo " FAILED: $CAT"; FAIL=$((FAIL + 1)); } +done + +echo "" +echo "=== F2 done: done=$DONE skipped=$SKIP failed=$FAIL ===" diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/batch_run_leadtime_f2.sh b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/batch_run_leadtime_f2.sh new file mode 100644 index 00000000..b9810bb4 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/batch_run_leadtime_f2.sh @@ -0,0 +1,77 @@ +#!/usr/bin/env bash +# Batch lead-time forecast sweep — F2 fixed R=0.07, 1 gauge holdout. +# Runs run_lead_time_forecast_sweep.py for all 21 catchments. +# Uses fixed R=0.07 via --hardcoded-r 0.07. +# Skips if both output CSVs already exist. +# +# After this completes, run route_lead_time_forecasts.py to route the +# per-catchment CSVs through T-route and produce routed_leadtime_*.parquet. +# +# Usage: +# nohup bash ~/da_1gauge_f2/2_assimilation/batch_run_leadtime_f2.sh \ +# > ~/leadtime_f2.log 2>&1 & + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PYTHON=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT="$SCRIPT_DIR/run_lead_time_forecast_sweep.py" + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +OBS_DIR=/home/svyas/catchment_ts_no_03463300_gapfilled +CFE_DIR=/mnt/disk2/suma_helen_poster/cfe_py +CONFIG_FILE=/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json +PARAM_BOUNDS=/mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder2_fixed_r007 +FORCING_DIR1=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings +FORCING_DIR2=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings + +echo "F2 lead-time sweep (R=0.07) — out: $OUT_DIR" +echo "" + +SKIP=0; DONE=0; FAIL=0 + +for CAT in "${CATS[@]}"; do + echo "===============================" + echo "=== $CAT ===" + echo "===============================" + + CAT_OUT="$OUT_DIR/$CAT" + DA_CSV="$CAT_OUT/${CAT}_lead_time_forecasts_da.csv" + OL_CSV="$CAT_OUT/${CAT}_lead_time_forecasts_openloop.csv" + + if [ ! -f "$CAT_OUT/${CAT}_best_params.json" ]; then + echo " WARNING: No best_params for $CAT — skipping" + SKIP=$((SKIP + 1)); continue + fi + + if [ -f "$DA_CSV" ] && [ -f "$OL_CSV" ]; then + echo " Lead-time CSVs already exist — skipping" + SKIP=$((SKIP + 1)); continue + fi + + "$PYTHON" "$SCRIPT" \ + --cat-id "$CAT" \ + --forcing-dir "$FORCING_DIR1" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --param-bounds "$PARAM_BOUNDS" \ + --out-dir "$OUT_DIR" \ + --test-forcing-dir1 "$FORCING_DIR1" \ + --test-forcing-dir2 "$FORCING_DIR2" \ + --enkf-members 20 \ + --base-step-h 6 \ + --hardcoded-r 0.07 \ + --prod-script "$SCRIPT_DIR/calibrate_catchment_cfe_da_v2.py" \ + && DONE=$((DONE + 1)) \ + || { echo " FAILED: $CAT"; FAIL=$((FAIL + 1)); } +done + +echo "" +echo "=== F2 lead-time sweep done: done=$DONE skipped=$SKIP failed=$FAIL ===" diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/batch_run_production_f2.sh b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/batch_run_production_f2.sh new file mode 100644 index 00000000..c0846fe7 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/batch_run_production_f2.sh @@ -0,0 +1,71 @@ +#!/usr/bin/env bash +# Batch production per-member run — F2 Fixed R=0.07, 1 gauge holdout. +# Runs run_production_per_member.py for all 21 catchments with --hardcoded-r 0.07. +# Skips if output CSV already exists. +# +# Usage: +# nohup bash batch_run_production_f2.sh > ~/logs/production_f2.log 2>&1 & + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PYTHON=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT="$SCRIPT_DIR/run_production_per_member.py" +PROD_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/1_distributed_cfe/calibrate_catchment_cfe_da_v2.py + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +OBS_DIR=/home/svyas/catchment_ts_no_03463300_gapfilled +CFE_DIR=/mnt/disk2/suma_helen_poster/cfe_py +CONFIG_FILE=/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json +PARAM_BOUNDS=/mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder2_fixed_r007 +FORCING_DIR1=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings +FORCING_DIR2=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings + +echo "F2 production per-member (R=0.07 fixed) — out: $OUT_DIR" +echo "" + +SKIP=0; DONE=0; FAIL=0 + +for CAT in "${CATS[@]}"; do + echo "===============================" + echo "=== $CAT ===" + echo "===============================" + + CAT_OUT="$OUT_DIR/$CAT" + OUT_FILE="$CAT_OUT/${CAT}_production_per_member.csv" + + if [ ! -f "$CAT_OUT/${CAT}_best_params.json" ]; then + echo " WARNING: No best_params for $CAT — skipping" + SKIP=$((SKIP + 1)); continue + fi + + if [ -f "$OUT_FILE" ]; then + echo " Production per-member CSV already exists — skipping" + SKIP=$((SKIP + 1)); continue + fi + + "$PYTHON" "$SCRIPT" \ + --cat-id "$CAT" \ + --forcing-dir "$FORCING_DIR1" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --param-bounds "$PARAM_BOUNDS" \ + --out-dir "$OUT_DIR" \ + --test-forcing-dir1 "$FORCING_DIR1" \ + --test-forcing-dir2 "$FORCING_DIR2" \ + --enkf-members 20 \ + --hardcoded-r 0.07 \ + --prod-script "$PROD_SCRIPT" \ + && DONE=$((DONE + 1)) \ + || { echo " FAILED: $CAT"; FAIL=$((FAIL + 1)); } +done + +echo "" +echo "=== F2 production per-member done: done=$DONE skipped=$SKIP failed=$FAIL ===" diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/batch_run_sensitivity_f2.sh b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/batch_run_sensitivity_f2.sh new file mode 100644 index 00000000..9b650275 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/batch_run_sensitivity_f2.sh @@ -0,0 +1,68 @@ +#!/usr/bin/env bash +# Batch perturbation sensitivity analysis — F2 Fixed R=0.07, 1 gauge holdout. +# Runs run_perturbation_sensitivity.py for all 21 catchments × 3 sources +# (init, forcing, process). Skips if output CSV already exists. +# +# Usage: +# nohup bash ~/da_1gauge_f2/2_assimilation/batch_run_sensitivity_f2.sh \ +# > ~/sensitivity_f2.log 2>&1 & + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PYTHON=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT="$SCRIPT_DIR/run_perturbation_sensitivity.py" + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +SOURCES=(init forcing process) + +OBS_DIR=/home/svyas/catchment_ts_no_03463300_gapfilled +CFE_DIR=/mnt/disk2/suma_helen_poster/cfe_py +CONFIG_FILE=/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json +PARAM_BOUNDS=/mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder2_fixed_r007 +FORCING_DIR1=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings +FORCING_DIR2=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings + +echo "F2 sensitivity analysis — out: $OUT_DIR" +echo "" + +SKIP=0; DONE=0; FAIL=0 + +for CAT in "${CATS[@]}"; do + for SRC in "${SOURCES[@]}"; do + OUT_FILE="$OUT_DIR/$CAT/${CAT}_sensitivity_${SRC}.csv" + + if [ -f "$OUT_FILE" ]; then + echo " [$CAT/$SRC] already exists — skipping" + SKIP=$((SKIP + 1)) + continue + fi + + echo "===============================" + echo "=== $CAT source=$SRC ===" + echo "===============================" + + "$PYTHON" "$SCRIPT" \ + --cat-id "$CAT" \ + --source "$SRC" \ + --forcing-dir "$FORCING_DIR1" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --param-bounds "$PARAM_BOUNDS" \ + --out-dir "$OUT_DIR" \ + --test-forcing-dir1 "$FORCING_DIR1" \ + --test-forcing-dir2 "$FORCING_DIR2" \ + && DONE=$((DONE + 1)) \ + || { echo " FAILED: $CAT/$SRC"; FAIL=$((FAIL + 1)); } + done +done + +echo "" +echo "=== F2 sensitivity done: done=$DONE skipped=$SKIP failed=$FAIL ===" diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/input_enkf_new.json b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/input_enkf_new.json new file mode 100644 index 00000000..482233bc --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/input_enkf_new.json @@ -0,0 +1 @@ +{"n": 1, "m": 1, "R": 0.07, "Q": 0, "smcmax" : 0.9394097311639178, "N": 1000, "P": 0.01, "D":2} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/new_EnKF.py b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/new_EnKF.py new file mode 100644 index 00000000..ff945c73 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/new_EnKF.py @@ -0,0 +1,53 @@ +# EnKF.py +import numpy as np + +class EnKF: + def __init__(self, n, m, R, Q, N, P, smcmax, D): + self.n = n # number of states + self.m = m # number of measurements + self.R = R # measurement noise + self.Q = Q # process noise + self.N = N # number of ensembles + self.P = P + self.state_estimates = np.zeros((self.n)) + self.covariance_matrices = np.zeros((self.n, self.n)) + self.current_step = 0 + self.smcmax = smcmax + self.D = D + self.storage_max_m = self.smcmax * self.D + self.storage_init = self.storage_max_m * 0.667 + self.ensembles = np.full((self.n, self.N), self.storage_init) + + def predict(self, F_results): + error_factor_sim = 0.005 + for i in range(self.N): + perturbation_factor_sim = np.random.standard_normal() + self.ensembles[:, i] = F_results[i] + self.ensembles[:, i] += (perturbation_factor_sim * F_results[i] * error_factor_sim) + return self.ensembles + + def update(self, yi, H_results): + if np.any(np.isnan(yi)): + return np.mean(self.ensembles, axis=1) # return current state estimate without updating + + y_ensembles = np.zeros((self.m, self.N)) + error_factor_sm = 0.03 + for i in range(self.N): + perturbation_factor_sm = np.random.standard_normal() + y_ensembles[:, i] = H_results[i] + y_ensembles[:, i] += (perturbation_factor_sm * H_results[i] * error_factor_sm) + + ensemble_mean = np.mean(self.ensembles, axis=1)[:, np.newaxis] + y_ensembles_mean = np.mean(y_ensembles, axis=1)[:, np.newaxis] + Pxy = (self.ensembles - ensemble_mean).dot((y_ensembles - y_ensembles_mean).T) / (self.N - 1) + Pyy = np.cov(y_ensembles, bias=True) + K = Pxy.dot(np.linalg.pinv(Pyy + self.R)) + self.ensembles += K.dot(yi - y_ensembles_mean) + return np.mean(self.ensembles, axis=1) + + def get_state_estimate(self): + return np.mean(self.ensembles, axis=1) + + def get_covariance(self): + ensemble_mean = np.mean(self.ensembles, axis=1)[:, np.newaxis] + return (self.ensembles - ensemble_mean).dot((self.ensembles - ensemble_mean).T) / (self.N - 1) diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/run_crossed_ensemble.py b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/run_crossed_ensemble.py new file mode 100644 index 00000000..325a64e2 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/run_crossed_ensemble.py @@ -0,0 +1,452 @@ +""" +run_crossed_ensemble.py -- 4b: 600-member crossed ensemble for Helene. + +Crosses the two DA-on perturbation arms: + Forcing arm : 30 met draws (from run_perturbation_da_on.py Phase 2A) + Hydro arm : 20 state draws (from run_perturbation_da_on.py Phase 2B) + +For each Helene issue time t0: + Run 600 members = 30 forcing draws x 20 hydro-state draws. + Each member (i, j) gets: + - Initial state = perturbed DA analysis state (draw j of hydro arm) + - Met forcing = perturbed precip/PET sequence (draw i of forcing arm) + +This gives the full uncertainty envelope combining both sources — +the operational probabilistic forecast picture. + +Output (per catchment): + //_crossed_ensemble.parquet + Columns: issue_time, lead_hour, member_0000..member_0599 + +Usage: + python3 run_crossed_ensemble.py \\ + --cat-id cat-1016300 \\ + --forcing-dir \\ + --obs-dir /mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance \\ + --cfe-dir /mnt/disk2/suma_helen_poster/cfe_py \\ + --config-file /mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json \\ + --out-dir /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble \\ + --test-forcing-dir1 \\ + --test-forcing-dir2 + +Strategy: + Re-runs DA analysis (same as Phase 1 of run_perturbation_da_on.py) to get + the per-issue-time analysis state. Then generates N_FORCING sets of perturbed + forcing sequences and N_HYDRO perturbed states. Runs all 600 combinations. + + To avoid re-running Phase 1, if the Phase 1 snapshot file exists from a + prior run_perturbation_da_on.py run, it is loaded directly. + +NOTE: All special characters kept ASCII to avoid encoding issues on GPU. +""" + +import argparse +import os +import sys +import json +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +N_FORCING = 30 +N_HYDRO = 20 +N_MEMBERS = N_FORCING * N_HYDRO # 600 + +TEST_START = "2024-08-24 00:00:00" # 1 month spinup before Helene window +TEST_END = "2024-10-31 23:00:00" + +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-30 06:00:00") +FORECAST_HOURS = 18 + +PRECIP_SIGMA = 0.15 +PET_SIGMA = 0.10 +STATE_FRAC = 0.05 + +CAT_ID = None +OBS_FILE = None +CFE_CONFIG_FILE = None +TEST_FORCING_FILE = None +OUT_DIR = None +HARDCODED_R = None +bmi_cfe = None + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(): + df = pd.read_csv(TEST_FORCING_FILE) + df["date"] = pd.to_datetime(df["time"]).dt.strftime("%Y-%m-%d %H:%M:%S") + df["total_precipitation"] = df["APCP_surface"] * 3600.0 + df["potential_evaporation"] = priestley_taylor_pet( + df["DSWRF_surface"].values, df["TMP_2maboveground"].values) + return df + + +def load_obs(): + df = pd.read_csv(OBS_FILE) + date_col = next((c for c in df.columns + if c.lower() in ("date", "time", "datetime", "timestamp")), None) + q_col = next((c for c in df.columns + if "qkrig" in c.lower() and "var" not in c.lower()), None) + var_col = next((c for c in df.columns if "var" in c.lower()), None) + df[date_col] = pd.to_datetime(df[date_col]) + df = df.sort_values(date_col).reset_index(drop=True) + date_strs = df[date_col].dt.strftime("%Y-%m-%d %H:%M:%S") + obs_dict = dict(zip(date_strs, df[q_col])) + var_vals = df[var_col] if var_col else pd.Series([0.07] * len(df)) + var_dict = dict(zip(date_strs, var_vals)) + return obs_dict, var_dict + + +def write_model_config(best_params): + """Write the BMI config file once; reused for every model instantiation.""" + with open(CFE_CONFIG_FILE) as f: + cfg = json.load(f) + cfg["forcing_file"] = TEST_FORCING_FILE + cfg["soil_params"]["bb"] = best_params["bb"] + cfg["soil_params"]["smcmax"] = best_params["smcmax"] + cfg["soil_params"]["satdk"] = best_params["satdk"] + cfg["slop"] = best_params["slop"] + cfg["max_gw_storage"] = best_params["max_gw_storage"] + cfg["expon"] = best_params["expon"] + cfg["Cgw"] = best_params["Cgw"] + cfg["K_lf"] = best_params["K_lf"] + cfg["K_nash"] = best_params["K_nash"] + cfg["partition_scheme"] = ("Schaake" if best_params["scheme"] <= 0.5 + else "Xinanjiang") + tmp_cfg = str(OUT_DIR / f"{CAT_ID}_bmi_config_crossed_stable.json") + with open(tmp_cfg, "w") as f: + json.dump(cfg, f) + return tmp_cfg + + +def make_model(cfg_path): + """Instantiate and initialize a CFE model from a pre-written config path.""" + m = bmi_cfe.BMI_CFE(cfg_file=cfg_path) + m.initialize() + return m + + +def get_state(m): + return { + "soil_m": m.soil_reservoir["storage_m"], + "gw_m": m.gw_reservoir["storage_m"], + "nash0": float(m.nash_storage[0]), + "nash1": float(m.nash_storage[1]), + } + + +def set_state(m, state): + m.soil_reservoir["storage_m"] = state["soil_m"] + m.gw_reservoir["storage_m"] = state["gw_m"] + m.nash_storage[0] = state["nash0"] + m.nash_storage[1] = state["nash1"] + + +def perturb_state(state, rng, frac=STATE_FRAC): + return { + "soil_m": max(state["soil_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "nash0": max(state["nash0"] * (1.0 + frac * rng.standard_normal()), 0.0), + "nash1": max(state["nash1"] * (1.0 + frac * rng.standard_normal()), 0.0), + } + + +def run_model_step(m, precip_mmh, pet_mmh): + m.set_value("atmosphere_water__time_integral_of_precipitation_mass_flux", + precip_mmh / 1000.0) + m.set_value("water_potential_evaporation_flux", + pet_mmh / 1000.0 / 3600.0) + m.update() + return m.get_value("land_surface_water__runoff_depth") * 1000.0 + + +def enkf_update(state, q_sim, y_obs, krig_var, rng): + R = HARDCODED_R if HARDCODED_R is not None else max( + (0.10 * max(y_obs, 0.0)) ** 2 + 0.001 * krig_var, 1e-6) + R = max(R, 1e-6) + P_yy = max(q_sim * 0.01, 1e-6) + K = P_yy / (P_yy + R) + scale = K * (y_obs - q_sim) / max(abs(q_sim), 1e-6) + return { + "soil_m": max(state["soil_m"] * (1.0 + scale), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + scale), 1e-6), + "nash0": max(state["nash0"] + state["nash0"] * scale, 0.0), + "nash1": max(state["nash1"] + state["nash1"] * scale, 0.0), + } + + +def run_phase1_da(cfg_path, df_test, obs_dict, var_dict, rng): + """Run DA through test period; return snapshot dict {timestamp -> state}.""" + print(f" Phase 1: DA analysis through test period...") + model = make_model(cfg_path) + snapshots = {} + for _, row in df_test.iterrows(): + date_str = row["date"] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + q_sim = run_model_step(model, P, E) + y_obs = obs_dict.get(date_str, np.nan) + if not np.isnan(y_obs): + krig_var = var_dict.get(date_str, 0.07) + state = get_state(model) + updated = enkf_update(state, q_sim, y_obs, krig_var, rng) + set_state(model, updated) + date_ts = pd.Timestamp(date_str) + if HELENE_START <= date_ts <= HELENE_END: + snapshots[date_ts] = get_state(model) + model.finalize() + print(f" Phase 1 done: {len(snapshots)} issue-time snapshots") + return snapshots + + +def main(): + global CAT_ID, OBS_FILE, CFE_CONFIG_FILE, TEST_FORCING_FILE + global OUT_DIR, HARDCODED_R, bmi_cfe + + parser = argparse.ArgumentParser() + parser.add_argument("--cat-id", required=True) + parser.add_argument("--forcing-dir", required=True) + parser.add_argument("--obs-dir", required=True) + parser.add_argument("--cfe-dir", required=True) + parser.add_argument("--config-file", required=True) + parser.add_argument("--out-dir", required=True) + parser.add_argument("--test-forcing-dir1", default=None) + parser.add_argument("--test-forcing-dir2", default=None) + parser.add_argument("--hardcoded-r", type=float, default=None, + help="Fixed R value for EnKF (e.g. 0.07). " + "Omit to use dynamic Vrugt R: R=(0.10*y)^2 + 0.001*krig_var") + parser.add_argument("--n-forcing", type=int, default=N_FORCING, + help="Number of met forcing draws (default 30)") + parser.add_argument("--n-hydro", type=int, default=N_HYDRO, + help="Number of hydro-state draws (default 20)") + args = parser.parse_args() + + CAT_ID = args.cat_id + OBS_FILE = os.path.join(args.obs_dir, f"{CAT_ID}.csv") + CFE_CONFIG_FILE = args.config_file + HARDCODED_R = args.hardcoded_r + OUT_DIR = Path(args.out_dir) / CAT_ID + OUT_DIR.mkdir(parents=True, exist_ok=True) + + n_fa = args.n_forcing + n_ha = args.n_hydro + n_total = n_fa * n_ha + print(f"[{CAT_ID}] Crossed ensemble: {n_fa} forcing x {n_ha} hydro = {n_total} members") + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + bmi_cfe = _bmi_cfe + + # Reuse combined forcing CSV from run_perturbation_da_on.py if it exists + # (avoids writing a duplicate large file to disk) + da_on_combined = (Path(args.out_dir).parent / "v2_perturbation_da_on" / + CAT_ID / f"{CAT_ID}_nwm_operational_combined.csv") + crossed_combined = OUT_DIR / f"{CAT_ID}_nwm_operational_combined.csv" + + if da_on_combined.exists(): + TEST_FORCING_FILE = str(da_on_combined) + print(f" Reusing combined forcing: {TEST_FORCING_FILE}") + elif crossed_combined.exists(): + TEST_FORCING_FILE = str(crossed_combined) + print(f" Reusing combined forcing: {TEST_FORCING_FILE}") + elif args.test_forcing_dir1 and args.test_forcing_dir2: + f1 = os.path.join(args.test_forcing_dir1, f"{CAT_ID}.csv") + f2 = os.path.join(args.test_forcing_dir2, f"{CAT_ID}.csv") + df1, df2 = pd.read_csv(f1), pd.read_csv(f2) + combined = (pd.concat([df1, df2], ignore_index=True) + .drop_duplicates(subset="time").sort_values("time")) + TEST_FORCING_FILE = str(crossed_combined) + combined.to_csv(TEST_FORCING_FILE, index=False) + print(f" Wrote combined forcing: {TEST_FORCING_FILE}") + else: + print("Need --test-forcing-dir1 and --test-forcing-dir2"); return + + best_params_file = OUT_DIR / f"{CAT_ID}_best_params.json" + if not best_params_file.exists(): + # Also check the DA results dir from run_perturbation_da_on.py + alt = Path(args.out_dir).parent / "v2_perturbation_da_on" / CAT_ID / f"{CAT_ID}_best_params.json" + if alt.exists(): + best_params_file = alt + else: + print(f"Missing best params at {best_params_file}"); return + with open(best_params_file) as f: + best_params = json.load(f)["best_parameters"] + + obs_dict, var_dict = load_obs() + df_forcing = load_test_forcing() + t_mask = (df_forcing["date"] >= TEST_START) & (df_forcing["date"] <= TEST_END) + df_test = df_forcing[t_mask].reset_index(drop=True) + + # Write BMI config once — reused for all 90,600 model instantiations + cfg_path = write_model_config(best_params) + print(f" BMI config written once: {cfg_path}") + + rng = np.random.default_rng(hash(CAT_ID) & 0x7fffffff) + + # Phase 1: get DA analysis snapshots + # Check if a prior run already saved them (skip re-running DA if so) + snap_cache = OUT_DIR / f"{CAT_ID}_da_snapshots.parquet" + if snap_cache.exists(): + print(f" Loading cached DA snapshots from {snap_cache}") + df_snaps = pd.read_parquet(snap_cache) + snapshots = { + pd.Timestamp(row["issue_time"]): { + "soil_m": row["soil_m"], "gw_m": row["gw_m"], + "nash0": row["nash0"], "nash1": row["nash1"], + } + for _, row in df_snaps.iterrows() + } + else: + snapshots = run_phase1_da(cfg_path, df_test, obs_dict, var_dict, rng) + + issue_times = sorted(snapshots.keys()) + t0_to_idx = {t: i for i, t in enumerate(df_test["date"].apply(pd.Timestamp))} + + # ------------------------------------------------------------------ # + # Crossed ensemble: for each issue_time, build n_fa forcing draws and # + # n_ha state draws, then run all n_fa x n_ha combinations. # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Running crossed ensemble ({n_total} members x " + f"{len(issue_times)} issue times)...") + + all_records = [] + all_hydro_draw_records = [] # (issue_time, hydro_draw_j, states) + all_forcing_draw_records = [] # (issue_time, forcing_draw_i, lead, P_pert, E_pert, scales) + + for t_idx, t0 in enumerate(issue_times): + snap = snapshots[t0] + t0_str = t0.strftime("%Y-%m-%d %H:%M:%S") + if t0 not in t0_to_idx: + continue + start_idx = t0_to_idx[t0] + + # Pre-draw n_fa forcing perturbation sequences + forcing_seqs = [] + for i in range(n_fa): + seq = [] + for lead in range(1, FORECAST_HOURS + 1): + fi = start_idx + lead + if fi >= len(df_test): + seq.append((0.0, 0.0)) + all_forcing_draw_records.append({ + "issue_time": t0_str, "forcing_draw_i": i, + "lead_hour": lead, + "P_pert_mmh": 0.0, "E_pert_mmh": 0.0, + "precip_scale": np.nan, "pet_scale": np.nan, + }) + continue + row = df_test.iloc[fi] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + mu_p = -0.5 * PRECIP_SIGMA ** 2 + P_pert = float(P * rng.lognormal(mu_p, PRECIP_SIGMA)) if P > 0 else 0.0 + E_pert = max(float(E * (1.0 + PET_SIGMA * rng.standard_normal())), 0.0) + seq.append((P_pert, E_pert)) + all_forcing_draw_records.append({ + "issue_time": t0_str, "forcing_draw_i": i, + "lead_hour": lead, + "P_pert_mmh": P_pert, "E_pert_mmh": E_pert, + "precip_scale": P_pert / P if P > 0 else np.nan, + "pet_scale": E_pert / E if E > 0 else np.nan, + }) + forcing_seqs.append(seq) + + # Pre-draw n_ha hydro-state perturbations + state_draws = [snap if j == 0 else perturb_state(snap, rng) + for j in range(n_ha)] + + # Record each hydro draw's actual initial state + for j, st in enumerate(state_draws): + all_hydro_draw_records.append({ + "issue_time": t0_str, "hydro_draw_j": j, + "soil_m": st["soil_m"], "gw_m": st["gw_m"], + "nash0": st["nash0"], "nash1": st["nash1"], + }) + + # q_matrix shape: (FORECAST_HOURS, n_fa, n_ha) + q_matrix = np.full((FORECAST_HOURS, n_fa, n_ha), np.nan) + + for i, f_seq in enumerate(forcing_seqs): + for j, state_j in enumerate(state_draws): + m = make_model(cfg_path) + set_state(m, state_j) + for lead in range(1, FORECAST_HOURS + 1): + P_pert, E_pert = f_seq[lead - 1] + q = run_model_step(m, P_pert, E_pert) + q_matrix[lead - 1, i, j] = q + m.finalize() + + # Flatten 600-member matrix into rows + for lead in range(1, FORECAST_HOURS + 1): + row_data = {"issue_time": t0_str, "lead_hour": lead} + mem_idx = 0 + for i in range(n_fa): + for j in range(n_ha): + row_data[f"member_{mem_idx:04d}"] = q_matrix[lead - 1, i, j] + mem_idx += 1 + all_records.append(row_data) + + if (t_idx + 1) % 10 == 0 or (t_idx + 1) == len(issue_times): + print(f" {t_idx + 1}/{len(issue_times)} issue times complete") + + df_out = pd.DataFrame(all_records) + out_path = OUT_DIR / f"{CAT_ID}_crossed_ensemble.parquet" + df_out.to_parquet(out_path, index=False) + print(f"Saved: {out_path}") + print(f" Shape: {df_out.shape} " + f"({df_out['issue_time'].nunique()} issue times x " + f"{FORECAST_HOURS} leads x {n_total} members)") + + # ------------------------------------------------------------------ # + # Provenance files — enable full ensemble traceability and restart # + # ------------------------------------------------------------------ # + + # 1. DA analysis snapshots — unperturbed state at each issue_time. + # Load any row and set_state() to restart from that init time. + snap_records = [ + {"issue_time": t.strftime("%Y-%m-%d %H:%M:%S"), **s} + for t, s in snapshots.items() + ] + snap_path = OUT_DIR / f"{CAT_ID}_da_snapshots.parquet" + pd.DataFrame(snap_records).to_parquet(snap_path, index=False) + print(f"Saved: {snap_path} ({len(snap_records)} snapshots)") + + # 2. Member manifest — decoder ring: member_col -> (forcing_draw_i, hydro_draw_j). + # member_k uses forcing draw k//n_ha and hydro draw k%n_ha. + manifest = [ + {"member": f"member_{k:04d}", "member_idx": k, + "forcing_draw_i": k // n_ha, "hydro_draw_j": k % n_ha} + for k in range(n_total) + ] + manifest_path = OUT_DIR / f"{CAT_ID}_member_manifest.csv" + pd.DataFrame(manifest).to_csv(manifest_path, index=False) + print(f"Saved: {manifest_path} ({n_total} members)") + + # 3. Hydro draw states — actual perturbed initial state for each of + # n_ha draws at every issue_time. Links member -> initial conditions. + hydro_path = OUT_DIR / f"{CAT_ID}_hydro_draw_states.parquet" + pd.DataFrame(all_hydro_draw_records).to_parquet(hydro_path, index=False) + print(f"Saved: {hydro_path} ({len(all_hydro_draw_records)} rows)") + + # 4. Forcing draw sequences — actual P/PET perturbation values for each + # of n_fa draws at every (issue_time, lead_hour). Links member -> forcing. + forcing_path = OUT_DIR / f"{CAT_ID}_forcing_draw_sequences.parquet" + pd.DataFrame(all_forcing_draw_records).to_parquet(forcing_path, index=False) + print(f"Saved: {forcing_path} ({len(all_forcing_draw_records)} rows)") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/run_lead_time_forecast_sweep.py b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/run_lead_time_forecast_sweep.py new file mode 100644 index 00000000..19f3b4d5 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/run_lead_time_forecast_sweep.py @@ -0,0 +1,508 @@ +""" +Forecast lead-time evaluation. + +For each issue-time t0 in a sampling schedule across the test period, fork the +ensemble at t0 and free-run 18 hours of forecast. Two parallel trajectories are +maintained throughout the test period: + + da: Production setup. DA on, forcing perturbed, process noise on. + openloop: No DA from t=0. Forcing perturbed, process noise on. + +At each issue time t0, the state of each scenario's ensemble is copied into a +forecast ensemble that free-runs 18 hours with: + - DA off (no observations consumed during the forecast window) + - process noise off (no obs to collapse toward, so no anti-collapse needed) + - forcing perturbed (lognormal precip, Gaussian PET — proxy for forecast + precip uncertainty in lieu of actual HEFS forecasts) + +Issue-time schedule: + - Base cadence: every --base-step-h hours across the full test period + (default 6h → ~1600 issue times for the year) + - Densified to hourly across --dense-start..--dense-end if provided + (default: 2024-09-24 → 2024-09-28, the Helene window) + +Outputs (stacked across all issue times): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv + Columns: issue_time, lead_hour, valid_time, member_00 .. member_19 + (units: mm/h) + +Post-processing pipeline (separate scripts, run after this): + route_lead_time_forecasts.py — pushes each (t0, lead_hour) forecast through + T-route to the gauge + plot_lead_time_decay.py — error-vs-lead-time curve, DA vs open-loop + +Imports EnKFAssimilator from calibrate_catchment_cfe_da_v2.py so the DA math +matches production exactly. +""" + +import argparse +import os +import sys +import json +import importlib.util +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +SPINUP_START = "2023-02-01 00:00:00" +SPINUP_END = "2023-09-30 23:00:00" +TEST_START = "2023-10-01 00:00:00" +TEST_END = "2024-10-31 23:00:00" + +# Forecast lead time (hours after each issue time) +FORECAST_LEAD_HOURS = 18 + +# Default densification window — Hurricane Helene +DEFAULT_DENSE_START = "2024-09-24 00:00:00" +DEFAULT_DENSE_END = "2024-09-28 23:00:00" + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(test_forcing_file): + df = pd.read_csv(test_forcing_file) + df['date'] = pd.to_datetime(df['time']).dt.strftime('%Y-%m-%d %H:%M:%S') + df['total_precipitation'] = df['APCP_surface'] * 3600.0 + df['potential_evaporation'] = priestley_taylor_pet( + df['DSWRF_surface'].values, df['TMP_2maboveground'].values + ) + return df + + +def import_enkf_class(script_path): + """Load EnKFAssimilator from the production script without triggering its main().""" + spec = importlib.util.spec_from_file_location("prod_v2", script_path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod.EnKFAssimilator + + +# ---------------- State snapshot / restore ----------------------------------- +# +# The EnKF only ever touches these 4 states, so they're the only states that +# diverge between members. Copy these from src → dst to fork the ensemble. + +def snapshot_states(models): + return [ + { + "soil_m": float(m.soil_reservoir["storage_m"]), + "gw_m": float(m.gw_reservoir["storage_m"]), + "nash0_m": float(m.nash_storage[0]), + "nash1_m": float(m.nash_storage[1]), + } + for m in models + ] + + +def restore_states(models, snapshot): + for m, s in zip(models, snapshot): + m.soil_reservoir["storage_m"] = s["soil_m"] + m.gw_reservoir["storage_m"] = s["gw_m"] + m.nash_storage[0] = s["nash0_m"] + m.nash_storage[1] = s["nash1_m"] + + +# ---------------- Issue-time schedule ---------------------------------------- + +def build_issue_time_schedule(dates_list, base_step_h, dense_start, dense_end): + """Return a sorted list of issue-time strings sampled from dates_list. + + Every base_step_h hours across the full test period, plus every hour inside + [dense_start, dense_end] if those are provided. + """ + dates_dt = pd.to_datetime(dates_list) + selected = set() + + # Base cadence (every base_step_h hours from the first date) + base_mask = (np.arange(len(dates_dt)) % base_step_h == 0) + for d in dates_dt[base_mask]: + selected.add(d.strftime('%Y-%m-%d %H:%M:%S')) + + # Densified hourly window + if dense_start is not None and dense_end is not None: + ds = pd.Timestamp(dense_start) + de = pd.Timestamp(dense_end) + dense_mask = (dates_dt >= ds) & (dates_dt <= de) + for d in dates_dt[dense_mask]: + selected.add(d.strftime('%Y-%m-%d %H:%M:%S')) + + return sorted(selected) + + +# ---------------- CFE helpers ------------------------------------------------ + +def build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, enkf, + apply_init_perturbation=True, member_zero_clean=True): + """Build N CFE BMI instances with optional init-state perturbation.""" + def custom_load_forcing(self_cfe): + df = load_test_forcing(test_forcing_file) + self_cfe.forcing_data = df.rename(columns={"date": "time"}) + + models = [] + for i in range(N): + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.load_forcing_file = custom_load_forcing.__get__(m) + m.initialize() + if apply_init_perturbation and not (member_zero_clean and i == 0): + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + r = enkf.init_state_perturb_frac + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1 + r * enkf.rng.standard_normal()), 1e-6, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1 + r * enkf.rng.standard_normal()), 1e-6, gw_max)) + jitter = max(sm_max * 1e-6, 1e-9) + m.nash_storage[0] = max(n00 + jitter * enkf.rng.standard_normal(), 0.0) + m.nash_storage[1] = max(n10 + jitter * enkf.rng.standard_normal(), 0.0) + models.append(m) + return models + + +def step_ensemble(models, p_arr, e_arr): + """Advance each member one hour with member-specific (P, PET). Returns Q (mm/h).""" + q = np.empty(len(models), dtype=float) + for i, m in enumerate(models): + m.set_value('atmosphere_water__time_integral_of_precipitation_mass_flux', + float(p_arr[i]) / 1000.0) + m.set_value('water_potential_evaporation_flux', + float(e_arr[i]) / 1000.0 / 3600.0) + m.update() + q[i] = m.get_value('land_surface_water__runoff_depth') * 1000.0 # m/h → mm/h + return q + + +def do_sanity_check(prod_models, fcst_models, dates_list, forcing_by_date, h, N): + """One-shot verification that the 4-state snapshot is sufficient to reproduce + next-hour forecast Q. + + At main-loop hour h (after prod's hour-h step has completed including DA + + process noise), snapshot prod, restore into fcst, then step BOTH ensembles + one hour using identical unperturbed forcing. If snapshot/restore captures + everything that matters, prod's and fcst's hour-(h+1) Q should match member- + by-member to machine precision. + + Note: this mutates prod's state (h+1 step with non-perturbed forcing and no + DA), so the caller should exit the script right after. + """ + if h + 1 >= len(dates_list): + print("[sanity-check] not enough horizon for h+1 step — skipping.") + return False + snap = snapshot_states(prod_models) + restore_states(fcst_models, snap) + + next_date = dates_list[h + 1] + p_next, e_next = forcing_by_date[next_date] + p_uniform = np.full(N, p_next, dtype=float) + e_uniform = np.full(N, e_next, dtype=float) + + q_prod = step_ensemble(prod_models, p_uniform, e_uniform) + q_fcst = step_ensemble(fcst_models, p_uniform, e_uniform) + + diff = q_prod - q_fcst + max_abs = float(np.max(np.abs(diff))) + print("[sanity-check] forecast Q at h+1 with identical unperturbed forcing:") + print(f" prod first 5 members: {q_prod[:5]}") + print(f" fcst first 5 members: {q_fcst[:5]}") + print(f" max |prod - fcst|: {max_abs:.3e} mm/h") + if max_abs < 1e-9: + print(" OK: snapshot/restore preserves state to machine precision.") + return True + print(" WARN: nonzero divergence — likely a state not in the 4-state snapshot.") + return False + + +def run_forecast(fcst_models, enkf_fcst, dates_list, forcing_by_date, + t0_idx, n_lead): + """Run an n_lead-hour free-forecast starting from the current state of fcst_models. + + Forcing is perturbed per member (lognormal precip + Gaussian PET, same as + production). No DA. No process noise. + + Returns: (lead_hours, valid_times, q_matrix shape (n_lead, N)). + """ + N = len(fcst_models) + q_matrix = np.full((n_lead, N), np.nan, dtype=float) + lead_hours = [] + valid_times = [] + for lead in range(1, n_lead + 1): + idx = t0_idx + lead + if idx >= len(dates_list): + break + valid_date = dates_list[idx] + p, e = forcing_by_date[valid_date] + p_arr, e_arr = enkf_fcst.perturb_forcing(p, e) + q = step_ensemble(fcst_models, p_arr, e_arr) + q_matrix[lead - 1, :] = q + lead_hours.append(lead) + valid_times.append(valid_date) + return lead_hours, valid_times, q_matrix + + +# ---------------- Main run --------------------------------------------------- + +def run(args, bmi_cfe, EnKFAssimilator): + cat_id = args.cat_id + out_dir = Path(args.out_dir) / cat_id + out_dir.mkdir(parents=True, exist_ok=True) + + # Combined test forcing (same idiom as the per-member script) + f1 = os.path.join(args.test_forcing_dir1, f'{cat_id}.csv') + f2 = os.path.join(args.test_forcing_dir2, f'{cat_id}.csv') + df1 = pd.read_csv(f1) + df2 = pd.read_csv(f2) + combined = pd.concat([df1, df2], ignore_index=True) + combined = combined.drop_duplicates(subset='time').sort_values('time') + test_forcing_file = str(out_dir / f'{cat_id}_nwm_operational_combined.csv') + combined.to_csv(test_forcing_file, index=False) + + # Pre-staged calibrated params + best_params_file = out_dir / f'{cat_id}_best_params.json' + if not best_params_file.exists(): + print(f"Missing {best_params_file}. Pre-stage from Run 3.") + return + with open(best_params_file) as f: + best = json.load(f)['best_parameters'] + + # Build temp CFE config with calibrated params + with open(args.config_file) as f: + cfg = json.load(f) + cfg['forcing_file'] = test_forcing_file + cfg['soil_params']['bb'] = best['bb'] + cfg['soil_params']['smcmax'] = best['smcmax'] + cfg['soil_params']['satdk'] = best['satdk'] + cfg['slop'] = best['slop'] + cfg['max_gw_storage'] = best['max_gw_storage'] + cfg['expon'] = best['expon'] + cfg['Cgw'] = best['Cgw'] + cfg['K_lf'] = best['K_lf'] + cfg['K_nash'] = best['K_nash'] + cfg['partition_scheme'] = "Schaake" if best['scheme'] <= 0.5 else "Xinanjiang" + tmp_cfg = str(out_dir / f'{cat_id}_bmi_config_temp_leadtime.json') + with open(tmp_cfg, 'w') as f: + json.dump(cfg, f) + + # Three EnKF instances: + # enkf_da — runs the DA-on trajectory; consumes obs each hour + # enkf_ol — runs the open-loop trajectory; never calls update_states + # enkf_fcst — used only inside the 18-hour free-runs (perturb_forcing only) + # All three share defaults with production. Distinct RNG seeds keep them + # independent so the openloop and forecast spreads are not coupled to DA. + obs_file = os.path.join(args.obs_dir, f'{cat_id}.csv') + seed_base = args.rng_seed if args.rng_seed is not None else 0 + enkf_da = EnKFAssimilator( + n_members=args.enkf_members, obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, use_vrugt_r=(not args.no_vrugt_r), + vrugt_alpha=args.vrugt_alpha, vrugt_scale=args.vrugt_scale, + rng_seed=seed_base if args.rng_seed is not None else None, + ) + enkf_ol = EnKFAssimilator( + n_members=args.enkf_members, obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, use_vrugt_r=(not args.no_vrugt_r), + vrugt_alpha=args.vrugt_alpha, vrugt_scale=args.vrugt_scale, + rng_seed=(seed_base + 1) if args.rng_seed is not None else None, + ) + enkf_fcst = EnKFAssimilator( + n_members=args.enkf_members, obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, use_vrugt_r=(not args.no_vrugt_r), + vrugt_alpha=args.vrugt_alpha, vrugt_scale=args.vrugt_scale, + rng_seed=(seed_base + 2) if args.rng_seed is not None else None, + ) + # Optional R override — replace every per-hour obs variance with a constant. + # Matches DualEarth/new_EnKF.py (R=0.07). Set via --hardcoded-r on the CLI. + # Effect: at storm peak ≈ existing Vrugt R, at low flow ≈ 15-17× larger → + # near-zero gain at low flow where kriging obs is noisy. + if args.hardcoded_r is not None: + for enkf in (enkf_da, enkf_ol, enkf_fcst): + for date in enkf.obs_var_dict: + enkf.obs_var_dict[date] = args.hardcoded_r + print(f"[lead-time] R hardcoded to {args.hardcoded_r} mm^2/h^2 " + f"(overrides Vrugt + kriging variance formula)") + + N = enkf_da.n_members + print(f"[lead-time] {cat_id} | N={N} | DA on/off both run | " + f"forecast lead = {FORECAST_LEAD_HOURS}h") + + # Build three ensembles (init-perturbed except the forecast ensemble, whose + # state will be overwritten at each issue time before each free-run) + prod_models = build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, enkf_da) + openloop_models = build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, enkf_ol) + fcst_models = build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, + enkf_fcst, apply_init_perturbation=False) + + # ----- Spinup with perturbed forcing on both trajectories ----- + df = load_test_forcing(test_forcing_file) + sp_mask = (df['date'] >= SPINUP_START) & (df['date'] <= SPINUP_END) + df_sp = df[sp_mask] + print(f"[lead-time] spinup: {len(df_sp)} hours") + for p, e in zip(df_sp['total_precipitation'], df_sp['potential_evaporation']): + p_da, e_da = enkf_da.perturb_forcing(p, e) + step_ensemble(prod_models, p_da, e_da) + enkf_da.add_process_noise(prod_models) + p_ol, e_ol = enkf_ol.perturb_forcing(p, e) + step_ensemble(openloop_models, p_ol, e_ol) + enkf_ol.add_process_noise(openloop_models) + + # ----- Test period: step both, fork forecasts at issue times ----- + t_mask = (df['date'] >= TEST_START) & (df['date'] <= TEST_END) + df_test = df[t_mask].reset_index(drop=True) + dates_list = list(df_test['date'].values) + forcing_by_date = dict(zip( + df_test['date'].values, + zip(df_test['total_precipitation'].values, + df_test['potential_evaporation'].values))) + + issue_times = build_issue_time_schedule( + dates_list, args.base_step_h, + args.dense_start, args.dense_end) + issue_set = set(issue_times) + print(f"[lead-time] test period: {len(dates_list)} hours | " + f"{len(issue_times)} issue times " + f"(base={args.base_step_h}h, dense={args.dense_start or 'none'}" + f"..{args.dense_end or 'none'})") + + # Pre-build the date→index map so forecast windows are fast to look up + date_to_idx = {d: i for i, d in enumerate(dates_list)} + + da_rows = [] # (issue_time, lead_hour, valid_time, *member_values) + ol_rows = [] + sanity_done = False + + for h, current_date in enumerate(dates_list): + p, e = forcing_by_date[current_date] + + # --- DA trajectory step --- + p_da, e_da = enkf_da.perturb_forcing(p, e) + q_da = step_ensemble(prod_models, p_da, e_da) + enkf_da.update_states(prod_models, current_date, q_da) + enkf_da.add_process_noise(prod_models) + + # --- Open-loop trajectory step (no DA) --- + p_ol, e_ol = enkf_ol.perturb_forcing(p, e) + q_ol = step_ensemble(openloop_models, p_ol, e_ol) + enkf_ol.add_process_noise(openloop_models) + + # --- One-shot sanity check at the first scheduled issue time --- + if args.sanity_check and not sanity_done and current_date in issue_set: + ok = do_sanity_check( + prod_models, fcst_models, dates_list, forcing_by_date, h, N) + sanity_done = True + print(f"[sanity-check] exiting (re-run without --sanity-check for " + f"the full sweep). Result: {'OK' if ok else 'FAIL'}") + for m in prod_models + openloop_models + fcst_models: + m.finalize() + sys.exit(0 if ok else 1) + + # --- Fork forecasts at scheduled issue times --- + if current_date in issue_set: + for scenario_label, src_models, rows_acc in [ + ('da', prod_models, da_rows), + ('openloop', openloop_models, ol_rows), + ]: + snap = snapshot_states(src_models) + restore_states(fcst_models, snap) + leads, valids, qm = run_forecast( + fcst_models, enkf_fcst, + dates_list, forcing_by_date, + t0_idx=h, n_lead=FORECAST_LEAD_HOURS) + for k, (lead, vt) in enumerate(zip(leads, valids)): + rows_acc.append((current_date, lead, vt, *qm[k, :])) + + for m in prod_models + openloop_models + fcst_models: + m.finalize() + + # ----- Save forecast CSVs ----- + cols = ['issue_time', 'lead_hour', 'valid_time'] + [f'member_{i:02d}' for i in range(N)] + da_path = out_dir / f'{cat_id}_lead_time_forecasts_da.csv' + ol_path = out_dir / f'{cat_id}_lead_time_forecasts_openloop.csv' + pd.DataFrame(da_rows, columns=cols).to_csv(da_path, index=False) + pd.DataFrame(ol_rows, columns=cols).to_csv(ol_path, index=False) + print(f"[lead-time] saved {da_path}") + print(f"[lead-time] saved {ol_path}") + + # Issue-time schedule (small, useful for downstream scripts) + sched_path = out_dir / f'{cat_id}_lead_time_issue_times.csv' + pd.DataFrame({'issue_time': issue_times}).to_csv(sched_path, index=False) + print(f"[lead-time] saved {sched_path}") + + print(f"[lead-time] DA assimilator: updates={enkf_da.n_updates} | " + f"avg Pyy={enkf_da.avg_pyy:.6f} | " + f"mass lost={enkf_da.total_overflow_lost_mm:.3f} mm") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', required=True) + parser.add_argument('--forcing-dir', required=True) + parser.add_argument('--obs-dir', required=True) + parser.add_argument('--cfe-dir', required=True) + parser.add_argument('--config-file', required=True) + parser.add_argument('--param-bounds', required=True, + help='Kept for symmetry; not read in this script') + parser.add_argument('--out-dir', required=True) + parser.add_argument('--test-forcing-dir1', required=True) + parser.add_argument('--test-forcing-dir2', required=True) + parser.add_argument('--enkf-members', type=int, default=20) + parser.add_argument('--enkf-obs-error-std', type=float, default=0.05) + parser.add_argument('--no-vrugt-r', action='store_true') + parser.add_argument('--vrugt-alpha', type=float, default=0.10) + parser.add_argument('--vrugt-scale', type=float, default=0.001) + parser.add_argument('--rng-seed', type=int, default=None) + parser.add_argument('--base-step-h', type=int, default=6, + help='Issue-time cadence in hours across the full test ' + 'period (default 6h)') + parser.add_argument('--dense-start', type=str, default=DEFAULT_DENSE_START, + help='Start of hourly-densification window ' + '(default 2024-09-24 00:00:00 — Helene)') + parser.add_argument('--dense-end', type=str, default=DEFAULT_DENSE_END, + help='End of hourly-densification window ' + '(default 2024-09-28 23:00:00 — Helene)') + parser.add_argument('--sanity-check', action='store_true', + help='At the first issue time, fork fcst_models and step ' + 'both prod and fcst one hour with identical unperturbed ' + 'forcing. Confirms the 4-state snapshot is sufficient ' + 'to reproduce next-hour Q to machine precision, then ' + 'exits. Use before kicking off the full sweep.') + parser.add_argument('--hardcoded-r', type=float, default=None, + help='If set, override the Vrugt + kriging-variance R formula ' + 'and use this constant value (mm^2/h^2) at every hour. ' + 'Matches the DualEarth/new_EnKF.py reference (R=0.07). ' + 'Suppresses DA at low flow where kriging is noisy while ' + 'preserving DA strength at storm peaks. Applied to all ' + 'three EnKF instances (da, openloop, fcst) for consistency.') + parser.add_argument('--prod-script', default=None, + help='Path to calibrate_catchment_cfe_da_v2.py for ' + 'importing EnKFAssimilator. Defaults to next-to-this-file.') + args = parser.parse_args() + + here = os.path.dirname(os.path.abspath(__file__)) + prod_script = args.prod_script or os.path.join(here, 'calibrate_catchment_cfe_da_v2.py') + if not os.path.exists(prod_script): + raise FileNotFoundError( + f"Could not find production script at {prod_script}. " + f"Pass --prod-script to override.") + EnKFAssimilator = import_enkf_class(prod_script) + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + + run(args, _bmi_cfe, EnKFAssimilator) + + +if __name__ == '__main__': + main() diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/run_perturbation_da_on.py b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/run_perturbation_da_on.py new file mode 100644 index 00000000..b9ef7655 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/run_perturbation_da_on.py @@ -0,0 +1,392 @@ +""" +run_perturbation_da_on.py — 2a/2b: ensemble spread with DA ON. + +Runs two separated perturbation arms for the Helene window (Sep 24-30), +with DA actively running during the analysis period: + + ARM A — forcing only (30 members): + DA analyzes states using Qkrig at every hour (DA ON). + At each issue_time: fork 30 members with ONLY met forcing perturbed + (different precip/PET draw per member). Hydro states = DA analysis + mean (same for all 30). Free-run 18 hours, no further DA. + + ARM B — hydro states only (20 members): + Same DA analysis. At each issue_time: fork 20 members with ONLY + hydro states perturbed (multiplicative noise on soil/GW/Nash from + the DA analysis ensemble). Forcing = deterministic. Free-run 18 hours. + +These two arms isolate "where does forecast spread come from WITH DA?" +— figures for the assimilation script section. + +Output (per catchment, per arm): + //_da_forcing_arm.csv + //_da_hydro_arm.csv + Columns: issue_time, lead_hour, member_00..member_29 (forcing arm) + issue_time, lead_hour, member_00..member_19 (hydro arm) + +Usage (run for each catchment, both arms): + python3 run_perturbation_da_on.py \\ + --cat-id cat-1016300 \\ + --forcing-dir \\ + --obs-dir /mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance \\ + --cfe-dir /mnt/disk2/suma_helen_poster/cfe_py \\ + --config-file /mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json \\ + --param-bounds /mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json \\ + --out-dir /mnt/disk2/suma_helen_poster/da_results/v2_perturbation_da_on \\ + --test-forcing-dir1 \\ + --test-forcing-dir2 \\ + --hardcoded-r 0.07 +""" + +import argparse +import os +import sys +import json +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +N_FORCING = 30 # members in forcing arm +N_HYDRO = 20 # members in hydro-state arm + +# Analysis period — run DA through entire test period +TEST_START = "2024-08-24 00:00:00" # 1 month spinup before Helene window +TEST_END = "2024-10-31 23:00:00" + +# Issue times for the Helene window (hourly, Sep 24-30) +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-30 06:00:00") +FORECAST_HOURS = 18 + +# Perturbation magnitudes (same as production) +PRECIP_SIGMA = 0.15 +PET_SIGMA = 0.10 +STATE_FRAC = 0.05 # hydro-state perturbation fraction + +# Set at runtime +CAT_ID = None +OBS_FILE = None +CFE_CONFIG_FILE = None +TEST_FORCING_FILE = None +OUT_DIR = None +HARDCODED_R = None +bmi_cfe = None + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(): + df = pd.read_csv(TEST_FORCING_FILE) + df["date"] = pd.to_datetime(df["time"]).dt.strftime("%Y-%m-%d %H:%M:%S") + df["total_precipitation"] = df["APCP_surface"] * 3600.0 + df["potential_evaporation"] = priestley_taylor_pet( + df["DSWRF_surface"].values, df["TMP_2maboveground"].values) + return df + + +def load_obs(): + df = pd.read_csv(OBS_FILE) + date_col = next((c for c in df.columns + if c.lower() in ("date", "time", "datetime", "timestamp")), None) + if date_col is None: + raise ValueError(f"No date column in {OBS_FILE}. Columns: {df.columns.tolist()}") + q_col = next((c for c in df.columns if "qkrig" in c.lower() and "var" not in c.lower()), None) + if q_col is None: + raise ValueError(f"No qkrig column in {OBS_FILE}. Columns: {df.columns.tolist()}") + var_col = next((c for c in df.columns if "var" in c.lower()), None) + df[date_col] = pd.to_datetime(df[date_col]) + df = df.sort_values(date_col).reset_index(drop=True) + date_strs = df[date_col].dt.strftime("%Y-%m-%d %H:%M:%S") + obs_dict = dict(zip(date_strs, df[q_col])) + var_vals = df[var_col] if var_col else pd.Series([0.07] * len(df)) + var_dict = dict(zip(date_strs, var_vals)) + print(f" Obs: date='{date_col}', q='{q_col}', var='{var_col}'") + return obs_dict, var_dict + + +def make_model(best_params): + with open(CFE_CONFIG_FILE) as f: + cfg = json.load(f) + cfg["forcing_file"] = TEST_FORCING_FILE + cfg["soil_params"]["bb"] = best_params["bb"] + cfg["soil_params"]["smcmax"] = best_params["smcmax"] + cfg["soil_params"]["satdk"] = best_params["satdk"] + cfg["slop"] = best_params["slop"] + cfg["max_gw_storage"] = best_params["max_gw_storage"] + cfg["expon"] = best_params["expon"] + cfg["Cgw"] = best_params["Cgw"] + cfg["K_lf"] = best_params["K_lf"] + cfg["K_nash"] = best_params["K_nash"] + cfg["partition_scheme"] = ("Schaake" if best_params["scheme"] <= 0.5 + else "Xinanjiang") + tmp_cfg = str(OUT_DIR / f"{CAT_ID}_bmi_config_temp_da_on.json") + with open(tmp_cfg, "w") as f: + json.dump(cfg, f) + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.initialize() + return m + + +def get_state(m): + return { + "soil_m": m.soil_reservoir["storage_m"], + "gw_m": m.gw_reservoir["storage_m"], + "nash0": float(m.nash_storage[0]), + "nash1": float(m.nash_storage[1]), + } + + +def set_state(m, state): + m.soil_reservoir["storage_m"] = state["soil_m"] + m.gw_reservoir["storage_m"] = state["gw_m"] + m.nash_storage[0] = state["nash0"] + m.nash_storage[1] = state["nash1"] + + +def perturb_state(state, rng, frac=STATE_FRAC): + sm_max = 1.0 # placeholder — will be clipped in set_state; actual max enforced by model + out = { + "soil_m": max(state["soil_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "nash0": max(state["nash0"] * (1.0 + frac * rng.standard_normal()), 0.0), + "nash1": max(state["nash1"] * (1.0 + frac * rng.standard_normal()), 0.0), + } + return out + + +def run_model_step(m, precip_mmh, pet_mmh): + m.set_value("atmosphere_water__time_integral_of_precipitation_mass_flux", + precip_mmh / 1000.0) + m.set_value("water_potential_evaporation_flux", + pet_mmh / 1000.0 / 3600.0) + m.update() + return m.get_value("land_surface_water__runoff_depth") * 1000.0 # m/h -> mm/h + + +def enkf_update(state, q_sim, y_obs, krig_var, rng, n_members=1): + """Simple scalar EnKF update on the mean analysis state.""" + if HARDCODED_R is not None: + R = HARDCODED_R + else: + alpha = 0.10 + R = (alpha * max(y_obs, 0.0)) ** 2 + 0.001 * krig_var + R = max(R, 1e-6) + P_yy = max(q_sim * 0.01, 1e-6) # rough prior variance on Q + K = P_yy / (P_yy + R) + innovation = y_obs - q_sim + # Apply correction to each state proportionally (simple scalar gain) + scale = K * innovation / max(abs(q_sim), 1e-6) + updated = { + "soil_m": max(state["soil_m"] * (1.0 + scale), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + scale), 1e-6), + "nash0": max(state["nash0"] + state["nash0"] * scale, 0.0), + "nash1": max(state["nash1"] + state["nash1"] * scale, 0.0), + } + return updated + + +def main(): + global CAT_ID, OBS_FILE, CFE_CONFIG_FILE, TEST_FORCING_FILE + global OUT_DIR, HARDCODED_R, bmi_cfe + + parser = argparse.ArgumentParser() + parser.add_argument("--cat-id", required=True) + parser.add_argument("--forcing-dir", required=True) + parser.add_argument("--obs-dir", required=True) + parser.add_argument("--cfe-dir", required=True) + parser.add_argument("--config-file", required=True) + parser.add_argument("--param-bounds", required=True) + parser.add_argument("--out-dir", required=True) + parser.add_argument("--test-forcing-dir1", default=None) + parser.add_argument("--test-forcing-dir2", default=None) + parser.add_argument("--hardcoded-r", type=float, default=0.07) + args = parser.parse_args() + + CAT_ID = args.cat_id + OBS_FILE = os.path.join(args.obs_dir, f"{CAT_ID}.csv") + CFE_CONFIG_FILE = args.config_file + HARDCODED_R = args.hardcoded_r + OUT_DIR = Path(args.out_dir) / CAT_ID + OUT_DIR.mkdir(parents=True, exist_ok=True) + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + bmi_cfe = _bmi_cfe + + # Build combined test forcing + if args.test_forcing_dir1 and args.test_forcing_dir2: + f1 = os.path.join(args.test_forcing_dir1, f"{CAT_ID}.csv") + f2 = os.path.join(args.test_forcing_dir2, f"{CAT_ID}.csv") + df1, df2 = pd.read_csv(f1), pd.read_csv(f2) + combined = (pd.concat([df1, df2], ignore_index=True) + .drop_duplicates(subset="time").sort_values("time")) + TEST_FORCING_FILE = str(OUT_DIR / f"{CAT_ID}_nwm_operational_combined.csv") + combined.to_csv(TEST_FORCING_FILE, index=False) + else: + print("Need --test-forcing-dir1 and --test-forcing-dir2"); return + + best_params_file = OUT_DIR / f"{CAT_ID}_best_params.json" + if not best_params_file.exists(): + print(f"Missing best params at {best_params_file}"); return + with open(best_params_file) as f: + best_params = json.load(f)["best_parameters"] + + obs_dict, var_dict = load_obs() + df_forcing = load_test_forcing() + t_mask = (df_forcing["date"] >= TEST_START) & (df_forcing["date"] <= TEST_END) + df_test = df_forcing[t_mask].reset_index(drop=True) + + rng = np.random.default_rng(hash(CAT_ID) & 0x7fffffff) + + # ------------------------------------------------------------------ # + # Phase 1: Run DA through test period, snapshot analysis state at # + # each Helene issue time # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Phase 1: DA analysis through test period...") + model = make_model(best_params) + + helene_snapshots = {} # issue_time -> analysis state dict + + for h, row in df_test.iterrows(): + date_str = row["date"] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + q_sim = run_model_step(model, P, E) + + date_ts = pd.Timestamp(date_str) + + # DA update if obs available + y_obs = obs_dict.get(date_str, np.nan) + if not np.isnan(y_obs): + krig_var = var_dict.get(date_str, 0.07) + state = get_state(model) + updated = enkf_update(state, q_sim, y_obs, krig_var, rng) + set_state(model, updated) + + # Snapshot if this is a Helene issue time + if HELENE_START <= date_ts <= HELENE_END: + helene_snapshots[date_ts] = get_state(model) + + model.finalize() + print(f" Snapshots saved: {len(helene_snapshots)} issue times") + + # ------------------------------------------------------------------ # + # Phase 2A: Forcing arm — 30 members, fixed analyzed state, varying # + # met forcing # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Phase 2A: Forcing arm ({N_FORCING} members)...") + forcing_records = [] + + issue_times = sorted(helene_snapshots.keys()) + t0_to_idx = {t: i for i, t in enumerate(df_test["date"].apply(pd.Timestamp))} + + for t0 in issue_times: + snap = helene_snapshots[t0] + t0_str = t0.strftime("%Y-%m-%d %H:%M:%S") + + # Find index in df_test + if t0 not in t0_to_idx: + continue + start_idx = t0_to_idx[t0] + + q_matrix = np.full((FORECAST_HOURS, N_FORCING), np.nan) + + for mem in range(N_FORCING): + m = make_model(best_params) + set_state(m, snap) # all members start from same DA analysis state + + for lead in range(1, FORECAST_HOURS + 1): + fi = start_idx + lead + if fi >= len(df_test): + break + row = df_test.iloc[fi] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + # Perturb forcing only + mu_p = -0.5 * PRECIP_SIGMA ** 2 + P_pert = float(P * rng.lognormal(mu_p, PRECIP_SIGMA)) if P > 0 else 0.0 + E_pert = max(float(E * (1.0 + PET_SIGMA * rng.standard_normal())), 0.0) + q = run_model_step(m, P_pert, E_pert) + q_matrix[lead - 1, mem] = q + + m.finalize() + + row_data = {"issue_time": t0_str} + for lead in range(1, FORECAST_HOURS + 1): + row_data["lead_hour"] = lead + for mem in range(N_FORCING): + row_data[f"member_{mem:02d}"] = q_matrix[lead - 1, mem] + forcing_records.append(dict(row_data)) + + df_forcing_arm = pd.DataFrame(forcing_records) + out_a = OUT_DIR / f"{CAT_ID}_da_forcing_arm.csv" + df_forcing_arm.to_csv(out_a, index=False) + print(f" Saved: {out_a}") + + # ------------------------------------------------------------------ # + # Phase 2B: Hydro-state arm — 20 members, deterministic forcing, # + # perturbed analysis states # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Phase 2B: Hydro-state arm ({N_HYDRO} members)...") + hydro_records = [] + + for t0 in issue_times: + snap = helene_snapshots[t0] + t0_str = t0.strftime("%Y-%m-%d %H:%M:%S") + + if t0 not in t0_to_idx: + continue + start_idx = t0_to_idx[t0] + + q_matrix = np.full((FORECAST_HOURS, N_HYDRO), np.nan) + + for mem in range(N_HYDRO): + # Perturb the DA analysis state for this member + if mem == 0: + state_m = snap # member 0 = unperturbed analysis mean + else: + state_m = perturb_state(snap, rng) + + m = make_model(best_params) + set_state(m, state_m) + + for lead in range(1, FORECAST_HOURS + 1): + fi = start_idx + lead + if fi >= len(df_test): + break + row = df_test.iloc[fi] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + # Deterministic forcing + q = run_model_step(m, P, E) + q_matrix[lead - 1, mem] = q + + m.finalize() + + row_data = {"issue_time": t0_str} + for lead in range(1, FORECAST_HOURS + 1): + row_data["lead_hour"] = lead + for mem in range(N_HYDRO): + row_data[f"member_{mem:02d}"] = q_matrix[lead - 1, mem] + hydro_records.append(dict(row_data)) + + df_hydro_arm = pd.DataFrame(hydro_records) + out_b = OUT_DIR / f"{CAT_ID}_da_hydro_arm.csv" + df_hydro_arm.to_csv(out_b, index=False) + print(f" Saved: {out_b}") + print(f"[{CAT_ID}] Done.") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/run_perturbation_sensitivity.py b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/run_perturbation_sensitivity.py new file mode 100644 index 00000000..57b8d185 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/run_perturbation_sensitivity.py @@ -0,0 +1,293 @@ +""" +Perturbation-source sensitivity analysis for the multi-catchment CFE DA experiment. + +Ensemble-spread attribution: run 20 CFE members through the test +period with ONLY ONE perturbation source active at a time, no DA, no spinup. + +Three sub-experiments (selected via --source): + init : perturb initial states once at t=0 (multiplicative ~5%) + forcing : perturb hourly precip (lognormal sigma=0.15) and PET (Gaussian sigma=0.10) + process : apply per-hour process noise on states (soil 0.2%, GW 0.15%, Nash 0.5%) + +Each sub-experiment runs N=20 CFE members. DA is always OFF. Spinup is skipped — each +member starts from CFE's default initial state (modified by the perturbation if --source +init). This is the cleanest setup for attributing test-period spread to its source. + +Output: + //_sensitivity_.csv + Columns: date, member_00, member_01, ..., member_19 + +Each member column is that member's hourly Q_sim (mm/h) over the test period. + +Reads the same calibrated best_params.json as the production script. Does not need +the kriging obs file (no DA), but takes the same --obs-dir argument so it can also +log Q_obs alongside for plotting reference. + +Usage: + python3 run_perturbation_sensitivity.py \ + --cat-id cat-1016300 \ + --source forcing \ + --forcing-dir \ + --obs-dir \ + --cfe-dir \ + --config-file \ + --param-bounds \ + --out-dir \ + --test-forcing-dir1 \ + --test-forcing-dir2 +""" + +import argparse +import os +import sys +import json +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 +N_MEMBERS = 20 + +# We only run the test window for this sensitivity analysis. +TEST_START = "2023-10-01 00:00:00" +TEST_END = "2024-10-31 23:00:00" + +# Perturbation magnitudes — same as production defaults +INIT_STATE_FRAC = 0.05 +PRECIP_PERTURB = 0.15 # lognormal sigma +PET_PERTURB = 0.10 # Gaussian sigma +SOIL_PROC_NOISE = 0.002 +GW_PROC_NOISE = 0.0015 +NASH_PROC_NOISE = 0.005 + +# Set at runtime +CAT_ID = None +FORCING_FILE = None +TEST_FORCING_FILE = None +OBS_FILE = None +CFE_CONFIG_FILE = None +OUT_DIR = None +bmi_cfe = None + + +# ---------- Forcing helpers (same as production) ---------- +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(): + df = pd.read_csv(TEST_FORCING_FILE) + df['date'] = pd.to_datetime(df['time']).dt.strftime('%Y-%m-%d %H:%M:%S') + df['total_precipitation'] = df['APCP_surface'] * 3600.0 # kg/m^2/s -> mm/h + df['potential_evaporation'] = priestley_taylor_pet( + df['DSWRF_surface'].values, + df['TMP_2maboveground'].values, + ) + return df + + +# ---------- Perturbation helpers ---------- +def perturb_precip_lognormal(P, rng, sigma=PRECIP_PERTURB): + """Lognormal multiplier with mean 1; preserves zero precip.""" + mu = -0.5 * sigma ** 2 + return float(P * rng.lognormal(mu, sigma)) + + +def perturb_pet_gaussian(E, rng, sigma=PET_PERTURB): + """Gaussian multiplicative noise on PET, clipped at 0.""" + return max(float(E * (1.0 + sigma * rng.standard_normal())), 0.0) + + +def apply_init_state_perturb(m, rng, frac=INIT_STATE_FRAC): + """One-time multiplicative perturbation on the four states (called per member at t=0).""" + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1.0 + frac * rng.standard_normal()), 1e-6, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1.0 + frac * rng.standard_normal()), 1e-6, gw_max)) + # Nash buckets often start at 0; small additive jitter so they're not identical. + jitter = max(sm_max * 1e-6, 1e-9) + m.nash_storage[0] = max(n00 + jitter * rng.standard_normal(), 0.0) + m.nash_storage[1] = max(n10 + jitter * rng.standard_normal(), 0.0) + + +def apply_process_noise(m, rng): + """Per-timestep process noise on all 4 states for one member.""" + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + nash_floor = max(sm_max * 1e-4, 1e-7) + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1.0 + SOIL_PROC_NOISE * rng.standard_normal()), 0.0, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1.0 + GW_PROC_NOISE * rng.standard_normal()), 0.0, gw_max)) + m.nash_storage[0] = max( + n00 * (1.0 + NASH_PROC_NOISE * rng.standard_normal()) + + nash_floor * rng.standard_normal(), 0.0) + m.nash_storage[1] = max( + n10 * (1.0 + NASH_PROC_NOISE * rng.standard_normal()) + + nash_floor * rng.standard_normal(), 0.0) + + +# ---------- Main per-source run ---------- +def run_sensitivity(source, best_param_dict): + """Run N members through the test period with only `source` perturbation active.""" + + def custom_load_forcing(self_cfe): + df = load_test_forcing() + self_cfe.forcing_data = df.rename(columns={"date": "time"}) + + # Build a per-catchment temporary CFE config with the calibrated params + with open(CFE_CONFIG_FILE) as f: + cfg = json.load(f) + cfg['forcing_file'] = TEST_FORCING_FILE + cfg['soil_params']['bb'] = best_param_dict['bb'] + cfg['soil_params']['smcmax'] = best_param_dict['smcmax'] + cfg['soil_params']['satdk'] = best_param_dict['satdk'] + cfg['slop'] = best_param_dict['slop'] + cfg['max_gw_storage'] = best_param_dict['max_gw_storage'] + cfg['expon'] = best_param_dict['expon'] + cfg['Cgw'] = best_param_dict['Cgw'] + cfg['K_lf'] = best_param_dict['K_lf'] + cfg['K_nash'] = best_param_dict['K_nash'] + cfg['partition_scheme'] = "Schaake" if best_param_dict['scheme'] <= 0.5 else "Xinanjiang" + + tmp_cfg = str(OUT_DIR / f'{CAT_ID}_bmi_config_temp_sensitivity_{source}.json') + with open(tmp_cfg, 'w') as f: + json.dump(cfg, f) + + # Deterministic seed per (catchment, source) so each run is reproducible + seed = hash((CAT_ID, source)) & 0x7fffffff + rng = np.random.default_rng(seed) + + print(f"[sensitivity] {CAT_ID} | source={source} | N={N_MEMBERS} | seed={seed}") + + # Build N members. Initial-state perturbation only happens when source == 'init'. + models = [] + for i in range(N_MEMBERS): + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.load_forcing_file = custom_load_forcing.__get__(m) + m.initialize() + if source == 'init' and i > 0: + apply_init_state_perturb(m, rng) + models.append(m) + + # Skip spinup. Go directly to the test period. + df = load_test_forcing() + t_mask = (df['date'] >= TEST_START) & (df['date'] <= TEST_END) + df_test = df[t_mask] + + # Pre-allocate per-member Q arrays + n_hours = len(df_test) + q_matrix = np.full((n_hours, N_MEMBERS), np.nan, dtype=float) + dates_out = df_test['date'].values + + for h, (p, e) in enumerate(zip(df_test['total_precipitation'], + df_test['potential_evaporation'])): + for i, m in enumerate(models): + if source == 'forcing': + p_i = perturb_precip_lognormal(p, rng) + e_i = perturb_pet_gaussian(e, rng) + else: + p_i = float(p) + e_i = float(e) + m.set_value('atmosphere_water__time_integral_of_precipitation_mass_flux', p_i / 1000) + m.set_value('water_potential_evaporation_flux', e_i / 1000 / 3600) + m.update() + q_matrix[h, i] = m.get_value('land_surface_water__runoff_depth') * 1000 # m/h -> mm/h + + # Process noise applied after each hour, but only when source == 'process' + if source == 'process': + for m in models: + apply_process_noise(m, rng) + + for m in models: + m.finalize() + + # Write per-member CSV: date + 20 columns + out = {'date': dates_out} + for i in range(N_MEMBERS): + out[f'member_{i:02d}'] = q_matrix[:, i] + df_out = pd.DataFrame(out) + out_path = OUT_DIR / f'{CAT_ID}_sensitivity_{source}.csv' + df_out.to_csv(out_path, index=False) + print(f"[sensitivity] saved {out_path}") + print(f"[sensitivity] mean ensemble spread (std across members) over test period: " + f"{float(q_matrix.std(axis=1).mean()):.5f} mm/h") + + +def main(): + global CAT_ID, FORCING_FILE, TEST_FORCING_FILE, OBS_FILE, CFE_CONFIG_FILE + global OUT_DIR, bmi_cfe + + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', required=True) + parser.add_argument('--source', required=True, + choices=['init', 'forcing', 'process'], + help="Which perturbation source to isolate") + parser.add_argument('--forcing-dir', required=True) + parser.add_argument('--obs-dir', required=True, + help='Kept for symmetry with production script; not used since DA is off') + parser.add_argument('--cfe-dir', required=True) + parser.add_argument('--config-file', required=True) + parser.add_argument('--param-bounds', required=True, + help='Kept for symmetry; not read in this script') + parser.add_argument('--out-dir', required=True) + parser.add_argument('--test-forcing-dir1', default=None) + parser.add_argument('--test-forcing-dir2', default=None) + args = parser.parse_args() + + CAT_ID = args.cat_id + FORCING_FILE = os.path.join(args.forcing_dir, f'{CAT_ID}.csv') + OBS_FILE = os.path.join(args.obs_dir, f'{CAT_ID}.csv') + CFE_CONFIG_FILE = args.config_file + OUT_DIR = Path(args.out_dir) / CAT_ID + OUT_DIR.mkdir(parents=True, exist_ok=True) + + # Build combined test forcing CSV (same idiom as production) + if args.test_forcing_dir1 and args.test_forcing_dir2: + f1 = os.path.join(args.test_forcing_dir1, f'{CAT_ID}.csv') + f2 = os.path.join(args.test_forcing_dir2, f'{CAT_ID}.csv') + if os.path.exists(f1) and os.path.exists(f2): + df1 = pd.read_csv(f1) + df2 = pd.read_csv(f2) + combined = pd.concat([df1, df2], ignore_index=True) + combined = combined.drop_duplicates(subset='time').sort_values('time') + TEST_FORCING_FILE = str(OUT_DIR / f'{CAT_ID}_nwm_operational_combined.csv') + combined.to_csv(TEST_FORCING_FILE, index=False) + else: + print(f"Warning: test forcing files not found for {CAT_ID}, exiting") + return + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + bmi_cfe = _bmi_cfe + + best_params_file = OUT_DIR / f'{CAT_ID}_best_params.json' + if not best_params_file.exists(): + print(f"No best params file at {best_params_file}. Pre-stage from Run 3.") + return + with open(best_params_file) as f: + saved = json.load(f) + best_param_dict = saved['best_parameters'] + + run_sensitivity(args.source, best_param_dict) + + +if __name__ == '__main__': + main() diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/run_production_per_member.py b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/run_production_per_member.py new file mode 100644 index 00000000..01c70dab --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/2_assimilation/run_production_per_member.py @@ -0,0 +1,300 @@ +""" +Run the production DA pipeline for ONE catchment and save per-member streamflow. + +Identical to calibrate_catchment_cfe_da_v2.py in algorithm (init perturbation + +forcing perturbation + process noise + true EnKF with Vrugt R), but writes one +column per ensemble member instead of only the ensemble mean. This is for +the per-member factor-decomposition plot. + +To keep the comparison apples-to-apples with the sensitivity runs (which skip +spinup), this script also SKIPS the spinup loop. Each member starts directly +from CFE's default initial state, plus the t=0 init perturbation. + +Output: + //_production_per_member.csv + Columns: date, member_00, member_01, ..., member_19 + +Imports EnKFAssimilator from calibrate_catchment_cfe_da_v2 to guarantee the +DA math matches production exactly. + +Usage: + python3 run_production_per_member.py \ + --cat-id cat-1016300 \ + --forcing-dir /mnt/.../nwm_retro_catchment_forcings \ + --obs-dir /mnt/.../catchment_ts_03463300_with_variance \ + --cfe-dir /mnt/.../cfe_py \ + --config-file /mnt/.../cat_03463300_bmi_config_cfe.json \ + --param-bounds /mnt/.../CFE_parameter_bounds.json \ + --out-dir /mnt/.../v2_production_per_member \ + --test-forcing-dir1 ... \ + --test-forcing-dir2 ... \ + --enkf-members 20 +""" + +import argparse +import os +import sys +import json +import importlib.util +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +TEST_START = "2023-10-01 00:00:00" +TEST_END = "2024-10-31 23:00:00" + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(test_forcing_file): + df = pd.read_csv(test_forcing_file) + df['date'] = pd.to_datetime(df['time']).dt.strftime('%Y-%m-%d %H:%M:%S') + df['total_precipitation'] = df['APCP_surface'] * 3600.0 + df['potential_evaporation'] = priestley_taylor_pet( + df['DSWRF_surface'].values, df['TMP_2maboveground'].values + ) + return df + + +def import_enkf_class(script_path): + """Load EnKFAssimilator from the production script without triggering its main().""" + spec = importlib.util.spec_from_file_location("prod_v2", script_path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod.EnKFAssimilator + + +def run(args, bmi_cfe, EnKFAssimilator): + cat_id = args.cat_id + out_dir = Path(args.out_dir) / cat_id + out_dir.mkdir(parents=True, exist_ok=True) + + # Combined test forcing + f1 = os.path.join(args.test_forcing_dir1, f'{cat_id}.csv') + f2 = os.path.join(args.test_forcing_dir2, f'{cat_id}.csv') + df1 = pd.read_csv(f1) + df2 = pd.read_csv(f2) + combined = pd.concat([df1, df2], ignore_index=True) + combined = combined.drop_duplicates(subset='time').sort_values('time') + test_forcing_file = str(out_dir / f'{cat_id}_nwm_operational_combined.csv') + combined.to_csv(test_forcing_file, index=False) + + # Load best_params + best_params_file = out_dir / f'{cat_id}_best_params.json' + if not best_params_file.exists(): + print(f"Missing {best_params_file}. Pre-stage from Run 3.") + return + with open(best_params_file) as f: + saved = json.load(f) + best = saved['best_parameters'] + + # Build temp CFE config with calibrated params + with open(args.config_file) as f: + cfg = json.load(f) + cfg['forcing_file'] = test_forcing_file + cfg['soil_params']['bb'] = best['bb'] + cfg['soil_params']['smcmax'] = best['smcmax'] + cfg['soil_params']['satdk'] = best['satdk'] + cfg['slop'] = best['slop'] + cfg['max_gw_storage'] = best['max_gw_storage'] + cfg['expon'] = best['expon'] + cfg['Cgw'] = best['Cgw'] + cfg['K_lf'] = best['K_lf'] + cfg['K_nash'] = best['K_nash'] + cfg['partition_scheme'] = "Schaake" if best['scheme'] <= 0.5 else "Xinanjiang" + tmp_cfg = str(out_dir / f'{cat_id}_bmi_config_temp_per_member.json') + with open(tmp_cfg, 'w') as f: + json.dump(cfg, f) + + # Build the EnKF assimilator (same defaults as production) + obs_file = os.path.join(args.obs_dir, f'{cat_id}.csv') + enkf = EnKFAssimilator( + n_members=args.enkf_members, + obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, + use_vrugt_r=(not args.no_vrugt_r), + vrugt_alpha=args.vrugt_alpha, + vrugt_scale=args.vrugt_scale, + rng_seed=args.rng_seed, + ) + N = enkf.n_members + print(f"[per-member] {cat_id} | N={N} | use_vrugt_r={enkf.use_vrugt_r}") + if getattr(args, 'hardcoded_r', None) is not None: + enkf.obs_var_dict = {k: args.hardcoded_r for k in enkf.obs_var_dict} + print(f"[per-member] obs_var_dict overridden: R = {args.hardcoded_r} (fixed)") + + # Custom forcing loader for CFE BMI + def custom_load_forcing(self_cfe): + df = load_test_forcing(test_forcing_file) + self_cfe.forcing_data = df.rename(columns={"date": "time"}) + + # Build N CFE instances; perturb initial states (member 0 left clean) + models = [] + init_states_per_member = [] # capture initial-state snapshot after perturbation + for i in range(N): + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.load_forcing_file = custom_load_forcing.__get__(m) + m.initialize() + if i > 0: + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + r = enkf.init_state_perturb_frac + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1 + r * enkf.rng.standard_normal()), 1e-6, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1 + r * enkf.rng.standard_normal()), 1e-6, gw_max)) + jitter = max(sm_max * 1e-6, 1e-9) + m.nash_storage[0] = max(n00 + jitter * enkf.rng.standard_normal(), 0.0) + m.nash_storage[1] = max(n10 + jitter * enkf.rng.standard_normal(), 0.0) + # Snapshot the (possibly-perturbed) initial state for this member + init_states_per_member.append({ + "soil_m": float(m.soil_reservoir["storage_m"]), + "gw_m": float(m.gw_reservoir["storage_m"]), + "nash0_m": float(m.nash_storage[0]), + "nash1_m": float(m.nash_storage[1]), + "soil_max_m": float(m.soil_reservoir["storage_max_m"]), + "gw_max_m": float(m.gw_reservoir["storage_max_m"]), + }) + models.append(m) + + # SKIP SPINUP — go directly to the test period (matches sensitivity-script setup) + df = load_test_forcing(test_forcing_file) + t_mask = (df['date'] >= TEST_START) & (df['date'] <= TEST_END) + df_test = df[t_mask] + n_hours = len(df_test) + q_matrix = np.full((n_hours, N), np.nan, dtype=float) + precip_matrix = np.full((n_hours, N), np.nan, dtype=float) # mm/h per member + pet_matrix = np.full((n_hours, N), np.nan, dtype=float) # mm/h per member + dates_out = df_test['date'].values + + for h, (p, e, current_date) in enumerate(zip( + df_test['total_precipitation'], + df_test['potential_evaporation'], + df_test['date'])): + + # Perturb forcing per member + p_arr, e_arr = enkf.perturb_forcing(p, e) + # Capture the per-member perturbed forcing for later plotting/diagnostics + precip_matrix[h, :] = p_arr + pet_matrix[h, :] = e_arr + + # Advance each member one hour + ensemble_q = np.empty(N, dtype=float) + for i, m in enumerate(models): + m.set_value('atmosphere_water__time_integral_of_precipitation_mass_flux', + float(p_arr[i]) / 1000) + m.set_value('water_potential_evaporation_flux', + float(e_arr[i]) / 1000 / 3600) + m.update() + ensemble_q[i] = m.get_value('land_surface_water__runoff_depth') * 1000 # mm/h + + # Record pre-analysis forecast per member (this is the "actual production" value) + q_matrix[h, :] = ensemble_q + + # DA step (uses production EnKFAssimilator) + enkf.update_states(models, current_date, ensemble_q) + + # Process noise after DA + enkf.add_process_noise(models) + + for m in models: + m.finalize() + + # ----- Save per-member CSVs ----- + def _save_matrix(matrix, suffix): + out = {'date': dates_out} + for i in range(N): + out[f'member_{i:02d}'] = matrix[:, i] + df_out = pd.DataFrame(out) + out_path = out_dir / f'{cat_id}_production_per_member_{suffix}.csv' + df_out.to_csv(out_path, index=False) + return out_path + + # Streamflow outputs + q_path = out_dir / f'{cat_id}_production_per_member.csv' + out = {'date': dates_out} + for i in range(N): + out[f'member_{i:02d}'] = q_matrix[:, i] + pd.DataFrame(out).to_csv(q_path, index=False) + print(f"[per-member] saved {q_path}") + + # Perturbed forcings (so each member's actual inputs are recoverable) + precip_path = _save_matrix(precip_matrix, "precip") + pet_path = _save_matrix(pet_matrix, "pet") + print(f"[per-member] saved {precip_path}") + print(f"[per-member] saved {pet_path}") + + # Initial states (small, JSON is fine) + init_states_path = out_dir / f'{cat_id}_production_per_member_initial_states.json' + with open(init_states_path, 'w') as f: + json.dump({ + "catchment_id": cat_id, + "n_members": N, + "members": {f"member_{i:02d}": init_states_per_member[i] for i in range(N)}, + }, f, indent=2) + print(f"[per-member] saved {init_states_path}") + + print(f"[per-member] ensemble mean spread over test period: " + f"{float(q_matrix.std(axis=1).mean()):.5f} mm/h") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', required=True) + parser.add_argument('--forcing-dir', required=True) + parser.add_argument('--obs-dir', required=True) + parser.add_argument('--cfe-dir', required=True) + parser.add_argument('--config-file', required=True) + parser.add_argument('--param-bounds', required=True, + help='Kept for symmetry; not read in this script') + parser.add_argument('--out-dir', required=True) + parser.add_argument('--test-forcing-dir1', required=True) + parser.add_argument('--test-forcing-dir2', required=True) + parser.add_argument('--enkf-members', type=int, default=20) + parser.add_argument('--enkf-obs-error-std', type=float, default=0.05) + parser.add_argument('--no-vrugt-r', action='store_true') + parser.add_argument('--hardcoded-r', type=float, default=None, + help='Fix R to this constant value for all timesteps (overrides Vrugt formula)') + parser.add_argument('--vrugt-alpha', type=float, default=0.10) + parser.add_argument('--vrugt-scale', type=float, default=0.001) + parser.add_argument('--rng-seed', type=int, default=None) + parser.add_argument('--prod-script', default=None, + help='Path to calibrate_catchment_cfe_da_v2.py for importing ' + 'EnKFAssimilator. Defaults to the script located next to this file.') + args = parser.parse_args() + + # Locate the production script (default: same dir as this script) + if args.prod_script is None: + here = os.path.dirname(os.path.abspath(__file__)) + prod_script = os.path.join(here, 'calibrate_catchment_cfe_da_v2.py') + else: + prod_script = args.prod_script + if not os.path.exists(prod_script): + raise FileNotFoundError(f"Could not find production script at {prod_script}. " + f"Pass --prod-script to override.") + EnKFAssimilator = import_enkf_class(prod_script) + + # Dynamic CFE import (same idiom as production script) + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + + run(args, _bmi_cfe, EnKFAssimilator) + + +if __name__ == '__main__': + main() diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/3_routing/route_det_f2.sh b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/3_routing/route_det_f2.sh new file mode 100644 index 00000000..8dcf87c6 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/3_routing/route_det_f2.sh @@ -0,0 +1,28 @@ +#!/usr/bin/env bash +# F2 fixed R=0.07 (1 gauge holdout) — deterministic T-route routing. +# Routes per-catchment _test_results.csv (sim_mm_h column) through Muskingum-Cunge +# to gauge 03463300 and writes routed_Q_test.csv + KGE/NSE summary. +# +# Usage: +# bash route_det_f2.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +ROUTE_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/2_troute_routing/run_route.py + +F2_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder2_fixed_r007 +GPKG=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +echo "[F2] Deterministic T-route routing..." +echo " da-dir : $F2_DIR" +echo " out-dir: $F2_DIR" + +$TROUTE "$ROUTE_SCRIPT" \ + --gpkg "$GPKG" \ + --da-dir "$F2_DIR" \ + --out-dir "$F2_DIR" \ + --usgs-csv "$USGS_CSV" + +echo "[F2] Deterministic routing done. Output: $F2_DIR/routed_Q_test.csv" diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/3_routing/route_ensemble_f2.sh b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/3_routing/route_ensemble_f2.sh new file mode 100644 index 00000000..287b2ce1 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/3_routing/route_ensemble_f2.sh @@ -0,0 +1,30 @@ +#!/usr/bin/env bash +# F2 Fixed R=0.07 (1 gauge holdout) — crossed-ensemble T-route routing. +# Routes all 600 crossed-ensemble members through Muskingum-Cunge +# to gauge 03463300 and writes routed_crossed_ensemble.parquet. +# +# Must be run AFTER crossed ensemble parquets exist in F2_DIR. +# +# Usage: +# bash route_ensemble_f2.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +ROUTE_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/2_troute_routing/run_route_crossed_ensemble.py + +F2_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder2_fixed_r007 +GPKG=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +echo "[F2] Crossed-ensemble T-route routing..." +echo " ensemble-dir: $F2_DIR" +echo " out-dir : $F2_DIR" + +$TROUTE "$ROUTE_SCRIPT" \ + --gpkg "$GPKG" \ + --ensemble-dir "$F2_DIR" \ + --out-dir "$F2_DIR" \ + --usgs-csv "$USGS_CSV" + +echo "[F2] Ensemble routing done. Output: $F2_DIR/routed_crossed_ensemble.parquet" diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/3_routing/route_leadtime_f2.sh b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/3_routing/route_leadtime_f2.sh new file mode 100644 index 00000000..23f4fb53 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/3_routing/route_leadtime_f2.sh @@ -0,0 +1,31 @@ +#!/usr/bin/env bash +# F2 fixed R=0.07 (1 gauge holdout) — route lead-time forecast CSVs. +# Routes all 21-catchment lead-time forecast CSVs through Muskingum-Cunge +# to gauge 03463300. Writes routed_leadtime_da_full.parquet and +# routed_leadtime_openloop_full.parquet for gauge-level error decay analysis. +# +# Run AFTER batch_run_leadtime_f2.sh finishes. +# +# Usage: +# bash route_leadtime_f2.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +ROUTE_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/2_troute_routing/run_route_leadtime_forecasts.py + +F2_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder2_fixed_r007 +GPKG=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder2_fixed_r007_leadtime_routed + +echo "[F2] Routing lead-time forecasts through T-route..." +echo " forecast-dir: $F2_DIR" +echo " out-dir : $OUT_DIR" + +$TROUTE "$ROUTE_SCRIPT" \ + --gpkg "$GPKG" \ + --forecast-dir "$F2_DIR" \ + --out-dir "$OUT_DIR" + +echo "[F2] Lead-time routing done. Outputs in: $OUT_DIR" +ls "$OUT_DIR"/*.parquet 2>/dev/null || echo " (no parquets found)" diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4a_error_decay/plot_forecast_error_fixed_target.py b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4a_error_decay/plot_forecast_error_fixed_target.py new file mode 100644 index 00000000..1460ed21 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4a_error_decay/plot_forecast_error_fixed_target.py @@ -0,0 +1,217 @@ +""" +plot_forecast_error_fixed_target.py — 4a: error decay, fixed-target-time view. + +For each target verification time T in the Helene peak window: + Collect all forecasts that verify AT T: + issue_time = T - lead_hour*1h, for lead in 1..18 + error[lead] = ensemble_mean(q at T, initialized T-lead) - USGS_obs(T) + +This gives the correct operational picture: + - lead 1 = initialized 1 hr before T (DA just ran -> small error) + - lead 18 = initialized 18 hr before T (DA long ago -> error ~ open loop) + +Two panels: + Top : signed error (m³/s) vs lead hour, one curve per target time + Bot : same for open-loop +Plus a summary panel: mean across all target times, DA vs OL. + +Outputs: + /error_fixed_target_helene.png (per-target spaghetti, DA vs OL) + /error_fixed_target_mean.png (mean across targets, DA vs OL) +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.cm as cm + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast_hardcoded_r_routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +# Target verification times: hourly through the Helene peak window +TARGET_START = pd.Timestamp("2024-09-26 18:00:00") +TARGET_END = pd.Timestamp("2024-09-28 06:00:00") + + +def load_parquet(path): + df = pd.read_parquet(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + df["ens_mean"] = df[member_cols].mean(axis=1) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + return df[["issue_time", "lead_hour", "valid_time", "ens_mean"]] + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() + or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def build_fixed_target_errors(df, obs_series, target_times): + """ + For each target time T and each lead L (1..18): + issue_time = T - L hours + error = ensemble_mean at (issue_time, lead=L) - obs(T) + Returns dict: target_time -> {lead: error} + """ + # Index df by (issue_time, lead_hour) for fast lookup + df_idx = df.set_index(["issue_time", "lead_hour"])["ens_mean"] + + results = {} + for T in target_times: + obs_val = obs_series.get(T, np.nan) + if np.isnan(obs_val): + continue + curve = {} + for lead in range(1, 19): + t0 = T - pd.Timedelta(hours=lead) + try: + q_fc = df_idx.loc[(t0, lead)] + curve[lead] = float(q_fc) - obs_val + except KeyError: + curve[lead] = np.nan + results[T] = curve + return results + + +def plot_spaghetti(ax, error_dict, color_da, label_prefix, linestyle="-", lw=0.9, alpha=0.35): + """Plot one thin line per target time + thick mean.""" + leads = list(range(1, 19)) + all_curves = [] + target_times = sorted(error_dict.keys()) + cmap = cm.get_cmap("plasma", len(target_times)) + + for i, T in enumerate(target_times): + curve = [error_dict[T].get(l, np.nan) for l in leads] + ax.plot(leads, curve, + color=cmap(i), lw=lw, alpha=alpha, + linestyle=linestyle, zorder=2) + all_curves.append(curve) + + if all_curves: + mean_curve = np.nanmean(np.array(all_curves), axis=0) + ax.plot(leads, mean_curve, + color=color_da, lw=2.8, alpha=0.95, + linestyle=linestyle, zorder=5, + label=f"{label_prefix} — mean across targets") + return all_curves + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--route-dir", default=DEFAULT_ROUTE_DIR) + parser.add_argument("--out-dir", default=None) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--da-name", default="routed_leadtime_da_full.parquet") + parser.add_argument("--ol-name", default="routed_leadtime_openloop_full.parquet") + parser.add_argument("--target-start", default=str(TARGET_START)) + parser.add_argument("--target-end", default=str(TARGET_END)) + args = parser.parse_args() + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + print("Loading parquets...") + da = load_parquet(os.path.join(args.route_dir, args.da_name)) + ol = load_parquet(os.path.join(args.route_dir, args.ol_name)) + obs = load_usgs(args.usgs_csv) + + target_times = pd.date_range(args.target_start, args.target_end, freq="1h") + print(f" Target verification times: {len(target_times)} " + f"({target_times[0]} → {target_times[-1]})") + + print("Building fixed-target error tables...") + da_errors = build_fixed_target_errors(da, obs, target_times) + ol_errors = build_fixed_target_errors(ol, obs, target_times) + print(f" Targets with obs: DA={len(da_errors)} OL={len(ol_errors)}") + + leads = list(range(1, 19)) + + # ---- Spaghetti: per-target-time curves, DA vs OL ---- + fig, (ax_da, ax_ol) = plt.subplots(2, 1, figsize=(13, 10), sharex=True, sharey=True) + + ax_da.axhline(0, color="black", lw=0.8, linestyle="--", alpha=0.4) + plot_spaghetti(ax_da, da_errors, "tab:purple", "DA", linestyle="-") + ax_da.set_ylabel("Error: forecast − USGS obs (m³/s)", fontsize=10) + ax_da.set_title("DA — error at each lead for fixed target times (Helene peak window)", fontsize=11) + ax_da.grid(True, alpha=0.2) + ax_da.legend(fontsize=9) + + ax_ol.axhline(0, color="black", lw=0.8, linestyle="--", alpha=0.4) + plot_spaghetti(ax_ol, ol_errors, "tab:gray", "Open-loop", linestyle="--") + ax_ol.set_ylabel("Error: forecast − USGS obs (m³/s)", fontsize=10) + ax_ol.set_xlabel("Forecast lead hour (hours before target)", fontsize=11) + ax_ol.set_title("Open-loop — error at each lead for fixed target times", fontsize=11) + ax_ol.set_xticks(leads) + ax_ol.grid(True, alpha=0.2) + ax_ol.legend(fontsize=9) + + fig.suptitle( + "Forecast error vs lead time — fixed verification time, Helene peak window\n" + f"Target times: {TARGET_START.strftime('%b %d %H UTC')} → " + f"{TARGET_END.strftime('%b %d %H UTC')} | USGS 03463300", + fontsize=12, y=1.01, + ) + plt.tight_layout() + out1 = os.path.join(out_dir, "error_fixed_target_helene.png") + plt.savefig(out1, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + # ---- Summary: mean across all target times, DA vs OL overlaid ---- + fig, ax = plt.subplots(figsize=(13, 6)) + ax.axhline(0, color="black", lw=0.8, linestyle="--", alpha=0.4) + + def mean_curve(error_dict): + arr = np.array([ + [error_dict[T].get(l, np.nan) for l in leads] + for T in sorted(error_dict.keys()) + ]) + return np.nanmean(arr, axis=0), np.nanstd(arr, axis=0) + + da_mean, da_std = mean_curve(da_errors) + ol_mean, ol_std = mean_curve(ol_errors) + + ax.fill_between(leads, da_mean - da_std, da_mean + da_std, + color="tab:purple", alpha=0.15, zorder=2) + ax.fill_between(leads, ol_mean - ol_std, ol_mean + ol_std, + color="tab:gray", alpha=0.15, zorder=2) + ax.plot(leads, da_mean, color="tab:purple", lw=2.6, marker="o", + zorder=5, label="DA — mean error (±1 std shaded)") + ax.plot(leads, ol_mean, color="tab:gray", lw=2.6, marker="s", + linestyle="--", zorder=5, label="Open-loop — mean error (±1 std shaded)") + + ax.set_xlabel("Forecast lead hour (hours before target verification time)", fontsize=11) + ax.set_ylabel("Mean error: forecast − USGS obs (m³/s)", fontsize=11) + ax.set_xticks(leads) + ax.set_title( + "Mean forecast error vs lead time — fixed verification time, Helene peak window\n" + f"Target: {TARGET_START.strftime('%b %d %H UTC')} → " + f"{TARGET_END.strftime('%b %d %H UTC')} | " + "Lead 1 = init 1 hr before target | Lead 18 = init 18 hr before target", + fontsize=11, + ) + ax.legend(fontsize=10, loc="lower right", frameon=True, framealpha=0.92) + ax.grid(True, alpha=0.25) + plt.tight_layout() + out2 = os.path.join(out_dir, "error_fixed_target_mean.png") + plt.savefig(out2, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4a_error_decay/plot_forecast_error_per_init.py b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4a_error_decay/plot_forecast_error_per_init.py new file mode 100644 index 00000000..691245d0 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4a_error_decay/plot_forecast_error_per_init.py @@ -0,0 +1,203 @@ +""" +plot_forecast_error_per_init.py + +Error decay by initialization time for the Helene window (Sep 24-28 2024). + +For each initialization time t0 in the Helene window: + error[lead] = ensemble_mean(q_gauge_m3s at t0+lead) - USGS_obs(t0+lead) + +Plotted as: + DA : thin colored lines (one per init time, colored by date) + thick mean across all + OL : thin gray dashed lines + thick gray dashed mean + +x-axis: forecast lead hour (1 -> 18) +y-axis: signed error (m³/s), positive = forecast too high + +The expected signal: DA error is small at lead 1 (just assimilated), grows +and converges toward the OL error curve by lead 18. + +Inputs: + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + +Obs: + /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +Outputs: + /forecast_error_per_init_helene.png (signed error) + /forecast_rmse_per_init_helene.png (absolute error / RMSE per init) +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.patches as mpatches + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast_hardcoded_r_routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +# Initialization times to show — Helene window +HELENE_INIT_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_INIT_END = pd.Timestamp("2024-09-28 23:00:00") + +# One color per init date +DATE_COLORS = { + "2024-09-24": "#1f77b4", + "2024-09-25": "#ff7f0e", + "2024-09-26": "#2ca02c", + "2024-09-27": "#d62728", + "2024-09-28": "#9467bd", +} + + +def load_parquet(path): + df = pd.read_parquet(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + df["ens_mean"] = df[member_cols].mean(axis=1) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + return df[["issue_time", "lead_hour", "valid_time", "ens_mean"]] + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() + or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def build_error_table(df, obs_series, init_start, init_end): + """Return DataFrame: issue_time, lead_hour, ens_mean, obs, error.""" + df = df[(df["issue_time"] >= init_start) & (df["issue_time"] <= init_end)].copy() + df["obs"] = df["valid_time"].map(obs_series) + df["error"] = df["ens_mean"] - df["obs"] + return df.dropna(subset=["obs", "error"]) + + +def plot_error(ax, err_df, color, alpha_thin, lw_thin, linestyle, label_prefix): + """Plot individual init-time error curves + thick mean curve.""" + leads = sorted(err_df["lead_hour"].unique()) + all_curves = [] + + for t0, grp in err_df.groupby("issue_time"): + date_str = str(pd.Timestamp(t0).date()) + c = DATE_COLORS.get(date_str, color) + grp_sorted = grp.sort_values("lead_hour") + # Align to leads grid — some may be missing + curve = grp_sorted.set_index("lead_hour")["error"].reindex(leads).values + ax.plot(leads, curve, + color=c, lw=lw_thin, alpha=alpha_thin, + linestyle=linestyle, zorder=2) + all_curves.append(curve) + + if all_curves: + mean_curve = np.nanmean(np.array(all_curves), axis=0) + ax.plot(leads, mean_curve, + color=color, lw=2.8, alpha=0.95, + linestyle=linestyle, zorder=5, + label=f"{label_prefix} — mean across all init times") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--route-dir", default=DEFAULT_ROUTE_DIR) + parser.add_argument("--out-dir", default=None) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--da-name", default="routed_leadtime_da_full.parquet") + parser.add_argument("--ol-name", default="routed_leadtime_openloop_full.parquet") + args = parser.parse_args() + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + print("Loading parquets...") + da = load_parquet(os.path.join(args.route_dir, args.da_name)) + ol = load_parquet(os.path.join(args.route_dir, args.ol_name)) + obs = load_usgs(args.usgs_csv) + print(f" DA issue_times: {da['issue_time'].nunique()} " + f"OL issue_times: {ol['issue_time'].nunique()}") + + da_err = build_error_table(da, obs, HELENE_INIT_START, HELENE_INIT_END) + ol_err = build_error_table(ol, obs, HELENE_INIT_START, HELENE_INIT_END) + print(f" DA init times in Helene window: {da_err['issue_time'].nunique()}") + + leads = sorted(da_err["lead_hour"].unique()) + + # ---- Signed error plot ---- + fig, ax = plt.subplots(figsize=(12, 6)) + ax.axhline(0, color="black", lw=0.8, linestyle="--", alpha=0.5, zorder=1) + + plot_error(ax, ol_err, color="tab:gray", alpha_thin=0.12, lw_thin=0.7, + linestyle="--", label_prefix="Open-loop") + plot_error(ax, da_err, color="tab:purple", alpha_thin=0.18, lw_thin=0.8, + linestyle="-", label_prefix="DA") + + # Date-color legend patches + patches = [mpatches.Patch(color=c, label=f"Init {d}") + for d, c in DATE_COLORS.items()] + handles, labels = ax.get_legend_handles_labels() + ax.legend(handles=handles + patches, + fontsize=8, loc="upper left", frameon=True, framealpha=0.9, ncol=2) + + ax.set_xlabel("Forecast lead hour", fontsize=11) + ax.set_ylabel("Error: forecast − USGS obs (m³/s)", fontsize=11) + ax.set_xticks(np.arange(1, 19)) + ax.set_title( + "Forecast error vs lead time — per initialization time, Helene window\n" + "DA (purple solid) vs Open-loop (gray dashed) | Sep 24–28 2024 | USGS 03463300", + fontsize=11, + ) + ax.grid(True, alpha=0.25, lw=0.4) + plt.tight_layout() + out1 = os.path.join(out_dir, "forecast_error_per_init_helene.png") + plt.savefig(out1, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + # ---- Absolute error (|error|) averaged per lead — cleaner summary ---- + fig, ax = plt.subplots(figsize=(12, 6)) + + def mean_abs_error_by_lead(err_df): + return err_df.groupby("lead_hour")["error"].apply( + lambda x: float(np.nanmean(np.abs(x))) + ) + + da_mae = mean_abs_error_by_lead(da_err) + ol_mae = mean_abs_error_by_lead(ol_err) + + ax.plot(da_mae.index, da_mae.values, + color="tab:purple", lw=2.4, marker="o", label="DA — mean |error|") + ax.plot(ol_mae.index, ol_mae.values, + color="tab:gray", lw=2.4, marker="s", linestyle="--", + label="Open-loop — mean |error|") + + ax.set_xlabel("Forecast lead hour", fontsize=11) + ax.set_ylabel("Mean |error| vs USGS obs (m³/s)", fontsize=11) + ax.set_xticks(np.arange(1, 19)) + ax.set_title( + "Mean absolute forecast error vs lead time — Helene window\n" + "DA (purple) vs Open-loop (gray) | Sep 24–28 2024 | USGS 03463300", + fontsize=11, + ) + ax.legend(fontsize=10, loc="upper left", frameon=True, framealpha=0.9) + ax.grid(True, alpha=0.25, lw=0.4) + plt.tight_layout() + out2 = os.path.join(out_dir, "forecast_mae_per_lead_helene.png") + plt.savefig(out2, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4a_error_decay/plot_lead_time_decay.py b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4a_error_decay/plot_lead_time_decay.py new file mode 100644 index 00000000..7da0b64c --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4a_error_decay/plot_lead_time_decay.py @@ -0,0 +1,174 @@ +""" +Catchment-level error-vs-lead-time decay curve + ensemble spread by lead. + +Two-panel figure for ONE catchment: + TOP — RMSE of ensemble-mean forecast at each lead hour (1..18) vs the + catchment's kriging observation. DA solid, open-loop dashed, with + a shaded band showing the min/max of per-member RMSE. + BOTTOM — Mean ensemble spread (std-dev across 20 members, averaged across + all issue times) at each lead hour. Tells you whether forcing + perturbation alone keeps the forecast ensemble diverse during the + 18-hour free-run — useful for inspecting individual members + without needing the full spaghetti view. + +This is the catchment-level analog of the gauge-level decay curve. +Routing to the gauge is a separate post-step (route_lead_time_forecasts.py ++ a gauge-level decay script); this script lets us look at the catchment-level +signal without T-route in the loop. + +Inputs (from run_lead_time_forecast_sweep.py): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv +Obs (from production EnKF run): + //_test_results.csv + +Output: + //_lead_time_decay.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" + +DEFAULT_LEADTIME_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast" +DEFAULT_DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" + +# Module-level globals set by main() so the helpers see them +CAT = "cat-1016300" +LEADTIME_DIR = DEFAULT_LEADTIME_DIR +DA_DIR = DEFAULT_DA_DIR +OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay.png") + + +def load_forecasts(path): + """Return (df, member_cols).""" + if not os.path.exists(path): + raise FileNotFoundError(path) + df = pd.read_csv(path, parse_dates=['issue_time', 'valid_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + return df, member_cols + + +OBS_DIR = None # set by main() when --obs-dir is supplied + + +def load_obs(): + if OBS_DIR is not None: + obs_path = os.path.join(OBS_DIR, f"{CAT}.csv") + df = pd.read_csv(obs_path) + time_col = 'datetime' if 'datetime' in df.columns else 'date' + df[time_col] = pd.to_datetime(df[time_col]) + col = 'qkrig' if 'qkrig' in df.columns else 'obs_mm_h' + return df.set_index(time_col)[col].rename('obs_mm_h') + obs_path = os.path.join(DA_DIR, CAT, f"{CAT}_test_results.csv") + df = pd.read_csv(obs_path, parse_dates=['date']) + return df.set_index('date')['obs_mm_h'] + + +def metrics_by_lead(df, member_cols, obs_series): + """For each lead hour, return (rmse_mean, rmse_min_member, rmse_max_member, + mean_ensemble_std).""" + df = df.copy() + df['obs'] = df['valid_time'].map(obs_series) + df = df.dropna(subset=['obs']) + df['ens_mean'] = df[member_cols].mean(axis=1) + df['ens_std'] = df[member_cols].std(axis=1) + + leads = sorted(df['lead_hour'].unique()) + rmse_mean = np.full(len(leads), np.nan) + rmse_min_member = np.full(len(leads), np.nan) + rmse_max_member = np.full(len(leads), np.nan) + mean_ens_std = np.full(len(leads), np.nan) + for i, L in enumerate(leads): + sub = df[df['lead_hour'] == L] + if len(sub) == 0: + continue + err = sub['ens_mean'].values - sub['obs'].values + rmse_mean[i] = float(np.sqrt(np.mean(err ** 2))) + per_member_rmse = np.array([ + np.sqrt(np.mean((sub[c].values - sub['obs'].values) ** 2)) + for c in member_cols + ]) + rmse_min_member[i] = float(np.nanmin(per_member_rmse)) + rmse_max_member[i] = float(np.nanmax(per_member_rmse)) + mean_ens_std[i] = float(np.nanmean(sub['ens_std'].values)) + return (np.asarray(leads), rmse_mean, rmse_min_member, + rmse_max_member, mean_ens_std) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', default="cat-1016300") + parser.add_argument('--leadtime-dir', default=DEFAULT_LEADTIME_DIR, + help='Dir holding /_lead_time_forecasts_{da,openloop}.csv') + parser.add_argument('--da-dir', default=DEFAULT_DA_DIR, + help='Dir holding /_test_results.csv (for obs)') + parser.add_argument('--obs-dir', default=None, + help='Kriging obs dir holding .csv with qkrig column; ' + 'overrides --da-dir for obs loading') + args = parser.parse_args() + + global CAT, LEADTIME_DIR, DA_DIR, OBS_DIR, OUT_PNG + CAT = args.cat_id + LEADTIME_DIR = args.leadtime_dir + DA_DIR = args.da_dir + OBS_DIR = args.obs_dir + OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay.png") + + da_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_da.csv") + ol_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_openloop.csv") + df_da, m_cols = load_forecasts(da_path) + df_ol, _ = load_forecasts(ol_path) + obs_series = load_obs() + + leads_da, rmse_da, rmse_da_lo, rmse_da_hi, std_da = metrics_by_lead( + df_da, m_cols, obs_series) + leads_ol, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = metrics_by_lead( + df_ol, m_cols, obs_series) + + fig, (ax_top, ax_bot) = plt.subplots(2, 1, figsize=(10, 10), sharex=True, + gridspec_kw={"height_ratios": [1.4, 1.0]}) + + # ----- TOP: RMSE decay curve ----- + ax_top.fill_between(leads_da, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2, + label="DA — member spread (min/max RMSE)") + ax_top.fill_between(leads_ol, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2, + label="Open-loop — member spread") + ax_top.plot(leads_da, rmse_da, color=DA_COLOR, lw=2.4, marker='o', + zorder=4, label="DA on (ensemble-mean RMSE)") + ax_top.plot(leads_ol, rmse_ol, color=OL_COLOR, lw=2.4, marker='s', + linestyle='--', zorder=4, label="Open-loop (ensemble-mean RMSE)") + ax_top.set_ylabel("RMSE vs Qkrig obs (mm/h)", fontsize=11) + ax_top.set_title( + f"Forecast lead-time error decay — {CAT}\n" + f"Issue times pooled across test period (Oct 2023 – Oct 2024)", + fontsize=12, + ) + ax_top.grid(True, alpha=0.25) + ax_top.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + # ----- BOTTOM: ensemble spread by lead hour ----- + ax_bot.plot(leads_da, std_da, color=DA_COLOR, lw=2.2, marker='o', + label="DA — mean ensemble std") + ax_bot.plot(leads_ol, std_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', label="Open-loop — mean ensemble std") + ax_bot.set_xlabel("Forecast lead time (hours)", fontsize=11) + ax_bot.set_ylabel("Mean ensemble std (mm/h)", fontsize=11) + ax_bot.set_xticks(np.arange(1, 19)) + ax_bot.grid(True, alpha=0.25) + ax_bot.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + plt.tight_layout() + plt.savefig(OUT_PNG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4a_error_decay/plot_lead_time_decay_by_regime.py b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4a_error_decay/plot_lead_time_decay_by_regime.py new file mode 100644 index 00000000..92b680b0 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4a_error_decay/plot_lead_time_decay_by_regime.py @@ -0,0 +1,214 @@ +""" +Lead-time decay curve, split by flow regime at issue time. + +The pooled lead-time curve (plot_lead_time_decay.py) showed DA losing to +open-loop across most lead hours, but the test period is ~99% low-flow. +This script splits the same forecast CSVs by the flow regime at the issue +time t0, so we can see whether DA helps when it matters (storms / Helene) +and hurts when it doesn't (low flow). + +Three regimes are partitioned on obs(t0) (kriging obs at the issue time): + + helene issue_time ∈ [2024-09-24, 2024-09-28] (the 5-day Helene window) + storm obs(t0) > 0.5 mm/h (high-flow issue times) + low_flow obs(t0) < 0.1 mm/h (low-flow issue times) + +Same 2-panel layout per regime: RMSE vs lead (top), ensemble spread vs lead +(bottom). Three regime columns × 2 metric rows in one figure. + +Inputs (from run_lead_time_forecast_sweep.py — no re-run needed): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv +Obs (from production EnKF): + //_test_results.csv + +Output: + //_lead_time_decay_by_regime.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" + +DEFAULT_LEADTIME_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast" +DEFAULT_DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" + +# Module-level globals set by main() so helpers see them +CAT = "cat-1016300" +LEADTIME_DIR = DEFAULT_LEADTIME_DIR +DA_DIR = DEFAULT_DA_DIR +OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay_by_regime.png") + +# Regime definitions +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-28 23:00:00") +STORM_OBS_THRESHOLD = 0.5 # mm/h +LOWFLOW_OBS_THRESHOLD = 0.1 # mm/h + + +def load_forecasts(path): + if not os.path.exists(path): + raise FileNotFoundError(path) + df = pd.read_csv(path, parse_dates=['issue_time', 'valid_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + return df, member_cols + + +OBS_DIR = None # set by main() when --obs-dir is supplied + + +def load_obs(): + if OBS_DIR is not None: + obs_path = os.path.join(OBS_DIR, f"{CAT}.csv") + df = pd.read_csv(obs_path) + time_col = 'datetime' if 'datetime' in df.columns else 'date' + df[time_col] = pd.to_datetime(df[time_col]) + col = 'qkrig' if 'qkrig' in df.columns else 'obs_mm_h' + return df.set_index(time_col)[col].rename('obs_mm_h') + obs_path = os.path.join(DA_DIR, CAT, f"{CAT}_test_results.csv") + df = pd.read_csv(obs_path, parse_dates=['date']) + return df.set_index('date')['obs_mm_h'] + + +def regime_mask(issue_times, obs_at_issue, regime): + """Boolean mask over issue_times for the named regime.""" + if regime == "helene": + return (issue_times >= HELENE_START) & (issue_times <= HELENE_END) + if regime == "storm": + return obs_at_issue > STORM_OBS_THRESHOLD + if regime == "low_flow": + return obs_at_issue < LOWFLOW_OBS_THRESHOLD + raise ValueError(regime) + + +def metrics_by_lead(df, member_cols, obs_series, issue_mask): + """Compute (leads, rmse_mean, rmse_min, rmse_max, mean_std) restricted + to issue times where issue_mask is True.""" + df = df.copy() + df = df[df['issue_time'].isin(issue_mask)] + if len(df) == 0: + return None + df['obs'] = df['valid_time'].map(obs_series) + df = df.dropna(subset=['obs']) + if len(df) == 0: + return None + df['ens_mean'] = df[member_cols].mean(axis=1) + df['ens_std'] = df[member_cols].std(axis=1) + + leads = sorted(df['lead_hour'].unique()) + rmse_mean = np.full(len(leads), np.nan) + rmse_lo = np.full(len(leads), np.nan) + rmse_hi = np.full(len(leads), np.nan) + std_mean = np.full(len(leads), np.nan) + for i, L in enumerate(leads): + sub = df[df['lead_hour'] == L] + if len(sub) == 0: + continue + err = sub['ens_mean'].values - sub['obs'].values + rmse_mean[i] = float(np.sqrt(np.mean(err ** 2))) + per_member = np.array([ + np.sqrt(np.mean((sub[c].values - sub['obs'].values) ** 2)) + for c in member_cols + ]) + rmse_lo[i] = float(np.nanmin(per_member)) + rmse_hi[i] = float(np.nanmax(per_member)) + std_mean[i] = float(np.nanmean(sub['ens_std'].values)) + return np.asarray(leads), rmse_mean, rmse_lo, rmse_hi, std_mean + + +def plot_pair(ax_rmse, ax_std, da_metrics, ol_metrics, + regime_label, n_issue_times): + if da_metrics is None or ol_metrics is None: + ax_rmse.set_title(f"{regime_label}\n(no issue times match)", fontsize=11) + ax_rmse.axis('off'); ax_std.axis('off') + return + leads, rmse_da, rmse_da_lo, rmse_da_hi, std_da = da_metrics + _, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = ol_metrics + + # RMSE panel + ax_rmse.fill_between(leads, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2) + ax_rmse.fill_between(leads, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2) + ax_rmse.plot(leads, rmse_da, color=DA_COLOR, lw=2.2, marker='o', + zorder=4, label="DA on") + ax_rmse.plot(leads, rmse_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', zorder=4, label="Open-loop") + ax_rmse.set_title(f"{regime_label}\n({n_issue_times} issue times)", fontsize=11) + ax_rmse.set_ylabel("RMSE (mm/h)", fontsize=10) + ax_rmse.grid(True, alpha=0.25) + ax_rmse.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + # Spread panel + ax_std.plot(leads, std_da, color=DA_COLOR, lw=2.0, marker='o', label="DA") + ax_std.plot(leads, std_ol, color=OL_COLOR, lw=2.0, marker='s', + linestyle='--', label="Open-loop") + ax_std.set_xlabel("Forecast lead time (hours)", fontsize=10) + ax_std.set_ylabel("Mean ensemble std (mm/h)", fontsize=10) + ax_std.set_xticks(np.arange(1, 19)) + ax_std.grid(True, alpha=0.25) + ax_std.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', default="cat-1016300") + parser.add_argument('--leadtime-dir', default=DEFAULT_LEADTIME_DIR, + help='Dir holding /_lead_time_forecasts_{da,openloop}.csv') + parser.add_argument('--da-dir', default=DEFAULT_DA_DIR, + help='Dir holding /_test_results.csv (for obs)') + parser.add_argument('--obs-dir', default=None, + help='Kriging obs dir holding .csv with qkrig column; ' + 'overrides --da-dir for obs loading') + args = parser.parse_args() + + global CAT, LEADTIME_DIR, DA_DIR, OBS_DIR, OUT_PNG + CAT = args.cat_id + LEADTIME_DIR = args.leadtime_dir + DA_DIR = args.da_dir + OBS_DIR = args.obs_dir + OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay_by_regime.png") + + da_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_da.csv") + ol_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_openloop.csv") + df_da, m_cols = load_forecasts(da_path) + df_ol, _ = load_forecasts(ol_path) + obs_series = load_obs() + + # Unique issue times across the run, with obs at each + issue_times = pd.to_datetime(sorted(df_da['issue_time'].unique())) + obs_at_issue = pd.Series(issue_times, index=issue_times).map(obs_series) + + regimes = [ + ("helene", "Helene window (Sep 24–28 2024)"), + ("storm", f"Storm hours (obs(t0) > {STORM_OBS_THRESHOLD} mm/h)"), + ("low_flow", f"Low flow (obs(t0) < {LOWFLOW_OBS_THRESHOLD} mm/h)"), + ] + + fig, axes = plt.subplots(2, 3, figsize=(18, 10), + gridspec_kw={"height_ratios": [1.4, 1.0]}) + for col, (key, label) in enumerate(regimes): + mask = regime_mask(issue_times, obs_at_issue, key) + kept = issue_times[mask] + kept_set = set(pd.to_datetime(kept)) + da_m = metrics_by_lead(df_da, m_cols, obs_series, kept_set) + ol_m = metrics_by_lead(df_ol, m_cols, obs_series, kept_set) + plot_pair(axes[0, col], axes[1, col], da_m, ol_m, label, len(kept)) + + fig.suptitle( + f"Lead-time error decay by flow regime — {CAT}\n" + f"Same forecast CSVs as the pooled view, partitioned on obs(t0).", + fontsize=13, y=0.995, + ) + plt.tight_layout(rect=(0, 0, 1, 0.97)) + plt.savefig(OUT_PNG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4a_error_decay/plot_lead_time_decay_gauge.py b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4a_error_decay/plot_lead_time_decay_gauge.py new file mode 100644 index 00000000..71ce54fc --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4a_error_decay/plot_lead_time_decay_gauge.py @@ -0,0 +1,311 @@ +""" +Gauge-level lead-time forecast decay curve. + +Routes 21-catchment forecast CSVs through T-route to USGS gauge 03463300 and +plots RMSE vs lead time (DA solid, open-loop dashed) against USGS observed Q. +Same 2-panel layout as the catchment-level plot, plus a regime split for the +Helene window. + +The routing step is done separately by route_lead_time_forecasts.py; +this script just reads the resulting parquets and produces the plot. + +Inputs: + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + Each: one row per (issue_time, lead_hour); member columns hold q_gauge_m3s. + (Long-format with explicit `member` and `q_gauge_m3s` columns also supported.) +Obs: + /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + USGS hourly Q in m³/s at gauge 03463300 (South Toe River near Celo, NC). + +Output: + /lead_time_decay_gauge_pooled.png + /lead_time_decay_gauge_by_regime.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" +OBS_COLOR = "black" + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/leadtime_troute_routing" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +# Watershed area at gauge 03463300 (sum of 21 catchments per GPKG hydrofabric) +WATERSHED_AREA_KM2 = 113.18 +# 1 mm/h depth × 113.18 km² = 113.18e3 m³/h = 31.439 m³/s +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-28 23:00:00") +STORM_OBS_THRESHOLD_M3S = 20.0 # gauge-level storm threshold (was 50 — too high + # for this 113 km² basin; only 2 issue times qualified) +LOWFLOW_OBS_THRESHOLD_M3S = 5.0 # gauge-level low-flow threshold + +USGS_HELENE_PEAK_M3S = 1886.0 # reference Sep 27 14:00 + + +def load_parquet_long(path): + """Load a routed parquet and normalize to long format: + columns = issue_time, lead_hour, member, q_gauge_m3s. + + Handles both wide format (member_00..member_19 columns) and long format + (explicit `member` + `q_gauge_m3s` columns). Auto-detects. + """ + if not os.path.exists(path): + raise FileNotFoundError(path) + df = pd.read_parquet(path) + + # Long format detection + if 'q_gauge_m3s' in df.columns and 'member' in df.columns: + df['issue_time'] = pd.to_datetime(df['issue_time']) + return df[['issue_time', 'lead_hour', 'member', 'q_gauge_m3s']] + + # Some variants might call it 'q_m3s' or similar + q_col_candidates = [c for c in df.columns if c.lower() in + ('q_gauge_m3s', 'q_m3s', 'q_gauge', 'q')] + if 'member' in df.columns and q_col_candidates: + qc = q_col_candidates[0] + df['issue_time'] = pd.to_datetime(df['issue_time']) + out = df[['issue_time', 'lead_hour', 'member', qc]].copy() + out = out.rename(columns={qc: 'q_gauge_m3s'}) + return out + + # Wide format: member_00..member_19 columns + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + if not member_cols: + raise ValueError(f"Couldn't identify member columns in {path}. " + f"Columns present: {df.columns.tolist()}") + df['issue_time'] = pd.to_datetime(df['issue_time']) + keep = ['issue_time', 'lead_hour'] + if 'valid_time' in df.columns: + keep.append('valid_time') + long = df[keep + member_cols].melt( + id_vars=keep, + value_vars=member_cols, + var_name='member', + value_name='q_gauge_m3s', + ) + return long + + +def load_usgs_obs(usgs_csv): + """Load USGS gauge obs and ensure m³/s. + + The CSV at gauge 03463300 has the column `QObs(mm/h)` (catchment-averaged + depth, not gauge discharge in m³/s). If the column name contains 'mm', we + convert mm/h → m³/s by multiplying by the watershed area factor: + m³/s = mm/h × WATERSHED_AREA_KM2 × 1000 / 3600 + Verified: 59.978 mm/h × 113.18 km² = 1886 m³/s (USGS Helene peak). + """ + df = pd.read_csv(usgs_csv) + date_candidates = [c for c in df.columns if c.lower() in + ('datetime', 'date', 'time', 'timestamp')] + q_candidates = [c for c in df.columns if 'q' in c.lower() + or 'flow' in c.lower() or 'discharge' in c.lower()] + if not date_candidates or not q_candidates: + raise ValueError(f"Could not identify date/Q columns in {usgs_csv}. " + f"Columns: {df.columns.tolist()}") + dc, qc = date_candidates[0], q_candidates[0] + df[dc] = pd.to_datetime(df[dc]) + series = df.set_index(dc)[qc].astype(float) + if 'mm' in qc.lower(): + series = series * MM_H_TO_M3_S + print(f" Converted obs from mm/h (column '{qc}') to m³/s " + f"using area = {WATERSHED_AREA_KM2} km² (× {MM_H_TO_M3_S:.4f})") + else: + print(f" Loaded obs as-is from column '{qc}' (assumed m³/s)") + return series + + +def metrics_by_lead(df_long, obs_series, issue_mask=None): + """Compute per-lead-hour: ensemble-mean RMSE, per-member min/max RMSE, + mean ensemble std. df_long must have issue_time, lead_hour, member, q_gauge_m3s. + """ + df = df_long.copy() + if issue_mask is not None: + df = df[df['issue_time'].isin(issue_mask)] + if len(df) == 0: + return None + df['valid_time'] = df['issue_time'] + pd.to_timedelta(df['lead_hour'], unit='h') + df['obs'] = df['valid_time'].map(obs_series) + df = df.dropna(subset=['obs']) + if len(df) == 0: + return None + + # Ensemble mean per (issue_time, lead_hour) + grouped = df.groupby(['issue_time', 'lead_hour']) + ens_mean = grouped['q_gauge_m3s'].mean().reset_index(name='ens_mean') + ens_std = grouped['q_gauge_m3s'].std().reset_index(name='ens_std') + obs_per = grouped['obs'].first().reset_index(name='obs') + panel = ens_mean.merge(ens_std, on=['issue_time', 'lead_hour']).merge( + obs_per, on=['issue_time', 'lead_hour']) + + leads = sorted(panel['lead_hour'].unique()) + rmse_mean = np.full(len(leads), np.nan) + rmse_lo = np.full(len(leads), np.nan) + rmse_hi = np.full(len(leads), np.nan) + std_mean = np.full(len(leads), np.nan) + for i, L in enumerate(leads): + psub = panel[panel['lead_hour'] == L] + if len(psub) == 0: + continue + err = psub['ens_mean'].values - psub['obs'].values + rmse_mean[i] = float(np.sqrt(np.mean(err ** 2))) + # Per-member RMSE — recompute from the underlying long frame + dsub = df[df['lead_hour'] == L] + member_rmses = [] + for m, g in dsub.groupby('member'): + err_m = g['q_gauge_m3s'].values - g['obs'].values + if len(err_m) >= 1: + member_rmses.append(float(np.sqrt(np.mean(err_m ** 2)))) + if member_rmses: + rmse_lo[i] = float(np.nanmin(member_rmses)) + rmse_hi[i] = float(np.nanmax(member_rmses)) + std_mean[i] = float(np.nanmean(psub['ens_std'].values)) + return np.asarray(leads), rmse_mean, rmse_lo, rmse_hi, std_mean + + +def plot_two_panel(out_path, leads_da, da_metrics, leads_ol, ol_metrics, title): + fig, (ax_top, ax_bot) = plt.subplots(2, 1, figsize=(10, 10), sharex=True, + gridspec_kw={"height_ratios": [1.4, 1.0]}) + leads, rmse_da, rmse_da_lo, rmse_da_hi, std_da = da_metrics + _, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = ol_metrics + + ax_top.fill_between(leads, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2, + label="DA — member spread (min/max RMSE)") + ax_top.fill_between(leads, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2, + label="Open-loop — member spread") + ax_top.plot(leads, rmse_da, color=DA_COLOR, lw=2.4, marker='o', + zorder=4, label="DA on (ensemble-mean RMSE)") + ax_top.plot(leads, rmse_ol, color=OL_COLOR, lw=2.4, marker='s', + linestyle='--', zorder=4, label="Open-loop (ensemble-mean RMSE)") + ax_top.set_ylabel("RMSE vs USGS obs (m³/s)", fontsize=11) + ax_top.set_title(title, fontsize=12) + ax_top.grid(True, alpha=0.25) + ax_top.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + ax_bot.plot(leads, std_da, color=DA_COLOR, lw=2.2, marker='o', + label="DA — mean ensemble std") + ax_bot.plot(leads, std_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', label="Open-loop — mean ensemble std") + ax_bot.set_xlabel("Forecast lead time (hours)", fontsize=11) + ax_bot.set_ylabel("Mean ensemble std (m³/s)", fontsize=11) + ax_bot.set_xticks(np.arange(1, 19)) + ax_bot.grid(True, alpha=0.25) + ax_bot.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + plt.tight_layout() + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + +def plot_three_panel_regime(out_path, regime_metrics): + fig, axes = plt.subplots(2, 3, figsize=(18, 10), + gridspec_kw={"height_ratios": [1.4, 1.0]}) + for col, (label, n_issue, da_m, ol_m) in enumerate(regime_metrics): + ax_rmse = axes[0, col] + ax_std = axes[1, col] + if da_m is None or ol_m is None: + ax_rmse.set_title(f"{label}\n(no issue times match)", fontsize=11) + ax_rmse.axis('off'); ax_std.axis('off') + continue + leads, rmse_da, rmse_da_lo, rmse_da_hi, std_da = da_m + _, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = ol_m + ax_rmse.fill_between(leads, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2) + ax_rmse.fill_between(leads, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2) + ax_rmse.plot(leads, rmse_da, color=DA_COLOR, lw=2.2, marker='o', + zorder=4, label="DA on") + ax_rmse.plot(leads, rmse_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', zorder=4, label="Open-loop") + ax_rmse.set_title(f"{label}\n({n_issue} issue times)", fontsize=11) + ax_rmse.set_ylabel("RMSE (m³/s)", fontsize=10) + ax_rmse.grid(True, alpha=0.25) + ax_rmse.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + ax_std.plot(leads, std_da, color=DA_COLOR, lw=2.0, marker='o', label="DA") + ax_std.plot(leads, std_ol, color=OL_COLOR, lw=2.0, marker='s', + linestyle='--', label="Open-loop") + ax_std.set_xlabel("Forecast lead time (hours)", fontsize=10) + ax_std.set_ylabel("Mean ensemble std (m³/s)", fontsize=10) + ax_std.set_xticks(np.arange(1, 19)) + ax_std.grid(True, alpha=0.25) + ax_std.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + fig.suptitle("Gauge-level lead-time error decay by regime — USGS 03463300", + fontsize=13, y=0.995) + plt.tight_layout(rect=(0, 0, 1, 0.97)) + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--route-dir', default=DEFAULT_ROUTE_DIR, + help='Dir holding routed_leadtime_{da,openloop}_full.parquet') + parser.add_argument('--out-dir', default=None, + help='Where to write the output PNGs. Defaults to --route-dir ' + '(which may not be writable if owned by another user — ' + 'pass an explicit path then).') + parser.add_argument('--usgs-csv', default=DEFAULT_USGS_CSV) + parser.add_argument('--da-name', default='routed_leadtime_da_full.parquet') + parser.add_argument('--ol-name', default='routed_leadtime_openloop_full.parquet') + args = parser.parse_args() + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + da_path = os.path.join(args.route_dir, args.da_name) + ol_path = os.path.join(args.route_dir, args.ol_name) + df_da = load_parquet_long(da_path) + df_ol = load_parquet_long(ol_path) + obs_series = load_usgs_obs(args.usgs_csv) + + print(f"DA rows: {len(df_da):,} | unique issue_times: {df_da['issue_time'].nunique()}") + print(f"OL rows: {len(df_ol):,} | unique issue_times: {df_ol['issue_time'].nunique()}") + print(f"USGS obs range: {obs_series.index.min()} .. {obs_series.index.max()} " + f"({len(obs_series):,} hours)") + print(f"USGS Helene peak (reference): {USGS_HELENE_PEAK_M3S:.0f} m³/s") + + # ----- Pooled (all issue times) ----- + da_pooled = metrics_by_lead(df_da, obs_series) + ol_pooled = metrics_by_lead(df_ol, obs_series) + out_pooled = os.path.join(out_dir, "lead_time_decay_gauge_pooled.png") + plot_two_panel(out_pooled, da_pooled[0], da_pooled, ol_pooled[0], ol_pooled, + title="Gauge-level lead-time decay — USGS 03463300\n" + "All issue times pooled (Oct 2023 – Oct 2024)") + + # ----- Regime split: Helene, storm, low-flow at gauge ----- + issue_times = pd.to_datetime(sorted(df_da['issue_time'].unique())) + obs_at_issue = pd.Series(issue_times, index=issue_times).map(obs_series) + + helene_mask = (issue_times >= HELENE_START) & (issue_times <= HELENE_END) + storm_mask = obs_at_issue > STORM_OBS_THRESHOLD_M3S + lowflow_mask = obs_at_issue < LOWFLOW_OBS_THRESHOLD_M3S + + regime_metrics = [] + for label, mask in [ + (f"Helene window (Sep 24–28 2024)", helene_mask), + (f"Storm hours (obs(t0) > {STORM_OBS_THRESHOLD_M3S:.0f} m³/s)", storm_mask), + (f"Low flow (obs(t0) < {LOWFLOW_OBS_THRESHOLD_M3S:.0f} m³/s)", lowflow_mask), + ]: + kept = set(pd.to_datetime(issue_times[mask])) + da_m = metrics_by_lead(df_da, obs_series, kept) + ol_m = metrics_by_lead(df_ol, obs_series, kept) + regime_metrics.append((label, len(kept), da_m, ol_m)) + + out_regime = os.path.join(out_dir, "lead_time_decay_gauge_by_regime.png") + plot_three_panel_regime(out_regime, regime_metrics) + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4a_error_decay/run_4a_f2.sh b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4a_error_decay/run_4a_f2.sh new file mode 100644 index 00000000..8044d7c2 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4a_error_decay/run_4a_f2.sh @@ -0,0 +1,78 @@ +#!/usr/bin/env bash +# F2 Fixed R=0.07 (1 gauge holdout) — 4a lead-time decay plots. +# +# Runs 3 plot scripts per catchment: +# plot_lead_time_decay.py — pooled RMSE vs lead (catchment-level) +# plot_lead_time_decay_by_regime.py — same split by flow regime +# plot_lead_time_decay_gauge.py — gauge-level (requires routed parquets) +# +# plot_lead_time_decay_gauge.py is skipped if the routed parquets don't exist. +# +# Usage: +# bash run_4a_f2.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +LEADTIME_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder2_fixed_r007 +DA_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder2_fixed_r007 +OBS_DIR=/home/svyas/catchment_ts_no_03463300_gapfilled +ROUTE_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder2_fixed_r007_leadtime_routed +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +echo "[F2-4a] Lead-time decay plots — leadtime dir: $LEADTIME_DIR" + +for CAT in "${CATS[@]}"; do + DA_CSV="$LEADTIME_DIR/$CAT/${CAT}_lead_time_forecasts_da.csv" + if [ ! -f "$DA_CSV" ]; then + echo " [$CAT] No lead-time CSV — skipping" + continue + fi + + echo " [$CAT] Plotting pooled decay..." + $TROUTE "$SCRIPT_DIR/plot_lead_time_decay.py" \ + --cat-id "$CAT" \ + --leadtime-dir "$LEADTIME_DIR" \ + --da-dir "$DA_DIR" \ + --obs-dir "$OBS_DIR" + + echo " [$CAT] Plotting decay by regime..." + $TROUTE "$SCRIPT_DIR/plot_lead_time_decay_by_regime.py" \ + --cat-id "$CAT" \ + --leadtime-dir "$LEADTIME_DIR" \ + --da-dir "$DA_DIR" \ + --obs-dir "$OBS_DIR" +done + +# Gauge-level decay (needs routed parquets) +if [ -f "$ROUTE_DIR/routed_leadtime_da_full.parquet" ]; then + echo "[F2-4a] Gauge-level decay (routed)..." + $TROUTE "$SCRIPT_DIR/plot_lead_time_decay_gauge.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" + + echo "[F2-4a] Fixed-target error plots..." + $TROUTE "$SCRIPT_DIR/plot_forecast_error_fixed_target.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" + + echo "[F2-4a] Per-init error plots..." + $TROUTE "$SCRIPT_DIR/plot_forecast_error_per_init.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" +else + echo "[F2-4a] Routed parquets not found — skipping gauge-level and error plots." + echo " Run route_leadtime_f2.sh first, then re-run this script." +fi + +echo "[F2-4a] Done." diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/plot_crossed_ensemble.py b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/plot_crossed_ensemble.py new file mode 100644 index 00000000..7cb00a93 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/plot_crossed_ensemble.py @@ -0,0 +1,229 @@ +""" +plot_crossed_ensemble.py -- 4b: 600-member crossed ensemble vs USGS obs. + +Loads the crossed ensemble parquet produced by run_crossed_ensemble.py: + //_crossed_ensemble.parquet + Columns: issue_time, lead_hour, member_0000..member_0599 + +Projects all members to valid_time = issue_time + lead_hour hours. +Aggregates at each valid_time across ALL members from ALL issue times +to build a probabilistic envelope (5th/25th/50th/75th/95th percentiles). + +Two output figures: + 4b_crossed_ensemble_helene.png + -- Full Helene window (Sep 24-30) probabilistic envelope + USGS obs. + Percentile shading in two layers (5-95 outer, 25-75 inner). + 4b_crossed_ensemble_peak.png + -- Zoomed to Helene peak (Sep 26 12 UTC -> Sep 28 06 UTC), + same shading + USGS obs + individual spaghetti from Sep 27 inits. + +Usage: + python3 plot_crossed_ensemble.py \\ + --ensemble-dir /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble \\ + --cat-id cat-1016300 \\ + --usgs-csv /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv \\ + --out-dir /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble/cat-1016300 +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import matplotlib.patches as mpatches + +DEFAULT_ENSEMBLE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" +DEFAULT_CAT_ID = "cat-1016300" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +PLOT_START = pd.Timestamp("2024-09-24 00:00:00") +PLOT_END = pd.Timestamp("2024-09-28 18:00:00") # cut trailing spike +PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +PEAK_END = pd.Timestamp("2024-09-28 06:00:00") +HELENE_PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +HELENE_PEAK_END = pd.Timestamp("2024-09-28 00:00:00") + + +def load_ensemble(path): + df = pd.read_parquet(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + return df, member_cols + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def build_percentile_envelope(df, member_cols, t_start, t_end): + """ + For each valid_time in [t_start, t_end], stack all member values + from all issue times and compute percentiles. + Returns DataFrame indexed by valid_time with columns p05, p25, p50, p75, p95. + """ + mask = (df["valid_time"] >= t_start) & (df["valid_time"] <= t_end) + sub = df[mask].copy() + sub["valid_time"] = sub["valid_time"].dt.floor("1h") + + records = [] + for vt, grp in sub.groupby("valid_time"): + vals = grp[member_cols].to_numpy(dtype=float).flatten() + vals = vals[~np.isnan(vals)] * MM_H_TO_M3_S + if len(vals) == 0: + continue + records.append({ + "valid_time": vt, + "p05": np.percentile(vals, 5), + "p25": np.percentile(vals, 25), + "p50": np.percentile(vals, 50), + "p75": np.percentile(vals, 75), + "p95": np.percentile(vals, 95), + "n": len(vals), + }) + return pd.DataFrame(records).set_index("valid_time").sort_index() + + +def _add_helene_band(ax, label=True): + ax.axvspan(HELENE_PEAK_START, HELENE_PEAK_END, + color="salmon", alpha=0.12, zorder=0) + if label: + ax.text(HELENE_PEAK_START + pd.Timedelta(hours=6), 0.97, "Helene peak", + transform=ax.get_xaxis_transform(), + fontsize=9, color="firebrick", ha="left", va="top") + + +def _format_xaxis(ax, major_interval=1, minor_hours=6): + ax.xaxis.set_major_locator(mdates.DayLocator(interval=major_interval)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.xaxis.set_minor_locator(mdates.HourLocator(interval=minor_hours)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.2, lw=0.4) + + +def plot_envelope(ax, env, obs_series, t_start, t_end, + ens_color="tab:blue", obs_color="black", helene_label=True): + _add_helene_band(ax, label=helene_label) + + # Outer band: 5-95 + ax.fill_between(env.index, env["p05"], env["p95"], + color=ens_color, alpha=0.15, zorder=2, + label="5th-95th percentile") + # Inner band: 25-75 + ax.fill_between(env.index, env["p25"], env["p75"], + color=ens_color, alpha=0.30, zorder=3, + label="25th-75th percentile") + # Median + ax.plot(env.index, env["p50"], + color=ens_color, lw=2.2, zorder=5, + label="Ensemble median (600 members)") + + # USGS obs + obs_w = obs_series.loc[t_start:t_end] + ax.plot(obs_w.index, obs_w.values, + color=obs_color, lw=1.8, zorder=6, label="USGS obs") + + ax.set_xlim(t_start, t_end) + ax.set_ylabel("Discharge (m3/s)", fontsize=11) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--ensemble-dir", default=DEFAULT_ENSEMBLE_DIR) + parser.add_argument("--cat-id", default=DEFAULT_CAT_ID) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--out-dir", default=None) + args = parser.parse_args() + + cat_dir = os.path.join(args.ensemble_dir, args.cat_id) + out_dir = args.out_dir or cat_dir + os.makedirs(out_dir, exist_ok=True) + + parquet_path = os.path.join(cat_dir, f"{args.cat_id}_crossed_ensemble.parquet") + print(f"Loading crossed ensemble: {parquet_path}") + df, member_cols = load_ensemble(parquet_path) + obs = load_usgs(args.usgs_csv) + print(f" {df['issue_time'].nunique()} issue times, " + f"{len(member_cols)} members per row") + + # ------------------------------------------------------------------ # + # Figure 4b-full: full Helene window probabilistic envelope # + # ------------------------------------------------------------------ # + print("Building percentile envelope (full window)...") + env_full = build_percentile_envelope(df, member_cols, PLOT_START, PLOT_END) + print(f" Valid times with data: {len(env_full)}, " + f"median members per timestep: {env_full['n'].median():.0f}") + + fig, ax = plt.subplots(figsize=(17, 7)) + plot_envelope(ax, env_full, obs, PLOT_START, PLOT_END) + ax.set_xlabel("Date", fontsize=11) + ax.set_title( + f"4b -- 600-member crossed ensemble (30 met x 20 hydro-state) | {args.cat_id}\n" + "Shaded: 5-95th (light) and 25-75th (dark) percentile across all members " + "and issue times | Sep 24-30 2024 | USGS 03463300", + fontsize=11, + ) + ax.legend(fontsize=10, loc="upper left", frameon=True, framealpha=0.9) + _format_xaxis(ax) + plt.tight_layout() + out1 = os.path.join(out_dir, f"{args.cat_id}_4b_crossed_ensemble_helene.png") + plt.savefig(out1, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + # ------------------------------------------------------------------ # + # Figure 4b-peak: zoomed to Helene peak window # + # ------------------------------------------------------------------ # + print("Building percentile envelope (peak window)...") + env_peak = build_percentile_envelope(df, member_cols, PEAK_START, PEAK_END) + + fig, (ax_main, ax_zoom) = plt.subplots( + 2, 1, figsize=(14, 11), + gridspec_kw={"height_ratios": [1, 1.4]}, sharex=False) + + # Top panel: full window context (smaller) + plot_envelope(ax_main, env_full, obs, PLOT_START, PLOT_END, helene_label=False) + ax_main.set_title("Full Helene window context (Sep 24 - Sep 28 18 UTC)", fontsize=10) + _format_xaxis(ax_main) + # Shade the zoom region on context panel + ax_main.axvspan(PEAK_START, PEAK_END, color="gold", alpha=0.18, zorder=1) + ax_main.text(PEAK_START + pd.Timedelta(hours=1), 0.97, "zoom", + transform=ax_main.get_xaxis_transform(), + fontsize=8, color="goldenrod", ha="left", va="top") + + # Bottom panel: peak zoom + plot_envelope(ax_zoom, env_peak, obs, PEAK_START, PEAK_END, helene_label=True) + ax_zoom.set_xlabel("Date", fontsize=11) + ax_zoom.set_title( + "Helene peak zoom (Sep 26 12 UTC - Sep 28 06 UTC)", fontsize=10) + _format_xaxis(ax_zoom, major_interval=1, minor_hours=3) + ax_zoom.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %HUTC")) + ax_zoom.legend(fontsize=10, loc="lower right", frameon=True, framealpha=0.9) + + fig.suptitle( + f"4b -- 600-member crossed ensemble vs USGS 03463300 | {args.cat_id}\n" + "30 met forcing draws x 20 DA analysis state draws | R=0.07 mm2/h2", + fontsize=12, y=1.01, + ) + plt.tight_layout() + out2 = os.path.join(out_dir, f"{args.cat_id}_4b_crossed_ensemble_peak.png") + plt.savefig(out2, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/plot_perturbation_category_shaded.py b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/plot_perturbation_category_shaded.py new file mode 100644 index 00000000..4ca72b26 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/plot_perturbation_category_shaded.py @@ -0,0 +1,160 @@ +""" +Per-catchment shaded ensemble-band plot, organized by perturbation category. +F2 fixed R=0.07 (1 gauge holdout). + +Three categories: + 1. Initial states (red) + 2. Meteorological forcings (blue) + 3. Hydrological states (green) + +For each category, all 20 ensemble members are shown as a shaded band +(min-max envelope fill) plus a thicker median line in the same color. +Qkrig observation overlaid in black. Hurricane Helene peak window shaded +in pink. Styled after a standard ensemble-forecast figure layout. + +Inputs (existing per-source sensitivity CSVs from run_perturbation_sensitivity.py): + //_sensitivity_init.csv (20 members, init only) + //_sensitivity_forcing.csv (20 members, forcing only) + //_sensitivity_process.csv (20 members, process noise only) + /.csv (Qkrig obs, gapfilled) + +Outputs: + //_perturbation_categories_linear.png + //_perturbation_categories_log.png +""" +import argparse +import os +import pandas as pd +import numpy as np +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +CAT = "cat-1016300" + +SEN_DIR = "/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder2_fixed_r007" +OBS_DIR = "/home/svyas/catchment_ts_no_03463300_gapfilled" + +OUT_LINEAR = os.path.join(SEN_DIR, CAT, f"{CAT}_perturbation_categories_linear.png") +OUT_LOG = os.path.join(SEN_DIR, CAT, f"{CAT}_perturbation_categories_log.png") + +PLOT_START = pd.Timestamp("2024-09-20") +PLOT_END = pd.Timestamp("2024-10-05") + +HELENE_START = pd.Timestamp("2024-09-26 12:00:00") +HELENE_END = pd.Timestamp("2024-09-28 00:00:00") + +CATEGORIES = [ + ("init", "Initial states only", "tab:red"), + ("forcing", "Meteorological forcings only", "tab:blue"), + ("process", "Hydrological states only", "tab:green"), +] + + +def load_members(source): + path = os.path.join(SEN_DIR, CAT, f"{CAT}_sensitivity_{source}.csv") + if not os.path.exists(path): + return None, None + df = pd.read_csv(path, parse_dates=["date"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + return df["date"].values, df[member_cols].to_numpy(dtype=float) + + +def load_obs(): + p = os.path.join(OBS_DIR, f"{CAT}.csv") + if not os.path.exists(p): + return None, None + df = pd.read_csv(p) + time_col = next((c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")), + df.columns[0]) + df[time_col] = pd.to_datetime(df[time_col]) + q_col = next((c for c in df.columns + if "qkrig" in c.lower() and "var" not in c.lower()), None) + if q_col is None: + q_col = next(c for c in df.columns + if c != time_col and pd.api.types.is_numeric_dtype(df[c])) + return df[time_col].values, df[q_col].values + + +def plot_panel(ax, obs_dates, obs_vals, log_y=False): + for source, label, color in CATEGORIES: + dates, q = load_members(source) + if dates is None: + continue + d = pd.to_datetime(dates) + mask = (d >= PLOT_START) & (d <= PLOT_END) + if mask.sum() == 0: + continue + q_window = q[mask, :] + qmin = np.nanmin(q_window, axis=1) + qmax = np.nanmax(q_window, axis=1) + median = np.nanmedian(q_window, axis=1) + + ax.fill_between(d[mask], qmin, qmax, + color=color, alpha=0.30, zorder=2, edgecolor="none") + ax.plot(d[mask], median, + color=color, lw=1.7, alpha=0.95, zorder=3, + label=f"{label} (N=20)") + + ax.axvspan(HELENE_START, HELENE_END, color="salmon", alpha=0.15, zorder=1) + + if obs_dates is not None: + od = pd.to_datetime(obs_dates) + om = (od >= PLOT_START) & (od <= PLOT_END) + ax.plot(od[om], obs_vals[om], + color="black", lw=1.4, label="Qkrig (obs)", zorder=4) + + ax.set_xlabel("Date", fontsize=10) + if log_y: + ax.set_yscale("log") + ax.set_ylabel(r"q (mm hr$^{-1}$) [log scale]", fontsize=10) + ax.set_ylim(0.01, None) + else: + ax.set_ylabel("Discharge (mm/h)", fontsize=10) + ax.tick_params(axis="x", rotation=30, labelsize=9) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=2)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.grid(True, alpha=0.2) + ax.legend(loc="upper left", fontsize=9, frameon=True, framealpha=0.9) + + +def main(): + global CAT, SEN_DIR, OUT_LINEAR, OUT_LOG + parser = argparse.ArgumentParser() + parser.add_argument("--cat-id", default=CAT) + args = parser.parse_args() + CAT = args.cat_id + OUT_LINEAR = os.path.join(SEN_DIR, CAT, f"{CAT}_perturbation_categories_linear.png") + OUT_LOG = os.path.join(SEN_DIR, CAT, f"{CAT}_perturbation_categories_log.png") + + obs_dates, obs_vals = load_obs() + + fig, ax = plt.subplots(1, 1, figsize=(20, 7)) + plot_panel(ax, obs_dates, obs_vals, log_y=False) + fig.suptitle( + f"Ensemble forecast by perturbation category — {CAT} — F2 (fixed R=0.07)\n" + "Sep 20 - Oct 5, 2024 (Hurricane Helene window) | " + "Shaded bands = min-max envelope across 20 members. Lines = ensemble median.", + fontsize=11, y=0.995, + ) + plt.tight_layout() + plt.savefig(OUT_LINEAR, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_LINEAR}") + + fig, ax = plt.subplots(1, 1, figsize=(20, 7)) + plot_panel(ax, obs_dates, obs_vals, log_y=True) + fig.suptitle( + f"Ensemble forecast by perturbation category — {CAT} — F2 (fixed R=0.07) — log scale\n" + "Sep 20 - Oct 5, 2024 (Hurricane Helene window) | " + "Shaded bands = min-max envelope across 20 members. Lines = ensemble median.", + fontsize=11, y=0.995, + ) + plt.tight_layout() + plt.savefig(OUT_LOG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_LOG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_combined.py b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_combined.py new file mode 100644 index 00000000..c6d531a0 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_combined.py @@ -0,0 +1,235 @@ +#!/usr/bin/env python3 +""" +plot_routed_ensemble_combined.py +Combined poster plot: 600-member routed ensemble envelope + +deterministic DA lines (Vrugt/no-Vrugt) vs USGS gauge obs at 03463300. + +Inputs: + --vrugt-csv : routed_Q_test.csv from vrugt_dynamic_routed/ + --novrugt-csv : routed_Q_test.csv from novrugt_dynamic_routed/ + --ensemble-pq : routed_crossed_ensemble.parquet from run_route_crossed_ensemble.py + --out-dir : output directory + +The ensemble parquet columns: issue_time, lead_hour, member_0000..member_0599 (Q in m3/s) + +Usage: + python3 plot_routed_ensemble_combined.py \ + --vrugt-csv /mnt/disk2/suma_helen_poster/da_results/vrugt_dynamic_routed/routed_Q_test.csv \ + --novrugt-csv /mnt/disk2/suma_helen_poster/da_results/novrugt_dynamic_routed/routed_Q_test.csv \ + --ensemble-pq /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble_routed/routed_crossed_ensemble.parquet \ + --out-dir /mnt/disk2/suma_helen_poster/da_results/comparison_plots +""" + +import argparse +import os + +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +HELENE_START = pd.Timestamp("2024-09-25 00:00:00") +HELENE_END = pd.Timestamp("2024-09-29 00:00:00") +ZOOM_START = pd.Timestamp("2024-09-24 00:00:00") +ZOOM_END = pd.Timestamp("2024-09-28 18:00:00") # cut trailing ensemble bump + +COLOR_OBS = "black" +COLOR_VRUGT = "#1f77b4" # blue +COLOR_NOVRUGT = "#d62728" # red +COLOR_ENS = "#1f77b4" # same blue family as Vrugt + + +def load_det_csv(path): + df = pd.read_csv(path, parse_dates=["date"]) + return df.set_index("date").sort_index() + + +def build_envelope(pq_path, zoom_start, zoom_end): + """ + Load routed ensemble parquet and build percentile envelope by valid_time. + Returns DataFrame indexed by valid_time with columns p05, p25, p50, p75, p95. + """ + df = pd.read_parquet(pq_path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = [c for c in df.columns if c.startswith("member_")] + + records = [] + for _, row in df.iterrows(): + t0 = row["issue_time"] + lead = int(row["lead_hour"]) + vt = t0 + pd.Timedelta(hours=lead) + vals = row[member_cols].values.astype(float) + records.append({"valid_time": vt, "vals": vals}) + + # Group by valid_time (floor to hour), stack all member values across issue_times + from collections import defaultdict + groups = defaultdict(list) + for rec in records: + vt = rec["valid_time"].floor("1h") + groups[vt].extend(rec["vals"].tolist()) + + rows = [] + for vt in sorted(groups): + if vt < zoom_start or vt > zoom_end: + continue + v = np.array(groups[vt]) + v = v[~np.isnan(v)] + if len(v) == 0: + continue + rows.append({ + "valid_time": vt, + "p05": np.percentile(v, 5), + "p25": np.percentile(v, 25), + "p50": np.percentile(v, 50), + "p75": np.percentile(v, 75), + "p95": np.percentile(v, 95), + }) + + return pd.DataFrame(rows).set_index("valid_time") + + +def compute_kge(obs, sim): + mask = ~(np.isnan(obs) | np.isnan(sim)) + o, s = obs[mask], sim[mask] + if len(o) < 2 or np.std(o) == 0: + return np.nan + r = np.corrcoef(o, s)[0, 1] + return 1.0 - np.sqrt((r-1)**2 + (np.std(s)/np.std(o)-1)**2 + (np.mean(s)/np.mean(o)-1)**2) + + +def _format_xaxis(ax, locator, fmt, rotation=30): + ax.xaxis.set_major_locator(locator) + ax.xaxis.set_major_formatter(mdates.DateFormatter(fmt)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=rotation, ha="right", fontsize=9) + + +def _add_helene_band(ax): + ax.axvspan(HELENE_START, HELENE_END, color="gold", alpha=0.15, zorder=0) + + +def plot_zoom(ax, dates, obs, sim_v, sim_nv, kge_v, kge_nv, env=None): + mask = (dates >= ZOOM_START) & (dates <= ZOOM_END) + dz = dates[mask] + + # Ensemble envelope (bottom layer) + if env is not None: + ax.fill_between(env.index, env["p05"], env["p95"], + color=COLOR_ENS, alpha=0.18, label="Ensemble 5-95th pct") + ax.fill_between(env.index, env["p25"], env["p75"], + color=COLOR_ENS, alpha=0.30, label="Ensemble 25-75th pct") + ax.plot(env.index, env["p50"], + color=COLOR_ENS, lw=1.2, linestyle="--", alpha=0.7, label="Ensemble median") + + # Deterministic DA lines + ax.plot(dz, sim_nv[mask], color=COLOR_NOVRUGT, lw=1.4, linestyle="--", + label=f"DA -- constant R KGE={kge_nv:.3f}", zorder=3) + ax.plot(dz, sim_v[mask], color=COLOR_VRUGT, lw=1.6, + label=f"DA -- dynamic R (Vrugt) KGE={kge_v:.3f}", zorder=4) + + # USGS obs on top + ax.plot(dz, obs[mask], color=COLOR_OBS, lw=1.6, label="USGS obs (gauge 03463300)", zorder=5) + + peak_usgs = np.nanmax(obs[mask]) + ax.axhline(peak_usgs, color=COLOR_OBS, lw=0.7, linestyle=":", alpha=0.5) + # Place annotation to the right of the peak to avoid legend overlap + ax.text(pd.Timestamp("2024-09-28 06:00:00"), peak_usgs * 1.02, + f"USGS peak {peak_usgs:.0f} m³/s", fontsize=8.5, color="black", + alpha=0.8, ha="right") + + _add_helene_band(ax) + _format_xaxis(ax, mdates.DayLocator(interval=1), "%b %d") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--vrugt-csv", required=True) + parser.add_argument("--novrugt-csv", required=True) + parser.add_argument("--ensemble-pq", required=True, + help="routed_crossed_ensemble.parquet from run_route_crossed_ensemble.py") + parser.add_argument("--out-dir", required=True) + args = parser.parse_args() + os.makedirs(args.out_dir, exist_ok=True) + + vrugt = load_det_csv(args.vrugt_csv) + novrugt = load_det_csv(args.novrugt_csv) + + obs = vrugt["Q_usgs_m3s"].values + dates = vrugt.index + sim_v = vrugt["Q_routed_m3s"].values + sim_nv = novrugt["Q_routed_m3s"].reindex(dates).values + + kge_v = compute_kge(obs, sim_v) + kge_nv = compute_kge(obs, sim_nv) + + print("Building ensemble envelope...") + env = build_envelope(args.ensemble_pq, ZOOM_START, ZOOM_END) + print(f" Envelope: {len(env)} timesteps, " + f"peak p95={env['p95'].max():.1f} m3/s, peak p50={env['p50'].max():.1f} m3/s") + + peak_usgs = np.nanmax(obs[(dates >= ZOOM_START) & (dates <= ZOOM_END)]) + in_env = env["p95"].max() >= peak_usgs + print(f" USGS peak {peak_usgs:.1f} m3/s {'IS' if in_env else 'IS NOT'} within p95 envelope") + + # ------------------------------------------------------------------ # + # Figure 1: Helene zoom only -- main poster panel + # ------------------------------------------------------------------ # + fig, ax = plt.subplots(figsize=(11, 5.5)) + plot_zoom(ax, dates, obs, sim_v, sim_nv, kge_v, kge_nv, env) + ax.set_xlabel("Date (UTC)", fontsize=11) + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_title( + "Hurricane Helene -- Probabilistic discharge forecast at gauge 03463300 (South Toe River)\n" + "600-member CFE ensemble + DA (Muskingum routing)", + fontsize=11) + ax.legend(fontsize=9, loc="lower right") + ax.grid(True, alpha=0.25) + plt.tight_layout() + out1 = os.path.join(args.out_dir, "helene_ensemble_vs_usgs.png") + plt.savefig(out1, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + # ------------------------------------------------------------------ # + # Figure 2: two-panel (full period + Helene zoom) + # ------------------------------------------------------------------ # + fig, (ax_top, ax_bot) = plt.subplots( + 2, 1, figsize=(14, 9), + gridspec_kw={"height_ratios": [1, 1.5]}) + + # Top: full test period (deterministic only, no ensemble for readability) + ax_top.plot(dates, obs, color=COLOR_OBS, lw=0.8, label="USGS obs", zorder=3) + ax_top.plot(dates, sim_v, color=COLOR_VRUGT, lw=0.9, + label=f"DA dynamic R (Vrugt) KGE={kge_v:.3f}", zorder=2) + ax_top.plot(dates, sim_nv, color=COLOR_NOVRUGT, lw=0.9, linestyle="--", + label=f"DA constant R KGE={kge_nv:.3f}", zorder=2) + ax_top.axvspan(ZOOM_START, ZOOM_END, color="gray", alpha=0.12, zorder=0) + ax_top.axvspan(HELENE_START, HELENE_END, color="gold", alpha=0.15, zorder=0) + ax_top.text(HELENE_START + pd.Timedelta(days=0.5), + np.nanmax(obs) * 0.88, "Helene", + fontsize=9, color="goldenrod", fontweight="bold") + ax_top.set_ylabel("Discharge (m³/s)", fontsize=10) + ax_top.set_title( + "Routed discharge at gauge 03463300 (South Toe River) -- Test period\n" + "DA with dynamic Vrugt R vs. constant R -- Muskingum routing", + fontsize=11) + ax_top.legend(fontsize=9, loc="upper left") + ax_top.grid(True, alpha=0.25) + _format_xaxis(ax_top, mdates.MonthLocator(interval=2), "%Y-%m") + + # Bottom: Helene zoom with ensemble + plot_zoom(ax_bot, dates, obs, sim_v, sim_nv, kge_v, kge_nv, env) + ax_bot.set_xlabel("Date (UTC)", fontsize=10) + ax_bot.set_ylabel("Discharge (m³/s)", fontsize=10) + ax_bot.set_title("Hurricane Helene window -- 600-member ensemble envelope", fontsize=10) + ax_bot.legend(fontsize=9, loc="lower right") + ax_bot.grid(True, alpha=0.25) + + plt.tight_layout() + out2 = os.path.join(args.out_dir, "helene_ensemble_twopanel.png") + plt.savefig(out2, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_vs_usgs.py b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_vs_usgs.py new file mode 100644 index 00000000..55c52d85 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_vs_usgs.py @@ -0,0 +1,176 @@ +#!/usr/bin/env python3 +""" +plot_routed_ensemble_vs_usgs.py +Plot T-route routed 600-member ensemble envelope vs USGS at gauge 03463300. + +Input: routed_crossed_ensemble.parquet from run_route_crossed_ensemble.py + Columns: issue_time, lead_hour, member_0000..member_0599 (Q in m3/s) + +Usage: + python3 plot_routed_ensemble_vs_usgs.py \ + --routed-pq /mnt/disk2/.../folder1_vrugt_routed/routed_crossed_ensemble.parquet \ + --usgs-csv /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv \ + --label "F1 Vrugt — 1 gauge holdout" \ + --out-dir ~/plots_f1_1gauge +""" + +import argparse +from collections import defaultdict +from pathlib import Path + +import matplotlib.dates as mdates +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd + +WATERSHED_AREA_KM2 = 113.18 # cat-1016300 drainage area +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +PLOT_START = pd.Timestamp("2024-09-24 00:00:00") +PLOT_END = pd.Timestamp("2024-09-28 18:00:00") +PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +PEAK_END = pd.Timestamp("2024-09-28 06:00:00") +HELENE_PEAK = pd.Timestamp("2024-09-27 12:00:00") + +COLOR_ENS = "#1f77b4" +COLOR_OBS = "black" + + +def build_envelope(pq_path, t_start, t_end): + df = pd.read_parquet(pq_path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = [c for c in df.columns if c.startswith("member_")] + + groups = defaultdict(list) + for _, row in df.iterrows(): + vt = (row["issue_time"] + pd.Timedelta(hours=int(row["lead_hour"]))).floor("1h") + if t_start <= vt <= t_end: + groups[vt].extend(row[member_cols].values.astype(float).tolist()) + + rows = [] + for vt in sorted(groups): + v = np.array(groups[vt]) + v = v[~np.isnan(v)] + if len(v) == 0: + continue + rows.append({"valid_time": vt, + "p05": np.percentile(v, 5), "p25": np.percentile(v, 25), + "p50": np.percentile(v, 50), "p75": np.percentile(v, 75), + "p95": np.percentile(v, 95)}) + return pd.DataFrame(rows).set_index("valid_time") + + +def load_usgs(path): + df = pd.read_csv(path) + print(f" USGS CSV columns: {list(df.columns)}") + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() or "discharge" in c.lower()) + print(f" Using date_col={date_col!r}, q_col={q_col!r}") + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series.sort_index() + + +def compute_kge(obs, sim): + mask = ~(np.isnan(obs) | np.isnan(sim)) + o, s = obs[mask], sim[mask] + if len(o) < 2 or np.std(o) == 0: + return np.nan + r = np.corrcoef(o, s)[0, 1] + return 1 - np.sqrt((r-1)**2 + (np.std(s)/np.std(o)-1)**2 + (np.mean(s)/np.mean(o)-1)**2) + + +def _shade(ax, env): + ax.fill_between(env.index, env["p05"], env["p95"], + color=COLOR_ENS, alpha=0.18, label="5th–95th percentile") + ax.fill_between(env.index, env["p25"], env["p75"], + color=COLOR_ENS, alpha=0.32, label="25th–75th percentile") + ax.plot(env.index, env["p50"], color=COLOR_ENS, lw=1.5, label="Ensemble median (600 members)") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--routed-pq", required=True) + parser.add_argument("--usgs-csv", required=True) + parser.add_argument("--label", default="F1 Vrugt — 1 gauge holdout") + parser.add_argument("--out-dir", required=True) + args = parser.parse_args() + + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + print("Building full-window envelope...") + env_full = build_envelope(args.routed_pq, PLOT_START, PLOT_END) + print("Building peak-window envelope...") + env_peak = build_envelope(args.routed_pq, PEAK_START, PEAK_END) + + print("Loading USGS obs...") + try: + obs_all = load_usgs(args.usgs_csv) + obs_full = obs_all.loc[PLOT_START:PLOT_END] + obs_peak = obs_all.loc[PEAK_START:PEAK_END] + print(f" Loaded {len(obs_all)} USGS obs rows; " + f"full window: {len(obs_full)}, peak window: {len(obs_peak)}") + except Exception as e: + print(f" Warning: USGS load failed ({e}) — plotting without obs") + obs_full = obs_peak = None + + kge = np.nan + if obs_full is not None and len(env_full) > 0: + med = env_full["p50"] + obs_hourly = obs_full.copy() + obs_hourly.index = obs_hourly.index.floor("1h") + obs_hourly = obs_hourly[~obs_hourly.index.duplicated(keep="first")] + aligned = obs_hourly.reindex(med.index) + print(f" KGE alignment: {aligned.notna().sum()} matched out of {len(med)} timesteps") + kge = compute_kge(aligned.values, med.values) + + # ── Figure 1: Full Helene window ────────────────────────────────────────── + fig, ax = plt.subplots(figsize=(12, 4.5)) + _shade(ax, env_full) + if obs_full is not None: + ax.plot(obs_full.index, obs_full.values, color=COLOR_OBS, lw=1.6, + label="USGS obs (gauge 03463300)") + ax.axvspan(PEAK_START, PEAK_END, color="salmon", alpha=0.12, zorder=0, label="Helene peak window") + ax.set_ylabel("Discharge (m³/s)") + ax.set_xlabel("Date") + ax.set_title(f"Routed 600-member ensemble vs USGS | {args.label}\n" + f"Ensemble median KGE = {kge:.3f} | Sep 24–28 2024") + ax.legend(fontsize=9, loc="upper left") + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + plt.tight_layout() + out = out_dir / "routed_ensemble_vs_usgs_full.png" + fig.savefig(out, dpi=150) + print(f"Saved: {out}") + plt.close(fig) + + # ── Figure 2: Peak zoom ─────────────────────────────────────────────────── + fig, ax = plt.subplots(figsize=(10, 4.5)) + _shade(ax, env_peak) + if obs_peak is not None: + ax.plot(obs_peak.index, obs_peak.values, color=COLOR_OBS, lw=1.6, + label="USGS obs (gauge 03463300)") + peak_val = float(obs_peak.max()) + ax.axhline(peak_val, color=COLOR_OBS, lw=0.7, ls=":", alpha=0.5) + ax.annotate(f"USGS peak\n{peak_val:.0f} m³/s", + xy=(HELENE_PEAK, peak_val), xytext=(10, -40), + textcoords="offset points", fontsize=8, color=COLOR_OBS) + ax.set_ylabel("Discharge (m³/s)") + ax.set_xlabel("Date") + ax.set_title(f"Helene peak zoom | {args.label}") + ax.legend(fontsize=9) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %HUTC")) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + plt.tight_layout() + out = out_dir / "routed_ensemble_vs_usgs_peak.png" + fig.savefig(out, dpi=150) + print(f"Saved: {out}") + plt.close(fig) + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/run_4b_crossed_f2.sh b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/run_4b_crossed_f2.sh new file mode 100644 index 00000000..cdccf171 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/run_4b_crossed_f2.sh @@ -0,0 +1,40 @@ +#!/usr/bin/env bash +# F2 Fixed-R=0.07 (1 gauge holdout) — 4b crossed ensemble plots. +# Runs plot_crossed_ensemble.py for all 21 catchments. +# Reads //_crossed_ensemble.parquet (already generated). +# +# Usage: +# bash run_4b_crossed_f2.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +DATA_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder2_fixed_r007 +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +echo "[F2-4b-crossed] Crossed ensemble plots — data: $DATA_DIR" + +for CAT in "${CATS[@]}"; do + PQ="$DATA_DIR/$CAT/${CAT}_crossed_ensemble.parquet" + if [ ! -f "$PQ" ]; then + echo " [$CAT] No crossed ensemble parquet — skipping" + continue + fi + echo " [$CAT] crossed ensemble..." + $TROUTE "$SCRIPT_DIR/plot_crossed_ensemble.py" \ + --ensemble-dir "$DATA_DIR" \ + --cat-id "$CAT" \ + --usgs-csv "$USGS_CSV" +done + +echo "[F2-4b-crossed] Done." diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/run_4b_f2.sh b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/run_4b_f2.sh new file mode 100644 index 00000000..2560f988 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/run_4b_f2.sh @@ -0,0 +1,39 @@ +#!/usr/bin/env bash +# F2 fixed R=0.07 (1 gauge holdout) — 4b routed ensemble plots. +# +# Runs: +# 1. plot_routed_ensemble_vs_usgs.py — F2 ensemble envelope vs USGS at outlet +# 2. plot_routed_ensemble_combined.py — F1 Vrugt vs F2 fixed R comparison +# +# Requires routed_Q_test.csv (both folders) and routed_crossed_ensemble.parquet +# (F2 dir) to be present. +# +# Usage: +# bash run_4b_f2.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +F1_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt +F2_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder2_fixed_r007 +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv +OUT_DIR="$F2_DIR" + +echo "[F2-4b] Routed ensemble vs USGS (F2 only)..." +$TROUTE "$SCRIPT_DIR/plot_routed_ensemble_vs_usgs.py" \ + --routed-pq "$F2_DIR/routed_crossed_ensemble.parquet" \ + --usgs-csv "$USGS_CSV" \ + --label "F2 Fixed R=0.07 — 1 gauge holdout" \ + --out-dir "$OUT_DIR" + +echo "[F2-4b] Combined comparison: F1 Vrugt vs F2 Fixed R..." +$TROUTE "$SCRIPT_DIR/plot_routed_ensemble_combined.py" \ + --vrugt-csv "$F1_DIR/routed_Q_test.csv" \ + --novrugt-csv "$F2_DIR/routed_Q_test.csv" \ + --ensemble-pq "$F2_DIR/routed_crossed_ensemble.parquet" \ + --out-dir "$OUT_DIR" + +echo "[F2-4b] Done. Outputs in: $OUT_DIR" +ls "$OUT_DIR"/*.png 2>/dev/null || echo " (no PNGs in out dir)" diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4c_timeseries/plot_forecast_spaghetti.py b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4c_timeseries/plot_forecast_spaghetti.py new file mode 100644 index 00000000..655edaef --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4c_timeseries/plot_forecast_spaghetti.py @@ -0,0 +1,176 @@ +""" +plot_forecast_spaghetti.py + +"All forecasts ending at this hour" — verification-time-fixed spaghetti view. + +Window: Sep 24 18 UTC -> Sep 30 06 UTC. +For every initialization time whose 18-hour forecast overlaps this window: + DA : shaded min/max band across 20 members + ensemble mean line + OL : ensemble mean line only (no shading), dashed + +Lines are colored by initialization DATE (7 colors, Sep 24-30) so the +temporal progression is readable. USGS obs overlaid in black. + +Inputs: + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + Wide format: issue_time, lead_hour, member_00..member_19 (q in m³/s) + +Obs: + /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +Outputs: + /forecast_spaghetti_helene.png +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import matplotlib.patches as mpatches + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast_hardcoded_r_routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +# Verification window — only trajectories whose valid_time falls here are shown +PLOT_START = pd.Timestamp("2024-09-24 18:00:00") +PLOT_END = pd.Timestamp("2024-09-30 06:00:00") + +HELENE_PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +HELENE_PEAK_END = pd.Timestamp("2024-09-28 00:00:00") + +# One color per init date (Sep 24-30) +DATE_COLORS = { + "2024-09-24": "#1f77b4", + "2024-09-25": "#ff7f0e", + "2024-09-26": "#2ca02c", + "2024-09-27": "#d62728", + "2024-09-28": "#9467bd", + "2024-09-29": "#8c564b", + "2024-09-30": "#e377c2", +} + + +def load_parquet(path): + df = pd.read_parquet(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + df["ens_mean"] = df[member_cols].mean(axis=1) + df["ens_min"] = df[member_cols].min(axis=1) + df["ens_max"] = df[member_cols].max(axis=1) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + return df + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() + or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--route-dir", default=DEFAULT_ROUTE_DIR) + parser.add_argument("--out-dir", default=None) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--da-name", default="routed_leadtime_da_full.parquet") + parser.add_argument("--ol-name", default="routed_leadtime_openloop_full.parquet") + args = parser.parse_args() + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + print("Loading parquets...") + da = load_parquet(os.path.join(args.route_dir, args.da_name)) + ol = load_parquet(os.path.join(args.route_dir, args.ol_name)) + obs = load_usgs(args.usgs_csv) + + # Keep only valid_times inside the plot window + da = da[(da["valid_time"] >= PLOT_START) & (da["valid_time"] <= PLOT_END)] + ol = ol[(ol["valid_time"] >= PLOT_START) & (ol["valid_time"] <= PLOT_END)] + + issue_times = sorted(da["issue_time"].unique()) + print(f" {len(issue_times)} initialization times contribute to this window") + + fig, ax = plt.subplots(figsize=(18, 6)) + + # Helene peak shading + ax.axvspan(HELENE_PEAK_START, HELENE_PEAK_END, + color="salmon", alpha=0.12, zorder=0) + ax.text(HELENE_PEAK_START + pd.Timedelta(hours=6), 1.0, "Helene peak", + transform=ax.get_xaxis_transform(), + fontsize=9, color="firebrick", ha="left", va="top") + + # Plot each initialization time's trajectory + for t0 in issue_times: + date_str = str(pd.Timestamp(t0).date()) + color = DATE_COLORS.get(date_str, "gray") + + da_t = da[da["issue_time"] == t0].sort_values("valid_time") + ol_t = ol[ol["issue_time"] == t0].sort_values("valid_time") + + if da_t.empty: + continue + + # DA: shaded band (min-max across 20 members) + ensemble mean + ax.fill_between(da_t["valid_time"], da_t["ens_min"], da_t["ens_max"], + color=color, alpha=0.06, zorder=2) + ax.plot(da_t["valid_time"], da_t["ens_mean"], + color=color, lw=0.8, alpha=0.55, zorder=3) + + # OL: ensemble mean only, dashed + if not ol_t.empty: + ax.plot(ol_t["valid_time"], ol_t["ens_mean"], + color=color, lw=0.6, alpha=0.30, linestyle="--", zorder=2) + + # USGS obs + obs_w = obs.loc[PLOT_START:PLOT_END] + ax.plot(obs_w.index, obs_w.values, + color="black", lw=1.8, zorder=6, label="USGS obs") + + # Legend: one patch per init date + obs + DA/OL style + patches = [mpatches.Patch(color=c, label=f"Init {d}") + for d, c in DATE_COLORS.items()] + patches.append(plt.Line2D([0], [0], color="black", lw=1.8, label="USGS obs")) + patches.append(plt.Line2D([0], [0], color="gray", lw=1.2, alpha=0.7, + label="DA — ens mean (solid) + spread (shaded)")) + patches.append(plt.Line2D([0], [0], color="gray", lw=1.0, linestyle="--", + alpha=0.5, label="OL — ens mean (dashed)")) + ax.legend(handles=patches, fontsize=8, loc="upper left", + frameon=True, framealpha=0.9, ncol=2) + + ax.set_xlim(PLOT_START, PLOT_END) + ax.set_xlabel("Date", fontsize=11) + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_title( + "All 18-hour forecast trajectories ending in this window — USGS 03463300\n" + "Sep 24 18 UTC → Sep 30 06 UTC | DA: shaded band + mean | OL: mean dashed", + fontsize=11, + ) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=1)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.xaxis.set_minor_locator(mdates.HourLocator(interval=6)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.2, lw=0.4) + plt.tight_layout() + + out_path = os.path.join(out_dir, "forecast_spaghetti_helene.png") + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4c_timeseries/plot_helene_issue_time_hydrograph.py b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4c_timeseries/plot_helene_issue_time_hydrograph.py new file mode 100644 index 00000000..c6bc5dd8 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4c_timeseries/plot_helene_issue_time_hydrograph.py @@ -0,0 +1,212 @@ +""" +Per-issue-time forecast hydrograph diagnostic. + +For two contrasting issue times — one storm-peak (Helene) and one low-flow — +show what the DA and open-loop ensembles actually produce over the 18-hour +forecast window, alongside the kriging observation. Designed to answer in one +picture: is DA over-shooting, collapsing, or oscillating compared to open-loop? + +Each panel shows: + - All 20 DA members as faint purple lines + median (thick purple) + - All 20 open-loop members as faint gray lines + median (thick gray, dashed) + - Kriging obs as a thick black line (context: t0-6 through t0+18) + - Vertical line at t0 + annotation with obs(t0) and lead-1 ensemble means + +Inputs (from run_lead_time_forecast_sweep.py): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv +Obs (from production EnKF): + //_test_results.csv + +Output: + //_issue_time_hydrograph_helene_vs_lowflow.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +DEFAULT_LEADTIME_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast" +DEFAULT_DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" + +# Module-level globals set by main() +CAT = "cat-1016300" +LEADTIME_DIR = DEFAULT_LEADTIME_DIR +DA_DIR = DEFAULT_DA_DIR + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" +OBS_COLOR = "black" + +# Two issue times — pick a Helene-peak one and a typical low-flow one +HELENE_T0 = pd.Timestamp("2024-09-26 12:00:00") +LOWFLOW_T0 = pd.Timestamp("2024-03-15 12:00:00") + +CONTEXT_HOURS_BEFORE = 6 # hours of obs context shown before t0 +LEAD_HOURS_AFTER = 18 # forecast horizon + +OUT_PNG = os.path.join( + LEADTIME_DIR, CAT, + f"{CAT}_issue_time_hydrograph_helene_vs_lowflow.png", +) + + +def load_forecasts(path): + df = pd.read_csv(path, parse_dates=['issue_time', 'valid_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + return df, member_cols + + +OBS_DIR = None # set by main() when --obs-dir is supplied + + +def load_obs(): + if OBS_DIR is not None: + obs_path = os.path.join(OBS_DIR, f"{CAT}.csv") + df = pd.read_csv(obs_path) + time_col = 'datetime' if 'datetime' in df.columns else 'date' + df[time_col] = pd.to_datetime(df[time_col]) + col = 'qkrig' if 'qkrig' in df.columns else 'obs_mm_h' + return df.set_index(time_col)[col].rename('obs_mm_h') + obs_path = os.path.join(DA_DIR, CAT, f"{CAT}_test_results.csv") + df = pd.read_csv(obs_path, parse_dates=['date']) + return df.set_index('date')['obs_mm_h'] + + +def slice_forecast(df, member_cols, t0): + """Return (valid_times, member_array shape (lead, N)) for a single issue time.""" + sub = df[df['issue_time'] == t0].sort_values('lead_hour') + if len(sub) == 0: + raise ValueError(f"No forecast rows for issue_time={t0}") + valid_times = pd.to_datetime(sub['valid_time'].values) + member_arr = sub[member_cols].to_numpy(dtype=float) + return valid_times, member_arr + + +def plot_panel(ax, t0, df_da, df_ol, member_cols, obs_series, title_extra=""): + # Forecast trajectories + vt_da, q_da = slice_forecast(df_da, member_cols, t0) + vt_ol, q_ol = slice_forecast(df_ol, member_cols, t0) + + # Obs context window + ctx_start = t0 - pd.Timedelta(hours=CONTEXT_HOURS_BEFORE) + ctx_end = t0 + pd.Timedelta(hours=LEAD_HOURS_AFTER) + obs_window = obs_series.loc[ctx_start:ctx_end] + + # Open-loop members (drawn first so DA paints on top) + for j in range(q_ol.shape[1]): + ax.plot(vt_ol, q_ol[:, j], color=OL_COLOR, lw=0.6, alpha=0.45, zorder=2) + # DA members + for j in range(q_da.shape[1]): + ax.plot(vt_da, q_da[:, j], color=DA_COLOR, lw=0.6, alpha=0.45, zorder=3) + # Medians + ax.plot(vt_ol, np.nanmedian(q_ol, axis=1), color=OL_COLOR, lw=2.6, + linestyle="--", zorder=4, label="Open-loop median") + ax.plot(vt_da, np.nanmedian(q_da, axis=1), color=DA_COLOR, lw=2.6, + zorder=5, label="DA median") + # Obs + ax.plot(obs_window.index, obs_window.values, color=OBS_COLOR, lw=1.8, + zorder=6, label="Qkrig (obs)") + + # t0 marker + annotation + ax.axvline(t0, color="black", lw=1.0, linestyle=":", alpha=0.6, zorder=1) + obs_t0 = float(obs_series.get(t0, np.nan)) + obs_lead1 = float(obs_series.get(t0 + pd.Timedelta(hours=1), np.nan)) + da_lead1_mean = float(np.nanmean(q_da[0, :])) + ol_lead1_mean = float(np.nanmean(q_ol[0, :])) + + info_text = ( + f"t0 = {t0.strftime('%Y-%m-%d %H:%M')}\n" + f"obs(t0) = {obs_t0:.3f} mm/h\n" + f"obs(t0+1) = {obs_lead1:.3f} mm/h\n" + f"DA mean(t0+1) = {da_lead1_mean:.3f} mm/h\n" + f"OL mean(t0+1) = {ol_lead1_mean:.3f} mm/h" + ) + ax.text(0.02, 0.97, info_text, + transform=ax.transAxes, + fontsize=9, va='top', ha='left', family='monospace', + bbox=dict(boxstyle='round,pad=0.4', facecolor='white', + alpha=0.85, edgecolor='gray')) + + ax.set_title(title_extra, fontsize=12) + ax.set_ylabel("Discharge (mm/h)", fontsize=10) + ax.set_xlabel("Date", fontsize=10) + ax.xaxis.set_major_locator(mdates.HourLocator(interval=3)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %H:%M")) + ax.tick_params(axis='x', rotation=30, labelsize=8) + ax.grid(True, alpha=0.2) + ax.legend(loc='upper right', fontsize=9, frameon=True, framealpha=0.9) + + +def snap_to_nearest(t0, available_times): + """Return the available issue_time closest to t0.""" + times = pd.DatetimeIndex(available_times) + idx = np.abs((times - t0).total_seconds()).argmin() + snapped = times[idx] + if snapped != t0: + print(f" Snapped {t0} -> {snapped} (nearest available issue time)") + return snapped + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', default="cat-1016300") + parser.add_argument('--leadtime-dir', default=DEFAULT_LEADTIME_DIR, + help='Dir holding /_lead_time_forecasts_{da,openloop}.csv') + parser.add_argument('--da-dir', default=DEFAULT_DA_DIR, + help='Dir holding /_test_results.csv (for obs)') + parser.add_argument('--obs-dir', default=None, + help='Kriging obs dir holding .csv with qkrig column; ' + 'overrides --da-dir for obs loading') + parser.add_argument('--helene-t0', default=None, + help='Issue time for the Helene panel (default: 2024-09-26 12:00:00, ' + 'snapped to nearest available)') + parser.add_argument('--lowflow-t0', default=None, + help='Issue time for the low-flow panel (default: 2024-03-15 00:00:00, ' + 'snapped to nearest available)') + args = parser.parse_args() + + global CAT, LEADTIME_DIR, DA_DIR, OBS_DIR, OUT_PNG + CAT = args.cat_id + LEADTIME_DIR = args.leadtime_dir + DA_DIR = args.da_dir + OBS_DIR = args.obs_dir + OUT_PNG = os.path.join(LEADTIME_DIR, CAT, + f"{CAT}_issue_time_hydrograph_helene_vs_lowflow.png") + + da_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_da.csv") + ol_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_openloop.csv") + df_da, m_cols = load_forecasts(da_path) + df_ol, _ = load_forecasts(ol_path) + obs_series = load_obs() + + available = df_da['issue_time'].unique() + helene_t0 = snap_to_nearest( + pd.Timestamp(args.helene_t0) if args.helene_t0 else HELENE_T0, available) + lowflow_t0 = snap_to_nearest( + pd.Timestamp(args.lowflow_t0) if args.lowflow_t0 else pd.Timestamp("2024-03-15 00:00:00"), + available) + + fig, (ax_h, ax_l) = plt.subplots(1, 2, figsize=(20, 7)) + + plot_panel(ax_h, helene_t0, df_da, df_ol, m_cols, obs_series, + title_extra=f"Helene peak issue time — {helene_t0.strftime('%Y-%m-%d %H:%M')}") + plot_panel(ax_l, lowflow_t0, df_da, df_ol, m_cols, obs_series, + title_extra=f"Low-flow issue time — {lowflow_t0.strftime('%Y-%m-%d %H:%M')}") + + fig.suptitle( + f"Per-issue-time forecast hydrograph — {CAT}\n" + "Faint lines = 20 ensemble members. Thick lines = ensemble medians. " + "Vertical dotted line = t0 (DA stops here for both scenarios).", + fontsize=12, y=1.00, + ) + plt.tight_layout() + plt.savefig(OUT_PNG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4c_timeseries/plot_lead_time_reconstructed_timeseries.py b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4c_timeseries/plot_lead_time_reconstructed_timeseries.py new file mode 100644 index 00000000..a6d03b03 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4c_timeseries/plot_lead_time_reconstructed_timeseries.py @@ -0,0 +1,263 @@ +""" +Reconstructed time series at gauge 03463300, from our routed +lead-time forecasts. Mirrors the methodology of plot_timeseries_hourly_ensemble.py +(qSpatialAR) but feeds from the EnKF + T-route pipeline instead of the CNN. + +For each target hour t, pool ALL forecasts that land on t (across many +issue_times × many lead_hours × 20 members). Compute median + 5th/95th +percentile envelope. Plot DA and open-loop reconstructions overlaid on USGS +obs, with the Helene window shaded. + +Input parquets (from run_route_troute_leadtime.py): + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + Wide format: issue_time, lead_hour, member_00..member_19 (q in m³/s) + +Obs: + /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +Output: + /lead_time_reconstructed_timeseries.png + Default plot window: 2024-09-10 → 2024-10-10. + +Each scenario's curve is annotated with NSE vs USGS obs over the plot window. +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" +OBS_COLOR = "black" + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast_hardcoded_r_routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +# Watershed area at gauge 03463300 (sum of 21 catchments per GPKG hydrofabric) +WATERSHED_AREA_KM2 = 113.18 +# mm/h depth → m³/s: × 113.18 km² × 1000 / 3600 = 31.439 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +PLOT_START_DEFAULT = pd.Timestamp("2024-09-10 00:00:00") +PLOT_END_DEFAULT = pd.Timestamp("2024-10-10 23:00:00") + +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-30 23:00:00") + + +def load_parquet_long(path): + df = pd.read_parquet(path) + df['issue_time'] = pd.to_datetime(df['issue_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + if not member_cols: + raise ValueError(f"No member columns in {path}") + long = df.melt( + id_vars=['issue_time', 'lead_hour'], + value_vars=member_cols, + var_name='member', + value_name='q_m3s', + ) + long['valid_time'] = (long['issue_time'] + + pd.to_timedelta(long['lead_hour'], unit='h')) + return long + + +def reconstruct(long_df): + """For each valid_time, pool all (issue_time × lead × member) forecasts and + return median, p05, p95, count.""" + grouped = long_df.groupby('valid_time')['q_m3s'].agg( + median='median', + p05=lambda s: s.quantile(0.05), + p95=lambda s: s.quantile(0.95), + count='count', + ).reset_index() + return grouped + + +def load_usgs(usgs_csv): + """Load USGS gauge obs and ensure m³/s. + + The CSV at gauge 03463300 has the column `QObs(mm/h)` — catchment-averaged + depth, not gauge discharge in m³/s. Auto-convert when column name contains + 'mm': m³/s = mm/h × WATERSHED_AREA_KM2 × 1000 / 3600. + """ + df = pd.read_csv(usgs_csv) + date_col = next((c for c in df.columns + if c.lower() in ('datetime', 'date', 'time', 'timestamp')), None) + q_col = next((c for c in df.columns + if 'q' in c.lower() or 'flow' in c.lower() + or 'discharge' in c.lower()), None) + if not date_col or not q_col: + raise ValueError(f"Couldn't id date/Q columns in {usgs_csv}: {df.columns.tolist()}") + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if 'mm' in q_col.lower(): + series = series * MM_H_TO_M3_S + print(f" Converted obs from mm/h (column '{q_col}') to m³/s " + f"using area = {WATERSHED_AREA_KM2} km² (× {MM_H_TO_M3_S:.4f})") + else: + print(f" Loaded obs as-is from column '{q_col}' (assumed m³/s)") + return series + + +def nse(obs, sim): + m = np.isfinite(obs) & np.isfinite(sim) + if m.sum() < 2: + return np.nan + o, s = obs[m], sim[m] + denom = ((o - o.mean()) ** 2).sum() + if denom < 1e-10: + return np.nan + return 1.0 - float(((o - s) ** 2).sum() / denom) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--route-dir', default=DEFAULT_ROUTE_DIR) + parser.add_argument('--out-dir', default=None, + help='Default: --route-dir') + parser.add_argument('--usgs-csv', default=DEFAULT_USGS_CSV) + parser.add_argument('--plot-start', default=str(PLOT_START_DEFAULT)) + parser.add_argument('--plot-end', default=str(PLOT_END_DEFAULT)) + parser.add_argument('--da-name', default='routed_leadtime_da_full.parquet') + parser.add_argument('--ol-name', default='routed_leadtime_openloop_full.parquet') + args = parser.parse_args() + + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + plot_start = pd.Timestamp(args.plot_start) + plot_end = pd.Timestamp(args.plot_end) + + print(f"Loading DA parquet: {args.da_name}") + da_long = load_parquet_long(os.path.join(args.route_dir, args.da_name)) + print(f"Loading OL parquet: {args.ol_name}") + ol_long = load_parquet_long(os.path.join(args.route_dir, args.ol_name)) + print(f"Loading USGS obs: {args.usgs_csv}") + obs = load_usgs(args.usgs_csv) + + print("Reconstructing DA time series (overlapping-leads pool)...") + da_rec = reconstruct(da_long) + print(f" {len(da_rec)} unique target hours") + print("Reconstructing OL time series...") + ol_rec = reconstruct(ol_long) + + # Clip to plot window and join obs + da_rec = da_rec[(da_rec['valid_time'] >= plot_start) & (da_rec['valid_time'] <= plot_end)].copy() + ol_rec = ol_rec[(ol_rec['valid_time'] >= plot_start) & (ol_rec['valid_time'] <= plot_end)].copy() + da_rec['obs'] = da_rec['valid_time'].map(obs) + ol_rec['obs'] = ol_rec['valid_time'].map(obs) + + nse_da = nse(da_rec['obs'].values, da_rec['median'].values) + nse_ol = nse(ol_rec['obs'].values, ol_rec['median'].values) + print(f"NSE (plot window) — DA: {nse_da:.3f} | OL: {nse_ol:.3f}") + + # ---- Plot ---- + fig, ax = plt.subplots(figsize=(16, 5.5)) + fig.patch.set_facecolor("white") + + # Helene shaded + ax.axvspan(HELENE_START, HELENE_END, + color="firebrick", alpha=0.08, zorder=0) + ax.text(HELENE_START + pd.Timedelta(hours=12), 1.0, "Helene", + transform=ax.get_xaxis_transform(), + fontsize=9, color="firebrick", ha="left", va="top", + fontweight="bold") + + # Open-loop band + median + ax.fill_between(ol_rec['valid_time'], ol_rec['p05'], ol_rec['p95'], + color=OL_COLOR, alpha=0.18, zorder=2, + label="Open-loop [5th–95th pctile envelope]") + ax.plot(ol_rec['valid_time'], ol_rec['median'], + color=OL_COLOR, lw=1.2, linestyle='--', zorder=3, + label=f"Open-loop median (NSE={nse_ol:.3f})") + + # DA band + median + ax.fill_between(da_rec['valid_time'], da_rec['p05'], da_rec['p95'], + color=DA_COLOR, alpha=0.22, zorder=4, + label="DA [5th–95th pctile envelope]") + ax.plot(da_rec['valid_time'], da_rec['median'], + color=DA_COLOR, lw=1.4, zorder=5, + label=f"DA median (NSE={nse_da:.3f})") + + # Observed as dots + obs_window = obs.loc[plot_start:plot_end] + finite = np.isfinite(obs_window.values) + ax.scatter(obs_window.index[finite], obs_window.values[finite], + s=6, color=OBS_COLOR, linewidths=0, zorder=6, + label="USGS obs") + + ax.set_xlabel("Date", fontsize=11) + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_title( + f"Reconstructed forecast time series at USGS 03463300\n" + f"Overlapping-leads pool from EnKF forecast ensemble " + f"({plot_start.date()} – {plot_end.date()})", + fontsize=12, + ) + ax.set_xlim(plot_start, plot_end) + ax.xaxis.set_major_locator(mdates.WeekdayLocator(interval=1)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.25, lw=0.4) + ax.legend(loc='upper left', fontsize=9, frameon=True, framealpha=0.92) + + plt.tight_layout() + out_path = os.path.join(out_dir, "lead_time_reconstructed_timeseries.png") + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + # Also produce a Helene-zoom version + helene_zoom_start = pd.Timestamp("2024-09-24 00:00:00") + helene_zoom_end = pd.Timestamp("2024-09-29 23:00:00") + da_z = da_rec[(da_rec['valid_time'] >= helene_zoom_start) + & (da_rec['valid_time'] <= helene_zoom_end)] + ol_z = ol_rec[(ol_rec['valid_time'] >= helene_zoom_start) + & (ol_rec['valid_time'] <= helene_zoom_end)] + nse_da_z = nse(da_z['obs'].values, da_z['median'].values) + nse_ol_z = nse(ol_z['obs'].values, ol_z['median'].values) + print(f"NSE (Helene zoom) — DA: {nse_da_z:.3f} | OL: {nse_ol_z:.3f}") + + fig, ax = plt.subplots(figsize=(14, 5.5)) + fig.patch.set_facecolor("white") + ax.fill_between(ol_z['valid_time'], ol_z['p05'], ol_z['p95'], + color=OL_COLOR, alpha=0.18, zorder=2, + label="Open-loop [5th–95th]") + ax.plot(ol_z['valid_time'], ol_z['median'], color=OL_COLOR, + lw=1.4, linestyle='--', zorder=3, + label=f"Open-loop median (NSE={nse_ol_z:.3f})") + ax.fill_between(da_z['valid_time'], da_z['p05'], da_z['p95'], + color=DA_COLOR, alpha=0.22, zorder=4, + label="DA [5th–95th]") + ax.plot(da_z['valid_time'], da_z['median'], color=DA_COLOR, + lw=1.6, zorder=5, + label=f"DA median (NSE={nse_da_z:.3f})") + obs_z = obs.loc[helene_zoom_start:helene_zoom_end] + finite = np.isfinite(obs_z.values) + ax.scatter(obs_z.index[finite], obs_z.values[finite], + s=12, color=OBS_COLOR, linewidths=0, zorder=6, + label="USGS obs") + ax.set_xlabel("Date", fontsize=11) + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_title( + "Reconstructed forecast time series at USGS 03463300 — Helene zoom\n" + "Sep 24 – 29, 2024 (overlapping-leads pool from EnKF forecast ensemble)", + fontsize=12, + ) + ax.xaxis.set_major_locator(mdates.HourLocator(interval=12)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %H:%M")) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.25, lw=0.4) + ax.legend(loc='upper left', fontsize=9, frameon=True, framealpha=0.92) + plt.tight_layout() + out_zoom = os.path.join(out_dir, "lead_time_reconstructed_timeseries_helene.png") + plt.savefig(out_zoom, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_zoom}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4c_timeseries/run_4c_f2.sh b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4c_timeseries/run_4c_f2.sh new file mode 100644 index 00000000..2d3da5cb --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/4c_timeseries/run_4c_f2.sh @@ -0,0 +1,51 @@ +#!/usr/bin/env bash +# F2 Fixed R=0.07 (1 gauge holdout) — 4c reconstructed timeseries plots. +# +# Reads routed_leadtime_da_full.parquet + routed_leadtime_openloop_full.parquet +# from the F2 leadtime route dir and writes output PNGs there. +# Must run AFTER route_lead_time_forecasts.py. +# +# Usage: +# bash run_4c_f2.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +LEADTIME_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder2_fixed_r007 +ROUTE_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder2_fixed_r007_leadtime_routed +OBS_DIR=/home/svyas/catchment_ts_no_03463300_gapfilled +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +echo "[F2-4c] Reconstructed timeseries from: $ROUTE_DIR" +$TROUTE "$SCRIPT_DIR/plot_lead_time_reconstructed_timeseries.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" + +echo "[F2-4c] Forecast spaghetti..." +$TROUTE "$SCRIPT_DIR/plot_forecast_spaghetti.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" + +echo "[F2-4c] Helene vs low-flow issue-time hydrograph (per catchment)..." +for CAT in "${CATS[@]}"; do + DA_CSV="$LEADTIME_DIR/$CAT/${CAT}_lead_time_forecasts_da.csv" + [ -f "$DA_CSV" ] || { echo " [$CAT] no leadtime CSV — skipping"; continue; } + echo " [$CAT] helene hydrograph..." + $TROUTE "$SCRIPT_DIR/plot_helene_issue_time_hydrograph.py" \ + --cat-id "$CAT" \ + --leadtime-dir "$LEADTIME_DIR" \ + --obs-dir "$OBS_DIR" +done + +echo "[F2-4c] Done. Outputs in: $ROUTE_DIR and $LEADTIME_DIR" +ls "$ROUTE_DIR"/*.png 2>/dev/null || echo " (no PNGs in route dir)" diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/replot_f2_arm_ymax.sh b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/replot_f2_arm_ymax.sh new file mode 100644 index 00000000..096f31f8 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/4_evaluation/replot_f2_arm_ymax.sh @@ -0,0 +1,51 @@ +#!/usr/bin/env bash +# Regenerate F2 2b hydro-state arm plots with y-axis capped at 2200 m³/s +# and the routed open-loop baseline overlaid. +# +# Usage: +# bash replot_f2_arm_ymax.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/plot_da_perturbation_arms.py +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv +OL_CSV=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/openloop/routed_Q_test.csv +F2_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder2_fixed_r007 + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +if [ ! -f "$OL_CSV" ]; then + echo "ERROR: routed open-loop not found: $OL_CSV" + exit 1 +fi + +SKIP=0; DONE=0; FAIL=0 + +for CAT in "${CATS[@]}"; do + HYDRO_CSV="$F2_DIR/$CAT/${CAT}_da_hydro_arm.csv" + + if [ ! -f "$HYDRO_CSV" ]; then + echo " [$CAT] No hydro arm CSV — skipping" + SKIP=$((SKIP + 1)); continue + fi + + echo " [$CAT] plotting (ymax=2200)..." + $TROUTE "$SCRIPT" \ + --arm-dir "$F2_DIR" \ + --cat-id "$CAT" \ + --usgs-csv "$USGS_CSV" \ + --ol-ts-csv "$OL_CSV" \ + --ymax 2200 \ + && DONE=$((DONE + 1)) \ + || { echo " FAILED: $CAT"; FAIL=$((FAIL + 1)); } +done + +echo "" +echo "=== F2 replot done: done=$DONE skipped=$SKIP failed=$FAIL ===" diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016279_best_params.json b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016279_best_params.json new file mode 100644 index 00000000..2b05b093 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016279_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016279", + "best_kge": 0.8206116996389099, + "best_parameters": { + "bb": 2.623111719150318, + "smcmax": 0.2, + "satdk": 9.781515618242848e-05, + "slop": 0.49040228658878326, + "max_gw_storage": 0.22502382651758357, + "expon": 3.9507553183626896, + "Cgw": 0.00028265811089655624, + "K_lf": 0.33708866739487964, + "K_nash": 0.7457712895070667, + "scheme": 0.024063276949108925 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016280_best_params.json b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016280_best_params.json new file mode 100644 index 00000000..577b7703 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016280_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016280", + "best_kge": 0.7845371451180185, + "best_parameters": { + "bb": 3.6853461244917503, + "smcmax": 0.2, + "satdk": 5.2992817070504554e-05, + "slop": 0.19547785108903476, + "max_gw_storage": 0.2491156800177733, + "expon": 2.254619375138279, + "Cgw": 0.000560475352848284, + "K_lf": 1.0, + "K_nash": 0.3230987974065112, + "scheme": 0.34480885239395503 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016281_best_params.json b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016281_best_params.json new file mode 100644 index 00000000..17346f8b --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016281_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016281", + "best_kge": 0.7653169268425993, + "best_parameters": { + "bb": 3.6065627281526345, + "smcmax": 0.2, + "satdk": 6.865558286144331e-05, + "slop": 0.16111639333094452, + "max_gw_storage": 0.17693006455477892, + "expon": 2.0040076166999006, + "Cgw": 0.00038006296817321816, + "K_lf": 0.7105167430808037, + "K_nash": 9.127e-08, + "scheme": 0.14298989888102898 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016282_best_params.json b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016282_best_params.json new file mode 100644 index 00000000..0e53037b --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016282_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016282", + "best_kge": 0.7868853451455085, + "best_parameters": { + "bb": 2.627220084315586, + "smcmax": 0.2, + "satdk": 7.889407420323658e-05, + "slop": 0.4136394925802219, + "max_gw_storage": 0.23896390723786864, + "expon": 7.187637073442723, + "Cgw": 9.498445725943469e-06, + "K_lf": 0.44532012273911126, + "K_nash": 0.30386284986679846, + "scheme": 0.4286174650536479 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016283_best_params.json b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016283_best_params.json new file mode 100644 index 00000000..57a06e77 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016283_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016283", + "best_kge": 0.7456494190014495, + "best_parameters": { + "bb": 2.500182877645966, + "smcmax": 0.2, + "satdk": 7.708052707732115e-05, + "slop": 0.43914051816633154, + "max_gw_storage": 0.17844520971245803, + "expon": 1.636624858961519, + "Cgw": 0.0007186031814004434, + "K_lf": 0.7604006750247336, + "K_nash": 0.935851019129656, + "scheme": 0.4295615557232241 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016300_best_params.json b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016300_best_params.json new file mode 100644 index 00000000..99b56bcf --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016300_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016300", + "best_kge": 0.7852863461640692, + "best_parameters": { + "bb": 2.1364476151747556, + "smcmax": 0.22502190033753489, + "satdk": 8.571865512107468e-05, + "slop": 3.06e-05, + "max_gw_storage": 0.25, + "expon": 5.11973472709136, + "Cgw": 0.0001071612486323186, + "K_lf": 0.25307693130285724, + "K_nash": 0.3553147233842257, + "scheme": 0.37275121762887575 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016301_best_params.json b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016301_best_params.json new file mode 100644 index 00000000..a4fc0214 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016301_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016301", + "best_kge": 0.7852546470254989, + "best_parameters": { + "bb": 2.4254415208275426, + "smcmax": 0.2, + "satdk": 7.587190338000397e-05, + "slop": 0.7895909975757629, + "max_gw_storage": 0.2175504021153558, + "expon": 2.748699844202327, + "Cgw": 0.0004359795527131938, + "K_lf": 0.7766373897901322, + "K_nash": 0.3599347814597565, + "scheme": 0.024075462516893846 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016302_best_params.json b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016302_best_params.json new file mode 100644 index 00000000..e9d65dff --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016302_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016302", + "best_kge": 0.7438367117368438, + "best_parameters": { + "bb": 2.0, + "smcmax": 0.2597182603107496, + "satdk": 0.00010927248622434935, + "slop": 0.74928833629841, + "max_gw_storage": 0.20590055515471217, + "expon": 2.6302920260523215, + "Cgw": 0.00036611733203253385, + "K_lf": 0.8653682881680377, + "K_nash": 0.023172715332685445, + "scheme": 1.621e-05 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016303_best_params.json b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016303_best_params.json new file mode 100644 index 00000000..6516c488 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016303_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016303", + "best_kge": 0.7717901438015875, + "best_parameters": { + "bb": 2.0011493088041257, + "smcmax": 0.25267740195882527, + "satdk": 0.00010267632057588901, + "slop": 0.7436120787827551, + "max_gw_storage": 0.25, + "expon": 3.40584201489453, + "Cgw": 0.0002855314649977274, + "K_lf": 2.22e-06, + "K_nash": 0.03462430327931909, + "scheme": 0.4006007187079474 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016304_best_params.json b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016304_best_params.json new file mode 100644 index 00000000..aeef2972 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016304_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016304", + "best_kge": 0.785816115160746, + "best_parameters": { + "bb": 2.627220084315586, + "smcmax": 0.2, + "satdk": 6.623684502570866e-05, + "slop": 0.7710190313472335, + "max_gw_storage": 0.23151061983908472, + "expon": 2.852593661529338, + "Cgw": 0.0003041527852735811, + "K_lf": 0.5192505545426351, + "K_nash": 0.77484040937936, + "scheme": 0.17801104705809254 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016305_best_params.json b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016305_best_params.json new file mode 100644 index 00000000..6bae6748 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016305_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016305", + "best_kge": 0.7744236016300623, + "best_parameters": { + "bb": 2.9787357329850916, + "smcmax": 0.2, + "satdk": 2.01708703759236e-05, + "slop": 0.2671853140329312, + "max_gw_storage": 0.20533656076576476, + "expon": 1.0, + "Cgw": 0.0012738648070631134, + "K_lf": 7.201e-07, + "K_nash": 0.20780337905391905, + "scheme": 0.14434010116001775 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016306_best_params.json b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016306_best_params.json new file mode 100644 index 00000000..5cf772c1 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016306_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016306", + "best_kge": 0.7210357839982497, + "best_parameters": { + "bb": 2.0, + "smcmax": 0.27768197452250365, + "satdk": 0.00011181892180957246, + "slop": 0.17080432119242306, + "max_gw_storage": 0.18472414564626138, + "expon": 1.0299420125608787, + "Cgw": 0.001262070070347035, + "K_lf": 1.049e-05, + "K_nash": 0.10556306992792686, + "scheme": 0.23455263978683152 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016307_best_params.json b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016307_best_params.json new file mode 100644 index 00000000..92a410bf --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016307_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016307", + "best_kge": 0.801808946394892, + "best_parameters": { + "bb": 3.4434053536999176, + "smcmax": 0.2, + "satdk": 5.6465464765815605e-05, + "slop": 0.4263217361110973, + "max_gw_storage": 0.25, + "expon": 4.337261095363485, + "Cgw": 0.0001245938975937281, + "K_lf": 0.845893106799097, + "K_nash": 0.2386693544648092, + "scheme": 0.13456997008570973 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016308_best_params.json b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016308_best_params.json new file mode 100644 index 00000000..b644d485 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016308_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016308", + "best_kge": 0.8061757655923723, + "best_parameters": { + "bb": 3.6702131149274404, + "smcmax": 0.2, + "satdk": 6.003079869425282e-05, + "slop": 0.41011011063560954, + "max_gw_storage": 0.23346373674800242, + "expon": 3.9123522328874523, + "Cgw": 0.000183853837621649, + "K_lf": 0.04420155753170629, + "K_nash": 0.4655759687155024, + "scheme": 0.2840640813933577 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016309_best_params.json b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016309_best_params.json new file mode 100644 index 00000000..c03c2c55 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016309_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016309", + "best_kge": 0.7407364262156276, + "best_parameters": { + "bb": 3.833847789992432, + "smcmax": 0.2, + "satdk": 5.364111110266205e-05, + "slop": 0.23910717165219175, + "max_gw_storage": 0.23812574692579214, + "expon": 3.010993385187965, + "Cgw": 0.0002958484512966826, + "K_lf": 0.3302542952661117, + "K_nash": 0.2737476576326264, + "scheme": 0.4849311470054365 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016310_best_params.json b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016310_best_params.json new file mode 100644 index 00000000..a193e2ad --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016310_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016310", + "best_kge": 0.7414257687945072, + "best_parameters": { + "bb": 4.77839967410216, + "smcmax": 0.2, + "satdk": 0.0005122349525110497, + "slop": 0.2343084566046406, + "max_gw_storage": 0.1972951075076314, + "expon": 2.689926753507914, + "Cgw": 0.00025277943884201336, + "K_lf": 0.10451650547499411, + "K_nash": 0.4766491757683344, + "scheme": 0.14703543348764395 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016311_best_params.json b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016311_best_params.json new file mode 100644 index 00000000..88e348ad --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016311_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016311", + "best_kge": 0.7370239364407825, + "best_parameters": { + "bb": 3.7224279985565554, + "smcmax": 0.2, + "satdk": 5.36848284289468e-05, + "slop": 0.2747382803914872, + "max_gw_storage": 0.20236610051190926, + "expon": 1.2298400851415858, + "Cgw": 0.0008686565105330384, + "K_lf": 0.5491364936227011, + "K_nash": 0.5637455013613035, + "scheme": 0.3172192584604878 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016312_best_params.json b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016312_best_params.json new file mode 100644 index 00000000..12e8f931 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016312_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016312", + "best_kge": 0.7651323332103412, + "best_parameters": { + "bb": 3.4163907102604436, + "smcmax": 0.2, + "satdk": 5.8900848701685016e-05, + "slop": 0.952690690782515, + "max_gw_storage": 0.20339558487059237, + "expon": 2.0393414268931, + "Cgw": 0.0005321846942954315, + "K_lf": 0.005616844228055001, + "K_nash": 0.10556197423792017, + "scheme": 0.16715606346334788 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016313_best_params.json b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016313_best_params.json new file mode 100644 index 00000000..b08b0bc6 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016313_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016313", + "best_kge": 0.7541333608718466, + "best_parameters": { + "bb": 2.0, + "smcmax": 0.3236473885603026, + "satdk": 0.0006067307257057383, + "slop": 1.0, + "max_gw_storage": 0.23803639343037397, + "expon": 4.052333909196705, + "Cgw": 0.0002604427783083463, + "K_lf": 1.0, + "K_nash": 0.013987580178250053, + "scheme": 0.8686112513355525 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016314_best_params.json b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016314_best_params.json new file mode 100644 index 00000000..26e34c1b --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016314_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016314", + "best_kge": 0.8237475701476421, + "best_parameters": { + "bb": 2.781255453729351, + "smcmax": 0.20819967102864187, + "satdk": 7.660284170122033e-05, + "slop": 0.41683958183300995, + "max_gw_storage": 0.23944464809927757, + "expon": 5.8409631177306265, + "Cgw": 5.492004041138577e-05, + "K_lf": 0.846968216875224, + "K_nash": 0.46122458817838097, + "scheme": 0.3779302228964909 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016315_best_params.json b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016315_best_params.json new file mode 100644 index 00000000..e3dbfa20 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder2_fixed_r_007/best_params/cat-1016315_best_params.json @@ -0,0 +1,16 @@ +{ + "catchment_id": "cat-1016315", + "best_kge": 0.8232409477622364, + "best_parameters": { + "bb": 2.6967758837666995, + "smcmax": 0.2, + "satdk": 7.062588287975823e-05, + "slop": 5.173e-07, + "max_gw_storage": 0.24187360050863188, + "expon": 4.042512273338185, + "Cgw": 0.0003098824042018927, + "K_lf": 0.1631469310671928, + "K_nash": 0.07301338855959635, + "scheme": 0.431430924337884 + } +} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/batch_run_crossed_f4.sh b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/batch_run_crossed_f4.sh new file mode 100644 index 00000000..4236dfe4 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/batch_run_crossed_f4.sh @@ -0,0 +1,65 @@ +#!/usr/bin/env bash +# Batch crossed ensemble for all 21 catchments — F4 direct variance. +# R = krig_var directly; uses --direct-variance flag in run_crossed_ensemble.py. +# Outputs _crossed_ensemble.parquet per catchment. +# +# Usage: +# nohup bash batch_run_crossed_f4.sh > ~/logs/crossed_f4.log 2>&1 & + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PYTHON=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT="$SCRIPT_DIR/run_crossed_ensemble.py" + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +OBS_DIR=/mnt/disk2/1400_sites_helene/catchment_ts_no_03463300_dynamic_variance +CFE_DIR=/mnt/disk2/suma_helen_poster/cfe_py +CONFIG_FILE=/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder4_dynamic_variance_direct +FORCING_DIR1=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings +FORCING_DIR2=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings + +echo "F4 Direct variance crossed ensemble — out: $OUT_DIR" +echo "" + +SKIP=0; DONE=0; FAIL=0 + +for CAT in "${CATS[@]}"; do + echo "===============================" + echo "=== $CAT ===" + echo "===============================" + + CAT_OUT="$OUT_DIR/$CAT" + + if [ ! -f "$CAT_OUT/${CAT}_best_params.json" ]; then + echo " WARNING: No best_params for $CAT — skipping" + SKIP=$((SKIP + 1)); continue + fi + + if [ -f "$CAT_OUT/${CAT}_crossed_ensemble.parquet" ]; then + echo " Crossed ensemble already exists — skipping" + SKIP=$((SKIP + 1)); continue + fi + + "$PYTHON" "$SCRIPT" \ + --cat-id "$CAT" \ + --forcing-dir "$FORCING_DIR1" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --out-dir "$OUT_DIR" \ + --test-forcing-dir1 "$FORCING_DIR1" \ + --test-forcing-dir2 "$FORCING_DIR2" \ + --direct-variance \ + && DONE=$((DONE + 1)) \ + || { echo " FAILED: $CAT"; FAIL=$((FAIL + 1)); } +done + +echo "" +echo "=== F4 crossed done: done=$DONE skipped=$SKIP failed=$FAIL ===" diff --git a/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/batch_run_da_on_all_cats.sh b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/batch_run_da_on_all_cats.sh new file mode 100644 index 00000000..15e9abfb --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/batch_run_da_on_all_cats.sh @@ -0,0 +1,82 @@ +#!/usr/bin/env bash +# Batch DA for all 21 catchments — F4 direct variance. +# R = krig_var directly (sigma^2 from Qkrig, no Vrugt scaling). +# Stages best_params.json from calibration results before running. +# +# Usage (from server, after SCPing this folder): +# bash batch_run_da_on_all_cats.sh + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PYTHON=/home/svyas/miniconda3/envs/troute/bin/python3 +DA_SCRIPT="$SCRIPT_DIR/run_perturbation_da_on.py" + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +OBS_DIR=/mnt/disk2/1400_sites_helene/catchment_ts_no_03463300_dynamic_variance +BEST_PARAMS_SRC=/mnt/disk2/suma_helen_poster/catchment_results_1gauge_heldout +CFE_DIR=/mnt/disk2/suma_helen_poster/cfe_py +CONFIG_FILE=/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json +PARAM_BOUNDS=/mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder4_dynamic_variance_direct +FORCING_DIR1=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings +FORCING_DIR2=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings + +mkdir -p "$OUT_DIR" +echo "F4 Direct variance DA — obs: $OBS_DIR" +echo " out: $OUT_DIR" +echo "" + +SKIP=0; DONE=0; FAIL=0 + +for CAT in "${CATS[@]}"; do + echo "===============================" + echo "=== $CAT ===" + echo "===============================" + + CAT_OUT="$OUT_DIR/$CAT" + mkdir -p "$CAT_OUT" + + # Stage best_params from calibration results + if [ ! -f "$CAT_OUT/${CAT}_best_params.json" ]; then + SRC="$BEST_PARAMS_SRC/$CAT/${CAT}_best_params.json" + if [ -f "$SRC" ]; then + cp "$SRC" "$CAT_OUT/" + echo " Staged best_params" + else + echo " WARNING: No best_params for $CAT — skipping" + SKIP=$((SKIP + 1)); continue + fi + else + echo " best_params already staged" + fi + + # Skip if both arm CSVs already exist + if [ -f "$CAT_OUT/${CAT}_da_forcing_arm.csv" ] && \ + [ -f "$CAT_OUT/${CAT}_da_hydro_arm.csv" ]; then + echo " Arms already complete — skipping" + SKIP=$((SKIP + 1)); continue + fi + + "$PYTHON" "$DA_SCRIPT" \ + --cat-id "$CAT" \ + --forcing-dir "$FORCING_DIR1" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --param-bounds "$PARAM_BOUNDS" \ + --out-dir "$OUT_DIR" \ + --test-forcing-dir1 "$FORCING_DIR1" \ + --test-forcing-dir2 "$FORCING_DIR2" \ + --direct-variance \ + && DONE=$((DONE + 1)) \ + || { echo " FAILED: $CAT"; FAIL=$((FAIL + 1)); } +done + +echo "" +echo "=== F4 done: done=$DONE skipped=$SKIP failed=$FAIL ===" diff --git a/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/batch_run_leadtime_f4.sh b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/batch_run_leadtime_f4.sh new file mode 100644 index 00000000..fdff757d --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/batch_run_leadtime_f4.sh @@ -0,0 +1,74 @@ +#!/usr/bin/env bash +# Batch lead-time forecast sweep — F4 direct variance, 1 gauge holdout. +# Runs run_lead_time_forecast_sweep.py for all 21 catchments with --direct-variance. +# Skips if both output CSVs already exist. +# +# After this completes, run route_lead_time_forecasts.py to route through T-route. +# +# Usage: +# nohup bash batch_run_leadtime_f4.sh > ~/logs/leadtime_f4.log 2>&1 & + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PYTHON=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT="$SCRIPT_DIR/run_lead_time_forecast_sweep.py" + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +OBS_DIR=/mnt/disk2/1400_sites_helene/catchment_ts_no_03463300_dynamic_variance +CFE_DIR=/mnt/disk2/suma_helen_poster/cfe_py +CONFIG_FILE=/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json +PARAM_BOUNDS=/mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder4_dynamic_variance_direct +FORCING_DIR1=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings +FORCING_DIR2=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings + +echo "F4 lead-time sweep (direct variance) — out: $OUT_DIR" +echo "" + +SKIP=0; DONE=0; FAIL=0 + +for CAT in "${CATS[@]}"; do + echo "===============================" + echo "=== $CAT ===" + echo "===============================" + + CAT_OUT="$OUT_DIR/$CAT" + DA_CSV="$CAT_OUT/${CAT}_lead_time_forecasts_da.csv" + OL_CSV="$CAT_OUT/${CAT}_lead_time_forecasts_openloop.csv" + + if [ ! -f "$CAT_OUT/${CAT}_best_params.json" ]; then + echo " WARNING: No best_params for $CAT — skipping" + SKIP=$((SKIP + 1)); continue + fi + + if [ -f "$DA_CSV" ] && [ -f "$OL_CSV" ]; then + echo " Lead-time CSVs already exist — skipping" + SKIP=$((SKIP + 1)); continue + fi + + "$PYTHON" "$SCRIPT" \ + --cat-id "$CAT" \ + --forcing-dir "$FORCING_DIR1" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --param-bounds "$PARAM_BOUNDS" \ + --out-dir "$OUT_DIR" \ + --test-forcing-dir1 "$FORCING_DIR1" \ + --test-forcing-dir2 "$FORCING_DIR2" \ + --enkf-members 20 \ + --base-step-h 6 \ + --direct-variance \ + --prod-script "$SCRIPT_DIR/../../../1_distributed_cfe/calibrate_catchment_cfe_da_v2.py" \ + && DONE=$((DONE + 1)) \ + || { echo " FAILED: $CAT"; FAIL=$((FAIL + 1)); } +done + +echo "" +echo "=== F4 lead-time sweep done: done=$DONE skipped=$SKIP failed=$FAIL ===" diff --git a/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/batch_run_production_f4.sh b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/batch_run_production_f4.sh new file mode 100644 index 00000000..f36fd074 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/batch_run_production_f4.sh @@ -0,0 +1,71 @@ +#!/usr/bin/env bash +# Batch production per-member run — F4 Direct variance, 1 gauge holdout. +# Runs run_production_per_member.py for all 21 catchments with --direct-variance. +# Skips if output CSV already exists. +# +# Usage: +# nohup bash batch_run_production_f4.sh > ~/logs/production_f4.log 2>&1 & + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PYTHON=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT="$SCRIPT_DIR/run_production_per_member.py" +PROD_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/1_distributed_cfe/calibrate_catchment_cfe_da_v2.py + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +OBS_DIR=/mnt/disk2/1400_sites_helene/catchment_ts_no_03463300_dynamic_variance +CFE_DIR=/mnt/disk2/suma_helen_poster/cfe_py +CONFIG_FILE=/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json +PARAM_BOUNDS=/mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder4_dynamic_variance_direct +FORCING_DIR1=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings +FORCING_DIR2=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings + +echo "F4 production per-member (direct variance) — out: $OUT_DIR" +echo "" + +SKIP=0; DONE=0; FAIL=0 + +for CAT in "${CATS[@]}"; do + echo "===============================" + echo "=== $CAT ===" + echo "===============================" + + CAT_OUT="$OUT_DIR/$CAT" + OUT_FILE="$CAT_OUT/${CAT}_production_per_member.csv" + + if [ ! -f "$CAT_OUT/${CAT}_best_params.json" ]; then + echo " WARNING: No best_params for $CAT — skipping" + SKIP=$((SKIP + 1)); continue + fi + + if [ -f "$OUT_FILE" ]; then + echo " Production per-member CSV already exists — skipping" + SKIP=$((SKIP + 1)); continue + fi + + "$PYTHON" "$SCRIPT" \ + --cat-id "$CAT" \ + --forcing-dir "$FORCING_DIR1" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --param-bounds "$PARAM_BOUNDS" \ + --out-dir "$OUT_DIR" \ + --test-forcing-dir1 "$FORCING_DIR1" \ + --test-forcing-dir2 "$FORCING_DIR2" \ + --enkf-members 20 \ + --direct-variance \ + --prod-script "$PROD_SCRIPT" \ + && DONE=$((DONE + 1)) \ + || { echo " FAILED: $CAT"; FAIL=$((FAIL + 1)); } +done + +echo "" +echo "=== F4 production per-member done: done=$DONE skipped=$SKIP failed=$FAIL ===" diff --git a/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/batch_run_sensitivity_f4.sh b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/batch_run_sensitivity_f4.sh new file mode 100644 index 00000000..96dacb6d --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/batch_run_sensitivity_f4.sh @@ -0,0 +1,68 @@ +#!/usr/bin/env bash +# Batch perturbation sensitivity analysis — F4 Direct variance, 1 gauge holdout. +# Runs run_perturbation_sensitivity.py for all 21 catchments × 3 sources +# (init, forcing, process). DA is OFF; only perturbation source differs. +# Skips if output CSV already exists. +# +# Usage: +# nohup bash batch_run_sensitivity_f4.sh > ~/logs/sensitivity_f4.log 2>&1 & + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PYTHON=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT="$SCRIPT_DIR/run_perturbation_sensitivity.py" + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +SOURCES=(init forcing process) + +OBS_DIR=/mnt/disk2/1400_sites_helene/catchment_ts_no_03463300_dynamic_variance +CFE_DIR=/mnt/disk2/suma_helen_poster/cfe_py +CONFIG_FILE=/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json +PARAM_BOUNDS=/mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder4_dynamic_variance_direct +FORCING_DIR1=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings +FORCING_DIR2=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings + +echo "F4 sensitivity analysis — out: $OUT_DIR" +echo "" + +SKIP=0; DONE=0; FAIL=0 + +for CAT in "${CATS[@]}"; do + for SRC in "${SOURCES[@]}"; do + OUT_FILE="$OUT_DIR/$CAT/${CAT}_sensitivity_${SRC}.csv" + + if [ -f "$OUT_FILE" ]; then + echo " [$CAT/$SRC] already exists — skipping" + SKIP=$((SKIP + 1)) + continue + fi + + echo "===============================" + echo "=== $CAT source=$SRC ===" + echo "===============================" + + "$PYTHON" "$SCRIPT" \ + --cat-id "$CAT" \ + --source "$SRC" \ + --forcing-dir "$FORCING_DIR1" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --param-bounds "$PARAM_BOUNDS" \ + --out-dir "$OUT_DIR" \ + --test-forcing-dir1 "$FORCING_DIR1" \ + --test-forcing-dir2 "$FORCING_DIR2" \ + && DONE=$((DONE + 1)) \ + || { echo " FAILED: $CAT/$SRC"; FAIL=$((FAIL + 1)); } + done +done + +echo "" +echo "=== F4 sensitivity done: done=$DONE skipped=$SKIP failed=$FAIL ===" diff --git a/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/input_enkf_new.json b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/input_enkf_new.json new file mode 100644 index 00000000..482233bc --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/input_enkf_new.json @@ -0,0 +1 @@ +{"n": 1, "m": 1, "R": 0.07, "Q": 0, "smcmax" : 0.9394097311639178, "N": 1000, "P": 0.01, "D":2} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/new_EnKF.py b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/new_EnKF.py new file mode 100644 index 00000000..ff945c73 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/new_EnKF.py @@ -0,0 +1,53 @@ +# EnKF.py +import numpy as np + +class EnKF: + def __init__(self, n, m, R, Q, N, P, smcmax, D): + self.n = n # number of states + self.m = m # number of measurements + self.R = R # measurement noise + self.Q = Q # process noise + self.N = N # number of ensembles + self.P = P + self.state_estimates = np.zeros((self.n)) + self.covariance_matrices = np.zeros((self.n, self.n)) + self.current_step = 0 + self.smcmax = smcmax + self.D = D + self.storage_max_m = self.smcmax * self.D + self.storage_init = self.storage_max_m * 0.667 + self.ensembles = np.full((self.n, self.N), self.storage_init) + + def predict(self, F_results): + error_factor_sim = 0.005 + for i in range(self.N): + perturbation_factor_sim = np.random.standard_normal() + self.ensembles[:, i] = F_results[i] + self.ensembles[:, i] += (perturbation_factor_sim * F_results[i] * error_factor_sim) + return self.ensembles + + def update(self, yi, H_results): + if np.any(np.isnan(yi)): + return np.mean(self.ensembles, axis=1) # return current state estimate without updating + + y_ensembles = np.zeros((self.m, self.N)) + error_factor_sm = 0.03 + for i in range(self.N): + perturbation_factor_sm = np.random.standard_normal() + y_ensembles[:, i] = H_results[i] + y_ensembles[:, i] += (perturbation_factor_sm * H_results[i] * error_factor_sm) + + ensemble_mean = np.mean(self.ensembles, axis=1)[:, np.newaxis] + y_ensembles_mean = np.mean(y_ensembles, axis=1)[:, np.newaxis] + Pxy = (self.ensembles - ensemble_mean).dot((y_ensembles - y_ensembles_mean).T) / (self.N - 1) + Pyy = np.cov(y_ensembles, bias=True) + K = Pxy.dot(np.linalg.pinv(Pyy + self.R)) + self.ensembles += K.dot(yi - y_ensembles_mean) + return np.mean(self.ensembles, axis=1) + + def get_state_estimate(self): + return np.mean(self.ensembles, axis=1) + + def get_covariance(self): + ensemble_mean = np.mean(self.ensembles, axis=1)[:, np.newaxis] + return (self.ensembles - ensemble_mean).dot((self.ensembles - ensemble_mean).T) / (self.N - 1) diff --git a/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/run_crossed_ensemble.py b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/run_crossed_ensemble.py new file mode 100644 index 00000000..bc1cec20 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/run_crossed_ensemble.py @@ -0,0 +1,463 @@ +""" +run_crossed_ensemble.py -- 4b: 600-member crossed ensemble for Helene. + +Crosses the two DA-on perturbation arms: + Forcing arm : 30 met draws (from run_perturbation_da_on.py Phase 2A) + Hydro arm : 20 state draws (from run_perturbation_da_on.py Phase 2B) + +For each Helene issue time t0: + Run 600 members = 30 forcing draws x 20 hydro-state draws. + Each member (i, j) gets: + - Initial state = perturbed DA analysis state (draw j of hydro arm) + - Met forcing = perturbed precip/PET sequence (draw i of forcing arm) + +This gives the full uncertainty envelope combining both sources — +the operational probabilistic forecast picture. + +Output (per catchment): + //_crossed_ensemble.parquet + Columns: issue_time, lead_hour, member_0000..member_0599 + +Usage: + python3 run_crossed_ensemble.py \\ + --cat-id cat-1016300 \\ + --forcing-dir \\ + --obs-dir /mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance \\ + --cfe-dir /mnt/disk2/suma_helen_poster/cfe_py \\ + --config-file /mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json \\ + --out-dir /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble \\ + --test-forcing-dir1 \\ + --test-forcing-dir2 + +Strategy: + Re-runs DA analysis (same as Phase 1 of run_perturbation_da_on.py) to get + the per-issue-time analysis state. Then generates N_FORCING sets of perturbed + forcing sequences and N_HYDRO perturbed states. Runs all 600 combinations. + + To avoid re-running Phase 1, if the Phase 1 snapshot file exists from a + prior run_perturbation_da_on.py run, it is loaded directly. + +NOTE: All special characters kept ASCII to avoid encoding issues on GPU. +""" + +import argparse +import os +import sys +import json +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +N_FORCING = 30 +N_HYDRO = 20 +N_MEMBERS = N_FORCING * N_HYDRO # 600 + +TEST_START = "2024-08-24 00:00:00" # 1 month spinup before Helene window +TEST_END = "2024-10-31 23:00:00" + +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-30 06:00:00") +FORECAST_HOURS = 18 + +PRECIP_SIGMA = 0.15 +PET_SIGMA = 0.10 +STATE_FRAC = 0.05 + +CAT_ID = None +OBS_FILE = None +CFE_CONFIG_FILE = None +TEST_FORCING_FILE = None +OUT_DIR = None +HARDCODED_R = None +DIRECT_VARIANCE = False # when True: R = krig_var directly (no Vrugt formula) +bmi_cfe = None + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(): + df = pd.read_csv(TEST_FORCING_FILE) + df["date"] = pd.to_datetime(df["time"]).dt.strftime("%Y-%m-%d %H:%M:%S") + df["total_precipitation"] = df["APCP_surface"] * 3600.0 + df["potential_evaporation"] = priestley_taylor_pet( + df["DSWRF_surface"].values, df["TMP_2maboveground"].values) + return df + + +def load_obs(): + df = pd.read_csv(OBS_FILE) + date_col = next((c for c in df.columns + if c.lower() in ("date", "time", "datetime", "timestamp")), None) + q_col = next((c for c in df.columns + if "qkrig" in c.lower() and "var" not in c.lower()), None) + var_col = next((c for c in df.columns if "var" in c.lower()), None) + df[date_col] = pd.to_datetime(df[date_col]) + df = df.sort_values(date_col).reset_index(drop=True) + date_strs = df[date_col].dt.strftime("%Y-%m-%d %H:%M:%S") + obs_dict = dict(zip(date_strs, df[q_col])) + var_vals = df[var_col] if var_col else pd.Series([0.07] * len(df)) + var_dict = dict(zip(date_strs, var_vals)) + return obs_dict, var_dict + + +def write_model_config(best_params): + """Write the BMI config file once; reused for every model instantiation.""" + with open(CFE_CONFIG_FILE) as f: + cfg = json.load(f) + cfg["forcing_file"] = TEST_FORCING_FILE + cfg["soil_params"]["bb"] = best_params["bb"] + cfg["soil_params"]["smcmax"] = best_params["smcmax"] + cfg["soil_params"]["satdk"] = best_params["satdk"] + cfg["slop"] = best_params["slop"] + cfg["max_gw_storage"] = best_params["max_gw_storage"] + cfg["expon"] = best_params["expon"] + cfg["Cgw"] = best_params["Cgw"] + cfg["K_lf"] = best_params["K_lf"] + cfg["K_nash"] = best_params["K_nash"] + cfg["partition_scheme"] = ("Schaake" if best_params["scheme"] <= 0.5 + else "Xinanjiang") + tmp_cfg = str(OUT_DIR / f"{CAT_ID}_bmi_config_crossed_stable.json") + with open(tmp_cfg, "w") as f: + json.dump(cfg, f) + return tmp_cfg + + +def make_model(cfg_path): + """Instantiate and initialize a CFE model from a pre-written config path.""" + m = bmi_cfe.BMI_CFE(cfg_file=cfg_path) + m.initialize() + return m + + +def get_state(m): + return { + "soil_m": m.soil_reservoir["storage_m"], + "gw_m": m.gw_reservoir["storage_m"], + "nash0": float(m.nash_storage[0]), + "nash1": float(m.nash_storage[1]), + } + + +def set_state(m, state): + m.soil_reservoir["storage_m"] = state["soil_m"] + m.gw_reservoir["storage_m"] = state["gw_m"] + m.nash_storage[0] = state["nash0"] + m.nash_storage[1] = state["nash1"] + + +def perturb_state(state, rng, frac=STATE_FRAC): + return { + "soil_m": max(state["soil_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "nash0": max(state["nash0"] * (1.0 + frac * rng.standard_normal()), 0.0), + "nash1": max(state["nash1"] * (1.0 + frac * rng.standard_normal()), 0.0), + } + + +def run_model_step(m, precip_mmh, pet_mmh): + m.set_value("atmosphere_water__time_integral_of_precipitation_mass_flux", + precip_mmh / 1000.0) + m.set_value("water_potential_evaporation_flux", + pet_mmh / 1000.0 / 3600.0) + m.update() + return m.get_value("land_surface_water__runoff_depth") * 1000.0 + + +def enkf_update(state, q_sim, y_obs, krig_var, rng): + if HARDCODED_R is not None: + R = HARDCODED_R + elif DIRECT_VARIANCE: + R = max(krig_var, 1e-6) + else: + R = max((0.10 * max(y_obs, 0.0)) ** 2 + 0.001 * krig_var, 1e-6) + R = max(R, 1e-6) + P_yy = max(q_sim * 0.01, 1e-6) + K = P_yy / (P_yy + R) + scale = K * (y_obs - q_sim) / max(abs(q_sim), 1e-6) + return { + "soil_m": max(state["soil_m"] * (1.0 + scale), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + scale), 1e-6), + "nash0": max(state["nash0"] + state["nash0"] * scale, 0.0), + "nash1": max(state["nash1"] + state["nash1"] * scale, 0.0), + } + + +def run_phase1_da(cfg_path, df_test, obs_dict, var_dict, rng): + """Run DA through test period; return snapshot dict {timestamp -> state}.""" + print(f" Phase 1: DA analysis through test period...") + model = make_model(cfg_path) + snapshots = {} + for _, row in df_test.iterrows(): + date_str = row["date"] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + q_sim = run_model_step(model, P, E) + y_obs = obs_dict.get(date_str, np.nan) + if not np.isnan(y_obs): + krig_var = var_dict.get(date_str, 0.07) + state = get_state(model) + updated = enkf_update(state, q_sim, y_obs, krig_var, rng) + set_state(model, updated) + date_ts = pd.Timestamp(date_str) + if HELENE_START <= date_ts <= HELENE_END: + snapshots[date_ts] = get_state(model) + model.finalize() + print(f" Phase 1 done: {len(snapshots)} issue-time snapshots") + return snapshots + + +def main(): + global CAT_ID, OBS_FILE, CFE_CONFIG_FILE, TEST_FORCING_FILE + global OUT_DIR, HARDCODED_R, DIRECT_VARIANCE, bmi_cfe + + parser = argparse.ArgumentParser() + parser.add_argument("--cat-id", required=True) + parser.add_argument("--forcing-dir", required=True) + parser.add_argument("--obs-dir", required=True) + parser.add_argument("--cfe-dir", required=True) + parser.add_argument("--config-file", required=True) + parser.add_argument("--out-dir", required=True) + parser.add_argument("--test-forcing-dir1", default=None) + parser.add_argument("--test-forcing-dir2", default=None) + parser.add_argument("--hardcoded-r", type=float, default=None, + help="Fixed R value for EnKF (e.g. 0.07). " + "Omit to use dynamic Vrugt R: R=(0.10*y)^2 + 0.001*krig_var") + parser.add_argument("--direct-variance", action="store_true", default=False, + help="Use krig_var directly as R (no Vrugt scaling)") + parser.add_argument("--rng-seed", type=int, default=None, + help="RNG seed for reproducibility (default: hash of cat-id)") + parser.add_argument("--n-forcing", type=int, default=N_FORCING, + help="Number of met forcing draws (default 30)") + parser.add_argument("--n-hydro", type=int, default=N_HYDRO, + help="Number of hydro-state draws (default 20)") + args = parser.parse_args() + + CAT_ID = args.cat_id + OBS_FILE = os.path.join(args.obs_dir, f"{CAT_ID}.csv") + CFE_CONFIG_FILE = args.config_file + HARDCODED_R = args.hardcoded_r + DIRECT_VARIANCE = args.direct_variance + OUT_DIR = Path(args.out_dir) / CAT_ID + OUT_DIR.mkdir(parents=True, exist_ok=True) + + n_fa = args.n_forcing + n_ha = args.n_hydro + n_total = n_fa * n_ha + print(f"[{CAT_ID}] Crossed ensemble: {n_fa} forcing x {n_ha} hydro = {n_total} members") + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + bmi_cfe = _bmi_cfe + + # Reuse combined forcing CSV from run_perturbation_da_on.py if it exists + # (avoids writing a duplicate large file to disk) + da_on_combined = (Path(args.out_dir).parent / "v2_perturbation_da_on" / + CAT_ID / f"{CAT_ID}_nwm_operational_combined.csv") + crossed_combined = OUT_DIR / f"{CAT_ID}_nwm_operational_combined.csv" + + if da_on_combined.exists(): + TEST_FORCING_FILE = str(da_on_combined) + print(f" Reusing combined forcing: {TEST_FORCING_FILE}") + elif crossed_combined.exists(): + TEST_FORCING_FILE = str(crossed_combined) + print(f" Reusing combined forcing: {TEST_FORCING_FILE}") + elif args.test_forcing_dir1 and args.test_forcing_dir2: + f1 = os.path.join(args.test_forcing_dir1, f"{CAT_ID}.csv") + f2 = os.path.join(args.test_forcing_dir2, f"{CAT_ID}.csv") + df1, df2 = pd.read_csv(f1), pd.read_csv(f2) + combined = (pd.concat([df1, df2], ignore_index=True) + .drop_duplicates(subset="time").sort_values("time")) + TEST_FORCING_FILE = str(crossed_combined) + combined.to_csv(TEST_FORCING_FILE, index=False) + print(f" Wrote combined forcing: {TEST_FORCING_FILE}") + else: + print("Need --test-forcing-dir1 and --test-forcing-dir2"); return + + best_params_file = OUT_DIR / f"{CAT_ID}_best_params.json" + if not best_params_file.exists(): + # Also check the DA results dir from run_perturbation_da_on.py + alt = Path(args.out_dir).parent / "v2_perturbation_da_on" / CAT_ID / f"{CAT_ID}_best_params.json" + if alt.exists(): + best_params_file = alt + else: + print(f"Missing best params at {best_params_file}"); return + with open(best_params_file) as f: + best_params = json.load(f)["best_parameters"] + + obs_dict, var_dict = load_obs() + df_forcing = load_test_forcing() + t_mask = (df_forcing["date"] >= TEST_START) & (df_forcing["date"] <= TEST_END) + df_test = df_forcing[t_mask].reset_index(drop=True) + + # Write BMI config once — reused for all 90,600 model instantiations + cfg_path = write_model_config(best_params) + print(f" BMI config written once: {cfg_path}") + + rng = np.random.default_rng(args.rng_seed if args.rng_seed is not None + else hash(CAT_ID) & 0x7fffffff) + + # Phase 1: get DA analysis snapshots + # Check if a prior run already saved them (skip re-running DA if so) + snap_cache = OUT_DIR / f"{CAT_ID}_da_snapshots.parquet" + if snap_cache.exists(): + print(f" Loading cached DA snapshots from {snap_cache}") + df_snaps = pd.read_parquet(snap_cache) + snapshots = { + pd.Timestamp(row["issue_time"]): { + "soil_m": row["soil_m"], "gw_m": row["gw_m"], + "nash0": row["nash0"], "nash1": row["nash1"], + } + for _, row in df_snaps.iterrows() + } + else: + snapshots = run_phase1_da(cfg_path, df_test, obs_dict, var_dict, rng) + + issue_times = sorted(snapshots.keys()) + t0_to_idx = {t: i for i, t in enumerate(df_test["date"].apply(pd.Timestamp))} + + # ------------------------------------------------------------------ # + # Crossed ensemble: for each issue_time, build n_fa forcing draws and # + # n_ha state draws, then run all n_fa x n_ha combinations. # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Running crossed ensemble ({n_total} members x " + f"{len(issue_times)} issue times)...") + + all_records = [] + all_hydro_draw_records = [] # (issue_time, hydro_draw_j, states) + all_forcing_draw_records = [] # (issue_time, forcing_draw_i, lead, P_pert, E_pert, scales) + + for t_idx, t0 in enumerate(issue_times): + snap = snapshots[t0] + t0_str = t0.strftime("%Y-%m-%d %H:%M:%S") + if t0 not in t0_to_idx: + continue + start_idx = t0_to_idx[t0] + + # Pre-draw n_fa forcing perturbation sequences + forcing_seqs = [] + for i in range(n_fa): + seq = [] + for lead in range(1, FORECAST_HOURS + 1): + fi = start_idx + lead + if fi >= len(df_test): + seq.append((0.0, 0.0)) + all_forcing_draw_records.append({ + "issue_time": t0_str, "forcing_draw_i": i, + "lead_hour": lead, + "P_pert_mmh": 0.0, "E_pert_mmh": 0.0, + "precip_scale": np.nan, "pet_scale": np.nan, + }) + continue + row = df_test.iloc[fi] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + mu_p = -0.5 * PRECIP_SIGMA ** 2 + P_pert = float(P * rng.lognormal(mu_p, PRECIP_SIGMA)) if P > 0 else 0.0 + E_pert = max(float(E * (1.0 + PET_SIGMA * rng.standard_normal())), 0.0) + seq.append((P_pert, E_pert)) + all_forcing_draw_records.append({ + "issue_time": t0_str, "forcing_draw_i": i, + "lead_hour": lead, + "P_pert_mmh": P_pert, "E_pert_mmh": E_pert, + "precip_scale": P_pert / P if P > 0 else np.nan, + "pet_scale": E_pert / E if E > 0 else np.nan, + }) + forcing_seqs.append(seq) + + # Pre-draw n_ha hydro-state perturbations + state_draws = [snap if j == 0 else perturb_state(snap, rng) + for j in range(n_ha)] + + # Record each hydro draw's actual initial state + for j, st in enumerate(state_draws): + all_hydro_draw_records.append({ + "issue_time": t0_str, "hydro_draw_j": j, + "soil_m": st["soil_m"], "gw_m": st["gw_m"], + "nash0": st["nash0"], "nash1": st["nash1"], + }) + + # q_matrix shape: (FORECAST_HOURS, n_fa, n_ha) + q_matrix = np.full((FORECAST_HOURS, n_fa, n_ha), np.nan) + + for i, f_seq in enumerate(forcing_seqs): + for j, state_j in enumerate(state_draws): + m = make_model(cfg_path) + set_state(m, state_j) + for lead in range(1, FORECAST_HOURS + 1): + P_pert, E_pert = f_seq[lead - 1] + q = run_model_step(m, P_pert, E_pert) + q_matrix[lead - 1, i, j] = q + m.finalize() + + # Flatten 600-member matrix into rows + for lead in range(1, FORECAST_HOURS + 1): + row_data = {"issue_time": t0_str, "lead_hour": lead} + mem_idx = 0 + for i in range(n_fa): + for j in range(n_ha): + row_data[f"member_{mem_idx:04d}"] = q_matrix[lead - 1, i, j] + mem_idx += 1 + all_records.append(row_data) + + if (t_idx + 1) % 10 == 0 or (t_idx + 1) == len(issue_times): + print(f" {t_idx + 1}/{len(issue_times)} issue times complete") + + df_out = pd.DataFrame(all_records) + out_path = OUT_DIR / f"{CAT_ID}_crossed_ensemble.parquet" + df_out.to_parquet(out_path, index=False) + print(f"Saved: {out_path}") + print(f" Shape: {df_out.shape} " + f"({df_out['issue_time'].nunique()} issue times x " + f"{FORECAST_HOURS} leads x {n_total} members)") + + # ------------------------------------------------------------------ # + # Provenance files — enable full ensemble traceability and restart # + # ------------------------------------------------------------------ # + + # 1. DA analysis snapshots — unperturbed state at each issue_time. + # Load any row and set_state() to restart from that init time. + snap_records = [ + {"issue_time": t.strftime("%Y-%m-%d %H:%M:%S"), **s} + for t, s in snapshots.items() + ] + snap_path = OUT_DIR / f"{CAT_ID}_da_snapshots.parquet" + pd.DataFrame(snap_records).to_parquet(snap_path, index=False) + print(f"Saved: {snap_path} ({len(snap_records)} snapshots)") + + # 2. Member manifest — decoder ring: member_col -> (forcing_draw_i, hydro_draw_j). + # member_k uses forcing draw k//n_ha and hydro draw k%n_ha. + manifest = [ + {"member": f"member_{k:04d}", "member_idx": k, + "forcing_draw_i": k // n_ha, "hydro_draw_j": k % n_ha} + for k in range(n_total) + ] + manifest_path = OUT_DIR / f"{CAT_ID}_member_manifest.csv" + pd.DataFrame(manifest).to_csv(manifest_path, index=False) + print(f"Saved: {manifest_path} ({n_total} members)") + + # 3. Hydro draw states — actual perturbed initial state for each of + # n_ha draws at every issue_time. Links member -> initial conditions. + hydro_path = OUT_DIR / f"{CAT_ID}_hydro_draw_states.parquet" + pd.DataFrame(all_hydro_draw_records).to_parquet(hydro_path, index=False) + print(f"Saved: {hydro_path} ({len(all_hydro_draw_records)} rows)") + + # 4. Forcing draw sequences — actual P/PET perturbation values for each + # of n_fa draws at every (issue_time, lead_hour). Links member -> forcing. + forcing_path = OUT_DIR / f"{CAT_ID}_forcing_draw_sequences.parquet" + pd.DataFrame(all_forcing_draw_records).to_parquet(forcing_path, index=False) + print(f"Saved: {forcing_path} ({len(all_forcing_draw_records)} rows)") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/run_lead_time_forecast_sweep.py b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/run_lead_time_forecast_sweep.py new file mode 100644 index 00000000..65bf45bb --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/run_lead_time_forecast_sweep.py @@ -0,0 +1,388 @@ +""" +Forecast lead-time evaluation for F4 (dynamic variance direct). + +For each issue-time t0 in a sampling schedule across the test period, fork the +ensemble at t0 and free-run 18 hours of forecast. Two parallel trajectories are +maintained throughout the test period: + + da: Production setup. DA on, forcing perturbed, process noise on. + R(t) = sigma2_krig directly from obs file variance column. + openloop: No DA from t=0. Forcing perturbed, process noise on. + +At each issue time t0, the state of each scenario's ensemble is copied into a +forecast ensemble that free-runs 18 hours with: + - DA off + - process noise off + - forcing perturbed (lognormal precip, Gaussian PET) + +Issue-time schedule: + - Base cadence: every --base-step-h hours (default 6h) + - Densified to hourly across the Helene window (2024-09-24 -> 2024-09-28) + +Outputs (stacked across all issue times): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv + Columns: issue_time, lead_hour, valid_time, member_00 .. member_19 (mm/h) +""" + +import argparse +import os +import sys +import json +import importlib.util +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +SPINUP_START = "2023-02-01 00:00:00" +SPINUP_END = "2023-09-30 23:00:00" +TEST_START = "2023-10-01 00:00:00" +TEST_END = "2024-10-31 23:00:00" + +FORECAST_LEAD_HOURS = 18 + +DEFAULT_DENSE_START = "2024-09-24 00:00:00" +DEFAULT_DENSE_END = "2024-09-28 23:00:00" + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(test_forcing_file): + df = pd.read_csv(test_forcing_file) + df['date'] = pd.to_datetime(df['time']).dt.strftime('%Y-%m-%d %H:%M:%S') + df['total_precipitation'] = df['APCP_surface'] * 3600.0 + df['potential_evaporation'] = priestley_taylor_pet( + df['DSWRF_surface'].values, df['TMP_2maboveground'].values + ) + return df + + +def import_enkf_class(script_path): + spec = importlib.util.spec_from_file_location("prod_v2", script_path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod.EnKFAssimilator + + +def snapshot_states(models): + return [ + { + "soil_m": float(m.soil_reservoir["storage_m"]), + "gw_m": float(m.gw_reservoir["storage_m"]), + "nash0_m": float(m.nash_storage[0]), + "nash1_m": float(m.nash_storage[1]), + } + for m in models + ] + + +def restore_states(models, snapshot): + for m, s in zip(models, snapshot): + m.soil_reservoir["storage_m"] = s["soil_m"] + m.gw_reservoir["storage_m"] = s["gw_m"] + m.nash_storage[0] = s["nash0_m"] + m.nash_storage[1] = s["nash1_m"] + + +def build_issue_time_schedule(dates_list, base_step_h, dense_start, dense_end): + dates_dt = pd.to_datetime(dates_list) + selected = set() + base_mask = (np.arange(len(dates_dt)) % base_step_h == 0) + for d in dates_dt[base_mask]: + selected.add(d.strftime('%Y-%m-%d %H:%M:%S')) + if dense_start is not None and dense_end is not None: + ds = pd.Timestamp(dense_start) + de = pd.Timestamp(dense_end) + dense_mask = (dates_dt >= ds) & (dates_dt <= de) + for d in dates_dt[dense_mask]: + selected.add(d.strftime('%Y-%m-%d %H:%M:%S')) + return sorted(selected) + + +def build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, enkf, + apply_init_perturbation=True, member_zero_clean=True): + def custom_load_forcing(self_cfe): + df = load_test_forcing(test_forcing_file) + self_cfe.forcing_data = df.rename(columns={"date": "time"}) + + models = [] + for i in range(N): + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.load_forcing_file = custom_load_forcing.__get__(m) + m.initialize() + if apply_init_perturbation and not (member_zero_clean and i == 0): + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + r = enkf.init_state_perturb_frac + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1 + r * enkf.rng.standard_normal()), 1e-6, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1 + r * enkf.rng.standard_normal()), 1e-6, gw_max)) + jitter = max(sm_max * 1e-6, 1e-9) + m.nash_storage[0] = max(n00 + jitter * enkf.rng.standard_normal(), 0.0) + m.nash_storage[1] = max(n10 + jitter * enkf.rng.standard_normal(), 0.0) + models.append(m) + return models + + +def step_ensemble(models, p_arr, e_arr): + q = np.empty(len(models), dtype=float) + for i, m in enumerate(models): + m.set_value('atmosphere_water__time_integral_of_precipitation_mass_flux', + float(p_arr[i]) / 1000.0) + m.set_value('water_potential_evaporation_flux', + float(e_arr[i]) / 1000.0 / 3600.0) + m.update() + q[i] = m.get_value('land_surface_water__runoff_depth') * 1000.0 + return q + + +def run_forecast(fcst_models, enkf_fcst, dates_list, forcing_by_date, + t0_idx, n_lead): + N = len(fcst_models) + q_matrix = np.full((n_lead, N), np.nan, dtype=float) + lead_hours = [] + valid_times = [] + for lead in range(1, n_lead + 1): + idx = t0_idx + lead + if idx >= len(dates_list): + break + valid_date = dates_list[idx] + p, e = forcing_by_date[valid_date] + p_arr, e_arr = enkf_fcst.perturb_forcing(p, e) + q = step_ensemble(fcst_models, p_arr, e_arr) + q_matrix[lead - 1, :] = q + lead_hours.append(lead) + valid_times.append(valid_date) + return lead_hours, valid_times, q_matrix + + +def apply_direct_variance(enkf_instances, obs_file): + """Override obs_var_dict with per-hour kriging variance from the obs file.""" + obs_df = pd.read_csv(obs_file) + t_col = next(c for c in obs_df.columns + if c.lower() in ('datetime', 'date', 'time', 'timestamp')) + v_col = next(c for c in obs_df.columns if 'var' in c.lower()) + obs_df[t_col] = pd.to_datetime(obs_df[t_col]) + obs_df = obs_df.set_index(t_col).sort_index() + var_dict = { + str(ts): float(v) + for ts, v in obs_df[v_col].items() + if pd.notna(v) and float(v) > 0 + } + for enkf in enkf_instances: + enkf.obs_var_dict = dict(var_dict) + print(f"[lead-time] R = direct kriging variance ({len(var_dict)} timesteps, " + f"range {min(var_dict.values()):.3e} - {max(var_dict.values()):.3e})") + + +def run(args, bmi_cfe, EnKFAssimilator): + cat_id = args.cat_id + out_dir = Path(args.out_dir) / cat_id + out_dir.mkdir(parents=True, exist_ok=True) + + f1 = os.path.join(args.test_forcing_dir1, f'{cat_id}.csv') + f2 = os.path.join(args.test_forcing_dir2, f'{cat_id}.csv') + df1 = pd.read_csv(f1) + df2 = pd.read_csv(f2) + combined = pd.concat([df1, df2], ignore_index=True) + combined = combined.drop_duplicates(subset='time').sort_values('time') + test_forcing_file = str(out_dir / f'{cat_id}_nwm_operational_combined.csv') + combined.to_csv(test_forcing_file, index=False) + + best_params_file = out_dir / f'{cat_id}_best_params.json' + if not best_params_file.exists(): + print(f"Missing {best_params_file}. Pre-stage from calibration run.") + return + with open(best_params_file) as f: + best = json.load(f)['best_parameters'] + + with open(args.config_file) as f: + cfg = json.load(f) + cfg['forcing_file'] = test_forcing_file + cfg['soil_params']['bb'] = best['bb'] + cfg['soil_params']['smcmax'] = best['smcmax'] + cfg['soil_params']['satdk'] = best['satdk'] + cfg['slop'] = best['slop'] + cfg['max_gw_storage'] = best['max_gw_storage'] + cfg['expon'] = best['expon'] + cfg['Cgw'] = best['Cgw'] + cfg['K_lf'] = best['K_lf'] + cfg['K_nash'] = best['K_nash'] + cfg['partition_scheme'] = "Schaake" if best['scheme'] <= 0.5 else "Xinanjiang" + tmp_cfg = str(out_dir / f'{cat_id}_bmi_config_temp_leadtime.json') + with open(tmp_cfg, 'w') as f: + json.dump(cfg, f) + + obs_file = os.path.join(args.obs_dir, f'{cat_id}.csv') + seed_base = args.rng_seed if args.rng_seed is not None else 0 + enkf_da = EnKFAssimilator( + n_members=args.enkf_members, obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, use_vrugt_r=(not args.no_vrugt_r), + vrugt_alpha=args.vrugt_alpha, vrugt_scale=args.vrugt_scale, + rng_seed=seed_base if args.rng_seed is not None else None, + ) + enkf_ol = EnKFAssimilator( + n_members=args.enkf_members, obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, use_vrugt_r=(not args.no_vrugt_r), + vrugt_alpha=args.vrugt_alpha, vrugt_scale=args.vrugt_scale, + rng_seed=(seed_base + 1) if args.rng_seed is not None else None, + ) + enkf_fcst = EnKFAssimilator( + n_members=args.enkf_members, obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, use_vrugt_r=(not args.no_vrugt_r), + vrugt_alpha=args.vrugt_alpha, vrugt_scale=args.vrugt_scale, + rng_seed=(seed_base + 2) if args.rng_seed is not None else None, + ) + + if args.hardcoded_r is not None: + for enkf in (enkf_da, enkf_ol, enkf_fcst): + for date in enkf.obs_var_dict: + enkf.obs_var_dict[date] = args.hardcoded_r + print(f"[lead-time] R hardcoded to {args.hardcoded_r}") + + if args.direct_variance: + apply_direct_variance((enkf_da, enkf_ol, enkf_fcst), obs_file) + + N = enkf_da.n_members + print(f"[lead-time] {cat_id} | N={N} | direct_variance={args.direct_variance} | " + f"forecast lead = {FORECAST_LEAD_HOURS}h") + + prod_models = build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, enkf_da) + openloop_models = build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, enkf_ol) + fcst_models = build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, + enkf_fcst, apply_init_perturbation=False) + + df = load_test_forcing(test_forcing_file) + sp_mask = (df['date'] >= SPINUP_START) & (df['date'] <= SPINUP_END) + df_sp = df[sp_mask] + print(f"[lead-time] spinup: {len(df_sp)} hours") + for p, e in zip(df_sp['total_precipitation'], df_sp['potential_evaporation']): + p_da, e_da = enkf_da.perturb_forcing(p, e) + step_ensemble(prod_models, p_da, e_da) + enkf_da.add_process_noise(prod_models) + p_ol, e_ol = enkf_ol.perturb_forcing(p, e) + step_ensemble(openloop_models, p_ol, e_ol) + enkf_ol.add_process_noise(openloop_models) + + t_mask = (df['date'] >= TEST_START) & (df['date'] <= TEST_END) + df_test = df[t_mask].reset_index(drop=True) + dates_list = list(df_test['date'].values) + forcing_by_date = dict(zip( + df_test['date'].values, + zip(df_test['total_precipitation'].values, + df_test['potential_evaporation'].values))) + + issue_times = build_issue_time_schedule( + dates_list, args.base_step_h, args.dense_start, args.dense_end) + issue_set = set(issue_times) + print(f"[lead-time] test period: {len(dates_list)} hours | " + f"{len(issue_times)} issue times " + f"(base={args.base_step_h}h, dense={args.dense_start or 'none'}" + f"..{args.dense_end or 'none'})") + + da_rows = [] + ol_rows = [] + + for h, current_date in enumerate(dates_list): + p, e = forcing_by_date[current_date] + + p_da, e_da = enkf_da.perturb_forcing(p, e) + q_da = step_ensemble(prod_models, p_da, e_da) + enkf_da.update_states(prod_models, current_date, q_da) + enkf_da.add_process_noise(prod_models) + + p_ol, e_ol = enkf_ol.perturb_forcing(p, e) + q_ol = step_ensemble(openloop_models, p_ol, e_ol) + enkf_ol.add_process_noise(openloop_models) + + if current_date in issue_set: + for scenario_label, src_models, rows_acc in [ + ('da', prod_models, da_rows), + ('openloop', openloop_models, ol_rows), + ]: + snap = snapshot_states(src_models) + restore_states(fcst_models, snap) + leads, valids, qm = run_forecast( + fcst_models, enkf_fcst, + dates_list, forcing_by_date, + t0_idx=h, n_lead=FORECAST_LEAD_HOURS) + for k, (lead, vt) in enumerate(zip(leads, valids)): + rows_acc.append((current_date, lead, vt, *qm[k, :])) + + for m in prod_models + openloop_models + fcst_models: + m.finalize() + + cols = (['issue_time', 'lead_hour', 'valid_time'] + + [f'member_{i:02d}' for i in range(N)]) + da_path = out_dir / f'{cat_id}_lead_time_forecasts_da.csv' + ol_path = out_dir / f'{cat_id}_lead_time_forecasts_openloop.csv' + pd.DataFrame(da_rows, columns=cols).to_csv(da_path, index=False) + pd.DataFrame(ol_rows, columns=cols).to_csv(ol_path, index=False) + print(f"[lead-time] saved {da_path}") + print(f"[lead-time] saved {ol_path}") + + sched_path = out_dir / f'{cat_id}_lead_time_issue_times.csv' + pd.DataFrame({'issue_time': issue_times}).to_csv(sched_path, index=False) + print(f"[lead-time] saved {sched_path}") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', required=True) + parser.add_argument('--forcing-dir', required=True) + parser.add_argument('--obs-dir', required=True) + parser.add_argument('--cfe-dir', required=True) + parser.add_argument('--config-file', required=True) + parser.add_argument('--param-bounds', required=True) + parser.add_argument('--out-dir', required=True) + parser.add_argument('--test-forcing-dir1', required=True) + parser.add_argument('--test-forcing-dir2', required=True) + parser.add_argument('--enkf-members', type=int, default=20) + parser.add_argument('--enkf-obs-error-std', type=float, default=0.05) + parser.add_argument('--no-vrugt-r', action='store_true') + parser.add_argument('--direct-variance', action='store_true', default=False, + help='Use per-hour kriging variance column from obs file as R.') + parser.add_argument('--hardcoded-r', type=float, default=None) + parser.add_argument('--vrugt-alpha', type=float, default=0.10) + parser.add_argument('--vrugt-scale', type=float, default=0.001) + parser.add_argument('--rng-seed', type=int, default=None) + parser.add_argument('--base-step-h', type=int, default=6) + parser.add_argument('--dense-start', type=str, default=DEFAULT_DENSE_START) + parser.add_argument('--dense-end', type=str, default=DEFAULT_DENSE_END) + parser.add_argument('--prod-script', default=None) + args = parser.parse_args() + + here = os.path.dirname(os.path.abspath(__file__)) + prod_script = args.prod_script or os.path.normpath( + os.path.join(here, '..', '..', '..', '1_distributed_cfe', + 'calibrate_catchment_cfe_da_v2.py')) + if not os.path.exists(prod_script): + raise FileNotFoundError( + f"Could not find production script at {prod_script}. " + f"Pass --prod-script to override.") + EnKFAssimilator = import_enkf_class(prod_script) + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + + run(args, _bmi_cfe, EnKFAssimilator) + + +if __name__ == '__main__': + main() diff --git a/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/run_perturbation_da_on.py b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/run_perturbation_da_on.py new file mode 100644 index 00000000..6ef27e96 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/run_perturbation_da_on.py @@ -0,0 +1,400 @@ +""" +run_perturbation_da_on.py — 2a/2b: ensemble spread with DA ON. + +Runs two separated perturbation arms for the Helene window (Sep 24-30), +with DA actively running during the analysis period: + + ARM A — forcing only (30 members): + DA analyzes states using Qkrig at every hour (DA ON). + At each issue_time: fork 30 members with ONLY met forcing perturbed + (different precip/PET draw per member). Hydro states = DA analysis + mean (same for all 30). Free-run 18 hours, no further DA. + + ARM B — hydro states only (20 members): + Same DA analysis. At each issue_time: fork 20 members with ONLY + hydro states perturbed (multiplicative noise on soil/GW/Nash from + the DA analysis ensemble). Forcing = deterministic. Free-run 18 hours. + +These two arms isolate "where does forecast spread come from WITH DA?" +— figures for the assimilation script section. + +Output (per catchment, per arm): + //_da_forcing_arm.csv + //_da_hydro_arm.csv + Columns: issue_time, lead_hour, member_00..member_29 (forcing arm) + issue_time, lead_hour, member_00..member_19 (hydro arm) + +Usage (run for each catchment, both arms): + python3 run_perturbation_da_on.py \\ + --cat-id cat-1016300 \\ + --forcing-dir \\ + --obs-dir /mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance \\ + --cfe-dir /mnt/disk2/suma_helen_poster/cfe_py \\ + --config-file /mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json \\ + --param-bounds /mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json \\ + --out-dir /mnt/disk2/suma_helen_poster/da_results/v2_perturbation_da_on \\ + --test-forcing-dir1 \\ + --test-forcing-dir2 \\ + --hardcoded-r 0.07 +""" + +import argparse +import os +import sys +import json +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +N_FORCING = 30 # members in forcing arm +N_HYDRO = 20 # members in hydro-state arm + +# Analysis period — run DA through entire test period +TEST_START = "2024-08-24 00:00:00" # 1 month spinup before Helene window +TEST_END = "2024-10-31 23:00:00" + +# Issue times for the Helene window (hourly, Sep 24-30) +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-30 06:00:00") +FORECAST_HOURS = 18 + +# Perturbation magnitudes (same as production) +PRECIP_SIGMA = 0.15 +PET_SIGMA = 0.10 +STATE_FRAC = 0.05 # hydro-state perturbation fraction + +# Set at runtime +CAT_ID = None +OBS_FILE = None +CFE_CONFIG_FILE = None +TEST_FORCING_FILE = None +OUT_DIR = None +HARDCODED_R = None +DIRECT_VARIANCE = False +bmi_cfe = None + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(): + df = pd.read_csv(TEST_FORCING_FILE) + df["date"] = pd.to_datetime(df["time"]).dt.strftime("%Y-%m-%d %H:%M:%S") + df["total_precipitation"] = df["APCP_surface"] * 3600.0 + df["potential_evaporation"] = priestley_taylor_pet( + df["DSWRF_surface"].values, df["TMP_2maboveground"].values) + return df + + +def load_obs(): + df = pd.read_csv(OBS_FILE) + date_col = next((c for c in df.columns + if c.lower() in ("date", "time", "datetime", "timestamp")), None) + if date_col is None: + raise ValueError(f"No date column in {OBS_FILE}. Columns: {df.columns.tolist()}") + q_col = next((c for c in df.columns if "qkrig" in c.lower() and "var" not in c.lower()), None) + if q_col is None: + raise ValueError(f"No qkrig column in {OBS_FILE}. Columns: {df.columns.tolist()}") + var_col = next((c for c in df.columns if "var" in c.lower()), None) + df[date_col] = pd.to_datetime(df[date_col]) + df = df.sort_values(date_col).reset_index(drop=True) + date_strs = df[date_col].dt.strftime("%Y-%m-%d %H:%M:%S") + obs_dict = dict(zip(date_strs, df[q_col])) + var_vals = df[var_col] if var_col else pd.Series([0.07] * len(df)) + var_dict = dict(zip(date_strs, var_vals)) + print(f" Obs: date='{date_col}', q='{q_col}', var='{var_col}'") + return obs_dict, var_dict + + +def make_model(best_params): + with open(CFE_CONFIG_FILE) as f: + cfg = json.load(f) + cfg["forcing_file"] = TEST_FORCING_FILE + cfg["soil_params"]["bb"] = best_params["bb"] + cfg["soil_params"]["smcmax"] = best_params["smcmax"] + cfg["soil_params"]["satdk"] = best_params["satdk"] + cfg["slop"] = best_params["slop"] + cfg["max_gw_storage"] = best_params["max_gw_storage"] + cfg["expon"] = best_params["expon"] + cfg["Cgw"] = best_params["Cgw"] + cfg["K_lf"] = best_params["K_lf"] + cfg["K_nash"] = best_params["K_nash"] + cfg["partition_scheme"] = ("Schaake" if best_params["scheme"] <= 0.5 + else "Xinanjiang") + tmp_cfg = str(OUT_DIR / f"{CAT_ID}_bmi_config_temp_da_on.json") + with open(tmp_cfg, "w") as f: + json.dump(cfg, f) + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.initialize() + return m + + +def get_state(m): + return { + "soil_m": m.soil_reservoir["storage_m"], + "gw_m": m.gw_reservoir["storage_m"], + "nash0": float(m.nash_storage[0]), + "nash1": float(m.nash_storage[1]), + } + + +def set_state(m, state): + m.soil_reservoir["storage_m"] = state["soil_m"] + m.gw_reservoir["storage_m"] = state["gw_m"] + m.nash_storage[0] = state["nash0"] + m.nash_storage[1] = state["nash1"] + + +def perturb_state(state, rng, frac=STATE_FRAC): + sm_max = 1.0 # placeholder — will be clipped in set_state; actual max enforced by model + out = { + "soil_m": max(state["soil_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "nash0": max(state["nash0"] * (1.0 + frac * rng.standard_normal()), 0.0), + "nash1": max(state["nash1"] * (1.0 + frac * rng.standard_normal()), 0.0), + } + return out + + +def run_model_step(m, precip_mmh, pet_mmh): + m.set_value("atmosphere_water__time_integral_of_precipitation_mass_flux", + precip_mmh / 1000.0) + m.set_value("water_potential_evaporation_flux", + pet_mmh / 1000.0 / 3600.0) + m.update() + return m.get_value("land_surface_water__runoff_depth") * 1000.0 # m/h -> mm/h + + +def enkf_update(state, q_sim, y_obs, krig_var, rng, n_members=1): + """Simple scalar EnKF update on the mean analysis state.""" + if HARDCODED_R is not None: + R = HARDCODED_R + elif DIRECT_VARIANCE: + R = max(krig_var, 1e-6) + else: + alpha = 0.10 + R = (alpha * max(y_obs, 0.0)) ** 2 + 0.001 * krig_var + R = max(R, 1e-6) + P_yy = max(q_sim * 0.01, 1e-6) # rough prior variance on Q + K = P_yy / (P_yy + R) + innovation = y_obs - q_sim + # Apply correction to each state proportionally (simple scalar gain) + scale = K * innovation / max(abs(q_sim), 1e-6) + updated = { + "soil_m": max(state["soil_m"] * (1.0 + scale), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + scale), 1e-6), + "nash0": max(state["nash0"] + state["nash0"] * scale, 0.0), + "nash1": max(state["nash1"] + state["nash1"] * scale, 0.0), + } + return updated + + +def main(): + global CAT_ID, OBS_FILE, CFE_CONFIG_FILE, TEST_FORCING_FILE + global OUT_DIR, HARDCODED_R, DIRECT_VARIANCE, bmi_cfe + + parser = argparse.ArgumentParser() + parser.add_argument("--cat-id", required=True) + parser.add_argument("--forcing-dir", required=True) + parser.add_argument("--obs-dir", required=True) + parser.add_argument("--cfe-dir", required=True) + parser.add_argument("--config-file", required=True) + parser.add_argument("--param-bounds", required=True) + parser.add_argument("--out-dir", required=True) + parser.add_argument("--test-forcing-dir1", default=None) + parser.add_argument("--test-forcing-dir2", default=None) + parser.add_argument("--hardcoded-r", type=float, default=None, + help="Fixed R (e.g. 0.07). Omit → Vrugt formula.") + parser.add_argument("--direct-variance", action="store_true", default=False, + help="Use krig_var directly as R (no Vrugt scaling)") + parser.add_argument("--rng-seed", type=int, default=None) + args = parser.parse_args() + + CAT_ID = args.cat_id + OBS_FILE = os.path.join(args.obs_dir, f"{CAT_ID}.csv") + CFE_CONFIG_FILE = args.config_file + HARDCODED_R = args.hardcoded_r + DIRECT_VARIANCE = args.direct_variance + OUT_DIR = Path(args.out_dir) / CAT_ID + OUT_DIR.mkdir(parents=True, exist_ok=True) + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + bmi_cfe = _bmi_cfe + + # Build combined test forcing + if args.test_forcing_dir1 and args.test_forcing_dir2: + f1 = os.path.join(args.test_forcing_dir1, f"{CAT_ID}.csv") + f2 = os.path.join(args.test_forcing_dir2, f"{CAT_ID}.csv") + df1, df2 = pd.read_csv(f1), pd.read_csv(f2) + combined = (pd.concat([df1, df2], ignore_index=True) + .drop_duplicates(subset="time").sort_values("time")) + TEST_FORCING_FILE = str(OUT_DIR / f"{CAT_ID}_nwm_operational_combined.csv") + combined.to_csv(TEST_FORCING_FILE, index=False) + else: + print("Need --test-forcing-dir1 and --test-forcing-dir2"); return + + best_params_file = OUT_DIR / f"{CAT_ID}_best_params.json" + if not best_params_file.exists(): + print(f"Missing best params at {best_params_file}"); return + with open(best_params_file) as f: + best_params = json.load(f)["best_parameters"] + + obs_dict, var_dict = load_obs() + df_forcing = load_test_forcing() + t_mask = (df_forcing["date"] >= TEST_START) & (df_forcing["date"] <= TEST_END) + df_test = df_forcing[t_mask].reset_index(drop=True) + + rng = np.random.default_rng(hash(CAT_ID) & 0x7fffffff) + + # ------------------------------------------------------------------ # + # Phase 1: Run DA through test period, snapshot analysis state at # + # each Helene issue time # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Phase 1: DA analysis through test period...") + model = make_model(best_params) + + helene_snapshots = {} # issue_time -> analysis state dict + + for h, row in df_test.iterrows(): + date_str = row["date"] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + q_sim = run_model_step(model, P, E) + + date_ts = pd.Timestamp(date_str) + + # DA update if obs available + y_obs = obs_dict.get(date_str, np.nan) + if not np.isnan(y_obs): + krig_var = var_dict.get(date_str, 0.07) + state = get_state(model) + updated = enkf_update(state, q_sim, y_obs, krig_var, rng) + set_state(model, updated) + + # Snapshot if this is a Helene issue time + if HELENE_START <= date_ts <= HELENE_END: + helene_snapshots[date_ts] = get_state(model) + + model.finalize() + print(f" Snapshots saved: {len(helene_snapshots)} issue times") + + # ------------------------------------------------------------------ # + # Phase 2A: Forcing arm — 30 members, fixed analyzed state, varying # + # met forcing # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Phase 2A: Forcing arm ({N_FORCING} members)...") + forcing_records = [] + + issue_times = sorted(helene_snapshots.keys()) + t0_to_idx = {t: i for i, t in enumerate(df_test["date"].apply(pd.Timestamp))} + + for t0 in issue_times: + snap = helene_snapshots[t0] + t0_str = t0.strftime("%Y-%m-%d %H:%M:%S") + + # Find index in df_test + if t0 not in t0_to_idx: + continue + start_idx = t0_to_idx[t0] + + q_matrix = np.full((FORECAST_HOURS, N_FORCING), np.nan) + + for mem in range(N_FORCING): + m = make_model(best_params) + set_state(m, snap) # all members start from same DA analysis state + + for lead in range(1, FORECAST_HOURS + 1): + fi = start_idx + lead + if fi >= len(df_test): + break + row = df_test.iloc[fi] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + # Perturb forcing only + mu_p = -0.5 * PRECIP_SIGMA ** 2 + P_pert = float(P * rng.lognormal(mu_p, PRECIP_SIGMA)) if P > 0 else 0.0 + E_pert = max(float(E * (1.0 + PET_SIGMA * rng.standard_normal())), 0.0) + q = run_model_step(m, P_pert, E_pert) + q_matrix[lead - 1, mem] = q + + m.finalize() + + row_data = {"issue_time": t0_str} + for lead in range(1, FORECAST_HOURS + 1): + row_data["lead_hour"] = lead + for mem in range(N_FORCING): + row_data[f"member_{mem:02d}"] = q_matrix[lead - 1, mem] + forcing_records.append(dict(row_data)) + + df_forcing_arm = pd.DataFrame(forcing_records) + out_a = OUT_DIR / f"{CAT_ID}_da_forcing_arm.csv" + df_forcing_arm.to_csv(out_a, index=False) + print(f" Saved: {out_a}") + + # ------------------------------------------------------------------ # + # Phase 2B: Hydro-state arm — 20 members, deterministic forcing, # + # perturbed analysis states # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Phase 2B: Hydro-state arm ({N_HYDRO} members)...") + hydro_records = [] + + for t0 in issue_times: + snap = helene_snapshots[t0] + t0_str = t0.strftime("%Y-%m-%d %H:%M:%S") + + if t0 not in t0_to_idx: + continue + start_idx = t0_to_idx[t0] + + q_matrix = np.full((FORECAST_HOURS, N_HYDRO), np.nan) + + for mem in range(N_HYDRO): + # Perturb the DA analysis state for this member + if mem == 0: + state_m = snap # member 0 = unperturbed analysis mean + else: + state_m = perturb_state(snap, rng) + + m = make_model(best_params) + set_state(m, state_m) + + for lead in range(1, FORECAST_HOURS + 1): + fi = start_idx + lead + if fi >= len(df_test): + break + row = df_test.iloc[fi] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + # Deterministic forcing + q = run_model_step(m, P, E) + q_matrix[lead - 1, mem] = q + + m.finalize() + + row_data = {"issue_time": t0_str} + for lead in range(1, FORECAST_HOURS + 1): + row_data["lead_hour"] = lead + for mem in range(N_HYDRO): + row_data[f"member_{mem:02d}"] = q_matrix[lead - 1, mem] + hydro_records.append(dict(row_data)) + + df_hydro_arm = pd.DataFrame(hydro_records) + out_b = OUT_DIR / f"{CAT_ID}_da_hydro_arm.csv" + df_hydro_arm.to_csv(out_b, index=False) + print(f" Saved: {out_b}") + print(f"[{CAT_ID}] Done.") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/run_perturbation_sensitivity.py b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/run_perturbation_sensitivity.py new file mode 100644 index 00000000..57b8d185 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/run_perturbation_sensitivity.py @@ -0,0 +1,293 @@ +""" +Perturbation-source sensitivity analysis for the multi-catchment CFE DA experiment. + +Ensemble-spread attribution: run 20 CFE members through the test +period with ONLY ONE perturbation source active at a time, no DA, no spinup. + +Three sub-experiments (selected via --source): + init : perturb initial states once at t=0 (multiplicative ~5%) + forcing : perturb hourly precip (lognormal sigma=0.15) and PET (Gaussian sigma=0.10) + process : apply per-hour process noise on states (soil 0.2%, GW 0.15%, Nash 0.5%) + +Each sub-experiment runs N=20 CFE members. DA is always OFF. Spinup is skipped — each +member starts from CFE's default initial state (modified by the perturbation if --source +init). This is the cleanest setup for attributing test-period spread to its source. + +Output: + //_sensitivity_.csv + Columns: date, member_00, member_01, ..., member_19 + +Each member column is that member's hourly Q_sim (mm/h) over the test period. + +Reads the same calibrated best_params.json as the production script. Does not need +the kriging obs file (no DA), but takes the same --obs-dir argument so it can also +log Q_obs alongside for plotting reference. + +Usage: + python3 run_perturbation_sensitivity.py \ + --cat-id cat-1016300 \ + --source forcing \ + --forcing-dir \ + --obs-dir \ + --cfe-dir \ + --config-file \ + --param-bounds \ + --out-dir \ + --test-forcing-dir1 \ + --test-forcing-dir2 +""" + +import argparse +import os +import sys +import json +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 +N_MEMBERS = 20 + +# We only run the test window for this sensitivity analysis. +TEST_START = "2023-10-01 00:00:00" +TEST_END = "2024-10-31 23:00:00" + +# Perturbation magnitudes — same as production defaults +INIT_STATE_FRAC = 0.05 +PRECIP_PERTURB = 0.15 # lognormal sigma +PET_PERTURB = 0.10 # Gaussian sigma +SOIL_PROC_NOISE = 0.002 +GW_PROC_NOISE = 0.0015 +NASH_PROC_NOISE = 0.005 + +# Set at runtime +CAT_ID = None +FORCING_FILE = None +TEST_FORCING_FILE = None +OBS_FILE = None +CFE_CONFIG_FILE = None +OUT_DIR = None +bmi_cfe = None + + +# ---------- Forcing helpers (same as production) ---------- +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(): + df = pd.read_csv(TEST_FORCING_FILE) + df['date'] = pd.to_datetime(df['time']).dt.strftime('%Y-%m-%d %H:%M:%S') + df['total_precipitation'] = df['APCP_surface'] * 3600.0 # kg/m^2/s -> mm/h + df['potential_evaporation'] = priestley_taylor_pet( + df['DSWRF_surface'].values, + df['TMP_2maboveground'].values, + ) + return df + + +# ---------- Perturbation helpers ---------- +def perturb_precip_lognormal(P, rng, sigma=PRECIP_PERTURB): + """Lognormal multiplier with mean 1; preserves zero precip.""" + mu = -0.5 * sigma ** 2 + return float(P * rng.lognormal(mu, sigma)) + + +def perturb_pet_gaussian(E, rng, sigma=PET_PERTURB): + """Gaussian multiplicative noise on PET, clipped at 0.""" + return max(float(E * (1.0 + sigma * rng.standard_normal())), 0.0) + + +def apply_init_state_perturb(m, rng, frac=INIT_STATE_FRAC): + """One-time multiplicative perturbation on the four states (called per member at t=0).""" + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1.0 + frac * rng.standard_normal()), 1e-6, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1.0 + frac * rng.standard_normal()), 1e-6, gw_max)) + # Nash buckets often start at 0; small additive jitter so they're not identical. + jitter = max(sm_max * 1e-6, 1e-9) + m.nash_storage[0] = max(n00 + jitter * rng.standard_normal(), 0.0) + m.nash_storage[1] = max(n10 + jitter * rng.standard_normal(), 0.0) + + +def apply_process_noise(m, rng): + """Per-timestep process noise on all 4 states for one member.""" + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + nash_floor = max(sm_max * 1e-4, 1e-7) + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1.0 + SOIL_PROC_NOISE * rng.standard_normal()), 0.0, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1.0 + GW_PROC_NOISE * rng.standard_normal()), 0.0, gw_max)) + m.nash_storage[0] = max( + n00 * (1.0 + NASH_PROC_NOISE * rng.standard_normal()) + + nash_floor * rng.standard_normal(), 0.0) + m.nash_storage[1] = max( + n10 * (1.0 + NASH_PROC_NOISE * rng.standard_normal()) + + nash_floor * rng.standard_normal(), 0.0) + + +# ---------- Main per-source run ---------- +def run_sensitivity(source, best_param_dict): + """Run N members through the test period with only `source` perturbation active.""" + + def custom_load_forcing(self_cfe): + df = load_test_forcing() + self_cfe.forcing_data = df.rename(columns={"date": "time"}) + + # Build a per-catchment temporary CFE config with the calibrated params + with open(CFE_CONFIG_FILE) as f: + cfg = json.load(f) + cfg['forcing_file'] = TEST_FORCING_FILE + cfg['soil_params']['bb'] = best_param_dict['bb'] + cfg['soil_params']['smcmax'] = best_param_dict['smcmax'] + cfg['soil_params']['satdk'] = best_param_dict['satdk'] + cfg['slop'] = best_param_dict['slop'] + cfg['max_gw_storage'] = best_param_dict['max_gw_storage'] + cfg['expon'] = best_param_dict['expon'] + cfg['Cgw'] = best_param_dict['Cgw'] + cfg['K_lf'] = best_param_dict['K_lf'] + cfg['K_nash'] = best_param_dict['K_nash'] + cfg['partition_scheme'] = "Schaake" if best_param_dict['scheme'] <= 0.5 else "Xinanjiang" + + tmp_cfg = str(OUT_DIR / f'{CAT_ID}_bmi_config_temp_sensitivity_{source}.json') + with open(tmp_cfg, 'w') as f: + json.dump(cfg, f) + + # Deterministic seed per (catchment, source) so each run is reproducible + seed = hash((CAT_ID, source)) & 0x7fffffff + rng = np.random.default_rng(seed) + + print(f"[sensitivity] {CAT_ID} | source={source} | N={N_MEMBERS} | seed={seed}") + + # Build N members. Initial-state perturbation only happens when source == 'init'. + models = [] + for i in range(N_MEMBERS): + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.load_forcing_file = custom_load_forcing.__get__(m) + m.initialize() + if source == 'init' and i > 0: + apply_init_state_perturb(m, rng) + models.append(m) + + # Skip spinup. Go directly to the test period. + df = load_test_forcing() + t_mask = (df['date'] >= TEST_START) & (df['date'] <= TEST_END) + df_test = df[t_mask] + + # Pre-allocate per-member Q arrays + n_hours = len(df_test) + q_matrix = np.full((n_hours, N_MEMBERS), np.nan, dtype=float) + dates_out = df_test['date'].values + + for h, (p, e) in enumerate(zip(df_test['total_precipitation'], + df_test['potential_evaporation'])): + for i, m in enumerate(models): + if source == 'forcing': + p_i = perturb_precip_lognormal(p, rng) + e_i = perturb_pet_gaussian(e, rng) + else: + p_i = float(p) + e_i = float(e) + m.set_value('atmosphere_water__time_integral_of_precipitation_mass_flux', p_i / 1000) + m.set_value('water_potential_evaporation_flux', e_i / 1000 / 3600) + m.update() + q_matrix[h, i] = m.get_value('land_surface_water__runoff_depth') * 1000 # m/h -> mm/h + + # Process noise applied after each hour, but only when source == 'process' + if source == 'process': + for m in models: + apply_process_noise(m, rng) + + for m in models: + m.finalize() + + # Write per-member CSV: date + 20 columns + out = {'date': dates_out} + for i in range(N_MEMBERS): + out[f'member_{i:02d}'] = q_matrix[:, i] + df_out = pd.DataFrame(out) + out_path = OUT_DIR / f'{CAT_ID}_sensitivity_{source}.csv' + df_out.to_csv(out_path, index=False) + print(f"[sensitivity] saved {out_path}") + print(f"[sensitivity] mean ensemble spread (std across members) over test period: " + f"{float(q_matrix.std(axis=1).mean()):.5f} mm/h") + + +def main(): + global CAT_ID, FORCING_FILE, TEST_FORCING_FILE, OBS_FILE, CFE_CONFIG_FILE + global OUT_DIR, bmi_cfe + + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', required=True) + parser.add_argument('--source', required=True, + choices=['init', 'forcing', 'process'], + help="Which perturbation source to isolate") + parser.add_argument('--forcing-dir', required=True) + parser.add_argument('--obs-dir', required=True, + help='Kept for symmetry with production script; not used since DA is off') + parser.add_argument('--cfe-dir', required=True) + parser.add_argument('--config-file', required=True) + parser.add_argument('--param-bounds', required=True, + help='Kept for symmetry; not read in this script') + parser.add_argument('--out-dir', required=True) + parser.add_argument('--test-forcing-dir1', default=None) + parser.add_argument('--test-forcing-dir2', default=None) + args = parser.parse_args() + + CAT_ID = args.cat_id + FORCING_FILE = os.path.join(args.forcing_dir, f'{CAT_ID}.csv') + OBS_FILE = os.path.join(args.obs_dir, f'{CAT_ID}.csv') + CFE_CONFIG_FILE = args.config_file + OUT_DIR = Path(args.out_dir) / CAT_ID + OUT_DIR.mkdir(parents=True, exist_ok=True) + + # Build combined test forcing CSV (same idiom as production) + if args.test_forcing_dir1 and args.test_forcing_dir2: + f1 = os.path.join(args.test_forcing_dir1, f'{CAT_ID}.csv') + f2 = os.path.join(args.test_forcing_dir2, f'{CAT_ID}.csv') + if os.path.exists(f1) and os.path.exists(f2): + df1 = pd.read_csv(f1) + df2 = pd.read_csv(f2) + combined = pd.concat([df1, df2], ignore_index=True) + combined = combined.drop_duplicates(subset='time').sort_values('time') + TEST_FORCING_FILE = str(OUT_DIR / f'{CAT_ID}_nwm_operational_combined.csv') + combined.to_csv(TEST_FORCING_FILE, index=False) + else: + print(f"Warning: test forcing files not found for {CAT_ID}, exiting") + return + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + bmi_cfe = _bmi_cfe + + best_params_file = OUT_DIR / f'{CAT_ID}_best_params.json' + if not best_params_file.exists(): + print(f"No best params file at {best_params_file}. Pre-stage from Run 3.") + return + with open(best_params_file) as f: + saved = json.load(f) + best_param_dict = saved['best_parameters'] + + run_sensitivity(args.source, best_param_dict) + + +if __name__ == '__main__': + main() diff --git a/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/run_production_per_member.py b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/run_production_per_member.py new file mode 100644 index 00000000..d9c1c4ac --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/2_assimilation/run_production_per_member.py @@ -0,0 +1,314 @@ +""" +Run the production DA pipeline for ONE catchment and save per-member streamflow. + +Identical to calibrate_catchment_cfe_da_v2.py in algorithm (init perturbation + +forcing perturbation + process noise + true EnKF with Vrugt R), but writes one +column per ensemble member instead of only the ensemble mean. This is for +the per-member factor-decomposition plot. + +To keep the comparison apples-to-apples with the sensitivity runs (which skip +spinup), this script also SKIPS the spinup loop. Each member starts directly +from CFE's default initial state, plus the t=0 init perturbation. + +Output: + //_production_per_member.csv + Columns: date, member_00, member_01, ..., member_19 + +Imports EnKFAssimilator from calibrate_catchment_cfe_da_v2 to guarantee the +DA math matches production exactly. + +Usage: + python3 run_production_per_member.py \ + --cat-id cat-1016300 \ + --forcing-dir /mnt/.../nwm_retro_catchment_forcings \ + --obs-dir /mnt/.../catchment_ts_03463300_with_variance \ + --cfe-dir /mnt/.../cfe_py \ + --config-file /mnt/.../cat_03463300_bmi_config_cfe.json \ + --param-bounds /mnt/.../CFE_parameter_bounds.json \ + --out-dir /mnt/.../v2_production_per_member \ + --test-forcing-dir1 ... \ + --test-forcing-dir2 ... \ + --enkf-members 20 +""" + +import argparse +import os +import sys +import json +import importlib.util +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +TEST_START = "2023-10-01 00:00:00" +TEST_END = "2024-10-31 23:00:00" + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(test_forcing_file): + df = pd.read_csv(test_forcing_file) + df['date'] = pd.to_datetime(df['time']).dt.strftime('%Y-%m-%d %H:%M:%S') + df['total_precipitation'] = df['APCP_surface'] * 3600.0 + df['potential_evaporation'] = priestley_taylor_pet( + df['DSWRF_surface'].values, df['TMP_2maboveground'].values + ) + return df + + +def import_enkf_class(script_path): + """Load EnKFAssimilator from the production script without triggering its main().""" + spec = importlib.util.spec_from_file_location("prod_v2", script_path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod.EnKFAssimilator + + +def run(args, bmi_cfe, EnKFAssimilator): + cat_id = args.cat_id + out_dir = Path(args.out_dir) / cat_id + out_dir.mkdir(parents=True, exist_ok=True) + + # Combined test forcing + f1 = os.path.join(args.test_forcing_dir1, f'{cat_id}.csv') + f2 = os.path.join(args.test_forcing_dir2, f'{cat_id}.csv') + df1 = pd.read_csv(f1) + df2 = pd.read_csv(f2) + combined = pd.concat([df1, df2], ignore_index=True) + combined = combined.drop_duplicates(subset='time').sort_values('time') + test_forcing_file = str(out_dir / f'{cat_id}_nwm_operational_combined.csv') + combined.to_csv(test_forcing_file, index=False) + + # Load best_params + best_params_file = out_dir / f'{cat_id}_best_params.json' + if not best_params_file.exists(): + print(f"Missing {best_params_file}. Pre-stage from Run 3.") + return + with open(best_params_file) as f: + saved = json.load(f) + best = saved['best_parameters'] + + # Build temp CFE config with calibrated params + with open(args.config_file) as f: + cfg = json.load(f) + cfg['forcing_file'] = test_forcing_file + cfg['soil_params']['bb'] = best['bb'] + cfg['soil_params']['smcmax'] = best['smcmax'] + cfg['soil_params']['satdk'] = best['satdk'] + cfg['slop'] = best['slop'] + cfg['max_gw_storage'] = best['max_gw_storage'] + cfg['expon'] = best['expon'] + cfg['Cgw'] = best['Cgw'] + cfg['K_lf'] = best['K_lf'] + cfg['K_nash'] = best['K_nash'] + cfg['partition_scheme'] = "Schaake" if best['scheme'] <= 0.5 else "Xinanjiang" + tmp_cfg = str(out_dir / f'{cat_id}_bmi_config_temp_per_member.json') + with open(tmp_cfg, 'w') as f: + json.dump(cfg, f) + + # Build the EnKF assimilator (same defaults as production) + obs_file = os.path.join(args.obs_dir, f'{cat_id}.csv') + enkf = EnKFAssimilator( + n_members=args.enkf_members, + obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, + use_vrugt_r=(not args.no_vrugt_r) and (not getattr(args, 'direct_variance', False)), + vrugt_alpha=args.vrugt_alpha, + vrugt_scale=args.vrugt_scale, + rng_seed=args.rng_seed, + ) + N = enkf.n_members + print(f"[per-member] {cat_id} | N={N} | use_vrugt_r={enkf.use_vrugt_r}") + if getattr(args, 'hardcoded_r', None) is not None: + enkf.obs_var_dict = {k: args.hardcoded_r for k in enkf.obs_var_dict} + print(f"[per-member] obs_var_dict overridden: R = {args.hardcoded_r} (fixed)") + if getattr(args, 'direct_variance', False): + obs_df = pd.read_csv(obs_file) + t_col = next(c for c in obs_df.columns if c.lower() in ('datetime', 'date', 'time', 'timestamp')) + v_col = next(c for c in obs_df.columns if 'var' in c.lower()) + obs_df[t_col] = pd.to_datetime(obs_df[t_col]) + obs_df = obs_df.set_index(t_col).sort_index() + enkf.obs_var_dict = { + str(ts): float(v) + for ts, v in obs_df[v_col].items() + if pd.notna(v) and float(v) > 0 + } + print(f"[per-member] obs_var_dict overridden: R = σ² direct ({len(enkf.obs_var_dict)} timesteps)") + + # Custom forcing loader for CFE BMI + def custom_load_forcing(self_cfe): + df = load_test_forcing(test_forcing_file) + self_cfe.forcing_data = df.rename(columns={"date": "time"}) + + # Build N CFE instances; perturb initial states (member 0 left clean) + models = [] + init_states_per_member = [] # capture initial-state snapshot after perturbation + for i in range(N): + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.load_forcing_file = custom_load_forcing.__get__(m) + m.initialize() + if i > 0: + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + r = enkf.init_state_perturb_frac + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1 + r * enkf.rng.standard_normal()), 1e-6, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1 + r * enkf.rng.standard_normal()), 1e-6, gw_max)) + jitter = max(sm_max * 1e-6, 1e-9) + m.nash_storage[0] = max(n00 + jitter * enkf.rng.standard_normal(), 0.0) + m.nash_storage[1] = max(n10 + jitter * enkf.rng.standard_normal(), 0.0) + # Snapshot the (possibly-perturbed) initial state for this member + init_states_per_member.append({ + "soil_m": float(m.soil_reservoir["storage_m"]), + "gw_m": float(m.gw_reservoir["storage_m"]), + "nash0_m": float(m.nash_storage[0]), + "nash1_m": float(m.nash_storage[1]), + "soil_max_m": float(m.soil_reservoir["storage_max_m"]), + "gw_max_m": float(m.gw_reservoir["storage_max_m"]), + }) + models.append(m) + + # SKIP SPINUP — go directly to the test period (matches sensitivity-script setup) + df = load_test_forcing(test_forcing_file) + t_mask = (df['date'] >= TEST_START) & (df['date'] <= TEST_END) + df_test = df[t_mask] + n_hours = len(df_test) + q_matrix = np.full((n_hours, N), np.nan, dtype=float) + precip_matrix = np.full((n_hours, N), np.nan, dtype=float) # mm/h per member + pet_matrix = np.full((n_hours, N), np.nan, dtype=float) # mm/h per member + dates_out = df_test['date'].values + + for h, (p, e, current_date) in enumerate(zip( + df_test['total_precipitation'], + df_test['potential_evaporation'], + df_test['date'])): + + # Perturb forcing per member + p_arr, e_arr = enkf.perturb_forcing(p, e) + # Capture the per-member perturbed forcing for later plotting/diagnostics + precip_matrix[h, :] = p_arr + pet_matrix[h, :] = e_arr + + # Advance each member one hour + ensemble_q = np.empty(N, dtype=float) + for i, m in enumerate(models): + m.set_value('atmosphere_water__time_integral_of_precipitation_mass_flux', + float(p_arr[i]) / 1000) + m.set_value('water_potential_evaporation_flux', + float(e_arr[i]) / 1000 / 3600) + m.update() + ensemble_q[i] = m.get_value('land_surface_water__runoff_depth') * 1000 # mm/h + + # Record pre-analysis forecast per member (this is the "actual production" value) + q_matrix[h, :] = ensemble_q + + # DA step (uses production EnKFAssimilator) + enkf.update_states(models, current_date, ensemble_q) + + # Process noise after DA + enkf.add_process_noise(models) + + for m in models: + m.finalize() + + # ----- Save per-member CSVs ----- + def _save_matrix(matrix, suffix): + out = {'date': dates_out} + for i in range(N): + out[f'member_{i:02d}'] = matrix[:, i] + df_out = pd.DataFrame(out) + out_path = out_dir / f'{cat_id}_production_per_member_{suffix}.csv' + df_out.to_csv(out_path, index=False) + return out_path + + # Streamflow outputs + q_path = out_dir / f'{cat_id}_production_per_member.csv' + out = {'date': dates_out} + for i in range(N): + out[f'member_{i:02d}'] = q_matrix[:, i] + pd.DataFrame(out).to_csv(q_path, index=False) + print(f"[per-member] saved {q_path}") + + # Perturbed forcings (so each member's actual inputs are recoverable) + precip_path = _save_matrix(precip_matrix, "precip") + pet_path = _save_matrix(pet_matrix, "pet") + print(f"[per-member] saved {precip_path}") + print(f"[per-member] saved {pet_path}") + + # Initial states (small, JSON is fine) + init_states_path = out_dir / f'{cat_id}_production_per_member_initial_states.json' + with open(init_states_path, 'w') as f: + json.dump({ + "catchment_id": cat_id, + "n_members": N, + "members": {f"member_{i:02d}": init_states_per_member[i] for i in range(N)}, + }, f, indent=2) + print(f"[per-member] saved {init_states_path}") + + print(f"[per-member] ensemble mean spread over test period: " + f"{float(q_matrix.std(axis=1).mean()):.5f} mm/h") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', required=True) + parser.add_argument('--forcing-dir', required=True) + parser.add_argument('--obs-dir', required=True) + parser.add_argument('--cfe-dir', required=True) + parser.add_argument('--config-file', required=True) + parser.add_argument('--param-bounds', required=True, + help='Kept for symmetry; not read in this script') + parser.add_argument('--out-dir', required=True) + parser.add_argument('--test-forcing-dir1', required=True) + parser.add_argument('--test-forcing-dir2', required=True) + parser.add_argument('--enkf-members', type=int, default=20) + parser.add_argument('--enkf-obs-error-std', type=float, default=0.05) + parser.add_argument('--no-vrugt-r', action='store_true') + parser.add_argument('--hardcoded-r', type=float, default=None, + help='Fix R to this constant value for all timesteps (overrides Vrugt formula)') + parser.add_argument('--direct-variance', action='store_true', default=False, + help='Use kriging variance column as R per timestep (R=σ²; overrides Vrugt formula)') + parser.add_argument('--vrugt-alpha', type=float, default=0.10) + parser.add_argument('--vrugt-scale', type=float, default=0.001) + parser.add_argument('--rng-seed', type=int, default=None) + parser.add_argument('--prod-script', default=None, + help='Path to calibrate_catchment_cfe_da_v2.py for importing ' + 'EnKFAssimilator. Defaults to the script located next to this file.') + args = parser.parse_args() + + # Locate the production script (default: same dir as this script) + if args.prod_script is None: + here = os.path.dirname(os.path.abspath(__file__)) + prod_script = os.path.join(here, 'calibrate_catchment_cfe_da_v2.py') + else: + prod_script = args.prod_script + if not os.path.exists(prod_script): + raise FileNotFoundError(f"Could not find production script at {prod_script}. " + f"Pass --prod-script to override.") + EnKFAssimilator = import_enkf_class(prod_script) + + # Dynamic CFE import (same idiom as production script) + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + + run(args, _bmi_cfe, EnKFAssimilator) + + +if __name__ == '__main__': + main() diff --git a/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/3_routing/route_det_f4.sh b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/3_routing/route_det_f4.sh new file mode 100644 index 00000000..daf5a675 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/3_routing/route_det_f4.sh @@ -0,0 +1,31 @@ +#!/usr/bin/env bash +# F4 dynamic variance direct (1 gauge holdout) — deterministic T-route routing. +# Routes per-catchment _test_results.csv (sim_mm_h column) through Muskingum-Cunge +# to gauge 03463300 and writes routed_Q_test.csv + KGE/NSE summary. +# +# Usage: +# bash route_det_f4.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +ROUTE_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/2_troute_routing/run_route.py + +F4_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder4_dynamic_variance_direct +GPKG=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv +KV_DIR=/mnt/disk2/1400_sites_helene/catchment_ts_no_03463300_dynamic_variance + +echo "[F4] Deterministic T-route routing..." +echo " da-dir : $F4_DIR" +echo " out-dir: $F4_DIR" +echo " kv-dir : $KV_DIR" + +$TROUTE "$ROUTE_SCRIPT" \ + --gpkg "$GPKG" \ + --da-dir "$F4_DIR" \ + --out-dir "$F4_DIR" \ + --usgs-csv "$USGS_CSV" \ + --kv-dir "$KV_DIR" + +echo "[F4] Deterministic routing done. Output: $F4_DIR/routed_Q_test.csv" diff --git a/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/3_routing/route_ensemble_f4.sh b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/3_routing/route_ensemble_f4.sh new file mode 100644 index 00000000..50724827 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/3_routing/route_ensemble_f4.sh @@ -0,0 +1,30 @@ +#!/usr/bin/env bash +# F4 Direct variance (1 gauge holdout) — crossed-ensemble T-route routing. +# Routes all 600 crossed-ensemble members through Muskingum-Cunge +# to gauge 03463300 and writes routed_crossed_ensemble.parquet. +# +# Must be run AFTER crossed ensemble parquets exist in F4_DIR. +# +# Usage: +# bash route_ensemble_f4.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +ROUTE_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/2_troute_routing/run_route_crossed_ensemble.py + +F4_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder4_dynamic_variance_direct +GPKG=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +echo "[F4] Crossed-ensemble T-route routing..." +echo " ensemble-dir: $F4_DIR" +echo " out-dir : $F4_DIR" + +$TROUTE "$ROUTE_SCRIPT" \ + --gpkg "$GPKG" \ + --ensemble-dir "$F4_DIR" \ + --out-dir "$F4_DIR" \ + --usgs-csv "$USGS_CSV" + +echo "[F4] Ensemble routing done. Output: $F4_DIR/routed_crossed_ensemble.parquet" diff --git a/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/3_routing/route_leadtime_f4.sh b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/3_routing/route_leadtime_f4.sh new file mode 100644 index 00000000..a2f666f3 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/3_routing/route_leadtime_f4.sh @@ -0,0 +1,31 @@ +#!/usr/bin/env bash +# F4 dynamic variance direct (1 gauge holdout) — route lead-time forecast CSVs. +# Routes all 21-catchment lead-time forecast CSVs through Muskingum-Cunge +# to gauge 03463300. Writes routed_leadtime_da_full.parquet and +# routed_leadtime_openloop_full.parquet for gauge-level error decay analysis. +# +# Run AFTER batch_run_leadtime_f4.sh finishes. +# +# Usage: +# bash route_leadtime_f4.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +ROUTE_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/2_troute_routing/run_route_leadtime_forecasts.py + +F4_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder4_dynamic_variance_direct +GPKG=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder4_dynamic_variance_direct_leadtime_routed + +echo "[F4] Routing lead-time forecasts through T-route..." +echo " forecast-dir: $F4_DIR" +echo " out-dir : $OUT_DIR" + +$TROUTE "$ROUTE_SCRIPT" \ + --gpkg "$GPKG" \ + --forecast-dir "$F4_DIR" \ + --out-dir "$OUT_DIR" + +echo "[F4] Lead-time routing done. Outputs in: $OUT_DIR" +ls "$OUT_DIR"/*.parquet 2>/dev/null || echo " (no parquets found)" diff --git a/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/plot_forecast_error_fixed_target.py b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/plot_forecast_error_fixed_target.py new file mode 100644 index 00000000..1460ed21 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/plot_forecast_error_fixed_target.py @@ -0,0 +1,217 @@ +""" +plot_forecast_error_fixed_target.py — 4a: error decay, fixed-target-time view. + +For each target verification time T in the Helene peak window: + Collect all forecasts that verify AT T: + issue_time = T - lead_hour*1h, for lead in 1..18 + error[lead] = ensemble_mean(q at T, initialized T-lead) - USGS_obs(T) + +This gives the correct operational picture: + - lead 1 = initialized 1 hr before T (DA just ran -> small error) + - lead 18 = initialized 18 hr before T (DA long ago -> error ~ open loop) + +Two panels: + Top : signed error (m³/s) vs lead hour, one curve per target time + Bot : same for open-loop +Plus a summary panel: mean across all target times, DA vs OL. + +Outputs: + /error_fixed_target_helene.png (per-target spaghetti, DA vs OL) + /error_fixed_target_mean.png (mean across targets, DA vs OL) +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.cm as cm + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast_hardcoded_r_routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +# Target verification times: hourly through the Helene peak window +TARGET_START = pd.Timestamp("2024-09-26 18:00:00") +TARGET_END = pd.Timestamp("2024-09-28 06:00:00") + + +def load_parquet(path): + df = pd.read_parquet(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + df["ens_mean"] = df[member_cols].mean(axis=1) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + return df[["issue_time", "lead_hour", "valid_time", "ens_mean"]] + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() + or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def build_fixed_target_errors(df, obs_series, target_times): + """ + For each target time T and each lead L (1..18): + issue_time = T - L hours + error = ensemble_mean at (issue_time, lead=L) - obs(T) + Returns dict: target_time -> {lead: error} + """ + # Index df by (issue_time, lead_hour) for fast lookup + df_idx = df.set_index(["issue_time", "lead_hour"])["ens_mean"] + + results = {} + for T in target_times: + obs_val = obs_series.get(T, np.nan) + if np.isnan(obs_val): + continue + curve = {} + for lead in range(1, 19): + t0 = T - pd.Timedelta(hours=lead) + try: + q_fc = df_idx.loc[(t0, lead)] + curve[lead] = float(q_fc) - obs_val + except KeyError: + curve[lead] = np.nan + results[T] = curve + return results + + +def plot_spaghetti(ax, error_dict, color_da, label_prefix, linestyle="-", lw=0.9, alpha=0.35): + """Plot one thin line per target time + thick mean.""" + leads = list(range(1, 19)) + all_curves = [] + target_times = sorted(error_dict.keys()) + cmap = cm.get_cmap("plasma", len(target_times)) + + for i, T in enumerate(target_times): + curve = [error_dict[T].get(l, np.nan) for l in leads] + ax.plot(leads, curve, + color=cmap(i), lw=lw, alpha=alpha, + linestyle=linestyle, zorder=2) + all_curves.append(curve) + + if all_curves: + mean_curve = np.nanmean(np.array(all_curves), axis=0) + ax.plot(leads, mean_curve, + color=color_da, lw=2.8, alpha=0.95, + linestyle=linestyle, zorder=5, + label=f"{label_prefix} — mean across targets") + return all_curves + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--route-dir", default=DEFAULT_ROUTE_DIR) + parser.add_argument("--out-dir", default=None) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--da-name", default="routed_leadtime_da_full.parquet") + parser.add_argument("--ol-name", default="routed_leadtime_openloop_full.parquet") + parser.add_argument("--target-start", default=str(TARGET_START)) + parser.add_argument("--target-end", default=str(TARGET_END)) + args = parser.parse_args() + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + print("Loading parquets...") + da = load_parquet(os.path.join(args.route_dir, args.da_name)) + ol = load_parquet(os.path.join(args.route_dir, args.ol_name)) + obs = load_usgs(args.usgs_csv) + + target_times = pd.date_range(args.target_start, args.target_end, freq="1h") + print(f" Target verification times: {len(target_times)} " + f"({target_times[0]} → {target_times[-1]})") + + print("Building fixed-target error tables...") + da_errors = build_fixed_target_errors(da, obs, target_times) + ol_errors = build_fixed_target_errors(ol, obs, target_times) + print(f" Targets with obs: DA={len(da_errors)} OL={len(ol_errors)}") + + leads = list(range(1, 19)) + + # ---- Spaghetti: per-target-time curves, DA vs OL ---- + fig, (ax_da, ax_ol) = plt.subplots(2, 1, figsize=(13, 10), sharex=True, sharey=True) + + ax_da.axhline(0, color="black", lw=0.8, linestyle="--", alpha=0.4) + plot_spaghetti(ax_da, da_errors, "tab:purple", "DA", linestyle="-") + ax_da.set_ylabel("Error: forecast − USGS obs (m³/s)", fontsize=10) + ax_da.set_title("DA — error at each lead for fixed target times (Helene peak window)", fontsize=11) + ax_da.grid(True, alpha=0.2) + ax_da.legend(fontsize=9) + + ax_ol.axhline(0, color="black", lw=0.8, linestyle="--", alpha=0.4) + plot_spaghetti(ax_ol, ol_errors, "tab:gray", "Open-loop", linestyle="--") + ax_ol.set_ylabel("Error: forecast − USGS obs (m³/s)", fontsize=10) + ax_ol.set_xlabel("Forecast lead hour (hours before target)", fontsize=11) + ax_ol.set_title("Open-loop — error at each lead for fixed target times", fontsize=11) + ax_ol.set_xticks(leads) + ax_ol.grid(True, alpha=0.2) + ax_ol.legend(fontsize=9) + + fig.suptitle( + "Forecast error vs lead time — fixed verification time, Helene peak window\n" + f"Target times: {TARGET_START.strftime('%b %d %H UTC')} → " + f"{TARGET_END.strftime('%b %d %H UTC')} | USGS 03463300", + fontsize=12, y=1.01, + ) + plt.tight_layout() + out1 = os.path.join(out_dir, "error_fixed_target_helene.png") + plt.savefig(out1, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + # ---- Summary: mean across all target times, DA vs OL overlaid ---- + fig, ax = plt.subplots(figsize=(13, 6)) + ax.axhline(0, color="black", lw=0.8, linestyle="--", alpha=0.4) + + def mean_curve(error_dict): + arr = np.array([ + [error_dict[T].get(l, np.nan) for l in leads] + for T in sorted(error_dict.keys()) + ]) + return np.nanmean(arr, axis=0), np.nanstd(arr, axis=0) + + da_mean, da_std = mean_curve(da_errors) + ol_mean, ol_std = mean_curve(ol_errors) + + ax.fill_between(leads, da_mean - da_std, da_mean + da_std, + color="tab:purple", alpha=0.15, zorder=2) + ax.fill_between(leads, ol_mean - ol_std, ol_mean + ol_std, + color="tab:gray", alpha=0.15, zorder=2) + ax.plot(leads, da_mean, color="tab:purple", lw=2.6, marker="o", + zorder=5, label="DA — mean error (±1 std shaded)") + ax.plot(leads, ol_mean, color="tab:gray", lw=2.6, marker="s", + linestyle="--", zorder=5, label="Open-loop — mean error (±1 std shaded)") + + ax.set_xlabel("Forecast lead hour (hours before target verification time)", fontsize=11) + ax.set_ylabel("Mean error: forecast − USGS obs (m³/s)", fontsize=11) + ax.set_xticks(leads) + ax.set_title( + "Mean forecast error vs lead time — fixed verification time, Helene peak window\n" + f"Target: {TARGET_START.strftime('%b %d %H UTC')} → " + f"{TARGET_END.strftime('%b %d %H UTC')} | " + "Lead 1 = init 1 hr before target | Lead 18 = init 18 hr before target", + fontsize=11, + ) + ax.legend(fontsize=10, loc="lower right", frameon=True, framealpha=0.92) + ax.grid(True, alpha=0.25) + plt.tight_layout() + out2 = os.path.join(out_dir, "error_fixed_target_mean.png") + plt.savefig(out2, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/plot_forecast_error_per_init.py b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/plot_forecast_error_per_init.py new file mode 100644 index 00000000..691245d0 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/plot_forecast_error_per_init.py @@ -0,0 +1,203 @@ +""" +plot_forecast_error_per_init.py + +Error decay by initialization time for the Helene window (Sep 24-28 2024). + +For each initialization time t0 in the Helene window: + error[lead] = ensemble_mean(q_gauge_m3s at t0+lead) - USGS_obs(t0+lead) + +Plotted as: + DA : thin colored lines (one per init time, colored by date) + thick mean across all + OL : thin gray dashed lines + thick gray dashed mean + +x-axis: forecast lead hour (1 -> 18) +y-axis: signed error (m³/s), positive = forecast too high + +The expected signal: DA error is small at lead 1 (just assimilated), grows +and converges toward the OL error curve by lead 18. + +Inputs: + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + +Obs: + /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +Outputs: + /forecast_error_per_init_helene.png (signed error) + /forecast_rmse_per_init_helene.png (absolute error / RMSE per init) +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.patches as mpatches + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast_hardcoded_r_routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +# Initialization times to show — Helene window +HELENE_INIT_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_INIT_END = pd.Timestamp("2024-09-28 23:00:00") + +# One color per init date +DATE_COLORS = { + "2024-09-24": "#1f77b4", + "2024-09-25": "#ff7f0e", + "2024-09-26": "#2ca02c", + "2024-09-27": "#d62728", + "2024-09-28": "#9467bd", +} + + +def load_parquet(path): + df = pd.read_parquet(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + df["ens_mean"] = df[member_cols].mean(axis=1) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + return df[["issue_time", "lead_hour", "valid_time", "ens_mean"]] + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() + or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def build_error_table(df, obs_series, init_start, init_end): + """Return DataFrame: issue_time, lead_hour, ens_mean, obs, error.""" + df = df[(df["issue_time"] >= init_start) & (df["issue_time"] <= init_end)].copy() + df["obs"] = df["valid_time"].map(obs_series) + df["error"] = df["ens_mean"] - df["obs"] + return df.dropna(subset=["obs", "error"]) + + +def plot_error(ax, err_df, color, alpha_thin, lw_thin, linestyle, label_prefix): + """Plot individual init-time error curves + thick mean curve.""" + leads = sorted(err_df["lead_hour"].unique()) + all_curves = [] + + for t0, grp in err_df.groupby("issue_time"): + date_str = str(pd.Timestamp(t0).date()) + c = DATE_COLORS.get(date_str, color) + grp_sorted = grp.sort_values("lead_hour") + # Align to leads grid — some may be missing + curve = grp_sorted.set_index("lead_hour")["error"].reindex(leads).values + ax.plot(leads, curve, + color=c, lw=lw_thin, alpha=alpha_thin, + linestyle=linestyle, zorder=2) + all_curves.append(curve) + + if all_curves: + mean_curve = np.nanmean(np.array(all_curves), axis=0) + ax.plot(leads, mean_curve, + color=color, lw=2.8, alpha=0.95, + linestyle=linestyle, zorder=5, + label=f"{label_prefix} — mean across all init times") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--route-dir", default=DEFAULT_ROUTE_DIR) + parser.add_argument("--out-dir", default=None) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--da-name", default="routed_leadtime_da_full.parquet") + parser.add_argument("--ol-name", default="routed_leadtime_openloop_full.parquet") + args = parser.parse_args() + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + print("Loading parquets...") + da = load_parquet(os.path.join(args.route_dir, args.da_name)) + ol = load_parquet(os.path.join(args.route_dir, args.ol_name)) + obs = load_usgs(args.usgs_csv) + print(f" DA issue_times: {da['issue_time'].nunique()} " + f"OL issue_times: {ol['issue_time'].nunique()}") + + da_err = build_error_table(da, obs, HELENE_INIT_START, HELENE_INIT_END) + ol_err = build_error_table(ol, obs, HELENE_INIT_START, HELENE_INIT_END) + print(f" DA init times in Helene window: {da_err['issue_time'].nunique()}") + + leads = sorted(da_err["lead_hour"].unique()) + + # ---- Signed error plot ---- + fig, ax = plt.subplots(figsize=(12, 6)) + ax.axhline(0, color="black", lw=0.8, linestyle="--", alpha=0.5, zorder=1) + + plot_error(ax, ol_err, color="tab:gray", alpha_thin=0.12, lw_thin=0.7, + linestyle="--", label_prefix="Open-loop") + plot_error(ax, da_err, color="tab:purple", alpha_thin=0.18, lw_thin=0.8, + linestyle="-", label_prefix="DA") + + # Date-color legend patches + patches = [mpatches.Patch(color=c, label=f"Init {d}") + for d, c in DATE_COLORS.items()] + handles, labels = ax.get_legend_handles_labels() + ax.legend(handles=handles + patches, + fontsize=8, loc="upper left", frameon=True, framealpha=0.9, ncol=2) + + ax.set_xlabel("Forecast lead hour", fontsize=11) + ax.set_ylabel("Error: forecast − USGS obs (m³/s)", fontsize=11) + ax.set_xticks(np.arange(1, 19)) + ax.set_title( + "Forecast error vs lead time — per initialization time, Helene window\n" + "DA (purple solid) vs Open-loop (gray dashed) | Sep 24–28 2024 | USGS 03463300", + fontsize=11, + ) + ax.grid(True, alpha=0.25, lw=0.4) + plt.tight_layout() + out1 = os.path.join(out_dir, "forecast_error_per_init_helene.png") + plt.savefig(out1, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + # ---- Absolute error (|error|) averaged per lead — cleaner summary ---- + fig, ax = plt.subplots(figsize=(12, 6)) + + def mean_abs_error_by_lead(err_df): + return err_df.groupby("lead_hour")["error"].apply( + lambda x: float(np.nanmean(np.abs(x))) + ) + + da_mae = mean_abs_error_by_lead(da_err) + ol_mae = mean_abs_error_by_lead(ol_err) + + ax.plot(da_mae.index, da_mae.values, + color="tab:purple", lw=2.4, marker="o", label="DA — mean |error|") + ax.plot(ol_mae.index, ol_mae.values, + color="tab:gray", lw=2.4, marker="s", linestyle="--", + label="Open-loop — mean |error|") + + ax.set_xlabel("Forecast lead hour", fontsize=11) + ax.set_ylabel("Mean |error| vs USGS obs (m³/s)", fontsize=11) + ax.set_xticks(np.arange(1, 19)) + ax.set_title( + "Mean absolute forecast error vs lead time — Helene window\n" + "DA (purple) vs Open-loop (gray) | Sep 24–28 2024 | USGS 03463300", + fontsize=11, + ) + ax.legend(fontsize=10, loc="upper left", frameon=True, framealpha=0.9) + ax.grid(True, alpha=0.25, lw=0.4) + plt.tight_layout() + out2 = os.path.join(out_dir, "forecast_mae_per_lead_helene.png") + plt.savefig(out2, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/plot_lead_time_decay.py b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/plot_lead_time_decay.py new file mode 100644 index 00000000..2b33fa82 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/plot_lead_time_decay.py @@ -0,0 +1,160 @@ +""" +Catchment-level error-vs-lead-time decay curve + ensemble spread by lead. + +Two-panel figure for ONE catchment: + TOP — RMSE of ensemble-mean forecast at each lead hour (1..18) vs the + catchment's kriging observation. DA solid, open-loop dashed, with + a shaded band showing the min/max of per-member RMSE. + BOTTOM — Mean ensemble spread (std-dev across 20 members, averaged across + all issue times) at each lead hour. Tells you whether forcing + perturbation alone keeps the forecast ensemble diverse during the + 18-hour free-run — useful for inspecting individual members + without needing the full spaghetti view. + +This is the catchment-level analog of the gauge-level decay curve. +Routing to the gauge is a separate post-step (route_lead_time_forecasts.py ++ a gauge-level decay script); this script lets us look at the catchment-level +signal without T-route in the loop. + +Inputs (from run_lead_time_forecast_sweep.py): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv +Obs (from production EnKF run): + //_test_results.csv + +Output: + //_lead_time_decay.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" + +DEFAULT_LEADTIME_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast" +DEFAULT_DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" + +# Module-level globals set by main() so the helpers see them +CAT = "cat-1016300" +LEADTIME_DIR = DEFAULT_LEADTIME_DIR +DA_DIR = DEFAULT_DA_DIR +OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay.png") + + +def load_forecasts(path): + """Return (df, member_cols).""" + if not os.path.exists(path): + raise FileNotFoundError(path) + df = pd.read_csv(path, parse_dates=['issue_time', 'valid_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + return df, member_cols + + +def load_obs(): + obs_path = os.path.join(DA_DIR, CAT, f"{CAT}_test_results.csv") + df = pd.read_csv(obs_path, parse_dates=['date']) + return df.set_index('date')['obs_mm_h'] + + +def metrics_by_lead(df, member_cols, obs_series): + """For each lead hour, return (rmse_mean, rmse_min_member, rmse_max_member, + mean_ensemble_std).""" + df = df.copy() + df['obs'] = df['valid_time'].map(obs_series) + df = df.dropna(subset=['obs']) + df['ens_mean'] = df[member_cols].mean(axis=1) + df['ens_std'] = df[member_cols].std(axis=1) + + leads = sorted(df['lead_hour'].unique()) + rmse_mean = np.full(len(leads), np.nan) + rmse_min_member = np.full(len(leads), np.nan) + rmse_max_member = np.full(len(leads), np.nan) + mean_ens_std = np.full(len(leads), np.nan) + for i, L in enumerate(leads): + sub = df[df['lead_hour'] == L] + if len(sub) == 0: + continue + err = sub['ens_mean'].values - sub['obs'].values + rmse_mean[i] = float(np.sqrt(np.mean(err ** 2))) + per_member_rmse = np.array([ + np.sqrt(np.mean((sub[c].values - sub['obs'].values) ** 2)) + for c in member_cols + ]) + rmse_min_member[i] = float(np.nanmin(per_member_rmse)) + rmse_max_member[i] = float(np.nanmax(per_member_rmse)) + mean_ens_std[i] = float(np.nanmean(sub['ens_std'].values)) + return (np.asarray(leads), rmse_mean, rmse_min_member, + rmse_max_member, mean_ens_std) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', default="cat-1016300") + parser.add_argument('--leadtime-dir', default=DEFAULT_LEADTIME_DIR, + help='Dir holding /_lead_time_forecasts_{da,openloop}.csv') + parser.add_argument('--da-dir', default=DEFAULT_DA_DIR, + help='Dir holding /_test_results.csv (for obs)') + args = parser.parse_args() + + global CAT, LEADTIME_DIR, DA_DIR, OUT_PNG + CAT = args.cat_id + LEADTIME_DIR = args.leadtime_dir + DA_DIR = args.da_dir + OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay.png") + + da_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_da.csv") + ol_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_openloop.csv") + df_da, m_cols = load_forecasts(da_path) + df_ol, _ = load_forecasts(ol_path) + obs_series = load_obs() + + leads_da, rmse_da, rmse_da_lo, rmse_da_hi, std_da = metrics_by_lead( + df_da, m_cols, obs_series) + leads_ol, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = metrics_by_lead( + df_ol, m_cols, obs_series) + + fig, (ax_top, ax_bot) = plt.subplots(2, 1, figsize=(10, 10), sharex=True, + gridspec_kw={"height_ratios": [1.4, 1.0]}) + + # ----- TOP: RMSE decay curve ----- + ax_top.fill_between(leads_da, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2, + label="DA — member spread (min/max RMSE)") + ax_top.fill_between(leads_ol, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2, + label="Open-loop — member spread") + ax_top.plot(leads_da, rmse_da, color=DA_COLOR, lw=2.4, marker='o', + zorder=4, label="DA on (ensemble-mean RMSE)") + ax_top.plot(leads_ol, rmse_ol, color=OL_COLOR, lw=2.4, marker='s', + linestyle='--', zorder=4, label="Open-loop (ensemble-mean RMSE)") + ax_top.set_ylabel("RMSE vs Qkrig obs (mm/h)", fontsize=11) + ax_top.set_title( + f"Forecast lead-time error decay — {CAT}\n" + f"Issue times pooled across test period (Oct 2023 – Oct 2024)", + fontsize=12, + ) + ax_top.grid(True, alpha=0.25) + ax_top.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + # ----- BOTTOM: ensemble spread by lead hour ----- + ax_bot.plot(leads_da, std_da, color=DA_COLOR, lw=2.2, marker='o', + label="DA — mean ensemble std") + ax_bot.plot(leads_ol, std_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', label="Open-loop — mean ensemble std") + ax_bot.set_xlabel("Forecast lead time (hours)", fontsize=11) + ax_bot.set_ylabel("Mean ensemble std (mm/h)", fontsize=11) + ax_bot.set_xticks(np.arange(1, 19)) + ax_bot.grid(True, alpha=0.25) + ax_bot.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + plt.tight_layout() + plt.savefig(OUT_PNG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/plot_lead_time_decay_by_regime.py b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/plot_lead_time_decay_by_regime.py new file mode 100644 index 00000000..ecf10408 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/plot_lead_time_decay_by_regime.py @@ -0,0 +1,200 @@ +""" +Lead-time decay curve, split by flow regime at issue time. + +The pooled lead-time curve (plot_lead_time_decay.py) showed DA losing to +open-loop across most lead hours, but the test period is ~99% low-flow. +This script splits the same forecast CSVs by the flow regime at the issue +time t0, so we can see whether DA helps when it matters (storms / Helene) +and hurts when it doesn't (low flow). + +Three regimes are partitioned on obs(t0) (kriging obs at the issue time): + + helene issue_time ∈ [2024-09-24, 2024-09-28] (the 5-day Helene window) + storm obs(t0) > 0.5 mm/h (high-flow issue times) + low_flow obs(t0) < 0.1 mm/h (low-flow issue times) + +Same 2-panel layout per regime: RMSE vs lead (top), ensemble spread vs lead +(bottom). Three regime columns × 2 metric rows in one figure. + +Inputs (from run_lead_time_forecast_sweep.py — no re-run needed): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv +Obs (from production EnKF): + //_test_results.csv + +Output: + //_lead_time_decay_by_regime.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" + +DEFAULT_LEADTIME_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast" +DEFAULT_DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" + +# Module-level globals set by main() so helpers see them +CAT = "cat-1016300" +LEADTIME_DIR = DEFAULT_LEADTIME_DIR +DA_DIR = DEFAULT_DA_DIR +OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay_by_regime.png") + +# Regime definitions +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-28 23:00:00") +STORM_OBS_THRESHOLD = 0.5 # mm/h +LOWFLOW_OBS_THRESHOLD = 0.1 # mm/h + + +def load_forecasts(path): + if not os.path.exists(path): + raise FileNotFoundError(path) + df = pd.read_csv(path, parse_dates=['issue_time', 'valid_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + return df, member_cols + + +def load_obs(): + obs_path = os.path.join(DA_DIR, CAT, f"{CAT}_test_results.csv") + df = pd.read_csv(obs_path, parse_dates=['date']) + return df.set_index('date')['obs_mm_h'] + + +def regime_mask(issue_times, obs_at_issue, regime): + """Boolean mask over issue_times for the named regime.""" + if regime == "helene": + return (issue_times >= HELENE_START) & (issue_times <= HELENE_END) + if regime == "storm": + return obs_at_issue > STORM_OBS_THRESHOLD + if regime == "low_flow": + return obs_at_issue < LOWFLOW_OBS_THRESHOLD + raise ValueError(regime) + + +def metrics_by_lead(df, member_cols, obs_series, issue_mask): + """Compute (leads, rmse_mean, rmse_min, rmse_max, mean_std) restricted + to issue times where issue_mask is True.""" + df = df.copy() + df = df[df['issue_time'].isin(issue_mask)] + if len(df) == 0: + return None + df['obs'] = df['valid_time'].map(obs_series) + df = df.dropna(subset=['obs']) + if len(df) == 0: + return None + df['ens_mean'] = df[member_cols].mean(axis=1) + df['ens_std'] = df[member_cols].std(axis=1) + + leads = sorted(df['lead_hour'].unique()) + rmse_mean = np.full(len(leads), np.nan) + rmse_lo = np.full(len(leads), np.nan) + rmse_hi = np.full(len(leads), np.nan) + std_mean = np.full(len(leads), np.nan) + for i, L in enumerate(leads): + sub = df[df['lead_hour'] == L] + if len(sub) == 0: + continue + err = sub['ens_mean'].values - sub['obs'].values + rmse_mean[i] = float(np.sqrt(np.mean(err ** 2))) + per_member = np.array([ + np.sqrt(np.mean((sub[c].values - sub['obs'].values) ** 2)) + for c in member_cols + ]) + rmse_lo[i] = float(np.nanmin(per_member)) + rmse_hi[i] = float(np.nanmax(per_member)) + std_mean[i] = float(np.nanmean(sub['ens_std'].values)) + return np.asarray(leads), rmse_mean, rmse_lo, rmse_hi, std_mean + + +def plot_pair(ax_rmse, ax_std, da_metrics, ol_metrics, + regime_label, n_issue_times): + if da_metrics is None or ol_metrics is None: + ax_rmse.set_title(f"{regime_label}\n(no issue times match)", fontsize=11) + ax_rmse.axis('off'); ax_std.axis('off') + return + leads, rmse_da, rmse_da_lo, rmse_da_hi, std_da = da_metrics + _, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = ol_metrics + + # RMSE panel + ax_rmse.fill_between(leads, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2) + ax_rmse.fill_between(leads, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2) + ax_rmse.plot(leads, rmse_da, color=DA_COLOR, lw=2.2, marker='o', + zorder=4, label="DA on") + ax_rmse.plot(leads, rmse_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', zorder=4, label="Open-loop") + ax_rmse.set_title(f"{regime_label}\n({n_issue_times} issue times)", fontsize=11) + ax_rmse.set_ylabel("RMSE (mm/h)", fontsize=10) + ax_rmse.grid(True, alpha=0.25) + ax_rmse.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + # Spread panel + ax_std.plot(leads, std_da, color=DA_COLOR, lw=2.0, marker='o', label="DA") + ax_std.plot(leads, std_ol, color=OL_COLOR, lw=2.0, marker='s', + linestyle='--', label="Open-loop") + ax_std.set_xlabel("Forecast lead time (hours)", fontsize=10) + ax_std.set_ylabel("Mean ensemble std (mm/h)", fontsize=10) + ax_std.set_xticks(np.arange(1, 19)) + ax_std.grid(True, alpha=0.25) + ax_std.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', default="cat-1016300") + parser.add_argument('--leadtime-dir', default=DEFAULT_LEADTIME_DIR, + help='Dir holding /_lead_time_forecasts_{da,openloop}.csv') + parser.add_argument('--da-dir', default=DEFAULT_DA_DIR, + help='Dir holding /_test_results.csv (for obs)') + args = parser.parse_args() + + global CAT, LEADTIME_DIR, DA_DIR, OUT_PNG + CAT = args.cat_id + LEADTIME_DIR = args.leadtime_dir + DA_DIR = args.da_dir + OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay_by_regime.png") + + da_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_da.csv") + ol_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_openloop.csv") + df_da, m_cols = load_forecasts(da_path) + df_ol, _ = load_forecasts(ol_path) + obs_series = load_obs() + + # Unique issue times across the run, with obs at each + issue_times = pd.to_datetime(sorted(df_da['issue_time'].unique())) + obs_at_issue = pd.Series(issue_times, index=issue_times).map(obs_series) + + regimes = [ + ("helene", "Helene window (Sep 24–28 2024)"), + ("storm", f"Storm hours (obs(t0) > {STORM_OBS_THRESHOLD} mm/h)"), + ("low_flow", f"Low flow (obs(t0) < {LOWFLOW_OBS_THRESHOLD} mm/h)"), + ] + + fig, axes = plt.subplots(2, 3, figsize=(18, 10), + gridspec_kw={"height_ratios": [1.4, 1.0]}) + for col, (key, label) in enumerate(regimes): + mask = regime_mask(issue_times, obs_at_issue, key) + kept = issue_times[mask] + kept_set = set(pd.to_datetime(kept)) + da_m = metrics_by_lead(df_da, m_cols, obs_series, kept_set) + ol_m = metrics_by_lead(df_ol, m_cols, obs_series, kept_set) + plot_pair(axes[0, col], axes[1, col], da_m, ol_m, label, len(kept)) + + fig.suptitle( + f"Lead-time error decay by flow regime — {CAT}\n" + f"Same forecast CSVs as the pooled view, partitioned on obs(t0).", + fontsize=13, y=0.995, + ) + plt.tight_layout(rect=(0, 0, 1, 0.97)) + plt.savefig(OUT_PNG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/plot_lead_time_decay_gauge.py b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/plot_lead_time_decay_gauge.py new file mode 100644 index 00000000..71ce54fc --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/plot_lead_time_decay_gauge.py @@ -0,0 +1,311 @@ +""" +Gauge-level lead-time forecast decay curve. + +Routes 21-catchment forecast CSVs through T-route to USGS gauge 03463300 and +plots RMSE vs lead time (DA solid, open-loop dashed) against USGS observed Q. +Same 2-panel layout as the catchment-level plot, plus a regime split for the +Helene window. + +The routing step is done separately by route_lead_time_forecasts.py; +this script just reads the resulting parquets and produces the plot. + +Inputs: + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + Each: one row per (issue_time, lead_hour); member columns hold q_gauge_m3s. + (Long-format with explicit `member` and `q_gauge_m3s` columns also supported.) +Obs: + /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + USGS hourly Q in m³/s at gauge 03463300 (South Toe River near Celo, NC). + +Output: + /lead_time_decay_gauge_pooled.png + /lead_time_decay_gauge_by_regime.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" +OBS_COLOR = "black" + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/leadtime_troute_routing" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +# Watershed area at gauge 03463300 (sum of 21 catchments per GPKG hydrofabric) +WATERSHED_AREA_KM2 = 113.18 +# 1 mm/h depth × 113.18 km² = 113.18e3 m³/h = 31.439 m³/s +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-28 23:00:00") +STORM_OBS_THRESHOLD_M3S = 20.0 # gauge-level storm threshold (was 50 — too high + # for this 113 km² basin; only 2 issue times qualified) +LOWFLOW_OBS_THRESHOLD_M3S = 5.0 # gauge-level low-flow threshold + +USGS_HELENE_PEAK_M3S = 1886.0 # reference Sep 27 14:00 + + +def load_parquet_long(path): + """Load a routed parquet and normalize to long format: + columns = issue_time, lead_hour, member, q_gauge_m3s. + + Handles both wide format (member_00..member_19 columns) and long format + (explicit `member` + `q_gauge_m3s` columns). Auto-detects. + """ + if not os.path.exists(path): + raise FileNotFoundError(path) + df = pd.read_parquet(path) + + # Long format detection + if 'q_gauge_m3s' in df.columns and 'member' in df.columns: + df['issue_time'] = pd.to_datetime(df['issue_time']) + return df[['issue_time', 'lead_hour', 'member', 'q_gauge_m3s']] + + # Some variants might call it 'q_m3s' or similar + q_col_candidates = [c for c in df.columns if c.lower() in + ('q_gauge_m3s', 'q_m3s', 'q_gauge', 'q')] + if 'member' in df.columns and q_col_candidates: + qc = q_col_candidates[0] + df['issue_time'] = pd.to_datetime(df['issue_time']) + out = df[['issue_time', 'lead_hour', 'member', qc]].copy() + out = out.rename(columns={qc: 'q_gauge_m3s'}) + return out + + # Wide format: member_00..member_19 columns + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + if not member_cols: + raise ValueError(f"Couldn't identify member columns in {path}. " + f"Columns present: {df.columns.tolist()}") + df['issue_time'] = pd.to_datetime(df['issue_time']) + keep = ['issue_time', 'lead_hour'] + if 'valid_time' in df.columns: + keep.append('valid_time') + long = df[keep + member_cols].melt( + id_vars=keep, + value_vars=member_cols, + var_name='member', + value_name='q_gauge_m3s', + ) + return long + + +def load_usgs_obs(usgs_csv): + """Load USGS gauge obs and ensure m³/s. + + The CSV at gauge 03463300 has the column `QObs(mm/h)` (catchment-averaged + depth, not gauge discharge in m³/s). If the column name contains 'mm', we + convert mm/h → m³/s by multiplying by the watershed area factor: + m³/s = mm/h × WATERSHED_AREA_KM2 × 1000 / 3600 + Verified: 59.978 mm/h × 113.18 km² = 1886 m³/s (USGS Helene peak). + """ + df = pd.read_csv(usgs_csv) + date_candidates = [c for c in df.columns if c.lower() in + ('datetime', 'date', 'time', 'timestamp')] + q_candidates = [c for c in df.columns if 'q' in c.lower() + or 'flow' in c.lower() or 'discharge' in c.lower()] + if not date_candidates or not q_candidates: + raise ValueError(f"Could not identify date/Q columns in {usgs_csv}. " + f"Columns: {df.columns.tolist()}") + dc, qc = date_candidates[0], q_candidates[0] + df[dc] = pd.to_datetime(df[dc]) + series = df.set_index(dc)[qc].astype(float) + if 'mm' in qc.lower(): + series = series * MM_H_TO_M3_S + print(f" Converted obs from mm/h (column '{qc}') to m³/s " + f"using area = {WATERSHED_AREA_KM2} km² (× {MM_H_TO_M3_S:.4f})") + else: + print(f" Loaded obs as-is from column '{qc}' (assumed m³/s)") + return series + + +def metrics_by_lead(df_long, obs_series, issue_mask=None): + """Compute per-lead-hour: ensemble-mean RMSE, per-member min/max RMSE, + mean ensemble std. df_long must have issue_time, lead_hour, member, q_gauge_m3s. + """ + df = df_long.copy() + if issue_mask is not None: + df = df[df['issue_time'].isin(issue_mask)] + if len(df) == 0: + return None + df['valid_time'] = df['issue_time'] + pd.to_timedelta(df['lead_hour'], unit='h') + df['obs'] = df['valid_time'].map(obs_series) + df = df.dropna(subset=['obs']) + if len(df) == 0: + return None + + # Ensemble mean per (issue_time, lead_hour) + grouped = df.groupby(['issue_time', 'lead_hour']) + ens_mean = grouped['q_gauge_m3s'].mean().reset_index(name='ens_mean') + ens_std = grouped['q_gauge_m3s'].std().reset_index(name='ens_std') + obs_per = grouped['obs'].first().reset_index(name='obs') + panel = ens_mean.merge(ens_std, on=['issue_time', 'lead_hour']).merge( + obs_per, on=['issue_time', 'lead_hour']) + + leads = sorted(panel['lead_hour'].unique()) + rmse_mean = np.full(len(leads), np.nan) + rmse_lo = np.full(len(leads), np.nan) + rmse_hi = np.full(len(leads), np.nan) + std_mean = np.full(len(leads), np.nan) + for i, L in enumerate(leads): + psub = panel[panel['lead_hour'] == L] + if len(psub) == 0: + continue + err = psub['ens_mean'].values - psub['obs'].values + rmse_mean[i] = float(np.sqrt(np.mean(err ** 2))) + # Per-member RMSE — recompute from the underlying long frame + dsub = df[df['lead_hour'] == L] + member_rmses = [] + for m, g in dsub.groupby('member'): + err_m = g['q_gauge_m3s'].values - g['obs'].values + if len(err_m) >= 1: + member_rmses.append(float(np.sqrt(np.mean(err_m ** 2)))) + if member_rmses: + rmse_lo[i] = float(np.nanmin(member_rmses)) + rmse_hi[i] = float(np.nanmax(member_rmses)) + std_mean[i] = float(np.nanmean(psub['ens_std'].values)) + return np.asarray(leads), rmse_mean, rmse_lo, rmse_hi, std_mean + + +def plot_two_panel(out_path, leads_da, da_metrics, leads_ol, ol_metrics, title): + fig, (ax_top, ax_bot) = plt.subplots(2, 1, figsize=(10, 10), sharex=True, + gridspec_kw={"height_ratios": [1.4, 1.0]}) + leads, rmse_da, rmse_da_lo, rmse_da_hi, std_da = da_metrics + _, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = ol_metrics + + ax_top.fill_between(leads, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2, + label="DA — member spread (min/max RMSE)") + ax_top.fill_between(leads, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2, + label="Open-loop — member spread") + ax_top.plot(leads, rmse_da, color=DA_COLOR, lw=2.4, marker='o', + zorder=4, label="DA on (ensemble-mean RMSE)") + ax_top.plot(leads, rmse_ol, color=OL_COLOR, lw=2.4, marker='s', + linestyle='--', zorder=4, label="Open-loop (ensemble-mean RMSE)") + ax_top.set_ylabel("RMSE vs USGS obs (m³/s)", fontsize=11) + ax_top.set_title(title, fontsize=12) + ax_top.grid(True, alpha=0.25) + ax_top.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + ax_bot.plot(leads, std_da, color=DA_COLOR, lw=2.2, marker='o', + label="DA — mean ensemble std") + ax_bot.plot(leads, std_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', label="Open-loop — mean ensemble std") + ax_bot.set_xlabel("Forecast lead time (hours)", fontsize=11) + ax_bot.set_ylabel("Mean ensemble std (m³/s)", fontsize=11) + ax_bot.set_xticks(np.arange(1, 19)) + ax_bot.grid(True, alpha=0.25) + ax_bot.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + plt.tight_layout() + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + +def plot_three_panel_regime(out_path, regime_metrics): + fig, axes = plt.subplots(2, 3, figsize=(18, 10), + gridspec_kw={"height_ratios": [1.4, 1.0]}) + for col, (label, n_issue, da_m, ol_m) in enumerate(regime_metrics): + ax_rmse = axes[0, col] + ax_std = axes[1, col] + if da_m is None or ol_m is None: + ax_rmse.set_title(f"{label}\n(no issue times match)", fontsize=11) + ax_rmse.axis('off'); ax_std.axis('off') + continue + leads, rmse_da, rmse_da_lo, rmse_da_hi, std_da = da_m + _, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = ol_m + ax_rmse.fill_between(leads, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2) + ax_rmse.fill_between(leads, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2) + ax_rmse.plot(leads, rmse_da, color=DA_COLOR, lw=2.2, marker='o', + zorder=4, label="DA on") + ax_rmse.plot(leads, rmse_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', zorder=4, label="Open-loop") + ax_rmse.set_title(f"{label}\n({n_issue} issue times)", fontsize=11) + ax_rmse.set_ylabel("RMSE (m³/s)", fontsize=10) + ax_rmse.grid(True, alpha=0.25) + ax_rmse.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + ax_std.plot(leads, std_da, color=DA_COLOR, lw=2.0, marker='o', label="DA") + ax_std.plot(leads, std_ol, color=OL_COLOR, lw=2.0, marker='s', + linestyle='--', label="Open-loop") + ax_std.set_xlabel("Forecast lead time (hours)", fontsize=10) + ax_std.set_ylabel("Mean ensemble std (m³/s)", fontsize=10) + ax_std.set_xticks(np.arange(1, 19)) + ax_std.grid(True, alpha=0.25) + ax_std.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + fig.suptitle("Gauge-level lead-time error decay by regime — USGS 03463300", + fontsize=13, y=0.995) + plt.tight_layout(rect=(0, 0, 1, 0.97)) + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--route-dir', default=DEFAULT_ROUTE_DIR, + help='Dir holding routed_leadtime_{da,openloop}_full.parquet') + parser.add_argument('--out-dir', default=None, + help='Where to write the output PNGs. Defaults to --route-dir ' + '(which may not be writable if owned by another user — ' + 'pass an explicit path then).') + parser.add_argument('--usgs-csv', default=DEFAULT_USGS_CSV) + parser.add_argument('--da-name', default='routed_leadtime_da_full.parquet') + parser.add_argument('--ol-name', default='routed_leadtime_openloop_full.parquet') + args = parser.parse_args() + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + da_path = os.path.join(args.route_dir, args.da_name) + ol_path = os.path.join(args.route_dir, args.ol_name) + df_da = load_parquet_long(da_path) + df_ol = load_parquet_long(ol_path) + obs_series = load_usgs_obs(args.usgs_csv) + + print(f"DA rows: {len(df_da):,} | unique issue_times: {df_da['issue_time'].nunique()}") + print(f"OL rows: {len(df_ol):,} | unique issue_times: {df_ol['issue_time'].nunique()}") + print(f"USGS obs range: {obs_series.index.min()} .. {obs_series.index.max()} " + f"({len(obs_series):,} hours)") + print(f"USGS Helene peak (reference): {USGS_HELENE_PEAK_M3S:.0f} m³/s") + + # ----- Pooled (all issue times) ----- + da_pooled = metrics_by_lead(df_da, obs_series) + ol_pooled = metrics_by_lead(df_ol, obs_series) + out_pooled = os.path.join(out_dir, "lead_time_decay_gauge_pooled.png") + plot_two_panel(out_pooled, da_pooled[0], da_pooled, ol_pooled[0], ol_pooled, + title="Gauge-level lead-time decay — USGS 03463300\n" + "All issue times pooled (Oct 2023 – Oct 2024)") + + # ----- Regime split: Helene, storm, low-flow at gauge ----- + issue_times = pd.to_datetime(sorted(df_da['issue_time'].unique())) + obs_at_issue = pd.Series(issue_times, index=issue_times).map(obs_series) + + helene_mask = (issue_times >= HELENE_START) & (issue_times <= HELENE_END) + storm_mask = obs_at_issue > STORM_OBS_THRESHOLD_M3S + lowflow_mask = obs_at_issue < LOWFLOW_OBS_THRESHOLD_M3S + + regime_metrics = [] + for label, mask in [ + (f"Helene window (Sep 24–28 2024)", helene_mask), + (f"Storm hours (obs(t0) > {STORM_OBS_THRESHOLD_M3S:.0f} m³/s)", storm_mask), + (f"Low flow (obs(t0) < {LOWFLOW_OBS_THRESHOLD_M3S:.0f} m³/s)", lowflow_mask), + ]: + kept = set(pd.to_datetime(issue_times[mask])) + da_m = metrics_by_lead(df_da, obs_series, kept) + ol_m = metrics_by_lead(df_ol, obs_series, kept) + regime_metrics.append((label, len(kept), da_m, ol_m)) + + out_regime = os.path.join(out_dir, "lead_time_decay_gauge_by_regime.png") + plot_three_panel_regime(out_regime, regime_metrics) + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/run_4a_f4.sh b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/run_4a_f4.sh new file mode 100644 index 00000000..d6ef9adc --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/run_4a_f4.sh @@ -0,0 +1,73 @@ +#!/usr/bin/env bash +# F4 Direct variance (1 gauge holdout) — 4a lead-time decay plots. +# +# Requires lead-time forecast CSVs generated by batch_run_leadtime_f4.sh first. +# Gauge-level decay plot is skipped if routed parquets don't exist yet. +# +# Usage: +# bash run_4a_f4.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +LEADTIME_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder4_dynamic_variance_direct +DA_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder4_dynamic_variance_direct +OBS_DIR=/mnt/disk2/1400_sites_helene/catchment_ts_no_03463300_dynamic_variance +ROUTE_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder4_dynamic_variance_direct_leadtime_routed +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +echo "[F4-4a] Lead-time decay plots — leadtime dir: $LEADTIME_DIR" + +for CAT in "${CATS[@]}"; do + DA_CSV="$LEADTIME_DIR/$CAT/${CAT}_lead_time_forecasts_da.csv" + if [ ! -f "$DA_CSV" ]; then + echo " [$CAT] No lead-time CSV — skipping" + continue + fi + + echo " [$CAT] Plotting pooled decay..." + $TROUTE "$SCRIPT_DIR/plot_lead_time_decay.py" \ + --cat-id "$CAT" \ + --leadtime-dir "$LEADTIME_DIR" \ + --da-dir "$DA_DIR" \ + --obs-dir "$OBS_DIR" + + echo " [$CAT] Plotting decay by regime..." + $TROUTE "$SCRIPT_DIR/plot_lead_time_decay_by_regime.py" \ + --cat-id "$CAT" \ + --leadtime-dir "$LEADTIME_DIR" \ + --da-dir "$DA_DIR" \ + --obs-dir "$OBS_DIR" +done + +if [ -f "$ROUTE_DIR/routed_leadtime_da_full.parquet" ]; then + echo "[F4-4a] Gauge-level decay (routed)..." + $TROUTE "$SCRIPT_DIR/plot_lead_time_decay_gauge.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" + + echo "[F4-4a] Fixed-target error plots..." + $TROUTE "$SCRIPT_DIR/plot_forecast_error_fixed_target.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" + + echo "[F4-4a] Per-init error plots..." + $TROUTE "$SCRIPT_DIR/plot_forecast_error_per_init.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" +else + echo "[F4-4a] Routed parquets not found — skipping gauge-level and error plots." + echo " Run route_leadtime_f4.sh first, then re-run this script." +fi + +echo "[F4-4a] Done." diff --git a/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_crossed_ensemble.py b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_crossed_ensemble.py new file mode 100644 index 00000000..7cb00a93 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_crossed_ensemble.py @@ -0,0 +1,229 @@ +""" +plot_crossed_ensemble.py -- 4b: 600-member crossed ensemble vs USGS obs. + +Loads the crossed ensemble parquet produced by run_crossed_ensemble.py: + //_crossed_ensemble.parquet + Columns: issue_time, lead_hour, member_0000..member_0599 + +Projects all members to valid_time = issue_time + lead_hour hours. +Aggregates at each valid_time across ALL members from ALL issue times +to build a probabilistic envelope (5th/25th/50th/75th/95th percentiles). + +Two output figures: + 4b_crossed_ensemble_helene.png + -- Full Helene window (Sep 24-30) probabilistic envelope + USGS obs. + Percentile shading in two layers (5-95 outer, 25-75 inner). + 4b_crossed_ensemble_peak.png + -- Zoomed to Helene peak (Sep 26 12 UTC -> Sep 28 06 UTC), + same shading + USGS obs + individual spaghetti from Sep 27 inits. + +Usage: + python3 plot_crossed_ensemble.py \\ + --ensemble-dir /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble \\ + --cat-id cat-1016300 \\ + --usgs-csv /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv \\ + --out-dir /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble/cat-1016300 +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import matplotlib.patches as mpatches + +DEFAULT_ENSEMBLE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" +DEFAULT_CAT_ID = "cat-1016300" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +PLOT_START = pd.Timestamp("2024-09-24 00:00:00") +PLOT_END = pd.Timestamp("2024-09-28 18:00:00") # cut trailing spike +PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +PEAK_END = pd.Timestamp("2024-09-28 06:00:00") +HELENE_PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +HELENE_PEAK_END = pd.Timestamp("2024-09-28 00:00:00") + + +def load_ensemble(path): + df = pd.read_parquet(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + return df, member_cols + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def build_percentile_envelope(df, member_cols, t_start, t_end): + """ + For each valid_time in [t_start, t_end], stack all member values + from all issue times and compute percentiles. + Returns DataFrame indexed by valid_time with columns p05, p25, p50, p75, p95. + """ + mask = (df["valid_time"] >= t_start) & (df["valid_time"] <= t_end) + sub = df[mask].copy() + sub["valid_time"] = sub["valid_time"].dt.floor("1h") + + records = [] + for vt, grp in sub.groupby("valid_time"): + vals = grp[member_cols].to_numpy(dtype=float).flatten() + vals = vals[~np.isnan(vals)] * MM_H_TO_M3_S + if len(vals) == 0: + continue + records.append({ + "valid_time": vt, + "p05": np.percentile(vals, 5), + "p25": np.percentile(vals, 25), + "p50": np.percentile(vals, 50), + "p75": np.percentile(vals, 75), + "p95": np.percentile(vals, 95), + "n": len(vals), + }) + return pd.DataFrame(records).set_index("valid_time").sort_index() + + +def _add_helene_band(ax, label=True): + ax.axvspan(HELENE_PEAK_START, HELENE_PEAK_END, + color="salmon", alpha=0.12, zorder=0) + if label: + ax.text(HELENE_PEAK_START + pd.Timedelta(hours=6), 0.97, "Helene peak", + transform=ax.get_xaxis_transform(), + fontsize=9, color="firebrick", ha="left", va="top") + + +def _format_xaxis(ax, major_interval=1, minor_hours=6): + ax.xaxis.set_major_locator(mdates.DayLocator(interval=major_interval)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.xaxis.set_minor_locator(mdates.HourLocator(interval=minor_hours)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.2, lw=0.4) + + +def plot_envelope(ax, env, obs_series, t_start, t_end, + ens_color="tab:blue", obs_color="black", helene_label=True): + _add_helene_band(ax, label=helene_label) + + # Outer band: 5-95 + ax.fill_between(env.index, env["p05"], env["p95"], + color=ens_color, alpha=0.15, zorder=2, + label="5th-95th percentile") + # Inner band: 25-75 + ax.fill_between(env.index, env["p25"], env["p75"], + color=ens_color, alpha=0.30, zorder=3, + label="25th-75th percentile") + # Median + ax.plot(env.index, env["p50"], + color=ens_color, lw=2.2, zorder=5, + label="Ensemble median (600 members)") + + # USGS obs + obs_w = obs_series.loc[t_start:t_end] + ax.plot(obs_w.index, obs_w.values, + color=obs_color, lw=1.8, zorder=6, label="USGS obs") + + ax.set_xlim(t_start, t_end) + ax.set_ylabel("Discharge (m3/s)", fontsize=11) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--ensemble-dir", default=DEFAULT_ENSEMBLE_DIR) + parser.add_argument("--cat-id", default=DEFAULT_CAT_ID) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--out-dir", default=None) + args = parser.parse_args() + + cat_dir = os.path.join(args.ensemble_dir, args.cat_id) + out_dir = args.out_dir or cat_dir + os.makedirs(out_dir, exist_ok=True) + + parquet_path = os.path.join(cat_dir, f"{args.cat_id}_crossed_ensemble.parquet") + print(f"Loading crossed ensemble: {parquet_path}") + df, member_cols = load_ensemble(parquet_path) + obs = load_usgs(args.usgs_csv) + print(f" {df['issue_time'].nunique()} issue times, " + f"{len(member_cols)} members per row") + + # ------------------------------------------------------------------ # + # Figure 4b-full: full Helene window probabilistic envelope # + # ------------------------------------------------------------------ # + print("Building percentile envelope (full window)...") + env_full = build_percentile_envelope(df, member_cols, PLOT_START, PLOT_END) + print(f" Valid times with data: {len(env_full)}, " + f"median members per timestep: {env_full['n'].median():.0f}") + + fig, ax = plt.subplots(figsize=(17, 7)) + plot_envelope(ax, env_full, obs, PLOT_START, PLOT_END) + ax.set_xlabel("Date", fontsize=11) + ax.set_title( + f"4b -- 600-member crossed ensemble (30 met x 20 hydro-state) | {args.cat_id}\n" + "Shaded: 5-95th (light) and 25-75th (dark) percentile across all members " + "and issue times | Sep 24-30 2024 | USGS 03463300", + fontsize=11, + ) + ax.legend(fontsize=10, loc="upper left", frameon=True, framealpha=0.9) + _format_xaxis(ax) + plt.tight_layout() + out1 = os.path.join(out_dir, f"{args.cat_id}_4b_crossed_ensemble_helene.png") + plt.savefig(out1, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + # ------------------------------------------------------------------ # + # Figure 4b-peak: zoomed to Helene peak window # + # ------------------------------------------------------------------ # + print("Building percentile envelope (peak window)...") + env_peak = build_percentile_envelope(df, member_cols, PEAK_START, PEAK_END) + + fig, (ax_main, ax_zoom) = plt.subplots( + 2, 1, figsize=(14, 11), + gridspec_kw={"height_ratios": [1, 1.4]}, sharex=False) + + # Top panel: full window context (smaller) + plot_envelope(ax_main, env_full, obs, PLOT_START, PLOT_END, helene_label=False) + ax_main.set_title("Full Helene window context (Sep 24 - Sep 28 18 UTC)", fontsize=10) + _format_xaxis(ax_main) + # Shade the zoom region on context panel + ax_main.axvspan(PEAK_START, PEAK_END, color="gold", alpha=0.18, zorder=1) + ax_main.text(PEAK_START + pd.Timedelta(hours=1), 0.97, "zoom", + transform=ax_main.get_xaxis_transform(), + fontsize=8, color="goldenrod", ha="left", va="top") + + # Bottom panel: peak zoom + plot_envelope(ax_zoom, env_peak, obs, PEAK_START, PEAK_END, helene_label=True) + ax_zoom.set_xlabel("Date", fontsize=11) + ax_zoom.set_title( + "Helene peak zoom (Sep 26 12 UTC - Sep 28 06 UTC)", fontsize=10) + _format_xaxis(ax_zoom, major_interval=1, minor_hours=3) + ax_zoom.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %HUTC")) + ax_zoom.legend(fontsize=10, loc="lower right", frameon=True, framealpha=0.9) + + fig.suptitle( + f"4b -- 600-member crossed ensemble vs USGS 03463300 | {args.cat_id}\n" + "30 met forcing draws x 20 DA analysis state draws | R=0.07 mm2/h2", + fontsize=12, y=1.01, + ) + plt.tight_layout() + out2 = os.path.join(out_dir, f"{args.cat_id}_4b_crossed_ensemble_peak.png") + plt.savefig(out2, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_perturbation_category_shaded.py b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_perturbation_category_shaded.py new file mode 100644 index 00000000..b845c52d --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_perturbation_category_shaded.py @@ -0,0 +1,165 @@ +""" +Per-catchment shaded ensemble-band plot, organized by perturbation category. + +Three categories: + 1. Initial states (red) + 2. Meteorological forcings (blue) + 3. Hydrological states (green) + +For each category, all 20 ensemble members are shown as a shaded band +(min-max envelope fill) plus a thicker median line in the same color. +Qkrig observation overlaid in black. Hurricane Helene peak window shaded +in pink. Styled after a standard ensemble-forecast figure layout. + +Inputs (existing per-source sensitivity CSVs from run_perturbation_sensitivity.py): + //_sensitivity_init.csv (20 members, init only) + //_sensitivity_forcing.csv (20 members, forcing only) + //_sensitivity_process.csv (20 members, process noise only) + //_test_results.csv (Qkrig obs) + +Outputs: + //_perturbation_categories_linear.png + //_perturbation_categories_log.png +""" +import os +import pandas as pd +import numpy as np +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import matplotlib.patches as mpatches + +CAT = "cat-1016300" + +SEN_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_sensitivity" +OBS_DIR = "/mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance" + +OUT_LINEAR = os.path.join(SEN_DIR, CAT, f"{CAT}_perturbation_categories_linear.png") +OUT_LOG = os.path.join(SEN_DIR, CAT, f"{CAT}_perturbation_categories_log.png") + +# Plot window — wider context, similar to the paper's Sep 10 - Oct 08 +PLOT_START = pd.Timestamp("2024-09-20") +PLOT_END = pd.Timestamp("2024-10-05") + +# Helene peak band +HELENE_START = pd.Timestamp("2024-09-26 12:00:00") +HELENE_END = pd.Timestamp("2024-09-28 00:00:00") + +# Category configuration: file suffix, display label, color +CATEGORIES = [ + ("init", "Initial states only", "tab:red"), + ("forcing", "Meteorological forcings only", "tab:blue"), + ("process", "Hydrological states only", "tab:green"), +] + + +def load_members(source): + path = os.path.join(SEN_DIR, CAT, f"{CAT}_sensitivity_{source}.csv") + if not os.path.exists(path): + return None, None + df = pd.read_csv(path, parse_dates=["date"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + return df["date"].values, df[member_cols].to_numpy(dtype=float) + + +def load_obs(): + p = os.path.join(OBS_DIR, f"{CAT}.csv") + if not os.path.exists(p): + return None, None + df = pd.read_csv(p, parse_dates=["date"]) + return df["date"].values, df["qkrig"].values + + +def plot_panel(ax, obs_dates, obs_vals, log_y=False): + handles_labels = [] # for the legend + + # Plot each category as a shaded band + median line + for source, label, color in CATEGORIES: + dates, q = load_members(source) + if dates is None: + continue + d = pd.to_datetime(dates) + mask = (d >= PLOT_START) & (d <= PLOT_END) + if mask.sum() == 0: + continue + # Min-max envelope across all 20 members per timestep (widest possible band). + # Bands are visually narrow even with min/max because perturbations are + # tuned for production EnKF stability, not for max visible spread. + q_window = q[mask, :] + qmin = np.nanmin(q_window, axis=1) + qmax = np.nanmax(q_window, axis=1) + median = np.nanmedian(q_window, axis=1) + + ax.fill_between(d[mask], qmin, qmax, + color=color, alpha=0.30, zorder=2, + edgecolor="none") + line, = ax.plot(d[mask], median, + color=color, lw=1.7, alpha=0.95, zorder=3, + label=f"{label} (N=20)") + handles_labels.append((line, label)) + + # Helene peak shaded band (vertical) + ax.axvspan(HELENE_START, HELENE_END, + color="salmon", alpha=0.15, zorder=1) + ax.text((HELENE_START + (HELENE_END - HELENE_START) / 2), + ax.get_ylim()[1] if not log_y else 1.0, + "Helene peak", + fontsize=9, color="salmon", + ha="center", va="bottom", zorder=3) + + # Observation + if obs_dates is not None: + od = pd.to_datetime(obs_dates) + om = (od >= PLOT_START) & (od <= PLOT_END) + ax.plot(od[om], obs_vals[om], + color="black", lw=1.4, label="Qkrig (obs)", zorder=4) + + ax.set_xlabel("Date", fontsize=10) + if log_y: + ax.set_yscale("log") + ax.set_ylabel(r"q (mm hr$^{-1}$) [log scale]", fontsize=10) + ax.set_ylim(0.01, None) + else: + ax.set_ylabel("Discharge (mm/h)", fontsize=10) + ax.tick_params(axis="x", rotation=30, labelsize=9) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=2)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.grid(True, alpha=0.2) + ax.legend(loc="upper left", fontsize=9, frameon=True, framealpha=0.9) + + +def main(): + obs_dates, obs_vals = load_obs() + + # ----- Linear-y ----- + fig, ax = plt.subplots(1, 1, figsize=(20, 7)) + plot_panel(ax, obs_dates, obs_vals, log_y=False) + fig.suptitle( + f"Ensemble forecast by perturbation category - {CAT} - " + f"Sep 20 - Oct 5, 2024 (Hurricane Helene window)\n" + "Shaded bands = min-max envelope across 20 members. " + "Lines = ensemble median. DA off in all three sub-experiments.", + fontsize=11, y=0.995, + ) + plt.tight_layout() + plt.savefig(OUT_LINEAR, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_LINEAR}") + + # ----- Log-y (paper style) ----- + fig, ax = plt.subplots(1, 1, figsize=(20, 7)) + plot_panel(ax, obs_dates, obs_vals, log_y=True) + fig.suptitle( + f"Ensemble forecast by perturbation category - {CAT} - log-scale q - " + f"Sep 20 - Oct 5, 2024 (Hurricane Helene window)\n" + "Shaded bands = min-max envelope across 20 members. " + "Lines = ensemble median. DA off in all three sub-experiments.", + fontsize=11, y=0.995, + ) + plt.tight_layout() + plt.savefig(OUT_LOG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_LOG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_combined.py b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_combined.py new file mode 100644 index 00000000..c6d531a0 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_combined.py @@ -0,0 +1,235 @@ +#!/usr/bin/env python3 +""" +plot_routed_ensemble_combined.py +Combined poster plot: 600-member routed ensemble envelope + +deterministic DA lines (Vrugt/no-Vrugt) vs USGS gauge obs at 03463300. + +Inputs: + --vrugt-csv : routed_Q_test.csv from vrugt_dynamic_routed/ + --novrugt-csv : routed_Q_test.csv from novrugt_dynamic_routed/ + --ensemble-pq : routed_crossed_ensemble.parquet from run_route_crossed_ensemble.py + --out-dir : output directory + +The ensemble parquet columns: issue_time, lead_hour, member_0000..member_0599 (Q in m3/s) + +Usage: + python3 plot_routed_ensemble_combined.py \ + --vrugt-csv /mnt/disk2/suma_helen_poster/da_results/vrugt_dynamic_routed/routed_Q_test.csv \ + --novrugt-csv /mnt/disk2/suma_helen_poster/da_results/novrugt_dynamic_routed/routed_Q_test.csv \ + --ensemble-pq /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble_routed/routed_crossed_ensemble.parquet \ + --out-dir /mnt/disk2/suma_helen_poster/da_results/comparison_plots +""" + +import argparse +import os + +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +HELENE_START = pd.Timestamp("2024-09-25 00:00:00") +HELENE_END = pd.Timestamp("2024-09-29 00:00:00") +ZOOM_START = pd.Timestamp("2024-09-24 00:00:00") +ZOOM_END = pd.Timestamp("2024-09-28 18:00:00") # cut trailing ensemble bump + +COLOR_OBS = "black" +COLOR_VRUGT = "#1f77b4" # blue +COLOR_NOVRUGT = "#d62728" # red +COLOR_ENS = "#1f77b4" # same blue family as Vrugt + + +def load_det_csv(path): + df = pd.read_csv(path, parse_dates=["date"]) + return df.set_index("date").sort_index() + + +def build_envelope(pq_path, zoom_start, zoom_end): + """ + Load routed ensemble parquet and build percentile envelope by valid_time. + Returns DataFrame indexed by valid_time with columns p05, p25, p50, p75, p95. + """ + df = pd.read_parquet(pq_path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = [c for c in df.columns if c.startswith("member_")] + + records = [] + for _, row in df.iterrows(): + t0 = row["issue_time"] + lead = int(row["lead_hour"]) + vt = t0 + pd.Timedelta(hours=lead) + vals = row[member_cols].values.astype(float) + records.append({"valid_time": vt, "vals": vals}) + + # Group by valid_time (floor to hour), stack all member values across issue_times + from collections import defaultdict + groups = defaultdict(list) + for rec in records: + vt = rec["valid_time"].floor("1h") + groups[vt].extend(rec["vals"].tolist()) + + rows = [] + for vt in sorted(groups): + if vt < zoom_start or vt > zoom_end: + continue + v = np.array(groups[vt]) + v = v[~np.isnan(v)] + if len(v) == 0: + continue + rows.append({ + "valid_time": vt, + "p05": np.percentile(v, 5), + "p25": np.percentile(v, 25), + "p50": np.percentile(v, 50), + "p75": np.percentile(v, 75), + "p95": np.percentile(v, 95), + }) + + return pd.DataFrame(rows).set_index("valid_time") + + +def compute_kge(obs, sim): + mask = ~(np.isnan(obs) | np.isnan(sim)) + o, s = obs[mask], sim[mask] + if len(o) < 2 or np.std(o) == 0: + return np.nan + r = np.corrcoef(o, s)[0, 1] + return 1.0 - np.sqrt((r-1)**2 + (np.std(s)/np.std(o)-1)**2 + (np.mean(s)/np.mean(o)-1)**2) + + +def _format_xaxis(ax, locator, fmt, rotation=30): + ax.xaxis.set_major_locator(locator) + ax.xaxis.set_major_formatter(mdates.DateFormatter(fmt)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=rotation, ha="right", fontsize=9) + + +def _add_helene_band(ax): + ax.axvspan(HELENE_START, HELENE_END, color="gold", alpha=0.15, zorder=0) + + +def plot_zoom(ax, dates, obs, sim_v, sim_nv, kge_v, kge_nv, env=None): + mask = (dates >= ZOOM_START) & (dates <= ZOOM_END) + dz = dates[mask] + + # Ensemble envelope (bottom layer) + if env is not None: + ax.fill_between(env.index, env["p05"], env["p95"], + color=COLOR_ENS, alpha=0.18, label="Ensemble 5-95th pct") + ax.fill_between(env.index, env["p25"], env["p75"], + color=COLOR_ENS, alpha=0.30, label="Ensemble 25-75th pct") + ax.plot(env.index, env["p50"], + color=COLOR_ENS, lw=1.2, linestyle="--", alpha=0.7, label="Ensemble median") + + # Deterministic DA lines + ax.plot(dz, sim_nv[mask], color=COLOR_NOVRUGT, lw=1.4, linestyle="--", + label=f"DA -- constant R KGE={kge_nv:.3f}", zorder=3) + ax.plot(dz, sim_v[mask], color=COLOR_VRUGT, lw=1.6, + label=f"DA -- dynamic R (Vrugt) KGE={kge_v:.3f}", zorder=4) + + # USGS obs on top + ax.plot(dz, obs[mask], color=COLOR_OBS, lw=1.6, label="USGS obs (gauge 03463300)", zorder=5) + + peak_usgs = np.nanmax(obs[mask]) + ax.axhline(peak_usgs, color=COLOR_OBS, lw=0.7, linestyle=":", alpha=0.5) + # Place annotation to the right of the peak to avoid legend overlap + ax.text(pd.Timestamp("2024-09-28 06:00:00"), peak_usgs * 1.02, + f"USGS peak {peak_usgs:.0f} m³/s", fontsize=8.5, color="black", + alpha=0.8, ha="right") + + _add_helene_band(ax) + _format_xaxis(ax, mdates.DayLocator(interval=1), "%b %d") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--vrugt-csv", required=True) + parser.add_argument("--novrugt-csv", required=True) + parser.add_argument("--ensemble-pq", required=True, + help="routed_crossed_ensemble.parquet from run_route_crossed_ensemble.py") + parser.add_argument("--out-dir", required=True) + args = parser.parse_args() + os.makedirs(args.out_dir, exist_ok=True) + + vrugt = load_det_csv(args.vrugt_csv) + novrugt = load_det_csv(args.novrugt_csv) + + obs = vrugt["Q_usgs_m3s"].values + dates = vrugt.index + sim_v = vrugt["Q_routed_m3s"].values + sim_nv = novrugt["Q_routed_m3s"].reindex(dates).values + + kge_v = compute_kge(obs, sim_v) + kge_nv = compute_kge(obs, sim_nv) + + print("Building ensemble envelope...") + env = build_envelope(args.ensemble_pq, ZOOM_START, ZOOM_END) + print(f" Envelope: {len(env)} timesteps, " + f"peak p95={env['p95'].max():.1f} m3/s, peak p50={env['p50'].max():.1f} m3/s") + + peak_usgs = np.nanmax(obs[(dates >= ZOOM_START) & (dates <= ZOOM_END)]) + in_env = env["p95"].max() >= peak_usgs + print(f" USGS peak {peak_usgs:.1f} m3/s {'IS' if in_env else 'IS NOT'} within p95 envelope") + + # ------------------------------------------------------------------ # + # Figure 1: Helene zoom only -- main poster panel + # ------------------------------------------------------------------ # + fig, ax = plt.subplots(figsize=(11, 5.5)) + plot_zoom(ax, dates, obs, sim_v, sim_nv, kge_v, kge_nv, env) + ax.set_xlabel("Date (UTC)", fontsize=11) + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_title( + "Hurricane Helene -- Probabilistic discharge forecast at gauge 03463300 (South Toe River)\n" + "600-member CFE ensemble + DA (Muskingum routing)", + fontsize=11) + ax.legend(fontsize=9, loc="lower right") + ax.grid(True, alpha=0.25) + plt.tight_layout() + out1 = os.path.join(args.out_dir, "helene_ensemble_vs_usgs.png") + plt.savefig(out1, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + # ------------------------------------------------------------------ # + # Figure 2: two-panel (full period + Helene zoom) + # ------------------------------------------------------------------ # + fig, (ax_top, ax_bot) = plt.subplots( + 2, 1, figsize=(14, 9), + gridspec_kw={"height_ratios": [1, 1.5]}) + + # Top: full test period (deterministic only, no ensemble for readability) + ax_top.plot(dates, obs, color=COLOR_OBS, lw=0.8, label="USGS obs", zorder=3) + ax_top.plot(dates, sim_v, color=COLOR_VRUGT, lw=0.9, + label=f"DA dynamic R (Vrugt) KGE={kge_v:.3f}", zorder=2) + ax_top.plot(dates, sim_nv, color=COLOR_NOVRUGT, lw=0.9, linestyle="--", + label=f"DA constant R KGE={kge_nv:.3f}", zorder=2) + ax_top.axvspan(ZOOM_START, ZOOM_END, color="gray", alpha=0.12, zorder=0) + ax_top.axvspan(HELENE_START, HELENE_END, color="gold", alpha=0.15, zorder=0) + ax_top.text(HELENE_START + pd.Timedelta(days=0.5), + np.nanmax(obs) * 0.88, "Helene", + fontsize=9, color="goldenrod", fontweight="bold") + ax_top.set_ylabel("Discharge (m³/s)", fontsize=10) + ax_top.set_title( + "Routed discharge at gauge 03463300 (South Toe River) -- Test period\n" + "DA with dynamic Vrugt R vs. constant R -- Muskingum routing", + fontsize=11) + ax_top.legend(fontsize=9, loc="upper left") + ax_top.grid(True, alpha=0.25) + _format_xaxis(ax_top, mdates.MonthLocator(interval=2), "%Y-%m") + + # Bottom: Helene zoom with ensemble + plot_zoom(ax_bot, dates, obs, sim_v, sim_nv, kge_v, kge_nv, env) + ax_bot.set_xlabel("Date (UTC)", fontsize=10) + ax_bot.set_ylabel("Discharge (m³/s)", fontsize=10) + ax_bot.set_title("Hurricane Helene window -- 600-member ensemble envelope", fontsize=10) + ax_bot.legend(fontsize=9, loc="lower right") + ax_bot.grid(True, alpha=0.25) + + plt.tight_layout() + out2 = os.path.join(args.out_dir, "helene_ensemble_twopanel.png") + plt.savefig(out2, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_vs_usgs.py b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_vs_usgs.py new file mode 100644 index 00000000..55c52d85 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_vs_usgs.py @@ -0,0 +1,176 @@ +#!/usr/bin/env python3 +""" +plot_routed_ensemble_vs_usgs.py +Plot T-route routed 600-member ensemble envelope vs USGS at gauge 03463300. + +Input: routed_crossed_ensemble.parquet from run_route_crossed_ensemble.py + Columns: issue_time, lead_hour, member_0000..member_0599 (Q in m3/s) + +Usage: + python3 plot_routed_ensemble_vs_usgs.py \ + --routed-pq /mnt/disk2/.../folder1_vrugt_routed/routed_crossed_ensemble.parquet \ + --usgs-csv /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv \ + --label "F1 Vrugt — 1 gauge holdout" \ + --out-dir ~/plots_f1_1gauge +""" + +import argparse +from collections import defaultdict +from pathlib import Path + +import matplotlib.dates as mdates +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd + +WATERSHED_AREA_KM2 = 113.18 # cat-1016300 drainage area +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +PLOT_START = pd.Timestamp("2024-09-24 00:00:00") +PLOT_END = pd.Timestamp("2024-09-28 18:00:00") +PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +PEAK_END = pd.Timestamp("2024-09-28 06:00:00") +HELENE_PEAK = pd.Timestamp("2024-09-27 12:00:00") + +COLOR_ENS = "#1f77b4" +COLOR_OBS = "black" + + +def build_envelope(pq_path, t_start, t_end): + df = pd.read_parquet(pq_path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = [c for c in df.columns if c.startswith("member_")] + + groups = defaultdict(list) + for _, row in df.iterrows(): + vt = (row["issue_time"] + pd.Timedelta(hours=int(row["lead_hour"]))).floor("1h") + if t_start <= vt <= t_end: + groups[vt].extend(row[member_cols].values.astype(float).tolist()) + + rows = [] + for vt in sorted(groups): + v = np.array(groups[vt]) + v = v[~np.isnan(v)] + if len(v) == 0: + continue + rows.append({"valid_time": vt, + "p05": np.percentile(v, 5), "p25": np.percentile(v, 25), + "p50": np.percentile(v, 50), "p75": np.percentile(v, 75), + "p95": np.percentile(v, 95)}) + return pd.DataFrame(rows).set_index("valid_time") + + +def load_usgs(path): + df = pd.read_csv(path) + print(f" USGS CSV columns: {list(df.columns)}") + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() or "discharge" in c.lower()) + print(f" Using date_col={date_col!r}, q_col={q_col!r}") + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series.sort_index() + + +def compute_kge(obs, sim): + mask = ~(np.isnan(obs) | np.isnan(sim)) + o, s = obs[mask], sim[mask] + if len(o) < 2 or np.std(o) == 0: + return np.nan + r = np.corrcoef(o, s)[0, 1] + return 1 - np.sqrt((r-1)**2 + (np.std(s)/np.std(o)-1)**2 + (np.mean(s)/np.mean(o)-1)**2) + + +def _shade(ax, env): + ax.fill_between(env.index, env["p05"], env["p95"], + color=COLOR_ENS, alpha=0.18, label="5th–95th percentile") + ax.fill_between(env.index, env["p25"], env["p75"], + color=COLOR_ENS, alpha=0.32, label="25th–75th percentile") + ax.plot(env.index, env["p50"], color=COLOR_ENS, lw=1.5, label="Ensemble median (600 members)") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--routed-pq", required=True) + parser.add_argument("--usgs-csv", required=True) + parser.add_argument("--label", default="F1 Vrugt — 1 gauge holdout") + parser.add_argument("--out-dir", required=True) + args = parser.parse_args() + + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + print("Building full-window envelope...") + env_full = build_envelope(args.routed_pq, PLOT_START, PLOT_END) + print("Building peak-window envelope...") + env_peak = build_envelope(args.routed_pq, PEAK_START, PEAK_END) + + print("Loading USGS obs...") + try: + obs_all = load_usgs(args.usgs_csv) + obs_full = obs_all.loc[PLOT_START:PLOT_END] + obs_peak = obs_all.loc[PEAK_START:PEAK_END] + print(f" Loaded {len(obs_all)} USGS obs rows; " + f"full window: {len(obs_full)}, peak window: {len(obs_peak)}") + except Exception as e: + print(f" Warning: USGS load failed ({e}) — plotting without obs") + obs_full = obs_peak = None + + kge = np.nan + if obs_full is not None and len(env_full) > 0: + med = env_full["p50"] + obs_hourly = obs_full.copy() + obs_hourly.index = obs_hourly.index.floor("1h") + obs_hourly = obs_hourly[~obs_hourly.index.duplicated(keep="first")] + aligned = obs_hourly.reindex(med.index) + print(f" KGE alignment: {aligned.notna().sum()} matched out of {len(med)} timesteps") + kge = compute_kge(aligned.values, med.values) + + # ── Figure 1: Full Helene window ────────────────────────────────────────── + fig, ax = plt.subplots(figsize=(12, 4.5)) + _shade(ax, env_full) + if obs_full is not None: + ax.plot(obs_full.index, obs_full.values, color=COLOR_OBS, lw=1.6, + label="USGS obs (gauge 03463300)") + ax.axvspan(PEAK_START, PEAK_END, color="salmon", alpha=0.12, zorder=0, label="Helene peak window") + ax.set_ylabel("Discharge (m³/s)") + ax.set_xlabel("Date") + ax.set_title(f"Routed 600-member ensemble vs USGS | {args.label}\n" + f"Ensemble median KGE = {kge:.3f} | Sep 24–28 2024") + ax.legend(fontsize=9, loc="upper left") + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + plt.tight_layout() + out = out_dir / "routed_ensemble_vs_usgs_full.png" + fig.savefig(out, dpi=150) + print(f"Saved: {out}") + plt.close(fig) + + # ── Figure 2: Peak zoom ─────────────────────────────────────────────────── + fig, ax = plt.subplots(figsize=(10, 4.5)) + _shade(ax, env_peak) + if obs_peak is not None: + ax.plot(obs_peak.index, obs_peak.values, color=COLOR_OBS, lw=1.6, + label="USGS obs (gauge 03463300)") + peak_val = float(obs_peak.max()) + ax.axhline(peak_val, color=COLOR_OBS, lw=0.7, ls=":", alpha=0.5) + ax.annotate(f"USGS peak\n{peak_val:.0f} m³/s", + xy=(HELENE_PEAK, peak_val), xytext=(10, -40), + textcoords="offset points", fontsize=8, color=COLOR_OBS) + ax.set_ylabel("Discharge (m³/s)") + ax.set_xlabel("Date") + ax.set_title(f"Helene peak zoom | {args.label}") + ax.legend(fontsize=9) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %HUTC")) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + plt.tight_layout() + out = out_dir / "routed_ensemble_vs_usgs_peak.png" + fig.savefig(out, dpi=150) + print(f"Saved: {out}") + plt.close(fig) + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/run_4b_crossed_f4.sh b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/run_4b_crossed_f4.sh new file mode 100644 index 00000000..4f588dc7 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/run_4b_crossed_f4.sh @@ -0,0 +1,40 @@ +#!/usr/bin/env bash +# F4 Direct variance (1 gauge holdout) — 4b crossed ensemble plots. +# Runs plot_crossed_ensemble.py for all 21 catchments. +# Reads //_crossed_ensemble.parquet (already generated). +# +# Usage: +# bash run_4b_crossed_f4.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +DATA_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder4_dynamic_variance_direct +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +echo "[F4-4b-crossed] Crossed ensemble plots — data: $DATA_DIR" + +for CAT in "${CATS[@]}"; do + PQ="$DATA_DIR/$CAT/${CAT}_crossed_ensemble.parquet" + if [ ! -f "$PQ" ]; then + echo " [$CAT] No crossed ensemble parquet — skipping" + continue + fi + echo " [$CAT] crossed ensemble..." + $TROUTE "$SCRIPT_DIR/plot_crossed_ensemble.py" \ + --ensemble-dir "$DATA_DIR" \ + --cat-id "$CAT" \ + --usgs-csv "$USGS_CSV" +done + +echo "[F4-4b-crossed] Done." diff --git a/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/run_4b_f4.sh b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/run_4b_f4.sh new file mode 100644 index 00000000..7b3318a0 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/run_4b_f4.sh @@ -0,0 +1,38 @@ +#!/usr/bin/env bash +# F4 Direct variance (1 gauge holdout) — 4b routed ensemble plots. +# +# Runs: +# 1. plot_routed_ensemble_vs_usgs.py — F4 ensemble envelope vs USGS at outlet +# 2. plot_routed_ensemble_combined.py — F1 Vrugt vs F4 direct variance comparison +# +# Requires routed_Q_test.csv and routed_crossed_ensemble.parquet in each DA dir. +# +# Usage: +# bash run_4b_f4.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +F1_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt +F4_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder4_dynamic_variance_direct +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv +OUT_DIR="$F4_DIR" + +echo "[F4-4b] Routed ensemble vs USGS (F4 only)..." +$TROUTE "$SCRIPT_DIR/plot_routed_ensemble_vs_usgs.py" \ + --routed-pq "$F4_DIR/routed_crossed_ensemble.parquet" \ + --usgs-csv "$USGS_CSV" \ + --label "F4 Direct variance — 1 gauge holdout" \ + --out-dir "$OUT_DIR" + +echo "[F4-4b] Combined comparison: F1 Vrugt vs F4 Direct variance..." +$TROUTE "$SCRIPT_DIR/plot_routed_ensemble_combined.py" \ + --vrugt-csv "$F1_DIR/routed_Q_test.csv" \ + --novrugt-csv "$F4_DIR/routed_Q_test.csv" \ + --ensemble-pq "$F4_DIR/routed_crossed_ensemble.parquet" \ + --out-dir "$OUT_DIR" + +echo "[F4-4b] Done. Outputs in: $OUT_DIR" +ls "$OUT_DIR"/*.png 2>/dev/null || echo " (no PNGs in out dir)" diff --git a/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4c_timeseries/plot_forecast_spaghetti.py b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4c_timeseries/plot_forecast_spaghetti.py new file mode 100644 index 00000000..655edaef --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4c_timeseries/plot_forecast_spaghetti.py @@ -0,0 +1,176 @@ +""" +plot_forecast_spaghetti.py + +"All forecasts ending at this hour" — verification-time-fixed spaghetti view. + +Window: Sep 24 18 UTC -> Sep 30 06 UTC. +For every initialization time whose 18-hour forecast overlaps this window: + DA : shaded min/max band across 20 members + ensemble mean line + OL : ensemble mean line only (no shading), dashed + +Lines are colored by initialization DATE (7 colors, Sep 24-30) so the +temporal progression is readable. USGS obs overlaid in black. + +Inputs: + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + Wide format: issue_time, lead_hour, member_00..member_19 (q in m³/s) + +Obs: + /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +Outputs: + /forecast_spaghetti_helene.png +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import matplotlib.patches as mpatches + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast_hardcoded_r_routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +# Verification window — only trajectories whose valid_time falls here are shown +PLOT_START = pd.Timestamp("2024-09-24 18:00:00") +PLOT_END = pd.Timestamp("2024-09-30 06:00:00") + +HELENE_PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +HELENE_PEAK_END = pd.Timestamp("2024-09-28 00:00:00") + +# One color per init date (Sep 24-30) +DATE_COLORS = { + "2024-09-24": "#1f77b4", + "2024-09-25": "#ff7f0e", + "2024-09-26": "#2ca02c", + "2024-09-27": "#d62728", + "2024-09-28": "#9467bd", + "2024-09-29": "#8c564b", + "2024-09-30": "#e377c2", +} + + +def load_parquet(path): + df = pd.read_parquet(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + df["ens_mean"] = df[member_cols].mean(axis=1) + df["ens_min"] = df[member_cols].min(axis=1) + df["ens_max"] = df[member_cols].max(axis=1) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + return df + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() + or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--route-dir", default=DEFAULT_ROUTE_DIR) + parser.add_argument("--out-dir", default=None) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--da-name", default="routed_leadtime_da_full.parquet") + parser.add_argument("--ol-name", default="routed_leadtime_openloop_full.parquet") + args = parser.parse_args() + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + print("Loading parquets...") + da = load_parquet(os.path.join(args.route_dir, args.da_name)) + ol = load_parquet(os.path.join(args.route_dir, args.ol_name)) + obs = load_usgs(args.usgs_csv) + + # Keep only valid_times inside the plot window + da = da[(da["valid_time"] >= PLOT_START) & (da["valid_time"] <= PLOT_END)] + ol = ol[(ol["valid_time"] >= PLOT_START) & (ol["valid_time"] <= PLOT_END)] + + issue_times = sorted(da["issue_time"].unique()) + print(f" {len(issue_times)} initialization times contribute to this window") + + fig, ax = plt.subplots(figsize=(18, 6)) + + # Helene peak shading + ax.axvspan(HELENE_PEAK_START, HELENE_PEAK_END, + color="salmon", alpha=0.12, zorder=0) + ax.text(HELENE_PEAK_START + pd.Timedelta(hours=6), 1.0, "Helene peak", + transform=ax.get_xaxis_transform(), + fontsize=9, color="firebrick", ha="left", va="top") + + # Plot each initialization time's trajectory + for t0 in issue_times: + date_str = str(pd.Timestamp(t0).date()) + color = DATE_COLORS.get(date_str, "gray") + + da_t = da[da["issue_time"] == t0].sort_values("valid_time") + ol_t = ol[ol["issue_time"] == t0].sort_values("valid_time") + + if da_t.empty: + continue + + # DA: shaded band (min-max across 20 members) + ensemble mean + ax.fill_between(da_t["valid_time"], da_t["ens_min"], da_t["ens_max"], + color=color, alpha=0.06, zorder=2) + ax.plot(da_t["valid_time"], da_t["ens_mean"], + color=color, lw=0.8, alpha=0.55, zorder=3) + + # OL: ensemble mean only, dashed + if not ol_t.empty: + ax.plot(ol_t["valid_time"], ol_t["ens_mean"], + color=color, lw=0.6, alpha=0.30, linestyle="--", zorder=2) + + # USGS obs + obs_w = obs.loc[PLOT_START:PLOT_END] + ax.plot(obs_w.index, obs_w.values, + color="black", lw=1.8, zorder=6, label="USGS obs") + + # Legend: one patch per init date + obs + DA/OL style + patches = [mpatches.Patch(color=c, label=f"Init {d}") + for d, c in DATE_COLORS.items()] + patches.append(plt.Line2D([0], [0], color="black", lw=1.8, label="USGS obs")) + patches.append(plt.Line2D([0], [0], color="gray", lw=1.2, alpha=0.7, + label="DA — ens mean (solid) + spread (shaded)")) + patches.append(plt.Line2D([0], [0], color="gray", lw=1.0, linestyle="--", + alpha=0.5, label="OL — ens mean (dashed)")) + ax.legend(handles=patches, fontsize=8, loc="upper left", + frameon=True, framealpha=0.9, ncol=2) + + ax.set_xlim(PLOT_START, PLOT_END) + ax.set_xlabel("Date", fontsize=11) + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_title( + "All 18-hour forecast trajectories ending in this window — USGS 03463300\n" + "Sep 24 18 UTC → Sep 30 06 UTC | DA: shaded band + mean | OL: mean dashed", + fontsize=11, + ) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=1)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.xaxis.set_minor_locator(mdates.HourLocator(interval=6)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.2, lw=0.4) + plt.tight_layout() + + out_path = os.path.join(out_dir, "forecast_spaghetti_helene.png") + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4c_timeseries/plot_helene_issue_time_hydrograph.py b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4c_timeseries/plot_helene_issue_time_hydrograph.py new file mode 100644 index 00000000..4e2e1dc2 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4c_timeseries/plot_helene_issue_time_hydrograph.py @@ -0,0 +1,198 @@ +""" +Per-issue-time forecast hydrograph diagnostic. + +For two contrasting issue times — one storm-peak (Helene) and one low-flow — +show what the DA and open-loop ensembles actually produce over the 18-hour +forecast window, alongside the kriging observation. Designed to answer in one +picture: is DA over-shooting, collapsing, or oscillating compared to open-loop? + +Each panel shows: + - All 20 DA members as faint purple lines + median (thick purple) + - All 20 open-loop members as faint gray lines + median (thick gray, dashed) + - Kriging obs as a thick black line (context: t0-6 through t0+18) + - Vertical line at t0 + annotation with obs(t0) and lead-1 ensemble means + +Inputs (from run_lead_time_forecast_sweep.py): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv +Obs (from production EnKF): + //_test_results.csv + +Output: + //_issue_time_hydrograph_helene_vs_lowflow.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +DEFAULT_LEADTIME_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast" +DEFAULT_DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" + +# Module-level globals set by main() +CAT = "cat-1016300" +LEADTIME_DIR = DEFAULT_LEADTIME_DIR +DA_DIR = DEFAULT_DA_DIR + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" +OBS_COLOR = "black" + +# Two issue times — pick a Helene-peak one and a typical low-flow one +HELENE_T0 = pd.Timestamp("2024-09-26 12:00:00") +LOWFLOW_T0 = pd.Timestamp("2024-03-15 12:00:00") + +CONTEXT_HOURS_BEFORE = 6 # hours of obs context shown before t0 +LEAD_HOURS_AFTER = 18 # forecast horizon + +OUT_PNG = os.path.join( + LEADTIME_DIR, CAT, + f"{CAT}_issue_time_hydrograph_helene_vs_lowflow.png", +) + + +def load_forecasts(path): + df = pd.read_csv(path, parse_dates=['issue_time', 'valid_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + return df, member_cols + + +def load_obs(): + obs_path = os.path.join(DA_DIR, CAT, f"{CAT}_test_results.csv") + df = pd.read_csv(obs_path, parse_dates=['date']) + return df.set_index('date')['obs_mm_h'] + + +def slice_forecast(df, member_cols, t0): + """Return (valid_times, member_array shape (lead, N)) for a single issue time.""" + sub = df[df['issue_time'] == t0].sort_values('lead_hour') + if len(sub) == 0: + raise ValueError(f"No forecast rows for issue_time={t0}") + valid_times = pd.to_datetime(sub['valid_time'].values) + member_arr = sub[member_cols].to_numpy(dtype=float) + return valid_times, member_arr + + +def plot_panel(ax, t0, df_da, df_ol, member_cols, obs_series, title_extra=""): + # Forecast trajectories + vt_da, q_da = slice_forecast(df_da, member_cols, t0) + vt_ol, q_ol = slice_forecast(df_ol, member_cols, t0) + + # Obs context window + ctx_start = t0 - pd.Timedelta(hours=CONTEXT_HOURS_BEFORE) + ctx_end = t0 + pd.Timedelta(hours=LEAD_HOURS_AFTER) + obs_window = obs_series.loc[ctx_start:ctx_end] + + # Open-loop members (drawn first so DA paints on top) + for j in range(q_ol.shape[1]): + ax.plot(vt_ol, q_ol[:, j], color=OL_COLOR, lw=0.6, alpha=0.45, zorder=2) + # DA members + for j in range(q_da.shape[1]): + ax.plot(vt_da, q_da[:, j], color=DA_COLOR, lw=0.6, alpha=0.45, zorder=3) + # Medians + ax.plot(vt_ol, np.nanmedian(q_ol, axis=1), color=OL_COLOR, lw=2.6, + linestyle="--", zorder=4, label="Open-loop median") + ax.plot(vt_da, np.nanmedian(q_da, axis=1), color=DA_COLOR, lw=2.6, + zorder=5, label="DA median") + # Obs + ax.plot(obs_window.index, obs_window.values, color=OBS_COLOR, lw=1.8, + zorder=6, label="Qkrig (obs)") + + # t0 marker + annotation + ax.axvline(t0, color="black", lw=1.0, linestyle=":", alpha=0.6, zorder=1) + obs_t0 = float(obs_series.get(t0, np.nan)) + obs_lead1 = float(obs_series.get(t0 + pd.Timedelta(hours=1), np.nan)) + da_lead1_mean = float(np.nanmean(q_da[0, :])) + ol_lead1_mean = float(np.nanmean(q_ol[0, :])) + + info_text = ( + f"t0 = {t0.strftime('%Y-%m-%d %H:%M')}\n" + f"obs(t0) = {obs_t0:.3f} mm/h\n" + f"obs(t0+1) = {obs_lead1:.3f} mm/h\n" + f"DA mean(t0+1) = {da_lead1_mean:.3f} mm/h\n" + f"OL mean(t0+1) = {ol_lead1_mean:.3f} mm/h" + ) + ax.text(0.02, 0.97, info_text, + transform=ax.transAxes, + fontsize=9, va='top', ha='left', family='monospace', + bbox=dict(boxstyle='round,pad=0.4', facecolor='white', + alpha=0.85, edgecolor='gray')) + + ax.set_title(title_extra, fontsize=12) + ax.set_ylabel("Discharge (mm/h)", fontsize=10) + ax.set_xlabel("Date", fontsize=10) + ax.xaxis.set_major_locator(mdates.HourLocator(interval=3)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %H:%M")) + ax.tick_params(axis='x', rotation=30, labelsize=8) + ax.grid(True, alpha=0.2) + ax.legend(loc='upper right', fontsize=9, frameon=True, framealpha=0.9) + + +def snap_to_nearest(t0, available_times): + """Return the available issue_time closest to t0.""" + times = pd.DatetimeIndex(available_times) + idx = np.abs((times - t0).total_seconds()).argmin() + snapped = times[idx] + if snapped != t0: + print(f" Snapped {t0} -> {snapped} (nearest available issue time)") + return snapped + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', default="cat-1016300") + parser.add_argument('--leadtime-dir', default=DEFAULT_LEADTIME_DIR, + help='Dir holding /_lead_time_forecasts_{da,openloop}.csv') + parser.add_argument('--da-dir', default=DEFAULT_DA_DIR, + help='Dir holding /_test_results.csv (for obs)') + parser.add_argument('--helene-t0', default=None, + help='Issue time for the Helene panel (default: 2024-09-26 12:00:00, ' + 'snapped to nearest available)') + parser.add_argument('--lowflow-t0', default=None, + help='Issue time for the low-flow panel (default: 2024-03-15 00:00:00, ' + 'snapped to nearest available)') + args = parser.parse_args() + + global CAT, LEADTIME_DIR, DA_DIR, OUT_PNG + CAT = args.cat_id + LEADTIME_DIR = args.leadtime_dir + DA_DIR = args.da_dir + OUT_PNG = os.path.join(LEADTIME_DIR, CAT, + f"{CAT}_issue_time_hydrograph_helene_vs_lowflow.png") + + da_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_da.csv") + ol_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_openloop.csv") + df_da, m_cols = load_forecasts(da_path) + df_ol, _ = load_forecasts(ol_path) + obs_series = load_obs() + + available = df_da['issue_time'].unique() + helene_t0 = snap_to_nearest( + pd.Timestamp(args.helene_t0) if args.helene_t0 else HELENE_T0, available) + lowflow_t0 = snap_to_nearest( + pd.Timestamp(args.lowflow_t0) if args.lowflow_t0 else pd.Timestamp("2024-03-15 00:00:00"), + available) + + fig, (ax_h, ax_l) = plt.subplots(1, 2, figsize=(20, 7)) + + plot_panel(ax_h, helene_t0, df_da, df_ol, m_cols, obs_series, + title_extra=f"Helene peak issue time — {helene_t0.strftime('%Y-%m-%d %H:%M')}") + plot_panel(ax_l, lowflow_t0, df_da, df_ol, m_cols, obs_series, + title_extra=f"Low-flow issue time — {lowflow_t0.strftime('%Y-%m-%d %H:%M')}") + + fig.suptitle( + f"Per-issue-time forecast hydrograph — {CAT}\n" + "Faint lines = 20 ensemble members. Thick lines = ensemble medians. " + "Vertical dotted line = t0 (DA stops here for both scenarios).", + fontsize=12, y=1.00, + ) + plt.tight_layout() + plt.savefig(OUT_PNG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4c_timeseries/plot_lead_time_reconstructed_timeseries.py b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4c_timeseries/plot_lead_time_reconstructed_timeseries.py new file mode 100644 index 00000000..a6d03b03 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4c_timeseries/plot_lead_time_reconstructed_timeseries.py @@ -0,0 +1,263 @@ +""" +Reconstructed time series at gauge 03463300, from our routed +lead-time forecasts. Mirrors the methodology of plot_timeseries_hourly_ensemble.py +(qSpatialAR) but feeds from the EnKF + T-route pipeline instead of the CNN. + +For each target hour t, pool ALL forecasts that land on t (across many +issue_times × many lead_hours × 20 members). Compute median + 5th/95th +percentile envelope. Plot DA and open-loop reconstructions overlaid on USGS +obs, with the Helene window shaded. + +Input parquets (from run_route_troute_leadtime.py): + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + Wide format: issue_time, lead_hour, member_00..member_19 (q in m³/s) + +Obs: + /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +Output: + /lead_time_reconstructed_timeseries.png + Default plot window: 2024-09-10 → 2024-10-10. + +Each scenario's curve is annotated with NSE vs USGS obs over the plot window. +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" +OBS_COLOR = "black" + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast_hardcoded_r_routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +# Watershed area at gauge 03463300 (sum of 21 catchments per GPKG hydrofabric) +WATERSHED_AREA_KM2 = 113.18 +# mm/h depth → m³/s: × 113.18 km² × 1000 / 3600 = 31.439 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +PLOT_START_DEFAULT = pd.Timestamp("2024-09-10 00:00:00") +PLOT_END_DEFAULT = pd.Timestamp("2024-10-10 23:00:00") + +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-30 23:00:00") + + +def load_parquet_long(path): + df = pd.read_parquet(path) + df['issue_time'] = pd.to_datetime(df['issue_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + if not member_cols: + raise ValueError(f"No member columns in {path}") + long = df.melt( + id_vars=['issue_time', 'lead_hour'], + value_vars=member_cols, + var_name='member', + value_name='q_m3s', + ) + long['valid_time'] = (long['issue_time'] + + pd.to_timedelta(long['lead_hour'], unit='h')) + return long + + +def reconstruct(long_df): + """For each valid_time, pool all (issue_time × lead × member) forecasts and + return median, p05, p95, count.""" + grouped = long_df.groupby('valid_time')['q_m3s'].agg( + median='median', + p05=lambda s: s.quantile(0.05), + p95=lambda s: s.quantile(0.95), + count='count', + ).reset_index() + return grouped + + +def load_usgs(usgs_csv): + """Load USGS gauge obs and ensure m³/s. + + The CSV at gauge 03463300 has the column `QObs(mm/h)` — catchment-averaged + depth, not gauge discharge in m³/s. Auto-convert when column name contains + 'mm': m³/s = mm/h × WATERSHED_AREA_KM2 × 1000 / 3600. + """ + df = pd.read_csv(usgs_csv) + date_col = next((c for c in df.columns + if c.lower() in ('datetime', 'date', 'time', 'timestamp')), None) + q_col = next((c for c in df.columns + if 'q' in c.lower() or 'flow' in c.lower() + or 'discharge' in c.lower()), None) + if not date_col or not q_col: + raise ValueError(f"Couldn't id date/Q columns in {usgs_csv}: {df.columns.tolist()}") + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if 'mm' in q_col.lower(): + series = series * MM_H_TO_M3_S + print(f" Converted obs from mm/h (column '{q_col}') to m³/s " + f"using area = {WATERSHED_AREA_KM2} km² (× {MM_H_TO_M3_S:.4f})") + else: + print(f" Loaded obs as-is from column '{q_col}' (assumed m³/s)") + return series + + +def nse(obs, sim): + m = np.isfinite(obs) & np.isfinite(sim) + if m.sum() < 2: + return np.nan + o, s = obs[m], sim[m] + denom = ((o - o.mean()) ** 2).sum() + if denom < 1e-10: + return np.nan + return 1.0 - float(((o - s) ** 2).sum() / denom) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--route-dir', default=DEFAULT_ROUTE_DIR) + parser.add_argument('--out-dir', default=None, + help='Default: --route-dir') + parser.add_argument('--usgs-csv', default=DEFAULT_USGS_CSV) + parser.add_argument('--plot-start', default=str(PLOT_START_DEFAULT)) + parser.add_argument('--plot-end', default=str(PLOT_END_DEFAULT)) + parser.add_argument('--da-name', default='routed_leadtime_da_full.parquet') + parser.add_argument('--ol-name', default='routed_leadtime_openloop_full.parquet') + args = parser.parse_args() + + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + plot_start = pd.Timestamp(args.plot_start) + plot_end = pd.Timestamp(args.plot_end) + + print(f"Loading DA parquet: {args.da_name}") + da_long = load_parquet_long(os.path.join(args.route_dir, args.da_name)) + print(f"Loading OL parquet: {args.ol_name}") + ol_long = load_parquet_long(os.path.join(args.route_dir, args.ol_name)) + print(f"Loading USGS obs: {args.usgs_csv}") + obs = load_usgs(args.usgs_csv) + + print("Reconstructing DA time series (overlapping-leads pool)...") + da_rec = reconstruct(da_long) + print(f" {len(da_rec)} unique target hours") + print("Reconstructing OL time series...") + ol_rec = reconstruct(ol_long) + + # Clip to plot window and join obs + da_rec = da_rec[(da_rec['valid_time'] >= plot_start) & (da_rec['valid_time'] <= plot_end)].copy() + ol_rec = ol_rec[(ol_rec['valid_time'] >= plot_start) & (ol_rec['valid_time'] <= plot_end)].copy() + da_rec['obs'] = da_rec['valid_time'].map(obs) + ol_rec['obs'] = ol_rec['valid_time'].map(obs) + + nse_da = nse(da_rec['obs'].values, da_rec['median'].values) + nse_ol = nse(ol_rec['obs'].values, ol_rec['median'].values) + print(f"NSE (plot window) — DA: {nse_da:.3f} | OL: {nse_ol:.3f}") + + # ---- Plot ---- + fig, ax = plt.subplots(figsize=(16, 5.5)) + fig.patch.set_facecolor("white") + + # Helene shaded + ax.axvspan(HELENE_START, HELENE_END, + color="firebrick", alpha=0.08, zorder=0) + ax.text(HELENE_START + pd.Timedelta(hours=12), 1.0, "Helene", + transform=ax.get_xaxis_transform(), + fontsize=9, color="firebrick", ha="left", va="top", + fontweight="bold") + + # Open-loop band + median + ax.fill_between(ol_rec['valid_time'], ol_rec['p05'], ol_rec['p95'], + color=OL_COLOR, alpha=0.18, zorder=2, + label="Open-loop [5th–95th pctile envelope]") + ax.plot(ol_rec['valid_time'], ol_rec['median'], + color=OL_COLOR, lw=1.2, linestyle='--', zorder=3, + label=f"Open-loop median (NSE={nse_ol:.3f})") + + # DA band + median + ax.fill_between(da_rec['valid_time'], da_rec['p05'], da_rec['p95'], + color=DA_COLOR, alpha=0.22, zorder=4, + label="DA [5th–95th pctile envelope]") + ax.plot(da_rec['valid_time'], da_rec['median'], + color=DA_COLOR, lw=1.4, zorder=5, + label=f"DA median (NSE={nse_da:.3f})") + + # Observed as dots + obs_window = obs.loc[plot_start:plot_end] + finite = np.isfinite(obs_window.values) + ax.scatter(obs_window.index[finite], obs_window.values[finite], + s=6, color=OBS_COLOR, linewidths=0, zorder=6, + label="USGS obs") + + ax.set_xlabel("Date", fontsize=11) + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_title( + f"Reconstructed forecast time series at USGS 03463300\n" + f"Overlapping-leads pool from EnKF forecast ensemble " + f"({plot_start.date()} – {plot_end.date()})", + fontsize=12, + ) + ax.set_xlim(plot_start, plot_end) + ax.xaxis.set_major_locator(mdates.WeekdayLocator(interval=1)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.25, lw=0.4) + ax.legend(loc='upper left', fontsize=9, frameon=True, framealpha=0.92) + + plt.tight_layout() + out_path = os.path.join(out_dir, "lead_time_reconstructed_timeseries.png") + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + # Also produce a Helene-zoom version + helene_zoom_start = pd.Timestamp("2024-09-24 00:00:00") + helene_zoom_end = pd.Timestamp("2024-09-29 23:00:00") + da_z = da_rec[(da_rec['valid_time'] >= helene_zoom_start) + & (da_rec['valid_time'] <= helene_zoom_end)] + ol_z = ol_rec[(ol_rec['valid_time'] >= helene_zoom_start) + & (ol_rec['valid_time'] <= helene_zoom_end)] + nse_da_z = nse(da_z['obs'].values, da_z['median'].values) + nse_ol_z = nse(ol_z['obs'].values, ol_z['median'].values) + print(f"NSE (Helene zoom) — DA: {nse_da_z:.3f} | OL: {nse_ol_z:.3f}") + + fig, ax = plt.subplots(figsize=(14, 5.5)) + fig.patch.set_facecolor("white") + ax.fill_between(ol_z['valid_time'], ol_z['p05'], ol_z['p95'], + color=OL_COLOR, alpha=0.18, zorder=2, + label="Open-loop [5th–95th]") + ax.plot(ol_z['valid_time'], ol_z['median'], color=OL_COLOR, + lw=1.4, linestyle='--', zorder=3, + label=f"Open-loop median (NSE={nse_ol_z:.3f})") + ax.fill_between(da_z['valid_time'], da_z['p05'], da_z['p95'], + color=DA_COLOR, alpha=0.22, zorder=4, + label="DA [5th–95th]") + ax.plot(da_z['valid_time'], da_z['median'], color=DA_COLOR, + lw=1.6, zorder=5, + label=f"DA median (NSE={nse_da_z:.3f})") + obs_z = obs.loc[helene_zoom_start:helene_zoom_end] + finite = np.isfinite(obs_z.values) + ax.scatter(obs_z.index[finite], obs_z.values[finite], + s=12, color=OBS_COLOR, linewidths=0, zorder=6, + label="USGS obs") + ax.set_xlabel("Date", fontsize=11) + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_title( + "Reconstructed forecast time series at USGS 03463300 — Helene zoom\n" + "Sep 24 – 29, 2024 (overlapping-leads pool from EnKF forecast ensemble)", + fontsize=12, + ) + ax.xaxis.set_major_locator(mdates.HourLocator(interval=12)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %H:%M")) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.25, lw=0.4) + ax.legend(loc='upper left', fontsize=9, frameon=True, framealpha=0.92) + plt.tight_layout() + out_zoom = os.path.join(out_dir, "lead_time_reconstructed_timeseries_helene.png") + plt.savefig(out_zoom, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_zoom}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4c_timeseries/run_4c_f4.sh b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4c_timeseries/run_4c_f4.sh new file mode 100644 index 00000000..d6187853 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder4_dynamic_variance_direct/4_evaluation/4c_timeseries/run_4c_f4.sh @@ -0,0 +1,50 @@ +#!/usr/bin/env bash +# F4 — Dynamic variance direct: 4c reconstructed timeseries plots +# +# Reads routed_leadtime_da_full.parquet + routed_leadtime_openloop_full.parquet +# from the F4 forecast route dir and writes two PNGs there. +# +# Usage: +# bash run_4c_f4.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +LEADTIME_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder4_dynamic_variance_direct +ROUTE_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder4_dynamic_variance_direct_leadtime_routed +OBS_DIR=/mnt/disk2/1400_sites_helene/catchment_ts_no_03463300_dynamic_variance +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +echo "[F4-4c] Reconstructed timeseries from: $ROUTE_DIR" +$TROUTE "$SCRIPT_DIR/plot_lead_time_reconstructed_timeseries.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" + +echo "[F4-4c] Forecast spaghetti..." +$TROUTE "$SCRIPT_DIR/plot_forecast_spaghetti.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" + +echo "[F4-4c] Helene vs low-flow issue-time hydrograph (per catchment)..." +for CAT in "${CATS[@]}"; do + DA_CSV="$LEADTIME_DIR/$CAT/${CAT}_lead_time_forecasts_da.csv" + [ -f "$DA_CSV" ] || { echo " [$CAT] no leadtime CSV — skipping"; continue; } + echo " [$CAT] helene hydrograph..." + $TROUTE "$SCRIPT_DIR/plot_helene_issue_time_hydrograph.py" \ + --cat-id "$CAT" \ + --leadtime-dir "$LEADTIME_DIR" \ + --obs-dir "$OBS_DIR" +done + +echo "[F4-4c] Done. Outputs in: $ROUTE_DIR" +ls "$ROUTE_DIR"/*.png 2>/dev/null || echo " (no PNGs in route dir)" diff --git a/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/batch_run_crossed_f5.sh b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/batch_run_crossed_f5.sh new file mode 100644 index 00000000..8848fae2 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/batch_run_crossed_f5.sh @@ -0,0 +1,65 @@ +#!/usr/bin/env bash +# Batch crossed ensemble for all 21 catchments — F5 re-kriging variance. +# R = krig_var directly (re-kriged obs variance); uses --direct-variance flag. +# Outputs _crossed_ensemble.parquet per catchment. +# +# Usage: +# nohup bash batch_run_crossed_f5.sh > ~/logs/crossed_f5.log 2>&1 & + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PYTHON=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT="$SCRIPT_DIR/run_crossed_ensemble.py" + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +OBS_DIR=/mnt/disk2/1400_sites_helene/catchment_ts_no_03463300_dynamic_variance_rekrig +CFE_DIR=/mnt/disk2/suma_helen_poster/cfe_py +CONFIG_FILE=/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder5_rekrig_variance_direct +FORCING_DIR1=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings +FORCING_DIR2=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings + +echo "F5 Re-kriging variance crossed ensemble — out: $OUT_DIR" +echo "" + +SKIP=0; DONE=0; FAIL=0 + +for CAT in "${CATS[@]}"; do + echo "===============================" + echo "=== $CAT ===" + echo "===============================" + + CAT_OUT="$OUT_DIR/$CAT" + + if [ ! -f "$CAT_OUT/${CAT}_best_params.json" ]; then + echo " WARNING: No best_params for $CAT — skipping" + SKIP=$((SKIP + 1)); continue + fi + + if [ -f "$CAT_OUT/${CAT}_crossed_ensemble.parquet" ]; then + echo " Crossed ensemble already exists — skipping" + SKIP=$((SKIP + 1)); continue + fi + + "$PYTHON" "$SCRIPT" \ + --cat-id "$CAT" \ + --forcing-dir "$FORCING_DIR1" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --out-dir "$OUT_DIR" \ + --test-forcing-dir1 "$FORCING_DIR1" \ + --test-forcing-dir2 "$FORCING_DIR2" \ + --direct-variance \ + && DONE=$((DONE + 1)) \ + || { echo " FAILED: $CAT"; FAIL=$((FAIL + 1)); } +done + +echo "" +echo "=== F5 crossed done: done=$DONE skipped=$SKIP failed=$FAIL ===" diff --git a/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/batch_run_da_on_all_cats.sh b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/batch_run_da_on_all_cats.sh new file mode 100644 index 00000000..3839a304 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/batch_run_da_on_all_cats.sh @@ -0,0 +1,80 @@ +#!/usr/bin/env bash +# Batch DA for all 21 catchments — F5 re-kriging variance. +# R = krig_var directly (re-kriged sigma^2 from Qkrig, no Vrugt scaling). +# Stages best_params.json from calibration results before running. +# +# Usage (from server, after SCPing this folder): +# bash batch_run_da_on_all_cats.sh + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PYTHON=/home/svyas/miniconda3/envs/troute/bin/python3 +DA_SCRIPT="$SCRIPT_DIR/run_perturbation_da_on.py" + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +OBS_DIR=/mnt/disk2/1400_sites_helene/catchment_ts_no_03463300_dynamic_variance_rekrig +BEST_PARAMS_SRC=/mnt/disk2/suma_helen_poster/catchment_results_1gauge_heldout +CFE_DIR=/mnt/disk2/suma_helen_poster/cfe_py +CONFIG_FILE=/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json +PARAM_BOUNDS=/mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder5_rekrig_variance_direct +FORCING_DIR1=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings +FORCING_DIR2=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings + +mkdir -p "$OUT_DIR" +echo "F5 Re-kriging variance DA — obs: $OBS_DIR" +echo " out: $OUT_DIR" +echo "" + +SKIP=0; DONE=0; FAIL=0 + +for CAT in "${CATS[@]}"; do + echo "===============================" + echo "=== $CAT ===" + echo "===============================" + + CAT_OUT="$OUT_DIR/$CAT" + mkdir -p "$CAT_OUT" + + if [ ! -f "$CAT_OUT/${CAT}_best_params.json" ]; then + SRC="$BEST_PARAMS_SRC/$CAT/${CAT}_best_params.json" + if [ -f "$SRC" ]; then + cp "$SRC" "$CAT_OUT/" + echo " Staged best_params" + else + echo " WARNING: No best_params for $CAT — skipping" + SKIP=$((SKIP + 1)); continue + fi + else + echo " best_params already staged" + fi + + if [ -f "$CAT_OUT/${CAT}_da_forcing_arm.csv" ] && \ + [ -f "$CAT_OUT/${CAT}_da_hydro_arm.csv" ]; then + echo " Arms already complete — skipping" + SKIP=$((SKIP + 1)); continue + fi + + "$PYTHON" "$DA_SCRIPT" \ + --cat-id "$CAT" \ + --forcing-dir "$FORCING_DIR1" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --param-bounds "$PARAM_BOUNDS" \ + --out-dir "$OUT_DIR" \ + --test-forcing-dir1 "$FORCING_DIR1" \ + --test-forcing-dir2 "$FORCING_DIR2" \ + --direct-variance \ + && DONE=$((DONE + 1)) \ + || { echo " FAILED: $CAT"; FAIL=$((FAIL + 1)); } +done + +echo "" +echo "=== F5 done: done=$DONE skipped=$SKIP failed=$FAIL ===" diff --git a/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/batch_run_leadtime_f5.sh b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/batch_run_leadtime_f5.sh new file mode 100644 index 00000000..84075e3c --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/batch_run_leadtime_f5.sh @@ -0,0 +1,74 @@ +#!/usr/bin/env bash +# Batch lead-time forecast sweep — F5 re-kriging variance, 1 gauge holdout. +# Runs run_lead_time_forecast_sweep.py for all 21 catchments with --direct-variance. +# Skips if both output CSVs already exist. +# +# After this completes, run route_leadtime_f5.sh to route through T-route. +# +# Usage: +# nohup bash batch_run_leadtime_f5.sh > ~/logs/leadtime_f5.log 2>&1 & + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PYTHON=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT="$SCRIPT_DIR/run_lead_time_forecast_sweep.py" + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +OBS_DIR=/mnt/disk2/1400_sites_helene/catchment_ts_no_03463300_dynamic_variance_rekrig +CFE_DIR=/mnt/disk2/suma_helen_poster/cfe_py +CONFIG_FILE=/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json +PARAM_BOUNDS=/mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder5_rekrig_variance_direct +FORCING_DIR1=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings +FORCING_DIR2=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings + +echo "F5 lead-time sweep (re-kriging variance) — out: $OUT_DIR" +echo "" + +SKIP=0; DONE=0; FAIL=0 + +for CAT in "${CATS[@]}"; do + echo "===============================" + echo "=== $CAT ===" + echo "===============================" + + CAT_OUT="$OUT_DIR/$CAT" + DA_CSV="$CAT_OUT/${CAT}_lead_time_forecasts_da.csv" + OL_CSV="$CAT_OUT/${CAT}_lead_time_forecasts_openloop.csv" + + if [ ! -f "$CAT_OUT/${CAT}_best_params.json" ]; then + echo " WARNING: No best_params for $CAT — skipping" + SKIP=$((SKIP + 1)); continue + fi + + if [ -f "$DA_CSV" ] && [ -f "$OL_CSV" ]; then + echo " Lead-time CSVs already exist — skipping" + SKIP=$((SKIP + 1)); continue + fi + + "$PYTHON" "$SCRIPT" \ + --cat-id "$CAT" \ + --forcing-dir "$FORCING_DIR1" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --param-bounds "$PARAM_BOUNDS" \ + --out-dir "$OUT_DIR" \ + --test-forcing-dir1 "$FORCING_DIR1" \ + --test-forcing-dir2 "$FORCING_DIR2" \ + --enkf-members 20 \ + --base-step-h 6 \ + --direct-variance \ + --prod-script "$SCRIPT_DIR/../../../1_distributed_cfe/calibrate_catchment_cfe_da_v2.py" \ + && DONE=$((DONE + 1)) \ + || { echo " FAILED: $CAT"; FAIL=$((FAIL + 1)); } +done + +echo "" +echo "=== F5 lead-time sweep done: done=$DONE skipped=$SKIP failed=$FAIL ===" diff --git a/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/batch_run_production_f5.sh b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/batch_run_production_f5.sh new file mode 100644 index 00000000..33a90e03 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/batch_run_production_f5.sh @@ -0,0 +1,72 @@ +#!/usr/bin/env bash +# Batch production per-member run — F5 re-kriging dynamic variance, 1 gauge holdout. +# Runs run_production_per_member.py for all 21 catchments with --direct-variance. +# Obs: re-kriging dynamic variance (different from F4's alpha-scaled variance). +# Skips if output CSV already exists. +# +# Usage: +# nohup bash batch_run_production_f5.sh > ~/logs/production_f5.log 2>&1 & + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PYTHON=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT="$SCRIPT_DIR/run_production_per_member.py" +PROD_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/1_distributed_cfe/calibrate_catchment_cfe_da_v2.py + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +OBS_DIR=/mnt/disk2/1400_sites_helene/catchment_ts_no_03463300_dynamic_variance_rekrig +CFE_DIR=/mnt/disk2/suma_helen_poster/cfe_py +CONFIG_FILE=/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json +PARAM_BOUNDS=/mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder5_rekrig_variance_direct +FORCING_DIR1=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings +FORCING_DIR2=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings + +echo "F5 production per-member (re-kriging direct variance) — out: $OUT_DIR" +echo "" + +SKIP=0; DONE=0; FAIL=0 + +for CAT in "${CATS[@]}"; do + echo "===============================" + echo "=== $CAT ===" + echo "===============================" + + CAT_OUT="$OUT_DIR/$CAT" + OUT_FILE="$CAT_OUT/${CAT}_production_per_member.csv" + + if [ ! -f "$CAT_OUT/${CAT}_best_params.json" ]; then + echo " WARNING: No best_params for $CAT — skipping" + SKIP=$((SKIP + 1)); continue + fi + + if [ -f "$OUT_FILE" ]; then + echo " Production per-member CSV already exists — skipping" + SKIP=$((SKIP + 1)); continue + fi + + "$PYTHON" "$SCRIPT" \ + --cat-id "$CAT" \ + --forcing-dir "$FORCING_DIR1" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --param-bounds "$PARAM_BOUNDS" \ + --out-dir "$OUT_DIR" \ + --test-forcing-dir1 "$FORCING_DIR1" \ + --test-forcing-dir2 "$FORCING_DIR2" \ + --enkf-members 20 \ + --direct-variance \ + --prod-script "$PROD_SCRIPT" \ + && DONE=$((DONE + 1)) \ + || { echo " FAILED: $CAT"; FAIL=$((FAIL + 1)); } +done + +echo "" +echo "=== F5 production per-member done: done=$DONE skipped=$SKIP failed=$FAIL ===" diff --git a/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/batch_run_sensitivity_f5.sh b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/batch_run_sensitivity_f5.sh new file mode 100644 index 00000000..24417ab2 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/batch_run_sensitivity_f5.sh @@ -0,0 +1,68 @@ +#!/usr/bin/env bash +# Batch perturbation sensitivity analysis — F5 re-kriged variance, 1 gauge holdout. +# Runs run_perturbation_sensitivity.py for all 21 catchments × 3 sources +# (init, forcing, process). Skips if output CSV already exists. +# +# Usage: +# nohup bash ~/da_1gauge_f5/2_assimilation/batch_run_sensitivity_f5.sh \ +# > ~/sensitivity_f5.log 2>&1 & + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PYTHON=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT="$SCRIPT_DIR/run_perturbation_sensitivity.py" + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +SOURCES=(init forcing process) + +OBS_DIR=/mnt/disk2/1400_sites_helene/catchment_ts_no_03463300_dynamic_variance_rekrig +CFE_DIR=/mnt/disk2/suma_helen_poster/cfe_py +CONFIG_FILE=/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json +PARAM_BOUNDS=/mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder5_rekrig_variance_direct +FORCING_DIR1=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings +FORCING_DIR2=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings + +echo "F5 sensitivity analysis — out: $OUT_DIR" +echo "" + +SKIP=0; DONE=0; FAIL=0 + +for CAT in "${CATS[@]}"; do + for SRC in "${SOURCES[@]}"; do + OUT_FILE="$OUT_DIR/$CAT/${CAT}_sensitivity_${SRC}.csv" + + if [ -f "$OUT_FILE" ]; then + echo " [$CAT/$SRC] already exists — skipping" + SKIP=$((SKIP + 1)) + continue + fi + + echo "===============================" + echo "=== $CAT source=$SRC ===" + echo "===============================" + + "$PYTHON" "$SCRIPT" \ + --cat-id "$CAT" \ + --source "$SRC" \ + --forcing-dir "$FORCING_DIR1" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --param-bounds "$PARAM_BOUNDS" \ + --out-dir "$OUT_DIR" \ + --test-forcing-dir1 "$FORCING_DIR1" \ + --test-forcing-dir2 "$FORCING_DIR2" \ + && DONE=$((DONE + 1)) \ + || { echo " FAILED: $CAT/$SRC"; FAIL=$((FAIL + 1)); } + done +done + +echo "" +echo "=== F5 sensitivity done: done=$DONE skipped=$SKIP failed=$FAIL ===" diff --git a/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/input_enkf_new.json b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/input_enkf_new.json new file mode 100644 index 00000000..482233bc --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/input_enkf_new.json @@ -0,0 +1 @@ +{"n": 1, "m": 1, "R": 0.07, "Q": 0, "smcmax" : 0.9394097311639178, "N": 1000, "P": 0.01, "D":2} \ No newline at end of file diff --git a/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/new_EnKF.py b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/new_EnKF.py new file mode 100644 index 00000000..ff945c73 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/new_EnKF.py @@ -0,0 +1,53 @@ +# EnKF.py +import numpy as np + +class EnKF: + def __init__(self, n, m, R, Q, N, P, smcmax, D): + self.n = n # number of states + self.m = m # number of measurements + self.R = R # measurement noise + self.Q = Q # process noise + self.N = N # number of ensembles + self.P = P + self.state_estimates = np.zeros((self.n)) + self.covariance_matrices = np.zeros((self.n, self.n)) + self.current_step = 0 + self.smcmax = smcmax + self.D = D + self.storage_max_m = self.smcmax * self.D + self.storage_init = self.storage_max_m * 0.667 + self.ensembles = np.full((self.n, self.N), self.storage_init) + + def predict(self, F_results): + error_factor_sim = 0.005 + for i in range(self.N): + perturbation_factor_sim = np.random.standard_normal() + self.ensembles[:, i] = F_results[i] + self.ensembles[:, i] += (perturbation_factor_sim * F_results[i] * error_factor_sim) + return self.ensembles + + def update(self, yi, H_results): + if np.any(np.isnan(yi)): + return np.mean(self.ensembles, axis=1) # return current state estimate without updating + + y_ensembles = np.zeros((self.m, self.N)) + error_factor_sm = 0.03 + for i in range(self.N): + perturbation_factor_sm = np.random.standard_normal() + y_ensembles[:, i] = H_results[i] + y_ensembles[:, i] += (perturbation_factor_sm * H_results[i] * error_factor_sm) + + ensemble_mean = np.mean(self.ensembles, axis=1)[:, np.newaxis] + y_ensembles_mean = np.mean(y_ensembles, axis=1)[:, np.newaxis] + Pxy = (self.ensembles - ensemble_mean).dot((y_ensembles - y_ensembles_mean).T) / (self.N - 1) + Pyy = np.cov(y_ensembles, bias=True) + K = Pxy.dot(np.linalg.pinv(Pyy + self.R)) + self.ensembles += K.dot(yi - y_ensembles_mean) + return np.mean(self.ensembles, axis=1) + + def get_state_estimate(self): + return np.mean(self.ensembles, axis=1) + + def get_covariance(self): + ensemble_mean = np.mean(self.ensembles, axis=1)[:, np.newaxis] + return (self.ensembles - ensemble_mean).dot((self.ensembles - ensemble_mean).T) / (self.N - 1) diff --git a/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/run_crossed_ensemble.py b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/run_crossed_ensemble.py new file mode 100644 index 00000000..bc1cec20 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/run_crossed_ensemble.py @@ -0,0 +1,463 @@ +""" +run_crossed_ensemble.py -- 4b: 600-member crossed ensemble for Helene. + +Crosses the two DA-on perturbation arms: + Forcing arm : 30 met draws (from run_perturbation_da_on.py Phase 2A) + Hydro arm : 20 state draws (from run_perturbation_da_on.py Phase 2B) + +For each Helene issue time t0: + Run 600 members = 30 forcing draws x 20 hydro-state draws. + Each member (i, j) gets: + - Initial state = perturbed DA analysis state (draw j of hydro arm) + - Met forcing = perturbed precip/PET sequence (draw i of forcing arm) + +This gives the full uncertainty envelope combining both sources — +the operational probabilistic forecast picture. + +Output (per catchment): + //_crossed_ensemble.parquet + Columns: issue_time, lead_hour, member_0000..member_0599 + +Usage: + python3 run_crossed_ensemble.py \\ + --cat-id cat-1016300 \\ + --forcing-dir \\ + --obs-dir /mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance \\ + --cfe-dir /mnt/disk2/suma_helen_poster/cfe_py \\ + --config-file /mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json \\ + --out-dir /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble \\ + --test-forcing-dir1 \\ + --test-forcing-dir2 + +Strategy: + Re-runs DA analysis (same as Phase 1 of run_perturbation_da_on.py) to get + the per-issue-time analysis state. Then generates N_FORCING sets of perturbed + forcing sequences and N_HYDRO perturbed states. Runs all 600 combinations. + + To avoid re-running Phase 1, if the Phase 1 snapshot file exists from a + prior run_perturbation_da_on.py run, it is loaded directly. + +NOTE: All special characters kept ASCII to avoid encoding issues on GPU. +""" + +import argparse +import os +import sys +import json +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +N_FORCING = 30 +N_HYDRO = 20 +N_MEMBERS = N_FORCING * N_HYDRO # 600 + +TEST_START = "2024-08-24 00:00:00" # 1 month spinup before Helene window +TEST_END = "2024-10-31 23:00:00" + +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-30 06:00:00") +FORECAST_HOURS = 18 + +PRECIP_SIGMA = 0.15 +PET_SIGMA = 0.10 +STATE_FRAC = 0.05 + +CAT_ID = None +OBS_FILE = None +CFE_CONFIG_FILE = None +TEST_FORCING_FILE = None +OUT_DIR = None +HARDCODED_R = None +DIRECT_VARIANCE = False # when True: R = krig_var directly (no Vrugt formula) +bmi_cfe = None + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(): + df = pd.read_csv(TEST_FORCING_FILE) + df["date"] = pd.to_datetime(df["time"]).dt.strftime("%Y-%m-%d %H:%M:%S") + df["total_precipitation"] = df["APCP_surface"] * 3600.0 + df["potential_evaporation"] = priestley_taylor_pet( + df["DSWRF_surface"].values, df["TMP_2maboveground"].values) + return df + + +def load_obs(): + df = pd.read_csv(OBS_FILE) + date_col = next((c for c in df.columns + if c.lower() in ("date", "time", "datetime", "timestamp")), None) + q_col = next((c for c in df.columns + if "qkrig" in c.lower() and "var" not in c.lower()), None) + var_col = next((c for c in df.columns if "var" in c.lower()), None) + df[date_col] = pd.to_datetime(df[date_col]) + df = df.sort_values(date_col).reset_index(drop=True) + date_strs = df[date_col].dt.strftime("%Y-%m-%d %H:%M:%S") + obs_dict = dict(zip(date_strs, df[q_col])) + var_vals = df[var_col] if var_col else pd.Series([0.07] * len(df)) + var_dict = dict(zip(date_strs, var_vals)) + return obs_dict, var_dict + + +def write_model_config(best_params): + """Write the BMI config file once; reused for every model instantiation.""" + with open(CFE_CONFIG_FILE) as f: + cfg = json.load(f) + cfg["forcing_file"] = TEST_FORCING_FILE + cfg["soil_params"]["bb"] = best_params["bb"] + cfg["soil_params"]["smcmax"] = best_params["smcmax"] + cfg["soil_params"]["satdk"] = best_params["satdk"] + cfg["slop"] = best_params["slop"] + cfg["max_gw_storage"] = best_params["max_gw_storage"] + cfg["expon"] = best_params["expon"] + cfg["Cgw"] = best_params["Cgw"] + cfg["K_lf"] = best_params["K_lf"] + cfg["K_nash"] = best_params["K_nash"] + cfg["partition_scheme"] = ("Schaake" if best_params["scheme"] <= 0.5 + else "Xinanjiang") + tmp_cfg = str(OUT_DIR / f"{CAT_ID}_bmi_config_crossed_stable.json") + with open(tmp_cfg, "w") as f: + json.dump(cfg, f) + return tmp_cfg + + +def make_model(cfg_path): + """Instantiate and initialize a CFE model from a pre-written config path.""" + m = bmi_cfe.BMI_CFE(cfg_file=cfg_path) + m.initialize() + return m + + +def get_state(m): + return { + "soil_m": m.soil_reservoir["storage_m"], + "gw_m": m.gw_reservoir["storage_m"], + "nash0": float(m.nash_storage[0]), + "nash1": float(m.nash_storage[1]), + } + + +def set_state(m, state): + m.soil_reservoir["storage_m"] = state["soil_m"] + m.gw_reservoir["storage_m"] = state["gw_m"] + m.nash_storage[0] = state["nash0"] + m.nash_storage[1] = state["nash1"] + + +def perturb_state(state, rng, frac=STATE_FRAC): + return { + "soil_m": max(state["soil_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "nash0": max(state["nash0"] * (1.0 + frac * rng.standard_normal()), 0.0), + "nash1": max(state["nash1"] * (1.0 + frac * rng.standard_normal()), 0.0), + } + + +def run_model_step(m, precip_mmh, pet_mmh): + m.set_value("atmosphere_water__time_integral_of_precipitation_mass_flux", + precip_mmh / 1000.0) + m.set_value("water_potential_evaporation_flux", + pet_mmh / 1000.0 / 3600.0) + m.update() + return m.get_value("land_surface_water__runoff_depth") * 1000.0 + + +def enkf_update(state, q_sim, y_obs, krig_var, rng): + if HARDCODED_R is not None: + R = HARDCODED_R + elif DIRECT_VARIANCE: + R = max(krig_var, 1e-6) + else: + R = max((0.10 * max(y_obs, 0.0)) ** 2 + 0.001 * krig_var, 1e-6) + R = max(R, 1e-6) + P_yy = max(q_sim * 0.01, 1e-6) + K = P_yy / (P_yy + R) + scale = K * (y_obs - q_sim) / max(abs(q_sim), 1e-6) + return { + "soil_m": max(state["soil_m"] * (1.0 + scale), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + scale), 1e-6), + "nash0": max(state["nash0"] + state["nash0"] * scale, 0.0), + "nash1": max(state["nash1"] + state["nash1"] * scale, 0.0), + } + + +def run_phase1_da(cfg_path, df_test, obs_dict, var_dict, rng): + """Run DA through test period; return snapshot dict {timestamp -> state}.""" + print(f" Phase 1: DA analysis through test period...") + model = make_model(cfg_path) + snapshots = {} + for _, row in df_test.iterrows(): + date_str = row["date"] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + q_sim = run_model_step(model, P, E) + y_obs = obs_dict.get(date_str, np.nan) + if not np.isnan(y_obs): + krig_var = var_dict.get(date_str, 0.07) + state = get_state(model) + updated = enkf_update(state, q_sim, y_obs, krig_var, rng) + set_state(model, updated) + date_ts = pd.Timestamp(date_str) + if HELENE_START <= date_ts <= HELENE_END: + snapshots[date_ts] = get_state(model) + model.finalize() + print(f" Phase 1 done: {len(snapshots)} issue-time snapshots") + return snapshots + + +def main(): + global CAT_ID, OBS_FILE, CFE_CONFIG_FILE, TEST_FORCING_FILE + global OUT_DIR, HARDCODED_R, DIRECT_VARIANCE, bmi_cfe + + parser = argparse.ArgumentParser() + parser.add_argument("--cat-id", required=True) + parser.add_argument("--forcing-dir", required=True) + parser.add_argument("--obs-dir", required=True) + parser.add_argument("--cfe-dir", required=True) + parser.add_argument("--config-file", required=True) + parser.add_argument("--out-dir", required=True) + parser.add_argument("--test-forcing-dir1", default=None) + parser.add_argument("--test-forcing-dir2", default=None) + parser.add_argument("--hardcoded-r", type=float, default=None, + help="Fixed R value for EnKF (e.g. 0.07). " + "Omit to use dynamic Vrugt R: R=(0.10*y)^2 + 0.001*krig_var") + parser.add_argument("--direct-variance", action="store_true", default=False, + help="Use krig_var directly as R (no Vrugt scaling)") + parser.add_argument("--rng-seed", type=int, default=None, + help="RNG seed for reproducibility (default: hash of cat-id)") + parser.add_argument("--n-forcing", type=int, default=N_FORCING, + help="Number of met forcing draws (default 30)") + parser.add_argument("--n-hydro", type=int, default=N_HYDRO, + help="Number of hydro-state draws (default 20)") + args = parser.parse_args() + + CAT_ID = args.cat_id + OBS_FILE = os.path.join(args.obs_dir, f"{CAT_ID}.csv") + CFE_CONFIG_FILE = args.config_file + HARDCODED_R = args.hardcoded_r + DIRECT_VARIANCE = args.direct_variance + OUT_DIR = Path(args.out_dir) / CAT_ID + OUT_DIR.mkdir(parents=True, exist_ok=True) + + n_fa = args.n_forcing + n_ha = args.n_hydro + n_total = n_fa * n_ha + print(f"[{CAT_ID}] Crossed ensemble: {n_fa} forcing x {n_ha} hydro = {n_total} members") + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + bmi_cfe = _bmi_cfe + + # Reuse combined forcing CSV from run_perturbation_da_on.py if it exists + # (avoids writing a duplicate large file to disk) + da_on_combined = (Path(args.out_dir).parent / "v2_perturbation_da_on" / + CAT_ID / f"{CAT_ID}_nwm_operational_combined.csv") + crossed_combined = OUT_DIR / f"{CAT_ID}_nwm_operational_combined.csv" + + if da_on_combined.exists(): + TEST_FORCING_FILE = str(da_on_combined) + print(f" Reusing combined forcing: {TEST_FORCING_FILE}") + elif crossed_combined.exists(): + TEST_FORCING_FILE = str(crossed_combined) + print(f" Reusing combined forcing: {TEST_FORCING_FILE}") + elif args.test_forcing_dir1 and args.test_forcing_dir2: + f1 = os.path.join(args.test_forcing_dir1, f"{CAT_ID}.csv") + f2 = os.path.join(args.test_forcing_dir2, f"{CAT_ID}.csv") + df1, df2 = pd.read_csv(f1), pd.read_csv(f2) + combined = (pd.concat([df1, df2], ignore_index=True) + .drop_duplicates(subset="time").sort_values("time")) + TEST_FORCING_FILE = str(crossed_combined) + combined.to_csv(TEST_FORCING_FILE, index=False) + print(f" Wrote combined forcing: {TEST_FORCING_FILE}") + else: + print("Need --test-forcing-dir1 and --test-forcing-dir2"); return + + best_params_file = OUT_DIR / f"{CAT_ID}_best_params.json" + if not best_params_file.exists(): + # Also check the DA results dir from run_perturbation_da_on.py + alt = Path(args.out_dir).parent / "v2_perturbation_da_on" / CAT_ID / f"{CAT_ID}_best_params.json" + if alt.exists(): + best_params_file = alt + else: + print(f"Missing best params at {best_params_file}"); return + with open(best_params_file) as f: + best_params = json.load(f)["best_parameters"] + + obs_dict, var_dict = load_obs() + df_forcing = load_test_forcing() + t_mask = (df_forcing["date"] >= TEST_START) & (df_forcing["date"] <= TEST_END) + df_test = df_forcing[t_mask].reset_index(drop=True) + + # Write BMI config once — reused for all 90,600 model instantiations + cfg_path = write_model_config(best_params) + print(f" BMI config written once: {cfg_path}") + + rng = np.random.default_rng(args.rng_seed if args.rng_seed is not None + else hash(CAT_ID) & 0x7fffffff) + + # Phase 1: get DA analysis snapshots + # Check if a prior run already saved them (skip re-running DA if so) + snap_cache = OUT_DIR / f"{CAT_ID}_da_snapshots.parquet" + if snap_cache.exists(): + print(f" Loading cached DA snapshots from {snap_cache}") + df_snaps = pd.read_parquet(snap_cache) + snapshots = { + pd.Timestamp(row["issue_time"]): { + "soil_m": row["soil_m"], "gw_m": row["gw_m"], + "nash0": row["nash0"], "nash1": row["nash1"], + } + for _, row in df_snaps.iterrows() + } + else: + snapshots = run_phase1_da(cfg_path, df_test, obs_dict, var_dict, rng) + + issue_times = sorted(snapshots.keys()) + t0_to_idx = {t: i for i, t in enumerate(df_test["date"].apply(pd.Timestamp))} + + # ------------------------------------------------------------------ # + # Crossed ensemble: for each issue_time, build n_fa forcing draws and # + # n_ha state draws, then run all n_fa x n_ha combinations. # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Running crossed ensemble ({n_total} members x " + f"{len(issue_times)} issue times)...") + + all_records = [] + all_hydro_draw_records = [] # (issue_time, hydro_draw_j, states) + all_forcing_draw_records = [] # (issue_time, forcing_draw_i, lead, P_pert, E_pert, scales) + + for t_idx, t0 in enumerate(issue_times): + snap = snapshots[t0] + t0_str = t0.strftime("%Y-%m-%d %H:%M:%S") + if t0 not in t0_to_idx: + continue + start_idx = t0_to_idx[t0] + + # Pre-draw n_fa forcing perturbation sequences + forcing_seqs = [] + for i in range(n_fa): + seq = [] + for lead in range(1, FORECAST_HOURS + 1): + fi = start_idx + lead + if fi >= len(df_test): + seq.append((0.0, 0.0)) + all_forcing_draw_records.append({ + "issue_time": t0_str, "forcing_draw_i": i, + "lead_hour": lead, + "P_pert_mmh": 0.0, "E_pert_mmh": 0.0, + "precip_scale": np.nan, "pet_scale": np.nan, + }) + continue + row = df_test.iloc[fi] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + mu_p = -0.5 * PRECIP_SIGMA ** 2 + P_pert = float(P * rng.lognormal(mu_p, PRECIP_SIGMA)) if P > 0 else 0.0 + E_pert = max(float(E * (1.0 + PET_SIGMA * rng.standard_normal())), 0.0) + seq.append((P_pert, E_pert)) + all_forcing_draw_records.append({ + "issue_time": t0_str, "forcing_draw_i": i, + "lead_hour": lead, + "P_pert_mmh": P_pert, "E_pert_mmh": E_pert, + "precip_scale": P_pert / P if P > 0 else np.nan, + "pet_scale": E_pert / E if E > 0 else np.nan, + }) + forcing_seqs.append(seq) + + # Pre-draw n_ha hydro-state perturbations + state_draws = [snap if j == 0 else perturb_state(snap, rng) + for j in range(n_ha)] + + # Record each hydro draw's actual initial state + for j, st in enumerate(state_draws): + all_hydro_draw_records.append({ + "issue_time": t0_str, "hydro_draw_j": j, + "soil_m": st["soil_m"], "gw_m": st["gw_m"], + "nash0": st["nash0"], "nash1": st["nash1"], + }) + + # q_matrix shape: (FORECAST_HOURS, n_fa, n_ha) + q_matrix = np.full((FORECAST_HOURS, n_fa, n_ha), np.nan) + + for i, f_seq in enumerate(forcing_seqs): + for j, state_j in enumerate(state_draws): + m = make_model(cfg_path) + set_state(m, state_j) + for lead in range(1, FORECAST_HOURS + 1): + P_pert, E_pert = f_seq[lead - 1] + q = run_model_step(m, P_pert, E_pert) + q_matrix[lead - 1, i, j] = q + m.finalize() + + # Flatten 600-member matrix into rows + for lead in range(1, FORECAST_HOURS + 1): + row_data = {"issue_time": t0_str, "lead_hour": lead} + mem_idx = 0 + for i in range(n_fa): + for j in range(n_ha): + row_data[f"member_{mem_idx:04d}"] = q_matrix[lead - 1, i, j] + mem_idx += 1 + all_records.append(row_data) + + if (t_idx + 1) % 10 == 0 or (t_idx + 1) == len(issue_times): + print(f" {t_idx + 1}/{len(issue_times)} issue times complete") + + df_out = pd.DataFrame(all_records) + out_path = OUT_DIR / f"{CAT_ID}_crossed_ensemble.parquet" + df_out.to_parquet(out_path, index=False) + print(f"Saved: {out_path}") + print(f" Shape: {df_out.shape} " + f"({df_out['issue_time'].nunique()} issue times x " + f"{FORECAST_HOURS} leads x {n_total} members)") + + # ------------------------------------------------------------------ # + # Provenance files — enable full ensemble traceability and restart # + # ------------------------------------------------------------------ # + + # 1. DA analysis snapshots — unperturbed state at each issue_time. + # Load any row and set_state() to restart from that init time. + snap_records = [ + {"issue_time": t.strftime("%Y-%m-%d %H:%M:%S"), **s} + for t, s in snapshots.items() + ] + snap_path = OUT_DIR / f"{CAT_ID}_da_snapshots.parquet" + pd.DataFrame(snap_records).to_parquet(snap_path, index=False) + print(f"Saved: {snap_path} ({len(snap_records)} snapshots)") + + # 2. Member manifest — decoder ring: member_col -> (forcing_draw_i, hydro_draw_j). + # member_k uses forcing draw k//n_ha and hydro draw k%n_ha. + manifest = [ + {"member": f"member_{k:04d}", "member_idx": k, + "forcing_draw_i": k // n_ha, "hydro_draw_j": k % n_ha} + for k in range(n_total) + ] + manifest_path = OUT_DIR / f"{CAT_ID}_member_manifest.csv" + pd.DataFrame(manifest).to_csv(manifest_path, index=False) + print(f"Saved: {manifest_path} ({n_total} members)") + + # 3. Hydro draw states — actual perturbed initial state for each of + # n_ha draws at every issue_time. Links member -> initial conditions. + hydro_path = OUT_DIR / f"{CAT_ID}_hydro_draw_states.parquet" + pd.DataFrame(all_hydro_draw_records).to_parquet(hydro_path, index=False) + print(f"Saved: {hydro_path} ({len(all_hydro_draw_records)} rows)") + + # 4. Forcing draw sequences — actual P/PET perturbation values for each + # of n_fa draws at every (issue_time, lead_hour). Links member -> forcing. + forcing_path = OUT_DIR / f"{CAT_ID}_forcing_draw_sequences.parquet" + pd.DataFrame(all_forcing_draw_records).to_parquet(forcing_path, index=False) + print(f"Saved: {forcing_path} ({len(all_forcing_draw_records)} rows)") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/run_lead_time_forecast_sweep.py b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/run_lead_time_forecast_sweep.py new file mode 100644 index 00000000..7b93772b --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/run_lead_time_forecast_sweep.py @@ -0,0 +1,389 @@ +""" +Forecast lead-time evaluation for F5 (re-kriged variance direct). + +For each issue-time t0 in a sampling schedule across the test period, fork the +ensemble at t0 and free-run 18 hours of forecast. Two parallel trajectories are +maintained throughout the test period: + + da: Production setup. DA on, forcing perturbed, process noise on. + R = per-hour kriging variance from the re-kriged obs file. + openloop: No DA from t=0. Forcing perturbed, process noise on. + +At each issue time t0, the state of each scenario's ensemble is copied into a +forecast ensemble that free-runs 18 hours with: + - DA off + - process noise off + - forcing perturbed (lognormal precip, Gaussian PET) + +Issue-time schedule: + - Base cadence: every --base-step-h hours (default 6h) + - Densified to hourly across the Helene window (2024-09-24 → 2024-09-28) + +Outputs (stacked across all issue times): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv + Columns: issue_time, lead_hour, valid_time, member_00 .. member_19 (mm/h) +""" + +import argparse +import os +import sys +import json +import importlib.util +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +SPINUP_START = "2023-02-01 00:00:00" +SPINUP_END = "2023-09-30 23:00:00" +TEST_START = "2023-10-01 00:00:00" +TEST_END = "2024-10-31 23:00:00" + +FORECAST_LEAD_HOURS = 18 + +DEFAULT_DENSE_START = "2024-09-24 00:00:00" +DEFAULT_DENSE_END = "2024-09-28 23:00:00" + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(test_forcing_file): + df = pd.read_csv(test_forcing_file) + df['date'] = pd.to_datetime(df['time']).dt.strftime('%Y-%m-%d %H:%M:%S') + df['total_precipitation'] = df['APCP_surface'] * 3600.0 + df['potential_evaporation'] = priestley_taylor_pet( + df['DSWRF_surface'].values, df['TMP_2maboveground'].values + ) + return df + + +def import_enkf_class(script_path): + spec = importlib.util.spec_from_file_location("prod_v2", script_path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod.EnKFAssimilator + + +def snapshot_states(models): + return [ + { + "soil_m": float(m.soil_reservoir["storage_m"]), + "gw_m": float(m.gw_reservoir["storage_m"]), + "nash0_m": float(m.nash_storage[0]), + "nash1_m": float(m.nash_storage[1]), + } + for m in models + ] + + +def restore_states(models, snapshot): + for m, s in zip(models, snapshot): + m.soil_reservoir["storage_m"] = s["soil_m"] + m.gw_reservoir["storage_m"] = s["gw_m"] + m.nash_storage[0] = s["nash0_m"] + m.nash_storage[1] = s["nash1_m"] + + +def build_issue_time_schedule(dates_list, base_step_h, dense_start, dense_end): + dates_dt = pd.to_datetime(dates_list) + selected = set() + base_mask = (np.arange(len(dates_dt)) % base_step_h == 0) + for d in dates_dt[base_mask]: + selected.add(d.strftime('%Y-%m-%d %H:%M:%S')) + if dense_start is not None and dense_end is not None: + ds = pd.Timestamp(dense_start) + de = pd.Timestamp(dense_end) + dense_mask = (dates_dt >= ds) & (dates_dt <= de) + for d in dates_dt[dense_mask]: + selected.add(d.strftime('%Y-%m-%d %H:%M:%S')) + return sorted(selected) + + +def build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, enkf, + apply_init_perturbation=True, member_zero_clean=True): + def custom_load_forcing(self_cfe): + df = load_test_forcing(test_forcing_file) + self_cfe.forcing_data = df.rename(columns={"date": "time"}) + + models = [] + for i in range(N): + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.load_forcing_file = custom_load_forcing.__get__(m) + m.initialize() + if apply_init_perturbation and not (member_zero_clean and i == 0): + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + r = enkf.init_state_perturb_frac + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1 + r * enkf.rng.standard_normal()), 1e-6, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1 + r * enkf.rng.standard_normal()), 1e-6, gw_max)) + jitter = max(sm_max * 1e-6, 1e-9) + m.nash_storage[0] = max(n00 + jitter * enkf.rng.standard_normal(), 0.0) + m.nash_storage[1] = max(n10 + jitter * enkf.rng.standard_normal(), 0.0) + models.append(m) + return models + + +def step_ensemble(models, p_arr, e_arr): + q = np.empty(len(models), dtype=float) + for i, m in enumerate(models): + m.set_value('atmosphere_water__time_integral_of_precipitation_mass_flux', + float(p_arr[i]) / 1000.0) + m.set_value('water_potential_evaporation_flux', + float(e_arr[i]) / 1000.0 / 3600.0) + m.update() + q[i] = m.get_value('land_surface_water__runoff_depth') * 1000.0 + return q + + +def run_forecast(fcst_models, enkf_fcst, dates_list, forcing_by_date, + t0_idx, n_lead): + N = len(fcst_models) + q_matrix = np.full((n_lead, N), np.nan, dtype=float) + lead_hours = [] + valid_times = [] + for lead in range(1, n_lead + 1): + idx = t0_idx + lead + if idx >= len(dates_list): + break + valid_date = dates_list[idx] + p, e = forcing_by_date[valid_date] + p_arr, e_arr = enkf_fcst.perturb_forcing(p, e) + q = step_ensemble(fcst_models, p_arr, e_arr) + q_matrix[lead - 1, :] = q + lead_hours.append(lead) + valid_times.append(valid_date) + return lead_hours, valid_times, q_matrix + + +def apply_direct_variance(enkf_instances, obs_file): + """Override obs_var_dict with per-hour kriging variance from the re-kriged obs file.""" + obs_df = pd.read_csv(obs_file) + t_col = next(c for c in obs_df.columns + if c.lower() in ('datetime', 'date', 'time', 'timestamp')) + v_col = next(c for c in obs_df.columns if 'var' in c.lower()) + obs_df[t_col] = pd.to_datetime(obs_df[t_col]) + obs_df = obs_df.set_index(t_col).sort_index() + var_dict = { + str(ts): float(v) + for ts, v in obs_df[v_col].items() + if pd.notna(v) and float(v) > 0 + } + for enkf in enkf_instances: + enkf.obs_var_dict = dict(var_dict) + print(f"[lead-time] R = direct kriging variance ({len(var_dict)} timesteps, " + f"range {min(var_dict.values()):.3e} – {max(var_dict.values()):.3e})") + + +def run(args, bmi_cfe, EnKFAssimilator): + cat_id = args.cat_id + out_dir = Path(args.out_dir) / cat_id + out_dir.mkdir(parents=True, exist_ok=True) + + f1 = os.path.join(args.test_forcing_dir1, f'{cat_id}.csv') + f2 = os.path.join(args.test_forcing_dir2, f'{cat_id}.csv') + df1 = pd.read_csv(f1) + df2 = pd.read_csv(f2) + combined = pd.concat([df1, df2], ignore_index=True) + combined = combined.drop_duplicates(subset='time').sort_values('time') + test_forcing_file = str(out_dir / f'{cat_id}_nwm_operational_combined.csv') + combined.to_csv(test_forcing_file, index=False) + + best_params_file = out_dir / f'{cat_id}_best_params.json' + if not best_params_file.exists(): + print(f"Missing {best_params_file}. Pre-stage from calibration run.") + return + with open(best_params_file) as f: + best = json.load(f)['best_parameters'] + + with open(args.config_file) as f: + cfg = json.load(f) + cfg['forcing_file'] = test_forcing_file + cfg['soil_params']['bb'] = best['bb'] + cfg['soil_params']['smcmax'] = best['smcmax'] + cfg['soil_params']['satdk'] = best['satdk'] + cfg['slop'] = best['slop'] + cfg['max_gw_storage'] = best['max_gw_storage'] + cfg['expon'] = best['expon'] + cfg['Cgw'] = best['Cgw'] + cfg['K_lf'] = best['K_lf'] + cfg['K_nash'] = best['K_nash'] + cfg['partition_scheme'] = "Schaake" if best['scheme'] <= 0.5 else "Xinanjiang" + tmp_cfg = str(out_dir / f'{cat_id}_bmi_config_temp_leadtime.json') + with open(tmp_cfg, 'w') as f: + json.dump(cfg, f) + + obs_file = os.path.join(args.obs_dir, f'{cat_id}.csv') + seed_base = args.rng_seed if args.rng_seed is not None else 0 + enkf_da = EnKFAssimilator( + n_members=args.enkf_members, obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, use_vrugt_r=(not args.no_vrugt_r), + vrugt_alpha=args.vrugt_alpha, vrugt_scale=args.vrugt_scale, + rng_seed=seed_base if args.rng_seed is not None else None, + ) + enkf_ol = EnKFAssimilator( + n_members=args.enkf_members, obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, use_vrugt_r=(not args.no_vrugt_r), + vrugt_alpha=args.vrugt_alpha, vrugt_scale=args.vrugt_scale, + rng_seed=(seed_base + 1) if args.rng_seed is not None else None, + ) + enkf_fcst = EnKFAssimilator( + n_members=args.enkf_members, obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, use_vrugt_r=(not args.no_vrugt_r), + vrugt_alpha=args.vrugt_alpha, vrugt_scale=args.vrugt_scale, + rng_seed=(seed_base + 2) if args.rng_seed is not None else None, + ) + + if args.hardcoded_r is not None: + for enkf in (enkf_da, enkf_ol, enkf_fcst): + for date in enkf.obs_var_dict: + enkf.obs_var_dict[date] = args.hardcoded_r + print(f"[lead-time] R hardcoded to {args.hardcoded_r}") + + if args.direct_variance: + apply_direct_variance((enkf_da, enkf_ol, enkf_fcst), obs_file) + + N = enkf_da.n_members + print(f"[lead-time] {cat_id} | N={N} | direct_variance={args.direct_variance} | " + f"forecast lead = {FORECAST_LEAD_HOURS}h") + + prod_models = build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, enkf_da) + openloop_models = build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, enkf_ol) + fcst_models = build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, + enkf_fcst, apply_init_perturbation=False) + + df = load_test_forcing(test_forcing_file) + sp_mask = (df['date'] >= SPINUP_START) & (df['date'] <= SPINUP_END) + df_sp = df[sp_mask] + print(f"[lead-time] spinup: {len(df_sp)} hours") + for p, e in zip(df_sp['total_precipitation'], df_sp['potential_evaporation']): + p_da, e_da = enkf_da.perturb_forcing(p, e) + step_ensemble(prod_models, p_da, e_da) + enkf_da.add_process_noise(prod_models) + p_ol, e_ol = enkf_ol.perturb_forcing(p, e) + step_ensemble(openloop_models, p_ol, e_ol) + enkf_ol.add_process_noise(openloop_models) + + t_mask = (df['date'] >= TEST_START) & (df['date'] <= TEST_END) + df_test = df[t_mask].reset_index(drop=True) + dates_list = list(df_test['date'].values) + forcing_by_date = dict(zip( + df_test['date'].values, + zip(df_test['total_precipitation'].values, + df_test['potential_evaporation'].values))) + + issue_times = build_issue_time_schedule( + dates_list, args.base_step_h, args.dense_start, args.dense_end) + issue_set = set(issue_times) + print(f"[lead-time] test period: {len(dates_list)} hours | " + f"{len(issue_times)} issue times " + f"(base={args.base_step_h}h, dense={args.dense_start or 'none'}" + f"..{args.dense_end or 'none'})") + + da_rows = [] + ol_rows = [] + + for h, current_date in enumerate(dates_list): + p, e = forcing_by_date[current_date] + + p_da, e_da = enkf_da.perturb_forcing(p, e) + q_da = step_ensemble(prod_models, p_da, e_da) + enkf_da.update_states(prod_models, current_date, q_da) + enkf_da.add_process_noise(prod_models) + + p_ol, e_ol = enkf_ol.perturb_forcing(p, e) + q_ol = step_ensemble(openloop_models, p_ol, e_ol) + enkf_ol.add_process_noise(openloop_models) + + if current_date in issue_set: + for scenario_label, src_models, rows_acc in [ + ('da', prod_models, da_rows), + ('openloop', openloop_models, ol_rows), + ]: + snap = snapshot_states(src_models) + restore_states(fcst_models, snap) + leads, valids, qm = run_forecast( + fcst_models, enkf_fcst, + dates_list, forcing_by_date, + t0_idx=h, n_lead=FORECAST_LEAD_HOURS) + for k, (lead, vt) in enumerate(zip(leads, valids)): + rows_acc.append((current_date, lead, vt, *qm[k, :])) + + for m in prod_models + openloop_models + fcst_models: + m.finalize() + + cols = (['issue_time', 'lead_hour', 'valid_time'] + + [f'member_{i:02d}' for i in range(N)]) + da_path = out_dir / f'{cat_id}_lead_time_forecasts_da.csv' + ol_path = out_dir / f'{cat_id}_lead_time_forecasts_openloop.csv' + pd.DataFrame(da_rows, columns=cols).to_csv(da_path, index=False) + pd.DataFrame(ol_rows, columns=cols).to_csv(ol_path, index=False) + print(f"[lead-time] saved {da_path}") + print(f"[lead-time] saved {ol_path}") + + sched_path = out_dir / f'{cat_id}_lead_time_issue_times.csv' + pd.DataFrame({'issue_time': issue_times}).to_csv(sched_path, index=False) + print(f"[lead-time] saved {sched_path}") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', required=True) + parser.add_argument('--forcing-dir', required=True) + parser.add_argument('--obs-dir', required=True) + parser.add_argument('--cfe-dir', required=True) + parser.add_argument('--config-file', required=True) + parser.add_argument('--param-bounds', required=True) + parser.add_argument('--out-dir', required=True) + parser.add_argument('--test-forcing-dir1', required=True) + parser.add_argument('--test-forcing-dir2', required=True) + parser.add_argument('--enkf-members', type=int, default=20) + parser.add_argument('--enkf-obs-error-std', type=float, default=0.05) + parser.add_argument('--no-vrugt-r', action='store_true') + parser.add_argument('--direct-variance', action='store_true', default=False, + help='Use per-hour kriging variance column from obs file as R ' + '(overrides Vrugt formula). Required for F5.') + parser.add_argument('--hardcoded-r', type=float, default=None) + parser.add_argument('--vrugt-alpha', type=float, default=0.10) + parser.add_argument('--vrugt-scale', type=float, default=0.001) + parser.add_argument('--rng-seed', type=int, default=None) + parser.add_argument('--base-step-h', type=int, default=6) + parser.add_argument('--dense-start', type=str, default=DEFAULT_DENSE_START) + parser.add_argument('--dense-end', type=str, default=DEFAULT_DENSE_END) + parser.add_argument('--prod-script', default=None) + args = parser.parse_args() + + here = os.path.dirname(os.path.abspath(__file__)) + prod_script = args.prod_script or os.path.normpath( + os.path.join(here, '..', '..', '..', '1_distributed_cfe', + 'calibrate_catchment_cfe_da_v2.py')) + if not os.path.exists(prod_script): + raise FileNotFoundError( + f"Could not find production script at {prod_script}. " + f"Pass --prod-script to override.") + EnKFAssimilator = import_enkf_class(prod_script) + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + + run(args, _bmi_cfe, EnKFAssimilator) + + +if __name__ == '__main__': + main() diff --git a/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/run_perturbation_da_on.py b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/run_perturbation_da_on.py new file mode 100644 index 00000000..6ef27e96 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/run_perturbation_da_on.py @@ -0,0 +1,400 @@ +""" +run_perturbation_da_on.py — 2a/2b: ensemble spread with DA ON. + +Runs two separated perturbation arms for the Helene window (Sep 24-30), +with DA actively running during the analysis period: + + ARM A — forcing only (30 members): + DA analyzes states using Qkrig at every hour (DA ON). + At each issue_time: fork 30 members with ONLY met forcing perturbed + (different precip/PET draw per member). Hydro states = DA analysis + mean (same for all 30). Free-run 18 hours, no further DA. + + ARM B — hydro states only (20 members): + Same DA analysis. At each issue_time: fork 20 members with ONLY + hydro states perturbed (multiplicative noise on soil/GW/Nash from + the DA analysis ensemble). Forcing = deterministic. Free-run 18 hours. + +These two arms isolate "where does forecast spread come from WITH DA?" +— figures for the assimilation script section. + +Output (per catchment, per arm): + //_da_forcing_arm.csv + //_da_hydro_arm.csv + Columns: issue_time, lead_hour, member_00..member_29 (forcing arm) + issue_time, lead_hour, member_00..member_19 (hydro arm) + +Usage (run for each catchment, both arms): + python3 run_perturbation_da_on.py \\ + --cat-id cat-1016300 \\ + --forcing-dir \\ + --obs-dir /mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance \\ + --cfe-dir /mnt/disk2/suma_helen_poster/cfe_py \\ + --config-file /mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json \\ + --param-bounds /mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json \\ + --out-dir /mnt/disk2/suma_helen_poster/da_results/v2_perturbation_da_on \\ + --test-forcing-dir1 \\ + --test-forcing-dir2 \\ + --hardcoded-r 0.07 +""" + +import argparse +import os +import sys +import json +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +N_FORCING = 30 # members in forcing arm +N_HYDRO = 20 # members in hydro-state arm + +# Analysis period — run DA through entire test period +TEST_START = "2024-08-24 00:00:00" # 1 month spinup before Helene window +TEST_END = "2024-10-31 23:00:00" + +# Issue times for the Helene window (hourly, Sep 24-30) +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-30 06:00:00") +FORECAST_HOURS = 18 + +# Perturbation magnitudes (same as production) +PRECIP_SIGMA = 0.15 +PET_SIGMA = 0.10 +STATE_FRAC = 0.05 # hydro-state perturbation fraction + +# Set at runtime +CAT_ID = None +OBS_FILE = None +CFE_CONFIG_FILE = None +TEST_FORCING_FILE = None +OUT_DIR = None +HARDCODED_R = None +DIRECT_VARIANCE = False +bmi_cfe = None + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(): + df = pd.read_csv(TEST_FORCING_FILE) + df["date"] = pd.to_datetime(df["time"]).dt.strftime("%Y-%m-%d %H:%M:%S") + df["total_precipitation"] = df["APCP_surface"] * 3600.0 + df["potential_evaporation"] = priestley_taylor_pet( + df["DSWRF_surface"].values, df["TMP_2maboveground"].values) + return df + + +def load_obs(): + df = pd.read_csv(OBS_FILE) + date_col = next((c for c in df.columns + if c.lower() in ("date", "time", "datetime", "timestamp")), None) + if date_col is None: + raise ValueError(f"No date column in {OBS_FILE}. Columns: {df.columns.tolist()}") + q_col = next((c for c in df.columns if "qkrig" in c.lower() and "var" not in c.lower()), None) + if q_col is None: + raise ValueError(f"No qkrig column in {OBS_FILE}. Columns: {df.columns.tolist()}") + var_col = next((c for c in df.columns if "var" in c.lower()), None) + df[date_col] = pd.to_datetime(df[date_col]) + df = df.sort_values(date_col).reset_index(drop=True) + date_strs = df[date_col].dt.strftime("%Y-%m-%d %H:%M:%S") + obs_dict = dict(zip(date_strs, df[q_col])) + var_vals = df[var_col] if var_col else pd.Series([0.07] * len(df)) + var_dict = dict(zip(date_strs, var_vals)) + print(f" Obs: date='{date_col}', q='{q_col}', var='{var_col}'") + return obs_dict, var_dict + + +def make_model(best_params): + with open(CFE_CONFIG_FILE) as f: + cfg = json.load(f) + cfg["forcing_file"] = TEST_FORCING_FILE + cfg["soil_params"]["bb"] = best_params["bb"] + cfg["soil_params"]["smcmax"] = best_params["smcmax"] + cfg["soil_params"]["satdk"] = best_params["satdk"] + cfg["slop"] = best_params["slop"] + cfg["max_gw_storage"] = best_params["max_gw_storage"] + cfg["expon"] = best_params["expon"] + cfg["Cgw"] = best_params["Cgw"] + cfg["K_lf"] = best_params["K_lf"] + cfg["K_nash"] = best_params["K_nash"] + cfg["partition_scheme"] = ("Schaake" if best_params["scheme"] <= 0.5 + else "Xinanjiang") + tmp_cfg = str(OUT_DIR / f"{CAT_ID}_bmi_config_temp_da_on.json") + with open(tmp_cfg, "w") as f: + json.dump(cfg, f) + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.initialize() + return m + + +def get_state(m): + return { + "soil_m": m.soil_reservoir["storage_m"], + "gw_m": m.gw_reservoir["storage_m"], + "nash0": float(m.nash_storage[0]), + "nash1": float(m.nash_storage[1]), + } + + +def set_state(m, state): + m.soil_reservoir["storage_m"] = state["soil_m"] + m.gw_reservoir["storage_m"] = state["gw_m"] + m.nash_storage[0] = state["nash0"] + m.nash_storage[1] = state["nash1"] + + +def perturb_state(state, rng, frac=STATE_FRAC): + sm_max = 1.0 # placeholder — will be clipped in set_state; actual max enforced by model + out = { + "soil_m": max(state["soil_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "nash0": max(state["nash0"] * (1.0 + frac * rng.standard_normal()), 0.0), + "nash1": max(state["nash1"] * (1.0 + frac * rng.standard_normal()), 0.0), + } + return out + + +def run_model_step(m, precip_mmh, pet_mmh): + m.set_value("atmosphere_water__time_integral_of_precipitation_mass_flux", + precip_mmh / 1000.0) + m.set_value("water_potential_evaporation_flux", + pet_mmh / 1000.0 / 3600.0) + m.update() + return m.get_value("land_surface_water__runoff_depth") * 1000.0 # m/h -> mm/h + + +def enkf_update(state, q_sim, y_obs, krig_var, rng, n_members=1): + """Simple scalar EnKF update on the mean analysis state.""" + if HARDCODED_R is not None: + R = HARDCODED_R + elif DIRECT_VARIANCE: + R = max(krig_var, 1e-6) + else: + alpha = 0.10 + R = (alpha * max(y_obs, 0.0)) ** 2 + 0.001 * krig_var + R = max(R, 1e-6) + P_yy = max(q_sim * 0.01, 1e-6) # rough prior variance on Q + K = P_yy / (P_yy + R) + innovation = y_obs - q_sim + # Apply correction to each state proportionally (simple scalar gain) + scale = K * innovation / max(abs(q_sim), 1e-6) + updated = { + "soil_m": max(state["soil_m"] * (1.0 + scale), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + scale), 1e-6), + "nash0": max(state["nash0"] + state["nash0"] * scale, 0.0), + "nash1": max(state["nash1"] + state["nash1"] * scale, 0.0), + } + return updated + + +def main(): + global CAT_ID, OBS_FILE, CFE_CONFIG_FILE, TEST_FORCING_FILE + global OUT_DIR, HARDCODED_R, DIRECT_VARIANCE, bmi_cfe + + parser = argparse.ArgumentParser() + parser.add_argument("--cat-id", required=True) + parser.add_argument("--forcing-dir", required=True) + parser.add_argument("--obs-dir", required=True) + parser.add_argument("--cfe-dir", required=True) + parser.add_argument("--config-file", required=True) + parser.add_argument("--param-bounds", required=True) + parser.add_argument("--out-dir", required=True) + parser.add_argument("--test-forcing-dir1", default=None) + parser.add_argument("--test-forcing-dir2", default=None) + parser.add_argument("--hardcoded-r", type=float, default=None, + help="Fixed R (e.g. 0.07). Omit → Vrugt formula.") + parser.add_argument("--direct-variance", action="store_true", default=False, + help="Use krig_var directly as R (no Vrugt scaling)") + parser.add_argument("--rng-seed", type=int, default=None) + args = parser.parse_args() + + CAT_ID = args.cat_id + OBS_FILE = os.path.join(args.obs_dir, f"{CAT_ID}.csv") + CFE_CONFIG_FILE = args.config_file + HARDCODED_R = args.hardcoded_r + DIRECT_VARIANCE = args.direct_variance + OUT_DIR = Path(args.out_dir) / CAT_ID + OUT_DIR.mkdir(parents=True, exist_ok=True) + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + bmi_cfe = _bmi_cfe + + # Build combined test forcing + if args.test_forcing_dir1 and args.test_forcing_dir2: + f1 = os.path.join(args.test_forcing_dir1, f"{CAT_ID}.csv") + f2 = os.path.join(args.test_forcing_dir2, f"{CAT_ID}.csv") + df1, df2 = pd.read_csv(f1), pd.read_csv(f2) + combined = (pd.concat([df1, df2], ignore_index=True) + .drop_duplicates(subset="time").sort_values("time")) + TEST_FORCING_FILE = str(OUT_DIR / f"{CAT_ID}_nwm_operational_combined.csv") + combined.to_csv(TEST_FORCING_FILE, index=False) + else: + print("Need --test-forcing-dir1 and --test-forcing-dir2"); return + + best_params_file = OUT_DIR / f"{CAT_ID}_best_params.json" + if not best_params_file.exists(): + print(f"Missing best params at {best_params_file}"); return + with open(best_params_file) as f: + best_params = json.load(f)["best_parameters"] + + obs_dict, var_dict = load_obs() + df_forcing = load_test_forcing() + t_mask = (df_forcing["date"] >= TEST_START) & (df_forcing["date"] <= TEST_END) + df_test = df_forcing[t_mask].reset_index(drop=True) + + rng = np.random.default_rng(hash(CAT_ID) & 0x7fffffff) + + # ------------------------------------------------------------------ # + # Phase 1: Run DA through test period, snapshot analysis state at # + # each Helene issue time # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Phase 1: DA analysis through test period...") + model = make_model(best_params) + + helene_snapshots = {} # issue_time -> analysis state dict + + for h, row in df_test.iterrows(): + date_str = row["date"] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + q_sim = run_model_step(model, P, E) + + date_ts = pd.Timestamp(date_str) + + # DA update if obs available + y_obs = obs_dict.get(date_str, np.nan) + if not np.isnan(y_obs): + krig_var = var_dict.get(date_str, 0.07) + state = get_state(model) + updated = enkf_update(state, q_sim, y_obs, krig_var, rng) + set_state(model, updated) + + # Snapshot if this is a Helene issue time + if HELENE_START <= date_ts <= HELENE_END: + helene_snapshots[date_ts] = get_state(model) + + model.finalize() + print(f" Snapshots saved: {len(helene_snapshots)} issue times") + + # ------------------------------------------------------------------ # + # Phase 2A: Forcing arm — 30 members, fixed analyzed state, varying # + # met forcing # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Phase 2A: Forcing arm ({N_FORCING} members)...") + forcing_records = [] + + issue_times = sorted(helene_snapshots.keys()) + t0_to_idx = {t: i for i, t in enumerate(df_test["date"].apply(pd.Timestamp))} + + for t0 in issue_times: + snap = helene_snapshots[t0] + t0_str = t0.strftime("%Y-%m-%d %H:%M:%S") + + # Find index in df_test + if t0 not in t0_to_idx: + continue + start_idx = t0_to_idx[t0] + + q_matrix = np.full((FORECAST_HOURS, N_FORCING), np.nan) + + for mem in range(N_FORCING): + m = make_model(best_params) + set_state(m, snap) # all members start from same DA analysis state + + for lead in range(1, FORECAST_HOURS + 1): + fi = start_idx + lead + if fi >= len(df_test): + break + row = df_test.iloc[fi] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + # Perturb forcing only + mu_p = -0.5 * PRECIP_SIGMA ** 2 + P_pert = float(P * rng.lognormal(mu_p, PRECIP_SIGMA)) if P > 0 else 0.0 + E_pert = max(float(E * (1.0 + PET_SIGMA * rng.standard_normal())), 0.0) + q = run_model_step(m, P_pert, E_pert) + q_matrix[lead - 1, mem] = q + + m.finalize() + + row_data = {"issue_time": t0_str} + for lead in range(1, FORECAST_HOURS + 1): + row_data["lead_hour"] = lead + for mem in range(N_FORCING): + row_data[f"member_{mem:02d}"] = q_matrix[lead - 1, mem] + forcing_records.append(dict(row_data)) + + df_forcing_arm = pd.DataFrame(forcing_records) + out_a = OUT_DIR / f"{CAT_ID}_da_forcing_arm.csv" + df_forcing_arm.to_csv(out_a, index=False) + print(f" Saved: {out_a}") + + # ------------------------------------------------------------------ # + # Phase 2B: Hydro-state arm — 20 members, deterministic forcing, # + # perturbed analysis states # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Phase 2B: Hydro-state arm ({N_HYDRO} members)...") + hydro_records = [] + + for t0 in issue_times: + snap = helene_snapshots[t0] + t0_str = t0.strftime("%Y-%m-%d %H:%M:%S") + + if t0 not in t0_to_idx: + continue + start_idx = t0_to_idx[t0] + + q_matrix = np.full((FORECAST_HOURS, N_HYDRO), np.nan) + + for mem in range(N_HYDRO): + # Perturb the DA analysis state for this member + if mem == 0: + state_m = snap # member 0 = unperturbed analysis mean + else: + state_m = perturb_state(snap, rng) + + m = make_model(best_params) + set_state(m, state_m) + + for lead in range(1, FORECAST_HOURS + 1): + fi = start_idx + lead + if fi >= len(df_test): + break + row = df_test.iloc[fi] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + # Deterministic forcing + q = run_model_step(m, P, E) + q_matrix[lead - 1, mem] = q + + m.finalize() + + row_data = {"issue_time": t0_str} + for lead in range(1, FORECAST_HOURS + 1): + row_data["lead_hour"] = lead + for mem in range(N_HYDRO): + row_data[f"member_{mem:02d}"] = q_matrix[lead - 1, mem] + hydro_records.append(dict(row_data)) + + df_hydro_arm = pd.DataFrame(hydro_records) + out_b = OUT_DIR / f"{CAT_ID}_da_hydro_arm.csv" + df_hydro_arm.to_csv(out_b, index=False) + print(f" Saved: {out_b}") + print(f"[{CAT_ID}] Done.") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/run_perturbation_sensitivity.py b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/run_perturbation_sensitivity.py new file mode 100644 index 00000000..57b8d185 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/run_perturbation_sensitivity.py @@ -0,0 +1,293 @@ +""" +Perturbation-source sensitivity analysis for the multi-catchment CFE DA experiment. + +Ensemble-spread attribution: run 20 CFE members through the test +period with ONLY ONE perturbation source active at a time, no DA, no spinup. + +Three sub-experiments (selected via --source): + init : perturb initial states once at t=0 (multiplicative ~5%) + forcing : perturb hourly precip (lognormal sigma=0.15) and PET (Gaussian sigma=0.10) + process : apply per-hour process noise on states (soil 0.2%, GW 0.15%, Nash 0.5%) + +Each sub-experiment runs N=20 CFE members. DA is always OFF. Spinup is skipped — each +member starts from CFE's default initial state (modified by the perturbation if --source +init). This is the cleanest setup for attributing test-period spread to its source. + +Output: + //_sensitivity_.csv + Columns: date, member_00, member_01, ..., member_19 + +Each member column is that member's hourly Q_sim (mm/h) over the test period. + +Reads the same calibrated best_params.json as the production script. Does not need +the kriging obs file (no DA), but takes the same --obs-dir argument so it can also +log Q_obs alongside for plotting reference. + +Usage: + python3 run_perturbation_sensitivity.py \ + --cat-id cat-1016300 \ + --source forcing \ + --forcing-dir \ + --obs-dir \ + --cfe-dir \ + --config-file \ + --param-bounds \ + --out-dir \ + --test-forcing-dir1 \ + --test-forcing-dir2 +""" + +import argparse +import os +import sys +import json +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 +N_MEMBERS = 20 + +# We only run the test window for this sensitivity analysis. +TEST_START = "2023-10-01 00:00:00" +TEST_END = "2024-10-31 23:00:00" + +# Perturbation magnitudes — same as production defaults +INIT_STATE_FRAC = 0.05 +PRECIP_PERTURB = 0.15 # lognormal sigma +PET_PERTURB = 0.10 # Gaussian sigma +SOIL_PROC_NOISE = 0.002 +GW_PROC_NOISE = 0.0015 +NASH_PROC_NOISE = 0.005 + +# Set at runtime +CAT_ID = None +FORCING_FILE = None +TEST_FORCING_FILE = None +OBS_FILE = None +CFE_CONFIG_FILE = None +OUT_DIR = None +bmi_cfe = None + + +# ---------- Forcing helpers (same as production) ---------- +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(): + df = pd.read_csv(TEST_FORCING_FILE) + df['date'] = pd.to_datetime(df['time']).dt.strftime('%Y-%m-%d %H:%M:%S') + df['total_precipitation'] = df['APCP_surface'] * 3600.0 # kg/m^2/s -> mm/h + df['potential_evaporation'] = priestley_taylor_pet( + df['DSWRF_surface'].values, + df['TMP_2maboveground'].values, + ) + return df + + +# ---------- Perturbation helpers ---------- +def perturb_precip_lognormal(P, rng, sigma=PRECIP_PERTURB): + """Lognormal multiplier with mean 1; preserves zero precip.""" + mu = -0.5 * sigma ** 2 + return float(P * rng.lognormal(mu, sigma)) + + +def perturb_pet_gaussian(E, rng, sigma=PET_PERTURB): + """Gaussian multiplicative noise on PET, clipped at 0.""" + return max(float(E * (1.0 + sigma * rng.standard_normal())), 0.0) + + +def apply_init_state_perturb(m, rng, frac=INIT_STATE_FRAC): + """One-time multiplicative perturbation on the four states (called per member at t=0).""" + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1.0 + frac * rng.standard_normal()), 1e-6, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1.0 + frac * rng.standard_normal()), 1e-6, gw_max)) + # Nash buckets often start at 0; small additive jitter so they're not identical. + jitter = max(sm_max * 1e-6, 1e-9) + m.nash_storage[0] = max(n00 + jitter * rng.standard_normal(), 0.0) + m.nash_storage[1] = max(n10 + jitter * rng.standard_normal(), 0.0) + + +def apply_process_noise(m, rng): + """Per-timestep process noise on all 4 states for one member.""" + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + nash_floor = max(sm_max * 1e-4, 1e-7) + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1.0 + SOIL_PROC_NOISE * rng.standard_normal()), 0.0, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1.0 + GW_PROC_NOISE * rng.standard_normal()), 0.0, gw_max)) + m.nash_storage[0] = max( + n00 * (1.0 + NASH_PROC_NOISE * rng.standard_normal()) + + nash_floor * rng.standard_normal(), 0.0) + m.nash_storage[1] = max( + n10 * (1.0 + NASH_PROC_NOISE * rng.standard_normal()) + + nash_floor * rng.standard_normal(), 0.0) + + +# ---------- Main per-source run ---------- +def run_sensitivity(source, best_param_dict): + """Run N members through the test period with only `source` perturbation active.""" + + def custom_load_forcing(self_cfe): + df = load_test_forcing() + self_cfe.forcing_data = df.rename(columns={"date": "time"}) + + # Build a per-catchment temporary CFE config with the calibrated params + with open(CFE_CONFIG_FILE) as f: + cfg = json.load(f) + cfg['forcing_file'] = TEST_FORCING_FILE + cfg['soil_params']['bb'] = best_param_dict['bb'] + cfg['soil_params']['smcmax'] = best_param_dict['smcmax'] + cfg['soil_params']['satdk'] = best_param_dict['satdk'] + cfg['slop'] = best_param_dict['slop'] + cfg['max_gw_storage'] = best_param_dict['max_gw_storage'] + cfg['expon'] = best_param_dict['expon'] + cfg['Cgw'] = best_param_dict['Cgw'] + cfg['K_lf'] = best_param_dict['K_lf'] + cfg['K_nash'] = best_param_dict['K_nash'] + cfg['partition_scheme'] = "Schaake" if best_param_dict['scheme'] <= 0.5 else "Xinanjiang" + + tmp_cfg = str(OUT_DIR / f'{CAT_ID}_bmi_config_temp_sensitivity_{source}.json') + with open(tmp_cfg, 'w') as f: + json.dump(cfg, f) + + # Deterministic seed per (catchment, source) so each run is reproducible + seed = hash((CAT_ID, source)) & 0x7fffffff + rng = np.random.default_rng(seed) + + print(f"[sensitivity] {CAT_ID} | source={source} | N={N_MEMBERS} | seed={seed}") + + # Build N members. Initial-state perturbation only happens when source == 'init'. + models = [] + for i in range(N_MEMBERS): + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.load_forcing_file = custom_load_forcing.__get__(m) + m.initialize() + if source == 'init' and i > 0: + apply_init_state_perturb(m, rng) + models.append(m) + + # Skip spinup. Go directly to the test period. + df = load_test_forcing() + t_mask = (df['date'] >= TEST_START) & (df['date'] <= TEST_END) + df_test = df[t_mask] + + # Pre-allocate per-member Q arrays + n_hours = len(df_test) + q_matrix = np.full((n_hours, N_MEMBERS), np.nan, dtype=float) + dates_out = df_test['date'].values + + for h, (p, e) in enumerate(zip(df_test['total_precipitation'], + df_test['potential_evaporation'])): + for i, m in enumerate(models): + if source == 'forcing': + p_i = perturb_precip_lognormal(p, rng) + e_i = perturb_pet_gaussian(e, rng) + else: + p_i = float(p) + e_i = float(e) + m.set_value('atmosphere_water__time_integral_of_precipitation_mass_flux', p_i / 1000) + m.set_value('water_potential_evaporation_flux', e_i / 1000 / 3600) + m.update() + q_matrix[h, i] = m.get_value('land_surface_water__runoff_depth') * 1000 # m/h -> mm/h + + # Process noise applied after each hour, but only when source == 'process' + if source == 'process': + for m in models: + apply_process_noise(m, rng) + + for m in models: + m.finalize() + + # Write per-member CSV: date + 20 columns + out = {'date': dates_out} + for i in range(N_MEMBERS): + out[f'member_{i:02d}'] = q_matrix[:, i] + df_out = pd.DataFrame(out) + out_path = OUT_DIR / f'{CAT_ID}_sensitivity_{source}.csv' + df_out.to_csv(out_path, index=False) + print(f"[sensitivity] saved {out_path}") + print(f"[sensitivity] mean ensemble spread (std across members) over test period: " + f"{float(q_matrix.std(axis=1).mean()):.5f} mm/h") + + +def main(): + global CAT_ID, FORCING_FILE, TEST_FORCING_FILE, OBS_FILE, CFE_CONFIG_FILE + global OUT_DIR, bmi_cfe + + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', required=True) + parser.add_argument('--source', required=True, + choices=['init', 'forcing', 'process'], + help="Which perturbation source to isolate") + parser.add_argument('--forcing-dir', required=True) + parser.add_argument('--obs-dir', required=True, + help='Kept for symmetry with production script; not used since DA is off') + parser.add_argument('--cfe-dir', required=True) + parser.add_argument('--config-file', required=True) + parser.add_argument('--param-bounds', required=True, + help='Kept for symmetry; not read in this script') + parser.add_argument('--out-dir', required=True) + parser.add_argument('--test-forcing-dir1', default=None) + parser.add_argument('--test-forcing-dir2', default=None) + args = parser.parse_args() + + CAT_ID = args.cat_id + FORCING_FILE = os.path.join(args.forcing_dir, f'{CAT_ID}.csv') + OBS_FILE = os.path.join(args.obs_dir, f'{CAT_ID}.csv') + CFE_CONFIG_FILE = args.config_file + OUT_DIR = Path(args.out_dir) / CAT_ID + OUT_DIR.mkdir(parents=True, exist_ok=True) + + # Build combined test forcing CSV (same idiom as production) + if args.test_forcing_dir1 and args.test_forcing_dir2: + f1 = os.path.join(args.test_forcing_dir1, f'{CAT_ID}.csv') + f2 = os.path.join(args.test_forcing_dir2, f'{CAT_ID}.csv') + if os.path.exists(f1) and os.path.exists(f2): + df1 = pd.read_csv(f1) + df2 = pd.read_csv(f2) + combined = pd.concat([df1, df2], ignore_index=True) + combined = combined.drop_duplicates(subset='time').sort_values('time') + TEST_FORCING_FILE = str(OUT_DIR / f'{CAT_ID}_nwm_operational_combined.csv') + combined.to_csv(TEST_FORCING_FILE, index=False) + else: + print(f"Warning: test forcing files not found for {CAT_ID}, exiting") + return + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + bmi_cfe = _bmi_cfe + + best_params_file = OUT_DIR / f'{CAT_ID}_best_params.json' + if not best_params_file.exists(): + print(f"No best params file at {best_params_file}. Pre-stage from Run 3.") + return + with open(best_params_file) as f: + saved = json.load(f) + best_param_dict = saved['best_parameters'] + + run_sensitivity(args.source, best_param_dict) + + +if __name__ == '__main__': + main() diff --git a/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/run_production_per_member.py b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/run_production_per_member.py new file mode 100644 index 00000000..d9c1c4ac --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/2_assimilation/run_production_per_member.py @@ -0,0 +1,314 @@ +""" +Run the production DA pipeline for ONE catchment and save per-member streamflow. + +Identical to calibrate_catchment_cfe_da_v2.py in algorithm (init perturbation + +forcing perturbation + process noise + true EnKF with Vrugt R), but writes one +column per ensemble member instead of only the ensemble mean. This is for +the per-member factor-decomposition plot. + +To keep the comparison apples-to-apples with the sensitivity runs (which skip +spinup), this script also SKIPS the spinup loop. Each member starts directly +from CFE's default initial state, plus the t=0 init perturbation. + +Output: + //_production_per_member.csv + Columns: date, member_00, member_01, ..., member_19 + +Imports EnKFAssimilator from calibrate_catchment_cfe_da_v2 to guarantee the +DA math matches production exactly. + +Usage: + python3 run_production_per_member.py \ + --cat-id cat-1016300 \ + --forcing-dir /mnt/.../nwm_retro_catchment_forcings \ + --obs-dir /mnt/.../catchment_ts_03463300_with_variance \ + --cfe-dir /mnt/.../cfe_py \ + --config-file /mnt/.../cat_03463300_bmi_config_cfe.json \ + --param-bounds /mnt/.../CFE_parameter_bounds.json \ + --out-dir /mnt/.../v2_production_per_member \ + --test-forcing-dir1 ... \ + --test-forcing-dir2 ... \ + --enkf-members 20 +""" + +import argparse +import os +import sys +import json +import importlib.util +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +TEST_START = "2023-10-01 00:00:00" +TEST_END = "2024-10-31 23:00:00" + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(test_forcing_file): + df = pd.read_csv(test_forcing_file) + df['date'] = pd.to_datetime(df['time']).dt.strftime('%Y-%m-%d %H:%M:%S') + df['total_precipitation'] = df['APCP_surface'] * 3600.0 + df['potential_evaporation'] = priestley_taylor_pet( + df['DSWRF_surface'].values, df['TMP_2maboveground'].values + ) + return df + + +def import_enkf_class(script_path): + """Load EnKFAssimilator from the production script without triggering its main().""" + spec = importlib.util.spec_from_file_location("prod_v2", script_path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod.EnKFAssimilator + + +def run(args, bmi_cfe, EnKFAssimilator): + cat_id = args.cat_id + out_dir = Path(args.out_dir) / cat_id + out_dir.mkdir(parents=True, exist_ok=True) + + # Combined test forcing + f1 = os.path.join(args.test_forcing_dir1, f'{cat_id}.csv') + f2 = os.path.join(args.test_forcing_dir2, f'{cat_id}.csv') + df1 = pd.read_csv(f1) + df2 = pd.read_csv(f2) + combined = pd.concat([df1, df2], ignore_index=True) + combined = combined.drop_duplicates(subset='time').sort_values('time') + test_forcing_file = str(out_dir / f'{cat_id}_nwm_operational_combined.csv') + combined.to_csv(test_forcing_file, index=False) + + # Load best_params + best_params_file = out_dir / f'{cat_id}_best_params.json' + if not best_params_file.exists(): + print(f"Missing {best_params_file}. Pre-stage from Run 3.") + return + with open(best_params_file) as f: + saved = json.load(f) + best = saved['best_parameters'] + + # Build temp CFE config with calibrated params + with open(args.config_file) as f: + cfg = json.load(f) + cfg['forcing_file'] = test_forcing_file + cfg['soil_params']['bb'] = best['bb'] + cfg['soil_params']['smcmax'] = best['smcmax'] + cfg['soil_params']['satdk'] = best['satdk'] + cfg['slop'] = best['slop'] + cfg['max_gw_storage'] = best['max_gw_storage'] + cfg['expon'] = best['expon'] + cfg['Cgw'] = best['Cgw'] + cfg['K_lf'] = best['K_lf'] + cfg['K_nash'] = best['K_nash'] + cfg['partition_scheme'] = "Schaake" if best['scheme'] <= 0.5 else "Xinanjiang" + tmp_cfg = str(out_dir / f'{cat_id}_bmi_config_temp_per_member.json') + with open(tmp_cfg, 'w') as f: + json.dump(cfg, f) + + # Build the EnKF assimilator (same defaults as production) + obs_file = os.path.join(args.obs_dir, f'{cat_id}.csv') + enkf = EnKFAssimilator( + n_members=args.enkf_members, + obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, + use_vrugt_r=(not args.no_vrugt_r) and (not getattr(args, 'direct_variance', False)), + vrugt_alpha=args.vrugt_alpha, + vrugt_scale=args.vrugt_scale, + rng_seed=args.rng_seed, + ) + N = enkf.n_members + print(f"[per-member] {cat_id} | N={N} | use_vrugt_r={enkf.use_vrugt_r}") + if getattr(args, 'hardcoded_r', None) is not None: + enkf.obs_var_dict = {k: args.hardcoded_r for k in enkf.obs_var_dict} + print(f"[per-member] obs_var_dict overridden: R = {args.hardcoded_r} (fixed)") + if getattr(args, 'direct_variance', False): + obs_df = pd.read_csv(obs_file) + t_col = next(c for c in obs_df.columns if c.lower() in ('datetime', 'date', 'time', 'timestamp')) + v_col = next(c for c in obs_df.columns if 'var' in c.lower()) + obs_df[t_col] = pd.to_datetime(obs_df[t_col]) + obs_df = obs_df.set_index(t_col).sort_index() + enkf.obs_var_dict = { + str(ts): float(v) + for ts, v in obs_df[v_col].items() + if pd.notna(v) and float(v) > 0 + } + print(f"[per-member] obs_var_dict overridden: R = σ² direct ({len(enkf.obs_var_dict)} timesteps)") + + # Custom forcing loader for CFE BMI + def custom_load_forcing(self_cfe): + df = load_test_forcing(test_forcing_file) + self_cfe.forcing_data = df.rename(columns={"date": "time"}) + + # Build N CFE instances; perturb initial states (member 0 left clean) + models = [] + init_states_per_member = [] # capture initial-state snapshot after perturbation + for i in range(N): + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.load_forcing_file = custom_load_forcing.__get__(m) + m.initialize() + if i > 0: + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + r = enkf.init_state_perturb_frac + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1 + r * enkf.rng.standard_normal()), 1e-6, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1 + r * enkf.rng.standard_normal()), 1e-6, gw_max)) + jitter = max(sm_max * 1e-6, 1e-9) + m.nash_storage[0] = max(n00 + jitter * enkf.rng.standard_normal(), 0.0) + m.nash_storage[1] = max(n10 + jitter * enkf.rng.standard_normal(), 0.0) + # Snapshot the (possibly-perturbed) initial state for this member + init_states_per_member.append({ + "soil_m": float(m.soil_reservoir["storage_m"]), + "gw_m": float(m.gw_reservoir["storage_m"]), + "nash0_m": float(m.nash_storage[0]), + "nash1_m": float(m.nash_storage[1]), + "soil_max_m": float(m.soil_reservoir["storage_max_m"]), + "gw_max_m": float(m.gw_reservoir["storage_max_m"]), + }) + models.append(m) + + # SKIP SPINUP — go directly to the test period (matches sensitivity-script setup) + df = load_test_forcing(test_forcing_file) + t_mask = (df['date'] >= TEST_START) & (df['date'] <= TEST_END) + df_test = df[t_mask] + n_hours = len(df_test) + q_matrix = np.full((n_hours, N), np.nan, dtype=float) + precip_matrix = np.full((n_hours, N), np.nan, dtype=float) # mm/h per member + pet_matrix = np.full((n_hours, N), np.nan, dtype=float) # mm/h per member + dates_out = df_test['date'].values + + for h, (p, e, current_date) in enumerate(zip( + df_test['total_precipitation'], + df_test['potential_evaporation'], + df_test['date'])): + + # Perturb forcing per member + p_arr, e_arr = enkf.perturb_forcing(p, e) + # Capture the per-member perturbed forcing for later plotting/diagnostics + precip_matrix[h, :] = p_arr + pet_matrix[h, :] = e_arr + + # Advance each member one hour + ensemble_q = np.empty(N, dtype=float) + for i, m in enumerate(models): + m.set_value('atmosphere_water__time_integral_of_precipitation_mass_flux', + float(p_arr[i]) / 1000) + m.set_value('water_potential_evaporation_flux', + float(e_arr[i]) / 1000 / 3600) + m.update() + ensemble_q[i] = m.get_value('land_surface_water__runoff_depth') * 1000 # mm/h + + # Record pre-analysis forecast per member (this is the "actual production" value) + q_matrix[h, :] = ensemble_q + + # DA step (uses production EnKFAssimilator) + enkf.update_states(models, current_date, ensemble_q) + + # Process noise after DA + enkf.add_process_noise(models) + + for m in models: + m.finalize() + + # ----- Save per-member CSVs ----- + def _save_matrix(matrix, suffix): + out = {'date': dates_out} + for i in range(N): + out[f'member_{i:02d}'] = matrix[:, i] + df_out = pd.DataFrame(out) + out_path = out_dir / f'{cat_id}_production_per_member_{suffix}.csv' + df_out.to_csv(out_path, index=False) + return out_path + + # Streamflow outputs + q_path = out_dir / f'{cat_id}_production_per_member.csv' + out = {'date': dates_out} + for i in range(N): + out[f'member_{i:02d}'] = q_matrix[:, i] + pd.DataFrame(out).to_csv(q_path, index=False) + print(f"[per-member] saved {q_path}") + + # Perturbed forcings (so each member's actual inputs are recoverable) + precip_path = _save_matrix(precip_matrix, "precip") + pet_path = _save_matrix(pet_matrix, "pet") + print(f"[per-member] saved {precip_path}") + print(f"[per-member] saved {pet_path}") + + # Initial states (small, JSON is fine) + init_states_path = out_dir / f'{cat_id}_production_per_member_initial_states.json' + with open(init_states_path, 'w') as f: + json.dump({ + "catchment_id": cat_id, + "n_members": N, + "members": {f"member_{i:02d}": init_states_per_member[i] for i in range(N)}, + }, f, indent=2) + print(f"[per-member] saved {init_states_path}") + + print(f"[per-member] ensemble mean spread over test period: " + f"{float(q_matrix.std(axis=1).mean()):.5f} mm/h") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', required=True) + parser.add_argument('--forcing-dir', required=True) + parser.add_argument('--obs-dir', required=True) + parser.add_argument('--cfe-dir', required=True) + parser.add_argument('--config-file', required=True) + parser.add_argument('--param-bounds', required=True, + help='Kept for symmetry; not read in this script') + parser.add_argument('--out-dir', required=True) + parser.add_argument('--test-forcing-dir1', required=True) + parser.add_argument('--test-forcing-dir2', required=True) + parser.add_argument('--enkf-members', type=int, default=20) + parser.add_argument('--enkf-obs-error-std', type=float, default=0.05) + parser.add_argument('--no-vrugt-r', action='store_true') + parser.add_argument('--hardcoded-r', type=float, default=None, + help='Fix R to this constant value for all timesteps (overrides Vrugt formula)') + parser.add_argument('--direct-variance', action='store_true', default=False, + help='Use kriging variance column as R per timestep (R=σ²; overrides Vrugt formula)') + parser.add_argument('--vrugt-alpha', type=float, default=0.10) + parser.add_argument('--vrugt-scale', type=float, default=0.001) + parser.add_argument('--rng-seed', type=int, default=None) + parser.add_argument('--prod-script', default=None, + help='Path to calibrate_catchment_cfe_da_v2.py for importing ' + 'EnKFAssimilator. Defaults to the script located next to this file.') + args = parser.parse_args() + + # Locate the production script (default: same dir as this script) + if args.prod_script is None: + here = os.path.dirname(os.path.abspath(__file__)) + prod_script = os.path.join(here, 'calibrate_catchment_cfe_da_v2.py') + else: + prod_script = args.prod_script + if not os.path.exists(prod_script): + raise FileNotFoundError(f"Could not find production script at {prod_script}. " + f"Pass --prod-script to override.") + EnKFAssimilator = import_enkf_class(prod_script) + + # Dynamic CFE import (same idiom as production script) + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + + run(args, _bmi_cfe, EnKFAssimilator) + + +if __name__ == '__main__': + main() diff --git a/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/3_routing/route_det_f5.sh b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/3_routing/route_det_f5.sh new file mode 100644 index 00000000..28763f04 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/3_routing/route_det_f5.sh @@ -0,0 +1,31 @@ +#!/usr/bin/env bash +# F5 re-kriging direct variance (1 gauge holdout) — deterministic T-route routing. +# Routes per-catchment _test_results.csv (sim_mm_h) through Muskingum-Cunge +# to gauge 03463300 and writes routed_Q_test.csv + KGE/NSE summary. +# +# Usage: +# bash route_det_f5.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +ROUTE_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/2_troute_routing/run_route.py + +F5_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder5_rekrig_variance_direct +GPKG=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv +KV_DIR=/mnt/disk2/1400_sites_helene/catchment_ts_no_03463300_dynamic_variance_rekrig + +echo "[F5] Deterministic T-route routing..." +echo " da-dir : $F5_DIR" +echo " out-dir: $F5_DIR" +echo " kv-dir : $KV_DIR" + +$TROUTE "$ROUTE_SCRIPT" \ + --gpkg "$GPKG" \ + --da-dir "$F5_DIR" \ + --out-dir "$F5_DIR" \ + --usgs-csv "$USGS_CSV" \ + --kv-dir "$KV_DIR" + +echo "[F5] Deterministic routing done. Output: $F5_DIR/routed_Q_test.csv" diff --git a/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/3_routing/route_ensemble_f5.sh b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/3_routing/route_ensemble_f5.sh new file mode 100644 index 00000000..64bb1c46 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/3_routing/route_ensemble_f5.sh @@ -0,0 +1,30 @@ +#!/usr/bin/env bash +# F5 re-kriging direct variance (1 gauge holdout) — crossed-ensemble T-route routing. +# Routes all 600 crossed-ensemble members through Muskingum-Cunge +# to gauge 03463300 and writes routed_crossed_ensemble.parquet. +# +# Must be run AFTER crossed ensemble parquets exist in F5_DIR. +# +# Usage: +# bash route_ensemble_f5.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +ROUTE_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/2_troute_routing/run_route_crossed_ensemble.py + +F5_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder5_rekrig_variance_direct +GPKG=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +echo "[F5] Crossed-ensemble T-route routing..." +echo " ensemble-dir: $F5_DIR" +echo " out-dir : $F5_DIR" + +$TROUTE "$ROUTE_SCRIPT" \ + --gpkg "$GPKG" \ + --ensemble-dir "$F5_DIR" \ + --out-dir "$F5_DIR" \ + --usgs-csv "$USGS_CSV" + +echo "[F5] Ensemble routing done. Output: $F5_DIR/routed_crossed_ensemble.parquet" diff --git a/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/3_routing/route_leadtime_f5.sh b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/3_routing/route_leadtime_f5.sh new file mode 100644 index 00000000..32b86286 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/3_routing/route_leadtime_f5.sh @@ -0,0 +1,31 @@ +#!/usr/bin/env bash +# F5 re-kriging variance (1 gauge holdout) — route lead-time forecast CSVs. +# Routes all 21-catchment lead-time forecast CSVs through Muskingum-Cunge +# to gauge 03463300. Writes routed_leadtime_da_full.parquet and +# routed_leadtime_openloop_full.parquet for gauge-level error decay analysis. +# +# Run AFTER batch_run_leadtime_f5.sh finishes. +# +# Usage: +# bash route_leadtime_f5.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +ROUTE_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/2_troute_routing/run_route_leadtime_forecasts.py + +F5_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder5_rekrig_variance_direct +GPKG=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder5_rekrig_variance_direct_leadtime_routed + +echo "[F5] Routing lead-time forecasts through T-route..." +echo " forecast-dir: $F5_DIR" +echo " out-dir : $OUT_DIR" + +$TROUTE "$ROUTE_SCRIPT" \ + --gpkg "$GPKG" \ + --forecast-dir "$F5_DIR" \ + --out-dir "$OUT_DIR" + +echo "[F5] Lead-time routing done. Outputs in: $OUT_DIR" +ls "$OUT_DIR"/*.parquet 2>/dev/null || echo " (no parquets found)" diff --git a/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4a_error_decay/plot_forecast_error_fixed_target.py b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4a_error_decay/plot_forecast_error_fixed_target.py new file mode 100644 index 00000000..c01379e0 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4a_error_decay/plot_forecast_error_fixed_target.py @@ -0,0 +1,204 @@ +""" +plot_forecast_error_fixed_target.py — 4a: error decay, fixed-target-time view. +F5 re-kriged variance direct (1 gauge holdout). + +For each target verification time T in the Helene peak window: + Collect all forecasts that verify AT T: + issue_time = T - lead_hour*1h, for lead in 1..18 + error[lead] = ensemble_mean(q at T, initialized T-lead) - USGS_obs(T) + +This gives the correct operational picture: + - lead 1 = initialized 1 hr before T (DA just ran -> small error) + - lead 18 = initialized 18 hr before T (DA long ago -> error ~ open loop) + +Two panels: + Top : signed error (m³/s) vs lead hour, one curve per target time + Bot : same for open-loop +Plus a summary panel: mean across all target times, DA vs OL. + +Outputs: + /error_fixed_target_helene.png (per-target spaghetti, DA vs OL) + /error_fixed_target_mean.png (mean across targets, DA vs OL) +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.cm as cm + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder5_rekrig_variance_direct_leadtime_routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +TARGET_START = pd.Timestamp("2024-09-26 18:00:00") +TARGET_END = pd.Timestamp("2024-09-28 06:00:00") + + +def load_parquet(path): + df = pd.read_parquet(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + df["ens_mean"] = df[member_cols].mean(axis=1) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + return df[["issue_time", "lead_hour", "valid_time", "ens_mean"]] + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() + or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def build_fixed_target_errors(df, obs_series, target_times): + df_idx = df.set_index(["issue_time", "lead_hour"])["ens_mean"] + results = {} + for T in target_times: + obs_val = obs_series.get(T, np.nan) + if np.isnan(obs_val): + continue + curve = {} + for lead in range(1, 19): + t0 = T - pd.Timedelta(hours=lead) + try: + q_fc = df_idx.loc[(t0, lead)] + curve[lead] = float(q_fc) - obs_val + except KeyError: + curve[lead] = np.nan + results[T] = curve + return results + + +def plot_spaghetti(ax, error_dict, color_da, label_prefix, linestyle="-", lw=0.9, alpha=0.35): + leads = list(range(1, 19)) + all_curves = [] + target_times = sorted(error_dict.keys()) + cmap = cm.get_cmap("plasma", len(target_times)) + + for i, T in enumerate(target_times): + curve = [error_dict[T].get(l, np.nan) for l in leads] + ax.plot(leads, curve, + color=cmap(i), lw=lw, alpha=alpha, + linestyle=linestyle, zorder=2) + all_curves.append(curve) + + if all_curves: + mean_curve = np.nanmean(np.array(all_curves), axis=0) + ax.plot(leads, mean_curve, + color=color_da, lw=2.8, alpha=0.95, + linestyle=linestyle, zorder=5, + label=f"{label_prefix} — mean across targets") + return all_curves + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--route-dir", default=DEFAULT_ROUTE_DIR) + parser.add_argument("--out-dir", default=None) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--da-name", default="routed_leadtime_da_full.parquet") + parser.add_argument("--ol-name", default="routed_leadtime_openloop_full.parquet") + parser.add_argument("--target-start", default=str(TARGET_START)) + parser.add_argument("--target-end", default=str(TARGET_END)) + args = parser.parse_args() + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + print("Loading parquets...") + da = load_parquet(os.path.join(args.route_dir, args.da_name)) + ol = load_parquet(os.path.join(args.route_dir, args.ol_name)) + obs = load_usgs(args.usgs_csv) + + target_times = pd.date_range(args.target_start, args.target_end, freq="1h") + print(f" Target verification times: {len(target_times)} " + f"({target_times[0]} → {target_times[-1]})") + + da_errors = build_fixed_target_errors(da, obs, target_times) + ol_errors = build_fixed_target_errors(ol, obs, target_times) + print(f" Targets with obs: DA={len(da_errors)} OL={len(ol_errors)}") + + leads = list(range(1, 19)) + + fig, (ax_da, ax_ol) = plt.subplots(2, 1, figsize=(13, 10), sharex=True, sharey=True) + + ax_da.axhline(0, color="black", lw=0.8, linestyle="--", alpha=0.4) + plot_spaghetti(ax_da, da_errors, "tab:purple", "DA", linestyle="-") + ax_da.set_ylabel("Error: forecast − USGS obs (m³/s)", fontsize=10) + ax_da.set_title("DA — error at each lead for fixed target times (Helene peak window)", fontsize=11) + ax_da.grid(True, alpha=0.2) + ax_da.legend(fontsize=9) + + ax_ol.axhline(0, color="black", lw=0.8, linestyle="--", alpha=0.4) + plot_spaghetti(ax_ol, ol_errors, "tab:gray", "Open-loop", linestyle="--") + ax_ol.set_ylabel("Error: forecast − USGS obs (m³/s)", fontsize=10) + ax_ol.set_xlabel("Forecast lead hour (hours before target)", fontsize=11) + ax_ol.set_title("Open-loop — error at each lead for fixed target times", fontsize=11) + ax_ol.set_xticks(leads) + ax_ol.grid(True, alpha=0.2) + ax_ol.legend(fontsize=9) + + fig.suptitle( + "Forecast error vs lead time — fixed verification time, Helene peak window\n" + f"F5 (re-kriged variance) | Target: {TARGET_START.strftime('%b %d %H UTC')} → " + f"{TARGET_END.strftime('%b %d %H UTC')} | USGS 03463300", + fontsize=12, y=1.01, + ) + plt.tight_layout() + out1 = os.path.join(out_dir, "error_fixed_target_helene.png") + plt.savefig(out1, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + fig, ax = plt.subplots(figsize=(13, 6)) + ax.axhline(0, color="black", lw=0.8, linestyle="--", alpha=0.4) + + def mean_curve(error_dict): + arr = np.array([ + [error_dict[T].get(l, np.nan) for l in leads] + for T in sorted(error_dict.keys()) + ]) + return np.nanmean(arr, axis=0), np.nanstd(arr, axis=0) + + da_mean, da_std = mean_curve(da_errors) + ol_mean, ol_std = mean_curve(ol_errors) + + ax.fill_between(leads, da_mean - da_std, da_mean + da_std, + color="tab:purple", alpha=0.15, zorder=2) + ax.fill_between(leads, ol_mean - ol_std, ol_mean + ol_std, + color="tab:gray", alpha=0.15, zorder=2) + ax.plot(leads, da_mean, color="tab:purple", lw=2.6, marker="o", + zorder=5, label="DA — mean error (±1 std shaded)") + ax.plot(leads, ol_mean, color="tab:gray", lw=2.6, marker="s", + linestyle="--", zorder=5, label="Open-loop — mean error (±1 std shaded)") + + ax.set_xlabel("Forecast lead hour (hours before target verification time)", fontsize=11) + ax.set_ylabel("Mean error: forecast − USGS obs (m³/s)", fontsize=11) + ax.set_xticks(leads) + ax.set_title( + "Mean forecast error vs lead time — fixed verification time, Helene peak window\n" + f"F5 (re-kriged variance) | Target: {TARGET_START.strftime('%b %d %H UTC')} → " + f"{TARGET_END.strftime('%b %d %H UTC')} | Lead 1 = init 1 hr before target", + fontsize=11, + ) + ax.legend(fontsize=10, loc="lower right", frameon=True, framealpha=0.92) + ax.grid(True, alpha=0.25) + plt.tight_layout() + out2 = os.path.join(out_dir, "error_fixed_target_mean.png") + plt.savefig(out2, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4a_error_decay/plot_forecast_error_per_init.py b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4a_error_decay/plot_forecast_error_per_init.py new file mode 100644 index 00000000..82637f38 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4a_error_decay/plot_forecast_error_per_init.py @@ -0,0 +1,197 @@ +""" +plot_forecast_error_per_init.py — F5 re-kriged variance direct (1 gauge holdout). + +Error decay by initialization time for the Helene window (Sep 24-28 2024). + +For each initialization time t0 in the Helene window: + error[lead] = ensemble_mean(q_gauge_m3s at t0+lead) - USGS_obs(t0+lead) + +Plotted as: + DA : thin colored lines (one per init time, colored by date) + thick mean + OL : thin gray dashed lines + thick gray dashed mean + +x-axis: forecast lead hour (1 -> 18) +y-axis: signed error (m³/s), positive = forecast too high + +The expected signal: DA error is small at lead 1 (just assimilated), grows +and converges toward the OL error curve by lead 18. + +Inputs: + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + +Obs: + /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +Outputs: + /forecast_error_per_init_helene.png (signed error) + /forecast_mae_per_lead_helene.png (mean absolute error per lead) +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.patches as mpatches + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder5_rekrig_variance_direct_leadtime_routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +HELENE_INIT_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_INIT_END = pd.Timestamp("2024-09-28 23:00:00") + +DATE_COLORS = { + "2024-09-24": "#1f77b4", + "2024-09-25": "#ff7f0e", + "2024-09-26": "#2ca02c", + "2024-09-27": "#d62728", + "2024-09-28": "#9467bd", +} + + +def load_parquet(path): + df = pd.read_parquet(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + df["ens_mean"] = df[member_cols].mean(axis=1) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + return df[["issue_time", "lead_hour", "valid_time", "ens_mean"]] + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() + or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def build_error_table(df, obs_series, init_start, init_end): + df = df[(df["issue_time"] >= init_start) & (df["issue_time"] <= init_end)].copy() + df["obs"] = df["valid_time"].map(obs_series) + df["error"] = df["ens_mean"] - df["obs"] + return df.dropna(subset=["obs", "error"]) + + +def plot_error(ax, err_df, color, alpha_thin, lw_thin, linestyle, label_prefix): + leads = sorted(err_df["lead_hour"].unique()) + all_curves = [] + + for t0, grp in err_df.groupby("issue_time"): + date_str = str(pd.Timestamp(t0).date()) + c = DATE_COLORS.get(date_str, color) + grp_sorted = grp.sort_values("lead_hour") + curve = grp_sorted.set_index("lead_hour")["error"].reindex(leads).values + ax.plot(leads, curve, + color=c, lw=lw_thin, alpha=alpha_thin, + linestyle=linestyle, zorder=2) + all_curves.append(curve) + + if all_curves: + mean_curve = np.nanmean(np.array(all_curves), axis=0) + ax.plot(leads, mean_curve, + color=color, lw=2.8, alpha=0.95, + linestyle=linestyle, zorder=5, + label=f"{label_prefix} — mean across all init times") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--route-dir", default=DEFAULT_ROUTE_DIR) + parser.add_argument("--out-dir", default=None) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--da-name", default="routed_leadtime_da_full.parquet") + parser.add_argument("--ol-name", default="routed_leadtime_openloop_full.parquet") + args = parser.parse_args() + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + print("Loading parquets...") + da = load_parquet(os.path.join(args.route_dir, args.da_name)) + ol = load_parquet(os.path.join(args.route_dir, args.ol_name)) + obs = load_usgs(args.usgs_csv) + print(f" DA issue_times: {da['issue_time'].nunique()} " + f"OL issue_times: {ol['issue_time'].nunique()}") + + da_err = build_error_table(da, obs, HELENE_INIT_START, HELENE_INIT_END) + ol_err = build_error_table(ol, obs, HELENE_INIT_START, HELENE_INIT_END) + print(f" DA init times in Helene window: {da_err['issue_time'].nunique()}") + + leads = sorted(da_err["lead_hour"].unique()) + + fig, ax = plt.subplots(figsize=(12, 6)) + ax.axhline(0, color="black", lw=0.8, linestyle="--", alpha=0.5, zorder=1) + + plot_error(ax, ol_err, color="tab:gray", alpha_thin=0.12, lw_thin=0.7, + linestyle="--", label_prefix="Open-loop") + plot_error(ax, da_err, color="tab:purple", alpha_thin=0.18, lw_thin=0.8, + linestyle="-", label_prefix="DA") + + patches = [mpatches.Patch(color=c, label=f"Init {d}") + for d, c in DATE_COLORS.items()] + handles, labels = ax.get_legend_handles_labels() + ax.legend(handles=handles + patches, + fontsize=8, loc="upper left", frameon=True, framealpha=0.9, ncol=2) + + ax.set_xlabel("Forecast lead hour", fontsize=11) + ax.set_ylabel("Error: forecast − USGS obs (m³/s)", fontsize=11) + ax.set_xticks(np.arange(1, 19)) + ax.set_title( + "Forecast error vs lead time — per initialization time, Helene window\n" + "F5 (re-kriged variance) | DA (purple solid) vs Open-loop (gray dashed) | " + "Sep 24–28 2024 | USGS 03463300", + fontsize=11, + ) + ax.grid(True, alpha=0.25, lw=0.4) + plt.tight_layout() + out1 = os.path.join(out_dir, "forecast_error_per_init_helene.png") + plt.savefig(out1, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + fig, ax = plt.subplots(figsize=(12, 6)) + + def mean_abs_error_by_lead(err_df): + return err_df.groupby("lead_hour")["error"].apply( + lambda x: float(np.nanmean(np.abs(x))) + ) + + da_mae = mean_abs_error_by_lead(da_err) + ol_mae = mean_abs_error_by_lead(ol_err) + + ax.plot(da_mae.index, da_mae.values, + color="tab:purple", lw=2.4, marker="o", label="DA — mean |error|") + ax.plot(ol_mae.index, ol_mae.values, + color="tab:gray", lw=2.4, marker="s", linestyle="--", + label="Open-loop — mean |error|") + + ax.set_xlabel("Forecast lead hour", fontsize=11) + ax.set_ylabel("Mean |error| vs USGS obs (m³/s)", fontsize=11) + ax.set_xticks(np.arange(1, 19)) + ax.set_title( + "Mean absolute forecast error vs lead time — Helene window\n" + "F5 (re-kriged variance) | DA (purple) vs Open-loop (gray) | " + "Sep 24–28 2024 | USGS 03463300", + fontsize=11, + ) + ax.legend(fontsize=10, loc="upper left", frameon=True, framealpha=0.9) + ax.grid(True, alpha=0.25, lw=0.4) + plt.tight_layout() + out2 = os.path.join(out_dir, "forecast_mae_per_lead_helene.png") + plt.savefig(out2, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4a_error_decay/plot_lead_time_decay.py b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4a_error_decay/plot_lead_time_decay.py new file mode 100644 index 00000000..93f91ca9 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4a_error_decay/plot_lead_time_decay.py @@ -0,0 +1,181 @@ +""" +Catchment-level lead-time forecast error decay curve — F5 (re-kriged variance). + +Two-panel figure for ONE catchment: + TOP — RMSE of ensemble-mean DA forecast vs Qkrig obs at each lead hour (1..18), + with a shaded band for min/max per-member RMSE. Open-loop dashed. + BOTTOM — Mean ensemble spread (std across 20 members) at each lead hour. + +Inputs (from run_lead_time_forecast_sweep.py): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv +Obs: + //_test_results.csv (obs_mm_h column) + +Output: + //_lead_time_decay.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" + +F5_DIR = "/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder5_rekrig_variance_direct" +DEFAULT_KRIG_OBS_DIR = "/mnt/disk2/1400_sites_helene/catchment_ts_no_03463300_dynamic_variance_rekrig" + +DEFAULT_LEADTIME_DIR = F5_DIR +DEFAULT_DA_DIR = F5_DIR + +CAT = "cat-1016300" +LEADTIME_DIR = DEFAULT_LEADTIME_DIR +DA_DIR = DEFAULT_DA_DIR +KRIG_OBS_DIR = DEFAULT_KRIG_OBS_DIR +OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay.png") + + +def load_forecasts(path): + if not os.path.exists(path): + raise FileNotFoundError(path) + df = pd.read_csv(path, parse_dates=['issue_time', 'valid_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + return df, member_cols + + +def load_obs(): + # Prefer _test_results.csv if it exists; otherwise read from kriged obs file. + test_path = os.path.join(DA_DIR, CAT, f"{CAT}_test_results.csv") + if os.path.exists(test_path): + df = pd.read_csv(test_path, parse_dates=['date']) + return df.set_index('date')['obs_mm_h'] + krig_path = os.path.join(KRIG_OBS_DIR, f"{CAT}.csv") + if not os.path.exists(krig_path): + raise FileNotFoundError( + f"No obs found: tried {test_path} and {krig_path}. " + f"Pass --krig-obs-dir to the re-kriged obs directory.") + df = pd.read_csv(krig_path) + print(f" kriged obs file columns: {list(df.columns)}") + t_col = next(c for c in df.columns + if c.lower() in ('datetime', 'date', 'time', 'timestamp')) + # Prefer columns that look like discharge obs (mm/h); exclude step/index columns + obs_col = next( + (c for c in df.columns + if c != t_col and 'var' not in c.lower() + and any(k in c.lower() for k in ('obs', 'q', 'krig', 'mm', 'flow', 'discharge')) + and pd.api.types.is_numeric_dtype(df[c])), + None, + ) + if obs_col is None: + # Fall back: any numeric column that isn't time-like or variance + obs_col = next(c for c in df.columns + if c != t_col and 'var' not in c.lower() + and 'step' not in c.lower() and 'time' not in c.lower() + and 'index' not in c.lower() + and pd.api.types.is_numeric_dtype(df[c])) + df[t_col] = pd.to_datetime(df[t_col]) + print(f" obs from kriged file, column '{obs_col}'") + return df.set_index(t_col)[obs_col].astype(float) + + +def metrics_by_lead(df, member_cols, obs_series): + df = df.copy() + df['obs'] = df['valid_time'].map(obs_series) + df = df.dropna(subset=['obs']) + df['ens_mean'] = df[member_cols].mean(axis=1) + df['ens_std'] = df[member_cols].std(axis=1) + + leads = sorted(df['lead_hour'].unique()) + rmse_mean = np.full(len(leads), np.nan) + rmse_min_member = np.full(len(leads), np.nan) + rmse_max_member = np.full(len(leads), np.nan) + mean_ens_std = np.full(len(leads), np.nan) + for i, L in enumerate(leads): + sub = df[df['lead_hour'] == L] + if len(sub) == 0: + continue + err = sub['ens_mean'].values - sub['obs'].values + rmse_mean[i] = float(np.sqrt(np.mean(err ** 2))) + per_member_rmse = np.array([ + np.sqrt(np.mean((sub[c].values - sub['obs'].values) ** 2)) + for c in member_cols + ]) + rmse_min_member[i] = float(np.nanmin(per_member_rmse)) + rmse_max_member[i] = float(np.nanmax(per_member_rmse)) + mean_ens_std[i] = float(np.nanmean(sub['ens_std'].values)) + return (np.asarray(leads), rmse_mean, rmse_min_member, + rmse_max_member, mean_ens_std) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', default="cat-1016300") + parser.add_argument('--leadtime-dir', default=DEFAULT_LEADTIME_DIR, + help='Dir holding /_lead_time_forecasts_{da,openloop}.csv') + parser.add_argument('--da-dir', default=DEFAULT_DA_DIR, + help='Dir holding /_test_results.csv (for obs_mm_h)') + parser.add_argument('--krig-obs-dir', default=DEFAULT_KRIG_OBS_DIR, + help='Fallback: dir holding per-catchment kriged obs CSVs ' + '(used when _test_results.csv is absent)') + args = parser.parse_args() + + global CAT, LEADTIME_DIR, DA_DIR, KRIG_OBS_DIR, OUT_PNG + CAT = args.cat_id + LEADTIME_DIR = args.leadtime_dir + DA_DIR = args.da_dir + KRIG_OBS_DIR = args.krig_obs_dir + OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay.png") + + da_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_da.csv") + ol_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_openloop.csv") + df_da, m_cols = load_forecasts(da_path) + df_ol, _ = load_forecasts(ol_path) + obs_series = load_obs() + + leads_da, rmse_da, rmse_da_lo, rmse_da_hi, std_da = metrics_by_lead( + df_da, m_cols, obs_series) + leads_ol, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = metrics_by_lead( + df_ol, m_cols, obs_series) + + fig, (ax_top, ax_bot) = plt.subplots(2, 1, figsize=(10, 10), sharex=True, + gridspec_kw={"height_ratios": [1.4, 1.0]}) + + ax_top.fill_between(leads_da, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2, + label="DA — member spread (min/max RMSE)") + ax_top.fill_between(leads_ol, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2, + label="Open-loop — member spread") + ax_top.plot(leads_da, rmse_da, color=DA_COLOR, lw=2.4, marker='o', + zorder=4, label="DA on (ensemble-mean RMSE)") + ax_top.plot(leads_ol, rmse_ol, color=OL_COLOR, lw=2.4, marker='s', + linestyle='--', zorder=4, label="Open-loop (ensemble-mean RMSE)") + ax_top.set_ylabel("RMSE vs Qkrig obs (mm/h)", fontsize=11) + ax_top.set_title( + f"Forecast lead-time error decay — {CAT}\n" + f"F5 (re-kriged variance direct) — test period Oct 2023 – Oct 2024", + fontsize=12, + ) + ax_top.grid(True, alpha=0.25) + ax_top.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + ax_bot.plot(leads_da, std_da, color=DA_COLOR, lw=2.2, marker='o', + label="DA — mean ensemble std") + ax_bot.plot(leads_ol, std_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', label="Open-loop — mean ensemble std") + ax_bot.set_xlabel("Forecast lead time (hours)", fontsize=11) + ax_bot.set_ylabel("Mean ensemble std (mm/h)", fontsize=11) + ax_bot.set_xticks(np.arange(1, 19)) + ax_bot.grid(True, alpha=0.25) + ax_bot.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + plt.tight_layout() + plt.savefig(OUT_PNG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4a_error_decay/plot_lead_time_decay_by_regime.py b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4a_error_decay/plot_lead_time_decay_by_regime.py new file mode 100644 index 00000000..221c540a --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4a_error_decay/plot_lead_time_decay_by_regime.py @@ -0,0 +1,217 @@ +""" +Lead-time decay curve, split by flow regime at issue time — F5 (re-kriged variance). + +Three regimes partitioned on obs(t0) (kriging obs at the issue time): + + helene issue_time in [2024-09-24, 2024-09-28] (5-day Helene window) + storm obs(t0) > 0.5 mm/h (high-flow issue times) + low_flow obs(t0) < 0.1 mm/h (low-flow issue times) + +Same 2-panel layout per regime: RMSE vs lead (top), ensemble spread vs lead +(bottom). Three regime columns x 2 metric rows in one figure. + +Inputs (from batch_run_f5_lead_time_sweep.sh): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv +Obs: + //_test_results.csv if present, else + /.csv (qkrig_mm_hr column) + +Output: + //_lead_time_decay_by_regime.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" + +F5_DIR = "/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder5_rekrig_variance_direct" +DEFAULT_KRIG_OBS_DIR = "/mnt/disk2/1400_sites_helene/catchment_ts_no_03463300_dynamic_variance_rekrig" + +DEFAULT_LEADTIME_DIR = F5_DIR +DEFAULT_DA_DIR = F5_DIR +DEFAULT_KRIG_OBS_DIR = DEFAULT_KRIG_OBS_DIR + +CAT = "cat-1016300" +LEADTIME_DIR = DEFAULT_LEADTIME_DIR +DA_DIR = DEFAULT_DA_DIR +KRIG_OBS_DIR = DEFAULT_KRIG_OBS_DIR +OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay_by_regime.png") + +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-28 23:00:00") +STORM_OBS_THRESHOLD = 0.5 # mm/h +LOWFLOW_OBS_THRESHOLD = 0.1 # mm/h + + +def load_forecasts(path): + if not os.path.exists(path): + raise FileNotFoundError(path) + df = pd.read_csv(path, parse_dates=['issue_time', 'valid_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + return df, member_cols + + +def load_obs(): + test_path = os.path.join(DA_DIR, CAT, f"{CAT}_test_results.csv") + if os.path.exists(test_path): + df = pd.read_csv(test_path, parse_dates=['date']) + return df.set_index('date')['obs_mm_h'] + krig_path = os.path.join(KRIG_OBS_DIR, f"{CAT}.csv") + if not os.path.exists(krig_path): + raise FileNotFoundError( + f"No obs found: tried {test_path} and {krig_path}. " + f"Pass --krig-obs-dir to the re-kriged obs directory.") + df = pd.read_csv(krig_path) + t_col = next(c for c in df.columns + if c.lower() in ('datetime', 'date', 'time', 'timestamp')) + obs_col = next( + (c for c in df.columns + if c != t_col and 'var' not in c.lower() + and any(k in c.lower() for k in ('obs', 'q', 'krig', 'mm', 'flow', 'discharge')) + and pd.api.types.is_numeric_dtype(df[c])), + None, + ) + if obs_col is None: + obs_col = next(c for c in df.columns + if c != t_col and 'var' not in c.lower() + and 'step' not in c.lower() and 'time' not in c.lower() + and 'index' not in c.lower() + and pd.api.types.is_numeric_dtype(df[c])) + df[t_col] = pd.to_datetime(df[t_col]) + return df.set_index(t_col)[obs_col].astype(float) + + +def regime_mask(issue_times, obs_at_issue, regime): + if regime == "helene": + return (issue_times >= HELENE_START) & (issue_times <= HELENE_END) + if regime == "storm": + return obs_at_issue > STORM_OBS_THRESHOLD + if regime == "low_flow": + return obs_at_issue < LOWFLOW_OBS_THRESHOLD + raise ValueError(regime) + + +def metrics_by_lead(df, member_cols, obs_series, issue_mask): + df = df.copy() + df = df[df['issue_time'].isin(issue_mask)] + if len(df) == 0: + return None + df['obs'] = df['valid_time'].map(obs_series) + df = df.dropna(subset=['obs']) + if len(df) == 0: + return None + df['ens_mean'] = df[member_cols].mean(axis=1) + df['ens_std'] = df[member_cols].std(axis=1) + + leads = sorted(df['lead_hour'].unique()) + rmse_mean = np.full(len(leads), np.nan) + rmse_lo = np.full(len(leads), np.nan) + rmse_hi = np.full(len(leads), np.nan) + std_mean = np.full(len(leads), np.nan) + for i, L in enumerate(leads): + sub = df[df['lead_hour'] == L] + if len(sub) == 0: + continue + err = sub['ens_mean'].values - sub['obs'].values + rmse_mean[i] = float(np.sqrt(np.mean(err ** 2))) + per_member = np.array([ + np.sqrt(np.mean((sub[c].values - sub['obs'].values) ** 2)) + for c in member_cols + ]) + rmse_lo[i] = float(np.nanmin(per_member)) + rmse_hi[i] = float(np.nanmax(per_member)) + std_mean[i] = float(np.nanmean(sub['ens_std'].values)) + return np.asarray(leads), rmse_mean, rmse_lo, rmse_hi, std_mean + + +def plot_pair(ax_rmse, ax_std, da_metrics, ol_metrics, regime_label, n_issue_times): + if da_metrics is None or ol_metrics is None: + ax_rmse.set_title(f"{regime_label}\n(no issue times match)", fontsize=11) + ax_rmse.axis('off'); ax_std.axis('off') + return + leads, rmse_da, rmse_da_lo, rmse_da_hi, std_da = da_metrics + _, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = ol_metrics + + ax_rmse.fill_between(leads, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2) + ax_rmse.fill_between(leads, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2) + ax_rmse.plot(leads, rmse_da, color=DA_COLOR, lw=2.2, marker='o', + zorder=4, label="DA on") + ax_rmse.plot(leads, rmse_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', zorder=4, label="Open-loop") + ax_rmse.set_title(f"{regime_label}\n({n_issue_times} issue times)", fontsize=11) + ax_rmse.set_ylabel("RMSE (mm/h)", fontsize=10) + ax_rmse.grid(True, alpha=0.25) + ax_rmse.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + ax_std.plot(leads, std_da, color=DA_COLOR, lw=2.0, marker='o', label="DA") + ax_std.plot(leads, std_ol, color=OL_COLOR, lw=2.0, marker='s', + linestyle='--', label="Open-loop") + ax_std.set_xlabel("Forecast lead time (hours)", fontsize=10) + ax_std.set_ylabel("Mean ensemble std (mm/h)", fontsize=10) + ax_std.set_xticks(np.arange(1, 19)) + ax_std.grid(True, alpha=0.25) + ax_std.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', default="cat-1016300") + parser.add_argument('--leadtime-dir', default=DEFAULT_LEADTIME_DIR) + parser.add_argument('--da-dir', default=DEFAULT_DA_DIR) + parser.add_argument('--krig-obs-dir', default=DEFAULT_KRIG_OBS_DIR, + help='Dir holding per-catchment re-kriged obs CSVs ' + '(used when _test_results.csv is absent)') + args = parser.parse_args() + + global CAT, LEADTIME_DIR, DA_DIR, KRIG_OBS_DIR, OUT_PNG + CAT = args.cat_id + LEADTIME_DIR = args.leadtime_dir + DA_DIR = args.da_dir + KRIG_OBS_DIR = args.krig_obs_dir + OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay_by_regime.png") + + da_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_da.csv") + ol_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_openloop.csv") + df_da, m_cols = load_forecasts(da_path) + df_ol, _ = load_forecasts(ol_path) + obs_series = load_obs() + + issue_times = pd.to_datetime(sorted(df_da['issue_time'].unique())) + obs_at_issue = pd.Series(issue_times, index=issue_times).map(obs_series) + + regimes = [ + ("helene", "Helene window (Sep 24-28 2024)"), + ("storm", f"Storm hours (obs(t0) > {STORM_OBS_THRESHOLD} mm/h)"), + ("low_flow", f"Low flow (obs(t0) < {LOWFLOW_OBS_THRESHOLD} mm/h)"), + ] + + fig, axes = plt.subplots(2, 3, figsize=(18, 10), + gridspec_kw={"height_ratios": [1.4, 1.0]}) + for col, (key, label) in enumerate(regimes): + mask = regime_mask(issue_times, obs_at_issue, key) + kept = issue_times[mask] + kept_set = set(pd.to_datetime(kept)) + da_m = metrics_by_lead(df_da, m_cols, obs_series, kept_set) + ol_m = metrics_by_lead(df_ol, m_cols, obs_series, kept_set) + plot_pair(axes[0, col], axes[1, col], da_m, ol_m, label, len(kept)) + + fig.suptitle( + f"Lead-time error decay by flow regime — {CAT}\n" + f"F5 (re-kriged variance direct) — test period Oct 2023 - Oct 2024", + fontsize=13, y=0.995, + ) + plt.tight_layout(rect=(0, 0, 1, 0.97)) + plt.savefig(OUT_PNG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4a_error_decay/plot_lead_time_decay_gauge.py b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4a_error_decay/plot_lead_time_decay_gauge.py new file mode 100644 index 00000000..0124e806 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4a_error_decay/plot_lead_time_decay_gauge.py @@ -0,0 +1,271 @@ +""" +Gauge-level lead-time forecast decay curve — F5 (re-kriged variance). + +Routes 21-catchment forecast CSVs through T-route to USGS gauge 03463300 and +plots RMSE vs lead time (DA solid, open-loop dashed) against USGS observed Q. +Same 2-panel layout as the catchment-level plot, plus a regime split for the +Helene window. + +The routing step is done separately by route_lead_time_forecasts.py; +this script reads the resulting parquets and produces the plot. + +Inputs: + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + Wide format: issue_time, lead_hour, member_00..member_19 (q_gauge_m3s) +Obs: + /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +Output: + /lead_time_decay_gauge_pooled.png + /lead_time_decay_gauge_by_regime.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder5_rekrig_variance_direct_leadtime_routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-28 23:00:00") +STORM_OBS_THRESHOLD_M3S = 20.0 +LOWFLOW_OBS_THRESHOLD_M3S = 5.0 + +USGS_HELENE_PEAK_M3S = 1886.0 + + +def load_parquet_long(path): + if not os.path.exists(path): + raise FileNotFoundError(path) + df = pd.read_parquet(path) + + if 'q_gauge_m3s' in df.columns and 'member' in df.columns: + df['issue_time'] = pd.to_datetime(df['issue_time']) + return df[['issue_time', 'lead_hour', 'member', 'q_gauge_m3s']] + + q_col_candidates = [c for c in df.columns if c.lower() in + ('q_gauge_m3s', 'q_m3s', 'q_gauge', 'q')] + if 'member' in df.columns and q_col_candidates: + qc = q_col_candidates[0] + df['issue_time'] = pd.to_datetime(df['issue_time']) + out = df[['issue_time', 'lead_hour', 'member', qc]].copy() + return out.rename(columns={qc: 'q_gauge_m3s'}) + + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + if not member_cols: + raise ValueError(f"Couldn't identify member columns in {path}. " + f"Columns: {df.columns.tolist()}") + df['issue_time'] = pd.to_datetime(df['issue_time']) + keep = ['issue_time', 'lead_hour'] + if 'valid_time' in df.columns: + keep.append('valid_time') + return df[keep + member_cols].melt( + id_vars=keep, value_vars=member_cols, + var_name='member', value_name='q_gauge_m3s') + + +def load_usgs_obs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_candidates = [c for c in df.columns if c.lower() in + ('datetime', 'date', 'time', 'timestamp')] + q_candidates = [c for c in df.columns if 'q' in c.lower() + or 'flow' in c.lower() or 'discharge' in c.lower()] + if not date_candidates or not q_candidates: + raise ValueError(f"Could not identify date/Q columns in {usgs_csv}. " + f"Columns: {df.columns.tolist()}") + dc, qc = date_candidates[0], q_candidates[0] + df[dc] = pd.to_datetime(df[dc]) + series = df.set_index(dc)[qc].astype(float) + if 'mm' in qc.lower(): + series = series * MM_H_TO_M3_S + print(f" Converted obs from mm/h (column '{qc}') to m3/s " + f"using area = {WATERSHED_AREA_KM2} km2 (x {MM_H_TO_M3_S:.4f})") + else: + print(f" Loaded obs as-is from column '{qc}' (assumed m3/s)") + return series + + +def metrics_by_lead(df_long, obs_series, issue_mask=None): + df = df_long.copy() + if issue_mask is not None: + df = df[df['issue_time'].isin(issue_mask)] + if len(df) == 0: + return None + df['valid_time'] = df['issue_time'] + pd.to_timedelta(df['lead_hour'], unit='h') + df['obs'] = df['valid_time'].map(obs_series) + df = df.dropna(subset=['obs']) + if len(df) == 0: + return None + + grouped = df.groupby(['issue_time', 'lead_hour']) + ens_mean = grouped['q_gauge_m3s'].mean().reset_index(name='ens_mean') + ens_std = grouped['q_gauge_m3s'].std().reset_index(name='ens_std') + obs_per = grouped['obs'].first().reset_index(name='obs') + panel = ens_mean.merge(ens_std, on=['issue_time', 'lead_hour']).merge( + obs_per, on=['issue_time', 'lead_hour']) + + leads = sorted(panel['lead_hour'].unique()) + rmse_mean = np.full(len(leads), np.nan) + rmse_lo = np.full(len(leads), np.nan) + rmse_hi = np.full(len(leads), np.nan) + std_mean = np.full(len(leads), np.nan) + for i, L in enumerate(leads): + psub = panel[panel['lead_hour'] == L] + if len(psub) == 0: + continue + err = psub['ens_mean'].values - psub['obs'].values + rmse_mean[i] = float(np.sqrt(np.mean(err ** 2))) + dsub = df[df['lead_hour'] == L] + member_rmses = [] + for m, g in dsub.groupby('member'): + err_m = g['q_gauge_m3s'].values - g['obs'].values + if len(err_m) >= 1: + member_rmses.append(float(np.sqrt(np.mean(err_m ** 2)))) + if member_rmses: + rmse_lo[i] = float(np.nanmin(member_rmses)) + rmse_hi[i] = float(np.nanmax(member_rmses)) + std_mean[i] = float(np.nanmean(psub['ens_std'].values)) + return np.asarray(leads), rmse_mean, rmse_lo, rmse_hi, std_mean + + +def plot_two_panel(out_path, da_metrics, ol_metrics, title): + fig, (ax_top, ax_bot) = plt.subplots(2, 1, figsize=(10, 10), sharex=True, + gridspec_kw={"height_ratios": [1.4, 1.0]}) + leads, rmse_da, rmse_da_lo, rmse_da_hi, std_da = da_metrics + _, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = ol_metrics + + ax_top.fill_between(leads, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2, + label="DA — member spread (min/max RMSE)") + ax_top.fill_between(leads, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2, + label="Open-loop — member spread") + ax_top.plot(leads, rmse_da, color=DA_COLOR, lw=2.4, marker='o', + zorder=4, label="DA on (ensemble-mean RMSE)") + ax_top.plot(leads, rmse_ol, color=OL_COLOR, lw=2.4, marker='s', + linestyle='--', zorder=4, label="Open-loop (ensemble-mean RMSE)") + ax_top.set_ylabel("RMSE vs USGS obs (m3/s)", fontsize=11) + ax_top.set_title(title, fontsize=12) + ax_top.grid(True, alpha=0.25) + ax_top.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + ax_bot.plot(leads, std_da, color=DA_COLOR, lw=2.2, marker='o', + label="DA — mean ensemble std") + ax_bot.plot(leads, std_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', label="Open-loop — mean ensemble std") + ax_bot.set_xlabel("Forecast lead time (hours)", fontsize=11) + ax_bot.set_ylabel("Mean ensemble std (m3/s)", fontsize=11) + ax_bot.set_xticks(np.arange(1, 19)) + ax_bot.grid(True, alpha=0.25) + ax_bot.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + plt.tight_layout() + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + +def plot_three_panel_regime(out_path, regime_metrics): + fig, axes = plt.subplots(2, 3, figsize=(18, 10), + gridspec_kw={"height_ratios": [1.4, 1.0]}) + for col, (label, n_issue, da_m, ol_m) in enumerate(regime_metrics): + ax_rmse = axes[0, col] + ax_std = axes[1, col] + if da_m is None or ol_m is None: + ax_rmse.set_title(f"{label}\n(no issue times match)", fontsize=11) + ax_rmse.axis('off'); ax_std.axis('off') + continue + leads, rmse_da, rmse_da_lo, rmse_da_hi, std_da = da_m + _, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = ol_m + ax_rmse.fill_between(leads, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2) + ax_rmse.fill_between(leads, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2) + ax_rmse.plot(leads, rmse_da, color=DA_COLOR, lw=2.2, marker='o', + zorder=4, label="DA on") + ax_rmse.plot(leads, rmse_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', zorder=4, label="Open-loop") + ax_rmse.set_title(f"{label}\n({n_issue} issue times)", fontsize=11) + ax_rmse.set_ylabel("RMSE (m3/s)", fontsize=10) + ax_rmse.grid(True, alpha=0.25) + ax_rmse.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + ax_std.plot(leads, std_da, color=DA_COLOR, lw=2.0, marker='o', label="DA") + ax_std.plot(leads, std_ol, color=OL_COLOR, lw=2.0, marker='s', + linestyle='--', label="Open-loop") + ax_std.set_xlabel("Forecast lead time (hours)", fontsize=10) + ax_std.set_ylabel("Mean ensemble std (m3/s)", fontsize=10) + ax_std.set_xticks(np.arange(1, 19)) + ax_std.grid(True, alpha=0.25) + ax_std.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + fig.suptitle("Gauge-level lead-time error decay by regime — USGS 03463300\n" + "F5 (re-kriged variance direct)", + fontsize=13, y=0.995) + plt.tight_layout(rect=(0, 0, 1, 0.97)) + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--route-dir', default=DEFAULT_ROUTE_DIR) + parser.add_argument('--out-dir', default=None) + parser.add_argument('--usgs-csv', default=DEFAULT_USGS_CSV) + parser.add_argument('--da-name', default='routed_leadtime_da_full.parquet') + parser.add_argument('--ol-name', default='routed_leadtime_openloop_full.parquet') + args = parser.parse_args() + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + da_path = os.path.join(args.route_dir, args.da_name) + ol_path = os.path.join(args.route_dir, args.ol_name) + df_da = load_parquet_long(da_path) + df_ol = load_parquet_long(ol_path) + obs_series = load_usgs_obs(args.usgs_csv) + + print(f"DA rows: {len(df_da):,} | unique issue_times: {df_da['issue_time'].nunique()}") + print(f"OL rows: {len(df_ol):,} | unique issue_times: {df_ol['issue_time'].nunique()}") + print(f"USGS Helene peak (reference): {USGS_HELENE_PEAK_M3S:.0f} m3/s") + + da_pooled = metrics_by_lead(df_da, obs_series) + ol_pooled = metrics_by_lead(df_ol, obs_series) + out_pooled = os.path.join(out_dir, "lead_time_decay_gauge_pooled.png") + plot_two_panel(out_pooled, da_pooled, ol_pooled, + title="Gauge-level lead-time decay — USGS 03463300\n" + "F5 (re-kriged variance direct) — all issue times pooled") + + issue_times = pd.to_datetime(sorted(df_da['issue_time'].unique())) + obs_at_issue = pd.Series(issue_times, index=issue_times).map(obs_series) + + helene_mask = (issue_times >= HELENE_START) & (issue_times <= HELENE_END) + storm_mask = obs_at_issue > STORM_OBS_THRESHOLD_M3S + lowflow_mask = obs_at_issue < LOWFLOW_OBS_THRESHOLD_M3S + + regime_metrics = [] + for label, mask in [ + ("Helene window (Sep 24-28 2024)", helene_mask), + (f"Storm hours (obs(t0) > {STORM_OBS_THRESHOLD_M3S:.0f} m3/s)", storm_mask), + (f"Low flow (obs(t0) < {LOWFLOW_OBS_THRESHOLD_M3S:.0f} m3/s)", lowflow_mask), + ]: + kept = set(pd.to_datetime(issue_times[mask])) + da_m = metrics_by_lead(df_da, obs_series, kept) + ol_m = metrics_by_lead(df_ol, obs_series, kept) + regime_metrics.append((label, len(kept), da_m, ol_m)) + + out_regime = os.path.join(out_dir, "lead_time_decay_gauge_by_regime.png") + plot_three_panel_regime(out_regime, regime_metrics) + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4a_error_decay/run_4a_f5.sh b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4a_error_decay/run_4a_f5.sh new file mode 100644 index 00000000..5ff06911 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4a_error_decay/run_4a_f5.sh @@ -0,0 +1,77 @@ +#!/usr/bin/env bash +# F5 re-kriged variance (1 gauge holdout) — 4a lead-time decay plots. +# +# Runs 3 plot scripts per catchment: +# plot_lead_time_decay.py -- pooled RMSE vs lead (catchment-level) +# plot_lead_time_decay_by_regime.py -- same split by flow regime +# plot_lead_time_decay_gauge.py -- gauge-level (requires routed parquets) +# +# F5 has no _test_results.csv per catchment, so obs are read from the +# re-kriged obs directory (--krig-obs-dir). +# +# Usage: +# bash run_4a_f5.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +F5_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder5_rekrig_variance_direct +KRIG_OBS_DIR=/mnt/disk2/1400_sites_helene/catchment_ts_no_03463300_dynamic_variance_rekrig +ROUTE_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder5_rekrig_variance_direct_leadtime_routed +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +echo "[F5-4a] Lead-time decay plots — leadtime dir: $F5_DIR" + +for CAT in "${CATS[@]}"; do + DA_CSV="$F5_DIR/$CAT/${CAT}_lead_time_forecasts_da.csv" + if [ ! -f "$DA_CSV" ]; then + echo " [$CAT] No lead-time CSV — skipping" + continue + fi + + echo " [$CAT] Plotting pooled decay..." + $TROUTE "$SCRIPT_DIR/plot_lead_time_decay.py" \ + --cat-id "$CAT" \ + --leadtime-dir "$F5_DIR" \ + --da-dir "$F5_DIR" \ + --krig-obs-dir "$KRIG_OBS_DIR" + + echo " [$CAT] Plotting decay by regime..." + $TROUTE "$SCRIPT_DIR/plot_lead_time_decay_by_regime.py" \ + --cat-id "$CAT" \ + --leadtime-dir "$F5_DIR" \ + --da-dir "$F5_DIR" \ + --krig-obs-dir "$KRIG_OBS_DIR" +done + +if [ -f "$ROUTE_DIR/routed_leadtime_da_full.parquet" ]; then + echo "[F5-4a] Gauge-level decay (routed)..." + $TROUTE "$SCRIPT_DIR/plot_lead_time_decay_gauge.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" + + echo "[F5-4a] Fixed-target error plots..." + $TROUTE "$SCRIPT_DIR/plot_forecast_error_fixed_target.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" + + echo "[F5-4a] Per-init error plots..." + $TROUTE "$SCRIPT_DIR/plot_forecast_error_per_init.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" +else + echo "[F5-4a] Routed parquets not found — skipping gauge-level and error plots." + echo " Run route_leadtime_f5.sh (in 3_routing/) first, then re-run this script." +fi + +echo "[F5-4a] Done." diff --git a/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_crossed_ensemble.py b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_crossed_ensemble.py new file mode 100644 index 00000000..7cb00a93 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_crossed_ensemble.py @@ -0,0 +1,229 @@ +""" +plot_crossed_ensemble.py -- 4b: 600-member crossed ensemble vs USGS obs. + +Loads the crossed ensemble parquet produced by run_crossed_ensemble.py: + //_crossed_ensemble.parquet + Columns: issue_time, lead_hour, member_0000..member_0599 + +Projects all members to valid_time = issue_time + lead_hour hours. +Aggregates at each valid_time across ALL members from ALL issue times +to build a probabilistic envelope (5th/25th/50th/75th/95th percentiles). + +Two output figures: + 4b_crossed_ensemble_helene.png + -- Full Helene window (Sep 24-30) probabilistic envelope + USGS obs. + Percentile shading in two layers (5-95 outer, 25-75 inner). + 4b_crossed_ensemble_peak.png + -- Zoomed to Helene peak (Sep 26 12 UTC -> Sep 28 06 UTC), + same shading + USGS obs + individual spaghetti from Sep 27 inits. + +Usage: + python3 plot_crossed_ensemble.py \\ + --ensemble-dir /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble \\ + --cat-id cat-1016300 \\ + --usgs-csv /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv \\ + --out-dir /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble/cat-1016300 +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import matplotlib.patches as mpatches + +DEFAULT_ENSEMBLE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" +DEFAULT_CAT_ID = "cat-1016300" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +PLOT_START = pd.Timestamp("2024-09-24 00:00:00") +PLOT_END = pd.Timestamp("2024-09-28 18:00:00") # cut trailing spike +PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +PEAK_END = pd.Timestamp("2024-09-28 06:00:00") +HELENE_PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +HELENE_PEAK_END = pd.Timestamp("2024-09-28 00:00:00") + + +def load_ensemble(path): + df = pd.read_parquet(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + return df, member_cols + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def build_percentile_envelope(df, member_cols, t_start, t_end): + """ + For each valid_time in [t_start, t_end], stack all member values + from all issue times and compute percentiles. + Returns DataFrame indexed by valid_time with columns p05, p25, p50, p75, p95. + """ + mask = (df["valid_time"] >= t_start) & (df["valid_time"] <= t_end) + sub = df[mask].copy() + sub["valid_time"] = sub["valid_time"].dt.floor("1h") + + records = [] + for vt, grp in sub.groupby("valid_time"): + vals = grp[member_cols].to_numpy(dtype=float).flatten() + vals = vals[~np.isnan(vals)] * MM_H_TO_M3_S + if len(vals) == 0: + continue + records.append({ + "valid_time": vt, + "p05": np.percentile(vals, 5), + "p25": np.percentile(vals, 25), + "p50": np.percentile(vals, 50), + "p75": np.percentile(vals, 75), + "p95": np.percentile(vals, 95), + "n": len(vals), + }) + return pd.DataFrame(records).set_index("valid_time").sort_index() + + +def _add_helene_band(ax, label=True): + ax.axvspan(HELENE_PEAK_START, HELENE_PEAK_END, + color="salmon", alpha=0.12, zorder=0) + if label: + ax.text(HELENE_PEAK_START + pd.Timedelta(hours=6), 0.97, "Helene peak", + transform=ax.get_xaxis_transform(), + fontsize=9, color="firebrick", ha="left", va="top") + + +def _format_xaxis(ax, major_interval=1, minor_hours=6): + ax.xaxis.set_major_locator(mdates.DayLocator(interval=major_interval)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.xaxis.set_minor_locator(mdates.HourLocator(interval=minor_hours)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.2, lw=0.4) + + +def plot_envelope(ax, env, obs_series, t_start, t_end, + ens_color="tab:blue", obs_color="black", helene_label=True): + _add_helene_band(ax, label=helene_label) + + # Outer band: 5-95 + ax.fill_between(env.index, env["p05"], env["p95"], + color=ens_color, alpha=0.15, zorder=2, + label="5th-95th percentile") + # Inner band: 25-75 + ax.fill_between(env.index, env["p25"], env["p75"], + color=ens_color, alpha=0.30, zorder=3, + label="25th-75th percentile") + # Median + ax.plot(env.index, env["p50"], + color=ens_color, lw=2.2, zorder=5, + label="Ensemble median (600 members)") + + # USGS obs + obs_w = obs_series.loc[t_start:t_end] + ax.plot(obs_w.index, obs_w.values, + color=obs_color, lw=1.8, zorder=6, label="USGS obs") + + ax.set_xlim(t_start, t_end) + ax.set_ylabel("Discharge (m3/s)", fontsize=11) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--ensemble-dir", default=DEFAULT_ENSEMBLE_DIR) + parser.add_argument("--cat-id", default=DEFAULT_CAT_ID) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--out-dir", default=None) + args = parser.parse_args() + + cat_dir = os.path.join(args.ensemble_dir, args.cat_id) + out_dir = args.out_dir or cat_dir + os.makedirs(out_dir, exist_ok=True) + + parquet_path = os.path.join(cat_dir, f"{args.cat_id}_crossed_ensemble.parquet") + print(f"Loading crossed ensemble: {parquet_path}") + df, member_cols = load_ensemble(parquet_path) + obs = load_usgs(args.usgs_csv) + print(f" {df['issue_time'].nunique()} issue times, " + f"{len(member_cols)} members per row") + + # ------------------------------------------------------------------ # + # Figure 4b-full: full Helene window probabilistic envelope # + # ------------------------------------------------------------------ # + print("Building percentile envelope (full window)...") + env_full = build_percentile_envelope(df, member_cols, PLOT_START, PLOT_END) + print(f" Valid times with data: {len(env_full)}, " + f"median members per timestep: {env_full['n'].median():.0f}") + + fig, ax = plt.subplots(figsize=(17, 7)) + plot_envelope(ax, env_full, obs, PLOT_START, PLOT_END) + ax.set_xlabel("Date", fontsize=11) + ax.set_title( + f"4b -- 600-member crossed ensemble (30 met x 20 hydro-state) | {args.cat_id}\n" + "Shaded: 5-95th (light) and 25-75th (dark) percentile across all members " + "and issue times | Sep 24-30 2024 | USGS 03463300", + fontsize=11, + ) + ax.legend(fontsize=10, loc="upper left", frameon=True, framealpha=0.9) + _format_xaxis(ax) + plt.tight_layout() + out1 = os.path.join(out_dir, f"{args.cat_id}_4b_crossed_ensemble_helene.png") + plt.savefig(out1, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + # ------------------------------------------------------------------ # + # Figure 4b-peak: zoomed to Helene peak window # + # ------------------------------------------------------------------ # + print("Building percentile envelope (peak window)...") + env_peak = build_percentile_envelope(df, member_cols, PEAK_START, PEAK_END) + + fig, (ax_main, ax_zoom) = plt.subplots( + 2, 1, figsize=(14, 11), + gridspec_kw={"height_ratios": [1, 1.4]}, sharex=False) + + # Top panel: full window context (smaller) + plot_envelope(ax_main, env_full, obs, PLOT_START, PLOT_END, helene_label=False) + ax_main.set_title("Full Helene window context (Sep 24 - Sep 28 18 UTC)", fontsize=10) + _format_xaxis(ax_main) + # Shade the zoom region on context panel + ax_main.axvspan(PEAK_START, PEAK_END, color="gold", alpha=0.18, zorder=1) + ax_main.text(PEAK_START + pd.Timedelta(hours=1), 0.97, "zoom", + transform=ax_main.get_xaxis_transform(), + fontsize=8, color="goldenrod", ha="left", va="top") + + # Bottom panel: peak zoom + plot_envelope(ax_zoom, env_peak, obs, PEAK_START, PEAK_END, helene_label=True) + ax_zoom.set_xlabel("Date", fontsize=11) + ax_zoom.set_title( + "Helene peak zoom (Sep 26 12 UTC - Sep 28 06 UTC)", fontsize=10) + _format_xaxis(ax_zoom, major_interval=1, minor_hours=3) + ax_zoom.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %HUTC")) + ax_zoom.legend(fontsize=10, loc="lower right", frameon=True, framealpha=0.9) + + fig.suptitle( + f"4b -- 600-member crossed ensemble vs USGS 03463300 | {args.cat_id}\n" + "30 met forcing draws x 20 DA analysis state draws | R=0.07 mm2/h2", + fontsize=12, y=1.01, + ) + plt.tight_layout() + out2 = os.path.join(out_dir, f"{args.cat_id}_4b_crossed_ensemble_peak.png") + plt.savefig(out2, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_perturbation_category_shaded.py b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_perturbation_category_shaded.py new file mode 100644 index 00000000..af2e9ac9 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_perturbation_category_shaded.py @@ -0,0 +1,163 @@ +""" +Per-catchment shaded ensemble-band plot, organized by perturbation category. +F5 re-kriged variance direct (1 gauge holdout). + +Three categories: + 1. Initial states (red) + 2. Meteorological forcings (blue) + 3. Hydrological states (green) + +For each category, all 20 ensemble members are shown as a shaded band +(min-max envelope fill) plus a thicker median line in the same color. +Re-kriged Qkrig observation overlaid in black. Hurricane Helene peak window +shaded in pink. + +Inputs (existing per-source sensitivity CSVs from run_perturbation_sensitivity.py): + //_sensitivity_init.csv (20 members, init only) + //_sensitivity_forcing.csv (20 members, forcing only) + //_sensitivity_process.csv (20 members, process noise only) + /.csv (re-kriged obs, qkrig_mm_hr column) + +Outputs: + //_perturbation_categories_linear.png + //_perturbation_categories_log.png +""" +import argparse +import os +import pandas as pd +import numpy as np +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +CAT = "cat-1016300" + +SEN_DIR = "/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder5_rekrig_variance_direct" +OBS_DIR = "/mnt/disk2/1400_sites_helene/catchment_ts_no_03463300_dynamic_variance_rekrig" + +OUT_LINEAR = os.path.join(SEN_DIR, CAT, f"{CAT}_perturbation_categories_linear.png") +OUT_LOG = os.path.join(SEN_DIR, CAT, f"{CAT}_perturbation_categories_log.png") + +PLOT_START = pd.Timestamp("2024-09-20") +PLOT_END = pd.Timestamp("2024-10-05") + +HELENE_START = pd.Timestamp("2024-09-26 12:00:00") +HELENE_END = pd.Timestamp("2024-09-28 00:00:00") + +CATEGORIES = [ + ("init", "Initial states only", "tab:red"), + ("forcing", "Meteorological forcings only", "tab:blue"), + ("process", "Hydrological states only", "tab:green"), +] + + +def load_members(source): + path = os.path.join(SEN_DIR, CAT, f"{CAT}_sensitivity_{source}.csv") + if not os.path.exists(path): + return None, None + df = pd.read_csv(path, parse_dates=["date"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + return df["date"].values, df[member_cols].to_numpy(dtype=float) + + +def load_obs(): + p = os.path.join(OBS_DIR, f"{CAT}.csv") + if not os.path.exists(p): + return None, None + df = pd.read_csv(p) + time_col = next((c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")), + df.columns[0]) + df[time_col] = pd.to_datetime(df[time_col]) + # rekrig files use qkrig_mm_hr column + q_col = next((c for c in df.columns + if "qkrig" in c.lower() and "var" not in c.lower()), None) + if q_col is None: + q_col = next(c for c in df.columns + if c != time_col and "var" not in c.lower() + and "step" not in c.lower() + and pd.api.types.is_numeric_dtype(df[c])) + return df[time_col].values, df[q_col].values + + +def plot_panel(ax, obs_dates, obs_vals, log_y=False): + for source, label, color in CATEGORIES: + dates, q = load_members(source) + if dates is None: + continue + d = pd.to_datetime(dates) + mask = (d >= PLOT_START) & (d <= PLOT_END) + if mask.sum() == 0: + continue + q_window = q[mask, :] + qmin = np.nanmin(q_window, axis=1) + qmax = np.nanmax(q_window, axis=1) + median = np.nanmedian(q_window, axis=1) + + ax.fill_between(d[mask], qmin, qmax, + color=color, alpha=0.30, zorder=2, edgecolor="none") + ax.plot(d[mask], median, + color=color, lw=1.7, alpha=0.95, zorder=3, + label=f"{label} (N=20)") + + ax.axvspan(HELENE_START, HELENE_END, color="salmon", alpha=0.15, zorder=1) + + if obs_dates is not None: + od = pd.to_datetime(obs_dates) + om = (od >= PLOT_START) & (od <= PLOT_END) + ax.plot(od[om], obs_vals[om], + color="black", lw=1.4, label="Qkrig re-kriged (obs)", zorder=4) + + ax.set_xlabel("Date", fontsize=10) + if log_y: + ax.set_yscale("log") + ax.set_ylabel(r"q (mm hr$^{-1}$) [log scale]", fontsize=10) + ax.set_ylim(0.01, None) + else: + ax.set_ylabel("Discharge (mm/h)", fontsize=10) + ax.tick_params(axis="x", rotation=30, labelsize=9) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=2)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.grid(True, alpha=0.2) + ax.legend(loc="upper left", fontsize=9, frameon=True, framealpha=0.9) + + +def main(): + global CAT, SEN_DIR, OUT_LINEAR, OUT_LOG + parser = argparse.ArgumentParser() + parser.add_argument("--cat-id", default=CAT) + args = parser.parse_args() + CAT = args.cat_id + OUT_LINEAR = os.path.join(SEN_DIR, CAT, f"{CAT}_perturbation_categories_linear.png") + OUT_LOG = os.path.join(SEN_DIR, CAT, f"{CAT}_perturbation_categories_log.png") + + obs_dates, obs_vals = load_obs() + + fig, ax = plt.subplots(1, 1, figsize=(20, 7)) + plot_panel(ax, obs_dates, obs_vals, log_y=False) + fig.suptitle( + f"Ensemble forecast by perturbation category — {CAT} — F5 (re-kriged variance)\n" + "Sep 20 - Oct 5, 2024 (Hurricane Helene window) | " + "Shaded bands = min-max envelope across 20 members. Lines = ensemble median.", + fontsize=11, y=0.995, + ) + plt.tight_layout() + plt.savefig(OUT_LINEAR, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_LINEAR}") + + fig, ax = plt.subplots(1, 1, figsize=(20, 7)) + plot_panel(ax, obs_dates, obs_vals, log_y=True) + fig.suptitle( + f"Ensemble forecast by perturbation category — {CAT} — F5 (re-kriged variance) — log scale\n" + "Sep 20 - Oct 5, 2024 (Hurricane Helene window) | " + "Shaded bands = min-max envelope across 20 members. Lines = ensemble median.", + fontsize=11, y=0.995, + ) + plt.tight_layout() + plt.savefig(OUT_LOG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_LOG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_combined.py b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_combined.py new file mode 100644 index 00000000..c6d531a0 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_combined.py @@ -0,0 +1,235 @@ +#!/usr/bin/env python3 +""" +plot_routed_ensemble_combined.py +Combined poster plot: 600-member routed ensemble envelope + +deterministic DA lines (Vrugt/no-Vrugt) vs USGS gauge obs at 03463300. + +Inputs: + --vrugt-csv : routed_Q_test.csv from vrugt_dynamic_routed/ + --novrugt-csv : routed_Q_test.csv from novrugt_dynamic_routed/ + --ensemble-pq : routed_crossed_ensemble.parquet from run_route_crossed_ensemble.py + --out-dir : output directory + +The ensemble parquet columns: issue_time, lead_hour, member_0000..member_0599 (Q in m3/s) + +Usage: + python3 plot_routed_ensemble_combined.py \ + --vrugt-csv /mnt/disk2/suma_helen_poster/da_results/vrugt_dynamic_routed/routed_Q_test.csv \ + --novrugt-csv /mnt/disk2/suma_helen_poster/da_results/novrugt_dynamic_routed/routed_Q_test.csv \ + --ensemble-pq /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble_routed/routed_crossed_ensemble.parquet \ + --out-dir /mnt/disk2/suma_helen_poster/da_results/comparison_plots +""" + +import argparse +import os + +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +HELENE_START = pd.Timestamp("2024-09-25 00:00:00") +HELENE_END = pd.Timestamp("2024-09-29 00:00:00") +ZOOM_START = pd.Timestamp("2024-09-24 00:00:00") +ZOOM_END = pd.Timestamp("2024-09-28 18:00:00") # cut trailing ensemble bump + +COLOR_OBS = "black" +COLOR_VRUGT = "#1f77b4" # blue +COLOR_NOVRUGT = "#d62728" # red +COLOR_ENS = "#1f77b4" # same blue family as Vrugt + + +def load_det_csv(path): + df = pd.read_csv(path, parse_dates=["date"]) + return df.set_index("date").sort_index() + + +def build_envelope(pq_path, zoom_start, zoom_end): + """ + Load routed ensemble parquet and build percentile envelope by valid_time. + Returns DataFrame indexed by valid_time with columns p05, p25, p50, p75, p95. + """ + df = pd.read_parquet(pq_path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = [c for c in df.columns if c.startswith("member_")] + + records = [] + for _, row in df.iterrows(): + t0 = row["issue_time"] + lead = int(row["lead_hour"]) + vt = t0 + pd.Timedelta(hours=lead) + vals = row[member_cols].values.astype(float) + records.append({"valid_time": vt, "vals": vals}) + + # Group by valid_time (floor to hour), stack all member values across issue_times + from collections import defaultdict + groups = defaultdict(list) + for rec in records: + vt = rec["valid_time"].floor("1h") + groups[vt].extend(rec["vals"].tolist()) + + rows = [] + for vt in sorted(groups): + if vt < zoom_start or vt > zoom_end: + continue + v = np.array(groups[vt]) + v = v[~np.isnan(v)] + if len(v) == 0: + continue + rows.append({ + "valid_time": vt, + "p05": np.percentile(v, 5), + "p25": np.percentile(v, 25), + "p50": np.percentile(v, 50), + "p75": np.percentile(v, 75), + "p95": np.percentile(v, 95), + }) + + return pd.DataFrame(rows).set_index("valid_time") + + +def compute_kge(obs, sim): + mask = ~(np.isnan(obs) | np.isnan(sim)) + o, s = obs[mask], sim[mask] + if len(o) < 2 or np.std(o) == 0: + return np.nan + r = np.corrcoef(o, s)[0, 1] + return 1.0 - np.sqrt((r-1)**2 + (np.std(s)/np.std(o)-1)**2 + (np.mean(s)/np.mean(o)-1)**2) + + +def _format_xaxis(ax, locator, fmt, rotation=30): + ax.xaxis.set_major_locator(locator) + ax.xaxis.set_major_formatter(mdates.DateFormatter(fmt)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=rotation, ha="right", fontsize=9) + + +def _add_helene_band(ax): + ax.axvspan(HELENE_START, HELENE_END, color="gold", alpha=0.15, zorder=0) + + +def plot_zoom(ax, dates, obs, sim_v, sim_nv, kge_v, kge_nv, env=None): + mask = (dates >= ZOOM_START) & (dates <= ZOOM_END) + dz = dates[mask] + + # Ensemble envelope (bottom layer) + if env is not None: + ax.fill_between(env.index, env["p05"], env["p95"], + color=COLOR_ENS, alpha=0.18, label="Ensemble 5-95th pct") + ax.fill_between(env.index, env["p25"], env["p75"], + color=COLOR_ENS, alpha=0.30, label="Ensemble 25-75th pct") + ax.plot(env.index, env["p50"], + color=COLOR_ENS, lw=1.2, linestyle="--", alpha=0.7, label="Ensemble median") + + # Deterministic DA lines + ax.plot(dz, sim_nv[mask], color=COLOR_NOVRUGT, lw=1.4, linestyle="--", + label=f"DA -- constant R KGE={kge_nv:.3f}", zorder=3) + ax.plot(dz, sim_v[mask], color=COLOR_VRUGT, lw=1.6, + label=f"DA -- dynamic R (Vrugt) KGE={kge_v:.3f}", zorder=4) + + # USGS obs on top + ax.plot(dz, obs[mask], color=COLOR_OBS, lw=1.6, label="USGS obs (gauge 03463300)", zorder=5) + + peak_usgs = np.nanmax(obs[mask]) + ax.axhline(peak_usgs, color=COLOR_OBS, lw=0.7, linestyle=":", alpha=0.5) + # Place annotation to the right of the peak to avoid legend overlap + ax.text(pd.Timestamp("2024-09-28 06:00:00"), peak_usgs * 1.02, + f"USGS peak {peak_usgs:.0f} m³/s", fontsize=8.5, color="black", + alpha=0.8, ha="right") + + _add_helene_band(ax) + _format_xaxis(ax, mdates.DayLocator(interval=1), "%b %d") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--vrugt-csv", required=True) + parser.add_argument("--novrugt-csv", required=True) + parser.add_argument("--ensemble-pq", required=True, + help="routed_crossed_ensemble.parquet from run_route_crossed_ensemble.py") + parser.add_argument("--out-dir", required=True) + args = parser.parse_args() + os.makedirs(args.out_dir, exist_ok=True) + + vrugt = load_det_csv(args.vrugt_csv) + novrugt = load_det_csv(args.novrugt_csv) + + obs = vrugt["Q_usgs_m3s"].values + dates = vrugt.index + sim_v = vrugt["Q_routed_m3s"].values + sim_nv = novrugt["Q_routed_m3s"].reindex(dates).values + + kge_v = compute_kge(obs, sim_v) + kge_nv = compute_kge(obs, sim_nv) + + print("Building ensemble envelope...") + env = build_envelope(args.ensemble_pq, ZOOM_START, ZOOM_END) + print(f" Envelope: {len(env)} timesteps, " + f"peak p95={env['p95'].max():.1f} m3/s, peak p50={env['p50'].max():.1f} m3/s") + + peak_usgs = np.nanmax(obs[(dates >= ZOOM_START) & (dates <= ZOOM_END)]) + in_env = env["p95"].max() >= peak_usgs + print(f" USGS peak {peak_usgs:.1f} m3/s {'IS' if in_env else 'IS NOT'} within p95 envelope") + + # ------------------------------------------------------------------ # + # Figure 1: Helene zoom only -- main poster panel + # ------------------------------------------------------------------ # + fig, ax = plt.subplots(figsize=(11, 5.5)) + plot_zoom(ax, dates, obs, sim_v, sim_nv, kge_v, kge_nv, env) + ax.set_xlabel("Date (UTC)", fontsize=11) + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_title( + "Hurricane Helene -- Probabilistic discharge forecast at gauge 03463300 (South Toe River)\n" + "600-member CFE ensemble + DA (Muskingum routing)", + fontsize=11) + ax.legend(fontsize=9, loc="lower right") + ax.grid(True, alpha=0.25) + plt.tight_layout() + out1 = os.path.join(args.out_dir, "helene_ensemble_vs_usgs.png") + plt.savefig(out1, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + # ------------------------------------------------------------------ # + # Figure 2: two-panel (full period + Helene zoom) + # ------------------------------------------------------------------ # + fig, (ax_top, ax_bot) = plt.subplots( + 2, 1, figsize=(14, 9), + gridspec_kw={"height_ratios": [1, 1.5]}) + + # Top: full test period (deterministic only, no ensemble for readability) + ax_top.plot(dates, obs, color=COLOR_OBS, lw=0.8, label="USGS obs", zorder=3) + ax_top.plot(dates, sim_v, color=COLOR_VRUGT, lw=0.9, + label=f"DA dynamic R (Vrugt) KGE={kge_v:.3f}", zorder=2) + ax_top.plot(dates, sim_nv, color=COLOR_NOVRUGT, lw=0.9, linestyle="--", + label=f"DA constant R KGE={kge_nv:.3f}", zorder=2) + ax_top.axvspan(ZOOM_START, ZOOM_END, color="gray", alpha=0.12, zorder=0) + ax_top.axvspan(HELENE_START, HELENE_END, color="gold", alpha=0.15, zorder=0) + ax_top.text(HELENE_START + pd.Timedelta(days=0.5), + np.nanmax(obs) * 0.88, "Helene", + fontsize=9, color="goldenrod", fontweight="bold") + ax_top.set_ylabel("Discharge (m³/s)", fontsize=10) + ax_top.set_title( + "Routed discharge at gauge 03463300 (South Toe River) -- Test period\n" + "DA with dynamic Vrugt R vs. constant R -- Muskingum routing", + fontsize=11) + ax_top.legend(fontsize=9, loc="upper left") + ax_top.grid(True, alpha=0.25) + _format_xaxis(ax_top, mdates.MonthLocator(interval=2), "%Y-%m") + + # Bottom: Helene zoom with ensemble + plot_zoom(ax_bot, dates, obs, sim_v, sim_nv, kge_v, kge_nv, env) + ax_bot.set_xlabel("Date (UTC)", fontsize=10) + ax_bot.set_ylabel("Discharge (m³/s)", fontsize=10) + ax_bot.set_title("Hurricane Helene window -- 600-member ensemble envelope", fontsize=10) + ax_bot.legend(fontsize=9, loc="lower right") + ax_bot.grid(True, alpha=0.25) + + plt.tight_layout() + out2 = os.path.join(args.out_dir, "helene_ensemble_twopanel.png") + plt.savefig(out2, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_vs_usgs.py b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_vs_usgs.py new file mode 100644 index 00000000..55c52d85 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_vs_usgs.py @@ -0,0 +1,176 @@ +#!/usr/bin/env python3 +""" +plot_routed_ensemble_vs_usgs.py +Plot T-route routed 600-member ensemble envelope vs USGS at gauge 03463300. + +Input: routed_crossed_ensemble.parquet from run_route_crossed_ensemble.py + Columns: issue_time, lead_hour, member_0000..member_0599 (Q in m3/s) + +Usage: + python3 plot_routed_ensemble_vs_usgs.py \ + --routed-pq /mnt/disk2/.../folder1_vrugt_routed/routed_crossed_ensemble.parquet \ + --usgs-csv /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv \ + --label "F1 Vrugt — 1 gauge holdout" \ + --out-dir ~/plots_f1_1gauge +""" + +import argparse +from collections import defaultdict +from pathlib import Path + +import matplotlib.dates as mdates +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd + +WATERSHED_AREA_KM2 = 113.18 # cat-1016300 drainage area +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +PLOT_START = pd.Timestamp("2024-09-24 00:00:00") +PLOT_END = pd.Timestamp("2024-09-28 18:00:00") +PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +PEAK_END = pd.Timestamp("2024-09-28 06:00:00") +HELENE_PEAK = pd.Timestamp("2024-09-27 12:00:00") + +COLOR_ENS = "#1f77b4" +COLOR_OBS = "black" + + +def build_envelope(pq_path, t_start, t_end): + df = pd.read_parquet(pq_path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = [c for c in df.columns if c.startswith("member_")] + + groups = defaultdict(list) + for _, row in df.iterrows(): + vt = (row["issue_time"] + pd.Timedelta(hours=int(row["lead_hour"]))).floor("1h") + if t_start <= vt <= t_end: + groups[vt].extend(row[member_cols].values.astype(float).tolist()) + + rows = [] + for vt in sorted(groups): + v = np.array(groups[vt]) + v = v[~np.isnan(v)] + if len(v) == 0: + continue + rows.append({"valid_time": vt, + "p05": np.percentile(v, 5), "p25": np.percentile(v, 25), + "p50": np.percentile(v, 50), "p75": np.percentile(v, 75), + "p95": np.percentile(v, 95)}) + return pd.DataFrame(rows).set_index("valid_time") + + +def load_usgs(path): + df = pd.read_csv(path) + print(f" USGS CSV columns: {list(df.columns)}") + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() or "discharge" in c.lower()) + print(f" Using date_col={date_col!r}, q_col={q_col!r}") + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series.sort_index() + + +def compute_kge(obs, sim): + mask = ~(np.isnan(obs) | np.isnan(sim)) + o, s = obs[mask], sim[mask] + if len(o) < 2 or np.std(o) == 0: + return np.nan + r = np.corrcoef(o, s)[0, 1] + return 1 - np.sqrt((r-1)**2 + (np.std(s)/np.std(o)-1)**2 + (np.mean(s)/np.mean(o)-1)**2) + + +def _shade(ax, env): + ax.fill_between(env.index, env["p05"], env["p95"], + color=COLOR_ENS, alpha=0.18, label="5th–95th percentile") + ax.fill_between(env.index, env["p25"], env["p75"], + color=COLOR_ENS, alpha=0.32, label="25th–75th percentile") + ax.plot(env.index, env["p50"], color=COLOR_ENS, lw=1.5, label="Ensemble median (600 members)") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--routed-pq", required=True) + parser.add_argument("--usgs-csv", required=True) + parser.add_argument("--label", default="F1 Vrugt — 1 gauge holdout") + parser.add_argument("--out-dir", required=True) + args = parser.parse_args() + + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + print("Building full-window envelope...") + env_full = build_envelope(args.routed_pq, PLOT_START, PLOT_END) + print("Building peak-window envelope...") + env_peak = build_envelope(args.routed_pq, PEAK_START, PEAK_END) + + print("Loading USGS obs...") + try: + obs_all = load_usgs(args.usgs_csv) + obs_full = obs_all.loc[PLOT_START:PLOT_END] + obs_peak = obs_all.loc[PEAK_START:PEAK_END] + print(f" Loaded {len(obs_all)} USGS obs rows; " + f"full window: {len(obs_full)}, peak window: {len(obs_peak)}") + except Exception as e: + print(f" Warning: USGS load failed ({e}) — plotting without obs") + obs_full = obs_peak = None + + kge = np.nan + if obs_full is not None and len(env_full) > 0: + med = env_full["p50"] + obs_hourly = obs_full.copy() + obs_hourly.index = obs_hourly.index.floor("1h") + obs_hourly = obs_hourly[~obs_hourly.index.duplicated(keep="first")] + aligned = obs_hourly.reindex(med.index) + print(f" KGE alignment: {aligned.notna().sum()} matched out of {len(med)} timesteps") + kge = compute_kge(aligned.values, med.values) + + # ── Figure 1: Full Helene window ────────────────────────────────────────── + fig, ax = plt.subplots(figsize=(12, 4.5)) + _shade(ax, env_full) + if obs_full is not None: + ax.plot(obs_full.index, obs_full.values, color=COLOR_OBS, lw=1.6, + label="USGS obs (gauge 03463300)") + ax.axvspan(PEAK_START, PEAK_END, color="salmon", alpha=0.12, zorder=0, label="Helene peak window") + ax.set_ylabel("Discharge (m³/s)") + ax.set_xlabel("Date") + ax.set_title(f"Routed 600-member ensemble vs USGS | {args.label}\n" + f"Ensemble median KGE = {kge:.3f} | Sep 24–28 2024") + ax.legend(fontsize=9, loc="upper left") + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + plt.tight_layout() + out = out_dir / "routed_ensemble_vs_usgs_full.png" + fig.savefig(out, dpi=150) + print(f"Saved: {out}") + plt.close(fig) + + # ── Figure 2: Peak zoom ─────────────────────────────────────────────────── + fig, ax = plt.subplots(figsize=(10, 4.5)) + _shade(ax, env_peak) + if obs_peak is not None: + ax.plot(obs_peak.index, obs_peak.values, color=COLOR_OBS, lw=1.6, + label="USGS obs (gauge 03463300)") + peak_val = float(obs_peak.max()) + ax.axhline(peak_val, color=COLOR_OBS, lw=0.7, ls=":", alpha=0.5) + ax.annotate(f"USGS peak\n{peak_val:.0f} m³/s", + xy=(HELENE_PEAK, peak_val), xytext=(10, -40), + textcoords="offset points", fontsize=8, color=COLOR_OBS) + ax.set_ylabel("Discharge (m³/s)") + ax.set_xlabel("Date") + ax.set_title(f"Helene peak zoom | {args.label}") + ax.legend(fontsize=9) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %HUTC")) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + plt.tight_layout() + out = out_dir / "routed_ensemble_vs_usgs_peak.png" + fig.savefig(out, dpi=150) + print(f"Saved: {out}") + plt.close(fig) + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4b_ensemble_vs_obs/run_4b_crossed_f5.sh b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4b_ensemble_vs_obs/run_4b_crossed_f5.sh new file mode 100644 index 00000000..82f4ca11 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4b_ensemble_vs_obs/run_4b_crossed_f5.sh @@ -0,0 +1,40 @@ +#!/usr/bin/env bash +# F5 re-kriged variance direct (1 gauge holdout) — 4b crossed ensemble plots. +# Runs plot_crossed_ensemble.py for all 21 catchments. +# Reads //_crossed_ensemble.parquet (already generated). +# +# Usage: +# bash run_4b_crossed_f5.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +DATA_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder5_rekrig_variance_direct +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +echo "[F5-4b-crossed] Crossed ensemble plots — data: $DATA_DIR" + +for CAT in "${CATS[@]}"; do + PQ="$DATA_DIR/$CAT/${CAT}_crossed_ensemble.parquet" + if [ ! -f "$PQ" ]; then + echo " [$CAT] No crossed ensemble parquet — skipping" + continue + fi + echo " [$CAT] crossed ensemble..." + $TROUTE "$SCRIPT_DIR/plot_crossed_ensemble.py" \ + --ensemble-dir "$DATA_DIR" \ + --cat-id "$CAT" \ + --usgs-csv "$USGS_CSV" +done + +echo "[F5-4b-crossed] Done." diff --git a/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4b_ensemble_vs_obs/run_4b_f5.sh b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4b_ensemble_vs_obs/run_4b_f5.sh new file mode 100644 index 00000000..bac0de06 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4b_ensemble_vs_obs/run_4b_f5.sh @@ -0,0 +1,66 @@ +#!/usr/bin/env bash +# F5 Re-kriging dynamic variance (1 gauge holdout) — 4b ensemble plots. +# +# Runs: +# 1. run_4b_crossed_f5.sh — per-catchment crossed ensemble PNGs +# 2. plot_routed_ensemble_vs_usgs.py — F5 ensemble envelope vs USGS at outlet +# 3. plot_routed_ensemble_combined.py — F1 Vrugt vs F5 re-kriging comparison +# +# Must be run AFTER: +# - crossed ensemble parquets exist (Kunal has done this) +# - routed_crossed_ensemble.parquet exists (Kunal has done this) +# - routed_Q_test.csv exists (run deterministic routing first) +# +# Usage: +# bash run_4b_f5.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +F1_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt +F5_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder5_rekrig_variance_direct +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +# 1. Per-catchment crossed ensemble plots +echo "[F5-4b] Per-catchment crossed ensemble plots..." +for CAT in "${CATS[@]}"; do + PQ="$F5_DIR/$CAT/${CAT}_crossed_ensemble.parquet" + if [ ! -f "$PQ" ]; then + echo " [$CAT] No parquet — skipping" + continue + fi + echo " [$CAT] crossed ensemble..." + $TROUTE "$SCRIPT_DIR/plot_crossed_ensemble.py" \ + --ensemble-dir "$F5_DIR" \ + --cat-id "$CAT" \ + --usgs-csv "$USGS_CSV" +done + +# 2. Routed ensemble vs USGS (F5 only) +echo "[F5-4b] Routed ensemble vs USGS (F5 only)..." +$TROUTE "$SCRIPT_DIR/plot_routed_ensemble_vs_usgs.py" \ + --routed-pq "$F5_DIR/routed_crossed_ensemble.parquet" \ + --usgs-csv "$USGS_CSV" \ + --label "F5 Re-kriging variance — 1 gauge holdout" \ + --out-dir "$F5_DIR" + +# 3. Combined: F1 Vrugt vs F5 re-kriging +echo "[F5-4b] Combined comparison: F1 Vrugt vs F5 Re-kriging..." +$TROUTE "$SCRIPT_DIR/plot_routed_ensemble_combined.py" \ + --vrugt-csv "$F1_DIR/routed_Q_test.csv" \ + --novrugt-csv "$F5_DIR/routed_Q_test.csv" \ + --ensemble-pq "$F5_DIR/routed_crossed_ensemble.parquet" \ + --out-dir "$F5_DIR" + +echo "[F5-4b] Done." +ls "$F5_DIR"/*.png 2>/dev/null || echo " (no root PNGs yet)" diff --git a/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4c_timeseries/plot_forecast_spaghetti.py b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4c_timeseries/plot_forecast_spaghetti.py new file mode 100644 index 00000000..f447d456 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4c_timeseries/plot_forecast_spaghetti.py @@ -0,0 +1,167 @@ +""" +plot_forecast_spaghetti.py — F5 (re-kriged variance direct) + +"All forecasts ending at this hour" — verification-time-fixed spaghetti view. + +Window: Sep 24 18 UTC -> Sep 30 06 UTC. +For every initialization time whose 18-hour forecast overlaps this window: + DA : shaded min/max band across 20 members + ensemble mean line + OL : ensemble mean line only (no shading), dashed + +Lines are colored by initialization DATE (7 colors, Sep 24-30). +USGS obs overlaid in black. + +Inputs: + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + Wide format: issue_time, lead_hour, member_00..member_19 (q in m3/s) + +Obs: + /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +Outputs: + /forecast_spaghetti_helene.png +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import matplotlib.patches as mpatches + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder5_rekrig_variance_direct_leadtime_routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +PLOT_START = pd.Timestamp("2024-09-24 18:00:00") +PLOT_END = pd.Timestamp("2024-09-30 06:00:00") + +HELENE_PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +HELENE_PEAK_END = pd.Timestamp("2024-09-28 00:00:00") + +DATE_COLORS = { + "2024-09-24": "#1f77b4", + "2024-09-25": "#ff7f0e", + "2024-09-26": "#2ca02c", + "2024-09-27": "#d62728", + "2024-09-28": "#9467bd", + "2024-09-29": "#8c564b", + "2024-09-30": "#e377c2", +} + + +def load_parquet(path): + df = pd.read_parquet(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + df["ens_mean"] = df[member_cols].mean(axis=1) + df["ens_min"] = df[member_cols].min(axis=1) + df["ens_max"] = df[member_cols].max(axis=1) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + return df + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() + or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--route-dir", default=DEFAULT_ROUTE_DIR) + parser.add_argument("--out-dir", default=None) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--da-name", default="routed_leadtime_da_full.parquet") + parser.add_argument("--ol-name", default="routed_leadtime_openloop_full.parquet") + args = parser.parse_args() + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + print("Loading parquets...") + da = load_parquet(os.path.join(args.route_dir, args.da_name)) + ol = load_parquet(os.path.join(args.route_dir, args.ol_name)) + obs = load_usgs(args.usgs_csv) + + da = da[(da["valid_time"] >= PLOT_START) & (da["valid_time"] <= PLOT_END)] + ol = ol[(ol["valid_time"] >= PLOT_START) & (ol["valid_time"] <= PLOT_END)] + + issue_times = sorted(da["issue_time"].unique()) + print(f" {len(issue_times)} initialization times contribute to this window") + + fig, ax = plt.subplots(figsize=(18, 6)) + + ax.axvspan(HELENE_PEAK_START, HELENE_PEAK_END, + color="salmon", alpha=0.12, zorder=0) + ax.text(HELENE_PEAK_START + pd.Timedelta(hours=6), 1.0, "Helene peak", + transform=ax.get_xaxis_transform(), + fontsize=9, color="firebrick", ha="left", va="top") + + for t0 in issue_times: + date_str = str(pd.Timestamp(t0).date()) + color = DATE_COLORS.get(date_str, "gray") + + da_t = da[da["issue_time"] == t0].sort_values("valid_time") + ol_t = ol[ol["issue_time"] == t0].sort_values("valid_time") + + if da_t.empty: + continue + + ax.fill_between(da_t["valid_time"], da_t["ens_min"], da_t["ens_max"], + color=color, alpha=0.06, zorder=2) + ax.plot(da_t["valid_time"], da_t["ens_mean"], + color=color, lw=0.8, alpha=0.55, zorder=3) + + if not ol_t.empty: + ax.plot(ol_t["valid_time"], ol_t["ens_mean"], + color=color, lw=0.6, alpha=0.30, linestyle="--", zorder=2) + + obs_w = obs.loc[PLOT_START:PLOT_END] + ax.plot(obs_w.index, obs_w.values, + color="black", lw=1.8, zorder=6, label="USGS obs") + + patches = [mpatches.Patch(color=c, label=f"Init {d}") + for d, c in DATE_COLORS.items()] + patches.append(plt.Line2D([0], [0], color="black", lw=1.8, label="USGS obs")) + patches.append(plt.Line2D([0], [0], color="gray", lw=1.2, alpha=0.7, + label="DA — ens mean (solid) + spread (shaded)")) + patches.append(plt.Line2D([0], [0], color="gray", lw=1.0, linestyle="--", + alpha=0.5, label="OL — ens mean (dashed)")) + ax.legend(handles=patches, fontsize=8, loc="upper left", + frameon=True, framealpha=0.9, ncol=2) + + ax.set_xlim(PLOT_START, PLOT_END) + ax.set_xlabel("Date", fontsize=11) + ax.set_ylabel("Discharge (m3/s)", fontsize=11) + ax.set_title( + "All 18-hour forecast trajectories ending in this window — USGS 03463300 — F5 (re-kriged variance)\n" + "Sep 24 18 UTC -> Sep 30 06 UTC | DA: shaded band + mean | OL: mean dashed", + fontsize=11, + ) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=1)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.xaxis.set_minor_locator(mdates.HourLocator(interval=6)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.2, lw=0.4) + plt.tight_layout() + + out_path = os.path.join(out_dir, "forecast_spaghetti_helene.png") + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4c_timeseries/plot_helene_issue_time_hydrograph.py b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4c_timeseries/plot_helene_issue_time_hydrograph.py new file mode 100644 index 00000000..15174787 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4c_timeseries/plot_helene_issue_time_hydrograph.py @@ -0,0 +1,208 @@ +""" +Per-issue-time forecast hydrograph diagnostic — F5 (re-kriged variance). + +For two contrasting issue times — one storm-peak (Helene) and one low-flow — +show what the DA and open-loop ensembles actually produce over the 18-hour +forecast window, alongside the kriging observation. + +Each panel shows: + - All 20 DA members as faint purple lines + median (thick purple) + - All 20 open-loop members as faint gray lines + median (thick gray, dashed) + - Kriging obs as a thick black line (context: t0-6 through t0+18) + - Vertical line at t0 + annotation with obs(t0) and lead-1 ensemble means + +Inputs (from batch_run_f5_lead_time_sweep.sh): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv +Obs: + //_test_results.csv if present, else + /.csv (qkrig_mm_hr column) + +Output: + //_issue_time_hydrograph_helene_vs_lowflow.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +F5_DIR = "/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder5_rekrig_variance_direct" +DEFAULT_KRIG_OBS_DIR = "/mnt/disk2/1400_sites_helene/catchment_ts_no_03463300_dynamic_variance_rekrig" + +DEFAULT_LEADTIME_DIR = F5_DIR +DEFAULT_DA_DIR = F5_DIR + +CAT = "cat-1016300" +LEADTIME_DIR = DEFAULT_LEADTIME_DIR +DA_DIR = DEFAULT_DA_DIR +KRIG_OBS_DIR = DEFAULT_KRIG_OBS_DIR + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" +OBS_COLOR = "black" + +HELENE_T0 = pd.Timestamp("2024-09-26 12:00:00") +LOWFLOW_T0 = pd.Timestamp("2024-03-15 12:00:00") + +CONTEXT_HOURS_BEFORE = 6 +LEAD_HOURS_AFTER = 18 + + +def load_forecasts(path): + df = pd.read_csv(path, parse_dates=['issue_time', 'valid_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + return df, member_cols + + +def load_obs(): + test_path = os.path.join(DA_DIR, CAT, f"{CAT}_test_results.csv") + if os.path.exists(test_path): + df = pd.read_csv(test_path, parse_dates=['date']) + return df.set_index('date')['obs_mm_h'] + krig_path = os.path.join(KRIG_OBS_DIR, f"{CAT}.csv") + if not os.path.exists(krig_path): + raise FileNotFoundError( + f"No obs found: tried {test_path} and {krig_path}. " + f"Pass --krig-obs-dir to the re-kriged obs directory.") + df = pd.read_csv(krig_path) + t_col = next(c for c in df.columns + if c.lower() in ('datetime', 'date', 'time', 'timestamp')) + obs_col = next( + (c for c in df.columns + if c != t_col and 'var' not in c.lower() + and any(k in c.lower() for k in ('obs', 'q', 'krig', 'mm', 'flow', 'discharge')) + and pd.api.types.is_numeric_dtype(df[c])), + None, + ) + if obs_col is None: + obs_col = next(c for c in df.columns + if c != t_col and 'var' not in c.lower() + and 'step' not in c.lower() and 'time' not in c.lower() + and 'index' not in c.lower() + and pd.api.types.is_numeric_dtype(df[c])) + df[t_col] = pd.to_datetime(df[t_col]) + return df.set_index(t_col)[obs_col].astype(float) + + +def slice_forecast(df, member_cols, t0): + sub = df[df['issue_time'] == t0].sort_values('lead_hour') + if len(sub) == 0: + raise ValueError(f"No forecast rows for issue_time={t0}") + valid_times = pd.to_datetime(sub['valid_time'].values) + member_arr = sub[member_cols].to_numpy(dtype=float) + return valid_times, member_arr + + +def plot_panel(ax, t0, df_da, df_ol, member_cols, obs_series, title_extra=""): + vt_da, q_da = slice_forecast(df_da, member_cols, t0) + vt_ol, q_ol = slice_forecast(df_ol, member_cols, t0) + + ctx_start = t0 - pd.Timedelta(hours=CONTEXT_HOURS_BEFORE) + ctx_end = t0 + pd.Timedelta(hours=LEAD_HOURS_AFTER) + obs_window = obs_series.loc[ctx_start:ctx_end] + + for j in range(q_ol.shape[1]): + ax.plot(vt_ol, q_ol[:, j], color=OL_COLOR, lw=0.6, alpha=0.45, zorder=2) + for j in range(q_da.shape[1]): + ax.plot(vt_da, q_da[:, j], color=DA_COLOR, lw=0.6, alpha=0.45, zorder=3) + ax.plot(vt_ol, np.nanmedian(q_ol, axis=1), color=OL_COLOR, lw=2.6, + linestyle="--", zorder=4, label="Open-loop median") + ax.plot(vt_da, np.nanmedian(q_da, axis=1), color=DA_COLOR, lw=2.6, + zorder=5, label="DA median") + ax.plot(obs_window.index, obs_window.values, color=OBS_COLOR, lw=1.8, + zorder=6, label="Qkrig (obs)") + + ax.axvline(t0, color="black", lw=1.0, linestyle=":", alpha=0.6, zorder=1) + obs_t0 = float(obs_series.get(t0, np.nan)) + obs_lead1 = float(obs_series.get(t0 + pd.Timedelta(hours=1), np.nan)) + da_lead1_mean = float(np.nanmean(q_da[0, :])) + ol_lead1_mean = float(np.nanmean(q_ol[0, :])) + + info_text = ( + f"t0 = {t0.strftime('%Y-%m-%d %H:%M')}\n" + f"obs(t0) = {obs_t0:.3f} mm/h\n" + f"obs(t0+1) = {obs_lead1:.3f} mm/h\n" + f"DA mean(t0+1) = {da_lead1_mean:.3f} mm/h\n" + f"OL mean(t0+1) = {ol_lead1_mean:.3f} mm/h" + ) + ax.text(0.02, 0.97, info_text, + transform=ax.transAxes, + fontsize=9, va='top', ha='left', family='monospace', + bbox=dict(boxstyle='round,pad=0.4', facecolor='white', + alpha=0.85, edgecolor='gray')) + + ax.set_title(title_extra, fontsize=12) + ax.set_ylabel("Discharge (mm/h)", fontsize=10) + ax.set_xlabel("Date", fontsize=10) + ax.xaxis.set_major_locator(mdates.HourLocator(interval=3)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %H:%M")) + ax.tick_params(axis='x', rotation=30, labelsize=8) + ax.grid(True, alpha=0.2) + ax.legend(loc='upper right', fontsize=9, frameon=True, framealpha=0.9) + + +def snap_to_nearest(t0, available_times): + times = pd.DatetimeIndex(available_times) + idx = np.abs((times - t0).total_seconds()).argmin() + snapped = times[idx] + if snapped != t0: + print(f" Snapped {t0} -> {snapped} (nearest available issue time)") + return snapped + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', default="cat-1016300") + parser.add_argument('--leadtime-dir', default=DEFAULT_LEADTIME_DIR) + parser.add_argument('--da-dir', default=DEFAULT_DA_DIR) + parser.add_argument('--krig-obs-dir', default=DEFAULT_KRIG_OBS_DIR, + help='Dir holding per-catchment re-kriged obs CSVs ' + '(used when _test_results.csv is absent)') + parser.add_argument('--helene-t0', default=None) + parser.add_argument('--lowflow-t0', default=None) + args = parser.parse_args() + + global CAT, LEADTIME_DIR, DA_DIR, KRIG_OBS_DIR + CAT = args.cat_id + LEADTIME_DIR = args.leadtime_dir + DA_DIR = args.da_dir + KRIG_OBS_DIR = args.krig_obs_dir + out_png = os.path.join(LEADTIME_DIR, CAT, + f"{CAT}_issue_time_hydrograph_helene_vs_lowflow.png") + + da_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_da.csv") + ol_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_openloop.csv") + df_da, m_cols = load_forecasts(da_path) + df_ol, _ = load_forecasts(ol_path) + obs_series = load_obs() + + available = df_da['issue_time'].unique() + helene_t0 = snap_to_nearest( + pd.Timestamp(args.helene_t0) if args.helene_t0 else HELENE_T0, available) + lowflow_t0 = snap_to_nearest( + pd.Timestamp(args.lowflow_t0) if args.lowflow_t0 else pd.Timestamp("2024-03-15 00:00:00"), + available) + + fig, (ax_h, ax_l) = plt.subplots(1, 2, figsize=(20, 7)) + + plot_panel(ax_h, helene_t0, df_da, df_ol, m_cols, obs_series, + title_extra=f"Helene peak issue time — {helene_t0.strftime('%Y-%m-%d %H:%M')}") + plot_panel(ax_l, lowflow_t0, df_da, df_ol, m_cols, obs_series, + title_extra=f"Low-flow issue time — {lowflow_t0.strftime('%Y-%m-%d %H:%M')}") + + fig.suptitle( + f"Per-issue-time forecast hydrograph — {CAT} — F5 (re-kriged variance)\n" + "Faint lines = 20 ensemble members. Thick lines = ensemble medians. " + "Vertical dotted line = t0 (DA stops here for both scenarios).", + fontsize=12, y=1.00, + ) + plt.tight_layout() + plt.savefig(out_png, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_png}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4c_timeseries/plot_lead_time_reconstructed_timeseries.py b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4c_timeseries/plot_lead_time_reconstructed_timeseries.py new file mode 100644 index 00000000..5e1b84dd --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4c_timeseries/plot_lead_time_reconstructed_timeseries.py @@ -0,0 +1,234 @@ +""" +Reconstructed time series at gauge 03463300 — F5 (re-kriged variance direct). + +For each target hour t, pool ALL forecasts that land on t (across many +issue_times x many lead_hours x 20 members). Compute median + 5th/95th +percentile envelope. Plot DA and open-loop reconstructions overlaid on USGS +obs, with the Helene window shaded. + +Input parquets (from route_lead_time_forecasts.py): + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + Wide format: issue_time, lead_hour, member_00..member_19 (q in m3/s) + +Obs: + /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +Output: + /lead_time_reconstructed_timeseries.png + /lead_time_reconstructed_timeseries_helene.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" +OBS_COLOR = "black" + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder5_rekrig_variance_direct_leadtime_routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +PLOT_START_DEFAULT = pd.Timestamp("2024-09-10 00:00:00") +PLOT_END_DEFAULT = pd.Timestamp("2024-10-10 23:00:00") + +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-30 23:00:00") + + +def load_parquet_long(path): + df = pd.read_parquet(path) + df['issue_time'] = pd.to_datetime(df['issue_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + if not member_cols: + raise ValueError(f"No member columns in {path}") + long = df.melt( + id_vars=['issue_time', 'lead_hour'], + value_vars=member_cols, + var_name='member', + value_name='q_m3s', + ) + long['valid_time'] = (long['issue_time'] + + pd.to_timedelta(long['lead_hour'], unit='h')) + return long + + +def reconstruct(long_df): + grouped = long_df.groupby('valid_time')['q_m3s'].agg( + median='median', + p05=lambda s: s.quantile(0.05), + p95=lambda s: s.quantile(0.95), + count='count', + ).reset_index() + return grouped + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next((c for c in df.columns + if c.lower() in ('datetime', 'date', 'time', 'timestamp')), None) + q_col = next((c for c in df.columns + if 'q' in c.lower() or 'flow' in c.lower() + or 'discharge' in c.lower()), None) + if not date_col or not q_col: + raise ValueError(f"Couldn't id date/Q columns in {usgs_csv}: {df.columns.tolist()}") + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if 'mm' in q_col.lower(): + series = series * MM_H_TO_M3_S + print(f" Converted obs from mm/h (column '{q_col}') to m3/s " + f"using area = {WATERSHED_AREA_KM2} km2 (x {MM_H_TO_M3_S:.4f})") + else: + print(f" Loaded obs as-is from column '{q_col}' (assumed m3/s)") + return series + + +def nse(obs, sim): + m = np.isfinite(obs) & np.isfinite(sim) + if m.sum() < 2: + return np.nan + o, s = obs[m], sim[m] + denom = ((o - o.mean()) ** 2).sum() + if denom < 1e-10: + return np.nan + return 1.0 - float(((o - s) ** 2).sum() / denom) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--route-dir', default=DEFAULT_ROUTE_DIR) + parser.add_argument('--out-dir', default=None) + parser.add_argument('--usgs-csv', default=DEFAULT_USGS_CSV) + parser.add_argument('--plot-start', default=str(PLOT_START_DEFAULT)) + parser.add_argument('--plot-end', default=str(PLOT_END_DEFAULT)) + parser.add_argument('--da-name', default='routed_leadtime_da_full.parquet') + parser.add_argument('--ol-name', default='routed_leadtime_openloop_full.parquet') + args = parser.parse_args() + + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + plot_start = pd.Timestamp(args.plot_start) + plot_end = pd.Timestamp(args.plot_end) + + print(f"Loading DA parquet: {args.da_name}") + da_long = load_parquet_long(os.path.join(args.route_dir, args.da_name)) + print(f"Loading OL parquet: {args.ol_name}") + ol_long = load_parquet_long(os.path.join(args.route_dir, args.ol_name)) + print(f"Loading USGS obs: {args.usgs_csv}") + obs = load_usgs(args.usgs_csv) + + print("Reconstructing DA time series...") + da_rec = reconstruct(da_long) + print(f" {len(da_rec)} unique target hours") + print("Reconstructing OL time series...") + ol_rec = reconstruct(ol_long) + + da_rec = da_rec[(da_rec['valid_time'] >= plot_start) & (da_rec['valid_time'] <= plot_end)].copy() + ol_rec = ol_rec[(ol_rec['valid_time'] >= plot_start) & (ol_rec['valid_time'] <= plot_end)].copy() + da_rec['obs'] = da_rec['valid_time'].map(obs) + ol_rec['obs'] = ol_rec['valid_time'].map(obs) + + nse_da = nse(da_rec['obs'].values, da_rec['median'].values) + nse_ol = nse(ol_rec['obs'].values, ol_rec['median'].values) + print(f"NSE (plot window) — DA: {nse_da:.3f} | OL: {nse_ol:.3f}") + + fig, ax = plt.subplots(figsize=(16, 5.5)) + fig.patch.set_facecolor("white") + + ax.axvspan(HELENE_START, HELENE_END, color="firebrick", alpha=0.08, zorder=0) + ax.text(HELENE_START + pd.Timedelta(hours=12), 1.0, "Helene", + transform=ax.get_xaxis_transform(), + fontsize=9, color="firebrick", ha="left", va="top", fontweight="bold") + + ax.fill_between(ol_rec['valid_time'], ol_rec['p05'], ol_rec['p95'], + color=OL_COLOR, alpha=0.18, zorder=2, + label="Open-loop [5th-95th pctile envelope]") + ax.plot(ol_rec['valid_time'], ol_rec['median'], + color=OL_COLOR, lw=1.2, linestyle='--', zorder=3, + label=f"Open-loop median (NSE={nse_ol:.3f})") + + ax.fill_between(da_rec['valid_time'], da_rec['p05'], da_rec['p95'], + color=DA_COLOR, alpha=0.22, zorder=4, + label="DA [5th-95th pctile envelope]") + ax.plot(da_rec['valid_time'], da_rec['median'], + color=DA_COLOR, lw=1.4, zorder=5, + label=f"DA median (NSE={nse_da:.3f})") + + obs_window = obs.loc[plot_start:plot_end] + finite = np.isfinite(obs_window.values) + ax.scatter(obs_window.index[finite], obs_window.values[finite], + s=6, color=OBS_COLOR, linewidths=0, zorder=6, label="USGS obs") + + ax.set_xlabel("Date", fontsize=11) + ax.set_ylabel("Discharge (m3/s)", fontsize=11) + ax.set_title( + f"Reconstructed forecast time series at USGS 03463300 — F5 (re-kriged variance)\n" + f"Overlapping-leads pool from EnKF forecast ensemble " + f"({plot_start.date()} - {plot_end.date()})", + fontsize=12, + ) + ax.set_xlim(plot_start, plot_end) + ax.xaxis.set_major_locator(mdates.WeekdayLocator(interval=1)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.25, lw=0.4) + ax.legend(loc='upper left', fontsize=9, frameon=True, framealpha=0.92) + + plt.tight_layout() + out_path = os.path.join(out_dir, "lead_time_reconstructed_timeseries.png") + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + helene_zoom_start = pd.Timestamp("2024-09-24 00:00:00") + helene_zoom_end = pd.Timestamp("2024-09-29 23:00:00") + da_z = da_rec[(da_rec['valid_time'] >= helene_zoom_start) + & (da_rec['valid_time'] <= helene_zoom_end)] + ol_z = ol_rec[(ol_rec['valid_time'] >= helene_zoom_start) + & (ol_rec['valid_time'] <= helene_zoom_end)] + nse_da_z = nse(da_z['obs'].values, da_z['median'].values) + nse_ol_z = nse(ol_z['obs'].values, ol_z['median'].values) + print(f"NSE (Helene zoom) — DA: {nse_da_z:.3f} | OL: {nse_ol_z:.3f}") + + fig, ax = plt.subplots(figsize=(14, 5.5)) + fig.patch.set_facecolor("white") + ax.fill_between(ol_z['valid_time'], ol_z['p05'], ol_z['p95'], + color=OL_COLOR, alpha=0.18, zorder=2, label="Open-loop [5th-95th]") + ax.plot(ol_z['valid_time'], ol_z['median'], color=OL_COLOR, + lw=1.4, linestyle='--', zorder=3, + label=f"Open-loop median (NSE={nse_ol_z:.3f})") + ax.fill_between(da_z['valid_time'], da_z['p05'], da_z['p95'], + color=DA_COLOR, alpha=0.22, zorder=4, label="DA [5th-95th]") + ax.plot(da_z['valid_time'], da_z['median'], color=DA_COLOR, + lw=1.6, zorder=5, label=f"DA median (NSE={nse_da_z:.3f})") + obs_z = obs.loc[helene_zoom_start:helene_zoom_end] + finite = np.isfinite(obs_z.values) + ax.scatter(obs_z.index[finite], obs_z.values[finite], + s=12, color=OBS_COLOR, linewidths=0, zorder=6, label="USGS obs") + ax.set_xlabel("Date", fontsize=11) + ax.set_ylabel("Discharge (m3/s)", fontsize=11) + ax.set_title( + "Reconstructed forecast time series at USGS 03463300 — Helene zoom\n" + "F5 (re-kriged variance) | Sep 24-29, 2024 | overlapping-leads pool", + fontsize=12, + ) + ax.xaxis.set_major_locator(mdates.HourLocator(interval=12)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %H:%M")) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.25, lw=0.4) + ax.legend(loc='upper left', fontsize=9, frameon=True, framealpha=0.92) + plt.tight_layout() + out_zoom = os.path.join(out_dir, "lead_time_reconstructed_timeseries_helene.png") + plt.savefig(out_zoom, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_zoom}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4c_timeseries/run_4c_f5.sh b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4c_timeseries/run_4c_f5.sh new file mode 100644 index 00000000..5a29efa3 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/folder5_rekrig_variance_direct/4_evaluation/4c_timeseries/run_4c_f5.sh @@ -0,0 +1,52 @@ +#!/usr/bin/env bash +# F5 re-kriged variance (1 gauge holdout) — 4c reconstructed timeseries plots. +# +# Reads routed_leadtime_da_full.parquet + routed_leadtime_openloop_full.parquet +# from the F5 leadtime route dir and writes PNGs there. +# Also runs per-catchment issue-time hydrograph plots. +# +# Usage: +# bash run_4c_f5.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +F5_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder5_rekrig_variance_direct +ROUTE_DIR=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder5_rekrig_variance_direct_leadtime_routed +KRIG_OBS_DIR=/mnt/disk2/1400_sites_helene/catchment_ts_no_03463300_dynamic_variance_rekrig +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +echo "[F5-4c] Reconstructed timeseries from: $ROUTE_DIR" +$TROUTE "$SCRIPT_DIR/plot_lead_time_reconstructed_timeseries.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" + +echo "[F5-4c] Forecast spaghetti..." +$TROUTE "$SCRIPT_DIR/plot_forecast_spaghetti.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" + +echo "[F5-4c] Helene vs low-flow issue-time hydrograph (per catchment)..." +for CAT in "${CATS[@]}"; do + DA_CSV="$F5_DIR/$CAT/${CAT}_lead_time_forecasts_da.csv" + [ -f "$DA_CSV" ] || { echo " [$CAT] no leadtime CSV — skipping"; continue; } + echo " [$CAT] helene hydrograph..." + $TROUTE "$SCRIPT_DIR/plot_helene_issue_time_hydrograph.py" \ + --cat-id "$CAT" \ + --leadtime-dir "$F5_DIR" \ + --da-dir "$F5_DIR" \ + --krig-obs-dir "$KRIG_OBS_DIR" +done + +echo "[F5-4c] Done. Outputs in: $ROUTE_DIR" +ls "$ROUTE_DIR"/*.png 2>/dev/null || echo " (no PNGs in route dir)" diff --git a/da_methods/test_1_heldout_gauge/gapfill_krig_obs.py b/da_methods/test_1_heldout_gauge/gapfill_krig_obs.py new file mode 100644 index 00000000..078e0736 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/gapfill_krig_obs.py @@ -0,0 +1,92 @@ +""" +gapfill_krig_obs.py + +Cleans per-catchment Qkrig CSVs: + 1. Drops all rows where date is 2019-01-01 (leading NaN day) + 2. Linear-interpolates any remaining isolated interior NaN hours + 3. Renames columns to match calibration script format: datetime, qkrig, variance + 4. Overwrites files in-place (or writes to --out-dir if given) + +Output columns: datetime, qkrig (mm/h), variance + +Usage: + python gapfill_krig_obs.py --obs-dir /mnt/disk2/1400_sites_helene/catchment_ts_no_03463300_with_variance + python gapfill_krig_obs.py --obs-dir --out-dir # non-destructive +""" + +import argparse +from pathlib import Path + +import pandas as pd + + +def gapfill(path_in: Path, path_out: Path) -> dict: + df = pd.read_csv(path_in) + df["time"] = pd.to_datetime(df["time"]) + + n_total = len(df) + + # Step 1 — drop leading 2019-01-01 day + df = df[df["time"].dt.date.astype(str) != "2019-01-01"].copy() + n_after_drop = len(df) + dropped = n_total - n_after_drop + + # Step 2 — linear-interpolate interior NaNs (limit=48 to avoid extrapolation) + nan_before = df["qkrig_mm_hr"].isna().sum() + df["qkrig_mm_hr"] = df["qkrig_mm_hr"].interpolate(method="linear", limit=48, limit_direction="both") + df["qkrig_variance"] = df["qkrig_variance"].interpolate(method="linear", limit=48, limit_direction="both") + nan_after = df["qkrig_mm_hr"].isna().sum() + + # Step 3 — rename to match calibration script expected format + df = df.reset_index(drop=True) + df = df.rename(columns={ + "time": "datetime", + "qkrig_mm_hr": "qkrig", + "qkrig_variance": "variance", + })[["datetime", "qkrig", "variance"]] + + path_out.parent.mkdir(parents=True, exist_ok=True) + df.to_csv(path_out, index=False) + + return {"dropped": dropped, "interpolated": nan_before - nan_after, "remaining_nan": nan_after} + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--obs-dir", required=True) + parser.add_argument("--out-dir", default=None, + help="Output directory. If omitted, overwrites files in-place.") + args = parser.parse_args() + + obs_dir = Path(args.obs_dir) + out_dir = Path(args.out_dir) if args.out_dir else obs_dir + + csvs = sorted(obs_dir.glob("cat-*.csv")) + if not csvs: + print(f"No cat-*.csv files found in {obs_dir}") + return + + print(f"Processing {len(csvs)} catchment files...") + total_interp = 0 + total_remaining = 0 + + for csv in csvs: + out_path = out_dir / csv.name + stats = gapfill(csv, out_path) + total_interp += stats["interpolated"] + total_remaining += stats["remaining_nan"] + if stats["remaining_nan"] > 0: + print(f" WARNING {csv.name}: {stats['remaining_nan']} NaNs remain after interpolation") + + print(f"\nDone.") + print(f" Dropped 2019-01-01 rows : yes (24 rows per file)") + print(f" Total hours interpolated: {total_interp}") + print(f" Total NaNs remaining : {total_remaining}") + if total_remaining > 0: + print(" WARNING: some NaNs could not be interpolated — check files above") + else: + print(" All series are now gap-free.") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_1_heldout_gauge/plot_krig_obs_comparison.py b/da_methods/test_1_heldout_gauge/plot_krig_obs_comparison.py new file mode 100644 index 00000000..c4980254 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/plot_krig_obs_comparison.py @@ -0,0 +1,154 @@ +""" +plot_krig_obs_comparison.py + +Compare Qkrig pseudo-observations used in two experiments at cat-1016300: + - 20% holdout: catchment_ts_03463300_with_variance + -> gauge 03463300 is HELD OUT; kriging built from ~80% of gauges + - 1-gauge holdout: catchment_ts_no_03463300_gapfilled + -> gauge 03463300 is HELD OUT; kriging built from all other gauges (~7,299) + +In both cases gauge 03463300 is withheld from the kriging network. +The difference: 20% holdout removes ~20% of all gauges (sparser network), +while 1-gauge holdout removes only one gauge (denser network). + +Overlays USGS gauge 03463300 observations for reference. + +Output: krig_obs_comparison_full.png (full test period) + krig_obs_comparison_helene.png (Helene window zoom) +""" + +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +# ── Paths ────────────────────────────────────────────────────────────────── +OBS_20PCT = "/mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance/cat-1016300.csv" +OBS_1GAUGE = "/home/svyas/catchment_ts_no_03463300_gapfilled/cat-1016300.csv" +USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" +OUT_DIR = "/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +HELENE_START = pd.Timestamp("2024-09-24") +HELENE_END = pd.Timestamp("2024-09-30") +PLOT_START = pd.Timestamp("2019-01-01") +PLOT_END = pd.Timestamp("2024-10-31") + + +def load_krig(path, label): + df = pd.read_csv(path) + time_col = next(c for c in df.columns if c.lower() in ('datetime', 'date', 'time', 'timestamp')) + df[time_col] = pd.to_datetime(df[time_col]) + df = df.set_index(time_col).sort_index() + q_col = next(c for c in df.columns if 'qkrig' in c.lower() and 'var' not in c.lower()) + print(f" [{label}] loaded {len(df)} rows, q_col='{q_col}'") + return df[q_col].astype(float) + + +def load_usgs(path): + df = pd.read_csv(path) + date_col = next(c for c in df.columns if c.lower() in ('datetime','date','time','timestamp')) + q_col = next(c for c in df.columns if 'q' in c.lower() or 'flow' in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + s = df.set_index(date_col)[q_col].astype(float).sort_index() + if 'mm' in q_col.lower(): + s = s * MM_H_TO_M3_S + print(f" [USGS] loaded {len(s)} rows, q_col='{q_col}'") + return s + + +def make_plot(q20, q1g, usgs, t_start, t_end, suffix, title): + mask20 = (q20.index >= t_start) & (q20.index <= t_end) + mask1g = (q1g.index >= t_start) & (q1g.index <= t_end) + usgs_m = (usgs.index >= t_start) & (usgs.index <= t_end) + + # Compute difference on common index + common = q20.index.intersection(q1g.index) + common = common[(common >= t_start) & (common <= t_end)] + diff = q20.reindex(common) - q1g.reindex(common) + + fig, (ax, ax2) = plt.subplots(2, 1, figsize=(14, 8), + gridspec_kw={'height_ratios': [3, 1]}, + sharex=True) + + # ── Top panel: both Qkrig series + USGS ────────────────────────────── + # Draw 1-gauge first, then 20% on top with dashed so both visible + ax.plot(q1g.index[mask1g], q1g.values[mask1g], + color='tomato', lw=1.2, alpha=0.9, zorder=2, + label='Qkrig — 1-gauge holdout (03463300 removed; ~7,299 gauges)') + ax.plot(q20.index[mask20], q20.values[mask20], + color='steelblue', lw=1.4, alpha=0.9, linestyle='--', zorder=3, + label='Qkrig — 20% holdout (03463300 removed; ~80% of gauges)') + ax.plot(usgs.index[usgs_m], usgs.values[usgs_m] / MM_H_TO_M3_S, + color='black', lw=1.6, zorder=4, + label='USGS obs 03463300 (converted to mm/h)') + + if t_end > HELENE_START: + hs = max(t_start, HELENE_START) + he = min(t_end, HELENE_END) + ax.axvspan(hs, he, color='salmon', alpha=0.12, zorder=0, label='Helene window') + + ax.set_ylabel('q (mm/h)', fontsize=11) + ax.set_title(title, fontsize=12) + ax.legend(fontsize=9, loc='upper left') + ax.grid(True, alpha=0.2) + + # ── Bottom panel: difference (20pct − 1gauge) ───────────────────────── + ax2.plot(common, diff.values, color='purple', lw=0.8, alpha=0.85) + ax2.axhline(0, color='black', lw=0.7, linestyle='--') + if t_end > HELENE_START: + hs = max(t_start, HELENE_START) + he = min(t_end, HELENE_END) + ax2.axvspan(hs, he, color='salmon', alpha=0.12, zorder=0) + ax2.set_ylabel('Difference\n(20% − 1-gauge)\nmm/h', fontsize=9) + ax2.set_xlabel('Date', fontsize=11) + ax2.grid(True, alpha=0.2) + + if suffix == 'helene': + ax2.xaxis.set_major_locator(mdates.DayLocator(interval=1)) + ax2.xaxis.set_major_formatter(mdates.DateFormatter('%b %d')) + else: + ax2.xaxis.set_major_locator(mdates.MonthLocator(interval=2)) + ax2.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m')) + plt.setp(ax2.xaxis.get_majorticklabels(), rotation=30, ha='right') + + plt.tight_layout() + out = os.path.join(OUT_DIR, f'krig_obs_comparison_{suffix}.png') + plt.savefig(out, dpi=150, bbox_inches='tight') + plt.close() + print(f'Saved: {out}') + + +def main(): + os.makedirs(OUT_DIR, exist_ok=True) + print('Loading Qkrig obs...') + q20 = load_krig(OBS_20PCT, '20pct holdout') + q1g = load_krig(OBS_1GAUGE, '1-gauge holdout') + print('Loading USGS obs...') + usgs = load_usgs(USGS_CSV) + + print('Plotting full test period...') + make_plot(q20, q1g, usgs, PLOT_START, PLOT_END, 'full', + 'Qkrig pseudo-observations at cat-1016300\n' + '20% holdout vs 1-gauge holdout — full test period (Oct 2023–Oct 2024)') + + print('Plotting Helene window...') + make_plot(q20, q1g, usgs, HELENE_START, HELENE_END, 'helene', + 'Qkrig pseudo-observations at cat-1016300\n' + '20% holdout vs 1-gauge holdout — Hurricane Helene (Sep 24–30, 2024)') + + # Print summary stats for the Helene window + h20 = q20[(q20.index >= HELENE_START) & (q20.index <= HELENE_END)] + h1g = q1g[(q1g.index >= HELENE_START) & (q1g.index <= HELENE_END)] + husgs= usgs[(usgs.index >= HELENE_START) & (usgs.index <= HELENE_END)] + print('\n── Helene window peak (mm/h) ────────────────') + print(f' Qkrig 20% holdout : {h20.max():.3f} mm/h') + print(f' Qkrig 1-gauge : {h1g.max():.3f} mm/h') + print(f' USGS (converted) : {(husgs.max()/MM_H_TO_M3_S):.3f} mm/h') + + +if __name__ == '__main__': + main() diff --git a/da_methods/test_1_heldout_gauge/run_hydro_arm_plots.sh b/da_methods/test_1_heldout_gauge/run_hydro_arm_plots.sh new file mode 100644 index 00000000..38762fb4 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/run_hydro_arm_plots.sh @@ -0,0 +1,58 @@ +#!/usr/bin/env bash +# Generate hydro-state arm plots (2b_hydro_arm_helene.png) for F1, F2, F4, F5 +# of the 1-gauge holdout experiment — all 21 catchments per experiment. +# +# Reads _da_hydro_arm.csv and _da_forcing_arm.csv from each experiment's dir. +# Saves PNGs alongside the arm CSVs in each catchment subfolder. +# +# Usage: +# bash run_hydro_arm_plots.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/plot_da_perturbation_arms.py +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +declare -A EXPS +EXPS["F1"]=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt +EXPS["F2"]=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder2_fixed_r007 +EXPS["F4"]=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder4_dynamic_variance_direct +EXPS["F5"]=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder5_rekrig_variance_direct + +for EXP in F1 F2 F4 F5; do + ARM_DIR="${EXPS[$EXP]}" + echo "" + echo "=======================================" + echo "=== $EXP arm-dir: $ARM_DIR ===" + echo "=======================================" + + for CAT in "${CATS[@]}"; do + HYDRO_CSV="$ARM_DIR/$CAT/${CAT}_da_hydro_arm.csv" + if [ ! -f "$HYDRO_CSV" ]; then + echo " [$EXP $CAT] No hydro arm CSV — skipping" + continue + fi + PNG="$ARM_DIR/$CAT/${CAT}_2b_hydro_arm_helene.png" + if [ -f "$PNG" ]; then + echo " [$EXP $CAT] Already exists — skipping" + continue + fi + echo " [$EXP $CAT] plotting..." + $TROUTE "$SCRIPT" \ + --arm-dir "$ARM_DIR" \ + --cat-id "$CAT" \ + --usgs-csv "$USGS_CSV" + done +done + +echo "" +echo "=== Hydro arm plots done ===" diff --git a/da_methods/test_1_heldout_gauge/run_hydro_arm_plots_openloop.sh b/da_methods/test_1_heldout_gauge/run_hydro_arm_plots_openloop.sh new file mode 100644 index 00000000..c61e2e30 --- /dev/null +++ b/da_methods/test_1_heldout_gauge/run_hydro_arm_plots_openloop.sh @@ -0,0 +1,69 @@ +#!/usr/bin/env bash +# Regenerate 2b hydro-state arm plots for F1, F2, F4, F5 with the routed +# open-loop baseline overlaid on each panel. +# +# Uses openloop/routed_Q_test.csv (Q_routed_m3s column, already in m³/s, +# correct timing via T-route) — same line on every catchment plot. +# +# Overwrites existing _2b_hydro_arm_helene.png files in each catchment subfolder. +# +# Usage: +# bash run_hydro_arm_plots_openloop.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/test_1_heldout_gauge/folder1_variance_scaled_vrugt/4_evaluation/plot_da_perturbation_arms.py +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv +OL_CSV=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/openloop/routed_Q_test.csv + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +declare -A EXPS +EXPS["F1"]=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder1_vrugt +EXPS["F2"]=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder2_fixed_r007 +EXPS["F4"]=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder4_dynamic_variance_direct +EXPS["F5"]=/mnt/disk2/suma_helen_poster/da_results_1gauge_heldout/folder5_rekrig_variance_direct + +if [ ! -f "$OL_CSV" ]; then + echo "ERROR: routed open-loop not found: $OL_CSV" + exit 1 +fi +echo "Open-loop: $OL_CSV" + +SKIP=0; DONE=0; FAIL=0 + +for EXP in F1 F2 F4 F5; do + ARM_DIR="${EXPS[$EXP]}" + echo "" + echo "=======================================" + echo "=== $EXP arm-dir: $ARM_DIR ===" + echo "=======================================" + + for CAT in "${CATS[@]}"; do + HYDRO_CSV="$ARM_DIR/$CAT/${CAT}_da_hydro_arm.csv" + + if [ ! -f "$HYDRO_CSV" ]; then + echo " [$EXP $CAT] No hydro arm CSV — skipping" + SKIP=$((SKIP + 1)); continue + fi + + echo " [$EXP $CAT] plotting..." + $TROUTE "$SCRIPT" \ + --arm-dir "$ARM_DIR" \ + --cat-id "$CAT" \ + --usgs-csv "$USGS_CSV" \ + --ol-ts-csv "$OL_CSV" \ + && DONE=$((DONE + 1)) \ + || { echo " FAILED: $EXP $CAT"; FAIL=$((FAIL + 1)); } + done +done + +echo "" +echo "=== Hydro arm plots (with routed open loop) done: done=$DONE skipped=$SKIP failed=$FAIL ===" diff --git a/da_methods/test_20pct_heldout_gauges/EXPERIMENTS.md b/da_methods/test_20pct_heldout_gauges/EXPERIMENTS.md new file mode 100644 index 00000000..e5c4e228 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/EXPERIMENTS.md @@ -0,0 +1,58 @@ +# DA Experiment Index + +Four experiments comparing observation error variance (R) formulations for CFE + EnKF +data assimilation at USGS gauge 03463300 (South Toe River Near Celo, NC). + +Each folder is self-contained with its own `1_calibrate/`, `2_assimilation/`, `3_routing/`, +`4_evaluation/` pipeline. Calibrated parameters (best_params.json) are shared across all +folders from the external `calibrate-cfe` repo — R formula is irrelevant at calibration time. + +--- + +## Results Summary + +| Folder | R Formula | Full KGE | Full NSE | Helene KGE | Helene NSE | Helene peak % USGS | +|---|---|---|---|---|---|---| +| [1 — Variance scaled Vrugt](folder1_variance_scaled_vrugt/) | `(0.10·y)² + 0.001·σ²_krig` | +0.277 | **+0.555** | +0.213 | **+0.459** | ~34% | +| [2 — Fixed R=0.07](folder2_fixed_r_007/) | `0.07` (constant) | **+0.503** | +0.163 | **+0.439** | -0.009 | **~65%** | +| [3 — Dynamic Vrugt seeded](folder3_dynamic_vrugt_seeded/) | `(0.10·y)² + 0.001·σ²_krig` (seed=42, spliced obs) | +0.261 | +0.485 | +0.189 | +0.372 | ~37% | +| [4 — Dynamic variance direct](folder4_dynamic_variance_direct/) | `σ²_krig` directly (seed=42) | +0.200 | +0.428 | +0.132 | +0.303 | ~35% | + +USGS Helene peak: 1885.7 m³/s + +**Finding:** Fixed R=0.07 (Folder 2) gives the best KGE and captures 65% of the Helene +peak — constant small R keeps Kalman gain high throughout the flood. The Vrugt formula +inflates R at high flows, dampening DA updates exactly when they matter most, so Folder 1 +captures only 34% of the peak despite having the best NSE. Raw σ²_krig (Folder 4) is +worst across all metrics. + +--- + +## Completeness Matrix + +| Step | Folder 1 | Folder 2 | Folder 3 | Folder 4 | +|---|---|---|---|---| +| 1. Calibrate | ✅ shared | ✅ shared | ✅ shared | ✅ shared | +| 2a. Forcing arm (30 members) | ✅ | ✅ (same as F1) | ✅ | ✅ | +| 2b. Hydro-state arm (20 members) | ✅ | ✅ (same as F1) | ✅ | ✅ | +| 2c. 18hr forecast cycles | ✅ | ✅ | ✅ | ✅ | +| 2d. 600-member crossed ensemble | ✅ | ✅ | ✅ | ✅ | +| 3. Route analysis trajectory | ✅ | ✅ | ✅ | ✅ | +| 3. Route 18hr forecast cycles | ✅ | ✅ | ✅ | ✅ | +| 3. Route crossed ensemble | ✅ | ✅ | ✅ | ✅ | +| 4a. Error decay | ✅ | ✅ | ✅ | ✅ | +| 4b. Ensemble vs obs | ✅ | ✅ | ✅ | ✅ | +| 4c. Reconstructed timeseries | ✅ | ✅ | ✅ | ✅ | + +--- + +## Server Data Paths + +| Folder | Analysis results | Forecast cycles | Ensemble | Routed analysis | +|---|---|---|---|---| +| 1 | `da_results/v2_true_enkf_vrugt/` | `v2_lead_time_forecast/` | `v2_crossed_ensemble_vrugt/` | `vrugt_dynamic_routed/` | +| 2 | `da_results/v2_fixed_r007_analysis/` | `v2_lead_time_forecast_hardcoded_r/` | `v2_crossed_ensemble/` | `fixed_r007_routed/` | +| 3 | `1400_sites_helene/da_results_dynamic_vrugt_seeded/` | `da_results/da_forecast_dynamic_vrugt_seeded/` | `da_results/da_crossed_dynamic_vrugt_seeded/` | `dynamic_vrugt_seeded_routed/` | +| 4 | `1400_sites_helene/da_results_dynamic_novrugt_seeded/` | `da_results/da_forecast_dynamic_novrugt_seeded/` | `da_results/da_crossed_dynamic_novrugt_seeded/` | `dynamic_novrugt_seeded_routed/` | + +All paths are under `/mnt/disk2/suma_helen_poster/` unless prefixed with `1400_sites_helene/`. diff --git a/da_methods/test_20pct_heldout_gauges/batch_run_openloop_20pct.sh b/da_methods/test_20pct_heldout_gauges/batch_run_openloop_20pct.sh new file mode 100644 index 00000000..92e95364 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/batch_run_openloop_20pct.sh @@ -0,0 +1,77 @@ +#!/usr/bin/env bash +# Open-loop CFE run (no DA) — 20% gauge holdout, all 21 catchments. +# Runs calibrate_catchment_cfe_da_v2.py WITHOUT --enkf-enabled. +# Params staged from da_results_dynamic_novrugt_seeded (20% holdout calibration). +# Obs: catchment_ts_03463300_dynamic_variance (03463300 in Qkrig, 22 others withheld). +# +# Usage: +# bash batch_run_openloop_20pct.sh > ~/logs/openloop_20pct.log 2>&1 + +set -euo pipefail + +PYTHON=/home/svyas/miniconda3/envs/troute/bin/python3 +PROD_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/1_distributed_cfe/calibrate_catchment_cfe_da_v2.py + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +PARAM_SRC=/mnt/disk2/1400_sites_helene/da_results_dynamic_novrugt_seeded +OBS_DIR=/mnt/disk2/1400_sites_helene/catchment_ts_03463300_dynamic_variance +CFE_DIR=/mnt/disk2/suma_helen_poster/cfe_py +CONFIG_FILE=/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json +PARAM_BOUNDS=/mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json +FORCING_DIR=/mnt/disk2/suma_helen_poster/nwm_retro_catchment_forcings +FORCING_DIR1=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings +FORCING_DIR2=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results/openloop_20pct_heldout + +mkdir -p "$OUT_DIR" +echo "Open-loop 20pct (no DA) — out: $OUT_DIR" +echo "" + +SKIP=0; DONE=0; FAIL=0 + +for CAT in "${CATS[@]}"; do + echo "===============================" + echo "=== $CAT ===" + echo "===============================" + + CAT_OUT="$OUT_DIR/$CAT" + OUT_FILE="$CAT_OUT/${CAT}_test_results.csv" + + PARAMS="$PARAM_SRC/$CAT/${CAT}_best_params.json" + if [ ! -f "$PARAMS" ]; then + echo " WARNING: No best_params for $CAT at $PARAMS — skipping" + SKIP=$((SKIP + 1)); continue + fi + + if [ -f "$OUT_FILE" ]; then + echo " Already exists — skipping" + SKIP=$((SKIP + 1)); continue + fi + + mkdir -p "$CAT_OUT" + cp "$PARAMS" "$CAT_OUT/" + + "$PYTHON" "$PROD_SCRIPT" \ + --cat-id "$CAT" \ + --forcing-dir "$FORCING_DIR" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --param-bounds "$PARAM_BOUNDS" \ + --out-dir "$OUT_DIR" \ + --test-forcing-dir1 "$FORCING_DIR1" \ + --test-forcing-dir2 "$FORCING_DIR2" \ + --enkf-members 1 \ + && DONE=$((DONE + 1)) \ + || { echo " FAILED: $CAT"; FAIL=$((FAIL + 1)); } +done + +echo "" +echo "=== Open-loop 20pct done: done=$DONE skipped=$SKIP failed=$FAIL ===" diff --git a/da_methods/test_20pct_heldout_gauges/compare_all_folders.py b/da_methods/test_20pct_heldout_gauges/compare_all_folders.py new file mode 100644 index 00000000..49152738 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/compare_all_folders.py @@ -0,0 +1,212 @@ +""" +compare_all_folders.py + +Four-way comparison of DA analysis trajectories routed to gauge 03463300, +one curve per R-formula experiment: + Folder 1 — Variance scaled using Vrugt R(t)=(0.10*y)^2 + 0.001*sigma2_krig + Folder 2 — Fixed R=0.07 R=0.07 (constant) + Folder 3 — Dynamic Vrugt seeded same Vrugt formula, spliced obs, seed=42 + Folder 4 — Dynamic variance direct R(t)=sigma2_krig directly + +Reads routed_Q_test.csv (columns: date, Q_routed_m3s, Q_usgs_m3s) from each folder. +Folder 2 deterministic route not available — shown as absent with a note. + +Outputs: + /compare_all_folders_full.png + /compare_all_folders_helene.png + /compare_all_folders_kge_table.csv + +Usage: + python3 compare_all_folders.py --out-dir /mnt/disk2/suma_helen_poster/da_results/comparison_plots +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +HELENE_START = pd.Timestamp("2024-09-24") +HELENE_END = pd.Timestamp("2024-09-29 23:00:00") + +BASE = "/mnt/disk2/suma_helen_poster/da_results" + +FOLDERS = [ + { + "label": "F1: Variance scaled Vrugt\nR(t)=(0.10·y)²+0.001·σ²_krig", + "short": "F1 Vrugt", + "color": "#1f77b4", + "lw": 2.0, + "csv": f"{BASE}/vrugt_dynamic_routed/routed_Q_test.csv", + }, + { + "label": "F2: Fixed R=0.07", + "short": "F2 R=0.07", + "color": "#ff7f0e", + "lw": 1.6, + "csv": f"{BASE}/fixed_r007_routed/routed_Q_test.csv", + }, + { + "label": "F3: Dynamic Vrugt seeded\n(spliced obs, seed=42)", + "short": "F3 Dyn Vrugt", + "color": "#2ca02c", + "lw": 1.6, + "csv": f"{BASE}/dynamic_vrugt_seeded_routed/routed_Q_test.csv", + }, + { + "label": "F4: Dynamic variance direct\nR(t)=σ²_krig", + "short": "F4 Direct σ²", + "color": "#d62728", + "lw": 1.4, + "csv": f"{BASE}/dynamic_novrugt_seeded_routed/routed_Q_test.csv", + }, + { + "label": "F5: Re-kriged variance\nR(t)=σ²_krig (re-kriged network)", + "short": "F5 Rekrig σ²", + "color": "#9467bd", + "lw": 2.2, + "csv": f"{BASE}/folder5_rekrig_variance_direct/routed_Q_test.csv", + }, +] + + +def kge(obs, sim): + mask = np.isfinite(obs) & np.isfinite(sim) + o, s = obs[mask], sim[mask] + if len(o) < 2 or np.std(o) == 0: + return np.nan + r = np.corrcoef(o, s)[0, 1] + return 1.0 - np.sqrt((r-1)**2 + (np.std(s)/np.std(o)-1)**2 + + (np.mean(s)/np.mean(o)-1)**2) + + +def nse(obs, sim): + mask = np.isfinite(obs) & np.isfinite(sim) + o, s = obs[mask], sim[mask] + if len(o) < 2: + return np.nan + denom = np.sum((o - o.mean())**2) + return 1.0 - np.sum((o - s)**2) / denom if denom > 0 else np.nan + + +def load_folder(cfg): + p = cfg["csv"] + if not os.path.exists(p): + print(f" MISSING: {p}") + return None + df = pd.read_csv(p, parse_dates=["date"]).set_index("date").sort_index() + return df + + +def plot_comparison(dfs, obs, dates, helene_mask, out_path, helene_only=False): + fig, ax = plt.subplots(figsize=(15, 6)) + + if helene_only: + plot_dates = dates[helene_mask] + obs_plot = obs[helene_mask] + else: + plot_dates = dates + obs_plot = obs + + ax.plot(plot_dates, obs_plot, + color="black", lw=2.2, zorder=6, label="USGS obs") + + for cfg, df in zip(FOLDERS, dfs): + if df is None: + continue + sim = df["Q_routed_m3s"].reindex(dates).values + if helene_only: + sim_plot = sim[helene_mask] + else: + sim_plot = sim + + kg = kge(obs[helene_mask] if helene_only else obs, + sim[helene_mask] if helene_only else sim) + ns = nse(obs[helene_mask] if helene_only else obs, + sim[helene_mask] if helene_only else sim) + + ax.plot(plot_dates, sim_plot, + color=cfg["color"], lw=cfg["lw"], zorder=4, + label=f"{cfg['short']} KGE={kg:+.3f} NSE={ns:+.3f}") + + if not helene_only: + ax.axvspan(HELENE_START, HELENE_END, color="gold", alpha=0.12, zorder=1) + + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_xlabel("Date (UTC)", fontsize=11) + title = ("Helene window — " if helene_only else "Full test period — ") + ax.set_title(title + "Four-folder R-formula comparison at gauge 03463300", fontsize=12) + ax.legend(fontsize=9, loc="upper left", framealpha=0.9) + ax.grid(True, alpha=0.22) + + if helene_only: + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=1)) + peak_usgs = np.nanmax(obs[helene_mask]) + ax.axhline(peak_usgs, color="black", lw=0.7, linestyle=":", alpha=0.5) + ax.text(HELENE_END - pd.Timedelta(hours=6), peak_usgs * 1.02, + f"USGS peak {peak_usgs:.0f} m³/s", + fontsize=8.5, ha="right", color="black") + else: + ax.xaxis.set_major_formatter(mdates.DateFormatter("%Y-%m")) + ax.xaxis.set_major_locator(mdates.MonthLocator(interval=2)) + + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right", fontsize=9) + plt.tight_layout() + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--out-dir", default=f"{BASE}/comparison_plots") + args = parser.parse_args() + os.makedirs(args.out_dir, exist_ok=True) + + print("Loading routed CSVs...") + dfs = [load_folder(cfg) for cfg in FOLDERS] + + # Use first available df for date index and USGS obs + ref = next(df for df in dfs if df is not None) + dates = ref.index + obs = ref["Q_usgs_m3s"].values if "Q_usgs_m3s" in ref.columns else np.full(len(dates), np.nan) + helene = ((dates >= HELENE_START) & (dates <= HELENE_END)) + + # ── Plots ────────────────────────────────────────────────────────────── + plot_comparison(dfs, obs, dates, helene, + os.path.join(args.out_dir, "compare_all_folders_full.png"), + helene_only=False) + plot_comparison(dfs, obs, dates, helene, + os.path.join(args.out_dir, "compare_all_folders_helene.png"), + helene_only=True) + + # ── KGE table ────────────────────────────────────────────────────────── + rows = [] + for cfg, df in zip(FOLDERS, dfs): + if df is None: + rows.append({"folder": cfg["short"], "full_kge": np.nan, + "full_nse": np.nan, "helene_kge": np.nan, + "helene_nse": np.nan, "helene_peak_m3s": np.nan}) + continue + sim = df["Q_routed_m3s"].reindex(dates).values + rows.append({ + "folder": cfg["short"], + "full_kge": round(kge(obs, sim), 3), + "full_nse": round(nse(obs, sim), 3), + "helene_kge": round(kge(obs[helene], sim[helene]), 3), + "helene_nse": round(nse(obs[helene], sim[helene]), 3), + "helene_peak_m3s": round(np.nanmax(sim[helene]), 1), + }) + + tbl = pd.DataFrame(rows) + tbl_path = os.path.join(args.out_dir, "compare_all_folders_kge_table.csv") + tbl.to_csv(tbl_path, index=False) + print(f"Saved: {tbl_path}") + print("\n" + tbl.to_string(index=False)) + print(f"\nUSGS Helene peak: {np.nanmax(obs[helene]):.1f} m³/s") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/figures/cross_folder/compare_all_folders_full.png b/da_methods/test_20pct_heldout_gauges/figures/cross_folder/compare_all_folders_full.png new file mode 100644 index 00000000..57321fad Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/cross_folder/compare_all_folders_full.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/cross_folder/compare_all_folders_helene.png b/da_methods/test_20pct_heldout_gauges/figures/cross_folder/compare_all_folders_helene.png new file mode 100644 index 00000000..38176461 Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/cross_folder/compare_all_folders_helene.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/cat-1016300_2a_forcing_arm_helene.png b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/cat-1016300_2a_forcing_arm_helene.png new file mode 100644 index 00000000..24684d0d Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/cat-1016300_2a_forcing_arm_helene.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/cat-1016300_2ab_arms_comparison.png b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/cat-1016300_2ab_arms_comparison.png new file mode 100644 index 00000000..b070fd6f Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/cat-1016300_2ab_arms_comparison.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/cat-1016300_2b_hydro_arm_helene.png b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/cat-1016300_2b_hydro_arm_helene.png new file mode 100644 index 00000000..5ad3182b Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/cat-1016300_2b_hydro_arm_helene.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/cat-1016300_per_member_factor_decomp.png b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/cat-1016300_per_member_factor_decomp.png new file mode 100644 index 00000000..c4ea84fd Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/cat-1016300_per_member_factor_decomp.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/cat-1016300_per_member_io_diagnostic.png b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/cat-1016300_per_member_io_diagnostic.png new file mode 100644 index 00000000..a0dcd0c1 Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/cat-1016300_per_member_io_diagnostic.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/cat-1016300_perturbation_categories_linear.png b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/cat-1016300_perturbation_categories_linear.png new file mode 100644 index 00000000..1fd42df6 Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/cat-1016300_perturbation_categories_linear.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/cat-1016300_perturbation_categories_log.png b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/cat-1016300_perturbation_categories_log.png new file mode 100644 index 00000000..4b4028b0 Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/cat-1016300_perturbation_categories_log.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/cat-1016300_production_ensemble_forecast_linear.png b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/cat-1016300_production_ensemble_forecast_linear.png new file mode 100644 index 00000000..ea3ff045 Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/cat-1016300_production_ensemble_forecast_linear.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/cat-1016300_production_ensemble_forecast_log.png b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/cat-1016300_production_ensemble_forecast_log.png new file mode 100644 index 00000000..6716545e Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/cat-1016300_production_ensemble_forecast_log.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/da_v2_vrugt_kge_table.png b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/da_v2_vrugt_kge_table.png new file mode 100644 index 00000000..941ee8e0 Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/da_v2_vrugt_kge_table.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/da_vs_qkrig_vs_usgs.png b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/da_vs_qkrig_vs_usgs.png new file mode 100644 index 00000000..6053b9af Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/da_vs_qkrig_vs_usgs.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/forecast_spaghetti_helene.png b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/forecast_spaghetti_helene.png new file mode 100644 index 00000000..a77e077f Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/forecast_spaghetti_helene.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/helene_da_v2_vrugt_vs_run3_grid.png b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/helene_da_v2_vrugt_vs_run3_grid.png new file mode 100644 index 00000000..692679ae Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/helene_da_v2_vrugt_vs_run3_grid.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/helene_da_v2_vrugt_vs_run3_grid_zoomed.png b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/helene_da_v2_vrugt_vs_run3_grid_zoomed.png new file mode 100644 index 00000000..a3fb0b47 Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/helene_da_v2_vrugt_vs_run3_grid_zoomed.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/helene_ensemble_twopanel.png b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/helene_ensemble_twopanel.png new file mode 100644 index 00000000..78d7dc5f Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/helene_ensemble_twopanel.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/helene_ensemble_vs_usgs.png b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/helene_ensemble_vs_usgs.png new file mode 100644 index 00000000..c515535f Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/helene_ensemble_vs_usgs.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/helene_sensitivity_spaghetti_appendix.png b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/helene_sensitivity_spaghetti_appendix.png new file mode 100644 index 00000000..690c41ce Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/helene_sensitivity_spaghetti_appendix.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/helene_sensitivity_spaghetti_main.png b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/helene_sensitivity_spaghetti_main.png new file mode 100644 index 00000000..bc239354 Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/helene_sensitivity_spaghetti_main.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/lead_time_reconstructed_timeseries.png b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/lead_time_reconstructed_timeseries.png new file mode 100644 index 00000000..e1a39773 Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/lead_time_reconstructed_timeseries.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/lead_time_reconstructed_timeseries_helene.png b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/lead_time_reconstructed_timeseries_helene.png new file mode 100644 index 00000000..46c1ac41 Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f1_variance_scaled_vrugt/lead_time_reconstructed_timeseries_helene.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f2_fixed_r_007/error_fixed_target_mean.png b/da_methods/test_20pct_heldout_gauges/figures/f2_fixed_r_007/error_fixed_target_mean.png new file mode 100644 index 00000000..654c571d Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f2_fixed_r_007/error_fixed_target_mean.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f3_dynamic_vrugt_seeded/f3_helene_ensemble_twopanel.png b/da_methods/test_20pct_heldout_gauges/figures/f3_dynamic_vrugt_seeded/f3_helene_ensemble_twopanel.png new file mode 100644 index 00000000..c2cfd14e Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f3_dynamic_vrugt_seeded/f3_helene_ensemble_twopanel.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f3_dynamic_vrugt_seeded/f3_helene_ensemble_vs_usgs.png b/da_methods/test_20pct_heldout_gauges/figures/f3_dynamic_vrugt_seeded/f3_helene_ensemble_vs_usgs.png new file mode 100644 index 00000000..a1c8d3ce Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f3_dynamic_vrugt_seeded/f3_helene_ensemble_vs_usgs.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f3_dynamic_vrugt_seeded/f3_lead_time_decay_gauge_by_regime.png b/da_methods/test_20pct_heldout_gauges/figures/f3_dynamic_vrugt_seeded/f3_lead_time_decay_gauge_by_regime.png new file mode 100644 index 00000000..23a087ff Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f3_dynamic_vrugt_seeded/f3_lead_time_decay_gauge_by_regime.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f3_dynamic_vrugt_seeded/f3_lead_time_decay_gauge_pooled.png b/da_methods/test_20pct_heldout_gauges/figures/f3_dynamic_vrugt_seeded/f3_lead_time_decay_gauge_pooled.png new file mode 100644 index 00000000..6bc2c6fd Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f3_dynamic_vrugt_seeded/f3_lead_time_decay_gauge_pooled.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f3_dynamic_vrugt_seeded/f3_reconstructed_timeseries.png b/da_methods/test_20pct_heldout_gauges/figures/f3_dynamic_vrugt_seeded/f3_reconstructed_timeseries.png new file mode 100644 index 00000000..33f2c4db Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f3_dynamic_vrugt_seeded/f3_reconstructed_timeseries.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f3_dynamic_vrugt_seeded/f3_reconstructed_timeseries_helene.png b/da_methods/test_20pct_heldout_gauges/figures/f3_dynamic_vrugt_seeded/f3_reconstructed_timeseries_helene.png new file mode 100644 index 00000000..86098be5 Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f3_dynamic_vrugt_seeded/f3_reconstructed_timeseries_helene.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f4_dynamic_variance_direct/f4_helene_ensemble_twopanel.png b/da_methods/test_20pct_heldout_gauges/figures/f4_dynamic_variance_direct/f4_helene_ensemble_twopanel.png new file mode 100644 index 00000000..f204b5d4 Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f4_dynamic_variance_direct/f4_helene_ensemble_twopanel.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f4_dynamic_variance_direct/f4_helene_ensemble_vs_usgs.png b/da_methods/test_20pct_heldout_gauges/figures/f4_dynamic_variance_direct/f4_helene_ensemble_vs_usgs.png new file mode 100644 index 00000000..8b247a29 Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f4_dynamic_variance_direct/f4_helene_ensemble_vs_usgs.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f4_dynamic_variance_direct/f4_lead_time_decay_gauge_by_regime.png b/da_methods/test_20pct_heldout_gauges/figures/f4_dynamic_variance_direct/f4_lead_time_decay_gauge_by_regime.png new file mode 100644 index 00000000..fe86cca5 Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f4_dynamic_variance_direct/f4_lead_time_decay_gauge_by_regime.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f4_dynamic_variance_direct/f4_lead_time_decay_gauge_pooled.png b/da_methods/test_20pct_heldout_gauges/figures/f4_dynamic_variance_direct/f4_lead_time_decay_gauge_pooled.png new file mode 100644 index 00000000..bf14d6c2 Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f4_dynamic_variance_direct/f4_lead_time_decay_gauge_pooled.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f4_dynamic_variance_direct/f4_reconstructed_timeseries.png b/da_methods/test_20pct_heldout_gauges/figures/f4_dynamic_variance_direct/f4_reconstructed_timeseries.png new file mode 100644 index 00000000..04625dfc Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f4_dynamic_variance_direct/f4_reconstructed_timeseries.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f4_dynamic_variance_direct/f4_reconstructed_timeseries_helene.png b/da_methods/test_20pct_heldout_gauges/figures/f4_dynamic_variance_direct/f4_reconstructed_timeseries_helene.png new file mode 100644 index 00000000..54a6dd30 Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f4_dynamic_variance_direct/f4_reconstructed_timeseries_helene.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/cat-1016300_2a_forcing_arm_helene.png b/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/cat-1016300_2a_forcing_arm_helene.png new file mode 100644 index 00000000..39641595 Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/cat-1016300_2a_forcing_arm_helene.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/cat-1016300_2ab_arms_comparison.png b/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/cat-1016300_2ab_arms_comparison.png new file mode 100644 index 00000000..981e0d0e Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/cat-1016300_2ab_arms_comparison.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/cat-1016300_2b_hydro_arm_helene.png b/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/cat-1016300_2b_hydro_arm_helene.png new file mode 100644 index 00000000..2880cf5a Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/cat-1016300_2b_hydro_arm_helene.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/cat-1016300_lead_time_decay.png b/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/cat-1016300_lead_time_decay.png new file mode 100644 index 00000000..c3a2c661 Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/cat-1016300_lead_time_decay.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/helene_ensemble_fan_f5.png b/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/helene_ensemble_fan_f5.png new file mode 100644 index 00000000..775e80e4 Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/helene_ensemble_fan_f5.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/lead_time_nse_gauge_f5_by_regime.png b/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/lead_time_nse_gauge_f5_by_regime.png new file mode 100644 index 00000000..0491872f Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/lead_time_nse_gauge_f5_by_regime.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/lead_time_nse_gauge_f5_pooled.png b/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/lead_time_nse_gauge_f5_pooled.png new file mode 100644 index 00000000..77a6b6df Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/lead_time_nse_gauge_f5_pooled.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/ol_comparison_full_log.png b/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/ol_comparison_full_log.png new file mode 100644 index 00000000..aaa9c4b3 Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/ol_comparison_full_log.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/ol_comparison_helene.png b/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/ol_comparison_helene.png new file mode 100644 index 00000000..871a3973 Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/ol_comparison_helene.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/view1_full_log_f2_vs_ol.png b/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/view1_full_log_f2_vs_ol.png new file mode 100644 index 00000000..09327218 Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/view1_full_log_f2_vs_ol.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/view2_helene_all_curves.png b/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/view2_helene_all_curves.png new file mode 100644 index 00000000..5b4adda4 Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/view2_helene_all_curves.png differ diff --git a/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/view3_full_linear_all.png b/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/view3_full_linear_all.png new file mode 100644 index 00000000..55796d80 Binary files /dev/null and b/da_methods/test_20pct_heldout_gauges/figures/f5_rekrig_variance_direct/view3_full_linear_all.png differ diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/2_assimilation/batch_run_crossed_all_cats_vrugt.sh b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/2_assimilation/batch_run_crossed_all_cats_vrugt.sh new file mode 100644 index 00000000..0eb09159 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/2_assimilation/batch_run_crossed_all_cats_vrugt.sh @@ -0,0 +1,79 @@ +#!/bin/bash +# batch_run_crossed_all_cats_vrugt.sh +# Runs run_crossed_ensemble.py for all 21 catchments using dynamic Vrugt R: +# R(t) = (0.10 * y_obs(t))^2 + 0.001 * krig_var(t) +# +# Outputs go to v2_crossed_ensemble_vrugt/ (separate from the R=0.07 ensemble). +# Run from the server: +# nohup bash ~/batch_run_crossed_all_cats_vrugt.sh \ +# > ~/crossed_vrugt.log 2>&1 & + +set -e + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +OBS_DIR=/mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance +CFE_DIR=/mnt/disk2/suma_helen_poster/cfe_py +CONFIG_FILE=/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble_vrugt +BEST_PARAMS_SRC=/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt +FORCING_DIR1=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings +FORCING_DIR2=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings + +PYTHON=/home/svyas/miniconda3/envs/troute/bin/python + +SKIP=0 +DONE=0 +FAIL=0 + +for CAT in "${CATS[@]}"; do + echo "" + echo "===============================" + echo "=== $CAT ===" + echo "===============================" + + CAT_OUT=$OUT_DIR/$CAT + mkdir -p "$CAT_OUT" + + # Stage best_params from v2_true_enkf_vrugt (Folder 1 analysis results) + if [ ! -f "$CAT_OUT/${CAT}_best_params.json" ]; then + SRC="$BEST_PARAMS_SRC/$CAT/${CAT}_best_params.json" + if [ -f "$SRC" ]; then + cp "$SRC" "$CAT_OUT/" + echo " Staged best_params from v2_true_enkf_vrugt" + else + echo " WARNING: No best_params for $CAT — skipping" + SKIP=$((SKIP + 1)) + continue + fi + fi + + # Skip if parquet already exists + if [ -f "$CAT_OUT/${CAT}_crossed_ensemble.parquet" ]; then + echo " Crossed ensemble already complete — skipping" + SKIP=$((SKIP + 1)) + continue + fi + + # No --hardcoded-r flag → default=None → uses Vrugt dynamic R + $PYTHON ~/run_crossed_ensemble.py \ + --cat-id "$CAT" \ + --forcing-dir "$FORCING_DIR1" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --out-dir "$OUT_DIR" \ + --test-forcing-dir1 "$FORCING_DIR1" \ + --test-forcing-dir2 "$FORCING_DIR2" \ + && DONE=$((DONE + 1)) \ + || { echo " FAILED: $CAT"; FAIL=$((FAIL + 1)); } +done + +echo "" +echo "=== Batch complete: done=$DONE skipped=$SKIP failed=$FAIL ===" diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/2_assimilation/batch_run_da_on_all_cats.sh b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/2_assimilation/batch_run_da_on_all_cats.sh new file mode 100644 index 00000000..8c3af28e --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/2_assimilation/batch_run_da_on_all_cats.sh @@ -0,0 +1,77 @@ +#!/bin/bash +# batch_run_da_on_all_cats.sh +# Runs run_perturbation_da_on.py for all 21 catchments. +# Stages best_params.json from v2_true_enkf_pn for each catchment. +# cat-1016300 is skipped (already done). + +set -e + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +OBS_DIR=/mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance +CFE_DIR=/mnt/disk2/suma_helen_poster/cfe_py +CONFIG_FILE=/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json +PARAM_BOUNDS=/mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results/v2_perturbation_da_on +BEST_PARAMS_SRC=/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_pn +FORCING_DIR1=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings +FORCING_DIR2=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings + +SKIP=0 +DONE=0 +FAIL=0 + +for CAT in "${CATS[@]}"; do + echo "" + echo "===============================" + echo "=== $CAT ===" + echo "===============================" + + CAT_OUT=$OUT_DIR/$CAT + mkdir -p "$CAT_OUT" + + # Stage best_params if not already present + if [ ! -f "$CAT_OUT/${CAT}_best_params.json" ]; then + SRC="$BEST_PARAMS_SRC/$CAT/${CAT}_best_params.json" + if [ -f "$SRC" ]; then + cp "$SRC" "$CAT_OUT/" + echo " Staged best_params from v2_true_enkf_pn" + else + echo " WARNING: No best_params for $CAT — skipping" + SKIP=$((SKIP + 1)) + continue + fi + else + echo " best_params already staged" + fi + + # Skip if both arm CSVs already exist + if [ -f "$CAT_OUT/${CAT}_da_forcing_arm.csv" ] && \ + [ -f "$CAT_OUT/${CAT}_da_hydro_arm.csv" ]; then + echo " Arms already complete — skipping" + SKIP=$((SKIP + 1)) + continue + fi + + python3 ~/run_perturbation_da_on.py \ + --cat-id "$CAT" \ + --forcing-dir "$FORCING_DIR1" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --param-bounds "$PARAM_BOUNDS" \ + --out-dir "$OUT_DIR" \ + --test-forcing-dir1 "$FORCING_DIR1" \ + --test-forcing-dir2 "$FORCING_DIR2" \ + --hardcoded-r 0.07 \ + && DONE=$((DONE + 1)) \ + || { echo " FAILED: $CAT"; FAIL=$((FAIL + 1)); } +done + +echo "" +echo "=== Batch complete: done=$DONE skipped=$SKIP failed=$FAIL ===" diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/2_assimilation/check_helene_precip_totals.py b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/2_assimilation/check_helene_precip_totals.py new file mode 100644 index 00000000..83c0cdce --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/2_assimilation/check_helene_precip_totals.py @@ -0,0 +1,96 @@ +""" +Sanity check: NWM operational precip totals during Hurricane Helene across +all 21 catchments of USGS gauge 03463300 (South Toe River near Celo, NC). + +Reads the precip_mm_h column from each catchment's _test_results.csv +(produced by the production EnKF run), sums hourly precip over the Helene +event window, and prints per-catchment totals + per-day breakdown. + +Reference event totals for the basin (Helene, Sep 25-28 2024): + - Asheville area: ~300-350 mm 3-day total + - Mt. Mitchell area: 500+ mm reported (South Toe basin is right here) + - General western NC: 250-500+ mm common at elevation + +If the NWM operational totals across the 21 catchments are < 200 mm typical, +the forcing is the dominant bottleneck and no state-space DA can manufacture +the missing runoff. That shifts the diagnosis from "DA tuning" +to "forcing quality / precip ensemble". +""" +import glob +import os +import pandas as pd +import numpy as np + +DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" + +EVENT_START = pd.Timestamp("2024-09-25 00:00:00") +EVENT_END = pd.Timestamp("2024-09-28 00:00:00") + +DAY_WINDOWS = [ + ("sep_25_mm", "2024-09-25 00:00:00", "2024-09-26 00:00:00"), + ("sep_26_mm", "2024-09-26 00:00:00", "2024-09-27 00:00:00"), + ("sep_27_mm", "2024-09-27 00:00:00", "2024-09-28 00:00:00"), +] + + +def main(): + cat_csvs = sorted(glob.glob(os.path.join(DA_DIR, "cat-*", "cat-*_test_results.csv"))) + if not cat_csvs: + print(f"No catchment CSVs found under {DA_DIR}.") + return + + rows = [] + for csv in cat_csvs: + cat_id = os.path.basename(csv).replace("_test_results.csv", "") + df = pd.read_csv(csv, parse_dates=["date"]) + if "precip_mm_h" not in df.columns: + print(f" WARN: {cat_id} missing precip_mm_h column — skipping") + continue + ev = df[(df["date"] >= EVENT_START) & (df["date"] < EVENT_END)] + total = float(ev["precip_mm_h"].sum()) + peak = float(ev["precip_mm_h"].max()) + peak_time = ev.loc[ev["precip_mm_h"].idxmax(), "date"] if len(ev) > 0 else pd.NaT + row = { + "cat_id": cat_id, + "total_3day_mm": total, + "peak_hourly_mm_h": peak, + "peak_time": peak_time, + } + for label, s, e in DAY_WINDOWS: + day = df[(df["date"] >= pd.Timestamp(s)) & (df["date"] < pd.Timestamp(e))] + row[label] = float(day["precip_mm_h"].sum()) + rows.append(row) + + out = pd.DataFrame(rows).sort_values("total_3day_mm", ascending=False) + + print("=" * 90) + print("NWM operational precip totals across 21 catchments — Helene window (Sep 25-27 2024)") + print("=" * 90) + with pd.option_context("display.max_rows", None, + "display.float_format", "{:.1f}".format, + "display.width", 160): + print(out.to_string(index=False)) + print("=" * 90) + print(f"Cross-catchment 3-day total (Sep 25-27):") + print(f" min : {out['total_3day_mm'].min():.1f} mm") + print(f" mean : {out['total_3day_mm'].mean():.1f} mm") + print(f" max : {out['total_3day_mm'].max():.1f} mm") + print() + print(f"Peak hourly intensity across catchments:") + print(f" min : {out['peak_hourly_mm_h'].min():.2f} mm/h") + print(f" mean : {out['peak_hourly_mm_h'].mean():.2f} mm/h") + print(f" max : {out['peak_hourly_mm_h'].max():.2f} mm/h") + print() + print("Reference (observed event totals):") + print(" - Western NC widely reported 250-500+ mm 3-day totals") + print(" - South Toe basin sits in the high-rainfall corridor") + print(" - Asheville (downstream, lower elevation) reported ~300-350 mm") + print("=" * 90) + + out_csv = os.path.join(DA_DIR, "_helene_precip_totals_all_cats.csv") + out.to_csv(out_csv, index=False) + print(f"\nSaved per-catchment totals: {out_csv}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/2_assimilation/input_enkf_new.json b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/2_assimilation/input_enkf_new.json new file mode 100644 index 00000000..482233bc --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/2_assimilation/input_enkf_new.json @@ -0,0 +1 @@ +{"n": 1, "m": 1, "R": 0.07, "Q": 0, "smcmax" : 0.9394097311639178, "N": 1000, "P": 0.01, "D":2} \ No newline at end of file diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/2_assimilation/new_EnKF.py b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/2_assimilation/new_EnKF.py new file mode 100644 index 00000000..ff945c73 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/2_assimilation/new_EnKF.py @@ -0,0 +1,53 @@ +# EnKF.py +import numpy as np + +class EnKF: + def __init__(self, n, m, R, Q, N, P, smcmax, D): + self.n = n # number of states + self.m = m # number of measurements + self.R = R # measurement noise + self.Q = Q # process noise + self.N = N # number of ensembles + self.P = P + self.state_estimates = np.zeros((self.n)) + self.covariance_matrices = np.zeros((self.n, self.n)) + self.current_step = 0 + self.smcmax = smcmax + self.D = D + self.storage_max_m = self.smcmax * self.D + self.storage_init = self.storage_max_m * 0.667 + self.ensembles = np.full((self.n, self.N), self.storage_init) + + def predict(self, F_results): + error_factor_sim = 0.005 + for i in range(self.N): + perturbation_factor_sim = np.random.standard_normal() + self.ensembles[:, i] = F_results[i] + self.ensembles[:, i] += (perturbation_factor_sim * F_results[i] * error_factor_sim) + return self.ensembles + + def update(self, yi, H_results): + if np.any(np.isnan(yi)): + return np.mean(self.ensembles, axis=1) # return current state estimate without updating + + y_ensembles = np.zeros((self.m, self.N)) + error_factor_sm = 0.03 + for i in range(self.N): + perturbation_factor_sm = np.random.standard_normal() + y_ensembles[:, i] = H_results[i] + y_ensembles[:, i] += (perturbation_factor_sm * H_results[i] * error_factor_sm) + + ensemble_mean = np.mean(self.ensembles, axis=1)[:, np.newaxis] + y_ensembles_mean = np.mean(y_ensembles, axis=1)[:, np.newaxis] + Pxy = (self.ensembles - ensemble_mean).dot((y_ensembles - y_ensembles_mean).T) / (self.N - 1) + Pyy = np.cov(y_ensembles, bias=True) + K = Pxy.dot(np.linalg.pinv(Pyy + self.R)) + self.ensembles += K.dot(yi - y_ensembles_mean) + return np.mean(self.ensembles, axis=1) + + def get_state_estimate(self): + return np.mean(self.ensembles, axis=1) + + def get_covariance(self): + ensemble_mean = np.mean(self.ensembles, axis=1)[:, np.newaxis] + return (self.ensembles - ensemble_mean).dot((self.ensembles - ensemble_mean).T) / (self.N - 1) diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/2_assimilation/run_crossed_ensemble.py b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/2_assimilation/run_crossed_ensemble.py new file mode 100644 index 00000000..325a64e2 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/2_assimilation/run_crossed_ensemble.py @@ -0,0 +1,452 @@ +""" +run_crossed_ensemble.py -- 4b: 600-member crossed ensemble for Helene. + +Crosses the two DA-on perturbation arms: + Forcing arm : 30 met draws (from run_perturbation_da_on.py Phase 2A) + Hydro arm : 20 state draws (from run_perturbation_da_on.py Phase 2B) + +For each Helene issue time t0: + Run 600 members = 30 forcing draws x 20 hydro-state draws. + Each member (i, j) gets: + - Initial state = perturbed DA analysis state (draw j of hydro arm) + - Met forcing = perturbed precip/PET sequence (draw i of forcing arm) + +This gives the full uncertainty envelope combining both sources — +the operational probabilistic forecast picture. + +Output (per catchment): + //_crossed_ensemble.parquet + Columns: issue_time, lead_hour, member_0000..member_0599 + +Usage: + python3 run_crossed_ensemble.py \\ + --cat-id cat-1016300 \\ + --forcing-dir \\ + --obs-dir /mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance \\ + --cfe-dir /mnt/disk2/suma_helen_poster/cfe_py \\ + --config-file /mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json \\ + --out-dir /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble \\ + --test-forcing-dir1 \\ + --test-forcing-dir2 + +Strategy: + Re-runs DA analysis (same as Phase 1 of run_perturbation_da_on.py) to get + the per-issue-time analysis state. Then generates N_FORCING sets of perturbed + forcing sequences and N_HYDRO perturbed states. Runs all 600 combinations. + + To avoid re-running Phase 1, if the Phase 1 snapshot file exists from a + prior run_perturbation_da_on.py run, it is loaded directly. + +NOTE: All special characters kept ASCII to avoid encoding issues on GPU. +""" + +import argparse +import os +import sys +import json +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +N_FORCING = 30 +N_HYDRO = 20 +N_MEMBERS = N_FORCING * N_HYDRO # 600 + +TEST_START = "2024-08-24 00:00:00" # 1 month spinup before Helene window +TEST_END = "2024-10-31 23:00:00" + +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-30 06:00:00") +FORECAST_HOURS = 18 + +PRECIP_SIGMA = 0.15 +PET_SIGMA = 0.10 +STATE_FRAC = 0.05 + +CAT_ID = None +OBS_FILE = None +CFE_CONFIG_FILE = None +TEST_FORCING_FILE = None +OUT_DIR = None +HARDCODED_R = None +bmi_cfe = None + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(): + df = pd.read_csv(TEST_FORCING_FILE) + df["date"] = pd.to_datetime(df["time"]).dt.strftime("%Y-%m-%d %H:%M:%S") + df["total_precipitation"] = df["APCP_surface"] * 3600.0 + df["potential_evaporation"] = priestley_taylor_pet( + df["DSWRF_surface"].values, df["TMP_2maboveground"].values) + return df + + +def load_obs(): + df = pd.read_csv(OBS_FILE) + date_col = next((c for c in df.columns + if c.lower() in ("date", "time", "datetime", "timestamp")), None) + q_col = next((c for c in df.columns + if "qkrig" in c.lower() and "var" not in c.lower()), None) + var_col = next((c for c in df.columns if "var" in c.lower()), None) + df[date_col] = pd.to_datetime(df[date_col]) + df = df.sort_values(date_col).reset_index(drop=True) + date_strs = df[date_col].dt.strftime("%Y-%m-%d %H:%M:%S") + obs_dict = dict(zip(date_strs, df[q_col])) + var_vals = df[var_col] if var_col else pd.Series([0.07] * len(df)) + var_dict = dict(zip(date_strs, var_vals)) + return obs_dict, var_dict + + +def write_model_config(best_params): + """Write the BMI config file once; reused for every model instantiation.""" + with open(CFE_CONFIG_FILE) as f: + cfg = json.load(f) + cfg["forcing_file"] = TEST_FORCING_FILE + cfg["soil_params"]["bb"] = best_params["bb"] + cfg["soil_params"]["smcmax"] = best_params["smcmax"] + cfg["soil_params"]["satdk"] = best_params["satdk"] + cfg["slop"] = best_params["slop"] + cfg["max_gw_storage"] = best_params["max_gw_storage"] + cfg["expon"] = best_params["expon"] + cfg["Cgw"] = best_params["Cgw"] + cfg["K_lf"] = best_params["K_lf"] + cfg["K_nash"] = best_params["K_nash"] + cfg["partition_scheme"] = ("Schaake" if best_params["scheme"] <= 0.5 + else "Xinanjiang") + tmp_cfg = str(OUT_DIR / f"{CAT_ID}_bmi_config_crossed_stable.json") + with open(tmp_cfg, "w") as f: + json.dump(cfg, f) + return tmp_cfg + + +def make_model(cfg_path): + """Instantiate and initialize a CFE model from a pre-written config path.""" + m = bmi_cfe.BMI_CFE(cfg_file=cfg_path) + m.initialize() + return m + + +def get_state(m): + return { + "soil_m": m.soil_reservoir["storage_m"], + "gw_m": m.gw_reservoir["storage_m"], + "nash0": float(m.nash_storage[0]), + "nash1": float(m.nash_storage[1]), + } + + +def set_state(m, state): + m.soil_reservoir["storage_m"] = state["soil_m"] + m.gw_reservoir["storage_m"] = state["gw_m"] + m.nash_storage[0] = state["nash0"] + m.nash_storage[1] = state["nash1"] + + +def perturb_state(state, rng, frac=STATE_FRAC): + return { + "soil_m": max(state["soil_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "nash0": max(state["nash0"] * (1.0 + frac * rng.standard_normal()), 0.0), + "nash1": max(state["nash1"] * (1.0 + frac * rng.standard_normal()), 0.0), + } + + +def run_model_step(m, precip_mmh, pet_mmh): + m.set_value("atmosphere_water__time_integral_of_precipitation_mass_flux", + precip_mmh / 1000.0) + m.set_value("water_potential_evaporation_flux", + pet_mmh / 1000.0 / 3600.0) + m.update() + return m.get_value("land_surface_water__runoff_depth") * 1000.0 + + +def enkf_update(state, q_sim, y_obs, krig_var, rng): + R = HARDCODED_R if HARDCODED_R is not None else max( + (0.10 * max(y_obs, 0.0)) ** 2 + 0.001 * krig_var, 1e-6) + R = max(R, 1e-6) + P_yy = max(q_sim * 0.01, 1e-6) + K = P_yy / (P_yy + R) + scale = K * (y_obs - q_sim) / max(abs(q_sim), 1e-6) + return { + "soil_m": max(state["soil_m"] * (1.0 + scale), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + scale), 1e-6), + "nash0": max(state["nash0"] + state["nash0"] * scale, 0.0), + "nash1": max(state["nash1"] + state["nash1"] * scale, 0.0), + } + + +def run_phase1_da(cfg_path, df_test, obs_dict, var_dict, rng): + """Run DA through test period; return snapshot dict {timestamp -> state}.""" + print(f" Phase 1: DA analysis through test period...") + model = make_model(cfg_path) + snapshots = {} + for _, row in df_test.iterrows(): + date_str = row["date"] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + q_sim = run_model_step(model, P, E) + y_obs = obs_dict.get(date_str, np.nan) + if not np.isnan(y_obs): + krig_var = var_dict.get(date_str, 0.07) + state = get_state(model) + updated = enkf_update(state, q_sim, y_obs, krig_var, rng) + set_state(model, updated) + date_ts = pd.Timestamp(date_str) + if HELENE_START <= date_ts <= HELENE_END: + snapshots[date_ts] = get_state(model) + model.finalize() + print(f" Phase 1 done: {len(snapshots)} issue-time snapshots") + return snapshots + + +def main(): + global CAT_ID, OBS_FILE, CFE_CONFIG_FILE, TEST_FORCING_FILE + global OUT_DIR, HARDCODED_R, bmi_cfe + + parser = argparse.ArgumentParser() + parser.add_argument("--cat-id", required=True) + parser.add_argument("--forcing-dir", required=True) + parser.add_argument("--obs-dir", required=True) + parser.add_argument("--cfe-dir", required=True) + parser.add_argument("--config-file", required=True) + parser.add_argument("--out-dir", required=True) + parser.add_argument("--test-forcing-dir1", default=None) + parser.add_argument("--test-forcing-dir2", default=None) + parser.add_argument("--hardcoded-r", type=float, default=None, + help="Fixed R value for EnKF (e.g. 0.07). " + "Omit to use dynamic Vrugt R: R=(0.10*y)^2 + 0.001*krig_var") + parser.add_argument("--n-forcing", type=int, default=N_FORCING, + help="Number of met forcing draws (default 30)") + parser.add_argument("--n-hydro", type=int, default=N_HYDRO, + help="Number of hydro-state draws (default 20)") + args = parser.parse_args() + + CAT_ID = args.cat_id + OBS_FILE = os.path.join(args.obs_dir, f"{CAT_ID}.csv") + CFE_CONFIG_FILE = args.config_file + HARDCODED_R = args.hardcoded_r + OUT_DIR = Path(args.out_dir) / CAT_ID + OUT_DIR.mkdir(parents=True, exist_ok=True) + + n_fa = args.n_forcing + n_ha = args.n_hydro + n_total = n_fa * n_ha + print(f"[{CAT_ID}] Crossed ensemble: {n_fa} forcing x {n_ha} hydro = {n_total} members") + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + bmi_cfe = _bmi_cfe + + # Reuse combined forcing CSV from run_perturbation_da_on.py if it exists + # (avoids writing a duplicate large file to disk) + da_on_combined = (Path(args.out_dir).parent / "v2_perturbation_da_on" / + CAT_ID / f"{CAT_ID}_nwm_operational_combined.csv") + crossed_combined = OUT_DIR / f"{CAT_ID}_nwm_operational_combined.csv" + + if da_on_combined.exists(): + TEST_FORCING_FILE = str(da_on_combined) + print(f" Reusing combined forcing: {TEST_FORCING_FILE}") + elif crossed_combined.exists(): + TEST_FORCING_FILE = str(crossed_combined) + print(f" Reusing combined forcing: {TEST_FORCING_FILE}") + elif args.test_forcing_dir1 and args.test_forcing_dir2: + f1 = os.path.join(args.test_forcing_dir1, f"{CAT_ID}.csv") + f2 = os.path.join(args.test_forcing_dir2, f"{CAT_ID}.csv") + df1, df2 = pd.read_csv(f1), pd.read_csv(f2) + combined = (pd.concat([df1, df2], ignore_index=True) + .drop_duplicates(subset="time").sort_values("time")) + TEST_FORCING_FILE = str(crossed_combined) + combined.to_csv(TEST_FORCING_FILE, index=False) + print(f" Wrote combined forcing: {TEST_FORCING_FILE}") + else: + print("Need --test-forcing-dir1 and --test-forcing-dir2"); return + + best_params_file = OUT_DIR / f"{CAT_ID}_best_params.json" + if not best_params_file.exists(): + # Also check the DA results dir from run_perturbation_da_on.py + alt = Path(args.out_dir).parent / "v2_perturbation_da_on" / CAT_ID / f"{CAT_ID}_best_params.json" + if alt.exists(): + best_params_file = alt + else: + print(f"Missing best params at {best_params_file}"); return + with open(best_params_file) as f: + best_params = json.load(f)["best_parameters"] + + obs_dict, var_dict = load_obs() + df_forcing = load_test_forcing() + t_mask = (df_forcing["date"] >= TEST_START) & (df_forcing["date"] <= TEST_END) + df_test = df_forcing[t_mask].reset_index(drop=True) + + # Write BMI config once — reused for all 90,600 model instantiations + cfg_path = write_model_config(best_params) + print(f" BMI config written once: {cfg_path}") + + rng = np.random.default_rng(hash(CAT_ID) & 0x7fffffff) + + # Phase 1: get DA analysis snapshots + # Check if a prior run already saved them (skip re-running DA if so) + snap_cache = OUT_DIR / f"{CAT_ID}_da_snapshots.parquet" + if snap_cache.exists(): + print(f" Loading cached DA snapshots from {snap_cache}") + df_snaps = pd.read_parquet(snap_cache) + snapshots = { + pd.Timestamp(row["issue_time"]): { + "soil_m": row["soil_m"], "gw_m": row["gw_m"], + "nash0": row["nash0"], "nash1": row["nash1"], + } + for _, row in df_snaps.iterrows() + } + else: + snapshots = run_phase1_da(cfg_path, df_test, obs_dict, var_dict, rng) + + issue_times = sorted(snapshots.keys()) + t0_to_idx = {t: i for i, t in enumerate(df_test["date"].apply(pd.Timestamp))} + + # ------------------------------------------------------------------ # + # Crossed ensemble: for each issue_time, build n_fa forcing draws and # + # n_ha state draws, then run all n_fa x n_ha combinations. # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Running crossed ensemble ({n_total} members x " + f"{len(issue_times)} issue times)...") + + all_records = [] + all_hydro_draw_records = [] # (issue_time, hydro_draw_j, states) + all_forcing_draw_records = [] # (issue_time, forcing_draw_i, lead, P_pert, E_pert, scales) + + for t_idx, t0 in enumerate(issue_times): + snap = snapshots[t0] + t0_str = t0.strftime("%Y-%m-%d %H:%M:%S") + if t0 not in t0_to_idx: + continue + start_idx = t0_to_idx[t0] + + # Pre-draw n_fa forcing perturbation sequences + forcing_seqs = [] + for i in range(n_fa): + seq = [] + for lead in range(1, FORECAST_HOURS + 1): + fi = start_idx + lead + if fi >= len(df_test): + seq.append((0.0, 0.0)) + all_forcing_draw_records.append({ + "issue_time": t0_str, "forcing_draw_i": i, + "lead_hour": lead, + "P_pert_mmh": 0.0, "E_pert_mmh": 0.0, + "precip_scale": np.nan, "pet_scale": np.nan, + }) + continue + row = df_test.iloc[fi] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + mu_p = -0.5 * PRECIP_SIGMA ** 2 + P_pert = float(P * rng.lognormal(mu_p, PRECIP_SIGMA)) if P > 0 else 0.0 + E_pert = max(float(E * (1.0 + PET_SIGMA * rng.standard_normal())), 0.0) + seq.append((P_pert, E_pert)) + all_forcing_draw_records.append({ + "issue_time": t0_str, "forcing_draw_i": i, + "lead_hour": lead, + "P_pert_mmh": P_pert, "E_pert_mmh": E_pert, + "precip_scale": P_pert / P if P > 0 else np.nan, + "pet_scale": E_pert / E if E > 0 else np.nan, + }) + forcing_seqs.append(seq) + + # Pre-draw n_ha hydro-state perturbations + state_draws = [snap if j == 0 else perturb_state(snap, rng) + for j in range(n_ha)] + + # Record each hydro draw's actual initial state + for j, st in enumerate(state_draws): + all_hydro_draw_records.append({ + "issue_time": t0_str, "hydro_draw_j": j, + "soil_m": st["soil_m"], "gw_m": st["gw_m"], + "nash0": st["nash0"], "nash1": st["nash1"], + }) + + # q_matrix shape: (FORECAST_HOURS, n_fa, n_ha) + q_matrix = np.full((FORECAST_HOURS, n_fa, n_ha), np.nan) + + for i, f_seq in enumerate(forcing_seqs): + for j, state_j in enumerate(state_draws): + m = make_model(cfg_path) + set_state(m, state_j) + for lead in range(1, FORECAST_HOURS + 1): + P_pert, E_pert = f_seq[lead - 1] + q = run_model_step(m, P_pert, E_pert) + q_matrix[lead - 1, i, j] = q + m.finalize() + + # Flatten 600-member matrix into rows + for lead in range(1, FORECAST_HOURS + 1): + row_data = {"issue_time": t0_str, "lead_hour": lead} + mem_idx = 0 + for i in range(n_fa): + for j in range(n_ha): + row_data[f"member_{mem_idx:04d}"] = q_matrix[lead - 1, i, j] + mem_idx += 1 + all_records.append(row_data) + + if (t_idx + 1) % 10 == 0 or (t_idx + 1) == len(issue_times): + print(f" {t_idx + 1}/{len(issue_times)} issue times complete") + + df_out = pd.DataFrame(all_records) + out_path = OUT_DIR / f"{CAT_ID}_crossed_ensemble.parquet" + df_out.to_parquet(out_path, index=False) + print(f"Saved: {out_path}") + print(f" Shape: {df_out.shape} " + f"({df_out['issue_time'].nunique()} issue times x " + f"{FORECAST_HOURS} leads x {n_total} members)") + + # ------------------------------------------------------------------ # + # Provenance files — enable full ensemble traceability and restart # + # ------------------------------------------------------------------ # + + # 1. DA analysis snapshots — unperturbed state at each issue_time. + # Load any row and set_state() to restart from that init time. + snap_records = [ + {"issue_time": t.strftime("%Y-%m-%d %H:%M:%S"), **s} + for t, s in snapshots.items() + ] + snap_path = OUT_DIR / f"{CAT_ID}_da_snapshots.parquet" + pd.DataFrame(snap_records).to_parquet(snap_path, index=False) + print(f"Saved: {snap_path} ({len(snap_records)} snapshots)") + + # 2. Member manifest — decoder ring: member_col -> (forcing_draw_i, hydro_draw_j). + # member_k uses forcing draw k//n_ha and hydro draw k%n_ha. + manifest = [ + {"member": f"member_{k:04d}", "member_idx": k, + "forcing_draw_i": k // n_ha, "hydro_draw_j": k % n_ha} + for k in range(n_total) + ] + manifest_path = OUT_DIR / f"{CAT_ID}_member_manifest.csv" + pd.DataFrame(manifest).to_csv(manifest_path, index=False) + print(f"Saved: {manifest_path} ({n_total} members)") + + # 3. Hydro draw states — actual perturbed initial state for each of + # n_ha draws at every issue_time. Links member -> initial conditions. + hydro_path = OUT_DIR / f"{CAT_ID}_hydro_draw_states.parquet" + pd.DataFrame(all_hydro_draw_records).to_parquet(hydro_path, index=False) + print(f"Saved: {hydro_path} ({len(all_hydro_draw_records)} rows)") + + # 4. Forcing draw sequences — actual P/PET perturbation values for each + # of n_fa draws at every (issue_time, lead_hour). Links member -> forcing. + forcing_path = OUT_DIR / f"{CAT_ID}_forcing_draw_sequences.parquet" + pd.DataFrame(all_forcing_draw_records).to_parquet(forcing_path, index=False) + print(f"Saved: {forcing_path} ({len(all_forcing_draw_records)} rows)") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/2_assimilation/run_lead_time_forecast_sweep.py b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/2_assimilation/run_lead_time_forecast_sweep.py new file mode 100644 index 00000000..19f3b4d5 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/2_assimilation/run_lead_time_forecast_sweep.py @@ -0,0 +1,508 @@ +""" +Forecast lead-time evaluation. + +For each issue-time t0 in a sampling schedule across the test period, fork the +ensemble at t0 and free-run 18 hours of forecast. Two parallel trajectories are +maintained throughout the test period: + + da: Production setup. DA on, forcing perturbed, process noise on. + openloop: No DA from t=0. Forcing perturbed, process noise on. + +At each issue time t0, the state of each scenario's ensemble is copied into a +forecast ensemble that free-runs 18 hours with: + - DA off (no observations consumed during the forecast window) + - process noise off (no obs to collapse toward, so no anti-collapse needed) + - forcing perturbed (lognormal precip, Gaussian PET — proxy for forecast + precip uncertainty in lieu of actual HEFS forecasts) + +Issue-time schedule: + - Base cadence: every --base-step-h hours across the full test period + (default 6h → ~1600 issue times for the year) + - Densified to hourly across --dense-start..--dense-end if provided + (default: 2024-09-24 → 2024-09-28, the Helene window) + +Outputs (stacked across all issue times): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv + Columns: issue_time, lead_hour, valid_time, member_00 .. member_19 + (units: mm/h) + +Post-processing pipeline (separate scripts, run after this): + route_lead_time_forecasts.py — pushes each (t0, lead_hour) forecast through + T-route to the gauge + plot_lead_time_decay.py — error-vs-lead-time curve, DA vs open-loop + +Imports EnKFAssimilator from calibrate_catchment_cfe_da_v2.py so the DA math +matches production exactly. +""" + +import argparse +import os +import sys +import json +import importlib.util +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +SPINUP_START = "2023-02-01 00:00:00" +SPINUP_END = "2023-09-30 23:00:00" +TEST_START = "2023-10-01 00:00:00" +TEST_END = "2024-10-31 23:00:00" + +# Forecast lead time (hours after each issue time) +FORECAST_LEAD_HOURS = 18 + +# Default densification window — Hurricane Helene +DEFAULT_DENSE_START = "2024-09-24 00:00:00" +DEFAULT_DENSE_END = "2024-09-28 23:00:00" + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(test_forcing_file): + df = pd.read_csv(test_forcing_file) + df['date'] = pd.to_datetime(df['time']).dt.strftime('%Y-%m-%d %H:%M:%S') + df['total_precipitation'] = df['APCP_surface'] * 3600.0 + df['potential_evaporation'] = priestley_taylor_pet( + df['DSWRF_surface'].values, df['TMP_2maboveground'].values + ) + return df + + +def import_enkf_class(script_path): + """Load EnKFAssimilator from the production script without triggering its main().""" + spec = importlib.util.spec_from_file_location("prod_v2", script_path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod.EnKFAssimilator + + +# ---------------- State snapshot / restore ----------------------------------- +# +# The EnKF only ever touches these 4 states, so they're the only states that +# diverge between members. Copy these from src → dst to fork the ensemble. + +def snapshot_states(models): + return [ + { + "soil_m": float(m.soil_reservoir["storage_m"]), + "gw_m": float(m.gw_reservoir["storage_m"]), + "nash0_m": float(m.nash_storage[0]), + "nash1_m": float(m.nash_storage[1]), + } + for m in models + ] + + +def restore_states(models, snapshot): + for m, s in zip(models, snapshot): + m.soil_reservoir["storage_m"] = s["soil_m"] + m.gw_reservoir["storage_m"] = s["gw_m"] + m.nash_storage[0] = s["nash0_m"] + m.nash_storage[1] = s["nash1_m"] + + +# ---------------- Issue-time schedule ---------------------------------------- + +def build_issue_time_schedule(dates_list, base_step_h, dense_start, dense_end): + """Return a sorted list of issue-time strings sampled from dates_list. + + Every base_step_h hours across the full test period, plus every hour inside + [dense_start, dense_end] if those are provided. + """ + dates_dt = pd.to_datetime(dates_list) + selected = set() + + # Base cadence (every base_step_h hours from the first date) + base_mask = (np.arange(len(dates_dt)) % base_step_h == 0) + for d in dates_dt[base_mask]: + selected.add(d.strftime('%Y-%m-%d %H:%M:%S')) + + # Densified hourly window + if dense_start is not None and dense_end is not None: + ds = pd.Timestamp(dense_start) + de = pd.Timestamp(dense_end) + dense_mask = (dates_dt >= ds) & (dates_dt <= de) + for d in dates_dt[dense_mask]: + selected.add(d.strftime('%Y-%m-%d %H:%M:%S')) + + return sorted(selected) + + +# ---------------- CFE helpers ------------------------------------------------ + +def build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, enkf, + apply_init_perturbation=True, member_zero_clean=True): + """Build N CFE BMI instances with optional init-state perturbation.""" + def custom_load_forcing(self_cfe): + df = load_test_forcing(test_forcing_file) + self_cfe.forcing_data = df.rename(columns={"date": "time"}) + + models = [] + for i in range(N): + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.load_forcing_file = custom_load_forcing.__get__(m) + m.initialize() + if apply_init_perturbation and not (member_zero_clean and i == 0): + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + r = enkf.init_state_perturb_frac + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1 + r * enkf.rng.standard_normal()), 1e-6, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1 + r * enkf.rng.standard_normal()), 1e-6, gw_max)) + jitter = max(sm_max * 1e-6, 1e-9) + m.nash_storage[0] = max(n00 + jitter * enkf.rng.standard_normal(), 0.0) + m.nash_storage[1] = max(n10 + jitter * enkf.rng.standard_normal(), 0.0) + models.append(m) + return models + + +def step_ensemble(models, p_arr, e_arr): + """Advance each member one hour with member-specific (P, PET). Returns Q (mm/h).""" + q = np.empty(len(models), dtype=float) + for i, m in enumerate(models): + m.set_value('atmosphere_water__time_integral_of_precipitation_mass_flux', + float(p_arr[i]) / 1000.0) + m.set_value('water_potential_evaporation_flux', + float(e_arr[i]) / 1000.0 / 3600.0) + m.update() + q[i] = m.get_value('land_surface_water__runoff_depth') * 1000.0 # m/h → mm/h + return q + + +def do_sanity_check(prod_models, fcst_models, dates_list, forcing_by_date, h, N): + """One-shot verification that the 4-state snapshot is sufficient to reproduce + next-hour forecast Q. + + At main-loop hour h (after prod's hour-h step has completed including DA + + process noise), snapshot prod, restore into fcst, then step BOTH ensembles + one hour using identical unperturbed forcing. If snapshot/restore captures + everything that matters, prod's and fcst's hour-(h+1) Q should match member- + by-member to machine precision. + + Note: this mutates prod's state (h+1 step with non-perturbed forcing and no + DA), so the caller should exit the script right after. + """ + if h + 1 >= len(dates_list): + print("[sanity-check] not enough horizon for h+1 step — skipping.") + return False + snap = snapshot_states(prod_models) + restore_states(fcst_models, snap) + + next_date = dates_list[h + 1] + p_next, e_next = forcing_by_date[next_date] + p_uniform = np.full(N, p_next, dtype=float) + e_uniform = np.full(N, e_next, dtype=float) + + q_prod = step_ensemble(prod_models, p_uniform, e_uniform) + q_fcst = step_ensemble(fcst_models, p_uniform, e_uniform) + + diff = q_prod - q_fcst + max_abs = float(np.max(np.abs(diff))) + print("[sanity-check] forecast Q at h+1 with identical unperturbed forcing:") + print(f" prod first 5 members: {q_prod[:5]}") + print(f" fcst first 5 members: {q_fcst[:5]}") + print(f" max |prod - fcst|: {max_abs:.3e} mm/h") + if max_abs < 1e-9: + print(" OK: snapshot/restore preserves state to machine precision.") + return True + print(" WARN: nonzero divergence — likely a state not in the 4-state snapshot.") + return False + + +def run_forecast(fcst_models, enkf_fcst, dates_list, forcing_by_date, + t0_idx, n_lead): + """Run an n_lead-hour free-forecast starting from the current state of fcst_models. + + Forcing is perturbed per member (lognormal precip + Gaussian PET, same as + production). No DA. No process noise. + + Returns: (lead_hours, valid_times, q_matrix shape (n_lead, N)). + """ + N = len(fcst_models) + q_matrix = np.full((n_lead, N), np.nan, dtype=float) + lead_hours = [] + valid_times = [] + for lead in range(1, n_lead + 1): + idx = t0_idx + lead + if idx >= len(dates_list): + break + valid_date = dates_list[idx] + p, e = forcing_by_date[valid_date] + p_arr, e_arr = enkf_fcst.perturb_forcing(p, e) + q = step_ensemble(fcst_models, p_arr, e_arr) + q_matrix[lead - 1, :] = q + lead_hours.append(lead) + valid_times.append(valid_date) + return lead_hours, valid_times, q_matrix + + +# ---------------- Main run --------------------------------------------------- + +def run(args, bmi_cfe, EnKFAssimilator): + cat_id = args.cat_id + out_dir = Path(args.out_dir) / cat_id + out_dir.mkdir(parents=True, exist_ok=True) + + # Combined test forcing (same idiom as the per-member script) + f1 = os.path.join(args.test_forcing_dir1, f'{cat_id}.csv') + f2 = os.path.join(args.test_forcing_dir2, f'{cat_id}.csv') + df1 = pd.read_csv(f1) + df2 = pd.read_csv(f2) + combined = pd.concat([df1, df2], ignore_index=True) + combined = combined.drop_duplicates(subset='time').sort_values('time') + test_forcing_file = str(out_dir / f'{cat_id}_nwm_operational_combined.csv') + combined.to_csv(test_forcing_file, index=False) + + # Pre-staged calibrated params + best_params_file = out_dir / f'{cat_id}_best_params.json' + if not best_params_file.exists(): + print(f"Missing {best_params_file}. Pre-stage from Run 3.") + return + with open(best_params_file) as f: + best = json.load(f)['best_parameters'] + + # Build temp CFE config with calibrated params + with open(args.config_file) as f: + cfg = json.load(f) + cfg['forcing_file'] = test_forcing_file + cfg['soil_params']['bb'] = best['bb'] + cfg['soil_params']['smcmax'] = best['smcmax'] + cfg['soil_params']['satdk'] = best['satdk'] + cfg['slop'] = best['slop'] + cfg['max_gw_storage'] = best['max_gw_storage'] + cfg['expon'] = best['expon'] + cfg['Cgw'] = best['Cgw'] + cfg['K_lf'] = best['K_lf'] + cfg['K_nash'] = best['K_nash'] + cfg['partition_scheme'] = "Schaake" if best['scheme'] <= 0.5 else "Xinanjiang" + tmp_cfg = str(out_dir / f'{cat_id}_bmi_config_temp_leadtime.json') + with open(tmp_cfg, 'w') as f: + json.dump(cfg, f) + + # Three EnKF instances: + # enkf_da — runs the DA-on trajectory; consumes obs each hour + # enkf_ol — runs the open-loop trajectory; never calls update_states + # enkf_fcst — used only inside the 18-hour free-runs (perturb_forcing only) + # All three share defaults with production. Distinct RNG seeds keep them + # independent so the openloop and forecast spreads are not coupled to DA. + obs_file = os.path.join(args.obs_dir, f'{cat_id}.csv') + seed_base = args.rng_seed if args.rng_seed is not None else 0 + enkf_da = EnKFAssimilator( + n_members=args.enkf_members, obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, use_vrugt_r=(not args.no_vrugt_r), + vrugt_alpha=args.vrugt_alpha, vrugt_scale=args.vrugt_scale, + rng_seed=seed_base if args.rng_seed is not None else None, + ) + enkf_ol = EnKFAssimilator( + n_members=args.enkf_members, obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, use_vrugt_r=(not args.no_vrugt_r), + vrugt_alpha=args.vrugt_alpha, vrugt_scale=args.vrugt_scale, + rng_seed=(seed_base + 1) if args.rng_seed is not None else None, + ) + enkf_fcst = EnKFAssimilator( + n_members=args.enkf_members, obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, use_vrugt_r=(not args.no_vrugt_r), + vrugt_alpha=args.vrugt_alpha, vrugt_scale=args.vrugt_scale, + rng_seed=(seed_base + 2) if args.rng_seed is not None else None, + ) + # Optional R override — replace every per-hour obs variance with a constant. + # Matches DualEarth/new_EnKF.py (R=0.07). Set via --hardcoded-r on the CLI. + # Effect: at storm peak ≈ existing Vrugt R, at low flow ≈ 15-17× larger → + # near-zero gain at low flow where kriging obs is noisy. + if args.hardcoded_r is not None: + for enkf in (enkf_da, enkf_ol, enkf_fcst): + for date in enkf.obs_var_dict: + enkf.obs_var_dict[date] = args.hardcoded_r + print(f"[lead-time] R hardcoded to {args.hardcoded_r} mm^2/h^2 " + f"(overrides Vrugt + kriging variance formula)") + + N = enkf_da.n_members + print(f"[lead-time] {cat_id} | N={N} | DA on/off both run | " + f"forecast lead = {FORECAST_LEAD_HOURS}h") + + # Build three ensembles (init-perturbed except the forecast ensemble, whose + # state will be overwritten at each issue time before each free-run) + prod_models = build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, enkf_da) + openloop_models = build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, enkf_ol) + fcst_models = build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, + enkf_fcst, apply_init_perturbation=False) + + # ----- Spinup with perturbed forcing on both trajectories ----- + df = load_test_forcing(test_forcing_file) + sp_mask = (df['date'] >= SPINUP_START) & (df['date'] <= SPINUP_END) + df_sp = df[sp_mask] + print(f"[lead-time] spinup: {len(df_sp)} hours") + for p, e in zip(df_sp['total_precipitation'], df_sp['potential_evaporation']): + p_da, e_da = enkf_da.perturb_forcing(p, e) + step_ensemble(prod_models, p_da, e_da) + enkf_da.add_process_noise(prod_models) + p_ol, e_ol = enkf_ol.perturb_forcing(p, e) + step_ensemble(openloop_models, p_ol, e_ol) + enkf_ol.add_process_noise(openloop_models) + + # ----- Test period: step both, fork forecasts at issue times ----- + t_mask = (df['date'] >= TEST_START) & (df['date'] <= TEST_END) + df_test = df[t_mask].reset_index(drop=True) + dates_list = list(df_test['date'].values) + forcing_by_date = dict(zip( + df_test['date'].values, + zip(df_test['total_precipitation'].values, + df_test['potential_evaporation'].values))) + + issue_times = build_issue_time_schedule( + dates_list, args.base_step_h, + args.dense_start, args.dense_end) + issue_set = set(issue_times) + print(f"[lead-time] test period: {len(dates_list)} hours | " + f"{len(issue_times)} issue times " + f"(base={args.base_step_h}h, dense={args.dense_start or 'none'}" + f"..{args.dense_end or 'none'})") + + # Pre-build the date→index map so forecast windows are fast to look up + date_to_idx = {d: i for i, d in enumerate(dates_list)} + + da_rows = [] # (issue_time, lead_hour, valid_time, *member_values) + ol_rows = [] + sanity_done = False + + for h, current_date in enumerate(dates_list): + p, e = forcing_by_date[current_date] + + # --- DA trajectory step --- + p_da, e_da = enkf_da.perturb_forcing(p, e) + q_da = step_ensemble(prod_models, p_da, e_da) + enkf_da.update_states(prod_models, current_date, q_da) + enkf_da.add_process_noise(prod_models) + + # --- Open-loop trajectory step (no DA) --- + p_ol, e_ol = enkf_ol.perturb_forcing(p, e) + q_ol = step_ensemble(openloop_models, p_ol, e_ol) + enkf_ol.add_process_noise(openloop_models) + + # --- One-shot sanity check at the first scheduled issue time --- + if args.sanity_check and not sanity_done and current_date in issue_set: + ok = do_sanity_check( + prod_models, fcst_models, dates_list, forcing_by_date, h, N) + sanity_done = True + print(f"[sanity-check] exiting (re-run without --sanity-check for " + f"the full sweep). Result: {'OK' if ok else 'FAIL'}") + for m in prod_models + openloop_models + fcst_models: + m.finalize() + sys.exit(0 if ok else 1) + + # --- Fork forecasts at scheduled issue times --- + if current_date in issue_set: + for scenario_label, src_models, rows_acc in [ + ('da', prod_models, da_rows), + ('openloop', openloop_models, ol_rows), + ]: + snap = snapshot_states(src_models) + restore_states(fcst_models, snap) + leads, valids, qm = run_forecast( + fcst_models, enkf_fcst, + dates_list, forcing_by_date, + t0_idx=h, n_lead=FORECAST_LEAD_HOURS) + for k, (lead, vt) in enumerate(zip(leads, valids)): + rows_acc.append((current_date, lead, vt, *qm[k, :])) + + for m in prod_models + openloop_models + fcst_models: + m.finalize() + + # ----- Save forecast CSVs ----- + cols = ['issue_time', 'lead_hour', 'valid_time'] + [f'member_{i:02d}' for i in range(N)] + da_path = out_dir / f'{cat_id}_lead_time_forecasts_da.csv' + ol_path = out_dir / f'{cat_id}_lead_time_forecasts_openloop.csv' + pd.DataFrame(da_rows, columns=cols).to_csv(da_path, index=False) + pd.DataFrame(ol_rows, columns=cols).to_csv(ol_path, index=False) + print(f"[lead-time] saved {da_path}") + print(f"[lead-time] saved {ol_path}") + + # Issue-time schedule (small, useful for downstream scripts) + sched_path = out_dir / f'{cat_id}_lead_time_issue_times.csv' + pd.DataFrame({'issue_time': issue_times}).to_csv(sched_path, index=False) + print(f"[lead-time] saved {sched_path}") + + print(f"[lead-time] DA assimilator: updates={enkf_da.n_updates} | " + f"avg Pyy={enkf_da.avg_pyy:.6f} | " + f"mass lost={enkf_da.total_overflow_lost_mm:.3f} mm") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', required=True) + parser.add_argument('--forcing-dir', required=True) + parser.add_argument('--obs-dir', required=True) + parser.add_argument('--cfe-dir', required=True) + parser.add_argument('--config-file', required=True) + parser.add_argument('--param-bounds', required=True, + help='Kept for symmetry; not read in this script') + parser.add_argument('--out-dir', required=True) + parser.add_argument('--test-forcing-dir1', required=True) + parser.add_argument('--test-forcing-dir2', required=True) + parser.add_argument('--enkf-members', type=int, default=20) + parser.add_argument('--enkf-obs-error-std', type=float, default=0.05) + parser.add_argument('--no-vrugt-r', action='store_true') + parser.add_argument('--vrugt-alpha', type=float, default=0.10) + parser.add_argument('--vrugt-scale', type=float, default=0.001) + parser.add_argument('--rng-seed', type=int, default=None) + parser.add_argument('--base-step-h', type=int, default=6, + help='Issue-time cadence in hours across the full test ' + 'period (default 6h)') + parser.add_argument('--dense-start', type=str, default=DEFAULT_DENSE_START, + help='Start of hourly-densification window ' + '(default 2024-09-24 00:00:00 — Helene)') + parser.add_argument('--dense-end', type=str, default=DEFAULT_DENSE_END, + help='End of hourly-densification window ' + '(default 2024-09-28 23:00:00 — Helene)') + parser.add_argument('--sanity-check', action='store_true', + help='At the first issue time, fork fcst_models and step ' + 'both prod and fcst one hour with identical unperturbed ' + 'forcing. Confirms the 4-state snapshot is sufficient ' + 'to reproduce next-hour Q to machine precision, then ' + 'exits. Use before kicking off the full sweep.') + parser.add_argument('--hardcoded-r', type=float, default=None, + help='If set, override the Vrugt + kriging-variance R formula ' + 'and use this constant value (mm^2/h^2) at every hour. ' + 'Matches the DualEarth/new_EnKF.py reference (R=0.07). ' + 'Suppresses DA at low flow where kriging is noisy while ' + 'preserving DA strength at storm peaks. Applied to all ' + 'three EnKF instances (da, openloop, fcst) for consistency.') + parser.add_argument('--prod-script', default=None, + help='Path to calibrate_catchment_cfe_da_v2.py for ' + 'importing EnKFAssimilator. Defaults to next-to-this-file.') + args = parser.parse_args() + + here = os.path.dirname(os.path.abspath(__file__)) + prod_script = args.prod_script or os.path.join(here, 'calibrate_catchment_cfe_da_v2.py') + if not os.path.exists(prod_script): + raise FileNotFoundError( + f"Could not find production script at {prod_script}. " + f"Pass --prod-script to override.") + EnKFAssimilator = import_enkf_class(prod_script) + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + + run(args, _bmi_cfe, EnKFAssimilator) + + +if __name__ == '__main__': + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/2_assimilation/run_perturbation_da_on.py b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/2_assimilation/run_perturbation_da_on.py new file mode 100644 index 00000000..b9ef7655 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/2_assimilation/run_perturbation_da_on.py @@ -0,0 +1,392 @@ +""" +run_perturbation_da_on.py — 2a/2b: ensemble spread with DA ON. + +Runs two separated perturbation arms for the Helene window (Sep 24-30), +with DA actively running during the analysis period: + + ARM A — forcing only (30 members): + DA analyzes states using Qkrig at every hour (DA ON). + At each issue_time: fork 30 members with ONLY met forcing perturbed + (different precip/PET draw per member). Hydro states = DA analysis + mean (same for all 30). Free-run 18 hours, no further DA. + + ARM B — hydro states only (20 members): + Same DA analysis. At each issue_time: fork 20 members with ONLY + hydro states perturbed (multiplicative noise on soil/GW/Nash from + the DA analysis ensemble). Forcing = deterministic. Free-run 18 hours. + +These two arms isolate "where does forecast spread come from WITH DA?" +— figures for the assimilation script section. + +Output (per catchment, per arm): + //_da_forcing_arm.csv + //_da_hydro_arm.csv + Columns: issue_time, lead_hour, member_00..member_29 (forcing arm) + issue_time, lead_hour, member_00..member_19 (hydro arm) + +Usage (run for each catchment, both arms): + python3 run_perturbation_da_on.py \\ + --cat-id cat-1016300 \\ + --forcing-dir \\ + --obs-dir /mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance \\ + --cfe-dir /mnt/disk2/suma_helen_poster/cfe_py \\ + --config-file /mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json \\ + --param-bounds /mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json \\ + --out-dir /mnt/disk2/suma_helen_poster/da_results/v2_perturbation_da_on \\ + --test-forcing-dir1 \\ + --test-forcing-dir2 \\ + --hardcoded-r 0.07 +""" + +import argparse +import os +import sys +import json +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +N_FORCING = 30 # members in forcing arm +N_HYDRO = 20 # members in hydro-state arm + +# Analysis period — run DA through entire test period +TEST_START = "2024-08-24 00:00:00" # 1 month spinup before Helene window +TEST_END = "2024-10-31 23:00:00" + +# Issue times for the Helene window (hourly, Sep 24-30) +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-30 06:00:00") +FORECAST_HOURS = 18 + +# Perturbation magnitudes (same as production) +PRECIP_SIGMA = 0.15 +PET_SIGMA = 0.10 +STATE_FRAC = 0.05 # hydro-state perturbation fraction + +# Set at runtime +CAT_ID = None +OBS_FILE = None +CFE_CONFIG_FILE = None +TEST_FORCING_FILE = None +OUT_DIR = None +HARDCODED_R = None +bmi_cfe = None + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(): + df = pd.read_csv(TEST_FORCING_FILE) + df["date"] = pd.to_datetime(df["time"]).dt.strftime("%Y-%m-%d %H:%M:%S") + df["total_precipitation"] = df["APCP_surface"] * 3600.0 + df["potential_evaporation"] = priestley_taylor_pet( + df["DSWRF_surface"].values, df["TMP_2maboveground"].values) + return df + + +def load_obs(): + df = pd.read_csv(OBS_FILE) + date_col = next((c for c in df.columns + if c.lower() in ("date", "time", "datetime", "timestamp")), None) + if date_col is None: + raise ValueError(f"No date column in {OBS_FILE}. Columns: {df.columns.tolist()}") + q_col = next((c for c in df.columns if "qkrig" in c.lower() and "var" not in c.lower()), None) + if q_col is None: + raise ValueError(f"No qkrig column in {OBS_FILE}. Columns: {df.columns.tolist()}") + var_col = next((c for c in df.columns if "var" in c.lower()), None) + df[date_col] = pd.to_datetime(df[date_col]) + df = df.sort_values(date_col).reset_index(drop=True) + date_strs = df[date_col].dt.strftime("%Y-%m-%d %H:%M:%S") + obs_dict = dict(zip(date_strs, df[q_col])) + var_vals = df[var_col] if var_col else pd.Series([0.07] * len(df)) + var_dict = dict(zip(date_strs, var_vals)) + print(f" Obs: date='{date_col}', q='{q_col}', var='{var_col}'") + return obs_dict, var_dict + + +def make_model(best_params): + with open(CFE_CONFIG_FILE) as f: + cfg = json.load(f) + cfg["forcing_file"] = TEST_FORCING_FILE + cfg["soil_params"]["bb"] = best_params["bb"] + cfg["soil_params"]["smcmax"] = best_params["smcmax"] + cfg["soil_params"]["satdk"] = best_params["satdk"] + cfg["slop"] = best_params["slop"] + cfg["max_gw_storage"] = best_params["max_gw_storage"] + cfg["expon"] = best_params["expon"] + cfg["Cgw"] = best_params["Cgw"] + cfg["K_lf"] = best_params["K_lf"] + cfg["K_nash"] = best_params["K_nash"] + cfg["partition_scheme"] = ("Schaake" if best_params["scheme"] <= 0.5 + else "Xinanjiang") + tmp_cfg = str(OUT_DIR / f"{CAT_ID}_bmi_config_temp_da_on.json") + with open(tmp_cfg, "w") as f: + json.dump(cfg, f) + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.initialize() + return m + + +def get_state(m): + return { + "soil_m": m.soil_reservoir["storage_m"], + "gw_m": m.gw_reservoir["storage_m"], + "nash0": float(m.nash_storage[0]), + "nash1": float(m.nash_storage[1]), + } + + +def set_state(m, state): + m.soil_reservoir["storage_m"] = state["soil_m"] + m.gw_reservoir["storage_m"] = state["gw_m"] + m.nash_storage[0] = state["nash0"] + m.nash_storage[1] = state["nash1"] + + +def perturb_state(state, rng, frac=STATE_FRAC): + sm_max = 1.0 # placeholder — will be clipped in set_state; actual max enforced by model + out = { + "soil_m": max(state["soil_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "nash0": max(state["nash0"] * (1.0 + frac * rng.standard_normal()), 0.0), + "nash1": max(state["nash1"] * (1.0 + frac * rng.standard_normal()), 0.0), + } + return out + + +def run_model_step(m, precip_mmh, pet_mmh): + m.set_value("atmosphere_water__time_integral_of_precipitation_mass_flux", + precip_mmh / 1000.0) + m.set_value("water_potential_evaporation_flux", + pet_mmh / 1000.0 / 3600.0) + m.update() + return m.get_value("land_surface_water__runoff_depth") * 1000.0 # m/h -> mm/h + + +def enkf_update(state, q_sim, y_obs, krig_var, rng, n_members=1): + """Simple scalar EnKF update on the mean analysis state.""" + if HARDCODED_R is not None: + R = HARDCODED_R + else: + alpha = 0.10 + R = (alpha * max(y_obs, 0.0)) ** 2 + 0.001 * krig_var + R = max(R, 1e-6) + P_yy = max(q_sim * 0.01, 1e-6) # rough prior variance on Q + K = P_yy / (P_yy + R) + innovation = y_obs - q_sim + # Apply correction to each state proportionally (simple scalar gain) + scale = K * innovation / max(abs(q_sim), 1e-6) + updated = { + "soil_m": max(state["soil_m"] * (1.0 + scale), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + scale), 1e-6), + "nash0": max(state["nash0"] + state["nash0"] * scale, 0.0), + "nash1": max(state["nash1"] + state["nash1"] * scale, 0.0), + } + return updated + + +def main(): + global CAT_ID, OBS_FILE, CFE_CONFIG_FILE, TEST_FORCING_FILE + global OUT_DIR, HARDCODED_R, bmi_cfe + + parser = argparse.ArgumentParser() + parser.add_argument("--cat-id", required=True) + parser.add_argument("--forcing-dir", required=True) + parser.add_argument("--obs-dir", required=True) + parser.add_argument("--cfe-dir", required=True) + parser.add_argument("--config-file", required=True) + parser.add_argument("--param-bounds", required=True) + parser.add_argument("--out-dir", required=True) + parser.add_argument("--test-forcing-dir1", default=None) + parser.add_argument("--test-forcing-dir2", default=None) + parser.add_argument("--hardcoded-r", type=float, default=0.07) + args = parser.parse_args() + + CAT_ID = args.cat_id + OBS_FILE = os.path.join(args.obs_dir, f"{CAT_ID}.csv") + CFE_CONFIG_FILE = args.config_file + HARDCODED_R = args.hardcoded_r + OUT_DIR = Path(args.out_dir) / CAT_ID + OUT_DIR.mkdir(parents=True, exist_ok=True) + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + bmi_cfe = _bmi_cfe + + # Build combined test forcing + if args.test_forcing_dir1 and args.test_forcing_dir2: + f1 = os.path.join(args.test_forcing_dir1, f"{CAT_ID}.csv") + f2 = os.path.join(args.test_forcing_dir2, f"{CAT_ID}.csv") + df1, df2 = pd.read_csv(f1), pd.read_csv(f2) + combined = (pd.concat([df1, df2], ignore_index=True) + .drop_duplicates(subset="time").sort_values("time")) + TEST_FORCING_FILE = str(OUT_DIR / f"{CAT_ID}_nwm_operational_combined.csv") + combined.to_csv(TEST_FORCING_FILE, index=False) + else: + print("Need --test-forcing-dir1 and --test-forcing-dir2"); return + + best_params_file = OUT_DIR / f"{CAT_ID}_best_params.json" + if not best_params_file.exists(): + print(f"Missing best params at {best_params_file}"); return + with open(best_params_file) as f: + best_params = json.load(f)["best_parameters"] + + obs_dict, var_dict = load_obs() + df_forcing = load_test_forcing() + t_mask = (df_forcing["date"] >= TEST_START) & (df_forcing["date"] <= TEST_END) + df_test = df_forcing[t_mask].reset_index(drop=True) + + rng = np.random.default_rng(hash(CAT_ID) & 0x7fffffff) + + # ------------------------------------------------------------------ # + # Phase 1: Run DA through test period, snapshot analysis state at # + # each Helene issue time # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Phase 1: DA analysis through test period...") + model = make_model(best_params) + + helene_snapshots = {} # issue_time -> analysis state dict + + for h, row in df_test.iterrows(): + date_str = row["date"] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + q_sim = run_model_step(model, P, E) + + date_ts = pd.Timestamp(date_str) + + # DA update if obs available + y_obs = obs_dict.get(date_str, np.nan) + if not np.isnan(y_obs): + krig_var = var_dict.get(date_str, 0.07) + state = get_state(model) + updated = enkf_update(state, q_sim, y_obs, krig_var, rng) + set_state(model, updated) + + # Snapshot if this is a Helene issue time + if HELENE_START <= date_ts <= HELENE_END: + helene_snapshots[date_ts] = get_state(model) + + model.finalize() + print(f" Snapshots saved: {len(helene_snapshots)} issue times") + + # ------------------------------------------------------------------ # + # Phase 2A: Forcing arm — 30 members, fixed analyzed state, varying # + # met forcing # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Phase 2A: Forcing arm ({N_FORCING} members)...") + forcing_records = [] + + issue_times = sorted(helene_snapshots.keys()) + t0_to_idx = {t: i for i, t in enumerate(df_test["date"].apply(pd.Timestamp))} + + for t0 in issue_times: + snap = helene_snapshots[t0] + t0_str = t0.strftime("%Y-%m-%d %H:%M:%S") + + # Find index in df_test + if t0 not in t0_to_idx: + continue + start_idx = t0_to_idx[t0] + + q_matrix = np.full((FORECAST_HOURS, N_FORCING), np.nan) + + for mem in range(N_FORCING): + m = make_model(best_params) + set_state(m, snap) # all members start from same DA analysis state + + for lead in range(1, FORECAST_HOURS + 1): + fi = start_idx + lead + if fi >= len(df_test): + break + row = df_test.iloc[fi] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + # Perturb forcing only + mu_p = -0.5 * PRECIP_SIGMA ** 2 + P_pert = float(P * rng.lognormal(mu_p, PRECIP_SIGMA)) if P > 0 else 0.0 + E_pert = max(float(E * (1.0 + PET_SIGMA * rng.standard_normal())), 0.0) + q = run_model_step(m, P_pert, E_pert) + q_matrix[lead - 1, mem] = q + + m.finalize() + + row_data = {"issue_time": t0_str} + for lead in range(1, FORECAST_HOURS + 1): + row_data["lead_hour"] = lead + for mem in range(N_FORCING): + row_data[f"member_{mem:02d}"] = q_matrix[lead - 1, mem] + forcing_records.append(dict(row_data)) + + df_forcing_arm = pd.DataFrame(forcing_records) + out_a = OUT_DIR / f"{CAT_ID}_da_forcing_arm.csv" + df_forcing_arm.to_csv(out_a, index=False) + print(f" Saved: {out_a}") + + # ------------------------------------------------------------------ # + # Phase 2B: Hydro-state arm — 20 members, deterministic forcing, # + # perturbed analysis states # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Phase 2B: Hydro-state arm ({N_HYDRO} members)...") + hydro_records = [] + + for t0 in issue_times: + snap = helene_snapshots[t0] + t0_str = t0.strftime("%Y-%m-%d %H:%M:%S") + + if t0 not in t0_to_idx: + continue + start_idx = t0_to_idx[t0] + + q_matrix = np.full((FORECAST_HOURS, N_HYDRO), np.nan) + + for mem in range(N_HYDRO): + # Perturb the DA analysis state for this member + if mem == 0: + state_m = snap # member 0 = unperturbed analysis mean + else: + state_m = perturb_state(snap, rng) + + m = make_model(best_params) + set_state(m, state_m) + + for lead in range(1, FORECAST_HOURS + 1): + fi = start_idx + lead + if fi >= len(df_test): + break + row = df_test.iloc[fi] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + # Deterministic forcing + q = run_model_step(m, P, E) + q_matrix[lead - 1, mem] = q + + m.finalize() + + row_data = {"issue_time": t0_str} + for lead in range(1, FORECAST_HOURS + 1): + row_data["lead_hour"] = lead + for mem in range(N_HYDRO): + row_data[f"member_{mem:02d}"] = q_matrix[lead - 1, mem] + hydro_records.append(dict(row_data)) + + df_hydro_arm = pd.DataFrame(hydro_records) + out_b = OUT_DIR / f"{CAT_ID}_da_hydro_arm.csv" + df_hydro_arm.to_csv(out_b, index=False) + print(f" Saved: {out_b}") + print(f"[{CAT_ID}] Done.") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/2_assimilation/run_perturbation_sensitivity.py b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/2_assimilation/run_perturbation_sensitivity.py new file mode 100644 index 00000000..57b8d185 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/2_assimilation/run_perturbation_sensitivity.py @@ -0,0 +1,293 @@ +""" +Perturbation-source sensitivity analysis for the multi-catchment CFE DA experiment. + +Ensemble-spread attribution: run 20 CFE members through the test +period with ONLY ONE perturbation source active at a time, no DA, no spinup. + +Three sub-experiments (selected via --source): + init : perturb initial states once at t=0 (multiplicative ~5%) + forcing : perturb hourly precip (lognormal sigma=0.15) and PET (Gaussian sigma=0.10) + process : apply per-hour process noise on states (soil 0.2%, GW 0.15%, Nash 0.5%) + +Each sub-experiment runs N=20 CFE members. DA is always OFF. Spinup is skipped — each +member starts from CFE's default initial state (modified by the perturbation if --source +init). This is the cleanest setup for attributing test-period spread to its source. + +Output: + //_sensitivity_.csv + Columns: date, member_00, member_01, ..., member_19 + +Each member column is that member's hourly Q_sim (mm/h) over the test period. + +Reads the same calibrated best_params.json as the production script. Does not need +the kriging obs file (no DA), but takes the same --obs-dir argument so it can also +log Q_obs alongside for plotting reference. + +Usage: + python3 run_perturbation_sensitivity.py \ + --cat-id cat-1016300 \ + --source forcing \ + --forcing-dir \ + --obs-dir \ + --cfe-dir \ + --config-file \ + --param-bounds \ + --out-dir \ + --test-forcing-dir1 \ + --test-forcing-dir2 +""" + +import argparse +import os +import sys +import json +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 +N_MEMBERS = 20 + +# We only run the test window for this sensitivity analysis. +TEST_START = "2023-10-01 00:00:00" +TEST_END = "2024-10-31 23:00:00" + +# Perturbation magnitudes — same as production defaults +INIT_STATE_FRAC = 0.05 +PRECIP_PERTURB = 0.15 # lognormal sigma +PET_PERTURB = 0.10 # Gaussian sigma +SOIL_PROC_NOISE = 0.002 +GW_PROC_NOISE = 0.0015 +NASH_PROC_NOISE = 0.005 + +# Set at runtime +CAT_ID = None +FORCING_FILE = None +TEST_FORCING_FILE = None +OBS_FILE = None +CFE_CONFIG_FILE = None +OUT_DIR = None +bmi_cfe = None + + +# ---------- Forcing helpers (same as production) ---------- +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(): + df = pd.read_csv(TEST_FORCING_FILE) + df['date'] = pd.to_datetime(df['time']).dt.strftime('%Y-%m-%d %H:%M:%S') + df['total_precipitation'] = df['APCP_surface'] * 3600.0 # kg/m^2/s -> mm/h + df['potential_evaporation'] = priestley_taylor_pet( + df['DSWRF_surface'].values, + df['TMP_2maboveground'].values, + ) + return df + + +# ---------- Perturbation helpers ---------- +def perturb_precip_lognormal(P, rng, sigma=PRECIP_PERTURB): + """Lognormal multiplier with mean 1; preserves zero precip.""" + mu = -0.5 * sigma ** 2 + return float(P * rng.lognormal(mu, sigma)) + + +def perturb_pet_gaussian(E, rng, sigma=PET_PERTURB): + """Gaussian multiplicative noise on PET, clipped at 0.""" + return max(float(E * (1.0 + sigma * rng.standard_normal())), 0.0) + + +def apply_init_state_perturb(m, rng, frac=INIT_STATE_FRAC): + """One-time multiplicative perturbation on the four states (called per member at t=0).""" + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1.0 + frac * rng.standard_normal()), 1e-6, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1.0 + frac * rng.standard_normal()), 1e-6, gw_max)) + # Nash buckets often start at 0; small additive jitter so they're not identical. + jitter = max(sm_max * 1e-6, 1e-9) + m.nash_storage[0] = max(n00 + jitter * rng.standard_normal(), 0.0) + m.nash_storage[1] = max(n10 + jitter * rng.standard_normal(), 0.0) + + +def apply_process_noise(m, rng): + """Per-timestep process noise on all 4 states for one member.""" + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + nash_floor = max(sm_max * 1e-4, 1e-7) + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1.0 + SOIL_PROC_NOISE * rng.standard_normal()), 0.0, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1.0 + GW_PROC_NOISE * rng.standard_normal()), 0.0, gw_max)) + m.nash_storage[0] = max( + n00 * (1.0 + NASH_PROC_NOISE * rng.standard_normal()) + + nash_floor * rng.standard_normal(), 0.0) + m.nash_storage[1] = max( + n10 * (1.0 + NASH_PROC_NOISE * rng.standard_normal()) + + nash_floor * rng.standard_normal(), 0.0) + + +# ---------- Main per-source run ---------- +def run_sensitivity(source, best_param_dict): + """Run N members through the test period with only `source` perturbation active.""" + + def custom_load_forcing(self_cfe): + df = load_test_forcing() + self_cfe.forcing_data = df.rename(columns={"date": "time"}) + + # Build a per-catchment temporary CFE config with the calibrated params + with open(CFE_CONFIG_FILE) as f: + cfg = json.load(f) + cfg['forcing_file'] = TEST_FORCING_FILE + cfg['soil_params']['bb'] = best_param_dict['bb'] + cfg['soil_params']['smcmax'] = best_param_dict['smcmax'] + cfg['soil_params']['satdk'] = best_param_dict['satdk'] + cfg['slop'] = best_param_dict['slop'] + cfg['max_gw_storage'] = best_param_dict['max_gw_storage'] + cfg['expon'] = best_param_dict['expon'] + cfg['Cgw'] = best_param_dict['Cgw'] + cfg['K_lf'] = best_param_dict['K_lf'] + cfg['K_nash'] = best_param_dict['K_nash'] + cfg['partition_scheme'] = "Schaake" if best_param_dict['scheme'] <= 0.5 else "Xinanjiang" + + tmp_cfg = str(OUT_DIR / f'{CAT_ID}_bmi_config_temp_sensitivity_{source}.json') + with open(tmp_cfg, 'w') as f: + json.dump(cfg, f) + + # Deterministic seed per (catchment, source) so each run is reproducible + seed = hash((CAT_ID, source)) & 0x7fffffff + rng = np.random.default_rng(seed) + + print(f"[sensitivity] {CAT_ID} | source={source} | N={N_MEMBERS} | seed={seed}") + + # Build N members. Initial-state perturbation only happens when source == 'init'. + models = [] + for i in range(N_MEMBERS): + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.load_forcing_file = custom_load_forcing.__get__(m) + m.initialize() + if source == 'init' and i > 0: + apply_init_state_perturb(m, rng) + models.append(m) + + # Skip spinup. Go directly to the test period. + df = load_test_forcing() + t_mask = (df['date'] >= TEST_START) & (df['date'] <= TEST_END) + df_test = df[t_mask] + + # Pre-allocate per-member Q arrays + n_hours = len(df_test) + q_matrix = np.full((n_hours, N_MEMBERS), np.nan, dtype=float) + dates_out = df_test['date'].values + + for h, (p, e) in enumerate(zip(df_test['total_precipitation'], + df_test['potential_evaporation'])): + for i, m in enumerate(models): + if source == 'forcing': + p_i = perturb_precip_lognormal(p, rng) + e_i = perturb_pet_gaussian(e, rng) + else: + p_i = float(p) + e_i = float(e) + m.set_value('atmosphere_water__time_integral_of_precipitation_mass_flux', p_i / 1000) + m.set_value('water_potential_evaporation_flux', e_i / 1000 / 3600) + m.update() + q_matrix[h, i] = m.get_value('land_surface_water__runoff_depth') * 1000 # m/h -> mm/h + + # Process noise applied after each hour, but only when source == 'process' + if source == 'process': + for m in models: + apply_process_noise(m, rng) + + for m in models: + m.finalize() + + # Write per-member CSV: date + 20 columns + out = {'date': dates_out} + for i in range(N_MEMBERS): + out[f'member_{i:02d}'] = q_matrix[:, i] + df_out = pd.DataFrame(out) + out_path = OUT_DIR / f'{CAT_ID}_sensitivity_{source}.csv' + df_out.to_csv(out_path, index=False) + print(f"[sensitivity] saved {out_path}") + print(f"[sensitivity] mean ensemble spread (std across members) over test period: " + f"{float(q_matrix.std(axis=1).mean()):.5f} mm/h") + + +def main(): + global CAT_ID, FORCING_FILE, TEST_FORCING_FILE, OBS_FILE, CFE_CONFIG_FILE + global OUT_DIR, bmi_cfe + + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', required=True) + parser.add_argument('--source', required=True, + choices=['init', 'forcing', 'process'], + help="Which perturbation source to isolate") + parser.add_argument('--forcing-dir', required=True) + parser.add_argument('--obs-dir', required=True, + help='Kept for symmetry with production script; not used since DA is off') + parser.add_argument('--cfe-dir', required=True) + parser.add_argument('--config-file', required=True) + parser.add_argument('--param-bounds', required=True, + help='Kept for symmetry; not read in this script') + parser.add_argument('--out-dir', required=True) + parser.add_argument('--test-forcing-dir1', default=None) + parser.add_argument('--test-forcing-dir2', default=None) + args = parser.parse_args() + + CAT_ID = args.cat_id + FORCING_FILE = os.path.join(args.forcing_dir, f'{CAT_ID}.csv') + OBS_FILE = os.path.join(args.obs_dir, f'{CAT_ID}.csv') + CFE_CONFIG_FILE = args.config_file + OUT_DIR = Path(args.out_dir) / CAT_ID + OUT_DIR.mkdir(parents=True, exist_ok=True) + + # Build combined test forcing CSV (same idiom as production) + if args.test_forcing_dir1 and args.test_forcing_dir2: + f1 = os.path.join(args.test_forcing_dir1, f'{CAT_ID}.csv') + f2 = os.path.join(args.test_forcing_dir2, f'{CAT_ID}.csv') + if os.path.exists(f1) and os.path.exists(f2): + df1 = pd.read_csv(f1) + df2 = pd.read_csv(f2) + combined = pd.concat([df1, df2], ignore_index=True) + combined = combined.drop_duplicates(subset='time').sort_values('time') + TEST_FORCING_FILE = str(OUT_DIR / f'{CAT_ID}_nwm_operational_combined.csv') + combined.to_csv(TEST_FORCING_FILE, index=False) + else: + print(f"Warning: test forcing files not found for {CAT_ID}, exiting") + return + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + bmi_cfe = _bmi_cfe + + best_params_file = OUT_DIR / f'{CAT_ID}_best_params.json' + if not best_params_file.exists(): + print(f"No best params file at {best_params_file}. Pre-stage from Run 3.") + return + with open(best_params_file) as f: + saved = json.load(f) + best_param_dict = saved['best_parameters'] + + run_sensitivity(args.source, best_param_dict) + + +if __name__ == '__main__': + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/2_assimilation/run_production_per_member.py b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/2_assimilation/run_production_per_member.py new file mode 100644 index 00000000..69424666 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/2_assimilation/run_production_per_member.py @@ -0,0 +1,295 @@ +""" +Run the production DA pipeline for ONE catchment and save per-member streamflow. + +Identical to calibrate_catchment_cfe_da_v2.py in algorithm (init perturbation + +forcing perturbation + process noise + true EnKF with Vrugt R), but writes one +column per ensemble member instead of only the ensemble mean. This is for +the per-member factor-decomposition plot. + +To keep the comparison apples-to-apples with the sensitivity runs (which skip +spinup), this script also SKIPS the spinup loop. Each member starts directly +from CFE's default initial state, plus the t=0 init perturbation. + +Output: + //_production_per_member.csv + Columns: date, member_00, member_01, ..., member_19 + +Imports EnKFAssimilator from calibrate_catchment_cfe_da_v2 to guarantee the +DA math matches production exactly. + +Usage: + python3 run_production_per_member.py \ + --cat-id cat-1016300 \ + --forcing-dir /mnt/.../nwm_retro_catchment_forcings \ + --obs-dir /mnt/.../catchment_ts_03463300_with_variance \ + --cfe-dir /mnt/.../cfe_py \ + --config-file /mnt/.../cat_03463300_bmi_config_cfe.json \ + --param-bounds /mnt/.../CFE_parameter_bounds.json \ + --out-dir /mnt/.../v2_production_per_member \ + --test-forcing-dir1 ... \ + --test-forcing-dir2 ... \ + --enkf-members 20 +""" + +import argparse +import os +import sys +import json +import importlib.util +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +TEST_START = "2023-10-01 00:00:00" +TEST_END = "2024-10-31 23:00:00" + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(test_forcing_file): + df = pd.read_csv(test_forcing_file) + df['date'] = pd.to_datetime(df['time']).dt.strftime('%Y-%m-%d %H:%M:%S') + df['total_precipitation'] = df['APCP_surface'] * 3600.0 + df['potential_evaporation'] = priestley_taylor_pet( + df['DSWRF_surface'].values, df['TMP_2maboveground'].values + ) + return df + + +def import_enkf_class(script_path): + """Load EnKFAssimilator from the production script without triggering its main().""" + spec = importlib.util.spec_from_file_location("prod_v2", script_path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod.EnKFAssimilator + + +def run(args, bmi_cfe, EnKFAssimilator): + cat_id = args.cat_id + out_dir = Path(args.out_dir) / cat_id + out_dir.mkdir(parents=True, exist_ok=True) + + # Combined test forcing + f1 = os.path.join(args.test_forcing_dir1, f'{cat_id}.csv') + f2 = os.path.join(args.test_forcing_dir2, f'{cat_id}.csv') + df1 = pd.read_csv(f1) + df2 = pd.read_csv(f2) + combined = pd.concat([df1, df2], ignore_index=True) + combined = combined.drop_duplicates(subset='time').sort_values('time') + test_forcing_file = str(out_dir / f'{cat_id}_nwm_operational_combined.csv') + combined.to_csv(test_forcing_file, index=False) + + # Load best_params + best_params_file = out_dir / f'{cat_id}_best_params.json' + if not best_params_file.exists(): + print(f"Missing {best_params_file}. Pre-stage from Run 3.") + return + with open(best_params_file) as f: + saved = json.load(f) + best = saved['best_parameters'] + + # Build temp CFE config with calibrated params + with open(args.config_file) as f: + cfg = json.load(f) + cfg['forcing_file'] = test_forcing_file + cfg['soil_params']['bb'] = best['bb'] + cfg['soil_params']['smcmax'] = best['smcmax'] + cfg['soil_params']['satdk'] = best['satdk'] + cfg['slop'] = best['slop'] + cfg['max_gw_storage'] = best['max_gw_storage'] + cfg['expon'] = best['expon'] + cfg['Cgw'] = best['Cgw'] + cfg['K_lf'] = best['K_lf'] + cfg['K_nash'] = best['K_nash'] + cfg['partition_scheme'] = "Schaake" if best['scheme'] <= 0.5 else "Xinanjiang" + tmp_cfg = str(out_dir / f'{cat_id}_bmi_config_temp_per_member.json') + with open(tmp_cfg, 'w') as f: + json.dump(cfg, f) + + # Build the EnKF assimilator (same defaults as production) + obs_file = os.path.join(args.obs_dir, f'{cat_id}.csv') + enkf = EnKFAssimilator( + n_members=args.enkf_members, + obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, + use_vrugt_r=(not args.no_vrugt_r), + vrugt_alpha=args.vrugt_alpha, + vrugt_scale=args.vrugt_scale, + rng_seed=args.rng_seed, + ) + N = enkf.n_members + print(f"[per-member] {cat_id} | N={N} | use_vrugt_r={enkf.use_vrugt_r}") + + # Custom forcing loader for CFE BMI + def custom_load_forcing(self_cfe): + df = load_test_forcing(test_forcing_file) + self_cfe.forcing_data = df.rename(columns={"date": "time"}) + + # Build N CFE instances; perturb initial states (member 0 left clean) + models = [] + init_states_per_member = [] # capture initial-state snapshot after perturbation + for i in range(N): + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.load_forcing_file = custom_load_forcing.__get__(m) + m.initialize() + if i > 0: + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + r = enkf.init_state_perturb_frac + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1 + r * enkf.rng.standard_normal()), 1e-6, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1 + r * enkf.rng.standard_normal()), 1e-6, gw_max)) + jitter = max(sm_max * 1e-6, 1e-9) + m.nash_storage[0] = max(n00 + jitter * enkf.rng.standard_normal(), 0.0) + m.nash_storage[1] = max(n10 + jitter * enkf.rng.standard_normal(), 0.0) + # Snapshot the (possibly-perturbed) initial state for this member + init_states_per_member.append({ + "soil_m": float(m.soil_reservoir["storage_m"]), + "gw_m": float(m.gw_reservoir["storage_m"]), + "nash0_m": float(m.nash_storage[0]), + "nash1_m": float(m.nash_storage[1]), + "soil_max_m": float(m.soil_reservoir["storage_max_m"]), + "gw_max_m": float(m.gw_reservoir["storage_max_m"]), + }) + models.append(m) + + # SKIP SPINUP — go directly to the test period (matches sensitivity-script setup) + df = load_test_forcing(test_forcing_file) + t_mask = (df['date'] >= TEST_START) & (df['date'] <= TEST_END) + df_test = df[t_mask] + n_hours = len(df_test) + q_matrix = np.full((n_hours, N), np.nan, dtype=float) + precip_matrix = np.full((n_hours, N), np.nan, dtype=float) # mm/h per member + pet_matrix = np.full((n_hours, N), np.nan, dtype=float) # mm/h per member + dates_out = df_test['date'].values + + for h, (p, e, current_date) in enumerate(zip( + df_test['total_precipitation'], + df_test['potential_evaporation'], + df_test['date'])): + + # Perturb forcing per member + p_arr, e_arr = enkf.perturb_forcing(p, e) + # Capture the per-member perturbed forcing for later plotting/diagnostics + precip_matrix[h, :] = p_arr + pet_matrix[h, :] = e_arr + + # Advance each member one hour + ensemble_q = np.empty(N, dtype=float) + for i, m in enumerate(models): + m.set_value('atmosphere_water__time_integral_of_precipitation_mass_flux', + float(p_arr[i]) / 1000) + m.set_value('water_potential_evaporation_flux', + float(e_arr[i]) / 1000 / 3600) + m.update() + ensemble_q[i] = m.get_value('land_surface_water__runoff_depth') * 1000 # mm/h + + # Record pre-analysis forecast per member (this is the "actual production" value) + q_matrix[h, :] = ensemble_q + + # DA step (uses production EnKFAssimilator) + enkf.update_states(models, current_date, ensemble_q) + + # Process noise after DA + enkf.add_process_noise(models) + + for m in models: + m.finalize() + + # ----- Save per-member CSVs ----- + def _save_matrix(matrix, suffix): + out = {'date': dates_out} + for i in range(N): + out[f'member_{i:02d}'] = matrix[:, i] + df_out = pd.DataFrame(out) + out_path = out_dir / f'{cat_id}_production_per_member_{suffix}.csv' + df_out.to_csv(out_path, index=False) + return out_path + + # Streamflow outputs + q_path = out_dir / f'{cat_id}_production_per_member.csv' + out = {'date': dates_out} + for i in range(N): + out[f'member_{i:02d}'] = q_matrix[:, i] + pd.DataFrame(out).to_csv(q_path, index=False) + print(f"[per-member] saved {q_path}") + + # Perturbed forcings (so each member's actual inputs are recoverable) + precip_path = _save_matrix(precip_matrix, "precip") + pet_path = _save_matrix(pet_matrix, "pet") + print(f"[per-member] saved {precip_path}") + print(f"[per-member] saved {pet_path}") + + # Initial states (small, JSON is fine) + init_states_path = out_dir / f'{cat_id}_production_per_member_initial_states.json' + with open(init_states_path, 'w') as f: + json.dump({ + "catchment_id": cat_id, + "n_members": N, + "members": {f"member_{i:02d}": init_states_per_member[i] for i in range(N)}, + }, f, indent=2) + print(f"[per-member] saved {init_states_path}") + + print(f"[per-member] ensemble mean spread over test period: " + f"{float(q_matrix.std(axis=1).mean()):.5f} mm/h") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', required=True) + parser.add_argument('--forcing-dir', required=True) + parser.add_argument('--obs-dir', required=True) + parser.add_argument('--cfe-dir', required=True) + parser.add_argument('--config-file', required=True) + parser.add_argument('--param-bounds', required=True, + help='Kept for symmetry; not read in this script') + parser.add_argument('--out-dir', required=True) + parser.add_argument('--test-forcing-dir1', required=True) + parser.add_argument('--test-forcing-dir2', required=True) + parser.add_argument('--enkf-members', type=int, default=20) + parser.add_argument('--enkf-obs-error-std', type=float, default=0.05) + parser.add_argument('--no-vrugt-r', action='store_true') + parser.add_argument('--vrugt-alpha', type=float, default=0.10) + parser.add_argument('--vrugt-scale', type=float, default=0.001) + parser.add_argument('--rng-seed', type=int, default=None) + parser.add_argument('--prod-script', default=None, + help='Path to calibrate_catchment_cfe_da_v2.py for importing ' + 'EnKFAssimilator. Defaults to the script located next to this file.') + args = parser.parse_args() + + # Locate the production script (default: same dir as this script) + if args.prod_script is None: + here = os.path.dirname(os.path.abspath(__file__)) + prod_script = os.path.join(here, 'calibrate_catchment_cfe_da_v2.py') + else: + prod_script = args.prod_script + if not os.path.exists(prod_script): + raise FileNotFoundError(f"Could not find production script at {prod_script}. " + f"Pass --prod-script to override.") + EnKFAssimilator = import_enkf_class(prod_script) + + # Dynamic CFE import (same idiom as production script) + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + + run(args, _bmi_cfe, EnKFAssimilator) + + +if __name__ == '__main__': + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/3_routing/route_det_f1.sh b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/3_routing/route_det_f1.sh new file mode 100644 index 00000000..400c348f --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/3_routing/route_det_f1.sh @@ -0,0 +1,28 @@ +#!/usr/bin/env bash +# F1 variance-scaled Vrugt (20pct gauge holdout) — deterministic T-route routing. +# Routes per-catchment _test_results.csv (sim_mm_h column) through Muskingum-Cunge +# to gauge 03463300 and writes routed_Q_test.csv + KGE/NSE summary. +# +# Usage: +# bash route_det_f1.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +ROUTE_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/2_troute_routing/run_route.py + +F1_DIR=/mnt/disk2/suma_helen_poster/da_results/v2_perturbation_da_on +GPKG=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +echo "[F1] Deterministic T-route routing..." +echo " da-dir : $F1_DIR" +echo " out-dir: $F1_DIR" + +$TROUTE "$ROUTE_SCRIPT" \ + --gpkg "$GPKG" \ + --da-dir "$F1_DIR" \ + --out-dir "$F1_DIR" \ + --usgs-csv "$USGS_CSV" + +echo "[F1] Deterministic routing done. Output: $F1_DIR/routed_Q_test.csv" diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/3_routing/route_ensemble_f1.sh b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/3_routing/route_ensemble_f1.sh new file mode 100644 index 00000000..3aeefc59 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/3_routing/route_ensemble_f1.sh @@ -0,0 +1,30 @@ +#!/usr/bin/env bash +# F1 variance-scaled Vrugt (20pct gauge holdout) — crossed-ensemble T-route routing. +# Routes all 600 crossed-ensemble members through Muskingum-Cunge +# to gauge 03463300 and writes routed_crossed_ensemble.parquet. +# +# Must be run AFTER crossed ensemble parquets exist in CROSSED_DIR. +# +# Usage: +# bash route_ensemble_f1.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +ROUTE_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/2_troute_routing/run_route_crossed_ensemble.py + +CROSSED_DIR=/mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble_vrugt +GPKG=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +echo "[F1] Crossed-ensemble T-route routing..." +echo " ensemble-dir: $CROSSED_DIR" +echo " out-dir : $CROSSED_DIR" + +$TROUTE "$ROUTE_SCRIPT" \ + --gpkg "$GPKG" \ + --ensemble-dir "$CROSSED_DIR" \ + --out-dir "$CROSSED_DIR" \ + --usgs-csv "$USGS_CSV" + +echo "[F1] Ensemble routing done. Output: $CROSSED_DIR/routed_crossed_ensemble.parquet" diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/3_routing/route_leadtime_f1.sh b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/3_routing/route_leadtime_f1.sh new file mode 100644 index 00000000..f48cef7e --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/3_routing/route_leadtime_f1.sh @@ -0,0 +1,31 @@ +#!/usr/bin/env bash +# F1 variance-scaled Vrugt (20pct gauge holdout) — route lead-time forecast CSVs. +# Routes all 21-catchment lead-time forecast CSVs through Muskingum-Cunge +# to gauge 03463300. Writes routed_leadtime_da_full.parquet and +# routed_leadtime_openloop_full.parquet for gauge-level error decay analysis. +# +# Run AFTER the lead-time sweep batch finishes. +# +# Usage: +# bash route_leadtime_f1.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +ROUTE_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/2_troute_routing/run_route_leadtime_forecasts.py + +F1_DIR=/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast +GPKG=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg +OUT_DIR=/mnt/disk2/suma_helen_poster/leadtime_troute_routing + +echo "[F1] Routing lead-time forecasts through T-route..." +echo " forecast-dir: $F1_DIR" +echo " out-dir : $OUT_DIR" + +$TROUTE "$ROUTE_SCRIPT" \ + --gpkg "$GPKG" \ + --forecast-dir "$F1_DIR" \ + --out-dir "$OUT_DIR" + +echo "[F1] Lead-time routing done. Outputs in: $OUT_DIR" +ls "$OUT_DIR"/*.parquet 2>/dev/null || echo " (no parquets found)" diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4a_error_decay/plot_lead_time_decay.py b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4a_error_decay/plot_lead_time_decay.py new file mode 100644 index 00000000..2b33fa82 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4a_error_decay/plot_lead_time_decay.py @@ -0,0 +1,160 @@ +""" +Catchment-level error-vs-lead-time decay curve + ensemble spread by lead. + +Two-panel figure for ONE catchment: + TOP — RMSE of ensemble-mean forecast at each lead hour (1..18) vs the + catchment's kriging observation. DA solid, open-loop dashed, with + a shaded band showing the min/max of per-member RMSE. + BOTTOM — Mean ensemble spread (std-dev across 20 members, averaged across + all issue times) at each lead hour. Tells you whether forcing + perturbation alone keeps the forecast ensemble diverse during the + 18-hour free-run — useful for inspecting individual members + without needing the full spaghetti view. + +This is the catchment-level analog of the gauge-level decay curve. +Routing to the gauge is a separate post-step (route_lead_time_forecasts.py ++ a gauge-level decay script); this script lets us look at the catchment-level +signal without T-route in the loop. + +Inputs (from run_lead_time_forecast_sweep.py): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv +Obs (from production EnKF run): + //_test_results.csv + +Output: + //_lead_time_decay.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" + +DEFAULT_LEADTIME_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast" +DEFAULT_DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" + +# Module-level globals set by main() so the helpers see them +CAT = "cat-1016300" +LEADTIME_DIR = DEFAULT_LEADTIME_DIR +DA_DIR = DEFAULT_DA_DIR +OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay.png") + + +def load_forecasts(path): + """Return (df, member_cols).""" + if not os.path.exists(path): + raise FileNotFoundError(path) + df = pd.read_csv(path, parse_dates=['issue_time', 'valid_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + return df, member_cols + + +def load_obs(): + obs_path = os.path.join(DA_DIR, CAT, f"{CAT}_test_results.csv") + df = pd.read_csv(obs_path, parse_dates=['date']) + return df.set_index('date')['obs_mm_h'] + + +def metrics_by_lead(df, member_cols, obs_series): + """For each lead hour, return (rmse_mean, rmse_min_member, rmse_max_member, + mean_ensemble_std).""" + df = df.copy() + df['obs'] = df['valid_time'].map(obs_series) + df = df.dropna(subset=['obs']) + df['ens_mean'] = df[member_cols].mean(axis=1) + df['ens_std'] = df[member_cols].std(axis=1) + + leads = sorted(df['lead_hour'].unique()) + rmse_mean = np.full(len(leads), np.nan) + rmse_min_member = np.full(len(leads), np.nan) + rmse_max_member = np.full(len(leads), np.nan) + mean_ens_std = np.full(len(leads), np.nan) + for i, L in enumerate(leads): + sub = df[df['lead_hour'] == L] + if len(sub) == 0: + continue + err = sub['ens_mean'].values - sub['obs'].values + rmse_mean[i] = float(np.sqrt(np.mean(err ** 2))) + per_member_rmse = np.array([ + np.sqrt(np.mean((sub[c].values - sub['obs'].values) ** 2)) + for c in member_cols + ]) + rmse_min_member[i] = float(np.nanmin(per_member_rmse)) + rmse_max_member[i] = float(np.nanmax(per_member_rmse)) + mean_ens_std[i] = float(np.nanmean(sub['ens_std'].values)) + return (np.asarray(leads), rmse_mean, rmse_min_member, + rmse_max_member, mean_ens_std) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', default="cat-1016300") + parser.add_argument('--leadtime-dir', default=DEFAULT_LEADTIME_DIR, + help='Dir holding /_lead_time_forecasts_{da,openloop}.csv') + parser.add_argument('--da-dir', default=DEFAULT_DA_DIR, + help='Dir holding /_test_results.csv (for obs)') + args = parser.parse_args() + + global CAT, LEADTIME_DIR, DA_DIR, OUT_PNG + CAT = args.cat_id + LEADTIME_DIR = args.leadtime_dir + DA_DIR = args.da_dir + OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay.png") + + da_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_da.csv") + ol_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_openloop.csv") + df_da, m_cols = load_forecasts(da_path) + df_ol, _ = load_forecasts(ol_path) + obs_series = load_obs() + + leads_da, rmse_da, rmse_da_lo, rmse_da_hi, std_da = metrics_by_lead( + df_da, m_cols, obs_series) + leads_ol, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = metrics_by_lead( + df_ol, m_cols, obs_series) + + fig, (ax_top, ax_bot) = plt.subplots(2, 1, figsize=(10, 10), sharex=True, + gridspec_kw={"height_ratios": [1.4, 1.0]}) + + # ----- TOP: RMSE decay curve ----- + ax_top.fill_between(leads_da, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2, + label="DA — member spread (min/max RMSE)") + ax_top.fill_between(leads_ol, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2, + label="Open-loop — member spread") + ax_top.plot(leads_da, rmse_da, color=DA_COLOR, lw=2.4, marker='o', + zorder=4, label="DA on (ensemble-mean RMSE)") + ax_top.plot(leads_ol, rmse_ol, color=OL_COLOR, lw=2.4, marker='s', + linestyle='--', zorder=4, label="Open-loop (ensemble-mean RMSE)") + ax_top.set_ylabel("RMSE vs Qkrig obs (mm/h)", fontsize=11) + ax_top.set_title( + f"Forecast lead-time error decay — {CAT}\n" + f"Issue times pooled across test period (Oct 2023 – Oct 2024)", + fontsize=12, + ) + ax_top.grid(True, alpha=0.25) + ax_top.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + # ----- BOTTOM: ensemble spread by lead hour ----- + ax_bot.plot(leads_da, std_da, color=DA_COLOR, lw=2.2, marker='o', + label="DA — mean ensemble std") + ax_bot.plot(leads_ol, std_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', label="Open-loop — mean ensemble std") + ax_bot.set_xlabel("Forecast lead time (hours)", fontsize=11) + ax_bot.set_ylabel("Mean ensemble std (mm/h)", fontsize=11) + ax_bot.set_xticks(np.arange(1, 19)) + ax_bot.grid(True, alpha=0.25) + ax_bot.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + plt.tight_layout() + plt.savefig(OUT_PNG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4a_error_decay/plot_lead_time_decay_by_regime.py b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4a_error_decay/plot_lead_time_decay_by_regime.py new file mode 100644 index 00000000..ecf10408 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4a_error_decay/plot_lead_time_decay_by_regime.py @@ -0,0 +1,200 @@ +""" +Lead-time decay curve, split by flow regime at issue time. + +The pooled lead-time curve (plot_lead_time_decay.py) showed DA losing to +open-loop across most lead hours, but the test period is ~99% low-flow. +This script splits the same forecast CSVs by the flow regime at the issue +time t0, so we can see whether DA helps when it matters (storms / Helene) +and hurts when it doesn't (low flow). + +Three regimes are partitioned on obs(t0) (kriging obs at the issue time): + + helene issue_time ∈ [2024-09-24, 2024-09-28] (the 5-day Helene window) + storm obs(t0) > 0.5 mm/h (high-flow issue times) + low_flow obs(t0) < 0.1 mm/h (low-flow issue times) + +Same 2-panel layout per regime: RMSE vs lead (top), ensemble spread vs lead +(bottom). Three regime columns × 2 metric rows in one figure. + +Inputs (from run_lead_time_forecast_sweep.py — no re-run needed): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv +Obs (from production EnKF): + //_test_results.csv + +Output: + //_lead_time_decay_by_regime.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" + +DEFAULT_LEADTIME_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast" +DEFAULT_DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" + +# Module-level globals set by main() so helpers see them +CAT = "cat-1016300" +LEADTIME_DIR = DEFAULT_LEADTIME_DIR +DA_DIR = DEFAULT_DA_DIR +OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay_by_regime.png") + +# Regime definitions +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-28 23:00:00") +STORM_OBS_THRESHOLD = 0.5 # mm/h +LOWFLOW_OBS_THRESHOLD = 0.1 # mm/h + + +def load_forecasts(path): + if not os.path.exists(path): + raise FileNotFoundError(path) + df = pd.read_csv(path, parse_dates=['issue_time', 'valid_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + return df, member_cols + + +def load_obs(): + obs_path = os.path.join(DA_DIR, CAT, f"{CAT}_test_results.csv") + df = pd.read_csv(obs_path, parse_dates=['date']) + return df.set_index('date')['obs_mm_h'] + + +def regime_mask(issue_times, obs_at_issue, regime): + """Boolean mask over issue_times for the named regime.""" + if regime == "helene": + return (issue_times >= HELENE_START) & (issue_times <= HELENE_END) + if regime == "storm": + return obs_at_issue > STORM_OBS_THRESHOLD + if regime == "low_flow": + return obs_at_issue < LOWFLOW_OBS_THRESHOLD + raise ValueError(regime) + + +def metrics_by_lead(df, member_cols, obs_series, issue_mask): + """Compute (leads, rmse_mean, rmse_min, rmse_max, mean_std) restricted + to issue times where issue_mask is True.""" + df = df.copy() + df = df[df['issue_time'].isin(issue_mask)] + if len(df) == 0: + return None + df['obs'] = df['valid_time'].map(obs_series) + df = df.dropna(subset=['obs']) + if len(df) == 0: + return None + df['ens_mean'] = df[member_cols].mean(axis=1) + df['ens_std'] = df[member_cols].std(axis=1) + + leads = sorted(df['lead_hour'].unique()) + rmse_mean = np.full(len(leads), np.nan) + rmse_lo = np.full(len(leads), np.nan) + rmse_hi = np.full(len(leads), np.nan) + std_mean = np.full(len(leads), np.nan) + for i, L in enumerate(leads): + sub = df[df['lead_hour'] == L] + if len(sub) == 0: + continue + err = sub['ens_mean'].values - sub['obs'].values + rmse_mean[i] = float(np.sqrt(np.mean(err ** 2))) + per_member = np.array([ + np.sqrt(np.mean((sub[c].values - sub['obs'].values) ** 2)) + for c in member_cols + ]) + rmse_lo[i] = float(np.nanmin(per_member)) + rmse_hi[i] = float(np.nanmax(per_member)) + std_mean[i] = float(np.nanmean(sub['ens_std'].values)) + return np.asarray(leads), rmse_mean, rmse_lo, rmse_hi, std_mean + + +def plot_pair(ax_rmse, ax_std, da_metrics, ol_metrics, + regime_label, n_issue_times): + if da_metrics is None or ol_metrics is None: + ax_rmse.set_title(f"{regime_label}\n(no issue times match)", fontsize=11) + ax_rmse.axis('off'); ax_std.axis('off') + return + leads, rmse_da, rmse_da_lo, rmse_da_hi, std_da = da_metrics + _, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = ol_metrics + + # RMSE panel + ax_rmse.fill_between(leads, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2) + ax_rmse.fill_between(leads, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2) + ax_rmse.plot(leads, rmse_da, color=DA_COLOR, lw=2.2, marker='o', + zorder=4, label="DA on") + ax_rmse.plot(leads, rmse_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', zorder=4, label="Open-loop") + ax_rmse.set_title(f"{regime_label}\n({n_issue_times} issue times)", fontsize=11) + ax_rmse.set_ylabel("RMSE (mm/h)", fontsize=10) + ax_rmse.grid(True, alpha=0.25) + ax_rmse.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + # Spread panel + ax_std.plot(leads, std_da, color=DA_COLOR, lw=2.0, marker='o', label="DA") + ax_std.plot(leads, std_ol, color=OL_COLOR, lw=2.0, marker='s', + linestyle='--', label="Open-loop") + ax_std.set_xlabel("Forecast lead time (hours)", fontsize=10) + ax_std.set_ylabel("Mean ensemble std (mm/h)", fontsize=10) + ax_std.set_xticks(np.arange(1, 19)) + ax_std.grid(True, alpha=0.25) + ax_std.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', default="cat-1016300") + parser.add_argument('--leadtime-dir', default=DEFAULT_LEADTIME_DIR, + help='Dir holding /_lead_time_forecasts_{da,openloop}.csv') + parser.add_argument('--da-dir', default=DEFAULT_DA_DIR, + help='Dir holding /_test_results.csv (for obs)') + args = parser.parse_args() + + global CAT, LEADTIME_DIR, DA_DIR, OUT_PNG + CAT = args.cat_id + LEADTIME_DIR = args.leadtime_dir + DA_DIR = args.da_dir + OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay_by_regime.png") + + da_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_da.csv") + ol_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_openloop.csv") + df_da, m_cols = load_forecasts(da_path) + df_ol, _ = load_forecasts(ol_path) + obs_series = load_obs() + + # Unique issue times across the run, with obs at each + issue_times = pd.to_datetime(sorted(df_da['issue_time'].unique())) + obs_at_issue = pd.Series(issue_times, index=issue_times).map(obs_series) + + regimes = [ + ("helene", "Helene window (Sep 24–28 2024)"), + ("storm", f"Storm hours (obs(t0) > {STORM_OBS_THRESHOLD} mm/h)"), + ("low_flow", f"Low flow (obs(t0) < {LOWFLOW_OBS_THRESHOLD} mm/h)"), + ] + + fig, axes = plt.subplots(2, 3, figsize=(18, 10), + gridspec_kw={"height_ratios": [1.4, 1.0]}) + for col, (key, label) in enumerate(regimes): + mask = regime_mask(issue_times, obs_at_issue, key) + kept = issue_times[mask] + kept_set = set(pd.to_datetime(kept)) + da_m = metrics_by_lead(df_da, m_cols, obs_series, kept_set) + ol_m = metrics_by_lead(df_ol, m_cols, obs_series, kept_set) + plot_pair(axes[0, col], axes[1, col], da_m, ol_m, label, len(kept)) + + fig.suptitle( + f"Lead-time error decay by flow regime — {CAT}\n" + f"Same forecast CSVs as the pooled view, partitioned on obs(t0).", + fontsize=13, y=0.995, + ) + plt.tight_layout(rect=(0, 0, 1, 0.97)) + plt.savefig(OUT_PNG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4a_error_decay/plot_lead_time_decay_gauge.py b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4a_error_decay/plot_lead_time_decay_gauge.py new file mode 100644 index 00000000..71ce54fc --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4a_error_decay/plot_lead_time_decay_gauge.py @@ -0,0 +1,311 @@ +""" +Gauge-level lead-time forecast decay curve. + +Routes 21-catchment forecast CSVs through T-route to USGS gauge 03463300 and +plots RMSE vs lead time (DA solid, open-loop dashed) against USGS observed Q. +Same 2-panel layout as the catchment-level plot, plus a regime split for the +Helene window. + +The routing step is done separately by route_lead_time_forecasts.py; +this script just reads the resulting parquets and produces the plot. + +Inputs: + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + Each: one row per (issue_time, lead_hour); member columns hold q_gauge_m3s. + (Long-format with explicit `member` and `q_gauge_m3s` columns also supported.) +Obs: + /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + USGS hourly Q in m³/s at gauge 03463300 (South Toe River near Celo, NC). + +Output: + /lead_time_decay_gauge_pooled.png + /lead_time_decay_gauge_by_regime.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" +OBS_COLOR = "black" + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/leadtime_troute_routing" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +# Watershed area at gauge 03463300 (sum of 21 catchments per GPKG hydrofabric) +WATERSHED_AREA_KM2 = 113.18 +# 1 mm/h depth × 113.18 km² = 113.18e3 m³/h = 31.439 m³/s +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-28 23:00:00") +STORM_OBS_THRESHOLD_M3S = 20.0 # gauge-level storm threshold (was 50 — too high + # for this 113 km² basin; only 2 issue times qualified) +LOWFLOW_OBS_THRESHOLD_M3S = 5.0 # gauge-level low-flow threshold + +USGS_HELENE_PEAK_M3S = 1886.0 # reference Sep 27 14:00 + + +def load_parquet_long(path): + """Load a routed parquet and normalize to long format: + columns = issue_time, lead_hour, member, q_gauge_m3s. + + Handles both wide format (member_00..member_19 columns) and long format + (explicit `member` + `q_gauge_m3s` columns). Auto-detects. + """ + if not os.path.exists(path): + raise FileNotFoundError(path) + df = pd.read_parquet(path) + + # Long format detection + if 'q_gauge_m3s' in df.columns and 'member' in df.columns: + df['issue_time'] = pd.to_datetime(df['issue_time']) + return df[['issue_time', 'lead_hour', 'member', 'q_gauge_m3s']] + + # Some variants might call it 'q_m3s' or similar + q_col_candidates = [c for c in df.columns if c.lower() in + ('q_gauge_m3s', 'q_m3s', 'q_gauge', 'q')] + if 'member' in df.columns and q_col_candidates: + qc = q_col_candidates[0] + df['issue_time'] = pd.to_datetime(df['issue_time']) + out = df[['issue_time', 'lead_hour', 'member', qc]].copy() + out = out.rename(columns={qc: 'q_gauge_m3s'}) + return out + + # Wide format: member_00..member_19 columns + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + if not member_cols: + raise ValueError(f"Couldn't identify member columns in {path}. " + f"Columns present: {df.columns.tolist()}") + df['issue_time'] = pd.to_datetime(df['issue_time']) + keep = ['issue_time', 'lead_hour'] + if 'valid_time' in df.columns: + keep.append('valid_time') + long = df[keep + member_cols].melt( + id_vars=keep, + value_vars=member_cols, + var_name='member', + value_name='q_gauge_m3s', + ) + return long + + +def load_usgs_obs(usgs_csv): + """Load USGS gauge obs and ensure m³/s. + + The CSV at gauge 03463300 has the column `QObs(mm/h)` (catchment-averaged + depth, not gauge discharge in m³/s). If the column name contains 'mm', we + convert mm/h → m³/s by multiplying by the watershed area factor: + m³/s = mm/h × WATERSHED_AREA_KM2 × 1000 / 3600 + Verified: 59.978 mm/h × 113.18 km² = 1886 m³/s (USGS Helene peak). + """ + df = pd.read_csv(usgs_csv) + date_candidates = [c for c in df.columns if c.lower() in + ('datetime', 'date', 'time', 'timestamp')] + q_candidates = [c for c in df.columns if 'q' in c.lower() + or 'flow' in c.lower() or 'discharge' in c.lower()] + if not date_candidates or not q_candidates: + raise ValueError(f"Could not identify date/Q columns in {usgs_csv}. " + f"Columns: {df.columns.tolist()}") + dc, qc = date_candidates[0], q_candidates[0] + df[dc] = pd.to_datetime(df[dc]) + series = df.set_index(dc)[qc].astype(float) + if 'mm' in qc.lower(): + series = series * MM_H_TO_M3_S + print(f" Converted obs from mm/h (column '{qc}') to m³/s " + f"using area = {WATERSHED_AREA_KM2} km² (× {MM_H_TO_M3_S:.4f})") + else: + print(f" Loaded obs as-is from column '{qc}' (assumed m³/s)") + return series + + +def metrics_by_lead(df_long, obs_series, issue_mask=None): + """Compute per-lead-hour: ensemble-mean RMSE, per-member min/max RMSE, + mean ensemble std. df_long must have issue_time, lead_hour, member, q_gauge_m3s. + """ + df = df_long.copy() + if issue_mask is not None: + df = df[df['issue_time'].isin(issue_mask)] + if len(df) == 0: + return None + df['valid_time'] = df['issue_time'] + pd.to_timedelta(df['lead_hour'], unit='h') + df['obs'] = df['valid_time'].map(obs_series) + df = df.dropna(subset=['obs']) + if len(df) == 0: + return None + + # Ensemble mean per (issue_time, lead_hour) + grouped = df.groupby(['issue_time', 'lead_hour']) + ens_mean = grouped['q_gauge_m3s'].mean().reset_index(name='ens_mean') + ens_std = grouped['q_gauge_m3s'].std().reset_index(name='ens_std') + obs_per = grouped['obs'].first().reset_index(name='obs') + panel = ens_mean.merge(ens_std, on=['issue_time', 'lead_hour']).merge( + obs_per, on=['issue_time', 'lead_hour']) + + leads = sorted(panel['lead_hour'].unique()) + rmse_mean = np.full(len(leads), np.nan) + rmse_lo = np.full(len(leads), np.nan) + rmse_hi = np.full(len(leads), np.nan) + std_mean = np.full(len(leads), np.nan) + for i, L in enumerate(leads): + psub = panel[panel['lead_hour'] == L] + if len(psub) == 0: + continue + err = psub['ens_mean'].values - psub['obs'].values + rmse_mean[i] = float(np.sqrt(np.mean(err ** 2))) + # Per-member RMSE — recompute from the underlying long frame + dsub = df[df['lead_hour'] == L] + member_rmses = [] + for m, g in dsub.groupby('member'): + err_m = g['q_gauge_m3s'].values - g['obs'].values + if len(err_m) >= 1: + member_rmses.append(float(np.sqrt(np.mean(err_m ** 2)))) + if member_rmses: + rmse_lo[i] = float(np.nanmin(member_rmses)) + rmse_hi[i] = float(np.nanmax(member_rmses)) + std_mean[i] = float(np.nanmean(psub['ens_std'].values)) + return np.asarray(leads), rmse_mean, rmse_lo, rmse_hi, std_mean + + +def plot_two_panel(out_path, leads_da, da_metrics, leads_ol, ol_metrics, title): + fig, (ax_top, ax_bot) = plt.subplots(2, 1, figsize=(10, 10), sharex=True, + gridspec_kw={"height_ratios": [1.4, 1.0]}) + leads, rmse_da, rmse_da_lo, rmse_da_hi, std_da = da_metrics + _, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = ol_metrics + + ax_top.fill_between(leads, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2, + label="DA — member spread (min/max RMSE)") + ax_top.fill_between(leads, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2, + label="Open-loop — member spread") + ax_top.plot(leads, rmse_da, color=DA_COLOR, lw=2.4, marker='o', + zorder=4, label="DA on (ensemble-mean RMSE)") + ax_top.plot(leads, rmse_ol, color=OL_COLOR, lw=2.4, marker='s', + linestyle='--', zorder=4, label="Open-loop (ensemble-mean RMSE)") + ax_top.set_ylabel("RMSE vs USGS obs (m³/s)", fontsize=11) + ax_top.set_title(title, fontsize=12) + ax_top.grid(True, alpha=0.25) + ax_top.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + ax_bot.plot(leads, std_da, color=DA_COLOR, lw=2.2, marker='o', + label="DA — mean ensemble std") + ax_bot.plot(leads, std_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', label="Open-loop — mean ensemble std") + ax_bot.set_xlabel("Forecast lead time (hours)", fontsize=11) + ax_bot.set_ylabel("Mean ensemble std (m³/s)", fontsize=11) + ax_bot.set_xticks(np.arange(1, 19)) + ax_bot.grid(True, alpha=0.25) + ax_bot.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + plt.tight_layout() + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + +def plot_three_panel_regime(out_path, regime_metrics): + fig, axes = plt.subplots(2, 3, figsize=(18, 10), + gridspec_kw={"height_ratios": [1.4, 1.0]}) + for col, (label, n_issue, da_m, ol_m) in enumerate(regime_metrics): + ax_rmse = axes[0, col] + ax_std = axes[1, col] + if da_m is None or ol_m is None: + ax_rmse.set_title(f"{label}\n(no issue times match)", fontsize=11) + ax_rmse.axis('off'); ax_std.axis('off') + continue + leads, rmse_da, rmse_da_lo, rmse_da_hi, std_da = da_m + _, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = ol_m + ax_rmse.fill_between(leads, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2) + ax_rmse.fill_between(leads, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2) + ax_rmse.plot(leads, rmse_da, color=DA_COLOR, lw=2.2, marker='o', + zorder=4, label="DA on") + ax_rmse.plot(leads, rmse_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', zorder=4, label="Open-loop") + ax_rmse.set_title(f"{label}\n({n_issue} issue times)", fontsize=11) + ax_rmse.set_ylabel("RMSE (m³/s)", fontsize=10) + ax_rmse.grid(True, alpha=0.25) + ax_rmse.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + ax_std.plot(leads, std_da, color=DA_COLOR, lw=2.0, marker='o', label="DA") + ax_std.plot(leads, std_ol, color=OL_COLOR, lw=2.0, marker='s', + linestyle='--', label="Open-loop") + ax_std.set_xlabel("Forecast lead time (hours)", fontsize=10) + ax_std.set_ylabel("Mean ensemble std (m³/s)", fontsize=10) + ax_std.set_xticks(np.arange(1, 19)) + ax_std.grid(True, alpha=0.25) + ax_std.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + fig.suptitle("Gauge-level lead-time error decay by regime — USGS 03463300", + fontsize=13, y=0.995) + plt.tight_layout(rect=(0, 0, 1, 0.97)) + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--route-dir', default=DEFAULT_ROUTE_DIR, + help='Dir holding routed_leadtime_{da,openloop}_full.parquet') + parser.add_argument('--out-dir', default=None, + help='Where to write the output PNGs. Defaults to --route-dir ' + '(which may not be writable if owned by another user — ' + 'pass an explicit path then).') + parser.add_argument('--usgs-csv', default=DEFAULT_USGS_CSV) + parser.add_argument('--da-name', default='routed_leadtime_da_full.parquet') + parser.add_argument('--ol-name', default='routed_leadtime_openloop_full.parquet') + args = parser.parse_args() + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + da_path = os.path.join(args.route_dir, args.da_name) + ol_path = os.path.join(args.route_dir, args.ol_name) + df_da = load_parquet_long(da_path) + df_ol = load_parquet_long(ol_path) + obs_series = load_usgs_obs(args.usgs_csv) + + print(f"DA rows: {len(df_da):,} | unique issue_times: {df_da['issue_time'].nunique()}") + print(f"OL rows: {len(df_ol):,} | unique issue_times: {df_ol['issue_time'].nunique()}") + print(f"USGS obs range: {obs_series.index.min()} .. {obs_series.index.max()} " + f"({len(obs_series):,} hours)") + print(f"USGS Helene peak (reference): {USGS_HELENE_PEAK_M3S:.0f} m³/s") + + # ----- Pooled (all issue times) ----- + da_pooled = metrics_by_lead(df_da, obs_series) + ol_pooled = metrics_by_lead(df_ol, obs_series) + out_pooled = os.path.join(out_dir, "lead_time_decay_gauge_pooled.png") + plot_two_panel(out_pooled, da_pooled[0], da_pooled, ol_pooled[0], ol_pooled, + title="Gauge-level lead-time decay — USGS 03463300\n" + "All issue times pooled (Oct 2023 – Oct 2024)") + + # ----- Regime split: Helene, storm, low-flow at gauge ----- + issue_times = pd.to_datetime(sorted(df_da['issue_time'].unique())) + obs_at_issue = pd.Series(issue_times, index=issue_times).map(obs_series) + + helene_mask = (issue_times >= HELENE_START) & (issue_times <= HELENE_END) + storm_mask = obs_at_issue > STORM_OBS_THRESHOLD_M3S + lowflow_mask = obs_at_issue < LOWFLOW_OBS_THRESHOLD_M3S + + regime_metrics = [] + for label, mask in [ + (f"Helene window (Sep 24–28 2024)", helene_mask), + (f"Storm hours (obs(t0) > {STORM_OBS_THRESHOLD_M3S:.0f} m³/s)", storm_mask), + (f"Low flow (obs(t0) < {LOWFLOW_OBS_THRESHOLD_M3S:.0f} m³/s)", lowflow_mask), + ]: + kept = set(pd.to_datetime(issue_times[mask])) + da_m = metrics_by_lead(df_da, obs_series, kept) + ol_m = metrics_by_lead(df_ol, obs_series, kept) + regime_metrics.append((label, len(kept), da_m, ol_m)) + + out_regime = os.path.join(out_dir, "lead_time_decay_gauge_by_regime.png") + plot_three_panel_regime(out_regime, regime_metrics) + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4a_error_decay/run_4a_f1.sh b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4a_error_decay/run_4a_f1.sh new file mode 100644 index 00000000..f5cc2379 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4a_error_decay/run_4a_f1.sh @@ -0,0 +1,80 @@ +#!/usr/bin/env bash +# F1 variance-scaled Vrugt (20pct gauge holdout) — 4a lead-time decay plots. +# +# Runs plot scripts per catchment: +# plot_lead_time_decay.py — pooled RMSE vs lead (catchment-level) +# plot_lead_time_decay_by_regime.py — same split by flow regime +# +# Gauge-level decay plots run once if routed parquets exist: +# plot_lead_time_decay_gauge.py — gauge-level (requires routed parquets) +# plot_forecast_error_fixed_target.py +# plot_forecast_error_per_init.py +# +# Usage: +# bash run_4a_f1.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +LEADTIME_DIR=/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast +DA_DIR=/mnt/disk2/suma_helen_poster/da_results/v2_perturbation_da_on +OBS_DIR=/mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance +ROUTE_DIR=/mnt/disk2/suma_helen_poster/leadtime_troute_routing +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +echo "[F1-4a] Lead-time decay plots — leadtime dir: $LEADTIME_DIR" + +for CAT in "${CATS[@]}"; do + DA_CSV="$LEADTIME_DIR/$CAT/${CAT}_lead_time_forecasts_da.csv" + if [ ! -f "$DA_CSV" ]; then + echo " [$CAT] No lead-time CSV — skipping" + continue + fi + + echo " [$CAT] Plotting pooled decay..." + $TROUTE "$SCRIPT_DIR/plot_lead_time_decay.py" \ + --cat-id "$CAT" \ + --leadtime-dir "$LEADTIME_DIR" \ + --da-dir "$DA_DIR" \ + --obs-dir "$OBS_DIR" + + echo " [$CAT] Plotting decay by regime..." + $TROUTE "$SCRIPT_DIR/plot_lead_time_decay_by_regime.py" \ + --cat-id "$CAT" \ + --leadtime-dir "$LEADTIME_DIR" \ + --da-dir "$DA_DIR" \ + --obs-dir "$OBS_DIR" +done + +# Gauge-level decay (needs routed parquets — run after route_leadtime_f1.sh) +if [ -f "$ROUTE_DIR/routed_leadtime_da_full.parquet" ]; then + echo "[F1-4a] Gauge-level decay (routed)..." + $TROUTE "$SCRIPT_DIR/plot_lead_time_decay_gauge.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" + + echo "[F1-4a] Fixed-target error plots..." + $TROUTE "$SCRIPT_DIR/plot_forecast_error_fixed_target.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" + + echo "[F1-4a] Per-init error plots..." + $TROUTE "$SCRIPT_DIR/plot_forecast_error_per_init.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" +else + echo "[F1-4a] Routed parquets not found — skipping gauge-level and error plots." + echo " Run route_leadtime_f1.sh first, then re-run this script." +fi + +echo "[F1-4a] Done." diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_crossed_ensemble.py b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_crossed_ensemble.py new file mode 100644 index 00000000..7cb00a93 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_crossed_ensemble.py @@ -0,0 +1,229 @@ +""" +plot_crossed_ensemble.py -- 4b: 600-member crossed ensemble vs USGS obs. + +Loads the crossed ensemble parquet produced by run_crossed_ensemble.py: + //_crossed_ensemble.parquet + Columns: issue_time, lead_hour, member_0000..member_0599 + +Projects all members to valid_time = issue_time + lead_hour hours. +Aggregates at each valid_time across ALL members from ALL issue times +to build a probabilistic envelope (5th/25th/50th/75th/95th percentiles). + +Two output figures: + 4b_crossed_ensemble_helene.png + -- Full Helene window (Sep 24-30) probabilistic envelope + USGS obs. + Percentile shading in two layers (5-95 outer, 25-75 inner). + 4b_crossed_ensemble_peak.png + -- Zoomed to Helene peak (Sep 26 12 UTC -> Sep 28 06 UTC), + same shading + USGS obs + individual spaghetti from Sep 27 inits. + +Usage: + python3 plot_crossed_ensemble.py \\ + --ensemble-dir /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble \\ + --cat-id cat-1016300 \\ + --usgs-csv /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv \\ + --out-dir /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble/cat-1016300 +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import matplotlib.patches as mpatches + +DEFAULT_ENSEMBLE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" +DEFAULT_CAT_ID = "cat-1016300" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +PLOT_START = pd.Timestamp("2024-09-24 00:00:00") +PLOT_END = pd.Timestamp("2024-09-28 18:00:00") # cut trailing spike +PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +PEAK_END = pd.Timestamp("2024-09-28 06:00:00") +HELENE_PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +HELENE_PEAK_END = pd.Timestamp("2024-09-28 00:00:00") + + +def load_ensemble(path): + df = pd.read_parquet(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + return df, member_cols + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def build_percentile_envelope(df, member_cols, t_start, t_end): + """ + For each valid_time in [t_start, t_end], stack all member values + from all issue times and compute percentiles. + Returns DataFrame indexed by valid_time with columns p05, p25, p50, p75, p95. + """ + mask = (df["valid_time"] >= t_start) & (df["valid_time"] <= t_end) + sub = df[mask].copy() + sub["valid_time"] = sub["valid_time"].dt.floor("1h") + + records = [] + for vt, grp in sub.groupby("valid_time"): + vals = grp[member_cols].to_numpy(dtype=float).flatten() + vals = vals[~np.isnan(vals)] * MM_H_TO_M3_S + if len(vals) == 0: + continue + records.append({ + "valid_time": vt, + "p05": np.percentile(vals, 5), + "p25": np.percentile(vals, 25), + "p50": np.percentile(vals, 50), + "p75": np.percentile(vals, 75), + "p95": np.percentile(vals, 95), + "n": len(vals), + }) + return pd.DataFrame(records).set_index("valid_time").sort_index() + + +def _add_helene_band(ax, label=True): + ax.axvspan(HELENE_PEAK_START, HELENE_PEAK_END, + color="salmon", alpha=0.12, zorder=0) + if label: + ax.text(HELENE_PEAK_START + pd.Timedelta(hours=6), 0.97, "Helene peak", + transform=ax.get_xaxis_transform(), + fontsize=9, color="firebrick", ha="left", va="top") + + +def _format_xaxis(ax, major_interval=1, minor_hours=6): + ax.xaxis.set_major_locator(mdates.DayLocator(interval=major_interval)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.xaxis.set_minor_locator(mdates.HourLocator(interval=minor_hours)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.2, lw=0.4) + + +def plot_envelope(ax, env, obs_series, t_start, t_end, + ens_color="tab:blue", obs_color="black", helene_label=True): + _add_helene_band(ax, label=helene_label) + + # Outer band: 5-95 + ax.fill_between(env.index, env["p05"], env["p95"], + color=ens_color, alpha=0.15, zorder=2, + label="5th-95th percentile") + # Inner band: 25-75 + ax.fill_between(env.index, env["p25"], env["p75"], + color=ens_color, alpha=0.30, zorder=3, + label="25th-75th percentile") + # Median + ax.plot(env.index, env["p50"], + color=ens_color, lw=2.2, zorder=5, + label="Ensemble median (600 members)") + + # USGS obs + obs_w = obs_series.loc[t_start:t_end] + ax.plot(obs_w.index, obs_w.values, + color=obs_color, lw=1.8, zorder=6, label="USGS obs") + + ax.set_xlim(t_start, t_end) + ax.set_ylabel("Discharge (m3/s)", fontsize=11) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--ensemble-dir", default=DEFAULT_ENSEMBLE_DIR) + parser.add_argument("--cat-id", default=DEFAULT_CAT_ID) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--out-dir", default=None) + args = parser.parse_args() + + cat_dir = os.path.join(args.ensemble_dir, args.cat_id) + out_dir = args.out_dir or cat_dir + os.makedirs(out_dir, exist_ok=True) + + parquet_path = os.path.join(cat_dir, f"{args.cat_id}_crossed_ensemble.parquet") + print(f"Loading crossed ensemble: {parquet_path}") + df, member_cols = load_ensemble(parquet_path) + obs = load_usgs(args.usgs_csv) + print(f" {df['issue_time'].nunique()} issue times, " + f"{len(member_cols)} members per row") + + # ------------------------------------------------------------------ # + # Figure 4b-full: full Helene window probabilistic envelope # + # ------------------------------------------------------------------ # + print("Building percentile envelope (full window)...") + env_full = build_percentile_envelope(df, member_cols, PLOT_START, PLOT_END) + print(f" Valid times with data: {len(env_full)}, " + f"median members per timestep: {env_full['n'].median():.0f}") + + fig, ax = plt.subplots(figsize=(17, 7)) + plot_envelope(ax, env_full, obs, PLOT_START, PLOT_END) + ax.set_xlabel("Date", fontsize=11) + ax.set_title( + f"4b -- 600-member crossed ensemble (30 met x 20 hydro-state) | {args.cat_id}\n" + "Shaded: 5-95th (light) and 25-75th (dark) percentile across all members " + "and issue times | Sep 24-30 2024 | USGS 03463300", + fontsize=11, + ) + ax.legend(fontsize=10, loc="upper left", frameon=True, framealpha=0.9) + _format_xaxis(ax) + plt.tight_layout() + out1 = os.path.join(out_dir, f"{args.cat_id}_4b_crossed_ensemble_helene.png") + plt.savefig(out1, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + # ------------------------------------------------------------------ # + # Figure 4b-peak: zoomed to Helene peak window # + # ------------------------------------------------------------------ # + print("Building percentile envelope (peak window)...") + env_peak = build_percentile_envelope(df, member_cols, PEAK_START, PEAK_END) + + fig, (ax_main, ax_zoom) = plt.subplots( + 2, 1, figsize=(14, 11), + gridspec_kw={"height_ratios": [1, 1.4]}, sharex=False) + + # Top panel: full window context (smaller) + plot_envelope(ax_main, env_full, obs, PLOT_START, PLOT_END, helene_label=False) + ax_main.set_title("Full Helene window context (Sep 24 - Sep 28 18 UTC)", fontsize=10) + _format_xaxis(ax_main) + # Shade the zoom region on context panel + ax_main.axvspan(PEAK_START, PEAK_END, color="gold", alpha=0.18, zorder=1) + ax_main.text(PEAK_START + pd.Timedelta(hours=1), 0.97, "zoom", + transform=ax_main.get_xaxis_transform(), + fontsize=8, color="goldenrod", ha="left", va="top") + + # Bottom panel: peak zoom + plot_envelope(ax_zoom, env_peak, obs, PEAK_START, PEAK_END, helene_label=True) + ax_zoom.set_xlabel("Date", fontsize=11) + ax_zoom.set_title( + "Helene peak zoom (Sep 26 12 UTC - Sep 28 06 UTC)", fontsize=10) + _format_xaxis(ax_zoom, major_interval=1, minor_hours=3) + ax_zoom.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %HUTC")) + ax_zoom.legend(fontsize=10, loc="lower right", frameon=True, framealpha=0.9) + + fig.suptitle( + f"4b -- 600-member crossed ensemble vs USGS 03463300 | {args.cat_id}\n" + "30 met forcing draws x 20 DA analysis state draws | R=0.07 mm2/h2", + fontsize=12, y=1.01, + ) + plt.tight_layout() + out2 = os.path.join(out_dir, f"{args.cat_id}_4b_crossed_ensemble_peak.png") + plt.savefig(out2, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_per_member_factor_decomposition.py b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_per_member_factor_decomposition.py new file mode 100644 index 00000000..880d7c51 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_per_member_factor_decomposition.py @@ -0,0 +1,154 @@ +""" +Per-member factor-decomposition plot for ONE catchment. + +For each of the 20 ensemble members, show the member's actual forecast +(production: init + forcing + process + DA all active) alongside the member's +trajectory under each ISOLATED perturbation source (init only / forcing only / +process noise only). The viewer reads each panel as 'why did THIS member give +THIS forecast — which perturbation source pushed it where?' + +Note: 'same member_i' across the four CSVs is a column-name correspondence +only; the underlying random draws are independent across the four runs. So +each panel is an illustrative comparison, not a rigorous matched-seed Shapley +decomposition. The story still reads correctly: each colored line shows what +ONE realization of that perturbation source produces, and the thick line +shows what the full production setup produces. + +Inputs (per catchment): + //_production_per_member.csv (20 cols) + //_sensitivity_init.csv (20 cols) + //_sensitivity_forcing.csv (20 cols) + //_sensitivity_process.csv (20 cols) + //_test_results.csv (for Qkrig obs) + +Output: + //_per_member_factor_decomp.png +""" +import os +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +# ---------- Configuration ---------- +CAT = "cat-1016300" # change here to do a different catchment + +PROD_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_production_per_member" +SEN_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_sensitivity" +DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" +OUT_PNG = os.path.join(PROD_DIR, CAT, f"{CAT}_per_member_factor_decomp.png") + +ZOOM_START = pd.Timestamp("2024-09-24") +ZOOM_END = pd.Timestamp("2024-09-28") + +PROD_COLOR = "tab:purple" +INIT_COLOR = "tab:red" +FORCING_COLOR = "tab:blue" +PROC_COLOR = "tab:green" +OBS_COLOR = "black" + + +def load_member_csv(path): + if not os.path.exists(path): + return None, None + df = pd.read_csv(path, parse_dates=["date"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + return df["date"].values, df[member_cols] + + +def load_obs(): + p = os.path.join(DA_DIR, CAT, f"{CAT}_test_results.csv") + if not os.path.exists(p): + return None, None + df = pd.read_csv(p, parse_dates=["date"]) + return df["date"].values, df["obs_mm_h"].values + + +def main(): + prod_dates, prod = load_member_csv(os.path.join(PROD_DIR, CAT, f"{CAT}_production_per_member.csv")) + init_dates, init = load_member_csv(os.path.join(SEN_DIR, CAT, f"{CAT}_sensitivity_init.csv")) + forc_dates, forc = load_member_csv(os.path.join(SEN_DIR, CAT, f"{CAT}_sensitivity_forcing.csv")) + proc_dates, proc = load_member_csv(os.path.join(SEN_DIR, CAT, f"{CAT}_sensitivity_process.csv")) + obs_dates, obs = load_obs() + + if prod is None: + raise FileNotFoundError( + f"Missing {PROD_DIR}/{CAT}/{CAT}_production_per_member.csv — " + "run run_production_per_member.py first.") + if init is None or forc is None or proc is None: + raise FileNotFoundError( + f"Missing one of the sensitivity CSVs under {SEN_DIR}/{CAT}/. " + "Run run_perturbation_sensitivity.py for all three sources first.") + + def to_mask(dates): + d = pd.to_datetime(dates) + return d, (d >= ZOOM_START) & (d <= ZOOM_END) + + pd_dates, pd_mask = to_mask(prod_dates) + id_dates, id_mask = to_mask(init_dates) + fd_dates, fd_mask = to_mask(forc_dates) + cd_dates, cd_mask = to_mask(proc_dates) + if obs_dates is not None: + od, om = to_mask(obs_dates) + else: + od, om = None, None + + N = prod.shape[1] + # Lay out 4 cols x 5 rows = 20 panels + n_cols = 4 + n_rows = (N + n_cols - 1) // n_cols + fig, axes = plt.subplots(n_rows, n_cols, figsize=(22, 4 * n_rows), sharex=True) + axes = axes.flatten() + + for i in range(N): + ax = axes[i] + col = f"member_{i:02d}" + + # Each member's INIT-only trajectory (thin red) + if col in init.columns: + ax.plot(id_dates[id_mask], init[col].values[id_mask], + color=INIT_COLOR, lw=0.9, alpha=0.9, label="Init only") + # FORCING-only (thin blue) + if col in forc.columns: + ax.plot(fd_dates[fd_mask], forc[col].values[fd_mask], + color=FORCING_COLOR,lw=0.9, alpha=0.9, label="Forcing only") + # PROCESS-only (thin green) + if col in proc.columns: + ax.plot(cd_dates[cd_mask], proc[col].values[cd_mask], + color=PROC_COLOR, lw=0.9, alpha=0.9, label="Process noise only") + # PRODUCTION (thick purple) — actual forecast with DA on, all perturbations + if col in prod.columns: + ax.plot(pd_dates[pd_mask], prod[col].values[pd_mask], + color=PROD_COLOR, lw=2.2, alpha=0.95, label="Production (DA on)") + # Qkrig observation (thick black dashed) + if od is not None: + ax.plot(od[om], obs[om], + color=OBS_COLOR, lw=1.4, linestyle="--", alpha=0.9, label="Qkrig (obs)") + + ax.set_title(f"{col}", fontsize=10) + ax.tick_params(axis="x", rotation=45, labelsize=7) + ax.tick_params(axis="y", labelsize=8) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=1)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + if i == 0: + ax.legend(fontsize=7, loc="upper left") + + # Hide any extra blank axes + for j in range(N, len(axes)): + axes[j].axis("off") + + fig.suptitle( + f"Per-member factor decomposition - {CAT} - Hurricane Helene peak (Sep 24-28, 2024)\n" + "Purple thick = production member (DA on, all perturbations) | " + "Red = init only | Blue = forcing only | Green = process noise only | " + "Black dashed = Qkrig obs", + fontsize=12, y=1.00, + ) + fig.text(0.005, 0.5, "Discharge (mm/h)", va="center", rotation="vertical", fontsize=12) + plt.tight_layout() + plt.savefig(OUT_PNG, dpi=140, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_per_member_input_output_diagnostic.py b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_per_member_input_output_diagnostic.py new file mode 100644 index 00000000..f6be63d8 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_per_member_input_output_diagnostic.py @@ -0,0 +1,177 @@ +""" +Per-member input/output diagnostic for ONE catchment. + +For each ensemble member, show a single combined panel (hydrograph-style) that +traces the inputs the member actually received and the output it produced: + + Each member panel uses an internal 2-row stack: + TOP sub-panel: perturbed precip (bars, left axis) + + perturbed PET (line, right twinx axis) + BOTTOM sub-panel: simulated streamflow Q (purple) + + Qkrig observation (black dashed) + Shared x-axis between the two sub-panels. + + Initial states for that member are shown in the title line. + +Layout: 5 rows x 4 columns = 20 member panels. + +Inputs (all from run_production_per_member.py): + //_production_per_member.csv (Q per member) + //_production_per_member_precip.csv (precip per member) + //_production_per_member_pet.csv (PET per member) + //_production_per_member_initial_states.json + //_test_results.csv (Qkrig obs) + +Output: + //_per_member_io_diagnostic.png +""" +import os +import json +import pandas as pd +import numpy as np +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import matplotlib.gridspec as gridspec + +# ----- Configuration ----- +CAT = "cat-1016300" + +PROD_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_production_per_member" +DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" +OUT_PNG = os.path.join(PROD_DIR, CAT, f"{CAT}_per_member_io_diagnostic.png") + +# Plot window (Helene 4-day zoom) +ZOOM_START = pd.Timestamp("2024-09-24") +ZOOM_END = pd.Timestamp("2024-09-28") + +PRECIP_COLOR = "tab:blue" +PET_COLOR = "tab:orange" +Q_COLOR = "tab:purple" +OBS_COLOR = "black" + +N_ROWS = 5 # grid rows of member panels +N_COLS = 4 # grid cols of member panels + + +def load_member_csv(path): + if not os.path.exists(path): + return None, None + df = pd.read_csv(path, parse_dates=["date"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + return df["date"].values, df[member_cols] + + +def main(): + q_dates, q_df = load_member_csv(os.path.join(PROD_DIR, CAT, f"{CAT}_production_per_member.csv")) + precip_dates, precip_df = load_member_csv(os.path.join(PROD_DIR, CAT, f"{CAT}_production_per_member_precip.csv")) + pet_dates, pet_df = load_member_csv(os.path.join(PROD_DIR, CAT, f"{CAT}_production_per_member_pet.csv")) + + init_states_path = os.path.join(PROD_DIR, CAT, f"{CAT}_production_per_member_initial_states.json") + if os.path.exists(init_states_path): + with open(init_states_path) as f: + init_states = json.load(f)["members"] + else: + init_states = {} + + obs_path = os.path.join(DA_DIR, CAT, f"{CAT}_test_results.csv") + if os.path.exists(obs_path): + obs_df = pd.read_csv(obs_path, parse_dates=["date"]) + obs_dates = pd.to_datetime(obs_df["date"].values) + obs_vals = obs_df["obs_mm_h"].values + else: + obs_dates = None + obs_vals = None + + if q_df is None or precip_df is None or pet_df is None: + raise FileNotFoundError( + "Missing one of the per-member CSVs. Run run_production_per_member.py " + "with the updated script that saves precip + PET + initial states.") + + member_cols = list(q_df.columns) + N = len(member_cols) + assert N <= N_ROWS * N_COLS, f"Grid too small for {N} members" + + qd = pd.to_datetime(q_dates) + pd_ = pd.to_datetime(precip_dates) + ed = pd.to_datetime(pet_dates) + qmask = (qd >= ZOOM_START) & (qd <= ZOOM_END) + pmask = (pd_ >= ZOOM_START) & (pd_ <= ZOOM_END) + emask = (ed >= ZOOM_START) & (ed <= ZOOM_END) + if obs_dates is not None: + omask = (obs_dates >= ZOOM_START) & (obs_dates <= ZOOM_END) + + # Compute global y-axis ranges so all member panels are directly comparable + p_ymax = float(np.nanmax(precip_df.values[pmask, :])) * 1.10 + e_ymax = float(np.nanmax(pet_df.values[emask, :])) * 1.10 + q_data_max = float(np.nanmax(q_df.values[qmask, :])) + if obs_dates is not None: + q_data_max = max(q_data_max, float(np.nanmax(obs_vals[omask]))) + q_ymax = q_data_max * 1.10 + + fig = plt.figure(figsize=(22, 19)) + outer = gridspec.GridSpec(N_ROWS, N_COLS, hspace=0.75, wspace=0.30) + + for i in range(N): + col = member_cols[i] + row, c = divmod(i, N_COLS) + # Each member cell is a 2-row inner grid (precip+PET on top, Q on bottom) + inner = gridspec.GridSpecFromSubplotSpec( + 2, 1, subplot_spec=outer[row, c], hspace=0.20, height_ratios=[1.0, 1.6]) + + # --- Top sub-panel: precip bars + PET line on twinx --- + ax_top = fig.add_subplot(inner[0]) + ax_top.bar(pd_[pmask], precip_df[col].values[pmask], + width=0.04, color=PRECIP_COLOR, alpha=0.85, label="P") + ax_top.set_ylim(0, p_ymax) + ax_top.set_ylabel("P (mm/h)", fontsize=8, color=PRECIP_COLOR) + ax_top.tick_params(axis="y", labelsize=7, labelcolor=PRECIP_COLOR) + ax_top.tick_params(axis="x", which="both", bottom=False, labelbottom=False) + ax_top.grid(True, alpha=0.15) + # PET on twinx (different scale) + ax_pet = ax_top.twinx() + ax_pet.plot(ed[emask], pet_df[col].values[emask], + color=PET_COLOR, lw=1.0, label="PET") + ax_pet.set_ylim(0, e_ymax) + ax_pet.set_ylabel("PET (mm/h)", fontsize=8, color=PET_COLOR) + ax_pet.tick_params(axis="y", labelsize=7, labelcolor=PET_COLOR) + # Title: member name on line 1, initial states on line 2 (avoids overlap with neighbors) + s = init_states.get(col, {}) + title = (f"{col}\n" + f"soil={s.get('soil_m', float('nan')):.3f} m | " + f"GW={s.get('gw_m', float('nan')):.4f} m | " + f"Nash[0]={s.get('nash0_m', float('nan')):.1e} " + f"Nash[1]={s.get('nash1_m', float('nan')):.1e}") + ax_top.set_title(title, fontsize=8, loc="left") + + # --- Bottom sub-panel: Q + obs (shares x-axis with top) --- + ax_bot = fig.add_subplot(inner[1], sharex=ax_top) + ax_bot.plot(qd[qmask], q_df[col].values[qmask], + color=Q_COLOR, lw=1.6, label=col) + if obs_dates is not None: + ax_bot.plot(obs_dates[omask], obs_vals[omask], + color=OBS_COLOR, lw=1.2, linestyle="--", label="Qkrig (obs)") + ax_bot.set_ylim(0, q_ymax) + ax_bot.set_ylabel("Q (mm/h)", fontsize=8) + ax_bot.tick_params(labelsize=7) + ax_bot.xaxis.set_major_locator(mdates.DayLocator(interval=1)) + ax_bot.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax_bot.grid(True, alpha=0.15) + if i == 0: + ax_bot.legend(fontsize=7, loc="upper left", ncol=2, + frameon=True, framealpha=0.85) + + fig.suptitle( + f"Per-member input/output diagnostic - {CAT} - " + f"Hurricane Helene peak (Sep 24-28, 2024)\n" + "Each member panel: TOP = perturbed precip (blue bars) + perturbed PET (orange line) | " + "BOTTOM = simulated Q (purple) vs Qkrig obs (black dashed). " + "Title shows initial states at t=0.", + fontsize=12, y=0.995, + ) + plt.savefig(OUT_PNG, dpi=140, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_perturbation_category_shaded.py b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_perturbation_category_shaded.py new file mode 100644 index 00000000..b845c52d --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_perturbation_category_shaded.py @@ -0,0 +1,165 @@ +""" +Per-catchment shaded ensemble-band plot, organized by perturbation category. + +Three categories: + 1. Initial states (red) + 2. Meteorological forcings (blue) + 3. Hydrological states (green) + +For each category, all 20 ensemble members are shown as a shaded band +(min-max envelope fill) plus a thicker median line in the same color. +Qkrig observation overlaid in black. Hurricane Helene peak window shaded +in pink. Styled after a standard ensemble-forecast figure layout. + +Inputs (existing per-source sensitivity CSVs from run_perturbation_sensitivity.py): + //_sensitivity_init.csv (20 members, init only) + //_sensitivity_forcing.csv (20 members, forcing only) + //_sensitivity_process.csv (20 members, process noise only) + //_test_results.csv (Qkrig obs) + +Outputs: + //_perturbation_categories_linear.png + //_perturbation_categories_log.png +""" +import os +import pandas as pd +import numpy as np +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import matplotlib.patches as mpatches + +CAT = "cat-1016300" + +SEN_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_sensitivity" +OBS_DIR = "/mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance" + +OUT_LINEAR = os.path.join(SEN_DIR, CAT, f"{CAT}_perturbation_categories_linear.png") +OUT_LOG = os.path.join(SEN_DIR, CAT, f"{CAT}_perturbation_categories_log.png") + +# Plot window — wider context, similar to the paper's Sep 10 - Oct 08 +PLOT_START = pd.Timestamp("2024-09-20") +PLOT_END = pd.Timestamp("2024-10-05") + +# Helene peak band +HELENE_START = pd.Timestamp("2024-09-26 12:00:00") +HELENE_END = pd.Timestamp("2024-09-28 00:00:00") + +# Category configuration: file suffix, display label, color +CATEGORIES = [ + ("init", "Initial states only", "tab:red"), + ("forcing", "Meteorological forcings only", "tab:blue"), + ("process", "Hydrological states only", "tab:green"), +] + + +def load_members(source): + path = os.path.join(SEN_DIR, CAT, f"{CAT}_sensitivity_{source}.csv") + if not os.path.exists(path): + return None, None + df = pd.read_csv(path, parse_dates=["date"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + return df["date"].values, df[member_cols].to_numpy(dtype=float) + + +def load_obs(): + p = os.path.join(OBS_DIR, f"{CAT}.csv") + if not os.path.exists(p): + return None, None + df = pd.read_csv(p, parse_dates=["date"]) + return df["date"].values, df["qkrig"].values + + +def plot_panel(ax, obs_dates, obs_vals, log_y=False): + handles_labels = [] # for the legend + + # Plot each category as a shaded band + median line + for source, label, color in CATEGORIES: + dates, q = load_members(source) + if dates is None: + continue + d = pd.to_datetime(dates) + mask = (d >= PLOT_START) & (d <= PLOT_END) + if mask.sum() == 0: + continue + # Min-max envelope across all 20 members per timestep (widest possible band). + # Bands are visually narrow even with min/max because perturbations are + # tuned for production EnKF stability, not for max visible spread. + q_window = q[mask, :] + qmin = np.nanmin(q_window, axis=1) + qmax = np.nanmax(q_window, axis=1) + median = np.nanmedian(q_window, axis=1) + + ax.fill_between(d[mask], qmin, qmax, + color=color, alpha=0.30, zorder=2, + edgecolor="none") + line, = ax.plot(d[mask], median, + color=color, lw=1.7, alpha=0.95, zorder=3, + label=f"{label} (N=20)") + handles_labels.append((line, label)) + + # Helene peak shaded band (vertical) + ax.axvspan(HELENE_START, HELENE_END, + color="salmon", alpha=0.15, zorder=1) + ax.text((HELENE_START + (HELENE_END - HELENE_START) / 2), + ax.get_ylim()[1] if not log_y else 1.0, + "Helene peak", + fontsize=9, color="salmon", + ha="center", va="bottom", zorder=3) + + # Observation + if obs_dates is not None: + od = pd.to_datetime(obs_dates) + om = (od >= PLOT_START) & (od <= PLOT_END) + ax.plot(od[om], obs_vals[om], + color="black", lw=1.4, label="Qkrig (obs)", zorder=4) + + ax.set_xlabel("Date", fontsize=10) + if log_y: + ax.set_yscale("log") + ax.set_ylabel(r"q (mm hr$^{-1}$) [log scale]", fontsize=10) + ax.set_ylim(0.01, None) + else: + ax.set_ylabel("Discharge (mm/h)", fontsize=10) + ax.tick_params(axis="x", rotation=30, labelsize=9) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=2)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.grid(True, alpha=0.2) + ax.legend(loc="upper left", fontsize=9, frameon=True, framealpha=0.9) + + +def main(): + obs_dates, obs_vals = load_obs() + + # ----- Linear-y ----- + fig, ax = plt.subplots(1, 1, figsize=(20, 7)) + plot_panel(ax, obs_dates, obs_vals, log_y=False) + fig.suptitle( + f"Ensemble forecast by perturbation category - {CAT} - " + f"Sep 20 - Oct 5, 2024 (Hurricane Helene window)\n" + "Shaded bands = min-max envelope across 20 members. " + "Lines = ensemble median. DA off in all three sub-experiments.", + fontsize=11, y=0.995, + ) + plt.tight_layout() + plt.savefig(OUT_LINEAR, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_LINEAR}") + + # ----- Log-y (paper style) ----- + fig, ax = plt.subplots(1, 1, figsize=(20, 7)) + plot_panel(ax, obs_dates, obs_vals, log_y=True) + fig.suptitle( + f"Ensemble forecast by perturbation category - {CAT} - log-scale q - " + f"Sep 20 - Oct 5, 2024 (Hurricane Helene window)\n" + "Shaded bands = min-max envelope across 20 members. " + "Lines = ensemble median. DA off in all three sub-experiments.", + fontsize=11, y=0.995, + ) + plt.tight_layout() + plt.savefig(OUT_LOG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_LOG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_production_ensemble_forecast.py b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_production_ensemble_forecast.py new file mode 100644 index 00000000..e0214018 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_production_ensemble_forecast.py @@ -0,0 +1,164 @@ +""" +Single-panel ensemble forecast plot for one catchment, paper-style. + +Plots all 20 production-per-member streamflow trajectories overlaid on a +single time-series panel, with the Hurricane Helene window highlighted by a +pink shaded vertical band and the Qkrig observation drawn on top in black. + +Ensemble forecast figure style (50/100/200 km +variogram-range ensemble panels): every member as a thin colored line, storm +window shaded, observation as a thick dark series, optional log-scaled y-axis. + +Inputs: + //_production_per_member.csv (date + 20 members) + //_test_results.csv (Qkrig obs) + +Outputs: + //_production_ensemble_forecast_linear.png + //_production_ensemble_forecast_log.png +""" + +import os +import pandas as pd +import numpy as np +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import matplotlib.patches as mpatches +from matplotlib import cm + +# ----- Configuration ----- +CAT = "cat-1016300" + +PROD_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_production_per_member" +DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" + +# Plot window: a few weeks around Helene so the storm is in context +PLOT_START = pd.Timestamp("2024-09-20") +PLOT_END = pd.Timestamp("2024-10-05") + +# Helene peak window highlighted in pink +HELENE_START = pd.Timestamp("2024-09-26 12:00:00") +HELENE_END = pd.Timestamp("2024-09-28 00:00:00") + +OUT_LINEAR = os.path.join(PROD_DIR, CAT, f"{CAT}_production_ensemble_forecast_linear.png") +OUT_LOG = os.path.join(PROD_DIR, CAT, f"{CAT}_production_ensemble_forecast_log.png") + + +def kge_score(obs, sim): + m = ~(np.isnan(obs) | np.isnan(sim)) + o, s = obs[m], sim[m] + if len(o) < 2: + return float("nan") + denom = float(np.sqrt(((o - o.mean()) ** 2).sum() * ((s - s.mean()) ** 2).sum())) + if denom == 0: + return float("nan") + r = float(((o - o.mean()) * (s - s.mean())).sum() / denom) + alpha = float(s.std() / o.std()) if o.std() != 0 else float("nan") + beta = float(s.mean() / o.mean()) if o.mean() != 0 else float("nan") + return 1.0 - float(np.sqrt((r - 1) ** 2 + (alpha - 1) ** 2 + (beta - 1) ** 2)) + + +def main(): + prod_path = os.path.join(PROD_DIR, CAT, f"{CAT}_production_per_member.csv") + obs_path = os.path.join(DA_DIR, CAT, f"{CAT}_test_results.csv") + + if not os.path.exists(prod_path): + raise FileNotFoundError( + f"Missing {prod_path}. Run run_production_per_member.py for {CAT} first.") + if not os.path.exists(obs_path): + print(f"Note: {obs_path} not found; will plot without observation overlay.") + + prod_df = pd.read_csv(prod_path, parse_dates=["date"]) + member_cols = sorted([c for c in prod_df.columns if c.startswith("member_")]) + member_arr = prod_df[member_cols].to_numpy(dtype=float) + dates = pd.to_datetime(prod_df["date"].values) + + if os.path.exists(obs_path): + obs_df = pd.read_csv(obs_path, parse_dates=["date"]) + obs_dates = pd.to_datetime(obs_df["date"].values) + obs_vals = obs_df["obs_mm_h"].values + else: + obs_dates = None + obs_vals = None + + # Compute per-member peak (within plot window) for the legend label + pw_mask = (dates >= PLOT_START) & (dates <= PLOT_END) + peaks = member_arr[pw_mask, :].max(axis=0) + + # Optional: compute per-member KGE vs obs across the plot window + member_kges = np.full(len(member_cols), np.nan) + if obs_dates is not None: + obs_series = pd.Series(obs_vals, index=obs_dates) + for i, _ in enumerate(member_cols): + member_series = pd.Series(member_arr[:, i], index=dates) + joined = pd.concat([obs_series, member_series], axis=1, join="inner").dropna() + if len(joined) >= 2: + member_kges[i] = kge_score(joined.iloc[:, 0].values, joined.iloc[:, 1].values) + + def plot_panel(ax, log_y=False): + colormap = cm.get_cmap("turbo", len(member_cols)) + # Thin colored lines per member + for i, col in enumerate(member_cols): + label = f"{col} | peak={peaks[i]:.1f} mm/h" + ax.plot(dates[pw_mask], member_arr[pw_mask, i], + color=colormap(i), lw=0.8, alpha=0.85, label=label, zorder=2) + + # Helene peak shaded band + ax.axvspan(HELENE_START, HELENE_END, color="salmon", alpha=0.18, zorder=1) + ax.text(HELENE_START + (HELENE_END - HELENE_START) / 2, + ax.get_ylim()[1] if not log_y else 1.0, + "Helene peak", fontsize=8, color="salmon", + ha="center", va="bottom", zorder=3) + + # Observation as thick black series + if obs_dates is not None: + obs_mask = (obs_dates >= PLOT_START) & (obs_dates <= PLOT_END) + ax.plot(obs_dates[obs_mask], obs_vals[obs_mask], + color="black", lw=1.6, label="Qkrig (obs)", zorder=4) + + ax.set_xlabel("Date", fontsize=10) + if log_y: + ax.set_yscale("log") + ax.set_ylabel(r"q (mm hr$^{-1}$) [log scale]", fontsize=10) + ax.set_ylim(0.01, max(member_arr.max(), (obs_vals.max() if obs_vals is not None else 0)) * 1.2) + else: + ax.set_ylabel("Discharge (mm/h)", fontsize=10) + + ax.tick_params(axis="x", rotation=30, labelsize=9) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=2)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.grid(True, alpha=0.2) + + # ---- LINEAR Y figure ---- + fig, ax = plt.subplots(1, 1, figsize=(20, 7)) + plot_panel(ax, log_y=False) + ax.legend(fontsize=6.5, loc="upper left", ncol=2, + frameon=True, framealpha=0.85, markerfirst=False) + fig.suptitle( + f"Production ensemble forecast (N=20) at {CAT} - " + f"Sep 20 - Oct 5, 2024 (Hurricane Helene window)", + fontsize=12, y=0.99, + ) + plt.tight_layout() + plt.savefig(OUT_LINEAR, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_LINEAR}") + + # ---- LOG Y figure (matches paper style) ---- + fig, ax = plt.subplots(1, 1, figsize=(20, 7)) + plot_panel(ax, log_y=True) + ax.legend(fontsize=6.5, loc="upper left", ncol=2, + frameon=True, framealpha=0.85, markerfirst=False) + fig.suptitle( + f"Production ensemble forecast (N=20) at {CAT} - log-scale q - " + f"Sep 20 - Oct 5, 2024 (Hurricane Helene window)", + fontsize=12, y=0.99, + ) + plt.tight_layout() + plt.savefig(OUT_LOG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_LOG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_combined.py b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_combined.py new file mode 100644 index 00000000..c6d531a0 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_combined.py @@ -0,0 +1,235 @@ +#!/usr/bin/env python3 +""" +plot_routed_ensemble_combined.py +Combined poster plot: 600-member routed ensemble envelope + +deterministic DA lines (Vrugt/no-Vrugt) vs USGS gauge obs at 03463300. + +Inputs: + --vrugt-csv : routed_Q_test.csv from vrugt_dynamic_routed/ + --novrugt-csv : routed_Q_test.csv from novrugt_dynamic_routed/ + --ensemble-pq : routed_crossed_ensemble.parquet from run_route_crossed_ensemble.py + --out-dir : output directory + +The ensemble parquet columns: issue_time, lead_hour, member_0000..member_0599 (Q in m3/s) + +Usage: + python3 plot_routed_ensemble_combined.py \ + --vrugt-csv /mnt/disk2/suma_helen_poster/da_results/vrugt_dynamic_routed/routed_Q_test.csv \ + --novrugt-csv /mnt/disk2/suma_helen_poster/da_results/novrugt_dynamic_routed/routed_Q_test.csv \ + --ensemble-pq /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble_routed/routed_crossed_ensemble.parquet \ + --out-dir /mnt/disk2/suma_helen_poster/da_results/comparison_plots +""" + +import argparse +import os + +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +HELENE_START = pd.Timestamp("2024-09-25 00:00:00") +HELENE_END = pd.Timestamp("2024-09-29 00:00:00") +ZOOM_START = pd.Timestamp("2024-09-24 00:00:00") +ZOOM_END = pd.Timestamp("2024-09-28 18:00:00") # cut trailing ensemble bump + +COLOR_OBS = "black" +COLOR_VRUGT = "#1f77b4" # blue +COLOR_NOVRUGT = "#d62728" # red +COLOR_ENS = "#1f77b4" # same blue family as Vrugt + + +def load_det_csv(path): + df = pd.read_csv(path, parse_dates=["date"]) + return df.set_index("date").sort_index() + + +def build_envelope(pq_path, zoom_start, zoom_end): + """ + Load routed ensemble parquet and build percentile envelope by valid_time. + Returns DataFrame indexed by valid_time with columns p05, p25, p50, p75, p95. + """ + df = pd.read_parquet(pq_path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = [c for c in df.columns if c.startswith("member_")] + + records = [] + for _, row in df.iterrows(): + t0 = row["issue_time"] + lead = int(row["lead_hour"]) + vt = t0 + pd.Timedelta(hours=lead) + vals = row[member_cols].values.astype(float) + records.append({"valid_time": vt, "vals": vals}) + + # Group by valid_time (floor to hour), stack all member values across issue_times + from collections import defaultdict + groups = defaultdict(list) + for rec in records: + vt = rec["valid_time"].floor("1h") + groups[vt].extend(rec["vals"].tolist()) + + rows = [] + for vt in sorted(groups): + if vt < zoom_start or vt > zoom_end: + continue + v = np.array(groups[vt]) + v = v[~np.isnan(v)] + if len(v) == 0: + continue + rows.append({ + "valid_time": vt, + "p05": np.percentile(v, 5), + "p25": np.percentile(v, 25), + "p50": np.percentile(v, 50), + "p75": np.percentile(v, 75), + "p95": np.percentile(v, 95), + }) + + return pd.DataFrame(rows).set_index("valid_time") + + +def compute_kge(obs, sim): + mask = ~(np.isnan(obs) | np.isnan(sim)) + o, s = obs[mask], sim[mask] + if len(o) < 2 or np.std(o) == 0: + return np.nan + r = np.corrcoef(o, s)[0, 1] + return 1.0 - np.sqrt((r-1)**2 + (np.std(s)/np.std(o)-1)**2 + (np.mean(s)/np.mean(o)-1)**2) + + +def _format_xaxis(ax, locator, fmt, rotation=30): + ax.xaxis.set_major_locator(locator) + ax.xaxis.set_major_formatter(mdates.DateFormatter(fmt)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=rotation, ha="right", fontsize=9) + + +def _add_helene_band(ax): + ax.axvspan(HELENE_START, HELENE_END, color="gold", alpha=0.15, zorder=0) + + +def plot_zoom(ax, dates, obs, sim_v, sim_nv, kge_v, kge_nv, env=None): + mask = (dates >= ZOOM_START) & (dates <= ZOOM_END) + dz = dates[mask] + + # Ensemble envelope (bottom layer) + if env is not None: + ax.fill_between(env.index, env["p05"], env["p95"], + color=COLOR_ENS, alpha=0.18, label="Ensemble 5-95th pct") + ax.fill_between(env.index, env["p25"], env["p75"], + color=COLOR_ENS, alpha=0.30, label="Ensemble 25-75th pct") + ax.plot(env.index, env["p50"], + color=COLOR_ENS, lw=1.2, linestyle="--", alpha=0.7, label="Ensemble median") + + # Deterministic DA lines + ax.plot(dz, sim_nv[mask], color=COLOR_NOVRUGT, lw=1.4, linestyle="--", + label=f"DA -- constant R KGE={kge_nv:.3f}", zorder=3) + ax.plot(dz, sim_v[mask], color=COLOR_VRUGT, lw=1.6, + label=f"DA -- dynamic R (Vrugt) KGE={kge_v:.3f}", zorder=4) + + # USGS obs on top + ax.plot(dz, obs[mask], color=COLOR_OBS, lw=1.6, label="USGS obs (gauge 03463300)", zorder=5) + + peak_usgs = np.nanmax(obs[mask]) + ax.axhline(peak_usgs, color=COLOR_OBS, lw=0.7, linestyle=":", alpha=0.5) + # Place annotation to the right of the peak to avoid legend overlap + ax.text(pd.Timestamp("2024-09-28 06:00:00"), peak_usgs * 1.02, + f"USGS peak {peak_usgs:.0f} m³/s", fontsize=8.5, color="black", + alpha=0.8, ha="right") + + _add_helene_band(ax) + _format_xaxis(ax, mdates.DayLocator(interval=1), "%b %d") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--vrugt-csv", required=True) + parser.add_argument("--novrugt-csv", required=True) + parser.add_argument("--ensemble-pq", required=True, + help="routed_crossed_ensemble.parquet from run_route_crossed_ensemble.py") + parser.add_argument("--out-dir", required=True) + args = parser.parse_args() + os.makedirs(args.out_dir, exist_ok=True) + + vrugt = load_det_csv(args.vrugt_csv) + novrugt = load_det_csv(args.novrugt_csv) + + obs = vrugt["Q_usgs_m3s"].values + dates = vrugt.index + sim_v = vrugt["Q_routed_m3s"].values + sim_nv = novrugt["Q_routed_m3s"].reindex(dates).values + + kge_v = compute_kge(obs, sim_v) + kge_nv = compute_kge(obs, sim_nv) + + print("Building ensemble envelope...") + env = build_envelope(args.ensemble_pq, ZOOM_START, ZOOM_END) + print(f" Envelope: {len(env)} timesteps, " + f"peak p95={env['p95'].max():.1f} m3/s, peak p50={env['p50'].max():.1f} m3/s") + + peak_usgs = np.nanmax(obs[(dates >= ZOOM_START) & (dates <= ZOOM_END)]) + in_env = env["p95"].max() >= peak_usgs + print(f" USGS peak {peak_usgs:.1f} m3/s {'IS' if in_env else 'IS NOT'} within p95 envelope") + + # ------------------------------------------------------------------ # + # Figure 1: Helene zoom only -- main poster panel + # ------------------------------------------------------------------ # + fig, ax = plt.subplots(figsize=(11, 5.5)) + plot_zoom(ax, dates, obs, sim_v, sim_nv, kge_v, kge_nv, env) + ax.set_xlabel("Date (UTC)", fontsize=11) + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_title( + "Hurricane Helene -- Probabilistic discharge forecast at gauge 03463300 (South Toe River)\n" + "600-member CFE ensemble + DA (Muskingum routing)", + fontsize=11) + ax.legend(fontsize=9, loc="lower right") + ax.grid(True, alpha=0.25) + plt.tight_layout() + out1 = os.path.join(args.out_dir, "helene_ensemble_vs_usgs.png") + plt.savefig(out1, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + # ------------------------------------------------------------------ # + # Figure 2: two-panel (full period + Helene zoom) + # ------------------------------------------------------------------ # + fig, (ax_top, ax_bot) = plt.subplots( + 2, 1, figsize=(14, 9), + gridspec_kw={"height_ratios": [1, 1.5]}) + + # Top: full test period (deterministic only, no ensemble for readability) + ax_top.plot(dates, obs, color=COLOR_OBS, lw=0.8, label="USGS obs", zorder=3) + ax_top.plot(dates, sim_v, color=COLOR_VRUGT, lw=0.9, + label=f"DA dynamic R (Vrugt) KGE={kge_v:.3f}", zorder=2) + ax_top.plot(dates, sim_nv, color=COLOR_NOVRUGT, lw=0.9, linestyle="--", + label=f"DA constant R KGE={kge_nv:.3f}", zorder=2) + ax_top.axvspan(ZOOM_START, ZOOM_END, color="gray", alpha=0.12, zorder=0) + ax_top.axvspan(HELENE_START, HELENE_END, color="gold", alpha=0.15, zorder=0) + ax_top.text(HELENE_START + pd.Timedelta(days=0.5), + np.nanmax(obs) * 0.88, "Helene", + fontsize=9, color="goldenrod", fontweight="bold") + ax_top.set_ylabel("Discharge (m³/s)", fontsize=10) + ax_top.set_title( + "Routed discharge at gauge 03463300 (South Toe River) -- Test period\n" + "DA with dynamic Vrugt R vs. constant R -- Muskingum routing", + fontsize=11) + ax_top.legend(fontsize=9, loc="upper left") + ax_top.grid(True, alpha=0.25) + _format_xaxis(ax_top, mdates.MonthLocator(interval=2), "%Y-%m") + + # Bottom: Helene zoom with ensemble + plot_zoom(ax_bot, dates, obs, sim_v, sim_nv, kge_v, kge_nv, env) + ax_bot.set_xlabel("Date (UTC)", fontsize=10) + ax_bot.set_ylabel("Discharge (m³/s)", fontsize=10) + ax_bot.set_title("Hurricane Helene window -- 600-member ensemble envelope", fontsize=10) + ax_bot.legend(fontsize=9, loc="lower right") + ax_bot.grid(True, alpha=0.25) + + plt.tight_layout() + out2 = os.path.join(args.out_dir, "helene_ensemble_twopanel.png") + plt.savefig(out2, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_vs_usgs.py b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_vs_usgs.py new file mode 100644 index 00000000..55c52d85 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_vs_usgs.py @@ -0,0 +1,176 @@ +#!/usr/bin/env python3 +""" +plot_routed_ensemble_vs_usgs.py +Plot T-route routed 600-member ensemble envelope vs USGS at gauge 03463300. + +Input: routed_crossed_ensemble.parquet from run_route_crossed_ensemble.py + Columns: issue_time, lead_hour, member_0000..member_0599 (Q in m3/s) + +Usage: + python3 plot_routed_ensemble_vs_usgs.py \ + --routed-pq /mnt/disk2/.../folder1_vrugt_routed/routed_crossed_ensemble.parquet \ + --usgs-csv /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv \ + --label "F1 Vrugt — 1 gauge holdout" \ + --out-dir ~/plots_f1_1gauge +""" + +import argparse +from collections import defaultdict +from pathlib import Path + +import matplotlib.dates as mdates +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd + +WATERSHED_AREA_KM2 = 113.18 # cat-1016300 drainage area +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +PLOT_START = pd.Timestamp("2024-09-24 00:00:00") +PLOT_END = pd.Timestamp("2024-09-28 18:00:00") +PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +PEAK_END = pd.Timestamp("2024-09-28 06:00:00") +HELENE_PEAK = pd.Timestamp("2024-09-27 12:00:00") + +COLOR_ENS = "#1f77b4" +COLOR_OBS = "black" + + +def build_envelope(pq_path, t_start, t_end): + df = pd.read_parquet(pq_path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = [c for c in df.columns if c.startswith("member_")] + + groups = defaultdict(list) + for _, row in df.iterrows(): + vt = (row["issue_time"] + pd.Timedelta(hours=int(row["lead_hour"]))).floor("1h") + if t_start <= vt <= t_end: + groups[vt].extend(row[member_cols].values.astype(float).tolist()) + + rows = [] + for vt in sorted(groups): + v = np.array(groups[vt]) + v = v[~np.isnan(v)] + if len(v) == 0: + continue + rows.append({"valid_time": vt, + "p05": np.percentile(v, 5), "p25": np.percentile(v, 25), + "p50": np.percentile(v, 50), "p75": np.percentile(v, 75), + "p95": np.percentile(v, 95)}) + return pd.DataFrame(rows).set_index("valid_time") + + +def load_usgs(path): + df = pd.read_csv(path) + print(f" USGS CSV columns: {list(df.columns)}") + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() or "discharge" in c.lower()) + print(f" Using date_col={date_col!r}, q_col={q_col!r}") + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series.sort_index() + + +def compute_kge(obs, sim): + mask = ~(np.isnan(obs) | np.isnan(sim)) + o, s = obs[mask], sim[mask] + if len(o) < 2 or np.std(o) == 0: + return np.nan + r = np.corrcoef(o, s)[0, 1] + return 1 - np.sqrt((r-1)**2 + (np.std(s)/np.std(o)-1)**2 + (np.mean(s)/np.mean(o)-1)**2) + + +def _shade(ax, env): + ax.fill_between(env.index, env["p05"], env["p95"], + color=COLOR_ENS, alpha=0.18, label="5th–95th percentile") + ax.fill_between(env.index, env["p25"], env["p75"], + color=COLOR_ENS, alpha=0.32, label="25th–75th percentile") + ax.plot(env.index, env["p50"], color=COLOR_ENS, lw=1.5, label="Ensemble median (600 members)") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--routed-pq", required=True) + parser.add_argument("--usgs-csv", required=True) + parser.add_argument("--label", default="F1 Vrugt — 1 gauge holdout") + parser.add_argument("--out-dir", required=True) + args = parser.parse_args() + + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + print("Building full-window envelope...") + env_full = build_envelope(args.routed_pq, PLOT_START, PLOT_END) + print("Building peak-window envelope...") + env_peak = build_envelope(args.routed_pq, PEAK_START, PEAK_END) + + print("Loading USGS obs...") + try: + obs_all = load_usgs(args.usgs_csv) + obs_full = obs_all.loc[PLOT_START:PLOT_END] + obs_peak = obs_all.loc[PEAK_START:PEAK_END] + print(f" Loaded {len(obs_all)} USGS obs rows; " + f"full window: {len(obs_full)}, peak window: {len(obs_peak)}") + except Exception as e: + print(f" Warning: USGS load failed ({e}) — plotting without obs") + obs_full = obs_peak = None + + kge = np.nan + if obs_full is not None and len(env_full) > 0: + med = env_full["p50"] + obs_hourly = obs_full.copy() + obs_hourly.index = obs_hourly.index.floor("1h") + obs_hourly = obs_hourly[~obs_hourly.index.duplicated(keep="first")] + aligned = obs_hourly.reindex(med.index) + print(f" KGE alignment: {aligned.notna().sum()} matched out of {len(med)} timesteps") + kge = compute_kge(aligned.values, med.values) + + # ── Figure 1: Full Helene window ────────────────────────────────────────── + fig, ax = plt.subplots(figsize=(12, 4.5)) + _shade(ax, env_full) + if obs_full is not None: + ax.plot(obs_full.index, obs_full.values, color=COLOR_OBS, lw=1.6, + label="USGS obs (gauge 03463300)") + ax.axvspan(PEAK_START, PEAK_END, color="salmon", alpha=0.12, zorder=0, label="Helene peak window") + ax.set_ylabel("Discharge (m³/s)") + ax.set_xlabel("Date") + ax.set_title(f"Routed 600-member ensemble vs USGS | {args.label}\n" + f"Ensemble median KGE = {kge:.3f} | Sep 24–28 2024") + ax.legend(fontsize=9, loc="upper left") + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + plt.tight_layout() + out = out_dir / "routed_ensemble_vs_usgs_full.png" + fig.savefig(out, dpi=150) + print(f"Saved: {out}") + plt.close(fig) + + # ── Figure 2: Peak zoom ─────────────────────────────────────────────────── + fig, ax = plt.subplots(figsize=(10, 4.5)) + _shade(ax, env_peak) + if obs_peak is not None: + ax.plot(obs_peak.index, obs_peak.values, color=COLOR_OBS, lw=1.6, + label="USGS obs (gauge 03463300)") + peak_val = float(obs_peak.max()) + ax.axhline(peak_val, color=COLOR_OBS, lw=0.7, ls=":", alpha=0.5) + ax.annotate(f"USGS peak\n{peak_val:.0f} m³/s", + xy=(HELENE_PEAK, peak_val), xytext=(10, -40), + textcoords="offset points", fontsize=8, color=COLOR_OBS) + ax.set_ylabel("Discharge (m³/s)") + ax.set_xlabel("Date") + ax.set_title(f"Helene peak zoom | {args.label}") + ax.legend(fontsize=9) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %HUTC")) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + plt.tight_layout() + out = out_dir / "routed_ensemble_vs_usgs_peak.png" + fig.savefig(out, dpi=150) + print(f"Saved: {out}") + plt.close(fig) + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/run_4b_crossed_f1.sh b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/run_4b_crossed_f1.sh new file mode 100644 index 00000000..ebf83e16 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/run_4b_crossed_f1.sh @@ -0,0 +1,40 @@ +#!/usr/bin/env bash +# F1 variance-scaled Vrugt (20pct gauge holdout) — 4b crossed ensemble plots. +# Runs plot_crossed_ensemble.py for all 21 catchments. +# Reads //_crossed_ensemble.parquet (already generated). +# +# Usage: +# bash run_4b_crossed_f1.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +DATA_DIR=/mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble_vrugt +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +echo "[F1-4b-crossed] Crossed ensemble plots — data: $DATA_DIR" + +for CAT in "${CATS[@]}"; do + PQ="$DATA_DIR/$CAT/${CAT}_crossed_ensemble.parquet" + if [ ! -f "$PQ" ]; then + echo " [$CAT] No crossed ensemble parquet — skipping" + continue + fi + echo " [$CAT] crossed ensemble..." + $TROUTE "$SCRIPT_DIR/plot_crossed_ensemble.py" \ + --ensemble-dir "$DATA_DIR" \ + --cat-id "$CAT" \ + --usgs-csv "$USGS_CSV" +done + +echo "[F1-4b-crossed] Done." diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/run_4b_f1.sh b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/run_4b_f1.sh new file mode 100644 index 00000000..c243781a --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4b_ensemble_vs_obs/run_4b_f1.sh @@ -0,0 +1,29 @@ +#!/usr/bin/env bash +# F1 variance-scaled Vrugt (20pct gauge holdout) — 4b routed ensemble plots. +# +# Runs: +# 1. plot_routed_ensemble_vs_usgs.py — F1 ensemble envelope vs USGS at outlet +# +# Requires routed_Q_test.csv and routed_crossed_ensemble.parquet in CROSSED_DIR. +# +# Usage: +# bash run_4b_f1.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +CROSSED_DIR=/mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble_vrugt +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv +OUT_DIR="$CROSSED_DIR" + +echo "[F1-4b] Routed ensemble vs USGS (F1 only)..." +$TROUTE "$SCRIPT_DIR/plot_routed_ensemble_vs_usgs.py" \ + --routed-pq "$CROSSED_DIR/routed_crossed_ensemble.parquet" \ + --usgs-csv "$USGS_CSV" \ + --label "F1 Variance-scaled Vrugt — 20pct gauge holdout" \ + --out-dir "$OUT_DIR" + +echo "[F1-4b] Done. Outputs in: $OUT_DIR" +ls "$OUT_DIR"/*.png 2>/dev/null || echo " (no PNGs in out dir)" diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4c_timeseries/plot_forecast_spaghetti.py b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4c_timeseries/plot_forecast_spaghetti.py new file mode 100644 index 00000000..655edaef --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4c_timeseries/plot_forecast_spaghetti.py @@ -0,0 +1,176 @@ +""" +plot_forecast_spaghetti.py + +"All forecasts ending at this hour" — verification-time-fixed spaghetti view. + +Window: Sep 24 18 UTC -> Sep 30 06 UTC. +For every initialization time whose 18-hour forecast overlaps this window: + DA : shaded min/max band across 20 members + ensemble mean line + OL : ensemble mean line only (no shading), dashed + +Lines are colored by initialization DATE (7 colors, Sep 24-30) so the +temporal progression is readable. USGS obs overlaid in black. + +Inputs: + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + Wide format: issue_time, lead_hour, member_00..member_19 (q in m³/s) + +Obs: + /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +Outputs: + /forecast_spaghetti_helene.png +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import matplotlib.patches as mpatches + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast_hardcoded_r_routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +# Verification window — only trajectories whose valid_time falls here are shown +PLOT_START = pd.Timestamp("2024-09-24 18:00:00") +PLOT_END = pd.Timestamp("2024-09-30 06:00:00") + +HELENE_PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +HELENE_PEAK_END = pd.Timestamp("2024-09-28 00:00:00") + +# One color per init date (Sep 24-30) +DATE_COLORS = { + "2024-09-24": "#1f77b4", + "2024-09-25": "#ff7f0e", + "2024-09-26": "#2ca02c", + "2024-09-27": "#d62728", + "2024-09-28": "#9467bd", + "2024-09-29": "#8c564b", + "2024-09-30": "#e377c2", +} + + +def load_parquet(path): + df = pd.read_parquet(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + df["ens_mean"] = df[member_cols].mean(axis=1) + df["ens_min"] = df[member_cols].min(axis=1) + df["ens_max"] = df[member_cols].max(axis=1) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + return df + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() + or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--route-dir", default=DEFAULT_ROUTE_DIR) + parser.add_argument("--out-dir", default=None) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--da-name", default="routed_leadtime_da_full.parquet") + parser.add_argument("--ol-name", default="routed_leadtime_openloop_full.parquet") + args = parser.parse_args() + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + print("Loading parquets...") + da = load_parquet(os.path.join(args.route_dir, args.da_name)) + ol = load_parquet(os.path.join(args.route_dir, args.ol_name)) + obs = load_usgs(args.usgs_csv) + + # Keep only valid_times inside the plot window + da = da[(da["valid_time"] >= PLOT_START) & (da["valid_time"] <= PLOT_END)] + ol = ol[(ol["valid_time"] >= PLOT_START) & (ol["valid_time"] <= PLOT_END)] + + issue_times = sorted(da["issue_time"].unique()) + print(f" {len(issue_times)} initialization times contribute to this window") + + fig, ax = plt.subplots(figsize=(18, 6)) + + # Helene peak shading + ax.axvspan(HELENE_PEAK_START, HELENE_PEAK_END, + color="salmon", alpha=0.12, zorder=0) + ax.text(HELENE_PEAK_START + pd.Timedelta(hours=6), 1.0, "Helene peak", + transform=ax.get_xaxis_transform(), + fontsize=9, color="firebrick", ha="left", va="top") + + # Plot each initialization time's trajectory + for t0 in issue_times: + date_str = str(pd.Timestamp(t0).date()) + color = DATE_COLORS.get(date_str, "gray") + + da_t = da[da["issue_time"] == t0].sort_values("valid_time") + ol_t = ol[ol["issue_time"] == t0].sort_values("valid_time") + + if da_t.empty: + continue + + # DA: shaded band (min-max across 20 members) + ensemble mean + ax.fill_between(da_t["valid_time"], da_t["ens_min"], da_t["ens_max"], + color=color, alpha=0.06, zorder=2) + ax.plot(da_t["valid_time"], da_t["ens_mean"], + color=color, lw=0.8, alpha=0.55, zorder=3) + + # OL: ensemble mean only, dashed + if not ol_t.empty: + ax.plot(ol_t["valid_time"], ol_t["ens_mean"], + color=color, lw=0.6, alpha=0.30, linestyle="--", zorder=2) + + # USGS obs + obs_w = obs.loc[PLOT_START:PLOT_END] + ax.plot(obs_w.index, obs_w.values, + color="black", lw=1.8, zorder=6, label="USGS obs") + + # Legend: one patch per init date + obs + DA/OL style + patches = [mpatches.Patch(color=c, label=f"Init {d}") + for d, c in DATE_COLORS.items()] + patches.append(plt.Line2D([0], [0], color="black", lw=1.8, label="USGS obs")) + patches.append(plt.Line2D([0], [0], color="gray", lw=1.2, alpha=0.7, + label="DA — ens mean (solid) + spread (shaded)")) + patches.append(plt.Line2D([0], [0], color="gray", lw=1.0, linestyle="--", + alpha=0.5, label="OL — ens mean (dashed)")) + ax.legend(handles=patches, fontsize=8, loc="upper left", + frameon=True, framealpha=0.9, ncol=2) + + ax.set_xlim(PLOT_START, PLOT_END) + ax.set_xlabel("Date", fontsize=11) + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_title( + "All 18-hour forecast trajectories ending in this window — USGS 03463300\n" + "Sep 24 18 UTC → Sep 30 06 UTC | DA: shaded band + mean | OL: mean dashed", + fontsize=11, + ) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=1)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.xaxis.set_minor_locator(mdates.HourLocator(interval=6)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.2, lw=0.4) + plt.tight_layout() + + out_path = os.path.join(out_dir, "forecast_spaghetti_helene.png") + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4c_timeseries/plot_helene_issue_time_hydrograph.py b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4c_timeseries/plot_helene_issue_time_hydrograph.py new file mode 100644 index 00000000..4e2e1dc2 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4c_timeseries/plot_helene_issue_time_hydrograph.py @@ -0,0 +1,198 @@ +""" +Per-issue-time forecast hydrograph diagnostic. + +For two contrasting issue times — one storm-peak (Helene) and one low-flow — +show what the DA and open-loop ensembles actually produce over the 18-hour +forecast window, alongside the kriging observation. Designed to answer in one +picture: is DA over-shooting, collapsing, or oscillating compared to open-loop? + +Each panel shows: + - All 20 DA members as faint purple lines + median (thick purple) + - All 20 open-loop members as faint gray lines + median (thick gray, dashed) + - Kriging obs as a thick black line (context: t0-6 through t0+18) + - Vertical line at t0 + annotation with obs(t0) and lead-1 ensemble means + +Inputs (from run_lead_time_forecast_sweep.py): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv +Obs (from production EnKF): + //_test_results.csv + +Output: + //_issue_time_hydrograph_helene_vs_lowflow.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +DEFAULT_LEADTIME_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast" +DEFAULT_DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" + +# Module-level globals set by main() +CAT = "cat-1016300" +LEADTIME_DIR = DEFAULT_LEADTIME_DIR +DA_DIR = DEFAULT_DA_DIR + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" +OBS_COLOR = "black" + +# Two issue times — pick a Helene-peak one and a typical low-flow one +HELENE_T0 = pd.Timestamp("2024-09-26 12:00:00") +LOWFLOW_T0 = pd.Timestamp("2024-03-15 12:00:00") + +CONTEXT_HOURS_BEFORE = 6 # hours of obs context shown before t0 +LEAD_HOURS_AFTER = 18 # forecast horizon + +OUT_PNG = os.path.join( + LEADTIME_DIR, CAT, + f"{CAT}_issue_time_hydrograph_helene_vs_lowflow.png", +) + + +def load_forecasts(path): + df = pd.read_csv(path, parse_dates=['issue_time', 'valid_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + return df, member_cols + + +def load_obs(): + obs_path = os.path.join(DA_DIR, CAT, f"{CAT}_test_results.csv") + df = pd.read_csv(obs_path, parse_dates=['date']) + return df.set_index('date')['obs_mm_h'] + + +def slice_forecast(df, member_cols, t0): + """Return (valid_times, member_array shape (lead, N)) for a single issue time.""" + sub = df[df['issue_time'] == t0].sort_values('lead_hour') + if len(sub) == 0: + raise ValueError(f"No forecast rows for issue_time={t0}") + valid_times = pd.to_datetime(sub['valid_time'].values) + member_arr = sub[member_cols].to_numpy(dtype=float) + return valid_times, member_arr + + +def plot_panel(ax, t0, df_da, df_ol, member_cols, obs_series, title_extra=""): + # Forecast trajectories + vt_da, q_da = slice_forecast(df_da, member_cols, t0) + vt_ol, q_ol = slice_forecast(df_ol, member_cols, t0) + + # Obs context window + ctx_start = t0 - pd.Timedelta(hours=CONTEXT_HOURS_BEFORE) + ctx_end = t0 + pd.Timedelta(hours=LEAD_HOURS_AFTER) + obs_window = obs_series.loc[ctx_start:ctx_end] + + # Open-loop members (drawn first so DA paints on top) + for j in range(q_ol.shape[1]): + ax.plot(vt_ol, q_ol[:, j], color=OL_COLOR, lw=0.6, alpha=0.45, zorder=2) + # DA members + for j in range(q_da.shape[1]): + ax.plot(vt_da, q_da[:, j], color=DA_COLOR, lw=0.6, alpha=0.45, zorder=3) + # Medians + ax.plot(vt_ol, np.nanmedian(q_ol, axis=1), color=OL_COLOR, lw=2.6, + linestyle="--", zorder=4, label="Open-loop median") + ax.plot(vt_da, np.nanmedian(q_da, axis=1), color=DA_COLOR, lw=2.6, + zorder=5, label="DA median") + # Obs + ax.plot(obs_window.index, obs_window.values, color=OBS_COLOR, lw=1.8, + zorder=6, label="Qkrig (obs)") + + # t0 marker + annotation + ax.axvline(t0, color="black", lw=1.0, linestyle=":", alpha=0.6, zorder=1) + obs_t0 = float(obs_series.get(t0, np.nan)) + obs_lead1 = float(obs_series.get(t0 + pd.Timedelta(hours=1), np.nan)) + da_lead1_mean = float(np.nanmean(q_da[0, :])) + ol_lead1_mean = float(np.nanmean(q_ol[0, :])) + + info_text = ( + f"t0 = {t0.strftime('%Y-%m-%d %H:%M')}\n" + f"obs(t0) = {obs_t0:.3f} mm/h\n" + f"obs(t0+1) = {obs_lead1:.3f} mm/h\n" + f"DA mean(t0+1) = {da_lead1_mean:.3f} mm/h\n" + f"OL mean(t0+1) = {ol_lead1_mean:.3f} mm/h" + ) + ax.text(0.02, 0.97, info_text, + transform=ax.transAxes, + fontsize=9, va='top', ha='left', family='monospace', + bbox=dict(boxstyle='round,pad=0.4', facecolor='white', + alpha=0.85, edgecolor='gray')) + + ax.set_title(title_extra, fontsize=12) + ax.set_ylabel("Discharge (mm/h)", fontsize=10) + ax.set_xlabel("Date", fontsize=10) + ax.xaxis.set_major_locator(mdates.HourLocator(interval=3)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %H:%M")) + ax.tick_params(axis='x', rotation=30, labelsize=8) + ax.grid(True, alpha=0.2) + ax.legend(loc='upper right', fontsize=9, frameon=True, framealpha=0.9) + + +def snap_to_nearest(t0, available_times): + """Return the available issue_time closest to t0.""" + times = pd.DatetimeIndex(available_times) + idx = np.abs((times - t0).total_seconds()).argmin() + snapped = times[idx] + if snapped != t0: + print(f" Snapped {t0} -> {snapped} (nearest available issue time)") + return snapped + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', default="cat-1016300") + parser.add_argument('--leadtime-dir', default=DEFAULT_LEADTIME_DIR, + help='Dir holding /_lead_time_forecasts_{da,openloop}.csv') + parser.add_argument('--da-dir', default=DEFAULT_DA_DIR, + help='Dir holding /_test_results.csv (for obs)') + parser.add_argument('--helene-t0', default=None, + help='Issue time for the Helene panel (default: 2024-09-26 12:00:00, ' + 'snapped to nearest available)') + parser.add_argument('--lowflow-t0', default=None, + help='Issue time for the low-flow panel (default: 2024-03-15 00:00:00, ' + 'snapped to nearest available)') + args = parser.parse_args() + + global CAT, LEADTIME_DIR, DA_DIR, OUT_PNG + CAT = args.cat_id + LEADTIME_DIR = args.leadtime_dir + DA_DIR = args.da_dir + OUT_PNG = os.path.join(LEADTIME_DIR, CAT, + f"{CAT}_issue_time_hydrograph_helene_vs_lowflow.png") + + da_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_da.csv") + ol_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_openloop.csv") + df_da, m_cols = load_forecasts(da_path) + df_ol, _ = load_forecasts(ol_path) + obs_series = load_obs() + + available = df_da['issue_time'].unique() + helene_t0 = snap_to_nearest( + pd.Timestamp(args.helene_t0) if args.helene_t0 else HELENE_T0, available) + lowflow_t0 = snap_to_nearest( + pd.Timestamp(args.lowflow_t0) if args.lowflow_t0 else pd.Timestamp("2024-03-15 00:00:00"), + available) + + fig, (ax_h, ax_l) = plt.subplots(1, 2, figsize=(20, 7)) + + plot_panel(ax_h, helene_t0, df_da, df_ol, m_cols, obs_series, + title_extra=f"Helene peak issue time — {helene_t0.strftime('%Y-%m-%d %H:%M')}") + plot_panel(ax_l, lowflow_t0, df_da, df_ol, m_cols, obs_series, + title_extra=f"Low-flow issue time — {lowflow_t0.strftime('%Y-%m-%d %H:%M')}") + + fig.suptitle( + f"Per-issue-time forecast hydrograph — {CAT}\n" + "Faint lines = 20 ensemble members. Thick lines = ensemble medians. " + "Vertical dotted line = t0 (DA stops here for both scenarios).", + fontsize=12, y=1.00, + ) + plt.tight_layout() + plt.savefig(OUT_PNG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4c_timeseries/plot_lead_time_reconstructed_timeseries.py b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4c_timeseries/plot_lead_time_reconstructed_timeseries.py new file mode 100644 index 00000000..a6d03b03 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4c_timeseries/plot_lead_time_reconstructed_timeseries.py @@ -0,0 +1,263 @@ +""" +Reconstructed time series at gauge 03463300, from our routed +lead-time forecasts. Mirrors the methodology of plot_timeseries_hourly_ensemble.py +(qSpatialAR) but feeds from the EnKF + T-route pipeline instead of the CNN. + +For each target hour t, pool ALL forecasts that land on t (across many +issue_times × many lead_hours × 20 members). Compute median + 5th/95th +percentile envelope. Plot DA and open-loop reconstructions overlaid on USGS +obs, with the Helene window shaded. + +Input parquets (from run_route_troute_leadtime.py): + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + Wide format: issue_time, lead_hour, member_00..member_19 (q in m³/s) + +Obs: + /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +Output: + /lead_time_reconstructed_timeseries.png + Default plot window: 2024-09-10 → 2024-10-10. + +Each scenario's curve is annotated with NSE vs USGS obs over the plot window. +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" +OBS_COLOR = "black" + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast_hardcoded_r_routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +# Watershed area at gauge 03463300 (sum of 21 catchments per GPKG hydrofabric) +WATERSHED_AREA_KM2 = 113.18 +# mm/h depth → m³/s: × 113.18 km² × 1000 / 3600 = 31.439 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +PLOT_START_DEFAULT = pd.Timestamp("2024-09-10 00:00:00") +PLOT_END_DEFAULT = pd.Timestamp("2024-10-10 23:00:00") + +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-30 23:00:00") + + +def load_parquet_long(path): + df = pd.read_parquet(path) + df['issue_time'] = pd.to_datetime(df['issue_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + if not member_cols: + raise ValueError(f"No member columns in {path}") + long = df.melt( + id_vars=['issue_time', 'lead_hour'], + value_vars=member_cols, + var_name='member', + value_name='q_m3s', + ) + long['valid_time'] = (long['issue_time'] + + pd.to_timedelta(long['lead_hour'], unit='h')) + return long + + +def reconstruct(long_df): + """For each valid_time, pool all (issue_time × lead × member) forecasts and + return median, p05, p95, count.""" + grouped = long_df.groupby('valid_time')['q_m3s'].agg( + median='median', + p05=lambda s: s.quantile(0.05), + p95=lambda s: s.quantile(0.95), + count='count', + ).reset_index() + return grouped + + +def load_usgs(usgs_csv): + """Load USGS gauge obs and ensure m³/s. + + The CSV at gauge 03463300 has the column `QObs(mm/h)` — catchment-averaged + depth, not gauge discharge in m³/s. Auto-convert when column name contains + 'mm': m³/s = mm/h × WATERSHED_AREA_KM2 × 1000 / 3600. + """ + df = pd.read_csv(usgs_csv) + date_col = next((c for c in df.columns + if c.lower() in ('datetime', 'date', 'time', 'timestamp')), None) + q_col = next((c for c in df.columns + if 'q' in c.lower() or 'flow' in c.lower() + or 'discharge' in c.lower()), None) + if not date_col or not q_col: + raise ValueError(f"Couldn't id date/Q columns in {usgs_csv}: {df.columns.tolist()}") + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if 'mm' in q_col.lower(): + series = series * MM_H_TO_M3_S + print(f" Converted obs from mm/h (column '{q_col}') to m³/s " + f"using area = {WATERSHED_AREA_KM2} km² (× {MM_H_TO_M3_S:.4f})") + else: + print(f" Loaded obs as-is from column '{q_col}' (assumed m³/s)") + return series + + +def nse(obs, sim): + m = np.isfinite(obs) & np.isfinite(sim) + if m.sum() < 2: + return np.nan + o, s = obs[m], sim[m] + denom = ((o - o.mean()) ** 2).sum() + if denom < 1e-10: + return np.nan + return 1.0 - float(((o - s) ** 2).sum() / denom) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--route-dir', default=DEFAULT_ROUTE_DIR) + parser.add_argument('--out-dir', default=None, + help='Default: --route-dir') + parser.add_argument('--usgs-csv', default=DEFAULT_USGS_CSV) + parser.add_argument('--plot-start', default=str(PLOT_START_DEFAULT)) + parser.add_argument('--plot-end', default=str(PLOT_END_DEFAULT)) + parser.add_argument('--da-name', default='routed_leadtime_da_full.parquet') + parser.add_argument('--ol-name', default='routed_leadtime_openloop_full.parquet') + args = parser.parse_args() + + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + plot_start = pd.Timestamp(args.plot_start) + plot_end = pd.Timestamp(args.plot_end) + + print(f"Loading DA parquet: {args.da_name}") + da_long = load_parquet_long(os.path.join(args.route_dir, args.da_name)) + print(f"Loading OL parquet: {args.ol_name}") + ol_long = load_parquet_long(os.path.join(args.route_dir, args.ol_name)) + print(f"Loading USGS obs: {args.usgs_csv}") + obs = load_usgs(args.usgs_csv) + + print("Reconstructing DA time series (overlapping-leads pool)...") + da_rec = reconstruct(da_long) + print(f" {len(da_rec)} unique target hours") + print("Reconstructing OL time series...") + ol_rec = reconstruct(ol_long) + + # Clip to plot window and join obs + da_rec = da_rec[(da_rec['valid_time'] >= plot_start) & (da_rec['valid_time'] <= plot_end)].copy() + ol_rec = ol_rec[(ol_rec['valid_time'] >= plot_start) & (ol_rec['valid_time'] <= plot_end)].copy() + da_rec['obs'] = da_rec['valid_time'].map(obs) + ol_rec['obs'] = ol_rec['valid_time'].map(obs) + + nse_da = nse(da_rec['obs'].values, da_rec['median'].values) + nse_ol = nse(ol_rec['obs'].values, ol_rec['median'].values) + print(f"NSE (plot window) — DA: {nse_da:.3f} | OL: {nse_ol:.3f}") + + # ---- Plot ---- + fig, ax = plt.subplots(figsize=(16, 5.5)) + fig.patch.set_facecolor("white") + + # Helene shaded + ax.axvspan(HELENE_START, HELENE_END, + color="firebrick", alpha=0.08, zorder=0) + ax.text(HELENE_START + pd.Timedelta(hours=12), 1.0, "Helene", + transform=ax.get_xaxis_transform(), + fontsize=9, color="firebrick", ha="left", va="top", + fontweight="bold") + + # Open-loop band + median + ax.fill_between(ol_rec['valid_time'], ol_rec['p05'], ol_rec['p95'], + color=OL_COLOR, alpha=0.18, zorder=2, + label="Open-loop [5th–95th pctile envelope]") + ax.plot(ol_rec['valid_time'], ol_rec['median'], + color=OL_COLOR, lw=1.2, linestyle='--', zorder=3, + label=f"Open-loop median (NSE={nse_ol:.3f})") + + # DA band + median + ax.fill_between(da_rec['valid_time'], da_rec['p05'], da_rec['p95'], + color=DA_COLOR, alpha=0.22, zorder=4, + label="DA [5th–95th pctile envelope]") + ax.plot(da_rec['valid_time'], da_rec['median'], + color=DA_COLOR, lw=1.4, zorder=5, + label=f"DA median (NSE={nse_da:.3f})") + + # Observed as dots + obs_window = obs.loc[plot_start:plot_end] + finite = np.isfinite(obs_window.values) + ax.scatter(obs_window.index[finite], obs_window.values[finite], + s=6, color=OBS_COLOR, linewidths=0, zorder=6, + label="USGS obs") + + ax.set_xlabel("Date", fontsize=11) + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_title( + f"Reconstructed forecast time series at USGS 03463300\n" + f"Overlapping-leads pool from EnKF forecast ensemble " + f"({plot_start.date()} – {plot_end.date()})", + fontsize=12, + ) + ax.set_xlim(plot_start, plot_end) + ax.xaxis.set_major_locator(mdates.WeekdayLocator(interval=1)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.25, lw=0.4) + ax.legend(loc='upper left', fontsize=9, frameon=True, framealpha=0.92) + + plt.tight_layout() + out_path = os.path.join(out_dir, "lead_time_reconstructed_timeseries.png") + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + # Also produce a Helene-zoom version + helene_zoom_start = pd.Timestamp("2024-09-24 00:00:00") + helene_zoom_end = pd.Timestamp("2024-09-29 23:00:00") + da_z = da_rec[(da_rec['valid_time'] >= helene_zoom_start) + & (da_rec['valid_time'] <= helene_zoom_end)] + ol_z = ol_rec[(ol_rec['valid_time'] >= helene_zoom_start) + & (ol_rec['valid_time'] <= helene_zoom_end)] + nse_da_z = nse(da_z['obs'].values, da_z['median'].values) + nse_ol_z = nse(ol_z['obs'].values, ol_z['median'].values) + print(f"NSE (Helene zoom) — DA: {nse_da_z:.3f} | OL: {nse_ol_z:.3f}") + + fig, ax = plt.subplots(figsize=(14, 5.5)) + fig.patch.set_facecolor("white") + ax.fill_between(ol_z['valid_time'], ol_z['p05'], ol_z['p95'], + color=OL_COLOR, alpha=0.18, zorder=2, + label="Open-loop [5th–95th]") + ax.plot(ol_z['valid_time'], ol_z['median'], color=OL_COLOR, + lw=1.4, linestyle='--', zorder=3, + label=f"Open-loop median (NSE={nse_ol_z:.3f})") + ax.fill_between(da_z['valid_time'], da_z['p05'], da_z['p95'], + color=DA_COLOR, alpha=0.22, zorder=4, + label="DA [5th–95th]") + ax.plot(da_z['valid_time'], da_z['median'], color=DA_COLOR, + lw=1.6, zorder=5, + label=f"DA median (NSE={nse_da_z:.3f})") + obs_z = obs.loc[helene_zoom_start:helene_zoom_end] + finite = np.isfinite(obs_z.values) + ax.scatter(obs_z.index[finite], obs_z.values[finite], + s=12, color=OBS_COLOR, linewidths=0, zorder=6, + label="USGS obs") + ax.set_xlabel("Date", fontsize=11) + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_title( + "Reconstructed forecast time series at USGS 03463300 — Helene zoom\n" + "Sep 24 – 29, 2024 (overlapping-leads pool from EnKF forecast ensemble)", + fontsize=12, + ) + ax.xaxis.set_major_locator(mdates.HourLocator(interval=12)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %H:%M")) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.25, lw=0.4) + ax.legend(loc='upper left', fontsize=9, frameon=True, framealpha=0.92) + plt.tight_layout() + out_zoom = os.path.join(out_dir, "lead_time_reconstructed_timeseries_helene.png") + plt.savefig(out_zoom, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_zoom}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4c_timeseries/run_4c_f1.sh b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4c_timeseries/run_4c_f1.sh new file mode 100644 index 00000000..28fbd3f0 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/4c_timeseries/run_4c_f1.sh @@ -0,0 +1,51 @@ +#!/usr/bin/env bash +# F1 variance-scaled Vrugt (20pct gauge holdout) — 4c reconstructed timeseries plots. +# +# Reads routed_leadtime_da_full.parquet + routed_leadtime_openloop_full.parquet +# from the F1 leadtime route dir and writes output PNGs there. +# Must run AFTER route_leadtime_f1.sh. +# +# Usage: +# bash run_4c_f1.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +LEADTIME_DIR=/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast +ROUTE_DIR=/mnt/disk2/suma_helen_poster/leadtime_troute_routing +OBS_DIR=/mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +echo "[F1-4c] Reconstructed timeseries from: $ROUTE_DIR" +$TROUTE "$SCRIPT_DIR/plot_lead_time_reconstructed_timeseries.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" + +echo "[F1-4c] Forecast spaghetti..." +$TROUTE "$SCRIPT_DIR/plot_forecast_spaghetti.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" + +echo "[F1-4c] Helene vs low-flow issue-time hydrograph (per catchment)..." +for CAT in "${CATS[@]}"; do + DA_CSV="$LEADTIME_DIR/$CAT/${CAT}_lead_time_forecasts_da.csv" + [ -f "$DA_CSV" ] || { echo " [$CAT] no leadtime CSV — skipping"; continue; } + echo " [$CAT] helene hydrograph..." + $TROUTE "$SCRIPT_DIR/plot_helene_issue_time_hydrograph.py" \ + --cat-id "$CAT" \ + --leadtime-dir "$LEADTIME_DIR" \ + --obs-dir "$OBS_DIR" +done + +echo "[F1-4c] Done. Outputs in: $ROUTE_DIR and $LEADTIME_DIR" +ls "$ROUTE_DIR"/*.png 2>/dev/null || echo " (no PNGs in route dir)" diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/compare_da_vs_qkrig_vs_usgs.py b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/compare_da_vs_qkrig_vs_usgs.py new file mode 100644 index 00000000..ae7bc03d --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/compare_da_vs_qkrig_vs_usgs.py @@ -0,0 +1,162 @@ +""" +compare_da_vs_qkrig_vs_usgs.py + +Three-way comparison at gauge 03463300: + DA-routed Q vs routed Qkrig vs USGS obs + +Reads the routed_Q_test.csv produced by run_route.py (which already contains +Q_routed_m3s, Q_usgs_m3s, and Q_krig_m3s columns) and prints KGE/NSE/peak +for the full test period and the Helene window separately. + +Also saves a two-panel comparison figure. + +Usage: + python3 compare_da_vs_qkrig_vs_usgs.py \ + --vrugt-csv /mnt/disk2/suma_helen_poster/da_results/vrugt_dynamic_routed/routed_Q_test.csv \ + --out-dir /mnt/disk2/suma_helen_poster/da_results/comparison_plots +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-29 23:00:00") + + +def kge(obs, sim): + mask = np.isfinite(obs) & np.isfinite(sim) + o, s = obs[mask], sim[mask] + if len(o) < 2 or np.std(o) == 0: + return np.nan + r = np.corrcoef(o, s)[0, 1] + return 1.0 - np.sqrt((r-1)**2 + (np.std(s)/np.std(o)-1)**2 + (np.mean(s)/np.mean(o)-1)**2) + + +def nse(obs, sim): + mask = np.isfinite(obs) & np.isfinite(sim) + o, s = obs[mask], sim[mask] + if len(o) < 2: + return np.nan + denom = np.sum((o - o.mean())**2) + return 1.0 - np.sum((o - s)**2) / denom if denom > 0 else np.nan + + +def peak_ratio(obs, sim): + return np.nanmax(sim) / np.nanmax(obs) + + +def print_stats(label, obs, sim, dates, window_name="full period"): + mask = np.isfinite(obs) & np.isfinite(sim) + print(f" {label:35s} KGE={kge(obs,sim):+.3f} NSE={nse(obs,sim):+.3f} " + f"peak_sim={np.nanmax(sim):.1f} peak_obs={np.nanmax(obs):.1f} " + f"ratio={peak_ratio(obs,sim):.2f}x [{window_name}]") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--vrugt-csv", required=True, + help="routed_Q_test.csv from vrugt_dynamic_routed/") + parser.add_argument("--out-dir", required=True) + args = parser.parse_args() + os.makedirs(args.out_dir, exist_ok=True) + + df = pd.read_csv(args.vrugt_csv, parse_dates=["date"]).set_index("date").sort_index() + + if "Q_usgs_m3s" not in df.columns: + raise ValueError("CSV missing Q_usgs_m3s — run run_route.py with --obs-csv") + if "Q_krig_m3s" not in df.columns: + raise ValueError("CSV missing Q_krig_m3s — need run_route.py with --kv-dir or --obs-csv") + + obs = df["Q_usgs_m3s"].values + da = df["Q_routed_m3s"].values + krig = df["Q_krig_m3s"].values + dates = df.index + + helene = (dates >= HELENE_START) & (dates <= HELENE_END) + + print("\n" + "="*85) + print("THREE-WAY COMPARISON — gauge 03463300 (South Toe River)") + print("="*85) + + print("\n[FULL TEST PERIOD]") + print_stats("DA-routed vs USGS", obs, da, dates, "full") + print_stats("Qkrig-routed vs USGS", obs, krig, dates, "full") + print_stats("DA-routed vs Qkrig-routed", krig, da, dates, "full") + + print("\n[HELENE WINDOW Sep 24–29]") + print_stats("DA-routed vs USGS", obs[helene], da[helene], dates[helene], "Helene") + print_stats("Qkrig-routed vs USGS", obs[helene], krig[helene], dates[helene], "Helene") + print_stats("DA-routed vs Qkrig-routed", krig[helene], da[helene], dates[helene], "Helene") + + print("\n[PEAK VALUES (Helene window)]") + print(f" USGS peak : {np.nanmax(obs[helene]):.1f} m³/s") + print(f" DA-routed peak : {np.nanmax(da[helene]):.1f} m³/s " + f"({np.nanmax(da[helene])/np.nanmax(obs[helene])*100:.0f}% of USGS)") + print(f" Qkrig-routed peak: {np.nanmax(krig[helene]):.1f} m³/s " + f"({np.nanmax(krig[helene])/np.nanmax(obs[helene])*100:.0f}% of USGS)") + print(f" DA vs Qkrig peak : DA is " + f"{'higher' if np.nanmax(da[helene]) > np.nanmax(krig[helene]) else 'lower'} " + f"by {abs(np.nanmax(da[helene])-np.nanmax(krig[helene])):.1f} m³/s") + print("="*85 + "\n") + + # ---- Plot ---- + fig, axes = plt.subplots(2, 1, figsize=(15, 9), + gridspec_kw={"height_ratios": [1, 1.6]}) + + # Top: full period + ax = axes[0] + ax.plot(dates, obs, color="black", lw=0.8, label="USGS obs", zorder=4) + ax.plot(dates, krig, color="#e6820e", lw=0.8, linestyle="--", + label=f"Qkrig-routed KGE={kge(obs,da):.3f}", zorder=2) + ax.plot(dates, da, color="#1f77b4", lw=0.9, + label=f"DA-routed KGE={kge(obs,da):.3f}", zorder=3) + ax.axvspan(HELENE_START, HELENE_END, color="gold", alpha=0.15, zorder=0) + ax.set_ylabel("Discharge (m³/s)", fontsize=10) + ax.set_title("Full test period — DA-routed vs Qkrig-routed vs USGS obs", fontsize=11) + ax.legend(fontsize=9, loc="upper left") + ax.grid(True, alpha=0.25) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%Y-%m")) + ax.xaxis.set_major_locator(mdates.MonthLocator(interval=2)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right", fontsize=8) + + # Bottom: Helene zoom + ax = axes[1] + dh = dates[helene] + kge_da_h = kge(obs[helene], da[helene]) + kge_krig_h = kge(obs[helene], krig[helene]) + nse_da_h = nse(obs[helene], da[helene]) + nse_krig_h = nse(obs[helene], krig[helene]) + + ax.plot(dh, obs[helene], color="black", lw=1.8, zorder=4, label="USGS obs") + ax.plot(dh, krig[helene], color="#e6820e", lw=1.4, linestyle="--", zorder=2, + label=f"Qkrig-routed KGE={kge_krig_h:.3f} NSE={nse_krig_h:.3f}") + ax.plot(dh, da[helene], color="#1f77b4", lw=1.6, zorder=3, + label=f"DA-routed KGE={kge_da_h:.3f} NSE={nse_da_h:.3f}") + + peak_usgs = np.nanmax(obs[helene]) + ax.axhline(peak_usgs, color="black", lw=0.6, linestyle=":", alpha=0.5) + ax.text(HELENE_END - pd.Timedelta(hours=6), peak_usgs * 1.02, + f"USGS peak {peak_usgs:.0f} m³/s", fontsize=8.5, ha="right", color="black") + + ax.set_ylabel("Discharge (m³/s)", fontsize=10) + ax.set_xlabel("Date (UTC)", fontsize=10) + ax.set_title("Helene window — does DA add value beyond Qkrig?", fontsize=11) + ax.legend(fontsize=9, loc="upper left") + ax.grid(True, alpha=0.25) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=1)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right", fontsize=9) + + plt.tight_layout() + out_path = os.path.join(args.out_dir, "da_vs_qkrig_vs_usgs.png") + plt.savefig(out_path, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/plot_da_perturbation_arms.py b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/plot_da_perturbation_arms.py new file mode 100644 index 00000000..001be9da --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/plot_da_perturbation_arms.py @@ -0,0 +1,375 @@ +""" +plot_da_perturbation_arms.py — 2a/2b: ensemble spread WITH DA on. + +Loads the two arm CSVs produced by run_perturbation_da_on.py: + //_da_forcing_arm.csv (30 members) + //_da_hydro_arm.csv (20 members) + +Plots: + Panel 2a — Forcing arm: 30-member spaghetti over Helene window + (spread = met forcing uncertainty with DA-corrected initial states) + Panel 2b — Hydro-state arm: 20-member spaghetti over Helene window + (spread = initial state uncertainty with deterministic forcing) + Optional: --ol-csv overlays the open-loop grand median as a + thick dashed gray line for comparison. + Panel 2c — Comparison: median ± spread envelope, both arms + USGS obs + +All trajectories projected to valid_time = issue_time + lead_hour hours. +Colored by initialization date (Sep 24-30). USGS obs in black. + +Outputs: + //_2a_forcing_arm_helene.png + //_2b_hydro_arm_helene.png + //_2ab_arms_comparison.png +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import matplotlib.patches as mpatches + +DEFAULT_ARM_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_perturbation_da_on" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" +DEFAULT_CAT_ID = "cat-1016300" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +PLOT_START = pd.Timestamp("2024-09-24 00:00:00") +PLOT_END = pd.Timestamp("2024-09-30 06:00:00") +HELENE_PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +HELENE_PEAK_END = pd.Timestamp("2024-09-28 00:00:00") + +DATE_COLORS = { + "2024-09-24": "#1f77b4", + "2024-09-25": "#ff7f0e", + "2024-09-26": "#2ca02c", + "2024-09-27": "#d62728", + "2024-09-28": "#9467bd", + "2024-09-29": "#8c564b", + "2024-09-30": "#e377c2", +} + + +def load_arm(path): + """Load arm CSV, compute valid_time, member columns.""" + df = pd.read_csv(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + return df, member_cols + + +def load_openloop(path): + """Load open-loop lead-time file (CSV or Parquet) and return grand median + indexed by valid_time. + + Accepts: + - routed_leadtime_openloop_full.parquet (T-route output, m³/s, recommended) + - cat-*_lead_time_forecasts_openloop.csv (unrouted CFE mm/h, single catchment) + + Returns a pd.Series (valid_time → median q in m³/s) clipped to plot window. + """ + if path.endswith(".parquet"): + df = pd.read_parquet(path) + else: + df = pd.read_csv(path) + + df["issue_time"] = pd.to_datetime(df["issue_time"]) + if "valid_time" in df.columns: + df["valid_time"] = pd.to_datetime(df["valid_time"]) + elif "lead_hour" in df.columns: + df["valid_time"] = (df["issue_time"] + + pd.to_timedelta(df["lead_hour"], unit="h")) + else: + raise ValueError("Open-loop file must have valid_time or lead_hour column") + + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + mask = (df["valid_time"] >= PLOT_START) & (df["valid_time"] <= PLOT_END) + df = df[mask].copy() + if df.empty: + return pd.Series(dtype=float) + + vals = df[member_cols].to_numpy(dtype=float) + # Only convert if values are clearly in mm/h (unrouted CFE output). + # Routed parquet is already in m³/s — do not convert. + if path.endswith(".csv") and np.nanmedian(vals[vals > 0]) < 5: + vals = vals * MM_H_TO_M3_S + + df["q_grand_median"] = np.nanmedian(vals, axis=1) + series = (df.groupby("valid_time")["q_grand_median"] + .median() + .sort_index()) + return series + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def _add_helene_band(ax): + ax.axvspan(HELENE_PEAK_START, HELENE_PEAK_END, + color="salmon", alpha=0.12, zorder=0) + ax.text(HELENE_PEAK_START + pd.Timedelta(hours=6), 1.0, "Helene peak", + transform=ax.get_xaxis_transform(), + fontsize=9, color="firebrick", ha="left", va="top") + + +def _format_xaxis(ax): + ax.set_xlim(PLOT_START, PLOT_END) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=1)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.xaxis.set_minor_locator(mdates.HourLocator(interval=6)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.2, lw=0.4) + + +def plot_spaghetti_arm(ax, df, member_cols, arm_color, obs_series, + title, label_stem, lw_thin=0.7, alpha_thin=0.35): + """ + Plot one thin trajectory per (issue_time, member) pair. + Each trajectory's x = valid_time values for that issue_time. + """ + _add_helene_band(ax) + + issue_times = sorted(df["issue_time"].unique()) + all_means = [] + + for t0 in issue_times: + date_str = str(pd.Timestamp(t0).date()) + color = DATE_COLORS.get(date_str, arm_color) + sub = df[df["issue_time"] == t0].sort_values("valid_time") + vt = sub["valid_time"].values + # Filter to plot window + mask = (sub["valid_time"] >= PLOT_START) & (sub["valid_time"] <= PLOT_END) + sub_w = sub[mask] + if sub_w.empty: + continue + vt_w = sub_w["valid_time"].values + + mem_vals = sub_w[member_cols].to_numpy(dtype=float) # shape (T, N_mem) + if "mm" not in label_stem.lower(): + mem_vals = mem_vals * MM_H_TO_M3_S # mm/h -> m³/s + + # Min-max envelope + median line + ax.fill_between(vt_w, + np.nanmin(mem_vals, axis=1), + np.nanmax(mem_vals, axis=1), + color=color, alpha=0.06, zorder=2) + ax.plot(vt_w, np.nanmedian(mem_vals, axis=1), + color=color, lw=lw_thin, alpha=alpha_thin + 0.1, zorder=3) + + all_means.append( + pd.Series(np.nanmedian(mem_vals, axis=1), index=vt_w)) + + # Overall mean across all issue times (interpolated to common grid) + if all_means: + full_idx = pd.date_range(PLOT_START, PLOT_END, freq="1h") + stacked = pd.concat(all_means, axis=1).reindex(full_idx) + grand_mean = stacked.mean(axis=1) + ax.plot(grand_mean.index, grand_mean.values, + color=arm_color, lw=2.4, alpha=0.95, zorder=5, + label=f"{label_stem} — mean across all forecasts") + + # USGS obs + obs_w = obs_series.loc[PLOT_START:PLOT_END] + ax.plot(obs_w.index, obs_w.values, + color="black", lw=1.8, zorder=6, label="USGS obs") + + ax.set_ylabel("Discharge (m³/s)", fontsize=10) + ax.set_title(title, fontsize=11) + _format_xaxis(ax) + + +def compute_envelope(df, member_cols): + """ + Aggregate all members across all issue_times by valid_time. + Returns DataFrame indexed by valid_time with columns: q_min, q_med, q_max. + """ + records = [] + for t0, grp in df.groupby("issue_time"): + mask = (grp["valid_time"] >= PLOT_START) & (grp["valid_time"] <= PLOT_END) + sub = grp[mask] + if sub.empty: + continue + vals = sub[member_cols].to_numpy(dtype=float) * MM_H_TO_M3_S + for i, row in enumerate(sub.itertuples()): + records.append({ + "valid_time": row.valid_time, + "q_min": np.nanmin(vals[i]), + "q_med": np.nanmedian(vals[i]), + "q_max": np.nanmax(vals[i]), + }) + if not records: + return pd.DataFrame(columns=["valid_time", "q_min", "q_med", "q_max"]) + + env_df = pd.DataFrame(records) + env_df = (env_df.groupby("valid_time") + .agg(q_min=("q_min", "min"), + q_med=("q_med", "mean"), + q_max=("q_max", "max")) + .reset_index() + .sort_values("valid_time")) + return env_df + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--arm-dir", default=DEFAULT_ARM_DIR) + parser.add_argument("--cat-id", default=DEFAULT_CAT_ID) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--out-dir", default=None) + parser.add_argument( + "--ol-csv", default=None, + help=( + "Path to open-loop lead-time forecast CSV " + "(e.g. cat-1016300_lead_time_forecasts_openloop.csv). " + "When provided, the grand median is overlaid on panel 2b " + "as a thick dashed gray line." + ), + ) + args = parser.parse_args() + + cat_dir = os.path.join(args.arm_dir, args.cat_id) + out_dir = args.out_dir or cat_dir + os.makedirs(out_dir, exist_ok=True) + + forcing_path = os.path.join(cat_dir, f"{args.cat_id}_da_forcing_arm.csv") + hydro_path = os.path.join(cat_dir, f"{args.cat_id}_da_hydro_arm.csv") + + print(f"Loading arm CSVs for {args.cat_id}...") + df_fa, fa_cols = load_arm(forcing_path) + df_ha, ha_cols = load_arm(hydro_path) + obs = load_usgs(args.usgs_csv) + print(f" Forcing arm: {df_fa['issue_time'].nunique()} issue times, " + f"{len(fa_cols)} members") + print(f" Hydro arm: {df_ha['issue_time'].nunique()} issue times, " + f"{len(ha_cols)} members") + + # ------------------------------------------------------------------ # + # Figure 2a: Forcing arm spaghetti # + # ------------------------------------------------------------------ # + fig, ax = plt.subplots(figsize=(17, 6)) + plot_spaghetti_arm( + ax, df_fa, fa_cols, + arm_color="tab:blue", obs_series=obs, + title=(f"2a — Forcing arm (DA on, {len(fa_cols)} members): " + "met forcing uncertainty | Sep 24-30 2024\n" + f"{args.cat_id} | Shaded = member min-max | " + "Line = median per init time | Thick = grand mean"), + label_stem="Forcing arm", + ) + # Date-color legend + patches = [mpatches.Patch(color=c, label=f"Init {d}") + for d, c in DATE_COLORS.items()] + handles, labels = ax.get_legend_handles_labels() + ax.legend(handles=handles + patches, + fontsize=8, loc="upper left", frameon=True, framealpha=0.9, ncol=2) + plt.tight_layout() + out_a = os.path.join(out_dir, f"{args.cat_id}_2a_forcing_arm_helene.png") + plt.savefig(out_a, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_a}") + + # ------------------------------------------------------------------ # + # Figure 2b: Hydro-state arm spaghetti # + # ------------------------------------------------------------------ # + fig, ax = plt.subplots(figsize=(17, 6)) + plot_spaghetti_arm( + ax, df_ha, ha_cols, + arm_color="tab:green", obs_series=obs, + title=(f"2b — Hydro-state arm (DA on, {len(ha_cols)} members): " + "initial state uncertainty | Sep 24-30 2024\n" + f"{args.cat_id} | Shaded = member min-max | " + "Line = median per init time | Thick = grand mean"), + label_stem="Hydro-state arm", + ) + + # Optional open-loop overlay + if args.ol_csv: + print(f"Loading open-loop CSV: {args.ol_csv}") + ol_series = load_openloop(args.ol_csv) + if not ol_series.empty: + ax.plot( + ol_series.index, ol_series.values, + color="black", lw=2.5, ls="--", alpha=0.95, zorder=7, + label="Open loop (no DA) — grand median", + ) + else: + print(" Warning: open-loop CSV yielded no data in plot window.") + + patches = [mpatches.Patch(color=c, label=f"Init {d}") + for d, c in DATE_COLORS.items()] + handles, labels = ax.get_legend_handles_labels() + ax.legend(handles=handles + patches, + fontsize=8, loc="upper left", frameon=True, framealpha=0.9, ncol=2) + plt.tight_layout() + out_b = os.path.join(out_dir, f"{args.cat_id}_2b_hydro_arm_helene.png") + plt.savefig(out_b, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_b}") + + # ------------------------------------------------------------------ # + # Figure 2ab: Comparison — envelopes overlaid, both arms # + # ------------------------------------------------------------------ # + print("Computing aggregated envelopes for comparison plot...") + env_fa = compute_envelope(df_fa, fa_cols) + env_ha = compute_envelope(df_ha, ha_cols) + + fig, ax = plt.subplots(figsize=(17, 7)) + _add_helene_band(ax) + + if not env_fa.empty: + ax.fill_between(env_fa["valid_time"], + env_fa["q_min"], env_fa["q_max"], + color="tab:blue", alpha=0.18, zorder=2, + label=f"Forcing arm spread (N={len(fa_cols)} members)") + ax.plot(env_fa["valid_time"], env_fa["q_med"], + color="tab:blue", lw=2.0, zorder=4, + label="Forcing arm — grand median") + + if not env_ha.empty: + ax.fill_between(env_ha["valid_time"], + env_ha["q_min"], env_ha["q_max"], + color="tab:green", alpha=0.18, zorder=2, + label=f"Hydro-state arm spread (N={len(ha_cols)} members)") + ax.plot(env_ha["valid_time"], env_ha["q_med"], + color="tab:green", lw=2.0, zorder=4, + label="Hydro-state arm — grand median") + + obs_w = obs.loc[PLOT_START:PLOT_END] + ax.plot(obs_w.index, obs_w.values, + color="black", lw=1.8, zorder=6, label="USGS obs") + + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_xlabel("Date", fontsize=11) + ax.set_title( + f"2a vs 2b — Forcing arm (blue) vs Hydro-state arm (green) | " + f"DA on | {args.cat_id}\n" + "Shaded = full member spread (min-max aggregated across all issue times). " + "Lines = grand median.", + fontsize=11, + ) + ax.legend(fontsize=9, loc="upper left", frameon=True, framealpha=0.9) + _format_xaxis(ax) + plt.tight_layout() + out_c = os.path.join(out_dir, f"{args.cat_id}_2ab_arms_comparison.png") + plt.savefig(out_c, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_c}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/plot_da_v2_vrugt_helene_grid.py b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/plot_da_v2_vrugt_helene_grid.py new file mode 100644 index 00000000..f7857bd3 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/plot_da_v2_vrugt_helene_grid.py @@ -0,0 +1,98 @@ +""" +Helene comparison grid: Run 3 (no DA) vs DA v2 (true EnKF + Vrugt R) vs Qkrig obs. +3 x 7 grid of all 21 catchments, zoomed to Sep 20 - Oct 5, 2024. + +Produces: helene_da_v2_vrugt_vs_run3_grid.png + +Expects per-catchment *_test_results.csv (columns: date, sim_mm_h, obs_mm_h) +in RUN3_DIR and DA_DIR. Output files are created by +calibrate_catchment_cfe_da_v2.py run_testing_period(). + +Run inside the same conda env that has pandas + matplotlib (e.g. troute). +""" +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import os + +CATS = [ + "cat-1016279", "cat-1016280", "cat-1016281", "cat-1016282", "cat-1016283", + "cat-1016300", "cat-1016301", "cat-1016302", "cat-1016303", "cat-1016304", + "cat-1016305", "cat-1016306", "cat-1016307", "cat-1016308", "cat-1016309", + "cat-1016310", "cat-1016311", "cat-1016312", "cat-1016313", "cat-1016314", + "cat-1016315", +] + +RUN3_DIR = "/mnt/disk2/suma_helen_poster/catchment_results_range100_run3" +DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" +OUT_PNG = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt/helene_da_v2_vrugt_vs_run3_grid.png" +DA_LABEL = "DA v2 (true EnKF + Vrugt R, N=20)" + +HELENE_START = pd.Timestamp("2024-09-20") +HELENE_END = pd.Timestamp("2024-10-05") + + +def kge(obs, sim): + """Kling-Gupta Efficiency on full overlap of obs and sim.""" + m = ~(pd.isna(obs) | pd.isna(sim)) + o, s = obs[m], sim[m] + if len(o) < 2: + return float('nan') + denom = (((o - o.mean()) ** 2).sum() * ((s - s.mean()) ** 2).sum()) ** 0.5 + if denom == 0: + return float('nan') + r = ((o - o.mean()) * (s - s.mean())).sum() / denom + alpha = s.std() / o.std() + beta = s.mean() / o.mean() if o.mean() != 0 else float('nan') + return 1 - ((r - 1) ** 2 + (alpha - 1) ** 2 + (beta - 1) ** 2) ** 0.5 + + +def main(): + fig, axes = plt.subplots(3, 7, figsize=(28, 12), sharex=True) + axes = axes.flatten() + + for i, cat in enumerate(CATS): + ax = axes[i] + run3_csv = os.path.join(RUN3_DIR, cat, f"{cat}_test_results.csv") + da_csv = os.path.join(DA_DIR, cat, f"{cat}_test_results.csv") + + if not (os.path.exists(run3_csv) and os.path.exists(da_csv)): + ax.set_title(f"{cat} (missing)") + continue + + df_run3 = pd.read_csv(run3_csv, parse_dates=["date"]) + df_da = pd.read_csv(da_csv, parse_dates=["date"]) + + mask3 = (df_run3["date"] >= HELENE_START) & (df_run3["date"] <= HELENE_END) + maskd = (df_da["date"] >= HELENE_START) & (df_da["date"] <= HELENE_END) + + kge_run3 = kge(df_run3["obs_mm_h"], df_run3["sim_mm_h"]) + kge_da = kge(df_da["obs_mm_h"], df_da["sim_mm_h"]) + + ax.plot(df_da.loc[maskd, "date"], df_da.loc[maskd, "obs_mm_h"], + color="black", lw=1.8, label="Qkrig (obs)", zorder=1) + ax.plot(df_run3.loc[mask3, "date"], df_run3.loc[mask3, "sim_mm_h"], + color="steelblue", lw=4.5, alpha=0.45, label=f"Run 3 KGE={kge_run3:.2f}", zorder=2) + ax.plot(df_da.loc[maskd, "date"], df_da.loc[maskd, "sim_mm_h"], + color="tomato", lw=1.3, label=f"DA KGE={kge_da:.2f}", zorder=3) + + ax.set_title(f"{cat}", fontsize=10) + ax.tick_params(axis="x", rotation=45, labelsize=7) + ax.tick_params(axis="y", labelsize=8) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=4)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.legend(fontsize=6, loc="upper right") + + fig.suptitle( + f"Hurricane Helene (Sep 20 - Oct 5, 2024) - {DA_LABEL} vs Run 3 baseline", + fontsize=14, y=1.00, + ) + fig.text(0.005, 0.5, "Discharge (mm/h)", va="center", rotation="vertical", fontsize=12) + plt.tight_layout() + plt.savefig(OUT_PNG, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/plot_da_v2_vrugt_helene_zoomed_grid.py b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/plot_da_v2_vrugt_helene_zoomed_grid.py new file mode 100644 index 00000000..1a3a94b6 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/plot_da_v2_vrugt_helene_zoomed_grid.py @@ -0,0 +1,100 @@ +""" +Hurricane Helene PEAK zoom: Run 3 (no DA) vs DA v2 (true EnKF + Vrugt R) vs Qkrig obs. +3 x 7 grid of all 21 catchments, zoomed to the 4-day peak window Sep 24 - Sep 28, 2024. + +Same data as plot_da_v2_vrugt_helene_grid.py, just a tighter x-axis to see +the peak detail. + +Produces: helene_da_v2_vrugt_vs_run3_grid_zoomed.png + +Run inside the same conda env that has pandas + matplotlib (e.g. troute). +""" +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import os + +CATS = [ + "cat-1016279", "cat-1016280", "cat-1016281", "cat-1016282", "cat-1016283", + "cat-1016300", "cat-1016301", "cat-1016302", "cat-1016303", "cat-1016304", + "cat-1016305", "cat-1016306", "cat-1016307", "cat-1016308", "cat-1016309", + "cat-1016310", "cat-1016311", "cat-1016312", "cat-1016313", "cat-1016314", + "cat-1016315", +] + +RUN3_DIR = "/mnt/disk2/suma_helen_poster/catchment_results_range100_run3" +DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" +OUT_PNG = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt/helene_da_v2_vrugt_vs_run3_grid_zoomed.png" +DA_LABEL = "DA v2 (true EnKF + Vrugt R, N=20)" + +# Tighter 4-day peak window +ZOOM_START = pd.Timestamp("2024-09-24") +ZOOM_END = pd.Timestamp("2024-09-28") + + +def kge(obs, sim): + """Kling-Gupta Efficiency on full overlap of obs and sim.""" + m = ~(pd.isna(obs) | pd.isna(sim)) + o, s = obs[m], sim[m] + if len(o) < 2: + return float('nan') + denom = (((o - o.mean()) ** 2).sum() * ((s - s.mean()) ** 2).sum()) ** 0.5 + if denom == 0: + return float('nan') + r = ((o - o.mean()) * (s - s.mean())).sum() / denom + alpha = s.std() / o.std() + beta = s.mean() / o.mean() if o.mean() != 0 else float('nan') + return 1 - ((r - 1) ** 2 + (alpha - 1) ** 2 + (beta - 1) ** 2) ** 0.5 + + +def main(): + fig, axes = plt.subplots(3, 7, figsize=(28, 12), sharex=True) + axes = axes.flatten() + + for i, cat in enumerate(CATS): + ax = axes[i] + run3_csv = os.path.join(RUN3_DIR, cat, f"{cat}_test_results.csv") + da_csv = os.path.join(DA_DIR, cat, f"{cat}_test_results.csv") + + if not (os.path.exists(run3_csv) and os.path.exists(da_csv)): + ax.set_title(f"{cat} (missing)") + continue + + df_run3 = pd.read_csv(run3_csv, parse_dates=["date"]) + df_da = pd.read_csv(da_csv, parse_dates=["date"]) + + mask3 = (df_run3["date"] >= ZOOM_START) & (df_run3["date"] <= ZOOM_END) + maskd = (df_da["date"] >= ZOOM_START) & (df_da["date"] <= ZOOM_END) + + # KGE on full test period (matches log-line value); displayed in legend + kge_run3 = kge(df_run3["obs_mm_h"], df_run3["sim_mm_h"]) + kge_da = kge(df_da["obs_mm_h"], df_da["sim_mm_h"]) + + ax.plot(df_da.loc[maskd, "date"], df_da.loc[maskd, "obs_mm_h"], + color="black", lw=1.8, label="Qkrig (obs)", zorder=1) + ax.plot(df_run3.loc[mask3, "date"], df_run3.loc[mask3, "sim_mm_h"], + color="steelblue", lw=4.0, alpha=0.45, label=f"Run 3 KGE={kge_run3:.2f}", zorder=2) + ax.plot(df_da.loc[maskd, "date"], df_da.loc[maskd, "sim_mm_h"], + color="tomato", lw=1.4, label=f"DA KGE={kge_da:.2f}", zorder=3) + + ax.set_title(f"{cat}", fontsize=10) + ax.tick_params(axis="x", rotation=45, labelsize=7) + ax.tick_params(axis="y", labelsize=8) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=1)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.xaxis.set_minor_locator(mdates.HourLocator(byhour=[6, 12, 18])) + ax.legend(fontsize=6, loc="upper right") + + fig.suptitle( + f"Hurricane Helene peak zoom (Sep 24 - Sep 28, 2024) - {DA_LABEL} vs Run 3 baseline", + fontsize=14, y=1.00, + ) + fig.text(0.005, 0.5, "Discharge (mm/h)", va="center", rotation="vertical", fontsize=12) + plt.tight_layout() + plt.savefig(OUT_PNG, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/plot_perturbation_sensitivity_spaghetti.py b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/plot_perturbation_sensitivity_spaghetti.py new file mode 100644 index 00000000..e46de257 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/plot_perturbation_sensitivity_spaghetti.py @@ -0,0 +1,157 @@ +""" +Plot perturbation-source sensitivity spaghetti plots. + +Reads the three per-source sensitivity CSVs per catchment and plots all 20 +member streamflow trajectories as colored bundles overlaid on one panel, +zoomed to the Hurricane Helene peak window (Sep 24 - Sep 28, 2024). + +Color legend: + red = initial state perturbation only + blue = forcing perturbation only (precip + PET) + green = process noise on hydrologic states only + +Produces: + helene_sensitivity_spaghetti_main.png (3-catchment subset for the main figure) + helene_sensitivity_spaghetti_appendix.png (full 3 x 7 grid of all 21 catchments) + +Requires the sensitivity runs to have completed first (see +run_perturbation_sensitivity.py). +""" +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import os + +# ---------- Configuration ---------- +ALL_CATS = [ + "cat-1016279", "cat-1016280", "cat-1016281", "cat-1016282", "cat-1016283", + "cat-1016300", "cat-1016301", "cat-1016302", "cat-1016303", "cat-1016304", + "cat-1016305", "cat-1016306", "cat-1016307", "cat-1016308", "cat-1016309", + "cat-1016310", "cat-1016311", "cat-1016312", "cat-1016313", "cat-1016314", + "cat-1016315", +] + +# Main figure: 3 catchments spanning the Run 3 KGE distribution. +# cat-1016311 = worst Run 3 (0.50); cat-1016300 = median (0.74); cat-1016302 = best (0.83) +MAIN_CATS = ["cat-1016311", "cat-1016300", "cat-1016302"] + +DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" # for Qkrig obs +SEN_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_sensitivity" # sensitivity CSVs +MAIN_PNG = "/mnt/disk2/suma_helen_poster/da_results/v2_sensitivity/helene_sensitivity_spaghetti_main.png" +APPENDIX_PNG = "/mnt/disk2/suma_helen_poster/da_results/v2_sensitivity/helene_sensitivity_spaghetti_appendix.png" + +# Helene 4-day peak zoom +ZOOM_START = pd.Timestamp("2024-09-24") +ZOOM_END = pd.Timestamp("2024-09-28") + +SOURCE_COLORS = { + "init": "tab:red", + "forcing": "tab:blue", + "process": "tab:green", +} +SOURCE_LABELS = { + "init": "Initial state only", + "forcing": "Forcing only (P, PET)", + "process": "Process noise on states only", +} + + +def load_sensitivity_csv(cat, source): + """Returns (dates, q_matrix) for the test period (member columns).""" + path = os.path.join(SEN_DIR, cat, f"{cat}_sensitivity_{source}.csv") + if not os.path.exists(path): + return None, None + df = pd.read_csv(path, parse_dates=["date"]) + member_cols = [c for c in df.columns if c.startswith("member_")] + return df["date"].values, df[member_cols].values + + +def load_qkrig_obs(cat): + """Loads the Qkrig observation timeseries from the production-run CSV.""" + path = os.path.join(DA_DIR, cat, f"{cat}_test_results.csv") + if not os.path.exists(path): + return None, None + df = pd.read_csv(path, parse_dates=["date"]) + return df["date"].values, df["obs_mm_h"].values + + +def plot_one_panel(ax, cat): + """Render the three colored bundles plus the obs line for one catchment.""" + # Each source: 20 thin colored lines + for source, color in SOURCE_COLORS.items(): + dates, q = load_sensitivity_csv(cat, source) + if dates is None: + continue + df_dates = pd.to_datetime(dates) + mask = (df_dates >= ZOOM_START) & (df_dates <= ZOOM_END) + if mask.sum() == 0: + continue + # Plot all 20 members as thin transparent lines, plus a thicker mean line + for i in range(q.shape[1]): + ax.plot(df_dates[mask], q[mask, i], + color=color, lw=0.5, alpha=0.35, zorder=1) + ax.plot(df_dates[mask], q[mask].mean(axis=1), + color=color, lw=1.6, alpha=0.95, zorder=2, + label=SOURCE_LABELS[source]) + + # Black observation overlay + obs_dates, obs = load_qkrig_obs(cat) + if obs_dates is not None: + od = pd.to_datetime(obs_dates) + mask = (od >= ZOOM_START) & (od <= ZOOM_END) + ax.plot(od[mask], obs[mask], + color="black", lw=1.4, label="Qkrig (obs)", zorder=3) + + ax.set_title(cat, fontsize=10) + ax.tick_params(axis="x", rotation=45, labelsize=7) + ax.tick_params(axis="y", labelsize=8) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=1)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + + +def make_main_figure(): + """3-catchment side-by-side panel for the main paper / poster figure.""" + fig, axes = plt.subplots(1, 3, figsize=(18, 5), sharex=True) + for ax, cat in zip(axes, MAIN_CATS): + plot_one_panel(ax, cat) + # Single legend on the rightmost panel (avoid clutter on the others) + axes[-1].legend(fontsize=8, loc="upper right") + axes[0].set_ylabel("Discharge (mm/h)", fontsize=11) + fig.suptitle( + "Ensemble spread by perturbation source — Hurricane Helene peak (Sep 24-28, 2024)\n" + "20 members per source. DA off. Worst / median / best Run 3 catchment.", + fontsize=12, y=1.02, + ) + plt.tight_layout() + plt.savefig(MAIN_PNG, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {MAIN_PNG}") + + +def make_appendix_figure(): + """Full 3 x 7 grid of all 21 catchments.""" + fig, axes = plt.subplots(3, 7, figsize=(28, 12), sharex=True) + axes = axes.flatten() + for i, cat in enumerate(ALL_CATS): + plot_one_panel(axes[i], cat) + # Single legend on the first panel + axes[0].legend(fontsize=7, loc="upper right") + fig.text(0.005, 0.5, "Discharge (mm/h)", va="center", rotation="vertical", fontsize=12) + fig.suptitle( + "Ensemble spread by perturbation source - Hurricane Helene peak (Sep 24-28, 2024) - all 21 sub-catchments\n" + "Red = initial state only | Blue = forcing only | Green = process noise only | Black = Qkrig obs", + fontsize=14, y=1.00, + ) + plt.tight_layout() + plt.savefig(APPENDIX_PNG, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {APPENDIX_PNG}") + + +def main(): + make_main_figure() + make_appendix_figure() + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/plot_perturbation_sensitivity_spread.py b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/plot_perturbation_sensitivity_spread.py new file mode 100644 index 00000000..cd771c6d --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/plot_perturbation_sensitivity_spread.py @@ -0,0 +1,162 @@ +""" +Plot ensemble-spread time series, one curve per perturbation source. + +For each catchment and each of the three sensitivity sub-experiments, compute the +hourly std-dev of streamflow across the 20 members. Plot all three as colored +lines on the same axes so the relative contribution of each source is directly +comparable hour by hour. + +This is the cleanest answer to "which perturbation source contributes most to +ensemble spread" — the spaghetti version is hard to read when the bundles +overlap. A single std-dev curve per source removes the visual clutter. + +Color legend: + red = std-dev across initial-state-only ensemble + blue = std-dev across forcing-only ensemble + green = std-dev across process-noise-only ensemble + +Produces: + helene_sensitivity_spread_main.png (3-catchment subset for the main figure) + helene_sensitivity_spread_appendix.png (full 3 x 7 grid of all 21 catchments) + +Consumes the same per-source CSVs produced by run_perturbation_sensitivity.py. +""" +import pandas as pd +import numpy as np +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import os + +ALL_CATS = [ + "cat-1016279", "cat-1016280", "cat-1016281", "cat-1016282", "cat-1016283", + "cat-1016300", "cat-1016301", "cat-1016302", "cat-1016303", "cat-1016304", + "cat-1016305", "cat-1016306", "cat-1016307", "cat-1016308", "cat-1016309", + "cat-1016310", "cat-1016311", "cat-1016312", "cat-1016313", "cat-1016314", + "cat-1016315", +] + +MAIN_CATS = ["cat-1016311", "cat-1016300", "cat-1016302"] + +DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" +SEN_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_sensitivity" +MAIN_PNG = "/mnt/disk2/suma_helen_poster/da_results/v2_sensitivity/helene_sensitivity_spread_main.png" +APPENDIX_PNG = "/mnt/disk2/suma_helen_poster/da_results/v2_sensitivity/helene_sensitivity_spread_appendix.png" + +ZOOM_START = pd.Timestamp("2024-09-24") +ZOOM_END = pd.Timestamp("2024-09-28") + +SOURCE_COLORS = { + "init": "tab:red", + "forcing": "tab:blue", + "process": "tab:green", +} +SOURCE_LABELS = { + "init": "Initial state perturbation", + "forcing": "Forcing perturbation (P, PET)", + "process": "Process noise on states", +} + + +def load_member_spread(cat, source): + """Returns (dates, std_per_hour) for the test period. + + std_per_hour is the std-dev of the 20 member streamflow values at each hour. + """ + path = os.path.join(SEN_DIR, cat, f"{cat}_sensitivity_{source}.csv") + if not os.path.exists(path): + return None, None + df = pd.read_csv(path, parse_dates=["date"]) + member_cols = [c for c in df.columns if c.startswith("member_")] + member_arr = df[member_cols].to_numpy(dtype=float) + # Hourly std-dev across members (sample std, ddof=1 to match Pyy convention) + std = member_arr.std(axis=1, ddof=1) + return df["date"].values, std + + +def load_qkrig_obs(cat): + """Loads the Qkrig observation timeseries from the production-run CSV.""" + path = os.path.join(DA_DIR, cat, f"{cat}_test_results.csv") + if not os.path.exists(path): + return None, None + df = pd.read_csv(path, parse_dates=["date"]) + return df["date"].values, df["obs_mm_h"].values + + +def plot_one_panel(ax, cat, show_obs=True): + """Render std-dev time series for the three sources plus optional obs reference.""" + for source, color in SOURCE_COLORS.items(): + dates, std = load_member_spread(cat, source) + if dates is None: + continue + d = pd.to_datetime(dates) + mask = (d >= ZOOM_START) & (d <= ZOOM_END) + if mask.sum() == 0: + continue + ax.plot(d[mask], std[mask], color=color, lw=1.6, + label=SOURCE_LABELS[source], zorder=2) + + if show_obs: + # Plot the observation on a secondary y-axis for context (storm timing) + obs_dates, obs = load_qkrig_obs(cat) + if obs_dates is not None: + od = pd.to_datetime(obs_dates) + mask = (od >= ZOOM_START) & (od <= ZOOM_END) + ax2 = ax.twinx() + ax2.plot(od[mask], obs[mask], + color="black", lw=1.0, alpha=0.35, zorder=1, label="Qkrig (obs, ref)") + ax2.set_ylabel("Qkrig (mm/h)", fontsize=8, color="0.4") + ax2.tick_params(axis="y", labelsize=7, colors="0.4") + ax2.spines["right"].set_color("0.7") + + ax.set_title(cat, fontsize=10) + ax.tick_params(axis="x", rotation=45, labelsize=7) + ax.tick_params(axis="y", labelsize=8) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=1)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.set_ylabel("Ensemble std-dev (mm/h)", fontsize=8) + + +def make_main_figure(): + fig, axes = plt.subplots(1, 3, figsize=(18, 5), sharex=True) + for ax, cat in zip(axes, MAIN_CATS): + plot_one_panel(ax, cat, show_obs=True) + axes[-1].legend(fontsize=8, loc="upper right") + fig.suptitle( + "Ensemble spread (std-dev across 20 members) by perturbation source\n" + "Hurricane Helene peak (Sep 24-28, 2024) | DA off | " + "worst / median / best Run 3 catchment\n" + "Grey line on right axis = Qkrig observation (for storm-timing context only)", + fontsize=11, y=1.06, + ) + plt.tight_layout() + plt.savefig(MAIN_PNG, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {MAIN_PNG}") + + +def make_appendix_figure(): + fig, axes = plt.subplots(3, 7, figsize=(28, 12), sharex=True) + axes = axes.flatten() + for i, cat in enumerate(ALL_CATS): + plot_one_panel(axes[i], cat, show_obs=True) + axes[0].legend(fontsize=7, loc="upper right") + fig.suptitle( + "Ensemble spread (hourly std-dev across 20 members) by perturbation source - " + "Hurricane Helene peak (Sep 24-28, 2024) - all 21 sub-catchments\n" + "Red = init state | Blue = forcing | Green = process noise | " + "Grey on right axis = Qkrig obs (storm-timing context only)", + fontsize=13, y=1.00, + ) + plt.tight_layout() + plt.savefig(APPENDIX_PNG, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {APPENDIX_PNG}") + + +def main(): + make_main_figure() + make_appendix_figure() + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/plot_vrugt_comparison.py b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/plot_vrugt_comparison.py new file mode 100644 index 00000000..d241bc04 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/4_evaluation/plot_vrugt_comparison.py @@ -0,0 +1,188 @@ +#!/usr/bin/env python3 +""" +plot_vrugt_comparison.py +Compare routed discharge: dynamic Vrugt R vs. constant R (no-Vrugt) +against USGS gauge obs at 03463300 (South Toe River Near Celo, NC). + +Reads routed_Q_test.csv from both run directories (produced by run_route.py +with --obs-csv). Produces: + - Two-panel figure: full test period + Helene zoom + - Single-panel Helene zoom only + +Usage: + python3 plot_vrugt_comparison.py \ + --vrugt-csv /mnt/disk2/suma_helen_poster/da_results/vrugt_dynamic_routed/routed_Q_test.csv \ + --novrugt-csv /mnt/disk2/suma_helen_poster/da_results/novrugt_dynamic_routed/routed_Q_test.csv \ + --out-dir /mnt/disk2/suma_helen_poster/da_results/comparison_plots +""" + +import argparse +import os + +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import matplotlib.patches as mpatches + +HELENE_START = pd.Timestamp("2024-09-25 00:00:00") +HELENE_END = pd.Timestamp("2024-09-29 00:00:00") +ZOOM_START = pd.Timestamp("2024-09-24 00:00:00") +ZOOM_END = pd.Timestamp("2024-09-30 00:00:00") + +COLOR_OBS = "black" +COLOR_VRUGT = "#1f77b4" # blue +COLOR_NOVRUGT = "#d62728" # red + + +def load_csv(path): + df = pd.read_csv(path, parse_dates=["date"]) + df = df.set_index("date").sort_index() + return df + + +def compute_kge(obs, sim): + mask = ~(np.isnan(obs) | np.isnan(sim)) + o, s = obs[mask], sim[mask] + if len(o) < 2 or np.std(o) == 0: + return np.nan + r = np.corrcoef(o, s)[0, 1] + return 1.0 - np.sqrt((r - 1)**2 + (np.std(s) / np.std(o) - 1)**2 + (np.mean(s) / np.mean(o) - 1)**2) + + +def compute_nse(obs, sim): + mask = ~(np.isnan(obs) | np.isnan(sim)) + o, s = obs[mask], sim[mask] + if len(o) < 2 or np.sum((o - np.mean(o))**2) == 0: + return np.nan + return 1.0 - np.sum((o - s)**2) / np.sum((o - np.mean(o))**2) + + +def _add_helene_band(ax): + ax.axvspan(HELENE_START, HELENE_END, + color="gold", alpha=0.18, zorder=0, label="_nolegend_") + + +def _format_xaxis(ax, locator, fmt, rotation=30): + ax.xaxis.set_major_locator(locator) + ax.xaxis.set_major_formatter(mdates.DateFormatter(fmt)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=rotation, ha="right", fontsize=9) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--vrugt-csv", required=True) + parser.add_argument("--novrugt-csv", required=True) + parser.add_argument("--out-dir", required=True) + args = parser.parse_args() + os.makedirs(args.out_dir, exist_ok=True) + + vrugt = load_csv(args.vrugt_csv) + novrugt = load_csv(args.novrugt_csv) + + # USGS obs from vrugt CSV (same timestamps, both runs used same --obs-csv) + obs = vrugt["Q_usgs_m3s"].values + dates = vrugt.index + + sim_v = vrugt["Q_routed_m3s"].values + sim_nv = novrugt["Q_routed_m3s"].reindex(dates).values + + kge_v = compute_kge(obs, sim_v) + nse_v = compute_nse(obs, sim_v) + kge_nv = compute_kge(obs, sim_nv) + nse_nv = compute_nse(obs, sim_nv) + + peak_usgs = np.nanmax(obs) + peak_v = np.nanmax(sim_v) + peak_nv = np.nanmax(sim_nv) + + print(f"USGS peak: {peak_usgs:.1f} m3/s") + print(f"Vrugt peak: {peak_v:.1f} m3/s KGE={kge_v:.3f} NSE={nse_v:.3f}") + print(f"No-Vrugt peak: {peak_nv:.1f} m3/s KGE={kge_nv:.3f} NSE={nse_nv:.3f}") + + # ------------------------------------------------------------------ # + # Figure 1: two-panel (full period + Helene zoom) + # ------------------------------------------------------------------ # + fig, (ax_full, ax_zoom) = plt.subplots( + 2, 1, figsize=(14, 9), + gridspec_kw={"height_ratios": [1, 1.4]}) + + # -- Top: full period -- + ax_full.plot(dates, obs, color=COLOR_OBS, lw=0.8, label="USGS obs (gauge 03463300)", zorder=3) + ax_full.plot(dates, sim_v, color=COLOR_VRUGT, lw=0.9, label=f"Vrugt (dynamic R) KGE={kge_v:.3f}", zorder=2) + ax_full.plot(dates, sim_nv, color=COLOR_NOVRUGT, lw=0.9, linestyle="--", + label=f"No-Vrugt (const R) KGE={kge_nv:.3f}", zorder=2) + _add_helene_band(ax_full) + ax_full.text(HELENE_START + pd.Timedelta(hours=12), ax_full.get_ylim()[1] * 0.85, + "Helene", fontsize=9, color="goldenrod", fontweight="bold") + ax_full.set_ylabel("Discharge (m³/s)", fontsize=10) + ax_full.set_title( + "Routed discharge at gauge 03463300 (South Toe River)\n" + "DA with dynamic Vrugt R vs. constant R -- Muskingum routing", + fontsize=11) + ax_full.legend(fontsize=9, loc="upper left") + ax_full.grid(True, alpha=0.25) + _format_xaxis(ax_full, mdates.MonthLocator(interval=2), "%Y-%m") + + # Add zoom indicator + ax_full.axvspan(ZOOM_START, ZOOM_END, color="gray", alpha=0.10, zorder=0) + ax_full.annotate("", xy=(ZOOM_END, ax_full.get_ylim()[1] * 0.5), + xytext=(ZOOM_START, ax_full.get_ylim()[1] * 0.5), + arrowprops=dict(arrowstyle="<->", color="gray", lw=1.0)) + + # -- Bottom: Helene zoom -- + mask_zoom = (dates >= ZOOM_START) & (dates <= ZOOM_END) + dz = dates[mask_zoom] + ax_zoom.plot(dz, obs[mask_zoom], color=COLOR_OBS, lw=1.2, label="USGS obs (gauge 03463300)", zorder=3) + ax_zoom.plot(dz, sim_v[mask_zoom], color=COLOR_VRUGT, lw=1.4, + label=f"Vrugt (dynamic R) KGE={kge_v:.3f} NSE={nse_v:.3f} peak={peak_v:.0f} m³/s", zorder=2) + ax_zoom.plot(dz, sim_nv[mask_zoom], color=COLOR_NOVRUGT, lw=1.4, linestyle="--", + label=f"No-Vrugt (const R) KGE={kge_nv:.3f} NSE={nse_nv:.3f} peak={peak_nv:.0f} m³/s", zorder=2) + _add_helene_band(ax_zoom) + ax_zoom.axhline(peak_usgs, color=COLOR_OBS, lw=0.7, linestyle=":", alpha=0.6) + ax_zoom.text(ZOOM_START + pd.Timedelta(hours=3), peak_usgs * 1.01, + f"USGS peak {peak_usgs:.0f} m³/s", fontsize=8, color="black", alpha=0.8) + ax_zoom.set_xlabel("Date (UTC)", fontsize=10) + ax_zoom.set_ylabel("Discharge (m³/s)", fontsize=10) + ax_zoom.set_title("Hurricane Helene window (Sep 24-30, 2024)", fontsize=10) + ax_zoom.legend(fontsize=9, loc="lower right") + ax_zoom.grid(True, alpha=0.25) + _format_xaxis(ax_zoom, mdates.DayLocator(interval=1), "%b %d") + + plt.tight_layout() + out1 = os.path.join(args.out_dir, "vrugt_vs_novrugt_twopanel.png") + plt.savefig(out1, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + # ------------------------------------------------------------------ # + # Figure 2: Helene zoom only (poster-ready single panel) + # ------------------------------------------------------------------ # + fig, ax = plt.subplots(figsize=(10, 5)) + ax.plot(dz, obs[mask_zoom], color=COLOR_OBS, lw=1.4, label="USGS obs (gauge 03463300)", zorder=3) + ax.plot(dz, sim_v[mask_zoom], color=COLOR_VRUGT, lw=1.6, + label=f"DA -- dynamic R (Vrugt) KGE={kge_v:.3f}", zorder=2) + ax.plot(dz, sim_nv[mask_zoom], color=COLOR_NOVRUGT, lw=1.6, linestyle="--", + label=f"DA -- constant R KGE={kge_nv:.3f}", zorder=2) + _add_helene_band(ax) + ax.axhline(peak_usgs, color=COLOR_OBS, lw=0.8, linestyle=":", alpha=0.5) + ax.text(ZOOM_START + pd.Timedelta(hours=3), peak_usgs * 1.02, + f"USGS peak {peak_usgs:.0f} m³/s", fontsize=9, color="black", alpha=0.75) + ax.set_xlabel("Date (UTC)", fontsize=11) + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_title( + "Hurricane Helene -- Routed discharge at gauge 03463300 (South Toe River)\n" + "CFE + DA (Muskingum routing): dynamic vs. constant observation error variance", + fontsize=11) + ax.legend(fontsize=10, loc="lower right") + ax.grid(True, alpha=0.25) + _format_xaxis(ax, mdates.DayLocator(interval=1), "%b %d") + plt.tight_layout() + out2 = os.path.join(args.out_dir, "vrugt_vs_novrugt_helene_zoom.png") + plt.savefig(out2, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/README.md b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/README.md new file mode 100644 index 00000000..03a43528 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder1_variance_scaled_vrugt/README.md @@ -0,0 +1,40 @@ +# Folder 1 — Variance Scaled Using Vrugt Formula + +## R Formula +``` +R(t) = (0.10 × y_obs(t))² + 0.001 × σ²_krig(t) +``` +Observation error variance scales with flow magnitude (Vrugt 2005 heteroscedastic). +σ²_krig from: `/mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance/` + +## Pipeline + +### 1. Calibrate +Shared best_params.json from external `calibrate-cfe` repo (R formula not involved). + +### 2. Assimilation +| Sub-step | Script | Server output | +|---|---|---| +| 2a. Forcing arm (30 members) | `run_perturbation_da_on.py` | `v2_perturbation_da_on/` | +| 2b. Hydro-state arm (20 members) | `run_perturbation_da_on.py` | `v2_perturbation_da_on/` | +| 18hr forecast cycles | `run_lead_time_forecast_sweep.py` | `v2_lead_time_forecast/` | +| 600-member crossed ensemble | `run_crossed_ensemble.py` (no --hardcoded-r) | `v2_crossed_ensemble_vrugt/` | + +### 3. Route +| Script | Input | Output | +|---|---|---| +| `run_route.py` | `v2_true_enkf_vrugt/` | `vrugt_dynamic_routed/routed_Q_test.csv` | +| `route_lead_time_forecasts.py` | `v2_lead_time_forecast/` | `v2_lead_time_forecast_routed/` | +| `run_route_crossed_ensemble.py` | `v2_crossed_ensemble_vrugt/` | `v2_crossed_ensemble_vrugt_routed/` | + +### 4. Evaluate +| Sub-step | Script | Status | +|---|---|---| +| 4a. Error decay | `plot_forecast_error_fixed_target.py` | ✅ | +| 4b. 600-member ensemble vs obs | `plot_production_ensemble_forecast.py` | ⏳ pending Vrugt ensemble run | +| 4c. Ensemble mean + spread (18hr cycles) | `plot_lead_time_decay_gauge.py` | ✅ | + +## Key Results (deterministic analysis routing) +- Full period KGE: **+0.277** +- Helene KGE: **+0.213** +- Helene peak: 752 m³/s (40% of USGS 1885.7 m³/s) diff --git a/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/batch_run_analysis_fixed_r.sh b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/batch_run_analysis_fixed_r.sh new file mode 100644 index 00000000..449d6a59 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/batch_run_analysis_fixed_r.sh @@ -0,0 +1,67 @@ +#!/usr/bin/env bash +# Batch: run CFE+EnKF analysis (R=0.07) for all 21 catchments → v2_fixed_r007_analysis/ +# Discovers catchments from the existing hardcoded-R lead-time forecast dir. +# After this completes, route with: +# conda run -n troute python3 da_methods/3_routing/run_route.py \ +# --gpkg /mnt/disk2/suma_helen_poster/gauge_03463300_network.gpkg \ +# --da-dir /mnt/disk2/suma_helen_poster/da_results/v2_fixed_r007_analysis \ +# --out-dir /mnt/disk2/suma_helen_poster/da_results/fixed_r007_routed \ +# --usgs-csv /mnt/disk2/suma_helen_poster/da_results/usgs_03463300_test.csv \ +# --kv-dir /mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance + +set -euo pipefail + +SCRIPT="$(dirname "$0")/run_analysis_fixed_r.py" +HARDCODED_R_DIR="/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast_hardcoded_r" +OBS_DIR="/mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance" +CFE_DIR="/mnt/disk2/suma_helen_poster/cfe_py" +CONFIG_FILE="/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json" +OUT_DIR="/mnt/disk2/suma_helen_poster/da_results/v2_fixed_r007_analysis" +LOG_DIR="$HOME/fixed_r007_analysis_logs" + +mkdir -p "$LOG_DIR" + +done=0; skipped=0; failed=0 + +for cat_dir in "$HARDCODED_R_DIR"/cat-*; do + cat_id=$(basename "$cat_dir") + forcing_file="$cat_dir/${cat_id}_nwm_operational_combined.csv" + params_file="$cat_dir/${cat_id}_best_params.json" + out_csv="$OUT_DIR/$cat_id/${cat_id}_test_results.csv" + + if [ -f "$out_csv" ]; then + echo "[skip] $cat_id — already done" + skipped=$((skipped + 1)) + continue + fi + + if [ ! -f "$forcing_file" ]; then + echo "[warn] $cat_id — forcing not found: $forcing_file" + failed=$((failed + 1)) + continue + fi + + if [ ! -f "$params_file" ]; then + echo "[warn] $cat_id — params not found: $params_file" + failed=$((failed + 1)) + continue + fi + + log_file="$LOG_DIR/${cat_id}.log" + echo "[run ] $cat_id" + python3 "$SCRIPT" \ + --cat-id "$cat_id" \ + --forcing-file "$forcing_file" \ + --obs-dir "$OBS_DIR" \ + --params-file "$params_file" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --out-dir "$OUT_DIR" \ + --hardcoded-r 0.07 \ + > "$log_file" 2>&1 \ + && { echo "[ok ] $cat_id"; done=$((done + 1)); } \ + || { echo "[FAIL] $cat_id — see $log_file"; failed=$((failed + 1)); } +done + +echo "" +echo "Summary: done=$done skipped=$skipped failed=$failed" diff --git a/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/batch_run_crossed_all_cats.sh b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/batch_run_crossed_all_cats.sh new file mode 100644 index 00000000..9d7e2448 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/batch_run_crossed_all_cats.sh @@ -0,0 +1,71 @@ +#!/bin/bash +# batch_run_crossed_all_cats.sh +# Runs run_crossed_ensemble.py for all 21 catchments. +# cat-1016300 is skipped (already done). + +set -e + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +OBS_DIR=/mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance +CFE_DIR=/mnt/disk2/suma_helen_poster/cfe_py +CONFIG_FILE=/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble +BEST_PARAMS_SRC=/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_pn +FORCING_DIR1=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings +FORCING_DIR2=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings + +SKIP=0 +DONE=0 +FAIL=0 + +for CAT in "${CATS[@]}"; do + echo "" + echo "===============================" + echo "=== $CAT ===" + echo "===============================" + + CAT_OUT=$OUT_DIR/$CAT + mkdir -p "$CAT_OUT" + + # Stage best_params directly from v2_true_enkf_pn calibration results + if [ ! -f "$CAT_OUT/${CAT}_best_params.json" ]; then + SRC="$BEST_PARAMS_SRC/$CAT/${CAT}_best_params.json" + if [ -f "$SRC" ]; then + cp "$SRC" "$CAT_OUT/" + echo " Staged best_params from v2_true_enkf_pn" + else + echo " WARNING: No best_params for $CAT in v2_true_enkf_pn — skipping" + SKIP=$((SKIP + 1)) + continue + fi + fi + + # Skip if parquet already exists + if [ -f "$CAT_OUT/${CAT}_crossed_ensemble.parquet" ]; then + echo " Crossed ensemble already complete — skipping" + SKIP=$((SKIP + 1)) + continue + fi + + python3 ~/run_crossed_ensemble.py \ + --cat-id "$CAT" \ + --forcing-dir "$FORCING_DIR1" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --out-dir "$OUT_DIR" \ + --test-forcing-dir1 "$FORCING_DIR1" \ + --test-forcing-dir2 "$FORCING_DIR2" \ + --hardcoded-r 0.07 \ + && DONE=$((DONE + 1)) \ + || { echo " FAILED: $CAT"; FAIL=$((FAIL + 1)); } +done + +echo "" +echo "=== Batch complete: done=$DONE skipped=$SKIP failed=$FAIL ===" diff --git a/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/input_enkf_new.json b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/input_enkf_new.json new file mode 100644 index 00000000..482233bc --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/input_enkf_new.json @@ -0,0 +1 @@ +{"n": 1, "m": 1, "R": 0.07, "Q": 0, "smcmax" : 0.9394097311639178, "N": 1000, "P": 0.01, "D":2} \ No newline at end of file diff --git a/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/new_EnKF.py b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/new_EnKF.py new file mode 100644 index 00000000..ff945c73 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/new_EnKF.py @@ -0,0 +1,53 @@ +# EnKF.py +import numpy as np + +class EnKF: + def __init__(self, n, m, R, Q, N, P, smcmax, D): + self.n = n # number of states + self.m = m # number of measurements + self.R = R # measurement noise + self.Q = Q # process noise + self.N = N # number of ensembles + self.P = P + self.state_estimates = np.zeros((self.n)) + self.covariance_matrices = np.zeros((self.n, self.n)) + self.current_step = 0 + self.smcmax = smcmax + self.D = D + self.storage_max_m = self.smcmax * self.D + self.storage_init = self.storage_max_m * 0.667 + self.ensembles = np.full((self.n, self.N), self.storage_init) + + def predict(self, F_results): + error_factor_sim = 0.005 + for i in range(self.N): + perturbation_factor_sim = np.random.standard_normal() + self.ensembles[:, i] = F_results[i] + self.ensembles[:, i] += (perturbation_factor_sim * F_results[i] * error_factor_sim) + return self.ensembles + + def update(self, yi, H_results): + if np.any(np.isnan(yi)): + return np.mean(self.ensembles, axis=1) # return current state estimate without updating + + y_ensembles = np.zeros((self.m, self.N)) + error_factor_sm = 0.03 + for i in range(self.N): + perturbation_factor_sm = np.random.standard_normal() + y_ensembles[:, i] = H_results[i] + y_ensembles[:, i] += (perturbation_factor_sm * H_results[i] * error_factor_sm) + + ensemble_mean = np.mean(self.ensembles, axis=1)[:, np.newaxis] + y_ensembles_mean = np.mean(y_ensembles, axis=1)[:, np.newaxis] + Pxy = (self.ensembles - ensemble_mean).dot((y_ensembles - y_ensembles_mean).T) / (self.N - 1) + Pyy = np.cov(y_ensembles, bias=True) + K = Pxy.dot(np.linalg.pinv(Pyy + self.R)) + self.ensembles += K.dot(yi - y_ensembles_mean) + return np.mean(self.ensembles, axis=1) + + def get_state_estimate(self): + return np.mean(self.ensembles, axis=1) + + def get_covariance(self): + ensemble_mean = np.mean(self.ensembles, axis=1)[:, np.newaxis] + return (self.ensembles - ensemble_mean).dot((self.ensembles - ensemble_mean).T) / (self.N - 1) diff --git a/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/run_analysis_fixed_r.py b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/run_analysis_fixed_r.py new file mode 100644 index 00000000..1994f179 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/run_analysis_fixed_r.py @@ -0,0 +1,231 @@ +""" +run_analysis_fixed_r.py -- DA analysis trajectory with constant R=0.07. + +Runs CFE + EnKF forward through the full test period for one catchment, +using a fixed observation-error variance (R=0.07) instead of the dynamic +Vrugt formula. Produces a _test_results.csv that run_route.py can read. + +Output: + //_test_results.csv + Columns: date, sim_mm_h, obs_mm_h, precip_mm_h + +Usage: + python3 run_analysis_fixed_r.py \\ + --cat-id cat-1016300 \\ + --forcing-file /mnt/disk2/.../cat-1016300_nwm_operational_combined.csv \\ + --obs-dir /mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance \\ + --params-file /mnt/disk2/.../cat-1016300_best_params.json \\ + --cfe-dir /mnt/disk2/suma_helen_poster/cfe_py \\ + --config-file /mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json \\ + --out-dir /mnt/disk2/suma_helen_poster/da_results/v2_fixed_r007_analysis \\ + --hardcoded-r 0.07 +""" + +import argparse +import json +import os +import sys +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +TEST_START = "2024-08-24 00:00:00" +TEST_END = "2024-10-31 23:00:00" + +CAT_ID = None +OBS_FILE = None +CFE_CONFIG_FILE = None +TEST_FORCING_FILE = None +OUT_DIR = None +HARDCODED_R = 0.07 +bmi_cfe = None + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(): + df = pd.read_csv(TEST_FORCING_FILE) + df["date"] = pd.to_datetime(df["time"]).dt.strftime("%Y-%m-%d %H:%M:%S") + df["total_precipitation"] = df["APCP_surface"] * 3600.0 + df["potential_evaporation"] = priestley_taylor_pet( + df["DSWRF_surface"].values, df["TMP_2maboveground"].values) + return df + + +def load_obs(): + df = pd.read_csv(OBS_FILE) + date_col = next((c for c in df.columns + if c.lower() in ("date", "time", "datetime", "timestamp")), None) + q_col = next((c for c in df.columns + if "qkrig" in c.lower() and "var" not in c.lower()), None) + var_col = next((c for c in df.columns if "var" in c.lower()), None) + df[date_col] = pd.to_datetime(df[date_col]) + df = df.sort_values(date_col).reset_index(drop=True) + date_strs = df[date_col].dt.strftime("%Y-%m-%d %H:%M:%S") + obs_dict = dict(zip(date_strs, df[q_col])) + var_vals = df[var_col] if var_col else pd.Series([0.07] * len(df)) + var_dict = dict(zip(date_strs, var_vals)) + return obs_dict, var_dict + + +def write_model_config(best_params): + with open(CFE_CONFIG_FILE) as f: + cfg = json.load(f) + cfg["forcing_file"] = TEST_FORCING_FILE + cfg["soil_params"]["bb"] = best_params["bb"] + cfg["soil_params"]["smcmax"] = best_params["smcmax"] + cfg["soil_params"]["satdk"] = best_params["satdk"] + cfg["slop"] = best_params["slop"] + cfg["max_gw_storage"] = best_params["max_gw_storage"] + cfg["expon"] = best_params["expon"] + cfg["Cgw"] = best_params["Cgw"] + cfg["K_lf"] = best_params["K_lf"] + cfg["K_nash"] = best_params["K_nash"] + cfg["partition_scheme"] = ("Schaake" if best_params["scheme"] <= 0.5 + else "Xinanjiang") + tmp_cfg = str(OUT_DIR / f"{CAT_ID}_bmi_config_fixed_r.json") + with open(tmp_cfg, "w") as f: + json.dump(cfg, f) + return tmp_cfg + + +def run_model_step(m, precip_mmh, pet_mmh): + m.set_value("atmosphere_water__time_integral_of_precipitation_mass_flux", + precip_mmh / 1000.0) + m.set_value("water_potential_evaporation_flux", + pet_mmh / 1000.0 / 3600.0) + m.update() + return m.get_value("land_surface_water__runoff_depth") * 1000.0 + + +def get_state(m): + return { + "soil_m": m.soil_reservoir["storage_m"], + "gw_m": m.gw_reservoir["storage_m"], + "nash0": float(m.nash_storage[0]), + "nash1": float(m.nash_storage[1]), + } + + +def set_state(m, state): + m.soil_reservoir["storage_m"] = state["soil_m"] + m.gw_reservoir["storage_m"] = state["gw_m"] + m.nash_storage[0] = state["nash0"] + m.nash_storage[1] = state["nash1"] + + +def enkf_update(state, q_sim, y_obs, krig_var): + R = HARDCODED_R if HARDCODED_R is not None else max( + (0.10 * max(y_obs, 0.0)) ** 2 + 0.001 * krig_var, 1e-6) + R = max(R, 1e-6) + P_yy = max(q_sim * 0.01, 1e-6) + K = P_yy / (P_yy + R) + scale = K * (y_obs - q_sim) / max(abs(q_sim), 1e-6) + return { + "soil_m": max(state["soil_m"] * (1.0 + scale), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + scale), 1e-6), + "nash0": max(state["nash0"] + state["nash0"] * scale, 0.0), + "nash1": max(state["nash1"] + state["nash1"] * scale, 0.0), + } + + +def run_da_analysis(cfg_path, df_test, obs_dict, var_dict): + """Run DA analysis; return per-timestep (date, sim_mm_h, obs_mm_h, precip_mm_h).""" + model = bmi_cfe.BMI_CFE(cfg_file=cfg_path) + model.initialize() + + records = [] + for _, row in df_test.iterrows(): + date_str = row["date"] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + + q_sim = run_model_step(model, P, E) + y_obs = obs_dict.get(date_str, np.nan) + + if not np.isnan(y_obs): + krig_var = var_dict.get(date_str, 0.07) + state = get_state(model) + updated = enkf_update(state, q_sim, y_obs, krig_var) + set_state(model, updated) + + records.append({ + "date": date_str, + "sim_mm_h": q_sim, + "obs_mm_h": y_obs, + "precip_mm_h": P, + }) + + model.finalize() + return pd.DataFrame(records) + + +def main(): + global CAT_ID, OBS_FILE, CFE_CONFIG_FILE, TEST_FORCING_FILE + global OUT_DIR, HARDCODED_R, bmi_cfe + + parser = argparse.ArgumentParser() + parser.add_argument("--cat-id", required=True) + parser.add_argument("--forcing-file", required=True, + help="Path to {cat}_nwm_operational_combined.csv") + parser.add_argument("--obs-dir", required=True, + help="Dir containing {cat}.csv with Qkrig + variance") + parser.add_argument("--params-file", required=True, + help="Path to {cat}_best_params.json") + parser.add_argument("--cfe-dir", required=True, + help="Directory containing bmi_cfe.py") + parser.add_argument("--config-file", required=True, + help="BMI config JSON template") + parser.add_argument("--out-dir", required=True) + parser.add_argument("--hardcoded-r", type=float, default=0.07, + help="Fixed R value for EnKF (default 0.07)") + args = parser.parse_args() + + CAT_ID = args.cat_id + TEST_FORCING_FILE = args.forcing_file + OBS_FILE = os.path.join(args.obs_dir, f"{CAT_ID}.csv") + CFE_CONFIG_FILE = args.config_file + HARDCODED_R = args.hardcoded_r + OUT_DIR = Path(args.out_dir) / CAT_ID + OUT_DIR.mkdir(parents=True, exist_ok=True) + + out_csv = OUT_DIR / f"{CAT_ID}_test_results.csv" + if out_csv.exists(): + print(f"[{CAT_ID}] already done — {out_csv}") + return + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + bmi_cfe = _bmi_cfe + + with open(args.params_file) as f: + raw = json.load(f) + best_params = raw.get("best_parameters", raw) + + obs_dict, var_dict = load_obs() + df_forcing = load_test_forcing() + t_mask = (df_forcing["date"] >= TEST_START) & (df_forcing["date"] <= TEST_END) + df_test = df_forcing[t_mask].reset_index(drop=True) + + print(f"[{CAT_ID}] R={HARDCODED_R} steps={len(df_test)}") + + cfg_path = write_model_config(best_params) + df_out = run_da_analysis(cfg_path, df_test, obs_dict, var_dict) + + df_out.to_csv(out_csv, index=False) + print(f"[{CAT_ID}] saved {out_csv} ({len(df_out)} rows)") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/run_crossed_ensemble.py b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/run_crossed_ensemble.py new file mode 100644 index 00000000..325a64e2 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/run_crossed_ensemble.py @@ -0,0 +1,452 @@ +""" +run_crossed_ensemble.py -- 4b: 600-member crossed ensemble for Helene. + +Crosses the two DA-on perturbation arms: + Forcing arm : 30 met draws (from run_perturbation_da_on.py Phase 2A) + Hydro arm : 20 state draws (from run_perturbation_da_on.py Phase 2B) + +For each Helene issue time t0: + Run 600 members = 30 forcing draws x 20 hydro-state draws. + Each member (i, j) gets: + - Initial state = perturbed DA analysis state (draw j of hydro arm) + - Met forcing = perturbed precip/PET sequence (draw i of forcing arm) + +This gives the full uncertainty envelope combining both sources — +the operational probabilistic forecast picture. + +Output (per catchment): + //_crossed_ensemble.parquet + Columns: issue_time, lead_hour, member_0000..member_0599 + +Usage: + python3 run_crossed_ensemble.py \\ + --cat-id cat-1016300 \\ + --forcing-dir \\ + --obs-dir /mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance \\ + --cfe-dir /mnt/disk2/suma_helen_poster/cfe_py \\ + --config-file /mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json \\ + --out-dir /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble \\ + --test-forcing-dir1 \\ + --test-forcing-dir2 + +Strategy: + Re-runs DA analysis (same as Phase 1 of run_perturbation_da_on.py) to get + the per-issue-time analysis state. Then generates N_FORCING sets of perturbed + forcing sequences and N_HYDRO perturbed states. Runs all 600 combinations. + + To avoid re-running Phase 1, if the Phase 1 snapshot file exists from a + prior run_perturbation_da_on.py run, it is loaded directly. + +NOTE: All special characters kept ASCII to avoid encoding issues on GPU. +""" + +import argparse +import os +import sys +import json +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +N_FORCING = 30 +N_HYDRO = 20 +N_MEMBERS = N_FORCING * N_HYDRO # 600 + +TEST_START = "2024-08-24 00:00:00" # 1 month spinup before Helene window +TEST_END = "2024-10-31 23:00:00" + +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-30 06:00:00") +FORECAST_HOURS = 18 + +PRECIP_SIGMA = 0.15 +PET_SIGMA = 0.10 +STATE_FRAC = 0.05 + +CAT_ID = None +OBS_FILE = None +CFE_CONFIG_FILE = None +TEST_FORCING_FILE = None +OUT_DIR = None +HARDCODED_R = None +bmi_cfe = None + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(): + df = pd.read_csv(TEST_FORCING_FILE) + df["date"] = pd.to_datetime(df["time"]).dt.strftime("%Y-%m-%d %H:%M:%S") + df["total_precipitation"] = df["APCP_surface"] * 3600.0 + df["potential_evaporation"] = priestley_taylor_pet( + df["DSWRF_surface"].values, df["TMP_2maboveground"].values) + return df + + +def load_obs(): + df = pd.read_csv(OBS_FILE) + date_col = next((c for c in df.columns + if c.lower() in ("date", "time", "datetime", "timestamp")), None) + q_col = next((c for c in df.columns + if "qkrig" in c.lower() and "var" not in c.lower()), None) + var_col = next((c for c in df.columns if "var" in c.lower()), None) + df[date_col] = pd.to_datetime(df[date_col]) + df = df.sort_values(date_col).reset_index(drop=True) + date_strs = df[date_col].dt.strftime("%Y-%m-%d %H:%M:%S") + obs_dict = dict(zip(date_strs, df[q_col])) + var_vals = df[var_col] if var_col else pd.Series([0.07] * len(df)) + var_dict = dict(zip(date_strs, var_vals)) + return obs_dict, var_dict + + +def write_model_config(best_params): + """Write the BMI config file once; reused for every model instantiation.""" + with open(CFE_CONFIG_FILE) as f: + cfg = json.load(f) + cfg["forcing_file"] = TEST_FORCING_FILE + cfg["soil_params"]["bb"] = best_params["bb"] + cfg["soil_params"]["smcmax"] = best_params["smcmax"] + cfg["soil_params"]["satdk"] = best_params["satdk"] + cfg["slop"] = best_params["slop"] + cfg["max_gw_storage"] = best_params["max_gw_storage"] + cfg["expon"] = best_params["expon"] + cfg["Cgw"] = best_params["Cgw"] + cfg["K_lf"] = best_params["K_lf"] + cfg["K_nash"] = best_params["K_nash"] + cfg["partition_scheme"] = ("Schaake" if best_params["scheme"] <= 0.5 + else "Xinanjiang") + tmp_cfg = str(OUT_DIR / f"{CAT_ID}_bmi_config_crossed_stable.json") + with open(tmp_cfg, "w") as f: + json.dump(cfg, f) + return tmp_cfg + + +def make_model(cfg_path): + """Instantiate and initialize a CFE model from a pre-written config path.""" + m = bmi_cfe.BMI_CFE(cfg_file=cfg_path) + m.initialize() + return m + + +def get_state(m): + return { + "soil_m": m.soil_reservoir["storage_m"], + "gw_m": m.gw_reservoir["storage_m"], + "nash0": float(m.nash_storage[0]), + "nash1": float(m.nash_storage[1]), + } + + +def set_state(m, state): + m.soil_reservoir["storage_m"] = state["soil_m"] + m.gw_reservoir["storage_m"] = state["gw_m"] + m.nash_storage[0] = state["nash0"] + m.nash_storage[1] = state["nash1"] + + +def perturb_state(state, rng, frac=STATE_FRAC): + return { + "soil_m": max(state["soil_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "nash0": max(state["nash0"] * (1.0 + frac * rng.standard_normal()), 0.0), + "nash1": max(state["nash1"] * (1.0 + frac * rng.standard_normal()), 0.0), + } + + +def run_model_step(m, precip_mmh, pet_mmh): + m.set_value("atmosphere_water__time_integral_of_precipitation_mass_flux", + precip_mmh / 1000.0) + m.set_value("water_potential_evaporation_flux", + pet_mmh / 1000.0 / 3600.0) + m.update() + return m.get_value("land_surface_water__runoff_depth") * 1000.0 + + +def enkf_update(state, q_sim, y_obs, krig_var, rng): + R = HARDCODED_R if HARDCODED_R is not None else max( + (0.10 * max(y_obs, 0.0)) ** 2 + 0.001 * krig_var, 1e-6) + R = max(R, 1e-6) + P_yy = max(q_sim * 0.01, 1e-6) + K = P_yy / (P_yy + R) + scale = K * (y_obs - q_sim) / max(abs(q_sim), 1e-6) + return { + "soil_m": max(state["soil_m"] * (1.0 + scale), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + scale), 1e-6), + "nash0": max(state["nash0"] + state["nash0"] * scale, 0.0), + "nash1": max(state["nash1"] + state["nash1"] * scale, 0.0), + } + + +def run_phase1_da(cfg_path, df_test, obs_dict, var_dict, rng): + """Run DA through test period; return snapshot dict {timestamp -> state}.""" + print(f" Phase 1: DA analysis through test period...") + model = make_model(cfg_path) + snapshots = {} + for _, row in df_test.iterrows(): + date_str = row["date"] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + q_sim = run_model_step(model, P, E) + y_obs = obs_dict.get(date_str, np.nan) + if not np.isnan(y_obs): + krig_var = var_dict.get(date_str, 0.07) + state = get_state(model) + updated = enkf_update(state, q_sim, y_obs, krig_var, rng) + set_state(model, updated) + date_ts = pd.Timestamp(date_str) + if HELENE_START <= date_ts <= HELENE_END: + snapshots[date_ts] = get_state(model) + model.finalize() + print(f" Phase 1 done: {len(snapshots)} issue-time snapshots") + return snapshots + + +def main(): + global CAT_ID, OBS_FILE, CFE_CONFIG_FILE, TEST_FORCING_FILE + global OUT_DIR, HARDCODED_R, bmi_cfe + + parser = argparse.ArgumentParser() + parser.add_argument("--cat-id", required=True) + parser.add_argument("--forcing-dir", required=True) + parser.add_argument("--obs-dir", required=True) + parser.add_argument("--cfe-dir", required=True) + parser.add_argument("--config-file", required=True) + parser.add_argument("--out-dir", required=True) + parser.add_argument("--test-forcing-dir1", default=None) + parser.add_argument("--test-forcing-dir2", default=None) + parser.add_argument("--hardcoded-r", type=float, default=None, + help="Fixed R value for EnKF (e.g. 0.07). " + "Omit to use dynamic Vrugt R: R=(0.10*y)^2 + 0.001*krig_var") + parser.add_argument("--n-forcing", type=int, default=N_FORCING, + help="Number of met forcing draws (default 30)") + parser.add_argument("--n-hydro", type=int, default=N_HYDRO, + help="Number of hydro-state draws (default 20)") + args = parser.parse_args() + + CAT_ID = args.cat_id + OBS_FILE = os.path.join(args.obs_dir, f"{CAT_ID}.csv") + CFE_CONFIG_FILE = args.config_file + HARDCODED_R = args.hardcoded_r + OUT_DIR = Path(args.out_dir) / CAT_ID + OUT_DIR.mkdir(parents=True, exist_ok=True) + + n_fa = args.n_forcing + n_ha = args.n_hydro + n_total = n_fa * n_ha + print(f"[{CAT_ID}] Crossed ensemble: {n_fa} forcing x {n_ha} hydro = {n_total} members") + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + bmi_cfe = _bmi_cfe + + # Reuse combined forcing CSV from run_perturbation_da_on.py if it exists + # (avoids writing a duplicate large file to disk) + da_on_combined = (Path(args.out_dir).parent / "v2_perturbation_da_on" / + CAT_ID / f"{CAT_ID}_nwm_operational_combined.csv") + crossed_combined = OUT_DIR / f"{CAT_ID}_nwm_operational_combined.csv" + + if da_on_combined.exists(): + TEST_FORCING_FILE = str(da_on_combined) + print(f" Reusing combined forcing: {TEST_FORCING_FILE}") + elif crossed_combined.exists(): + TEST_FORCING_FILE = str(crossed_combined) + print(f" Reusing combined forcing: {TEST_FORCING_FILE}") + elif args.test_forcing_dir1 and args.test_forcing_dir2: + f1 = os.path.join(args.test_forcing_dir1, f"{CAT_ID}.csv") + f2 = os.path.join(args.test_forcing_dir2, f"{CAT_ID}.csv") + df1, df2 = pd.read_csv(f1), pd.read_csv(f2) + combined = (pd.concat([df1, df2], ignore_index=True) + .drop_duplicates(subset="time").sort_values("time")) + TEST_FORCING_FILE = str(crossed_combined) + combined.to_csv(TEST_FORCING_FILE, index=False) + print(f" Wrote combined forcing: {TEST_FORCING_FILE}") + else: + print("Need --test-forcing-dir1 and --test-forcing-dir2"); return + + best_params_file = OUT_DIR / f"{CAT_ID}_best_params.json" + if not best_params_file.exists(): + # Also check the DA results dir from run_perturbation_da_on.py + alt = Path(args.out_dir).parent / "v2_perturbation_da_on" / CAT_ID / f"{CAT_ID}_best_params.json" + if alt.exists(): + best_params_file = alt + else: + print(f"Missing best params at {best_params_file}"); return + with open(best_params_file) as f: + best_params = json.load(f)["best_parameters"] + + obs_dict, var_dict = load_obs() + df_forcing = load_test_forcing() + t_mask = (df_forcing["date"] >= TEST_START) & (df_forcing["date"] <= TEST_END) + df_test = df_forcing[t_mask].reset_index(drop=True) + + # Write BMI config once — reused for all 90,600 model instantiations + cfg_path = write_model_config(best_params) + print(f" BMI config written once: {cfg_path}") + + rng = np.random.default_rng(hash(CAT_ID) & 0x7fffffff) + + # Phase 1: get DA analysis snapshots + # Check if a prior run already saved them (skip re-running DA if so) + snap_cache = OUT_DIR / f"{CAT_ID}_da_snapshots.parquet" + if snap_cache.exists(): + print(f" Loading cached DA snapshots from {snap_cache}") + df_snaps = pd.read_parquet(snap_cache) + snapshots = { + pd.Timestamp(row["issue_time"]): { + "soil_m": row["soil_m"], "gw_m": row["gw_m"], + "nash0": row["nash0"], "nash1": row["nash1"], + } + for _, row in df_snaps.iterrows() + } + else: + snapshots = run_phase1_da(cfg_path, df_test, obs_dict, var_dict, rng) + + issue_times = sorted(snapshots.keys()) + t0_to_idx = {t: i for i, t in enumerate(df_test["date"].apply(pd.Timestamp))} + + # ------------------------------------------------------------------ # + # Crossed ensemble: for each issue_time, build n_fa forcing draws and # + # n_ha state draws, then run all n_fa x n_ha combinations. # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Running crossed ensemble ({n_total} members x " + f"{len(issue_times)} issue times)...") + + all_records = [] + all_hydro_draw_records = [] # (issue_time, hydro_draw_j, states) + all_forcing_draw_records = [] # (issue_time, forcing_draw_i, lead, P_pert, E_pert, scales) + + for t_idx, t0 in enumerate(issue_times): + snap = snapshots[t0] + t0_str = t0.strftime("%Y-%m-%d %H:%M:%S") + if t0 not in t0_to_idx: + continue + start_idx = t0_to_idx[t0] + + # Pre-draw n_fa forcing perturbation sequences + forcing_seqs = [] + for i in range(n_fa): + seq = [] + for lead in range(1, FORECAST_HOURS + 1): + fi = start_idx + lead + if fi >= len(df_test): + seq.append((0.0, 0.0)) + all_forcing_draw_records.append({ + "issue_time": t0_str, "forcing_draw_i": i, + "lead_hour": lead, + "P_pert_mmh": 0.0, "E_pert_mmh": 0.0, + "precip_scale": np.nan, "pet_scale": np.nan, + }) + continue + row = df_test.iloc[fi] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + mu_p = -0.5 * PRECIP_SIGMA ** 2 + P_pert = float(P * rng.lognormal(mu_p, PRECIP_SIGMA)) if P > 0 else 0.0 + E_pert = max(float(E * (1.0 + PET_SIGMA * rng.standard_normal())), 0.0) + seq.append((P_pert, E_pert)) + all_forcing_draw_records.append({ + "issue_time": t0_str, "forcing_draw_i": i, + "lead_hour": lead, + "P_pert_mmh": P_pert, "E_pert_mmh": E_pert, + "precip_scale": P_pert / P if P > 0 else np.nan, + "pet_scale": E_pert / E if E > 0 else np.nan, + }) + forcing_seqs.append(seq) + + # Pre-draw n_ha hydro-state perturbations + state_draws = [snap if j == 0 else perturb_state(snap, rng) + for j in range(n_ha)] + + # Record each hydro draw's actual initial state + for j, st in enumerate(state_draws): + all_hydro_draw_records.append({ + "issue_time": t0_str, "hydro_draw_j": j, + "soil_m": st["soil_m"], "gw_m": st["gw_m"], + "nash0": st["nash0"], "nash1": st["nash1"], + }) + + # q_matrix shape: (FORECAST_HOURS, n_fa, n_ha) + q_matrix = np.full((FORECAST_HOURS, n_fa, n_ha), np.nan) + + for i, f_seq in enumerate(forcing_seqs): + for j, state_j in enumerate(state_draws): + m = make_model(cfg_path) + set_state(m, state_j) + for lead in range(1, FORECAST_HOURS + 1): + P_pert, E_pert = f_seq[lead - 1] + q = run_model_step(m, P_pert, E_pert) + q_matrix[lead - 1, i, j] = q + m.finalize() + + # Flatten 600-member matrix into rows + for lead in range(1, FORECAST_HOURS + 1): + row_data = {"issue_time": t0_str, "lead_hour": lead} + mem_idx = 0 + for i in range(n_fa): + for j in range(n_ha): + row_data[f"member_{mem_idx:04d}"] = q_matrix[lead - 1, i, j] + mem_idx += 1 + all_records.append(row_data) + + if (t_idx + 1) % 10 == 0 or (t_idx + 1) == len(issue_times): + print(f" {t_idx + 1}/{len(issue_times)} issue times complete") + + df_out = pd.DataFrame(all_records) + out_path = OUT_DIR / f"{CAT_ID}_crossed_ensemble.parquet" + df_out.to_parquet(out_path, index=False) + print(f"Saved: {out_path}") + print(f" Shape: {df_out.shape} " + f"({df_out['issue_time'].nunique()} issue times x " + f"{FORECAST_HOURS} leads x {n_total} members)") + + # ------------------------------------------------------------------ # + # Provenance files — enable full ensemble traceability and restart # + # ------------------------------------------------------------------ # + + # 1. DA analysis snapshots — unperturbed state at each issue_time. + # Load any row and set_state() to restart from that init time. + snap_records = [ + {"issue_time": t.strftime("%Y-%m-%d %H:%M:%S"), **s} + for t, s in snapshots.items() + ] + snap_path = OUT_DIR / f"{CAT_ID}_da_snapshots.parquet" + pd.DataFrame(snap_records).to_parquet(snap_path, index=False) + print(f"Saved: {snap_path} ({len(snap_records)} snapshots)") + + # 2. Member manifest — decoder ring: member_col -> (forcing_draw_i, hydro_draw_j). + # member_k uses forcing draw k//n_ha and hydro draw k%n_ha. + manifest = [ + {"member": f"member_{k:04d}", "member_idx": k, + "forcing_draw_i": k // n_ha, "hydro_draw_j": k % n_ha} + for k in range(n_total) + ] + manifest_path = OUT_DIR / f"{CAT_ID}_member_manifest.csv" + pd.DataFrame(manifest).to_csv(manifest_path, index=False) + print(f"Saved: {manifest_path} ({n_total} members)") + + # 3. Hydro draw states — actual perturbed initial state for each of + # n_ha draws at every issue_time. Links member -> initial conditions. + hydro_path = OUT_DIR / f"{CAT_ID}_hydro_draw_states.parquet" + pd.DataFrame(all_hydro_draw_records).to_parquet(hydro_path, index=False) + print(f"Saved: {hydro_path} ({len(all_hydro_draw_records)} rows)") + + # 4. Forcing draw sequences — actual P/PET perturbation values for each + # of n_fa draws at every (issue_time, lead_hour). Links member -> forcing. + forcing_path = OUT_DIR / f"{CAT_ID}_forcing_draw_sequences.parquet" + pd.DataFrame(all_forcing_draw_records).to_parquet(forcing_path, index=False) + print(f"Saved: {forcing_path} ({len(all_forcing_draw_records)} rows)") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/run_lead_time_forecast_sweep.py b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/run_lead_time_forecast_sweep.py new file mode 100644 index 00000000..19f3b4d5 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/run_lead_time_forecast_sweep.py @@ -0,0 +1,508 @@ +""" +Forecast lead-time evaluation. + +For each issue-time t0 in a sampling schedule across the test period, fork the +ensemble at t0 and free-run 18 hours of forecast. Two parallel trajectories are +maintained throughout the test period: + + da: Production setup. DA on, forcing perturbed, process noise on. + openloop: No DA from t=0. Forcing perturbed, process noise on. + +At each issue time t0, the state of each scenario's ensemble is copied into a +forecast ensemble that free-runs 18 hours with: + - DA off (no observations consumed during the forecast window) + - process noise off (no obs to collapse toward, so no anti-collapse needed) + - forcing perturbed (lognormal precip, Gaussian PET — proxy for forecast + precip uncertainty in lieu of actual HEFS forecasts) + +Issue-time schedule: + - Base cadence: every --base-step-h hours across the full test period + (default 6h → ~1600 issue times for the year) + - Densified to hourly across --dense-start..--dense-end if provided + (default: 2024-09-24 → 2024-09-28, the Helene window) + +Outputs (stacked across all issue times): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv + Columns: issue_time, lead_hour, valid_time, member_00 .. member_19 + (units: mm/h) + +Post-processing pipeline (separate scripts, run after this): + route_lead_time_forecasts.py — pushes each (t0, lead_hour) forecast through + T-route to the gauge + plot_lead_time_decay.py — error-vs-lead-time curve, DA vs open-loop + +Imports EnKFAssimilator from calibrate_catchment_cfe_da_v2.py so the DA math +matches production exactly. +""" + +import argparse +import os +import sys +import json +import importlib.util +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +SPINUP_START = "2023-02-01 00:00:00" +SPINUP_END = "2023-09-30 23:00:00" +TEST_START = "2023-10-01 00:00:00" +TEST_END = "2024-10-31 23:00:00" + +# Forecast lead time (hours after each issue time) +FORECAST_LEAD_HOURS = 18 + +# Default densification window — Hurricane Helene +DEFAULT_DENSE_START = "2024-09-24 00:00:00" +DEFAULT_DENSE_END = "2024-09-28 23:00:00" + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(test_forcing_file): + df = pd.read_csv(test_forcing_file) + df['date'] = pd.to_datetime(df['time']).dt.strftime('%Y-%m-%d %H:%M:%S') + df['total_precipitation'] = df['APCP_surface'] * 3600.0 + df['potential_evaporation'] = priestley_taylor_pet( + df['DSWRF_surface'].values, df['TMP_2maboveground'].values + ) + return df + + +def import_enkf_class(script_path): + """Load EnKFAssimilator from the production script without triggering its main().""" + spec = importlib.util.spec_from_file_location("prod_v2", script_path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod.EnKFAssimilator + + +# ---------------- State snapshot / restore ----------------------------------- +# +# The EnKF only ever touches these 4 states, so they're the only states that +# diverge between members. Copy these from src → dst to fork the ensemble. + +def snapshot_states(models): + return [ + { + "soil_m": float(m.soil_reservoir["storage_m"]), + "gw_m": float(m.gw_reservoir["storage_m"]), + "nash0_m": float(m.nash_storage[0]), + "nash1_m": float(m.nash_storage[1]), + } + for m in models + ] + + +def restore_states(models, snapshot): + for m, s in zip(models, snapshot): + m.soil_reservoir["storage_m"] = s["soil_m"] + m.gw_reservoir["storage_m"] = s["gw_m"] + m.nash_storage[0] = s["nash0_m"] + m.nash_storage[1] = s["nash1_m"] + + +# ---------------- Issue-time schedule ---------------------------------------- + +def build_issue_time_schedule(dates_list, base_step_h, dense_start, dense_end): + """Return a sorted list of issue-time strings sampled from dates_list. + + Every base_step_h hours across the full test period, plus every hour inside + [dense_start, dense_end] if those are provided. + """ + dates_dt = pd.to_datetime(dates_list) + selected = set() + + # Base cadence (every base_step_h hours from the first date) + base_mask = (np.arange(len(dates_dt)) % base_step_h == 0) + for d in dates_dt[base_mask]: + selected.add(d.strftime('%Y-%m-%d %H:%M:%S')) + + # Densified hourly window + if dense_start is not None and dense_end is not None: + ds = pd.Timestamp(dense_start) + de = pd.Timestamp(dense_end) + dense_mask = (dates_dt >= ds) & (dates_dt <= de) + for d in dates_dt[dense_mask]: + selected.add(d.strftime('%Y-%m-%d %H:%M:%S')) + + return sorted(selected) + + +# ---------------- CFE helpers ------------------------------------------------ + +def build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, enkf, + apply_init_perturbation=True, member_zero_clean=True): + """Build N CFE BMI instances with optional init-state perturbation.""" + def custom_load_forcing(self_cfe): + df = load_test_forcing(test_forcing_file) + self_cfe.forcing_data = df.rename(columns={"date": "time"}) + + models = [] + for i in range(N): + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.load_forcing_file = custom_load_forcing.__get__(m) + m.initialize() + if apply_init_perturbation and not (member_zero_clean and i == 0): + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + r = enkf.init_state_perturb_frac + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1 + r * enkf.rng.standard_normal()), 1e-6, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1 + r * enkf.rng.standard_normal()), 1e-6, gw_max)) + jitter = max(sm_max * 1e-6, 1e-9) + m.nash_storage[0] = max(n00 + jitter * enkf.rng.standard_normal(), 0.0) + m.nash_storage[1] = max(n10 + jitter * enkf.rng.standard_normal(), 0.0) + models.append(m) + return models + + +def step_ensemble(models, p_arr, e_arr): + """Advance each member one hour with member-specific (P, PET). Returns Q (mm/h).""" + q = np.empty(len(models), dtype=float) + for i, m in enumerate(models): + m.set_value('atmosphere_water__time_integral_of_precipitation_mass_flux', + float(p_arr[i]) / 1000.0) + m.set_value('water_potential_evaporation_flux', + float(e_arr[i]) / 1000.0 / 3600.0) + m.update() + q[i] = m.get_value('land_surface_water__runoff_depth') * 1000.0 # m/h → mm/h + return q + + +def do_sanity_check(prod_models, fcst_models, dates_list, forcing_by_date, h, N): + """One-shot verification that the 4-state snapshot is sufficient to reproduce + next-hour forecast Q. + + At main-loop hour h (after prod's hour-h step has completed including DA + + process noise), snapshot prod, restore into fcst, then step BOTH ensembles + one hour using identical unperturbed forcing. If snapshot/restore captures + everything that matters, prod's and fcst's hour-(h+1) Q should match member- + by-member to machine precision. + + Note: this mutates prod's state (h+1 step with non-perturbed forcing and no + DA), so the caller should exit the script right after. + """ + if h + 1 >= len(dates_list): + print("[sanity-check] not enough horizon for h+1 step — skipping.") + return False + snap = snapshot_states(prod_models) + restore_states(fcst_models, snap) + + next_date = dates_list[h + 1] + p_next, e_next = forcing_by_date[next_date] + p_uniform = np.full(N, p_next, dtype=float) + e_uniform = np.full(N, e_next, dtype=float) + + q_prod = step_ensemble(prod_models, p_uniform, e_uniform) + q_fcst = step_ensemble(fcst_models, p_uniform, e_uniform) + + diff = q_prod - q_fcst + max_abs = float(np.max(np.abs(diff))) + print("[sanity-check] forecast Q at h+1 with identical unperturbed forcing:") + print(f" prod first 5 members: {q_prod[:5]}") + print(f" fcst first 5 members: {q_fcst[:5]}") + print(f" max |prod - fcst|: {max_abs:.3e} mm/h") + if max_abs < 1e-9: + print(" OK: snapshot/restore preserves state to machine precision.") + return True + print(" WARN: nonzero divergence — likely a state not in the 4-state snapshot.") + return False + + +def run_forecast(fcst_models, enkf_fcst, dates_list, forcing_by_date, + t0_idx, n_lead): + """Run an n_lead-hour free-forecast starting from the current state of fcst_models. + + Forcing is perturbed per member (lognormal precip + Gaussian PET, same as + production). No DA. No process noise. + + Returns: (lead_hours, valid_times, q_matrix shape (n_lead, N)). + """ + N = len(fcst_models) + q_matrix = np.full((n_lead, N), np.nan, dtype=float) + lead_hours = [] + valid_times = [] + for lead in range(1, n_lead + 1): + idx = t0_idx + lead + if idx >= len(dates_list): + break + valid_date = dates_list[idx] + p, e = forcing_by_date[valid_date] + p_arr, e_arr = enkf_fcst.perturb_forcing(p, e) + q = step_ensemble(fcst_models, p_arr, e_arr) + q_matrix[lead - 1, :] = q + lead_hours.append(lead) + valid_times.append(valid_date) + return lead_hours, valid_times, q_matrix + + +# ---------------- Main run --------------------------------------------------- + +def run(args, bmi_cfe, EnKFAssimilator): + cat_id = args.cat_id + out_dir = Path(args.out_dir) / cat_id + out_dir.mkdir(parents=True, exist_ok=True) + + # Combined test forcing (same idiom as the per-member script) + f1 = os.path.join(args.test_forcing_dir1, f'{cat_id}.csv') + f2 = os.path.join(args.test_forcing_dir2, f'{cat_id}.csv') + df1 = pd.read_csv(f1) + df2 = pd.read_csv(f2) + combined = pd.concat([df1, df2], ignore_index=True) + combined = combined.drop_duplicates(subset='time').sort_values('time') + test_forcing_file = str(out_dir / f'{cat_id}_nwm_operational_combined.csv') + combined.to_csv(test_forcing_file, index=False) + + # Pre-staged calibrated params + best_params_file = out_dir / f'{cat_id}_best_params.json' + if not best_params_file.exists(): + print(f"Missing {best_params_file}. Pre-stage from Run 3.") + return + with open(best_params_file) as f: + best = json.load(f)['best_parameters'] + + # Build temp CFE config with calibrated params + with open(args.config_file) as f: + cfg = json.load(f) + cfg['forcing_file'] = test_forcing_file + cfg['soil_params']['bb'] = best['bb'] + cfg['soil_params']['smcmax'] = best['smcmax'] + cfg['soil_params']['satdk'] = best['satdk'] + cfg['slop'] = best['slop'] + cfg['max_gw_storage'] = best['max_gw_storage'] + cfg['expon'] = best['expon'] + cfg['Cgw'] = best['Cgw'] + cfg['K_lf'] = best['K_lf'] + cfg['K_nash'] = best['K_nash'] + cfg['partition_scheme'] = "Schaake" if best['scheme'] <= 0.5 else "Xinanjiang" + tmp_cfg = str(out_dir / f'{cat_id}_bmi_config_temp_leadtime.json') + with open(tmp_cfg, 'w') as f: + json.dump(cfg, f) + + # Three EnKF instances: + # enkf_da — runs the DA-on trajectory; consumes obs each hour + # enkf_ol — runs the open-loop trajectory; never calls update_states + # enkf_fcst — used only inside the 18-hour free-runs (perturb_forcing only) + # All three share defaults with production. Distinct RNG seeds keep them + # independent so the openloop and forecast spreads are not coupled to DA. + obs_file = os.path.join(args.obs_dir, f'{cat_id}.csv') + seed_base = args.rng_seed if args.rng_seed is not None else 0 + enkf_da = EnKFAssimilator( + n_members=args.enkf_members, obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, use_vrugt_r=(not args.no_vrugt_r), + vrugt_alpha=args.vrugt_alpha, vrugt_scale=args.vrugt_scale, + rng_seed=seed_base if args.rng_seed is not None else None, + ) + enkf_ol = EnKFAssimilator( + n_members=args.enkf_members, obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, use_vrugt_r=(not args.no_vrugt_r), + vrugt_alpha=args.vrugt_alpha, vrugt_scale=args.vrugt_scale, + rng_seed=(seed_base + 1) if args.rng_seed is not None else None, + ) + enkf_fcst = EnKFAssimilator( + n_members=args.enkf_members, obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, use_vrugt_r=(not args.no_vrugt_r), + vrugt_alpha=args.vrugt_alpha, vrugt_scale=args.vrugt_scale, + rng_seed=(seed_base + 2) if args.rng_seed is not None else None, + ) + # Optional R override — replace every per-hour obs variance with a constant. + # Matches DualEarth/new_EnKF.py (R=0.07). Set via --hardcoded-r on the CLI. + # Effect: at storm peak ≈ existing Vrugt R, at low flow ≈ 15-17× larger → + # near-zero gain at low flow where kriging obs is noisy. + if args.hardcoded_r is not None: + for enkf in (enkf_da, enkf_ol, enkf_fcst): + for date in enkf.obs_var_dict: + enkf.obs_var_dict[date] = args.hardcoded_r + print(f"[lead-time] R hardcoded to {args.hardcoded_r} mm^2/h^2 " + f"(overrides Vrugt + kriging variance formula)") + + N = enkf_da.n_members + print(f"[lead-time] {cat_id} | N={N} | DA on/off both run | " + f"forecast lead = {FORECAST_LEAD_HOURS}h") + + # Build three ensembles (init-perturbed except the forecast ensemble, whose + # state will be overwritten at each issue time before each free-run) + prod_models = build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, enkf_da) + openloop_models = build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, enkf_ol) + fcst_models = build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, + enkf_fcst, apply_init_perturbation=False) + + # ----- Spinup with perturbed forcing on both trajectories ----- + df = load_test_forcing(test_forcing_file) + sp_mask = (df['date'] >= SPINUP_START) & (df['date'] <= SPINUP_END) + df_sp = df[sp_mask] + print(f"[lead-time] spinup: {len(df_sp)} hours") + for p, e in zip(df_sp['total_precipitation'], df_sp['potential_evaporation']): + p_da, e_da = enkf_da.perturb_forcing(p, e) + step_ensemble(prod_models, p_da, e_da) + enkf_da.add_process_noise(prod_models) + p_ol, e_ol = enkf_ol.perturb_forcing(p, e) + step_ensemble(openloop_models, p_ol, e_ol) + enkf_ol.add_process_noise(openloop_models) + + # ----- Test period: step both, fork forecasts at issue times ----- + t_mask = (df['date'] >= TEST_START) & (df['date'] <= TEST_END) + df_test = df[t_mask].reset_index(drop=True) + dates_list = list(df_test['date'].values) + forcing_by_date = dict(zip( + df_test['date'].values, + zip(df_test['total_precipitation'].values, + df_test['potential_evaporation'].values))) + + issue_times = build_issue_time_schedule( + dates_list, args.base_step_h, + args.dense_start, args.dense_end) + issue_set = set(issue_times) + print(f"[lead-time] test period: {len(dates_list)} hours | " + f"{len(issue_times)} issue times " + f"(base={args.base_step_h}h, dense={args.dense_start or 'none'}" + f"..{args.dense_end or 'none'})") + + # Pre-build the date→index map so forecast windows are fast to look up + date_to_idx = {d: i for i, d in enumerate(dates_list)} + + da_rows = [] # (issue_time, lead_hour, valid_time, *member_values) + ol_rows = [] + sanity_done = False + + for h, current_date in enumerate(dates_list): + p, e = forcing_by_date[current_date] + + # --- DA trajectory step --- + p_da, e_da = enkf_da.perturb_forcing(p, e) + q_da = step_ensemble(prod_models, p_da, e_da) + enkf_da.update_states(prod_models, current_date, q_da) + enkf_da.add_process_noise(prod_models) + + # --- Open-loop trajectory step (no DA) --- + p_ol, e_ol = enkf_ol.perturb_forcing(p, e) + q_ol = step_ensemble(openloop_models, p_ol, e_ol) + enkf_ol.add_process_noise(openloop_models) + + # --- One-shot sanity check at the first scheduled issue time --- + if args.sanity_check and not sanity_done and current_date in issue_set: + ok = do_sanity_check( + prod_models, fcst_models, dates_list, forcing_by_date, h, N) + sanity_done = True + print(f"[sanity-check] exiting (re-run without --sanity-check for " + f"the full sweep). Result: {'OK' if ok else 'FAIL'}") + for m in prod_models + openloop_models + fcst_models: + m.finalize() + sys.exit(0 if ok else 1) + + # --- Fork forecasts at scheduled issue times --- + if current_date in issue_set: + for scenario_label, src_models, rows_acc in [ + ('da', prod_models, da_rows), + ('openloop', openloop_models, ol_rows), + ]: + snap = snapshot_states(src_models) + restore_states(fcst_models, snap) + leads, valids, qm = run_forecast( + fcst_models, enkf_fcst, + dates_list, forcing_by_date, + t0_idx=h, n_lead=FORECAST_LEAD_HOURS) + for k, (lead, vt) in enumerate(zip(leads, valids)): + rows_acc.append((current_date, lead, vt, *qm[k, :])) + + for m in prod_models + openloop_models + fcst_models: + m.finalize() + + # ----- Save forecast CSVs ----- + cols = ['issue_time', 'lead_hour', 'valid_time'] + [f'member_{i:02d}' for i in range(N)] + da_path = out_dir / f'{cat_id}_lead_time_forecasts_da.csv' + ol_path = out_dir / f'{cat_id}_lead_time_forecasts_openloop.csv' + pd.DataFrame(da_rows, columns=cols).to_csv(da_path, index=False) + pd.DataFrame(ol_rows, columns=cols).to_csv(ol_path, index=False) + print(f"[lead-time] saved {da_path}") + print(f"[lead-time] saved {ol_path}") + + # Issue-time schedule (small, useful for downstream scripts) + sched_path = out_dir / f'{cat_id}_lead_time_issue_times.csv' + pd.DataFrame({'issue_time': issue_times}).to_csv(sched_path, index=False) + print(f"[lead-time] saved {sched_path}") + + print(f"[lead-time] DA assimilator: updates={enkf_da.n_updates} | " + f"avg Pyy={enkf_da.avg_pyy:.6f} | " + f"mass lost={enkf_da.total_overflow_lost_mm:.3f} mm") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', required=True) + parser.add_argument('--forcing-dir', required=True) + parser.add_argument('--obs-dir', required=True) + parser.add_argument('--cfe-dir', required=True) + parser.add_argument('--config-file', required=True) + parser.add_argument('--param-bounds', required=True, + help='Kept for symmetry; not read in this script') + parser.add_argument('--out-dir', required=True) + parser.add_argument('--test-forcing-dir1', required=True) + parser.add_argument('--test-forcing-dir2', required=True) + parser.add_argument('--enkf-members', type=int, default=20) + parser.add_argument('--enkf-obs-error-std', type=float, default=0.05) + parser.add_argument('--no-vrugt-r', action='store_true') + parser.add_argument('--vrugt-alpha', type=float, default=0.10) + parser.add_argument('--vrugt-scale', type=float, default=0.001) + parser.add_argument('--rng-seed', type=int, default=None) + parser.add_argument('--base-step-h', type=int, default=6, + help='Issue-time cadence in hours across the full test ' + 'period (default 6h)') + parser.add_argument('--dense-start', type=str, default=DEFAULT_DENSE_START, + help='Start of hourly-densification window ' + '(default 2024-09-24 00:00:00 — Helene)') + parser.add_argument('--dense-end', type=str, default=DEFAULT_DENSE_END, + help='End of hourly-densification window ' + '(default 2024-09-28 23:00:00 — Helene)') + parser.add_argument('--sanity-check', action='store_true', + help='At the first issue time, fork fcst_models and step ' + 'both prod and fcst one hour with identical unperturbed ' + 'forcing. Confirms the 4-state snapshot is sufficient ' + 'to reproduce next-hour Q to machine precision, then ' + 'exits. Use before kicking off the full sweep.') + parser.add_argument('--hardcoded-r', type=float, default=None, + help='If set, override the Vrugt + kriging-variance R formula ' + 'and use this constant value (mm^2/h^2) at every hour. ' + 'Matches the DualEarth/new_EnKF.py reference (R=0.07). ' + 'Suppresses DA at low flow where kriging is noisy while ' + 'preserving DA strength at storm peaks. Applied to all ' + 'three EnKF instances (da, openloop, fcst) for consistency.') + parser.add_argument('--prod-script', default=None, + help='Path to calibrate_catchment_cfe_da_v2.py for ' + 'importing EnKFAssimilator. Defaults to next-to-this-file.') + args = parser.parse_args() + + here = os.path.dirname(os.path.abspath(__file__)) + prod_script = args.prod_script or os.path.join(here, 'calibrate_catchment_cfe_da_v2.py') + if not os.path.exists(prod_script): + raise FileNotFoundError( + f"Could not find production script at {prod_script}. " + f"Pass --prod-script to override.") + EnKFAssimilator = import_enkf_class(prod_script) + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + + run(args, _bmi_cfe, EnKFAssimilator) + + +if __name__ == '__main__': + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/run_perturbation_da_on.py b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/run_perturbation_da_on.py new file mode 100644 index 00000000..b9ef7655 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/run_perturbation_da_on.py @@ -0,0 +1,392 @@ +""" +run_perturbation_da_on.py — 2a/2b: ensemble spread with DA ON. + +Runs two separated perturbation arms for the Helene window (Sep 24-30), +with DA actively running during the analysis period: + + ARM A — forcing only (30 members): + DA analyzes states using Qkrig at every hour (DA ON). + At each issue_time: fork 30 members with ONLY met forcing perturbed + (different precip/PET draw per member). Hydro states = DA analysis + mean (same for all 30). Free-run 18 hours, no further DA. + + ARM B — hydro states only (20 members): + Same DA analysis. At each issue_time: fork 20 members with ONLY + hydro states perturbed (multiplicative noise on soil/GW/Nash from + the DA analysis ensemble). Forcing = deterministic. Free-run 18 hours. + +These two arms isolate "where does forecast spread come from WITH DA?" +— figures for the assimilation script section. + +Output (per catchment, per arm): + //_da_forcing_arm.csv + //_da_hydro_arm.csv + Columns: issue_time, lead_hour, member_00..member_29 (forcing arm) + issue_time, lead_hour, member_00..member_19 (hydro arm) + +Usage (run for each catchment, both arms): + python3 run_perturbation_da_on.py \\ + --cat-id cat-1016300 \\ + --forcing-dir \\ + --obs-dir /mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance \\ + --cfe-dir /mnt/disk2/suma_helen_poster/cfe_py \\ + --config-file /mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json \\ + --param-bounds /mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json \\ + --out-dir /mnt/disk2/suma_helen_poster/da_results/v2_perturbation_da_on \\ + --test-forcing-dir1 \\ + --test-forcing-dir2 \\ + --hardcoded-r 0.07 +""" + +import argparse +import os +import sys +import json +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +N_FORCING = 30 # members in forcing arm +N_HYDRO = 20 # members in hydro-state arm + +# Analysis period — run DA through entire test period +TEST_START = "2024-08-24 00:00:00" # 1 month spinup before Helene window +TEST_END = "2024-10-31 23:00:00" + +# Issue times for the Helene window (hourly, Sep 24-30) +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-30 06:00:00") +FORECAST_HOURS = 18 + +# Perturbation magnitudes (same as production) +PRECIP_SIGMA = 0.15 +PET_SIGMA = 0.10 +STATE_FRAC = 0.05 # hydro-state perturbation fraction + +# Set at runtime +CAT_ID = None +OBS_FILE = None +CFE_CONFIG_FILE = None +TEST_FORCING_FILE = None +OUT_DIR = None +HARDCODED_R = None +bmi_cfe = None + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(): + df = pd.read_csv(TEST_FORCING_FILE) + df["date"] = pd.to_datetime(df["time"]).dt.strftime("%Y-%m-%d %H:%M:%S") + df["total_precipitation"] = df["APCP_surface"] * 3600.0 + df["potential_evaporation"] = priestley_taylor_pet( + df["DSWRF_surface"].values, df["TMP_2maboveground"].values) + return df + + +def load_obs(): + df = pd.read_csv(OBS_FILE) + date_col = next((c for c in df.columns + if c.lower() in ("date", "time", "datetime", "timestamp")), None) + if date_col is None: + raise ValueError(f"No date column in {OBS_FILE}. Columns: {df.columns.tolist()}") + q_col = next((c for c in df.columns if "qkrig" in c.lower() and "var" not in c.lower()), None) + if q_col is None: + raise ValueError(f"No qkrig column in {OBS_FILE}. Columns: {df.columns.tolist()}") + var_col = next((c for c in df.columns if "var" in c.lower()), None) + df[date_col] = pd.to_datetime(df[date_col]) + df = df.sort_values(date_col).reset_index(drop=True) + date_strs = df[date_col].dt.strftime("%Y-%m-%d %H:%M:%S") + obs_dict = dict(zip(date_strs, df[q_col])) + var_vals = df[var_col] if var_col else pd.Series([0.07] * len(df)) + var_dict = dict(zip(date_strs, var_vals)) + print(f" Obs: date='{date_col}', q='{q_col}', var='{var_col}'") + return obs_dict, var_dict + + +def make_model(best_params): + with open(CFE_CONFIG_FILE) as f: + cfg = json.load(f) + cfg["forcing_file"] = TEST_FORCING_FILE + cfg["soil_params"]["bb"] = best_params["bb"] + cfg["soil_params"]["smcmax"] = best_params["smcmax"] + cfg["soil_params"]["satdk"] = best_params["satdk"] + cfg["slop"] = best_params["slop"] + cfg["max_gw_storage"] = best_params["max_gw_storage"] + cfg["expon"] = best_params["expon"] + cfg["Cgw"] = best_params["Cgw"] + cfg["K_lf"] = best_params["K_lf"] + cfg["K_nash"] = best_params["K_nash"] + cfg["partition_scheme"] = ("Schaake" if best_params["scheme"] <= 0.5 + else "Xinanjiang") + tmp_cfg = str(OUT_DIR / f"{CAT_ID}_bmi_config_temp_da_on.json") + with open(tmp_cfg, "w") as f: + json.dump(cfg, f) + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.initialize() + return m + + +def get_state(m): + return { + "soil_m": m.soil_reservoir["storage_m"], + "gw_m": m.gw_reservoir["storage_m"], + "nash0": float(m.nash_storage[0]), + "nash1": float(m.nash_storage[1]), + } + + +def set_state(m, state): + m.soil_reservoir["storage_m"] = state["soil_m"] + m.gw_reservoir["storage_m"] = state["gw_m"] + m.nash_storage[0] = state["nash0"] + m.nash_storage[1] = state["nash1"] + + +def perturb_state(state, rng, frac=STATE_FRAC): + sm_max = 1.0 # placeholder — will be clipped in set_state; actual max enforced by model + out = { + "soil_m": max(state["soil_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "nash0": max(state["nash0"] * (1.0 + frac * rng.standard_normal()), 0.0), + "nash1": max(state["nash1"] * (1.0 + frac * rng.standard_normal()), 0.0), + } + return out + + +def run_model_step(m, precip_mmh, pet_mmh): + m.set_value("atmosphere_water__time_integral_of_precipitation_mass_flux", + precip_mmh / 1000.0) + m.set_value("water_potential_evaporation_flux", + pet_mmh / 1000.0 / 3600.0) + m.update() + return m.get_value("land_surface_water__runoff_depth") * 1000.0 # m/h -> mm/h + + +def enkf_update(state, q_sim, y_obs, krig_var, rng, n_members=1): + """Simple scalar EnKF update on the mean analysis state.""" + if HARDCODED_R is not None: + R = HARDCODED_R + else: + alpha = 0.10 + R = (alpha * max(y_obs, 0.0)) ** 2 + 0.001 * krig_var + R = max(R, 1e-6) + P_yy = max(q_sim * 0.01, 1e-6) # rough prior variance on Q + K = P_yy / (P_yy + R) + innovation = y_obs - q_sim + # Apply correction to each state proportionally (simple scalar gain) + scale = K * innovation / max(abs(q_sim), 1e-6) + updated = { + "soil_m": max(state["soil_m"] * (1.0 + scale), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + scale), 1e-6), + "nash0": max(state["nash0"] + state["nash0"] * scale, 0.0), + "nash1": max(state["nash1"] + state["nash1"] * scale, 0.0), + } + return updated + + +def main(): + global CAT_ID, OBS_FILE, CFE_CONFIG_FILE, TEST_FORCING_FILE + global OUT_DIR, HARDCODED_R, bmi_cfe + + parser = argparse.ArgumentParser() + parser.add_argument("--cat-id", required=True) + parser.add_argument("--forcing-dir", required=True) + parser.add_argument("--obs-dir", required=True) + parser.add_argument("--cfe-dir", required=True) + parser.add_argument("--config-file", required=True) + parser.add_argument("--param-bounds", required=True) + parser.add_argument("--out-dir", required=True) + parser.add_argument("--test-forcing-dir1", default=None) + parser.add_argument("--test-forcing-dir2", default=None) + parser.add_argument("--hardcoded-r", type=float, default=0.07) + args = parser.parse_args() + + CAT_ID = args.cat_id + OBS_FILE = os.path.join(args.obs_dir, f"{CAT_ID}.csv") + CFE_CONFIG_FILE = args.config_file + HARDCODED_R = args.hardcoded_r + OUT_DIR = Path(args.out_dir) / CAT_ID + OUT_DIR.mkdir(parents=True, exist_ok=True) + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + bmi_cfe = _bmi_cfe + + # Build combined test forcing + if args.test_forcing_dir1 and args.test_forcing_dir2: + f1 = os.path.join(args.test_forcing_dir1, f"{CAT_ID}.csv") + f2 = os.path.join(args.test_forcing_dir2, f"{CAT_ID}.csv") + df1, df2 = pd.read_csv(f1), pd.read_csv(f2) + combined = (pd.concat([df1, df2], ignore_index=True) + .drop_duplicates(subset="time").sort_values("time")) + TEST_FORCING_FILE = str(OUT_DIR / f"{CAT_ID}_nwm_operational_combined.csv") + combined.to_csv(TEST_FORCING_FILE, index=False) + else: + print("Need --test-forcing-dir1 and --test-forcing-dir2"); return + + best_params_file = OUT_DIR / f"{CAT_ID}_best_params.json" + if not best_params_file.exists(): + print(f"Missing best params at {best_params_file}"); return + with open(best_params_file) as f: + best_params = json.load(f)["best_parameters"] + + obs_dict, var_dict = load_obs() + df_forcing = load_test_forcing() + t_mask = (df_forcing["date"] >= TEST_START) & (df_forcing["date"] <= TEST_END) + df_test = df_forcing[t_mask].reset_index(drop=True) + + rng = np.random.default_rng(hash(CAT_ID) & 0x7fffffff) + + # ------------------------------------------------------------------ # + # Phase 1: Run DA through test period, snapshot analysis state at # + # each Helene issue time # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Phase 1: DA analysis through test period...") + model = make_model(best_params) + + helene_snapshots = {} # issue_time -> analysis state dict + + for h, row in df_test.iterrows(): + date_str = row["date"] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + q_sim = run_model_step(model, P, E) + + date_ts = pd.Timestamp(date_str) + + # DA update if obs available + y_obs = obs_dict.get(date_str, np.nan) + if not np.isnan(y_obs): + krig_var = var_dict.get(date_str, 0.07) + state = get_state(model) + updated = enkf_update(state, q_sim, y_obs, krig_var, rng) + set_state(model, updated) + + # Snapshot if this is a Helene issue time + if HELENE_START <= date_ts <= HELENE_END: + helene_snapshots[date_ts] = get_state(model) + + model.finalize() + print(f" Snapshots saved: {len(helene_snapshots)} issue times") + + # ------------------------------------------------------------------ # + # Phase 2A: Forcing arm — 30 members, fixed analyzed state, varying # + # met forcing # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Phase 2A: Forcing arm ({N_FORCING} members)...") + forcing_records = [] + + issue_times = sorted(helene_snapshots.keys()) + t0_to_idx = {t: i for i, t in enumerate(df_test["date"].apply(pd.Timestamp))} + + for t0 in issue_times: + snap = helene_snapshots[t0] + t0_str = t0.strftime("%Y-%m-%d %H:%M:%S") + + # Find index in df_test + if t0 not in t0_to_idx: + continue + start_idx = t0_to_idx[t0] + + q_matrix = np.full((FORECAST_HOURS, N_FORCING), np.nan) + + for mem in range(N_FORCING): + m = make_model(best_params) + set_state(m, snap) # all members start from same DA analysis state + + for lead in range(1, FORECAST_HOURS + 1): + fi = start_idx + lead + if fi >= len(df_test): + break + row = df_test.iloc[fi] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + # Perturb forcing only + mu_p = -0.5 * PRECIP_SIGMA ** 2 + P_pert = float(P * rng.lognormal(mu_p, PRECIP_SIGMA)) if P > 0 else 0.0 + E_pert = max(float(E * (1.0 + PET_SIGMA * rng.standard_normal())), 0.0) + q = run_model_step(m, P_pert, E_pert) + q_matrix[lead - 1, mem] = q + + m.finalize() + + row_data = {"issue_time": t0_str} + for lead in range(1, FORECAST_HOURS + 1): + row_data["lead_hour"] = lead + for mem in range(N_FORCING): + row_data[f"member_{mem:02d}"] = q_matrix[lead - 1, mem] + forcing_records.append(dict(row_data)) + + df_forcing_arm = pd.DataFrame(forcing_records) + out_a = OUT_DIR / f"{CAT_ID}_da_forcing_arm.csv" + df_forcing_arm.to_csv(out_a, index=False) + print(f" Saved: {out_a}") + + # ------------------------------------------------------------------ # + # Phase 2B: Hydro-state arm — 20 members, deterministic forcing, # + # perturbed analysis states # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Phase 2B: Hydro-state arm ({N_HYDRO} members)...") + hydro_records = [] + + for t0 in issue_times: + snap = helene_snapshots[t0] + t0_str = t0.strftime("%Y-%m-%d %H:%M:%S") + + if t0 not in t0_to_idx: + continue + start_idx = t0_to_idx[t0] + + q_matrix = np.full((FORECAST_HOURS, N_HYDRO), np.nan) + + for mem in range(N_HYDRO): + # Perturb the DA analysis state for this member + if mem == 0: + state_m = snap # member 0 = unperturbed analysis mean + else: + state_m = perturb_state(snap, rng) + + m = make_model(best_params) + set_state(m, state_m) + + for lead in range(1, FORECAST_HOURS + 1): + fi = start_idx + lead + if fi >= len(df_test): + break + row = df_test.iloc[fi] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + # Deterministic forcing + q = run_model_step(m, P, E) + q_matrix[lead - 1, mem] = q + + m.finalize() + + row_data = {"issue_time": t0_str} + for lead in range(1, FORECAST_HOURS + 1): + row_data["lead_hour"] = lead + for mem in range(N_HYDRO): + row_data[f"member_{mem:02d}"] = q_matrix[lead - 1, mem] + hydro_records.append(dict(row_data)) + + df_hydro_arm = pd.DataFrame(hydro_records) + out_b = OUT_DIR / f"{CAT_ID}_da_hydro_arm.csv" + df_hydro_arm.to_csv(out_b, index=False) + print(f" Saved: {out_b}") + print(f"[{CAT_ID}] Done.") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/run_perturbation_sensitivity.py b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/run_perturbation_sensitivity.py new file mode 100644 index 00000000..57b8d185 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/run_perturbation_sensitivity.py @@ -0,0 +1,293 @@ +""" +Perturbation-source sensitivity analysis for the multi-catchment CFE DA experiment. + +Ensemble-spread attribution: run 20 CFE members through the test +period with ONLY ONE perturbation source active at a time, no DA, no spinup. + +Three sub-experiments (selected via --source): + init : perturb initial states once at t=0 (multiplicative ~5%) + forcing : perturb hourly precip (lognormal sigma=0.15) and PET (Gaussian sigma=0.10) + process : apply per-hour process noise on states (soil 0.2%, GW 0.15%, Nash 0.5%) + +Each sub-experiment runs N=20 CFE members. DA is always OFF. Spinup is skipped — each +member starts from CFE's default initial state (modified by the perturbation if --source +init). This is the cleanest setup for attributing test-period spread to its source. + +Output: + //_sensitivity_.csv + Columns: date, member_00, member_01, ..., member_19 + +Each member column is that member's hourly Q_sim (mm/h) over the test period. + +Reads the same calibrated best_params.json as the production script. Does not need +the kriging obs file (no DA), but takes the same --obs-dir argument so it can also +log Q_obs alongside for plotting reference. + +Usage: + python3 run_perturbation_sensitivity.py \ + --cat-id cat-1016300 \ + --source forcing \ + --forcing-dir \ + --obs-dir \ + --cfe-dir \ + --config-file \ + --param-bounds \ + --out-dir \ + --test-forcing-dir1 \ + --test-forcing-dir2 +""" + +import argparse +import os +import sys +import json +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 +N_MEMBERS = 20 + +# We only run the test window for this sensitivity analysis. +TEST_START = "2023-10-01 00:00:00" +TEST_END = "2024-10-31 23:00:00" + +# Perturbation magnitudes — same as production defaults +INIT_STATE_FRAC = 0.05 +PRECIP_PERTURB = 0.15 # lognormal sigma +PET_PERTURB = 0.10 # Gaussian sigma +SOIL_PROC_NOISE = 0.002 +GW_PROC_NOISE = 0.0015 +NASH_PROC_NOISE = 0.005 + +# Set at runtime +CAT_ID = None +FORCING_FILE = None +TEST_FORCING_FILE = None +OBS_FILE = None +CFE_CONFIG_FILE = None +OUT_DIR = None +bmi_cfe = None + + +# ---------- Forcing helpers (same as production) ---------- +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(): + df = pd.read_csv(TEST_FORCING_FILE) + df['date'] = pd.to_datetime(df['time']).dt.strftime('%Y-%m-%d %H:%M:%S') + df['total_precipitation'] = df['APCP_surface'] * 3600.0 # kg/m^2/s -> mm/h + df['potential_evaporation'] = priestley_taylor_pet( + df['DSWRF_surface'].values, + df['TMP_2maboveground'].values, + ) + return df + + +# ---------- Perturbation helpers ---------- +def perturb_precip_lognormal(P, rng, sigma=PRECIP_PERTURB): + """Lognormal multiplier with mean 1; preserves zero precip.""" + mu = -0.5 * sigma ** 2 + return float(P * rng.lognormal(mu, sigma)) + + +def perturb_pet_gaussian(E, rng, sigma=PET_PERTURB): + """Gaussian multiplicative noise on PET, clipped at 0.""" + return max(float(E * (1.0 + sigma * rng.standard_normal())), 0.0) + + +def apply_init_state_perturb(m, rng, frac=INIT_STATE_FRAC): + """One-time multiplicative perturbation on the four states (called per member at t=0).""" + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1.0 + frac * rng.standard_normal()), 1e-6, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1.0 + frac * rng.standard_normal()), 1e-6, gw_max)) + # Nash buckets often start at 0; small additive jitter so they're not identical. + jitter = max(sm_max * 1e-6, 1e-9) + m.nash_storage[0] = max(n00 + jitter * rng.standard_normal(), 0.0) + m.nash_storage[1] = max(n10 + jitter * rng.standard_normal(), 0.0) + + +def apply_process_noise(m, rng): + """Per-timestep process noise on all 4 states for one member.""" + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + nash_floor = max(sm_max * 1e-4, 1e-7) + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1.0 + SOIL_PROC_NOISE * rng.standard_normal()), 0.0, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1.0 + GW_PROC_NOISE * rng.standard_normal()), 0.0, gw_max)) + m.nash_storage[0] = max( + n00 * (1.0 + NASH_PROC_NOISE * rng.standard_normal()) + + nash_floor * rng.standard_normal(), 0.0) + m.nash_storage[1] = max( + n10 * (1.0 + NASH_PROC_NOISE * rng.standard_normal()) + + nash_floor * rng.standard_normal(), 0.0) + + +# ---------- Main per-source run ---------- +def run_sensitivity(source, best_param_dict): + """Run N members through the test period with only `source` perturbation active.""" + + def custom_load_forcing(self_cfe): + df = load_test_forcing() + self_cfe.forcing_data = df.rename(columns={"date": "time"}) + + # Build a per-catchment temporary CFE config with the calibrated params + with open(CFE_CONFIG_FILE) as f: + cfg = json.load(f) + cfg['forcing_file'] = TEST_FORCING_FILE + cfg['soil_params']['bb'] = best_param_dict['bb'] + cfg['soil_params']['smcmax'] = best_param_dict['smcmax'] + cfg['soil_params']['satdk'] = best_param_dict['satdk'] + cfg['slop'] = best_param_dict['slop'] + cfg['max_gw_storage'] = best_param_dict['max_gw_storage'] + cfg['expon'] = best_param_dict['expon'] + cfg['Cgw'] = best_param_dict['Cgw'] + cfg['K_lf'] = best_param_dict['K_lf'] + cfg['K_nash'] = best_param_dict['K_nash'] + cfg['partition_scheme'] = "Schaake" if best_param_dict['scheme'] <= 0.5 else "Xinanjiang" + + tmp_cfg = str(OUT_DIR / f'{CAT_ID}_bmi_config_temp_sensitivity_{source}.json') + with open(tmp_cfg, 'w') as f: + json.dump(cfg, f) + + # Deterministic seed per (catchment, source) so each run is reproducible + seed = hash((CAT_ID, source)) & 0x7fffffff + rng = np.random.default_rng(seed) + + print(f"[sensitivity] {CAT_ID} | source={source} | N={N_MEMBERS} | seed={seed}") + + # Build N members. Initial-state perturbation only happens when source == 'init'. + models = [] + for i in range(N_MEMBERS): + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.load_forcing_file = custom_load_forcing.__get__(m) + m.initialize() + if source == 'init' and i > 0: + apply_init_state_perturb(m, rng) + models.append(m) + + # Skip spinup. Go directly to the test period. + df = load_test_forcing() + t_mask = (df['date'] >= TEST_START) & (df['date'] <= TEST_END) + df_test = df[t_mask] + + # Pre-allocate per-member Q arrays + n_hours = len(df_test) + q_matrix = np.full((n_hours, N_MEMBERS), np.nan, dtype=float) + dates_out = df_test['date'].values + + for h, (p, e) in enumerate(zip(df_test['total_precipitation'], + df_test['potential_evaporation'])): + for i, m in enumerate(models): + if source == 'forcing': + p_i = perturb_precip_lognormal(p, rng) + e_i = perturb_pet_gaussian(e, rng) + else: + p_i = float(p) + e_i = float(e) + m.set_value('atmosphere_water__time_integral_of_precipitation_mass_flux', p_i / 1000) + m.set_value('water_potential_evaporation_flux', e_i / 1000 / 3600) + m.update() + q_matrix[h, i] = m.get_value('land_surface_water__runoff_depth') * 1000 # m/h -> mm/h + + # Process noise applied after each hour, but only when source == 'process' + if source == 'process': + for m in models: + apply_process_noise(m, rng) + + for m in models: + m.finalize() + + # Write per-member CSV: date + 20 columns + out = {'date': dates_out} + for i in range(N_MEMBERS): + out[f'member_{i:02d}'] = q_matrix[:, i] + df_out = pd.DataFrame(out) + out_path = OUT_DIR / f'{CAT_ID}_sensitivity_{source}.csv' + df_out.to_csv(out_path, index=False) + print(f"[sensitivity] saved {out_path}") + print(f"[sensitivity] mean ensemble spread (std across members) over test period: " + f"{float(q_matrix.std(axis=1).mean()):.5f} mm/h") + + +def main(): + global CAT_ID, FORCING_FILE, TEST_FORCING_FILE, OBS_FILE, CFE_CONFIG_FILE + global OUT_DIR, bmi_cfe + + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', required=True) + parser.add_argument('--source', required=True, + choices=['init', 'forcing', 'process'], + help="Which perturbation source to isolate") + parser.add_argument('--forcing-dir', required=True) + parser.add_argument('--obs-dir', required=True, + help='Kept for symmetry with production script; not used since DA is off') + parser.add_argument('--cfe-dir', required=True) + parser.add_argument('--config-file', required=True) + parser.add_argument('--param-bounds', required=True, + help='Kept for symmetry; not read in this script') + parser.add_argument('--out-dir', required=True) + parser.add_argument('--test-forcing-dir1', default=None) + parser.add_argument('--test-forcing-dir2', default=None) + args = parser.parse_args() + + CAT_ID = args.cat_id + FORCING_FILE = os.path.join(args.forcing_dir, f'{CAT_ID}.csv') + OBS_FILE = os.path.join(args.obs_dir, f'{CAT_ID}.csv') + CFE_CONFIG_FILE = args.config_file + OUT_DIR = Path(args.out_dir) / CAT_ID + OUT_DIR.mkdir(parents=True, exist_ok=True) + + # Build combined test forcing CSV (same idiom as production) + if args.test_forcing_dir1 and args.test_forcing_dir2: + f1 = os.path.join(args.test_forcing_dir1, f'{CAT_ID}.csv') + f2 = os.path.join(args.test_forcing_dir2, f'{CAT_ID}.csv') + if os.path.exists(f1) and os.path.exists(f2): + df1 = pd.read_csv(f1) + df2 = pd.read_csv(f2) + combined = pd.concat([df1, df2], ignore_index=True) + combined = combined.drop_duplicates(subset='time').sort_values('time') + TEST_FORCING_FILE = str(OUT_DIR / f'{CAT_ID}_nwm_operational_combined.csv') + combined.to_csv(TEST_FORCING_FILE, index=False) + else: + print(f"Warning: test forcing files not found for {CAT_ID}, exiting") + return + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + bmi_cfe = _bmi_cfe + + best_params_file = OUT_DIR / f'{CAT_ID}_best_params.json' + if not best_params_file.exists(): + print(f"No best params file at {best_params_file}. Pre-stage from Run 3.") + return + with open(best_params_file) as f: + saved = json.load(f) + best_param_dict = saved['best_parameters'] + + run_sensitivity(args.source, best_param_dict) + + +if __name__ == '__main__': + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/run_production_per_member.py b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/run_production_per_member.py new file mode 100644 index 00000000..d9c1c4ac --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/run_production_per_member.py @@ -0,0 +1,314 @@ +""" +Run the production DA pipeline for ONE catchment and save per-member streamflow. + +Identical to calibrate_catchment_cfe_da_v2.py in algorithm (init perturbation + +forcing perturbation + process noise + true EnKF with Vrugt R), but writes one +column per ensemble member instead of only the ensemble mean. This is for +the per-member factor-decomposition plot. + +To keep the comparison apples-to-apples with the sensitivity runs (which skip +spinup), this script also SKIPS the spinup loop. Each member starts directly +from CFE's default initial state, plus the t=0 init perturbation. + +Output: + //_production_per_member.csv + Columns: date, member_00, member_01, ..., member_19 + +Imports EnKFAssimilator from calibrate_catchment_cfe_da_v2 to guarantee the +DA math matches production exactly. + +Usage: + python3 run_production_per_member.py \ + --cat-id cat-1016300 \ + --forcing-dir /mnt/.../nwm_retro_catchment_forcings \ + --obs-dir /mnt/.../catchment_ts_03463300_with_variance \ + --cfe-dir /mnt/.../cfe_py \ + --config-file /mnt/.../cat_03463300_bmi_config_cfe.json \ + --param-bounds /mnt/.../CFE_parameter_bounds.json \ + --out-dir /mnt/.../v2_production_per_member \ + --test-forcing-dir1 ... \ + --test-forcing-dir2 ... \ + --enkf-members 20 +""" + +import argparse +import os +import sys +import json +import importlib.util +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +TEST_START = "2023-10-01 00:00:00" +TEST_END = "2024-10-31 23:00:00" + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(test_forcing_file): + df = pd.read_csv(test_forcing_file) + df['date'] = pd.to_datetime(df['time']).dt.strftime('%Y-%m-%d %H:%M:%S') + df['total_precipitation'] = df['APCP_surface'] * 3600.0 + df['potential_evaporation'] = priestley_taylor_pet( + df['DSWRF_surface'].values, df['TMP_2maboveground'].values + ) + return df + + +def import_enkf_class(script_path): + """Load EnKFAssimilator from the production script without triggering its main().""" + spec = importlib.util.spec_from_file_location("prod_v2", script_path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod.EnKFAssimilator + + +def run(args, bmi_cfe, EnKFAssimilator): + cat_id = args.cat_id + out_dir = Path(args.out_dir) / cat_id + out_dir.mkdir(parents=True, exist_ok=True) + + # Combined test forcing + f1 = os.path.join(args.test_forcing_dir1, f'{cat_id}.csv') + f2 = os.path.join(args.test_forcing_dir2, f'{cat_id}.csv') + df1 = pd.read_csv(f1) + df2 = pd.read_csv(f2) + combined = pd.concat([df1, df2], ignore_index=True) + combined = combined.drop_duplicates(subset='time').sort_values('time') + test_forcing_file = str(out_dir / f'{cat_id}_nwm_operational_combined.csv') + combined.to_csv(test_forcing_file, index=False) + + # Load best_params + best_params_file = out_dir / f'{cat_id}_best_params.json' + if not best_params_file.exists(): + print(f"Missing {best_params_file}. Pre-stage from Run 3.") + return + with open(best_params_file) as f: + saved = json.load(f) + best = saved['best_parameters'] + + # Build temp CFE config with calibrated params + with open(args.config_file) as f: + cfg = json.load(f) + cfg['forcing_file'] = test_forcing_file + cfg['soil_params']['bb'] = best['bb'] + cfg['soil_params']['smcmax'] = best['smcmax'] + cfg['soil_params']['satdk'] = best['satdk'] + cfg['slop'] = best['slop'] + cfg['max_gw_storage'] = best['max_gw_storage'] + cfg['expon'] = best['expon'] + cfg['Cgw'] = best['Cgw'] + cfg['K_lf'] = best['K_lf'] + cfg['K_nash'] = best['K_nash'] + cfg['partition_scheme'] = "Schaake" if best['scheme'] <= 0.5 else "Xinanjiang" + tmp_cfg = str(out_dir / f'{cat_id}_bmi_config_temp_per_member.json') + with open(tmp_cfg, 'w') as f: + json.dump(cfg, f) + + # Build the EnKF assimilator (same defaults as production) + obs_file = os.path.join(args.obs_dir, f'{cat_id}.csv') + enkf = EnKFAssimilator( + n_members=args.enkf_members, + obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, + use_vrugt_r=(not args.no_vrugt_r) and (not getattr(args, 'direct_variance', False)), + vrugt_alpha=args.vrugt_alpha, + vrugt_scale=args.vrugt_scale, + rng_seed=args.rng_seed, + ) + N = enkf.n_members + print(f"[per-member] {cat_id} | N={N} | use_vrugt_r={enkf.use_vrugt_r}") + if getattr(args, 'hardcoded_r', None) is not None: + enkf.obs_var_dict = {k: args.hardcoded_r for k in enkf.obs_var_dict} + print(f"[per-member] obs_var_dict overridden: R = {args.hardcoded_r} (fixed)") + if getattr(args, 'direct_variance', False): + obs_df = pd.read_csv(obs_file) + t_col = next(c for c in obs_df.columns if c.lower() in ('datetime', 'date', 'time', 'timestamp')) + v_col = next(c for c in obs_df.columns if 'var' in c.lower()) + obs_df[t_col] = pd.to_datetime(obs_df[t_col]) + obs_df = obs_df.set_index(t_col).sort_index() + enkf.obs_var_dict = { + str(ts): float(v) + for ts, v in obs_df[v_col].items() + if pd.notna(v) and float(v) > 0 + } + print(f"[per-member] obs_var_dict overridden: R = σ² direct ({len(enkf.obs_var_dict)} timesteps)") + + # Custom forcing loader for CFE BMI + def custom_load_forcing(self_cfe): + df = load_test_forcing(test_forcing_file) + self_cfe.forcing_data = df.rename(columns={"date": "time"}) + + # Build N CFE instances; perturb initial states (member 0 left clean) + models = [] + init_states_per_member = [] # capture initial-state snapshot after perturbation + for i in range(N): + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.load_forcing_file = custom_load_forcing.__get__(m) + m.initialize() + if i > 0: + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + r = enkf.init_state_perturb_frac + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1 + r * enkf.rng.standard_normal()), 1e-6, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1 + r * enkf.rng.standard_normal()), 1e-6, gw_max)) + jitter = max(sm_max * 1e-6, 1e-9) + m.nash_storage[0] = max(n00 + jitter * enkf.rng.standard_normal(), 0.0) + m.nash_storage[1] = max(n10 + jitter * enkf.rng.standard_normal(), 0.0) + # Snapshot the (possibly-perturbed) initial state for this member + init_states_per_member.append({ + "soil_m": float(m.soil_reservoir["storage_m"]), + "gw_m": float(m.gw_reservoir["storage_m"]), + "nash0_m": float(m.nash_storage[0]), + "nash1_m": float(m.nash_storage[1]), + "soil_max_m": float(m.soil_reservoir["storage_max_m"]), + "gw_max_m": float(m.gw_reservoir["storage_max_m"]), + }) + models.append(m) + + # SKIP SPINUP — go directly to the test period (matches sensitivity-script setup) + df = load_test_forcing(test_forcing_file) + t_mask = (df['date'] >= TEST_START) & (df['date'] <= TEST_END) + df_test = df[t_mask] + n_hours = len(df_test) + q_matrix = np.full((n_hours, N), np.nan, dtype=float) + precip_matrix = np.full((n_hours, N), np.nan, dtype=float) # mm/h per member + pet_matrix = np.full((n_hours, N), np.nan, dtype=float) # mm/h per member + dates_out = df_test['date'].values + + for h, (p, e, current_date) in enumerate(zip( + df_test['total_precipitation'], + df_test['potential_evaporation'], + df_test['date'])): + + # Perturb forcing per member + p_arr, e_arr = enkf.perturb_forcing(p, e) + # Capture the per-member perturbed forcing for later plotting/diagnostics + precip_matrix[h, :] = p_arr + pet_matrix[h, :] = e_arr + + # Advance each member one hour + ensemble_q = np.empty(N, dtype=float) + for i, m in enumerate(models): + m.set_value('atmosphere_water__time_integral_of_precipitation_mass_flux', + float(p_arr[i]) / 1000) + m.set_value('water_potential_evaporation_flux', + float(e_arr[i]) / 1000 / 3600) + m.update() + ensemble_q[i] = m.get_value('land_surface_water__runoff_depth') * 1000 # mm/h + + # Record pre-analysis forecast per member (this is the "actual production" value) + q_matrix[h, :] = ensemble_q + + # DA step (uses production EnKFAssimilator) + enkf.update_states(models, current_date, ensemble_q) + + # Process noise after DA + enkf.add_process_noise(models) + + for m in models: + m.finalize() + + # ----- Save per-member CSVs ----- + def _save_matrix(matrix, suffix): + out = {'date': dates_out} + for i in range(N): + out[f'member_{i:02d}'] = matrix[:, i] + df_out = pd.DataFrame(out) + out_path = out_dir / f'{cat_id}_production_per_member_{suffix}.csv' + df_out.to_csv(out_path, index=False) + return out_path + + # Streamflow outputs + q_path = out_dir / f'{cat_id}_production_per_member.csv' + out = {'date': dates_out} + for i in range(N): + out[f'member_{i:02d}'] = q_matrix[:, i] + pd.DataFrame(out).to_csv(q_path, index=False) + print(f"[per-member] saved {q_path}") + + # Perturbed forcings (so each member's actual inputs are recoverable) + precip_path = _save_matrix(precip_matrix, "precip") + pet_path = _save_matrix(pet_matrix, "pet") + print(f"[per-member] saved {precip_path}") + print(f"[per-member] saved {pet_path}") + + # Initial states (small, JSON is fine) + init_states_path = out_dir / f'{cat_id}_production_per_member_initial_states.json' + with open(init_states_path, 'w') as f: + json.dump({ + "catchment_id": cat_id, + "n_members": N, + "members": {f"member_{i:02d}": init_states_per_member[i] for i in range(N)}, + }, f, indent=2) + print(f"[per-member] saved {init_states_path}") + + print(f"[per-member] ensemble mean spread over test period: " + f"{float(q_matrix.std(axis=1).mean()):.5f} mm/h") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', required=True) + parser.add_argument('--forcing-dir', required=True) + parser.add_argument('--obs-dir', required=True) + parser.add_argument('--cfe-dir', required=True) + parser.add_argument('--config-file', required=True) + parser.add_argument('--param-bounds', required=True, + help='Kept for symmetry; not read in this script') + parser.add_argument('--out-dir', required=True) + parser.add_argument('--test-forcing-dir1', required=True) + parser.add_argument('--test-forcing-dir2', required=True) + parser.add_argument('--enkf-members', type=int, default=20) + parser.add_argument('--enkf-obs-error-std', type=float, default=0.05) + parser.add_argument('--no-vrugt-r', action='store_true') + parser.add_argument('--hardcoded-r', type=float, default=None, + help='Fix R to this constant value for all timesteps (overrides Vrugt formula)') + parser.add_argument('--direct-variance', action='store_true', default=False, + help='Use kriging variance column as R per timestep (R=σ²; overrides Vrugt formula)') + parser.add_argument('--vrugt-alpha', type=float, default=0.10) + parser.add_argument('--vrugt-scale', type=float, default=0.001) + parser.add_argument('--rng-seed', type=int, default=None) + parser.add_argument('--prod-script', default=None, + help='Path to calibrate_catchment_cfe_da_v2.py for importing ' + 'EnKFAssimilator. Defaults to the script located next to this file.') + args = parser.parse_args() + + # Locate the production script (default: same dir as this script) + if args.prod_script is None: + here = os.path.dirname(os.path.abspath(__file__)) + prod_script = os.path.join(here, 'calibrate_catchment_cfe_da_v2.py') + else: + prod_script = args.prod_script + if not os.path.exists(prod_script): + raise FileNotFoundError(f"Could not find production script at {prod_script}. " + f"Pass --prod-script to override.") + EnKFAssimilator = import_enkf_class(prod_script) + + # Dynamic CFE import (same idiom as production script) + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + + run(args, _bmi_cfe, EnKFAssimilator) + + +if __name__ == '__main__': + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/run_sweep_all_21_hardcoded_r.py b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/run_sweep_all_21_hardcoded_r.py new file mode 100644 index 00000000..f5024cab --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/2_assimilation/run_sweep_all_21_hardcoded_r.py @@ -0,0 +1,134 @@ +""" +Launch hardcoded-R (R=0.07) lead-time sweep for all 21 catchments. + +- Discovers catchment IDs from the existing Vrugt-R output dir +- Pre-stages best_params.json from run3 into the hardcoded-R out dir +- Skips catchments that already completed +- Launches the sweep with bounded parallelism (each catchment runs ~60 BMI + instances, so 4-5 concurrent is usually safe on this box) +- Each sweep's stdout/stderr goes to ~/leadtime_logs/sweep_.log + +Run: + python3 run_sweep_all_21_hardcoded_r.py --concurrent 5 + nohup python3 run_sweep_all_21_hardcoded_r.py --concurrent 5 \ + > ~/sweep_all_21_master.log 2>&1 & + # then monitor: + tail -f ~/sweep_all_21_master.log + ls ~/leadtime_logs/ + +Each catchment ~6-12 hr. With concurrency=5, total wall clock ~2-4 batches +× 6-12 hr each ≈ overnight + a few hours. +""" +import argparse +import os +import subprocess +import sys +from concurrent.futures import ThreadPoolExecutor, as_completed +from pathlib import Path + +EXISTING_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast" +OUT_DIR_ROOT = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast_hardcoded_r" +RUN3_DIR = "/mnt/disk2/suma_helen_poster/catchment_results_range100_run3" +SWEEP_SCRIPT = os.path.expanduser("~/run_lead_time_forecast_sweep.py") +LOG_DIR = os.path.expanduser("~/leadtime_logs") + +COMMON_ARGS = [ + "--forcing-dir", "/mnt/disk2/suma_helen_poster/nwm_retro_catchment_forcings", + "--obs-dir", "/mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance", + "--cfe-dir", "/mnt/disk2/suma_helen_poster/cfe_py", + "--config-file", "/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json", + "--param-bounds", "/mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json", + "--out-dir", OUT_DIR_ROOT, + "--test-forcing-dir1", + "/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings", + "--test-forcing-dir2", + "/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings", + "--enkf-members", "20", + "--hardcoded-r", "0.07", +] + + +def stage_params(cat_id): + out_cat_dir = Path(OUT_DIR_ROOT) / cat_id + out_cat_dir.mkdir(parents=True, exist_ok=True) + src = Path(RUN3_DIR) / cat_id / f"{cat_id}_best_params.json" + dst = out_cat_dir / f"{cat_id}_best_params.json" + if dst.exists(): + return True + if not src.exists(): + print(f" [warn] no run3 params for {cat_id} at {src}") + return False + dst.write_text(src.read_text()) + return True + + +def run_one(cat_id): + if not stage_params(cat_id): + return cat_id, -1, "missing params" + log_path = Path(LOG_DIR) / f"sweep_{cat_id}.log" + cmd = ["python3", SWEEP_SCRIPT, "--cat-id", cat_id, *COMMON_ARGS] + with open(log_path, "w") as logf: + result = subprocess.run(cmd, stdout=logf, stderr=subprocess.STDOUT) + return cat_id, result.returncode, str(log_path) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--concurrent', type=int, default=5, + help='Max concurrent sweeps. Each uses ~60 BMI instances + ' + 'Python overhead, so 5 ≈ 300 BMIs in flight. ' + 'Drop to 3 if memory is tight.') + parser.add_argument('--dry-run', action='store_true', + help='Print plan and exit without launching.') + args = parser.parse_args() + + os.makedirs(LOG_DIR, exist_ok=True) + + # Discover catchments + catchments = sorted([ + d.name for d in Path(EXISTING_DIR).iterdir() + if d.is_dir() and d.name.startswith('cat-') + ]) + print(f"[plan] {len(catchments)} catchments discovered under {EXISTING_DIR}") + + # Filter to those not yet done + to_run = [] + for c in catchments: + done_marker = Path(OUT_DIR_ROOT) / c / f"{c}_lead_time_forecasts_da.csv" + if done_marker.exists(): + print(f" [skip] {c} — already complete") + else: + to_run.append(c) + print(f"[plan] {len(to_run)} to run | concurrent={args.concurrent}") + + if args.dry_run: + print("[dry-run] would launch:", to_run) + return + + if not to_run: + print("[plan] nothing to do — all catchments complete.") + return + + # Launch with bounded parallelism. ThreadPoolExecutor is fine since each + # task is a subprocess.run that releases the GIL while waiting. + with ThreadPoolExecutor(max_workers=args.concurrent) as ex: + futures = {ex.submit(run_one, c): c for c in to_run} + completed = 0 + failed = [] + for fut in as_completed(futures): + cat_id, rc, log = fut.result() + completed += 1 + status = "OK" if rc == 0 else f"FAIL (rc={rc})" + print(f"[done {completed}/{len(to_run)}] {cat_id} {status} — log: {log}", + flush=True) + if rc != 0: + failed.append(cat_id) + + print(f"\n[summary] {len(to_run) - len(failed)} succeeded, {len(failed)} failed") + if failed: + print(f" failed catchments: {failed}") + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/3_routing/route_det_f2.sh b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/3_routing/route_det_f2.sh new file mode 100644 index 00000000..674bdef6 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/3_routing/route_det_f2.sh @@ -0,0 +1,31 @@ +#!/usr/bin/env bash +# F2 fixed R=0.07 (20pct gauge holdout) — deterministic T-route routing. +# Routes per-catchment _test_results.csv (sim_mm_h column) through Muskingum-Cunge +# to gauge 03463300 and writes routed_Q_test.csv + KGE/NSE summary. +# +# Usage: +# bash route_det_f2.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +ROUTE_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/2_troute_routing/run_route.py + +F2_DA_DIR=/mnt/disk2/suma_helen_poster/da_results/v2_fixed_r007_analysis +F2_OUT_DIR=/mnt/disk2/suma_helen_poster/da_results/fixed_r007_routed +GPKG=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +mkdir -p "$F2_OUT_DIR" + +echo "[F2] Deterministic T-route routing..." +echo " da-dir : $F2_DA_DIR" +echo " out-dir: $F2_OUT_DIR" + +$TROUTE "$ROUTE_SCRIPT" \ + --gpkg "$GPKG" \ + --da-dir "$F2_DA_DIR" \ + --out-dir "$F2_OUT_DIR" \ + --usgs-csv "$USGS_CSV" + +echo "[F2] Deterministic routing done. Output: $F2_OUT_DIR/routed_Q_test.csv" diff --git a/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/3_routing/route_ensemble_f2.sh b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/3_routing/route_ensemble_f2.sh new file mode 100644 index 00000000..8311bffa --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/3_routing/route_ensemble_f2.sh @@ -0,0 +1,30 @@ +#!/usr/bin/env bash +# F2 fixed R=0.07 (20pct gauge holdout) — crossed-ensemble T-route routing. +# Routes all 600 crossed-ensemble members through Muskingum-Cunge +# to gauge 03463300 and writes routed_crossed_ensemble.parquet. +# +# Must be run AFTER crossed ensemble parquets exist in CROSSED_DIR. +# +# Usage: +# bash route_ensemble_f2.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +ROUTE_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/2_troute_routing/run_route_crossed_ensemble.py + +CROSSED_DIR=/mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble +GPKG=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +echo "[F2] Crossed-ensemble T-route routing..." +echo " ensemble-dir: $CROSSED_DIR" +echo " out-dir : $CROSSED_DIR" + +$TROUTE "$ROUTE_SCRIPT" \ + --gpkg "$GPKG" \ + --ensemble-dir "$CROSSED_DIR" \ + --out-dir "$CROSSED_DIR" \ + --usgs-csv "$USGS_CSV" + +echo "[F2] Ensemble routing done. Output: $CROSSED_DIR/routed_crossed_ensemble.parquet" diff --git a/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/3_routing/route_leadtime_f2.sh b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/3_routing/route_leadtime_f2.sh new file mode 100644 index 00000000..ba9d1b77 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/3_routing/route_leadtime_f2.sh @@ -0,0 +1,31 @@ +#!/usr/bin/env bash +# F2 fixed R=0.07 (20pct gauge holdout) — route lead-time forecast CSVs. +# Routes all 21-catchment lead-time forecast CSVs through Muskingum-Cunge +# to gauge 03463300. Writes routed_leadtime_da_full.parquet and +# routed_leadtime_openloop_full.parquet for gauge-level error decay analysis. +# +# Run AFTER the lead-time sweep batch finishes. +# +# Usage: +# bash route_leadtime_f2.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +ROUTE_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/2_troute_routing/run_route_leadtime_forecasts.py + +F2_DIR=/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast +GPKG=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg +OUT_DIR=/mnt/disk2/suma_helen_poster/leadtime_troute_routing_f2 + +echo "[F2] Routing lead-time forecasts through T-route..." +echo " forecast-dir: $F2_DIR" +echo " out-dir : $OUT_DIR" + +$TROUTE "$ROUTE_SCRIPT" \ + --gpkg "$GPKG" \ + --forecast-dir "$F2_DIR" \ + --out-dir "$OUT_DIR" + +echo "[F2] Lead-time routing done. Outputs in: $OUT_DIR" +ls "$OUT_DIR"/*.parquet 2>/dev/null || echo " (no parquets found)" diff --git a/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4a_error_decay/plot_forecast_error_fixed_target.py b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4a_error_decay/plot_forecast_error_fixed_target.py new file mode 100644 index 00000000..1460ed21 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4a_error_decay/plot_forecast_error_fixed_target.py @@ -0,0 +1,217 @@ +""" +plot_forecast_error_fixed_target.py — 4a: error decay, fixed-target-time view. + +For each target verification time T in the Helene peak window: + Collect all forecasts that verify AT T: + issue_time = T - lead_hour*1h, for lead in 1..18 + error[lead] = ensemble_mean(q at T, initialized T-lead) - USGS_obs(T) + +This gives the correct operational picture: + - lead 1 = initialized 1 hr before T (DA just ran -> small error) + - lead 18 = initialized 18 hr before T (DA long ago -> error ~ open loop) + +Two panels: + Top : signed error (m³/s) vs lead hour, one curve per target time + Bot : same for open-loop +Plus a summary panel: mean across all target times, DA vs OL. + +Outputs: + /error_fixed_target_helene.png (per-target spaghetti, DA vs OL) + /error_fixed_target_mean.png (mean across targets, DA vs OL) +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.cm as cm + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast_hardcoded_r_routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +# Target verification times: hourly through the Helene peak window +TARGET_START = pd.Timestamp("2024-09-26 18:00:00") +TARGET_END = pd.Timestamp("2024-09-28 06:00:00") + + +def load_parquet(path): + df = pd.read_parquet(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + df["ens_mean"] = df[member_cols].mean(axis=1) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + return df[["issue_time", "lead_hour", "valid_time", "ens_mean"]] + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() + or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def build_fixed_target_errors(df, obs_series, target_times): + """ + For each target time T and each lead L (1..18): + issue_time = T - L hours + error = ensemble_mean at (issue_time, lead=L) - obs(T) + Returns dict: target_time -> {lead: error} + """ + # Index df by (issue_time, lead_hour) for fast lookup + df_idx = df.set_index(["issue_time", "lead_hour"])["ens_mean"] + + results = {} + for T in target_times: + obs_val = obs_series.get(T, np.nan) + if np.isnan(obs_val): + continue + curve = {} + for lead in range(1, 19): + t0 = T - pd.Timedelta(hours=lead) + try: + q_fc = df_idx.loc[(t0, lead)] + curve[lead] = float(q_fc) - obs_val + except KeyError: + curve[lead] = np.nan + results[T] = curve + return results + + +def plot_spaghetti(ax, error_dict, color_da, label_prefix, linestyle="-", lw=0.9, alpha=0.35): + """Plot one thin line per target time + thick mean.""" + leads = list(range(1, 19)) + all_curves = [] + target_times = sorted(error_dict.keys()) + cmap = cm.get_cmap("plasma", len(target_times)) + + for i, T in enumerate(target_times): + curve = [error_dict[T].get(l, np.nan) for l in leads] + ax.plot(leads, curve, + color=cmap(i), lw=lw, alpha=alpha, + linestyle=linestyle, zorder=2) + all_curves.append(curve) + + if all_curves: + mean_curve = np.nanmean(np.array(all_curves), axis=0) + ax.plot(leads, mean_curve, + color=color_da, lw=2.8, alpha=0.95, + linestyle=linestyle, zorder=5, + label=f"{label_prefix} — mean across targets") + return all_curves + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--route-dir", default=DEFAULT_ROUTE_DIR) + parser.add_argument("--out-dir", default=None) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--da-name", default="routed_leadtime_da_full.parquet") + parser.add_argument("--ol-name", default="routed_leadtime_openloop_full.parquet") + parser.add_argument("--target-start", default=str(TARGET_START)) + parser.add_argument("--target-end", default=str(TARGET_END)) + args = parser.parse_args() + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + print("Loading parquets...") + da = load_parquet(os.path.join(args.route_dir, args.da_name)) + ol = load_parquet(os.path.join(args.route_dir, args.ol_name)) + obs = load_usgs(args.usgs_csv) + + target_times = pd.date_range(args.target_start, args.target_end, freq="1h") + print(f" Target verification times: {len(target_times)} " + f"({target_times[0]} → {target_times[-1]})") + + print("Building fixed-target error tables...") + da_errors = build_fixed_target_errors(da, obs, target_times) + ol_errors = build_fixed_target_errors(ol, obs, target_times) + print(f" Targets with obs: DA={len(da_errors)} OL={len(ol_errors)}") + + leads = list(range(1, 19)) + + # ---- Spaghetti: per-target-time curves, DA vs OL ---- + fig, (ax_da, ax_ol) = plt.subplots(2, 1, figsize=(13, 10), sharex=True, sharey=True) + + ax_da.axhline(0, color="black", lw=0.8, linestyle="--", alpha=0.4) + plot_spaghetti(ax_da, da_errors, "tab:purple", "DA", linestyle="-") + ax_da.set_ylabel("Error: forecast − USGS obs (m³/s)", fontsize=10) + ax_da.set_title("DA — error at each lead for fixed target times (Helene peak window)", fontsize=11) + ax_da.grid(True, alpha=0.2) + ax_da.legend(fontsize=9) + + ax_ol.axhline(0, color="black", lw=0.8, linestyle="--", alpha=0.4) + plot_spaghetti(ax_ol, ol_errors, "tab:gray", "Open-loop", linestyle="--") + ax_ol.set_ylabel("Error: forecast − USGS obs (m³/s)", fontsize=10) + ax_ol.set_xlabel("Forecast lead hour (hours before target)", fontsize=11) + ax_ol.set_title("Open-loop — error at each lead for fixed target times", fontsize=11) + ax_ol.set_xticks(leads) + ax_ol.grid(True, alpha=0.2) + ax_ol.legend(fontsize=9) + + fig.suptitle( + "Forecast error vs lead time — fixed verification time, Helene peak window\n" + f"Target times: {TARGET_START.strftime('%b %d %H UTC')} → " + f"{TARGET_END.strftime('%b %d %H UTC')} | USGS 03463300", + fontsize=12, y=1.01, + ) + plt.tight_layout() + out1 = os.path.join(out_dir, "error_fixed_target_helene.png") + plt.savefig(out1, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + # ---- Summary: mean across all target times, DA vs OL overlaid ---- + fig, ax = plt.subplots(figsize=(13, 6)) + ax.axhline(0, color="black", lw=0.8, linestyle="--", alpha=0.4) + + def mean_curve(error_dict): + arr = np.array([ + [error_dict[T].get(l, np.nan) for l in leads] + for T in sorted(error_dict.keys()) + ]) + return np.nanmean(arr, axis=0), np.nanstd(arr, axis=0) + + da_mean, da_std = mean_curve(da_errors) + ol_mean, ol_std = mean_curve(ol_errors) + + ax.fill_between(leads, da_mean - da_std, da_mean + da_std, + color="tab:purple", alpha=0.15, zorder=2) + ax.fill_between(leads, ol_mean - ol_std, ol_mean + ol_std, + color="tab:gray", alpha=0.15, zorder=2) + ax.plot(leads, da_mean, color="tab:purple", lw=2.6, marker="o", + zorder=5, label="DA — mean error (±1 std shaded)") + ax.plot(leads, ol_mean, color="tab:gray", lw=2.6, marker="s", + linestyle="--", zorder=5, label="Open-loop — mean error (±1 std shaded)") + + ax.set_xlabel("Forecast lead hour (hours before target verification time)", fontsize=11) + ax.set_ylabel("Mean error: forecast − USGS obs (m³/s)", fontsize=11) + ax.set_xticks(leads) + ax.set_title( + "Mean forecast error vs lead time — fixed verification time, Helene peak window\n" + f"Target: {TARGET_START.strftime('%b %d %H UTC')} → " + f"{TARGET_END.strftime('%b %d %H UTC')} | " + "Lead 1 = init 1 hr before target | Lead 18 = init 18 hr before target", + fontsize=11, + ) + ax.legend(fontsize=10, loc="lower right", frameon=True, framealpha=0.92) + ax.grid(True, alpha=0.25) + plt.tight_layout() + out2 = os.path.join(out_dir, "error_fixed_target_mean.png") + plt.savefig(out2, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4a_error_decay/plot_forecast_error_per_init.py b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4a_error_decay/plot_forecast_error_per_init.py new file mode 100644 index 00000000..691245d0 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4a_error_decay/plot_forecast_error_per_init.py @@ -0,0 +1,203 @@ +""" +plot_forecast_error_per_init.py + +Error decay by initialization time for the Helene window (Sep 24-28 2024). + +For each initialization time t0 in the Helene window: + error[lead] = ensemble_mean(q_gauge_m3s at t0+lead) - USGS_obs(t0+lead) + +Plotted as: + DA : thin colored lines (one per init time, colored by date) + thick mean across all + OL : thin gray dashed lines + thick gray dashed mean + +x-axis: forecast lead hour (1 -> 18) +y-axis: signed error (m³/s), positive = forecast too high + +The expected signal: DA error is small at lead 1 (just assimilated), grows +and converges toward the OL error curve by lead 18. + +Inputs: + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + +Obs: + /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +Outputs: + /forecast_error_per_init_helene.png (signed error) + /forecast_rmse_per_init_helene.png (absolute error / RMSE per init) +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.patches as mpatches + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast_hardcoded_r_routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +# Initialization times to show — Helene window +HELENE_INIT_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_INIT_END = pd.Timestamp("2024-09-28 23:00:00") + +# One color per init date +DATE_COLORS = { + "2024-09-24": "#1f77b4", + "2024-09-25": "#ff7f0e", + "2024-09-26": "#2ca02c", + "2024-09-27": "#d62728", + "2024-09-28": "#9467bd", +} + + +def load_parquet(path): + df = pd.read_parquet(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + df["ens_mean"] = df[member_cols].mean(axis=1) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + return df[["issue_time", "lead_hour", "valid_time", "ens_mean"]] + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() + or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def build_error_table(df, obs_series, init_start, init_end): + """Return DataFrame: issue_time, lead_hour, ens_mean, obs, error.""" + df = df[(df["issue_time"] >= init_start) & (df["issue_time"] <= init_end)].copy() + df["obs"] = df["valid_time"].map(obs_series) + df["error"] = df["ens_mean"] - df["obs"] + return df.dropna(subset=["obs", "error"]) + + +def plot_error(ax, err_df, color, alpha_thin, lw_thin, linestyle, label_prefix): + """Plot individual init-time error curves + thick mean curve.""" + leads = sorted(err_df["lead_hour"].unique()) + all_curves = [] + + for t0, grp in err_df.groupby("issue_time"): + date_str = str(pd.Timestamp(t0).date()) + c = DATE_COLORS.get(date_str, color) + grp_sorted = grp.sort_values("lead_hour") + # Align to leads grid — some may be missing + curve = grp_sorted.set_index("lead_hour")["error"].reindex(leads).values + ax.plot(leads, curve, + color=c, lw=lw_thin, alpha=alpha_thin, + linestyle=linestyle, zorder=2) + all_curves.append(curve) + + if all_curves: + mean_curve = np.nanmean(np.array(all_curves), axis=0) + ax.plot(leads, mean_curve, + color=color, lw=2.8, alpha=0.95, + linestyle=linestyle, zorder=5, + label=f"{label_prefix} — mean across all init times") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--route-dir", default=DEFAULT_ROUTE_DIR) + parser.add_argument("--out-dir", default=None) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--da-name", default="routed_leadtime_da_full.parquet") + parser.add_argument("--ol-name", default="routed_leadtime_openloop_full.parquet") + args = parser.parse_args() + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + print("Loading parquets...") + da = load_parquet(os.path.join(args.route_dir, args.da_name)) + ol = load_parquet(os.path.join(args.route_dir, args.ol_name)) + obs = load_usgs(args.usgs_csv) + print(f" DA issue_times: {da['issue_time'].nunique()} " + f"OL issue_times: {ol['issue_time'].nunique()}") + + da_err = build_error_table(da, obs, HELENE_INIT_START, HELENE_INIT_END) + ol_err = build_error_table(ol, obs, HELENE_INIT_START, HELENE_INIT_END) + print(f" DA init times in Helene window: {da_err['issue_time'].nunique()}") + + leads = sorted(da_err["lead_hour"].unique()) + + # ---- Signed error plot ---- + fig, ax = plt.subplots(figsize=(12, 6)) + ax.axhline(0, color="black", lw=0.8, linestyle="--", alpha=0.5, zorder=1) + + plot_error(ax, ol_err, color="tab:gray", alpha_thin=0.12, lw_thin=0.7, + linestyle="--", label_prefix="Open-loop") + plot_error(ax, da_err, color="tab:purple", alpha_thin=0.18, lw_thin=0.8, + linestyle="-", label_prefix="DA") + + # Date-color legend patches + patches = [mpatches.Patch(color=c, label=f"Init {d}") + for d, c in DATE_COLORS.items()] + handles, labels = ax.get_legend_handles_labels() + ax.legend(handles=handles + patches, + fontsize=8, loc="upper left", frameon=True, framealpha=0.9, ncol=2) + + ax.set_xlabel("Forecast lead hour", fontsize=11) + ax.set_ylabel("Error: forecast − USGS obs (m³/s)", fontsize=11) + ax.set_xticks(np.arange(1, 19)) + ax.set_title( + "Forecast error vs lead time — per initialization time, Helene window\n" + "DA (purple solid) vs Open-loop (gray dashed) | Sep 24–28 2024 | USGS 03463300", + fontsize=11, + ) + ax.grid(True, alpha=0.25, lw=0.4) + plt.tight_layout() + out1 = os.path.join(out_dir, "forecast_error_per_init_helene.png") + plt.savefig(out1, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + # ---- Absolute error (|error|) averaged per lead — cleaner summary ---- + fig, ax = plt.subplots(figsize=(12, 6)) + + def mean_abs_error_by_lead(err_df): + return err_df.groupby("lead_hour")["error"].apply( + lambda x: float(np.nanmean(np.abs(x))) + ) + + da_mae = mean_abs_error_by_lead(da_err) + ol_mae = mean_abs_error_by_lead(ol_err) + + ax.plot(da_mae.index, da_mae.values, + color="tab:purple", lw=2.4, marker="o", label="DA — mean |error|") + ax.plot(ol_mae.index, ol_mae.values, + color="tab:gray", lw=2.4, marker="s", linestyle="--", + label="Open-loop — mean |error|") + + ax.set_xlabel("Forecast lead hour", fontsize=11) + ax.set_ylabel("Mean |error| vs USGS obs (m³/s)", fontsize=11) + ax.set_xticks(np.arange(1, 19)) + ax.set_title( + "Mean absolute forecast error vs lead time — Helene window\n" + "DA (purple) vs Open-loop (gray) | Sep 24–28 2024 | USGS 03463300", + fontsize=11, + ) + ax.legend(fontsize=10, loc="upper left", frameon=True, framealpha=0.9) + ax.grid(True, alpha=0.25, lw=0.4) + plt.tight_layout() + out2 = os.path.join(out_dir, "forecast_mae_per_lead_helene.png") + plt.savefig(out2, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4a_error_decay/plot_lead_time_decay.py b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4a_error_decay/plot_lead_time_decay.py new file mode 100644 index 00000000..2b33fa82 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4a_error_decay/plot_lead_time_decay.py @@ -0,0 +1,160 @@ +""" +Catchment-level error-vs-lead-time decay curve + ensemble spread by lead. + +Two-panel figure for ONE catchment: + TOP — RMSE of ensemble-mean forecast at each lead hour (1..18) vs the + catchment's kriging observation. DA solid, open-loop dashed, with + a shaded band showing the min/max of per-member RMSE. + BOTTOM — Mean ensemble spread (std-dev across 20 members, averaged across + all issue times) at each lead hour. Tells you whether forcing + perturbation alone keeps the forecast ensemble diverse during the + 18-hour free-run — useful for inspecting individual members + without needing the full spaghetti view. + +This is the catchment-level analog of the gauge-level decay curve. +Routing to the gauge is a separate post-step (route_lead_time_forecasts.py ++ a gauge-level decay script); this script lets us look at the catchment-level +signal without T-route in the loop. + +Inputs (from run_lead_time_forecast_sweep.py): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv +Obs (from production EnKF run): + //_test_results.csv + +Output: + //_lead_time_decay.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" + +DEFAULT_LEADTIME_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast" +DEFAULT_DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" + +# Module-level globals set by main() so the helpers see them +CAT = "cat-1016300" +LEADTIME_DIR = DEFAULT_LEADTIME_DIR +DA_DIR = DEFAULT_DA_DIR +OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay.png") + + +def load_forecasts(path): + """Return (df, member_cols).""" + if not os.path.exists(path): + raise FileNotFoundError(path) + df = pd.read_csv(path, parse_dates=['issue_time', 'valid_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + return df, member_cols + + +def load_obs(): + obs_path = os.path.join(DA_DIR, CAT, f"{CAT}_test_results.csv") + df = pd.read_csv(obs_path, parse_dates=['date']) + return df.set_index('date')['obs_mm_h'] + + +def metrics_by_lead(df, member_cols, obs_series): + """For each lead hour, return (rmse_mean, rmse_min_member, rmse_max_member, + mean_ensemble_std).""" + df = df.copy() + df['obs'] = df['valid_time'].map(obs_series) + df = df.dropna(subset=['obs']) + df['ens_mean'] = df[member_cols].mean(axis=1) + df['ens_std'] = df[member_cols].std(axis=1) + + leads = sorted(df['lead_hour'].unique()) + rmse_mean = np.full(len(leads), np.nan) + rmse_min_member = np.full(len(leads), np.nan) + rmse_max_member = np.full(len(leads), np.nan) + mean_ens_std = np.full(len(leads), np.nan) + for i, L in enumerate(leads): + sub = df[df['lead_hour'] == L] + if len(sub) == 0: + continue + err = sub['ens_mean'].values - sub['obs'].values + rmse_mean[i] = float(np.sqrt(np.mean(err ** 2))) + per_member_rmse = np.array([ + np.sqrt(np.mean((sub[c].values - sub['obs'].values) ** 2)) + for c in member_cols + ]) + rmse_min_member[i] = float(np.nanmin(per_member_rmse)) + rmse_max_member[i] = float(np.nanmax(per_member_rmse)) + mean_ens_std[i] = float(np.nanmean(sub['ens_std'].values)) + return (np.asarray(leads), rmse_mean, rmse_min_member, + rmse_max_member, mean_ens_std) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', default="cat-1016300") + parser.add_argument('--leadtime-dir', default=DEFAULT_LEADTIME_DIR, + help='Dir holding /_lead_time_forecasts_{da,openloop}.csv') + parser.add_argument('--da-dir', default=DEFAULT_DA_DIR, + help='Dir holding /_test_results.csv (for obs)') + args = parser.parse_args() + + global CAT, LEADTIME_DIR, DA_DIR, OUT_PNG + CAT = args.cat_id + LEADTIME_DIR = args.leadtime_dir + DA_DIR = args.da_dir + OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay.png") + + da_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_da.csv") + ol_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_openloop.csv") + df_da, m_cols = load_forecasts(da_path) + df_ol, _ = load_forecasts(ol_path) + obs_series = load_obs() + + leads_da, rmse_da, rmse_da_lo, rmse_da_hi, std_da = metrics_by_lead( + df_da, m_cols, obs_series) + leads_ol, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = metrics_by_lead( + df_ol, m_cols, obs_series) + + fig, (ax_top, ax_bot) = plt.subplots(2, 1, figsize=(10, 10), sharex=True, + gridspec_kw={"height_ratios": [1.4, 1.0]}) + + # ----- TOP: RMSE decay curve ----- + ax_top.fill_between(leads_da, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2, + label="DA — member spread (min/max RMSE)") + ax_top.fill_between(leads_ol, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2, + label="Open-loop — member spread") + ax_top.plot(leads_da, rmse_da, color=DA_COLOR, lw=2.4, marker='o', + zorder=4, label="DA on (ensemble-mean RMSE)") + ax_top.plot(leads_ol, rmse_ol, color=OL_COLOR, lw=2.4, marker='s', + linestyle='--', zorder=4, label="Open-loop (ensemble-mean RMSE)") + ax_top.set_ylabel("RMSE vs Qkrig obs (mm/h)", fontsize=11) + ax_top.set_title( + f"Forecast lead-time error decay — {CAT}\n" + f"Issue times pooled across test period (Oct 2023 – Oct 2024)", + fontsize=12, + ) + ax_top.grid(True, alpha=0.25) + ax_top.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + # ----- BOTTOM: ensemble spread by lead hour ----- + ax_bot.plot(leads_da, std_da, color=DA_COLOR, lw=2.2, marker='o', + label="DA — mean ensemble std") + ax_bot.plot(leads_ol, std_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', label="Open-loop — mean ensemble std") + ax_bot.set_xlabel("Forecast lead time (hours)", fontsize=11) + ax_bot.set_ylabel("Mean ensemble std (mm/h)", fontsize=11) + ax_bot.set_xticks(np.arange(1, 19)) + ax_bot.grid(True, alpha=0.25) + ax_bot.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + plt.tight_layout() + plt.savefig(OUT_PNG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4a_error_decay/plot_lead_time_decay_by_regime.py b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4a_error_decay/plot_lead_time_decay_by_regime.py new file mode 100644 index 00000000..ecf10408 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4a_error_decay/plot_lead_time_decay_by_regime.py @@ -0,0 +1,200 @@ +""" +Lead-time decay curve, split by flow regime at issue time. + +The pooled lead-time curve (plot_lead_time_decay.py) showed DA losing to +open-loop across most lead hours, but the test period is ~99% low-flow. +This script splits the same forecast CSVs by the flow regime at the issue +time t0, so we can see whether DA helps when it matters (storms / Helene) +and hurts when it doesn't (low flow). + +Three regimes are partitioned on obs(t0) (kriging obs at the issue time): + + helene issue_time ∈ [2024-09-24, 2024-09-28] (the 5-day Helene window) + storm obs(t0) > 0.5 mm/h (high-flow issue times) + low_flow obs(t0) < 0.1 mm/h (low-flow issue times) + +Same 2-panel layout per regime: RMSE vs lead (top), ensemble spread vs lead +(bottom). Three regime columns × 2 metric rows in one figure. + +Inputs (from run_lead_time_forecast_sweep.py — no re-run needed): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv +Obs (from production EnKF): + //_test_results.csv + +Output: + //_lead_time_decay_by_regime.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" + +DEFAULT_LEADTIME_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast" +DEFAULT_DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" + +# Module-level globals set by main() so helpers see them +CAT = "cat-1016300" +LEADTIME_DIR = DEFAULT_LEADTIME_DIR +DA_DIR = DEFAULT_DA_DIR +OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay_by_regime.png") + +# Regime definitions +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-28 23:00:00") +STORM_OBS_THRESHOLD = 0.5 # mm/h +LOWFLOW_OBS_THRESHOLD = 0.1 # mm/h + + +def load_forecasts(path): + if not os.path.exists(path): + raise FileNotFoundError(path) + df = pd.read_csv(path, parse_dates=['issue_time', 'valid_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + return df, member_cols + + +def load_obs(): + obs_path = os.path.join(DA_DIR, CAT, f"{CAT}_test_results.csv") + df = pd.read_csv(obs_path, parse_dates=['date']) + return df.set_index('date')['obs_mm_h'] + + +def regime_mask(issue_times, obs_at_issue, regime): + """Boolean mask over issue_times for the named regime.""" + if regime == "helene": + return (issue_times >= HELENE_START) & (issue_times <= HELENE_END) + if regime == "storm": + return obs_at_issue > STORM_OBS_THRESHOLD + if regime == "low_flow": + return obs_at_issue < LOWFLOW_OBS_THRESHOLD + raise ValueError(regime) + + +def metrics_by_lead(df, member_cols, obs_series, issue_mask): + """Compute (leads, rmse_mean, rmse_min, rmse_max, mean_std) restricted + to issue times where issue_mask is True.""" + df = df.copy() + df = df[df['issue_time'].isin(issue_mask)] + if len(df) == 0: + return None + df['obs'] = df['valid_time'].map(obs_series) + df = df.dropna(subset=['obs']) + if len(df) == 0: + return None + df['ens_mean'] = df[member_cols].mean(axis=1) + df['ens_std'] = df[member_cols].std(axis=1) + + leads = sorted(df['lead_hour'].unique()) + rmse_mean = np.full(len(leads), np.nan) + rmse_lo = np.full(len(leads), np.nan) + rmse_hi = np.full(len(leads), np.nan) + std_mean = np.full(len(leads), np.nan) + for i, L in enumerate(leads): + sub = df[df['lead_hour'] == L] + if len(sub) == 0: + continue + err = sub['ens_mean'].values - sub['obs'].values + rmse_mean[i] = float(np.sqrt(np.mean(err ** 2))) + per_member = np.array([ + np.sqrt(np.mean((sub[c].values - sub['obs'].values) ** 2)) + for c in member_cols + ]) + rmse_lo[i] = float(np.nanmin(per_member)) + rmse_hi[i] = float(np.nanmax(per_member)) + std_mean[i] = float(np.nanmean(sub['ens_std'].values)) + return np.asarray(leads), rmse_mean, rmse_lo, rmse_hi, std_mean + + +def plot_pair(ax_rmse, ax_std, da_metrics, ol_metrics, + regime_label, n_issue_times): + if da_metrics is None or ol_metrics is None: + ax_rmse.set_title(f"{regime_label}\n(no issue times match)", fontsize=11) + ax_rmse.axis('off'); ax_std.axis('off') + return + leads, rmse_da, rmse_da_lo, rmse_da_hi, std_da = da_metrics + _, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = ol_metrics + + # RMSE panel + ax_rmse.fill_between(leads, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2) + ax_rmse.fill_between(leads, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2) + ax_rmse.plot(leads, rmse_da, color=DA_COLOR, lw=2.2, marker='o', + zorder=4, label="DA on") + ax_rmse.plot(leads, rmse_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', zorder=4, label="Open-loop") + ax_rmse.set_title(f"{regime_label}\n({n_issue_times} issue times)", fontsize=11) + ax_rmse.set_ylabel("RMSE (mm/h)", fontsize=10) + ax_rmse.grid(True, alpha=0.25) + ax_rmse.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + # Spread panel + ax_std.plot(leads, std_da, color=DA_COLOR, lw=2.0, marker='o', label="DA") + ax_std.plot(leads, std_ol, color=OL_COLOR, lw=2.0, marker='s', + linestyle='--', label="Open-loop") + ax_std.set_xlabel("Forecast lead time (hours)", fontsize=10) + ax_std.set_ylabel("Mean ensemble std (mm/h)", fontsize=10) + ax_std.set_xticks(np.arange(1, 19)) + ax_std.grid(True, alpha=0.25) + ax_std.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', default="cat-1016300") + parser.add_argument('--leadtime-dir', default=DEFAULT_LEADTIME_DIR, + help='Dir holding /_lead_time_forecasts_{da,openloop}.csv') + parser.add_argument('--da-dir', default=DEFAULT_DA_DIR, + help='Dir holding /_test_results.csv (for obs)') + args = parser.parse_args() + + global CAT, LEADTIME_DIR, DA_DIR, OUT_PNG + CAT = args.cat_id + LEADTIME_DIR = args.leadtime_dir + DA_DIR = args.da_dir + OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay_by_regime.png") + + da_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_da.csv") + ol_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_openloop.csv") + df_da, m_cols = load_forecasts(da_path) + df_ol, _ = load_forecasts(ol_path) + obs_series = load_obs() + + # Unique issue times across the run, with obs at each + issue_times = pd.to_datetime(sorted(df_da['issue_time'].unique())) + obs_at_issue = pd.Series(issue_times, index=issue_times).map(obs_series) + + regimes = [ + ("helene", "Helene window (Sep 24–28 2024)"), + ("storm", f"Storm hours (obs(t0) > {STORM_OBS_THRESHOLD} mm/h)"), + ("low_flow", f"Low flow (obs(t0) < {LOWFLOW_OBS_THRESHOLD} mm/h)"), + ] + + fig, axes = plt.subplots(2, 3, figsize=(18, 10), + gridspec_kw={"height_ratios": [1.4, 1.0]}) + for col, (key, label) in enumerate(regimes): + mask = regime_mask(issue_times, obs_at_issue, key) + kept = issue_times[mask] + kept_set = set(pd.to_datetime(kept)) + da_m = metrics_by_lead(df_da, m_cols, obs_series, kept_set) + ol_m = metrics_by_lead(df_ol, m_cols, obs_series, kept_set) + plot_pair(axes[0, col], axes[1, col], da_m, ol_m, label, len(kept)) + + fig.suptitle( + f"Lead-time error decay by flow regime — {CAT}\n" + f"Same forecast CSVs as the pooled view, partitioned on obs(t0).", + fontsize=13, y=0.995, + ) + plt.tight_layout(rect=(0, 0, 1, 0.97)) + plt.savefig(OUT_PNG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4a_error_decay/plot_lead_time_decay_gauge.py b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4a_error_decay/plot_lead_time_decay_gauge.py new file mode 100644 index 00000000..71ce54fc --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4a_error_decay/plot_lead_time_decay_gauge.py @@ -0,0 +1,311 @@ +""" +Gauge-level lead-time forecast decay curve. + +Routes 21-catchment forecast CSVs through T-route to USGS gauge 03463300 and +plots RMSE vs lead time (DA solid, open-loop dashed) against USGS observed Q. +Same 2-panel layout as the catchment-level plot, plus a regime split for the +Helene window. + +The routing step is done separately by route_lead_time_forecasts.py; +this script just reads the resulting parquets and produces the plot. + +Inputs: + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + Each: one row per (issue_time, lead_hour); member columns hold q_gauge_m3s. + (Long-format with explicit `member` and `q_gauge_m3s` columns also supported.) +Obs: + /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + USGS hourly Q in m³/s at gauge 03463300 (South Toe River near Celo, NC). + +Output: + /lead_time_decay_gauge_pooled.png + /lead_time_decay_gauge_by_regime.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" +OBS_COLOR = "black" + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/leadtime_troute_routing" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +# Watershed area at gauge 03463300 (sum of 21 catchments per GPKG hydrofabric) +WATERSHED_AREA_KM2 = 113.18 +# 1 mm/h depth × 113.18 km² = 113.18e3 m³/h = 31.439 m³/s +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-28 23:00:00") +STORM_OBS_THRESHOLD_M3S = 20.0 # gauge-level storm threshold (was 50 — too high + # for this 113 km² basin; only 2 issue times qualified) +LOWFLOW_OBS_THRESHOLD_M3S = 5.0 # gauge-level low-flow threshold + +USGS_HELENE_PEAK_M3S = 1886.0 # reference Sep 27 14:00 + + +def load_parquet_long(path): + """Load a routed parquet and normalize to long format: + columns = issue_time, lead_hour, member, q_gauge_m3s. + + Handles both wide format (member_00..member_19 columns) and long format + (explicit `member` + `q_gauge_m3s` columns). Auto-detects. + """ + if not os.path.exists(path): + raise FileNotFoundError(path) + df = pd.read_parquet(path) + + # Long format detection + if 'q_gauge_m3s' in df.columns and 'member' in df.columns: + df['issue_time'] = pd.to_datetime(df['issue_time']) + return df[['issue_time', 'lead_hour', 'member', 'q_gauge_m3s']] + + # Some variants might call it 'q_m3s' or similar + q_col_candidates = [c for c in df.columns if c.lower() in + ('q_gauge_m3s', 'q_m3s', 'q_gauge', 'q')] + if 'member' in df.columns and q_col_candidates: + qc = q_col_candidates[0] + df['issue_time'] = pd.to_datetime(df['issue_time']) + out = df[['issue_time', 'lead_hour', 'member', qc]].copy() + out = out.rename(columns={qc: 'q_gauge_m3s'}) + return out + + # Wide format: member_00..member_19 columns + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + if not member_cols: + raise ValueError(f"Couldn't identify member columns in {path}. " + f"Columns present: {df.columns.tolist()}") + df['issue_time'] = pd.to_datetime(df['issue_time']) + keep = ['issue_time', 'lead_hour'] + if 'valid_time' in df.columns: + keep.append('valid_time') + long = df[keep + member_cols].melt( + id_vars=keep, + value_vars=member_cols, + var_name='member', + value_name='q_gauge_m3s', + ) + return long + + +def load_usgs_obs(usgs_csv): + """Load USGS gauge obs and ensure m³/s. + + The CSV at gauge 03463300 has the column `QObs(mm/h)` (catchment-averaged + depth, not gauge discharge in m³/s). If the column name contains 'mm', we + convert mm/h → m³/s by multiplying by the watershed area factor: + m³/s = mm/h × WATERSHED_AREA_KM2 × 1000 / 3600 + Verified: 59.978 mm/h × 113.18 km² = 1886 m³/s (USGS Helene peak). + """ + df = pd.read_csv(usgs_csv) + date_candidates = [c for c in df.columns if c.lower() in + ('datetime', 'date', 'time', 'timestamp')] + q_candidates = [c for c in df.columns if 'q' in c.lower() + or 'flow' in c.lower() or 'discharge' in c.lower()] + if not date_candidates or not q_candidates: + raise ValueError(f"Could not identify date/Q columns in {usgs_csv}. " + f"Columns: {df.columns.tolist()}") + dc, qc = date_candidates[0], q_candidates[0] + df[dc] = pd.to_datetime(df[dc]) + series = df.set_index(dc)[qc].astype(float) + if 'mm' in qc.lower(): + series = series * MM_H_TO_M3_S + print(f" Converted obs from mm/h (column '{qc}') to m³/s " + f"using area = {WATERSHED_AREA_KM2} km² (× {MM_H_TO_M3_S:.4f})") + else: + print(f" Loaded obs as-is from column '{qc}' (assumed m³/s)") + return series + + +def metrics_by_lead(df_long, obs_series, issue_mask=None): + """Compute per-lead-hour: ensemble-mean RMSE, per-member min/max RMSE, + mean ensemble std. df_long must have issue_time, lead_hour, member, q_gauge_m3s. + """ + df = df_long.copy() + if issue_mask is not None: + df = df[df['issue_time'].isin(issue_mask)] + if len(df) == 0: + return None + df['valid_time'] = df['issue_time'] + pd.to_timedelta(df['lead_hour'], unit='h') + df['obs'] = df['valid_time'].map(obs_series) + df = df.dropna(subset=['obs']) + if len(df) == 0: + return None + + # Ensemble mean per (issue_time, lead_hour) + grouped = df.groupby(['issue_time', 'lead_hour']) + ens_mean = grouped['q_gauge_m3s'].mean().reset_index(name='ens_mean') + ens_std = grouped['q_gauge_m3s'].std().reset_index(name='ens_std') + obs_per = grouped['obs'].first().reset_index(name='obs') + panel = ens_mean.merge(ens_std, on=['issue_time', 'lead_hour']).merge( + obs_per, on=['issue_time', 'lead_hour']) + + leads = sorted(panel['lead_hour'].unique()) + rmse_mean = np.full(len(leads), np.nan) + rmse_lo = np.full(len(leads), np.nan) + rmse_hi = np.full(len(leads), np.nan) + std_mean = np.full(len(leads), np.nan) + for i, L in enumerate(leads): + psub = panel[panel['lead_hour'] == L] + if len(psub) == 0: + continue + err = psub['ens_mean'].values - psub['obs'].values + rmse_mean[i] = float(np.sqrt(np.mean(err ** 2))) + # Per-member RMSE — recompute from the underlying long frame + dsub = df[df['lead_hour'] == L] + member_rmses = [] + for m, g in dsub.groupby('member'): + err_m = g['q_gauge_m3s'].values - g['obs'].values + if len(err_m) >= 1: + member_rmses.append(float(np.sqrt(np.mean(err_m ** 2)))) + if member_rmses: + rmse_lo[i] = float(np.nanmin(member_rmses)) + rmse_hi[i] = float(np.nanmax(member_rmses)) + std_mean[i] = float(np.nanmean(psub['ens_std'].values)) + return np.asarray(leads), rmse_mean, rmse_lo, rmse_hi, std_mean + + +def plot_two_panel(out_path, leads_da, da_metrics, leads_ol, ol_metrics, title): + fig, (ax_top, ax_bot) = plt.subplots(2, 1, figsize=(10, 10), sharex=True, + gridspec_kw={"height_ratios": [1.4, 1.0]}) + leads, rmse_da, rmse_da_lo, rmse_da_hi, std_da = da_metrics + _, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = ol_metrics + + ax_top.fill_between(leads, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2, + label="DA — member spread (min/max RMSE)") + ax_top.fill_between(leads, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2, + label="Open-loop — member spread") + ax_top.plot(leads, rmse_da, color=DA_COLOR, lw=2.4, marker='o', + zorder=4, label="DA on (ensemble-mean RMSE)") + ax_top.plot(leads, rmse_ol, color=OL_COLOR, lw=2.4, marker='s', + linestyle='--', zorder=4, label="Open-loop (ensemble-mean RMSE)") + ax_top.set_ylabel("RMSE vs USGS obs (m³/s)", fontsize=11) + ax_top.set_title(title, fontsize=12) + ax_top.grid(True, alpha=0.25) + ax_top.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + ax_bot.plot(leads, std_da, color=DA_COLOR, lw=2.2, marker='o', + label="DA — mean ensemble std") + ax_bot.plot(leads, std_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', label="Open-loop — mean ensemble std") + ax_bot.set_xlabel("Forecast lead time (hours)", fontsize=11) + ax_bot.set_ylabel("Mean ensemble std (m³/s)", fontsize=11) + ax_bot.set_xticks(np.arange(1, 19)) + ax_bot.grid(True, alpha=0.25) + ax_bot.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + plt.tight_layout() + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + +def plot_three_panel_regime(out_path, regime_metrics): + fig, axes = plt.subplots(2, 3, figsize=(18, 10), + gridspec_kw={"height_ratios": [1.4, 1.0]}) + for col, (label, n_issue, da_m, ol_m) in enumerate(regime_metrics): + ax_rmse = axes[0, col] + ax_std = axes[1, col] + if da_m is None or ol_m is None: + ax_rmse.set_title(f"{label}\n(no issue times match)", fontsize=11) + ax_rmse.axis('off'); ax_std.axis('off') + continue + leads, rmse_da, rmse_da_lo, rmse_da_hi, std_da = da_m + _, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = ol_m + ax_rmse.fill_between(leads, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2) + ax_rmse.fill_between(leads, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2) + ax_rmse.plot(leads, rmse_da, color=DA_COLOR, lw=2.2, marker='o', + zorder=4, label="DA on") + ax_rmse.plot(leads, rmse_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', zorder=4, label="Open-loop") + ax_rmse.set_title(f"{label}\n({n_issue} issue times)", fontsize=11) + ax_rmse.set_ylabel("RMSE (m³/s)", fontsize=10) + ax_rmse.grid(True, alpha=0.25) + ax_rmse.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + ax_std.plot(leads, std_da, color=DA_COLOR, lw=2.0, marker='o', label="DA") + ax_std.plot(leads, std_ol, color=OL_COLOR, lw=2.0, marker='s', + linestyle='--', label="Open-loop") + ax_std.set_xlabel("Forecast lead time (hours)", fontsize=10) + ax_std.set_ylabel("Mean ensemble std (m³/s)", fontsize=10) + ax_std.set_xticks(np.arange(1, 19)) + ax_std.grid(True, alpha=0.25) + ax_std.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + fig.suptitle("Gauge-level lead-time error decay by regime — USGS 03463300", + fontsize=13, y=0.995) + plt.tight_layout(rect=(0, 0, 1, 0.97)) + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--route-dir', default=DEFAULT_ROUTE_DIR, + help='Dir holding routed_leadtime_{da,openloop}_full.parquet') + parser.add_argument('--out-dir', default=None, + help='Where to write the output PNGs. Defaults to --route-dir ' + '(which may not be writable if owned by another user — ' + 'pass an explicit path then).') + parser.add_argument('--usgs-csv', default=DEFAULT_USGS_CSV) + parser.add_argument('--da-name', default='routed_leadtime_da_full.parquet') + parser.add_argument('--ol-name', default='routed_leadtime_openloop_full.parquet') + args = parser.parse_args() + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + da_path = os.path.join(args.route_dir, args.da_name) + ol_path = os.path.join(args.route_dir, args.ol_name) + df_da = load_parquet_long(da_path) + df_ol = load_parquet_long(ol_path) + obs_series = load_usgs_obs(args.usgs_csv) + + print(f"DA rows: {len(df_da):,} | unique issue_times: {df_da['issue_time'].nunique()}") + print(f"OL rows: {len(df_ol):,} | unique issue_times: {df_ol['issue_time'].nunique()}") + print(f"USGS obs range: {obs_series.index.min()} .. {obs_series.index.max()} " + f"({len(obs_series):,} hours)") + print(f"USGS Helene peak (reference): {USGS_HELENE_PEAK_M3S:.0f} m³/s") + + # ----- Pooled (all issue times) ----- + da_pooled = metrics_by_lead(df_da, obs_series) + ol_pooled = metrics_by_lead(df_ol, obs_series) + out_pooled = os.path.join(out_dir, "lead_time_decay_gauge_pooled.png") + plot_two_panel(out_pooled, da_pooled[0], da_pooled, ol_pooled[0], ol_pooled, + title="Gauge-level lead-time decay — USGS 03463300\n" + "All issue times pooled (Oct 2023 – Oct 2024)") + + # ----- Regime split: Helene, storm, low-flow at gauge ----- + issue_times = pd.to_datetime(sorted(df_da['issue_time'].unique())) + obs_at_issue = pd.Series(issue_times, index=issue_times).map(obs_series) + + helene_mask = (issue_times >= HELENE_START) & (issue_times <= HELENE_END) + storm_mask = obs_at_issue > STORM_OBS_THRESHOLD_M3S + lowflow_mask = obs_at_issue < LOWFLOW_OBS_THRESHOLD_M3S + + regime_metrics = [] + for label, mask in [ + (f"Helene window (Sep 24–28 2024)", helene_mask), + (f"Storm hours (obs(t0) > {STORM_OBS_THRESHOLD_M3S:.0f} m³/s)", storm_mask), + (f"Low flow (obs(t0) < {LOWFLOW_OBS_THRESHOLD_M3S:.0f} m³/s)", lowflow_mask), + ]: + kept = set(pd.to_datetime(issue_times[mask])) + da_m = metrics_by_lead(df_da, obs_series, kept) + ol_m = metrics_by_lead(df_ol, obs_series, kept) + regime_metrics.append((label, len(kept), da_m, ol_m)) + + out_regime = os.path.join(out_dir, "lead_time_decay_gauge_by_regime.png") + plot_three_panel_regime(out_regime, regime_metrics) + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4a_error_decay/run_4a_f2.sh b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4a_error_decay/run_4a_f2.sh new file mode 100644 index 00000000..80ca8dd3 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4a_error_decay/run_4a_f2.sh @@ -0,0 +1,80 @@ +#!/usr/bin/env bash +# F2 fixed R=0.07 (20pct gauge holdout) — 4a lead-time decay plots. +# +# Runs plot scripts per catchment: +# plot_lead_time_decay.py — pooled RMSE vs lead (catchment-level) +# plot_lead_time_decay_by_regime.py — same split by flow regime +# +# Gauge-level decay plots run once if routed parquets exist: +# plot_lead_time_decay_gauge.py — gauge-level (requires routed parquets) +# plot_forecast_error_fixed_target.py +# plot_forecast_error_per_init.py +# +# Usage: +# bash run_4a_f2.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +LEADTIME_DIR=/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast +DA_DIR=/mnt/disk2/suma_helen_poster/da_results/v2_fixed_r007_analysis +OBS_DIR=/mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance +ROUTE_DIR=/mnt/disk2/suma_helen_poster/leadtime_troute_routing_f2 +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +echo "[F2-4a] Lead-time decay plots — leadtime dir: $LEADTIME_DIR" + +for CAT in "${CATS[@]}"; do + DA_CSV="$LEADTIME_DIR/$CAT/${CAT}_lead_time_forecasts_da.csv" + if [ ! -f "$DA_CSV" ]; then + echo " [$CAT] No lead-time CSV — skipping" + continue + fi + + echo " [$CAT] Plotting pooled decay..." + $TROUTE "$SCRIPT_DIR/plot_lead_time_decay.py" \ + --cat-id "$CAT" \ + --leadtime-dir "$LEADTIME_DIR" \ + --da-dir "$DA_DIR" \ + --obs-dir "$OBS_DIR" + + echo " [$CAT] Plotting decay by regime..." + $TROUTE "$SCRIPT_DIR/plot_lead_time_decay_by_regime.py" \ + --cat-id "$CAT" \ + --leadtime-dir "$LEADTIME_DIR" \ + --da-dir "$DA_DIR" \ + --obs-dir "$OBS_DIR" +done + +# Gauge-level decay (needs routed parquets — run after route_leadtime_f2.sh) +if [ -f "$ROUTE_DIR/routed_leadtime_da_full.parquet" ]; then + echo "[F2-4a] Gauge-level decay (routed)..." + $TROUTE "$SCRIPT_DIR/plot_lead_time_decay_gauge.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" + + echo "[F2-4a] Fixed-target error plots..." + $TROUTE "$SCRIPT_DIR/plot_forecast_error_fixed_target.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" + + echo "[F2-4a] Per-init error plots..." + $TROUTE "$SCRIPT_DIR/plot_forecast_error_per_init.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" +else + echo "[F2-4a] Routed parquets not found — skipping gauge-level and error plots." + echo " Run route_leadtime_f2.sh first, then re-run this script." +fi + +echo "[F2-4a] Done." diff --git a/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/plot_crossed_ensemble.py b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/plot_crossed_ensemble.py new file mode 100644 index 00000000..7cb00a93 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/plot_crossed_ensemble.py @@ -0,0 +1,229 @@ +""" +plot_crossed_ensemble.py -- 4b: 600-member crossed ensemble vs USGS obs. + +Loads the crossed ensemble parquet produced by run_crossed_ensemble.py: + //_crossed_ensemble.parquet + Columns: issue_time, lead_hour, member_0000..member_0599 + +Projects all members to valid_time = issue_time + lead_hour hours. +Aggregates at each valid_time across ALL members from ALL issue times +to build a probabilistic envelope (5th/25th/50th/75th/95th percentiles). + +Two output figures: + 4b_crossed_ensemble_helene.png + -- Full Helene window (Sep 24-30) probabilistic envelope + USGS obs. + Percentile shading in two layers (5-95 outer, 25-75 inner). + 4b_crossed_ensemble_peak.png + -- Zoomed to Helene peak (Sep 26 12 UTC -> Sep 28 06 UTC), + same shading + USGS obs + individual spaghetti from Sep 27 inits. + +Usage: + python3 plot_crossed_ensemble.py \\ + --ensemble-dir /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble \\ + --cat-id cat-1016300 \\ + --usgs-csv /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv \\ + --out-dir /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble/cat-1016300 +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import matplotlib.patches as mpatches + +DEFAULT_ENSEMBLE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" +DEFAULT_CAT_ID = "cat-1016300" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +PLOT_START = pd.Timestamp("2024-09-24 00:00:00") +PLOT_END = pd.Timestamp("2024-09-28 18:00:00") # cut trailing spike +PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +PEAK_END = pd.Timestamp("2024-09-28 06:00:00") +HELENE_PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +HELENE_PEAK_END = pd.Timestamp("2024-09-28 00:00:00") + + +def load_ensemble(path): + df = pd.read_parquet(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + return df, member_cols + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def build_percentile_envelope(df, member_cols, t_start, t_end): + """ + For each valid_time in [t_start, t_end], stack all member values + from all issue times and compute percentiles. + Returns DataFrame indexed by valid_time with columns p05, p25, p50, p75, p95. + """ + mask = (df["valid_time"] >= t_start) & (df["valid_time"] <= t_end) + sub = df[mask].copy() + sub["valid_time"] = sub["valid_time"].dt.floor("1h") + + records = [] + for vt, grp in sub.groupby("valid_time"): + vals = grp[member_cols].to_numpy(dtype=float).flatten() + vals = vals[~np.isnan(vals)] * MM_H_TO_M3_S + if len(vals) == 0: + continue + records.append({ + "valid_time": vt, + "p05": np.percentile(vals, 5), + "p25": np.percentile(vals, 25), + "p50": np.percentile(vals, 50), + "p75": np.percentile(vals, 75), + "p95": np.percentile(vals, 95), + "n": len(vals), + }) + return pd.DataFrame(records).set_index("valid_time").sort_index() + + +def _add_helene_band(ax, label=True): + ax.axvspan(HELENE_PEAK_START, HELENE_PEAK_END, + color="salmon", alpha=0.12, zorder=0) + if label: + ax.text(HELENE_PEAK_START + pd.Timedelta(hours=6), 0.97, "Helene peak", + transform=ax.get_xaxis_transform(), + fontsize=9, color="firebrick", ha="left", va="top") + + +def _format_xaxis(ax, major_interval=1, minor_hours=6): + ax.xaxis.set_major_locator(mdates.DayLocator(interval=major_interval)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.xaxis.set_minor_locator(mdates.HourLocator(interval=minor_hours)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.2, lw=0.4) + + +def plot_envelope(ax, env, obs_series, t_start, t_end, + ens_color="tab:blue", obs_color="black", helene_label=True): + _add_helene_band(ax, label=helene_label) + + # Outer band: 5-95 + ax.fill_between(env.index, env["p05"], env["p95"], + color=ens_color, alpha=0.15, zorder=2, + label="5th-95th percentile") + # Inner band: 25-75 + ax.fill_between(env.index, env["p25"], env["p75"], + color=ens_color, alpha=0.30, zorder=3, + label="25th-75th percentile") + # Median + ax.plot(env.index, env["p50"], + color=ens_color, lw=2.2, zorder=5, + label="Ensemble median (600 members)") + + # USGS obs + obs_w = obs_series.loc[t_start:t_end] + ax.plot(obs_w.index, obs_w.values, + color=obs_color, lw=1.8, zorder=6, label="USGS obs") + + ax.set_xlim(t_start, t_end) + ax.set_ylabel("Discharge (m3/s)", fontsize=11) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--ensemble-dir", default=DEFAULT_ENSEMBLE_DIR) + parser.add_argument("--cat-id", default=DEFAULT_CAT_ID) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--out-dir", default=None) + args = parser.parse_args() + + cat_dir = os.path.join(args.ensemble_dir, args.cat_id) + out_dir = args.out_dir or cat_dir + os.makedirs(out_dir, exist_ok=True) + + parquet_path = os.path.join(cat_dir, f"{args.cat_id}_crossed_ensemble.parquet") + print(f"Loading crossed ensemble: {parquet_path}") + df, member_cols = load_ensemble(parquet_path) + obs = load_usgs(args.usgs_csv) + print(f" {df['issue_time'].nunique()} issue times, " + f"{len(member_cols)} members per row") + + # ------------------------------------------------------------------ # + # Figure 4b-full: full Helene window probabilistic envelope # + # ------------------------------------------------------------------ # + print("Building percentile envelope (full window)...") + env_full = build_percentile_envelope(df, member_cols, PLOT_START, PLOT_END) + print(f" Valid times with data: {len(env_full)}, " + f"median members per timestep: {env_full['n'].median():.0f}") + + fig, ax = plt.subplots(figsize=(17, 7)) + plot_envelope(ax, env_full, obs, PLOT_START, PLOT_END) + ax.set_xlabel("Date", fontsize=11) + ax.set_title( + f"4b -- 600-member crossed ensemble (30 met x 20 hydro-state) | {args.cat_id}\n" + "Shaded: 5-95th (light) and 25-75th (dark) percentile across all members " + "and issue times | Sep 24-30 2024 | USGS 03463300", + fontsize=11, + ) + ax.legend(fontsize=10, loc="upper left", frameon=True, framealpha=0.9) + _format_xaxis(ax) + plt.tight_layout() + out1 = os.path.join(out_dir, f"{args.cat_id}_4b_crossed_ensemble_helene.png") + plt.savefig(out1, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + # ------------------------------------------------------------------ # + # Figure 4b-peak: zoomed to Helene peak window # + # ------------------------------------------------------------------ # + print("Building percentile envelope (peak window)...") + env_peak = build_percentile_envelope(df, member_cols, PEAK_START, PEAK_END) + + fig, (ax_main, ax_zoom) = plt.subplots( + 2, 1, figsize=(14, 11), + gridspec_kw={"height_ratios": [1, 1.4]}, sharex=False) + + # Top panel: full window context (smaller) + plot_envelope(ax_main, env_full, obs, PLOT_START, PLOT_END, helene_label=False) + ax_main.set_title("Full Helene window context (Sep 24 - Sep 28 18 UTC)", fontsize=10) + _format_xaxis(ax_main) + # Shade the zoom region on context panel + ax_main.axvspan(PEAK_START, PEAK_END, color="gold", alpha=0.18, zorder=1) + ax_main.text(PEAK_START + pd.Timedelta(hours=1), 0.97, "zoom", + transform=ax_main.get_xaxis_transform(), + fontsize=8, color="goldenrod", ha="left", va="top") + + # Bottom panel: peak zoom + plot_envelope(ax_zoom, env_peak, obs, PEAK_START, PEAK_END, helene_label=True) + ax_zoom.set_xlabel("Date", fontsize=11) + ax_zoom.set_title( + "Helene peak zoom (Sep 26 12 UTC - Sep 28 06 UTC)", fontsize=10) + _format_xaxis(ax_zoom, major_interval=1, minor_hours=3) + ax_zoom.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %HUTC")) + ax_zoom.legend(fontsize=10, loc="lower right", frameon=True, framealpha=0.9) + + fig.suptitle( + f"4b -- 600-member crossed ensemble vs USGS 03463300 | {args.cat_id}\n" + "30 met forcing draws x 20 DA analysis state draws | R=0.07 mm2/h2", + fontsize=12, y=1.01, + ) + plt.tight_layout() + out2 = os.path.join(out_dir, f"{args.cat_id}_4b_crossed_ensemble_peak.png") + plt.savefig(out2, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_combined.py b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_combined.py new file mode 100644 index 00000000..c6d531a0 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_combined.py @@ -0,0 +1,235 @@ +#!/usr/bin/env python3 +""" +plot_routed_ensemble_combined.py +Combined poster plot: 600-member routed ensemble envelope + +deterministic DA lines (Vrugt/no-Vrugt) vs USGS gauge obs at 03463300. + +Inputs: + --vrugt-csv : routed_Q_test.csv from vrugt_dynamic_routed/ + --novrugt-csv : routed_Q_test.csv from novrugt_dynamic_routed/ + --ensemble-pq : routed_crossed_ensemble.parquet from run_route_crossed_ensemble.py + --out-dir : output directory + +The ensemble parquet columns: issue_time, lead_hour, member_0000..member_0599 (Q in m3/s) + +Usage: + python3 plot_routed_ensemble_combined.py \ + --vrugt-csv /mnt/disk2/suma_helen_poster/da_results/vrugt_dynamic_routed/routed_Q_test.csv \ + --novrugt-csv /mnt/disk2/suma_helen_poster/da_results/novrugt_dynamic_routed/routed_Q_test.csv \ + --ensemble-pq /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble_routed/routed_crossed_ensemble.parquet \ + --out-dir /mnt/disk2/suma_helen_poster/da_results/comparison_plots +""" + +import argparse +import os + +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +HELENE_START = pd.Timestamp("2024-09-25 00:00:00") +HELENE_END = pd.Timestamp("2024-09-29 00:00:00") +ZOOM_START = pd.Timestamp("2024-09-24 00:00:00") +ZOOM_END = pd.Timestamp("2024-09-28 18:00:00") # cut trailing ensemble bump + +COLOR_OBS = "black" +COLOR_VRUGT = "#1f77b4" # blue +COLOR_NOVRUGT = "#d62728" # red +COLOR_ENS = "#1f77b4" # same blue family as Vrugt + + +def load_det_csv(path): + df = pd.read_csv(path, parse_dates=["date"]) + return df.set_index("date").sort_index() + + +def build_envelope(pq_path, zoom_start, zoom_end): + """ + Load routed ensemble parquet and build percentile envelope by valid_time. + Returns DataFrame indexed by valid_time with columns p05, p25, p50, p75, p95. + """ + df = pd.read_parquet(pq_path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = [c for c in df.columns if c.startswith("member_")] + + records = [] + for _, row in df.iterrows(): + t0 = row["issue_time"] + lead = int(row["lead_hour"]) + vt = t0 + pd.Timedelta(hours=lead) + vals = row[member_cols].values.astype(float) + records.append({"valid_time": vt, "vals": vals}) + + # Group by valid_time (floor to hour), stack all member values across issue_times + from collections import defaultdict + groups = defaultdict(list) + for rec in records: + vt = rec["valid_time"].floor("1h") + groups[vt].extend(rec["vals"].tolist()) + + rows = [] + for vt in sorted(groups): + if vt < zoom_start or vt > zoom_end: + continue + v = np.array(groups[vt]) + v = v[~np.isnan(v)] + if len(v) == 0: + continue + rows.append({ + "valid_time": vt, + "p05": np.percentile(v, 5), + "p25": np.percentile(v, 25), + "p50": np.percentile(v, 50), + "p75": np.percentile(v, 75), + "p95": np.percentile(v, 95), + }) + + return pd.DataFrame(rows).set_index("valid_time") + + +def compute_kge(obs, sim): + mask = ~(np.isnan(obs) | np.isnan(sim)) + o, s = obs[mask], sim[mask] + if len(o) < 2 or np.std(o) == 0: + return np.nan + r = np.corrcoef(o, s)[0, 1] + return 1.0 - np.sqrt((r-1)**2 + (np.std(s)/np.std(o)-1)**2 + (np.mean(s)/np.mean(o)-1)**2) + + +def _format_xaxis(ax, locator, fmt, rotation=30): + ax.xaxis.set_major_locator(locator) + ax.xaxis.set_major_formatter(mdates.DateFormatter(fmt)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=rotation, ha="right", fontsize=9) + + +def _add_helene_band(ax): + ax.axvspan(HELENE_START, HELENE_END, color="gold", alpha=0.15, zorder=0) + + +def plot_zoom(ax, dates, obs, sim_v, sim_nv, kge_v, kge_nv, env=None): + mask = (dates >= ZOOM_START) & (dates <= ZOOM_END) + dz = dates[mask] + + # Ensemble envelope (bottom layer) + if env is not None: + ax.fill_between(env.index, env["p05"], env["p95"], + color=COLOR_ENS, alpha=0.18, label="Ensemble 5-95th pct") + ax.fill_between(env.index, env["p25"], env["p75"], + color=COLOR_ENS, alpha=0.30, label="Ensemble 25-75th pct") + ax.plot(env.index, env["p50"], + color=COLOR_ENS, lw=1.2, linestyle="--", alpha=0.7, label="Ensemble median") + + # Deterministic DA lines + ax.plot(dz, sim_nv[mask], color=COLOR_NOVRUGT, lw=1.4, linestyle="--", + label=f"DA -- constant R KGE={kge_nv:.3f}", zorder=3) + ax.plot(dz, sim_v[mask], color=COLOR_VRUGT, lw=1.6, + label=f"DA -- dynamic R (Vrugt) KGE={kge_v:.3f}", zorder=4) + + # USGS obs on top + ax.plot(dz, obs[mask], color=COLOR_OBS, lw=1.6, label="USGS obs (gauge 03463300)", zorder=5) + + peak_usgs = np.nanmax(obs[mask]) + ax.axhline(peak_usgs, color=COLOR_OBS, lw=0.7, linestyle=":", alpha=0.5) + # Place annotation to the right of the peak to avoid legend overlap + ax.text(pd.Timestamp("2024-09-28 06:00:00"), peak_usgs * 1.02, + f"USGS peak {peak_usgs:.0f} m³/s", fontsize=8.5, color="black", + alpha=0.8, ha="right") + + _add_helene_band(ax) + _format_xaxis(ax, mdates.DayLocator(interval=1), "%b %d") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--vrugt-csv", required=True) + parser.add_argument("--novrugt-csv", required=True) + parser.add_argument("--ensemble-pq", required=True, + help="routed_crossed_ensemble.parquet from run_route_crossed_ensemble.py") + parser.add_argument("--out-dir", required=True) + args = parser.parse_args() + os.makedirs(args.out_dir, exist_ok=True) + + vrugt = load_det_csv(args.vrugt_csv) + novrugt = load_det_csv(args.novrugt_csv) + + obs = vrugt["Q_usgs_m3s"].values + dates = vrugt.index + sim_v = vrugt["Q_routed_m3s"].values + sim_nv = novrugt["Q_routed_m3s"].reindex(dates).values + + kge_v = compute_kge(obs, sim_v) + kge_nv = compute_kge(obs, sim_nv) + + print("Building ensemble envelope...") + env = build_envelope(args.ensemble_pq, ZOOM_START, ZOOM_END) + print(f" Envelope: {len(env)} timesteps, " + f"peak p95={env['p95'].max():.1f} m3/s, peak p50={env['p50'].max():.1f} m3/s") + + peak_usgs = np.nanmax(obs[(dates >= ZOOM_START) & (dates <= ZOOM_END)]) + in_env = env["p95"].max() >= peak_usgs + print(f" USGS peak {peak_usgs:.1f} m3/s {'IS' if in_env else 'IS NOT'} within p95 envelope") + + # ------------------------------------------------------------------ # + # Figure 1: Helene zoom only -- main poster panel + # ------------------------------------------------------------------ # + fig, ax = plt.subplots(figsize=(11, 5.5)) + plot_zoom(ax, dates, obs, sim_v, sim_nv, kge_v, kge_nv, env) + ax.set_xlabel("Date (UTC)", fontsize=11) + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_title( + "Hurricane Helene -- Probabilistic discharge forecast at gauge 03463300 (South Toe River)\n" + "600-member CFE ensemble + DA (Muskingum routing)", + fontsize=11) + ax.legend(fontsize=9, loc="lower right") + ax.grid(True, alpha=0.25) + plt.tight_layout() + out1 = os.path.join(args.out_dir, "helene_ensemble_vs_usgs.png") + plt.savefig(out1, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + # ------------------------------------------------------------------ # + # Figure 2: two-panel (full period + Helene zoom) + # ------------------------------------------------------------------ # + fig, (ax_top, ax_bot) = plt.subplots( + 2, 1, figsize=(14, 9), + gridspec_kw={"height_ratios": [1, 1.5]}) + + # Top: full test period (deterministic only, no ensemble for readability) + ax_top.plot(dates, obs, color=COLOR_OBS, lw=0.8, label="USGS obs", zorder=3) + ax_top.plot(dates, sim_v, color=COLOR_VRUGT, lw=0.9, + label=f"DA dynamic R (Vrugt) KGE={kge_v:.3f}", zorder=2) + ax_top.plot(dates, sim_nv, color=COLOR_NOVRUGT, lw=0.9, linestyle="--", + label=f"DA constant R KGE={kge_nv:.3f}", zorder=2) + ax_top.axvspan(ZOOM_START, ZOOM_END, color="gray", alpha=0.12, zorder=0) + ax_top.axvspan(HELENE_START, HELENE_END, color="gold", alpha=0.15, zorder=0) + ax_top.text(HELENE_START + pd.Timedelta(days=0.5), + np.nanmax(obs) * 0.88, "Helene", + fontsize=9, color="goldenrod", fontweight="bold") + ax_top.set_ylabel("Discharge (m³/s)", fontsize=10) + ax_top.set_title( + "Routed discharge at gauge 03463300 (South Toe River) -- Test period\n" + "DA with dynamic Vrugt R vs. constant R -- Muskingum routing", + fontsize=11) + ax_top.legend(fontsize=9, loc="upper left") + ax_top.grid(True, alpha=0.25) + _format_xaxis(ax_top, mdates.MonthLocator(interval=2), "%Y-%m") + + # Bottom: Helene zoom with ensemble + plot_zoom(ax_bot, dates, obs, sim_v, sim_nv, kge_v, kge_nv, env) + ax_bot.set_xlabel("Date (UTC)", fontsize=10) + ax_bot.set_ylabel("Discharge (m³/s)", fontsize=10) + ax_bot.set_title("Hurricane Helene window -- 600-member ensemble envelope", fontsize=10) + ax_bot.legend(fontsize=9, loc="lower right") + ax_bot.grid(True, alpha=0.25) + + plt.tight_layout() + out2 = os.path.join(args.out_dir, "helene_ensemble_twopanel.png") + plt.savefig(out2, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_vs_usgs.py b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_vs_usgs.py new file mode 100644 index 00000000..55c52d85 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_vs_usgs.py @@ -0,0 +1,176 @@ +#!/usr/bin/env python3 +""" +plot_routed_ensemble_vs_usgs.py +Plot T-route routed 600-member ensemble envelope vs USGS at gauge 03463300. + +Input: routed_crossed_ensemble.parquet from run_route_crossed_ensemble.py + Columns: issue_time, lead_hour, member_0000..member_0599 (Q in m3/s) + +Usage: + python3 plot_routed_ensemble_vs_usgs.py \ + --routed-pq /mnt/disk2/.../folder1_vrugt_routed/routed_crossed_ensemble.parquet \ + --usgs-csv /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv \ + --label "F1 Vrugt — 1 gauge holdout" \ + --out-dir ~/plots_f1_1gauge +""" + +import argparse +from collections import defaultdict +from pathlib import Path + +import matplotlib.dates as mdates +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd + +WATERSHED_AREA_KM2 = 113.18 # cat-1016300 drainage area +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +PLOT_START = pd.Timestamp("2024-09-24 00:00:00") +PLOT_END = pd.Timestamp("2024-09-28 18:00:00") +PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +PEAK_END = pd.Timestamp("2024-09-28 06:00:00") +HELENE_PEAK = pd.Timestamp("2024-09-27 12:00:00") + +COLOR_ENS = "#1f77b4" +COLOR_OBS = "black" + + +def build_envelope(pq_path, t_start, t_end): + df = pd.read_parquet(pq_path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = [c for c in df.columns if c.startswith("member_")] + + groups = defaultdict(list) + for _, row in df.iterrows(): + vt = (row["issue_time"] + pd.Timedelta(hours=int(row["lead_hour"]))).floor("1h") + if t_start <= vt <= t_end: + groups[vt].extend(row[member_cols].values.astype(float).tolist()) + + rows = [] + for vt in sorted(groups): + v = np.array(groups[vt]) + v = v[~np.isnan(v)] + if len(v) == 0: + continue + rows.append({"valid_time": vt, + "p05": np.percentile(v, 5), "p25": np.percentile(v, 25), + "p50": np.percentile(v, 50), "p75": np.percentile(v, 75), + "p95": np.percentile(v, 95)}) + return pd.DataFrame(rows).set_index("valid_time") + + +def load_usgs(path): + df = pd.read_csv(path) + print(f" USGS CSV columns: {list(df.columns)}") + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() or "discharge" in c.lower()) + print(f" Using date_col={date_col!r}, q_col={q_col!r}") + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series.sort_index() + + +def compute_kge(obs, sim): + mask = ~(np.isnan(obs) | np.isnan(sim)) + o, s = obs[mask], sim[mask] + if len(o) < 2 or np.std(o) == 0: + return np.nan + r = np.corrcoef(o, s)[0, 1] + return 1 - np.sqrt((r-1)**2 + (np.std(s)/np.std(o)-1)**2 + (np.mean(s)/np.mean(o)-1)**2) + + +def _shade(ax, env): + ax.fill_between(env.index, env["p05"], env["p95"], + color=COLOR_ENS, alpha=0.18, label="5th–95th percentile") + ax.fill_between(env.index, env["p25"], env["p75"], + color=COLOR_ENS, alpha=0.32, label="25th–75th percentile") + ax.plot(env.index, env["p50"], color=COLOR_ENS, lw=1.5, label="Ensemble median (600 members)") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--routed-pq", required=True) + parser.add_argument("--usgs-csv", required=True) + parser.add_argument("--label", default="F1 Vrugt — 1 gauge holdout") + parser.add_argument("--out-dir", required=True) + args = parser.parse_args() + + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + print("Building full-window envelope...") + env_full = build_envelope(args.routed_pq, PLOT_START, PLOT_END) + print("Building peak-window envelope...") + env_peak = build_envelope(args.routed_pq, PEAK_START, PEAK_END) + + print("Loading USGS obs...") + try: + obs_all = load_usgs(args.usgs_csv) + obs_full = obs_all.loc[PLOT_START:PLOT_END] + obs_peak = obs_all.loc[PEAK_START:PEAK_END] + print(f" Loaded {len(obs_all)} USGS obs rows; " + f"full window: {len(obs_full)}, peak window: {len(obs_peak)}") + except Exception as e: + print(f" Warning: USGS load failed ({e}) — plotting without obs") + obs_full = obs_peak = None + + kge = np.nan + if obs_full is not None and len(env_full) > 0: + med = env_full["p50"] + obs_hourly = obs_full.copy() + obs_hourly.index = obs_hourly.index.floor("1h") + obs_hourly = obs_hourly[~obs_hourly.index.duplicated(keep="first")] + aligned = obs_hourly.reindex(med.index) + print(f" KGE alignment: {aligned.notna().sum()} matched out of {len(med)} timesteps") + kge = compute_kge(aligned.values, med.values) + + # ── Figure 1: Full Helene window ────────────────────────────────────────── + fig, ax = plt.subplots(figsize=(12, 4.5)) + _shade(ax, env_full) + if obs_full is not None: + ax.plot(obs_full.index, obs_full.values, color=COLOR_OBS, lw=1.6, + label="USGS obs (gauge 03463300)") + ax.axvspan(PEAK_START, PEAK_END, color="salmon", alpha=0.12, zorder=0, label="Helene peak window") + ax.set_ylabel("Discharge (m³/s)") + ax.set_xlabel("Date") + ax.set_title(f"Routed 600-member ensemble vs USGS | {args.label}\n" + f"Ensemble median KGE = {kge:.3f} | Sep 24–28 2024") + ax.legend(fontsize=9, loc="upper left") + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + plt.tight_layout() + out = out_dir / "routed_ensemble_vs_usgs_full.png" + fig.savefig(out, dpi=150) + print(f"Saved: {out}") + plt.close(fig) + + # ── Figure 2: Peak zoom ─────────────────────────────────────────────────── + fig, ax = plt.subplots(figsize=(10, 4.5)) + _shade(ax, env_peak) + if obs_peak is not None: + ax.plot(obs_peak.index, obs_peak.values, color=COLOR_OBS, lw=1.6, + label="USGS obs (gauge 03463300)") + peak_val = float(obs_peak.max()) + ax.axhline(peak_val, color=COLOR_OBS, lw=0.7, ls=":", alpha=0.5) + ax.annotate(f"USGS peak\n{peak_val:.0f} m³/s", + xy=(HELENE_PEAK, peak_val), xytext=(10, -40), + textcoords="offset points", fontsize=8, color=COLOR_OBS) + ax.set_ylabel("Discharge (m³/s)") + ax.set_xlabel("Date") + ax.set_title(f"Helene peak zoom | {args.label}") + ax.legend(fontsize=9) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %HUTC")) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + plt.tight_layout() + out = out_dir / "routed_ensemble_vs_usgs_peak.png" + fig.savefig(out, dpi=150) + print(f"Saved: {out}") + plt.close(fig) + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/run_4b_crossed_f2.sh b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/run_4b_crossed_f2.sh new file mode 100644 index 00000000..f39fb7b5 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/run_4b_crossed_f2.sh @@ -0,0 +1,40 @@ +#!/usr/bin/env bash +# F2 fixed R=0.07 (20pct gauge holdout) — 4b crossed ensemble plots. +# Runs plot_crossed_ensemble.py for all 21 catchments. +# Reads //_crossed_ensemble.parquet (already generated). +# +# Usage: +# bash run_4b_crossed_f2.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +DATA_DIR=/mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +echo "[F2-4b-crossed] Crossed ensemble plots — data: $DATA_DIR" + +for CAT in "${CATS[@]}"; do + PQ="$DATA_DIR/$CAT/${CAT}_crossed_ensemble.parquet" + if [ ! -f "$PQ" ]; then + echo " [$CAT] No crossed ensemble parquet — skipping" + continue + fi + echo " [$CAT] crossed ensemble..." + $TROUTE "$SCRIPT_DIR/plot_crossed_ensemble.py" \ + --ensemble-dir "$DATA_DIR" \ + --cat-id "$CAT" \ + --usgs-csv "$USGS_CSV" +done + +echo "[F2-4b-crossed] Done." diff --git a/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/run_4b_f2.sh b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/run_4b_f2.sh new file mode 100644 index 00000000..b4f3f665 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4b_ensemble_vs_obs/run_4b_f2.sh @@ -0,0 +1,40 @@ +#!/usr/bin/env bash +# F2 fixed R=0.07 (20pct gauge holdout) — 4b routed ensemble plots. +# +# Runs: +# 1. plot_routed_ensemble_vs_usgs.py — F2 ensemble envelope vs USGS at outlet +# 2. plot_routed_ensemble_combined.py — F1 Vrugt vs F2 fixed R comparison +# +# Requires routed_Q_test.csv (both folders) and routed_crossed_ensemble.parquet +# (CROSSED_DIR) to be present. +# +# Usage: +# bash run_4b_f2.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +F1_DA_DIR=/mnt/disk2/suma_helen_poster/da_results/v2_perturbation_da_on +F2_DA_DIR=/mnt/disk2/suma_helen_poster/da_results/v2_fixed_r007_analysis +CROSSED_DIR=/mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv +OUT_DIR="$CROSSED_DIR" + +echo "[F2-4b] Routed ensemble vs USGS (F2 only)..." +$TROUTE "$SCRIPT_DIR/plot_routed_ensemble_vs_usgs.py" \ + --routed-pq "$CROSSED_DIR/routed_crossed_ensemble.parquet" \ + --usgs-csv "$USGS_CSV" \ + --label "F2 Fixed R=0.07 — 20pct gauge holdout" \ + --out-dir "$OUT_DIR" + +echo "[F2-4b] Combined comparison: F1 Vrugt vs F2 Fixed R..." +$TROUTE "$SCRIPT_DIR/plot_routed_ensemble_combined.py" \ + --vrugt-csv "$F1_DA_DIR/routed_Q_test.csv" \ + --novrugt-csv "$F2_DA_DIR/routed_Q_test.csv" \ + --ensemble-pq "$CROSSED_DIR/routed_crossed_ensemble.parquet" \ + --out-dir "$OUT_DIR" + +echo "[F2-4b] Done. Outputs in: $OUT_DIR" +ls "$OUT_DIR"/*.png 2>/dev/null || echo " (no PNGs in out dir)" diff --git a/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4c_timeseries/plot_forecast_spaghetti.py b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4c_timeseries/plot_forecast_spaghetti.py new file mode 100644 index 00000000..655edaef --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4c_timeseries/plot_forecast_spaghetti.py @@ -0,0 +1,176 @@ +""" +plot_forecast_spaghetti.py + +"All forecasts ending at this hour" — verification-time-fixed spaghetti view. + +Window: Sep 24 18 UTC -> Sep 30 06 UTC. +For every initialization time whose 18-hour forecast overlaps this window: + DA : shaded min/max band across 20 members + ensemble mean line + OL : ensemble mean line only (no shading), dashed + +Lines are colored by initialization DATE (7 colors, Sep 24-30) so the +temporal progression is readable. USGS obs overlaid in black. + +Inputs: + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + Wide format: issue_time, lead_hour, member_00..member_19 (q in m³/s) + +Obs: + /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +Outputs: + /forecast_spaghetti_helene.png +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import matplotlib.patches as mpatches + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast_hardcoded_r_routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +# Verification window — only trajectories whose valid_time falls here are shown +PLOT_START = pd.Timestamp("2024-09-24 18:00:00") +PLOT_END = pd.Timestamp("2024-09-30 06:00:00") + +HELENE_PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +HELENE_PEAK_END = pd.Timestamp("2024-09-28 00:00:00") + +# One color per init date (Sep 24-30) +DATE_COLORS = { + "2024-09-24": "#1f77b4", + "2024-09-25": "#ff7f0e", + "2024-09-26": "#2ca02c", + "2024-09-27": "#d62728", + "2024-09-28": "#9467bd", + "2024-09-29": "#8c564b", + "2024-09-30": "#e377c2", +} + + +def load_parquet(path): + df = pd.read_parquet(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + df["ens_mean"] = df[member_cols].mean(axis=1) + df["ens_min"] = df[member_cols].min(axis=1) + df["ens_max"] = df[member_cols].max(axis=1) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + return df + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() + or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--route-dir", default=DEFAULT_ROUTE_DIR) + parser.add_argument("--out-dir", default=None) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--da-name", default="routed_leadtime_da_full.parquet") + parser.add_argument("--ol-name", default="routed_leadtime_openloop_full.parquet") + args = parser.parse_args() + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + print("Loading parquets...") + da = load_parquet(os.path.join(args.route_dir, args.da_name)) + ol = load_parquet(os.path.join(args.route_dir, args.ol_name)) + obs = load_usgs(args.usgs_csv) + + # Keep only valid_times inside the plot window + da = da[(da["valid_time"] >= PLOT_START) & (da["valid_time"] <= PLOT_END)] + ol = ol[(ol["valid_time"] >= PLOT_START) & (ol["valid_time"] <= PLOT_END)] + + issue_times = sorted(da["issue_time"].unique()) + print(f" {len(issue_times)} initialization times contribute to this window") + + fig, ax = plt.subplots(figsize=(18, 6)) + + # Helene peak shading + ax.axvspan(HELENE_PEAK_START, HELENE_PEAK_END, + color="salmon", alpha=0.12, zorder=0) + ax.text(HELENE_PEAK_START + pd.Timedelta(hours=6), 1.0, "Helene peak", + transform=ax.get_xaxis_transform(), + fontsize=9, color="firebrick", ha="left", va="top") + + # Plot each initialization time's trajectory + for t0 in issue_times: + date_str = str(pd.Timestamp(t0).date()) + color = DATE_COLORS.get(date_str, "gray") + + da_t = da[da["issue_time"] == t0].sort_values("valid_time") + ol_t = ol[ol["issue_time"] == t0].sort_values("valid_time") + + if da_t.empty: + continue + + # DA: shaded band (min-max across 20 members) + ensemble mean + ax.fill_between(da_t["valid_time"], da_t["ens_min"], da_t["ens_max"], + color=color, alpha=0.06, zorder=2) + ax.plot(da_t["valid_time"], da_t["ens_mean"], + color=color, lw=0.8, alpha=0.55, zorder=3) + + # OL: ensemble mean only, dashed + if not ol_t.empty: + ax.plot(ol_t["valid_time"], ol_t["ens_mean"], + color=color, lw=0.6, alpha=0.30, linestyle="--", zorder=2) + + # USGS obs + obs_w = obs.loc[PLOT_START:PLOT_END] + ax.plot(obs_w.index, obs_w.values, + color="black", lw=1.8, zorder=6, label="USGS obs") + + # Legend: one patch per init date + obs + DA/OL style + patches = [mpatches.Patch(color=c, label=f"Init {d}") + for d, c in DATE_COLORS.items()] + patches.append(plt.Line2D([0], [0], color="black", lw=1.8, label="USGS obs")) + patches.append(plt.Line2D([0], [0], color="gray", lw=1.2, alpha=0.7, + label="DA — ens mean (solid) + spread (shaded)")) + patches.append(plt.Line2D([0], [0], color="gray", lw=1.0, linestyle="--", + alpha=0.5, label="OL — ens mean (dashed)")) + ax.legend(handles=patches, fontsize=8, loc="upper left", + frameon=True, framealpha=0.9, ncol=2) + + ax.set_xlim(PLOT_START, PLOT_END) + ax.set_xlabel("Date", fontsize=11) + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_title( + "All 18-hour forecast trajectories ending in this window — USGS 03463300\n" + "Sep 24 18 UTC → Sep 30 06 UTC | DA: shaded band + mean | OL: mean dashed", + fontsize=11, + ) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=1)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.xaxis.set_minor_locator(mdates.HourLocator(interval=6)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.2, lw=0.4) + plt.tight_layout() + + out_path = os.path.join(out_dir, "forecast_spaghetti_helene.png") + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4c_timeseries/plot_helene_issue_time_hydrograph.py b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4c_timeseries/plot_helene_issue_time_hydrograph.py new file mode 100644 index 00000000..4e2e1dc2 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4c_timeseries/plot_helene_issue_time_hydrograph.py @@ -0,0 +1,198 @@ +""" +Per-issue-time forecast hydrograph diagnostic. + +For two contrasting issue times — one storm-peak (Helene) and one low-flow — +show what the DA and open-loop ensembles actually produce over the 18-hour +forecast window, alongside the kriging observation. Designed to answer in one +picture: is DA over-shooting, collapsing, or oscillating compared to open-loop? + +Each panel shows: + - All 20 DA members as faint purple lines + median (thick purple) + - All 20 open-loop members as faint gray lines + median (thick gray, dashed) + - Kriging obs as a thick black line (context: t0-6 through t0+18) + - Vertical line at t0 + annotation with obs(t0) and lead-1 ensemble means + +Inputs (from run_lead_time_forecast_sweep.py): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv +Obs (from production EnKF): + //_test_results.csv + +Output: + //_issue_time_hydrograph_helene_vs_lowflow.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +DEFAULT_LEADTIME_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast" +DEFAULT_DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" + +# Module-level globals set by main() +CAT = "cat-1016300" +LEADTIME_DIR = DEFAULT_LEADTIME_DIR +DA_DIR = DEFAULT_DA_DIR + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" +OBS_COLOR = "black" + +# Two issue times — pick a Helene-peak one and a typical low-flow one +HELENE_T0 = pd.Timestamp("2024-09-26 12:00:00") +LOWFLOW_T0 = pd.Timestamp("2024-03-15 12:00:00") + +CONTEXT_HOURS_BEFORE = 6 # hours of obs context shown before t0 +LEAD_HOURS_AFTER = 18 # forecast horizon + +OUT_PNG = os.path.join( + LEADTIME_DIR, CAT, + f"{CAT}_issue_time_hydrograph_helene_vs_lowflow.png", +) + + +def load_forecasts(path): + df = pd.read_csv(path, parse_dates=['issue_time', 'valid_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + return df, member_cols + + +def load_obs(): + obs_path = os.path.join(DA_DIR, CAT, f"{CAT}_test_results.csv") + df = pd.read_csv(obs_path, parse_dates=['date']) + return df.set_index('date')['obs_mm_h'] + + +def slice_forecast(df, member_cols, t0): + """Return (valid_times, member_array shape (lead, N)) for a single issue time.""" + sub = df[df['issue_time'] == t0].sort_values('lead_hour') + if len(sub) == 0: + raise ValueError(f"No forecast rows for issue_time={t0}") + valid_times = pd.to_datetime(sub['valid_time'].values) + member_arr = sub[member_cols].to_numpy(dtype=float) + return valid_times, member_arr + + +def plot_panel(ax, t0, df_da, df_ol, member_cols, obs_series, title_extra=""): + # Forecast trajectories + vt_da, q_da = slice_forecast(df_da, member_cols, t0) + vt_ol, q_ol = slice_forecast(df_ol, member_cols, t0) + + # Obs context window + ctx_start = t0 - pd.Timedelta(hours=CONTEXT_HOURS_BEFORE) + ctx_end = t0 + pd.Timedelta(hours=LEAD_HOURS_AFTER) + obs_window = obs_series.loc[ctx_start:ctx_end] + + # Open-loop members (drawn first so DA paints on top) + for j in range(q_ol.shape[1]): + ax.plot(vt_ol, q_ol[:, j], color=OL_COLOR, lw=0.6, alpha=0.45, zorder=2) + # DA members + for j in range(q_da.shape[1]): + ax.plot(vt_da, q_da[:, j], color=DA_COLOR, lw=0.6, alpha=0.45, zorder=3) + # Medians + ax.plot(vt_ol, np.nanmedian(q_ol, axis=1), color=OL_COLOR, lw=2.6, + linestyle="--", zorder=4, label="Open-loop median") + ax.plot(vt_da, np.nanmedian(q_da, axis=1), color=DA_COLOR, lw=2.6, + zorder=5, label="DA median") + # Obs + ax.plot(obs_window.index, obs_window.values, color=OBS_COLOR, lw=1.8, + zorder=6, label="Qkrig (obs)") + + # t0 marker + annotation + ax.axvline(t0, color="black", lw=1.0, linestyle=":", alpha=0.6, zorder=1) + obs_t0 = float(obs_series.get(t0, np.nan)) + obs_lead1 = float(obs_series.get(t0 + pd.Timedelta(hours=1), np.nan)) + da_lead1_mean = float(np.nanmean(q_da[0, :])) + ol_lead1_mean = float(np.nanmean(q_ol[0, :])) + + info_text = ( + f"t0 = {t0.strftime('%Y-%m-%d %H:%M')}\n" + f"obs(t0) = {obs_t0:.3f} mm/h\n" + f"obs(t0+1) = {obs_lead1:.3f} mm/h\n" + f"DA mean(t0+1) = {da_lead1_mean:.3f} mm/h\n" + f"OL mean(t0+1) = {ol_lead1_mean:.3f} mm/h" + ) + ax.text(0.02, 0.97, info_text, + transform=ax.transAxes, + fontsize=9, va='top', ha='left', family='monospace', + bbox=dict(boxstyle='round,pad=0.4', facecolor='white', + alpha=0.85, edgecolor='gray')) + + ax.set_title(title_extra, fontsize=12) + ax.set_ylabel("Discharge (mm/h)", fontsize=10) + ax.set_xlabel("Date", fontsize=10) + ax.xaxis.set_major_locator(mdates.HourLocator(interval=3)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %H:%M")) + ax.tick_params(axis='x', rotation=30, labelsize=8) + ax.grid(True, alpha=0.2) + ax.legend(loc='upper right', fontsize=9, frameon=True, framealpha=0.9) + + +def snap_to_nearest(t0, available_times): + """Return the available issue_time closest to t0.""" + times = pd.DatetimeIndex(available_times) + idx = np.abs((times - t0).total_seconds()).argmin() + snapped = times[idx] + if snapped != t0: + print(f" Snapped {t0} -> {snapped} (nearest available issue time)") + return snapped + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', default="cat-1016300") + parser.add_argument('--leadtime-dir', default=DEFAULT_LEADTIME_DIR, + help='Dir holding /_lead_time_forecasts_{da,openloop}.csv') + parser.add_argument('--da-dir', default=DEFAULT_DA_DIR, + help='Dir holding /_test_results.csv (for obs)') + parser.add_argument('--helene-t0', default=None, + help='Issue time for the Helene panel (default: 2024-09-26 12:00:00, ' + 'snapped to nearest available)') + parser.add_argument('--lowflow-t0', default=None, + help='Issue time for the low-flow panel (default: 2024-03-15 00:00:00, ' + 'snapped to nearest available)') + args = parser.parse_args() + + global CAT, LEADTIME_DIR, DA_DIR, OUT_PNG + CAT = args.cat_id + LEADTIME_DIR = args.leadtime_dir + DA_DIR = args.da_dir + OUT_PNG = os.path.join(LEADTIME_DIR, CAT, + f"{CAT}_issue_time_hydrograph_helene_vs_lowflow.png") + + da_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_da.csv") + ol_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_openloop.csv") + df_da, m_cols = load_forecasts(da_path) + df_ol, _ = load_forecasts(ol_path) + obs_series = load_obs() + + available = df_da['issue_time'].unique() + helene_t0 = snap_to_nearest( + pd.Timestamp(args.helene_t0) if args.helene_t0 else HELENE_T0, available) + lowflow_t0 = snap_to_nearest( + pd.Timestamp(args.lowflow_t0) if args.lowflow_t0 else pd.Timestamp("2024-03-15 00:00:00"), + available) + + fig, (ax_h, ax_l) = plt.subplots(1, 2, figsize=(20, 7)) + + plot_panel(ax_h, helene_t0, df_da, df_ol, m_cols, obs_series, + title_extra=f"Helene peak issue time — {helene_t0.strftime('%Y-%m-%d %H:%M')}") + plot_panel(ax_l, lowflow_t0, df_da, df_ol, m_cols, obs_series, + title_extra=f"Low-flow issue time — {lowflow_t0.strftime('%Y-%m-%d %H:%M')}") + + fig.suptitle( + f"Per-issue-time forecast hydrograph — {CAT}\n" + "Faint lines = 20 ensemble members. Thick lines = ensemble medians. " + "Vertical dotted line = t0 (DA stops here for both scenarios).", + fontsize=12, y=1.00, + ) + plt.tight_layout() + plt.savefig(OUT_PNG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4c_timeseries/plot_lead_time_reconstructed_timeseries.py b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4c_timeseries/plot_lead_time_reconstructed_timeseries.py new file mode 100644 index 00000000..a6d03b03 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4c_timeseries/plot_lead_time_reconstructed_timeseries.py @@ -0,0 +1,263 @@ +""" +Reconstructed time series at gauge 03463300, from our routed +lead-time forecasts. Mirrors the methodology of plot_timeseries_hourly_ensemble.py +(qSpatialAR) but feeds from the EnKF + T-route pipeline instead of the CNN. + +For each target hour t, pool ALL forecasts that land on t (across many +issue_times × many lead_hours × 20 members). Compute median + 5th/95th +percentile envelope. Plot DA and open-loop reconstructions overlaid on USGS +obs, with the Helene window shaded. + +Input parquets (from run_route_troute_leadtime.py): + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + Wide format: issue_time, lead_hour, member_00..member_19 (q in m³/s) + +Obs: + /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +Output: + /lead_time_reconstructed_timeseries.png + Default plot window: 2024-09-10 → 2024-10-10. + +Each scenario's curve is annotated with NSE vs USGS obs over the plot window. +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" +OBS_COLOR = "black" + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast_hardcoded_r_routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +# Watershed area at gauge 03463300 (sum of 21 catchments per GPKG hydrofabric) +WATERSHED_AREA_KM2 = 113.18 +# mm/h depth → m³/s: × 113.18 km² × 1000 / 3600 = 31.439 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +PLOT_START_DEFAULT = pd.Timestamp("2024-09-10 00:00:00") +PLOT_END_DEFAULT = pd.Timestamp("2024-10-10 23:00:00") + +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-30 23:00:00") + + +def load_parquet_long(path): + df = pd.read_parquet(path) + df['issue_time'] = pd.to_datetime(df['issue_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + if not member_cols: + raise ValueError(f"No member columns in {path}") + long = df.melt( + id_vars=['issue_time', 'lead_hour'], + value_vars=member_cols, + var_name='member', + value_name='q_m3s', + ) + long['valid_time'] = (long['issue_time'] + + pd.to_timedelta(long['lead_hour'], unit='h')) + return long + + +def reconstruct(long_df): + """For each valid_time, pool all (issue_time × lead × member) forecasts and + return median, p05, p95, count.""" + grouped = long_df.groupby('valid_time')['q_m3s'].agg( + median='median', + p05=lambda s: s.quantile(0.05), + p95=lambda s: s.quantile(0.95), + count='count', + ).reset_index() + return grouped + + +def load_usgs(usgs_csv): + """Load USGS gauge obs and ensure m³/s. + + The CSV at gauge 03463300 has the column `QObs(mm/h)` — catchment-averaged + depth, not gauge discharge in m³/s. Auto-convert when column name contains + 'mm': m³/s = mm/h × WATERSHED_AREA_KM2 × 1000 / 3600. + """ + df = pd.read_csv(usgs_csv) + date_col = next((c for c in df.columns + if c.lower() in ('datetime', 'date', 'time', 'timestamp')), None) + q_col = next((c for c in df.columns + if 'q' in c.lower() or 'flow' in c.lower() + or 'discharge' in c.lower()), None) + if not date_col or not q_col: + raise ValueError(f"Couldn't id date/Q columns in {usgs_csv}: {df.columns.tolist()}") + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if 'mm' in q_col.lower(): + series = series * MM_H_TO_M3_S + print(f" Converted obs from mm/h (column '{q_col}') to m³/s " + f"using area = {WATERSHED_AREA_KM2} km² (× {MM_H_TO_M3_S:.4f})") + else: + print(f" Loaded obs as-is from column '{q_col}' (assumed m³/s)") + return series + + +def nse(obs, sim): + m = np.isfinite(obs) & np.isfinite(sim) + if m.sum() < 2: + return np.nan + o, s = obs[m], sim[m] + denom = ((o - o.mean()) ** 2).sum() + if denom < 1e-10: + return np.nan + return 1.0 - float(((o - s) ** 2).sum() / denom) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--route-dir', default=DEFAULT_ROUTE_DIR) + parser.add_argument('--out-dir', default=None, + help='Default: --route-dir') + parser.add_argument('--usgs-csv', default=DEFAULT_USGS_CSV) + parser.add_argument('--plot-start', default=str(PLOT_START_DEFAULT)) + parser.add_argument('--plot-end', default=str(PLOT_END_DEFAULT)) + parser.add_argument('--da-name', default='routed_leadtime_da_full.parquet') + parser.add_argument('--ol-name', default='routed_leadtime_openloop_full.parquet') + args = parser.parse_args() + + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + plot_start = pd.Timestamp(args.plot_start) + plot_end = pd.Timestamp(args.plot_end) + + print(f"Loading DA parquet: {args.da_name}") + da_long = load_parquet_long(os.path.join(args.route_dir, args.da_name)) + print(f"Loading OL parquet: {args.ol_name}") + ol_long = load_parquet_long(os.path.join(args.route_dir, args.ol_name)) + print(f"Loading USGS obs: {args.usgs_csv}") + obs = load_usgs(args.usgs_csv) + + print("Reconstructing DA time series (overlapping-leads pool)...") + da_rec = reconstruct(da_long) + print(f" {len(da_rec)} unique target hours") + print("Reconstructing OL time series...") + ol_rec = reconstruct(ol_long) + + # Clip to plot window and join obs + da_rec = da_rec[(da_rec['valid_time'] >= plot_start) & (da_rec['valid_time'] <= plot_end)].copy() + ol_rec = ol_rec[(ol_rec['valid_time'] >= plot_start) & (ol_rec['valid_time'] <= plot_end)].copy() + da_rec['obs'] = da_rec['valid_time'].map(obs) + ol_rec['obs'] = ol_rec['valid_time'].map(obs) + + nse_da = nse(da_rec['obs'].values, da_rec['median'].values) + nse_ol = nse(ol_rec['obs'].values, ol_rec['median'].values) + print(f"NSE (plot window) — DA: {nse_da:.3f} | OL: {nse_ol:.3f}") + + # ---- Plot ---- + fig, ax = plt.subplots(figsize=(16, 5.5)) + fig.patch.set_facecolor("white") + + # Helene shaded + ax.axvspan(HELENE_START, HELENE_END, + color="firebrick", alpha=0.08, zorder=0) + ax.text(HELENE_START + pd.Timedelta(hours=12), 1.0, "Helene", + transform=ax.get_xaxis_transform(), + fontsize=9, color="firebrick", ha="left", va="top", + fontweight="bold") + + # Open-loop band + median + ax.fill_between(ol_rec['valid_time'], ol_rec['p05'], ol_rec['p95'], + color=OL_COLOR, alpha=0.18, zorder=2, + label="Open-loop [5th–95th pctile envelope]") + ax.plot(ol_rec['valid_time'], ol_rec['median'], + color=OL_COLOR, lw=1.2, linestyle='--', zorder=3, + label=f"Open-loop median (NSE={nse_ol:.3f})") + + # DA band + median + ax.fill_between(da_rec['valid_time'], da_rec['p05'], da_rec['p95'], + color=DA_COLOR, alpha=0.22, zorder=4, + label="DA [5th–95th pctile envelope]") + ax.plot(da_rec['valid_time'], da_rec['median'], + color=DA_COLOR, lw=1.4, zorder=5, + label=f"DA median (NSE={nse_da:.3f})") + + # Observed as dots + obs_window = obs.loc[plot_start:plot_end] + finite = np.isfinite(obs_window.values) + ax.scatter(obs_window.index[finite], obs_window.values[finite], + s=6, color=OBS_COLOR, linewidths=0, zorder=6, + label="USGS obs") + + ax.set_xlabel("Date", fontsize=11) + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_title( + f"Reconstructed forecast time series at USGS 03463300\n" + f"Overlapping-leads pool from EnKF forecast ensemble " + f"({plot_start.date()} – {plot_end.date()})", + fontsize=12, + ) + ax.set_xlim(plot_start, plot_end) + ax.xaxis.set_major_locator(mdates.WeekdayLocator(interval=1)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.25, lw=0.4) + ax.legend(loc='upper left', fontsize=9, frameon=True, framealpha=0.92) + + plt.tight_layout() + out_path = os.path.join(out_dir, "lead_time_reconstructed_timeseries.png") + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + # Also produce a Helene-zoom version + helene_zoom_start = pd.Timestamp("2024-09-24 00:00:00") + helene_zoom_end = pd.Timestamp("2024-09-29 23:00:00") + da_z = da_rec[(da_rec['valid_time'] >= helene_zoom_start) + & (da_rec['valid_time'] <= helene_zoom_end)] + ol_z = ol_rec[(ol_rec['valid_time'] >= helene_zoom_start) + & (ol_rec['valid_time'] <= helene_zoom_end)] + nse_da_z = nse(da_z['obs'].values, da_z['median'].values) + nse_ol_z = nse(ol_z['obs'].values, ol_z['median'].values) + print(f"NSE (Helene zoom) — DA: {nse_da_z:.3f} | OL: {nse_ol_z:.3f}") + + fig, ax = plt.subplots(figsize=(14, 5.5)) + fig.patch.set_facecolor("white") + ax.fill_between(ol_z['valid_time'], ol_z['p05'], ol_z['p95'], + color=OL_COLOR, alpha=0.18, zorder=2, + label="Open-loop [5th–95th]") + ax.plot(ol_z['valid_time'], ol_z['median'], color=OL_COLOR, + lw=1.4, linestyle='--', zorder=3, + label=f"Open-loop median (NSE={nse_ol_z:.3f})") + ax.fill_between(da_z['valid_time'], da_z['p05'], da_z['p95'], + color=DA_COLOR, alpha=0.22, zorder=4, + label="DA [5th–95th]") + ax.plot(da_z['valid_time'], da_z['median'], color=DA_COLOR, + lw=1.6, zorder=5, + label=f"DA median (NSE={nse_da_z:.3f})") + obs_z = obs.loc[helene_zoom_start:helene_zoom_end] + finite = np.isfinite(obs_z.values) + ax.scatter(obs_z.index[finite], obs_z.values[finite], + s=12, color=OBS_COLOR, linewidths=0, zorder=6, + label="USGS obs") + ax.set_xlabel("Date", fontsize=11) + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_title( + "Reconstructed forecast time series at USGS 03463300 — Helene zoom\n" + "Sep 24 – 29, 2024 (overlapping-leads pool from EnKF forecast ensemble)", + fontsize=12, + ) + ax.xaxis.set_major_locator(mdates.HourLocator(interval=12)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %H:%M")) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.25, lw=0.4) + ax.legend(loc='upper left', fontsize=9, frameon=True, framealpha=0.92) + plt.tight_layout() + out_zoom = os.path.join(out_dir, "lead_time_reconstructed_timeseries_helene.png") + plt.savefig(out_zoom, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_zoom}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4c_timeseries/run_4c_f2.sh b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4c_timeseries/run_4c_f2.sh new file mode 100644 index 00000000..aa58a291 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/4_evaluation/4c_timeseries/run_4c_f2.sh @@ -0,0 +1,51 @@ +#!/usr/bin/env bash +# F2 fixed R=0.07 (20pct gauge holdout) — 4c reconstructed timeseries plots. +# +# Reads routed_leadtime_da_full.parquet + routed_leadtime_openloop_full.parquet +# from the F2 leadtime route dir and writes output PNGs there. +# Must run AFTER route_leadtime_f2.sh. +# +# Usage: +# bash run_4c_f2.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +LEADTIME_DIR=/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast +ROUTE_DIR=/mnt/disk2/suma_helen_poster/leadtime_troute_routing_f2 +OBS_DIR=/mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +echo "[F2-4c] Reconstructed timeseries from: $ROUTE_DIR" +$TROUTE "$SCRIPT_DIR/plot_lead_time_reconstructed_timeseries.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" + +echo "[F2-4c] Forecast spaghetti..." +$TROUTE "$SCRIPT_DIR/plot_forecast_spaghetti.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" + +echo "[F2-4c] Helene vs low-flow issue-time hydrograph (per catchment)..." +for CAT in "${CATS[@]}"; do + DA_CSV="$LEADTIME_DIR/$CAT/${CAT}_lead_time_forecasts_da.csv" + [ -f "$DA_CSV" ] || { echo " [$CAT] no leadtime CSV — skipping"; continue; } + echo " [$CAT] helene hydrograph..." + $TROUTE "$SCRIPT_DIR/plot_helene_issue_time_hydrograph.py" \ + --cat-id "$CAT" \ + --leadtime-dir "$LEADTIME_DIR" \ + --obs-dir "$OBS_DIR" +done + +echo "[F2-4c] Done. Outputs in: $ROUTE_DIR and $LEADTIME_DIR" +ls "$ROUTE_DIR"/*.png 2>/dev/null || echo " (no PNGs in route dir)" diff --git a/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/README.md b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/README.md new file mode 100644 index 00000000..fd32e569 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder2_fixed_r_007/README.md @@ -0,0 +1,45 @@ +# Folder 2 — Fixed R = 0.07 + +## R Formula +``` +R = 0.07 (constant, all timesteps) +``` +Constant observation error variance — no flow-magnitude or kriging-variance dependence. + +## Pipeline + +### 1. Calibrate +Shared best_params.json from external `calibrate-cfe` repo (same as all folders). + +### 2. Assimilation +| Sub-step | Script | Server output | +|---|---|---| +| 2a. Forcing arm (30 members, R=0.07) | `run_perturbation_da_on.py --hardcoded-r 0.07` | `v2_perturbation_da_on/` | +| 2b. Hydro-state arm (20 members, R=0.07) | `run_perturbation_da_on.py --hardcoded-r 0.07` | `v2_perturbation_da_on/` | +| 18hr forecast cycles | `run_lead_time_forecast_sweep.py --hardcoded-r 0.07` | `v2_lead_time_forecast_hardcoded_r/` | +| 600-member crossed ensemble | `batch_run_crossed_all_cats.sh` (--hardcoded-r 0.07) | `v2_crossed_ensemble/` | + +> Note: 2a/2b arm runs default to R=0.07 — same figures as Folder 1 arms. + +### 3. Route +| Script | Input | Output | Status | +|---|---|---|---| +| `run_route.py` | `v2_fixed_r007_analysis/` | `fixed_r007_routed/routed_Q_test.csv` | ✅ | +| `route_lead_time_forecasts.py` | `v2_lead_time_forecast_hardcoded_r/` | `v2_lead_time_forecast_hardcoded_r_routed/` | ✅ | +| `run_route_crossed_ensemble.py` | `v2_crossed_ensemble/` | `v2_crossed_ensemble_routed/` | ✅ | + +### 4. Evaluate +| Sub-step | Status | +|---|---| +| 4a. Error decay | ✅ `v2_lead_time_forecast_hardcoded_r_routed/error_fixed_target_mean.png` | +| 4b. 600-member ensemble vs obs | ✅ `v2_crossed_ensemble_routed/` | +| 4c. Ensemble mean + spread | ✅ `v2_lead_time_forecast_hardcoded_r_routed/lead_time_reconstructed_timeseries*.png` | + +## Key Results (deterministic analysis routing) +- Full period KGE: **+0.503** NSE: **+0.163** +- Helene KGE: **+0.439** NSE: **-0.009** +- Helene peak: 1227.1 m³/s (65% of USGS 1885.7 m³/s) + +## Data Source +`/mnt/disk2/suma_helen_poster/da_results/v2_fixed_r007_analysis/` (analysis trajectory) +`/mnt/disk2/suma_helen_poster/da_results/fixed_r007_routed/` (routed output) diff --git a/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/2_assimilation/batch_run_all_f3.sh b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/2_assimilation/batch_run_all_f3.sh new file mode 100644 index 00000000..c663f0c9 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/2_assimilation/batch_run_all_f3.sh @@ -0,0 +1,123 @@ +#!/usr/bin/env bash +# Folder 3 — Dynamic Vrugt seeded: full 2_assimilation pipeline +# +# Runs 3 stages for all 21 catchments: +# Stage A: 2a/2b perturbation arms (run_perturbation_da_on.py, Vrugt R, seed=42) +# Stage B: 2c 18hr forecast cycles (run_lead_time_forecast_sweep.py, Vrugt R, seed=42) +# Stage C: 2d 600-member ensemble (run_crossed_ensemble.py, Vrugt R, seed=42) +# +# Usage: +# bash batch_run_all_f3.sh # runs all 3 stages +# bash batch_run_all_f3.sh arms # stage A only +# bash batch_run_all_f3.sh forecast # stage B only +# bash batch_run_all_f3.sh ensemble # stage C only + +set -euo pipefail + +STAGE="${1:-all}" + +SCRIPT_DIR="$(dirname "$0")" +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python + +# Paths +DA_SRC="/mnt/disk2/1400_sites_helene/da_results_dynamic_vrugt_seeded" +OBS_DIR="/mnt/disk2/1400_sites_helene/catchment_ts_03463300_spliced_dyn_helene" +CFE_DIR="/mnt/disk2/suma_helen_poster/cfe_py" +CONFIG_FILE="/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json" +PARAM_BOUNDS="/mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json" +FORCING_DIR="/mnt/disk2/suma_helen_poster/nwm_retro_catchment_forcings" +TEST_FORCING_DIR1="/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings" +TEST_FORCING_DIR2="/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings" + +OUT_ARMS="/mnt/disk2/1400_sites_helene/da_arms_dynamic_vrugt_seeded" +OUT_FORECAST="/mnt/disk2/1400_sites_helene/da_forecast_dynamic_vrugt_seeded" +OUT_ENSEMBLE="/mnt/disk2/1400_sites_helene/da_crossed_dynamic_vrugt_seeded" + +LOG_DIR="$HOME/f3_logs" +mkdir -p "$LOG_DIR" + +CATS=$(ls -d "$DA_SRC"/cat-* 2>/dev/null | xargs -I{} basename {}) + +# ── Stage A: 2a/2b perturbation arms ───────────────────────────────────────── +run_arms() { + echo "[F3] Stage A: perturbation arms (Vrugt R, seed=42)" + for cat_id in $CATS; do + params="$DA_SRC/$cat_id/${cat_id}_best_params.json" + [ -f "$params" ] || { echo " [skip] $cat_id — no params"; continue; } + mkdir -p "$OUT_ARMS/$cat_id" + cp "$params" "$OUT_ARMS/$cat_id/" + log="$LOG_DIR/arms_${cat_id}.log" + echo " [run] $cat_id" + $TROUTE "$SCRIPT_DIR/run_perturbation_da_on.py" \ + --cat-id "$cat_id" \ + --forcing-dir "$FORCING_DIR" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --param-bounds "$PARAM_BOUNDS" \ + --out-dir "$OUT_ARMS" \ + --test-forcing-dir1 "$TEST_FORCING_DIR1" \ + --test-forcing-dir2 "$TEST_FORCING_DIR2" \ + --rng-seed 42 \ + > "$log" 2>&1 && echo " [ok ] $cat_id" || echo " [FAIL] $cat_id — see $log" + done +} + +# ── Stage B: 2c 18hr forecast cycles ───────────────────────────────────────── +run_forecast() { + echo "[F3] Stage B: 18hr forecast cycles (Vrugt R, seed=42)" + for cat_id in $CATS; do + params="$DA_SRC/$cat_id/${cat_id}_best_params.json" + [ -f "$params" ] || { echo " [skip] $cat_id — no params"; continue; } + mkdir -p "$OUT_FORECAST/$cat_id" + cp "$params" "$OUT_FORECAST/$cat_id/" + log="$LOG_DIR/forecast_${cat_id}.log" + echo " [run] $cat_id" + $TROUTE "$SCRIPT_DIR/run_lead_time_forecast_sweep.py" \ + --cat-id "$cat_id" \ + --forcing-dir "$FORCING_DIR" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --param-bounds "$PARAM_BOUNDS" \ + --out-dir "$OUT_FORECAST" \ + --test-forcing-dir1 "$TEST_FORCING_DIR1" \ + --test-forcing-dir2 "$TEST_FORCING_DIR2" \ + --rng-seed 42 \ + > "$log" 2>&1 && echo " [ok ] $cat_id" || echo " [FAIL] $cat_id — see $log" + done +} + +# ── Stage C: 2d 600-member crossed ensemble ─────────────────────────────────── +run_ensemble() { + echo "[F3] Stage C: 600-member crossed ensemble (Vrugt R, seed=42)" + for cat_id in $CATS; do + params="$DA_SRC/$cat_id/${cat_id}_best_params.json" + [ -f "$params" ] || { echo " [skip] $cat_id — no params"; continue; } + mkdir -p "$OUT_ENSEMBLE/$cat_id" + cp "$params" "$OUT_ENSEMBLE/$cat_id/" + log="$LOG_DIR/ensemble_${cat_id}.log" + echo " [run] $cat_id" + $TROUTE "$SCRIPT_DIR/run_crossed_ensemble.py" \ + --cat-id "$cat_id" \ + --forcing-dir "$FORCING_DIR" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --out-dir "$OUT_ENSEMBLE" \ + --test-forcing-dir1 "$TEST_FORCING_DIR1" \ + --test-forcing-dir2 "$TEST_FORCING_DIR2" \ + --rng-seed 42 \ + > "$log" 2>&1 && echo " [ok ] $cat_id" || echo " [FAIL] $cat_id — see $log" + done +} + +case "$STAGE" in + arms) run_arms ;; + forecast) run_forecast ;; + ensemble) run_ensemble ;; + all) run_arms; run_forecast; run_ensemble ;; + *) echo "Usage: $0 [arms|forecast|ensemble|all]"; exit 1 ;; +esac + +echo "[F3] Done." diff --git a/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/2_assimilation/input_enkf_new.json b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/2_assimilation/input_enkf_new.json new file mode 100644 index 00000000..482233bc --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/2_assimilation/input_enkf_new.json @@ -0,0 +1 @@ +{"n": 1, "m": 1, "R": 0.07, "Q": 0, "smcmax" : 0.9394097311639178, "N": 1000, "P": 0.01, "D":2} \ No newline at end of file diff --git a/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/2_assimilation/new_EnKF.py b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/2_assimilation/new_EnKF.py new file mode 100644 index 00000000..ff945c73 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/2_assimilation/new_EnKF.py @@ -0,0 +1,53 @@ +# EnKF.py +import numpy as np + +class EnKF: + def __init__(self, n, m, R, Q, N, P, smcmax, D): + self.n = n # number of states + self.m = m # number of measurements + self.R = R # measurement noise + self.Q = Q # process noise + self.N = N # number of ensembles + self.P = P + self.state_estimates = np.zeros((self.n)) + self.covariance_matrices = np.zeros((self.n, self.n)) + self.current_step = 0 + self.smcmax = smcmax + self.D = D + self.storage_max_m = self.smcmax * self.D + self.storage_init = self.storage_max_m * 0.667 + self.ensembles = np.full((self.n, self.N), self.storage_init) + + def predict(self, F_results): + error_factor_sim = 0.005 + for i in range(self.N): + perturbation_factor_sim = np.random.standard_normal() + self.ensembles[:, i] = F_results[i] + self.ensembles[:, i] += (perturbation_factor_sim * F_results[i] * error_factor_sim) + return self.ensembles + + def update(self, yi, H_results): + if np.any(np.isnan(yi)): + return np.mean(self.ensembles, axis=1) # return current state estimate without updating + + y_ensembles = np.zeros((self.m, self.N)) + error_factor_sm = 0.03 + for i in range(self.N): + perturbation_factor_sm = np.random.standard_normal() + y_ensembles[:, i] = H_results[i] + y_ensembles[:, i] += (perturbation_factor_sm * H_results[i] * error_factor_sm) + + ensemble_mean = np.mean(self.ensembles, axis=1)[:, np.newaxis] + y_ensembles_mean = np.mean(y_ensembles, axis=1)[:, np.newaxis] + Pxy = (self.ensembles - ensemble_mean).dot((y_ensembles - y_ensembles_mean).T) / (self.N - 1) + Pyy = np.cov(y_ensembles, bias=True) + K = Pxy.dot(np.linalg.pinv(Pyy + self.R)) + self.ensembles += K.dot(yi - y_ensembles_mean) + return np.mean(self.ensembles, axis=1) + + def get_state_estimate(self): + return np.mean(self.ensembles, axis=1) + + def get_covariance(self): + ensemble_mean = np.mean(self.ensembles, axis=1)[:, np.newaxis] + return (self.ensembles - ensemble_mean).dot((self.ensembles - ensemble_mean).T) / (self.N - 1) diff --git a/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/2_assimilation/run_crossed_ensemble.py b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/2_assimilation/run_crossed_ensemble.py new file mode 100644 index 00000000..bc1cec20 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/2_assimilation/run_crossed_ensemble.py @@ -0,0 +1,463 @@ +""" +run_crossed_ensemble.py -- 4b: 600-member crossed ensemble for Helene. + +Crosses the two DA-on perturbation arms: + Forcing arm : 30 met draws (from run_perturbation_da_on.py Phase 2A) + Hydro arm : 20 state draws (from run_perturbation_da_on.py Phase 2B) + +For each Helene issue time t0: + Run 600 members = 30 forcing draws x 20 hydro-state draws. + Each member (i, j) gets: + - Initial state = perturbed DA analysis state (draw j of hydro arm) + - Met forcing = perturbed precip/PET sequence (draw i of forcing arm) + +This gives the full uncertainty envelope combining both sources — +the operational probabilistic forecast picture. + +Output (per catchment): + //_crossed_ensemble.parquet + Columns: issue_time, lead_hour, member_0000..member_0599 + +Usage: + python3 run_crossed_ensemble.py \\ + --cat-id cat-1016300 \\ + --forcing-dir \\ + --obs-dir /mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance \\ + --cfe-dir /mnt/disk2/suma_helen_poster/cfe_py \\ + --config-file /mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json \\ + --out-dir /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble \\ + --test-forcing-dir1 \\ + --test-forcing-dir2 + +Strategy: + Re-runs DA analysis (same as Phase 1 of run_perturbation_da_on.py) to get + the per-issue-time analysis state. Then generates N_FORCING sets of perturbed + forcing sequences and N_HYDRO perturbed states. Runs all 600 combinations. + + To avoid re-running Phase 1, if the Phase 1 snapshot file exists from a + prior run_perturbation_da_on.py run, it is loaded directly. + +NOTE: All special characters kept ASCII to avoid encoding issues on GPU. +""" + +import argparse +import os +import sys +import json +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +N_FORCING = 30 +N_HYDRO = 20 +N_MEMBERS = N_FORCING * N_HYDRO # 600 + +TEST_START = "2024-08-24 00:00:00" # 1 month spinup before Helene window +TEST_END = "2024-10-31 23:00:00" + +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-30 06:00:00") +FORECAST_HOURS = 18 + +PRECIP_SIGMA = 0.15 +PET_SIGMA = 0.10 +STATE_FRAC = 0.05 + +CAT_ID = None +OBS_FILE = None +CFE_CONFIG_FILE = None +TEST_FORCING_FILE = None +OUT_DIR = None +HARDCODED_R = None +DIRECT_VARIANCE = False # when True: R = krig_var directly (no Vrugt formula) +bmi_cfe = None + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(): + df = pd.read_csv(TEST_FORCING_FILE) + df["date"] = pd.to_datetime(df["time"]).dt.strftime("%Y-%m-%d %H:%M:%S") + df["total_precipitation"] = df["APCP_surface"] * 3600.0 + df["potential_evaporation"] = priestley_taylor_pet( + df["DSWRF_surface"].values, df["TMP_2maboveground"].values) + return df + + +def load_obs(): + df = pd.read_csv(OBS_FILE) + date_col = next((c for c in df.columns + if c.lower() in ("date", "time", "datetime", "timestamp")), None) + q_col = next((c for c in df.columns + if "qkrig" in c.lower() and "var" not in c.lower()), None) + var_col = next((c for c in df.columns if "var" in c.lower()), None) + df[date_col] = pd.to_datetime(df[date_col]) + df = df.sort_values(date_col).reset_index(drop=True) + date_strs = df[date_col].dt.strftime("%Y-%m-%d %H:%M:%S") + obs_dict = dict(zip(date_strs, df[q_col])) + var_vals = df[var_col] if var_col else pd.Series([0.07] * len(df)) + var_dict = dict(zip(date_strs, var_vals)) + return obs_dict, var_dict + + +def write_model_config(best_params): + """Write the BMI config file once; reused for every model instantiation.""" + with open(CFE_CONFIG_FILE) as f: + cfg = json.load(f) + cfg["forcing_file"] = TEST_FORCING_FILE + cfg["soil_params"]["bb"] = best_params["bb"] + cfg["soil_params"]["smcmax"] = best_params["smcmax"] + cfg["soil_params"]["satdk"] = best_params["satdk"] + cfg["slop"] = best_params["slop"] + cfg["max_gw_storage"] = best_params["max_gw_storage"] + cfg["expon"] = best_params["expon"] + cfg["Cgw"] = best_params["Cgw"] + cfg["K_lf"] = best_params["K_lf"] + cfg["K_nash"] = best_params["K_nash"] + cfg["partition_scheme"] = ("Schaake" if best_params["scheme"] <= 0.5 + else "Xinanjiang") + tmp_cfg = str(OUT_DIR / f"{CAT_ID}_bmi_config_crossed_stable.json") + with open(tmp_cfg, "w") as f: + json.dump(cfg, f) + return tmp_cfg + + +def make_model(cfg_path): + """Instantiate and initialize a CFE model from a pre-written config path.""" + m = bmi_cfe.BMI_CFE(cfg_file=cfg_path) + m.initialize() + return m + + +def get_state(m): + return { + "soil_m": m.soil_reservoir["storage_m"], + "gw_m": m.gw_reservoir["storage_m"], + "nash0": float(m.nash_storage[0]), + "nash1": float(m.nash_storage[1]), + } + + +def set_state(m, state): + m.soil_reservoir["storage_m"] = state["soil_m"] + m.gw_reservoir["storage_m"] = state["gw_m"] + m.nash_storage[0] = state["nash0"] + m.nash_storage[1] = state["nash1"] + + +def perturb_state(state, rng, frac=STATE_FRAC): + return { + "soil_m": max(state["soil_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "nash0": max(state["nash0"] * (1.0 + frac * rng.standard_normal()), 0.0), + "nash1": max(state["nash1"] * (1.0 + frac * rng.standard_normal()), 0.0), + } + + +def run_model_step(m, precip_mmh, pet_mmh): + m.set_value("atmosphere_water__time_integral_of_precipitation_mass_flux", + precip_mmh / 1000.0) + m.set_value("water_potential_evaporation_flux", + pet_mmh / 1000.0 / 3600.0) + m.update() + return m.get_value("land_surface_water__runoff_depth") * 1000.0 + + +def enkf_update(state, q_sim, y_obs, krig_var, rng): + if HARDCODED_R is not None: + R = HARDCODED_R + elif DIRECT_VARIANCE: + R = max(krig_var, 1e-6) + else: + R = max((0.10 * max(y_obs, 0.0)) ** 2 + 0.001 * krig_var, 1e-6) + R = max(R, 1e-6) + P_yy = max(q_sim * 0.01, 1e-6) + K = P_yy / (P_yy + R) + scale = K * (y_obs - q_sim) / max(abs(q_sim), 1e-6) + return { + "soil_m": max(state["soil_m"] * (1.0 + scale), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + scale), 1e-6), + "nash0": max(state["nash0"] + state["nash0"] * scale, 0.0), + "nash1": max(state["nash1"] + state["nash1"] * scale, 0.0), + } + + +def run_phase1_da(cfg_path, df_test, obs_dict, var_dict, rng): + """Run DA through test period; return snapshot dict {timestamp -> state}.""" + print(f" Phase 1: DA analysis through test period...") + model = make_model(cfg_path) + snapshots = {} + for _, row in df_test.iterrows(): + date_str = row["date"] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + q_sim = run_model_step(model, P, E) + y_obs = obs_dict.get(date_str, np.nan) + if not np.isnan(y_obs): + krig_var = var_dict.get(date_str, 0.07) + state = get_state(model) + updated = enkf_update(state, q_sim, y_obs, krig_var, rng) + set_state(model, updated) + date_ts = pd.Timestamp(date_str) + if HELENE_START <= date_ts <= HELENE_END: + snapshots[date_ts] = get_state(model) + model.finalize() + print(f" Phase 1 done: {len(snapshots)} issue-time snapshots") + return snapshots + + +def main(): + global CAT_ID, OBS_FILE, CFE_CONFIG_FILE, TEST_FORCING_FILE + global OUT_DIR, HARDCODED_R, DIRECT_VARIANCE, bmi_cfe + + parser = argparse.ArgumentParser() + parser.add_argument("--cat-id", required=True) + parser.add_argument("--forcing-dir", required=True) + parser.add_argument("--obs-dir", required=True) + parser.add_argument("--cfe-dir", required=True) + parser.add_argument("--config-file", required=True) + parser.add_argument("--out-dir", required=True) + parser.add_argument("--test-forcing-dir1", default=None) + parser.add_argument("--test-forcing-dir2", default=None) + parser.add_argument("--hardcoded-r", type=float, default=None, + help="Fixed R value for EnKF (e.g. 0.07). " + "Omit to use dynamic Vrugt R: R=(0.10*y)^2 + 0.001*krig_var") + parser.add_argument("--direct-variance", action="store_true", default=False, + help="Use krig_var directly as R (no Vrugt scaling)") + parser.add_argument("--rng-seed", type=int, default=None, + help="RNG seed for reproducibility (default: hash of cat-id)") + parser.add_argument("--n-forcing", type=int, default=N_FORCING, + help="Number of met forcing draws (default 30)") + parser.add_argument("--n-hydro", type=int, default=N_HYDRO, + help="Number of hydro-state draws (default 20)") + args = parser.parse_args() + + CAT_ID = args.cat_id + OBS_FILE = os.path.join(args.obs_dir, f"{CAT_ID}.csv") + CFE_CONFIG_FILE = args.config_file + HARDCODED_R = args.hardcoded_r + DIRECT_VARIANCE = args.direct_variance + OUT_DIR = Path(args.out_dir) / CAT_ID + OUT_DIR.mkdir(parents=True, exist_ok=True) + + n_fa = args.n_forcing + n_ha = args.n_hydro + n_total = n_fa * n_ha + print(f"[{CAT_ID}] Crossed ensemble: {n_fa} forcing x {n_ha} hydro = {n_total} members") + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + bmi_cfe = _bmi_cfe + + # Reuse combined forcing CSV from run_perturbation_da_on.py if it exists + # (avoids writing a duplicate large file to disk) + da_on_combined = (Path(args.out_dir).parent / "v2_perturbation_da_on" / + CAT_ID / f"{CAT_ID}_nwm_operational_combined.csv") + crossed_combined = OUT_DIR / f"{CAT_ID}_nwm_operational_combined.csv" + + if da_on_combined.exists(): + TEST_FORCING_FILE = str(da_on_combined) + print(f" Reusing combined forcing: {TEST_FORCING_FILE}") + elif crossed_combined.exists(): + TEST_FORCING_FILE = str(crossed_combined) + print(f" Reusing combined forcing: {TEST_FORCING_FILE}") + elif args.test_forcing_dir1 and args.test_forcing_dir2: + f1 = os.path.join(args.test_forcing_dir1, f"{CAT_ID}.csv") + f2 = os.path.join(args.test_forcing_dir2, f"{CAT_ID}.csv") + df1, df2 = pd.read_csv(f1), pd.read_csv(f2) + combined = (pd.concat([df1, df2], ignore_index=True) + .drop_duplicates(subset="time").sort_values("time")) + TEST_FORCING_FILE = str(crossed_combined) + combined.to_csv(TEST_FORCING_FILE, index=False) + print(f" Wrote combined forcing: {TEST_FORCING_FILE}") + else: + print("Need --test-forcing-dir1 and --test-forcing-dir2"); return + + best_params_file = OUT_DIR / f"{CAT_ID}_best_params.json" + if not best_params_file.exists(): + # Also check the DA results dir from run_perturbation_da_on.py + alt = Path(args.out_dir).parent / "v2_perturbation_da_on" / CAT_ID / f"{CAT_ID}_best_params.json" + if alt.exists(): + best_params_file = alt + else: + print(f"Missing best params at {best_params_file}"); return + with open(best_params_file) as f: + best_params = json.load(f)["best_parameters"] + + obs_dict, var_dict = load_obs() + df_forcing = load_test_forcing() + t_mask = (df_forcing["date"] >= TEST_START) & (df_forcing["date"] <= TEST_END) + df_test = df_forcing[t_mask].reset_index(drop=True) + + # Write BMI config once — reused for all 90,600 model instantiations + cfg_path = write_model_config(best_params) + print(f" BMI config written once: {cfg_path}") + + rng = np.random.default_rng(args.rng_seed if args.rng_seed is not None + else hash(CAT_ID) & 0x7fffffff) + + # Phase 1: get DA analysis snapshots + # Check if a prior run already saved them (skip re-running DA if so) + snap_cache = OUT_DIR / f"{CAT_ID}_da_snapshots.parquet" + if snap_cache.exists(): + print(f" Loading cached DA snapshots from {snap_cache}") + df_snaps = pd.read_parquet(snap_cache) + snapshots = { + pd.Timestamp(row["issue_time"]): { + "soil_m": row["soil_m"], "gw_m": row["gw_m"], + "nash0": row["nash0"], "nash1": row["nash1"], + } + for _, row in df_snaps.iterrows() + } + else: + snapshots = run_phase1_da(cfg_path, df_test, obs_dict, var_dict, rng) + + issue_times = sorted(snapshots.keys()) + t0_to_idx = {t: i for i, t in enumerate(df_test["date"].apply(pd.Timestamp))} + + # ------------------------------------------------------------------ # + # Crossed ensemble: for each issue_time, build n_fa forcing draws and # + # n_ha state draws, then run all n_fa x n_ha combinations. # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Running crossed ensemble ({n_total} members x " + f"{len(issue_times)} issue times)...") + + all_records = [] + all_hydro_draw_records = [] # (issue_time, hydro_draw_j, states) + all_forcing_draw_records = [] # (issue_time, forcing_draw_i, lead, P_pert, E_pert, scales) + + for t_idx, t0 in enumerate(issue_times): + snap = snapshots[t0] + t0_str = t0.strftime("%Y-%m-%d %H:%M:%S") + if t0 not in t0_to_idx: + continue + start_idx = t0_to_idx[t0] + + # Pre-draw n_fa forcing perturbation sequences + forcing_seqs = [] + for i in range(n_fa): + seq = [] + for lead in range(1, FORECAST_HOURS + 1): + fi = start_idx + lead + if fi >= len(df_test): + seq.append((0.0, 0.0)) + all_forcing_draw_records.append({ + "issue_time": t0_str, "forcing_draw_i": i, + "lead_hour": lead, + "P_pert_mmh": 0.0, "E_pert_mmh": 0.0, + "precip_scale": np.nan, "pet_scale": np.nan, + }) + continue + row = df_test.iloc[fi] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + mu_p = -0.5 * PRECIP_SIGMA ** 2 + P_pert = float(P * rng.lognormal(mu_p, PRECIP_SIGMA)) if P > 0 else 0.0 + E_pert = max(float(E * (1.0 + PET_SIGMA * rng.standard_normal())), 0.0) + seq.append((P_pert, E_pert)) + all_forcing_draw_records.append({ + "issue_time": t0_str, "forcing_draw_i": i, + "lead_hour": lead, + "P_pert_mmh": P_pert, "E_pert_mmh": E_pert, + "precip_scale": P_pert / P if P > 0 else np.nan, + "pet_scale": E_pert / E if E > 0 else np.nan, + }) + forcing_seqs.append(seq) + + # Pre-draw n_ha hydro-state perturbations + state_draws = [snap if j == 0 else perturb_state(snap, rng) + for j in range(n_ha)] + + # Record each hydro draw's actual initial state + for j, st in enumerate(state_draws): + all_hydro_draw_records.append({ + "issue_time": t0_str, "hydro_draw_j": j, + "soil_m": st["soil_m"], "gw_m": st["gw_m"], + "nash0": st["nash0"], "nash1": st["nash1"], + }) + + # q_matrix shape: (FORECAST_HOURS, n_fa, n_ha) + q_matrix = np.full((FORECAST_HOURS, n_fa, n_ha), np.nan) + + for i, f_seq in enumerate(forcing_seqs): + for j, state_j in enumerate(state_draws): + m = make_model(cfg_path) + set_state(m, state_j) + for lead in range(1, FORECAST_HOURS + 1): + P_pert, E_pert = f_seq[lead - 1] + q = run_model_step(m, P_pert, E_pert) + q_matrix[lead - 1, i, j] = q + m.finalize() + + # Flatten 600-member matrix into rows + for lead in range(1, FORECAST_HOURS + 1): + row_data = {"issue_time": t0_str, "lead_hour": lead} + mem_idx = 0 + for i in range(n_fa): + for j in range(n_ha): + row_data[f"member_{mem_idx:04d}"] = q_matrix[lead - 1, i, j] + mem_idx += 1 + all_records.append(row_data) + + if (t_idx + 1) % 10 == 0 or (t_idx + 1) == len(issue_times): + print(f" {t_idx + 1}/{len(issue_times)} issue times complete") + + df_out = pd.DataFrame(all_records) + out_path = OUT_DIR / f"{CAT_ID}_crossed_ensemble.parquet" + df_out.to_parquet(out_path, index=False) + print(f"Saved: {out_path}") + print(f" Shape: {df_out.shape} " + f"({df_out['issue_time'].nunique()} issue times x " + f"{FORECAST_HOURS} leads x {n_total} members)") + + # ------------------------------------------------------------------ # + # Provenance files — enable full ensemble traceability and restart # + # ------------------------------------------------------------------ # + + # 1. DA analysis snapshots — unperturbed state at each issue_time. + # Load any row and set_state() to restart from that init time. + snap_records = [ + {"issue_time": t.strftime("%Y-%m-%d %H:%M:%S"), **s} + for t, s in snapshots.items() + ] + snap_path = OUT_DIR / f"{CAT_ID}_da_snapshots.parquet" + pd.DataFrame(snap_records).to_parquet(snap_path, index=False) + print(f"Saved: {snap_path} ({len(snap_records)} snapshots)") + + # 2. Member manifest — decoder ring: member_col -> (forcing_draw_i, hydro_draw_j). + # member_k uses forcing draw k//n_ha and hydro draw k%n_ha. + manifest = [ + {"member": f"member_{k:04d}", "member_idx": k, + "forcing_draw_i": k // n_ha, "hydro_draw_j": k % n_ha} + for k in range(n_total) + ] + manifest_path = OUT_DIR / f"{CAT_ID}_member_manifest.csv" + pd.DataFrame(manifest).to_csv(manifest_path, index=False) + print(f"Saved: {manifest_path} ({n_total} members)") + + # 3. Hydro draw states — actual perturbed initial state for each of + # n_ha draws at every issue_time. Links member -> initial conditions. + hydro_path = OUT_DIR / f"{CAT_ID}_hydro_draw_states.parquet" + pd.DataFrame(all_hydro_draw_records).to_parquet(hydro_path, index=False) + print(f"Saved: {hydro_path} ({len(all_hydro_draw_records)} rows)") + + # 4. Forcing draw sequences — actual P/PET perturbation values for each + # of n_fa draws at every (issue_time, lead_hour). Links member -> forcing. + forcing_path = OUT_DIR / f"{CAT_ID}_forcing_draw_sequences.parquet" + pd.DataFrame(all_forcing_draw_records).to_parquet(forcing_path, index=False) + print(f"Saved: {forcing_path} ({len(all_forcing_draw_records)} rows)") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/2_assimilation/run_lead_time_forecast_sweep.py b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/2_assimilation/run_lead_time_forecast_sweep.py new file mode 100644 index 00000000..19f3b4d5 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/2_assimilation/run_lead_time_forecast_sweep.py @@ -0,0 +1,508 @@ +""" +Forecast lead-time evaluation. + +For each issue-time t0 in a sampling schedule across the test period, fork the +ensemble at t0 and free-run 18 hours of forecast. Two parallel trajectories are +maintained throughout the test period: + + da: Production setup. DA on, forcing perturbed, process noise on. + openloop: No DA from t=0. Forcing perturbed, process noise on. + +At each issue time t0, the state of each scenario's ensemble is copied into a +forecast ensemble that free-runs 18 hours with: + - DA off (no observations consumed during the forecast window) + - process noise off (no obs to collapse toward, so no anti-collapse needed) + - forcing perturbed (lognormal precip, Gaussian PET — proxy for forecast + precip uncertainty in lieu of actual HEFS forecasts) + +Issue-time schedule: + - Base cadence: every --base-step-h hours across the full test period + (default 6h → ~1600 issue times for the year) + - Densified to hourly across --dense-start..--dense-end if provided + (default: 2024-09-24 → 2024-09-28, the Helene window) + +Outputs (stacked across all issue times): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv + Columns: issue_time, lead_hour, valid_time, member_00 .. member_19 + (units: mm/h) + +Post-processing pipeline (separate scripts, run after this): + route_lead_time_forecasts.py — pushes each (t0, lead_hour) forecast through + T-route to the gauge + plot_lead_time_decay.py — error-vs-lead-time curve, DA vs open-loop + +Imports EnKFAssimilator from calibrate_catchment_cfe_da_v2.py so the DA math +matches production exactly. +""" + +import argparse +import os +import sys +import json +import importlib.util +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +SPINUP_START = "2023-02-01 00:00:00" +SPINUP_END = "2023-09-30 23:00:00" +TEST_START = "2023-10-01 00:00:00" +TEST_END = "2024-10-31 23:00:00" + +# Forecast lead time (hours after each issue time) +FORECAST_LEAD_HOURS = 18 + +# Default densification window — Hurricane Helene +DEFAULT_DENSE_START = "2024-09-24 00:00:00" +DEFAULT_DENSE_END = "2024-09-28 23:00:00" + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(test_forcing_file): + df = pd.read_csv(test_forcing_file) + df['date'] = pd.to_datetime(df['time']).dt.strftime('%Y-%m-%d %H:%M:%S') + df['total_precipitation'] = df['APCP_surface'] * 3600.0 + df['potential_evaporation'] = priestley_taylor_pet( + df['DSWRF_surface'].values, df['TMP_2maboveground'].values + ) + return df + + +def import_enkf_class(script_path): + """Load EnKFAssimilator from the production script without triggering its main().""" + spec = importlib.util.spec_from_file_location("prod_v2", script_path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod.EnKFAssimilator + + +# ---------------- State snapshot / restore ----------------------------------- +# +# The EnKF only ever touches these 4 states, so they're the only states that +# diverge between members. Copy these from src → dst to fork the ensemble. + +def snapshot_states(models): + return [ + { + "soil_m": float(m.soil_reservoir["storage_m"]), + "gw_m": float(m.gw_reservoir["storage_m"]), + "nash0_m": float(m.nash_storage[0]), + "nash1_m": float(m.nash_storage[1]), + } + for m in models + ] + + +def restore_states(models, snapshot): + for m, s in zip(models, snapshot): + m.soil_reservoir["storage_m"] = s["soil_m"] + m.gw_reservoir["storage_m"] = s["gw_m"] + m.nash_storage[0] = s["nash0_m"] + m.nash_storage[1] = s["nash1_m"] + + +# ---------------- Issue-time schedule ---------------------------------------- + +def build_issue_time_schedule(dates_list, base_step_h, dense_start, dense_end): + """Return a sorted list of issue-time strings sampled from dates_list. + + Every base_step_h hours across the full test period, plus every hour inside + [dense_start, dense_end] if those are provided. + """ + dates_dt = pd.to_datetime(dates_list) + selected = set() + + # Base cadence (every base_step_h hours from the first date) + base_mask = (np.arange(len(dates_dt)) % base_step_h == 0) + for d in dates_dt[base_mask]: + selected.add(d.strftime('%Y-%m-%d %H:%M:%S')) + + # Densified hourly window + if dense_start is not None and dense_end is not None: + ds = pd.Timestamp(dense_start) + de = pd.Timestamp(dense_end) + dense_mask = (dates_dt >= ds) & (dates_dt <= de) + for d in dates_dt[dense_mask]: + selected.add(d.strftime('%Y-%m-%d %H:%M:%S')) + + return sorted(selected) + + +# ---------------- CFE helpers ------------------------------------------------ + +def build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, enkf, + apply_init_perturbation=True, member_zero_clean=True): + """Build N CFE BMI instances with optional init-state perturbation.""" + def custom_load_forcing(self_cfe): + df = load_test_forcing(test_forcing_file) + self_cfe.forcing_data = df.rename(columns={"date": "time"}) + + models = [] + for i in range(N): + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.load_forcing_file = custom_load_forcing.__get__(m) + m.initialize() + if apply_init_perturbation and not (member_zero_clean and i == 0): + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + r = enkf.init_state_perturb_frac + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1 + r * enkf.rng.standard_normal()), 1e-6, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1 + r * enkf.rng.standard_normal()), 1e-6, gw_max)) + jitter = max(sm_max * 1e-6, 1e-9) + m.nash_storage[0] = max(n00 + jitter * enkf.rng.standard_normal(), 0.0) + m.nash_storage[1] = max(n10 + jitter * enkf.rng.standard_normal(), 0.0) + models.append(m) + return models + + +def step_ensemble(models, p_arr, e_arr): + """Advance each member one hour with member-specific (P, PET). Returns Q (mm/h).""" + q = np.empty(len(models), dtype=float) + for i, m in enumerate(models): + m.set_value('atmosphere_water__time_integral_of_precipitation_mass_flux', + float(p_arr[i]) / 1000.0) + m.set_value('water_potential_evaporation_flux', + float(e_arr[i]) / 1000.0 / 3600.0) + m.update() + q[i] = m.get_value('land_surface_water__runoff_depth') * 1000.0 # m/h → mm/h + return q + + +def do_sanity_check(prod_models, fcst_models, dates_list, forcing_by_date, h, N): + """One-shot verification that the 4-state snapshot is sufficient to reproduce + next-hour forecast Q. + + At main-loop hour h (after prod's hour-h step has completed including DA + + process noise), snapshot prod, restore into fcst, then step BOTH ensembles + one hour using identical unperturbed forcing. If snapshot/restore captures + everything that matters, prod's and fcst's hour-(h+1) Q should match member- + by-member to machine precision. + + Note: this mutates prod's state (h+1 step with non-perturbed forcing and no + DA), so the caller should exit the script right after. + """ + if h + 1 >= len(dates_list): + print("[sanity-check] not enough horizon for h+1 step — skipping.") + return False + snap = snapshot_states(prod_models) + restore_states(fcst_models, snap) + + next_date = dates_list[h + 1] + p_next, e_next = forcing_by_date[next_date] + p_uniform = np.full(N, p_next, dtype=float) + e_uniform = np.full(N, e_next, dtype=float) + + q_prod = step_ensemble(prod_models, p_uniform, e_uniform) + q_fcst = step_ensemble(fcst_models, p_uniform, e_uniform) + + diff = q_prod - q_fcst + max_abs = float(np.max(np.abs(diff))) + print("[sanity-check] forecast Q at h+1 with identical unperturbed forcing:") + print(f" prod first 5 members: {q_prod[:5]}") + print(f" fcst first 5 members: {q_fcst[:5]}") + print(f" max |prod - fcst|: {max_abs:.3e} mm/h") + if max_abs < 1e-9: + print(" OK: snapshot/restore preserves state to machine precision.") + return True + print(" WARN: nonzero divergence — likely a state not in the 4-state snapshot.") + return False + + +def run_forecast(fcst_models, enkf_fcst, dates_list, forcing_by_date, + t0_idx, n_lead): + """Run an n_lead-hour free-forecast starting from the current state of fcst_models. + + Forcing is perturbed per member (lognormal precip + Gaussian PET, same as + production). No DA. No process noise. + + Returns: (lead_hours, valid_times, q_matrix shape (n_lead, N)). + """ + N = len(fcst_models) + q_matrix = np.full((n_lead, N), np.nan, dtype=float) + lead_hours = [] + valid_times = [] + for lead in range(1, n_lead + 1): + idx = t0_idx + lead + if idx >= len(dates_list): + break + valid_date = dates_list[idx] + p, e = forcing_by_date[valid_date] + p_arr, e_arr = enkf_fcst.perturb_forcing(p, e) + q = step_ensemble(fcst_models, p_arr, e_arr) + q_matrix[lead - 1, :] = q + lead_hours.append(lead) + valid_times.append(valid_date) + return lead_hours, valid_times, q_matrix + + +# ---------------- Main run --------------------------------------------------- + +def run(args, bmi_cfe, EnKFAssimilator): + cat_id = args.cat_id + out_dir = Path(args.out_dir) / cat_id + out_dir.mkdir(parents=True, exist_ok=True) + + # Combined test forcing (same idiom as the per-member script) + f1 = os.path.join(args.test_forcing_dir1, f'{cat_id}.csv') + f2 = os.path.join(args.test_forcing_dir2, f'{cat_id}.csv') + df1 = pd.read_csv(f1) + df2 = pd.read_csv(f2) + combined = pd.concat([df1, df2], ignore_index=True) + combined = combined.drop_duplicates(subset='time').sort_values('time') + test_forcing_file = str(out_dir / f'{cat_id}_nwm_operational_combined.csv') + combined.to_csv(test_forcing_file, index=False) + + # Pre-staged calibrated params + best_params_file = out_dir / f'{cat_id}_best_params.json' + if not best_params_file.exists(): + print(f"Missing {best_params_file}. Pre-stage from Run 3.") + return + with open(best_params_file) as f: + best = json.load(f)['best_parameters'] + + # Build temp CFE config with calibrated params + with open(args.config_file) as f: + cfg = json.load(f) + cfg['forcing_file'] = test_forcing_file + cfg['soil_params']['bb'] = best['bb'] + cfg['soil_params']['smcmax'] = best['smcmax'] + cfg['soil_params']['satdk'] = best['satdk'] + cfg['slop'] = best['slop'] + cfg['max_gw_storage'] = best['max_gw_storage'] + cfg['expon'] = best['expon'] + cfg['Cgw'] = best['Cgw'] + cfg['K_lf'] = best['K_lf'] + cfg['K_nash'] = best['K_nash'] + cfg['partition_scheme'] = "Schaake" if best['scheme'] <= 0.5 else "Xinanjiang" + tmp_cfg = str(out_dir / f'{cat_id}_bmi_config_temp_leadtime.json') + with open(tmp_cfg, 'w') as f: + json.dump(cfg, f) + + # Three EnKF instances: + # enkf_da — runs the DA-on trajectory; consumes obs each hour + # enkf_ol — runs the open-loop trajectory; never calls update_states + # enkf_fcst — used only inside the 18-hour free-runs (perturb_forcing only) + # All three share defaults with production. Distinct RNG seeds keep them + # independent so the openloop and forecast spreads are not coupled to DA. + obs_file = os.path.join(args.obs_dir, f'{cat_id}.csv') + seed_base = args.rng_seed if args.rng_seed is not None else 0 + enkf_da = EnKFAssimilator( + n_members=args.enkf_members, obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, use_vrugt_r=(not args.no_vrugt_r), + vrugt_alpha=args.vrugt_alpha, vrugt_scale=args.vrugt_scale, + rng_seed=seed_base if args.rng_seed is not None else None, + ) + enkf_ol = EnKFAssimilator( + n_members=args.enkf_members, obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, use_vrugt_r=(not args.no_vrugt_r), + vrugt_alpha=args.vrugt_alpha, vrugt_scale=args.vrugt_scale, + rng_seed=(seed_base + 1) if args.rng_seed is not None else None, + ) + enkf_fcst = EnKFAssimilator( + n_members=args.enkf_members, obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, use_vrugt_r=(not args.no_vrugt_r), + vrugt_alpha=args.vrugt_alpha, vrugt_scale=args.vrugt_scale, + rng_seed=(seed_base + 2) if args.rng_seed is not None else None, + ) + # Optional R override — replace every per-hour obs variance with a constant. + # Matches DualEarth/new_EnKF.py (R=0.07). Set via --hardcoded-r on the CLI. + # Effect: at storm peak ≈ existing Vrugt R, at low flow ≈ 15-17× larger → + # near-zero gain at low flow where kriging obs is noisy. + if args.hardcoded_r is not None: + for enkf in (enkf_da, enkf_ol, enkf_fcst): + for date in enkf.obs_var_dict: + enkf.obs_var_dict[date] = args.hardcoded_r + print(f"[lead-time] R hardcoded to {args.hardcoded_r} mm^2/h^2 " + f"(overrides Vrugt + kriging variance formula)") + + N = enkf_da.n_members + print(f"[lead-time] {cat_id} | N={N} | DA on/off both run | " + f"forecast lead = {FORECAST_LEAD_HOURS}h") + + # Build three ensembles (init-perturbed except the forecast ensemble, whose + # state will be overwritten at each issue time before each free-run) + prod_models = build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, enkf_da) + openloop_models = build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, enkf_ol) + fcst_models = build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, + enkf_fcst, apply_init_perturbation=False) + + # ----- Spinup with perturbed forcing on both trajectories ----- + df = load_test_forcing(test_forcing_file) + sp_mask = (df['date'] >= SPINUP_START) & (df['date'] <= SPINUP_END) + df_sp = df[sp_mask] + print(f"[lead-time] spinup: {len(df_sp)} hours") + for p, e in zip(df_sp['total_precipitation'], df_sp['potential_evaporation']): + p_da, e_da = enkf_da.perturb_forcing(p, e) + step_ensemble(prod_models, p_da, e_da) + enkf_da.add_process_noise(prod_models) + p_ol, e_ol = enkf_ol.perturb_forcing(p, e) + step_ensemble(openloop_models, p_ol, e_ol) + enkf_ol.add_process_noise(openloop_models) + + # ----- Test period: step both, fork forecasts at issue times ----- + t_mask = (df['date'] >= TEST_START) & (df['date'] <= TEST_END) + df_test = df[t_mask].reset_index(drop=True) + dates_list = list(df_test['date'].values) + forcing_by_date = dict(zip( + df_test['date'].values, + zip(df_test['total_precipitation'].values, + df_test['potential_evaporation'].values))) + + issue_times = build_issue_time_schedule( + dates_list, args.base_step_h, + args.dense_start, args.dense_end) + issue_set = set(issue_times) + print(f"[lead-time] test period: {len(dates_list)} hours | " + f"{len(issue_times)} issue times " + f"(base={args.base_step_h}h, dense={args.dense_start or 'none'}" + f"..{args.dense_end or 'none'})") + + # Pre-build the date→index map so forecast windows are fast to look up + date_to_idx = {d: i for i, d in enumerate(dates_list)} + + da_rows = [] # (issue_time, lead_hour, valid_time, *member_values) + ol_rows = [] + sanity_done = False + + for h, current_date in enumerate(dates_list): + p, e = forcing_by_date[current_date] + + # --- DA trajectory step --- + p_da, e_da = enkf_da.perturb_forcing(p, e) + q_da = step_ensemble(prod_models, p_da, e_da) + enkf_da.update_states(prod_models, current_date, q_da) + enkf_da.add_process_noise(prod_models) + + # --- Open-loop trajectory step (no DA) --- + p_ol, e_ol = enkf_ol.perturb_forcing(p, e) + q_ol = step_ensemble(openloop_models, p_ol, e_ol) + enkf_ol.add_process_noise(openloop_models) + + # --- One-shot sanity check at the first scheduled issue time --- + if args.sanity_check and not sanity_done and current_date in issue_set: + ok = do_sanity_check( + prod_models, fcst_models, dates_list, forcing_by_date, h, N) + sanity_done = True + print(f"[sanity-check] exiting (re-run without --sanity-check for " + f"the full sweep). Result: {'OK' if ok else 'FAIL'}") + for m in prod_models + openloop_models + fcst_models: + m.finalize() + sys.exit(0 if ok else 1) + + # --- Fork forecasts at scheduled issue times --- + if current_date in issue_set: + for scenario_label, src_models, rows_acc in [ + ('da', prod_models, da_rows), + ('openloop', openloop_models, ol_rows), + ]: + snap = snapshot_states(src_models) + restore_states(fcst_models, snap) + leads, valids, qm = run_forecast( + fcst_models, enkf_fcst, + dates_list, forcing_by_date, + t0_idx=h, n_lead=FORECAST_LEAD_HOURS) + for k, (lead, vt) in enumerate(zip(leads, valids)): + rows_acc.append((current_date, lead, vt, *qm[k, :])) + + for m in prod_models + openloop_models + fcst_models: + m.finalize() + + # ----- Save forecast CSVs ----- + cols = ['issue_time', 'lead_hour', 'valid_time'] + [f'member_{i:02d}' for i in range(N)] + da_path = out_dir / f'{cat_id}_lead_time_forecasts_da.csv' + ol_path = out_dir / f'{cat_id}_lead_time_forecasts_openloop.csv' + pd.DataFrame(da_rows, columns=cols).to_csv(da_path, index=False) + pd.DataFrame(ol_rows, columns=cols).to_csv(ol_path, index=False) + print(f"[lead-time] saved {da_path}") + print(f"[lead-time] saved {ol_path}") + + # Issue-time schedule (small, useful for downstream scripts) + sched_path = out_dir / f'{cat_id}_lead_time_issue_times.csv' + pd.DataFrame({'issue_time': issue_times}).to_csv(sched_path, index=False) + print(f"[lead-time] saved {sched_path}") + + print(f"[lead-time] DA assimilator: updates={enkf_da.n_updates} | " + f"avg Pyy={enkf_da.avg_pyy:.6f} | " + f"mass lost={enkf_da.total_overflow_lost_mm:.3f} mm") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', required=True) + parser.add_argument('--forcing-dir', required=True) + parser.add_argument('--obs-dir', required=True) + parser.add_argument('--cfe-dir', required=True) + parser.add_argument('--config-file', required=True) + parser.add_argument('--param-bounds', required=True, + help='Kept for symmetry; not read in this script') + parser.add_argument('--out-dir', required=True) + parser.add_argument('--test-forcing-dir1', required=True) + parser.add_argument('--test-forcing-dir2', required=True) + parser.add_argument('--enkf-members', type=int, default=20) + parser.add_argument('--enkf-obs-error-std', type=float, default=0.05) + parser.add_argument('--no-vrugt-r', action='store_true') + parser.add_argument('--vrugt-alpha', type=float, default=0.10) + parser.add_argument('--vrugt-scale', type=float, default=0.001) + parser.add_argument('--rng-seed', type=int, default=None) + parser.add_argument('--base-step-h', type=int, default=6, + help='Issue-time cadence in hours across the full test ' + 'period (default 6h)') + parser.add_argument('--dense-start', type=str, default=DEFAULT_DENSE_START, + help='Start of hourly-densification window ' + '(default 2024-09-24 00:00:00 — Helene)') + parser.add_argument('--dense-end', type=str, default=DEFAULT_DENSE_END, + help='End of hourly-densification window ' + '(default 2024-09-28 23:00:00 — Helene)') + parser.add_argument('--sanity-check', action='store_true', + help='At the first issue time, fork fcst_models and step ' + 'both prod and fcst one hour with identical unperturbed ' + 'forcing. Confirms the 4-state snapshot is sufficient ' + 'to reproduce next-hour Q to machine precision, then ' + 'exits. Use before kicking off the full sweep.') + parser.add_argument('--hardcoded-r', type=float, default=None, + help='If set, override the Vrugt + kriging-variance R formula ' + 'and use this constant value (mm^2/h^2) at every hour. ' + 'Matches the DualEarth/new_EnKF.py reference (R=0.07). ' + 'Suppresses DA at low flow where kriging is noisy while ' + 'preserving DA strength at storm peaks. Applied to all ' + 'three EnKF instances (da, openloop, fcst) for consistency.') + parser.add_argument('--prod-script', default=None, + help='Path to calibrate_catchment_cfe_da_v2.py for ' + 'importing EnKFAssimilator. Defaults to next-to-this-file.') + args = parser.parse_args() + + here = os.path.dirname(os.path.abspath(__file__)) + prod_script = args.prod_script or os.path.join(here, 'calibrate_catchment_cfe_da_v2.py') + if not os.path.exists(prod_script): + raise FileNotFoundError( + f"Could not find production script at {prod_script}. " + f"Pass --prod-script to override.") + EnKFAssimilator = import_enkf_class(prod_script) + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + + run(args, _bmi_cfe, EnKFAssimilator) + + +if __name__ == '__main__': + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/2_assimilation/run_perturbation_da_on.py b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/2_assimilation/run_perturbation_da_on.py new file mode 100644 index 00000000..6ef27e96 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/2_assimilation/run_perturbation_da_on.py @@ -0,0 +1,400 @@ +""" +run_perturbation_da_on.py — 2a/2b: ensemble spread with DA ON. + +Runs two separated perturbation arms for the Helene window (Sep 24-30), +with DA actively running during the analysis period: + + ARM A — forcing only (30 members): + DA analyzes states using Qkrig at every hour (DA ON). + At each issue_time: fork 30 members with ONLY met forcing perturbed + (different precip/PET draw per member). Hydro states = DA analysis + mean (same for all 30). Free-run 18 hours, no further DA. + + ARM B — hydro states only (20 members): + Same DA analysis. At each issue_time: fork 20 members with ONLY + hydro states perturbed (multiplicative noise on soil/GW/Nash from + the DA analysis ensemble). Forcing = deterministic. Free-run 18 hours. + +These two arms isolate "where does forecast spread come from WITH DA?" +— figures for the assimilation script section. + +Output (per catchment, per arm): + //_da_forcing_arm.csv + //_da_hydro_arm.csv + Columns: issue_time, lead_hour, member_00..member_29 (forcing arm) + issue_time, lead_hour, member_00..member_19 (hydro arm) + +Usage (run for each catchment, both arms): + python3 run_perturbation_da_on.py \\ + --cat-id cat-1016300 \\ + --forcing-dir \\ + --obs-dir /mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance \\ + --cfe-dir /mnt/disk2/suma_helen_poster/cfe_py \\ + --config-file /mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json \\ + --param-bounds /mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json \\ + --out-dir /mnt/disk2/suma_helen_poster/da_results/v2_perturbation_da_on \\ + --test-forcing-dir1 \\ + --test-forcing-dir2 \\ + --hardcoded-r 0.07 +""" + +import argparse +import os +import sys +import json +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +N_FORCING = 30 # members in forcing arm +N_HYDRO = 20 # members in hydro-state arm + +# Analysis period — run DA through entire test period +TEST_START = "2024-08-24 00:00:00" # 1 month spinup before Helene window +TEST_END = "2024-10-31 23:00:00" + +# Issue times for the Helene window (hourly, Sep 24-30) +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-30 06:00:00") +FORECAST_HOURS = 18 + +# Perturbation magnitudes (same as production) +PRECIP_SIGMA = 0.15 +PET_SIGMA = 0.10 +STATE_FRAC = 0.05 # hydro-state perturbation fraction + +# Set at runtime +CAT_ID = None +OBS_FILE = None +CFE_CONFIG_FILE = None +TEST_FORCING_FILE = None +OUT_DIR = None +HARDCODED_R = None +DIRECT_VARIANCE = False +bmi_cfe = None + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(): + df = pd.read_csv(TEST_FORCING_FILE) + df["date"] = pd.to_datetime(df["time"]).dt.strftime("%Y-%m-%d %H:%M:%S") + df["total_precipitation"] = df["APCP_surface"] * 3600.0 + df["potential_evaporation"] = priestley_taylor_pet( + df["DSWRF_surface"].values, df["TMP_2maboveground"].values) + return df + + +def load_obs(): + df = pd.read_csv(OBS_FILE) + date_col = next((c for c in df.columns + if c.lower() in ("date", "time", "datetime", "timestamp")), None) + if date_col is None: + raise ValueError(f"No date column in {OBS_FILE}. Columns: {df.columns.tolist()}") + q_col = next((c for c in df.columns if "qkrig" in c.lower() and "var" not in c.lower()), None) + if q_col is None: + raise ValueError(f"No qkrig column in {OBS_FILE}. Columns: {df.columns.tolist()}") + var_col = next((c for c in df.columns if "var" in c.lower()), None) + df[date_col] = pd.to_datetime(df[date_col]) + df = df.sort_values(date_col).reset_index(drop=True) + date_strs = df[date_col].dt.strftime("%Y-%m-%d %H:%M:%S") + obs_dict = dict(zip(date_strs, df[q_col])) + var_vals = df[var_col] if var_col else pd.Series([0.07] * len(df)) + var_dict = dict(zip(date_strs, var_vals)) + print(f" Obs: date='{date_col}', q='{q_col}', var='{var_col}'") + return obs_dict, var_dict + + +def make_model(best_params): + with open(CFE_CONFIG_FILE) as f: + cfg = json.load(f) + cfg["forcing_file"] = TEST_FORCING_FILE + cfg["soil_params"]["bb"] = best_params["bb"] + cfg["soil_params"]["smcmax"] = best_params["smcmax"] + cfg["soil_params"]["satdk"] = best_params["satdk"] + cfg["slop"] = best_params["slop"] + cfg["max_gw_storage"] = best_params["max_gw_storage"] + cfg["expon"] = best_params["expon"] + cfg["Cgw"] = best_params["Cgw"] + cfg["K_lf"] = best_params["K_lf"] + cfg["K_nash"] = best_params["K_nash"] + cfg["partition_scheme"] = ("Schaake" if best_params["scheme"] <= 0.5 + else "Xinanjiang") + tmp_cfg = str(OUT_DIR / f"{CAT_ID}_bmi_config_temp_da_on.json") + with open(tmp_cfg, "w") as f: + json.dump(cfg, f) + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.initialize() + return m + + +def get_state(m): + return { + "soil_m": m.soil_reservoir["storage_m"], + "gw_m": m.gw_reservoir["storage_m"], + "nash0": float(m.nash_storage[0]), + "nash1": float(m.nash_storage[1]), + } + + +def set_state(m, state): + m.soil_reservoir["storage_m"] = state["soil_m"] + m.gw_reservoir["storage_m"] = state["gw_m"] + m.nash_storage[0] = state["nash0"] + m.nash_storage[1] = state["nash1"] + + +def perturb_state(state, rng, frac=STATE_FRAC): + sm_max = 1.0 # placeholder — will be clipped in set_state; actual max enforced by model + out = { + "soil_m": max(state["soil_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "nash0": max(state["nash0"] * (1.0 + frac * rng.standard_normal()), 0.0), + "nash1": max(state["nash1"] * (1.0 + frac * rng.standard_normal()), 0.0), + } + return out + + +def run_model_step(m, precip_mmh, pet_mmh): + m.set_value("atmosphere_water__time_integral_of_precipitation_mass_flux", + precip_mmh / 1000.0) + m.set_value("water_potential_evaporation_flux", + pet_mmh / 1000.0 / 3600.0) + m.update() + return m.get_value("land_surface_water__runoff_depth") * 1000.0 # m/h -> mm/h + + +def enkf_update(state, q_sim, y_obs, krig_var, rng, n_members=1): + """Simple scalar EnKF update on the mean analysis state.""" + if HARDCODED_R is not None: + R = HARDCODED_R + elif DIRECT_VARIANCE: + R = max(krig_var, 1e-6) + else: + alpha = 0.10 + R = (alpha * max(y_obs, 0.0)) ** 2 + 0.001 * krig_var + R = max(R, 1e-6) + P_yy = max(q_sim * 0.01, 1e-6) # rough prior variance on Q + K = P_yy / (P_yy + R) + innovation = y_obs - q_sim + # Apply correction to each state proportionally (simple scalar gain) + scale = K * innovation / max(abs(q_sim), 1e-6) + updated = { + "soil_m": max(state["soil_m"] * (1.0 + scale), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + scale), 1e-6), + "nash0": max(state["nash0"] + state["nash0"] * scale, 0.0), + "nash1": max(state["nash1"] + state["nash1"] * scale, 0.0), + } + return updated + + +def main(): + global CAT_ID, OBS_FILE, CFE_CONFIG_FILE, TEST_FORCING_FILE + global OUT_DIR, HARDCODED_R, DIRECT_VARIANCE, bmi_cfe + + parser = argparse.ArgumentParser() + parser.add_argument("--cat-id", required=True) + parser.add_argument("--forcing-dir", required=True) + parser.add_argument("--obs-dir", required=True) + parser.add_argument("--cfe-dir", required=True) + parser.add_argument("--config-file", required=True) + parser.add_argument("--param-bounds", required=True) + parser.add_argument("--out-dir", required=True) + parser.add_argument("--test-forcing-dir1", default=None) + parser.add_argument("--test-forcing-dir2", default=None) + parser.add_argument("--hardcoded-r", type=float, default=None, + help="Fixed R (e.g. 0.07). Omit → Vrugt formula.") + parser.add_argument("--direct-variance", action="store_true", default=False, + help="Use krig_var directly as R (no Vrugt scaling)") + parser.add_argument("--rng-seed", type=int, default=None) + args = parser.parse_args() + + CAT_ID = args.cat_id + OBS_FILE = os.path.join(args.obs_dir, f"{CAT_ID}.csv") + CFE_CONFIG_FILE = args.config_file + HARDCODED_R = args.hardcoded_r + DIRECT_VARIANCE = args.direct_variance + OUT_DIR = Path(args.out_dir) / CAT_ID + OUT_DIR.mkdir(parents=True, exist_ok=True) + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + bmi_cfe = _bmi_cfe + + # Build combined test forcing + if args.test_forcing_dir1 and args.test_forcing_dir2: + f1 = os.path.join(args.test_forcing_dir1, f"{CAT_ID}.csv") + f2 = os.path.join(args.test_forcing_dir2, f"{CAT_ID}.csv") + df1, df2 = pd.read_csv(f1), pd.read_csv(f2) + combined = (pd.concat([df1, df2], ignore_index=True) + .drop_duplicates(subset="time").sort_values("time")) + TEST_FORCING_FILE = str(OUT_DIR / f"{CAT_ID}_nwm_operational_combined.csv") + combined.to_csv(TEST_FORCING_FILE, index=False) + else: + print("Need --test-forcing-dir1 and --test-forcing-dir2"); return + + best_params_file = OUT_DIR / f"{CAT_ID}_best_params.json" + if not best_params_file.exists(): + print(f"Missing best params at {best_params_file}"); return + with open(best_params_file) as f: + best_params = json.load(f)["best_parameters"] + + obs_dict, var_dict = load_obs() + df_forcing = load_test_forcing() + t_mask = (df_forcing["date"] >= TEST_START) & (df_forcing["date"] <= TEST_END) + df_test = df_forcing[t_mask].reset_index(drop=True) + + rng = np.random.default_rng(hash(CAT_ID) & 0x7fffffff) + + # ------------------------------------------------------------------ # + # Phase 1: Run DA through test period, snapshot analysis state at # + # each Helene issue time # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Phase 1: DA analysis through test period...") + model = make_model(best_params) + + helene_snapshots = {} # issue_time -> analysis state dict + + for h, row in df_test.iterrows(): + date_str = row["date"] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + q_sim = run_model_step(model, P, E) + + date_ts = pd.Timestamp(date_str) + + # DA update if obs available + y_obs = obs_dict.get(date_str, np.nan) + if not np.isnan(y_obs): + krig_var = var_dict.get(date_str, 0.07) + state = get_state(model) + updated = enkf_update(state, q_sim, y_obs, krig_var, rng) + set_state(model, updated) + + # Snapshot if this is a Helene issue time + if HELENE_START <= date_ts <= HELENE_END: + helene_snapshots[date_ts] = get_state(model) + + model.finalize() + print(f" Snapshots saved: {len(helene_snapshots)} issue times") + + # ------------------------------------------------------------------ # + # Phase 2A: Forcing arm — 30 members, fixed analyzed state, varying # + # met forcing # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Phase 2A: Forcing arm ({N_FORCING} members)...") + forcing_records = [] + + issue_times = sorted(helene_snapshots.keys()) + t0_to_idx = {t: i for i, t in enumerate(df_test["date"].apply(pd.Timestamp))} + + for t0 in issue_times: + snap = helene_snapshots[t0] + t0_str = t0.strftime("%Y-%m-%d %H:%M:%S") + + # Find index in df_test + if t0 not in t0_to_idx: + continue + start_idx = t0_to_idx[t0] + + q_matrix = np.full((FORECAST_HOURS, N_FORCING), np.nan) + + for mem in range(N_FORCING): + m = make_model(best_params) + set_state(m, snap) # all members start from same DA analysis state + + for lead in range(1, FORECAST_HOURS + 1): + fi = start_idx + lead + if fi >= len(df_test): + break + row = df_test.iloc[fi] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + # Perturb forcing only + mu_p = -0.5 * PRECIP_SIGMA ** 2 + P_pert = float(P * rng.lognormal(mu_p, PRECIP_SIGMA)) if P > 0 else 0.0 + E_pert = max(float(E * (1.0 + PET_SIGMA * rng.standard_normal())), 0.0) + q = run_model_step(m, P_pert, E_pert) + q_matrix[lead - 1, mem] = q + + m.finalize() + + row_data = {"issue_time": t0_str} + for lead in range(1, FORECAST_HOURS + 1): + row_data["lead_hour"] = lead + for mem in range(N_FORCING): + row_data[f"member_{mem:02d}"] = q_matrix[lead - 1, mem] + forcing_records.append(dict(row_data)) + + df_forcing_arm = pd.DataFrame(forcing_records) + out_a = OUT_DIR / f"{CAT_ID}_da_forcing_arm.csv" + df_forcing_arm.to_csv(out_a, index=False) + print(f" Saved: {out_a}") + + # ------------------------------------------------------------------ # + # Phase 2B: Hydro-state arm — 20 members, deterministic forcing, # + # perturbed analysis states # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Phase 2B: Hydro-state arm ({N_HYDRO} members)...") + hydro_records = [] + + for t0 in issue_times: + snap = helene_snapshots[t0] + t0_str = t0.strftime("%Y-%m-%d %H:%M:%S") + + if t0 not in t0_to_idx: + continue + start_idx = t0_to_idx[t0] + + q_matrix = np.full((FORECAST_HOURS, N_HYDRO), np.nan) + + for mem in range(N_HYDRO): + # Perturb the DA analysis state for this member + if mem == 0: + state_m = snap # member 0 = unperturbed analysis mean + else: + state_m = perturb_state(snap, rng) + + m = make_model(best_params) + set_state(m, state_m) + + for lead in range(1, FORECAST_HOURS + 1): + fi = start_idx + lead + if fi >= len(df_test): + break + row = df_test.iloc[fi] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + # Deterministic forcing + q = run_model_step(m, P, E) + q_matrix[lead - 1, mem] = q + + m.finalize() + + row_data = {"issue_time": t0_str} + for lead in range(1, FORECAST_HOURS + 1): + row_data["lead_hour"] = lead + for mem in range(N_HYDRO): + row_data[f"member_{mem:02d}"] = q_matrix[lead - 1, mem] + hydro_records.append(dict(row_data)) + + df_hydro_arm = pd.DataFrame(hydro_records) + out_b = OUT_DIR / f"{CAT_ID}_da_hydro_arm.csv" + df_hydro_arm.to_csv(out_b, index=False) + print(f" Saved: {out_b}") + print(f"[{CAT_ID}] Done.") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/2_assimilation/run_perturbation_sensitivity.py b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/2_assimilation/run_perturbation_sensitivity.py new file mode 100644 index 00000000..57b8d185 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/2_assimilation/run_perturbation_sensitivity.py @@ -0,0 +1,293 @@ +""" +Perturbation-source sensitivity analysis for the multi-catchment CFE DA experiment. + +Ensemble-spread attribution: run 20 CFE members through the test +period with ONLY ONE perturbation source active at a time, no DA, no spinup. + +Three sub-experiments (selected via --source): + init : perturb initial states once at t=0 (multiplicative ~5%) + forcing : perturb hourly precip (lognormal sigma=0.15) and PET (Gaussian sigma=0.10) + process : apply per-hour process noise on states (soil 0.2%, GW 0.15%, Nash 0.5%) + +Each sub-experiment runs N=20 CFE members. DA is always OFF. Spinup is skipped — each +member starts from CFE's default initial state (modified by the perturbation if --source +init). This is the cleanest setup for attributing test-period spread to its source. + +Output: + //_sensitivity_.csv + Columns: date, member_00, member_01, ..., member_19 + +Each member column is that member's hourly Q_sim (mm/h) over the test period. + +Reads the same calibrated best_params.json as the production script. Does not need +the kriging obs file (no DA), but takes the same --obs-dir argument so it can also +log Q_obs alongside for plotting reference. + +Usage: + python3 run_perturbation_sensitivity.py \ + --cat-id cat-1016300 \ + --source forcing \ + --forcing-dir \ + --obs-dir \ + --cfe-dir \ + --config-file \ + --param-bounds \ + --out-dir \ + --test-forcing-dir1 \ + --test-forcing-dir2 +""" + +import argparse +import os +import sys +import json +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 +N_MEMBERS = 20 + +# We only run the test window for this sensitivity analysis. +TEST_START = "2023-10-01 00:00:00" +TEST_END = "2024-10-31 23:00:00" + +# Perturbation magnitudes — same as production defaults +INIT_STATE_FRAC = 0.05 +PRECIP_PERTURB = 0.15 # lognormal sigma +PET_PERTURB = 0.10 # Gaussian sigma +SOIL_PROC_NOISE = 0.002 +GW_PROC_NOISE = 0.0015 +NASH_PROC_NOISE = 0.005 + +# Set at runtime +CAT_ID = None +FORCING_FILE = None +TEST_FORCING_FILE = None +OBS_FILE = None +CFE_CONFIG_FILE = None +OUT_DIR = None +bmi_cfe = None + + +# ---------- Forcing helpers (same as production) ---------- +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(): + df = pd.read_csv(TEST_FORCING_FILE) + df['date'] = pd.to_datetime(df['time']).dt.strftime('%Y-%m-%d %H:%M:%S') + df['total_precipitation'] = df['APCP_surface'] * 3600.0 # kg/m^2/s -> mm/h + df['potential_evaporation'] = priestley_taylor_pet( + df['DSWRF_surface'].values, + df['TMP_2maboveground'].values, + ) + return df + + +# ---------- Perturbation helpers ---------- +def perturb_precip_lognormal(P, rng, sigma=PRECIP_PERTURB): + """Lognormal multiplier with mean 1; preserves zero precip.""" + mu = -0.5 * sigma ** 2 + return float(P * rng.lognormal(mu, sigma)) + + +def perturb_pet_gaussian(E, rng, sigma=PET_PERTURB): + """Gaussian multiplicative noise on PET, clipped at 0.""" + return max(float(E * (1.0 + sigma * rng.standard_normal())), 0.0) + + +def apply_init_state_perturb(m, rng, frac=INIT_STATE_FRAC): + """One-time multiplicative perturbation on the four states (called per member at t=0).""" + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1.0 + frac * rng.standard_normal()), 1e-6, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1.0 + frac * rng.standard_normal()), 1e-6, gw_max)) + # Nash buckets often start at 0; small additive jitter so they're not identical. + jitter = max(sm_max * 1e-6, 1e-9) + m.nash_storage[0] = max(n00 + jitter * rng.standard_normal(), 0.0) + m.nash_storage[1] = max(n10 + jitter * rng.standard_normal(), 0.0) + + +def apply_process_noise(m, rng): + """Per-timestep process noise on all 4 states for one member.""" + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + nash_floor = max(sm_max * 1e-4, 1e-7) + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1.0 + SOIL_PROC_NOISE * rng.standard_normal()), 0.0, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1.0 + GW_PROC_NOISE * rng.standard_normal()), 0.0, gw_max)) + m.nash_storage[0] = max( + n00 * (1.0 + NASH_PROC_NOISE * rng.standard_normal()) + + nash_floor * rng.standard_normal(), 0.0) + m.nash_storage[1] = max( + n10 * (1.0 + NASH_PROC_NOISE * rng.standard_normal()) + + nash_floor * rng.standard_normal(), 0.0) + + +# ---------- Main per-source run ---------- +def run_sensitivity(source, best_param_dict): + """Run N members through the test period with only `source` perturbation active.""" + + def custom_load_forcing(self_cfe): + df = load_test_forcing() + self_cfe.forcing_data = df.rename(columns={"date": "time"}) + + # Build a per-catchment temporary CFE config with the calibrated params + with open(CFE_CONFIG_FILE) as f: + cfg = json.load(f) + cfg['forcing_file'] = TEST_FORCING_FILE + cfg['soil_params']['bb'] = best_param_dict['bb'] + cfg['soil_params']['smcmax'] = best_param_dict['smcmax'] + cfg['soil_params']['satdk'] = best_param_dict['satdk'] + cfg['slop'] = best_param_dict['slop'] + cfg['max_gw_storage'] = best_param_dict['max_gw_storage'] + cfg['expon'] = best_param_dict['expon'] + cfg['Cgw'] = best_param_dict['Cgw'] + cfg['K_lf'] = best_param_dict['K_lf'] + cfg['K_nash'] = best_param_dict['K_nash'] + cfg['partition_scheme'] = "Schaake" if best_param_dict['scheme'] <= 0.5 else "Xinanjiang" + + tmp_cfg = str(OUT_DIR / f'{CAT_ID}_bmi_config_temp_sensitivity_{source}.json') + with open(tmp_cfg, 'w') as f: + json.dump(cfg, f) + + # Deterministic seed per (catchment, source) so each run is reproducible + seed = hash((CAT_ID, source)) & 0x7fffffff + rng = np.random.default_rng(seed) + + print(f"[sensitivity] {CAT_ID} | source={source} | N={N_MEMBERS} | seed={seed}") + + # Build N members. Initial-state perturbation only happens when source == 'init'. + models = [] + for i in range(N_MEMBERS): + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.load_forcing_file = custom_load_forcing.__get__(m) + m.initialize() + if source == 'init' and i > 0: + apply_init_state_perturb(m, rng) + models.append(m) + + # Skip spinup. Go directly to the test period. + df = load_test_forcing() + t_mask = (df['date'] >= TEST_START) & (df['date'] <= TEST_END) + df_test = df[t_mask] + + # Pre-allocate per-member Q arrays + n_hours = len(df_test) + q_matrix = np.full((n_hours, N_MEMBERS), np.nan, dtype=float) + dates_out = df_test['date'].values + + for h, (p, e) in enumerate(zip(df_test['total_precipitation'], + df_test['potential_evaporation'])): + for i, m in enumerate(models): + if source == 'forcing': + p_i = perturb_precip_lognormal(p, rng) + e_i = perturb_pet_gaussian(e, rng) + else: + p_i = float(p) + e_i = float(e) + m.set_value('atmosphere_water__time_integral_of_precipitation_mass_flux', p_i / 1000) + m.set_value('water_potential_evaporation_flux', e_i / 1000 / 3600) + m.update() + q_matrix[h, i] = m.get_value('land_surface_water__runoff_depth') * 1000 # m/h -> mm/h + + # Process noise applied after each hour, but only when source == 'process' + if source == 'process': + for m in models: + apply_process_noise(m, rng) + + for m in models: + m.finalize() + + # Write per-member CSV: date + 20 columns + out = {'date': dates_out} + for i in range(N_MEMBERS): + out[f'member_{i:02d}'] = q_matrix[:, i] + df_out = pd.DataFrame(out) + out_path = OUT_DIR / f'{CAT_ID}_sensitivity_{source}.csv' + df_out.to_csv(out_path, index=False) + print(f"[sensitivity] saved {out_path}") + print(f"[sensitivity] mean ensemble spread (std across members) over test period: " + f"{float(q_matrix.std(axis=1).mean()):.5f} mm/h") + + +def main(): + global CAT_ID, FORCING_FILE, TEST_FORCING_FILE, OBS_FILE, CFE_CONFIG_FILE + global OUT_DIR, bmi_cfe + + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', required=True) + parser.add_argument('--source', required=True, + choices=['init', 'forcing', 'process'], + help="Which perturbation source to isolate") + parser.add_argument('--forcing-dir', required=True) + parser.add_argument('--obs-dir', required=True, + help='Kept for symmetry with production script; not used since DA is off') + parser.add_argument('--cfe-dir', required=True) + parser.add_argument('--config-file', required=True) + parser.add_argument('--param-bounds', required=True, + help='Kept for symmetry; not read in this script') + parser.add_argument('--out-dir', required=True) + parser.add_argument('--test-forcing-dir1', default=None) + parser.add_argument('--test-forcing-dir2', default=None) + args = parser.parse_args() + + CAT_ID = args.cat_id + FORCING_FILE = os.path.join(args.forcing_dir, f'{CAT_ID}.csv') + OBS_FILE = os.path.join(args.obs_dir, f'{CAT_ID}.csv') + CFE_CONFIG_FILE = args.config_file + OUT_DIR = Path(args.out_dir) / CAT_ID + OUT_DIR.mkdir(parents=True, exist_ok=True) + + # Build combined test forcing CSV (same idiom as production) + if args.test_forcing_dir1 and args.test_forcing_dir2: + f1 = os.path.join(args.test_forcing_dir1, f'{CAT_ID}.csv') + f2 = os.path.join(args.test_forcing_dir2, f'{CAT_ID}.csv') + if os.path.exists(f1) and os.path.exists(f2): + df1 = pd.read_csv(f1) + df2 = pd.read_csv(f2) + combined = pd.concat([df1, df2], ignore_index=True) + combined = combined.drop_duplicates(subset='time').sort_values('time') + TEST_FORCING_FILE = str(OUT_DIR / f'{CAT_ID}_nwm_operational_combined.csv') + combined.to_csv(TEST_FORCING_FILE, index=False) + else: + print(f"Warning: test forcing files not found for {CAT_ID}, exiting") + return + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + bmi_cfe = _bmi_cfe + + best_params_file = OUT_DIR / f'{CAT_ID}_best_params.json' + if not best_params_file.exists(): + print(f"No best params file at {best_params_file}. Pre-stage from Run 3.") + return + with open(best_params_file) as f: + saved = json.load(f) + best_param_dict = saved['best_parameters'] + + run_sensitivity(args.source, best_param_dict) + + +if __name__ == '__main__': + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/2_assimilation/run_production_per_member.py b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/2_assimilation/run_production_per_member.py new file mode 100644 index 00000000..d9c1c4ac --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/2_assimilation/run_production_per_member.py @@ -0,0 +1,314 @@ +""" +Run the production DA pipeline for ONE catchment and save per-member streamflow. + +Identical to calibrate_catchment_cfe_da_v2.py in algorithm (init perturbation + +forcing perturbation + process noise + true EnKF with Vrugt R), but writes one +column per ensemble member instead of only the ensemble mean. This is for +the per-member factor-decomposition plot. + +To keep the comparison apples-to-apples with the sensitivity runs (which skip +spinup), this script also SKIPS the spinup loop. Each member starts directly +from CFE's default initial state, plus the t=0 init perturbation. + +Output: + //_production_per_member.csv + Columns: date, member_00, member_01, ..., member_19 + +Imports EnKFAssimilator from calibrate_catchment_cfe_da_v2 to guarantee the +DA math matches production exactly. + +Usage: + python3 run_production_per_member.py \ + --cat-id cat-1016300 \ + --forcing-dir /mnt/.../nwm_retro_catchment_forcings \ + --obs-dir /mnt/.../catchment_ts_03463300_with_variance \ + --cfe-dir /mnt/.../cfe_py \ + --config-file /mnt/.../cat_03463300_bmi_config_cfe.json \ + --param-bounds /mnt/.../CFE_parameter_bounds.json \ + --out-dir /mnt/.../v2_production_per_member \ + --test-forcing-dir1 ... \ + --test-forcing-dir2 ... \ + --enkf-members 20 +""" + +import argparse +import os +import sys +import json +import importlib.util +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +TEST_START = "2023-10-01 00:00:00" +TEST_END = "2024-10-31 23:00:00" + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(test_forcing_file): + df = pd.read_csv(test_forcing_file) + df['date'] = pd.to_datetime(df['time']).dt.strftime('%Y-%m-%d %H:%M:%S') + df['total_precipitation'] = df['APCP_surface'] * 3600.0 + df['potential_evaporation'] = priestley_taylor_pet( + df['DSWRF_surface'].values, df['TMP_2maboveground'].values + ) + return df + + +def import_enkf_class(script_path): + """Load EnKFAssimilator from the production script without triggering its main().""" + spec = importlib.util.spec_from_file_location("prod_v2", script_path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod.EnKFAssimilator + + +def run(args, bmi_cfe, EnKFAssimilator): + cat_id = args.cat_id + out_dir = Path(args.out_dir) / cat_id + out_dir.mkdir(parents=True, exist_ok=True) + + # Combined test forcing + f1 = os.path.join(args.test_forcing_dir1, f'{cat_id}.csv') + f2 = os.path.join(args.test_forcing_dir2, f'{cat_id}.csv') + df1 = pd.read_csv(f1) + df2 = pd.read_csv(f2) + combined = pd.concat([df1, df2], ignore_index=True) + combined = combined.drop_duplicates(subset='time').sort_values('time') + test_forcing_file = str(out_dir / f'{cat_id}_nwm_operational_combined.csv') + combined.to_csv(test_forcing_file, index=False) + + # Load best_params + best_params_file = out_dir / f'{cat_id}_best_params.json' + if not best_params_file.exists(): + print(f"Missing {best_params_file}. Pre-stage from Run 3.") + return + with open(best_params_file) as f: + saved = json.load(f) + best = saved['best_parameters'] + + # Build temp CFE config with calibrated params + with open(args.config_file) as f: + cfg = json.load(f) + cfg['forcing_file'] = test_forcing_file + cfg['soil_params']['bb'] = best['bb'] + cfg['soil_params']['smcmax'] = best['smcmax'] + cfg['soil_params']['satdk'] = best['satdk'] + cfg['slop'] = best['slop'] + cfg['max_gw_storage'] = best['max_gw_storage'] + cfg['expon'] = best['expon'] + cfg['Cgw'] = best['Cgw'] + cfg['K_lf'] = best['K_lf'] + cfg['K_nash'] = best['K_nash'] + cfg['partition_scheme'] = "Schaake" if best['scheme'] <= 0.5 else "Xinanjiang" + tmp_cfg = str(out_dir / f'{cat_id}_bmi_config_temp_per_member.json') + with open(tmp_cfg, 'w') as f: + json.dump(cfg, f) + + # Build the EnKF assimilator (same defaults as production) + obs_file = os.path.join(args.obs_dir, f'{cat_id}.csv') + enkf = EnKFAssimilator( + n_members=args.enkf_members, + obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, + use_vrugt_r=(not args.no_vrugt_r) and (not getattr(args, 'direct_variance', False)), + vrugt_alpha=args.vrugt_alpha, + vrugt_scale=args.vrugt_scale, + rng_seed=args.rng_seed, + ) + N = enkf.n_members + print(f"[per-member] {cat_id} | N={N} | use_vrugt_r={enkf.use_vrugt_r}") + if getattr(args, 'hardcoded_r', None) is not None: + enkf.obs_var_dict = {k: args.hardcoded_r for k in enkf.obs_var_dict} + print(f"[per-member] obs_var_dict overridden: R = {args.hardcoded_r} (fixed)") + if getattr(args, 'direct_variance', False): + obs_df = pd.read_csv(obs_file) + t_col = next(c for c in obs_df.columns if c.lower() in ('datetime', 'date', 'time', 'timestamp')) + v_col = next(c for c in obs_df.columns if 'var' in c.lower()) + obs_df[t_col] = pd.to_datetime(obs_df[t_col]) + obs_df = obs_df.set_index(t_col).sort_index() + enkf.obs_var_dict = { + str(ts): float(v) + for ts, v in obs_df[v_col].items() + if pd.notna(v) and float(v) > 0 + } + print(f"[per-member] obs_var_dict overridden: R = σ² direct ({len(enkf.obs_var_dict)} timesteps)") + + # Custom forcing loader for CFE BMI + def custom_load_forcing(self_cfe): + df = load_test_forcing(test_forcing_file) + self_cfe.forcing_data = df.rename(columns={"date": "time"}) + + # Build N CFE instances; perturb initial states (member 0 left clean) + models = [] + init_states_per_member = [] # capture initial-state snapshot after perturbation + for i in range(N): + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.load_forcing_file = custom_load_forcing.__get__(m) + m.initialize() + if i > 0: + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + r = enkf.init_state_perturb_frac + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1 + r * enkf.rng.standard_normal()), 1e-6, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1 + r * enkf.rng.standard_normal()), 1e-6, gw_max)) + jitter = max(sm_max * 1e-6, 1e-9) + m.nash_storage[0] = max(n00 + jitter * enkf.rng.standard_normal(), 0.0) + m.nash_storage[1] = max(n10 + jitter * enkf.rng.standard_normal(), 0.0) + # Snapshot the (possibly-perturbed) initial state for this member + init_states_per_member.append({ + "soil_m": float(m.soil_reservoir["storage_m"]), + "gw_m": float(m.gw_reservoir["storage_m"]), + "nash0_m": float(m.nash_storage[0]), + "nash1_m": float(m.nash_storage[1]), + "soil_max_m": float(m.soil_reservoir["storage_max_m"]), + "gw_max_m": float(m.gw_reservoir["storage_max_m"]), + }) + models.append(m) + + # SKIP SPINUP — go directly to the test period (matches sensitivity-script setup) + df = load_test_forcing(test_forcing_file) + t_mask = (df['date'] >= TEST_START) & (df['date'] <= TEST_END) + df_test = df[t_mask] + n_hours = len(df_test) + q_matrix = np.full((n_hours, N), np.nan, dtype=float) + precip_matrix = np.full((n_hours, N), np.nan, dtype=float) # mm/h per member + pet_matrix = np.full((n_hours, N), np.nan, dtype=float) # mm/h per member + dates_out = df_test['date'].values + + for h, (p, e, current_date) in enumerate(zip( + df_test['total_precipitation'], + df_test['potential_evaporation'], + df_test['date'])): + + # Perturb forcing per member + p_arr, e_arr = enkf.perturb_forcing(p, e) + # Capture the per-member perturbed forcing for later plotting/diagnostics + precip_matrix[h, :] = p_arr + pet_matrix[h, :] = e_arr + + # Advance each member one hour + ensemble_q = np.empty(N, dtype=float) + for i, m in enumerate(models): + m.set_value('atmosphere_water__time_integral_of_precipitation_mass_flux', + float(p_arr[i]) / 1000) + m.set_value('water_potential_evaporation_flux', + float(e_arr[i]) / 1000 / 3600) + m.update() + ensemble_q[i] = m.get_value('land_surface_water__runoff_depth') * 1000 # mm/h + + # Record pre-analysis forecast per member (this is the "actual production" value) + q_matrix[h, :] = ensemble_q + + # DA step (uses production EnKFAssimilator) + enkf.update_states(models, current_date, ensemble_q) + + # Process noise after DA + enkf.add_process_noise(models) + + for m in models: + m.finalize() + + # ----- Save per-member CSVs ----- + def _save_matrix(matrix, suffix): + out = {'date': dates_out} + for i in range(N): + out[f'member_{i:02d}'] = matrix[:, i] + df_out = pd.DataFrame(out) + out_path = out_dir / f'{cat_id}_production_per_member_{suffix}.csv' + df_out.to_csv(out_path, index=False) + return out_path + + # Streamflow outputs + q_path = out_dir / f'{cat_id}_production_per_member.csv' + out = {'date': dates_out} + for i in range(N): + out[f'member_{i:02d}'] = q_matrix[:, i] + pd.DataFrame(out).to_csv(q_path, index=False) + print(f"[per-member] saved {q_path}") + + # Perturbed forcings (so each member's actual inputs are recoverable) + precip_path = _save_matrix(precip_matrix, "precip") + pet_path = _save_matrix(pet_matrix, "pet") + print(f"[per-member] saved {precip_path}") + print(f"[per-member] saved {pet_path}") + + # Initial states (small, JSON is fine) + init_states_path = out_dir / f'{cat_id}_production_per_member_initial_states.json' + with open(init_states_path, 'w') as f: + json.dump({ + "catchment_id": cat_id, + "n_members": N, + "members": {f"member_{i:02d}": init_states_per_member[i] for i in range(N)}, + }, f, indent=2) + print(f"[per-member] saved {init_states_path}") + + print(f"[per-member] ensemble mean spread over test period: " + f"{float(q_matrix.std(axis=1).mean()):.5f} mm/h") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', required=True) + parser.add_argument('--forcing-dir', required=True) + parser.add_argument('--obs-dir', required=True) + parser.add_argument('--cfe-dir', required=True) + parser.add_argument('--config-file', required=True) + parser.add_argument('--param-bounds', required=True, + help='Kept for symmetry; not read in this script') + parser.add_argument('--out-dir', required=True) + parser.add_argument('--test-forcing-dir1', required=True) + parser.add_argument('--test-forcing-dir2', required=True) + parser.add_argument('--enkf-members', type=int, default=20) + parser.add_argument('--enkf-obs-error-std', type=float, default=0.05) + parser.add_argument('--no-vrugt-r', action='store_true') + parser.add_argument('--hardcoded-r', type=float, default=None, + help='Fix R to this constant value for all timesteps (overrides Vrugt formula)') + parser.add_argument('--direct-variance', action='store_true', default=False, + help='Use kriging variance column as R per timestep (R=σ²; overrides Vrugt formula)') + parser.add_argument('--vrugt-alpha', type=float, default=0.10) + parser.add_argument('--vrugt-scale', type=float, default=0.001) + parser.add_argument('--rng-seed', type=int, default=None) + parser.add_argument('--prod-script', default=None, + help='Path to calibrate_catchment_cfe_da_v2.py for importing ' + 'EnKFAssimilator. Defaults to the script located next to this file.') + args = parser.parse_args() + + # Locate the production script (default: same dir as this script) + if args.prod_script is None: + here = os.path.dirname(os.path.abspath(__file__)) + prod_script = os.path.join(here, 'calibrate_catchment_cfe_da_v2.py') + else: + prod_script = args.prod_script + if not os.path.exists(prod_script): + raise FileNotFoundError(f"Could not find production script at {prod_script}. " + f"Pass --prod-script to override.") + EnKFAssimilator = import_enkf_class(prod_script) + + # Dynamic CFE import (same idiom as production script) + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + + run(args, _bmi_cfe, EnKFAssimilator) + + +if __name__ == '__main__': + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/3_routing/route_det_f3.sh b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/3_routing/route_det_f3.sh new file mode 100644 index 00000000..2facdf1c --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/3_routing/route_det_f3.sh @@ -0,0 +1,28 @@ +#!/usr/bin/env bash +# F3 dynamic Vrugt seeded (20pct gauge holdout) — deterministic T-route routing. +# Routes per-catchment _test_results.csv (sim_mm_h column) through Muskingum-Cunge +# to gauge 03463300 and writes routed_Q_test.csv + KGE/NSE summary. +# +# Usage: +# bash route_det_f3.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +ROUTE_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/2_troute_routing/run_route.py + +F3_DIR=/mnt/disk2/1400_sites_helene/da_arms_dynamic_vrugt_seeded +GPKG=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +echo "[F3] Deterministic T-route routing..." +echo " da-dir : $F3_DIR" +echo " out-dir: $F3_DIR" + +$TROUTE "$ROUTE_SCRIPT" \ + --gpkg "$GPKG" \ + --da-dir "$F3_DIR" \ + --out-dir "$F3_DIR" \ + --usgs-csv "$USGS_CSV" + +echo "[F3] Deterministic routing done. Output: $F3_DIR/routed_Q_test.csv" diff --git a/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/3_routing/route_ensemble_f3.sh b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/3_routing/route_ensemble_f3.sh new file mode 100644 index 00000000..b15adcd5 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/3_routing/route_ensemble_f3.sh @@ -0,0 +1,30 @@ +#!/usr/bin/env bash +# F3 dynamic Vrugt seeded (20pct gauge holdout) — crossed-ensemble T-route routing. +# Routes all 600 crossed-ensemble members through Muskingum-Cunge +# to gauge 03463300 and writes routed_crossed_ensemble.parquet. +# +# Must be run AFTER crossed ensemble parquets exist in CROSSED_DIR. +# +# Usage: +# bash route_ensemble_f3.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +ROUTE_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/2_troute_routing/run_route_crossed_ensemble.py + +CROSSED_DIR=/mnt/disk2/1400_sites_helene/da_crossed_dynamic_vrugt_seeded +GPKG=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +echo "[F3] Crossed-ensemble T-route routing..." +echo " ensemble-dir: $CROSSED_DIR" +echo " out-dir : $CROSSED_DIR" + +$TROUTE "$ROUTE_SCRIPT" \ + --gpkg "$GPKG" \ + --ensemble-dir "$CROSSED_DIR" \ + --out-dir "$CROSSED_DIR" \ + --usgs-csv "$USGS_CSV" + +echo "[F3] Ensemble routing done. Output: $CROSSED_DIR/routed_crossed_ensemble.parquet" diff --git a/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/3_routing/route_leadtime_f3.sh b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/3_routing/route_leadtime_f3.sh new file mode 100644 index 00000000..02ef36a1 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/3_routing/route_leadtime_f3.sh @@ -0,0 +1,31 @@ +#!/usr/bin/env bash +# F3 dynamic Vrugt seeded (20pct gauge holdout) — route lead-time forecast CSVs. +# Routes all 21-catchment lead-time forecast CSVs through Muskingum-Cunge +# to gauge 03463300. Writes routed_leadtime_da_full.parquet and +# routed_leadtime_openloop_full.parquet for gauge-level error decay analysis. +# +# Run AFTER the lead-time sweep batch finishes. +# +# Usage: +# bash route_leadtime_f3.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +ROUTE_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/2_troute_routing/run_route_leadtime_forecasts.py + +F3_DIR=/mnt/disk2/1400_sites_helene/da_forecast_dynamic_vrugt_seeded +GPKG=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg +OUT_DIR=/mnt/disk2/suma_helen_poster/leadtime_troute_routing_f3 + +echo "[F3] Routing lead-time forecasts through T-route..." +echo " forecast-dir: $F3_DIR" +echo " out-dir : $OUT_DIR" + +$TROUTE "$ROUTE_SCRIPT" \ + --gpkg "$GPKG" \ + --forecast-dir "$F3_DIR" \ + --out-dir "$OUT_DIR" + +echo "[F3] Lead-time routing done. Outputs in: $OUT_DIR" +ls "$OUT_DIR"/*.parquet 2>/dev/null || echo " (no parquets found)" diff --git a/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4a_error_decay/plot_forecast_error_fixed_target.py b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4a_error_decay/plot_forecast_error_fixed_target.py new file mode 100644 index 00000000..1460ed21 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4a_error_decay/plot_forecast_error_fixed_target.py @@ -0,0 +1,217 @@ +""" +plot_forecast_error_fixed_target.py — 4a: error decay, fixed-target-time view. + +For each target verification time T in the Helene peak window: + Collect all forecasts that verify AT T: + issue_time = T - lead_hour*1h, for lead in 1..18 + error[lead] = ensemble_mean(q at T, initialized T-lead) - USGS_obs(T) + +This gives the correct operational picture: + - lead 1 = initialized 1 hr before T (DA just ran -> small error) + - lead 18 = initialized 18 hr before T (DA long ago -> error ~ open loop) + +Two panels: + Top : signed error (m³/s) vs lead hour, one curve per target time + Bot : same for open-loop +Plus a summary panel: mean across all target times, DA vs OL. + +Outputs: + /error_fixed_target_helene.png (per-target spaghetti, DA vs OL) + /error_fixed_target_mean.png (mean across targets, DA vs OL) +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.cm as cm + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast_hardcoded_r_routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +# Target verification times: hourly through the Helene peak window +TARGET_START = pd.Timestamp("2024-09-26 18:00:00") +TARGET_END = pd.Timestamp("2024-09-28 06:00:00") + + +def load_parquet(path): + df = pd.read_parquet(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + df["ens_mean"] = df[member_cols].mean(axis=1) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + return df[["issue_time", "lead_hour", "valid_time", "ens_mean"]] + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() + or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def build_fixed_target_errors(df, obs_series, target_times): + """ + For each target time T and each lead L (1..18): + issue_time = T - L hours + error = ensemble_mean at (issue_time, lead=L) - obs(T) + Returns dict: target_time -> {lead: error} + """ + # Index df by (issue_time, lead_hour) for fast lookup + df_idx = df.set_index(["issue_time", "lead_hour"])["ens_mean"] + + results = {} + for T in target_times: + obs_val = obs_series.get(T, np.nan) + if np.isnan(obs_val): + continue + curve = {} + for lead in range(1, 19): + t0 = T - pd.Timedelta(hours=lead) + try: + q_fc = df_idx.loc[(t0, lead)] + curve[lead] = float(q_fc) - obs_val + except KeyError: + curve[lead] = np.nan + results[T] = curve + return results + + +def plot_spaghetti(ax, error_dict, color_da, label_prefix, linestyle="-", lw=0.9, alpha=0.35): + """Plot one thin line per target time + thick mean.""" + leads = list(range(1, 19)) + all_curves = [] + target_times = sorted(error_dict.keys()) + cmap = cm.get_cmap("plasma", len(target_times)) + + for i, T in enumerate(target_times): + curve = [error_dict[T].get(l, np.nan) for l in leads] + ax.plot(leads, curve, + color=cmap(i), lw=lw, alpha=alpha, + linestyle=linestyle, zorder=2) + all_curves.append(curve) + + if all_curves: + mean_curve = np.nanmean(np.array(all_curves), axis=0) + ax.plot(leads, mean_curve, + color=color_da, lw=2.8, alpha=0.95, + linestyle=linestyle, zorder=5, + label=f"{label_prefix} — mean across targets") + return all_curves + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--route-dir", default=DEFAULT_ROUTE_DIR) + parser.add_argument("--out-dir", default=None) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--da-name", default="routed_leadtime_da_full.parquet") + parser.add_argument("--ol-name", default="routed_leadtime_openloop_full.parquet") + parser.add_argument("--target-start", default=str(TARGET_START)) + parser.add_argument("--target-end", default=str(TARGET_END)) + args = parser.parse_args() + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + print("Loading parquets...") + da = load_parquet(os.path.join(args.route_dir, args.da_name)) + ol = load_parquet(os.path.join(args.route_dir, args.ol_name)) + obs = load_usgs(args.usgs_csv) + + target_times = pd.date_range(args.target_start, args.target_end, freq="1h") + print(f" Target verification times: {len(target_times)} " + f"({target_times[0]} → {target_times[-1]})") + + print("Building fixed-target error tables...") + da_errors = build_fixed_target_errors(da, obs, target_times) + ol_errors = build_fixed_target_errors(ol, obs, target_times) + print(f" Targets with obs: DA={len(da_errors)} OL={len(ol_errors)}") + + leads = list(range(1, 19)) + + # ---- Spaghetti: per-target-time curves, DA vs OL ---- + fig, (ax_da, ax_ol) = plt.subplots(2, 1, figsize=(13, 10), sharex=True, sharey=True) + + ax_da.axhline(0, color="black", lw=0.8, linestyle="--", alpha=0.4) + plot_spaghetti(ax_da, da_errors, "tab:purple", "DA", linestyle="-") + ax_da.set_ylabel("Error: forecast − USGS obs (m³/s)", fontsize=10) + ax_da.set_title("DA — error at each lead for fixed target times (Helene peak window)", fontsize=11) + ax_da.grid(True, alpha=0.2) + ax_da.legend(fontsize=9) + + ax_ol.axhline(0, color="black", lw=0.8, linestyle="--", alpha=0.4) + plot_spaghetti(ax_ol, ol_errors, "tab:gray", "Open-loop", linestyle="--") + ax_ol.set_ylabel("Error: forecast − USGS obs (m³/s)", fontsize=10) + ax_ol.set_xlabel("Forecast lead hour (hours before target)", fontsize=11) + ax_ol.set_title("Open-loop — error at each lead for fixed target times", fontsize=11) + ax_ol.set_xticks(leads) + ax_ol.grid(True, alpha=0.2) + ax_ol.legend(fontsize=9) + + fig.suptitle( + "Forecast error vs lead time — fixed verification time, Helene peak window\n" + f"Target times: {TARGET_START.strftime('%b %d %H UTC')} → " + f"{TARGET_END.strftime('%b %d %H UTC')} | USGS 03463300", + fontsize=12, y=1.01, + ) + plt.tight_layout() + out1 = os.path.join(out_dir, "error_fixed_target_helene.png") + plt.savefig(out1, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + # ---- Summary: mean across all target times, DA vs OL overlaid ---- + fig, ax = plt.subplots(figsize=(13, 6)) + ax.axhline(0, color="black", lw=0.8, linestyle="--", alpha=0.4) + + def mean_curve(error_dict): + arr = np.array([ + [error_dict[T].get(l, np.nan) for l in leads] + for T in sorted(error_dict.keys()) + ]) + return np.nanmean(arr, axis=0), np.nanstd(arr, axis=0) + + da_mean, da_std = mean_curve(da_errors) + ol_mean, ol_std = mean_curve(ol_errors) + + ax.fill_between(leads, da_mean - da_std, da_mean + da_std, + color="tab:purple", alpha=0.15, zorder=2) + ax.fill_between(leads, ol_mean - ol_std, ol_mean + ol_std, + color="tab:gray", alpha=0.15, zorder=2) + ax.plot(leads, da_mean, color="tab:purple", lw=2.6, marker="o", + zorder=5, label="DA — mean error (±1 std shaded)") + ax.plot(leads, ol_mean, color="tab:gray", lw=2.6, marker="s", + linestyle="--", zorder=5, label="Open-loop — mean error (±1 std shaded)") + + ax.set_xlabel("Forecast lead hour (hours before target verification time)", fontsize=11) + ax.set_ylabel("Mean error: forecast − USGS obs (m³/s)", fontsize=11) + ax.set_xticks(leads) + ax.set_title( + "Mean forecast error vs lead time — fixed verification time, Helene peak window\n" + f"Target: {TARGET_START.strftime('%b %d %H UTC')} → " + f"{TARGET_END.strftime('%b %d %H UTC')} | " + "Lead 1 = init 1 hr before target | Lead 18 = init 18 hr before target", + fontsize=11, + ) + ax.legend(fontsize=10, loc="lower right", frameon=True, framealpha=0.92) + ax.grid(True, alpha=0.25) + plt.tight_layout() + out2 = os.path.join(out_dir, "error_fixed_target_mean.png") + plt.savefig(out2, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4a_error_decay/plot_forecast_error_per_init.py b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4a_error_decay/plot_forecast_error_per_init.py new file mode 100644 index 00000000..691245d0 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4a_error_decay/plot_forecast_error_per_init.py @@ -0,0 +1,203 @@ +""" +plot_forecast_error_per_init.py + +Error decay by initialization time for the Helene window (Sep 24-28 2024). + +For each initialization time t0 in the Helene window: + error[lead] = ensemble_mean(q_gauge_m3s at t0+lead) - USGS_obs(t0+lead) + +Plotted as: + DA : thin colored lines (one per init time, colored by date) + thick mean across all + OL : thin gray dashed lines + thick gray dashed mean + +x-axis: forecast lead hour (1 -> 18) +y-axis: signed error (m³/s), positive = forecast too high + +The expected signal: DA error is small at lead 1 (just assimilated), grows +and converges toward the OL error curve by lead 18. + +Inputs: + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + +Obs: + /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +Outputs: + /forecast_error_per_init_helene.png (signed error) + /forecast_rmse_per_init_helene.png (absolute error / RMSE per init) +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.patches as mpatches + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast_hardcoded_r_routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +# Initialization times to show — Helene window +HELENE_INIT_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_INIT_END = pd.Timestamp("2024-09-28 23:00:00") + +# One color per init date +DATE_COLORS = { + "2024-09-24": "#1f77b4", + "2024-09-25": "#ff7f0e", + "2024-09-26": "#2ca02c", + "2024-09-27": "#d62728", + "2024-09-28": "#9467bd", +} + + +def load_parquet(path): + df = pd.read_parquet(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + df["ens_mean"] = df[member_cols].mean(axis=1) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + return df[["issue_time", "lead_hour", "valid_time", "ens_mean"]] + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() + or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def build_error_table(df, obs_series, init_start, init_end): + """Return DataFrame: issue_time, lead_hour, ens_mean, obs, error.""" + df = df[(df["issue_time"] >= init_start) & (df["issue_time"] <= init_end)].copy() + df["obs"] = df["valid_time"].map(obs_series) + df["error"] = df["ens_mean"] - df["obs"] + return df.dropna(subset=["obs", "error"]) + + +def plot_error(ax, err_df, color, alpha_thin, lw_thin, linestyle, label_prefix): + """Plot individual init-time error curves + thick mean curve.""" + leads = sorted(err_df["lead_hour"].unique()) + all_curves = [] + + for t0, grp in err_df.groupby("issue_time"): + date_str = str(pd.Timestamp(t0).date()) + c = DATE_COLORS.get(date_str, color) + grp_sorted = grp.sort_values("lead_hour") + # Align to leads grid — some may be missing + curve = grp_sorted.set_index("lead_hour")["error"].reindex(leads).values + ax.plot(leads, curve, + color=c, lw=lw_thin, alpha=alpha_thin, + linestyle=linestyle, zorder=2) + all_curves.append(curve) + + if all_curves: + mean_curve = np.nanmean(np.array(all_curves), axis=0) + ax.plot(leads, mean_curve, + color=color, lw=2.8, alpha=0.95, + linestyle=linestyle, zorder=5, + label=f"{label_prefix} — mean across all init times") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--route-dir", default=DEFAULT_ROUTE_DIR) + parser.add_argument("--out-dir", default=None) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--da-name", default="routed_leadtime_da_full.parquet") + parser.add_argument("--ol-name", default="routed_leadtime_openloop_full.parquet") + args = parser.parse_args() + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + print("Loading parquets...") + da = load_parquet(os.path.join(args.route_dir, args.da_name)) + ol = load_parquet(os.path.join(args.route_dir, args.ol_name)) + obs = load_usgs(args.usgs_csv) + print(f" DA issue_times: {da['issue_time'].nunique()} " + f"OL issue_times: {ol['issue_time'].nunique()}") + + da_err = build_error_table(da, obs, HELENE_INIT_START, HELENE_INIT_END) + ol_err = build_error_table(ol, obs, HELENE_INIT_START, HELENE_INIT_END) + print(f" DA init times in Helene window: {da_err['issue_time'].nunique()}") + + leads = sorted(da_err["lead_hour"].unique()) + + # ---- Signed error plot ---- + fig, ax = plt.subplots(figsize=(12, 6)) + ax.axhline(0, color="black", lw=0.8, linestyle="--", alpha=0.5, zorder=1) + + plot_error(ax, ol_err, color="tab:gray", alpha_thin=0.12, lw_thin=0.7, + linestyle="--", label_prefix="Open-loop") + plot_error(ax, da_err, color="tab:purple", alpha_thin=0.18, lw_thin=0.8, + linestyle="-", label_prefix="DA") + + # Date-color legend patches + patches = [mpatches.Patch(color=c, label=f"Init {d}") + for d, c in DATE_COLORS.items()] + handles, labels = ax.get_legend_handles_labels() + ax.legend(handles=handles + patches, + fontsize=8, loc="upper left", frameon=True, framealpha=0.9, ncol=2) + + ax.set_xlabel("Forecast lead hour", fontsize=11) + ax.set_ylabel("Error: forecast − USGS obs (m³/s)", fontsize=11) + ax.set_xticks(np.arange(1, 19)) + ax.set_title( + "Forecast error vs lead time — per initialization time, Helene window\n" + "DA (purple solid) vs Open-loop (gray dashed) | Sep 24–28 2024 | USGS 03463300", + fontsize=11, + ) + ax.grid(True, alpha=0.25, lw=0.4) + plt.tight_layout() + out1 = os.path.join(out_dir, "forecast_error_per_init_helene.png") + plt.savefig(out1, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + # ---- Absolute error (|error|) averaged per lead — cleaner summary ---- + fig, ax = plt.subplots(figsize=(12, 6)) + + def mean_abs_error_by_lead(err_df): + return err_df.groupby("lead_hour")["error"].apply( + lambda x: float(np.nanmean(np.abs(x))) + ) + + da_mae = mean_abs_error_by_lead(da_err) + ol_mae = mean_abs_error_by_lead(ol_err) + + ax.plot(da_mae.index, da_mae.values, + color="tab:purple", lw=2.4, marker="o", label="DA — mean |error|") + ax.plot(ol_mae.index, ol_mae.values, + color="tab:gray", lw=2.4, marker="s", linestyle="--", + label="Open-loop — mean |error|") + + ax.set_xlabel("Forecast lead hour", fontsize=11) + ax.set_ylabel("Mean |error| vs USGS obs (m³/s)", fontsize=11) + ax.set_xticks(np.arange(1, 19)) + ax.set_title( + "Mean absolute forecast error vs lead time — Helene window\n" + "DA (purple) vs Open-loop (gray) | Sep 24–28 2024 | USGS 03463300", + fontsize=11, + ) + ax.legend(fontsize=10, loc="upper left", frameon=True, framealpha=0.9) + ax.grid(True, alpha=0.25, lw=0.4) + plt.tight_layout() + out2 = os.path.join(out_dir, "forecast_mae_per_lead_helene.png") + plt.savefig(out2, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4a_error_decay/plot_lead_time_decay.py b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4a_error_decay/plot_lead_time_decay.py new file mode 100644 index 00000000..2b33fa82 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4a_error_decay/plot_lead_time_decay.py @@ -0,0 +1,160 @@ +""" +Catchment-level error-vs-lead-time decay curve + ensemble spread by lead. + +Two-panel figure for ONE catchment: + TOP — RMSE of ensemble-mean forecast at each lead hour (1..18) vs the + catchment's kriging observation. DA solid, open-loop dashed, with + a shaded band showing the min/max of per-member RMSE. + BOTTOM — Mean ensemble spread (std-dev across 20 members, averaged across + all issue times) at each lead hour. Tells you whether forcing + perturbation alone keeps the forecast ensemble diverse during the + 18-hour free-run — useful for inspecting individual members + without needing the full spaghetti view. + +This is the catchment-level analog of the gauge-level decay curve. +Routing to the gauge is a separate post-step (route_lead_time_forecasts.py ++ a gauge-level decay script); this script lets us look at the catchment-level +signal without T-route in the loop. + +Inputs (from run_lead_time_forecast_sweep.py): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv +Obs (from production EnKF run): + //_test_results.csv + +Output: + //_lead_time_decay.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" + +DEFAULT_LEADTIME_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast" +DEFAULT_DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" + +# Module-level globals set by main() so the helpers see them +CAT = "cat-1016300" +LEADTIME_DIR = DEFAULT_LEADTIME_DIR +DA_DIR = DEFAULT_DA_DIR +OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay.png") + + +def load_forecasts(path): + """Return (df, member_cols).""" + if not os.path.exists(path): + raise FileNotFoundError(path) + df = pd.read_csv(path, parse_dates=['issue_time', 'valid_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + return df, member_cols + + +def load_obs(): + obs_path = os.path.join(DA_DIR, CAT, f"{CAT}_test_results.csv") + df = pd.read_csv(obs_path, parse_dates=['date']) + return df.set_index('date')['obs_mm_h'] + + +def metrics_by_lead(df, member_cols, obs_series): + """For each lead hour, return (rmse_mean, rmse_min_member, rmse_max_member, + mean_ensemble_std).""" + df = df.copy() + df['obs'] = df['valid_time'].map(obs_series) + df = df.dropna(subset=['obs']) + df['ens_mean'] = df[member_cols].mean(axis=1) + df['ens_std'] = df[member_cols].std(axis=1) + + leads = sorted(df['lead_hour'].unique()) + rmse_mean = np.full(len(leads), np.nan) + rmse_min_member = np.full(len(leads), np.nan) + rmse_max_member = np.full(len(leads), np.nan) + mean_ens_std = np.full(len(leads), np.nan) + for i, L in enumerate(leads): + sub = df[df['lead_hour'] == L] + if len(sub) == 0: + continue + err = sub['ens_mean'].values - sub['obs'].values + rmse_mean[i] = float(np.sqrt(np.mean(err ** 2))) + per_member_rmse = np.array([ + np.sqrt(np.mean((sub[c].values - sub['obs'].values) ** 2)) + for c in member_cols + ]) + rmse_min_member[i] = float(np.nanmin(per_member_rmse)) + rmse_max_member[i] = float(np.nanmax(per_member_rmse)) + mean_ens_std[i] = float(np.nanmean(sub['ens_std'].values)) + return (np.asarray(leads), rmse_mean, rmse_min_member, + rmse_max_member, mean_ens_std) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', default="cat-1016300") + parser.add_argument('--leadtime-dir', default=DEFAULT_LEADTIME_DIR, + help='Dir holding /_lead_time_forecasts_{da,openloop}.csv') + parser.add_argument('--da-dir', default=DEFAULT_DA_DIR, + help='Dir holding /_test_results.csv (for obs)') + args = parser.parse_args() + + global CAT, LEADTIME_DIR, DA_DIR, OUT_PNG + CAT = args.cat_id + LEADTIME_DIR = args.leadtime_dir + DA_DIR = args.da_dir + OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay.png") + + da_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_da.csv") + ol_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_openloop.csv") + df_da, m_cols = load_forecasts(da_path) + df_ol, _ = load_forecasts(ol_path) + obs_series = load_obs() + + leads_da, rmse_da, rmse_da_lo, rmse_da_hi, std_da = metrics_by_lead( + df_da, m_cols, obs_series) + leads_ol, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = metrics_by_lead( + df_ol, m_cols, obs_series) + + fig, (ax_top, ax_bot) = plt.subplots(2, 1, figsize=(10, 10), sharex=True, + gridspec_kw={"height_ratios": [1.4, 1.0]}) + + # ----- TOP: RMSE decay curve ----- + ax_top.fill_between(leads_da, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2, + label="DA — member spread (min/max RMSE)") + ax_top.fill_between(leads_ol, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2, + label="Open-loop — member spread") + ax_top.plot(leads_da, rmse_da, color=DA_COLOR, lw=2.4, marker='o', + zorder=4, label="DA on (ensemble-mean RMSE)") + ax_top.plot(leads_ol, rmse_ol, color=OL_COLOR, lw=2.4, marker='s', + linestyle='--', zorder=4, label="Open-loop (ensemble-mean RMSE)") + ax_top.set_ylabel("RMSE vs Qkrig obs (mm/h)", fontsize=11) + ax_top.set_title( + f"Forecast lead-time error decay — {CAT}\n" + f"Issue times pooled across test period (Oct 2023 – Oct 2024)", + fontsize=12, + ) + ax_top.grid(True, alpha=0.25) + ax_top.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + # ----- BOTTOM: ensemble spread by lead hour ----- + ax_bot.plot(leads_da, std_da, color=DA_COLOR, lw=2.2, marker='o', + label="DA — mean ensemble std") + ax_bot.plot(leads_ol, std_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', label="Open-loop — mean ensemble std") + ax_bot.set_xlabel("Forecast lead time (hours)", fontsize=11) + ax_bot.set_ylabel("Mean ensemble std (mm/h)", fontsize=11) + ax_bot.set_xticks(np.arange(1, 19)) + ax_bot.grid(True, alpha=0.25) + ax_bot.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + plt.tight_layout() + plt.savefig(OUT_PNG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4a_error_decay/plot_lead_time_decay_by_regime.py b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4a_error_decay/plot_lead_time_decay_by_regime.py new file mode 100644 index 00000000..ecf10408 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4a_error_decay/plot_lead_time_decay_by_regime.py @@ -0,0 +1,200 @@ +""" +Lead-time decay curve, split by flow regime at issue time. + +The pooled lead-time curve (plot_lead_time_decay.py) showed DA losing to +open-loop across most lead hours, but the test period is ~99% low-flow. +This script splits the same forecast CSVs by the flow regime at the issue +time t0, so we can see whether DA helps when it matters (storms / Helene) +and hurts when it doesn't (low flow). + +Three regimes are partitioned on obs(t0) (kriging obs at the issue time): + + helene issue_time ∈ [2024-09-24, 2024-09-28] (the 5-day Helene window) + storm obs(t0) > 0.5 mm/h (high-flow issue times) + low_flow obs(t0) < 0.1 mm/h (low-flow issue times) + +Same 2-panel layout per regime: RMSE vs lead (top), ensemble spread vs lead +(bottom). Three regime columns × 2 metric rows in one figure. + +Inputs (from run_lead_time_forecast_sweep.py — no re-run needed): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv +Obs (from production EnKF): + //_test_results.csv + +Output: + //_lead_time_decay_by_regime.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" + +DEFAULT_LEADTIME_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast" +DEFAULT_DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" + +# Module-level globals set by main() so helpers see them +CAT = "cat-1016300" +LEADTIME_DIR = DEFAULT_LEADTIME_DIR +DA_DIR = DEFAULT_DA_DIR +OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay_by_regime.png") + +# Regime definitions +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-28 23:00:00") +STORM_OBS_THRESHOLD = 0.5 # mm/h +LOWFLOW_OBS_THRESHOLD = 0.1 # mm/h + + +def load_forecasts(path): + if not os.path.exists(path): + raise FileNotFoundError(path) + df = pd.read_csv(path, parse_dates=['issue_time', 'valid_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + return df, member_cols + + +def load_obs(): + obs_path = os.path.join(DA_DIR, CAT, f"{CAT}_test_results.csv") + df = pd.read_csv(obs_path, parse_dates=['date']) + return df.set_index('date')['obs_mm_h'] + + +def regime_mask(issue_times, obs_at_issue, regime): + """Boolean mask over issue_times for the named regime.""" + if regime == "helene": + return (issue_times >= HELENE_START) & (issue_times <= HELENE_END) + if regime == "storm": + return obs_at_issue > STORM_OBS_THRESHOLD + if regime == "low_flow": + return obs_at_issue < LOWFLOW_OBS_THRESHOLD + raise ValueError(regime) + + +def metrics_by_lead(df, member_cols, obs_series, issue_mask): + """Compute (leads, rmse_mean, rmse_min, rmse_max, mean_std) restricted + to issue times where issue_mask is True.""" + df = df.copy() + df = df[df['issue_time'].isin(issue_mask)] + if len(df) == 0: + return None + df['obs'] = df['valid_time'].map(obs_series) + df = df.dropna(subset=['obs']) + if len(df) == 0: + return None + df['ens_mean'] = df[member_cols].mean(axis=1) + df['ens_std'] = df[member_cols].std(axis=1) + + leads = sorted(df['lead_hour'].unique()) + rmse_mean = np.full(len(leads), np.nan) + rmse_lo = np.full(len(leads), np.nan) + rmse_hi = np.full(len(leads), np.nan) + std_mean = np.full(len(leads), np.nan) + for i, L in enumerate(leads): + sub = df[df['lead_hour'] == L] + if len(sub) == 0: + continue + err = sub['ens_mean'].values - sub['obs'].values + rmse_mean[i] = float(np.sqrt(np.mean(err ** 2))) + per_member = np.array([ + np.sqrt(np.mean((sub[c].values - sub['obs'].values) ** 2)) + for c in member_cols + ]) + rmse_lo[i] = float(np.nanmin(per_member)) + rmse_hi[i] = float(np.nanmax(per_member)) + std_mean[i] = float(np.nanmean(sub['ens_std'].values)) + return np.asarray(leads), rmse_mean, rmse_lo, rmse_hi, std_mean + + +def plot_pair(ax_rmse, ax_std, da_metrics, ol_metrics, + regime_label, n_issue_times): + if da_metrics is None or ol_metrics is None: + ax_rmse.set_title(f"{regime_label}\n(no issue times match)", fontsize=11) + ax_rmse.axis('off'); ax_std.axis('off') + return + leads, rmse_da, rmse_da_lo, rmse_da_hi, std_da = da_metrics + _, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = ol_metrics + + # RMSE panel + ax_rmse.fill_between(leads, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2) + ax_rmse.fill_between(leads, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2) + ax_rmse.plot(leads, rmse_da, color=DA_COLOR, lw=2.2, marker='o', + zorder=4, label="DA on") + ax_rmse.plot(leads, rmse_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', zorder=4, label="Open-loop") + ax_rmse.set_title(f"{regime_label}\n({n_issue_times} issue times)", fontsize=11) + ax_rmse.set_ylabel("RMSE (mm/h)", fontsize=10) + ax_rmse.grid(True, alpha=0.25) + ax_rmse.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + # Spread panel + ax_std.plot(leads, std_da, color=DA_COLOR, lw=2.0, marker='o', label="DA") + ax_std.plot(leads, std_ol, color=OL_COLOR, lw=2.0, marker='s', + linestyle='--', label="Open-loop") + ax_std.set_xlabel("Forecast lead time (hours)", fontsize=10) + ax_std.set_ylabel("Mean ensemble std (mm/h)", fontsize=10) + ax_std.set_xticks(np.arange(1, 19)) + ax_std.grid(True, alpha=0.25) + ax_std.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', default="cat-1016300") + parser.add_argument('--leadtime-dir', default=DEFAULT_LEADTIME_DIR, + help='Dir holding /_lead_time_forecasts_{da,openloop}.csv') + parser.add_argument('--da-dir', default=DEFAULT_DA_DIR, + help='Dir holding /_test_results.csv (for obs)') + args = parser.parse_args() + + global CAT, LEADTIME_DIR, DA_DIR, OUT_PNG + CAT = args.cat_id + LEADTIME_DIR = args.leadtime_dir + DA_DIR = args.da_dir + OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay_by_regime.png") + + da_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_da.csv") + ol_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_openloop.csv") + df_da, m_cols = load_forecasts(da_path) + df_ol, _ = load_forecasts(ol_path) + obs_series = load_obs() + + # Unique issue times across the run, with obs at each + issue_times = pd.to_datetime(sorted(df_da['issue_time'].unique())) + obs_at_issue = pd.Series(issue_times, index=issue_times).map(obs_series) + + regimes = [ + ("helene", "Helene window (Sep 24–28 2024)"), + ("storm", f"Storm hours (obs(t0) > {STORM_OBS_THRESHOLD} mm/h)"), + ("low_flow", f"Low flow (obs(t0) < {LOWFLOW_OBS_THRESHOLD} mm/h)"), + ] + + fig, axes = plt.subplots(2, 3, figsize=(18, 10), + gridspec_kw={"height_ratios": [1.4, 1.0]}) + for col, (key, label) in enumerate(regimes): + mask = regime_mask(issue_times, obs_at_issue, key) + kept = issue_times[mask] + kept_set = set(pd.to_datetime(kept)) + da_m = metrics_by_lead(df_da, m_cols, obs_series, kept_set) + ol_m = metrics_by_lead(df_ol, m_cols, obs_series, kept_set) + plot_pair(axes[0, col], axes[1, col], da_m, ol_m, label, len(kept)) + + fig.suptitle( + f"Lead-time error decay by flow regime — {CAT}\n" + f"Same forecast CSVs as the pooled view, partitioned on obs(t0).", + fontsize=13, y=0.995, + ) + plt.tight_layout(rect=(0, 0, 1, 0.97)) + plt.savefig(OUT_PNG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4a_error_decay/plot_lead_time_decay_gauge.py b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4a_error_decay/plot_lead_time_decay_gauge.py new file mode 100644 index 00000000..71ce54fc --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4a_error_decay/plot_lead_time_decay_gauge.py @@ -0,0 +1,311 @@ +""" +Gauge-level lead-time forecast decay curve. + +Routes 21-catchment forecast CSVs through T-route to USGS gauge 03463300 and +plots RMSE vs lead time (DA solid, open-loop dashed) against USGS observed Q. +Same 2-panel layout as the catchment-level plot, plus a regime split for the +Helene window. + +The routing step is done separately by route_lead_time_forecasts.py; +this script just reads the resulting parquets and produces the plot. + +Inputs: + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + Each: one row per (issue_time, lead_hour); member columns hold q_gauge_m3s. + (Long-format with explicit `member` and `q_gauge_m3s` columns also supported.) +Obs: + /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + USGS hourly Q in m³/s at gauge 03463300 (South Toe River near Celo, NC). + +Output: + /lead_time_decay_gauge_pooled.png + /lead_time_decay_gauge_by_regime.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" +OBS_COLOR = "black" + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/leadtime_troute_routing" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +# Watershed area at gauge 03463300 (sum of 21 catchments per GPKG hydrofabric) +WATERSHED_AREA_KM2 = 113.18 +# 1 mm/h depth × 113.18 km² = 113.18e3 m³/h = 31.439 m³/s +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-28 23:00:00") +STORM_OBS_THRESHOLD_M3S = 20.0 # gauge-level storm threshold (was 50 — too high + # for this 113 km² basin; only 2 issue times qualified) +LOWFLOW_OBS_THRESHOLD_M3S = 5.0 # gauge-level low-flow threshold + +USGS_HELENE_PEAK_M3S = 1886.0 # reference Sep 27 14:00 + + +def load_parquet_long(path): + """Load a routed parquet and normalize to long format: + columns = issue_time, lead_hour, member, q_gauge_m3s. + + Handles both wide format (member_00..member_19 columns) and long format + (explicit `member` + `q_gauge_m3s` columns). Auto-detects. + """ + if not os.path.exists(path): + raise FileNotFoundError(path) + df = pd.read_parquet(path) + + # Long format detection + if 'q_gauge_m3s' in df.columns and 'member' in df.columns: + df['issue_time'] = pd.to_datetime(df['issue_time']) + return df[['issue_time', 'lead_hour', 'member', 'q_gauge_m3s']] + + # Some variants might call it 'q_m3s' or similar + q_col_candidates = [c for c in df.columns if c.lower() in + ('q_gauge_m3s', 'q_m3s', 'q_gauge', 'q')] + if 'member' in df.columns and q_col_candidates: + qc = q_col_candidates[0] + df['issue_time'] = pd.to_datetime(df['issue_time']) + out = df[['issue_time', 'lead_hour', 'member', qc]].copy() + out = out.rename(columns={qc: 'q_gauge_m3s'}) + return out + + # Wide format: member_00..member_19 columns + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + if not member_cols: + raise ValueError(f"Couldn't identify member columns in {path}. " + f"Columns present: {df.columns.tolist()}") + df['issue_time'] = pd.to_datetime(df['issue_time']) + keep = ['issue_time', 'lead_hour'] + if 'valid_time' in df.columns: + keep.append('valid_time') + long = df[keep + member_cols].melt( + id_vars=keep, + value_vars=member_cols, + var_name='member', + value_name='q_gauge_m3s', + ) + return long + + +def load_usgs_obs(usgs_csv): + """Load USGS gauge obs and ensure m³/s. + + The CSV at gauge 03463300 has the column `QObs(mm/h)` (catchment-averaged + depth, not gauge discharge in m³/s). If the column name contains 'mm', we + convert mm/h → m³/s by multiplying by the watershed area factor: + m³/s = mm/h × WATERSHED_AREA_KM2 × 1000 / 3600 + Verified: 59.978 mm/h × 113.18 km² = 1886 m³/s (USGS Helene peak). + """ + df = pd.read_csv(usgs_csv) + date_candidates = [c for c in df.columns if c.lower() in + ('datetime', 'date', 'time', 'timestamp')] + q_candidates = [c for c in df.columns if 'q' in c.lower() + or 'flow' in c.lower() or 'discharge' in c.lower()] + if not date_candidates or not q_candidates: + raise ValueError(f"Could not identify date/Q columns in {usgs_csv}. " + f"Columns: {df.columns.tolist()}") + dc, qc = date_candidates[0], q_candidates[0] + df[dc] = pd.to_datetime(df[dc]) + series = df.set_index(dc)[qc].astype(float) + if 'mm' in qc.lower(): + series = series * MM_H_TO_M3_S + print(f" Converted obs from mm/h (column '{qc}') to m³/s " + f"using area = {WATERSHED_AREA_KM2} km² (× {MM_H_TO_M3_S:.4f})") + else: + print(f" Loaded obs as-is from column '{qc}' (assumed m³/s)") + return series + + +def metrics_by_lead(df_long, obs_series, issue_mask=None): + """Compute per-lead-hour: ensemble-mean RMSE, per-member min/max RMSE, + mean ensemble std. df_long must have issue_time, lead_hour, member, q_gauge_m3s. + """ + df = df_long.copy() + if issue_mask is not None: + df = df[df['issue_time'].isin(issue_mask)] + if len(df) == 0: + return None + df['valid_time'] = df['issue_time'] + pd.to_timedelta(df['lead_hour'], unit='h') + df['obs'] = df['valid_time'].map(obs_series) + df = df.dropna(subset=['obs']) + if len(df) == 0: + return None + + # Ensemble mean per (issue_time, lead_hour) + grouped = df.groupby(['issue_time', 'lead_hour']) + ens_mean = grouped['q_gauge_m3s'].mean().reset_index(name='ens_mean') + ens_std = grouped['q_gauge_m3s'].std().reset_index(name='ens_std') + obs_per = grouped['obs'].first().reset_index(name='obs') + panel = ens_mean.merge(ens_std, on=['issue_time', 'lead_hour']).merge( + obs_per, on=['issue_time', 'lead_hour']) + + leads = sorted(panel['lead_hour'].unique()) + rmse_mean = np.full(len(leads), np.nan) + rmse_lo = np.full(len(leads), np.nan) + rmse_hi = np.full(len(leads), np.nan) + std_mean = np.full(len(leads), np.nan) + for i, L in enumerate(leads): + psub = panel[panel['lead_hour'] == L] + if len(psub) == 0: + continue + err = psub['ens_mean'].values - psub['obs'].values + rmse_mean[i] = float(np.sqrt(np.mean(err ** 2))) + # Per-member RMSE — recompute from the underlying long frame + dsub = df[df['lead_hour'] == L] + member_rmses = [] + for m, g in dsub.groupby('member'): + err_m = g['q_gauge_m3s'].values - g['obs'].values + if len(err_m) >= 1: + member_rmses.append(float(np.sqrt(np.mean(err_m ** 2)))) + if member_rmses: + rmse_lo[i] = float(np.nanmin(member_rmses)) + rmse_hi[i] = float(np.nanmax(member_rmses)) + std_mean[i] = float(np.nanmean(psub['ens_std'].values)) + return np.asarray(leads), rmse_mean, rmse_lo, rmse_hi, std_mean + + +def plot_two_panel(out_path, leads_da, da_metrics, leads_ol, ol_metrics, title): + fig, (ax_top, ax_bot) = plt.subplots(2, 1, figsize=(10, 10), sharex=True, + gridspec_kw={"height_ratios": [1.4, 1.0]}) + leads, rmse_da, rmse_da_lo, rmse_da_hi, std_da = da_metrics + _, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = ol_metrics + + ax_top.fill_between(leads, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2, + label="DA — member spread (min/max RMSE)") + ax_top.fill_between(leads, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2, + label="Open-loop — member spread") + ax_top.plot(leads, rmse_da, color=DA_COLOR, lw=2.4, marker='o', + zorder=4, label="DA on (ensemble-mean RMSE)") + ax_top.plot(leads, rmse_ol, color=OL_COLOR, lw=2.4, marker='s', + linestyle='--', zorder=4, label="Open-loop (ensemble-mean RMSE)") + ax_top.set_ylabel("RMSE vs USGS obs (m³/s)", fontsize=11) + ax_top.set_title(title, fontsize=12) + ax_top.grid(True, alpha=0.25) + ax_top.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + ax_bot.plot(leads, std_da, color=DA_COLOR, lw=2.2, marker='o', + label="DA — mean ensemble std") + ax_bot.plot(leads, std_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', label="Open-loop — mean ensemble std") + ax_bot.set_xlabel("Forecast lead time (hours)", fontsize=11) + ax_bot.set_ylabel("Mean ensemble std (m³/s)", fontsize=11) + ax_bot.set_xticks(np.arange(1, 19)) + ax_bot.grid(True, alpha=0.25) + ax_bot.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + plt.tight_layout() + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + +def plot_three_panel_regime(out_path, regime_metrics): + fig, axes = plt.subplots(2, 3, figsize=(18, 10), + gridspec_kw={"height_ratios": [1.4, 1.0]}) + for col, (label, n_issue, da_m, ol_m) in enumerate(regime_metrics): + ax_rmse = axes[0, col] + ax_std = axes[1, col] + if da_m is None or ol_m is None: + ax_rmse.set_title(f"{label}\n(no issue times match)", fontsize=11) + ax_rmse.axis('off'); ax_std.axis('off') + continue + leads, rmse_da, rmse_da_lo, rmse_da_hi, std_da = da_m + _, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = ol_m + ax_rmse.fill_between(leads, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2) + ax_rmse.fill_between(leads, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2) + ax_rmse.plot(leads, rmse_da, color=DA_COLOR, lw=2.2, marker='o', + zorder=4, label="DA on") + ax_rmse.plot(leads, rmse_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', zorder=4, label="Open-loop") + ax_rmse.set_title(f"{label}\n({n_issue} issue times)", fontsize=11) + ax_rmse.set_ylabel("RMSE (m³/s)", fontsize=10) + ax_rmse.grid(True, alpha=0.25) + ax_rmse.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + ax_std.plot(leads, std_da, color=DA_COLOR, lw=2.0, marker='o', label="DA") + ax_std.plot(leads, std_ol, color=OL_COLOR, lw=2.0, marker='s', + linestyle='--', label="Open-loop") + ax_std.set_xlabel("Forecast lead time (hours)", fontsize=10) + ax_std.set_ylabel("Mean ensemble std (m³/s)", fontsize=10) + ax_std.set_xticks(np.arange(1, 19)) + ax_std.grid(True, alpha=0.25) + ax_std.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + fig.suptitle("Gauge-level lead-time error decay by regime — USGS 03463300", + fontsize=13, y=0.995) + plt.tight_layout(rect=(0, 0, 1, 0.97)) + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--route-dir', default=DEFAULT_ROUTE_DIR, + help='Dir holding routed_leadtime_{da,openloop}_full.parquet') + parser.add_argument('--out-dir', default=None, + help='Where to write the output PNGs. Defaults to --route-dir ' + '(which may not be writable if owned by another user — ' + 'pass an explicit path then).') + parser.add_argument('--usgs-csv', default=DEFAULT_USGS_CSV) + parser.add_argument('--da-name', default='routed_leadtime_da_full.parquet') + parser.add_argument('--ol-name', default='routed_leadtime_openloop_full.parquet') + args = parser.parse_args() + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + da_path = os.path.join(args.route_dir, args.da_name) + ol_path = os.path.join(args.route_dir, args.ol_name) + df_da = load_parquet_long(da_path) + df_ol = load_parquet_long(ol_path) + obs_series = load_usgs_obs(args.usgs_csv) + + print(f"DA rows: {len(df_da):,} | unique issue_times: {df_da['issue_time'].nunique()}") + print(f"OL rows: {len(df_ol):,} | unique issue_times: {df_ol['issue_time'].nunique()}") + print(f"USGS obs range: {obs_series.index.min()} .. {obs_series.index.max()} " + f"({len(obs_series):,} hours)") + print(f"USGS Helene peak (reference): {USGS_HELENE_PEAK_M3S:.0f} m³/s") + + # ----- Pooled (all issue times) ----- + da_pooled = metrics_by_lead(df_da, obs_series) + ol_pooled = metrics_by_lead(df_ol, obs_series) + out_pooled = os.path.join(out_dir, "lead_time_decay_gauge_pooled.png") + plot_two_panel(out_pooled, da_pooled[0], da_pooled, ol_pooled[0], ol_pooled, + title="Gauge-level lead-time decay — USGS 03463300\n" + "All issue times pooled (Oct 2023 – Oct 2024)") + + # ----- Regime split: Helene, storm, low-flow at gauge ----- + issue_times = pd.to_datetime(sorted(df_da['issue_time'].unique())) + obs_at_issue = pd.Series(issue_times, index=issue_times).map(obs_series) + + helene_mask = (issue_times >= HELENE_START) & (issue_times <= HELENE_END) + storm_mask = obs_at_issue > STORM_OBS_THRESHOLD_M3S + lowflow_mask = obs_at_issue < LOWFLOW_OBS_THRESHOLD_M3S + + regime_metrics = [] + for label, mask in [ + (f"Helene window (Sep 24–28 2024)", helene_mask), + (f"Storm hours (obs(t0) > {STORM_OBS_THRESHOLD_M3S:.0f} m³/s)", storm_mask), + (f"Low flow (obs(t0) < {LOWFLOW_OBS_THRESHOLD_M3S:.0f} m³/s)", lowflow_mask), + ]: + kept = set(pd.to_datetime(issue_times[mask])) + da_m = metrics_by_lead(df_da, obs_series, kept) + ol_m = metrics_by_lead(df_ol, obs_series, kept) + regime_metrics.append((label, len(kept), da_m, ol_m)) + + out_regime = os.path.join(out_dir, "lead_time_decay_gauge_by_regime.png") + plot_three_panel_regime(out_regime, regime_metrics) + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4a_error_decay/run_4a_f3.sh b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4a_error_decay/run_4a_f3.sh new file mode 100644 index 00000000..bb5a9449 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4a_error_decay/run_4a_f3.sh @@ -0,0 +1,80 @@ +#!/usr/bin/env bash +# F3 dynamic Vrugt seeded (20pct gauge holdout) — 4a lead-time decay plots. +# +# Runs plot scripts per catchment: +# plot_lead_time_decay.py — pooled RMSE vs lead (catchment-level) +# plot_lead_time_decay_by_regime.py — same split by flow regime +# +# Gauge-level decay plots run once if routed parquets exist: +# plot_lead_time_decay_gauge.py — gauge-level (requires routed parquets) +# plot_forecast_error_fixed_target.py +# plot_forecast_error_per_init.py +# +# Usage: +# bash run_4a_f3.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +LEADTIME_DIR=/mnt/disk2/1400_sites_helene/da_forecast_dynamic_vrugt_seeded +DA_DIR=/mnt/disk2/1400_sites_helene/da_arms_dynamic_vrugt_seeded +OBS_DIR=/mnt/disk2/1400_sites_helene/catchment_ts_03463300_spliced_dyn_helene +ROUTE_DIR=/mnt/disk2/suma_helen_poster/leadtime_troute_routing_f3 +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +echo "[F3-4a] Lead-time decay plots — leadtime dir: $LEADTIME_DIR" + +for CAT in "${CATS[@]}"; do + DA_CSV="$LEADTIME_DIR/$CAT/${CAT}_lead_time_forecasts_da.csv" + if [ ! -f "$DA_CSV" ]; then + echo " [$CAT] No lead-time CSV — skipping" + continue + fi + + echo " [$CAT] Plotting pooled decay..." + $TROUTE "$SCRIPT_DIR/plot_lead_time_decay.py" \ + --cat-id "$CAT" \ + --leadtime-dir "$LEADTIME_DIR" \ + --da-dir "$DA_DIR" \ + --obs-dir "$OBS_DIR" + + echo " [$CAT] Plotting decay by regime..." + $TROUTE "$SCRIPT_DIR/plot_lead_time_decay_by_regime.py" \ + --cat-id "$CAT" \ + --leadtime-dir "$LEADTIME_DIR" \ + --da-dir "$DA_DIR" \ + --obs-dir "$OBS_DIR" +done + +# Gauge-level decay (needs routed parquets — run after route_leadtime_f3.sh) +if [ -f "$ROUTE_DIR/routed_leadtime_da_full.parquet" ]; then + echo "[F3-4a] Gauge-level decay (routed)..." + $TROUTE "$SCRIPT_DIR/plot_lead_time_decay_gauge.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" + + echo "[F3-4a] Fixed-target error plots..." + $TROUTE "$SCRIPT_DIR/plot_forecast_error_fixed_target.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" + + echo "[F3-4a] Per-init error plots..." + $TROUTE "$SCRIPT_DIR/plot_forecast_error_per_init.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" +else + echo "[F3-4a] Routed parquets not found — skipping gauge-level and error plots." + echo " Run route_leadtime_f3.sh first, then re-run this script." +fi + +echo "[F3-4a] Done." diff --git a/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4b_ensemble_vs_obs/plot_crossed_ensemble.py b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4b_ensemble_vs_obs/plot_crossed_ensemble.py new file mode 100644 index 00000000..7cb00a93 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4b_ensemble_vs_obs/plot_crossed_ensemble.py @@ -0,0 +1,229 @@ +""" +plot_crossed_ensemble.py -- 4b: 600-member crossed ensemble vs USGS obs. + +Loads the crossed ensemble parquet produced by run_crossed_ensemble.py: + //_crossed_ensemble.parquet + Columns: issue_time, lead_hour, member_0000..member_0599 + +Projects all members to valid_time = issue_time + lead_hour hours. +Aggregates at each valid_time across ALL members from ALL issue times +to build a probabilistic envelope (5th/25th/50th/75th/95th percentiles). + +Two output figures: + 4b_crossed_ensemble_helene.png + -- Full Helene window (Sep 24-30) probabilistic envelope + USGS obs. + Percentile shading in two layers (5-95 outer, 25-75 inner). + 4b_crossed_ensemble_peak.png + -- Zoomed to Helene peak (Sep 26 12 UTC -> Sep 28 06 UTC), + same shading + USGS obs + individual spaghetti from Sep 27 inits. + +Usage: + python3 plot_crossed_ensemble.py \\ + --ensemble-dir /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble \\ + --cat-id cat-1016300 \\ + --usgs-csv /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv \\ + --out-dir /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble/cat-1016300 +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import matplotlib.patches as mpatches + +DEFAULT_ENSEMBLE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" +DEFAULT_CAT_ID = "cat-1016300" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +PLOT_START = pd.Timestamp("2024-09-24 00:00:00") +PLOT_END = pd.Timestamp("2024-09-28 18:00:00") # cut trailing spike +PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +PEAK_END = pd.Timestamp("2024-09-28 06:00:00") +HELENE_PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +HELENE_PEAK_END = pd.Timestamp("2024-09-28 00:00:00") + + +def load_ensemble(path): + df = pd.read_parquet(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + return df, member_cols + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def build_percentile_envelope(df, member_cols, t_start, t_end): + """ + For each valid_time in [t_start, t_end], stack all member values + from all issue times and compute percentiles. + Returns DataFrame indexed by valid_time with columns p05, p25, p50, p75, p95. + """ + mask = (df["valid_time"] >= t_start) & (df["valid_time"] <= t_end) + sub = df[mask].copy() + sub["valid_time"] = sub["valid_time"].dt.floor("1h") + + records = [] + for vt, grp in sub.groupby("valid_time"): + vals = grp[member_cols].to_numpy(dtype=float).flatten() + vals = vals[~np.isnan(vals)] * MM_H_TO_M3_S + if len(vals) == 0: + continue + records.append({ + "valid_time": vt, + "p05": np.percentile(vals, 5), + "p25": np.percentile(vals, 25), + "p50": np.percentile(vals, 50), + "p75": np.percentile(vals, 75), + "p95": np.percentile(vals, 95), + "n": len(vals), + }) + return pd.DataFrame(records).set_index("valid_time").sort_index() + + +def _add_helene_band(ax, label=True): + ax.axvspan(HELENE_PEAK_START, HELENE_PEAK_END, + color="salmon", alpha=0.12, zorder=0) + if label: + ax.text(HELENE_PEAK_START + pd.Timedelta(hours=6), 0.97, "Helene peak", + transform=ax.get_xaxis_transform(), + fontsize=9, color="firebrick", ha="left", va="top") + + +def _format_xaxis(ax, major_interval=1, minor_hours=6): + ax.xaxis.set_major_locator(mdates.DayLocator(interval=major_interval)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.xaxis.set_minor_locator(mdates.HourLocator(interval=minor_hours)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.2, lw=0.4) + + +def plot_envelope(ax, env, obs_series, t_start, t_end, + ens_color="tab:blue", obs_color="black", helene_label=True): + _add_helene_band(ax, label=helene_label) + + # Outer band: 5-95 + ax.fill_between(env.index, env["p05"], env["p95"], + color=ens_color, alpha=0.15, zorder=2, + label="5th-95th percentile") + # Inner band: 25-75 + ax.fill_between(env.index, env["p25"], env["p75"], + color=ens_color, alpha=0.30, zorder=3, + label="25th-75th percentile") + # Median + ax.plot(env.index, env["p50"], + color=ens_color, lw=2.2, zorder=5, + label="Ensemble median (600 members)") + + # USGS obs + obs_w = obs_series.loc[t_start:t_end] + ax.plot(obs_w.index, obs_w.values, + color=obs_color, lw=1.8, zorder=6, label="USGS obs") + + ax.set_xlim(t_start, t_end) + ax.set_ylabel("Discharge (m3/s)", fontsize=11) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--ensemble-dir", default=DEFAULT_ENSEMBLE_DIR) + parser.add_argument("--cat-id", default=DEFAULT_CAT_ID) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--out-dir", default=None) + args = parser.parse_args() + + cat_dir = os.path.join(args.ensemble_dir, args.cat_id) + out_dir = args.out_dir or cat_dir + os.makedirs(out_dir, exist_ok=True) + + parquet_path = os.path.join(cat_dir, f"{args.cat_id}_crossed_ensemble.parquet") + print(f"Loading crossed ensemble: {parquet_path}") + df, member_cols = load_ensemble(parquet_path) + obs = load_usgs(args.usgs_csv) + print(f" {df['issue_time'].nunique()} issue times, " + f"{len(member_cols)} members per row") + + # ------------------------------------------------------------------ # + # Figure 4b-full: full Helene window probabilistic envelope # + # ------------------------------------------------------------------ # + print("Building percentile envelope (full window)...") + env_full = build_percentile_envelope(df, member_cols, PLOT_START, PLOT_END) + print(f" Valid times with data: {len(env_full)}, " + f"median members per timestep: {env_full['n'].median():.0f}") + + fig, ax = plt.subplots(figsize=(17, 7)) + plot_envelope(ax, env_full, obs, PLOT_START, PLOT_END) + ax.set_xlabel("Date", fontsize=11) + ax.set_title( + f"4b -- 600-member crossed ensemble (30 met x 20 hydro-state) | {args.cat_id}\n" + "Shaded: 5-95th (light) and 25-75th (dark) percentile across all members " + "and issue times | Sep 24-30 2024 | USGS 03463300", + fontsize=11, + ) + ax.legend(fontsize=10, loc="upper left", frameon=True, framealpha=0.9) + _format_xaxis(ax) + plt.tight_layout() + out1 = os.path.join(out_dir, f"{args.cat_id}_4b_crossed_ensemble_helene.png") + plt.savefig(out1, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + # ------------------------------------------------------------------ # + # Figure 4b-peak: zoomed to Helene peak window # + # ------------------------------------------------------------------ # + print("Building percentile envelope (peak window)...") + env_peak = build_percentile_envelope(df, member_cols, PEAK_START, PEAK_END) + + fig, (ax_main, ax_zoom) = plt.subplots( + 2, 1, figsize=(14, 11), + gridspec_kw={"height_ratios": [1, 1.4]}, sharex=False) + + # Top panel: full window context (smaller) + plot_envelope(ax_main, env_full, obs, PLOT_START, PLOT_END, helene_label=False) + ax_main.set_title("Full Helene window context (Sep 24 - Sep 28 18 UTC)", fontsize=10) + _format_xaxis(ax_main) + # Shade the zoom region on context panel + ax_main.axvspan(PEAK_START, PEAK_END, color="gold", alpha=0.18, zorder=1) + ax_main.text(PEAK_START + pd.Timedelta(hours=1), 0.97, "zoom", + transform=ax_main.get_xaxis_transform(), + fontsize=8, color="goldenrod", ha="left", va="top") + + # Bottom panel: peak zoom + plot_envelope(ax_zoom, env_peak, obs, PEAK_START, PEAK_END, helene_label=True) + ax_zoom.set_xlabel("Date", fontsize=11) + ax_zoom.set_title( + "Helene peak zoom (Sep 26 12 UTC - Sep 28 06 UTC)", fontsize=10) + _format_xaxis(ax_zoom, major_interval=1, minor_hours=3) + ax_zoom.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %HUTC")) + ax_zoom.legend(fontsize=10, loc="lower right", frameon=True, framealpha=0.9) + + fig.suptitle( + f"4b -- 600-member crossed ensemble vs USGS 03463300 | {args.cat_id}\n" + "30 met forcing draws x 20 DA analysis state draws | R=0.07 mm2/h2", + fontsize=12, y=1.01, + ) + plt.tight_layout() + out2 = os.path.join(out_dir, f"{args.cat_id}_4b_crossed_ensemble_peak.png") + plt.savefig(out2, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4b_ensemble_vs_obs/plot_perturbation_category_shaded.py b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4b_ensemble_vs_obs/plot_perturbation_category_shaded.py new file mode 100644 index 00000000..b845c52d --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4b_ensemble_vs_obs/plot_perturbation_category_shaded.py @@ -0,0 +1,165 @@ +""" +Per-catchment shaded ensemble-band plot, organized by perturbation category. + +Three categories: + 1. Initial states (red) + 2. Meteorological forcings (blue) + 3. Hydrological states (green) + +For each category, all 20 ensemble members are shown as a shaded band +(min-max envelope fill) plus a thicker median line in the same color. +Qkrig observation overlaid in black. Hurricane Helene peak window shaded +in pink. Styled after a standard ensemble-forecast figure layout. + +Inputs (existing per-source sensitivity CSVs from run_perturbation_sensitivity.py): + //_sensitivity_init.csv (20 members, init only) + //_sensitivity_forcing.csv (20 members, forcing only) + //_sensitivity_process.csv (20 members, process noise only) + //_test_results.csv (Qkrig obs) + +Outputs: + //_perturbation_categories_linear.png + //_perturbation_categories_log.png +""" +import os +import pandas as pd +import numpy as np +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import matplotlib.patches as mpatches + +CAT = "cat-1016300" + +SEN_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_sensitivity" +OBS_DIR = "/mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance" + +OUT_LINEAR = os.path.join(SEN_DIR, CAT, f"{CAT}_perturbation_categories_linear.png") +OUT_LOG = os.path.join(SEN_DIR, CAT, f"{CAT}_perturbation_categories_log.png") + +# Plot window — wider context, similar to the paper's Sep 10 - Oct 08 +PLOT_START = pd.Timestamp("2024-09-20") +PLOT_END = pd.Timestamp("2024-10-05") + +# Helene peak band +HELENE_START = pd.Timestamp("2024-09-26 12:00:00") +HELENE_END = pd.Timestamp("2024-09-28 00:00:00") + +# Category configuration: file suffix, display label, color +CATEGORIES = [ + ("init", "Initial states only", "tab:red"), + ("forcing", "Meteorological forcings only", "tab:blue"), + ("process", "Hydrological states only", "tab:green"), +] + + +def load_members(source): + path = os.path.join(SEN_DIR, CAT, f"{CAT}_sensitivity_{source}.csv") + if not os.path.exists(path): + return None, None + df = pd.read_csv(path, parse_dates=["date"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + return df["date"].values, df[member_cols].to_numpy(dtype=float) + + +def load_obs(): + p = os.path.join(OBS_DIR, f"{CAT}.csv") + if not os.path.exists(p): + return None, None + df = pd.read_csv(p, parse_dates=["date"]) + return df["date"].values, df["qkrig"].values + + +def plot_panel(ax, obs_dates, obs_vals, log_y=False): + handles_labels = [] # for the legend + + # Plot each category as a shaded band + median line + for source, label, color in CATEGORIES: + dates, q = load_members(source) + if dates is None: + continue + d = pd.to_datetime(dates) + mask = (d >= PLOT_START) & (d <= PLOT_END) + if mask.sum() == 0: + continue + # Min-max envelope across all 20 members per timestep (widest possible band). + # Bands are visually narrow even with min/max because perturbations are + # tuned for production EnKF stability, not for max visible spread. + q_window = q[mask, :] + qmin = np.nanmin(q_window, axis=1) + qmax = np.nanmax(q_window, axis=1) + median = np.nanmedian(q_window, axis=1) + + ax.fill_between(d[mask], qmin, qmax, + color=color, alpha=0.30, zorder=2, + edgecolor="none") + line, = ax.plot(d[mask], median, + color=color, lw=1.7, alpha=0.95, zorder=3, + label=f"{label} (N=20)") + handles_labels.append((line, label)) + + # Helene peak shaded band (vertical) + ax.axvspan(HELENE_START, HELENE_END, + color="salmon", alpha=0.15, zorder=1) + ax.text((HELENE_START + (HELENE_END - HELENE_START) / 2), + ax.get_ylim()[1] if not log_y else 1.0, + "Helene peak", + fontsize=9, color="salmon", + ha="center", va="bottom", zorder=3) + + # Observation + if obs_dates is not None: + od = pd.to_datetime(obs_dates) + om = (od >= PLOT_START) & (od <= PLOT_END) + ax.plot(od[om], obs_vals[om], + color="black", lw=1.4, label="Qkrig (obs)", zorder=4) + + ax.set_xlabel("Date", fontsize=10) + if log_y: + ax.set_yscale("log") + ax.set_ylabel(r"q (mm hr$^{-1}$) [log scale]", fontsize=10) + ax.set_ylim(0.01, None) + else: + ax.set_ylabel("Discharge (mm/h)", fontsize=10) + ax.tick_params(axis="x", rotation=30, labelsize=9) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=2)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.grid(True, alpha=0.2) + ax.legend(loc="upper left", fontsize=9, frameon=True, framealpha=0.9) + + +def main(): + obs_dates, obs_vals = load_obs() + + # ----- Linear-y ----- + fig, ax = plt.subplots(1, 1, figsize=(20, 7)) + plot_panel(ax, obs_dates, obs_vals, log_y=False) + fig.suptitle( + f"Ensemble forecast by perturbation category - {CAT} - " + f"Sep 20 - Oct 5, 2024 (Hurricane Helene window)\n" + "Shaded bands = min-max envelope across 20 members. " + "Lines = ensemble median. DA off in all three sub-experiments.", + fontsize=11, y=0.995, + ) + plt.tight_layout() + plt.savefig(OUT_LINEAR, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_LINEAR}") + + # ----- Log-y (paper style) ----- + fig, ax = plt.subplots(1, 1, figsize=(20, 7)) + plot_panel(ax, obs_dates, obs_vals, log_y=True) + fig.suptitle( + f"Ensemble forecast by perturbation category - {CAT} - log-scale q - " + f"Sep 20 - Oct 5, 2024 (Hurricane Helene window)\n" + "Shaded bands = min-max envelope across 20 members. " + "Lines = ensemble median. DA off in all three sub-experiments.", + fontsize=11, y=0.995, + ) + plt.tight_layout() + plt.savefig(OUT_LOG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_LOG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_combined.py b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_combined.py new file mode 100644 index 00000000..c6d531a0 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_combined.py @@ -0,0 +1,235 @@ +#!/usr/bin/env python3 +""" +plot_routed_ensemble_combined.py +Combined poster plot: 600-member routed ensemble envelope + +deterministic DA lines (Vrugt/no-Vrugt) vs USGS gauge obs at 03463300. + +Inputs: + --vrugt-csv : routed_Q_test.csv from vrugt_dynamic_routed/ + --novrugt-csv : routed_Q_test.csv from novrugt_dynamic_routed/ + --ensemble-pq : routed_crossed_ensemble.parquet from run_route_crossed_ensemble.py + --out-dir : output directory + +The ensemble parquet columns: issue_time, lead_hour, member_0000..member_0599 (Q in m3/s) + +Usage: + python3 plot_routed_ensemble_combined.py \ + --vrugt-csv /mnt/disk2/suma_helen_poster/da_results/vrugt_dynamic_routed/routed_Q_test.csv \ + --novrugt-csv /mnt/disk2/suma_helen_poster/da_results/novrugt_dynamic_routed/routed_Q_test.csv \ + --ensemble-pq /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble_routed/routed_crossed_ensemble.parquet \ + --out-dir /mnt/disk2/suma_helen_poster/da_results/comparison_plots +""" + +import argparse +import os + +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +HELENE_START = pd.Timestamp("2024-09-25 00:00:00") +HELENE_END = pd.Timestamp("2024-09-29 00:00:00") +ZOOM_START = pd.Timestamp("2024-09-24 00:00:00") +ZOOM_END = pd.Timestamp("2024-09-28 18:00:00") # cut trailing ensemble bump + +COLOR_OBS = "black" +COLOR_VRUGT = "#1f77b4" # blue +COLOR_NOVRUGT = "#d62728" # red +COLOR_ENS = "#1f77b4" # same blue family as Vrugt + + +def load_det_csv(path): + df = pd.read_csv(path, parse_dates=["date"]) + return df.set_index("date").sort_index() + + +def build_envelope(pq_path, zoom_start, zoom_end): + """ + Load routed ensemble parquet and build percentile envelope by valid_time. + Returns DataFrame indexed by valid_time with columns p05, p25, p50, p75, p95. + """ + df = pd.read_parquet(pq_path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = [c for c in df.columns if c.startswith("member_")] + + records = [] + for _, row in df.iterrows(): + t0 = row["issue_time"] + lead = int(row["lead_hour"]) + vt = t0 + pd.Timedelta(hours=lead) + vals = row[member_cols].values.astype(float) + records.append({"valid_time": vt, "vals": vals}) + + # Group by valid_time (floor to hour), stack all member values across issue_times + from collections import defaultdict + groups = defaultdict(list) + for rec in records: + vt = rec["valid_time"].floor("1h") + groups[vt].extend(rec["vals"].tolist()) + + rows = [] + for vt in sorted(groups): + if vt < zoom_start or vt > zoom_end: + continue + v = np.array(groups[vt]) + v = v[~np.isnan(v)] + if len(v) == 0: + continue + rows.append({ + "valid_time": vt, + "p05": np.percentile(v, 5), + "p25": np.percentile(v, 25), + "p50": np.percentile(v, 50), + "p75": np.percentile(v, 75), + "p95": np.percentile(v, 95), + }) + + return pd.DataFrame(rows).set_index("valid_time") + + +def compute_kge(obs, sim): + mask = ~(np.isnan(obs) | np.isnan(sim)) + o, s = obs[mask], sim[mask] + if len(o) < 2 or np.std(o) == 0: + return np.nan + r = np.corrcoef(o, s)[0, 1] + return 1.0 - np.sqrt((r-1)**2 + (np.std(s)/np.std(o)-1)**2 + (np.mean(s)/np.mean(o)-1)**2) + + +def _format_xaxis(ax, locator, fmt, rotation=30): + ax.xaxis.set_major_locator(locator) + ax.xaxis.set_major_formatter(mdates.DateFormatter(fmt)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=rotation, ha="right", fontsize=9) + + +def _add_helene_band(ax): + ax.axvspan(HELENE_START, HELENE_END, color="gold", alpha=0.15, zorder=0) + + +def plot_zoom(ax, dates, obs, sim_v, sim_nv, kge_v, kge_nv, env=None): + mask = (dates >= ZOOM_START) & (dates <= ZOOM_END) + dz = dates[mask] + + # Ensemble envelope (bottom layer) + if env is not None: + ax.fill_between(env.index, env["p05"], env["p95"], + color=COLOR_ENS, alpha=0.18, label="Ensemble 5-95th pct") + ax.fill_between(env.index, env["p25"], env["p75"], + color=COLOR_ENS, alpha=0.30, label="Ensemble 25-75th pct") + ax.plot(env.index, env["p50"], + color=COLOR_ENS, lw=1.2, linestyle="--", alpha=0.7, label="Ensemble median") + + # Deterministic DA lines + ax.plot(dz, sim_nv[mask], color=COLOR_NOVRUGT, lw=1.4, linestyle="--", + label=f"DA -- constant R KGE={kge_nv:.3f}", zorder=3) + ax.plot(dz, sim_v[mask], color=COLOR_VRUGT, lw=1.6, + label=f"DA -- dynamic R (Vrugt) KGE={kge_v:.3f}", zorder=4) + + # USGS obs on top + ax.plot(dz, obs[mask], color=COLOR_OBS, lw=1.6, label="USGS obs (gauge 03463300)", zorder=5) + + peak_usgs = np.nanmax(obs[mask]) + ax.axhline(peak_usgs, color=COLOR_OBS, lw=0.7, linestyle=":", alpha=0.5) + # Place annotation to the right of the peak to avoid legend overlap + ax.text(pd.Timestamp("2024-09-28 06:00:00"), peak_usgs * 1.02, + f"USGS peak {peak_usgs:.0f} m³/s", fontsize=8.5, color="black", + alpha=0.8, ha="right") + + _add_helene_band(ax) + _format_xaxis(ax, mdates.DayLocator(interval=1), "%b %d") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--vrugt-csv", required=True) + parser.add_argument("--novrugt-csv", required=True) + parser.add_argument("--ensemble-pq", required=True, + help="routed_crossed_ensemble.parquet from run_route_crossed_ensemble.py") + parser.add_argument("--out-dir", required=True) + args = parser.parse_args() + os.makedirs(args.out_dir, exist_ok=True) + + vrugt = load_det_csv(args.vrugt_csv) + novrugt = load_det_csv(args.novrugt_csv) + + obs = vrugt["Q_usgs_m3s"].values + dates = vrugt.index + sim_v = vrugt["Q_routed_m3s"].values + sim_nv = novrugt["Q_routed_m3s"].reindex(dates).values + + kge_v = compute_kge(obs, sim_v) + kge_nv = compute_kge(obs, sim_nv) + + print("Building ensemble envelope...") + env = build_envelope(args.ensemble_pq, ZOOM_START, ZOOM_END) + print(f" Envelope: {len(env)} timesteps, " + f"peak p95={env['p95'].max():.1f} m3/s, peak p50={env['p50'].max():.1f} m3/s") + + peak_usgs = np.nanmax(obs[(dates >= ZOOM_START) & (dates <= ZOOM_END)]) + in_env = env["p95"].max() >= peak_usgs + print(f" USGS peak {peak_usgs:.1f} m3/s {'IS' if in_env else 'IS NOT'} within p95 envelope") + + # ------------------------------------------------------------------ # + # Figure 1: Helene zoom only -- main poster panel + # ------------------------------------------------------------------ # + fig, ax = plt.subplots(figsize=(11, 5.5)) + plot_zoom(ax, dates, obs, sim_v, sim_nv, kge_v, kge_nv, env) + ax.set_xlabel("Date (UTC)", fontsize=11) + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_title( + "Hurricane Helene -- Probabilistic discharge forecast at gauge 03463300 (South Toe River)\n" + "600-member CFE ensemble + DA (Muskingum routing)", + fontsize=11) + ax.legend(fontsize=9, loc="lower right") + ax.grid(True, alpha=0.25) + plt.tight_layout() + out1 = os.path.join(args.out_dir, "helene_ensemble_vs_usgs.png") + plt.savefig(out1, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + # ------------------------------------------------------------------ # + # Figure 2: two-panel (full period + Helene zoom) + # ------------------------------------------------------------------ # + fig, (ax_top, ax_bot) = plt.subplots( + 2, 1, figsize=(14, 9), + gridspec_kw={"height_ratios": [1, 1.5]}) + + # Top: full test period (deterministic only, no ensemble for readability) + ax_top.plot(dates, obs, color=COLOR_OBS, lw=0.8, label="USGS obs", zorder=3) + ax_top.plot(dates, sim_v, color=COLOR_VRUGT, lw=0.9, + label=f"DA dynamic R (Vrugt) KGE={kge_v:.3f}", zorder=2) + ax_top.plot(dates, sim_nv, color=COLOR_NOVRUGT, lw=0.9, linestyle="--", + label=f"DA constant R KGE={kge_nv:.3f}", zorder=2) + ax_top.axvspan(ZOOM_START, ZOOM_END, color="gray", alpha=0.12, zorder=0) + ax_top.axvspan(HELENE_START, HELENE_END, color="gold", alpha=0.15, zorder=0) + ax_top.text(HELENE_START + pd.Timedelta(days=0.5), + np.nanmax(obs) * 0.88, "Helene", + fontsize=9, color="goldenrod", fontweight="bold") + ax_top.set_ylabel("Discharge (m³/s)", fontsize=10) + ax_top.set_title( + "Routed discharge at gauge 03463300 (South Toe River) -- Test period\n" + "DA with dynamic Vrugt R vs. constant R -- Muskingum routing", + fontsize=11) + ax_top.legend(fontsize=9, loc="upper left") + ax_top.grid(True, alpha=0.25) + _format_xaxis(ax_top, mdates.MonthLocator(interval=2), "%Y-%m") + + # Bottom: Helene zoom with ensemble + plot_zoom(ax_bot, dates, obs, sim_v, sim_nv, kge_v, kge_nv, env) + ax_bot.set_xlabel("Date (UTC)", fontsize=10) + ax_bot.set_ylabel("Discharge (m³/s)", fontsize=10) + ax_bot.set_title("Hurricane Helene window -- 600-member ensemble envelope", fontsize=10) + ax_bot.legend(fontsize=9, loc="lower right") + ax_bot.grid(True, alpha=0.25) + + plt.tight_layout() + out2 = os.path.join(args.out_dir, "helene_ensemble_twopanel.png") + plt.savefig(out2, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_vs_usgs.py b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_vs_usgs.py new file mode 100644 index 00000000..55c52d85 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_vs_usgs.py @@ -0,0 +1,176 @@ +#!/usr/bin/env python3 +""" +plot_routed_ensemble_vs_usgs.py +Plot T-route routed 600-member ensemble envelope vs USGS at gauge 03463300. + +Input: routed_crossed_ensemble.parquet from run_route_crossed_ensemble.py + Columns: issue_time, lead_hour, member_0000..member_0599 (Q in m3/s) + +Usage: + python3 plot_routed_ensemble_vs_usgs.py \ + --routed-pq /mnt/disk2/.../folder1_vrugt_routed/routed_crossed_ensemble.parquet \ + --usgs-csv /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv \ + --label "F1 Vrugt — 1 gauge holdout" \ + --out-dir ~/plots_f1_1gauge +""" + +import argparse +from collections import defaultdict +from pathlib import Path + +import matplotlib.dates as mdates +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd + +WATERSHED_AREA_KM2 = 113.18 # cat-1016300 drainage area +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +PLOT_START = pd.Timestamp("2024-09-24 00:00:00") +PLOT_END = pd.Timestamp("2024-09-28 18:00:00") +PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +PEAK_END = pd.Timestamp("2024-09-28 06:00:00") +HELENE_PEAK = pd.Timestamp("2024-09-27 12:00:00") + +COLOR_ENS = "#1f77b4" +COLOR_OBS = "black" + + +def build_envelope(pq_path, t_start, t_end): + df = pd.read_parquet(pq_path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = [c for c in df.columns if c.startswith("member_")] + + groups = defaultdict(list) + for _, row in df.iterrows(): + vt = (row["issue_time"] + pd.Timedelta(hours=int(row["lead_hour"]))).floor("1h") + if t_start <= vt <= t_end: + groups[vt].extend(row[member_cols].values.astype(float).tolist()) + + rows = [] + for vt in sorted(groups): + v = np.array(groups[vt]) + v = v[~np.isnan(v)] + if len(v) == 0: + continue + rows.append({"valid_time": vt, + "p05": np.percentile(v, 5), "p25": np.percentile(v, 25), + "p50": np.percentile(v, 50), "p75": np.percentile(v, 75), + "p95": np.percentile(v, 95)}) + return pd.DataFrame(rows).set_index("valid_time") + + +def load_usgs(path): + df = pd.read_csv(path) + print(f" USGS CSV columns: {list(df.columns)}") + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() or "discharge" in c.lower()) + print(f" Using date_col={date_col!r}, q_col={q_col!r}") + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series.sort_index() + + +def compute_kge(obs, sim): + mask = ~(np.isnan(obs) | np.isnan(sim)) + o, s = obs[mask], sim[mask] + if len(o) < 2 or np.std(o) == 0: + return np.nan + r = np.corrcoef(o, s)[0, 1] + return 1 - np.sqrt((r-1)**2 + (np.std(s)/np.std(o)-1)**2 + (np.mean(s)/np.mean(o)-1)**2) + + +def _shade(ax, env): + ax.fill_between(env.index, env["p05"], env["p95"], + color=COLOR_ENS, alpha=0.18, label="5th–95th percentile") + ax.fill_between(env.index, env["p25"], env["p75"], + color=COLOR_ENS, alpha=0.32, label="25th–75th percentile") + ax.plot(env.index, env["p50"], color=COLOR_ENS, lw=1.5, label="Ensemble median (600 members)") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--routed-pq", required=True) + parser.add_argument("--usgs-csv", required=True) + parser.add_argument("--label", default="F1 Vrugt — 1 gauge holdout") + parser.add_argument("--out-dir", required=True) + args = parser.parse_args() + + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + print("Building full-window envelope...") + env_full = build_envelope(args.routed_pq, PLOT_START, PLOT_END) + print("Building peak-window envelope...") + env_peak = build_envelope(args.routed_pq, PEAK_START, PEAK_END) + + print("Loading USGS obs...") + try: + obs_all = load_usgs(args.usgs_csv) + obs_full = obs_all.loc[PLOT_START:PLOT_END] + obs_peak = obs_all.loc[PEAK_START:PEAK_END] + print(f" Loaded {len(obs_all)} USGS obs rows; " + f"full window: {len(obs_full)}, peak window: {len(obs_peak)}") + except Exception as e: + print(f" Warning: USGS load failed ({e}) — plotting without obs") + obs_full = obs_peak = None + + kge = np.nan + if obs_full is not None and len(env_full) > 0: + med = env_full["p50"] + obs_hourly = obs_full.copy() + obs_hourly.index = obs_hourly.index.floor("1h") + obs_hourly = obs_hourly[~obs_hourly.index.duplicated(keep="first")] + aligned = obs_hourly.reindex(med.index) + print(f" KGE alignment: {aligned.notna().sum()} matched out of {len(med)} timesteps") + kge = compute_kge(aligned.values, med.values) + + # ── Figure 1: Full Helene window ────────────────────────────────────────── + fig, ax = plt.subplots(figsize=(12, 4.5)) + _shade(ax, env_full) + if obs_full is not None: + ax.plot(obs_full.index, obs_full.values, color=COLOR_OBS, lw=1.6, + label="USGS obs (gauge 03463300)") + ax.axvspan(PEAK_START, PEAK_END, color="salmon", alpha=0.12, zorder=0, label="Helene peak window") + ax.set_ylabel("Discharge (m³/s)") + ax.set_xlabel("Date") + ax.set_title(f"Routed 600-member ensemble vs USGS | {args.label}\n" + f"Ensemble median KGE = {kge:.3f} | Sep 24–28 2024") + ax.legend(fontsize=9, loc="upper left") + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + plt.tight_layout() + out = out_dir / "routed_ensemble_vs_usgs_full.png" + fig.savefig(out, dpi=150) + print(f"Saved: {out}") + plt.close(fig) + + # ── Figure 2: Peak zoom ─────────────────────────────────────────────────── + fig, ax = plt.subplots(figsize=(10, 4.5)) + _shade(ax, env_peak) + if obs_peak is not None: + ax.plot(obs_peak.index, obs_peak.values, color=COLOR_OBS, lw=1.6, + label="USGS obs (gauge 03463300)") + peak_val = float(obs_peak.max()) + ax.axhline(peak_val, color=COLOR_OBS, lw=0.7, ls=":", alpha=0.5) + ax.annotate(f"USGS peak\n{peak_val:.0f} m³/s", + xy=(HELENE_PEAK, peak_val), xytext=(10, -40), + textcoords="offset points", fontsize=8, color=COLOR_OBS) + ax.set_ylabel("Discharge (m³/s)") + ax.set_xlabel("Date") + ax.set_title(f"Helene peak zoom | {args.label}") + ax.legend(fontsize=9) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %HUTC")) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + plt.tight_layout() + out = out_dir / "routed_ensemble_vs_usgs_peak.png" + fig.savefig(out, dpi=150) + print(f"Saved: {out}") + plt.close(fig) + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4b_ensemble_vs_obs/run_4b_crossed_f3.sh b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4b_ensemble_vs_obs/run_4b_crossed_f3.sh new file mode 100644 index 00000000..97d0bd83 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4b_ensemble_vs_obs/run_4b_crossed_f3.sh @@ -0,0 +1,40 @@ +#!/usr/bin/env bash +# F3 dynamic Vrugt seeded (20pct gauge holdout) — 4b crossed ensemble plots. +# Runs plot_crossed_ensemble.py for all 21 catchments. +# Reads //_crossed_ensemble.parquet (already generated). +# +# Usage: +# bash run_4b_crossed_f3.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +DATA_DIR=/mnt/disk2/1400_sites_helene/da_crossed_dynamic_vrugt_seeded +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +echo "[F3-4b-crossed] Crossed ensemble plots — data: $DATA_DIR" + +for CAT in "${CATS[@]}"; do + PQ="$DATA_DIR/$CAT/${CAT}_crossed_ensemble.parquet" + if [ ! -f "$PQ" ]; then + echo " [$CAT] No crossed ensemble parquet — skipping" + continue + fi + echo " [$CAT] crossed ensemble..." + $TROUTE "$SCRIPT_DIR/plot_crossed_ensemble.py" \ + --ensemble-dir "$DATA_DIR" \ + --cat-id "$CAT" \ + --usgs-csv "$USGS_CSV" +done + +echo "[F3-4b-crossed] Done." diff --git a/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4b_ensemble_vs_obs/run_4b_f3.sh b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4b_ensemble_vs_obs/run_4b_f3.sh new file mode 100644 index 00000000..33041359 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4b_ensemble_vs_obs/run_4b_f3.sh @@ -0,0 +1,40 @@ +#!/usr/bin/env bash +# F3 dynamic Vrugt seeded (20pct gauge holdout) — 4b routed ensemble plots. +# +# Runs: +# 1. plot_routed_ensemble_vs_usgs.py — F3 ensemble envelope vs USGS at outlet +# 2. plot_routed_ensemble_combined.py — F1 Vrugt vs F3 dynamic Vrugt seeded comparison +# +# Requires routed_Q_test.csv (both folders) and routed_crossed_ensemble.parquet +# (CROSSED_DIR) to be present. +# +# Usage: +# bash run_4b_f3.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +F1_DA_DIR=/mnt/disk2/suma_helen_poster/da_results/v2_perturbation_da_on +F3_DA_DIR=/mnt/disk2/1400_sites_helene/da_arms_dynamic_vrugt_seeded +CROSSED_DIR=/mnt/disk2/1400_sites_helene/da_crossed_dynamic_vrugt_seeded +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv +OUT_DIR="$CROSSED_DIR" + +echo "[F3-4b] Routed ensemble vs USGS (F3 only)..." +$TROUTE "$SCRIPT_DIR/plot_routed_ensemble_vs_usgs.py" \ + --routed-pq "$CROSSED_DIR/routed_crossed_ensemble.parquet" \ + --usgs-csv "$USGS_CSV" \ + --label "F3 Dynamic Vrugt seeded — 20pct gauge holdout" \ + --out-dir "$OUT_DIR" + +echo "[F3-4b] Combined comparison: F1 Vrugt vs F3 Dynamic Vrugt seeded..." +$TROUTE "$SCRIPT_DIR/plot_routed_ensemble_combined.py" \ + --vrugt-csv "$F1_DA_DIR/routed_Q_test.csv" \ + --novrugt-csv "$F3_DA_DIR/routed_Q_test.csv" \ + --ensemble-pq "$CROSSED_DIR/routed_crossed_ensemble.parquet" \ + --out-dir "$OUT_DIR" + +echo "[F3-4b] Done. Outputs in: $OUT_DIR" +ls "$OUT_DIR"/*.png 2>/dev/null || echo " (no PNGs in out dir)" diff --git a/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4c_timeseries/plot_forecast_spaghetti.py b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4c_timeseries/plot_forecast_spaghetti.py new file mode 100644 index 00000000..655edaef --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4c_timeseries/plot_forecast_spaghetti.py @@ -0,0 +1,176 @@ +""" +plot_forecast_spaghetti.py + +"All forecasts ending at this hour" — verification-time-fixed spaghetti view. + +Window: Sep 24 18 UTC -> Sep 30 06 UTC. +For every initialization time whose 18-hour forecast overlaps this window: + DA : shaded min/max band across 20 members + ensemble mean line + OL : ensemble mean line only (no shading), dashed + +Lines are colored by initialization DATE (7 colors, Sep 24-30) so the +temporal progression is readable. USGS obs overlaid in black. + +Inputs: + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + Wide format: issue_time, lead_hour, member_00..member_19 (q in m³/s) + +Obs: + /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +Outputs: + /forecast_spaghetti_helene.png +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import matplotlib.patches as mpatches + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast_hardcoded_r_routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +# Verification window — only trajectories whose valid_time falls here are shown +PLOT_START = pd.Timestamp("2024-09-24 18:00:00") +PLOT_END = pd.Timestamp("2024-09-30 06:00:00") + +HELENE_PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +HELENE_PEAK_END = pd.Timestamp("2024-09-28 00:00:00") + +# One color per init date (Sep 24-30) +DATE_COLORS = { + "2024-09-24": "#1f77b4", + "2024-09-25": "#ff7f0e", + "2024-09-26": "#2ca02c", + "2024-09-27": "#d62728", + "2024-09-28": "#9467bd", + "2024-09-29": "#8c564b", + "2024-09-30": "#e377c2", +} + + +def load_parquet(path): + df = pd.read_parquet(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + df["ens_mean"] = df[member_cols].mean(axis=1) + df["ens_min"] = df[member_cols].min(axis=1) + df["ens_max"] = df[member_cols].max(axis=1) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + return df + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() + or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--route-dir", default=DEFAULT_ROUTE_DIR) + parser.add_argument("--out-dir", default=None) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--da-name", default="routed_leadtime_da_full.parquet") + parser.add_argument("--ol-name", default="routed_leadtime_openloop_full.parquet") + args = parser.parse_args() + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + print("Loading parquets...") + da = load_parquet(os.path.join(args.route_dir, args.da_name)) + ol = load_parquet(os.path.join(args.route_dir, args.ol_name)) + obs = load_usgs(args.usgs_csv) + + # Keep only valid_times inside the plot window + da = da[(da["valid_time"] >= PLOT_START) & (da["valid_time"] <= PLOT_END)] + ol = ol[(ol["valid_time"] >= PLOT_START) & (ol["valid_time"] <= PLOT_END)] + + issue_times = sorted(da["issue_time"].unique()) + print(f" {len(issue_times)} initialization times contribute to this window") + + fig, ax = plt.subplots(figsize=(18, 6)) + + # Helene peak shading + ax.axvspan(HELENE_PEAK_START, HELENE_PEAK_END, + color="salmon", alpha=0.12, zorder=0) + ax.text(HELENE_PEAK_START + pd.Timedelta(hours=6), 1.0, "Helene peak", + transform=ax.get_xaxis_transform(), + fontsize=9, color="firebrick", ha="left", va="top") + + # Plot each initialization time's trajectory + for t0 in issue_times: + date_str = str(pd.Timestamp(t0).date()) + color = DATE_COLORS.get(date_str, "gray") + + da_t = da[da["issue_time"] == t0].sort_values("valid_time") + ol_t = ol[ol["issue_time"] == t0].sort_values("valid_time") + + if da_t.empty: + continue + + # DA: shaded band (min-max across 20 members) + ensemble mean + ax.fill_between(da_t["valid_time"], da_t["ens_min"], da_t["ens_max"], + color=color, alpha=0.06, zorder=2) + ax.plot(da_t["valid_time"], da_t["ens_mean"], + color=color, lw=0.8, alpha=0.55, zorder=3) + + # OL: ensemble mean only, dashed + if not ol_t.empty: + ax.plot(ol_t["valid_time"], ol_t["ens_mean"], + color=color, lw=0.6, alpha=0.30, linestyle="--", zorder=2) + + # USGS obs + obs_w = obs.loc[PLOT_START:PLOT_END] + ax.plot(obs_w.index, obs_w.values, + color="black", lw=1.8, zorder=6, label="USGS obs") + + # Legend: one patch per init date + obs + DA/OL style + patches = [mpatches.Patch(color=c, label=f"Init {d}") + for d, c in DATE_COLORS.items()] + patches.append(plt.Line2D([0], [0], color="black", lw=1.8, label="USGS obs")) + patches.append(plt.Line2D([0], [0], color="gray", lw=1.2, alpha=0.7, + label="DA — ens mean (solid) + spread (shaded)")) + patches.append(plt.Line2D([0], [0], color="gray", lw=1.0, linestyle="--", + alpha=0.5, label="OL — ens mean (dashed)")) + ax.legend(handles=patches, fontsize=8, loc="upper left", + frameon=True, framealpha=0.9, ncol=2) + + ax.set_xlim(PLOT_START, PLOT_END) + ax.set_xlabel("Date", fontsize=11) + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_title( + "All 18-hour forecast trajectories ending in this window — USGS 03463300\n" + "Sep 24 18 UTC → Sep 30 06 UTC | DA: shaded band + mean | OL: mean dashed", + fontsize=11, + ) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=1)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.xaxis.set_minor_locator(mdates.HourLocator(interval=6)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.2, lw=0.4) + plt.tight_layout() + + out_path = os.path.join(out_dir, "forecast_spaghetti_helene.png") + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4c_timeseries/plot_helene_issue_time_hydrograph.py b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4c_timeseries/plot_helene_issue_time_hydrograph.py new file mode 100644 index 00000000..4e2e1dc2 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4c_timeseries/plot_helene_issue_time_hydrograph.py @@ -0,0 +1,198 @@ +""" +Per-issue-time forecast hydrograph diagnostic. + +For two contrasting issue times — one storm-peak (Helene) and one low-flow — +show what the DA and open-loop ensembles actually produce over the 18-hour +forecast window, alongside the kriging observation. Designed to answer in one +picture: is DA over-shooting, collapsing, or oscillating compared to open-loop? + +Each panel shows: + - All 20 DA members as faint purple lines + median (thick purple) + - All 20 open-loop members as faint gray lines + median (thick gray, dashed) + - Kriging obs as a thick black line (context: t0-6 through t0+18) + - Vertical line at t0 + annotation with obs(t0) and lead-1 ensemble means + +Inputs (from run_lead_time_forecast_sweep.py): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv +Obs (from production EnKF): + //_test_results.csv + +Output: + //_issue_time_hydrograph_helene_vs_lowflow.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +DEFAULT_LEADTIME_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast" +DEFAULT_DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" + +# Module-level globals set by main() +CAT = "cat-1016300" +LEADTIME_DIR = DEFAULT_LEADTIME_DIR +DA_DIR = DEFAULT_DA_DIR + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" +OBS_COLOR = "black" + +# Two issue times — pick a Helene-peak one and a typical low-flow one +HELENE_T0 = pd.Timestamp("2024-09-26 12:00:00") +LOWFLOW_T0 = pd.Timestamp("2024-03-15 12:00:00") + +CONTEXT_HOURS_BEFORE = 6 # hours of obs context shown before t0 +LEAD_HOURS_AFTER = 18 # forecast horizon + +OUT_PNG = os.path.join( + LEADTIME_DIR, CAT, + f"{CAT}_issue_time_hydrograph_helene_vs_lowflow.png", +) + + +def load_forecasts(path): + df = pd.read_csv(path, parse_dates=['issue_time', 'valid_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + return df, member_cols + + +def load_obs(): + obs_path = os.path.join(DA_DIR, CAT, f"{CAT}_test_results.csv") + df = pd.read_csv(obs_path, parse_dates=['date']) + return df.set_index('date')['obs_mm_h'] + + +def slice_forecast(df, member_cols, t0): + """Return (valid_times, member_array shape (lead, N)) for a single issue time.""" + sub = df[df['issue_time'] == t0].sort_values('lead_hour') + if len(sub) == 0: + raise ValueError(f"No forecast rows for issue_time={t0}") + valid_times = pd.to_datetime(sub['valid_time'].values) + member_arr = sub[member_cols].to_numpy(dtype=float) + return valid_times, member_arr + + +def plot_panel(ax, t0, df_da, df_ol, member_cols, obs_series, title_extra=""): + # Forecast trajectories + vt_da, q_da = slice_forecast(df_da, member_cols, t0) + vt_ol, q_ol = slice_forecast(df_ol, member_cols, t0) + + # Obs context window + ctx_start = t0 - pd.Timedelta(hours=CONTEXT_HOURS_BEFORE) + ctx_end = t0 + pd.Timedelta(hours=LEAD_HOURS_AFTER) + obs_window = obs_series.loc[ctx_start:ctx_end] + + # Open-loop members (drawn first so DA paints on top) + for j in range(q_ol.shape[1]): + ax.plot(vt_ol, q_ol[:, j], color=OL_COLOR, lw=0.6, alpha=0.45, zorder=2) + # DA members + for j in range(q_da.shape[1]): + ax.plot(vt_da, q_da[:, j], color=DA_COLOR, lw=0.6, alpha=0.45, zorder=3) + # Medians + ax.plot(vt_ol, np.nanmedian(q_ol, axis=1), color=OL_COLOR, lw=2.6, + linestyle="--", zorder=4, label="Open-loop median") + ax.plot(vt_da, np.nanmedian(q_da, axis=1), color=DA_COLOR, lw=2.6, + zorder=5, label="DA median") + # Obs + ax.plot(obs_window.index, obs_window.values, color=OBS_COLOR, lw=1.8, + zorder=6, label="Qkrig (obs)") + + # t0 marker + annotation + ax.axvline(t0, color="black", lw=1.0, linestyle=":", alpha=0.6, zorder=1) + obs_t0 = float(obs_series.get(t0, np.nan)) + obs_lead1 = float(obs_series.get(t0 + pd.Timedelta(hours=1), np.nan)) + da_lead1_mean = float(np.nanmean(q_da[0, :])) + ol_lead1_mean = float(np.nanmean(q_ol[0, :])) + + info_text = ( + f"t0 = {t0.strftime('%Y-%m-%d %H:%M')}\n" + f"obs(t0) = {obs_t0:.3f} mm/h\n" + f"obs(t0+1) = {obs_lead1:.3f} mm/h\n" + f"DA mean(t0+1) = {da_lead1_mean:.3f} mm/h\n" + f"OL mean(t0+1) = {ol_lead1_mean:.3f} mm/h" + ) + ax.text(0.02, 0.97, info_text, + transform=ax.transAxes, + fontsize=9, va='top', ha='left', family='monospace', + bbox=dict(boxstyle='round,pad=0.4', facecolor='white', + alpha=0.85, edgecolor='gray')) + + ax.set_title(title_extra, fontsize=12) + ax.set_ylabel("Discharge (mm/h)", fontsize=10) + ax.set_xlabel("Date", fontsize=10) + ax.xaxis.set_major_locator(mdates.HourLocator(interval=3)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %H:%M")) + ax.tick_params(axis='x', rotation=30, labelsize=8) + ax.grid(True, alpha=0.2) + ax.legend(loc='upper right', fontsize=9, frameon=True, framealpha=0.9) + + +def snap_to_nearest(t0, available_times): + """Return the available issue_time closest to t0.""" + times = pd.DatetimeIndex(available_times) + idx = np.abs((times - t0).total_seconds()).argmin() + snapped = times[idx] + if snapped != t0: + print(f" Snapped {t0} -> {snapped} (nearest available issue time)") + return snapped + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', default="cat-1016300") + parser.add_argument('--leadtime-dir', default=DEFAULT_LEADTIME_DIR, + help='Dir holding /_lead_time_forecasts_{da,openloop}.csv') + parser.add_argument('--da-dir', default=DEFAULT_DA_DIR, + help='Dir holding /_test_results.csv (for obs)') + parser.add_argument('--helene-t0', default=None, + help='Issue time for the Helene panel (default: 2024-09-26 12:00:00, ' + 'snapped to nearest available)') + parser.add_argument('--lowflow-t0', default=None, + help='Issue time for the low-flow panel (default: 2024-03-15 00:00:00, ' + 'snapped to nearest available)') + args = parser.parse_args() + + global CAT, LEADTIME_DIR, DA_DIR, OUT_PNG + CAT = args.cat_id + LEADTIME_DIR = args.leadtime_dir + DA_DIR = args.da_dir + OUT_PNG = os.path.join(LEADTIME_DIR, CAT, + f"{CAT}_issue_time_hydrograph_helene_vs_lowflow.png") + + da_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_da.csv") + ol_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_openloop.csv") + df_da, m_cols = load_forecasts(da_path) + df_ol, _ = load_forecasts(ol_path) + obs_series = load_obs() + + available = df_da['issue_time'].unique() + helene_t0 = snap_to_nearest( + pd.Timestamp(args.helene_t0) if args.helene_t0 else HELENE_T0, available) + lowflow_t0 = snap_to_nearest( + pd.Timestamp(args.lowflow_t0) if args.lowflow_t0 else pd.Timestamp("2024-03-15 00:00:00"), + available) + + fig, (ax_h, ax_l) = plt.subplots(1, 2, figsize=(20, 7)) + + plot_panel(ax_h, helene_t0, df_da, df_ol, m_cols, obs_series, + title_extra=f"Helene peak issue time — {helene_t0.strftime('%Y-%m-%d %H:%M')}") + plot_panel(ax_l, lowflow_t0, df_da, df_ol, m_cols, obs_series, + title_extra=f"Low-flow issue time — {lowflow_t0.strftime('%Y-%m-%d %H:%M')}") + + fig.suptitle( + f"Per-issue-time forecast hydrograph — {CAT}\n" + "Faint lines = 20 ensemble members. Thick lines = ensemble medians. " + "Vertical dotted line = t0 (DA stops here for both scenarios).", + fontsize=12, y=1.00, + ) + plt.tight_layout() + plt.savefig(OUT_PNG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4c_timeseries/plot_lead_time_reconstructed_timeseries.py b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4c_timeseries/plot_lead_time_reconstructed_timeseries.py new file mode 100644 index 00000000..a6d03b03 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4c_timeseries/plot_lead_time_reconstructed_timeseries.py @@ -0,0 +1,263 @@ +""" +Reconstructed time series at gauge 03463300, from our routed +lead-time forecasts. Mirrors the methodology of plot_timeseries_hourly_ensemble.py +(qSpatialAR) but feeds from the EnKF + T-route pipeline instead of the CNN. + +For each target hour t, pool ALL forecasts that land on t (across many +issue_times × many lead_hours × 20 members). Compute median + 5th/95th +percentile envelope. Plot DA and open-loop reconstructions overlaid on USGS +obs, with the Helene window shaded. + +Input parquets (from run_route_troute_leadtime.py): + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + Wide format: issue_time, lead_hour, member_00..member_19 (q in m³/s) + +Obs: + /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +Output: + /lead_time_reconstructed_timeseries.png + Default plot window: 2024-09-10 → 2024-10-10. + +Each scenario's curve is annotated with NSE vs USGS obs over the plot window. +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" +OBS_COLOR = "black" + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast_hardcoded_r_routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +# Watershed area at gauge 03463300 (sum of 21 catchments per GPKG hydrofabric) +WATERSHED_AREA_KM2 = 113.18 +# mm/h depth → m³/s: × 113.18 km² × 1000 / 3600 = 31.439 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +PLOT_START_DEFAULT = pd.Timestamp("2024-09-10 00:00:00") +PLOT_END_DEFAULT = pd.Timestamp("2024-10-10 23:00:00") + +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-30 23:00:00") + + +def load_parquet_long(path): + df = pd.read_parquet(path) + df['issue_time'] = pd.to_datetime(df['issue_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + if not member_cols: + raise ValueError(f"No member columns in {path}") + long = df.melt( + id_vars=['issue_time', 'lead_hour'], + value_vars=member_cols, + var_name='member', + value_name='q_m3s', + ) + long['valid_time'] = (long['issue_time'] + + pd.to_timedelta(long['lead_hour'], unit='h')) + return long + + +def reconstruct(long_df): + """For each valid_time, pool all (issue_time × lead × member) forecasts and + return median, p05, p95, count.""" + grouped = long_df.groupby('valid_time')['q_m3s'].agg( + median='median', + p05=lambda s: s.quantile(0.05), + p95=lambda s: s.quantile(0.95), + count='count', + ).reset_index() + return grouped + + +def load_usgs(usgs_csv): + """Load USGS gauge obs and ensure m³/s. + + The CSV at gauge 03463300 has the column `QObs(mm/h)` — catchment-averaged + depth, not gauge discharge in m³/s. Auto-convert when column name contains + 'mm': m³/s = mm/h × WATERSHED_AREA_KM2 × 1000 / 3600. + """ + df = pd.read_csv(usgs_csv) + date_col = next((c for c in df.columns + if c.lower() in ('datetime', 'date', 'time', 'timestamp')), None) + q_col = next((c for c in df.columns + if 'q' in c.lower() or 'flow' in c.lower() + or 'discharge' in c.lower()), None) + if not date_col or not q_col: + raise ValueError(f"Couldn't id date/Q columns in {usgs_csv}: {df.columns.tolist()}") + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if 'mm' in q_col.lower(): + series = series * MM_H_TO_M3_S + print(f" Converted obs from mm/h (column '{q_col}') to m³/s " + f"using area = {WATERSHED_AREA_KM2} km² (× {MM_H_TO_M3_S:.4f})") + else: + print(f" Loaded obs as-is from column '{q_col}' (assumed m³/s)") + return series + + +def nse(obs, sim): + m = np.isfinite(obs) & np.isfinite(sim) + if m.sum() < 2: + return np.nan + o, s = obs[m], sim[m] + denom = ((o - o.mean()) ** 2).sum() + if denom < 1e-10: + return np.nan + return 1.0 - float(((o - s) ** 2).sum() / denom) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--route-dir', default=DEFAULT_ROUTE_DIR) + parser.add_argument('--out-dir', default=None, + help='Default: --route-dir') + parser.add_argument('--usgs-csv', default=DEFAULT_USGS_CSV) + parser.add_argument('--plot-start', default=str(PLOT_START_DEFAULT)) + parser.add_argument('--plot-end', default=str(PLOT_END_DEFAULT)) + parser.add_argument('--da-name', default='routed_leadtime_da_full.parquet') + parser.add_argument('--ol-name', default='routed_leadtime_openloop_full.parquet') + args = parser.parse_args() + + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + plot_start = pd.Timestamp(args.plot_start) + plot_end = pd.Timestamp(args.plot_end) + + print(f"Loading DA parquet: {args.da_name}") + da_long = load_parquet_long(os.path.join(args.route_dir, args.da_name)) + print(f"Loading OL parquet: {args.ol_name}") + ol_long = load_parquet_long(os.path.join(args.route_dir, args.ol_name)) + print(f"Loading USGS obs: {args.usgs_csv}") + obs = load_usgs(args.usgs_csv) + + print("Reconstructing DA time series (overlapping-leads pool)...") + da_rec = reconstruct(da_long) + print(f" {len(da_rec)} unique target hours") + print("Reconstructing OL time series...") + ol_rec = reconstruct(ol_long) + + # Clip to plot window and join obs + da_rec = da_rec[(da_rec['valid_time'] >= plot_start) & (da_rec['valid_time'] <= plot_end)].copy() + ol_rec = ol_rec[(ol_rec['valid_time'] >= plot_start) & (ol_rec['valid_time'] <= plot_end)].copy() + da_rec['obs'] = da_rec['valid_time'].map(obs) + ol_rec['obs'] = ol_rec['valid_time'].map(obs) + + nse_da = nse(da_rec['obs'].values, da_rec['median'].values) + nse_ol = nse(ol_rec['obs'].values, ol_rec['median'].values) + print(f"NSE (plot window) — DA: {nse_da:.3f} | OL: {nse_ol:.3f}") + + # ---- Plot ---- + fig, ax = plt.subplots(figsize=(16, 5.5)) + fig.patch.set_facecolor("white") + + # Helene shaded + ax.axvspan(HELENE_START, HELENE_END, + color="firebrick", alpha=0.08, zorder=0) + ax.text(HELENE_START + pd.Timedelta(hours=12), 1.0, "Helene", + transform=ax.get_xaxis_transform(), + fontsize=9, color="firebrick", ha="left", va="top", + fontweight="bold") + + # Open-loop band + median + ax.fill_between(ol_rec['valid_time'], ol_rec['p05'], ol_rec['p95'], + color=OL_COLOR, alpha=0.18, zorder=2, + label="Open-loop [5th–95th pctile envelope]") + ax.plot(ol_rec['valid_time'], ol_rec['median'], + color=OL_COLOR, lw=1.2, linestyle='--', zorder=3, + label=f"Open-loop median (NSE={nse_ol:.3f})") + + # DA band + median + ax.fill_between(da_rec['valid_time'], da_rec['p05'], da_rec['p95'], + color=DA_COLOR, alpha=0.22, zorder=4, + label="DA [5th–95th pctile envelope]") + ax.plot(da_rec['valid_time'], da_rec['median'], + color=DA_COLOR, lw=1.4, zorder=5, + label=f"DA median (NSE={nse_da:.3f})") + + # Observed as dots + obs_window = obs.loc[plot_start:plot_end] + finite = np.isfinite(obs_window.values) + ax.scatter(obs_window.index[finite], obs_window.values[finite], + s=6, color=OBS_COLOR, linewidths=0, zorder=6, + label="USGS obs") + + ax.set_xlabel("Date", fontsize=11) + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_title( + f"Reconstructed forecast time series at USGS 03463300\n" + f"Overlapping-leads pool from EnKF forecast ensemble " + f"({plot_start.date()} – {plot_end.date()})", + fontsize=12, + ) + ax.set_xlim(plot_start, plot_end) + ax.xaxis.set_major_locator(mdates.WeekdayLocator(interval=1)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.25, lw=0.4) + ax.legend(loc='upper left', fontsize=9, frameon=True, framealpha=0.92) + + plt.tight_layout() + out_path = os.path.join(out_dir, "lead_time_reconstructed_timeseries.png") + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + # Also produce a Helene-zoom version + helene_zoom_start = pd.Timestamp("2024-09-24 00:00:00") + helene_zoom_end = pd.Timestamp("2024-09-29 23:00:00") + da_z = da_rec[(da_rec['valid_time'] >= helene_zoom_start) + & (da_rec['valid_time'] <= helene_zoom_end)] + ol_z = ol_rec[(ol_rec['valid_time'] >= helene_zoom_start) + & (ol_rec['valid_time'] <= helene_zoom_end)] + nse_da_z = nse(da_z['obs'].values, da_z['median'].values) + nse_ol_z = nse(ol_z['obs'].values, ol_z['median'].values) + print(f"NSE (Helene zoom) — DA: {nse_da_z:.3f} | OL: {nse_ol_z:.3f}") + + fig, ax = plt.subplots(figsize=(14, 5.5)) + fig.patch.set_facecolor("white") + ax.fill_between(ol_z['valid_time'], ol_z['p05'], ol_z['p95'], + color=OL_COLOR, alpha=0.18, zorder=2, + label="Open-loop [5th–95th]") + ax.plot(ol_z['valid_time'], ol_z['median'], color=OL_COLOR, + lw=1.4, linestyle='--', zorder=3, + label=f"Open-loop median (NSE={nse_ol_z:.3f})") + ax.fill_between(da_z['valid_time'], da_z['p05'], da_z['p95'], + color=DA_COLOR, alpha=0.22, zorder=4, + label="DA [5th–95th]") + ax.plot(da_z['valid_time'], da_z['median'], color=DA_COLOR, + lw=1.6, zorder=5, + label=f"DA median (NSE={nse_da_z:.3f})") + obs_z = obs.loc[helene_zoom_start:helene_zoom_end] + finite = np.isfinite(obs_z.values) + ax.scatter(obs_z.index[finite], obs_z.values[finite], + s=12, color=OBS_COLOR, linewidths=0, zorder=6, + label="USGS obs") + ax.set_xlabel("Date", fontsize=11) + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_title( + "Reconstructed forecast time series at USGS 03463300 — Helene zoom\n" + "Sep 24 – 29, 2024 (overlapping-leads pool from EnKF forecast ensemble)", + fontsize=12, + ) + ax.xaxis.set_major_locator(mdates.HourLocator(interval=12)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %H:%M")) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.25, lw=0.4) + ax.legend(loc='upper left', fontsize=9, frameon=True, framealpha=0.92) + plt.tight_layout() + out_zoom = os.path.join(out_dir, "lead_time_reconstructed_timeseries_helene.png") + plt.savefig(out_zoom, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_zoom}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4c_timeseries/run_4c_f3.sh b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4c_timeseries/run_4c_f3.sh new file mode 100644 index 00000000..01fe09ec --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/4_evaluation/4c_timeseries/run_4c_f3.sh @@ -0,0 +1,24 @@ +#!/usr/bin/env bash +# F3 — Dynamic Vrugt seeded: 4c reconstructed timeseries plots +# +# Reads routed_leadtime_da_full.parquet + routed_leadtime_openloop_full.parquet +# from the F3 forecast route dir and writes two PNGs there. +# +# Usage: +# bash run_4c_f3.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +ROUTE_DIR="/mnt/disk2/suma_helen_poster/da_results/da_forecast_dynamic_vrugt_seeded_routed" +USGS_CSV="/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +echo "[F3-4c] Reconstructed timeseries from: $ROUTE_DIR" +$TROUTE "$SCRIPT_DIR/plot_lead_time_reconstructed_timeseries.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" + +echo "[F3-4c] Done. Outputs:" +ls "$ROUTE_DIR"/lead_time_reconstructed_timeseries*.png 2>/dev/null || echo " (no PNGs found)" diff --git a/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/README.md b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/README.md new file mode 100644 index 00000000..86ea330d --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder3_dynamic_vrugt_seeded/README.md @@ -0,0 +1,68 @@ +# Folder 3 — Dynamic Vrugt Seeded + +## R Formula +``` +R(t) = (0.10 × y_obs(t))² + 0.001 × σ²_krig(t) +``` +Same Vrugt 2005 heteroscedastic formula as Folder 1, but using dynamic kriging +variance and a fixed RNG seed. See the controlled-comparison note below. + +## Design Differences vs Folder 1 + +| | Folder 1 | Folder 3 | +|---|---|---| +| R formula | Vrugt | Vrugt (same) | +| RNG seed | `hash(cat_id)` | `42` (fixed) | +| σ²_krig source | `catchment_ts_03463300_with_variance` (static ~4.17) | `catchment_ts_03463300_spliced_dyn_helene` (dynamic, ~1.6 at Helene peak) | + +> **Controlled comparison caveat:** F3 uses a different RNG seed and a different +> kriging variance dataset than F1. The Qkrig flow observations are identical across +> both datasets — only σ²_krig differs. Since the Vrugt formula weights σ²_krig by +> 0.001, the variance difference has negligible effect on R. The remaining gap in KGE +> between F1 (+0.277) and F3 (+0.261) is explained by the seed difference, not by any +> change in the R formula. **F3 vs F4 is the clean controlled comparison** for this +> obs dataset and seed. See the main [README](../README.md) for full details. + +## Key Results + +| Metric | Value | +|---|---| +| Full period KGE | **+0.261** | +| Full period NSE | **+0.485** | +| Helene KGE | **+0.189** | +| Helene NSE | **+0.372** | +| Helene peak (routed) | **693.8 m³/s (37% of USGS 1885.7)** | +| Ensemble p95 at Helene peak | 759 m³/s | +| Ensemble p50 at Helene peak | 484 m³/s | + +## Pipeline Status + +| Step | Status | Server path | +|---|---|---| +| 1. Calibrate | ✅ shared | `1400_sites_helene/da_results_dynamic_vrugt_seeded/{cat}/{cat}_best_params.json` | +| 2a/2b. Perturbation arms | ✅ | `suma_helen_poster/da_results/da_arms_dynamic_vrugt_seeded/` | +| 2c. 18hr forecast cycles | ✅ | `suma_helen_poster/da_results/da_forecast_dynamic_vrugt_seeded/` | +| 2d. 600-member ensemble | ✅ | `suma_helen_poster/da_results/da_crossed_dynamic_vrugt_seeded/` | +| 3. Route analysis | ✅ | `suma_helen_poster/da_results/dynamic_vrugt_seeded_routed/` | +| 3. Route forecast cycles | ✅ | `suma_helen_poster/da_results/da_forecast_dynamic_vrugt_seeded_routed/` | +| 3. Route ensemble | ✅ | `suma_helen_poster/da_results/da_crossed_dynamic_vrugt_seeded_routed/` | +| 4a. Error decay | ✅ | `f3_lead_time_decay_gauge_pooled.png` / `_by_regime.png` | +| 4b. Ensemble vs obs | ✅ | `f3_helene_ensemble_vs_usgs.png` / `f3_helene_ensemble_twopanel.png` | +| 4c. Reconstructed timeseries | ✅ | `f3_reconstructed_timeseries.png` / `_helene.png` | + +All server paths are under `/mnt/disk2/` unless prefixed with `1400_sites_helene/`. + +## Running This Experiment + +```bash +# On the server — all 3 stages +bash 2_assimilation/batch_run_all_f3.sh + +# Or stage by stage +bash 2_assimilation/batch_run_all_f3.sh arms # 2a/2b +bash 2_assimilation/batch_run_all_f3.sh forecast # 2c +bash 2_assimilation/batch_run_all_f3.sh ensemble # 2d + +# 4c timeseries plots +bash 4_evaluation/4c_timeseries/run_4c_f3.sh +``` diff --git a/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/2_assimilation/batch_run_all_f4.sh b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/2_assimilation/batch_run_all_f4.sh new file mode 100644 index 00000000..fdd92b0d --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/2_assimilation/batch_run_all_f4.sh @@ -0,0 +1,126 @@ +#!/usr/bin/env bash +# Folder 4 — Dynamic variance direct: full 2_assimilation pipeline +# +# Runs 3 stages for all 21 catchments: +# Stage A: 2a/2b perturbation arms (run_perturbation_da_on.py, R=krig_var, seed=42) +# Stage B: 2c 18hr forecast cycles (run_lead_time_forecast_sweep.py, --no-vrugt-r, seed=42) +# Stage C: 2d 600-member ensemble (run_crossed_ensemble.py, --direct-variance, seed=42) +# +# Usage: +# bash batch_run_all_f4.sh # runs all 3 stages +# bash batch_run_all_f4.sh arms # stage A only +# bash batch_run_all_f4.sh forecast # stage B only +# bash batch_run_all_f4.sh ensemble # stage C only + +set -euo pipefail + +STAGE="${1:-all}" + +SCRIPT_DIR="$(dirname "$0")" +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python + +# Paths +DA_SRC="/mnt/disk2/1400_sites_helene/da_results_dynamic_novrugt_seeded" +OBS_DIR="/mnt/disk2/1400_sites_helene/catchment_ts_03463300_spliced_dyn_helene" +CFE_DIR="/mnt/disk2/suma_helen_poster/cfe_py" +CONFIG_FILE="/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json" +PARAM_BOUNDS="/mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json" +FORCING_DIR="/mnt/disk2/suma_helen_poster/nwm_retro_catchment_forcings" +TEST_FORCING_DIR1="/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings" +TEST_FORCING_DIR2="/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings" + +OUT_ARMS="/mnt/disk2/1400_sites_helene/da_arms_dynamic_novrugt_seeded" +OUT_FORECAST="/mnt/disk2/1400_sites_helene/da_forecast_dynamic_novrugt_seeded" +OUT_ENSEMBLE="/mnt/disk2/1400_sites_helene/da_crossed_dynamic_novrugt_seeded" + +LOG_DIR="$HOME/f4_logs" +mkdir -p "$LOG_DIR" + +CATS=$(ls -d "$DA_SRC"/cat-* 2>/dev/null | xargs -I{} basename {}) + +# ── Stage A: 2a/2b perturbation arms ───────────────────────────────────────── +run_arms() { + echo "[F4] Stage A: perturbation arms (R=krig_var direct, seed=42)" + for cat_id in $CATS; do + params="$DA_SRC/$cat_id/${cat_id}_best_params.json" + [ -f "$params" ] || { echo " [skip] $cat_id — no params"; continue; } + mkdir -p "$OUT_ARMS/$cat_id" + cp "$params" "$OUT_ARMS/$cat_id/" + log="$LOG_DIR/arms_${cat_id}.log" + echo " [run] $cat_id" + $TROUTE "$SCRIPT_DIR/run_perturbation_da_on.py" \ + --cat-id "$cat_id" \ + --forcing-dir "$FORCING_DIR" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --param-bounds "$PARAM_BOUNDS" \ + --out-dir "$OUT_ARMS" \ + --test-forcing-dir1 "$TEST_FORCING_DIR1" \ + --test-forcing-dir2 "$TEST_FORCING_DIR2" \ + --direct-variance \ + --rng-seed 42 \ + > "$log" 2>&1 && echo " [ok ] $cat_id" || echo " [FAIL] $cat_id — see $log" + done +} + +# ── Stage B: 2c 18hr forecast cycles ───────────────────────────────────────── +run_forecast() { + echo "[F4] Stage B: 18hr forecast cycles (R=krig_var direct, seed=42)" + for cat_id in $CATS; do + params="$DA_SRC/$cat_id/${cat_id}_best_params.json" + [ -f "$params" ] || { echo " [skip] $cat_id — no params"; continue; } + mkdir -p "$OUT_FORECAST/$cat_id" + cp "$params" "$OUT_FORECAST/$cat_id/" + log="$LOG_DIR/forecast_${cat_id}.log" + echo " [run] $cat_id" + $TROUTE "$SCRIPT_DIR/run_lead_time_forecast_sweep.py" \ + --cat-id "$cat_id" \ + --forcing-dir "$FORCING_DIR" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --param-bounds "$PARAM_BOUNDS" \ + --out-dir "$OUT_FORECAST" \ + --test-forcing-dir1 "$TEST_FORCING_DIR1" \ + --test-forcing-dir2 "$TEST_FORCING_DIR2" \ + --no-vrugt-r \ + --rng-seed 42 \ + > "$log" 2>&1 && echo " [ok ] $cat_id" || echo " [FAIL] $cat_id — see $log" + done +} + +# ── Stage C: 2d 600-member crossed ensemble ─────────────────────────────────── +run_ensemble() { + echo "[F4] Stage C: 600-member crossed ensemble (R=krig_var direct, seed=42)" + for cat_id in $CATS; do + params="$DA_SRC/$cat_id/${cat_id}_best_params.json" + [ -f "$params" ] || { echo " [skip] $cat_id — no params"; continue; } + mkdir -p "$OUT_ENSEMBLE/$cat_id" + cp "$params" "$OUT_ENSEMBLE/$cat_id/" + log="$LOG_DIR/ensemble_${cat_id}.log" + echo " [run] $cat_id" + $TROUTE "$SCRIPT_DIR/run_crossed_ensemble.py" \ + --cat-id "$cat_id" \ + --forcing-dir "$FORCING_DIR" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --out-dir "$OUT_ENSEMBLE" \ + --test-forcing-dir1 "$TEST_FORCING_DIR1" \ + --test-forcing-dir2 "$TEST_FORCING_DIR2" \ + --direct-variance \ + --rng-seed 42 \ + > "$log" 2>&1 && echo " [ok ] $cat_id" || echo " [FAIL] $cat_id — see $log" + done +} + +case "$STAGE" in + arms) run_arms ;; + forecast) run_forecast ;; + ensemble) run_ensemble ;; + all) run_arms; run_forecast; run_ensemble ;; + *) echo "Usage: $0 [arms|forecast|ensemble|all]"; exit 1 ;; +esac + +echo "[F4] Done." diff --git a/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/2_assimilation/input_enkf_new.json b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/2_assimilation/input_enkf_new.json new file mode 100644 index 00000000..482233bc --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/2_assimilation/input_enkf_new.json @@ -0,0 +1 @@ +{"n": 1, "m": 1, "R": 0.07, "Q": 0, "smcmax" : 0.9394097311639178, "N": 1000, "P": 0.01, "D":2} \ No newline at end of file diff --git a/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/2_assimilation/new_EnKF.py b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/2_assimilation/new_EnKF.py new file mode 100644 index 00000000..ff945c73 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/2_assimilation/new_EnKF.py @@ -0,0 +1,53 @@ +# EnKF.py +import numpy as np + +class EnKF: + def __init__(self, n, m, R, Q, N, P, smcmax, D): + self.n = n # number of states + self.m = m # number of measurements + self.R = R # measurement noise + self.Q = Q # process noise + self.N = N # number of ensembles + self.P = P + self.state_estimates = np.zeros((self.n)) + self.covariance_matrices = np.zeros((self.n, self.n)) + self.current_step = 0 + self.smcmax = smcmax + self.D = D + self.storage_max_m = self.smcmax * self.D + self.storage_init = self.storage_max_m * 0.667 + self.ensembles = np.full((self.n, self.N), self.storage_init) + + def predict(self, F_results): + error_factor_sim = 0.005 + for i in range(self.N): + perturbation_factor_sim = np.random.standard_normal() + self.ensembles[:, i] = F_results[i] + self.ensembles[:, i] += (perturbation_factor_sim * F_results[i] * error_factor_sim) + return self.ensembles + + def update(self, yi, H_results): + if np.any(np.isnan(yi)): + return np.mean(self.ensembles, axis=1) # return current state estimate without updating + + y_ensembles = np.zeros((self.m, self.N)) + error_factor_sm = 0.03 + for i in range(self.N): + perturbation_factor_sm = np.random.standard_normal() + y_ensembles[:, i] = H_results[i] + y_ensembles[:, i] += (perturbation_factor_sm * H_results[i] * error_factor_sm) + + ensemble_mean = np.mean(self.ensembles, axis=1)[:, np.newaxis] + y_ensembles_mean = np.mean(y_ensembles, axis=1)[:, np.newaxis] + Pxy = (self.ensembles - ensemble_mean).dot((y_ensembles - y_ensembles_mean).T) / (self.N - 1) + Pyy = np.cov(y_ensembles, bias=True) + K = Pxy.dot(np.linalg.pinv(Pyy + self.R)) + self.ensembles += K.dot(yi - y_ensembles_mean) + return np.mean(self.ensembles, axis=1) + + def get_state_estimate(self): + return np.mean(self.ensembles, axis=1) + + def get_covariance(self): + ensemble_mean = np.mean(self.ensembles, axis=1)[:, np.newaxis] + return (self.ensembles - ensemble_mean).dot((self.ensembles - ensemble_mean).T) / (self.N - 1) diff --git a/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/2_assimilation/run_crossed_ensemble.py b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/2_assimilation/run_crossed_ensemble.py new file mode 100644 index 00000000..bc1cec20 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/2_assimilation/run_crossed_ensemble.py @@ -0,0 +1,463 @@ +""" +run_crossed_ensemble.py -- 4b: 600-member crossed ensemble for Helene. + +Crosses the two DA-on perturbation arms: + Forcing arm : 30 met draws (from run_perturbation_da_on.py Phase 2A) + Hydro arm : 20 state draws (from run_perturbation_da_on.py Phase 2B) + +For each Helene issue time t0: + Run 600 members = 30 forcing draws x 20 hydro-state draws. + Each member (i, j) gets: + - Initial state = perturbed DA analysis state (draw j of hydro arm) + - Met forcing = perturbed precip/PET sequence (draw i of forcing arm) + +This gives the full uncertainty envelope combining both sources — +the operational probabilistic forecast picture. + +Output (per catchment): + //_crossed_ensemble.parquet + Columns: issue_time, lead_hour, member_0000..member_0599 + +Usage: + python3 run_crossed_ensemble.py \\ + --cat-id cat-1016300 \\ + --forcing-dir \\ + --obs-dir /mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance \\ + --cfe-dir /mnt/disk2/suma_helen_poster/cfe_py \\ + --config-file /mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json \\ + --out-dir /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble \\ + --test-forcing-dir1 \\ + --test-forcing-dir2 + +Strategy: + Re-runs DA analysis (same as Phase 1 of run_perturbation_da_on.py) to get + the per-issue-time analysis state. Then generates N_FORCING sets of perturbed + forcing sequences and N_HYDRO perturbed states. Runs all 600 combinations. + + To avoid re-running Phase 1, if the Phase 1 snapshot file exists from a + prior run_perturbation_da_on.py run, it is loaded directly. + +NOTE: All special characters kept ASCII to avoid encoding issues on GPU. +""" + +import argparse +import os +import sys +import json +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +N_FORCING = 30 +N_HYDRO = 20 +N_MEMBERS = N_FORCING * N_HYDRO # 600 + +TEST_START = "2024-08-24 00:00:00" # 1 month spinup before Helene window +TEST_END = "2024-10-31 23:00:00" + +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-30 06:00:00") +FORECAST_HOURS = 18 + +PRECIP_SIGMA = 0.15 +PET_SIGMA = 0.10 +STATE_FRAC = 0.05 + +CAT_ID = None +OBS_FILE = None +CFE_CONFIG_FILE = None +TEST_FORCING_FILE = None +OUT_DIR = None +HARDCODED_R = None +DIRECT_VARIANCE = False # when True: R = krig_var directly (no Vrugt formula) +bmi_cfe = None + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(): + df = pd.read_csv(TEST_FORCING_FILE) + df["date"] = pd.to_datetime(df["time"]).dt.strftime("%Y-%m-%d %H:%M:%S") + df["total_precipitation"] = df["APCP_surface"] * 3600.0 + df["potential_evaporation"] = priestley_taylor_pet( + df["DSWRF_surface"].values, df["TMP_2maboveground"].values) + return df + + +def load_obs(): + df = pd.read_csv(OBS_FILE) + date_col = next((c for c in df.columns + if c.lower() in ("date", "time", "datetime", "timestamp")), None) + q_col = next((c for c in df.columns + if "qkrig" in c.lower() and "var" not in c.lower()), None) + var_col = next((c for c in df.columns if "var" in c.lower()), None) + df[date_col] = pd.to_datetime(df[date_col]) + df = df.sort_values(date_col).reset_index(drop=True) + date_strs = df[date_col].dt.strftime("%Y-%m-%d %H:%M:%S") + obs_dict = dict(zip(date_strs, df[q_col])) + var_vals = df[var_col] if var_col else pd.Series([0.07] * len(df)) + var_dict = dict(zip(date_strs, var_vals)) + return obs_dict, var_dict + + +def write_model_config(best_params): + """Write the BMI config file once; reused for every model instantiation.""" + with open(CFE_CONFIG_FILE) as f: + cfg = json.load(f) + cfg["forcing_file"] = TEST_FORCING_FILE + cfg["soil_params"]["bb"] = best_params["bb"] + cfg["soil_params"]["smcmax"] = best_params["smcmax"] + cfg["soil_params"]["satdk"] = best_params["satdk"] + cfg["slop"] = best_params["slop"] + cfg["max_gw_storage"] = best_params["max_gw_storage"] + cfg["expon"] = best_params["expon"] + cfg["Cgw"] = best_params["Cgw"] + cfg["K_lf"] = best_params["K_lf"] + cfg["K_nash"] = best_params["K_nash"] + cfg["partition_scheme"] = ("Schaake" if best_params["scheme"] <= 0.5 + else "Xinanjiang") + tmp_cfg = str(OUT_DIR / f"{CAT_ID}_bmi_config_crossed_stable.json") + with open(tmp_cfg, "w") as f: + json.dump(cfg, f) + return tmp_cfg + + +def make_model(cfg_path): + """Instantiate and initialize a CFE model from a pre-written config path.""" + m = bmi_cfe.BMI_CFE(cfg_file=cfg_path) + m.initialize() + return m + + +def get_state(m): + return { + "soil_m": m.soil_reservoir["storage_m"], + "gw_m": m.gw_reservoir["storage_m"], + "nash0": float(m.nash_storage[0]), + "nash1": float(m.nash_storage[1]), + } + + +def set_state(m, state): + m.soil_reservoir["storage_m"] = state["soil_m"] + m.gw_reservoir["storage_m"] = state["gw_m"] + m.nash_storage[0] = state["nash0"] + m.nash_storage[1] = state["nash1"] + + +def perturb_state(state, rng, frac=STATE_FRAC): + return { + "soil_m": max(state["soil_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "nash0": max(state["nash0"] * (1.0 + frac * rng.standard_normal()), 0.0), + "nash1": max(state["nash1"] * (1.0 + frac * rng.standard_normal()), 0.0), + } + + +def run_model_step(m, precip_mmh, pet_mmh): + m.set_value("atmosphere_water__time_integral_of_precipitation_mass_flux", + precip_mmh / 1000.0) + m.set_value("water_potential_evaporation_flux", + pet_mmh / 1000.0 / 3600.0) + m.update() + return m.get_value("land_surface_water__runoff_depth") * 1000.0 + + +def enkf_update(state, q_sim, y_obs, krig_var, rng): + if HARDCODED_R is not None: + R = HARDCODED_R + elif DIRECT_VARIANCE: + R = max(krig_var, 1e-6) + else: + R = max((0.10 * max(y_obs, 0.0)) ** 2 + 0.001 * krig_var, 1e-6) + R = max(R, 1e-6) + P_yy = max(q_sim * 0.01, 1e-6) + K = P_yy / (P_yy + R) + scale = K * (y_obs - q_sim) / max(abs(q_sim), 1e-6) + return { + "soil_m": max(state["soil_m"] * (1.0 + scale), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + scale), 1e-6), + "nash0": max(state["nash0"] + state["nash0"] * scale, 0.0), + "nash1": max(state["nash1"] + state["nash1"] * scale, 0.0), + } + + +def run_phase1_da(cfg_path, df_test, obs_dict, var_dict, rng): + """Run DA through test period; return snapshot dict {timestamp -> state}.""" + print(f" Phase 1: DA analysis through test period...") + model = make_model(cfg_path) + snapshots = {} + for _, row in df_test.iterrows(): + date_str = row["date"] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + q_sim = run_model_step(model, P, E) + y_obs = obs_dict.get(date_str, np.nan) + if not np.isnan(y_obs): + krig_var = var_dict.get(date_str, 0.07) + state = get_state(model) + updated = enkf_update(state, q_sim, y_obs, krig_var, rng) + set_state(model, updated) + date_ts = pd.Timestamp(date_str) + if HELENE_START <= date_ts <= HELENE_END: + snapshots[date_ts] = get_state(model) + model.finalize() + print(f" Phase 1 done: {len(snapshots)} issue-time snapshots") + return snapshots + + +def main(): + global CAT_ID, OBS_FILE, CFE_CONFIG_FILE, TEST_FORCING_FILE + global OUT_DIR, HARDCODED_R, DIRECT_VARIANCE, bmi_cfe + + parser = argparse.ArgumentParser() + parser.add_argument("--cat-id", required=True) + parser.add_argument("--forcing-dir", required=True) + parser.add_argument("--obs-dir", required=True) + parser.add_argument("--cfe-dir", required=True) + parser.add_argument("--config-file", required=True) + parser.add_argument("--out-dir", required=True) + parser.add_argument("--test-forcing-dir1", default=None) + parser.add_argument("--test-forcing-dir2", default=None) + parser.add_argument("--hardcoded-r", type=float, default=None, + help="Fixed R value for EnKF (e.g. 0.07). " + "Omit to use dynamic Vrugt R: R=(0.10*y)^2 + 0.001*krig_var") + parser.add_argument("--direct-variance", action="store_true", default=False, + help="Use krig_var directly as R (no Vrugt scaling)") + parser.add_argument("--rng-seed", type=int, default=None, + help="RNG seed for reproducibility (default: hash of cat-id)") + parser.add_argument("--n-forcing", type=int, default=N_FORCING, + help="Number of met forcing draws (default 30)") + parser.add_argument("--n-hydro", type=int, default=N_HYDRO, + help="Number of hydro-state draws (default 20)") + args = parser.parse_args() + + CAT_ID = args.cat_id + OBS_FILE = os.path.join(args.obs_dir, f"{CAT_ID}.csv") + CFE_CONFIG_FILE = args.config_file + HARDCODED_R = args.hardcoded_r + DIRECT_VARIANCE = args.direct_variance + OUT_DIR = Path(args.out_dir) / CAT_ID + OUT_DIR.mkdir(parents=True, exist_ok=True) + + n_fa = args.n_forcing + n_ha = args.n_hydro + n_total = n_fa * n_ha + print(f"[{CAT_ID}] Crossed ensemble: {n_fa} forcing x {n_ha} hydro = {n_total} members") + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + bmi_cfe = _bmi_cfe + + # Reuse combined forcing CSV from run_perturbation_da_on.py if it exists + # (avoids writing a duplicate large file to disk) + da_on_combined = (Path(args.out_dir).parent / "v2_perturbation_da_on" / + CAT_ID / f"{CAT_ID}_nwm_operational_combined.csv") + crossed_combined = OUT_DIR / f"{CAT_ID}_nwm_operational_combined.csv" + + if da_on_combined.exists(): + TEST_FORCING_FILE = str(da_on_combined) + print(f" Reusing combined forcing: {TEST_FORCING_FILE}") + elif crossed_combined.exists(): + TEST_FORCING_FILE = str(crossed_combined) + print(f" Reusing combined forcing: {TEST_FORCING_FILE}") + elif args.test_forcing_dir1 and args.test_forcing_dir2: + f1 = os.path.join(args.test_forcing_dir1, f"{CAT_ID}.csv") + f2 = os.path.join(args.test_forcing_dir2, f"{CAT_ID}.csv") + df1, df2 = pd.read_csv(f1), pd.read_csv(f2) + combined = (pd.concat([df1, df2], ignore_index=True) + .drop_duplicates(subset="time").sort_values("time")) + TEST_FORCING_FILE = str(crossed_combined) + combined.to_csv(TEST_FORCING_FILE, index=False) + print(f" Wrote combined forcing: {TEST_FORCING_FILE}") + else: + print("Need --test-forcing-dir1 and --test-forcing-dir2"); return + + best_params_file = OUT_DIR / f"{CAT_ID}_best_params.json" + if not best_params_file.exists(): + # Also check the DA results dir from run_perturbation_da_on.py + alt = Path(args.out_dir).parent / "v2_perturbation_da_on" / CAT_ID / f"{CAT_ID}_best_params.json" + if alt.exists(): + best_params_file = alt + else: + print(f"Missing best params at {best_params_file}"); return + with open(best_params_file) as f: + best_params = json.load(f)["best_parameters"] + + obs_dict, var_dict = load_obs() + df_forcing = load_test_forcing() + t_mask = (df_forcing["date"] >= TEST_START) & (df_forcing["date"] <= TEST_END) + df_test = df_forcing[t_mask].reset_index(drop=True) + + # Write BMI config once — reused for all 90,600 model instantiations + cfg_path = write_model_config(best_params) + print(f" BMI config written once: {cfg_path}") + + rng = np.random.default_rng(args.rng_seed if args.rng_seed is not None + else hash(CAT_ID) & 0x7fffffff) + + # Phase 1: get DA analysis snapshots + # Check if a prior run already saved them (skip re-running DA if so) + snap_cache = OUT_DIR / f"{CAT_ID}_da_snapshots.parquet" + if snap_cache.exists(): + print(f" Loading cached DA snapshots from {snap_cache}") + df_snaps = pd.read_parquet(snap_cache) + snapshots = { + pd.Timestamp(row["issue_time"]): { + "soil_m": row["soil_m"], "gw_m": row["gw_m"], + "nash0": row["nash0"], "nash1": row["nash1"], + } + for _, row in df_snaps.iterrows() + } + else: + snapshots = run_phase1_da(cfg_path, df_test, obs_dict, var_dict, rng) + + issue_times = sorted(snapshots.keys()) + t0_to_idx = {t: i for i, t in enumerate(df_test["date"].apply(pd.Timestamp))} + + # ------------------------------------------------------------------ # + # Crossed ensemble: for each issue_time, build n_fa forcing draws and # + # n_ha state draws, then run all n_fa x n_ha combinations. # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Running crossed ensemble ({n_total} members x " + f"{len(issue_times)} issue times)...") + + all_records = [] + all_hydro_draw_records = [] # (issue_time, hydro_draw_j, states) + all_forcing_draw_records = [] # (issue_time, forcing_draw_i, lead, P_pert, E_pert, scales) + + for t_idx, t0 in enumerate(issue_times): + snap = snapshots[t0] + t0_str = t0.strftime("%Y-%m-%d %H:%M:%S") + if t0 not in t0_to_idx: + continue + start_idx = t0_to_idx[t0] + + # Pre-draw n_fa forcing perturbation sequences + forcing_seqs = [] + for i in range(n_fa): + seq = [] + for lead in range(1, FORECAST_HOURS + 1): + fi = start_idx + lead + if fi >= len(df_test): + seq.append((0.0, 0.0)) + all_forcing_draw_records.append({ + "issue_time": t0_str, "forcing_draw_i": i, + "lead_hour": lead, + "P_pert_mmh": 0.0, "E_pert_mmh": 0.0, + "precip_scale": np.nan, "pet_scale": np.nan, + }) + continue + row = df_test.iloc[fi] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + mu_p = -0.5 * PRECIP_SIGMA ** 2 + P_pert = float(P * rng.lognormal(mu_p, PRECIP_SIGMA)) if P > 0 else 0.0 + E_pert = max(float(E * (1.0 + PET_SIGMA * rng.standard_normal())), 0.0) + seq.append((P_pert, E_pert)) + all_forcing_draw_records.append({ + "issue_time": t0_str, "forcing_draw_i": i, + "lead_hour": lead, + "P_pert_mmh": P_pert, "E_pert_mmh": E_pert, + "precip_scale": P_pert / P if P > 0 else np.nan, + "pet_scale": E_pert / E if E > 0 else np.nan, + }) + forcing_seqs.append(seq) + + # Pre-draw n_ha hydro-state perturbations + state_draws = [snap if j == 0 else perturb_state(snap, rng) + for j in range(n_ha)] + + # Record each hydro draw's actual initial state + for j, st in enumerate(state_draws): + all_hydro_draw_records.append({ + "issue_time": t0_str, "hydro_draw_j": j, + "soil_m": st["soil_m"], "gw_m": st["gw_m"], + "nash0": st["nash0"], "nash1": st["nash1"], + }) + + # q_matrix shape: (FORECAST_HOURS, n_fa, n_ha) + q_matrix = np.full((FORECAST_HOURS, n_fa, n_ha), np.nan) + + for i, f_seq in enumerate(forcing_seqs): + for j, state_j in enumerate(state_draws): + m = make_model(cfg_path) + set_state(m, state_j) + for lead in range(1, FORECAST_HOURS + 1): + P_pert, E_pert = f_seq[lead - 1] + q = run_model_step(m, P_pert, E_pert) + q_matrix[lead - 1, i, j] = q + m.finalize() + + # Flatten 600-member matrix into rows + for lead in range(1, FORECAST_HOURS + 1): + row_data = {"issue_time": t0_str, "lead_hour": lead} + mem_idx = 0 + for i in range(n_fa): + for j in range(n_ha): + row_data[f"member_{mem_idx:04d}"] = q_matrix[lead - 1, i, j] + mem_idx += 1 + all_records.append(row_data) + + if (t_idx + 1) % 10 == 0 or (t_idx + 1) == len(issue_times): + print(f" {t_idx + 1}/{len(issue_times)} issue times complete") + + df_out = pd.DataFrame(all_records) + out_path = OUT_DIR / f"{CAT_ID}_crossed_ensemble.parquet" + df_out.to_parquet(out_path, index=False) + print(f"Saved: {out_path}") + print(f" Shape: {df_out.shape} " + f"({df_out['issue_time'].nunique()} issue times x " + f"{FORECAST_HOURS} leads x {n_total} members)") + + # ------------------------------------------------------------------ # + # Provenance files — enable full ensemble traceability and restart # + # ------------------------------------------------------------------ # + + # 1. DA analysis snapshots — unperturbed state at each issue_time. + # Load any row and set_state() to restart from that init time. + snap_records = [ + {"issue_time": t.strftime("%Y-%m-%d %H:%M:%S"), **s} + for t, s in snapshots.items() + ] + snap_path = OUT_DIR / f"{CAT_ID}_da_snapshots.parquet" + pd.DataFrame(snap_records).to_parquet(snap_path, index=False) + print(f"Saved: {snap_path} ({len(snap_records)} snapshots)") + + # 2. Member manifest — decoder ring: member_col -> (forcing_draw_i, hydro_draw_j). + # member_k uses forcing draw k//n_ha and hydro draw k%n_ha. + manifest = [ + {"member": f"member_{k:04d}", "member_idx": k, + "forcing_draw_i": k // n_ha, "hydro_draw_j": k % n_ha} + for k in range(n_total) + ] + manifest_path = OUT_DIR / f"{CAT_ID}_member_manifest.csv" + pd.DataFrame(manifest).to_csv(manifest_path, index=False) + print(f"Saved: {manifest_path} ({n_total} members)") + + # 3. Hydro draw states — actual perturbed initial state for each of + # n_ha draws at every issue_time. Links member -> initial conditions. + hydro_path = OUT_DIR / f"{CAT_ID}_hydro_draw_states.parquet" + pd.DataFrame(all_hydro_draw_records).to_parquet(hydro_path, index=False) + print(f"Saved: {hydro_path} ({len(all_hydro_draw_records)} rows)") + + # 4. Forcing draw sequences — actual P/PET perturbation values for each + # of n_fa draws at every (issue_time, lead_hour). Links member -> forcing. + forcing_path = OUT_DIR / f"{CAT_ID}_forcing_draw_sequences.parquet" + pd.DataFrame(all_forcing_draw_records).to_parquet(forcing_path, index=False) + print(f"Saved: {forcing_path} ({len(all_forcing_draw_records)} rows)") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/2_assimilation/run_lead_time_forecast_sweep.py b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/2_assimilation/run_lead_time_forecast_sweep.py new file mode 100644 index 00000000..19f3b4d5 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/2_assimilation/run_lead_time_forecast_sweep.py @@ -0,0 +1,508 @@ +""" +Forecast lead-time evaluation. + +For each issue-time t0 in a sampling schedule across the test period, fork the +ensemble at t0 and free-run 18 hours of forecast. Two parallel trajectories are +maintained throughout the test period: + + da: Production setup. DA on, forcing perturbed, process noise on. + openloop: No DA from t=0. Forcing perturbed, process noise on. + +At each issue time t0, the state of each scenario's ensemble is copied into a +forecast ensemble that free-runs 18 hours with: + - DA off (no observations consumed during the forecast window) + - process noise off (no obs to collapse toward, so no anti-collapse needed) + - forcing perturbed (lognormal precip, Gaussian PET — proxy for forecast + precip uncertainty in lieu of actual HEFS forecasts) + +Issue-time schedule: + - Base cadence: every --base-step-h hours across the full test period + (default 6h → ~1600 issue times for the year) + - Densified to hourly across --dense-start..--dense-end if provided + (default: 2024-09-24 → 2024-09-28, the Helene window) + +Outputs (stacked across all issue times): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv + Columns: issue_time, lead_hour, valid_time, member_00 .. member_19 + (units: mm/h) + +Post-processing pipeline (separate scripts, run after this): + route_lead_time_forecasts.py — pushes each (t0, lead_hour) forecast through + T-route to the gauge + plot_lead_time_decay.py — error-vs-lead-time curve, DA vs open-loop + +Imports EnKFAssimilator from calibrate_catchment_cfe_da_v2.py so the DA math +matches production exactly. +""" + +import argparse +import os +import sys +import json +import importlib.util +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +SPINUP_START = "2023-02-01 00:00:00" +SPINUP_END = "2023-09-30 23:00:00" +TEST_START = "2023-10-01 00:00:00" +TEST_END = "2024-10-31 23:00:00" + +# Forecast lead time (hours after each issue time) +FORECAST_LEAD_HOURS = 18 + +# Default densification window — Hurricane Helene +DEFAULT_DENSE_START = "2024-09-24 00:00:00" +DEFAULT_DENSE_END = "2024-09-28 23:00:00" + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(test_forcing_file): + df = pd.read_csv(test_forcing_file) + df['date'] = pd.to_datetime(df['time']).dt.strftime('%Y-%m-%d %H:%M:%S') + df['total_precipitation'] = df['APCP_surface'] * 3600.0 + df['potential_evaporation'] = priestley_taylor_pet( + df['DSWRF_surface'].values, df['TMP_2maboveground'].values + ) + return df + + +def import_enkf_class(script_path): + """Load EnKFAssimilator from the production script without triggering its main().""" + spec = importlib.util.spec_from_file_location("prod_v2", script_path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod.EnKFAssimilator + + +# ---------------- State snapshot / restore ----------------------------------- +# +# The EnKF only ever touches these 4 states, so they're the only states that +# diverge between members. Copy these from src → dst to fork the ensemble. + +def snapshot_states(models): + return [ + { + "soil_m": float(m.soil_reservoir["storage_m"]), + "gw_m": float(m.gw_reservoir["storage_m"]), + "nash0_m": float(m.nash_storage[0]), + "nash1_m": float(m.nash_storage[1]), + } + for m in models + ] + + +def restore_states(models, snapshot): + for m, s in zip(models, snapshot): + m.soil_reservoir["storage_m"] = s["soil_m"] + m.gw_reservoir["storage_m"] = s["gw_m"] + m.nash_storage[0] = s["nash0_m"] + m.nash_storage[1] = s["nash1_m"] + + +# ---------------- Issue-time schedule ---------------------------------------- + +def build_issue_time_schedule(dates_list, base_step_h, dense_start, dense_end): + """Return a sorted list of issue-time strings sampled from dates_list. + + Every base_step_h hours across the full test period, plus every hour inside + [dense_start, dense_end] if those are provided. + """ + dates_dt = pd.to_datetime(dates_list) + selected = set() + + # Base cadence (every base_step_h hours from the first date) + base_mask = (np.arange(len(dates_dt)) % base_step_h == 0) + for d in dates_dt[base_mask]: + selected.add(d.strftime('%Y-%m-%d %H:%M:%S')) + + # Densified hourly window + if dense_start is not None and dense_end is not None: + ds = pd.Timestamp(dense_start) + de = pd.Timestamp(dense_end) + dense_mask = (dates_dt >= ds) & (dates_dt <= de) + for d in dates_dt[dense_mask]: + selected.add(d.strftime('%Y-%m-%d %H:%M:%S')) + + return sorted(selected) + + +# ---------------- CFE helpers ------------------------------------------------ + +def build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, enkf, + apply_init_perturbation=True, member_zero_clean=True): + """Build N CFE BMI instances with optional init-state perturbation.""" + def custom_load_forcing(self_cfe): + df = load_test_forcing(test_forcing_file) + self_cfe.forcing_data = df.rename(columns={"date": "time"}) + + models = [] + for i in range(N): + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.load_forcing_file = custom_load_forcing.__get__(m) + m.initialize() + if apply_init_perturbation and not (member_zero_clean and i == 0): + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + r = enkf.init_state_perturb_frac + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1 + r * enkf.rng.standard_normal()), 1e-6, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1 + r * enkf.rng.standard_normal()), 1e-6, gw_max)) + jitter = max(sm_max * 1e-6, 1e-9) + m.nash_storage[0] = max(n00 + jitter * enkf.rng.standard_normal(), 0.0) + m.nash_storage[1] = max(n10 + jitter * enkf.rng.standard_normal(), 0.0) + models.append(m) + return models + + +def step_ensemble(models, p_arr, e_arr): + """Advance each member one hour with member-specific (P, PET). Returns Q (mm/h).""" + q = np.empty(len(models), dtype=float) + for i, m in enumerate(models): + m.set_value('atmosphere_water__time_integral_of_precipitation_mass_flux', + float(p_arr[i]) / 1000.0) + m.set_value('water_potential_evaporation_flux', + float(e_arr[i]) / 1000.0 / 3600.0) + m.update() + q[i] = m.get_value('land_surface_water__runoff_depth') * 1000.0 # m/h → mm/h + return q + + +def do_sanity_check(prod_models, fcst_models, dates_list, forcing_by_date, h, N): + """One-shot verification that the 4-state snapshot is sufficient to reproduce + next-hour forecast Q. + + At main-loop hour h (after prod's hour-h step has completed including DA + + process noise), snapshot prod, restore into fcst, then step BOTH ensembles + one hour using identical unperturbed forcing. If snapshot/restore captures + everything that matters, prod's and fcst's hour-(h+1) Q should match member- + by-member to machine precision. + + Note: this mutates prod's state (h+1 step with non-perturbed forcing and no + DA), so the caller should exit the script right after. + """ + if h + 1 >= len(dates_list): + print("[sanity-check] not enough horizon for h+1 step — skipping.") + return False + snap = snapshot_states(prod_models) + restore_states(fcst_models, snap) + + next_date = dates_list[h + 1] + p_next, e_next = forcing_by_date[next_date] + p_uniform = np.full(N, p_next, dtype=float) + e_uniform = np.full(N, e_next, dtype=float) + + q_prod = step_ensemble(prod_models, p_uniform, e_uniform) + q_fcst = step_ensemble(fcst_models, p_uniform, e_uniform) + + diff = q_prod - q_fcst + max_abs = float(np.max(np.abs(diff))) + print("[sanity-check] forecast Q at h+1 with identical unperturbed forcing:") + print(f" prod first 5 members: {q_prod[:5]}") + print(f" fcst first 5 members: {q_fcst[:5]}") + print(f" max |prod - fcst|: {max_abs:.3e} mm/h") + if max_abs < 1e-9: + print(" OK: snapshot/restore preserves state to machine precision.") + return True + print(" WARN: nonzero divergence — likely a state not in the 4-state snapshot.") + return False + + +def run_forecast(fcst_models, enkf_fcst, dates_list, forcing_by_date, + t0_idx, n_lead): + """Run an n_lead-hour free-forecast starting from the current state of fcst_models. + + Forcing is perturbed per member (lognormal precip + Gaussian PET, same as + production). No DA. No process noise. + + Returns: (lead_hours, valid_times, q_matrix shape (n_lead, N)). + """ + N = len(fcst_models) + q_matrix = np.full((n_lead, N), np.nan, dtype=float) + lead_hours = [] + valid_times = [] + for lead in range(1, n_lead + 1): + idx = t0_idx + lead + if idx >= len(dates_list): + break + valid_date = dates_list[idx] + p, e = forcing_by_date[valid_date] + p_arr, e_arr = enkf_fcst.perturb_forcing(p, e) + q = step_ensemble(fcst_models, p_arr, e_arr) + q_matrix[lead - 1, :] = q + lead_hours.append(lead) + valid_times.append(valid_date) + return lead_hours, valid_times, q_matrix + + +# ---------------- Main run --------------------------------------------------- + +def run(args, bmi_cfe, EnKFAssimilator): + cat_id = args.cat_id + out_dir = Path(args.out_dir) / cat_id + out_dir.mkdir(parents=True, exist_ok=True) + + # Combined test forcing (same idiom as the per-member script) + f1 = os.path.join(args.test_forcing_dir1, f'{cat_id}.csv') + f2 = os.path.join(args.test_forcing_dir2, f'{cat_id}.csv') + df1 = pd.read_csv(f1) + df2 = pd.read_csv(f2) + combined = pd.concat([df1, df2], ignore_index=True) + combined = combined.drop_duplicates(subset='time').sort_values('time') + test_forcing_file = str(out_dir / f'{cat_id}_nwm_operational_combined.csv') + combined.to_csv(test_forcing_file, index=False) + + # Pre-staged calibrated params + best_params_file = out_dir / f'{cat_id}_best_params.json' + if not best_params_file.exists(): + print(f"Missing {best_params_file}. Pre-stage from Run 3.") + return + with open(best_params_file) as f: + best = json.load(f)['best_parameters'] + + # Build temp CFE config with calibrated params + with open(args.config_file) as f: + cfg = json.load(f) + cfg['forcing_file'] = test_forcing_file + cfg['soil_params']['bb'] = best['bb'] + cfg['soil_params']['smcmax'] = best['smcmax'] + cfg['soil_params']['satdk'] = best['satdk'] + cfg['slop'] = best['slop'] + cfg['max_gw_storage'] = best['max_gw_storage'] + cfg['expon'] = best['expon'] + cfg['Cgw'] = best['Cgw'] + cfg['K_lf'] = best['K_lf'] + cfg['K_nash'] = best['K_nash'] + cfg['partition_scheme'] = "Schaake" if best['scheme'] <= 0.5 else "Xinanjiang" + tmp_cfg = str(out_dir / f'{cat_id}_bmi_config_temp_leadtime.json') + with open(tmp_cfg, 'w') as f: + json.dump(cfg, f) + + # Three EnKF instances: + # enkf_da — runs the DA-on trajectory; consumes obs each hour + # enkf_ol — runs the open-loop trajectory; never calls update_states + # enkf_fcst — used only inside the 18-hour free-runs (perturb_forcing only) + # All three share defaults with production. Distinct RNG seeds keep them + # independent so the openloop and forecast spreads are not coupled to DA. + obs_file = os.path.join(args.obs_dir, f'{cat_id}.csv') + seed_base = args.rng_seed if args.rng_seed is not None else 0 + enkf_da = EnKFAssimilator( + n_members=args.enkf_members, obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, use_vrugt_r=(not args.no_vrugt_r), + vrugt_alpha=args.vrugt_alpha, vrugt_scale=args.vrugt_scale, + rng_seed=seed_base if args.rng_seed is not None else None, + ) + enkf_ol = EnKFAssimilator( + n_members=args.enkf_members, obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, use_vrugt_r=(not args.no_vrugt_r), + vrugt_alpha=args.vrugt_alpha, vrugt_scale=args.vrugt_scale, + rng_seed=(seed_base + 1) if args.rng_seed is not None else None, + ) + enkf_fcst = EnKFAssimilator( + n_members=args.enkf_members, obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, use_vrugt_r=(not args.no_vrugt_r), + vrugt_alpha=args.vrugt_alpha, vrugt_scale=args.vrugt_scale, + rng_seed=(seed_base + 2) if args.rng_seed is not None else None, + ) + # Optional R override — replace every per-hour obs variance with a constant. + # Matches DualEarth/new_EnKF.py (R=0.07). Set via --hardcoded-r on the CLI. + # Effect: at storm peak ≈ existing Vrugt R, at low flow ≈ 15-17× larger → + # near-zero gain at low flow where kriging obs is noisy. + if args.hardcoded_r is not None: + for enkf in (enkf_da, enkf_ol, enkf_fcst): + for date in enkf.obs_var_dict: + enkf.obs_var_dict[date] = args.hardcoded_r + print(f"[lead-time] R hardcoded to {args.hardcoded_r} mm^2/h^2 " + f"(overrides Vrugt + kriging variance formula)") + + N = enkf_da.n_members + print(f"[lead-time] {cat_id} | N={N} | DA on/off both run | " + f"forecast lead = {FORECAST_LEAD_HOURS}h") + + # Build three ensembles (init-perturbed except the forecast ensemble, whose + # state will be overwritten at each issue time before each free-run) + prod_models = build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, enkf_da) + openloop_models = build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, enkf_ol) + fcst_models = build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, + enkf_fcst, apply_init_perturbation=False) + + # ----- Spinup with perturbed forcing on both trajectories ----- + df = load_test_forcing(test_forcing_file) + sp_mask = (df['date'] >= SPINUP_START) & (df['date'] <= SPINUP_END) + df_sp = df[sp_mask] + print(f"[lead-time] spinup: {len(df_sp)} hours") + for p, e in zip(df_sp['total_precipitation'], df_sp['potential_evaporation']): + p_da, e_da = enkf_da.perturb_forcing(p, e) + step_ensemble(prod_models, p_da, e_da) + enkf_da.add_process_noise(prod_models) + p_ol, e_ol = enkf_ol.perturb_forcing(p, e) + step_ensemble(openloop_models, p_ol, e_ol) + enkf_ol.add_process_noise(openloop_models) + + # ----- Test period: step both, fork forecasts at issue times ----- + t_mask = (df['date'] >= TEST_START) & (df['date'] <= TEST_END) + df_test = df[t_mask].reset_index(drop=True) + dates_list = list(df_test['date'].values) + forcing_by_date = dict(zip( + df_test['date'].values, + zip(df_test['total_precipitation'].values, + df_test['potential_evaporation'].values))) + + issue_times = build_issue_time_schedule( + dates_list, args.base_step_h, + args.dense_start, args.dense_end) + issue_set = set(issue_times) + print(f"[lead-time] test period: {len(dates_list)} hours | " + f"{len(issue_times)} issue times " + f"(base={args.base_step_h}h, dense={args.dense_start or 'none'}" + f"..{args.dense_end or 'none'})") + + # Pre-build the date→index map so forecast windows are fast to look up + date_to_idx = {d: i for i, d in enumerate(dates_list)} + + da_rows = [] # (issue_time, lead_hour, valid_time, *member_values) + ol_rows = [] + sanity_done = False + + for h, current_date in enumerate(dates_list): + p, e = forcing_by_date[current_date] + + # --- DA trajectory step --- + p_da, e_da = enkf_da.perturb_forcing(p, e) + q_da = step_ensemble(prod_models, p_da, e_da) + enkf_da.update_states(prod_models, current_date, q_da) + enkf_da.add_process_noise(prod_models) + + # --- Open-loop trajectory step (no DA) --- + p_ol, e_ol = enkf_ol.perturb_forcing(p, e) + q_ol = step_ensemble(openloop_models, p_ol, e_ol) + enkf_ol.add_process_noise(openloop_models) + + # --- One-shot sanity check at the first scheduled issue time --- + if args.sanity_check and not sanity_done and current_date in issue_set: + ok = do_sanity_check( + prod_models, fcst_models, dates_list, forcing_by_date, h, N) + sanity_done = True + print(f"[sanity-check] exiting (re-run without --sanity-check for " + f"the full sweep). Result: {'OK' if ok else 'FAIL'}") + for m in prod_models + openloop_models + fcst_models: + m.finalize() + sys.exit(0 if ok else 1) + + # --- Fork forecasts at scheduled issue times --- + if current_date in issue_set: + for scenario_label, src_models, rows_acc in [ + ('da', prod_models, da_rows), + ('openloop', openloop_models, ol_rows), + ]: + snap = snapshot_states(src_models) + restore_states(fcst_models, snap) + leads, valids, qm = run_forecast( + fcst_models, enkf_fcst, + dates_list, forcing_by_date, + t0_idx=h, n_lead=FORECAST_LEAD_HOURS) + for k, (lead, vt) in enumerate(zip(leads, valids)): + rows_acc.append((current_date, lead, vt, *qm[k, :])) + + for m in prod_models + openloop_models + fcst_models: + m.finalize() + + # ----- Save forecast CSVs ----- + cols = ['issue_time', 'lead_hour', 'valid_time'] + [f'member_{i:02d}' for i in range(N)] + da_path = out_dir / f'{cat_id}_lead_time_forecasts_da.csv' + ol_path = out_dir / f'{cat_id}_lead_time_forecasts_openloop.csv' + pd.DataFrame(da_rows, columns=cols).to_csv(da_path, index=False) + pd.DataFrame(ol_rows, columns=cols).to_csv(ol_path, index=False) + print(f"[lead-time] saved {da_path}") + print(f"[lead-time] saved {ol_path}") + + # Issue-time schedule (small, useful for downstream scripts) + sched_path = out_dir / f'{cat_id}_lead_time_issue_times.csv' + pd.DataFrame({'issue_time': issue_times}).to_csv(sched_path, index=False) + print(f"[lead-time] saved {sched_path}") + + print(f"[lead-time] DA assimilator: updates={enkf_da.n_updates} | " + f"avg Pyy={enkf_da.avg_pyy:.6f} | " + f"mass lost={enkf_da.total_overflow_lost_mm:.3f} mm") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', required=True) + parser.add_argument('--forcing-dir', required=True) + parser.add_argument('--obs-dir', required=True) + parser.add_argument('--cfe-dir', required=True) + parser.add_argument('--config-file', required=True) + parser.add_argument('--param-bounds', required=True, + help='Kept for symmetry; not read in this script') + parser.add_argument('--out-dir', required=True) + parser.add_argument('--test-forcing-dir1', required=True) + parser.add_argument('--test-forcing-dir2', required=True) + parser.add_argument('--enkf-members', type=int, default=20) + parser.add_argument('--enkf-obs-error-std', type=float, default=0.05) + parser.add_argument('--no-vrugt-r', action='store_true') + parser.add_argument('--vrugt-alpha', type=float, default=0.10) + parser.add_argument('--vrugt-scale', type=float, default=0.001) + parser.add_argument('--rng-seed', type=int, default=None) + parser.add_argument('--base-step-h', type=int, default=6, + help='Issue-time cadence in hours across the full test ' + 'period (default 6h)') + parser.add_argument('--dense-start', type=str, default=DEFAULT_DENSE_START, + help='Start of hourly-densification window ' + '(default 2024-09-24 00:00:00 — Helene)') + parser.add_argument('--dense-end', type=str, default=DEFAULT_DENSE_END, + help='End of hourly-densification window ' + '(default 2024-09-28 23:00:00 — Helene)') + parser.add_argument('--sanity-check', action='store_true', + help='At the first issue time, fork fcst_models and step ' + 'both prod and fcst one hour with identical unperturbed ' + 'forcing. Confirms the 4-state snapshot is sufficient ' + 'to reproduce next-hour Q to machine precision, then ' + 'exits. Use before kicking off the full sweep.') + parser.add_argument('--hardcoded-r', type=float, default=None, + help='If set, override the Vrugt + kriging-variance R formula ' + 'and use this constant value (mm^2/h^2) at every hour. ' + 'Matches the DualEarth/new_EnKF.py reference (R=0.07). ' + 'Suppresses DA at low flow where kriging is noisy while ' + 'preserving DA strength at storm peaks. Applied to all ' + 'three EnKF instances (da, openloop, fcst) for consistency.') + parser.add_argument('--prod-script', default=None, + help='Path to calibrate_catchment_cfe_da_v2.py for ' + 'importing EnKFAssimilator. Defaults to next-to-this-file.') + args = parser.parse_args() + + here = os.path.dirname(os.path.abspath(__file__)) + prod_script = args.prod_script or os.path.join(here, 'calibrate_catchment_cfe_da_v2.py') + if not os.path.exists(prod_script): + raise FileNotFoundError( + f"Could not find production script at {prod_script}. " + f"Pass --prod-script to override.") + EnKFAssimilator = import_enkf_class(prod_script) + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + + run(args, _bmi_cfe, EnKFAssimilator) + + +if __name__ == '__main__': + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/2_assimilation/run_perturbation_da_on.py b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/2_assimilation/run_perturbation_da_on.py new file mode 100644 index 00000000..6ef27e96 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/2_assimilation/run_perturbation_da_on.py @@ -0,0 +1,400 @@ +""" +run_perturbation_da_on.py — 2a/2b: ensemble spread with DA ON. + +Runs two separated perturbation arms for the Helene window (Sep 24-30), +with DA actively running during the analysis period: + + ARM A — forcing only (30 members): + DA analyzes states using Qkrig at every hour (DA ON). + At each issue_time: fork 30 members with ONLY met forcing perturbed + (different precip/PET draw per member). Hydro states = DA analysis + mean (same for all 30). Free-run 18 hours, no further DA. + + ARM B — hydro states only (20 members): + Same DA analysis. At each issue_time: fork 20 members with ONLY + hydro states perturbed (multiplicative noise on soil/GW/Nash from + the DA analysis ensemble). Forcing = deterministic. Free-run 18 hours. + +These two arms isolate "where does forecast spread come from WITH DA?" +— figures for the assimilation script section. + +Output (per catchment, per arm): + //_da_forcing_arm.csv + //_da_hydro_arm.csv + Columns: issue_time, lead_hour, member_00..member_29 (forcing arm) + issue_time, lead_hour, member_00..member_19 (hydro arm) + +Usage (run for each catchment, both arms): + python3 run_perturbation_da_on.py \\ + --cat-id cat-1016300 \\ + --forcing-dir \\ + --obs-dir /mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance \\ + --cfe-dir /mnt/disk2/suma_helen_poster/cfe_py \\ + --config-file /mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json \\ + --param-bounds /mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json \\ + --out-dir /mnt/disk2/suma_helen_poster/da_results/v2_perturbation_da_on \\ + --test-forcing-dir1 \\ + --test-forcing-dir2 \\ + --hardcoded-r 0.07 +""" + +import argparse +import os +import sys +import json +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +N_FORCING = 30 # members in forcing arm +N_HYDRO = 20 # members in hydro-state arm + +# Analysis period — run DA through entire test period +TEST_START = "2024-08-24 00:00:00" # 1 month spinup before Helene window +TEST_END = "2024-10-31 23:00:00" + +# Issue times for the Helene window (hourly, Sep 24-30) +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-30 06:00:00") +FORECAST_HOURS = 18 + +# Perturbation magnitudes (same as production) +PRECIP_SIGMA = 0.15 +PET_SIGMA = 0.10 +STATE_FRAC = 0.05 # hydro-state perturbation fraction + +# Set at runtime +CAT_ID = None +OBS_FILE = None +CFE_CONFIG_FILE = None +TEST_FORCING_FILE = None +OUT_DIR = None +HARDCODED_R = None +DIRECT_VARIANCE = False +bmi_cfe = None + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(): + df = pd.read_csv(TEST_FORCING_FILE) + df["date"] = pd.to_datetime(df["time"]).dt.strftime("%Y-%m-%d %H:%M:%S") + df["total_precipitation"] = df["APCP_surface"] * 3600.0 + df["potential_evaporation"] = priestley_taylor_pet( + df["DSWRF_surface"].values, df["TMP_2maboveground"].values) + return df + + +def load_obs(): + df = pd.read_csv(OBS_FILE) + date_col = next((c for c in df.columns + if c.lower() in ("date", "time", "datetime", "timestamp")), None) + if date_col is None: + raise ValueError(f"No date column in {OBS_FILE}. Columns: {df.columns.tolist()}") + q_col = next((c for c in df.columns if "qkrig" in c.lower() and "var" not in c.lower()), None) + if q_col is None: + raise ValueError(f"No qkrig column in {OBS_FILE}. Columns: {df.columns.tolist()}") + var_col = next((c for c in df.columns if "var" in c.lower()), None) + df[date_col] = pd.to_datetime(df[date_col]) + df = df.sort_values(date_col).reset_index(drop=True) + date_strs = df[date_col].dt.strftime("%Y-%m-%d %H:%M:%S") + obs_dict = dict(zip(date_strs, df[q_col])) + var_vals = df[var_col] if var_col else pd.Series([0.07] * len(df)) + var_dict = dict(zip(date_strs, var_vals)) + print(f" Obs: date='{date_col}', q='{q_col}', var='{var_col}'") + return obs_dict, var_dict + + +def make_model(best_params): + with open(CFE_CONFIG_FILE) as f: + cfg = json.load(f) + cfg["forcing_file"] = TEST_FORCING_FILE + cfg["soil_params"]["bb"] = best_params["bb"] + cfg["soil_params"]["smcmax"] = best_params["smcmax"] + cfg["soil_params"]["satdk"] = best_params["satdk"] + cfg["slop"] = best_params["slop"] + cfg["max_gw_storage"] = best_params["max_gw_storage"] + cfg["expon"] = best_params["expon"] + cfg["Cgw"] = best_params["Cgw"] + cfg["K_lf"] = best_params["K_lf"] + cfg["K_nash"] = best_params["K_nash"] + cfg["partition_scheme"] = ("Schaake" if best_params["scheme"] <= 0.5 + else "Xinanjiang") + tmp_cfg = str(OUT_DIR / f"{CAT_ID}_bmi_config_temp_da_on.json") + with open(tmp_cfg, "w") as f: + json.dump(cfg, f) + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.initialize() + return m + + +def get_state(m): + return { + "soil_m": m.soil_reservoir["storage_m"], + "gw_m": m.gw_reservoir["storage_m"], + "nash0": float(m.nash_storage[0]), + "nash1": float(m.nash_storage[1]), + } + + +def set_state(m, state): + m.soil_reservoir["storage_m"] = state["soil_m"] + m.gw_reservoir["storage_m"] = state["gw_m"] + m.nash_storage[0] = state["nash0"] + m.nash_storage[1] = state["nash1"] + + +def perturb_state(state, rng, frac=STATE_FRAC): + sm_max = 1.0 # placeholder — will be clipped in set_state; actual max enforced by model + out = { + "soil_m": max(state["soil_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + frac * rng.standard_normal()), 1e-6), + "nash0": max(state["nash0"] * (1.0 + frac * rng.standard_normal()), 0.0), + "nash1": max(state["nash1"] * (1.0 + frac * rng.standard_normal()), 0.0), + } + return out + + +def run_model_step(m, precip_mmh, pet_mmh): + m.set_value("atmosphere_water__time_integral_of_precipitation_mass_flux", + precip_mmh / 1000.0) + m.set_value("water_potential_evaporation_flux", + pet_mmh / 1000.0 / 3600.0) + m.update() + return m.get_value("land_surface_water__runoff_depth") * 1000.0 # m/h -> mm/h + + +def enkf_update(state, q_sim, y_obs, krig_var, rng, n_members=1): + """Simple scalar EnKF update on the mean analysis state.""" + if HARDCODED_R is not None: + R = HARDCODED_R + elif DIRECT_VARIANCE: + R = max(krig_var, 1e-6) + else: + alpha = 0.10 + R = (alpha * max(y_obs, 0.0)) ** 2 + 0.001 * krig_var + R = max(R, 1e-6) + P_yy = max(q_sim * 0.01, 1e-6) # rough prior variance on Q + K = P_yy / (P_yy + R) + innovation = y_obs - q_sim + # Apply correction to each state proportionally (simple scalar gain) + scale = K * innovation / max(abs(q_sim), 1e-6) + updated = { + "soil_m": max(state["soil_m"] * (1.0 + scale), 1e-6), + "gw_m": max(state["gw_m"] * (1.0 + scale), 1e-6), + "nash0": max(state["nash0"] + state["nash0"] * scale, 0.0), + "nash1": max(state["nash1"] + state["nash1"] * scale, 0.0), + } + return updated + + +def main(): + global CAT_ID, OBS_FILE, CFE_CONFIG_FILE, TEST_FORCING_FILE + global OUT_DIR, HARDCODED_R, DIRECT_VARIANCE, bmi_cfe + + parser = argparse.ArgumentParser() + parser.add_argument("--cat-id", required=True) + parser.add_argument("--forcing-dir", required=True) + parser.add_argument("--obs-dir", required=True) + parser.add_argument("--cfe-dir", required=True) + parser.add_argument("--config-file", required=True) + parser.add_argument("--param-bounds", required=True) + parser.add_argument("--out-dir", required=True) + parser.add_argument("--test-forcing-dir1", default=None) + parser.add_argument("--test-forcing-dir2", default=None) + parser.add_argument("--hardcoded-r", type=float, default=None, + help="Fixed R (e.g. 0.07). Omit → Vrugt formula.") + parser.add_argument("--direct-variance", action="store_true", default=False, + help="Use krig_var directly as R (no Vrugt scaling)") + parser.add_argument("--rng-seed", type=int, default=None) + args = parser.parse_args() + + CAT_ID = args.cat_id + OBS_FILE = os.path.join(args.obs_dir, f"{CAT_ID}.csv") + CFE_CONFIG_FILE = args.config_file + HARDCODED_R = args.hardcoded_r + DIRECT_VARIANCE = args.direct_variance + OUT_DIR = Path(args.out_dir) / CAT_ID + OUT_DIR.mkdir(parents=True, exist_ok=True) + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + bmi_cfe = _bmi_cfe + + # Build combined test forcing + if args.test_forcing_dir1 and args.test_forcing_dir2: + f1 = os.path.join(args.test_forcing_dir1, f"{CAT_ID}.csv") + f2 = os.path.join(args.test_forcing_dir2, f"{CAT_ID}.csv") + df1, df2 = pd.read_csv(f1), pd.read_csv(f2) + combined = (pd.concat([df1, df2], ignore_index=True) + .drop_duplicates(subset="time").sort_values("time")) + TEST_FORCING_FILE = str(OUT_DIR / f"{CAT_ID}_nwm_operational_combined.csv") + combined.to_csv(TEST_FORCING_FILE, index=False) + else: + print("Need --test-forcing-dir1 and --test-forcing-dir2"); return + + best_params_file = OUT_DIR / f"{CAT_ID}_best_params.json" + if not best_params_file.exists(): + print(f"Missing best params at {best_params_file}"); return + with open(best_params_file) as f: + best_params = json.load(f)["best_parameters"] + + obs_dict, var_dict = load_obs() + df_forcing = load_test_forcing() + t_mask = (df_forcing["date"] >= TEST_START) & (df_forcing["date"] <= TEST_END) + df_test = df_forcing[t_mask].reset_index(drop=True) + + rng = np.random.default_rng(hash(CAT_ID) & 0x7fffffff) + + # ------------------------------------------------------------------ # + # Phase 1: Run DA through test period, snapshot analysis state at # + # each Helene issue time # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Phase 1: DA analysis through test period...") + model = make_model(best_params) + + helene_snapshots = {} # issue_time -> analysis state dict + + for h, row in df_test.iterrows(): + date_str = row["date"] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + q_sim = run_model_step(model, P, E) + + date_ts = pd.Timestamp(date_str) + + # DA update if obs available + y_obs = obs_dict.get(date_str, np.nan) + if not np.isnan(y_obs): + krig_var = var_dict.get(date_str, 0.07) + state = get_state(model) + updated = enkf_update(state, q_sim, y_obs, krig_var, rng) + set_state(model, updated) + + # Snapshot if this is a Helene issue time + if HELENE_START <= date_ts <= HELENE_END: + helene_snapshots[date_ts] = get_state(model) + + model.finalize() + print(f" Snapshots saved: {len(helene_snapshots)} issue times") + + # ------------------------------------------------------------------ # + # Phase 2A: Forcing arm — 30 members, fixed analyzed state, varying # + # met forcing # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Phase 2A: Forcing arm ({N_FORCING} members)...") + forcing_records = [] + + issue_times = sorted(helene_snapshots.keys()) + t0_to_idx = {t: i for i, t in enumerate(df_test["date"].apply(pd.Timestamp))} + + for t0 in issue_times: + snap = helene_snapshots[t0] + t0_str = t0.strftime("%Y-%m-%d %H:%M:%S") + + # Find index in df_test + if t0 not in t0_to_idx: + continue + start_idx = t0_to_idx[t0] + + q_matrix = np.full((FORECAST_HOURS, N_FORCING), np.nan) + + for mem in range(N_FORCING): + m = make_model(best_params) + set_state(m, snap) # all members start from same DA analysis state + + for lead in range(1, FORECAST_HOURS + 1): + fi = start_idx + lead + if fi >= len(df_test): + break + row = df_test.iloc[fi] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + # Perturb forcing only + mu_p = -0.5 * PRECIP_SIGMA ** 2 + P_pert = float(P * rng.lognormal(mu_p, PRECIP_SIGMA)) if P > 0 else 0.0 + E_pert = max(float(E * (1.0 + PET_SIGMA * rng.standard_normal())), 0.0) + q = run_model_step(m, P_pert, E_pert) + q_matrix[lead - 1, mem] = q + + m.finalize() + + row_data = {"issue_time": t0_str} + for lead in range(1, FORECAST_HOURS + 1): + row_data["lead_hour"] = lead + for mem in range(N_FORCING): + row_data[f"member_{mem:02d}"] = q_matrix[lead - 1, mem] + forcing_records.append(dict(row_data)) + + df_forcing_arm = pd.DataFrame(forcing_records) + out_a = OUT_DIR / f"{CAT_ID}_da_forcing_arm.csv" + df_forcing_arm.to_csv(out_a, index=False) + print(f" Saved: {out_a}") + + # ------------------------------------------------------------------ # + # Phase 2B: Hydro-state arm — 20 members, deterministic forcing, # + # perturbed analysis states # + # ------------------------------------------------------------------ # + print(f"[{CAT_ID}] Phase 2B: Hydro-state arm ({N_HYDRO} members)...") + hydro_records = [] + + for t0 in issue_times: + snap = helene_snapshots[t0] + t0_str = t0.strftime("%Y-%m-%d %H:%M:%S") + + if t0 not in t0_to_idx: + continue + start_idx = t0_to_idx[t0] + + q_matrix = np.full((FORECAST_HOURS, N_HYDRO), np.nan) + + for mem in range(N_HYDRO): + # Perturb the DA analysis state for this member + if mem == 0: + state_m = snap # member 0 = unperturbed analysis mean + else: + state_m = perturb_state(snap, rng) + + m = make_model(best_params) + set_state(m, state_m) + + for lead in range(1, FORECAST_HOURS + 1): + fi = start_idx + lead + if fi >= len(df_test): + break + row = df_test.iloc[fi] + P = float(row["total_precipitation"]) + E = float(row["potential_evaporation"]) + # Deterministic forcing + q = run_model_step(m, P, E) + q_matrix[lead - 1, mem] = q + + m.finalize() + + row_data = {"issue_time": t0_str} + for lead in range(1, FORECAST_HOURS + 1): + row_data["lead_hour"] = lead + for mem in range(N_HYDRO): + row_data[f"member_{mem:02d}"] = q_matrix[lead - 1, mem] + hydro_records.append(dict(row_data)) + + df_hydro_arm = pd.DataFrame(hydro_records) + out_b = OUT_DIR / f"{CAT_ID}_da_hydro_arm.csv" + df_hydro_arm.to_csv(out_b, index=False) + print(f" Saved: {out_b}") + print(f"[{CAT_ID}] Done.") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/2_assimilation/run_perturbation_sensitivity.py b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/2_assimilation/run_perturbation_sensitivity.py new file mode 100644 index 00000000..57b8d185 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/2_assimilation/run_perturbation_sensitivity.py @@ -0,0 +1,293 @@ +""" +Perturbation-source sensitivity analysis for the multi-catchment CFE DA experiment. + +Ensemble-spread attribution: run 20 CFE members through the test +period with ONLY ONE perturbation source active at a time, no DA, no spinup. + +Three sub-experiments (selected via --source): + init : perturb initial states once at t=0 (multiplicative ~5%) + forcing : perturb hourly precip (lognormal sigma=0.15) and PET (Gaussian sigma=0.10) + process : apply per-hour process noise on states (soil 0.2%, GW 0.15%, Nash 0.5%) + +Each sub-experiment runs N=20 CFE members. DA is always OFF. Spinup is skipped — each +member starts from CFE's default initial state (modified by the perturbation if --source +init). This is the cleanest setup for attributing test-period spread to its source. + +Output: + //_sensitivity_.csv + Columns: date, member_00, member_01, ..., member_19 + +Each member column is that member's hourly Q_sim (mm/h) over the test period. + +Reads the same calibrated best_params.json as the production script. Does not need +the kriging obs file (no DA), but takes the same --obs-dir argument so it can also +log Q_obs alongside for plotting reference. + +Usage: + python3 run_perturbation_sensitivity.py \ + --cat-id cat-1016300 \ + --source forcing \ + --forcing-dir \ + --obs-dir \ + --cfe-dir \ + --config-file \ + --param-bounds \ + --out-dir \ + --test-forcing-dir1 \ + --test-forcing-dir2 +""" + +import argparse +import os +import sys +import json +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 +N_MEMBERS = 20 + +# We only run the test window for this sensitivity analysis. +TEST_START = "2023-10-01 00:00:00" +TEST_END = "2024-10-31 23:00:00" + +# Perturbation magnitudes — same as production defaults +INIT_STATE_FRAC = 0.05 +PRECIP_PERTURB = 0.15 # lognormal sigma +PET_PERTURB = 0.10 # Gaussian sigma +SOIL_PROC_NOISE = 0.002 +GW_PROC_NOISE = 0.0015 +NASH_PROC_NOISE = 0.005 + +# Set at runtime +CAT_ID = None +FORCING_FILE = None +TEST_FORCING_FILE = None +OBS_FILE = None +CFE_CONFIG_FILE = None +OUT_DIR = None +bmi_cfe = None + + +# ---------- Forcing helpers (same as production) ---------- +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(): + df = pd.read_csv(TEST_FORCING_FILE) + df['date'] = pd.to_datetime(df['time']).dt.strftime('%Y-%m-%d %H:%M:%S') + df['total_precipitation'] = df['APCP_surface'] * 3600.0 # kg/m^2/s -> mm/h + df['potential_evaporation'] = priestley_taylor_pet( + df['DSWRF_surface'].values, + df['TMP_2maboveground'].values, + ) + return df + + +# ---------- Perturbation helpers ---------- +def perturb_precip_lognormal(P, rng, sigma=PRECIP_PERTURB): + """Lognormal multiplier with mean 1; preserves zero precip.""" + mu = -0.5 * sigma ** 2 + return float(P * rng.lognormal(mu, sigma)) + + +def perturb_pet_gaussian(E, rng, sigma=PET_PERTURB): + """Gaussian multiplicative noise on PET, clipped at 0.""" + return max(float(E * (1.0 + sigma * rng.standard_normal())), 0.0) + + +def apply_init_state_perturb(m, rng, frac=INIT_STATE_FRAC): + """One-time multiplicative perturbation on the four states (called per member at t=0).""" + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1.0 + frac * rng.standard_normal()), 1e-6, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1.0 + frac * rng.standard_normal()), 1e-6, gw_max)) + # Nash buckets often start at 0; small additive jitter so they're not identical. + jitter = max(sm_max * 1e-6, 1e-9) + m.nash_storage[0] = max(n00 + jitter * rng.standard_normal(), 0.0) + m.nash_storage[1] = max(n10 + jitter * rng.standard_normal(), 0.0) + + +def apply_process_noise(m, rng): + """Per-timestep process noise on all 4 states for one member.""" + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + nash_floor = max(sm_max * 1e-4, 1e-7) + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1.0 + SOIL_PROC_NOISE * rng.standard_normal()), 0.0, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1.0 + GW_PROC_NOISE * rng.standard_normal()), 0.0, gw_max)) + m.nash_storage[0] = max( + n00 * (1.0 + NASH_PROC_NOISE * rng.standard_normal()) + + nash_floor * rng.standard_normal(), 0.0) + m.nash_storage[1] = max( + n10 * (1.0 + NASH_PROC_NOISE * rng.standard_normal()) + + nash_floor * rng.standard_normal(), 0.0) + + +# ---------- Main per-source run ---------- +def run_sensitivity(source, best_param_dict): + """Run N members through the test period with only `source` perturbation active.""" + + def custom_load_forcing(self_cfe): + df = load_test_forcing() + self_cfe.forcing_data = df.rename(columns={"date": "time"}) + + # Build a per-catchment temporary CFE config with the calibrated params + with open(CFE_CONFIG_FILE) as f: + cfg = json.load(f) + cfg['forcing_file'] = TEST_FORCING_FILE + cfg['soil_params']['bb'] = best_param_dict['bb'] + cfg['soil_params']['smcmax'] = best_param_dict['smcmax'] + cfg['soil_params']['satdk'] = best_param_dict['satdk'] + cfg['slop'] = best_param_dict['slop'] + cfg['max_gw_storage'] = best_param_dict['max_gw_storage'] + cfg['expon'] = best_param_dict['expon'] + cfg['Cgw'] = best_param_dict['Cgw'] + cfg['K_lf'] = best_param_dict['K_lf'] + cfg['K_nash'] = best_param_dict['K_nash'] + cfg['partition_scheme'] = "Schaake" if best_param_dict['scheme'] <= 0.5 else "Xinanjiang" + + tmp_cfg = str(OUT_DIR / f'{CAT_ID}_bmi_config_temp_sensitivity_{source}.json') + with open(tmp_cfg, 'w') as f: + json.dump(cfg, f) + + # Deterministic seed per (catchment, source) so each run is reproducible + seed = hash((CAT_ID, source)) & 0x7fffffff + rng = np.random.default_rng(seed) + + print(f"[sensitivity] {CAT_ID} | source={source} | N={N_MEMBERS} | seed={seed}") + + # Build N members. Initial-state perturbation only happens when source == 'init'. + models = [] + for i in range(N_MEMBERS): + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.load_forcing_file = custom_load_forcing.__get__(m) + m.initialize() + if source == 'init' and i > 0: + apply_init_state_perturb(m, rng) + models.append(m) + + # Skip spinup. Go directly to the test period. + df = load_test_forcing() + t_mask = (df['date'] >= TEST_START) & (df['date'] <= TEST_END) + df_test = df[t_mask] + + # Pre-allocate per-member Q arrays + n_hours = len(df_test) + q_matrix = np.full((n_hours, N_MEMBERS), np.nan, dtype=float) + dates_out = df_test['date'].values + + for h, (p, e) in enumerate(zip(df_test['total_precipitation'], + df_test['potential_evaporation'])): + for i, m in enumerate(models): + if source == 'forcing': + p_i = perturb_precip_lognormal(p, rng) + e_i = perturb_pet_gaussian(e, rng) + else: + p_i = float(p) + e_i = float(e) + m.set_value('atmosphere_water__time_integral_of_precipitation_mass_flux', p_i / 1000) + m.set_value('water_potential_evaporation_flux', e_i / 1000 / 3600) + m.update() + q_matrix[h, i] = m.get_value('land_surface_water__runoff_depth') * 1000 # m/h -> mm/h + + # Process noise applied after each hour, but only when source == 'process' + if source == 'process': + for m in models: + apply_process_noise(m, rng) + + for m in models: + m.finalize() + + # Write per-member CSV: date + 20 columns + out = {'date': dates_out} + for i in range(N_MEMBERS): + out[f'member_{i:02d}'] = q_matrix[:, i] + df_out = pd.DataFrame(out) + out_path = OUT_DIR / f'{CAT_ID}_sensitivity_{source}.csv' + df_out.to_csv(out_path, index=False) + print(f"[sensitivity] saved {out_path}") + print(f"[sensitivity] mean ensemble spread (std across members) over test period: " + f"{float(q_matrix.std(axis=1).mean()):.5f} mm/h") + + +def main(): + global CAT_ID, FORCING_FILE, TEST_FORCING_FILE, OBS_FILE, CFE_CONFIG_FILE + global OUT_DIR, bmi_cfe + + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', required=True) + parser.add_argument('--source', required=True, + choices=['init', 'forcing', 'process'], + help="Which perturbation source to isolate") + parser.add_argument('--forcing-dir', required=True) + parser.add_argument('--obs-dir', required=True, + help='Kept for symmetry with production script; not used since DA is off') + parser.add_argument('--cfe-dir', required=True) + parser.add_argument('--config-file', required=True) + parser.add_argument('--param-bounds', required=True, + help='Kept for symmetry; not read in this script') + parser.add_argument('--out-dir', required=True) + parser.add_argument('--test-forcing-dir1', default=None) + parser.add_argument('--test-forcing-dir2', default=None) + args = parser.parse_args() + + CAT_ID = args.cat_id + FORCING_FILE = os.path.join(args.forcing_dir, f'{CAT_ID}.csv') + OBS_FILE = os.path.join(args.obs_dir, f'{CAT_ID}.csv') + CFE_CONFIG_FILE = args.config_file + OUT_DIR = Path(args.out_dir) / CAT_ID + OUT_DIR.mkdir(parents=True, exist_ok=True) + + # Build combined test forcing CSV (same idiom as production) + if args.test_forcing_dir1 and args.test_forcing_dir2: + f1 = os.path.join(args.test_forcing_dir1, f'{CAT_ID}.csv') + f2 = os.path.join(args.test_forcing_dir2, f'{CAT_ID}.csv') + if os.path.exists(f1) and os.path.exists(f2): + df1 = pd.read_csv(f1) + df2 = pd.read_csv(f2) + combined = pd.concat([df1, df2], ignore_index=True) + combined = combined.drop_duplicates(subset='time').sort_values('time') + TEST_FORCING_FILE = str(OUT_DIR / f'{CAT_ID}_nwm_operational_combined.csv') + combined.to_csv(TEST_FORCING_FILE, index=False) + else: + print(f"Warning: test forcing files not found for {CAT_ID}, exiting") + return + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + bmi_cfe = _bmi_cfe + + best_params_file = OUT_DIR / f'{CAT_ID}_best_params.json' + if not best_params_file.exists(): + print(f"No best params file at {best_params_file}. Pre-stage from Run 3.") + return + with open(best_params_file) as f: + saved = json.load(f) + best_param_dict = saved['best_parameters'] + + run_sensitivity(args.source, best_param_dict) + + +if __name__ == '__main__': + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/2_assimilation/run_production_per_member.py b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/2_assimilation/run_production_per_member.py new file mode 100644 index 00000000..d9c1c4ac --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/2_assimilation/run_production_per_member.py @@ -0,0 +1,314 @@ +""" +Run the production DA pipeline for ONE catchment and save per-member streamflow. + +Identical to calibrate_catchment_cfe_da_v2.py in algorithm (init perturbation + +forcing perturbation + process noise + true EnKF with Vrugt R), but writes one +column per ensemble member instead of only the ensemble mean. This is for +the per-member factor-decomposition plot. + +To keep the comparison apples-to-apples with the sensitivity runs (which skip +spinup), this script also SKIPS the spinup loop. Each member starts directly +from CFE's default initial state, plus the t=0 init perturbation. + +Output: + //_production_per_member.csv + Columns: date, member_00, member_01, ..., member_19 + +Imports EnKFAssimilator from calibrate_catchment_cfe_da_v2 to guarantee the +DA math matches production exactly. + +Usage: + python3 run_production_per_member.py \ + --cat-id cat-1016300 \ + --forcing-dir /mnt/.../nwm_retro_catchment_forcings \ + --obs-dir /mnt/.../catchment_ts_03463300_with_variance \ + --cfe-dir /mnt/.../cfe_py \ + --config-file /mnt/.../cat_03463300_bmi_config_cfe.json \ + --param-bounds /mnt/.../CFE_parameter_bounds.json \ + --out-dir /mnt/.../v2_production_per_member \ + --test-forcing-dir1 ... \ + --test-forcing-dir2 ... \ + --enkf-members 20 +""" + +import argparse +import os +import sys +import json +import importlib.util +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +TEST_START = "2023-10-01 00:00:00" +TEST_END = "2024-10-31 23:00:00" + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(test_forcing_file): + df = pd.read_csv(test_forcing_file) + df['date'] = pd.to_datetime(df['time']).dt.strftime('%Y-%m-%d %H:%M:%S') + df['total_precipitation'] = df['APCP_surface'] * 3600.0 + df['potential_evaporation'] = priestley_taylor_pet( + df['DSWRF_surface'].values, df['TMP_2maboveground'].values + ) + return df + + +def import_enkf_class(script_path): + """Load EnKFAssimilator from the production script without triggering its main().""" + spec = importlib.util.spec_from_file_location("prod_v2", script_path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod.EnKFAssimilator + + +def run(args, bmi_cfe, EnKFAssimilator): + cat_id = args.cat_id + out_dir = Path(args.out_dir) / cat_id + out_dir.mkdir(parents=True, exist_ok=True) + + # Combined test forcing + f1 = os.path.join(args.test_forcing_dir1, f'{cat_id}.csv') + f2 = os.path.join(args.test_forcing_dir2, f'{cat_id}.csv') + df1 = pd.read_csv(f1) + df2 = pd.read_csv(f2) + combined = pd.concat([df1, df2], ignore_index=True) + combined = combined.drop_duplicates(subset='time').sort_values('time') + test_forcing_file = str(out_dir / f'{cat_id}_nwm_operational_combined.csv') + combined.to_csv(test_forcing_file, index=False) + + # Load best_params + best_params_file = out_dir / f'{cat_id}_best_params.json' + if not best_params_file.exists(): + print(f"Missing {best_params_file}. Pre-stage from Run 3.") + return + with open(best_params_file) as f: + saved = json.load(f) + best = saved['best_parameters'] + + # Build temp CFE config with calibrated params + with open(args.config_file) as f: + cfg = json.load(f) + cfg['forcing_file'] = test_forcing_file + cfg['soil_params']['bb'] = best['bb'] + cfg['soil_params']['smcmax'] = best['smcmax'] + cfg['soil_params']['satdk'] = best['satdk'] + cfg['slop'] = best['slop'] + cfg['max_gw_storage'] = best['max_gw_storage'] + cfg['expon'] = best['expon'] + cfg['Cgw'] = best['Cgw'] + cfg['K_lf'] = best['K_lf'] + cfg['K_nash'] = best['K_nash'] + cfg['partition_scheme'] = "Schaake" if best['scheme'] <= 0.5 else "Xinanjiang" + tmp_cfg = str(out_dir / f'{cat_id}_bmi_config_temp_per_member.json') + with open(tmp_cfg, 'w') as f: + json.dump(cfg, f) + + # Build the EnKF assimilator (same defaults as production) + obs_file = os.path.join(args.obs_dir, f'{cat_id}.csv') + enkf = EnKFAssimilator( + n_members=args.enkf_members, + obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, + use_vrugt_r=(not args.no_vrugt_r) and (not getattr(args, 'direct_variance', False)), + vrugt_alpha=args.vrugt_alpha, + vrugt_scale=args.vrugt_scale, + rng_seed=args.rng_seed, + ) + N = enkf.n_members + print(f"[per-member] {cat_id} | N={N} | use_vrugt_r={enkf.use_vrugt_r}") + if getattr(args, 'hardcoded_r', None) is not None: + enkf.obs_var_dict = {k: args.hardcoded_r for k in enkf.obs_var_dict} + print(f"[per-member] obs_var_dict overridden: R = {args.hardcoded_r} (fixed)") + if getattr(args, 'direct_variance', False): + obs_df = pd.read_csv(obs_file) + t_col = next(c for c in obs_df.columns if c.lower() in ('datetime', 'date', 'time', 'timestamp')) + v_col = next(c for c in obs_df.columns if 'var' in c.lower()) + obs_df[t_col] = pd.to_datetime(obs_df[t_col]) + obs_df = obs_df.set_index(t_col).sort_index() + enkf.obs_var_dict = { + str(ts): float(v) + for ts, v in obs_df[v_col].items() + if pd.notna(v) and float(v) > 0 + } + print(f"[per-member] obs_var_dict overridden: R = σ² direct ({len(enkf.obs_var_dict)} timesteps)") + + # Custom forcing loader for CFE BMI + def custom_load_forcing(self_cfe): + df = load_test_forcing(test_forcing_file) + self_cfe.forcing_data = df.rename(columns={"date": "time"}) + + # Build N CFE instances; perturb initial states (member 0 left clean) + models = [] + init_states_per_member = [] # capture initial-state snapshot after perturbation + for i in range(N): + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.load_forcing_file = custom_load_forcing.__get__(m) + m.initialize() + if i > 0: + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + r = enkf.init_state_perturb_frac + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1 + r * enkf.rng.standard_normal()), 1e-6, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1 + r * enkf.rng.standard_normal()), 1e-6, gw_max)) + jitter = max(sm_max * 1e-6, 1e-9) + m.nash_storage[0] = max(n00 + jitter * enkf.rng.standard_normal(), 0.0) + m.nash_storage[1] = max(n10 + jitter * enkf.rng.standard_normal(), 0.0) + # Snapshot the (possibly-perturbed) initial state for this member + init_states_per_member.append({ + "soil_m": float(m.soil_reservoir["storage_m"]), + "gw_m": float(m.gw_reservoir["storage_m"]), + "nash0_m": float(m.nash_storage[0]), + "nash1_m": float(m.nash_storage[1]), + "soil_max_m": float(m.soil_reservoir["storage_max_m"]), + "gw_max_m": float(m.gw_reservoir["storage_max_m"]), + }) + models.append(m) + + # SKIP SPINUP — go directly to the test period (matches sensitivity-script setup) + df = load_test_forcing(test_forcing_file) + t_mask = (df['date'] >= TEST_START) & (df['date'] <= TEST_END) + df_test = df[t_mask] + n_hours = len(df_test) + q_matrix = np.full((n_hours, N), np.nan, dtype=float) + precip_matrix = np.full((n_hours, N), np.nan, dtype=float) # mm/h per member + pet_matrix = np.full((n_hours, N), np.nan, dtype=float) # mm/h per member + dates_out = df_test['date'].values + + for h, (p, e, current_date) in enumerate(zip( + df_test['total_precipitation'], + df_test['potential_evaporation'], + df_test['date'])): + + # Perturb forcing per member + p_arr, e_arr = enkf.perturb_forcing(p, e) + # Capture the per-member perturbed forcing for later plotting/diagnostics + precip_matrix[h, :] = p_arr + pet_matrix[h, :] = e_arr + + # Advance each member one hour + ensemble_q = np.empty(N, dtype=float) + for i, m in enumerate(models): + m.set_value('atmosphere_water__time_integral_of_precipitation_mass_flux', + float(p_arr[i]) / 1000) + m.set_value('water_potential_evaporation_flux', + float(e_arr[i]) / 1000 / 3600) + m.update() + ensemble_q[i] = m.get_value('land_surface_water__runoff_depth') * 1000 # mm/h + + # Record pre-analysis forecast per member (this is the "actual production" value) + q_matrix[h, :] = ensemble_q + + # DA step (uses production EnKFAssimilator) + enkf.update_states(models, current_date, ensemble_q) + + # Process noise after DA + enkf.add_process_noise(models) + + for m in models: + m.finalize() + + # ----- Save per-member CSVs ----- + def _save_matrix(matrix, suffix): + out = {'date': dates_out} + for i in range(N): + out[f'member_{i:02d}'] = matrix[:, i] + df_out = pd.DataFrame(out) + out_path = out_dir / f'{cat_id}_production_per_member_{suffix}.csv' + df_out.to_csv(out_path, index=False) + return out_path + + # Streamflow outputs + q_path = out_dir / f'{cat_id}_production_per_member.csv' + out = {'date': dates_out} + for i in range(N): + out[f'member_{i:02d}'] = q_matrix[:, i] + pd.DataFrame(out).to_csv(q_path, index=False) + print(f"[per-member] saved {q_path}") + + # Perturbed forcings (so each member's actual inputs are recoverable) + precip_path = _save_matrix(precip_matrix, "precip") + pet_path = _save_matrix(pet_matrix, "pet") + print(f"[per-member] saved {precip_path}") + print(f"[per-member] saved {pet_path}") + + # Initial states (small, JSON is fine) + init_states_path = out_dir / f'{cat_id}_production_per_member_initial_states.json' + with open(init_states_path, 'w') as f: + json.dump({ + "catchment_id": cat_id, + "n_members": N, + "members": {f"member_{i:02d}": init_states_per_member[i] for i in range(N)}, + }, f, indent=2) + print(f"[per-member] saved {init_states_path}") + + print(f"[per-member] ensemble mean spread over test period: " + f"{float(q_matrix.std(axis=1).mean()):.5f} mm/h") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', required=True) + parser.add_argument('--forcing-dir', required=True) + parser.add_argument('--obs-dir', required=True) + parser.add_argument('--cfe-dir', required=True) + parser.add_argument('--config-file', required=True) + parser.add_argument('--param-bounds', required=True, + help='Kept for symmetry; not read in this script') + parser.add_argument('--out-dir', required=True) + parser.add_argument('--test-forcing-dir1', required=True) + parser.add_argument('--test-forcing-dir2', required=True) + parser.add_argument('--enkf-members', type=int, default=20) + parser.add_argument('--enkf-obs-error-std', type=float, default=0.05) + parser.add_argument('--no-vrugt-r', action='store_true') + parser.add_argument('--hardcoded-r', type=float, default=None, + help='Fix R to this constant value for all timesteps (overrides Vrugt formula)') + parser.add_argument('--direct-variance', action='store_true', default=False, + help='Use kriging variance column as R per timestep (R=σ²; overrides Vrugt formula)') + parser.add_argument('--vrugt-alpha', type=float, default=0.10) + parser.add_argument('--vrugt-scale', type=float, default=0.001) + parser.add_argument('--rng-seed', type=int, default=None) + parser.add_argument('--prod-script', default=None, + help='Path to calibrate_catchment_cfe_da_v2.py for importing ' + 'EnKFAssimilator. Defaults to the script located next to this file.') + args = parser.parse_args() + + # Locate the production script (default: same dir as this script) + if args.prod_script is None: + here = os.path.dirname(os.path.abspath(__file__)) + prod_script = os.path.join(here, 'calibrate_catchment_cfe_da_v2.py') + else: + prod_script = args.prod_script + if not os.path.exists(prod_script): + raise FileNotFoundError(f"Could not find production script at {prod_script}. " + f"Pass --prod-script to override.") + EnKFAssimilator = import_enkf_class(prod_script) + + # Dynamic CFE import (same idiom as production script) + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + + run(args, _bmi_cfe, EnKFAssimilator) + + +if __name__ == '__main__': + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/3_routing/route_det_f4.sh b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/3_routing/route_det_f4.sh new file mode 100644 index 00000000..b2ebedc0 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/3_routing/route_det_f4.sh @@ -0,0 +1,28 @@ +#!/usr/bin/env bash +# F4 dynamic variance direct (20pct gauge holdout) — deterministic T-route routing. +# Routes per-catchment _test_results.csv (sim_mm_h column) through Muskingum-Cunge +# to gauge 03463300 and writes routed_Q_test.csv + KGE/NSE summary. +# +# Usage: +# bash route_det_f4.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +ROUTE_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/2_troute_routing/run_route.py + +F4_DIR=/mnt/disk2/1400_sites_helene/da_arms_dynamic_novrugt_seeded +GPKG=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +echo "[F4] Deterministic T-route routing..." +echo " da-dir : $F4_DIR" +echo " out-dir: $F4_DIR" + +$TROUTE "$ROUTE_SCRIPT" \ + --gpkg "$GPKG" \ + --da-dir "$F4_DIR" \ + --out-dir "$F4_DIR" \ + --usgs-csv "$USGS_CSV" + +echo "[F4] Deterministic routing done. Output: $F4_DIR/routed_Q_test.csv" diff --git a/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/3_routing/route_ensemble_f4.sh b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/3_routing/route_ensemble_f4.sh new file mode 100644 index 00000000..fd42c4d8 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/3_routing/route_ensemble_f4.sh @@ -0,0 +1,30 @@ +#!/usr/bin/env bash +# F4 dynamic variance direct (20pct gauge holdout) — crossed-ensemble T-route routing. +# Routes all 600 crossed-ensemble members through Muskingum-Cunge +# to gauge 03463300 and writes routed_crossed_ensemble.parquet. +# +# Must be run AFTER crossed ensemble parquets exist in CROSSED_DIR. +# +# Usage: +# bash route_ensemble_f4.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +ROUTE_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/2_troute_routing/run_route_crossed_ensemble.py + +CROSSED_DIR=/mnt/disk2/1400_sites_helene/da_crossed_dynamic_novrugt_seeded +GPKG=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +echo "[F4] Crossed-ensemble T-route routing..." +echo " ensemble-dir: $CROSSED_DIR" +echo " out-dir : $CROSSED_DIR" + +$TROUTE "$ROUTE_SCRIPT" \ + --gpkg "$GPKG" \ + --ensemble-dir "$CROSSED_DIR" \ + --out-dir "$CROSSED_DIR" \ + --usgs-csv "$USGS_CSV" + +echo "[F4] Ensemble routing done. Output: $CROSSED_DIR/routed_crossed_ensemble.parquet" diff --git a/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/3_routing/route_leadtime_f4.sh b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/3_routing/route_leadtime_f4.sh new file mode 100644 index 00000000..b9f1b898 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/3_routing/route_leadtime_f4.sh @@ -0,0 +1,31 @@ +#!/usr/bin/env bash +# F4 dynamic variance direct (20pct gauge holdout) — route lead-time forecast CSVs. +# Routes all 21-catchment lead-time forecast CSVs through Muskingum-Cunge +# to gauge 03463300. Writes routed_leadtime_da_full.parquet and +# routed_leadtime_openloop_full.parquet for gauge-level error decay analysis. +# +# Run AFTER the lead-time sweep batch finishes. +# +# Usage: +# bash route_leadtime_f4.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +ROUTE_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/2_troute_routing/run_route_leadtime_forecasts.py + +F4_DIR=/mnt/disk2/1400_sites_helene/da_forecast_dynamic_novrugt_seeded +GPKG=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg +OUT_DIR=/mnt/disk2/suma_helen_poster/leadtime_troute_routing_f4 + +echo "[F4] Routing lead-time forecasts through T-route..." +echo " forecast-dir: $F4_DIR" +echo " out-dir : $OUT_DIR" + +$TROUTE "$ROUTE_SCRIPT" \ + --gpkg "$GPKG" \ + --forecast-dir "$F4_DIR" \ + --out-dir "$OUT_DIR" + +echo "[F4] Lead-time routing done. Outputs in: $OUT_DIR" +ls "$OUT_DIR"/*.parquet 2>/dev/null || echo " (no parquets found)" diff --git a/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/plot_forecast_error_fixed_target.py b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/plot_forecast_error_fixed_target.py new file mode 100644 index 00000000..1460ed21 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/plot_forecast_error_fixed_target.py @@ -0,0 +1,217 @@ +""" +plot_forecast_error_fixed_target.py — 4a: error decay, fixed-target-time view. + +For each target verification time T in the Helene peak window: + Collect all forecasts that verify AT T: + issue_time = T - lead_hour*1h, for lead in 1..18 + error[lead] = ensemble_mean(q at T, initialized T-lead) - USGS_obs(T) + +This gives the correct operational picture: + - lead 1 = initialized 1 hr before T (DA just ran -> small error) + - lead 18 = initialized 18 hr before T (DA long ago -> error ~ open loop) + +Two panels: + Top : signed error (m³/s) vs lead hour, one curve per target time + Bot : same for open-loop +Plus a summary panel: mean across all target times, DA vs OL. + +Outputs: + /error_fixed_target_helene.png (per-target spaghetti, DA vs OL) + /error_fixed_target_mean.png (mean across targets, DA vs OL) +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.cm as cm + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast_hardcoded_r_routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +# Target verification times: hourly through the Helene peak window +TARGET_START = pd.Timestamp("2024-09-26 18:00:00") +TARGET_END = pd.Timestamp("2024-09-28 06:00:00") + + +def load_parquet(path): + df = pd.read_parquet(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + df["ens_mean"] = df[member_cols].mean(axis=1) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + return df[["issue_time", "lead_hour", "valid_time", "ens_mean"]] + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() + or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def build_fixed_target_errors(df, obs_series, target_times): + """ + For each target time T and each lead L (1..18): + issue_time = T - L hours + error = ensemble_mean at (issue_time, lead=L) - obs(T) + Returns dict: target_time -> {lead: error} + """ + # Index df by (issue_time, lead_hour) for fast lookup + df_idx = df.set_index(["issue_time", "lead_hour"])["ens_mean"] + + results = {} + for T in target_times: + obs_val = obs_series.get(T, np.nan) + if np.isnan(obs_val): + continue + curve = {} + for lead in range(1, 19): + t0 = T - pd.Timedelta(hours=lead) + try: + q_fc = df_idx.loc[(t0, lead)] + curve[lead] = float(q_fc) - obs_val + except KeyError: + curve[lead] = np.nan + results[T] = curve + return results + + +def plot_spaghetti(ax, error_dict, color_da, label_prefix, linestyle="-", lw=0.9, alpha=0.35): + """Plot one thin line per target time + thick mean.""" + leads = list(range(1, 19)) + all_curves = [] + target_times = sorted(error_dict.keys()) + cmap = cm.get_cmap("plasma", len(target_times)) + + for i, T in enumerate(target_times): + curve = [error_dict[T].get(l, np.nan) for l in leads] + ax.plot(leads, curve, + color=cmap(i), lw=lw, alpha=alpha, + linestyle=linestyle, zorder=2) + all_curves.append(curve) + + if all_curves: + mean_curve = np.nanmean(np.array(all_curves), axis=0) + ax.plot(leads, mean_curve, + color=color_da, lw=2.8, alpha=0.95, + linestyle=linestyle, zorder=5, + label=f"{label_prefix} — mean across targets") + return all_curves + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--route-dir", default=DEFAULT_ROUTE_DIR) + parser.add_argument("--out-dir", default=None) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--da-name", default="routed_leadtime_da_full.parquet") + parser.add_argument("--ol-name", default="routed_leadtime_openloop_full.parquet") + parser.add_argument("--target-start", default=str(TARGET_START)) + parser.add_argument("--target-end", default=str(TARGET_END)) + args = parser.parse_args() + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + print("Loading parquets...") + da = load_parquet(os.path.join(args.route_dir, args.da_name)) + ol = load_parquet(os.path.join(args.route_dir, args.ol_name)) + obs = load_usgs(args.usgs_csv) + + target_times = pd.date_range(args.target_start, args.target_end, freq="1h") + print(f" Target verification times: {len(target_times)} " + f"({target_times[0]} → {target_times[-1]})") + + print("Building fixed-target error tables...") + da_errors = build_fixed_target_errors(da, obs, target_times) + ol_errors = build_fixed_target_errors(ol, obs, target_times) + print(f" Targets with obs: DA={len(da_errors)} OL={len(ol_errors)}") + + leads = list(range(1, 19)) + + # ---- Spaghetti: per-target-time curves, DA vs OL ---- + fig, (ax_da, ax_ol) = plt.subplots(2, 1, figsize=(13, 10), sharex=True, sharey=True) + + ax_da.axhline(0, color="black", lw=0.8, linestyle="--", alpha=0.4) + plot_spaghetti(ax_da, da_errors, "tab:purple", "DA", linestyle="-") + ax_da.set_ylabel("Error: forecast − USGS obs (m³/s)", fontsize=10) + ax_da.set_title("DA — error at each lead for fixed target times (Helene peak window)", fontsize=11) + ax_da.grid(True, alpha=0.2) + ax_da.legend(fontsize=9) + + ax_ol.axhline(0, color="black", lw=0.8, linestyle="--", alpha=0.4) + plot_spaghetti(ax_ol, ol_errors, "tab:gray", "Open-loop", linestyle="--") + ax_ol.set_ylabel("Error: forecast − USGS obs (m³/s)", fontsize=10) + ax_ol.set_xlabel("Forecast lead hour (hours before target)", fontsize=11) + ax_ol.set_title("Open-loop — error at each lead for fixed target times", fontsize=11) + ax_ol.set_xticks(leads) + ax_ol.grid(True, alpha=0.2) + ax_ol.legend(fontsize=9) + + fig.suptitle( + "Forecast error vs lead time — fixed verification time, Helene peak window\n" + f"Target times: {TARGET_START.strftime('%b %d %H UTC')} → " + f"{TARGET_END.strftime('%b %d %H UTC')} | USGS 03463300", + fontsize=12, y=1.01, + ) + plt.tight_layout() + out1 = os.path.join(out_dir, "error_fixed_target_helene.png") + plt.savefig(out1, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + # ---- Summary: mean across all target times, DA vs OL overlaid ---- + fig, ax = plt.subplots(figsize=(13, 6)) + ax.axhline(0, color="black", lw=0.8, linestyle="--", alpha=0.4) + + def mean_curve(error_dict): + arr = np.array([ + [error_dict[T].get(l, np.nan) for l in leads] + for T in sorted(error_dict.keys()) + ]) + return np.nanmean(arr, axis=0), np.nanstd(arr, axis=0) + + da_mean, da_std = mean_curve(da_errors) + ol_mean, ol_std = mean_curve(ol_errors) + + ax.fill_between(leads, da_mean - da_std, da_mean + da_std, + color="tab:purple", alpha=0.15, zorder=2) + ax.fill_between(leads, ol_mean - ol_std, ol_mean + ol_std, + color="tab:gray", alpha=0.15, zorder=2) + ax.plot(leads, da_mean, color="tab:purple", lw=2.6, marker="o", + zorder=5, label="DA — mean error (±1 std shaded)") + ax.plot(leads, ol_mean, color="tab:gray", lw=2.6, marker="s", + linestyle="--", zorder=5, label="Open-loop — mean error (±1 std shaded)") + + ax.set_xlabel("Forecast lead hour (hours before target verification time)", fontsize=11) + ax.set_ylabel("Mean error: forecast − USGS obs (m³/s)", fontsize=11) + ax.set_xticks(leads) + ax.set_title( + "Mean forecast error vs lead time — fixed verification time, Helene peak window\n" + f"Target: {TARGET_START.strftime('%b %d %H UTC')} → " + f"{TARGET_END.strftime('%b %d %H UTC')} | " + "Lead 1 = init 1 hr before target | Lead 18 = init 18 hr before target", + fontsize=11, + ) + ax.legend(fontsize=10, loc="lower right", frameon=True, framealpha=0.92) + ax.grid(True, alpha=0.25) + plt.tight_layout() + out2 = os.path.join(out_dir, "error_fixed_target_mean.png") + plt.savefig(out2, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/plot_forecast_error_per_init.py b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/plot_forecast_error_per_init.py new file mode 100644 index 00000000..691245d0 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/plot_forecast_error_per_init.py @@ -0,0 +1,203 @@ +""" +plot_forecast_error_per_init.py + +Error decay by initialization time for the Helene window (Sep 24-28 2024). + +For each initialization time t0 in the Helene window: + error[lead] = ensemble_mean(q_gauge_m3s at t0+lead) - USGS_obs(t0+lead) + +Plotted as: + DA : thin colored lines (one per init time, colored by date) + thick mean across all + OL : thin gray dashed lines + thick gray dashed mean + +x-axis: forecast lead hour (1 -> 18) +y-axis: signed error (m³/s), positive = forecast too high + +The expected signal: DA error is small at lead 1 (just assimilated), grows +and converges toward the OL error curve by lead 18. + +Inputs: + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + +Obs: + /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +Outputs: + /forecast_error_per_init_helene.png (signed error) + /forecast_rmse_per_init_helene.png (absolute error / RMSE per init) +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.patches as mpatches + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast_hardcoded_r_routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +# Initialization times to show — Helene window +HELENE_INIT_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_INIT_END = pd.Timestamp("2024-09-28 23:00:00") + +# One color per init date +DATE_COLORS = { + "2024-09-24": "#1f77b4", + "2024-09-25": "#ff7f0e", + "2024-09-26": "#2ca02c", + "2024-09-27": "#d62728", + "2024-09-28": "#9467bd", +} + + +def load_parquet(path): + df = pd.read_parquet(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + df["ens_mean"] = df[member_cols].mean(axis=1) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + return df[["issue_time", "lead_hour", "valid_time", "ens_mean"]] + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() + or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def build_error_table(df, obs_series, init_start, init_end): + """Return DataFrame: issue_time, lead_hour, ens_mean, obs, error.""" + df = df[(df["issue_time"] >= init_start) & (df["issue_time"] <= init_end)].copy() + df["obs"] = df["valid_time"].map(obs_series) + df["error"] = df["ens_mean"] - df["obs"] + return df.dropna(subset=["obs", "error"]) + + +def plot_error(ax, err_df, color, alpha_thin, lw_thin, linestyle, label_prefix): + """Plot individual init-time error curves + thick mean curve.""" + leads = sorted(err_df["lead_hour"].unique()) + all_curves = [] + + for t0, grp in err_df.groupby("issue_time"): + date_str = str(pd.Timestamp(t0).date()) + c = DATE_COLORS.get(date_str, color) + grp_sorted = grp.sort_values("lead_hour") + # Align to leads grid — some may be missing + curve = grp_sorted.set_index("lead_hour")["error"].reindex(leads).values + ax.plot(leads, curve, + color=c, lw=lw_thin, alpha=alpha_thin, + linestyle=linestyle, zorder=2) + all_curves.append(curve) + + if all_curves: + mean_curve = np.nanmean(np.array(all_curves), axis=0) + ax.plot(leads, mean_curve, + color=color, lw=2.8, alpha=0.95, + linestyle=linestyle, zorder=5, + label=f"{label_prefix} — mean across all init times") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--route-dir", default=DEFAULT_ROUTE_DIR) + parser.add_argument("--out-dir", default=None) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--da-name", default="routed_leadtime_da_full.parquet") + parser.add_argument("--ol-name", default="routed_leadtime_openloop_full.parquet") + args = parser.parse_args() + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + print("Loading parquets...") + da = load_parquet(os.path.join(args.route_dir, args.da_name)) + ol = load_parquet(os.path.join(args.route_dir, args.ol_name)) + obs = load_usgs(args.usgs_csv) + print(f" DA issue_times: {da['issue_time'].nunique()} " + f"OL issue_times: {ol['issue_time'].nunique()}") + + da_err = build_error_table(da, obs, HELENE_INIT_START, HELENE_INIT_END) + ol_err = build_error_table(ol, obs, HELENE_INIT_START, HELENE_INIT_END) + print(f" DA init times in Helene window: {da_err['issue_time'].nunique()}") + + leads = sorted(da_err["lead_hour"].unique()) + + # ---- Signed error plot ---- + fig, ax = plt.subplots(figsize=(12, 6)) + ax.axhline(0, color="black", lw=0.8, linestyle="--", alpha=0.5, zorder=1) + + plot_error(ax, ol_err, color="tab:gray", alpha_thin=0.12, lw_thin=0.7, + linestyle="--", label_prefix="Open-loop") + plot_error(ax, da_err, color="tab:purple", alpha_thin=0.18, lw_thin=0.8, + linestyle="-", label_prefix="DA") + + # Date-color legend patches + patches = [mpatches.Patch(color=c, label=f"Init {d}") + for d, c in DATE_COLORS.items()] + handles, labels = ax.get_legend_handles_labels() + ax.legend(handles=handles + patches, + fontsize=8, loc="upper left", frameon=True, framealpha=0.9, ncol=2) + + ax.set_xlabel("Forecast lead hour", fontsize=11) + ax.set_ylabel("Error: forecast − USGS obs (m³/s)", fontsize=11) + ax.set_xticks(np.arange(1, 19)) + ax.set_title( + "Forecast error vs lead time — per initialization time, Helene window\n" + "DA (purple solid) vs Open-loop (gray dashed) | Sep 24–28 2024 | USGS 03463300", + fontsize=11, + ) + ax.grid(True, alpha=0.25, lw=0.4) + plt.tight_layout() + out1 = os.path.join(out_dir, "forecast_error_per_init_helene.png") + plt.savefig(out1, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + # ---- Absolute error (|error|) averaged per lead — cleaner summary ---- + fig, ax = plt.subplots(figsize=(12, 6)) + + def mean_abs_error_by_lead(err_df): + return err_df.groupby("lead_hour")["error"].apply( + lambda x: float(np.nanmean(np.abs(x))) + ) + + da_mae = mean_abs_error_by_lead(da_err) + ol_mae = mean_abs_error_by_lead(ol_err) + + ax.plot(da_mae.index, da_mae.values, + color="tab:purple", lw=2.4, marker="o", label="DA — mean |error|") + ax.plot(ol_mae.index, ol_mae.values, + color="tab:gray", lw=2.4, marker="s", linestyle="--", + label="Open-loop — mean |error|") + + ax.set_xlabel("Forecast lead hour", fontsize=11) + ax.set_ylabel("Mean |error| vs USGS obs (m³/s)", fontsize=11) + ax.set_xticks(np.arange(1, 19)) + ax.set_title( + "Mean absolute forecast error vs lead time — Helene window\n" + "DA (purple) vs Open-loop (gray) | Sep 24–28 2024 | USGS 03463300", + fontsize=11, + ) + ax.legend(fontsize=10, loc="upper left", frameon=True, framealpha=0.9) + ax.grid(True, alpha=0.25, lw=0.4) + plt.tight_layout() + out2 = os.path.join(out_dir, "forecast_mae_per_lead_helene.png") + plt.savefig(out2, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/plot_lead_time_decay.py b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/plot_lead_time_decay.py new file mode 100644 index 00000000..2b33fa82 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/plot_lead_time_decay.py @@ -0,0 +1,160 @@ +""" +Catchment-level error-vs-lead-time decay curve + ensemble spread by lead. + +Two-panel figure for ONE catchment: + TOP — RMSE of ensemble-mean forecast at each lead hour (1..18) vs the + catchment's kriging observation. DA solid, open-loop dashed, with + a shaded band showing the min/max of per-member RMSE. + BOTTOM — Mean ensemble spread (std-dev across 20 members, averaged across + all issue times) at each lead hour. Tells you whether forcing + perturbation alone keeps the forecast ensemble diverse during the + 18-hour free-run — useful for inspecting individual members + without needing the full spaghetti view. + +This is the catchment-level analog of the gauge-level decay curve. +Routing to the gauge is a separate post-step (route_lead_time_forecasts.py ++ a gauge-level decay script); this script lets us look at the catchment-level +signal without T-route in the loop. + +Inputs (from run_lead_time_forecast_sweep.py): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv +Obs (from production EnKF run): + //_test_results.csv + +Output: + //_lead_time_decay.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" + +DEFAULT_LEADTIME_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast" +DEFAULT_DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" + +# Module-level globals set by main() so the helpers see them +CAT = "cat-1016300" +LEADTIME_DIR = DEFAULT_LEADTIME_DIR +DA_DIR = DEFAULT_DA_DIR +OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay.png") + + +def load_forecasts(path): + """Return (df, member_cols).""" + if not os.path.exists(path): + raise FileNotFoundError(path) + df = pd.read_csv(path, parse_dates=['issue_time', 'valid_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + return df, member_cols + + +def load_obs(): + obs_path = os.path.join(DA_DIR, CAT, f"{CAT}_test_results.csv") + df = pd.read_csv(obs_path, parse_dates=['date']) + return df.set_index('date')['obs_mm_h'] + + +def metrics_by_lead(df, member_cols, obs_series): + """For each lead hour, return (rmse_mean, rmse_min_member, rmse_max_member, + mean_ensemble_std).""" + df = df.copy() + df['obs'] = df['valid_time'].map(obs_series) + df = df.dropna(subset=['obs']) + df['ens_mean'] = df[member_cols].mean(axis=1) + df['ens_std'] = df[member_cols].std(axis=1) + + leads = sorted(df['lead_hour'].unique()) + rmse_mean = np.full(len(leads), np.nan) + rmse_min_member = np.full(len(leads), np.nan) + rmse_max_member = np.full(len(leads), np.nan) + mean_ens_std = np.full(len(leads), np.nan) + for i, L in enumerate(leads): + sub = df[df['lead_hour'] == L] + if len(sub) == 0: + continue + err = sub['ens_mean'].values - sub['obs'].values + rmse_mean[i] = float(np.sqrt(np.mean(err ** 2))) + per_member_rmse = np.array([ + np.sqrt(np.mean((sub[c].values - sub['obs'].values) ** 2)) + for c in member_cols + ]) + rmse_min_member[i] = float(np.nanmin(per_member_rmse)) + rmse_max_member[i] = float(np.nanmax(per_member_rmse)) + mean_ens_std[i] = float(np.nanmean(sub['ens_std'].values)) + return (np.asarray(leads), rmse_mean, rmse_min_member, + rmse_max_member, mean_ens_std) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', default="cat-1016300") + parser.add_argument('--leadtime-dir', default=DEFAULT_LEADTIME_DIR, + help='Dir holding /_lead_time_forecasts_{da,openloop}.csv') + parser.add_argument('--da-dir', default=DEFAULT_DA_DIR, + help='Dir holding /_test_results.csv (for obs)') + args = parser.parse_args() + + global CAT, LEADTIME_DIR, DA_DIR, OUT_PNG + CAT = args.cat_id + LEADTIME_DIR = args.leadtime_dir + DA_DIR = args.da_dir + OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay.png") + + da_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_da.csv") + ol_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_openloop.csv") + df_da, m_cols = load_forecasts(da_path) + df_ol, _ = load_forecasts(ol_path) + obs_series = load_obs() + + leads_da, rmse_da, rmse_da_lo, rmse_da_hi, std_da = metrics_by_lead( + df_da, m_cols, obs_series) + leads_ol, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = metrics_by_lead( + df_ol, m_cols, obs_series) + + fig, (ax_top, ax_bot) = plt.subplots(2, 1, figsize=(10, 10), sharex=True, + gridspec_kw={"height_ratios": [1.4, 1.0]}) + + # ----- TOP: RMSE decay curve ----- + ax_top.fill_between(leads_da, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2, + label="DA — member spread (min/max RMSE)") + ax_top.fill_between(leads_ol, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2, + label="Open-loop — member spread") + ax_top.plot(leads_da, rmse_da, color=DA_COLOR, lw=2.4, marker='o', + zorder=4, label="DA on (ensemble-mean RMSE)") + ax_top.plot(leads_ol, rmse_ol, color=OL_COLOR, lw=2.4, marker='s', + linestyle='--', zorder=4, label="Open-loop (ensemble-mean RMSE)") + ax_top.set_ylabel("RMSE vs Qkrig obs (mm/h)", fontsize=11) + ax_top.set_title( + f"Forecast lead-time error decay — {CAT}\n" + f"Issue times pooled across test period (Oct 2023 – Oct 2024)", + fontsize=12, + ) + ax_top.grid(True, alpha=0.25) + ax_top.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + # ----- BOTTOM: ensemble spread by lead hour ----- + ax_bot.plot(leads_da, std_da, color=DA_COLOR, lw=2.2, marker='o', + label="DA — mean ensemble std") + ax_bot.plot(leads_ol, std_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', label="Open-loop — mean ensemble std") + ax_bot.set_xlabel("Forecast lead time (hours)", fontsize=11) + ax_bot.set_ylabel("Mean ensemble std (mm/h)", fontsize=11) + ax_bot.set_xticks(np.arange(1, 19)) + ax_bot.grid(True, alpha=0.25) + ax_bot.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + plt.tight_layout() + plt.savefig(OUT_PNG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/plot_lead_time_decay_by_regime.py b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/plot_lead_time_decay_by_regime.py new file mode 100644 index 00000000..ecf10408 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/plot_lead_time_decay_by_regime.py @@ -0,0 +1,200 @@ +""" +Lead-time decay curve, split by flow regime at issue time. + +The pooled lead-time curve (plot_lead_time_decay.py) showed DA losing to +open-loop across most lead hours, but the test period is ~99% low-flow. +This script splits the same forecast CSVs by the flow regime at the issue +time t0, so we can see whether DA helps when it matters (storms / Helene) +and hurts when it doesn't (low flow). + +Three regimes are partitioned on obs(t0) (kriging obs at the issue time): + + helene issue_time ∈ [2024-09-24, 2024-09-28] (the 5-day Helene window) + storm obs(t0) > 0.5 mm/h (high-flow issue times) + low_flow obs(t0) < 0.1 mm/h (low-flow issue times) + +Same 2-panel layout per regime: RMSE vs lead (top), ensemble spread vs lead +(bottom). Three regime columns × 2 metric rows in one figure. + +Inputs (from run_lead_time_forecast_sweep.py — no re-run needed): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv +Obs (from production EnKF): + //_test_results.csv + +Output: + //_lead_time_decay_by_regime.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" + +DEFAULT_LEADTIME_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast" +DEFAULT_DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" + +# Module-level globals set by main() so helpers see them +CAT = "cat-1016300" +LEADTIME_DIR = DEFAULT_LEADTIME_DIR +DA_DIR = DEFAULT_DA_DIR +OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay_by_regime.png") + +# Regime definitions +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-28 23:00:00") +STORM_OBS_THRESHOLD = 0.5 # mm/h +LOWFLOW_OBS_THRESHOLD = 0.1 # mm/h + + +def load_forecasts(path): + if not os.path.exists(path): + raise FileNotFoundError(path) + df = pd.read_csv(path, parse_dates=['issue_time', 'valid_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + return df, member_cols + + +def load_obs(): + obs_path = os.path.join(DA_DIR, CAT, f"{CAT}_test_results.csv") + df = pd.read_csv(obs_path, parse_dates=['date']) + return df.set_index('date')['obs_mm_h'] + + +def regime_mask(issue_times, obs_at_issue, regime): + """Boolean mask over issue_times for the named regime.""" + if regime == "helene": + return (issue_times >= HELENE_START) & (issue_times <= HELENE_END) + if regime == "storm": + return obs_at_issue > STORM_OBS_THRESHOLD + if regime == "low_flow": + return obs_at_issue < LOWFLOW_OBS_THRESHOLD + raise ValueError(regime) + + +def metrics_by_lead(df, member_cols, obs_series, issue_mask): + """Compute (leads, rmse_mean, rmse_min, rmse_max, mean_std) restricted + to issue times where issue_mask is True.""" + df = df.copy() + df = df[df['issue_time'].isin(issue_mask)] + if len(df) == 0: + return None + df['obs'] = df['valid_time'].map(obs_series) + df = df.dropna(subset=['obs']) + if len(df) == 0: + return None + df['ens_mean'] = df[member_cols].mean(axis=1) + df['ens_std'] = df[member_cols].std(axis=1) + + leads = sorted(df['lead_hour'].unique()) + rmse_mean = np.full(len(leads), np.nan) + rmse_lo = np.full(len(leads), np.nan) + rmse_hi = np.full(len(leads), np.nan) + std_mean = np.full(len(leads), np.nan) + for i, L in enumerate(leads): + sub = df[df['lead_hour'] == L] + if len(sub) == 0: + continue + err = sub['ens_mean'].values - sub['obs'].values + rmse_mean[i] = float(np.sqrt(np.mean(err ** 2))) + per_member = np.array([ + np.sqrt(np.mean((sub[c].values - sub['obs'].values) ** 2)) + for c in member_cols + ]) + rmse_lo[i] = float(np.nanmin(per_member)) + rmse_hi[i] = float(np.nanmax(per_member)) + std_mean[i] = float(np.nanmean(sub['ens_std'].values)) + return np.asarray(leads), rmse_mean, rmse_lo, rmse_hi, std_mean + + +def plot_pair(ax_rmse, ax_std, da_metrics, ol_metrics, + regime_label, n_issue_times): + if da_metrics is None or ol_metrics is None: + ax_rmse.set_title(f"{regime_label}\n(no issue times match)", fontsize=11) + ax_rmse.axis('off'); ax_std.axis('off') + return + leads, rmse_da, rmse_da_lo, rmse_da_hi, std_da = da_metrics + _, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = ol_metrics + + # RMSE panel + ax_rmse.fill_between(leads, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2) + ax_rmse.fill_between(leads, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2) + ax_rmse.plot(leads, rmse_da, color=DA_COLOR, lw=2.2, marker='o', + zorder=4, label="DA on") + ax_rmse.plot(leads, rmse_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', zorder=4, label="Open-loop") + ax_rmse.set_title(f"{regime_label}\n({n_issue_times} issue times)", fontsize=11) + ax_rmse.set_ylabel("RMSE (mm/h)", fontsize=10) + ax_rmse.grid(True, alpha=0.25) + ax_rmse.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + # Spread panel + ax_std.plot(leads, std_da, color=DA_COLOR, lw=2.0, marker='o', label="DA") + ax_std.plot(leads, std_ol, color=OL_COLOR, lw=2.0, marker='s', + linestyle='--', label="Open-loop") + ax_std.set_xlabel("Forecast lead time (hours)", fontsize=10) + ax_std.set_ylabel("Mean ensemble std (mm/h)", fontsize=10) + ax_std.set_xticks(np.arange(1, 19)) + ax_std.grid(True, alpha=0.25) + ax_std.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', default="cat-1016300") + parser.add_argument('--leadtime-dir', default=DEFAULT_LEADTIME_DIR, + help='Dir holding /_lead_time_forecasts_{da,openloop}.csv') + parser.add_argument('--da-dir', default=DEFAULT_DA_DIR, + help='Dir holding /_test_results.csv (for obs)') + args = parser.parse_args() + + global CAT, LEADTIME_DIR, DA_DIR, OUT_PNG + CAT = args.cat_id + LEADTIME_DIR = args.leadtime_dir + DA_DIR = args.da_dir + OUT_PNG = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_decay_by_regime.png") + + da_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_da.csv") + ol_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_openloop.csv") + df_da, m_cols = load_forecasts(da_path) + df_ol, _ = load_forecasts(ol_path) + obs_series = load_obs() + + # Unique issue times across the run, with obs at each + issue_times = pd.to_datetime(sorted(df_da['issue_time'].unique())) + obs_at_issue = pd.Series(issue_times, index=issue_times).map(obs_series) + + regimes = [ + ("helene", "Helene window (Sep 24–28 2024)"), + ("storm", f"Storm hours (obs(t0) > {STORM_OBS_THRESHOLD} mm/h)"), + ("low_flow", f"Low flow (obs(t0) < {LOWFLOW_OBS_THRESHOLD} mm/h)"), + ] + + fig, axes = plt.subplots(2, 3, figsize=(18, 10), + gridspec_kw={"height_ratios": [1.4, 1.0]}) + for col, (key, label) in enumerate(regimes): + mask = regime_mask(issue_times, obs_at_issue, key) + kept = issue_times[mask] + kept_set = set(pd.to_datetime(kept)) + da_m = metrics_by_lead(df_da, m_cols, obs_series, kept_set) + ol_m = metrics_by_lead(df_ol, m_cols, obs_series, kept_set) + plot_pair(axes[0, col], axes[1, col], da_m, ol_m, label, len(kept)) + + fig.suptitle( + f"Lead-time error decay by flow regime — {CAT}\n" + f"Same forecast CSVs as the pooled view, partitioned on obs(t0).", + fontsize=13, y=0.995, + ) + plt.tight_layout(rect=(0, 0, 1, 0.97)) + plt.savefig(OUT_PNG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/plot_lead_time_decay_gauge.py b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/plot_lead_time_decay_gauge.py new file mode 100644 index 00000000..71ce54fc --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/plot_lead_time_decay_gauge.py @@ -0,0 +1,311 @@ +""" +Gauge-level lead-time forecast decay curve. + +Routes 21-catchment forecast CSVs through T-route to USGS gauge 03463300 and +plots RMSE vs lead time (DA solid, open-loop dashed) against USGS observed Q. +Same 2-panel layout as the catchment-level plot, plus a regime split for the +Helene window. + +The routing step is done separately by route_lead_time_forecasts.py; +this script just reads the resulting parquets and produces the plot. + +Inputs: + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + Each: one row per (issue_time, lead_hour); member columns hold q_gauge_m3s. + (Long-format with explicit `member` and `q_gauge_m3s` columns also supported.) +Obs: + /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + USGS hourly Q in m³/s at gauge 03463300 (South Toe River near Celo, NC). + +Output: + /lead_time_decay_gauge_pooled.png + /lead_time_decay_gauge_by_regime.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" +OBS_COLOR = "black" + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/leadtime_troute_routing" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +# Watershed area at gauge 03463300 (sum of 21 catchments per GPKG hydrofabric) +WATERSHED_AREA_KM2 = 113.18 +# 1 mm/h depth × 113.18 km² = 113.18e3 m³/h = 31.439 m³/s +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-28 23:00:00") +STORM_OBS_THRESHOLD_M3S = 20.0 # gauge-level storm threshold (was 50 — too high + # for this 113 km² basin; only 2 issue times qualified) +LOWFLOW_OBS_THRESHOLD_M3S = 5.0 # gauge-level low-flow threshold + +USGS_HELENE_PEAK_M3S = 1886.0 # reference Sep 27 14:00 + + +def load_parquet_long(path): + """Load a routed parquet and normalize to long format: + columns = issue_time, lead_hour, member, q_gauge_m3s. + + Handles both wide format (member_00..member_19 columns) and long format + (explicit `member` + `q_gauge_m3s` columns). Auto-detects. + """ + if not os.path.exists(path): + raise FileNotFoundError(path) + df = pd.read_parquet(path) + + # Long format detection + if 'q_gauge_m3s' in df.columns and 'member' in df.columns: + df['issue_time'] = pd.to_datetime(df['issue_time']) + return df[['issue_time', 'lead_hour', 'member', 'q_gauge_m3s']] + + # Some variants might call it 'q_m3s' or similar + q_col_candidates = [c for c in df.columns if c.lower() in + ('q_gauge_m3s', 'q_m3s', 'q_gauge', 'q')] + if 'member' in df.columns and q_col_candidates: + qc = q_col_candidates[0] + df['issue_time'] = pd.to_datetime(df['issue_time']) + out = df[['issue_time', 'lead_hour', 'member', qc]].copy() + out = out.rename(columns={qc: 'q_gauge_m3s'}) + return out + + # Wide format: member_00..member_19 columns + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + if not member_cols: + raise ValueError(f"Couldn't identify member columns in {path}. " + f"Columns present: {df.columns.tolist()}") + df['issue_time'] = pd.to_datetime(df['issue_time']) + keep = ['issue_time', 'lead_hour'] + if 'valid_time' in df.columns: + keep.append('valid_time') + long = df[keep + member_cols].melt( + id_vars=keep, + value_vars=member_cols, + var_name='member', + value_name='q_gauge_m3s', + ) + return long + + +def load_usgs_obs(usgs_csv): + """Load USGS gauge obs and ensure m³/s. + + The CSV at gauge 03463300 has the column `QObs(mm/h)` (catchment-averaged + depth, not gauge discharge in m³/s). If the column name contains 'mm', we + convert mm/h → m³/s by multiplying by the watershed area factor: + m³/s = mm/h × WATERSHED_AREA_KM2 × 1000 / 3600 + Verified: 59.978 mm/h × 113.18 km² = 1886 m³/s (USGS Helene peak). + """ + df = pd.read_csv(usgs_csv) + date_candidates = [c for c in df.columns if c.lower() in + ('datetime', 'date', 'time', 'timestamp')] + q_candidates = [c for c in df.columns if 'q' in c.lower() + or 'flow' in c.lower() or 'discharge' in c.lower()] + if not date_candidates or not q_candidates: + raise ValueError(f"Could not identify date/Q columns in {usgs_csv}. " + f"Columns: {df.columns.tolist()}") + dc, qc = date_candidates[0], q_candidates[0] + df[dc] = pd.to_datetime(df[dc]) + series = df.set_index(dc)[qc].astype(float) + if 'mm' in qc.lower(): + series = series * MM_H_TO_M3_S + print(f" Converted obs from mm/h (column '{qc}') to m³/s " + f"using area = {WATERSHED_AREA_KM2} km² (× {MM_H_TO_M3_S:.4f})") + else: + print(f" Loaded obs as-is from column '{qc}' (assumed m³/s)") + return series + + +def metrics_by_lead(df_long, obs_series, issue_mask=None): + """Compute per-lead-hour: ensemble-mean RMSE, per-member min/max RMSE, + mean ensemble std. df_long must have issue_time, lead_hour, member, q_gauge_m3s. + """ + df = df_long.copy() + if issue_mask is not None: + df = df[df['issue_time'].isin(issue_mask)] + if len(df) == 0: + return None + df['valid_time'] = df['issue_time'] + pd.to_timedelta(df['lead_hour'], unit='h') + df['obs'] = df['valid_time'].map(obs_series) + df = df.dropna(subset=['obs']) + if len(df) == 0: + return None + + # Ensemble mean per (issue_time, lead_hour) + grouped = df.groupby(['issue_time', 'lead_hour']) + ens_mean = grouped['q_gauge_m3s'].mean().reset_index(name='ens_mean') + ens_std = grouped['q_gauge_m3s'].std().reset_index(name='ens_std') + obs_per = grouped['obs'].first().reset_index(name='obs') + panel = ens_mean.merge(ens_std, on=['issue_time', 'lead_hour']).merge( + obs_per, on=['issue_time', 'lead_hour']) + + leads = sorted(panel['lead_hour'].unique()) + rmse_mean = np.full(len(leads), np.nan) + rmse_lo = np.full(len(leads), np.nan) + rmse_hi = np.full(len(leads), np.nan) + std_mean = np.full(len(leads), np.nan) + for i, L in enumerate(leads): + psub = panel[panel['lead_hour'] == L] + if len(psub) == 0: + continue + err = psub['ens_mean'].values - psub['obs'].values + rmse_mean[i] = float(np.sqrt(np.mean(err ** 2))) + # Per-member RMSE — recompute from the underlying long frame + dsub = df[df['lead_hour'] == L] + member_rmses = [] + for m, g in dsub.groupby('member'): + err_m = g['q_gauge_m3s'].values - g['obs'].values + if len(err_m) >= 1: + member_rmses.append(float(np.sqrt(np.mean(err_m ** 2)))) + if member_rmses: + rmse_lo[i] = float(np.nanmin(member_rmses)) + rmse_hi[i] = float(np.nanmax(member_rmses)) + std_mean[i] = float(np.nanmean(psub['ens_std'].values)) + return np.asarray(leads), rmse_mean, rmse_lo, rmse_hi, std_mean + + +def plot_two_panel(out_path, leads_da, da_metrics, leads_ol, ol_metrics, title): + fig, (ax_top, ax_bot) = plt.subplots(2, 1, figsize=(10, 10), sharex=True, + gridspec_kw={"height_ratios": [1.4, 1.0]}) + leads, rmse_da, rmse_da_lo, rmse_da_hi, std_da = da_metrics + _, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = ol_metrics + + ax_top.fill_between(leads, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2, + label="DA — member spread (min/max RMSE)") + ax_top.fill_between(leads, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2, + label="Open-loop — member spread") + ax_top.plot(leads, rmse_da, color=DA_COLOR, lw=2.4, marker='o', + zorder=4, label="DA on (ensemble-mean RMSE)") + ax_top.plot(leads, rmse_ol, color=OL_COLOR, lw=2.4, marker='s', + linestyle='--', zorder=4, label="Open-loop (ensemble-mean RMSE)") + ax_top.set_ylabel("RMSE vs USGS obs (m³/s)", fontsize=11) + ax_top.set_title(title, fontsize=12) + ax_top.grid(True, alpha=0.25) + ax_top.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + ax_bot.plot(leads, std_da, color=DA_COLOR, lw=2.2, marker='o', + label="DA — mean ensemble std") + ax_bot.plot(leads, std_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', label="Open-loop — mean ensemble std") + ax_bot.set_xlabel("Forecast lead time (hours)", fontsize=11) + ax_bot.set_ylabel("Mean ensemble std (m³/s)", fontsize=11) + ax_bot.set_xticks(np.arange(1, 19)) + ax_bot.grid(True, alpha=0.25) + ax_bot.legend(fontsize=9, loc='upper left', frameon=True, framealpha=0.92) + + plt.tight_layout() + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + +def plot_three_panel_regime(out_path, regime_metrics): + fig, axes = plt.subplots(2, 3, figsize=(18, 10), + gridspec_kw={"height_ratios": [1.4, 1.0]}) + for col, (label, n_issue, da_m, ol_m) in enumerate(regime_metrics): + ax_rmse = axes[0, col] + ax_std = axes[1, col] + if da_m is None or ol_m is None: + ax_rmse.set_title(f"{label}\n(no issue times match)", fontsize=11) + ax_rmse.axis('off'); ax_std.axis('off') + continue + leads, rmse_da, rmse_da_lo, rmse_da_hi, std_da = da_m + _, rmse_ol, rmse_ol_lo, rmse_ol_hi, std_ol = ol_m + ax_rmse.fill_between(leads, rmse_da_lo, rmse_da_hi, + color=DA_COLOR, alpha=0.18, zorder=2) + ax_rmse.fill_between(leads, rmse_ol_lo, rmse_ol_hi, + color=OL_COLOR, alpha=0.18, zorder=2) + ax_rmse.plot(leads, rmse_da, color=DA_COLOR, lw=2.2, marker='o', + zorder=4, label="DA on") + ax_rmse.plot(leads, rmse_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', zorder=4, label="Open-loop") + ax_rmse.set_title(f"{label}\n({n_issue} issue times)", fontsize=11) + ax_rmse.set_ylabel("RMSE (m³/s)", fontsize=10) + ax_rmse.grid(True, alpha=0.25) + ax_rmse.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + ax_std.plot(leads, std_da, color=DA_COLOR, lw=2.0, marker='o', label="DA") + ax_std.plot(leads, std_ol, color=OL_COLOR, lw=2.0, marker='s', + linestyle='--', label="Open-loop") + ax_std.set_xlabel("Forecast lead time (hours)", fontsize=10) + ax_std.set_ylabel("Mean ensemble std (m³/s)", fontsize=10) + ax_std.set_xticks(np.arange(1, 19)) + ax_std.grid(True, alpha=0.25) + ax_std.legend(fontsize=8, loc='upper left', frameon=True, framealpha=0.92) + + fig.suptitle("Gauge-level lead-time error decay by regime — USGS 03463300", + fontsize=13, y=0.995) + plt.tight_layout(rect=(0, 0, 1, 0.97)) + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--route-dir', default=DEFAULT_ROUTE_DIR, + help='Dir holding routed_leadtime_{da,openloop}_full.parquet') + parser.add_argument('--out-dir', default=None, + help='Where to write the output PNGs. Defaults to --route-dir ' + '(which may not be writable if owned by another user — ' + 'pass an explicit path then).') + parser.add_argument('--usgs-csv', default=DEFAULT_USGS_CSV) + parser.add_argument('--da-name', default='routed_leadtime_da_full.parquet') + parser.add_argument('--ol-name', default='routed_leadtime_openloop_full.parquet') + args = parser.parse_args() + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + da_path = os.path.join(args.route_dir, args.da_name) + ol_path = os.path.join(args.route_dir, args.ol_name) + df_da = load_parquet_long(da_path) + df_ol = load_parquet_long(ol_path) + obs_series = load_usgs_obs(args.usgs_csv) + + print(f"DA rows: {len(df_da):,} | unique issue_times: {df_da['issue_time'].nunique()}") + print(f"OL rows: {len(df_ol):,} | unique issue_times: {df_ol['issue_time'].nunique()}") + print(f"USGS obs range: {obs_series.index.min()} .. {obs_series.index.max()} " + f"({len(obs_series):,} hours)") + print(f"USGS Helene peak (reference): {USGS_HELENE_PEAK_M3S:.0f} m³/s") + + # ----- Pooled (all issue times) ----- + da_pooled = metrics_by_lead(df_da, obs_series) + ol_pooled = metrics_by_lead(df_ol, obs_series) + out_pooled = os.path.join(out_dir, "lead_time_decay_gauge_pooled.png") + plot_two_panel(out_pooled, da_pooled[0], da_pooled, ol_pooled[0], ol_pooled, + title="Gauge-level lead-time decay — USGS 03463300\n" + "All issue times pooled (Oct 2023 – Oct 2024)") + + # ----- Regime split: Helene, storm, low-flow at gauge ----- + issue_times = pd.to_datetime(sorted(df_da['issue_time'].unique())) + obs_at_issue = pd.Series(issue_times, index=issue_times).map(obs_series) + + helene_mask = (issue_times >= HELENE_START) & (issue_times <= HELENE_END) + storm_mask = obs_at_issue > STORM_OBS_THRESHOLD_M3S + lowflow_mask = obs_at_issue < LOWFLOW_OBS_THRESHOLD_M3S + + regime_metrics = [] + for label, mask in [ + (f"Helene window (Sep 24–28 2024)", helene_mask), + (f"Storm hours (obs(t0) > {STORM_OBS_THRESHOLD_M3S:.0f} m³/s)", storm_mask), + (f"Low flow (obs(t0) < {LOWFLOW_OBS_THRESHOLD_M3S:.0f} m³/s)", lowflow_mask), + ]: + kept = set(pd.to_datetime(issue_times[mask])) + da_m = metrics_by_lead(df_da, obs_series, kept) + ol_m = metrics_by_lead(df_ol, obs_series, kept) + regime_metrics.append((label, len(kept), da_m, ol_m)) + + out_regime = os.path.join(out_dir, "lead_time_decay_gauge_by_regime.png") + plot_three_panel_regime(out_regime, regime_metrics) + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/run_4a_f4.sh b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/run_4a_f4.sh new file mode 100644 index 00000000..c013172d --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4a_error_decay/run_4a_f4.sh @@ -0,0 +1,80 @@ +#!/usr/bin/env bash +# F4 dynamic variance direct (20pct gauge holdout) — 4a lead-time decay plots. +# +# Runs plot scripts per catchment: +# plot_lead_time_decay.py — pooled RMSE vs lead (catchment-level) +# plot_lead_time_decay_by_regime.py — same split by flow regime +# +# Gauge-level decay plots run once if routed parquets exist: +# plot_lead_time_decay_gauge.py — gauge-level (requires routed parquets) +# plot_forecast_error_fixed_target.py +# plot_forecast_error_per_init.py +# +# Usage: +# bash run_4a_f4.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +LEADTIME_DIR=/mnt/disk2/1400_sites_helene/da_forecast_dynamic_novrugt_seeded +DA_DIR=/mnt/disk2/1400_sites_helene/da_arms_dynamic_novrugt_seeded +OBS_DIR=/mnt/disk2/1400_sites_helene/catchment_ts_03463300_spliced_dyn_helene +ROUTE_DIR=/mnt/disk2/suma_helen_poster/leadtime_troute_routing_f4 +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +echo "[F4-4a] Lead-time decay plots — leadtime dir: $LEADTIME_DIR" + +for CAT in "${CATS[@]}"; do + DA_CSV="$LEADTIME_DIR/$CAT/${CAT}_lead_time_forecasts_da.csv" + if [ ! -f "$DA_CSV" ]; then + echo " [$CAT] No lead-time CSV — skipping" + continue + fi + + echo " [$CAT] Plotting pooled decay..." + $TROUTE "$SCRIPT_DIR/plot_lead_time_decay.py" \ + --cat-id "$CAT" \ + --leadtime-dir "$LEADTIME_DIR" \ + --da-dir "$DA_DIR" \ + --obs-dir "$OBS_DIR" + + echo " [$CAT] Plotting decay by regime..." + $TROUTE "$SCRIPT_DIR/plot_lead_time_decay_by_regime.py" \ + --cat-id "$CAT" \ + --leadtime-dir "$LEADTIME_DIR" \ + --da-dir "$DA_DIR" \ + --obs-dir "$OBS_DIR" +done + +# Gauge-level decay (needs routed parquets — run after route_leadtime_f4.sh) +if [ -f "$ROUTE_DIR/routed_leadtime_da_full.parquet" ]; then + echo "[F4-4a] Gauge-level decay (routed)..." + $TROUTE "$SCRIPT_DIR/plot_lead_time_decay_gauge.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" + + echo "[F4-4a] Fixed-target error plots..." + $TROUTE "$SCRIPT_DIR/plot_forecast_error_fixed_target.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" + + echo "[F4-4a] Per-init error plots..." + $TROUTE "$SCRIPT_DIR/plot_forecast_error_per_init.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" +else + echo "[F4-4a] Routed parquets not found — skipping gauge-level and error plots." + echo " Run route_leadtime_f4.sh first, then re-run this script." +fi + +echo "[F4-4a] Done." diff --git a/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_crossed_ensemble.py b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_crossed_ensemble.py new file mode 100644 index 00000000..7cb00a93 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_crossed_ensemble.py @@ -0,0 +1,229 @@ +""" +plot_crossed_ensemble.py -- 4b: 600-member crossed ensemble vs USGS obs. + +Loads the crossed ensemble parquet produced by run_crossed_ensemble.py: + //_crossed_ensemble.parquet + Columns: issue_time, lead_hour, member_0000..member_0599 + +Projects all members to valid_time = issue_time + lead_hour hours. +Aggregates at each valid_time across ALL members from ALL issue times +to build a probabilistic envelope (5th/25th/50th/75th/95th percentiles). + +Two output figures: + 4b_crossed_ensemble_helene.png + -- Full Helene window (Sep 24-30) probabilistic envelope + USGS obs. + Percentile shading in two layers (5-95 outer, 25-75 inner). + 4b_crossed_ensemble_peak.png + -- Zoomed to Helene peak (Sep 26 12 UTC -> Sep 28 06 UTC), + same shading + USGS obs + individual spaghetti from Sep 27 inits. + +Usage: + python3 plot_crossed_ensemble.py \\ + --ensemble-dir /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble \\ + --cat-id cat-1016300 \\ + --usgs-csv /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv \\ + --out-dir /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble/cat-1016300 +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import matplotlib.patches as mpatches + +DEFAULT_ENSEMBLE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" +DEFAULT_CAT_ID = "cat-1016300" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +PLOT_START = pd.Timestamp("2024-09-24 00:00:00") +PLOT_END = pd.Timestamp("2024-09-28 18:00:00") # cut trailing spike +PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +PEAK_END = pd.Timestamp("2024-09-28 06:00:00") +HELENE_PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +HELENE_PEAK_END = pd.Timestamp("2024-09-28 00:00:00") + + +def load_ensemble(path): + df = pd.read_parquet(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + return df, member_cols + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def build_percentile_envelope(df, member_cols, t_start, t_end): + """ + For each valid_time in [t_start, t_end], stack all member values + from all issue times and compute percentiles. + Returns DataFrame indexed by valid_time with columns p05, p25, p50, p75, p95. + """ + mask = (df["valid_time"] >= t_start) & (df["valid_time"] <= t_end) + sub = df[mask].copy() + sub["valid_time"] = sub["valid_time"].dt.floor("1h") + + records = [] + for vt, grp in sub.groupby("valid_time"): + vals = grp[member_cols].to_numpy(dtype=float).flatten() + vals = vals[~np.isnan(vals)] * MM_H_TO_M3_S + if len(vals) == 0: + continue + records.append({ + "valid_time": vt, + "p05": np.percentile(vals, 5), + "p25": np.percentile(vals, 25), + "p50": np.percentile(vals, 50), + "p75": np.percentile(vals, 75), + "p95": np.percentile(vals, 95), + "n": len(vals), + }) + return pd.DataFrame(records).set_index("valid_time").sort_index() + + +def _add_helene_band(ax, label=True): + ax.axvspan(HELENE_PEAK_START, HELENE_PEAK_END, + color="salmon", alpha=0.12, zorder=0) + if label: + ax.text(HELENE_PEAK_START + pd.Timedelta(hours=6), 0.97, "Helene peak", + transform=ax.get_xaxis_transform(), + fontsize=9, color="firebrick", ha="left", va="top") + + +def _format_xaxis(ax, major_interval=1, minor_hours=6): + ax.xaxis.set_major_locator(mdates.DayLocator(interval=major_interval)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.xaxis.set_minor_locator(mdates.HourLocator(interval=minor_hours)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.2, lw=0.4) + + +def plot_envelope(ax, env, obs_series, t_start, t_end, + ens_color="tab:blue", obs_color="black", helene_label=True): + _add_helene_band(ax, label=helene_label) + + # Outer band: 5-95 + ax.fill_between(env.index, env["p05"], env["p95"], + color=ens_color, alpha=0.15, zorder=2, + label="5th-95th percentile") + # Inner band: 25-75 + ax.fill_between(env.index, env["p25"], env["p75"], + color=ens_color, alpha=0.30, zorder=3, + label="25th-75th percentile") + # Median + ax.plot(env.index, env["p50"], + color=ens_color, lw=2.2, zorder=5, + label="Ensemble median (600 members)") + + # USGS obs + obs_w = obs_series.loc[t_start:t_end] + ax.plot(obs_w.index, obs_w.values, + color=obs_color, lw=1.8, zorder=6, label="USGS obs") + + ax.set_xlim(t_start, t_end) + ax.set_ylabel("Discharge (m3/s)", fontsize=11) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--ensemble-dir", default=DEFAULT_ENSEMBLE_DIR) + parser.add_argument("--cat-id", default=DEFAULT_CAT_ID) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--out-dir", default=None) + args = parser.parse_args() + + cat_dir = os.path.join(args.ensemble_dir, args.cat_id) + out_dir = args.out_dir or cat_dir + os.makedirs(out_dir, exist_ok=True) + + parquet_path = os.path.join(cat_dir, f"{args.cat_id}_crossed_ensemble.parquet") + print(f"Loading crossed ensemble: {parquet_path}") + df, member_cols = load_ensemble(parquet_path) + obs = load_usgs(args.usgs_csv) + print(f" {df['issue_time'].nunique()} issue times, " + f"{len(member_cols)} members per row") + + # ------------------------------------------------------------------ # + # Figure 4b-full: full Helene window probabilistic envelope # + # ------------------------------------------------------------------ # + print("Building percentile envelope (full window)...") + env_full = build_percentile_envelope(df, member_cols, PLOT_START, PLOT_END) + print(f" Valid times with data: {len(env_full)}, " + f"median members per timestep: {env_full['n'].median():.0f}") + + fig, ax = plt.subplots(figsize=(17, 7)) + plot_envelope(ax, env_full, obs, PLOT_START, PLOT_END) + ax.set_xlabel("Date", fontsize=11) + ax.set_title( + f"4b -- 600-member crossed ensemble (30 met x 20 hydro-state) | {args.cat_id}\n" + "Shaded: 5-95th (light) and 25-75th (dark) percentile across all members " + "and issue times | Sep 24-30 2024 | USGS 03463300", + fontsize=11, + ) + ax.legend(fontsize=10, loc="upper left", frameon=True, framealpha=0.9) + _format_xaxis(ax) + plt.tight_layout() + out1 = os.path.join(out_dir, f"{args.cat_id}_4b_crossed_ensemble_helene.png") + plt.savefig(out1, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + # ------------------------------------------------------------------ # + # Figure 4b-peak: zoomed to Helene peak window # + # ------------------------------------------------------------------ # + print("Building percentile envelope (peak window)...") + env_peak = build_percentile_envelope(df, member_cols, PEAK_START, PEAK_END) + + fig, (ax_main, ax_zoom) = plt.subplots( + 2, 1, figsize=(14, 11), + gridspec_kw={"height_ratios": [1, 1.4]}, sharex=False) + + # Top panel: full window context (smaller) + plot_envelope(ax_main, env_full, obs, PLOT_START, PLOT_END, helene_label=False) + ax_main.set_title("Full Helene window context (Sep 24 - Sep 28 18 UTC)", fontsize=10) + _format_xaxis(ax_main) + # Shade the zoom region on context panel + ax_main.axvspan(PEAK_START, PEAK_END, color="gold", alpha=0.18, zorder=1) + ax_main.text(PEAK_START + pd.Timedelta(hours=1), 0.97, "zoom", + transform=ax_main.get_xaxis_transform(), + fontsize=8, color="goldenrod", ha="left", va="top") + + # Bottom panel: peak zoom + plot_envelope(ax_zoom, env_peak, obs, PEAK_START, PEAK_END, helene_label=True) + ax_zoom.set_xlabel("Date", fontsize=11) + ax_zoom.set_title( + "Helene peak zoom (Sep 26 12 UTC - Sep 28 06 UTC)", fontsize=10) + _format_xaxis(ax_zoom, major_interval=1, minor_hours=3) + ax_zoom.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %HUTC")) + ax_zoom.legend(fontsize=10, loc="lower right", frameon=True, framealpha=0.9) + + fig.suptitle( + f"4b -- 600-member crossed ensemble vs USGS 03463300 | {args.cat_id}\n" + "30 met forcing draws x 20 DA analysis state draws | R=0.07 mm2/h2", + fontsize=12, y=1.01, + ) + plt.tight_layout() + out2 = os.path.join(out_dir, f"{args.cat_id}_4b_crossed_ensemble_peak.png") + plt.savefig(out2, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_perturbation_category_shaded.py b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_perturbation_category_shaded.py new file mode 100644 index 00000000..b845c52d --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_perturbation_category_shaded.py @@ -0,0 +1,165 @@ +""" +Per-catchment shaded ensemble-band plot, organized by perturbation category. + +Three categories: + 1. Initial states (red) + 2. Meteorological forcings (blue) + 3. Hydrological states (green) + +For each category, all 20 ensemble members are shown as a shaded band +(min-max envelope fill) plus a thicker median line in the same color. +Qkrig observation overlaid in black. Hurricane Helene peak window shaded +in pink. Styled after a standard ensemble-forecast figure layout. + +Inputs (existing per-source sensitivity CSVs from run_perturbation_sensitivity.py): + //_sensitivity_init.csv (20 members, init only) + //_sensitivity_forcing.csv (20 members, forcing only) + //_sensitivity_process.csv (20 members, process noise only) + //_test_results.csv (Qkrig obs) + +Outputs: + //_perturbation_categories_linear.png + //_perturbation_categories_log.png +""" +import os +import pandas as pd +import numpy as np +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import matplotlib.patches as mpatches + +CAT = "cat-1016300" + +SEN_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_sensitivity" +OBS_DIR = "/mnt/disk2/1400_sites_helene/catchment_ts_03463300_with_variance" + +OUT_LINEAR = os.path.join(SEN_DIR, CAT, f"{CAT}_perturbation_categories_linear.png") +OUT_LOG = os.path.join(SEN_DIR, CAT, f"{CAT}_perturbation_categories_log.png") + +# Plot window — wider context, similar to the paper's Sep 10 - Oct 08 +PLOT_START = pd.Timestamp("2024-09-20") +PLOT_END = pd.Timestamp("2024-10-05") + +# Helene peak band +HELENE_START = pd.Timestamp("2024-09-26 12:00:00") +HELENE_END = pd.Timestamp("2024-09-28 00:00:00") + +# Category configuration: file suffix, display label, color +CATEGORIES = [ + ("init", "Initial states only", "tab:red"), + ("forcing", "Meteorological forcings only", "tab:blue"), + ("process", "Hydrological states only", "tab:green"), +] + + +def load_members(source): + path = os.path.join(SEN_DIR, CAT, f"{CAT}_sensitivity_{source}.csv") + if not os.path.exists(path): + return None, None + df = pd.read_csv(path, parse_dates=["date"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + return df["date"].values, df[member_cols].to_numpy(dtype=float) + + +def load_obs(): + p = os.path.join(OBS_DIR, f"{CAT}.csv") + if not os.path.exists(p): + return None, None + df = pd.read_csv(p, parse_dates=["date"]) + return df["date"].values, df["qkrig"].values + + +def plot_panel(ax, obs_dates, obs_vals, log_y=False): + handles_labels = [] # for the legend + + # Plot each category as a shaded band + median line + for source, label, color in CATEGORIES: + dates, q = load_members(source) + if dates is None: + continue + d = pd.to_datetime(dates) + mask = (d >= PLOT_START) & (d <= PLOT_END) + if mask.sum() == 0: + continue + # Min-max envelope across all 20 members per timestep (widest possible band). + # Bands are visually narrow even with min/max because perturbations are + # tuned for production EnKF stability, not for max visible spread. + q_window = q[mask, :] + qmin = np.nanmin(q_window, axis=1) + qmax = np.nanmax(q_window, axis=1) + median = np.nanmedian(q_window, axis=1) + + ax.fill_between(d[mask], qmin, qmax, + color=color, alpha=0.30, zorder=2, + edgecolor="none") + line, = ax.plot(d[mask], median, + color=color, lw=1.7, alpha=0.95, zorder=3, + label=f"{label} (N=20)") + handles_labels.append((line, label)) + + # Helene peak shaded band (vertical) + ax.axvspan(HELENE_START, HELENE_END, + color="salmon", alpha=0.15, zorder=1) + ax.text((HELENE_START + (HELENE_END - HELENE_START) / 2), + ax.get_ylim()[1] if not log_y else 1.0, + "Helene peak", + fontsize=9, color="salmon", + ha="center", va="bottom", zorder=3) + + # Observation + if obs_dates is not None: + od = pd.to_datetime(obs_dates) + om = (od >= PLOT_START) & (od <= PLOT_END) + ax.plot(od[om], obs_vals[om], + color="black", lw=1.4, label="Qkrig (obs)", zorder=4) + + ax.set_xlabel("Date", fontsize=10) + if log_y: + ax.set_yscale("log") + ax.set_ylabel(r"q (mm hr$^{-1}$) [log scale]", fontsize=10) + ax.set_ylim(0.01, None) + else: + ax.set_ylabel("Discharge (mm/h)", fontsize=10) + ax.tick_params(axis="x", rotation=30, labelsize=9) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=2)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.grid(True, alpha=0.2) + ax.legend(loc="upper left", fontsize=9, frameon=True, framealpha=0.9) + + +def main(): + obs_dates, obs_vals = load_obs() + + # ----- Linear-y ----- + fig, ax = plt.subplots(1, 1, figsize=(20, 7)) + plot_panel(ax, obs_dates, obs_vals, log_y=False) + fig.suptitle( + f"Ensemble forecast by perturbation category - {CAT} - " + f"Sep 20 - Oct 5, 2024 (Hurricane Helene window)\n" + "Shaded bands = min-max envelope across 20 members. " + "Lines = ensemble median. DA off in all three sub-experiments.", + fontsize=11, y=0.995, + ) + plt.tight_layout() + plt.savefig(OUT_LINEAR, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_LINEAR}") + + # ----- Log-y (paper style) ----- + fig, ax = plt.subplots(1, 1, figsize=(20, 7)) + plot_panel(ax, obs_dates, obs_vals, log_y=True) + fig.suptitle( + f"Ensemble forecast by perturbation category - {CAT} - log-scale q - " + f"Sep 20 - Oct 5, 2024 (Hurricane Helene window)\n" + "Shaded bands = min-max envelope across 20 members. " + "Lines = ensemble median. DA off in all three sub-experiments.", + fontsize=11, y=0.995, + ) + plt.tight_layout() + plt.savefig(OUT_LOG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_LOG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_combined.py b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_combined.py new file mode 100644 index 00000000..c6d531a0 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_combined.py @@ -0,0 +1,235 @@ +#!/usr/bin/env python3 +""" +plot_routed_ensemble_combined.py +Combined poster plot: 600-member routed ensemble envelope + +deterministic DA lines (Vrugt/no-Vrugt) vs USGS gauge obs at 03463300. + +Inputs: + --vrugt-csv : routed_Q_test.csv from vrugt_dynamic_routed/ + --novrugt-csv : routed_Q_test.csv from novrugt_dynamic_routed/ + --ensemble-pq : routed_crossed_ensemble.parquet from run_route_crossed_ensemble.py + --out-dir : output directory + +The ensemble parquet columns: issue_time, lead_hour, member_0000..member_0599 (Q in m3/s) + +Usage: + python3 plot_routed_ensemble_combined.py \ + --vrugt-csv /mnt/disk2/suma_helen_poster/da_results/vrugt_dynamic_routed/routed_Q_test.csv \ + --novrugt-csv /mnt/disk2/suma_helen_poster/da_results/novrugt_dynamic_routed/routed_Q_test.csv \ + --ensemble-pq /mnt/disk2/suma_helen_poster/da_results/v2_crossed_ensemble_routed/routed_crossed_ensemble.parquet \ + --out-dir /mnt/disk2/suma_helen_poster/da_results/comparison_plots +""" + +import argparse +import os + +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +HELENE_START = pd.Timestamp("2024-09-25 00:00:00") +HELENE_END = pd.Timestamp("2024-09-29 00:00:00") +ZOOM_START = pd.Timestamp("2024-09-24 00:00:00") +ZOOM_END = pd.Timestamp("2024-09-28 18:00:00") # cut trailing ensemble bump + +COLOR_OBS = "black" +COLOR_VRUGT = "#1f77b4" # blue +COLOR_NOVRUGT = "#d62728" # red +COLOR_ENS = "#1f77b4" # same blue family as Vrugt + + +def load_det_csv(path): + df = pd.read_csv(path, parse_dates=["date"]) + return df.set_index("date").sort_index() + + +def build_envelope(pq_path, zoom_start, zoom_end): + """ + Load routed ensemble parquet and build percentile envelope by valid_time. + Returns DataFrame indexed by valid_time with columns p05, p25, p50, p75, p95. + """ + df = pd.read_parquet(pq_path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = [c for c in df.columns if c.startswith("member_")] + + records = [] + for _, row in df.iterrows(): + t0 = row["issue_time"] + lead = int(row["lead_hour"]) + vt = t0 + pd.Timedelta(hours=lead) + vals = row[member_cols].values.astype(float) + records.append({"valid_time": vt, "vals": vals}) + + # Group by valid_time (floor to hour), stack all member values across issue_times + from collections import defaultdict + groups = defaultdict(list) + for rec in records: + vt = rec["valid_time"].floor("1h") + groups[vt].extend(rec["vals"].tolist()) + + rows = [] + for vt in sorted(groups): + if vt < zoom_start or vt > zoom_end: + continue + v = np.array(groups[vt]) + v = v[~np.isnan(v)] + if len(v) == 0: + continue + rows.append({ + "valid_time": vt, + "p05": np.percentile(v, 5), + "p25": np.percentile(v, 25), + "p50": np.percentile(v, 50), + "p75": np.percentile(v, 75), + "p95": np.percentile(v, 95), + }) + + return pd.DataFrame(rows).set_index("valid_time") + + +def compute_kge(obs, sim): + mask = ~(np.isnan(obs) | np.isnan(sim)) + o, s = obs[mask], sim[mask] + if len(o) < 2 or np.std(o) == 0: + return np.nan + r = np.corrcoef(o, s)[0, 1] + return 1.0 - np.sqrt((r-1)**2 + (np.std(s)/np.std(o)-1)**2 + (np.mean(s)/np.mean(o)-1)**2) + + +def _format_xaxis(ax, locator, fmt, rotation=30): + ax.xaxis.set_major_locator(locator) + ax.xaxis.set_major_formatter(mdates.DateFormatter(fmt)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=rotation, ha="right", fontsize=9) + + +def _add_helene_band(ax): + ax.axvspan(HELENE_START, HELENE_END, color="gold", alpha=0.15, zorder=0) + + +def plot_zoom(ax, dates, obs, sim_v, sim_nv, kge_v, kge_nv, env=None): + mask = (dates >= ZOOM_START) & (dates <= ZOOM_END) + dz = dates[mask] + + # Ensemble envelope (bottom layer) + if env is not None: + ax.fill_between(env.index, env["p05"], env["p95"], + color=COLOR_ENS, alpha=0.18, label="Ensemble 5-95th pct") + ax.fill_between(env.index, env["p25"], env["p75"], + color=COLOR_ENS, alpha=0.30, label="Ensemble 25-75th pct") + ax.plot(env.index, env["p50"], + color=COLOR_ENS, lw=1.2, linestyle="--", alpha=0.7, label="Ensemble median") + + # Deterministic DA lines + ax.plot(dz, sim_nv[mask], color=COLOR_NOVRUGT, lw=1.4, linestyle="--", + label=f"DA -- constant R KGE={kge_nv:.3f}", zorder=3) + ax.plot(dz, sim_v[mask], color=COLOR_VRUGT, lw=1.6, + label=f"DA -- dynamic R (Vrugt) KGE={kge_v:.3f}", zorder=4) + + # USGS obs on top + ax.plot(dz, obs[mask], color=COLOR_OBS, lw=1.6, label="USGS obs (gauge 03463300)", zorder=5) + + peak_usgs = np.nanmax(obs[mask]) + ax.axhline(peak_usgs, color=COLOR_OBS, lw=0.7, linestyle=":", alpha=0.5) + # Place annotation to the right of the peak to avoid legend overlap + ax.text(pd.Timestamp("2024-09-28 06:00:00"), peak_usgs * 1.02, + f"USGS peak {peak_usgs:.0f} m³/s", fontsize=8.5, color="black", + alpha=0.8, ha="right") + + _add_helene_band(ax) + _format_xaxis(ax, mdates.DayLocator(interval=1), "%b %d") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--vrugt-csv", required=True) + parser.add_argument("--novrugt-csv", required=True) + parser.add_argument("--ensemble-pq", required=True, + help="routed_crossed_ensemble.parquet from run_route_crossed_ensemble.py") + parser.add_argument("--out-dir", required=True) + args = parser.parse_args() + os.makedirs(args.out_dir, exist_ok=True) + + vrugt = load_det_csv(args.vrugt_csv) + novrugt = load_det_csv(args.novrugt_csv) + + obs = vrugt["Q_usgs_m3s"].values + dates = vrugt.index + sim_v = vrugt["Q_routed_m3s"].values + sim_nv = novrugt["Q_routed_m3s"].reindex(dates).values + + kge_v = compute_kge(obs, sim_v) + kge_nv = compute_kge(obs, sim_nv) + + print("Building ensemble envelope...") + env = build_envelope(args.ensemble_pq, ZOOM_START, ZOOM_END) + print(f" Envelope: {len(env)} timesteps, " + f"peak p95={env['p95'].max():.1f} m3/s, peak p50={env['p50'].max():.1f} m3/s") + + peak_usgs = np.nanmax(obs[(dates >= ZOOM_START) & (dates <= ZOOM_END)]) + in_env = env["p95"].max() >= peak_usgs + print(f" USGS peak {peak_usgs:.1f} m3/s {'IS' if in_env else 'IS NOT'} within p95 envelope") + + # ------------------------------------------------------------------ # + # Figure 1: Helene zoom only -- main poster panel + # ------------------------------------------------------------------ # + fig, ax = plt.subplots(figsize=(11, 5.5)) + plot_zoom(ax, dates, obs, sim_v, sim_nv, kge_v, kge_nv, env) + ax.set_xlabel("Date (UTC)", fontsize=11) + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_title( + "Hurricane Helene -- Probabilistic discharge forecast at gauge 03463300 (South Toe River)\n" + "600-member CFE ensemble + DA (Muskingum routing)", + fontsize=11) + ax.legend(fontsize=9, loc="lower right") + ax.grid(True, alpha=0.25) + plt.tight_layout() + out1 = os.path.join(args.out_dir, "helene_ensemble_vs_usgs.png") + plt.savefig(out1, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {out1}") + + # ------------------------------------------------------------------ # + # Figure 2: two-panel (full period + Helene zoom) + # ------------------------------------------------------------------ # + fig, (ax_top, ax_bot) = plt.subplots( + 2, 1, figsize=(14, 9), + gridspec_kw={"height_ratios": [1, 1.5]}) + + # Top: full test period (deterministic only, no ensemble for readability) + ax_top.plot(dates, obs, color=COLOR_OBS, lw=0.8, label="USGS obs", zorder=3) + ax_top.plot(dates, sim_v, color=COLOR_VRUGT, lw=0.9, + label=f"DA dynamic R (Vrugt) KGE={kge_v:.3f}", zorder=2) + ax_top.plot(dates, sim_nv, color=COLOR_NOVRUGT, lw=0.9, linestyle="--", + label=f"DA constant R KGE={kge_nv:.3f}", zorder=2) + ax_top.axvspan(ZOOM_START, ZOOM_END, color="gray", alpha=0.12, zorder=0) + ax_top.axvspan(HELENE_START, HELENE_END, color="gold", alpha=0.15, zorder=0) + ax_top.text(HELENE_START + pd.Timedelta(days=0.5), + np.nanmax(obs) * 0.88, "Helene", + fontsize=9, color="goldenrod", fontweight="bold") + ax_top.set_ylabel("Discharge (m³/s)", fontsize=10) + ax_top.set_title( + "Routed discharge at gauge 03463300 (South Toe River) -- Test period\n" + "DA with dynamic Vrugt R vs. constant R -- Muskingum routing", + fontsize=11) + ax_top.legend(fontsize=9, loc="upper left") + ax_top.grid(True, alpha=0.25) + _format_xaxis(ax_top, mdates.MonthLocator(interval=2), "%Y-%m") + + # Bottom: Helene zoom with ensemble + plot_zoom(ax_bot, dates, obs, sim_v, sim_nv, kge_v, kge_nv, env) + ax_bot.set_xlabel("Date (UTC)", fontsize=10) + ax_bot.set_ylabel("Discharge (m³/s)", fontsize=10) + ax_bot.set_title("Hurricane Helene window -- 600-member ensemble envelope", fontsize=10) + ax_bot.legend(fontsize=9, loc="lower right") + ax_bot.grid(True, alpha=0.25) + + plt.tight_layout() + out2 = os.path.join(args.out_dir, "helene_ensemble_twopanel.png") + plt.savefig(out2, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {out2}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_vs_usgs.py b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_vs_usgs.py new file mode 100644 index 00000000..55c52d85 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/plot_routed_ensemble_vs_usgs.py @@ -0,0 +1,176 @@ +#!/usr/bin/env python3 +""" +plot_routed_ensemble_vs_usgs.py +Plot T-route routed 600-member ensemble envelope vs USGS at gauge 03463300. + +Input: routed_crossed_ensemble.parquet from run_route_crossed_ensemble.py + Columns: issue_time, lead_hour, member_0000..member_0599 (Q in m3/s) + +Usage: + python3 plot_routed_ensemble_vs_usgs.py \ + --routed-pq /mnt/disk2/.../folder1_vrugt_routed/routed_crossed_ensemble.parquet \ + --usgs-csv /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv \ + --label "F1 Vrugt — 1 gauge holdout" \ + --out-dir ~/plots_f1_1gauge +""" + +import argparse +from collections import defaultdict +from pathlib import Path + +import matplotlib.dates as mdates +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd + +WATERSHED_AREA_KM2 = 113.18 # cat-1016300 drainage area +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +PLOT_START = pd.Timestamp("2024-09-24 00:00:00") +PLOT_END = pd.Timestamp("2024-09-28 18:00:00") +PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +PEAK_END = pd.Timestamp("2024-09-28 06:00:00") +HELENE_PEAK = pd.Timestamp("2024-09-27 12:00:00") + +COLOR_ENS = "#1f77b4" +COLOR_OBS = "black" + + +def build_envelope(pq_path, t_start, t_end): + df = pd.read_parquet(pq_path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = [c for c in df.columns if c.startswith("member_")] + + groups = defaultdict(list) + for _, row in df.iterrows(): + vt = (row["issue_time"] + pd.Timedelta(hours=int(row["lead_hour"]))).floor("1h") + if t_start <= vt <= t_end: + groups[vt].extend(row[member_cols].values.astype(float).tolist()) + + rows = [] + for vt in sorted(groups): + v = np.array(groups[vt]) + v = v[~np.isnan(v)] + if len(v) == 0: + continue + rows.append({"valid_time": vt, + "p05": np.percentile(v, 5), "p25": np.percentile(v, 25), + "p50": np.percentile(v, 50), "p75": np.percentile(v, 75), + "p95": np.percentile(v, 95)}) + return pd.DataFrame(rows).set_index("valid_time") + + +def load_usgs(path): + df = pd.read_csv(path) + print(f" USGS CSV columns: {list(df.columns)}") + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() or "discharge" in c.lower()) + print(f" Using date_col={date_col!r}, q_col={q_col!r}") + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series.sort_index() + + +def compute_kge(obs, sim): + mask = ~(np.isnan(obs) | np.isnan(sim)) + o, s = obs[mask], sim[mask] + if len(o) < 2 or np.std(o) == 0: + return np.nan + r = np.corrcoef(o, s)[0, 1] + return 1 - np.sqrt((r-1)**2 + (np.std(s)/np.std(o)-1)**2 + (np.mean(s)/np.mean(o)-1)**2) + + +def _shade(ax, env): + ax.fill_between(env.index, env["p05"], env["p95"], + color=COLOR_ENS, alpha=0.18, label="5th–95th percentile") + ax.fill_between(env.index, env["p25"], env["p75"], + color=COLOR_ENS, alpha=0.32, label="25th–75th percentile") + ax.plot(env.index, env["p50"], color=COLOR_ENS, lw=1.5, label="Ensemble median (600 members)") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--routed-pq", required=True) + parser.add_argument("--usgs-csv", required=True) + parser.add_argument("--label", default="F1 Vrugt — 1 gauge holdout") + parser.add_argument("--out-dir", required=True) + args = parser.parse_args() + + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + print("Building full-window envelope...") + env_full = build_envelope(args.routed_pq, PLOT_START, PLOT_END) + print("Building peak-window envelope...") + env_peak = build_envelope(args.routed_pq, PEAK_START, PEAK_END) + + print("Loading USGS obs...") + try: + obs_all = load_usgs(args.usgs_csv) + obs_full = obs_all.loc[PLOT_START:PLOT_END] + obs_peak = obs_all.loc[PEAK_START:PEAK_END] + print(f" Loaded {len(obs_all)} USGS obs rows; " + f"full window: {len(obs_full)}, peak window: {len(obs_peak)}") + except Exception as e: + print(f" Warning: USGS load failed ({e}) — plotting without obs") + obs_full = obs_peak = None + + kge = np.nan + if obs_full is not None and len(env_full) > 0: + med = env_full["p50"] + obs_hourly = obs_full.copy() + obs_hourly.index = obs_hourly.index.floor("1h") + obs_hourly = obs_hourly[~obs_hourly.index.duplicated(keep="first")] + aligned = obs_hourly.reindex(med.index) + print(f" KGE alignment: {aligned.notna().sum()} matched out of {len(med)} timesteps") + kge = compute_kge(aligned.values, med.values) + + # ── Figure 1: Full Helene window ────────────────────────────────────────── + fig, ax = plt.subplots(figsize=(12, 4.5)) + _shade(ax, env_full) + if obs_full is not None: + ax.plot(obs_full.index, obs_full.values, color=COLOR_OBS, lw=1.6, + label="USGS obs (gauge 03463300)") + ax.axvspan(PEAK_START, PEAK_END, color="salmon", alpha=0.12, zorder=0, label="Helene peak window") + ax.set_ylabel("Discharge (m³/s)") + ax.set_xlabel("Date") + ax.set_title(f"Routed 600-member ensemble vs USGS | {args.label}\n" + f"Ensemble median KGE = {kge:.3f} | Sep 24–28 2024") + ax.legend(fontsize=9, loc="upper left") + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + plt.tight_layout() + out = out_dir / "routed_ensemble_vs_usgs_full.png" + fig.savefig(out, dpi=150) + print(f"Saved: {out}") + plt.close(fig) + + # ── Figure 2: Peak zoom ─────────────────────────────────────────────────── + fig, ax = plt.subplots(figsize=(10, 4.5)) + _shade(ax, env_peak) + if obs_peak is not None: + ax.plot(obs_peak.index, obs_peak.values, color=COLOR_OBS, lw=1.6, + label="USGS obs (gauge 03463300)") + peak_val = float(obs_peak.max()) + ax.axhline(peak_val, color=COLOR_OBS, lw=0.7, ls=":", alpha=0.5) + ax.annotate(f"USGS peak\n{peak_val:.0f} m³/s", + xy=(HELENE_PEAK, peak_val), xytext=(10, -40), + textcoords="offset points", fontsize=8, color=COLOR_OBS) + ax.set_ylabel("Discharge (m³/s)") + ax.set_xlabel("Date") + ax.set_title(f"Helene peak zoom | {args.label}") + ax.legend(fontsize=9) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %HUTC")) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + plt.tight_layout() + out = out_dir / "routed_ensemble_vs_usgs_peak.png" + fig.savefig(out, dpi=150) + print(f"Saved: {out}") + plt.close(fig) + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/run_4b_crossed_f4.sh b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/run_4b_crossed_f4.sh new file mode 100644 index 00000000..7de57f45 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/run_4b_crossed_f4.sh @@ -0,0 +1,40 @@ +#!/usr/bin/env bash +# F4 dynamic variance direct (20pct gauge holdout) — 4b crossed ensemble plots. +# Runs plot_crossed_ensemble.py for all 21 catchments. +# Reads //_crossed_ensemble.parquet (already generated). +# +# Usage: +# bash run_4b_crossed_f4.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +DATA_DIR=/mnt/disk2/1400_sites_helene/da_crossed_dynamic_novrugt_seeded +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +echo "[F4-4b-crossed] Crossed ensemble plots — data: $DATA_DIR" + +for CAT in "${CATS[@]}"; do + PQ="$DATA_DIR/$CAT/${CAT}_crossed_ensemble.parquet" + if [ ! -f "$PQ" ]; then + echo " [$CAT] No crossed ensemble parquet — skipping" + continue + fi + echo " [$CAT] crossed ensemble..." + $TROUTE "$SCRIPT_DIR/plot_crossed_ensemble.py" \ + --ensemble-dir "$DATA_DIR" \ + --cat-id "$CAT" \ + --usgs-csv "$USGS_CSV" +done + +echo "[F4-4b-crossed] Done." diff --git a/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/run_4b_f4.sh b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/run_4b_f4.sh new file mode 100644 index 00000000..1bde9583 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4b_ensemble_vs_obs/run_4b_f4.sh @@ -0,0 +1,40 @@ +#!/usr/bin/env bash +# F4 dynamic variance direct (20pct gauge holdout) — 4b routed ensemble plots. +# +# Runs: +# 1. plot_routed_ensemble_vs_usgs.py — F4 ensemble envelope vs USGS at outlet +# 2. plot_routed_ensemble_combined.py — F1 Vrugt vs F4 dynamic variance direct comparison +# +# Requires routed_Q_test.csv (both folders) and routed_crossed_ensemble.parquet +# (CROSSED_DIR) to be present. +# +# Usage: +# bash run_4b_f4.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +F1_DA_DIR=/mnt/disk2/suma_helen_poster/da_results/v2_perturbation_da_on +F4_DA_DIR=/mnt/disk2/1400_sites_helene/da_arms_dynamic_novrugt_seeded +CROSSED_DIR=/mnt/disk2/1400_sites_helene/da_crossed_dynamic_novrugt_seeded +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv +OUT_DIR="$CROSSED_DIR" + +echo "[F4-4b] Routed ensemble vs USGS (F4 only)..." +$TROUTE "$SCRIPT_DIR/plot_routed_ensemble_vs_usgs.py" \ + --routed-pq "$CROSSED_DIR/routed_crossed_ensemble.parquet" \ + --usgs-csv "$USGS_CSV" \ + --label "F4 Dynamic variance direct — 20pct gauge holdout" \ + --out-dir "$OUT_DIR" + +echo "[F4-4b] Combined comparison: F1 Vrugt vs F4 Dynamic variance direct..." +$TROUTE "$SCRIPT_DIR/plot_routed_ensemble_combined.py" \ + --vrugt-csv "$F1_DA_DIR/routed_Q_test.csv" \ + --novrugt-csv "$F4_DA_DIR/routed_Q_test.csv" \ + --ensemble-pq "$CROSSED_DIR/routed_crossed_ensemble.parquet" \ + --out-dir "$OUT_DIR" + +echo "[F4-4b] Done. Outputs in: $OUT_DIR" +ls "$OUT_DIR"/*.png 2>/dev/null || echo " (no PNGs in out dir)" diff --git a/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4c_timeseries/plot_forecast_spaghetti.py b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4c_timeseries/plot_forecast_spaghetti.py new file mode 100644 index 00000000..655edaef --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4c_timeseries/plot_forecast_spaghetti.py @@ -0,0 +1,176 @@ +""" +plot_forecast_spaghetti.py + +"All forecasts ending at this hour" — verification-time-fixed spaghetti view. + +Window: Sep 24 18 UTC -> Sep 30 06 UTC. +For every initialization time whose 18-hour forecast overlaps this window: + DA : shaded min/max band across 20 members + ensemble mean line + OL : ensemble mean line only (no shading), dashed + +Lines are colored by initialization DATE (7 colors, Sep 24-30) so the +temporal progression is readable. USGS obs overlaid in black. + +Inputs: + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + Wide format: issue_time, lead_hour, member_00..member_19 (q in m³/s) + +Obs: + /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +Outputs: + /forecast_spaghetti_helene.png +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import matplotlib.patches as mpatches + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast_hardcoded_r_routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +# Verification window — only trajectories whose valid_time falls here are shown +PLOT_START = pd.Timestamp("2024-09-24 18:00:00") +PLOT_END = pd.Timestamp("2024-09-30 06:00:00") + +HELENE_PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +HELENE_PEAK_END = pd.Timestamp("2024-09-28 00:00:00") + +# One color per init date (Sep 24-30) +DATE_COLORS = { + "2024-09-24": "#1f77b4", + "2024-09-25": "#ff7f0e", + "2024-09-26": "#2ca02c", + "2024-09-27": "#d62728", + "2024-09-28": "#9467bd", + "2024-09-29": "#8c564b", + "2024-09-30": "#e377c2", +} + + +def load_parquet(path): + df = pd.read_parquet(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + df["ens_mean"] = df[member_cols].mean(axis=1) + df["ens_min"] = df[member_cols].min(axis=1) + df["ens_max"] = df[member_cols].max(axis=1) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + return df + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() + or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--route-dir", default=DEFAULT_ROUTE_DIR) + parser.add_argument("--out-dir", default=None) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--da-name", default="routed_leadtime_da_full.parquet") + parser.add_argument("--ol-name", default="routed_leadtime_openloop_full.parquet") + args = parser.parse_args() + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + print("Loading parquets...") + da = load_parquet(os.path.join(args.route_dir, args.da_name)) + ol = load_parquet(os.path.join(args.route_dir, args.ol_name)) + obs = load_usgs(args.usgs_csv) + + # Keep only valid_times inside the plot window + da = da[(da["valid_time"] >= PLOT_START) & (da["valid_time"] <= PLOT_END)] + ol = ol[(ol["valid_time"] >= PLOT_START) & (ol["valid_time"] <= PLOT_END)] + + issue_times = sorted(da["issue_time"].unique()) + print(f" {len(issue_times)} initialization times contribute to this window") + + fig, ax = plt.subplots(figsize=(18, 6)) + + # Helene peak shading + ax.axvspan(HELENE_PEAK_START, HELENE_PEAK_END, + color="salmon", alpha=0.12, zorder=0) + ax.text(HELENE_PEAK_START + pd.Timedelta(hours=6), 1.0, "Helene peak", + transform=ax.get_xaxis_transform(), + fontsize=9, color="firebrick", ha="left", va="top") + + # Plot each initialization time's trajectory + for t0 in issue_times: + date_str = str(pd.Timestamp(t0).date()) + color = DATE_COLORS.get(date_str, "gray") + + da_t = da[da["issue_time"] == t0].sort_values("valid_time") + ol_t = ol[ol["issue_time"] == t0].sort_values("valid_time") + + if da_t.empty: + continue + + # DA: shaded band (min-max across 20 members) + ensemble mean + ax.fill_between(da_t["valid_time"], da_t["ens_min"], da_t["ens_max"], + color=color, alpha=0.06, zorder=2) + ax.plot(da_t["valid_time"], da_t["ens_mean"], + color=color, lw=0.8, alpha=0.55, zorder=3) + + # OL: ensemble mean only, dashed + if not ol_t.empty: + ax.plot(ol_t["valid_time"], ol_t["ens_mean"], + color=color, lw=0.6, alpha=0.30, linestyle="--", zorder=2) + + # USGS obs + obs_w = obs.loc[PLOT_START:PLOT_END] + ax.plot(obs_w.index, obs_w.values, + color="black", lw=1.8, zorder=6, label="USGS obs") + + # Legend: one patch per init date + obs + DA/OL style + patches = [mpatches.Patch(color=c, label=f"Init {d}") + for d, c in DATE_COLORS.items()] + patches.append(plt.Line2D([0], [0], color="black", lw=1.8, label="USGS obs")) + patches.append(plt.Line2D([0], [0], color="gray", lw=1.2, alpha=0.7, + label="DA — ens mean (solid) + spread (shaded)")) + patches.append(plt.Line2D([0], [0], color="gray", lw=1.0, linestyle="--", + alpha=0.5, label="OL — ens mean (dashed)")) + ax.legend(handles=patches, fontsize=8, loc="upper left", + frameon=True, framealpha=0.9, ncol=2) + + ax.set_xlim(PLOT_START, PLOT_END) + ax.set_xlabel("Date", fontsize=11) + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_title( + "All 18-hour forecast trajectories ending in this window — USGS 03463300\n" + "Sep 24 18 UTC → Sep 30 06 UTC | DA: shaded band + mean | OL: mean dashed", + fontsize=11, + ) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=1)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.xaxis.set_minor_locator(mdates.HourLocator(interval=6)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.2, lw=0.4) + plt.tight_layout() + + out_path = os.path.join(out_dir, "forecast_spaghetti_helene.png") + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4c_timeseries/plot_helene_issue_time_hydrograph.py b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4c_timeseries/plot_helene_issue_time_hydrograph.py new file mode 100644 index 00000000..4e2e1dc2 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4c_timeseries/plot_helene_issue_time_hydrograph.py @@ -0,0 +1,198 @@ +""" +Per-issue-time forecast hydrograph diagnostic. + +For two contrasting issue times — one storm-peak (Helene) and one low-flow — +show what the DA and open-loop ensembles actually produce over the 18-hour +forecast window, alongside the kriging observation. Designed to answer in one +picture: is DA over-shooting, collapsing, or oscillating compared to open-loop? + +Each panel shows: + - All 20 DA members as faint purple lines + median (thick purple) + - All 20 open-loop members as faint gray lines + median (thick gray, dashed) + - Kriging obs as a thick black line (context: t0-6 through t0+18) + - Vertical line at t0 + annotation with obs(t0) and lead-1 ensemble means + +Inputs (from run_lead_time_forecast_sweep.py): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv +Obs (from production EnKF): + //_test_results.csv + +Output: + //_issue_time_hydrograph_helene_vs_lowflow.png +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +DEFAULT_LEADTIME_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast" +DEFAULT_DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_true_enkf_vrugt" + +# Module-level globals set by main() +CAT = "cat-1016300" +LEADTIME_DIR = DEFAULT_LEADTIME_DIR +DA_DIR = DEFAULT_DA_DIR + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" +OBS_COLOR = "black" + +# Two issue times — pick a Helene-peak one and a typical low-flow one +HELENE_T0 = pd.Timestamp("2024-09-26 12:00:00") +LOWFLOW_T0 = pd.Timestamp("2024-03-15 12:00:00") + +CONTEXT_HOURS_BEFORE = 6 # hours of obs context shown before t0 +LEAD_HOURS_AFTER = 18 # forecast horizon + +OUT_PNG = os.path.join( + LEADTIME_DIR, CAT, + f"{CAT}_issue_time_hydrograph_helene_vs_lowflow.png", +) + + +def load_forecasts(path): + df = pd.read_csv(path, parse_dates=['issue_time', 'valid_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + return df, member_cols + + +def load_obs(): + obs_path = os.path.join(DA_DIR, CAT, f"{CAT}_test_results.csv") + df = pd.read_csv(obs_path, parse_dates=['date']) + return df.set_index('date')['obs_mm_h'] + + +def slice_forecast(df, member_cols, t0): + """Return (valid_times, member_array shape (lead, N)) for a single issue time.""" + sub = df[df['issue_time'] == t0].sort_values('lead_hour') + if len(sub) == 0: + raise ValueError(f"No forecast rows for issue_time={t0}") + valid_times = pd.to_datetime(sub['valid_time'].values) + member_arr = sub[member_cols].to_numpy(dtype=float) + return valid_times, member_arr + + +def plot_panel(ax, t0, df_da, df_ol, member_cols, obs_series, title_extra=""): + # Forecast trajectories + vt_da, q_da = slice_forecast(df_da, member_cols, t0) + vt_ol, q_ol = slice_forecast(df_ol, member_cols, t0) + + # Obs context window + ctx_start = t0 - pd.Timedelta(hours=CONTEXT_HOURS_BEFORE) + ctx_end = t0 + pd.Timedelta(hours=LEAD_HOURS_AFTER) + obs_window = obs_series.loc[ctx_start:ctx_end] + + # Open-loop members (drawn first so DA paints on top) + for j in range(q_ol.shape[1]): + ax.plot(vt_ol, q_ol[:, j], color=OL_COLOR, lw=0.6, alpha=0.45, zorder=2) + # DA members + for j in range(q_da.shape[1]): + ax.plot(vt_da, q_da[:, j], color=DA_COLOR, lw=0.6, alpha=0.45, zorder=3) + # Medians + ax.plot(vt_ol, np.nanmedian(q_ol, axis=1), color=OL_COLOR, lw=2.6, + linestyle="--", zorder=4, label="Open-loop median") + ax.plot(vt_da, np.nanmedian(q_da, axis=1), color=DA_COLOR, lw=2.6, + zorder=5, label="DA median") + # Obs + ax.plot(obs_window.index, obs_window.values, color=OBS_COLOR, lw=1.8, + zorder=6, label="Qkrig (obs)") + + # t0 marker + annotation + ax.axvline(t0, color="black", lw=1.0, linestyle=":", alpha=0.6, zorder=1) + obs_t0 = float(obs_series.get(t0, np.nan)) + obs_lead1 = float(obs_series.get(t0 + pd.Timedelta(hours=1), np.nan)) + da_lead1_mean = float(np.nanmean(q_da[0, :])) + ol_lead1_mean = float(np.nanmean(q_ol[0, :])) + + info_text = ( + f"t0 = {t0.strftime('%Y-%m-%d %H:%M')}\n" + f"obs(t0) = {obs_t0:.3f} mm/h\n" + f"obs(t0+1) = {obs_lead1:.3f} mm/h\n" + f"DA mean(t0+1) = {da_lead1_mean:.3f} mm/h\n" + f"OL mean(t0+1) = {ol_lead1_mean:.3f} mm/h" + ) + ax.text(0.02, 0.97, info_text, + transform=ax.transAxes, + fontsize=9, va='top', ha='left', family='monospace', + bbox=dict(boxstyle='round,pad=0.4', facecolor='white', + alpha=0.85, edgecolor='gray')) + + ax.set_title(title_extra, fontsize=12) + ax.set_ylabel("Discharge (mm/h)", fontsize=10) + ax.set_xlabel("Date", fontsize=10) + ax.xaxis.set_major_locator(mdates.HourLocator(interval=3)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %H:%M")) + ax.tick_params(axis='x', rotation=30, labelsize=8) + ax.grid(True, alpha=0.2) + ax.legend(loc='upper right', fontsize=9, frameon=True, framealpha=0.9) + + +def snap_to_nearest(t0, available_times): + """Return the available issue_time closest to t0.""" + times = pd.DatetimeIndex(available_times) + idx = np.abs((times - t0).total_seconds()).argmin() + snapped = times[idx] + if snapped != t0: + print(f" Snapped {t0} -> {snapped} (nearest available issue time)") + return snapped + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', default="cat-1016300") + parser.add_argument('--leadtime-dir', default=DEFAULT_LEADTIME_DIR, + help='Dir holding /_lead_time_forecasts_{da,openloop}.csv') + parser.add_argument('--da-dir', default=DEFAULT_DA_DIR, + help='Dir holding /_test_results.csv (for obs)') + parser.add_argument('--helene-t0', default=None, + help='Issue time for the Helene panel (default: 2024-09-26 12:00:00, ' + 'snapped to nearest available)') + parser.add_argument('--lowflow-t0', default=None, + help='Issue time for the low-flow panel (default: 2024-03-15 00:00:00, ' + 'snapped to nearest available)') + args = parser.parse_args() + + global CAT, LEADTIME_DIR, DA_DIR, OUT_PNG + CAT = args.cat_id + LEADTIME_DIR = args.leadtime_dir + DA_DIR = args.da_dir + OUT_PNG = os.path.join(LEADTIME_DIR, CAT, + f"{CAT}_issue_time_hydrograph_helene_vs_lowflow.png") + + da_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_da.csv") + ol_path = os.path.join(LEADTIME_DIR, CAT, f"{CAT}_lead_time_forecasts_openloop.csv") + df_da, m_cols = load_forecasts(da_path) + df_ol, _ = load_forecasts(ol_path) + obs_series = load_obs() + + available = df_da['issue_time'].unique() + helene_t0 = snap_to_nearest( + pd.Timestamp(args.helene_t0) if args.helene_t0 else HELENE_T0, available) + lowflow_t0 = snap_to_nearest( + pd.Timestamp(args.lowflow_t0) if args.lowflow_t0 else pd.Timestamp("2024-03-15 00:00:00"), + available) + + fig, (ax_h, ax_l) = plt.subplots(1, 2, figsize=(20, 7)) + + plot_panel(ax_h, helene_t0, df_da, df_ol, m_cols, obs_series, + title_extra=f"Helene peak issue time — {helene_t0.strftime('%Y-%m-%d %H:%M')}") + plot_panel(ax_l, lowflow_t0, df_da, df_ol, m_cols, obs_series, + title_extra=f"Low-flow issue time — {lowflow_t0.strftime('%Y-%m-%d %H:%M')}") + + fig.suptitle( + f"Per-issue-time forecast hydrograph — {CAT}\n" + "Faint lines = 20 ensemble members. Thick lines = ensemble medians. " + "Vertical dotted line = t0 (DA stops here for both scenarios).", + fontsize=12, y=1.00, + ) + plt.tight_layout() + plt.savefig(OUT_PNG, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4c_timeseries/plot_lead_time_reconstructed_timeseries.py b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4c_timeseries/plot_lead_time_reconstructed_timeseries.py new file mode 100644 index 00000000..a6d03b03 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4c_timeseries/plot_lead_time_reconstructed_timeseries.py @@ -0,0 +1,263 @@ +""" +Reconstructed time series at gauge 03463300, from our routed +lead-time forecasts. Mirrors the methodology of plot_timeseries_hourly_ensemble.py +(qSpatialAR) but feeds from the EnKF + T-route pipeline instead of the CNN. + +For each target hour t, pool ALL forecasts that land on t (across many +issue_times × many lead_hours × 20 members). Compute median + 5th/95th +percentile envelope. Plot DA and open-loop reconstructions overlaid on USGS +obs, with the Helene window shaded. + +Input parquets (from run_route_troute_leadtime.py): + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + Wide format: issue_time, lead_hour, member_00..member_19 (q in m³/s) + +Obs: + /mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +Output: + /lead_time_reconstructed_timeseries.png + Default plot window: 2024-09-10 → 2024-10-10. + +Each scenario's curve is annotated with NSE vs USGS obs over the plot window. +""" +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" +OBS_COLOR = "black" + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results/v2_lead_time_forecast_hardcoded_r_routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +# Watershed area at gauge 03463300 (sum of 21 catchments per GPKG hydrofabric) +WATERSHED_AREA_KM2 = 113.18 +# mm/h depth → m³/s: × 113.18 km² × 1000 / 3600 = 31.439 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +PLOT_START_DEFAULT = pd.Timestamp("2024-09-10 00:00:00") +PLOT_END_DEFAULT = pd.Timestamp("2024-10-10 23:00:00") + +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-30 23:00:00") + + +def load_parquet_long(path): + df = pd.read_parquet(path) + df['issue_time'] = pd.to_datetime(df['issue_time']) + member_cols = sorted([c for c in df.columns if c.startswith('member_')]) + if not member_cols: + raise ValueError(f"No member columns in {path}") + long = df.melt( + id_vars=['issue_time', 'lead_hour'], + value_vars=member_cols, + var_name='member', + value_name='q_m3s', + ) + long['valid_time'] = (long['issue_time'] + + pd.to_timedelta(long['lead_hour'], unit='h')) + return long + + +def reconstruct(long_df): + """For each valid_time, pool all (issue_time × lead × member) forecasts and + return median, p05, p95, count.""" + grouped = long_df.groupby('valid_time')['q_m3s'].agg( + median='median', + p05=lambda s: s.quantile(0.05), + p95=lambda s: s.quantile(0.95), + count='count', + ).reset_index() + return grouped + + +def load_usgs(usgs_csv): + """Load USGS gauge obs and ensure m³/s. + + The CSV at gauge 03463300 has the column `QObs(mm/h)` — catchment-averaged + depth, not gauge discharge in m³/s. Auto-convert when column name contains + 'mm': m³/s = mm/h × WATERSHED_AREA_KM2 × 1000 / 3600. + """ + df = pd.read_csv(usgs_csv) + date_col = next((c for c in df.columns + if c.lower() in ('datetime', 'date', 'time', 'timestamp')), None) + q_col = next((c for c in df.columns + if 'q' in c.lower() or 'flow' in c.lower() + or 'discharge' in c.lower()), None) + if not date_col or not q_col: + raise ValueError(f"Couldn't id date/Q columns in {usgs_csv}: {df.columns.tolist()}") + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if 'mm' in q_col.lower(): + series = series * MM_H_TO_M3_S + print(f" Converted obs from mm/h (column '{q_col}') to m³/s " + f"using area = {WATERSHED_AREA_KM2} km² (× {MM_H_TO_M3_S:.4f})") + else: + print(f" Loaded obs as-is from column '{q_col}' (assumed m³/s)") + return series + + +def nse(obs, sim): + m = np.isfinite(obs) & np.isfinite(sim) + if m.sum() < 2: + return np.nan + o, s = obs[m], sim[m] + denom = ((o - o.mean()) ** 2).sum() + if denom < 1e-10: + return np.nan + return 1.0 - float(((o - s) ** 2).sum() / denom) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--route-dir', default=DEFAULT_ROUTE_DIR) + parser.add_argument('--out-dir', default=None, + help='Default: --route-dir') + parser.add_argument('--usgs-csv', default=DEFAULT_USGS_CSV) + parser.add_argument('--plot-start', default=str(PLOT_START_DEFAULT)) + parser.add_argument('--plot-end', default=str(PLOT_END_DEFAULT)) + parser.add_argument('--da-name', default='routed_leadtime_da_full.parquet') + parser.add_argument('--ol-name', default='routed_leadtime_openloop_full.parquet') + args = parser.parse_args() + + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + plot_start = pd.Timestamp(args.plot_start) + plot_end = pd.Timestamp(args.plot_end) + + print(f"Loading DA parquet: {args.da_name}") + da_long = load_parquet_long(os.path.join(args.route_dir, args.da_name)) + print(f"Loading OL parquet: {args.ol_name}") + ol_long = load_parquet_long(os.path.join(args.route_dir, args.ol_name)) + print(f"Loading USGS obs: {args.usgs_csv}") + obs = load_usgs(args.usgs_csv) + + print("Reconstructing DA time series (overlapping-leads pool)...") + da_rec = reconstruct(da_long) + print(f" {len(da_rec)} unique target hours") + print("Reconstructing OL time series...") + ol_rec = reconstruct(ol_long) + + # Clip to plot window and join obs + da_rec = da_rec[(da_rec['valid_time'] >= plot_start) & (da_rec['valid_time'] <= plot_end)].copy() + ol_rec = ol_rec[(ol_rec['valid_time'] >= plot_start) & (ol_rec['valid_time'] <= plot_end)].copy() + da_rec['obs'] = da_rec['valid_time'].map(obs) + ol_rec['obs'] = ol_rec['valid_time'].map(obs) + + nse_da = nse(da_rec['obs'].values, da_rec['median'].values) + nse_ol = nse(ol_rec['obs'].values, ol_rec['median'].values) + print(f"NSE (plot window) — DA: {nse_da:.3f} | OL: {nse_ol:.3f}") + + # ---- Plot ---- + fig, ax = plt.subplots(figsize=(16, 5.5)) + fig.patch.set_facecolor("white") + + # Helene shaded + ax.axvspan(HELENE_START, HELENE_END, + color="firebrick", alpha=0.08, zorder=0) + ax.text(HELENE_START + pd.Timedelta(hours=12), 1.0, "Helene", + transform=ax.get_xaxis_transform(), + fontsize=9, color="firebrick", ha="left", va="top", + fontweight="bold") + + # Open-loop band + median + ax.fill_between(ol_rec['valid_time'], ol_rec['p05'], ol_rec['p95'], + color=OL_COLOR, alpha=0.18, zorder=2, + label="Open-loop [5th–95th pctile envelope]") + ax.plot(ol_rec['valid_time'], ol_rec['median'], + color=OL_COLOR, lw=1.2, linestyle='--', zorder=3, + label=f"Open-loop median (NSE={nse_ol:.3f})") + + # DA band + median + ax.fill_between(da_rec['valid_time'], da_rec['p05'], da_rec['p95'], + color=DA_COLOR, alpha=0.22, zorder=4, + label="DA [5th–95th pctile envelope]") + ax.plot(da_rec['valid_time'], da_rec['median'], + color=DA_COLOR, lw=1.4, zorder=5, + label=f"DA median (NSE={nse_da:.3f})") + + # Observed as dots + obs_window = obs.loc[plot_start:plot_end] + finite = np.isfinite(obs_window.values) + ax.scatter(obs_window.index[finite], obs_window.values[finite], + s=6, color=OBS_COLOR, linewidths=0, zorder=6, + label="USGS obs") + + ax.set_xlabel("Date", fontsize=11) + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_title( + f"Reconstructed forecast time series at USGS 03463300\n" + f"Overlapping-leads pool from EnKF forecast ensemble " + f"({plot_start.date()} – {plot_end.date()})", + fontsize=12, + ) + ax.set_xlim(plot_start, plot_end) + ax.xaxis.set_major_locator(mdates.WeekdayLocator(interval=1)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.25, lw=0.4) + ax.legend(loc='upper left', fontsize=9, frameon=True, framealpha=0.92) + + plt.tight_layout() + out_path = os.path.join(out_dir, "lead_time_reconstructed_timeseries.png") + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + # Also produce a Helene-zoom version + helene_zoom_start = pd.Timestamp("2024-09-24 00:00:00") + helene_zoom_end = pd.Timestamp("2024-09-29 23:00:00") + da_z = da_rec[(da_rec['valid_time'] >= helene_zoom_start) + & (da_rec['valid_time'] <= helene_zoom_end)] + ol_z = ol_rec[(ol_rec['valid_time'] >= helene_zoom_start) + & (ol_rec['valid_time'] <= helene_zoom_end)] + nse_da_z = nse(da_z['obs'].values, da_z['median'].values) + nse_ol_z = nse(ol_z['obs'].values, ol_z['median'].values) + print(f"NSE (Helene zoom) — DA: {nse_da_z:.3f} | OL: {nse_ol_z:.3f}") + + fig, ax = plt.subplots(figsize=(14, 5.5)) + fig.patch.set_facecolor("white") + ax.fill_between(ol_z['valid_time'], ol_z['p05'], ol_z['p95'], + color=OL_COLOR, alpha=0.18, zorder=2, + label="Open-loop [5th–95th]") + ax.plot(ol_z['valid_time'], ol_z['median'], color=OL_COLOR, + lw=1.4, linestyle='--', zorder=3, + label=f"Open-loop median (NSE={nse_ol_z:.3f})") + ax.fill_between(da_z['valid_time'], da_z['p05'], da_z['p95'], + color=DA_COLOR, alpha=0.22, zorder=4, + label="DA [5th–95th]") + ax.plot(da_z['valid_time'], da_z['median'], color=DA_COLOR, + lw=1.6, zorder=5, + label=f"DA median (NSE={nse_da_z:.3f})") + obs_z = obs.loc[helene_zoom_start:helene_zoom_end] + finite = np.isfinite(obs_z.values) + ax.scatter(obs_z.index[finite], obs_z.values[finite], + s=12, color=OBS_COLOR, linewidths=0, zorder=6, + label="USGS obs") + ax.set_xlabel("Date", fontsize=11) + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_title( + "Reconstructed forecast time series at USGS 03463300 — Helene zoom\n" + "Sep 24 – 29, 2024 (overlapping-leads pool from EnKF forecast ensemble)", + fontsize=12, + ) + ax.xaxis.set_major_locator(mdates.HourLocator(interval=12)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d %H:%M")) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.25, lw=0.4) + ax.legend(loc='upper left', fontsize=9, frameon=True, framealpha=0.92) + plt.tight_layout() + out_zoom = os.path.join(out_dir, "lead_time_reconstructed_timeseries_helene.png") + plt.savefig(out_zoom, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_zoom}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4c_timeseries/run_4c_f4.sh b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4c_timeseries/run_4c_f4.sh new file mode 100644 index 00000000..f52278c9 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/4_evaluation/4c_timeseries/run_4c_f4.sh @@ -0,0 +1,24 @@ +#!/usr/bin/env bash +# F4 — Dynamic variance direct: 4c reconstructed timeseries plots +# +# Reads routed_leadtime_da_full.parquet + routed_leadtime_openloop_full.parquet +# from the F4 forecast route dir and writes two PNGs there. +# +# Usage: +# bash run_4c_f4.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +ROUTE_DIR="/mnt/disk2/suma_helen_poster/da_results/da_forecast_dynamic_novrugt_seeded_routed" +USGS_CSV="/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +echo "[F4-4c] Reconstructed timeseries from: $ROUTE_DIR" +$TROUTE "$SCRIPT_DIR/plot_lead_time_reconstructed_timeseries.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" + +echo "[F4-4c] Done. Outputs:" +ls "$ROUTE_DIR"/lead_time_reconstructed_timeseries*.png 2>/dev/null || echo " (no PNGs found)" diff --git a/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/README.md b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/README.md new file mode 100644 index 00000000..d9e570ef --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder4_dynamic_variance_direct/README.md @@ -0,0 +1,82 @@ +# Folder 4 — Dynamic Variance Direct + +## R Formula +``` +R(t) = σ²_krig(t) +``` +Raw kriging variance used directly as observation error — no flow-magnitude scaling +and no Vrugt formula. RNG seed = 42. + +## Design Differences vs Other Folders + +| | Folder 1 | Folder 2 | Folder 3 | Folder 4 | +|---|---|---|---|---| +| R formula | Vrugt | Fixed 0.07 | Vrugt | **Direct σ²** | +| RNG seed | hash(cat_id) | hash(cat_id) | 42 | **42** | +| σ²_krig source | static ~4.17 | static ~4.17 | dynamic | **dynamic** | + +> **Controlled comparison:** F4 shares the same RNG seed (42), same obs dataset +> (`spliced_dyn_helene`), and same 21 catchments as F3. The only difference between +> F3 and F4 is the R formula — making **F3 vs F4 the clean controlled comparison** +> for this experimental setup. +> +> F4 cannot be cleanly compared to F1 or F2 because both the seed and the R formula +> differ. See the main [README](../README.md) for full details. + +## Why F4 Performs Worst + +With `R = σ²_krig`, the Kalman gain is: +``` +K = P / (P + σ²_krig) +``` +During Helene, `σ²_krig` is dynamic (~1.6–2.5 at peak in spliced obs, vs ~4.17 +static). Even with lower variance, `σ²_krig` is still much larger than the fixed +R=0.07 used by F2. This keeps K lower than F2 throughout the flood, weakening DA +updates at the most critical timesteps. The absence of the flow-magnitude term means +there is no systematic relationship between R and observed flow — the DA strength is +driven entirely by the spatial kriging uncertainty, which is not well-correlated with +the hydrological need for correction. + +## Key Results + +| Metric | Value | +|---|---| +| Full period KGE | **+0.200** | +| Full period NSE | **+0.428** | +| Helene KGE | **+0.132** | +| Helene NSE | **+0.303** | +| Helene peak (routed) | **667.9 m³/s (35% of USGS 1885.7)** | +| Ensemble p95 at Helene peak | 744 m³/s | +| Ensemble p50 at Helene peak | 552 m³/s | + +## Pipeline Status + +| Step | Status | Server path | +|---|---|---| +| 1. Calibrate | ✅ shared | `1400_sites_helene/da_results_dynamic_novrugt_seeded/{cat}/{cat}_best_params.json` | +| 2a/2b. Perturbation arms | ✅ | `suma_helen_poster/da_results/da_arms_dynamic_novrugt_seeded/` | +| 2c. 18hr forecast cycles | ✅ | `suma_helen_poster/da_results/da_forecast_dynamic_novrugt_seeded/` | +| 2d. 600-member ensemble | ✅ | `suma_helen_poster/da_results/da_crossed_dynamic_novrugt_seeded/` | +| 3. Route analysis | ✅ | `suma_helen_poster/da_results/dynamic_novrugt_seeded_routed/` | +| 3. Route forecast cycles | ✅ | `suma_helen_poster/da_results/da_forecast_dynamic_novrugt_seeded_routed/` | +| 3. Route ensemble | ✅ | `suma_helen_poster/da_results/da_crossed_dynamic_novrugt_seeded_routed/` | +| 4a. Error decay | ✅ | `f4_lead_time_decay_gauge_pooled.png` / `_by_regime.png` | +| 4b. Ensemble vs obs | ✅ | `f4_helene_ensemble_vs_usgs.png` / `f4_helene_ensemble_twopanel.png` | +| 4c. Reconstructed timeseries | ✅ | `f4_reconstructed_timeseries.png` / `_helene.png` | + +All server paths are under `/mnt/disk2/` unless prefixed with `1400_sites_helene/`. + +## Running This Experiment + +```bash +# On the server — all 3 stages +bash 2_assimilation/batch_run_all_f4.sh + +# Or stage by stage +bash 2_assimilation/batch_run_all_f4.sh arms # 2a/2b +bash 2_assimilation/batch_run_all_f4.sh forecast # 2c +bash 2_assimilation/batch_run_all_f4.sh ensemble # 2d + +# 4c timeseries plots +bash 4_evaluation/4c_timeseries/run_4c_f4.sh +``` diff --git a/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/2_assimilation/batch_run_f5_20pct.sh b/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/2_assimilation/batch_run_f5_20pct.sh new file mode 100644 index 00000000..a65a1122 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/2_assimilation/batch_run_f5_20pct.sh @@ -0,0 +1,87 @@ +#!/usr/bin/env bash +# F5 re-kriged variance (20% gauge holdout) — DA assimilation for all 21 catchments. +# +# R(t) = sigma2_rekrig (re-kriged per-hour variogram variance, --no-vrugt-r) +# Obs : /mnt/disk2/1400_sites_helene/catchment_ts_03463300_dynamic_variance_rekrig/ +# Params sourced from F4 DA results (shared calibration). +# +# Usage: +# bash batch_run_f5_20pct.sh > ~/logs/f5_20pct.log 2>&1 + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PYTHON=/home/svyas/miniconda3/envs/troute/bin/python3 +PROD_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/1_distributed_cfe/calibrate_catchment_cfe_da_v2.py + +CATS=( + cat-1016279 cat-1016280 cat-1016281 cat-1016282 cat-1016283 + cat-1016300 + cat-1016301 cat-1016302 cat-1016303 cat-1016304 cat-1016305 + cat-1016306 cat-1016307 cat-1016308 cat-1016309 cat-1016310 + cat-1016311 cat-1016312 cat-1016313 cat-1016314 cat-1016315 +) + +# Re-kriged obs (per-hour variogram refit, 03463300 withheld) +OBS_DIR=/mnt/disk2/1400_sites_helene/catchment_ts_03463300_dynamic_variance_rekrig + +# Params source — F4 results carry best_params for all 21 catchments +PARAM_SRC=/mnt/disk2/1400_sites_helene/da_results_dynamic_novrugt_seeded + +CFE_DIR=/mnt/disk2/suma_helen_poster/cfe_py +CONFIG_FILE=/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json +PARAM_BOUNDS=/mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json +FORCING_DIR=/mnt/disk2/suma_helen_poster/nwm_retro_catchment_forcings +FORCING_DIR1=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings +FORCING_DIR2=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings + +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results/folder5_rekrig_variance_direct + +mkdir -p "$OUT_DIR" +echo "[F5-20pct] Re-kriged variance DA — out: $OUT_DIR" +echo "" + +SKIP=0; DONE=0; FAIL=0 + +for CAT in "${CATS[@]}"; do + echo "===============================" + echo "=== $CAT ===" + echo "===============================" + + CAT_OUT="$OUT_DIR/$CAT" + OUT_FILE="$CAT_OUT/${CAT}_test_results.csv" + + # Stage best_params from F4 param source + PARAMS="$PARAM_SRC/$CAT/${CAT}_best_params.json" + if [ ! -f "$PARAMS" ]; then + echo " WARNING: No best_params for $CAT at $PARAMS — skipping" + SKIP=$((SKIP + 1)); continue + fi + + if [ -f "$OUT_FILE" ]; then + echo " Already exists — skipping" + SKIP=$((SKIP + 1)); continue + fi + + mkdir -p "$CAT_OUT" + cp "$PARAMS" "$CAT_OUT/" + + "$PYTHON" "$PROD_SCRIPT" \ + --cat-id "$CAT" \ + --forcing-dir "$FORCING_DIR" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --param-bounds "$PARAM_BOUNDS" \ + --out-dir "$OUT_DIR" \ + --test-forcing-dir1 "$FORCING_DIR1" \ + --test-forcing-dir2 "$FORCING_DIR2" \ + --enkf-enabled \ + --enkf-members 20 \ + --no-vrugt-r \ + && DONE=$((DONE + 1)) \ + || { echo " FAILED: $CAT"; FAIL=$((FAIL + 1)); } +done + +echo "" +echo "=== F5-20pct done: done=$DONE skipped=$SKIP failed=$FAIL ===" diff --git a/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/2_assimilation/batch_run_f5_analysis_20pct.sh b/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/2_assimilation/batch_run_f5_analysis_20pct.sh new file mode 100644 index 00000000..b3d19438 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/2_assimilation/batch_run_f5_analysis_20pct.sh @@ -0,0 +1,97 @@ +#!/usr/bin/env bash +# F5 re-kriged variance (20% gauge holdout) — analysis pipeline. +# +# Stage A: perturbation arms (run_perturbation_da_on.py, --direct-variance, seed=42) +# Stage B: lead time forecast (run_lead_time_forecast_sweep.py, --no-vrugt-r, seed=42) +# +# Run AFTER batch_run_f5_20pct.sh (production DA) completes. +# +# Usage: +# bash batch_run_f5_analysis_20pct.sh arms +# bash batch_run_f5_analysis_20pct.sh forecast +# bash batch_run_f5_analysis_20pct.sh all + +set -euo pipefail + +STAGE="${1:-all}" + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR=/home/svyas + +DA_SRC=/mnt/disk2/1400_sites_helene/da_results_dynamic_novrugt_seeded +OBS_DIR=/mnt/disk2/1400_sites_helene/catchment_ts_03463300_dynamic_variance_rekrig +CFE_DIR=/mnt/disk2/suma_helen_poster/cfe_py +CONFIG_FILE=/mnt/disk2/suma_helen_poster/run_gpu/cat_03463300_bmi_config_cfe.json +PARAM_BOUNDS=/mnt/disk2/suma_helen_poster/run_gpu/CFE_parameter_bounds.json +FORCING_DIR=/mnt/disk2/suma_helen_poster/nwm_retro_catchment_forcings +TEST_FORCING_DIR1=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2023_2024_feb/forcings +TEST_FORCING_DIR2=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_nwmoperational/03463300/2024_feb_2025_sep/forcings + +OUT_ARMS=/mnt/disk2/suma_helen_poster/da_results/da_arms_f5_rekrig_20pct +OUT_FORECAST=/mnt/disk2/suma_helen_poster/da_results/da_forecast_f5_rekrig_20pct + +LOG_DIR=$HOME/f5_20pct_logs +mkdir -p "$LOG_DIR" + +CATS=$(ls -d "$DA_SRC"/cat-* 2>/dev/null | xargs -I{} basename {}) + +# ── Stage A: perturbation arms ──────────────────────────────────────────────── +run_arms() { + echo "[F5-20pct] Stage A: perturbation arms (re-kriged variance direct, seed=42)" + for cat_id in $CATS; do + params="$DA_SRC/$cat_id/${cat_id}_best_params.json" + [ -f "$params" ] || { echo " [skip] $cat_id — no params"; continue; } + mkdir -p "$OUT_ARMS/$cat_id" + cp "$params" "$OUT_ARMS/$cat_id/" + log="$LOG_DIR/arms_${cat_id}.log" + echo " [run] $cat_id" + $TROUTE "$SCRIPT_DIR/run_perturbation_da_on.py" \ + --cat-id "$cat_id" \ + --forcing-dir "$FORCING_DIR" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --param-bounds "$PARAM_BOUNDS" \ + --out-dir "$OUT_ARMS" \ + --test-forcing-dir1 "$TEST_FORCING_DIR1" \ + --test-forcing-dir2 "$TEST_FORCING_DIR2" \ + --direct-variance \ + --rng-seed 42 \ + > "$log" 2>&1 && echo " [ok ] $cat_id" || echo " [FAIL] $cat_id — see $log" + done +} + +# ── Stage B: lead time forecast sweep ──────────────────────────────────────── +run_forecast() { + echo "[F5-20pct] Stage B: lead time forecast sweep (re-kriged variance, seed=42)" + for cat_id in $CATS; do + params="$DA_SRC/$cat_id/${cat_id}_best_params.json" + [ -f "$params" ] || { echo " [skip] $cat_id — no params"; continue; } + mkdir -p "$OUT_FORECAST/$cat_id" + cp "$params" "$OUT_FORECAST/$cat_id/" + log="$LOG_DIR/forecast_${cat_id}.log" + echo " [run] $cat_id" + $TROUTE "$SCRIPT_DIR/run_lead_time_forecast_sweep.py" \ + --cat-id "$cat_id" \ + --forcing-dir "$FORCING_DIR" \ + --obs-dir "$OBS_DIR" \ + --cfe-dir "$CFE_DIR" \ + --config-file "$CONFIG_FILE" \ + --param-bounds "$PARAM_BOUNDS" \ + --out-dir "$OUT_FORECAST" \ + --test-forcing-dir1 "$TEST_FORCING_DIR1" \ + --test-forcing-dir2 "$TEST_FORCING_DIR2" \ + --no-vrugt-r \ + --rng-seed 42 \ + > "$log" 2>&1 && echo " [ok ] $cat_id" || echo " [FAIL] $cat_id — see $log" + done +} + +case "$STAGE" in + arms) run_arms ;; + forecast) run_forecast ;; + all) run_arms; run_forecast ;; + *) echo "Usage: $0 [arms|forecast|all]"; exit 1 ;; +esac + +echo "[F5-20pct] Analysis done." diff --git a/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/2_assimilation/plot_da_perturbation_arms.py b/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/2_assimilation/plot_da_perturbation_arms.py new file mode 100644 index 00000000..283d5f9d --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/2_assimilation/plot_da_perturbation_arms.py @@ -0,0 +1,355 @@ +""" +plot_da_perturbation_arms.py — 2a/2b: ensemble spread WITH DA on. + +Loads the two arm CSVs produced by run_perturbation_da_on.py: + //_da_forcing_arm.csv (30 members) + //_da_hydro_arm.csv (20 members) + +Plots: + Panel 2a — Forcing arm: 30-member spaghetti over Helene window + (spread = met forcing uncertainty with DA-corrected initial states) + Panel 2b — Hydro-state arm: 20-member spaghetti over Helene window + (spread = initial state uncertainty with deterministic forcing) + Optional: --ol-csv overlays the open-loop grand median as a + thick dashed gray line for comparison. + Panel 2c — Comparison: median ± spread envelope, both arms + USGS obs + +All trajectories projected to valid_time = issue_time + lead_hour hours. +Colored by initialization date (Sep 24-30). USGS obs in black. + +Outputs: + //_2a_forcing_arm_helene.png + //_2b_hydro_arm_helene.png + //_2ab_arms_comparison.png + +Run on server (troute env): + python3 plot_da_perturbation_arms.py \\ + --cat-id cat-1016300 \\ + --arm-dir /mnt/disk2/suma_helen_poster/da_results/da_arms_f5_rekrig_20pct \\ + --ol-csv /mnt/disk2/suma_helen_poster/da_results/da_forecast_f5_rekrig_20pct/routed/routed_leadtime_openloop_full.parquet +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates +import matplotlib.patches as mpatches + +DEFAULT_ARM_DIR = "/mnt/disk2/suma_helen_poster/da_results/da_arms_f5_rekrig_20pct" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" +DEFAULT_CAT_ID = "cat-1016300" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +PLOT_START = pd.Timestamp("2024-09-24 00:00:00") +PLOT_END = pd.Timestamp("2024-09-30 06:00:00") +HELENE_PEAK_START = pd.Timestamp("2024-09-26 12:00:00") +HELENE_PEAK_END = pd.Timestamp("2024-09-28 00:00:00") + +DATE_COLORS = { + "2024-09-24": "#1f77b4", + "2024-09-25": "#ff7f0e", + "2024-09-26": "#2ca02c", + "2024-09-27": "#d62728", + "2024-09-28": "#9467bd", + "2024-09-29": "#8c564b", + "2024-09-30": "#e377c2", +} + + +def load_arm(path): + df = pd.read_csv(path) + df["issue_time"] = pd.to_datetime(df["issue_time"]) + df["valid_time"] = df["issue_time"] + pd.to_timedelta(df["lead_hour"], unit="h") + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + return df, member_cols + + +def load_openloop(path): + """Load open-loop lead-time file (CSV or Parquet) and return grand median + indexed by valid_time. + + Accepts: + - routed_leadtime_openloop_full.parquet (T-route output, m³/s, recommended) + - cat-*_lead_time_forecasts_openloop.csv (unrouted CFE mm/h, single catchment) + + Returns a pd.Series (valid_time → median q in m³/s) clipped to plot window. + """ + if path.endswith(".parquet"): + df = pd.read_parquet(path) + else: + df = pd.read_csv(path) + + df["issue_time"] = pd.to_datetime(df["issue_time"]) + if "valid_time" in df.columns: + df["valid_time"] = pd.to_datetime(df["valid_time"]) + elif "lead_hour" in df.columns: + df["valid_time"] = (df["issue_time"] + + pd.to_timedelta(df["lead_hour"], unit="h")) + else: + raise ValueError("Open-loop file must have valid_time or lead_hour column") + + member_cols = sorted([c for c in df.columns if c.startswith("member_")]) + mask = (df["valid_time"] >= PLOT_START) & (df["valid_time"] <= PLOT_END) + df = df[mask].copy() + if df.empty: + return pd.Series(dtype=float) + + vals = df[member_cols].to_numpy(dtype=float) + # Only convert if values are clearly in mm/h (unrouted CFE output). + # Routed parquet is already in m³/s — do not convert. + if path.endswith(".csv") and np.nanmedian(vals[vals > 0]) < 5: + vals = vals * MM_H_TO_M3_S + + df["q_grand_median"] = np.nanmedian(vals, axis=1) + series = (df.groupby("valid_time")["q_grand_median"] + .median() + .sort_index()) + return series + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ("datetime", "date", "time", "timestamp")) + q_col = next(c for c in df.columns + if "q" in c.lower() or "flow" in c.lower() or "discharge" in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if "mm" in q_col.lower(): + series = series * MM_H_TO_M3_S + return series + + +def _add_helene_band(ax): + ax.axvspan(HELENE_PEAK_START, HELENE_PEAK_END, + color="salmon", alpha=0.12, zorder=0) + ax.text(HELENE_PEAK_START + pd.Timedelta(hours=6), 1.0, "Helene peak", + transform=ax.get_xaxis_transform(), + fontsize=9, color="firebrick", ha="left", va="top") + + +def _format_xaxis(ax): + ax.set_xlim(PLOT_START, PLOT_END) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=1)) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.xaxis.set_minor_locator(mdates.HourLocator(interval=6)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=30, ha="right") + ax.grid(True, alpha=0.2, lw=0.4) + + +def plot_spaghetti_arm(ax, df, member_cols, arm_color, obs_series, + title, label_stem, lw_thin=0.7, alpha_thin=0.35): + _add_helene_band(ax) + + issue_times = sorted(df["issue_time"].unique()) + all_means = [] + + for t0 in issue_times: + date_str = str(pd.Timestamp(t0).date()) + color = DATE_COLORS.get(date_str, arm_color) + sub = df[df["issue_time"] == t0].sort_values("valid_time") + mask = (sub["valid_time"] >= PLOT_START) & (sub["valid_time"] <= PLOT_END) + sub_w = sub[mask] + if sub_w.empty: + continue + vt_w = sub_w["valid_time"].values + + mem_vals = sub_w[member_cols].to_numpy(dtype=float) + if "mm" not in label_stem.lower(): + mem_vals = mem_vals * MM_H_TO_M3_S + + ax.fill_between(vt_w, + np.nanmin(mem_vals, axis=1), + np.nanmax(mem_vals, axis=1), + color=color, alpha=0.06, zorder=2) + ax.plot(vt_w, np.nanmedian(mem_vals, axis=1), + color=color, lw=lw_thin, alpha=alpha_thin + 0.1, zorder=3) + + all_means.append( + pd.Series(np.nanmedian(mem_vals, axis=1), index=vt_w)) + + if all_means: + full_idx = pd.date_range(PLOT_START, PLOT_END, freq="1h") + stacked = pd.concat(all_means, axis=1).reindex(full_idx) + grand_mean = stacked.mean(axis=1) + ax.plot(grand_mean.index, grand_mean.values, + color=arm_color, lw=2.4, alpha=0.95, zorder=5, + label=f"{label_stem} — mean across all forecasts") + + obs_w = obs_series.loc[PLOT_START:PLOT_END] + ax.plot(obs_w.index, obs_w.values, + color="black", lw=1.8, zorder=6, label="USGS obs") + + ax.set_ylabel("Discharge (m³/s)", fontsize=10) + ax.set_title(title, fontsize=11) + _format_xaxis(ax) + + +def compute_envelope(df, member_cols): + records = [] + for t0, grp in df.groupby("issue_time"): + mask = (grp["valid_time"] >= PLOT_START) & (grp["valid_time"] <= PLOT_END) + sub = grp[mask] + if sub.empty: + continue + vals = sub[member_cols].to_numpy(dtype=float) * MM_H_TO_M3_S + for i, row in enumerate(sub.itertuples()): + records.append({ + "valid_time": row.valid_time, + "q_min": np.nanmin(vals[i]), + "q_med": np.nanmedian(vals[i]), + "q_max": np.nanmax(vals[i]), + }) + if not records: + return pd.DataFrame(columns=["valid_time", "q_min", "q_med", "q_max"]) + + env_df = pd.DataFrame(records) + env_df = (env_df.groupby("valid_time") + .agg(q_min=("q_min", "min"), + q_med=("q_med", "mean"), + q_max=("q_max", "max")) + .reset_index() + .sort_values("valid_time")) + return env_df + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--arm-dir", default=DEFAULT_ARM_DIR) + parser.add_argument("--cat-id", default=DEFAULT_CAT_ID) + parser.add_argument("--usgs-csv", default=DEFAULT_USGS_CSV) + parser.add_argument("--out-dir", default=None) + parser.add_argument( + "--ol-csv", default=None, + help=( + "Path to open-loop lead-time forecast (CSV or Parquet). " + "When provided, the grand median is overlaid on panel 2b " + "as a thick dashed black line." + ), + ) + args = parser.parse_args() + + cat_dir = os.path.join(args.arm_dir, args.cat_id) + out_dir = args.out_dir or cat_dir + os.makedirs(out_dir, exist_ok=True) + + forcing_path = os.path.join(cat_dir, f"{args.cat_id}_da_forcing_arm.csv") + hydro_path = os.path.join(cat_dir, f"{args.cat_id}_da_hydro_arm.csv") + + print(f"Loading arm CSVs for {args.cat_id}...") + df_fa, fa_cols = load_arm(forcing_path) + df_ha, ha_cols = load_arm(hydro_path) + obs = load_usgs(args.usgs_csv) + print(f" Forcing arm: {df_fa['issue_time'].nunique()} issue times, " + f"{len(fa_cols)} members") + print(f" Hydro arm: {df_ha['issue_time'].nunique()} issue times, " + f"{len(ha_cols)} members") + + # ── Figure 2a: Forcing arm ──────────────────────────────────────────────── + fig, ax = plt.subplots(figsize=(17, 6)) + plot_spaghetti_arm( + ax, df_fa, fa_cols, + arm_color="tab:blue", obs_series=obs, + title=(f"2a — Forcing arm (DA on, {len(fa_cols)} members): " + "met forcing uncertainty | Sep 24-30 2024\n" + f"{args.cat_id} | Shaded = member min-max | " + "Line = median per init time | Thick = grand mean"), + label_stem="Forcing arm", + ) + patches = [mpatches.Patch(color=c, label=f"Init {d}") + for d, c in DATE_COLORS.items()] + handles, _ = ax.get_legend_handles_labels() + ax.legend(handles=handles + patches, + fontsize=8, loc="upper left", frameon=True, framealpha=0.9, ncol=2) + plt.tight_layout() + out_a = os.path.join(out_dir, f"{args.cat_id}_2a_forcing_arm_helene.png") + plt.savefig(out_a, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_a}") + + # ── Figure 2b: Hydro-state arm ──────────────────────────────────────────── + fig, ax = plt.subplots(figsize=(17, 6)) + plot_spaghetti_arm( + ax, df_ha, ha_cols, + arm_color="tab:green", obs_series=obs, + title=(f"2b — Hydro-state arm (DA on, {len(ha_cols)} members): " + "initial state uncertainty | Sep 24-30 2024\n" + f"{args.cat_id} | Shaded = member min-max | " + "Line = median per init time | Thick = grand mean"), + label_stem="Hydro-state arm", + ) + if args.ol_csv: + print(f"Loading open-loop: {args.ol_csv}") + ol_series = load_openloop(args.ol_csv) + if not ol_series.empty: + ax.plot( + ol_series.index, ol_series.values, + color="black", lw=2.5, ls="--", alpha=0.95, zorder=7, + label="Open loop (no DA) — grand median", + ) + else: + print(" Warning: open-loop yielded no data in plot window.") + + patches = [mpatches.Patch(color=c, label=f"Init {d}") + for d, c in DATE_COLORS.items()] + handles, _ = ax.get_legend_handles_labels() + ax.legend(handles=handles + patches, + fontsize=8, loc="upper left", frameon=True, framealpha=0.9, ncol=2) + plt.tight_layout() + out_b = os.path.join(out_dir, f"{args.cat_id}_2b_hydro_arm_helene.png") + plt.savefig(out_b, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_b}") + + # ── Figure 2ab: Comparison ──────────────────────────────────────────────── + print("Computing aggregated envelopes for comparison plot...") + env_fa = compute_envelope(df_fa, fa_cols) + env_ha = compute_envelope(df_ha, ha_cols) + + fig, ax = plt.subplots(figsize=(17, 7)) + _add_helene_band(ax) + + if not env_fa.empty: + ax.fill_between(env_fa["valid_time"], env_fa["q_min"], env_fa["q_max"], + color="tab:blue", alpha=0.18, zorder=2, + label=f"Forcing arm spread (N={len(fa_cols)} members)") + ax.plot(env_fa["valid_time"], env_fa["q_med"], + color="tab:blue", lw=2.0, zorder=4, + label="Forcing arm — grand median") + + if not env_ha.empty: + ax.fill_between(env_ha["valid_time"], env_ha["q_min"], env_ha["q_max"], + color="tab:green", alpha=0.18, zorder=2, + label=f"Hydro-state arm spread (N={len(ha_cols)} members)") + ax.plot(env_ha["valid_time"], env_ha["q_med"], + color="tab:green", lw=2.0, zorder=4, + label="Hydro-state arm — grand median") + + obs_w = obs.loc[PLOT_START:PLOT_END] + ax.plot(obs_w.index, obs_w.values, + color="black", lw=1.8, zorder=6, label="USGS obs") + + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_xlabel("Date", fontsize=11) + ax.set_title( + f"2a vs 2b — Forcing arm (blue) vs Hydro-state arm (green) | " + f"DA on | {args.cat_id}\n" + "Shaded = full member spread (min-max aggregated across all issue times). " + "Lines = grand median.", + fontsize=11, + ) + ax.legend(fontsize=9, loc="upper left", frameon=True, framealpha=0.9) + _format_xaxis(ax) + plt.tight_layout() + out_c = os.path.join(out_dir, f"{args.cat_id}_2ab_arms_comparison.png") + plt.savefig(out_c, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_c}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/2_assimilation/run_lead_time_forecast_sweep.py b/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/2_assimilation/run_lead_time_forecast_sweep.py new file mode 100644 index 00000000..19f3b4d5 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/2_assimilation/run_lead_time_forecast_sweep.py @@ -0,0 +1,508 @@ +""" +Forecast lead-time evaluation. + +For each issue-time t0 in a sampling schedule across the test period, fork the +ensemble at t0 and free-run 18 hours of forecast. Two parallel trajectories are +maintained throughout the test period: + + da: Production setup. DA on, forcing perturbed, process noise on. + openloop: No DA from t=0. Forcing perturbed, process noise on. + +At each issue time t0, the state of each scenario's ensemble is copied into a +forecast ensemble that free-runs 18 hours with: + - DA off (no observations consumed during the forecast window) + - process noise off (no obs to collapse toward, so no anti-collapse needed) + - forcing perturbed (lognormal precip, Gaussian PET — proxy for forecast + precip uncertainty in lieu of actual HEFS forecasts) + +Issue-time schedule: + - Base cadence: every --base-step-h hours across the full test period + (default 6h → ~1600 issue times for the year) + - Densified to hourly across --dense-start..--dense-end if provided + (default: 2024-09-24 → 2024-09-28, the Helene window) + +Outputs (stacked across all issue times): + //_lead_time_forecasts_da.csv + //_lead_time_forecasts_openloop.csv + Columns: issue_time, lead_hour, valid_time, member_00 .. member_19 + (units: mm/h) + +Post-processing pipeline (separate scripts, run after this): + route_lead_time_forecasts.py — pushes each (t0, lead_hour) forecast through + T-route to the gauge + plot_lead_time_decay.py — error-vs-lead-time curve, DA vs open-loop + +Imports EnKFAssimilator from calibrate_catchment_cfe_da_v2.py so the DA math +matches production exactly. +""" + +import argparse +import os +import sys +import json +import importlib.util +import numpy as np +import pandas as pd +from pathlib import Path + +ALBEDO = 0.20 +ALPHA_PT = 1.26 + +SPINUP_START = "2023-02-01 00:00:00" +SPINUP_END = "2023-09-30 23:00:00" +TEST_START = "2023-10-01 00:00:00" +TEST_END = "2024-10-31 23:00:00" + +# Forecast lead time (hours after each issue time) +FORECAST_LEAD_HOURS = 18 + +# Default densification window — Hurricane Helene +DEFAULT_DENSE_START = "2024-09-24 00:00:00" +DEFAULT_DENSE_END = "2024-09-28 23:00:00" + + +def priestley_taylor_pet(srad_wm2, T_kelvin, alpha=ALPHA_PT): + T = T_kelvin - 273.15 + delta = 4098 * (0.6108 * np.exp(17.27 * T / (T + 237.3))) / (T + 237.3) ** 2 + gamma = 0.0638 + Rn_mj = np.maximum((1.0 - ALBEDO) * srad_wm2, 0.0) * 0.0036 + lam = 2.501 - 0.002361 * T + pet = alpha * (delta / (delta + gamma)) * Rn_mj / lam + return np.maximum(pet, 0.0) + + +def load_test_forcing(test_forcing_file): + df = pd.read_csv(test_forcing_file) + df['date'] = pd.to_datetime(df['time']).dt.strftime('%Y-%m-%d %H:%M:%S') + df['total_precipitation'] = df['APCP_surface'] * 3600.0 + df['potential_evaporation'] = priestley_taylor_pet( + df['DSWRF_surface'].values, df['TMP_2maboveground'].values + ) + return df + + +def import_enkf_class(script_path): + """Load EnKFAssimilator from the production script without triggering its main().""" + spec = importlib.util.spec_from_file_location("prod_v2", script_path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod.EnKFAssimilator + + +# ---------------- State snapshot / restore ----------------------------------- +# +# The EnKF only ever touches these 4 states, so they're the only states that +# diverge between members. Copy these from src → dst to fork the ensemble. + +def snapshot_states(models): + return [ + { + "soil_m": float(m.soil_reservoir["storage_m"]), + "gw_m": float(m.gw_reservoir["storage_m"]), + "nash0_m": float(m.nash_storage[0]), + "nash1_m": float(m.nash_storage[1]), + } + for m in models + ] + + +def restore_states(models, snapshot): + for m, s in zip(models, snapshot): + m.soil_reservoir["storage_m"] = s["soil_m"] + m.gw_reservoir["storage_m"] = s["gw_m"] + m.nash_storage[0] = s["nash0_m"] + m.nash_storage[1] = s["nash1_m"] + + +# ---------------- Issue-time schedule ---------------------------------------- + +def build_issue_time_schedule(dates_list, base_step_h, dense_start, dense_end): + """Return a sorted list of issue-time strings sampled from dates_list. + + Every base_step_h hours across the full test period, plus every hour inside + [dense_start, dense_end] if those are provided. + """ + dates_dt = pd.to_datetime(dates_list) + selected = set() + + # Base cadence (every base_step_h hours from the first date) + base_mask = (np.arange(len(dates_dt)) % base_step_h == 0) + for d in dates_dt[base_mask]: + selected.add(d.strftime('%Y-%m-%d %H:%M:%S')) + + # Densified hourly window + if dense_start is not None and dense_end is not None: + ds = pd.Timestamp(dense_start) + de = pd.Timestamp(dense_end) + dense_mask = (dates_dt >= ds) & (dates_dt <= de) + for d in dates_dt[dense_mask]: + selected.add(d.strftime('%Y-%m-%d %H:%M:%S')) + + return sorted(selected) + + +# ---------------- CFE helpers ------------------------------------------------ + +def build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, enkf, + apply_init_perturbation=True, member_zero_clean=True): + """Build N CFE BMI instances with optional init-state perturbation.""" + def custom_load_forcing(self_cfe): + df = load_test_forcing(test_forcing_file) + self_cfe.forcing_data = df.rename(columns={"date": "time"}) + + models = [] + for i in range(N): + m = bmi_cfe.BMI_CFE(cfg_file=tmp_cfg) + m.load_forcing_file = custom_load_forcing.__get__(m) + m.initialize() + if apply_init_perturbation and not (member_zero_clean and i == 0): + sm_max = m.soil_reservoir["storage_max_m"] + gw_max = m.gw_reservoir["storage_max_m"] + sm0 = m.soil_reservoir["storage_m"] + gw0 = m.gw_reservoir["storage_m"] + n00 = float(m.nash_storage[0]) + n10 = float(m.nash_storage[1]) + r = enkf.init_state_perturb_frac + m.soil_reservoir["storage_m"] = float(np.clip( + sm0 * (1 + r * enkf.rng.standard_normal()), 1e-6, sm_max)) + m.gw_reservoir["storage_m"] = float(np.clip( + gw0 * (1 + r * enkf.rng.standard_normal()), 1e-6, gw_max)) + jitter = max(sm_max * 1e-6, 1e-9) + m.nash_storage[0] = max(n00 + jitter * enkf.rng.standard_normal(), 0.0) + m.nash_storage[1] = max(n10 + jitter * enkf.rng.standard_normal(), 0.0) + models.append(m) + return models + + +def step_ensemble(models, p_arr, e_arr): + """Advance each member one hour with member-specific (P, PET). Returns Q (mm/h).""" + q = np.empty(len(models), dtype=float) + for i, m in enumerate(models): + m.set_value('atmosphere_water__time_integral_of_precipitation_mass_flux', + float(p_arr[i]) / 1000.0) + m.set_value('water_potential_evaporation_flux', + float(e_arr[i]) / 1000.0 / 3600.0) + m.update() + q[i] = m.get_value('land_surface_water__runoff_depth') * 1000.0 # m/h → mm/h + return q + + +def do_sanity_check(prod_models, fcst_models, dates_list, forcing_by_date, h, N): + """One-shot verification that the 4-state snapshot is sufficient to reproduce + next-hour forecast Q. + + At main-loop hour h (after prod's hour-h step has completed including DA + + process noise), snapshot prod, restore into fcst, then step BOTH ensembles + one hour using identical unperturbed forcing. If snapshot/restore captures + everything that matters, prod's and fcst's hour-(h+1) Q should match member- + by-member to machine precision. + + Note: this mutates prod's state (h+1 step with non-perturbed forcing and no + DA), so the caller should exit the script right after. + """ + if h + 1 >= len(dates_list): + print("[sanity-check] not enough horizon for h+1 step — skipping.") + return False + snap = snapshot_states(prod_models) + restore_states(fcst_models, snap) + + next_date = dates_list[h + 1] + p_next, e_next = forcing_by_date[next_date] + p_uniform = np.full(N, p_next, dtype=float) + e_uniform = np.full(N, e_next, dtype=float) + + q_prod = step_ensemble(prod_models, p_uniform, e_uniform) + q_fcst = step_ensemble(fcst_models, p_uniform, e_uniform) + + diff = q_prod - q_fcst + max_abs = float(np.max(np.abs(diff))) + print("[sanity-check] forecast Q at h+1 with identical unperturbed forcing:") + print(f" prod first 5 members: {q_prod[:5]}") + print(f" fcst first 5 members: {q_fcst[:5]}") + print(f" max |prod - fcst|: {max_abs:.3e} mm/h") + if max_abs < 1e-9: + print(" OK: snapshot/restore preserves state to machine precision.") + return True + print(" WARN: nonzero divergence — likely a state not in the 4-state snapshot.") + return False + + +def run_forecast(fcst_models, enkf_fcst, dates_list, forcing_by_date, + t0_idx, n_lead): + """Run an n_lead-hour free-forecast starting from the current state of fcst_models. + + Forcing is perturbed per member (lognormal precip + Gaussian PET, same as + production). No DA. No process noise. + + Returns: (lead_hours, valid_times, q_matrix shape (n_lead, N)). + """ + N = len(fcst_models) + q_matrix = np.full((n_lead, N), np.nan, dtype=float) + lead_hours = [] + valid_times = [] + for lead in range(1, n_lead + 1): + idx = t0_idx + lead + if idx >= len(dates_list): + break + valid_date = dates_list[idx] + p, e = forcing_by_date[valid_date] + p_arr, e_arr = enkf_fcst.perturb_forcing(p, e) + q = step_ensemble(fcst_models, p_arr, e_arr) + q_matrix[lead - 1, :] = q + lead_hours.append(lead) + valid_times.append(valid_date) + return lead_hours, valid_times, q_matrix + + +# ---------------- Main run --------------------------------------------------- + +def run(args, bmi_cfe, EnKFAssimilator): + cat_id = args.cat_id + out_dir = Path(args.out_dir) / cat_id + out_dir.mkdir(parents=True, exist_ok=True) + + # Combined test forcing (same idiom as the per-member script) + f1 = os.path.join(args.test_forcing_dir1, f'{cat_id}.csv') + f2 = os.path.join(args.test_forcing_dir2, f'{cat_id}.csv') + df1 = pd.read_csv(f1) + df2 = pd.read_csv(f2) + combined = pd.concat([df1, df2], ignore_index=True) + combined = combined.drop_duplicates(subset='time').sort_values('time') + test_forcing_file = str(out_dir / f'{cat_id}_nwm_operational_combined.csv') + combined.to_csv(test_forcing_file, index=False) + + # Pre-staged calibrated params + best_params_file = out_dir / f'{cat_id}_best_params.json' + if not best_params_file.exists(): + print(f"Missing {best_params_file}. Pre-stage from Run 3.") + return + with open(best_params_file) as f: + best = json.load(f)['best_parameters'] + + # Build temp CFE config with calibrated params + with open(args.config_file) as f: + cfg = json.load(f) + cfg['forcing_file'] = test_forcing_file + cfg['soil_params']['bb'] = best['bb'] + cfg['soil_params']['smcmax'] = best['smcmax'] + cfg['soil_params']['satdk'] = best['satdk'] + cfg['slop'] = best['slop'] + cfg['max_gw_storage'] = best['max_gw_storage'] + cfg['expon'] = best['expon'] + cfg['Cgw'] = best['Cgw'] + cfg['K_lf'] = best['K_lf'] + cfg['K_nash'] = best['K_nash'] + cfg['partition_scheme'] = "Schaake" if best['scheme'] <= 0.5 else "Xinanjiang" + tmp_cfg = str(out_dir / f'{cat_id}_bmi_config_temp_leadtime.json') + with open(tmp_cfg, 'w') as f: + json.dump(cfg, f) + + # Three EnKF instances: + # enkf_da — runs the DA-on trajectory; consumes obs each hour + # enkf_ol — runs the open-loop trajectory; never calls update_states + # enkf_fcst — used only inside the 18-hour free-runs (perturb_forcing only) + # All three share defaults with production. Distinct RNG seeds keep them + # independent so the openloop and forecast spreads are not coupled to DA. + obs_file = os.path.join(args.obs_dir, f'{cat_id}.csv') + seed_base = args.rng_seed if args.rng_seed is not None else 0 + enkf_da = EnKFAssimilator( + n_members=args.enkf_members, obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, use_vrugt_r=(not args.no_vrugt_r), + vrugt_alpha=args.vrugt_alpha, vrugt_scale=args.vrugt_scale, + rng_seed=seed_base if args.rng_seed is not None else None, + ) + enkf_ol = EnKFAssimilator( + n_members=args.enkf_members, obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, use_vrugt_r=(not args.no_vrugt_r), + vrugt_alpha=args.vrugt_alpha, vrugt_scale=args.vrugt_scale, + rng_seed=(seed_base + 1) if args.rng_seed is not None else None, + ) + enkf_fcst = EnKFAssimilator( + n_members=args.enkf_members, obs_error_std=args.enkf_obs_error_std, + obs_file=obs_file, use_vrugt_r=(not args.no_vrugt_r), + vrugt_alpha=args.vrugt_alpha, vrugt_scale=args.vrugt_scale, + rng_seed=(seed_base + 2) if args.rng_seed is not None else None, + ) + # Optional R override — replace every per-hour obs variance with a constant. + # Matches DualEarth/new_EnKF.py (R=0.07). Set via --hardcoded-r on the CLI. + # Effect: at storm peak ≈ existing Vrugt R, at low flow ≈ 15-17× larger → + # near-zero gain at low flow where kriging obs is noisy. + if args.hardcoded_r is not None: + for enkf in (enkf_da, enkf_ol, enkf_fcst): + for date in enkf.obs_var_dict: + enkf.obs_var_dict[date] = args.hardcoded_r + print(f"[lead-time] R hardcoded to {args.hardcoded_r} mm^2/h^2 " + f"(overrides Vrugt + kriging variance formula)") + + N = enkf_da.n_members + print(f"[lead-time] {cat_id} | N={N} | DA on/off both run | " + f"forecast lead = {FORECAST_LEAD_HOURS}h") + + # Build three ensembles (init-perturbed except the forecast ensemble, whose + # state will be overwritten at each issue time before each free-run) + prod_models = build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, enkf_da) + openloop_models = build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, enkf_ol) + fcst_models = build_models(N, bmi_cfe, tmp_cfg, test_forcing_file, + enkf_fcst, apply_init_perturbation=False) + + # ----- Spinup with perturbed forcing on both trajectories ----- + df = load_test_forcing(test_forcing_file) + sp_mask = (df['date'] >= SPINUP_START) & (df['date'] <= SPINUP_END) + df_sp = df[sp_mask] + print(f"[lead-time] spinup: {len(df_sp)} hours") + for p, e in zip(df_sp['total_precipitation'], df_sp['potential_evaporation']): + p_da, e_da = enkf_da.perturb_forcing(p, e) + step_ensemble(prod_models, p_da, e_da) + enkf_da.add_process_noise(prod_models) + p_ol, e_ol = enkf_ol.perturb_forcing(p, e) + step_ensemble(openloop_models, p_ol, e_ol) + enkf_ol.add_process_noise(openloop_models) + + # ----- Test period: step both, fork forecasts at issue times ----- + t_mask = (df['date'] >= TEST_START) & (df['date'] <= TEST_END) + df_test = df[t_mask].reset_index(drop=True) + dates_list = list(df_test['date'].values) + forcing_by_date = dict(zip( + df_test['date'].values, + zip(df_test['total_precipitation'].values, + df_test['potential_evaporation'].values))) + + issue_times = build_issue_time_schedule( + dates_list, args.base_step_h, + args.dense_start, args.dense_end) + issue_set = set(issue_times) + print(f"[lead-time] test period: {len(dates_list)} hours | " + f"{len(issue_times)} issue times " + f"(base={args.base_step_h}h, dense={args.dense_start or 'none'}" + f"..{args.dense_end or 'none'})") + + # Pre-build the date→index map so forecast windows are fast to look up + date_to_idx = {d: i for i, d in enumerate(dates_list)} + + da_rows = [] # (issue_time, lead_hour, valid_time, *member_values) + ol_rows = [] + sanity_done = False + + for h, current_date in enumerate(dates_list): + p, e = forcing_by_date[current_date] + + # --- DA trajectory step --- + p_da, e_da = enkf_da.perturb_forcing(p, e) + q_da = step_ensemble(prod_models, p_da, e_da) + enkf_da.update_states(prod_models, current_date, q_da) + enkf_da.add_process_noise(prod_models) + + # --- Open-loop trajectory step (no DA) --- + p_ol, e_ol = enkf_ol.perturb_forcing(p, e) + q_ol = step_ensemble(openloop_models, p_ol, e_ol) + enkf_ol.add_process_noise(openloop_models) + + # --- One-shot sanity check at the first scheduled issue time --- + if args.sanity_check and not sanity_done and current_date in issue_set: + ok = do_sanity_check( + prod_models, fcst_models, dates_list, forcing_by_date, h, N) + sanity_done = True + print(f"[sanity-check] exiting (re-run without --sanity-check for " + f"the full sweep). Result: {'OK' if ok else 'FAIL'}") + for m in prod_models + openloop_models + fcst_models: + m.finalize() + sys.exit(0 if ok else 1) + + # --- Fork forecasts at scheduled issue times --- + if current_date in issue_set: + for scenario_label, src_models, rows_acc in [ + ('da', prod_models, da_rows), + ('openloop', openloop_models, ol_rows), + ]: + snap = snapshot_states(src_models) + restore_states(fcst_models, snap) + leads, valids, qm = run_forecast( + fcst_models, enkf_fcst, + dates_list, forcing_by_date, + t0_idx=h, n_lead=FORECAST_LEAD_HOURS) + for k, (lead, vt) in enumerate(zip(leads, valids)): + rows_acc.append((current_date, lead, vt, *qm[k, :])) + + for m in prod_models + openloop_models + fcst_models: + m.finalize() + + # ----- Save forecast CSVs ----- + cols = ['issue_time', 'lead_hour', 'valid_time'] + [f'member_{i:02d}' for i in range(N)] + da_path = out_dir / f'{cat_id}_lead_time_forecasts_da.csv' + ol_path = out_dir / f'{cat_id}_lead_time_forecasts_openloop.csv' + pd.DataFrame(da_rows, columns=cols).to_csv(da_path, index=False) + pd.DataFrame(ol_rows, columns=cols).to_csv(ol_path, index=False) + print(f"[lead-time] saved {da_path}") + print(f"[lead-time] saved {ol_path}") + + # Issue-time schedule (small, useful for downstream scripts) + sched_path = out_dir / f'{cat_id}_lead_time_issue_times.csv' + pd.DataFrame({'issue_time': issue_times}).to_csv(sched_path, index=False) + print(f"[lead-time] saved {sched_path}") + + print(f"[lead-time] DA assimilator: updates={enkf_da.n_updates} | " + f"avg Pyy={enkf_da.avg_pyy:.6f} | " + f"mass lost={enkf_da.total_overflow_lost_mm:.3f} mm") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--cat-id', required=True) + parser.add_argument('--forcing-dir', required=True) + parser.add_argument('--obs-dir', required=True) + parser.add_argument('--cfe-dir', required=True) + parser.add_argument('--config-file', required=True) + parser.add_argument('--param-bounds', required=True, + help='Kept for symmetry; not read in this script') + parser.add_argument('--out-dir', required=True) + parser.add_argument('--test-forcing-dir1', required=True) + parser.add_argument('--test-forcing-dir2', required=True) + parser.add_argument('--enkf-members', type=int, default=20) + parser.add_argument('--enkf-obs-error-std', type=float, default=0.05) + parser.add_argument('--no-vrugt-r', action='store_true') + parser.add_argument('--vrugt-alpha', type=float, default=0.10) + parser.add_argument('--vrugt-scale', type=float, default=0.001) + parser.add_argument('--rng-seed', type=int, default=None) + parser.add_argument('--base-step-h', type=int, default=6, + help='Issue-time cadence in hours across the full test ' + 'period (default 6h)') + parser.add_argument('--dense-start', type=str, default=DEFAULT_DENSE_START, + help='Start of hourly-densification window ' + '(default 2024-09-24 00:00:00 — Helene)') + parser.add_argument('--dense-end', type=str, default=DEFAULT_DENSE_END, + help='End of hourly-densification window ' + '(default 2024-09-28 23:00:00 — Helene)') + parser.add_argument('--sanity-check', action='store_true', + help='At the first issue time, fork fcst_models and step ' + 'both prod and fcst one hour with identical unperturbed ' + 'forcing. Confirms the 4-state snapshot is sufficient ' + 'to reproduce next-hour Q to machine precision, then ' + 'exits. Use before kicking off the full sweep.') + parser.add_argument('--hardcoded-r', type=float, default=None, + help='If set, override the Vrugt + kriging-variance R formula ' + 'and use this constant value (mm^2/h^2) at every hour. ' + 'Matches the DualEarth/new_EnKF.py reference (R=0.07). ' + 'Suppresses DA at low flow where kriging is noisy while ' + 'preserving DA strength at storm peaks. Applied to all ' + 'three EnKF instances (da, openloop, fcst) for consistency.') + parser.add_argument('--prod-script', default=None, + help='Path to calibrate_catchment_cfe_da_v2.py for ' + 'importing EnKFAssimilator. Defaults to next-to-this-file.') + args = parser.parse_args() + + here = os.path.dirname(os.path.abspath(__file__)) + prod_script = args.prod_script or os.path.join(here, 'calibrate_catchment_cfe_da_v2.py') + if not os.path.exists(prod_script): + raise FileNotFoundError( + f"Could not find production script at {prod_script}. " + f"Pass --prod-script to override.") + EnKFAssimilator = import_enkf_class(prod_script) + + sys.path.insert(0, args.cfe_dir) + import bmi_cfe as _bmi_cfe + + run(args, _bmi_cfe, EnKFAssimilator) + + +if __name__ == '__main__': + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/3_routing/route_det_f5_20pct.sh b/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/3_routing/route_det_f5_20pct.sh new file mode 100644 index 00000000..c98352c0 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/3_routing/route_det_f5_20pct.sh @@ -0,0 +1,29 @@ +#!/usr/bin/env bash +# F5 re-kriged variance (20% gauge holdout) — deterministic T-route routing. +# Routes per-catchment _test_results.csv through Muskingum-Cunge to gauge 03463300. +# +# Run AFTER batch_run_f5_20pct.sh completes. +# +# Usage: +# bash route_det_f5_20pct.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +ROUTE_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/2_troute_routing/run_route.py + +F5_DIR=/mnt/disk2/suma_helen_poster/da_results/folder5_rekrig_variance_direct +GPKG=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +echo "[F5-20pct] Deterministic T-route routing..." +echo " da-dir : $F5_DIR" +echo " out-dir: $F5_DIR" + +$TROUTE "$ROUTE_SCRIPT" \ + --gpkg "$GPKG" \ + --da-dir "$F5_DIR" \ + --out-dir "$F5_DIR" \ + --usgs-csv "$USGS_CSV" + +echo "[F5-20pct] Done. Output: $F5_DIR/routed_Q_test.csv" diff --git a/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/3_routing/route_leadtime_f5.sh b/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/3_routing/route_leadtime_f5.sh new file mode 100644 index 00000000..5d07194d --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/3_routing/route_leadtime_f5.sh @@ -0,0 +1,31 @@ +#!/usr/bin/env bash +# F5 re-kriged variance (20% holdout) — route lead-time forecast CSVs. +# Routes 21-catchment lead-time forecast CSVs through Muskingum-Cunge +# to gauge 03463300. Writes routed_leadtime_da_full.parquet and +# routed_leadtime_openloop_full.parquet for gauge-level NSE decay analysis. +# +# Run AFTER batch_run_f5_analysis_20pct.sh forecast completes. +# +# Usage: +# bash route_leadtime_f5.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +ROUTE_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/2_troute_routing/run_route_leadtime_forecasts.py + +F5_DIR=/mnt/disk2/suma_helen_poster/da_results/da_forecast_f5_rekrig_20pct +GPKG=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg +OUT_DIR=$F5_DIR/routed + +echo "[F5-20pct] Routing lead-time forecasts through T-Route..." +echo " forecast-dir: $F5_DIR" +echo " out-dir : $OUT_DIR" + +$TROUTE "$ROUTE_SCRIPT" \ + --gpkg "$GPKG" \ + --forecast-dir "$F5_DIR" \ + --out-dir "$OUT_DIR" + +echo "[F5-20pct] Lead-time routing done. Outputs in: $OUT_DIR" +ls "$OUT_DIR"/*.parquet 2>/dev/null || echo " (no parquets found)" diff --git a/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/4_evaluation/4a_error_decay/plot_lead_time_nse_gauge_f5.py b/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/4_evaluation/4a_error_decay/plot_lead_time_nse_gauge_f5.py new file mode 100644 index 00000000..dc8499f9 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/4_evaluation/4a_error_decay/plot_lead_time_nse_gauge_f5.py @@ -0,0 +1,234 @@ +""" +plot_lead_time_nse_gauge_f5.py + +Gauge-level lead-time NSE decay for F5 (re-kriged variance, 20% holdout). + +Reads routed lead-time forecast parquets produced by run_lead_time_forecast_sweep.py +(routed through run_route_leadtime_forecasts.py) and plots NSE vs lead time for +the DA-on (F5) and open-loop ensembles at USGS gauge 03463300. + +NSE per lead time is computed by pooling all (issue_time, valid_time) pairs for +that lead hour across the full test period, then: + NSE = 1 - sum((obs - ens_mean)^2) / sum((obs - mean(obs))^2) + +A second panel shows the regime split (Helene window vs. storm hours vs. low flow). + +Inputs: + /routed_leadtime_da_full.parquet + /routed_leadtime_openloop_full.parquet + +Output: + /lead_time_nse_gauge_f5_pooled.png + /lead_time_nse_gauge_f5_by_regime.png + +Run on server: + python3 plot_lead_time_nse_gauge_f5.py + python3 plot_lead_time_nse_gauge_f5.py --route-dir /path/to/routed --out-dir /path/to/out +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt + +DA_COLOR = "tab:purple" +OL_COLOR = "tab:gray" + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results/da_forecast_f5_rekrig_20pct/routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3_S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-28 23:00:00") +STORM_THRESHOLD_M3S = 20.0 +LOWFLOW_THRESHOLD_M3S = 5.0 + + +# ── I/O helpers ─────────────────────────────────────────────────────────────── + +def load_parquet(path): + if not os.path.exists(path): + raise FileNotFoundError(path) + df = pd.read_parquet(path) + df['issue_time'] = pd.to_datetime(df['issue_time']) + + if 'q_gauge_m3s' in df.columns and 'member' in df.columns: + return df[['issue_time', 'lead_hour', 'member', 'q_gauge_m3s']] + + member_cols = sorted(c for c in df.columns if c.startswith('member_')) + if not member_cols: + raise ValueError(f"Cannot identify member columns in {path}. " + f"Columns: {df.columns.tolist()}") + keep = ['issue_time', 'lead_hour'] + long = df[keep + member_cols].melt( + id_vars=keep, value_vars=member_cols, + var_name='member', value_name='q_gauge_m3s') + return long + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns if c.lower() in ('datetime', 'date', 'time', 'timestamp')) + q_col = next(c for c in df.columns if 'q' in c.lower() or 'flow' in c.lower() + or 'discharge' in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if 'mm' in q_col.lower(): + series = series * MM_H_TO_M3_S + print(f" Converted obs from mm/h ('{q_col}') → m³/s (×{MM_H_TO_M3_S:.4f})") + return series + + +# ── NSE computation ─────────────────────────────────────────────────────────── + +def nse_by_lead(df_long, obs_series, issue_mask=None): + """Return (leads, nse_values) pooling all issue times for each lead hour.""" + df = df_long.copy() + if issue_mask is not None: + df = df[df['issue_time'].isin(issue_mask)] + if len(df) == 0: + return None, None + + df['valid_time'] = df['issue_time'] + pd.to_timedelta(df['lead_hour'], unit='h') + df['obs'] = df['valid_time'].map(obs_series) + df = df.dropna(subset=['obs']) + if len(df) == 0: + return None, None + + ens_mean = (df.groupby(['issue_time', 'lead_hour'])['q_gauge_m3s'] + .mean().reset_index(name='ens_mean')) + obs_per = (df.groupby(['issue_time', 'lead_hour'])['obs'] + .first().reset_index(name='obs')) + panel = ens_mean.merge(obs_per, on=['issue_time', 'lead_hour']) + + leads = sorted(panel['lead_hour'].unique()) + nse_vals = [] + for L in leads: + sub = panel[panel['lead_hour'] == L] + obs_v = sub['obs'].values + sim_v = sub['ens_mean'].values + denom = np.sum((obs_v - obs_v.mean()) ** 2) + nse_L = 1.0 - np.sum((obs_v - sim_v) ** 2) / denom if denom > 0 else np.nan + nse_vals.append(nse_L) + + return np.asarray(leads), np.asarray(nse_vals) + + +# ── Plots ────────────────────────────────────────────────────────────────────── + +def plot_pooled(out_path, leads_da, nse_da, leads_ol, nse_ol): + fig, ax = plt.subplots(figsize=(10, 5)) + + ax.plot(leads_da, nse_da, color=DA_COLOR, lw=2.4, marker='o', + zorder=4, label="F5 DA on (re-kriged σ²)") + ax.plot(leads_ol, nse_ol, color=OL_COLOR, lw=2.4, marker='s', + linestyle='--', zorder=4, label="Open loop (no DA)") + + ax.axhline(0, color='black', lw=0.8, linestyle=':', alpha=0.5) + ax.set_xlabel("Forecast lead time (hours)", fontsize=11) + ax.set_ylabel("NSE (ensemble-mean vs USGS obs)", fontsize=11) + ax.set_title("Gauge-level lead-time NSE decay — F5 (re-kriged σ²), 20% holdout\n" + "USGS 03463300 | all issue times pooled (Oct 2023 – Oct 2024)", fontsize=12) + ax.set_xticks(np.arange(0, int(max(leads_da)) + 1, 1)) + ax.set_ylim(-0.1, 1.05) + ax.grid(True, alpha=0.25) + ax.legend(fontsize=10, loc='lower left', frameon=True, framealpha=0.92) + + plt.tight_layout() + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + +def plot_by_regime(out_path, regime_results): + n = len(regime_results) + fig, axes = plt.subplots(1, n, figsize=(6 * n, 5), sharey=True) + if n == 1: + axes = [axes] + + for ax, (label, n_issue, res_da, res_ol) in zip(axes, regime_results): + if res_da[0] is None: + ax.set_title(f"{label}\n(no issue times)", fontsize=10) + ax.axis('off') + continue + leads_da, nse_da = res_da + leads_ol, nse_ol = res_ol + ax.plot(leads_da, nse_da, color=DA_COLOR, lw=2.2, marker='o', + label="F5 DA on") + ax.plot(leads_ol, nse_ol, color=OL_COLOR, lw=2.2, marker='s', + linestyle='--', label="Open loop") + ax.axhline(0, color='black', lw=0.8, linestyle=':', alpha=0.5) + ax.set_title(f"{label}\n({n_issue} issue times)", fontsize=10) + ax.set_xlabel("Lead time (hours)", fontsize=10) + ax.set_xticks(np.arange(0, int(max(leads_da)) + 1, 1)) + ax.grid(True, alpha=0.25) + ax.legend(fontsize=9, loc='lower left', frameon=True, framealpha=0.9) + + axes[0].set_ylabel("NSE (ensemble-mean vs USGS obs)", fontsize=11) + fig.suptitle("Lead-time NSE decay by regime — F5 (re-kriged σ²), USGS 03463300, 20% holdout", + fontsize=12, y=1.01) + plt.tight_layout() + plt.savefig(out_path, dpi=160, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + + +# ── Main ────────────────────────────────────────────────────────────────────── + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--route-dir', default=DEFAULT_ROUTE_DIR) + parser.add_argument('--out-dir', default=None) + parser.add_argument('--usgs-csv', default=DEFAULT_USGS_CSV) + parser.add_argument('--da-name', default='routed_leadtime_da_full.parquet') + parser.add_argument('--ol-name', default='routed_leadtime_openloop_full.parquet') + args = parser.parse_args() + + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + print("Loading parquets...") + df_da = load_parquet(os.path.join(args.route_dir, args.da_name)) + df_ol = load_parquet(os.path.join(args.route_dir, args.ol_name)) + obs = load_usgs(args.usgs_csv) + + print(f"DA rows: {len(df_da):,} | OL rows: {len(df_ol):,}") + + # ---- Pooled NSE ---- + leads_da, nse_da = nse_by_lead(df_da, obs) + leads_ol, nse_ol = nse_by_lead(df_ol, obs) + plot_pooled( + os.path.join(out_dir, "lead_time_nse_gauge_f5_pooled.png"), + leads_da, nse_da, leads_ol, nse_ol) + + # ---- Regime split ---- + issue_times = pd.to_datetime(sorted(df_da['issue_time'].unique())) + obs_at_issue = pd.Series(issue_times, index=issue_times).map(obs) + + helene_mask = (issue_times >= HELENE_START) & (issue_times <= HELENE_END) + storm_mask = obs_at_issue > STORM_THRESHOLD_M3S + lowflow_mask = obs_at_issue < LOWFLOW_THRESHOLD_M3S + + regime_results = [] + for label, mask in [ + (f"Helene window (Sep 24–28 2024)", helene_mask), + (f"Storm hours (obs > {STORM_THRESHOLD_M3S:.0f} m³/s)", storm_mask), + (f"Low flow (obs < {LOWFLOW_THRESHOLD_M3S:.0f} m³/s)", lowflow_mask), + ]: + kept = set(pd.to_datetime(issue_times[mask])) + res_da = nse_by_lead(df_da, obs, kept) + res_ol = nse_by_lead(df_ol, obs, kept) + regime_results.append((label, len(kept), res_da, res_ol)) + + plot_by_regime( + os.path.join(out_dir, "lead_time_nse_gauge_f5_by_regime.png"), + regime_results) + + print("Done.") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/4_evaluation/4a_error_decay/run_4a_f5.sh b/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/4_evaluation/4a_error_decay/run_4a_f5.sh new file mode 100644 index 00000000..bfadfaf6 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/4_evaluation/4a_error_decay/run_4a_f5.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +# F5 re-kriged variance (20% holdout) — gauge-level lead-time NSE evaluation. +# Run AFTER the lead-time forecast parquets have been routed via T-Route. +# +# Usage: +# bash run_4a_f5.sh + +set -euo pipefail + +PYTHON=/home/svyas/miniconda3/envs/troute/bin/python3 +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +ROUTE_DIR=/mnt/disk2/suma_helen_poster/da_results/da_forecast_f5_rekrig_20pct/routed +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv +OUT_DIR=/mnt/disk2/suma_helen_poster/da_results/da_forecast_f5_rekrig_20pct/figures + +mkdir -p "$OUT_DIR" + +echo "[F5-20pct] Lead-time NSE evaluation..." +$PYTHON "$SCRIPT_DIR/plot_lead_time_nse_gauge_f5.py" \ + --route-dir "$ROUTE_DIR" \ + --usgs-csv "$USGS_CSV" \ + --out-dir "$OUT_DIR" + +echo "[F5-20pct] Done. Figures in: $OUT_DIR" diff --git a/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/4_evaluation/4b_helene_ensemble_fan/plot_helene_fan_f5.py b/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/4_evaluation/4b_helene_ensemble_fan/plot_helene_fan_f5.py new file mode 100644 index 00000000..37eb8dec --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/4_evaluation/4b_helene_ensemble_fan/plot_helene_fan_f5.py @@ -0,0 +1,236 @@ +""" +plot_helene_fan_f5.py + +Helene-window ensemble fan plot for F5 (re-kriged variance, 20% holdout) +at USGS gauge 03463300. + +For each daily init time (issue_time) in Sep 24–30 2024: + - shaded band = member min–max envelope + - thin line = ensemble median +Grand mean across all init times shown as thick line. +USGS obs = black solid. Open-loop grand median = black dashed. + +Reads routed parquets produced by run_route_leadtime_forecasts.py. + +Usage: + python3 plot_helene_fan_f5.py + python3 plot_helene_fan_f5.py --route-dir /path/to/routed --out-dir /path/to/out +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.dates as mdates + +DEFAULT_ROUTE_DIR = "/mnt/disk2/suma_helen_poster/da_results/da_forecast_f5_rekrig_20pct/routed" +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" + +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-30 23:00:00") + +WATERSHED_AREA_KM2 = 113.18 +MM_H_TO_M3S = WATERSHED_AREA_KM2 * 1000.0 / 3600.0 + +# One colour per init date (7 days) +INIT_COLORS = [ + "#1f77b4", # Sep-24 blue + "#ff7f0e", # Sep-25 orange + "#2ca02c", # Sep-26 green + "#d62728", # Sep-27 red + "#9467bd", # Sep-28 purple + "#8c564b", # Sep-29 brown + "#e377c2", # Sep-30 pink +] + + +# ── helpers ─────────────────────────────────────────────────────────────────── + +def load_parquet(path): + if not os.path.exists(path): + raise FileNotFoundError(path) + df = pd.read_parquet(path) + df['issue_time'] = pd.to_datetime(df['issue_time']) + + if 'q_gauge_m3s' in df.columns and 'member' in df.columns: + return df[['issue_time', 'lead_hour', 'member', 'q_gauge_m3s']] + + member_cols = sorted(c for c in df.columns if c.startswith('member_')) + if not member_cols: + raise ValueError(f"Cannot identify member columns in {path}. Columns: {df.columns.tolist()}") + keep = ['issue_time', 'lead_hour'] + long = df[keep + member_cols].melt( + id_vars=keep, value_vars=member_cols, + var_name='member', value_name='q_gauge_m3s') + return long + + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns if c.lower() in ('datetime', 'date', 'time', 'timestamp')) + q_col = next(c for c in df.columns if 'q' in c.lower() or 'flow' in c.lower() + or 'discharge' in c.lower()) + df[date_col] = pd.to_datetime(df[date_col]) + series = df.set_index(date_col)[q_col].astype(float) + if 'mm' in q_col.lower(): + series = series * MM_H_TO_M3S + print(f" Converted obs from mm/h ('{q_col}') → m³/s (×{MM_H_TO_M3S:.4f})") + return series + + +def add_valid_time(df): + df = df.copy() + df['valid_time'] = df['issue_time'] + pd.to_timedelta(df['lead_hour'], unit='h') + return df + + +# ── main plot ───────────────────────────────────────────────────────────────── + +def plot_fan(df_da, df_ol, obs, out_path): + df_da = add_valid_time(df_da) + df_ol = add_valid_time(df_ol) + + # Use one representative init time per day: 00:00 UTC of each day in Sep 24-30 + # This gives a single 20-member fan per day (wide spread), matching the reference style. + day_range = pd.date_range(HELENE_START.normalize(), HELENE_END.normalize(), freq='D') + plot_end = HELENE_END + pd.Timedelta(hours=24) + obs_h = obs[(obs.index >= HELENE_START) & (obs.index <= plot_end)] + + # For each day pick the closest available issue_time to 00:00 + all_issue = pd.to_datetime(sorted(df_da['issue_time'].unique())) + + fig, ax = plt.subplots(figsize=(14, 5)) + + # ---- Helene window background shading ---- + ax.axvspan(HELENE_START, HELENE_END, color='#ffcccc', alpha=0.35, zorder=0) + + # ---- per-day fan: single representative init time → full 20-member spread ---- + all_mean_series = [] + legend_handles = [] + for i, day in enumerate(day_range): + color = INIT_COLORS[i % len(INIT_COLORS)] + + # pick closest issue_time to 00:00 of that day + diffs = np.abs((all_issue - day).total_seconds()) + t0 = all_issue[diffs.argmin()] + + sub = df_da[df_da['issue_time'] == t0].copy() + sub = add_valid_time(sub) + if sub.empty: + continue + + piv = sub.pivot_table(index='valid_time', columns='member', + values='q_gauge_m3s', aggfunc='mean').sort_index() + # clip negatives (routing artefacts) + piv = piv.clip(lower=0) + vt = piv.index + mn = piv.min(axis=1).values + mx = piv.max(axis=1).values + med = piv.median(axis=1).values + mean_ = piv.mean(axis=1).values + + # draw all 20 member lines faint, then median thick + for col in piv.columns: + ax.plot(vt, piv[col].values, color=color, lw=0.5, alpha=0.25, zorder=2) + ax.fill_between(vt, mn, mx, color=color, alpha=0.30, linewidth=0, zorder=2) + line, = ax.plot(vt, med, color=color, lw=1.6, alpha=0.95, zorder=3) + legend_handles.append((color, day.strftime("Init %Y-%m-%d"))) + + all_mean_series.append(pd.Series(mean_, index=vt, name=day)) + + # ---- grand mean across all init days (thick green) ---- + if all_mean_series: + grand = pd.concat(all_mean_series, axis=1).mean(axis=1).sort_index() + ax.plot(grand.index, grand.values, color='#2ca02c', lw=3.2, zorder=5, + label="F5 DA — mean across all forecasts") + + # ---- open loop: pick one representative init per day, show grand median ---- + ol_meds = [] + for day in day_range: + diffs = np.abs((all_issue - day).total_seconds()) + t0 = all_issue[diffs.argmin()] + sub = df_ol[df_ol['issue_time'] == t0].copy() + sub = add_valid_time(sub) + if sub.empty: + continue + piv = sub.pivot_table(index='valid_time', columns='member', + values='q_gauge_m3s', aggfunc='mean').sort_index() + ol_meds.append(piv.clip(lower=0).median(axis=1)) + if ol_meds: + ol_grand = pd.concat(ol_meds, axis=1).median(axis=1).sort_index() + ax.plot(ol_grand.index, ol_grand.values, + color='black', lw=1.8, linestyle='--', zorder=4, + label="Open loop (no DA) — grand median") + + # ---- USGS obs ---- + ax.plot(obs_h.index, obs_h.values, color='black', lw=2.4, zorder=6, label="USGS obs") + + # ---- Helene peak annotation ---- + if len(obs_h) > 0: + peak_t = obs_h.idxmax() + peak_v = obs_h.max() + ax.annotate("Helene peak", + xy=(peak_t, peak_v), + xytext=(peak_t + pd.Timedelta(hours=10), peak_v * 0.96), + fontsize=9, color='#d62728', + arrowprops=dict(arrowstyle='->', color='#d62728', lw=1.2)) + + # ---- legend ---- + import matplotlib.patches as mpatches + day_patches = [mpatches.Patch(color=INIT_COLORS[i % len(INIT_COLORS)], alpha=0.8, + label=label) + for i, (_, label) in enumerate(legend_handles)] + extra = [ + plt.Line2D([0], [0], color='#2ca02c', lw=3.0, + label="F5 DA — mean across all forecasts"), + plt.Line2D([0], [0], color='black', lw=1.8, linestyle='--', + label="Open loop (no DA) — grand median"), + plt.Line2D([0], [0], color='black', lw=2.4, label="USGS obs"), + ] + ax.legend(handles=day_patches + extra, + fontsize=9, loc='upper left', framealpha=0.92, ncol=1) + + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_xlabel("Date (UTC)", fontsize=11) + ax.set_title( + "F5 (re-kriged σ²) — ensemble forecast fan: initial state uncertainty | Sep 24–30 2024\n" + "USGS 03463300 | Shaded = member min–max | Line = median per init time | Thick = grand mean", + fontsize=11) + ax.set_xlim(HELENE_START, plot_end) + ax.grid(True, alpha=0.22) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=1)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=20, ha='right', fontsize=9) + plt.tight_layout() + plt.savefig(out_path, dpi=160, bbox_inches='tight') + plt.close() + print(f"Saved: {out_path}") + + +# ── entry point ─────────────────────────────────────────────────────────────── + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--route-dir', default=DEFAULT_ROUTE_DIR) + parser.add_argument('--usgs-csv', default=DEFAULT_USGS_CSV) + parser.add_argument('--out-dir', default=None) + parser.add_argument('--da-name', default='routed_leadtime_da_full.parquet') + parser.add_argument('--ol-name', default='routed_leadtime_openloop_full.parquet') + args = parser.parse_args() + + out_dir = args.out_dir or args.route_dir + os.makedirs(out_dir, exist_ok=True) + + print("Loading parquets...") + df_da = load_parquet(os.path.join(args.route_dir, args.da_name)) + df_ol = load_parquet(os.path.join(args.route_dir, args.ol_name)) + obs = load_usgs(args.usgs_csv) + + out_path = os.path.join(out_dir, "helene_ensemble_fan_f5.png") + plot_fan(df_da, df_ol, obs, out_path) + print("Done.") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/4_evaluation/4b_helene_ensemble_fan/plot_hydro_arm_with_openloop_f5.py b/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/4_evaluation/4b_helene_ensemble_fan/plot_hydro_arm_with_openloop_f5.py new file mode 100644 index 00000000..4427fc41 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/4_evaluation/4b_helene_ensemble_fan/plot_hydro_arm_with_openloop_f5.py @@ -0,0 +1,364 @@ +""" +plot_hydro_arm_with_openloop_f5.py + +Plots the 2b hydro-state arm (DA on, 20 members: initial state uncertainty) +for cat-1016300 during Hurricane Helene (Sep 24-30 2024), with open-loop +grand median overlaid as a black dashed line. + +Reads: + - //_da_hydro_arm.csv (all 21 catchments, mm/h) + - /routed_leadtime_openloop_full.parquet (open loop m³/s at gauge) + - USGS obs CSV + +Routes arm CSVs through T-Route Muskingum-Cunge to gauge 03463300, +one init time per Helene day (Sep 24-30), 20 members each. + +Run on server (troute env): + python3 plot_hydro_arm_with_openloop_f5.py +""" + +import argparse +import os +import sys +import sqlite3 +import types as _types +from functools import partial + +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib.patches as mpatches +import matplotlib.dates as mdates + +# ── Coverage stub required for numba/troute on Python 3.10 ─────────────────── +_stub = _types.ModuleType('coverage.types') +for _cls in ['Tracer', 'TTraceData', 'TShouldTraceFn', 'TFileDisposition', + 'TShouldStartContextFn', 'TWarnFn', 'TTraceFn']: + setattr(_stub, _cls, type(_cls, (), {})) +sys.modules['coverage.types'] = _stub + +import troute.nhd_network as nhd_network +from troute.routing.fast_reach.mc_reach import compute_network_structured + +# ── Constants ───────────────────────────────────────────────────────────────── +DEFAULT_ARMS_DIR = "/mnt/disk2/suma_helen_poster/da_results/da_arms_f5_rekrig_20pct" +DEFAULT_ROUTE_DIR = ("/mnt/disk2/suma_helen_poster/da_results/" + "da_forecast_f5_rekrig_20pct/routed") +DEFAULT_USGS_CSV = "/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv" +DEFAULT_GPKG = ("/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/" + "gage-03463300_subset.gpkg") +DEFAULT_OUT_DIR = None # falls back to --arms-dir + +CATS = [ + 'cat-1016279', 'cat-1016280', 'cat-1016281', 'cat-1016282', 'cat-1016283', + 'cat-1016300', 'cat-1016301', 'cat-1016302', 'cat-1016303', 'cat-1016304', + 'cat-1016305', 'cat-1016306', 'cat-1016307', 'cat-1016308', 'cat-1016309', + 'cat-1016310', 'cat-1016311', 'cat-1016312', 'cat-1016313', 'cat-1016314', + 'cat-1016315', +] + +TERMINAL_INT = 1016283 # wb-1016283 = gauge 03463300 +DT = 3600.0 +QTS_SUBDIVISIONS = 1 +WATERSHED_AREA_KM2 = 113.18 # total drainage area to gauge 03463300 + +HELENE_START = pd.Timestamp("2024-09-24 00:00:00") +HELENE_END = pd.Timestamp("2024-09-30 23:00:00") + +INIT_COLORS = [ + "#1f77b4", # Sep-24 + "#ff7f0e", # Sep-25 + "#2ca02c", # Sep-26 + "#d62728", # Sep-27 + "#9467bd", # Sep-28 + "#8c564b", # Sep-29 + "#e377c2", # Sep-30 +] + + +# ── T-Route helpers ─────────────────────────────────────────────────────────── + +def read_network(gpkg): + con = sqlite3.connect(gpkg) + fp_attr = pd.read_sql( + 'SELECT link, "to", BtmWdth, TopWdth, TopWdthCC, n, nCC, ChSlp, So, Length_m ' + 'FROM "flowpath-attributes"', con) + fp = pd.read_sql('SELECT divide_id, areasqkm FROM flowpaths', con) + con.close() + return fp_attr, fp + + +def build_connections(fp_attr): + link_set = {int(r[3:]) for r in fp_attr['link']} # "wb-XXXX" + connections = {} + for _, row in fp_attr.iterrows(): + us = int(row['link'][3:]) + ds = int(row['to'][4:]) # "nex-XXXX" + connections[us] = [ds] if ds in link_set else [] + return connections + + +def build_param_df(fp_attr): + rows = [{'seg_id': int(r['link'][3:]), + 'dt': float(DT), + 'bw': float(r['BtmWdth']), 'tw': float(r['TopWdth']), + 'twcc': float(r['TopWdthCC']), 'dx': float(r['Length_m']), + 'n': float(r['n']), 'ncc': float(r['nCC']), + 'cs': float(r['ChSlp']), 's0': float(r['So']), + 'alt': 0.0} + for _, r in fp_attr.iterrows()] + return pd.DataFrame(rows).set_index('seg_id').sort_index().astype('float32') + + +def route_one(reaches_wTypes, upstreams, param_df, q0_df, + qlat_arr, nts, terminal_pos): + """Route one member's qlat_arr (n_segs × nts) through MC; return Q at terminal.""" + e1i = np.zeros(0, dtype='int32') + e1f = np.zeros(0, dtype='float32') + e2f = np.zeros((0, nts), dtype='float32') + e00f32 = np.zeros((0, 0), dtype='float32') + e00f64 = np.zeros((0, 0), dtype='float64') + e00i32 = np.zeros((0, 0), dtype='int32') + + results = compute_network_structured( + nts, DT, QTS_SUBDIVISIONS, + reaches_wTypes, upstreams, + param_df.index.values.astype('int64'), + param_df.columns.values, + param_df.values, + q0_df.values.astype('float32'), + qlat_arr.astype('float32'), + [], e00f64, {}, e00i32, False, + '2024-09-24_00:00:00', + e2f, e1i, e1i, e1i, e1f, e1f, 0.0, + e2f, e1i, e1f, e1f, e1f, e1f, e1f, + e2f, e1i, e1f, e1f, e1f, e1f, e1f, + e2f, e1i, e1i, [], e1i, e1i, e1f, e1i, e1i, + e1i, e1i, e1f, e1i, e1f, e1i, e1i, e00f32, + ) + seg_ids = np.asarray(results[0]) + fvd = np.asarray(results[1]) + Q_var = fvd[terminal_pos, :nts] + Q_time = fvd[terminal_pos, 0::3] + return Q_var if Q_var.max() > Q_time.max() else Q_time + + +def route_issue_time(t0, arm_dfs, member_cols, lead_hours, + reaches_wTypes, upstreams, param_df, q0_df, + seg_ids_sorted, area_map, terminal_pos): + """Route all 20 members for one issue time. Returns DataFrame index=valid_time.""" + n_segs = len(param_df) + n_leads = len(lead_hours) + n_members = len(member_cols) + + qlat_cube = np.zeros((n_members, n_segs, n_leads), dtype='float32') + for cat, df in arm_dfs.items(): + sid = int(cat[4:]) # "cat-XXXX" + area_m2 = area_map.get(sid, 0.0) + if sid not in seg_ids_sorted: + continue + seg_pos = int(np.where(seg_ids_sorted == sid)[0][0]) + sub = df[df['issue_time'] == t0].sort_values('lead_hour') + if sub.empty: + continue + vals = np.maximum(sub[member_cols].to_numpy(dtype='float32'), 0.0) + vals = vals / 1000.0 / 3600.0 * area_m2 # mm/h → m³/s + qlat_cube[:, seg_pos, :vals.shape[0]] = vals.T + + Q_members = np.full((n_leads, n_members), np.nan, dtype='float32') + for m in range(n_members): + Q_m = route_one(reaches_wTypes, upstreams, param_df, q0_df, + qlat_cube[m], n_leads, terminal_pos) + Q_members[:, m] = Q_m[:n_leads] + + valid_times = pd.DatetimeIndex([t0 + pd.Timedelta(hours=int(h)) for h in lead_hours]) + return pd.DataFrame(np.maximum(Q_members, 0.0), + columns=member_cols, index=valid_times) + + +# ── USGS loader ─────────────────────────────────────────────────────────────── + +def load_usgs(usgs_csv): + df = pd.read_csv(usgs_csv) + date_col = next(c for c in df.columns + if c.lower() in ('datetime', 'date', 'time', 'timestamp')) + q_col = next(c for c in df.columns + if any(k in c.lower() for k in ('q', 'flow', 'discharge'))) + df[date_col] = pd.to_datetime(df[date_col]) + obs = df.set_index(date_col)[q_col].astype(float) + if 'mm' in q_col.lower(): + obs = obs * WATERSHED_AREA_KM2 * 1000.0 / 3600.0 # mm/h → m³/s + return obs + + +# ── Open loop helper ────────────────────────────────────────────────────────── + +def load_ol_grand_median(route_dir): + """Return open-loop grand median (m³/s) indexed by valid_time for Helene.""" + path = os.path.join(route_dir, 'routed_leadtime_openloop_full.parquet') + df_ol = pd.read_parquet(path) + df_ol['issue_time'] = pd.to_datetime(df_ol['issue_time']) + + helene = df_ol[(df_ol['issue_time'] >= HELENE_START) & + (df_ol['issue_time'] <= HELENE_END)] + + if 'q_gauge_m3s' in helene.columns: + helene = helene.copy() + helene['valid_time'] = (helene['issue_time'] + + pd.to_timedelta(helene['lead_hour'], unit='h')) + grand = helene.groupby('valid_time')['q_gauge_m3s'].median() + else: + mem_cols = sorted(c for c in helene.columns if c.startswith('member_')) + helene = helene.copy() + helene['ens_mean'] = helene[mem_cols].mean(axis=1) + helene['valid_time'] = (helene['issue_time'] + + pd.to_timedelta(helene['lead_hour'], unit='h')) + grand = helene.groupby('valid_time')['ens_mean'].median() + + return grand.clip(lower=0).sort_index() + + +# ── Main ────────────────────────────────────────────────────────────────────── + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--arms-dir', default=DEFAULT_ARMS_DIR) + parser.add_argument('--route-dir', default=DEFAULT_ROUTE_DIR) + parser.add_argument('--usgs-csv', default=DEFAULT_USGS_CSV) + parser.add_argument('--gpkg', default=DEFAULT_GPKG) + parser.add_argument('--out-dir', default=DEFAULT_OUT_DIR) + args = parser.parse_args() + out_dir = args.out_dir or args.arms_dir + os.makedirs(out_dir, exist_ok=True) + + # ── Load hydro arm CSVs ─────────────────────────────────────────────────── + print("Loading hydro arm CSVs...") + arm_dfs = {} + for cat in CATS: + p = os.path.join(args.arms_dir, cat, f"{cat}_da_hydro_arm.csv") + if os.path.exists(p): + df = pd.read_csv(p) + df['issue_time'] = pd.to_datetime(df['issue_time']) + arm_dfs[cat] = df + if 'cat-1016300' not in arm_dfs: + raise FileNotFoundError(f"cat-1016300_da_hydro_arm.csv not found in {args.arms_dir}") + print(f" Loaded {len(arm_dfs)}/21 catchments") + + ref_df = arm_dfs['cat-1016300'] + member_cols = sorted(c for c in ref_df.columns if c.startswith('member_')) + all_issues = pd.to_datetime(sorted(ref_df['issue_time'].unique())) + lead_hours = sorted(ref_df['lead_hour'].unique()) + print(f" {len(all_issues)} issue times | {len(lead_hours)} leads | {len(member_cols)} members") + + # ── Build T-Route network (once) ────────────────────────────────────────── + print("Building T-Route network from GPKG...") + fp_attr, fp = read_network(args.gpkg) + connections = build_connections(fp_attr) + rconn = nhd_network.reverse_network(connections) + path_func = partial(nhd_network.split_at_junction, rconn) + reach_list = nhd_network.dfs_decomposition(rconn, path_func) + reaches_wTypes = [(r, 0) for r in reach_list] + upstreams = dict(rconn) + param_df = build_param_df(fp_attr) + n_segs = len(param_df) + seg_ids_sorted = param_df.index.values + + area_map = {int(r['divide_id'][4:]): r['areasqkm'] * 1e6 + for _, r in fp.iterrows() + if r['divide_id'] and str(r['divide_id']).startswith('cat-')} + + terminal_pos = int(np.where(seg_ids_sorted == TERMINAL_INT)[0][0]) + q0_df = pd.DataFrame(np.zeros((n_segs, 3), dtype='float32'), + index=param_df.index, columns=['qu0', 'qd0', 'h0']) + print(f" {n_segs} segments | terminal wb-{TERMINAL_INT} at index {terminal_pos}") + + # ── Load open loop + USGS ───────────────────────────────────────────────── + print("Loading open loop parquet...") + ol_grand = load_ol_grand_median(args.route_dir) + print("Loading USGS obs...") + usgs = load_usgs(args.usgs_csv) + + # ── Route one init time per Helene day + plot ───────────────────────────── + print("Routing arms through T-Route (7 init times × 20 members)...") + day_range = pd.date_range(HELENE_START.normalize(), HELENE_END.normalize(), freq='D') + + HELENE_PEAK_START = pd.Timestamp("2024-09-26 12:00:00") + HELENE_PEAK_END = pd.Timestamp("2024-09-28 00:00:00") + + fig, ax = plt.subplots(figsize=(17, 6)) + ax.axvspan(HELENE_PEAK_START, HELENE_PEAK_END, + color='salmon', alpha=0.12, zorder=0, label='_nolegend_') + ax.text(HELENE_PEAK_START + pd.Timedelta(hours=6), 1.0, "Helene peak", + transform=ax.get_xaxis_transform(), + fontsize=9, color='firebrick', ha='left', va='top') + + legend_patches = [] + + for i, day in enumerate(day_range): + color = INIT_COLORS[i % len(INIT_COLORS)] + diffs = np.abs((all_issues - day).total_seconds()) + t0 = all_issues[diffs.argmin()] + print(f" Init {day.date()} → matched {t0} ...") + + piv = route_issue_time( + t0, arm_dfs, member_cols, lead_hours, + reaches_wTypes, upstreams, param_df, q0_df, + seg_ids_sorted, area_map, terminal_pos) + + vt = piv.index + mn = piv.min(axis=1) + mx = piv.max(axis=1) + med = piv.median(axis=1) + + for col in piv.columns: + ax.plot(vt, piv[col].values, color=color, lw=0.4, alpha=0.18, zorder=2) + ax.fill_between(vt, mn, mx, color=color, alpha=0.22, linewidth=0, zorder=2) + ax.plot(vt, med.values, color=color, lw=1.8, alpha=0.92, zorder=3) + + legend_patches.append(mpatches.Patch(color=color, alpha=0.85, + label=f"Init {day.strftime('%Y-%m-%d')}")) + + # ── Open loop grand median ──────────────────────────────────────────────── + ol_mask = ((ol_grand.index >= HELENE_START) & + (ol_grand.index <= HELENE_END + pd.Timedelta(hours=24))) + ax.plot(ol_grand[ol_mask].index, ol_grand[ol_mask].values, + color='black', lw=2.0, linestyle='--', zorder=5, + label="Open loop (no DA) — grand median") + + # ── USGS obs ────────────────────────────────────────────────────────────── + obs_mask = ((usgs.index >= HELENE_START) & + (usgs.index <= HELENE_END + pd.Timedelta(hours=24))) + ax.plot(usgs[obs_mask].index, usgs[obs_mask].values, + color='black', lw=2.6, zorder=6, label="USGS obs") + + # ── Legend ──────────────────────────────────────────────────────────────── + extra_lines = [ + plt.Line2D([0], [0], color='black', lw=2.0, linestyle='--', + label="Open loop (no DA) — grand median"), + plt.Line2D([0], [0], color='black', lw=2.6, label="USGS obs"), + ] + ax.legend(handles=legend_patches + extra_lines, + fontsize=8, loc='upper left', framealpha=0.92, ncol=2) + + ax.set_ylabel("Discharge (m³/s)", fontsize=11) + ax.set_xlabel("Date (UTC)", fontsize=11) + ax.set_title( + "2b — Hydro-state arm (DA on, 20 members): initial state uncertainty | " + "Sep 24–30 2024\n" + "cat-1016300 | Shaded = member min–max | Line = median per init time", + fontsize=11) + ax.set_xlim(HELENE_START, HELENE_END + pd.Timedelta(hours=24)) + ax.grid(True, alpha=0.22) + ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d")) + ax.xaxis.set_major_locator(mdates.DayLocator(interval=1)) + plt.setp(ax.xaxis.get_majorticklabels(), rotation=20, ha='right', fontsize=9) + + plt.tight_layout() + out_path = os.path.join(out_dir, "cat-1016300_2b_hydro_arm_helene.png") + plt.savefig(out_path, dpi=150, bbox_inches='tight') + plt.close() + print(f"Saved: {out_path}") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/4_evaluation/4c_catchment_da_check/plot_catchment_da_check_f5.py b/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/4_evaluation/4c_catchment_da_check/plot_catchment_da_check_f5.py new file mode 100644 index 00000000..33f4c255 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/4_evaluation/4c_catchment_da_check/plot_catchment_da_check_f5.py @@ -0,0 +1,126 @@ +""" +plot_catchment_da_check_f5.py + +Catchment-level sanity check for F5 (re-kriged variance, 20% holdout). + +Reads _test_results.csv (DA ensemble mean output from Step 1) and plots: + - sim_mm_h : DA ensemble mean streamflow prediction + - obs_mm_h : Qkrig observation DA was trying to match + - precip : rainfall (inverted second axis) + +Two panels per catchment: + Panel 1 : full test period (Oct 2023 - Oct 2024) + Panel 2 : Helene zoom (Sep 20 - Oct 5 2024) + +NSE at catchment level (sim vs Qkrig) reported in title. +This figure checks DA in isolation — before routing — so any +remaining error at the gauge can be attributed to routing vs DA. + +Run on server: + python3 plot_catchment_da_check_f5.py + python3 plot_catchment_da_check_f5.py --da-dir /path/to/da --out-dir /path/to/out +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt + +DEFAULT_DA_DIR = "/mnt/disk2/suma_helen_poster/da_results/folder5_rekrig_variance_direct" +DEFAULT_OUT_DIR = None # falls back to DEFAULT_DA_DIR + +CAT_ID = "cat-1016300" # focal catchment draining to gauge 03463300 + +HELENE_ZOOM_START = pd.Timestamp("2024-09-20") +HELENE_ZOOM_END = pd.Timestamp("2024-10-05") +HELENE_PEAK_START = pd.Timestamp("2024-09-24") +HELENE_PEAK_END = pd.Timestamp("2024-09-30") + +SIM_COLOR = "tomato" +OBS_COLOR = "black" + + +def nse(obs, sim): + mask = ~(np.isnan(obs) | np.isnan(sim)) + o, s = obs[mask], sim[mask] + if len(o) < 2: + return float("nan") + denom = ((o - o.mean()) ** 2).sum() + return float(1.0 - ((o - s) ** 2).sum() / denom) if denom > 0 else float("nan") + + +def plot_catchment(cat_id, df, out_dir): + dates = pd.to_datetime(df["date"]) + sim = df["sim_mm_h"].values.astype(float) + obs = df["obs_mm_h"].values.astype(float) + precip = df["precip_mm_h"].values.astype(float) + + nse_full = nse(obs, sim) + + helene_mask = (dates >= HELENE_ZOOM_START) & (dates <= HELENE_ZOOM_END) + nse_helene = nse(obs[helene_mask], sim[helene_mask]) + + fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(18, 9)) + fig.suptitle( + f"{cat_id} | F5 catchment-level DA check (re-kriged σ², 20% holdout)\n" + f"NSE = {nse_full:.3f} (full year) NSE = {nse_helene:.3f} (Helene zoom)", + fontsize=12, + ) + + # ── Panel 1: full period ────────────────────────────────────────── + ax1.plot(dates, sim, color=SIM_COLOR, lw=1.2, label="DA ensemble mean [mm/h]") + ax1.plot(dates, obs, color=OBS_COLOR, lw=0.8, alpha=0.7, label="Qkrig obs [mm/h]") + ax1.set_ylabel("Streamflow (mm/h)", fontsize=10) + ax1.set_title("Full test period: Oct 2023 – Oct 2024", fontsize=10) + ax1.legend(fontsize=9, loc="upper left") + ax1.grid(True, alpha=0.2) + ax1_twin = ax1.twinx() + ax1_twin.bar(dates, precip, color="steelblue", alpha=0.35, width=0.04) + ax1_twin.set_ylim([precip.max() * 4, 0]) + ax1_twin.set_ylabel("Precip (mm/h)", fontsize=9, color="steelblue") + + # ── Panel 2: Helene zoom ────────────────────────────────────────── + d_h = dates[helene_mask] + sim_h = sim[helene_mask] + obs_h = obs[helene_mask] + pre_h = precip[helene_mask] + + ax2.axvspan(HELENE_PEAK_START, HELENE_PEAK_END, color="#ffcccc", alpha=0.35, zorder=0) + ax2.plot(d_h, sim_h, color=SIM_COLOR, lw=2.0, label="DA ensemble mean [mm/h]") + ax2.plot(d_h, obs_h, color=OBS_COLOR, lw=1.5, label="Qkrig obs [mm/h]") + ax2.set_ylabel("Streamflow (mm/h)", fontsize=10) + ax2.set_title("Helene zoom: Sep 20 – Oct 5 2024 (pink = Sep 24–30)", fontsize=10) + ax2.legend(fontsize=9, loc="upper left") + ax2.grid(True, alpha=0.2) + ax2_twin = ax2.twinx() + ax2_twin.bar(d_h, pre_h, color="steelblue", alpha=0.40, width=0.04) + ax2_twin.set_ylim([pre_h.max() * 4, 0]) + ax2_twin.set_ylabel("Precip (mm/h)", fontsize=9, color="steelblue") + + plt.tight_layout() + out_path = os.path.join(out_dir, f"{cat_id}_catchment_da_check.png") + plt.savefig(out_path, dpi=150, bbox_inches="tight") + plt.close() + print(f" saved: {out_path} (NSE={nse_full:.3f})") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--da-dir", default=DEFAULT_DA_DIR) + parser.add_argument("--out-dir", default=DEFAULT_OUT_DIR) + args = parser.parse_args() + + out_dir = args.out_dir or args.da_dir + os.makedirs(out_dir, exist_ok=True) + + csv_path = os.path.join(args.da_dir, CAT_ID, f"{CAT_ID}_test_results.csv") + if not os.path.exists(csv_path): + raise FileNotFoundError(csv_path) + df = pd.read_csv(csv_path) + plot_catchment(CAT_ID, df, out_dir) + print("Done.") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/4_evaluation/4d_routed_gauge_check/plot_routed_gauge_check_f5.py b/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/4_evaluation/4d_routed_gauge_check/plot_routed_gauge_check_f5.py new file mode 100644 index 00000000..0c2f8797 --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/folder5_rekrig_variance_direct/4_evaluation/4d_routed_gauge_check/plot_routed_gauge_check_f5.py @@ -0,0 +1,122 @@ +""" +plot_routed_gauge_check_f5.py + +Gauge-level evaluation for F5 (re-kriged variance, 20% holdout). + +Reads routed_Q_test.csv (T-Route output from Step 2) and plots +DA-routed streamflow vs USGS observed discharge at gauge 03463300. +Units are m³/s — directly comparable to the USGS gauge reading. + +Two panels: + Panel 1 : full test period (Oct 2023 – Oct 2024) + Panel 2 : Helene zoom (Sep 20 – Oct 5 2024) + +NSE and KGE reported in title. + +Run on server: + python3 plot_routed_gauge_check_f5.py + python3 plot_routed_gauge_check_f5.py --routed-csv /path/to/routed_Q_test.csv +""" + +import argparse +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt + +DEFAULT_ROUTED_CSV = "/mnt/disk2/suma_helen_poster/da_results/folder5_rekrig_variance_direct/routed_Q_test.csv" +DEFAULT_OUT_DIR = "/mnt/disk2/suma_helen_poster/da_results/folder5_rekrig_variance_direct" + +HELENE_ZOOM_START = pd.Timestamp("2024-09-20") +HELENE_ZOOM_END = pd.Timestamp("2024-10-05") +HELENE_PEAK_START = pd.Timestamp("2024-09-24") +HELENE_PEAK_END = pd.Timestamp("2024-09-30") + +DA_COLOR = "tab:purple" +OBS_COLOR = "black" + + +def nse(obs, sim): + mask = ~(np.isnan(obs) | np.isnan(sim)) + o, s = obs[mask], sim[mask] + if len(o) < 2: + return float("nan") + denom = ((o - o.mean()) ** 2).sum() + return float(1.0 - ((o - s) ** 2).sum() / denom) if denom > 0 else float("nan") + + +def kge(obs, sim): + mask = ~(np.isnan(obs) | np.isnan(sim)) + o, s = obs[mask], sim[mask] + if len(o) < 2: + return float("nan") + r = float(np.corrcoef(o, s)[0, 1]) + alpha = float(s.std() / o.std()) if o.std() > 0 else float("nan") + beta = float(s.mean() / o.mean()) if o.mean() > 0 else float("nan") + return 1.0 - float(np.sqrt((r - 1) ** 2 + (alpha - 1) ** 2 + (beta - 1) ** 2)) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--routed-csv", default=DEFAULT_ROUTED_CSV) + parser.add_argument("--out-dir", default=DEFAULT_OUT_DIR) + args = parser.parse_args() + + os.makedirs(args.out_dir, exist_ok=True) + + df = pd.read_csv(args.routed_csv) + df["date"] = pd.to_datetime(df["date"]) + df = df.set_index("date").sort_index() + + sim_col = next(c for c in df.columns if "routed" in c.lower() or "sim" in c.lower()) + obs_col = next(c for c in df.columns if "usgs" in c.lower() or "obs" in c.lower()) + + sim = df[sim_col].values.astype(float) + obs = df[obs_col].values.astype(float) + dates = df.index + + nse_full = nse(obs, sim) + kge_full = kge(obs, sim) + + helene_mask = (dates >= HELENE_ZOOM_START) & (dates <= HELENE_ZOOM_END) + nse_helene = nse(obs[helene_mask], sim[helene_mask]) + + fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(18, 9)) + fig.suptitle( + f"Gauge 03463300 | F5 routed DA output vs USGS (re-kriged σ², 20% holdout)\n" + f"NSE = {nse_full:.3f} KGE = {kge_full:.3f} (full year) " + f"NSE = {nse_helene:.3f} (Helene zoom)", + fontsize=12, + ) + + # ── Panel 1: full period ────────────────────────────────────────── + ax1.plot(dates, sim, color=DA_COLOR, lw=1.2, label="F5 DA routed [m³/s]") + ax1.plot(dates, obs, color=OBS_COLOR, lw=0.8, alpha=0.7, label="USGS obs [m³/s]") + ax1.set_ylabel("Discharge (m³/s)", fontsize=10) + ax1.set_title("Full test period: Oct 2023 – Oct 2024", fontsize=10) + ax1.legend(fontsize=9, loc="upper left") + ax1.grid(True, alpha=0.2) + + # ── Panel 2: Helene zoom ────────────────────────────────────────── + d_h = dates[helene_mask] + sim_h = sim[helene_mask] + obs_h = obs[helene_mask] + + ax2.axvspan(HELENE_PEAK_START, HELENE_PEAK_END, color="#ffcccc", alpha=0.35, zorder=0) + ax2.plot(d_h, sim_h, color=DA_COLOR, lw=2.0, label="F5 DA routed [m³/s]") + ax2.plot(d_h, obs_h, color=OBS_COLOR, lw=1.5, label="USGS obs [m³/s]") + ax2.set_ylabel("Discharge (m³/s)", fontsize=10) + ax2.set_title("Helene zoom: Sep 20 – Oct 5 2024 (pink = Sep 24–30)", fontsize=10) + ax2.legend(fontsize=9, loc="upper left") + ax2.grid(True, alpha=0.2) + + plt.tight_layout() + out_path = os.path.join(args.out_dir, "cat-1016300_routed_gauge_check_f5.png") + plt.savefig(out_path, dpi=150, bbox_inches="tight") + plt.close() + print(f"Saved: {out_path}") + print(f"NSE={nse_full:.3f} KGE={kge_full:.3f} (full year) NSE={nse_helene:.3f} (Helene)") + + +if __name__ == "__main__": + main() diff --git a/da_methods/test_20pct_heldout_gauges/route_openloop_20pct.sh b/da_methods/test_20pct_heldout_gauges/route_openloop_20pct.sh new file mode 100644 index 00000000..26b24acd --- /dev/null +++ b/da_methods/test_20pct_heldout_gauges/route_openloop_20pct.sh @@ -0,0 +1,27 @@ +#!/usr/bin/env bash +# Route open-loop 20% holdout outputs through T-route to gauge 03463300. +# Run AFTER batch_run_openloop_20pct.sh completes. +# +# Usage: +# bash route_openloop_20pct.sh + +set -euo pipefail + +TROUTE=/home/svyas/miniconda3/envs/troute/bin/python3 +ROUTE_SCRIPT=/home/svyas/cfe_DA_bmi/da_methods/2_troute_routing/run_route.py + +OL_DIR=/mnt/disk2/suma_helen_poster/da_results/openloop_20pct_heldout +GPKG=/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg +USGS_CSV=/mnt/disk2/suma_helen_poster/03463300_usgs_hourly_2018_2024.csv + +echo "[OL-20pct] T-route routing..." +echo " da-dir : $OL_DIR" +echo " out : $OL_DIR/routed_Q_test.csv" + +$TROUTE "$ROUTE_SCRIPT" \ + --gpkg "$GPKG" \ + --da-dir "$OL_DIR" \ + --out-dir "$OL_DIR" \ + --usgs-csv "$USGS_CSV" + +echo "[OL-20pct] Done." diff --git a/example_run_da_CFE/bmi_cfe.py b/example_run_da_CFE/bmi_cfe.py index 9de06a8f..f0f6a312 100644 --- a/example_run_da_CFE/bmi_cfe.py +++ b/example_run_da_CFE/bmi_cfe.py @@ -699,7 +699,7 @@ def set_value(self, var_name, value): src : array_like Array of new values. """ - # JMFrame -- Fixing a slight issue with the self._var_name_units_map + # fix: _var_name_units_map lookup # This is a temporary fix (20230703), # but a permanent solution would be to figure out how to use # get_var_name and setattr with dictionaries. @@ -725,11 +725,11 @@ def set_value_at_indices(self, name, inds, src): # val = self.get_value_ptr(name) # val.flat[inds] = src - #JMFrame: chances are that the index will be zero, so let's include that logic + # index may be zero; include that logic if np.array(self.get_value(name)).flatten().shape[0] == 1: self.set_value(name, src) else: - # JMFrame: Need to set the value with the updated array with new index value + # set value with updated array at new index val = self.get_value_ptr(name) for i in inds.shape: val.flatten()[inds[i]] = src[i] @@ -747,7 +747,7 @@ def get_var_nbytes(self, long_var_name): int Size of data array in bytes. """ - # JMFrame NOTE: Had to import sys for this function + # sys imported for this function return sys.getsizeof(self.get_value_ptr(long_var_name)) #------------------------------------------------------------ @@ -766,7 +766,7 @@ def get_value_at_indices(self, var_name, dest, indices): array_like Values at indices. """ - #JMFrame: chances are that the index will be zero, so let's include that logic + # index may be zero; include that logic if np.array(self.get_value(var_name)).flatten().shape[0] == 1: return self.get_value(var_name) else: diff --git a/example_run_da_CFE/bmi_cfe_perturb_ens.py b/example_run_da_CFE/bmi_cfe_perturb_ens.py index 9c305162..35ff1404 100644 --- a/example_run_da_CFE/bmi_cfe_perturb_ens.py +++ b/example_run_da_CFE/bmi_cfe_perturb_ens.py @@ -899,11 +899,11 @@ def set_value_at_indices(self, name, inds, src): # val = self.get_value_ptr(name) # val.flat[inds] = src - #JMFrame: chances are that the index will be zero, so let's include that logic + # index may be zero; include that logic if np.array(self.get_value(name)).flatten().shape[0] == 1: self.set_value(name, src) else: - # JMFrame: Need to set the value with the updated array with new index value + # set value with updated array at new index val = self.get_value_ptr(name) for i in inds.shape: val.flatten()[inds[i]] = src[i] @@ -921,7 +921,7 @@ def get_var_nbytes(self, long_var_name): int Size of data array in bytes. """ - # JMFrame NOTE: Had to import sys for this function + # sys imported for this function return sys.getsizeof(self.get_value_ptr(long_var_name)) #------------------------------------------------------------ @@ -940,7 +940,7 @@ def get_value_at_indices(self, var_name, dest, indices): array_like Values at indices. """ - #JMFrame: chances are that the index will be zero, so let's include that logic + # index may be zero; include that logic if np.array(self.get_value(var_name)).flatten().shape[0] == 1: return self.get_value(var_name) else: