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 separate_by_real_and_imag(data, real_only):
for name, field in data:
from pytools.obj_array import log_shape, is_obj_array
ls = log_shape(field)
if is_obj_array(field):
assert len(ls) == 1
from pytools.obj_array import (
oarray_real_copy, oarray_imag_copy,
with_object_array_or_scalar)
if field[0].dtype.kind == "c":
if real_only:
yield (name,
with_object_array_or_scalar(oarray_real_copy, field))
else:
yield (name+"_r",
with_object_array_or_scalar(oarray_real_copy, field))
yield (name+"_i",
with_object_array_or_scalar(oarray_imag_copy, field))
else:
yield (name, field)
else:
# ls == ()
if field.dtype.kind == "c":
yield (name+"_r", field.real.copy())
yield (name+"_i", field.imag.copy())
else:
yield (name, field)
def separate_by_real_and_imag(data, real_only):
for name, field in data:
from pytools.obj_array import log_shape
ls = log_shape(field)
if ls != () and ls[0] > 1:
assert len(ls) == 1
from pytools.obj_array import (oarray_real_copy, oarray_imag_copy,
with_object_array_or_scalar)
if field[0].dtype.kind == "c":
if real_only:
yield (name,
with_object_array_or_scalar(oarray_real_copy,
field))
else:
yield (name + "_r",
with_object_array_or_scalar(oarray_real_copy,
field))
yield (name + "_i",
with_object_array_or_scalar(oarray_imag_copy,
field))
else:
yield (name, field)
else:
# ls == ()
if field.dtype.kind == "c":
yield (name + "_r", field.real.copy())
yield (name + "_i", field.imag.copy())
else:
yield (name, field)
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 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 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 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 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, 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 join_fields(*args):
from pytools.obj_array import make_obj_array, log_shape
from pymbolic.geometric_algebra import MultiVector, bit_count
res_list = []
for arg in args:
if isinstance(arg, list):
res_list.extend(arg)
elif isinstance(arg, MultiVector):
for grade in arg.all_grades():
for bits in range(2 ** arg.space.dimensions):
if bit_count(bits) == grade:
res_list.append(arg.data.get(bits, 0))
elif isinstance(arg, np.ndarray):
if log_shape(arg) == ():
res_list.append(arg)
else:
res_list.extend(arg.flat)
else:
res_list.append(arg)
return make_obj_array(res_list)
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)
