def test_step_matrix_vector_state(show_matrix=True, show_dag=False):
from leap.step_matrix import StepMatrixFinder
from pymbolic import var
component_id = 'y'
code = euler(component_id, show_dag)
J = np.diag([-3, -2, -1]) # noqa
def rhs_sym(t, y):
return J.dot(y)
finder = StepMatrixFinder(code,
function_map={"<func>" + component_id: rhs_sym},
variables=["<state>" + component_id])
mat = finder.get_phase_step_matrix("primary",
shapes={"<state>" + component_id: 3})
if show_matrix:
print('Variables: %s' % finder.variables)
from pytools import indices_in_shape
for i in indices_in_shape(mat.shape):
print(i, mat[i])
# XXX: brittle
dt = var("<dt>")
true_mat = np.eye(3, dtype=np.object) + dt * J
assert (mat == true_mat).all()
def map_numpy_array(self, expr):
import numpy
result = numpy.empty(expr.shape, dtype=object)
from pytools import indices_in_shape
for i in indices_in_shape(expr.shape):
result[i] = self.rec(expr[i])
return result
def with_object_array_or_scalar_n_args(f, *args):
oarray_arg_indices = []
for i, arg in enumerate(args):
if is_obj_array(arg):
oarray_arg_indices.append(i)
if not oarray_arg_indices:
return f(*args)
leading_oa_index = oarray_arg_indices[0]
ls = log_shape(args[leading_oa_index])
if ls != ():
from pytools import indices_in_shape
result = np.zeros(ls, dtype=object)
new_args = list(args)
for i in indices_in_shape(ls):
for arg_i in oarray_arg_indices:
new_args[arg_i] = args[arg_i][i]
result[i] = f(*new_args)
return result
else:
return f(*args)
def map_numpy_array(self, expr):
if not self.visit(expr):
return
from pytools import indices_in_shape
for i in indices_in_shape(expr.shape):
self.rec(expr[i])
def make_common_subexpression(field, prefix=None):
from pytools.obj_array import log_shape
from hedge.tools import is_zero
from pymbolic.primitives import CommonSubexpression
ls = log_shape(field)
if ls != ():
from pytools import indices_in_shape
result = numpy.zeros(ls, dtype=object)
for i in indices_in_shape(ls):
if prefix is not None:
component_prefix = prefix+"_".join(str(i_i) for i_i in i)
else:
component_prefix = None
if is_zero(field[i]):
result[i] = 0
else:
result[i] = CommonSubexpression(field[i], component_prefix)
return result
else:
if is_zero(field):
return 0
else:
return CommonSubexpression(field, prefix)
def with_object_array_or_scalar_n_args(f, *args):
oarray_arg_indices = []
for i, arg in enumerate(args):
if is_obj_array(arg):
oarray_arg_indices.append(i)
if not oarray_arg_indices:
return f(*args)
leading_oa_index = oarray_arg_indices[0]
ls = log_shape(args[leading_oa_index])
if ls != ():
from pytools import indices_in_shape
result = np.zeros(ls, dtype=object)
new_args = list(args)
for i in indices_in_shape(ls):
for arg_i in oarray_arg_indices:
new_args[arg_i] = args[arg_i][i]
result[i] = f(*new_args)
return result
else:
return f(*args)
def map_numpy_array(self, expr, *args, **kwargs):
import numpy
result = numpy.empty(expr.shape, dtype=object)
from pytools import indices_in_shape
for i in indices_in_shape(expr.shape):
result[i] = self.rec(expr[i], *args, **kwargs)
return result
def __call__(self, a):
from pytools import indices_in_shape
# assumes nonempty, which is reasonable
return max(
abs(self.scalar_kernel(a[i]))
for i in indices_in_shape(a.shape))
def map_numpy_array(self, expr):
if not self.visit(expr):
return
from pytools import indices_in_shape
for i in indices_in_shape(expr.shape):
self.rec(expr[i])
def __call__(self, a):
from pytools import indices_in_shape
# assumes nonempty, which is reasonable
return max(
abs(self.scalar_kernel(a[i]))
for i in indices_in_shape(a.shape))
def make_common_subexpression(field, prefix=None):
try:
from pytools.obj_array import log_shape
except ImportError:
have_obj_array = False
else:
have_obj_array = True
if have_obj_array:
ls = log_shape(field)
if have_obj_array and ls != ():
from pytools import indices_in_shape
result = numpy.zeros(ls, dtype=object)
for i in indices_in_shape(ls):
if prefix is not None:
component_prefix = prefix+"_".join(str(i_i) for i_i in i)
else:
component_prefix = None
if is_constant(field[i]):
result[i] = field[i]
else:
result[i] = CommonSubexpression(field[i], component_prefix)
return result
else:
if is_constant(field):
return field
else:
return CommonSubexpression(field, prefix)
def to_obj_array(ary):
ls = log_shape(ary)
result = np.empty(ls, dtype=object)
from pytools import indices_in_shape
for i in indices_in_shape(ls):
result[i] = ary[i]
return result
def apply_mask(field):
from hedge.tools import log_shape
ls = log_shape(field)
result = discr.volume_empty(ls)
from pytools import indices_in_shape
for i in indices_in_shape(ls):
result[i] = mask * field[i]
return result
def map_numpy_array(self, expr, *args, **kwargs):
if not self.visit(expr, *args, **kwargs):
return
from pytools import indices_in_shape
for i in indices_in_shape(expr.shape):
self.rec(expr[i], *args, **kwargs)
self.post_visit(expr, *args, **kwargs)
def to_obj_array(ary):
ls = log_shape(ary)
result = np.empty(ls, dtype=object)
from pytools import indices_in_shape
for i in indices_in_shape(ls):
result[i] = ary[i]
return result
def grad_S(kernel, arg, dim):
from pytools.obj_array import log_shape
arg_shape = log_shape(arg)
result = np.zeros(arg_shape + (dim, ), dtype=object)
from pytools import indices_in_shape
for i in indices_in_shape(arg_shape):
for j in range(dim):
result[i + (j, )] = IntGdTarget(kernel, arg[i], j)
return result
def map_numpy_array(self, expr, *args, **kwargs):
if not self.visit(expr, *args, **kwargs):
return
from pytools import indices_in_shape
for i in indices_in_shape(expr.shape):
self.rec(expr[i], *args, **kwargs)
self.post_visit(expr, *args, **kwargs)
def grad_D(kernel, arg, dim):
from pytools.obj_array import log_shape
arg_shape = log_shape(arg)
result = np.zeros(arg_shape+(dim,), dtype=object)
from pytools import indices_in_shape
for i in indices_in_shape(arg_shape):
for j in range(dim):
result[i+(j,)] = IntGdMixed(kernel, arg[i], j)
return result
def apply_mask(field):
from grudge.tools import log_shape
ls = log_shape(field)
result = discr.volume_empty(ls)
from pytools import indices_in_shape
for i in indices_in_shape(ls):
result[i] = mask * field[i]
return result
def ptwise_mul(a, b):
from pytools.obj_array import log_shape
a_log_shape = log_shape(a)
b_log_shape = log_shape(b)
from pytools import indices_in_shape
if a_log_shape == ():
result = np.empty(b_log_shape, dtype=object)
for i in indices_in_shape(b_log_shape):
result[i] = a * b[i]
elif b_log_shape == ():
result = np.empty(a_log_shape, dtype=object)
for i in indices_in_shape(a_log_shape):
result[i] = a[i] * b
else:
raise ValueError("ptwise_mul can't handle two non-scalars")
return result
def __call__(self, a, b):
from pytools import indices_in_shape
assert a.shape == b.shape
result = 0
for i in indices_in_shape(a.shape):
result += self.scalar_kernel(a[i], b[i])
return result
def __call__(self, a, b):
from pytools import indices_in_shape
assert a.shape == b.shape
result = 0
for i in indices_in_shape(a.shape):
result += self.scalar_kernel(a[i], b[i])
return result
def ptwise_dot(logdims1, logdims2, a1, a2):
a1_log_shape = a1.shape[:logdims1]
a2_log_shape = a1.shape[:logdims2]
assert a1_log_shape[-1] == a2_log_shape[0]
len_k = a2_log_shape[0]
result = np.empty(a1_log_shape[:-1] + a2_log_shape[1:], dtype=object)
from pytools import indices_in_shape
for a1_i in indices_in_shape(a1_log_shape[:-1]):
for a2_i in indices_in_shape(a2_log_shape[1:]):
result[a1_i + a2_i] = sum(a1[a1_i + (k, )] * a2[(k, ) + a2_i]
for k in range(len_k))
if result.shape == ():
return result[()]
else:
return result
def ptwise_mul(a, b):
from pytools.obj_array import log_shape
a_log_shape = log_shape(a)
b_log_shape = log_shape(b)
from pytools import indices_in_shape
if a_log_shape == ():
result = np.empty(b_log_shape, dtype=object)
for i in indices_in_shape(b_log_shape):
result[i] = a*b[i]
elif b_log_shape == ():
result = np.empty(a_log_shape, dtype=object)
for i in indices_in_shape(a_log_shape):
result[i] = a[i]*b
else:
raise ValueError("ptwise_mul can't handle two non-scalars")
return result
def __call__(self, *args):
from pytools import indices_in_shape, single_valued
oa_shape = single_valued(ary.shape for fac, ary in args)
result = numpy.zeros(oa_shape, dtype=object)
for i in indices_in_shape(oa_shape):
args_i = [(fac, ary[i]) for fac, ary in args]
result[i] = self.scalar_kernel(*args_i)
return result
def __call__(self, *args):
from pytools import indices_in_shape, single_valued
oa_shape = single_valued(ary.shape for fac, ary in args)
result = numpy.zeros(oa_shape, dtype=object)
for i in indices_in_shape(oa_shape):
args_i = [(fac, ary[i]) for fac, ary in args]
result[i] = self.scalar_kernel(*args_i)
return result
def make_sym_array(name, shape):
vfld = Variable(name)
if shape == ():
return vfld
import numpy as np
result = np.zeros(shape, dtype=object)
from pytools import indices_in_shape
for i in indices_in_shape(shape):
result[i] = vfld[i]
return result
def make_sym_array(name, shape, var_factory=Variable):
vfld = var_factory(name)
if shape == ():
return vfld
import numpy as np
result = np.zeros(shape, dtype=object)
from pytools import indices_in_shape
for i in indices_in_shape(shape):
result[i] = vfld.index(i)
return result
def ptwise_dot(logdims1, logdims2, a1, a2):
a1_log_shape = a1.shape[:logdims1]
a2_log_shape = a1.shape[:logdims2]
assert a1_log_shape[-1] == a2_log_shape[0]
len_k = a2_log_shape[0]
result = np.empty(a1_log_shape[:-1]+a2_log_shape[1:], dtype=object)
from pytools import indices_in_shape
for a1_i in indices_in_shape(a1_log_shape[:-1]):
for a2_i in indices_in_shape(a2_log_shape[1:]):
result[a1_i+a2_i] = sum(
a1[a1_i+(k,)] * a2[(k,)+a2_i]
for k in xrange(len_k)
)
if result.shape == ():
return result[()]
else:
return result
def map_numpy_array(self, expr, enclosing_prec, *args, **kwargs):
import numpy
from pytools import indices_in_shape
str_array = numpy.zeros(expr.shape, dtype="object")
max_length = 0
for i in indices_in_shape(expr.shape):
s = self.rec(expr[i], PREC_NONE, *args, **kwargs)
max_length = max(len(s), max_length)
str_array[i] = s.replace("\n", "\n ")
if len(expr.shape) == 1 and max_length < 15:
return "array(%s)" % ", ".join(str_array)
else:
lines = [" %s: %s\n" % (
",".join(str(i_i) for i_i in i), str_array[i])
for i in indices_in_shape(expr.shape)]
if max_length > 70:
splitter = " " + "-"*75 + "\n"
return "array(\n%s)" % splitter.join(lines)
else:
return "array(\n%s)" % "".join(lines)
def map_numpy_array(self, expr, enclosing_prec, *args, **kwargs):
import numpy
from pytools import indices_in_shape
str_array = numpy.zeros(expr.shape, dtype="object")
max_length = 0
for i in indices_in_shape(expr.shape):
s = self.rec(expr[i], PREC_NONE, *args, **kwargs)
max_length = max(len(s), max_length)
str_array[i] = s.replace("\n", "\n ")
if len(expr.shape) == 1 and max_length < 15:
return "array(%s)" % ", ".join(str_array)
else:
lines = [" %s: %s\n" % (
",".join(str(i_i) for i_i in i), str_array[i])
for i in indices_in_shape(expr.shape)]
if max_length > 70:
splitter = " " + "-"*75 + "\n"
return "array(\n%s)" % splitter.join(lines)
else:
return "array(\n%s)" % "".join(lines)
def __call__(self, discr, t, fields, x, make_empty):
result = self.make_func(discr)(t=numpy.float64(t), x=x, fields=fields)
# make sure we return no scalars in the result
from pytools.obj_array import log_shape, is_obj_array
if is_obj_array(result):
from pytools import indices_in_shape
from hedge.optemplate.tools import is_scalar
for i in indices_in_shape(log_shape(result)):
if is_scalar(result[i]):
result[i] = make_empty().fill(result[i])
return result
def __call__(self, discr, t, fields, x, make_empty):
result = self.make_func(discr)(
t=numpy.float64(t), x=x, fields=fields)
# make sure we return no scalars in the result
from pytools.obj_array import log_shape, is_obj_array
if is_obj_array(result):
from pytools import indices_in_shape
from hedge.optemplate.tools import is_scalar
for i in indices_in_shape(log_shape(result)):
if is_scalar(result[i]):
result[i] = make_empty().fill(result[i])
return result
def make_sym_array(name, shape, var_factory=None):
if var_factory is None:
var_factory = Variable
vfld = var_factory(name)
if shape == ():
return vfld
import numpy as np
result = np.zeros(shape, dtype=object)
from pytools import indices_in_shape
for i in indices_in_shape(shape):
result[i] = vfld.index(i)
return result
def count_dofs(vec):
try:
dtype = vec.dtype
size = vec.size
shape = vec.shape
except AttributeError:
from warnings import warn
warn("could not count dofs of vector")
return 0
if dtype == object:
from pytools import indices_in_shape
return sum(count_dofs(vec[i]) for i in indices_in_shape(vec.shape))
else:
return size
def split(self, whole_vol_vector):
from pytools import indices_in_shape
from hedge.tools import log_shape
ls = log_shape(whole_vol_vector)
if ls != ():
result = [numpy.zeros(ls, dtype=object)
for part_emb in self._embeddings()]
for p, part_emb in enumerate(self._embeddings()):
for i in indices_in_shape(ls):
result[p][i] = whole_vol_vector[part_emb]
return result
else:
return [whole_vol_vector[part_emb]
for part_emb in self._embeddings()]
def with_object_array_or_scalar(f, field, obj_array_only=False):
if obj_array_only:
if is_obj_array(field):
ls = field.shape
else:
ls = ()
else:
ls = log_shape(field)
if ls != ():
from pytools import indices_in_shape
result = np.zeros(ls, dtype=object)
for i in indices_in_shape(ls):
result[i] = f(field[i])
return result
else:
return f(field)
def with_object_array_or_scalar(f, field, obj_array_only=False):
if obj_array_only:
if is_obj_array(field):
ls = field.shape
else:
ls = ()
else:
ls = log_shape(field)
if ls != ():
from pytools import indices_in_shape
result = np.zeros(ls, dtype=object)
for i in indices_in_shape(ls):
result[i] = f(field[i])
return result
else:
return f(field)
def count_dofs(vec):
try:
dtype = vec.dtype
size = vec.size
shape = vec.shape
except AttributeError:
from warnings import warn
warn("could not count dofs of vector")
return 0
if dtype == object:
from pytools import indices_in_shape
return sum(count_dofs(vec[i])
for i in indices_in_shape(vec.shape))
else:
return size
def split(self, whole_vol_vector):
from pytools import indices_in_shape
from hedge.tools import log_shape
ls = log_shape(whole_vol_vector)
if ls != ():
result = [
numpy.zeros(ls, dtype=object)
for part_emb in self._embeddings()
]
for p, part_emb in enumerate(self._embeddings()):
for i in indices_in_shape(ls):
result[p][i] = whole_vol_vector[part_emb]
return result
else:
return [
whole_vol_vector[part_emb] for part_emb in self._embeddings()
]
def subscripts_and_names(self):
sep_shape = self.sep_shape()
if not sep_shape:
return None
def unwrap_1d_indices(idx):
# This allows these indices to work on Python sequences, too, not
# just numpy arrays.
if len(idx) == 1:
return idx[0]
else:
return idx
from pytools import indices_in_shape
return [(unwrap_1d_indices(i), self.name + "".join("_s%d" % sub_i
for sub_i in i))
for i in indices_in_shape(sep_shape)]
def subscripts_and_names(self):
sep_shape = self.sep_shape()
if not sep_shape:
return None
def unwrap_1d_indices(idx):
# This allows these indices to work on Python sequences, too, not
# just numpy arrays.
if len(idx) == 1:
return idx[0]
else:
return idx
from pytools import indices_in_shape
return [
(unwrap_1d_indices(i),
self.name + "".join("_s%d" % sub_i for sub_i in i))
for i in indices_in_shape(sep_shape)]
def reassemble(self, parts_vol_vectors):
from pytools import single_valued, indices_in_shape
from hedge.tools import log_shape
ls = single_valued(log_shape(pvv) for pvv in parts_vol_vectors)
def remap_scalar_field(idx):
result = self.whole_discr.volume_zeros()
for part_emb, part_vol_vector in zip(self._embeddings(),
parts_vol_vectors):
result[part_emb] = part_vol_vector[idx]
return result
if ls != ():
result = numpy.zeros(ls, dtype=object)
for i in indices_in_shape(ls):
result[i] = remap_scalar_field(i)
return result
else:
return remap_scalar_field(())
def basis(self):
""""
:returns: a :class:`list` containing functions that realize
a high-order interpolation basis on the :attr:`points`.
"""
from pytools import indices_in_shape
from scipy.special import eval_chebyt
def eval_basis(ind, x):
result = 1
for i in range(self.dim):
coord = (x[i] - self.center[i])/(self.h/2)
result *= eval_chebyt(ind[i], coord)
return result
from functools import partial
return [
partial(eval_basis, ind)
for ind in indices_in_shape((self.npoints,)*self.dim)]
def reassemble(self, parts_vol_vectors):
from pytools import single_valued, indices_in_shape
from hedge.tools import log_shape
ls = single_valued(log_shape(pvv) for pvv in parts_vol_vectors)
def remap_scalar_field(idx):
result = self.whole_discr.volume_zeros()
for part_emb, part_vol_vector in zip(
self._embeddings(), parts_vol_vectors):
result[part_emb] = part_vol_vector[idx]
return result
if ls != ():
result = numpy.zeros(ls, dtype=object)
for i in indices_in_shape(ls):
result[i] = remap_scalar_field(i)
return result
else:
return remap_scalar_field(())
def basis(self):
"""
:returns: a :class:`list` containing functions that realize
a high-order interpolation basis on the :py:attr:`points`.
"""
from pytools import indices_in_shape
from scipy.special import eval_chebyt
def eval_basis(ind, x):
result = 1
for i in range(self.dim):
coord = (x[i] - self.center[i]) / (self.h / 2)
result *= eval_chebyt(ind[i], coord)
return result
from functools import partial
return [
partial(eval_basis, ind)
for ind in indices_in_shape((self.npoints, ) * self.dim)
]
def generate_linearized_array(array, value):
from pytools import product
size = product(shape_ax for shape_ax in array.shape)
if not isinstance(size, int):
raise LoopyError("cannot produce literal for array '%s': "
"shape is not a compile-time constant"
% array.name)
strides = []
data = np.zeros(size, array.dtype.numpy_dtype)
from loopy.kernel.array import FixedStrideArrayDimTag
for i, dim_tag in enumerate(array.dim_tags):
if isinstance(dim_tag, FixedStrideArrayDimTag):
if not isinstance(dim_tag.stride, int):
raise LoopyError("cannot produce literal for array '%s': "
"stride along axis %d (1-based) is not a "
"compile-time constant"
% (array.name, i+1))
strides.append(dim_tag.stride)
else:
raise LoopyError("cannot produce literal for array '%s': "
"dim_tag type '%s' not supported"
% (array.name, type(dim_tag).__name__))
assert array.offset == 0
from pytools import indices_in_shape
for ituple in indices_in_shape(value.shape):
i = sum(i_ax * strd_ax for i_ax, strd_ax in zip(ituple, strides))
data[i] = value[ituple]
return data
def generate_linearized_array(array, value):
from pytools import product
size = product(shape_ax for shape_ax in array.shape)
if not isinstance(size, int):
raise LoopyError("cannot produce literal for array '%s': "
"shape is not a compile-time constant"
% array.name)
strides = []
data = np.zeros(size, array.dtype.numpy_dtype)
from loopy.kernel.array import FixedStrideArrayDimTag
for i, dim_tag in enumerate(array.dim_tags):
if isinstance(dim_tag, FixedStrideArrayDimTag):
if not isinstance(dim_tag.stride, int):
raise LoopyError("cannot produce literal for array '%s': "
"stride along axis %d (1-based) is not a "
"compile-time constant"
% (array.name, i+1))
strides.append(dim_tag.stride)
else:
raise LoopyError("cannot produce literal for array '%s': "
"dim_tag type '%s' not supported"
% (array.name, type(dim_tag).__name__))
assert array.offset == 0
from pytools import indices_in_shape
for ituple in indices_in_shape(value.shape):
i = sum(i_ax * strd_ax for i_ax, strd_ax in zip(ituple, strides))
data[i] = value[ituple]
return data
def make_common_subexpression(field, prefix=None, scope=None):
"""Wrap *field* in a :class:`CommonSubexpression` with
*prefix*. If *field* is a :mod:`numpy` object array,
each individual entry is instead wrapped. If *field* is a
:class:`pymbolic.geometric_algebra.MultiVector`, each
coefficient is individually wrapped.
See :class:`CommonSubexpression` for the meaning of *prefix*
and *scope*.
"""
if isinstance(field, CommonSubexpression) and (
scope is None or scope == cse_scope.EVALUATION
or field.scope == scope):
# Don't re-wrap
return field
try:
from pytools.obj_array import log_shape
except ImportError:
have_obj_array = False
else:
have_obj_array = True
if have_obj_array:
ls = log_shape(field)
from pymbolic.geometric_algebra import MultiVector
if isinstance(field, MultiVector):
new_data = {}
for bits, coeff in six.iteritems(field.data):
if prefix is not None:
blade_str = field.space.blade_bits_to_str(bits, "")
component_prefix = prefix+"_"+blade_str
else:
component_prefix = None
new_data[bits] = make_common_subexpression(
coeff, component_prefix, scope)
return MultiVector(new_data, field.space)
elif have_obj_array and ls != ():
from pytools import indices_in_shape
result = numpy.zeros(ls, dtype=object)
for i in indices_in_shape(ls):
if prefix is not None:
component_prefix = prefix+"_".join(str(i_i) for i_i in i)
else:
component_prefix = None
if is_constant(field[i]):
result[i] = field[i]
else:
result[i] = make_common_subexpression(
field[i], component_prefix, scope)
return result
else:
if is_constant(field):
return field
else:
return CommonSubexpression(field, prefix, scope)
def make_common_subexpression(field, prefix=None, scope=None):
"""Wrap *field* in a :class:`CommonSubexpression` with
*prefix*. If *field* is a :mod:`numpy` object array,
each individual entry is instead wrapped. If *field* is a
:class:`pymbolic.geometric_algebra.MultiVector`, each
coefficient is individually wrapped.
See :class:`CommonSubexpression` for the meaning of *prefix*
and *scope*.
"""
if isinstance(field, CommonSubexpression) and (scope is None or scope
== cse_scope.EVALUATION
or field.scope == scope):
# Don't re-wrap
return field
try:
from pytools.obj_array import log_shape
except ImportError:
have_obj_array = False
else:
have_obj_array = True
if have_obj_array:
ls = log_shape(field)
from pymbolic.geometric_algebra import MultiVector
if isinstance(field, MultiVector):
new_data = {}
for bits, coeff in six.iteritems(field.data):
if prefix is not None:
blade_str = field.space.blade_bits_to_str(bits, "")
component_prefix = prefix + "_" + blade_str
else:
component_prefix = None
new_data[bits] = make_common_subexpression(coeff, component_prefix,
scope)
return MultiVector(new_data, field.space)
elif have_obj_array and ls != ():
from pytools import indices_in_shape
result = numpy.zeros(ls, dtype=object)
for i in indices_in_shape(ls):
if prefix is not None:
component_prefix = prefix + "_".join(str(i_i) for i_i in i)
else:
component_prefix = None
if is_constant(field[i]):
result[i] = field[i]
else:
result[i] = make_common_subexpression(field[i],
component_prefix, scope)
return result
else:
if is_constant(field):
return field
else:
return CommonSubexpression(field, prefix, scope)
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 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.
.. note:: This function does some heavy-lifting computation in Python,
which we intend to optimize in the future. In particular, we plan
to shift the batched linear solves and basis evaluations to
:mod:`loopy`.
TODO: make linear solvers available as :mod:`loopy` callables.
TODO: make :mod:`modepy` emit :mod:`loopy` callables for basis evaluation.
"""
if not isinstance(dof_vec, cl.array.Array):
raise TypeError("non-array passed to interpolator")
assert len(elements_to_sources_lookup.discr.groups) == 1
assert len(elements_to_sources_lookup.discr.mesh.groups) == 1
degroup = elements_to_sources_lookup.discr.groups[0]
megroup = elements_to_sources_lookup.discr.mesh.groups[0]
if not degroup.is_affine:
raise ValueError(
"interpolation requires global-to-local map, "
"which is only available for affinely mapped elements")
mesh = elements_to_sources_lookup.discr.mesh
dim = elements_to_sources_lookup.discr.dim
template_simplex = mesh.groups[0].vertex_unit_coordinates().T
# -------------------------------------------------------
# Inversely map source points with a global-to-local map.
#
# 1. For each element, solve for the affine map.
#
# 2. Apply the map to corresponding source points.
#
# This step computes `unit_sources`, the list of inversely
# mapped source points.
sources_in_element_starts = \
elements_to_sources_lookup.sources_in_element_starts.get(queue)
sources_in_element_lists = \
elements_to_sources_lookup.sources_in_element_lists.get(queue)
tree = elements_to_sources_lookup.tree.get(queue)
unit_sources_host = make_obj_array(
[np.zeros_like(srccrd) for srccrd in tree.sources])
for iel in range(degroup.nelements):
vertex_ids = megroup.vertex_indices[iel]
vertices = mesh.vertices[:, vertex_ids]
afa, afb = compute_affine_transform(vertices, template_simplex)
beg = sources_in_element_starts[iel]
end = sources_in_element_starts[iel + 1]
source_ids_in_el = sources_in_element_lists[beg:end]
sources_in_el = np.vstack(
[tree.sources[iaxis][source_ids_in_el] for iaxis in range(dim)])
ivmapped_el_sources = afa @ sources_in_el + afb.reshape([dim, 1])
for iaxis in range(dim):
unit_sources_host[iaxis][source_ids_in_el] = \
ivmapped_el_sources[iaxis, :]
unit_sources = make_obj_array(
[cl.array.to_device(queue, usc) for usc in unit_sources_host])
# -----------------------------------------------------
# Carry out evaluations in the local (template) frames.
#
# 1. Assemble a resampling matrix for each element, with
# the basis functions and the local source points.
#
# 2. For each element, perform matvec on the resampling
# matrix and the local DoF coefficients.
#
# This step assumes `unit_sources` computed on device, so
# that the previous step can be swapped with a kernel without
# interrupting the followed computation.
mapped_sources = np.vstack([usc.get(queue) for usc in unit_sources])
basis_funcs = degroup.basis()
arr_ctx = PyOpenCLArrayContext(queue)
dof_vec_view = unflatten(arr_ctx, elements_to_sources_lookup.discr,
dof_vec)[0]
dof_vec_view = dof_vec_view.get()
sym_shape = dof_vec.shape[:-1]
source_vec = np.zeros(sym_shape + (tree.nsources, ))
for iel in range(degroup.nelements):
beg = sources_in_element_starts[iel]
end = sources_in_element_starts[iel + 1]
source_ids_in_el = sources_in_element_lists[beg:end]
mapped_sources_in_el = mapped_sources[:, source_ids_in_el]
local_dof_vec = dof_vec_view[..., iel, :]
# resampling matrix built from Vandermonde matrices
import modepy as mp
rsplm = mp.resampling_matrix(basis=basis_funcs,
new_nodes=mapped_sources_in_el,
old_nodes=degroup.unit_nodes)
if len(sym_shape) == 0:
local_coeffs = local_dof_vec
source_vec[source_ids_in_el] = rsplm @ local_coeffs
else:
from pytools import indices_in_shape
for sym_id in indices_in_shape(sym_shape):
source_vec[sym_id + (source_ids_in_el, )] = \
rsplm @ local_dof_vec[sym_id]
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
def test_step_matrix(method, show_matrix=True, show_dag=False):
component_id = 'y'
code = method.generate()
if show_dag:
from dagrt.language import show_dependency_graph
show_dependency_graph(code)
from dagrt.exec_numpy import NumpyInterpreter
from leap.step_matrix import StepMatrixFinder
from pymbolic import var
# {{{ build matrix
def rhs_sym(t, y):
return var("lambda") * y
finder = StepMatrixFinder(code,
function_map={"<func>" + component_id: rhs_sym})
mat = finder.get_phase_step_matrix("primary")
if show_matrix:
print('Variables: %s' % finder.variables)
from pytools import indices_in_shape
for i in indices_in_shape(mat.shape):
print(i, mat[i])
# }}}
dt = 0.1
lambda_ = -0.4
def rhs(t, y):
return lambda_ * y
interp = NumpyInterpreter(code,
function_map={"<func>" + component_id: rhs})
interp.set_up(t_start=0, dt_start=dt, context={component_id: 15})
assert interp.next_phase == "initial"
for event in interp.run_single_step():
pass
assert interp.next_phase == "primary"
start_values = np.array([interp.context[v] for v in finder.variables])
for event in interp.run_single_step():
pass
assert interp.next_phase == "primary"
stop_values = np.array([interp.context[v] for v in finder.variables])
from dagrt.expression import EvaluationMapper
concrete_mat = EvaluationMapper({
"lambda": lambda_,
"<dt>": dt,
}, {})(mat)
stop_values_from_mat = concrete_mat.dot(start_values)
rel_err = (
la.norm(stop_values - stop_values_from_mat) / # noqa: W504
la.norm(stop_values))
assert rel_err < 1e-12