"""
This script builds on the Dataset and DataArray tutorials, showing how these
objects now have interpolation methods built in (used to be an OMFIT patch)
See official documentation @ http://xarray.pydata.org/en/stable/user-guide/interpolation.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')
# now we introduce a new, higher resolution radial grid
rho_fine = linspace(0, 1, 101)
te_fine = te.interp(rho=rho_fine, method='cubic', kwargs=dict(fill_value='extrapolate'))
print(te_fine[:4]) # keeps name and attrs
print(ne * te)
# OMFITx.End()
fig, ax = subplots()
te.plot(marker='o', mec='none')
te_fine.plot(marker='x')
# OMFITx.End()
# operations on discrepent grids will subselect common grid points
pe_sparse = ne * te_fine
print(pe_sparse)
# have to interpolate up/down as you prefer
pe_fine = ne.interp_like(te_fine) * te_fine
print(pe_fine)
# again, the magic happens when dims don't match up and you don't know/control shapes
pe_oscil_fine = pe.interp_like(te_fine)
fig, ax = subplots()
pe_oscil_fine.plot()
