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

  • get_sparsity –

    TODO: Currently, only correct for full nodal fields. Incomplete fields are

  • 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).

get_sparsity classmethod ¤

get_sparsity(
    block_wise: bool = False,
) -> csr_matrix | list[list[csr_matrix]]

TODO: Currently, only correct for full nodal fields. Incomplete fields are missing in-element couplings. Other fields are only connected to themselves (diagonal entries)! Do not use blindly for mixed meshes.

Create a sparsity pattern automatically from this Compound class and its attached mesh.

Nodal fields are fully coupled within elements. All other fields (Local, Shared) are only connected to themselves (diagonal entries). For non-trivial couplings, you must create the sparsity pattern manually.

Parameters:

  • block_wise ¤
    (bool, default: False ) –

    If True, return the pattern as a list of lists of sparse matrices corresponding to the compound fields/blocks. Stacked fields are one block.

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.