def _apply_inverse_mass_operator(dcoll: DiscretizationCollection, dd_out, dd_in, vec): if isinstance(vec, np.ndarray): return obj_array_vectorize( lambda vi: _apply_inverse_mass_operator(dcoll, dd_out, dd_in, vi), 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) 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 _apply_mass_operator(dcoll: DiscretizationCollection, dd_out, dd_in, vec): if isinstance(vec, np.ndarray): return obj_array_vectorize( lambda vi: _apply_mass_operator(dcoll, dd_out, dd_in, vi), 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) 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 rand(): from meshmode.dof_array import DOFArray return DOFArray( actx, tuple( actx.from_numpy(np.random.rand(grp.nelements, grp.nunit_dofs)) for grp in discr.discr_from_dd("vol").groups))
def test_flatten_unflatten(actx_factory): actx = actx_factory() ambient_dim = 2 from meshmode.mesh.generation import generate_regular_rect_mesh mesh = generate_regular_rect_mesh(a=(-0.5, ) * ambient_dim, b=(+0.5, ) * ambient_dim, n=(3, ) * ambient_dim, order=1) discr = Discretization(actx, mesh, PolynomialWarpAndBlendGroupFactory(3)) a = np.random.randn(discr.ndofs) from meshmode.dof_array import flatten, unflatten a_round_trip = actx.to_numpy( flatten(unflatten(actx, discr, actx.from_numpy(a)))) assert np.array_equal(a, a_round_trip) from meshmode.dof_array import flatten_to_numpy, unflatten_from_numpy a_round_trip = flatten_to_numpy(actx, unflatten_from_numpy(actx, discr, a)) assert np.array_equal(a, a_round_trip) x = thaw(discr.nodes(), actx) avg_mass = DOFArray( actx, tuple([(np.pi + actx.zeros((grp.nelements, 1), a.dtype)) for grp in discr.groups])) c = MyContainer(name="flatten", mass=avg_mass, momentum=make_obj_array([x, x, x]), enthalpy=x) from meshmode.dof_array import unflatten_like c_round_trip = unflatten_like(actx, flatten(c), c) assert flat_norm(c - c_round_trip) < 1.0e-8
def _apply_stiffness_transpose_operator(dcoll: DiscretizationCollection, dd_out, dd_in, vec, xyz_axis): from grudge.geometry import \ inverse_surface_metric_derivative, 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) inverse_jac_t = actx.np.stack([ inverse_surface_metric_derivative(actx, dcoll, rst_axis, xyz_axis, dd=dd_in) for rst_axis in range(dcoll.dim) ]) return DOFArray( actx, data=tuple( actx.einsum( "dij,ej,ej,dej->ei", reference_stiffness_transpose_matrix( actx, out_element_group=out_grp, in_element_group=in_grp), ae_i, vec_i, inv_jac_t_i, arg_names=("ref_stiffT_mat", "jac", "vec", "inv_jac_t"), tagged=(FirstAxisIsElementsTag(), )) for out_grp, in_grp, vec_i, ae_i, inv_jac_t_i in zip( out_discr.groups, in_discr.groups, vec, area_elements, inverse_jac_t)))
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 _single_axis_derivative_kernel(actx, out_discr, in_discr, get_diff_mat, inv_jac_mat, xyz_axis, vec, *, metric_in_matvec): # This gets used from both the strong and the weak derivative. These differ # in three ways: # - which differentiation matrix gets used, # - whether inv_jac_mat is pre-multiplied by a factor that includes the # area element, and # - 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". 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(), )) for out_grp, in_grp, vec_i, ijm_i in zip( out_discr.groups, in_discr.groups, vec, inv_jac_mat)))
def _interleave_dof_arrays(self, ary1, ary2): if not isinstance(ary1, DOFArray) or not isinstance(ary2, DOFArray): raise TypeError("non-array passed to connection") @memoize_in(self.array_context, (CenterGranularityConnection, "interleave")) def prg(): from meshmode.array_context import make_loopy_program return make_loopy_program( """{[iel, idof]: 0<=iel<nelements and 0<=idof<nunit_dofs}""", """ dst[iel, 2*idof] = src1[iel, idof] dst[iel, 2*idof + 1] = src2[iel, idof] """, [ lp.GlobalArg("src1", shape="(nelements, nunit_dofs)"), lp.GlobalArg("src2", shape="(nelements, nunit_dofs)"), lp.GlobalArg("dst", shape="(nelements, 2*nunit_dofs)"), "...", ], name="interleave") results = [] for grp, src1, src2 in zip(self.discr.groups, ary1, ary2): if src1.dtype != src2.dtype: raise ValueError("dtype mismatch in inputs") result = self.array_context.empty( (grp.nelements, 2 * grp.nunit_dofs), dtype=src1.dtype) self.array_context.call_loopy(prg(), src1=src1, src2=src2, dst=result, nelements=grp.nelements, nunit_dofs=grp.nunit_dofs) results.append(result) return DOFArray(self.array_context, tuple(results))
def nodes(self): r""" :returns: object array of shape ``(ambient_dim,)`` containing :class:`~meshmode.dof_array.DOFArray`\ s of node coordinates. """ actx = self._setup_actx @memoize_in(actx, (Discretization, "nodes_prg")) def prg(): return make_loopy_program("""{[iel,idof,j]: 0<=iel<nelements and 0<=idof<ndiscr_nodes and 0<=j<nmesh_nodes}""", """ result[iel, idof] = \ sum(j, resampling_mat[idof, j] * nodes[iel, j]) """, name="nodes") return make_obj_array([ _DOFArray.from_list(None, [ actx.freeze( actx.call_loopy( prg(), resampling_mat=actx.from_numpy( grp.from_mesh_interp_matrix()), nodes=actx.from_numpy( grp.mesh_el_group.nodes[iaxis]))["result"]) for grp in self.groups ]) for iaxis in range(self.ambient_dim) ])
def evaluate_sphere_eigf(actx, discr, m: int, n: int) -> DOFArray: assert discr.ambient_dim == 3 # {{{ get spherical coordinates from arraycontext import thaw x, y, z = thaw(discr.nodes(), actx) theta = actx.np.arctan2(actx.np.sqrt(x**2 + y**2), z) phi = actx.np.arctan2(y, x) # }}} # {{{ evaluate Y^m_n from scipy.special import sph_harm # pylint: disable=no-name-in-module y_mn = [] for gtheta, gphi in zip(theta, phi): result = sph_harm(m, n, actx.to_numpy(gphi), actx.to_numpy(gtheta)) y_mn.append(actx.from_numpy(result.real.copy())) # }}} return DOFArray(actx, tuple(y_mn))
def num_reference_derivative(self, ref_axes, vec): actx = vec.array_context @memoize_in(actx, (Discretization, "reference_derivative_prg")) def prg(): return make_loopy_program( """{[iel,idof,j]: 0<=iel<nelements and 0<=idof,j<nunit_dofs}""", "result[iel,idof] = sum(j, diff_mat[idof, j] * vec[iel, j])", name="diff") def get_mat(grp): mat = None for ref_axis in ref_axes: next_mat = grp.diff_matrices()[ref_axis] if mat is None: mat = next_mat else: mat = np.dot(next_mat, mat) return mat return _DOFArray.from_list(actx, [ actx.call_loopy(prg(), diff_mat=actx.from_numpy(get_mat(grp)), vec=vec[grp.index])["result"] for grp in self.groups ])
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 test_elementwise_reductions(actx_factory): actx = actx_factory() from mesh_data import BoxMeshBuilder builder = BoxMeshBuilder(ambient_dim=1) nelements = 4 mesh = builder.get_mesh(nelements, builder.mesh_order) dcoll = DiscretizationCollection(actx, mesh, order=builder.order) x = thaw(dcoll.nodes(), actx) def f(x): return actx.np.sin(x[0]) field = f(x) mins = [] maxs = [] sums = [] for gidx, grp_f in enumerate(field): min_res = np.empty(grp_f.shape) max_res = np.empty(grp_f.shape) sum_res = np.empty(grp_f.shape) for eidx in range(dcoll._volume_discr.groups[gidx].nelements): element_data = actx.to_numpy(grp_f[eidx]) min_res[eidx, :] = np.min(element_data) max_res[eidx, :] = np.max(element_data) sum_res[eidx, :] = np.sum(element_data) mins.append(actx.from_numpy(min_res)) maxs.append(actx.from_numpy(max_res)) sums.append(actx.from_numpy(sum_res)) from meshmode.dof_array import DOFArray, flat_norm ref_mins = DOFArray(actx, data=tuple(mins)) ref_maxs = DOFArray(actx, data=tuple(maxs)) ref_sums = DOFArray(actx, data=tuple(sums)) elem_mins = op.elementwise_min(dcoll, field) elem_maxs = op.elementwise_max(dcoll, field) elem_sums = op.elementwise_sum(dcoll, field) assert flat_norm(elem_mins - ref_mins, ord=np.inf) < 1.e-15 assert flat_norm(elem_maxs - ref_maxs, ord=np.inf) < 1.e-15 assert flat_norm(elem_sums - ref_sums, ord=np.inf) < 1.e-15
def test_modal_coefficients_by_projection(actx_factory, quad_group_factory): group_cls = SimplexElementGroup modal_group_factory = ModalSimplexGroupFactory actx = actx_factory() order = 10 m_order = 5 # Make a regular rectangle mesh mesh = mgen.generate_regular_rect_mesh(a=(0, 0), b=(5, 3), npoints_per_axis=(10, 6), order=order, group_cls=group_cls) # Make discretizations nodal_disc = Discretization(actx, mesh, quad_group_factory(order)) modal_disc = Discretization(actx, mesh, modal_group_factory(m_order)) # Make connections one using quadrature projection nodal_to_modal_conn_quad = NodalToModalDiscretizationConnection( nodal_disc, modal_disc, allow_approximate_quad=True) def f(x): return 2 * actx.np.sin(5 * x) x_nodal = thaw(nodal_disc.nodes()[0], actx) nodal_f = f(x_nodal) # Compute modal coefficients we expect to get import modepy as mp grp, = nodal_disc.groups shape = mp.Simplex(grp.dim) space = mp.space_for_shape(shape, order=m_order) basis = mp.orthonormal_basis_for_space(space, shape) quad = grp.quadrature_rule() nodal_f_data = actx.to_numpy(nodal_f[0]) vdm = mp.vandermonde(basis.functions, quad.nodes) w_diag = np.diag(quad.weights) modal_data = [] for _, nodal_data in enumerate(nodal_f_data): # Compute modal data in each element: V.T * W * nodal_data elem_modal_f = np.dot(vdm.T, np.dot(w_diag, nodal_data)) modal_data.append(elem_modal_f) modal_data = actx.from_numpy(np.asarray(modal_data)) modal_f_expected = DOFArray(actx, data=(modal_data, )) # Map nodal coefficients using the quadrature-based projection modal_f_computed = nodal_to_modal_conn_quad(nodal_f) err = flat_norm(modal_f_expected - modal_f_computed) assert err <= 1e-13
def _compute_characteristic_lengthscales(): return freeze( DOFArray( actx, data=tuple( # Scale each group array of geometric factors by the # corresponding group non-geometric factor cng * geo_facts for cng, geo_facts in zip( dt_non_geometric_factors(dcoll), thaw(dt_geometric_factors(dcoll), actx)))))
def _get_ref_data(field): mins = [] maxs = [] sums = [] for grp_f in field: min_res = np.empty(grp_f.shape) max_res = np.empty(grp_f.shape) sum_res = np.empty(grp_f.shape) for eidx in range(dcoll.mesh.nelements): element_data = actx.to_numpy(grp_f[eidx]) min_res[eidx, :] = np.min(element_data) max_res[eidx, :] = np.max(element_data) sum_res[eidx, :] = np.sum(element_data) mins.append(actx.from_numpy(min_res)) maxs.append(actx.from_numpy(max_res)) sums.append(actx.from_numpy(sum_res)) min_field = DOFArray(actx, data=tuple(mins)) max_field = DOFArray(actx, data=tuple(maxs)) sums_field = DOFArray(actx, data=tuple(sums)) return min_field, max_field, sums_field
def _new_array(self, actx, creation_func, dtype=None): if dtype is None: dtype = self.real_dtype elif dtype == "c": dtype = self.complex_dtype else: dtype = np.dtype(dtype) return _DOFArray.from_list(actx, [ creation_func(shape=(grp.nelements, grp.nunit_dofs), dtype=dtype) for grp in self.groups ])
def _apply_elementwise_reduction(op_name: str, dcoll: DiscretizationCollection, *args) -> DOFArray: r"""Returns a vector of DOFs with all entries on each element set to the reduction operation *op_name* over all degrees of freedom. :arg \*args: Arguments for the reduction operator, such as *dd* and *vec*. :returns: a :class:`~meshmode.dof_array.DOFArray` or object arrary of :class:`~meshmode.dof_array.DOFArray`s. """ 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 isinstance(vec, np.ndarray): return obj_array_vectorize( lambda vi: _apply_elementwise_reduction(op_name, dcoll, dd, vi), vec) actx = vec.array_context @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 _map_elementwise_reduction(self, reduction_name, expr): @memoize_in(self.places, "elementwise_node_" + reduction_name) def node_knl(): from meshmode.array_context import make_loopy_program return make_loopy_program("""{[iel, idof, jdof]: 0<=iel<nelements and 0<=idof, jdof<ndofs}""", """ result[iel, idof] = %s(jdof, operand[iel, jdof]) """ % reduction_name, name="nodewise_reduce") @memoize_in(self.places, "elementwise_" + reduction_name) def element_knl(): from meshmode.array_context import make_loopy_program return make_loopy_program("""{[iel, jdof]: 0<=iel<nelements and 0<=jdof<ndofs} """, """ result[iel, 0] = %s(jdof, operand[iel, jdof]) """ % reduction_name, name="elementwise_reduce") discr = self.places.get_discretization(expr.dofdesc.geometry, expr.dofdesc.discr_stage) operand = self.rec(expr.operand) assert operand.shape == (len(discr.groups), ) def _reduce(knl, result): for grp in discr.groups: self.array_context.call_loopy(knl, operand=operand[grp.index], result=result[grp.index]) return result dtype = operand.entry_dtype granularity = expr.dofdesc.granularity if granularity is sym.GRANULARITY_NODE: return _reduce(node_knl(), discr.empty(self.array_context, dtype=dtype)) elif granularity is sym.GRANULARITY_ELEMENT: result = DOFArray( self.array_context, tuple([ self.array_context.empty((grp.nelements, 1), dtype=dtype) for grp in discr.groups ])) return _reduce(element_knl(), result) else: raise ValueError(f"unsupported granularity: {granularity}")
def _apply_face_mass_operator(dcoll: DiscretizationCollection, dd, vec): if isinstance(vec, np.ndarray): return obj_array_vectorize( lambda vi: _apply_face_mass_operator(dcoll, dd, vi), 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 @memoize_in(actx, (_apply_face_mass_operator, "face_mass_knl")) def prg(): t_unit = make_loopy_program([ "{[iel]: 0 <= iel < nelements}", "{[f]: 0 <= f < nfaces}", "{[idof]: 0 <= idof < nvol_nodes}", "{[jdof]: 0 <= jdof < nface_nodes}" ], """ result[iel, idof] = sum(f, sum(jdof, mat[idof, f, jdof] * jac_surf[f, iel, jdof] * vec[f, iel, jdof])) """, name="face_mass") import loopy as lp from meshmode.transform_metadata import (ConcurrentElementInameTag, ConcurrentDOFInameTag) return lp.tag_inames(t_unit, { "iel": ConcurrentElementInameTag(), "idof": ConcurrentDOFInameTag() }) assert len(face_discr.groups) == len(volm_discr.groups) surf_area_elements = area_element(actx, dcoll, dd=dd) return DOFArray( actx, data=tuple( actx.call_loopy( prg(), mat=reference_face_mass_matrix(actx, face_element_group=afgrp, vol_element_group=vgrp, dtype=dtype), jac_surf=surf_ae_i.reshape(vgrp.mesh_el_group.nfaces, vgrp.nelements, afgrp.nunit_dofs), vec=vec_i.reshape(vgrp.mesh_el_group.nfaces, vgrp.nelements, afgrp.nunit_dofs))["result"] for vgrp, afgrp, vec_i, surf_ae_i in zip(volm_discr.groups, face_discr.groups, vec, surf_area_elements)))
def map_insn_loopy_kernel(self, insn, profile_data=None): kdescr = insn.kernel_descriptor discr = self.discrwb.discr_from_dd(kdescr.governing_dd) dof_array_kwargs = {} other_kwargs = {} for name, expr in kdescr.input_mappings.items(): v = self.rec(expr) if isinstance(v, DOFArray): dof_array_kwargs[name] = v else: other_kwargs[name] = v for name in kdescr.scalar_args(): v = other_kwargs[name] if isinstance(v, (int, float)): other_kwargs[name] = discr.real_dtype.type(v) elif isinstance(v, complex): other_kwargs[name] = discr.complex_dtype.type(v) elif isinstance(v, np.number): pass else: raise ValueError( "unrecognized scalar type for variable '%s': %s" % (name, type(v))) result = {} for grp in discr.groups: kwargs = other_kwargs.copy() kwargs["nelements"] = grp.nelements kwargs["nunit_dofs"] = grp.nunit_dofs for name, ary in dof_array_kwargs.items(): kwargs[name] = ary[grp.index] knl_result = self.array_context.call_loopy(kdescr.loopy_kernel, **kwargs) for name, val in knl_result.items(): result.setdefault(name, []).append(val) result = { name: DOFArray(self.array_context, tuple(val)) for name, val in result.items() } return list(result.items()), []
def _interleave_dof_arrays(self, ary1, ary2): if not isinstance(ary1, DOFArray) or not isinstance(ary2, DOFArray): raise TypeError("non-array passed to connection") if ary1.array_context is not ary2.array_context: raise ValueError("array context of the two arguments must match") actx = ary1.array_context @memoize_in(actx, (CenterGranularityConnection, "interleave")) def prg(): from arraycontext import make_loopy_program t_unit = make_loopy_program( "{[iel, idof]: 0 <= iel < nelements and 0 <= idof < nunit_dofs}", """ result[iel, 2*idof] = ary1[iel, idof] result[iel, 2*idof + 1] = ary2[iel, idof] """, [ lp.GlobalArg("ary1", shape="(nelements, nunit_dofs)"), lp.GlobalArg("ary2", shape="(nelements, nunit_dofs)"), lp.GlobalArg("result", shape="(nelements, 2*nunit_dofs)"), ... ], name="interleave") from meshmode.transform_metadata import (ConcurrentElementInameTag, ConcurrentDOFInameTag) return lp.tag_inames( t_unit, { "iel": ConcurrentElementInameTag(), "idof": ConcurrentDOFInameTag() }) results = [] for grp, subary1, subary2 in zip(self.discr.groups, ary1, ary2): if subary1.dtype != subary2.dtype: raise ValueError( "dtype mismatch in inputs: " f"'{subary1.dtype.name}' and '{subary2.dtype.name}'") result = actx.call_loopy(prg(), ary1=subary1, ary2=subary2, nelements=grp.nelements, nunit_dofs=grp.nunit_dofs)["result"] results.append(result) return DOFArray(actx, tuple(results))
def apply_spectral_filter(actx, modal_field, discr, cutoff, mode_response_function): r"""Apply the spectral filter, defined by the *mode_response_function*. This routine returns filtered data in the modal basis, which has been applied using a user-provided *mode_response_function* to dampen modes beyond the user-provided *cutoff*. Parameters ---------- actx: :class:`arraycontext.ArrayContext` A :class:`arraycontext.ArrayContext` associated with an array of degrees of freedom modal_field: numpy.ndarray DOFArray or object array of DOFArrays denoting the modal data discr: :class:`meshmode.discretization.Discretization` A :class:`meshmode.discretization.Discretization` describing the volume discretization the *modal_field* comes from. cutoff: int Mode cutoff beyond which the filter will be applied, and below which the filter will preserve. mode_response_function: A function that returns a filter weight for for each mode id. Returns ------- modal_field: :class:`meshmode.dof_array.DOFArray` DOFArray or object array of DOFArrays """ from meshmode.transform_metadata import FirstAxisIsElementsTag return DOFArray( actx, tuple(actx.einsum("j,ej->ej", make_spectral_filter( actx, group=grp, cutoff=cutoff, mode_response_function=mode_response_function ), vec_i, arg_names=("filter", "vec"), tagged=(FirstAxisIsElementsTag(),)) for grp, vec_i in zip(discr.groups, modal_field)) )
def _visualize_refinement(actx: PyOpenCLArrayContext, discr, niter, stage_nr, stage_name, flags, visualize=False): if not visualize: return if stage_nr not in (1, 2): raise ValueError("unexpected stage number") flags = flags.get() logger.info("for stage %s: splitting %d/%d stage-%d elements", stage_name, np.sum(flags), discr.mesh.nelements, stage_nr) from meshmode.discretization.visualization import make_visualizer vis = make_visualizer(actx, discr, 3) assert len(flags) == discr.mesh.nelements flags = flags.astype(bool) nodes_flags_template = discr.zeros(actx) nodes_flags = [] for grp in discr.groups: meg = grp.mesh_el_group nodes_flags_grp = actx.to_numpy(nodes_flags_template[grp.index]) nodes_flags_grp[flags[meg.element_nr_base:meg.nelements + meg.element_nr_base]] = 1 nodes_flags.append(actx.from_numpy(nodes_flags_grp)) nodes_flags = DOFArray(actx, tuple(nodes_flags)) vis_data = [ ("refine_flags", nodes_flags), ] if 0: from pytential import sym, bind bdry_normals = bind(discr, sym.normal( discr.ambient_dim))(actx).as_vector(dtype=object) vis_data.append(("bdry_normals", bdry_normals), ) vis.write_vtk_file(f"refinement-{stage_name}-{niter:03d}.vtu", vis_data)
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 map_insn_diff_batch_assign(self, insn, profile_data=None): field = self.rec(insn.field) repr_op = insn.operators[0] assert repr_op.dd_in.domain_tag == repr_op.dd_out.domain_tag in_discr = self.discrwb.discr_from_dd(repr_op.dd_in) out_discr = self.discrwb.discr_from_dd(repr_op.dd_out) prg = self._elwise_linear_loopy_prg() result = [] for name, op in zip(insn.names, insn.operators): group_results = [] for in_grp, out_grp in zip(in_discr.groups, out_discr.groups): assert in_grp.nelements == out_grp.nelements if in_grp.nelements == 0: continue # Cache operator cache_key = "diff_batch", in_grp, out_grp, tuple(insn.operators),\ field.entry_dtype try: matrices_dev = self.bound_op.operator_data_cache[cache_key] except KeyError: matrices_dev = [ self.array_context.from_numpy(mat) for mat in repr_op.matrices(out_grp, in_grp) ] self.bound_op.operator_data_cache[cache_key] = matrices_dev group_results.append( self.array_context.call_loopy( prg, mat=matrices_dev[op.rst_axis], vec=field[in_grp.index])["result"]) result.append( (name, DOFArray(self.array_context, tuple(group_results)))) return result, []
def quad_weights(self): """:returns: A :class:`~meshmode.dof_array.DOFArray` with quadrature weights. """ actx = self._setup_actx @memoize_in(actx, (Discretization, "quad_weights_prg")) def prg(): return make_loopy_program( "{[iel,idof]: 0<=iel<nelements and 0<=idof<nunit_dofs}", "result[iel,idof] = weights[idof]", name="quad_weights") return _DOFArray.from_list(None, [ actx.freeze( actx.call_loopy( prg(), weights=actx.from_numpy(grp.weights), nelements=grp.nelements, )["result"]) for grp in self.groups ])
def _compute_local_gradient(dcoll: DiscretizationCollection, vec, xyz_axis): from grudge.geometry import inverse_surface_metric_derivative discr = dcoll.discr_from_dd(dof_desc.DD_VOLUME) actx = vec.array_context inverse_jac_t = actx.np.stack([ inverse_surface_metric_derivative(actx, dcoll, rst_axis, xyz_axis) for rst_axis in range(dcoll.dim) ]) return DOFArray(actx, data=tuple( actx.einsum("dei,dij,ej->ei", inv_jac_t_i, reference_derivative_matrices(actx, grp), vec_i, arg_names=("inv_jac_t", "ref_diff_mat", "vec"), tagged=(FirstAxisIsElementsTag(), )) for grp, vec_i, inv_jac_t_i in zip( discr.groups, vec, inverse_jac_t)))
def test_node_reduction(ctx_factory): ctx = ctx_factory() queue = cl.CommandQueue(ctx) actx = PyOpenCLArrayContext(queue) # {{{ build discretization target_order = 4 nelements = 32 mesh = make_curve_mesh(starfish, np.linspace(0.0, 1.0, nelements + 1), target_order) discr = Discretization(actx, mesh, InterpolatoryQuadratureSimplexGroupFactory(target_order)) # }}} # {{{ test # create a shuffled [1, nelements + 1] array ary = [] el_nr_base = 0 for grp in discr.groups: x = 1 + np.arange(el_nr_base, grp.nelements) np.random.shuffle(x) ary.append(actx.freeze(actx.from_numpy(x.reshape(-1, 1)))) el_nr_base += grp.nelements from meshmode.dof_array import DOFArray ary = DOFArray(actx, tuple(ary)) for func, expected in [ (sym.NodeSum, nelements * (nelements + 1) // 2), (sym.NodeMax, nelements), (sym.NodeMin, 1), ]: r = bind(discr, func(sym.var("x")))(actx, x=ary) assert abs(r - expected) < 1.0e-15, r
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)))