123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410 |
- """
- R code printer
- The RCodePrinter converts single SymPy expressions into single R expressions,
- using the functions defined in math.h where possible.
- """
- from typing import Any, Dict as tDict
- from sympy.printing.codeprinter import CodePrinter
- from sympy.printing.precedence import precedence, PRECEDENCE
- from sympy.sets.fancysets import Range
- # dictionary mapping SymPy function to (argument_conditions, C_function).
- # Used in RCodePrinter._print_Function(self)
- known_functions = {
- #"Abs": [(lambda x: not x.is_integer, "fabs")],
- "Abs": "abs",
- "sin": "sin",
- "cos": "cos",
- "tan": "tan",
- "asin": "asin",
- "acos": "acos",
- "atan": "atan",
- "atan2": "atan2",
- "exp": "exp",
- "log": "log",
- "erf": "erf",
- "sinh": "sinh",
- "cosh": "cosh",
- "tanh": "tanh",
- "asinh": "asinh",
- "acosh": "acosh",
- "atanh": "atanh",
- "floor": "floor",
- "ceiling": "ceiling",
- "sign": "sign",
- "Max": "max",
- "Min": "min",
- "factorial": "factorial",
- "gamma": "gamma",
- "digamma": "digamma",
- "trigamma": "trigamma",
- "beta": "beta",
- "sqrt": "sqrt", # To enable automatic rewrite
- }
- # These are the core reserved words in the R language. Taken from:
- # https://cran.r-project.org/doc/manuals/r-release/R-lang.html#Reserved-words
- reserved_words = ['if',
- 'else',
- 'repeat',
- 'while',
- 'function',
- 'for',
- 'in',
- 'next',
- 'break',
- 'TRUE',
- 'FALSE',
- 'NULL',
- 'Inf',
- 'NaN',
- 'NA',
- 'NA_integer_',
- 'NA_real_',
- 'NA_complex_',
- 'NA_character_',
- 'volatile']
- class RCodePrinter(CodePrinter):
- """A printer to convert SymPy expressions to strings of R code"""
- printmethod = "_rcode"
- language = "R"
- _default_settings = {
- 'order': None,
- 'full_prec': 'auto',
- 'precision': 15,
- 'user_functions': {},
- 'human': True,
- 'contract': True,
- 'dereference': set(),
- 'error_on_reserved': False,
- 'reserved_word_suffix': '_',
- } # type: tDict[str, Any]
- _operators = {
- 'and': '&',
- 'or': '|',
- 'not': '!',
- }
- _relationals = {
- } # type: tDict[str, str]
- def __init__(self, settings={}):
- CodePrinter.__init__(self, settings)
- self.known_functions = dict(known_functions)
- userfuncs = settings.get('user_functions', {})
- self.known_functions.update(userfuncs)
- self._dereference = set(settings.get('dereference', []))
- self.reserved_words = set(reserved_words)
- def _rate_index_position(self, p):
- return p*5
- def _get_statement(self, codestring):
- return "%s;" % codestring
- def _get_comment(self, text):
- return "// {}".format(text)
- def _declare_number_const(self, name, value):
- return "{} = {};".format(name, value)
- def _format_code(self, lines):
- return self.indent_code(lines)
- def _traverse_matrix_indices(self, mat):
- rows, cols = mat.shape
- return ((i, j) for i in range(rows) for j in range(cols))
- def _get_loop_opening_ending(self, indices):
- """Returns a tuple (open_lines, close_lines) containing lists of codelines
- """
- open_lines = []
- close_lines = []
- loopstart = "for (%(var)s in %(start)s:%(end)s){"
- for i in indices:
- # R arrays start at 1 and end at dimension
- open_lines.append(loopstart % {
- 'var': self._print(i.label),
- 'start': self._print(i.lower+1),
- 'end': self._print(i.upper + 1)})
- close_lines.append("}")
- return open_lines, close_lines
- def _print_Pow(self, expr):
- if "Pow" in self.known_functions:
- return self._print_Function(expr)
- PREC = precedence(expr)
- if expr.exp == -1:
- return '1.0/%s' % (self.parenthesize(expr.base, PREC))
- elif expr.exp == 0.5:
- return 'sqrt(%s)' % self._print(expr.base)
- else:
- return '%s^%s' % (self.parenthesize(expr.base, PREC),
- self.parenthesize(expr.exp, PREC))
- def _print_Rational(self, expr):
- p, q = int(expr.p), int(expr.q)
- return '%d.0/%d.0' % (p, q)
- def _print_Indexed(self, expr):
- inds = [ self._print(i) for i in expr.indices ]
- return "%s[%s]" % (self._print(expr.base.label), ", ".join(inds))
- def _print_Idx(self, expr):
- return self._print(expr.label)
- def _print_Exp1(self, expr):
- return "exp(1)"
- def _print_Pi(self, expr):
- return 'pi'
- def _print_Infinity(self, expr):
- return 'Inf'
- def _print_NegativeInfinity(self, expr):
- return '-Inf'
- def _print_Assignment(self, expr):
- from sympy.codegen.ast import Assignment
- from sympy.matrices.expressions.matexpr import MatrixSymbol
- from sympy.tensor.indexed import IndexedBase
- lhs = expr.lhs
- rhs = expr.rhs
- # We special case assignments that take multiple lines
- #if isinstance(expr.rhs, Piecewise):
- # from sympy.functions.elementary.piecewise import Piecewise
- # # Here we modify Piecewise so each expression is now
- # # an Assignment, and then continue on the print.
- # expressions = []
- # conditions = []
- # for (e, c) in rhs.args:
- # expressions.append(Assignment(lhs, e))
- # conditions.append(c)
- # temp = Piecewise(*zip(expressions, conditions))
- # return self._print(temp)
- #elif isinstance(lhs, MatrixSymbol):
- if isinstance(lhs, MatrixSymbol):
- # Here we form an Assignment for each element in the array,
- # printing each one.
- lines = []
- for (i, j) in self._traverse_matrix_indices(lhs):
- temp = Assignment(lhs[i, j], rhs[i, j])
- code0 = self._print(temp)
- lines.append(code0)
- return "\n".join(lines)
- elif self._settings["contract"] and (lhs.has(IndexedBase) or
- rhs.has(IndexedBase)):
- # Here we check if there is looping to be done, and if so
- # print the required loops.
- return self._doprint_loops(rhs, lhs)
- else:
- lhs_code = self._print(lhs)
- rhs_code = self._print(rhs)
- return self._get_statement("%s = %s" % (lhs_code, rhs_code))
- def _print_Piecewise(self, expr):
- # This method is called only for inline if constructs
- # Top level piecewise is handled in doprint()
- if expr.args[-1].cond == True:
- last_line = "%s" % self._print(expr.args[-1].expr)
- else:
- last_line = "ifelse(%s,%s,NA)" % (self._print(expr.args[-1].cond), self._print(expr.args[-1].expr))
- code=last_line
- for e, c in reversed(expr.args[:-1]):
- code= "ifelse(%s,%s," % (self._print(c), self._print(e))+code+")"
- return(code)
- def _print_ITE(self, expr):
- from sympy.functions import Piecewise
- return self._print(expr.rewrite(Piecewise))
- def _print_MatrixElement(self, expr):
- return "{}[{}]".format(self.parenthesize(expr.parent, PRECEDENCE["Atom"],
- strict=True), expr.j + expr.i*expr.parent.shape[1])
- def _print_Symbol(self, expr):
- name = super()._print_Symbol(expr)
- if expr in self._dereference:
- return '(*{})'.format(name)
- else:
- return name
- def _print_Relational(self, expr):
- lhs_code = self._print(expr.lhs)
- rhs_code = self._print(expr.rhs)
- op = expr.rel_op
- return "{} {} {}".format(lhs_code, op, rhs_code)
- def _print_AugmentedAssignment(self, expr):
- lhs_code = self._print(expr.lhs)
- op = expr.op
- rhs_code = self._print(expr.rhs)
- return "{} {} {};".format(lhs_code, op, rhs_code)
- def _print_For(self, expr):
- target = self._print(expr.target)
- if isinstance(expr.iterable, Range):
- start, stop, step = expr.iterable.args
- else:
- raise NotImplementedError("Only iterable currently supported is Range")
- body = self._print(expr.body)
- return ('for ({target} = {start}; {target} < {stop}; {target} += '
- '{step}) {{\n{body}\n}}').format(target=target, start=start,
- stop=stop, step=step, body=body)
- def indent_code(self, code):
- """Accepts a string of code or a list of code lines"""
- if isinstance(code, str):
- code_lines = self.indent_code(code.splitlines(True))
- return ''.join(code_lines)
- tab = " "
- inc_token = ('{', '(', '{\n', '(\n')
- dec_token = ('}', ')')
- code = [ line.lstrip(' \t') for line in code ]
- increase = [ int(any(map(line.endswith, inc_token))) for line in code ]
- decrease = [ int(any(map(line.startswith, dec_token)))
- for line in code ]
- pretty = []
- level = 0
- for n, line in enumerate(code):
- if line in ('', '\n'):
- pretty.append(line)
- continue
- level -= decrease[n]
- pretty.append("%s%s" % (tab*level, line))
- level += increase[n]
- return pretty
- def rcode(expr, assign_to=None, **settings):
- """Converts an expr to a string of r code
- Parameters
- ==========
- expr : Expr
- A SymPy expression to be converted.
- assign_to : optional
- When given, the argument is used as the name of the variable to which
- the expression is assigned. Can be a string, ``Symbol``,
- ``MatrixSymbol``, or ``Indexed`` type. This is helpful in case of
- line-wrapping, or for expressions that generate multi-line statements.
- precision : integer, optional
- The precision for numbers such as pi [default=15].
- user_functions : dict, optional
- A dictionary where the keys are string representations of either
- ``FunctionClass`` or ``UndefinedFunction`` instances and the values
- are their desired R string representations. Alternatively, the
- dictionary value can be a list of tuples i.e. [(argument_test,
- rfunction_string)] or [(argument_test, rfunction_formater)]. See below
- for examples.
- human : bool, optional
- If True, the result is a single string that may contain some constant
- declarations for the number symbols. If False, the same information is
- returned in a tuple of (symbols_to_declare, not_supported_functions,
- code_text). [default=True].
- contract: bool, optional
- If True, ``Indexed`` instances are assumed to obey tensor contraction
- rules and the corresponding nested loops over indices are generated.
- Setting contract=False will not generate loops, instead the user is
- responsible to provide values for the indices in the code.
- [default=True].
- Examples
- ========
- >>> from sympy import rcode, symbols, Rational, sin, ceiling, Abs, Function
- >>> x, tau = symbols("x, tau")
- >>> rcode((2*tau)**Rational(7, 2))
- '8*sqrt(2)*tau^(7.0/2.0)'
- >>> rcode(sin(x), assign_to="s")
- 's = sin(x);'
- Simple custom printing can be defined for certain types by passing a
- dictionary of {"type" : "function"} to the ``user_functions`` kwarg.
- Alternatively, the dictionary value can be a list of tuples i.e.
- [(argument_test, cfunction_string)].
- >>> custom_functions = {
- ... "ceiling": "CEIL",
- ... "Abs": [(lambda x: not x.is_integer, "fabs"),
- ... (lambda x: x.is_integer, "ABS")],
- ... "func": "f"
- ... }
- >>> func = Function('func')
- >>> rcode(func(Abs(x) + ceiling(x)), user_functions=custom_functions)
- 'f(fabs(x) + CEIL(x))'
- or if the R-function takes a subset of the original arguments:
- >>> rcode(2**x + 3**x, user_functions={'Pow': [
- ... (lambda b, e: b == 2, lambda b, e: 'exp2(%s)' % e),
- ... (lambda b, e: b != 2, 'pow')]})
- 'exp2(x) + pow(3, x)'
- ``Piecewise`` expressions are converted into conditionals. If an
- ``assign_to`` variable is provided an if statement is created, otherwise
- the ternary operator is used. Note that if the ``Piecewise`` lacks a
- default term, represented by ``(expr, True)`` then an error will be thrown.
- This is to prevent generating an expression that may not evaluate to
- anything.
- >>> from sympy import Piecewise
- >>> expr = Piecewise((x + 1, x > 0), (x, True))
- >>> print(rcode(expr, assign_to=tau))
- tau = ifelse(x > 0,x + 1,x);
- Support for loops is provided through ``Indexed`` types. With
- ``contract=True`` these expressions will be turned into loops, whereas
- ``contract=False`` will just print the assignment expression that should be
- looped over:
- >>> from sympy import Eq, IndexedBase, Idx
- >>> len_y = 5
- >>> y = IndexedBase('y', shape=(len_y,))
- >>> t = IndexedBase('t', shape=(len_y,))
- >>> Dy = IndexedBase('Dy', shape=(len_y-1,))
- >>> i = Idx('i', len_y-1)
- >>> e=Eq(Dy[i], (y[i+1]-y[i])/(t[i+1]-t[i]))
- >>> rcode(e.rhs, assign_to=e.lhs, contract=False)
- 'Dy[i] = (y[i + 1] - y[i])/(t[i + 1] - t[i]);'
- Matrices are also supported, but a ``MatrixSymbol`` of the same dimensions
- must be provided to ``assign_to``. Note that any expression that can be
- generated normally can also exist inside a Matrix:
- >>> from sympy import Matrix, MatrixSymbol
- >>> mat = Matrix([x**2, Piecewise((x + 1, x > 0), (x, True)), sin(x)])
- >>> A = MatrixSymbol('A', 3, 1)
- >>> print(rcode(mat, A))
- A[0] = x^2;
- A[1] = ifelse(x > 0,x + 1,x);
- A[2] = sin(x);
- """
- return RCodePrinter(settings).doprint(expr, assign_to)
- def print_rcode(expr, **settings):
- """Prints R representation of the given expression."""
- print(rcode(expr, **settings))
|