Coverage for /usr/lib/python3/dist-packages/sympy/printing/rcode.py: 28%
144 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-14 15:55 +0200
1"""
2R code printer
4The RCodePrinter converts single SymPy expressions into single R expressions,
5using the functions defined in math.h where possible.
9"""
11from __future__ import annotations
12from typing import Any
14from sympy.core.numbers import equal_valued
15from sympy.printing.codeprinter import CodePrinter
16from sympy.printing.precedence import precedence, PRECEDENCE
17from sympy.sets.fancysets import Range
19# dictionary mapping SymPy function to (argument_conditions, C_function).
20# Used in RCodePrinter._print_Function(self)
21known_functions = {
22 #"Abs": [(lambda x: not x.is_integer, "fabs")],
23 "Abs": "abs",
24 "sin": "sin",
25 "cos": "cos",
26 "tan": "tan",
27 "asin": "asin",
28 "acos": "acos",
29 "atan": "atan",
30 "atan2": "atan2",
31 "exp": "exp",
32 "log": "log",
33 "erf": "erf",
34 "sinh": "sinh",
35 "cosh": "cosh",
36 "tanh": "tanh",
37 "asinh": "asinh",
38 "acosh": "acosh",
39 "atanh": "atanh",
40 "floor": "floor",
41 "ceiling": "ceiling",
42 "sign": "sign",
43 "Max": "max",
44 "Min": "min",
45 "factorial": "factorial",
46 "gamma": "gamma",
47 "digamma": "digamma",
48 "trigamma": "trigamma",
49 "beta": "beta",
50 "sqrt": "sqrt", # To enable automatic rewrite
51}
53# These are the core reserved words in the R language. Taken from:
54# https://cran.r-project.org/doc/manuals/r-release/R-lang.html#Reserved-words
56reserved_words = ['if',
57 'else',
58 'repeat',
59 'while',
60 'function',
61 'for',
62 'in',
63 'next',
64 'break',
65 'TRUE',
66 'FALSE',
67 'NULL',
68 'Inf',
69 'NaN',
70 'NA',
71 'NA_integer_',
72 'NA_real_',
73 'NA_complex_',
74 'NA_character_',
75 'volatile']
78class RCodePrinter(CodePrinter):
79 """A printer to convert SymPy expressions to strings of R code"""
80 printmethod = "_rcode"
81 language = "R"
83 _default_settings: dict[str, Any] = {
84 'order': None,
85 'full_prec': 'auto',
86 'precision': 15,
87 'user_functions': {},
88 'human': True,
89 'contract': True,
90 'dereference': set(),
91 'error_on_reserved': False,
92 'reserved_word_suffix': '_',
93 }
94 _operators = {
95 'and': '&',
96 'or': '|',
97 'not': '!',
98 }
100 _relationals: dict[str, str] = {}
102 def __init__(self, settings={}):
103 CodePrinter.__init__(self, settings)
104 self.known_functions = dict(known_functions)
105 userfuncs = settings.get('user_functions', {})
106 self.known_functions.update(userfuncs)
107 self._dereference = set(settings.get('dereference', []))
108 self.reserved_words = set(reserved_words)
110 def _rate_index_position(self, p):
111 return p*5
113 def _get_statement(self, codestring):
114 return "%s;" % codestring
116 def _get_comment(self, text):
117 return "// {}".format(text)
119 def _declare_number_const(self, name, value):
120 return "{} = {};".format(name, value)
122 def _format_code(self, lines):
123 return self.indent_code(lines)
125 def _traverse_matrix_indices(self, mat):
126 rows, cols = mat.shape
127 return ((i, j) for i in range(rows) for j in range(cols))
129 def _get_loop_opening_ending(self, indices):
130 """Returns a tuple (open_lines, close_lines) containing lists of codelines
131 """
132 open_lines = []
133 close_lines = []
134 loopstart = "for (%(var)s in %(start)s:%(end)s){"
135 for i in indices:
136 # R arrays start at 1 and end at dimension
137 open_lines.append(loopstart % {
138 'var': self._print(i.label),
139 'start': self._print(i.lower+1),
140 'end': self._print(i.upper + 1)})
141 close_lines.append("}")
142 return open_lines, close_lines
144 def _print_Pow(self, expr):
145 if "Pow" in self.known_functions:
146 return self._print_Function(expr)
147 PREC = precedence(expr)
148 if equal_valued(expr.exp, -1):
149 return '1.0/%s' % (self.parenthesize(expr.base, PREC))
150 elif equal_valued(expr.exp, 0.5):
151 return 'sqrt(%s)' % self._print(expr.base)
152 else:
153 return '%s^%s' % (self.parenthesize(expr.base, PREC),
154 self.parenthesize(expr.exp, PREC))
157 def _print_Rational(self, expr):
158 p, q = int(expr.p), int(expr.q)
159 return '%d.0/%d.0' % (p, q)
161 def _print_Indexed(self, expr):
162 inds = [ self._print(i) for i in expr.indices ]
163 return "%s[%s]" % (self._print(expr.base.label), ", ".join(inds))
165 def _print_Idx(self, expr):
166 return self._print(expr.label)
168 def _print_Exp1(self, expr):
169 return "exp(1)"
171 def _print_Pi(self, expr):
172 return 'pi'
174 def _print_Infinity(self, expr):
175 return 'Inf'
177 def _print_NegativeInfinity(self, expr):
178 return '-Inf'
180 def _print_Assignment(self, expr):
181 from sympy.codegen.ast import Assignment
183 from sympy.matrices.expressions.matexpr import MatrixSymbol
184 from sympy.tensor.indexed import IndexedBase
185 lhs = expr.lhs
186 rhs = expr.rhs
187 # We special case assignments that take multiple lines
188 #if isinstance(expr.rhs, Piecewise):
189 # from sympy.functions.elementary.piecewise import Piecewise
190 # # Here we modify Piecewise so each expression is now
191 # # an Assignment, and then continue on the print.
192 # expressions = []
193 # conditions = []
194 # for (e, c) in rhs.args:
195 # expressions.append(Assignment(lhs, e))
196 # conditions.append(c)
197 # temp = Piecewise(*zip(expressions, conditions))
198 # return self._print(temp)
199 #elif isinstance(lhs, MatrixSymbol):
200 if isinstance(lhs, MatrixSymbol):
201 # Here we form an Assignment for each element in the array,
202 # printing each one.
203 lines = []
204 for (i, j) in self._traverse_matrix_indices(lhs):
205 temp = Assignment(lhs[i, j], rhs[i, j])
206 code0 = self._print(temp)
207 lines.append(code0)
208 return "\n".join(lines)
209 elif self._settings["contract"] and (lhs.has(IndexedBase) or
210 rhs.has(IndexedBase)):
211 # Here we check if there is looping to be done, and if so
212 # print the required loops.
213 return self._doprint_loops(rhs, lhs)
214 else:
215 lhs_code = self._print(lhs)
216 rhs_code = self._print(rhs)
217 return self._get_statement("%s = %s" % (lhs_code, rhs_code))
219 def _print_Piecewise(self, expr):
220 # This method is called only for inline if constructs
221 # Top level piecewise is handled in doprint()
222 if expr.args[-1].cond == True:
223 last_line = "%s" % self._print(expr.args[-1].expr)
224 else:
225 last_line = "ifelse(%s,%s,NA)" % (self._print(expr.args[-1].cond), self._print(expr.args[-1].expr))
226 code=last_line
227 for e, c in reversed(expr.args[:-1]):
228 code= "ifelse(%s,%s," % (self._print(c), self._print(e))+code+")"
229 return(code)
231 def _print_ITE(self, expr):
232 from sympy.functions import Piecewise
233 return self._print(expr.rewrite(Piecewise))
235 def _print_MatrixElement(self, expr):
236 return "{}[{}]".format(self.parenthesize(expr.parent, PRECEDENCE["Atom"],
237 strict=True), expr.j + expr.i*expr.parent.shape[1])
239 def _print_Symbol(self, expr):
240 name = super()._print_Symbol(expr)
241 if expr in self._dereference:
242 return '(*{})'.format(name)
243 else:
244 return name
246 def _print_Relational(self, expr):
247 lhs_code = self._print(expr.lhs)
248 rhs_code = self._print(expr.rhs)
249 op = expr.rel_op
250 return "{} {} {}".format(lhs_code, op, rhs_code)
252 def _print_AugmentedAssignment(self, expr):
253 lhs_code = self._print(expr.lhs)
254 op = expr.op
255 rhs_code = self._print(expr.rhs)
256 return "{} {} {};".format(lhs_code, op, rhs_code)
258 def _print_For(self, expr):
259 target = self._print(expr.target)
260 if isinstance(expr.iterable, Range):
261 start, stop, step = expr.iterable.args
262 else:
263 raise NotImplementedError("Only iterable currently supported is Range")
264 body = self._print(expr.body)
265 return 'for({target} in seq(from={start}, to={stop}, by={step}){{\n{body}\n}}'.format(target=target, start=start,
266 stop=stop-1, step=step, body=body)
269 def indent_code(self, code):
270 """Accepts a string of code or a list of code lines"""
272 if isinstance(code, str):
273 code_lines = self.indent_code(code.splitlines(True))
274 return ''.join(code_lines)
276 tab = " "
277 inc_token = ('{', '(', '{\n', '(\n')
278 dec_token = ('}', ')')
280 code = [ line.lstrip(' \t') for line in code ]
282 increase = [ int(any(map(line.endswith, inc_token))) for line in code ]
283 decrease = [ int(any(map(line.startswith, dec_token)))
284 for line in code ]
286 pretty = []
287 level = 0
288 for n, line in enumerate(code):
289 if line in ('', '\n'):
290 pretty.append(line)
291 continue
292 level -= decrease[n]
293 pretty.append("%s%s" % (tab*level, line))
294 level += increase[n]
295 return pretty
298def rcode(expr, assign_to=None, **settings):
299 """Converts an expr to a string of r code
301 Parameters
302 ==========
304 expr : Expr
305 A SymPy expression to be converted.
306 assign_to : optional
307 When given, the argument is used as the name of the variable to which
308 the expression is assigned. Can be a string, ``Symbol``,
309 ``MatrixSymbol``, or ``Indexed`` type. This is helpful in case of
310 line-wrapping, or for expressions that generate multi-line statements.
311 precision : integer, optional
312 The precision for numbers such as pi [default=15].
313 user_functions : dict, optional
314 A dictionary where the keys are string representations of either
315 ``FunctionClass`` or ``UndefinedFunction`` instances and the values
316 are their desired R string representations. Alternatively, the
317 dictionary value can be a list of tuples i.e. [(argument_test,
318 rfunction_string)] or [(argument_test, rfunction_formater)]. See below
319 for examples.
320 human : bool, optional
321 If True, the result is a single string that may contain some constant
322 declarations for the number symbols. If False, the same information is
323 returned in a tuple of (symbols_to_declare, not_supported_functions,
324 code_text). [default=True].
325 contract: bool, optional
326 If True, ``Indexed`` instances are assumed to obey tensor contraction
327 rules and the corresponding nested loops over indices are generated.
328 Setting contract=False will not generate loops, instead the user is
329 responsible to provide values for the indices in the code.
330 [default=True].
332 Examples
333 ========
335 >>> from sympy import rcode, symbols, Rational, sin, ceiling, Abs, Function
336 >>> x, tau = symbols("x, tau")
337 >>> rcode((2*tau)**Rational(7, 2))
338 '8*sqrt(2)*tau^(7.0/2.0)'
339 >>> rcode(sin(x), assign_to="s")
340 's = sin(x);'
342 Simple custom printing can be defined for certain types by passing a
343 dictionary of {"type" : "function"} to the ``user_functions`` kwarg.
344 Alternatively, the dictionary value can be a list of tuples i.e.
345 [(argument_test, cfunction_string)].
347 >>> custom_functions = {
348 ... "ceiling": "CEIL",
349 ... "Abs": [(lambda x: not x.is_integer, "fabs"),
350 ... (lambda x: x.is_integer, "ABS")],
351 ... "func": "f"
352 ... }
353 >>> func = Function('func')
354 >>> rcode(func(Abs(x) + ceiling(x)), user_functions=custom_functions)
355 'f(fabs(x) + CEIL(x))'
357 or if the R-function takes a subset of the original arguments:
359 >>> rcode(2**x + 3**x, user_functions={'Pow': [
360 ... (lambda b, e: b == 2, lambda b, e: 'exp2(%s)' % e),
361 ... (lambda b, e: b != 2, 'pow')]})
362 'exp2(x) + pow(3, x)'
364 ``Piecewise`` expressions are converted into conditionals. If an
365 ``assign_to`` variable is provided an if statement is created, otherwise
366 the ternary operator is used. Note that if the ``Piecewise`` lacks a
367 default term, represented by ``(expr, True)`` then an error will be thrown.
368 This is to prevent generating an expression that may not evaluate to
369 anything.
371 >>> from sympy import Piecewise
372 >>> expr = Piecewise((x + 1, x > 0), (x, True))
373 >>> print(rcode(expr, assign_to=tau))
374 tau = ifelse(x > 0,x + 1,x);
376 Support for loops is provided through ``Indexed`` types. With
377 ``contract=True`` these expressions will be turned into loops, whereas
378 ``contract=False`` will just print the assignment expression that should be
379 looped over:
381 >>> from sympy import Eq, IndexedBase, Idx
382 >>> len_y = 5
383 >>> y = IndexedBase('y', shape=(len_y,))
384 >>> t = IndexedBase('t', shape=(len_y,))
385 >>> Dy = IndexedBase('Dy', shape=(len_y-1,))
386 >>> i = Idx('i', len_y-1)
387 >>> e=Eq(Dy[i], (y[i+1]-y[i])/(t[i+1]-t[i]))
388 >>> rcode(e.rhs, assign_to=e.lhs, contract=False)
389 'Dy[i] = (y[i + 1] - y[i])/(t[i + 1] - t[i]);'
391 Matrices are also supported, but a ``MatrixSymbol`` of the same dimensions
392 must be provided to ``assign_to``. Note that any expression that can be
393 generated normally can also exist inside a Matrix:
395 >>> from sympy import Matrix, MatrixSymbol
396 >>> mat = Matrix([x**2, Piecewise((x + 1, x > 0), (x, True)), sin(x)])
397 >>> A = MatrixSymbol('A', 3, 1)
398 >>> print(rcode(mat, A))
399 A[0] = x^2;
400 A[1] = ifelse(x > 0,x + 1,x);
401 A[2] = sin(x);
403 """
405 return RCodePrinter(settings).doprint(expr, assign_to)
408def print_rcode(expr, **settings):
409 """Prints R representation of the given expression."""
410 print(rcode(expr, **settings))