Global Sensitivity Analysis of the SAC-SMA Rainfall-Runoff Model¶

import numpy as np
import pandas as pd
import matplotlib
#matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
from scipy.stats.qmc import Sobol
import warnings
import sys
import os
warnings.filterwarnings('ignore')

plt.rcParams.update({
    'font.family': 'serif',
    'font.size': 11,
    'figure.dpi': 120,
    'savefig.dpi': 200,
    'savefig.bbox': 'tight'
})

# ShapleyX
from shapleyx import rshdmr
from shapleyx.utilities.mc_shapley import (
    GaussianCopulaUniform,
    shapley_effects,
)

# SAC-SMA model
data_dir = os.path.join(os.path.dirname(os.path.abspath('')), 'sac_sma_data')
sys.path.insert(0, data_dir)
from sacramento import sacsma

print('All imports OK')

df_forcing = pd.read_csv(f'{data_dir}/st_helens_forcing_original.csv', parse_dates=['date'], dayfirst=True)
print(f'Forcing data: {len(df_forcing)} daily timesteps')
print(f'Date range: {df_forcing.date.iloc[0].date()} to {df_forcing.date.iloc[-1].date()}')

CATCHMENT_AREA_M2 = 120_500_000.0

prcp = df_forcing['rainfall'].values
pet  = df_forcing['pet'].values
q_obs = df_forcing['Q_CUMEC'].values

def flow_mm_to_cumec(flow_mm):
    return flow_mm * CATCHMENT_AREA_M2 / 1000.0 / 86400.0

WARMUP_DAYS = 365
print(f'Warm-up: {WARMUP_DAYS} days')
print(f'Evaluation: day {WARMUP_DAYS+1} to {len(df_forcing)} ({len(df_forcing)-WARMUP_DAYS} timesteps)')

def nse_from_params(params, warmup=WARMUP_DAYS):
    flow_mm = sacsma(prcp, pet, params)
    sim = flow_mm_to_cumec(flow_mm[warmup:])
    obs = q_obs[warmup:]
    residual = obs - sim
    numerator = np.sum(residual ** 2)
    denominator = np.sum((obs - np.mean(obs)) ** 2)
    return 1.0 - numerator / denominator if denominator > 0 else -np.inf

# Warm-up (triggers Numba JIT)
params_default = np.array([44.13, 52.76, 80.87, 112.70, 52.81, 0.084, 0.490,
                           0.0148, 0.133, 35.31, 2.532, 0.00578, 0.424, 0.0093])
nse_default = nse_from_params(params_default)
print(f'NSE at default parameters: {nse_default:.4f}')

N_PARAMS = 14
N_SAMPLES = 4096

param_names = ['Uztwm', 'Uzfwm', 'Lztwm', 'Lzfpm', 'Lzfsm',
               'Adimp', 'Uzk', 'Lzpk', 'Lzsk', 'Zperc',
               'Rexp', 'Pctim', 'Pfree', 'Side']

param_lower = np.array([30, 10, 20, 40, 15, 0, 0.2, 0.001, 0.03, 0, 0, 0, 0, 0])
param_upper = np.array([125, 75, 300, 600, 300, 0.2, 0.9, 0.03, 0.8, 80, 5, 0.05, 0.8, 0.8])
param_ranges = param_upper - param_lower

print(f'Parameter space: {N_PARAMS} dimensions x {N_SAMPLES} Sobol samples')
for i, name in enumerate(param_names):
    print(f'  {name:>6s}: [{param_lower[i]:8.3f}, {param_upper[i]:8.3f}]')

sobol_gen = Sobol(d=N_PARAMS, scramble=True, seed=42)
X_uniform = sobol_gen.random(n=N_SAMPLES)

# Cranley-Patterson rotation
rng = np.random.default_rng(123)
shift = rng.uniform(0, 1, size=N_PARAMS)
X_uniform = (X_uniform + shift) % 1.0

X_params = param_lower + X_uniform * param_ranges
print(f'Sobol samples shape: {X_params.shape}')
print(f'First 3 samples:')
for i in range(min(3, N_SAMPLES)):
    print(f'  [{i}] ' + ', '.join(f'{v:.2f}' for v in X_params[i]))

print('Evaluating 4096 parameter sets...')
Y_nse = np.array([nse_from_params(X_params[i]) for i in range(N_SAMPLES)])

valid = Y_nse > -np.inf
print(f'NSE distribution ({valid.sum()} valid samples):')
Yv = Y_nse[valid]
print(f'  Mean:   {Yv.mean():.4f}')
print(f'  Std:    {Yv.std():.4f}')
print(f'  Min:    {Yv.min():.4f}')
print(f'  Max:    {Yv.max():.4f}')
print(f'  Median: {np.median(Yv):.4f}')
print(f'  NSE > 0: {(Yv > 0).mean()*100:.1f}% (better than mean-flow benchmark)')

