sympycore/heads/polynomial.py

   1
   2 __all__ = ['SPARSE_POLY', 'DENSE_POLY']
   3
   4 from .base import Head, heads
   5
   6 from ..core import init_module, Pair, Expr
   7 init_module.import_heads()
   8 init_module.import_numbers()
   9
  10 class SparsepolyHead(Head):
  11     """
  12     SparsepolyHead is a head for sparse polynomials represented as a
  13     dictionary of exponent and coefficient pairs, data can be dict or
  14     frozenset.
  15     """
  16
  17     def __repr__(self): return 'SPARSE_POLY'
  18
  19     def data_to_str_and_precedence(self, cls, data):
  20         return EXP_COEFF_DICT.data_to_str_and_precedence(cls, Pair(cls.variables, data))
  21
  22     def reevaluate(self, cls, data):
  23         return cls(self, data)
  24
  25     def to_lowlevel(self, cls, data, pair):
  26         return EXP_COEFF_DICT.to_lowlevel(cls, Pair(cls.variables, data), pair)
  27
  28     def add(self, cls, lhs, rhs):
  29         rhead, rdata = rhs.pair
  30         if rhead is not self:
  31             rhs = rhead.to_SPARSE_POLY(cls, rdata, rhs)
  32             return lhs + rhs
  33         return NotImplemented
  34
  35     inplace_add = add
  36
  37     def sub_number(self, cls, lhs, rhs):
  38         return lhs + (-rhs)
  39
  40     def sub(self, cls, lhs, rhs):
  41         return lhs + (-rhs)
  42
  43     def commutative_mul(self, cls, lhs, rhs):
  44         rhead, rdata = rhs.pair
  45         if rhead is not self:
  46             rhs = rhead.to_SPARSE_POLY(cls, rdata, rhs)
  47             return lhs * rhs
  48         return NotImplemented
  49
  50     inplace_commutative_mul = commutative_mul
  51
  52     def term_coeff(self, cls, expr):
  53         return expr, 1
  54
  55     def integrate_indefinite_index(self, cls, data, expr, index):
  56         """
  57         Return indefinite integral of expr with respect to index-th variable.
  58         data is expr.data, index is integer.
  59         """
  60         new_data = {}
  61         for exp, coeff in data.iteritems():
  62             new_exp = exp.copy()
  63             new_exp[index] += 1
  64             new_coeff = number_div(cls.ring, coeff, new_exp[index])
  65             new_data[new_exp] = new_coeff
  66         return cls(new_data)
  67
  68     def diff_index(self, cls, data, expr, index, order):
  69         """
  70         Return the order-th derivative of expr with respect to index-th variable.
  71         data is expr.data, index is integer.
  72         """
  73         if order==0:
  74             return expr
  75         new_data = {}
  76         for exp, coeff in data.iteritems():
  77             if order > exp[index]:
  78                 continue
  79             new_exp = exp.copy()
  80             new_exp[index] -= order
  81             new_coeff = coeff
  82             for i in range(order):
  83                 new_coeff = new_coeff * (exp[index] - i)
  84             new_data[new_exp] = new_coeff
  85         return cls(new_data)
  86
  87     def expand(self, cls, expr):
  88         new_data = {}
  89         for exp, coeff in expr.data.iteritems():
  90             if isinstance(coeff, Expr):
  91                 new_data[exp] = coeff.expand()
  92             else:
  93                 new_data[exp] = coeff
  94         return cls(new_data)
  95
  96     def evalf(self, cls, expr, n):
  97         new_data = {}
  98         for exp, coeff in expr.data.iteritems():
  99             if isinstance(coeff, Expr):
 100                 new_data[exp] = coeff.evalf(n)
 101             else:
 102                 new_data[exp] = coeff
 103         return cls(new_data)
 104
 105
 106 class DensepolyHead(Head):
 107     """
 108     DensepolyHead is a head for dense polynomials represented
 109     as n-tuple of coefficients, data is a 2-tuple (symbol, coeffseq).
 110     """
 111     def __repr__(self): return 'DENSE_POLY'
 112
 113     def data_to_str_and_precedence(self, cls, (symbol, data)):
 114         if not isinstance(symbol, cls):
 115             symbol = cls(SYMBOL, symbol)
 116         terms = []
 117         for exp, coeff in enumerate(data):
 118             if coeff:
 119                 terms.append(cls(TERM_COEFF, (cls(POW, (symbol, exp)), coeff)))
 120         return ADD.data_to_str_and_precedence(cls, terms)
 121
 122 SPARSE_POLY = SparsepolyHead()
 123 DENSE_POLY = DensepolyHead()