earth2studio.io.XarrayBackend

class earth2studio.io.XarrayBackend(coords={}, **xr_kwargs)

An xarray backed IO object.

Parameters:

• coords (CoordSystem) – Coordinates to initialize the xarray Dataset with. Must be a complete set of coordinates, i.e., the Dataset object should be viewed as (mostly) immutable with the given set of coordinates.

• xr_kwargs (dict) – Optional keyword arguments to pass to the xarray.Dataset constructor.

add_array(coords, array_name, data=None, **xr_kwargs)

Add an array to the existing xarray Dataset.

Parameters:

• coords (CoordSystem) – Ordered dict of coordinate information.

• array_name (str) – Name to add to xarray Dataset for the new array.

• data (torch.Tensor | list[torch.Tensor], optional) – Optional data to initialize the array with. If None, then the array is NaN initialized (xarray default). Can also pass a list of tensors, which must match in length to the list of array_names passed. If a list of tensors is passed, it is assumed that each tensor share coords.

• xr_kwargs (Any) – Optional keyword arguments passed to xr.DataArray constructor.

Return type:

None

write(x, coords, array_name)

Write data to the current xarray Dataset using the passed array_name.

Parameters:

• x (torch.Tensor | list[torch.Tensor]) – Tensor(s) to be written to xarray dataset.

• coords (OrderedDict) – Coordinates of the passed data.

• array_name (str | list[str]) – Name(s) of the array(s) that will be written to.

Return type:

None