glsl.py 20 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556
  1. from typing import Set as tSet
  2. from sympy.core import Basic, S
  3. from sympy.core.function import Lambda
  4. from sympy.printing.codeprinter import CodePrinter
  5. from sympy.printing.precedence import precedence
  6. from functools import reduce
  7. known_functions = {
  8. 'Abs': 'abs',
  9. 'sin': 'sin',
  10. 'cos': 'cos',
  11. 'tan': 'tan',
  12. 'acos': 'acos',
  13. 'asin': 'asin',
  14. 'atan': 'atan',
  15. 'atan2': 'atan',
  16. 'ceiling': 'ceil',
  17. 'floor': 'floor',
  18. 'sign': 'sign',
  19. 'exp': 'exp',
  20. 'log': 'log',
  21. 'add': 'add',
  22. 'sub': 'sub',
  23. 'mul': 'mul',
  24. 'pow': 'pow'
  25. }
  26. class GLSLPrinter(CodePrinter):
  27. """
  28. Rudimentary, generic GLSL printing tools.
  29. Additional settings:
  30. 'use_operators': Boolean (should the printer use operators for +,-,*, or functions?)
  31. """
  32. _not_supported = set() # type: tSet[Basic]
  33. printmethod = "_glsl"
  34. language = "GLSL"
  35. _default_settings = {
  36. 'use_operators': True,
  37. 'zero': 0,
  38. 'mat_nested': False,
  39. 'mat_separator': ',\n',
  40. 'mat_transpose': False,
  41. 'array_type': 'float',
  42. 'glsl_types': True,
  43. 'order': None,
  44. 'full_prec': 'auto',
  45. 'precision': 9,
  46. 'user_functions': {},
  47. 'human': True,
  48. 'allow_unknown_functions': False,
  49. 'contract': True,
  50. 'error_on_reserved': False,
  51. 'reserved_word_suffix': '_',
  52. }
  53. def __init__(self, settings={}):
  54. CodePrinter.__init__(self, settings)
  55. self.known_functions = dict(known_functions)
  56. userfuncs = settings.get('user_functions', {})
  57. self.known_functions.update(userfuncs)
  58. def _rate_index_position(self, p):
  59. return p*5
  60. def _get_statement(self, codestring):
  61. return "%s;" % codestring
  62. def _get_comment(self, text):
  63. return "// {}".format(text)
  64. def _declare_number_const(self, name, value):
  65. return "float {} = {};".format(name, value)
  66. def _format_code(self, lines):
  67. return self.indent_code(lines)
  68. def indent_code(self, code):
  69. """Accepts a string of code or a list of code lines"""
  70. if isinstance(code, str):
  71. code_lines = self.indent_code(code.splitlines(True))
  72. return ''.join(code_lines)
  73. tab = " "
  74. inc_token = ('{', '(', '{\n', '(\n')
  75. dec_token = ('}', ')')
  76. code = [line.lstrip(' \t') for line in code]
  77. increase = [int(any(map(line.endswith, inc_token))) for line in code]
  78. decrease = [int(any(map(line.startswith, dec_token))) for line in code]
  79. pretty = []
  80. level = 0
  81. for n, line in enumerate(code):
  82. if line in ('', '\n'):
  83. pretty.append(line)
  84. continue
  85. level -= decrease[n]
  86. pretty.append("%s%s" % (tab*level, line))
  87. level += increase[n]
  88. return pretty
  89. def _print_MatrixBase(self, mat):
  90. mat_separator = self._settings['mat_separator']
  91. mat_transpose = self._settings['mat_transpose']
  92. column_vector = (mat.rows == 1) if mat_transpose else (mat.cols == 1)
  93. A = mat.transpose() if mat_transpose != column_vector else mat
  94. glsl_types = self._settings['glsl_types']
  95. array_type = self._settings['array_type']
  96. array_size = A.cols*A.rows
  97. array_constructor = "{}[{}]".format(array_type, array_size)
  98. if A.cols == 1:
  99. return self._print(A[0]);
  100. if A.rows <= 4 and A.cols <= 4 and glsl_types:
  101. if A.rows == 1:
  102. return "vec{}{}".format(
  103. A.cols, A.table(self,rowstart='(',rowend=')')
  104. )
  105. elif A.rows == A.cols:
  106. return "mat{}({})".format(
  107. A.rows, A.table(self,rowsep=', ',
  108. rowstart='',rowend='')
  109. )
  110. else:
  111. return "mat{}x{}({})".format(
  112. A.cols, A.rows,
  113. A.table(self,rowsep=', ',
  114. rowstart='',rowend='')
  115. )
  116. elif S.One in A.shape:
  117. return "{}({})".format(
  118. array_constructor,
  119. A.table(self,rowsep=mat_separator,rowstart='',rowend='')
  120. )
  121. elif not self._settings['mat_nested']:
  122. return "{}(\n{}\n) /* a {}x{} matrix */".format(
  123. array_constructor,
  124. A.table(self,rowsep=mat_separator,rowstart='',rowend=''),
  125. A.rows, A.cols
  126. )
  127. elif self._settings['mat_nested']:
  128. return "{}[{}][{}](\n{}\n)".format(
  129. array_type, A.rows, A.cols,
  130. A.table(self,rowsep=mat_separator,rowstart='float[](',rowend=')')
  131. )
  132. def _print_SparseRepMatrix(self, mat):
  133. # do not allow sparse matrices to be made dense
  134. return self._print_not_supported(mat)
  135. def _traverse_matrix_indices(self, mat):
  136. mat_transpose = self._settings['mat_transpose']
  137. if mat_transpose:
  138. rows,cols = mat.shape
  139. else:
  140. cols,rows = mat.shape
  141. return ((i, j) for i in range(cols) for j in range(rows))
  142. def _print_MatrixElement(self, expr):
  143. # print('begin _print_MatrixElement')
  144. nest = self._settings['mat_nested'];
  145. glsl_types = self._settings['glsl_types'];
  146. mat_transpose = self._settings['mat_transpose'];
  147. if mat_transpose:
  148. cols,rows = expr.parent.shape
  149. i,j = expr.j,expr.i
  150. else:
  151. rows,cols = expr.parent.shape
  152. i,j = expr.i,expr.j
  153. pnt = self._print(expr.parent)
  154. if glsl_types and ((rows <= 4 and cols <=4) or nest):
  155. return "{}[{}][{}]".format(pnt, i, j)
  156. else:
  157. return "{}[{}]".format(pnt, i + j*rows)
  158. def _print_list(self, expr):
  159. l = ', '.join(self._print(item) for item in expr)
  160. glsl_types = self._settings['glsl_types']
  161. array_type = self._settings['array_type']
  162. array_size = len(expr)
  163. array_constructor = '{}[{}]'.format(array_type, array_size)
  164. if array_size <= 4 and glsl_types:
  165. return 'vec{}({})'.format(array_size, l)
  166. else:
  167. return '{}({})'.format(array_constructor, l)
  168. _print_tuple = _print_list
  169. _print_Tuple = _print_list
  170. def _get_loop_opening_ending(self, indices):
  171. open_lines = []
  172. close_lines = []
  173. loopstart = "for (int %(varble)s=%(start)s; %(varble)s<%(end)s; %(varble)s++){"
  174. for i in indices:
  175. # GLSL arrays start at 0 and end at dimension-1
  176. open_lines.append(loopstart % {
  177. 'varble': self._print(i.label),
  178. 'start': self._print(i.lower),
  179. 'end': self._print(i.upper + 1)})
  180. close_lines.append("}")
  181. return open_lines, close_lines
  182. def _print_Function_with_args(self, func, func_args):
  183. if func in self.known_functions:
  184. cond_func = self.known_functions[func]
  185. func = None
  186. if isinstance(cond_func, str):
  187. func = cond_func
  188. else:
  189. for cond, func in cond_func:
  190. if cond(func_args):
  191. break
  192. if func is not None:
  193. try:
  194. return func(*[self.parenthesize(item, 0) for item in func_args])
  195. except TypeError:
  196. return '{}({})'.format(func, self.stringify(func_args, ", "))
  197. elif isinstance(func, Lambda):
  198. # inlined function
  199. return self._print(func(*func_args))
  200. else:
  201. return self._print_not_supported(func)
  202. def _print_Piecewise(self, expr):
  203. from sympy.codegen.ast import Assignment
  204. if expr.args[-1].cond != True:
  205. # We need the last conditional to be a True, otherwise the resulting
  206. # function may not return a result.
  207. raise ValueError("All Piecewise expressions must contain an "
  208. "(expr, True) statement to be used as a default "
  209. "condition. Without one, the generated "
  210. "expression may not evaluate to anything under "
  211. "some condition.")
  212. lines = []
  213. if expr.has(Assignment):
  214. for i, (e, c) in enumerate(expr.args):
  215. if i == 0:
  216. lines.append("if (%s) {" % self._print(c))
  217. elif i == len(expr.args) - 1 and c == True:
  218. lines.append("else {")
  219. else:
  220. lines.append("else if (%s) {" % self._print(c))
  221. code0 = self._print(e)
  222. lines.append(code0)
  223. lines.append("}")
  224. return "\n".join(lines)
  225. else:
  226. # The piecewise was used in an expression, need to do inline
  227. # operators. This has the downside that inline operators will
  228. # not work for statements that span multiple lines (Matrix or
  229. # Indexed expressions).
  230. ecpairs = ["((%s) ? (\n%s\n)\n" % (self._print(c),
  231. self._print(e))
  232. for e, c in expr.args[:-1]]
  233. last_line = ": (\n%s\n)" % self._print(expr.args[-1].expr)
  234. return ": ".join(ecpairs) + last_line + " ".join([")"*len(ecpairs)])
  235. def _print_Idx(self, expr):
  236. return self._print(expr.label)
  237. def _print_Indexed(self, expr):
  238. # calculate index for 1d array
  239. dims = expr.shape
  240. elem = S.Zero
  241. offset = S.One
  242. for i in reversed(range(expr.rank)):
  243. elem += expr.indices[i]*offset
  244. offset *= dims[i]
  245. return "{}[{}]".format(
  246. self._print(expr.base.label),
  247. self._print(elem)
  248. )
  249. def _print_Pow(self, expr):
  250. PREC = precedence(expr)
  251. if expr.exp == -1:
  252. return '1.0/%s' % (self.parenthesize(expr.base, PREC))
  253. elif expr.exp == 0.5:
  254. return 'sqrt(%s)' % self._print(expr.base)
  255. else:
  256. try:
  257. e = self._print(float(expr.exp))
  258. except TypeError:
  259. e = self._print(expr.exp)
  260. return self._print_Function_with_args('pow', (
  261. self._print(expr.base),
  262. e
  263. ))
  264. def _print_int(self, expr):
  265. return str(float(expr))
  266. def _print_Rational(self, expr):
  267. return "{}.0/{}.0".format(expr.p, expr.q)
  268. def _print_Relational(self, expr):
  269. lhs_code = self._print(expr.lhs)
  270. rhs_code = self._print(expr.rhs)
  271. op = expr.rel_op
  272. return "{} {} {}".format(lhs_code, op, rhs_code)
  273. def _print_Add(self, expr, order=None):
  274. if self._settings['use_operators']:
  275. return CodePrinter._print_Add(self, expr, order=order)
  276. terms = expr.as_ordered_terms()
  277. def partition(p,l):
  278. return reduce(lambda x, y: (x[0]+[y], x[1]) if p(y) else (x[0], x[1]+[y]), l, ([], []))
  279. def add(a,b):
  280. return self._print_Function_with_args('add', (a, b))
  281. # return self.known_functions['add']+'(%s, %s)' % (a,b)
  282. neg, pos = partition(lambda arg: arg.could_extract_minus_sign(), terms)
  283. if pos:
  284. s = pos = reduce(lambda a,b: add(a,b), map(lambda t: self._print(t),pos))
  285. else:
  286. s = pos = self._print(self._settings['zero'])
  287. if neg:
  288. # sum the absolute values of the negative terms
  289. neg = reduce(lambda a,b: add(a,b), map(lambda n: self._print(-n),neg))
  290. # then subtract them from the positive terms
  291. s = self._print_Function_with_args('sub', (pos,neg))
  292. # s = self.known_functions['sub']+'(%s, %s)' % (pos,neg)
  293. return s
  294. def _print_Mul(self, expr, **kwargs):
  295. if self._settings['use_operators']:
  296. return CodePrinter._print_Mul(self, expr, **kwargs)
  297. terms = expr.as_ordered_factors()
  298. def mul(a,b):
  299. # return self.known_functions['mul']+'(%s, %s)' % (a,b)
  300. return self._print_Function_with_args('mul', (a,b))
  301. s = reduce(lambda a,b: mul(a,b), map(lambda t: self._print(t), terms))
  302. return s
  303. def glsl_code(expr,assign_to=None,**settings):
  304. """Converts an expr to a string of GLSL code
  305. Parameters
  306. ==========
  307. expr : Expr
  308. A SymPy expression to be converted.
  309. assign_to : optional
  310. When given, the argument is used for naming the variable or variables
  311. to which the expression is assigned. Can be a string, ``Symbol``,
  312. ``MatrixSymbol`` or ``Indexed`` type object. In cases where ``expr``
  313. would be printed as an array, a list of string or ``Symbol`` objects
  314. can also be passed.
  315. This is helpful in case of line-wrapping, or for expressions that
  316. generate multi-line statements. It can also be used to spread an array-like
  317. expression into multiple assignments.
  318. use_operators: bool, optional
  319. If set to False, then *,/,+,- operators will be replaced with functions
  320. mul, add, and sub, which must be implemented by the user, e.g. for
  321. implementing non-standard rings or emulated quad/octal precision.
  322. [default=True]
  323. glsl_types: bool, optional
  324. Set this argument to ``False`` in order to avoid using the ``vec`` and ``mat``
  325. types. The printer will instead use arrays (or nested arrays).
  326. [default=True]
  327. mat_nested: bool, optional
  328. GLSL version 4.3 and above support nested arrays (arrays of arrays). Set this to ``True``
  329. to render matrices as nested arrays.
  330. [default=False]
  331. mat_separator: str, optional
  332. By default, matrices are rendered with newlines using this separator,
  333. making them easier to read, but less compact. By removing the newline
  334. this option can be used to make them more vertically compact.
  335. [default=',\n']
  336. mat_transpose: bool, optional
  337. GLSL's matrix multiplication implementation assumes column-major indexing.
  338. By default, this printer ignores that convention. Setting this option to
  339. ``True`` transposes all matrix output.
  340. [default=False]
  341. array_type: str, optional
  342. The GLSL array constructor type.
  343. [default='float']
  344. precision : integer, optional
  345. The precision for numbers such as pi [default=15].
  346. user_functions : dict, optional
  347. A dictionary where keys are ``FunctionClass`` instances and values are
  348. their string representations. Alternatively, the dictionary value can
  349. be a list of tuples i.e. [(argument_test, js_function_string)]. See
  350. below for examples.
  351. human : bool, optional
  352. If True, the result is a single string that may contain some constant
  353. declarations for the number symbols. If False, the same information is
  354. returned in a tuple of (symbols_to_declare, not_supported_functions,
  355. code_text). [default=True].
  356. contract: bool, optional
  357. If True, ``Indexed`` instances are assumed to obey tensor contraction
  358. rules and the corresponding nested loops over indices are generated.
  359. Setting contract=False will not generate loops, instead the user is
  360. responsible to provide values for the indices in the code.
  361. [default=True].
  362. Examples
  363. ========
  364. >>> from sympy import glsl_code, symbols, Rational, sin, ceiling, Abs
  365. >>> x, tau = symbols("x, tau")
  366. >>> glsl_code((2*tau)**Rational(7, 2))
  367. '8*sqrt(2)*pow(tau, 3.5)'
  368. >>> glsl_code(sin(x), assign_to="float y")
  369. 'float y = sin(x);'
  370. Various GLSL types are supported:
  371. >>> from sympy import Matrix, glsl_code
  372. >>> glsl_code(Matrix([1,2,3]))
  373. 'vec3(1, 2, 3)'
  374. >>> glsl_code(Matrix([[1, 2],[3, 4]]))
  375. 'mat2(1, 2, 3, 4)'
  376. Pass ``mat_transpose = True`` to switch to column-major indexing:
  377. >>> glsl_code(Matrix([[1, 2],[3, 4]]), mat_transpose = True)
  378. 'mat2(1, 3, 2, 4)'
  379. By default, larger matrices get collapsed into float arrays:
  380. >>> print(glsl_code( Matrix([[1,2,3,4,5],[6,7,8,9,10]]) ))
  381. float[10](
  382. 1, 2, 3, 4, 5,
  383. 6, 7, 8, 9, 10
  384. ) /* a 2x5 matrix */
  385. The type of array constructor used to print GLSL arrays can be controlled
  386. via the ``array_type`` parameter:
  387. >>> glsl_code(Matrix([1,2,3,4,5]), array_type='int')
  388. 'int[5](1, 2, 3, 4, 5)'
  389. Passing a list of strings or ``symbols`` to the ``assign_to`` parameter will yield
  390. a multi-line assignment for each item in an array-like expression:
  391. >>> x_struct_members = symbols('x.a x.b x.c x.d')
  392. >>> print(glsl_code(Matrix([1,2,3,4]), assign_to=x_struct_members))
  393. x.a = 1;
  394. x.b = 2;
  395. x.c = 3;
  396. x.d = 4;
  397. This could be useful in cases where it's desirable to modify members of a
  398. GLSL ``Struct``. It could also be used to spread items from an array-like
  399. expression into various miscellaneous assignments:
  400. >>> misc_assignments = ('x[0]', 'x[1]', 'float y', 'float z')
  401. >>> print(glsl_code(Matrix([1,2,3,4]), assign_to=misc_assignments))
  402. x[0] = 1;
  403. x[1] = 2;
  404. float y = 3;
  405. float z = 4;
  406. Passing ``mat_nested = True`` instead prints out nested float arrays, which are
  407. supported in GLSL 4.3 and above.
  408. >>> mat = Matrix([
  409. ... [ 0, 1, 2],
  410. ... [ 3, 4, 5],
  411. ... [ 6, 7, 8],
  412. ... [ 9, 10, 11],
  413. ... [12, 13, 14]])
  414. >>> print(glsl_code( mat, mat_nested = True ))
  415. float[5][3](
  416. float[]( 0, 1, 2),
  417. float[]( 3, 4, 5),
  418. float[]( 6, 7, 8),
  419. float[]( 9, 10, 11),
  420. float[](12, 13, 14)
  421. )
  422. Custom printing can be defined for certain types by passing a dictionary of
  423. "type" : "function" to the ``user_functions`` kwarg. Alternatively, the
  424. dictionary value can be a list of tuples i.e. [(argument_test,
  425. js_function_string)].
  426. >>> custom_functions = {
  427. ... "ceiling": "CEIL",
  428. ... "Abs": [(lambda x: not x.is_integer, "fabs"),
  429. ... (lambda x: x.is_integer, "ABS")]
  430. ... }
  431. >>> glsl_code(Abs(x) + ceiling(x), user_functions=custom_functions)
  432. 'fabs(x) + CEIL(x)'
  433. If further control is needed, addition, subtraction, multiplication and
  434. division operators can be replaced with ``add``, ``sub``, and ``mul``
  435. functions. This is done by passing ``use_operators = False``:
  436. >>> x,y,z = symbols('x,y,z')
  437. >>> glsl_code(x*(y+z), use_operators = False)
  438. 'mul(x, add(y, z))'
  439. >>> glsl_code(x*(y+z*(x-y)**z), use_operators = False)
  440. 'mul(x, add(y, mul(z, pow(sub(x, y), z))))'
  441. ``Piecewise`` expressions are converted into conditionals. If an
  442. ``assign_to`` variable is provided an if statement is created, otherwise
  443. the ternary operator is used. Note that if the ``Piecewise`` lacks a
  444. default term, represented by ``(expr, True)`` then an error will be thrown.
  445. This is to prevent generating an expression that may not evaluate to
  446. anything.
  447. >>> from sympy import Piecewise
  448. >>> expr = Piecewise((x + 1, x > 0), (x, True))
  449. >>> print(glsl_code(expr, tau))
  450. if (x > 0) {
  451. tau = x + 1;
  452. }
  453. else {
  454. tau = x;
  455. }
  456. Support for loops is provided through ``Indexed`` types. With
  457. ``contract=True`` these expressions will be turned into loops, whereas
  458. ``contract=False`` will just print the assignment expression that should be
  459. looped over:
  460. >>> from sympy import Eq, IndexedBase, Idx
  461. >>> len_y = 5
  462. >>> y = IndexedBase('y', shape=(len_y,))
  463. >>> t = IndexedBase('t', shape=(len_y,))
  464. >>> Dy = IndexedBase('Dy', shape=(len_y-1,))
  465. >>> i = Idx('i', len_y-1)
  466. >>> e=Eq(Dy[i], (y[i+1]-y[i])/(t[i+1]-t[i]))
  467. >>> glsl_code(e.rhs, assign_to=e.lhs, contract=False)
  468. 'Dy[i] = (y[i + 1] - y[i])/(t[i + 1] - t[i]);'
  469. >>> from sympy import Matrix, MatrixSymbol
  470. >>> mat = Matrix([x**2, Piecewise((x + 1, x > 0), (x, True)), sin(x)])
  471. >>> A = MatrixSymbol('A', 3, 1)
  472. >>> print(glsl_code(mat, A))
  473. A[0][0] = pow(x, 2.0);
  474. if (x > 0) {
  475. A[1][0] = x + 1;
  476. }
  477. else {
  478. A[1][0] = x;
  479. }
  480. A[2][0] = sin(x);
  481. """
  482. return GLSLPrinter(settings).doprint(expr,assign_to)
  483. def print_glsl(expr, **settings):
  484. """Prints the GLSL representation of the given expression.
  485. See GLSLPrinter init function for settings.
  486. """
  487. print(glsl_code(expr, **settings))