Repository: scipy/scipy
Pull request: ENH: `special.rel_entr`: Avoid overflow before computing log
<!-- 
Thanks for contributing a pull request! Please ensure that
your PR satisfies the checklist before submitting:
http://scipy.github.io/devdocs/dev/contributor/development_workflow.html#checklist-before-submitting-a-pr

Also, please name and describe your PR as you would write a
commit message:
http://scipy.github.io/devdocs/dev/contributor/development_workflow.html#writing-the-commit-message

Depending on your changes, you can skip CI operations and save time and energy: 
http://scipy.github.io/devdocs/dev/contributor/continuous_integration.html#skipping

Note that we are a team of volunteers; we appreciate your
patience during the review process.

Again, thanks for contributing!
-->

#### Reference issue
Fixes #20710

#### What does this implement/fix?
This PR improves the accuracy of `rel_entr()` in two ways:

1. If x / y overflows, underflows, or results in a subnormal number, the result is computed with the formula `x * (log(x) - log(y))`.
2. If x / y is close to 1, then a formula based on log1p is used instead. This improves accuracy for small rel_entr() values.

#### Performance regression

This is slower than the original `rel_entr()`.

Example benchmark script:

```
from scipy.special import rel_entr
import numpy as np
import timeit

x = np.random.rand(1_000_000)
y = np.random.rand(1_000_000)

def func():
    rel_entr(x, y)

print(timeit.timeit(func, number=100))
```

Prints 1.958 on main vs 3.587 on my branch.

#### Additional information

The log1p formula was contributed by @fancidev. 

I didn't understand why it worked, so here's a derivation of the formula:

```
rel_entry(x, y) = x * log(x / y)
rel_entry(x, y) = x * log((x - y + y) / y)
rel_entry(x, y) = x * log(1 + (x - y) / y)
rel_entry(x, y) = x * log1p((x - y) / y)
```

cc @mdhaber 

Problem statement:
ENH: `special.rel_entr`: avoid premature overflow
### Describe your issue.

Looks like there's opportunity for improvement in `scipy.special.rel_entr`. Presumably, it takes the ratio `x/y` before taking the logarithm because it's faster than taking the difference of two logs, but this can overflow when the correct result is finite. Perhaps a compromise is to check for the overflow and, in that case, take the time to compute the two logs separately.

This came up gh-20673 as a Hypothesis counterexample rather than a real use case, which is why I consider it an enhancement request.

### Reproducing Code Example

```python
import numpy as np
from scipy import special
x = np.asarray([2., 3., 4.])
y = np.asarray(np.finfo(np.float64).tiny)
special.rel_entr(x, y)  # array([1418.17913143, 2128.48509246,           inf])
x * (np.log(x) - np.log(y))  # array([1418.17913143, 2128.48509246, 2839.13085157])
```


### Error message

```shell
-
```


### SciPy/NumPy/Python version and system information

```shell
main
```

Hints from issue discussion:
It is clearly desirable to avoid premature overflow.

If x and y are greater than one, taking the log and then subtracting might reduce the accuracy, as the domain of logarithm is “larger” than its image. On the other hand, accuracy is not compromised if the arguments are less than one, which may be typical if rel_entr is applied to two probability values.
@fancidev

> If x and y are greater than one, taking the log and then subtracting might reduce the accuracy, as the domain of logarithm is “larger” than its image. On the other hand, accuracy is not compromised if the arguments are less than one, which may be typical if rel_entr is applied to two probability values.

I checked this idea experimentally, and in general, the division approach has better accuracy than the logarithm approach. With three exceptions, the `x * log(x / y)` approach has equal or better accuracy than the `x * (log(x) - log(y))` approach. 

The three exceptions:

