codeprinter.py 35 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875
  1. from __future__ import annotations
  2. from typing import Any
  3. from functools import wraps
  4. from sympy.core import Add, Mul, Pow, S, sympify, Float
  5. from sympy.core.basic import Basic
  6. from sympy.core.expr import UnevaluatedExpr
  7. from sympy.core.function import Lambda
  8. from sympy.core.mul import _keep_coeff
  9. from sympy.core.sorting import default_sort_key
  10. from sympy.core.symbol import Symbol
  11. from sympy.functions.elementary.complexes import re
  12. from sympy.printing.str import StrPrinter
  13. from sympy.printing.precedence import precedence, PRECEDENCE
  14. class requires:
  15. """ Decorator for registering requirements on print methods. """
  16. def __init__(self, **kwargs):
  17. self._req = kwargs
  18. def __call__(self, method):
  19. def _method_wrapper(self_, *args, **kwargs):
  20. for k, v in self._req.items():
  21. getattr(self_, k).update(v)
  22. return method(self_, *args, **kwargs)
  23. return wraps(method)(_method_wrapper)
  24. class AssignmentError(Exception):
  25. """
  26. Raised if an assignment variable for a loop is missing.
  27. """
  28. pass
  29. def _convert_python_lists(arg):
  30. if isinstance(arg, list):
  31. from sympy.codegen.abstract_nodes import List
  32. return List(*(_convert_python_lists(e) for e in arg))
  33. elif isinstance(arg, tuple):
  34. return tuple(_convert_python_lists(e) for e in arg)
  35. else:
  36. return arg
  37. class CodePrinter(StrPrinter):
  38. """
  39. The base class for code-printing subclasses.
  40. """
  41. _operators = {
  42. 'and': '&&',
  43. 'or': '||',
  44. 'not': '!',
  45. }
  46. _default_settings: dict[str, Any] = {
  47. 'order': None,
  48. 'full_prec': 'auto',
  49. 'error_on_reserved': False,
  50. 'reserved_word_suffix': '_',
  51. 'human': True,
  52. 'inline': False,
  53. 'allow_unknown_functions': False,
  54. }
  55. # Functions which are "simple" to rewrite to other functions that
  56. # may be supported
  57. # function_to_rewrite : (function_to_rewrite_to, iterable_with_other_functions_required)
  58. _rewriteable_functions = {
  59. 'cot': ('tan', []),
  60. 'csc': ('sin', []),
  61. 'sec': ('cos', []),
  62. 'acot': ('atan', []),
  63. 'acsc': ('asin', []),
  64. 'asec': ('acos', []),
  65. 'coth': ('exp', []),
  66. 'csch': ('exp', []),
  67. 'sech': ('exp', []),
  68. 'acoth': ('log', []),
  69. 'acsch': ('log', []),
  70. 'asech': ('log', []),
  71. 'catalan': ('gamma', []),
  72. 'fibonacci': ('sqrt', []),
  73. 'lucas': ('sqrt', []),
  74. 'beta': ('gamma', []),
  75. 'sinc': ('sin', ['Piecewise']),
  76. 'Mod': ('floor', []),
  77. 'factorial': ('gamma', []),
  78. 'factorial2': ('gamma', ['Piecewise']),
  79. 'subfactorial': ('uppergamma', []),
  80. 'RisingFactorial': ('gamma', ['Piecewise']),
  81. 'FallingFactorial': ('gamma', ['Piecewise']),
  82. 'binomial': ('gamma', []),
  83. 'frac': ('floor', []),
  84. 'Max': ('Piecewise', []),
  85. 'Min': ('Piecewise', []),
  86. 'Heaviside': ('Piecewise', []),
  87. 'erf2': ('erf', []),
  88. 'erfc': ('erf', []),
  89. 'Li': ('li', []),
  90. 'Ei': ('li', []),
  91. 'dirichlet_eta': ('zeta', []),
  92. 'riemann_xi': ('zeta', ['gamma']),
  93. }
  94. def __init__(self, settings=None):
  95. super().__init__(settings=settings)
  96. if not hasattr(self, 'reserved_words'):
  97. self.reserved_words = set()
  98. def _handle_UnevaluatedExpr(self, expr):
  99. return expr.replace(re, lambda arg: arg if isinstance(
  100. arg, UnevaluatedExpr) and arg.args[0].is_real else re(arg))
  101. def doprint(self, expr, assign_to=None):
  102. """
  103. Print the expression as code.
  104. Parameters
  105. ----------
  106. expr : Expression
  107. The expression to be printed.
  108. assign_to : Symbol, string, MatrixSymbol, list of strings or Symbols (optional)
  109. If provided, the printed code will set the expression to a variable or multiple variables
  110. with the name or names given in ``assign_to``.
  111. """
  112. from sympy.matrices.expressions.matexpr import MatrixSymbol
  113. from sympy.codegen.ast import CodeBlock, Assignment
  114. def _handle_assign_to(expr, assign_to):
  115. if assign_to is None:
  116. return sympify(expr)
  117. if isinstance(assign_to, (list, tuple)):
  118. if len(expr) != len(assign_to):
  119. raise ValueError('Failed to assign an expression of length {} to {} variables'.format(len(expr), len(assign_to)))
  120. return CodeBlock(*[_handle_assign_to(lhs, rhs) for lhs, rhs in zip(expr, assign_to)])
  121. if isinstance(assign_to, str):
  122. if expr.is_Matrix:
  123. assign_to = MatrixSymbol(assign_to, *expr.shape)
  124. else:
  125. assign_to = Symbol(assign_to)
  126. elif not isinstance(assign_to, Basic):
  127. raise TypeError("{} cannot assign to object of type {}".format(
  128. type(self).__name__, type(assign_to)))
  129. return Assignment(assign_to, expr)
  130. expr = _convert_python_lists(expr)
  131. expr = _handle_assign_to(expr, assign_to)
  132. # Remove re(...) nodes due to UnevaluatedExpr.is_real always is None:
  133. expr = self._handle_UnevaluatedExpr(expr)
  134. # keep a set of expressions that are not strictly translatable to Code
  135. # and number constants that must be declared and initialized
  136. self._not_supported = set()
  137. self._number_symbols = set()
  138. lines = self._print(expr).splitlines()
  139. # format the output
  140. if self._settings["human"]:
  141. frontlines = []
  142. if self._not_supported:
  143. frontlines.append(self._get_comment(
  144. "Not supported in {}:".format(self.language)))
  145. for expr in sorted(self._not_supported, key=str):
  146. frontlines.append(self._get_comment(type(expr).__name__))
  147. for name, value in sorted(self._number_symbols, key=str):
  148. frontlines.append(self._declare_number_const(name, value))
  149. lines = frontlines + lines
  150. lines = self._format_code(lines)
  151. result = "\n".join(lines)
  152. else:
  153. lines = self._format_code(lines)
  154. num_syms = {(k, self._print(v)) for k, v in self._number_symbols}
  155. result = (num_syms, self._not_supported, "\n".join(lines))
  156. self._not_supported = set()
  157. self._number_symbols = set()
  158. return result
  159. def _doprint_loops(self, expr, assign_to=None):
  160. # Here we print an expression that contains Indexed objects, they
  161. # correspond to arrays in the generated code. The low-level implementation
  162. # involves looping over array elements and possibly storing results in temporary
  163. # variables or accumulate it in the assign_to object.
  164. if self._settings.get('contract', True):
  165. from sympy.tensor import get_contraction_structure
  166. # Setup loops over non-dummy indices -- all terms need these
  167. indices = self._get_expression_indices(expr, assign_to)
  168. # Setup loops over dummy indices -- each term needs separate treatment
  169. dummies = get_contraction_structure(expr)
  170. else:
  171. indices = []
  172. dummies = {None: (expr,)}
  173. openloop, closeloop = self._get_loop_opening_ending(indices)
  174. # terms with no summations first
  175. if None in dummies:
  176. text = StrPrinter.doprint(self, Add(*dummies[None]))
  177. else:
  178. # If all terms have summations we must initialize array to Zero
  179. text = StrPrinter.doprint(self, 0)
  180. # skip redundant assignments (where lhs == rhs)
  181. lhs_printed = self._print(assign_to)
  182. lines = []
  183. if text != lhs_printed:
  184. lines.extend(openloop)
  185. if assign_to is not None:
  186. text = self._get_statement("%s = %s" % (lhs_printed, text))
  187. lines.append(text)
  188. lines.extend(closeloop)
  189. # then terms with summations
  190. for d in dummies:
  191. if isinstance(d, tuple):
  192. indices = self._sort_optimized(d, expr)
  193. openloop_d, closeloop_d = self._get_loop_opening_ending(
  194. indices)
  195. for term in dummies[d]:
  196. if term in dummies and not ([list(f.keys()) for f in dummies[term]]
  197. == [[None] for f in dummies[term]]):
  198. # If one factor in the term has it's own internal
  199. # contractions, those must be computed first.
  200. # (temporary variables?)
  201. raise NotImplementedError(
  202. "FIXME: no support for contractions in factor yet")
  203. else:
  204. # We need the lhs expression as an accumulator for
  205. # the loops, i.e
  206. #
  207. # for (int d=0; d < dim; d++){
  208. # lhs[] = lhs[] + term[][d]
  209. # } ^.................. the accumulator
  210. #
  211. # We check if the expression already contains the
  212. # lhs, and raise an exception if it does, as that
  213. # syntax is currently undefined. FIXME: What would be
  214. # a good interpretation?
  215. if assign_to is None:
  216. raise AssignmentError(
  217. "need assignment variable for loops")
  218. if term.has(assign_to):
  219. raise ValueError("FIXME: lhs present in rhs,\
  220. this is undefined in CodePrinter")
  221. lines.extend(openloop)
  222. lines.extend(openloop_d)
  223. text = "%s = %s" % (lhs_printed, StrPrinter.doprint(
  224. self, assign_to + term))
  225. lines.append(self._get_statement(text))
  226. lines.extend(closeloop_d)
  227. lines.extend(closeloop)
  228. return "\n".join(lines)
  229. def _get_expression_indices(self, expr, assign_to):
  230. from sympy.tensor import get_indices
  231. rinds, junk = get_indices(expr)
  232. linds, junk = get_indices(assign_to)
  233. # support broadcast of scalar
  234. if linds and not rinds:
  235. rinds = linds
  236. if rinds != linds:
  237. raise ValueError("lhs indices must match non-dummy"
  238. " rhs indices in %s" % expr)
  239. return self._sort_optimized(rinds, assign_to)
  240. def _sort_optimized(self, indices, expr):
  241. from sympy.tensor.indexed import Indexed
  242. if not indices:
  243. return []
  244. # determine optimized loop order by giving a score to each index
  245. # the index with the highest score are put in the innermost loop.
  246. score_table = {}
  247. for i in indices:
  248. score_table[i] = 0
  249. arrays = expr.atoms(Indexed)
  250. for arr in arrays:
  251. for p, ind in enumerate(arr.indices):
  252. try:
  253. score_table[ind] += self._rate_index_position(p)
  254. except KeyError:
  255. pass
  256. return sorted(indices, key=lambda x: score_table[x])
  257. def _rate_index_position(self, p):
  258. """function to calculate score based on position among indices
  259. This method is used to sort loops in an optimized order, see
  260. CodePrinter._sort_optimized()
  261. """
  262. raise NotImplementedError("This function must be implemented by "
  263. "subclass of CodePrinter.")
  264. def _get_statement(self, codestring):
  265. """Formats a codestring with the proper line ending."""
  266. raise NotImplementedError("This function must be implemented by "
  267. "subclass of CodePrinter.")
  268. def _get_comment(self, text):
  269. """Formats a text string as a comment."""
  270. raise NotImplementedError("This function must be implemented by "
  271. "subclass of CodePrinter.")
  272. def _declare_number_const(self, name, value):
  273. """Declare a numeric constant at the top of a function"""
  274. raise NotImplementedError("This function must be implemented by "
  275. "subclass of CodePrinter.")
  276. def _format_code(self, lines):
  277. """Take in a list of lines of code, and format them accordingly.
  278. This may include indenting, wrapping long lines, etc..."""
  279. raise NotImplementedError("This function must be implemented by "
  280. "subclass of CodePrinter.")
  281. def _get_loop_opening_ending(self, indices):
  282. """Returns a tuple (open_lines, close_lines) containing lists
  283. of codelines"""
  284. raise NotImplementedError("This function must be implemented by "
  285. "subclass of CodePrinter.")
  286. def _print_Dummy(self, expr):
  287. if expr.name.startswith('Dummy_'):
  288. return '_' + expr.name
  289. else:
  290. return '%s_%d' % (expr.name, expr.dummy_index)
  291. def _print_CodeBlock(self, expr):
  292. return '\n'.join([self._print(i) for i in expr.args])
  293. def _print_String(self, string):
  294. return str(string)
  295. def _print_QuotedString(self, arg):
  296. return '"%s"' % arg.text
  297. def _print_Comment(self, string):
  298. return self._get_comment(str(string))
  299. def _print_Assignment(self, expr):
  300. from sympy.codegen.ast import Assignment
  301. from sympy.functions.elementary.piecewise import Piecewise
  302. from sympy.matrices.expressions.matexpr import MatrixSymbol
  303. from sympy.tensor.indexed import IndexedBase
  304. lhs = expr.lhs
  305. rhs = expr.rhs
  306. # We special case assignments that take multiple lines
  307. if isinstance(expr.rhs, Piecewise):
  308. # Here we modify Piecewise so each expression is now
  309. # an Assignment, and then continue on the print.
  310. expressions = []
  311. conditions = []
  312. for (e, c) in rhs.args:
  313. expressions.append(Assignment(lhs, e))
  314. conditions.append(c)
  315. temp = Piecewise(*zip(expressions, conditions))
  316. return self._print(temp)
  317. elif isinstance(lhs, MatrixSymbol):
  318. # Here we form an Assignment for each element in the array,
  319. # printing each one.
  320. lines = []
  321. for (i, j) in self._traverse_matrix_indices(lhs):
  322. temp = Assignment(lhs[i, j], rhs[i, j])
  323. code0 = self._print(temp)
  324. lines.append(code0)
  325. return "\n".join(lines)
  326. elif self._settings.get("contract", False) and (lhs.has(IndexedBase) or
  327. rhs.has(IndexedBase)):
  328. # Here we check if there is looping to be done, and if so
  329. # print the required loops.
  330. return self._doprint_loops(rhs, lhs)
  331. else:
  332. lhs_code = self._print(lhs)
  333. rhs_code = self._print(rhs)
  334. return self._get_statement("%s = %s" % (lhs_code, rhs_code))
  335. def _print_AugmentedAssignment(self, expr):
  336. lhs_code = self._print(expr.lhs)
  337. rhs_code = self._print(expr.rhs)
  338. return self._get_statement("{} {} {}".format(
  339. *(self._print(arg) for arg in [lhs_code, expr.op, rhs_code])))
  340. def _print_FunctionCall(self, expr):
  341. return '%s(%s)' % (
  342. expr.name,
  343. ', '.join((self._print(arg) for arg in expr.function_args)))
  344. def _print_Variable(self, expr):
  345. return self._print(expr.symbol)
  346. def _print_Symbol(self, expr):
  347. name = super()._print_Symbol(expr)
  348. if name in self.reserved_words:
  349. if self._settings['error_on_reserved']:
  350. msg = ('This expression includes the symbol "{}" which is a '
  351. 'reserved keyword in this language.')
  352. raise ValueError(msg.format(name))
  353. return name + self._settings['reserved_word_suffix']
  354. else:
  355. return name
  356. def _can_print(self, name):
  357. """ Check if function ``name`` is either a known function or has its own
  358. printing method. Used to check if rewriting is possible."""
  359. return name in self.known_functions or getattr(self, '_print_{}'.format(name), False)
  360. def _print_Function(self, expr):
  361. if expr.func.__name__ in self.known_functions:
  362. cond_func = self.known_functions[expr.func.__name__]
  363. if isinstance(cond_func, str):
  364. return "%s(%s)" % (cond_func, self.stringify(expr.args, ", "))
  365. else:
  366. for cond, func in cond_func:
  367. if cond(*expr.args):
  368. break
  369. if func is not None:
  370. try:
  371. return func(*[self.parenthesize(item, 0) for item in expr.args])
  372. except TypeError:
  373. return "%s(%s)" % (func, self.stringify(expr.args, ", "))
  374. elif hasattr(expr, '_imp_') and isinstance(expr._imp_, Lambda):
  375. # inlined function
  376. return self._print(expr._imp_(*expr.args))
  377. elif expr.func.__name__ in self._rewriteable_functions:
  378. # Simple rewrite to supported function possible
  379. target_f, required_fs = self._rewriteable_functions[expr.func.__name__]
  380. if self._can_print(target_f) and all(self._can_print(f) for f in required_fs):
  381. return self._print(expr.rewrite(target_f))
  382. if expr.is_Function and self._settings.get('allow_unknown_functions', False):
  383. return '%s(%s)' % (self._print(expr.func), ', '.join(map(self._print, expr.args)))
  384. else:
  385. return self._print_not_supported(expr)
  386. _print_Expr = _print_Function
  387. # Don't inherit the str-printer method for Heaviside to the code printers
  388. _print_Heaviside = None
  389. def _print_NumberSymbol(self, expr):
  390. if self._settings.get("inline", False):
  391. return self._print(Float(expr.evalf(self._settings["precision"])))
  392. else:
  393. # A Number symbol that is not implemented here or with _printmethod
  394. # is registered and evaluated
  395. self._number_symbols.add((expr,
  396. Float(expr.evalf(self._settings["precision"]))))
  397. return str(expr)
  398. def _print_Catalan(self, expr):
  399. return self._print_NumberSymbol(expr)
  400. def _print_EulerGamma(self, expr):
  401. return self._print_NumberSymbol(expr)
  402. def _print_GoldenRatio(self, expr):
  403. return self._print_NumberSymbol(expr)
  404. def _print_TribonacciConstant(self, expr):
  405. return self._print_NumberSymbol(expr)
  406. def _print_Exp1(self, expr):
  407. return self._print_NumberSymbol(expr)
  408. def _print_Pi(self, expr):
  409. return self._print_NumberSymbol(expr)
  410. def _print_And(self, expr):
  411. PREC = precedence(expr)
  412. return (" %s " % self._operators['and']).join(self.parenthesize(a, PREC)
  413. for a in sorted(expr.args, key=default_sort_key))
  414. def _print_Or(self, expr):
  415. PREC = precedence(expr)
  416. return (" %s " % self._operators['or']).join(self.parenthesize(a, PREC)
  417. for a in sorted(expr.args, key=default_sort_key))
  418. def _print_Xor(self, expr):
  419. if self._operators.get('xor') is None:
  420. return self._print(expr.to_nnf())
  421. PREC = precedence(expr)
  422. return (" %s " % self._operators['xor']).join(self.parenthesize(a, PREC)
  423. for a in expr.args)
  424. def _print_Equivalent(self, expr):
  425. if self._operators.get('equivalent') is None:
  426. return self._print(expr.to_nnf())
  427. PREC = precedence(expr)
  428. return (" %s " % self._operators['equivalent']).join(self.parenthesize(a, PREC)
  429. for a in expr.args)
  430. def _print_Not(self, expr):
  431. PREC = precedence(expr)
  432. return self._operators['not'] + self.parenthesize(expr.args[0], PREC)
  433. def _print_BooleanFunction(self, expr):
  434. return self._print(expr.to_nnf())
  435. def _print_Mul(self, expr):
  436. prec = precedence(expr)
  437. c, e = expr.as_coeff_Mul()
  438. if c < 0:
  439. expr = _keep_coeff(-c, e)
  440. sign = "-"
  441. else:
  442. sign = ""
  443. a = [] # items in the numerator
  444. b = [] # items that are in the denominator (if any)
  445. pow_paren = [] # Will collect all pow with more than one base element and exp = -1
  446. if self.order not in ('old', 'none'):
  447. args = expr.as_ordered_factors()
  448. else:
  449. # use make_args in case expr was something like -x -> x
  450. args = Mul.make_args(expr)
  451. # Gather args for numerator/denominator
  452. for item in args:
  453. if item.is_commutative and item.is_Pow and item.exp.is_Rational and item.exp.is_negative:
  454. if item.exp != -1:
  455. b.append(Pow(item.base, -item.exp, evaluate=False))
  456. else:
  457. if len(item.args[0].args) != 1 and isinstance(item.base, Mul): # To avoid situations like #14160
  458. pow_paren.append(item)
  459. b.append(Pow(item.base, -item.exp))
  460. else:
  461. a.append(item)
  462. a = a or [S.One]
  463. if len(a) == 1 and sign == "-":
  464. # Unary minus does not have a SymPy class, and hence there's no
  465. # precedence weight associated with it, Python's unary minus has
  466. # an operator precedence between multiplication and exponentiation,
  467. # so we use this to compute a weight.
  468. a_str = [self.parenthesize(a[0], 0.5*(PRECEDENCE["Pow"]+PRECEDENCE["Mul"]))]
  469. else:
  470. a_str = [self.parenthesize(x, prec) for x in a]
  471. b_str = [self.parenthesize(x, prec) for x in b]
  472. # To parenthesize Pow with exp = -1 and having more than one Symbol
  473. for item in pow_paren:
  474. if item.base in b:
  475. b_str[b.index(item.base)] = "(%s)" % b_str[b.index(item.base)]
  476. if not b:
  477. return sign + '*'.join(a_str)
  478. elif len(b) == 1:
  479. return sign + '*'.join(a_str) + "/" + b_str[0]
  480. else:
  481. return sign + '*'.join(a_str) + "/(%s)" % '*'.join(b_str)
  482. def _print_not_supported(self, expr):
  483. try:
  484. self._not_supported.add(expr)
  485. except TypeError:
  486. # not hashable
  487. pass
  488. return self.emptyPrinter(expr)
  489. # The following can not be simply translated into C or Fortran
  490. _print_Basic = _print_not_supported
  491. _print_ComplexInfinity = _print_not_supported
  492. _print_Derivative = _print_not_supported
  493. _print_ExprCondPair = _print_not_supported
  494. _print_GeometryEntity = _print_not_supported
  495. _print_Infinity = _print_not_supported
  496. _print_Integral = _print_not_supported
  497. _print_Interval = _print_not_supported
  498. _print_AccumulationBounds = _print_not_supported
  499. _print_Limit = _print_not_supported
  500. _print_MatrixBase = _print_not_supported
  501. _print_DeferredVector = _print_not_supported
  502. _print_NaN = _print_not_supported
  503. _print_NegativeInfinity = _print_not_supported
  504. _print_Order = _print_not_supported
  505. _print_RootOf = _print_not_supported
  506. _print_RootsOf = _print_not_supported
  507. _print_RootSum = _print_not_supported
  508. _print_Uniform = _print_not_supported
  509. _print_Unit = _print_not_supported
  510. _print_Wild = _print_not_supported
  511. _print_WildFunction = _print_not_supported
  512. _print_Relational = _print_not_supported
  513. # Code printer functions. These are included in this file so that they can be
  514. # imported in the top-level __init__.py without importing the sympy.codegen
  515. # module.
  516. def ccode(expr, assign_to=None, standard='c99', **settings):
  517. """Converts an expr to a string of c code
  518. Parameters
  519. ==========
  520. expr : Expr
  521. A SymPy expression to be converted.
  522. assign_to : optional
  523. When given, the argument is used as the name of the variable to which
  524. the expression is assigned. Can be a string, ``Symbol``,
  525. ``MatrixSymbol``, or ``Indexed`` type. This is helpful in case of
  526. line-wrapping, or for expressions that generate multi-line statements.
  527. standard : str, optional
  528. String specifying the standard. If your compiler supports a more modern
  529. standard you may set this to 'c99' to allow the printer to use more math
  530. functions. [default='c89'].
  531. precision : integer, optional
  532. The precision for numbers such as pi [default=17].
  533. user_functions : dict, optional
  534. A dictionary where the keys are string representations of either
  535. ``FunctionClass`` or ``UndefinedFunction`` instances and the values
  536. are their desired C string representations. Alternatively, the
  537. dictionary value can be a list of tuples i.e. [(argument_test,
  538. cfunction_string)] or [(argument_test, cfunction_formater)]. See below
  539. for examples.
  540. dereference : iterable, optional
  541. An iterable of symbols that should be dereferenced in the printed code
  542. expression. These would be values passed by address to the function.
  543. For example, if ``dereference=[a]``, the resulting code would print
  544. ``(*a)`` instead of ``a``.
  545. human : bool, optional
  546. If True, the result is a single string that may contain some constant
  547. declarations for the number symbols. If False, the same information is
  548. returned in a tuple of (symbols_to_declare, not_supported_functions,
  549. code_text). [default=True].
  550. contract: bool, optional
  551. If True, ``Indexed`` instances are assumed to obey tensor contraction
  552. rules and the corresponding nested loops over indices are generated.
  553. Setting contract=False will not generate loops, instead the user is
  554. responsible to provide values for the indices in the code.
  555. [default=True].
  556. Examples
  557. ========
  558. >>> from sympy import ccode, symbols, Rational, sin, ceiling, Abs, Function
  559. >>> x, tau = symbols("x, tau")
  560. >>> expr = (2*tau)**Rational(7, 2)
  561. >>> ccode(expr)
  562. '8*M_SQRT2*pow(tau, 7.0/2.0)'
  563. >>> ccode(expr, math_macros={})
  564. '8*sqrt(2)*pow(tau, 7.0/2.0)'
  565. >>> ccode(sin(x), assign_to="s")
  566. 's = sin(x);'
  567. >>> from sympy.codegen.ast import real, float80
  568. >>> ccode(expr, type_aliases={real: float80})
  569. '8*M_SQRT2l*powl(tau, 7.0L/2.0L)'
  570. Simple custom printing can be defined for certain types by passing a
  571. dictionary of {"type" : "function"} to the ``user_functions`` kwarg.
  572. Alternatively, the dictionary value can be a list of tuples i.e.
  573. [(argument_test, cfunction_string)].
  574. >>> custom_functions = {
  575. ... "ceiling": "CEIL",
  576. ... "Abs": [(lambda x: not x.is_integer, "fabs"),
  577. ... (lambda x: x.is_integer, "ABS")],
  578. ... "func": "f"
  579. ... }
  580. >>> func = Function('func')
  581. >>> ccode(func(Abs(x) + ceiling(x)), standard='C89', user_functions=custom_functions)
  582. 'f(fabs(x) + CEIL(x))'
  583. or if the C-function takes a subset of the original arguments:
  584. >>> ccode(2**x + 3**x, standard='C99', user_functions={'Pow': [
  585. ... (lambda b, e: b == 2, lambda b, e: 'exp2(%s)' % e),
  586. ... (lambda b, e: b != 2, 'pow')]})
  587. 'exp2(x) + pow(3, x)'
  588. ``Piecewise`` expressions are converted into conditionals. If an
  589. ``assign_to`` variable is provided an if statement is created, otherwise
  590. the ternary operator is used. Note that if the ``Piecewise`` lacks a
  591. default term, represented by ``(expr, True)`` then an error will be thrown.
  592. This is to prevent generating an expression that may not evaluate to
  593. anything.
  594. >>> from sympy import Piecewise
  595. >>> expr = Piecewise((x + 1, x > 0), (x, True))
  596. >>> print(ccode(expr, tau, standard='C89'))
  597. if (x > 0) {
  598. tau = x + 1;
  599. }
  600. else {
  601. tau = x;
  602. }
  603. Support for loops is provided through ``Indexed`` types. With
  604. ``contract=True`` these expressions will be turned into loops, whereas
  605. ``contract=False`` will just print the assignment expression that should be
  606. looped over:
  607. >>> from sympy import Eq, IndexedBase, Idx
  608. >>> len_y = 5
  609. >>> y = IndexedBase('y', shape=(len_y,))
  610. >>> t = IndexedBase('t', shape=(len_y,))
  611. >>> Dy = IndexedBase('Dy', shape=(len_y-1,))
  612. >>> i = Idx('i', len_y-1)
  613. >>> e=Eq(Dy[i], (y[i+1]-y[i])/(t[i+1]-t[i]))
  614. >>> ccode(e.rhs, assign_to=e.lhs, contract=False, standard='C89')
  615. 'Dy[i] = (y[i + 1] - y[i])/(t[i + 1] - t[i]);'
  616. Matrices are also supported, but a ``MatrixSymbol`` of the same dimensions
  617. must be provided to ``assign_to``. Note that any expression that can be
  618. generated normally can also exist inside a Matrix:
  619. >>> from sympy import Matrix, MatrixSymbol
  620. >>> mat = Matrix([x**2, Piecewise((x + 1, x > 0), (x, True)), sin(x)])
  621. >>> A = MatrixSymbol('A', 3, 1)
  622. >>> print(ccode(mat, A, standard='C89'))
  623. A[0] = pow(x, 2);
  624. if (x > 0) {
  625. A[1] = x + 1;
  626. }
  627. else {
  628. A[1] = x;
  629. }
  630. A[2] = sin(x);
  631. """
  632. from sympy.printing.c import c_code_printers
  633. return c_code_printers[standard.lower()](settings).doprint(expr, assign_to)
  634. def print_ccode(expr, **settings):
  635. """Prints C representation of the given expression."""
  636. print(ccode(expr, **settings))
  637. def fcode(expr, assign_to=None, **settings):
  638. """Converts an expr to a string of fortran code
  639. Parameters
  640. ==========
  641. expr : Expr
  642. A SymPy expression to be converted.
  643. assign_to : optional
  644. When given, the argument is used as the name of the variable to which
  645. the expression is assigned. Can be a string, ``Symbol``,
  646. ``MatrixSymbol``, or ``Indexed`` type. This is helpful in case of
  647. line-wrapping, or for expressions that generate multi-line statements.
  648. precision : integer, optional
  649. DEPRECATED. Use type_mappings instead. The precision for numbers such
  650. as pi [default=17].
  651. user_functions : dict, optional
  652. A dictionary where keys are ``FunctionClass`` instances and values are
  653. their string representations. Alternatively, the dictionary value can
  654. be a list of tuples i.e. [(argument_test, cfunction_string)]. See below
  655. for examples.
  656. human : bool, optional
  657. If True, the result is a single string that may contain some constant
  658. declarations for the number symbols. If False, the same information is
  659. returned in a tuple of (symbols_to_declare, not_supported_functions,
  660. code_text). [default=True].
  661. contract: bool, optional
  662. If True, ``Indexed`` instances are assumed to obey tensor contraction
  663. rules and the corresponding nested loops over indices are generated.
  664. Setting contract=False will not generate loops, instead the user is
  665. responsible to provide values for the indices in the code.
  666. [default=True].
  667. source_format : optional
  668. The source format can be either 'fixed' or 'free'. [default='fixed']
  669. standard : integer, optional
  670. The Fortran standard to be followed. This is specified as an integer.
  671. Acceptable standards are 66, 77, 90, 95, 2003, and 2008. Default is 77.
  672. Note that currently the only distinction internally is between
  673. standards before 95, and those 95 and after. This may change later as
  674. more features are added.
  675. name_mangling : bool, optional
  676. If True, then the variables that would become identical in
  677. case-insensitive Fortran are mangled by appending different number
  678. of ``_`` at the end. If False, SymPy Will not interfere with naming of
  679. variables. [default=True]
  680. Examples
  681. ========
  682. >>> from sympy import fcode, symbols, Rational, sin, ceiling, floor
  683. >>> x, tau = symbols("x, tau")
  684. >>> fcode((2*tau)**Rational(7, 2))
  685. ' 8*sqrt(2.0d0)*tau**(7.0d0/2.0d0)'
  686. >>> fcode(sin(x), assign_to="s")
  687. ' s = sin(x)'
  688. Custom printing can be defined for certain types by passing a dictionary of
  689. "type" : "function" to the ``user_functions`` kwarg. Alternatively, the
  690. dictionary value can be a list of tuples i.e. [(argument_test,
  691. cfunction_string)].
  692. >>> custom_functions = {
  693. ... "ceiling": "CEIL",
  694. ... "floor": [(lambda x: not x.is_integer, "FLOOR1"),
  695. ... (lambda x: x.is_integer, "FLOOR2")]
  696. ... }
  697. >>> fcode(floor(x) + ceiling(x), user_functions=custom_functions)
  698. ' CEIL(x) + FLOOR1(x)'
  699. ``Piecewise`` expressions are converted into conditionals. If an
  700. ``assign_to`` variable is provided an if statement is created, otherwise
  701. the ternary operator is used. Note that if the ``Piecewise`` lacks a
  702. default term, represented by ``(expr, True)`` then an error will be thrown.
  703. This is to prevent generating an expression that may not evaluate to
  704. anything.
  705. >>> from sympy import Piecewise
  706. >>> expr = Piecewise((x + 1, x > 0), (x, True))
  707. >>> print(fcode(expr, tau))
  708. if (x > 0) then
  709. tau = x + 1
  710. else
  711. tau = x
  712. end if
  713. Support for loops is provided through ``Indexed`` types. With
  714. ``contract=True`` these expressions will be turned into loops, whereas
  715. ``contract=False`` will just print the assignment expression that should be
  716. looped over:
  717. >>> from sympy import Eq, IndexedBase, Idx
  718. >>> len_y = 5
  719. >>> y = IndexedBase('y', shape=(len_y,))
  720. >>> t = IndexedBase('t', shape=(len_y,))
  721. >>> Dy = IndexedBase('Dy', shape=(len_y-1,))
  722. >>> i = Idx('i', len_y-1)
  723. >>> e=Eq(Dy[i], (y[i+1]-y[i])/(t[i+1]-t[i]))
  724. >>> fcode(e.rhs, assign_to=e.lhs, contract=False)
  725. ' Dy(i) = (y(i + 1) - y(i))/(t(i + 1) - t(i))'
  726. Matrices are also supported, but a ``MatrixSymbol`` of the same dimensions
  727. must be provided to ``assign_to``. Note that any expression that can be
  728. generated normally can also exist inside a Matrix:
  729. >>> from sympy import Matrix, MatrixSymbol
  730. >>> mat = Matrix([x**2, Piecewise((x + 1, x > 0), (x, True)), sin(x)])
  731. >>> A = MatrixSymbol('A', 3, 1)
  732. >>> print(fcode(mat, A))
  733. A(1, 1) = x**2
  734. if (x > 0) then
  735. A(2, 1) = x + 1
  736. else
  737. A(2, 1) = x
  738. end if
  739. A(3, 1) = sin(x)
  740. """
  741. from sympy.printing.fortran import FCodePrinter
  742. return FCodePrinter(settings).doprint(expr, assign_to)
  743. def print_fcode(expr, **settings):
  744. """Prints the Fortran representation of the given expression.
  745. See fcode for the meaning of the optional arguments.
  746. """
  747. print(fcode(expr, **settings))
  748. def cxxcode(expr, assign_to=None, standard='c++11', **settings):
  749. """ C++ equivalent of :func:`~.ccode`. """
  750. from sympy.printing.cxx import cxx_code_printers
  751. return cxx_code_printers[standard.lower()](settings).doprint(expr, assign_to)