jscode.py 12 KB

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