sympycore/heads/apply.py

   1
   2 __all__ = ['APPLY']
   3
   4 from .base import Head, heads_precedence
   5 from .functional import FunctionalHead
   6
   7 from ..core import init_module, Expr
   8 init_module.import_heads()
   9 init_module.import_numbers()
  10 init_module.import_lowlevel_operations()
  11
  12 class ApplyHead(FunctionalHead):
  13     """
  14     ApplyHead is a head for n-ary apply operation,
  15     data is a (1+n)-tuple of expression operands
  16     """
  17     op_mth = '__call__'
  18     parenthesis = '()'
  19
  20     def is_data_ok(self, cls, data):
  21         if isinstance(data, tuple) and len(data)==2:
  22             func, args = data
  23             fcls = cls.get_function_algebra()
  24             if not isinstance(func, fcls):
  25                 return '%s data[0] must be %s instance but got %s' % (self, fcls, type(func))
  26             msg = func.head.is_data_ok(type(func), func.data)
  27             if msg:
  28                 return '%s data=%r: %s' % (func.head, func.pair, msg)
  29             if type(args) is tuple:
  30                 for i,a in enumerate(args):
  31                     acls = func.get_argument_algebra(i)
  32                     if not isinstance(a, Expr) or type(a) is not acls:
  33                         return '%s data[1][%s] must be %s instance but got %s' % (self, i, acls, type(a))
  34             else:
  35                 return '%s data[1] must be tuple but got %s' % (self, type(args))
  36         else:
  37             return '%s data instance must be 2-tuple' % (self)
  38
  39     def __repr__(self): return 'APPLY'
  40
  41     def new(self, cls, (func, args), evaluate=True):
  42         if not isinstance(func, Expr):
  43             fcls = cls.get_function_algebra()
  44             func = fcls(CALLABLE, func)
  45         return cls(APPLY, (func, args))
  46
  47     def commutative_mul_number(self, cls, lhs, rhs):
  48         return term_coeff_new(cls, (lhs, rhs))
  49
  50     def pow(self, cls, base, exp):
  51         return POW.new(cls, (base, exp))
  52
  53     pow_number = pow
  54
  55     def scan(self, proc, cls, data, target):
  56         f, args = data
  57         f.head.scan(proc, cls, f.data, target)
  58         for arg in args:
  59             arg.head.scan(proc, cls, arg.data, target)
  60         proc(cls, self, data, target)
  61
  62     def walk(self, func, cls, data, target):
  63         f, args = data
  64         new_f = f.head.walk(func, cls, f.data, f)
  65         new_args = []
  66         flag = f is not new_f
  67         for arg in args:
  68             h, d = arg.pair
  69             new_arg = arg.head.walk(func, cls, arg.data, arg)
  70             if arg is not new_arg:
  71                 flag = True
  72             new_args.append(new_arg)
  73         if flag:
  74             new_args = tuple(new_args)
  75             if new_f.head is CALLABLE:
  76                 r = new_f.data(*new_args)
  77             else:
  78                 r = cls(APPLY, (new_f, new_args))
  79             return func(cls, r.head, r.data, r)
  80         return func(cls, self, data, target)
  81
  82     def expand(self, cls, expr):
  83         f, args = expr.data
  84         new_f = f
  85         new_args = [a.expand() for a in args]
  86         return new_f (*new_args)
  87         return expr
  88
  89     def expand_intpow(self, cls, base, exp):
  90         return cls(POW, (base, exp))
  91
  92     def diff(self, cls, data, expr, symbol, order, cache={}):
  93         key = (expr, symbol, order)
  94         result = cache.get(key)
  95         if result is not None:
  96             return result
  97         if isinstance(order, inttypes):
  98             pass
  99         key1 = (expr, symbol, 1)
 100         result = cache.get(key1)
 101         if result is None:
 102             func, args = data
 103             fcls = cls.get_function_algebra()
 104             dcls = fcls.get_differential_algebra()
 105             if len(args)==1:
 106                 arg = args[0]
 107                 da = arg.head.diff(cls, arg.data, arg, symbol, 1, cache=cache)
 108                 if symbol not in da.symbols_data:
 109                     # argument is linear with respect to symbol
 110                     df = func.head.fdiff(fcls, func.data, func, 0, order)
 111                     if df is NotImplemented:
 112                         if isinstance(order, int):
 113                             df = func.head.fdiff(fcls, func.data, func, 0, 1)
 114                             result = df(*args) * da
 115                             if order>1:
 116                                 result = result.head.diff(cls, result.data, result, symbol, order-1, cache=cache)
 117                         else:
 118                             d = dcls(FDIFF, cls(NUMBER, 0))**order
 119                             df = fcls(APPLY, (d, (func,)))
 120                             result = df(*args) * da**order
 121                     else:
 122                         result = df(*args) * da**order
 123                 elif isinstance(order, int):
 124                     df = func.head.fdiff(fcls, func.data, func, 0, 1)
 125                     result = df(*args) * da
 126                     if order>1:
 127                         result = result.head.diff(cls, result.data, result, symbol, order-1, cache=cache)
 128                 else:
 129                     d = dcls(DIFF, cls(NUMBER, 0))**order
 130                     result = cls(APPLY, (d, (expr,)))
 131                 cache[key] = result
 132                 return result
 133             else:
 134                 result = cls(NUMBER, 0)
 135                 for i in range(len(args)):
 136                     arg = args[i]
 137                     da = arg.head.diff(cls, arg.data, arg, symbol, 1, cache=cache)
 138                     df = func.head.fdiff(fcls, func.data, func, i, 1)
 139                     result += df(*args) * da
 140             cache[key1] = result
 141         if order>1:
 142             result = result.head.diff(cls, result.data, result, symbol, order-1, cache=cache)
 143             cache[key] = result
 144         return result
 145
 146     def apply(self, cls, data, func, args):
 147         return cls(APPLY, (func, args))
 148
 149     def integrate_indefinite(self, cls, data, expr, x):
 150         if x not in expr.symbols_data:
 151             return expr * cls(SYMBOL, x)
 152         return cls(INTEGRAL_INDEFINITE, (expr, x))
 153
 154     def integrate_definite(self, cls, data, expr, x, a, b):
 155         if x not in expr.symbols_data:
 156             return expr * (b - a)
 157         return cls(INTEGRAL_DEFINITE, (expr, x, a, b))
 158
 159
 160
 161 APPLY = ApplyHead()