Ticket #11888: trac_11888_v5.patch

File trac_11888_v5.patch, 7.9 KB (added by benjaminfjones, 8 years ago)

adds lambert_w function

  • sage/functions/all.py

    # HG changeset patch
    # User Benjamin Jones <benjaminfjones@gmail.com>
    # Date 1328501390 21600
    # Node ID 8dce01b3732625ea69fba25b5887c73a97ebeb55
    # Parent  7f21a9a76fab333da15bf01dd9fc5972816d4894
    Trac 11888: add the lambert_w symbolic function
    
    diff --git a/sage/functions/all.py b/sage/functions/all.py
    a b  
    2121                    arg, 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)
    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
     6from sage.symbolic.constants import e as const_e
     7
     8from sage.libs.mpmath import utils as mpmath_utils
     9from sage.structure.coerce import parent as sage_structure_coerce_parent
     10from sage.symbolic.expression import Expression
     11from sage.rings.real_double import RDF
     12from sage.rings.complex_double import CDF
     13from sage.rings.all import Integer
    514
    615class Function_exp(GinacFunction):
    716    def __init__(self):
     
    429438            dilog(-1/2*I)
    430439            sage: conjugate(dilog(2))
    431440            conjugate(dilog(2))
    432         """       
     441        """
    433442        GinacFunction.__init__(self, 'dilog',
    434443                conversions=dict(maxima='li[2]'))
    435444
    436445dilog = Function_dilog()
    437446
     447
     448class 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
     515    def __init__(self):
     516        r"""
     517        See the docstring for :meth:`Function_lambert_w`.
     518
     519        EXAMPLES::
     520
     521            sage: lambert_w(0, 1.0)
     522            0.567143290409784
     523        """
     524        BuiltinFunction.__init__(self, "lambert_w", nargs=2,
     525                                 conversions={'maxima':'lambert_w',
     526                                              'mathematica':'ProductLog',
     527                                              'maple':'LambertW'})
     528        self._simplifications = {Integer(0):Integer(0), const_e:Integer(1),
     529                                 -1/const_e:Integer(-1)}
     530
     531    def __call__(self, *args, **kwds):
     532        r"""
     533        Custom call method allows the user to pass one argument or two. If
     534        one argument is passed, we assume it is ``z`` and that ``n=0``.
     535
     536        EXAMPLES::
     537
     538            sage: lambert_w(1)
     539            lambert_w(1)
     540            sage: lambert_w(1, 2)
     541            lambert_w(1, 2)
     542        """
     543        if len(args) == 2:
     544            return BuiltinFunction.__call__(self, *args, **kwds)
     545        elif len(args) == 1:
     546            return BuiltinFunction.__call__(self, 0, args[0], **kwds)
     547        else:
     548            raise TypeError("lambert_w takes either one or two arguments.")
     549
     550    def _eval_(self, n, z):
     551        """
     552        EXAMPLES::
     553
     554            sage: lambert_w(6.0)
     555            1.43240477589830
     556            sage: lambert_w(1)
     557            lambert_w(1)
     558            sage: lambert_w(x+1)
     559            lambert_w(x + 1)
     560
     561        There are three special values which are automatically simplified::
     562
     563            sage: lambert_w(0)
     564            0
     565            sage: lambert_w(e)
     566            1
     567            sage: lambert_w(-1/e)
     568            -1
     569            sage: lambert_w(SR(-1/e))
     570            -1
     571            sage: lambert_w(SR(0))   
     572            0
     573
     574        The special values only hold on the principal branch::
     575
     576            sage: lambert_w(1,e)
     577            lambert_w(1, e)
     578            sage: lambert_w(1, e.n())
     579            -0.532092121986380 + 4.59715801330257*I
     580        """
     581        if not isinstance(z, Expression):
     582            if is_inexact(z):
     583                return self._evalf_(n, z, parent=sage_structure_coerce_parent(z))
     584            elif z == Integer(0):
     585                return Integer(0)
     586        elif n == 0:
     587            if z.is_trivial_zero():
     588                return Integer(0)
     589            elif (z-const_e).is_trivial_zero():
     590                return Integer(1)
     591            elif (z+1/const_e).is_trivial_zero():
     592                return Integer(-1)
     593        return None
     594
     595    def _evalf_(self, n, z, parent=None):
     596        """
     597        EXAMPLES::
     598
     599            sage: N(lambert_w(1))
     600            0.567143290409784
     601            sage: lambert_w(RealField(100)(1))
     602            0.56714329040978387299996866221
     603
     604        SciPy is used to evaluate for float, RDF, and CDF inputs::
     605
     606            sage: lambert_w(RDF(1))
     607            0.56714329041
     608        """
     609        R = parent or sage_structure_coerce_parent(z)
     610        if R is float or R is complex or R is RDF or R is CDF:
     611            import scipy.special
     612            return scipy.special.lambertw(z, n)
     613        else:
     614            import mpmath
     615            return mpmath_utils.call(mpmath.lambertw, z, n, parent=R)
     616
     617    def _derivative_(self, n, z, diff_param=None):
     618        """
     619        The derivative of `W_n(x)` is `W_n(x)/(x \cdot W_n(x) + x)`.
     620
     621        EXAMPLES::
     622
     623            sage: x = var('x')
     624            sage: derivative(lambert_w(x), x)
     625            lambert_w(x)/(x*lambert_w(x) + x)
     626        """
     627        return lambert_w(n, z)/(z*lambert_w(n, z)+z)
     628
     629    def _print_(self, n, z):
     630        """
     631        Custom _print_ method to avoid printing the branch number if
     632        it is zero.
     633
     634        EXAMPLES::
     635
     636            sage: lambert_w(1)
     637            lambert_w(1)
     638            sage: lambert_w(0,x)
     639            lambert_w(x)
     640        """
     641        if n == 0:
     642            return "lambert_w(%s)" % z
     643        else:
     644            return "lambert_w(%s, %s)" % (n,z)
     645
     646    def _print_latex_(self, n, z):
     647        """
     648        Custom _print_latex_ method to avoid printing the branch
     649        number if it is zero.
     650
     651        EXAMPLES::
     652
     653            sage: latex(lambert_w(1))
     654            \operatorname{W_0}(1)
     655            sage: latex(lambert_w(0,x))
     656            \operatorname{W_0}(x)
     657            sage: latex(lambert_w(1,x))
     658            \operatorname{W_{1}}(x)
     659        """
     660        if n == 0:
     661            return r"\operatorname{W_0}(%s)" % z
     662        else:
     663            return r"\operatorname{W_{%s}}(%s)" % (n,z)
     664
     665lambert_w = Function_lambert_w()