Repository: scipy/scipy
Pull request: MAINT: special.logsumexp: preserve tiny imaginary components
#### Reference issue
Addresses https://github.com/scipy/scipy/pull/21622#discussion_r1781742660
Closes gh-21709

#### What does this implement/fix?
https://github.com/scipy/scipy/pull/21622#discussion_r1781742660 noted that `logsumexp` woule lose precision for tiny imaginary components; gh-21709 reported it as a bug. This fixes that.

Problem statement:
BUG: logsumexp returning incorrect results in Scipy 1.15.0.dev0
### Describe your issue.

`logsumexp` appears to have broken between 1.14.1 and 1.15.0dev. The difference is in how the complex element is handled. This is producing failures in statsmodels testing, and will in real-world code that uses approximate numerical derivatives.


### Reproducing Code Example

```python
import numpy as np
from scipy.special import logsumexp

np.set_printoptions(precision=24)
y = np.array([[[0.+2.2204460492503132e-17j]]])
tmp = np.array([[[0.+0.0000000000000000e+00j]],
       [[0.+2.2204460492503132e-17j]]])
print("logsumexp")
print(logsumexp(tmp, axis=0))
print("y - logsumexp")
print(y - logsumexp(tmp, axis=0))
print("exp(y - logsumexp)")
print(np.exp(y - logsumexp(tmp, axis=0)))
```


### Error message

```shell
On the released version, this returns
logsumexp
[[0.6931471805599453+1.1102230246251566e-17j]]
y - logsumexp
[[[-0.6931471805599453+1.1102230246251566e-17j]]]
exp(y - logsumexp)
[[[0.5+5.551115123125783e-18j]]]


On dev this returns

logsumexp
[[0.6931471805599453+0.j]]
y - logsumexp
[[[-0.6931471805599453+2.2204460492503132e-17j]]]
exp(y - logsumexp)
[[[0.5+1.1102230246251566e-17j]]]
```


### SciPy/NumPy/Python version and system information

