Coverage for /usr/lib/python3/dist-packages/sympy/printing/julia.py: 21%

1"""

2Julia code printer

4The `JuliaCodePrinter` converts SymPy expressions into Julia expressions.

6A complete code generator, which uses `julia_code` extensively, can be found

7in `sympy.utilities.codegen`. The `codegen` module can be used to generate

8complete source code files.

10"""

12from __future__ import annotations

13from typing import Any

15from sympy.core import Mul, Pow, S, Rational

16from sympy.core.mul import _keep_coeff

17from sympy.core.numbers import equal_valued

18from sympy.printing.codeprinter import CodePrinter

19from sympy.printing.precedence import precedence, PRECEDENCE

20from re import search

22# List of known functions. First, those that have the same name in

23# SymPy and Julia. This is almost certainly incomplete!

24known_fcns_src1 = ["sin", "cos", "tan", "cot", "sec", "csc",

25 "asin", "acos", "atan", "acot", "asec", "acsc",

26 "sinh", "cosh", "tanh", "coth", "sech", "csch",

27 "asinh", "acosh", "atanh", "acoth", "asech", "acsch",

28 "sinc", "atan2", "sign", "floor", "log", "exp",

29 "cbrt", "sqrt", "erf", "erfc", "erfi",

30 "factorial", "gamma", "digamma", "trigamma",

31 "polygamma", "beta",

32 "airyai", "airyaiprime", "airybi", "airybiprime",

33 "besselj", "bessely", "besseli", "besselk",

34 "erfinv", "erfcinv"]

35# These functions have different names ("SymPy": "Julia"), more

36# generally a mapping to (argument_conditions, julia_function).

37known_fcns_src2 = {

38 "Abs": "abs",

39 "ceiling": "ceil",

40 "conjugate": "conj",

41 "hankel1": "hankelh1",

42 "hankel2": "hankelh2",

43 "im": "imag",

44 "re": "real"

45}

48class JuliaCodePrinter(CodePrinter):

49 """

50 A printer to convert expressions to strings of Julia code.

51 """

52 printmethod = "_julia"

53 language = "Julia"

55 _operators = {

56 'and': '&&',

57 'or': '||',

58 'not': '!',

59 }

61 _default_settings: dict[str, Any] = {

62 'order': None,

63 'full_prec': 'auto',

64 'precision': 17,

65 'user_functions': {},

66 'human': True,

67 'allow_unknown_functions': False,

68 'contract': True,

69 'inline': True,

70 }

71 # Note: contract is for expressing tensors as loops (if True), or just

72 # assignment (if False). FIXME: this should be looked a more carefully

73 # for Julia.

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

76 super().__init__(settings)

77 self.known_functions = dict(zip(known_fcns_src1, known_fcns_src1))

78 self.known_functions.update(dict(known_fcns_src2))

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

80 self.known_functions.update(userfuncs)

83 def _rate_index_position(self, p):

84 return p*5

87 def _get_statement(self, codestring):

88 return "%s" % codestring

91 def _get_comment(self, text):

92 return "# {}".format(text)

95 def _declare_number_const(self, name, value):

96 return "const {} = {}".format(name, value)

99 def _format_code(self, lines):

100 return self.indent_code(lines)

101

102

103 def _traverse_matrix_indices(self, mat):

104 # Julia uses Fortran order (column-major)

105 rows, cols = mat.shape

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

107

108

109 def _get_loop_opening_ending(self, indices):

110 open_lines = []

111 close_lines = []

112 for i in indices:

113 # Julia arrays start at 1 and end at dimension

114 var, start, stop = map(self._print,

115 [i.label, i.lower + 1, i.upper + 1])

116 open_lines.append("for %s = %s:%s" % (var, start, stop))

117 close_lines.append("end")

118 return open_lines, close_lines

119

120

121 def _print_Mul(self, expr):

122 # print complex numbers nicely in Julia

123 if (expr.is_number and expr.is_imaginary and

124 expr.as_coeff_Mul()[0].is_integer):

125 return "%sim" % self._print(-S.ImaginaryUnit*expr)

126

127 # cribbed from str.py

128 prec = precedence(expr)

129

130 c, e = expr.as_coeff_Mul()

131 if c < 0:

132 expr = _keep_coeff(-c, e)

133 sign = "-"

134 else:

135 sign = ""

136

137 a = [] # items in the numerator

138 b = [] # items that are in the denominator (if any)

139

140 pow_paren = [] # Will collect all pow with more than one base element and exp = -1

141

142 if self.order not in ('old', 'none'):

