Port to new obj_array functions

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)

pytential/symbolic/matrix.py

@@ -401,8 +401,8 @@ def map_interpolation(self, expr):
                     expr.from_dd.geometry, expr.from_dd.discr_stage)
             template_ary = actx.thaw(discr.nodes()[0])
 
-            from pytools.obj_array import make_obj_array
-            return make_obj_array([
+            from pytools import obj_array
+            return obj_array.new_1d([
                 actx.to_numpy(flatten(
                     conn(unflatten(template_ary, actx.from_numpy(o), actx)),
                     actx))

pytential/symbolic/old_diffop_primitives.py

@@ -1,30 +1,26 @@
 # {{{ differential operators on layer potentials
 
 def grad_S(kernel, arg, dim):
-    from pytools.obj_array import log_shape
-    arg_shape = log_shape(arg)
-    result = np.zeros(arg_shape+(dim,), dtype=object)
+    result = np.zeros(arg.shape+(dim,), dtype=object)
     from pytools import indices_in_shape
-    for i in indices_in_shape(arg_shape):
+    for i in indices_in_shape(arg.shape):
         for j in range(dim):
             result[i+(j,)] = IntGdTarget(kernel, arg[i], j)
     return result
 
 
 def grad_D(kernel, arg, dim):
-    from pytools.obj_array import log_shape
-    arg_shape = log_shape(arg)
-    result = np.zeros(arg_shape+(dim,), dtype=object)
+    result = np.zeros(arg.shape+(dim,), dtype=object)
     from pytools import indices_in_shape
-    for i in indices_in_shape(arg_shape):
+    for i in indices_in_shape(arg.shape):
         for j in range(dim):
             result[i+(j,)] = IntGdMixed(kernel, arg[i], j)
     return result
 
 
 def tangential_surf_grad_source_S(kernel, arg, dim=3):
-    from pytools.obj_array import make_obj_array
-    return make_obj_array([
+    from pytools import obj_array
+    return obj_array.new_1d([
         IntGdSource(kernel, arg,
             ds_direction=make_tangent(i, dim, "src"))
         for i in range(dim-1)])
@@ -42,15 +38,15 @@ def curl_curl_S_volume(k, arg):
     # By vector identity, this is grad div S volume + k^2 S_k(arg),
     # since S_k(arg) satisfies a Helmholtz equation.
 
-    from pytools.obj_array import make_obj_array
+    from pytools import obj_array
 
     def swap_min_first(i, j):
         if i < j:
             return i, j
         else:
             return j, i
 
-    return make_obj_array([
+    return obj_array.new_1d([
         sum(IntGd2Target(k, arg[m], *swap_min_first(m, n)) for m in range(3))
         for n in range(3)]) + k**2*S(k, arg)
 
@@ -133,8 +129,8 @@ def project_to_tangential(xyz_vec, which=None):
 
 def surf_n_cross(tangential_vec):
     assert len(tangential_vec) == 2
-    from pytools.obj_array import make_obj_array
-    return make_obj_array([-tangential_vec[1], tangential_vec[0]])
+    from pytools import obj_array
+    return obj_array.new_1d([-tangential_vec[1], tangential_vec[0]])
 
 # }}}
 

pytential/symbolic/pde/cahn_hilliard.py

@@ -83,11 +83,13 @@ def Sn_G(i, density):
                         qbx_forced_limit="avg",
                         op_map=partial(sym.normal_derivative, 2))
 
-        d = sym.make_obj_array([
+        from pytools import obj_array
+
+        d = obj_array.new_1d([
             0.5*sig1,
             0.5*lam2**2*sig1 - 0.5*sig2
             ])
-        a = sym.make_obj_array([
+        a = obj_array.new_1d([
             # A11
             Sn_G(1, sig1) + c*S_G(1, sig1)
             # A12

pytential/symbolic/pde/maxwell/generalized_debye.py

@@ -151,8 +151,8 @@ def volume_field_base(self, r, q, j):
         E = 1j*k*A - grad_phi - curl_S_volume(k, m)
         H = curl_S_volume(k, j) + 1j*k*Q - grad_psi
 
-        from pytools.obj_array import flat_obj_array
-        return flat_obj_array(E, H)
+        from pytools import obj_array
+        return obj_array.flat(E, H)
 
     def integral_equation(self, *args, **kwargs):
         nxnxE, ndotH = self.boundary_field(*args)
@@ -162,8 +162,8 @@ def integral_equation(self, *args, **kwargs):
         if kwargs:
             raise TypeError("invalid keyword argument(s)")
 
-        from pytools.obj_array import make_obj_array
-        eh_op = make_obj_array([
+        from pytools import obj_array
+        eh_op = obj_array.new_1d([
             2*_debye_S0_surf_div(nxnxE),
             -ndotH,
             ]) + fix
@@ -176,8 +176,8 @@ def integral_equation(self, *args, **kwargs):
         E_minus_grad_phi = 1j*k*A - curl_S_volume(k, m)
 
         from hellskitchen.fmm import DifferenceKernel
-        from pytools.obj_array import flat_obj_array
-        return flat_obj_array(
+        from pytools import obj_array
+        return obj_array.flat(
                 eh_op,
                 # FIXME: These are inefficient. They compute a full volume field,
                 # but only actually use the line part of it.
@@ -220,8 +220,8 @@ def prepare_rhs(self, e_inc, h_inc):
                     self.h_on_spanning_surface_symbols)]
                 )
 
-        from pytools.obj_array import make_obj_array
-        return make_obj_array(result)
+        from pytools import obj_array
+        return obj_array.new_1d(result)
 
 
     def harmonic_vector_field_current(self, hvf_coefficients):
@@ -260,10 +260,10 @@ def inv_rank_one_coeff(u):
             r_coeff = inv_rank_one_coeff(r_tilde)
             q_coeff = inv_rank_one_coeff(q_tilde)
 
-            from pytools.obj_array import flat_obj_array
+            from pytools import obj_array
             factors = self.cluster_points()
 
-            fix = flat_obj_array(
+            fix = obj_array.flat(
                     factors[0]*s_ones*r_coeff,
                     factors[1]*Ones()*q_coeff,
                     )
@@ -374,10 +374,10 @@ def inv_rank_one_coeff(u):
             r_coeff = inv_rank_one_coeff(r_tilde)
             q_coeff = inv_rank_one_coeff(q_tilde)
 
-            from pytools.obj_array import flat_obj_array
+            from pytools import obj_array
             factors = self.cluster_points()
 
-            fix = flat_obj_array(
+            fix = obj_array.flat(
                     factors[0]*s_ones*(r_coeff),
                     factors[1]*Ones()*(q_coeff),
                     )