Add lambertw_pvlib (#2723)

* add lambertw_pvlib

* whatsnew

* formatting

* Apply suggestions from code review

Co-authored-by: Echedey Luis <80125792+echedey-ls@users.noreply.github.com>

* from review

* move to ivtools, make private

* Apply suggestions from code review

Co-authored-by: Kevin Anderson <kevin.anderso@gmail.com>

* accept float, restrict to x>=0, simplify test

* formatting

* small < 100

* Apply suggestions from code review

Co-authored-by: Echedey Luis <80125792+echedey-ls@users.noreply.github.com>

* fix test loop

---------

Co-authored-by: Echedey Luis <80125792+echedey-ls@users.noreply.github.com>
Co-authored-by: Kevin Anderson <kevin.anderso@gmail.com>

Author: 3 people

docs/sphinx/source/whatsnew/v0.15.1.rst

@@ -30,7 +30,8 @@ Bug fixes
 
 Enhancements
 ~~~~~~~~~~~~
-* Use ``k`` and ``cap_adjustment`` from :py:func:`pvlib.pvsystem.Array.module_parameters` in :py:func:`pvlib.pvsystem.PVSystem.pvwatts_dc`
+* Use ``k`` and ``cap_adjustment`` from :py:func:`pvlib.pvsystem.Array.module_parameters`
+  in :py:func:`pvlib.pvsystem.PVSystem.pvwatts_dc`
   (:issue:`2714`, :pull:`2715`)
 * Include `ross` and `faiman_rad` in the allowed models within 
   :py:meth:`pvlib.pvsystem.PVSystem.get_cell_temperature` (:issue:`2625`, :pull:`2631`)  

pvlib/ivtools/utils.py

@@ -544,3 +544,63 @@ def astm_e1036(v, i, imax_limits=(0.75, 1.15), vmax_limits=(0.75, 1.15),
     result['mp_fit'] = mp_fit
 
     return result
+
+
+def _log_lambertw(logx):
+    r'''Computes W(x) starting from log(x).
+
+    Parameters
+    ----------
+    logx : numeric
+
+    Returns
+    -------
+    numeric
+        Lambert's W(x)
+
+    '''
+    # handles overflow cases, but results in nan for x <= 1
+    w = logx - np.log(logx)  # initial guess, w = log(x) - log(log(x))
+
+    for _ in range(0, 3):
+        # Newton's. Halley's is not substantially faster or more accurate
+        # because f''(w) = -1 / (w**2) is small for large w
+        w = w * (1. - np.log(w) + logx) / (1. + w)
+    return w
+
+
+def _lambertw_pvlib(x):
+    r'''Lambert's W function principal branch, :math:`W_0(x)`, for
+    :math:`x>=0`.
+
+    Parameters
+    ----------
+    x : float or np.array
+        Must be real numbers.
+
+    Returns
+    -------
+    float or np.array
+
+    '''
+    localx = np.asarray(x, float)
+    w = np.full_like(localx, np.nan)
+    small = localx <= 100
+    # for large x, solve 0 = f(w) = w + log(w) - log(x) using Newton's
+    # w will contain nan for these numbers due to log(w) = log(log(x))
+    w[~small] = _log_lambertw(np.log(localx[~small]))
+
+    # for small x, solve 0 = g(w) = w * exp(w) - x using Halley's method
+    if np.any(small):
+        z = localx[small]
+        temp = np.log1p(localx[small])
+        g = temp - np.log1p(temp)
+        for _ in range(0, 3):
+            expg = np.exp(g)
+            g_expg_z = g*expg - z
+            g_p1 = g + 1
+            g = g - g_expg_z * g_p1 / \
+                (expg * g_p1**2 - 0.5*(g + 2)*g_expg_z)
+        w[small] = g
+
+    return w[0] if w.shape == 1 else w

tests/ivtools/test_utils.py

@@ -3,6 +3,7 @@
 import pytest
 from pvlib.ivtools.utils import _numdiff, rectify_iv_curve, astm_e1036
 from pvlib.ivtools.utils import _schumaker_qspline
+from pvlib.ivtools.utils import _lambertw_pvlib
 
 from tests.conftest import TESTS_DATA_DIR
 
@@ -171,3 +172,23 @@ def test_astm_e1036_fit_points(v_array, i_array):
                 'ff': 0.7520255886236707}
     result.pop('mp_fit')
     assert result == pytest.approx(expected)
+
+
+def test_lambertw_pvlib():
+    test_x = np.array([0., 1.e-10, 1., 10., 100., 1.e+10, 1.e+100, 1.e+300])
+    # known solution from scipy.special.lambertw
+    # scipy 1.7.1, python 3.13.1, numpy 2.3.5
+    expected = np.array([
+        0.0000000000000000e+00, 9.9999999989999997e-11, 5.6714329040978384e-01,
+        1.7455280027406994e+00, 3.3856301402900502e+00, 2.0028685413304952e+01,
+        2.2484310644511851e+02, 6.8424720862976085e+02])
+    result = _lambertw_pvlib(test_x)
+    assert np.allclose(result, expected, rtol=1e-14)
+    # with float input
+    for x, known in zip(test_x[[1, 5]], expected[[1, 5]]):
+        result = _lambertw_pvlib(x)
+        assert np.isclose(result, known)
+    # with 1d array
+    for x, known in zip(test_x[[1, 5]], expected[[1, 5]]):
+        result = _lambertw_pvlib(np.array([x]))
+        assert np.isclose(result, known)