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)))
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)))
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
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
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)
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)
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)
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)
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)))
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)
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))
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)
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))
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
def __getattr__(self, name): return map_array_container(lambda ary: getattr(ary, name), self)
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)))
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)