Coverage for /usr/lib/python3/dist-packages/sympy/printing/glsl.py: 20%

205 statements  

« prev     ^ index     » next       coverage.py v7.9.1, created at 2025-06-14 15:55 +0200

1from __future__ import annotations 

2 

3from sympy.core import Basic, S 

4from sympy.core.function import Lambda 

5from sympy.core.numbers import equal_valued 

6from sympy.printing.codeprinter import CodePrinter 

7from sympy.printing.precedence import precedence 

8from functools import reduce 

9 

10known_functions = { 

11 'Abs': 'abs', 

12 'sin': 'sin', 

13 'cos': 'cos', 

14 'tan': 'tan', 

15 'acos': 'acos', 

16 'asin': 'asin', 

17 'atan': 'atan', 

18 'atan2': 'atan', 

19 'ceiling': 'ceil', 

20 'floor': 'floor', 

21 'sign': 'sign', 

22 'exp': 'exp', 

23 'log': 'log', 

24 'add': 'add', 

25 'sub': 'sub', 

26 'mul': 'mul', 

27 'pow': 'pow' 

28} 

29 

30class GLSLPrinter(CodePrinter): 

31 """ 

32 Rudimentary, generic GLSL printing tools. 

33 

34 Additional settings: 

35 'use_operators': Boolean (should the printer use operators for +,-,*, or functions?) 

36 """ 

37 _not_supported: set[Basic] = set() 

38 printmethod = "_glsl" 

39 language = "GLSL" 

40 

41 _default_settings = { 

42 'use_operators': True, 

43 'zero': 0, 

44 'mat_nested': False, 

45 'mat_separator': ',\n', 

46 'mat_transpose': False, 

47 'array_type': 'float', 

48 'glsl_types': True, 

49 

50 'order': None, 

51 'full_prec': 'auto', 

52 'precision': 9, 

53 'user_functions': {}, 

54 'human': True, 

55 'allow_unknown_functions': False, 

56 'contract': True, 

57 'error_on_reserved': False, 

58 'reserved_word_suffix': '_', 

59 } 

60 

61 def __init__(self, settings={}): 

62 CodePrinter.__init__(self, settings) 

63 self.known_functions = dict(known_functions) 

64 userfuncs = settings.get('user_functions', {}) 

65 self.known_functions.update(userfuncs) 

66 

67 def _rate_index_position(self, p): 

68 return p*5 

69 

70 def _get_statement(self, codestring): 

71 return "%s;" % codestring 

72 

73 def _get_comment(self, text): 

74 return "// {}".format(text) 

75 

76 def _declare_number_const(self, name, value): 

77 return "float {} = {};".format(name, value) 

78 

79 def _format_code(self, lines): 

80 return self.indent_code(lines) 

81 

82 def indent_code(self, code): 

83 """Accepts a string of code or a list of code lines""" 

84 

85 if isinstance(code, str): 

86 code_lines = self.indent_code(code.splitlines(True)) 

87 return ''.join(code_lines) 

88 

89 tab = " " 

90 inc_token = ('{', '(', '{\n', '(\n') 

91 dec_token = ('}', ')') 

92 

93 code = [line.lstrip(' \t') for line in code] 

94 

95 increase = [int(any(map(line.endswith, inc_token))) for line in code] 

96 decrease = [int(any(map(line.startswith, dec_token))) for line in code] 

97 

98 pretty = [] 

99 level = 0 

100 for n, line in enumerate(code): 

101 if line in ('', '\n'): 

102 pretty.append(line) 

103 continue 

104 level -= decrease[n] 

105 pretty.append("%s%s" % (tab*level, line)) 

106 level += increase[n] 

107 return pretty 

108 

109 def _print_MatrixBase(self, mat): 

110 mat_separator = self._settings['mat_separator'] 

111 mat_transpose = self._settings['mat_transpose'] 

112 column_vector = (mat.rows == 1) if mat_transpose else (mat.cols == 1) 

