Array Containers

class arraycontext.ArrayContainer(*args, **kwargs)[source]

A protocol for generic containers of 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.

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 and as a marker of array container attributes for dataclass_array_container(). As a protocol, it is not intended as a superclass.

class arraycontext.ArithArrayContainer(*args, **kwargs)[source]

A sub-protocol of ArrayContainer that supports basic arithmetic.

class arraycontext.ArrayContainerT

A type variable with a lower bound of ArrayContainer.

exception arraycontext.NotAnArrayContainerError[source]

TypeError subclass raised when an array container is expected.

Serialization/deserialization

arraycontext.SerializationKey

alias of Hashable

arraycontext.SerializedContainer

alias of Sequence[tuple[Hashable, ArrayOrContainer]]

arraycontext.is_array_container_type(cls: type) → bool[source]
Returns:

True if the type cls has a registered implementation of serialize_container(), or if it is an ArrayContainer.

Warning

Not all instances of a type that this function labels an array container must automatically be array containers. For example, while this function will say that numpy.ndarray is an array container type, only object arrays actually are array containers.

arraycontext.serialize_container(ary: ArrayContainer) → SerializedContainer[source]
arraycontext.serialize_container(ary: numpy.ndarray) → SerializedContainer

Serialize the array container into a sequence over its components.

The order of the components and their identifiers are entirely under the control of the container class. However, the order is required to be deterministic, i.e. two calls to serialize_container() on array containers of the same types with the same number of sub-arrays must result in a sequence with the keys in the same order.

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:

a Sequence 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: ArrayContainerT, serialized: SerializedContainer) → ArrayContainerT[source]
arraycontext.deserialize_container(template: numpy.ndarray, serialized: SerializedContainer) → numpy.ndarray

Deserialize a sequence into an array container following a template.

Parameters:

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

  • serialized – a sequence that mirrors the output of serialize_container().

Context retrieval

arraycontext.get_container_context_opt(ary: ArrayContainer) → ArrayContext | None[source]

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: ArrayContainer) → ArrayContext | None[source]

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.

Raises an error if no array container is found.

arraycontext.get_container_context_recursively_opt(ary: ArrayContainer) → ArrayContext | None[source]

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.

Returns None if no array context was found.

MultiVector support

arraycontext.register_multivector_as_array_container() → None[source]

Registers MultiVector as an ArrayContainer. This function may be called multiple times. The second and subsequent calls have no effect.

Canonical locations for type annotations

class arraycontext.container.ArrayContainerT
Canonical:

arraycontext.ArrayContainerT

class arraycontext.container.ArrayOrContainerT
Canonical:

arraycontext.ArrayOrContainerT

class arraycontext.container.SerializationKey
Canonical:

arraycontext.SerializationKey

class arraycontext.container.SerializedContainer
Canonical:

arraycontext.SerializedContainer

Containers with arithmetic

arraycontext.with_container_arithmetic(*, number_bcasts_across: bool | None = None, bcasts_across_obj_array: bool | None = None, container_types_bcast_across: tuple[type, ...] | None = None, arithmetic: bool = True, matmul: bool = False, bitwise: bool = False, shift: bool = False, _cls_has_array_context_attr: bool | None = None, eq_comparison: bool | None = None, rel_comparison: bool | None = None, bcast_number: bool | None = None, bcast_obj_array: bool | None = None, bcast_numpy_array: bool = False, _bcast_actx_array_type: bool | None = None, bcast_container_types: tuple[type, ...] | None = None) → Callable[[type], type][source]

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

Parameters:

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

  • bcasts_across_obj_array – If True, this container will be broadcast across numpy object arrays (with the object array as the ‘outer’ structure). Add numpy.ndarray to container_types_bcast_across to achieve the ‘reverse’ broadcasting.

  • container_types_bcast_across – A sequence of container types that will broadcast across this container, with this container as the ‘outer’ structure. numpy.ndarray is permitted to be part of this sequence to indicate that object arrays (and only object arrays) will be broadcast. In this case, bcasts_across_obj_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. If None (the default), this value is set based on whether the class has an array_context attribute. 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[source]

A class decorator that makes the class to which it is applied an 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_type(). This includes some support for type annotations:

  • a typing.Union of array containers is considered an array container.

  • other type annotations, e.g. typing.Optional, are not considered array containers, even if they wrap one.

Note

When type annotations are strings (e.g. because of from __future__ import annotations), this function relies on inspect.get_annotations() (with eval_str=True) to obtain type annotations. This means that cls must live in a module that is importable.

Traversing containers

arraycontext.map_array_container(f: Callable[[Any], Any], ary: Array | ArrayContainer) → Array | ArrayContainer[source]

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[source]

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: Array | ArrayContainer, leaf_class: type | None = None) → Array | ArrayContainer[source]

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, leaf_class: type | None = None) → Any[source]

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

arraycontext.map_reduce_array_container(reduce_func: Callable[[Iterable[Any]], Any], map_func: Callable[[Any], Any], ary: ArrayOrContainerT) → Array[source]

Perform a map-reduce over array containers.

Parameters:
arraycontext.multimap_reduce_array_container(reduce_func: Callable[[Iterable[Any]], Any], map_func: Callable[[...], Any], *args: Any) → Array | ArrayContainer[source]

Perform a map-reduce over multiple array containers.