143 args = expr.as_ordered_factors()

144 else:

145 # use make_args in case expr was something like -x -> x

146 args = Mul.make_args(expr)

147

148 # Gather args for numerator/denominator

149 for item in args:

150 if (item.is_commutative and item.is_Pow and item.exp.is_Rational

151 and item.exp.is_negative):

152 if item.exp != -1:

153 b.append(Pow(item.base, -item.exp, evaluate=False))

154 else:

155 if len(item.args[0].args) != 1 and isinstance(item.base, Mul): # To avoid situations like #14160

156 pow_paren.append(item)

157 b.append(Pow(item.base, -item.exp))

158 elif item.is_Rational and item is not S.Infinity and item.p == 1:

159 # Save the Rational type in julia Unless the numerator is 1.

160 # For example:

161 # julia_code(Rational(3, 7)*x) --> (3 // 7) * x

162 # julia_code(x/3) --> x / 3 but not x * (1 // 3)

163 b.append(Rational(item.q))

164 else:

165 a.append(item)

166

167 a = a or [S.One]

168

169 a_str = [self.parenthesize(x, prec) for x in a]

170 b_str = [self.parenthesize(x, prec) for x in b]

171

172 # To parenthesize Pow with exp = -1 and having more than one Symbol

173 for item in pow_paren:

174 if item.base in b:

175 b_str[b.index(item.base)] = "(%s)" % b_str[b.index(item.base)]

176

177 # from here it differs from str.py to deal with "*" and ".*"

178 def multjoin(a, a_str):

179 # here we probably are assuming the constants will come first

180 r = a_str[0]

181 for i in range(1, len(a)):

182 mulsym = '*' if a[i-1].is_number else '.*'

183 r = "%s %s %s" % (r, mulsym, a_str[i])

184 return r

185

186 if not b:

187 return sign + multjoin(a, a_str)

188 elif len(b) == 1:

189 divsym = '/' if b[0].is_number else './'

190 return "%s %s %s" % (sign+multjoin(a, a_str), divsym, b_str[0])

191 else:

192 divsym = '/' if all(bi.is_number for bi in b) else './'

193 return "%s %s (%s)" % (sign + multjoin(a, a_str), divsym, multjoin(b, b_str))

194

195 def _print_Relational(self, expr):

196 lhs_code = self._print(expr.lhs)

197 rhs_code = self._print(expr.rhs)

198 op = expr.rel_op

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

200

201 def _print_Pow(self, expr):

202 powsymbol = '^' if all(x.is_number for x in expr.args) else '.^'

203

204 PREC = precedence(expr)

205

206 if equal_valued(expr.exp, 0.5):

207 return "sqrt(%s)" % self._print(expr.base)

208

209 if expr.is_commutative:

210 if equal_valued(expr.exp, -0.5):

211 sym = '/' if expr.base.is_number else './'

212 return "1 %s sqrt(%s)" % (sym, self._print(expr.base))

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

214 sym = '/' if expr.base.is_number else './'

215 return "1 %s %s" % (sym, self.parenthesize(expr.base, PREC))

216

217 return '%s %s %s' % (self.parenthesize(expr.base, PREC), powsymbol,

218 self.parenthesize(expr.exp, PREC))

219

220

221 def _print_MatPow(self, expr):

222 PREC = precedence(expr)

223 return '%s ^ %s' % (self.parenthesize(expr.base, PREC),

224 self.parenthesize(expr.exp, PREC))

225

226

227 def _print_Pi(self, expr):

228 if self._settings["inline"]:

229 return "pi"

230 else:

231 return super()._print_NumberSymbol(expr)

232

233

234 def _print_ImaginaryUnit(self, expr):

235 return "im"

236

237

238 def _print_Exp1(self, expr):

239 if self._settings["inline"]:

240 return "e"

241 else:

242 return super()._print_NumberSymbol(expr)

243

244

245 def _print_EulerGamma(self, expr):

246 if self._settings["inline"]:

247 return "eulergamma"

248 else:

249 return super()._print_NumberSymbol(expr)

250

251

252 def _print_Catalan(self, expr):

253 if self._settings["inline"]:

254 return "catalan"

255 else:

256 return super()._print_NumberSymbol(expr)

257

258

259 def _print_GoldenRatio(self, expr):

260 if self._settings["inline"]:

261 return "golden"

262 else:

263 return super()._print_NumberSymbol(expr)

264

265

266 def _print_Assignment(self, expr):

267 from sympy.codegen.ast import Assignment

268 from sympy.functions.elementary.piecewise import Piecewise

269 from sympy.tensor.indexed import IndexedBase