113 A = mat.transpose() if mat_transpose != column_vector else mat 

114 

115 glsl_types = self._settings['glsl_types'] 

116 array_type = self._settings['array_type'] 

117 array_size = A.cols*A.rows 

118 array_constructor = "{}[{}]".format(array_type, array_size) 

119 

120 if A.cols == 1: 

121 return self._print(A[0]); 

122 if A.rows <= 4 and A.cols <= 4 and glsl_types: 

123 if A.rows == 1: 

124 return "vec{}{}".format( 

125 A.cols, A.table(self,rowstart='(',rowend=')') 

126 ) 

127 elif A.rows == A.cols: 

128 return "mat{}({})".format( 

129 A.rows, A.table(self,rowsep=', ', 

130 rowstart='',rowend='') 

131 ) 

132 else: 

133 return "mat{}x{}({})".format( 

134 A.cols, A.rows, 

135 A.table(self,rowsep=', ', 

136 rowstart='',rowend='') 

137 ) 

138 elif S.One in A.shape: 

139 return "{}({})".format( 

140 array_constructor, 

141 A.table(self,rowsep=mat_separator,rowstart='',rowend='') 

142 ) 

143 elif not self._settings['mat_nested']: 

144 return "{}(\n{}\n) /* a {}x{} matrix */".format( 

145 array_constructor, 

146 A.table(self,rowsep=mat_separator,rowstart='',rowend=''), 

147 A.rows, A.cols 

148 ) 

149 elif self._settings['mat_nested']: 

150 return "{}[{}][{}](\n{}\n)".format( 

151 array_type, A.rows, A.cols, 

152 A.table(self,rowsep=mat_separator,rowstart='float[](',rowend=')') 

153 ) 

154 

155 def _print_SparseRepMatrix(self, mat): 

156 # do not allow sparse matrices to be made dense 

157 return self._print_not_supported(mat) 

158 

159 def _traverse_matrix_indices(self, mat): 

160 mat_transpose = self._settings['mat_transpose'] 

161 if mat_transpose: 

162 rows,cols = mat.shape 

163 else: 

164 cols,rows = mat.shape 

165 return ((i, j) for i in range(cols) for j in range(rows)) 

166 

167 def _print_MatrixElement(self, expr): 

168 # print('begin _print_MatrixElement') 

169 nest = self._settings['mat_nested']; 

170 glsl_types = self._settings['glsl_types']; 

171 mat_transpose = self._settings['mat_transpose']; 

172 if mat_transpose: 

173 cols,rows = expr.parent.shape 

174 i,j = expr.j,expr.i 

175 else: 

176 rows,cols = expr.parent.shape 

177 i,j = expr.i,expr.j 

178 pnt = self._print(expr.parent) 

179 if glsl_types and ((rows <= 4 and cols <=4) or nest): 

180 return "{}[{}][{}]".format(pnt, i, j) 

181 else: 

182 return "{}[{}]".format(pnt, i + j*rows) 

183 

184 def _print_list(self, expr): 

185 l = ', '.join(self._print(item) for item in expr) 

186 glsl_types = self._settings['glsl_types'] 

187 array_type = self._settings['array_type'] 

188 array_size = len(expr) 

189 array_constructor = '{}[{}]'.format(array_type, array_size) 

190 

191 if array_size <= 4 and glsl_types: 

192 return 'vec{}({})'.format(array_size, l) 

193 else: 

194 return '{}({})'.format(array_constructor, l) 

195 

196 _print_tuple = _print_list 

197 _print_Tuple = _print_list 

198 

199 def _get_loop_opening_ending(self, indices): 

200 open_lines = [] 

201 close_lines = [] 

202 loopstart = "for (int %(varble)s=%(start)s; %(varble)s<%(end)s; %(varble)s++){" 

203 for i in indices: 

204 # GLSL arrays start at 0 and end at dimension-1 

