From 2092d8e9d5d1e76fe69a0417bea1add112674b0e Mon Sep 17 00:00:00 2001 From: Alexandru Fikl Date: Wed, 20 Oct 2021 21:47:39 -0500 Subject: [PATCH 01/14] remove usage of meshmode.dof_array.flatten --- pytential/linalg/proxy.py | 47 +++++++-------- pytential/qbx/__init__.py | 33 ++++++---- pytential/qbx/geometry.py | 39 ++++++------ pytential/qbx/refinement.py | 11 ++-- pytential/qbx/target_assoc.py | 16 ++--- pytential/qbx/utils.py | 13 ++-- pytential/source.py | 29 +++++---- pytential/symbolic/execution.py | 13 ++-- pytential/symbolic/matrix.py | 104 +++++++++++++++++++++----------- pytential/unregularized.py | 53 +++++++++------- 10 files changed, 212 insertions(+), 146 deletions(-) diff --git a/pytential/linalg/proxy.py b/pytential/linalg/proxy.py index 1af52e4a0..c6898e25d 100644 --- a/pytential/linalg/proxy.py +++ b/pytential/linalg/proxy.py @@ -26,7 +26,7 @@ import numpy as np import numpy.linalg as la -from arraycontext import PyOpenCLArrayContext +from arraycontext import PyOpenCLArrayContext, flatten from meshmode.discretization import Discretization from pytools import memoize_in @@ -76,10 +76,10 @@ def partition_by_nodes( builder = TreeBuilder(actx.context) - from arraycontext import thaw - from meshmode.dof_array import flatten tree, _ = builder(actx.queue, - flatten(thaw(discr.nodes(), actx)), + particles=actx.np.reshape( + flatten(discr.nodes(), actx) + (discr.ambient_dim, -1)), max_particles_in_box=max_particles_in_box, kind=tree_kind) @@ -244,8 +244,7 @@ def prg(): %(insns)s end """ % dict(insns=insns), [ - lp.GlobalArg("sources", None, - shape=(ndim, "nsources"), dim_tags="sep,C"), + lp.GlobalArg("sources", None, shape=(ndim, "nsources")), lp.ValueArg("nsources", np.int64), ... ], @@ -330,9 +329,9 @@ def __call__(self, # {{{ get proxy centers and radii - from arraycontext import thaw - from meshmode.dof_array import flatten - sources = flatten(thaw(discr.nodes(), actx)) + sources = actx.np.reshape( + flatten(discr.nodes(), actx), + (self.ambient_dim, -1)) knl = self.get_centers_knl(actx) _, (centers_dev,) = knl(actx.queue, @@ -404,8 +403,7 @@ def prg(): proxy_radius[irange] = radius_factor * rblk end """, [ - lp.GlobalArg("sources", None, - shape=(ndim, "nsources"), dim_tags="sep,C"), + lp.GlobalArg("sources", None, shape=(ndim, "nsources")), lp.ValueArg("nsources", np.int64), lp.ValueArg("radius_factor", np.float64), ... @@ -462,12 +460,9 @@ def prg(): proxy_radius[irange] = radius_factor * rqbx end """, [ - lp.GlobalArg("sources", None, - shape=(ndim, "nsources"), dim_tags="sep,C"), - lp.GlobalArg("center_int", None, - shape=(ndim, "nsources"), dim_tags="sep,C"), - lp.GlobalArg("center_ext", None, - shape=(ndim, "nsources"), dim_tags="sep,C"), + lp.GlobalArg("sources", None, shape=(ndim, "nsources")), + lp.GlobalArg("center_int", None, shape=(ndim, "nsources")), + lp.GlobalArg("center_ext", None, shape=(ndim, "nsources")), lp.ValueArg("nsources", np.int64), lp.ValueArg("radius_factor", np.float64), ... @@ -510,11 +505,12 @@ def __call__(self, center_ext = bind(self.places, sym.expansion_centers( self.ambient_dim, +1, dofdesc=source_dd))(actx) - from meshmode.dof_array import flatten return super().__call__(actx, source_dd, indices, - expansion_radii=flatten(radii), - center_int=flatten(center_int), - center_ext=flatten(center_ext), + expansion_radii=flatten(radii, actx), + center_int=actx.np.reshape( + flatten(center_int, actx), (self.ambient_dim, -1)), + center_ext=actx.np.reshape( + flatten(center_ext, actx), (self.ambient_dim, -1)), **kwargs) # }}} @@ -545,8 +541,7 @@ def prg(): """ result[idim, i] = ary[idim, srcindices[i]] """, [ - lp.GlobalArg("ary", None, - shape=(discr.ambient_dim, "ndofs"), dim_tags="sep,C"), + lp.GlobalArg("ary", None, shape=(discr.ambient_dim, "ndofs")), lp.ValueArg("ndofs", np.int64), ...], name="picker_knl", @@ -560,10 +555,10 @@ def prg(): return knl - from arraycontext import thaw - from meshmode.dof_array import flatten _, (sources,) = prg()(actx.queue, - ary=flatten(thaw(discr.nodes(), actx)), + ary=actx.np.reshape( + flatten(discr.nodes(), actx), + (discr.ambient_dim, -1)), srcindices=pxy.srcindices.indices) # }}} diff --git a/pytential/qbx/__init__.py b/pytential/qbx/__init__.py index 74ad6e772..1cffca2a4 100644 --- a/pytential/qbx/__init__.py +++ b/pytential/qbx/__init__.py @@ -23,8 +23,7 @@ import numpy as np import pyopencl as cl -from arraycontext import PyOpenCLArrayContext, freeze, thaw -from meshmode.dof_array import flatten, unflatten +from arraycontext import PyOpenCLArrayContext, thaw, flatten, unflatten from pytools import memoize_method, memoize_in, single_valued from pytential.qbx.target_assoc import QBXTargetAssociationFailedException @@ -683,7 +682,8 @@ def exec_compute_potential_insn_fmm(self, actx: PyOpenCLArrayContext, from meshmode.discretization import Discretization if isinstance(target_discr, Discretization): - result = unflatten(actx, target_discr, result) + template_ary = thaw(target_discr.nodes()[0], actx) + result = unflatten(template_ary, result, actx) results.append((o.name, result)) @@ -779,7 +779,9 @@ def exec_compute_potential_insn_direct(self, actx, insn, bound_expr, evaluate, def _flat_nodes(dofdesc): discr = bound_expr.places.get_discretization( dofdesc.geometry, dofdesc.discr_stage) - return freeze(flatten(thaw(discr.nodes(), actx), strict=False), actx) + return actx.freeze(actx.np.reshape( + flatten(discr.nodes(), actx), (self.ambient_dim, -1) + )) @memoize_in(bound_expr.places, (QBXLayerPotentialSource, "flat_expansion_radii")) @@ -788,7 +790,7 @@ def _flat_expansion_radii(dofdesc): bound_expr.places, sym.expansion_radii(self.ambient_dim, dofdesc=dofdesc), )(actx) - return freeze(flatten(radii), actx) + return actx.freeze(flatten(radii, actx)) @memoize_in(bound_expr.places, (QBXLayerPotentialSource, "flat_centers")) @@ -797,11 +799,19 @@ def _flat_centers(dofdesc, qbx_forced_limit): sym.expansion_centers( self.ambient_dim, qbx_forced_limit, dofdesc=dofdesc), )(actx) - return freeze(flatten(centers), actx) + return actx.freeze(actx.np.reshape( + flatten(centers, actx), (self.ambient_dim, -1) + )) kernel_args = {} for arg_name, arg_expr in insn.kernel_arguments.items(): - kernel_args[arg_name] = flatten(evaluate(arg_expr), strict=False) + value = evaluate(arg_expr) + + if isinstance(value, np.ndarray): + kernel_args[arg_name] = actx.np.reshape( + flatten(value, actx), (value.size, -1)) + else: + kernel_args[arg_name] = flatten(value, actx) flat_strengths = _get_flat_strengths_from_densities( actx, bound_expr.places, evaluate, insn.densities, @@ -868,7 +878,8 @@ def _flat_centers(dofdesc, qbx_forced_limit): for i, o in outputs: result = output_for_each_kernel[o.target_kernel_index] if isinstance(target_discr, Discretization): - result = unflatten(actx, target_discr, result) + template_ary = thaw(target_discr.nodes()[0], actx) + result = unflatten(template_ary, result, actx) results[i] = (o.name, result) @@ -944,7 +955,8 @@ def _flat_centers(dofdesc, qbx_forced_limit): for i, o in outputs: result = output_for_each_kernel[o.target_kernel_index] if isinstance(target_discr, Discretization): - result = unflatten(actx, target_discr, result) + template_ary = thaw(target_discr.nodes()[0], actx) + result = unflatten(template_ary, result, actx) results[i] = (o.name, result) @@ -967,8 +979,7 @@ def _get_flat_strengths_from_densities( )(actx) strength_vecs = [waa * evaluate(density) for density in densities] - from meshmode.dof_array import flatten - return [flatten(strength) for strength in strength_vecs] + return [flatten(strength, actx) for strength in strength_vecs] # }}} diff --git a/pytential/qbx/geometry.py b/pytential/qbx/geometry.py index f3623edcb..648e177a8 100644 --- a/pytential/qbx/geometry.py +++ b/pytential/qbx/geometry.py @@ -28,8 +28,7 @@ from pytools import memoize_method, log_process from pytools.obj_array import obj_array_vectorize -from arraycontext import PyOpenCLArrayContext, thaw -from meshmode.dof_array import flatten +from arraycontext import PyOpenCLArrayContext, flatten from boxtree.tools import DeviceDataRecord from boxtree.pyfmmlib_integration import FMMLibRotationDataInterface @@ -433,10 +432,14 @@ def flat_centers(self): """ from pytential import bind, sym + actx = self.array_context centers = bind(self.places, sym.interleaved_expansion_centers( self.ambient_dim, - dofdesc=self.source_dd.to_stage1()))(self.array_context) - return obj_array_vectorize(self.array_context.freeze, flatten(centers)) + dofdesc=self.source_dd.to_stage1()))(actx) + + return actx.freeze(actx.np.reshape( + flatten(centers, actx), (centers.size, -1) + )) @memoize_method def flat_expansion_radii(self): @@ -447,14 +450,14 @@ def flat_expansion_radii(self): """ from pytential import bind, sym + actx = self.array_context radii = bind(self.places, sym.expansion_radii( self.ambient_dim, granularity=sym.GRANULARITY_CENTER, - dofdesc=self.source_dd.to_stage1()))( - self.array_context) + dofdesc=self.source_dd.to_stage1()))(actx) - return self.array_context.freeze(flatten(radii)) + return actx.freeze(flatten(radii, actx)) # }}} @@ -465,7 +468,7 @@ def target_info(self): """Return a :class:`TargetInfo`. |cached|""" code_getter = self.code_getter - queue = self.array_context.queue + actx = self.array_context ntargets = self.ncenters target_discr_starts = [] @@ -475,23 +478,22 @@ def target_info(self): target_discr_starts.append(ntargets) - targets = cl.array.empty( - self.cl_context, (self.ambient_dim, ntargets), - self.coord_dtype) + targets = actx.empty((self.ambient_dim, ntargets), self.coord_dtype) code_getter.copy_targets_kernel()( - queue, + actx.queue, targets=targets[:, :self.ncenters], points=self.flat_centers()) for start, (target_discr, _) in zip( target_discr_starts, self.target_discrs_and_qbx_sides): code_getter.copy_targets_kernel()( - queue, + actx.queue, targets=targets[:, start:start+target_discr.ndofs], - points=flatten( - thaw(target_discr.nodes(), self.array_context), - strict=False) + points=actx.np.reshape( + flatten(target_discr.nodes(), actx), + (target_discr.ambient_dim, -1) + ) ) return TargetInfo( @@ -566,8 +568,9 @@ def tree(self): refine_weights.finish() tree, _ = code_getter.build_tree()(queue, - particles=flatten(thaw( - quad_stage2_discr.nodes(), self.array_context)), + particles=actx.np.reshape( + flatten(quad_stage2_discr.nodes(), actx), + (quad_stage2_discr.ambient_dim, -1)), targets=target_info.targets, target_radii=target_radii, max_leaf_refine_weight=lpot_source._max_leaf_refine_weight, diff --git a/pytential/qbx/refinement.py b/pytential/qbx/refinement.py index 7d6d6e66c..6265ed4c1 100644 --- a/pytential/qbx/refinement.py +++ b/pytential/qbx/refinement.py @@ -29,8 +29,8 @@ import loopy as lp from loopy.version import MOST_RECENT_LANGUAGE_VERSION -from arraycontext import PyOpenCLArrayContext -from meshmode.dof_array import flatten, DOFArray +from arraycontext import PyOpenCLArrayContext, flatten +from meshmode.dof_array import DOFArray from pytools import memoize_method from boxtree.area_query import AreaQueryElementwiseTemplate @@ -312,7 +312,8 @@ def check_expansion_disks_undisturbed_by_sources(self, center_danger_zone_radii = flatten( bind(stage1_density_discr, sym.expansion_radii(stage1_density_discr.ambient_dim, - granularity=sym.GRANULARITY_CENTER))(self.array_context)) + granularity=sym.GRANULARITY_CENTER))(self.array_context), + self.array_context) evt = knl( *unwrap_args( @@ -372,7 +373,7 @@ def check_sufficient_source_quadrature_resolution(self, bind(stage2_density_discr, sym._source_danger_zone_radii( stage2_density_discr.ambient_dim, dofdesc=dd)) - (self.array_context)) + (self.array_context), self.array_context) unwrap_args = AreaQueryElementwiseTemplate.unwrap_args evt = knl( @@ -410,7 +411,7 @@ def check_element_prop_threshold(self, element_property, threshold, refine_flags if debug: npanels_to_refine_prev = cl.array.sum(refine_flags).get() - element_property = flatten(element_property) + element_property = flatten(element_property, self.array_context) evt, out = knl(self.queue, element_property=element_property, diff --git a/pytential/qbx/target_assoc.py b/pytential/qbx/target_assoc.py index f244fda4d..6c572619e 100644 --- a/pytential/qbx/target_assoc.py +++ b/pytential/qbx/target_assoc.py @@ -33,8 +33,7 @@ from cgen import Enum -from arraycontext import PyOpenCLArrayContext -from meshmode.dof_array import flatten +from arraycontext import PyOpenCLArrayContext, flatten from pytential.qbx.utils import ( QBX_TREE_C_PREAMBLE, QBX_TREE_MAKO_DEFS, TreeWranglerBase, TreeCodeContainerMixin) @@ -527,8 +526,9 @@ def mark_targets(self, places, dofdesc, tunnel_radius_by_source = flatten( bind(places, - sym._close_target_tunnel_radii(ambient_dim, dofdesc=dofdesc)) - (self.array_context)) + sym._close_target_tunnel_radii(ambient_dim, dofdesc=dofdesc), + )(self.array_context), + self.array_context) # Target-marking algorithm (TGTMARK): # @@ -627,7 +627,8 @@ def find_centers(self, places, dofdesc, granularity=sym.GRANULARITY_CENTER, dofdesc=dofdesc))(self.array_context) expansion_radii_by_center_with_tolerance = flatten( - expansion_radii_by_center * (1 + target_association_tolerance)) + expansion_radii_by_center * (1 + target_association_tolerance), + self.array_context) # Idea: # @@ -724,8 +725,9 @@ def mark_panels_for_refinement(self, places, dofdesc, tunnel_radius_by_source = flatten( bind(places, - sym._close_target_tunnel_radii(ambient_dim, dofdesc=dofdesc)) - (self.array_context)) + sym._close_target_tunnel_radii(ambient_dim, dofdesc=dofdesc), + )(self.array_context), + self.array_context) # see (TGTMARK) above for algorithm. diff --git a/pytential/qbx/utils.py b/pytential/qbx/utils.py index 298bfa013..5f99b3fc7 100644 --- a/pytential/qbx/utils.py +++ b/pytential/qbx/utils.py @@ -279,12 +279,15 @@ def _make_centers(discr): stage1_density_discr = stage1_density_discrs[0] density_discr = density_discrs[0] - from arraycontext import thaw - from meshmode.dof_array import flatten - sources = flatten(thaw(density_discr.nodes(), actx)) - centers = flatten(_make_centers(stage1_density_discr)) + from arraycontext import flatten + sources = actx.np.reshape( + flatten(density_discr.nodes(), actx), + (density_discr.ambient_dim, -1)) + centers = actx.np.reshape( + flatten(_make_centers(stage1_density_discr), actx), + (stage1_density_discr.ambient_dim, -1)) targets = [ - flatten(thaw(tgt.nodes(), actx), strict=False) + actx.np.reshape(flatten(tgt.nodes(), actx), (tgt.ambient_dim, -1)) for tgt in targets_list] queue = actx.queue diff --git a/pytential/source.py b/pytential/source.py index 34d2e2e9f..8d57adc80 100644 --- a/pytential/source.py +++ b/pytential/source.py @@ -25,7 +25,6 @@ from pytools import memoize_in -from meshmode.dof_array import flatten from sumpy.fmm import UnableToCollectTimingData @@ -176,17 +175,19 @@ def exec_compute_potential_insn(self, actx, insn, bound_expr, evaluate, p2p = self.get_p2p(actx, source_kernels=insn.source_kernels, target_kernels=insn.target_kernels) - from arraycontext import thaw + from arraycontext import thaw, flatten, unflatten evt, output_for_each_kernel = p2p(actx.queue, - flatten(thaw(target_discr.nodes(), actx), strict=False), - self._nodes, - strengths, **kernel_args) + targets=actx.np.reshape( + flatten(target_discr.nodes(), actx), + (target_discr.ambient_dim, -1)), + sources=self._nodes, + strength=strengths, **kernel_args) from meshmode.discretization import Discretization result = output_for_each_kernel[o.target_kernel_index] if isinstance(target_discr, Discretization): - from meshmode.dof_array import unflatten - result = unflatten(actx, target_discr, result) + template_ary = thaw(target_discr.nodes()[0], actx) + result = unflatten(template_ary, result, actx) results.append((o.name, result)) @@ -290,7 +291,7 @@ def get_fmm_expansion_wrangler_extra_kwargs( def reorder_sources(source_array): if isinstance(source_array, cl.array.Array): - return (source_array + return source_array[].with_queue(None) .with_queue(queue) [tree_user_source_ids] .with_queue(None)) @@ -301,16 +302,20 @@ def reorder_sources(source_array): source_extra_kwargs = {} from sumpy.tools import gather_arguments, gather_source_arguments - from pytools.obj_array import obj_array_vectorize + from arraycontext import flatten, map_array_container for func, var_dict in [ (gather_arguments, kernel_extra_kwargs), (gather_source_arguments, source_extra_kwargs), ]: for arg in func(target_kernels): - var_dict[arg.name] = obj_array_vectorize( - reorder_sources, - flatten(evaluator(arguments[arg.name]), strict=False)) + value = evaluator(arguments[arg.name]) + if isinstance(value, np.ndarray): + value = actx.np.reshape(flatten(value, actx), (value.size, -1)) + else: + value = flatten(value, actx) + + var_dict[arg.name] = reorder_sources(value) return kernel_extra_kwargs, source_extra_kwargs diff --git a/pytential/symbolic/execution.py b/pytential/symbolic/execution.py index 5695d3987..5e5ae8651 100644 --- a/pytential/symbolic/execution.py +++ b/pytential/symbolic/execution.py @@ -34,7 +34,7 @@ import pyopencl.clmath # noqa from arraycontext import PyOpenCLArrayContext, thaw, freeze -from meshmode.dof_array import DOFArray, flatten +from meshmode.dof_array import DOFArray from pytools import memoize_in, memoize_method from pytential.qbx.cost import AbstractQBXCostModel @@ -470,9 +470,11 @@ def flatten(self, ary): if not self._operator_uses_obj_array: ary = [ary] + from arraycontext import flatten result = self.array_context.empty(self.total_dofs, self.dtype) for res_i, (start, end) in zip(ary, self.starts_and_ends): - result[start:end] = flatten(thaw(res_i, self.array_context)) + result[start:end] = flatten(res_i, self.array_context) + return result def unflatten(self, ary): @@ -480,10 +482,13 @@ def unflatten(self, ary): components = [] for discr, (start, end) in zip(self.discrs, self.starts_and_ends): component = ary[start:end] + from meshmode.discretization import Discretization if isinstance(discr, Discretization): - from meshmode.dof_array import unflatten - component = unflatten(self.array_context, discr, component) + from arraycontext import unflatten + template_ary = thaw(discr.nodes()[0], self.array_context) + component = unflatten(template_ary, component, self.array_context) + components.append(component) if self._operator_uses_obj_array: diff --git a/pytential/symbolic/matrix.py b/pytential/symbolic/matrix.py index 949bf9706..139ada1a7 100644 --- a/pytential/symbolic/matrix.py +++ b/pytential/symbolic/matrix.py @@ -27,9 +27,8 @@ from sys import intern -from arraycontext import thaw from pytools import memoize_method -from meshmode.dof_array import flatten, flatten_to_numpy, unflatten_from_numpy +from arraycontext import thaw, flatten, unflatten from pytential.symbolic.mappers import EvaluationMapperBase @@ -56,9 +55,13 @@ def _get_layer_potential_args(actx, places, expr, context=None, include_args=Non if include_args is not None and arg_name not in include_args: continue - kernel_args[arg_name] = flatten( - bind(places, arg_expr)(actx, **context), - strict=False) + value = bind(places, arg_expr)(actx, **context) + if isinstance(value, np.ndarray): + value = actx.np.reshape(flatten(value, actx), (value.size, -1)) + else: + value = flatten(value, actx) + + kernel_args[arg_name] = value return kernel_args @@ -198,6 +201,7 @@ def map_num_reference_derivative(self, expr): if self.is_kind_matrix(rec_operand): raise NotImplementedError("derivatives") + actx = self.array_context dofdesc = expr.dofdesc op = sym.NumReferenceDerivative( ref_axes=expr.ref_axes, @@ -205,18 +209,20 @@ def map_num_reference_derivative(self, expr): dofdesc=dofdesc) discr = self.places.get_discretization(dofdesc.geometry, dofdesc.discr_stage) - rec_operand = unflatten_from_numpy(self.array_context, discr, rec_operand) - return flatten_to_numpy(self.array_context, - bind(self.places, op)(self.array_context, u=rec_operand) - ) + template_ary = thaw(discr.nodes()[0], actx) + rec_operand = unflatten(template_ary, actx.from_numpy(rec_operand), actx) + + return actx.to_numpy(flatten( + bind(self.places, op)(self.array_context, u=rec_operand), + actx)) def map_node_coordinate_component(self, expr): from pytential import bind, sym op = sym.NodeCoordinateComponent(expr.ambient_axis, dofdesc=expr.dofdesc) - return flatten_to_numpy(self.array_context, - bind(self.places, op)(self.array_context) - ) + + actx = self.array_context + return actx.to_numpy(flatten(bind(self.places, op)(actx), actx)) def map_call(self, expr): arg, = expr.parameters @@ -229,10 +235,11 @@ def map_call(self, expr): if isinstance(rec_arg, Number): return getattr(np, expr.function.name)(rec_arg) else: - from arraycontext import from_numpy - rec_arg = from_numpy(rec_arg, self.array_context) - result = getattr(self.array_context.np, expr.function.name)(rec_arg) - return flatten_to_numpy(self.array_context, result, strict=False) + actx = self.array_context + + rec_arg = actx.from_numpy(rec_arg, actx) + result = getattr(actx.np, expr.function.name)(rec_arg) + return actx.to_numpy(flatten(result, actx), actx) # }}} @@ -332,9 +339,15 @@ def map_interpolation(self, expr): conn = self.places.get_connection(expr.from_dd, expr.to_dd) discr = self.places.get_discretization( expr.from_dd.geometry, expr.from_dd.discr_stage) - - operand = unflatten_from_numpy(actx, discr, operand) - return flatten_to_numpy(actx, conn(operand)) + template_ary = thaw(discr.nodes()[0], actx) + + from pytools.obj_array import make_obj_array + return make_obj_array([ + actx.to_numpy(flatten( + conn(unflatten(template_ary, actx.from_numpy(o), actx)), + actx)) + for o in operand + ]) elif isinstance(operand, np.ndarray) and operand.ndim == 2: cache = self.places._get_cache(MatrixBuilderDirectResamplerCacheKey) key = (expr.from_dd.geometry, @@ -402,17 +415,23 @@ def map_int_g(self, expr): dofdesc=expr.target))(actx) _, (mat,) = mat_gen(actx.queue, - targets=flatten(thaw(target_discr.nodes(), actx), strict=False), - sources=flatten(thaw(source_discr.nodes(), actx)), - centers=flatten(centers), - expansion_radii=flatten(radii), + targets=actx.np.reshape( + flatten(target_discr.nodes(), actx), + (target_discr.ambient_dim, -1)), + sources=actx.np.reshape( + flatten(source_discr.nodes(), actx), + (source_discr.ambient_dim, -1)), + centers=actx.np.reshape( + flatten(centers, actx), + (target_discr.ambient_dim, -1)), + expansion_radii=flatten(radii, actx), **kernel_args) mat = actx.to_numpy(mat) waa = bind(self.places, sym.weights_and_area_elements( source_discr.ambient_dim, dofdesc=expr.source))(actx) - mat[:, :] *= flatten_to_numpy(actx, waa) + mat[:, :] *= actx.to_numpy(flatten(waa, actx)) result += mat @ rec_density @@ -471,8 +490,12 @@ def map_int_g(self, expr): exclude_self=self.exclude_self) _, (mat,) = mat_gen(actx.queue, - targets=flatten(thaw(target_discr.nodes(), actx), strict=False), - sources=flatten(thaw(source_discr.nodes(), actx), strict=False), + targets=actx.np.reshape( + flatten(target_discr.nodes(), actx), + (target_discr.ambient_dim, -1)), + sources=actx.np.reshape( + flatten(source_discr.nodes(), actx), + (source_discr.ambient_dim, -1)), **kernel_args) mat = actx.to_numpy(mat) @@ -483,7 +506,7 @@ def map_int_g(self, expr): source_discr.ambient_dim, dofdesc=expr.source))(actx) - mat[:, :] *= flatten_to_numpy(actx, waa) + mat[:, :] *= actx.to_numpy(flatten(waa, actx)) result += mat @ rec_density @@ -545,10 +568,16 @@ def map_int_g(self, expr): dofdesc=expr.target))(actx) _, (mat,) = mat_gen(actx.queue, - targets=flatten(thaw(target_discr.nodes(), actx), strict=False), - sources=flatten(thaw(source_discr.nodes(), actx)), - centers=flatten(centers), - expansion_radii=flatten(radii), + targets=actx.np.reshape( + flatten(target_discr.nodes(), actx), + (target_discr.ambient_dim, -1)), + sources=actx.np.reshape( + flatten(source_discr.nodes(), actx), + (source_discr.ambient_dim, -1)), + centers=actx.np.reshape( + flatten(centers, actx), + (target_discr.ambient_dim, -1)), + expansion_radii=flatten(radii, actx), tgtindices=tgtindices, srcindices=srcindices, **kernel_args) @@ -557,7 +586,8 @@ def map_int_g(self, expr): bind(self.places, sym.weights_and_area_elements( source_discr.ambient_dim, - dofdesc=expr.source))(actx)) + dofdesc=expr.source))(actx), + actx) mat *= waa[srcindices] result += actx.to_numpy(mat) * rec_density @@ -616,8 +646,12 @@ def map_int_g(self, expr): exclude_self=self.exclude_self) _, (mat,) = mat_gen(actx.queue, - targets=flatten(thaw(target_discr.nodes(), actx), strict=False), - sources=flatten(thaw(source_discr.nodes(), actx), strict=False), + targets=actx.np.reshape( + flatten(target_discr.nodes(), actx), + (target_discr.ambient_dim, -1)), + sources=actx.np.reshape( + flatten(source_discr.nodes(), actx), + (source_discr.ambient_dim, -1)), tgtindices=tgtindices, srcindices=srcindices, **kernel_args) @@ -628,7 +662,7 @@ def map_int_g(self, expr): waa = bind(self.places, sym.weights_and_area_elements( source_discr.ambient_dim, dofdesc=expr.source))(actx) - waa = flatten(waa) + waa = flatten(waa, actx) mat *= waa[srcindices] diff --git a/pytential/unregularized.py b/pytential/unregularized.py index fcd6aacd2..5967f1d45 100644 --- a/pytential/unregularized.py +++ b/pytential/unregularized.py @@ -30,7 +30,7 @@ from loopy.version import MOST_RECENT_LANGUAGE_VERSION from pytools import memoize_method -from arraycontext import PyOpenCLArrayContext, thaw +from arraycontext import PyOpenCLArrayContext, thaw, flatten, unflatten from boxtree.tools import DeviceDataRecord from pytential.source import LayerPotentialSourceBase @@ -140,16 +140,20 @@ def exec_compute_potential_insn_direct(self, actx: PyOpenCLArrayContext, insn, bound_expr, evaluate): kernel_args = {} - from meshmode.dof_array import flatten, unflatten - for arg_name, arg_expr in insn.kernel_arguments.items(): - kernel_args[arg_name] = flatten(evaluate(arg_expr)) + value = evaluate(arg_expr) + if isinstance(value, np.ndarray): + value = actx.np.reshape(flatten(value, actx), (value.size, -1)) + else: + value = flatten(value, actx) + + kernel_args[arg_name] = value from pytential import bind, sym waa = bind(bound_expr.places, sym.weights_and_area_elements( self.ambient_dim, dofdesc=insn.source))(actx) strengths = [waa * evaluate(density) for density in insn.densities] - flat_strengths = [flatten(strength) for strength in strengths] + flat_strengths = [flatten(strength, actx) for strength in strengths] results = [] p2p = None @@ -163,14 +167,19 @@ def exec_compute_potential_insn_direct(self, actx: PyOpenCLArrayContext, target_kernels=insn.target_kernels) evt, output_for_each_kernel = p2p(actx.queue, - flatten(thaw(target_discr.nodes(), actx), strict=False), - flatten(thaw(self.density_discr.nodes(), actx)), - flat_strengths, **kernel_args) + targets=actx.np.reshape( + flatten(target_discr.nodes(), actx), + (target_discr.ambient_dim, -1)), + sources=actx.np.reshape( + flatten(self.density_discr.nodes(), actx) + (self.density_discr.ambient_dim, -1)), + strength=flat_strengths, **kernel_args) from meshmode.discretization import Discretization result = output_for_each_kernel[o.target_kernel_index] if isinstance(target_discr, Discretization): - result = unflatten(actx, target_discr, result) + template_ary = thaw(target_discr.nodes()[0], actx) + result = unflatten(template_ary, result, actx) results.append((o.name, result)) @@ -239,8 +248,7 @@ def exec_compute_potential_insn_fmm(self, actx: PyOpenCLArrayContext, self.ambient_dim, dofdesc=insn.source))(actx) strengths = [waa * evaluate(density) for density in insn.densities] - from meshmode.dof_array import flatten - flat_strengths = [flatten(strength) for strength in strengths] + flat_strengths = [flatten(strength, actx) for strength in strengths] fmm_kernel = self.get_fmm_kernel(insn.target_kernels) output_and_expansion_dtype = ( @@ -284,8 +292,8 @@ def exec_compute_potential_insn_fmm(self, actx: PyOpenCLArrayContext, from meshmode.discretization import Discretization if isinstance(target_discr, Discretization): - from meshmode.dof_array import unflatten - result = unflatten(actx, target_discr, result) + template_ary = thaw(target_discr.nodes()[0], actx) + result = unflatten(template_ary, result, actx) results.append((o.name, result)) @@ -403,21 +411,21 @@ def tree(self): lpot_src = self.lpot_source target_info = self.target_info() - queue = self.array_context.queue + actx = self.array_context nsources = lpot_src.density_discr.ndofs nparticles = nsources + target_info.ntargets - refine_weights = cl.array.zeros(queue, nparticles, dtype=np.int32) + refine_weights = cl.array.zeros(actx.queue, nparticles, dtype=np.int32) refine_weights[:nsources] = 1 refine_weights.finish() MAX_LEAF_REFINE_WEIGHT = 32 # noqa - from meshmode.dof_array import flatten - tree, _ = code_getter.build_tree(queue, - particles=flatten( - thaw(lpot_src.density_discr.nodes(), self.array_context)), + tree, _ = code_getter.build_tree(actx.queue, + particles=actx.np.reshape( + flatten(lpot_src.density_discr.nodes(), actx), + (lpot_src.ambient_dim, -1)), targets=target_info.targets, max_leaf_refine_weight=MAX_LEAF_REFINE_WEIGHT, refine_weights=refine_weights, @@ -445,14 +453,13 @@ def target_info(self): (lpot_src.ambient_dim, ntargets), self.coord_dtype) - from meshmode.dof_array import flatten for start, target_discr in zip(target_discr_starts, target_discrs): code_getter.copy_targets_kernel()( self.array_context.queue, targets=targets[:, start:start+target_discr.ndofs], - points=flatten( - thaw(target_discr.nodes(), self.array_context), - strict=False) + points=self.array_context.np.reshape( + flatten(target_discr.nodes(), self.array_context), + (target_discr.ambient_dim, -1)) ) return _TargetInfo( From 430b9ff55b56d86339d29046fce861fdd9da1765 Mon Sep 17 00:00:00 2001 From: Alexandru Fikl Date: Fri, 22 Oct 2021 16:47:58 -0500 Subject: [PATCH 02/14] fix some flattening issues --- pytential/linalg/proxy.py | 2 +- pytential/qbx/__init__.py | 34 ++++++------ pytential/qbx/geometry.py | 4 +- pytential/qbx/interactions.py | 99 ++++++++++++++++++++++++----------- pytential/qbx/utils.py | 19 +++---- pytential/source.py | 48 +++++++++++------ pytential/symbolic/matrix.py | 36 ++++++------- pytential/unregularized.py | 2 +- 8 files changed, 147 insertions(+), 97 deletions(-) diff --git a/pytential/linalg/proxy.py b/pytential/linalg/proxy.py index c6898e25d..601b6a145 100644 --- a/pytential/linalg/proxy.py +++ b/pytential/linalg/proxy.py @@ -78,7 +78,7 @@ def partition_by_nodes( tree, _ = builder(actx.queue, particles=actx.np.reshape( - flatten(discr.nodes(), actx) + flatten(discr.nodes(), actx), (discr.ambient_dim, -1)), max_particles_in_box=max_particles_in_box, kind=tree_kind) diff --git a/pytential/qbx/__init__.py b/pytential/qbx/__init__.py index 1cffca2a4..5d5c944d4 100644 --- a/pytential/qbx/__init__.py +++ b/pytential/qbx/__init__.py @@ -621,7 +621,7 @@ def exec_compute_potential_insn_fmm(self, actx: PyOpenCLArrayContext, # FIXME don't compute *all* output kernels on all targets--respect that # some target discretizations may only be asking for derivatives (e.g.) - flat_strengths = _get_flat_strengths_from_densities( + flat_strengths = get_flat_strengths_from_densities( actx, bound_expr.places, evaluate, insn.densities, dofdesc=insn.source) @@ -682,7 +682,10 @@ def exec_compute_potential_insn_fmm(self, actx: PyOpenCLArrayContext, from meshmode.discretization import Discretization if isinstance(target_discr, Discretization): - template_ary = thaw(target_discr.nodes()[0], actx) + template_ary = actx.np.astype( + thaw(target_discr.nodes()[0], actx), + result.dtype) + result = unflatten(template_ary, result, actx) results.append((o.name, result)) @@ -803,17 +806,10 @@ def _flat_centers(dofdesc, qbx_forced_limit): flatten(centers, actx), (self.ambient_dim, -1) )) - kernel_args = {} - for arg_name, arg_expr in insn.kernel_arguments.items(): - value = evaluate(arg_expr) - - if isinstance(value, np.ndarray): - kernel_args[arg_name] = actx.np.reshape( - flatten(value, actx), (value.size, -1)) - else: - kernel_args[arg_name] = flatten(value, actx) - - flat_strengths = _get_flat_strengths_from_densities( + from pytential.source import evaluate_kernel_arguments + flat_kernel_args = evaluate_kernel_arguments( + actx, evaluate, insn.kernel_arguments, flat=True) + flat_strengths = get_flat_strengths_from_densities( actx, bound_expr.places, evaluate, insn.densities, dofdesc=insn.source) @@ -873,12 +869,14 @@ def _flat_centers(dofdesc, qbx_forced_limit): centers=_flat_centers(target_name, qbx_forced_limit), strengths=flat_strengths, expansion_radii=_flat_expansion_radii(target_name), - **kernel_args) + **flat_kernel_args) for i, o in outputs: result = output_for_each_kernel[o.target_kernel_index] if isinstance(target_discr, Discretization): - template_ary = thaw(target_discr.nodes()[0], actx) + template_ary = actx.np.astype( + thaw(target_discr.nodes()[0], actx), + result.dtype) result = unflatten(template_ary, result, actx) results[i] = (o.name, result) @@ -902,7 +900,7 @@ def _flat_centers(dofdesc, qbx_forced_limit): targets=flat_target_nodes, sources=flat_source_nodes, strength=flat_strengths, - **kernel_args) + **flat_kernel_args) target_discrs_and_qbx_sides = ((target_discr, qbx_forced_limit),) geo_data = self.qbx_fmm_geometry_data( @@ -936,7 +934,7 @@ def _flat_centers(dofdesc, qbx_forced_limit): qbx_center_numbers = tgt_to_qbx_center[qbx_tgt_numbers] qbx_center_numbers.finish() - tgt_subset_kwargs = kernel_args.copy() + tgt_subset_kwargs = flat_kernel_args.copy() for i, res_i in enumerate(output_for_each_kernel): tgt_subset_kwargs[f"result_{i}"] = res_i @@ -970,7 +968,7 @@ def _flat_centers(dofdesc, qbx_forced_limit): # }}} -def _get_flat_strengths_from_densities( +def get_flat_strengths_from_densities( actx, places, evaluate, densities, dofdesc=None): from pytential import bind, sym waa = bind( diff --git a/pytential/qbx/geometry.py b/pytential/qbx/geometry.py index 648e177a8..838f2e2ac 100644 --- a/pytential/qbx/geometry.py +++ b/pytential/qbx/geometry.py @@ -568,8 +568,8 @@ def tree(self): refine_weights.finish() tree, _ = code_getter.build_tree()(queue, - particles=actx.np.reshape( - flatten(quad_stage2_discr.nodes(), actx), + particles=self.array_context.np.reshape( + flatten(quad_stage2_discr.nodes(), self.array_context), (quad_stage2_discr.ambient_dim, -1)), targets=target_info.targets, target_radii=target_radii, diff --git a/pytential/qbx/interactions.py b/pytential/qbx/interactions.py index f08366a03..3c2100a9e 100644 --- a/pytential/qbx/interactions.py +++ b/pytential/qbx/interactions.py @@ -46,9 +46,8 @@ def get_kernel(self): from sumpy.tools import gather_loopy_source_arguments arguments = ( [ - lp.GlobalArg("sources", None, shape=(self.dim, "nsources"), - dim_tags="sep,c"), - lp.GlobalArg("strengths", None, dim_tags="sep,c", + lp.GlobalArg("sources", None, shape=(self.dim, "nsources")), + lp.GlobalArg("strengths", None, dim_tags="sep,C", shape="strength_count, nsources"), lp.GlobalArg("qbx_center_to_target_box", None, shape=None), @@ -56,14 +55,13 @@ def get_kernel(self): None, shape=None), lp.GlobalArg("box_source_starts,box_source_counts_nonchild", None, shape=None), - lp.GlobalArg("qbx_centers", None, shape="dim, ncenters", - dim_tags="sep,c"), + lp.GlobalArg("qbx_centers", None, shape=("dim", "ncenters")), lp.GlobalArg("qbx_expansion_radii", None, shape="ncenters"), lp.GlobalArg("qbx_expansions", None, shape=("ncenters", ncoeffs)), lp.ValueArg("ncenters", np.int32), lp.ValueArg("nsources", np.int32), - "..." + ... ] + gather_loopy_source_arguments( self.source_kernels + (self.expansion,)) ) @@ -113,6 +111,7 @@ def get_kernel(self): silenced_warnings="write_race(write_expn*)", fixed_parameters=dict(dim=self.dim, strength_count=self.strength_count), + default_offset=lp.auto, lang_version=MOST_RECENT_LANGUAGE_VERSION) for knl in self.source_kernels: @@ -122,15 +121,28 @@ def get_kernel(self): return loopy_knl @memoize_method - def get_optimized_kernel(self): + def get_optimized_kernel(self, is_sources_obj_array, is_centers_obj_array): # FIXME knl = self.get_kernel() + + if is_sources_obj_array: + knl = lp.tag_array_axes(knl, "sources", "sep,C") + if is_centers_obj_array: + knl = lp.tag_array_axes(knl, "qbx_centers", "sep,C") + knl = lp.split_iname(knl, "itgt_center", 16, outer_tag="g.0") knl = self._allow_redundant_execution_of_knl_scaling(knl) return knl def __call__(self, queue, **kwargs): - return self.get_cached_optimized_kernel()(queue, **kwargs) + sources = kwargs.pop("sources") + qbx_centers = kwargs.pop("qbx_centers") + + from sumpy.tools import is_obj_array_like + return self.get_cached_optimized_kernel( + is_sources_obj_array=is_obj_array_like(sources), + is_centers_obj_array=is_obj_array_like(qbx_centers), + )(queue, sources=sources, qbx_centers=qbx_centers, **kwargs) # }}} @@ -203,8 +215,7 @@ def get_kernel(self): lp.ValueArg("src_rscale", None), lp.GlobalArg("src_box_starts, src_box_lists", None, shape=None, strides=(1,)), - lp.GlobalArg("qbx_centers", None, shape="dim, ncenters", - dim_tags="sep,c"), + lp.GlobalArg("qbx_centers", None, shape="dim, ncenters"), lp.GlobalArg("qbx_expansion_radii", None, shape="ncenters"), lp.ValueArg("aligned_nboxes,nsrc_level_boxes", np.int32), lp.ValueArg("src_base_ibox", np.int32), @@ -212,7 +223,7 @@ def get_kernel(self): shape=("nsrc_level_boxes", ncoeff_src), offset=lp.auto), lp.GlobalArg("qbx_expansions", None, shape=("ncenters", ncoeff_tgt)), - "..." + ... ] + gather_loopy_arguments([self.src_expansion, self.tgt_expansion]), name=self.name, assumptions="ncenters>=1", silenced_warnings="write_race(write_expn*)", @@ -227,9 +238,13 @@ def get_kernel(self): return loopy_knl @memoize_method - def get_optimized_kernel(self): + def get_optimized_kernel(self, is_centers_obj_array): # FIXME knl = self.get_kernel() + + if is_centers_obj_array: + knl = lp.tag_array_axes(knl, "qbx_centers", "sep,C") + knl = lp.split_iname(knl, "icenter", 16, outer_tag="g.0") knl = self._allow_redundant_execution_of_knl_scaling(knl) return knl @@ -240,10 +255,13 @@ def __call__(self, queue, **kwargs): # meaningfully inferred. Make the type of rscale explicit. src_rscale = centers.dtype.type(kwargs.pop("src_rscale")) - return self.get_cached_optimized_kernel()(queue, - centers=centers, - src_rscale=src_rscale, - **kwargs) + from sumpy.tools import is_obj_array_like + qbx_centers = kwargs.pop("qbx_centers") + return self.get_cached_optimized_kernel( + is_centers_obj_array=is_obj_array_like(qbx_centers), + )(queue, centers=centers, + qbx_centers=qbx_centers, + src_rscale=src_rscale, **kwargs) # }}} @@ -308,13 +326,12 @@ def get_kernel(self): offset=lp.auto), lp.GlobalArg("centers", None, shape="dim, naligned_boxes"), lp.ValueArg("src_rscale", None), - lp.GlobalArg("qbx_centers", None, shape="dim, ncenters", - dim_tags="sep,c"), + lp.GlobalArg("qbx_centers", None, shape="dim, ncenters"), lp.GlobalArg("qbx_expansion_radii", None, shape="ncenters"), lp.ValueArg("naligned_boxes,target_base_ibox,nboxes", np.int32), lp.GlobalArg("expansions", None, shape=("nboxes", ncoeff_src), offset=lp.auto), - "..." + ... ] + gather_loopy_arguments([self.src_expansion, self.tgt_expansion]), name=self.name, assumptions="ncenters>=1", @@ -330,9 +347,13 @@ def get_kernel(self): return loopy_knl @memoize_method - def get_optimized_kernel(self): + def get_optimized_kernel(self, is_centers_obj_array): # FIXME knl = self.get_kernel() + + if is_centers_obj_array: + knl = lp.tag_array_axes(knl, "qbx_centers", "sep,C") + knl = lp.split_iname(knl, "icenter", 16, outer_tag="g.0") knl = self._allow_redundant_execution_of_knl_scaling(knl) return knl @@ -343,10 +364,13 @@ def __call__(self, queue, **kwargs): # meaningfully inferred. Make the type of rscale explicit. src_rscale = centers.dtype.type(kwargs.pop("src_rscale")) - return self.get_cached_optimized_kernel()(queue, - centers=centers, - src_rscale=src_rscale, - **kwargs) + from sumpy.tools import is_obj_array_like + qbx_centers = kwargs.pop("qbx_centers") + return self.get_cached_optimized_kernel( + is_centers_obj_array=is_obj_array_like(qbx_centers), + )(queue, centers=centers, + qbx_centers=qbx_centers, + src_rscale=src_rscale, **kwargs) # }}} @@ -402,18 +426,16 @@ def get_kernel(self): end """], [ - lp.GlobalArg("result", None, shape="nresults, ntargets", + lp.GlobalArg("result", None, shape=("nresults", "ntargets"), dim_tags="sep,C"), - lp.GlobalArg("qbx_centers", None, shape="dim, ncenters", - dim_tags="sep,c"), + lp.GlobalArg("qbx_centers", None, shape=("dim", "ncenters")), lp.GlobalArg("center_to_targets_starts,center_to_targets_lists", None, shape=None), lp.GlobalArg("qbx_expansions", None, shape=("ncenters", ncoeffs)), - lp.GlobalArg("targets", None, shape=(self.dim, "ntargets"), - dim_tags="sep,C"), + lp.GlobalArg("targets", None, shape=(self.dim, "ntargets")), lp.ValueArg("ncenters,ntargets", np.int32), - "..." + ... ] + [arg.loopy_arg for arg in self.expansion.get_args()], name=self.name, assumptions="nglobal_qbx_centers>=1", @@ -428,15 +450,28 @@ def get_kernel(self): return loopy_knl @memoize_method - def get_optimized_kernel(self): + def get_optimized_kernel(self, is_targets_obj_array, is_centers_obj_array): # FIXME knl = self.get_kernel() + + if is_targets_obj_array: + knl = lp.tag_array_axes(knl, "targets", "sep,C") + if is_centers_obj_array: + knl = lp.tag_array_axes(knl, "qbx_centers", "sep,C") + knl = lp.tag_inames(knl, dict(iglobal_center="g.0")) knl = self._allow_redundant_execution_of_knl_scaling(knl) return knl def __call__(self, queue, **kwargs): - return self.get_cached_optimized_kernel()(queue, **kwargs) + targets = kwargs.pop("targets") + qbx_centers = kwargs.pop("qbx_centers") + + from sumpy.tools import is_obj_array_like + return self.get_cached_optimized_kernel( + is_targets_obj_array=is_obj_array_like(targets), + is_centers_obj_array=is_obj_array_like(qbx_centers), + )(queue, targets=targets, qbx_centers=qbx_centers, **kwargs) # }}} diff --git a/pytential/qbx/utils.py b/pytential/qbx/utils.py index 5f99b3fc7..29514bd14 100644 --- a/pytential/qbx/utils.py +++ b/pytential/qbx/utils.py @@ -276,19 +276,20 @@ def _make_centers(discr): return bind(discr, sym.interleaved_expansion_centers( discr.ambient_dim))(actx) + def _flatten_nodes(ary): + if isinstance(ary, np.ndarray): + return actx.np.reshape(flatten(ary, actx), (ary.size, -1)) + else: + assert ary.shape[0] == density_discr.ambient_dim + return ary + stage1_density_discr = stage1_density_discrs[0] density_discr = density_discrs[0] from arraycontext import flatten - sources = actx.np.reshape( - flatten(density_discr.nodes(), actx), - (density_discr.ambient_dim, -1)) - centers = actx.np.reshape( - flatten(_make_centers(stage1_density_discr), actx), - (stage1_density_discr.ambient_dim, -1)) - targets = [ - actx.np.reshape(flatten(tgt.nodes(), actx), (tgt.ambient_dim, -1)) - for tgt in targets_list] + sources =_flatten_nodes(density_discr.nodes()) + centers = _flatten_nodes(_make_centers(stage1_density_discr)) + targets = [_flatten_nodes(tgt.nodes()) for tgt in targets_list] queue = actx.queue particles = tuple( diff --git a/pytential/source.py b/pytential/source.py index 8d57adc80..ff284867a 100644 --- a/pytential/source.py +++ b/pytential/source.py @@ -24,6 +24,7 @@ import pyopencl as cl from pytools import memoize_in +from arraycontext import thaw, flatten, unflatten from sumpy.fmm import UnableToCollectTimingData @@ -83,6 +84,26 @@ def p2p(target_kernels, source_kernels): # {{{ point potential source +def evaluate_kernel_arguments(actx, evaluate, kernel_arguments, flat=True): + from meshmode.dof_array import DOFArray + + kernel_args = {} + for arg_name, arg_expr in kernel_arguments.items(): + value = evaluate(arg_expr) + + if flat: + if isinstance(value, np.ndarray): + value = actx.np.reshape(flatten(value, actx), (value.size, -1)) + elif isinstance(value, DOFArray): + value = flatten(value, actx) + else: + pass + + kernel_args[arg_name] = value + + return kernel_args + + class PointPotentialSource(_SumpyP2PMixin, PotentialSource): """ .. attribute:: nodes @@ -158,10 +179,8 @@ def exec_compute_potential_insn(self, actx, insn, bound_expr, evaluate, p2p = None - kernel_args = {} - for arg_name, arg_expr in insn.kernel_arguments.items(): - kernel_args[arg_name] = evaluate(arg_expr) - + kernel_args = evaluate_kernel_arguments( + actx, evaluate, insn.kernel_arguments, flat=False) strengths = [evaluate(density) for density in insn.densities] # FIXME: Do this all at once @@ -175,7 +194,6 @@ def exec_compute_potential_insn(self, actx, insn, bound_expr, evaluate, p2p = self.get_p2p(actx, source_kernels=insn.source_kernels, target_kernels=insn.target_kernels) - from arraycontext import thaw, flatten, unflatten evt, output_for_each_kernel = p2p(actx.queue, targets=actx.np.reshape( flatten(target_discr.nodes(), actx), @@ -287,14 +305,9 @@ def get_fmm_expansion_wrangler_extra_kwargs( # This contains things like the Helmholtz parameter k or # the normal directions for double layers. - queue = actx.queue - def reorder_sources(source_array): if isinstance(source_array, cl.array.Array): - return source_array[].with_queue(None) - .with_queue(queue) - [tree_user_source_ids] - .with_queue(None)) + return source_array[tree_user_source_ids].with_queue(None) else: return source_array @@ -302,7 +315,8 @@ def reorder_sources(source_array): source_extra_kwargs = {} from sumpy.tools import gather_arguments, gather_source_arguments - from arraycontext import flatten, map_array_container + from arraycontext import map_array_container + from meshmode.dof_array import DOFArray for func, var_dict in [ (gather_arguments, kernel_extra_kwargs), @@ -311,11 +325,15 @@ def reorder_sources(source_array): for arg in func(target_kernels): value = evaluator(arguments[arg.name]) if isinstance(value, np.ndarray): - value = actx.np.reshape(flatten(value, actx), (value.size, -1)) + value = map_array_container( + lambda ary: reorder_sources(flatten(value, actx)), + value) + elif isinstance(value, DOFArray): + value = reorder_sources(flatten(value, actx)) else: - value = flatten(value, actx) + pass - var_dict[arg.name] = reorder_sources(value) + var_dict[arg.name] = value return kernel_extra_kwargs, source_extra_kwargs diff --git a/pytential/symbolic/matrix.py b/pytential/symbolic/matrix.py index 139ada1a7..5cb09b742 100644 --- a/pytential/symbolic/matrix.py +++ b/pytential/symbolic/matrix.py @@ -41,29 +41,27 @@ def is_zero(x): def _get_layer_potential_args(actx, places, expr, context=None, include_args=None): """ - :arg mapper: a :class:`~pytential.symbolic.matrix.MatrixBuilderBase`. - :arg expr: symbolic layer potential expression. - - :return: a mapping of kernel arguments evaluated by the *mapper*. + :arg expr: symbolic layer potential expression containing the kernel arguments. + :arg include_args: subset of the kernel arguments to evaluate. """ + from pytential import bind if context is None: context = {} - kernel_args = {} - for arg_name, arg_expr in expr.kernel_arguments.items(): - if include_args is not None and arg_name not in include_args: - continue - - value = bind(places, arg_expr)(actx, **context) - if isinstance(value, np.ndarray): - value = actx.np.reshape(flatten(value, actx), (value.size, -1)) - else: - value = flatten(value, actx) - - kernel_args[arg_name] = value + if include_args is not None: + kernel_arguments = { + k: v for k, v in expr.kernel_arguments.items() + if k in include_args + } + else: + kernel_arguments = expr.kernel_arguments - return kernel_args + from pytential.source import evaluate_kernel_arguments + return evaluate_kernel_arguments( + actx, + lambda expr: bind(places, expr)(actx, **context), + kernel_arguments, flat=True) # }}} @@ -237,9 +235,9 @@ def map_call(self, expr): else: actx = self.array_context - rec_arg = actx.from_numpy(rec_arg, actx) + rec_arg = actx.from_numpy(rec_arg) result = getattr(actx.np, expr.function.name)(rec_arg) - return actx.to_numpy(flatten(result, actx), actx) + return actx.to_numpy(flatten(result, actx)) # }}} diff --git a/pytential/unregularized.py b/pytential/unregularized.py index 5967f1d45..238176b28 100644 --- a/pytential/unregularized.py +++ b/pytential/unregularized.py @@ -171,7 +171,7 @@ def exec_compute_potential_insn_direct(self, actx: PyOpenCLArrayContext, flatten(target_discr.nodes(), actx), (target_discr.ambient_dim, -1)), sources=actx.np.reshape( - flatten(self.density_discr.nodes(), actx) + flatten(self.density_discr.nodes(), actx), (self.density_discr.ambient_dim, -1)), strength=flat_strengths, **kernel_args) From ede42100cbc3ad96d042ccb81e8239421ff7ad45 Mon Sep 17 00:00:00 2001 From: Alexandru Fikl Date: Fri, 22 Oct 2021 16:51:29 -0500 Subject: [PATCH 03/14] bump requirements to fixed branches --- requirements.txt | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/requirements.txt b/requirements.txt index 1532b3cc3..9bd896d9b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -9,7 +9,7 @@ git+https://github.com/inducer/pyopencl.git#egg=pyopencl git+https://github.com/inducer/islpy.git#egg=islpy git+https://github.com/inducer/loopy.git#egg=loopy git+https://github.com/inducer/boxtree.git#egg=boxtree -git+https://github.com/inducer/arraycontext.git#egg=arraycontext +git+https://github.com/alexfikl/arraycontext.git@add-astype#egg=arraycontext git+https://github.com/inducer/meshmode.git#egg=meshmode -git+https://github.com/inducer/sumpy.git#egg=sumpy +git+https://github.com/alexfikl/sumpy.git@dir-auto-offset#egg=sumpy git+https://github.com/inducer/pyfmmlib.git#egg=pyfmmlib From 78536b9a4e45b33a24498bff3508ea1313f0ac2e Mon Sep 17 00:00:00 2001 From: Alexandru Fikl Date: Fri, 22 Oct 2021 16:55:21 -0500 Subject: [PATCH 04/14] fix some flake8 errors --- pytential/qbx/geometry.py | 2 -- pytential/qbx/refinement.py | 2 +- pytential/qbx/target_assoc.py | 6 ++++-- pytential/qbx/utils.py | 2 +- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/pytential/qbx/geometry.py b/pytential/qbx/geometry.py index 838f2e2ac..1f482d425 100644 --- a/pytential/qbx/geometry.py +++ b/pytential/qbx/geometry.py @@ -26,8 +26,6 @@ import pyopencl.array # noqa from pytools import memoize_method, log_process -from pytools.obj_array import obj_array_vectorize - from arraycontext import PyOpenCLArrayContext, flatten from boxtree.tools import DeviceDataRecord diff --git a/pytential/qbx/refinement.py b/pytential/qbx/refinement.py index 6265ed4c1..e85e66283 100644 --- a/pytential/qbx/refinement.py +++ b/pytential/qbx/refinement.py @@ -313,7 +313,7 @@ def check_expansion_disks_undisturbed_by_sources(self, bind(stage1_density_discr, sym.expansion_radii(stage1_density_discr.ambient_dim, granularity=sym.GRANULARITY_CENTER))(self.array_context), - self.array_context) + self.array_context) evt = knl( *unwrap_args( diff --git a/pytential/qbx/target_assoc.py b/pytential/qbx/target_assoc.py index 6c572619e..4b54f8272 100644 --- a/pytential/qbx/target_assoc.py +++ b/pytential/qbx/target_assoc.py @@ -525,7 +525,8 @@ def mark_targets(self, places, dofdesc, axis.with_queue(self.queue)[source_slice] for axis in tree.sources] tunnel_radius_by_source = flatten( - bind(places, + bind( + places, sym._close_target_tunnel_radii(ambient_dim, dofdesc=dofdesc), )(self.array_context), self.array_context) @@ -724,7 +725,8 @@ def mark_panels_for_refinement(self, places, dofdesc, axis.with_queue(self.queue)[source_slice] for axis in tree.sources] tunnel_radius_by_source = flatten( - bind(places, + bind( + places, sym._close_target_tunnel_radii(ambient_dim, dofdesc=dofdesc), )(self.array_context), self.array_context) diff --git a/pytential/qbx/utils.py b/pytential/qbx/utils.py index 29514bd14..d427ac012 100644 --- a/pytential/qbx/utils.py +++ b/pytential/qbx/utils.py @@ -287,7 +287,7 @@ def _flatten_nodes(ary): density_discr = density_discrs[0] from arraycontext import flatten - sources =_flatten_nodes(density_discr.nodes()) + sources = _flatten_nodes(density_discr.nodes()) centers = _flatten_nodes(_make_centers(stage1_density_discr)) targets = [_flatten_nodes(tgt.nodes()) for tgt in targets_list] From 0c3f97278c5f26c9792d2f78592326fbb005c5fb Mon Sep 17 00:00:00 2001 From: Alexandru Fikl Date: Sun, 24 Oct 2021 09:48:37 -0500 Subject: [PATCH 05/14] update unflatten usage --- pytential/qbx/__init__.py | 15 +++++---------- pytential/source.py | 2 +- pytential/symbolic/execution.py | 4 +++- pytential/unregularized.py | 4 ++-- requirements.txt | 4 ++-- 5 files changed, 13 insertions(+), 16 deletions(-) diff --git a/pytential/qbx/__init__.py b/pytential/qbx/__init__.py index 5d5c944d4..bd377157e 100644 --- a/pytential/qbx/__init__.py +++ b/pytential/qbx/__init__.py @@ -682,11 +682,8 @@ def exec_compute_potential_insn_fmm(self, actx: PyOpenCLArrayContext, from meshmode.discretization import Discretization if isinstance(target_discr, Discretization): - template_ary = actx.np.astype( - thaw(target_discr.nodes()[0], actx), - result.dtype) - - result = unflatten(template_ary, result, actx) + template_ary = thaw(target_discr.nodes()[0], actx) + result = unflatten(template_ary, result, actx, strict=False) results.append((o.name, result)) @@ -874,10 +871,8 @@ def _flat_centers(dofdesc, qbx_forced_limit): for i, o in outputs: result = output_for_each_kernel[o.target_kernel_index] if isinstance(target_discr, Discretization): - template_ary = actx.np.astype( - thaw(target_discr.nodes()[0], actx), - result.dtype) - result = unflatten(template_ary, result, actx) + template_ary = thaw(target_discr.nodes()[0], actx) + result = unflatten(template_ary, result, actx, strict=False) results[i] = (o.name, result) @@ -954,7 +949,7 @@ def _flat_centers(dofdesc, qbx_forced_limit): result = output_for_each_kernel[o.target_kernel_index] if isinstance(target_discr, Discretization): template_ary = thaw(target_discr.nodes()[0], actx) - result = unflatten(template_ary, result, actx) + result = unflatten(template_ary, result, actx, strict=False) results[i] = (o.name, result) diff --git a/pytential/source.py b/pytential/source.py index ff284867a..ce4b28d81 100644 --- a/pytential/source.py +++ b/pytential/source.py @@ -205,7 +205,7 @@ def exec_compute_potential_insn(self, actx, insn, bound_expr, evaluate, result = output_for_each_kernel[o.target_kernel_index] if isinstance(target_discr, Discretization): template_ary = thaw(target_discr.nodes()[0], actx) - result = unflatten(template_ary, result, actx) + result = unflatten(template_ary, result, actx, strict=False) results.append((o.name, result)) diff --git a/pytential/symbolic/execution.py b/pytential/symbolic/execution.py index 5e5ae8651..df4bf1775 100644 --- a/pytential/symbolic/execution.py +++ b/pytential/symbolic/execution.py @@ -487,7 +487,9 @@ def unflatten(self, ary): if isinstance(discr, Discretization): from arraycontext import unflatten template_ary = thaw(discr.nodes()[0], self.array_context) - component = unflatten(template_ary, component, self.array_context) + component = unflatten( + template_ary, component, self.array_context, + strict=False) components.append(component) diff --git a/pytential/unregularized.py b/pytential/unregularized.py index 238176b28..276e858ac 100644 --- a/pytential/unregularized.py +++ b/pytential/unregularized.py @@ -179,7 +179,7 @@ def exec_compute_potential_insn_direct(self, actx: PyOpenCLArrayContext, result = output_for_each_kernel[o.target_kernel_index] if isinstance(target_discr, Discretization): template_ary = thaw(target_discr.nodes()[0], actx) - result = unflatten(template_ary, result, actx) + result = unflatten(template_ary, result, actx, strict=False) results.append((o.name, result)) @@ -293,7 +293,7 @@ def exec_compute_potential_insn_fmm(self, actx: PyOpenCLArrayContext, from meshmode.discretization import Discretization if isinstance(target_discr, Discretization): template_ary = thaw(target_discr.nodes()[0], actx) - result = unflatten(template_ary, result, actx) + result = unflatten(template_ary, result, actx, strict=False) results.append((o.name, result)) diff --git a/requirements.txt b/requirements.txt index 9bd896d9b..1532b3cc3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -9,7 +9,7 @@ git+https://github.com/inducer/pyopencl.git#egg=pyopencl git+https://github.com/inducer/islpy.git#egg=islpy git+https://github.com/inducer/loopy.git#egg=loopy git+https://github.com/inducer/boxtree.git#egg=boxtree -git+https://github.com/alexfikl/arraycontext.git@add-astype#egg=arraycontext +git+https://github.com/inducer/arraycontext.git#egg=arraycontext git+https://github.com/inducer/meshmode.git#egg=meshmode -git+https://github.com/alexfikl/sumpy.git@dir-auto-offset#egg=sumpy +git+https://github.com/inducer/sumpy.git#egg=sumpy git+https://github.com/inducer/pyfmmlib.git#egg=pyfmmlib From 5ad5c8f08fbfda6d1b31294b77d1b2110ae39b60 Mon Sep 17 00:00:00 2001 From: Alexandru Fikl Date: Sun, 24 Oct 2021 10:00:08 -0500 Subject: [PATCH 06/14] port examples away from meshmode.dof_array.flatten --- examples/laplace-dirichlet-3d.py | 18 ++++++------------ examples/layerpot-3d.py | 6 +++--- examples/layerpot.py | 16 +++++++++------- 3 files changed, 18 insertions(+), 22 deletions(-) diff --git a/examples/laplace-dirichlet-3d.py b/examples/laplace-dirichlet-3d.py index c499cac4c..fac418165 100644 --- a/examples/laplace-dirichlet-3d.py +++ b/examples/laplace-dirichlet-3d.py @@ -112,23 +112,17 @@ def main(mesh_name="torus", visualize=False): # {{{ fix rhs and solve - from meshmode.dof_array import flatten, unflatten + from arraycontext import flatten, unflatten nodes = thaw(density_discr.nodes(), actx) source = np.array([rout, 0, 0]) def u_incoming_func(x): - from pytools.obj_array import obj_array_vectorize - x = obj_array_vectorize(actx.to_numpy, flatten(x)) - x = np.array(list(x)) - # return 1/cl.clmath.sqrt( (x[0] - source[0])**2 - # +(x[1] - source[1])**2 - # +(x[2] - source[2])**2 ) - return 1.0/la.norm(x - source[:, None], axis=0) - - bc = unflatten(actx, - density_discr, - actx.from_numpy(u_incoming_func(nodes))) + x = actx.to_numpy(flatten(x, actx)).reshape(x.size, -1) + dists = 1.0/la.norm(x - source[:, None], axis=0) + return unflatten(x[0], actx.from_numpy(dists), actx) + + bc = u_incoming_func(nodes) bvp_rhs = bind(places, sqrt_w*sym.var("bc"))(actx, bc=bc) from pytential.solve import gmres diff --git a/examples/layerpot-3d.py b/examples/layerpot-3d.py index 255a58d0e..702ecccdb 100644 --- a/examples/layerpot-3d.py +++ b/examples/layerpot-3d.py @@ -86,12 +86,12 @@ def main(mesh_name="ellipsoid"): sigma = actx.np.cos(mode_nr*angle) if 0: - from meshmode.dof_array import flatten, unflatten - sigma = flatten(0 * angle) + from arraycontext import flatten, unflatten + sigma = flatten(0 * angle, actx) from random import randrange for _ in range(5): sigma[randrange(len(sigma))] = 1 - sigma = unflatten(actx, density_discr, sigma) + sigma = unflatten(angle, sigma, actx) if isinstance(kernel, HelmholtzKernel): for i, elem in np.ndenumerate(sigma): diff --git a/examples/layerpot.py b/examples/layerpot.py index 3f5c458bc..f190b5389 100644 --- a/examples/layerpot.py +++ b/examples/layerpot.py @@ -82,12 +82,12 @@ def op(**kwargs): sigma = actx.np.cos(mode_nr*angle) if 0: - from meshmode.dof_array import flatten, unflatten - sigma = flatten(0 * angle) + from arraycontext import flatten, unflatten + sigma = flatten(0 * angle, actx) from random import randrange for _ in range(5): sigma[randrange(len(sigma))] = 1 - sigma = unflatten(actx, density_discr, sigma) + sigma = unflatten(angle, sigma, actx) if isinstance(kernel, HelmholtzKernel): for i, elem in np.ndenumerate(sigma): @@ -121,11 +121,13 @@ def op(**kwargs): if enable_mayavi: # {{{ plot boundary field - from meshmode.dof_array import flatten_to_numpy + from arraycontext import flatten - fld_on_bdry = flatten_to_numpy( - actx, bound_bdry_op(actx, sigma=sigma, k=k)) - nodes_host = flatten_to_numpy(actx, density_discr.nodes()) + fld_on_bdry = actx.to_numpy( + flatten(bound_bdry_op(actx, sigma=sigma, k=k), actx)) + nodes_host = actx.to_numpy( + flatten(density_discr.nodes(), actx) + ).reshape(density_discr.ambient_dim, -1) mlab.points3d(nodes_host[0], nodes_host[1], fld_on_bdry.real, scale_factor=0.03) From 4702b8d63cdf94333b014f44ea5a5a40f6d96162 Mon Sep 17 00:00:00 2001 From: Alexandru Fikl Date: Sun, 24 Oct 2021 10:51:58 -0500 Subject: [PATCH 07/14] port tests away from meshmode.dof_array.flatten --- test/test_cost_model.py | 4 +- test/test_global_qbx.py | 102 ++++++++++++++--------------- test/test_layer_pot.py | 5 +- test/test_layer_pot_eigenvalues.py | 35 +++++----- test/test_layer_pot_identity.py | 19 +++--- test/test_linalg_proxy.py | 40 ++++++----- test/test_matrix.py | 29 ++++---- test/test_scalar_int_eq.py | 24 ++++--- test/test_stokes.py | 10 +-- test/test_symbolic.py | 33 +++++----- test/test_target_specific_qbx.py | 11 ++-- 11 files changed, 158 insertions(+), 154 deletions(-) diff --git a/test/test_cost_model.py b/test/test_cost_model.py index f5b8c2a50..00cbae6fe 100644 --- a/test/test_cost_model.py +++ b/test/test_cost_model.py @@ -1,6 +1,6 @@ __copyright__ = """ - Copyright (C) 2018 Matt Wala - Copyright (C) 2019 Hao Gao +Copyright (C) 2018 Matt Wala +Copyright (C) 2019 Hao Gao """ __license__ = """ diff --git a/test/test_global_qbx.py b/test/test_global_qbx.py index 65c79debe..cc0dfd5cc 100644 --- a/test/test_global_qbx.py +++ b/test/test_global_qbx.py @@ -30,7 +30,7 @@ import numpy as np import numpy.linalg as la -from arraycontext import thaw +from arraycontext import flatten from pytential import GeometryCollection, bind, sym from pytential.qbx import QBXLayerPotentialSource import meshmode.mesh.generation as mgen @@ -52,22 +52,6 @@ FAR_TARGET_DIST_FROM_SOURCE = 10 -# {{{ utils - -def dof_array_to_numpy(actx, ary): - """Converts DOFArrays (or object arrays of DOFArrays) to NumPy arrays. - Object arrays get turned into multidimensional arrays. - """ - from pytools.obj_array import obj_array_vectorize - from meshmode.dof_array import flatten - arr = obj_array_vectorize(actx.to_numpy, flatten(ary)) - if arr.dtype.char == "O": - arr = np.array(list(arr)) - return arr - -# }}} - - # {{{ source refinement checker @dataclass @@ -124,34 +108,40 @@ def run_source_refinement_test(actx_factory, mesh, order, # }}} dd = places.auto_source + ambient_dim = places.ambient_dim stage1_density_discr = places.get_discretization(dd.geometry) - stage1_density_nodes = dof_array_to_numpy(actx, - thaw(stage1_density_discr.nodes(), actx)) + stage1_density_nodes = actx.to_numpy( + flatten(stage1_density_discr.nodes(), actx) + ).reshape(ambient_dim, -1) quad_stage2_density_discr = places.get_discretization( dd.geometry, sym.QBX_SOURCE_QUAD_STAGE2) - quad_stage2_density_nodes = dof_array_to_numpy(actx, - thaw(quad_stage2_density_discr.nodes(), actx)) - - int_centers = dof_array_to_numpy(actx, - bind(places, - sym.expansion_centers(lpot_source.ambient_dim, -1))(actx)) - ext_centers = dof_array_to_numpy(actx, - bind(places, - sym.expansion_centers(lpot_source.ambient_dim, +1))(actx)) - expansion_radii = dof_array_to_numpy(actx, - bind(places, sym.expansion_radii(lpot_source.ambient_dim))(actx)) + quad_stage2_density_nodes = actx.to_numpy( + flatten(quad_stage2_density_discr.nodes(), actx) + ).reshape(ambient_dim, -1) + + int_centers = actx.to_numpy(flatten( + bind(places, sym.expansion_centers(ambient_dim, -1))(actx), actx) + ).reshape(ambient_dim, -1) + ext_centers = actx.to_numpy(flatten( + bind(places, sym.expansion_centers(ambient_dim, +1))(actx), actx) + ).reshape(ambient_dim, -1) + expansion_radii = actx.to_numpy(flatten( + bind(places, sym.expansion_radii(ambient_dim))(actx), actx) + ) dd = dd.copy(granularity=sym.GRANULARITY_ELEMENT) - source_danger_zone_radii = dof_array_to_numpy(actx, - bind(places, - sym._source_danger_zone_radii( - lpot_source.ambient_dim, dofdesc=dd.to_stage2()))(actx)) - quad_res = dof_array_to_numpy(actx, + source_danger_zone_radii = actx.to_numpy(flatten( + bind( + places, + sym._source_danger_zone_radii(ambient_dim, dofdesc=dd.to_stage2()) + )(actx), actx) + ) + quad_res = actx.to_numpy(flatten( bind(places, - sym._quad_resolution( - lpot_source.ambient_dim, dofdesc=dd))(actx)) + sym._quad_resolution(ambient_dim, dofdesc=dd))(actx), actx) + ) # {{{ check if satisfying criteria @@ -276,10 +266,13 @@ def test_target_association(actx_factory, curve_name, curve_f, nelements, from pyopencl.clrandom import PhiloxGenerator rng = PhiloxGenerator(actx.context, seed=RNG_SEED) + ambient_dim = places.ambient_dim dd = places.auto_source.to_stage1() - centers = dof_array_to_numpy(actx, - bind(places, sym.interleaved_expansion_centers( - lpot_source.ambient_dim, dofdesc=dd))(actx)) + + centers = actx.to_numpy(flatten( + bind(places, + sym.interleaved_expansion_centers(ambient_dim, dofdesc=dd))(actx), + actx)).reshape(ambient_dim, -1) density_discr = places.get_discretization(dd.geometry) @@ -288,15 +281,18 @@ def test_target_association(actx_factory, curve_name, curve_f, nelements, dtype=np.float64, a=0.01, b=1.0) ) - tunnel_radius = dof_array_to_numpy(actx, - bind(places, sym._close_target_tunnel_radii( - lpot_source.ambient_dim, dofdesc=dd))(actx)) + tunnel_radius = actx.to_numpy(flatten( + bind(places, sym._close_target_tunnel_radii(ambient_dim, dofdesc=dd))(actx), + actx)) def targets_from_sources(sign, dist, dim=2): - nodes = dof_array_to_numpy(actx, - bind(places, sym.nodes(dim, dofdesc=dd))(actx).as_vector(object)) - normals = dof_array_to_numpy(actx, - bind(places, sym.normal(dim, dofdesc=dd))(actx).as_vector(object)) + nodes = actx.to_numpy(flatten( + bind(places, sym.nodes(dim, dofdesc=dd))(actx).as_vector(), actx) + ).reshape(dim, -1) + normals = actx.to_numpy(flatten( + bind(places, sym.normal(dim, dofdesc=dd))(actx).as_vector(), actx) + ).reshape(dim, -1) + return actx.from_numpy(nodes + normals * sign * dist) from pytential.target import PointsTarget @@ -346,11 +342,13 @@ def targets_from_sources(sign, dist, dim=2): target_association_tolerance=1e-10) .get(queue=actx.queue)) - expansion_radii = dof_array_to_numpy(actx, - bind(places, sym.expansion_radii( - lpot_source.ambient_dim, - granularity=sym.GRANULARITY_CENTER))(actx)) - surf_targets = dof_array_to_numpy(actx, thaw(density_discr.nodes(), actx)) + expansion_radii = actx.to_numpy(flatten( + bind(places, sym.expansion_radii(ambient_dim, + granularity=sym.GRANULARITY_CENTER))(actx), actx) + ) + surf_targets = actx.to_numpy( + flatten(density_discr.nodes(), actx) + ).reshape(ambient_dim, -1) int_targets = actx.to_numpy(int_targets.nodes()) ext_targets = actx.to_numpy(ext_targets.nodes()) diff --git a/test/test_layer_pot.py b/test/test_layer_pot.py index bd91b1b08..61b9da8fa 100644 --- a/test/test_layer_pot.py +++ b/test/test_layer_pot.py @@ -25,7 +25,7 @@ import numpy as np -from arraycontext import thaw +from arraycontext import thaw, flatten from pytential import bind, sym, norm from pytential import GeometryCollection import meshmode.mesh.generation as mgen @@ -373,8 +373,7 @@ def test_unregularized_with_ones_kernel(actx_factory): auto_where=(places.auto_source, "target_non_self"))( actx, sigma=sigma) - from meshmode.dof_array import flatten - assert np.allclose(actx.to_numpy(flatten(result_self)), 2 * np.pi) + assert np.allclose(actx.to_numpy(flatten(result_self, actx)), 2 * np.pi) assert np.allclose(actx.to_numpy(result_nonself), 2 * np.pi) diff --git a/test/test_layer_pot_eigenvalues.py b/test/test_layer_pot_eigenvalues.py index 2ec593ce2..fb9d4797b 100644 --- a/test/test_layer_pot_eigenvalues.py +++ b/test/test_layer_pot_eigenvalues.py @@ -25,7 +25,7 @@ import numpy as np -from arraycontext import thaw +from arraycontext import thaw, flatten, unflatten from pytential import bind, sym, norm from pytential import GeometryCollection import meshmode.mesh.generation as mgen @@ -69,6 +69,7 @@ def test_ellipse_eigenvalues(actx_factory, ellipse_aspect, mode_nr, qbx_order, actx = actx_factory() + ambient_dim = 2 target_order = 8 from meshmode.discretization import Discretization @@ -113,7 +114,6 @@ def test_ellipse_eigenvalues(actx_factory, ellipse_aspect, mode_nr, qbx_order, places = GeometryCollection(qbx) density_discr = places.get_discretization(places.auto_source.geometry) - from meshmode.dof_array import flatten nodes = thaw(density_discr.nodes(), actx) if visualize: @@ -123,9 +123,12 @@ def test_ellipse_eigenvalues(actx_factory, ellipse_aspect, mode_nr, qbx_order, normals = bind(places, sym.normal(qbx.ambient_dim))(actx).as_vector(object) - nodes_h = np.array([actx.to_numpy(axis) for axis in flatten(nodes)]) - centers_h = np.array([actx.to_numpy(axis) for axis in flatten(centers)]) - normals_h = np.array([actx.to_numpy(axis) for axis in flatten(normals)]) + nodes_h = actx.to_numpy( + flatten(nodes, actx)).reshape(ambient_dim, -1) + centers_h = actx.to_numpy( + flatten(centers, actx)).reshape(ambient_dim, -1) + normals_h = actx.to_numpy( + flatten(normals, actx)).reshape(ambient_dim, -1) pt.plot(nodes_h[0], nodes_h[1], "x-") pt.plot(centers_h[0], centers_h[1], "o") @@ -160,9 +163,12 @@ def test_ellipse_eigenvalues(actx_factory, ellipse_aspect, mode_nr, qbx_order, s_sigma_ref = s_eigval*J*sigma if 0: - #pt.plot(s_sigma.get(), label="result") - #pt.plot(s_sigma_ref.get(), label="ref") - pt.plot(actx.to_numpy(flatten(s_sigma_ref - s_sigma)), label="err") + s_sigma_h = actx.to_numpy(flatten(s_sigma, actx)) + s_sigma_ref_h = actx.to_numpy(flatten(s_sigma_ref, actx)) + + # pt.plot(s_sigma_h, label="Result") + # pt.plot(s_sigma_ref_h, label="Reference") + pt.plot(s_sigma_ref_h - s_sigma_h, label="Error") pt.legend() pt.show() @@ -189,8 +195,8 @@ def test_ellipse_eigenvalues(actx_factory, ellipse_aspect, mode_nr, qbx_order, d_sigma_ref = d_eigval*sigma if 0: - pt.plot(actx.to_numpy(flatten(d_sigma)), label="result") - pt.plot(actx.to_numpy(flatten(d_sigma_ref)), label="ref") + pt.plot(actx.to_numpy(flatten(d_sigma, actx)), label="Result") + pt.plot(actx.to_numpy(flatten(d_sigma_ref, actx)), label="Reference") pt.legend() pt.show() @@ -293,20 +299,19 @@ def rel_err(comp, ref): ) places = GeometryCollection(qbx) - from meshmode.dof_array import flatten, unflatten - density_discr = places.get_discretization(places.auto_source.geometry) nodes = thaw(density_discr.nodes(), actx) r = actx.np.sqrt(nodes[0]*nodes[0] + nodes[1]*nodes[1] + nodes[2]*nodes[2]) phi = actx.np.arccos(nodes[2]/r) theta = actx.np.arctan2(nodes[0], nodes[1]) - ymn = unflatten(actx, density_discr, + ymn = unflatten(theta, actx.from_numpy( special.sph_harm( mode_m, mode_n, - actx.to_numpy(flatten(theta)), - actx.to_numpy(flatten(phi))))) + actx.to_numpy(flatten(theta, actx)), + actx.to_numpy(flatten(phi, actx)))), + actx) from sumpy.kernel import LaplaceKernel lap_knl = LaplaceKernel(3) diff --git a/test/test_layer_pot_identity.py b/test/test_layer_pot_identity.py index e30d2a5a9..82e7c6643 100644 --- a/test/test_layer_pot_identity.py +++ b/test/test_layer_pot_identity.py @@ -25,7 +25,7 @@ import numpy as np import numpy.linalg as la -from arraycontext import thaw +from arraycontext import flatten, unflatten from pytential import bind, sym, norm from pytential import GeometryCollection import meshmode.mesh.generation as mgen @@ -328,12 +328,13 @@ def test_identity_convergence(actx_factory, case, visualize=False): # {{{ compute values of a solution to the PDE density_discr = places.get_discretization(places.auto_source.geometry) + ambient_dim = places.ambient_dim - from meshmode.dof_array import flatten, unflatten - nodes_host = [actx.to_numpy(axis) - for axis in flatten(thaw(density_discr.nodes(), actx))] + nodes_host = actx.to_numpy( + flatten(density_discr.nodes(), actx) + ).reshape(ambient_dim, -1) normal = bind(places, sym.normal(d))(actx).as_vector(object) - normal_host = [actx.to_numpy(axis)for axis in flatten(normal)] + normal_host = actx.to_numpy(flatten(normal, actx)).reshape(ambient_dim, -1) if k != 0: if d == 2: @@ -369,11 +370,9 @@ def test_identity_convergence(actx_factory, case, visualize=False): # }}} - u_dev = unflatten(actx, density_discr, actx.from_numpy(u)) - dn_u_dev = unflatten(actx, density_discr, actx.from_numpy(dn_u)) - from pytools.obj_array import make_obj_array, obj_array_vectorize - grad_u_dev = unflatten(actx, density_discr, - obj_array_vectorize(actx.from_numpy, make_obj_array(grad_u))) + u_dev = unflatten(normal[0], actx.from_numpy(u), actx) + dn_u_dev = unflatten(normal[0], actx.from_numpy(dn_u), actx) + grad_u_dev = unflatten(normal, actx.from_numpy(grad_u.ravel()), actx) key = (case.qbx_order, case.geometry.mesh_name, resolution, case.expr.zero_op_name) diff --git a/test/test_linalg_proxy.py b/test/test_linalg_proxy.py index 425400923..0ffb827f7 100644 --- a/test/test_linalg_proxy.py +++ b/test/test_linalg_proxy.py @@ -26,6 +26,7 @@ import numpy as np import numpy.linalg as la +from arraycontext import thaw, flatten, unflatten from pytential import bind, sym from pytential import GeometryCollection from pytential.linalg import ProxyGenerator, QBXProxyGenerator @@ -64,8 +65,9 @@ def plot_proxy_geometry( pt.clf() if ambient_dim == 2: - from meshmode.dof_array import flatten_to_numpy - sources = np.stack(flatten_to_numpy(actx, discr.nodes())) + sources = actx.to_numpy( + flatten(discr.nodes(), actx) + ).reshape(ambient_dim, -1) if pxy is not None: proxies = np.stack(pxy.points) @@ -73,15 +75,15 @@ def plot_proxy_geometry( pxyranges = pxy.indices.ranges if with_qbx_centers: - ci = np.stack(flatten_to_numpy(actx, - bind(places, sym.expansion_centers(ambient_dim, -1))(actx) - )) - ce = np.stack(flatten_to_numpy(actx, - bind(places, sym.expansion_centers(ambient_dim, +1))(actx) - )) - r = flatten_to_numpy(actx, - bind(places, sym.expansion_radii(ambient_dim))(actx) - ) + ci = actx.to_numpy(flatten( + bind(places, sym.expansion_centers(ambient_dim, -1))(actx), + actx)).reshape(ambient_dim, -1) + ce = actx.to_numpy(flatten( + bind(places, sym.expansion_centers(ambient_dim, +1))(actx), + actx)).reshape(ambient_dim, -1) + r = actx.to_numpy(flatten( + bind(places, sym.expansion_radii(ambient_dim))(actx), + actx)) fig = pt.figure(figsize=(10, 8), dpi=300) if indices.indices.shape[0] != discr.ndofs: @@ -120,8 +122,8 @@ def plot_proxy_geometry( isrc = indices.block_indices(i) marker[isrc] = 10.0 * (i + 1.0) - from meshmode.dof_array import unflatten_from_numpy - marker_dev = unflatten_from_numpy(actx, discr, marker) + template_ary = thaw(discr.nodes()[0], actx) + marker_dev = unflatten(template_ary, actx.from_numpy(marker), actx) vis = make_visualizer(actx, discr) vis.write_vtk_file(f"test_proxy_geometry_{suffix}.vtu", [ @@ -137,7 +139,7 @@ def plot_proxy_geometry( marker[indices.indices] = 0.0 marker[isrc] = -42.0 marker[inbr] = +42.0 - marker_dev = unflatten_from_numpy(actx, discr, marker) + marker_dev = unflatten(template_ary, actx.from_numpy(marker), actx) vis.write_vtk_file( f"test_proxy_geometry_{suffix}_neighbor_{i:04d}.vtu", @@ -264,10 +266,11 @@ def test_proxy_generator(actx_factory, case, radius_factor=case.proxy_radius_factor) pxy = generator(actx, places.auto_source, srcindices).to_numpy(actx) - from meshmode.dof_array import flatten_to_numpy pxypoints = np.stack(pxy.points) pxycenters = np.stack(pxy.centers) - sources = np.stack(flatten_to_numpy(actx, density_discr.nodes())) + sources = actx.to_numpy( + flatten(density_discr.nodes(), actx) + ).reshape(places.ambient_dim, -1) for i in range(srcindices.nblocks): isrc = pxy.srcindices.block_indices(i) @@ -337,10 +340,11 @@ def test_neighbor_points(actx_factory, case, from pytential.linalg import gather_block_neighbor_points nbrindices = gather_block_neighbor_points(actx, density_discr, pxy) - from meshmode.dof_array import flatten_to_numpy pxy = pxy.to_numpy(actx) pxycenters = np.stack(pxy.centers) - nodes = np.vstack(flatten_to_numpy(actx, density_discr.nodes())) + nodes = actx.to_numpy( + flatten(density_discr.nodes(), actx) + ).reshape(places.ambient_dim, -1) for i in range(srcindices.nblocks): isrc = pxy.srcindices.block_indices(i) diff --git a/test/test_matrix.py b/test/test_matrix.py index 1481e57af..37e17f3c0 100644 --- a/test/test_matrix.py +++ b/test/test_matrix.py @@ -29,6 +29,7 @@ import numpy as np import numpy.linalg as la +from arraycontext import thaw, flatten, unflatten from pytential import bind, sym from pytential import GeometryCollection from pytools.obj_array import make_obj_array @@ -153,23 +154,23 @@ def test_build_matrix(actx_factory, k, curve_fn, op_type, visualize=False): # {{{ check - from meshmode.dof_array import unflatten_from_numpy, flatten_to_numpy - np.random.seed(12) + template_ary = thaw(density_discr.nodes()[0], actx) + for i in range(5): if isinstance(sym_u, np.ndarray): - u = make_obj_array([ - np.random.randn(density_discr.ndofs) - for _ in range(len(sym_u)) + u = np.random.randn(len(sym_u), density_discr.ndofs) + u_dev = make_obj_array([ + unflatten(template_ary, ui, actx) for ui in u ]) else: u = np.random.randn(density_discr.ndofs) - u_dev = unflatten_from_numpy(actx, density_discr, u) + u_dev = unflatten(template_ary, u, actx) - res_matvec = np.hstack(flatten_to_numpy(actx, - bound_op(actx, u=u_dev, **case.knl_concrete_kwargs) - )) - res_mat = mat.dot(np.hstack(u)) + res_matvec = actx.to_numpy(flatten( + bound_op(actx, u=u_dev, **case.knl_concrete_kwargs), + actx)) + res_mat = mat @ u abs_err = la.norm(res_mat - res_matvec, np.inf) rel_err = abs_err / la.norm(res_matvec, np.inf) @@ -252,10 +253,10 @@ def test_build_matrix_conditioning(actx_factory, side, op_type, visualize=False) if side == +1 and op_type == "double": # NOTE: this adds the "mean" to remove the nullspace for the operator # See `pytential.symbolic.pde.scalar` for the equivalent formulation - from meshmode.dof_array import flatten_to_numpy - w = flatten_to_numpy(actx, - bind(places, sym.sqrt_jac_q_weight(places.ambient_dim)**2)(actx) - ) + w = actx.to_numpy(flatten( + bind(places, sym.sqrt_jac_q_weight(places.ambient_dim)**2)(actx), + actx)) + w = np.tile(w.reshape(-1, 1), w.size).T kappa = la.cond(mat + w) diff --git a/test/test_scalar_int_eq.py b/test/test_scalar_int_eq.py index 8ad16eff5..d3085306d 100644 --- a/test/test_scalar_int_eq.py +++ b/test/test_scalar_int_eq.py @@ -25,8 +25,8 @@ import numpy as np import numpy.linalg as la +from arraycontext import flatten from meshmode.discretization.visualization import make_visualizer -from meshmode.dof_array import flatten_to_numpy from sumpy.kernel import LaplaceKernel, HelmholtzKernel, BiharmonicKernel @@ -148,8 +148,12 @@ def run_int_eq_test(actx, # show geometry, centers, normals if ambient_dim == 2: - nodes = flatten_to_numpy(actx, density_discr.nodes()) - normals = flatten_to_numpy(actx, normals) + nodes = actx.to_numpy( + flatten(density_discr.nodes(), actx) + ).reshape(ambient_dim, -1) + normals = actx.to_numpy( + flatten(normals, actx) + ).reshape(ambient_dim, -1) pt.plot(nodes[0], nodes[1], "x-") pt.quiver(nodes[0], nodes[1], normals[0], normals[1]) @@ -287,9 +291,9 @@ def run_int_eq_test(actx, err = test_via_bdry - test_direct - err = flatten_to_numpy(actx, err, strict=False) - test_direct = flatten_to_numpy(actx, test_direct, strict=False) - test_via_bdry = flatten_to_numpy(actx, test_via_bdry, strict=False) + err = actx.to_numpy(flatten(err, actx)) + test_direct = actx.to_numpy(flatten(test_direct, actx)) + test_via_bdry = actx.to_numpy(flatten(test_via_bdry, actx)) # {{{ remove effect of net source charge @@ -330,8 +334,8 @@ def run_int_eq_test(actx, actx, charges=source_charges_dev, **case.knl_concrete_kwargs) grad_err = grad_from_src - grad_ref - grad_ref = flatten_to_numpy(actx, grad_ref[0]) - grad_err = flatten_to_numpy(actx, grad_err[0]) + grad_ref = actx.to_numpy(flatten(grad_ref[0], actx)) + grad_err = actx.to_numpy(flatten(grad_err[0], actx)) rel_grad_err_inf = la.norm(grad_err, np.inf) / la.norm(grad_ref, np.inf) logger.info("rel_grad_err_inf: %.5e", rel_grad_err_inf) @@ -356,8 +360,8 @@ def run_int_eq_test(actx, auto_where=("point_source", case.name))( actx, charges=source_charges_dev, **case.knl_concrete_kwargs) - tang_deriv_from_src = flatten_to_numpy(actx, tang_deriv_from_src) - tang_deriv_ref = flatten_to_numpy(actx, tang_deriv_ref) + tang_deriv_from_src = actx.to_numpy(flatten(tang_deriv_from_src, actx)) + tang_deriv_ref = actx.to_numpy(flatten(tang_deriv_ref, actx)) td_err = tang_deriv_from_src - tang_deriv_ref if visualize: diff --git a/test/test_stokes.py b/test/test_stokes.py index 3cae8f6f0..c6b5aed86 100644 --- a/test/test_stokes.py +++ b/test/test_stokes.py @@ -25,6 +25,7 @@ import numpy as np +from arraycontext import flatten from pytential import GeometryCollection, bind, sym from meshmode.discretization import Discretization from meshmode.discretization.poly_element import \ @@ -345,14 +346,7 @@ def run_stokes_identity(actx_factory, case, identity, resolution, visualize=Fals type(identity).__name__.lower(), places.ambient_dim, resolution) if places.ambient_dim == 2: - from meshmode.dof_array import flatten_to_numpy - result = flatten_to_numpy(actx, result) - if not isinstance(ref_result[0], (int, float)): - ref_result = flatten_to_numpy(actx, ref_result) - else: - ref_result = [ - c * np.ones_like(r) - for c, r in zip(ref_result, result)] + result = actx.to_numpy(flatten(result, actx)) import matplotlib.pyplot as plt fig = plt.figure() diff --git a/test/test_symbolic.py b/test/test_symbolic.py index 6df55498e..40947c828 100644 --- a/test/test_symbolic.py +++ b/test/test_symbolic.py @@ -26,11 +26,7 @@ import numpy as np import numpy.linalg as la -import pyopencl as cl -import pyopencl.array -import pyopencl.clmath - -from arraycontext import thaw +from arraycontext import thaw, flatten, unflatten import meshmode.mesh.generation as mgen from meshmode.discretization import Discretization from meshmode.discretization.poly_element import \ @@ -155,10 +151,10 @@ def test_tangential_onb(actx_factory): for i in range(nvecs) for j in range(nvecs)]) )(actx) - from meshmode.dof_array import flatten - orth_check = flatten(orth_check) for orth_i in orth_check: - assert (cl.clmath.fabs(orth_i) < 1e-13).get().all() + assert actx.to_numpy( + actx.np.all(actx.np.abs(orth_i) < 1e-13) + ) # make sure tangential_onb is orthogonal to normal orth_check = bind(discr, sym.make_obj_array([ @@ -166,9 +162,10 @@ def test_tangential_onb(actx_factory): for i in range(nvecs)]) )(actx) - orth_check = flatten(orth_check) for orth_i in orth_check: - assert (cl.clmath.fabs(orth_i) < 1e-13).get().all() + assert actx.to_numpy( + actx.np.all(actx.np.abs(orth_i) < 1e-13) + ) # }}} @@ -261,22 +258,22 @@ def test_interpolation(actx_factory, name, source_discr_stage, target_granularit op_sym = sym.sin(sym.interp(from_dd, to_dd, sigma_sym)) bound_op = bind(places, op_sym, auto_where=where) - from meshmode.dof_array import flatten, unflatten - def discr_and_nodes(stage): density_discr = places.get_discretization(where.geometry, stage) - return density_discr, np.array([ - actx.to_numpy(flatten(axis)) - for axis in thaw(density_discr.nodes(), actx)]) + return density_discr, actx.to_numpy( + flatten(density_discr.nodes(), actx) + ).reshape(density_discr.ambient_dim, -1) _, target_nodes = discr_and_nodes(sym.QBX_SOURCE_QUAD_STAGE2) source_discr, source_nodes = discr_and_nodes(source_discr_stage) sigma_target = np.sin(la.norm(target_nodes, axis=0)) sigma_dev = unflatten( - actx, source_discr, - actx.from_numpy(la.norm(source_nodes, axis=0))) - sigma_target_interp = actx.to_numpy(flatten(bound_op(actx, sigma=sigma_dev))) + thaw(source_discr.nodes()[0], actx), + actx.from_numpy(la.norm(source_nodes, axis=0)), actx) + sigma_target_interp = actx.to_numpy( + flatten(bound_op(actx, sigma=sigma_dev), actx) + ) if name in ("default", "default_explicit", "stage2", "quad"): error = la.norm(sigma_target_interp - sigma_target) / la.norm(sigma_target) diff --git a/test/test_target_specific_qbx.py b/test/test_target_specific_qbx.py index 0664a0b6c..09f8facea 100644 --- a/test/test_target_specific_qbx.py +++ b/test/test_target_specific_qbx.py @@ -24,7 +24,7 @@ import numpy as np -from arraycontext import thaw +from arraycontext import thaw, flatten from pytential import GeometryCollection, bind, sym from sumpy.kernel import LaplaceKernel, HelmholtzKernel @@ -194,12 +194,15 @@ def test_target_specific_qbx(actx_factory, op, helmholtz_k, qbx_order): expr = op(kernel, u_sym, qbx_forced_limit=-1, **kernel_kwargs) - from meshmode.dof_array import flatten bound_op = bind(places, expr) - pot_ref = actx.to_numpy(flatten(bound_op(actx, u=u_dev, k=helmholtz_k))) + pot_ref = actx.to_numpy( + flatten(bound_op(actx, u=u_dev, k=helmholtz_k), actx) + ) bound_op = bind(places, expr, auto_where="qbx_target_specific") - pot_tsqbx = actx.to_numpy(flatten(bound_op(actx, u=u_dev, k=helmholtz_k))) + pot_tsqbx = actx.to_numpy( + flatten(bound_op(actx, u=u_dev, k=helmholtz_k), actx) + ) assert np.allclose(pot_tsqbx, pot_ref, atol=1e-13, rtol=1e-13) From 9c76c296273797b5f903267bf636c40aee66cae4 Mon Sep 17 00:00:00 2001 From: Alexandru Fikl Date: Sun, 24 Oct 2021 12:46:18 -0500 Subject: [PATCH 08/14] fix some failing examples --- examples/laplace-dirichlet-3d.py | 4 ++-- pytential/qbx/direct.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/laplace-dirichlet-3d.py b/examples/laplace-dirichlet-3d.py index fac418165..3ea9e36b0 100644 --- a/examples/laplace-dirichlet-3d.py +++ b/examples/laplace-dirichlet-3d.py @@ -117,8 +117,8 @@ def main(mesh_name="torus", visualize=False): source = np.array([rout, 0, 0]) def u_incoming_func(x): - x = actx.to_numpy(flatten(x, actx)).reshape(x.size, -1) - dists = 1.0/la.norm(x - source[:, None], axis=0) + x_flat = actx.to_numpy(flatten(x, actx)).reshape(x.size, -1) + dists = 1.0/la.norm(x_flat - source[:, None], axis=0) return unflatten(x[0], actx.from_numpy(dists), actx) diff --git a/pytential/qbx/direct.py b/pytential/qbx/direct.py index bbb0f7aa1..4c7b037c0 100644 --- a/pytential/qbx/direct.py +++ b/pytential/qbx/direct.py @@ -50,7 +50,7 @@ def get_kernel(self): lp.GlobalArg("targets", None, shape=(self.dim, "ntargets_total"), order="C"), lp.GlobalArg("center", None, - shape=(self.dim, "ncenters_total"), dim_tags="sep,C"), + shape=(self.dim, "ncenters_total")), lp.GlobalArg("expansion_radii", None, shape="ncenters_total"), lp.GlobalArg("qbx_tgt_numbers", None, From a65c5100b6237167fc8a630f63180d53f4908728 Mon Sep 17 00:00:00 2001 From: Alexandru Fikl Date: Sun, 24 Oct 2021 14:55:41 -0500 Subject: [PATCH 09/14] fix bug introduced in reordering sources --- pytential/source.py | 10 ++++++---- test/test_layer_pot.py | 7 ++++--- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/pytential/source.py b/pytential/source.py index ce4b28d81..802566aeb 100644 --- a/pytential/source.py +++ b/pytential/source.py @@ -307,7 +307,9 @@ def get_fmm_expansion_wrangler_extra_kwargs( def reorder_sources(source_array): if isinstance(source_array, cl.array.Array): - return source_array[tree_user_source_ids].with_queue(None) + return actx.freeze( + actx.thaw(source_array)[tree_user_source_ids] + ) else: return source_array @@ -325,9 +327,9 @@ def reorder_sources(source_array): for arg in func(target_kernels): value = evaluator(arguments[arg.name]) if isinstance(value, np.ndarray): - value = map_array_container( - lambda ary: reorder_sources(flatten(value, actx)), - value) + value = actx.np.stack(map_array_container( + lambda ary: reorder_sources(flatten(ary, actx)), + value)) elif isinstance(value, DOFArray): value = reorder_sources(flatten(value, actx)) else: diff --git a/test/test_layer_pot.py b/test/test_layer_pot.py index 61b9da8fa..ed1b7b70c 100644 --- a/test/test_layer_pot.py +++ b/test/test_layer_pot.py @@ -126,9 +126,10 @@ def test_off_surface_eval(actx_factory, use_fmm, visualize=False): fld_in_vol = bind(places, op)(actx, sigma=sigma) fld_in_vol_exact = -1 - err = actx.np.fabs(fld_in_vol - fld_in_vol_exact) - linf_err = actx.to_numpy(err).max() - print("l_inf error:", linf_err) + linf_err = actx.to_numpy( + actx.np.linalg.norm(fld_in_vol - fld_in_vol_exact, ord=np.inf) + ) + logger.info("l_inf error: %.12e", linf_err) if visualize: fplot.show_scalar_in_matplotlib(actx.to_numpy(fld_in_vol)) From 4653891cd56cd1e704b85f86f6be77f0a9994c37 Mon Sep 17 00:00:00 2001 From: Alexandru Fikl Date: Sun, 24 Oct 2021 15:41:45 -0500 Subject: [PATCH 10/14] fix some more unflatten mismatches --- test/test_layer_pot_eigenvalues.py | 2 +- test/test_layer_pot_identity.py | 9 ++++++--- test/test_matrix.py | 7 ++++--- 3 files changed, 11 insertions(+), 7 deletions(-) diff --git a/test/test_layer_pot_eigenvalues.py b/test/test_layer_pot_eigenvalues.py index fb9d4797b..8ecf42c4a 100644 --- a/test/test_layer_pot_eigenvalues.py +++ b/test/test_layer_pot_eigenvalues.py @@ -311,7 +311,7 @@ def rel_err(comp, ref): mode_m, mode_n, actx.to_numpy(flatten(theta, actx)), actx.to_numpy(flatten(phi, actx)))), - actx) + actx, strict=False) from sumpy.kernel import LaplaceKernel lap_knl = LaplaceKernel(3) diff --git a/test/test_layer_pot_identity.py b/test/test_layer_pot_identity.py index 82e7c6643..5590ab3b2 100644 --- a/test/test_layer_pot_identity.py +++ b/test/test_layer_pot_identity.py @@ -370,9 +370,12 @@ def test_identity_convergence(actx_factory, case, visualize=False): # }}} - u_dev = unflatten(normal[0], actx.from_numpy(u), actx) - dn_u_dev = unflatten(normal[0], actx.from_numpy(dn_u), actx) - grad_u_dev = unflatten(normal, actx.from_numpy(grad_u.ravel()), actx) + u_dev = unflatten( + normal[0], actx.from_numpy(u), actx, strict=False) + dn_u_dev = unflatten( + normal[0], actx.from_numpy(dn_u), actx, strict=False) + grad_u_dev = unflatten( + normal, actx.from_numpy(grad_u.ravel()), actx, strict=False) key = (case.qbx_order, case.geometry.mesh_name, resolution, case.expr.zero_op_name) diff --git a/test/test_matrix.py b/test/test_matrix.py index 37e17f3c0..9d60b421a 100644 --- a/test/test_matrix.py +++ b/test/test_matrix.py @@ -161,16 +161,17 @@ def test_build_matrix(actx_factory, k, curve_fn, op_type, visualize=False): if isinstance(sym_u, np.ndarray): u = np.random.randn(len(sym_u), density_discr.ndofs) u_dev = make_obj_array([ - unflatten(template_ary, ui, actx) for ui in u + unflatten(template_ary, actx.from_numpy(ui), actx, strict=False) + for ui in u ]) else: u = np.random.randn(density_discr.ndofs) - u_dev = unflatten(template_ary, u, actx) + u_dev = unflatten(template_ary, actx.from_numpy(u), actx, strict=False) res_matvec = actx.to_numpy(flatten( bound_op(actx, u=u_dev, **case.knl_concrete_kwargs), actx)) - res_mat = mat @ u + res_mat = mat @ u.ravel() abs_err = la.norm(res_mat - res_matvec, np.inf) rel_err = abs_err / la.norm(res_matvec, np.inf) From cce29954a810104bed4292dd2c3c507730e5617a Mon Sep 17 00:00:00 2001 From: Alexandru Fikl Date: Wed, 12 Jan 2022 22:08:04 -0600 Subject: [PATCH 11/14] use the new leaf_class argument to flatten --- pytential/linalg/proxy.py | 19 ++++++----------- pytential/qbx/__init__.py | 9 +++----- pytential/qbx/geometry.py | 19 +++++++---------- pytential/qbx/utils.py | 16 ++++++-------- pytential/source.py | 31 ++++++--------------------- pytential/symbolic/matrix.py | 41 ++++++++++-------------------------- pytential/unregularized.py | 39 +++++++++++++++------------------- 7 files changed, 58 insertions(+), 116 deletions(-) diff --git a/pytential/linalg/proxy.py b/pytential/linalg/proxy.py index 601b6a145..a01f4d958 100644 --- a/pytential/linalg/proxy.py +++ b/pytential/linalg/proxy.py @@ -28,6 +28,7 @@ from arraycontext import PyOpenCLArrayContext, flatten from meshmode.discretization import Discretization +from meshmode.dof_array import DOFArray from pytools import memoize_in from pytential.linalg.utils import BlockIndexRanges @@ -77,9 +78,7 @@ def partition_by_nodes( builder = TreeBuilder(actx.context) tree, _ = builder(actx.queue, - particles=actx.np.reshape( - flatten(discr.nodes(), actx), - (discr.ambient_dim, -1)), + particles=flatten(discr.nodes(), actx, leaf_class=DOFArray), max_particles_in_box=max_particles_in_box, kind=tree_kind) @@ -329,9 +328,7 @@ def __call__(self, # {{{ get proxy centers and radii - sources = actx.np.reshape( - flatten(discr.nodes(), actx), - (self.ambient_dim, -1)) + sources = flatten(discr.nodes(), actx, leaf_class=DOFArray) knl = self.get_centers_knl(actx) _, (centers_dev,) = knl(actx.queue, @@ -507,10 +504,8 @@ def __call__(self, return super().__call__(actx, source_dd, indices, expansion_radii=flatten(radii, actx), - center_int=actx.np.reshape( - flatten(center_int, actx), (self.ambient_dim, -1)), - center_ext=actx.np.reshape( - flatten(center_ext, actx), (self.ambient_dim, -1)), + center_int=flatten(center_int, actx, leaf_class=DOFArray), + center_ext=flatten(center_ext, actx, leaf_class=DOFArray), **kwargs) # }}} @@ -556,9 +551,7 @@ def prg(): return knl _, (sources,) = prg()(actx.queue, - ary=actx.np.reshape( - flatten(discr.nodes(), actx), - (discr.ambient_dim, -1)), + ary=flatten(discr.nodes(), actx, leaf_class=DOFArray), srcindices=pxy.srcindices.indices) # }}} diff --git a/pytential/qbx/__init__.py b/pytential/qbx/__init__.py index bd377157e..a89848463 100644 --- a/pytential/qbx/__init__.py +++ b/pytential/qbx/__init__.py @@ -24,6 +24,7 @@ import pyopencl as cl from arraycontext import PyOpenCLArrayContext, thaw, flatten, unflatten +from meshmode.dof_array import DOFArray from pytools import memoize_method, memoize_in, single_valued from pytential.qbx.target_assoc import QBXTargetAssociationFailedException @@ -779,9 +780,7 @@ def exec_compute_potential_insn_direct(self, actx, insn, bound_expr, evaluate, def _flat_nodes(dofdesc): discr = bound_expr.places.get_discretization( dofdesc.geometry, dofdesc.discr_stage) - return actx.freeze(actx.np.reshape( - flatten(discr.nodes(), actx), (self.ambient_dim, -1) - )) + return actx.freeze(flatten(discr.nodes(), actx, leaf_class=DOFArray)) @memoize_in(bound_expr.places, (QBXLayerPotentialSource, "flat_expansion_radii")) @@ -799,9 +798,7 @@ def _flat_centers(dofdesc, qbx_forced_limit): sym.expansion_centers( self.ambient_dim, qbx_forced_limit, dofdesc=dofdesc), )(actx) - return actx.freeze(actx.np.reshape( - flatten(centers, actx), (self.ambient_dim, -1) - )) + return actx.freeze(flatten(centers, actx, leaf_class=DOFArray)) from pytential.source import evaluate_kernel_arguments flat_kernel_args = evaluate_kernel_arguments( diff --git a/pytential/qbx/geometry.py b/pytential/qbx/geometry.py index 1f482d425..72c247a18 100644 --- a/pytential/qbx/geometry.py +++ b/pytential/qbx/geometry.py @@ -27,6 +27,7 @@ from pytools import memoize_method, log_process from arraycontext import PyOpenCLArrayContext, flatten +from meshmode.dof_array import DOFArray from boxtree.tools import DeviceDataRecord from boxtree.pyfmmlib_integration import FMMLibRotationDataInterface @@ -435,9 +436,7 @@ def flat_centers(self): self.ambient_dim, dofdesc=self.source_dd.to_stage1()))(actx) - return actx.freeze(actx.np.reshape( - flatten(centers, actx), (centers.size, -1) - )) + return actx.freeze(flatten(centers, actx, leaf_class=DOFArray)) @memoize_method def flat_expansion_radii(self): @@ -488,10 +487,7 @@ def target_info(self): actx.queue, targets=targets[:, start:start+target_discr.ndofs], - points=actx.np.reshape( - flatten(target_discr.nodes(), actx), - (target_discr.ambient_dim, -1) - ) + points=flatten(target_discr.nodes(), actx, leaf_class=DOFArray), ) return TargetInfo( @@ -536,7 +532,8 @@ def tree(self): lpot_source = self.lpot_source target_info = self.target_info() - queue = self.array_context.queue + actx = self.array_context + queue = actx.queue from pytential import sym quad_stage2_discr = self.places.get_discretization( @@ -566,9 +563,9 @@ def tree(self): refine_weights.finish() tree, _ = code_getter.build_tree()(queue, - particles=self.array_context.np.reshape( - flatten(quad_stage2_discr.nodes(), self.array_context), - (quad_stage2_discr.ambient_dim, -1)), + particles=flatten( + quad_stage2_discr.nodes(), actx, leaf_class=DOFArray + ), targets=target_info.targets, target_radii=target_radii, max_leaf_refine_weight=lpot_source._max_leaf_refine_weight, diff --git a/pytential/qbx/utils.py b/pytential/qbx/utils.py index d427ac012..a0d260c1e 100644 --- a/pytential/qbx/utils.py +++ b/pytential/qbx/utils.py @@ -28,6 +28,7 @@ from pytools import memoize_method, log_process from arraycontext import PyOpenCLArrayContext +from meshmode.dof_array import DOFArray from boxtree.tree import Tree from boxtree.pyfmmlib_integration import FMMLibRotationDataInterface @@ -276,20 +277,15 @@ def _make_centers(discr): return bind(discr, sym.interleaved_expansion_centers( discr.ambient_dim))(actx) - def _flatten_nodes(ary): - if isinstance(ary, np.ndarray): - return actx.np.reshape(flatten(ary, actx), (ary.size, -1)) - else: - assert ary.shape[0] == density_discr.ambient_dim - return ary - stage1_density_discr = stage1_density_discrs[0] density_discr = density_discrs[0] from arraycontext import flatten - sources = _flatten_nodes(density_discr.nodes()) - centers = _flatten_nodes(_make_centers(stage1_density_discr)) - targets = [_flatten_nodes(tgt.nodes()) for tgt in targets_list] + sources = flatten(density_discr.nodes(), actx, leaf_class=DOFArray) + centers = flatten(_make_centers(stage1_density_discr), actx, leaf_class=DOFArray) + targets = [ + flatten(tgt.nodes(), actx, leaf_class=DOFArray) + for tgt in targets_list] queue = actx.queue particles = tuple( diff --git a/pytential/source.py b/pytential/source.py index 802566aeb..8bd5387ea 100644 --- a/pytential/source.py +++ b/pytential/source.py @@ -25,6 +25,7 @@ from pytools import memoize_in from arraycontext import thaw, flatten, unflatten +from meshmode.dof_array import DOFArray from sumpy.fmm import UnableToCollectTimingData @@ -85,20 +86,12 @@ def p2p(target_kernels, source_kernels): # {{{ point potential source def evaluate_kernel_arguments(actx, evaluate, kernel_arguments, flat=True): - from meshmode.dof_array import DOFArray - kernel_args = {} for arg_name, arg_expr in kernel_arguments.items(): value = evaluate(arg_expr) if flat: - if isinstance(value, np.ndarray): - value = actx.np.reshape(flatten(value, actx), (value.size, -1)) - elif isinstance(value, DOFArray): - value = flatten(value, actx) - else: - pass - + value = flatten(value, actx, leaf_class=DOFArray) kernel_args[arg_name] = value return kernel_args @@ -195,9 +188,7 @@ def exec_compute_potential_insn(self, actx, insn, bound_expr, evaluate, target_kernels=insn.target_kernels) evt, output_for_each_kernel = p2p(actx.queue, - targets=actx.np.reshape( - flatten(target_discr.nodes(), actx), - (target_discr.ambient_dim, -1)), + targets=flatten(target_discr.nodes(), actx, leaf_class=DOFArray), sources=self._nodes, strength=strengths, **kernel_args) @@ -318,24 +309,16 @@ def reorder_sources(source_array): from sumpy.tools import gather_arguments, gather_source_arguments from arraycontext import map_array_container - from meshmode.dof_array import DOFArray for func, var_dict in [ (gather_arguments, kernel_extra_kwargs), (gather_source_arguments, source_extra_kwargs), ]: for arg in func(target_kernels): - value = evaluator(arguments[arg.name]) - if isinstance(value, np.ndarray): - value = actx.np.stack(map_array_container( - lambda ary: reorder_sources(flatten(ary, actx)), - value)) - elif isinstance(value, DOFArray): - value = reorder_sources(flatten(value, actx)) - else: - pass - - var_dict[arg.name] = value + var_dict[arg.name] = map_array_container( + lambda ary: reorder_sources(flatten(ary, actx)), + evaluator(arguments[arg.name]), + leaf_class=DOFArray) return kernel_extra_kwargs, source_extra_kwargs diff --git a/pytential/symbolic/matrix.py b/pytential/symbolic/matrix.py index 5cb09b742..4ce1d5011 100644 --- a/pytential/symbolic/matrix.py +++ b/pytential/symbolic/matrix.py @@ -29,6 +29,7 @@ from pytools import memoize_method from arraycontext import thaw, flatten, unflatten +from meshmode.dof_array import DOFArray from pytential.symbolic.mappers import EvaluationMapperBase @@ -413,15 +414,9 @@ def map_int_g(self, expr): dofdesc=expr.target))(actx) _, (mat,) = mat_gen(actx.queue, - targets=actx.np.reshape( - flatten(target_discr.nodes(), actx), - (target_discr.ambient_dim, -1)), - sources=actx.np.reshape( - flatten(source_discr.nodes(), actx), - (source_discr.ambient_dim, -1)), - centers=actx.np.reshape( - flatten(centers, actx), - (target_discr.ambient_dim, -1)), + targets=flatten(target_discr.nodes(), actx, leaf_class=DOFArray), + sources=flatten(source_discr.nodes(), actx, leaf_class=DOFArray), + centers=flatten(centers, actx, leaf_class=DOFArray), expansion_radii=flatten(radii, actx), **kernel_args) mat = actx.to_numpy(mat) @@ -488,12 +483,8 @@ def map_int_g(self, expr): exclude_self=self.exclude_self) _, (mat,) = mat_gen(actx.queue, - targets=actx.np.reshape( - flatten(target_discr.nodes(), actx), - (target_discr.ambient_dim, -1)), - sources=actx.np.reshape( - flatten(source_discr.nodes(), actx), - (source_discr.ambient_dim, -1)), + targets=flatten(target_discr.nodes(), actx, leaf_class=DOFArray), + sources=flatten(source_discr.nodes(), actx, leaf_class=DOFArray), **kernel_args) mat = actx.to_numpy(mat) @@ -566,15 +557,9 @@ def map_int_g(self, expr): dofdesc=expr.target))(actx) _, (mat,) = mat_gen(actx.queue, - targets=actx.np.reshape( - flatten(target_discr.nodes(), actx), - (target_discr.ambient_dim, -1)), - sources=actx.np.reshape( - flatten(source_discr.nodes(), actx), - (source_discr.ambient_dim, -1)), - centers=actx.np.reshape( - flatten(centers, actx), - (target_discr.ambient_dim, -1)), + targets=flatten(target_discr.nodes(), actx, leaf_class=DOFArray), + sources=flatten(source_discr.nodes(), actx, leaf_class=DOFArray), + centers=flatten(centers, actx, leaf_class=DOFArray), expansion_radii=flatten(radii, actx), tgtindices=tgtindices, srcindices=srcindices, @@ -644,12 +629,8 @@ def map_int_g(self, expr): exclude_self=self.exclude_self) _, (mat,) = mat_gen(actx.queue, - targets=actx.np.reshape( - flatten(target_discr.nodes(), actx), - (target_discr.ambient_dim, -1)), - sources=actx.np.reshape( - flatten(source_discr.nodes(), actx), - (source_discr.ambient_dim, -1)), + targets=flatten(target_discr.nodes(), actx, leaf_class=DOFArray), + sources=flatten(source_discr.nodes(), actx, leaf_class=DOFArray), tgtindices=tgtindices, srcindices=srcindices, **kernel_args) diff --git a/pytential/unregularized.py b/pytential/unregularized.py index 276e858ac..ef714c89e 100644 --- a/pytential/unregularized.py +++ b/pytential/unregularized.py @@ -31,6 +31,7 @@ from pytools import memoize_method from arraycontext import PyOpenCLArrayContext, thaw, flatten, unflatten +from meshmode.dof_array import DOFArray from boxtree.tools import DeviceDataRecord from pytential.source import LayerPotentialSourceBase @@ -141,13 +142,9 @@ def exec_compute_potential_insn_direct(self, actx: PyOpenCLArrayContext, kernel_args = {} for arg_name, arg_expr in insn.kernel_arguments.items(): - value = evaluate(arg_expr) - if isinstance(value, np.ndarray): - value = actx.np.reshape(flatten(value, actx), (value.size, -1)) - else: - value = flatten(value, actx) - - kernel_args[arg_name] = value + kernel_args[arg_name] = flatten( + evaluate(arg_expr), actx, leaf_class=DOFArray + ) from pytential import bind, sym waa = bind(bound_expr.places, sym.weights_and_area_elements( @@ -167,12 +164,10 @@ def exec_compute_potential_insn_direct(self, actx: PyOpenCLArrayContext, target_kernels=insn.target_kernels) evt, output_for_each_kernel = p2p(actx.queue, - targets=actx.np.reshape( - flatten(target_discr.nodes(), actx), - (target_discr.ambient_dim, -1)), - sources=actx.np.reshape( - flatten(self.density_discr.nodes(), actx), - (self.density_discr.ambient_dim, -1)), + targets=flatten(target_discr.nodes(), actx, leaf_class=DOFArray), + sources=flatten( + self.density_discr.nodes(), actx, leaf_class=DOFArray + ), strength=flat_strengths, **kernel_args) from meshmode.discretization import Discretization @@ -423,9 +418,9 @@ def tree(self): MAX_LEAF_REFINE_WEIGHT = 32 # noqa tree, _ = code_getter.build_tree(actx.queue, - particles=actx.np.reshape( - flatten(lpot_src.density_discr.nodes(), actx), - (lpot_src.ambient_dim, -1)), + particles=flatten( + lpot_src.density_discr.nodes(), actx, leaf_class=DOFArray + ), targets=target_info.targets, max_leaf_refine_weight=MAX_LEAF_REFINE_WEIGHT, refine_weights=refine_weights, @@ -447,19 +442,19 @@ def target_info(self): target_discr_starts.append(ntargets) ntargets += target_discr.ndofs + actx = self.array_context target_discr_starts.append(ntargets) - targets = self.array_context.empty( + targets = actx.empty( (lpot_src.ambient_dim, ntargets), self.coord_dtype) for start, target_discr in zip(target_discr_starts, target_discrs): - code_getter.copy_targets_kernel()( - self.array_context.queue, + code_getter.copy_targets_kernel()(actx.queue, targets=targets[:, start:start+target_discr.ndofs], - points=self.array_context.np.reshape( - flatten(target_discr.nodes(), self.array_context), - (target_discr.ambient_dim, -1)) + points=flatten( + target_discr.nodes(), actx, leaf_class=DOFArray + ), ) return _TargetInfo( From c1a92a02d356ebcffaa59161e59a59cf1fe7cf71 Mon Sep 17 00:00:00 2001 From: Alexandru Fikl Date: Wed, 12 Jan 2022 22:08:36 -0600 Subject: [PATCH 12/14] point requirements to arraycontext branch --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 1532b3cc3..df76004e3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -9,7 +9,7 @@ git+https://github.com/inducer/pyopencl.git#egg=pyopencl git+https://github.com/inducer/islpy.git#egg=islpy git+https://github.com/inducer/loopy.git#egg=loopy git+https://github.com/inducer/boxtree.git#egg=boxtree -git+https://github.com/inducer/arraycontext.git#egg=arraycontext +git+https://github.com/alexfikl/arraycontext.git@flatten-with-leaf-class#egg=arraycontext git+https://github.com/inducer/meshmode.git#egg=meshmode git+https://github.com/inducer/sumpy.git#egg=sumpy git+https://github.com/inducer/pyfmmlib.git#egg=pyfmmlib From f9d5fcdfb5f0fd47b60c629de8eeba796786eeca Mon Sep 17 00:00:00 2001 From: Alexandru Fikl Date: Thu, 13 Jan 2022 09:04:27 -0600 Subject: [PATCH 13/14] fix issues from using leaf_class --- examples/laplace-dirichlet-3d.py | 9 +++------ examples/layerpot-3d.py | 8 +++++--- examples/layerpot.py | 9 +++++---- pytential/linalg/proxy.py | 23 ++++++++++++++++------- pytential/qbx/__init__.py | 8 ++++---- pytential/qbx/direct.py | 2 +- pytential/qbx/geometry.py | 6 +++--- pytential/qbx/interactions.py | 22 ++++++++++++++-------- pytential/source.py | 13 ++++++++----- 9 files changed, 59 insertions(+), 41 deletions(-) diff --git a/examples/laplace-dirichlet-3d.py b/examples/laplace-dirichlet-3d.py index 3ea9e36b0..4178fb576 100644 --- a/examples/laplace-dirichlet-3d.py +++ b/examples/laplace-dirichlet-3d.py @@ -112,15 +112,12 @@ def main(mesh_name="torus", visualize=False): # {{{ fix rhs and solve - from arraycontext import flatten, unflatten nodes = thaw(density_discr.nodes(), actx) - source = np.array([rout, 0, 0]) + source = np.array([rout, 0, 0], dtype=object) def u_incoming_func(x): - x_flat = actx.to_numpy(flatten(x, actx)).reshape(x.size, -1) - dists = 1.0/la.norm(x_flat - source[:, None], axis=0) - - return unflatten(x[0], actx.from_numpy(dists), actx) + dists = x - source + return 1.0 / actx.np.sqrt(sum(dists**2)) bc = u_incoming_func(nodes) bvp_rhs = bind(places, sqrt_w*sym.var("bc"))(actx, bc=bc) diff --git a/examples/layerpot-3d.py b/examples/layerpot-3d.py index 702ecccdb..9a57fb8d2 100644 --- a/examples/layerpot-3d.py +++ b/examples/layerpot-3d.py @@ -84,14 +84,16 @@ def main(mesh_name="ellipsoid"): op = sym.D(kernel, sym.var("sigma"), qbx_forced_limit=None) #op = sym.S(kernel, sym.var("sigma"), qbx_forced_limit=None) - sigma = actx.np.cos(mode_nr*angle) if 0: - from arraycontext import flatten, unflatten - sigma = flatten(0 * angle, actx) from random import randrange + sigma = actx.zeros(density_discr.ndofs, angle.entry_dtype) for _ in range(5): sigma[randrange(len(sigma))] = 1 + + from arraycontext import unflatten sigma = unflatten(angle, sigma, actx) + else: + sigma = actx.np.cos(mode_nr*angle) if isinstance(kernel, HelmholtzKernel): for i, elem in np.ndenumerate(sigma): diff --git a/examples/layerpot.py b/examples/layerpot.py index f190b5389..58293d1e6 100644 --- a/examples/layerpot.py +++ b/examples/layerpot.py @@ -80,14 +80,16 @@ def op(**kwargs): return sym.D(kernel, sym.var("sigma"), **kwargs) #op = sym.S(kernel, sym.var("sigma"), qbx_forced_limit=None, **kwargs) - sigma = actx.np.cos(mode_nr*angle) if 0: - from arraycontext import flatten, unflatten - sigma = flatten(0 * angle, actx) from random import randrange + sigma = actx.zeros(density_discr.ndofs, angle.entry_dtype) for _ in range(5): sigma[randrange(len(sigma))] = 1 + + from arraycontext import unflatten sigma = unflatten(angle, sigma, actx) + else: + sigma = actx.np.cos(mode_nr*angle) if isinstance(kernel, HelmholtzKernel): for i, elem in np.ndenumerate(sigma): @@ -122,7 +124,6 @@ def op(**kwargs): # {{{ plot boundary field from arraycontext import flatten - fld_on_bdry = actx.to_numpy( flatten(bound_bdry_op(actx, sigma=sigma, k=k), actx)) nodes_host = actx.to_numpy( diff --git a/pytential/linalg/proxy.py b/pytential/linalg/proxy.py index a01f4d958..cdfa6006d 100644 --- a/pytential/linalg/proxy.py +++ b/pytential/linalg/proxy.py @@ -77,8 +77,11 @@ def partition_by_nodes( builder = TreeBuilder(actx.context) + from arraycontext import thaw tree, _ = builder(actx.queue, - particles=flatten(discr.nodes(), actx, leaf_class=DOFArray), + particles=flatten( + thaw(discr.nodes(), actx), actx, leaf_class=DOFArray + ), max_particles_in_box=max_particles_in_box, kind=tree_kind) @@ -243,7 +246,8 @@ def prg(): %(insns)s end """ % dict(insns=insns), [ - lp.GlobalArg("sources", None, shape=(ndim, "nsources")), + lp.GlobalArg("sources", None, + shape=(ndim, "nsources"), dim_tags="sep,C"), lp.ValueArg("nsources", np.int64), ... ], @@ -400,7 +404,8 @@ def prg(): proxy_radius[irange] = radius_factor * rblk end """, [ - lp.GlobalArg("sources", None, shape=(ndim, "nsources")), + lp.GlobalArg("sources", None, + shape=(ndim, "nsources"), dim_tags="sep,C"), lp.ValueArg("nsources", np.int64), lp.ValueArg("radius_factor", np.float64), ... @@ -457,9 +462,12 @@ def prg(): proxy_radius[irange] = radius_factor * rqbx end """, [ - lp.GlobalArg("sources", None, shape=(ndim, "nsources")), - lp.GlobalArg("center_int", None, shape=(ndim, "nsources")), - lp.GlobalArg("center_ext", None, shape=(ndim, "nsources")), + lp.GlobalArg("sources", None, + shape=(ndim, "nsources"), dim_tags="sep,C"), + lp.GlobalArg("center_int", None, + shape=(ndim, "nsources"), dim_tags="sep,C"), + lp.GlobalArg("center_ext", None, + shape=(ndim, "nsources"), dim_tags="sep,C"), lp.ValueArg("nsources", np.int64), lp.ValueArg("radius_factor", np.float64), ... @@ -536,7 +544,8 @@ def prg(): """ result[idim, i] = ary[idim, srcindices[i]] """, [ - lp.GlobalArg("ary", None, shape=(discr.ambient_dim, "ndofs")), + lp.GlobalArg("ary", None, + shape=(discr.ambient_dim, "ndofs"), dim_tags="sep,C"), lp.ValueArg("ndofs", np.int64), ...], name="picker_knl", diff --git a/pytential/qbx/__init__.py b/pytential/qbx/__init__.py index a89848463..57d4d5431 100644 --- a/pytential/qbx/__init__.py +++ b/pytential/qbx/__init__.py @@ -23,7 +23,7 @@ import numpy as np import pyopencl as cl -from arraycontext import PyOpenCLArrayContext, thaw, flatten, unflatten +from arraycontext import PyOpenCLArrayContext, thaw, freeze, flatten, unflatten from meshmode.dof_array import DOFArray from pytools import memoize_method, memoize_in, single_valued @@ -780,7 +780,7 @@ def exec_compute_potential_insn_direct(self, actx, insn, bound_expr, evaluate, def _flat_nodes(dofdesc): discr = bound_expr.places.get_discretization( dofdesc.geometry, dofdesc.discr_stage) - return actx.freeze(flatten(discr.nodes(), actx, leaf_class=DOFArray)) + return freeze(flatten(discr.nodes(), actx, leaf_class=DOFArray), actx) @memoize_in(bound_expr.places, (QBXLayerPotentialSource, "flat_expansion_radii")) @@ -789,7 +789,7 @@ def _flat_expansion_radii(dofdesc): bound_expr.places, sym.expansion_radii(self.ambient_dim, dofdesc=dofdesc), )(actx) - return actx.freeze(flatten(radii, actx)) + return freeze(flatten(radii, actx), actx) @memoize_in(bound_expr.places, (QBXLayerPotentialSource, "flat_centers")) @@ -798,7 +798,7 @@ def _flat_centers(dofdesc, qbx_forced_limit): sym.expansion_centers( self.ambient_dim, qbx_forced_limit, dofdesc=dofdesc), )(actx) - return actx.freeze(flatten(centers, actx, leaf_class=DOFArray)) + return freeze(flatten(centers, actx, leaf_class=DOFArray), actx) from pytential.source import evaluate_kernel_arguments flat_kernel_args = evaluate_kernel_arguments( diff --git a/pytential/qbx/direct.py b/pytential/qbx/direct.py index 4c7b037c0..bbb0f7aa1 100644 --- a/pytential/qbx/direct.py +++ b/pytential/qbx/direct.py @@ -50,7 +50,7 @@ def get_kernel(self): lp.GlobalArg("targets", None, shape=(self.dim, "ntargets_total"), order="C"), lp.GlobalArg("center", None, - shape=(self.dim, "ncenters_total")), + shape=(self.dim, "ncenters_total"), dim_tags="sep,C"), lp.GlobalArg("expansion_radii", None, shape="ncenters_total"), lp.GlobalArg("qbx_tgt_numbers", None, diff --git a/pytential/qbx/geometry.py b/pytential/qbx/geometry.py index 72c247a18..8bdfe8df2 100644 --- a/pytential/qbx/geometry.py +++ b/pytential/qbx/geometry.py @@ -26,7 +26,7 @@ import pyopencl.array # noqa from pytools import memoize_method, log_process -from arraycontext import PyOpenCLArrayContext, flatten +from arraycontext import PyOpenCLArrayContext, flatten, freeze from meshmode.dof_array import DOFArray from boxtree.tools import DeviceDataRecord @@ -436,7 +436,7 @@ def flat_centers(self): self.ambient_dim, dofdesc=self.source_dd.to_stage1()))(actx) - return actx.freeze(flatten(centers, actx, leaf_class=DOFArray)) + return freeze(flatten(centers, actx, leaf_class=DOFArray), actx) @memoize_method def flat_expansion_radii(self): @@ -454,7 +454,7 @@ def flat_expansion_radii(self): granularity=sym.GRANULARITY_CENTER, dofdesc=self.source_dd.to_stage1()))(actx) - return actx.freeze(flatten(radii, actx)) + return freeze(flatten(radii, actx), actx) # }}} diff --git a/pytential/qbx/interactions.py b/pytential/qbx/interactions.py index 3c2100a9e..36b4bcfb7 100644 --- a/pytential/qbx/interactions.py +++ b/pytential/qbx/interactions.py @@ -46,7 +46,8 @@ def get_kernel(self): from sumpy.tools import gather_loopy_source_arguments arguments = ( [ - lp.GlobalArg("sources", None, shape=(self.dim, "nsources")), + lp.GlobalArg("sources", None, + shape=(self.dim, "nsources"), dim_tags="sep,C"), lp.GlobalArg("strengths", None, dim_tags="sep,C", shape="strength_count, nsources"), lp.GlobalArg("qbx_center_to_target_box", @@ -55,7 +56,8 @@ def get_kernel(self): None, shape=None), lp.GlobalArg("box_source_starts,box_source_counts_nonchild", None, shape=None), - lp.GlobalArg("qbx_centers", None, shape=("dim", "ncenters")), + lp.GlobalArg("qbx_centers", None, + shape=("dim", "ncenters"), dim_tags="sep,C"), lp.GlobalArg("qbx_expansion_radii", None, shape="ncenters"), lp.GlobalArg("qbx_expansions", None, shape=("ncenters", ncoeffs)), @@ -215,7 +217,8 @@ def get_kernel(self): lp.ValueArg("src_rscale", None), lp.GlobalArg("src_box_starts, src_box_lists", None, shape=None, strides=(1,)), - lp.GlobalArg("qbx_centers", None, shape="dim, ncenters"), + lp.GlobalArg("qbx_centers", None, + shape="dim, ncenters", dim_tags="sep,C"), lp.GlobalArg("qbx_expansion_radii", None, shape="ncenters"), lp.ValueArg("aligned_nboxes,nsrc_level_boxes", np.int32), lp.ValueArg("src_base_ibox", np.int32), @@ -326,7 +329,8 @@ def get_kernel(self): offset=lp.auto), lp.GlobalArg("centers", None, shape="dim, naligned_boxes"), lp.ValueArg("src_rscale", None), - lp.GlobalArg("qbx_centers", None, shape="dim, ncenters"), + lp.GlobalArg("qbx_centers", None, + shape="dim, ncenters", dim_tags="sep,C"), lp.GlobalArg("qbx_expansion_radii", None, shape="ncenters"), lp.ValueArg("naligned_boxes,target_base_ibox,nboxes", np.int32), lp.GlobalArg("expansions", None, @@ -426,14 +430,16 @@ def get_kernel(self): end """], [ - lp.GlobalArg("result", None, shape=("nresults", "ntargets"), - dim_tags="sep,C"), - lp.GlobalArg("qbx_centers", None, shape=("dim", "ncenters")), + lp.GlobalArg("result", None, + shape=("nresults", "ntargets"), dim_tags="sep,C"), + lp.GlobalArg("qbx_centers", None, + shape=("dim", "ncenters"), dim_tags="sep,C"), lp.GlobalArg("center_to_targets_starts,center_to_targets_lists", None, shape=None), lp.GlobalArg("qbx_expansions", None, shape=("ncenters", ncoeffs)), - lp.GlobalArg("targets", None, shape=(self.dim, "ntargets")), + lp.GlobalArg("targets", None, + shape=(self.dim, "ntargets"), dim_tags="sep,C"), lp.ValueArg("ncenters,ntargets", np.int32), ... ] + [arg.loopy_arg for arg in self.expansion.get_args()], diff --git a/pytential/source.py b/pytential/source.py index 8bd5387ea..b44730128 100644 --- a/pytential/source.py +++ b/pytential/source.py @@ -296,8 +296,11 @@ def get_fmm_expansion_wrangler_extra_kwargs( # This contains things like the Helmholtz parameter k or # the normal directions for double layers. - def reorder_sources(source_array): - if isinstance(source_array, cl.array.Array): + def flatten_and_reorder_sources(source_array): + if isinstance(source_array, DOFArray): + source_array = flatten(source_array, actx) + + if isinstance(source_array, actx.array_types): return actx.freeze( actx.thaw(source_array)[tree_user_source_ids] ) @@ -308,15 +311,15 @@ def reorder_sources(source_array): source_extra_kwargs = {} from sumpy.tools import gather_arguments, gather_source_arguments - from arraycontext import map_array_container + from arraycontext import rec_map_array_container for func, var_dict in [ (gather_arguments, kernel_extra_kwargs), (gather_source_arguments, source_extra_kwargs), ]: for arg in func(target_kernels): - var_dict[arg.name] = map_array_container( - lambda ary: reorder_sources(flatten(ary, actx)), + var_dict[arg.name] = rec_map_array_container( + flatten_and_reorder_sources, evaluator(arguments[arg.name]), leaf_class=DOFArray) From 162942751177059d99bb0bf131fcb05c7dbe9421 Mon Sep 17 00:00:00 2001 From: Alexandru Fikl Date: Thu, 13 Jan 2022 15:18:49 -0600 Subject: [PATCH 14/14] redirect requirements.txt to their rightful place --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index df76004e3..1532b3cc3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -9,7 +9,7 @@ git+https://github.com/inducer/pyopencl.git#egg=pyopencl git+https://github.com/inducer/islpy.git#egg=islpy git+https://github.com/inducer/loopy.git#egg=loopy git+https://github.com/inducer/boxtree.git#egg=boxtree -git+https://github.com/alexfikl/arraycontext.git@flatten-with-leaf-class#egg=arraycontext +git+https://github.com/inducer/arraycontext.git#egg=arraycontext git+https://github.com/inducer/meshmode.git#egg=meshmode git+https://github.com/inducer/sumpy.git#egg=sumpy git+https://github.com/inducer/pyfmmlib.git#egg=pyfmmlib