Repository: scipy/scipy
Pull request: MAINT: special.logsumexp: improve behavior with complex infinities
#### Reference issue
Closes gh-22601

#### What does this implement/fix?
gh-22601 reported that `special.logsumexp` does not always return the correct complex infinity. (Although there are some differences between the old behavior and the new behavior, it was never correct in general.) This improves that behavior.

#### Additional information
Skipped a test with torch, ultimately because
```python3
import torch as xp
a = xp.asarray([1 + 0j])
b = xp.asarray([xp.inf + 0j])
print(a - b, -b + a)
# tensor([-inf+nanj]) tensor([-inf+0.j])
```
Should this be reported?

---

I'm not sure if I wrote the math out somewhere before, so here is what the current implementation does:

```math
\texttt{logsumexp}(x, b) = \log \left( b_1 \exp(x_1) + \cdots + b_n \exp(x_n) \right)
```

We want to rearrange things to avoid overflow, so we separate this as:

```math
\texttt{logsumexp}(x, b) = \log \left( m \exp(x') + \sum_{i : x_i \neq x'} \exp(x_i)  \right)
```

where $x'$ is the element of $x_i$ with maximum real component and 

```math
m = \sum_{i : x_i = x'} b_i
```

Now we can divide and multiply by $m \exp(x')$ to get:

```math
\texttt{logsumexp}(x, b) = \log \left( \left(  1 + \frac{1}{m} \sum_{i : x_i \neq x'} \frac{\exp(x_i)}{\exp(x')} \right) m \exp(x') \right)
```

Simplifying:
```math
\texttt{logsumexp}(x, b) = \texttt{log1p} \left( \frac{1}{m} \sum_{i : x_i \neq x'} \exp(x_i - x') \right) + \log(m) + x'
```

The issue with this that led to gh-22601 is that when $x'$ has infinite real part, $x_i - x'$ can break down. There are four cases:
- Only $x'$ has infinite (positive or negative) real part. No problem because $x_i - x'$ will just have negative infinite real part and the exponentiation will turn this to zero.
- $x'$ has infinite positive real part, and at least one of $x_i$ has negative real part. Also not a problem, because again, $x_i - x'$ will just have negative infinite real part; no NaNs.
- $x'$ and at least one of the $x_i$ has infinite negative real part. Really, this requires that *all* the $x_i$ also have negative infinite real part because $x'$ has the maximum real part. The correct result of `logsumexp` in this case is the logarithm of `0`. This will have negative real part, and I have trouble imagining a case in which the phase matters. In the past, I think we've always returned `-inf + 0j` in this case. What we do now is avoid subtracting $x_i - x'$. The sum of exponentials ends up as `0 + 0j` rather than NaN, and the final result is $\log(m) + x'$. This has the property that `logsumexp([-inf + y j]*m) = -inf + y j` rather than `-inf + 0j`, which doesn't make *less* sense to me than the old behavior. 
- $x'$ and at least one of the $x_i$ has infinite positive real part. Previously, we made an adjustment to prevent the subtraction from producing NaNs, figuring (correctly for real input) that the final $+ x'$ would make the result positive infinity as it should be. What this didn't account for was the complex case, for which we can have complex infinities that are aligned, anti-aligned, or somewhere in-between. The correct result depends on which is the case, and we have never correctly distinguished between them in general, so now we let the subtraction produce NaNs, and the result is `nan + nan j`. It *might* be better to produce `inf + nanj`, but even a light breeze will turn this unicorn into `nan + nanj`, so we don't bother.

Problem statement:
BUG: special.logsumexp: inconsistency in phase when one element is complex infinity
### Describe your issue.

This is an edge case, but there is a difference between scipy and naiive numpy implementation of `logsumexp`. The following code produces different outputs:

### Reproducing Code Example

```python
import numpy as np
from numpy import inf
from scipy.special import logsumexp

def check(x):
    y_scipy = logsumexp(x, 1)
    y_numpy = np.log(np.sum(np.exp(x), 1))

    print("numpy:", y_numpy)
    print("scipy:", y_scipy)

x = np.array([[  inf + 0.0000j, -0.9116 + 0.0000j],
              [ -inf + 0.0000j,    -inf + 0.0000j]])
check(x)

x = np.array([[  inf + 0.0000j, -0.9116 + 0.0000j],
              [ -inf + 0.0000j,    -inf + 0.7533j]])
check(x)

x = np.array([[  inf + 0.0000j, -0.9116 + 0.7533j],
              [ -inf + 0.0000j,    -inf + 0.0000j]])
check(x)

x = np.array([[  inf + 0.0000j, -0.9116 + 0.7533j],
              [ -inf + 0.0000j,    -inf + 0.7533j]])
```

### Current (incorrect) output
Scipy version 1.15.0 and newer prints:
```shell

numpy: [ inf+0.j -inf+0.j]
scipy: [ inf+0.j -inf+0.j]
numpy: [ inf+0.j -inf+0.j]
scipy: [ inf+0.j     -inf+0.7533j]
numpy: [ inf+0.j -inf+0.j]
scipy: [ inf+0.20946886j -inf+0.j        ]
```

### Expected output
Scipy version 1.14.1 and older prints:
```
numpy: [ inf+0.j -inf+0.j]
scipy: [ inf+0.j -inf+0.j]
numpy: [ inf+0.j -inf+0.j]
scipy: [ inf+0.j -inf+0.j]
numpy: [ inf+0.j -inf+0.j]
scipy: [ inf+0.j -inf+0.j]
```

### Versions

Errors are present in scipy v1.15.0 and newer irrespective of NumPy and Python versions.

Hints from issue discussion:
SciPy 1.5.0 is very old now and unlikely to be fixed. Is this still a problem with the latest release 1.15?
I apologize for my typos, it 1.15.0 and newer. I edited the description, thanks!
I'll need to take a closer look, but when you're looking at edge cases (e.g. involving infinities), it's important to consider the correct answer. `np.log(np.exp(a) + np.exp(b))` won't necessarily get you that due to overflow. Keep in mind that `exp(-inf + 0j) = exp(-inf + 0.7533j) = 0`.
The NumPy version is indeed naiive and based on a reduction of the binary operator. IIRC, the scipy version takes a more global approach and is likely better in general. 
@charris Just to clarify, the top post is comparing against `np.log(np.sum(np.exp(x)))`, so the reduction (`sum`) is based on binary addition. Were you thinking of a reduction of `logaddexp`? Unfortunately, we can't compare against `logaddexp` here because it doesn't support complex input (https://github.com/numpy/numpy/issues/5652#issuecomment-2353894558).
There are only two unique cases of discrepancies in the top post. Here is a MRE of the first.

```python3
import numpy as np
from scipy import special
special.logsumexp([ -np.inf + 0.0000j,    -np.inf + 0.7533j])
# (-inf+0.7533j)
```
Both of the numbers in the argument are the logarithm of zero, so any result with `-inf` real component is reasonable. The way to reason about these cases is to recognize that $\exp(x + y j) = \exp(x) \exp(yj) = r \exp(\theta j)$ and think of this as a magnitude and direction in the complex plane. In both cases, $r = \exp(-\inf) = 0$, so we have two different numbers with zero magnitude; one has phase $0$ and the other has phase $0.7533$. What should the phase of the result be? 

One could argue that the phase should be `0.7533 / 2` if `-inf ==-inf` is `True`. There is another argument that the imaginary component is indeterminate. But I don't think the old result is any *more* correct than the new one.

The second case is:
```python3
special.logsumexp([  np.inf + 0.0000j, -0.9116 + 0.7533j])
# (inf+0.20946885578578167j)
```
The best answer in this case would indeed be `inf + 0.0000j` because we have a number with infinite magnitude and zero phases added to a finite number in the first quadrant. The resultant number will have infinite magnitude and phase on the positive side of zero. 

I'll have to look into exactly why the implementation is giving a nonzero phase here. But the short story is probably that the identities we use to get the best result for finite numbers break down for infinities, so we get some weirdness if there is not special care for degenerate input. This is not unique to `logsumexp`; there are plenty of similar cases in `scipy.special` that we are just starting to work out (gh-22539, gh-22514). Even the old `logsumexp` had issues of the same type.

```python3
logsumexp([[  inf + np.pi*1j, -0.9116 + 0.7533j]])
# np.complex128(inf+2.356194490192345j)
logsumexp([[  inf + np.pi*1j, -0.9116 - 0.7533j]])
# np.complex128(inf+2.356194490192345j) 
```

---

There is a line:
```python3
# Arithmetic between infinities will introduce NaNs.
# `+ a_max` at the end naturally corrects for removing them here.
shift = xp.where(xp.isfinite(a_max), a_max, xp.asarray(0, dtype=a_max.dtype))
```

This does improve the result when there are multiple infinities. And when there are multiple infinites, the phase is typically indeterminate, so it does no harm. But you have chosen a case in which there is only one infinity. So we'll have to do something more principled here. Maybe instead of setting to zero we leave the imaginary part.

@Aidyn-A how did you run across this? Are you interested in the phase of complex infinities in an application, or were you hunting for edge cases? 
>@Aidyn-A how did you run across this? Are you interested in the phase of complex infinities in an application, or were you hunting for edge cases?

Fortunately it was not a real-world application. I have discovered it in PyTorch unit-test:
https://github.com/pytorch/pytorch/blob/4106aa33eb2946bf67f4ffd2ad9f2dcb52ed2384/test/test_reductions.py#L492-L509
It used to be successful until the scipy update.

Commit message: Merge pull request #120 from mdhaber/gh22607

Simplify logsumexp edge case handling

Patch to review:
diff --git a/scipy/special/_logsumexp.py b/scipy/special/_logsumexp.py
index 03a43348c9c2..6cd4621bf973 100644
--- a/scipy/special/_logsumexp.py
+++ b/scipy/special/_logsumexp.py
@@ -110,8 +110,25 @@ def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
     axis = tuple(range(a.ndim)) if axis is None else axis
 
     if xp_size(a) != 0:
+        with np.errstate(over='ignore'):
+            # Where result is infinite, we use the direct logsumexp calculation to
+            # delegate edge case handling to the behavior of `xp.log` and `xp.exp`,
+            # which should follow the C99 standard for complex values.
+            b_exp_a = xp.exp(a) if b is None else b * xp.exp(a)
+            sum = xp.sum(b_exp_a, axis=axis, keepdims=True)
+            sgn_inf = _sign(sum, xp) if return_sign else None
+            sum = xp.abs(sum) if return_sign else sum
+            out_inf = xp.log(sum)
+
         with np.errstate(divide='ignore', invalid='ignore'):  # log of zero is OK
             out, sgn = _logsumexp(a, b, axis=axis, return_sign=return_sign, xp=xp)
+
+        # Replace infinite results. This probably could be done with an
+        # `apply_where`-like strategy to avoid redundant calculation, but currently
+        # `apply_where` itself is only for elementwise functions.
+        out_finite = xp.isfinite(out)
+        out = xp.where(out_finite, out, out_inf)
+        sgn = xp.where(out_finite, sgn, sgn_inf) if return_sign else sgn
     else:
         shape = np.asarray(a.shape)  # NumPy is convenient for shape manipulation
         shape[axis] = 1
@@ -203,12 +220,8 @@ def _logsumexp(a, b, axis, return_sign, xp):
     m = (xp.sum(i_max_dt, axis=axis, keepdims=True, dtype=a.dtype) if b is None
          else xp.sum(b * i_max_dt, axis=axis, keepdims=True, dtype=a.dtype))
 
-    # Arithmetic between infinities will introduce NaNs.
-    # `+ a_max` at the end naturally corrects for removing them here.
-    shift = xp.where(xp.isfinite(a_max), a_max, 0.)
-
     # Shift, exponentiate, scale, and sum
-    exp = b * xp.exp(a - shift) if b is not None else xp.exp(a - shift)
+    exp = b * xp.exp(a - a_max) if b is not None else xp.exp(a - a_max)
     s = xp.sum(exp, axis=axis, keepdims=True, dtype=exp.dtype)
     s = xp.where(s == 0, s, s/m)
 
diff --git a/scipy/special/tests/test_logsumexp.py b/scipy/special/tests/test_logsumexp.py
index 28537ffbe571..dadb3aab0ddb 100644
--- a/scipy/special/tests/test_logsumexp.py
+++ b/scipy/special/tests/test_logsumexp.py
@@ -1,3 +1,4 @@
+import itertools as it
 import math
 import pytest
 
@@ -9,10 +10,11 @@
 
 from scipy.special import log_softmax, logsumexp, softmax
 from scipy.special._logsumexp import _wrap_radians
+from scipy.stats.tests.test_stats import skip_xp_backends
+
 from scipy._lib.array_api_extra.testing import lazy_xp_function
 
 
-skip_xp_backends = pytest.mark.skip_xp_backends
 
 
 dtypes = ['float32', 'float64', 'int32', 'int64', 'complex64', 'complex128']
@@ -33,11 +35,13 @@ def test_wrap_radians(xp):
     xp_assert_close(res, ref, atol=0)
 
 
+# numpy warning filters don't work for dask (dask/dask#3245)
+# (also we should not expect the numpy warning filter to work for any Array API
+# library)
+@pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning")
+@pytest.mark.filterwarnings("ignore:divide by zero encountered:RuntimeWarning")
+@pytest.mark.filterwarnings("ignore:overflow encountered:RuntimeWarning")
 class TestLogSumExp:
-    # numpy warning filters don't work for dask (dask/dask#3245)
-    # (also we should not expect the numpy warning filter to work for any Array API
-    # library)
-    @pytest.mark.filterwarnings("ignore:divide by zero encountered in log")
     def test_logsumexp(self, xp):
         # Test with zero-size array
         a = xp.asarray([])
@@ -117,8 +121,6 @@ def test_logsumexp_sign(self, xp):
         xp_assert_close(r, xp.asarray(1.))
         xp_assert_equal(s, xp.asarray(-1.))
 
-    @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning:dask")
-    @pytest.mark.filterwarnings("ignore:divide by zero encountered:RuntimeWarning:dask")
     def test_logsumexp_sign_zero(self, xp):
         a = xp.asarray([1, 1])
         b = xp.asarray([1, -1])
@@ -266,6 +268,37 @@ def test_gh21709_small_imaginary(self, xp):
         xp_assert_close(xp.imag(res), xp.imag(ref), atol=0, rtol=1e-15)
 
 
+    @pytest.mark.parametrize('x,y', it.product(
+        [
+            -np.inf,
+            np.inf,
+            complex(-np.inf, 0.),
+            complex(-np.inf, -0.),
+            complex(-np.inf, np.inf),
+            complex(-np.inf, -np.inf),
+            complex(np.inf, 0.),
+            complex(np.inf, -0.),
+            complex(np.inf, np.inf),
+            complex(np.inf, -np.inf),
+            # Phase in each quadrant.
+            complex(-np.inf, 0.7533),
+            complex(-np.inf, 2.3562),
+            complex(-np.inf, 3.9270),
+            complex(-np.inf, 5.4978),
+            complex(np.inf, 0.7533),
+            complex(np.inf, 2.3562),
+            complex(np.inf, 3.9270),
+            complex(np.inf, 5.4978),
+        ], repeat=2)
+    )
+    def test_gh22601_infinite_elements(self, x, y, xp):
+        # Test that `logsumexp` does reasonable things in the presence of
+        # real and complex infinities.
+        res = logsumexp(xp.asarray([x, y]))
+        ref = xp.log(xp.sum(xp.exp(xp.asarray([x, y]))))
+        xp_assert_equal(res, ref)
+
+
 class TestSoftmax:
     def test_softmax_fixtures(self, xp):
         xp_assert_close(softmax(xp.asarray([1000., 0., 0., 0.])),


Code review comments:
diff_hunk: @@ -110,8 +110,25 @@ def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
     axis = tuple(range(a.ndim)) if axis is None else axis
 
     if xp_size(a) != 0:
+        with np.errstate(over='ignore'):
path: scipy/special/_logsumexp.py
text: @user1:
```suggestion
        with np.errstate(divide='ignore', invalid='ignore', over='ignore'):
```

I think this would fix the test failures.