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/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 0f41ce5a4..17a22918c 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__ = """ @@ -61,17 +68,17 @@ 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 import grudge.dof_desc as dof_desc from grudge.dof_desc import DD_VOLUME_ALL, FACE_RESTR_ALL -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, @@ -86,15 +93,64 @@ elementwise_integral, ) -from grudge.trace_pair import ( # noqa: F401 +from grudge.trace_pair import ( + project_tracepair, + tracepair_with_discr_tag, interior_trace_pair, interior_trace_pairs, + local_interior_trace_pair, + inter_volume_trace_pairs, + local_inter_volume_trace_pairs, cross_rank_trace_pairs, + cross_rank_inter_volume_trace_pairs, bdry_trace_pair, bv_trace_pair ) +__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", + + "project_tracepair", + "tracepair_with_discr_tag", + "interior_trace_pair", + "interior_trace_pairs", + "local_interior_trace_pair", + "inter_volume_trace_pairs", + "local_inter_volume_trace_pairs", + "cross_rank_trace_pairs", + "cross_rank_inter_volume_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( @@ -157,90 +213,6 @@ 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 = DD_VOLUME_ALL - elif len(args) == 2: - dd, vecs = args - else: - raise TypeError("invalid number of arguments") - - if not isinstance(vecs, np.ndarray): - # vecs is not an object array -> treat as array container - return map_array_container( - partial(_div_helper, dcoll, diff_func, dd), vecs) - - assert vecs.dtype == object - - if vecs.size: - sample_vec = vecs[(0,)*vecs.ndim] - - if isinstance(sample_vec, np.ndarray): - assert sample_vec.dtype == object - # vecs is an object array containing further object arrays - # -> treat as array container - return map_array_container( - partial(_div_helper, dcoll, diff_func, dd), vecs) - - if vecs.shape[-1] != dcoll.ambient_dim: - raise ValueError("last/innermost dimension of *vecs* argument doesn't match " - "ambient dimension") - - div_result_shape = vecs.shape[:-1] - - if len(div_result_shape) == 0: - return sum(diff_func(dd, i, vec_i) for i, vec_i in enumerate(vecs)) - else: - result = np.zeros(div_result_shape, dtype=object) - for idx in np.ndindex(div_result_shape): - result[idx] = sum( - diff_func(dd, i, vec_i) for i, vec_i in enumerate(vecs[idx])) - return result - - -def _grad_helper(dcoll, scalar_grad, *args, nested): - if len(args) == 1: - vec, = args - dd_in = dof_desc.DD_VOLUME_ALL - elif len(args) == 2: - dd_in, vec = args - else: - raise TypeError("invalid number of arguments") - - if isinstance(vec, np.ndarray): - # Occasionally, data structures coming from *mirgecom* will - # contain empty object arrays as placeholders for fields. - # For example, species mass fractions is an empty object array when - # running in a single-species configuration. - # This hack here ensures that these empty arrays, at the very least, - # have their shape updated when applying the gradient operator - if vec.size == 0: - return vec.reshape(vec.shape + (dcoll.ambient_dim,)) - - # For containers with ndarray data (such as momentum/velocity), - # the gradient is matrix-valued, so we compute the gradient for - # each component. If requested (via 'not nested'), return a matrix of - # derivatives by stacking the results. - grad = obj_array_vectorize( - lambda el: _grad_helper( - dcoll, scalar_grad, dd_in, el, nested=nested), vec) - if nested: - return grad - else: - return np.stack(grad, axis=0) - - if not isinstance(vec, DOFArray): - return map_array_container( - partial(_grad_helper, dcoll, scalar_grad, dd_in, nested=nested), vec) - - return scalar_grad(dcoll, dd_in, vec) - -# }}} - - # {{{ Derivative operators def _reference_derivative_matrices(actx: ArrayContext, @@ -297,8 +269,12 @@ 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 = DD_VOLUME_ALL + from grudge.tools import rec_map_subarrays + return rec_map_subarrays( + partial(_strong_scalar_grad, dcoll, dd_in), + (), (dcoll.ambient_dim,), + vec, scalar_cls=DOFArray, return_nested=nested,) def local_d_dx( @@ -349,10 +325,13 @@ 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) + from grudge.tools import rec_map_subarrays + return rec_map_subarrays( + lambda vec: sum( + local_d_dx(dcoll, i, vec_i) + for i, vec_i in enumerate(vec)), + (dcoll.ambient_dim,), (), + vecs, scalar_cls=DOFArray) # }}} @@ -443,7 +422,19 @@ 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 = DD_VOLUME_ALL + elif len(args) == 2: + dd_in, vecs = args + else: + raise TypeError("invalid number of arguments") + + from grudge.tools import rec_map_subarrays + return rec_map_subarrays( + 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: @@ -535,11 +526,21 @@ 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 = DD_VOLUME_ALL + elif len(args) == 2: + dd_in, vecs = args + else: + raise TypeError("invalid number of arguments") + + from grudge.tools import rec_map_subarrays + return rec_map_subarrays( + lambda vec: sum( + weak_local_d_dx(dcoll, dd_in, i, vec_i) + for i, vec_i in enumerate(vec)), + (dcoll.ambient_dim,), (), + vecs, scalar_cls=DOFArray) # }}} 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/grudge/tools.py b/grudge/tools.py index 9723f9249..3427ba1cb 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" @@ -25,7 +27,10 @@ """ import numpy as np -from pytools import levi_civita +from pytools import levi_civita, product +from typing import Tuple, Callable, Optional, Union, Any +from functools import partial +from arraycontext.container import ArrayOrContainerT def is_zero(x): @@ -235,4 +240,133 @@ def build_jacobian(actx, f, base_state, stepsize): return mat +def map_subarrays( + f: Callable[[Any], Any], + in_shape: Tuple[int, ...], out_shape: Tuple[int, ...], + ary: Any, *, return_nested=False) -> Any: + """ + 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` 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`. + + *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) + + 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:: + + map_subarrays(f, (2,), (2, 2), ary) + + 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:: + + 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: + 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) + 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) + 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: + 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( + f: Callable[[Any], Any], + in_shape: Tuple[int, ...], + out_shape: Tuple[int, ...], + ary: ArrayOrContainerT, *, + 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 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 + 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( + 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, scalar_cls=scalar_cls, + return_nested=return_nested), + ary) + + # vim: foldmethod=marker diff --git a/grudge/trace_pair.py b/grudge/trace_pair.py index 7d9f216bf..73e9c068e 100644 --- a/grudge/trace_pair.py +++ b/grudge/trace_pair.py @@ -2,13 +2,18 @@ Trace Pairs ^^^^^^^^^^^ -Container class ---------------- +Container class and auxiliary functionality +------------------------------------------- .. autoclass:: TracePair -Boundary traces ---------------- +.. currentmodule:: grudge.op + +.. autoclass:: project_tracepair +.. autoclass:: tracepair_with_discr_tag + +Boundary trace functions +------------------------ .. autofunction:: bdry_trace_pair .. autofunction:: bv_trace_pair @@ -836,4 +841,27 @@ def start_comm(remote_part_id): # }}} +# {{{ 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 diff --git a/test/test_grudge.py b/test/test_grudge.py index c0a562cef..d7e788560 100644 --- a/test/test_grudge.py +++ b/test/test_grudge.py @@ -727,6 +727,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: @@ -737,35 +739,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) diff --git a/test/test_grudge_sym_old.py b/test/test_grudge_sym_old.py index b59c98b18..ee6a00671 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 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