MAINT: stats.boxcox_llf: refactor for simplicity (scipy#22835)

mdhaber · web-flow · commit bb27534ad5ca · 2025-04-14T22:47:07.000+01:00
* MAINT: stats.boxcox_llf: refactor for simplicity

* MAINT: stats.boxcox_llf: compensate for NumPy &lt;2.0
diff --git a/scipy/stats/_morestats.py b/scipy/stats/_morestats.py
@@ -864,15 +864,11 @@ def _log_mean(logx, axis):
 
 def _log_var(logx, xp, axis):
     # compute log of variance of x from log(x)
-    logmean = _log_mean(logx, axis=axis)
-    # get complex dtype with component dtypes same as `logx` dtype;
-    dtype = xp.result_type(logx.dtype, 1j)
-    pij = xp.full(logx.shape, pi * 1j, dtype=dtype)
-    logxmu = special.logsumexp(xp.stack((logx, logmean + pij)), axis=0)
-    return (
-        xp.real(xp.asarray(special.logsumexp(2 * logxmu, axis=axis)))
-        - math.log(logx.shape[axis])
-    )
+    logmean = xp.broadcast_to(_log_mean(logx, axis=axis), logx.shape)
+    ones = xp.ones_like(logx)
+    logxmu, _ = special.logsumexp(xp.stack((logx, logmean), axis=0), axis=0,
+                                  b=xp.stack((ones, -ones), axis=0), return_sign=True)
+    return special.logsumexp(2 * logxmu, axis=axis) - math.log(logx.shape[axis])
 
 
 def boxcox_llf(lmb, data, *, axis=0, keepdims=False, nan_policy='propagate'):
@@ -991,7 +987,7 @@ def boxcox_llf(lmb, data, *, axis=0, keepdims=False, nan_policy='propagate'):
                           result_to_tuple=lambda x: (x,))
 def _boxcox_llf(data, axis=0, *, lmb):
     xp = array_namespace(data)
-    data = xp_promote(data, force_floating=True, xp=xp)
+    lmb, data = xp_promote(lmb, data, force_floating=True, xp=xp)
     N = data.shape[axis]
     if N == 0:
         return _get_nan(data, xp=xp)
@@ -1010,7 +1006,7 @@ def _boxcox_llf(data, axis=0, *, lmb):
         logvar = _log_var(logx, xp, axis) - 2 * math.log(abs(lmb))
 
     res = (lmb - 1) * xp.sum(logdata, axis=axis) - N/2 * logvar
-    res = xp.astype(res, data.dtype)
+    res = xp.astype(res, data.dtype, copy=False)  # compensate for NumPy <2.0
     res = res[()] if res.ndim == 0 else res
     return res
 
diff --git a/scipy/stats/tests/test_morestats.py b/scipy/stats/tests/test_morestats.py
@@ -2049,10 +2049,9 @@ def test_instability_gh20021(self, xp):
     def test_axis(self, xp):
         data = xp.asarray([[100, 200], [300, 400]])
         llf_axis_0 = stats.boxcox_llf(1, data, axis=0)
-        data_axes_swapped = xp.moveaxis(data, 0, -1)
         llf_0 = xp.asarray([
-                stats.boxcox_llf(1, data_axes_swapped[0, :]),
-                stats.boxcox_llf(1, data_axes_swapped[1, :]),
+            stats.boxcox_llf(1, data[:, 0]),
+            stats.boxcox_llf(1, data[:, 1]),
         ])
         xp_assert_close(llf_axis_0, llf_0)
         llf_axis_1 = stats.boxcox_llf(1, data, axis=1)
