Ticket #11888: trac_11888_v3.patch

File trac_11888_v3.patch, 9.3 KB (added by benjaminfjones, 8 years ago)

adds lambert_w and lambert_w_branch functions

  • sage/functions/all.py

    # HG changeset patch
    # User Benjamin Jones <benjaminfjones@gmail.com>
    # Date 1328501390 21600
    # Node ID 42b52b609add82f1af3993dd393fd3fbbfd065a2
    # 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    - ``z`` - a complex number
     459
     460    - ``n`` - an integer. `n=0` corresponds to the principal branch.
     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(1.0, 0)
     475        0.567143290409784
     476        sage: lambert_w_branch(-1,0).n()
     477        -0.318131505204764 + 1.33723570143069*I
     478        sage: lambert_w_branch(-1.5 + 5*I, 0)
     479        1.17418016254171 + 1.10651494102011*I
     480
     481    Evaluation of other branches::
     482
     483        sage: lambert_w_branch(1.0, 2)
     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(5, 0)
     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(1.0, 0)
     507            0.567143290409784
     508
     509        """
     510        BuiltinFunction.__init__(self, "lambert_w_branch", nargs=2, latex_name=r'W')
     511
     512    def _eval_(self, z, n):
     513        """
     514        EXAMPLES::
     515
     516            sage: lambert_w_branch(0, 0)
     517            lambert_w_branch(0, 0)
     518            sage: x = var('x')
     519            sage: lambert_w_branch(x, 1)
     520            lambert_w_branch(x, 1)
     521            sage: lambert_w_branch(0.0, 0)
     522            0.000000000000000
     523        """
     524        if not isinstance(z, Expression) and is_inexact(z):
     525            return self._evalf_(z, n, parent=parent(z))
     526
     527        return None
     528
     529    def _evalf_(self, z, n, parent=None):
     530        """
     531        EXAMPLES::
     532
     533            sage: N(lambert_w_branch(1,0))
     534            0.567143290409784
     535            sage: lambert_w_branch(RealField(100)(1),0)
     536            0.56714329040978387299996866221
     537
     538        SciPy is used to evalute for RDF or CDF inputs::
     539
     540            sage: lambert_w_branch(RDF(1),0)
     541            0.56714329041
     542        """
     543
     544        if parent(z) == RDF or parent(z) == CDF:
     545            import scipy.special
     546            return scipy.special.lambertw(z, n)
     547        else:
     548            import mpmath
     549            return mpmath_utils.call(mpmath.lambertw, z, n, parent=parent)
     550
     551    def _derivative_(self, z, n, diff_param=None):
     552        """
     553        The derivative of `W_n(x)` is `W_n(x)/(x \cdot W_n(x) + x)`.
     554
     555        EXAMPLES::
     556
     557            sage: x = var('x')
     558            sage: derivative(lambert_w(x), x)
     559            lambert_w_branch(x, 0)/(x*lambert_w_branch(x, 0) + x)
     560        """
     561        return lambert_w(z, n)/(z*lambert_w(z, n)+z)
     562
     563lambert_w_branch = Function_lambert_w_branch()
     564
     565def lambert_w(z, *args, **kwds):
     566    """
     567    The principal branch of the Lambert W function. This is a wrapper for
     568    lambert_w_branch(z,0).
     569
     570    INPUT:
     571
     572    - ``z`` - complex number
     573
     574    OUTPUT:
     575
     576    - The value ``W_0(z)`` of the principal branch of the lambert_w function. The output
     577      satisfies the equation
     578
     579    .. math::
     580
     581        z = W_0(z) e^{W_0(z)}
     582   
     583    ALGORITHM:
     584
     585    Numerical evaluation is handled using the mpmath and SciPy libraries.
     586
     587    REFERENCES:
     588
     589    - http://en.wikipedia.org/wiki/Lambert_W_function
     590
     591    EXAMPLES::
     592
     593        sage: lambert_w(1.0)
     594        0.567143290409784
     595        sage: lambert_w(2)
     596        lambert_w_branch(2, 0)
     597        sage: lambert_w(2).n()
     598        0.852605502013726
     599    """
     600    if not args:
     601        return lambert_w_branch(z, 0, **kwds)
     602    elif len(args) > 1:
     603        raise ArgumentError("lambert_w takes at most two arguments.")
     604    else:
     605        return lambert_w_branch(z, args[0], **kwds)
     606
     607symbol_table['functions']['lambert_w'] = lambert_w
  • sage/symbolic/random_tests.py

    diff --git a/sage/symbolic/random_tests.py b/sage/symbolic/random_tests.py
    a b  
    2020        dickman_rho, dilog, dirac_delta, elliptic_e, elliptic_ec,
    2121        elliptic_eu, elliptic_f, elliptic_kc, elliptic_pi, erf, exp,
    2222        factorial, floor, heaviside, imag_part, integrate,
    23         kronecker_delta, log, polylog, real_part, sec, sech, sgn, sin,
    24         sinh, tan, tanh, unit_step, zeta, zetaderiv]
     23        kronecker_delta, lambert_w_branch, log, polylog, real_part,
     24        sec, sech, sgn, sin, sinh, tan, tanh, unit_step, zeta,
     25        zetaderiv]
    2526
    2627    Note that this doctest will fail whenever a Pynac function is added or
    2728    removed.  In that case, it is very likely that the doctests for
     
    234235        sage: from sage.symbolic.random_tests import *
    235236        sage: set_random_seed(2)
    236237        sage: random_expr(50, nvars=3, coeff_generator=CDF.random_element)
    237         (euler_gamma - v3^(-e) + (v2 - factorial(-e/v2))^(((2.85879036573 - 1.18163393202*I)*v2 + (2.85879036573 - 1.18163393202*I)*v3)*pi - 0.247786879678 + 0.931826724898*I)*arccsc((0.891138386848 - 0.0936820840629*I)/v1) - (0.553423153995 - 0.5481180572*I)*v3 + 0.149683576515 - 0.155746451854*I)*v1 + arccsch(pi + e)*elliptic_f(khinchin*v2, 1.4656989704 + 0.863754357069*I)
     238        -v1*elliptic_kc((0.0666829501658 + 0.206976992303*I)/(v3 +
     239         e))/v3 + heaviside(arccosh((abs(v2 -
     240         floor(-v3))^(-real_part(elliptic_e(-0.703991792631 +
     241         0.750156228797*I, -(1.21734510331 -
     242         1.22580558833*I)*pi*v1)))*real_part(arccsc((0.891138386848 -
     243         0.0936820840629*I)/v1))*e^(imag_part(elliptic_e(-0.703991792631
     244         + 0.750156228797*I, -(1.21734510331 -
     245         1.22580558833*I)*pi*v1))*arctan2(imag_part(v2) -
     246         imag_part(floor(-v3)), real_part(v2) -
     247         real_part(floor(-v3))))*cos(-real_part(elliptic_e(-0.703991792631
     248         + 0.750156228797*I, -(1.21734510331 -
     249         1.22580558833*I)*pi*v1))*arctan2(imag_part(v2) -
     250         imag_part(floor(-v3)), real_part(v2) - real_part(floor(-v3)))
     251         - imag_part(elliptic_e(-0.703991792631 + 0.750156228797*I,
     252         -(1.21734510331 - 1.22580558833*I)*pi*v1))*log(abs(v2 -
     253         floor(-v3)))) - abs(v2 -
     254         floor(-v3))^(-real_part(elliptic_e(-0.703991792631 +
     255         0.750156228797*I, -(1.21734510331 -
     256         1.22580558833*I)*pi*v1)))*imag_part(arccsc((0.891138386848 -
     257         0.0936820840629*I)/v1))*e^(imag_part(elliptic_e(-0.703991792631
     258         + 0.750156228797*I, -(1.21734510331 -
     259         1.22580558833*I)*pi*v1))*arctan2(imag_part(v2) -
     260         imag_part(floor(-v3)), real_part(v2) -
     261         real_part(floor(-v3))))*sin(-real_part(elliptic_e(-0.703991792631
     262         + 0.750156228797*I, -(1.21734510331 -
     263         1.22580558833*I)*pi*v1))*arctan2(imag_part(v2) -
     264         imag_part(floor(-v3)), real_part(v2) - real_part(floor(-v3)))
     265         - imag_part(elliptic_e(-0.703991792631 + 0.750156228797*I,
     266         -(1.21734510331 - 1.22580558833*I)*pi*v1))*log(abs(v2 -
     267         floor(-v3)))) - v3^(-0.48519994364 - 0.485764091302*I) -
     268         0.0909404921682*real_part(v3)/(real_part(v3)^2 +
     269         imag_part(v3)^2) -
     270         0.281538203756*imag_part(v3)/(real_part(v3)^2 +
     271         imag_part(v3)^2) + 0.647983235144*real_part(v1) +
     272         1.20665952957*imag_part(v1))*v1))
    238273        sage: random_expr(5, verbose=True)
    239         About to apply dirac_delta to [1]
    240         About to apply arccsch to [0]
    241         About to apply <built-in function add> to [0, arccsch(0)]
    242         arccsch(0)
     274        About to apply <built-in function mul> to [v1, e]
     275        About to apply <built-in function div> to [v1*e, 1]
     276        v1*e
    243277    """
    244278    vars = [(1.0, sage.calculus.calculus.var('v%d' % (n+1))) for n in range(nvars)]
    245279    if ncoeffs is None: