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     """
  10     Hash of a sequence, that *depends* on the order of elements.
  11     """
  12     # make this more robust:
  13     m = 2
  14     for x in args:
  15         m = hash(m + 1001 ^ hash(x))
  16     return m
  17
  18 def compare_lists(a, b):
  19     """
  20     Compare two sequences.
  21
  22     Sequences are equal even with a *different* order of elements.
  23     """
  24
  25     return set(a) == set(b)
  26
  27 class Basic(object):
  28
  29     def __new__(cls, type, args):
  30         obj = object.__new__(cls)
  31         obj.type = type
  32         obj._args = tuple(args)
  33         return obj
  34
  35     def __repr__(self):
  36         return str(self)
  37
  38     def __hash__(self):
  39         return hash_seq(self.args)
  40
  41     @property
  42     def args(self):
  43         return self._args
  44
  45     def as_coeff_rest(self):
  46         return (Integer(1), self)
  47
  48     def as_base_exp(self):
  49         return (self, Integer(1))
  50
  51     def expand(self):
  52         return self
  53
  54     def __add__(x, y):
  55         return Add((x, y))
  56
  57     def __radd__(x, y):
  58         return x.__add__(y)
  59
  60     def __sub__(x, y):
  61         return Add((x, -y))
  62
  63     def __rsub__(x, y):
  64         return Add((y, -x))
  65
  66     def __mul__(x, y):
  67         return Mul((x, y))
  68
  69     def __rmul__(x, y):
  70         return Mul((y, x))
  71
  72     def __div__(x, y):
  73         return Mul((x, Pow((y, Integer(-1)))))
  74
  75     def __rdiv__(x, y):
  76         return Mul((y, Pow((x, Integer(-1)))))
  77
  78     def __pow__(x, y):
  79         return Pow((x, y))
  80
  81     def __rpow__(x, y):
  82         return Pow((y, x))
  83
  84     def __neg__(x):
  85         return Mul((Integer(-1), x))
  86
  87     def __pos__(x):
  88         return x
  89
  90     def __ne__(self, x):
  91         return not self.__eq__(x)
  92
  93     def __eq__(self, o):
  94         o = sympify(o)
  95         if o.type == self.type:
  96             return self.args == o.args
  97         else:
  98             return False
  99
 100
 101 class Integer(Basic):
 102
 103     def __new__(cls, i):
 104         obj = Basic.__new__(cls, INTEGER, [])
 105         obj.i = i
 106         return obj
 107
 108     def __hash__(self):
 109         return hash(self.i)
 110
 111     def __eq__(self, o):
 112         o = sympify(o)
 113         if o.type == INTEGER:
 114             return self.i == o.i
 115         else:
 116             return False
 117
 118     def __str__(self):
 119         return str(self.i)
 120
 121     def __add__(self, o):
 122         o = sympify(o)
 123         if o.type == INTEGER:
 124             return Integer(self.i+o.i)
 125         return Basic.__add__(self, o)
 126
 127     def __mul__(self, o):
 128         o = sympify(o)
 129         if o.type == INTEGER:
 130             return Integer(self.i*o.i)
 131         return Basic.__mul__(self, o)
 132
 133
 134 class Symbol(Basic):
 135
 136     def __new__(cls, name):
 137         obj = Basic.__new__(cls, SYMBOL, [])
 138         obj.name = name
 139         return obj
 140
 141     def __hash__(self):
 142         return hash(self.name)
 143
 144     def __eq__(self, o):
 145         o = sympify(o)
 146         if o.type == SYMBOL:
 147             return self.name == o.name
 148         return False
 149
 150     def __str__(self):
 151         return self.name
 152
 153
 154 class Add(Basic):
 155
 156     def __new__(cls, args, canonicalize=True):
 157         if canonicalize == False:
 158             obj = Basic.__new__(cls, ADD, args)
 159             return obj
 160         args = [sympify(x) for x in args]
 161         return Add.canonicalize(args)
 162
 163     @classmethod
 164     def canonicalize(cls, args):
 165         use_glib = 0
 166         if use_glib:
 167             from csympy import HashTable
 168             d = HashTable()
 169         else:
 170             d = {}
 171         num = Integer(0)
 172         for a in args:
 173             if a.type == INTEGER:
 174                 num += a
 175             elif a.type == ADD:
 176                 for b in a.args:
 177                     if b.type == INTEGER:
 178                         num += b
 179                     else:
 180                         coeff, key = b.as_coeff_rest()
 181                         if key in d:
 182                             d[key] += coeff
 183                         else:
 184                             d[key] = coeff
 185             else:
 186                 coeff, key = a.as_coeff_rest()
 187                 if key in d:
 188                     d[key] += coeff
 189                 else:
 190                     d[key] = coeff
 191         if len(d)==0:
 192             return num
 193         args = []
 194         for a, b in d.iteritems():
 195             args.append(Mul((a, b)))
 196         if num.i != 0:
 197             args.insert(0, num)
 198         if len(args) == 1:
 199             return args[0]
 200         else:
 201             return Add(args, False)
 202
 203     def __eq__(self, o):
 204         o = sympify(o)
 205         if o.type == ADD:
 206             return compare_lists(self.args, o.args)
 207         else:
 208             return False
 209
 210     def __str__(self):
 211         s = str(self.args[0])
 212         if self.args[0].type == ADD:
 213             s = "(%s)" % str(s)
 214         for x in self.args[1:]:
 215             s = "%s + %s" % (s, str(x))
 216             if x.type == ADD:
 217                 s = "(%s)" % s
 218         return s
 219
 220     def __hash__(self):
 221         a = list(self.args[:])
 222         a.sort(key=hash)
 223         return hash_seq(a)
 224
 225     def expand(self):
 226         r = Integer(0)
 227         for term in self.args:
 228             r += term.expand()
 229         return r
 230
 231 class Mul(Basic):
 232
 233     def __new__(cls, args, canonicalize=True):
 234         if canonicalize == False:
 235             obj = Basic.__new__(cls, MUL, args)
 236             return obj
 237         args = [sympify(x) for x in args]
 238         return Mul.canonicalize(args)
 239
 240     @classmethod
 241     def canonicalize(cls, args):
 242         use_glib = 0
 243         if use_glib:
 244             from csympy import HashTable
 245             d = HashTable()
 246         else:
 247             d = {}
 248         num = Integer(1)
 249         for a in args:
 250             if a.type == INTEGER:
 251                 num *= a
 252             elif a.type == MUL:
 253                 for b in a.args:
 254                     if b.type == INTEGER:
 255                         num *= b
 256                     else:
 257                         key, coeff = b.as_base_exp()
 258                         if key in d:
 259                             d[key] += coeff
 260                         else:
 261                             d[key] = coeff
 262             else:
 263                 key, coeff = a.as_base_exp()
 264                 if key in d:
 265                     d[key] += coeff
 266                 else:
 267                     d[key] = coeff
 268         if num.i == 0 or len(d)==0:
 269             return num
 270         args = []
 271         for a, b in d.iteritems():
 272             args.append(Pow((a, b)))
 273         if num.i != 1:
 274             args.insert(0, num)
 275         if len(args) == 1:
 276             return args[0]
 277         else:
 278             return Mul(args, False)
 279
 280     def __hash__(self):
 281         a = list(self.args[:])
 282         a.sort(key=hash)
 283         return hash_seq(a)
 284
 285     def __eq__(self, o):
 286         o = sympify(o)
 287         if o.type == MUL:
 288             return compare_lists(self.args, o.args)
 289         else:
 290             return False
 291
 292
 293     def as_coeff_rest(self):
 294         if self.args[0].type == INTEGER:
 295             return self.as_two_terms()
 296         return (Integer(1), self)
 297
 298     def as_two_terms(self):
 299         return (self.args[0], Mul(self.args[1:]))
 300
 301
 302     def __str__(self):
 303         s = str(self.args[0])
 304         if self.args[0].type in [ADD, MUL]:
 305             s = "(%s)" % str(s)
 306         for x in self.args[1:]:
 307             if x.type in [ADD, MUL]:
 308                 s = "%s * (%s)" % (s, str(x))
 309             else:
 310                 s = "%s*%s" % (s, str(x))
 311         return s
 312
 313     @classmethod
 314     def expand_two(self, a, b):
 315         """
 316         Both a and b are assumed to be expanded.
 317         """
 318         if a.type == ADD and b.type == ADD:
 319             r = Integer(0)
 320             for x in a.args:
 321                 for y in b.args:
 322                     r += x*y
 323             return r
 324         if a.type == ADD:
 325             r = Integer(0)
 326             for x in a.args:
 327                 r += x*b
 328             return r
 329         if b.type == ADD:
 330             r = Integer(0)
 331             for y in b.args:
 332                 r += a*y
 333             return r
 334         return a*b
 335
 336     def expand(self):
 337         a, b = self.as_two_terms()
 338         r = Mul.expand_two(a, b)
 339         if r == self:
 340             a = a.expand()
 341             b = b.expand()
 342             return Mul.expand_two(a, b)
 343         else:
 344             return r.expand()
 345
 346 class Pow(Basic):
 347
 348     def __new__(cls, args, canonicalize=True):
 349         if canonicalize == False:
 350             obj = Basic.__new__(cls, POW, args)
 351             return obj
 352         args = [sympify(x) for x in args]
 353         return Pow.canonicalize(args)
 354
 355     @classmethod
 356     def canonicalize(cls, args):
 357         base, exp = args
 358         if base.type == INTEGER:
 359             if base.i == 0:
 360                 return Integer(0)
 361             if base.i == 1:
 362                 return Integer(1)
 363         if exp.type == INTEGER:
 364             if exp.i == 0:
 365                 return Integer(1)
 366             if exp.i == 1:
 367                 return base
 368         if base.type == POW:
 369             return Pow((base.args[0], base.args[1]*exp))
 370         return Pow(args, False)
 371
 372     def __str__(self):
 373         s = str(self.args[0])
 374         if self.args[0].type == ADD:
 375             s = "(%s)" % s
 376         if self.args[1].type == ADD:
 377             s = "%s^(%s)" % (s, str(self.args[1]))
 378         else:
 379             s = "%s^%s" % (s, str(self.args[1]))
 380         return s
 381
 382     def as_base_exp(self):
 383         return self.args
 384
 385     def expand(self):
 386         base, exp = self.args
 387         if base.type == ADD and exp.type == INTEGER:
 388             n = exp.i
 389             m = len(base.args)
 390             print "multi:", m, n
 391             d = multinomial_coefficients(m, n)
 392             r = Integer(0)
 393             print "assembly"
 394             for powers, coeff in d.iteritems():
 395                 t = [Integer(coeff)]
 396                 for x, p in zip(base.args, powers):
 397                     t.append(Pow((x, p)))
 398                 t = Mul(t)
 399                 if r.type == ADD:
 400                     add_args = list(r.args) + [t]
 401                     # XXX: this needs to be checked, i.e. for some more
 402                     # complicated examples one should still simplify things
 403                     r = Add(add_args, False)
 404                 else:
 405                     r = r + t
 406             print "done"
 407             return r
 408         return self
 409
 410 def sympify(x):
 411     if isinstance(x, int):
 412         return Integer(x)
 413     return x
 414
 415 def var(s):
 416     """
 417         Create a symbolic variable with the name *s*.
 418
 419         INPUT:
 420             s -- a string, either a single variable name, or
 421                  a space separated list of variable names, or
 422                  a list of variable names.
 423
 424         NOTE: The new variable is both returned and automatically injected into
 425         the parent's *global* namespace.  It's recommended not to use "var" in
 426         library code, it is better to use symbols() instead.
 427
 428         EXAMPLES:
 429         We define some symbolic variables:
 430             >>> var('m')
 431             m
 432             >>> var('n xx yy zz')
 433             (n, xx, yy, zz)
 434             >>> n
 435             n
 436
 437     """
 438     import re
 439     import inspect
 440     frame = inspect.currentframe().f_back
 441
 442     try:
 443         if not isinstance(s, list):
 444             s = re.split('\s|,', s)
 445
 446         res = []
 447
 448         for t in s:
 449             # skip empty strings
 450             if not t:
 451                 continue
 452             sym = Symbol(t)
 453             frame.f_globals[t] = sym
 454             res.append(sym)
 455
 456         res = tuple(res)
 457         if len(res) == 0:   # var('')
 458             res = None
 459         elif len(res) == 1: # var('x')
 460             res = res[0]
 461                             # otherwise var('a b ...')
 462         return res
 463
 464     finally:
 465         # we should explicitly break cyclic dependencies as stated in inspect
 466         # doc
 467         del frame
 468
 469 def binomial_coefficients(n):
 470     """Return a dictionary containing pairs {(k1,k2) : C_kn} where
 471     C_kn are binomial coefficients and n=k1+k2."""
 472     d = {(0, n):1, (n, 0):1}
 473     a = 1
 474     for k in xrange(1, n//2+1):
 475         a = (a * (n-k+1))//k
 476         d[k, n-k] = d[n-k, k] = a
 477     return d
 478
 479 def binomial_coefficients_list(n):
 480     """ Return a list of binomial coefficients as rows of the Pascal's
 481     triangle.
 482     """
 483     d = [1] * (n+1)
 484     a = 1
 485     for k in xrange(1, n//2+1):
 486         a = (a * (n-k+1))//k
 487         d[k] = d[n-k] = a
 488     return d
 489
 490 def multinomial_coefficients(m, n, _tuple=tuple, _zip=zip):
 491     """Return a dictionary containing pairs ``{(k1,k2,..,km) : C_kn}``
 492     where ``C_kn`` are multinomial coefficients such that
 493     ``n=k1+k2+..+km``.
 494
 495     For example:
 496
 497     >>> print multinomial_coefficients(2,5)
 498     {(3, 2): 10, (1, 4): 5, (2, 3): 10, (5, 0): 1, (0, 5): 1, (4, 1): 5}
 499
 500     The algorithm is based on the following result:
 501
 502        Consider a polynomial and it's ``m``-th exponent::
 503
 504          P(x) = sum_{i=0}^m p_i x^k
 505          P(x)^n = sum_{k=0}^{m n} a(n,k) x^k
 506
 507        The coefficients ``a(n,k)`` can be computed using the
 508        J.C.P. Miller Pure Recurrence [see D.E.Knuth, Seminumerical
 509        Algorithms, The art of Computer Programming v.2, Addison
 510        Wesley, Reading, 1981;]::
 511
 512          a(n,k) = 1/(k p_0) sum_{i=1}^m p_i ((n+1)i-k) a(n,k-i),
 513
 514        where ``a(n,0) = p_0^n``.
 515     """
 516
 517     if m==2:
 518         return binomial_coefficients(n)
 519     symbols = [(0,)*i + (1,) + (0,)*(m-i-1) for i in range(m)]
 520     s0 = symbols[0]
 521     p0 = [_tuple(aa-bb for aa,bb in _zip(s,s0)) for s in symbols]
 522     r = {_tuple(aa*n for aa in s0):1}
 523     r_get = r.get
 524     r_update = r.update
 525     l = [0] * (n*(m-1)+1)
 526     l[0] = r.items()
 527     for k in xrange(1, n*(m-1)+1):
 528         d = {}
 529         d_get = d.get
 530         for i in xrange(1, min(m,k+1)):
 531             nn = (n+1)*i-k
 532             if not nn:
 533                 continue
 534             t = p0[i]
 535             for t2, c2 in l[k-i]:
 536                 tt = _tuple([aa+bb for aa,bb in _zip(t2,t)])
 537                 cc = nn * c2
 538                 b = d_get(tt)
 539                 if b is None:
 540                     d[tt] = cc
 541                 else:
 542                     cc = b + cc
 543                     if cc:
 544                         d[tt] = cc
 545                     else:
 546                         del d[tt]
 547         r1 = [(t, c//k) for (t, c) in d.iteritems()]
 548         l[k] = r1
 549         r_update(r1)
 550     return r