def cell_normal(self, o): if self._preserve_types[o._ufl_typecode_]: return o domain = o.ufl_domain() gdim = domain.geometric_dimension() tdim = domain.topological_dimension() if tdim == gdim - 1: # n-manifold embedded in n-1 space i = Index() J = self.jacobian(Jacobian(domain)) if tdim == 2: # Surface in 3D t0 = as_vector(J[i, 0], i) t1 = as_vector(J[i, 1], i) cell_normal = cross_expr(t0, t1) elif tdim == 1: # Line in 2D (cell normal is 'up' for a line pointing # to the 'right') cell_normal = as_vector((-J[1, 0], J[0, 0])) else: error("Cell normal not implemented for tdim %d, gdim %d" % (tdim, gdim)) # Return normalized vector, sign corrected by cell # orientation co = CellOrientation(domain) return co * cell_normal / sqrt(cell_normal[i] * cell_normal[i]) else: error("What do you want cell normal in gdim={0}, tdim={1} to be?". format(gdim, tdim))
def cell_normal(self, o): if self._preserve_types[o._ufl_typecode_]: return o domain = o.ufl_domain() gdim = domain.geometric_dimension() tdim = domain.topological_dimension() if tdim == gdim - 1: # n-manifold embedded in n-1 space i = Index() J = self.jacobian(Jacobian(domain)) if tdim == 2: # Surface in 3D t0 = as_vector(J[i, 0], i) t1 = as_vector(J[i, 1], i) cell_normal = cross_expr(t0, t1) elif tdim == 1: # Line in 2D (cell normal is 'up' for a line pointing # to the 'right') cell_normal = as_vector((-J[1, 0], J[0, 0])) else: error("Cell normal not implemented for tdim %d, gdim %d" % (tdim, gdim)) # Return normalized vector, sign corrected by cell # orientation co = CellOrientation(domain) return co * cell_normal / sqrt(cell_normal[i]*cell_normal[i]) else: error("What do you want cell normal in gdim={0}, tdim={1} to be?".format(gdim, tdim))
def curl(self, o, a): # o = curl a = "[a.dx(1), -a.dx(0)]" if a.ufl_shape == () # o = curl a = "cross(nabla, (a0, a1, 0))[2]" if a.ufl_shape == (2,) # o = curl a = "cross(nabla, a)" if a.ufl_shape == (3,) def c(i, j): return a[j].dx(i) - a[i].dx(j) sh = a.ufl_shape if sh == (): return as_vector((a.dx(1), -a.dx(0))) if sh == (2,): return c(0, 1) if sh == (3,): return as_vector((c(1, 2), c(2, 0), c(0, 1))) error("Invalid shape %s of curl argument." % (sh,))
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 cross_expr(a, b): assert len(a) == 3 assert len(b) == 3 def c(i, j): return a[i] * b[j] - a[j] * b[i] return as_vector((c(1, 2), c(2, 0), c(0, 1)))
def curl(self, o, a): # o = curl a = "[a.dx(1), -a.dx(0)]" if a.shape() == () # o = curl a = "cross(nabla, (a0, a1, 0))[2]" if a.shape() == (2,) # o = curl a = "cross(nabla, a)" if a.shape() == (3,) Da = Grad(a) def c(i, j): #return a[j].dx(i) - a[i].dx(j) return Da[j,i] - Da[i,j] sh = a.shape() if sh == (): #return as_vector((a.dx(1), -a.dx(0))) return as_vector((Da[1], -Da[0])) if sh == (2,): return c(0,1) if sh == (3,): return as_vector((c(1,2), c(2,0), c(0,1))) error("Invalid shape %s of curl argument." % (sh,))
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 apply_single_function_pullbacks(r, element): """Apply an appropriate pullback to something in physical space :arg r: An expression wrapped in ReferenceValue. :arg element: The element this expression lives in. :returns: a pulled back expression.""" mapping = element.mapping() if r.ufl_shape != element.reference_value_shape(): error("Expecting reference space expression with shape '%s', got '%s'" % (element.reference_value_shape(), r.ufl_shape)) if mapping in {"physical", "identity", "contravariant Piola", "covariant Piola", "double contravariant Piola", "double covariant Piola", "L2 Piola"}: # Base case in recursion through elements. If the element # advertises a mapping we know how to handle, do that # directly. f = apply_known_single_pullback(r, element) if f.ufl_shape != element.value_shape(): error("Expecting pulled back expression with shape '%s', got '%s'" % (element.value_shape(), f.ufl_shape)) return f elif mapping in {"symmetries", "undefined"}: # Need to pull back each unique piece of the reference space thing gsh = element.value_shape() rsh = r.ufl_shape if mapping == "symmetries": subelem = element.sub_elements()[0] fcm = element.flattened_sub_element_mapping() offsets = (product(subelem.reference_value_shape()) * i for i in fcm) elements = repeat(subelem) else: elements = sub_elements_with_mappings(element) # Python >= 3.8 has an initial keyword argument to # accumulate, but 3.7 does not. offsets = chain([0], accumulate(product(e.reference_value_shape()) for e in elements)) rflat = as_vector([r[idx] for idx in numpy.ndindex(rsh)]) g_components = [] # For each unique piece in reference space, apply the appropriate pullback for offset, subelem in zip(offsets, elements): sub_rsh = subelem.reference_value_shape() rm = product(sub_rsh) rsub = [rflat[offset + i] for i in range(rm)] rsub = as_tensor(numpy.asarray(rsub).reshape(sub_rsh)) rmapped = apply_single_function_pullbacks(rsub, subelem) # Flatten into the pulled back expression for the whole thing g_components.extend([rmapped[idx] for idx in numpy.ndindex(rmapped.ufl_shape)]) # And reshape appropriately f = as_tensor(numpy.asarray(g_components).reshape(gsh)) if f.ufl_shape != element.value_shape(): error("Expecting pulled back expression with shape '%s', got '%s'" % (element.value_shape(), f.ufl_shape)) return f else: error("Unhandled mapping type '%s'" % mapping)
def diag_vector(A): """UFL operator: Take the diagonal part of rank 2 tensor A and return as a vector. See also diag.""" # TODO: Make a compound type for this operator # Get and check dimensions ufl_assert(A.rank() == 2, "Expecting rank 2 tensor.") m, n = A.shape() ufl_assert(m == n, "Can only take diagonal of square tensors.") # Return diagonal vector return as_vector([A[i, i] for i in range(n)])
def diag_vector(A): """UFL operator: Take the diagonal part of rank 2 tensor A and return as a vector. See also diag.""" # TODO: Make a compound type for this operator # Get and check dimensions ufl_assert(A.rank() == 2, "Expecting rank 2 tensor.") m, n = A.shape() ufl_assert(m == n, "Can only take diagonal of square tensors.") # Return diagonal vector return as_vector([A[i,i] for i in range(n)])
def diag_vector(A): """UFL operator: Take the diagonal part of rank 2 tensor *A* and return as a vector. See also ``diag``.""" # TODO: Make a compound type for this operator # Get and check dimensions if len(A.ufl_shape) != 2: error("Expecting rank 2 tensor.") m, n = A.ufl_shape if m != n: error("Can only take diagonal of square tensors.") # Return diagonal vector return as_vector([A[i, i] for i in range(n)])
def facet_normal(self, o): if self._preserve_types[o._ufl_typecode_]: return o domain = o.ufl_domain() tdim = domain.topological_dimension() if tdim == 1: # Special-case 1D (possibly immersed), for which we say # that n is just in the direction of J. J = self.jacobian(Jacobian(domain)) # dx/dX ndir = J[:, 0] gdim = domain.geometric_dimension() if gdim == 1: nlen = abs(ndir[0]) else: i = Index() nlen = sqrt(ndir[i] * ndir[i]) rn = ReferenceNormal(domain) # +/- 1.0 here n = rn[0] * ndir / nlen r = n else: # Recall that the covariant Piola transform u -> J^(-T)*u # preserves tangential components. The normal vector is # characterised by having zero tangential component in # reference and physical space. Jinv = self.jacobian_inverse(JacobianInverse(domain)) i, j = indices(2) rn = ReferenceNormal(domain) # compute signed, unnormalised normal; note transpose ndir = as_vector(Jinv[j, i] * rn[j], i) # normalise i = Index() n = ndir / sqrt(ndir[i] * ndir[i]) r = n if r.ufl_shape != o.ufl_shape: error("Inconsistent dimensions (in=%d, out=%d)." % (o.ufl_shape[0], r.ufl_shape[0])) return r
def facet_normal(self, o): if self._preserve_types[o._ufl_typecode_]: return o domain = o.ufl_domain() tdim = domain.topological_dimension() if tdim == 1: # Special-case 1D (possibly immersed), for which we say # that n is just in the direction of J. J = self.jacobian(Jacobian(domain)) # dx/dX ndir = J[:, 0] gdim = domain.geometric_dimension() if gdim == 1: nlen = abs(ndir[0]) else: i = Index() nlen = sqrt(ndir[i]*ndir[i]) rn = ReferenceNormal(domain) # +/- 1.0 here n = rn[0] * ndir / nlen r = n else: # Recall that the covariant Piola transform u -> J^(-T)*u # preserves tangential components. The normal vector is # characterised by having zero tangential component in # reference and physical space. Jinv = self.jacobian_inverse(JacobianInverse(domain)) i, j = indices(2) rn = ReferenceNormal(domain) # compute signed, unnormalised normal; note transpose ndir = as_vector(Jinv[j, i] * rn[j], i) # normalise i = Index() n = ndir / sqrt(ndir[i]*ndir[i]) r = n if r.ufl_shape != o.ufl_shape: error("Inconsistent dimensions (in=%d, out=%d)." % (o.ufl_shape[0], r.ufl_shape[0])) return r
def cross(self, o, a, b): def c(i, j): return Product(a[i], b[j]) - Product(a[j], b[i]) return as_vector((c(1, 2), c(2, 0), c(0, 1)))
def perp(v): "UFL operator: Take the perp of *v*, i.e. :math:`(-v_1, +v_0)`." v = as_ufl(v) if v.ufl_shape != (2,): error("Expecting a 2D vector expression.") return as_vector((-v[1], v[0]))
def split(v): """UFL operator: If v is a Coefficient or Argument in a mixed space, returns a tuple with the function components corresponding to the subelements.""" # Default range is all of v begin = 0 end = None if isinstance(v, Indexed): # Special case: split previous output of split again # Consistent with simple element, just return function in a tuple return (v, ) elif isinstance(v, ListTensor): # Special case: split previous output of split again ops = v.ufl_operands if all(isinstance(comp, Indexed) for comp in ops): args = [comp.ufl_operands[0] for comp in ops] if all(args[0] == args[i] for i in range(1, len(args))): # Get innermost terminal here and its element v = args[0] # Get relevant range of v components begin, = ops[0].ufl_operands[1] end, = ops[-1].ufl_operands[1] begin = int(begin) end = int(end) + 1 else: error("Don't know how to split %s." % (v, )) else: error("Don't know how to split %s." % (v, )) # Special case: simple element, just return function in a tuple element = v.ufl_element() if not isinstance(element, MixedElement): assert end is None return (v, ) if isinstance(element, TensorElement): if element.symmetry(): error("Split not implemented for symmetric tensor elements.") if len(v.ufl_shape) != 1: error( "Don't know how to split tensor valued mixed functions without flattened index space." ) # Compute value size and set default range end value_size = product(element.value_shape()) if end is None: end = value_size else: # Recursively dive into mixedelement in to subelement # corresponding to beginning of range j = begin while True: sub_i, j = element.extract_subelement_component(j) element = element.sub_elements()[sub_i] # Then break when we find the subelement that covers the whole range if product(element.value_shape()) == (end - begin): break # Build expressions representing the subfunction of v for each subelement offset = begin sub_functions = [] for i, e in enumerate(element.sub_elements()): # Get shape, size, indices, and v components # corresponding to subelement value shape = e.value_shape() strides = shape_to_strides(shape) rank = len(shape) sub_size = product(shape) subindices = [ flatten_multiindex(c, strides) for c in compute_indices(shape) ] components = [v[k + offset] for k in subindices] # Shape components into same shape as subelement if rank == 0: subv, = components elif rank <= 1: subv = as_vector(components) elif rank == 2: subv = as_matrix([ components[i * shape[1]:(i + 1) * shape[1]] for i in range(shape[0]) ]) else: error( "Don't know how to split functions with sub functions of rank %d." % rank) offset += sub_size sub_functions.append(subv) if end != offset: error( "Function splitting failed to extract components for whole intended range. Something is wrong." ) return tuple(sub_functions)
def perp(v): "UFL operator: Take the perp of v, i.e. (-v1, +v0)." v = as_ufl(v) ufl_assert(v.shape() == (2,), "Expecting a 2D vector expression.") return as_vector((-v[1],v[0]))
def perp(v): "UFL operator: Take the perp of v, i.e. (-v1, +v0)." v = as_ufl(v) ufl_assert(v.shape() == (2, ), "Expecting a 2D vector expression.") return as_vector((-v[1], v[0]))
def perp(v): "UFL operator: Take the perp of *v*, i.e. :math:`(-v_1, +v_0)`." v = as_ufl(v) if v.ufl_shape != (2, ): error("Expecting a 2D vector expression.") return as_vector((-v[1], v[0]))
def split(v): """UFL operator: If v is a Coefficient or Argument in a mixed space, returns a tuple with the function components corresponding to the subelements.""" # Default range is all of v begin = 0 end = None if isinstance(v, Indexed): # Special case: split previous output of split again # Consistent with simple element, just return function in a tuple return (v,) elif isinstance(v, ListTensor): # Special case: split previous output of split again ops = v.ufl_operands if all(isinstance(comp, Indexed) for comp in ops): args = [comp.ufl_operands[0] for comp in ops] if all(args[0] == args[i] for i in range(1, len(args))): # Get innermost terminal here and its element v = args[0] # Get relevant range of v components begin, = ops[0].ufl_operands[1] end, = ops[-1].ufl_operands[1] begin = int(begin) end = int(end) + 1 else: error("Don't know how to split %s." % (v,)) else: error("Don't know how to split %s." % (v,)) # Special case: simple element, just return function in a tuple element = v.ufl_element() if not isinstance(element, MixedElement): assert end is None return (v,) if isinstance(element, TensorElement): if element.symmetry(): error("Split not implemented for symmetric tensor elements.") if len(v.ufl_shape) != 1: error("Don't know how to split tensor valued mixed functions without flattened index space.") # Compute value size and set default range end value_size = product(element.value_shape()) if end is None: end = value_size else: # Recursively dive into mixedelement in to subelement # corresponding to beginning of range j = begin while True: sub_i, j = element.extract_subelement_component(j) element = element.sub_elements()[sub_i] # Then break when we find the subelement that covers the whole range if product(element.value_shape()) == (end - begin): break # Build expressions representing the subfunction of v for each subelement offset = begin sub_functions = [] for i, e in enumerate(element.sub_elements()): # Get shape, size, indices, and v components # corresponding to subelement value shape = e.value_shape() strides = shape_to_strides(shape) rank = len(shape) sub_size = product(shape) subindices = [flatten_multiindex(c, strides) for c in compute_indices(shape)] components = [v[k + offset] for k in subindices] # Shape components into same shape as subelement if rank == 0: subv, = components elif rank <= 1: subv = as_vector(components) elif rank == 2: subv = as_matrix([components[i*shape[1]: (i+1)*shape[1]] for i in range(shape[0])]) else: error("Don't know how to split functions with sub functions of rank %d." % rank) offset += sub_size sub_functions.append(subv) if end != offset: error("Function splitting failed to extract components for whole intended range. Something is wrong.") return tuple(sub_functions)
def apply_single_function_pullbacks(g): element = g.ufl_element() mapping = element.mapping() r = ReferenceValue(g) gsh = g.ufl_shape rsh = r.ufl_shape if mapping == "physical": # TODO: Is this right for immersed things assert gsh == rsh return r # Shortcut the "identity" case which includes Expression and # Constant from dolfin that may be ill-formed without a domain # (until we get that fixed) if mapping == "identity": assert rsh == gsh return r gsize = product(gsh) rsize = product(rsh) # Create some geometric objects for reuse domain = g.ufl_domain() J = Jacobian(domain) detJ = JacobianDeterminant(domain) Jinv = JacobianInverse(domain) # Create contravariant transform for reuse (note that detJ is the # _signed_ (pseudo-)determinant) transform_hdiv = (1.0 / detJ) * J # Shortcut simple cases for a more efficient representation, # including directly Piola-mapped elements and mixed elements of # any combination of affinely mapped elements without symmetries if mapping == "symmetries": fcm = element.flattened_sub_element_mapping() assert gsize >= rsize assert len(fcm) == gsize assert sorted(set(fcm)) == sorted(range(rsize)) g_components = [r[fcm[i]] for i in range(gsize)] g_components = reshape_to_nested_list(g_components, gsh) f = as_tensor(g_components) assert f.ufl_shape == g.ufl_shape return f elif mapping == "contravariant Piola": assert transform_hdiv.ufl_shape == (gsize, rsize) i, j = indices(2) f = as_vector(transform_hdiv[i, j] * r[j], i) # f = as_tensor(transform_hdiv[i, j]*r[k,j], (k,i)) # FIXME: Handle Vector(Piola) here? assert f.ufl_shape == g.ufl_shape return f elif mapping == "covariant Piola": assert Jinv.ufl_shape == (rsize, gsize) i, j = indices(2) f = as_vector(Jinv[j, i] * r[j], i) # f = as_tensor(Jinv[j, i]*r[k,j], (k,i)) # FIXME: Handle Vector(Piola) here? assert f.ufl_shape == g.ufl_shape return f elif mapping == "double covariant Piola": i, j, m, n = indices(4) f = as_tensor(Jinv[m, i] * r[m, n] * Jinv[n, j], (i, j)) assert f.ufl_shape == g.ufl_shape return f elif mapping == "double contravariant Piola": i, j, m, n = indices(4) f = as_tensor( (1.0 / detJ) * (1.0 / detJ) * J[i, m] * r[m, n] * J[j, n], (i, j)) assert f.ufl_shape == g.ufl_shape return f elif mapping == "L2 Piola": assert rsh == gsh return r / detJ # By placing components in a list and using as_vector at the end, # we're assuming below that both global function g and its # reference value r have vector shape, which is the case for most # elements with the exceptions: # - TensorElements # - All cases with scalar subelements and without symmetries # are covered by the shortcut above # (ONLY IF REFERENCE VALUE SHAPE PRESERVES TENSOR RANK) # - All cases with scalar subelements and without symmetries are # covered by the shortcut above g_components = [None] * gsize gpos = 0 rpos = 0 r = as_vector([r[idx] for idx in numpy.ndindex(r.ufl_shape)]) for subelm in sub_elements_with_mappings(element): gm = product(subelm.value_shape()) rm = product(subelm.reference_value_shape()) mp = subelm.mapping() if mp == "identity": assert gm == rm for i in range(gm): g_components[gpos + i] = r[rpos + i] elif mp == "symmetries": """ tensor_element.value_shape() == (2,2) tensor_element.reference_value_shape() == (3,) tensor_element.symmetry() == { (1,0): (0,1) } tensor_element.component_mapping() == { (0,0): 0, (0,1): 1, (1,0): 1, (1,1): 2 } tensor_element.flattened_component_mapping() == { 0: 0, 1: 1, 2: 1, 3: 2 } """ fcm = subelm.flattened_sub_element_mapping() assert gm >= rm assert len(fcm) == gm assert sorted(set(fcm)) == sorted(range(rm)) for i in range(gm): g_components[gpos + i] = r[rpos + fcm[i]] elif mp == "contravariant Piola": assert transform_hdiv.ufl_shape == (gm, rm) # Get reference value vector corresponding to this subelement: rv = as_vector([r[rpos + k] for k in range(rm)]) # Apply transform with IndexSum over j for each row j = Index() for i in range(gm): g_components[gpos + i] = transform_hdiv[i, j] * rv[j] elif mp == "covariant Piola": assert Jinv.ufl_shape == (rm, gm) # Get reference value vector corresponding to this subelement: rv = as_vector([r[rpos + k] for k in range(rm)]) # Apply transform with IndexSum over j for each row j = Index() for i in range(gm): g_components[gpos + i] = Jinv[j, i] * rv[j] elif mp == "double covariant Piola": # components are flatten, map accordingly rv = as_vector([r[rpos + k] for k in range(rm)]) (gdim, _) = subelm.value_shape() (rdim, _) = subelm.reference_value_shape() for i in range(gdim): for j in range(gdim): gv = 0 # int times Index is not allowed. so sum by hand for m in range(rdim): for n in range(rdim): gv += Jinv[m, i] * rv[m * rdim + n] * Jinv[n, j] g_components[gpos + i * gdim + j] = gv elif mp == "double contravariant Piola": # components are flatten, map accordingly rv = as_vector([r[rpos + k] for k in range(rm)]) (gdim, _) = subelm.value_shape() (rdim, _) = subelm.reference_value_shape() for i in range(gdim): for j in range(gdim): gv = 0 # int times Index is not allowed. so sum by hand for m in range(rdim): for n in range(rdim): gv += ((1.0 / detJ) * (1.0 / detJ) * J[i, m] * rv[m * rdim + n] * J[j, n]) g_components[gpos + i * gdim + j] = gv elif mp == "L2 Piola": assert gm == rm for i in range(gm): g_components[gpos + i] = r[rpos + i] / detJ else: error("Unknown subelement mapping type %s for element %s." % (mp, str(subelm))) gpos += gm rpos += rm # Wrap up components in a vector, must return same shape as input # function g f = as_tensor(numpy.asarray(g_components).reshape(gsh)) assert f.ufl_shape == g.ufl_shape return f
def split(v): """UFL operator: If v is a Coefficient or Argument in a mixed space, returns a tuple with the function components corresponding to the subelements.""" # Special case: simple element, just return function in a tuple element = v.element() if not isinstance(element, MixedElement): return (v,) if isinstance(element, TensorElement): s = element.symmetry() if s: # FIXME: How should this be defined? Should we return one subfunction # for each value component or only for those not mapped to another? # I think split should ignore the symmetry. error("Split not implemented for symmetric tensor elements.") # Compute value size value_size = product(element.value_shape()) actual_value_size = value_size # Extract sub coefficient offset = 0 sub_functions = [] for i, e in enumerate(element.sub_elements()): shape = e.value_shape() rank = len(shape) if rank == 0: # This subelement is a scalar, always maps to a single value subv = v[offset] offset += 1 elif rank == 1: # This subelement is a vector, always maps to a sequence of values sub_size, = shape components = [v[j] for j in range(offset, offset + sub_size)] subv = as_vector(components) offset += sub_size elif rank == 2: # This subelement is a tensor, possibly with symmetries, slightly more complicated... # Size of this subvalue sub_size = product(shape) # If this subelement is a symmetric element, subtract symmetric components s = None if isinstance(e, TensorElement): s = e.symmetry() s = s or EmptyDict # If we do this, we must fix the size computation in MixedElement.__init__ as well #actual_value_size -= len(s) #sub_size -= len(s) #print s # Build list of lists of value components components = [] for ii in range(shape[0]): row = [] for jj in range(shape[1]): # Map component (i,j) through symmetry mapping c = (ii, jj) c = s.get(c, c) i, j = c # Extract component c of this subvalue from global tensor v if v.rank() == 1: # Mapping into a flattened vector k = offset + i*shape[1] + j component = v[k] #print "k, offset, i, j, shape, component", k, offset, i, j, shape, component elif v.rank() == 2: # Mapping into a concatenated tensor (is this a figment of my imagination?) error("Not implemented.") row_offset, col_offset = 0, 0 # TODO k = (row_offset + i, col_offset + j) component = v[k] row.append(component) components.append(row) # Make a matrix of the components subv = as_matrix(components) offset += sub_size else: # TODO: Handle rank > 2? Or is there such a thing? error("Don't know how to split functions with sub functions of rank %d (yet)." % rank) #for indices in compute_indices(shape): # #k = offset + sum(i*s for (i,s) in izip(indices, shape[1:] + (1,))) # vs.append(v[indices]) sub_functions.append(subv) ufl_assert(actual_value_size == offset, "Logic breach in function splitting.") return tuple(sub_functions)
def split(v): """UFL operator: If v is a Coefficient or Argument in a mixed space, returns a tuple with the function components corresponding to the subelements.""" # Special case: simple element, just return function in a tuple element = v.element() if not isinstance(element, MixedElement): return (v, ) if isinstance(element, TensorElement): s = element.symmetry() if s: # FIXME: How should this be defined? Should we return one subfunction # for each value component or only for those not mapped to another? # I think split should ignore the symmetry. error("Split not implemented for symmetric tensor elements.") # Compute value size value_size = product(element.value_shape()) actual_value_size = value_size # Extract sub coefficient offset = 0 sub_functions = [] for i, e in enumerate(element.sub_elements()): shape = e.value_shape() rank = len(shape) if rank == 0: # This subelement is a scalar, always maps to a single value subv = v[offset] offset += 1 elif rank == 1: # This subelement is a vector, always maps to a sequence of values sub_size, = shape components = [v[j] for j in range(offset, offset + sub_size)] subv = as_vector(components) offset += sub_size elif rank == 2: # This subelement is a tensor, possibly with symmetries, slightly more complicated... # Size of this subvalue sub_size = product(shape) # If this subelement is a symmetric element, subtract symmetric components s = None if isinstance(e, TensorElement): s = e.symmetry() s = s or EmptyDict # If we do this, we must fix the size computation in MixedElement.__init__ as well #actual_value_size -= len(s) #sub_size -= len(s) #print s # Build list of lists of value components components = [] for ii in range(shape[0]): row = [] for jj in range(shape[1]): # Map component (i,j) through symmetry mapping c = (ii, jj) c = s.get(c, c) i, j = c # Extract component c of this subvalue from global tensor v if v.rank() == 1: # Mapping into a flattened vector k = offset + i * shape[1] + j component = v[k] #print "k, offset, i, j, shape, component", k, offset, i, j, shape, component elif v.rank() == 2: # Mapping into a concatenated tensor (is this a figment of my imagination?) error("Not implemented.") row_offset, col_offset = 0, 0 # TODO k = (row_offset + i, col_offset + j) component = v[k] row.append(component) components.append(row) # Make a matrix of the components subv = as_matrix(components) offset += sub_size else: # TODO: Handle rank > 2? Or is there such a thing? error( "Don't know how to split functions with sub functions of rank %d (yet)." % rank) #for indices in compute_indices(shape): # #k = offset + sum(i*s for (i,s) in izip(indices, shape[1:] + (1,))) # vs.append(v[indices]) sub_functions.append(subv) ufl_assert(actual_value_size == offset, "Logic breach in function splitting.") return tuple(sub_functions)