Coverage for /usr/lib/python3/dist-packages/sympy/printing/numpy.py: 53%

1from sympy.core import S

2from sympy.core.function import Lambda

3from sympy.core.power import Pow

4from .pycode import PythonCodePrinter, _known_functions_math, _print_known_const, _print_known_func, _unpack_integral_limits, ArrayPrinter

5from .codeprinter import CodePrinter

8_not_in_numpy = 'erf erfc factorial gamma loggamma'.split()

9_in_numpy = [(k, v) for k, v in _known_functions_math.items() if k not in _not_in_numpy]

10_known_functions_numpy = dict(_in_numpy, **{

11 'acos': 'arccos',

12 'acosh': 'arccosh',

13 'asin': 'arcsin',

14 'asinh': 'arcsinh',

15 'atan': 'arctan',

16 'atan2': 'arctan2',

17 'atanh': 'arctanh',

18 'exp2': 'exp2',

19 'sign': 'sign',

20 'logaddexp': 'logaddexp',

21 'logaddexp2': 'logaddexp2',

22})

23_known_constants_numpy = {

24 'Exp1': 'e',

25 'Pi': 'pi',

26 'EulerGamma': 'euler_gamma',

27 'NaN': 'nan',

28 'Infinity': 'PINF',

29 'NegativeInfinity': 'NINF'

30}

32_numpy_known_functions = {k: 'numpy.' + v for k, v in _known_functions_numpy.items()}

33_numpy_known_constants = {k: 'numpy.' + v for k, v in _known_constants_numpy.items()}

35class NumPyPrinter(ArrayPrinter, PythonCodePrinter):

36 """

37 Numpy printer which handles vectorized piecewise functions,

38 logical operators, etc.

39 """

41 _module = 'numpy'

42 _kf = _numpy_known_functions

43 _kc = _numpy_known_constants

45 def __init__(self, settings=None):

46 """

47 `settings` is passed to CodePrinter.__init__()

48 `module` specifies the array module to use, currently 'NumPy', 'CuPy'

49 or 'JAX'.

50 """

51 self.language = "Python with {}".format(self._module)

52 self.printmethod = "_{}code".format(self._module)

54 self._kf = {**PythonCodePrinter._kf, **self._kf}

56 super().__init__(settings=settings)

59 def _print_seq(self, seq):

60 "General sequence printer: converts to tuple"

61 # Print tuples here instead of lists because numba supports

62 # tuples in nopython mode.

63 delimiter=', '

64 return '({},)'.format(delimiter.join(self._print(item) for item in seq))

66 def _print_MatMul(self, expr):

67 "Matrix multiplication printer"

68 if expr.as_coeff_matrices()[0] is not S.One:

69 expr_list = expr.as_coeff_matrices()[1]+[(expr.as_coeff_matrices()[0])]

70 return '({})'.format(').dot('.join(self._print(i) for i in expr_list))

71 return '({})'.format(').dot('.join(self._print(i) for i in expr.args))

73 def _print_MatPow(self, expr):

74 "Matrix power printer"

75 return '{}({}, {})'.format(self._module_format(self._module + '.linalg.matrix_power'),

76 self._print(expr.args[0]), self._print(expr.args[1]))

78 def _print_Inverse(self, expr):

79 "Matrix inverse printer"

80 return '{}({})'.format(self._module_format(self._module + '.linalg.inv'),

81 self._print(expr.args[0]))

83 def _print_DotProduct(self, expr):

84 # DotProduct allows any shape order, but numpy.dot does matrix

85 # multiplication, so we have to make sure it gets 1 x n by n x 1.

86 arg1, arg2 = expr.args

87 if arg1.shape[0] != 1:

88 arg1 = arg1.T

89 if arg2.shape[1] != 1:

90 arg2 = arg2.T

92 return "%s(%s, %s)" % (self._module_format(self._module + '.dot'),

93 self._print(arg1),

94 self._print(arg2))

96 def _print_MatrixSolve(self, expr):

97 return "%s(%s, %s)" % (self._module_format(self._module + '.linalg.solve'),

98 self._print(expr.matrix),

99 self._print(expr.vector))

100

101 def _print_ZeroMatrix(self, expr):

102 return '{}({})'.format(self._module_format(self._module + '.zeros'),

103 self._print(expr.shape))

104

105 def _print_OneMatrix(self, expr):

