SCRIPTS XARRAY Plotting_Extras

"""
This script builds on the Dataset and DataArray tutorials, showing how these
objects now have smart plotting methods built in.

We have already been using the plot methods, and seen it is smart enough to change the plot
as the dimensionality of the data changes (1D -> line plot, 2D -> color mesh plot, ND -> histogram reductions)

The xarray module adds some nice "bonus" key word arguments on top of the usual matplotlib options too.

See official documentation @ http://xarray.pydata.org/en/stable/user-guide/plotting.html

"""

# reuse the previous tutorial profiles
rho = linspace(0, 1, 21)
temp = 5 * (1 - rho ** 5)
psi_n = DataArray(rho ** 2, coords=[('rho', rho)], name='psi_n')
te = DataArray(
    temp, dims=['rho'], coords={'rho': rho, 'psi_n': psi_n}, name='T_e', attrs={'units': 'keV', 'long_name': 'Electron Temperature'}
)
ne = DataArray(1e19 * (1 - rho), coords=[('rho', rho)], name='n_e', attrs=dict(units='m^-3'))
time = arange(1000, 3001, 20)
ne_avg = DataArray(1e19 + 1e18 * sin(time * 2 * pi / 1000), coords=[('time', time)])
pe = (ne_avg * te).rename('P_e')


# the robust key word sets limits based on a percentile instead of the extrema,
# which is helpful when a few outliers are ruining the scale of your plot
fig, axs = subplots(2, figsize=(8, 9))
pe[0, 0] = 1e99  # outlier
im = pe.T.plot(ax=axs[0])  # color scale swamps all the important stuff
im = pe.T.plot(ax=axs[1], robust=True)  # no longer swamped
fig.tight_layout()
print(im)  # quadmesh -> can use pcolormesh key word args in call of .plot method