Remove usage of meshmode.dof_array.flatten and friends

examples/laplace-dirichlet-3d.py

@@ -112,23 +112,14 @@ def main(mesh_name="torus", visualize=False):
 
     # {{{ fix rhs and solve
 
-    from meshmode.dof_array 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):
-        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)))
+        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)
 
     from pytential.solve import gmres

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 meshmode.dof_array import flatten, unflatten
-        sigma = flatten(0 * angle)
         from random import randrange
+        sigma = actx.zeros(density_discr.ndofs, angle.entry_dtype)
         for _ in range(5):
             sigma[randrange(len(sigma))] = 1
-        sigma = unflatten(actx, density_discr, sigma)
+
+        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):

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 meshmode.dof_array import flatten, unflatten
-        sigma = flatten(0 * angle)
         from random import randrange
+        sigma = actx.zeros(density_discr.ndofs, angle.entry_dtype)
         for _ in range(5):
             sigma[randrange(len(sigma))] = 1
-        sigma = unflatten(actx, density_discr, sigma)
+
+        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):
@@ -121,11 +123,12 @@ def op(**kwargs):
     if enable_mayavi:
         # {{{ plot boundary field
 
-        from meshmode.dof_array import flatten_to_numpy
-
-        fld_on_bdry = flatten_to_numpy(
-                actx, bound_bdry_op(actx, sigma=sigma, k=k))
-        nodes_host = flatten_to_numpy(actx, density_discr.nodes())
+        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(
+                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)

pytential/linalg/proxy.py

@@ -26,8 +26,9 @@
 import numpy as np
 import numpy.linalg as la
 
-from arraycontext import PyOpenCLArrayContext
+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,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=flatten(
+                    thaw(discr.nodes(), actx), actx, leaf_class=DOFArray
+                    ),
                 max_particles_in_box=max_particles_in_box,
                 kind=tree_kind)
 
@@ -330,9 +332,7 @@ 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 = flatten(discr.nodes(), actx, leaf_class=DOFArray)
 
         knl = self.get_centers_knl(actx)
         _, (centers_dev,) = knl(actx.queue,
@@ -510,11 +510,10 @@ 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=flatten(center_int, actx, leaf_class=DOFArray),
+                center_ext=flatten(center_ext, actx, leaf_class=DOFArray),
                 **kwargs)
 
 # }}}
@@ -560,10 +559,8 @@ 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=flatten(discr.nodes(), actx, leaf_class=DOFArray),
             srcindices=pxy.srcindices.indices)
 
     # }}}

pytential/qbx/__init__.py

@@ -23,8 +23,8 @@
 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, freeze, flatten, unflatten
+from meshmode.dof_array import DOFArray
 
 from pytools import memoize_method, memoize_in, single_valued
 from pytential.qbx.target_assoc import QBXTargetAssociationFailedException
@@ -622,7 +622,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)
 
@@ -683,7 +683,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, strict=False)
 
             results.append((o.name, result))
 
@@ -779,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 freeze(flatten(thaw(discr.nodes(), actx), strict=False), actx)
+            return freeze(flatten(discr.nodes(), actx, leaf_class=DOFArray), actx)
 
         @memoize_in(bound_expr.places,
                 (QBXLayerPotentialSource, "flat_expansion_radii"))
@@ -788,7 +789,7 @@ def _flat_expansion_radii(dofdesc):
                     bound_expr.places,
                     sym.expansion_radii(self.ambient_dim, dofdesc=dofdesc),
                     )(actx)
-            return freeze(flatten(radii), actx)
+            return freeze(flatten(radii, actx), actx)
 
         @memoize_in(bound_expr.places,
                 (QBXLayerPotentialSource, "flat_centers"))
@@ -797,13 +798,12 @@ 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 freeze(flatten(centers, actx, leaf_class=DOFArray), actx)
 
-        kernel_args = {}
-        for arg_name, arg_expr in insn.kernel_arguments.items():
-            kernel_args[arg_name] = flatten(evaluate(arg_expr), strict=False)
-
-        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)
 
@@ -863,12 +863,13 @@ 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):
-                    result = unflatten(actx, target_discr, result)
+                    template_ary = thaw(target_discr.nodes()[0], actx)
+                    result = unflatten(template_ary, result, actx, strict=False)
 
                 results[i] = (o.name, result)
 
@@ -891,7 +892,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(
@@ -925,7 +926,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
 
@@ -944,7 +945,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, strict=False)
 
                 results[i] = (o.name, result)
 
@@ -958,7 +960,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(
@@ -967,8 +969,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]
 
 # }}}
 

pytential/qbx/geometry.py

@@ -26,10 +26,8 @@
 import pyopencl.array  # noqa
 
 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, freeze
+from meshmode.dof_array import DOFArray
 
 from boxtree.tools import DeviceDataRecord
 from boxtree.pyfmmlib_integration import FMMLibRotationDataInterface
@@ -433,10 +431,12 @@ 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 freeze(flatten(centers, actx, leaf_class=DOFArray), actx)
 
     @memoize_method
     def flat_expansion_radii(self):
@@ -447,14 +447,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 freeze(flatten(radii, actx), actx)
 
     # }}}
 
@@ -465,7 +465,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 +475,19 @@ 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=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,8 +563,9 @@ def tree(self):
         refine_weights.finish()
 
         tree, _ = code_getter.build_tree()(queue,
-                particles=flatten(thaw(
-                    quad_stage2_discr.nodes(), self.array_context)),
+                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,