Add vdot to PyOpenCLArrayContext

arraycontext/impl/pyopencl/fake_numpy.py

@@ -162,6 +162,27 @@ def _rec_ravel(a):
 
  return rec_map_array_container(_rec_ravel, a)
 
+ def vdot(self, x, y, dtype=None):
+ import pyopencl.array as cl_array
+ from arraycontext import is_array_container, serialize_container
+
+ def _rec_vdot(xi, yi):
+ if is_array_container(xi):
+ assert type(xi) == type(yi)
+ return sum(_rec_vdot(subxi, subyi)
+ for (_, subxi), (_, subyi) in zip(
+ serialize_container(xi), serialize_container(yi)
+ ))
+ else:
+ result = cl_array.vdot(xi, yi,
+ dtype=dtype, queue=self._array_context.queue)
+ if not self._array_context._force_device_scalars:
+ result = result.get()[()]
+
+ return result
+
+ return _rec_vdot(x, y)
+
 # }}}
 
 

test/test_arraycontext.py

@@ -188,6 +188,22 @@ def _thaw_dofarray(ary, actx):
 
 # {{{ assert_close_to_numpy*
 
+def randn(shape, dtype):
+ rng = np.random.default_rng()
+ dtype = np.dtype(dtype)
+
+ if dtype.kind == "c":
+ dtype = np.dtype(f"<f{dtype.itemsize // 2}")
+ return rng.standard_normal(shape, dtype) \
+ + 1j * rng.standard_normal(shape, dtype)
+ elif dtype.kind == "f":
+ return rng.standard_normal(shape, dtype)
+ elif dtype.kind == "i":
+ return rng.integers(0, 128, shape, dtype)
+ else:
+ raise TypeError(dtype.kind)
+
+
 def assert_close_to_numpy(actx, op, args):
  assert np.allclose(
  actx.to_numpy(
@@ -208,9 +224,12 @@ def assert_close_to_numpy_in_containers(actx, op, args):
  DOFArray(actx, (actx.from_numpy(arg),))
  if isinstance(arg, np.ndarray) else arg
  for arg in args]
- actx_result = actx.to_numpy(op(actx.np, *dofarray_args)[0])
 
- assert np.allclose(actx_result, ref_result)
+ actx_result = op(actx.np, *dofarray_args)
+ if isinstance(actx_result, DOFArray):
+ actx_result = actx_result[0]
+
+ assert np.allclose(actx.to_numpy(actx_result), ref_result)
 
  # }}}
 
@@ -219,9 +238,12 @@ def assert_close_to_numpy_in_containers(actx, op, args):
  obj_array_args = [
  make_obj_array([arg]) if isinstance(arg, DOFArray) else arg
  for arg in dofarray_args]
- obj_array_result = actx.to_numpy(op(actx.np, *obj_array_args)[0][0])
 
- assert np.allclose(obj_array_result, ref_result)
+ obj_array_result = op(actx.np, *obj_array_args)
+ if isinstance(obj_array_result, np.ndarray):
+ obj_array_result = obj_array_result[0][0]
+
+ assert np.allclose(actx.to_numpy(obj_array_result), ref_result)
 
  # }}}
 
@@ -230,34 +252,44 @@ def assert_close_to_numpy_in_containers(actx, op, args):
 
 # {{{ np.function same as numpy
 
-@pytest.mark.parametrize(("sym_name", "n_args"), [
- ("sin", 1),
- ("exp", 1),
- ("arctan2", 2),
- ("minimum", 2),
- ("maximum", 2),
- ("where", 3),
- ("conj", 1),
+@pytest.mark.parametrize(("sym_name", "n_args", "dtype"), [
+ ("sin", 1, np.float64),
+ ("sin", 1, np.complex128),
+ ("exp", 1, np.float64),
+ ("arctan2", 2, np.float64),
+ ("minimum", 2, np.float64),
+ ("maximum", 2, np.float64),
+ ("where", 3, np.float64),
+ ("conj", 1, np.float64),
+ ("conj", 1, np.complex128),
+ ("vdot", 2, np.float64),
+ ("vdot", 2, np.complex128),
+ ("abs", 1, np.float64),
+ ("abs", 1, np.complex128),
  ])
-def test_array_context_np_workalike(actx_factory, sym_name, n_args):
+def test_array_context_np_workalike(actx_factory, sym_name, n_args, dtype):
  actx = actx_factory()
+ if not hasattr(actx.np, sym_name):
+ pytest.skip(f"'{sym_name}' not implemented on '{type(actx).__name__}'")
+
+ ndofs = 512
+ args = [randn(ndofs, dtype) for i in range(n_args)]
 
- ndofs = 5000
- args = [np.random.randn(ndofs) for i in range(n_args)]
  assert_close_to_numpy_in_containers(
  actx, lambda _np, *_args: getattr(_np, sym_name)(*_args), args)
 
 
-@pytest.mark.parametrize(("sym_name", "n_args"), [
- # ("empty_like", 1), # NOTE: fails np.allclose, obviously
- ("zeros_like", 1),
- ("ones_like", 1),
+@pytest.mark.parametrize(("sym_name", "n_args", "dtype"), [
+ ("zeros_like", 1, np.float64),
+ ("zeros_like", 1, np.complex128),
+ ("ones_like", 1, np.float64),
+ ("ones_like", 1, np.complex128),
  ])
-def test_array_context_np_like(actx_factory, sym_name, n_args):
+def test_array_context_np_like(actx_factory, sym_name, n_args, dtype):
  actx = actx_factory()
 
- ndofs = 5000
- args = [np.random.randn(ndofs) for i in range(n_args)]
+ ndofs = 512
+ args = [randn(ndofs, dtype) for i in range(n_args)]
  assert_close_to_numpy(
  actx, lambda _np, *_args: getattr(_np, sym_name)(*_args), args)
 
@@ -460,6 +492,7 @@ def get_imag(ary):
 
 
 # {{{ reductions same as numpy
+
 @pytest.mark.parametrize("op", ["sum", "min", "max"])
 def test_dof_array_reductions_same_as_numpy(actx_factory, op):
  actx = actx_factory()
@@ -780,10 +813,12 @@ def test_numpy_conversion(actx_factory):
 # }}}
 
 
+# {{{ test actx.np.linalg.norm
+
 @pytest.mark.parametrize("norm_ord", [2, np.inf])
 def test_norm_complex(actx_factory, norm_ord):
  actx = actx_factory()
- a = np.random.randn(2000) + 1j * np.random.randn(2000)
+ a = randn(2000, np.complex128)
 
  norm_a_ref = np.linalg.norm(a, norm_ord)
  norm_a = actx.np.linalg.norm(actx.from_numpy(a), norm_ord)
@@ -809,6 +844,8 @@ def test_norm_ord_none(actx_factory, ndim):
 
  np.testing.assert_allclose(actx.to_numpy(norm_a), norm_a_ref)
 
+# }}}
+
 
 # {{{ test_actx_compile helpers
 
@@ -828,8 +865,6 @@ def scale_and_orthogonalize(alpha, vel):
  vel)
  return Velocity2D(-scaled_vel.v, scaled_vel.u, actx)
 
-# }}}
-
 
 def test_actx_compile(actx_factory):
  from arraycontext import (to_numpy, from_numpy)
@@ -848,6 +883,10 @@ def test_actx_compile(actx_factory):
  np.testing.assert_allclose(result.u, -3.14*v_y)
  np.testing.assert_allclose(result.v, 3.14*v_x)
 
+# }}}
+
+
+# {{{ test_container_equality
 
 def test_container_equality(actx_factory):
  actx = actx_factory()
@@ -865,17 +904,10 @@ def test_container_equality(actx_factory):
 
  assert isinstance(bcast_dc_of_dofs == bcast_dc_of_dofs_2, MyContainerDOFBcast)
 
+# }}}
 
-def test_abs_complex(actx_factory):
- actx = actx_factory()
- a = np.random.randn(2000) + 1j * np.random.randn(2000)
-
- abs_a_ref = np.abs(a)
- abs_a = actx.np.abs(actx.from_numpy(a))
-
- assert abs_a.dtype == abs_a_ref.dtype
- np.testing.assert_allclose(actx.to_numpy(abs_a), abs_a_ref)
 
+# {{{ test_leaf_array_type_broadcasting
 
 @with_container_arithmetic(
  bcast_obj_array=True,
@@ -925,6 +957,8 @@ def _actx_allows_scalar_broadcast(actx):
  np.testing.assert_allclose(actx.to_numpy(bar.u[0]),
  actx.to_numpy(quuz.u[0]))
 
+# }}}
+
 
 if __name__ == "__main__":
  import sys