106 return '{}({})'.format(self._module_format(self._module + '.ones'),

107 self._print(expr.shape))

108

109 def _print_FunctionMatrix(self, expr):

110 from sympy.abc import i, j

111 lamda = expr.lamda

112 if not isinstance(lamda, Lambda):

113 lamda = Lambda((i, j), lamda(i, j))

114 return '{}(lambda {}: {}, {})'.format(self._module_format(self._module + '.fromfunction'),

115 ', '.join(self._print(arg) for arg in lamda.args[0]),

116 self._print(lamda.args[1]), self._print(expr.shape))

117

118 def _print_HadamardProduct(self, expr):

119 func = self._module_format(self._module + '.multiply')

120 return ''.join('{}({}, '.format(func, self._print(arg)) \

121 for arg in expr.args[:-1]) + "{}{}".format(self._print(expr.args[-1]),

122 ')' * (len(expr.args) - 1))

123

124 def _print_KroneckerProduct(self, expr):

125 func = self._module_format(self._module + '.kron')

126 return ''.join('{}({}, '.format(func, self._print(arg)) \

127 for arg in expr.args[:-1]) + "{}{}".format(self._print(expr.args[-1]),

128 ')' * (len(expr.args) - 1))

129

130 def _print_Adjoint(self, expr):

131 return '{}({}({}))'.format(

132 self._module_format(self._module + '.conjugate'),

133 self._module_format(self._module + '.transpose'),

134 self._print(expr.args[0]))

135

136 def _print_DiagonalOf(self, expr):

137 vect = '{}({})'.format(

138 self._module_format(self._module + '.diag'),

139 self._print(expr.arg))

140 return '{}({}, (-1, 1))'.format(

141 self._module_format(self._module + '.reshape'), vect)

142

143 def _print_DiagMatrix(self, expr):

144 return '{}({})'.format(self._module_format(self._module + '.diagflat'),

145 self._print(expr.args[0]))

146

147 def _print_DiagonalMatrix(self, expr):

148 return '{}({}, {}({}, {}))'.format(self._module_format(self._module + '.multiply'),

149 self._print(expr.arg), self._module_format(self._module + '.eye'),

150 self._print(expr.shape[0]), self._print(expr.shape[1]))

151

152 def _print_Piecewise(self, expr):

153 "Piecewise function printer"

154 from sympy.logic.boolalg import ITE, simplify_logic

155 def print_cond(cond):

156 """ Problem having an ITE in the cond. """

157 if cond.has(ITE):

158 return self._print(simplify_logic(cond))

159 else:

160 return self._print(cond)

161 exprs = '[{}]'.format(','.join(self._print(arg.expr) for arg in expr.args))

162 conds = '[{}]'.format(','.join(print_cond(arg.cond) for arg in expr.args))

163 # If [default_value, True] is a (expr, cond) sequence in a Piecewise object

164 # it will behave the same as passing the 'default' kwarg to select()

165 # *as long as* it is the last element in expr.args.

166 # If this is not the case, it may be triggered prematurely.

167 return '{}({}, {}, default={})'.format(

168 self._module_format(self._module + '.select'), conds, exprs,

169 self._print(S.NaN))

170

171 def _print_Relational(self, expr):

172 "Relational printer for Equality and Unequality"

173 op = {

174 '==' :'equal',

175 '!=' :'not_equal',

176 '<' :'less',

177 '<=' :'less_equal',

178 '>' :'greater',

179 '>=' :'greater_equal',

180 }

181 if expr.rel_op in op:

182 lhs = self._print(expr.lhs)

183 rhs = self._print(expr.rhs)

184 return '{op}({lhs}, {rhs})'.format(op=self._module_format(self._module + '.'+op[expr.rel_op]),

185 lhs=lhs, rhs=rhs)

186 return super()._print_Relational(expr)

187

188 def _print_And(self, expr):

189 "Logical And printer"

190 # We have to override LambdaPrinter because it uses Python 'and' keyword.

191 # If LambdaPrinter didn't define it, we could use StrPrinter's

192 # version of the function and add 'logical_and' to NUMPY_TRANSLATIONS.

193 return '{}.reduce(({}))'.format(self._module_format(self._module + '.logical_and'), ','.join(self._print(i) for i in expr.args))

194

195 def _print_Or(self, expr):

196 "Logical Or printer"

