3
\$\begingroup\$

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()
\$\endgroup\$
3
  • 1
    \$\begingroup\$ Anyway, this will not run: your velfun expects four parameters but you only ever pass three. \$\endgroup\$ – Reinderien May 14 at 14:23
  • \$\begingroup\$ 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 \$\endgroup\$ – ben May 15 at 6:59
  • \$\begingroup\$ 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. \$\endgroup\$ – Lutz Lehmann May 23 at 7:08
2
\$\begingroup\$

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.

plot

\$\endgroup\$
1
  • 1
    \$\begingroup\$ 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. \$\endgroup\$ – Andrew May 14 at 20:38
2
\$\begingroup\$

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)
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.