""" Runge-Kutta implementations (for particle tracing)
For tutorial of particle tracing, see :ref:`tracing`.
In a general term, Runge-Kutta solves
.. math::
dy/dt = f(t, y)
The functions
- :func:`rk4_progress`
- :func:`rk5_progress`
- :func:`rkf45_progress`
will integrate the above generic differential equation by 1 step.
Classes
- :class:`RK_EqM`
- :class:`RK4_EqM`
- :class:`RK5_EqM`
- :class:`RKF45_EqM`
- :class:`RKF45_BackEqM`
provide the forward integrator of the equation of motion.
.. math::
dv_i/dt = q/m F_i
.. math::
dx_i/dt = v_i
"""
import numpy as np
import logging
_logger = logging.getLogger(__name__)
[docs]class ElectricForce:
def __init__(self, mass, charge, electric_field):
"""
:param mass: The mass of the particle, in kg.
:param charge: The charge of the particle, in C.
:param vector3dfield: An instance of :class:`irfpy.util.fields.Vector3dField`
representing the electric field (E).
"""
self._field = electric_field
self._q_per_m = charge / mass
def __call__(self, t, x, y, z, vx, vy, vz):
ex, ey, ez = self._field.get(x, y, z)
return self._q_per_m * ex, self._q_per_m * ey, self._q_per_m * ez
[docs]class MagneticForce:
def __init__(self, mass, charge, magnetic_field):
"""
:param mass: The mass of the particle, in kg.
:param charge: The charge of the particle, in C.
:param vector3dfield: An instance of :class:`irfpy.util.fields.Vector3dField`
representing the magnetic field (B).
>>> m = 1
>>> q = 1
>>> import irfpy.util.fields
>>> bfield = irfpy.util.fields.UniformVector3dField(0, 0, 1) # Bz=1
>>> vel = [1, 0, 0] # Vx=1
>>> forcefield = MagneticForce(m, q, bfield)
>>> f = forcefield(0, 0, 0, 0, vel[0], vel[1], vel[2])
>>> print(f)
[ 0. -1. 0.]
# A bit more complex example
>>> bfield = irfpy.util.fields.UniformVector3dField(-1, 2, 6)
>>> vel = [-2, -3, -4]
>>> forcefield = MagneticForce(m, q, bfield)
>>> f = forcefield(0, 0, 0, 0, vel[0], vel[1], vel[2])
>>> print(f)
[-10. 16. -7.]
"""
self._field = magnetic_field
self._q_per_m = charge / mass
def __call__(self, t, x, y, z, vx, vy, vz):
# b = self._field.get(x, y, z)
# v = [vx, vy, vz]
# return np.cross(v, b) * self._q_per_m # cross takes time...
Bx, By, Bz = self._field.get(x, y, z)
return self._q_per_m * np.array([vy * Bz - vz * By, vz * Bx - vx * Bz, vx * By - vy * Bx])
[docs]class ElectroMagneticForce:
def __init__(self, mass, charge, electric_field, magnetic_field):
self._ef = ElectricForce(mass, charge, electric_field)
self._mf = MagneticForce(mass, charge, magnetic_field)
def __call__(self, t, x, y, z, vx, vy, vz):
return self._ef(t, x, y, z, vx, vy, vz) + self._mf(t, x, y, z, vx, vy, vz)
[docs]class GravityForce:
""" Gravity force class to be used for :class:`RK`.
"""
gravity_constant = 6.674e-11
def __init__(self, center_mass):
"""
:param center_mass: Mass of the central object (or equivalent mass) in kg.
"""
self._GM = center_mass * self.gravity_constant
def __call__(self, t, x, y, z, vx, vy, vz):
r2 = x * x + y * y + z * z
force_per_mass = -self._GM / r2
r = np.sqrt(r2)
return force_per_mass * x / r, force_per_mass * y / r, force_per_mass * z / r
[docs]class RK_EqM:
r""" Classical RK.
.. math::
y_{n+1} = y_n + \frac{h}{6} (k_1 + 2 k_2 + 2 k_3 + k_4)
where
.. math::
k_1 = f(t_n, y_n) \\
k_2 = f(t_n + h/2, y_n + h/2 k_1) \\
k_3 = f(t_n + h/2, y_n + h/2 k_2) \\
k_4 = f(t_n + h, y_n + h k_3)
*Example*
Here we try to understand the Io motion around Jupiter.
The mass of Juptier is 1.90e27 kg and Io for 8.93e22.
>>> massj = 1.90e27
Io is at 4.217e5 km with velocity of 17.334 km/s at t=0.
>>> x0 = np.array([4.217e8, 0, 0])
>>> v0 = np.array([0, 17.334e3, 0])
The force per mass (N/kg) is define as a :class:`GravityForce` class.
>>> gforce = GravityForce(massj)
The :class:`RK_EqM` class is instanced then.
>>> rk = RK_EqM(x0, v0, gforce)
Let's put the time resolution as 3600 s.
>>> rk.progress(3600)
>>> print(rk.pos[-1]) # doctest: +NORMALIZE_WHITESPACE
[4.17087687e+08 6.21744701e+07 0.00000000e+00]
>>> print(rk.vel[-1]) # doctest: +NORMALIZE_WHITESPACE
[-2557.7194899 17144.4089548 0. ]
>>> print(rk.t[-1])
3600
After 21 hours, the Io goes to the other side (-x with -vy).
>>> for i in range(20):
... rk.progress(3600)
>>> print(rk.pos[-1]) # doctest: +NORMALIZE_WHITESPACE
[-4.20854860e+08 1.22957706e+07 0.00000000e+00]
>>> print(rk.vel[-1]) # doctest: +NORMALIZE_WHITESPACE
[ -506.56403226 -17353.95549486 0. ]
>>> print(rk.t[-1])
75600
"""
def __init__(self, x0, v0, forcefield):
"""
:param x0: Initial position, as numpy array with the shape of (3,)
:param v0: Initial velocity, as numpy array with the shape of (3,)
:param forcefield: Force field acting to the particle per mass. This should be a
function(-like) that must get 7 floats (t, x, y, z, vx, vy, vz) and return
the 3 floats (fx, fy, fz).
"""
if x0.shape != (3,):
raise ValueError('x0 should be (3,) shape array')
if v0.shape != (3,):
raise ValueError('v0 should be (3,) shape array')
self._forcefield = forcefield
self.pos = [x0, ]
self.vel = [v0, ]
self.t = [0, ]
[docs] def progress(self, dt):
""" Progress the particle by dt.
"""
xn = self.pos[-1][0]
yn = self.pos[-1][1]
zn = self.pos[-1][2]
vxn = self.vel[-1][0]
vyn = self.vel[-1][1]
vzn = self.vel[-1][2]
tn = self.t[-1]
### k1 = f(tn, yn)
fx, fy, fz = self._forcefield(tn, xn, yn, zn, vxn, vyn, vzn)
k1x = vxn
k1y = vyn
k1z = vzn
k1vx = fx
k1vy = fy
k1vz = fz
### k2 = f(tn + h/2, yn + h/2 k1)
fx, fy, fz = self._forcefield(tn + dt / 2.,
xn + dt / 2. * k1x,
yn + dt / 2. * k1y,
zn + dt / 2. * k1z,
vxn + dt / 2. * k1vx,
vyn + dt / 2. * k1vy,
vzn + dt / 2. * k1vz)
k2x = vxn + dt / 2. * k1vx
k2y = vyn + dt / 2. * k1vy
k2z = vzn + dt / 2. * k1vz
k2vx = fx
k2vy = fy
k2vz = fz
### k3 = f(tn + h/2, yn + h/2 k2)
fx, fy, fz = self._forcefield(tn + dt / 2.,
xn + dt / 2. * k2x,
yn + dt / 2. * k2y,
zn + dt / 2. * k2z,
vxn + dt / 2. * k2vx,
vyn + dt / 2. * k2vy,
vzn + dt / 2. * k2vz)
k3x = vxn + dt / 2. * k2vx
k3y = vyn + dt / 2. * k2vy
k3z = vzn + dt / 2. * k2vz
k3vx = fx
k3vy = fy
k3vz = fz
### k4 = f(tn + h, yn + h k3)
fx, fy, fz = self._forcefield(tn + dt,
xn + dt * k3x,
yn + dt * k3y,
zn + dt * k3z,
vxn + dt * k3vx,
vyn + dt * k3vy,
vzn + dt * k3vz)
k4x = vxn + dt * k3vx
k4y = vyn + dt * k3vy
k4z = vzn + dt * k3vz
k4vx = fx
k4vy = fy
k4vz = fz
xn1 = xn + dt / 6. * (k1x + 2 * k2x + 2 * k3x + k4x)
yn1 = yn + dt / 6. * (k1y + 2 * k2y + 2 * k3y + k4y)
zn1 = zn + dt / 6. * (k1z + 2 * k2z + 2 * k3z + k4z)
vxn1 = vxn + dt / 6. * (k1vx + 2 * k2vx + 2 * k3vx + k4vx)
vyn1 = vyn + dt / 6. * (k1vy + 2 * k2vy + 2 * k3vy + k4vy)
vzn1 = vzn + dt / 6. * (k1vz + 2 * k2vz + 2 * k3vz + k4vz)
tn1 = tn + dt
self.pos.append(np.array([xn1, yn1, zn1]))
self.vel.append(np.array([vxn1, vyn1, vzn1]))
self.t.append(tn1)
[docs]def rk4_progress(tn, yk, dt, dfunc):
""" Integration by RK4 (6 order).
"""
h = dt
### k1
#k_1 = h f(t_k, y_k) \\
k1 = dfunc(tn, yk) * h
### k2
#k_2 = h f(t_k + h/4, y_k + 1/4 k_1) \\
k2 = dfunc(tn + h / 4., yk + k1 / 4.) * h
### k3
#k_3 = h f(t_k + \frac{3}{8}h, y_k + \frac{3}{32}k_1 + \frac{9}{32}k_2) \\
k3 = dfunc(tn + 3 / 8. * h, yk + 3 / 32. * k1 + 9 / 32. * k2) * h
### k4
#k_4 = h f(t_k + \frac{12}{13}h, y_k + \frac{1932}{2197}k_1 - \frac{7200}{2197}k_2 + \frac{7296}{2197} k_3) \\
k4 = dfunc(tn + 12 / 13. * h, yk + 1932 / 2197. * k1 - 7200 / 2197. * k2 + 7296 / 2197. * k3) * h
### k5
#k_5 = h f(t_k + h, y_k + \frac{439}{216}k_1 - 8 k_2 + \frac{3680}{513} k_3 - \frac{845}{4104} k_4) \\
k5 = dfunc(tn + h, yk + 439 / 216. * k1 - 8 * k2 + 3680 / 513. * k3 - 845 / 4104. * k4) * h
#k_6 = h f(t_k + \frac{1}{2}h, y_k - \frac{8}{27}k_1 + 2 k_2 - \frac{3544}{2565} k_3 + \frac{1859}{4104} k_4 - \frac{11}{40} k_5) \\
k6 = dfunc(tn + h / 2., yk - 8 / 27. * k1 + 2 * k2 - 3544 / 2565. * k3 + 1859 / 4104. * k4 - 11 / 40. * k5) * h
# y_{k+1} = y_k + \frac{25}{216}k_1 + \frac{1408}{2565}k_3 + \frac{2197}{4101}k_4 - \frac{1}{5}k_5
yk1 = yk + 25 / 216. * k1 + 1408 / 2565. * k3 + 2197 / 4101. * k4 - k5 / 5.
tn1 = tn + h
return tn1, yk1
[docs]class RK4_EqM:
r""" Runge-Kutta with the order of 4 using 6 k-s.
.. math::
y_{k+1} = y_k + \frac{25}{216}k_1 + \frac{1408}{2565}k_3 + \frac{2197}{4101}k_4 - \frac{1}{5}k_5
where
.. math::
k_1 = h f(t_k, y_k) \\
k_2 = h f(t_k + h/4, y_k + k_1/4) \\
k_3 = h f(t_k + \frac{3}{8}h, y_k + \frac{3}{32}k_1 + \frac{9}{32}k_2) \\
k_4 = h f(t_k + \frac{12}{13}h, y_k + \frac{1932}{2197}k_1 - \frac{7200}{2197}k_2 + \frac{7296}{2197} k_3) \\
k_5 = h f(t_k + h, y_k + \frac{439}{216}k_1 - 8 k_2 + \frac{3680}{513} k_3 - \frac{845}{4104} k_4) \\
k_6 = h f(t_k + \frac{1}{2}h, y_k - \frac{8}{27}k_1 + 2 k_2 - \frac{3544}{2565} k_3 + \frac{1859}{4104} k_4 - \frac{11}{40} k_5) \\
*Example*
Here we try to understand the Io motion around Jupiter.
The mass of Juptier is 1.90e27 kg and Io for 8.93e22.
>>> massj = 1.90e27
Io is at 4.217e5 km with velocity of 17.334 km/s at t=0.
>>> x0 = np.array([4.217e8, 0, 0])
>>> v0 = np.array([0, 17.334e3, 0])
The force per mass (N/kg) is define as a :class:`GravityForce` class.
>>> gforce = GravityForce(massj)
The :class:`RK4_EqM` class is instanced then.
>>> rk = RK4_EqM(x0, v0, gforce)
Let's put the time resolution as 3600 s.
>>> rk.progress(3600)
>>> print(rk.pos[-1]) # doctest: +NORMALIZE_WHITESPACE
[4.17084349e+08 6.21989712e+07 0.00000000e+00]
>>> print(rk.vel[-1]) # doctest: +NORMALIZE_WHITESPACE
[-2558.70808725 17144.27369938 0. ]
>>> print(rk.t[-1])
3600
After 21 hours, the Io goes to the other side (-x with -vy).
>>> for i in range(20):
... rk.progress(3600)
>>> print(rk.pos[-1]) # doctest: +NORMALIZE_WHITESPACE
[-4.20899082e+08 1.18280482e+07 0.00000000e+00]
>>> print(rk.vel[-1])
[ -487.75930547 -17353.750109 0. ]
>>> print(rk.t[-1])
75600
"""
def __init__(self, x0, v0, forcefield):
"""
:param x0: Initial position, as numpy array with the shape of (3,)
:param v0: Initial velocity, as numpy array with the shape of (3,)
:param forcefield: Force field acting to the particle per mass. This should be a
function(-like) that must get 7 floats (t, x, y, z, vx, vy, vz) and return
the 3 floats (fx, fy, fz).
"""
if x0.shape != (3,):
raise ValueError('x0 should be (3,) shape array')
if v0.shape != (3,):
raise ValueError('v0 should be (3,) shape array')
self._forcefield = forcefield
self.pos = [x0, ]
self.vel = [v0, ]
self.t = [0, ]
def __f(self, t, xyzvxyz):
""" 6 dimensional function for integration.
:param t: Time
:param xyzvxyz: 6 dimensional position + velocity
:returns: 6 dimensional array, representing k_i vector.
"""
x, y, z, vx, vy, vz = xyzvxyz
retx = vx
rety = vy
retz = vz
force = self._forcefield(t, x, y, z, vx, vy, vz)
retvx, retvy, retvz = force
return np.array([retx, rety, retz, retvx, retvy, retvz])
[docs] def progress(self, dt):
""" Progress the particle by dt.
"""
xn = self.pos[-1][0]
yn = self.pos[-1][1]
zn = self.pos[-1][2]
vxn = self.vel[-1][0]
vyn = self.vel[-1][1]
vzn = self.vel[-1][2]
tn = self.t[-1]
tn1, yk1 = rk4_progress(tn, np.array([xn, yn, zn, vxn, vyn, vzn]), dt, self.__f)
self.pos.append(yk1[:3])
self.vel.append(yk1[3:])
self.t.append(tn1)
[docs]def rk5_progress(tn, yk, dt, dfunc):
h = dt
### k1
#k_1 = h f(t_k, y_k) \\
k1 = dfunc(tn, yk) * h
### k2
#k_2 = h f(t_k + h/4, y_k + 1/4 k_1) \\
k2 = dfunc(tn + h / 4., yk + k1 / 4.) * h
### k3
#k_3 = h f(t_k + \frac{3}{8}h, y_k + \frac{3}{32}k_1 + \frac{9}{32}k_2) \\
k3 = dfunc(tn + 3 / 8. * h, yk + 3 / 32. * k1 + 9 / 32. * k2) * h
### k4
#k_4 = h f(t_k + \frac{12}{13}h, y_k + \frac{1932}{2197}k_1 - \frac{7200}{2197}k_2 + \frac{7296}{2197} k_3) \\
k4 = dfunc(tn + 12 / 13. * h, yk + 1932 / 2197. * k1 - 7200 / 2197. * k2 + 7296 / 2197. * k3) * h
### k5
#k_5 = h f(t_k + h, y_k + \frac{439}{216}k_1 - 8 k_2 + \frac{3680}{513} k_3 - \frac{845}{4104} k_4) \\
k5 = dfunc(tn + h, yk + 439 / 216. * k1 - 8 * k2 + 3680 / 513. * k3 - 845 / 4104. * k4) * h
#k_6 = h f(t_k + \frac{1}{2}h, y_k - \frac{8}{27}k_1 + 2 k_2 - \frac{3544}{2565} k_3 + \frac{1859}{4104} k_4 - \frac{11}{40} k_5) \\
k6 = dfunc(tn + h / 2., yk - 8 / 27. * k1 + 2 * k2 - 3544 / 2565. * k3 + 1859 / 4104. * k4 - 11 / 40. * k5) * h
# y_{k+1} = y_k + \frac{16}{135} k_1 + \frac{6656}{12825} k_3 + \frac{28561}{56430} k_4 - \frac{9}{50} k_5 + \frac{2}{55} k_6
yk1 = yk + 16 / 135. * k1 + 6656 / 12825. * k3 + 28561 / 56430. * k4 - 9 / 50. * k5 + 2 / 55. * k6
tn1 = tn + h
return tn1, yk1
[docs]class RK5_EqM:
r""" Runge-Kutta with the order of 4 using 6 k-s.
.. math::
y_{k+1} = y_k + \frac{16}{135} k_1 + \frac{6656}{12825} k_3 + \frac{28561}{56430} k_4 - \frac{9}{50} k_5 + \frac{2}{55} k_6
where
.. math::
k_1 = h f(t_k, y_k) \\
k_2 = h f(t_k + h/4, y_k + k_1/4) \\
k_3 = h f(t_k + \frac{3}{8}h, y_k + \frac{3}{32}k_1 + \frac{9}{32}k_2) \\
k_4 = h f(t_k + \frac{12}{13}h, y_k + \frac{1932}{2197}k_1 - \frac{7200}{2197}k_2 + \frac{7296}{2197} k_3) \\
k_5 = h f(t_k + h, y_k + \frac{439}{216}k_1 - 8 k_2 + \frac{3680}{513} k_3 - \frac{845}{4104} k_4) \\
k_6 = h f(t_k + \frac{1}{2}h, y_k - \frac{8}{27}k_1 + 2 k_2 - \frac{3544}{2565} k_3 + \frac{1859}{4104} k_4 - \frac{11}{40} k_5) \\
*Example*
Here we try to understand the Io motion around Jupiter.
The mass of Juptier is 1.90e27 kg and Io for 8.93e22.
>>> massj = 1.90e27
Io is at 4.217e5 km with velocity of 17.334 km/s at t=0.
>>> x0 = np.array([4.217e8, 0, 0])
>>> v0 = np.array([0, 17.334e3, 0])
The force per mass (N/kg) is define as a :class:`GravityForce` class.
>>> gforce = GravityForce(massj)
The :class:`RK5_EqM` class is instanced then.
>>> rk = RK5_EqM(x0, v0, gforce)
Let's put the time resolution as 3600 s.
>>> rk.progress(3600)
>>> print(rk.pos[-1]) # doctest: +NORMALIZE_WHITESPACE
[4.17087705e+08 6.21747255e+07 0.00000000e+00]
>>> print(rk.vel[-1]) # doctest: +NORMALIZE_WHITESPACE
[-2557.71081377 17144.41010741 0. ]
>>> print(rk.t[-1])
3600
After 21 hours, the Io goes to the other side (-x with -vy).
>>> for i in range(20):
... rk.progress(3600)
>>> print(rk.pos[-1]) # doctest: +NORMALIZE_WHITESPACE
[-4.20862773e+08 1.23126860e+07 0.00000000e+00]
>>> print(rk.vel[-1])
[ -507.29360568 -17353.63323363 0. ]
>>> print(rk.t[-1])
75600
"""
def __init__(self, x0, v0, forcefield):
"""
:param x0: Initial position, as numpy array with the shape of (3,)
:param v0: Initial velocity, as numpy array with the shape of (3,)
:param forcefield: Force field acting to the particle per mass. This should be a
function(-like) that must get 7 floats (t, x, y, z, vx, vy, vz) and return
the 3 floats (fx, fy, fz).
"""
if x0.shape != (3,):
raise ValueError('x0 should be (3,) shape array')
if v0.shape != (3,):
raise ValueError('v0 should be (3,) shape array')
self._forcefield = forcefield
self.pos = [x0, ]
self.vel = [v0, ]
self.t = [0, ]
def __f(self, t, xyzvxyz):
""" 6 dimensional function for integration.
:param t: Time
:param xyzvxyz: 6 dimensional position + velocity
:returns: 6 dimensional array, representing k_i vector.
"""
x, y, z, vx, vy, vz = xyzvxyz
retx = vx
rety = vy
retz = vz
force = self._forcefield(t, x, y, z, vx, vy, vz)
retvx, retvy, retvz = force
return np.array([retx, rety, retz, retvx, retvy, retvz])
[docs] def progress(self, dt):
""" Progress the particle by dt.
"""
xn = self.pos[-1][0]
yn = self.pos[-1][1]
zn = self.pos[-1][2]
vxn = self.vel[-1][0]
vyn = self.vel[-1][1]
vzn = self.vel[-1][2]
tn = self.t[-1]
tn1, yk1 = rk5_progress(tn, np.array([xn, yn, zn, vxn, vyn, vzn]), dt, self.__f)
self.pos.append(yk1[:3])
self.vel.append(yk1[3:])
self.t.append(tn1)
[docs]def rk45_progress(tn, yk, h, dfunc):
""" 4- and 5-dimensional Runge-Kutta solver.
The following two calss are equivalent.
.. code-block:: py
rk4, rk5 = rk45_progress(tn, yk, h, dfunc)
rk4, rk5 = rk4_progress(tn, yk, h, dfunc), rk5_progress(tn, yk, h, dfunc)
But for performance reasons, the results of :func:`rk4_progress` and :func:`rk5_rk5_progress` is
calculaated & returned together by this function.
:param tn: The time step to be fed into dfunc
:param yk: The y coordinate, an (N,) shaped array
:param h: Delta to proceed. Sclara.
:param dfunc: Differential function. It should get the argument (t, y0, y1, ... yN-1), returning the (N,) shaped np-array.
:returns: Two np-array. 4Dim and 5Dim.
"""
k1 = dfunc(tn, yk) * h
k2 = dfunc(tn + h / 4., yk + k1 / 4.) * h
k3 = dfunc(tn + 3 / 8. * h, yk + 3 / 32. * k1 + 9 / 32. * k2) * h
k4 = dfunc(tn + 12 / 13. * h, yk + 1932 / 2197. * k1 - 7200 / 2197. * k2 + 7296 / 2197. * k3) * h
k5 = dfunc(tn + h, yk + 439 / 216. * k1 - 8 * k2 + 3680 / 513. * k3 - 845 / 4104. * k4) * h
k6 = dfunc(tn + h / 2., yk - 8 / 27. * k1 + 2 * k2 - 3544 / 2565. * k3 + 1859 / 4104. * k4 - 11 / 40. * k5) * h
yk4 = yk + 25 / 216. * k1 + 1408 / 2565. * k3 + 2197 / 4101. * k4 - k5 / 5.
yk5 = yk + 16 / 135. * k1 + 6656 / 12825. * k3 + 28561 / 56430. * k4 - 9 / 50. * k5 + 2 / 55. * k6
return yk4, yk5
[docs]def rkf45_progress(tn, yk, dt, dfunc, dtmin=None, dtmax=None, tol=1e-5, reltol=1e-8):
""" Single step for RKF.
:param tn: The time. To be fed to ``dfunc``.
:param yk: The y-value at time *tn*.
:param dt: Time step to evaluate first. The exact time step is defined automatically.
:param dfunc: Function to be integrated. This should eat *tn* and *y* at *tn* (array), returning the *dy/dt* value (array).
:param dtmin: The minimum time step for integration. Default is *dt*/1000.
:param dtmax: The maximum time step for integration. Default is *dt* * 1000.
:keyword tol: Absolute tolerance. Scalar or array with the same shape of ``yk``.
:keyword reltol: Relative tolerance. Scalar or array with the same shape of ``yk``.
*Algorithm*
- Obtain *rk4* = :func:`rk4_progress` and *rk5* = :func:`rk5_progress`.
- Error is *E=|rk5 - rk4|*
- Tolerance is the bigger one out of *tol* and *reltol* * *(rk5 - yk)*.
- Compare error and tolerance.
- If error > tolerance: *tn* is halven, then recursive call. (Too bad approximation)
- If error / tolerance is between 0.2 and 1: *rk5* is the progressed value.
- If error < 0.2 * tolerance: *tn* is increased by 1.5 times, then recursive call. (Too good approximation)
"""
logger = logging.getLogger(__name__ + '.rkf45_progress')
if dtmin is None:
dtmin = dt / 1000.
if dtmax is None:
dtmax = dt * 1000.
if dt > dtmax:
h = dtmax
elif dt < dtmin:
h = dtmin
else:
h = dt
abstolarr = tol + np.zeros_like(yk) # Absolute tolerance is now array
reltolarr = reltol + np.zeros_like(yk) # Relative tolerance is now array
yk4, yk5 = rk45_progress(tn, yk, h, dfunc)
# err = np.sqrt(((yk5 - yk4) ** 2).sum())
err = np.abs(yk5 - yk4) # Absolute error
abstolarr_rel = np.abs(yk5 - yk) * reltolarr # Absolute tolerance derived from the relative tolerance
tolarr = np.max([abstolarr, abstolarr_rel], axis=0) # The tolerance is the larger one for the rel & abs
np.where(np.isnan(tolarr), np.inf, tolarr)
logger.info('T={} : DT={} : MAXERR(Err/Tor)@{}={}'.format(tn, dt, (err/tolarr).argmax(), max(err/tolarr)))
# logger.debug('RK5 : {}'.format(yk5))
# logger.debug('RK4 : {}'.format(yk4))
# logger.debug('ERR : {}'.format(err))
# logger.debug('TOL(REL): {}'.format(abstolarr_rel))
# logger.debug('TOL(ABS): {}'.format(abstolarr))
# logger.debug('TOL : {}'.format(tolarr))
# logger.debug('ERRRATIO: {}'.format(err/tolarr))
# logger.debug('TOOLONG?: {}'.format((err > tolarr).any()))
# logger.debug('TOOSHORT: {}'.format((err <= tolarr / 5.0).all()))
if (err > tolarr).any():
### NG
# dt_new = 0.8 * dt * np.power(tol * dt / err, 0.25)
dt_new = dt * 0.5 # Halven
if dt_new <= dtmin:
dt_new = dtmin
logger.warning("MIN DT IS USED")
tn5, yn5 = rk5_progress(tn, yk, dt_new, dfunc)
tn4, yn4 = rk4_progress(tn, yk, dt_new, dfunc)
return tn5, yn5, np.abs(tn5 - tn4)
else:
# Call recursively
return rkf45_progress(tn, yk, dt_new, dfunc, dtmin=dtmin, dtmax=dtmax, tol=tol, reltol=reltol)
elif (err <= tolarr / 5.0).all():
dt_new = dt * 1.5 # More coarse resolution should be fine.
if dt_new >= dtmax:
dt_new = dtmax
logger.warning("MAX DT IS USED")
tn5, yn5 = rk5_progress(tn, yk, dt_new, dfunc)
tn4, yn4 = rk4_progress(tn, yk, dt_new, dfunc)
return tn5, yn5, np.abs(tn5 - tn4)
return rkf45_progress(tn, yk, dt_new, dfunc, dtmin=dtmin, dtmax=dtmax, tol=tol, reltol=reltol)
tn1 = tn + h
return tn1, yk5, err
[docs]class RKF45_EqM:
""" Solve equation of motion by RKF.
>>> from irfpy.util import fields
>>> efield = fields.UniformVector3dField(0, 1e-6, 0)
>>> eforce = ElectricForce(1.67e-27, 1.6e-19, efield)
>>> x0 = np.array([0, 0, 0])
>>> v0 = np.array([1, 0, 0])
>>> rk = RKF45_EqM(x0, v0, eforce, 30e-3, tolx=0.01, tolv=0.001)
>>> rk.progress()
>>> rk.progress()
"""
def __init__(self, x0, v0, forcefield, dt0, dtmin=None, dtmax=None, tolx=1e-5, tolv=1e-5, reltolx=1e-8, reltolv=1e-8):
"""
:keyword dtmin: Minimum h allowed. Default is 0.01*dt0.
:keyword dtmax: Maximum h allowed. Default is 100*dt0.
"""
if x0.shape != (3,):
raise ValueError('x0 should be (3,) shape array')
if v0.shape != (3,):
raise ValueError('v0 should be (3,) shape array')
self._forcefield = forcefield
self.pos = [x0, ]
self.vel = [v0, ]
self.t = [0, ]
self.err = [np.zeros(6),]
self._dtnow = dt0
self._dtmin = dtmin
self._dtmax = dtmax
self._tol = np.array([tolx, tolx, tolx, tolv, tolv, tolv])
self._reltol = np.array([reltolx, reltolx, reltolx, reltolv, reltolv, reltolv])
self.increase_rate = 1.01 # The time step is multiplied every time the calculation is conducted.
if dtmin is None:
self._dtmin = dt0 * 0.001
if dtmax is None:
self._dtmax = dt0 * 1000
def __f(self, t, xyzvxyz):
""" 6 dimensional function for integration.
:param t: Time
:param xyzvxyz: 6 dimensional position + velocity
:returns: 6 dimensional array, representing k_i vector.
"""
x, y, z, vx, vy, vz = xyzvxyz
retx = vx
rety = vy
retz = vz
force = self._forcefield(t, x, y, z, vx, vy, vz)
retvx, retvy, retvz = force
return np.array([retx, rety, retz, retvx, retvy, retvz])
[docs] def progress(self):
""" Progress the particle by dt.
"""
xn = self.pos[-1][0]
yn = self.pos[-1][1]
zn = self.pos[-1][2]
vxn = self.vel[-1][0]
vyn = self.vel[-1][1]
vzn = self.vel[-1][2]
tn = self.t[-1]
tn1, yk1, tol = rkf45_progress(tn, np.array([xn, yn, zn, vxn, vyn, vzn]), self._dtnow, self.__f,
dtmin=self._dtmin, dtmax=self._dtmax, tol=self._tol, reltol=self._reltol)
self.pos.append(yk1[:3])
self.vel.append(yk1[3:])
self._dtnow = (tn1 - self.t[-1]) * self.increase_rate # Slightly increasing time step in time
self.t.append(tn1)
self.err.append(tol)
def __str__(self):
t = self.t[-1]
x = self.pos[-1]
v = self.vel[-1]
return "t={} x=[{} {} {}] v=[{} {} {}]".format(t, x[0], x[1], x[2], v[0], v[1], v[2])
[docs]class RKF45_BackEqM:
""" Backward integration of the equation of motion.
>>> from irfpy.util import fields
>>> efield = fields.UniformVector3dField(0, 1e-6, 0)
>>> eforce = ElectricForce(1.67e-27, 1.6e-19, efield)
>>> x0 = np.array([0, 0, 0])
>>> v0 = np.array([1, 0, 0])
>>> rk = RKF45_BackEqM(x0, v0, eforce, 30e-3, tolx=0.000001, tolv=0.00001)
>>> rk.retrogress()
>>> rk.retrogress()
>>> rk.retrogress()
"""
def __init__(self, x0, v0, forcefield, dt0, dtmin=None, dtmax=None, tolx=1e-5, tolv=1e-5, reltolx=1e-8, reltolv=1e-8):
"""
:param dt0: The time step, must be positive!
:keyword dtmin: Minimum h allowed. Default is 0.001*dt0.
:keyword dtmax: Maximum h allowed. Default is 1000*dt0.
"""
if x0.shape != (3,):
raise ValueError('x0 should be (3,) shape array')
if v0.shape != (3,):
raise ValueError('v0 should be (3,) shape array')
self._forcefield = forcefield
self.pos = [x0, ]
self.vel = [v0, ]
self.t = [0, ]
self.err = [np.zeros(6),]
self._dtnow = np.abs(dt0) # Implementation-wise, the dt0 should be positive.
self._dtmin = dtmin
self._dtmax = dtmax
self._tol = np.array([tolx, tolx, tolx, tolv, tolv, tolv])
self._reltol = np.array([reltolx, reltolx, reltolx, reltolv, reltolv, reltolv])
self.increase_rate = 1.1 # The time step is multiplied every time the calculation is conducted.
if dtmin is None:
self._dtmin = dt0 * 0.001
if dtmax is None:
self._dtmax = dt0 * 1000
def __f(self, t, xyzvxyz):
""" 6 dimensional function for integration.
:param t: Time
:param xyzvxyz: 6 dimensional position + velocity
:returns: 6 dimensional array, representing k_i vector.
"""
x, y, z, vx, vy, vz = xyzvxyz
retx = -vx
rety = -vy
retz = -vz
force = self._forcefield(t, x, y, z, vx, vy, vz)
retvx, retvy, retvz = force
return np.array([retx, rety, retz, -retvx, -retvy, -retvz])
[docs] def retrogress(self):
""" Progress the particle by dt.
"""
xn = self.pos[-1][0]
yn = self.pos[-1][1]
zn = self.pos[-1][2]
vxn = self.vel[-1][0]
vyn = self.vel[-1][1]
vzn = self.vel[-1][2]
tn = -self.t[-1]
tn1, yk1, tol = rkf45_progress(tn, np.array([xn, yn, zn, vxn, vyn, vzn]), self._dtnow, self.__f,
dtmin=self._dtmin, dtmax=self._dtmax, tol=self._tol, reltol=self._reltol)
self.pos.append(yk1[:3])
self.vel.append(yk1[3:])
self._dtnow = (tn1 + self.t[-1]) * self.increase_rate # Slightly increasing time step in time
self.t.append(-tn1)
self.err.append(tol)
def __str__(self):
t = self.t[-1]
x = self.pos[-1]
v = self.vel[-1]
return "t={} x=[{} {} {}] v=[{} {} {}]".format(t, x[0], x[1], x[2], v[0], v[1], v[2])