FMM driver

boxtree.fmm.drive_fmm(actx: ArrayContext, wrangler: DistributedExpansionWranglerMixin, src_weight_vecs: Sequence[Array], *, global_src_idx_all_ranks: None = None, global_tgt_idx_all_ranks: None = None) → None

boxtree.fmm.drive_fmm(actx: ArrayContext, wrangler: ExpansionWranglerInterface, src_weight_vecs: Sequence[Array], *, global_src_idx_all_ranks: Sequence[Array] | None = None, global_tgt_idx_all_ranks: Sequence[Array] | None = None) → Array

Top-level driver routine for a fast multipole calculation.

In part, this is intended as a template for custom FMMs, in the sense that you may copy and paste its source code as a starting point.

Nonetheless, many common applications (such as point-to-point FMMs) can be covered by supplying the right wrangler to this routine.

Parameters:

• wrangler – An object exhibiting the ExpansionWranglerInterface. For distributed implementation, this wrangler should be a subclass of boxtree.distributed.calculation.DistributedExpansionWranglerMixin.

• src_weight_vecs – A sequence of source ‘density/weights/charges’. Passed unmodified to the wrangler. For a distributed implementation, this argument is only significant on the root rank, and worker ranks still need to supply a dummy vector.

• global_src_idx_all_ranks – (only used in the distributed implementation) A sequence of length nranks, where the i-th entry is an array representing the global indices of sources in the local tree on rank i. Each entry can be returned from generate_local_tree(). This argument is only significant on the root rank.

• global_tgt_idx_all_ranks – (only used in the distributed implementation) A sequence of length nranks, where the i-th entry is an array representing the global indices of targets in the local tree on rank i. Each entry can be returned from generate_local_tree(). This argument is only significant on the root rank.

Returns:

the potentials computed by wrangler. For the distributed implementation, the potentials are gathered and returned on the root rank. This function returns None on the worker ranks.

class boxtree.fmm.TreeIndependentDataForWrangler

An object that can be used to store information for efficient wrangler execution that depends on the kernel but not the tree and/or the traversal.

Examples of such data include generated code for carrying out translations.

Note

Instances of this type should not hold a reference (and thereby be specific to) a boxtree.Tree instance. Their purpose is to host caches for generated translation code that is reusable across trees. It is OK for these instances to be specific to a given kernel (or set of kernels).

class boxtree.fmm.ExpansionWranglerInterface(tree_indep: TreeIndependentDataForWrangler, traversal: FMMTraversalInfo)

Abstract expansion handling interface for use with drive_fmm().

See boxtree.constant_one.ConstantOneExpansionWrangler (source) for a very simple sample implementation.

Note

Wranglers may hold a reference (and thereby be specific to) a boxtree.Tree instance. TreeIndependentDataForWrangler exists to hold data that is more broadly reusable.

Changed in version 2018.1: Changed (a subset of) functions to return timing data.

Changed in version 2022.1: Removed timing data that should be handled by the ArrayContext.

tree_indep: TreeIndependentDataForWrangler

An instance of (a typically wrangler-dependent subclass of) TreeIndependentDataForWrangler.

traversal: FMMTraversalInfo

An instance of FMMTraversalInfo.

property tree: Tree

The Tree on which this wrangler operates, same as self.traversal.tree.

Particle ordering

abstractmethod reorder_sources(source_array: Array) → Array

Returns:

a copy of source_array in tree source order. source_array is in user source order.

abstractmethod reorder_potentials(potentials: Array) → Array

Parameters:

potentials – an array of potentials in tree target order.

Returns:

a copy of potentials in user target order.

Views into arrays of expansions

abstractmethod multipole_expansions_view(mpole_exps: Array, level: int) → tuple[int, Array]

Returns:

a tuple (box_start, view), where box_start is the number of the first box on level, and view is an array of the multipole expansions of the boxes on level, indexed such that the expansion of box ibox is view[ibox - box_start].