270 # Copied from codeprinter, but remove special MatrixSymbol treatment

271 lhs = expr.lhs

272 rhs = expr.rhs

273 # We special case assignments that take multiple lines

274 if not self._settings["inline"] and isinstance(expr.rhs, Piecewise):

275 # Here we modify Piecewise so each expression is now

276 # an Assignment, and then continue on the print.

277 expressions = []

278 conditions = []

279 for (e, c) in rhs.args:

280 expressions.append(Assignment(lhs, e))

281 conditions.append(c)

282 temp = Piecewise(*zip(expressions, conditions))

283 return self._print(temp)

284 if self._settings["contract"] and (lhs.has(IndexedBase) or

285 rhs.has(IndexedBase)):

286 # Here we check if there is looping to be done, and if so

287 # print the required loops.

288 return self._doprint_loops(rhs, lhs)

289 else:

290 lhs_code = self._print(lhs)

291 rhs_code = self._print(rhs)

292 return self._get_statement("%s = %s" % (lhs_code, rhs_code))

293

294

295 def _print_Infinity(self, expr):

296 return 'Inf'

297

298

299 def _print_NegativeInfinity(self, expr):

300 return '-Inf'

301

302

303 def _print_NaN(self, expr):

304 return 'NaN'

305

306

307 def _print_list(self, expr):

308 return 'Any[' + ', '.join(self._print(a) for a in expr) + ']'

309

310

311 def _print_tuple(self, expr):

312 if len(expr) == 1:

313 return "(%s,)" % self._print(expr[0])

314 else:

315 return "(%s)" % self.stringify(expr, ", ")

316 _print_Tuple = _print_tuple

317

318

319 def _print_BooleanTrue(self, expr):

320 return "true"

321

322

323 def _print_BooleanFalse(self, expr):

324 return "false"

325

326

327 def _print_bool(self, expr):

328 return str(expr).lower()

329

330

331 # Could generate quadrature code for definite Integrals?

332 #_print_Integral = _print_not_supported

333

334

335 def _print_MatrixBase(self, A):

336 # Handle zero dimensions:

337 if S.Zero in A.shape:

338 return 'zeros(%s, %s)' % (A.rows, A.cols)

339 elif (A.rows, A.cols) == (1, 1):

340 return "[%s]" % A[0, 0]

341 elif A.rows == 1:

342 return "[%s]" % A.table(self, rowstart='', rowend='', colsep=' ')

343 elif A.cols == 1:

344 # note .table would unnecessarily equispace the rows

345 return "[%s]" % ", ".join([self._print(a) for a in A])

346 return "[%s]" % A.table(self, rowstart='', rowend='',

347 rowsep=';\n', colsep=' ')

348

349

350 def _print_SparseRepMatrix(self, A):

351 from sympy.matrices import Matrix

352 L = A.col_list();

353 # make row vectors of the indices and entries

354 I = Matrix([k[0] + 1 for k in L])

355 J = Matrix([k[1] + 1 for k in L])

356 AIJ = Matrix([k[2] for k in L])

357 return "sparse(%s, %s, %s, %s, %s)" % (self._print(I), self._print(J),

358 self._print(AIJ), A.rows, A.cols)

359

360

361 def _print_MatrixElement(self, expr):

362 return self.parenthesize(expr.parent, PRECEDENCE["Atom"], strict=True) \

363 + '[%s,%s]' % (expr.i + 1, expr.j + 1)

364

365

366 def _print_MatrixSlice(self, expr):

367 def strslice(x, lim):

368 l = x[0] + 1

369 h = x[1]

370 step = x[2]

371 lstr = self._print(l)

372 hstr = 'end' if h == lim else self._print(h)

373 if step == 1:

374 if l == 1 and h == lim:

375 return ':'

376 if l == h:

377 return lstr

378 else:

379 return lstr + ':' + hstr

380 else:

381 return ':'.join((lstr, self._print(step), hstr))

382 return (self._print(expr.parent) + '[' +

383 strslice(expr.rowslice, expr.parent.shape[0]) + ',' +

384 strslice(expr.colslice, expr.parent.shape[1]) + ']')

385

386

387 def _print_Indexed(self, expr):

388 inds = [ self._print(i) for i in expr.indices ]

389 return "%s[%s]" % (self._print(expr.base.label), ",".join(inds))

390

391

392 def _print_Idx(self, expr):

393 return self._print(expr.label)

394

395

396 def _print_Identity(self, expr):

397 return "eye(%s)" % self._print(expr.shape[0])

398