Parameters:
arraycontext.rec_map_reduce_array_container(reduce_func: Callable[[Iterable[Any]], Any], map_func: Callable[[Any], Any], ary: Array | ArrayContainer, leaf_class: type | None = None) → Array | ArrayContainer[source]

Perform a map-reduce over array containers recursively.

Parameters:

  • reduce_func – callable used to reduce over the components of ary (and those of its sub-containers) if ary is a ArrayContainer. Must be associative.

  • map_func – callable used to map a single array of type arraycontext.ArrayContext.array_types. Returns an array of the same type or a scalar.

Note

The traversal order is unspecified. reduce_func must be associative in order to guarantee a sensible result. This is because reduce_func may be called on subsets of the component arrays, and then again (potentially multiple times) on the results. As an example, consider a container made up of two sub-containers, subcontainer0 and subcontainer1, that each contain two component arrays, array0 and array1. The same result must be computed whether traversing recursively:

reduce_func([
    reduce_func([
        map_func(subcontainer0.array0),
        map_func(subcontainer0.array1)]),
    reduce_func([
        map_func(subcontainer1.array0),
        map_func(subcontainer1.array1)])])

reducing all of the arrays at once:

reduce_func([
    map_func(subcontainer0.array0),
    map_func(subcontainer0.array1),
    map_func(subcontainer1.array0),
    map_func(subcontainer1.array1)])

or any other such traversal.

arraycontext.rec_multimap_reduce_array_container(reduce_func: Callable[[Iterable[Any]], Any], map_func: Callable[[...], Any], *args: Any, leaf_class: type | None = None) → Array | ArrayContainer[source]

Perform a map-reduce over multiple array containers recursively.

Parameters:

  • reduce_func – callable used to reduce over the components of any ArrayContainers in *args (and those of their sub-containers). Must be associative.

  • map_func – callable used to map a single array of type arraycontext.ArrayContext.array_types. Returns an array of the same type or a scalar.

Note

The traversal order is unspecified. reduce_func must be associative in order to guarantee a sensible result. See rec_map_reduce_array_container() for additional details.

arraycontext.stringify_array_container_tree(ary: Array | ArrayContainer) → str[source]
Returns:

a string for an ASCII tree representation of the array container, similar to asciitree.

Traversing decorators

arraycontext.mapped_over_array_containers(f: Callable[[Array | ArrayContainer], Array | ArrayContainer] | None = None, leaf_class: type | None = None) → Callable[[Array | ArrayContainer], Array | ArrayContainer] | Callable[[Callable[[Any], Any]], Callable[[Array | ArrayContainer], Array | ArrayContainer]][source]

Decorator around rec_map_array_container().

arraycontext.multimapped_over_array_containers(f: Callable[[...], Any] | None = None, leaf_class: type | None = None) → Callable[[...], Any] | Callable[[Callable[[...], Any]], Callable[[...], Any]][source]

Decorator around rec_multimap_array_container().

Freezing and thawing

arraycontext.freeze(ary: ArrayOrContainerT, actx: ArrayContext | None = None) → ArrayOrContainerT[source]

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: ArrayOrContainerT, actx: ArrayContext) → ArrayOrContainerT[source]

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.

Flattening and unflattening

arraycontext.flatten(ary: Array | ArrayContainer, actx: ArrayContext, *, leaf_class: type | None = None) → Any[source]

Convert all arrays in the ArrayContainer into single flat array of a type arraycontext.ArrayContext.array_types.

The operation requires arraycontext.ArrayContext.np to have ravel and concatenate methods implemented. The order in which the individual leaf arrays appear in the final array is dependent on the order given by serialize_container().

If leaf_class is given, then unflatten() will not be able to recover the original ary.

Parameters:

leaf_class – an ArrayContainer class on which the recursion is stopped (subclasses are not considered). If given, only the entries of this type are flattened and the rest of the tree structure is left as is. By default, the recursion is stopped when a non-ArrayContainer is found, which results in the whole input container ary being flattened.

arraycontext.unflatten(template: ArrayOrContainerT, ary: Array, actx: ArrayContext, *, strict: bool = True) → ArrayOrContainerT[source]

Unflatten an array ary produced by flatten() back into an ArrayContainer.

The order and sizes of each slice into ary are determined by the array container template.

Parameters:

  • ary – a flat one-dimensional array with a size that matches the number of entries in template.

  • strict – if True additional dtype and stride checking is performed on the unflattened array. Otherwise, these checks are skipped.

arraycontext.flat_size_and_dtype(ary: Array | ArrayContainer) → tuple[int, dtype[Any] | None][source]
Returns:

a tuple (size, dtype) that would be the length and numpy.dtype of the one-dimensional array returned by flatten().

Numpy conversion

arraycontext.from_numpy(ary: ndarray | int | float | complex | generic, actx: ArrayContext) → Array | ArrayContainer | int | float | complex | generic[source]

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: Array | ArrayContainer, actx: ArrayContext) → Array | ArrayContainer[source]

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

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

Algebraic operations

arraycontext.outer(a: Any, b: Any) → Any[source]

Compute the outer product of a and b while allowing either of them to be an ArrayContainer.

Tweaks the behavior of numpy.outer() to return a lower-dimensional object if either/both of a and b are scalars (whereas numpy.outer() always returns a matrix). Here the definition of “scalar” includes all non-array-container types and any scalar-like array container types.

If a and b are both array containers, the result will have the same type as a. If both are array containers and neither is an object array, they must have the same type.