streamline_sample

Sample script of stream lines

'''Sample script of stream lines
'''
import matplotlib.pyplot as plt
import numpy as np

import irfpy.util.streamline


def draw_vortex(ax):

    # In x, y plane, vector field at (x, y) is (-y, x).
    def vector_field_vortex(pos):
        return np.array([-pos[1], pos[0], 0])

    # Initialize
    strm = irfpy.util.streamline.SimpleTracing((1, 0, 0), vector_field_vortex, 0.01)

    for t in range(400):
        strm.step_forward()


    # Backward
    strm2 = irfpy.util.streamline.SimpleTracing((1, 0, 0), vector_field_vortex, -0.01)

    for t in range(200):
        strm2.step_forward()

    ax.plot(strm.xlist, strm.ylist, 'b')
    ax.plot(strm2.xlist, strm2.ylist, 'r')
    ax.set_aspect('equal')


def draw_dipole(ax):

    # In x, z plane, the vector field is r = 2cos t and t = sin t.
    def vector_field_dipole(pos):
        x, y, z = pos
        r2 = x * x + z * z
        th = np.arccos(z / np.sqrt(r2))
        r = 2 * np.cos(th)
        t = np.sin(th)
        lx = r * np.sin(th) + t * np.cos(th)
        lz = r * np.cos(th) - t * np.sin(th)
        return np.array([lx, 0, lz])


    strm = irfpy.util.streamline.SimpleTracing((0.5, 0, 0), vector_field_dipole, 0.001)
    # The vectof_field length is of the order of 1.
    # dt is 0.001, thus, the segment length per step is about 0.001.

    for t in range(30000):
        strm.step_forward()
        if strm.x ** 2 + strm.y ** 2 < 0.001 ** 2:
            break
    ax.plot(strm.xlist, strm.zlist, 'b')

    strm = irfpy.util.streamline.SimpleTracing((0.5, 0, 0), vector_field_dipole, -0.001)
    for t in range(30000):
        strm.step_forward()
        if strm.x ** 2 + strm.y ** 2 < 0.001 ** 2:
            break
    ax.plot(strm.xlist, strm.zlist, 'r-x')


    # Indeed, direction only matters for vector field.
    # Thus, the following is more physical, but practically rather
    # difficult to handle, as the segment per step changes by r^2.

    # In x, z plane, the vector field is r = 2cos t / r^2 and t = sin t / r^2.
    def vector_field_dipole2(pos):
        x, y, z = pos
        r2 = x * x + z * z
        th = np.arccos(z / np.sqrt(r2))
        r = 2 * np.cos(th) / r2
        t = np.sin(th) / r2
        lx = r * np.sin(th) + t * np.cos(th)
        lz = r * np.cos(th) - t * np.sin(th)
        return np.array([lx, 0, lz])

    strm = pyana.util.streamline.SimpleTracing((0.5, 0, 0), vector_field_dipole2, 0.001)
    # The vectof_field length is of the order of 1.
    # dt is 0.001, thus, the segment length per step is about 0.001.

    for t in range(30000):
        strm.step_forward()
        if strm.x ** 2 + strm.y ** 2 < 0.1 ** 2:
            break
    ax.plot(strm.xlist, strm.zlist, 'bo')


    ax.set_aspect('equal')


def draw_multidipole(ax):
    # In x, z plane, the vector field is r = 2cos t and t = sin t.
    def vector_field_dipole(pos):
        x, y, z = pos
        r2 = x * x + z * z
        th = np.arccos(z / np.sqrt(r2))
        r = 2 * np.cos(th)
        t = np.sin(th)
        lx = r * np.sin(th) + t * np.cos(th)
        lz = r * np.cos(th) - t * np.sin(th)
        return np.array([lx, 0, lz])

    for x0 in np.arange(0.1, 3, 0.3):
        print(x0)
        strm = irfpy.util.streamline.SimpleTracing((x0, 0, 0), vector_field_dipole, 0.001)
        for t in range(30000):
            strm.step_forward()
            if strm.x ** 2 + strm.y ** 2 < 0.001 ** 2:
                break
        ax.plot(strm.xlist, strm.zlist, 'b')


def main():
    fig = plt.figure()

    ax1 = fig.add_subplot(221)
    draw_vortex(ax1)

    ax2 = fig.add_subplot(222)
    draw_dipole(ax2)

    ax3 = fig.add_subplot(223)
    draw_multidipole(ax3)

    fig.savefig('streamline_sample.png')

if __name__ == "__main__":
    main()