| 447 | |
| 448 | class Function_lambert_w(BuiltinFunction): |
| 449 | r""" |
| 450 | The integral branches of the Lambert W function `W_n(z)`. |
| 451 | |
| 452 | This function satisfies the equation |
| 453 | |
| 454 | .. math:: |
| 455 | |
| 456 | z = W_n(z) e^{W_n(z)} |
| 457 | |
| 458 | INPUT: |
| 459 | |
| 460 | - ``n`` - an integer. `n=0` corresponds to the principal branch. |
| 461 | |
| 462 | - ``z`` - a complex number |
| 463 | |
| 464 | If called with a single argument, that argument is ``z`` and the branch ``n`` is |
| 465 | assumed to be 0 (the principal branch). |
| 466 | |
| 467 | ALGORITHM: |
| 468 | |
| 469 | Numerical evaluation is handled using the mpmath and SciPy libraries. |
| 470 | |
| 471 | REFERENCES: |
| 472 | |
| 473 | - http://en.wikipedia.org/wiki/Lambert_W_function |
| 474 | |
| 475 | EXAMPLES: |
| 476 | |
| 477 | Evaluation of the principal branch:: |
| 478 | |
| 479 | sage: lambert_w(1.0) |
| 480 | 0.567143290409784 |
| 481 | sage: lambert_w(-1).n() |
| 482 | -0.318131505204764 + 1.33723570143069*I |
| 483 | sage: lambert_w(-1.5 + 5*I) |
| 484 | 1.17418016254171 + 1.10651494102011*I |
| 485 | |
| 486 | Evaluation of other branches:: |
| 487 | |
| 488 | sage: lambert_w(2, 1.0) |
| 489 | -2.40158510486800 + 10.7762995161151*I |
| 490 | |
| 491 | Solutions to certain exponential equations are returned in terms of lambert_w:: |
| 492 | |
| 493 | sage: S = solve(e^(5*x)+x==0, x, to_poly_solve=True) |
| 494 | sage: z = S[0].rhs(); z |
| 495 | -1/5*lambert_w(5) |
| 496 | sage: N(z) |
| 497 | -0.265344933048440 |
| 498 | |
| 499 | Check the defining equation numerically at `z=5`:: |
| 500 | |
| 501 | sage: N(lambert_w(5)*exp(lambert_w(5)) - 5) |
| 502 | 0.000000000000000 |
| 503 | |
| 504 | There are several special values of the principal branch which |
| 505 | are automatically simplified:: |
| 506 | |
| 507 | sage: lambert_w(0) |
| 508 | 0 |
| 509 | sage: lambert_w(e) |
| 510 | 1 |
| 511 | sage: lambert_w(-1/e) |
| 512 | -1 |
| 513 | """ |
| 514 | |
| 515 | def __init__(self): |
| 516 | r""" |
| 517 | See the docstring for :meth:`Function_lambert_w`. |
| 518 | |
| 519 | EXAMPLES:: |
| 520 | |
| 521 | sage: lambert_w(0, 1.0) |
| 522 | 0.567143290409784 |
| 523 | """ |
| 524 | BuiltinFunction.__init__(self, "lambert_w", nargs=2, |
| 525 | conversions={'maxima':'lambert_w', |
| 526 | 'mathematica':'ProductLog', |
| 527 | 'maple':'LambertW'}) |
| 528 | self._simplifications = {Integer(0):Integer(0), const_e:Integer(1), |
| 529 | -1/const_e:Integer(-1)} |
| 530 | |
| 531 | def __call__(self, *args, **kwds): |
| 532 | r""" |
| 533 | Custom call method allows the user to pass one argument or two. If |
| 534 | one argument is passed, we assume it is ``z`` and that ``n=0``. |
| 535 | |
| 536 | EXAMPLES:: |
| 537 | |
| 538 | sage: lambert_w(1) |
| 539 | lambert_w(1) |
| 540 | sage: lambert_w(1, 2) |
| 541 | lambert_w(1, 2) |
| 542 | """ |
| 543 | if len(args) == 2: |
| 544 | return BuiltinFunction.__call__(self, *args, **kwds) |
| 545 | elif len(args) == 1: |
| 546 | return BuiltinFunction.__call__(self, 0, args[0], **kwds) |
| 547 | else: |
| 548 | raise TypeError("lambert_w takes either one or two arguments.") |
| 549 | |
| 550 | def _eval_(self, n, z): |
| 551 | """ |
| 552 | EXAMPLES:: |
| 553 | |
| 554 | sage: lambert_w(6.0) |
| 555 | 1.43240477589830 |
| 556 | sage: lambert_w(1) |
| 557 | lambert_w(1) |
| 558 | sage: lambert_w(x+1) |
| 559 | lambert_w(x + 1) |
| 560 | |
| 561 | There are three special values which are automatically simplified:: |
| 562 | |
| 563 | sage: lambert_w(0) |
| 564 | 0 |
| 565 | sage: lambert_w(e) |
| 566 | 1 |
| 567 | sage: lambert_w(-1/e) |
| 568 | -1 |
| 569 | sage: lambert_w(SR(-1/e)) |
| 570 | -1 |
| 571 | sage: lambert_w(SR(0)) |
| 572 | 0 |
| 573 | |
| 574 | The special values only hold on the principal branch:: |
| 575 | |
| 576 | sage: lambert_w(1,e) |
| 577 | lambert_w(1, e) |
| 578 | sage: lambert_w(1, e.n()) |
| 579 | -0.532092121986380 + 4.59715801330257*I |
| 580 | """ |
| 581 | if not isinstance(z, Expression): |
| 582 | if is_inexact(z): |
| 583 | return self._evalf_(n, z, parent=sage_structure_coerce_parent(z)) |
| 584 | elif z == Integer(0): |
| 585 | return Integer(0) |
| 586 | elif n == 0: |
| 587 | if z.is_trivial_zero(): |
| 588 | return Integer(0) |
| 589 | elif (z-const_e).is_trivial_zero(): |
| 590 | return Integer(1) |
| 591 | elif (z+1/const_e).is_trivial_zero(): |
| 592 | return Integer(-1) |
| 593 | return None |
| 594 | |
| 595 | def _evalf_(self, n, z, parent=None): |
| 596 | """ |
| 597 | EXAMPLES:: |
| 598 | |
| 599 | sage: N(lambert_w(1)) |
| 600 | 0.567143290409784 |
| 601 | sage: lambert_w(RealField(100)(1)) |
| 602 | 0.56714329040978387299996866221 |
| 603 | |
| 604 | SciPy is used to evaluate for float, RDF, and CDF inputs:: |
| 605 | |
| 606 | sage: lambert_w(RDF(1)) |
| 607 | 0.56714329041 |
| 608 | """ |
| 609 | R = parent or sage_structure_coerce_parent(z) |
| 610 | if R is float or R is complex or R is RDF or R is CDF: |
| 611 | import scipy.special |
| 612 | return scipy.special.lambertw(z, n) |
| 613 | else: |
| 614 | import mpmath |
| 615 | return mpmath_utils.call(mpmath.lambertw, z, n, parent=R) |
| 616 | |
| 617 | def _derivative_(self, n, z, diff_param=None): |
| 618 | """ |
| 619 | The derivative of `W_n(x)` is `W_n(x)/(x \cdot W_n(x) + x)`. |
| 620 | |
| 621 | EXAMPLES:: |
| 622 | |
| 623 | sage: x = var('x') |
| 624 | sage: derivative(lambert_w(x), x) |
| 625 | lambert_w(x)/(x*lambert_w(x) + x) |
| 626 | """ |
| 627 | return lambert_w(n, z)/(z*lambert_w(n, z)+z) |
| 628 | |
| 629 | def _print_(self, n, z): |
| 630 | """ |
| 631 | Custom _print_ method to avoid printing the branch number if |
| 632 | it is zero. |
| 633 | |
| 634 | EXAMPLES:: |
| 635 | |
| 636 | sage: lambert_w(1) |
| 637 | lambert_w(1) |
| 638 | sage: lambert_w(0,x) |
| 639 | lambert_w(x) |
| 640 | """ |
| 641 | if n == 0: |
| 642 | return "lambert_w(%s)" % z |
| 643 | else: |
| 644 | return "lambert_w(%s, %s)" % (n,z) |
| 645 | |
| 646 | def _print_latex_(self, n, z): |
| 647 | """ |
| 648 | Custom _print_latex_ method to avoid printing the branch |
| 649 | number if it is zero. |
| 650 | |
| 651 | EXAMPLES:: |
| 652 | |
| 653 | sage: latex(lambert_w(1)) |
| 654 | \operatorname{W_0}(1) |
| 655 | sage: latex(lambert_w(0,x)) |
| 656 | \operatorname{W_0}(x) |
| 657 | sage: latex(lambert_w(1,x)) |
| 658 | \operatorname{W_{1}}(x) |
| 659 | """ |
| 660 | if n == 0: |
| 661 | return r"\operatorname{W_0}(%s)" % z |
| 662 | else: |
| 663 | return r"\operatorname{W_{%s}}(%s)" % (n,z) |
| 664 | |
| 665 | lambert_w = Function_lambert_w() |