Source code for irfpy.util.streamline

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')