示例#1
0
文件: op.py 项目: nchristensen/grudge
def _apply_face_mass_operator(dcoll: DiscretizationCollection, dd, vec):
    if not isinstance(vec, DOFArray):
        return map_array_container(
            partial(_apply_face_mass_operator, dcoll, dd), vec)

    from grudge.geometry import area_element

    volm_discr = dcoll.discr_from_dd(dof_desc.DD_VOLUME)
    face_discr = dcoll.discr_from_dd(dd)
    dtype = vec.entry_dtype
    actx = vec.array_context

    assert len(face_discr.groups) == len(volm_discr.groups)
    surf_area_elements = area_element(
        actx,
        dcoll,
        dd=dd,
        _use_geoderiv_connection=actx.supports_nonscalar_broadcasting)

    return DOFArray(
        actx,
        data=tuple(
            actx.einsum("ifj,fej,fej->ei",
                        reference_face_mass_matrix(actx,
                                                   face_element_group=afgrp,
                                                   vol_element_group=vgrp,
                                                   dtype=dtype),
                        surf_ae_i.reshape(vgrp.mesh_el_group.nfaces,
                                          vgrp.nelements, -1),
                        vec_i.reshape(vgrp.mesh_el_group.nfaces,
                                      vgrp.nelements, afgrp.nunit_dofs),
                        arg_names=("ref_face_mass_mat", "jac_surf", "vec"),
                        tagged=(FirstAxisIsElementsTag(), )) for vgrp, afgrp,
            vec_i, surf_ae_i in zip(volm_discr.groups, face_discr.groups, vec,
                                    surf_area_elements)))
示例#2
0
文件: op.py 项目: nchristensen/grudge
def _apply_mass_operator(dcoll: DiscretizationCollection, dd_out, dd_in, vec):
    if not isinstance(vec, DOFArray):
        return map_array_container(
            partial(_apply_mass_operator, dcoll, dd_out, dd_in), vec)

    from grudge.geometry import area_element

    in_discr = dcoll.discr_from_dd(dd_in)
    out_discr = dcoll.discr_from_dd(dd_out)

    actx = vec.array_context
    area_elements = area_element(
        actx,
        dcoll,
        dd=dd_in,
        _use_geoderiv_connection=actx.supports_nonscalar_broadcasting)
    return DOFArray(
        actx,
        data=tuple(
            actx.einsum("ij,ej,ej->ei",
                        reference_mass_matrix(actx,
                                              out_element_group=out_grp,
                                              in_element_group=in_grp),
                        ae_i,
                        vec_i,
                        arg_names=("mass_mat", "jac", "vec"),
                        tagged=(FirstAxisIsElementsTag(), ))
            for in_grp, out_grp, ae_i, vec_i in zip(
                in_discr.groups, out_discr.groups, area_elements, vec)))
示例#3
0
文件: op.py 项目: nchristensen/grudge
def _div_helper(dcoll, diff_func, *args):
    if len(args) == 1:
        vecs, = args
        dd = dof_desc.DOFDesc("vol", dof_desc.DISCR_TAG_BASE)
    elif len(args) == 2:
        dd, vecs = args
    else:
        raise TypeError("invalid number of arguments")

    if not isinstance(vecs, np.ndarray):
        # vecs is not an object array -> treat as array container
        return map_array_container(partial(_div_helper, dcoll, diff_func, dd),
                                   vecs)

    assert vecs.dtype == object

    if vecs.size:
        sample_vec = vecs[(0, ) * vecs.ndim]

        if isinstance(sample_vec, np.ndarray):
            assert sample_vec.dtype == object
            # vecs is an object array containing further object arrays
            # -> treat as array container
            return map_array_container(
                partial(_div_helper, dcoll, diff_func, dd), vecs)

    if vecs.shape[-1] != dcoll.ambient_dim:
        raise ValueError(
            "last/innermost dimension of *vecs* argument doesn't match "
            "ambient dimension")

    div_result_shape = vecs.shape[:-1]

    if len(div_result_shape) == 0:
        return sum(diff_func(dd, i, vec_i) for i, vec_i in enumerate(vecs))
    else:
        result = np.zeros(div_result_shape, dtype=object)
        for idx in np.ndindex(div_result_shape):
            result[idx] = sum(
                diff_func(dd, i, vec_i) for i, vec_i in enumerate(vecs[idx]))
        return result
