some new features

.venv/lib/python3.12/site-packages/scipy/special/_lambertw.py

@@ -0,0 +1,149 @@
+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)