diff --git a/devito/ir/support/space.py b/devito/ir/support/space.py index a1c0f7b7a8..1c760128b5 100644 --- a/devito/ir/support/space.py +++ b/devito/ir/support/space.py @@ -912,8 +912,30 @@ def switch(self, d0, d1, direction=None): return IterationSpace(intervals, sub_iterators, directions) - def translate(self, d, v0=0, v1=None): - intervals = self.intervals.translate(d, v0, v1) + def translate(self, maybe_dim, v0=0, v1=None): + """ + Translate the IterationSpace along a given set of Dimensions by + given offsets. + + This method can be invoked in three different flavours: + + * `ispace.translate(dim, v0, v1)`; + * `ispace.translate(dims, v0, v1)` + Same as before, but applies the translation to all Dimensions + in `dims`; + * `ispace.translate({d0: (v0, v1), d1: v2, ...})` + This is essentially a shortcut to avoid calling `translate` in + a loop at the call site. + """ + if not maybe_dim: + return self + elif isinstance(maybe_dim, dict): + intervals = self.intervals + for d, v in maybe_dim.items(): + intervals = intervals.translate(d, *as_tuple(v)) + else: + intervals = self.intervals.translate(maybe_dim, v0, v1) + return IterationSpace(intervals, self.sub_iterators, self.directions) def reset(self):