def test_indexing_to_component(): assert 0 == flatten_multiindex((), shape_to_strides(())) assert 0 == flatten_multiindex((0,), shape_to_strides((2,))) assert 1 == flatten_multiindex((1,), shape_to_strides((2,))) assert 3 == flatten_multiindex((1, 1), shape_to_strides((2, 2))) for i in range(5): for j in range(3): for k in range(2): assert 15*k+5*j+i == flatten_multiindex((k, j, i), shape_to_strides((2, 3, 5)))
def extract_subelement_component(self, i): """Extract direct subelement index and subelement relative component index for a given component index.""" if isinstance(i, int): i = (i, ) self._check_component(i) # Select between indexing modes if len(self.value_shape()) == 1: # Indexing into a long vector of flattened subelement # shapes j, = i # Find subelement for this index for sub_element_index, e in enumerate(self._sub_elements): sh = e.value_shape() si = product(sh) if j < si: break j -= si if j < 0: error("Moved past last value component!") # Convert index into a shape tuple st = shape_to_strides(sh) component = unflatten_index(j, st) else: # Indexing into a multidimensional tensor where subelement # index is first axis sub_element_index = i[0] if sub_element_index >= len(self._sub_elements): error("Illegal component index (dimension %d)." % sub_element_index) component = i[1:] return (sub_element_index, component)
def extract_subelement_reference_component(self, i): """Extract direct subelement index and subelement relative reference_component index for a given reference_component index.""" if isinstance(i, int): i = (i, ) self._check_reference_component(i) # Select between indexing modes assert len(self.reference_value_shape()) == 1 # Indexing into a long vector of flattened subelement shapes j, = i # Find subelement for this index for sub_element_index, e in enumerate(self._sub_elements): sh = e.reference_value_shape() si = product(sh) if j < si: break j -= si if j < 0: error("Moved past last value reference_component!") # Convert index into a shape tuple st = shape_to_strides(sh) reference_component = unflatten_index(j, st) return (sub_element_index, reference_component)
def extract_subelement_component(self, i): """Extract direct subelement index and subelement relative component index for a given component index.""" if isinstance(i, int): i = (i,) self._check_component(i) # Select between indexing modes if len(self.value_shape()) == 1: # Indexing into a long vector of flattened subelement # shapes j, = i # Find subelement for this index for sub_element_index, e in enumerate(self._sub_elements): sh = e.value_shape() si = product(sh) if j < si: break j -= si if j < 0: error("Moved past last value component!") # Convert index into a shape tuple st = shape_to_strides(sh) component = unflatten_index(j, st) else: # Indexing into a multidimensional tensor where subelement # index is first axis sub_element_index = i[0] if sub_element_index >= len(self._sub_elements): error("Illegal component index (dimension %d)." % sub_element_index) component = i[1:] return (sub_element_index, component)
def extract_subelement_reference_component(self, i): """Extract direct subelement index and subelement relative reference_component index for a given reference_component index.""" if isinstance(i, int): i = (i,) self._check_reference_component(i) # Select between indexing modes assert len(self.reference_value_shape()) == 1 # Indexing into a long vector of flattened subelement shapes j, = i # Find subelement for this index for sub_element_index, e in enumerate(self._sub_elements): sh = e.reference_value_shape() si = product(sh) if j < si: break j -= si if j < 0: error("Moved past last value reference_component!") # Convert index into a shape tuple st = shape_to_strides(sh) reference_component = unflatten_index(j, st) return (sub_element_index, reference_component)
def test_flatten_multiindex_to_multiindex(): sh = (2, 3, 5) strides = shape_to_strides(sh) for i in range(sh[2]): for j in range(sh[1]): for k in range(sh[0]): index = (k, j, i) c = flatten_multiindex(index, strides) index2 = unflatten_index(c, strides) assert index == index2
def product(self, o, ops): if len(ops) != 2: error("Expecting two operands.") # Get the simple cases out of the way if len(ops[0]) == 1: # True scalar * something a, = ops[0] return [Product(a, b) for b in ops[1]] elif len(ops[1]) == 1: # Something * true scalar b, = ops[1] return [Product(a, b) for a in ops[0]] # Neither of operands are true scalars, this is the tricky part o0, o1 = o.ufl_operands # Get shapes and index shapes fi = o.ufl_free_indices fi0 = o0.ufl_free_indices fi1 = o1.ufl_free_indices fid = o.ufl_index_dimensions fid0 = o0.ufl_index_dimensions fid1 = o1.ufl_index_dimensions # Need to map each return component to one component of o0 and one component of o1. indices = compute_indices(fid) # Compute which component of o0 is used in component (comp,ind) of o # Compute strides within free index spaces ist0 = shape_to_strides(fid0) ist1 = shape_to_strides(fid1) # Map o0 and o1 indices to o indices indmap0 = [fi.index(i) for i in fi0] indmap1 = [fi.index(i) for i in fi1] indks = [(flatten_multiindex([ind[i] for i in indmap0], ist0), flatten_multiindex([ind[i] for i in indmap1], ist1)) for ind in indices] # Build products for scalar components results = [Product(ops[0][k0], ops[1][k1]) for k0, k1 in indks] return results
def ufl_evaluate(self, x, component, derivatives): """Function used by ufl to evaluate the Expression""" assert derivatives == () # TODO: Handle derivatives if component: shape = self.ufl_shape assert len(shape) == len(component) value_size = product(shape) index = flatten_multiindex(component, shape_to_strides(shape)) values = numpy.zeros(value_size) # FIXME: use a function with a return value self(*x, values=values) return values[index] else: # Scalar evaluation return self(*x)
def ufl_evaluate(self, x, component, derivatives): """Function used by ufl to evaluate the Function""" import numpy import ufl assert derivatives == () # TODO: Handle derivatives if component: shape = self.shape() assert len(shape) == len(component) value_size = ufl.common.product(shape) index = flatten_multiindex(component, shape_to_strides(shape)) values = numpy.zeros(value_size) self(*x, values=values) return values[index] else: # Scalar evaluation return self(*x)
def symmetry(self): """Return the symmetry dict, which is a mapping :math:`c_0 \\to c_1` meaning that component :math:`c_0` is represented by component :math:`c_1`. A component is a tuple of one or more ints.""" # Build symmetry map from symmetries of subelements sm = {} # Base index of the current subelement into mixed value j = 0 for e in self._sub_elements: sh = e.value_shape() st = shape_to_strides(sh) # Map symmetries of subelement into index space of this # element for c0, c1 in e.symmetry().items(): j0 = flatten_multiindex(c0, st) + j j1 = flatten_multiindex(c1, st) + j sm[(j0,)] = (j1,) # Update base index for next element j += product(sh) if j != product(self.value_shape()): error("Size mismatch in symmetry algorithm.") return sm or EmptyDict
def symmetry(self): """Return the symmetry dict, which is a mapping :math:`c_0 \\to c_1` meaning that component :math:`c_0` is represented by component :math:`c_1`. A component is a tuple of one or more ints.""" # Build symmetry map from symmetries of subelements sm = {} # Base index of the current subelement into mixed value j = 0 for e in self._sub_elements: sh = e.value_shape() st = shape_to_strides(sh) # Map symmetries of subelement into index space of this # element for c0, c1 in e.symmetry().items(): j0 = flatten_multiindex(c0, st) + j j1 = flatten_multiindex(c1, st) + j sm[(j0, )] = (j1, ) # Update base index for next element j += product(sh) if j != product(self.value_shape()): error("Size mismatch in symmetry algorithm.") return sm or EmptyDict
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 map_component_tensor_arg_components(tensor): """Build integer list mapping between flattended components of tensor and its underlying indexed subexpression.""" assert isinstance(tensor, ComponentTensor) # AKA tensor = as_tensor(indexed, multiindex) indexed, multiindex = tensor.ufl_operands e1 = indexed e2 = tensor # e2 = as_tensor(e1, multiindex) mi = [i for i in multiindex if isinstance(i, Index)] # Get tensor and index shapes sh1 = e1.ufl_shape # (sh)ape of e1 sh2 = e2.ufl_shape # (sh)ape of e2 fi1 = e1.ufl_free_indices # (f)ree (i)ndices of e1 fi2 = e2.ufl_free_indices # ... fid1 = e1.ufl_index_dimensions # (f)ree (i)ndex (d)imensions of e1 fid2 = e2.ufl_index_dimensions # ... # Compute total shape (tsh) of e1 and e2 tsh1 = sh1 + fid1 tsh2 = sh2 + fid2 r1 = len(tsh1) # 'total rank' or e1 r2 = len(tsh2) # ... str1 = shape_to_strides(tsh1) assert not sh1 assert sh2 assert len(mi) == len(multiindex) assert product(tsh1) == product(tsh2) assert fi1 assert all(i in fi1 for i in fi2) nmui = len(multiindex) assert nmui == len(sh2) # Build map from fi2/fid2 position (-offset nmui) to fi1/fid1 position p2_to_p1_map = [None] * r2 for k, i in enumerate(fi2): p2_to_p1_map[k + nmui] = fi1.index(i) # Build map from fi1/fid1 position to mi position for k, i in enumerate(mi): p2_to_p1_map[k] = fi1.index(mi[k].count()) # Build map from flattened e1 component to flattened e2 component perm2 = compute_indices(tsh2) ni2 = product(tsh2) # Situation: e2 = as_tensor(e1, mi) d2 = [None] * ni2 p1 = [None] * r1 for c2, p2 in enumerate(perm2): for k2, k1 in enumerate(p2_to_p1_map): p1[k1] = p2[k2] c1 = flatten_multiindex(p1, str1) d2[c2] = c1 # Consistency checks assert all(isinstance(x, int) for x in d2) assert len(set(d2)) == len(d2) return d2
def map_component_tensor_arg_components(tensor): """Build a map from flattened components to subexpression. Builds integer list mapping between flattended components of tensor and its underlying indexed subexpression.""" assert isinstance(tensor, ComponentTensor) # AKA tensor = as_tensor(indexed, multiindex) indexed, multiindex = tensor.ufl_operands e1 = indexed e2 = tensor # e2 = as_tensor(e1, multiindex) mi = [i for i in multiindex if isinstance(i, Index)] # Get tensor and index shapes sh1 = e1.ufl_shape # (sh)ape of e1 sh2 = e2.ufl_shape # (sh)ape of e2 fi1 = e1.ufl_free_indices # (f)ree (i)ndices of e1 fi2 = e2.ufl_free_indices # ... fid1 = e1.ufl_index_dimensions # (f)ree (i)ndex (d)imensions of e1 fid2 = e2.ufl_index_dimensions # ... # Compute total shape (tsh) of e1 and e2 tsh1 = sh1 + fid1 tsh2 = sh2 + fid2 r1 = len(tsh1) # 'total rank' or e1 r2 = len(tsh2) # ... str1 = shape_to_strides(tsh1) assert not sh1 assert sh2 assert len(mi) == len(multiindex) assert ufl.product(tsh1) == ufl.product(tsh2) assert fi1 assert all(i in fi1 for i in fi2) nmui = len(multiindex) assert nmui == len(sh2) # Build map from fi2/fid2 position (-offset nmui) to fi1/fid1 position p2_to_p1_map = [None] * r2 for k, i in enumerate(fi2): p2_to_p1_map[k + nmui] = fi1.index(i) # Build map from fi1/fid1 position to mi position for k, i in enumerate(mi): p2_to_p1_map[k] = fi1.index(mi[k].count()) # Build map from flattened e1 component to flattened e2 component perm2 = compute_indices(tsh2) ni2 = ufl.product(tsh2) # Situation: e2 = as_tensor(e1, mi) d2 = [None] * ni2 p1 = [None] * r1 for c2, p2 in enumerate(perm2): for k2, k1 in enumerate(p2_to_p1_map): p1[k1] = p2[k2] c1 = flatten_multiindex(p1, str1) d2[c2] = c1 # Consistency checks assert all(isinstance(x, int) for x in d2) assert len(set(d2)) == len(d2) return d2
def map_indexed_arg_components(indexed): """Build a map from flattened components to subexpression. Builds integer list mapping between flattened components of indexed expression and its underlying tensor-valued subexpression.""" assert isinstance(indexed, Indexed) # AKA indexed = tensor[multiindex] tensor, multiindex = indexed.ufl_operands # AKA e1 = e2[multiindex] # (this renaming is historical, but kept for consistency with all the variables *1,*2 below) e2 = tensor e1 = indexed # Get tensor and index shape sh1 = e1.ufl_shape sh2 = e2.ufl_shape fi1 = e1.ufl_free_indices fi2 = e2.ufl_free_indices fid1 = e1.ufl_index_dimensions fid2 = e2.ufl_index_dimensions # Compute regular and total shape tsh1 = sh1 + fid1 tsh2 = sh2 + fid2 # r1 = len(tsh1) r2 = len(tsh2) # str1 = shape_to_strides(tsh1) str2 = shape_to_strides(tsh2) assert not sh1 assert sh2 # Must have shape to be indexed in the first place assert ufl.product(tsh1) <= ufl.product(tsh2) # Build map from fi2/fid2 position (-offset nmui) to fi1/fid1 position ind2_to_ind1_map = [None] * len(fi2) for k, i in enumerate(fi2): ind2_to_ind1_map[k] = fi1.index(i) # Build map from fi1/fid1 position to mi position nmui = len(multiindex) multiindex_to_ind1_map = [None] * nmui for k, i in enumerate(multiindex): if isinstance(i, Index): multiindex_to_ind1_map[k] = fi1.index(i.count()) # Build map from flattened e1 component to flattened e2 component perm1 = compute_indices(tsh1) ni1 = ufl.product(tsh1) # Situation: e1 = e2[mi] d1 = [None] * ni1 p2 = [None] * r2 assert len(sh2) == nmui for k, i in enumerate(multiindex): if isinstance(i, FixedIndex): p2[k] = int(i) for c1, p1 in enumerate(perm1): for k, i in enumerate(multiindex): if isinstance(i, Index): p2[k] = p1[multiindex_to_ind1_map[k]] for k, i in enumerate(ind2_to_ind1_map): p2[nmui + k] = p1[i] c2 = flatten_multiindex(p2, str2) d1[c1] = c2 # Consistency checks assert all(isinstance(x, int) for x in d1) assert len(set(d1)) == len(d1) return d1
def test_shape_to_strides(): assert () == shape_to_strides(()) assert (1,) == shape_to_strides((3,)) assert (2, 1) == shape_to_strides((3, 2)) assert (4, 1) == shape_to_strides((3, 4)) assert (12, 4, 1) == shape_to_strides((6, 3, 4))
def test_index_flattening(): from ufl.utils.indexflattening import (shape_to_strides, flatten_multiindex, unflatten_index) # Scalar shape s = () st = shape_to_strides(s) assert st == () c = () q = flatten_multiindex(c, st) c2 = unflatten_index(q, st) assert q == 0 assert c2 == () # Vector shape s = (2,) st = shape_to_strides(s) assert st == (1,) for i in range(s[0]): c = (i,) q = flatten_multiindex(c, st) c2 = unflatten_index(q, st) assert c == c2 # Tensor shape s = (2, 3) st = shape_to_strides(s) assert st == (3, 1) for i in range(s[0]): for j in range(s[1]): c = (i, j) q = flatten_multiindex(c, st) c2 = unflatten_index(q, st) assert c == c2 # Rank 3 tensor shape s = (2, 3, 4) st = shape_to_strides(s) assert st == (12, 4, 1) for i in range(s[0]): for j in range(s[1]): for k in range(s[2]): c = (i, j, k) q = flatten_multiindex(c, st) c2 = unflatten_index(q, st) assert c == c2 # Taylor-Hood example: # pressure element is index 3: c = (3,) # get flat index: i = flatten_multiindex(c, shape_to_strides((4,))) # remove offset: i -= 3 # map back to scalar component: c2 = unflatten_index(i, shape_to_strides(())) assert () == c2 # vector element y-component is index 1: c = (1,) # get flat index: i = flatten_multiindex(c, shape_to_strides((4,))) # remove offset: i -= 0 # map back to vector component: c2 = unflatten_index(i, shape_to_strides((3,))) assert (1,) == c2 # Try a tensor/vector element: mixed_shape = (6,) ts = (2, 2) vs = (2,) offset = 4 # product(ts) # vector element y-component is index offset+1: c = (offset + 1,) # get flat index: i = flatten_multiindex(c, shape_to_strides(mixed_shape)) # remove offset: i -= offset # map back to vector component: c2 = unflatten_index(i, shape_to_strides(vs)) assert (1,) == c2 for k in range(4): # tensor element (1,1)-component is index 3: c = (k,) # get flat index: i = flatten_multiindex(c, shape_to_strides(mixed_shape)) # remove offset: i -= 0 # map back to tensor component: c2 = unflatten_index(i, shape_to_strides(ts)) assert (k//2, k % 2) == c2
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 test_shape_to_strides(): assert () == shape_to_strides(()) assert (1, ) == shape_to_strides((3, )) assert (2, 1) == shape_to_strides((3, 2)) assert (4, 1) == shape_to_strides((3, 4)) assert (12, 4, 1) == shape_to_strides((6, 3, 4))
def map_indexed_arg_components(indexed): """Build integer list mapping between flattended components of indexed expression and its underlying tensor-valued subexpression.""" assert isinstance(indexed, Indexed) # AKA indexed = tensor[multiindex] tensor, multiindex = indexed.ufl_operands # AKA e1 = e2[multiindex] # (this renaming is historical, but kept for consistency with all the variables *1,*2 below) e2 = tensor e1 = indexed # Get tensor and index shape sh1 = e1.ufl_shape sh2 = e2.ufl_shape fi1 = e1.ufl_free_indices fi2 = e2.ufl_free_indices fid1 = e1.ufl_index_dimensions fid2 = e2.ufl_index_dimensions # Compute regular and total shape tsh1 = sh1 + fid1 tsh2 = sh2 + fid2 # r1 = len(tsh1) r2 = len(tsh2) # str1 = shape_to_strides(tsh1) str2 = shape_to_strides(tsh2) assert not sh1 assert sh2 # Must have shape to be indexed in the first place assert product(tsh1) <= product(tsh2) # Build map from fi2/fid2 position (-offset nmui) to fi1/fid1 position ind2_to_ind1_map = [None] * len(fi2) for k, i in enumerate(fi2): ind2_to_ind1_map[k] = fi1.index(i) # Build map from fi1/fid1 position to mi position nmui = len(multiindex) multiindex_to_ind1_map = [None] * nmui for k, i in enumerate(multiindex): if isinstance(i, Index): multiindex_to_ind1_map[k] = fi1.index(i.count()) # Build map from flattened e1 component to flattened e2 component perm1 = compute_indices(tsh1) ni1 = product(tsh1) # Situation: e1 = e2[mi] d1 = [None] * ni1 p2 = [None] * r2 assert len(sh2) == nmui for k, i in enumerate(multiindex): if isinstance(i, FixedIndex): p2[k] = int(i) for c1, p1 in enumerate(perm1): for k, i in enumerate(multiindex): if isinstance(i, Index): p2[k] = p1[multiindex_to_ind1_map[k]] for k, i in enumerate(ind2_to_ind1_map): p2[nmui + k] = p1[i] c2 = flatten_multiindex(p2, str2) d1[c1] = c2 # Consistency checks assert all(isinstance(x, int) for x in d1) assert len(set(d1)) == len(d1) return d1
def test_index_flattening(): from ufl.utils.indexflattening import (shape_to_strides, flatten_multiindex, unflatten_index) # Scalar shape s = () st = shape_to_strides(s) assert st == () c = () q = flatten_multiindex(c, st) c2 = unflatten_index(q, st) assert q == 0 assert c2 == () # Vector shape s = (2, ) st = shape_to_strides(s) assert st == (1, ) for i in range(s[0]): c = (i, ) q = flatten_multiindex(c, st) c2 = unflatten_index(q, st) assert c == c2 # Tensor shape s = (2, 3) st = shape_to_strides(s) assert st == (3, 1) for i in range(s[0]): for j in range(s[1]): c = (i, j) q = flatten_multiindex(c, st) c2 = unflatten_index(q, st) assert c == c2 # Rank 3 tensor shape s = (2, 3, 4) st = shape_to_strides(s) assert st == (12, 4, 1) for i in range(s[0]): for j in range(s[1]): for k in range(s[2]): c = (i, j, k) q = flatten_multiindex(c, st) c2 = unflatten_index(q, st) assert c == c2 # Taylor-Hood example: # pressure element is index 3: c = (3, ) # get flat index: i = flatten_multiindex(c, shape_to_strides((4, ))) # remove offset: i -= 3 # map back to scalar component: c2 = unflatten_index(i, shape_to_strides(())) assert () == c2 # vector element y-component is index 1: c = (1, ) # get flat index: i = flatten_multiindex(c, shape_to_strides((4, ))) # remove offset: i -= 0 # map back to vector component: c2 = unflatten_index(i, shape_to_strides((3, ))) assert (1, ) == c2 # Try a tensor/vector element: mixed_shape = (6, ) ts = (2, 2) vs = (2, ) offset = 4 # product(ts) # vector element y-component is index offset+1: c = (offset + 1, ) # get flat index: i = flatten_multiindex(c, shape_to_strides(mixed_shape)) # remove offset: i -= offset # map back to vector component: c2 = unflatten_index(i, shape_to_strides(vs)) assert (1, ) == c2 for k in range(4): # tensor element (1,1)-component is index 3: c = (k, ) # get flat index: i = flatten_multiindex(c, shape_to_strides(mixed_shape)) # remove offset: i -= 0 # map back to tensor component: c2 = unflatten_index(i, shape_to_strides(ts)) assert (k // 2, k % 2) == c2