示例#4
0
def componentwise_norms(discr, fields, order=np.inf):
    """Return the *order*-norm for each component of *fields*.

    .. note::
        This is a collective routine and must be called by all MPI ranks.
    """
    if not isinstance(fields, DOFArray):
        return map_array_container(
            partial(componentwise_norms, discr, order=order), fields)
    if len(fields) > 0:
        return discr.norm(fields, order)
    else:
        # FIXME: This work-around for #575 can go away after #569
        return 0
示例#5
0
def filter_modally(dcoll, dd, cutoff, mode_resp_func, field):
    """Stand-alone procedural interface to spectral filtering.

    For each element group in the discretization, and restriction,
    This routine generates:

    * a filter operator:
        - *cutoff* filters only modes above this mode id
        - *mode_resp_func* function returns a filter coefficient
            for a given mode
        - memoized into the array context

    * a filtered solution wherein the filter is applied to *field*.

    Parameters
    ----------
    dcoll: :class:`grudge.discretization.DiscretizationCollection`
        Grudge discretization with boundaries object
    dd: :class:`grudge.dof_desc.DOFDesc` or as accepted by
        :func:`grudge.dof_desc.as_dofdesc`
        Describe the type of DOF vector on which to operate.
    cutoff: int
        Mode below which *field* will not be filtered
    mode_resp_func:
        Modal response function returns a filter coefficient for input mode id
    field: :class:`mirgecom.fluid.ConservedVars`
        An array container containing the relevant field(s) to filter.

    Returns
    -------
    result: :class:`mirgecom.fluid.ConservedVars`
        An array container containing the filtered field(s).
    """
    if not isinstance(field, DOFArray):
        return map_array_container(
            partial(filter_modally, dcoll, dd, cutoff, mode_resp_func), field)

    actx = field.array_context
    dd = dof_desc.as_dofdesc(dd)
    dd_modal = dof_desc.DD_VOLUME_MODAL
    discr = dcoll.discr_from_dd(dd)

    modal_map = dcoll.connection_from_dds(dd, dd_modal)
    nodal_map = dcoll.connection_from_dds(dd_modal, dd)
    field = modal_map(field)
    field = apply_spectral_filter(actx, field, discr, cutoff, mode_resp_func)
    return nodal_map(field)
示例#6
0
    def _unflatten_from_numpy(subary):
        if isinstance(subary, np.ndarray) and subary.dtype.char != "O":
            subary = actx.from_numpy(subary)

        # FIXME: this is doing the recursion itself instead of just using
        # `rec_map_dof_array_container` like `flatten_to_numpy` to catch
        # non-object ndarrays, which `is_array_container` considers as as
        # containers and tries to serialize.
        from arraycontext import map_array_container, is_array_container
        if is_array_container(subary):
            return map_array_container(_unflatten_from_numpy, subary)
        else:
            return _unflatten_dof_array(actx,
                                        subary,
                                        group_shapes,
                                        group_starts,
                                        strict=strict)
示例#7
0
    def inverse_mass(self, vec):
        if not isinstance(vec, DOFArray):
            return map_array_container(self.inverse_mass, vec)

        actx = vec.array_context
        dtype = vec.entry_dtype
        discr = self.volume_discr

        return DOFArray(
            actx,
            data=tuple(
                actx.einsum("ij,ej->ei",
                            self.get_inverse_mass_matrix(grp, dtype),
                            vec_i,
                            arg_names=("mass_inv_mat", "vec"),
                            tagged=(FirstAxisIsElementsTag(), ))
                for grp, vec_i in zip(discr.groups, vec))) / thaw(
                    self.vol_jacobian(), actx)
