Ticket #11888: trac_11888_v2.patch
File trac_11888_v2.patch, 8.8 KB (added by , 8 years ago) 


sage/functions/all.py
# HG changeset patch # User Benjamin Jones <benjaminfjones@gmail.com> # Date 1328501390 21600 # Node ID f4852bbf918bd1385bbca27e644f987c155bddcc # Parent c239be1054e01526a1b0b62da6691061b9dd5587 Trac 11888: add the lambert_w and lambert_w_branch symbolic function diff git a/sage/functions/all.py b/sage/functions/all.py
a b 21 21 real_part, real, 22 22 imag_part, imag, imaginary, conjugate) 23 23 24 from log import (exp, log, ln, polylog, dilog )24 from log import (exp, log, ln, polylog, dilog, lambert_w, lambert_w_branch) 25 25 26 26 27 27 from transcendental import (exponential_integral_1, 
sage/functions/log.py
diff git a/sage/functions/log.py b/sage/functions/log.py
a b 1 1 """ 2 2 Logarithmic functions 3 3 """ 4 from sage.symbolic.function import GinacFunction 4 from sage.symbolic.function import GinacFunction, BuiltinFunction, is_inexact 5 6 from sage.libs.mpmath import utils as mpmath_utils 7 from sage.structure.coerce import parent 8 from sage.symbolic.expression import Expression 9 from sage.rings.real_double import RDF 10 from sage.rings.complex_double import CDF 5 11 6 12 class Function_exp(GinacFunction): 7 13 def __init__(self): … … 435 441 436 442 dilog = Function_dilog() 437 443 444 445 class Function_lambert_w_branch(BuiltinFunction): 446 r""" 447 The integral branches of the Lambert W function `W_n(z)`. 448 449 This function satisfies the equation 450 451 .. math:: 452 453 z = W_n(z) e^{W_n(z)} 454 455 INPUT: 456 457  ``z``  a complex number 458 459  ``n``  an integer. `n=0` corresponds to the principle branch. 460 461 ALGORITHM: 462 463 Numerical evaluation is handled using the mpmath library. 464 465 REFERENCES: 466 467  http://en.wikipedia.org/wiki/Lambert_W_function 468 469 EXAMPLES:: 470 471 sage: lambert_w(1.0) 472 0.567143290409784 473 sage: lambert_w(1).n() 474 0.318131505204764 + 1.33723570143069*I 475 sage: lambert_w(1.5 + 5*I) 476 1.17418016254171 + 1.10651494102011*I 477 478 sage: lambert_w(RealField(100)(1)) 479 0.56714329040978387299996866221 480 481 sage: S = solve(e^(5*x)+x==0, x, to_poly_solve=True) 482 sage: z = S[0].rhs(); z 483 1/5*lambert_w(5) 484 sage: N(z) 485 0.265344933048440 486 487 Check the defining equation numerically at `z=5`:: 488 489 sage: N(lambert_w(5)*exp(lambert_w(5))  5) 490 0.000000000000000 491 """ 492 493 def __init__(self): 494 """ 495 See the docstring for :meth:`Function_lambert_w`. 496 497 EXAMPLES:: 498 499 sage: lambert_w(1.0) 500 0.567143290409784 501 502 """ 503 BuiltinFunction.__init__(self, "lambert_w_branch", nargs=2, latex_name=r'W') 504 505 def _eval_(self, z, n): 506 """ 507 EXAMPLES:: 508 509 sage: lambert_w(0) 510 lambert_w(0) 511 sage: x = var('x') 512 sage: lambert_w(x) 513 lambert_w(x) 514 sage: lambert_w(0.0) 515 0.000000000000000 516 517 """ 518 if not isinstance(z, Expression) and is_inexact(z): 519 return self._evalf_(z, n, parent=parent(z)) 520 521 return None 522 523 def _evalf_(self, z, n, parent=None): 524 """ 525 EXAMPLES:: 526 527 sage: N(lambert_w(1)) 528 0.567143290409784 529 sage: lambert_w(RealField(100)(1)) 530 0.56714329040978387299996866221 531 532 """ 533 534 if parent(z) == RDF or parent(z) == CDF: 535 import scipy.special 536 return scipy.special.lambertw(z, n) 537 else: 538 import mpmath 539 return mpmath_utils.call(mpmath.lambertw, z, n, parent=parent) 540 541 def _derivative_(self, z, n, diff_param=None): 542 """ 543 The derivative of `W_n(x)` is `W_n(x)/(x \cdot W_n(x) + x)`. 544 545 EXAMPLES:: 546 547 sage: x = var('x') 548 sage: derivative(lambert_w(x), x) 549 lambert_w(x)/(x*lambert_w(x) + x) 550 551 """ 552 return lambert_w(z, n)/(z*lambert_w(z, n)+z) 553 554 lambert_w_branch = Function_lambert_w_branch() 555 556 def lambert_w(z, *args, **kwds): 557 """ 558 The Lambert w function. This is a wrapper for lambert_w_branch. 559 """ 560 if not args: 561 return lambert_w_branch(z, 0, **kwds) 562 elif len(args) > 1: 563 raise ArgumentError("lambert_w takes at most two arguments.") 564 else: 565 return lambert_w_branch(z, args[0], **kwds) 
sage/symbolic/random_tests.py
diff git a/sage/symbolic/random_tests.py b/sage/symbolic/random_tests.py
a b 14 14 15 15 sage: from sage.symbolic.random_tests import _mk_full_functions 16 16 sage: [f for (one,f,arity) in _mk_full_functions()] 17 [Ei, abs, arccos, arccosh, arccot, arccoth, arccsc, arccsch, 18 arcsec , arcsech, arcsin, arcsinh, arctan, arctan2, arctanh,19 binomial, ceil, conjugate, cos, cosh, cot, coth, csc, csch,20 di ckman_rho, dilog, dirac_delta, elliptic_e, elliptic_ec,21 elliptic_ eu, elliptic_f, elliptic_kc, elliptic_pi, erf, exp,22 factorial, floor, heaviside, imag_part, integrate,23 kronecker_delta, log, polylog, real_part, sec, sech, sgn, sin,24 sinh, tan, tanh, unit_step, zeta,zetaderiv]17 [Ei, abs, arccos, arccosh, arccot, arccoth, arccsc, arccsch, arcsec, 18 arcsech, arcsin, arcsinh, arctan, arctan2, arctanh, binomial, ceil, 19 conjugate, cos, cosh, cot, coth, csc, csch, dickman_rho, dilog, 20 dirac_delta, elliptic_e, elliptic_ec, elliptic_eu, elliptic_f, 21 elliptic_kc, elliptic_pi, erf, exp, factorial, floor, heaviside, 22 imag_part, integrate, kronecker_delta, lambert_w, log, polylog, 23 real_part, sec, sech, sgn, sin, sinh, tan, tanh, unit_step, zeta, 24 zetaderiv] 25 25 26 26 Note that this doctest will fail whenever a Pynac function is added or 27 27 removed. In that case, it is very likely that the doctests for … … 234 234 sage: from sage.symbolic.random_tests import * 235 235 sage: set_random_seed(2) 236 236 sage: random_expr(50, nvars=3, coeff_generator=CDF.random_element) 237 (euler_gamma  v3^(e) + (v2  factorial(e/v2))^(((2.85879036573  1.18163393202*I)*v2 + (2.85879036573  1.18163393202*I)*v3)*pi  0.247786879678 + 0.931826724898*I)*arccsc((0.891138386848  0.0936820840629*I)/v1)  (0.553423153995  0.5481180572*I)*v3 + 0.149683576515  0.155746451854*I)*v1 + arccsch(pi + e)*elliptic_f(khinchin*v2, 1.4656989704 + 0.863754357069*I) 237 v1*elliptic_kc((0.0666829501658 + 0.206976992303*I)/(v3 + e))/v3 + 238 heaviside(arccosh((abs(v2  239 floor(v3))^(real_part(elliptic_e(0.703991792631 + 0.750156228797*I, 240 (1.21734510331  1.22580558833*I)*pi*v1)))*e^(arctan2(imag_part(v2)  241 imag_part(floor(v3)), real_part(v2)  242 real_part(floor(v3)))*imag_part(elliptic_e(0.703991792631 + 243 0.750156228797*I, (1.21734510331  244 1.22580558833*I)*pi*v1)))*sin(log(abs(v2  245 floor(v3)))*imag_part(elliptic_e(0.703991792631 + 0.750156228797*I, 246 (1.21734510331  1.22580558833*I)*pi*v1))  arctan2(imag_part(v2)  247 imag_part(floor(v3)), real_part(v2)  248 real_part(floor(v3)))*real_part(elliptic_e(0.703991792631 + 249 0.750156228797*I, (1.21734510331  250 1.22580558833*I)*pi*v1)))*imag_part(arccsc((0.891138386848  251 0.0936820840629*I)/v1))  abs(v2  252 floor(v3))^(real_part(elliptic_e(0.703991792631 + 0.750156228797*I, 253 (1.21734510331  1.22580558833*I)*pi*v1)))*e^(arctan2(imag_part(v2)  254 imag_part(floor(v3)), real_part(v2)  255 real_part(floor(v3)))*imag_part(elliptic_e(0.703991792631 + 256 0.750156228797*I, (1.21734510331  257 1.22580558833*I)*pi*v1)))*cos(log(abs(v2  258 floor(v3)))*imag_part(elliptic_e(0.703991792631 + 0.750156228797*I, 259 (1.21734510331  1.22580558833*I)*pi*v1))  arctan2(imag_part(v2)  260 imag_part(floor(v3)), real_part(v2)  261 real_part(floor(v3)))*real_part(elliptic_e(0.703991792631 + 262 0.750156228797*I, (1.21734510331  263 1.22580558833*I)*pi*v1)))*real_part(arccsc((0.891138386848  264 0.0936820840629*I)/v1)) + v3^(0.48519994364  0.485764091302*I) + 265 0.0909404921682*real_part(v3)/(real_part(v3)^2 + imag_part(v3)^2) + 266 0.281538203756*imag_part(v3)/(real_part(v3)^2 + imag_part(v3)^2)  267 0.647983235144*real_part(v1)  1.20665952957*imag_part(v1))*v1)) 238 268 sage: random_expr(5, verbose=True) 239 About to apply dirac_delta to [1] 240 About to apply arccsch to [0] 241 About to apply <builtin function add> to [0, arccsch(0)] 242 arccsch(0) 269 About to apply <builtin function mul> to [v1, e] 270 About to apply <builtin function div> to [v1*e, 1] 271 v1*e 243 272 """ 244 273 vars = [(1.0, sage.calculus.calculus.var('v%d' % (n+1))) for n in range(nvars)] 245 274 if ncoeffs is None: