Repository: scipy/scipy

Issue:
scipy.special.logsumexp for complex input with return_sign=True under numpy 2.0
### Describe your issue.

https://github.com/numpy/numpy/pull/25441 changes the definition of `np.sign` for complex input in numpy 2.0. When exploring the effect of this change, I found that it leads to strange results for `scipy.special.logsumexp` with `return_sign = True`.

I've demonstrated this below by patching `np.sign` to the new definition. I'm not sure how to interpret the meaning of the `sign` output here for `logsumexp`, but it seems like the change upstream might break some assumptions made by downstream callers of `special.logsumexp`.

Perhaps `scipy` should change how `sign` is computed here, so that it is always $\pm 1$, even for complex input?

### Reproducing Code Example

```python
import numpy as np
from scipy.special import logsumexp

def newsign(x):
  x = np.asarray(x)
  return np.where(x == 0, x, (x / np.abs(x)).astype(x.dtype))

def test():
  rng = np.random.default_rng(0)
  x = rng.normal(size=10) + 1j * rng.normal(size=10)
  print("sign =", np.sign(x.sum()))
  print("log(sum(exp(x))) =", np.log(np.sum(np.exp(x))))
  result, s = logsumexp(x, return_sign=True)
  print("logsumexp result:", result, s)
  print("reconstructed sum(exp(x)):", np.exp(result) * s)
  print("manual sum(exp(x))       :", np.sum(np.exp(x)))
  print()

print("NumPy<2.0 behavior")
test()

_np_sign = np.sign
np.sign = newsign
try:
  print("NumPy>=2.0 behavior")
  test()
finally:
  np.sign = _np_sign
```


### Error message

```shell
NumPy<2.0 behavior
sign = (1+0j)
log(sum(exp(x))) = (2.385269178468705-0.5887388825917717j)
logsumexp result: (2.385269178468705-0.5887388825917715j) (1+0j)
reconstructed sum(exp(x)): (9.033280080568113-6.0317984614767175j)
manual sum(exp(x))       : (9.033280080568113-6.031798461476719j)

NumPy>=2.0 behavior
sign = (0.18452108892318747-0.9828285546027857j)
log(sum(exp(x))) = (2.385269178468705-0.5887388825917717j)
logsumexp result: (2.385269178468705-0.6332187823262332j) (0.9990109323452934-0.044465234223914865j)
reconstructed sum(exp(x)): (8.461679537367544-6.8104861133957915j)
manual sum(exp(x))       : (9.033280080568113-6.031798461476719j)
```


### SciPy/NumPy/Python version and system information

```shell
1.11.4 1.23.5 sys.version_info(major=3, minor=10, micro=12, releaselevel='final', serial=0)
Build Dependencies:
  blas:
    detection method: pkgconfig
    found: true
    include directory: /usr/local/include
    lib directory: /usr/local/lib
    name: openblas
    openblas configuration: USE_64BITINT= DYNAMIC_ARCH=1 DYNAMIC_OLDER= NO_CBLAS=
      NO_LAPACK= NO_LAPACKE= NO_AFFINITY=1 USE_OPENMP= HASWELL MAX_THREADS=2
    pc file directory: /usr/local/lib/pkgconfig
    version: 0.3.21.dev
  lapack:
    detection method: pkgconfig
    found: true
    include directory: /usr/local/include
    lib directory: /usr/local/lib
    name: openblas
    openblas configuration: USE_64BITINT= DYNAMIC_ARCH=1 DYNAMIC_OLDER= NO_CBLAS=
      NO_LAPACK= NO_LAPACKE= NO_AFFINITY=1 USE_OPENMP= HASWELL MAX_THREADS=2
    pc file directory: /usr/local/lib/pkgconfig
    version: 0.3.21.dev
  pybind11:
    detection method: config-tool
    include directory: unknown
    name: pybind11
    version: 2.11.0
Compilers:
  c:
    commands: cc
    linker: ld.bfd
    name: gcc
    version: 10.2.1
  c++:
    commands: c++
    linker: ld.bfd
    name: gcc
    version: 10.2.1
  cython:
    commands: cython
    linker: cython
    name: cython
    version: 0.29.36
  fortran:
    commands: gfortran
    linker: ld.bfd
    name: gcc
    version: 10.2.1
  pythran:
    include directory: ../../tmp/pip-build-env-c6c8ru56/overlay/lib/python3.10/site-packages/pythran
    version: 0.14.0
Machine Information:
  build:
    cpu: x86_64
    endian: little
    family: x86_64
    system: linux
  cross-compiled: true
  host:
    cpu: x86_64
    endian: little
    family: x86_64
    system: linux
Python Information:
  path: /opt/python/cp310-cp310/bin/python
  version: '3.10'
```


Discussion:
It doesn't matter that for the fact that there is a change with the new definition of the sign, but note that it is defined as `x / abs(x)`, not the other way around.
Though perhaps you were trying to make things correct at https://github.com/scipy/scipy/blob/906e6cd9bb022e78a07b34636bee64813ce56afa/scipy/special/_logsumexp.py#L118

That part would certainly need adjustment for complex. With it, one would in principle get a real-only return value, with the correct sign.
Yeah, it looks like that is the root of the numerical issue I'm seeing.

But note that the new `np.sign` definition also conflicts with the documented behavior of `scipy.special.logsumexp`, which [states](https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.logsumexp.html):

