#
# module to perform 4th Order Runge-Kutta with adaptive
# timestepping on systems of the form:
#
# -->
# d x -->
# ------ = func( x , t)
# dt
#
# R. Orvedahl 9-6-2014
# sample RHS routine:
# # RHS of the equations of motion, i.e. defines f(x,t):
# # xdot = f(x, t)
# # x is the input vector and t is time
# # return value must be a numpy array
# def rhs(x, t):
#
# xdot = x*t
#
# return numpy.array(xdot)
import math
import numpy as np
# integrate the equations of motion:
# ydot = f(y, t)
# using 4th order R-K method with an adaptive stepsize, to try to
# achieve the relative error err.
#
# if err < 0, then we don't do adaptive stepping, but rather
# we always walk at the input dt
#
# X0 --> intial state vector
# t0 --> initial time
# dt --> initial timestep
# tmax --> integrate until t=tmax
# err --> relative error for adaptive stepsize
# rhs --> right hand side of equations of motion
# S1, S2 --> adaptive timestepping parameters
def intRK4(X0, t0, dt, tmax, err, rhs, S1=0.9, S2=4.0):
# keep track of how many steps were taken
steps = 0
# convert x0 to numpy array
X0 = np.array(X0)
# initial conditions
t = t0
x = X0
# store the history for plotting
tpoints = [t]
xpoints = [x]
# start with the old timestep
dtNew = dt
while (t < tmax):
if (err > 0.0):
# adaptive stepping
# iteration loop -- keep trying to take a step until
# we achieve our desired error
relError = 1.e10
while (relError > err):
dt = dtNew
if t+dt > tmax:
dt = tmax-t
# take 2 half steps
xtmp = RK4_singlestep(x, t, 0.5*dt, rhs)
xnew = RK4_singlestep(xtmp, t+0.5*dt, 0.5*dt, rhs)
# now take just a single step to cover dt
xsingle = RK4_singlestep(x, t, dt, rhs)
# {x,y,u,v}double should be more accurate that
# {x,y,u,v}single, since it used smaller steps.
# estimate the relative error now
relError = np.amax(np.abs((xnew-xsingle)/xnew))
# adaptive timestep algorithm from Garcia (Eqs. 3.30
# and 3.31)
dtEst = dt*np.abs(err/relError)**0.2
dtNew = np.amin(np.amax(S1*dtEst, dt/S2), S2*dt)
# keep track of step count
steps += 1
else:
if t+dt > tmax:
dt = tmax-t
# take just a single step to cover dt
xnew = RK4_singlestep(x, t, dt, rhs)
# keep track of step count
steps += 1
t += dt
# store t in single array
# xpoints is a list of arrays where each array is the size of the
# system to solve. So if there are 4 unknowns in the system, each
# xpoints element will be a size=4 array
tpoints.append(t)
xpoints.append(xnew)
# set for the next step
x = xnew
# return the trajectory
tpoints = np.array(tpoints); xpoints = np.array(xpoints)
return tpoints, xpoints, steps
# take a single RK-4 timestep from t to t+dt for the system ydot = rhs
def RK4_singlestep(X0, t, dt, rhs):
# X0 is an array that holds the initial state
# t is current time
# dt is stepsize
# rhs is right hand side of the system to solve
# x holds the new state
x = X0
# get the RHS at several points
xdot1 = rhs(x, t)
xdot2 = rhs(x+0.5*dt*xdot1, t+0.5*dt)
xdot3 = rhs(x+0.5*dt*xdot2, t+0.5*dt)
xdot4 = rhs(x+dt*xdot3, t+dt)
# advance
xnew = x + (dt/6.0)*(xdot1 + 2.0*xdot2 + 2.0*xdot3 + xdot4)
return xnew
##########################################################################
# test the RK4 routine
##########################################################################
if __name__ == "__main__":
t0 = 0.0
tmax = 1.0
dt = 0.01
err = 1.e-8
# RHS of the equations of motion, i.e. defines f(x,t):
# xdot = f(x, t)
# x is the input vector and t is time
def rhs(x, t):
# unpack incoming array
xx = x[0]
yy = x[1]
# evalute the RHS for each component
xdot = -2*xx + t + 4.
ydot = np.exp(-.5*t)
# store output as numpy array
fdot = []
fdot.append(xdot)
fdot.append(ydot)
return np.array(fdot)
# initial conditions
X0 = []
X0.append(1.0)
X0.append(4.0)
# exact solutions for above initial conditions
def exact(t):
xsol = -0.75*np.exp(-2.*t) + 0.5*t + 1.75
ysol = -2.*np.exp(-.5*t) + 6.
fsol = []
fsol.append(xsol)
fsol.append(ysol)
return fsol
# Integrate:
# returns tpts as an array and xpts as a list of arrays
tpts, xpts, steps = intRK4(X0, t0, dt, tmax, err, rhs, S1=0.9, S2=4.0)
if (err < 0.):
print "\nNumerical (Classical):"
else:
print "\nNumerical (Adaptive):"
print "\ttmax, x(tmax), y(tmax)"
# -1 index is solution at tmax and the 0,1 indices are the x,y solutions
print "\t",tpts[-1], xpts[-1][0], xpts[-1][1]
print "\tNumber of steps: ",steps
print "\nExact:"
print "\ttmax, x(tmax), y(tmax)"
xx, yy = exact(tpts[-1])
print "\t",tpts[-1], xx, yy
print