ouyangwenyu / hydromodel Goto Github PK

def xaj(
    p_and_e,
    params: Union[np.array, list],
    return_state=False,
    warmup_length=365,
    **kwargs,
) -> Union[tuple, np.array]:
    """
    run XAJ model

    Parameters
    ----------
    p_and_e
        prcp and pet; sequence-first (time is the first dim) 3-d np array: [time, basin, feature=2]
    params
        parameters of XAJ model for basin(s);
        2-dim variable -- [basin, parameter]:
        the parameters are B IM UM LM DM C SM EX KI KG A THETA CI CG (notice the sequence)
    return_state
        if True, return state values, mainly for warmup periods
    warmup_length
        hydro models need a warm-up period to get good initial state values
    kwargs
        route_method
            now we provide two ways: "CSL" (recession constant + lag time) and "MZ" (method from mizuRoute)
        source_type
            default is "sources" and it will call "sources" function; the other is "sources5mm",
            and we will divide the runoff to some <5mm pieces according to the books in this case
        source_book
            When source_type is "sources5mm" there are two implementions for dividing sources,
            as the methods in "ShuiWenYuBao" and "GongChengShuiWenXue"" are different.
            Hence, both are provided, and the default is the former.

    Returns
    -------
    Union[np.array, tuple]
        streamflow or (streamflow, states)
    """
    # default values for some function parameters
    model_name = kwargs.get("name", "xaj")
    source_type = kwargs.get("source_type", "sources")
    source_book = kwargs.get("source_book", "HF")
    pr_file = kwargs.get("param_range_file", None)
    model_param_dict = read_model_param_dict(pr_file)
    # params
    param_ranges = model_param_dict[model_name]["param_range"]
    if model_name == "xaj":
        route_method = "CSL"
    elif model_name == "xaj_mz":
        route_method = "MZ"
    else:
        raise NotImplementedError(
            "We don't provide this route method now! Please use 'CSL' or 'MZ'!"
        )
    if np.isnan(params).any():
        raise ValueError(
            "Parameters contain NaN values. Please check your opt algorithm"
        )
    xaj_params = [
        (value[1] - value[0]) * params[:, i] + value[0]
        for i, (key, value) in enumerate(param_ranges.items())
    ]

    k = xaj_params[0]
    b = xaj_params[1]
    im = xaj_params[2]
    um = xaj_params[3]
    lm = xaj_params[4]
    dm = xaj_params[5]
    c = xaj_params[6]
    sm = xaj_params[7]
    ex = xaj_params[8]
    ki_ = xaj_params[9]
    kg_ = xaj_params[10]
    # ki+kg should be smaller than 1; if not, we scale them
    ki = np.where(ki_ + kg_ < 1.0, ki_, (1.0 - PRECISION) / (ki_ + kg_) * ki_)
    kg = np.where(ki_ + kg_ < 1.0, kg_, (1.0 - PRECISION) / (ki_ + kg_) * kg_)
    if route_method == "CSL":
        cs = xaj_params[11]
        l = xaj_params[12]
    elif route_method == "MZ":
        # we will use routing method from mizuRoute -- http://www.geosci-model-dev.net/9/2223/2016/
        a = xaj_params[11]
        theta = xaj_params[12]
        # make it as a parameter
        kernel_size = int(xaj_params[15])
    else:
        raise NotImplementedError(
            "We don't provide this route method now! Please use 'CSL' or 'MZ'!"
        )
    ci = xaj_params[13]
    cg = xaj_params[14]

    # initialize state values
    if warmup_length > 0:
        p_and_e_warmup = p_and_e[0:warmup_length, :, :]
        _q, _e, *w0, s0, fr0, qi0, qg0 = xaj(
            p_and_e_warmup,
            params,
            return_state=True,
            warmup_length=0,
            **kwargs,
        )
    else:
        w0 = (0.5 * um, 0.5 * lm, 0.5 * dm)
        s0 = 0.5 * sm
        fr0 = np.full(ex.shape, 0.1)
        qi0 = np.full(ci.shape, 0.1)
        qg0 = np.full(cg.shape, 0.1)

    # state_variables
    inputs = p_and_e[warmup_length:, :, :]
    runoff_ims_ = np.full(inputs.shape[:2], 0.0)
    rss_ = np.full(inputs.shape[:2], 0.0)
    ris_ = np.full(inputs.shape[:2], 0.0)
    rgs_ = np.full(inputs.shape[:2], 0.0)
    es_ = np.full(inputs.shape[:2], 0.0)
    for i in range(inputs.shape[0]):
        if i == 0:
            '''
            这个函数实现了新安江模型中的单步径流产生过程。在水文学中，径流是指雨水通过地表流向河流、湖泊等水体的过程，是水文循环中重要的组成部分之一。
            '''
            (r, rim, e, pe), w = generation(
                inputs[i, :, :], k, b, im, um, lm, dm, c, *w0
            )  ###Runoff（径流） Runoff from impermeable areas（不透水面径流） Evapotranspiration（蒸发蒸腾量） Net Precipitation（净降水）
            if source_type == "sources":
                (rs, ri, rg), (s, fr) = sources(
                    pe, r, sm, ex, ki, kg, s0, fr0, book=source_book
                )
            elif source_type == "sources5mm":
                (rs, ri, rg), (s, fr) = sources5mm(
                    pe, r, sm, ex, ki, kg, s0, fr0, book=source_book
                ) ## 总的地表径流  总的壤中流  总的地下径流  模拟得到的最终的自由水蓄水容量深度  模拟得到的最终的初始面积。
            else:
                raise NotImplementedError("No such divide-sources method")
        else:
            (r, rim, e, pe), w = generation(
                inputs[i, :, :], k, b, im, um, lm, dm, c, *w
            )
            if source_type == "sources":
                (rs, ri, rg), (s, fr) = sources(
                    pe, r, sm, ex, ki, kg, s, fr, book=source_book
                )
            elif source_type == "sources5mm":
                (rs, ri, rg), (s, fr) = sources5mm(
                    pe, r, sm, ex, ki, kg, s, fr, book=source_book
                )
            else:
                raise NotImplementedError("No such divide-sources method")
        # impevious part is pe * im
        runoff_ims_[i, :] = rim
        # so for non-imprvious part, the result should be corrected
        rss_[i, :] = rs * (1 - im)  ###  im 是透水率（impermeability）
        ris_[i, :] = ri * (1 - im)
        rgs_[i, :] = rg * (1 - im)
        es_[i, :] = e
    # seq, batch, feature
    runoff_im = np.expand_dims(runoff_ims_, axis=2)
    rss = np.expand_dims(rss_, axis=2)
    es = np.expand_dims(es_, axis=2)

    qs = np.full(inputs.shape[:2], 0.0)
    if route_method == "CSL":   ###qs 是输出的流量（汇流后的径流）。qt 是通过简单叠加 qs_、qi 和 qg 计算得到的流量。lag 是汇流时段的滞后时间。cs 是汇流过程中的汇流常数。
        qt = np.full(inputs.shape[:2], 0.0)
        for i in range(inputs.shape[0]):
            if i == 0:
                qi = linear_reservoir(ris_[i], ci, qi0) #地表径流
                qg = linear_reservoir(rgs_[i], cg, qg0) #地下径流
            else:
                qi = linear_reservoir(ris_[i], ci, qi)
                qg = linear_reservoir(rgs_[i], cg, qg)
            qs_ = rss_[i]
            qt[i, :] = qs_ + qi + qg
        for j in range(len(l)):
            lag = int(l[j])
            for i in range(lag):
                qs[i, j] = qt[i, j]
            for i in range(lag, inputs.shape[0]):
                qs[i, j] = cs[j] * qs[i - 1, j] + (1 - cs[j]) * qt[i - lag, j]
    elif route_method == "MZ":
        rout_a = a.repeat(rss.shape[0]).reshape(rss.shape)
        rout_b = theta.repeat(rss.shape[0]).reshape(rss.shape)
        conv_uh = uh_gamma(rout_a, rout_b, kernel_size)
        qs_ = uh_conv(runoff_im + rss, conv_uh)
        for i in range(inputs.shape[0]):
            if i == 0:
                qi = linear_reservoir(ris_[i], ci, qi0)
                qg = linear_reservoir(rgs_[i], cg, qg0)
            else:
                qi = linear_reservoir(ris_[i], ci, qi)
                qg = linear_reservoir(rgs_[i], cg, qg)
            qs[i, :] = qs_[i, :, 0] + qi + qg
    else:
        raise NotImplementedError(
            "We don't provide this route method now! Please use 'CS' or 'MZ'!"
        )

    # seq, batch, feature
    q_sim = np.expand_dims(qs, axis=2)
    if return_state:
        return q_sim, es, *w, s, fr, qi, qg

   ******马斯京根算法集成在这里进行上游的汇流演算******
    q_sim = q_sim + q_mk1 + q_mk2 + q_mk3
 *************************************************************
    return q_sim, es

