diff --git a/grudge/op.py b/grudge/op.py index c64c51b1d7a1dfb2576bdd06e07bab36ac5138ba..c62732d1e7ddca7d9d0a3974b61065f27e35eeef 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -48,7 +48,7 @@ THE SOFTWARE. from numbers import Number -from pytools import memoize_on_first_arg +from pytools import memoize_on_first_arg, keyed_memoize_in import numpy as np # noqa from pytools.obj_array import obj_array_vectorize, make_obj_array @@ -60,6 +60,8 @@ import grudge.dof_desc as dof_desc from meshmode.mesh import BTAG_ALL, BTAG_NONE, BTAG_PARTITION # noqa from meshmode.dof_array import freeze, flatten, unflatten +import loopy as lp + from grudge.symbolic.primitives import TracePair @@ -310,6 +312,68 @@ def mass(dcoll, *args): return _bound_mass(dcoll, dd)(u=vec) +@memoize_on_first_arg +def _elwise_linear_loopy_prg(actx): + result = make_loopy_program( + """{[iel, idof, j]: + 0<=ielij", weights, vand_inv_t, o_vand) + + return actx.freeze( + actx.from_numpy(get_ref_mass_mat(out_element_group, + in_element_group)) + ) + + +def _apply_mass_operator(dcoll, dd, vec): + pass + + +def mass_operator(dcoll, *args): + + if len(args) == 1: + vec, = args + dd = sym.DOFDesc("vol", sym.QTAG_NONE) + elif len(args) == 2: + dd, vec = args + else: + raise TypeError("invalid number of arguments") + + if isinstance(vec, np.ndarray): + return obj_array_vectorize( + lambda el: mass_operator(dcoll, dd, el), vec + ) + + return _apply_mass_operator(dcoll, dd, vec) + + @memoize_on_first_arg def _bound_inverse_mass(dcoll): return bind(dcoll, sym.InverseMassOperator()(sym.Variable("u")),