197 # We have to override LambdaPrinter because it uses Python 'or' keyword.

198 # If LambdaPrinter didn't define it, we could use StrPrinter's

199 # version of the function and add 'logical_or' to NUMPY_TRANSLATIONS.

200 return '{}.reduce(({}))'.format(self._module_format(self._module + '.logical_or'), ','.join(self._print(i) for i in expr.args))

201

202 def _print_Not(self, expr):

203 "Logical Not printer"

204 # We have to override LambdaPrinter because it uses Python 'not' keyword.

205 # If LambdaPrinter didn't define it, we would still have to define our

206 # own because StrPrinter doesn't define it.

207 return '{}({})'.format(self._module_format(self._module + '.logical_not'), ','.join(self._print(i) for i in expr.args))

208

209 def _print_Pow(self, expr, rational=False):

210 # XXX Workaround for negative integer power error

211 if expr.exp.is_integer and expr.exp.is_negative:

212 expr = Pow(expr.base, expr.exp.evalf(), evaluate=False)

213 return self._hprint_Pow(expr, rational=rational, sqrt=self._module + '.sqrt')

214

215 def _print_Min(self, expr):

216 return '{}(({}), axis=0)'.format(self._module_format(self._module + '.amin'), ','.join(self._print(i) for i in expr.args))

217

218 def _print_Max(self, expr):

219 return '{}(({}), axis=0)'.format(self._module_format(self._module + '.amax'), ','.join(self._print(i) for i in expr.args))

220

221 def _print_arg(self, expr):

222 return "%s(%s)" % (self._module_format(self._module + '.angle'), self._print(expr.args[0]))

223

224 def _print_im(self, expr):

225 return "%s(%s)" % (self._module_format(self._module + '.imag'), self._print(expr.args[0]))

226

227 def _print_Mod(self, expr):

228 return "%s(%s)" % (self._module_format(self._module + '.mod'), ', '.join(

229 (self._print(arg) for arg in expr.args)))

230

231 def _print_re(self, expr):

232 return "%s(%s)" % (self._module_format(self._module + '.real'), self._print(expr.args[0]))

233

234 def _print_sinc(self, expr):

235 return "%s(%s)" % (self._module_format(self._module + '.sinc'), self._print(expr.args[0]/S.Pi))

236

237 def _print_MatrixBase(self, expr):

238 func = self.known_functions.get(expr.__class__.__name__, None)

239 if func is None:

240 func = self._module_format(self._module + '.array')

241 return "%s(%s)" % (func, self._print(expr.tolist()))

242

243 def _print_Identity(self, expr):

244 shape = expr.shape

245 if all(dim.is_Integer for dim in shape):

246 return "%s(%s)" % (self._module_format(self._module + '.eye'), self._print(expr.shape[0]))

247 else:

248 raise NotImplementedError("Symbolic matrix dimensions are not yet supported for identity matrices")

249

250 def _print_BlockMatrix(self, expr):

251 return '{}({})'.format(self._module_format(self._module + '.block'),

252 self._print(expr.args[0].tolist()))

253

254 def _print_NDimArray(self, expr):

255 if len(expr.shape) == 1:

256 return self._module + '.array(' + self._print(expr.args[0]) + ')'

257 if len(expr.shape) == 2:

258 return self._print(expr.tomatrix())

259 # Should be possible to extend to more dimensions

260 return CodePrinter._print_not_supported(self, expr)

261

262 _add = "add"

263 _einsum = "einsum"

264 _transpose = "transpose"

265 _ones = "ones"

266 _zeros = "zeros"

267

268 _print_lowergamma = CodePrinter._print_not_supported

269 _print_uppergamma = CodePrinter._print_not_supported

270 _print_fresnelc = CodePrinter._print_not_supported

271 _print_fresnels = CodePrinter._print_not_supported

272

273for func in _numpy_known_functions:

274 setattr(NumPyPrinter, f'_print_{func}', _print_known_func)

275

276for const in _numpy_known_constants:

277 setattr(NumPyPrinter, f'_print_{const}', _print_known_const)

278

279

