Lambert's W function

w_function.py

#!/usr/bin/env python3
import math
# from sys import stderr

def w_function_dx(x):
    return (x + 1) * math.exp(x) 

def w_function(x, x0):
    return x * math.exp(x) - x0

def next_approx(x, x0):
    return x - (w_function(x, x0) / w_function_dx(x))

def W(x, epsilon=1e-16):
    """
    Calculation of W function using Newton-Raphson method
    Quadratic convergence
    https://en.wikipedia.org/wiki/Newton%27s_method

    >>> W(0.0)
    0.0
    >>> W(2.0)
    0.8526055020137254
    >>> W(math.e)
    1.0
    >>> W(3.0)
    1.0499088949640398
    >>> W(2.0 * (math.e ** 2))
    2.0
    >>> W(3.0 * (math.e ** 3))
    3.0
    >>> x = math.exp(W(math.log(25.0)))
    >>> y = math.fabs((x ** x) - 25.0)
    >>> y < 2e-14
    True
    """
    x0 = x
    while True:
        # print(f"{'-' * 30} {x0 = } {x = }", file=stderr)
        next = next_approx(x, x0)
        if abs(next - x) <= epsilon:
            return next
        x = next

if __name__ == '__main__':
    from doctest import testmod
    testmod()