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 anArrayContainer
:serialize_container()
for serialization, which gives the components of the array.deserialize_container()
for deserialization, which constructs a container from a set of components.get_container_context_opt()
retrieves theArrayContext
from a container, if it has one.
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 typenumpy.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.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.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 anArrayContainer
.
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 theArrayContainer
. Components can themselves beArrayContainer
s, 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 anArrayContext
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 anArrayContext
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 anArrayContainer
. 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). Addnumpy.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 asabs()
.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 anarray_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 inmeshmode.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, sowith_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 ofserialize_container()
anddeserialize_container()
. cls must be adataclass()
.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.
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
ArrayContainer
s, 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
ArrayContainer
s.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
ArrayContainer
s, 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
ArrayContainer
s.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:
reduce_func – callable used to reduce over the components of ary if ary is an
ArrayContainer
. The callable should be associative, as forrec_map_reduce_array_container()
.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.
- 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:
reduce_func – callable used to reduce over the components of any
ArrayContainer
s in *args. The callable should be associative, as forrec_map_reduce_array_container()
.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.
- 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
ArrayContainer
s 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()
.
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()
edArrayContainer
.
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()
edArrayContainer
.
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 forthaw()
.Note
This function has the reverse argument order from the original function in
meshmode
. This was necessary becausesingledispatch()
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 typearraycontext.ArrayContext.array_types
.The operation requires
arraycontext.ArrayContext.np
to haveravel
andconcatenate
methods implemented. The order in which the individual leaf arrays appear in the final array is dependent on the order given byserialize_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 anArrayContainer
.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.
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 theArrayContainer
to the base array type ofArrayContext
.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
tonumpy
using the providedArrayContext
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 (whereasnumpy.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.