Loading data with yt_xarray

This notebook demonstrates how to initialize a yt dataset object from an open xarray dataset.

After describing the sample data that is used, the notebook covers:

• Loading all fields

• Overview of yt datasets

• Loading a subset of fields

• Loading method and memory usage

sample data

We’ll be using some random sample data in this notebook (as well as many of the others), generated from a convenience function, yt_xarray.sample_data.load_random_xr_data(). To use it, we have to supply two dictionaries: one containing fieldnames mapped to the dimension names and a second containing the starting value, end value and number of elements for each dimension:

from yt_xarray.sample_data import load_random_xr_data

fields = {'temperature': ('x', 'y', 'z'), 'pressure': ('x', 'y', 'z')}
dims = {'x': (0,1,15), 'y': (0, 1, 10), 'z': (0, 1, 15)}
ds = load_random_xr_data(fields, dims, length_unit='m')
ds

<xarray.Dataset>
Dimensions:      (x: 15, y: 10, z: 15)
Coordinates:
  * x            (x) float64 0.0 0.07143 0.1429 0.2143 ... 0.8571 0.9286 1.0
  * y            (y) float64 0.0 0.1111 0.2222 0.3333 ... 0.7778 0.8889 1.0
  * z            (z) float64 0.0 0.07143 0.1429 0.2143 ... 0.8571 0.9286 1.0
Data variables:
    temperature  (x, y, z) float64 0.7337 0.2377 0.9107 ... 0.3052 0.8424 0.615
    pressure     (x, y, z) float64 0.8171 0.6735 0.5087 ... 0.5845 0.2743 0.3072
Attributes:
    geospatial_vertical_units:  m

xarray.Dataset

• Dimensions:
  • x: 15
  • y: 10
  • z: 15
• Coordinates: (3)
  • x
    (x)
    float64
    0.0 0.07143 0.1429 ... 0.9286 1.0

    array([0.      , 0.071429, 0.142857, 0.214286, 0.285714, 0.357143, 0.428571,
           0.5     , 0.571429, 0.642857, 0.714286, 0.785714, 0.857143, 0.928571,
           1.      ])

  • y
    (y)
    float64
    0.0 0.1111 0.2222 ... 0.8889 1.0

    array([0.      , 0.111111, 0.222222, 0.333333, 0.444444, 0.555556, 0.666667,
           0.777778, 0.888889, 1.      ])

  • z
    (z)
    float64
    0.0 0.07143 0.1429 ... 0.9286 1.0

    array([0.      , 0.071429, 0.142857, 0.214286, 0.285714, 0.357143, 0.428571,
           0.5     , 0.571429, 0.642857, 0.714286, 0.785714, 0.857143, 0.928571,
           1.      ])

