Hydrological Modeling: NCRS-PDM¶

cfg = Config(**config)
cfg.get_data()

geoms = gpd.read_parquet(cfg.file_paths.aoi_parquet)
area_sqm = geoms.to_crs(5070).area.iloc[0]
geoms

f"Upstream drainage area: {area_sqm / 1e6:.2f} km²"

nlcd = xr.open_dataset(cfg.file_paths.nlcd[0])
cmap, norm, levels = gh.cover_legends()
cover = nlcd.cover_2016
_ = cover.plot.imshow(cmap=cmap, levels=levels, cbar_kwargs={"ticks": levels[:-1]})

nlcd_meta = gh.helpers.nlcd_helper()
pd.Series(nlcd_meta["classes"])

urban_pct = cover.isin(range(21, 25)).sum() / cover.notnull().sum() * 100
f"{urban_pct.item():.2f} % of the area is urbanized"

snow_fraction = data.snow.resample(time="1YE").mean() / data.prcp.resample(time="1YE").mean()
snow_fraction = snow_fraction.mean().item() * 100
f"{snow_fraction:.1f}% of the precipitation is snow"

qobs_ds = xr.open_dataset(cfg.file_paths.streamflow[0])
qobs = qobs_ds.sel(station_id=f"USGS-{geoms.iloc[0].gage_id}").to_dataframe().discharge
qobs.index = pd.to_datetime(qobs.index.date)
qobs = qobs.loc[qobs.index.intersection(data.time.values)]
qobs

runoff = np.zeros_like(prcp)
dv_soil = np.zeros_like(prcp)
src_soil = np.zeros_like(prcp)

dt = 60 * 60 * 24
lpt2cms = 1e-3 * area_sqm / dt
ms, sb, a = 0.5, 1000, 1

s_soil, s_snow = np.float64(0), np.float64(0)
t_range = t_rain - t_snow
cerr_max = np.float64(0)
for t in range(prcp.shape[0]):
    # Snowmelt (mm/day)
    if tmin[t] > t_rain:
        snm = (s_snow + snow[t]) * ms
    elif tmin[t] < t_snow:
        snm = np.float64(0)
    else:
        snm = (s_snow + snow[t]) * ms * (t_rain - tmin[t]) / t_range

    # Liquid water (mm/day)
    dv_soil[t] = snow[t] - snm
    s_snow = s_snow + dv_soil[t]
    liquid = prcp[t] - snow[t] + snm

    # Soil wetting (mm/day)
    m = s_soil * (2 * sb - a * s_soil) / (2 * sb * (sb - s_soil))
    wetting = np.minimum(
        (
            liquid
            + sb * np.sqrt(np.square(m + 1) - 2 * a * m)
            - np.sqrt(
                np.square(liquid + (m + 1) * sb) - 2 * a * m * np.square(sb) - 2 * a * sb * liquid
            )
        )
        / a,
        liquid,
    )

    # Evapotranspiration (mm/day)
    moisture = wetting + s_soil
    et = np.minimum(
        moisture * (pet[t] + sb - np.sqrt(np.square(pet[t] + sb) - 2 * a * sb * pet[t])) / (a * sb),
        moisture,
    )

    # Soil storage (mm/day)
    ds = wetting - et
    s_soil = s_soil + ds
    dv_soil[t] += ds

    # Total runoff: surface + subsurface (mm/day)
    runoff[t] = liquid - wetting
    src_soil[t] = prcp[t] - et
    cerr_max = np.maximum(cerr_max, np.abs(dv_soil[t] - 0.0 + runoff[t] - src_soil[t]))

runoff_cms = runoff * lpt2cms
results = pd.DataFrame({"runoff (cms)": runoff_cms}, index=data.time.values)
f"Water balance error: {cerr_max:.2e} mm/day"

rc_sim = hs.mean_monthly(results["runoff (cms)"], index_abbr=True, cms=True)
rc_obs = hs.mean_monthly(qobs, index_abbr=True, cms=True)

fig, ax = plt.subplots()
rc_obs.plot(ax=ax, ls="--", label="Observed")
rc_sim.plot(ax=ax, ls="-", label="Simulated")
ax.set_ylabel("Regime Curve (cms)")
ax.set_xlabel("Date")
ax.legend()
_ = ax.set_xmargin(0)

gamma, kb, kd = 0.5, 0.07, 0.57

