sympycore/expr.py

   1 # -*- coding: latin-1 -*-
   2 """
   3   This module implements extension type Expr that holds two Python
   4   objects, head and data, in a pair attribute.
   5
   6   When adding new features to Expr class, make sure that these are
   7   added also to extension type Expr in src/expr_ext.c.
   8
   9 C Expr:
  10
  11 >>> from sympycore.expr_ext import *
  12 >>> %timeit Expr(1,2)
  13 10000000 loops, best of 3: 179 ns per loop
  14 >>> e=Expr(1,2)
  15 >>> %timeit h = e.head
  16 10000000 loops, best of 3: 113 ns per loop
  17 >>> %timeit h,d = e.pair
  18 10000000 loops, best of 3: 141 ns per loop
  19
  20 Python Expr:
  21
  22 >>> from sympycore.expr import *
  23 >>> %timeit Expr(1,2)
  24 1000000 loops, best of 3: 988 ns per loop
  25 >>> e=Expr(1,2)
  26 >>> %timeit h = e.head
  27 10000000 loops, best of 3: 119 ns per loop
  28 >>> %timeit h,d = e.pair
  29 10000000 loops, best of 3: 139 ns per loop
  30 """
  31 # Author: Pearu Peterson
  32 # Created: March 2008
  33
  34 __all__ = ['Expr', 'Pair']
  35
  36 def init_module(m):
  37     from .core import heads
  38     for n,h in heads.iterNameValue(): setattr(m, n, h)
  39     from .arithmetic.numbers import inttypes, numbertypes, try_power, gcd
  40     m.inttypes = inttypes
  41     m.numbertypes = numbertypes
  42     m.try_power = try_power
  43     m.gcd = gcd
  44
  45 class Expr(object):
  46     """Represents an symbolic expression in a pair form: (head, data)
  47
  48 The pair (head, data) is saved in an attribute ``pair``. The parts of
  49 a pair, head and data, can be also accessed via ``head`` and ``data``
  50 attributes, respectively. All three attributes are read-only.
  51
  52 The head part is assumed to be an immutable object.  The data part can
  53 be either an immutable object or Python dictionary.  In the former
  54 case, the hash value of Expr instance is defined as::
  55
  56     hash((<Expr>.head, <Expr>.
  57
  58 Otherwise, if ``data`` contains a Python dictionary, then the hash
  59 value is defined as::
  60
  61     hash((<Expr>.head, frozenset(<Expr>.data.items())
  62
  63 If ``data`` is a Python list, then the hash value is::
  64
  65   hash((<Expr>.head, tuple(<Expr>.data)))
  66
  67 WARNING: the hash value of an Expr instance is computed (and cached)
  68 when it is used as a key to Python dictionary. This means that the
  69 instance content MUST NOT be changed after the hash is computed.  To
  70 check if it is safe to change the ``data`` dictionary, use
  71 ``is_writable`` attribute that returns True when the hash value has
  72 not been computed::
  73
  74     <Expr>.is_writable -> True or False
  75
  76 There are two ways to access the parts of a Expr instance from
  77 Python::
  78
  79     a = Expr(<head>, <data>)
  80     head, data = a.head, a.data     - for backward compatibility
  81     head, data = a.pair             - fastest way
  82
  83 When Expr constructor is called with one argument, say ``x``, then
  84 ``<Expr subclass>.convert(x)`` will be returned.
  85
  86 This is Python version of Expr type.
  87 """
  88
  89     __slots__ = ['head', 'data', 'pair', '_hash']
  90
  91     def __init__(self, *args):
  92         if len(args)==1:
  93             obj = self.convert(args[0])
  94             self.pair = obj.pair
  95             self._hash = obj._hash
  96         elif len(args)==2:
  97             self.pair = args
  98             self._hash = None
  99         else:
 100             raise TypeError("%s requires 1 or 2 arguments but got %r" % (type(self), len(args)))
 101         msg = self.head.is_data_ok(type(self), self.data)
 102         if msg:
 103             msg = '%s(head=%s, data=%s): %s' % (type(self).__name__, self.head, self.data, msg)
 104             #print msg
 105             raise TypeError(msg)
 106
 107     def __repr__(self):
 108         return '%s%r' % (type(self).__name__, self.pair)
 109
 110     def __hash__(self):
 111         """ Compute hash value.
 112
 113         Different from expr_ext.Expr, an exception is raised when data
 114         dictionary values contain dictionaries.
 115         """
 116         h = self._hash
 117         if h is None:
 118             pair = self.pair
 119             obj = self.as_lowlevel()
 120             if obj is not pair:
 121                 h = hash(obj)
 122             else:
 123                 head, data = pair
 124                 tdata = type(data)
 125                 if tdata is dict:
 126                     h = hash((head, frozenset(data.iteritems())))
 127                 elif tdata is list:
 128                     h = hash((head, tuple(data)))
 129                 else:
 130                     h = hash(pair)
 131             self._hash = h
 132         return h
 133
 134     @property
 135     def is_writable(self, _writable_types = (list, dict)):
 136         if self._hash is None:
 137             data = self.pair[1]
 138             tdata = type(data)
 139             if tdata in _writable_types:
 140                 return True
 141             if tdata is Pair:
 142                 return data.is_writable
 143         return False
 144
 145     @property
 146     def head(self):
 147         return self.pair[0]
 148
 149     @property
 150     def data(self):
 151         return self.pair[1]
 152
 153     # Pickle support:
 154     def _sethash(self, hashvalue):
 155         """ Set hash value for the object.
 156
 157         If hashvalue==-1, then the hash value will be reset.
 158
 159         Used by pickle support in sympycore.core._reconstruct. DO NOT
 160         use this method directly.
 161         """
 162         if hashvalue==-1:
 163             self._hash = None
 164         else:
 165             self._hash = hashvalue
 166
 167     def __reduce__(self):
 168         # see also _reconstruct function in sympycore/core.py
 169         version = 3
 170         from sympycore.core import _reconstruct
 171         if version==1:
 172             hashvalue = self._hash
 173             if hashvalue is None:
 174                 hashvalue = -1
 175             state = (type(self), self.pair, hashvalue)
 176         elif version==2 or version==3:
 177             hashvalue = self._hash
 178             if hashvalue is None:
 179                 hashvalue = -1
 180             cls = type(self)
 181             typ = type(cls)
 182             try:
 183                 args = typ.__getinitargs__(cls)
 184             except AttributeError:
 185                 args = None
 186             if args is None:
 187                 # either metaclass does not define __getinitargs__ method
 188                 # or cls has no metaclass
 189                 state = (cls, self.pair, hashvalue)
 190             else:
 191                 state = ((typ, args), self.pair, hashvalue)
 192         else:
 193             raise NotImplementedError('pickle state version %s' % (version))
 194         return  _reconstruct, (version, state)
 195
 196     def __nonzero__(self):
 197         # Note that `not cls(MUL, [])` would return True while `cls(MUL, [])==1`.
 198         # So, must use as_lowlevel:
 199         obj = self.as_lowlevel()
 200         if obj is not self.pair:
 201             return not not obj
 202         return not not self.data
 203
 204     def as_lowlevel(self):
 205         """ Return self as low-level object instance that will be used
 206         in comparison and in hash computation.
 207
 208         By default, as_lowlevel uses heads to_lowlevel(cls, data, pair)
 209         method but since as_lowlevel is a most frequently called
 210         method then for some heads the corresponding code is copied
 211         here. The default return value is a pair tuple for composite
 212         objects and data part for atomic objects. The as_lowlevel
 213         method may also return an Expr instance but not self
 214         (otherwise infinite recurrsion will occur).
 215
 216         See __hash__, __nonzero__ method for more details how the
 217         results of as_lowlevel method are interpreted.
 218         """
 219         head, data = pair = self.pair
 220         if head is NUMBER or head is SYMBOL or head is SPECIAL:
 221             return data
 222         elif head is MUL or head is DIV:
 223             n = len(data)
 224             if n==0:
 225                 return 1
 226             if n==1:
 227                 return data[0]
 228         elif head is ADD or head is SUB:
 229             n = len(data)
 230             if n==0:
 231                 return 0
 232             if n==1:
 233                 return data[0]
 234         elif head is POW:
 235             base, exp = data
 236             if exp==0 or base==1:
 237                 return 1
 238             if exp==1:
 239                 return base
 240         elif head is TERM_COEFF:
 241             term, coeff = data
 242             if coeff==0:
 243                 return 0
 244             if coeff==1:
 245                 return term
 246             if term==1:
 247                 return coeff
 248         elif head is TERM_COEFF_DICT:
 249             n = len(data)
 250             if n==0:
 251                 return 0
 252             if n==1:
 253                 return type(self)(TERM_COEFF, dict_get_item(data))
 254         elif head is BASE_EXP_DICT:
 255             n = len(data)
 256             if n==0:
 257                 return 1
 258             if n==1:
 259                 return type(self)(POW, dict_get_item(data))
 260         else:
 261             return head.to_lowlevel(type(self), data, pair)
 262         return pair
 263
 264     def term_coeff(self):
 265         head, data = self.pair
 266         if head is TERM_COEFF:
 267             return data
 268         if head is BASE_EXP_DICT:
 269             cls = type(self)
 270             coeff = base_exp_dict_get_coefficient(cls, data)
 271             if coeff is not None:
 272                 d = data.copy()
 273                 del d[coeff]
 274                 r = BASE_EXP_DICT.new(cls, d)
 275                 return r, coeff.data
 276                 t, c = r.head.term_coeff(cls, r)
 277                 return t, c * coeff
 278             return self, 1
 279         if head is NUMBER:
 280             return self, data
 281         return self, 1
 282
 283     def base_exp(self):
 284         head, data = self.pair
 285         if head==POW:
 286             return data
 287         return self, 1
 288
 289     for _item in dict(__eq__ = '==', __ne__ = '!=',
 290                       __lt__ = '<', __le__ = '<=',
 291                       __gt__ = '>', __ge__ = '>=',
 292                       ).items():
 293         exec '''
 294 def %s(self, other):
 295     if type(self) is type(other):
 296         other = other.as_lowlevel()
 297     return self.as_lowlevel() %s other
 298 ''' % _item
 299
 300
 301 class Pair(Expr):
 302     """ Holds a pair that may contain list and dict second element.
 303     """
 304     def __init__(self, *args):
 305         if len(args)==2:
 306             self.pair = args
 307             self._hash = None
 308         else:
 309             raise TypeError("%s requires 2 arguments but got %r" % (type(self), len(args)))
 310
 311     def __eq__(self, other):
 312         return self.pair == other
 313
 314     def __len__(self):
 315         return 2
 316
 317     def __getitem__(self, index):
 318         return self.pair[index]
 319
 320
 321 def dict_mul_item(Algebra, d, key, value):
 322     c = d.get(key)
 323     if c is None:
 324         d[key] = value
 325     else:
 326         d[key] = c * value
 327
 328 base_exp_dict_mul_item = dict_mul_item
 329
 330 def term_coeff_dict_mul_item(Algebra, d, key, value):
 331     return dict_mul_item(Algebra, d, key, value)
 332
 333 def dict_add_item(Algebra, d, key, value):
 334     c = d.get(key)
 335     if c is None:
 336         if value:
 337             d[key] = value
 338     else:
 339         c = c + value
 340         if c:
 341             d[key] = c
 342         else:
 343             del d[key]
 344
 345 def base_exp_dict_get_coefficient(Algebra, d):
 346     for k, v in d.iteritems():
 347         if v is 1 and k.head is NUMBER:
 348             return k
 349     return
 350
 351 def base_exp_dict_add_item(Algebra, d, base, exp):
 352     """
 353     d is a dictonary that will be updated with base, exp pair.
 354     Base part must be an Algebra instance while exp part
 355     can either be a number instance or an expression.
 356
 357     The dictionary items (base, exp) must satisfy the following
 358     conditions:
 359       1) numeric (base, exp) must be completely evaluated, that is,
 360       if base is numeric then exp is either 1 or rational with denominator
 361       part that is not a root of base.
 362       2) all numeric items with exp==1 must be combined to one item
 363       (coefficient)
 364       3) items with base==1 and exp==0 must not be present in the dictionary
 365     """
 366     base_head, base_data = base.pair
 367     value = d.get(base)
 368     if type(exp) is Algebra and exp.head is NUMBER:
 369         exp = exp.data
 370     if value is None:
 371         if base_head is NUMBER:
 372             assert base
 373             if base==1:
 374                 return
 375             if exp==1:
 376                 coeff = base_exp_dict_get_coefficient(Algebra, d)
 377                 if coeff is None:
 378                     d[base] = 1
 379                 else:
 380                     del d[coeff]
 381                     coeff = coeff * base
 382                     if coeff==1:
 383                         return
 384                     d[coeff] = 1
 385             else:
 386                 assert not isinstance(exp, inttypes),`base, exp`
 387                 d[base] = exp
 388         else:
 389             assert exp
 390             d[base] = exp
 391     else:
 392         value = value + exp
 393         if type(value) is Algebra and value.head is NUMBER:
 394             value = value.data
 395         if value:
 396             if base_head is NUMBER and isinstance(value, numbertypes):
 397                 del d[base]
 398                 r, l = try_power(base_data, value)
 399                 if r!=1:
 400                     base_exp_dict_add_item(Algebra, d, Algebra(NUMBER, r), 1)
 401                 elif len(l)==1 and l[0]==base_data:
 402                     d[base] = value
 403                 else:
 404                     for b, e in l:
 405                         base_exp_dict_add_item(Algebra, d, Algebra(NUMBER, b), e)
 406             elif base_head is TERM_COEFF and isinstance(value, inttypes):
 407                 del d[base]
 408                 term, coeff = base_data
 409                 base_exp_dict_add_item(Algebra, d, term, value)
 410                 base_exp_dict_add_item(Algebra, d, Algebra(NUMBER, coeff), value)
 411             else:
 412                 d[base] = value
 413         else:
 414             del d[base]
 415
 416 def term_coeff_dict_add_item(Algebra, d, key, value):
 417     khead, kdata = key.pair
 418     if khead is TERM_COEFF:
 419         term, coeff = kdata
 420         dict_add_item(Algebra, d, term, value * coeff)
 421     else:
 422         assert khead not in [TERM_COEFF_DICT],`Algebra, key.pair`
 423         dict_add_item(Algebra, d, key, value)
 424
 425 def term_coeff_dict_add_dict(Algebra, dict1, dict2):
 426     for key, value in dict2.iteritems():
 427         term_coeff_dict_add_item(Algebra, dict1, key, value)
 428
 429 def dict_get_item(d):
 430     return d.items()[0]
 431
 432 def dict_add_dict(Algebra, dict1, dict2):
 433     for key, value in dict2.iteritems():
 434         dict_add_item(Algebra, dict1, key, value)
 435
 436 def base_exp_dict_add_dict(Algebra, dict1, dict2):
 437     for key, value in dict2.iteritems():
 438         base_exp_dict_add_item(Algebra, dict1, key, value)
 439
 440 def base_exp_dict_sub_item(Algebra, dict, key, value):
 441     return base_exp_dict_add_item(Algebra, dict, key, -value)
 442
 443 def base_exp_dict_sub_dict(Algebra, dict1, dict2):
 444     for key, value in dict2.iteritems():
 445         base_exp_dict_sub_item(Algebra, dict1, key, value)
 446
 447 def dict_sub_dict(Algebra, dict1, dict2):
 448     for key, value in dict2.iteritems():
 449         dict_add_item(Algebra, dict1, key, -value)
 450
 451 def base_exp_dict_mul_dict(Algebra, d, dict1, dict2):
 452     for t1,c1 in dict1.iteritems():
 453         for t2,c2 in dict2.iteritems():
 454             t = t1 * t2
 455             t, c = t.term_coeff()
 456             c12 = c * c1 * c2
 457             base_exp_dict_add_item(Algebra, d, t, c12)
 458
 459 def exp_coeff_dict_mul_dict(Algebra, d, dict1, dict2):
 460     for t1,c1 in dict1.iteritems():
 461         for t2,c2 in dict2.iteritems():
 462             t = t1 + t2
 463             c12 = c1 * c2
 464             dict_add_item(Algebra, d, t, c12)
 465
 466 def dict_mul_dict(Algebra, d, dict1, dict2):
 467     for t1,c1 in dict1.iteritems():
 468         for t2,c2 in dict2.iteritems():
 469             t = t1 * t2
 470             #t, c = t.term_coeff()
 471             c12 = c1 * c2
 472             dict_add_item(Algebra, d, t, c12)
 473
 474 def dict_mul_value(Algebra, d, value):
 475     if value==1:
 476         return
 477     for t, c in d.items():
 478         c *= value
 479         if c:
 480             d[t] = c
 481         else:
 482             del d[t]
 483
 484 term_coeff_dict_mul_dict = dict_mul_dict
 485 base_exp_dict_mul_value = dict_mul_value
 486 term_coeff_dict_mul_value = dict_mul_value
 487
 488 def term_coeff_new(Algebra, data):
 489     term, coeff = data
 490     if coeff==1:
 491         return term
 492     if term==1:
 493         return Algebra(NUMBER, coeff)
 494     if coeff==0:
 495         return Algebra(NUMBER, 0)
 496     h,d = term.pair
 497     if h is TERM_COEFF:
 498         t, c = d
 499         return term_coeff_new(Algebra, (t, c * coeff))
 500     if h is NUMBER:
 501         return Algebra(NUMBER, d * coeff)
 502     return Algebra(TERM_COEFF, data)
 503
 504 def term_coeff_dict_new(Algebra, data):
 505     """
 506     Return canonicalized TERM_COEFF_DICT expression from data.
 507     """
 508     n = len(data)
 509     if n>1:
 510         return Algebra(TERM_COEFF_DICT, data)
 511     if n==0:
 512         return Algebra(NUMBER, 0)
 513     return term_coeff_new(Algebra, dict_get_item(data))
 514
 515 def term_coeff_dict(Algebra, expr):
 516     """
 517     Return canonicalized TERM_COEFF_DICT expression from existing
 518     expression:
 519
 520       * if expr has no data, return 0
 521       * if expr data has one item, return TERM_COEFF expression
 522     """
 523     data = expr.data
 524     n = len(data)
 525     if n>1:
 526         return expr
 527     if n==0:
 528         return Algebra(NUMBER, 0)
 529     return term_coeff_new(Algebra, dict_get_item(data))
 530
 531 def pow_new(Algebra, data):
 532     base, exp = data
 533     if exp==1:
 534         return base
 535     if base==1 or exp==0:
 536         return Algebra(NUMBER, 1)
 537     #if type(exp) is Algebra and exp.head is NUMBER:
 538     #    data = base, exp.data
 539     return Algebra(POW, data)
 540
 541 def base_exp_dict_new(Algebra, data):
 542     n = len(data)
 543     if n==0:
 544         return Algebra(NUMBER, 1)
 545     if n==1:
 546         return pow_new(Algebra, dict_get_item(data))
 547     coeff = base_exp_dict_get_coefficient(Algebra, data)
 548     if coeff is None:
 549         return Algebra(BASE_EXP_DICT, data)
 550     del data[coeff]
 551     return term_coeff_new(Algebra, (base_exp_dict_new(Algebra, data), coeff.data))
 552
 553 def add_new(Algebra, data):
 554     n = len(data)
 555     if n==0:
 556         return Algebra(NUMBER, 0)
 557     if n==1:
 558         return data[0]
 559     return Algebra(ADD, data)
 560
 561 def add(Algebra, expr):
 562     data = expr.data
 563     n = len(data)
 564     if n==0:
 565         return Algebra(NUMBER, 0)
 566     if n==1:
 567         return data[0]
 568     return expr
 569
 570 def mul_new(Algebra, data):
 571     n = len(data)
 572     if n==0:
 573         return Algebra(NUMBER, 1)
 574     if n==1:
 575         return data[0]
 576     return Algebra(MUL, data)