399 def _print_HadamardProduct(self, expr):

400 return ' .* '.join([self.parenthesize(arg, precedence(expr))

401 for arg in expr.args])

402

403 def _print_HadamardPower(self, expr):

404 PREC = precedence(expr)

405 return '.**'.join([

406 self.parenthesize(expr.base, PREC),

407 self.parenthesize(expr.exp, PREC)

408 ])

409

410 def _print_Rational(self, expr):

411 if expr.q == 1:

412 return str(expr.p)

413 return "%s // %s" % (expr.p, expr.q)

414

415 # Note: as of 2022, Julia doesn't have spherical Bessel functions

416 def _print_jn(self, expr):

417 from sympy.functions import sqrt, besselj

418 x = expr.argument

419 expr2 = sqrt(S.Pi/(2*x))*besselj(expr.order + S.Half, x)

420 return self._print(expr2)

421

422

423 def _print_yn(self, expr):

424 from sympy.functions import sqrt, bessely

425 x = expr.argument

426 expr2 = sqrt(S.Pi/(2*x))*bessely(expr.order + S.Half, x)

427 return self._print(expr2)

428

429

430 def _print_Piecewise(self, expr):

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

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

433 # function may not return a result.

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

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

436 "condition. Without one, the generated "

437 "expression may not evaluate to anything under "

438 "some condition.")

439 lines = []

440 if self._settings["inline"]:

441 # Express each (cond, expr) pair in a nested Horner form:

442 # (condition) .* (expr) + (not cond) .* (<others>)

443 # Expressions that result in multiple statements won't work here.

