Skip to content

Compound¤

tatva.compound.Compound ¤

Compound(arr: Array | None = None, **kwargs)

A compound array/state.

A helper class to create a compound state with multiple fields. It simplifies packing and unpacking into and from a flat array. Useful to manage fields while working with a flat array for the solver.

Parameters:

  • arr ¤

    (Array | None, default: None ) –

    The flat data array. If None, initializes to zeros.

Examples:

Create a compound state with fields::

class MyState(Compound):
    u = field(shape=(N, 3))
    phi = field(shape=(N,), default_factory=lambda: jnp.ones(N))

state = MyState()

Use state.arr to access the flat array, and state.u, state.phi to access the individual fields.

You can use iterator unpacking to directly unpack the fields from the state::

u, phi = MyState(arr)

Methods:

  • tree_flatten –
  • tree_unflatten –
  • get_layout –

    MPI only. Get the local layout information for this compound state. Requires

  • at –

    Helper for functional updates to fields. Usage: `new_state =

  • flatten –

    Return the flat array representation of the state. Same as state.arr.

Attributes:

fields class-attribute instance-attribute ¤

fields: tuple[tuple[str, Field], ...] = ()

arr instance-attribute ¤

arr: Array

size class-attribute instance-attribute ¤

size: int = 0

tree_flatten ¤

tree_flatten() -> tuple[tuple[Array], Any]

tree_unflatten classmethod ¤

tree_unflatten(
    aux_data: Any, children: tuple[Array]
) -> Self

get_layout classmethod ¤

get_layout() -> _LocalLayout

MPI only. Get the local layout information for this compound state. Requires that the class has been initialized with MPI context (mesh, partition_info, comm).

at ¤

at(name: str) -> _CompoundAtHelper[Self]

Helper for functional updates to fields. Usage: new_state = state.at('field_name').set(new_value).

Parameters:

  • name ¤
    (str) –

    The name of the field to update.

flatten ¤

flatten() -> Array

Return the flat array representation of the state. Same as state.arr.

tatva.compound.field ¤

Classes:

  • FieldSize –
  • Field –

    A descriptor to define fields in the State class.

  • FieldStackedView –

    A descriptor to define fields that are sub-fields of a stacked field in the State class.

Functions:

Attributes:

T_Compound module-attribute ¤

T_Compound = TypeVar('T_Compound', bound='Compound')

FieldSize ¤

Attributes:

AUTO class-attribute instance-attribute ¤
AUTO = -1

Field ¤

Field(
    shape: tuple[int, ...],
    default_factory: Callable | None = None,
    field_type: _FieldType | FieldType = FieldType.LOCAL,
    *,
    _slice: slice | None = None,
)

A descriptor to define fields in the State class.

Methods:

Attributes:

default_factory instance-attribute ¤
default_factory: Callable | None = default_factory
field_type instance-attribute ¤
field_type: _FieldType | FieldType = field_type
shape instance-attribute ¤
shape: tuple[int, ...] = shape
size instance-attribute ¤
size: int = int(prod(shape))
indices ¤
indices(
    arg: int
    | slice
    | Array
    | tuple[int | slice | Array, ...],
) -> Array

FieldStackedView ¤

FieldStackedView(
    shape: tuple[int, ...],
    default_factory: Callable | None,
    parent_field: _FieldBase,
    parent_slice: tuple[slice, ...],
    field_type: _FieldType | FieldType = FieldType.LOCAL,
)

A descriptor to define fields that are sub-fields of a stacked field in the State class.

Methods:

Attributes:

default_factory instance-attribute ¤
default_factory = default_factory
field_type instance-attribute ¤
field_type = field_type
shape instance-attribute ¤
shape: tuple[int, ...] = shape
size instance-attribute ¤
size: int = int(prod(shape))
indices ¤
indices(
    arg: int
    | slice
    | Array
    | tuple[int | slice | Array, ...],
) -> Array

field ¤

field(
    shape: tuple[int, ...],
    default_factory: Callable | None = None,
    field_type: FieldType | _FieldType = FieldType.LOCAL,
) -> Field

tatva.compound.stack_fields ¤

stack_fields(
    *fields: str, axis: int = -1
) -> Callable[[type[T_Compound]], type[T_Compound]]

Class decorator to stack compatible fields into a shared contiguous layout.