Repository: optuna/optuna

Issue:
Wrong logpdf for scaled truncated norm function (TPESampler)
I think, I found a bug in the implementation of "logpdf(x,a,b,loc,scale)" in _truncnorm.py. This function is used for the Parzen estimator in the TPE Sampler with continuous values. I think this might lead to a wrong behavior of the TPESampler.

### Expected behavior
I would expect that the logpdf is a correctly normalized pdf. This means especially that the integral over np.exp(logpdf(x, a, b, loc, scale)) from a to b should be always 1.

(I assume the "logpdf(x,a,b,loc,scale)" function should represent the log pdf of a truncated Gaussian with borders "a" & "b", positioned at "loc" with standard deviation "scale".)

**Actual behavior**: 
The integral is only 1 if the standard deviation is 1. Otherwise it has a different value. See "steps to reproduce below".

**Further effects**
This has then effects on higher levels of the code. Especially the resulting Parzen estimator has no normalized pdf. As the TPESampler computes the ratio of the pdfs, this likely leads to a wrong behavior. As the difference is dependent on the standard deviation of the components, I don't think these cancel out. See "steps to reproduce" about the integral of the Parzen Estimator's pdfs. 

### The bug
I think the problem is the following. The function logpdf ( https://github.com/optuna/optuna/blob/master/optuna/samplers/_tpe/_truncnorm.py#L218 ) normalizes the variable 'x' and then computes the logpdf of a normal variable (loc=0, scale=1).
```
x = (x - loc) / scale
x, a, b = np.atleast_1d(x, a, b)
out = _norm_logpdf(x) - _log_gauss_mass(a, b)
```
This seems to assume that the pdf of a scaled normal distribution is "normpdf((x-loc)/scale)" where normpdf has loc=0 and scale=1. However according to https://en.wikipedia.org/wiki/Normal_distribution#General_normal_distribution 
the actual formula is "1/scale * normpdf((x-loc)/scale)" 

For a Gaussian, the solution should be similar to adding another term based on the scale, i.e.:
```
out = _norm_logpdf(x) - _log_gauss_mass(a, b) - np.log(scale)
```
However, I am not exactly sure if this also works for a truncated Gaussian.

### Tests
There are tests for truncnorm.logpdf but they all seem to assume that the standard deviation is 1, and only vary the truncation borders:
https://github.com/optuna/optuna/blob/master/tests/samplers_tests/tpe_tests/test_truncnorm.py





### Environment

- Optuna version:3.5.0.dev
- Python version:3.9.13
- OS:Windows-10-10.0.19045-SP0


### Steps to reproduce

**truncnorm**

``` 
from optuna.samplers._tpe._truncnorm import logpdf
import matplotlib.pyplot as plt
import numpy as np
from scipy.integrate import quad

integral = quad(lambda t: np.exp(logpdf(t, -10, 10, loc=0, scale=1)), -10, 10)
integral2 = quad(lambda t: np.exp(logpdf(t, -10, 10, loc=0, scale=2)), -10, 10)

print(f"Int = {integral[0]} +-{integral[1]}")
print(f"Int = {integral2[0]} +-{integral2[1]}")
```
Output:
```
Int = 1.0 +-8.671029987439099e-10
Int = 1.9999988533937123 +-1.7336641289222854e-09
```

**Parzen estimator**
I added some logging calls in "_sample" from the TPESampler class. Here, the TPESampler computes two parzen estimators. According to my understanding these should be normalized pdfs, i.e. also integrate to 1
```
        log_likelihoods_below = mpe_below.log_pdf(samples_below)
        log_likelihoods_above = mpe_above.log_pdf(samples_below)
        int_below = quad(lambda x: np.exp(mpe_below.log_pdf({'x': np.array([x])})), search_space['x'].low, search_space['x'].high)
        int_above = quad(lambda x: np.exp(mpe_above.log_pdf({'x': np.array([x])})), search_space['x'].low, search_space['x'].high)
        _logger.info(f"int_below = {int_below[0]}")
        _logger.info(f"int_above = {int_above[0]}")
```
Due to the _truncnorm.logpdf, they also don't integrate to 1. The output is:
```
[I 2023-11-13 19:13:23,575] int_below = 6.666666666666666
[I 2023-11-13 19:13:23,575] int_above = 2.099087317390244
```
(The problem is just a simple minimization of a quadratic function of 'x')



Discussion:
An easy way to reproduce the problem is by comparing the values to scipy.stats.truncnorm:

```
import numpy as np
from scipy.stats import truncnorm as scipy_truncnorm
from optuna.samplers._tpe import _truncnorm as optuna_truncnorm

print(f"scipy pdf(0, a=-10, b=10, scale=1): {np.exp(scipy_truncnorm.logpdf(0, a=-10, b=10, scale=1))}")
print(f"optuna pdf(0, a=-10, b=10, scale=1): {np.exp(optuna_truncnorm.logpdf(0, a=-10, b=10, scale=1))}")
print("---")
print(f"scipy pdf(0, a=-10, b=10, scale=2): {np.exp(scipy_truncnorm.logpdf(0, a=-10, b=10, scale=2))}")
print(f"optuna pdf(0, a=-10, b=10, scale=2): {np.exp(optuna_truncnorm.logpdf(0, a=-10, b=10, scale=2))}")
```

Output:
```
scipy pdf(0, a=-10, b=10, scale=1): 0.3989422804014327
optuna pdf(0, a=-10, b=10, scale=1): [0.39894228]
---
scipy pdf(0, a=-10, b=10, scale=2): 0.19947114020071632
optuna pdf(0, a=-10, b=10, scale=2): [0.39894228]
```

(According to the headers optuna's _truncnorm.py seems to originate from scipy. Maybe a bug that was fixed in an earlier version of scipy? I looked through the issues in scipy for truncnorm but could not find an issue which directly seems to relate to this problem.)
Thank you for your bug report! Your suggestion should work for a truncated Gaussian. I'll make a PR.

```
out = _norm_logpdf(x) - _log_gauss_mass(a, b) - np.log(scale)
```

In my understanding, this bug luckily does not affect the sampling results.

Resolving patch:
diff --git a/optuna/samplers/_tpe/_truncnorm.py b/optuna/samplers/_tpe/_truncnorm.py
--- a/optuna/samplers/_tpe/_truncnorm.py
+++ b/optuna/samplers/_tpe/_truncnorm.py
@@ -226,7 +226,7 @@ def logpdf(
 
     x, a, b = np.atleast_1d(x, a, b)
 
-    out = _norm_logpdf(x) - _log_gauss_mass(a, b)
+    out = _norm_logpdf(x) - _log_gauss_mass(a, b) - np.log(scale)
 
     x, a, b = np.broadcast_arrays(x, a, b)
     out[(x < a) | (b < x)] = -np.inf


Test patch:
diff --git a/tests/samplers_tests/tpe_tests/test_truncnorm.py b/tests/samplers_tests/tpe_tests/test_truncnorm.py
--- a/tests/samplers_tests/tpe_tests/test_truncnorm.py
+++ b/tests/samplers_tests/tpe_tests/test_truncnorm.py
@@ -35,12 +35,14 @@ def test_ppf(a: float, b: float) -> None:
     "a,b",
     [(-np.inf, np.inf), (-10, +10), (-1, +1), (-1e-3, +1e-3), (10, 100), (-100, -10), (0, 0)],
 )
-def test_logpdf(a: float, b: float) -> None:
+@pytest.mark.parametrize("loc", [-10, 0, 10])
+@pytest.mark.parametrize("scale", [0.1, 1, 10])
+def test_logpdf(a: float, b: float, loc: float, scale: float) -> None:
     for x in np.concatenate(
         [np.linspace(np.max([a, -100]), np.min([b, 100]), num=1000), np.array([-2000.0, +2000.0])]
     ):
-        assert truncnorm_ours.logpdf(x, a, b) == pytest.approx(
-            truncnorm_scipy.logpdf(x, a, b), nan_ok=True
+        assert truncnorm_ours.logpdf(x, a, b, loc, scale) == pytest.approx(
+            truncnorm_scipy.logpdf(x, a, b, loc, scale), nan_ok=True
         ), f"logpdf(x={x}, a={a}, b={b})"