• Data variables: (2)
  • temperature
    (x, y, z)
    float64
    0.7337 0.2377 ... 0.8424 0.615

    array([[[0.73368176, 0.23771448, 0.91073601, ..., 0.06008264,
             0.37625964, 0.60894307],
            [0.36703254, 0.92475816, 0.93867247, ..., 0.62492092,
             0.85110374, 0.01710773],
            [0.00786893, 0.24076778, 0.15011879, ..., 0.84437393,
             0.90841068, 0.80679204],
            ...,
            [0.68457997, 0.4605649 , 0.79889932, ..., 0.30221566,
             0.03109809, 0.33582512],
            [0.05650869, 0.20030877, 0.90223706, ..., 0.90971329,
             0.45483192, 0.65580303],
            [0.84728364, 0.77436776, 0.37467561, ..., 0.73817668,
             0.95293866, 0.02186096]],
    
           [[0.26920323, 0.57245454, 0.63481885, ..., 0.16790371,
             0.1605442 , 0.92581818],
            [0.49346315, 0.44162591, 0.3513115 , ..., 0.82997637,
             0.26915307, 0.69966989],
            [0.25534113, 0.6505062 , 0.71681105, ..., 0.79076677,
             0.80730701, 0.33001465],
    ...
            [0.52114317, 0.80826312, 0.6081363 , ..., 0.9506118 ,
             0.01655698, 0.62626934],
            [0.50851549, 0.34243155, 0.30978383, ..., 0.40368426,
             0.65488025, 0.03426049],
            [0.07107822, 0.07294193, 0.68502822, ..., 0.37187616,
             0.48171863, 0.32071365]],
    
           [[0.00980528, 0.53736763, 0.9977264 , ..., 0.69486939,
             0.16251459, 0.79911676],
            [0.63242795, 0.62261564, 0.01691862, ..., 0.34824451,
             0.55688345, 0.14296885],
            [0.70497576, 0.04419448, 0.91866438, ..., 0.33666275,
             0.4068478 , 0.12225796],
            ...,
            [0.27302668, 0.07709793, 0.82416857, ..., 0.06691517,
             0.03271343, 0.55392403],
            [0.03522629, 0.22326103, 0.2900984 , ..., 0.3656976 ,
             0.25798972, 0.18221715],
            [0.58971867, 0.07627362, 0.61026897, ..., 0.30522095,
             0.84243193, 0.6150061 ]]])

  • pressure
    (x, y, z)
    float64
    0.8171 0.6735 ... 0.2743 0.3072

    array([[[0.81705536, 0.67346801, 0.50869439, ..., 0.27714732,
             0.61226064, 0.85082597],
            [0.9426367 , 0.41442298, 0.91748896, ..., 0.56614286,
             0.72885531, 0.89973352],
            [0.65620274, 0.12359577, 0.42757551, ..., 0.49788247,
             0.55293363, 0.54619041],
            ...,
            [0.59284083, 0.733435  , 0.98857121, ..., 0.2209063 ,
             0.97476519, 0.79449075],
            [0.45258129, 0.6306918 , 0.82110543, ..., 0.92057936,
             0.11464806, 0.98857959],
            [0.18825568, 0.59885988, 0.82740492, ..., 0.54189993,
             0.58514944, 0.39845225]],
    
           [[0.76941162, 0.78807526, 0.25807742, ..., 0.23660976,
             0.16401969, 0.67811219],
            [0.37219013, 0.24312284, 0.29684515, ..., 0.1505512 ,
             0.96448087, 0.99265141],
            [0.58479592, 0.041415  , 0.05089004, ..., 0.15931219,
             0.54076498, 0.71717521],
    ...
            [0.22967799, 0.94408679, 0.47757772, ..., 0.18990161,
             0.70186946, 0.67768263],
            [0.33497562, 0.92862798, 0.26213819, ..., 0.95788614,
             0.56837299, 0.35687302],
            [0.64178113, 0.47886679, 0.41504985, ..., 0.28395306,
             0.07680298, 0.36630436]],
    
           [[0.90369568, 0.05583482, 0.91647535, ..., 0.43616768,
             0.12088932, 0.87967023],
            [0.20886172, 0.97878674, 0.8707838 , ..., 0.383281  ,
             0.57368865, 0.05284215],
            [0.71899016, 0.47950988, 0.74057308, ..., 0.30676444,
             0.51546102, 0.95487417],
            ...,
            [0.68292096, 0.96613135, 0.46183927, ..., 0.2367633 ,
             0.27646835, 0.68366993],
            [0.21992429, 0.60791978, 0.3722718 , ..., 0.70672664,
             0.6636147 , 0.94569748],
            [0.16862236, 0.42773392, 0.83714861, ..., 0.58446992,
             0.2743106 , 0.30723962]]])

• Attributes: (1)

  geospatial_vertical_units :

  m

While we’re using random sample data here, note that yt_xarray provides a simple wrapper of the standard xarray open_dataset function that will check yt’s test_data_dir for data if the file is not found in the local path. It is used in the same way as xarray:

ds = yt_xarray.open_dataset("path/to/your/dataset.nc")

Loading all fields

The primary way of loading data into yt is by creation of a yt dataset object.

To create a yt dataset that loads all the data variables:

import yt_xarray
import yt

yt_ds = ds.yt.load_grid()

yt_xarray : [INFO ] 2023-02-06 12:24:07,825:  Inferred geometry type is cartesian. To override, use ds.yt.set_geometry
yt_xarray : [INFO ] 2023-02-06 12:24:07,826:  Attempting to detect if yt_xarray will require field interpolation:
yt_xarray : [INFO ] 2023-02-06 12:24:07,827:      Cartesian geometry on uniform grid: yt_xarray will not interpolate.
yt : [INFO     ] 2023-02-06 12:24:07,927 Parameters: current_time              = 0.0
yt : [INFO     ] 2023-02-06 12:24:07,928 Parameters: domain_dimensions         = [15 10 15]
yt : [INFO     ] 2023-02-06 12:24:07,929 Parameters: domain_left_edge          = [-0.03571429 -0.05555556 -0.03571429]
yt : [INFO     ] 2023-02-06 12:24:07,930 Parameters: domain_right_edge         = [1.03571429 1.05555556 1.03571429]
yt : [INFO     ] 2023-02-06 12:24:07,931 Parameters: cosmological_simulation   = 0

note that this yt dataset actually maintains references to the open xarray dataset! Data will be loaded into yt only as needed.

A brief overview of yt datasets

