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

Author: cwhanse

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

@@ -26,9 +26,6 @@ Enhancements
 * 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`)
-* Add :py:func:`pvlib.tools._lambertw_pvlib` to speed up calculations when using
-  LambertW single diode equation methods.
-  (:discuss:`2720`, :pull:`2723`)  
 
 Documentation
 ~~~~~~~~~~~~~

pvlib/ivtools/utils.py

@@ -552,7 +552,6 @@ def _log_lambertw(logx):
     Parameters
     ----------
     logx : numeric
-        Log(x) of
 
     Returns
     -------
@@ -571,29 +570,30 @@ def _log_lambertw(logx):
 
 
 def _lambertw_pvlib(x):
-    r'''Lambert's W function principal branch, :math:`W_0(x)`, for :math:`x`
-    real valued.
+    r'''Lambert's W function principal branch, :math:`W_0(x)`, for
+    :math:`x>=0`.
 
     Parameters
     ----------
-    x : np.array
+    x : float or np.array
         Must be real numbers.
 
     Returns
     -------
-    np.array
+    float or np.array
 
     '''
-    w = np.full_like(x, np.nan)
-    small = x <= 10
+    localx = np.asarray(x, float)
+    w = np.full_like(localx, np.nan)
+    small = localx <= 10
     # for large x, solve 0 = f(w) = w + log(w) - log(x) using Newton's
-    w[~small] = _log_lambertw(np.log(x[~small]))
-
     # 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 any(small):
-        z = x[small]
-        temp = np.log(x[small] + 1)
+    if np.any(small):
+        z = localx[small]
+        temp = np.log(localx[small] + 1)
         g = temp - np.log(temp + 1)
         for _ in range(0, 3):
             expg = np.exp(g)
@@ -603,4 +603,4 @@ def _lambertw_pvlib(x):
                 (expg * g_p1**2 - 0.5*(g + 2)*g_expg_z)
         w[small] = g
 
-    return w
+    return w[0] if w.shape==1 else w

tests/ivtools/test_utils.py

@@ -175,26 +175,20 @@ def test_astm_e1036_fit_points(v_array, i_array):
 
 
 def test_lambertw_pvlib():
-    test_exp = np.arange(-10., 300, step=10)
-    test_x = 10.**test_exp
+    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([
-        9.9999999989999997e-11, 5.6714329040978384e-01,
-        2.0028685413304952e+01, 4.2306755091738395e+01,
-        6.4904633770046118e+01, 8.7630277151947183e+01,
-        1.1042491882731335e+02, 1.3326278259180333e+02,
-        1.5613026581351718e+02, 1.7901931374150624e+02,
-        2.0192476320084489e+02, 2.2484310644511851e+02,
-        2.4777185185877809e+02, 2.7070916610249782e+02,
-        2.9365366103997610e+02, 3.1660426041503479e+02,
-        3.3956011295458728e+02, 3.6252053376149752e+02,
-        3.8548496362161768e+02, 4.0845294003314166e+02,
-        4.3142407612718210e+02, 4.5439804503371403e+02,
-        4.7737456808796901e+02, 5.0035340579834485e+02,
-        5.2333435083468805e+02, 5.4631722251791496e+02,
-        5.6930186244110166e+02, 5.9228813095427859e+02,
-        6.1527590431628334e+02, 6.3826507236734335e+02,
-        6.6125553661218726e+02])
+        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, k in zip([1.e-10, 1.e+10], [1, 5]):
+        result = _lambertw_pvlib(x)
+        assert np.isclose(result, expected[k])
+    # with 1d array
+    for x, k in zip([1.e-10, 1.e+10], [1, 5]):
+        result = _lambertw_pvlib(np.array([x]))
+        assert np.isclose(result, expected[k])