From 3af070c3a7a2075254ff93ef49b0570477c5c8b7 Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Sun, 17 Apr 2022 12:01:05 -0500 Subject: [PATCH 01/15] Bump another tolerance in test_mass_mat_trig in sym old tests --- test/test_grudge_sym_old.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_grudge_sym_old.py b/test/test_grudge_sym_old.py index ca6c70d5c..f211a316e 100644 --- a/test/test_grudge_sym_old.py +++ b/test/test_grudge_sym_old.py @@ -148,7 +148,7 @@ def _get_variables_on(dd): num_integral_2 = np.dot(f_volm, actx.to_numpy(flatten(mass_op(f=ones_quad)))) err_2 = abs(num_integral_2 - true_integral) - assert err_2 < 2e-9, err_2 + assert err_2 < 5e-9, err_2 if discr_tag is dof_desc.DISCR_TAG_BASE: # NOTE: `integral` always makes a square mass matrix and From a6164cf22f0c2ac0646e060a699fe038e1d9626a Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Mon, 11 Apr 2022 19:30:20 -0500 Subject: [PATCH 02/15] Add a helper to enable use of compiled LSRK45 --- examples/wave/wave-op-mpi.py | 52 ++++++++++++++++++++---------------- grudge/shortcuts.py | 33 +++++++++++++++++++++-- test/test_grudge.py | 40 ++++++++++++++------------- 3 files changed, 81 insertions(+), 44 deletions(-) diff --git a/examples/wave/wave-op-mpi.py b/examples/wave/wave-op-mpi.py index ab6c41e5e..e33240443 100644 --- a/examples/wave/wave-op-mpi.py +++ b/examples/wave/wave-op-mpi.py @@ -30,7 +30,7 @@ import pyopencl.tools as cl_tools from arraycontext import ( - thaw, freeze, + thaw, with_container_arithmetic, dataclass_array_container ) @@ -45,7 +45,7 @@ from grudge.dof_desc import as_dofdesc, DOFDesc, DISCR_TAG_BASE, DISCR_TAG_QUAD from grudge.trace_pair import TracePair from grudge.discretization import DiscretizationCollection -from grudge.shortcuts import make_visualizer, rk4_step +from grudge.shortcuts import make_visualizer, compiled_lsrk45_step import grudge.op as op @@ -57,7 +57,8 @@ # {{{ wave equation bits -@with_container_arithmetic(bcast_obj_array=True, rel_comparison=True) +@with_container_arithmetic(bcast_obj_array=True, rel_comparison=True, + _cls_has_array_context_attr=True) @dataclass_array_container @dataclass(frozen=True) class WaveState: @@ -251,7 +252,8 @@ def main(ctx_factory, dim=2, order=3, c = 1 # FIXME: Sketchy, empirically determined fudge factor - dt = actx.to_numpy(0.45 * estimate_rk4_timestep(actx, dcoll, c)) + # 5/4 to account for larger LSRK45 stability region + dt = actx.to_numpy(0.45 * estimate_rk4_timestep(actx, dcoll, c)) * 5/4 vis = make_visualizer(dcoll) @@ -271,25 +273,32 @@ def rhs(t, w): istep = 0 while t < t_final: start = time.time() - if lazy: - fields = thaw(freeze(fields, actx), actx) - fields = rk4_step(fields, t, dt, compiled_rhs) - - l2norm = actx.to_numpy(op.norm(dcoll, fields.u, 2)) + fields = compiled_lsrk45_step(actx, fields, t, dt, compiled_rhs) if istep % 10 == 0: stop = time.time() - linfnorm = actx.to_numpy(op.norm(dcoll, fields.u, np.inf)) - nodalmax = actx.to_numpy(op.nodal_max(dcoll, "vol", fields.u)) - nodalmin = actx.to_numpy(op.nodal_min(dcoll, "vol", fields.u)) - if comm.rank == 0: - logger.info(f"step: {istep} t: {t} " - f"L2: {l2norm} " - f"Linf: {linfnorm} " - f"sol max: {nodalmax} " - f"sol min: {nodalmin} " - f"wall: {stop-start} ") + if args.no_diagnostics: + if comm.rank == 0: + logger.info(f"step: {istep} t: {t} " + f"wall: {stop-start} ") + else: + l2norm = actx.to_numpy(op.norm(dcoll, fields.u, 2)) + + # NOTE: These are here to ensure the solution is bounded for the + # time interval specified + assert l2norm < 1 + + linfnorm = actx.to_numpy(op.norm(dcoll, fields.u, np.inf)) + nodalmax = actx.to_numpy(op.nodal_max(dcoll, "vol", fields.u)) + nodalmin = actx.to_numpy(op.nodal_min(dcoll, "vol", fields.u)) + if comm.rank == 0: + logger.info(f"step: {istep} t: {t} " + f"L2: {l2norm} " + f"Linf: {linfnorm} " + f"sol max: {nodalmax} " + f"sol min: {nodalmin} " + f"wall: {stop-start} ") if visualize: vis.write_parallel_vtk_file( comm, @@ -304,10 +313,6 @@ def rhs(t, w): t += dt istep += 1 - # NOTE: These are here to ensure the solution is bounded for the - # time interval specified - assert l2norm < 1 - if __name__ == "__main__": import argparse @@ -320,6 +325,7 @@ def rhs(t, w): help="switch to a lazy computation mode") parser.add_argument("--quad", action="store_true") parser.add_argument("--nonaffine", action="store_true") + parser.add_argument("--no-diagnostics", action="store_true") args = parser.parse_args() diff --git a/grudge/shortcuts.py b/grudge/shortcuts.py index 52c77097d..0aca64a58 100644 --- a/grudge/shortcuts.py +++ b/grudge/shortcuts.py @@ -1,5 +1,3 @@ -"""Minimal example of a grudge driver.""" - __copyright__ = "Copyright (C) 2009 Andreas Kloeckner" __license__ = """ @@ -23,6 +21,9 @@ """ +from pytools import memoize_in + + def rk4_step(y, t, h, f): k1 = f(t, y) k2 = f(t+h/2, y + h/2*k1) @@ -31,6 +32,34 @@ def rk4_step(y, t, h, f): return y + h/6*(k1 + 2*k2 + 2*k3 + k4) +def _lsrk45_update(y, a, b, h, rhs_val, residual=0): + residual = a*residual + h*rhs_val + y = y + b * residual + from pytools.obj_array import make_obj_array + return make_obj_array([y, residual]) + + +def compiled_lsrk45_step(actx, y, t, h, f): + from leap.rk import LSRK4MethodBuilder + + @memoize_in(actx, (compiled_lsrk45_step, "update")) + def get_state_updater(): + return actx.compile(_lsrk45_update) + + update = get_state_updater() + + residual = None + + for a, b, c in LSRK4MethodBuilder.coeffs: # pylint: disable=not-an-iterable + rhs_val = f(t + c*h, y) + if residual is None: + y, residual = update(y, a, b, h, rhs_val) + else: + y, residual = update(y, a, b, h, rhs_val, residual) + + return y + + def set_up_rk4(field_var_name, dt, fields, rhs, t_start=0.0): from leap.rk import LSRK4MethodBuilder from dagrt.codegen import PythonCodeGenerator diff --git a/test/test_grudge.py b/test/test_grudge.py index 06f0710cb..7e469d468 100644 --- a/test/test_grudge.py +++ b/test/test_grudge.py @@ -724,6 +724,8 @@ def u_analytic(x, t=0): def rhs(t, u): return adv_operator.operator(t, u) + compiled_rhs = actx.compile(rhs) + if dim == 3: final_time = 0.1 else: @@ -734,35 +736,35 @@ def rhs(t, u): h_max = h_max_from_volume(dcoll, dim=dcoll.ambient_dim) dt = actx.to_numpy(dt_factor * h_max/order**2) nsteps = (final_time // dt) + 1 - dt = final_time/nsteps + 1e-15 - - from grudge.shortcuts import set_up_rk4 - dt_stepper = set_up_rk4("u", dt, u, rhs) - - last_u = None + tol = 1e-14 + dt = final_time/nsteps + tol - from grudge.shortcuts import make_visualizer + from grudge.shortcuts import make_visualizer, compiled_lsrk45_step vis = make_visualizer(dcoll) step = 0 + t = 0 - for event in dt_stepper.run(t_end=final_time): - if isinstance(event, dt_stepper.StateComputed): - step += 1 - logger.debug("[%04d] t = %.5f", step, event.t) + while t < final_time - tol: + step += 1 + logger.debug("[%04d] t = %.5f", step, t) + + u = compiled_lsrk45_step(actx, u, t, dt, compiled_rhs) + + if visualize: + vis.write_vtk_file( + "fld-%s-%04d.vtu" % (mesh_par, step), + [("u", u)] + ) - last_t = event.t - last_u = event.state_component + t += dt - if visualize: - vis.write_vtk_file( - "fld-%s-%04d.vtu" % (mesh_par, step), - [("u", event.state_component)] - ) + if t + dt >= final_time - tol: + dt = final_time-t error_l2 = op.norm( dcoll, - last_u - u_analytic(nodes, t=last_t), + u - u_analytic(nodes, t=t), 2 ) logger.info("h_max %.5e error %.5e", actx.to_numpy(h_max), error_l2) From 5f2acca2ad3657e45e6acb22ed448e3a6076d228 Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Tue, 26 Apr 2022 14:50:26 -0500 Subject: [PATCH 03/15] op: List exported symbols in __all__ --- grudge/op.py | 46 ++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 42 insertions(+), 4 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index 70bd1caf7..5ed6e5953 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -67,10 +67,10 @@ import grudge.dof_desc as dof_desc -from grudge.interpolation import interp # noqa: F401 -from grudge.projection import project # noqa: F401 +from grudge.interpolation import interp +from grudge.projection import project -from grudge.reductions import ( # noqa: F401 +from grudge.reductions import ( norm, nodal_sum, nodal_min, @@ -85,7 +85,7 @@ elementwise_integral, ) -from grudge.trace_pair import ( # noqa: F401 +from grudge.trace_pair import ( interior_trace_pair, interior_trace_pairs, connected_ranks, @@ -95,6 +95,44 @@ ) +__all__ = ( + "project", + "interp", + + "norm", + "nodal_sum", + "nodal_min", + "nodal_max", + "nodal_sum_loc", + "nodal_min_loc", + "nodal_max_loc", + "integral", + "elementwise_sum", + "elementwise_max", + "elementwise_min", + "elementwise_integral", + + "interior_trace_pair", + "interior_trace_pairs", + "connected_ranks", + "cross_rank_trace_pairs", + "bdry_trace_pair", + "bv_trace_pair", + + "local_grad", + "local_d_dx", + "local_div", + + "weak_local_grad", + "weak_local_d_dx", + "weak_local_div", + + "mass", + "inverse_mass", + "face_mass", + ) + + # {{{ common derivative "kernels" def _single_axis_derivative_kernel( From 9f0e4f7db095096fac73cc98734b35427791f259 Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Tue, 26 Apr 2022 14:57:18 -0500 Subject: [PATCH 04/15] Add project_tracepair, tracepair_with_discr_tag --- grudge/models/euler.py | 10 +--------- grudge/op.py | 13 +++++++++++++ grudge/trace_pair.py | 32 ++++++++++++++++++++++++++++++-- 3 files changed, 44 insertions(+), 11 deletions(-) diff --git a/grudge/models/euler.py b/grudge/models/euler.py index e3cc0b6b5..98acbc1ef 100644 --- a/grudge/models/euler.py +++ b/grudge/models/euler.py @@ -318,18 +318,10 @@ def operator(self, t, q): qtag = self.qtag dq = DOFDesc("vol", qtag) df = DOFDesc("all_faces", qtag) - df_int = DOFDesc("int_faces", qtag) def interp_to_quad(u): return op.project(dcoll, "vol", dq, u) - def interp_to_quad_surf(u): - return TracePair( - df_int, - interior=op.project(dcoll, "int_faces", df_int, u.int), - exterior=op.project(dcoll, "int_faces", df_int, u.ext) - ) - # Compute volume fluxes volume_fluxes = op.weak_local_div( dcoll, dq, @@ -341,7 +333,7 @@ def interp_to_quad_surf(u): sum( euler_numerical_flux( dcoll, - interp_to_quad_surf(tpair), + op.tracepair_with_discr_tag(dcoll, qtag, tpair), gamma=gamma, lf_stabilization=self.lf_stabilization ) for tpair in op.interior_trace_pairs(dcoll, q) diff --git a/grudge/op.py b/grudge/op.py index 5ed6e5953..d2e7a5446 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -22,6 +22,13 @@ .. autofunction:: mass .. autofunction:: inverse_mass .. autofunction:: face_mass + +Working around documentation tool akwardness +-------------------------------------------- + +.. class:: TracePair + + See :class:`grudge.trace_pair.TracePair`. """ __copyright__ = """ @@ -86,8 +93,11 @@ ) from grudge.trace_pair import ( + project_tracepair, + tracepair_with_discr_tag, interior_trace_pair, interior_trace_pairs, + local_interior_trace_pair, connected_ranks, cross_rank_trace_pairs, bdry_trace_pair, @@ -112,8 +122,11 @@ "elementwise_min", "elementwise_integral", + "project_tracepair", + "tracepair_with_discr_tag", "interior_trace_pair", "interior_trace_pairs", + "local_interior_trace_pair", "connected_ranks", "cross_rank_trace_pairs", "bdry_trace_pair", diff --git a/grudge/trace_pair.py b/grudge/trace_pair.py index 872028d74..f84dfdc58 100644 --- a/grudge/trace_pair.py +++ b/grudge/trace_pair.py @@ -2,11 +2,16 @@ Trace Pairs ^^^^^^^^^^^ -Container class ---------------- +Container class and auxiliary functionality +------------------------------------------- .. autoclass:: TracePair +.. currentmodule:: grudge.op + +.. autoclass:: project_tracepair +.. autoclass:: tracepair_with_discr_tag + Boundary trace functions ------------------------ @@ -530,4 +535,27 @@ def cross_rank_trace_pairs( # }}} +# {{{ project_tracepair + +def project_tracepair( + dcoll: DiscretizationCollection, new_dd: dof_desc.DOFDesc, + tpair: TracePair) -> TracePair: + r"""Return a new :class:`TracePair` :func:`~grudge.op.project`\ 'ed + onto *new_dd*. + """ + return TracePair( + new_dd, + interior=project(dcoll, tpair.dd, new_dd, tpair.int), + exterior=project(dcoll, tpair.dd, new_dd, tpair.ext) + ) + + +def tracepair_with_discr_tag(dcoll, discr_tag, tpair: TracePair) -> TracePair: + r"""Return a new :class:`TracePair` :func:`~grudge.op.project`\ 'ed + onto a :class:`~grudge.dof_desc.DOFDesc` with discretization tag *discr_tag*. + """ + return project_tracepair(dcoll, tpair.dd.with_discr_tag(discr_tag), tpair) + +# }}} + # vim: foldmethod=marker From dce90d62663c7246bdc89a73dbef7b01df86adaa Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Thu, 3 Mar 2022 17:31:16 -0600 Subject: [PATCH 05/15] generalize div/grad helpers so that logic can be reused --- grudge/op.py | 113 ++++++++++++++++++++++++++++----------------------- 1 file changed, 62 insertions(+), 51 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index d2e7a5446..255515644 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -210,66 +210,45 @@ def _gradient_kernel(actx, out_discr, in_discr, get_diff_mat, inv_jac_mat, vec, # {{{ common derivative "helpers" -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") - +def _div_helper(ambient_dim, component_div, scalar_type, vecs): 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) + partial(_div_helper, ambient_dim, component_div, scalar_type), 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.size and not isinstance(vecs[(0,)*vecs.ndim], scalar_type): + # vecs is an object array containing further object arrays + # -> treat as array container + return map_array_container( + partial(_div_helper, ambient_dim, component_div, scalar_type), vecs) - if vecs.shape[-1] != dcoll.ambient_dim: + if vecs.shape[-1] != 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)) + return component_div(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])) + result[idx] = component_div(vecs[idx]) return result -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): +def _grad_helper(ambient_dim, component_grad, scalar_type, vecs, nested): + if isinstance(vecs, 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,)) + if vecs.size == 0: + return vecs.reshape(vecs.shape + (ambient_dim,)) # For containers with ndarray data (such as momentum/velocity), # the gradient is matrix-valued, so we compute the gradient for @@ -277,17 +256,20 @@ def _grad_helper(dcoll, scalar_grad, *args, nested): # derivatives by stacking the results. grad = obj_array_vectorize( lambda el: _grad_helper( - dcoll, scalar_grad, dd_in, el, nested=nested), vec) + ambient_dim, component_grad, scalar_type, el, nested=nested), vecs) if nested: return grad else: return np.stack(grad, axis=0) - if not isinstance(vec, DOFArray): + if not isinstance(vecs, scalar_type): return map_array_container( - partial(_grad_helper, dcoll, scalar_grad, dd_in, nested=nested), vec) + partial( + _grad_helper, ambient_dim, component_grad, scalar_type, + nested=nested), + vecs) - return scalar_grad(dcoll, dd_in, vec) + return component_grad(vecs) # }}} @@ -348,8 +330,13 @@ def local_grad( :class:`~meshmode.dof_array.DOFArray`\ s or :class:`~arraycontext.container.ArrayContainer`\ of object arrays. """ - - return _grad_helper(dcoll, _strong_scalar_grad, vec, nested=nested) + dd_in = dof_desc.DOFDesc("vol", dof_desc.DISCR_TAG_BASE) + return _grad_helper( + dcoll.ambient_dim, + partial(_strong_scalar_grad, dcoll, dd_in), + DOFArray, + vec, + nested=nested) def local_d_dx( @@ -400,10 +387,12 @@ def local_div(dcoll: DiscretizationCollection, vecs) -> ArrayOrContainerT: :returns: a :class:`~meshmode.dof_array.DOFArray` or an :class:`~arraycontext.container.ArrayContainer` of them. """ - return _div_helper( - dcoll, - lambda dd, i, subvec: local_d_dx(dcoll, i, subvec), - vecs) + def component_div(vec): + return sum( + local_d_dx(dcoll, i, vec_i) + for i, vec_i in enumerate(vec)) + + return _div_helper(dcoll.ambient_dim, component_div, DOFArray, vecs) # }}} @@ -494,7 +483,20 @@ def weak_local_grad( :class:`~meshmode.dof_array.DOFArray`\ s or :class:`~arraycontext.container.ArrayContainer`\ of object arrays. """ - return _grad_helper(dcoll, _weak_scalar_grad, *args, nested=nested) + if len(args) == 1: + vecs, = args + dd_in = dof_desc.DOFDesc("vol", dof_desc.DISCR_TAG_BASE) + elif len(args) == 2: + dd_in, vecs = args + else: + raise TypeError("invalid number of arguments") + + return _grad_helper( + dcoll.ambient_dim, + partial(_weak_scalar_grad, dcoll, dd_in), + DOFArray, + vecs, + nested=nested) def weak_local_d_dx(dcoll: DiscretizationCollection, *args) -> ArrayOrContainerT: @@ -586,11 +588,20 @@ def weak_local_div(dcoll: DiscretizationCollection, *args) -> ArrayOrContainerT: :returns: a :class:`~meshmode.dof_array.DOFArray` or an :class:`~arraycontext.container.ArrayContainer` like *vec*. """ - return _div_helper( - dcoll, - lambda dd, i, subvec: weak_local_d_dx(dcoll, dd, i, subvec), - *args - ) + if len(args) == 1: + vecs, = args + dd_in = dof_desc.DOFDesc("vol", dof_desc.DISCR_TAG_BASE) + elif len(args) == 2: + dd_in, vecs = args + else: + raise TypeError("invalid number of arguments") + + def component_div(vec): + return sum( + weak_local_d_dx(dcoll, dd_in, i, vec_i) + for i, vec_i in enumerate(vec)) + + return _div_helper(dcoll.ambient_dim, component_div, DOFArray, vecs) # }}} From 765d4585c2cb7eead404f5bcfdb2e6f1e0b197a8 Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Thu, 3 Mar 2022 18:49:40 -0600 Subject: [PATCH 06/15] allow more flexibility in deciding whether something is a scalar --- grudge/op.py | 32 ++++++++++++++++++++------------ 1 file changed, 20 insertions(+), 12 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index 255515644..45595977c 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -210,19 +210,19 @@ def _gradient_kernel(actx, out_discr, in_discr, get_diff_mat, inv_jac_mat, vec, # {{{ common derivative "helpers" -def _div_helper(ambient_dim, component_div, scalar_type, vecs): +def _div_helper(ambient_dim, component_div, is_scalar, vecs): if not isinstance(vecs, np.ndarray): # vecs is not an object array -> treat as array container return map_array_container( - partial(_div_helper, ambient_dim, component_div, scalar_type), vecs) + partial(_div_helper, ambient_dim, component_div, is_scalar), vecs) assert vecs.dtype == object - if vecs.size and not isinstance(vecs[(0,)*vecs.ndim], scalar_type): + if vecs.size and not is_scalar(vecs[(0,)*vecs.ndim]): # vecs is an object array containing further object arrays # -> treat as array container return map_array_container( - partial(_div_helper, ambient_dim, component_div, scalar_type), vecs) + partial(_div_helper, ambient_dim, component_div, is_scalar), vecs) if vecs.shape[-1] != ambient_dim: raise ValueError("last/innermost dimension of *vecs* argument doesn't match " @@ -239,7 +239,7 @@ def _div_helper(ambient_dim, component_div, scalar_type, vecs): return result -def _grad_helper(ambient_dim, component_grad, scalar_type, vecs, nested): +def _grad_helper(ambient_dim, component_grad, is_scalar, vecs, nested): if isinstance(vecs, np.ndarray): # Occasionally, data structures coming from *mirgecom* will # contain empty object arrays as placeholders for fields. @@ -256,16 +256,16 @@ def _grad_helper(ambient_dim, component_grad, scalar_type, vecs, nested): # derivatives by stacking the results. grad = obj_array_vectorize( lambda el: _grad_helper( - ambient_dim, component_grad, scalar_type, el, nested=nested), vecs) + ambient_dim, component_grad, is_scalar, el, nested=nested), vecs) if nested: return grad else: return np.stack(grad, axis=0) - if not isinstance(vecs, scalar_type): + if not is_scalar(vecs): return map_array_container( partial( - _grad_helper, ambient_dim, component_grad, scalar_type, + _grad_helper, ambient_dim, component_grad, is_scalar, nested=nested), vecs) @@ -334,7 +334,7 @@ def local_grad( return _grad_helper( dcoll.ambient_dim, partial(_strong_scalar_grad, dcoll, dd_in), - DOFArray, + lambda v: isinstance(v, DOFArray), vec, nested=nested) @@ -392,7 +392,11 @@ def component_div(vec): local_d_dx(dcoll, i, vec_i) for i, vec_i in enumerate(vec)) - return _div_helper(dcoll.ambient_dim, component_div, DOFArray, vecs) + return _div_helper( + dcoll.ambient_dim, + component_div, + lambda v: isinstance(v, DOFArray), + vecs) # }}} @@ -494,7 +498,7 @@ def weak_local_grad( return _grad_helper( dcoll.ambient_dim, partial(_weak_scalar_grad, dcoll, dd_in), - DOFArray, + lambda v: isinstance(v, DOFArray), vecs, nested=nested) @@ -601,7 +605,11 @@ def component_div(vec): weak_local_d_dx(dcoll, dd_in, i, vec_i) for i, vec_i in enumerate(vec)) - return _div_helper(dcoll.ambient_dim, component_div, DOFArray, vecs) + return _div_helper( + dcoll.ambient_dim, + component_div, + lambda v: isinstance(v, DOFArray), + vecs) # }}} From dd9bfb6f3733728d4f80e2d990d99f6d64be4469 Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Fri, 4 Mar 2022 14:23:53 -0600 Subject: [PATCH 07/15] move div/grad helpers to tools --- grudge/op.py | 78 +++++-------------------------------------------- grudge/tools.py | 66 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 74 insertions(+), 70 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index 45595977c..0077f4a07 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -208,72 +208,6 @@ def _gradient_kernel(actx, out_discr, in_discr, get_diff_mat, inv_jac_mat, vec, # }}} -# {{{ common derivative "helpers" - -def _div_helper(ambient_dim, component_div, is_scalar, vecs): - if not isinstance(vecs, np.ndarray): - # vecs is not an object array -> treat as array container - return map_array_container( - partial(_div_helper, ambient_dim, component_div, is_scalar), vecs) - - assert vecs.dtype == object - - if vecs.size and not is_scalar(vecs[(0,)*vecs.ndim]): - # vecs is an object array containing further object arrays - # -> treat as array container - return map_array_container( - partial(_div_helper, ambient_dim, component_div, is_scalar), vecs) - - if vecs.shape[-1] != 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 component_div(vecs) - else: - result = np.zeros(div_result_shape, dtype=object) - for idx in np.ndindex(div_result_shape): - result[idx] = component_div(vecs[idx]) - return result - - -def _grad_helper(ambient_dim, component_grad, is_scalar, vecs, nested): - if isinstance(vecs, 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 vecs.size == 0: - return vecs.reshape(vecs.shape + (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( - ambient_dim, component_grad, is_scalar, el, nested=nested), vecs) - if nested: - return grad - else: - return np.stack(grad, axis=0) - - if not is_scalar(vecs): - return map_array_container( - partial( - _grad_helper, ambient_dim, component_grad, is_scalar, - nested=nested), - vecs) - - return component_grad(vecs) - -# }}} - - # {{{ Derivative operators def _reference_derivative_matrices(actx: ArrayContext, @@ -331,7 +265,8 @@ def local_grad( :class:`~arraycontext.container.ArrayContainer`\ of object arrays. """ dd_in = dof_desc.DOFDesc("vol", dof_desc.DISCR_TAG_BASE) - return _grad_helper( + from grudge.tools import container_grad + return container_grad( dcoll.ambient_dim, partial(_strong_scalar_grad, dcoll, dd_in), lambda v: isinstance(v, DOFArray), @@ -392,7 +327,8 @@ def component_div(vec): local_d_dx(dcoll, i, vec_i) for i, vec_i in enumerate(vec)) - return _div_helper( + from grudge.tools import container_div + return container_div( dcoll.ambient_dim, component_div, lambda v: isinstance(v, DOFArray), @@ -495,7 +431,8 @@ def weak_local_grad( else: raise TypeError("invalid number of arguments") - return _grad_helper( + from grudge.tools import container_grad + return container_grad( dcoll.ambient_dim, partial(_weak_scalar_grad, dcoll, dd_in), lambda v: isinstance(v, DOFArray), @@ -605,7 +542,8 @@ def component_div(vec): weak_local_d_dx(dcoll, dd_in, i, vec_i) for i, vec_i in enumerate(vec)) - return _div_helper( + from grudge.tools import container_div + return container_div( dcoll.ambient_dim, component_div, lambda v: isinstance(v, DOFArray), diff --git a/grudge/tools.py b/grudge/tools.py index 9723f9249..f6fe9ec49 100644 --- a/grudge/tools.py +++ b/grudge/tools.py @@ -235,4 +235,70 @@ def build_jacobian(actx, f, base_state, stepsize): return mat +# {{{ common derivative "helpers" + +def container_div(ambient_dim, component_div, is_scalar, vecs): + if not isinstance(vecs, np.ndarray): + # vecs is not an object array -> treat as array container + return map_array_container( + partial(container_div, ambient_dim, component_div, is_scalar), vecs) + + assert vecs.dtype == object + + if vecs.size and not is_scalar(vecs[(0,)*vecs.ndim]): + # vecs is an object array containing further object arrays + # -> treat as array container + return map_array_container( + partial(container_div, ambient_dim, component_div, is_scalar), vecs) + + if vecs.shape[-1] != 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 component_div(vecs) + else: + result = np.zeros(div_result_shape, dtype=object) + for idx in np.ndindex(div_result_shape): + result[idx] = component_div(vecs[idx]) + return result + + +def container_grad(ambient_dim, component_grad, is_scalar, vecs, nested): + if isinstance(vecs, 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 vecs.size == 0: + return vecs.reshape(vecs.shape + (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: container_grad( + ambient_dim, component_grad, is_scalar, el, nested=nested), vecs) + if nested: + return grad + else: + return np.stack(grad, axis=0) + + if not is_scalar(vecs): + return map_array_container( + partial( + container_grad, ambient_dim, component_grad, is_scalar, + nested=nested), + vecs) + + return component_grad(vecs) + +# }}} + + # vim: foldmethod=marker From 5b3463e9c7ba96e9742dcd7142d6c31e23da8a7a Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Mon, 7 Mar 2022 14:54:26 -0600 Subject: [PATCH 08/15] refactor div/grad helpers into one function --- grudge/op.py | 68 ++++++++++++------------ grudge/tools.py | 138 +++++++++++++++++++++++++++++------------------- 2 files changed, 118 insertions(+), 88 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index 0077f4a07..5b7c10f23 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -68,7 +68,7 @@ from grudge.discretization import DiscretizationCollection from pytools import keyed_memoize_in -from pytools.obj_array import obj_array_vectorize, make_obj_array +from pytools.obj_array import make_obj_array import numpy as np @@ -265,13 +265,14 @@ def local_grad( :class:`~arraycontext.container.ArrayContainer`\ of object arrays. """ dd_in = dof_desc.DOFDesc("vol", dof_desc.DISCR_TAG_BASE) - from grudge.tools import container_grad - return container_grad( - dcoll.ambient_dim, - partial(_strong_scalar_grad, dcoll, dd_in), - lambda v: isinstance(v, DOFArray), - vec, - nested=nested) + from grudge.tools import rec_map_subarrays + return rec_map_subarrays( + f=partial(_strong_scalar_grad, dcoll, dd_in), + in_shape=(), + out_shape=(dcoll.ambient_dim,), + is_scalar=lambda v: isinstance(v, DOFArray), + return_nested=nested, + ary=vec) def local_d_dx( @@ -322,17 +323,15 @@ def local_div(dcoll: DiscretizationCollection, vecs) -> ArrayOrContainerT: :returns: a :class:`~meshmode.dof_array.DOFArray` or an :class:`~arraycontext.container.ArrayContainer` of them. """ - def component_div(vec): - return sum( + from grudge.tools import rec_map_subarrays + return rec_map_subarrays( + f=lambda vec: sum( local_d_dx(dcoll, i, vec_i) - for i, vec_i in enumerate(vec)) - - from grudge.tools import container_div - return container_div( - dcoll.ambient_dim, - component_div, - lambda v: isinstance(v, DOFArray), - vecs) + for i, vec_i in enumerate(vec)), + in_shape=(dcoll.ambient_dim,), + out_shape=(), + is_scalar=lambda v: isinstance(v, DOFArray), + ary=vecs) # }}} @@ -431,13 +430,14 @@ def weak_local_grad( else: raise TypeError("invalid number of arguments") - from grudge.tools import container_grad - return container_grad( - dcoll.ambient_dim, - partial(_weak_scalar_grad, dcoll, dd_in), - lambda v: isinstance(v, DOFArray), - vecs, - nested=nested) + from grudge.tools import rec_map_subarrays + return rec_map_subarrays( + f=partial(_weak_scalar_grad, dcoll, dd_in), + in_shape=(), + out_shape=(dcoll.ambient_dim,), + is_scalar=lambda v: isinstance(v, DOFArray), + return_nested=nested, + ary=vecs) def weak_local_d_dx(dcoll: DiscretizationCollection, *args) -> ArrayOrContainerT: @@ -537,17 +537,15 @@ def weak_local_div(dcoll: DiscretizationCollection, *args) -> ArrayOrContainerT: else: raise TypeError("invalid number of arguments") - def component_div(vec): - return sum( + from grudge.tools import rec_map_subarrays + return rec_map_subarrays( + f=lambda vec: sum( weak_local_d_dx(dcoll, dd_in, i, vec_i) - for i, vec_i in enumerate(vec)) - - from grudge.tools import container_div - return container_div( - dcoll.ambient_dim, - component_div, - lambda v: isinstance(v, DOFArray), - vecs) + for i, vec_i in enumerate(vec)), + in_shape=(dcoll.ambient_dim,), + out_shape=(), + is_scalar=lambda v: isinstance(v, DOFArray), + ary=vecs) # }}} diff --git a/grudge/tools.py b/grudge/tools.py index f6fe9ec49..72aaae8dd 100644 --- a/grudge/tools.py +++ b/grudge/tools.py @@ -1,5 +1,7 @@ """ .. autofunction:: build_jacobian +.. autofunction:: map_subarrays +.. autofunction:: rec_map_subarrays """ __copyright__ = "Copyright (C) 2007 Andreas Kloeckner" @@ -26,6 +28,7 @@ import numpy as np from pytools import levi_civita +from functools import partial def is_zero(x): @@ -235,70 +238,99 @@ def build_jacobian(actx, f, base_state, stepsize): return mat -# {{{ common derivative "helpers" +def map_subarrays(f, in_shape, out_shape, ary, *, return_nested=False): + """ + Map a function *f* over a :class:`numpy.ndarray`, applying it to subarrays + of shape *in_shape* individually. -def container_div(ambient_dim, component_div, is_scalar, vecs): - if not isinstance(vecs, np.ndarray): - # vecs is not an object array -> treat as array container - return map_array_container( - partial(container_div, ambient_dim, component_div, is_scalar), vecs) + Example 1: given a function *f* that maps arrays of shape ``(2, 2)`` to scalars + and an input array *ary* of shape ``(3, 2, 2)``, the call:: - assert vecs.dtype == object + map_subarrays(f, (2, 2), (), ary) - if vecs.size and not is_scalar(vecs[(0,)*vecs.ndim]): - # vecs is an object array containing further object arrays - # -> treat as array container - return map_array_container( - partial(container_div, ambient_dim, component_div, is_scalar), vecs) + will produce an array of shape ``(3,)`` containing the results of calling *f* on + the 3 subarrays of shape ``(2, 2)`` in *ary*. + + Example 2: given a function *f* that maps arrays of shape ``(2,)`` to arrays of + shape ``(2, 2)`` and an input array *ary* of shape ``(3, 2)``, the call:: - if vecs.shape[-1] != ambient_dim: - raise ValueError("last/innermost dimension of *vecs* argument doesn't match " - "ambient dimension") + map_subarrays(f, (2,), (2, 2), ary) - div_result_shape = vecs.shape[:-1] + will produce an array of shape ``(3, 2, 2)`` containing the results of calling + *f* on the 3 subarrays of shape ``(2,)`` in *ary*. The call:: - if len(div_result_shape) == 0: - return component_div(vecs) + map_subarrays(f, (2,), (2, 2), ary, return_nested=True) + + will instead produce an array of shape ``(3,)`` with each entry containing an + array of shape ``(2, 2)``. + + :param f: the function to be called. + :param in_shape: the shape of any inputs to *f*. + :param out_shape: the shape of the result of calling *f* on an array of shape + *in_shape*. + :param ary: a :class:`numpy.ndarray` instance. + :param return_nested: if *out_shape* is nontrivial, this flag indicates whether + to return a nested array (containing one entry for each application of *f*), + or to return a single, higher-dimensional array. + + :returns: an array with the subarrays of shape *in_shape* replaced with + subarrays of shape *out_shape* resulting from the application of *f*. + """ + if not isinstance(ary, np.ndarray): + if len(in_shape) != 0: + raise ValueError(f"found scalar, expected array of shape {in_shape}") + return f(ary) else: - result = np.zeros(div_result_shape, dtype=object) - for idx in np.ndindex(div_result_shape): - result[idx] = component_div(vecs[idx]) - return result - - -def container_grad(ambient_dim, component_grad, is_scalar, vecs, nested): - if isinstance(vecs, 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 vecs.size == 0: - return vecs.reshape(vecs.shape + (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: container_grad( - ambient_dim, component_grad, is_scalar, el, nested=nested), vecs) - if nested: - return grad + if ( + ary.ndim < len(in_shape) + or ary.shape[ary.ndim-len(in_shape):] != in_shape): + raise ValueError( + f"array of shape {ary.shape} is incompatible with function " + f"expecting input shape {in_shape}") + base_shape = ary.shape[:ary.ndim-len(in_shape)] + if len(base_shape) == 0: + return f(ary) else: - return np.stack(grad, axis=0) + in_slice = tuple(slice(0, n) for n in in_shape) + out_slice = tuple(slice(0, n) for n in out_shape) + if return_nested: + result = np.empty(base_shape, dtype=object) + for idx in np.ndindex(base_shape): + result[idx] = f(ary[idx + in_slice]) + return result + else: + result = np.empty(base_shape + out_shape, dtype=object) + for idx in np.ndindex(base_shape): + result[idx + out_slice] = f(ary[idx + in_slice]) + return result + + +def rec_map_subarrays( + f, in_shape, out_shape, is_scalar, ary, *, return_nested=False): + """ + Map a function *f* over an object *ary*, applying it to any subarrays + of shape *in_shape* individually. + + Array container version of :func:`map_subarrays`. - if not is_scalar(vecs): + :param is_scalar: a function that indicates whether a given object is to be + treated as a scalar or not. + """ + def is_array_of_scalars(a): + return ( + isinstance(a, np.ndarray) and a.dtype == object + and all(is_scalar(a[idx]) for idx in np.ndindex(a.shape))) + + if is_scalar(ary) or is_array_of_scalars(ary): + return map_subarrays( + f, in_shape, out_shape, ary, return_nested=return_nested) + else: + from arraycontext import map_array_container return map_array_container( partial( - container_grad, ambient_dim, component_grad, is_scalar, - nested=nested), - vecs) - - return component_grad(vecs) - -# }}} + rec_map_subarrays, f, in_shape, out_shape, is_scalar, + return_nested=return_nested), + ary) # vim: foldmethod=marker From 60d4fa0d8334ba439b21e1cd83c0fcb86ae94890 Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Sun, 27 Mar 2022 21:37:54 -0500 Subject: [PATCH 09/15] Rework *map_subarrays docs, add types, use leaf_cls --- grudge/op.py | 14 ++++--------- grudge/tools.py | 53 ++++++++++++++++++++++++++++--------------------- 2 files changed, 34 insertions(+), 33 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index 5b7c10f23..bdb77cde4 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -270,9 +270,7 @@ def local_grad( f=partial(_strong_scalar_grad, dcoll, dd_in), in_shape=(), out_shape=(dcoll.ambient_dim,), - is_scalar=lambda v: isinstance(v, DOFArray), - return_nested=nested, - ary=vec) + ary=vec, leaf_cls=DOFArray, return_nested=nested,) def local_d_dx( @@ -330,8 +328,7 @@ def local_div(dcoll: DiscretizationCollection, vecs) -> ArrayOrContainerT: for i, vec_i in enumerate(vec)), in_shape=(dcoll.ambient_dim,), out_shape=(), - is_scalar=lambda v: isinstance(v, DOFArray), - ary=vecs) + ary=vecs, leaf_cls=DOFArray) # }}} @@ -435,9 +432,7 @@ def weak_local_grad( f=partial(_weak_scalar_grad, dcoll, dd_in), in_shape=(), out_shape=(dcoll.ambient_dim,), - is_scalar=lambda v: isinstance(v, DOFArray), - return_nested=nested, - ary=vecs) + ary=vecs, leaf_cls=DOFArray, return_nested=nested) def weak_local_d_dx(dcoll: DiscretizationCollection, *args) -> ArrayOrContainerT: @@ -544,8 +539,7 @@ def weak_local_div(dcoll: DiscretizationCollection, *args) -> ArrayOrContainerT: for i, vec_i in enumerate(vec)), in_shape=(dcoll.ambient_dim,), out_shape=(), - is_scalar=lambda v: isinstance(v, DOFArray), - ary=vecs) + ary=vecs, leaf_cls=DOFArray) # }}} diff --git a/grudge/tools.py b/grudge/tools.py index 72aaae8dd..ab7f5e02d 100644 --- a/grudge/tools.py +++ b/grudge/tools.py @@ -28,6 +28,7 @@ import numpy as np from pytools import levi_civita +from typing import Tuple, Callable, Optional from functools import partial @@ -238,12 +239,20 @@ def build_jacobian(actx, f, base_state, stepsize): return mat -def map_subarrays(f, in_shape, out_shape, ary, *, return_nested=False): +def map_subarrays( + f: Callable[[np.ndarray], np.ndarray], + in_shape: Tuple[int, ...], out_shape: Tuple[int, ...], + ary: np.ndarray, *, return_nested=False) -> np.ndarray: """ - Map a function *f* over a :class:`numpy.ndarray`, applying it to subarrays - of shape *in_shape* individually. + Apply a function *f* to subarrrays of shape *in_shape* of an + :class:`numpy.ndarray`, typically (but not necessarily) of dtype + :class:`object`. Return an :class:`numpy.ndarray` of the same dtype, + with the corresponding subarrays replaced by the return values of *f*, + and with the shape adapted to reflect *out_shape*. - Example 1: given a function *f* that maps arrays of shape ``(2, 2)`` to scalars + Similar to :class:`numpy.vectorize`. + + *Example 1:* given a function *f* that maps arrays of shape ``(2, 2)`` to scalars and an input array *ary* of shape ``(3, 2, 2)``, the call:: map_subarrays(f, (2, 2), (), ary) @@ -251,7 +260,7 @@ def map_subarrays(f, in_shape, out_shape, ary, *, return_nested=False): will produce an array of shape ``(3,)`` containing the results of calling *f* on the 3 subarrays of shape ``(2, 2)`` in *ary*. - Example 2: given a function *f* that maps arrays of shape ``(2,)`` to arrays of + *Example 2:* given a function *f* that maps arrays of shape ``(2,)`` to arrays of shape ``(2, 2)`` and an input array *ary* of shape ``(3, 2)``, the call:: map_subarrays(f, (2,), (2, 2), ary) @@ -294,41 +303,39 @@ def map_subarrays(f, in_shape, out_shape, ary, *, return_nested=False): in_slice = tuple(slice(0, n) for n in in_shape) out_slice = tuple(slice(0, n) for n in out_shape) if return_nested: - result = np.empty(base_shape, dtype=object) + result = np.empty(base_shape, dtype=ary.dtype) for idx in np.ndindex(base_shape): result[idx] = f(ary[idx + in_slice]) return result else: - result = np.empty(base_shape + out_shape, dtype=object) + result = np.empty(base_shape + out_shape, dtype=ary.dtype) for idx in np.ndindex(base_shape): result[idx + out_slice] = f(ary[idx + in_slice]) return result def rec_map_subarrays( - f, in_shape, out_shape, is_scalar, ary, *, return_nested=False): + f: Callable[[np.ndarray], np.ndarray], + in_shape: Tuple[int, ...], + out_shape: Tuple[int, ...], + ary: np.ndarray, *, + leaf_cls: Optional[type] = None, + return_nested: bool = False) -> np.ndarray: + r""" + Like :func:`map_subarrays`, but with support for + :class:`arraycontext.ArrayContainer`\ s. + + :param leaf_cls: An array container of this type will be considered a leaf + and will be passed to *f* without further destructuring. """ - Map a function *f* over an object *ary*, applying it to any subarrays - of shape *in_shape* individually. - - Array container version of :func:`map_subarrays`. - - :param is_scalar: a function that indicates whether a given object is to be - treated as a scalar or not. - """ - def is_array_of_scalars(a): - return ( - isinstance(a, np.ndarray) and a.dtype == object - and all(is_scalar(a[idx]) for idx in np.ndindex(a.shape))) - - if is_scalar(ary) or is_array_of_scalars(ary): + if type(ary) is leaf_cls or isinstance(ary, np.ndarray): return map_subarrays( f, in_shape, out_shape, ary, return_nested=return_nested) else: from arraycontext import map_array_container return map_array_container( partial( - rec_map_subarrays, f, in_shape, out_shape, is_scalar, + rec_map_subarrays, f, in_shape, out_shape, leaf_cls=leaf_cls, return_nested=return_nested), ary) From 2334aa00036c50b942a6bd89381d427e9ad98786 Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Wed, 27 Apr 2022 10:18:21 -0500 Subject: [PATCH 10/15] Revert "Rework *map_subarrays docs, add types, use leaf_cls" This reverts commit 47d2d6b0dc76bdc01be7168cc79e13da72540820. --- grudge/op.py | 14 +++++++++---- grudge/tools.py | 53 +++++++++++++++++++++---------------------------- 2 files changed, 33 insertions(+), 34 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index bdb77cde4..5b7c10f23 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -270,7 +270,9 @@ def local_grad( f=partial(_strong_scalar_grad, dcoll, dd_in), in_shape=(), out_shape=(dcoll.ambient_dim,), - ary=vec, leaf_cls=DOFArray, return_nested=nested,) + is_scalar=lambda v: isinstance(v, DOFArray), + return_nested=nested, + ary=vec) def local_d_dx( @@ -328,7 +330,8 @@ def local_div(dcoll: DiscretizationCollection, vecs) -> ArrayOrContainerT: for i, vec_i in enumerate(vec)), in_shape=(dcoll.ambient_dim,), out_shape=(), - ary=vecs, leaf_cls=DOFArray) + is_scalar=lambda v: isinstance(v, DOFArray), + ary=vecs) # }}} @@ -432,7 +435,9 @@ def weak_local_grad( f=partial(_weak_scalar_grad, dcoll, dd_in), in_shape=(), out_shape=(dcoll.ambient_dim,), - ary=vecs, leaf_cls=DOFArray, return_nested=nested) + is_scalar=lambda v: isinstance(v, DOFArray), + return_nested=nested, + ary=vecs) def weak_local_d_dx(dcoll: DiscretizationCollection, *args) -> ArrayOrContainerT: @@ -539,7 +544,8 @@ def weak_local_div(dcoll: DiscretizationCollection, *args) -> ArrayOrContainerT: for i, vec_i in enumerate(vec)), in_shape=(dcoll.ambient_dim,), out_shape=(), - ary=vecs, leaf_cls=DOFArray) + is_scalar=lambda v: isinstance(v, DOFArray), + ary=vecs) # }}} diff --git a/grudge/tools.py b/grudge/tools.py index ab7f5e02d..72aaae8dd 100644 --- a/grudge/tools.py +++ b/grudge/tools.py @@ -28,7 +28,6 @@ import numpy as np from pytools import levi_civita -from typing import Tuple, Callable, Optional from functools import partial @@ -239,20 +238,12 @@ def build_jacobian(actx, f, base_state, stepsize): return mat -def map_subarrays( - f: Callable[[np.ndarray], np.ndarray], - in_shape: Tuple[int, ...], out_shape: Tuple[int, ...], - ary: np.ndarray, *, return_nested=False) -> np.ndarray: +def map_subarrays(f, in_shape, out_shape, ary, *, return_nested=False): """ - Apply a function *f* to subarrrays of shape *in_shape* of an - :class:`numpy.ndarray`, typically (but not necessarily) of dtype - :class:`object`. Return an :class:`numpy.ndarray` of the same dtype, - with the corresponding subarrays replaced by the return values of *f*, - and with the shape adapted to reflect *out_shape*. + Map a function *f* over a :class:`numpy.ndarray`, applying it to subarrays + of shape *in_shape* individually. - Similar to :class:`numpy.vectorize`. - - *Example 1:* given a function *f* that maps arrays of shape ``(2, 2)`` to scalars + Example 1: given a function *f* that maps arrays of shape ``(2, 2)`` to scalars and an input array *ary* of shape ``(3, 2, 2)``, the call:: map_subarrays(f, (2, 2), (), ary) @@ -260,7 +251,7 @@ def map_subarrays( will produce an array of shape ``(3,)`` containing the results of calling *f* on the 3 subarrays of shape ``(2, 2)`` in *ary*. - *Example 2:* given a function *f* that maps arrays of shape ``(2,)`` to arrays of + Example 2: given a function *f* that maps arrays of shape ``(2,)`` to arrays of shape ``(2, 2)`` and an input array *ary* of shape ``(3, 2)``, the call:: map_subarrays(f, (2,), (2, 2), ary) @@ -303,39 +294,41 @@ def map_subarrays( in_slice = tuple(slice(0, n) for n in in_shape) out_slice = tuple(slice(0, n) for n in out_shape) if return_nested: - result = np.empty(base_shape, dtype=ary.dtype) + result = np.empty(base_shape, dtype=object) for idx in np.ndindex(base_shape): result[idx] = f(ary[idx + in_slice]) return result else: - result = np.empty(base_shape + out_shape, dtype=ary.dtype) + result = np.empty(base_shape + out_shape, dtype=object) for idx in np.ndindex(base_shape): result[idx + out_slice] = f(ary[idx + in_slice]) return result def rec_map_subarrays( - f: Callable[[np.ndarray], np.ndarray], - in_shape: Tuple[int, ...], - out_shape: Tuple[int, ...], - ary: np.ndarray, *, - leaf_cls: Optional[type] = None, - return_nested: bool = False) -> np.ndarray: - r""" - Like :func:`map_subarrays`, but with support for - :class:`arraycontext.ArrayContainer`\ s. - - :param leaf_cls: An array container of this type will be considered a leaf - and will be passed to *f* without further destructuring. + f, in_shape, out_shape, is_scalar, ary, *, return_nested=False): """ - if type(ary) is leaf_cls or isinstance(ary, np.ndarray): + Map a function *f* over an object *ary*, applying it to any subarrays + of shape *in_shape* individually. + + Array container version of :func:`map_subarrays`. + + :param is_scalar: a function that indicates whether a given object is to be + treated as a scalar or not. + """ + def is_array_of_scalars(a): + return ( + isinstance(a, np.ndarray) and a.dtype == object + and all(is_scalar(a[idx]) for idx in np.ndindex(a.shape))) + + if is_scalar(ary) or is_array_of_scalars(ary): return map_subarrays( f, in_shape, out_shape, ary, return_nested=return_nested) else: from arraycontext import map_array_container return map_array_container( partial( - rec_map_subarrays, f, in_shape, out_shape, leaf_cls=leaf_cls, + rec_map_subarrays, f, in_shape, out_shape, is_scalar, return_nested=return_nested), ary) From 82b200ebabfee56473e2dbf97c58905668b89429 Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Wed, 27 Apr 2022 12:58:25 -0500 Subject: [PATCH 11/15] partially unrevert changes --- grudge/op.py | 17 +++++++---------- grudge/tools.py | 39 ++++++++++++++++++++++++++------------- 2 files changed, 33 insertions(+), 23 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index 5b7c10f23..18ccd2661 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -270,9 +270,8 @@ def local_grad( f=partial(_strong_scalar_grad, dcoll, dd_in), in_shape=(), out_shape=(dcoll.ambient_dim,), - is_scalar=lambda v: isinstance(v, DOFArray), - return_nested=nested, - ary=vec) + ary=vec, is_scalar=lambda v: isinstance(v, DOFArray), + return_nested=nested,) def local_d_dx( @@ -330,8 +329,8 @@ def local_div(dcoll: DiscretizationCollection, vecs) -> ArrayOrContainerT: for i, vec_i in enumerate(vec)), in_shape=(dcoll.ambient_dim,), out_shape=(), - is_scalar=lambda v: isinstance(v, DOFArray), - ary=vecs) + ary=vecs, + is_scalar=lambda v: isinstance(v, DOFArray)) # }}} @@ -435,9 +434,8 @@ def weak_local_grad( f=partial(_weak_scalar_grad, dcoll, dd_in), in_shape=(), out_shape=(dcoll.ambient_dim,), - is_scalar=lambda v: isinstance(v, DOFArray), - return_nested=nested, - ary=vecs) + ary=vecs, is_scalar=lambda v: isinstance(v, DOFArray), + return_nested=nested) def weak_local_d_dx(dcoll: DiscretizationCollection, *args) -> ArrayOrContainerT: @@ -544,8 +542,7 @@ def weak_local_div(dcoll: DiscretizationCollection, *args) -> ArrayOrContainerT: for i, vec_i in enumerate(vec)), in_shape=(dcoll.ambient_dim,), out_shape=(), - is_scalar=lambda v: isinstance(v, DOFArray), - ary=vecs) + ary=vecs, is_scalar=lambda v: isinstance(v, DOFArray)) # }}} diff --git a/grudge/tools.py b/grudge/tools.py index 72aaae8dd..74889b880 100644 --- a/grudge/tools.py +++ b/grudge/tools.py @@ -28,7 +28,9 @@ import numpy as np from pytools import levi_civita +from typing import Tuple, Callable, Optional, Any from functools import partial +from arraycontext.container import ArrayOrContainerT def is_zero(x): @@ -238,12 +240,20 @@ def build_jacobian(actx, f, base_state, stepsize): return mat -def map_subarrays(f, in_shape, out_shape, ary, *, return_nested=False): +def map_subarrays( + f: Callable[[Any], Any], + in_shape: Tuple[int, ...], out_shape: Tuple[int, ...], + ary: Any, *, return_nested=False) -> Any: """ - Map a function *f* over a :class:`numpy.ndarray`, applying it to subarrays - of shape *in_shape* individually. + Apply a function *f* to subarrays of shape *in_shape* of an + :class:`numpy.ndarray`, typically (but not necessarily) of dtype + :class:`object`. Return an :class:`numpy.ndarray` of the same dtype, + with the corresponding subarrays replaced by the return values of *f*, + and with the shape adapted to reflect *out_shape*. - Example 1: given a function *f* that maps arrays of shape ``(2, 2)`` to scalars + Similar to :class:`numpy.vectorize`. + + *Example 1:* given a function *f* that maps arrays of shape ``(2, 2)`` to scalars and an input array *ary* of shape ``(3, 2, 2)``, the call:: map_subarrays(f, (2, 2), (), ary) @@ -251,7 +261,7 @@ def map_subarrays(f, in_shape, out_shape, ary, *, return_nested=False): will produce an array of shape ``(3,)`` containing the results of calling *f* on the 3 subarrays of shape ``(2, 2)`` in *ary*. - Example 2: given a function *f* that maps arrays of shape ``(2,)`` to arrays of + *Example 2:* given a function *f* that maps arrays of shape ``(2,)`` to arrays of shape ``(2, 2)`` and an input array *ary* of shape ``(3, 2)``, the call:: map_subarrays(f, (2,), (2, 2), ary) @@ -299,19 +309,22 @@ def map_subarrays(f, in_shape, out_shape, ary, *, return_nested=False): result[idx] = f(ary[idx + in_slice]) return result else: - result = np.empty(base_shape + out_shape, dtype=object) + result = np.empty(base_shape + out_shape, dtype=ary.dtype) for idx in np.ndindex(base_shape): result[idx + out_slice] = f(ary[idx + in_slice]) return result def rec_map_subarrays( - f, in_shape, out_shape, is_scalar, ary, *, return_nested=False): - """ - Map a function *f* over an object *ary*, applying it to any subarrays - of shape *in_shape* individually. - - Array container version of :func:`map_subarrays`. + f: Callable[[Any], Any], + in_shape: Tuple[int, ...], + out_shape: Tuple[int, ...], + ary: ArrayOrContainerT, *, + is_scalar: Optional[Callable[[Any], bool]] = None, + return_nested: bool = False) -> ArrayOrContainerT: + r""" + Like :func:`map_subarrays`, but with support for + :class:`arraycontext.ArrayContainer`\ s. :param is_scalar: a function that indicates whether a given object is to be treated as a scalar or not. @@ -328,7 +341,7 @@ def is_array_of_scalars(a): from arraycontext import map_array_container return map_array_container( partial( - rec_map_subarrays, f, in_shape, out_shape, is_scalar, + rec_map_subarrays, f, in_shape, out_shape, is_scalar=is_scalar, return_nested=return_nested), ary) From 683d11d1eb02e33ee6963caa7d10222ea7e22a77 Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Wed, 27 Apr 2022 13:15:13 -0500 Subject: [PATCH 12/15] assume numpy scalar types are scalars and use scalar_cls instead of is_scalar --- grudge/op.py | 11 ++++------- grudge/tools.py | 21 +++++++++++++++------ 2 files changed, 19 insertions(+), 13 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index 18ccd2661..1607cba61 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -270,8 +270,7 @@ def local_grad( f=partial(_strong_scalar_grad, dcoll, dd_in), in_shape=(), out_shape=(dcoll.ambient_dim,), - ary=vec, is_scalar=lambda v: isinstance(v, DOFArray), - return_nested=nested,) + ary=vec, scalar_cls=DOFArray, return_nested=nested,) def local_d_dx( @@ -329,8 +328,7 @@ def local_div(dcoll: DiscretizationCollection, vecs) -> ArrayOrContainerT: for i, vec_i in enumerate(vec)), in_shape=(dcoll.ambient_dim,), out_shape=(), - ary=vecs, - is_scalar=lambda v: isinstance(v, DOFArray)) + ary=vecs, scalar_cls=DOFArray) # }}} @@ -434,8 +432,7 @@ def weak_local_grad( f=partial(_weak_scalar_grad, dcoll, dd_in), in_shape=(), out_shape=(dcoll.ambient_dim,), - ary=vecs, is_scalar=lambda v: isinstance(v, DOFArray), - return_nested=nested) + ary=vecs, scalar_cls=DOFArray, return_nested=nested) def weak_local_d_dx(dcoll: DiscretizationCollection, *args) -> ArrayOrContainerT: @@ -542,7 +539,7 @@ def weak_local_div(dcoll: DiscretizationCollection, *args) -> ArrayOrContainerT: for i, vec_i in enumerate(vec)), in_shape=(dcoll.ambient_dim,), out_shape=(), - ary=vecs, is_scalar=lambda v: isinstance(v, DOFArray)) + ary=vecs, scalar_cls=DOFArray) # }}} diff --git a/grudge/tools.py b/grudge/tools.py index 74889b880..246c21d81 100644 --- a/grudge/tools.py +++ b/grudge/tools.py @@ -320,19 +320,28 @@ def rec_map_subarrays( in_shape: Tuple[int, ...], out_shape: Tuple[int, ...], ary: ArrayOrContainerT, *, - is_scalar: Optional[Callable[[Any], bool]] = None, + scalar_cls: Optional[Union[type, Tuple[type]]] = None, return_nested: bool = False) -> ArrayOrContainerT: r""" Like :func:`map_subarrays`, but with support for :class:`arraycontext.ArrayContainer`\ s. - :param is_scalar: a function that indicates whether a given object is to be - treated as a scalar or not. + :param scalar_cls: An array container of this type will be considered a scalar + and arrays of it will be passed to *f* without further destructuring. """ + if scalar_cls is not None: + def is_scalar(x): + return np.isscalar(x) or isinstance(x, scalar_cls) + else: + def is_scalar(x): + return np.isscalar(x) + def is_array_of_scalars(a): return ( - isinstance(a, np.ndarray) and a.dtype == object - and all(is_scalar(a[idx]) for idx in np.ndindex(a.shape))) + isinstance(a, np.ndarray) + and ( + a.dtype != object + or all(is_scalar(a[idx]) for idx in np.ndindex(a.shape)))) if is_scalar(ary) or is_array_of_scalars(ary): return map_subarrays( @@ -341,7 +350,7 @@ def is_array_of_scalars(a): from arraycontext import map_array_container return map_array_container( partial( - rec_map_subarrays, f, in_shape, out_shape, is_scalar=is_scalar, + rec_map_subarrays, f, in_shape, out_shape, scalar_cls=scalar_cls, return_nested=return_nested), ary) From cc03864eeb3f644d0dc7d2b52c8ab0d1e892f95f Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Wed, 27 Apr 2022 14:09:43 -0500 Subject: [PATCH 13/15] deduce result dtype --- grudge/tools.py | 44 +++++++++++++++++++++++++++++--------------- 1 file changed, 29 insertions(+), 15 deletions(-) diff --git a/grudge/tools.py b/grudge/tools.py index 246c21d81..3427ba1cb 100644 --- a/grudge/tools.py +++ b/grudge/tools.py @@ -27,8 +27,8 @@ """ import numpy as np -from pytools import levi_civita -from typing import Tuple, Callable, Optional, Any +from pytools import levi_civita, product +from typing import Tuple, Callable, Optional, Union, Any from functools import partial from arraycontext.container import ArrayOrContainerT @@ -247,9 +247,9 @@ def map_subarrays( """ Apply a function *f* to subarrays of shape *in_shape* of an :class:`numpy.ndarray`, typically (but not necessarily) of dtype - :class:`object`. Return an :class:`numpy.ndarray` of the same dtype, - with the corresponding subarrays replaced by the return values of *f*, - and with the shape adapted to reflect *out_shape*. + :class:`object`. Return an :class:`numpy.ndarray` with the corresponding + subarrays replaced by the return values of *f*, and with the shape adapted + to reflect *out_shape*. Similar to :class:`numpy.vectorize`. @@ -300,19 +300,33 @@ def map_subarrays( base_shape = ary.shape[:ary.ndim-len(in_shape)] if len(base_shape) == 0: return f(ary) + elif product(base_shape) == 0: + if return_nested: + return np.empty(base_shape, dtype=object) + else: + return np.empty(base_shape + out_shape, dtype=object) else: in_slice = tuple(slice(0, n) for n in in_shape) - out_slice = tuple(slice(0, n) for n in out_shape) - if return_nested: - result = np.empty(base_shape, dtype=object) - for idx in np.ndindex(base_shape): - result[idx] = f(ary[idx + in_slice]) - return result + result_entries = np.empty(base_shape, dtype=object) + for idx in np.ndindex(base_shape): + result_entries[idx] = f(ary[idx + in_slice]) + if len(out_shape) == 0: + out_entry_template = result_entries.flat[0] + if np.isscalar(out_entry_template): + return result_entries.astype(type(out_entry_template)) + else: + return result_entries else: - result = np.empty(base_shape + out_shape, dtype=ary.dtype) - for idx in np.ndindex(base_shape): - result[idx + out_slice] = f(ary[idx + in_slice]) - return result + if return_nested: + return result_entries + else: + out_slice = tuple(slice(0, n) for n in out_shape) + out_entry_template = result_entries.flat[0] + result = np.empty( + base_shape + out_shape, dtype=out_entry_template.dtype) + for idx in np.ndindex(base_shape): + result[idx + out_slice] = result_entries[idx] + return result def rec_map_subarrays( From 5b9c3a345eec5fff6b558b1581fb857904bd30a1 Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Wed, 27 Apr 2022 15:59:37 -0500 Subject: [PATCH 14/15] add tests for map_subarrays and rec_map_subarrays --- test/test_tools.py | 202 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 202 insertions(+) create mode 100644 test/test_tools.py diff --git a/test/test_tools.py b/test/test_tools.py new file mode 100644 index 000000000..fb1785faf --- /dev/null +++ b/test/test_tools.py @@ -0,0 +1,202 @@ +__copyright__ = """ +Copyright (C) 2022 University of Illinois Board of Trustees +""" + +__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. +""" + +import numpy as np +import numpy.linalg as la # noqa + +from dataclasses import dataclass + +from grudge.array_context import PytestPyOpenCLArrayContextFactory +from arraycontext import pytest_generate_tests_for_array_contexts +pytest_generate_tests = pytest_generate_tests_for_array_contexts( + [PytestPyOpenCLArrayContextFactory]) + +from pytools.obj_array import make_obj_array + +import pytest + +import logging + +logger = logging.getLogger(__name__) + + +# {{{ map_subarrays and rec_map_subarrays + +@dataclass(frozen=True, eq=True) +class _DummyScalar: + val: int + + +def test_map_subarrays(): + """Test map_subarrays.""" + from grudge.tools import map_subarrays + + # Scalar + result = map_subarrays( + lambda x: np.array([x, 2*x]), (), (2,), 1) + assert result.dtype == int + assert np.all(result == np.array([1, 2])) + + # Scalar, nested + result = map_subarrays( + lambda x: np.array([x, 2*x]), (), (2,), 1, return_nested=True) + assert result.dtype == int + assert np.all(result == np.array([1, 2])) # Same as non-nested + + # in_shape is whole array + result = map_subarrays( + lambda x: 2*x, (2,), (2,), np.array([1, 2])) + assert result.dtype == int + assert np.all(result == np.array([2, 4])) + + # in_shape is whole array, nested + result = map_subarrays( + lambda x: 2*x, (2,), (2,), np.array([1, 2]), return_nested=True) + assert result.dtype == int + assert np.all(result == np.array([2, 4])) # Same as non-nested + + # len(out_shape) == 0 + result = map_subarrays( + np.sum, (2,), (), np.array([[1, 2], [2, 4]])) + assert result.dtype == int + assert np.all(result == np.array([3, 6])) + + # len(out_shape) == 0, nested output + result = map_subarrays( + np.sum, (2,), (), np.array([[1, 2], [2, 4]]), return_nested=True) + assert np.all(result == np.array([3, 6])) # Same as non-nested output + + # len(out_shape) == 0, non-numerical scalars + result = map_subarrays( + lambda x: _DummyScalar(x[0].val + x[1].val), (2,), (), + np.array([ + [_DummyScalar(1), _DummyScalar(2)], + [_DummyScalar(2), _DummyScalar(4)]])) + assert result.dtype == object + assert result.shape == (2,) + assert result[0] == _DummyScalar(3) + assert result[1] == _DummyScalar(6) + + # len(out_shape) != 0 + result = map_subarrays( + lambda x: np.array([x, 2*x]), (), (2,), np.array([1, 2])) + assert result.dtype == int + assert np.all(result == np.array([[1, 2], [2, 4]])) + + # len(out_shape) != 0, nested + result = map_subarrays( + lambda x: np.array([x, 2*x]), (), (2,), np.array([1, 2]), return_nested=True) + assert result.dtype == object + assert result.shape == (2,) + assert np.all(result[0] == np.array([1, 2])) + assert np.all(result[1] == np.array([2, 4])) + + # len(out_shape) != 0, non-numerical scalars + result = map_subarrays( + lambda x: np.array([_DummyScalar(x), _DummyScalar(2*x)]), (), (2,), + np.array([1, 2])) + assert result.dtype == object + assert result.shape == (2, 2) + assert np.all(result[0] == np.array([_DummyScalar(1), _DummyScalar(2)])) + assert np.all(result[1] == np.array([_DummyScalar(2), _DummyScalar(4)])) + + # Zero-size input array + result = map_subarrays( + lambda x: np.array([x, 2*x]), (), (2,), np.empty((2, 0))) + assert result.dtype == object + assert result.shape == (2, 0, 2) + + # Zero-size input array, nested + result = map_subarrays( + lambda x: np.array([x, 2*x]), (), (2,), np.empty((2, 0)), + return_nested=True) + assert result.dtype == object + assert result.shape == (2, 0) + + +def test_rec_map_subarrays(): + """Test rec_map_subarrays.""" + from grudge.tools import rec_map_subarrays + + # Scalar + result = rec_map_subarrays( + lambda x: np.array([x, 2*x]), (), (2,), 1) + assert result.dtype == int + assert np.all(result == np.array([1, 2])) + + # Scalar, non-numerical + result = rec_map_subarrays( + lambda x: np.array([x.val, 2*x.val]), (), (2,), _DummyScalar(1), + scalar_cls=_DummyScalar) + assert result.dtype == int + assert np.all(result == np.array([1, 2])) + + # Array of scalars + result = rec_map_subarrays( + np.sum, (2,), (), np.array([[1, 2], [2, 4]])) + assert result.dtype == int + assert np.all(result == np.array([3, 6])) + + # Array of scalars, non-numerical + result = rec_map_subarrays( + lambda x: x[0].val + x[1].val, (2,), (), + np.array([ + [_DummyScalar(1), _DummyScalar(2)], + [_DummyScalar(2), _DummyScalar(4)]]), + scalar_cls=_DummyScalar) + assert result.dtype == int + assert np.all(result == np.array([3, 6])) + + # Array container + result = rec_map_subarrays( + np.sum, (2,), (), make_obj_array([np.array([1, 2]), np.array([2, 4])])) + assert result.dtype == object + assert result[0] == 3 + assert result[1] == 6 + + # Array container, non-numerical scalars + result = rec_map_subarrays( + lambda x: x[0].val + x[1], (2,), (), + make_obj_array([ + np.array([_DummyScalar(1), 2]), + np.array([_DummyScalar(2), 4])]), + scalar_cls=_DummyScalar) + assert result.dtype == object + assert result[0] == 3 + assert result[1] == 6 + +# }}} + + +# You can test individual routines by typing +# $ python test_tools.py 'test_routine()' + +if __name__ == "__main__": + import sys + if len(sys.argv) > 1: + exec(sys.argv[1]) + else: + pytest.main([__file__]) + +# vim: fdm=marker From d4b21e0abdd106bd7f268e76a5660a6446dfa7b2 Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Wed, 27 Apr 2022 16:38:28 -0500 Subject: [PATCH 15/15] cosmetic change --- grudge/op.py | 28 ++++++++++++---------------- 1 file changed, 12 insertions(+), 16 deletions(-) diff --git a/grudge/op.py b/grudge/op.py index 1607cba61..07b94f022 100644 --- a/grudge/op.py +++ b/grudge/op.py @@ -267,10 +267,9 @@ def local_grad( dd_in = dof_desc.DOFDesc("vol", dof_desc.DISCR_TAG_BASE) from grudge.tools import rec_map_subarrays return rec_map_subarrays( - f=partial(_strong_scalar_grad, dcoll, dd_in), - in_shape=(), - out_shape=(dcoll.ambient_dim,), - ary=vec, scalar_cls=DOFArray, return_nested=nested,) + partial(_strong_scalar_grad, dcoll, dd_in), + (), (dcoll.ambient_dim,), + vec, scalar_cls=DOFArray, return_nested=nested,) def local_d_dx( @@ -323,12 +322,11 @@ def local_div(dcoll: DiscretizationCollection, vecs) -> ArrayOrContainerT: """ from grudge.tools import rec_map_subarrays return rec_map_subarrays( - f=lambda vec: sum( + lambda vec: sum( local_d_dx(dcoll, i, vec_i) for i, vec_i in enumerate(vec)), - in_shape=(dcoll.ambient_dim,), - out_shape=(), - ary=vecs, scalar_cls=DOFArray) + (dcoll.ambient_dim,), (), + vecs, scalar_cls=DOFArray) # }}} @@ -429,10 +427,9 @@ def weak_local_grad( from grudge.tools import rec_map_subarrays return rec_map_subarrays( - f=partial(_weak_scalar_grad, dcoll, dd_in), - in_shape=(), - out_shape=(dcoll.ambient_dim,), - ary=vecs, scalar_cls=DOFArray, return_nested=nested) + partial(_weak_scalar_grad, dcoll, dd_in), + (), (dcoll.ambient_dim,), + vecs, scalar_cls=DOFArray, return_nested=nested) def weak_local_d_dx(dcoll: DiscretizationCollection, *args) -> ArrayOrContainerT: @@ -534,12 +531,11 @@ def weak_local_div(dcoll: DiscretizationCollection, *args) -> ArrayOrContainerT: from grudge.tools import rec_map_subarrays return rec_map_subarrays( - f=lambda vec: sum( + lambda vec: sum( weak_local_d_dx(dcoll, dd_in, i, vec_i) for i, vec_i in enumerate(vec)), - in_shape=(dcoll.ambient_dim,), - out_shape=(), - ary=vecs, scalar_cls=DOFArray) + (dcoll.ambient_dim,), (), + vecs, scalar_cls=DOFArray) # }}}