Now that we have a yt dataset, let’s do a quick overview of a yt dataset:

Field Tuples

You can check the available fields in a yt dataset with:

yt_ds.field_list

[('stream', 'pressure'), ('stream', 'temperature')]

fields in yt include both a field type and a field name. yt_xarray relies on yt’s “stream” frontend infrastructure, so all of our xarray fields end up with a field type of “stream”. This is important because when referring to fields, you’ll need to supply both the field type and the field name.

To construct a SlicePlot, for example:

slc = yt.SlicePlot(yt_ds, "z", ("stream", "temperature"))
slc.set_log(("stream", "temperature"), False)
slc.show()

yt : [INFO     ] 2023-02-06 12:24:08,113 xlim = -0.035714 1.035714
yt : [INFO     ] 2023-02-06 12:24:08,114 ylim = -0.055556 1.055556
yt : [INFO     ] 2023-02-06 12:24:08,115 xlim = -0.035714 1.035714
yt : [INFO     ] 2023-02-06 12:24:08,115 ylim = -0.055556 1.055556
yt : [INFO     ] 2023-02-06 12:24:08,121 Making a fixed resolution buffer of (('stream', 'temperature')) 800 by 800

Domain extent and units

yt also has some useful attributes to quickly check the domain extents:

print(yt_ds.domain_center)

[0.5 0.5 0.5] code_length

print(yt_ds.domain_width)

[1.07142857 1.11111111 1.07142857] code_length

You’ll notice that the output above are unyt arrays. The “code_length” refers to the representative length of your volume. You can view a unyt array with different unyts with:

print(yt_ds.domain_width.to('m'))

[1.07142857 1.11111111 1.07142857] m

The above sample dataset sets an attribute, geospatial_vertical_units:

ds.geospatial_vertical_units

'm'

which yt_xarray sets as the dataset length_unit:

yt_ds.length_unit

unyt_quantity(1., 'm')

without this attribute, you can set it explicitly so that yt will know the dimensions of your input coordinates:

yt_ds = ds.yt.load_grid(length_unit='km')

yt_xarray : [INFO ] 2023-02-06 12:24:08,945:  Attempting to detect if yt_xarray will require field interpolation:
yt_xarray : [INFO ] 2023-02-06 12:24:08,946:      Cartesian geometry on uniform grid: yt_xarray will not interpolate.
yt : [INFO     ] 2023-02-06 12:24:09,045 Parameters: current_time              = 0.0
yt : [INFO     ] 2023-02-06 12:24:09,046 Parameters: domain_dimensions         = [15 10 15]
yt : [INFO     ] 2023-02-06 12:24:09,047 Parameters: domain_left_edge          = [-0.03571429 -0.05555556 -0.03571429]
yt : [INFO     ] 2023-02-06 12:24:09,048 Parameters: domain_right_edge         = [1.03571429 1.05555556 1.03571429]
yt : [INFO     ] 2023-02-06 12:24:09,049 Parameters: cosmological_simulation   = 0

yt_ds.domain_width.to('m')

unyt_array([1071.42857143, 1111.11111111, 1071.42857143], 'm')

for more on units, unyt and yt, you can read more here.

You also might notice that the domain width in the yt dataset is slightly larger than that in the xarray dataset:

print(yt_ds.domain_width)

[1.07142857 1.11111111 1.07142857] code_length

[ds.coords[dim].max().values.item() - ds.coords[dim].min().values.item() for dim in ds.dims]

[1.0, 1.0, 1.0]

This is related to how yt_xarray builds cells from node values, see here for an explanation.

Loading a subset of fields

When you call ds.yt.load_grid() without arguments, it attempts to grab references to all of the available fields. If you only want to work with a subset of fields, you can supply the fields argument:

yt_ds = ds.yt.load_grid(fields=('temperature',))

yt_xarray : [INFO ] 2023-02-06 12:24:09,089:  Attempting to detect if yt_xarray will require field interpolation:
yt_xarray : [INFO ] 2023-02-06 12:24:09,091:      Cartesian geometry on uniform grid: yt_xarray will not interpolate.
yt : [INFO     ] 2023-02-06 12:24:09,164 Parameters: current_time              = 0.0
yt : [INFO     ] 2023-02-06 12:24:09,165 Parameters: domain_dimensions         = [15 10 15]
yt : [INFO     ] 2023-02-06 12:24:09,166 Parameters: domain_left_edge          = [-0.03571429 -0.05555556 -0.03571429]
yt : [INFO     ] 2023-02-06 12:24:09,168 Parameters: domain_right_edge         = [1.03571429 1.05555556 1.03571429]
yt : [INFO     ] 2023-02-06 12:24:09,169 Parameters: cosmological_simulation   = 0

