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

Code before:
    >>> ax.set_ylabel('Yeo-Johnson log-likelihood')

    Now add some probability plots to show that where the log-likelihood is
    maximized the data transformed with `yeojohnson` looks closest to normal:

    >>> locs = [3, 10, 4]  # 'lower left', 'center', 'lower right'
    >>> for lmbda, loc in zip([-1, lmbda_optimal, 9], locs):
    ...     xt = stats.yeojohnson(x, lmbda=lmbda)
    ...     (osm, osr), (slope, intercept, r_sq) = stats.probplot(xt)
    ...     ax_inset = inset_axes(ax, width="20%", height="20%", loc=loc)
    ...     ax_inset.plot(osm, osr, 'c.', osm, slope*osm + intercept, 'k-')
    ...     ax_inset.set_xticklabels([])
    ...     ax_inset.set_yticklabels([])
    ...     ax_inset.set_title(r'$\lambda=%1.2f$' % lmbda)

    >>> 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)
    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
    log_sigma = xpx.apply_where(finite_variance, (sigma,), xp.log, fill_value=-xp.inf)

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

    return loglike


@xp_capabilities(np_only=True)
def yeojohnson_normmax(x, brack=None):
    """Compute optimal Yeo-Johnson transform parameter.

Review comment (suggestion):
Should the keys not be strings?
```suggestion
        kwargs['keepdims'] = keepdims
    if nan_policy != 'propagate':
        kwargs['nan_policy'] = nan_policy
```

Code after:
    >>> ax.set_ylabel('Yeo-Johnson log-likelihood')

    Now add some probability plots to show that where the log-likelihood is
    maximized the data transformed with `yeojohnson` looks closest to normal:

    >>> locs = [3, 10, 4]  # 'lower left', 'center', 'lower right'
    >>> for lmbda, loc in zip([-1, lmbda_optimal, 9], locs):
    ...     xt = stats.yeojohnson(x, lmbda=lmbda)
    ...     (osm, osr), (slope, intercept, r_sq) = stats.probplot(xt)
    ...     ax_inset = inset_axes(ax, width="20%", height="20%", loc=loc)
    ...     ax_inset.plot(osm, osr, 'c.', osm, slope*osm + intercept, 'k-')
    ...     ax_inset.set_xticklabels([])
    ...     ax_inset.set_yticklabels([])
    ...     ax_inset.set_title(r'$\lambda=%1.2f$' % lmbda)

    >>> 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)
    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
    log_sigma = xpx.apply_where(finite_variance, (sigma,), xp.log, fill_value=-xp.inf)

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

    return loglike


@xp_capabilities(np_only=True)
def yeojohnson_normmax(x, brack=None):
    """Compute optimal Yeo-Johnson transform parameter.