abstractmethod local_expansions_view(local_exps: Array, level: int) → tuple[int, Array]

Returns:

a tuple (box_start, view), where box_start is the number of the first box on level, and view is an array of the local expansions of the boxes on level, indexed such that the expansion of box ibox is view[ibox - box_start].

Translations

abstractmethod form_multipoles(actx: ArrayContext, level_start_source_box_nrs: Array, source_boxes: Array, src_weight_vecs: Sequence[Array]) → Array

Returns:

an expansions array containing multipole expansions in source_boxes due to sources with src_weight_vecs. All other expansions must be zero.

abstractmethod coarsen_multipoles(actx: ArrayContext, level_start_source_parent_box_nrs: Array, source_parent_boxes: Array, mpoles: Array) → Array

For each box in source_parent_boxes, gather (and translate) the box’s children’s multipole expansions in mpoles and add the resulting expansion into the box’s multipole expansion in mpoles.

Returns:

the updated mpoles.

abstractmethod eval_direct(actx: ArrayContext, target_boxes: Array, neighbor_sources_starts: Array, neighbor_sources_lists: Array, src_weight_vecs: Sequence[Array]) → Array

For each box in target_boxes, evaluate the influence of the neighbor sources due to src_weight_vecs, which use CSR-like interaction list storage and are indexed like target_boxes.

Returns:

a new potential array.

abstractmethod multipole_to_local(actx: ArrayContext, level_start_target_or_target_parent_box_nrs: Array, target_or_target_parent_boxes: Array, starts: Array, lists: Array, mpole_exps: Array) → Array

For each box in target_or_target_parent_boxes, translate and add the influence of the multipole expansion in mpole_exps into a new array of local expansions. starts and lists use CSR-like interaction list storage, and starts is indexed like target_or_target_parent_boxes.

Returns:

a new (local) expansion array.

abstractmethod eval_multipoles(actx: ArrayContext, target_boxes_by_source_level: ObjectArray1D[Array], from_sep_smaller_by_level: ObjectArray1D[BuiltList], mpole_exps: Array) → Array

For a level i, for each box in target_boxes_by_source_level[i], evaluate the multipole expansion in mpole_exps in the nearby boxes given in from_sep_smaller_by_level, and return a new potential array. starts and lists in from_sep_smaller_by_level[i] use CSR-like interaction list storage and starts is indexed like target_boxes_by_source_level[i].

Returns:

a new potential array.

abstractmethod form_locals(actx: ArrayContext, level_start_target_or_target_parent_box_nrs: Array, target_or_target_parent_boxes: Array, starts: Array, lists: Array, src_weight_vecs: Sequence[Array]) → Array

For each box in target_or_target_parent_boxes, form local expansions due to the sources in the nearby boxes given in starts and lists, and return a new local expansion array. starts and lists use CSR-like interaction list storage and starts is indexed like target_or_target_parent_boxes.

Returns:

a new local expansion array.

abstractmethod refine_locals(actx: ArrayContext, level_start_target_or_target_parent_box_nrs: Array, target_or_target_parent_boxes: Array, local_exps: Array) → Array

For each box in target_or_target_parent_boxes, translate the box’s parent’s local expansion in local_exps and add the resulting expansion into the box’s local expansion in local_exps.

Returns:

an updated local expansion array local_exps.

abstractmethod eval_locals(actx: ArrayContext, level_start_target_box_nrs: Array, target_boxes: Array, local_exps: Array) → Array

For each box in target_boxes, evaluate the local expansion in local_exps and return a new potential array.

Returns:

a new potential array.

abstractmethod finalize_potentials(actx: ArrayContext, potentials: Array) → Array

Postprocess the reordered potentials. This is where global scaling factors could be applied. This is distinct from reorder_potentials() because some derived FMMs (notably the QBX FMM) do their own reordering.

Distributed-memory communication

distribute_source_weights(actx: ArrayContext, src_weight_vecs: Sequence[Array] | None, src_idx_all_ranks: Sequence[Array] | None) → Sequence[Array]

