diff --git a/examples/simple-dg.py b/examples/simple-dg.py index 2858bb4c234e77c3f7c84fb2f4aeafe2af8a487e..b945687758f172194d859560bc1723064d56971b 100644 --- a/examples/simple-dg.py +++ b/examples/simple-dg.py @@ -455,10 +455,9 @@ def bump(actx, discr, t=0): nodes[1] - source_center[1], ]) - exp = (actx.special_func("exp")) return ( np.cos(source_omega*t) - * exp( + * actx.np.exp( -np.dot(center_dist, center_dist) / source_width**2)) diff --git a/meshmode/array_context.py b/meshmode/array_context.py index 5ed7d88cd68558647fc80851cd52f166b07e0aae..26da731eaa4a6e44873bd66d5a514a18b1e2e8dc 100644 --- a/meshmode/array_context.py +++ b/meshmode/array_context.py @@ -52,6 +52,24 @@ def make_loopy_program(domains, statements, kernel_data=["..."], name=None): # {{{ ArrayContext +class _FakeNumpyNamespace: + def __init__(self, array_context): + self._array_context = array_context + + def __getattr__(self, name): + def f(*args): + actx = self._array_context + # FIXME: Maybe involve loopy type inference? + result = actx.empty(args[0].shape, args[0].dtype) + prg = actx._get_scalar_func_loopy_program(name, len(args)) + actx.call_loopy(prg, out=result, + **{"inp%d" % i: arg for i, arg in enumerate(args)}) + return result + + from pytools.obj_array import obj_array_vectorized_n_args + return obj_array_vectorized_n_args(f) + + class ArrayContext: """An interface that allows a :class:`Discretization` to create and interact with arrays of degrees of freedom without fully specifying their types. @@ -63,13 +81,29 @@ class ArrayContext: .. automethod:: from_numpy .. automethod:: to_numpy .. automethod:: call_loopy - .. automethod:: special_func + .. attribute:: np + + Provides access to a namespace that serves as a work-alike to + :mod:`numpy`. The actual level of functionality provided is up to the + individual array context implementation, however the functions and + objects available under this namespace must not behave differently + from :mod:`numpy`. + + As a baseline, special functions available through :mod:`loopy` + (e.g. ``sin``, ``exp``) are accessible through this interface. + + Callables accessible through this namespace vectorize over object + arrays, including :class:`meshmode.dof_array.DOFArray`. + .. automethod:: freeze .. automethod:: thaw .. versionadded:: 2020.2 """ + def __init__(self): + self.np = _FakeNumpyNamespace(self) + def empty(self, shape, dtype): raise NotImplementedError @@ -112,7 +146,7 @@ class ArrayContext: raise NotImplementedError @memoize_method - def _get_special_func_loopy_program(self, name, nargs): + def _get_scalar_func_loopy_program(self, name, nargs): from pymbolic import var iel = var("iel") idof = var("idof") @@ -126,25 +160,6 @@ class ArrayContext: ], name="actx_special_%s" % name) - @memoize_method - def special_func(self, name): - """Returns a callable for the special function *name*, where *name* is a - (potentially dotted) function name resolvable by :mod:`loopy`. - - The returned callable will vectorize over object arrays, including - :class:`meshmode.dof_array.DOFArray`. - """ - def f(*args): - # FIXME: Maybe involve loopy type inference? - result = self.empty(args[0].shape, args[0].dtype) - prg = self._get_special_func_loopy_program(name, len(args)) - self.call_loopy(prg, out=result, - **{"inp%d" % i: arg for i, arg in enumerate(args)}) - return result - - from pytools.obj_array import obj_array_vectorized_n_args - return obj_array_vectorized_n_args(f) - def freeze(self, array): """Return a version of the context-defined array *array* that is 'frozen', i.e. suitable for long-term storage and reuse. Frozen arrays @@ -194,6 +209,7 @@ class PyOpenCLArrayContext(ArrayContext): """ def __init__(self, queue, allocator=None): + super().__init__() self.context = queue.context self.queue = queue self.allocator = allocator diff --git a/test/test_chained.py b/test/test_chained.py index 91e610596bae66eabe1a96176e54c5f80f79db1e..51c1e12c494f32af90e8a7fa98724188a04bc2ca 100644 --- a/test/test_chained.py +++ b/test/test_chained.py @@ -214,11 +214,9 @@ def test_chained_connection(ctx_factory, ndim, visualize=False): ChainedDiscretizationConnection chained = ChainedDiscretizationConnection(connections) - sin = actx.special_func("sin") - def f(x): - from six.moves import reduce - return 0.1 * reduce(lambda x, y: x * sin(5 * y), x) + from functools import reduce + return 0.1 * reduce(lambda x, y: x * actx.np.sin(5 * y), x) x = thaw(actx, connections[0].from_discr.nodes()) fx = f(x) @@ -249,11 +247,9 @@ def test_chained_full_resample_matrix(ctx_factory, ndim, visualize=False): ChainedDiscretizationConnection chained = ChainedDiscretizationConnection(connections) - sin = actx.special_func("sin") - def f(x): from six.moves import reduce - return 0.1 * reduce(lambda x, y: x * sin(5 * y), x) + return 0.1 * reduce(lambda x, y: x * actx.np.sin(5 * y), x) resample_mat = actx.to_numpy(make_full_resample_matrix(actx, chained)) @@ -320,11 +316,9 @@ def test_chained_to_direct(ctx_factory, ndim, chain_type, to_element_indices[i] += 1 assert np.min(to_element_indices) > 0 - sin = actx.special_func("sin") - def f(x): from six.moves import reduce - return 0.1 * reduce(lambda x, y: x * sin(5 * y), x) + return 0.1 * reduce(lambda x, y: x * actx.np.sin(5 * y), x) x = thaw(actx, connections[0].from_discr.nodes()) fx = f(x) @@ -394,13 +388,11 @@ def test_reversed_chained_connection(ctx_factory, ndim, mesh_name): from_nodes = thaw(actx, chained.from_discr.nodes()) to_nodes = thaw(actx, chained.to_discr.nodes()) - cos = actx.special_func("cos") - from_x = 0 to_x = 0 for d in range(ndim): - from_x += cos(from_nodes[d]) ** (d + 1) - to_x += cos(to_nodes[d]) ** (d + 1) + from_x += actx.np.cos(from_nodes[d]) ** (d + 1) + to_x += actx.np.cos(to_nodes[d]) ** (d + 1) from_interp = reverse(to_x) diff --git a/test/test_meshmode.py b/test/test_meshmode.py index d4952e3813e019b7046b838d920ec7dd744cfa2b..a888beff918cc0ef1ccc002fa2f9878123999911 100644 --- a/test/test_meshmode.py +++ b/test/test_meshmode.py @@ -235,10 +235,8 @@ def test_boundary_interpolation(ctx_factory, group_factory, boundary_tag, order = 4 - sin = actx.special_func("sin") - def f(x): - return 0.1*sin(30*x) + return 0.1*actx.np.sin(30*x) for mesh_par in mesh_pars: # {{{ get mesh @@ -331,10 +329,8 @@ def test_all_faces_interpolation(ctx_factory, mesh_name, dim, mesh_pars, order = 4 - sin = actx.special_func("sin") - def f(x): - return 0.1*sin(30*x) + return 0.1*actx.np.sin(30*x) for mesh_par in mesh_pars: # {{{ get mesh @@ -450,10 +446,8 @@ def test_opposite_face_interpolation(ctx_factory, group_factory, order = 5 - sin = actx.special_func("sin") - def f(x): - return 0.1*sin(30*x) + return 0.1*actx.np.sin(30*x) for mesh_par in mesh_pars: # {{{ get mesh @@ -787,10 +781,8 @@ def test_sanity_qhull_nd(ctx_factory, dim, order): from meshmode.discretization.connection import make_same_mesh_connection cnx = make_same_mesh_connection(actx, high_discr, low_discr) - sin = actx.special_func("sin") - def f(x): - return 0.1*sin(x) + return 0.1*actx.np.sin(x) x_low = thaw(actx, low_discr.nodes()[0]) f_low = f(x_low) diff --git a/test/test_partition.py b/test/test_partition.py index f34fe3e89cffdd4daf28cf76b5cd3d49ea0953a4..1e8e6c7fb5d556807375dc3e132b4e1b1c3f3a09 100644 --- a/test/test_partition.py +++ b/test/test_partition.py @@ -81,10 +81,8 @@ def test_partition_interpolation(ctx_factory, dim, mesh_pars, continue eoc_rec[i, j] = EOCRecorder() - sin = actx.special_func("sin") - def f(x): - return 10.*sin(50.*x) + return 10.*actx.np.sin(50.*x) for n in mesh_pars: from meshmode.mesh.generation import generate_warped_rect_mesh @@ -415,10 +413,8 @@ def _test_data_transfer(mpi_comm, actx, local_bdry_conns, remote_to_local_bdry_conns, connected_parts): from mpi4py import MPI - sin = actx.special_func("sin") - def f(x): - return 10*sin(20.*x) + return 10*actx.np.sin(20.*x) ''' Here is a simplified example of what happens from diff --git a/test/test_refinement.py b/test/test_refinement.py index 3162461046f5e0fb67d712216c39be4a783a2452..336c8a93cf4530c404e38dc5018ac921b7e685f8 100644 --- a/test/test_refinement.py +++ b/test/test_refinement.py @@ -198,8 +198,6 @@ def test_refinement_connection( from meshmode.discretization.connection import ( make_refinement_connection, check_connection) - sin = actx.special_func("sin") - from pytools.convergence import EOCRecorder eoc_rec = EOCRecorder() @@ -239,7 +237,7 @@ def test_refinement_connection( factor = 9 for iaxis in range(len(x)): - result = result * sin(factor * (x[iaxis]/mesh_ext[iaxis])) + result = result * actx.np.sin(factor * (x[iaxis]/mesh_ext[iaxis])) return result