# 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 real_part, real, imag_part, imag, imaginary, conjugate) from log import (exp, log, ln, polylog, dilog) from log import (exp, log, ln, polylog, dilog, lambert_w, lambert_w_branch) from transcendental import (exponential_integral_1,
• ## sage/functions/log.py

diff --git a/sage/functions/log.py b/sage/functions/log.py
 a """ Logarithmic functions """ from sage.symbolic.function import GinacFunction from sage.symbolic.function import GinacFunction, BuiltinFunction, is_inexact from sage.libs.mpmath import utils as mpmath_utils from sage.structure.coerce import parent from sage.symbolic.expression import Expression from sage.rings.real_double import RDF from sage.rings.complex_double import CDF class Function_exp(GinacFunction): def __init__(self): dilog = Function_dilog() class Function_lambert_w_branch(BuiltinFunction): r""" The integral branches of the Lambert W function W_n(z). This function satisfies the equation .. math:: z = W_n(z) e^{W_n(z)} INPUT: - z - a complex number - n - an integer. n=0 corresponds to the principle branch. ALGORITHM: Numerical evaluation is handled using the mpmath library. REFERENCES: - http://en.wikipedia.org/wiki/Lambert_W_function EXAMPLES:: sage: lambert_w(1.0) 0.567143290409784 sage: lambert_w(-1).n() -0.318131505204764 + 1.33723570143069*I sage: lambert_w(-1.5 + 5*I) 1.17418016254171 + 1.10651494102011*I sage: lambert_w(RealField(100)(1)) 0.56714329040978387299996866221 sage: S = solve(e^(5*x)+x==0, x, to_poly_solve=True) sage: z = S[0].rhs(); z -1/5*lambert_w(5) sage: N(z) -0.265344933048440 Check the defining equation numerically at z=5:: sage: N(lambert_w(5)*exp(lambert_w(5)) - 5) 0.000000000000000 """ def __init__(self): """ See the docstring for :meth:Function_lambert_w. EXAMPLES:: sage: lambert_w(1.0) 0.567143290409784 """ BuiltinFunction.__init__(self, "lambert_w_branch", nargs=2, latex_name=r'W') def _eval_(self, z, n): """ EXAMPLES:: sage: lambert_w(0) lambert_w(0) sage: x = var('x') sage: lambert_w(x) lambert_w(x) sage: lambert_w(0.0) 0.000000000000000 """ if not isinstance(z, Expression) and is_inexact(z): return self._evalf_(z, n, parent=parent(z)) return None def _evalf_(self, z, n, parent=None): """ EXAMPLES:: sage: N(lambert_w(1)) 0.567143290409784 sage: lambert_w(RealField(100)(1)) 0.56714329040978387299996866221 """ if parent(z) == RDF or parent(z) == CDF: import scipy.special return scipy.special.lambertw(z, n) else: import mpmath return mpmath_utils.call(mpmath.lambertw, z, n, parent=parent) def _derivative_(self, z, n, diff_param=None): """ The derivative of W_n(x) is W_n(x)/(x \cdot W_n(x) + x). EXAMPLES:: sage: x = var('x') sage: derivative(lambert_w(x), x) lambert_w(x)/(x*lambert_w(x) + x) """ return lambert_w(z, n)/(z*lambert_w(z, n)+z) lambert_w_branch = Function_lambert_w_branch() def lambert_w(z, *args, **kwds): """ The Lambert w function. This is a wrapper for lambert_w_branch. """ if not args: return lambert_w_branch(z, 0, **kwds) elif len(args) > 1: raise ArgumentError("lambert_w takes at most two arguments.") else: 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 sage: from sage.symbolic.random_tests import _mk_full_functions sage: [f for (one,f,arity) in _mk_full_functions()] [Ei, abs, arccos, arccosh, arccot, arccoth, arccsc, arccsch, arcsec, arcsech, arcsin, arcsinh, arctan, arctan2, arctanh, binomial, ceil, conjugate, cos, cosh, cot, coth, csc, csch, dickman_rho, dilog, dirac_delta, elliptic_e, elliptic_ec, elliptic_eu, elliptic_f, elliptic_kc, elliptic_pi, erf, exp, factorial, floor, heaviside, imag_part, integrate, kronecker_delta, log, polylog, real_part, sec, sech, sgn, sin, sinh, tan, tanh, unit_step, zeta, zetaderiv] [Ei, abs, arccos, arccosh, arccot, arccoth, arccsc, arccsch, arcsec, arcsech, arcsin, arcsinh, arctan, arctan2, arctanh, binomial, ceil, conjugate, cos, cosh, cot, coth, csc, csch, dickman_rho, dilog, dirac_delta, elliptic_e, elliptic_ec, elliptic_eu, elliptic_f, elliptic_kc, elliptic_pi, erf, exp, factorial, floor, heaviside, imag_part, integrate, kronecker_delta, lambert_w, log, polylog, real_part, sec, sech, sgn, sin, sinh, tan, tanh, unit_step, zeta, zetaderiv] Note that this doctest will fail whenever a Pynac function is added or removed.  In that case, it is very likely that the doctests for sage: from sage.symbolic.random_tests import * sage: set_random_seed(2) sage: random_expr(50, nvars=3, coeff_generator=CDF.random_element) (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) -v1*elliptic_kc((0.0666829501658 + 0.206976992303*I)/(v3 + e))/v3 + heaviside(arccosh(-(abs(v2 - floor(-v3))^(-real_part(elliptic_e(-0.703991792631 + 0.750156228797*I, -(1.21734510331 - 1.22580558833*I)*pi*v1)))*e^(arctan2(imag_part(v2) - imag_part(floor(-v3)), real_part(v2) - real_part(floor(-v3)))*imag_part(elliptic_e(-0.703991792631 + 0.750156228797*I, -(1.21734510331 - 1.22580558833*I)*pi*v1)))*sin(-log(abs(v2 - floor(-v3)))*imag_part(elliptic_e(-0.703991792631 + 0.750156228797*I, -(1.21734510331 - 1.22580558833*I)*pi*v1)) - arctan2(imag_part(v2) - imag_part(floor(-v3)), real_part(v2) - real_part(floor(-v3)))*real_part(elliptic_e(-0.703991792631 + 0.750156228797*I, -(1.21734510331 - 1.22580558833*I)*pi*v1)))*imag_part(arccsc((0.891138386848 - 0.0936820840629*I)/v1)) - abs(v2 - floor(-v3))^(-real_part(elliptic_e(-0.703991792631 + 0.750156228797*I, -(1.21734510331 - 1.22580558833*I)*pi*v1)))*e^(arctan2(imag_part(v2) - imag_part(floor(-v3)), real_part(v2) - real_part(floor(-v3)))*imag_part(elliptic_e(-0.703991792631 + 0.750156228797*I, -(1.21734510331 - 1.22580558833*I)*pi*v1)))*cos(-log(abs(v2 - floor(-v3)))*imag_part(elliptic_e(-0.703991792631 + 0.750156228797*I, -(1.21734510331 - 1.22580558833*I)*pi*v1)) - arctan2(imag_part(v2) - imag_part(floor(-v3)), real_part(v2) - real_part(floor(-v3)))*real_part(elliptic_e(-0.703991792631 + 0.750156228797*I, -(1.21734510331 - 1.22580558833*I)*pi*v1)))*real_part(arccsc((0.891138386848 - 0.0936820840629*I)/v1)) + v3^(-0.48519994364 - 0.485764091302*I) + 0.0909404921682*real_part(v3)/(real_part(v3)^2 + imag_part(v3)^2) + 0.281538203756*imag_part(v3)/(real_part(v3)^2 + imag_part(v3)^2) - 0.647983235144*real_part(v1) - 1.20665952957*imag_part(v1))*v1)) sage: random_expr(5, verbose=True) About to apply dirac_delta to [1] About to apply arccsch to [0] About to apply to [0, arccsch(0)] arccsch(0) About to apply to [v1, e] About to apply to [v1*e, 1] v1*e """ vars = [(1.0, sage.calculus.calculus.var('v%d' % (n+1))) for n in range(nvars)] if ncoeffs is None: