def _set_function(f, self):
if self.is_integer:
#for positive integers:
if self.args == (1, S.Infinity, 1):
expr = f.expr
if not isinstance(expr, Expr):
return
x = f.variables[0]
if not expr.free_symbols - {x}:
step = expr.coeff(x)
c = expr.subs(x, 0)
if c.is_Integer and step.is_Integer and expr == step * x + c:
if self is S.Naturals:
c += step
if step > 0:
return Range(c, S.Infinity, step)
return Range(c, S.NegativeInfinity, step)
return
expr = f.expr
if not isinstance(expr, Expr):
return
n = f.variables[0]
# f(x) + c and f(-x) + c cover the same integers
# so choose the form that has the fewest negatives
c = f(0)
fx = f(n) - c
f_x = f(-n) - c
neg_count = lambda e: sum(_._coeff_isneg() for _ in Add.make_args(e))
if neg_count(f_x) < neg_count(fx):
expr = f_x + c
a = Wild('a', exclude=[n])
b = Wild('b', exclude=[n])
match = expr.match(a * n + b)
if match and match[a]:
# canonical shift
b = match[b]
if abs(match[a]) == 1:
nonint = []
for bi in Add.make_args(b):
if not bi.is_integer:
nonint.append(bi)
b = Add(*nonint)
if b.is_number and match[a].is_real:
mod = b % match[a]
reps = dict([(m, m.args[0]) for m in mod.atoms(Mod)
if not m.args[0].is_real])
mod = mod.xreplace(reps)
expr = match[a] * n + mod
else:
expr = match[a] * n + b
if expr != f.expr:
return ImageSet(Lambda(n, expr), S.Integers)
def _set_function(f, self):
expr = f.expr
if not isinstance(expr, Expr):
return
x = f.variables[0]
if not expr.free_symbols - {x}:
step = expr.coeff(x)
c = expr.subs(x, 0)
if c.is_Integer and step.is_Integer and expr == step * x + c:
if self is S.Naturals:
c += step
if step > 0:
return Range(c, S.Infinity, step)
return Range(c, S.NegativeInfinity, step)
def _set_function(f, self): # noqa:F811
expr = f.expr
if not isinstance(expr, Expr):
return
x = f.variables[0]
if not expr.free_symbols - {x}:
if expr == abs(x):
if self is S.Naturals:
return self
return S.Naturals0
step = expr.coeff(x)
c = expr.subs(x, 0)
if c.is_Integer and step.is_Integer and expr == step * x + c:
if self is S.Naturals:
c += step
if step > 0:
if step == 1:
if c == 0:
return S.Naturals0
elif c == 1:
return S.Naturals
return Range(c, oo, step)
return Range(c, -oo, step)
def _set_function(f, self): # noqa:F811
if not self:
return S.EmptySet
if not isinstance(f.expr, Expr):
return
if self.size == 1:
return FiniteSet(f(self[0]))
if f is S.IdentityFunction:
return self
x = f.variables[0]
expr = f.expr
# handle f that is linear in f's variable
if x not in expr.free_symbols or x in expr.diff(x).free_symbols:
return
if self.start.is_finite:
F = f(self.step * x + self.start) # for i in range(len(self))
else:
F = f(-self.step * x + self[-1])
F = expand_mul(F)
if F != expr:
return imageset(x, F, Range(self.size))
def test_ccode_For():
f = For(x, Range(0, 10, 2), [aug_assign(y, '*', x)])
assert ccode(f) == ("for (x = 0; x < 10; x += 2) {\n" " y *= x;\n" "}")
def test_rcode_For():
f = For(x, Range(0, 10, 2), [aug_assign(y, "*", x)])
sol = rcode(f)
assert sol == ("for (x = 0; x < 10; x += 2) {\n" " y *= x;\n" "}")