SCRIPTS XARRAY DatasetΒΆ

"""
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)