示例#8
0
def project(dcoll: DiscretizationCollection, src, tgt,
            vec) -> ArrayOrContainerT:
    """Project from one discretization to another, e.g. from the
    volume to the boundary, or from the base to the an overintegrated
    quadrature discretization.

    :arg src: a :class:`~grudge.dof_desc.DOFDesc`, or a value convertible to one.
    :arg tgt: a :class:`~grudge.dof_desc.DOFDesc`, or a value convertible to one.
    :arg vec: a :class:`~meshmode.dof_array.DOFArray` or an
        :class:`~arraycontext.container.ArrayContainer` of them.
    :returns: a :class:`~meshmode.dof_array.DOFArray` or an
        :class:`~arraycontext.container.ArrayContainer` like *vec*.
    """
    src = as_dofdesc(src)
    tgt = as_dofdesc(tgt)

    if isinstance(vec, Number) or src == tgt:
        return vec

    if not isinstance(vec, DOFArray):
        return map_array_container(partial(project, dcoll, src, tgt), vec)

    return dcoll.connection_from_dds(src, tgt)(vec)
示例#9
0
    def face_mass(self, vec):
        if not isinstance(vec, DOFArray):
            return map_array_container(self.face_mass, vec)

        actx = vec.array_context
        dtype = vec.entry_dtype

        @memoize_in(self, "face_mass_knl")
        def knl():
            return make_loopy_program(
                """{[iel,idof,f,j]:
                    0<=iel<nelements and
                    0<=f<nfaces and
                    0<=idof<nvol_nodes and
                    0<=j<nface_nodes}""", "result[iel,idof] = "
                "sum(f, sum(j, mat[idof, f, j] * vec[f, iel, j]))",
                name="face_mass")

        all_faces_conn = self.get_connection("vol", "all_faces")
        all_faces_discr = all_faces_conn.to_discr
        vol_discr = all_faces_conn.from_discr

        fj = thaw(self.face_jacobian("all_faces"), vec.array_context)
        vec = vec * fj

        assert len(all_faces_discr.groups) == len(vol_discr.groups)

        return DOFArray(
            actx,
            data=tuple(
                actx.call_loopy(
                    knl(),
                    mat=self.get_local_face_mass_matrix(afgrp, volgrp, dtype),
                    vec=vec_i.reshape(volgrp.mesh_el_group.nfaces, volgrp.
                                      nelements, afgrp.nunit_dofs))["result"]
                for afgrp, volgrp, vec_i in zip(all_faces_discr.groups,
                                                vol_discr.groups, vec)))
示例#10
0
文件: op.py 项目: nchristensen/grudge
def _grad_helper(dcoll, scalar_grad, *args, nested):
    if len(args) == 1:
        vec, = args
        dd_in = dof_desc.DOFDesc("vol", dof_desc.DISCR_TAG_BASE)
    elif len(args) == 2:
        dd_in, vec = args
    else:
        raise TypeError("invalid number of arguments")

    if isinstance(vec, np.ndarray):
        # Occasionally, data structures coming from *mirgecom* will
        # contain empty object arrays as placeholders for fields.
        # For example, species mass fractions is an empty object array when
        # running in a single-species configuration.
        # This hack here ensures that these empty arrays, at the very least,
        # have their shape updated when applying the gradient operator
        if vec.size == 0:
            return vec.reshape(vec.shape + (dcoll.ambient_dim, ))

        # For containers with ndarray data (such as momentum/velocity),
        # the gradient is matrix-valued, so we compute the gradient for
        # each component. If requested (via 'not nested'), return a matrix of
        # derivatives by stacking the results.
        grad = obj_array_vectorize(
            lambda el: _grad_helper(
                dcoll, scalar_grad, dd_in, el, nested=nested), vec)
        if nested:
            return grad
        else:
            return np.stack(grad, axis=0)

    if not isinstance(vec, DOFArray):
        return map_array_container(
            partial(_grad_helper, dcoll, scalar_grad, dd_in, nested=nested),
            vec)

    return scalar_grad(dcoll, dd_in, vec)