fig, ax = plt.subplots(figsize=(8, 4))
ax.hist(Yv, bins=40, color='steelblue', edgecolor='white', alpha=0.8)
ax.axvline(Yv.mean(), color='firebrick', ls='--', lw=1.5, label=f'Mean = {Yv.mean():.3f}')
ax.axvline(0, color='grey', ls=':', lw=1, label='NSE = 0')
ax.set_xlabel('NSE')
ax.set_ylabel('Frequency')
ax.set_title('Distribution of NSE Across 4096 Sobol Samples')
ax.legend(fontsize=9)
plt.tight_layout()
plt.savefig('nse_distribution.eps')
plt.savefig('nse_distribution.svg')
plt.show()

df_train = pd.DataFrame(X_params, columns=param_names)
df_train['Y'] = Y_nse

print('Training RS-HDMR surrogate (streaming OMP-CV)...')
analyzer = rshdmr(
    data_file=df_train,
    polys=[10, 8, 2],
    method='omp_cv',
    n_iter=400,
    test_size=0.25,
    verbose=True,
    resampling=False,
    n_jobs=4,
)

sobol, shap, total = analyzer.run_all()
print('\nSurrogate training complete.')

display_col = 'scaled effect' if 'scaled effect' in shap.columns else 'effect'
print(f'RS-HDMR Shapley Effects — Independent Inputs:')
print(shap.to_string(float_format=lambda x: f'{x:.6f}'))

shap_sorted = shap.sort_values(display_col, ascending=True)
fig, ax = plt.subplots(figsize=(8, 5))
colors = plt.cm.Blues(np.linspace(0.4, 0.9, len(shap_sorted)))
ax.barh(range(len(shap_sorted)), shap_sorted[display_col].values, color=colors)
ax.set_yticks(range(len(shap_sorted)))
ax.set_yticklabels(shap_sorted['label'].values)
ax.set_xlabel('Shapley Effect (fraction of explained variance)')
ax.set_title('RS-HDMR Shapley Effects — Independent Inputs')
ax.axvline(0, color='grey', lw=0.8)
plt.tight_layout()
plt.savefig('shapley_independent.png')
plt.show()

# Vrugt et al. (2006) Table 2 — 13x13 correlation matrix
# Vrugt order: UZTWM, UZFWM, UZK, PCTIM, ADIMP, ZPERC, REXP,
#               LZTWM, LZFSM, LZFPM, LZSK, LZPK, PFREE
corr_vrugt = np.array([
    [ 1.00, -0.02, -0.04,  0.19,  0.36, -0.01,  0.11, -0.84, -0.17,  0.28,  0.17,  0.16,  0.57],
    [-0.02,  1.00, -0.53,  0.05,  0.01,  0.15,  0.11, -0.01, -0.08,  0.22, -0.15, -0.18, -0.33],
    [-0.04, -0.53,  1.00, -0.01, -0.41, -0.13, -0.03,  0.07, -0.11, -0.01, -0.02,  0.08,  0.08],
    [ 0.19,  0.05, -0.01,  1.00, -0.16,  0.01,  0.03, -0.14,  0.01,  0.02, -0.08,  0.02, -0.04],
    [ 0.36,  0.01, -0.41, -0.16,  1.00,  0.07,  0.04, -0.21, -0.15,  0.09,  0.17,  0.15,  0.31],
    [-0.01,  0.15, -0.13,  0.01,  0.07,  1.00,  0.05,  0.06, -0.06, -0.10, -0.13,  0.02,  0.00],
    [ 0.11,  0.11, -0.03,  0.03,  0.04,  0.05,  1.00, -0.03,  0.65,  0.68,  0.24, -0.10, -0.12],
    [-0.84, -0.01,  0.07, -0.14, -0.21,  0.06, -0.03,  1.00,  0.28, -0.16, -0.14, -0.23, -0.63],
    [-0.17, -0.08, -0.11,  0.01, -0.15, -0.06,  0.65,  0.28,  1.00,  0.41, -0.33, -0.55, -0.36],
    [ 0.28,  0.22, -0.01,  0.02,  0.09, -0.10,  0.68, -0.16,  0.41,  1.00,  0.06, -0.43, -0.26],
    [ 0.17, -0.15, -0.02, -0.08,  0.17, -0.13,  0.24, -0.14, -0.33,  0.06,  1.00,  0.60,  0.41],
    [ 0.16, -0.18,  0.08,  0.02,  0.15,  0.02, -0.10, -0.23, -0.55, -0.43,  0.60,  1.00,  0.56],
    [ 0.57, -0.33,  0.08, -0.04,  0.31,  0.00, -0.12, -0.63, -0.36, -0.26,  0.41,  0.56,  1.00],
])