444 ecpairs = ["({}) ? ({}) :".format

445 (self._print(c), self._print(e))

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

447 elast = " (%s)" % self._print(expr.args[-1].expr)

448 pw = "\n".join(ecpairs) + elast

449 # Note: current need these outer brackets for 2*pw. Would be

450 # nicer to teach parenthesize() to do this for us when needed!

451 return "(" + pw + ")"

452 else:

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

454 if i == 0:

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

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

457 lines.append("else")

458 else:

459 lines.append("elseif (%s)" % self._print(c))

460 code0 = self._print(e)

461 lines.append(code0)

462 if i == len(expr.args) - 1:

463 lines.append("end")

464 return "\n".join(lines)

465

466 def _print_MatMul(self, expr):

467 c, m = expr.as_coeff_mmul()

468

469 sign = ""

470 if c.is_number:

471 re, im = c.as_real_imag()

472 if im.is_zero and re.is_negative:

473 expr = _keep_coeff(-c, m)

474 sign = "-"

475 elif re.is_zero and im.is_negative:

476 expr = _keep_coeff(-c, m)

477 sign = "-"

478

479 return sign + ' * '.join(

480 (self.parenthesize(arg, precedence(expr)) for arg in expr.args)

481 )

482

483

484 def indent_code(self, code):

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

486

487 # code mostly copied from ccode

488 if isinstance(code, str):

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

490 return ''.join(code_lines)

491

492 tab = " "

493 inc_regex = ('^function ', '^if ', '^elseif ', '^else$', '^for ')

494 dec_regex = ('^end$', '^elseif ', '^else$')

495

496 # pre-strip left-space from the code

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

498

499 increase = [ int(any(search(re, line) for re in inc_regex))

500 for line in code ]

501 decrease = [ int(any(search(re, line) for re in dec_regex))

502 for line in code ]

503

504 pretty = []

505 level = 0

506 for n, line in enumerate(code):

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

508 pretty.append(line)

509 continue

510 level -= decrease[n]

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

512 level += increase[n]

513 return pretty

514

515

516def julia_code(expr, assign_to=None, **settings):

517 r"""Converts `expr` to a string of Julia code.

518

519 Parameters

520 ==========

521

522 expr : Expr

523 A SymPy expression to be converted.

524 assign_to : optional

525 When given, the argument is used as the name of the variable to which

526 the expression is assigned. Can be a string, ``Symbol``,

527 ``MatrixSymbol``, or ``Indexed`` type. This can be helpful for

528 expressions that generate multi-line statements.

529 precision : integer, optional

530 The precision for numbers such as pi [default=16].

531 user_functions : dict, optional

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

533 their string representations. Alternatively, the dictionary value can

534 be a list of tuples i.e. [(argument_test, cfunction_string)]. See

535 below for examples.

536 human : bool, optional

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

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

539 returned in a tuple of (symbols_to_declare, not_supported_functions,

540 code_text). [default=True].

541 contract: bool, optional

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

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

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

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

546 [default=True].

547 inline: bool, optional

548 If True, we try to create single-statement code instead of multiple

549 statements. [default=True].

550

551 Examples

552 ========

553

554 >>> from sympy import julia_code, symbols, sin, pi

555 >>> x = symbols('x')

556 >>> julia_code(sin(x).series(x).removeO())

557 'x .^ 5 / 120 - x .^ 3 / 6 + x'

558

559 >>> from sympy import Rational, ceiling

560 >>> x, y, tau = symbols("x, y, tau")

561 >>> julia_code((2*tau)**Rational(7, 2))

562 '8 * sqrt(2) * tau .^ (7 // 2)'

563

564 Note that element-wise (Hadamard) operations are used by default between

565 symbols. This is because its possible in Julia to write "vectorized"

566 code. It is harmless if the values are scalars.

567

568 >>> julia_code(sin(pi*x*y), assign_to="s")

569 's = sin(pi * x .* y)'

570

571 If you need a matrix product "*" or matrix power "^", you can specify the

572 symbol as a ``MatrixSymbol``.

573

574 >>> from sympy import Symbol, MatrixSymbol

575 >>> n = Symbol('n', integer=True, positive=True)

576 >>> A = MatrixSymbol('A', n, n)

577 >>> julia_code(3*pi*A**3)

578 '(3 * pi) * A ^ 3'

579

580 This class uses several rules to decide which symbol to use a product.

581 Pure numbers use "*", Symbols use ".*" and MatrixSymbols use "*".

582 A HadamardProduct can be used to specify componentwise multiplication ".*"

583 of two MatrixSymbols. There is currently there is no easy way to specify

584 scalar symbols, so sometimes the code might have some minor cosmetic

585 issues. For example, suppose x and y are scalars and A is a Matrix, then

586 while a human programmer might write "(x^2*y)*A^3", we generate:

587

588 >>> julia_code(x**2*y*A**3)

589 '(x .^ 2 .* y) * A ^ 3'

590

591 Matrices are supported using Julia inline notation. When using

592 ``assign_to`` with matrices, the name can be specified either as a string

593 or as a ``MatrixSymbol``. The dimensions must align in the latter case.

594

595 >>> from sympy import Matrix, MatrixSymbol

596 >>> mat = Matrix([[x**2, sin(x), ceiling(x)]])

597 >>> julia_code(mat, assign_to='A')

598 'A = [x .^ 2 sin(x) ceil(x)]'

599

600 ``Piecewise`` expressions are implemented with logical masking by default.

601 Alternatively, you can pass "inline=False" to use if-else conditionals.

602 Note that if the ``Piecewise`` lacks a default term, represented by

603 ``(expr, True)`` then an error will be thrown. This is to prevent

604 generating an expression that may not evaluate to anything.

605

606 >>> from sympy import Piecewise

607 >>> pw = Piecewise((x + 1, x > 0), (x, True))

608 >>> julia_code(pw, assign_to=tau)

609 'tau = ((x > 0) ? (x + 1) : (x))'

610

611 Note that any expression that can be generated normally can also exist

612 inside a Matrix:

613

614 >>> mat = Matrix([[x**2, pw, sin(x)]])

615 >>> julia_code(mat, assign_to='A')

616 'A = [x .^ 2 ((x > 0) ? (x + 1) : (x)) sin(x)]'

617

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

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

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

621 cfunction_string)]. This can be used to call a custom Julia function.

622

623 >>> from sympy import Function

624 >>> f = Function('f')

625 >>> g = Function('g')

626 >>> custom_functions = {

627 ... "f": "existing_julia_fcn",

628 ... "g": [(lambda x: x.is_Matrix, "my_mat_fcn"),

629 ... (lambda x: not x.is_Matrix, "my_fcn")]

630 ... }

631 >>> mat = Matrix([[1, x]])

632 >>> julia_code(f(x) + g(x) + g(mat), user_functions=custom_functions)

633 'existing_julia_fcn(x) + my_fcn(x) + my_mat_fcn([1 x])'

634

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

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

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

638 looped over:

639

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

641 >>> len_y = 5

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

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

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

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

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

647 >>> julia_code(e.rhs, assign_to=e.lhs, contract=False)

648 'Dy[i] = (y[i + 1] - y[i]) ./ (t[i + 1] - t[i])'

649 """

650 return JuliaCodePrinter(settings).doprint(expr, assign_to)

651

652

653def print_julia_code(expr, **settings):

654 """Prints the Julia representation of the given expression.

655

656 See `julia_code` for the meaning of the optional arguments.

657 """

658 print(julia_code(expr, **settings))