diff --git a/examples/wave/wave-op-var-velocity.py b/examples/wave/wave-op-var-velocity.py
index 2d8c6903dd7d311f41567a6433b73599d87e3337..758c7d38ac6269e4c7ba5672542c07500b5832d4 100644
--- a/examples/wave/wave-op-var-velocity.py
+++ b/examples/wave/wave-op-var-velocity.py
@@ -90,7 +90,8 @@ def wave_operator(dcoll, c, w):
op.inverse_mass(dcoll,
flat_obj_array(
-op.weak_local_div(dcoll, dd_quad, c_quad*v_quad),
- -op.weak_local_grad(dcoll, dd_quad, c_quad*u_quad)
+ -op.weak_local_grad(dcoll, dd_quad, c_quad*u_quad) \
+ # pylint: disable=E1130
)
+ # noqa: W504
op.face_mass(dcoll,
diff --git a/grudge/op.py b/grudge/op.py
index 44b9ede79227bfc43bead56be82d1778f98af072..fe39cfa855661954614722d96dd07a97d8b441ef 100644
--- a/grudge/op.py
+++ b/grudge/op.py
@@ -260,14 +260,26 @@ def _compute_local_gradient(dcoll, vec, xyz_axis):
)
-def local_grad(dcoll, vec):
+def local_grad(dcoll, vec, *, nested=False):
r"""Return the element-local gradient of the volume function represented by
*vec*.
:arg dcoll: a :class:`grudge.discretization.DiscretizationCollection`.
- :arg vec: a :class:`~meshmode.dof_array.DOFArray`
- :returns: an object array of :class:`~meshmode.dof_array.DOFArray`\ s
+ :arg vec: a :class:`~meshmode.dof_array.DOFArray` or object array of
+ `~meshmode.dof_array.DOFArray`
+ :arg nested: return nested object arrays instead of a single multidimensional
+ array if *vec* is non-scalar
+ :returns: an object array (possibly nested) of
+ :class:`~meshmode.dof_array.DOFArray`\ s
"""
+ if isinstance(vec, np.ndarray):
+ grad = obj_array_vectorize(
+ lambda el: local_grad(dcoll, el, nested=nested), vec)
+ if nested:
+ return grad
+ else:
+ return np.stack(grad, axis=0)
+
return make_obj_array([_compute_local_gradient(dcoll, vec, xyz_axis)
for xyz_axis in range(dcoll.dim)])
@@ -290,15 +302,19 @@ def _div_helper(dcoll, diff_func, vecs):
raise TypeError("argument must be an object array")
assert vecs.dtype == object
- if vecs.shape[-1] != dcoll.ambient_dim:
- raise ValueError("last dimension of *vecs* argument must match "
- "ambient dimension")
+ if isinstance(vecs[(0,)*vecs.ndim], np.ndarray):
+ div_shape = vecs.shape
+ else:
+ if vecs.shape[-1] != dcoll.ambient_dim:
+ raise ValueError("last dimension of *vecs* argument doesn't match "
+ "ambient dimension")
+ div_shape = vecs.shape[:-1]
- if len(vecs.shape) == 1:
+ if len(div_shape) == 0:
return sum(diff_func(i, vec_i) for i, vec_i in enumerate(vecs))
else:
- result = np.zeros(vecs.shape[:-1], dtype=object)
- for idx in np.ndindex(vecs.shape[:-1]):
+ result = np.zeros(div_shape, dtype=object)
+ for idx in np.ndindex(div_shape):
result[idx] = sum(
diff_func(i, vec_i) for i, vec_i in enumerate(vecs[idx]))
return result
@@ -407,7 +423,7 @@ def _apply_stiffness_transpose_operator(dcoll, dd_out, dd_in, vec, xyz_axis):
)
-def weak_local_grad(dcoll, *args):
+def weak_local_grad(dcoll, *args, nested=False):
r"""Return the element-local weak gradient of the volume function
represented by *vec*.
@@ -416,8 +432,12 @@ def weak_local_grad(dcoll, *args):
:arg dcoll: a :class:`grudge.discretization.DiscretizationCollection`.
:arg dd: a :class:`~grudge.dof_desc.DOFDesc`, or a value convertible to one.
Defaults to the base volume discretization if not provided.
- :arg vec: a :class:`~meshmode.dof_array.DOFArray`
- :returns: an object array of :class:`~meshmode.dof_array.DOFArray`\ s
+ :arg vec: a :class:`~meshmode.dof_array.DOFArray` or object array of
+ `~meshmode.dof_array.DOFArray`
+ :arg nested: return nested object arrays instead of a single multidimensional
+ array if *vec* is non-scalar
+ :returns: an object array (possibly nested) of
+ :class:`~meshmode.dof_array.DOFArray`\ s
"""
if len(args) == 1:
vec, = args
@@ -427,6 +447,14 @@ def weak_local_grad(dcoll, *args):
else:
raise TypeError("invalid number of arguments")
+ if isinstance(vec, np.ndarray):
+ grad = obj_array_vectorize(
+ lambda el: weak_local_grad(dcoll, dd, el, nested=nested), vec)
+ if nested:
+ return grad
+ else:
+ return np.stack(grad, axis=0)
+
return make_obj_array(
[_apply_stiffness_transpose_operator(dcoll,
dof_desc.DD_VOLUME,
diff --git a/test/test_op.py b/test/test_op.py
new file mode 100644
index 0000000000000000000000000000000000000000..d74cdfe6ce7a7682dfb12aca506d68f53807199c
--- /dev/null
+++ b/test/test_op.py
@@ -0,0 +1,286 @@
+__copyright__ = "Copyright (C) 2021 University of Illinois Board of Trustees"
+
+__license__ = """
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
+"""
+
+
+import numpy as np
+
+import meshmode.mesh.generation as mgen
+from meshmode.dof_array import thaw
+
+from pytools.obj_array import make_obj_array
+
+from grudge import op, DiscretizationCollection
+from grudge.dof_desc import DOFDesc
+
+import pytest
+from meshmode.array_context import ( # noqa
+ pytest_generate_tests_for_pyopencl_array_context
+ as pytest_generate_tests)
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+# {{{ gradient
+
+@pytest.mark.parametrize("form", ["strong", "weak"])
+@pytest.mark.parametrize("dim", [1, 2, 3])
+@pytest.mark.parametrize("order", [2, 3])
+@pytest.mark.parametrize(("vectorize", "nested"), [
+ (False, False),
+ (True, False),
+ (True, True)
+ ])
+def test_gradient(actx_factory, form, dim, order, vectorize, nested,
+ visualize=False):
+ actx = actx_factory()
+
+ from pytools.convergence import EOCRecorder
+ eoc_rec = EOCRecorder()
+
+ for n in [4, 6, 8]:
+ mesh = mgen.generate_regular_rect_mesh(
+ a=(-1,)*dim, b=(1,)*dim,
+ nelements_per_axis=(n,)*dim)
+
+ dcoll = DiscretizationCollection(actx, mesh, order=order)
+
+ def f(x):
+ result = dcoll.zeros(actx) + 1
+ for i in range(dim-1):
+ result *= actx.np.sin(np.pi*x[i])
+ result *= actx.np.cos(np.pi/2*x[dim-1])
+ return result
+
+ def grad_f(x):
+ result = make_obj_array([dcoll.zeros(actx) + 1 for _ in range(dim)])
+ for i in range(dim-1):
+ for j in range(i):
+ result[i] *= actx.np.sin(np.pi*x[j])
+ result[i] *= np.pi*actx.np.cos(np.pi*x[i])
+ for j in range(i+1, dim-1):
+ result[i] *= actx.np.sin(np.pi*x[j])
+ result[i] *= actx.np.cos(np.pi/2*x[dim-1])
+ for j in range(dim-1):
+ result[dim-1] *= actx.np.sin(np.pi*x[j])
+ result[dim-1] *= -np.pi/2*actx.np.sin(np.pi/2*x[dim-1])
+ return result
+
+ x = thaw(actx, op.nodes(dcoll))
+
+ if vectorize:
+ u = make_obj_array([(i+1)*f(x) for i in range(dim)])
+ else:
+ u = f(x)
+
+ def get_flux(u_tpair):
+ dd = u_tpair.dd
+ dd_allfaces = dd.with_dtag("all_faces")
+ normal = thaw(actx, op.normal(dcoll, dd))
+ u_avg = u_tpair.avg
+ if vectorize:
+ if nested:
+ flux = make_obj_array([u_avg_i * normal for u_avg_i in u_avg])
+ else:
+ flux = np.outer(u_avg, normal)
+ else:
+ flux = u_avg * normal
+ return op.project(dcoll, dd, dd_allfaces, flux)
+
+ dd_allfaces = DOFDesc("all_faces")
+
+ if form == "strong":
+ grad_u = (
+ op.local_grad(dcoll, u, nested=nested)
+ # No flux terms because u doesn't have inter-el jumps
+ )
+ elif form == "weak":
+ grad_u = op.inverse_mass(dcoll,
+ -op.weak_local_grad(dcoll, u, nested=nested) # pylint: disable=E1130
+ + # noqa: W504
+ op.face_mass(dcoll,
+ dd_allfaces,
+ # Note: no boundary flux terms here because u_ext == u_int == 0
+ get_flux(op.interior_trace_pair(dcoll, u)))
+ )
+ else:
+ raise ValueError("Invalid form argument.")
+
+ if vectorize:
+ expected_grad_u = make_obj_array(
+ [(i+1)*grad_f(x) for i in range(dim)])
+ if not nested:
+ expected_grad_u = np.stack(expected_grad_u, axis=0)
+ else:
+ expected_grad_u = grad_f(x)
+
+ if visualize:
+ from grudge.shortcuts import make_visualizer
+ vis = make_visualizer(dcoll, vis_order=order if dim == 3 else dim+3)
+
+ filename = (f"test_gradient_{form}_{dim}_{order}"
+ f"{'_vec' if vectorize else ''}{'_nested' if nested else ''}.vtu")
+ vis.write_vtk_file(filename, [
+ ("u", u),
+ ("grad_u", grad_u),
+ ("expected_grad_u", expected_grad_u),
+ ], overwrite=True)
+
+ rel_linf_err = (
+ op.norm(dcoll, grad_u - expected_grad_u, np.inf)
+ / op.norm(dcoll, expected_grad_u, np.inf))
+ eoc_rec.add_data_point(1./n, rel_linf_err)
+
+ print("L^inf error:")
+ print(eoc_rec)
+ assert(eoc_rec.order_estimate() >= order - 0.5
+ or eoc_rec.max_error() < 1e-11)
+
+# }}}
+
+
+# {{{ divergence
+
+@pytest.mark.parametrize("form", ["strong", "weak"])
+@pytest.mark.parametrize("dim", [1, 2, 3])
+@pytest.mark.parametrize("order", [2, 3])
+@pytest.mark.parametrize(("vectorize", "nested"), [
+ (False, False),
+ (True, False),
+ (True, True)
+ ])
+def test_divergence(actx_factory, form, dim, order, vectorize, nested,
+ visualize=False):
+ actx = actx_factory()
+
+ from pytools.convergence import EOCRecorder
+ eoc_rec = EOCRecorder()
+
+ for n in [4, 6, 8]:
+ mesh = mgen.generate_regular_rect_mesh(
+ a=(-1,)*dim, b=(1,)*dim,
+ nelements_per_axis=(n,)*dim)
+
+ dcoll = DiscretizationCollection(actx, mesh, order=order)
+
+ def f(x):
+ result = make_obj_array([dcoll.zeros(actx) + (i+1) for i in range(dim)])
+ for i in range(dim-1):
+ result *= actx.np.sin(np.pi*x[i])
+ result *= actx.np.cos(np.pi/2*x[dim-1])
+ return result
+
+ def div_f(x):
+ result = dcoll.zeros(actx)
+ for i in range(dim-1):
+ deriv = dcoll.zeros(actx) + (i+1)
+ for j in range(i):
+ deriv *= actx.np.sin(np.pi*x[j])
+ deriv *= np.pi*actx.np.cos(np.pi*x[i])
+ for j in range(i+1, dim-1):
+ deriv *= actx.np.sin(np.pi*x[j])
+ deriv *= actx.np.cos(np.pi/2*x[dim-1])
+ result += deriv
+ deriv = dcoll.zeros(actx) + dim
+ for j in range(dim-1):
+ deriv *= actx.np.sin(np.pi*x[j])
+ deriv *= -np.pi/2*actx.np.sin(np.pi/2*x[dim-1])
+ result += deriv
+ return result
+
+ x = thaw(actx, op.nodes(dcoll))
+
+ if vectorize:
+ u = make_obj_array([(i+1)*f(x) for i in range(dim)])
+ if not nested:
+ u = np.stack(u, axis=0)
+ else:
+ u = f(x)
+
+ def get_flux(u_tpair):
+ dd = u_tpair.dd
+ dd_allfaces = dd.with_dtag("all_faces")
+ normal = thaw(actx, op.normal(dcoll, dd))
+ flux = u_tpair.avg @ normal
+ return op.project(dcoll, dd, dd_allfaces, flux)
+
+ dd_allfaces = DOFDesc("all_faces")
+
+ if form == "strong":
+ div_u = (
+ op.local_div(dcoll, u)
+ # No flux terms because u doesn't have inter-el jumps
+ )
+ elif form == "weak":
+ div_u = op.inverse_mass(dcoll,
+ -op.weak_local_div(dcoll, u)
+ + # noqa: W504
+ op.face_mass(dcoll,
+ dd_allfaces,
+ # Note: no boundary flux terms here because u_ext == u_int == 0
+ get_flux(op.interior_trace_pair(dcoll, u)))
+ )
+ else:
+ raise ValueError("Invalid form argument.")
+
+ if vectorize:
+ expected_div_u = make_obj_array([(i+1)*div_f(x) for i in range(dim)])
+ else:
+ expected_div_u = div_f(x)
+
+ if visualize:
+ from grudge.shortcuts import make_visualizer
+ vis = make_visualizer(dcoll, vis_order=order if dim == 3 else dim+3)
+
+ filename = (f"test_divergence_{form}_{dim}_{order}"
+ f"{'_vec' if vectorize else ''}{'_nested' if nested else ''}.vtu")
+ vis.write_vtk_file(filename, [
+ ("u", u),
+ ("div_u", div_u),
+ ("expected_div_u", expected_div_u),
+ ], overwrite=True)
+
+ rel_linf_err = (
+ op.norm(dcoll, div_u - expected_div_u, np.inf)
+ / op.norm(dcoll, expected_div_u, np.inf))
+ eoc_rec.add_data_point(1./n, rel_linf_err)
+
+ print("L^inf error:")
+ print(eoc_rec)
+ assert(eoc_rec.order_estimate() >= order - 0.5
+ or eoc_rec.max_error() < 1e-11)
+
+# }}}
+
+
+# You can test individual routines by typing
+# $ python test_grudge.py 'test_routine()'
+
+if __name__ == "__main__":
+ import sys
+ if len(sys.argv) > 1:
+ exec(sys.argv[1])
+ else:
+ pytest.main([__file__])
+
+# vim: fdm=marker