Repository: scipy/scipy
File: scipy/stats/_morestats.py (Python)

Code before:
    >>> plt.show()

    """
    # _axis_nan_policy decorator does not currently support these for lazy arrays.
    # We want to run tests with lazy backends, so don't pass the arguments explicitly
    # unless necessary.
    kwargs = {}
    if keepdims is not False:
        kwargs['keepdims'] = keepdims
    if nan_policy != 'propagate':
        kwargs['nan_policy'] = nan_policy
    res = _yeojohnson_llf(data, lmb=lmb, axis=axis, **kwargs)
    return res[()] if res.ndim == 0 else res


@_axis_nan_policy_factory(lambda x: x, n_outputs=1, default_axis=0,
                          result_to_tuple=lambda x, _: (x,))
def _yeojohnson_llf(data, *, lmb, axis=0):
    xp = array_namespace(data)
    dtype = xp_result_type(lmb, data, force_floating=True, xp=xp)
    data = xp.asarray(data, dtype=dtype)

    n = data.shape[axis]
    if n == 0:
        return _get_nan(data, xp=xp)
    eps = xp.finfo(dtype).eps
    pos = data >= 0  # binary mask

    if xp.all(pos):
        if abs(lmb) < eps:
            logvar = xp.log(xp.var(xp.log1p(data), axis=axis))
        else:
            logvar = _log_var(lmb * xp.log1p(data), xp, axis) - 2 * math.log(abs(lmb))

    elif xp.all(~pos):
        if abs(lmb - 2) < eps:
            logvar = xp.log(xp.var(xp.log1p(-data), axis=axis))
        else:
            logvar = _log_var((2 - lmb) * xp.log1p(-data), xp, axis) - 2 * math.log(
                abs(2 - lmb)
            )

    else:  # mixed positive and negative data
        dtypemix = xp.result_type(dtype, 1j)
        logyj = xp.zeros_like(data, dtype=dtypemix)

        # x >= 0
        if abs(lmb) < eps:
            logyj = xpx.at(logyj)[pos].set(xp.log(xp.log1p(data[pos]) + 0j))
        else:  # lmbda != 0
            logm1_pos = xp.full_like(data[pos], pi * 1j, dtype=dtypemix)

Review comment (bug):
Does something go wrong when `abs(lmb) < eps`, or is this an optimization for performance (as opposed to `lmb == 0`)?

Code after:
    >>> plt.show()

    """
    # _axis_nan_policy decorator does not currently support these for lazy arrays.
    # We want to run tests with lazy backends, so don't pass the arguments explicitly
    # unless necessary.
    kwargs = {}
    if keepdims is not False:
        kwargs['keepdims'] = keepdims
    if nan_policy != 'propagate':
        kwargs['nan_policy'] = nan_policy
    res = _yeojohnson_llf(data, lmb=lmb, axis=axis, **kwargs)
    return res[()] if res.ndim == 0 else res


@_axis_nan_policy_factory(lambda x: x, n_outputs=1, default_axis=0,
                          result_to_tuple=lambda x, _: (x,))
def _yeojohnson_llf(data, *, lmb, axis=0):
    xp = array_namespace(data)
    dtype = xp_result_type(lmb, data, force_floating=True, xp=xp)
    data = xp.asarray(data, dtype=dtype)

    n = data.shape[axis]
    if n == 0:
        return _get_nan(data, xp=xp)
    eps = xp.finfo(dtype).eps
    # Special case all-positive/negative data to avoid overflow and precision loss
    pos = data >= 0  # binary mask

    # There exists numerical instability when abs(lmb) or abs(lmb - 2) is very small
    if xp.all(pos):
        if abs(lmb) < eps:
            logvar = xp.log(xp.var(xp.log1p(data), axis=axis))
        else:
            logvar = _log_var(lmb * xp.log1p(data), xp, axis) - 2 * math.log(abs(lmb))

    elif xp.all(~pos):
        if abs(lmb - 2) < eps:
            logvar = xp.log(xp.var(xp.log1p(-data), axis=axis))
        else:
            logvar = _log_var((2 - lmb) * xp.log1p(-data), xp, axis) - 2 * math.log(
                abs(2 - lmb)
            )

    # overflow/precision loss not reported for mixed data; calculate `logvar` directly
    else:  # mixed positive and negative data
        y = _yeojohnson_transform(data, lmb, xp=xp)
        sigma = xp.var(y, axis=axis)

        # Suppress RuntimeWarning raised by np.log when the variance is too low
        finite_variance = sigma >= xp.finfo(sigma.dtype).smallest_normal