示例#11
0
文件: op.py 项目: nchristensen/grudge
def _apply_inverse_mass_operator(dcoll: DiscretizationCollection, dd_out,
                                 dd_in, vec):
    if not isinstance(vec, DOFArray):
        return map_array_container(
            partial(_apply_inverse_mass_operator, dcoll, dd_out, dd_in), vec)

    from grudge.geometry import area_element

    if dd_out != dd_in:
        raise ValueError("Cannot compute inverse of a mass matrix mapping "
                         "between different element groups; inverse is not "
                         "guaranteed to be well-defined")

    actx = vec.array_context
    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)
    group_data = []
    for grp, jac_inv, vec_i in zip(discr.groups, inv_area_elements, vec):

        ref_mass_inverse = reference_inverse_mass_matrix(actx,
                                                         element_group=grp)

        group_data.append(
            # 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,
                        ref_mass_inverse,
                        vec_i,
                        tagged=(FirstAxisIsElementsTag(), )))

    return DOFArray(actx, data=tuple(group_data))
示例#12
0
文件: op.py 项目: nchristensen/grudge
def local_d_dx(dcoll: DiscretizationCollection, xyz_axis,
               vec) -> ArrayOrContainerT:
    r"""Return the element-local derivative along axis *xyz_axis* of a
    function :math:`f` represented by *vec*:

    .. math::

        \frac{\partial f}{\partial \lbrace x,y,z\rbrace}\Big|_E

    :arg xyz_axis: an integer indicating the axis along which the derivative
        is taken.
    :arg vec: a :class:`~meshmode.dof_array.DOFArray` or an
        :class:`~arraycontext.container.ArrayContainer` of them.
    :returns: a :class:`~meshmode.dof_array.DOFArray` or an
        :class:`~arraycontext.container.ArrayContainer` of them.
    """
    if not isinstance(vec, DOFArray):
        return map_array_container(partial(local_d_dx, dcoll, xyz_axis), vec)

    discr = dcoll.discr_from_dd(dof_desc.DD_VOLUME)
    actx = vec.array_context

    from grudge.geometry import inverse_surface_metric_derivative_mat
    inverse_jac_mat = inverse_surface_metric_derivative_mat(
        actx,
        dcoll,
        _use_geoderiv_connection=actx.supports_nonscalar_broadcasting)

    return _single_axis_derivative_kernel(actx,
                                          discr,
                                          discr,
                                          _reference_derivative_matrices,
                                          inverse_jac_mat,
                                          xyz_axis,
                                          vec,
                                          metric_in_matvec=False)
示例#13
0
def _apply_elementwise_reduction(op_name: str, dcoll: DiscretizationCollection,
                                 *args) -> ArrayOrContainerT:
    r"""Returns a vector of DOFs with all entries on each element set
    to the reduction operation *op_name* over all degrees of freedom.

    May be called with ``(vec)`` or ``(dd, vec)``.

    :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` or an
        :class:`~arraycontext.container.ArrayContainer`.
    :returns: a :class:`~meshmode.dof_array.DOFArray` or an
        :class:`~arraycontext.container.ArrayContainer`.
    """
    if len(args) == 1:
        vec, = args
        dd = dof_desc.DOFDesc("vol", dof_desc.DISCR_TAG_BASE)
    elif len(args) == 2:
        dd, vec = args
    else:
        raise TypeError("invalid number of arguments")

    dd = dof_desc.as_dofdesc(dd)

    if not isinstance(vec, DOFArray):
        return map_array_container(
            partial(_apply_elementwise_reduction, op_name, dcoll, dd), vec)

    actx = vec.array_context

    if actx.supports_nonscalar_broadcasting:
        return DOFArray(actx,
                        data=tuple(
                            actx.np.broadcast_to((getattr(actx.np, op_name)(
                                vec_i, axis=1).reshape(-1, 1)), vec_i.shape)
                            for vec_i in vec))
    else:

        @memoize_in(
            actx,
            (_apply_elementwise_reduction, "elementwise_%s_prg" % op_name))
        def elementwise_prg():
            # FIXME: This computes the reduction value redundantly for each
            # output DOF.
            t_unit = make_loopy_program([
                "{[iel]: 0 <= iel < nelements}",
                "{[idof, jdof]: 0 <= idof, jdof < ndofs}"
            ],
                                        """
                    result[iel, idof] = %s(jdof, operand[iel, jdof])
                """ % op_name,
                                        name="grudge_elementwise_%s_knl" %
                                        op_name)
            import loopy as lp
            from meshmode.transform_metadata import (ConcurrentElementInameTag,
                                                     ConcurrentDOFInameTag)
            return lp.tag_inames(
                t_unit, {
                    "iel": ConcurrentElementInameTag(),
                    "idof": ConcurrentDOFInameTag()
                })

        return DOFArray(
            actx,
            data=tuple(
                actx.call_loopy(elementwise_prg(), operand=vec_i)["result"]
                for vec_i in vec))