205 open_lines.append(loopstart % { 

206 'varble': self._print(i.label), 

207 'start': self._print(i.lower), 

208 'end': self._print(i.upper + 1)}) 

209 close_lines.append("}") 

210 return open_lines, close_lines 

211 

212 def _print_Function_with_args(self, func, func_args): 

213 if func in self.known_functions: 

214 cond_func = self.known_functions[func] 

215 func = None 

216 if isinstance(cond_func, str): 

217 func = cond_func 

218 else: 

219 for cond, func in cond_func: 

220 if cond(func_args): 

221 break 

222 if func is not None: 

223 try: 

224 return func(*[self.parenthesize(item, 0) for item in func_args]) 

225 except TypeError: 

226 return '{}({})'.format(func, self.stringify(func_args, ", ")) 

227 elif isinstance(func, Lambda): 

228 # inlined function 

229 return self._print(func(*func_args)) 

230 else: 

231 return self._print_not_supported(func) 

232 

233 def _print_Piecewise(self, expr): 

234 from sympy.codegen.ast import Assignment 

235 if expr.args[-1].cond != True: 

236 # We need the last conditional to be a True, otherwise the resulting 

237 # function may not return a result. 

238 raise ValueError("All Piecewise expressions must contain an " 

239 "(expr, True) statement to be used as a default " 

240 "condition. Without one, the generated " 

241 "expression may not evaluate to anything under " 

242 "some condition.") 

243 lines = [] 

244 if expr.has(Assignment): 

245 for i, (e, c) in enumerate(expr.args): 

246 if i == 0: 

247 lines.append("if (%s) {" % self._print(c)) 

248 elif i == len(expr.args) - 1 and c == True: 

249 lines.append("else {") 

250 else: 

251 lines.append("else if (%s) {" % self._print(c)) 

252 code0 = self._print(e) 

253 lines.append(code0) 

254 lines.append("}") 

255 return "\n".join(lines) 

256 else: 

257 # The piecewise was used in an expression, need to do inline 

258 # operators. This has the downside that inline operators will 

259 # not work for statements that span multiple lines (Matrix or 

260 # Indexed expressions). 

261 ecpairs = ["((%s) ? (\n%s\n)\n" % (self._print(c), 

262 self._print(e)) 

263 for e, c in expr.args[:-1]] 

264 last_line = ": (\n%s\n)" % self._print(expr.args[-1].expr) 

265 return ": ".join(ecpairs) + last_line + " ".join([")"*len(ecpairs)]) 

266 

267 def _print_Idx(self, expr): 

268 return self._print(expr.label) 

269 

270 def _print_Indexed(self, expr): 

271 # calculate index for 1d array 

272 dims = expr.shape 

273 elem = S.Zero 

274 offset = S.One 

275 for i in reversed(range(expr.rank)): 

276 elem += expr.indices[i]*offset 

277 offset *= dims[i] 

278 return "{}[{}]".format( 

279 self._print(expr.base.label), 

280 self._print(elem) 

281 ) 

282 

283 def _print_Pow(self, expr): 

284 PREC = precedence(expr) 

285 if equal_valued(expr.exp, -1): 

286 return '1.0/%s' % (self.parenthesize(expr.base, PREC)) 

287 elif equal_valued(expr.exp, 0.5): 

288 return 'sqrt(%s)' % self._print(expr.base) 

289 else: 

290 try: 

291 e = self._print(float(expr.exp)) 

292 except TypeError: 

293 e = self._print(expr.exp) 

294 return self._print_Function_with_args('pow', ( 

295 self._print(expr.base), 

296 e 

297 )) 

298 

299 def _print_int(self, expr): 

300 return str(float(expr)) 

301 

302 def _print_Rational(self, expr): 

303 return "{}.0/{}.0".format(expr.p, expr.q) 

304 

305 def _print_Relational(self, expr): 

306 lhs_code = self._print(expr.lhs) 

