import math
    import numpy as np
    import matplotlib.pyplot as plt
    from particle import Particle
    
    n  = 0.020      # Manning の祖度係数
    Q  = 500        # 流量 [cu.m/s]
    Hd = 12         # 下流端の水位 [m]
    zd = 0          # 下流端の河床高 [m]
    dx = 200        # 距離刻み [m]
    dt = 10         # 時間刻み [s]
    So = 1 / 700    # 河床勾配
    B  = 100        # 川幅 [m]
    X  = 10000      # 水路長 [m]
    d  = 0.01       # 粒径 [cm]
    void = 0.4      # 空隙率
    
    sz = int(X / dx) + 1
    xs = np.linspace(0, X, sz)
    zs = xs * So + zd
    Hs = np.empty(sz)
    qss = np.empty(sz)
    cs = np.zeros(sz)
    dc_dt =np.zeros(sz)
    
    # 汎用変数
    p10_3 = 10 / 3
    q = Q / B
    qq = q * q
    qq_2g = qq / 19.6
    nnqq = n * n * qq
    dx_2 = dx / 2
    K = nnqq * dx_2
    q_dx = 100 * q / dx # cm/s
    dt_ve = dt / (1 - void)
    
    p = Particle(d)
    sgd = p.sgd
    wf = p.wf
    
    def calc_qss(isFirst=False):
    
        # 境界条件
        h = Hd - zs[0]
        E = Hd + qq_2g / h**2
        Sf = nnqq / math.pow(h, p10_3)
        ts = 98000 * h * Sf / sgd
    
        Hs[0]  = Hd
        qss[0] = p.Itakura(ts)
    
        for i, H, z in zip(range(1, sz), Hs[1:], zs[1:]):
            cnst = E + Sf * dx_2 - z
            if not isFirst: h = H - z 
            while True:
                vv_2g = qq_2g / (h * h)        # 速度水頭
                Sfdx_2 = K / math.pow(h, p10_3)    # 損失水頭
                er = h + vv_2g - Sfdx_2 - cnst # 残差
                if abs(er) < 1e-4: break
                # ニュートン法による補正
                h -= er / (1 + (p10_3 * Sfdx_2 - 2 * vv_2g) / h)
            H = h + z
            E = H + vv_2g
            Sf = Sfdx_2 / dx_2
            ts = 98000 * h * Sf / sgd
    
            Hs[i]  = H
            qss[i] = p.Itakura(ts) # cu.cm/s/sq.cm
    
    tms_plt = [t * 3600 for t in [1, 4, 12, 24]]
    cnv = dt_ve / 100
    
    plt.plot(xs, zs, label="initial")
    
    t = dt; isFirst = True
    while t <= tms_plt[-1]:
    
        calc_qss(isFirst)
    
        cs[-1] = qss[-1] / wf
        c_up = cs[-1]
        for i, qs, c, H, z in zip(range(sz - 2, -1, -1),
                qss[-2::-1], cs[-2::-1], Hs[-2::-1], zs[-2::-1]):
            w = qs - wf * c
            dc_dt[i] = (w - q_dx * (c - c_up)) / (100 * (H - z))
            zs[i]   -= cnv * w # 河床高の更新
            c_up = c
    
        cs += dt * dc_dt # 浮遊砂濃度の更新
    
        if t in tms_plt:
            lbl = f"{t // 3600} hr"
            plt.plot(xs, zs, label=lbl)
    
        t += dt
        isFirst = False
    
    plt.xlabel("$x$ (m)")
    plt.ylabel("$z_b$ (m)")
    plt.legend(loc="lower right")
    plt.grid()
    plt.show()
download