def __init__(self, gc, i):
self.scalar = True
component = (i, )
shape = gc.ufl_shape
all_indices, _, _ = create_slice_indices(component, shape,
gc.ufl_free_indices)
mi = MultiIndex(all_indices)
Indexed.__init__(self, gc, mi)
if gc.gf.scalar:
assert i == 0
self.__impl__ = gc.gf
self.gf = gc
else:
self.__impl__ = gc.gf[i]
self.gf = None
def indexed(self, o):
A, jj = o.operands()
A2, Ap = self.visit(A)
o = self.reuse_if_possible(o, A2, jj)
if isinstance(Ap, Zero):
op = self._make_zero_diff(o)
else:
r = Ap.rank() - len(jj)
if r:
ii = indices(r)
op = Indexed(Ap, jj._indices + ii)
op = as_tensor(op, ii)
else:
op = Indexed(Ap, jj)
return (o, op)
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 expr(self, o, *ops):
r = self.reuse_if_possible(o, *ops)
c = self.components.peek()
if c:
#if isinstance(r, ListTensor):
# pass # TODO: If c has FixedIndex objects, extract subtensor, or evt. move this functionality from ListTensor.__getitem__ to Indexed.__new__
# Take component
r = Indexed(r, c)
return r
def indexed(self, o, *dops):
if o in self._mapping:
return self._mapping[o]
else:
assert len(dops) == 2
assert isinstance(dops[1], MultiIndex)
if dops[0] in self._mapping:
replaced_ufl_operand_0 = self._mapping[dops[0]]
else:
replaced_ufl_operand_0 = map_integrand_dags(self, dops[0])
return Indexed(replaced_ufl_operand_0, dops[1])
def _getitem(self, key):
# Analyse key, getting rid of slices and the ellipsis
r = self.rank()
indices, axis_indices = analyse_key(key, r)
# Special case for foo[...] => foo
if len(indices) == len(axis_indices):
return self
# Special case for simplifying ({ai}_i)[i] -> ai
if isinstance(self, ComponentTensor):
if tuple(indices) == tuple(self._indices):
return self._expression
# Index self, yielding scalar valued expressions
a = Indexed(self, indices)
# Make a tensor from components designated by axis indices
if axis_indices:
a = as_tensor(a, axis_indices)
# TODO: Should we apply IndexSum or as_tensor first?
# Apply sum for each repeated index
ri = repeated_indices(self.free_indices() + indices)
for i in ri:
a = IndexSum(a, i)
# Check for zero (last so we can get indices etc from a)
if isinstance(self, Zero):
shape = a.shape()
fi = a.free_indices()
idims = subdict(a.index_dimensions(), fi)
a = Zero(shape, fi, idims)
return a
def terminal(self, o):
r = o
c = self.components.peek()
if c:
r = Indexed(r, c)
return r