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