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