Ticket #11888: trac_11888_v4.patch

File trac_11888_v4.patch, 5.6 KB (added by benjaminfjones, 8 years ago)

addressed reviewer issues, changed order of arguments to be consistant with Mma/Maple?

  • sage/functions/all.py

    # HG changeset patch
    # User Benjamin Jones <benjaminfjones@gmail.com>
    # Date 1328501390 21600
    # Node ID 3c8e890562e91a479b8a71dbee088b0f14deb08e
    # Parent  eba013370ff69c553c44b2ef9f71a148c911f5a5
    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  
    2121                    real_part, real,
    2222                    imag_part, imag, imaginary, conjugate)
    2323
    24 from log import (exp, log, ln, polylog, dilog)
     24from log import (exp, log, ln, polylog, dilog, lambert_w, lambert_w_branch)
    2525
    2626
    2727from transcendental import (exponential_integral_1,
  • sage/functions/log.py

    diff --git a/sage/functions/log.py b/sage/functions/log.py
    a b  
    11"""
    22Logarithmic functions
    33"""
    4 from sage.symbolic.function import GinacFunction
     4from sage.symbolic.function import GinacFunction, BuiltinFunction, is_inexact
     5from sage.symbolic.pynac import symbol_table
     6
     7from sage.libs.mpmath import utils as mpmath_utils
     8from sage.structure.coerce import parent
     9from sage.symbolic.expression import Expression
     10from sage.rings.real_double import RDF
     11from sage.rings.complex_double import CDF
    512
    613class Function_exp(GinacFunction):
    714    def __init__(self):
     
    435442
    436443dilog = Function_dilog()
    437444
     445
     446class 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
     565lambert_w_branch = Function_lambert_w_branch()
     566
     567def 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
     609symbol_table['functions']['lambert_w'] = lambert_w