307 rhs_code = self._print(expr.rhs) 

308 op = expr.rel_op 

309 return "{} {} {}".format(lhs_code, op, rhs_code) 

310 

311 def _print_Add(self, expr, order=None): 

312 if self._settings['use_operators']: 

313 return CodePrinter._print_Add(self, expr, order=order) 

314 

315 terms = expr.as_ordered_terms() 

316 

317 def partition(p,l): 

318 return reduce(lambda x, y: (x[0]+[y], x[1]) if p(y) else (x[0], x[1]+[y]), l, ([], [])) 

319 def add(a,b): 

320 return self._print_Function_with_args('add', (a, b)) 

321 # return self.known_functions['add']+'(%s, %s)' % (a,b) 

322 neg, pos = partition(lambda arg: arg.could_extract_minus_sign(), terms) 

323 if pos: 

324 s = pos = reduce(lambda a,b: add(a,b), (self._print(t) for t in pos)) 

325 else: 

326 s = pos = self._print(self._settings['zero']) 

327 

328 if neg: 

329 # sum the absolute values of the negative terms 

330 neg = reduce(lambda a,b: add(a,b), (self._print(-n) for n in neg)) 

331 # then subtract them from the positive terms 

332 s = self._print_Function_with_args('sub', (pos,neg)) 

333 # s = self.known_functions['sub']+'(%s, %s)' % (pos,neg) 

334 return s 

335 

336 def _print_Mul(self, expr, **kwargs): 

337 if self._settings['use_operators']: 

338 return CodePrinter._print_Mul(self, expr, **kwargs) 

339 terms = expr.as_ordered_factors() 

340 def mul(a,b): 

341 # return self.known_functions['mul']+'(%s, %s)' % (a,b) 

342 return self._print_Function_with_args('mul', (a,b)) 

343 

344 s = reduce(lambda a,b: mul(a,b), (self._print(t) for t in terms)) 

345 return s 

346 

347def glsl_code(expr,assign_to=None,**settings): 

