集成马斯京根算法的想法。

[Smile-L-up]

Description

您好，很感谢您能分享这么优秀且有意义的项目。但是实际应用中，非源头水域水库的入库计算的需求感觉还是比较大的。为了适用不是源头的流域，需要考虑上游的汇流情况。想着是不是可以将马斯京根算法集成到这个项目，以适应复杂水域的情况，以下是自己的一个简单想法及简单实现，不知道对不对。
Describe the feature (e.g., new functions/tutorials) you would like to propose.
Tell us what can be achieved with this new feature and what's the expected outcome.

Source code

##马斯京根算法的计算流程。

class MuskingumReach:
    
    def __init__(self, k, x):
        self.k = k
        self.x = x

    def route_hydrograph(self, inflows, dt):
        outflows = list()
        outflows.append((inflows[0]))

        c0 = ((dt / self.k) - (2 * self.x)) / ((2 * (1 - self.x)) + (dt / self.k))
        c1 = ((dt / self.k) + (2 * self.x)) / ((2 * (1 - self.x)) + (dt / self.k))
        c2 = ((2 * (1 - self.x)) - (dt / self.k)) / ((2 * (1 - self.x)) + (dt / self.k))
        for i in range(len(inflows) - 1):
            q_out = (c0 * inflows[i + 1]) + (c1 * inflows[i]) + (c2 * outflows[i])
            q_out = max(min(inflows), q_out)
            outflows.append(q_out)
        return outflows

这个是仓库代码的XAJ模型部分，将马斯京根算法集成到xaj里面，汇总得到最终的汇流返回。

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

不知道这种方式是否可行？不知道作者大大什么时候可以实现相关的半分布式模型？可以加您个联系方式吗？想请教一下您？

[OuyangWenyu]

谢谢您的建议，我们已经有同学把马斯京根放到这个仓库的其他下游fork仓库了，不过还需要整个测试下，等没什么问题了，就会合并到这个仓库的main分支里面；如果有联系需要可以发邮件，我的GitHub主页有