One of the limitations to yt_xarrayat present is that the fields loaded into yt must have the same dimensions.

If for example, we had a dataset with a mix of 2d and 3d and 4d (space + time) variables:

fields = {'temperature': ('x', 'y', 'z'),
          'pressure': ('x', 'y', 'z'),
          'precip': ('x', 'y', 'time'),
          'cumulative_precip': ('x', 'y')}
dims = {'x': (0,1,14), 'y': (0, 1, 10), 'z': (0, 1, 15), 'time': (0, 1, 11)}
ds = load_random_xr_data(fields, dims, length_unit='m')

and simply try to load the full dataset with

ds.yt.load_grid()

then we will get an error due to the mismatching dimensions of the variables. In this case, we can load both the temperature and pressure fields with:

yt_ds = ds.yt.load_grid(fields=('temperature', 'pressure'))

yt_xarray : [INFO ] 2023-02-06 12:24:09,202:  Inferred geometry type is cartesian. To override, use ds.yt.set_geometry
yt_xarray : [INFO ] 2023-02-06 12:24:09,203:  Attempting to detect if yt_xarray will require field interpolation:
yt_xarray : [INFO ] 2023-02-06 12:24:09,204:      Cartesian geometry on uniform grid: yt_xarray will not interpolate.
yt : [INFO     ] 2023-02-06 12:24:09,266 Parameters: current_time              = 0.0
yt : [INFO     ] 2023-02-06 12:24:09,267 Parameters: domain_dimensions         = [14 10 15]
yt : [INFO     ] 2023-02-06 12:24:09,268 Parameters: domain_left_edge          = [-0.03846154 -0.05555556 -0.03571429]
yt : [INFO     ] 2023-02-06 12:24:09,269 Parameters: domain_right_edge         = [1.03846154 1.05555556 1.03571429]
yt : [INFO     ] 2023-02-06 12:24:09,270 Parameters: cosmological_simulation   = 0

Furthermore, when loading variables with a time dimension, you should select the time index to load:

yt_ds = ds.yt.load_grid(fields=('precip',), sel_dict={'time':1})

yt_xarray : [INFO ] 2023-02-06 12:24:09,285:  Attempting to detect if yt_xarray will require field interpolation:
yt_xarray : [INFO ] 2023-02-06 12:24:09,287:      Cartesian geometry on uniform grid: yt_xarray will not interpolate.
yt : [INFO     ] 2023-02-06 12:24:09,344 Parameters: current_time              = 0.1
yt : [INFO     ] 2023-02-06 12:24:09,345 Parameters: domain_dimensions         = [14 10  1]
yt : [INFO     ] 2023-02-06 12:24:09,347 Parameters: domain_left_edge          = [-0.03846154 -0.05555556 -0.5       ]
yt : [INFO     ] 2023-02-06 12:24:09,348 Parameters: domain_right_edge         = [1.03846154 1.05555556 0.5       ]
yt : [INFO     ] 2023-02-06 12:24:09,349 Parameters: cosmological_simulation   = 0

value, location = yt_ds.find_min(('stream', 'precip'))

yt : [INFO     ] 2023-02-06 12:24:09,494 min value is 1.64191e-02 at 0.9230769230769234 0.3333333333333334 0.0000000000000000

Loading method and memory usage

When you use ds.yt.load_grid, by default yt_xarray will use references to the open xarray dataset handle internally within yt to avoid copying data. If your dataset fits into memory without a problem, then you can supply an additional use_callable=False argument to direct yt_xarray to instead work in memory on copies of data.

yt_ds = ds.yt.load_grid(fields=('temperature',), use_callable=False)

yt_xarray : [INFO ] 2023-02-06 12:24:09,500:  Attempting to detect if yt_xarray will require field interpolation:
yt_xarray : [INFO ] 2023-02-06 12:24:09,502:      Cartesian geometry on uniform grid: yt_xarray will not interpolate.
yt : [INFO     ] 2023-02-06 12:24:09,552 Parameters: current_time              = 0.0
yt : [INFO     ] 2023-02-06 12:24:09,553 Parameters: domain_dimensions         = [14 10 15]
yt : [INFO     ] 2023-02-06 12:24:09,554 Parameters: domain_left_edge          = [-0.03846154 -0.05555556 -0.03571429]
yt : [INFO     ] 2023-02-06 12:24:09,556 Parameters: domain_right_edge         = [1.03846154 1.05555556 1.03571429]
yt : [INFO     ] 2023-02-06 12:24:09,557 Parameters: cosmological_simulation   = 0

and data will be copied from the xarray dataset during creation of the yt dataset.