add Experiment folder: EnKF DA scripts (run_cfe_enkf_wide_spread, for…

Experiment/README.md

@@ -0,0 +1,85 @@
+## Setup — build the environment (once)
+Requires **Python 3.10 with t-route**. 
+`troute-network` / `troute-routing` / `troute-config` and `bmipy`.
+```bash
+conda create -n t-route python=3.10 -y && conda activate t-route
+pip install numpy==1.26.4 pandas==2.2.0 pyarrow matplotlib netcdf4 pyyaml bmipy
+pip install -r requirements.txt                 # repo extras, run from the repo root
+
+# t-route routing engine — clone and build per its README:
+#   https://github.com/NOAA-OWP/t-route
+```
+The CFE model code is needed for step 1 (the EnKF/forecast). On this server it is at
+`/mnt/disk2/suma_helen_poster/cfe_py`.
+
+## Each shell (before running)
+```bash
+conda activate t-route
+PY=python                                       # your t-route env's python
+export MPLBACKEND=Agg OMP_NUM_THREADS=1
+cd <repo>/da_methods/forecast_experiment        # where you cloned it
+```
+
+> Paths below are gauge **03463300** (also the built-in defaults). For a different gauge,
+> change the paths in each command, and edit `CATS` / `TERMINAL_INT` / `WATERSHED_AREA_KM2`
+> at the top of `route_forecast.py`.
+
+## Step 1 — EnKF + 18 h forecasts (DA + open loop) **and** member_q trajectories, all 21 catchments
+Writes per-catchment forecast parquets to `--out-dir/<cat>/`, and (via `--save-members`) the
+per-member runoff trajectories to `--da-traj-dir` / `--ol-traj-dir` (read by the routing warm-up).
+Assimilation always uses the **NWM** operational forcing (`--test-forcing-dir1/2`); the 18 h
+forecast leads use **HRRR** (`--forecast-forcing-dir`) — drop it to fall back to NWM perfect-forcing.
+```bash
+$PY run_cfe_enkf_forecast_wide_spread.py \
+    --cat-id all --n-members 600 --rng-seed 42 --lead-hours 18 \
+    --issue-start "2024-09-20 00:00:00" --issue-end "2024-09-29 23:00:00" \
+    --obs-dir      /mnt/disk2/1400_sites_helene/catchment_ts_03463300_dynamic_variance_rekrig \
+    --params-dir   /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 \
+    --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 \
+    --forecast-forcing-dir /mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/output_03463300_hrrr \
+    --save-members \
+    --da-traj-dir  /mnt/disk2/1400_sites_helene/da_results_enkf600_wide_spread_traj_da \
+    --ol-traj-dir  /mnt/disk2/1400_sites_helene/da_results_enkf600_wide_spread_traj_ol \
+    --out-dir      /mnt/disk2/1400_sites_helene/da_results_enkf600_forecast18h_hrrr
+```
+
+## Step 2 — route every member to the gauge
+Writes `routed_forecast_leadtime_{da,openloop}.parquet`
+(columns: issue_time, lead_hour, member_0000..0599; discharge in m³/s) to `--out-dir`.
+The warm-up trajectories are the NWM-assimilation `member_q` (same for NWM or HRRR forecasts).
+```bash
+$PY route_forecast.py \
+    --scenario both --warmup-mode trajectory --warmup-h 48 --procs 20 \
+    --gpkg        /mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/gage-03463300_subset.gpkg \
+    --fc-dir      /mnt/disk2/1400_sites_helene/da_results_enkf600_forecast18h_hrrr \
+    --da-traj-dir /mnt/disk2/1400_sites_helene/da_results_enkf600_wide_spread_traj_da \
+    --ol-traj-dir /mnt/disk2/1400_sites_helene/da_results_enkf600_wide_spread_traj_ol \
+    --out-dir     /mnt/disk2/1400_sites_helene/da_results_enkf600_forecast18h_hrrr/routed
+```
+
+---
+- To run unattended/in the background, prefix a command with `nohup ... > run.log 2>&1 &`.
+
+## CFE parameters — where they come from / how to re-calibrate
+
+The `--params-dir` in Step 1 points to pre-calibrated CFE best-parameter JSON files
+(`<cat-id>_best_params.json`, one per catchment).  Two sets are available:
+
+| Set | Calibration | Use when |
+|-----|-------------|----------|
+| `held_in_gauge/params/` | gauge 03463300 **included** in calibration | default — best skill at the target gauge |
+| `held_out_gauge/params/` | gauge 03463300 **withheld** | evaluating out-of-sample generalization |
+
+Both sets live in the companion calibration repo.  To **re-run or extend calibration**,
+follow the workflow in:
+
+> **[calibrate-cfe PR #25](https://github.com/NWC-CUAHSI-Summer-Institute/calibrate-cfe/pull/25)**
+> (`multicatchment-calibration-held-in-out` branch)
+
+That repo uses DDS (Dynamically Dimensioned Search, 1 000 iterations via spotpy) against
+the 2020–2022 retrospective period and writes one `<cat-id>_best_params.json` per catchment.
+Clone it, checkout the branch above, and follow its README.
+

Experiment/route_forecast.py

@@ -0,0 +1,294 @@
+#!/usr/bin/env python3
+import argparse
+import os
+import sqlite3
+import sys
+import time
+import types as _types
+from functools import partial
+from multiprocessing import Pool
+
+import numpy as np
+import pandas as pd
+
+# coverage.types stub -- required so numba / t-route import under Python 3.10.
+# MUST run before troute is imported.
+_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
+
+# ---------------------------------------------------------------------------
+# Gauge / network constants (GAUGE-SPECIFIC -- change these to swap the gauge)
+# ---------------------------------------------------------------------------
+TERMINAL_INT       = 1016283     # outlet segment wb-<id> = gauge 03463300
+WATERSHED_AREA_KM2 = 113.18      # total basin area (sum of the divides)
+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",
+]
+
+# Routing constants (not gauge-specific)
+DT               = 3600.0        # 1-hour timestep (seconds)
+QTS_SUBDIVISIONS = 1
+
+# Default I/O (all overridable on the CLI)
+FC_DIR      = "/mnt/disk2/1400_sites_helene/da_results_enkf600_forecast18h_wide_spread"
+DA_TRAJ_DIR = "/mnt/disk2/1400_sites_helene/da_results_enkf600_wide_spread_traj_da"
+OL_TRAJ_DIR = "/mnt/disk2/1400_sites_helene/da_results_enkf600_wide_spread_traj_ol"
+OUT_DIR     = "/mnt/disk2/1400_sites_helene/da_results_enkf600_forecast18h_wide_spread/routed"
+GPKG        = ("/mnt/disk1/usgs_streamflow_allgauges/subdaily_15min/test/"
+               "gage-03463300_subset.gpkg")
+
+# globals shared with fork()ed workers (copy-on-write, read-only)
+G = {}
+
+
+# ===========================================================================
+# 1. CHANNEL NETWORK  (build the Muskingum-Cunge network from the GeoPackage)
+# ===========================================================================
+def seg_int(wb_id):
+    """'wb-1016283' -> 1016283."""
+    return int(wb_id[3:])
+
+
+def read_network(gpkg):
+    """Read the flowpath-attributes (channel params + topology) and flowpaths
+    (divide_id + area) tables from the GeoPackage."""
+    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):
+    """upstream segment -> [downstream segment] (empty list at the outlet)."""
+    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):
+    """Per-segment Muskingum-Cunge channel parameters, indexed by seg_id."""
+    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()]
+    return pd.DataFrame(rows).set_index("seg_id").sort_index().astype("float32")
+
+
+def build_network(gpkg):
+    """Assemble everything route_flow() needs: reach decomposition, upstream map,
+    param table, terminal position, per-catchment segment position, and the
+    mm/h -> m^3/s factor per catchment (= divide_area_m2 / 1000 / 3600)."""
+    fp_attr, fp = read_network(gpkg)
+    connections = build_connections(fp_attr)
+    rconn = nhd_network.reverse_network(connections)
+    reach_list = nhd_network.dfs_decomposition(
+        rconn, partial(nhd_network.split_at_junction, rconn))
+    param_df = build_param_df(fp_attr)
+    seg_ids = 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-")}
+    net = {
+        "reaches_wTypes": [(r, 0) for r in reach_list],
+        "upstreams": dict(rconn),
+        "param_df": param_df,
+        "n_segs": len(param_df),
+        "terminal_pos": int(np.where(seg_ids == TERMINAL_INT)[0][0]),
+        "seg_pos": {int(c[4:]): int(np.where(seg_ids == int(c[4:]))[0][0]) for c in CATS},
+        "factor": {int(c[4:]): area_map[int(c[4:])] / 1000.0 / 3600.0 for c in CATS},
+        "q0_df": pd.DataFrame(np.zeros((len(param_df), 3), dtype="float32"),
+                              index=param_df.index, columns=["qu0", "qd0", "h0"]),
+    }
+    return net
+
+
+# ===========================================================================
+# 2. ROUTING  (one member through the whole network)
+# ===========================================================================
+def route_flow(net, qlat_arr, nts):
+    """Route one member's lateral inflow (n_segs x nts, m^3/s) through the network
+    and return the terminal FLOW (m^3/s) for all nts steps.
+    compute_network_structured returns (n_segs, nts*3) interleaved
+    (flow, velocity, depth) per timestep; flow is the 0::3 stride."""
+    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")
+    res = compute_network_structured(
+        nts, DT, QTS_SUBDIVISIONS,
+        net["reaches_wTypes"], net["upstreams"],
+        net["param_df"].index.values.astype("int64"),
+        net["param_df"].columns.values,
+        net["param_df"].values,
+        net["q0_df"].values.astype("float32"),
+        qlat_arr.astype("float32"),
+        [], e00f64, {}, e00i32, False,
+        "2024-09-20_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,
+    )
+    fvd = np.asarray(res[1])
+    return fvd[net["terminal_pos"], 0::3]
+
+
+def route_issue(k):
+    """Route all members for issue index k. Returns (k, (n_leads, n_members)).
+
+    Builds qlat (n_members, n_segs, warmup+leads): for each catchment segment, the
+    warm-up block (real antecedent runoff in 'trajectory' mode, or a constant
+    lead-1 proxy) followed by the 18 forecast leads; routes each member; slices
+    the warm-up off, keeping the leads on a channel that's already full."""
+    net = G["net"]
+    n_segs, n_members = net["n_segs"], G["n_members"]
+    n_leads, warmup, mode = G["n_leads"], G["warmup_h"], G["warmup_mode"]
+    n_route = warmup + n_leads
+    qlat = np.zeros((n_members, n_segs, n_route), dtype="float32")
+
+    for sid in G["sids"]:
+        pos = net["seg_pos"][sid]
+        f = net["factor"][sid]
+        leads = G["fc"][sid][k].T * f                    # (n_members, n_leads) m^3/s
+        if warmup > 0:
+            if mode == "trajectory":
+                t_idx = G["traj_t0_idx"][k]              # row of t0 in the trajectory
+                qlat[:, pos, :warmup] = G["traj"][sid][t_idx - warmup + 1:t_idx + 1].T * f
+            else:  # constant lead-1 proxy
+                qlat[:, pos, :warmup] = leads[:, 0:1]
+        qlat[:, pos, warmup:] = leads
+
+    out = np.empty((n_leads, n_members), dtype="float32")
+    for m in range(n_members):
+        Q = route_flow(net, qlat[m], n_route)
+        out[:, m] = Q[warmup:warmup + n_leads]
+    return k, out
+
+
+# ===========================================================================
+# 3. SCENARIO DRIVER  (load forecasts/trajectories, route, save)
+# ===========================================================================
+def run_scenario(tag, fc_name, traj_dir, net, args):
+    t0 = time.time()
+    print(f"[{tag}] loading forecasts ({args.warmup_mode} warm-up) for {len(CATS)} catchments...")
+    member_cols, issues, n_leads = None, None, None
+    fc, traj, traj_times = {}, {}, None
+    for cat in CATS:
+        sid = int(cat[4:])
+        f = pd.read_parquet(os.path.join(args.fc_dir, cat,
+                                         f"{cat}_forecast_leadtime_{fc_name}.parquet"))
+        if member_cols is None:
+            member_cols = [c for c in f.columns if c.startswith("member_")]
+            issues = sorted(f["issue_time"].unique().tolist())
+            n_leads = int(f["lead_hour"].max())
+            if args.max_issues > 0:
+                issues = issues[:args.max_issues]
+        sub = f[f["issue_time"].isin(issues)].sort_values(["issue_time", "lead_hour"])
+        # (n_issues, n_leads, n_members) -- member columns are the SAME order in every catchment
+        fc[sid] = (sub[member_cols].values.astype("float32")
+                   .reshape(len(issues), n_leads, len(member_cols)))
+        if args.warmup_mode == "trajectory":
+            tr = pd.read_parquet(os.path.join(traj_dir, cat, f"{cat}_member_q.parquet"))
+            if traj_times is None:
+                traj_times = tr["time"].astype(str).tolist()
+            traj[sid] = tr[member_cols].values.astype("float32")   # (n_hours, n_members)
+
+    G.update({
+        "net": net, "fc": fc, "sids": [int(c[4:]) for c in CATS],
+        "n_members": len(member_cols), "n_leads": n_leads,
+        "warmup_h": args.warmup_h, "warmup_mode": args.warmup_mode,
+    })
+    if args.warmup_mode == "trajectory":
+        t_row = {t: i for i, t in enumerate(traj_times)}
+        missing = [t for t in issues if t not in t_row]
+        if missing:
+            raise SystemExit(f"[{tag}] {len(missing)} issue times absent from {traj_dir} "
+                             f"member_q (e.g. {missing[:2]}); cannot trajectory-warm-up.")
+        G["traj"] = traj
+        G["traj_t0_idx"] = [t_row[t] for t in issues]
+
+    print(f"[{tag}] routing {len(issues)} issues x {len(member_cols)} members "
+          f"({args.warmup_h}h {args.warmup_mode} warm-up + {n_leads} leads each), "
+          f"{args.procs} procs...")
+    results = {}
+    with Pool(args.procs) as pool:
+        for i, (k, out) in enumerate(pool.imap_unordered(route_issue,
+                                                         range(len(issues)), chunksize=2)):
+            results[k] = out
+            if (i + 1) % 24 == 0:
+                print(f"  [{tag}] {i + 1}/{len(issues)} issues ({time.time() - t0:.0f}s)")
+
+    meta, rows = [], []
+    for k, t in enumerate(issues):
+        for lead in range(1, n_leads + 1):
+            meta.append((t, lead))
+            rows.append(results[k][lead - 1])
+    out_df = pd.concat([pd.DataFrame(meta, columns=["issue_time", "lead_hour"]),
+                        pd.DataFrame(np.vstack(rows), columns=member_cols)], axis=1)
+    os.makedirs(args.out_dir, exist_ok=True)
+    pq = os.path.join(args.out_dir, f"routed_forecast_leadtime_{fc_name}.parquet")
+    out_df.to_parquet(pq, index=False)
+    print(f"[{tag}] saved {pq} ({len(out_df)} rows, {len(member_cols)} members) "
+          f"in {time.time() - t0:.0f}s")
+    G.pop("traj", None)
+    G.pop("traj_t0_idx", None)
+
+
+# ===========================================================================
+# 4. CLI
+# ===========================================================================
+def main():
+    p = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    p.add_argument("--scenario", default="both", choices=["da", "openloop", "both"])
+    p.add_argument("--warmup-mode", default="trajectory", choices=["trajectory", "constant"],
+                   help="trajectory = actual antecedent member_q (correct); "
+                        "constant = hold lead-1 inflow (no trajectory files needed)")
+    p.add_argument("--warmup-h", type=int, default=48)
+    p.add_argument("--procs", type=int, default=20)
+    p.add_argument("--max-issues", type=int, default=0,
+                   help="route only the first N issue times (0 = all; for smoke tests)")
+    p.add_argument("--gpkg", default=GPKG)
+    p.add_argument("--fc-dir", default=FC_DIR)
+    p.add_argument("--da-traj-dir", default=DA_TRAJ_DIR)
+    p.add_argument("--ol-traj-dir", default=OL_TRAJ_DIR)
+    p.add_argument("--out-dir", default=OUT_DIR)
+    args = p.parse_args()
+
+    print("Building channel network from GPKG...")
+    net = build_network(args.gpkg)
+    print(f"  {net['n_segs']} segments | terminal wb-{TERMINAL_INT} at pos {net['terminal_pos']}")
+
+    if args.scenario in ("da", "both"):
+        run_scenario("DA", "da", args.da_traj_dir, net, args)
+    if args.scenario in ("openloop", "both"):
+        run_scenario("OL", "openloop", args.ol_traj_dir, net, args)
+
+
+if __name__ == "__main__":
+    main()