> sgn ndarray
> If return_sign is True, this will be an array of floating-point numbers matching res and +1, 0, or -1 depending on the sign of the result. If False, only one result is returned.

So either way, there is a bug in the implementation that needs to be fixed, and if we decide to stick with the new behavior of `np.sign`, the documentation will need to be updated as well.
> note that it is defined as x / abs(x), not the other way around.

Thanks, I updated my repro to be more accurate. Note that the spirit of the original issue is unchanged by this.
Thinking about this more, I think the best fix would be to avoid the use of `np.sign`, and stick with the behavior as currently documented, where `sign` is a float that is either +1, -1, or 0. That will be backward compatible with the current behavior under the numpy 2.0 upgrade, and is less confusing than scalaing the output by a complex unit vector.

Alternatively, we could raise an error when `return_sign` is used with complex input, because it's not really necessary: unlike log(float), log(complex) is well-defined for negative numbers.
I don't agree with using -1, +1 for complex, but definitely can see the case for raising for complex, since as you note the reason for the float case is just to avoid a log of negative numbers, which is not an issue for complex. However, I can see a point of separating the amplitude from the sign, in which case the current code almost works - just needs to do a divide.
@rgommers I think we need some input from you or other maintainers on a design question here. As I see it, there are three approaches to fixing this issue:

1. Make the `sign` output of `logsumexp` always match the convention used by `np.sign`. This would mean that scipy would return different results depending on which numpy version is installed.
2. Make the `sign` output of `logsumexp` always equal `x / abs(s)`, which is the convention in numpy 2.0.
3. Choose to raise an exception when `return_sign` is used with complex input, because it is only really necessary for floating point where `log(negative)` returns `NaN`.

Do you have thoughts on which approach we should choose?
xref gh-4859 for why `return_sign` exists in the first place (has to do with weights with incorrect values). It's a fairly weak rationale, and these kinds of keywords we now understand to be a bad idea, so ideally `return_sign` shouldn't exist at all. It's also not used in the SciPy code base at all, so the choice won't affect us much.

Given the choices here, I suspect that (2) is slightly nicer than (1), and both are better than raising (that may make it harder to write generic code). I think (2) aligns with your discussion on the PR, right?
(2) sounds good to me. I'll update my PR.

Resolving patch:
diff --git a/scipy/special/_logsumexp.py b/scipy/special/_logsumexp.py
--- a/scipy/special/_logsumexp.py
+++ b/scipy/special/_logsumexp.py
@@ -40,11 +40,14 @@ def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
     res : ndarray
         The result, ``np.log(np.sum(np.exp(a)))`` calculated in a numerically
         more stable way. If `b` is given then ``np.log(np.sum(b*np.exp(a)))``
-        is returned.
+        is returned. If ``return_sign`` is True, ``res`` contains the log of
+        the absolute value of the argument.
     sgn : ndarray
-        If return_sign is True, this will be an array of floating-point
-        numbers matching res and +1, 0, or -1 depending on the sign
-        of the result. If False, only one result is returned.
+        If ``return_sign`` is True, this will be an array of floating-point
+        numbers matching res containing +1, 0, -1 (for real-valued inputs)
+        or a complex phase (for complex inputs). This gives the sign of the
+        argument of the logarithm in ``res``.
+        If ``return_sign`` is False, only one result is returned.
 
     See Also
     --------
@@ -97,7 +100,9 @@ def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
             a = a + 0.  # promote to at least float
             a[b == 0] = -np.inf
 
-    a_max = np.amax(a, axis=axis, keepdims=True)
+    # Scale by real part for complex inputs, because this affects
+    # the magnitude of the exponential.
+    a_max = np.amax(a.real, axis=axis, keepdims=True)
 
     if a_max.ndim > 0:
         a_max[~np.isfinite(a_max)] = 0
@@ -114,8 +119,12 @@ def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
     with np.errstate(divide='ignore'):
         s = np.sum(tmp, axis=axis, keepdims=keepdims)
         if return_sign:
-            sgn = np.sign(s)
-            s *= sgn  # /= makes more sense but we need zero -> zero
+            # For complex, use the numpy>=2.0 convention for sign.
+            if np.issubdtype(s.dtype, np.complexfloating):
+                sgn = s / np.where(s == 0, 1, abs(s))
+            else:
+                sgn = np.sign(s)
+            s = abs(s)
         out = np.log(s)
 
     if not keepdims:


Test patch:
diff --git a/scipy/special/tests/test_logsumexp.py b/scipy/special/tests/test_logsumexp.py
--- a/scipy/special/tests/test_logsumexp.py
+++ b/scipy/special/tests/test_logsumexp.py
@@ -115,6 +115,19 @@ def test_logsumexp_sign_shape():
     assert_equal(r.shape, (1,3))
 
 
+def test_logsumexp_complex_sign():
+    a = np.array([1 + 1j, 2 - 1j, -2 + 3j])
+
+    r, s = logsumexp(a, return_sign=True)
+
+    expected_sumexp = np.exp(a).sum()
+    # This is the numpy>=2.0 convention for np.sign
+    expected_sign = expected_sumexp / abs(expected_sumexp)
+
+    assert_allclose(s, expected_sign)
+    assert_allclose(s * np.exp(r), expected_sumexp)
+
+
 def test_logsumexp_shape():
     a = np.ones((1, 2, 3, 4))
     b = np.ones_like(a)