| 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 | Integration (of the principal branch) is evaluated using Maxima:: |
| 515 | |
| 516 | sage: integrate(lambert_w(x), x) |
| 517 | (lambert_w(x)^2 - lambert_w(x) + 1)*x/lambert_w(x) |
| 518 | sage: integrate(lambert_w(x), x, 0, 1) |
| 519 | (lambert_w(1)^2 - lambert_w(1) + 1)/lambert_w(1) - 1 |
| 520 | sage: integrate(lambert_w(x), x, 0, 1.0) |
| 521 | 0.330366124762 |
| 522 | |
| 523 | Warning: The integral of a non-principal branch is not implemented, and |
| 524 | neither is numerical integration using GSL. The :meth:`numerical_integral` |
| 525 | function does work if you pass a lambda function:: |
| 526 | |
| 527 | sage: numerical_integral(lambda x: lambert_w(x), 0, 1) |
| 528 | (0.33036612476168054, 3.667800782666048e-15) |
| 529 | """ |
| 530 | |
| 531 | def __init__(self): |
| 532 | r""" |
| 533 | See the docstring for :meth:`Function_lambert_w`. |
| 534 | |
| 535 | EXAMPLES:: |
| 536 | |
| 537 | sage: lambert_w(0, 1.0) |
| 538 | 0.567143290409784 |
| 539 | """ |
| 540 | BuiltinFunction.__init__(self, "lambert_w", nargs=2, |
| 541 | conversions={'mathematica':'ProductLog', |
| 542 | 'maple':'LambertW'}) |
| 543 | |
| 544 | def __call__(self, *args, **kwds): |
| 545 | r""" |
| 546 | Custom call method allows the user to pass one argument or two. If |
| 547 | one argument is passed, we assume it is ``z`` and that ``n=0``. |
| 548 | |
| 549 | EXAMPLES:: |
| 550 | |
| 551 | sage: lambert_w(1) |
| 552 | lambert_w(1) |
| 553 | sage: lambert_w(1, 2) |
| 554 | lambert_w(1, 2) |
| 555 | """ |
| 556 | if len(args) == 2: |
| 557 | return BuiltinFunction.__call__(self, *args, **kwds) |
| 558 | elif len(args) == 1: |
| 559 | return BuiltinFunction.__call__(self, 0, args[0], **kwds) |
| 560 | else: |
| 561 | raise TypeError("lambert_w takes either one or two arguments") |
| 562 | |
| 563 | def _eval_(self, n, z): |
| 564 | """ |
| 565 | EXAMPLES:: |
| 566 | |
| 567 | sage: lambert_w(6.0) |
| 568 | 1.43240477589830 |
| 569 | sage: lambert_w(1) |
| 570 | lambert_w(1) |
| 571 | sage: lambert_w(x+1) |
| 572 | lambert_w(x + 1) |
| 573 | |
| 574 | There are several special values which are automatically simplified:: |
| 575 | |
| 576 | sage: lambert_w(0) |
| 577 | 0 |
| 578 | sage: lambert_w(e) |
| 579 | 1 |
| 580 | sage: lambert_w(-1/e) |
| 581 | -1 |
| 582 | sage: lambert_w(SR(0)) |
| 583 | 0 |
| 584 | sage: lambert_w(-pi/2) |
| 585 | 1/2*I*pi |
| 586 | |
| 587 | The special values only hold on the principal branch:: |
| 588 | |
| 589 | sage: lambert_w(1,e) |
| 590 | lambert_w(1, e) |
| 591 | sage: lambert_w(1, e.n()) |
| 592 | -0.532092121986380 + 4.59715801330257*I |
| 593 | |
| 594 | TESTS: |
| 595 | |
| 596 | When automatic simplification occurs, the parent of the output |
| 597 | value should be either the same as the parent of the input, or |
| 598 | a Sage type:: |
| 599 | |
| 600 | sage: parent(lambert_w(int(0))) |
| 601 | <type 'int'> |
| 602 | sage: parent(lambert_w(Integer(0))) |
| 603 | Integer Ring |
| 604 | sage: parent(lambert_w(e)) |
| 605 | Integer Ring |
| 606 | """ |
| 607 | if not isinstance(z, Expression): |
| 608 | if is_inexact(z): |
| 609 | return self._evalf_(n, z, parent=sage_structure_coerce_parent(z)) |
| 610 | elif n == 0 and z == 0: |
| 611 | return sage_structure_coerce_parent(z)(Integer(0)) |
| 612 | elif n == 0: |
| 613 | if z.is_trivial_zero(): |
| 614 | return sage_structure_coerce_parent(z)(Integer(0)) |
| 615 | elif (z-const_e).is_trivial_zero(): |
| 616 | return sage_structure_coerce_parent(z)(Integer(1)) |
| 617 | elif (z+1/const_e).is_trivial_zero(): |
| 618 | return sage_structure_coerce_parent(z)(Integer(-1)) |
| 619 | elif (z+const_pi/2).is_trivial_zero(): |
| 620 | return sage_structure_coerce_parent(z)(const_pi/2*const_I) |
| 621 | return None |
| 622 | |
| 623 | def _evalf_(self, n, z, parent=None): |
| 624 | """ |
| 625 | EXAMPLES:: |
| 626 | |
| 627 | sage: N(lambert_w(1)) |
| 628 | 0.567143290409784 |
| 629 | sage: lambert_w(RealField(100)(1)) |
| 630 | 0.56714329040978387299996866221 |
| 631 | |
| 632 | SciPy is used to evaluate for float, complex, RDF, and CDF |
| 633 | inputs:: |
| 634 | |
| 635 | sage: lambert_w(RDF(1)) |
| 636 | 0.56714329041 |
| 637 | """ |
| 638 | R = parent or sage_structure_coerce_parent(z) |
| 639 | if R is float or R is complex or R is RDF or R is CDF: |
| 640 | import scipy.special |
| 641 | return scipy.special.lambertw(z, n) |
| 642 | else: |
| 643 | import mpmath |
| 644 | return mpmath_utils.call(mpmath.lambertw, z, n, parent=R) |
| 645 | |
| 646 | def _derivative_(self, n, z, diff_param=None): |
| 647 | r""" |
| 648 | The derivative of `W_n(x)` is `W_n(x)/(x \cdot W_n(x) + x)`. |
| 649 | |
| 650 | EXAMPLES:: |
| 651 | |
| 652 | sage: x = var('x') |
| 653 | sage: derivative(lambert_w(x), x) |
| 654 | lambert_w(x)/(x*lambert_w(x) + x) |
| 655 | |
| 656 | Note that this expression is not defined at two points: `-1/e` and `0.` The |
| 657 | derivative is not defined at `-1/e`, but does exist at 0:: |
| 658 | |
| 659 | sage: limit(diff(lambert_w(x),x),x=0) |
| 660 | 1 |
| 661 | |
| 662 | TESTS: |
| 663 | |
| 664 | Make sure that we can only take the derivative with respect to z and not n:: |
| 665 | |
| 666 | sage: n, z = var("n z") |
| 667 | sage: diff(lambert_w(n, z), z) |
| 668 | lambert_w(n, z)/(z*lambert_w(n, z) + z) |
| 669 | sage: diff(lambert_w(z), z) |
| 670 | lambert_w(z)/(z*lambert_w(z) + z) |
| 671 | sage: diff(lambert_w(n, z), n) |
| 672 | Traceback (most recent call last): |
| 673 | ... |
| 674 | ValueError: derivative not defined with respect to the branch number |
| 675 | sage: diff(lambert_w(2, z), n) |
| 676 | 0 |
| 677 | |
| 678 | """ |
| 679 | if diff_param != 1: |
| 680 | raise ValueError("derivative not defined with respect to the branch number") |
| 681 | return lambert_w(n, z)/(z*lambert_w(n, z)+z) |
| 682 | |
| 683 | def _maxima_init_evaled_(self, n, z): |
| 684 | """ |
| 685 | EXAMPLES: |
| 686 | |
| 687 | These are indirect doctests for this function.:: |
| 688 | |
| 689 | sage: lambert_w(0, x)._maxima_() |
| 690 | lambert_w(x) |
| 691 | sage: lambert_w(1, x)._maxima_() |
| 692 | Traceback (most recent call last): |
| 693 | ... |
| 694 | NotImplementedError: non-principal branch lambert_w[1](x) is not implemented in Maxima |
| 695 | """ |
| 696 | if n == 0: |
| 697 | return "lambert_w(%s)" % z |
| 698 | else: |
| 699 | raise NotImplementedError("non-principal branch lambert_w[%s](%s) is not implemented in Maxima" % (n, z)) |
| 700 | |
| 701 | |
| 702 | def _print_(self, n, z): |
| 703 | """ |
| 704 | Custom _print_ method to avoid printing the branch number if |
| 705 | it is zero. |
| 706 | |
| 707 | EXAMPLES:: |
| 708 | |
| 709 | sage: lambert_w(1) |
| 710 | lambert_w(1) |
| 711 | sage: lambert_w(0,x) |
| 712 | lambert_w(x) |
| 713 | """ |
| 714 | if n == 0: |
| 715 | return "lambert_w(%s)" % z |
| 716 | else: |
| 717 | return "lambert_w(%s, %s)" % (n,z) |
| 718 | |
| 719 | def _print_latex_(self, n, z): |
| 720 | r""" |
| 721 | Custom _print_latex_ method to avoid printing the branch |
| 722 | number if it is zero. |
| 723 | |
| 724 | EXAMPLES:: |
| 725 | |
| 726 | sage: latex(lambert_w(1)) |
| 727 | \operatorname{W_0}(1) |
| 728 | sage: latex(lambert_w(0,x)) |
| 729 | \operatorname{W_0}(x) |
| 730 | sage: latex(lambert_w(1,x)) |
| 731 | \operatorname{W_{1}}(x) |
| 732 | """ |
| 733 | if n == 0: |
| 734 | return r"\operatorname{W_0}(%s)" % z |
| 735 | else: |
| 736 | return r"\operatorname{W_{%s}}(%s)" % (n,z) |
| 737 | |
| 738 | lambert_w = Function_lambert_w() |