Question

I’m starting to learn Python by writing an ODE solver. I would like to handle transparently both one- or many-variable input functions. Here is my code for one step of Euler’s method:

def euler(h, t, y, f):
    return (y + h*f for y,f in zip(y,f(t,y)))

Now I define two functions, f1 and f2 like this:

def f1(t,y):
    return -2*t*y

def f2(t,y):
    x, y = y #is rebinding usually ok, or confusing?
    return (x - t*y, y + x/t)

When I test them, that’s what (obviously) happens

>>> list(euler(0.01, 1, (1,2), f2))
[0.99, 2.03]
>>> list(euler(0.01, 1, 1, f1))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 2, in euler
TypeError: zip argument #1 must support iteration

I would like the solver to handle transparently if the function given works on one or more variables, but haven’t found a cool way to do it. A way I found was

import operator as op
def euler(h, t, y, f):
    if op.isNumberType(y):
        return (y + h*f(t,y),)
    return (y + h*f for y,f in zip(y,f(t,y)))

But now I passed a float and returned an iterable, so list(euler(...)) can suceed. However, I can’t call, for example, f(t,euler(...)).

Is there a way to handle a singleton sequence as a primitive type, or a primitive as a singleton sequence without endless checking? By “endless checking” I mean, having to check just in a few places, and not all over my code. Or should I just suck it up and make f(t,y) expects a sequence instead of a numeric?

Thanks for any help, and tips about my coding are welcome, too!

Answer 1

It’s very easy to make f(t, y) use sequence only:

def euler(h, t, y, f):
    return (y + h*v for y,v in zip(y,f(t,y)))

def f1(t,status):
    x, = status
    return -2*t*x,

def f2(t,status):
    x, y = status
    return x - t*y, y + x/t


print list(euler(0.01, 1, (1,2), f2))
print list(euler(0.01, 1, (1,), f1))