n_res_nash = 3
nt = runoff.shape[0]
q_sub = np.zeros_like(runoff)
q_sur = np.zeros_like(runoff)

s_sub_new = np.zeros(n_res_nash)
s_sur_new = np.float64(0)
outflow = np.zeros_like(runoff)
inflow = np.zeros_like(runoff)
cerr_max = np.float64(0)
for t in range(1, nt):
    # Subsurface runoff using Nash cascade and linear storage (mm/day)
    dv_sum = dv_soil[t] * lpt2cms
    q_sub_init = runoff[t] * (1.0 - gamma)
    for ii in range(n_res_nash):
        s_sub_old = np.float64(s_sub_new[ii])
        srg = s_sub_old + q_sub_init
        s_sub_new[ii] = (1.0 - kb) * srg
        q_sub_init = kb * srg
        dv_sum += (s_sub_new[ii] - s_sub_old) * lpt2cms

    # Surface runoff (mm/day); inflow should be converted to mm/day
    q_sur[t] = inflow[t] / lpt2cms + runoff[t] * gamma
    s_sur_old = np.float64(s_sur_new)
    srd = s_sur_old + q_sur[t]
    s_sur_new = (1.0 - kd) * srd
    q_sur[t] = kd * srd * lpt2cms
    dv_sum += (s_sur_new - s_sur_old) * lpt2cms

    # Convert to cms so all outputs are in the same unit (cms)
    q_sub[t] = q_sub_init * lpt2cms

    # Compute total outflow
    outflow[t] = q_sur[t] + q_sub[t]
    cerr_max = np.maximum(cerr_max, np.abs(dv_sum - inflow[t] + outflow[t] - src_soil[t] * lpt2cms))

results["outflow (cms)"] = outflow
results["q_sub (cms)"] = q_sub
results["q_sur (cms)"] = q_sur
results["inflow (cms)"] = inflow
f"Water balance error: {cerr_max:.2e} mm/day"

rc_sim = hs.mean_monthly(results["outflow (cms)"], index_abbr=True, cms=True)
rc_obs = hs.mean_monthly(qobs, index_abbr=True, cms=True)

fig, ax = plt.subplots()
rc_obs.plot(ax=ax, ls="--", label="Observed")
rc_sim.plot(ax=ax, ls="-", label="Simulated")
ax.set_ylabel("Regime Curve (cms)")
ax.set_xlabel("Date")
ax.legend()
_ = ax.set_xmargin(0)

def compute_kge(params: FloatArray, *args: Any) -> float:
    """Objective function to minimize."""
    prcp, pet, tmin, snow, qobs_arr, lpt2cms = args
    inflow = np.zeros_like(prcp)
    outflow, *_ = nrcs_pdm(inflow, prcp, pet, tmin, snow, params, lpt2cms)
    return -utils.compute_kge(outflow[365:], qobs_arr[365:])


cal_idx = int(len(prcp) * 0.7)

args = (
    prcp[:cal_idx],
    pet[:cal_idx],
    tmin[:cal_idx],
    snow[:cal_idx],
    utils.to_numpy(qobs)[:cal_idx],
    lpt2cms,
)
bounds = [(0.01, 0.99), (50.0, 2000.0), (0.01, 1.99), (0.01, 0.99), (0.01, 0.14), (0.14, 1)]
result = sceua.minimize(
    compute_kge,
    bounds=bounds,
    args=args,
    seed=42,
    n_complexes=3 * len(bounds),
    pca_freq=10,
    max_workers=8,
)

ms, sb_cal, a, gamma, kb, kd = result.x
print(
    f"ms = {ms:.2f}, sb = {sb_cal:.1f}, a = {a:.2f}, gamma = {gamma:.2f}, kb = {kb:.3f}, kd = {kd:.3f}"
)

inflow = np.zeros_like(prcp)
outflow, q_sur, q_sub = nrcs_pdm(inflow, prcp, pet, tmin, snow, result.x, lpt2cms)
q = pd.DataFrame(
    {
        "Simulated (mm/day)": outflow[cal_idx:],
        "Observed (mm/day)": utils.to_numpy(qobs)[cal_idx:],
        "Surface Runoff (mm/day)": q_sur[cal_idx:],
        "Subsurface Runoff (mm/day)": q_sub[cal_idx:],
    },
    index=qobs.index[cal_idx:],
)
q = q / lpt2cms
hs.plot.signatures(
    q[["Simulated (mm/day)", "Observed (mm/day)"]], title=f"KGE (validation) = {-result.fun:.2f}"
)

