Anonymous

import numpy
def iterate_scenario_1(Q23i, QX, V2i, VAL):
    if Q23i > 0.001:
        if V2i > VAL/2.:
            return QX
        else:
            return 0
    elif Q23i < 0.001:
        if V2i < 2 * VAL:
            return 0
        else:
            return QX

def set_QX(V24i, C24i, k24, QX1, QX2):
    if V24i < 1.5 * C24i/k24: return QX1
    else: return QX2

def model_sewer(**kwargs):

    iterations = int(kwargs.get('iterations', 144))
    scenario   = int(kwargs.get('scenario', 1))

    CM12 = float(kwargs.get('CM12', 25))
    CM24 = float(kwargs.get('CM24', 30))
    k12  = float(kwargs.get('k12', 70))
    k24  = float(kwargs.get('k24', 30))
    VAL  = float(kwargs.get('VAL', 1))
    QX1  = float(kwargs.get('QX1', 60))
    QX2  = float(kwargs.get('QX2', 120))

    b12 = float(kwargs.get('b12', 10./k12))
    b24 =  float(kwargs.get('b24', 5./k24))
    
    V12 = numpy.zeros((iterations,))
    V24 = numpy.zeros((iterations,))
    V2  = numpy.zeros((iterations,)); V2[0] = VAL
    C12 = numpy.zeros((iterations,)); C12[0] = CM12
    C24 = numpy.zeros((iterations,)); C24[0] = CM24
    Q12 = numpy.zeros((iterations,))
    Q23 = numpy.zeros((iterations,))
    Q24 = numpy.zeros((iterations,))

    R12 = numpy.ones((iterations,)) * 20.
    R12[36:36+72] = 50.

    R24 = numpy.ones((iterations,)) * 10.
    R24[36:36+72] = 30.

    if kwargs.get('switch'):
        tmp = R24[:]
        R24 = R12[:]
        R12 = tmp[:]
        del tmp

    # Q: is this right, for scenario1?
    QX = QX1

    delta_t = 1./144
    for i in range(iterations - 1):
        if V12[i] < C12[i]/k12:
            Q12[i] = k12 * V12[i] 
        else:
            Q12[i] = C12[i]
        
        if V24[i] < C24[i]/k24:
            Q24[i] = k24 * V24[i]
        else:
            Q24[i] = C24[i]


        if scenario == 2:
            QX = set_QX(V24[i], C24[i], k24, QX1, QX2)

        Q23[i + 1] = iterate_scenario_1(Q23[i], QX, V2[i], VAL)

        # ADDED: max() so cant go below zero.
        V12[i + 1] = V12[i] + (R12[i] - Q12[i]) * delta_t
        V24[i + 1] = V24[i] + (R24[i] - Q24[i]) * delta_t

        V2[i + 1]  = max(V2[i] + (Q12[i] + Q24[i] - Q23[i]) * delta_t, 0.)

        if V2[i + 1] < VAL:
            C12[i + 1] = CM12
            C24[i + 1] = CM24
        else:
            C12[i + 1] = CM12 - b12 * ( V2[i + 1] - VAL )
            C24[i + 1] = CM24 - b24 * ( V2[i + 1] - VAL )

    if 't' in kwargs:
        t = int(kwargs['t'])
        return {'L12': Q12[t], 'L23': Q23[t], 'L24': Q24[t]}
    return {'L12': Q12, 'L23': Q23, 'L24': Q24}\
            , {'V12': V12, 'V2': V2, 'V24': V24} \
            , {'R12': R12, 'R24': R24 }

import pylab
def plot(result, vs, rs, filename=None):
    labels = result.keys()

    pylab.subplot(311)
    for l in labels:
        pylab.plot(result[l])
    pylab.legend(['Q' + l[1:] for l in labels])

    pylab.subplot(312)
    labels = vs.keys()
    for l in labels:
        pylab.plot(vs[l])
    pylab.legend(labels)

    pylab.subplot(313)
    labels = rs.keys()
    for l in labels:
        pylab.plot(rs[l])
    pylab.legend(labels)

    if not filename:
        pylab.show()
    else:
        r = random.random()
        fn = '/var/www/ms_tmp/model%f.png' % r
        pylab.savefig(fn)
        pylab.close()
        return fn


if __name__ == "__main__":
    import pylab
    qs, vs, rs = model_sewer(C23=3, iterations=144, scenario=1)
    plot(qs, vs, rs, False)