apps120807_mhddata.mhddata_equator_plane

''' A simple sample to plot density contour and velocity vector in the eq. plane.

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

import bisect

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

from pyana.pep import mhddata

def main(pp=None):

    if pp == None:
        pp = mhddata.PlasmaParameter1205()

    x = pp.xlist()
    y = pp.ylist()
    z = pp.zlist()

    n = pp.nlist()
    lnn = np.log10(n)
    lnn[np.isinf(lnn)] = np.nan

    vx = pp.vxlist()
    vy = pp.vylist()
    vz = pp.vzlist()

    t = pp.tlist()

    # 1-D coordinate value
    xc = x[:, 0, 0]
    yc = y[0, :, 0]
    zc = z[0, 0, :]

    # z=0 index
    z0idx = bisect.bisect(zc, 0) - 1

    # Slice
    x_z0 = x[:, :, z0idx]   # 200, 200
    y_z0 = y[:, :, z0idx]   # 200, 200
    z_z0 = z[:, :, z0idx]   # 200, 200
    n_z0 = n[:, :, z0idx]
    lnn_z0 = lnn[:, :, z0idx]
    vx_z0 = vx[:, :, z0idx]
    vy_z0 = vy[:, :, z0idx]
    vz_z0 = vz[:, :, z0idx]
    t_z0 = t[:, :, z0idx]

    # Plot density/velocity vector
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.contourf(xc, yc, lnn_z0.T, 8)

#    Something wrong??
#    ax.quiver(x_z0[::5, ::5], y_z0[::5, ::5], vx_z0[::5, ::5], vy_z0[::5, ::5], units='xy')
    tmax= t_z0.max()

    scale = 1e-3  # 100 km/s corresponds to 0.1 Rg
    for xx, yy, uu, vv, tt in zip(x_z0[::10, ::10].flatten(), y_z0[::10, ::10].flatten(), vx_z0[::10, ::10].flatten(), vy_z0[::10, ::10].flatten(), t_z0[::10, ::10].flatten()):
#        ax.arrow(xx, yy, uu * scale, vv * scale)
        arr = Arrow(xx, yy, uu * scale, vv * scale, width=tt / tmax)
        ax.add_patch(arr)

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

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

    plt.savefig('mhddata_equator_plane.png')

    return pp

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