feat: port to new obj_array functions

Author: alexfikl

experiments/poisson.py

@@ -12,7 +12,7 @@
 from meshmode.discretization.visualization import make_visualizer
 
 from pytential import bind, sym, norm  # noqa
-from pytools.obj_array import make_obj_array
+from pytools import obj_array
 
 import pytential.symbolic.primitives as p
 
@@ -196,7 +196,7 @@ def get_kernel():
                     p.area_element(mesh.ambient_dim, mesh.dim))
                 (queue)).get())
 
-    centers = make_obj_array([ci.copy().reshape(vol_discr.nnodes) for ci in targets])
+    centers = obj_array.new_1d([ci.copy().reshape(vol_discr.nnodes) for ci in targets])
     centers[2][:] = center_dist
 
     print(center_dist)

experiments/qbx-tangential-deriv-jump.py

@@ -34,10 +34,10 @@ def main():
     qbx = QBXLayerPotentialSource(density_discr, 4*target_order,
             qbx_order, fmm_order=False)
 
-    from pytools.obj_array import join_fields
+    from pytools import obj_array
     sig_sym = sym.var("sig")
     knl = LaplaceKernel(2)
-    op = join_fields(
+    op = obj_array.flat(
             sym.tangential_derivative(mesh.ambient_dim,
                 sym.D(knl, sig_sym, qbx_forced_limit=+1)).as_scalar(),
             sym.tangential_derivative(mesh.ambient_dim,

experiments/stokes-2d-interior.py

@@ -33,8 +33,8 @@ def main(nelements):
     logging.basicConfig(level=logging.INFO)
 
     def get_obj_array(obj_array):
-        from pytools.obj_array import make_obj_array
-        return make_obj_array([
+        from pytools import obj_array
+        return obj_array.new_1d([
                 ary.get()
                 for ary in obj_array
                 ])
@@ -72,7 +72,6 @@ def get_obj_array(obj_array):
 
     from sumpy.kernel import LaplaceKernel
     from pytential.symbolic.stokes import StressletWrapper
-    from pytools.obj_array import make_obj_array
     dim=2
     cse = sym.cse
 

experiments/two-domain-helmholtz.py

@@ -27,7 +27,7 @@
 import pyopencl.array  # noqa
 import pyopencl.clmath  # noqa
 
-from pytools.obj_array import make_obj_array
+from pytools import obj_array
 
 from meshmode.discretization import Discretization
 from meshmode.discretization.poly_element import \
@@ -105,7 +105,7 @@ def main():
     bound_pde_op = bind(qbx, pde_op.operator(op_unknown_sym))
 
     # in inner domain
-    sources_1 = make_obj_array(list(np.array([
+    sources_1 = obj_array.new_1d(list(np.array([
         [-1.5, 0.5]
         ]).T.copy()))
     strengths_1 = np.array([1])

pytential/linalg/direct_solver_symbolic.py

@@ -20,7 +20,7 @@
 THE SOFTWARE.
 """
 
-from pytools.obj_array import make_obj_array
+from pytools import obj_array
 
 from pytential.symbolic.mappers import (
         IdentityMapper, OperatorCollector, LocationTagger)
@@ -44,7 +44,7 @@ class PROXY_SKELETONIZATION_TARGET:             # noqa: N801
 
 def prepare_expr(places, exprs, auto_where=None):
     from pytential.symbolic.execution import _prepare_expr
-    return make_obj_array([
+    return obj_array.new_1d([
         _prepare_expr(places, expr, auto_where=auto_where)
         for expr in exprs])
 
@@ -60,7 +60,7 @@ def _prepare_expr(expr):
 
         return expr
 
-    return make_obj_array([_prepare_expr(expr) for expr in exprs])
+    return obj_array.new_1d([_prepare_expr(expr) for expr in exprs])
 
 # }}}
 

pytential/qbx/__init__.py

@@ -485,10 +485,9 @@ def drive_cost_model(
                     calibration_params
                 )
 
-            from pytools.obj_array import obj_array_vectorize
-            from functools import partial
+            from pytools import obj_array
             return (
-                    obj_array_vectorize(
+                    obj_array.vectorize(
                         partial(wrangler.finalize_potentials,
                             template_ary=strengths[0]),
                         wrangler.full_output_zeros(strengths[0])),
@@ -518,7 +517,6 @@ def _dispatch_compute_potential_insn(self, actx, insn, bound_expr,
 
     @memoize_method
     def _tree_indep_data_for_wrangler(self, source_kernels, target_kernels):
-        from functools import partial
         base_kernel = single_valued(kernel.get_base_kernel() for
             kernel in source_kernels)
         mpole_expn_class = \

pytential/qbx/fmm.py

@@ -143,8 +143,8 @@ def output_zeros(self, template_ary):
 
         nqbtl = self.geo_data.non_qbx_box_target_lists()
 
-        from pytools.obj_array import make_obj_array
-        return make_obj_array([
+        from pytools import obj_array
+        return obj_array.new_1d([
                 cl.array.zeros(
                     template_ary.queue,
                     nqbtl.nfiltered_targets,
@@ -596,8 +596,8 @@ def reorder_and_finalize_potentials(x):
         # potential back into a CL array.
         return wrangler.finalize_potentials(x[tree.sorted_target_ids], template_ary)
 
-    from pytools.obj_array import obj_array_vectorize
-    result = obj_array_vectorize(
+    from pytools import obj_array
+    result = obj_array.vectorize(
             reorder_and_finalize_potentials, all_potentials_in_tree_order)
 
     # }}}

pytential/qbx/fmmlib.py

@@ -209,16 +209,16 @@ def output_zeros(self):
 
         nqbtl = self.geo_data.non_qbx_box_target_lists()
 
-        from pytools.obj_array import make_obj_array
-        return make_obj_array([
+        from pytools import obj_array
+        return obj_array.new_1d([
                 np.zeros(nqbtl.nfiltered_targets, self.tree_indep.dtype)
                 for k in self.tree_indep.outputs])
 
     def full_output_zeros(self, template_ary):
         """This includes QBX and non-QBX targets."""
 
-        from pytools.obj_array import make_obj_array
-        return make_obj_array([
+        from pytools import obj_array
+        return obj_array.new_1d([
                 np.zeros(self.tree.ntargets, self.tree_indep.dtype)
                 for k in self.tree_indep.outputs])
 

pytential/symbolic/compiler.py

@@ -394,7 +394,7 @@ def _get_next_step(
 
     # {{{ make sure results do not get discarded
 
-    from pytools.obj_array import obj_array_vectorize
+    from pytools import obj_array
 
     from pytential.symbolic.mappers import DependencyMapper
     dm = DependencyMapper(composite_leaves=False)
@@ -409,7 +409,7 @@ def remove_result_variable(result_expr):
             assert isinstance(var, Variable)
             discardable_vars.discard(var.name)
 
-    obj_array_vectorize(remove_result_variable, result)
+    obj_array.vectorize(remove_result_variable, result)
 
     # }}}
 

pytential/symbolic/execution.py

@@ -492,8 +492,8 @@ def unflatten(self, ary):
             components.append(component)
 
         if self._operator_uses_obj_array:
-            from pytools.obj_array import make_obj_array
-            return make_obj_array(components)
+            from pytools import obj_array
+            return obj_array.new_1d(components)
         else:
             return components[0]
 
@@ -651,8 +651,8 @@ def execute(code: Code, exec_mapper, pre_assign_check=None) -> np.ndarray:
             assert target in assignees
             context[target] = value
 
-    from pytools.obj_array import obj_array_vectorize
-    return obj_array_vectorize(exec_mapper, code.result)
+    from pytools import obj_array
+    return obj_array.vectorize(exec_mapper, code.result)
 
 # }}}
 
@@ -967,8 +967,8 @@ def build_matrix(actx, places, exprs, input_exprs, domains=None,
     exprs = _prepare_expr(places, exprs, auto_where=auto_where)
 
     if not (isinstance(exprs, np.ndarray) and exprs.dtype.char == "O"):
-        from pytools.obj_array import make_obj_array
-        exprs = make_obj_array([exprs])
+        from pytools import obj_array
+        exprs = obj_array.new_1d([exprs])
 
     try:
         input_exprs = list(input_exprs)