import math
    import numpy as np
    import matplotlib.pyplot as plt
    
    n = 0.025       # Manning の祖度係数
    Q = 1500        # 流量 [cu.m/s]
    h = 2.5         # 下流端の水深 [m]
    xs = (          # 距離 [m]
           0,  500, 1000, 1200, 1800,
        2100, 2500, 3000, 3300, 3800)
    zs = (          # 河床高 [m]
        0.0, 0.5, 0.9, 0.8, 2.0,
        2.3, 3.0, 3.0, 3.5, 4.0)
    Bs = (          # 川幅 [m]
        300, 320, 280, 250, 300,
        300, 320, 350, 300, 250)
    
    sz = 10
    Hs = np.empty(sz)
    invSf = np.empty(sz)
    
    # 汎用変数
    p = 10 / 3
    QQ = Q * Q
    QQ_2g = QQ / 19.6
    nnQQ = n * n * QQ
    
    # 境界条件
    H = zs[0] + h; Hs[0] = H
    BB = Bs[0]**2
    E = H + QQ_2g / (BB * h**2)
    Sf = nnQQ / (BB * math.pow(h, p)); invSf[0] = 1 / Sf
    
    for i, x, z, B in zip(range(1, sz), xs[1:], zs[1:], Bs[1:]):
        BB = B * B
        qq_2g = QQ_2g / BB
    
        dx_2 = (x - xs[i-1]) / 2
        K = nnQQ / BB * dx_2
        cnst = E + Sf * dx_2 - z
        # 水深の初期推定値として直下の計算点の水深を当てる
        while True:
            vv_2g = qq_2g / (h * h)        # 速度水頭
            Sfdx_2 = K / math.pow(h, p)    # 損失水頭
            er = h + vv_2g - Sfdx_2 - cnst # 残差
            if abs(er) < 1e-4: break
            # ニュートン法による補正
            h -= er / (1 + (p * Sfdx_2 - 2 * vv_2g) / h)
        H = h + z; Hs[i] = H
        E = H + vv_2g
        Sf = Sfdx_2 / dx_2; invSf[i] = 1 / Sf
    
    plt.plot(xs, zs, label="$z_b$")
    plt.plot(xs, Hs, label="$H$")
    plt.xlabel("$x$ (m)")
    plt.ylabel("$z_b, H$ (m)")
    plt.legend()
    plt.grid()
    plt.show()
    
    print("|$x$<br>[m]|$z$<br>[m]|$H$<br>[m]|$h$<br>[m]|$1/S_f$|")
    print("|---|---|---|---|")
    for i in range(sz):
        h = Hs[i] - zs[i]
        print("|{:.0f}|{:.1f}|{:.3f}|{:.3f}|{:.0f}|".format(
            xs[i], zs[i], Hs[i], h, invSf[i]
        ))
download