Ticket #11888: trac_11888_v2.patch

File trac_11888_v2.patch, 8.8 KB (added by benjaminfjones, 8 years ago)

proof of concept patch (not ready for review)

  • sage/functions/all.py

    # HG changeset patch
    # User Benjamin Jones <benjaminfjones@gmail.com>
    # Date 1328501390 21600
    # Node ID f4852bbf918bd1385bbca27e644f987c155bddcc
    # Parent  c239be1054e01526a1b0b62da6691061b9dd5587
    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
     5
     6from sage.libs.mpmath import utils as mpmath_utils
     7from sage.structure.coerce import parent
     8from sage.symbolic.expression import Expression
     9from sage.rings.real_double import RDF
     10from sage.rings.complex_double import CDF
    511
    612class Function_exp(GinacFunction):
    713    def __init__(self):
     
    435441
    436442dilog = Function_dilog()
    437443
     444
     445class Function_lambert_w_branch(BuiltinFunction):
     446    r"""
     447    The integral branches of the Lambert W function `W_n(z)`.
     448
     449    This function satisfies the equation
     450
     451    .. math::
     452
     453        z = W_n(z) e^{W_n(z)}
     454
     455    INPUT:
     456
     457    - ``z`` - a complex number
     458
     459    - ``n`` - an integer. `n=0` corresponds to the principle branch.
     460
     461    ALGORITHM:
     462
     463    Numerical evaluation is handled using the mpmath library.
     464
     465    REFERENCES:
     466
     467    - http://en.wikipedia.org/wiki/Lambert_W_function
     468
     469    EXAMPLES::
     470
     471        sage: lambert_w(1.0)
     472        0.567143290409784
     473        sage: lambert_w(-1).n()
     474        -0.318131505204764 + 1.33723570143069*I
     475        sage: lambert_w(-1.5 + 5*I)
     476        1.17418016254171 + 1.10651494102011*I
     477
     478        sage: lambert_w(RealField(100)(1))
     479        0.56714329040978387299996866221
     480
     481        sage: S = solve(e^(5*x)+x==0, x, to_poly_solve=True)
     482        sage: z = S[0].rhs(); z
     483        -1/5*lambert_w(5)
     484        sage: N(z)
     485        -0.265344933048440
     486
     487    Check the defining equation numerically at `z=5`::
     488
     489        sage: N(lambert_w(5)*exp(lambert_w(5)) - 5)
     490        0.000000000000000
     491    """
     492
     493    def __init__(self):
     494        """
     495        See the docstring for :meth:`Function_lambert_w`.
     496
     497        EXAMPLES::
     498
     499            sage: lambert_w(1.0)
     500            0.567143290409784
     501
     502        """
     503        BuiltinFunction.__init__(self, "lambert_w_branch", nargs=2, latex_name=r'W')
     504
     505    def _eval_(self, z, n):
     506        """
     507        EXAMPLES::
     508
     509            sage: lambert_w(0)
     510            lambert_w(0)
     511            sage: x = var('x')
     512            sage: lambert_w(x)
     513            lambert_w(x)
     514            sage: lambert_w(0.0)
     515            0.000000000000000
     516
     517        """
     518        if not isinstance(z, Expression) and is_inexact(z):
     519            return self._evalf_(z, n, parent=parent(z))
     520
     521        return None
     522
     523    def _evalf_(self, z, n, parent=None):
     524        """
     525        EXAMPLES::
     526
     527            sage: N(lambert_w(1))
     528            0.567143290409784
     529            sage: lambert_w(RealField(100)(1))
     530            0.56714329040978387299996866221
     531
     532        """
     533
     534        if parent(z) == RDF or parent(z) == CDF:
     535            import scipy.special
     536            return scipy.special.lambertw(z, n)
     537        else:
     538            import mpmath
     539            return mpmath_utils.call(mpmath.lambertw, z, n, parent=parent)
     540
     541    def _derivative_(self, z, n, diff_param=None):
     542        """
     543        The derivative of `W_n(x)` is `W_n(x)/(x \cdot W_n(x) + x)`.
     544
     545        EXAMPLES::
     546
     547            sage: x = var('x')
     548            sage: derivative(lambert_w(x), x)
     549            lambert_w(x)/(x*lambert_w(x) + x)
     550
     551        """
     552        return lambert_w(z, n)/(z*lambert_w(z, n)+z)
     553
     554lambert_w_branch = Function_lambert_w_branch()
     555
     556def lambert_w(z, *args, **kwds):
     557    """
     558    The Lambert w function. This is a wrapper for lambert_w_branch.
     559    """
     560    if not args:
     561        return lambert_w_branch(z, 0, **kwds)
     562    elif len(args) > 1:
     563        raise ArgumentError("lambert_w takes at most two arguments.")
     564    else:
     565        return lambert_w_branch(z, args[0], **kwds)
  • sage/symbolic/random_tests.py

    diff --git a/sage/symbolic/random_tests.py b/sage/symbolic/random_tests.py
    a b  
    1414
    1515        sage: from sage.symbolic.random_tests import _mk_full_functions
    1616        sage: [f for (one,f,arity) in _mk_full_functions()]
    17         [Ei, abs, arccos, arccosh, arccot, arccoth, arccsc, arccsch,
    18         arcsec, arcsech, arcsin, arcsinh, arctan, arctan2, arctanh,
    19         binomial, ceil, conjugate, cos, cosh, cot, coth, csc, csch,
    20         dickman_rho, dilog, dirac_delta, elliptic_e, elliptic_ec,
    21         elliptic_eu, elliptic_f, elliptic_kc, elliptic_pi, erf, exp,
    22         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]
     17        [Ei, abs, arccos, arccosh, arccot, arccoth, arccsc, arccsch, arcsec,
     18        arcsech, arcsin, arcsinh, arctan, arctan2, arctanh, binomial, ceil,
     19        conjugate, cos, cosh, cot, coth, csc, csch, dickman_rho, dilog,
     20        dirac_delta, elliptic_e, elliptic_ec, elliptic_eu, elliptic_f,
     21        elliptic_kc, elliptic_pi, erf, exp, factorial, floor, heaviside,
     22        imag_part, integrate, kronecker_delta, lambert_w, log, polylog,
     23        real_part, sec, sech, sgn, sin, sinh, tan, tanh, unit_step, zeta,
     24        zetaderiv]
    2525
    2626    Note that this doctest will fail whenever a Pynac function is added or
    2727    removed.  In that case, it is very likely that the doctests for
     
    234234        sage: from sage.symbolic.random_tests import *
    235235        sage: set_random_seed(2)
    236236        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)
     237        -v1*elliptic_kc((0.0666829501658 + 0.206976992303*I)/(v3 + e))/v3 +
     238        heaviside(arccosh(-(abs(v2 -
     239        floor(-v3))^(-real_part(elliptic_e(-0.703991792631 + 0.750156228797*I,
     240        -(1.21734510331 - 1.22580558833*I)*pi*v1)))*e^(arctan2(imag_part(v2) -
     241        imag_part(floor(-v3)), real_part(v2) -
     242        real_part(floor(-v3)))*imag_part(elliptic_e(-0.703991792631 +
     243        0.750156228797*I, -(1.21734510331 -
     244        1.22580558833*I)*pi*v1)))*sin(-log(abs(v2 -
     245        floor(-v3)))*imag_part(elliptic_e(-0.703991792631 + 0.750156228797*I,
     246        -(1.21734510331 - 1.22580558833*I)*pi*v1)) - arctan2(imag_part(v2) -
     247        imag_part(floor(-v3)), real_part(v2) -
     248        real_part(floor(-v3)))*real_part(elliptic_e(-0.703991792631 +
     249        0.750156228797*I, -(1.21734510331 -
     250        1.22580558833*I)*pi*v1)))*imag_part(arccsc((0.891138386848 -
     251        0.0936820840629*I)/v1)) - abs(v2 -
     252        floor(-v3))^(-real_part(elliptic_e(-0.703991792631 + 0.750156228797*I,
     253        -(1.21734510331 - 1.22580558833*I)*pi*v1)))*e^(arctan2(imag_part(v2) -
     254        imag_part(floor(-v3)), real_part(v2) -
     255        real_part(floor(-v3)))*imag_part(elliptic_e(-0.703991792631 +
     256        0.750156228797*I, -(1.21734510331 -
     257        1.22580558833*I)*pi*v1)))*cos(-log(abs(v2 -
     258        floor(-v3)))*imag_part(elliptic_e(-0.703991792631 + 0.750156228797*I,
     259        -(1.21734510331 - 1.22580558833*I)*pi*v1)) - arctan2(imag_part(v2) -
     260        imag_part(floor(-v3)), real_part(v2) -
     261        real_part(floor(-v3)))*real_part(elliptic_e(-0.703991792631 +
     262        0.750156228797*I, -(1.21734510331 -
     263        1.22580558833*I)*pi*v1)))*real_part(arccsc((0.891138386848 -
     264        0.0936820840629*I)/v1)) + v3^(-0.48519994364 - 0.485764091302*I) +
     265        0.0909404921682*real_part(v3)/(real_part(v3)^2 + imag_part(v3)^2) +
     266        0.281538203756*imag_part(v3)/(real_part(v3)^2 + imag_part(v3)^2) -
     267        0.647983235144*real_part(v1) - 1.20665952957*imag_part(v1))*v1))
    238268        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)
     269        About to apply <built-in function mul> to [v1, e]
     270        About to apply <built-in function div> to [v1*e, 1]
     271        v1*e
    243272    """
    244273    vars = [(1.0, sage.calculus.calculus.var('v%d' % (n+1))) for n in range(nvars)]
    245274    if ncoeffs is None: