earth2studio.io.KVBackend#
- class earth2studio.io.KVBackend(device='cpu')[source]#
A key-value (dict) backend.
- Parameters:
device (str = 'cpu') – Device to keep array tensors.
- add_array(coords, array_name, data=None)[source]#
Add an array to the existing KV store.
- Parameters:
coords (CoordSystem) – Ordered dict of coordinate information.
array_name (str) – Name to add to kv store for the new array. Can optionally be a list of array_names.
data (torch.Tensor | list[torch.Tensor], optional) – Optional data to initialize the array with. If None, then the KV store array is initialized with torch.float32 zeros. 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.
- Return type:
None
- write(x, coords, array_name)[source]#
Write data to the current KV store using the passed array_name.
- Parameters:
x (torch.Tensor | list[torch.Tensor]) – Tensor(s) to be written to KV store.
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