fig, ax1 = plt.subplots()

ax1.set_xlabel("No. of Evaluations")
ax1.set_ylabel(r"$\min\ f(x)$")
ax1.plot(np.minimum.accumulate(result.funv))

ax2 = ax1.twinx()
ax2.set_yscale("log")
ax2.set_ylabel(r"$\sum \left[ f(x) - f_\mathrm{min} \right]$")
func_min = -1
ax2.plot(np.cumsum(result.funv - func_min), color="r")
_ = fig.legend(
    ["Convergence", "Cumulative Regret"],
    loc="right",
    bbox_to_anchor=(1, 0.5),
    bbox_transform=ax1.transAxes,
)

sp = ProblemSpec(
    {
        "names": ["ms", "sb", "a", "gamma", "kb", "kd"],
        "num_vars": len(bounds),
        "bounds": bounds,
    }
)
seed = 42
samples = sp.sample_sobol(2**10)
print(f"Number of samples: {len(samples.samples):,}")
kges = Parallel(n_jobs=8, return_as="generator")(
    delayed(compute_kge)(params, prcp, pet, tmin, snow, utils.to_numpy(qobs), lpt2cms)
    for params in samples.samples
)
sp.set_results(np.asarray(list(kges)))
si = sp.analyze_sobol()
total, *_ = si.to_df()
_ = total.plot.barh(y="ST", xerr="ST_conf", capsize=4, legend=False, figsize=(5, 3))

def compute_kge_no_sb(params: FloatArray, *args: Any) -> float:
    """Objective function to minimize."""
    prcp, pet, tmin, snow, sb, qobs_arr, lpt2cms = args
    params_full = np.array([params[0], sb, *params[1:]])
    inflow = np.zeros_like(prcp)
    outflow, *_ = nrcs_pdm(inflow, prcp, pet, tmin, snow, params_full, lpt2cms)
    return -utils.compute_kge(outflow[365:], qobs_arr[365:])


cal_idx = int(len(prcp) * 0.7)

args = (
    prcp[:cal_idx],
    pet[:cal_idx],
    tmin[:cal_idx],
    snow[:cal_idx],
    sb_ap,
    utils.to_numpy(qobs)[:cal_idx],
    lpt2cms,
)
bounds = [(0.01, 0.99), (0.01, 1.99), (0.01, 0.99), (0.01, 0.14), (0.14, 1)]
result = sceua.minimize(
    compute_kge_no_sb,
    bounds=bounds,
    args=args,
    seed=42,
    n_complexes=30 * len(bounds),
    pca_freq=10,
    max_workers=8,
)

params = np.array([result.x[0], sb_ap, *result.x[1:]])
print(
    f"ms = {params[0]:.2f}, sb = {sb_ap:.1f}, a = {params[2]:.2f}, gamma = {params[3]:.2f}, kb = {params[4]:.3f}, kd = {params[5]:.3f}"
)
inflow = np.zeros_like(prcp)
outflow, q_sur, q_sub = nrcs_pdm(inflow, prcp, pet, tmin, snow, params, lpt2cms)
q_no_sb = pd.DataFrame(
    {
        "Simulated (mm/day)": outflow[cal_idx:],
        "Observed (mm/day)": utils.to_numpy(qobs)[cal_idx:],
        "Surface Runoff (mm/day)": q_sur[cal_idx:],
        "Subsurface Runoff (mm/day)": q_sub[cal_idx:],
    },
    index=qobs.index[cal_idx:],
)
q_no_sb = q_no_sb / lpt2cms
hs.plot.signatures(
    q_no_sb[["Simulated (mm/day)", "Observed (mm/day)"]], title=f"KGE = {-result.fun:.2f}"
)

_, axs = plt.subplots(2, 1, figsize=(10, 6), sharex=True, sharey=True)
q_no_sb.loc[q_no_sb.index.year == 2016].plot(ax=axs[0])
axs[0].set_title(f"$S_b$ = {int(sb_ap)} mm (a priori estimate)")
q.loc[q.index.year == 2016].plot(ax=axs[1])
_ = axs[1].set_title(f"$S_b$ = {int(sb_cal)} mm (calibrated)")