XAJ-SLW Model (Songliao)¶
xaj_slw is a XinAnJiang variant developed for the Songliao basin, using SMS3 and LAG3 routing (ported from hydromodeljava). It is the only registered model whose internal routing works in discharge (m^3/s); UnifiedSimulator converts its qsim back to runoff depth (mm) so results are comparable with other models and observations.
- Registered name:
xaj_slw - Parameters: 26
- Requires:
basin_area(km^2) andtime_interval_hoursat simulation time - Routing: SMS3 + LAG3 (storage-lag-weighted)
API Reference¶
Copyright (c) 2023-2024 Wenyu Ouyang. All rights reserved.
calculate_net_precipitation(precipitation, potential_evapotranspiration, kc)
¶
Calculate net precipitation and actual evapotranspiration.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
precipitation |
ndarray |
Precipitation array. |
required |
potential_evapotranspiration |
ndarray |
Potential evapotranspiration array. |
required |
kc |
ndarray |
Evapotranspiration coefficient array. |
required |
Returns:
| Type | Description |
|---|---|
Tuple[numpy.ndarray, numpy.ndarray] |
Tuple containing net precipitation and actual evapotranspiration (pe, edt). |
Source code in hydromodel/models/xaj_slw.py
@jit(nopython=True)
def calculate_net_precipitation(
precipitation: np.ndarray,
potential_evapotranspiration: np.ndarray,
kc: np.ndarray,
) -> Tuple[np.ndarray, np.ndarray]:
"""Calculate net precipitation and actual evapotranspiration.
Args:
precipitation: Precipitation array.
potential_evapotranspiration: Potential evapotranspiration array.
kc: Evapotranspiration coefficient array.
Returns:
Tuple containing net precipitation and actual evapotranspiration (pe, edt).
"""
edt = kc * potential_evapotranspiration # Actual evapotranspiration
pe = precipitation - edt # Net precipitation
return pe, edt
lag3_routing_vectorized(rs, ri, rg, time_interval, basin_area, ci, cg, lag, cs, kk, x, mp, qsp=0.0, qip=0.0, qgp=0.0, qsig_initial=None, qx_initial=None, return_states=False)
¶
Vectorized LAG3 routing model.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
rs |
ndarray |
Surface runoff. |
required |
ri |
ndarray |
Interflow runoff. |
required |
rg |
ndarray |
Groundwater runoff. |
required |
time_interval |
float |
Time interval. |
required |
basin_area |
float |
Basin area. |
required |
ci |
float |
Interflow recession coefficient. |
required |
cg |
float |
Groundwater recession coefficient. |
required |
lag |
float |
Lag time. |
required |
cs |
float |
Surface runoff recession coefficient. |
required |
kk |
float |
Muskingum K parameter. |
required |
x |
float |
Muskingum X parameter. |
required |
mp |
int |
Number of river reaches. |
required |
qsp |
float |
Initial surface flow. |
0.0 |
qip |
float |
Initial interflow. |
0.0 |
qgp |
float |
Initial groundwater flow. |
0.0 |
qsig_initial |
ndarray |
Initial qsig state array. |
None |
qx_initial |
ndarray |
Initial qx state array. |
None |
return_states |
bool |
Whether to return state variables. |
False |
Returns:
| Type | Description |
|---|---|
Union[numpy.ndarray, Tuple[numpy.ndarray, Dict[str, numpy.ndarray]]] |
If return_states is False, returns final flow array. If return_states is True, returns (flow_array, state_dict) tuple. |
Source code in hydromodel/models/xaj_slw.py
def lag3_routing_vectorized(
rs: np.ndarray,
ri: np.ndarray,
rg: np.ndarray,
time_interval: float,
basin_area: float,
ci: float,
cg: float,
lag: float,
cs: float,
kk: float,
x: float,
mp: int,
qsp: float = 0.0,
qip: float = 0.0,
qgp: float = 0.0,
qsig_initial: np.ndarray = None,
qx_initial: np.ndarray = None,
return_states: bool = False,
) -> Union[np.ndarray, Tuple[np.ndarray, Dict[str, np.ndarray]]]:
"""Vectorized LAG3 routing model.
Args:
rs: Surface runoff.
ri: Interflow runoff.
rg: Groundwater runoff.
time_interval: Time interval.
basin_area: Basin area.
ci: Interflow recession coefficient.
cg: Groundwater recession coefficient.
lag: Lag time.
cs: Surface runoff recession coefficient.
kk: Muskingum K parameter.
x: Muskingum X parameter.
mp: Number of river reaches.
qsp: Initial surface flow.
qip: Initial interflow.
qgp: Initial groundwater flow.
qsig_initial: Initial qsig state array.
qx_initial: Initial qx state array.
return_states: Whether to return state variables.
Returns:
If return_states is False, returns final flow array.
If return_states is True, returns (flow_array, state_dict) tuple.
"""
# === 诊断日志：LAG3 输入参数 ===
logger.info(f"=== LAG3 汇流计算开始 ===")
logger.info(f"rs 类型: {type(rs)}, 形状: {rs.shape if hasattr(rs, 'shape') else 'no shape'}, 值: {rs if not hasattr(rs, 'shape') else rs[:5] if len(rs) > 5 else rs}")
logger.info(f"ri 类型: {type(ri)}, 形状: {ri.shape if hasattr(ri, 'shape') else 'no shape'}")
logger.info(f"rg 类型: {type(rg)}, 形状: {rg.shape if hasattr(rg, 'shape') else 'no shape'}")
logger.info(f"lag 参数: {lag}, cs: {cs}, ci: {ci}, cg: {cg}")
# 防御性检查：确保 rs, ri, rg 是数组
if not hasattr(rs, '__len__'):
raise TypeError(f"rs 必须是数组类型，但收到 {type(rs)}: {rs}")
if not hasattr(ri, '__len__'):
raise TypeError(f"ri 必须是数组类型，但收到 {type(ri)}: {ri}")
if not hasattr(rg, '__len__'):
raise TypeError(f"rg 必须是数组类型，但收到 {type(rg)}: {rg}")
time_steps = len(rs)
t_steps = int(round(lag / time_interval)) # 新的滞时步数计算
t = max(t_steps, 0) # 确保非负值
# 单位转换系数
cp = basin_area / time_interval / 3.6
# 参数时段转换
ci = np.power(ci, time_interval / 24.0)
cg = np.power(cg, time_interval / 24.0)
# 初始化输出数组
qs = np.zeros(time_steps)
qi = np.zeros(time_steps)
qg = np.zeros(time_steps)
qsig = np.zeros(time_steps + t)
# 初始化QSIG数组
if qsig_initial is None:
qsig_initial = np.zeros(max(t, 3))
if lag <= 1:
# 防御性检查：避免空数组导致索引越界
if len(qsig) > 0 and len(qsig_initial) > 0:
qsig[0] = qsig_initial[0]
else:
for i in range(t):
if len(qsig_initial) >= t:
qsig[i] = qsig_initial[i]
else:
if i < len(qsig_initial):
qsig[i] = qsig_initial[i]
else:
qsig[i] = qsig[i - 1]
# 初始化马斯京根参数
fkt = kk - kk * x + 0.5 * time_interval
c0 = (0.5 * time_interval - kk * x) / fkt
c1 = (kk * x + 0.5 * time_interval) / fkt
c2 = (kk - kk * x - 0.5 * time_interval) / fkt
# 初始化QX数组
qx = np.zeros(mp + 1)
if qx_initial is not None:
for i in range(mp + 1):
if i < len(qx_initial):
qx[i] = qx_initial[i]
else:
qx[i] = qx[i - 1] if i > 0 else 0.0
# 主循环计算
qip_curr = qip
qgp_curr = qgp
qsig1 = qsig[t - 1] if lag > 1 else qsig[0]
for i in range(time_steps):
# 计算三水源汇流
qgp_curr = qgp_curr * cg + rg[i] * (1.0 - cg) * cp
qip_curr = qip_curr * ci + ri[i] * (1.0 - ci) * cp
qsp_curr = rs[i] * cp
# 存储结果
qg[i] = qgp_curr
qi[i] = qip_curr
qs[i] = qsp_curr
# 计算总入流并更新QSIG
qsig1 = qsig1 * cs + (qgp_curr + qip_curr + qsp_curr) * (1.0 - cs)
qtsig = qsig1
# 使用LCHCO进行马斯京根演算
qsig[i + t] = lchco_vectorized(mp, qtsig, qx, c0, c1, c2)
# 提取最终结果
q_routing = qsig[:time_steps]
# 确保非负值
q_routing = np.maximum(q_routing, 0.0)
if return_states:
# 返回最终状态
final_states = {
"qsig_final": qsig[time_steps:],
"qx_final": qx.copy(),
"qsig1_final": qsig1,
}
return q_routing, final_states
else:
return q_routing
lchco_vectorized(mp, rq, qx, c0, c1, c2)
¶
Vectorized LCHCO calculation.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
mp |
int |
Number of river reaches. |
required |
rq |
float |
Inflow rate. |
required |
qx |
ndarray |
Flow array. |
required |
c0 |
float |
Muskingum coefficient 0. |
required |
c1 |
float |
Muskingum coefficient 1. |
required |
c2 |
float |
Muskingum coefficient 2. |
required |
Returns:
| Type | Description |
|---|---|
float |
Calculated flow value. |
Source code in hydromodel/models/xaj_slw.py
def lchco_vectorized(
mp: int, rq: float, qx: np.ndarray, c0: float, c1: float, c2: float
) -> float:
"""Vectorized LCHCO calculation.
Args:
mp: Number of river reaches.
rq: Inflow rate.
qx: Flow array.
c0: Muskingum coefficient 0.
c1: Muskingum coefficient 1.
c2: Muskingum coefficient 2.
Returns:
Calculated flow value.
"""
im = mp + 1
if im == 1:
qx[int(im - 1)] = rq
else:
for j in range(1, int(im)):
q1 = rq
q2 = qx[j - 1]
q3 = qx[j]
qx[j - 1] = rq
rq = c0 * q1 + c1 * q2 + c2 * q3
qx[int(im - 1)] = rq
return rq
load_sms_lag_data_from_json(sms_json_path, lag_json_path, default_evap)
¶
Load XAJ model data from SMS and LAG JSON files.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
sms_json_path |
str |
Path to SMS_3 model JSON file. |
required |
lag_json_path |
str |
Path to LAG_3 model JSON file. |
required |
default_evap |
float |
Default evapotranspiration value. |
required |
Returns:
| Type | Description |
|---|---|
Tuple containing |
|
Source code in hydromodel/models/xaj_slw.py
def load_sms_lag_data_from_json(
sms_json_path: str,
lag_json_path: str,
default_evap: float,
) -> Tuple[np.ndarray, np.ndarray, List[str], str]:
"""Load XAJ model data from SMS and LAG JSON files.
Args:
sms_json_path: Path to SMS_3 model JSON file.
lag_json_path: Path to LAG_3 model JSON file.
default_evap: Default evapotranspiration value.
Returns:
Tuple containing:
- p_and_e: Precipitation and evapotranspiration data [time, basin=1, feature=2]
- parameters: Model parameters [basin=1, parameter=26]
- time_dates: List of time dates
- start_time: Start time
- es: Monthly evapotranspiration array
"""
# 读取SMS JSON文件
with open(sms_json_path, "r", encoding="utf-8") as f:
sms_data = json.load(f)
# 读取LAG JSON文件
with open(lag_json_path, "r", encoding="utf-8") as f:
lag_data = json.load(f)
# 解析时间序列和降雨数据
dt = sms_data["dt"]
rain = sms_data["rain"]
start_time = lag_data.get("start", dt[0] if dt else "")
# 解析月蒸发量数组
es = np.array(
sms_data.get(
"ES",
[
23.8,
23.2,
33.7,
53,
66.6,
78.9,
122.3,
113.9,
90.8,
62.9,
46.2,
34.1,
],
)
)
# 构建p_and_e数组 [time, basin=1, feature=2]
time_steps = len(rain)
p_and_e = np.zeros((time_steps, 1, 2))
p_and_e[:, 0, 0] = rain # 降雨数据
# 根据ES数组计算蒸散发值
if "ES" in sms_data:
# 获取时间间隔参数
time_interval = float(sms_data.get("clen", 1.0))
# 计算每个时间步的蒸散发值
evap_values = np.zeros(time_steps)
for i in range(time_steps):
# 从时间字符串中提取月份，默认为8月
try:
if dt and i < len(dt):
time_str = dt[i]
# 尝试解析时间字符串获取月份
if ":" in time_str: # 包含时间的格式
month = int(time_str.split("-")[1]) # 提取月份
else: # 只有日期的格式
month = int(time_str.split("-")[1]) # 提取月份
else:
month = 8 # 默认月份为8
except:
month = 8 # 解析失败时使用默认月份
# 根据月份确定天数
if month in [4, 6, 9, 11]:
iday = 30
elif month == 2:
iday = 28
else:
iday = 31
# 计算蒸散发值：ES[month-1] / (IDAY * 24.0 / T)
em = es[month - 1] / (iday * 24.0 / time_interval)
evap_values[i] = em
else:
# 如果没有ES数组，使用输入的蒸散发数值
evap_values = np.full(time_steps, default_evap)
p_and_e[:, 0, 1] = evap_values
# 构建参数数组 [basin=1, parameter=26]
# 参数顺序: [WUP, WLP, WDP, SP, FRP, WM, WUMx, WLMx, KC, B, C, IM,
# SM, EX, KG, KI, CS, CI, CG, LAG, KK, X, MP, QSP, QIP, QGP]
parameters = np.array(
[
[
float(sms_data["WUP"]), # 0: Initial upper layer tension water
float(sms_data["WLP"]), # 1: Initial lower layer tension water
float(sms_data["WDP"]), # 2: Initial deep layer tension water
float(sms_data["SP"]), # 3: Initial free water storage
float(sms_data["FRP"]), # 4: Initial runoff basin_area ratio
float(sms_data["WM"]), # 5: Total tension water capacity
float(sms_data["WUMx"]), # 6: Upper layer capacity ratio
float(sms_data["WLMx"]), # 7: Lower layer capacity ratio
float(sms_data["K"]), # 8: Evaporation coefficient (KC)
float(
sms_data["B"]
), # 9: Exponent of tension water capacity curve
float(
sms_data["C"]
), # 10: Deep evapotranspiration coefficient
float(sms_data["IM"]), # 11: Impervious basin_area ratio
float(sms_data["SM"]), # 12: Average free water capacity
float(
sms_data["EX"]
), # 13: Exponent of free water capacity curve
float(sms_data["KG"]), # 14: Groundwater outflow coefficient
float(sms_data["KI"]), # 15: Interflow outflow coefficient
float(lag_data["CS"]), # 16: Channel system recession constant
float(
lag_data["CI"]
), # 17: Lower interflow recession constant
float(
lag_data["CG"]
), # 18: Groundwater storage recession constant
float(lag_data["LAG"]), # 19: Lag time
float(lag_data["KK"]), # 20: Muskingum K parameter
float(lag_data["X"]), # 21: Muskingum X parameter
float(lag_data["MP"]), # 22: Number of Muskingum reaches
float(lag_data["QSP"]), # 23: Initial surface flow
float(lag_data["QIP"]), # 24: Initial interflow
float(lag_data["QGP"]), # 25: Initial groundwater flow
]
]
)
return p_and_e, parameters, dt, start_time, es
sms3_runoff_generation_vectorized(precipitation, evapotranspiration, wu, wl, wd, s, fr, wm, wumx, wlmx, kc, b, c, im, sm, ex, kg, ki, time_interval, time_steps)
¶
Vectorized SMS3 runoff generation model.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
precipitation |
ndarray |
Precipitation time series. |
required |
evapotranspiration |
ndarray |
Evapotranspiration time series. |
required |
wu |
ndarray |
Initial upper layer tension water content. |
required |
wl |
ndarray |
Initial lower layer tension water content. |
required |
wd |
ndarray |
Initial deep layer tension water content. |
required |
s |
ndarray |
Initial free water storage. |
required |
fr |
ndarray |
Initial runoff area ratio. |
required |
wm |
float |
Total tension water capacity. |
required |
wumx |
float |
Upper layer capacity ratio. |
required |
wlmx |
float |
Lower layer capacity ratio. |
required |
kc |
float |
Evapotranspiration coefficient. |
required |
b |
float |
Exponent of tension water capacity curve. |
required |
c |
float |
Deep evapotranspiration coefficient. |
required |
im |
float |
Impervious area ratio. |
required |
sm |
float |
Average free water capacity. |
required |
ex |
float |
Exponent of free water capacity curve. |
required |
kg |
float |
Groundwater outflow coefficient. |
required |
ki |
float |
Interflow outflow coefficient. |
required |
time_interval |
float |
Time interval. |
required |
time_steps |
int |
Number of time steps. |
required |
Returns:
| Type | Description |
|---|---|
Tuple containing runoff results and state variables |
(wu_out, wl_out, wd_out, s_out, fr_out, rs, ri, rg, runoff_total) |
Source code in hydromodel/models/xaj_slw.py
def sms3_runoff_generation_vectorized(
precipitation: np.ndarray,
evapotranspiration: np.ndarray,
wu: np.ndarray,
wl: np.ndarray,
wd: np.ndarray,
s: np.ndarray,
fr: np.ndarray,
wm: float,
wumx: float,
wlmx: float,
kc: float,
b: float,
c: float,
im: float,
sm: float,
ex: float,
kg: float,
ki: float,
time_interval: float,
time_steps: int,
) -> Tuple[
np.ndarray,
np.ndarray,
np.ndarray,
np.ndarray,
np.ndarray,
np.ndarray,
np.ndarray,
np.ndarray,
np.ndarray,
]:
"""Vectorized SMS3 runoff generation model.
Args:
precipitation: Precipitation time series.
evapotranspiration: Evapotranspiration time series.
wu: Initial upper layer tension water content.
wl: Initial lower layer tension water content.
wd: Initial deep layer tension water content.
s: Initial free water storage.
fr: Initial runoff area ratio.
wm: Total tension water capacity.
wumx: Upper layer capacity ratio.
wlmx: Lower layer capacity ratio.
kc: Evapotranspiration coefficient.
b: Exponent of tension water capacity curve.
c: Deep evapotranspiration coefficient.
im: Impervious area ratio.
sm: Average free water capacity.
ex: Exponent of free water capacity curve.
kg: Groundwater outflow coefficient.
ki: Interflow outflow coefficient.
time_interval: Time interval.
time_steps: Number of time steps.
Returns:
Tuple containing runoff results and state variables:
(wu_out, wl_out, wd_out, s_out, fr_out, rs, ri, rg, runoff_total)
"""
# Calculate derived parameters
wum = wumx * wm
wlm = (1.0 - wumx) * wlmx * wm
wdm = wm - wum - wlm
wmm = (1.0 + b) * wm / (1.0 - im)
smm = (1.0 + ex) * sm
# Adjust KG and KI
if kg + ki > 0.9:
tmp = (kg + ki - 0.9) / (kg + ki)
kg = kg - kg * tmp
ki = ki - ki * tmp
# Calculate HGI and adjust KG and KI
hgi = (1.0 - np.power((1.0 - kg - ki), (time_interval / 24.0))) / (kg + ki)
kg = hgi * kg
ki = hgi * ki
# 初始化输出数组
wu_out = np.zeros(time_steps)
wl_out = np.zeros(time_steps)
wd_out = np.zeros(time_steps)
s_out = np.zeros(time_steps)
fr_out = np.zeros(time_steps)
rs = np.zeros(time_steps) # 产流数组长度不变
ri = np.zeros(time_steps)
rg = np.zeros(time_steps)
runoff_total = np.zeros(time_steps)
# 设置初始值
wu_curr = wu[0]
wl_curr = wl[0]
wd_curr = wd[0]
s_curr = s[0]
fr_curr = fr[0]
# 存储初始值
wu_out[0] = wu_curr
wl_out[0] = wl_curr
wd_out[0] = wd_curr
s_out[0] = s_curr
fr_out[0] = fr_curr
div = 5.0
# 主循环
for i in range(time_steps - 1):
ek = kc * evapotranspiration[i]
pe = precipitation[i] - ek
# 产流计算
w_curr = wu_curr + wl_curr + wd_curr
if pe >= 2 * div:
nd = int(np.floor(pe / div))
ped = np.full(nd, div)
ped[-1] = pe - (nd - 1) * div
else:
nd = 1
ped = np.array([pe])
rd = np.zeros(nd)
# YIELD1计算
if pe <= 0.0:
r = 0.0
if wu_curr + pe >= 0.0:
wu_curr = wu_curr + pe
else:
eu = wu_curr + ek + pe
wu_curr = 0.0
el = (ek - eu) * wl_curr / wlm
if wl_curr < c * wlm:
el = c * (ek - eu)
if wl_curr - el < 0.0:
ed = el - wl_curr
el = wl_curr
wl_curr = 0.0
wd_curr = wd_curr - ed
else:
wl_curr = wl_curr - el
w_curr = wu_curr + wl_curr + wd_curr
else:
a = 0.0
if wm - w_curr < 0.0001:
a = wmm
else:
a = wmm * (
1.0 - np.power((1.0 - w_curr / wm), (1.0 / (1.0 + b)))
)
r = 0.0
peds = 0.0
for j in range(nd):
a = a + ped[j]
peds = peds + ped[j]
ri_temp = r
r = peds - wm + w_curr
if a < wmm:
r = r + wm * np.power((1.0 - a / wmm), (1.0 + b))
rd[j] = r - ri_temp
if wu_curr + pe - r <= wum:
wu_curr = wu_curr + pe - r
else:
if wu_curr + wl_curr + pe - r - wum >= wlm:
wu_curr = wum
wl_curr = wlm
wd_curr = w_curr + peds - r - wu_curr - wl_curr
if wd_curr > wdm:
wd_curr = wdm
else:
wl_curr = wu_curr + wl_curr + pe - r - wum
wu_curr = wum
w_curr = wu_curr + wl_curr + wd_curr
# DIVI31计算
if pe <= 0.0:
rs_curr = 0.0
rg_curr = s_curr * kg * fr_curr
ri_curr = s_curr * ki * fr_curr
s_curr = s_curr * (1.0 - kg - ki)
else:
rb = im * pe
kid = (1.0 - np.power((1.0 - (kg + ki)), (1.0 / nd))) / (kg + ki)
kgd = kid * kg
kid = kid * ki
rs_curr = 0.0
ri_curr = 0.0
rg_curr = 0.0
for j in range(nd):
td = rd[j] - im * ped[j]
x = fr_curr
if ped[j] > 0:
fr_curr = td / ped[j]
s_curr = x * s_curr / max(fr_curr, 0.0001)
rr = 0
if s_curr >= sm:
rr = (ped[j] + s_curr - sm) * fr_curr
else:
au = smm * (
1.0 - np.power((1.0 - s_curr / sm), (1.0 / (1.0 + ex)))
)
if au + ped[j] < smm:
rr = (
ped[j]
- sm
+ s_curr
+ sm
* np.power((1.0 - (ped[j] + au) / smm), (1.0 + ex))
) * fr_curr
else:
rr = (ped[j] + s_curr - sm) * fr_curr
rs_curr = rr + rs_curr
s_curr = ped[j] - rr / max(fr_curr, 0.0001) + s_curr
rg_curr = s_curr * kgd * fr_curr + rg_curr
ri_curr = s_curr * kid * fr_curr + ri_curr
s_curr = s_curr * (1.0 - kid - kgd)
rs_curr = rs_curr + rb
# 确保非负值
rs_curr = max(0.0, rs_curr)
ri_curr = max(0.0, ri_curr)
rg_curr = max(0.0, rg_curr)
# 存储结果
wu_out[i + 1] = wu_curr
wl_out[i + 1] = wl_curr
wd_out[i + 1] = wd_curr
s_out[i + 1] = s_curr
fr_out[i + 1] = fr_curr
rs[i] = rs_curr
ri[i] = ri_curr
rg[i] = rg_curr
runoff_total[i] = rs_curr + ri_curr + rg_curr
return wu_out, wl_out, wd_out, s_out, fr_out, rs, ri, rg, runoff_total
xaj_slw(p_and_e, parameters, warmup_length=365, return_state=False, return_warmup_states=False, normalized_params='auto', **kwargs)
¶
Vectorized XinAnJiang Songliao hydrological model using SMS3 and LAG3 algorithms.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
p_and_e |
ndarray |
Precipitation and potential evapotranspiration data, 3D array [time, basin, feature=2]. Feature 0 is precipitation, feature 1 is potential evapotranspiration. |
required |
parameters |
ndarray |
Model parameters, 2D array [basin, parameter]. Parameter order: [WUP, WLP, WDP, SP, FRP, WM, WUMx, WLMx, KC, B, C, IM, SM, EX, KG, KI, CS, CI, CG, LAG, KK, X, MP, QSP, QIP, QGP]. |
required |
warmup_length |
int |
Warmup period length, default is 365. |
365 |
return_state |
bool |
If True, returns internal state variables, default is False. |
False |
return_warmup_states |
bool |
If True, returns initial states after warmup, default is False. |
False |
normalized_params |
Union[bool, str] |
Parameter format specification: - "auto": Automatically detect parameter format (default) - True: Parameters are normalized (0-1 range), convert to original scale - False: Parameters are in original scale, use directly |
'auto' |
**kwargs |
Other keyword arguments: - time_interval_hours (float): Time interval in hours, default is 3.0 - basin_area (float): REQUIRED - Basin area in km² Used for converting runoff depth (mm) to discharge (m³/s) - initial_states (dict): Initial state values to override after warmup |
{} |
Returns:
| Type | Description |
|---|---|
Results in different formats based on parameters |
|
Exceptions:
| Type | Description |
|---|---|
ValueError |
If basin_area is not provided in kwargs. |
Note
The basin_area parameter is REQUIRED for XAJ-SLW model because it needs to convert runoff depth (mm) to discharge (m³/s) internally. For CAMELS datasets, basin_area is automatically retrieved from basin attributes. For custom datasets, ensure basin_area is included in your basin configuration or passed explicitly.
Source code in hydromodel/models/xaj_slw.py
def xaj_slw(
p_and_e: np.ndarray,
parameters: np.ndarray,
warmup_length: int = 365,
return_state: bool = False,
return_warmup_states: bool = False,
normalized_params: Union[bool, str] = "auto",
**kwargs,
) -> Union[
np.ndarray,
Tuple[
np.ndarray,
np.ndarray,
np.ndarray,
np.ndarray,
np.ndarray,
np.ndarray,
np.ndarray,
np.ndarray,
np.ndarray,
],
Tuple[
np.ndarray,
Dict[str, np.ndarray],
],
]:
"""Vectorized XinAnJiang Songliao hydrological model using SMS3 and LAG3 algorithms.
Args:
p_and_e: Precipitation and potential evapotranspiration data, 3D array [time, basin, feature=2].
Feature 0 is precipitation, feature 1 is potential evapotranspiration.
parameters: Model parameters, 2D array [basin, parameter].
Parameter order: [WUP, WLP, WDP, SP, FRP, WM, WUMx, WLMx, KC, B, C, IM,
SM, EX, KG, KI, CS, CI, CG, LAG, KK, X, MP, QSP, QIP, QGP].
warmup_length: Warmup period length, default is 365.
return_state: If True, returns internal state variables, default is False.
return_warmup_states: If True, returns initial states after warmup, default is False.
normalized_params: Parameter format specification:
- "auto": Automatically detect parameter format (default)
- True: Parameters are normalized (0-1 range), convert to original scale
- False: Parameters are in original scale, use directly
**kwargs: Other keyword arguments:
- time_interval_hours (float): Time interval in hours, default is 3.0
- basin_area (float): **REQUIRED** - Basin area in km²
Used for converting runoff depth (mm) to discharge (m³/s)
- initial_states (dict): Initial state values to override after warmup
Returns:
Results in different formats based on parameters:
- If return_state is False and return_warmup_states is False:
Returns QSim array [time, basin, 1]
- If return_state is False and return_warmup_states is True:
Returns (QSim, warmup_states) tuple, where warmup_states contains all state variables
- If return_state is True and return_warmup_states is False:
Returns (QSim, runoffSim, rs, ri, rg, pe, wu, wl, wd) tuple
- If return_state is True and return_warmup_states is True:
Returns (QSim, runoffSim, rs, ri, rg, pe, wu, wl, wd, warmup_states) tuple
Raises:
ValueError: If basin_area is not provided in kwargs.
Note:
The basin_area parameter is REQUIRED for XAJ-SLW model because it needs to
convert runoff depth (mm) to discharge (m³/s) internally. For CAMELS datasets,
basin_area is automatically retrieved from basin attributes. For custom datasets,
ensure basin_area is included in your basin configuration or passed explicitly.
"""
time_steps, num_basins, _ = p_and_e.shape
time_interval = kwargs.get("time_interval_hours", 3.0)
basin_area = kwargs.get("basin_area", None) # km^2
# === 诊断日志：xaj_slw 函数入口 ===
logger.info(f"=== XAJ_SLW 模型计算开始 ===")
logger.info(f"p_and_e 形状: {p_and_e.shape}, 类型: {type(p_and_e)}")
logger.info(f"parameters 形状: {parameters.shape}, 类型: {type(parameters)}")
logger.info(f"warmup_length: {warmup_length}")
logger.info(f"basin_area: {basin_area}")
logger.info(f"kwargs 参数: {list(kwargs.keys())}")
if "initial_states" in kwargs and kwargs["initial_states"] is not None:
logger.info(f"initial_states: {kwargs['initial_states']}")
if basin_area is None:
raise ValueError(
"basin_area must be provided for XAJ-SLW model. "
"It is required for converting runoff depth (mm) to discharge (m³/s). "
"Please ensure basin_area is included in basin_config when initializing UnifiedSimulator, "
"or pass it explicitly in kwargs."
)
# Process parameters using unified parameter handling
processed_parameters = parameters.copy()
if normalized_params != False:
model_param_dict = get_model_param_config("xaj_slw", kwargs)
param_ranges = model_param_dict["param_range"]
processed_parameters = process_parameters(
parameters, param_ranges, normalized=normalized_params
)
# Extract parameters - all are [basin] arrays
wup = processed_parameters[:, 0] # Initial upper layer tension water
wlp = processed_parameters[:, 1] # Initial lower layer tension water
wdp = processed_parameters[:, 2] # Initial deep layer tension water
sp = processed_parameters[:, 3] # Initial free water storage
frp = processed_parameters[:, 4] # Initial runoff basin_area ratio
wm = processed_parameters[:, 5] # Total tension water capacity
wumx = processed_parameters[:, 6] # Upper layer capacity ratio
wlmx = processed_parameters[:, 7] # Lower layer capacity ratio
kc = processed_parameters[:, 8] # Evaporation coefficient
b = processed_parameters[:, 9] # Exponent of tension water capacity curve
c = processed_parameters[:, 10] # Deep evapotranspiration coefficient
im = processed_parameters[:, 11] # Impervious basin_area ratio
sm = processed_parameters[:, 12] # Average free water capacity
ex = processed_parameters[:, 13] # Exponent of free water capacity curve
kg = processed_parameters[:, 14] # Groundwater outflow coefficient
ki = processed_parameters[:, 15] # Interflow outflow coefficient
cs = processed_parameters[:, 16] # Channel system recession constant
ci = processed_parameters[:, 17] # Lower interflow recession constant
cg = processed_parameters[:, 18] # Groundwater storage recession constant
lag = processed_parameters[:, 19] # Lag time
kk = processed_parameters[:, 20] # Muskingum K parameter
x = processed_parameters[:, 21] # Muskingum X parameter
mp = processed_parameters[:, 22] # Number of Muskingum reaches
qsp = processed_parameters[:, 23] # Initial surface flow
qip = processed_parameters[:, 24] # Initial interflow
qgp = processed_parameters[:, 25] # Initial groundwater flow
# Handle warmup period
if warmup_length > 0:
p_and_e_warmup = p_and_e[0:warmup_length, :, :]
# Remove initial_states from kwargs for warmup period to avoid applying override during warmup
warmup_kwargs = {
k: v for k, v in kwargs.items() if k != "initial_states"
}
# 初始化LAG状态变量用于预热期
warmup_lag_states = {}
for basin_idx in range(num_basins):
warmup_lag_states[f"qsig_initial_{basin_idx}"] = np.zeros(
max(int(lag[basin_idx]), 3)
)
warmup_lag_states[f"qx_initial_{basin_idx}"] = np.zeros(
int(mp[basin_idx]) + 1
)
warmup_kwargs["lag_initial_states"] = warmup_lag_states
# 运行预热期，获取所有状态变量（包括LAG状态）
warmup_result = xaj_slw(
p_and_e_warmup,
processed_parameters,
warmup_length=0,
return_state=True,
normalized_params=False,
**warmup_kwargs,
)
# 解析预热期结果
if len(warmup_result) >= 12:
(
_,
_,
_,
_,
_,
_,
wu_final,
wl_final,
wd_final,
s_final,
fr_final,
warmup_lag_final_states,
) = warmup_result[:12]
else:
# 如果没有LAG状态，使用默认值
_, _, _, _, _, _, wu_final, wl_final, wd_final = warmup_result[:9]
warmup_lag_final_states = {}
for basin_idx in range(num_basins):
warmup_lag_final_states[f"qsig_initial_{basin_idx}"] = (
np.zeros(max(int(lag[basin_idx]), 3))
)
warmup_lag_final_states[f"qx_initial_{basin_idx}"] = np.zeros(
int(mp[basin_idx]) + 1
)
# 获取s和fr的最终状态
warmup_s_final = np.zeros(num_basins)
warmup_fr_final = np.zeros(num_basins)
for basin_idx in range(num_basins):
# Handle both [time, basin] and [time, basin, feature] shapes.
warmup_s_final[basin_idx] = np.asarray(s_final[-1, basin_idx]).item()
warmup_fr_final[basin_idx] = np.asarray(fr_final[-1, basin_idx]).item()
# 使用预热期结果作为初始条件
wu0 = wu_final[-1, :, 0] # [basin]
wl0 = wl_final[-1, :, 0] # [basin]
wd0 = wd_final[-1, :, 0] # [basin]
s0 = warmup_s_final # [basin]
fr0 = warmup_fr_final # [basin]
kwargs["lag_initial_states"] = warmup_lag_final_states
else:
# 默认状态从参数中获取
wu0 = wup.copy() # 上层张力水含量
wl0 = wlp.copy() # 下层张力水含量
wd0 = wdp.copy() # 深层张力水含量
s0 = sp.copy() # 自由水蓄量
fr0 = frp.copy() # 产流面积系数
# Apply initial state overrides if provided (only after warmup in main call)
initial_states = kwargs.get("initial_states", None)
if initial_states is not None:
# SMS states - 标量状态
if "wu0" in initial_states:
wu0[0] = float(initial_states["wu0"])
if "wl0" in initial_states:
wl0[0] = float(initial_states["wl0"])
if "wd0" in initial_states:
wd0[0] = float(initial_states["wd0"])
if "s0" in initial_states:
s0[0] = float(initial_states["s0"])
if "fr0" in initial_states:
fr0[0] = float(initial_states["fr0"])
# LAG states - 数组状态
if "qsig" in initial_states:
qsig_value = initial_states["qsig"]
# 检查是否为数组类型（避免对标量调用 len()）
if hasattr(qsig_value, '__len__') and not isinstance(qsig_value, (str, float, int)):
qsig_len = max(int(lag[0]), 6)
if len(qsig_value) >= qsig_len:
kwargs["lag_initial_states"] = {
"qsig_initial_0": qsig_value[:qsig_len].copy()
}
if "qx" in initial_states:
qx_value = initial_states["qx"]
# 检查是否为数组类型（避免对标量调用 len()）
if hasattr(qx_value, '__len__') and not isinstance(qx_value, (str, float, int)):
qx_len = int(mp[0]) + 1
if len(qx_value) >= qx_len:
if "lag_initial_states" not in kwargs:
kwargs["lag_initial_states"] = {}
kwargs["lag_initial_states"]["qx_initial_0"] = qx_value[:qx_len].copy()
# Save warmup states before applying overrides (for return_warmup_states)
warmup_states = None
if return_warmup_states:
warmup_states = {
# SMS states (scalar)
"wu0": float(wu0[0]),
"wl0": float(wl0[0]),
"wd0": float(wd0[0]),
"s0": float(s0[0]),
"fr0": float(fr0[0]),
# LAG states (array)
"qsig": warmup_lag_final_states["qsig_initial_0"], # 保持数组形式
"qx": warmup_lag_final_states["qx_initial_0"],
}
inputs = p_and_e[warmup_length:, :, :]
actual_time_steps = inputs.shape[0]
# Initialize output arrays
q_sim = np.zeros((actual_time_steps, num_basins))
runoff_sim = np.zeros((actual_time_steps, num_basins))
rs_out = np.zeros((actual_time_steps, num_basins))
ri_out = np.zeros((actual_time_steps, num_basins))
rg_out = np.zeros((actual_time_steps, num_basins))
pe_out = np.zeros((actual_time_steps, num_basins))
wu_out = np.zeros((actual_time_steps, num_basins))
wl_out = np.zeros((actual_time_steps, num_basins))
wd_out = np.zeros((actual_time_steps, num_basins))
s_out = np.zeros((actual_time_steps, num_basins))
fr_out = np.zeros((actual_time_steps, num_basins))
# Initialize LAG final states dictionary
final_lag_states = {}
# Process each basin
for basin_idx in range(num_basins):
# Extract time series for this basin
prcp = inputs[:, basin_idx, 0]
pet = inputs[:, basin_idx, 1]
# Calculate net precipitation
pe = prcp - pet
# Initial states for this basin
wu_init = np.array([wu0[basin_idx]])
wl_init = np.array([wl0[basin_idx]])
wd_init = np.array([wd0[basin_idx]])
s_init = np.array([s0[basin_idx]])
fr_init = np.array([fr0[basin_idx]])
# === 诊断日志：SMS3 调用前 ===
logger.info(f"=== Basin {basin_idx}: 调用 SMS3 产流计算 ===")
logger.info(f" prcp 形状: {prcp.shape}, 类型: {type(prcp)}")
logger.info(f" pet 形状: {pet.shape}, 类型: {type(pet)}")
logger.info(f" pe 形状: {pe.shape}, 类型: {type(pe)}")
logger.info(f" 初始状态: wu_init={wu_init}, wl_init={wl_init}, wd_init={wd_init}")
logger.info(f" s_init={s_init}, fr_init={fr_init}")
# Run SMS3 runoff generation
(
wu_new,
wl_new,
wd_new,
s_new,
fr_new,
rs_basin,
ri_basin,
rg_basin,
runoff_basin,
) = sms3_runoff_generation_vectorized(
prcp,
pet,
wu_init,
wl_init,
wd_init,
s_init,
fr_init,
wm[basin_idx],
wumx[basin_idx],
wlmx[basin_idx],
kc[basin_idx],
b[basin_idx],
c[basin_idx],
im[basin_idx],
sm[basin_idx],
ex[basin_idx],
kg[basin_idx],
ki[basin_idx],
time_interval,
actual_time_steps,
)
# 诊断日志：检查 SMS3 输出
if len(rs_basin) == 0:
logger.warning(
f"Basin {basin_idx}: SMS3 returned empty array! "
f"actual_time_steps={actual_time_steps}, "
f"rs_basin.shape={rs_basin.shape}"
)
# === 诊断日志：LAG3 调用前 ===
logger.info(f"=== Basin {basin_idx}: 调用 LAG3 汇流计算 ===")
logger.info(f" rs_basin 形状: {rs_basin.shape}, 类型: {type(rs_basin)}")
logger.info(f" ri_basin 形状: {ri_basin.shape}, 类型: {type(ri_basin)}")
logger.info(f" rg_basin 形状: {rg_basin.shape}, 类型: {type(rg_basin)}")
logger.info(f" LAG3 参数: lag={lag[basin_idx]}, cs={cs[basin_idx]}, ci={ci[basin_idx]}, cg={cg[basin_idx]}")
logger.info(f" basin_area={basin_area}, time_interval={time_interval}")
# Run LAG3 routing
lag_initial_states = kwargs.get("lag_initial_states", None)
if lag_initial_states is not None:
qsig_initial = lag_initial_states.get(
f"qsig_initial_{basin_idx}",
np.zeros(max(int(lag[basin_idx]), 3)),
)
qx_initial = lag_initial_states.get(
f"qx_initial_{basin_idx}", np.zeros(int(mp[basin_idx]) + 1)
)
else:
qsig_initial = np.zeros(max(int(lag[basin_idx]), 3))
qx_initial = np.zeros(int(mp[basin_idx]) + 1)
# Run LAG3 routing and collect final states
q_basin, lag_final_states = lag3_routing_vectorized(
rs_basin,
ri_basin,
rg_basin,
time_interval,
basin_area,
ci[basin_idx],
cg[basin_idx],
lag[basin_idx],
cs[basin_idx],
kk[basin_idx],
x[basin_idx],
int(mp[basin_idx]),
qsp[basin_idx],
qip[basin_idx],
qgp[basin_idx],
qsig_initial,
qx_initial,
return_states=True,
)
# Store LAG final states for this basin
final_lag_states[f"qsig_initial_{basin_idx}"] = lag_final_states[
"qsig_final"
]
final_lag_states[f"qx_initial_{basin_idx}"] = lag_final_states[
"qx_final"
]
# Store results
q_sim[:, basin_idx] = q_basin
runoff_sim[:, basin_idx] = runoff_basin
rs_out[:, basin_idx] = rs_basin
ri_out[:, basin_idx] = ri_basin
rg_out[:, basin_idx] = rg_basin
pe_out[:, basin_idx] = pe
wu_out[:, basin_idx] = wu_new
wl_out[:, basin_idx] = wl_new
wd_out[:, basin_idx] = wd_new
s_out[:, basin_idx] = s_new
fr_out[:, basin_idx] = fr_new
# Ensure non-negative discharge
q_sim = np.maximum(q_sim, 0.0)
# Format outputs to match DHF interface: [seq, batch, feature]
q_sim = np.expand_dims(q_sim, axis=2)
runoff_sim = np.expand_dims(runoff_sim, axis=2)
rs_out = np.expand_dims(rs_out, axis=2)
ri_out = np.expand_dims(ri_out, axis=2)
rg_out = np.expand_dims(rg_out, axis=2)
pe_out = np.expand_dims(pe_out, axis=2)
wu_out = np.expand_dims(wu_out, axis=2)
wl_out = np.expand_dims(wl_out, axis=2)
wd_out = np.expand_dims(wd_out, axis=2)
s_out = np.expand_dims(s_out, axis=2)
fr_out = np.expand_dims(fr_out, axis=2)
if return_state:
result = (
q_sim,
runoff_sim,
rs_out,
ri_out,
rg_out,
pe_out,
wu_out,
wl_out,
wd_out,
s_out,
fr_out,
final_lag_states, # 返回最终的LAG状态
)
if return_warmup_states and warmup_states is not None:
return result + (warmup_states,)
else:
return result
else:
if return_warmup_states and warmup_states is not None:
return q_sim, warmup_states
else:
return q_sim