示例#14
0
文件: direct.py 项目: MTCam/meshmode
    def __call__(self, ary):
        from meshmode.dof_array import DOFArray
        if is_array_container(ary) and not isinstance(ary, DOFArray):
            return map_array_container(self, ary)

        if not isinstance(ary, DOFArray):
            raise TypeError("non-array passed to discretization connection")

        if ary.shape != (len(self.from_discr.groups), ):
            raise ValueError("invalid shape of incoming resampling data")

        actx = ary.array_context

        @memoize_in(actx,
                    (DirectDiscretizationConnection, "resample_by_mat_knl"))
        def mat_knl():
            knl = make_loopy_program(
                """{[iel, idof, j]:
                    0<=iel<nelements and
                    0<=idof<n_to_nodes and
                    0<=j<n_from_nodes}""",
                "result[to_element_indices[iel], idof] \
                    = sum(j, resample_mat[idof, j] \
                    * ary[from_element_indices[iel], j])", [
                    lp.GlobalArg("result",
                                 None,
                                 shape="nelements_result, n_to_nodes",
                                 offset=lp.auto),
                    lp.GlobalArg("ary",
                                 None,
                                 shape="nelements_vec, n_from_nodes",
                                 offset=lp.auto),
                    lp.ValueArg("nelements_result", np.int32),
                    lp.ValueArg("nelements_vec", np.int32),
                    "...",
                ],
                name="resample_by_mat")

            return knl

        @memoize_in(actx,
                    (DirectDiscretizationConnection, "resample_by_picking_knl")
                    )
        def pick_knl():
            knl = make_loopy_program(
                """{[iel, idof]:
                    0<=iel<nelements and
                    0<=idof<n_to_nodes}""",
                "result[to_element_indices[iel], idof] \
                    = ary[from_element_indices[iel], pick_list[idof]]", [
                    lp.GlobalArg("result",
                                 None,
                                 shape="nelements_result, n_to_nodes",
                                 offset=lp.auto),
                    lp.GlobalArg("ary",
                                 None,
                                 shape="nelements_vec, n_from_nodes",
                                 offset=lp.auto),
                    lp.ValueArg("nelements_result", np.int32),
                    lp.ValueArg("nelements_vec", np.int32),
                    lp.ValueArg("n_from_nodes", np.int32),
                    "...",
                ],
                name="resample_by_picking")

            return knl

        if self.is_surjective:
            result = self.to_discr.empty(actx, dtype=ary.entry_dtype)
        else:
            result = self.to_discr.zeros(actx, dtype=ary.entry_dtype)

        for i_tgrp, cgrp in enumerate(self.groups):
            for i_batch, batch in enumerate(cgrp.batches):
                if not len(batch.from_element_indices):
                    continue

                point_pick_indices = self._resample_point_pick_indices(
                    actx, i_tgrp, i_batch)

                if point_pick_indices is None:
                    actx.call_loopy(
                        mat_knl(),
                        resample_mat=self._resample_matrix(
                            actx, i_tgrp, i_batch),
                        result=result[i_tgrp],
                        ary=ary[batch.from_group_index],
                        from_element_indices=batch.from_element_indices,
                        to_element_indices=batch.to_element_indices)

                else:
                    actx.call_loopy(
                        pick_knl(),
                        pick_list=point_pick_indices,
                        result=result[i_tgrp],
                        ary=ary[batch.from_group_index],
                        from_element_indices=batch.from_element_indices,
                        to_element_indices=batch.to_element_indices)

        return result
