sympy.py

   1 BASIC   = 0
   2 SYMBOL  = 1
   3 ADD     = 2
   4 MUL     = 3
   5 POW     = 4
   6 INTEGER = 5
   7
   8 def hash_seq(args):
   9     # make this more robust:
  10     m = 2
  11     for x in args:
  12         m = hash(m + 1001 ^ hash(x))
  13     return m
  14
  15 class Basic(object):
  16
  17     def __new__(cls, type, args):
  18         obj = object.__new__(cls)
  19         obj.type = type
  20         obj._args = tuple(args)
  21         return obj
  22
  23     def __repr__(self):
  24         return str(self)
  25
  26     def __hash__(self):
  27         return hash_seq(self.args)
  28
  29     @property
  30     def args(self):
  31         return self._args
  32
  33     def as_coeff_rest(self):
  34         return (Integer(1), self)
  35
  36     def as_base_exp(self):
  37         return (self, Integer(1))
  38
  39     def __add__(x, y):
  40         return Add((x, y))
  41
  42     def __radd__(x, y):
  43         return x.__add__(y)
  44
  45     def __sub__(x, y):
  46         return Add((x, -y))
  47
  48     def __rsub__(x, y):
  49         return Add((y, -x))
  50
  51     def __mul__(x, y):
  52         return Mul((x, y))
  53
  54     def __rmul__(x, y):
  55         return Mul((y, x))
  56
  57     def __div__(x, y):
  58         return Mul((x, Pow((y, Integer(-1)))))
  59
  60     def __rdiv__(x, y):
  61         return Mul((y, Pow((x, Integer(-1)))))
  62
  63     def __pow__(x, y):
  64         return Pow((x, y))
  65
  66     def __rpow__(x, y):
  67         return Pow((y, x))
  68
  69     def __neg__(x):
  70         return Mul((Integer(-1), x))
  71
  72     def __pos__(x):
  73         return x
  74
  75
  76 class Integer(Basic):
  77
  78     def __new__(cls, i):
  79         obj = Basic.__new__(cls, INTEGER, [])
  80         obj.i = i
  81         return obj
  82
  83     def __eq__(self, o):
  84         o = sympify(o)
  85         if o.type == INTEGER:
  86             return self.i == o.i
  87         else:
  88             return False
  89
  90     def __str__(self):
  91         return str(self.i)
  92
  93     def __add__(self, o):
  94         o = sympify(o)
  95         if o.type == INTEGER:
  96             return Integer(self.i+o.i)
  97         return Basic.__add__(self, o)
  98
  99     def __mul__(self, o):
 100         o = sympify(o)
 101         if o.type == INTEGER:
 102             return Integer(self.i*o.i)
 103         return Basic.__mul__(self, o)
 104
 105
 106 class Symbol(Basic):
 107
 108     def __new__(cls, name):
 109         obj = Basic.__new__(cls, SYMBOL, [])
 110         obj.name = name
 111         return obj
 112
 113     def __hash__(self):
 114         return hash(self.name)
 115
 116     def __str__(self):
 117         return self.name
 118
 119
 120 class Add(Basic):
 121
 122     def __new__(cls, args, canonicalize=True):
 123         if canonicalize == False:
 124             obj = Basic.__new__(cls, ADD, args)
 125             return obj
 126         args = [sympify(x) for x in args]
 127         return Add.canonicalize(args)
 128
 129     @classmethod
 130     def canonicalize(cls, args):
 131         d = {}
 132         for a in args:
 133             if a.type == ADD:
 134                 for b in a.args:
 135                     coeff, key = b.as_coeff_rest()
 136                     if key in d:
 137                         d[key] += coeff
 138                     else:
 139                         d[key] = coeff
 140             else:
 141                 coeff, key = a.as_coeff_rest()
 142                 if key in d:
 143                     d[key] += coeff
 144                 else:
 145                     d[key] = coeff
 146         print d
 147         args = []
 148         for a, b in d.iteritems():
 149             args.append(Mul((a, b)))
 150
 151         return Add(args, False)
 152
 153     def __str__(self):
 154         s = str(self.args[0])
 155         if self.args[0].type == ADD:
 156             s = "(%s)" % str(s)
 157         for x in self.args[1:]:
 158             s = "%s + %s" % (s, str(x))
 159             if x.type == ADD:
 160                 s = "(%s)" % s
 161         return s
 162
 163 class Mul(Basic):
 164
 165     def __new__(cls, args, canonicalize=True):
 166         if canonicalize == False:
 167             obj = Basic.__new__(cls, MUL, args)
 168             return obj
 169         args = [sympify(x) for x in args]
 170         return Mul.canonicalize(args)
 171
 172     @classmethod
 173     def canonicalize(cls, args):
 174         d = {}
 175         num = Integer(1)
 176         for a in args:
 177             if a.type == INTEGER:
 178                 num *= a
 179             elif a.type == MUL:
 180                 for b in a.args:
 181                     coeff, key = b.as_base_exp()
 182                     if key in d:
 183                         d[key] += coeff
 184                     else:
 185                         d[key] = coeff
 186             else:
 187                 coeff, key = a.as_base_exp()
 188                 if key in d:
 189                     d[key] += coeff
 190                 else:
 191                     d[key] = coeff
 192         if num.i == 0 or len(d)==0:
 193             return num
 194         args = []
 195         for a, b in d.iteritems():
 196             args.append(Pow((b, a)))
 197         if num.i != 1:
 198             args.insert(0, num)
 199         if len(args) == 1:
 200             return args[0]
 201         else:
 202             return Mul(args, False)
 203
 204     def __hash__(self):
 205         a = list(self.args[:])
 206         a.sort(key=hash)
 207         return hash_seq(a)
 208
 209     def __eq__(self, o):
 210         o = sympify(o)
 211         if o.type == MUL:
 212             a = list(self.args[:])
 213             a.sort(key=hash)
 214             b = list(o.args[:])
 215             b.sort(key=hash)
 216             return a == b
 217         else:
 218             return False
 219
 220
 221     def as_coeff_rest(self):
 222         if self.args[0].type == INTEGER:
 223             return (self.args[0], Mul(self.args[1:]))
 224         return (Integer(1), self)
 225
 226     def __str__(self):
 227         s = str(self.args[0])
 228         if self.args[0].type == MUL:
 229             s = "(%s)" % str(s)
 230         for x in self.args[1:]:
 231             s = "%s*%s" % (s, str(x))
 232             if x.type == MUL:
 233                 s = "(%s)" % s
 234         return s
 235
 236 class Pow(Basic):
 237
 238     def __new__(cls, args, canonicalize=True):
 239         if canonicalize == False:
 240             obj = Basic.__new__(cls, POW, args)
 241             return obj
 242         args = [sympify(x) for x in args]
 243         return Pow.canonicalize(args)
 244
 245     @classmethod
 246     def canonicalize(cls, args):
 247         base, exp = args
 248         if exp.type == INTEGER:
 249             if exp.i == 0:
 250                 return Integer(1)
 251             if exp.i == 1:
 252                 return base
 253         return Pow(args, False)
 254
 255     def __str__(self):
 256         s = str(self.args[0])
 257         if self.args[0].type == ADD:
 258             s = "(%s)" % s
 259         if self.args[1].type == ADD:
 260             s = "%s^(%s)" % (s, str(self.args[1]))
 261         else:
 262             s = "%s^%s" % (s, str(self.args[1]))
 263
 264         return s
 265
 266 def sympify(x):
 267     if isinstance(x, int):
 268         return Integer(x)
 269     return x