Array Containers

class arraycontext.ArrayContainerT

TypeVar for array container-like objects.

class arraycontext.ArrayContainer

A generic container for the array type supported by the ArrayContext.

The functionality required for the container to operated is supplied via functools.singledispatch(). Implementations of the following functions need to be registered for a type serving as an ArrayContainer:

This allows enumeration of the component arrays in a container and the construction of modified containers from an iterable of those component arrays. is_array_container() will return True for types that have a container serialization function registered.

Packages may register their own types as array containers. They must not register other types (e.g. list) as array containers. The type numpy.ndarray is considered an array container, but only arrays with dtype object may be used as such. (This is so because object arrays cannot be distinguished from non-object arrays via their type.)

The container and its serialization interface has goals and uses approaches similar to JAX’s PyTrees, however its implementation differs a bit.

Note

This class is used in type annotation. Inheriting from it confers no special meaning or behavior.

Serialization/deserialization

arraycontext.is_array_container(ary: Any)bool
Returns

True if the instance ary has a registered implementation of serialize_container().

arraycontext.serialize_container(ary: arraycontext.container.ArrayContainer)Iterable[Tuple[Any, Any]]
arraycontext.serialize_container(ary: numpy.ndarray)Iterable[Tuple[Any, Any]]

Serialize the array container into an iterable over its components.

The order of the components and their identifiers are entirely under the control of the container class.

If ary is mutable, the serialization function is not required to ensure that the serialization result reflects the array state at the time of the call to serialize_container().

Returns

an Iterable of 2-tuples where the first entry is an identifier for the component and the second entry is an array-like component of the ArrayContainer. Components can themselves be ArrayContainers, allowing for arbitrarily nested structures. The identifiers need to be hashable but are otherwise treated as opaque.

arraycontext.deserialize_container(template: Any, iterable: Iterable[Tuple[Any, Any]])Any
arraycontext.deserialize_container(template: numpy.ndarray, iterable: Iterable[Tuple[Any, Any]])numpy.ndarray

Deserialize an iterable into an array container.

Parameters
  • template – an instance of an existing object that can be used to aid in the deserialization. For a similar choice see __array_finalize__.

  • iterable – an iterable that mirrors the output of serialize_container().

Context retrieval

arraycontext.get_container_context(ary: arraycontext.container.ArrayContainer)Optional[arraycontext.context.ArrayContext]

Retrieves the ArrayContext from the container, if any.

This function is not recursive, so it will only search at the root level of the container. For the recursive version, see get_container_context_recursively().

arraycontext.get_container_context_recursively(ary: Any)Optional[arraycontext.context.ArrayContext]

Walks the ArrayContainer hierarchy to find an ArrayContext associated with it.

If different components that have different array contexts are found at any level, an assertion error is raised.

Containers with arithmetic

arraycontext.with_container_arithmetic(*, bcast_number: bool = True, bcast_obj_array: Optional[bool] = None, bcast_numpy_array: bool = False, bcast_container_types: Optional[Tuple[type, ...]] = None, arithmetic: bool = True, matmul: bool = False, bitwise: bool = False, shift: bool = False, _cls_has_array_context_attr: bool = False, eq_comparison: Optional[bool] = None, rel_comparison: Optional[bool] = None)Callable[[type], type]

A class decorator that implements built-in operators for array containers by propagating the operations to the elements of the container.

Parameters
  • bcast_number – If True, numbers broadcast over the container (with the container as the ‘outer’ structure).

  • bcast_obj_array – If True, numpy object arrays broadcast over the container. (with the container as the ‘inner’ structure)

  • bcast_numpy_array – If True, any numpy.ndarray will broadcast over the container. (with the container as the ‘inner’ structure) If this is set to True, bcast_obj_array must also be True.

  • bcast_container_types – A sequence of container types that will broadcast over this container (with this container as the ‘outer’ structure). numpy.ndarray is permitted to be part of this sequence to indicate that, in such broadcasting situations, this container should be the ‘outer’ structure. In this case, bcast_obj_array (and consequently bcast_numpy_array) must be False.

  • arithmetic – Implement the conventional arithmetic operators, including **, divmod(), and //. Also includes + and - as well as abs().

  • bitwise – If True, implement bitwise and, or, not, and inversion.

  • shift – If True, implement bit shifts.

  • eq_comparison – If True, implement == and !=.

  • rel_comparison – If True, implement <, <=, >, >=. In that case, if eq_comparison is unspecified, it is also set to True.

  • _cls_has_array_context_attr – A flag indicating whether the decorated class has an array_context attribute. If so, and if __debug__ is True, an additional check is performed in binary operators to ensure that both containers use the same array context. Consider this argument an unstable interface. It may disappear at any moment.

