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

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.1, n=50, tmax=500)
    sweep(go, rngs)