Source code for sumpy.toys

__copyright__ = """
Copyright (C) 2017 Andreas Kloeckner
Copyright (C) 2017 Matt Wala
"""

__license__ = """
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
"""

from pytools import memoize_method
from numbers import Number
from sumpy.kernel import TargetTransformationRemover

import numpy as np  # noqa: F401
import loopy as lp  # noqa: F401
import pyopencl as cl

import logging
logger = logging.getLogger(__name__)

__doc__ = """

This module provides a convenient interface for numerical experiments with
local and multipole expansions.

.. autoclass:: ToyContext
.. autoclass:: PotentialSource
.. autoclass:: ConstantPotential
.. autoclass:: PointSources

These functions manipulate these potentials:

.. autofunction:: multipole_expand
.. autofunction:: local_expand
.. autofunction:: logplot
.. autofunction:: combine_inner_outer
.. autofunction:: combine_halfspace
.. autofunction:: combine_halfspace_and_outer

These functions help with plotting:

.. autofunction:: draw_box
.. autofunction:: draw_circle
.. autofunction:: draw_annotation
.. autofunction:: draw_schematic

These are created behind the scenes and are not typically directly instantiated
by users:

.. autoclass:: OneOnBallPotential
.. autoclass:: HalfspaceOnePotential
.. autoclass:: ExpansionPotentialSource
.. autoclass:: MultipoleExpansion
.. autoclass:: LocalExpansion
.. autoclass:: Sum
.. autoclass:: Product
.. autoclass:: SchematicVisitor

"""


# {{{ context

