solid func, needs testing

1 files changed, 16 insertions, 6 deletions

diff --git a/throw.py b/throw.py
index e36b151..b144eaa 100644
--- a/throw.py
+++ b/throw.py
@@ -143,20 +143,26 @@ def interpret(soln, max_tension):
     
     return phi, phi1, v_tan, tension, moment, momentum, max_arg
 
-def solver(func, init_vars, timespace=(0,600,10**5), max_tension=max_tension):
+def solver(func, init_vars, timespace=(0,600,10**5), max_tension=max_tension,
+    returns=['time', 'variables', 'slices', 'maxes']):
     for n in itertools.count(1):
         time, soln = solve_fun(func, init_vars, timespace)
         
-        try: phi, phi1, v_tan, tension, moment, momentum, max_arg = \
-                interpret(soln, max_tension)
+        try: variables, max_arg = interpret(soln, max_tension)
         except TimespaceTooSmall:
             timespace = map(lambda x: x*.2*1.5**n, timespace)
             continue
         break
     
-    phi, phi1, v_tan, tension, moment, momentum, time = \
-        map(lambda x: numpy.ndarray.__getslice__(x, 0, max_arg+1),
-            [phi, phi1, v_tan, tension, moment, momentum, time])
+    getslice = lambda x: numpy.ndarray.__getslice__(x, 0, max_arg+1)
+    getmax   = lambda x: numpy.ndarray.__getitem__( x,    max_arg)
+    slices = map(getslice, (time,) + variables)
+    maxes  = map(getmax,   (time,) + variables)
+    
+    out = []
+    for item in returns: out.append(vars()[item])
+    
+    return tuple(out)
 
 def plot(time, soln, graphs=['all'], i=7):
     phi, phi1, v_tan, tension, moment, momentum, max_arg = \
@@ -253,6 +259,10 @@ def opt_fun(guesses, f, func, init_vars):
         map(lambda x: numpy.ndarray.__getslice__(x, 0, max_arg+1),
             [phi, phi1, v_tan, tension, moment, momentum, time])
     
+    (time, soln,
+        (phi, phi1, v_tan, tension, moment, momentum, max_arg),
+        (
+    
     # Parameter values.
     f.write("%s, " % r_roll)
     f.write("%s, " % length)