def iet_lower_dimensions(iet):
"""
Replace all DerivedDimensions within the ``iet``'s expressions with
lower-level symbolic objects (other Dimensions or Symbols).
* Array indices involving SteppingDimensions are turned into ModuloDimensions.
Example: ``u[t+1, x] = u[t, x] + 1 >>> u[t1, x] = u[t0, x] + 1``
* Array indices involving ConditionalDimensions used are turned into
integer-division expressions.
Example: ``u[t_sub, x] = u[time, x] >>> u[time / 4, x] = u[time, x]``
"""
# Lower SteppingDimensions
for i in FindNodes(Iteration).visit(iet):
if not i.uindices:
# Be quick: avoid uselessy reconstructing nodes
continue
# In an expression, there could be `u[t+1, ...]` and `v[t+1, ...]`, where
# `u` and `v` are TimeFunction with circular time buffers (save=None) *but*
# different modulo extent. The `t+1` indices above are therefore conceptually
# different, so they will be replaced with the proper ModuloDimension through
# two different calls to `xreplace`
groups = as_mapper(i.uindices, lambda d: d.modulo)
for k, v in groups.items():
mapper = {d.origin: d for d in v}
rule = lambda i: i.function.is_TimeFunction and i.function._time_size == k
replacer = lambda i: xreplace_indices(i, mapper, rule)
iet = XSubs(replacer=replacer).visit(iet)
# Lower ConditionalDimensions
cdims = [d for d in FindSymbols('free-symbols').visit(iet)
if isinstance(d, ConditionalDimension)]
mapper = {d: IntDiv(d.index, d.factor) for d in cdims}
iet = XSubs(mapper).visit(iet)
return iet
def _lower_stepping_dims(iet):
"""
Lower SteppingDimensions: index functions involving SteppingDimensions are
turned into ModuloDimensions.
Examples
--------
u[t+1, x] = u[t, x] + 1
becomes
u[t1, x] = u[t0, x] + 1
"""
for i in FindNodes(Iteration).visit(iet):
if not i.uindices:
# Be quick: avoid uselessy reconstructing nodes
continue
# In an expression, there could be `u[t+1, ...]` and `v[t+1, ...]`, where
# `u` and `v` are TimeFunction with circular time buffers (save=None) *but*
# different modulo extent. The `t+1` indices above are therefore conceptually
# different, so they will be replaced with the proper ModuloDimension through
# two different calls to `xreplace`
mindices = [d for d in i.uindices if d.is_Modulo]
groups = as_mapper(mindices, lambda d: d.modulo)
for k, v in groups.items():
mapper = {d.origin: d for d in v}
rule = lambda i: i.function.is_TimeFunction and i.function._time_size == k
replacer = lambda i: xreplace_indices(i, mapper, rule)
iet = XSubs(replacer=replacer).visit(iet)
return iet
def _lower_conditional_dims(iet):
"""
Lower ConditionalDimensions: index functions involving ConditionalDimensions
are turned into integer-division expressions.
Examples
--------
u[t_sub, x] = u[time, x]
becomes
u[time / 4, x] = u[time, x]
"""
cdims = [d for d in FindSymbols('free-symbols').visit(iet)
if isinstance(d, ConditionalDimension)]
mapper = {d: IntDiv(d.index, d.factor) for d in cdims}
iet = XSubs(mapper).visit(iet)
return iet