示例#15
0
 def __getattr__(self, name):
     return map_array_container(lambda ary: getattr(ary, name), self)
示例#16
0
    def __call__(self, ary):
        from meshmode.dof_array import DOFArray
        if is_array_container(ary) and not isinstance(ary, DOFArray):
            return map_array_container(self, ary)

        if not isinstance(ary, DOFArray):
            raise TypeError("non-array passed to discretization connection")

        if ary.shape != (len(self.from_discr.groups), ):
            raise ValueError("invalid shape of incoming resampling data")

        actx = ary.array_context

        @memoize_in(actx, (L2ProjectionInverseDiscretizationConnection,
                           "conn_projection_knl"))
        def kproj():
            return make_loopy_program(
                [
                    "{[iel_init]: 0 <= iel_init < n_to_elements}",
                    "{[idof_init]: 0 <= idof_init < n_to_nodes}",
                    "{[iel]: 0 <= iel < nelements}",
                    "{[i_quad]: 0 <= i_quad < n_to_nodes}",
                    "{[ibasis]: 0 <= ibasis < n_to_nodes}"
                ],
                """
                    result[iel_init, idof_init] = 0 {id=init}
                    ... gbarrier {id=barrier, dep=init}
                    result[to_element_indices[iel], ibasis] =               \
                        result[to_element_indices[iel], ibasis] +           \
                        sum(i_quad, ary[from_element_indices[iel], i_quad]  \
                                    * basis_tabulation[ibasis, i_quad]      \
                                    * weights[i_quad]) {dep=barrier}
                """, [
                    lp.GlobalArg("ary",
                                 None,
                                 shape=("n_from_elements", "n_from_nodes")),
                    lp.GlobalArg(
                        "result", None, shape=("n_to_elements", "n_to_nodes")),
                    lp.GlobalArg("basis_tabulation",
                                 None,
                                 shape=("n_to_nodes", "n_to_nodes")),
                    lp.GlobalArg("weights", None, shape="n_from_nodes"),
                    lp.ValueArg("n_from_elements", np.int32),
                    lp.ValueArg("n_from_nodes", np.int32),
                    lp.ValueArg("n_to_elements", np.int32),
                    lp.ValueArg("n_to_nodes", np.int32), "..."
                ],
                name="conn_projection_knl")

        # compute weights on each refinement of the reference element
        weights = self._batch_weights(actx)

        # perform dot product (on reference element) to get basis coefficients
        c_group_data = []
        for igrp, cgrp in enumerate(self.conn.groups):
            c_batch_data = []
            for ibatch, batch in enumerate(cgrp.batches):
                sgrp = self.from_discr.groups[batch.from_group_index]

                # Generate the basis tabulation matrix
                tabulations = []
                for basis_fn in sgrp.basis_obj().functions:
                    tabulations.append(
                        basis_fn(batch.result_unit_nodes).flatten())
                tabulations = actx.from_numpy(np.asarray(tabulations))

                # NOTE: batch.*_element_indices are reversed here because
                # they are from the original forward connection, but
                # we are going in reverse here. a bit confusing, but
                # saves on recreating the connection groups and batches.
                c_batch_data.append(
                    actx.call_loopy(
                        kproj(),
                        ary=ary[sgrp.index],
                        basis_tabulation=tabulations,
                        weights=weights[igrp, ibatch],
                        from_element_indices=batch.to_element_indices,
                        to_element_indices=batch.from_element_indices,
                        n_to_elements=self.to_discr.groups[igrp].nelements,
                        n_to_nodes=self.to_discr.groups[igrp].nunit_dofs,
                    )["result"])

            c_group_data.append(sum(c_batch_data))
        coefficients = DOFArray(actx, data=tuple(c_group_data))

        @keyed_memoize_in(actx, (L2ProjectionInverseDiscretizationConnection,
                                 "vandermonde_matrix"),
                          lambda grp: grp.discretization_key())
        def vandermonde_matrix(grp):
            from modepy import vandermonde
            vdm = vandermonde(grp.basis_obj().functions, grp.unit_nodes)
            return actx.from_numpy(vdm)

        return DOFArray(
            actx,
            data=tuple(
                actx.einsum("ij,ej->ei",
                            vandermonde_matrix(grp),
                            c_i,
                            arg_names=("vdm", "coeffs"),
                            tagged=(FirstAxisIsElementsTag(), ))
                for grp, c_i in zip(self.to_discr.groups, coefficients)))
