diff --git a/pytential/qbx/__init__.py b/pytential/qbx/__init__.py index 492230332..3f787666b 100644 --- a/pytential/qbx/__init__.py +++ b/pytential/qbx/__init__.py @@ -459,11 +459,11 @@ def cost_model_compute_potential_insn(self, actx, insn, bound_expr, evaluate, :arg calibration_params: a :class:`dict` of calibration parameters, mapping from parameter names to calibration values. - :arg per_box: if *true*, cost model result will be a :class:`numpy.ndarray` - or :class:`pyopencl.array.Array` with shape of the number of boxes, where - the ith entry is the sum of the cost of all stages for box i. If *false*, - cost model result will be a :class:`dict`, mapping from the stage name to - predicted cost of the stage for all boxes. + :arg per_box: if *True*, cost model result will be an array with shape + of the number of boxes, where the ith entry is the sum of the cost + of all stages for box i. If *False*, cost model result will be a + :class:`dict`, mapping from the stage name to predicted cost of the + stage for all boxes. :returns: whatever :meth:`exec_compute_potential_insn_fmm` returns. """ diff --git a/pytential/qbx/cost.py b/pytential/qbx/cost.py index b8ba17a91..28432ee47 100644 --- a/pytential/qbx/cost.py +++ b/pytential/qbx/cost.py @@ -207,14 +207,12 @@ def process_form_qbxl(self, actx: PyOpenCLArrayContext, geo_data, p2qbxl_cost, :arg geo_data: a :class:`pytential.qbx.geometry.QBXFMMGeometryData` object. :arg p2qbxl_cost: a :class:`numpy.float64` constant representing the cost of adding a source to a QBX local expansion. - :arg ndirect_sources_per_target_box: a :class:`numpy.ndarray` or - :class:`pyopencl.array.Array` of shape ``(ntarget_boxes,)``, with the - *i*th entry representing the number of direct evaluation sources (list 1, - list 3 close and list 4 close) for ``target_boxes[i]``. - :return: a :class:`numpy.ndarray` or :class:`pyopencl.array.Array` of shape - ``(ntarget_boxes,)``, with the *i*th entry representing the cost of - adding all direct evaluation sources to QBX local expansions of centers - in ``target_boxes[i]``. + :arg ndirect_sources_per_target_box: an array of shape ``(ntarget_boxes,)``, + with the *i*th entry representing the number of direct evaluation + sources (list 1, list 3 close and list 4 close) for ``target_boxes[i]``. + :return: an array of shape ``(ntarget_boxes,)``, with the *i*th entry + representing the cost of adding all direct evaluation sources to + QBX local expansions of centers in ``target_boxes[i]``. """ pass @@ -222,13 +220,12 @@ def process_form_qbxl(self, actx: PyOpenCLArrayContext, geo_data, p2qbxl_cost, def process_m2qbxl(self, actx: PyOpenCLArrayContext, geo_data, m2qbxl_cost): """ :arg geo_data: a :class:`pytential.qbx.geometry.QBXFMMGeometryData` object. - :arg m2qbxl_cost: a :class:`numpy.ndarray` or :class:`pyopencl.array.Array` - of shape ``(nlevels,)`` where the *i*th entry represents the translation - cost from multipole expansion at level *i* to a QBX center. - :return: a :class:`numpy.ndarray` or :class:`pyopencl.array.Array` of shape - ``(ntarget_boxes,)``, with the *i*th entry representing the cost of - translating multipole expansions of list 3 boxes at all source levels to - all QBX centers in ``target_boxes[i]``. + :arg m2qbxl_cost: an array of shape ``(nlevels,)`` where the *i*th entry + represents the translation cost from multipole expansion at level + *i* to a QBX center. + :return: an array of shape ``(ntarget_boxes,)``, with the *i*th entry + representing the cost of translating multipole expansions of list + 3 boxes at all source levels to all QBX centers in ``target_boxes[i]``. """ pass @@ -236,12 +233,12 @@ def process_m2qbxl(self, actx: PyOpenCLArrayContext, geo_data, m2qbxl_cost): def process_l2qbxl(self, actx: PyOpenCLArrayContext, geo_data, l2qbxl_cost): """ :arg geo_data: a :class:`pytential.qbx.geometry.QBXFMMGeometryData` object. - :arg l2qbxl_cost: a :class:`numpy.ndarray` or :class:`pyopencl.array.Array` - of shape ``(nlevels,)`` where each entry represents the translation - cost from a box local expansion to a QBX local expansion. - :return: a :class:`numpy.ndarray` or :class:`pyopencl.array.Array` of shape - ``(ntarget_boxes,)``, with each entry representing the cost of - translating box local expansions to all QBX local expansions. + :arg l2qbxl_cost: an array of shape ``(nlevels,)`` where each entry + represents the translation cost from a box local expansion to a QBX + local expansion. + :return: an array of shape ``(ntarget_boxes,)``, with each entry + representing the cost of translating box local expansions to all + QBX local expansions. """ pass @@ -251,10 +248,9 @@ def process_eval_qbxl(self, actx: PyOpenCLArrayContext, geo_data, qbxl2p_cost): :arg geo_data: a :class:`pytential.qbx.geometry.QBXFMMGeometryData` object. :arg qbxl2p_cost: a :class:`numpy.float64` constant, representing the evaluation cost of a target from its QBX local expansion. - :return: a :class:`numpy.ndarray` or :class:`pyopencl.array.Array` of shape - ``(ntarget_boxes,)``, with the *i*th entry representing the cost of - evaluating all targets associated with QBX centers in ``target_boxes[i]`` - from QBX local expansions. + :return: an array of shape ``(ntarget_boxes,)``, with the *i*th entry + representing the cost of evaluating all targets associated with QBX + centers in ``target_boxes[i]`` from QBX local expansions. """ pass @@ -267,14 +263,13 @@ def process_eval_target_specific_qbxl(self, actx: PyOpenCLArrayContext, :arg p2p_tsqbx_cost: a :class:`numpy.float64` constant representing the evaluation cost of a target from a direct evaluation source of the target box containing the expansion center. - :arg ndirect_sources_per_target_box: a :class:`numpy.ndarray` or - :class:`pyopencl.array.Array` of shape ``(ntarget_boxes,)``, with the - *i*th entry representing the number of direct evaluation sources + :arg ndirect_sources_per_target_box: an array of shape ``(ntarget_boxes,)``, + with the *i*th entry representing the number of direct evaluation sources (list 1, list 3 close and list 4 close) for ``target_boxes[i]``. - :return: a :class:`numpy.ndarray` or :class:`pyopencl.array.Array` of shape - ``(ntarget_boxes,)``, with the *i*th entry representing the evaluation - cost of all targets associated with centers in ``target_boxes[i]`` from - the direct evaluation sources of ``target_boxes[i]``. + :return: an array of shape ``(ntarget_boxes,)``, with the *i*th entry + representing the evaluation cost of all targets associated with + centers in ``target_boxes[i]`` from the direct evaluation sources of + ``target_boxes[i]``. """ pass @@ -578,9 +573,7 @@ def estimate_kernel_specific_calibration_params( class QBXCostModel(AbstractQBXCostModel, FMMCostModel): - """This class is an implementation of interface :class:`AbstractQBXCostModel` - using :mod:`pyopencl`. - """ + """An implementation of interface :class:`AbstractQBXCostModel`.""" def __init__( self, translation_cost_model_factory=make_pde_aware_translation_cost_model): @@ -670,12 +663,12 @@ def get_nqbx_centers_per_tgt_box(self, actx: PyOpenCLArrayContext, geo_data, weights=None): """ :arg geo_data: a :class:`pytential.qbx.geometry.QBXFMMGeometryData` object. - :arg weights: a :class:`pyopencl.array.Array` of shape ``(ncenters,)`` with - particle_id_dtype, the weight of each center in user order. - :return: a :class:`pyopencl.array.Array` of shape ``(ntarget_boxes,)`` with - type *particle_id_dtype* where the *i*th entry represents the number of - `geo_data.global_qbx_centers` in ``target_boxes[i]``, optionally weighted - by *weights*. + :arg weights: an array of shape ``(ncenters,)`` with ``particle_id_dtype``, + the weight of each center in user order. + :return: an array of shape ``(ntarget_boxes,)`` with type + *particle_id_dtype* where the *i*th entry represents the number of + `geo_data.global_qbx_centers` in ``target_boxes[i]``, optionally + weighted by *weights*. """ traversal = geo_data.traversal() tree = geo_data.tree() diff --git a/pytential/qbx/fmmlib.py b/pytential/qbx/fmmlib.py index 04194b03b..4ec09d50f 100644 --- a/pytential/qbx/fmmlib.py +++ b/pytential/qbx/fmmlib.py @@ -23,8 +23,6 @@ import numpy as np from pytools import memoize_method -import pyopencl as cl # noqa -import pyopencl.array # noqa: F401 from boxtree.pyfmmlib_integration import ( Kernel, @@ -33,6 +31,8 @@ from sumpy.kernel import ( LaplaceKernel, HelmholtzKernel, AxisTargetDerivative, DirectionalSourceDerivative) + +from pytential.array_context import PyOpenCLArrayContext import pytential.qbx.target_specific as ts @@ -44,11 +44,11 @@ class QBXFMMLibTreeIndependentDataForWrangler(FMMLibTreeIndependentDataForWrangler): - def __init__(self, cl_context, *, + def __init__(self, actx: PyOpenCLArrayContext, *, multipole_expansion_factory, local_expansion_factory, qbx_local_expansion_factory, target_kernels, _use_target_specific_qbx): - self.cl_context = cl_context + self._setup_actx = actx self.multipole_expansion_factory = multipole_expansion_factory self.local_expansion_factory = local_expansion_factory self.qbx_local_expansion_factory = qbx_local_expansion_factory @@ -171,12 +171,11 @@ def __init__(self, tree_indep, geo_data, dtype, dipole_vec = None if tree_indep.source_deriv_name is not None: - with cl.CommandQueue(tree_indep.cl_context) as queue: - dipole_vec = np.array([ - d_i.get(queue=queue) - for d_i in source_extra_kwargs[ - tree_indep.source_deriv_name]], - order="F") + dipole_vec = np.array([ + tree_indep._setup_actx.to_numpy(d_i) + for d_i in source_extra_kwargs[ + tree_indep.source_deriv_name]], + order="F") def inner_fmm_level_to_order(tree, level): if helmholtz_k == 0: @@ -199,6 +198,10 @@ def inner_fmm_level_to_order(tree, level): fmm_level_to_order=inner_fmm_level_to_order, rotation_data=geo_data) + @property + def _setup_actx(self): + return self.tree_indep._setup_actx + # {{{ data vector helpers def output_zeros(self): @@ -214,7 +217,7 @@ def output_zeros(self): np.zeros(nqbtl.nfiltered_targets, self.tree_indep.dtype) for k in self.tree_indep.outputs]) - def full_output_zeros(self, template_ary): + def full_output_zeros(self, actx: PyOpenCLArrayContext): """This includes QBX and non-QBX targets.""" from pytools.obj_array import make_obj_array @@ -283,7 +286,7 @@ def qbx_local_expansion_zeros(self): @log_process(logger) @return_timing_data - def form_global_qbx_locals(self, src_weight_vecs): + def form_global_qbx_locals(self, actx: PyOpenCLArrayContext, src_weight_vecs): src_weights, = src_weight_vecs if self.tree_indep.using_tsqbx: return self.qbx_local_expansion_zeros() @@ -340,7 +343,7 @@ def form_global_qbx_locals(self, src_weight_vecs): @log_process(logger) @return_timing_data - def translate_box_multipoles_to_qbx_local(self, multipole_exps): + def translate_box_multipoles_to_qbx_local(self, actx, multipole_exps): qbx_exps = self.qbx_local_expansion_zeros() geo_data = self.geo_data @@ -453,7 +456,7 @@ def translate_box_multipoles_to_qbx_local(self, multipole_exps): @log_process(logger) @return_timing_data - def translate_box_local_to_qbx_local(self, local_exps): + def translate_box_local_to_qbx_local(self, actx, local_exps): qbx_expansions = self.qbx_local_expansion_zeros() geo_data = self.geo_data @@ -546,7 +549,7 @@ def translate_box_local_to_qbx_local(self, local_exps): @log_process(logger) @return_timing_data - def eval_qbx_expansions(self, qbx_expansions): + def eval_qbx_expansions(self, actx, qbx_expansions): output = self.full_output_zeros(template_ary=qbx_expansions) geo_data = self.geo_data @@ -581,7 +584,7 @@ def eval_qbx_expansions(self, qbx_expansions): @log_process(logger) @return_timing_data - def eval_target_specific_qbx_locals(self, src_weight_vecs): + def eval_target_specific_qbx_locals(self, actx, src_weight_vecs): src_weights, = src_weight_vecs if not self.tree_indep.using_tsqbx: return self.full_output_zeros(template_ary=src_weights) @@ -634,9 +637,9 @@ def eval_target_specific_qbx_locals(self, src_weight_vecs): return output - def finalize_potentials(self, potential, template_ary): - potential = super().finalize_potentials(potential, template_ary) - return cl.array.to_device(template_ary.queue, potential) + def finalize_potentials(self, actx, potential): + potential = super().finalize_potentials(actx, potential) + return actx.from_numpy(potential) # }}} diff --git a/pytential/qbx/refinement.py b/pytential/qbx/refinement.py index 35fde2f6b..426d12be2 100644 --- a/pytential/qbx/refinement.py +++ b/pytential/qbx/refinement.py @@ -483,8 +483,7 @@ def make_empty_refine_flags(actx, density_discr): :arg density_discr: An instance of a :class:`meshmode.discretization.Discretization`. - :returns: A :class:`pyopencl.array.Array` suitable for use as refine flags, - initialized to zero. + :returns: an array suitable for use as refine flags, initialized to zero. """ result = actx.zeros(density_discr.mesh.nelements, np.int32) result.finish() diff --git a/pytential/source.py b/pytential/source.py index b581bf732..661f421de 100644 --- a/pytential/source.py +++ b/pytential/source.py @@ -102,7 +102,7 @@ class PointPotentialSource(_SumpyP2PMixin, PotentialSource): """ .. attribute:: nodes - An :class:`pyopencl.array.Array` of shape ``[ambient_dim, ndofs]``. + An array of shape ``[ambient_dim, ndofs]``. .. attribute:: ndofs diff --git a/pytential/symbolic/execution.py b/pytential/symbolic/execution.py index eaa718152..2aa1f1710 100644 --- a/pytential/symbolic/execution.py +++ b/pytential/symbolic/execution.py @@ -709,9 +709,9 @@ def cost_per_box(self, calibration_params, **kwargs): `estimate_kernel_specific_calibration_params`, or a :class:`str` "constant_one". :return: a :class:`dict` mapping from instruction to per-box cost. Each - per-box cost is represented by a :class:`numpy.ndarray` or - :class:`pyopencl.array.Array` of shape (nboxes,), where the ith entry - represents the cost of all stages for box i. + per-box cost is represented by an array of shape ``(nboxes,)``, where + the :math:`i`-th entry represents the cost of all stages for box + :math:`i`. """ array_context = _find_array_context_from_args_in_context(kwargs) @@ -733,7 +733,7 @@ def scipy_op( to be a key in :attr:`places`. :returns: An object that (mostly) satisfies the :class:`scipy.sparse.linalg.LinearOperator` protocol, except for - accepting and returning :class:`pyopencl.array.Array` arrays. + arrays supported by *actx*. """ if isinstance(self.code.result, np.ndarray): diff --git a/pytential/symbolic/matrix.py b/pytential/symbolic/matrix.py index 030412156..4d3bee1b2 100644 --- a/pytential/symbolic/matrix.py +++ b/pytential/symbolic/matrix.py @@ -73,7 +73,6 @@ class MatrixBuilderBase(EvaluationMapperBase): def __init__(self, actx, dep_expr, other_dep_exprs, dep_source, dep_discr, places, context): """ - :arg queue: a :class:`pyopencl.CommandQueue`. :arg dep_expr: symbolic expression for the input block column that the builder is evaluating. :arg other_dep_exprs: symbolic expressions for the remaining input diff --git a/pytential/symbolic/primitives.py b/pytential/symbolic/primitives.py index a4a652c71..709732cba 100644 --- a/pytential/symbolic/primitives.py +++ b/pytential/symbolic/primitives.py @@ -77,14 +77,14 @@ This can be converted to an object array by calling: :meth:`pymbolic.geometric_algebra.MultiVector.as_vector`. -:class:`pyopencl.array.Array` and :class:`meshmode.dof_array.DOFArray` instances -hold per-node degrees of freedom (and only those). Such instances do *not* occur -on the symbolic side of :mod:`pytential` at all. They're only visible either as -bound inputs (see :func:`pytential.bind`) or outputs of evaluation. Which one is -used depends on the meaning of the data being represented. If the data is -associated with a :class:`~meshmode.discretization.Discretization`, then -:class:`~meshmode.dof_array.DOFArray` is used and otherwise -:class:`~pyopencl.array.Array` is used. +:class:`meshmode.dof_array.DOFArray` instances hold per-node degrees of freedom +(and only those). Such instances do *not* occur on the symbolic side of +:mod:`pytential` at all. They're only visible either as bound inputs (see +:func:`pytential.bind`) or outputs of evaluation. Which one is used depends on +the meaning of the data being represented. If the data is associated with a +:class:`~meshmode.discretization.Discretization`, then +:class:`~meshmode.dof_array.DOFArray` is used and otherwise a base array +type for the underlying :class:`~arraycontext.ArrayContext` is used. Diagnostics ^^^^^^^^^^^ diff --git a/test/test_global_qbx.py b/test/test_global_qbx.py index 2cea9844f..663654a3f 100644 --- a/test/test_global_qbx.py +++ b/test/test_global_qbx.py @@ -396,14 +396,14 @@ def targets_from_sources(sign, dist, dim=2): target_association_code_container, associate_targets_to_qbx_centers) code_container = target_association_code_container(actx) - target_assoc = ( + target_assoc = actx.to_numpy( associate_targets_to_qbx_centers( places, places.auto_source, code_container.get_wrangler(actx), target_discrs, target_association_tolerance=1e-10) - ).get(queue=actx.queue) + ) expansion_radii = actx.to_numpy(flatten( bind(places, sym.expansion_radii(ambient_dim,