mathematica.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353
  1. """
  2. Mathematica code printer
  3. """
  4. from typing import Any, Dict as tDict, Set as tSet, Tuple as tTuple
  5. from sympy.core import Basic, Expr, Float
  6. from sympy.core.sorting import default_sort_key
  7. from sympy.printing.codeprinter import CodePrinter
  8. from sympy.printing.precedence import precedence
  9. # Used in MCodePrinter._print_Function(self)
  10. known_functions = {
  11. "exp": [(lambda x: True, "Exp")],
  12. "log": [(lambda x: True, "Log")],
  13. "sin": [(lambda x: True, "Sin")],
  14. "cos": [(lambda x: True, "Cos")],
  15. "tan": [(lambda x: True, "Tan")],
  16. "cot": [(lambda x: True, "Cot")],
  17. "sec": [(lambda x: True, "Sec")],
  18. "csc": [(lambda x: True, "Csc")],
  19. "asin": [(lambda x: True, "ArcSin")],
  20. "acos": [(lambda x: True, "ArcCos")],
  21. "atan": [(lambda x: True, "ArcTan")],
  22. "acot": [(lambda x: True, "ArcCot")],
  23. "asec": [(lambda x: True, "ArcSec")],
  24. "acsc": [(lambda x: True, "ArcCsc")],
  25. "atan2": [(lambda *x: True, "ArcTan")],
  26. "sinh": [(lambda x: True, "Sinh")],
  27. "cosh": [(lambda x: True, "Cosh")],
  28. "tanh": [(lambda x: True, "Tanh")],
  29. "coth": [(lambda x: True, "Coth")],
  30. "sech": [(lambda x: True, "Sech")],
  31. "csch": [(lambda x: True, "Csch")],
  32. "asinh": [(lambda x: True, "ArcSinh")],
  33. "acosh": [(lambda x: True, "ArcCosh")],
  34. "atanh": [(lambda x: True, "ArcTanh")],
  35. "acoth": [(lambda x: True, "ArcCoth")],
  36. "asech": [(lambda x: True, "ArcSech")],
  37. "acsch": [(lambda x: True, "ArcCsch")],
  38. "sinc": [(lambda x: True, "Sinc")],
  39. "conjugate": [(lambda x: True, "Conjugate")],
  40. "Max": [(lambda *x: True, "Max")],
  41. "Min": [(lambda *x: True, "Min")],
  42. "erf": [(lambda x: True, "Erf")],
  43. "erf2": [(lambda *x: True, "Erf")],
  44. "erfc": [(lambda x: True, "Erfc")],
  45. "erfi": [(lambda x: True, "Erfi")],
  46. "erfinv": [(lambda x: True, "InverseErf")],
  47. "erfcinv": [(lambda x: True, "InverseErfc")],
  48. "erf2inv": [(lambda *x: True, "InverseErf")],
  49. "expint": [(lambda *x: True, "ExpIntegralE")],
  50. "Ei": [(lambda x: True, "ExpIntegralEi")],
  51. "fresnelc": [(lambda x: True, "FresnelC")],
  52. "fresnels": [(lambda x: True, "FresnelS")],
  53. "gamma": [(lambda x: True, "Gamma")],
  54. "uppergamma": [(lambda *x: True, "Gamma")],
  55. "polygamma": [(lambda *x: True, "PolyGamma")],
  56. "loggamma": [(lambda x: True, "LogGamma")],
  57. "beta": [(lambda *x: True, "Beta")],
  58. "Ci": [(lambda x: True, "CosIntegral")],
  59. "Si": [(lambda x: True, "SinIntegral")],
  60. "Chi": [(lambda x: True, "CoshIntegral")],
  61. "Shi": [(lambda x: True, "SinhIntegral")],
  62. "li": [(lambda x: True, "LogIntegral")],
  63. "factorial": [(lambda x: True, "Factorial")],
  64. "factorial2": [(lambda x: True, "Factorial2")],
  65. "subfactorial": [(lambda x: True, "Subfactorial")],
  66. "catalan": [(lambda x: True, "CatalanNumber")],
  67. "harmonic": [(lambda *x: True, "HarmonicNumber")],
  68. "lucas": [(lambda x: True, "LucasL")],
  69. "RisingFactorial": [(lambda *x: True, "Pochhammer")],
  70. "FallingFactorial": [(lambda *x: True, "FactorialPower")],
  71. "laguerre": [(lambda *x: True, "LaguerreL")],
  72. "assoc_laguerre": [(lambda *x: True, "LaguerreL")],
  73. "hermite": [(lambda *x: True, "HermiteH")],
  74. "jacobi": [(lambda *x: True, "JacobiP")],
  75. "gegenbauer": [(lambda *x: True, "GegenbauerC")],
  76. "chebyshevt": [(lambda *x: True, "ChebyshevT")],
  77. "chebyshevu": [(lambda *x: True, "ChebyshevU")],
  78. "legendre": [(lambda *x: True, "LegendreP")],
  79. "assoc_legendre": [(lambda *x: True, "LegendreP")],
  80. "mathieuc": [(lambda *x: True, "MathieuC")],
  81. "mathieus": [(lambda *x: True, "MathieuS")],
  82. "mathieucprime": [(lambda *x: True, "MathieuCPrime")],
  83. "mathieusprime": [(lambda *x: True, "MathieuSPrime")],
  84. "stieltjes": [(lambda x: True, "StieltjesGamma")],
  85. "elliptic_e": [(lambda *x: True, "EllipticE")],
  86. "elliptic_f": [(lambda *x: True, "EllipticE")],
  87. "elliptic_k": [(lambda x: True, "EllipticK")],
  88. "elliptic_pi": [(lambda *x: True, "EllipticPi")],
  89. "zeta": [(lambda *x: True, "Zeta")],
  90. "dirichlet_eta": [(lambda x: True, "DirichletEta")],
  91. "riemann_xi": [(lambda x: True, "RiemannXi")],
  92. "besseli": [(lambda *x: True, "BesselI")],
  93. "besselj": [(lambda *x: True, "BesselJ")],
  94. "besselk": [(lambda *x: True, "BesselK")],
  95. "bessely": [(lambda *x: True, "BesselY")],
  96. "hankel1": [(lambda *x: True, "HankelH1")],
  97. "hankel2": [(lambda *x: True, "HankelH2")],
  98. "airyai": [(lambda x: True, "AiryAi")],
  99. "airybi": [(lambda x: True, "AiryBi")],
  100. "airyaiprime": [(lambda x: True, "AiryAiPrime")],
  101. "airybiprime": [(lambda x: True, "AiryBiPrime")],
  102. "polylog": [(lambda *x: True, "PolyLog")],
  103. "lerchphi": [(lambda *x: True, "LerchPhi")],
  104. "gcd": [(lambda *x: True, "GCD")],
  105. "lcm": [(lambda *x: True, "LCM")],
  106. "jn": [(lambda *x: True, "SphericalBesselJ")],
  107. "yn": [(lambda *x: True, "SphericalBesselY")],
  108. "hyper": [(lambda *x: True, "HypergeometricPFQ")],
  109. "meijerg": [(lambda *x: True, "MeijerG")],
  110. "appellf1": [(lambda *x: True, "AppellF1")],
  111. "DiracDelta": [(lambda x: True, "DiracDelta")],
  112. "Heaviside": [(lambda x: True, "HeavisideTheta")],
  113. "KroneckerDelta": [(lambda *x: True, "KroneckerDelta")],
  114. "sqrt": [(lambda x: True, "Sqrt")], # For automatic rewrites
  115. }
  116. class MCodePrinter(CodePrinter):
  117. """A printer to convert Python expressions to
  118. strings of the Wolfram's Mathematica code
  119. """
  120. printmethod = "_mcode"
  121. language = "Wolfram Language"
  122. _default_settings = {
  123. 'order': None,
  124. 'full_prec': 'auto',
  125. 'precision': 15,
  126. 'user_functions': {},
  127. 'human': True,
  128. 'allow_unknown_functions': False,
  129. } # type: tDict[str, Any]
  130. _number_symbols = set() # type: tSet[tTuple[Expr, Float]]
  131. _not_supported = set() # type: tSet[Basic]
  132. def __init__(self, settings={}):
  133. """Register function mappings supplied by user"""
  134. CodePrinter.__init__(self, settings)
  135. self.known_functions = dict(known_functions)
  136. userfuncs = settings.get('user_functions', {}).copy()
  137. for k, v in userfuncs.items():
  138. if not isinstance(v, list):
  139. userfuncs[k] = [(lambda *x: True, v)]
  140. self.known_functions.update(userfuncs)
  141. def _format_code(self, lines):
  142. return lines
  143. def _print_Pow(self, expr):
  144. PREC = precedence(expr)
  145. return '%s^%s' % (self.parenthesize(expr.base, PREC),
  146. self.parenthesize(expr.exp, PREC))
  147. def _print_Mul(self, expr):
  148. PREC = precedence(expr)
  149. c, nc = expr.args_cnc()
  150. res = super()._print_Mul(expr.func(*c))
  151. if nc:
  152. res += '*'
  153. res += '**'.join(self.parenthesize(a, PREC) for a in nc)
  154. return res
  155. def _print_Relational(self, expr):
  156. lhs_code = self._print(expr.lhs)
  157. rhs_code = self._print(expr.rhs)
  158. op = expr.rel_op
  159. return "{} {} {}".format(lhs_code, op, rhs_code)
  160. # Primitive numbers
  161. def _print_Zero(self, expr):
  162. return '0'
  163. def _print_One(self, expr):
  164. return '1'
  165. def _print_NegativeOne(self, expr):
  166. return '-1'
  167. def _print_Half(self, expr):
  168. return '1/2'
  169. def _print_ImaginaryUnit(self, expr):
  170. return 'I'
  171. # Infinity and invalid numbers
  172. def _print_Infinity(self, expr):
  173. return 'Infinity'
  174. def _print_NegativeInfinity(self, expr):
  175. return '-Infinity'
  176. def _print_ComplexInfinity(self, expr):
  177. return 'ComplexInfinity'
  178. def _print_NaN(self, expr):
  179. return 'Indeterminate'
  180. # Mathematical constants
  181. def _print_Exp1(self, expr):
  182. return 'E'
  183. def _print_Pi(self, expr):
  184. return 'Pi'
  185. def _print_GoldenRatio(self, expr):
  186. return 'GoldenRatio'
  187. def _print_TribonacciConstant(self, expr):
  188. expanded = expr.expand(func=True)
  189. PREC = precedence(expr)
  190. return self.parenthesize(expanded, PREC)
  191. def _print_EulerGamma(self, expr):
  192. return 'EulerGamma'
  193. def _print_Catalan(self, expr):
  194. return 'Catalan'
  195. def _print_list(self, expr):
  196. return '{' + ', '.join(self.doprint(a) for a in expr) + '}'
  197. _print_tuple = _print_list
  198. _print_Tuple = _print_list
  199. def _print_ImmutableDenseMatrix(self, expr):
  200. return self.doprint(expr.tolist())
  201. def _print_ImmutableSparseMatrix(self, expr):
  202. def print_rule(pos, val):
  203. return '{} -> {}'.format(
  204. self.doprint((pos[0]+1, pos[1]+1)), self.doprint(val))
  205. def print_data():
  206. items = sorted(expr.todok().items(), key=default_sort_key)
  207. return '{' + \
  208. ', '.join(print_rule(k, v) for k, v in items) + \
  209. '}'
  210. def print_dims():
  211. return self.doprint(expr.shape)
  212. return 'SparseArray[{}, {}]'.format(print_data(), print_dims())
  213. def _print_ImmutableDenseNDimArray(self, expr):
  214. return self.doprint(expr.tolist())
  215. def _print_ImmutableSparseNDimArray(self, expr):
  216. def print_string_list(string_list):
  217. return '{' + ', '.join(a for a in string_list) + '}'
  218. def to_mathematica_index(*args):
  219. """Helper function to change Python style indexing to
  220. Pathematica indexing.
  221. Python indexing (0, 1 ... n-1)
  222. -> Mathematica indexing (1, 2 ... n)
  223. """
  224. return tuple(i + 1 for i in args)
  225. def print_rule(pos, val):
  226. """Helper function to print a rule of Mathematica"""
  227. return '{} -> {}'.format(self.doprint(pos), self.doprint(val))
  228. def print_data():
  229. """Helper function to print data part of Mathematica
  230. sparse array.
  231. It uses the fourth notation ``SparseArray[data,{d1,d2,...}]``
  232. from
  233. https://reference.wolfram.com/language/ref/SparseArray.html
  234. ``data`` must be formatted with rule.
  235. """
  236. return print_string_list(
  237. [print_rule(
  238. to_mathematica_index(*(expr._get_tuple_index(key))),
  239. value)
  240. for key, value in sorted(expr._sparse_array.items())]
  241. )
  242. def print_dims():
  243. """Helper function to print dimensions part of Mathematica
  244. sparse array.
  245. It uses the fourth notation ``SparseArray[data,{d1,d2,...}]``
  246. from
  247. https://reference.wolfram.com/language/ref/SparseArray.html
  248. """
  249. return self.doprint(expr.shape)
  250. return 'SparseArray[{}, {}]'.format(print_data(), print_dims())
  251. def _print_Function(self, expr):
  252. if expr.func.__name__ in self.known_functions:
  253. cond_mfunc = self.known_functions[expr.func.__name__]
  254. for cond, mfunc in cond_mfunc:
  255. if cond(*expr.args):
  256. return "%s[%s]" % (mfunc, self.stringify(expr.args, ", "))
  257. elif expr.func.__name__ in self._rewriteable_functions:
  258. # Simple rewrite to supported function possible
  259. target_f, required_fs = self._rewriteable_functions[expr.func.__name__]
  260. if self._can_print(target_f) and all(self._can_print(f) for f in required_fs):
  261. return self._print(expr.rewrite(target_f))
  262. return expr.func.__name__ + "[%s]" % self.stringify(expr.args, ", ")
  263. _print_MinMaxBase = _print_Function
  264. def _print_LambertW(self, expr):
  265. if len(expr.args) == 1:
  266. return "ProductLog[{}]".format(self._print(expr.args[0]))
  267. return "ProductLog[{}, {}]".format(
  268. self._print(expr.args[1]), self._print(expr.args[0]))
  269. def _print_Integral(self, expr):
  270. if len(expr.variables) == 1 and not expr.limits[0][1:]:
  271. args = [expr.args[0], expr.variables[0]]
  272. else:
  273. args = expr.args
  274. return "Hold[Integrate[" + ', '.join(self.doprint(a) for a in args) + "]]"
  275. def _print_Sum(self, expr):
  276. return "Hold[Sum[" + ', '.join(self.doprint(a) for a in expr.args) + "]]"
  277. def _print_Derivative(self, expr):
  278. dexpr = expr.expr
  279. dvars = [i[0] if i[1] == 1 else i for i in expr.variable_count]
  280. return "Hold[D[" + ', '.join(self.doprint(a) for a in [dexpr] + dvars) + "]]"
  281. def _get_comment(self, text):
  282. return "(* {} *)".format(text)
  283. def mathematica_code(expr, **settings):
  284. r"""Converts an expr to a string of the Wolfram Mathematica code
  285. Examples
  286. ========
  287. >>> from sympy import mathematica_code as mcode, symbols, sin
  288. >>> x = symbols('x')
  289. >>> mcode(sin(x).series(x).removeO())
  290. '(1/120)*x^5 - 1/6*x^3 + x'
  291. """
  292. return MCodePrinter(settings).doprint(expr)