lambdarepr.py 6.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210
  1. from .pycode import (
  2. PythonCodePrinter,
  3. MpmathPrinter,
  4. )
  5. from .numpy import NumPyPrinter # NumPyPrinter is imported for backward compatibility
  6. from sympy.core.sorting import default_sort_key
  7. __all__ = [
  8. 'PythonCodePrinter',
  9. 'MpmathPrinter', # MpmathPrinter is published for backward compatibility
  10. 'NumPyPrinter',
  11. 'LambdaPrinter',
  12. 'NumPyPrinter',
  13. 'IntervalPrinter',
  14. 'lambdarepr',
  15. ]
  16. class LambdaPrinter(PythonCodePrinter):
  17. """
  18. This printer converts expressions into strings that can be used by
  19. lambdify.
  20. """
  21. printmethod = "_lambdacode"
  22. def _print_And(self, expr):
  23. result = ['(']
  24. for arg in sorted(expr.args, key=default_sort_key):
  25. result.extend(['(', self._print(arg), ')'])
  26. result.append(' and ')
  27. result = result[:-1]
  28. result.append(')')
  29. return ''.join(result)
  30. def _print_Or(self, expr):
  31. result = ['(']
  32. for arg in sorted(expr.args, key=default_sort_key):
  33. result.extend(['(', self._print(arg), ')'])
  34. result.append(' or ')
  35. result = result[:-1]
  36. result.append(')')
  37. return ''.join(result)
  38. def _print_Not(self, expr):
  39. result = ['(', 'not (', self._print(expr.args[0]), '))']
  40. return ''.join(result)
  41. def _print_BooleanTrue(self, expr):
  42. return "True"
  43. def _print_BooleanFalse(self, expr):
  44. return "False"
  45. def _print_ITE(self, expr):
  46. result = [
  47. '((', self._print(expr.args[1]),
  48. ') if (', self._print(expr.args[0]),
  49. ') else (', self._print(expr.args[2]), '))'
  50. ]
  51. return ''.join(result)
  52. def _print_NumberSymbol(self, expr):
  53. return str(expr)
  54. def _print_Pow(self, expr, **kwargs):
  55. # XXX Temporary workaround. Should Python math printer be
  56. # isolated from PythonCodePrinter?
  57. return super(PythonCodePrinter, self)._print_Pow(expr, **kwargs)
  58. # numexpr works by altering the string passed to numexpr.evaluate
  59. # rather than by populating a namespace. Thus a special printer...
  60. class NumExprPrinter(LambdaPrinter):
  61. # key, value pairs correspond to SymPy name and numexpr name
  62. # functions not appearing in this dict will raise a TypeError
  63. printmethod = "_numexprcode"
  64. _numexpr_functions = {
  65. 'sin' : 'sin',
  66. 'cos' : 'cos',
  67. 'tan' : 'tan',
  68. 'asin': 'arcsin',
  69. 'acos': 'arccos',
  70. 'atan': 'arctan',
  71. 'atan2' : 'arctan2',
  72. 'sinh' : 'sinh',
  73. 'cosh' : 'cosh',
  74. 'tanh' : 'tanh',
  75. 'asinh': 'arcsinh',
  76. 'acosh': 'arccosh',
  77. 'atanh': 'arctanh',
  78. 'ln' : 'log',
  79. 'log': 'log',
  80. 'exp': 'exp',
  81. 'sqrt' : 'sqrt',
  82. 'Abs' : 'abs',
  83. 'conjugate' : 'conj',
  84. 'im' : 'imag',
  85. 're' : 'real',
  86. 'where' : 'where',
  87. 'complex' : 'complex',
  88. 'contains' : 'contains',
  89. }
  90. def _print_ImaginaryUnit(self, expr):
  91. return '1j'
  92. def _print_seq(self, seq, delimiter=', '):
  93. # simplified _print_seq taken from pretty.py
  94. s = [self._print(item) for item in seq]
  95. if s:
  96. return delimiter.join(s)
  97. else:
  98. return ""
  99. def _print_Function(self, e):
  100. func_name = e.func.__name__
  101. nstr = self._numexpr_functions.get(func_name, None)
  102. if nstr is None:
  103. # check for implemented_function
  104. if hasattr(e, '_imp_'):
  105. return "(%s)" % self._print(e._imp_(*e.args))
  106. else:
  107. raise TypeError("numexpr does not support function '%s'" %
  108. func_name)
  109. return "%s(%s)" % (nstr, self._print_seq(e.args))
  110. def _print_Piecewise(self, expr):
  111. "Piecewise function printer"
  112. exprs = [self._print(arg.expr) for arg in expr.args]
  113. conds = [self._print(arg.cond) for arg in expr.args]
  114. # If [default_value, True] is a (expr, cond) sequence in a Piecewise object
  115. # it will behave the same as passing the 'default' kwarg to select()
  116. # *as long as* it is the last element in expr.args.
  117. # If this is not the case, it may be triggered prematurely.
  118. ans = []
  119. parenthesis_count = 0
  120. is_last_cond_True = False
  121. for cond, expr in zip(conds, exprs):
  122. if cond == 'True':
  123. ans.append(expr)
  124. is_last_cond_True = True
  125. break
  126. else:
  127. ans.append('where(%s, %s, ' % (cond, expr))
  128. parenthesis_count += 1
  129. if not is_last_cond_True:
  130. # simplest way to put a nan but raises
  131. # 'RuntimeWarning: invalid value encountered in log'
  132. ans.append('log(-1)')
  133. return ''.join(ans) + ')' * parenthesis_count
  134. def _print_ITE(self, expr):
  135. from sympy.functions.elementary.piecewise import Piecewise
  136. return self._print(expr.rewrite(Piecewise))
  137. def blacklisted(self, expr):
  138. raise TypeError("numexpr cannot be used with %s" %
  139. expr.__class__.__name__)
  140. # blacklist all Matrix printing
  141. _print_SparseRepMatrix = \
  142. _print_MutableSparseMatrix = \
  143. _print_ImmutableSparseMatrix = \
  144. _print_Matrix = \
  145. _print_DenseMatrix = \
  146. _print_MutableDenseMatrix = \
  147. _print_ImmutableMatrix = \
  148. _print_ImmutableDenseMatrix = \
  149. blacklisted
  150. # blacklist some Python expressions
  151. _print_list = \
  152. _print_tuple = \
  153. _print_Tuple = \
  154. _print_dict = \
  155. _print_Dict = \
  156. blacklisted
  157. def doprint(self, expr):
  158. lstr = super().doprint(expr)
  159. return "evaluate('%s', truediv=True)" % lstr
  160. class IntervalPrinter(MpmathPrinter, LambdaPrinter):
  161. """Use ``lambda`` printer but print numbers as ``mpi`` intervals. """
  162. def _print_Integer(self, expr):
  163. return "mpi('%s')" % super(PythonCodePrinter, self)._print_Integer(expr)
  164. def _print_Rational(self, expr):
  165. return "mpi('%s')" % super(PythonCodePrinter, self)._print_Rational(expr)
  166. def _print_Half(self, expr):
  167. return "mpi('%s')" % super(PythonCodePrinter, self)._print_Rational(expr)
  168. def _print_Pow(self, expr):
  169. return super(MpmathPrinter, self)._print_Pow(expr, rational=True)
  170. for k in NumExprPrinter._numexpr_functions:
  171. setattr(NumExprPrinter, '_print_%s' % k, NumExprPrinter._print_Function)
  172. def lambdarepr(expr, **settings):
  173. """
  174. Returns a string usable for lambdifying.
  175. """
  176. return LambdaPrinter(settings).doprint(expr)