diff --git a/examples/wave/wave-eager-mpi.py b/examples/wave/wave-op-mpi.py similarity index 85% rename from examples/wave/wave-eager-mpi.py rename to examples/wave/wave-op-mpi.py index b6ca9ae386874607fdaf31533952a013dbbe8c1e..19193ca068dee7cdff6d2728cb47ae3c73ccba2c 100644 --- a/examples/wave/wave-eager-mpi.py +++ b/examples/wave/wave-op-mpi.py @@ -32,8 +32,8 @@ from meshmode.dof_array import thaw from meshmode.mesh import BTAG_ALL, BTAG_NONE # noqa -from grudge.eager import ( - EagerDGDiscretization, interior_trace_pair, cross_rank_trace_pairs) +from grudge.discretization import DGDiscretizationWithBoundaries +import grudge.op as op from grudge.shortcuts import make_visualizer from grudge.symbolic.primitives import TracePair from mpi4py import MPI @@ -45,7 +45,7 @@ def wave_flux(discr, c, w_tpair): u = w_tpair[0] v = w_tpair[1:] - normal = thaw(u.int.array_context, discr.normal(w_tpair.dd)) + normal = thaw(u.int.array_context, op.normal(discr, w_tpair.dd)) flux_weak = flat_obj_array( np.dot(v.avg, normal), @@ -59,32 +59,32 @@ def wave_flux(discr, c, w_tpair): 0.5*normal*v_jump, ) - return discr.project(w_tpair.dd, "all_faces", c*flux_weak) + return op.project(discr, w_tpair.dd, "all_faces", c*flux_weak) def wave_operator(discr, c, w): u = w[0] v = w[1:] - dir_u = discr.project("vol", BTAG_ALL, u) - dir_v = discr.project("vol", BTAG_ALL, v) + dir_u = op.project(discr, "vol", BTAG_ALL, u) + dir_v = op.project(discr, "vol", BTAG_ALL, v) dir_bval = flat_obj_array(dir_u, dir_v) dir_bc = flat_obj_array(-dir_u, dir_v) return ( - discr.inverse_mass( + op.inverse_mass(discr, flat_obj_array( - -c*discr.weak_div(v), - -c*discr.weak_grad(u) + -c*op.weak_div(discr, v), + -c*op.weak_grad(discr, u) ) + # noqa: W504 - discr.face_mass( - wave_flux(discr, c=c, w_tpair=interior_trace_pair(discr, w)) + op.face_mass(discr, + wave_flux(discr, c=c, w_tpair=op.interior_trace_pair(discr, w)) + wave_flux(discr, c=c, w_tpair=TracePair( BTAG_ALL, interior=dir_bval, exterior=dir_bc)) + sum( wave_flux(discr, c=c, w_tpair=tpair) - for tpair in cross_rank_trace_pairs(discr, w)) + for tpair in op.cross_rank_trace_pairs(discr, w)) ) ) ) @@ -105,7 +105,7 @@ def bump(actx, discr, t=0): source_width = 0.05 source_omega = 3 - nodes = thaw(actx, discr.nodes()) + nodes = thaw(actx, op.nodes(discr)) center_dist = flat_obj_array([ nodes[i] - source_center[i] for i in range(discr.dim) @@ -152,7 +152,7 @@ def main(): order = 3 - discr = EagerDGDiscretization(actx, local_mesh, order=order, + discr = DGDiscretizationWithBoundaries(actx, local_mesh, order=order, mpi_communicator=comm) if dim == 2: @@ -181,7 +181,7 @@ def main(): fields = rk4_step(fields, t, dt, rhs) if istep % 10 == 0: - print(istep, t, discr.norm(fields[0])) + print(istep, t, op.norm(discr, fields[0], p=2)) vis.write_parallel_vtk_file( comm, f"fld-wave-eager-mpi-{{rank:03d}}-{istep:04d}.vtu", diff --git a/examples/wave/wave-eager-var-velocity.py b/examples/wave/wave-op-var-velocity.py similarity index 84% rename from examples/wave/wave-eager-var-velocity.py rename to examples/wave/wave-op-var-velocity.py index 194a1d654e602d5add77ea6d5722849001ca87b0..496b180289b048a4e749b1394ad1bc63223925d4 100644 --- a/examples/wave/wave-eager-var-velocity.py +++ b/examples/wave/wave-op-var-velocity.py @@ -32,7 +32,8 @@ from meshmode.dof_array import thaw from meshmode.mesh import BTAG_ALL, BTAG_NONE # noqa -from grudge.eager import EagerDGDiscretization, interior_trace_pair +from grudge.discretization import DGDiscretizationWithBoundaries +import grudge.op as op from grudge.shortcuts import make_visualizer from grudge.symbolic.primitives import TracePair, QTAG_NONE, DOFDesc @@ -46,7 +47,7 @@ def wave_flux(discr, c, w_tpair): u = w_tpair[0] v = w_tpair[1:] - normal = thaw(u.int.array_context, discr.normal(dd)) + normal = thaw(u.int.array_context, op.normal(discr, dd)) flux_weak = flat_obj_array( np.dot(v.avg, normal), @@ -61,39 +62,39 @@ def wave_flux(discr, c, w_tpair): # FIXME this flux is only correct for continuous c dd_allfaces_quad = dd_quad.with_dtag("all_faces") - c_quad = discr.project("vol", dd_quad, c) - flux_quad = discr.project(dd, dd_quad, flux_weak) + c_quad = op.project(discr, "vol", dd_quad, c) + flux_quad = op.project(discr, dd, dd_quad, flux_weak) - return discr.project(dd_quad, dd_allfaces_quad, c_quad*flux_quad) + return op.project(discr, dd_quad, dd_allfaces_quad, c_quad*flux_quad) def wave_operator(discr, c, w): u = w[0] v = w[1:] - dir_u = discr.project("vol", BTAG_ALL, u) - dir_v = discr.project("vol", BTAG_ALL, v) + dir_u = op.project(discr, "vol", BTAG_ALL, u) + dir_v = op.project(discr, "vol", BTAG_ALL, v) dir_bval = flat_obj_array(dir_u, dir_v) dir_bc = flat_obj_array(-dir_u, dir_v) dd_quad = DOFDesc("vol", "vel_prod") - c_quad = discr.project("vol", dd_quad, c) - w_quad = discr.project("vol", dd_quad, w) + c_quad = op.project(discr, "vol", dd_quad, c) + w_quad = op.project(discr, "vol", dd_quad, w) u_quad = w_quad[0] v_quad = w_quad[1:] dd_allfaces_quad = DOFDesc("all_faces", "vel_prod") return ( - discr.inverse_mass( + op.inverse_mass(discr, flat_obj_array( - -discr.weak_div(dd_quad, c_quad*v_quad), - -discr.weak_grad(dd_quad, c_quad*u_quad) + -op.weak_div(discr, dd_quad, c_quad*v_quad), + -op.weak_grad(discr, dd_quad, c_quad*u_quad) ) + # noqa: W504 - discr.face_mass( + op.face_mass(discr, dd_allfaces_quad, - wave_flux(discr, c=c, w_tpair=interior_trace_pair(discr, w)) + wave_flux(discr, c=c, w_tpair=op.interior_trace_pair(discr, w)) + wave_flux(discr, c=c, w_tpair=TracePair( BTAG_ALL, interior=dir_bval, exterior=dir_bc)) )) @@ -117,7 +118,7 @@ def bump(actx, discr, t=0, width=0.05, center=None): center = center[:discr.dim] source_omega = 3 - nodes = thaw(actx, discr.nodes()) + nodes = thaw(actx, op.nodes(discr)) center_dist = flat_obj_array([ nodes[i] - center[i] for i in range(discr.dim) @@ -159,7 +160,7 @@ def main(): from meshmode.discretization.poly_element import \ QuadratureSimplexGroupFactory, \ PolynomialWarpAndBlendGroupFactory - discr = EagerDGDiscretization(actx, mesh, + discr = DGDiscretizationWithBoundaries(actx, mesh, quad_tag_to_group_factory={ QTAG_NONE: PolynomialWarpAndBlendGroupFactory(order), "vel_prod": QuadratureSimplexGroupFactory(3*order), @@ -185,7 +186,7 @@ def main(): fields = rk4_step(fields, t, dt, rhs) if istep % 10 == 0: - print(istep, t, discr.norm(fields[0])) + print(istep, t, op.norm(discr, fields[0], p=2)) vis.write_vtk_file("fld-wave-eager-var-velocity-%04d.vtu" % istep, [ ("c", c), diff --git a/examples/wave/wave-eager.py b/examples/wave/wave-op.py similarity index 83% rename from examples/wave/wave-eager.py rename to examples/wave/wave-op.py index a8c4438217f5a09c03ddd9c00f3e1209b2edd49a..cac70d6b5e24f2ce18071f7d1174fa2b8ca0d6ea 100644 --- a/examples/wave/wave-eager.py +++ b/examples/wave/wave-op.py @@ -32,7 +32,8 @@ from meshmode.dof_array import thaw from meshmode.mesh import BTAG_ALL, BTAG_NONE # noqa -from grudge.eager import EagerDGDiscretization, interior_trace_pair +from grudge.discretization import DGDiscretizationWithBoundaries +import grudge.op as op from grudge.shortcuts import make_visualizer from grudge.symbolic.primitives import TracePair @@ -43,7 +44,7 @@ def wave_flux(discr, c, w_tpair): u = w_tpair[0] v = w_tpair[1:] - normal = thaw(u.int.array_context, discr.normal(w_tpair.dd)) + normal = thaw(u.int.array_context, op.normal(discr, w_tpair.dd)) flux_weak = flat_obj_array( np.dot(v.avg, normal), @@ -56,27 +57,27 @@ def wave_flux(discr, c, w_tpair): 0.5*normal*np.dot(normal, v.ext-v.int), ) - return discr.project(w_tpair.dd, "all_faces", c*flux_weak) + return op.project(discr, w_tpair.dd, "all_faces", c*flux_weak) def wave_operator(discr, c, w): u = w[0] v = w[1:] - dir_u = discr.project("vol", BTAG_ALL, u) - dir_v = discr.project("vol", BTAG_ALL, v) + dir_u = op.project(discr, "vol", BTAG_ALL, u) + dir_v = op.project(discr, "vol", BTAG_ALL, v) dir_bval = flat_obj_array(dir_u, dir_v) dir_bc = flat_obj_array(-dir_u, dir_v) return ( - discr.inverse_mass( + op.inverse_mass(discr, flat_obj_array( - -c*discr.weak_div(v), - -c*discr.weak_grad(u) + -c*op.weak_div(discr, v), + -c*op.weak_grad(discr, u) ) + # noqa: W504 - discr.face_mass( - wave_flux(discr, c=c, w_tpair=interior_trace_pair(discr, w)) + op.face_mass(discr, + wave_flux(discr, c=c, w_tpair=op.interior_trace_pair(discr, w)) + wave_flux(discr, c=c, w_tpair=TracePair( BTAG_ALL, interior=dir_bval, exterior=dir_bc)) )) @@ -98,7 +99,7 @@ def bump(actx, discr, t=0): source_width = 0.05 source_omega = 3 - nodes = thaw(actx, discr.nodes()) + nodes = thaw(actx, op.nodes(discr)) center_dist = flat_obj_array([ nodes[i] - source_center[i] for i in range(discr.dim) @@ -137,7 +138,7 @@ def main(): print("%d elements" % mesh.nelements) - discr = EagerDGDiscretization(actx, mesh, order=order) + discr = DGDiscretizationWithBoundaries(actx, mesh, order=order) fields = flat_obj_array( bump(actx, discr), @@ -156,8 +157,8 @@ def main(): fields = rk4_step(fields, t, dt, rhs) if istep % 10 == 0: - print(f"step: {istep} t: {t} L2: {discr.norm(fields[0])} " - f"sol max: {discr.nodal_max('vol', fields[0])}") + print(f"step: {istep} t: {t} L2: {op.norm(discr, fields[0], 2)} " + f"sol max: {op.nodal_max(discr, 'vol', fields[0])}") vis.write_vtk_file("fld-wave-eager-%04d.vtu" % istep, [ ("u", fields[0]), diff --git a/grudge/eager.py b/grudge/eager.py index 7d9eb1d884844cfe34f8d3425df8b3e6bc44824e..516b56290bdb5ee67e705a35524eb1668cbd81f0 100644 --- a/grudge/eager.py +++ b/grudge/eager.py @@ -20,26 +20,8 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ - -import numpy as np # noqa -from pytools import memoize_method -from pytools.obj_array import obj_array_vectorize, make_obj_array -import pyopencl.array as cla # noqa -from grudge import sym, bind - -from meshmode.mesh import BTAG_ALL, BTAG_NONE, BTAG_PARTITION # noqa -from meshmode.dof_array import freeze, flatten, unflatten - +import grudge.op as op from grudge.discretization import DGDiscretizationWithBoundaries -from grudge.symbolic.primitives import TracePair - -from numbers import Number - -__doc__ = """ -.. autoclass:: EagerDGDiscretization -.. autofunction:: interior_trace_pair -.. autofunction:: cross_rank_trace_pairs -""" class EagerDGDiscretization(DGDiscretizationWithBoundaries): @@ -47,425 +29,68 @@ class EagerDGDiscretization(DGDiscretizationWithBoundaries): Inherits from :class:`~grudge.discretization.DGDiscretizationWithBoundaries`. .. automethod:: __init__ - .. automethod:: project - .. automethod:: nodes - - .. automethod:: grad - .. automethod:: d_dx - .. automethod:: div - - .. automethod:: weak_grad - .. automethod:: weak_d_dx - .. automethod:: weak_div - - .. automethod:: normal - .. automethod:: mass - .. automethod:: inverse_mass - .. automethod:: face_mass - - .. automethod:: norm - .. automethod:: nodal_sum - .. automethod:: nodal_min - .. automethod:: nodal_max """ - def interp(self, src, tgt, vec): + def __init__(self, *args, **kwargs): from warnings import warn - warn("using 'interp' is deprecated, use 'project' instead.", + warn("EagerDGDiscretization is deprecated and will go away in 2022. " + "Use the base DGDiscretizationWithBoundaries with grudge.op " + "instead.", DeprecationWarning, stacklevel=2) - return self.project(src, tgt, vec) + super().__init__(*args, **kwargs) def project(self, src, tgt, vec): - """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.sym.DOFDesc`, or a value convertible to one - :arg tgt: a :class:`~grudge.sym.DOFDesc`, or a value convertible to one - :arg vec: a :class:`~meshmode.dof_array.DOFArray` - """ - src = sym.as_dofdesc(src) - tgt = sym.as_dofdesc(tgt) - if src == tgt: - return vec - - if isinstance(vec, np.ndarray): - return obj_array_vectorize( - lambda el: self.project(src, tgt, el), vec) - - if isinstance(vec, Number): - return vec - - return self.connection_from_dds(src, tgt)(vec) + return op.project(self, src, tgt, vec) def nodes(self, dd=None): - r"""Return the nodes of a discretization. - - :arg dd: a :class:`~grudge.sym.DOFDesc`, or a value convertible to one. - Defaults to the base volume discretization. - :returns: an object array of :class:`~meshmode.dof_array.DOFArray`\ s - """ - if dd is None: - return self._volume_discr.nodes() - else: - return self.discr_from_dd(dd).nodes() - - # {{{ derivatives + return op.nodes(self, dd) - @memoize_method - def _bound_grad(self): - return bind(self, sym.nabla(self.dim) * sym.Variable("u"), local_only=True) + def normal(self, dd): + return op.normal(self, dd) def grad(self, vec): - r"""Return the gradient of the volume function represented by *vec*. - - :arg vec: a :class:`~meshmode.dof_array.DOFArray` - :returns: an object array of :class:`~meshmode.dof_array.DOFArray`\ s - """ - return self._bound_grad()(u=vec) - - @memoize_method - def _bound_d_dx(self, xyz_axis): - return bind(self, sym.nabla(self.dim)[xyz_axis] * sym.Variable("u"), - local_only=True) + return op.grad(self, vec) def d_dx(self, xyz_axis, vec): - r"""Return the derivative along axis *xyz_axis* of the volume function - represented by *vec*. - - :arg xyz_axis: an integer indicating the axis along which the derivative - is taken - :arg vec: a :class:`~meshmode.dof_array.DOFArray` - :returns: a :class:`~meshmode.dof_array.DOFArray`\ s - """ - return self._bound_d_dx(xyz_axis)(u=vec) - - def _div_helper(self, diff_func, vecs): - if not isinstance(vecs, np.ndarray): - raise TypeError("argument must be an object array") - assert vecs.dtype == np.object - - if vecs.shape[-1] != self.ambient_dim: - raise ValueError("last dimension of *vecs* argument must match " - "ambient dimension") - - if len(vecs.shape) == 1: - return sum(diff_func(i, vec_i) for i, vec_i in enumerate(vecs)) - else: - result = np.zeros(vecs.shape[:-1], dtype=object) - for idx in np.ndindex(vecs.shape[:-1]): - result[idx] = sum( - diff_func(i, vec_i) for i, vec_i in enumerate(vecs[idx])) - return result + return op.d_dx(self, xyz_axis, vec) def div(self, vecs): - r"""Return the divergence of the vector volume function - represented by *vecs*. - - :arg vec: an object array of - a :class:`~meshmode.dof_array.DOFArray`\ s, - where the last axis of the array must have length - matching the volume dimension. - :returns: a :class:`~meshmode.dof_array.DOFArray` - """ - - return self._div_helper( - lambda i, subvec: self.d_dx(i, subvec), - vecs) - - @memoize_method - def _bound_weak_grad(self, dd): - return bind(self, - sym.stiffness_t(self.dim, dd_in=dd) * sym.Variable("u", dd), - local_only=True) + return op.d_dx(self, vecs) def weak_grad(self, *args): - r"""Return the "weak gradient" of the volume function represented by - *vec*. - - May be called with ``(vecs)`` or ``(dd, vecs)``. - - :arg dd: a :class:`~grudge.sym.DOFDesc`, or a value convertible to one. - Defaults to the base volume discretization if not provided. - :arg vec: a :class:`~meshmode.dof_array.DOFArray` - :returns: an object array of :class:`~meshmode.dof_array.DOFArray`\ s - """ - 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") - - return self._bound_weak_grad(dd)(u=vec) - - @memoize_method - def _bound_weak_d_dx(self, dd, xyz_axis): - return bind(self, - sym.stiffness_t(self.dim, dd_in=dd)[xyz_axis] - * sym.Variable("u", dd), - local_only=True) + return op.weak_grad(self, *args) def weak_d_dx(self, *args): - r"""Return the derivative along axis *xyz_axis* of the volume function - represented by *vec*. - - May be called with ``(xyz_axis, vecs)`` or ``(dd, xyz_axis, vecs)``. - - :arg xyz_axis: an integer indicating the axis along which the derivative - is taken - :arg vec: a :class:`~meshmode.dof_array.DOFArray` - :returns: a :class:`~meshmode.dof_array.DOFArray`\ s - """ - if len(args) == 2: - xyz_axis, vec = args - dd = sym.DOFDesc("vol", sym.QTAG_NONE) - elif len(args) == 3: - dd, xyz_axis, vec = args - else: - raise TypeError("invalid number of arguments") - - return self._bound_weak_d_dx(dd, xyz_axis)(u=vec) + return op.weak_d_dx(self, *args) def weak_div(self, *args): - r"""Return the "weak divergence" of the vector volume function - represented by *vecs*. - - May be called with ``(vecs)`` or ``(dd, vecs)``. - - :arg dd: a :class:`~grudge.sym.DOFDesc`, or a value convertible to one. - Defaults to the base volume discretization if not provided. - :arg vec: a object array of - a :class:`~meshmode.dof_array.DOFArray`\ s, - where the last axis of the array must have length - matching the volume dimension. - :returns: a :class:`~meshmode.dof_array.DOFArray` - """ - if len(args) == 1: - vecs, = args - dd = sym.DOFDesc("vol", sym.QTAG_NONE) - elif len(args) == 2: - dd, vecs = args - else: - raise TypeError("invalid number of arguments") - - return self._div_helper( - lambda i, subvec: self.weak_d_dx(dd, i, subvec), - vecs) - - # }}} - - @memoize_method - def normal(self, dd): - """Get unit normal to specified surface discretization, *dd*. - - :arg dd: a :class:`~grudge.sym.DOFDesc` as the surface discretization. - :returns: an object array of :class:`~meshmode.dof_array.DOFArray`. - """ - surface_discr = self.discr_from_dd(dd) - actx = surface_discr._setup_actx - return freeze( - bind(self, - sym.normal(dd, surface_discr.ambient_dim, surface_discr.dim), - local_only=True) - (array_context=actx)) - - @memoize_method - def _bound_mass(self, dd): - return bind(self, sym.MassOperator(dd_in=dd)(sym.Variable("u", dd)), - local_only=True) + return op.weak_div(self, *args) def mass(self, *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: self.mass(dd, el), vec) - - return self._bound_mass(dd)(u=vec) - - @memoize_method - def _bound_inverse_mass(self): - return bind(self, sym.InverseMassOperator()(sym.Variable("u")), - local_only=True) + return op.mass(self, *args) def inverse_mass(self, vec): - if isinstance(vec, np.ndarray): - return obj_array_vectorize( - lambda el: self.inverse_mass(el), vec) - - return self._bound_inverse_mass()(u=vec) - - @memoize_method - def _bound_face_mass(self, dd): - u = sym.Variable("u", dd=dd) - return bind(self, sym.FaceMassOperator(dd_in=dd)(u), local_only=True) + return op.inverse_mass(self, vec) def face_mass(self, *args): - if len(args) == 1: - vec, = args - dd = sym.DOFDesc("all_faces", 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: self.face_mass(dd, el), vec) - - return self._bound_face_mass(dd)(u=vec) - - @memoize_method - def _norm(self, p, dd): - return bind(self, - sym.norm(p, sym.var("arg", dd=dd), dd=dd), - local_only=True) + return op.face_mass(self, *args) def norm(self, vec, p=2, dd=None): - if dd is None: - dd = "vol" - - dd = sym.as_dofdesc(dd) - - if isinstance(vec, np.ndarray): - if p == 2: - return sum( - self.norm(vec[idx], dd=dd)**2 - for idx in np.ndindex(vec.shape))**0.5 - elif p == np.inf: - return max( - self.norm(vec[idx], np.inf, dd=dd) - for idx in np.ndindex(vec.shape)) - else: - raise ValueError("unsupported norm order") - - return self._norm(p, dd)(arg=vec) - - @memoize_method - def _nodal_reduction(self, operator, dd): - return bind(self, operator(dd)(sym.var("arg")), local_only=True) + return op.norm(self, vec, p, dd) def nodal_sum(self, dd, vec): - return self._nodal_reduction(sym.NodalSum, dd)(arg=vec) + return op.nodal_sum(self, dd, vec) def nodal_min(self, dd, vec): - return self._nodal_reduction(sym.NodalMin, dd)(arg=vec) + return op.nodal_min(self, dd, vec) def nodal_max(self, dd, vec): - return self._nodal_reduction(sym.NodalMax, dd)(arg=vec) - - @memoize_method - def connected_ranks(self): - from meshmode.distributed import get_connected_partitions - return get_connected_partitions(self._volume_discr.mesh) - - -def interior_trace_pair(discrwb, vec): - """Return a :class:`grudge.sym.TracePair` for the interior faces of - *discrwb*. - """ - i = discrwb.project("vol", "int_faces", vec) - - def get_opposite_face(el): - if isinstance(el, Number): - return el - else: - return discrwb.opposite_face_connection()(el) - - e = obj_array_vectorize(get_opposite_face, i) - - return TracePair("int_faces", interior=i, exterior=e) - - -# {{{ distributed-memory functionality - -class _RankBoundaryCommunication: - base_tag = 1273 - - def __init__(self, discrwb, remote_rank, vol_field, tag=None): - self.tag = self.base_tag - if tag is not None: - self.tag += tag - - self.discrwb = discrwb - self.array_context = vol_field.array_context - self.remote_btag = BTAG_PARTITION(remote_rank) - - self.bdry_discr = discrwb.discr_from_dd(self.remote_btag) - self.local_dof_array = discrwb.project("vol", self.remote_btag, vol_field) - - local_data = self.array_context.to_numpy(flatten(self.local_dof_array)) - - comm = self.discrwb.mpi_communicator - - self.send_req = comm.Isend( - local_data, remote_rank, tag=self.tag) - - self.remote_data_host = np.empty_like(local_data) - self.recv_req = comm.Irecv(self.remote_data_host, remote_rank, self.tag) - - def finish(self): - self.recv_req.Wait() - - actx = self.array_context - remote_dof_array = unflatten(self.array_context, self.bdry_discr, - actx.from_numpy(self.remote_data_host)) - - bdry_conn = self.discrwb.get_distributed_boundary_swap_connection( - sym.as_dofdesc(sym.DTAG_BOUNDARY(self.remote_btag))) - swapped_remote_dof_array = bdry_conn(remote_dof_array) - - self.send_req.Wait() - - return TracePair(self.remote_btag, - interior=self.local_dof_array, - exterior=swapped_remote_dof_array) - - -def _cross_rank_trace_pairs_scalar_field(discrwb, vec, tag=None): - if isinstance(vec, Number): - return [TracePair(BTAG_PARTITION(remote_rank), interior=vec, exterior=vec) - for remote_rank in discrwb.connected_ranks()] - else: - rbcomms = [_RankBoundaryCommunication(discrwb, remote_rank, vec, tag=tag) - for remote_rank in discrwb.connected_ranks()] - return [rbcomm.finish() for rbcomm in rbcomms] - - -def cross_rank_trace_pairs(discrwb, vec, tag=None): - if isinstance(vec, np.ndarray): - - n, = vec.shape - result = {} - for ivec in range(n): - for rank_tpair in _cross_rank_trace_pairs_scalar_field( - discrwb, vec[ivec]): - assert isinstance(rank_tpair.dd.domain_tag, sym.DTAG_BOUNDARY) - assert isinstance(rank_tpair.dd.domain_tag.tag, BTAG_PARTITION) - result[rank_tpair.dd.domain_tag.tag.part_nr, ivec] = rank_tpair - - return [ - TracePair( - dd=sym.as_dofdesc(sym.DTAG_BOUNDARY(BTAG_PARTITION(remote_rank))), - interior=make_obj_array([ - result[remote_rank, i].int for i in range(n)]), - exterior=make_obj_array([ - result[remote_rank, i].ext for i in range(n)]) - ) - for remote_rank in discrwb.connected_ranks()] - else: - return _cross_rank_trace_pairs_scalar_field(discrwb, vec, tag=tag) + return op.nodal_max(self, dd, vec) -# }}} +connected_ranks = op.connected_ranks +interior_trace_pair = op.interior_trace_pair +cross_rank_trace_pairs = op.cross_rank_trace_pairs # vim: foldmethod=marker diff --git a/grudge/op.py b/grudge/op.py new file mode 100644 index 0000000000000000000000000000000000000000..c763f2bc8b65cc235168b4970095ed1d8c000151 --- /dev/null +++ b/grudge/op.py @@ -0,0 +1,506 @@ +""" +.. autofunction:: project +.. autofunction:: nodes + +.. autofunction:: grad +.. autofunction:: d_dx +.. autofunction:: div + +.. autofunction:: weak_grad +.. autofunction:: weak_d_dx +.. autofunction:: weak_div + +.. autofunction:: normal +.. autofunction:: mass +.. autofunction:: inverse_mass +.. autofunction:: face_mass + +.. autofunction:: norm +.. autofunction:: nodal_sum +.. autofunction:: nodal_min +.. autofunction:: nodal_max + +.. autofunction:: interior_trace_pair +.. autofunction:: cross_rank_trace_pairs +""" + +__copyright__ = "Copyright (C) 2021 Andreas Kloeckner" + +__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 numbers import Number +from pytools import memoize_on_first_arg + +import numpy as np # noqa +from pytools.obj_array import obj_array_vectorize, make_obj_array +import pyopencl.array as cla # noqa +from grudge import sym, bind + +from meshmode.mesh import BTAG_ALL, BTAG_NONE, BTAG_PARTITION # noqa +from meshmode.dof_array import freeze, flatten, unflatten + +from grudge.symbolic.primitives import TracePair + + +# def interp(discrwb, src, tgt, vec): +# from warnings import warn +# warn("using 'interp' is deprecated, use 'project' instead.", +# DeprecationWarning, stacklevel=2) +# +# return discrwb.project(src, tgt, vec) + + +def project(discrwb, src, tgt, vec): + """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.sym.DOFDesc`, or a value convertible to one + :arg tgt: a :class:`~grudge.sym.DOFDesc`, or a value convertible to one + :arg vec: a :class:`~meshmode.dof_array.DOFArray` + """ + src = sym.as_dofdesc(src) + tgt = sym.as_dofdesc(tgt) + if src == tgt: + return vec + + if isinstance(vec, np.ndarray): + return obj_array_vectorize( + lambda el: project(discrwb, src, tgt, el), vec) + + if isinstance(vec, Number): + return vec + + return discrwb.connection_from_dds(src, tgt)(vec) + + +# {{{ geometric properties + +def nodes(discrwb, dd=None): + r"""Return the nodes of a discretization. + + :arg dd: a :class:`~grudge.sym.DOFDesc`, or a value convertible to one. + Defaults to the base volume discretization. + :returns: an object array of :class:`~meshmode.dof_array.DOFArray`\ s + """ + if dd is None: + return discrwb._volume_discr.nodes() + else: + return discrwb.discr_from_dd(dd).nodes() + + +@memoize_on_first_arg +def normal(discrwb, dd): + """Get unit normal to specified surface discretization, *dd*. + + :arg dd: a :class:`~grudge.sym.DOFDesc` as the surface discretization. + :returns: an object array of :class:`~meshmode.dof_array.DOFArray`. + """ + surface_discr = discrwb.discr_from_dd(dd) + actx = surface_discr._setup_actx + return freeze( + bind(discrwb, + sym.normal(dd, surface_discr.ambient_dim, surface_discr.dim), + local_only=True) + (array_context=actx)) + +# }}} + + +# {{{ derivatives + +@memoize_on_first_arg +def _bound_grad(discrwb): + return bind(discrwb, sym.nabla(discrwb.dim) * sym.Variable("u"), local_only=True) + + +def grad(discrwb, vec): + r"""Return the gradient of the volume function represented by *vec*. + + :arg vec: a :class:`~meshmode.dof_array.DOFArray` + :returns: an object array of :class:`~meshmode.dof_array.DOFArray`\ s + """ + return _bound_grad(discrwb)(u=vec) + + +@memoize_on_first_arg +def _bound_d_dx(discrwb, xyz_axis): + return bind(discrwb, sym.nabla(discrwb.dim)[xyz_axis] * sym.Variable("u"), + local_only=True) + + +def d_dx(discrwb, xyz_axis, vec): + r"""Return the derivative along axis *xyz_axis* of the volume function + represented by *vec*. + + :arg xyz_axis: an integer indicating the axis along which the derivative + is taken + :arg vec: a :class:`~meshmode.dof_array.DOFArray` + :returns: a :class:`~meshmode.dof_array.DOFArray`\ s + """ + return _bound_d_dx(discrwb, xyz_axis)(u=vec) + + +def _div_helper(discrwb, diff_func, vecs): + if not isinstance(vecs, np.ndarray): + raise TypeError("argument must be an object array") + assert vecs.dtype == object + + if vecs.shape[-1] != discrwb.ambient_dim: + raise ValueError("last dimension of *vecs* argument must match " + "ambient dimension") + + if len(vecs.shape) == 1: + return sum(diff_func(i, vec_i) for i, vec_i in enumerate(vecs)) + else: + result = np.zeros(vecs.shape[:-1], dtype=object) + for idx in np.ndindex(vecs.shape[:-1]): + result[idx] = sum( + diff_func(i, vec_i) for i, vec_i in enumerate(vecs[idx])) + return result + + +def div(discrwb, vecs): + r"""Return the divergence of the vector volume function + represented by *vecs*. + + :arg vec: an object array of + a :class:`~meshmode.dof_array.DOFArray`\ s, + where the last axis of the array must have length + matching the volume dimension. + :returns: a :class:`~meshmode.dof_array.DOFArray` + """ + + return _div_helper(discrwb, + lambda i, subvec: d_dx(discrwb, i, subvec), + vecs) + + +@memoize_on_first_arg +def _bound_weak_grad(discrwb, dd): + return bind(discrwb, + sym.stiffness_t(discrwb.dim, dd_in=dd) * sym.Variable("u", dd), + local_only=True) + + +def weak_grad(discrwb, *args): + r"""Return the "weak gradient" of the volume function represented by + *vec*. + + May be called with ``(vecs)`` or ``(dd, vecs)``. + + :arg dd: a :class:`~grudge.sym.DOFDesc`, or a value convertible to one. + Defaults to the base volume discretization if not provided. + :arg vec: a :class:`~meshmode.dof_array.DOFArray` + :returns: an object array of :class:`~meshmode.dof_array.DOFArray`\ s + """ + 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") + + return _bound_weak_grad(discrwb, dd)(u=vec) + + +@memoize_on_first_arg +def _bound_weak_d_dx(discrwb, dd, xyz_axis): + return bind(discrwb, + sym.stiffness_t(discrwb.dim, dd_in=dd)[xyz_axis] + * sym.Variable("u", dd), + local_only=True) + + +def weak_d_dx(discrwb, *args): + r"""Return the derivative along axis *xyz_axis* of the volume function + represented by *vec*. + + May be called with ``(xyz_axis, vecs)`` or ``(dd, xyz_axis, vecs)``. + + :arg xyz_axis: an integer indicating the axis along which the derivative + is taken + :arg vec: a :class:`~meshmode.dof_array.DOFArray` + :returns: a :class:`~meshmode.dof_array.DOFArray`\ s + """ + if len(args) == 2: + xyz_axis, vec = args + dd = sym.DOFDesc("vol", sym.QTAG_NONE) + elif len(args) == 3: + dd, xyz_axis, vec = args + else: + raise TypeError("invalid number of arguments") + + return _bound_weak_d_dx(discrwb, dd, xyz_axis)(u=vec) + + +def weak_div(discrwb, *args): + r"""Return the "weak divergence" of the vector volume function + represented by *vecs*. + + May be called with ``(vecs)`` or ``(dd, vecs)``. + + :arg dd: a :class:`~grudge.sym.DOFDesc`, or a value convertible to one. + Defaults to the base volume discretization if not provided. + :arg vec: a object array of + a :class:`~meshmode.dof_array.DOFArray`\ s, + where the last axis of the array must have length + matching the volume dimension. + :returns: a :class:`~meshmode.dof_array.DOFArray` + """ + if len(args) == 1: + vecs, = args + dd = sym.DOFDesc("vol", sym.QTAG_NONE) + elif len(args) == 2: + dd, vecs = args + else: + raise TypeError("invalid number of arguments") + + return _div_helper(discrwb, + lambda i, subvec: weak_d_dx(discrwb, dd, i, subvec), + vecs) + +# }}} + + +# {{{ mass-like + +@memoize_on_first_arg +def _bound_mass(discrwb, dd): + return bind(discrwb, sym.MassOperator(dd_in=dd)(sym.Variable("u", dd)), + local_only=True) + + +def mass(discrwb, *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(discrwb, dd, el), vec) + + return _bound_mass(discrwb, dd)(u=vec) + + +@memoize_on_first_arg +def _bound_inverse_mass(discrwb): + return bind(discrwb, sym.InverseMassOperator()(sym.Variable("u")), + local_only=True) + + +def inverse_mass(discrwb, vec): + if isinstance(vec, np.ndarray): + return obj_array_vectorize( + lambda el: inverse_mass(discrwb, el), vec) + + return _bound_inverse_mass(discrwb)(u=vec) + + +@memoize_on_first_arg +def _bound_face_mass(discrwb, dd): + u = sym.Variable("u", dd=dd) + return bind(discrwb, sym.FaceMassOperator(dd_in=dd)(u), local_only=True) + + +def face_mass(discrwb, *args): + if len(args) == 1: + vec, = args + dd = sym.DOFDesc("all_faces", 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: face_mass(discrwb, dd, el), vec) + + return _bound_face_mass(discrwb, dd)(u=vec) + +# }}} + + +# {{{ reductions + +@memoize_on_first_arg +def _norm(discrwb, p, dd): + return bind(discrwb, + sym.norm(p, sym.var("arg", dd=dd), dd=dd), + local_only=True) + + +def norm(discrwb, vec, p, dd=None): + if dd is None: + dd = "vol" + + dd = sym.as_dofdesc(dd) + + if isinstance(vec, np.ndarray): + if p == 2: + return sum( + norm(discrwb, vec[idx], dd=dd)**2 + for idx in np.ndindex(vec.shape))**0.5 + elif p == np.inf: + return max( + norm(discrwb, vec[idx], np.inf, dd=dd) + for idx in np.ndindex(vec.shape)) + else: + raise ValueError("unsupported norm order") + + return _norm(discrwb, p, dd)(arg=vec) + + +@memoize_on_first_arg +def _nodal_reduction(discrwb, operator, dd): + return bind(discrwb, operator(dd)(sym.var("arg")), local_only=True) + + +def nodal_sum(discrwb, dd, vec): + return _nodal_reduction(discrwb, sym.NodalSum, dd)(arg=vec) + + +def nodal_min(discrwb, dd, vec): + return _nodal_reduction(discrwb, sym.NodalMin, dd)(arg=vec) + + +def nodal_max(discrwb, dd, vec): + return _nodal_reduction(discrwb, sym.NodalMax, dd)(arg=vec) + +# }}} + + +@memoize_on_first_arg +def connected_ranks(discrwb): + from meshmode.distributed import get_connected_partitions + return get_connected_partitions(discrwb._volume_discr.mesh) + + +# {{{ interior_trace_pair + +def interior_trace_pair(discrwb, vec): + """Return a :class:`grudge.sym.TracePair` for the interior faces of + *discrwb*. + """ + i = project(discrwb, "vol", "int_faces", vec) + + def get_opposite_face(el): + if isinstance(el, Number): + return el + else: + return discrwb.opposite_face_connection()(el) + + e = obj_array_vectorize(get_opposite_face, i) + + return TracePair("int_faces", interior=i, exterior=e) + +# }}} + + +# {{{ distributed-memory functionality + +class _RankBoundaryCommunication: + base_tag = 1273 + + def __init__(self, discrwb, remote_rank, vol_field, tag=None): + self.tag = self.base_tag + if tag is not None: + self.tag += tag + + self.discrwb = discrwb + self.array_context = vol_field.array_context + self.remote_btag = BTAG_PARTITION(remote_rank) + + self.bdry_discr = discrwb.discr_from_dd(self.remote_btag) + self.local_dof_array = project(discrwb, "vol", self.remote_btag, vol_field) + + local_data = self.array_context.to_numpy(flatten(self.local_dof_array)) + + comm = self.discrwb.mpi_communicator + + self.send_req = comm.Isend( + local_data, remote_rank, tag=self.tag) + + self.remote_data_host = np.empty_like(local_data) + self.recv_req = comm.Irecv(self.remote_data_host, remote_rank, self.tag) + + def finish(self): + self.recv_req.Wait() + + actx = self.array_context + remote_dof_array = unflatten(self.array_context, self.bdry_discr, + actx.from_numpy(self.remote_data_host)) + + bdry_conn = self.discrwb.get_distributed_boundary_swap_connection( + sym.as_dofdesc(sym.DTAG_BOUNDARY(self.remote_btag))) + swapped_remote_dof_array = bdry_conn(remote_dof_array) + + self.send_req.Wait() + + return TracePair(self.remote_btag, + interior=self.local_dof_array, + exterior=swapped_remote_dof_array) + + +def _cross_rank_trace_pairs_scalar_field(discrwb, vec, tag=None): + if isinstance(vec, Number): + return [TracePair(BTAG_PARTITION(remote_rank), interior=vec, exterior=vec) + for remote_rank in connected_ranks(discrwb)] + else: + rbcomms = [_RankBoundaryCommunication(discrwb, remote_rank, vec, tag=tag) + for remote_rank in connected_ranks(discrwb)] + return [rbcomm.finish() for rbcomm in rbcomms] + + +def cross_rank_trace_pairs(discrwb, vec, tag=None): + if isinstance(vec, np.ndarray): + + n, = vec.shape + result = {} + for ivec in range(n): + for rank_tpair in _cross_rank_trace_pairs_scalar_field( + discrwb, vec[ivec]): + assert isinstance(rank_tpair.dd.domain_tag, sym.DTAG_BOUNDARY) + assert isinstance(rank_tpair.dd.domain_tag.tag, BTAG_PARTITION) + result[rank_tpair.dd.domain_tag.tag.part_nr, ivec] = rank_tpair + + return [ + TracePair( + dd=sym.as_dofdesc(sym.DTAG_BOUNDARY(BTAG_PARTITION(remote_rank))), + interior=make_obj_array([ + result[remote_rank, i].int for i in range(n)]), + exterior=make_obj_array([ + result[remote_rank, i].ext for i in range(n)]) + ) + for remote_rank in connected_ranks(discrwb)] + else: + return _cross_rank_trace_pairs_scalar_field(discrwb, vec, tag=tag) + +# }}} + + +# vim: foldmethod=marker