def argument(self, obj):
Q = obj.ufl_function_space()
dom = Q.ufl_domain()
sub_elements = obj.ufl_element().sub_elements()
# If not a mixed element, do nothing
if (len(sub_elements) == 0):
return obj
# Split into sub-elements, creating appropriate space for each
args = []
for i, sub_elem in enumerate(sub_elements):
Q_i = FunctionSpace(dom, sub_elem)
a = Argument(Q_i, obj.number(), part=obj.part())
indices = [()]
for m in a.ufl_shape:
indices = [(k + (j,)) for k in indices for j in range(m)]
if (i == self.idx[obj.number()]):
args += [a[j] for j in indices]
else:
args += [Zero() for j in indices]
return as_vector(args)
def Dn(f):
"""UFL operator: Take the directional derivative of *f* in the
facet normal direction, Dn(f) := dot(grad(f), n)."""
f = as_ufl(f)
if is_cellwise_constant(f):
return Zero(f.ufl_shape, f.ufl_free_indices, f.ufl_index_dimensions)
return dot(grad(f), FacetNormal(f.ufl_domain()))
class AssignmentBase(Operator):
"""Base class for UFL augmented assignments."""
__slots__ = ("ufl_shape", "_symbol", "_ast", "_visit")
_identity = Zero()
def __init__(self, lhs, rhs):
operands = map(ufl.as_ufl, (lhs, rhs))
super(AssignmentBase, self).__init__(operands)
self.ufl_shape = lhs.ufl_shape
# Sub function assignment, we've put a Zero in the lhs
# indicating we should do nothing.
if type(lhs) is Zero:
return
if not (isinstance(lhs, function.Function)
or isinstance(lhs, DummyFunction)):
raise TypeError("Can only assign to a Function")
def __str__(self):
return (" %s " % self._symbol).join(map(str, self.ufl_operands))
def __repr__(self):
return "%s(%s)" % (self.__class__.__name__, ", ".join(
repr(o) for o in self.ufl_operands))
@property
def ast(self):
return self._ast(_ast(self.ufl_operands[0]),
_ast(self.ufl_operands[1]))
def argument(self, arg):
"""Split an argument into its constituent spaces."""
from itertools import product
if isinstance(arg.function_space(), functionspace.MixedFunctionSpace):
# Look up the split argument in cache since we want it unique
if arg in self._args:
return self._args[arg]
args = []
for i, fs in enumerate(arg.function_space().split()):
# Build the sub-space Argument (not part of the mixed
# space).
a = Argument(fs, arg.number(), part=arg.part())
# Produce indexing iterator.
# For scalar-valued spaces this results in the empty
# tuple (), which returns the Argument when indexing.
# For vector-valued spaces, this is just
# range(a.ufl_shape[0])
# For tensor-valued spaces, it's the nested loop of
# range(x for x in a.ufl_shape).
#
# Each of the indexed things is then scalar-valued
# which we turn into a ufl Vector.
indices = product(*map(range, a.ufl_shape))
if self._idx[arg.number()] == i:
args += [a[idx] for idx in indices]
else:
args += [Zero() for _ in indices]
self._args[arg] = as_vector(args)
return self._args[arg]
return arg
def __new__(cls, *expressions):
# All lists and tuples should already be unwrapped in
# as_tensor
if any(not isinstance(e, Expr) for e in expressions):
error("Expecting only UFL expressions in ListTensor constructor.")
# Get properties of the first expression
e0 = expressions[0]
sh = e0.ufl_shape
fi = e0.ufl_free_indices
fid = e0.ufl_index_dimensions
# Obviously, each subexpression must have the same shape
if any(sh != e.ufl_shape for e in expressions[1:]):
error(
"Cannot create a tensor by joining subexpressions with different shapes."
)
if any(fi != e.ufl_free_indices for e in expressions[1:]):
error(
"Cannot create a tensor where the components have different free indices."
)
if any(fid != e.ufl_index_dimensions for e in expressions[1:]):
error(
"Cannot create a tensor where the components have different free index dimensions."
)
# Simplify to Zero if possible
if all(isinstance(e, Zero) for e in expressions):
shape = (len(expressions), ) + sh
return Zero(shape, fi, fid)
return Operator.__new__(cls)
def __new__(cls, f):
# Return zero if expression is trivially constant
if is_cellwise_constant(f):
dim = f.ufl_domain().topological_dimension()
return Zero(f.ufl_shape + (dim,), f.ufl_free_indices,
f.ufl_index_dimensions)
return CompoundDerivative.__new__(cls)
def __new__(cls, f):
# Return zero if expression is trivially constant
if is_cellwise_constant(f):
dim = find_geometric_dimension(f)
return Zero((dim,) + f.ufl_shape, f.ufl_free_indices,
f.ufl_index_dimensions)
return CompoundDerivative.__new__(cls)
def __new__(cls, *expressions):
# All lists and tuples should already be unwrapped in as_tensor
if any(not isinstance(e, Expr) for e in expressions):
error("Expecting only UFL expressions in ListTensor constructor.")
# Get properties of the first expression
e0 = expressions[0]
sh = e0.shape()
fi = e0.free_indices()
idim = e0.index_dimensions()
# Obviously, each subextpression must have the same shape
if any(sh != e.shape() for e in expressions):
error("ListTensor assumption 1 failed, "\
"please report this incident as a potential bug.")
# Are these assumptions correct? Need to think
# through the listtensor concept and free indices.
# Are there any cases where it makes sense to even
# have any free indices here?
if any(set(fi) - set(e.free_indices()) for e in expressions):
error("ListTensor assumption 2 failed, "\
"please report this incident as a potential bug.")
if any(idim != e.index_dimensions() for e in expressions):
error("ListTensor assumption 3 failed, "\
"please report this incident as a potential bug.")
# Simplify to Zero if possible
if all(isinstance(e, Zero) for e in expressions):
shape = (len(expressions), ) + sh
return Zero(shape, fi, idim)
return WrapperType.__new__(cls)
def argument(self, o):
from ufl import as_vector
from ufl.constantvalue import Zero
from dolfin import Argument
from numpy import ndindex
V = o.function_space()
if V.num_sub_spaces() == 0:
# Not on a mixed space, just return ourselves.
return o
# Already seen this argument, return the cached version.
if o in self._args:
return self._args[o]
args = []
for i, V_i_sub in enumerate(V.split()):
# Walk over the subspaces and build a vector that is zero
# where the argument does not match the one we're looking
# for and is just the non-mixed argument when we do want
# it.
V_i = V_i_sub.collapse()
a = Argument(V_i, o.number(), part=o.part())
indices = ndindex(a.ufl_shape)
if self.idx[o.number()] == i:
args += [a[j] for j in indices]
else:
args += [Zero() for j in indices]
self._args[o] = as_vector(args)
return self._args[o]
def __new__(cls, f):
# Return zero if expression is trivially constant
dim = f.geometric_dimension()
if f.is_cellwise_constant():
free_indices = f.free_indices()
index_dimensions = subdict(f.index_dimensions(), free_indices)
return Zero((dim,) + f.shape(), free_indices, index_dimensions)
return CompoundDerivative.__new__(cls)
def Dn(f):
"""UFL operator: Take the directional derivative of f in the
facet normal direction, Dn(f) := dot(grad(f), n)."""
f = as_ufl(f)
cell = f.cell()
if cell is None: # FIXME: Rather if f.is_cellwise_constant()?
return Zero(f.shape(), f.free_indices(), f.index_dimensions())
return dot(grad(f), cell.n)
def __new__(cls, a, b):
ash, bsh = a.shape(), b.shape()
if isinstance(a, Zero) or isinstance(b, Zero):
free_indices, index_dimensions = merge_indices(a, b)
return Zero(ash + bsh, free_indices, index_dimensions)
if ash == () or bsh == ():
return a * b
return CompoundTensorOperator.__new__(cls)
def test_real(self):
z0 = Zero()
z1 = as_ufl(1.0)
z2 = ComplexValue(1j)
z3 = ComplexValue(1+1j)
assert Real(z1) == z1
assert Real(z3) == z1
assert Real(z2) == z0
def __new__(cls, a, b):
ash, bsh = a.ufl_shape, b.ufl_shape
if isinstance(a, Zero) or isinstance(b, Zero):
fi, fid = merge_nonoverlapping_indices(a, b)
return Zero(ash + bsh, fi, fid)
if ash == () or bsh == ():
return a * b
return CompoundTensorOperator.__new__(cls)
def __new__(cls, f, v):
# Return zero if expression is trivially independent of variable
if isinstance(f, Terminal):
free_indices = tuple(set(f.free_indices()) ^ set(v.free_indices()))
index_dimensions = mergedicts([f.index_dimensions(), v.index_dimensions()])
index_dimensions = subdict(index_dimensions, free_indices)
return Zero(f.shape() + v.shape(), free_indices, index_dimensions)
return Derivative.__new__(cls)
def __new__(cls, a, b):
ash, bsh = a.shape(), b.shape()
ufl_assert(ash == bsh, "Shape mismatch.")
if isinstance(a, Zero) or isinstance(b, Zero):
free_indices, index_dimensions = merge_indices(a, b)
return Zero((), free_indices, index_dimensions)
if ash == ():
return a*b
return CompoundTensorOperator.__new__(cls)
def __new__(cls, f):
if f.ufl_free_indices:
error("Free indices in the divergence argument is not allowed.")
# Return zero if expression is trivially constant
if is_cellwise_constant(f):
return Zero(f.ufl_shape[1:]) # No free indices asserted above
return CompoundDerivative.__new__(cls)
def _filter(o: Expr, self: Memoizer) -> Expr:
if not isinstance(o, Expr):
raise AssertionError(f"Cannot handle term with type {type(o)}")
if self.predicate(o):
return Zero(shape=o.ufl_shape,
free_indices=o.ufl_free_indices,
index_dimensions=o.ufl_index_dimensions)
else:
return ufl_reuse_if_untouched(o, *map(self, o.ufl_operands))
def __new__(cls, A):
sh = A.shape()
ufl_assert(len(sh) == 2, "Symmetric part of tensor with rank != 2 is undefined.")
if sh[0] != sh[1]:
error("Cannot take symmetric part of rectangular matrix with dimensions %s." % repr(sh))
ufl_assert(not A.free_indices(), "Not expecting free indices in Sym.")
if isinstance(A, Zero):
return Zero(A.shape(), A.free_indices(), A.index_dimensions())
return CompoundTensorOperator.__new__(cls)
def test_imag(self):
z0 = Zero()
z1 = as_ufl(1.0)
z2 = as_ufl(1j)
z3 = ComplexValue(1+1j)
assert Imag(z2) == z1
assert Imag(z3) == z1
assert Imag(z1) == z0
def _getitem(self, component):
# Treat component consistently as tuple below
if not isinstance(component, tuple):
component = (component, )
shape = self.ufl_shape
# Analyse slices (:) and Ellipsis (...)
all_indices, slice_indices, repeated_indices = create_slice_indices(
component, shape, self.ufl_free_indices)
# Check that we have the right number of indices for a tensor with
# this shape
if len(shape) != len(all_indices):
error(
"Invalid number of indices {0} for expression of rank {1}.".format(
len(all_indices), len(shape)))
# Special case for simplifying foo[...] => foo, foo[:] => foo or
# similar
if len(slice_indices) == len(all_indices):
return self
# Special case for simplifying as_tensor(ai,(i,))[i] => ai
if isinstance(self, ComponentTensor):
if all_indices == self.indices().indices():
return self.ufl_operands[0]
# Apply all indices to index self, yielding a scalar valued
# expression
mi = MultiIndex(all_indices)
a = Indexed(self, mi)
# TODO: I think applying as_tensor after index sums results in
# cleaner expression graphs.
# If the Ellipsis or any slices were found, wrap as tensor valued
# with the slice indices created at the top here
if slice_indices:
a = as_tensor(a, slice_indices)
# If any repeated indices were found, apply implicit summation
# over those
for i in repeated_indices:
mi = MultiIndex((i, ))
a = IndexSum(a, mi)
# Check for zero (last so we can get indices etc from a, could
# possibly be done faster by checking early instead)
if isinstance(self, Zero):
shape = a.ufl_shape
fi = a.ufl_free_indices
fid = a.ufl_index_dimensions
a = Zero(shape, fi, fid)
return a
def argument(self, obj):
if (obj.part() is not None):
# Mixed element built from MixedFunctionSpace,
# whose sub-function spaces are indexed by obj.part()
if len(obj.ufl_shape) == 0:
if (obj.part() == self.idx[obj.number()]):
return obj
else:
return Zero()
else:
indices = [()]
for m in obj.ufl_shape:
indices = [(k + (j, )) for k in indices for j in range(m)]
if (obj.part() == self.idx[obj.number()]):
return as_vector([obj[j] for j in indices])
else:
return as_vector([Zero() for j in indices])
else:
# Mixed element built from MixedElement,
# whose sub-elements need their function space to be created
Q = obj.ufl_function_space()
dom = Q.ufl_domain()
sub_elements = obj.ufl_element().sub_elements()
# If not a mixed element, do nothing
if (len(sub_elements) == 0):
return obj
args = []
for i, sub_elem in enumerate(sub_elements):
Q_i = FunctionSpace(dom, sub_elem)
a = Argument(Q_i, obj.number(), part=obj.part())
indices = [()]
for m in a.ufl_shape:
indices = [(k + (j, )) for k in indices for j in range(m)]
if (i == self.idx[obj.number()]):
args += [a[j] for j in indices]
else:
args += [Zero() for j in indices]
return as_vector(args)
def _ufl_expr_reconstruct_(self, op):
"Return a new object of the same type with new operands."
if is_cellwise_constant(op):
if op.ufl_shape != self.ufl_operands[0].ufl_shape:
error("Operand shape mismatch in NablaGrad reconstruct.")
if self.ufl_operands[0].ufl_free_indices != op.ufl_free_indices:
error("Free index mismatch in NablaGrad reconstruct.")
return Zero(self.ufl_shape, self.ufl_free_indices,
self.ufl_index_dimensions)
return self._ufl_class_(op)
def reconstruct(self, op):
"Return a new object of the same type with new operands."
if op.is_cellwise_constant():
ufl_assert(op.shape() == self._f.shape(),
"Operand shape mismatch in NablaGrad reconstruct.")
ufl_assert(self._f.free_indices() == op.free_indices(),
"Free index mismatch in NablaGrad reconstruct.")
return Zero(self.shape(), self.free_indices(),
self.index_dimensions())
return self.__class__._uflclass(op)
def __new__(cls, A):
# Checks
if len(A.ufl_shape) != 2:
error("Trace of tensor with rank != 2 is undefined.")
# Simplification
if isinstance(A, Zero):
return Zero((), A.ufl_free_indices, A.ufl_index_dimensions)
return CompoundTensorOperator.__new__(cls)
def __new__(cls, a, b):
ash, bsh = a.shape(), b.shape()
ufl_assert(len(ash) == 1 and ash == bsh,
"Cross product requires arguments of rank 1.")
if isinstance(a, Zero) or isinstance(b, Zero):
free_indices, index_dimensions = merge_indices(a, b)
return Zero(ash, free_indices, index_dimensions)
return CompoundTensorOperator.__new__(cls)
def __new__(cls, f):
ufl_assert(not f.free_indices(), \
"TODO: Taking divergence of an expression with free indices,"\
"should this be a valid expression? Please provide examples!")
# Return zero if expression is trivially constant
if f.is_cellwise_constant():
return Zero(f.shape()[1:]) # No free indices asserted above
return CompoundDerivative.__new__(cls)
def __new__(cls, expression, indices):
# Simplify
if isinstance(expression, Zero):
fi, fid, sh = remove_indices(expression.ufl_free_indices,
expression.ufl_index_dimensions,
[ind.count() for ind in indices])
return Zero(sh, fi, fid)
# Construct
return Operator.__new__(cls)
def __new__(cls, expression, indices):
if isinstance(expression, Zero):
if isinstance(indices, MultiIndex):
indices = tuple(indices)
elif not isinstance(indices, tuple):
indices = (indices, )
dims = expression.index_dimensions()
shape = tuple(dims[i] for i in indices)
fi = tuple(set(expression.free_indices()) - set(indices))
idim = dict((i, dims[i]) for i in fi)
return Zero(shape, fi, idim)
return WrapperType.__new__(cls)
def __new__(cls, a):
a = as_ufl(a)
# Simplification
if isinstance(a, Zero):
return a
if isinstance(a, (Real, Imag, Abs)):
return Zero(a.ufl_shape, a.ufl_free_indices, a.ufl_index_dimensions)
if isinstance(a, ScalarValue):
return as_ufl(a.imag())
return Operator.__new__(cls)