1. Overflow. If x / y == inf, then log subtraction is more accurate. (This is the lower-right white area of the graph below - here error is infinite for the ratio approach.)
2. Underflow. If x / y == 0, then log subtraction is more accurate. (This is the upper-left white area of the graph below - here error is infinite for the ratio approach.)
3. Subnormal numbers. If x / y is a [subnormal number](https://en.wikipedia.org/wiki/Subnormal_number), then log subtraction is more accurate. (This is the diagonal band of yellow squares above the main diagonal. This has finite error, but is still quite inaccurate.)

Other than these three exceptions, the division approach is more accurate.

Here's a comparison of the error of the two methods:

![ratio](https://github.com/scipy/scipy/assets/666402/efcd2fa5-d259-463a-bd4a-416f21407cef)

![logdiff](https://github.com/scipy/scipy/assets/666402/6a2af61d-ea5a-454a-9c32-2bb2fd2e9877)

(Note: Errors have been clipped at 95% to make the graphs readable. Some cells have thousands of ULP of error. See full notebook for details.)



With that in mind, here is the logic I would propose:

```
# assume all of the corner cases have already been handled
ratio = x / y
if ratio == inf or ratio < np.finfo(np.float64).tiny:
    return x * (log(x) - log(y))
else:
    return x * log(ratio)
```

The `ratio == inf` check handles case 1. The `ratio < np.finfo(np.float64).tiny` check handles cases 2 and 3.

If you look only at the area where x < 1 and y < 1, which is the bottom-left quadrant of the graph, then the `x * log(x / y)` approach still has equal or better accuracy than the `x * (log(x) - log(y))` approach.

[Notebook source](https://gist.github.com/nickodell/fd7af4afdf03aaece462c5c638d23d66)
Very interesting analysis @nickodell ! I suspect the "clipped" values with large error could be related to `x` and `y` that are close.

To illustrate this, I tried a contrived example:

```python
import numpy as np
from mpmath import mp
mp.dps = 1000
x = 1234
y = np.nextafter(x, 100)
print(f'{np.log(y/x)=}')
print(f'{np.log(y) - np.log(x)=}')
print(f'{float(mp.log(y)-mp.log(x))=}')
print(f'{float(mp.log(mp.mpf(y)/x))=}')
print(f'{np.log1p((y-x)/x)=}')
```

Output:

```
np.log(y/x)=-2.2204460492503136e-16
np.log(y) - np.log(x)=-8.881784197001252e-16
float(mp.log(y)-mp.log(x))=-1.842574355293615e-16
float(mp.log(mp.mpf(y)/x))=-1.842574355293615e-16
np.log1p((y-x)/x)=-1.842574355293615e-16
```

So for those two numbers that are really close, `log(y/x)` is off by 20%, `log(y)-log(x)` is off by 380%, and `log1p` gives the correct answer.

The contrived example shows that if we insist on the accuracy from a numerical perspective, there's a fair bit of work to do.
That's a really smart idea. Here's another graph of function error for the log1p method. (I used `np.log1p((x - y) / y)` rather than `np.log1p((y-x)/x)`, though.)

![Untitled](https://github.com/scipy/scipy/assets/666402/1eb7a38a-e74a-4650-951a-a07ef436264b)

So I could piece together a more accurate `rel_entr()` from three sources: log1p for the diagonal, where x and y are very close, ratio for everywhere else, and log difference for places where there is under/overflow.

Commit message: DOC: Remove docs which require exact float match

Because rel_entr can be calculated in one of two ways, the relative
entropy formula is no longer accurate for x and y which are similar to
one another. However, the numpy calculation is still a useful inuition
for what the calculation is doing.

Patch to review:
diff --git a/scipy/special/_convex_analysis.pxd b/scipy/special/_convex_analysis.pxd
index 2fe95b95ccaf..aff10cb0562f 100644
--- a/scipy/special/_convex_analysis.pxd
+++ b/scipy/special/_convex_analysis.pxd
@@ -1,4 +1,6 @@
 from libc.math cimport log, fabs, expm1, log1p, isnan, NAN, INFINITY
+from libc.float cimport DBL_MIN
+import cython
 
 cdef inline double entr(double x) noexcept nogil:
     if isnan(x):
@@ -20,15 +22,26 @@ cdef inline double kl_div(double x, double y) noexcept nogil:
     else:
         return INFINITY
 
+@cython.cdivision(True)
 cdef inline double rel_entr(double x, double y) noexcept nogil:
+    cdef double ratio
     if isnan(x) or isnan(y):
         return NAN
-    elif x > 0 and y > 0:
-        return x * log(x / y)
-    elif x == 0 and y >= 0:
-        return 0
-    else:
+    if x <= 0 or y <= 0:
+        if x == 0 and y >= 0:
+            return 0
         return INFINITY
+    ratio = x / y
+    if 0.5 < ratio < 2:
+        # When x and y are close, this is more accurate
+        return x * log1p((x - y) / y)
+    if DBL_MIN < ratio < INFINITY:
+        # There are no underflow/overflow issues
+        return x * log(ratio)
+    # x and y are so far apart that taking x / y
+    # results in either an underflow, overflow,
+    # or subnormal number. Do the logarithm first
+    return x * (log(x) - log(y))
 
 cdef inline double huber(double delta, double r) noexcept nogil:
     if delta < 0:
diff --git a/scipy/special/tests/test_basic.py b/scipy/special/tests/test_basic.py
index ca8845f7a3e2..0ee5d90d722a 100644
--- a/scipy/special/tests/test_basic.py
+++ b/scipy/special/tests/test_basic.py
@@ -4152,6 +4152,46 @@ def xfunc(x, y):
     assert_func_equal(special.rel_entr, w, z, rtol=1e-13, atol=1e-13)
 
 
+def test_rel_entr_gh_20710_near_zero():
+    # Check accuracy of inputs which are very close
+    inputs = np.array([
+        # x, y
+        (0.9456657713430001, 0.9456657713430094),
+        (0.48066098564791515, 0.48066098564794774),
+        (0.786048657854401, 0.7860486578542367),
+    ])
+    # Known values produced using `x * mpmath.log(x / y)` with dps=30
+    expected = [
+        -9.325873406851269e-15,
+        -3.258504577274724e-14,
+        1.6431300764454033e-13,
+    ]
+    x = inputs[:, 0]
+    y = inputs[:, 1]
+    assert_allclose(special.rel_entr(x, y), expected, rtol=1e-13, atol=0)
+
+
+def test_rel_entr_gh_20710_overflow():
+    inputs = np.array([
+        # x, y
+        # Overflow
+        (4, 2.22e-308),
+        # Underflow
+        (1e-200, 1e+200),
+        # Subnormal
+        (2.22e-308, 1e15),
+    ])
+    # Known values produced using `x * mpmath.log(x / y)` with dps=30
+    expected = [
+        2839.139983229607,
+        -9.210340371976183e-198,
+        -1.6493212008074475e-305,
+    ]
+    x = inputs[:, 0]
+    y = inputs[:, 1]
+    assert_allclose(special.rel_entr(x, y), expected, rtol=1e-13, atol=0)
+
+
 def test_huber():
     assert_equal(special.huber(-1, 1.5), np.inf)
     assert_allclose(special.huber(2, 1.5), 0.5 * np.square(1.5))
diff --git a/scipy/stats/_entropy.py b/scipy/stats/_entropy.py
index 1e6c8d9e1963..6562ade0aebb 100644
--- a/scipy/stats/_entropy.py
+++ b/scipy/stats/_entropy.py
@@ -123,9 +123,9 @@ def entropy(pk: np.typing.ArrayLike,
 
     >>> D = entropy(pk, qk, base=base)
     >>> D
-    0.7369655941662062
-    >>> D == np.sum(pk * np.log(pk/qk)) / np.log(base)
-    True
+    0.736965594166206
+    >>> np.sum(pk * np.log(pk/qk)) / np.log(base)
+    0.736965594166206
 
     The cross entropy can be calculated as the sum of the entropy and
     relative entropy`:


Code review comments:
diff_hunk: @@ -123,9 +123,9 @@ def entropy(pk: np.typing.ArrayLike,
 
     >>> D = entropy(pk, qk, base=base)
     >>> D
-    0.7369655941662062
-    >>> D == np.sum(pk * np.log(pk/qk)) / np.log(base)
-    True
+    0.736965594166206
+    >>> np.sum(pk * np.log(pk/qk)) / np.log(base)
path: scipy/stats/_entropy.py
text: @user1:
How about using the actual value that you're getting for `D` to avoid the "slight lie" and then using `np.allclose` with a very small `rtol` (around twice machine epsilon?) to replace the exact check?

@author:
Here's a suggested fix: 2d7b443.

I used a value which is about five times bigger than machine epsilon, for simplicity. I also set atol to zero, as it defaults to 1e-8 otherwise.

PS: I can rebase this to remove the revert commits later.

@user1:
No need to rebase. I’m going to squash merge. How about `4e-16` for the `rtol`? `1e-15` feels big for something which should never be off by more than 1 ulp. 

@author:
See 3da2b2d.