diff --git a/examples/tensor-product-examples/acoustic_pulse.py b/examples/tensor-product-examples/acoustic_pulse.py new file mode 100644 index 000000000..9520945f3 --- /dev/null +++ b/examples/tensor-product-examples/acoustic_pulse.py @@ -0,0 +1,272 @@ +__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. +""" + + +from meshmode.mesh import TensorProductElementGroup +import numpy as np + +import pyopencl as cl +import pyopencl.tools as cl_tools + +from grudge.array_context import ( + PyOpenCLArrayContext, + PytatoPyOpenCLArrayContext +) +from grudge.models.euler import ( + ConservedEulerField, + EulerOperator, + InviscidWallBC +) +from grudge.shortcuts import rk4_step + +from meshmode.mesh import BTAG_ALL + +from pytools.obj_array import make_obj_array + +import grudge.op as op + +import logging +logger = logging.getLogger(__name__) + + +def gaussian_profile( + x_vec, t=0, rho0=1.0, rhoamp=1.0, p0=1.0, gamma=1.4, + center=None, velocity=None): + + dim = len(x_vec) + if center is None: + center = np.zeros(shape=(dim,)) + if velocity is None: + velocity = np.zeros(shape=(dim,)) + + lump_loc = center + t * velocity + + # coordinates relative to lump center + rel_center = make_obj_array( + [x_vec[i] - lump_loc[i] for i in range(dim)] + ) + actx = x_vec[0].array_context + r = actx.np.sqrt(np.dot(rel_center, rel_center)) + expterm = rhoamp * actx.np.exp(1 - r ** 2) + + mass = expterm + rho0 + mom = velocity * mass + energy = (p0 / (gamma - 1.0)) + np.dot(mom, mom) / (2.0 * mass) + + return ConservedEulerField(mass=mass, energy=energy, momentum=mom) + + +def make_pulse(amplitude, r0, w, r): + dim = len(r) + r_0 = np.zeros(dim) + r_0 = r_0 + r0 + rel_center = make_obj_array( + [r[i] - r_0[i] for i in range(dim)] + ) + actx = r[0].array_context + rms2 = w * w + r2 = np.dot(rel_center, rel_center) / rms2 + return amplitude * actx.np.exp(-.5 * r2) + + +def acoustic_pulse_condition(x_vec, t=0): + dim = len(x_vec) + vel = np.zeros(shape=(dim,)) + orig = np.zeros(shape=(dim,)) + uniform_gaussian = gaussian_profile( + x_vec, t=t, center=orig, velocity=vel, rhoamp=0.0) + + amplitude = 1.0 + width = 0.1 + pulse = make_pulse(amplitude, orig, width, x_vec) + + return ConservedEulerField( + mass=uniform_gaussian.mass, + energy=uniform_gaussian.energy + pulse, + momentum=uniform_gaussian.momentum + ) + + +def run_acoustic_pulse(actx, + order=3, + final_time=1, + resolution=16, + overintegration=False, + visualize=False, + rotate_mesh=False): + + # eos-related parameters + gamma = 1.4 + + # {{{ discretization + + from meshmode.mesh.generation import generate_regular_rect_mesh + + dim = 2 + box_ll = -0.5 + box_ur = 0.5 + mesh = generate_regular_rect_mesh( + a=(box_ll,)*dim, + b=(box_ur,)*dim, + nelements_per_axis=(resolution,)*dim, + group_cls=TensorProductElementGroup) + + if rotate_mesh: + from meshmode.mesh.processing import affine_map + alpha = .3 + rot_mat = np.array([ + [np.cos(alpha), np.sin(alpha)], + [-np.sin(alpha), np.cos(alpha)] + ]) + mesh = affine_map(mesh, A=rot_mat) + + from grudge import DiscretizationCollection + from grudge.dof_desc import DISCR_TAG_BASE, DISCR_TAG_QUAD + from meshmode.discretization.poly_element import \ + LegendreGaussLobattoTensorProductGroupFactory as LGL + + exp_name = f"fld-acoustic-pulse-N{order}-K{resolution}" + if overintegration: + exp_name += "-overintegrated" + quad_tag = DISCR_TAG_QUAD + else: + quad_tag = None + + dcoll = DiscretizationCollection( + actx, mesh, + discr_tag_to_group_factory={ + DISCR_TAG_BASE: LGL(order) + } + ) + + # }}} + + # {{{ Euler operator + + euler_operator = EulerOperator( + dcoll, + bdry_conditions={BTAG_ALL: InviscidWallBC()}, + flux_type="lf", + gamma=gamma, + quadrature_tag=quad_tag + ) + + def rhs(t, q): + return euler_operator.operator(t, q) + + compiled_rhs = actx.compile(rhs) + + from grudge.dt_utils import h_min_from_volume + + cfl = 0.125 + cn = 0.5*(order + 1)**2 + dt = cfl * actx.to_numpy(h_min_from_volume(dcoll)) / cn + + fields = acoustic_pulse_condition(actx.thaw(dcoll.nodes())) + + logger.info("Timestep size: %g", dt) + + # }}} + + from grudge.shortcuts import make_visualizer + + vis = make_visualizer(dcoll) + + # {{{ time stepping + + step = 0 + t = 0.0 + while t < final_time: + if step % 10 == 0: + norm_q = actx.to_numpy(op.norm(dcoll, fields, 2)) + logger.info("[%04d] t = %.5f |q| = %.5e", step, t, norm_q) + if visualize: + vis.write_vtk_file( + f"{exp_name}-{step:04d}.vtu", + [ + ("rho", fields.mass), + ("energy", fields.energy), + ("momentum", fields.momentum) + ] + ) + assert norm_q < 5 + + fields = actx.thaw(actx.freeze(fields)) + fields = rk4_step(fields, t, dt, compiled_rhs) + t += dt + step += 1 + + # }}} + + +def main(ctx_factory, order=3, final_time=1, resolution=16, + overintegration=False, visualize=False, lazy=False): + cl_ctx = ctx_factory() + queue = cl.CommandQueue(cl_ctx) + + if lazy: + actx = PytatoPyOpenCLArrayContext( + queue, + allocator=cl_tools.MemoryPool(cl_tools.ImmediateAllocator(queue)), + ) + else: + actx = PyOpenCLArrayContext( + queue, + allocator=cl_tools.MemoryPool(cl_tools.ImmediateAllocator(queue)), + force_device_scalars=False + ) + + run_acoustic_pulse( + actx, + order=order, + resolution=resolution, + overintegration=overintegration, + final_time=final_time, + visualize=visualize + ) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument("--order", default=3, type=int) + parser.add_argument("--tfinal", default=0.1, type=float) + parser.add_argument("--resolution", default=16, type=int) + parser.add_argument("--oi", action="store_true", + help="use overintegration") + parser.add_argument("--visualize", action="store_true", + help="write out vtk output") + parser.add_argument("--lazy", action="store_true", + help="switch to a lazy computation mode") + args = parser.parse_args() + + logging.basicConfig(level=logging.INFO) + main(cl.create_some_context, + order=args.order, + final_time=args.tfinal, + resolution=args.resolution, + overintegration=args.oi, + visualize=args.visualize, + lazy=args.lazy) diff --git a/examples/tp-transform-cartoon.py b/examples/tp-transform-cartoon.py new file mode 100644 index 000000000..f26a58737 --- /dev/null +++ b/examples/tp-transform-cartoon.py @@ -0,0 +1,102 @@ +import loopy as lp + +import meshmode.mesh.generation as mgen + +import numpy as np +import pyopencl as cl +import pytato as pt + +from grudge import op +from grudge.array_context import OutputIsTensorProductDOFArrayOrdered +from grudge.discretization import make_discretization_collection + +from meshmode.array_context import PytatoPyOpenCLArrayContext + + +class PytatoTensorProductArrayContext(PytatoPyOpenCLArrayContext): + def transform_dag(self, dag): + if "dag_dots" not in dir(self): + self.dag_dots = [] + + self.dag_dots.append(pt.get_dot_graph(dag)) + + return super().transform_dag(dag) + + def transform_loopy_program(self, t_unit): + knl = t_unit.default_entrypoint + + # {{{ adjust strides according to tensor product structure + if knl.tags_of_type(OutputIsTensorProductDOFArrayOrdered): + new_args = [] + for arg in knl.args: + if arg.is_output: + arg = arg.copy(dim_tags=( + f"N{len(arg.shape)-1}," + + ",".join(f"N{i}" + for i in range(len(arg.shape)-1)) + )) + + new_args.append(arg) + + knl = knl.copy(args=new_args) + # }}} + + # {{{ prefetch + # }}} + + # {{{ tile + # }}} + + # FIXME: remove this (eventually) + knl = lp.set_options(knl, insert_gbarriers=True) + t_unit = t_unit.with_kernel(knl) + self.dev_code = lp.generate_code_v2(t_unit).device_code() + + return super().transform_loopy_program(t_unit) + + +def main(): + order = 1 + + ctx = cl.create_some_context() + queue = cl.CommandQueue(ctx) + actx = PytatoTensorProductArrayContext(queue) + + dim = 2 + res = 2 + + from meshmode.mesh import TensorProductElementGroup + from meshmode.discretization.poly_element import \ + LegendreGaussLobattoTensorProductGroupFactory as LGL + + mesh = mgen.generate_regular_rect_mesh( + a=(-1,)*dim, b=(1,)*dim, + nelements_per_axis=(res,)*dim, + group_cls=TensorProductElementGroup) + + import grudge.dof_desc as dd + dcoll = make_discretization_collection( + actx, + mesh, + discr_tag_to_group_factory={ + dd.DISCR_TAG_BASE: LGL(order)}) + + def f(x): + result = dcoll.zeros(actx) + 1 + 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 + + + x = actx.thaw(dcoll.nodes()) + + u = f(x) + + grad_u = op.local_grad(dcoll, u) + grad_u = actx.np.stack(grad_u)[0] + pt.show_dot_graph(grad_u) + +if __name__ == "__main__": + main() + diff --git a/grudge/array_context.py b/grudge/array_context.py index c5672178d..0ad5b038c 100644 --- a/grudge/array_context.py +++ b/grudge/array_context.py @@ -36,6 +36,10 @@ FrozenSet) from dataclasses import dataclass from pytools.tag import Tag + +from grudge.transform.metadata import OutputIsTensorProductDOFArrayOrdered + +from meshmode.transform_metadata import DiscretizationDOFAxisTag from meshmode.array_context import ( PyOpenCLArrayContext as _PyOpenCLArrayContextBase, PytatoPyOpenCLArrayContext as _PytatoPyOpenCLArrayContextBase) @@ -105,11 +109,14 @@ import pyopencl.tools from mpi4py import MPI +# }}} + + +# {{{ pyopencl class PyOpenCLArrayContext(_PyOpenCLArrayContextBase): """Inherits from :class:`meshmode.array_context.PyOpenCLArrayContext`. Extends it - to understand :mod:`grudge`-specific transform metadata. (Of which there isn't - any, for now.) + to understand :mod:`grudge`-specific transform metadata. """ def __init__(self, queue: "pyopencl.CommandQueue", allocator: Optional["pyopencl.tools.AllocatorBase"] = None, @@ -124,6 +131,33 @@ def __init__(self, queue: "pyopencl.CommandQueue", super().__init__(queue, allocator, wait_event_queue_length, force_device_scalars) + def transform_loopy_program(self, t_unit): + knl = t_unit.default_entrypoint + + # {{{ process tensor product specific metadata + + # NOTE: This differs from the lazy case b/c we don't have access to axis + # tags that can be manipulated pre-execution. In eager, we update + # strides/loop nest ordering for the output array + if knl.tags_of_type(OutputIsTensorProductDOFArrayOrdered): + new_args = [] + for arg in knl.args: + if arg.is_output: + arg = arg.copy(dim_tags=( + f"N{len(arg.shape)-1}," + + ",".join(f"N{i}" + for i in range(len(arg.shape)-1)) + )) + + new_args.append(arg) + + knl = knl.copy(args=new_args) + t_unit = t_unit.with_kernel(knl) + + # }}} + + return super().transform_loopy_program(t_unit) + # }}} @@ -131,8 +165,7 @@ def __init__(self, queue: "pyopencl.CommandQueue", class PytatoPyOpenCLArrayContext(_PytatoPyOpenCLArrayContextBase): """Inherits from :class:`meshmode.array_context.PytatoPyOpenCLArrayContext`. - Extends it to understand :mod:`grudge`-specific transform metadata. (Of - which there isn't any, for now.) + Extends it to understand :mod:`grudge`-specific transform metadata. """ def __init__(self, queue, allocator=None, *, diff --git a/grudge/geometry/metrics.py b/grudge/geometry/metrics.py index 260f219b3..c5eff00a1 100644 --- a/grudge/geometry/metrics.py +++ b/grudge/geometry/metrics.py @@ -527,6 +527,7 @@ def inverse_surface_metric_derivative_mat( @memoize_in(dcoll, (inverse_surface_metric_derivative_mat, dd, times_area_element, _use_geoderiv_connection)) def _inv_surf_metric_deriv(): + if times_area_element: multiplier = area_element(actx, dcoll, dd=dd, _use_geoderiv_connection=_use_geoderiv_connection) diff --git a/grudge/op.py b/grudge/op.py index f5781f4be..fa9cad155 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -75,13 +75,29 @@ from functools import partial from meshmode.dof_array import DOFArray -from meshmode.transform_metadata import (FirstAxisIsElementsTag, - DiscretizationDOFAxisTag, - DiscretizationElementAxisTag, - DiscretizationFaceAxisTag) +from meshmode.discretization.poly_element import ( + TensorProductElementGroupBase as TensorProductElementGroup, + SimplexElementGroupBase as SimplexElementGroup) +from meshmode.transform_metadata import ( + DiscretizationAmbientDimAxisTag, + FirstAxisIsElementsTag, + DiscretizationDOFAxisTag, + DiscretizationElementAxisTag, + DiscretizationFaceAxisTag) + +from modepy.tools import ( + reshape_array_for_tensor_product_space as fold, + unreshape_array_for_tensor_product_space as unfold) from grudge.discretization import DiscretizationCollection from grudge.dof_desc import as_dofdesc +from grudge.transform.metadata import ( + OutputIsTensorProductDOFArrayOrdered, + TensorProductDOFAxisTag, + TensorProductOperatorAxisTag, + ReferenceTensorProductMassOperatorTag as MassMatrix1DTag, + ReferenceTensorProductInverseMassOperatorTag as InverseMassMatrix1DTag +) from pytools import keyed_memoize_in from pytools.obj_array import make_obj_array @@ -168,6 +184,8 @@ # {{{ common derivative "kernels" +# {{{ single axis derivative + def _single_axis_derivative_kernel( actx, out_discr, in_discr, get_diff_mat, inv_jac_mat, xyz_axis, vec, *, metric_in_matvec): @@ -179,77 +197,353 @@ def _single_axis_derivative_kernel( # - whether the chain rule terms ("inv_jac_mat") sit outside (strong) # or inside (weak) the matrix-vector product that carries out the # derivative, cf. "metric_in_matvec". + + + # {{{ tensor product single axis derivative + + def compute_tensor_product_derivative(actx, grp, get_diff_mat, vec, ijm, + xyz_axis, metric_in_matvec): + + vec = fold(grp.space, vec) + + if metric_in_matvec: + stiff_1d, mass_1d = get_diff_mat(actx, grp, grp) + + apply_mass_axes = set(range(grp.dim)) - {xyz_axis} + + for ax in apply_mass_axes: + vec_mass_applied = single_axis_operator_application( + actx, grp.dim, mass_1d, ax, vec, + tagged=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered(),), + arg_names=("mass_1d", "vec") + ) + + ref_weak_derivative = unfold( + grp.space, + single_axis_operator_application( + actx, grp.dim, stiff_1d, xyz_axis, vec_mass_applied, + tagged=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered(),), + arg_names=("stiff_1d", "vec_with_mass_applied")) + ) + + derivative = actx.einsum( + 'rej,ej->ej', + ijm[xyz_axis], + ref_weak_derivative, + tagged=(FirstAxisIsElementsTag(),), + arg_names=("inv_jac_t", "ref_weak_derivative") + ) + + else: + diff_mat = get_diff_mat(actx, grp, grp) + + ref_derivative = unfold( + grp.space, + single_axis_operator_application( + actx, grp.dim, diff_mat, xyz_axis, vec, + tagged=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered(),), + arg_names=("diff_mat", "vec")) + ) + + derivative = actx.einsum( + 'rej,ej->ej', + ijm[xyz_axis], + ref_derivative, + tagged=(FirstAxisIsElementsTag(),), + arg_names=("inv_jac_t", "ref_derivs") + ) + + return derivative + + # }}} + + + # {{{ simplicial single axis derivative + + def compute_simplicial_derivative(actx, in_grp, out_grp, + get_diff_mat, vec, ijm, + xyz_axis, metric_in_matvec): + # r for rst axis + return actx.einsum( + "rej,rij,ej->ei" if metric_in_matvec else "rei,rij,ej->ei", + ijm[xyz_axis], + get_diff_mat( + actx, + out_element_group=out_grp, + in_element_group=in_grp), + vec, + arg_names=("inv_jac_t", "ref_stiffT_mat", "vec", ), + tagged=(FirstAxisIsElementsTag(),)) + + # }}} + + return DOFArray( actx, data=tuple( - # r for rst axis - actx.einsum("rej,rij,ej->ei" if metric_in_matvec else "rei,rij,ej->ei", - ijm_i[xyz_axis], - get_diff_mat( - actx, - out_element_group=out_grp, - in_element_group=in_grp), - vec_i, - arg_names=("inv_jac_t", "ref_stiffT_mat", "vec", ), - tagged=(FirstAxisIsElementsTag(),)) - + compute_tensor_product_derivative(actx, in_grp, get_diff_mat, vec_i, + ijm_i, xyz_axis, metric_in_matvec) + if isinstance(in_grp, TensorProductElementGroup) + else compute_simplicial_derivative(actx, in_grp, out_grp, + get_diff_mat, vec_i, ijm_i, + xyz_axis, metric_in_matvec) for out_grp, in_grp, vec_i, ijm_i in zip( out_discr.groups, in_discr.groups, vec, inv_jac_mat))) +# }}} + + +# {{{ gradient def _gradient_kernel(actx, out_discr, in_discr, get_diff_mat, inv_jac_mat, vec, *, metric_in_matvec): # See _single_axis_derivative_kernel for comments on the usage scenarios # (both strong and weak derivative) and their differences. + + # {{{ tensor product grad + + def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm, + metric_in_matvec): + """ + Applies 1D operators one-axis-at-a-time to tensor-product discretized + DOF data. + """ + + # reshape vector to expose tensor product structure + vec = fold(grp.space, vec) + + if metric_in_matvec: + stiff_1d, mass_1d = get_diff_mat(actx, grp, grp) + + grad = [] + for xyz_axis in range(grp.dim): + grad.append(vec) + apply_mass_axes = set(range(grp.dim)) - {xyz_axis} + + # apply mass operators + for ax in apply_mass_axes: + grad[xyz_axis] = single_axis_operator_application( + actx, grp.dim, mass_1d, ax, grad[xyz_axis], + tagged=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered(),), + arg_names=("mass_1d", f"vec_{xyz_axis}") + ) + + # apply stiffness operator and unfold + grad[xyz_axis] = unfold( + grp.space, + single_axis_operator_application( + actx, grp.dim, stiff_1d, xyz_axis, grad[xyz_axis], + tagged=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered(),), + arg_names=("stiff_1d", f"vec_{xyz_axis}")) + ) + + else: + diff_mat = get_diff_mat(actx, grp, grp) + + grad = [] + for xyz_axis in range(grp.dim): + grad.append(vec) + grad[xyz_axis] = unfold( + grp.space, + single_axis_operator_application( + actx, grp.dim, diff_mat, xyz_axis, grad[xyz_axis], + tagged=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered(),), + arg_names=("diff_mat", f"vec_{xyz_axis}") + ) + ) + + grad = actx.np.stack(grad) + return tag_axes( + actx, + { + 0: DiscretizationAmbientDimAxisTag(), + 1: DiscretizationElementAxisTag(), + 2: DiscretizationDOFAxisTag() + }, + actx.einsum( + "xrej,rej->xej", + ijm, + grad, + tagged=(FirstAxisIsElementsTag(),), + arg_names=("inv_jac_t", f"vec") + )) + + # }}} + + + # {{{ simplicial grad + + def compute_simplicial_grad(actx, in_grp, out_grp, get_diff_mat, vec_i, + ijm_i, metric_in_matvec): + return actx.einsum( + "xrej,rij,ej->xei" if metric_in_matvec else "xrei,rij,ej->xei", + ijm_i, + get_diff_mat( + actx, + out_element_group=out_grp, + in_element_group=in_grp + ), + vec_i, + arg_names=("inv_jac_t", "ref_stiffT_mat", "vec"), + tagged=(FirstAxisIsElementsTag(),)) + + # }}} + + per_group_grads = [ - # r for rst axis - # x for xyz axis - actx.einsum("xrej,rij,ej->xei" if metric_in_matvec else "xrei,rij,ej->xei", - ijm_i, - get_diff_mat( - actx, - out_element_group=out_grp, - in_element_group=in_grp - ), - vec_i, - arg_names=("inv_jac_t", "ref_stiffT_mat", "vec"), - tagged=(FirstAxisIsElementsTag(),)) + compute_tensor_product_grad(actx, in_grp, get_diff_mat, vec_i, ijm_i, + metric_in_matvec) + if isinstance(in_grp, TensorProductElementGroup) + else compute_simplicial_grad(actx, in_grp, out_grp, get_diff_mat, vec_i, + ijm_i, metric_in_matvec) + for out_grp, in_grp, vec_i, ijm_i in zip( out_discr.groups, in_discr.groups, vec, - inv_jac_mat)] + inv_jac_mat) + ] return make_obj_array([ DOFArray( actx, data=tuple([pgg_i[xyz_axis] for pgg_i in per_group_grads])) for xyz_axis in range(out_discr.ambient_dim)]) +# }}} + + +# {{{ divergence def _divergence_kernel(actx, out_discr, in_discr, get_diff_mat, inv_jac_mat, vec, *, metric_in_matvec): # See _single_axis_derivative_kernel for comments on the usage scenarios # (both strong and weak derivative) and their differences. + + + # {{{ tensor product div + + def compute_tensor_product_div(actx, grp, diff_mat, vec, ijm): + """ + Exploits tensor product structure to reduce complexity. See + `_gradient_kernel.compute_tensor_product_grad` for more details. + """ + + vec = make_obj_array([ + fold(grp.space, vec[func_axis]) + for func_axis in range(vec.shape[0]) + ]) + + if metric_in_matvec: + stiff_1d, mass_1d = get_diff_mat(actx, grp, grp) + + partials = [] + for func_axis in range(vec.shape[0]): + ref = [] + for xyz_axis in range(grp.dim): + ref.append(vec[func_axis]) + + apply_mass_axes = set(range(grp.dim)) - {xyz_axis} + for ax in apply_mass_axes: + ref[xyz_axis] = single_axis_operator_application( + actx, grp.dim, mass_1d, ax, ref[xyz_axis], + tagged=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered(),), + arg_names=("mass_1d", f"vec_{func_axis}_{xyz_axis}") + ) + + ref[xyz_axis] = single_axis_operator_application( + actx, grp.dim, stiff_1d, xyz_axis, ref[xyz_axis], + tagged=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered(),), + arg_names=("stiff_1d", f"vec_{func_axis}_{xyz_axis}") + ) + + partials.append(ref) + + else: + diff_mat = get_diff_mat(actx, grp, grp) + + partials = [] + for func_axis in range(vec.shape[0]): + ref = [] + for xyz_axis in range(grp.dim): + ref.append(vec[func_axis]) + + ref[xyz_axis] = single_axis_operator_application( + actx, grp.dim, diff_mat, xyz_axis, ref[xyz_axis], + tagged=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered(),), + arg_names=("diff_mat", f"vec_{func_axis}_{xyz_axis}") + ) + + partials.append(ref) + + partials = actx.np.stack([ + unfold(grp.space, partials[func_axis][xyz_axis]) + for func_axis in range(grp.dim) + for xyz_axis in range(grp.dim) + ]) + partials = partials.reshape(grp.dim, grp.dim, *partials.shape[-2:]) + + div = actx.einsum( + 'xrej,xrej->ej', + ijm, + partials, + arg_names=("inv_jac_t", "partials"), + tagged=(FirstAxisIsElementsTag(),) + ) + + return div + + + # }}} + + + # {{{ simplicial div + + def compute_simplicial_div(actx, in_grp, out_grp, get_diff_mat, vec_i, + ijm_i, metric_in_matvec): + return actx.einsum( + "xrej,rij,xej->ei" if metric_in_matvec else "xrei,rij,xej->ei", + ijm_i, + get_diff_mat( + actx, + out_element_group=out_grp, + in_element_group=in_grp + ), + vec_i, + arg_names=("inv_jac_t", "ref_stiffT_mat", "vec"), + tagged=(FirstAxisIsElementsTag(),)) + + # }}} + + per_group_divs = [ + + compute_tensor_product_div(actx, in_grp, get_diff_mat, vec_i, ijm_i) + if isinstance(in_grp, TensorProductElementGroup) + # r for rst axis # x for xyz axis - actx.einsum("xrej,rij,xej->ei" if metric_in_matvec else "xrei,rij,xej->ei", - ijm_i, - get_diff_mat( - actx, - out_element_group=out_grp, - in_element_group=in_grp - ), - vec_i, - arg_names=("inv_jac_t", "ref_stiffT_mat", "vec"), - tagged=(FirstAxisIsElementsTag(),)) + else compute_simplicial_div(actx, in_grp, out_grp, get_diff_mat, vec_i, + ijm_i, metric_in_matvec) + for out_grp, in_grp, vec_i, ijm_i in zip( out_discr.groups, in_discr.groups, vec, - inv_jac_mat)] + inv_jac_mat) + ] return DOFArray(actx, data=tuple(per_group_divs)) # }}} +# }}} + # {{{ Derivative operators @@ -263,12 +557,36 @@ def _reference_derivative_matrices(actx: ArrayContext, actx, _reference_derivative_matrices, lambda grp: grp.discretization_key()) def get_ref_derivative_mats(grp): - from meshmode.discretization.poly_element import diff_matrices - return actx.freeze( - actx.tag_axis( - 1, DiscretizationDOFAxisTag(), - actx.from_numpy( - np.asarray(diff_matrices(grp))))) + + if isinstance(grp, TensorProductElementGroup): + import modepy as mp + import numpy.linalg as la + + #FIXME: Can be gotten rid of by updating meshmode + nodes1d = grp.unit_nodes_1d + bases_1d = grp.bases_1d() + + vdm_1d = mp.vandermonde(bases_1d.functions, nodes1d) + vdm_p_1d = mp.vandermonde(bases_1d.gradients, nodes1d)[0] + + diff_mat = actx.from_numpy(vdm_p_1d @ la.inv(vdm_1d)) + + from arraycontext.metadata import NameHint + return actx.freeze(actx.tag(NameHint("tp_diff_mat_1d"), diff_mat)) + + elif isinstance(grp, SimplexElementGroup): + from meshmode.discretization.poly_element import diff_matrices + + return actx.freeze( + actx.tag_axis( + 1, DiscretizationDOFAxisTag(), + actx.from_numpy( + np.asarray(diff_matrices(grp))))) + + else: + raise TypeError("grp must be either a TensorProductElementGroup or" + f" a SimplexElementGroup. Found {grp}") + return get_ref_derivative_mats(out_element_group) @@ -439,7 +757,35 @@ def get_ref_stiffness_transpose_mat(out_grp, in_grp): from meshmode.discretization.poly_element import \ mass_matrix, diff_matrices + # {{{ tensor product case + + if isinstance(out_grp, TensorProductElementGroup): + import modepy as mp + import numpy.linalg as la + + # FIXME: can be gotten rid of by updating meshmode operators + basis_1d = out_grp.bases_1d() + nodes_1d = out_grp.unit_nodes_1d + + vdm = mp.vandermonde(basis_1d.functions, nodes_1d) + vdm_p = mp.vandermonde(basis_1d.gradients, nodes_1d)[0] + + mass_1d = la.inv(vdm @ vdm.T) + diff_mat = la.solve(vdm.T, vdm_p.T).T + + stiff_1d = actx.freeze( + actx.from_numpy(np.asarray(diff_mat.T @ mass_1d.T))) + + mass_1d = actx.freeze( + actx.tag(MassMatrix1DTag(), + actx.from_numpy(np.asarray(mass_1d)))) + + return (stiff_1d, mass_1d) + + # }}} + mmat = mass_matrix(out_grp) + return actx.freeze( actx.tag_axis(1, DiscretizationDOFAxisTag(), actx.from_numpy( @@ -467,6 +813,7 @@ def get_ref_stiffness_transpose_mat(out_grp, in_grp): ).copy() # contigify the array ) ) + return get_ref_stiffness_transpose_mat(out_element_group, in_element_group) @@ -777,14 +1124,30 @@ def reference_inverse_mass_matrix(actx: ArrayContext, element_group): lambda grp: grp.discretization_key()) def get_ref_inv_mass_mat(grp): from modepy import inverse_mass_matrix - basis = grp.basis_obj() - return actx.freeze( - actx.tag_axis(0, DiscretizationDOFAxisTag(), - actx.from_numpy( - np.asarray( - inverse_mass_matrix(basis.functions, grp.unit_nodes), - order="C")))) + if isinstance(grp, TensorProductElementGroup): + + basis_1d = grp.bases_1d() + nodes_1d = grp.unit_nodes_1d + + inv_mass_1d = inverse_mass_matrix(basis_1d.functions, nodes_1d) + inv_mass_1d = actx.from_numpy(np.asarray(inv_mass_1d)) + + return actx.freeze(actx.tag(InverseMassMatrix1DTag(), inv_mass_1d)) + + elif isinstance(grp, SimplexElementGroup): + + basis = grp.basis_obj() + + return actx.freeze( + actx.tag_axis(0, DiscretizationDOFAxisTag(), + actx.from_numpy( + np.asarray( + inverse_mass_matrix(basis.functions, grp.unit_nodes), + order="C")))) + else: + raise TypeError("grp must be either a TensorProductElementGroup or" + f" a SimplexElementGroup. Found {grp}") return get_ref_inv_mass_mat(element_group) @@ -809,15 +1172,50 @@ def _apply_inverse_mass_operator( discr = dcoll.discr_from_dd(dd_in) inv_area_elements = 1./area_element(actx, dcoll, dd=dd_in, _use_geoderiv_connection=actx.supports_nonscalar_broadcasting) + + + def apply_to_tensor_product_elements(grp, jac_inv, vec, ref_inv_mass): + + vec = fold(grp.space, vec) + + for xyz_axis in range(grp.dim): + vec = single_axis_operator_application( + actx, grp.dim, ref_inv_mass, xyz_axis, vec, + tagged=(FirstAxisIsElementsTag(), + OutputIsTensorProductDOFArrayOrdered(),), + arg_names=("ref_inv_mass_1d", "vec")) + + vec = unfold(grp.space, vec) + + return actx.einsum( + "ei,ei->ei", + jac_inv, + vec, + tagged=(FirstAxisIsElementsTag(),) + ) + + + def apply_to_simplicial_elements(jac_inv, vec, ref_inv_mass): + + # Based on https://arxiv.org/pdf/1608.03836.pdf + # true_Minv ~ ref_Minv * ref_M * (1/jac_det) * ref_Minv + return actx.einsum( + "ei,ij,ej->ei", + jac_inv, + ref_inv_mass, + vec, + tagged=(FirstAxisIsElementsTag(),)) + + group_data = [ - # Based on https://arxiv.org/pdf/1608.03836.pdf - # true_Minv ~ ref_Minv * ref_M * (1/jac_det) * ref_Minv - actx.einsum("ei,ij,ej->ei", - jac_inv, - reference_inverse_mass_matrix(actx, element_group=grp), - vec_i, - tagged=(FirstAxisIsElementsTag(),)) - for grp, jac_inv, vec_i in zip(discr.groups, inv_area_elements, vec)] + apply_to_tensor_product_elements( + grp, jac_inv, vec_i, + reference_inverse_mass_matrix(actx, element_group=grp)) + if isinstance(grp, TensorProductElementGroup) else + apply_to_simplicial_elements(jac_inv, vec_i, + reference_inverse_mass_matrix(actx, element_group=grp)) + for grp, jac_inv, vec_i in zip(discr.groups, inv_area_elements, vec) + ] return DOFArray(actx, data=tuple(group_data)) @@ -1064,4 +1462,58 @@ def face_mass(dcoll: DiscretizationCollection, *args) -> ArrayOrContainer: # }}} +# {{{ general single axis operator application + +def single_axis_operator_application(actx, dim, operator, axis, vec, + arg_names=None, tagged=None): + """ + Used for applying 1D operators to a single axis of a tensor of DOF data. + """ + + if not isinstance(arg_names, tuple) and arg_names is not None: + raise TypeError("arg_names must be a tuple.") + if not isinstance(tagged, tuple) and tagged is not None: + raise TypeError("tagged must be a tuple.") + + # {{{ ensure axes are properly tagged + + vec = actx.tag_axis(0, DiscretizationElementAxisTag(), vec) + vec = tag_axes( + actx, + { i: TensorProductDOFAxisTag(i-1) for i in range(1, dim+1) }, + vec + ) + + operator = tag_axes( + actx, + { i: TensorProductOperatorAxisTag() for i in range(2) }, + operator + ) + + # }}} + + # {{{ einsum spec construction + + # 3D grad example spec using formula below: + # assume operator is a differentiation operator + # x-axis (axis = 0) contraction: ij,ejop->eiop + # y-axis (axis = 1) contraction: ij,eajp->eaip + # z-axis (axis = 2) contraction: ij,eabj->eabi + operator_spec = 'ij' + data_spec = f'e{"abcdefghklmn"[:axis]}j{"opqrstuvwxyz"[:dim-axis-1]}' + out_spec = f'e{"abcdefghklmn"[:axis]}i{"opqrstuvwxyz"[:dim-axis-1]}' + + spec = operator_spec + ',' + data_spec + '->' + out_spec + + # }}} + + return tag_axes( + actx, + { i: TensorProductDOFAxisTag(i-1) for i in range(1, dim+1) }, + actx.einsum(spec, operator, vec, arg_names=arg_names, tagged=tagged) + ) + +# }}} + + # vim: foldmethod=marker diff --git a/grudge/transform/metadata.py b/grudge/transform/metadata.py new file mode 100644 index 000000000..d5634ebc0 --- /dev/null +++ b/grudge/transform/metadata.py @@ -0,0 +1,71 @@ +__copyright__ = "Copyright (C) 2024 Addison Alvey-Blanco" + +__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. +""" + +from pytools.tag import IgnoredForPropagationTag, Tag, tag_dataclass +from meshmode.transform_metadata import DiscretizationDOFAxisTag + + +# {{{ tensor product specific metadata + +class OutputIsTensorProductDOFArrayOrdered(Tag): + # FIXME: REMOVE THIS + # /!\ THIS IS TEMPORARY AND WILL GO AWAY /!\ + """ + Signify that the strides will not be of order "C" or "F". + + Used to specify strides for eager einsums. + """ + pass + + +@tag_dataclass +class TensorProductDOFAxisTag(DiscretizationDOFAxisTag): + """ + Tag an axis as being an axis containing the DOFs of a tensor-product + discretization. Used to signify the relative update speed of an axis for + transformation (i.e. loop nest ordering) purposes. + """ + iaxis: int + + +class TensorProductOperatorAxisTag(IgnoredForPropagationTag): + """ + Signify that an axis belongs to a 1D operator. No tags will be propagated + along an axis tagged with this tag. + """ + pass + + +class ReferenceTensorProductMassOperatorTag(Tag): + """ + Used in DAG transformation to realize algebraic simplification of 1D + inverse mass operator times mass operator. + """ + pass + + +class ReferenceTensorProductInverseMassOperatorTag(Tag): + """ + See MassMatrix1d. + """ + +# }}} diff --git a/requirements.txt b/requirements.txt index 2107e5aeb..86e4a9237 100644 --- a/requirements.txt +++ b/requirements.txt @@ -14,7 +14,7 @@ git+https://github.com/inducer/meshmode.git#egg=meshmode git+https://github.com/inducer/pyvisfile.git#egg=pyvisfile git+https://github.com/inducer/pymetis.git#egg=pymetis git+https://github.com/illinois-ceesd/logpyle.git#egg=logpyle -git+https://github.com/inducer/pytato.git#egg=pytato +git+https://github.com/a-alveyblanc/pytato.git@implement-f-ordered-reshapes # for test_wave_dt_estimate sympy diff --git a/test/test_grudge.py b/test/test_grudge.py index 547ce8a2c..fd225cca0 100644 --- a/test/test_grudge.py +++ b/test/test_grudge.py @@ -34,6 +34,7 @@ from meshmode import _acf # noqa: F401 from meshmode.dof_array import flat_norm +from meshmode.mesh import SimplexElementGroup, TensorProductElementGroup import meshmode.mesh.generation as mgen from pytools.obj_array import flat_obj_array @@ -428,31 +429,66 @@ def df(x, axis): # {{{ divergence theorem -def test_2d_gauss_theorem(actx_factory): +@pytest.mark.parametrize("group_cls", + [SimplexElementGroup, TensorProductElementGroup]) +def test_2d_gauss_theorem(actx_factory, group_cls): """Verify Gauss's theorem explicitly on a mesh""" - pytest.importorskip("meshpy") + actx = actx_factory() - from meshpy.geometry import make_circle, GeometryBuilder - from meshpy.triangle import MeshInfo, build + if group_cls is SimplexElementGroup: + pytest.importorskip("meshpy") - geob = GeometryBuilder() - geob.add_geometry(*make_circle(1)) - mesh_info = MeshInfo() - geob.set(mesh_info) + from meshpy.geometry import make_circle, GeometryBuilder + from meshpy.triangle import MeshInfo, build - mesh_info = build(mesh_info) + geob = GeometryBuilder() + geob.add_geometry(*make_circle(1)) + mesh_info = MeshInfo() + geob.set(mesh_info) - from meshmode.mesh.io import from_meshpy - from meshmode.mesh import BTAG_ALL + mesh_info = build(mesh_info) - mesh = from_meshpy(mesh_info, order=1) + from meshmode.mesh.io import from_meshpy + from meshmode.mesh import BTAG_ALL - actx = actx_factory() + mesh = from_meshpy(mesh_info, order=1) - dcoll = DiscretizationCollection(actx, mesh, order=2) - volm_disc = dcoll.discr_from_dd(dof_desc.DD_VOLUME) - x_volm = actx.thaw(volm_disc.nodes()) + dcoll = DiscretizationCollection(actx, mesh, order=2) + volm_disc = dcoll.discr_from_dd(dof_desc.DD_VOLUME) + x_volm = actx.thaw(volm_disc.nodes()) + + elif group_cls is TensorProductElementGroup: + from meshmode.mesh.generation import generate_regular_rect_mesh + + dim = 2 + mesh = generate_regular_rect_mesh( + (-1,)*dim, (1,)*dim, nelements_per_axis=(2,)*dim, + group_cls=TensorProductElementGroup) + + alpha = 0.3 + rot_mat = np.array([ + [np.cos(alpha), np.sin(alpha)], + [-np.sin(alpha), np.cos(alpha)] + ]) + + from meshmode.mesh.processing import affine_map + mesh = affine_map(mesh, A=rot_mat) + + from meshmode.discretization.poly_element import \ + LegendreGaussLobattoTensorProductGroupFactory as LGL + dcoll = DiscretizationCollection( + actx, mesh, discr_tag_to_group_factory={ + dof_desc.DISCR_TAG_BASE: LGL(order=2) + } + ) + + volm_disc = dcoll.discr_from_dd(dof_desc.DD_VOLUME) + x_volm = actx.thaw(volm_disc.nodes()) + + else: + raise AssertionError('group_cls must be SimplexElementGroup or ' + f'TensorProductElementGroup. Found {group_cls}') def f(x): return flat_obj_array( @@ -463,6 +499,7 @@ def f(x): f_volm = f(x_volm) int_1 = op.integral(dcoll, "vol", op.local_div(dcoll, f_volm)) + from grudge.dof_desc import BTAG_ALL prj_f = op.project(dcoll, "vol", BTAG_ALL, f_volm) normal = geo.normal(actx, dcoll, BTAG_ALL) int_2 = op.integral(dcoll, BTAG_ALL, prj_f.dot(normal)) diff --git a/test/test_metrics.py b/test/test_metrics.py index 586b093a5..7ab4b5ffe 100644 --- a/test/test_metrics.py +++ b/test/test_metrics.py @@ -23,6 +23,9 @@ THE SOFTWARE. """ +from meshmode.discretization.poly_element import LegendreGaussLobattoTensorProductGroupFactory +from meshmode.mesh import SimplexElementGroup, TensorProductElementGroup +from meshmode.mesh.processing import affine_map import numpy as np from grudge.array_context import ( @@ -51,12 +54,18 @@ @pytest.mark.parametrize("dim", [2, 3]) @pytest.mark.parametrize("nonaffine", [False, True]) @pytest.mark.parametrize("use_quad", [False, True]) -def test_inverse_metric(actx_factory, dim, nonaffine, use_quad): +@pytest.mark.parametrize("group_cls", [ + SimplexElementGroup, + TensorProductElementGroup +]) +def test_inverse_metric(actx_factory, dim, nonaffine, use_quad, group_cls): actx = actx_factory() order = 3 - mesh = mgen.generate_regular_rect_mesh(a=(-0.5,)*dim, b=(0.5,)*dim, - nelements_per_axis=(6,)*dim, order=order) + mesh = mgen.generate_regular_rect_mesh( + a=(-0.5,)*dim, b=(0.5,)*dim, + nelements_per_axis=(6,)*dim, order=order, + group_cls=group_cls) if nonaffine: def m(x): @@ -73,19 +82,36 @@ def m(x): from meshmode.mesh.processing import map_mesh mesh = map_mesh(mesh, m) + else: + alpha = 0.3 + rot_mat = np.array([ + [np.cos(alpha), np.sin(alpha), 0], + [-np.sin(alpha), np.cos(alpha), 0], + [0, 0, 1], + ])[:dim, :dim] + + mesh = affine_map(mesh, A=rot_mat) from grudge.dof_desc import as_dofdesc, DISCR_TAG_BASE, DISCR_TAG_QUAD from meshmode.discretization.poly_element import \ QuadratureSimplexGroupFactory, \ default_simplex_group_factory - dcoll = DiscretizationCollection( - actx, mesh, - discr_tag_to_group_factory={ + if group_cls is SimplexElementGroup: + discr_tag_to_group_factory = { DISCR_TAG_BASE: default_simplex_group_factory(base_dim=dim, order=order), DISCR_TAG_QUAD: QuadratureSimplexGroupFactory(2*order + 1), } - ) + elif group_cls is TensorProductElementGroup: + discr_tag_to_group_factory = { + DISCR_TAG_BASE: LegendreGaussLobattoTensorProductGroupFactory(order=order), + DISCR_TAG_QUAD: LegendreGaussLobattoTensorProductGroupFactory(order=3*order), + } + else: + raise AssertionError() + + dcoll = DiscretizationCollection( + actx, mesh, discr_tag_to_group_factory=discr_tag_to_group_factory) from grudge.geometry import \ forward_metric_derivative_mat, inverse_metric_derivative_mat diff --git a/test/test_op.py b/test/test_op.py index f04f25567..c12b674f0 100644 --- a/test/test_op.py +++ b/test/test_op.py @@ -21,6 +21,7 @@ """ +from meshmode.mesh.processing import affine_map import numpy as np import meshmode.mesh.generation as mgen @@ -30,6 +31,8 @@ from grudge import op, geometry as geo, DiscretizationCollection from grudge.dof_desc import DOFDesc +from meshmode.mesh import SimplexElementGroup, TensorProductElementGroup + import pytest from grudge.array_context import PytestPyOpenCLArrayContextFactory @@ -44,6 +47,10 @@ # {{{ gradient +@pytest.mark.parametrize("group_cls", [ + SimplexElementGroup, + TensorProductElementGroup +]) @pytest.mark.parametrize("form", ["strong", "weak"]) @pytest.mark.parametrize("dim", [1, 2, 3]) @pytest.mark.parametrize("order", [2, 3]) @@ -53,7 +60,7 @@ (True, True) ]) def test_gradient(actx_factory, form, dim, order, vectorize, nested, - visualize=False): + group_cls, visualize=False): actx = actx_factory() from pytools.convergence import EOCRecorder @@ -61,10 +68,40 @@ def test_gradient(actx_factory, form, dim, order, vectorize, nested, for n in [4, 6, 8]: 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) + + if group_cls is TensorProductElementGroup: + # no reason to test 1D tensor product elements + if dim == 1: + pytest.skip() + + import grudge.dof_desc as dd + from meshmode.discretization.poly_element import \ + LegendreGaussLobattoTensorProductGroupFactory as LGL + + dcoll = DiscretizationCollection( + actx, + mesh, + discr_tag_to_group_factory={ + dd.DISCR_TAG_BASE: LGL(order)}) + + elif group_cls is SimplexElementGroup: + dcoll = DiscretizationCollection(actx, mesh, order=order) + + else: + raise AssertionError('Expecting TensorProductElementGroup or ' + f'SimplexElementGroup. Found {group_cls}') - dcoll = DiscretizationCollection(actx, mesh, order=order) + alpha = 0.3 + rot_mat = np.array([ + [np.cos(alpha), np.sin(alpha), 0], + [-np.sin(alpha), np.cos(alpha), 0], + [0, 0, 1], + ])[:dim, :dim] + + mesh = affine_map(mesh, A=rot_mat) def f(x): result = dcoll.zeros(actx) + 1 @@ -164,6 +201,10 @@ def get_flux(u_tpair): # {{{ divergence +@pytest.mark.parametrize("group_cls", [ + SimplexElementGroup, + TensorProductElementGroup +]) @pytest.mark.parametrize("form", ["strong", "weak"]) @pytest.mark.parametrize("dim", [1, 2, 3]) @pytest.mark.parametrize("order", [2, 3]) @@ -173,7 +214,7 @@ def get_flux(u_tpair): (True, True) ]) def test_divergence(actx_factory, form, dim, order, vectorize, nested, - visualize=False): + group_cls, visualize=False): actx = actx_factory() from pytools.convergence import EOCRecorder @@ -181,11 +222,40 @@ def test_divergence(actx_factory, form, dim, order, vectorize, nested, for n in [4, 6, 8]: 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) + + if group_cls is TensorProductElementGroup: + # no reason to test 1D tensor product elements + if dim == 1: + return + + import grudge.dof_desc as dd + from meshmode.discretization.poly_element import \ + LegendreGaussLobattoTensorProductGroupFactory as LGL + + dcoll = DiscretizationCollection( + actx, + mesh, + discr_tag_to_group_factory={ + dd.DISCR_TAG_BASE: LGL(order)}) + + elif group_cls is SimplexElementGroup: + dcoll = DiscretizationCollection(actx, mesh, order=order) + + else: + raise AssertionError('Expecting TensorProductElementGroup or ' + f'SimplexElementGroup. Found {group_cls}') - dcoll = DiscretizationCollection(actx, mesh, order=order) + alpha = 0.3 + rot_mat = np.array([ + [np.cos(alpha), np.sin(alpha), 0], + [-np.sin(alpha), np.cos(alpha), 0], + [0, 0, 1], + ])[:dim, :dim] + mesh = affine_map(mesh, A=rot_mat) def f(x): result = make_obj_array([dcoll.zeros(actx) + (i+1) for i in range(dim)]) for i in range(dim-1):