# Map Vrugt order -> our parameter order
# Ours: Uztwm(0), Uzfwm(1), Lztwm(2), Lzfpm(3), Lzfsm(4), Adimp(5), Uzk(6),
#       Lzpk(7),  Lzsk(8),  Zperc(9), Rexp(10), Pctim(11), Pfree(12), Side(13)
vrugt_to_our = [0, 1, 6, 11, 5, 9, 10, 2, 4, 3, 8, 7, 12]

corr_14 = np.eye(N_PARAMS)
for i, vi in enumerate(vrugt_to_our):
    for j, vj in enumerate(vrugt_to_our):
        corr_14[vi, vj] = corr_vrugt[i, j]

# Ensure symmetry and positive definiteness
corr_14 = (corr_14 + corr_14.T) / 2
eigvals = np.linalg.eigvalsh(corr_14)
if eigvals.min() < 0:
    corr_14 += np.eye(N_PARAMS) * (-eigvals.min() + 1e-8)

print('Correlation matrix:')
pd.set_option('display.precision', 3)
display(pd.DataFrame(corr_14, index=param_names, columns=param_names))

print('\nStrongest correlations:')
for i in range(N_PARAMS):
    for j in range(i+1, N_PARAMS):
        if abs(corr_14[i, j]) > 0.3:
            print(f'  {param_names[i]:>6s} <-> {param_names[j]:<6s}: rho = {corr_14[i, j]:+.3f}')

print('Computing MC Shapley Effects with correlated inputs...')
print(f'  k_max = 3, N = 2000, B = 500\n')

df_mc = analyzer.get_mc_shapley(
    corr=corr_14,
    N=2000,
    method='exhaustive',
    B=500,
    alpha=0.05,
    random_state=42,
    progress=True,
    k_max=3,
)

print('\nMC Shapley Effects (correlated inputs):')
display(df_mc)

N_direct = 1000
B_direct = 100

def sacsma_nse_wrapper(x):
    """SAC-SMA NSE function taking model parameter vector.
    Note: GaussianCopulaUniform handles the [0,1] → physical mapping."""
    return nse_from_params(x)

print('Computing MC Shapley directly on SAC-SMA model...')
print(f'  Distribution: GaussianCopulaUniform (Vrugt correlations)')
print(f'  N = {N_direct}, B = {B_direct}, k_max = 2')
print(f'  Estimated evals: ~{105 * N_direct * 2:,}')

eff_dir, sh_dir, var_dir, lower_dir, upper_dir = shapley_effects(
    sacsma_nse_wrapper,
    joint_correlated,
    N=N_direct,
    method='exhaustive',
    B=B_direct,
    alpha=0.05,
    random_state=42,
    progress=True,
    k_max=2,
)

df_direct = pd.DataFrame({
    'variable': param_names,
    'effect': eff_dir,
    'value': sh_dir,
    'total_var': var_dir,
    'lower': lower_dir,
    'upper': upper_dir,
})
print('\nDone. Top 5 parameters (direct SAC-SMA):')
print(df_direct.sort_values('effect', ascending=False).head(10).to_string(index=False))

# Compute surrogate MC Shapley at same settings for fair comparison
print('Computing surrogate MC Shapley at same N=1000, k_max=2...')
eff_surr, sh_surr, var_surr, lower_surr, upper_surr = shapley_effects(
    analyzer.predict,
    joint_correlated,
    N=N_direct,
    method='exhaustive',
    B=B_direct,
    alpha=0.05,
    random_state=42,
    progress=True,
    k_max=2,
)

df_surrogate = pd.DataFrame({
    'variable': param_names,
    'effect': eff_surr,
    'value': sh_surr,
    'total_var': var_surr,
    'lower': lower_surr,
    'upper': upper_surr,
})

# Merge into comparison table
df_comp_direct = df_direct[['variable', 'effect', 'lower', 'upper']].copy()
df_comp_direct.columns = ['Variable', 'SAC-SMA', 'CI_low', 'CI_high']
df_comp_surr = df_surrogate[['variable', 'effect']].copy()
df_comp_surr.columns = ['Variable', 'Surrogate']

df_compare = df_comp_direct.merge(df_comp_surr, on='Variable')
df_compare['Delta'] = df_compare['SAC-SMA'] - df_compare['Surrogate']
df_compare = df_compare.sort_values('SAC-SMA', ascending=False).reset_index(drop=True)

print('\n=== Direct SAC-SMA vs Surrogate MC Shapley (N=1000, k_max=2) ===')
print(df_compare.to_string(index=False))
print(f'\nSum SAC-SMA: {df_compare["SAC-SMA"].sum():.4f}')
print(f'Sum Surrogate: {df_compare["Surrogate"].sum():.4f}')
print(f'Max |Delta|: {df_compare["Delta"].abs().max():.4f}')

