"""
This script builds on the DataArray tutorial, showing how to collect multiple DataArrays
into a larger Dataset object.
See official documentation @ http://xarray.pydata.org/en/stable/user-guide/data-structures.html#dataset
defaultVars parameters
----------------------
:param verbose: bool. Print / display steps along the way
:param do_save: bool. Save data in netcdf on disk and read it back
"""
defaultVars(verbose=True, do_save=True)
# we will reuse the DataArray tutorial temperature profile
# 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'}
)
# Make another DataArray
ne = DataArray(1e19 * (1 - rho), coords=[('rho', rho)], name='n_e', attrs=dict(units='m^-3'))
# Now create a Dataset from these two DataArrays
ds = Dataset({'T_e': te, 'n_e': ne})
if verbose:
# xarray has a nice succinct printed representation
print(ds)
# Get the coordinates and all the DataArrays in this Dataset
print('DataArrays in our Dataset')
print(list(ds.keys()))
print('Variables in our Dataset')
print(list(ds.variables.keys()))
print('Dimensions in our Dataset')
print(list(ds.dims.keys()))
print('Coordinates in our Dataset')
print(list(ds.coords.keys()))
# OMFITx.End()
# Datasets are essentially in-memory netcdf files
if do_save:
ds.to_netcdf('xarray_demo.nc')
ds_disk = xarray.open_dataset('xarray_demo.nc')
if verbose:
print(ds_disk)
