Python scipy.special.lambertw() Examples

def test_wrightomega_region2():
    # This region gets less coverage in the TestSystematic test
    x = np.linspace(-2, 1)
    y = np.linspace(-2*np.pi, -1)
    x, y = np.meshgrid(x, y)
    z = (x + 1j*y).flatten()

    dataset = []
    for z0 in z:
        dataset.append((z0, complex(_mpmath_wrightomega(z0, 25))))
    dataset = np.asarray(dataset)

    FuncData(sc.wrightomega, dataset, 0, 1, rtol=1e-15).check()


# ------------------------------------------------------------------------------
# lambertw
# ------------------------------------------------------------------------------ 

def prox_max_entropy(X, step, gamma=1, type="relative"):
    """Proximal operator for maximum entropy regularization.

    g(x) = gamma sum_i x_i ln(x_i)

    has the analytical solution of gamma W(1/gamma exp((X-gamma)/gamma)), where
    W is the Lambert W function.

    If type == 'relative', the penalty is expressed in units of the function value;
    if type == 'absolute', it's expressed in units of the variable `X`.
    """
    from scipy.special import lambertw

    assert type in ["relative", "absolute"]
    if type == "relative":
        gamma_ = _step_gamma(step, gamma)
    else:
        gamma_ = gamma
    # minimize entropy: return gamma_ * np.real(lambertw(np.exp((X - gamma_) / gamma_) / gamma_))
    above = X > 0
    X[above] = gamma_ * np.real(lambertw(np.exp(X[above] / gamma_ - 1) / gamma_))
    return X 

def test_ufunc():
    assert_array_almost_equal(
        lambertw(r_[0., e, 1.]), r_[0., 1., 0.567143290409783873]) 

def test_lambertw(self):
        assert_mpmath_equal(lambda x, k: sc.lambertw(x, int(k)),
                            lambda x, k: mpmath.lambertw(x, int(k)),
                            [Arg(), IntArg(0, 10)]) 

def test_lambertw(self):
        # python-list/2010-December/594592.html
        xxroot = fixed_point(lambda xx: np.exp(-2.0*xx)/2.0, 1.0,
                args=(), xtol=1e-12, maxiter=500)
        assert_allclose(xxroot, np.exp(-2.0*xxroot)/2.0)
        assert_allclose(xxroot, lambertw(1)/2) 

def test_ufunc():
    assert_array_almost_equal(
        lambertw(r_[0., e, 1.]), r_[0., 1., 0.567143290409783873]) 

def test_lambertw_ufunc_loop_selection():
    # see https://github.com/scipy/scipy/issues/4895
    dt = np.dtype(np.complex128)
    assert_equal(lambertw(0, 0, 0).dtype, dt)
    assert_equal(lambertw([0], 0, 0).dtype, dt)
    assert_equal(lambertw(0, [0], 0).dtype, dt)
    assert_equal(lambertw(0, 0, [0]).dtype, dt)
    assert_equal(lambertw([0], [0], [0]).dtype, dt) 

def _mpmath_wrightomega(z, dps):
    with mpmath.workdps(dps):
        z = mpmath.mpc(z)
        unwind = mpmath.ceil((z.imag - mpmath.pi)/(2*mpmath.pi))
        res = mpmath.lambertw(mpmath.exp(z), unwind)
    return res 

def test_lambertw_real(self):
        assert_mpmath_equal(lambda x, k: sc.lambertw(x, int(k.real)),
                            lambda x, k: mpmath.lambertw(x, int(k.real)),
                            [ComplexArg(-np.inf, np.inf), IntArg(0, 10)],
                            rtol=1e-13, nan_ok=False) 

def main():
    parser = argparse.ArgumentParser(
        description="Estimates parameter for scheduled sampling.")
    parser.add_argument("--value", type=float, required=True,
                        help="The value the threshold should achieve.")
    parser.add_argument("--step", type=int, required=True,
                        help="Step when you want to achieve the value.")
    args = parser.parse_args()

    x = args.step
    c = args.value

    coeff = c * np.exp(lambertw((1 - c) / c * x)) / (1 - c)

    print(coeff.real) 

def lambertw(*args, **kwargs):
            from scipy.special import lambertw
            return lambertw(*args, **kwargs) 

def erf(*args, **kwargs):
            from scipy.special import erf
            return erf(*args, **kwargs)


#    from scipy.special import lambertw, ellipe, gammaincc, gamma # fluids
#    from scipy.special import i1, i0, k1, k0, iv # ht
#    from scipy.special import hyp2f1    
#    if erf is None:
#        from scipy.special import erf 

def w_d(z, delta):
    # Eq. 9
    if delta < _EPS:
        return z
    return np.sign(z) * np.sqrt(np.real(special.lambertw(delta * z ** 2)) / delta) 

def _lambertw_i_from_v(resistance_shunt, resistance_series, nNsVth, voltage,
                       saturation_current, photocurrent):
    try:
        from scipy.special import lambertw
    except ImportError:
        raise ImportError('This function requires scipy')

    # Record if inputs were all scalar
    output_is_scalar = all(map(np.isscalar,
                               [resistance_shunt, resistance_series, nNsVth,
                                voltage, saturation_current, photocurrent]))

    # This transforms Gsh=1/Rsh, including ideal Rsh=np.inf into Gsh=0., which
    #  is generally more numerically stable
    conductance_shunt = 1. / resistance_shunt

    # Ensure that we are working with read-only views of numpy arrays
    # Turns Series into arrays so that we don't have to worry about
    #  multidimensional broadcasting failing
    Gsh, Rs, a, V, I0, IL = \
        np.broadcast_arrays(conductance_shunt, resistance_series, nNsVth,
                            voltage, saturation_current, photocurrent)

    # Intitalize output I (V might not be float64)
    I = np.full_like(V, np.nan, dtype=np.float64)           # noqa: E741, N806

    # Determine indices where 0 < Rs requires implicit model solution
    idx_p = 0. < Rs

    # Determine indices where 0 = Rs allows explicit model solution
    idx_z = 0. == Rs

    # Explicit solutions where Rs=0
    if np.any(idx_z):
        I[idx_z] = IL[idx_z] - I0[idx_z] * np.expm1(V[idx_z] / a[idx_z]) - \
                   Gsh[idx_z] * V[idx_z]

    # Only compute using LambertW if there are cases with Rs>0
    # Does NOT handle possibility of overflow, github issue 298
    if np.any(idx_p):
        # LambertW argument, cannot be float128, may overflow to np.inf
        argW = Rs[idx_p] * I0[idx_p] / (
                    a[idx_p] * (Rs[idx_p] * Gsh[idx_p] + 1.)) * \
               np.exp((Rs[idx_p] * (IL[idx_p] + I0[idx_p]) + V[idx_p]) /
                      (a[idx_p] * (Rs[idx_p] * Gsh[idx_p] + 1.)))

        # lambertw typically returns complex value with zero imaginary part
        # may overflow to np.inf
        lambertwterm = lambertw(argW).real

        # Eqn. 2 in Jain and Kapoor, 2004
        #  I = -V/(Rs + Rsh) - (a/Rs)*lambertwterm + Rsh*(IL + I0)/(Rs + Rsh)
        # Recast in terms of Gsh=1/Rsh for better numerical stability.
        I[idx_p] = (IL[idx_p] + I0[idx_p] - V[idx_p] * Gsh[idx_p]) / \
                   (Rs[idx_p] * Gsh[idx_p] + 1.) - (
                               a[idx_p] / Rs[idx_p]) * lambertwterm

    if output_is_scalar:
        return I.item()
    else:
        return I