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)