Runge-Kutta with wrapper

Home > Routine

See the codes for each example at the end of the post.

Example 1

Let us solve the differential equations \frac{dx}{dt}=-x. The following wrapper uses Runge-Kutta solver from scipy. The purpose of the wrapper is to compute the specified points \{(x_i,t_i)\}.

  1. Input time_checkpoints into RK45_wrapper. time_checkpoints is the list [t1,t2,…,tN] whose values of x we want to compute.
  2. x0 is the initial value, i.e. the value of x at t1. This is the initial condition for the initial value problem.
  3. stepsize. This sets the maximum step size between tk and t(k+1). For example, if tk=1 and t(k+1)=2, stepsize can be set to 0.1 and the values will be computed at 1.1, 1.2, …, 2.0. By default stepsize = None.  When set to None, the stepsize is a-hundreth of tk and t(k+1) interval.

As seen in the figure below, the green “check points” are the values that we specify for the wrapper to compute via time_checkpoints. Note that t_checkpoint_new will be computed, since the solver DOES NOT exactly computes the values at t_checkpoint. However, by setting the stepsize as small as possible, t_checkpoint_new will be closer to t_checkpoint. We print max checkpt diff to show the difference between t_checkpoint_new and t_checkpoint. In this example, using the default stepsize is sufficiently close, i.e. max checkpt diff=0.

Red circles are the solutions generated by RK45 solver, and blue curve is the solution plotted analytically.

rk45wrapper

Example 2

Now we try out similarly with pendulum equation example like the example here. This shows that essentially the same method works for multi-valued function as well. The equation is given by

\ddot{\theta}(t)+ b \cdot \dot{\theta}(t) +c\cdot sin(\theta (t)) = 0

which can be linearized as a system of equations:

\dot{\theta}(t) = \omega(t)
\dot{\omega}(t) = -b\cdot \omega (t) - c\cdot sin(\theta (t))

We get the following. Each color correspond to one component in the linearized system. The scatter dots are solutions solved by RK45, the cross and circle are respectively the checkpoints we want to compute and the same checkpoints produced by RK45.

rk45wrap2.JPG

Code for Example 1

import numpy as np
from scipy.integrate import RK45
import matplotlib.pyplot as plt

def f(t,y):
	# dy/dt = f(t,y)
	return np.array([-1*y[0]])

x0 = np.array([1.0]) # initial value
# t0 = 0  # initial time
# t1 = 10 # final time

def RK45_wrapper(x0, time_checkpoints, f, stepsize=None):

	if stepsize is None:
		stepsize = 0.01*(time_checkpoints[-1]-time_checkpoints[0])
	t_set = [time_checkpoints[0]]
	x_set = [x0]
	t_checkpoint_new = [time_checkpoints[0]]
	x_checkpoint = [x0]

	for i in range(len(time_checkpoints)-1):	
		t1 = time_checkpoints[i+1]
		if i>0:
			x0 = ox.y
			t0 = ox.t
			
		else:
			t0 = time_checkpoints[i]
		ox = RK45(f,t0,x0,t1,max_step=stepsize)
		
		while ox.t < t1:
			msg = ox.step()
			t_set.append(ox.t)
			x_set.append(ox.y)
		x_checkpoint.append(ox.y)
		t_checkpoint_new.append(ox.t)
	return x_set, t_set, x_checkpoint, t_checkpoint_new

# all data by DE (differential equation)
time_checkpoints = np.linspace(0,10,11)
x_set, t_set, x_checkpoint, t_checkpoint_new = RK45_wrapper(x0, time_checkpoints, f, stepsize=None)

t_checkpt_diff = []
for t,t_new in zip(time_checkpoints, t_checkpoint_new):
	t_checkpt_diff.append(abs(t-t_new))
print("max checkpt diff = ", np.max(t_checkpt_diff))

# all data, analytic
t0 = time_checkpoints[0]
t1 = time_checkpoints[-1]
t_set2 = np.linspace(t0,t1,100)
f_int = lambda t: np.exp(-t)
x_set2 = [f_int(t_set2[i]) for i in range(len(t_set2))]

fig = plt.figure()
ax=fig.add_subplot(111)
ax.scatter(t_set,x_set,facecolors='none',edgecolors='r',label="all data by DE")
ax.scatter(t_checkpoint_new,x_checkpoint,20,'g',label="check points")
ax.plot(t_set2,x_set2,label="all data, analytic")
ax.set_xlabel("t")
ax.set_ylabel("x")
ax.legend()
plt.show()

Code for Example 2

import numpy as np
from scipy.integrate import RK45
import matplotlib.pyplot as plt

def f(t,y, f_args):
	theta, omega = y
	dydt = [omega, -f_args[0]*omega - f_args[1]*np.sin(theta)]
	return dydt

def RK45_wrapper(x0, time_checkpoints, f, f_args, stepsize=None):
	if stepsize is None:
		stepsize = 0.01*(time_checkpoints[-1]-time_checkpoints[0])
	t_set = [time_checkpoints[0]]
	x_set = [x0]
	t_checkpoint_new = [time_checkpoints[0]]
	x_checkpoint = [x0]

	f_temp = lambda t,y : f(t,y, f_args)

	for i in range(len(time_checkpoints)-1):	
		t1 = time_checkpoints[i+1]
		if i>0:
			x0 = ox.y
			t0 = ox.t
			
		else:
			t0 = time_checkpoints[i]
		ox = RK45(f_temp,t0,x0,t1,max_step=stepsize)
		
		while ox.t < t1:
			msg = ox.step()
			t_set.append(ox.t)
			x_set.append(ox.y)
		x_checkpoint.append(ox.y)
		t_checkpoint_new.append(ox.t)
	return x_set, t_set, x_checkpoint, t_checkpoint_new

b = 0.25
c = 5.0
f_args = [b,c]
x0 = [np.pi - 0.1, 0.0]
time_checkpoints = np.linspace(0, 10, 11)
x_set, t_set, x_checkpoint, t_checkpoint_new = RK45_wrapper(x0, time_checkpoints, f, f_args, stepsize=None)

# print(np.array(x_set).shape) # (112, 2)
fig=plt.figure()
ax = fig.add_subplot(111)
ax.scatter(t_set,[x[0] for x in x_set],3,label="RK45 x[0]")
ax.scatter(t_set,[x[1] for x in x_set],3,label="RK45 x[1]")
ax.scatter(t_checkpoint_new,[x[0] for x in x_checkpoint],50,label="output ckpt x[1]",facecolors='none',edgecolors='b')
ax.scatter(t_checkpoint_new,[x[1] for x in x_checkpoint],50,label="output ckpt x[1]",facecolors='none',edgecolors='r')
ax.scatter(time_checkpoints,[x[0] for x in x_checkpoint],30,label="manual ckpt x[1]",marker="x",color="b")
ax.scatter(time_checkpoints,[x[1] for x in x_checkpoint],30,label="manual ckpt x[1]",marker="x",color="r")

ax.set_xlabel("x")
ax.set_ylabel("y")
ax.legend()
plt.show()