idmsweep.py

from matplotlib import pyplot as plt
from scipy import signal
from scipy import sqrt, linspace, exp, pi, polymul, sinh, cosh, arange, ones, size, array, diff, insert

from idmlin import IDMLin


def linfty(re, n=25, tmax=100, dt=0.1, timeplot=False):
    def go(t, a, b, s, v, fmt=None):
        idm = IDMLin(t=t, a=a, b=b, s0=s)
        idm.go(v, re)

        sys = idm.sys()
        T = linspace(0, tmax, num=tmax/dt)
        U = ones(size(T))
        A = [0]*n
        Us = ones((len(T), n))

        for i in range(n):
            _, U, _ = signal.lsim(sys, U, T)
            A[i] = max(U)
            Us[:,i] = U

        if timeplot:
            if fmt is None:
                return T, array(Us)
            else:
                return T, array(Us), fmt
        else:
            if fmt is None:
                return arange(n)+1, A
            else:
                return arange(n)+1, A, fmt
    return go


def idmgo(t, a, b, s, v, fmt=None):
    idm = IDMLin(t=t, a=a, b=b, s0=s)
    res = arange(100)/100.
    idm.go(v, res)
    if fmt is None:
        return res, idm.os()
    else:
        return res, idm.os(), fmt

def sweep(fn, rngs, axis=None, separate=True):
    defaults = [x["default"] for x in rngs]

    if not separate:
        plt.plot(*fn(*defaults, fmt='k'))

    for i, r in enumerate(rngs):
        if "values" in r and r["values"] is not None:
            if separate:
                plt.plot(*fn(*defaults, fmt='k'))

            print r["name"], defaults[i], r["values"]
            args = list(defaults)
            for dot, arg in zip(r["dots"], r["values"]):
                args[i] = arg
                plt.plot(*fn(*args, fmt=r["color"]+dot))

            if axis is not None:
                plt.axis(axis)
            if separate:
                plt.show()
    if not separate:
        plt.show()

if __name__ == "__main__":

    '''
    ### What happens when you vary a?
    ### ( enhance )

    avals = arange(20)/20. + 0.2
    dots = ["-"] * len(avals)

    arng = {"name": "a",   "color": "", "dots": dots, "default": 1.0, "values": avals }
    trng = {"name": "T",   "default": 1.5}
    brng = {"name": "b",   "default": 3.0}
    srng = {"name": "s_0", "default": 2.0}
    vrng = {"name": "v_0", "default": 30.}

    rngs = [trng, arng, brng, srng, vrng]

    sweep(idmgo, rngs, axis=[0,0.1,-.05,.45])
    '''

    '''
    ### What happens when you vary IDM parameters?
    dots = (':', '--', '-', '.-')

    trng = {"name": "T",   "color": "r", "dots": dots, "default": 1.5, "values": (0.5, 1, 2, 2.5) }
    arng = {"name": "a",   "color": "g", "dots": dots, "default": 1.0, "values": (0.3, 0.5, 0.8, 1.5) }
    brng = {"name": "b",   "color": "b", "dots": dots, "default": 3.0, "values": (1, 2, 4, 5) }
    srng = {"name": "s_0", "color": "m", "dots": dots, "default": 2.0, "values": (0.5, 1, 2.5, 3) }
    vrng = {"name": "v_0", "color": "c", "dots": dots, "default": 30., "values": (10, 20, 40, 50) }

    rngs = [trng, arng, brng, srng, vrng]

    sweep(idmgo, rngs, axis=[0,1,-.2,.2])
    '''

    '''
    ### What happens when you vary IDM parameters?
    dots = (':', '--', '-', '.-')

    trng = {"name": "T",   "color": "r", "dots": dots, "default": 1.5, "values": (0.8, 0.9, 1, 1.1) }
    arng = {"name": "a",   "color": "g", "dots": dots, "default": 1.0, "values": (0.3, 0.4, 0.5, 0.6) }
    brng = {"name": "b",   "color": "b", "dots": dots, "default": 3.0, "values": (1, 5, 10, 20) }
    srng = {"name": "s_0", "color": "m", "dots": dots, "default": 2.0, "values": (3, 5, 10, 20) }
    vrng = {"name": "v_0", "color": "c", "dots": dots, "default": 30., "values": (10, 20, 40, 50) }

    rngs = [trng, arng, brng, srng, vrng]

    go = linfty(0.5, n=50, tmax=500)
    sweep(go, rngs)
    '''

    go = linfty(0.05, n=80, tmax=500, timeplot=True)
    T, Us = go(1.5, 0.3, 3, 2, 30)

    f, axarr = plt.subplots(2, sharex=True)
    axarr[0].plot(T, Us)
    axarr[1].plot(T, diff(insert(Us, 0, 1, axis=1)))
    plt.show()

    '''
    go = linfty(0.5, n=500, tmax=2000, dt=1, timeplot=False)
    plt.plot(*go(1.5, 0.57, 3, 2, 30))
    plt.show()
    '''