xbatcher.BatchSchema

class xbatcher.BatchSchema(ds: xarray.core.dataset.Dataset | xarray.core.dataarray.DataArray, input_dims: dict[collections.abc.Hashable, int], input_overlap: Optional[dict[collections.abc.Hashable, int]] = None, batch_dims: Optional[dict[collections.abc.Hashable, int]] = None, concat_input_bins: bool = True, preload_batch: bool = True)

A representation of the indices and stacking/transposing parameters needed to generator batches from Xarray DataArrays and Datasets using xbatcher.BatchGenerator.

Parameters:

dsxarray.Dataset or xarray.DataArray

The data to iterate over. Unlike for the BatchGenerator, the data is not retained as a class attribute for the BatchSchema.

input_dimsdict

A dictionary specifying the size of the inputs in each dimension, e.g. {'lat': 30, 'lon': 30} These are the dimensions the ML library will see. All other dimensions will be stacked into one dimension called sample.

input_overlapdict, optional

A dictionary specifying the overlap along each dimension e.g. {'lat': 3, 'lon': 3}

batch_dimsdict, optional

A dictionary specifying the size of the batch along each dimension e.g. {'time': 10}. These will always be iterated over.

concat_input_dimsbool, optional

If True, the dimension chunks specified in input_dims will be concatenated and stacked into the sample dimension. The batch index will be included as a new level input_batch in the sample coordinate. If False, the dimension chunks specified in input_dims will be iterated over.

preload_batchbool, optional

If True, each batch will be loaded into memory before reshaping / processing, triggering any dask arrays to be computed.

Notes

The BatchSchema is experimental and subject to change without notice.

__init__(ds: xarray.core.dataset.Dataset | xarray.core.dataarray.DataArray, input_dims: dict[collections.abc.Hashable, int], input_overlap: Optional[dict[collections.abc.Hashable, int]] = None, batch_dims: Optional[dict[collections.abc.Hashable, int]] = None, concat_input_bins: bool = True, preload_batch: bool = True)

Methods

__init__(ds, input_dims[, input_overlap, ...])
to_file(out_file_name)	Dumps the JSON representation of the BatchSchema object to a file.
to_json()	Dump the BatchSchema properties to a JSON file.