| 445 | |
| 446 | class Function_lambert_w_branch(BuiltinFunction): |
| 447 | r""" |
| 448 | The integral branches of the Lambert W function `W_n(z)`. |
| 449 | |
| 450 | This function satisfies the equation |
| 451 | |
| 452 | .. math:: |
| 453 | |
| 454 | z = W_n(z) e^{W_n(z)} |
| 455 | |
| 456 | INPUT: |
| 457 | |
| 458 | - ``n`` - an integer. `n=0` corresponds to the principal branch. |
| 459 | |
| 460 | - ``z`` - a complex number |
| 461 | |
| 462 | ALGORITHM: |
| 463 | |
| 464 | Numerical evaluation is handled using the mpmath and SciPy libraries. |
| 465 | |
| 466 | REFERENCES: |
| 467 | |
| 468 | - http://en.wikipedia.org/wiki/Lambert_W_function |
| 469 | |
| 470 | EXAMPLES: |
| 471 | |
| 472 | Evaluation of the principal branch:: |
| 473 | |
| 474 | sage: lambert_w_branch(0, 1.0) |
| 475 | 0.567143290409784 |
| 476 | sage: lambert_w_branch(0, -1).n() |
| 477 | -0.318131505204764 + 1.33723570143069*I |
| 478 | sage: lambert_w_branch(0, -1.5 + 5*I) |
| 479 | 1.17418016254171 + 1.10651494102011*I |
| 480 | |
| 481 | Evaluation of other branches:: |
| 482 | |
| 483 | sage: lambert_w_branch(2, 1.0) |
| 484 | -2.40158510486800 + 10.7762995161151*I |
| 485 | |
| 486 | Solutions to certain exponential equations are returned in terms of lambert_w:: |
| 487 | |
| 488 | sage: S = solve(e^(5*x)+x==0, x, to_poly_solve=True) |
| 489 | sage: z = S[0].rhs(); z |
| 490 | -1/5*lambert_w_branch(0, 5) |
| 491 | sage: N(z) |
| 492 | -0.265344933048440 |
| 493 | |
| 494 | Check the defining equation numerically at `z=5`:: |
| 495 | |
| 496 | sage: N(lambert_w(5)*exp(lambert_w(5)) - 5) |
| 497 | 0.000000000000000 |
| 498 | """ |
| 499 | |
| 500 | def __init__(self): |
| 501 | """ |
| 502 | See the docstring for :meth:`Function_lambert_w`. |
| 503 | |
| 504 | EXAMPLES:: |
| 505 | |
| 506 | sage: lambert_w_branch(0, 1.0) |
| 507 | 0.567143290409784 |
| 508 | |
| 509 | """ |
| 510 | BuiltinFunction.__init__(self, "lambert_w_branch", nargs=2, |
| 511 | latex_name=r'W', conversions={'mathematica':'ProductLog', |
| 512 | 'maple':'LambertW'}) |
| 513 | |
| 514 | def _eval_(self, n, z): |
| 515 | """ |
| 516 | EXAMPLES:: |
| 517 | |
| 518 | sage: lambert_w_branch(0, 0) |
| 519 | lambert_w_branch(0, 0) |
| 520 | sage: x = var('x') |
| 521 | sage: lambert_w_branch(1, x) |
| 522 | lambert_w_branch(1, x) |
| 523 | sage: lambert_w_branch(0, 0.0) |
| 524 | 0.000000000000000 |
| 525 | """ |
| 526 | if not isinstance(z, Expression) and is_inexact(z): |
| 527 | return self._evalf_(n, z, parent=parent(z)) |
| 528 | |
| 529 | return None |
| 530 | |
| 531 | def _evalf_(self, n, z, parent=None): |
| 532 | """ |
| 533 | EXAMPLES:: |
| 534 | |
| 535 | sage: N(lambert_w_branch(0, 1)) |
| 536 | 0.567143290409784 |
| 537 | sage: lambert_w_branch(0, RealField(100)(1)) |
| 538 | 0.56714329040978387299996866221 |
| 539 | |
| 540 | SciPy is used to evaluate for RDF or CDF inputs:: |
| 541 | |
| 542 | sage: lambert_w_branch(0, RDF(1)) |
| 543 | 0.56714329041 |
| 544 | """ |
| 545 | |
| 546 | if parent(z) == RDF or parent(z) == CDF: |
| 547 | import scipy.special |
| 548 | return scipy.special.lambertw(z, n) |
| 549 | else: |
| 550 | import mpmath |
| 551 | return mpmath_utils.call(mpmath.lambertw, z, n, parent=parent) |
| 552 | |
| 553 | def _derivative_(self, n, z, diff_param=None): |
| 554 | """ |
| 555 | The derivative of `W_n(x)` is `W_n(x)/(x \cdot W_n(x) + x)`. |
| 556 | |
| 557 | EXAMPLES:: |
| 558 | |
| 559 | sage: x = var('x') |
| 560 | sage: derivative(lambert_w(x), x) |
| 561 | lambert_w_branch(0, x)/(x*lambert_w_branch(0, x) + x) |
| 562 | """ |
| 563 | return lambert_w_branch(n, z)/(z*lambert_w_branch(n, z)+z) |
| 564 | |
| 565 | lambert_w_branch = Function_lambert_w_branch() |
| 566 | |
| 567 | def lambert_w(z, *args, **kwds): |
| 568 | """ |
| 569 | The principal branch of the Lambert W function. This is a wrapper for |
| 570 | lambert_w_branch(0, z). |
| 571 | |
| 572 | INPUT: |
| 573 | |
| 574 | - ``z`` - complex number |
| 575 | |
| 576 | OUTPUT: |
| 577 | |
| 578 | - The value ``W_0(z)`` of the principal branch of the lambert_w function. The output |
| 579 | satisfies the equation |
| 580 | |
| 581 | .. math:: |
| 582 | |
| 583 | z = W_0(z) e^{W_0(z)} |
| 584 | |
| 585 | ALGORITHM: |
| 586 | |
| 587 | Numerical evaluation is handled using the mpmath and SciPy libraries. |
| 588 | |
| 589 | REFERENCES: |
| 590 | |
| 591 | - http://en.wikipedia.org/wiki/Lambert_W_function |
| 592 | |
| 593 | EXAMPLES:: |
| 594 | |
| 595 | sage: lambert_w(1.0) |
| 596 | 0.567143290409784 |
| 597 | sage: lambert_w(2) |
| 598 | lambert_w_branch(0, 2) |
| 599 | sage: lambert_w(2).n() |
| 600 | 0.852605502013726 |
| 601 | """ |
| 602 | if not args: |
| 603 | return lambert_w_branch(0, z, **kwds) |
| 604 | elif len(args) > 1: |
| 605 | raise ArgumentError("lambert_w takes at most two arguments.") |
| 606 | else: |
| 607 | return lambert_w_branch(args[0], z, **kwds) |
| 608 | |
| 609 | symbol_table['functions']['lambert_w'] = lambert_w |