```shell
release
1.14.1 2.1.2 sys.version_info(major=3, minor=11, micro=0, releaselevel='final', serial=0)
Build Dependencies:
  blas:
    detection method: pkgconfig
    found: true
    include directory: C:/Users/runneradmin/AppData/Local/Temp/cibw-run-rrcld3x_/cp311-win_amd64/build/venv/Lib/site-packages/scipy_openblas32/include
    lib directory: C:/Users/runneradmin/AppData/Local/Temp/cibw-run-rrcld3x_/cp311-win_amd64/build/venv/Lib/site-packages/scipy_openblas32/lib
    name: scipy-openblas
    openblas configuration: OpenBLAS 0.3.27.dev DYNAMIC_ARCH NO_AFFINITY Zen MAX_THREADS=24
    pc file directory: D:/a/scipy/scipy
    version: 0.3.27.dev
  lapack:
    detection method: pkgconfig
    found: true
    include directory: C:/Users/runneradmin/AppData/Local/Temp/cibw-run-rrcld3x_/cp311-win_amd64/build/venv/Lib/site-packages/scipy_openblas32/include
    lib directory: C:/Users/runneradmin/AppData/Local/Temp/cibw-run-rrcld3x_/cp311-win_amd64/build/venv/Lib/site-packages/scipy_openblas32/lib
    name: scipy-openblas
    openblas configuration: OpenBLAS 0.3.27.dev DYNAMIC_ARCH NO_AFFINITY Zen MAX_THREADS=24
    pc file directory: D:/a/scipy/scipy
    version: 0.3.27.dev
  pybind11:
    detection method: config-tool
    include directory: unknown
    name: pybind11
    version: 2.12.0
Compilers:
  c:
    commands: cc
    linker: ld.bfd
    name: gcc
    version: 10.3.0
  c++:
    commands: c++
    linker: ld.bfd
    name: gcc
    version: 10.3.0
  cython:
    commands: cython
    linker: cython
    name: cython
    version: 3.0.11
  fortran:
    commands: gfortran
    linker: ld.bfd
    name: gcc
    version: 10.3.0
  pythran:
    include directory: C:\Users\runneradmin\AppData\Local\Temp\pip-build-env-utnpaygc\overlay\Lib\site-packages/pythran
    version: 0.16.1
Machine Information:
  build:
    cpu: x86_64
    endian: little
    family: x86_64
    system: windows
  cross-compiled: false
  host:
    cpu: x86_64
    endian: little
    family: x86_64
    system: windows
Python Information:
  path: C:\Users\runneradmin\AppData\Local\Temp\cibw-run-rrcld3x_\cp311-win_amd64\build\venv\Scripts\python.exe
  version: '3.11'

dev

1.15.0.dev0+git20241010.0396bdd 2.0.2 sys.version_info(major=3, minor=12, micro=7, releaselevel='final', serial=0)
Build Dependencies:
  blas:
    detection method: pkgconfig
    found: true
    include directory: C:/Users/runneradmin/AppData/Local/Temp/cibw-run-ca9fpofl/cp312-win_amd64/build/venv/Lib/site-packages/scipy_openblas32/include
    lib directory: C:/Users/runneradmin/AppData/Local/Temp/cibw-run-ca9fpofl/cp312-win_amd64/build/venv/Lib/site-packages/scipy_openblas32/lib
    name: scipy-openblas
    openblas configuration: OpenBLAS 0.3.28 DYNAMIC_ARCH NO_AFFINITY Haswell MAX_THREADS=24
    pc file directory: D:/a/scipy/scipy
    version: 0.3.28
  lapack:
    detection method: pkgconfig
    found: true
    include directory: C:/Users/runneradmin/AppData/Local/Temp/cibw-run-ca9fpofl/cp312-win_amd64/build/venv/Lib/site-packages/scipy_openblas32/include
    lib directory: C:/Users/runneradmin/AppData/Local/Temp/cibw-run-ca9fpofl/cp312-win_amd64/build/venv/Lib/site-packages/scipy_openblas32/lib
    name: scipy-openblas
    openblas configuration: OpenBLAS 0.3.28 DYNAMIC_ARCH NO_AFFINITY Haswell MAX_THREADS=24
    pc file directory: D:/a/scipy/scipy
    version: 0.3.28
  pybind11:
    detection method: config-tool
    include directory: unknown
    name: pybind11
    version: 2.13.6
Compilers:
  c:
    commands: cc
    linker: ld.bfd
    name: gcc
    version: 10.3.0
  c++:
    commands: c++
    linker: ld.bfd
    name: gcc
    version: 10.3.0
  cython:
    commands: cython
    linker: cython
    name: cython
    version: 3.0.11
  fortran:
    commands: gfortran
    linker: ld.bfd
    name: gcc
    version: 10.3.0
  pythran:
    include directory: C:\Users\runneradmin\AppData\Local\Temp\pip-build-env-aypk19gl\overlay\Lib\site-packages/pythran
    version: 0.16.1
Machine Information:
  build:
    cpu: x86_64
    endian: little
    family: x86_64
    system: windows
  cross-compiled: false
  host:
    cpu: x86_64
    endian: little
    family: x86_64
    system: windows
Python Information:
  path: C:\Users\runneradmin\AppData\Local\Temp\cibw-run-ca9fpofl\cp312-win_amd64\build\venv\Scripts\python.exe
  version: '3.12'
```

Hints from issue discussion:
https://github.com/scipy/scipy/pull/21622 perhaps?
I woudl guess nightly with that PR would be out already?

Still seeing the failure as of this morning.
IIRC you want the relative precision for small imaginary components to stay as it was?

> #21622 perhaps?

Yes, could be that, specifically https://github.com/scipy/scipy/pull/21622#discussion_r1781742660. *Yes; the initial calculation appears to preserve the precision, but it is lost when ensuring that enforcing the branch cut convention. A solution is to leave the imaginary component alone if it is already within bounds.*
@mdhaber Yes - in this case it is being used as part of a numerical derivative which (approx.) evaluates `f(x + EPS * 1j)` and f(x) of the function, then divides this by `EPS * 1j`. Precision of small arguments is essential here.

Commit message: MAINT: special.logsumexp: preserve tiny imaginary components

