Array Containers

class arraycontext.ArrayT

TypeVar for arrays.

class arraycontext.ContainerT

TypeVar for array container-like objects.

class arraycontext.ArrayOrContainerT

TypeVar for arrays or array container-like objects.

class arraycontext.ArrayContainer[source]

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

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

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

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

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][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: Any) Optional[arraycontext.context.ArrayContext][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.

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.

Containers with arithmetic

arraycontext.with_container_arithmetic(*, bcast_number: bool = True, _bcast_actx_array_type: Optional[bool] = None, 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][source]

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_actx_array_type – If True, instances of base array types of the container’s array context are broadcasted over the container. Can be True only if the container has _cls_has_array_context_attr set. Defaulted to bcast_number if _cls_has_array_context_attr is set, else False.

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

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: Union[arraycontext.container.ArrayT, arraycontext.container.ContainerT]) Union[arraycontext.container.ArrayT, arraycontext.container.ContainerT][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: Union[arraycontext.container.ArrayT, arraycontext.container.ContainerT]) Union[arraycontext.container.ArrayT, arraycontext.container.ContainerT][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) 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: Union[arraycontext.container.ArrayT, arraycontext.container.ContainerT]) Any[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) Any[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: Union[arraycontext.container.ArrayT, arraycontext.container.ContainerT]) Any[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) Any[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.

Traversing decorators

arraycontext.mapped_over_array_containers(f: Callable[[Any], Any]) Callable[[Union[arraycontext.container.ArrayT, arraycontext.container.ContainerT]], Union[arraycontext.container.ArrayT, arraycontext.container.ContainerT]][source]

Decorator around rec_map_array_container().

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

Decorator around rec_multimap_array_container().

Freezing and thawing

arraycontext.freeze(ary: Union[arraycontext.container.ArrayT, arraycontext.container.ContainerT], actx: Optional[arraycontext.context.ArrayContext] = None) Union[arraycontext.container.ArrayT, arraycontext.container.ContainerT][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: Union[arraycontext.container.ArrayT, arraycontext.container.ContainerT], actx: arraycontext.context.ArrayContext) Union[arraycontext.container.ArrayT, arraycontext.container.ContainerT][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.

Numpy conversion

arraycontext.from_numpy(ary: Any, actx: arraycontext.context.ArrayContext) Any[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: Any, actx: arraycontext.context.ArrayContext) Any[source]

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

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