280_known_functions_scipy_special = {

281 'Ei': 'expi',

282 'erf': 'erf',

283 'erfc': 'erfc',

284 'besselj': 'jv',

285 'bessely': 'yv',

286 'besseli': 'iv',

287 'besselk': 'kv',

288 'cosm1': 'cosm1',

289 'powm1': 'powm1',

290 'factorial': 'factorial',

291 'gamma': 'gamma',

292 'loggamma': 'gammaln',

293 'digamma': 'psi',

294 'polygamma': 'polygamma',

295 'RisingFactorial': 'poch',

296 'jacobi': 'eval_jacobi',

297 'gegenbauer': 'eval_gegenbauer',

298 'chebyshevt': 'eval_chebyt',

299 'chebyshevu': 'eval_chebyu',

300 'legendre': 'eval_legendre',

301 'hermite': 'eval_hermite',

302 'laguerre': 'eval_laguerre',

303 'assoc_laguerre': 'eval_genlaguerre',

304 'beta': 'beta',

305 'LambertW' : 'lambertw',

306}

307

308_known_constants_scipy_constants = {

309 'GoldenRatio': 'golden_ratio',

310 'Pi': 'pi',

311}

312_scipy_known_functions = {k : "scipy.special." + v for k, v in _known_functions_scipy_special.items()}

313_scipy_known_constants = {k : "scipy.constants." + v for k, v in _known_constants_scipy_constants.items()}

314

315class SciPyPrinter(NumPyPrinter):

316

317 _kf = {**NumPyPrinter._kf, **_scipy_known_functions}

318 _kc = {**NumPyPrinter._kc, **_scipy_known_constants}

319

320 def __init__(self, settings=None):

321 super().__init__(settings=settings)

322 self.language = "Python with SciPy and NumPy"

323

324 def _print_SparseRepMatrix(self, expr):

325 i, j, data = [], [], []

326 for (r, c), v in expr.todok().items():

327 i.append(r)

328 j.append(c)

329 data.append(v)

330

331 return "{name}(({data}, ({i}, {j})), shape={shape})".format(

332 name=self._module_format('scipy.sparse.coo_matrix'),

333 data=data, i=i, j=j, shape=expr.shape

334 )

335

336 _print_ImmutableSparseMatrix = _print_SparseRepMatrix

337

338 # SciPy's lpmv has a different order of arguments from assoc_legendre

339 def _print_assoc_legendre(self, expr):

340 return "{0}({2}, {1}, {3})".format(

341 self._module_format('scipy.special.lpmv'),

342 self._print(expr.args[0]),

343 self._print(expr.args[1]),

344 self._print(expr.args[2]))

345

346 def _print_lowergamma(self, expr):

347 return "{0}({2})*{1}({2}, {3})".format(

348 self._module_format('scipy.special.gamma'),

349 self._module_format('scipy.special.gammainc'),

350 self._print(expr.args[0]),

351 self._print(expr.args[1]))

352

353 def _print_uppergamma(self, expr):

354 return "{0}({2})*{1}({2}, {3})".format(

355 self._module_format('scipy.special.gamma'),

356 self._module_format('scipy.special.gammaincc'),

357 self._print(expr.args[0]),

358 self._print(expr.args[1]))

359

360 def _print_betainc(self, expr):

361 betainc = self._module_format('scipy.special.betainc')

362 beta = self._module_format('scipy.special.beta')

363 args = [self._print(arg) for arg in expr.args]

364 return f"({betainc}({args[0]}, {args[1]}, {args[3]}) - {betainc}({args[0]}, {args[1]}, {args[2]})) \

365 * {beta}({args[0]}, {args[1]})"

366

367 def _print_betainc_regularized(self, expr):

368 return "{0}({1}, {2}, {4}) - {0}({1}, {2}, {3})".format(

369 self._module_format('scipy.special.betainc'),

370 self._print(expr.args[0]),

371 self._print(expr.args[1]),

372 self._print(expr.args[2]),

373 self._print(expr.args[3]))

374

375 def _print_fresnels(self, expr):

376 return "{}({})[0]".format(

377 self._module_format("scipy.special.fresnel"),

378 self._print(expr.args[0]))

379

380 def _print_fresnelc(self, expr):

381 return "{}({})[1]".format(

382 self._module_format("scipy.special.fresnel"),

383 self._print(expr.args[0]))

384

385 def _print_airyai(self, expr):

386 return "{}({})[0]".format(

387 self._module_format("scipy.special.airy"),

388 self._print(expr.args[0]))

389

390 def _print_airyaiprime(self, expr):

