diff --git a/grudge/op.py b/grudge/op.py index f9e762f14..9c17e5cc4 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -93,6 +93,7 @@ from grudge.dof_desc import ( DD_VOLUME_ALL, DISCR_TAG_BASE, + DISCR_TAG_QUAD, FACE_RESTR_ALL, DOFDesc, VolumeDomainTag, @@ -819,6 +820,68 @@ def _apply_inverse_mass_operator( return DOFArray(actx, data=tuple(group_data)) +def _apply_inverse_mass_operator_quad( + dcoll: DiscretizationCollection, dd, vec): + if not isinstance(vec, DOFArray): + return map_array_container( + partial(_apply_inverse_mass_operator_quad, dcoll, dd), vec + ) + + from grudge.geometry import area_element + + actx = vec.array_context + dd_quad = dd.with_discr_tag(DISCR_TAG_QUAD) + dd_base = dd.with_discr_tag(DISCR_TAG_BASE) + discr_quad = dcoll.discr_from_dd(dd_quad) + discr_base = dcoll.discr_from_dd(dd_base) + + # Based on https://arxiv.org/pdf/1608.03836.pdf + # true_Minv ~ ref_Minv * ref_M * (1/jac_det) * ref_Minv + # Overintegration version of action on *vec*: + # true_Minv ~ ref_Minv * (ref_M)_qtb * (1/Jac)_quad * P(Minv*vec) + # P => projection to quadrature, qti => quad-to-base + + # Compute 1/Jac on quadrature discr + inv_area_elements = 1/area_element( + actx, dcoll, dd=dd_quad, + _use_geoderiv_connection=actx.supports_nonscalar_broadcasting) + + def apply_minv_to_vec(vec, ref_inv_mass): + return actx.einsum( + "ij,ej->ei", + ref_inv_mass, + vec, + tagged=(FirstAxisIsElementsTag(),)) + + # The rest of wadg + def apply_rest_of_wadg(mm_inv, mm, vec): + return actx.einsum( + "ni,ij,ej->en", + mm_inv, + mm, + vec, + tagged=(FirstAxisIsElementsTag(),)) + + stage1_group_data = [ + apply_minv_to_vec( + vec_i, reference_inverse_mass_matrix(actx, element_group=grp)) + for grp, vec_i in zip(discr_base.groups, vec) + ] + stage2 = inv_area_elements * project( + dcoll, dd_base, dd_quad, + DOFArray(actx, data=tuple(stage1_group_data))) + + wadg_group_data = [ + apply_rest_of_wadg( + reference_inverse_mass_matrix(actx, out_grp), + reference_mass_matrix(actx, out_grp, in_grp), vec_i) + for in_grp, out_grp, vec_i in zip( + discr_quad.groups, discr_base.groups, stage2) + ] + + return DOFArray(actx, data=tuple(wadg_group_data)) + + def inverse_mass(dcoll: DiscretizationCollection, *args) -> ArrayOrContainer: r"""Return the action of the DG mass matrix inverse on a vector (or vectors) of :class:`~meshmode.dof_array.DOFArray`\ s, *vec*. @@ -866,6 +929,9 @@ def inverse_mass(dcoll: DiscretizationCollection, *args) -> ArrayOrContainer: else: raise TypeError("invalid number of arguments") + if dd.uses_quadrature(): + return _apply_inverse_mass_operator_quad(dcoll, dd, vec) + return _apply_inverse_mass_operator(dcoll, dd, dd, vec) # }}} diff --git a/test/test_grudge.py b/test/test_grudge.py index 8a5a5a678..57710c0f8 100644 --- a/test/test_grudge.py +++ b/test/test_grudge.py @@ -301,12 +301,14 @@ def f(x): ) else: dd_base_vol = dof_desc.as_dofdesc( - dof_desc.DTAG_VOLUME_ALL, dof_desc.DISCR_TAG_BASE) + dof_desc.DTAG_VOLUME_ALL, dof_desc.DISCR_TAG_BASE) dd_quad_vol = dof_desc.as_dofdesc( - dof_desc.DTAG_VOLUME_ALL, dof_desc.DISCR_TAG_QUAD) + dof_desc.DTAG_VOLUME_ALL, dof_desc.DISCR_TAG_QUAD) + # Uses _apply_inverse_mass_operator_quad if overintegrate f_inv = op.inverse_mass( - dcoll, op.mass(dcoll, dd_quad_vol, - op.project(dcoll, dd_base_vol, dd_quad_vol, f_volm))) + dcoll, dd_quad_vol, op.mass( + dcoll, dd_quad_vol, + op.project(dcoll, dd_base_vol, dd_quad_vol, f_volm))) inv_error = actx.to_numpy( op.norm(dcoll, f_volm - f_inv, 2) / op.norm(dcoll, f_volm, 2)) diff --git a/test/test_op.py b/test/test_op.py index 17f49a074..64d5b6ca5 100644 --- a/test/test_op.py +++ b/test/test_op.py @@ -32,7 +32,7 @@ InterpolatoryEdgeClusteredGroupFactory, QuadratureGroupFactory, ) -from meshmode.mesh import BTAG_ALL +from meshmode.mesh import BTAG_ALL, TensorProductElementGroup from pytools.obj_array import make_obj_array from grudge import geometry, op @@ -61,14 +61,26 @@ @pytest.mark.parametrize("dim", [1, 2, 3]) @pytest.mark.parametrize("order", [2, 3]) @pytest.mark.parametrize("warp_mesh", [False, True]) +@pytest.mark.parametrize("tpe", [False, True]) @pytest.mark.parametrize(("vectorize", "nested"), [ (False, False), (True, False), (True, True) ]) -def test_gradient(actx_factory, form, dim, order, vectorize, nested, +def test_gradient(actx_factory, form, dim, order, tpe, vectorize, nested, warp_mesh, visualize=False): actx = actx_factory() + group_cls = None + if tpe: + if dim == 1: + pytest.skip("Skipping 1D test for tensor product elements.") + group_cls = TensorProductElementGroup + + def vectorize_if_requested(vec): + if vectorize: + return make_obj_array([(i+1)*vec for i in range(dim)]) + else: + return vec from pytools.convergence import EOCRecorder eoc_rec = EOCRecorder() @@ -78,19 +90,36 @@ def test_gradient(actx_factory, form, dim, order, vectorize, nested, if dim == 1: pytest.skip("warped mesh in 1D not implemented") mesh = mgen.generate_warped_rect_mesh( - dim=dim, order=order, nelements_side=n) + dim=dim, order=order, nelements_side=n, + group_cls=group_cls) else: mesh = mgen.generate_regular_rect_mesh( - a=(-1,)*dim, b=(1,)*dim, - nelements_per_axis=(n,)*dim) + a=(-1,)*dim, b=(1,)*dim, + nelements_per_axis=(n,)*dim, + group_cls=group_cls) dcoll = make_discretization_collection( actx, mesh, discr_tag_to_group_factory={ - DISCR_TAG_BASE: InterpolatoryEdgeClusteredGroupFactory(order), + DISCR_TAG_BASE: + InterpolatoryEdgeClusteredGroupFactory(order), DISCR_TAG_QUAD: QuadratureGroupFactory(3 * order) }) + if "overint" in form: + quad_discr_tag = DISCR_TAG_QUAD + else: + quad_discr_tag = DISCR_TAG_BASE + + allfaces_dd_quad = as_dofdesc(FACE_RESTR_ALL, quad_discr_tag) + vol_dd_base = as_dofdesc(DTAG_VOLUME_ALL) + vol_dd_quad = vol_dd_base.with_discr_tag(quad_discr_tag) + bdry_dd_base = as_dofdesc(BTAG_ALL) + bdry_dd_quad = bdry_dd_base.with_discr_tag(quad_discr_tag) + + x_base = actx.thaw(dcoll.nodes(dd=vol_dd_base)) + bdry_x_base = actx.thaw(dcoll.nodes(bdry_dd_base)) + def f(x): result = 1 for i in range(dim-1): @@ -112,12 +141,6 @@ def grad_f(x): result[dim-1] = result[dim-1] * (-np.pi/2*actx.np.sin(np.pi/2*x[dim-1])) return result - def vectorize_if_requested(vec): - if vectorize: - return make_obj_array([(i+1)*vec for i in range(dim)]) - else: - return vec - def get_flux(u_tpair, dcoll=dcoll): dd = u_tpair.dd dd_allfaces = dd.with_domain_tag( @@ -134,57 +157,47 @@ def get_flux(u_tpair, dcoll=dcoll): flux = u_avg * normal return op.project(dcoll, dd, dd_allfaces, flux) - x = actx.thaw(dcoll.nodes()) - u = vectorize_if_requested(f(x)) - - bdry_dd_base = as_dofdesc(BTAG_ALL) - bdry_x = actx.thaw(dcoll.nodes(bdry_dd_base)) - bdry_u = vectorize_if_requested(f(bdry_x)) + u_base = vectorize_if_requested(f(x_base)) + bdry_u_base = vectorize_if_requested(f(bdry_x_base)) if form == "strong": grad_u = ( - op.local_grad(dcoll, u, nested=nested) + op.local_grad(dcoll, u_base, nested=nested) # No flux terms because u doesn't have inter-el jumps ) elif form.startswith("weak"): assert form in ["weak", "weak-overint"] - if "overint" in form: - quad_discr_tag = DISCR_TAG_QUAD - else: - quad_discr_tag = DISCR_TAG_BASE - - allfaces_dd_base = as_dofdesc(FACE_RESTR_ALL, quad_discr_tag) - vol_dd_base = as_dofdesc(DTAG_VOLUME_ALL) - vol_dd_quad = vol_dd_base.with_discr_tag(quad_discr_tag) - bdry_dd_quad = bdry_dd_base.with_discr_tag(quad_discr_tag) - allfaces_dd_quad = allfaces_dd_base.with_discr_tag(quad_discr_tag) - - grad_u = op.inverse_mass(dcoll, - -op.weak_local_grad(dcoll, vol_dd_quad, - op.project(dcoll, vol_dd_base, vol_dd_quad, u), - nested=nested) + + # Uses _apply_inverse_mass_operator_quad (if overint) + grad_u = op.inverse_mass( + dcoll, vol_dd_quad, + -op.weak_local_grad( + dcoll, vol_dd_quad, + op.project(dcoll, vol_dd_base, vol_dd_quad, u_base), + nested=nested) + - op.face_mass(dcoll, - allfaces_dd_quad, - sum(get_flux( - op.project_tracepair(dcoll, allfaces_dd_quad, utpair)) - for utpair in op.interior_trace_pairs( - dcoll, u, volume_dd=vol_dd_base)) - + get_flux( - op.project_tracepair(dcoll, bdry_dd_quad, - bv_trace_pair(dcoll, bdry_dd_base, u, bdry_u))) - ) + op.face_mass(dcoll, allfaces_dd_quad, + sum(get_flux( + op.project_tracepair(dcoll, allfaces_dd_quad, + utpair)) + for utpair in op.interior_trace_pairs( + dcoll, u_base, volume_dd=vol_dd_base)) + + get_flux( + op.project_tracepair( + dcoll, bdry_dd_quad, + bv_trace_pair(dcoll, bdry_dd_base, u_base, + bdry_u_base)))) ) else: raise ValueError("Invalid form argument.") if vectorize: expected_grad_u = make_obj_array( - [(i+1)*grad_f(x) for i in range(dim)]) + [(i+1)*grad_f(x_base) 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) + expected_grad_u = grad_f(x_base) if visualize: # the code below does not handle the vectorized case @@ -196,7 +209,7 @@ def get_flux(u_tpair, dcoll=dcoll): filename = (f"test_gradient_{form}_{dim}_{order}" f"{'_vec' if vectorize else ''}{'_nested' if nested else ''}.vtu") vis.write_vtk_file(filename, [ - ("u", u), + ("u", u_base), ("grad_u", grad_u), ("expected_grad_u", expected_grad_u), ], overwrite=True) @@ -224,31 +237,55 @@ def get_flux(u_tpair, dcoll=dcoll): (True, False), (True, True) ]) +@pytest.mark.parametrize("overint", [True, False]) +@pytest.mark.parametrize("tpe", [False, True]) def test_divergence(actx_factory, form, dim, order, vectorize, nested, - visualize=False): + overint, tpe, visualize=False): actx = actx_factory() + if dim == 1 and tpe: + pytest.skip("Skipping 1D test for tensor product elements.") from pytools.convergence import EOCRecorder eoc_rec = EOCRecorder() for n in [4, 6, 8]: + group_cls = TensorProductElementGroup if tpe else None mesh = mgen.generate_regular_rect_mesh( - a=(-1,)*dim, b=(1,)*dim, - nelements_per_axis=(n,)*dim) + a=(-1,)*dim, b=(1,)*dim, + nelements_per_axis=(n,)*dim, + group_cls=group_cls) + + dcoll = make_discretization_collection( + actx, mesh, discr_tag_to_group_factory={ + DISCR_TAG_BASE: + InterpolatoryEdgeClusteredGroupFactory(order), + DISCR_TAG_QUAD: QuadratureGroupFactory(3 * order) + }) + + if overint: + quad_discr_tag = DISCR_TAG_QUAD + else: + quad_discr_tag = DISCR_TAG_BASE + + allfaces_dd_quad = as_dofdesc(FACE_RESTR_ALL, quad_discr_tag) + vol_dd_base = as_dofdesc(DTAG_VOLUME_ALL) + vol_dd_quad = vol_dd_base.with_discr_tag(quad_discr_tag) - dcoll = make_discretization_collection(actx, mesh, order=order) + x_base = actx.thaw(dcoll.nodes()) def f(x, dcoll=dcoll): - result = make_obj_array([dcoll.zeros(actx) + (i+1) for i in range(dim)]) + zeros = 0 * x[0] + result = make_obj_array([zeros + (i+1) for i in range(dim)]) for i in range(dim-1): result = result * actx.np.sin(np.pi*x[i]) result = result * actx.np.cos(np.pi/2*x[dim-1]) return result def div_f(x, dcoll=dcoll): - result = dcoll.zeros(actx) + zeros = 0*x[0] + result = 1*zeros for i in range(dim-1): - deriv = dcoll.zeros(actx) + (i+1) + deriv = 1*zeros + (i+1) for j in range(i): deriv = deriv * actx.np.sin(np.pi*x[j]) deriv = deriv * np.pi*actx.np.cos(np.pi*x[i]) @@ -256,21 +293,19 @@ def div_f(x, dcoll=dcoll): deriv = deriv * actx.np.sin(np.pi*x[j]) deriv = deriv * actx.np.cos(np.pi/2*x[dim-1]) result = result + deriv - deriv = dcoll.zeros(actx) + dim + deriv = 1*zeros + dim for j in range(dim-1): deriv = deriv * actx.np.sin(np.pi*x[j]) deriv = deriv * (-np.pi/2*actx.np.sin(np.pi/2*x[dim-1])) result = result + deriv return result - x = actx.thaw(dcoll.nodes()) - if vectorize: - u = make_obj_array([(i+1)*f(x) for i in range(dim)]) + u_base = make_obj_array([(i+1)*f(x_base) for i in range(dim)]) if not nested: - u = np.stack(u, axis=0) + u_base = np.stack(u_base, axis=0) else: - u = f(x) + u_base = f(x_base) def get_flux(u_tpair, dcoll=dcoll): dd = u_tpair.dd @@ -281,40 +316,45 @@ def get_flux(u_tpair, dcoll=dcoll): flux = u_tpair.avg @ normal return op.project(dcoll, dd, dd_allfaces, flux) - dd_allfaces = as_dofdesc(FACE_RESTR_ALL) - if form == "strong": div_u = ( - op.local_div(dcoll, u) + op.local_div(dcoll, u_base) # 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) + div_u = op.inverse_mass( + dcoll, vol_dd_quad, + -op.weak_local_div( + dcoll, vol_dd_quad, + op.project(dcoll, vol_dd_base, vol_dd_quad, u_base)) + - op.face_mass(dcoll, - dd_allfaces, - # Note: no boundary flux terms here because u_ext == u_int == 0 - sum(get_flux(utpair) - for utpair in op.interior_trace_pairs(dcoll, u)) + op.face_mass( + dcoll, allfaces_dd_quad, + # Note: no boundary flux terms here because + # u_ext == u_int == 0 + sum(get_flux(op.project_tracepair( + dcoll, allfaces_dd_quad, utpair)) + for utpair in op.interior_trace_pairs(dcoll, u_base)) ) ) else: raise ValueError("Invalid form argument.") if vectorize: - expected_div_u = make_obj_array([(i+1)*div_f(x) for i in range(dim)]) + expected_div_u = make_obj_array([(i+1)*div_f(x_base) + for i in range(dim)]) else: - expected_div_u = div_f(x) + expected_div_u = div_f(x_base) if visualize: from grudge.shortcuts import make_visualizer - vis = make_visualizer(dcoll, vis_order=order if dim == 3 else dim+3) + 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), + ("u", u_base), ("div_u", div_u), ("expected_div_u", expected_div_u), ], overwrite=True)