"""
This script builds on the Dataset and DataArray tutorials, showing how these
objects smart enough to adjust/align dims for mathematical operations
See official documentation @ http://xarray.pydata.org/en/stable/user-guide/combining.html
"""
# we will reuse the previous tutorial profiles
# Example: A 1d array, such as a time history
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'))
# example 1: multiply a time series by a radial profile
time = arange(1000, 3001, 20)
ne_avg = DataArray(1e19 + 1e18 * sin(time * 2 * pi / 1000), coords=[('time', time)])
print(ne_avg)
print(te)
pe = (ne_avg * te).rename('P_e')
print(pe)
fig, ax = subplots()
pe.plot()
# OMFITx.End()
# Example 2: multiply a time dependent profile and a single profile
temp_2d = 5 * (1 - rho ** 5)[:, None] * (time / 1000)[None, :]
te_ramp = DataArray(temp_2d, coords=[('rho', rho), ('time', time)])
# pe_2d = ne.values * te_ramp # numpy doesn't know how to do this
pe_ramp = ne * te_ramp # xarray knows to line things up along rho
pe_ramp.name = 'P_e'
fig, ax = subplots()
pe_ramp.plot() # different orientation
# Data can be transposed using the numpy-like .T method, but a more general method .transpose lets you explicitly
# name the order of dimensions you would like. This is powerful, as you do not need to know if it is right or wrong...
# Unlike .T, if things are "correct" already there is no harm in calling .transpose with your desired order
fig, ax = subplots()
pe_ramp.transpose('time', 'rho').plot() # "correct" orientation
# Extra Credit: built in methods operate along dims
print(pe_ramp.mean(dim='rho'))
