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 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 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 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 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
