Ticket #11888: trac_11888_v6.patch

File trac_11888_v6.patch, 8.2 KB (added by benjaminfjones, 8 years ago)

added custom latex printing and Maxima initialization

  • sage/functions/all.py

    # HG changeset patch
    # User Benjamin Jones <benjaminfjones@gmail.com>
    # Date 1328501390 21600
    # Node ID 55066e5e5e1307090f7d4bc39460948d05246a51
    # 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={'mathematica':'ProductLog',
     526                                              'maple':'LambertW'})
     527
     528    def __call__(self, *args, **kwds):
     529        r"""
     530        Custom call method allows the user to pass one argument or two. If
     531        one argument is passed, we assume it is ``z`` and that ``n=0``.
     532
     533        EXAMPLES::
     534
     535            sage: lambert_w(1)
     536            lambert_w(1)
     537            sage: lambert_w(1, 2)
     538            lambert_w(1, 2)
     539        """
     540        if len(args) == 2:
     541            return BuiltinFunction.__call__(self, *args, **kwds)
     542        elif len(args) == 1:
     543            return BuiltinFunction.__call__(self, 0, args[0], **kwds)
     544        else:
     545            raise TypeError("lambert_w takes either one or two arguments.")
     546
     547    def _eval_(self, n, z):
     548        """
     549        EXAMPLES::
     550
     551            sage: lambert_w(6.0)
     552            1.43240477589830
     553            sage: lambert_w(1)
     554            lambert_w(1)
     555            sage: lambert_w(x+1)
     556            lambert_w(x + 1)
     557
     558        There are three special values which are automatically simplified::
     559
     560            sage: lambert_w(0)
     561            0
     562            sage: lambert_w(e)
     563            1
     564            sage: lambert_w(-1/e)
     565            -1
     566            sage: lambert_w(SR(-1/e))
     567            -1
     568            sage: lambert_w(SR(0))   
     569            0
     570
     571        The special values only hold on the principal branch::
     572
     573            sage: lambert_w(1,e)
     574            lambert_w(1, e)
     575            sage: lambert_w(1, e.n())
     576            -0.532092121986380 + 4.59715801330257*I
     577        """
     578        if not isinstance(z, Expression):
     579            if is_inexact(z):
     580                return self._evalf_(n, z, parent=sage_structure_coerce_parent(z))
     581            elif n == 0 and z == 0:
     582                return sage_structure_coerce_parent(z)(0)
     583        elif n == 0:
     584            if z.is_trivial_zero():
     585                return sage_structure_coerce_parent(z)(0)
     586            elif (z-const_e).is_trivial_zero():
     587                return sage_structure_coerce_parent(z)(1)
     588            elif (z+1/const_e).is_trivial_zero():
     589                return sage_structure_coerce_parent(z)(-1)
     590        return None
     591
     592    def _evalf_(self, n, z, parent=None):
     593        """
     594        EXAMPLES::
     595
     596            sage: N(lambert_w(1))
     597            0.567143290409784
     598            sage: lambert_w(RealField(100)(1))
     599            0.56714329040978387299996866221
     600
     601        SciPy is used to evaluate for float, RDF, and CDF inputs::
     602
     603            sage: lambert_w(RDF(1))
     604            0.56714329041
     605        """
     606        R = parent or sage_structure_coerce_parent(z)
     607        if R is float or R is complex or R is RDF or R is CDF:
     608            import scipy.special
     609            return scipy.special.lambertw(z, n)
     610        else:
     611            import mpmath
     612            return mpmath_utils.call(mpmath.lambertw, z, n, parent=R)
     613
     614    def _derivative_(self, n, z, diff_param=None):
     615        """
     616        The derivative of `W_n(x)` is `W_n(x)/(x \cdot W_n(x) + x)`.
     617
     618        EXAMPLES::
     619
     620            sage: x = var('x')
     621            sage: derivative(lambert_w(x), x)
     622            lambert_w(x)/(x*lambert_w(x) + x)
     623        """
     624        return lambert_w(n, z)/(z*lambert_w(n, z)+z)
     625
     626    def _maxima_init_evaled_(self, n, z):
     627        """
     628        EXAMPLES:
     629
     630        These are indirect doctests for this function.::
     631
     632            sage: lambert_w(0, x)._maxima_()
     633            lambert_w(x)
     634            sage: lambert_w(1, x)._maxima_()
     635            ...
     636        """
     637        if n == 0:
     638            return "lambert_w(%s)" % z
     639        else:
     640            raise NotImplementedError("Non-principal branch lambert_w[%s](%s) is not implemented in Maxima" % (n, z))
     641
     642
     643    def _print_(self, n, z):
     644        """
     645        Custom _print_ method to avoid printing the branch number if
     646        it is zero.
     647
     648        EXAMPLES::
     649
     650            sage: lambert_w(1)
     651            lambert_w(1)
     652            sage: lambert_w(0,x)
     653            lambert_w(x)
     654        """
     655        if n == 0:
     656            return "lambert_w(%s)" % z
     657        else:
     658            return "lambert_w(%s, %s)" % (n,z)
     659
     660    def _print_latex_(self, n, z):
     661        """
     662        Custom _print_latex_ method to avoid printing the branch
     663        number if it is zero.
     664
     665        EXAMPLES::
     666
     667            sage: latex(lambert_w(1))
     668            \operatorname{W_0}(1)
     669            sage: latex(lambert_w(0,x))
     670            \operatorname{W_0}(x)
     671            sage: latex(lambert_w(1,x))
     672            \operatorname{W_{1}}(x)
     673        """
     674        if n == 0:
     675            return r"\operatorname{W_0}(%s)" % z
     676        else:
     677            return r"\operatorname{W_{%s}}(%s)" % (n,z)
     678
     679lambert_w = Function_lambert_w()