Patch to review:
diff --git a/scipy/special/_logsumexp.py b/scipy/special/_logsumexp.py
index b649ed48b59f..d85c43e59e11 100644
--- a/scipy/special/_logsumexp.py
+++ b/scipy/special/_logsumexp.py
@@ -125,11 +125,11 @@ def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
     if xp.isdtype(out.dtype, 'complex floating'):
         if return_sign:
             real = xp.real(sgn)
-            imag = xp_float_to_complex(_wrap_radians(xp.imag(sgn)))
+            imag = xp_float_to_complex(_wrap_radians(xp.imag(sgn), xp))
             sgn = real + imag*1j
         else:
             real = xp.real(out)
-            imag = xp_float_to_complex(_wrap_radians(xp.imag(out)))
+            imag = xp_float_to_complex(_wrap_radians(xp.imag(out), xp))
             out = real + imag*1j
 
     # Deal with shape details - reducing dimensions and convert 0-D to scalar for NumPy
@@ -141,9 +141,13 @@ def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
     return (out, sgn) if return_sign else out
 
 
-def _wrap_radians(x):
+def _wrap_radians(x, xp):
     # Wrap radians to -pi, pi interval
-    return (x + math.pi) % (2 * math.pi) - math.pi
+    out = (x + math.pi) % (2 * math.pi) - math.pi
+    # preserve relative precision
+    no_wrap = xp.abs(x) < xp.pi
+    out[no_wrap] = x[no_wrap]
+    return out
 
 
 def _elements_and_indices_with_max_real(a, axis=-1, xp=None):
diff --git a/scipy/special/tests/test_logsumexp.py b/scipy/special/tests/test_logsumexp.py
index 380e181c046f..e8c0f9f37aad 100644
--- a/scipy/special/tests/test_logsumexp.py
+++ b/scipy/special/tests/test_logsumexp.py
@@ -229,6 +229,19 @@ def test_gh21610(self, xp, dtype):
         xp_assert_close(out, xp.real(xp.log(ref)))
         xp_assert_close(sgn, ref/xp.abs(ref))
 
+    def test_gh21709_small_imaginary(self, xp):
+        # Test that `logsumexp` does not lose relative precision of
+        # small imaginary components
+        x = xp.asarray([0, 0.+2.2204460492503132e-17j])
+        res = logsumexp(x)
+        # from mpmath import mp
+        # mp.dps = 100
+        # x, y = mp.mpc(0), mp.mpc('0', '2.2204460492503132e-17')
+        # ref = complex(mp.log(mp.exp(x) + mp.exp(y)))
+        ref = xp.asarray(0.6931471805599453+1.1102230246251566e-17j)
+        xp_assert_close(xp.real(res), xp.real(ref))
+        xp_assert_close(xp.imag(res), xp.imag(ref), atol=0, rtol=1e-15)
+
 
 class TestSoftmax:
     def test_softmax_fixtures(self):


Code review comments:
diff_hunk: @@ -141,9 +141,13 @@ def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
     return (out, sgn) if return_sign else out
 
 
-def _wrap_radians(x):
+def _wrap_radians(x, xp):
     # Wrap radians to -pi, pi interval
-    return (x + math.pi) % (2 * math.pi) - math.pi
+    out = (x + math.pi) % (2 * math.pi) - math.pi
path: scipy/special/_logsumexp.py
text: @author:
Currently, this produces output in $[-\pi, \pi)$.
With the change in the last commit, it would be $[-\pi, \pi]$.
If we want $(-\pi, \pi]$, this should be:
```suggestion
    out = -((-x + math.pi) % (2 * math.pi) - math.pi)
```
If this looks right, I'll test `_wrap_radians` separately with:

```python3
@user1
@user2.mark.usefixtures("skip_xp_backends")
@user2.mark.skip_xp_backends('jax.numpy',
                              reason="JAX arrays do not support item assignment")
def test_wrap_radians(xp):
    x = xp.asarray([-math.pi-1, -math.pi, -1, -1e-300,
                    0, 1e-300, 1, math.pi, math.pi+1])
    ref = xp.asarray([math.pi-1, math.pi, -1, -1e-300,
                    0, 1e-300, 1, math.pi, -math.pi+1])
    res = _wrap_radians(x, xp)
    xp_assert_close(res, ref)
```