ENH: use minimax approximations for real branches of lambertw

[kandersolar]

This PR is an alternative to #116 based on the method described in [1]. It is much faster than #116, and ~14x faster than the current scipy implementation in my local testing.

I gave the complex and real portions separate detail namespaces. Not sure that's the preferred style.

[1] Toshio Fukushima, Precise and fast computation of Lambert W function by piecewise minimax rational function approximation with variable transformation, 2020. Preprint. https://doi.org/10.13140/RG.2.2.30264.37128

[steppi]

Thanks @kandersolar. This is looking good. I don't think it's necessary to split the implementations into separate detail_real and detail_complex namespaces like this. Usually we just add separate overloads, and detail namespace is reserved for helper functions that we don't want to be part of the public API. The separation I'd suggested before was just a workaround for potential template ambiguity issues that could of occurred, but that won't be a problem here.

[steppi]

Thanks @kandersolar. This is very nice! There is one hiccup though. If we add the real overload directly in the ufunc in SciPy, it will lead to backwards incompatible behavior. Since only a complex in, complex out overload is shipped, one still gets valid results when evaluating non-real branches at real numbers

lambertw(2, k=2)
Out[2]: np.complex128(-1.7022590055576041+10.839808676359299j)

or evaluating real branches on the branch cut

In [3]: lambertw(-3, k=0)
Out[3]: np.complex128(0.46699785792566023+1.8217398230084245j)

there will need to be a deprecation period in order to add the real overload to the lambertw ufunc, and the process will be slightly involved.

In the meantime, to get most of the benefit I think you can add a check to the complex lambertw to see ifk is 0 or -1, the imaginary part of z is exactly zero, and the real part is off the branch cut, then delegate to the real implementation. There will be some overhead incurred compared to shipping the real overload, but at least the underlying calculations will be done with real floats using the fast method when special.lambertw is passed real z.

[steppi]

A few more minor comments.

[kandersolar]

Thanks @steppi for the review! Regarding this comment:

In the meantime, to get most of the benefit I think you can add a check to the complex lambertw ...

I think I see how it would work, but I wonder if it would be tidier overall to ship all overloads in xsf and have scipy.special.lambertw handle the backwards compatibility? For example, instead of this:

https://github.com/scipy/scipy/blob/8a4633fa0e01d62e9ccdd06ebe5bb30551cfa056/scipy/special/_lambertw.py#L149

Have something like this:

w = _lambertw(z, k, tol)
return np.astype(w, np.complex128)  # subject to whatever deprecation machinery is added

[steppi]

Thanks @steppi for the review! Regarding this comment:

In the meantime, to get most of the benefit I think you can add a check to the complex lambertw ...

I think I see how it would work, but I wonder if it would be tidier overall to ship all overloads in xsf and have scipy.special.lambertw handle the backwards compatibility? For example, instead of this:

https://github.com/scipy/scipy/blob/8a4633fa0e01d62e9ccdd06ebe5bb30551cfa056/scipy/special/_lambertw.py#L149

Have something like this:

w = _lambertw(z, k, tol)
return np.astype(w, np.complex128)  # subject to whatever deprecation machinery is added

I think xsf should have all of the overloads in the public API regardless. I'm not sure I understand your suggestion though. The ufunc machinery used in SciPy handles all of the dtype conversions. The reason for needing the deprecation process is that currently there are inputs that lead to non-NaN results which will result in NaN after adding the real overload to SciPy. The process will involve adding a keyword arg to SciPy's lambertw to switch between the current complex-only behavior and the proposed real to real, complex to complex behavior.

Using the nice new real implementation in the complex implementation when relevant seems like it would be a net win regardless of what SciPy does. This PR is already a nice unit of work on its own though, so I could do that in a follow-up.

[steppi]

I think this will be good to merge after running clang-format. I'd like to explore the benefits of using the new implementation in the complex-valued implementation for cases where it maps points on the real line to the real line,, but I don't think it's necessary for you to do in this PR.

[kandersolar]

The reason for needing the deprecation process is that currently there are inputs that lead to non-NaN results which will result in NaN after adding the real overload to SciPy.

Ok, this made it click. I should have read your original comment more carefully. I'll try it out in this PR.

[kandersolar]

Seems like it works!

In [10]: x = np.linspace(0, 1, 1000)

In [11]: %timeit scipy.special.lambertw(x, k=0)
21.7 μs ± 52.9 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

In [12]: x = np.linspace(-5, -4, 1000)

In [13]: %timeit scipy.special.lambertw(x, k=0)
398 μs ± 514 ns per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

[kandersolar]

The test failures are in scipy_special_tests and seem out of my depth. I'd be happy to revert to the "good to merge" state earlier if that makes things easier. Thanks for all the help here :)

[steppi]

The test failures are in scipy_special_tests and seem out of my depth. I'd be happy to revert to the "good to merge" state earlier if that makes things easier. Thanks for all the help here :)

The failures look real and seem worth addressing. I'll look into it.