"""
xarray (xarray.pydata.org) is a convenient container for data
that provides more functionality than a simple array or dictionary.
The basic object is an xarray DataArray, which is an array with dimensions,
attributes, plotting methods, etc.
See official documentation at http://xarray.pydata.org/en/stable/user-guide/data-structures.html#dataarray
defaultVars parameters
----------------------
:param verbose: bool. Print / display steps along the way
"""
defaultVars(verbose=True)
# Example: A 1d array, such as a time history
rho = linspace(0, 1, 21)
temp = 5 * (1 - rho ** 5)
# Create a DataArray. There are a few equivalent ways
te = DataArray(temp, coords=[('rho', rho)])
# or
te = DataArray(temp, dims='rho', coords=[rho])
# check it out
if verbose:
print(te)
print(f'values = \n{te.values}')
print(f'dims = \n{te.dims}')
print(f'coords = \n{te.coords}')
print(f'attrs = \n{te.attrs}')
# OMFITx.End()
# Add some attributes
te.name = 'T_e'
te.attrs['units'] = 'keV'
# Now we can plot it and display some information about it.
if verbose:
fig, axs = plt.subplots(2, figsize=(8, 9))
te.plot(ax=axs[0])
# OMFITx.End()
# the same physical quantity can be represented in different coordinates
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'}
)
te_psi = te.swap_dims({'rho': 'psi_n'})
if verbose:
print(te)
print(te_psi)
te_psi.plot(ax=axs[1])
fig.tight_layout()