Used by the distributed implementation for transferring needed source weights from root rank to each worker rank in the communicator.

This method needs to be called collectively by all ranks in the communicator.

Parameters:

• src_weight_vecs – a sequence of arrays, each with length nsources, representing the weights of sources on the root rank. None on worker ranks.

• src_idx_all_ranks – a sequence of length nranks, including the root rank, where the i-th entry is an array of indices, of which src_weight_vecs to be sent from the root rank to rank i. Each entry can be generated by generate_local_tree(). None on worker ranks.

Returns:

received source weights of the current rank (including the root rank).

gather_potential_results(actx: ArrayContext, potentials: Array, tgt_idx_all_ranks: Sequence[Array] | None) → Array | None

Used by the distributed implementation for gathering calculated potentials from all worker ranks in the communicator to the root rank.

This method needs to be called collectively by all ranks in the communicator.

Parameters:

• potentials – calculated potentials on each rank. This argument is significant on all ranks, including the root rank.

• tgt_idx_all_ranks – a sequence of length nranks, where the i-th entry is an array of the global potential indices of potentials from rank i. This argument is only significant on the root rank.

Returns:

Gathered potentials on the root rank. None on worker ranks.

communicate_mpoles(actx: ArrayContext, mpole_exps: Array, return_stats: bool = False) → dict[str, Any] | None

Used by the distributed implementation for forming the complete multipole expansions from the partial multipole expansions.

This function accepts partial multipole expansions in the argument mpole_exps, and modifies mpole_exps in place with the communicated and reduced multipole expansions.

This function needs to be called collectively by all ranks in the communicator.

Returns:

Statistics of the communication if return_stats is True. None otherwise.

Integrates boxtree with pyfmmlib.

class boxtree.pyfmmlib_integration.Kernel(*values)

Bases: Enum

HELMHOLTZ = 2

Helmholtz equation kernel.

LAPLACE = 1

Laplace equation kernel.

class boxtree.pyfmmlib_integration.FMMLibTreeIndependentDataForWrangler(dim: int, kernel: Kernel, *, ifgrad: bool = False)

Bases: TreeIndependentDataForWrangler

__init__(dim: int, kernel: Kernel, *, ifgrad: bool = False) → None

class boxtree.pyfmmlib_integration.FMMLibExpansionWrangler(tree_indep: FMMLibTreeIndependentDataForWrangler, traversal: FMMTraversalInfo, *, helmholtz_k: float | None = None, fmm_level_to_order: Callable[[Tree, int], int] | None = None, dipole_vec: onp.Array2D[inexact[Any]] | None = None, dipoles_already_reordered: bool = False, order: int | None = None, optimized_m2l_precomputation_memory_cutoff_bytes: int = 100000000, rotation_data: FMMLibRotationDataInterface | None = None)

Bases: ExpansionWranglerInterface

Implements the boxtree.fmm.ExpansionWranglerInterface by using pyfmmlib.

Internal bits

class boxtree.pyfmmlib_integration.FMMLibRotationDataInterface

Abstract interface for additional, optional data for precomputation of rotation matrices passed to the expansion wrangler.

abstractmethod m2l_rotation_lists() → onp.Array1D[integer[Any]]

Returns:

a numpy array mapping entries of List 2 to rotation classes.

abstractmethod m2l_rotation_angles() → onp.Array1D[floating[Any]]

Returns:

a numpy array mapping List 2 rotation classes to rotation angles.

class boxtree.pyfmmlib_integration.FMMLibRotationData(array_context: ArrayContext, trav: FMMTraversalInfo)

An implementation of the FMMLibRotationDataInterface.

trav: FMMTraversalInfo

tree: Tree

__init__(array_context: ArrayContext, trav: FMMTraversalInfo) → None

class boxtree.pyfmmlib_integration.FMMLibRotationDataNotSuppliedWarning