Runge Kutta 2 Python ODE fluid flow

Question

I am trying to solve a set of differential equations in x,y,z, I have made simple kutta 2 code in python before but not with 3 variables. I am not sure if what i have done is an acceptable way to adjust for 3. Here is what i have so far: (i have not entered the exact solution as i do not know it)

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# Define function to compute velocity field

A = 0.2
B = 1.0
C = 0.20


def velfun(x,y,z,t):
    
    xvelocity =  (A*np.sin(z) +  C*np.cos(x))
    yvelocity =  (B*np.sin(x) + A*np.cos(z))
    zvelocity =  (C*np.sin(y) + B*np.cos(x))
    return xvelocity, yvelocity, zvelocity


# Set stopping time
# and the step length
    
Nstep = 10000   
h     = 0.01

# Create arrays to store x and t values

x     = np.zeros(Nstep);
y     = np.zeros(Nstep);
z     = np.zeros(Nstep);

# Set the initial condition

x[0] = 1
y[0] = 0
z[0] = 0

# Carry out steps in a loop

for k in range(1,Nstep):
    
    # Provisional Euler step
    
    xx    = x[k-1]
    yy    = y[k-1]
    zz    = z[k-1]
    ux,uy,uz = velfun(xx,yy,zz)
    xp    = x[k-1] + h*ux
    yp    = y[k-1] + h*uy
    zp    = z[k-1] + h*uz
    # Compute velocity at provisional point
    
    uxp,uyp,uzp = velfun(xp,yp,zp)

    # Compute average velocity
    
    uxa = 0.5*(ux + uxp)
    uya = 0.5*(uy + uyp)
    uza = 0.5*(uz + uzp)
    
    # Make final Euler step
    # using average velocity
    
    x[k] = x[k-1] + h*uxa
    y[k] = y[k-1] + h*uya
    z[k] = z[k-1] + h*uza
    


    
# Exact solution
    


# Plot results
fig = plt.figure()
ax = plt.axes(projection='3d')
ax = plt.axes(projection='3d')
ax.scatter3D(x,y,z,'b',label='x (with RK2)')

plt.show()

Answer 1

This is an interesting problem but I'm afraid that my feedback will be somewhat limited:

• Consider vectorizing your velfun to be represented with matrix multiplications
• Your x/y/z should be combined into one matrix
• You should add axis labels, and some form of colour indicator to make your graph more legible - here I've added an arbitrary "start-to-finish progress" colour
• Rather than label, you likely meant to use title

The code I'm suggesting produces exactly the same results as yours. The problem is that I have no idea whether those are correct. I suggest that you re-post on https://physics.stackexchange.com/ with more details on your physics problem to get confirmation as to whether your RK is implemented correctly. You should include a picture of your graph.

Suggested

from typing import Tuple

import numpy as np
import matplotlib.pyplot as plt

A = 0.2
B = 1.0
C = 0.20

N_STEP = 10_000  # stopping time
H = 0.01         # step length


def new() -> np.ndarray:
    M = np.array([
        [0, 0, A],
        [B, 0, 0],
        [0, C, 0],
    ])
    N = np.array([
        [C, 0, 0],
        [0, 0, A],
        [B, 0, 0],
    ])

    # Define function to compute velocity field
    def velfun(p: np.ndarray) -> np.ndarray:
        return M@np.sin(p) + N@np.cos(p)

    X = np.empty((3, N_STEP))
    a = np.array((1, 0, 0), dtype=np.float64)
    X[:, 0] = a

    for k in range(1, N_STEP):
        # Provisional Euler step
        u = velfun(a)
        p = a + H*u

        # Compute velocity at provisional point
        up = velfun(p)

        # Compute average velocity
        ua = 0.5*(u + up)

        # Make final Euler step
        # using average velocity
        a += H*ua
        X[:, k] = a

    return X


def plot(x, y, z):
    # Plot results
    fig = plt.figure()
    ax = plt.axes(projection='3d')
    ax.set_title('x (with RK2)')
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')
    colour = np.linspace(0, 1, len(x))

    ax.scatter3D(x, y, z, s=0.5, c=colour)


def main():
    # plot(*old())
    plot(*new())
    plt.show()


if __name__ == '__main__':
    main()

Be warned: I have very little understanding of what this code is actually doing, so my variable names are bad.

Answer 2

A few minor things on the RK front

for k in range(1,Nstep): implies Nstep-1 actual steps taken. I don't count the initial value as a step taken. This may or may not be intentional.

---

The RK code itself looks like it's the second order Heun's method $$ u^{k+1} = u^{k} + \frac{1}{2}\left(hf(u^{k}) + hf(u^{k} + hf(u^k))\right). $$

---

I would vectorize all similar triplets

x[k] = x[k-1] + h*uxa
y[k] = y[k-1] + h*uya
z[k] = z[k-1] + h*uza

to this end, see Reinderien's answer!

---

I think it's better to not elaborate on the RK formulae.

The classical form has two versions,

\begin{align} U^{k}_1 &= u^{k} \\ U^k_2 &= u^{k} + hf(u^{k}) \\ &\\ u^{k+1} &= u^{k} + \frac{1}{2}\left(hf(U^k_1) + hf(U^k_2)\right) \end{align}

and

\begin{align} k_1 &= f(u^{k}) \\ k_2 &= f(u^{k} + hk_1) \\ &\\ u^{k+1} &= u^{k} + \frac{1}{2}\left(k_1 + k_2\right). \end{align}

A similar structure appearing in the code would be more welcome to my eyes. I think it's safe to rename velfun to f. Then it could go something like

U_1 = u[k-1]
U_2 = u[k-1] + h*f(U_1)

u[k] = u[k-1] + 0.5*h*(f(U_1) + f(U_2))

Another common possibility is

k_1 = f(u[k-1])
k_2 = f(u[k-1] + h*k_1)

u[k] = u[k-1] + 0.5*h*(k_1 + k_2)