apps120807_mhddata.mhddata_bsimple_trace2

''' A simple sample to plot MHD B filed data in 2D.  Using ``pyana.util.streamline``

The result is similar to :mod:`mhddata_bsimple_trace`, but probably more stable.
At least no strange failure cannot be found.

.. image:: ../../../src/scripts/apps120807_mhddata/mhddata_bsimple_trace2.png
'''
import logging
logger = logging.getLogger('mhddata_bsimple_trace2.py')
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
from pyana.util import streamline


def trace(mf, pos):
    ''' Return the traced coordinates.
    '''

    def bfield(pos):
        ''' Return the bfield direction in x-z plane
        '''
        bx, by, bz = mf.interpolate3d(pos[0], pos[1], pos[2])
        babs = np.sqrt(bx * bx + by * by + bz * bz)
        return np.array((bx / babs, by / babs, bz / babs))
    

    tracer = streamline.SimpleTracing(pos, bfield, 0.02)
    for t in range(1000):
        try:
            tracer.step_forward()
        except IndexError:
            logger.warn('Index error. Stop calculation')
            break

        if tracer.x ** 2 + tracer.y ** 2 + tracer.z ** 2 < 1:
            break

    tracer2 = streamline.SimpleTracing(pos, bfield, -0.02)
    for t in range(1000):
        try:
            tracer2.step_forward()
        except IndexError:
            logger.warn('Index error. Stop calculation')
            break

        if tracer2.x ** 2 + tracer2.y ** 2 + tracer2.z ** 2 < 1:
            break

    xlist = np.concatenate([tracer2.xlist[::-1], tracer.xlist[1:]])
    zlist = np.concatenate([tracer2.zlist[::-1], tracer.zlist[1:]])

    return xlist, zlist
    

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)
        mf.set_interpolate_strategy('none')
        pos = np.array([1.2 * np.sin(lat * np.pi / 180.), 0, 1.2 * np.cos(lat * np.pi / 180.)])

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

        mf.set_interpolate_strategy('project')
        pos = np.array([1.2 * np.sin(lat * np.pi / 180.), 0, 1.2 * np.cos(lat * np.pi / 180.)])

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

    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_trace2.png')