Ticket #11888: trac_11888_v9.patch

File trac_11888_v9.patch, 11.8 KB (added by dsm, 7 years ago)

post-review version of lambertw sf

  • sage/functions/all.py

    # HG changeset patch
    # User Benjamin Jones <benjaminfjones@gmail.com>
    # Date 1328501390 21600
    # Node ID b97e821d4512e0477250778aedb0094d1e771ee7
    # Parent  5299aa8232d7acb68d37162f5bef2068c8a68d54
    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, pi as const_pi, I as const_I
     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    Integration (of the principal branch) is evaluated using Maxima::
     515
     516        sage: integrate(lambert_w(x), x)
     517        (lambert_w(x)^2 - lambert_w(x) + 1)*x/lambert_w(x)
     518        sage: integrate(lambert_w(x), x, 0, 1)
     519        (lambert_w(1)^2 - lambert_w(1) + 1)/lambert_w(1) - 1
     520        sage: integrate(lambert_w(x), x, 0, 1.0)
     521        0.330366124762
     522
     523    Warning: The integral of a non-principal branch is not implemented, and
     524    neither is numerical integration using GSL. The :meth:`numerical_integral`
     525    function does work if you pass a lambda function::
     526
     527        sage: numerical_integral(lambda x: lambert_w(x), 0, 1)
     528        (0.33036612476168054, 3.667800782666048e-15)
     529    """
     530
     531    def __init__(self):
     532        r"""
     533        See the docstring for :meth:`Function_lambert_w`.
     534
     535        EXAMPLES::
     536
     537            sage: lambert_w(0, 1.0)
     538            0.567143290409784
     539        """
     540        BuiltinFunction.__init__(self, "lambert_w", nargs=2,
     541                                 conversions={'mathematica':'ProductLog',
     542                                              'maple':'LambertW'})
     543
     544    def __call__(self, *args, **kwds):
     545        r"""
     546        Custom call method allows the user to pass one argument or two. If
     547        one argument is passed, we assume it is ``z`` and that ``n=0``.
     548
     549        EXAMPLES::
     550
     551            sage: lambert_w(1)
     552            lambert_w(1)
     553            sage: lambert_w(1, 2)
     554            lambert_w(1, 2)
     555        """
     556        if len(args) == 2:
     557            return BuiltinFunction.__call__(self, *args, **kwds)
     558        elif len(args) == 1:
     559            return BuiltinFunction.__call__(self, 0, args[0], **kwds)
     560        else:
     561            raise TypeError("lambert_w takes either one or two arguments.")
     562
     563    def _eval_(self, n, z):
     564        """
     565        EXAMPLES::
     566
     567            sage: lambert_w(6.0)
     568            1.43240477589830
     569            sage: lambert_w(1)
     570            lambert_w(1)
     571            sage: lambert_w(x+1)
     572            lambert_w(x + 1)
     573
     574        There are several special values which are automatically simplified::
     575
     576            sage: lambert_w(0)
     577            0
     578            sage: lambert_w(e)
     579            1
     580            sage: lambert_w(-1/e)
     581            -1
     582            sage: lambert_w(SR(0))
     583            0
     584            sage: lambert_w(-pi/2)
     585            1/2*I*pi
     586
     587        The special values only hold on the principal branch::
     588
     589            sage: lambert_w(1,e)
     590            lambert_w(1, e)
     591            sage: lambert_w(1, e.n())
     592            -0.532092121986380 + 4.59715801330257*I
     593
     594        TESTS:
     595
     596        When automatic simplication occurs, the parent of the output value should be
     597        either the same as the parent of the input, or a Sage type::
     598
     599            sage: parent(lambert_w(int(0)))
     600            <type 'int'>
     601            sage: parent(lambert_w(Integer(0)))
     602            Integer Ring
     603            sage: parent(lambert_w(e))
     604            Integer Ring
     605        """
     606        if not isinstance(z, Expression):
     607            if is_inexact(z):
     608                return self._evalf_(n, z, parent=sage_structure_coerce_parent(z))
     609            elif n == 0 and z == 0:
     610                return sage_structure_coerce_parent(z)(Integer(0))
     611        elif n == 0:
     612            if z.is_trivial_zero():
     613                return sage_structure_coerce_parent(z)(Integer(0))
     614            elif (z-const_e).is_trivial_zero():
     615                return sage_structure_coerce_parent(z)(Integer(1))
     616            elif (z+1/const_e).is_trivial_zero():
     617                return sage_structure_coerce_parent(z)(Integer(-1))
     618            elif (z+const_pi/2).is_trivial_zero():
     619                return sage_structure_coerce_parent(z)(const_pi/2*const_I)
     620        return None
     621
     622    def _evalf_(self, n, z, parent=None):
     623        """
     624        EXAMPLES::
     625
     626            sage: N(lambert_w(1))
     627            0.567143290409784
     628            sage: lambert_w(RealField(100)(1))
     629            0.56714329040978387299996866221
     630
     631        SciPy is used to evaluate for float, complex, RDF, and CDF
     632        inputs::
     633
     634            sage: lambert_w(RDF(1))
     635            0.56714329041
     636        """
     637        R = parent or sage_structure_coerce_parent(z)
     638        if R is float or R is complex or R is RDF or R is CDF:
     639            import scipy.special
     640            return scipy.special.lambertw(z, n)
     641        else:
     642            import mpmath
     643            return mpmath_utils.call(mpmath.lambertw, z, n, parent=R)
     644
     645    def _derivative_(self, n, z, diff_param=None):
     646        """
     647        The derivative of `W_n(x)` is `W_n(x)/(x \cdot W_n(x) + x)`.
     648
     649        EXAMPLES::
     650
     651            sage: x = var('x')
     652            sage: derivative(lambert_w(x), x)
     653            lambert_w(x)/(x*lambert_w(x) + x)
     654
     655        Note that this expression is not defined at two points: `-1/e` and `0.`  The
     656        derivative is not defined at `-1/e`, but does exist at 0::
     657
     658            sage: limit(diff(lambert_w(x),x),x=0)
     659            1
     660
     661        TESTS:
     662       
     663        Make sure that we can only take the derivative with respect to z and not n::
     664
     665            sage: n, z = var("n z")
     666            sage: diff(lambert_w(n, z), z)
     667            lambert_w(n, z)/(z*lambert_w(n, z) + z)
     668            sage: diff(lambert_w(z), z)
     669            lambert_w(z)/(z*lambert_w(z) + z)
     670            sage: diff(lambert_w(n, z), n)
     671            Traceback (most recent call last):
     672            ...
     673            ValueError: Derivative not defined with respect to the branch number.
     674            sage: diff(lambert_w(2, z), n)
     675            0
     676           
     677        """
     678        if diff_param != 1:
     679            raise ValueError("Derivative not defined with respect to the branch number.")
     680        return lambert_w(n, z)/(z*lambert_w(n, z)+z)
     681
     682    def _maxima_init_evaled_(self, n, z):
     683        """
     684        EXAMPLES:
     685
     686        These are indirect doctests for this function.::
     687
     688            sage: lambert_w(0, x)._maxima_()
     689            lambert_w(x)
     690            sage: lambert_w(1, x)._maxima_()
     691            Traceback (most recent call last):
     692            ...
     693            NotImplementedError: Non-principal branch lambert_w[1](x) is not implemented in Maxima
     694        """
     695        if n == 0:
     696            return "lambert_w(%s)" % z
     697        else:
     698            raise NotImplementedError("Non-principal branch lambert_w[%s](%s) is not implemented in Maxima" % (n, z))
     699
     700
     701    def _print_(self, n, z):
     702        """
     703        Custom _print_ method to avoid printing the branch number if
     704        it is zero.
     705
     706        EXAMPLES::
     707
     708            sage: lambert_w(1)
     709            lambert_w(1)
     710            sage: lambert_w(0,x)
     711            lambert_w(x)
     712        """
     713        if n == 0:
     714            return "lambert_w(%s)" % z
     715        else:
     716            return "lambert_w(%s, %s)" % (n,z)
     717
     718    def _print_latex_(self, n, z):
     719        """
     720        Custom _print_latex_ method to avoid printing the branch
     721        number if it is zero.
     722
     723        EXAMPLES::
     724
     725            sage: latex(lambert_w(1))
     726            \operatorname{W_0}(1)
     727            sage: latex(lambert_w(0,x))
     728            \operatorname{W_0}(x)
     729            sage: latex(lambert_w(1,x))
     730            \operatorname{W_{1}}(x)
     731        """
     732        if n == 0:
     733            return r"\operatorname{W_0}(%s)" % z
     734        else:
     735            return r"\operatorname{W_{%s}}(%s)" % (n,z)
     736
     737lambert_w = Function_lambert_w()
  • sage/interfaces/maxima_lib.py

    diff --git a/sage/interfaces/maxima_lib.py b/sage/interfaces/maxima_lib.py
    a b  
    10771077    sage.functions.log.exp : "%EXP",
    10781078    sage.functions.log.ln : "%LOG",
    10791079    sage.functions.log.log : "%LOG",
     1080    sage.functions.log.lambert_w : "%LAMBERT_W",
    10801081    sage.functions.other.factorial : "MFACTORIAL",
    10811082    sage.functions.other.erf : "%ERF",
    10821083    sage.functions.other.gamma_inc : "%GAMMA_INCOMPLETE",
     
    11661167max_array=EclObject("ARRAY")
    11671168mdiff=EclObject("%DERIVATIVE")
    11681169max_gamma_incomplete=sage_op_dict[sage.functions.other.gamma_inc]
     1170max_lambert_w=sage_op_dict[sage.functions.log.lambert_w]
    11691171
    11701172def mrat_to_sage(expr):
    11711173    r"""
     
    13551357    sage.functions.log.polylog : lambda N,X : [[mqapply],[[max_li, max_array],N],X],
    13561358    sage.functions.other.psi1 : lambda X : [[mqapply],[[max_psi, max_array],0],X],
    13571359    sage.functions.other.psi2 : lambda N,X : [[mqapply],[[max_psi, max_array],N],X],
    1358     sage.functions.other.Ei : lambda X : [[max_gamma_incomplete], 0, X]
     1360    sage.functions.other.Ei : lambda X : [[max_gamma_incomplete], 0, X],
     1361    sage.functions.log.lambert_w : lambda N,X : [[max_lambert_w], X] if N==EclObject(0) else [[mqapply],[[max_lambert_w, max_array],N],X]
    13591362}
    13601363
    13611364