diff --git a/test/test_interpolation.py b/test/test_interpolation.py index bcc74a483bd74c38374567151752b5c48629ab85..d38cd85d601f90d12ef967e4a574299336cb752b 100644 --- a/test/test_interpolation.py +++ b/test/test_interpolation.py @@ -29,8 +29,9 @@ from meshmode.mesh.generation import generate_regular_rect_mesh from meshmode.discretization import Discretization from meshmode.discretization.poly_element import ( PolynomialWarpAndBlendGroupFactory) -from volumential.interpolation import (ElementsToSourcesLookupBuilder, - interpolate_from_meshmode) +from volumential.interpolation import ( + ElementsToSourcesLookupBuilder, LeavesToNodesLookupBuilder, + interpolate_from_meshmode, interpolate_to_meshmode) from volumential.geometry import BoundingBoxFactory, BoxFMMGeometryFactory @@ -153,6 +154,88 @@ def drive_test_from_meshmode_interpolation_2d( return resid +def drive_test_to_meshmode_exact_interpolation_2d( + cl_ctx, queue, + degree, nel_1d, n_levels, q_order, + a=-0.5, b=0.5, seed=0): + """ + meshmode mesh control: nel_1d, degree + volumential mesh control: n_levels, q_order + """ + dim = 2 + + mesh = generate_regular_rect_mesh( + a=(a, a), + b=(b, b), + n=(nel_1d, nel_1d)) + + arr_ctx = PyOpenCLArrayContext(queue) + group_factory = PolynomialWarpAndBlendGroupFactory(order=degree) + discr = Discretization(arr_ctx, mesh, group_factory) + + bbox_fac = BoundingBoxFactory(dim=2) + boxfmm_fac = BoxFMMGeometryFactory( + cl_ctx, dim=dim, order=q_order, + nlevels=n_levels, bbox_getter=bbox_fac, + expand_to_hold_mesh=mesh, mesh_padding_factor=0.) + boxgeo = boxfmm_fac(queue) + lookup_fac = LeavesToNodesLookupBuilder( + cl_ctx, trav=boxgeo.trav, discr=discr) + lookup, evt = lookup_fac(queue) + + # algebraically exact interpolation + func = random_polynomial_func(dim, degree, seed) + + tree = boxgeo.tree.get(queue) + potential = func(np.vstack(tree.sources)) + res = interpolate_to_meshmode(queue, potential, lookup).get(queue) + + ref = eval_func_on_discr_nodes(queue, discr, func).get(queue) + return np.allclose(ref, res) + + +def drive_test_to_meshmode_interpolation_2d( + cl_ctx, queue, + degree, nel_1d, n_levels, q_order, + a=-0.5, b=0.5, seed=0): + """ + meshmode mesh control: nel_1d, degree + volumential mesh control: n_levels, q_order + """ + dim = 2 + + mesh = generate_regular_rect_mesh( + a=(a, a), + b=(b, b), + n=(nel_1d, nel_1d)) + + arr_ctx = PyOpenCLArrayContext(queue) + group_factory = PolynomialWarpAndBlendGroupFactory(order=degree) + discr = Discretization(arr_ctx, mesh, group_factory) + + bbox_fac = BoundingBoxFactory(dim=2) + boxfmm_fac = BoxFMMGeometryFactory( + cl_ctx, dim=dim, order=q_order, + nlevels=n_levels, bbox_getter=bbox_fac, + expand_to_hold_mesh=mesh, mesh_padding_factor=0.) + boxgeo = boxfmm_fac(queue) + lookup_fac = LeavesToNodesLookupBuilder( + cl_ctx, trav=boxgeo.trav, discr=discr) + lookup, evt = lookup_fac(queue) + + def func(pts): + x, y = pts + return np.sin(x + y) * x + + tree = boxgeo.tree.get(queue) + potential = func(np.vstack(tree.sources)) + res = interpolate_to_meshmode(queue, potential, lookup).get(queue) + + ref = eval_func_on_discr_nodes(queue, discr, func).get(queue) + resid = np.linalg.norm(ref - res, ord=np.inf) + return resid + + @pytest.mark.parametrize("params", [ [1, 8, 3, 1], [1, 8, 5, 2], [1, 8, 7, 3], [2, 16, 3, 1], [2, 32, 5, 2], [2, 64, 7, 3], @@ -174,6 +257,27 @@ def test_from_meshmode_interpolation_2d_nonexact(ctx_factory, params): cl_ctx, queue, *params) < 1e-3 +@pytest.mark.parametrize("params", [ + [1, 8, 3, 2], [1, 8, 5, 2], [1, 8, 7, 3], + [2, 16, 3, 3], [3, 32, 5, 4], [4, 64, 7, 5], + ]) +def test_to_meshmode_interpolation_2d_exact(ctx_factory, params): + cl_ctx = ctx_factory() + queue = cl.CommandQueue(cl_ctx) + assert drive_test_to_meshmode_exact_interpolation_2d( + cl_ctx, queue, *params) + + +@pytest.mark.parametrize("params", [ + [1, 32, 7, 2], [2, 64, 5, 3], [3, 64, 6, 4] + ]) +def test_to_meshmode_interpolation_2d_nonexact(ctx_factory, params): + cl_ctx = ctx_factory() + queue = cl.CommandQueue(cl_ctx) + assert drive_test_to_meshmode_interpolation_2d( + cl_ctx, queue, *params) < 1e-3 + + if __name__ == '__main__': cl_ctx = cl.create_some_context() queue = cl.CommandQueue(cl_ctx) diff --git a/volumential/geometry.py b/volumential/geometry.py index c3d334a9546e80f045d78dbf50bd7058396414b1..e454ef8596d13b293daffc175dfec3d6b94fbb96 100644 --- a/volumential/geometry.py +++ b/volumential/geometry.py @@ -273,11 +273,11 @@ class BoxFMMGeometryFactory(): q_points = self._get_q_points(queue) - tree, _ = tb(queue, - particles=q_points, targets=q_points, bbox=self._bbox, - max_particles_in_box=( - self.n_q_points_per_cell * (2**self.dim) - 1), - kind="adaptive-level-restricted") + tree, _ = tb(queue, particles=q_points, targets=q_points, + bbox=self._bbox, max_particles_in_box=( + self.n_q_points_per_cell * (2**self.dim) * (2**self.dim) + - 1), + kind="adaptive-level-restricted") from boxtree.traversal import FMMTraversalBuilder tg = FMMTraversalBuilder(self.cl_context) diff --git a/volumential/interpolation.py b/volumential/interpolation.py index dd3dfac0d593ea59a3536f01311145a3d7e89574..4a737cecc5c5933d1af0c8056c0cb566e7bcc952 100644 --- a/volumential/interpolation.py +++ b/volumential/interpolation.py @@ -25,13 +25,12 @@ THE SOFTWARE. import numpy as np import pyopencl as cl import loopy as lp -from pyopencl.algorithm import KeyValueSorter from pytools import memoize_method, ProcessLogger from pytools.obj_array import make_obj_array -from boxtree.area_query import AreaQueryBuilder from boxtree.tools import DeviceDataRecord from meshmode.array_context import PyOpenCLArrayContext -from meshmode.dof_array import unflatten +from meshmode.dof_array import unflatten, flatten, thaw +from volumential.volume_fmm import interpolate_volume_potential import logging logger = logging.getLogger(__name__) @@ -104,9 +103,11 @@ class ElementsToSourcesLookup(DeviceDataRecord): class LeavesToNodesLookup(DeviceDataRecord): """ - .. attribute:: tree + .. attribute:: trav - The :class:`boxtree.Tree` instance representing the box mesh. + The :class:`boxtree.FMMTraversalInfo` instance representing the + box mesh with metadata needed for interpolation. It contains a + reference to the underlying tree as `trav.tree`. .. attribute:: discr @@ -127,9 +128,13 @@ class LeavesToNodesLookup(DeviceDataRecord): .. note:: Only leaf boxes have non-empty entries in this table. Nonetheless, this list is indexed by the global box index. - .. attribute:: box_nodes_in_element_lists + .. attribute:: nodes_in_leaf_lists - Indices into :attr:`tree.sources`. + Indices into :attr:`discr.nodes()`. + + .. note:: Unlike :class:`ElementsToSourcesLookup`, lists are not disjoint + in the leaves-to-nodes lookup. :mod:`volumential` automatically computes + the average contribution from overlapping boxes. .. automethod:: get """ @@ -159,7 +164,10 @@ class ElementsToSourcesLookupBuilder: self.tree = tree self.discr = discr + from pyopencl.algorithm import KeyValueSorter self.key_value_sorter = KeyValueSorter(context) + + from boxtree.area_query import AreaQueryBuilder self.area_query_builder = AreaQueryBuilder(self.context) # {{{ kernel generation @@ -323,6 +331,53 @@ class ElementsToSourcesLookupBuilder: # }}} End elements-to-sources lookup builder +# {{{ leaves-to-nodes lookup builder + +class LeavesToNodesLookupBuilder: + """Given a :mod:`meshmod` mesh and a :mod:`boxtree.Tree`, both discretizing + the same bounding box, this class helps to build a look-up table from + leaf boxes to mesh nodes that are positioned inside the box. + """ + + def __init__(self, context, trav, discr): + """ + :arg trav: a :class:`boxtree.FMMTraversalInfo` + :arg discr: a :class: `meshmode.discretization.Discretization` + + Boxes and elements can be non-aligned as long as the domains + (bounding boxes) are the same. + """ + assert trav.tree.dimensions == discr.dim + self.dim = discr.dim + self.context = context + self.trav = trav + self.discr = discr + + from boxtree.area_query import LeavesToBallsLookupBuilder + self.leaves_to_balls_lookup_builder = \ + LeavesToBallsLookupBuilder(self.context) + + def __call__(self, queue, tol=1e-12, wait_for=None): + """ + :arg queue: a :class:`pyopencl.CommandQueue` + :tol: nodes close enough to the boundary will be treated as + lying on the boundary, whose interpolated values are averaged. + """ + arr_ctx = PyOpenCLArrayContext(queue) + nodes = flatten(thaw(arr_ctx, self.discr.nodes())) + radii = cl.array.zeros_like(nodes[0]) + tol + + lbl_lookup, evt = self.leaves_to_balls_lookup_builder( + queue, self.trav.tree, nodes, radii, wait_for=wait_for) + + return LeavesToNodesLookup( + trav=self.trav, discr=self.discr, + nodes_in_leaf_starts=lbl_lookup.balls_near_box_starts, + nodes_in_leaf_lists=lbl_lookup.balls_near_box_lists), evt + +# }}} End leaves-to-nodes lookup builder + + # {{{ transform helper def compute_affine_transform(source_simplex, target_simplex): @@ -378,11 +433,14 @@ def invert_affine_transform(mat_a, disp_b): # {{{ from meshmode interpolation -def interpolate_from_meshmode(queue, dof_vec, elements_to_sources_lookup): - """ +def interpolate_from_meshmode(queue, dof_vec, elements_to_sources_lookup, + order="tree"): + """Interpolate a DoF vector from :mod:`meshmode`. + :arg dof_vec: a DoF vector representing a field in :mod:`meshmode` of shape ``(..., nnodes)``. - :arg elements_to_sources_lookup: a :class:`ElementsToSourcesLookup` + :arg elements_to_sources_lookup: a :class:`ElementsToSourcesLookup`. + :arg order: order of the output potential, either "tree" or "user". .. note:: This function currently supports meshes with just one element group. Also, the element group must be simplex-based. @@ -501,6 +559,65 @@ def interpolate_from_meshmode(queue, dof_vec, elements_to_sources_lookup): source_vec = cl.array.to_device(queue, source_vec) + if order == "tree": + pass # no need to do anything + elif order == "user": + source_vec = source_vec[tree.sorted_target_ids] # into user order + else: + raise ValueError(f"order must be 'tree' or 'user' (got {order}).") + return source_vec # }}} End from meshmode interpolation + + +# {{{ to meshmode interpolation + +def interpolate_to_meshmode(queue, potential, leaves_to_nodes_lookup, + order="tree"): + """ + :arg potential: a DoF vector representing a field in :mod:`volumential`, + in tree order. + :arg leaves_to_nodes_lookup: a :class:`LeavesToNodesLookup`. + :arg order: order of the input potential, either "tree" or "user". + + :returns: a :class:`pyopencl.Array` of shape (nnodes, 1) containing the + interpolated data. + """ + if order == "tree": + potential_in_tree_order = True + elif order == "user": + potential_in_tree_order = False + else: + raise ValueError(f"order must be 'tree' or 'user' (got {order}).") + + arr_ctx = PyOpenCLArrayContext(queue) + target_points = flatten(thaw( + arr_ctx, leaves_to_nodes_lookup.discr.nodes())) + + traversal = leaves_to_nodes_lookup.trav + tree = leaves_to_nodes_lookup.trav.tree + + dim = tree.dimensions + + # infer q_order from tree + pts_per_box = tree.ntargets // traversal.ntarget_boxes + print(pts_per_box, traversal.ntarget_boxes, tree.ntargets) + assert pts_per_box * traversal.ntarget_boxes == tree.ntargets + + # allow for +/- 0.25 floating point error + q_order = int(pts_per_box**(1 / dim) + 0.25) + assert q_order**dim == pts_per_box + + interp_p = interpolate_volume_potential( + target_points=target_points, traversal=traversal, + wrangler=None, potential=potential, + potential_in_tree_order=potential_in_tree_order, + dim=dim, tree=tree, queue=queue, q_order=q_order, + dtype=potential.dtype, lbl_lookup=None, + balls_near_box_starts=leaves_to_nodes_lookup.nodes_in_leaf_starts, + balls_near_box_lists=leaves_to_nodes_lookup.nodes_in_leaf_lists) + + return interp_p + +# }}} End to meshmode interpolation diff --git a/volumential/volume_fmm.py b/volumential/volume_fmm.py index 5695b228207cad15a2cabcb6dd5d7d6f68232664..409f904a84b0747be234d300378fbe1b38668c63 100644 --- a/volumential/volume_fmm.py +++ b/volumential/volume_fmm.py @@ -27,28 +27,24 @@ __doc__ = """ .. autofunction:: interpolate_volume_potential """ -import logging - -logger = logging.getLogger(__name__) - +import numpy as np import pyopencl as cl from boxtree.fmm import TimingRecorder from volumential.expansion_wrangler_interface import ExpansionWranglerInterface from volumential.expansion_wrangler_fpnd import ( FPNDSumpyExpansionWrangler, FPNDFMMLibExpansionWrangler) -import numpy as np +import logging +logger = logging.getLogger(__name__) + + +# {{{ volume FMM driver -def drive_volume_fmm( - traversal, - expansion_wrangler, - src_weights, - src_func, - direct_evaluation=False, - timing_data=None, - **kwargs -): +def drive_volume_fmm(traversal, expansion_wrangler, src_weights, src_func, + direct_evaluation=False, timing_data=None, + reorder_sources=True, reorder_potentials=True, + **kwargs): """ Top-level driver routine for volume potential calculation via fast multiple method. @@ -70,6 +66,10 @@ def drive_volume_fmm( Passed unmodified to *expansion_wrangler*. :arg src_func: Source 'density/weights/charges' function. Passed unmodified to *expansion_wrangler*. + :arg reorder_sources: Whether sources are in user order (if True, sources + are reordered into tree order before conducting FMM). + :arg reorder_potentials: Whether potentials should be in user order (if True, + potentials are reordered into user order before return). Returns the potentials computed by *expansion_wrangler*. """ @@ -92,10 +92,10 @@ def drive_volume_fmm( if isinstance(src_func, cl.array.Array): src_func = src_func.get(wrangler.queue) - logger.debug("reorder source weights") - - src_weights = wrangler.reorder_sources(src_weights) - src_func = wrangler.reorder_targets(src_func) + if reorder_sources: + logger.debug("reorder source weights") + src_weights = wrangler.reorder_sources(src_weights) + src_func = wrangler.reorder_targets(src_func) # {{{ Construct local multipoles @@ -142,8 +142,9 @@ def drive_volume_fmm( # 'list1_only' takes precedence over 'exclude_list1' if 'list1_only' in kwargs and kwargs['list1_only']: - logger.debug("reorder potentials") - result = wrangler.reorder_potentials(potentials) + if reorder_potentials: + logger.debug("reorder potentials") + result = wrangler.reorder_potentials(potentials) logger.debug("finalize potentials") result = wrangler.finalize_potentials(result) @@ -214,8 +215,9 @@ def drive_volume_fmm( # list 4 assert traversal.from_sep_close_bigger_starts is None - logger.debug("reorder potentials") - result = wrangler.reorder_potentials(potentials) + if reorder_potentials: + logger.debug("reorder potentials") + result = wrangler.reorder_potentials(potentials) logger.debug("finalize potentials") result = wrangler.finalize_potentials(result) @@ -334,8 +336,9 @@ def drive_volume_fmm( # }}} # {{{ Reorder potentials - logger.debug("reorder potentials") - result = wrangler.reorder_potentials(potentials) + if reorder_potentials: + logger.debug("reorder potentials") + result = wrangler.reorder_potentials(potentials) logger.debug("finalize potentials") result = wrangler.finalize_potentials(result) @@ -348,6 +351,10 @@ def drive_volume_fmm( return result +# }}} End volume FMM driver + + +# {{{ free form interpolation of potentials def compute_barycentric_lagrange_params(q_order): @@ -366,22 +373,36 @@ def compute_barycentric_lagrange_params(q_order): return (interp_points, interp_weights) -def interpolate_volume_potential( - target_points, traversal, wrangler, potential, target_radii=None, lbl_lookup=None -): +def interpolate_volume_potential(target_points, traversal, wrangler, potential, + potential_in_tree_order=False, + target_radii=None, lbl_lookup=None, **kwargs): """ Interpolate the volume potential, only works for tensor-product quadrature formulae. target_points and potential should be an cl array. - wrangler is used only for general info (nothing sumpy kernel specific) + :arg wrangler: Used only for general info (nothing sumpy kernel specific). May also + be None if the needed information is passed by kwargs. + :arg potential_in_tree_order: Whether the potential is in tree order (as opposed to + in user order). + :lbl_lookup: a :class:`boxtree.LeavesToBallsLookup` that has the lookup information + for target points. Can be None if the lookup lists are provided separately in + kwargs. If it is None and no other information is provided, the lookup will be + built from scratch. """ + if wrangler is not None: + dim = next(iter(wrangler.near_field_table.values()))[0].dim + tree = wrangler.tree + queue = wrangler.queue + q_order = wrangler.quad_order + dtype = wrangler.dtype + else: + dim = kwargs['dim'] + tree = kwargs['tree'] + queue = kwargs['queue'] + q_order = kwargs['q_order'] + dtype = kwargs['dtype'] - dim = next(iter(wrangler.near_field_table.values()))[0].dim - tree = wrangler.tree - queue = wrangler.queue - q_order = wrangler.quad_order ctx = queue.context - dtype = wrangler.dtype coord_dtype = tree.coord_dtype n_points = len(target_points[0]) @@ -391,27 +412,35 @@ def interpolate_volume_potential( traversal = traversal.to_device(queue) assert dim == len(target_points) + evt = None - # Building the lookup takes O(n*log(n)) - if lbl_lookup is None: - from boxtree.area_query import LeavesToBallsLookupBuilder + if ("balls_near_box_starts" in kwargs) and ("balls_near_box_lists" in kwargs): + balls_near_box_starts = kwargs["balls_near_box_starts"] + balls_near_box_lists = kwargs["balls_near_box_lists"] - lookup_builder = LeavesToBallsLookupBuilder(ctx) + else: + # Building the lookup takes O(n*log(n)) + if lbl_lookup is None: + from boxtree.area_query import LeavesToBallsLookupBuilder + lookup_builder = LeavesToBallsLookupBuilder(ctx) - if target_radii is None: - import pyopencl as cl + if target_radii is None: + # Set this number small enough so that all points found + # are inside the box + target_radii = cl.array.to_device( + queue, np.ones(n_points, dtype=coord_dtype) * 1e-12) - # Set this number small enough so that all points found - # are inside the box - target_radii = cl.array.to_device( - queue, np.ones(n_points, dtype=coord_dtype) * 1e-12 - ) + lbl_lookup, evt = lookup_builder( + queue, tree, target_points, target_radii) - lbl_lookup, evt = lookup_builder(queue, tree, target_points, target_radii) + balls_near_box_starts = lbl_lookup.balls_near_box_starts + balls_near_box_lists = lbl_lookup.balls_near_box_lists pout = cl.array.zeros(queue, n_points, dtype=dtype) multiplicity = cl.array.zeros(queue, n_points, dtype=dtype) - pout.add_event(evt) + + if evt is not None: + pout.add_event(evt) # map all boxes to a template [0,1]^2 so that the interpolation # weights and modes can be precomputed @@ -588,6 +617,13 @@ def interpolate_volume_potential( else: raise NotImplementedError + if potential_in_tree_order: + user_mode_ids = cl.array.arange( + queue, len(tree.user_source_ids), dtype=tree.user_source_ids.dtype) + else: + # fetching from user_source_ids converts potential to tree order + user_mode_ids = tree.user_source_ids + lpknl = loopy.set_options(lpknl, return_dict=True) lpknl = loopy.fix_parameters(lpknl, dim=int(dim), q_order=int(q_order)) lpknl = loopy.split_iname(lpknl, "tbox", 128, outer_tag="g.0", inner_tag="l.0") @@ -597,14 +633,14 @@ def interpolate_volume_potential( multiplicity=multiplicity, box_centers=tree.box_centers, box_levels=tree.box_levels, - balls_near_box_starts=lbl_lookup.balls_near_box_starts, - balls_near_box_lists=lbl_lookup.balls_near_box_lists, + balls_near_box_starts=balls_near_box_starts, + balls_near_box_lists=balls_near_box_lists, barycentric_lagrange_weights=blweights, barycentric_lagrange_points=blpoints, box_target_starts=tree.box_target_starts, box_target_counts_cumul=tree.box_target_counts_cumul, potential=potential, - user_mode_ids=tree.user_source_ids, + user_mode_ids=user_mode_ids, target_boxes=traversal.target_boxes, root_extent=tree.root_extent, n_tgt_boxes=len(traversal.target_boxes), @@ -617,5 +653,6 @@ def interpolate_volume_potential( return pout / multiplicity +# }}} End free form interpolation of potentials # vim: filetype=pyopencl:fdm=marker