How to use positional arguments appropriately for a complex problem

Question

I am revisiting a school project, which I did not complete to my satisfaction. Namely, I wrote an algorithm that takes an ALMOST arbitrary size set of equations and solves them iteratively. The problem being the "almost" part. Essentially, it must have at least two equations, and will not solve for a single one. This is because, I believe, that I don't understand how to use positional arguments correctly.

Below, in the main method, I define two functions y_prime and z_prime. If I pass them both, I get a beautiful graph of my solutions. But, if I only pass y_prime along with its initial conditions and the solution vector to the rungekutta() function, things go haywire:

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd


def rungekutta(dt, y, t, *funcs):
    """
    The following code was written in order to
    reproduce the classic 4th order Runge-Kutta numerical
    method of solving a system of differential equations.
    The aim was to not only apply this to the budworm deforestation
    model developed by Ludwig et al, but also to create an algorithm
    that is generic enough to accept a wide range of ODEs and
    systems of ODEs.

    :param dt: time step "Delta t"
    :param y: The solution vector at the last time step
    :param t: The time at the last time step
    :param funcs: the vector field dy/dt = f(t,y)
    :return: The solution vector for the next time step
    """

    k1 = [dt * f(*y, t) for f in funcs]
    args = [y_n + 0.5 * k_1 for y_n, k_1 in zip((*y, t), (*k1, dt))]
    k2 = [dt * f(*args) for f in funcs]
    args = [y_n + 0.5 * k_2 for y_n, k_2 in zip((*y, t), (*k2, dt))]
    k3 = [dt * f(*args) for f in funcs]
    args = [y_n + k_3 for y_n, k_3 in zip((*y, t), (*k3, dt))]
    k4 = [dt * f(*args) for f in funcs]

    return [y_n + (k_1 + 2 * k_2 + 2 * k_3 + k_4) / 6 for y_n, k_1, k_2, k_3, k_4 in
            zip(y, k1, k2, k3, k4)]


if __name__ == '__main__':

    def y_prime(y, z, t):
        return -t * y

    def z_prime(y, z, t):
        return z

    t_0 = -10
    t_n = 10
    dt = .05

    steps = int((t_n - t_0) / dt)

    y_soln = [0] * steps
    z_soln = [0] * steps
    time = np.arange(t_0, t_n, dt)

    y_soln[0] = 1.928749848e-22
    z_soln[0] = .0000453999297625

    for i in np.arange(1, steps):
        y_soln[i] = rungekutta(dt, y_soln[i-1], time[i-1], y_prime)

The first error I received, when trying to pass a single equation was:

Traceback (most recent call last):
  File "C:/Users/wesle/PycharmProjects/Budworms/RK4v2.py", line 57, in 
    y_soln[i] = rungekutta(dt, y_soln[i-1], time[i-1], y_prime, z_prime)
  File "C:/Users/wesle/PycharmProjects/Budworms/RK4v2.py", line 23, in rungekutta
    k1 = [dt * f(*y, t) for f in funcs]
  File "C:/Users/wesle/PycharmProjects/Budworms/RK4v2.py", line 23, in 
    k1 = [dt * f(*y, t) for f in funcs]
TypeError: y_prime() argument after * must be an iterable, not float

This was because, I think, I have "y_soln" as a positional argument, but now there is only one and it is no longer iterable. So, I made it a tuple of 1 when I passed it in the main method:

for i in np.arange(1, steps):
    y_soln[i] = rungekutta(dt, (y_soln[i-1],), time[i-1], y_prime)

That bit me in the butt, however, because now I am passing a tuple into my y_prime equation, when what it really needs is a float:

Traceback (most recent call last):
  File "C:/Users/wesle/PycharmProjects/Budworms/RK4v2.py", line 57, in 
    y_soln[i] = rungekutta(dt, (y_soln[i-1],), time[i-1], y_prime)
  File "C:/Users/wesle/PycharmProjects/Budworms/RK4v2.py", line 23, in rungekutta
    k1 = [dt * f(*y, t) for f in funcs]
  File "C:/Users/wesle/PycharmProjects/Budworms/RK4v2.py", line 23, in 
    k1 = [dt * f(*y, t) for f in funcs]
  File "C:/Users/wesle/PycharmProjects/Budworms/RK4v2.py", line 38, in y_prime
    return -t * y
TypeError: can't multiply sequence by non-int of type 'numpy.float64'

My only work-around so far has been to solve an extra, random equation like $y= y'$ in addition to whatever equation I'm interested in. This seems pretty inefficient though.

So, it seems like I'm damned if I do, or damned if I don't. Is there any remedy to this?

EDIT If you want to see the code actually work, replace this:

  for i in np.arange(1, steps):
        y_soln[i] = rungekutta(dt, (y_soln[i-1],), time[i-1], y_prime)

with the instance where I pass both equations and their solution vectors to the function:

for i in np.arange(1, steps):
    y_soln[i], z_soln[i] = rungekutta(dt, (y_soln[i-1], z_soln[i-1]), time[i-1], y_prime, z_prime)

user_3696204 · Answer

My solution ended up being to convert all lists to numpy arrays, which allowed me to take advantage of the built in element-wise scalar addition and multiplication. This made computing the "k" values much less cumbersome and convoluted:

def rk4(dt, t, field, y_0):
    """
    :param dt: float - the timestep
    :param t: array - the time mesh
    :param field: method - the vector field y' = f(t, y)
    :param y_0: array - contains initial conditions
    :return: ndarray - solution
    """

    # Initialize solution matrix. Each row is the solution to the system
    # for a given time step. Each column is the full solution for a single
    # equation.
    y = np.asarray(len(t) * [y_0])

    for i in np.arange(len(t) - 1):
        k1 = dt * field(t[i], y[i])
        k2 = dt * field(t[i] + 0.5 * dt, y[i] + 0.5 * k1)
        k3 = dt * field(t[i] + 0.5 * dt, y[i] + 0.5 * k2)
        k4 = dt * field(t[i] + dt, y[i] + k3)
        y[i + 1] = y[i] + (k1 + 2 * k2 + 2 * k3 + k4) / 6

    return y

How to use positional arguments appropriately for a complex problem

1 Answers