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)