__copyright__ = """
Copyright (C) 2013 Andreas Kloeckner
Copyright (C) 2018 Alexandru Fikl
"""
__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 typing import Optional
from pymbolic.mapper.evaluator import (
EvaluationMapper as PymbolicEvaluationMapper)
import numpy as np
import pyopencl as cl
import pyopencl.array # noqa
import pyopencl.clmath # noqa
from arraycontext import PyOpenCLArrayContext, thaw, freeze
from meshmode.dof_array import DOFArray, flatten
from pytools import memoize_in, memoize_method
from pytential.qbx.cost import AbstractQBXCostModel
from pytential import sym
import logging
logger = logging.getLogger(__name__)
__doc__ = """
.. autoclass :: BoundExpression
"""
# FIXME caches: fix up queues
class EvaluationMapperCSECacheKey:
"""Serves as a unique key for the common subexpression cache in
:meth:`GeometryCollection._get_cache`.
"""
class EvaluationMapperBoundOpCacheKey:
"""Serves as a unique key for the bound operator cache in
:meth:`GeometryCollection._get_cache`.
"""
# {{{ evaluation mapper base (shared, between actual eval and cost model)
def mesh_el_view(mesh, group_nr, global_array):
"""Return a view of *global_array* of shape
``(..., mesh.groups[group_nr].nelements)``
where *global_array* is of shape ``(..., nelements)``,
where *nelements* is the global (per-mesh) element count.
"""
group = mesh.groups[group_nr]
return global_array[
..., group.element_nr_base:group.element_nr_base + group.nelements] \
.reshape(
global_array.shape[:-1]
+ (group.nelements,))
class EvaluationMapperBase(PymbolicEvaluationMapper):
def __init__(self, bound_expr, actx: PyOpenCLArrayContext, context=None,
target_geometry=None,
target_points=None, target_normals=None, target_tangents=None):
if context is None:
context = {}
PymbolicEvaluationMapper.__init__(self, context)
self.bound_expr = bound_expr
self.places = bound_expr.places
self.array_context = actx
from arraycontext import PyOpenCLArrayContext
if not isinstance(actx, PyOpenCLArrayContext):
raise NotImplementedError("evaluation with non-PyOpenCL array context")
self.queue = actx.queue
# {{{ map_XXX
def _map_minmax(self, func, inherited_func, expr):
ev_children = [self.rec(ch) for ch in expr.children]
from functools import reduce
from meshmode.dof_array import DOFArray
if any(isinstance(ch, (cl.array.Array, DOFArray)) for ch in ev_children):
return reduce(func, ev_children)
else:
return inherited_func(expr)
def map_max(self, expr):
return self._map_minmax(
self.array_context.np.maximum,
super().map_max,
expr)
def map_min(self, expr):
return self._map_minmax(
self.array_context.np.minimum,
super().map_min,
expr)
def map_node_sum(self, expr):
# FIXME: make less CL specific
queue = self.array_context.queue
return sum(
cl.array.sum(grp_ary, queue=queue).get()[()]
for grp_ary in self.rec(expr.operand))
def map_node_max(self, expr):
# FIXME: make less CL specific
queue = self.array_context.queue
return max(
cl.array.max(grp_ary, queue=queue).get()[()]
for grp_ary in self.rec(expr.operand))
def map_node_min(self, expr):
# FIXME: make less CL specific
queue = self.array_context.queue
return min(
cl.array.min(grp_ary, queue=queue).get()[()]
for grp_ary in self.rec(expr.operand))
def _map_elementwise_reduction(self, reduction_name, expr):
@memoize_in(self.places, "elementwise_node_"+reduction_name)
def node_knl():
from arraycontext import make_loopy_program
return make_loopy_program(
"""{[iel, idof, jdof]:
0<=iel<nelements and
0<=idof, jdof<ndofs}""",
"""
result[iel, idof] = %s(jdof, operand[iel, jdof])
""" % reduction_name,
name="nodewise_reduce")
@memoize_in(self.places, "elementwise_"+reduction_name)
def element_knl():
from arraycontext import make_loopy_program
return make_loopy_program(
"""{[iel, jdof]:
0<=iel<nelements and
0<=jdof<ndofs}
""",
"""
result[iel, 0] = %s(jdof, operand[iel, jdof])
""" % reduction_name,
name="elementwise_reduce")
discr = self.places.get_discretization(
expr.dofdesc.geometry, expr.dofdesc.discr_stage)
operand = self.rec(expr.operand)
assert operand.shape == (len(discr.groups),)
def _reduce(knl, result):
for grp in discr.groups:
self.array_context.call_loopy(knl,
operand=operand[grp.index],
result=result[grp.index])
return result
dtype = operand.entry_dtype
granularity = expr.dofdesc.granularity
if granularity is sym.GRANULARITY_NODE:
return _reduce(node_knl(),
discr.empty(self.array_context, dtype=dtype))
elif granularity is sym.GRANULARITY_ELEMENT:
result = DOFArray(self.array_context, tuple([
self.array_context.empty((grp.nelements, 1), dtype=dtype)
for grp in discr.groups
]))
return _reduce(element_knl(), result)
else:
raise ValueError(f"unsupported granularity: {granularity}")
def map_elementwise_sum(self, expr):
return self._map_elementwise_reduction("sum", expr)
def map_elementwise_min(self, expr):
return self._map_elementwise_reduction("min", expr)
def map_elementwise_max(self, expr):
return self._map_elementwise_reduction("max", expr)
def map_ones(self, expr):
discr = self.places.get_discretization(
expr.dofdesc.geometry, expr.dofdesc.discr_stage)
result = discr.empty(actx=self.array_context, dtype=discr.real_dtype)
for grp_ary in result:
grp_ary.fill(1)
return result
def map_node_coordinate_component(self, expr):
discr = self.places.get_discretization(
expr.dofdesc.geometry, expr.dofdesc.discr_stage)
from arraycontext import thaw
x = discr.nodes()[expr.ambient_axis]
return thaw(x, self.array_context)
def map_num_reference_derivative(self, expr):
from pytools import flatten
ref_axes = flatten([axis] * mult for axis, mult in expr.ref_axes)
from meshmode.discretization import num_reference_derivative
discr = self.places.get_discretization(
expr.dofdesc.geometry, expr.dofdesc.discr_stage)
return num_reference_derivative(discr, ref_axes, self.rec(expr.operand))
def map_q_weight(self, expr):
from arraycontext import thaw
discr = self.places.get_discretization(
expr.dofdesc.geometry, expr.dofdesc.discr_stage)
return thaw(discr.quad_weights(), self.array_context)
def map_inverse(self, expr):
bound_op_cache = self.bound_expr.places._get_cache(
EvaluationMapperBoundOpCacheKey)
try:
bound_op = bound_op_cache[expr]
except KeyError:
bound_op = bind(
expr.expression,
self.places.get_geometry(expr.dofdesc.geometry),
self.bound_expr.iprec)
bound_op_cache[expr] = bound_op
scipy_op = bound_op.scipy_op(expr.variable_name, expr.dofdesc,
**{var_name: self.rec(var_expr)
for var_name, var_expr in expr.extra_vars.items()})
from pytential.solve import gmres
rhs = self.rec(expr.rhs)
result = gmres(scipy_op, rhs)
return result
def map_interpolation(self, expr):
operand = self.rec(expr.operand)
if isinstance(operand, (cl.array.Array, list, np.ndarray, DOFArray)):
conn = self.places.get_connection(expr.from_dd, expr.to_dd)
return conn(operand)
elif isinstance(operand, (int, float, complex, np.number)):
return operand
else:
raise TypeError("cannot interpolate `{}`".format(type(operand)))
def map_common_subexpression(self, expr):
if expr.scope == sym.cse_scope.EXPRESSION:
cache = self.bound_expr._get_cache(EvaluationMapperCSECacheKey)
elif expr.scope == sym.cse_scope.DISCRETIZATION:
cache = self.places._get_cache(EvaluationMapperCSECacheKey)
else:
return self.rec(expr.child)
from numbers import Number
try:
rec = cache[expr.child]
if (expr.scope == sym.cse_scope.DISCRETIZATION
and not isinstance(rec, Number)):
rec = thaw(rec, self.array_context)
except KeyError:
cached_rec = rec = self.rec(expr.child)
if (expr.scope == sym.cse_scope.DISCRETIZATION
and not isinstance(rec, Number)):
cached_rec = freeze(cached_rec, self.array_context)
cache[expr.child] = cached_rec
return rec
# }}}
def exec_assign(self, actx: PyOpenCLArrayContext, insn, bound_expr, evaluate):
return [(name, evaluate(expr))
for name, expr in zip(insn.names, insn.exprs)]
def exec_compute_potential_insn(
self, actx: PyOpenCLArrayContext, insn, bound_expr, evaluate):
raise NotImplementedError
# {{{ functions
def apply_real(self, args):
arg, = args
from pytools.obj_array import obj_array_real
return obj_array_real(self.rec(arg))
def apply_imag(self, args):
arg, = args
from pytools.obj_array import obj_array_imag
return obj_array_imag(self.rec(arg))
def apply_conj(self, args):
arg, = args
return self.rec(arg).conj()
def apply_abs(self, args):
arg, = args
return abs(self.rec(arg))
# }}}
def map_call(self, expr):
from pytential.symbolic.primitives import (
EvalMapperFunction, NumpyMathFunction)
if isinstance(expr.function, EvalMapperFunction):
return getattr(self, f"apply_{expr.function.name}")(expr.parameters)
elif isinstance(expr.function, NumpyMathFunction):
args = [self.rec(arg) for arg in expr.parameters]
from numbers import Number
if all(isinstance(arg, Number) for arg in args):
return getattr(np, expr.function.name)(*args)
else:
return getattr(self.array_context.np, expr.function.name)(*args)
else:
return super().map_call(expr)
# }}}
# {{{ evaluation mapper
class EvaluationMapper(EvaluationMapperBase):
def __init__(self, bound_expr, actx, context=None,
timing_data=None):
EvaluationMapperBase.__init__(self, bound_expr, actx, context)
self.timing_data = timing_data
def exec_compute_potential_insn(
self, actx: PyOpenCLArrayContext, insn, bound_expr, evaluate):
source = bound_expr.places.get_geometry(insn.source.geometry)
return_timing_data = self.timing_data is not None
result, timing_data = (
source.exec_compute_potential_insn(
actx, insn, bound_expr, evaluate, return_timing_data))
if return_timing_data:
# The compiler ensures this.
assert insn not in self.timing_data
self.timing_data[insn] = timing_data
return result
# }}}
# {{{ cost model evaluation mapper
class CostModelMapper(EvaluationMapperBase):
"""Mapper for evaluating cost models.
This executes everything *except* the layer potential operator. Instead of
executing the operator, the cost model gets run and the cost
data is collected.
.. attribute:: kernel_to_calibration_params
Can either be a :class:`str` "constant_one", which uses the constant 1.0 as
calibration parameters for all stages of all kernels, or be a :class:`dict`,
which maps from kernels to the calibration parameters, returned from
`estimate_kernel_specific_calibration_params`.
"""
def __init__(self, bound_expr, actx,
kernel_to_calibration_params, per_box,
context=None,
target_geometry=None,
target_points=None, target_normals=None, target_tangents=None):
if context is None:
context = {}
EvaluationMapperBase.__init__(
self, bound_expr, actx, context,
target_geometry,
target_points,
target_normals,
target_tangents)
self.kernel_to_calibration_params = kernel_to_calibration_params
self.modeled_cost = {}
self.metadata = {}
self.per_box = per_box
def exec_compute_potential_insn(
self, actx: PyOpenCLArrayContext, insn, bound_expr, evaluate):
source = bound_expr.places.get_geometry(insn.source.geometry)
knls = frozenset(knl for knl in insn.target_kernels)
if (isinstance(self.kernel_to_calibration_params, str)
and self.kernel_to_calibration_params == "constant_one"):
calibration_params = \
AbstractQBXCostModel.get_unit_calibration_params()
else:
calibration_params = self.kernel_to_calibration_params[knls]
result, (cost_model_result, metadata) = \
source.cost_model_compute_potential_insn(
actx, insn, bound_expr, evaluate, calibration_params,
self.per_box)
# The compiler ensures this.
assert insn not in self.modeled_cost
self.modeled_cost[insn] = cost_model_result
self.metadata[insn] = metadata
return result
def get_modeled_cost(self):
return self.modeled_cost, self.metadata
# }}}
# {{{ scipy-like mat-vec op
class MatVecOp:
"""A :class:`scipy.sparse.linalg.LinearOperator` work-alike.
Exposes a :mod:`pytential` operator as a generic matrix operation,
i.e., given :math:`x`, compute :math:`Ax`.
.. attribute:: shape
.. attribute:: dtype
.. automethod:: matvec
"""
def __init__(self,
bound_expr, actx: PyOpenCLArrayContext,
arg_name, dtype, total_dofs, discrs, starts_and_ends, extra_args):
self.bound_expr = bound_expr
self.array_context = actx
self.arg_name = arg_name
self.dtype = dtype
self.total_dofs = total_dofs
self.discrs = discrs
self.starts_and_ends = starts_and_ends
self.extra_args = extra_args
@property
def shape(self):
return (self.total_dofs, self.total_dofs)
@property
def _operator_uses_obj_array(self):
return len(self.discrs) > 1
def flatten(self, ary):
# Return a flat version of *ary*. The returned value is suitable for
# use with solvers whose API expects a one-dimensional array.
if not self._operator_uses_obj_array:
ary = [ary]
result = self.array_context.empty(self.total_dofs, self.dtype)
for res_i, (start, end) in zip(ary, self.starts_and_ends):
result[start:end] = flatten(thaw(res_i, self.array_context))
return result
def unflatten(self, ary):
# Convert a flat version of *ary* into a structured version.
components = []
for discr, (start, end) in zip(self.discrs, self.starts_and_ends):
component = ary[start:end]
from meshmode.discretization import Discretization
if isinstance(discr, Discretization):
from meshmode.dof_array import unflatten
component = unflatten(self.array_context, discr, component)
components.append(component)
if self._operator_uses_obj_array:
from pytools.obj_array import make_obj_array
return make_obj_array(components)
else:
return components[0]
def matvec(self, x):
# Three types of inputs are supported:
# * flat NumPy arrays
# => output is a flat NumPy array
# * flat PyOpenCL arrays
# => output is a flat PyOpenCL array
# * structured arrays (object arrays/DOFArrays)
# => output has same structure as input
if isinstance(x, np.ndarray) and x.dtype.char != "O":
x = self.array_context.from_numpy(x)
flat = True
host = True
assert x.shape == (self.total_dofs,)
elif isinstance(x, cl.array.Array):
flat = True
host = False
assert x.shape == (self.total_dofs,)
elif isinstance(x, np.ndarray) and x.dtype.char == "O":
flat = False
host = False
elif isinstance(x, DOFArray):
flat = False
host = False
else:
raise ValueError("unsupported input type")
args = self.extra_args.copy()
args[self.arg_name] = self.unflatten(x) if flat else x
result = self.bound_expr(self.array_context, **args)
if flat:
result = self.flatten(result)
if host:
result = self.array_context.to_numpy(result)
return result
# }}}
# {{{ expression prep
def _prepare_domains(nresults, places, domains, default_domain):
"""
:arg nresults: number of results.
:arg places: a :class:`~pytential.GeometryCollection`.
:arg domains: recommended domains.
:arg default_domain: default value for domains which are not provided.
:return: a list of domains for each result. If domains is `None`, each
element in the list is *default_domain*. If *domains* is a scalar
(i.e., not a *list* or *tuple*), each element in the list is
*domains*. Otherwise, *domains* is returned as is.
"""
if domains is None:
dom_name = default_domain
return nresults * [dom_name]
elif not isinstance(domains, (list, tuple)):
dom_name = domains
return nresults * [dom_name]
domains = [sym.as_dofdesc(d) for d in domains]
assert len(domains) == nresults
return domains
def _prepare_auto_where(auto_where, places=None):
"""
:arg auto_where: a 2-tuple, single identifier or `None` used as a hint
to determine the default geometries.
:arg places: a :class:`GeometryCollection`,
whose :attr:`GeometryCollection.auto_where` is used by default if
provided and `auto_where` is `None`.
:return: a tuple ``(source, target)`` of
:class:`~pytential.symbolic.primitives.DOFDescriptor`s denoting
the default source and target geometries.
"""
if auto_where is None:
if places is None:
auto_source = sym.DEFAULT_SOURCE
auto_target = sym.DEFAULT_TARGET
else:
auto_source, auto_target = places.auto_where
elif isinstance(auto_where, (list, tuple)):
auto_source, auto_target = auto_where
else:
auto_source = auto_where
auto_target = auto_source
return (sym.as_dofdesc(auto_source), sym.as_dofdesc(auto_target))
def _prepare_expr(places, expr, auto_where=None):
"""
:arg places: :class:`~pytential.GeometryCollection`.
:arg expr: a symbolic expression.
:return: processed symbolic expressions, tagged with the appropriate
`where` identifier from places, etc.
"""
from pytential.source import LayerPotentialSourceBase
from pytential.symbolic.mappers import (
ToTargetTagger,
DerivativeBinder)
auto_source, auto_target = _prepare_auto_where(auto_where, places=places)
expr = ToTargetTagger(auto_source, auto_target)(expr)
expr = DerivativeBinder()(expr)
for name, place in places.places.items():
if isinstance(place, LayerPotentialSourceBase):
expr = place.preprocess_optemplate(name, places, expr)
from pytential.symbolic.mappers import InterpolationPreprocessor
expr = InterpolationPreprocessor(places)(expr)
return expr
# }}}
# {{{ geometry collection
def _is_valid_identifier(name):
import keyword
return name.isidentifier() and not keyword.iskeyword(name)
class _GeometryCollectionDiscretizationCacheKey:
"""Serves as a unique key for the discretization cache in
:meth:`GeometryCollection._get_cache`.
"""
class _GeometryCollectionConnectionCacheKey:
"""Serves as a unique key for the connection cache in
:meth:`GeometryCollection._get_cache`.
"""
[docs]class GeometryCollection:
"""A mapping from symbolic identifiers ("place IDs", typically strings)
to 'geometries', where a geometry can be a
:class:`pytential.source.PotentialSource`
or a :class:`pytential.target.TargetBase`.
This class is meant to hold a specific combination of sources and targets
serve to host caches of information derived from them, e.g. FMM trees
of subsets of them, as well as related common subexpressions such as
metric terms.
.. automethod:: get_geometry
.. automethod:: get_connection
.. automethod:: get_discretization
.. automethod:: copy
.. automethod:: merge
Refinement of :class:`pytential.qbx.QBXLayerPotentialSource` entries is
performed on demand, or it may be performed by explcitly calling
:func:`pytential.qbx.refinement.refine_geometry_collection`,
which allows more customization of the refinement process through
parameters.
"""
def __init__(self, places, auto_where=None):
"""
:arg places: a scalar, tuple of or mapping of symbolic names to
geometry objects. Supported objects are
:class:`~pytential.source.PotentialSource`,
:class:`~pytential.target.TargetBase` and
:class:`~meshmode.discretization.Discretization`. If this is
a mapping, the keys that are strings must be valid Python identifiers.
:arg auto_where: location identifier for each geometry object, used
to denote specific discretizations, e.g. in the case where
*places* is a :class:`~pytential.source.LayerPotentialSourceBase`.
By default, we assume
:class:`~pytential.symbolic.primitives.DEFAULT_SOURCE` and
:class:`~pytential.symbolic.primitives.DEFAULT_TARGET` for
sources and targets, respectively.
"""
from pytential.target import TargetBase
from pytential.source import PotentialSource
from pytential.qbx import QBXLayerPotentialSource
from meshmode.discretization import Discretization
# {{{ construct dict
self.places = {}
self.caches = {}
auto_source, auto_target = _prepare_auto_where(auto_where)
if isinstance(places, QBXLayerPotentialSource):
self.places[auto_source.geometry] = places
auto_target = auto_source
elif isinstance(places, TargetBase):
self.places[auto_target.geometry] = places
auto_source = auto_target
if isinstance(places, (Discretization, PotentialSource)):
self.places[auto_source.geometry] = places
self.places[auto_target.geometry] = places
elif isinstance(places, tuple):
source_discr, target_discr = places
self.places[auto_source.geometry] = source_discr
self.places[auto_target.geometry] = target_discr
else:
self.places = places
self.auto_where = (auto_source, auto_target)
# }}}
# {{{ validate
# check allowed identifiers
for name in self.places:
if not isinstance(name, str):
continue
if not _is_valid_identifier(name):
raise ValueError(f"'{name}' is not a valid identifier")
# check allowed types
for p in self.places.values():
if not isinstance(p, (PotentialSource, TargetBase, Discretization)):
raise TypeError("Values in 'places' must be discretization, targets "
"or layer potential sources.")
# check ambient_dim
from pytools import is_single_valued
ambient_dims = [p.ambient_dim for p in self.places.values()]
if not is_single_valued(ambient_dims):
raise RuntimeError("All 'places' must have the same ambient dimension.")
self.ambient_dim = ambient_dims[0]
# }}}
@property
def auto_source(self):
return self.auto_where[0]
@property
def auto_target(self):
return self.auto_where[1]
# {{{ cache handling
def _get_cache(self, name):
return self.caches.setdefault(name, {})
def _get_discr_from_cache(self, geometry, discr_stage):
cache = self._get_cache(_GeometryCollectionDiscretizationCacheKey)
key = (geometry, discr_stage)
if key not in cache:
raise KeyError("cached discretization does not exist on '{}'"
"for stage '{}'".format(geometry, discr_stage))
return cache[key]
def _add_discr_to_cache(self, discr, geometry, discr_stage):
cache = self._get_cache(_GeometryCollectionDiscretizationCacheKey)
key = (geometry, discr_stage)
if key in cache:
raise RuntimeError("trying to overwrite the cache")
cache[key] = discr
def _get_conn_from_cache(self, geometry, from_stage, to_stage):
cache = self._get_cache(_GeometryCollectionConnectionCacheKey)
key = (geometry, from_stage, to_stage)
if key not in cache:
raise KeyError("cached connection does not exist on '{}' "
"from '{}' to '{}'".format(geometry, from_stage, to_stage))
return cache[key]
def _add_conn_to_cache(self, conn, geometry, from_stage, to_stage):
cache = self._get_cache(_GeometryCollectionConnectionCacheKey)
key = (geometry, from_stage, to_stage)
if key in cache:
raise RuntimeError("trying to overwrite the cache")
cache[key] = conn
def _get_qbx_discretization(self, geometry, discr_stage):
lpot_source = self.get_geometry(geometry)
try:
discr = self._get_discr_from_cache(geometry, discr_stage)
except KeyError:
from pytential import sym
from pytential.qbx.refinement import _refine_for_global_qbx
# NOTE: this adds the required discretizations to the cache
dofdesc = sym.DOFDescriptor(geometry, discr_stage)
_refine_for_global_qbx(self, dofdesc,
lpot_source.refiner_code_container.get_wrangler(),
_copy_collection=False)
discr = self._get_discr_from_cache(geometry, discr_stage)
return discr
# }}}
[docs] def get_connection(self, from_dd, to_dd):
from pytential.symbolic.dof_connection import connection_from_dds
return connection_from_dds(self, from_dd, to_dd)
[docs] def get_discretization(self, geometry, discr_stage=None):
"""
:arg dofdesc: a :class:`~pytential.symbolic.primitives.DOFDescriptor`
specifying the desired discretization.
:return: a geometry object in the collection corresponding to the
key *dofdesc*. If it is a
:class:`~pytential.source.LayerPotentialSourceBase`, we look for
the corresponding :class:`~meshmode.discretization.Discretization`
in its attributes instead.
"""
if discr_stage is None:
discr_stage = sym.QBX_SOURCE_STAGE1
discr = self.get_geometry(geometry)
from pytential.qbx import QBXLayerPotentialSource
from pytential.source import LayerPotentialSourceBase
if isinstance(discr, QBXLayerPotentialSource):
return self._get_qbx_discretization(geometry, discr_stage)
elif isinstance(discr, LayerPotentialSourceBase):
return discr.density_discr
else:
return discr
[docs] def get_geometry(self, geometry):
try:
return self.places[geometry]
except KeyError:
raise KeyError("geometry not in the collection: '{}'".format(
geometry))
[docs] def copy(self, places=None, auto_where=None):
places = self.places if places is None else places
return type(self)(
places=places.copy(),
auto_where=self.auto_where if auto_where is None else auto_where)
[docs] def merge(self, places):
"""Merges two geometry collections and returns the new collection.
:arg places: A :class:`dict` or :class:`GeometryCollection` to
merge with the current collection. If it is empty, a copy of the
current collection is returned.
"""
new_places = self.places.copy()
if places:
if isinstance(places, GeometryCollection):
places = places.places
new_places.update(places)
return self.copy(places=new_places)
def __repr__(self):
return "{}({})".format(type(self).__name__, repr(self.places))
def __str__(self):
return "{}({})".format(type(self).__name__, str(self.places))
# }}}
# {{{ bound expression
def _find_array_context_from_args_in_context(context, supplied_array_context=None):
from arraycontext import PyOpenCLArrayContext
array_contexts = []
if supplied_array_context is not None:
if not isinstance(supplied_array_context, PyOpenCLArrayContext):
raise TypeError(
"first argument (if supplied) must be a PyOpenCLArrayContext, "
f"got '{type(supplied_array_context).__name__}'")
array_contexts.append(supplied_array_context)
del supplied_array_context
def look_for_array_contexts(ary):
if isinstance(ary, DOFArray):
if ary.array_context is not None:
array_contexts.append(ary.array_context)
elif isinstance(ary, np.ndarray) and ary.dtype.char == "O":
for idx in np.ndindex(ary.shape):
look_for_array_contexts(ary[idx])
else:
pass
for val in context.values():
look_for_array_contexts(val)
if array_contexts:
from pytools import is_single_valued
if not is_single_valued(array_contexts):
raise ValueError("arguments do not agree on an array context")
array_context = array_contexts[0]
else:
array_context = None
if not isinstance(array_context, PyOpenCLArrayContext):
raise TypeError(
"array context (derived from arguments) is not a "
f"PyOpenCLArrayContext: '{type(array_context).__name__}'")
return array_context
[docs]class BoundExpression:
"""An expression readied for evaluation by binding it to a
:class:`~pytential.GeometryCollection`.
.. automethod :: cost_per_stage
.. automethod :: cost_per_box
.. automethod :: scipy_op
.. automethod :: eval
.. automethod :: __call__
.. attribute :: places
Created by calling :func:`pytential.bind`.
"""
def __init__(self, places, sym_op_expr):
self.places = places
self.sym_op_expr = sym_op_expr
self.caches = {}
@property
@memoize_method
def code(self):
from pytential.symbolic.compiler import OperatorCompiler
return OperatorCompiler(self.places)(self.sym_op_expr)
def _get_cache(self, name):
return self.caches.setdefault(name, {})
[docs] def cost_per_stage(self, calibration_params, **kwargs):
"""
:arg queue: a :class:`pyopencl.CommandQueue` object.
:arg calibration_params: either a :class:`dict` returned by
`estimate_kernel_specific_calibration_params`, or a :class:`str`
"constant_one".
:return: a :class:`dict` mapping from instruction to per-stage cost. Each
per-stage cost is represented by a :class:`dict` mapping from the stage
name to the predicted time.
"""
array_context = _find_array_context_from_args_in_context(kwargs)
if array_context is None:
raise ValueError("unable to figure array context from arguments")
cost_model_mapper = CostModelMapper(
self, array_context, calibration_params, per_box=False, context=kwargs
)
self.code.execute(cost_model_mapper)
return cost_model_mapper.get_modeled_cost()
[docs] def cost_per_box(self, calibration_params, **kwargs):
"""
:arg queue: a :class:`pyopencl.CommandQueue` object.
:arg calibration_params: either a :class:`dict` returned by
`estimate_kernel_specific_calibration_params`, or a :class:`str`
"constant_one".
:return: a :class:`dict` mapping from instruction to per-box cost. Each
per-box cost is represented by a :class:`numpy.ndarray` or
:class:`pyopencl.array.Array` of shape (nboxes,), where the ith entry
represents the cost of all stages for box i.
"""
array_context = _find_array_context_from_args_in_context(kwargs)
cost_model_mapper = CostModelMapper(
self, array_context, calibration_params, per_box=True, context=kwargs
)
self.code.execute(cost_model_mapper)
return cost_model_mapper.get_modeled_cost()
[docs] def scipy_op(
self, actx: PyOpenCLArrayContext, arg_name, dtype,
domains=None, **extra_args):
"""
:arg domains: a list of discretization identifiers or
*None* values indicating the domains on which each component of the
solution vector lives. *None* values indicate that the component
is a scalar. If *domains* is *None*,
:class:`~pytential.symbolic.primitives.DEFAULT_TARGET` is required
to be a key in :attr:`places`.
:returns: An object that (mostly) satisfies the
:class:`scipy.sparse.linalg.LinearOperator` protocol, except for
accepting and returning :class:`pyopencl.array.Array` arrays.
"""
if isinstance(self.code.result, np.ndarray):
nresults = len(self.code.result)
else:
nresults = 1
domains = _prepare_domains(nresults,
self.places, domains, self.places.auto_target)
total_dofs = 0
discrs = []
starts_and_ends = []
for dom_name in domains:
if dom_name is None:
discr = None
size = 1
else:
discr = self.places.get_discretization(
dom_name.geometry, dom_name.discr_stage)
size = discr.ndofs
discrs.append(discr)
starts_and_ends.append((total_dofs, total_dofs+size))
total_dofs += size
# Hidden assumption: Number of input components
# equals number of output compoments. But IMO that's
# fair, since these operators are usually only used
# for linear system solving, in which case the assumption
# has to be true.
return MatVecOp(self, actx,
arg_name, dtype, total_dofs, discrs, starts_and_ends, extra_args)
[docs] def eval(self, context=None, timing_data=None,
array_context: Optional[PyOpenCLArrayContext] = None):
"""Evaluate the expression in *self*, using the
:class:`pyopencl.CommandQueue` *queue* and the
input variables given in the dictionary *context*.
:arg timing_data: A dictionary into which timing
data will be inserted during evaluation.
(experimental)
:arg array_context: only needs to be supplied if no instances of
:class:`~meshmode.dof_array.DOFArray` with a
:class:`~arraycontext.PyOpenCLArrayContext`
are supplied as part of *context*.
:returns: the value of the expression, as a scalar,
:class:`pyopencl.array.Array`, or an object array of these.
"""
if context is None:
context = {}
# NOTE: avoid compiling any code if the expression is long lived
# and already nicely cached in the collection from a previous run
import pymbolic.primitives as prim
if isinstance(self.sym_op_expr, prim.CommonSubexpression) \
and self.sym_op_expr.scope == sym.cse_scope.DISCRETIZATION:
cache = self.places._get_cache(EvaluationMapperCSECacheKey)
from numbers import Number
expr = self.sym_op_expr
if expr.child in cache:
value = cache[expr.child]
if (expr.scope == sym.cse_scope.DISCRETIZATION
and not isinstance(value, Number)):
value = thaw(value, array_context)
return value
array_context = _find_array_context_from_args_in_context(
context, array_context)
exec_mapper = EvaluationMapper(
self, array_context, context, timing_data=timing_data)
return self.code.execute(exec_mapper)
[docs] def __call__(self, *args, **kwargs):
"""Evaluate the expression in *self*, using the
:class:`pyopencl.CommandQueue` *queue* and the
input variables given in the dictionary *context*.
:returns: the value of the expression, as a scalar,
:class:`meshmode.dof_array.DOFArray`, or an object array of
these.
"""
array_context = None
if len(args) == 1:
array_context, = args
if not isinstance(array_context, PyOpenCLArrayContext):
raise TypeError(
"first positional argument (if given) must be a "
f"PyOpenCLArrayContext: '{type(array_context).__name__}'")
elif not args:
pass
else:
raise TypeError("More than one positional argument supplied. "
"None or an PyOpenCLArrayContext expected.")
return self.eval(kwargs, array_context=array_context)
[docs]def bind(places, expr, auto_where=None):
"""
:arg places: a :class:`pytential.GeometryCollection`.
Alternatively, any list or mapping that is a valid argument for its
constructor can also be used.
:arg auto_where: for simple source-to-self or source-to-target
evaluations, find 'where' attributes automatically.
:arg expr: one or multiple expressions consisting of primitives
form :mod:`pytential.symbolic.primitives` (aka ``pytential.sym``).
Multiple expressions can be combined into one object to pass here
in the form of a :mod:`numpy` object array
:returns: a :class:`pytential.symbolic.execution.BoundExpression`
"""
if not isinstance(places, GeometryCollection):
places = GeometryCollection(places, auto_where=auto_where)
auto_where = places.auto_where
expr = _prepare_expr(places, expr, auto_where=auto_where)
return BoundExpression(places, expr)
# }}}
# {{{ matrix building
def _bmat(blocks, dtypes):
from pytools import single_valued
from pytential.symbolic.matrix import is_zero
nrows = blocks.shape[0]
ncolumns = blocks.shape[1]
# "block row starts"/"block column starts"
brs = np.cumsum([0]
+ [single_valued(blocks[ibrow, ibcol].shape[0]
for ibcol in range(ncolumns)
if not is_zero(blocks[ibrow, ibcol]))
for ibrow in range(nrows)])
bcs = np.cumsum([0]
+ [single_valued(blocks[ibrow, ibcol].shape[1]
for ibrow in range(nrows)
if not is_zero(blocks[ibrow, ibcol]))
for ibcol in range(ncolumns)])
result = np.zeros((brs[-1], bcs[-1]),
dtype=np.find_common_type(dtypes, []))
for ibcol in range(ncolumns):
for ibrow in range(nrows):
result[brs[ibrow]:brs[ibrow + 1], bcs[ibcol]:bcs[ibcol + 1]] = \
blocks[ibrow, ibcol]
return result
def build_matrix(actx, places, exprs, input_exprs, domains=None,
auto_where=None, context=None):
"""
:arg actx: a :class:`~arraycontext.PyOpenCLArrayContext`.
:arg places: a :class:`pytential.GeometryCollection`.
Alternatively, any list or mapping that is a valid argument for its
constructor can also be used.
:arg exprs: an array of expressions corresponding to the output block
rows of the matrix. May also be a single expression.
:arg input_exprs: an array of expressions corresponding to the
input block columns of the matrix. May also be a single expression.
:arg domains: a list of discretization identifiers (see 'places') or
*None* values indicating the domains on which each component of the
solution vector lives. *None* values indicate that the component
is a scalar. If *None*, *auto_where* or, if it is not provided,
:class:`~pytential.symbolic.primitives.DEFAULT_SOURCE` is required
to be a key in :attr:`places`.
:arg auto_where: For simple source-to-self or source-to-target
evaluations, find 'where' attributes automatically.
"""
if context is None:
context = {}
from pytential import GeometryCollection
if not isinstance(places, GeometryCollection):
places = GeometryCollection(places, auto_where=auto_where)
exprs = _prepare_expr(places, exprs, auto_where=auto_where)
if not (isinstance(exprs, np.ndarray) and exprs.dtype.char == "O"):
from pytools.obj_array import make_obj_array
exprs = make_obj_array([exprs])
try:
input_exprs = list(input_exprs)
except TypeError:
# not iterable, wrap in a list
input_exprs = [input_exprs]
domains = _prepare_domains(len(input_exprs),
places, domains, places.auto_source)
from pytential.symbolic.matrix import MatrixBuilder, is_zero
nblock_rows = len(exprs)
nblock_columns = len(input_exprs)
blocks = np.zeros((nblock_rows, nblock_columns), dtype=object)
dtypes = []
for ibcol in range(nblock_columns):
dep_source = places.get_geometry(domains[ibcol].geometry)
dep_discr = places.get_discretization(
domains[ibcol].geometry, domains[ibcol].discr_stage)
mbuilder = MatrixBuilder(
actx,
dep_expr=input_exprs[ibcol],
other_dep_exprs=(input_exprs[:ibcol]
+ input_exprs[ibcol + 1:]),
dep_source=dep_source,
dep_discr=dep_discr,
places=places,
context=context)
for ibrow in range(nblock_rows):
block = mbuilder(exprs[ibrow])
assert is_zero(block) or isinstance(block, np.ndarray)
blocks[ibrow, ibcol] = block
if isinstance(block, np.ndarray):
dtypes.append(block.dtype)
return actx.from_numpy(_bmat(blocks, dtypes))
# }}}
# vim: foldmethod=marker