示例#17
0
文件: op.py 项目: nchristensen/grudge
def weak_local_d_dx(dcoll: DiscretizationCollection,
                    *args) -> ArrayOrContainerT:
    r"""Return the element-local weak derivative along axis *xyz_axis* of the
    volume function represented by *vec*.

    May be called with ``(xyz_axis, vec)`` or ``(dd_in, xyz_axis, vec)``.

    Specifically, this function computes the volume contribution of the
    weak derivative in the :math:`i`-th component (specified by *xyz_axis*)
    of a function :math:`f`, in each element :math:`E`, with respect to polynomial
    test functions :math:`\phi`:

    .. math::

        \int_E \partial_i\phi\,f\,\mathrm{d}x \sim
        \mathbf{D}_{E,i}^T \mathbf{M}_{E}^T\mathbf{f}|_E,

    where :math:`\mathbf{D}_{E,i}` is the polynomial differentiation matrix on
    an :math:`E` for the :math:`i`-th spatial coordinate, :math:`\mathbf{M}_E`
    is the elemental mass matrix (see :func:`mass` for more information), and
    :math:`\mathbf{f}|_E` is a vector of coefficients for :math:`f` on :math:`E`.

    :arg dd_in: a :class:`~grudge.dof_desc.DOFDesc`, or a value convertible to one.
        Defaults to the base volume discretization if not provided.
    :arg xyz_axis: an integer indicating the axis along which the derivative
        is taken.
    :arg vec: a :class:`~meshmode.dof_array.DOFArray` or an
        :class:`~arraycontext.container.ArrayContainer` of them.
    :returns: a :class:`~meshmode.dof_array.DOFArray` or an
        :class:`~arraycontext.container.ArrayContainer` of them.
    """
    if len(args) == 2:
        xyz_axis, vec = args
        dd_in = dof_desc.DOFDesc("vol", dof_desc.DISCR_TAG_BASE)
    elif len(args) == 3:
        dd_in, xyz_axis, vec = args
    else:
        raise TypeError("invalid number of arguments")

    if not isinstance(vec, DOFArray):
        return map_array_container(
            partial(weak_local_d_dx, dcoll, dd_in, xyz_axis), vec)

    from grudge.geometry import inverse_surface_metric_derivative_mat

    in_discr = dcoll.discr_from_dd(dd_in)
    out_discr = dcoll.discr_from_dd(dof_desc.DD_VOLUME)

    actx = vec.array_context
    inverse_jac_mat = inverse_surface_metric_derivative_mat(
        actx,
        dcoll,
        dd=dd_in,
        times_area_element=True,
        _use_geoderiv_connection=actx.supports_nonscalar_broadcasting)

    return _single_axis_derivative_kernel(
        actx,
        out_discr,
        in_discr,
        _reference_stiffness_transpose_matrix,
        inverse_jac_mat,
        xyz_axis,
        vec,
        metric_in_matvec=True)