Implement a State object #181
Description
Currently, the state is represented by a plain python dictionary containing key-value pairs with temporal information ("date", "step", "previous_step"), coordinates ("latitudes", "longitudes") and "fields". It is the most central object in the codebase and many separate components interact with it, having to implement the logic to do so.
This leads to significant code duplication and it is error prone because we are interacting with the same object inconsistently.
I suggest introducing a new State object, which will implement interfaces for instantiating it:
@classmethod
def from_dict(cls, state_dict: StateDict) -> "State":
"""Create a state object from a dictionary."""
@classmethod
def from_fieldlist(
cls,
fieldlist: ekd.FieldList | ekd.Source,
namerules=None,
time_strict: bool = True,
check_coords: bool = True,
) -> "State":
"""Create a state object from a fieldlist."""
@classmethod
def from_tensor(cls, tensor: np.ndarray, fieldnames: list[str], ...) -> "State":
"""Create a state object from a tensor."""and their corresponding inverses:
def to_dict(self) -> StateDict:
"""Return the state as a dictionary."""
def to_fieldlist(self, overrides={}) -> ekd.FieldList:
"""Return the state as a fieldlist."""
def to_tensor(self) -> np.ndarray:
"""Return the state as a tensor."""as well as some useful methods and properties to operate with it, such as:
def concatenate(self, other: "State") -> "State":
"""Concatenate the state with another state."""
def combine(self, other: "State") -> "State":
"""Combine the state's fields with another state."""
@cached_property
def shape(self):
"""Return the shape of the state."""
def __getitem__(self, key: slice | np.ndarray) -> "State":
"""Return a new state with a subset of the data."""
def __setitem__(self, key: slice | np.ndarray, value: "State") -> "State":
"""Update a subset of the state by assigning values from another state."""I've implemented a prototype of such a class but I wanted to start a discussion and get your feedback on this. I have also a prototype of a class that is basically a thin wrapper around an xarray object (and it can keep all metadata of each field correctly).
Let me know what you think 😉