Each operator class also includes the “reverse” operators if applicable.

Note

To generate the code implementing the operators, this function relies on class methods _deserialize_init_arrays_code and _serialize_init_arrays_code. This interface should be considered undocumented and subject to change, however if you are curious, you may look at its implementation in meshmode.dof_array.DOFArray. For a simple structure type, the implementation might look like this:

@classmethod
def _serialize_init_arrays_code(cls, instance_name):
    return {"u": f"{instance_name}.u", "v": f"{instance_name}.v"}

@classmethod
def _deserialize_init_arrays_code(cls, tmpl_instance_name, args):
    return f"u={args['u']}, v={args['v']}"

dataclass_array_container() automatically generates an appropriate implementation of these methods, so with_container_arithmetic() should nest “outside” :func:dataclass_array_container`.

Containers based on dataclasses

arraycontext.dataclass_array_container(cls: type)type

A class decorator that makes the class to which it is applied a ArrayContainer by registering appropriate implementations of serialize_container() and deserialize_container(). cls must be a dataclass().

Attributes that are not array containers are allowed. In order to decide whether an attribute is an array container, the declared attribute type is checked by the criteria from is_array_container().

Traversing containers

arraycontext.map_array_container(f: Callable[[Any], Any], ary: arraycontext.container.ArrayContainerT)arraycontext.container.ArrayContainerT

Applies f to all components of an ArrayContainer.

Works similarly to obj_array_vectorize(), but on arbitrary containers.

For a recursive version, see rec_map_array_container().

Parameters

ary – a (potentially nested) structure of ArrayContainers, or an instance of a base array type.

arraycontext.multimap_array_container(f: Callable[[...], Any], *args: Any)Any

Applies f to the components of multiple ArrayContainers.

Works similarly to obj_array_vectorize_n_args(), but on arbitrary containers. The containers must all have the same type, which will also be the return type.

For a recursive version, see rec_multimap_array_container().

Parameters

args – all ArrayContainer arguments must be of the same type and with the same structure (same number of components, etc.).

arraycontext.rec_map_array_container(f: Callable[[Any], Any], ary: arraycontext.container.ArrayContainerT)arraycontext.container.ArrayContainerT

Applies f recursively to an ArrayContainer.

For a non-recursive version see map_array_container().

Parameters

ary – a (potentially nested) structure of ArrayContainers, or an instance of a base array type.

arraycontext.rec_multimap_array_container(f: Callable[[...], Any], *args: Any)Any

Applies f recursively to multiple ArrayContainers.

For a non-recursive version see multimap_array_container().

Parameters

args – all ArrayContainer arguments must be of the same type and with the same structure (same number of components, etc.).

Traversing decorators

arraycontext.mapped_over_array_containers(f: Callable[[Any], Any])Callable[[arraycontext.container.ArrayContainerT], arraycontext.container.ArrayContainerT]

Decorator around rec_map_array_container().

arraycontext.multimapped_over_array_containers(f: Callable[[...], Any])Callable[[...], Any]

Decorator around rec_multimap_array_container().

Freezing and thawing

arraycontext.freeze(ary: arraycontext.container.ArrayContainerT, actx: Optional[arraycontext.context.ArrayContext] = None)arraycontext.container.ArrayContainerT

Freezes recursively by going through all components of the ArrayContainer ary.

Parameters

ary – a thaw()ed ArrayContainer.

Array container types may use functools.singledispatch() .register to register additional implementations.

See ArrayContext.thaw().

arraycontext.thaw(ary: arraycontext.container.ArrayContainerT, actx: arraycontext.context.ArrayContext)arraycontext.container.ArrayContainerT

Thaws recursively by going through all components of the ArrayContainer ary.

Parameters

ary – a freeze()ed ArrayContainer.

Array container types may use functools.singledispatch() .register to register additional implementations.

See ArrayContext.thaw().

Serves as the registration point (using singledispatch() .register to register additional implementations for thaw().

Note

This function has the reverse argument order from the original function in meshmode. This was necessary because singledispatch() only dispatches on the first argument.

Numpy conversion

arraycontext.from_numpy(ary: Any, actx: arraycontext.context.ArrayContext)Any

Convert all numpy arrays in the ArrayContainer to the base array type of ArrayContext.

The conversion is done using arraycontext.ArrayContext.from_numpy().

arraycontext.to_numpy(ary: Any, actx: arraycontext.context.ArrayContext)Any

Convert all arrays in the ArrayContainer to numpy using the provided ArrayContext actx.

The conversion is done using arraycontext.ArrayContext.to_numpy().