jscode.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339
  1. """
  2. Javascript code printer
  3. The JavascriptCodePrinter converts single SymPy expressions into single
  4. Javascript expressions, using the functions defined in the Javascript
  5. Math object where possible.
  6. """
  7. from __future__ import annotations
  8. from typing import Any
  9. from sympy.core import S
  10. from sympy.core.numbers import equal_valued
  11. from sympy.printing.codeprinter import CodePrinter
  12. from sympy.printing.precedence import precedence, PRECEDENCE
  13. # dictionary mapping SymPy function to (argument_conditions, Javascript_function).
  14. # Used in JavascriptCodePrinter._print_Function(self)
  15. known_functions = {
  16. 'Abs': 'Math.abs',
  17. 'acos': 'Math.acos',
  18. 'acosh': 'Math.acosh',
  19. 'asin': 'Math.asin',
  20. 'asinh': 'Math.asinh',
  21. 'atan': 'Math.atan',
  22. 'atan2': 'Math.atan2',
  23. 'atanh': 'Math.atanh',
  24. 'ceiling': 'Math.ceil',
  25. 'cos': 'Math.cos',
  26. 'cosh': 'Math.cosh',
  27. 'exp': 'Math.exp',
  28. 'floor': 'Math.floor',
  29. 'log': 'Math.log',
  30. 'Max': 'Math.max',
  31. 'Min': 'Math.min',
  32. 'sign': 'Math.sign',
  33. 'sin': 'Math.sin',
  34. 'sinh': 'Math.sinh',
  35. 'tan': 'Math.tan',
  36. 'tanh': 'Math.tanh',
  37. }
  38. class JavascriptCodePrinter(CodePrinter):
  39. """"A Printer to convert Python expressions to strings of JavaScript code
  40. """
  41. printmethod = '_javascript'
  42. language = 'JavaScript'
  43. _default_settings: dict[str, Any] = {
  44. 'order': None,
  45. 'full_prec': 'auto',
  46. 'precision': 17,
  47. 'user_functions': {},
  48. 'human': True,
  49. 'allow_unknown_functions': False,
  50. 'contract': True,
  51. }
  52. def __init__(self, settings={}):
  53. CodePrinter.__init__(self, settings)
  54. self.known_functions = dict(known_functions)
  55. userfuncs = settings.get('user_functions', {})
  56. self.known_functions.update(userfuncs)
  57. def _rate_index_position(self, p):
  58. return p*5
  59. def _get_statement(self, codestring):
  60. return "%s;" % codestring
  61. def _get_comment(self, text):
  62. return "// {}".format(text)
  63. def _declare_number_const(self, name, value):
  64. return "var {} = {};".format(name, value.evalf(self._settings['precision']))
  65. def _format_code(self, lines):
  66. return self.indent_code(lines)
  67. def _traverse_matrix_indices(self, mat):
  68. rows, cols = mat.shape
  69. return ((i, j) for i in range(rows) for j in range(cols))
  70. def _get_loop_opening_ending(self, indices):
  71. open_lines = []
  72. close_lines = []
  73. loopstart = "for (var %(varble)s=%(start)s; %(varble)s<%(end)s; %(varble)s++){"
  74. for i in indices:
  75. # Javascript arrays start at 0 and end at dimension-1
  76. open_lines.append(loopstart % {
  77. 'varble': self._print(i.label),
  78. 'start': self._print(i.lower),
  79. 'end': self._print(i.upper + 1)})
  80. close_lines.append("}")
  81. return open_lines, close_lines
  82. def _print_Pow(self, expr):
  83. PREC = precedence(expr)
  84. if equal_valued(expr.exp, -1):
  85. return '1/%s' % (self.parenthesize(expr.base, PREC))
  86. elif equal_valued(expr.exp, 0.5):
  87. return 'Math.sqrt(%s)' % self._print(expr.base)
  88. elif expr.exp == S.One/3:
  89. return 'Math.cbrt(%s)' % self._print(expr.base)
  90. else:
  91. return 'Math.pow(%s, %s)' % (self._print(expr.base),
  92. self._print(expr.exp))
  93. def _print_Rational(self, expr):
  94. p, q = int(expr.p), int(expr.q)
  95. return '%d/%d' % (p, q)
  96. def _print_Mod(self, expr):
  97. num, den = expr.args
  98. PREC = precedence(expr)
  99. snum, sden = [self.parenthesize(arg, PREC) for arg in expr.args]
  100. # % is remainder (same sign as numerator), not modulo (same sign as
  101. # denominator), in js. Hence, % only works as modulo if both numbers
  102. # have the same sign
  103. if (num.is_nonnegative and den.is_nonnegative or
  104. num.is_nonpositive and den.is_nonpositive):
  105. return f"{snum} % {sden}"
  106. return f"(({snum} % {sden}) + {sden}) % {sden}"
  107. def _print_Relational(self, expr):
  108. lhs_code = self._print(expr.lhs)
  109. rhs_code = self._print(expr.rhs)
  110. op = expr.rel_op
  111. return "{} {} {}".format(lhs_code, op, rhs_code)
  112. def _print_Indexed(self, expr):
  113. # calculate index for 1d array
  114. dims = expr.shape
  115. elem = S.Zero
  116. offset = S.One
  117. for i in reversed(range(expr.rank)):
  118. elem += expr.indices[i]*offset
  119. offset *= dims[i]
  120. return "%s[%s]" % (self._print(expr.base.label), self._print(elem))
  121. def _print_Idx(self, expr):
  122. return self._print(expr.label)
  123. def _print_Exp1(self, expr):
  124. return "Math.E"
  125. def _print_Pi(self, expr):
  126. return 'Math.PI'
  127. def _print_Infinity(self, expr):
  128. return 'Number.POSITIVE_INFINITY'
  129. def _print_NegativeInfinity(self, expr):
  130. return 'Number.NEGATIVE_INFINITY'
  131. def _print_Piecewise(self, expr):
  132. from sympy.codegen.ast import Assignment
  133. if expr.args[-1].cond != True:
  134. # We need the last conditional to be a True, otherwise the resulting
  135. # function may not return a result.
  136. raise ValueError("All Piecewise expressions must contain an "
  137. "(expr, True) statement to be used as a default "
  138. "condition. Without one, the generated "
  139. "expression may not evaluate to anything under "
  140. "some condition.")
  141. lines = []
  142. if expr.has(Assignment):
  143. for i, (e, c) in enumerate(expr.args):
  144. if i == 0:
  145. lines.append("if (%s) {" % self._print(c))
  146. elif i == len(expr.args) - 1 and c == True:
  147. lines.append("else {")
  148. else:
  149. lines.append("else if (%s) {" % self._print(c))
  150. code0 = self._print(e)
  151. lines.append(code0)
  152. lines.append("}")
  153. return "\n".join(lines)
  154. else:
  155. # The piecewise was used in an expression, need to do inline
  156. # operators. This has the downside that inline operators will
  157. # not work for statements that span multiple lines (Matrix or
  158. # Indexed expressions).
  159. ecpairs = ["((%s) ? (\n%s\n)\n" % (self._print(c), self._print(e))
  160. for e, c in expr.args[:-1]]
  161. last_line = ": (\n%s\n)" % self._print(expr.args[-1].expr)
  162. return ": ".join(ecpairs) + last_line + " ".join([")"*len(ecpairs)])
  163. def _print_MatrixElement(self, expr):
  164. return "{}[{}]".format(self.parenthesize(expr.parent,
  165. PRECEDENCE["Atom"], strict=True),
  166. expr.j + expr.i*expr.parent.shape[1])
  167. def indent_code(self, code):
  168. """Accepts a string of code or a list of code lines"""
  169. if isinstance(code, str):
  170. code_lines = self.indent_code(code.splitlines(True))
  171. return ''.join(code_lines)
  172. tab = " "
  173. inc_token = ('{', '(', '{\n', '(\n')
  174. dec_token = ('}', ')')
  175. code = [ line.lstrip(' \t') for line in code ]
  176. increase = [ int(any(map(line.endswith, inc_token))) for line in code ]
  177. decrease = [ int(any(map(line.startswith, dec_token)))
  178. for line in code ]
  179. pretty = []
  180. level = 0
  181. for n, line in enumerate(code):
  182. if line in ('', '\n'):
  183. pretty.append(line)
  184. continue
  185. level -= decrease[n]
  186. pretty.append("%s%s" % (tab*level, line))
  187. level += increase[n]
  188. return pretty
  189. def jscode(expr, assign_to=None, **settings):
  190. """Converts an expr to a string of javascript code
  191. Parameters
  192. ==========
  193. expr : Expr
  194. A SymPy expression to be converted.
  195. assign_to : optional
  196. When given, the argument is used as the name of the variable to which
  197. the expression is assigned. Can be a string, ``Symbol``,
  198. ``MatrixSymbol``, or ``Indexed`` type. This is helpful in case of
  199. line-wrapping, or for expressions that generate multi-line statements.
  200. precision : integer, optional
  201. The precision for numbers such as pi [default=15].
  202. user_functions : dict, optional
  203. A dictionary where keys are ``FunctionClass`` instances and values are
  204. their string representations. Alternatively, the dictionary value can
  205. be a list of tuples i.e. [(argument_test, js_function_string)]. See
  206. below for examples.
  207. human : bool, optional
  208. If True, the result is a single string that may contain some constant
  209. declarations for the number symbols. If False, the same information is
  210. returned in a tuple of (symbols_to_declare, not_supported_functions,
  211. code_text). [default=True].
  212. contract: bool, optional
  213. If True, ``Indexed`` instances are assumed to obey tensor contraction
  214. rules and the corresponding nested loops over indices are generated.
  215. Setting contract=False will not generate loops, instead the user is
  216. responsible to provide values for the indices in the code.
  217. [default=True].
  218. Examples
  219. ========
  220. >>> from sympy import jscode, symbols, Rational, sin, ceiling, Abs
  221. >>> x, tau = symbols("x, tau")
  222. >>> jscode((2*tau)**Rational(7, 2))
  223. '8*Math.sqrt(2)*Math.pow(tau, 7/2)'
  224. >>> jscode(sin(x), assign_to="s")
  225. 's = Math.sin(x);'
  226. Custom printing can be defined for certain types by passing a dictionary of
  227. "type" : "function" to the ``user_functions`` kwarg. Alternatively, the
  228. dictionary value can be a list of tuples i.e. [(argument_test,
  229. js_function_string)].
  230. >>> custom_functions = {
  231. ... "ceiling": "CEIL",
  232. ... "Abs": [(lambda x: not x.is_integer, "fabs"),
  233. ... (lambda x: x.is_integer, "ABS")]
  234. ... }
  235. >>> jscode(Abs(x) + ceiling(x), user_functions=custom_functions)
  236. 'fabs(x) + CEIL(x)'
  237. ``Piecewise`` expressions are converted into conditionals. If an
  238. ``assign_to`` variable is provided an if statement is created, otherwise
  239. the ternary operator is used. Note that if the ``Piecewise`` lacks a
  240. default term, represented by ``(expr, True)`` then an error will be thrown.
  241. This is to prevent generating an expression that may not evaluate to
  242. anything.
  243. >>> from sympy import Piecewise
  244. >>> expr = Piecewise((x + 1, x > 0), (x, True))
  245. >>> print(jscode(expr, tau))
  246. if (x > 0) {
  247. tau = x + 1;
  248. }
  249. else {
  250. tau = x;
  251. }
  252. Support for loops is provided through ``Indexed`` types. With
  253. ``contract=True`` these expressions will be turned into loops, whereas
  254. ``contract=False`` will just print the assignment expression that should be
  255. looped over:
  256. >>> from sympy import Eq, IndexedBase, Idx
  257. >>> len_y = 5
  258. >>> y = IndexedBase('y', shape=(len_y,))
  259. >>> t = IndexedBase('t', shape=(len_y,))
  260. >>> Dy = IndexedBase('Dy', shape=(len_y-1,))
  261. >>> i = Idx('i', len_y-1)
  262. >>> e=Eq(Dy[i], (y[i+1]-y[i])/(t[i+1]-t[i]))
  263. >>> jscode(e.rhs, assign_to=e.lhs, contract=False)
  264. 'Dy[i] = (y[i + 1] - y[i])/(t[i + 1] - t[i]);'
  265. Matrices are also supported, but a ``MatrixSymbol`` of the same dimensions
  266. must be provided to ``assign_to``. Note that any expression that can be
  267. generated normally can also exist inside a Matrix:
  268. >>> from sympy import Matrix, MatrixSymbol
  269. >>> mat = Matrix([x**2, Piecewise((x + 1, x > 0), (x, True)), sin(x)])
  270. >>> A = MatrixSymbol('A', 3, 1)
  271. >>> print(jscode(mat, A))
  272. A[0] = Math.pow(x, 2);
  273. if (x > 0) {
  274. A[1] = x + 1;
  275. }
  276. else {
  277. A[1] = x;
  278. }
  279. A[2] = Math.sin(x);
  280. """
  281. return JavascriptCodePrinter(settings).doprint(expr, assign_to)
  282. def print_jscode(expr, **settings):
  283. """Prints the Javascript representation of the given expression.
  284. See jscode for the meaning of the optional arguments.
  285. """
  286. print(jscode(expr, **settings))