import optuna

def objective(trial):
    x = trial.suggest_float("x", -10, 10)
    y = trial.suggest_float("y", -10, 10)
    z = trial.suggest_float("z", -10, 10, step=0.5)
    return x ** 2 + y ** 2 + z ** 2

sampler = optuna.samplers.TPESampler(seed=42)
study = optuna.create_study(sampler=sampler)
study.optimize(objective, n_trials=2000)

import math
import numpy as np
import matplotlib.pyplot as plt
import mpmath
import optuna


_norm_pdf_C = math.sqrt(2 * math.pi)
_norm_pdf_logC = math.log(_norm_pdf_C)
_ndtri_exp_approx_C = math.sqrt(3) / math.pi
_log_2 = math.log(2)


def _ndtri_exp(y: np.ndarray) -> np.ndarray:
    # Flip the sign of y close to zero for better numerical stability and flip back the sign later.
    flipped = y > -1e-2
    z = y.copy()
    z[flipped] = np.log(-np.expm1(y[flipped]))
    x = np.empty_like(y)
    if (small_inds := np.nonzero(z < -5))[0].size:
        x[small_inds] = -np.sqrt(-2.0 * (z[small_inds] + _norm_pdf_logC))
    if (moderate_inds := np.nonzero(z >= -5))[0].size:
        x[moderate_inds] = -_ndtri_exp_approx_C * np.log(np.expm1(-z[moderate_inds]))

    for _ in range(100):
        log_ndtr_x = optuna.samplers._tpe._truncnorm._log_ndtr(x)
        log_norm_pdf_x = -0.5 * x**2 - _norm_pdf_logC
        # NOTE(nabenabe): Use exp(log_ndtr_x - log_norm_pdf_x) instead of ndtr_x / norm_pdf_x for
        # numerical stability.
        dx = (log_ndtr_x - z) * np.exp(log_ndtr_x - log_norm_pdf_x)
        x -= dx
        if np.all(np.abs(dx) < 1e-8 * np.abs(x)):  # NOTE: rtol controls the precision.
            # Equivalent to np.isclose with atol=0.0 and rtol=1e-8.
            break
    x[flipped] *= -1
    # NOTE(nabe): x[y == 0.0] = np.inf, x[np.isneginf(y)] = -np.inf are necessary for the accurate
    # computation, but we omit them as the ppf applies clipping, removing the need for them.
    return x


def _ndtri_exp_single_mp(y):
    a = -1e9
    b = +1e9
    for _ in range(1000):
        m = (a + b) / 2
        if mpmath.log(mpmath.ncdf(m)) < y:
            a = m
        else:
            b = m
    return (a + b) / 2

mpmath.mp.dps = 100

y = np.asarray([-(10**i) for i in np.arange(-50, 10, 0.1)])
x_master = np.asarray([optuna.samplers._tpe._truncnorm._ndtri_exp_single(yi) for yi in y])
x_pr = _ndtri_exp(y)
x_mp = np.asarray([_ndtri_exp_single_mp(yi) for yi in y])
plt.plot(-y, x_master - x_mp, label="master")
plt.plot(-y, x_pr - x_mp, label="PR")
plt.xscale("log")
plt.legend()
plt.savefig("6229.png")