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

Code before:
    # 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)
            logyj = xpx.at(logyj)[pos].set(
                special.logsumexp([lmb * xp.log1p(data[pos]), logm1_pos], axis=0)
                - xp.log(lmb + 0j)
            )


Review comment (performance):
Presumably the mixed case handles the case of all positive or all negative data, too, so these are performance optimizations?

JAX JIT will not be able to run with data-dependent code like this. Would it be so bad to always use the mixed case code?

Eventually I wanted to vectorize w.r.t. `lmb`, too, so we won't even be able to branch on `lmb`. But refactoring for that can wait for another day.

Code after:
    # 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
        logvar = xpx.apply_where(finite_variance, (sigma,), xp.log, fill_value=-xp.inf)

    loglike = (-n / 2 * logvar
               + (lmb - 1) * xp.sum(xp.sign(data) * xp.log1p(xp.abs(data)), axis=axis))