348 """Converts an expr to a string of GLSL code 

349 

350 Parameters 

351 ========== 

352 

353 expr : Expr 

354 A SymPy expression to be converted. 

355 assign_to : optional 

356 When given, the argument is used for naming the variable or variables 

357 to which the expression is assigned. Can be a string, ``Symbol``, 

358 ``MatrixSymbol`` or ``Indexed`` type object. In cases where ``expr`` 

359 would be printed as an array, a list of string or ``Symbol`` objects 

360 can also be passed. 

361 

362 This is helpful in case of line-wrapping, or for expressions that 

363 generate multi-line statements. It can also be used to spread an array-like 

364 expression into multiple assignments. 

365 use_operators: bool, optional 

366 If set to False, then *,/,+,- operators will be replaced with functions 

367 mul, add, and sub, which must be implemented by the user, e.g. for 

368 implementing non-standard rings or emulated quad/octal precision. 

369 [default=True] 

370 glsl_types: bool, optional 

371 Set this argument to ``False`` in order to avoid using the ``vec`` and ``mat`` 

372 types. The printer will instead use arrays (or nested arrays). 

373 [default=True] 

374 mat_nested: bool, optional 

375 GLSL version 4.3 and above support nested arrays (arrays of arrays). Set this to ``True`` 

376 to render matrices as nested arrays. 

377 [default=False] 

378 mat_separator: str, optional 

379 By default, matrices are rendered with newlines using this separator, 

380 making them easier to read, but less compact. By removing the newline 

381 this option can be used to make them more vertically compact. 

382 [default=',\n'] 

383 mat_transpose: bool, optional 

384 GLSL's matrix multiplication implementation assumes column-major indexing. 

385 By default, this printer ignores that convention. Setting this option to 

386 ``True`` transposes all matrix output. 

387 [default=False] 

388 array_type: str, optional 

389 The GLSL array constructor type. 

390 [default='float'] 

391 precision : integer, optional 

392 The precision for numbers such as pi [default=15]. 

393 user_functions : dict, optional 

394 A dictionary where keys are ``FunctionClass`` instances and values are 

395 their string representations. Alternatively, the dictionary value can 

396 be a list of tuples i.e. [(argument_test, js_function_string)]. See 

397 below for examples. 

398 human : bool, optional 

399 If True, the result is a single string that may contain some constant 

400 declarations for the number symbols. If False, the same information is 

401 returned in a tuple of (symbols_to_declare, not_supported_functions, 

402 code_text). [default=True]. 

403 contract: bool, optional 

404 If True, ``Indexed`` instances are assumed to obey tensor contraction 

405 rules and the corresponding nested loops over indices are generated. 

406 Setting contract=False will not generate loops, instead the user is 

407 responsible to provide values for the indices in the code. 

408 [default=True]. 

409 

410 Examples 

411 ======== 

412 

413 >>> from sympy import glsl_code, symbols, Rational, sin, ceiling, Abs 

414 >>> x, tau = symbols("x, tau") 

415 >>> glsl_code((2*tau)**Rational(7, 2)) 

416 '8*sqrt(2)*pow(tau, 3.5)' 

417 >>> glsl_code(sin(x), assign_to="float y") 

418 'float y = sin(x);' 

419 

420 Various GLSL types are supported: 

421 >>> from sympy import Matrix, glsl_code 

422 >>> glsl_code(Matrix([1,2,3])) 

423 'vec3(1, 2, 3)' 

424 

425 >>> glsl_code(Matrix([[1, 2],[3, 4]])) 

426 'mat2(1, 2, 3, 4)' 

427 

428 Pass ``mat_transpose = True`` to switch to column-major indexing: 

429 >>> glsl_code(Matrix([[1, 2],[3, 4]]), mat_transpose = True) 

430 'mat2(1, 3, 2, 4)' 

431 

432 By default, larger matrices get collapsed into float arrays: 

433 >>> print(glsl_code( Matrix([[1,2,3,4,5],[6,7,8,9,10]]) )) 

434 float[10]( 

435 1, 2, 3, 4, 5, 

436 6, 7, 8, 9, 10 

437 ) /* a 2x5 matrix */ 

438 

439 The type of array constructor used to print GLSL arrays can be controlled 

440 via the ``array_type`` parameter: 

441 >>> glsl_code(Matrix([1,2,3,4,5]), array_type='int') 

442 'int[5](1, 2, 3, 4, 5)' 

443 

444 Passing a list of strings or ``symbols`` to the ``assign_to`` parameter will yield 

445 a multi-line assignment for each item in an array-like expression: 

446 >>> x_struct_members = symbols('x.a x.b x.c x.d') 

447 >>> print(glsl_code(Matrix([1,2,3,4]), assign_to=x_struct_members)) 

448 x.a = 1; 

449 x.b = 2; 

450 x.c = 3; 

451 x.d = 4; 

452 

453 This could be useful in cases where it's desirable to modify members of a 

454 GLSL ``Struct``. It could also be used to spread items from an array-like 

455 expression into various miscellaneous assignments: 

456 >>> misc_assignments = ('x[0]', 'x[1]', 'float y', 'float z') 

457 >>> print(glsl_code(Matrix([1,2,3,4]), assign_to=misc_assignments)) 

458 x[0] = 1; 

459 x[1] = 2; 

460 float y = 3; 

461 float z = 4; 

462 

463 Passing ``mat_nested = True`` instead prints out nested float arrays, which are 

464 supported in GLSL 4.3 and above. 

465 >>> mat = Matrix([ 

466 ... [ 0, 1, 2], 

467 ... [ 3, 4, 5], 

468 ... [ 6, 7, 8], 

469 ... [ 9, 10, 11], 

470 ... [12, 13, 14]]) 

471 >>> print(glsl_code( mat, mat_nested = True )) 

472 float[5][3]( 

473 float[]( 0, 1, 2), 

474 float[]( 3, 4, 5), 

475 float[]( 6, 7, 8), 

476 float[]( 9, 10, 11), 

477 float[](12, 13, 14) 

478 ) 

479 

480 

481 

482 Custom printing can be defined for certain types by passing a dictionary of 

483 "type" : "function" to the ``user_functions`` kwarg. Alternatively, the 

484 dictionary value can be a list of tuples i.e. [(argument_test, 

485 js_function_string)]. 

486 

487 >>> custom_functions = { 

488 ... "ceiling": "CEIL", 

489 ... "Abs": [(lambda x: not x.is_integer, "fabs"), 

490 ... (lambda x: x.is_integer, "ABS")] 

491 ... } 

492 >>> glsl_code(Abs(x) + ceiling(x), user_functions=custom_functions) 

493 'fabs(x) + CEIL(x)' 

494 

495 If further control is needed, addition, subtraction, multiplication and 

496 division operators can be replaced with ``add``, ``sub``, and ``mul`` 

497 functions. This is done by passing ``use_operators = False``: 

498 

499 >>> x,y,z = symbols('x,y,z') 

500 >>> glsl_code(x*(y+z), use_operators = False) 

501 'mul(x, add(y, z))' 

502 >>> glsl_code(x*(y+z*(x-y)**z), use_operators = False) 

503 'mul(x, add(y, mul(z, pow(sub(x, y), z))))' 

504 

505 ``Piecewise`` expressions are converted into conditionals. If an 

506 ``assign_to`` variable is provided an if statement is created, otherwise 

507 the ternary operator is used. Note that if the ``Piecewise`` lacks a 

508 default term, represented by ``(expr, True)`` then an error will be thrown. 

509 This is to prevent generating an expression that may not evaluate to 

510 anything. 

511 

512 >>> from sympy import Piecewise 

513 >>> expr = Piecewise((x + 1, x > 0), (x, True)) 

514 >>> print(glsl_code(expr, tau)) 

515 if (x > 0) { 

516 tau = x + 1; 

517 } 

518 else { 

519 tau = x; 

520 } 

521 

522 Support for loops is provided through ``Indexed`` types. With 

523 ``contract=True`` these expressions will be turned into loops, whereas 

524 ``contract=False`` will just print the assignment expression that should be 

525 looped over: 

526 

527 >>> from sympy import Eq, IndexedBase, Idx 

528 >>> len_y = 5 

529 >>> y = IndexedBase('y', shape=(len_y,)) 

530 >>> t = IndexedBase('t', shape=(len_y,)) 

531 >>> Dy = IndexedBase('Dy', shape=(len_y-1,)) 

532 >>> i = Idx('i', len_y-1) 

533 >>> e=Eq(Dy[i], (y[i+1]-y[i])/(t[i+1]-t[i])) 

534 >>> glsl_code(e.rhs, assign_to=e.lhs, contract=False) 

535 'Dy[i] = (y[i + 1] - y[i])/(t[i + 1] - t[i]);' 

536 

537 >>> from sympy import Matrix, MatrixSymbol 

538 >>> mat = Matrix([x**2, Piecewise((x + 1, x > 0), (x, True)), sin(x)]) 

539 >>> A = MatrixSymbol('A', 3, 1) 

540 >>> print(glsl_code(mat, A)) 

541 A[0][0] = pow(x, 2.0); 

542 if (x > 0) { 

543 A[1][0] = x + 1; 

544 } 

545 else { 

546 A[1][0] = x; 

547 } 

548 A[2][0] = sin(x); 

549 """ 

550 return GLSLPrinter(settings).doprint(expr,assign_to) 

551 

552def print_glsl(expr, **settings): 

553 """Prints the GLSL representation of the given expression. 

554 

555 See GLSLPrinter init function for settings. 

556 """ 

557 print(glsl_code(expr, **settings))