io.KVBackend

Import path: earth2studio.io.KVBackend

View source on GitHub

Documentation

A key-value (dict) backend.

Parameters:

• device (str, default: 'cpu' ) –

  Device to keep array tensors.

add_array

add_array(
    coords: CoordSystem,
    array_name: str | list[str],
    data: Tensor | list[Tensor] = None,
) -> None

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 (Tensor | list[Tensor], default: None ) –

  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.

write

write(
    x: Tensor | list[Tensor],
    coords: CoordSystem,
    array_name: str | list[str],
) -> None

Write data to the current KV store using the passed array_name.

Parameters:

• x (Tensor | list[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.