sympycore/heads/term_coeff_dict.py

   1
   2 __all__ = ['TERM_COEFF_DICT']
   3
   4 from .base import heads, heads_precedence, ArithmeticHead, Pair
   5
   6 from ..core import init_module, Expr
   7 init_module.import_heads()
   8 init_module.import_numbers()
   9 init_module.import_lowlevel_operations()
  10
  11 @init_module
  12 def _init(module):
  13     from ..arithmetic.number_theory import multinomial_coefficients
  14     module.multinomial_coefficients = multinomial_coefficients
  15
  16 class TermCoeffDictHead(ArithmeticHead):
  17
  18     def is_data_ok(self, cls, data):
  19         if type(data) is dict:
  20             n = len(data)
  21             #if n<=1:
  22             #    return 'data dictonary should have more than 1 item'
  23             for item in data.iteritems():
  24                 msg = TERM_COEFF.is_data_ok(cls, item, allow_number_term=True)
  25                 if msg:
  26                     return 'TERM_COEFF data=%s: %s' % (item, msg) #pragma: no cover
  27         else:
  28             return 'data must be dict instance but got %s' % (type(data)) #pragma: no cover
  29         return
  30
  31     def __repr__(self): return 'TERM_COEFF_DICT'
  32
  33     def data_to_str_and_precedence(self, cls, term_coeff_dict):
  34         r = [cls(TERM_COEFF, tc) for tc in term_coeff_dict.items()]
  35         return ADD.data_to_str_and_precedence(cls, r)
  36
  37     def new(self, cls, data, evaluate=True):
  38         return term_coeff_dict_new(cls, data)
  39
  40     def reevaluate(self, cls, data):
  41         r = cls(NUMBER, 0)
  42         for term, coeff in data.iteritems():
  43             r += term * coeff
  44         return r
  45
  46     def to_ADD(self, Algebra, data, expr):
  47         return add_new(Algebra, [term_coeff_new(Algebra, term_coeff) for term_coeff in data.iteritems()])
  48
  49     def to_EXP_COEFF_DICT(self, cls, data, expr, variables=None):
  50         if variables is None:
  51             r = cls(EXP_COEFF_DICT, Pair((), {}))
  52         else:
  53             r = cls(EXP_COEFF_DICT, Pair(variables, {}))
  54         for term, coeff in data.iteritems():
  55             r += term * coeff
  56         return r
  57
  58     def term_coeff(self, cls, expr):
  59         term_coeff_dict = expr.data
  60         if len(term_coeff_dict)==1:
  61             return dict_get_item(term_coeff_dict)
  62         return expr, 1
  63
  64     def inplace_add(self, cls, lhs, rhs):
  65         if lhs.is_writable:
  66             data = lhs.data
  67         else:
  68             data = lhs.data.copy()
  69
  70         self.add(cls, data, rhs, inplace=True)
  71
  72         return term_coeff_dict_new(cls, data)
  73
  74     def add(self, cls, lhs, rhs, inplace=False):
  75         if inplace:
  76             data = lhs
  77         else:
  78             data = lhs.data.copy()
  79
  80         h2, d2 = rhs.pair
  81         if h2 is NUMBER:
  82             if d2 != 0:
  83                 dict_add_item(cls, data, cls(NUMBER, 1), d2)
  84         elif h2 is SYMBOL:
  85             dict_add_item(cls, data, rhs, 1)
  86         elif h2 is TERM_COEFF:
  87             term, coeff = d2
  88             dict_add_item(cls, data, term, coeff)
  89         elif h2 is TERM_COEFF_DICT:
  90             dict_add_dict(cls, data, d2)
  91         elif h2 is ADD:
  92             for op in d2:
  93                 term, coeff = op.term_coeff()
  94                 dict_add_item(cls, data, term, coeff)
  95         elif h2 is SUB or h2 is NEG or h2 is POS:
  96             raise NotImplementedError(`self, rhs.pair`)
  97         elif h2 is BASE_EXP_DICT:
  98             c = base_exp_dict_get_coefficient(cls, d2)
  99             if c is not None:
 100                 d = d2.copy()
 101                 del d[c]
 102                 t = BASE_EXP_DICT.new(cls, d)
 103                 if t.head is BASE_EXP_DICT:
 104                     dict_add_item(cls, data, t, c)
 105                 else:
 106                     self.add(cls, data, t * c, inplace=True)
 107             else:
 108                 dict_add_item(cls, data, rhs, 1)
 109         else:
 110             dict_add_item(cls, data, rhs, 1)
 111
 112         if inplace:
 113             return lhs
 114
 115         return term_coeff_dict_new(cls, data)
 116
 117     def add_number(self, cls, lhs, rhs):
 118         if rhs==0:
 119             return lhs
 120         data = lhs.data.copy()
 121         term_coeff_dict_add_item(cls, data, cls(NUMBER, 1), rhs)
 122         return term_coeff_dict_new(cls, data)
 123
 124     def sub_number(self, cls, lhs, rhs):
 125         if rhs==0:
 126             return lhs
 127         data = lhs.data.copy()
 128         term_coeff_dict_add_item(cls, data, cls(NUMBER, 1), -rhs)
 129         return term_coeff_dict_new(cls, data)
 130
 131     def sub(self, cls, lhs, rhs):
 132         d = lhs.data.copy()
 133         self.add(cls, d, -rhs, inplace=True)
 134         return term_coeff_dict_new(cls, d)
 135
 136     def commutative_mul(self, cls, lhs, rhs):
 137         rhead, rdata = rhs.pair
 138         if rhead is NUMBER:
 139             if rdata==0:
 140                 return rhs
 141             data = lhs.data.copy()
 142             dict_mul_value(cls, data, rdata)
 143             return term_coeff_dict_new(cls, data)
 144         if rhead is TERM_COEFF:
 145             term, coeff = rdata
 146             return (lhs * term) * coeff
 147         if rhead is POW:
 148             base, exp = rdata
 149             if lhs==base:
 150                 return POW.new(cls, (lhs, exp + 1))
 151             return cls(BASE_EXP_DICT, {lhs:1, base:exp})
 152         if rhead is SYMBOL or rhead is APPLY:
 153             return cls(BASE_EXP_DICT, {lhs:1, rhs:1})
 154         if rhead is TERM_COEFF_DICT:
 155             if rdata==lhs.data:
 156                 return cls(POW, (lhs, 2))
 157             return cls(BASE_EXP_DICT, {lhs:1, rhs:1})
 158         if rhead is BASE_EXP_DICT:
 159             data = rdata.copy()
 160             dict_add_item(cls, data, lhs, 1)
 161             return BASE_EXP_DICT.new(cls, data)
 162         if rhead is ADD:
 163             return cls(BASE_EXP_DICT, {lhs:1, rhs:1})
 164         return ArithmeticHead.commutative_mul(self, cls, lhs, rhs)
 165
 166     def commutative_mul_number(self, cls, lhs, rhs):
 167         if rhs==0:
 168             return cls(NUMBER, 0)
 169         data = lhs.data.copy()
 170         dict_mul_value(cls, data, rhs)
 171         return cls(self, data)
 172
 173     non_commutative_mul_number = commutative_mul_number
 174
 175     def commutative_rdiv_number(self, cls, lhs, rhs):
 176         return term_coeff_new(cls, (cls(POW, (lhs, -1)), rhs))
 177
 178     def commutative_div(self, cls, lhs, rhs):
 179         rhead, rdata = rhs.pair
 180         if rhead is NUMBER:
 181             return self.commutative_div_number(cls, lhs, rdata)
 182         if rhead is TERM_COEFF_DICT:
 183             if lhs.data == rdata:
 184                 return cls(NUMBER, 1)
 185             return cls(BASE_EXP_DICT, {lhs:1, rhs:-1})
 186         if rhead is SYMBOL or rhead is APPLY:
 187             return cls(BASE_EXP_DICT, {lhs:1, rhs:-1})
 188         if rhead is TERM_COEFF:
 189             term, coeff = rdata
 190             return number_div(cls, 1, coeff) * (lhs / term)
 191         if rhead is POW:
 192             base, exp = rdata
 193             if lhs==base:
 194                 return pow_new(cls, (lhs, 1-exp))
 195             return cls(BASE_EXP_DICT, {lhs:1, base:-exp})
 196         if rhead is BASE_EXP_DICT:
 197             data = {lhs:1}
 198             for base, exp in rdata.iteritems():
 199                 base_exp_dict_add_item(cls, data, base, -exp)
 200             return base_exp_dict_new(cls, data)
 201         return ArithmeticHead.commutative_div(self, cls, lhs, rhs)
 202
 203
 204     def pow(self, cls, base, exp):
 205         if exp==0: return cls(NUMBER, 1)
 206         if exp==1: return base
 207         d = base.data
 208         if len(d)==1:
 209             t,c = dict_get_item(d)
 210             t,c = t**exp, c**exp
 211             if t==1: return cls(NUMBER, c)
 212             if c==1: return t
 213             return cls(TERM_COEFF, (t, c))
 214         return POW.new(cls, (base, exp))
 215
 216     pow_number = pow
 217
 218     def neg(self, cls, expr):
 219         d = expr.data.copy()
 220         for key in d:
 221             d[key] = -d[key]
 222         return cls(TERM_COEFF_DICT, d)
 223
 224     def expand(self, cls, expr):
 225         d = {}
 226         for t, c in expr.data.items():
 227             self.add(cls, d, t.expand() * c, inplace=True)
 228         return term_coeff_dict_new(cls, d)
 229
 230     def expand_intpow(self, cls, expr, intexp):
 231         if intexp<0:
 232             return cls(POW, (expr, intexp))
 233         if intexp==0:
 234             return cls(NUMBER, 1)
 235         if intexp==1:
 236             return expr
 237         term_coeff_list = [(term.base_exp(), coeff) for term, coeff in expr.data.items()]
 238         mdata = multinomial_coefficients(len(term_coeff_list), intexp)
 239         d = {}
 240         for e,c in mdata.iteritems():
 241             new_coeff = c
 242             df = {}
 243             for e_i, ((base, exp), coeff) in zip(e, term_coeff_list):
 244                 if e_i:
 245                     if e_i==1:
 246                         base_exp_dict_add_item(cls, df, base, exp)
 247                         if coeff is not 1:
 248                             new_coeff *= coeff
 249                     else:
 250                         base_exp_dict_add_item(cls, df, base, exp*e_i)
 251                         if coeff is not 1:
 252                             new_coeff *= coeff ** e_i
 253             new_term = base_exp_dict_new(cls, df)
 254             term_coeff_dict_add_item(cls, d, new_term, new_coeff)
 255         return term_coeff_dict_new(cls, d)
 256
 257     def walk(self, func, cls, data, target):
 258         d = {}
 259         flag = False
 260         add = self.add
 261         for t, c in data.iteritems():
 262             t1 = t.head.walk(func, cls, t.data, t)
 263             if isinstance(c, Expr):
 264                 c1 = c.head.walk(func, cls, c.data, c)
 265             else:
 266                 c1 = NUMBER.walk(func, cls, c, c)
 267             if t1 is not t or c1 is not c:
 268                 flag = True
 269             add(cls, d, t1 * c1, inplace=True)
 270         if flag:
 271             r = term_coeff_dict_new(cls, d)
 272             return func(cls, r.head, r.data, r)
 273         return func(cls, self, data, target)
 274
 275     def scan(self, proc, cls, data, target):
 276         for t, c in data.iteritems():
 277             t.head.scan(proc, cls, t.data, target)
 278             if isinstance(c, Expr):
 279                 c.head.scan(proc, cls, c.data, target)
 280             else:
 281                 NUMBER.scan(proc, cls, c, target)
 282         proc(cls, self, data, target)
 283
 284     def diff(self, cls, data, expr, symbol, order, cache={}):
 285         key = (expr, symbol, order)
 286         result = cache.get(key)
 287         if result is not None:
 288             return result
 289         if result is None:
 290             d = {}
 291             result = cls(NUMBER, 0)
 292             for term, coeff in data.iteritems():
 293                 result += term.head.diff(cls, term.data, term, symbol, order, cache=cache) * coeff
 294             key1 = (expr, symbol, 1)
 295         cache[key] = result
 296         return result
 297
 298     def apply(self, cls, data, func, args):
 299         result = cls(NUMBER, 0)
 300         for term, coeff in data.iteritems():
 301             result += term.head.apply(cls, term.data, term, args) * coeff
 302         return result
 303
 304     def integrate_indefinite(self, cls, data, expr, x):
 305         result = cls(NUMBER, 0)
 306         for term, coeff in data.iteritems():
 307             result += term.head.integrate_indefinite(cls, term.data, term, x) * coeff
 308         return result
 309
 310     def integrate_definite(self, cls, data, expr, x, a, b):
 311         result = cls(NUMBER, 0)
 312         for term, coeff in data.iteritems():
 313             result += term.head.integrate_definite(cls, term.data, term, x, a, b) * coeff
 314         return result
 315
 316     def algebra_pos(self, Algebra, expr):
 317         return expr
 318
 319     def algebra_neg(self, Algebra, expr):
 320         if Algebra.algebra_options.get('evaluate_addition'):
 321             d = expr.data.copy()
 322             for key in d:
 323                 d[key] = -d[key]
 324             return Algebra(TERM_COEFF_DICT, d)
 325         return Algebra(NEG, expr)
 326
 327     def algebra_add_number(self, Algebra, lhs, rhs, inplace):
 328         if not rhs:
 329             return lhs
 330         if inplace:
 331             term_coeff_dict_add_item(Algebra, lhs.data, Algebra(NUMBER, 1), rhs)
 332             return term_coeff_dict(Algebra, lhs)
 333         d = lhs.data.copy()
 334         term_coeff_dict_add_item(Algebra, d, Algebra(NUMBER, 1), rhs)
 335         return term_coeff_dict_new(Algebra, d)
 336
 337     def algebra_add(self, Algebra, lhs, rhs, inplace):
 338         rhead, rdata = rhs.pair
 339         if Algebra.algebra_options.get('evaluate_addition'):
 340             ldata = lhs.data
 341             if rhead is ADD or rhead is EXP_COEFF_DICT or rhead is MUL or rhead is NEG:
 342                 rhs = rhead.to_TERM_COEFF_DICT(Algebra, rdata, rhs)
 343                 rhead, rdata = rhs.pair
 344             if rhead is NUMBER:
 345                 if not rdata:
 346                     return lhs
 347                 rterm, rcoeff = Algebra(NUMBER, 1), rdata
 348             elif rhead is SYMBOL:
 349                 rterm, rcoeff = rhs, 1
 350             elif rhead is TERM_COEFF:
 351                 rterm, rcoeff = rdata
 352             elif rhead is TERM_COEFF_DICT:
 353                 if inplace:
 354                     term_coeff_dict_add_dict(Algebra, ldata, rdata)
 355                     return term_coeff_dict(Algebra, lhs)
 356                 d = ldata.copy()
 357                 term_coeff_dict_add_dict(Algebra, d, rdata)
 358                 return term_coeff_dict_new(Algebra, d)
 359             else:
 360                 return super(type(self), self).algebra_add(Algebra, lhs, rhs, inplace)
 361
 362             if inplace:
 363                 term_coeff_dict_add_item(Algebra, ldata, rterm, rcoeff)
 364                 return term_coeff_dict(Algebra, lhs)
 365             d = ldata.copy()
 366             term_coeff_dict_add_item(Algebra, d, rterm, rcoeff)
 367             return term_coeff_dict_new(Algebra, d)
 368         else:
 369             return TERM_COEFF_DICT.to_ADD(Algebra, lhs.data, lhs) + rhs
 370         return super(type(self), self).algebra_add(Algebra, lhs, rhs, inplace)
 371
 372     def algebra_mul_number(self, Algebra, lhs, rhs, inplace):
 373         if Algebra.algebra_options.get('evaluate_addition'):
 374             if not rhs:
 375                 return Algebra(NUMBER, 0)
 376             if rhs==1:
 377                 return lhs
 378             if inplace:
 379                 term_coeff_dict_mul_value(Algebra, lhs.data, rhs)
 380                 return lhs
 381             d = lhs.data.copy()
 382             term_coeff_dict_mul_value(Algebra, d, rhs)
 383             return Algebra(TERM_COEFF_DICT, d)
 384         return Algebra(MUL, [lhs, Algebra(NUMBER, rhs)])
 385
 386     def algebra_mul(self, Algebra, lhs, rhs, inplace):
 387         rhead, rdata = rhs.pair
 388         if rhead is NUMBER:
 389             return self.algebra_mul_number(Algebra, lhs, rdata, inplace)
 390         return super(type(self), self).algebra_mul(Algebra, lhs, rhs, inplace)
 391
 392     def algebra_div_number(self, Algebra, lhs, rhs, inplace):
 393         if Algebra.algebra_options.get('evaluate_addition'):
 394             if rhs==1:
 395                 return lhs
 396             d1 = gcd(*lhs.data.values())
 397             d2 = gcd(d1, rhs)
 398             d3 = rhs / d2
 399             d = {}
 400             rd = 0
 401             for t,c in lhs.data.items():
 402                 c /= d2
 403                 q, c = divmod(c, d3)
 404                 if c:
 405                     d[t] = c
 406                 rd += q
 407             s = term_coeff_dict_new(Algebra, d)
 408             if rhs==d2:
 409                 assert rd==0,`lsh, rhs, s,rd`
 410                 return s
 411             return Algebra(DIV, [s, Algebra(NUMBER, d3)]) + rd
 412         return Algebra(DIV, [lhs, Algebra(NUMBER, rhs)])
 413
 414     def algebra_div(self, Algebra, lhs, rhs, inplace):
 415         rhead, rdata = rhs.pair
 416         if rhead is NUMBER:
 417             return self.algebra_div_number(Algebra, lhs, rdata, inplace)
 418         return super(type(self), self).algebra_div(Algebra, lhs, rhs, inplace)
 419
 420 TERM_COEFF_DICT = TermCoeffDictHead()