# Runge Kutta 2 Python ODE fluid flow

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()

• Anyway, this will not run: your velfun expects four parameters but you only ever pass three. – Reinderien May 14 at 14:23
• its not a typo and i have adjusted the velfun, the code runs i am just unsure if the solution is in any way accurate – ben May 15 at 6:59
• Cross-posted from math.stackexchange.com/questions/4137970/…. The system should probably be the Arnold-Beltrami-Childress ABC dynamic, with variables and constants cyclic in all positions. – Lutz Lehmann May 23 at 7:08

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.

• Very nice code especially if you haven't ever done something similar. a is some point in space. velfun(a) is the velocity at that point. H is the step size, i.e. length of time between two successive numerical approximations. I would use h for the latter. x instead of a for a point in space. If you are interested, a possible next step is writing code that takes the Butcher tableau of an explicit method (the tableau is then strictly lower triangular), and yields the method itself as a function. – Andrew May 14 at 20:38

## 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)