Explore the relative error in `scipy.special.lambertw`

lambertw_error.py

from __future__ import division, print_function, absolute_import

from time import time
from multiprocessing import Pool

import mpmath
import numpy as np
import matplotlib.pyplot as plt
import scipy
import scipy.special as sc


def mpmath_lambertw(z):
    return complex(mpmath.lambertw(z))


def main():
    print("scipy.__version__ = {}".format(scipy.__version__))

    t0 = time()

    x = np.linspace(-0.75, -0.55, 400)
    y = np.linspace(-0.05, 0.05, 400)
    x, y = np.meshgrid(x, y)
    z = x + 1j*y

    p = Pool(8)
    lambertw_std = p.map(mpmath_lambertw, z.flatten())
    lambertw_std = np.array(lambertw_std).reshape(x.shape)

    lambertw = sc.lambertw(z)
    err = abs((lambertw - lambertw_std)/lambertw_std)
    print("Max error is {}".format(np.max(err)))
    print("Elapsed time was {}".format(time() - t0))

    plt.pcolormesh(x, y, err)
    plt.colorbar()
    plt.title("Relative error")
    plt.show()


if __name__ == '__main__':
    main()