state

StateMemoryMapping = dict[str, dict | ArrayLike | None] module-attribute

OptionalQuantityType = Quantity | None module-attribute

StateElementType = dict[str, Quantity | OptionalQuantityType | Local | dict[str, Any]] module-attribute

State dataclass

Base class for state objects in models.

A State groups a collection of (possibly nested) Quantities in a dataclass.

This baseclass implements common initialization functions and serialization.

Typical usage example:

    class MyState(State):
        pass

    my_state = MyState.zeros(quantity_factory)

    # ...

    my_state.to_netcdf()

__init__()

empty(quantity_factory, *, data_dimensions=None) classmethod

Allocate all quantities. Do not expect 0 on values, values are random.

Parameters:

Name	Type	Description	Default
quantity_factory	QuantityFactory	factory, expected to be defined on the Grid dimensions e.g. without data dimensions.	required
data_dimensions	dict[str, int] | None	extra data dimensions required for any field with data dimensions. Dict of name/size pair.	None

zeros(quantity_factory, *, data_dimensions=None) classmethod

Allocate all quantities and fill their value to zeros

Parameters:

Name	Type	Description	Default
quantity_factory	QuantityFactory	factory, expected to be defined on the Grid dimensions e.g. without data dimensions.	required
data_dimensions	dict[str, int] | None	extra data dimensions required for any field with data dimensions. Dict of name/size pair.	None

ones(quantity_factory, *, data_dimensions=None) classmethod

Allocate all quantities and fill their value to ones

Parameters:

Name	Type	Description	Default
quantity_factory	QuantityFactory	factory, expected to be defined on the Grid dimensions e.g. without data dimensions.	required
data_dimensions	dict[str, int] | None	extra data dimensions required for any field with data dimensions. Dict of name/size pair.	None

full(quantity_factory, value, *, data_dimensions=None) classmethod

Allocate all quantities and fill them with the input value

Parameters:

Name	Type	Description	Default
quantity_factory	QuantityFactory	factory, expected to be defined on the Grid dimensions e.g. without data dimensions.	required
value	Number	number to initialize the buffers with.	required
data_dimensions	dict[str, int] | None	extra data dimensions required for any field with data dimensions. Dict of name/size pair.	None

copy_memory(quantity_factory, memory_map, *, data_dimensions=None) classmethod

Allocate all quantities and fill their value based on the given memory map. See update_from_memory.

Parameters:

Name	Type	Description	Default
quantity_factory	QuantityFactory	factory, expected to be defined on the Grid dimensions e.g. without data dimensions.	required
memory_map	StateMemoryMapping	Dict of name/buffer. See update_from_memory.	required
data_dimensions	dict[str, int] | None	extra data dimensions required for any field with data dimensions. Dict of name/size pair.	None

move_memory(quantity_factory, memory_map, *, data_dimensions=None, check_shape_and_strides=True) classmethod

Allocate all quantities and move memory based on on the given memory map. See update_move_memory.

Parameters:

Name	Type	Description	Default
quantity_factory	QuantityFactory	factory, expected to be defined on the Grid dimensions e.g. without data dimensions.	required
memory_map	StateMemoryMapping	Dict of name/buffer. See update_from_memory.	required
data_dimensions	dict[str, int] | None	extra data dimensions required for any field with data dimensions. Dict of name/size pair.	None
check_shape_and_strides	bool	Check for every given buffer that the shape & strides match the previously allocated memory.	True

fill(value)

update_copy_memory(memory_map)

Copy data into the Quantities carried by the state.

The memory map must follow the dataclass naming convention, e.g.

@dataclass
class MyState:
    @dataclass
    class InnerA
        a: Quantity

    inner_a: InnerA
    b: Quantity

will update with a dictionary looking like

{
    "inner_a":
    {
        "a": Quantity(...)
    }
    "b": Quantity(...)

}

The memory map can be sparse.

update_move_memory(memory_map, *, check_shape_and_strides=True)

Move memory into the Quantities carried by the state. Memory is moved rather than copied (e.g. buffers are swapped)

The memory map must follow the dataclass naming convention, e.g.

@dataclass
class MyState:
    @dataclass
    class InnerA
        a: Quantity

    inner_a: InnerA
    b: Quantity

will update with a dictionary looking like

{
    "inner_a":
    {
        "a": Quantity(...)
    }
    "b": Quantity(...)

}

The memory map can be sparse.

Parameters:

Name	Type	Description	Default
memory_map	StateMemoryMapping	Dictionary of keys to buffers. Buffers must be np.ArrayLike	required
check_shape_and_strides	bool	check that the given buffers have the same shape and strides as the original quantity	True

to_xarray()

Format the State into a xr.DataTree.

to_netcdf(directory_path=None, postfix='')

Save state to NetCDF. Can be reloaded with update_from_netcdf.

If applicable, will save separate NetCDF files for each running rank.

The file names are deduced from the class name, and post fix with rank number in the case of a multi-process use.

Parameters:

Name	Type	Description	Default
directory_path	Path | None	directory to save the netcdf in	None

update_from_netcdf(directory_path, postfix='')

This is a mirror of the to_netcdf method NOT a generic NetCDF loader. It expects the NetCDF to be named with the auto-naming scheme of to_netcdf.

Parameters:

Name	Type	Description	Default
directory_path	Path	directory carrying the netcdf saved with to_netcdf	required

LocalState dataclass

Bases: State

__init__()

make_locals(quantity_factory, *, data_dimensions=None) classmethod

Allocate all elements as Locals. Do not expect 0 on values, values are random.

Parameters:

Name	Type	Description	Default
quantity_factory	QuantityFactory	factory, expected to be defined on the Grid dimensions e.g. without data dimensions.	required
data_dimensions	dict[str, int] | None	extra data dimensions required for any field with data dimensions. Dict of name/size pair.	None

make_as_state(quantity_factory, *, data_dimensions=None) classmethod

Allow LocalState to be allocate as if it was a regular State, e.g. with Quantities.

This behavior is useful for testing and interfacing, but should not be used in the regular numerical workflow, hence the warning.

zeros(quantity_factory, *, data_dimensions=None) classmethod

empty(quantity_factory, *, data_dimensions=None) classmethod

ones(quantity_factory, *, data_dimensions=None) classmethod

full(quantity_factory, value, *, data_dimensions=None) classmethod