| 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, |
| 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 three 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 | |
| 585 | The special values only hold on the principal branch:: |
| 586 | |
| 587 | sage: lambert_w(1,e) |
| 588 | lambert_w(1, e) |
| 589 | sage: lambert_w(1, e.n()) |
| 590 | -0.532092121986380 + 4.59715801330257*I |
| 591 | |
| 592 | TESTS: |
| 593 | |
| 594 | When automatic simplication occurs, the parent of the output value should be |
| 595 | either the same as the parent of the input, or a Sage type:: |
| 596 | |
| 597 | sage: parent(lambert_w(int(0))) |
| 598 | <type 'int'> |
| 599 | sage: parent(lambert_w(Integer(0))) |
| 600 | Integer Ring |
| 601 | sage: parent(lambert_w(e)) |
| 602 | Integer Ring |
| 603 | """ |
| 604 | if not isinstance(z, Expression): |
| 605 | if is_inexact(z): |
| 606 | return self._evalf_(n, z, parent=sage_structure_coerce_parent(z)) |
| 607 | elif n == 0 and z == 0: |
| 608 | return sage_structure_coerce_parent(z)(Integer(0)) |
| 609 | elif n == 0: |
| 610 | if z.is_trivial_zero(): |
| 611 | return sage_structure_coerce_parent(z)(Integer(0)) |
| 612 | elif (z-const_e).is_trivial_zero(): |
| 613 | return sage_structure_coerce_parent(z)(Integer(1)) |
| 614 | elif (z+1/const_e).is_trivial_zero(): |
| 615 | return sage_structure_coerce_parent(z)(Integer(-1)) |
| 616 | return None |
| 617 | |
| 618 | def _evalf_(self, n, z, parent=None): |
| 619 | """ |
| 620 | EXAMPLES:: |
| 621 | |
| 622 | sage: N(lambert_w(1)) |
| 623 | 0.567143290409784 |
| 624 | sage: lambert_w(RealField(100)(1)) |
| 625 | 0.56714329040978387299996866221 |
| 626 | |
| 627 | SciPy is used to evaluate for float, RDF, and CDF inputs:: |
| 628 | |
| 629 | sage: lambert_w(RDF(1)) |
| 630 | 0.56714329041 |
| 631 | """ |
| 632 | R = parent or sage_structure_coerce_parent(z) |
| 633 | if R is float or R is complex or R is RDF or R is CDF: |
| 634 | import scipy.special |
| 635 | return scipy.special.lambertw(z, n) |
| 636 | else: |
| 637 | import mpmath |
| 638 | return mpmath_utils.call(mpmath.lambertw, z, n, parent=R) |
| 639 | |
| 640 | def _derivative_(self, n, z, diff_param=None): |
| 641 | """ |
| 642 | The derivative of `W_n(x)` is `W_n(x)/(x \cdot W_n(x) + x)`. |
| 643 | |
| 644 | EXAMPLES:: |
| 645 | |
| 646 | sage: x = var('x') |
| 647 | sage: derivative(lambert_w(x), x) |
| 648 | lambert_w(x)/(x*lambert_w(x) + x) |
| 649 | """ |
| 650 | return lambert_w(n, z)/(z*lambert_w(n, z)+z) |
| 651 | |
| 652 | def _maxima_init_evaled_(self, n, z): |
| 653 | """ |
| 654 | EXAMPLES: |
| 655 | |
| 656 | These are indirect doctests for this function.:: |
| 657 | |
| 658 | sage: lambert_w(0, x)._maxima_() |
| 659 | lambert_w(x) |
| 660 | sage: lambert_w(1, x)._maxima_() |
| 661 | Traceback (most recent call last): |
| 662 | ... |
| 663 | NotImplementedError: Non-principal branch lambert_w[1](x) is not implemented in Maxima |
| 664 | """ |
| 665 | if n == 0: |
| 666 | return "lambert_w(%s)" % z |
| 667 | else: |
| 668 | raise NotImplementedError("Non-principal branch lambert_w[%s](%s) is not implemented in Maxima" % (n, z)) |
| 669 | |
| 670 | |
| 671 | def _print_(self, n, z): |
| 672 | """ |
| 673 | Custom _print_ method to avoid printing the branch number if |
| 674 | it is zero. |
| 675 | |
| 676 | EXAMPLES:: |
| 677 | |
| 678 | sage: lambert_w(1) |
| 679 | lambert_w(1) |
| 680 | sage: lambert_w(0,x) |
| 681 | lambert_w(x) |
| 682 | """ |
| 683 | if n == 0: |
| 684 | return "lambert_w(%s)" % z |
| 685 | else: |
| 686 | return "lambert_w(%s, %s)" % (n,z) |
| 687 | |
| 688 | def _print_latex_(self, n, z): |
| 689 | """ |
| 690 | Custom _print_latex_ method to avoid printing the branch |
| 691 | number if it is zero. |
| 692 | |
| 693 | EXAMPLES:: |
| 694 | |
| 695 | sage: latex(lambert_w(1)) |
| 696 | \operatorname{W_0}(1) |
| 697 | sage: latex(lambert_w(0,x)) |
| 698 | \operatorname{W_0}(x) |
| 699 | sage: latex(lambert_w(1,x)) |
| 700 | \operatorname{W_{1}}(x) |
| 701 | """ |
| 702 | if n == 0: |
| 703 | return r"\operatorname{W_0}(%s)" % z |
| 704 | else: |
| 705 | return r"\operatorname{W_{%s}}(%s)" % (n,z) |
| 706 | |
| 707 | lambert_w = Function_lambert_w() |