Generalize sweep

idmsweep.py

-from idmlin import IDMLin
 from matplotlib import pyplot as plt
-from numpy import arange
+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)
@@ -11,44 +43,48 @@ def idmgo(t, a, b, s, v, fmt=None):
     else:
         return res, idm.os(), fmt
 
-def sweep():
-    names = ("T", "a", "b", "s_0", "v_0")
-    defaults = (1.5, 1, 3, 2, 30)
-
-    trng = [0.5, 1, 2, 2.5]
-    arng = [0.3, 0.5, 0.8, 1.5]
-    brng = [1, 2, 4, 5]
-    srng = [0.5, 1, 2.5, 3]
-    vrng = [10, 20, 40, 50]
-
-    rngs = [trng, arng, brng, srng, vrng]
-
-    colors = ['r', 'g', 'b', 'm', 'c']
-    dots = [':', '--', '-', '.-']
-
-    '''
-    for i in range(len(rngs)):
-        plt.plot(*idmgo(*defaults, fmt='k'))
-        print names[i], rngs[i][:2], defaults[i], rngs[i][2:]
-        args = list(defaults)
-        for n, arg in enumerate(rngs[i]):
-            args[i] = arg
-            plt.plot(*idmgo(*args, fmt=colors[i]+dots[n]))
-
-        plt.axis([0,1,-.2,.2])
-        plt.show()
-    '''
-
-    plt.plot(*idmgo(*defaults, fmt='k'))
-    print "a = ", 
-    for arg in range(20):
-        args = list(defaults)
-        a = arg/20. + 0.2
-        print a,
-        args[1] = a
-        plt.plot(*idmgo(*args))
-    print
-    plt.axis([0,0.1,-.05,.45])
-    plt.show()
-
-sweep()
+def sweep(fn, rngs):
+    defaults = [x["default"] for x in rngs]
+
+    for i, r in enumerate(rngs):
+        if "values" in r and r["values"] is not None:
+            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))
+
+            # plt.axis([0,1,-.2,.2])
+            plt.show()
+
+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)
+
+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)
+
+'''
+go = linfty(0.1, n=20, tmax=200)
+sweep(go, rngs)
+'''