391 return "{}({})[1]".format(

392 self._module_format("scipy.special.airy"),

393 self._print(expr.args[0]))

394

395 def _print_airybi(self, expr):

396 return "{}({})[2]".format(

397 self._module_format("scipy.special.airy"),

398 self._print(expr.args[0]))

399

400 def _print_airybiprime(self, expr):

401 return "{}({})[3]".format(

402 self._module_format("scipy.special.airy"),

403 self._print(expr.args[0]))

404

405 def _print_bernoulli(self, expr):

406 # scipy's bernoulli is inconsistent with SymPy's so rewrite

407 return self._print(expr._eval_rewrite_as_zeta(*expr.args))

408

409 def _print_harmonic(self, expr):

410 return self._print(expr._eval_rewrite_as_zeta(*expr.args))

411

412 def _print_Integral(self, e):

413 integration_vars, limits = _unpack_integral_limits(e)

414

415 if len(limits) == 1:

416 # nicer (but not necessary) to prefer quad over nquad for 1D case

417 module_str = self._module_format("scipy.integrate.quad")

418 limit_str = "%s, %s" % tuple(map(self._print, limits[0]))

419 else:

420 module_str = self._module_format("scipy.integrate.nquad")

421 limit_str = "({})".format(", ".join(

422 "(%s, %s)" % tuple(map(self._print, l)) for l in limits))

423

424 return "{}(lambda {}: {}, {})[0]".format(

425 module_str,

426 ", ".join(map(self._print, integration_vars)),

427 self._print(e.args[0]),

428 limit_str)

429

430 def _print_Si(self, expr):

431 return "{}({})[0]".format(

432 self._module_format("scipy.special.sici"),

433 self._print(expr.args[0]))

434

435 def _print_Ci(self, expr):

436 return "{}({})[1]".format(

437 self._module_format("scipy.special.sici"),

438 self._print(expr.args[0]))

439

440for func in _scipy_known_functions:

441 setattr(SciPyPrinter, f'_print_{func}', _print_known_func)

442

443for const in _scipy_known_constants:

444 setattr(SciPyPrinter, f'_print_{const}', _print_known_const)

445

446

447_cupy_known_functions = {k : "cupy." + v for k, v in _known_functions_numpy.items()}

448_cupy_known_constants = {k : "cupy." + v for k, v in _known_constants_numpy.items()}

449

450class CuPyPrinter(NumPyPrinter):

451 """

452 CuPy printer which handles vectorized piecewise functions,

453 logical operators, etc.

454 """

455

456 _module = 'cupy'

457 _kf = _cupy_known_functions

458 _kc = _cupy_known_constants

459

460 def __init__(self, settings=None):

461 super().__init__(settings=settings)

462

463for func in _cupy_known_functions:

464 setattr(CuPyPrinter, f'_print_{func}', _print_known_func)

465

466for const in _cupy_known_constants:

467 setattr(CuPyPrinter, f'_print_{const}', _print_known_const)

468

469

470_jax_known_functions = {k: 'jax.numpy.' + v for k, v in _known_functions_numpy.items()}

471_jax_known_constants = {k: 'jax.numpy.' + v for k, v in _known_constants_numpy.items()}

472

473class JaxPrinter(NumPyPrinter):

474 """

475 JAX printer which handles vectorized piecewise functions,

476 logical operators, etc.

477 """

478 _module = "jax.numpy"

479

480 _kf = _jax_known_functions

481 _kc = _jax_known_constants

482

483 def __init__(self, settings=None):

484 super().__init__(settings=settings)

485

486 # These need specific override to allow for the lack of "jax.numpy.reduce"

487 def _print_And(self, expr):

488 "Logical And printer"

489 return "{}({}.asarray([{}]), axis=0)".format(

490 self._module_format(self._module + ".all"),

491 self._module_format(self._module),

492 ",".join(self._print(i) for i in expr.args),

493 )

494

495 def _print_Or(self, expr):

496 "Logical Or printer"

497 return "{}({}.asarray([{}]), axis=0)".format(

498 self._module_format(self._module + ".any"),

499 self._module_format(self._module),

500 ",".join(self._print(i) for i in expr.args),

501 )

502

503for func in _jax_known_functions:

504 setattr(JaxPrinter, f'_print_{func}', _print_known_func)

505

506for const in _jax_known_constants:

507 setattr(JaxPrinter, f'_print_{const}', _print_known_const)