[docs]class ToyContext: """This class functions as a container for generated code and 'behind-the-scenes' information. .. automethod:: __init__ """
[docs] def __init__(self, cl_context, kernel, mpole_expn_class=None, local_expn_class=None, expansion_factory=None, extra_source_kwargs=None, extra_kernel_kwargs=None): self.cl_context = cl_context self.queue = cl.CommandQueue(self.cl_context) self.kernel = kernel self.no_target_deriv_kernel = TargetTransformationRemover()(kernel) if expansion_factory is None: from sumpy.expansion import DefaultExpansionFactory expansion_factory = DefaultExpansionFactory() if mpole_expn_class is None: mpole_expn_class = \ expansion_factory.get_multipole_expansion_class(kernel) if local_expn_class is None: local_expn_class = \ expansion_factory.get_local_expansion_class(kernel) self.mpole_expn_class = mpole_expn_class self.local_expn_class = local_expn_class if extra_source_kwargs is None: extra_source_kwargs = {} if extra_kernel_kwargs is None: extra_kernel_kwargs = {} self.extra_source_kwargs = extra_source_kwargs self.extra_kernel_kwargs = extra_kernel_kwargs extra_source_and_kernel_kwargs = extra_source_kwargs.copy() extra_source_and_kernel_kwargs.update(extra_kernel_kwargs) self.extra_source_and_kernel_kwargs = extra_source_and_kernel_kwargs
@memoize_method def get_p2p(self): from sumpy.p2p import P2P return P2P(self.cl_context, (self.kernel,), exclude_self=False) @memoize_method def get_p2m(self, order): from sumpy import P2EFromSingleBox return P2EFromSingleBox(self.cl_context, self.mpole_expn_class(self.no_target_deriv_kernel, order), kernels=(self.kernel,)) @memoize_method def get_p2l(self, order): from sumpy import P2EFromSingleBox return P2EFromSingleBox(self.cl_context, self.local_expn_class(self.no_target_deriv_kernel, order), kernels=(self.kernel,)) @memoize_method def get_m2p(self, order): from sumpy import E2PFromSingleBox return E2PFromSingleBox(self.cl_context, self.mpole_expn_class(self.no_target_deriv_kernel, order), (self.kernel,)) @memoize_method def get_l2p(self, order): from sumpy import E2PFromSingleBox return E2PFromSingleBox(self.cl_context, self.local_expn_class(self.no_target_deriv_kernel, order), (self.kernel,)) @memoize_method def get_m2m(self, from_order, to_order): from sumpy import E2EFromCSR return E2EFromCSR(self.cl_context, self.mpole_expn_class(self.no_target_deriv_kernel, from_order), self.mpole_expn_class(self.no_target_deriv_kernel, to_order)) @memoize_method def get_m2l(self, from_order, to_order): from sumpy import E2EFromCSR return E2EFromCSR(self.cl_context, self.mpole_expn_class(self.no_target_deriv_kernel, from_order), self.local_expn_class(self.no_target_deriv_kernel, to_order)) @memoize_method def get_l2l(self, from_order, to_order): from sumpy import E2EFromCSR return E2EFromCSR(self.cl_context, self.local_expn_class(self.no_target_deriv_kernel, from_order), self.local_expn_class(self.no_target_deriv_kernel, to_order))
# }}} # {{{ helpers def _p2e(psource, center, rscale, order, p2e, expn_class, expn_kwargs): source_boxes = np.array([0], dtype=np.int32) box_source_starts = np.array([0], dtype=np.int32) box_source_counts_nonchild = np.array( [psource.points.shape[-1]], dtype=np.int32) toy_ctx = psource.toy_ctx center = np.asarray(center) centers = np.array(center, dtype=np.float64).reshape( toy_ctx.kernel.dim, 1) evt, (coeffs,) = p2e(toy_ctx.queue, source_boxes=source_boxes, box_source_starts=box_source_starts, box_source_counts_nonchild=box_source_counts_nonchild, centers=centers, sources=psource.points, strengths=(psource.weights,), rscale=rscale, nboxes=1, tgt_base_ibox=0, #flags="print_hl_cl", out_host=True, **toy_ctx.extra_source_and_kernel_kwargs) return expn_class(toy_ctx, center, rscale, order, coeffs[0], derived_from=psource, **expn_kwargs) def _e2p(psource, targets, e2p): ntargets = targets.shape[-1] boxes = np.array([0], dtype=np.int32) box_target_starts = np.array([0], dtype=np.int32) box_target_counts_nonchild = np.array([ntargets], dtype=np.int32) toy_ctx = psource.toy_ctx centers = np.array(psource.center, dtype=np.float64).reshape( toy_ctx.kernel.dim, 1) coeffs = np.array([psource.coeffs]) evt, (pot,) = e2p( toy_ctx.queue, src_expansions=coeffs, src_base_ibox=0, target_boxes=boxes, box_target_starts=box_target_starts, box_target_counts_nonchild=box_target_counts_nonchild, centers=centers, rscale=psource.rscale, targets=targets, #flags="print_hl_cl", out_host=True, **toy_ctx.extra_kernel_kwargs) return pot def _e2e(psource, to_center, to_rscale, to_order, e2e, expn_class, expn_kwargs): toy_ctx = psource.toy_ctx target_boxes = np.array([1], dtype=np.int32) src_box_starts = np.array([0, 1], dtype=np.int32) src_box_lists = np.array([0], dtype=np.int32) centers = (np.array( [ # box 0: source psource.center, # box 1: target to_center, ], dtype=np.float64)).T.copy() coeffs = np.array([psource.coeffs]) evt, (to_coeffs,) = e2e( toy_ctx.queue, src_expansions=coeffs, src_base_ibox=0, tgt_base_ibox=0, ntgt_level_boxes=2, target_boxes=target_boxes, src_box_starts=src_box_starts, src_box_lists=src_box_lists, centers=centers, src_rscale=psource.rscale, tgt_rscale=to_rscale, #flags="print_hl_cl", out_host=True, **toy_ctx.extra_kernel_kwargs) return expn_class(toy_ctx, to_center, to_rscale, to_order, to_coeffs[1], derived_from=psource, **expn_kwargs) # }}} # {{{ potential source classes
[docs]class PotentialSource: """A base class for all classes representing potentials that can be evaluated anywhere in space. .. automethod:: eval Supports (lazy) arithmetic: .. automethod:: __neg__ .. automethod:: __add__ .. automethod:: __radd__ .. automethod:: __sub__ .. automethod:: __rsub__ .. automethod:: __mul__ .. automethod:: __rmul__ """ def __init__(self, toy_ctx): self.toy_ctx = toy_ctx
[docs] def eval(self, targets): raise NotImplementedError()
[docs] def __neg__(self): return -1*self
[docs] def __add__(self, other): if isinstance(other, (Number, np.number)): other = ConstantPotential(self.toy_ctx, other) elif not isinstance(other, PotentialSource): return NotImplemented return Sum((self, other))
__radd__ = __add__
[docs] def __sub__(self, other): return self.__add__(-other)
[docs] def __rsub__(self, other): return (-self).__add__(other)
[docs] def __mul__(self, other): if isinstance(other, (Number, np.number)): other = ConstantPotential(self.toy_ctx, other) elif not isinstance(other, PotentialSource): return NotImplemented return Product((self, other))
__rmul__ = __mul__
[docs]class ConstantPotential(PotentialSource): """ .. automethod:: __init__ """
[docs] def __init__(self, toy_ctx, value): super().__init__(toy_ctx) self.value = np.array(value)
def eval(self, targets): pot = np.empty(targets.shape[-1], dtype=self.value.dtype) pot.fill(self.value) return pot
[docs]class OneOnBallPotential(PotentialSource): """ .. automethod:: __init__ """
[docs] def __init__(self, toy_ctx, center, radius): super().__init__(toy_ctx) self.center = np.asarray(center) self.radius = radius
def eval(self, targets): dist_vec = targets - self.center[:, np.newaxis] return (np.sum(dist_vec**2, axis=0) < self.radius**2).astype(np.float64)
[docs]class HalfspaceOnePotential(PotentialSource): """ .. automethod:: __init__ """
[docs] def __init__(self, toy_ctx, center, axis, side=1): super().__init__(toy_ctx) self.center = np.asarray(center) self.axis = axis self.side = side
def eval(self, targets): return ( (self.side*(targets[self.axis] - self.center[self.axis])) >= 0 ).astype(np.float64)
[docs]class PointSources(PotentialSource): """ .. attribute:: points ``[ndim, npoints]`` .. automethod:: __init__ """
[docs] def __init__(self, toy_ctx, points, weights, center=None): super().__init__(toy_ctx) self.points = points self.weights = weights self._center = center
def eval(self, targets): evt, (potential,) = self.toy_ctx.get_p2p()( self.toy_ctx.queue, targets, self.points, [self.weights], out_host=True, **self.toy_ctx.extra_source_and_kernel_kwargs) return potential @property def center(self): if self._center is not None: return self._center return np.average(self.points, axis=1)
[docs]class ExpansionPotentialSource(PotentialSource): """ .. attribute:: radius Not used mathematically. Just for visualization, purely advisory. .. attribute:: text_kwargs Passed to :func:`matplotlib.pyplot.annotate`. Used for customizing the expansion label. Changing the label text is supported by passing the kwarg *s*. Just for visualization, purely advisory. """ def __init__(self, toy_ctx, center, rscale, order, coeffs, derived_from, radius=None, expn_style=None, text_kwargs=None): super().__init__(toy_ctx) self.center = np.asarray(center) self.rscale = rscale self.order = order self.coeffs = coeffs self.derived_from = derived_from self.radius = radius self.expn_style = expn_style self.text_kwargs = text_kwargs def with_coeffs(self, coeffs): return type(self)(self.toy_ctx, self.center, self.rscale, self.order, coeffs, self.derived_from, radius=self.radius, expn_style=self.expn_style, text_kwargs=self.text_kwargs)
[docs]class MultipoleExpansion(ExpansionPotentialSource): def eval(self, targets): return _e2p(self, targets, self.toy_ctx.get_m2p(self.order))
[docs]class LocalExpansion(ExpansionPotentialSource): def eval(self, targets): return _e2p(self, targets, self.toy_ctx.get_l2p(self.order))
class PotentialExpressionNode(PotentialSource): def __init__(self, psources): from pytools import single_valued super().__init__( single_valued(psource.toy_ctx for psource in psources)) self.psources = psources @property def center(self): for psource in self.psources: try: return psource.center except AttributeError: pass raise ValueError("no psource with a center found")
[docs]class Sum(PotentialExpressionNode): def eval(self, targets): result = 0 for psource in self.psources: result = result + psource.eval(targets) return result
[docs]class Product(PotentialExpressionNode): def eval(self, targets): result = 1 for psource in self.psources: result = result * psource.eval(targets) return result
# }}}
[docs]def multipole_expand(psource, center, order=None, rscale=1, **expn_kwargs): if isinstance(psource, PointSources): if order is None: raise ValueError("order may not be None") return _p2e(psource, center, rscale, order, psource.toy_ctx.get_p2m(order), MultipoleExpansion, expn_kwargs) elif isinstance(psource, MultipoleExpansion): if order is None: order = psource.order return _e2e(psource, center, rscale, order, psource.toy_ctx.get_m2m(psource.order, order), MultipoleExpansion, expn_kwargs) else: raise TypeError("do not know how to expand '%s'" % type(psource).__name__)
[docs]def local_expand(psource, center, order=None, rscale=1, **expn_kwargs): if isinstance(psource, PointSources): if order is None: raise ValueError("order may not be None") return _p2e(psource, center, rscale, order, psource.toy_ctx.get_p2l(order), LocalExpansion, expn_kwargs) elif isinstance(psource, MultipoleExpansion): if order is None: order = psource.order return _e2e(psource, center, rscale, order, psource.toy_ctx.get_m2l(psource.order, order), LocalExpansion, expn_kwargs) elif isinstance(psource, LocalExpansion): if order is None: order = psource.order return _e2e(psource, center, rscale, order, psource.toy_ctx.get_l2l(psource.order, order), LocalExpansion, expn_kwargs) else: raise TypeError("do not know how to expand '%s'" % type(psource).__name__)
[docs]def logplot(fp, psource, **kwargs): fp.show_scalar_in_matplotlib( np.log10(np.abs(psource.eval(fp.points) + 1e-15)), **kwargs)
[docs]def combine_inner_outer(psource_inner, psource_outer, radius, center=None): if center is None: center = psource_inner.center if radius is None: radius = psource_inner.radius ball_one = OneOnBallPotential(psource_inner.toy_ctx, center, radius) return ( psource_inner * ball_one + psource_outer * (1 - ball_one))
[docs]def combine_halfspace(psource_pos, psource_neg, axis, center=None): if center is None: center = psource_pos.center halfspace_one = HalfspaceOnePotential(psource_pos.toy_ctx, center, axis) return ( psource_pos * halfspace_one + psource_neg * (1-halfspace_one))
[docs]def combine_halfspace_and_outer(psource_pos, psource_neg, psource_outer, axis, radius=None, center=None): if center is None: center = psource_pos.center if radius is None: center = psource_pos.radius return combine_inner_outer( combine_halfspace(psource_pos, psource_neg, axis, center), psource_outer, radius, center)
def l_inf(psource, radius, center=None, npoints=100, debug=False): if center is None: center = psource.center restr = psource * OneOnBallPotential(psource.toy_ctx, center, radius) from sumpy.visualization import FieldPlotter fp = FieldPlotter(center, extent=2*radius, npoints=npoints) z = restr.eval(fp.points) if debug: fp.show_scalar_in_matplotlib( np.log10(np.abs(z + 1e-15))) import matplotlib.pyplot as pt pt.colorbar() pt.show() return np.max(np.abs(z)) # {{{ schematic visualization
[docs]def draw_box(el, eh, **kwargs): import matplotlib.pyplot as pt import matplotlib.patches as mpatches from matplotlib.path import Path pathdata = [ (Path.MOVETO, (el[0], el[1])), (Path.LINETO, (eh[0], el[1])), (Path.LINETO, (eh[0], eh[1])), (Path.LINETO, (el[0], eh[1])), (Path.CLOSEPOLY, (el[0], el[1])), ] codes, verts = zip(*pathdata) path = Path(verts, codes) patch = mpatches.PathPatch(path, **kwargs) pt.gca().add_patch(patch)
[docs]def draw_circle(center, radius, **kwargs): center = np.asarray(center) import matplotlib.pyplot as plt plt.gca().add_patch(plt.Circle((center[0], center[1]), radius, **kwargs))
def draw_point(loc, **kwargs): import matplotlib.pyplot as plt plt.plot(*loc, marker="o", **kwargs)
[docs]def draw_annotation(to_pt, from_pt, label, arrowprops=None, **kwargs): """ :arg to_pt: Head of arrow :arg from_pt: Tail of arrow :arg label: Annotation label :arg arrowprops: Passed to arrowprops :arg kwargs: Passed to annotate """ if arrowprops is None: arrowprops = {} import matplotlib.pyplot as plt my_arrowprops = dict( facecolor="black", edgecolor="black", arrowstyle="->") my_arrowprops.update(arrowprops) plt.gca().annotate(label, xy=to_pt, xytext=from_pt, arrowprops=my_arrowprops, **kwargs)
[docs]class SchematicVisitor: def __init__(self, default_expn_style="circle"): self.default_expn_style = default_expn_style def rec(self, psource): getattr(self, "visit_"+type(psource).__name__.lower())(psource) def visit_pointsources(self, psource): import matplotlib.pyplot as plt plt.plot(psource.points[0], psource.points[1], "o", label="source") def visit_sum(self, psource): for ps in psource.psources: self.rec(ps) visit_product = visit_sum def visit_multipoleexpansion(self, psource): expn_style = self.default_expn_style if psource.expn_style is not None: expn_style = psource.expn_style if psource.radius is not None: if expn_style == "box": r2 = psource.radius / np.sqrt(2) draw_box(psource.center - r2, psource.center + r2, fill=None) elif expn_style == "circle": draw_circle(psource.center, psource.radius, fill=None) else: raise ValueError("unknown expn_style: %s" % self.expn_style) if psource.derived_from is None: return # Draw an annotation of the form # # ------> M text_kwargs = dict( verticalalignment="center", horizontalalignment="center") label = "${}_{{{}}}$".format( type(psource).__name__[0].lower().replace("l", "\\ell"), psource.order) if psource.text_kwargs is not None: psource_text_kwargs_copy = psource.text_kwargs.copy() label = psource_text_kwargs_copy.pop("s", label) text_kwargs.update(psource_text_kwargs_copy) shrinkB = 0 # noqa if isinstance(psource.derived_from, ExpansionPotentialSource): # Avoid overlapping the tail of the arrow with any expansion labels that # are present at the tail. import matplotlib as mpl font_size = mpl.rcParams["font.size"] shrinkB = 7/8 * font_size # noqa arrowprops = dict(shrinkB=shrinkB, arrowstyle="<|-") draw_annotation(psource.derived_from.center, psource.center, label, arrowprops, **text_kwargs) self.rec(psource.derived_from) visit_localexpansion = visit_multipoleexpansion
[docs]def draw_schematic(psource, **kwargs): SchematicVisitor(**kwargs).rec(psource) import matplotlib.pyplot as plt plt.gca().set_aspect("equal") plt.tight_layout()
# }}} # vim: foldmethod=marker