def test_expressions():
x = gem.Variable("x", (3, 4))
y = gem.Variable("y", (4, ))
i, j = gem.indices(2)
xij = x[i, j]
yj = y[j]
assert xij == gem.Indexed(x, (i, j))
assert yj == gem.Indexed(y, (j, ))
assert xij + yj == gem.Sum(xij, yj)
assert xij * yj == gem.Product(xij, yj)
assert xij - yj == gem.Sum(xij, gem.Product(gem.Literal(-1), yj))
assert xij / yj == gem.Division(xij, yj)
assert xij + 1 == gem.Sum(xij, gem.Literal(1))
assert 1 + xij == gem.Sum(gem.Literal(1), xij)
assert (xij + y).shape == (4, )
assert (x @ y).shape == (3, )
assert x.T.shape == (4, 3)
with pytest.raises(ValueError):
xij.T @ y
with pytest.raises(ValueError):
xij + "foo"
def conditional(self, o, condition, then, else_):
assert condition.shape == ()
shape, = set([then.shape, else_.shape])
indices = gem.indices(len(shape))
return gem.ComponentTensor(
gem.Conditional(condition, gem.Indexed(then, indices),
gem.Indexed(else_, indices)), indices)
def matvec(table):
i, j = gem.indices(2)
value_indices = self.get_value_indices()
table = gem.Indexed(table, (j, ) + value_indices)
val = gem.ComponentTensor(gem.IndexSum(M[i, j] * table, (j, )),
(i, ) + value_indices)
# Eliminate zeros
return gem.optimise.aggressive_unroll(val)
def dual_basis(self):
# Outer product the dual bases of the factors
qs, pss = zip(*(factor.dual_basis for factor in self.factors))
ps = TensorPointSet(pss)
qis = tuple(q[gem.indices(len(q.shape))] for q in qs)
indices = tuple(chain(*(q.index_ordering() for q in qis)))
Q = gem.ComponentTensor(reduce(gem.Product, qis), indices)
return Q, ps
def compile_to_gem(expr, translator):
"""Compile a single pointwise expression to GEM.
:arg expr: The expression to compile.
:arg translator: a :class:`Translator` instance.
:returns: A (lvalue, rvalue) pair of preprocessed GEM."""
if not isinstance(expr, Assign):
raise ValueError(
f"Don't know how to assign expression of type {type(expr)}")
spaces = tuple(c.function_space() for c in expr.coefficients)
if any(
type(s.ufl_element()) is ufl.MixedElement for s in spaces
if s is not None):
raise ValueError("Not expecting a mixed space at this point, "
"did you forget to index a function with .sub(...)?")
if len(set(s.ufl_element() for s in spaces if s is not None)) != 1:
raise ValueError("All coefficients must be defined on the same space")
lvalue = expr.lvalue
rvalue = expr.rvalue
broadcast = all(
isinstance(c, firedrake.Constant)
for c in expr.rcoefficients) and rvalue.ufl_shape == ()
if not broadcast and lvalue.ufl_shape != rvalue.ufl_shape:
try:
rvalue = reshape(rvalue, lvalue.ufl_shape)
except ValueError:
raise ValueError(
"Mismatching shapes between lvalue and rvalue in pointwise assignment"
)
rvalue, = map_expr_dags(LowerCompoundAlgebra(), [rvalue])
try:
lvalue, rvalue = map_expr_dags(translator, [lvalue, rvalue])
except (AssertionError, ValueError):
raise ValueError("Mismatching shapes in pointwise assignment. "
"For intrinsically vector-/tensor-valued spaces make "
"sure you're not using shaped Constants or literals.")
indices = gem.indices(len(lvalue.shape))
if not broadcast:
if rvalue.shape != lvalue.shape:
raise ValueError(
"Mismatching shapes in pointwise assignment. "
"For intrinsically vector-/tensor-valued spaces make "
"sure you're not using shaped Constants or literals.")
rvalue = gem.Indexed(rvalue, indices)
lvalue = gem.Indexed(lvalue, indices)
if isinstance(expr, IAdd):
rvalue = gem.Sum(lvalue, rvalue)
elif isinstance(expr, ISub):
rvalue = gem.Sum(lvalue, gem.Product(gem.Literal(-1), rvalue))
elif isinstance(expr, IMul):
rvalue = gem.Product(lvalue, rvalue)
elif isinstance(expr, IDiv):
rvalue = gem.Division(lvalue, rvalue)
return preprocess_gem([lvalue, rvalue])
def dual_basis(self):
# Return Q with x.indices already a free index for the
# consumer to use
# expensive numerical extraction is done once per element
# instance, but the point set must be created every time we
# build the dual.
Q, pts = self._dual_basis
x = PointSet(pts)
assert len(x.indices) == 1
assert Q.shape[1] == x.indices[0].extent
i, *js = gem.indices(len(Q.shape) - 1)
Q = gem.ComponentTensor(gem.Indexed(Q, (i, *x.indices, *js)), (i, *js))
return Q, x
def dual_evaluation(self, fn):
'''Get a GEM expression for performing the dual basis evaluation at
the nodes of the reference element. Currently only works for flat
elements: tensor elements are implemented in
:class:`TensorFiniteElement`.
:param fn: Callable representing the function to dual evaluate.
Callable should take in an :class:`AbstractPointSet` and
return a GEM expression for evaluation of the function at
those points.
:returns: A tuple ``(dual_evaluation_gem_expression, basis_indices)``
where the given ``basis_indices`` are those needed to form a
return expression for the code which is compiled from
``dual_evaluation_gem_expression`` (alongside any argument
multiindices already encoded within ``fn``)
'''
Q, x = self.dual_basis
expr = fn(x)
# Apply targeted sum factorisation and delta elimination to
# the expression
sum_indices, factors = delta_elimination(*traverse_product(expr))
expr = sum_factorise(sum_indices, factors)
# NOTE: any shape indices in the expression are because the
# expression is tensor valued.
assert expr.shape == Q.shape[len(Q.shape) - len(expr.shape):]
shape_indices = gem.indices(len(expr.shape))
basis_indices = gem.indices(len(Q.shape) - len(expr.shape))
Qi = Q[basis_indices + shape_indices]
expri = expr[shape_indices]
evaluation = gem.IndexSum(Qi * expri, x.indices + shape_indices)
# Now we want to factorise over the new contraction with x,
# ignoring any shape indices to avoid hitting the sum-
# factorisation index limit (this is a bit of a hack).
# Really need to do a more targeted job here.
evaluation = gem.optimise.contraction(evaluation, shape_indices)
return evaluation, basis_indices
def index_sum(self, o, summand, indices):
index, = indices
indices = gem.indices(len(summand.shape))
return gem.ComponentTensor(
gem.IndexSum(gem.Indexed(summand, indices), (index, )), indices)
def abs(self, o, expr):
indices = gem.indices(len(expr.shape))
return gem.ComponentTensor(
gem.MathFunction('abs', gem.Indexed(expr, indices)), indices)
def sum(self, o, *ops):
shape, = set(o.shape for o in ops)
indices = gem.indices(len(shape))
return gem.ComponentTensor(
gem.Sum(*[gem.Indexed(op, indices) for op in ops]), indices)
