apps120807_mhddata.mhddata_bsimple_trace

''' A simple sample to plot MHD B filed data in 2D

.. image:: ../../../src/scripts/apps120807_mhddata/mhddata_bsimple_trace.png
'''

import bisect

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Circle, Arrow

import scipy.integrate

from pyana.pep import mhddata

def trace(mf, pos):
    ''' Return the traced coordinates.
    '''
    # Try tracing.
    #    dx / dt = Bx (x, y)
    #    dy / dt = By (x, y)
    def b(pos, t):
        if (pos * pos).sum() < 1.05 * 1.05:
            return -pos * np.sign(t)
        elif (pos * pos).sum() > 7 * 7:
            return pos * np.sign(t)
        try:
            bx, by, bz = mf.interpolate3d(pos[0], 0, pos[1])
        except:
            return pos
        return np.array((bx, bz))

    ret = scipy.integrate.odeint(b, pos, np.arange(0, 1, 0.0001))   # (N, 2) shpae
    rret = (ret * ret).sum(1)
    mask = np.logical_or(rret < 1.05 * 1.05, rret > 7 * 7)
    ret = np.ma.masked_array(ret, mask=np.array([mask, mask]).T)

    ret2 = scipy.integrate.odeint(b, pos, np.arange(0, -1, -0.0001))
    rret2 = (ret2 * ret2).sum(1)
    mask2 = np.logical_or(rret2 < 1.05 * 1.05, rret2 > 7 * 7)
    ret2 = np.ma.masked_array(ret2, mask=np.array([mask2, mask2]).T)

    ret2 = ret2[::-1, :]
    ret = np.ma.concatenate([ret2, ret], axis=0)
    return ret

def main(mf=None, ax=None):

    if mf == None:
        mf = mhddata.MagneticField1201()

    if ax == None:
        fig = plt.figure()
        ax = fig.add_subplot(111)

    for lat in np.arange(-180, 181, 5):
        print(lat)
        pos = np.array([1.2 * np.sin(lat * np.pi / 180.), 1.2 * np.cos(lat * np.pi / 180.)])

        ret = trace(mf, pos)
        ax.plot(ret[:, 0], ret[:, 1], 'r')

    ax.set_xlim(-5, 5)
    ax.set_ylim(-5, 5)
    ax.set_aspect('equal')

    # Circle
    gany = Circle((0, 0), 1., facecolor='w', edgecolor='k')
    ax.add_patch(gany)

    return mf, ax

if __name__ == "__main__":
    retval = main()

    plt.savefig('./mhddata_bsimple_trace.png')