Skip to content

state

LocalState dataclass

Bases: State

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

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.

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

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

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

to_xarray()

Format the State into a xr.DataTree.

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_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

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

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.

Example

Typical usage looks something like this:

    class MyState(State):
        pass

    my_state = MyState.zeros(quantity_factory)

    # ...

    my_state.to_netcdf()

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

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

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

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

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

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

to_xarray()

Format the State into a xr.DataTree.

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_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

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

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