# Bar chart comparison: SAC-SMA vs Surrogate
import matplotlib.pyplot as plt
import numpy as np

N_PLOT = 10
plot_df = df_compare.head(N_PLOT)

fig, ax = plt.subplots(figsize=(10, 5))
x = np.arange(N_PLOT)
width = 0.35

ax.bar(x - width/2, plot_df['SAC-SMA'], width, label='SAC-SMA (direct)',
       color='#4C72B0', edgecolor='white', linewidth=0.5)
ax.bar(x + width/2, plot_df['Surrogate'], width, label='RS-HDMR Surrogate',
       color='#DD8452', edgecolor='white', linewidth=0.5)

ax.set_xticks(x)
ax.set_xticklabels(plot_df['Variable'], rotation=45, ha='right')
ax.set_ylabel('Shapley Effect')
ax.set_title('Direct SAC-SMA vs Surrogate MC Shapley Effects (N=1000, k_max=2)')
ax.legend(frameon=False)
fig.tight_layout()
plt.show()

# Save for the case study report
fig.savefig('anal_surr_comparison.png', dpi=200, bbox_inches='tight')
print('Saved: anal_surr_comparison.png')

# Extract RS-HDMR effects
indep_col = 'scaled effect' if 'scaled effect' in shap.columns else 'effect'
indep_dict = {}
for _, row in shap.iterrows():
    label = row['label']
    # ShapleyX may label like "Uztwm_0" — match by prefix
    matched = False
    for p in param_names:
        if label.startswith(p):
            indep_dict[p] = row[indep_col]
            matched = True
            break
    if not matched:
        indep_dict[label] = row[indep_col]

df_comp = pd.DataFrame({
    'Parameter': param_names,
    'Independent': [indep_dict.get(p, 0) for p in param_names],
    'Correlated': df_mc['effect'].values,
})

if 'lower' in df_mc.columns:
    df_comp['CI_lower'] = df_mc['lower'].values
    df_comp['CI_upper'] = df_mc['upper'].values

# Normalise independent to sum to 1
df_comp['Independent'] /= df_comp['Independent'].sum()
df_comp['Change'] = df_comp['Correlated'] - df_comp['Independent']

df_comp_sorted = df_comp.sort_values('Correlated', ascending=False)
print('Comparison: Independent vs Correlated Shapley Effects')
display(df_comp_sorted.round(4))

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 6))

y = np.arange(len(param_names))
order = np.argsort(df_comp['Correlated'].values)

# Left: side-by-side bars
for col, color, label, offset in [
    ('Independent', 'royalblue', 'Independent (RS-HDMR)', -0.17),
    ('Correlated', 'firebrick', 'Correlated (MC Shapley)', 0.17),
]:
    vals = df_comp[col].values[order]
    ax1.barh(y + offset, vals, 0.3, label=label,
             color=color, alpha=0.75, edgecolor='white', linewidth=0.5)

ax1.set_yticks(y)
ax1.set_yticklabels(np.array(param_names)[order])
ax1.set_xlabel('Shapley Effect')
ax1.set_title('Independent vs Correlated')
ax1.legend(fontsize=9, loc='lower right')
ax1.axvline(0, color='grey', lw=0.8)

# Right: delta
changes = df_comp['Change'].values[order]
clrs = ['firebrick' if c > 0 else 'royalblue' for c in changes]
ax2.barh(y, changes, 0.35, color=clrs, alpha=0.75,
         edgecolor='white', linewidth=0.5)
ax2.axvline(0, color='grey', lw=0.8)
ax2.set_yticks([])
ax2.set_xlabel('Δ (Correlated − Independent)')
ax2.set_title('Change Due to Correlation')

plt.tight_layout()
plt.savefig('shapley_comparison.png')
plt.show()

print('Top 5 parameters by sensitivity:')

print('\n  Under INDEPENDENT inputs (RS-HDMR):')
top5_indep = df_comp_sorted.sort_values('Independent', ascending=False).head(5)
for _, row in top5_indep.iterrows():
    print(f'    {row["Parameter"]:>6s}: {row["Independent"]:.4f}')

print('\n  Under CORRELATED inputs (MC Shapley):')
has_ci = 'CI_lower' in df_comp.columns
for _, row in df_comp_sorted.head(5).iterrows():
    ci = ''
    if has_ci:
        ci = f' [{row["CI_lower"]:.4f}, {row["CI_upper"]:.4f}]'
    print(f'    {row["Parameter"]:>6s}: {row["Correlated"]:.4f}{ci}')

print('\n  Largest positive change (gained importance with correlation):')
print(df_comp_sorted.sort_values('Change', ascending=False).head(3)[['Parameter','Change']].to_string(index=False))

print('\n  Largest negative change (lost importance with correlation):')
print(df_comp_sorted.sort_values('Change', ascending=True).head(3)[['Parameter','Change']].to_string(index=False))