150 lines
3.9 KiB
Python
150 lines
3.9 KiB
Python
from ._ufuncs import _lambertw
|
|
|
|
import numpy as np
|
|
|
|
|
|
def lambertw(z, k=0, tol=1e-8):
|
|
r"""
|
|
lambertw(z, k=0, tol=1e-8)
|
|
|
|
Lambert W function.
|
|
|
|
The Lambert W function `W(z)` is defined as the inverse function
|
|
of ``w * exp(w)``. In other words, the value of ``W(z)`` is
|
|
such that ``z = W(z) * exp(W(z))`` for any complex number
|
|
``z``.
|
|
|
|
The Lambert W function is a multivalued function with infinitely
|
|
many branches. Each branch gives a separate solution of the
|
|
equation ``z = w exp(w)``. Here, the branches are indexed by the
|
|
integer `k`.
|
|
|
|
Parameters
|
|
----------
|
|
z : array_like
|
|
Input argument.
|
|
k : int, optional
|
|
Branch index.
|
|
tol : float, optional
|
|
Evaluation tolerance.
|
|
|
|
Returns
|
|
-------
|
|
w : array
|
|
`w` will have the same shape as `z`.
|
|
|
|
See Also
|
|
--------
|
|
wrightomega : the Wright Omega function
|
|
|
|
Notes
|
|
-----
|
|
All branches are supported by `lambertw`:
|
|
|
|
* ``lambertw(z)`` gives the principal solution (branch 0)
|
|
* ``lambertw(z, k)`` gives the solution on branch `k`
|
|
|
|
The Lambert W function has two partially real branches: the
|
|
principal branch (`k = 0`) is real for real ``z > -1/e``, and the
|
|
``k = -1`` branch is real for ``-1/e < z < 0``. All branches except
|
|
``k = 0`` have a logarithmic singularity at ``z = 0``.
|
|
|
|
**Possible issues**
|
|
|
|
The evaluation can become inaccurate very close to the branch point
|
|
at ``-1/e``. In some corner cases, `lambertw` might currently
|
|
fail to converge, or can end up on the wrong branch.
|
|
|
|
**Algorithm**
|
|
|
|
Halley's iteration is used to invert ``w * exp(w)``, using a first-order
|
|
asymptotic approximation (O(log(w)) or `O(w)`) as the initial estimate.
|
|
|
|
The definition, implementation and choice of branches is based on [2]_.
|
|
|
|
References
|
|
----------
|
|
.. [1] https://en.wikipedia.org/wiki/Lambert_W_function
|
|
.. [2] Corless et al, "On the Lambert W function", Adv. Comp. Math. 5
|
|
(1996) 329-359.
|
|
https://cs.uwaterloo.ca/research/tr/1993/03/W.pdf
|
|
|
|
Examples
|
|
--------
|
|
The Lambert W function is the inverse of ``w exp(w)``:
|
|
|
|
>>> import numpy as np
|
|
>>> from scipy.special import lambertw
|
|
>>> w = lambertw(1)
|
|
>>> w
|
|
(0.56714329040978384+0j)
|
|
>>> w * np.exp(w)
|
|
(1.0+0j)
|
|
|
|
Any branch gives a valid inverse:
|
|
|
|
>>> w = lambertw(1, k=3)
|
|
>>> w
|
|
(-2.8535817554090377+17.113535539412148j)
|
|
>>> w*np.exp(w)
|
|
(1.0000000000000002+1.609823385706477e-15j)
|
|
|
|
**Applications to equation-solving**
|
|
|
|
The Lambert W function may be used to solve various kinds of
|
|
equations. We give two examples here.
|
|
|
|
First, the function can be used to solve implicit equations of the
|
|
form
|
|
|
|
:math:`x = a + b e^{c x}`
|
|
|
|
for :math:`x`. We assume :math:`c` is not zero. After a little
|
|
algebra, the equation may be written
|
|
|
|
:math:`z e^z = -b c e^{a c}`
|
|
|
|
where :math:`z = c (a - x)`. :math:`z` may then be expressed using
|
|
the Lambert W function
|
|
|
|
:math:`z = W(-b c e^{a c})`
|
|
|
|
giving
|
|
|
|
:math:`x = a - W(-b c e^{a c})/c`
|
|
|
|
For example,
|
|
|
|
>>> a = 3
|
|
>>> b = 2
|
|
>>> c = -0.5
|
|
|
|
The solution to :math:`x = a + b e^{c x}` is:
|
|
|
|
>>> x = a - lambertw(-b*c*np.exp(a*c))/c
|
|
>>> x
|
|
(3.3707498368978794+0j)
|
|
|
|
Verify that it solves the equation:
|
|
|
|
>>> a + b*np.exp(c*x)
|
|
(3.37074983689788+0j)
|
|
|
|
The Lambert W function may also be used find the value of the infinite
|
|
power tower :math:`z^{z^{z^{\ldots}}}`:
|
|
|
|
>>> def tower(z, n):
|
|
... if n == 0:
|
|
... return z
|
|
... return z ** tower(z, n-1)
|
|
...
|
|
>>> tower(0.5, 100)
|
|
0.641185744504986
|
|
>>> -lambertw(-np.log(0.5)) / np.log(0.5)
|
|
(0.64118574450498589+0j)
|
|
"""
|
|
# TODO: special expert should inspect this
|
|
# interception; better place to do it?
|
|
k = np.asarray(k, dtype=np.dtype("long"))
|
|
return _lambertw(z, k, tol)
|