WM
WUM
WLM
KC
B
SMD
EX
SKIGD
KID
IMP
C
CID
CGD
N
NK

     if route_method == "CSL":
        qt = np.full(inputs.shape[:2], 0.0)
        for i in range(inputs.shape[0]):
            if i == 0:
                qi = linear_reservoir(ris_[i], ci, qi0)
                qg = linear_reservoir(rgs_[i], cg, qg0)
            else:
                qi = linear_reservoir(ris_[i], ci, qi)
                qg = linear_reservoir(rgs_[i], cg, qg)
            qs_ = rss_[i]
            qt[i, :] = qs_ + qi + qg
        for j in range(inputs.shape[1]):
            qs[:, j] = nash_iu(qt[:, j], N, NK, inputs.shape[0])

        # for j in range(len(l)):
        #     lag = int(l[j])
        #     for i in range(lag):
        #         qs[i, j] = qt[i, j]
        #     for i in range(lag, inputs.shape[0]):
        #         qs[i, j] = cs[j] * qs[i - 1, j] + (1 - cs[j]) * qt[i - lag, j]
def nash_iu(r, k, n, time_steps):
    """
    计算纳什瞬时单位线输出
    r: 输入径流 (输入时间序列)
    k: 储水库时间常数
    n: 储水库数量
    time_steps: 总时间步数
    """
    from scipy.special import gamma, gammaincc

    t = np.arange(time_steps)
    unit_hydrograph = (t ** (n - 1) * np.exp(-t / k)) / (k ** n * gamma(n))

    outflow = np.convolve(r, unit_hydrograph)[:time_steps]
    return outflow

default={
            "name": "SCE_UA",
            "random_seed": 2048,
            # these params are just for test
            "rep": 1500,
            "ngs": 22,
            "kstop": 5,
            "peps": 0.1,
            "pcento": 0.1,
        },