Added method of intersection for ImageSet with Integer as baseset

sympy/sets/fancysets.py

@@ -2,11 +2,12 @@
 
 from sympy.core.basic import Basic
 from sympy.core.compatibility import as_int, with_metaclass
-from sympy.sets.sets import Set, Interval, Intersection
+from sympy.sets.sets import Set, Interval, Intersection, EmptySet
 from sympy.core.singleton import Singleton, S
 from sympy.core.symbol import symbols
 from sympy.core.sympify import sympify
 from sympy.core.decorators import deprecated
+from sympy.core.function import Lambda
 
 
 class Naturals(with_metaclass(Singleton, Set)):
@@ -147,6 +148,26 @@ def _sup(self):
     def _boundary(self):
         return self
 
+    def _eval_imageset(self, f):
+        from sympy import Wild
+        expr = f.expr
+        if len(f.variables) > 1:
+            return
+        n = f.variables[0]
+
+        a = Wild('a')
+        b = Wild('b')
+
+        match = expr.match(a*n + b)
+        if match[a].is_negative:
+            expr = -expr
+
+        match = expr.match(a*n + b)
+        if match[a] is S.One and match[b].is_integer:
+            expr = expr - match[b]
+
+        return ImageSet(Lambda(n, expr), S.Integers)
+
 
 class Reals(with_metaclass(Singleton, Interval)):
 
@@ -223,6 +244,33 @@ def _contains(self, other):
     def is_iterable(self):
         return self.base_set.is_iterable
 
+    def _intersect(self, other):
+        from sympy import Dummy
+        from sympy.solvers.diophantine import diophantine
+        from sympy.sets.sets import imageset
+        if self.base_set is S.Integers:
+            if isinstance(other, ImageSet) and other.base_set is S.Integers:
+                f, g = self.lamda.expr, other.lamda.expr
+                n, m = self.lamda.variables[0], other.lamda.variables[0]
+
+                # Diophantine sorts the solutions according to the alphabetic
+                # order of the variable names, since the result should not depend
+                # on the variable name, they are replaced by the dummy variables
+                # below
+                a, b = Dummy('a'), Dummy('b')
+                f, g = f.subs(n, a), g.subs(m, b)
+                solns_set = diophantine(f - g)
+                if solns_set == set():
+                    return EmptySet()
+                solns = list(diophantine(f - g))
+                if len(solns) == 1:
+                    t = list(solns[0][0].free_symbols)[0]
+                else:
+                    return None
+
+                # since 'a' < 'b'
+                return imageset(Lambda(t, f.subs(a, solns[0][0])), S.Integers)
+
 
 @deprecated(useinstead="ImageSet", issue=7057, deprecated_since_version="0.7.4")
 def TransformationSet(*args, **kwargs):

sympy/sets/tests/test_fancysets.py

@@ -1,7 +1,7 @@
 from sympy.sets.fancysets import ImageSet, Range
-from sympy.sets.sets import FiniteSet, Interval, imageset
+from sympy.sets.sets import FiniteSet, Interval, imageset, EmptySet
 from sympy import (S, Symbol, Lambda, symbols, cos, sin, pi, oo, Basic,
-        Rational, sqrt)
+        Rational, sqrt, Eq)
 from sympy.utilities.pytest import XFAIL
 import itertools
 
@@ -168,3 +168,43 @@ def test_intersections():
     assert -5 in S.Integers.intersect(Interval(-oo, 3))
     assert all(x.is_Integer
             for x in take(10, S.Integers.intersect(Interval(3, oo)) ))
+
+
+def test_infinitely_indexed_set_1():
+    from sympy.abc import n, m, t
+    assert imageset(Lambda(n, n), S.Integers) == imageset(Lambda(m, m), S.Integers)
+
+    assert imageset(Lambda(n, 2*n), S.Integers).intersect(imageset(Lambda(m, 2*m + 1), S.Integers)) == \
+            EmptySet()
+
+    assert imageset(Lambda(n, 2*n), S.Integers).intersect(imageset(Lambda(n, 2*n + 1), S.Integers)) == \
+            EmptySet()
+
+    assert imageset(Lambda(m, 2*m), S.Integers).intersect(imageset(Lambda(n, 3*n), S.Integers)) == \
+            ImageSet(Lambda(t, 6*t), S.Integers)
+
+
+def test_infinitely_indexed_set_2():
+    from sympy import exp
+    from sympy.abc import n
+    a = Symbol('a', integer=True)
+    assert imageset(Lambda(n, n), S.Integers) == imageset(Lambda(n, n + a), S.Integers)
+    assert imageset(Lambda(n, n), S.Integers) == imageset(Lambda(n, -n + a), S.Integers)
+    assert imageset(Lambda(n, -6*n), S.Integers) == ImageSet(Lambda(n, 6*n), S.Integers)
+    assert imageset(Lambda(n, 2*n + pi), S.Integers) == ImageSet(Lambda(n, 2*n + pi), S.Integers)
+    assert imageset(Lambda(n, pi*n + pi), S.Integers) == ImageSet(Lambda(n, pi*n + pi), S.Integers)
+    assert imageset(Lambda(n, exp(n)), S.Integers) != imageset(Lambda(n, n), S.Integers)
+
+
+@XFAIL
+def test_infinitely_indexed_failed_diophantine():
+    from sympy.abc import n, m, t
+    assert imageset(Lambda(m, 2*pi*m), S.Integers).intersect(imageset(Lambda(n, 3*pi*n), S.Integers)) == \
+            ImageSet(Lambda(t, -6*pi*t), S.Integers)
+
+
+@XFAIL
+def test_infinitely_indexed_set_3():
+    from sympy.abc import n
+    assert imageset(Lambda(n, 2*n + 1), S.Integers) == imageset(Lambda(n, 2*n - 1), S.Integers)
+    assert imageset(Lambda(n, 3*n + 2), S.Integers) == imageset(Lambda(n, 3*n - 1), S.Integers)