Skip to content

Compound¤

tatva.compound.Compound ¤

Compound(arr: Array | None = None)

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:

Attributes:

  • fields –

    Built-in immutable sequence.

  • size –

    int([x]) -> integer

fields class-attribute ¤

fields = ()

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

size class-attribute ¤

size = 0

int([x]) -> integer int(x, base=10) -> integer

Convert a number or string to an integer, or return 0 if no arguments are given. If x is a number, return x.int(). For floating-point numbers, this truncates towards zero.

If x is not a number or if base is given, then x must be a string, bytes, or bytearray instance representing an integer literal in the given base. The literal can be preceded by '+' or '-' and be surrounded by whitespace. The base defaults to 10. Valid bases are 0 and 2-36. Base 0 means to interpret the base from the string as an integer literal.

int('0b100', base=0) 4

tree_flatten ¤

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

tree_unflatten classmethod ¤

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