r''' Tracing a field to make a stream line
.. codeauthor:: Yoshifumi Futaana
Any vector field, :math:`\mathbf{B}(x, y, z)`, is
traced from a given point :math:`(x_0, y_0, z_0)`.
See also the sample script at :mod:`streamline_sample`.
'''
import numpy as np
[docs]class SimpleTracing:
r''' Simple tracing
The differential equations are
.. math::
dx = Bx(x, y, z) dt
dy = By(x, y, z) dt
dz = Bz(x, y, z) dt
For numerical calculation, we use Runge-Kutta 4.
.. math::
r_{i+1} = r_i + \frac{\Delta t}{6}(k_1 + 2k_2 + 2k_3 + k_4)
where r denotes (x, y, z). Each coefficients are
.. math::
k_1 = B(r_i)
k_2 = B(r_i + \frac{\Delta t}{2}k_1)
k_3 = B(r_i + \frac{\Delta t}{2}k_2)
k_4 = B(r_i + \Delta t k_3)
If one wants to backward trace, you may instance another object
that gives -vector_field in the constructor.
An example is shown.
Assume a vector field as
.. math::
Bx = x
By = 0
Bz = \sqrt{x^2+z^2}
The vector_field should be a function that returns a vector, i.e.,
>>> vfield = lambda p: np.array((p[0], 0, np.sqrt(p[0]**2+p[2]**2)))
The vector_field should take np.array, returning np.array() object.
>>> tracer = SimpleTracing([1, 0, 1], vfield, 0.001)
>>> tracer.step_forward()
>>> print(tracer.x, tracer.y, tracer.z, tracer.tlist[-1]) # doctest: +NORMALIZE_WHITESPACE
1.0010005001667084 0.0 1.0014150674450009 0.001
'''
def __init__(self, initpos, vector_field, dt):
self.vector_field = vector_field
self.dt = dt
self.x = initpos[0]
self.y = initpos[1]
self.z = initpos[2]
self.tlist = [0,]
self.xlist = [self.x,]
self.ylist = [self.y,]
self.zlist = [self.z,]
[docs] def step_forward(self):
h = self.dt
x, y, z = self.x, self.y, self.z
k1 = self.vector_field((x, y, z))
k2 = self.vector_field((x + h / 2. * k1[0], y + h / 2. * k1[1], z + h / 2. * k1[2]))
k3 = self.vector_field((x + h / 2. * k2[0], y + h / 2. * k2[1], z + h / 2. * k2[2]))
k4 = self.vector_field((x + h * k3[0], y + h * k3[1], z + h * k3[2]))
xnew = x + h / 6. * (k1[0] + 2 * k2[0] + 2 * k3[0] + k4[0])
ynew = y + h / 6. * (k1[1] + 2 * k2[1] + 2 * k3[1] + k4[1])
znew = z + h / 6. * (k1[2] + 2 * k2[2] + 2 * k3[2] + k4[2])
self.x, self.y, self.z = xnew, ynew, znew
self.xlist.append(xnew)
self.ylist.append(ynew)
self.zlist.append(znew)
self.tlist.append(self.tlist[-1] + h)
import unittest
import doctest
[docs]def doctests():
return unittest.TestSuite((
doctest.DocTestSuite(),
))
if __name__ == '__main__':
unittest.main(defaultTest='doctests')