o }oiA@sdZddlmZddlZddlmZddlmZddlZddl m Z ddl m Z ddl mZer?dd lmZddlZddlZndd lmZed Zed ZeeZd%ddZGdddejjZGdddZ  d&d'd#d$ZdS)(aNotations in this Gaussian process implementation X_train: Observed parameter values with the shape of (len(trials), len(params)). y_train: Observed objective values with the shape of (len(trials), ). x: (Possibly batched) parameter value(s) to evaluate with the shape of (..., len(params)). cov_fX_fX: Kernel matrix X = V[f(X)] with the shape of (len(trials), len(trials)). cov_fx_fX: Kernel matrix Cov[f(x), f(X)] with the shape of (..., len(trials)). cov_fx_fx: Kernel scalar value x = V[f(x)]. This value is constant for the Matern 5/2 kernel. cov_Y_Y_inv: The inverse of the covariance matrix (V[f(X) + noise_var])^-1 with the shape of (len(trials), len(trials)). cov_Y_Y_inv_Y: `cov_Y_Y_inv @ y` with the shape of (len(trials), ). max_Y: The maximum of Y (Note that we transform the objective values such that it is maximized.) sqd: The squared differences of each dimension between two points. is_categorical: A boolean array with the shape of (len(params), ). If is_categorical[i] is True, the i-th parameter is categorical. ) annotationsN)Any) TYPE_CHECKING)*single_blas_thread_if_scipy_v1_15_or_newer) optuna_warn) get_logger)Callable) _LazyImportscipytorchvalues np.ndarrayreturnc Cs~t|}t|r |Stdtj|dd}t|t|tjt||tjdddt|tj t||tj dddS)NzDClip non-finite values to the min/max finite values for GP fittings.r)axis) npisfiniteallranyclipwheremininfmax)r is_values_finite is_any_finiterA/home/ubuntu/.local/lib/python3.10/site-packages/optuna/_gp/gp.pywarn_and_convert_inf/s  "$rc@s(eZdZed ddZed d d Zd S)Matern52Kernelctxrsquared_distance torch.TensorrcCsLtd|}t| }|d||d}d|d|}|||S)a This method calculates `exp(-sqrt5d) * (1/3 * sqrt5d ** 2 + sqrt5d + 1)` where `sqrt5d = sqrt(5 * squared_distance)`. Please note that automatic differentiation by PyTorch does not work well at `squared_distance = 0` due to zero division, so we manually save the derivative, i.e., `-5/6 * (1 + sqrt5d) * exp(-sqrt5d)`, for the exact derivative calculation. Notice that the derivative of this function is taken w.r.t. d**2, but not w.r.t. d. g?g)r sqrtexpsave_for_backward)r r!sqrt5dexp_partvalderivrrrforward@s   zMatern52Kernel.forwardgradcCs|j\}||S)z Let x be squared_distance, f(x) be forward(ctx, x), and g(f) be a provided function, then deriv := df/dx, grad := dg/df, and deriv * grad = df/dx * dg/df = dg/dx. ) saved_tensors)r r-r+rrrbackwardSszMatern52Kernel.backwardN)r rr!r"rr")r rr-r"rr")__name__ __module__ __qualname__ staticmethodr,r/rrrrr?s  rc@s\eZdZd(d d Zed)d dZd*ddZ d+d,ddZd-d.ddZd/ddZ d0d&d'Z dS)1 GPRegressoris_categoricalr"X_trainy_traininverse_squared_lengthscales kernel_scale noise_varrNonecCs||_||_||_|d|d|_|jr/|jd|jfdktj |jd|jf<d|_ d|_ ||_ ||_ ||_dS)N.r)_is_categorical_X_train_y_train unsqueezesquare__squared_X_diffrtyper float64 _cov_Y_Y_chol_cov_Y_Y_inv_Yr8r9r:)selfr5r6r7r8r9r:rrr__init__^s   zGPRegressor.__init__r cCsdt|jS)Ng?)rr%r8detachcpunumpy)rHrrr length_scalesvszGPRegressor.length_scalescCs |jdur |jdusJdt|}Wdn1s'wY|t |j j d|j 7<tj|}tjj|jtjj||jdddd}t||_t||_|j|_d|j_|j|_d|j_|j |_ d|j _dS)Nz(Cannot call cache_matrix more than once.rT)lowerF)rFrGr no_gradkernelrJrKrLr diag_indicesr?shaper:itemlinalgcholeskyr solve_triangularTr@ from_numpyr8r-r9)rHcov_Y_Y cov_Y_Y_chol cov_Y_Y_inv_Yrrr _cache_matrixzs* $       zGPRegressor._cache_matrixNX1torch.Tensor | NoneX2cCs|dur|dus J|j}n3|dur|j}|jdkr||n |d|d}|jrA|d|jfdktj |d|jf<| |j }t ||jS)am Return the kernel matrix with the shape of (..., n_A, n_B) given X1 and X2 each with the shapes of (..., n_A, len(params)) and (..., n_B, len(params)). If x1 and x2 have the shape of (len(params), ), kernel(x1, x2) is computed as: kernel_scale * Matern52Kernel.apply( sqd(x1, x2) @ inverse_squared_lengthscales ) where if x1[i] is continuous, sqd(x1, x2)[i] = (x1[i] - x2[i]) ** 2 and if x1[i] is categorical, sqd(x1, x2)[i] = int(x1[i] != x2[i]). Note that the distance for categorical parameters is the Hamming distance. Nr$r<r=.r)rCr?ndimrArBr>rrDr rEmatmulr8rapplyr9)rHr]r_sqdsqdistrrrrPs *  zGPRegressor.kernelFxjointbool!tuple[torch.Tensor, torch.Tensor]c Cs|jdur |jdusJd|jdk}|s|n|d}tj||}|j}tjj|jtjj|jj |dddddd}|rb|rFJd|||}|| | d d } | j d d d  d n|j}|tj||} | d |r|d| dfS|| fS) a) This method computes the posterior mean and variance given the points `x` where both mean and variance tensors will have the shape of x.shape[:-1]. If ``joint=True``, the joint posterior will be computed. The posterior mean and variance are computed as: mean = cov_fx_fX @ inv(cov_fX_fX + noise_var * I) @ y, and var = cov_fx_fx - cov_fx_fX @ inv(cov_fX_fX + noise_var * I) @ cov_fx_fX.T. Please note that we clamp the variance to avoid negative values due to numerical errors. Nz+Call cache_matrix before calling posterior.r$rTF)upperleftz3Call posterior with joint=False for a single point.r<)dim1dim2r)rFrGr`rAr rTvecdotrPrVrWra transposediagonal clamp_min_r9squeeze) rHrerfis_single_pointx_ cov_fx_fXmeanV cov_fx_fxvar_rrr posteriors*      zGPRegressor.posteriorcCs|jjd}d|tdtj}||jtj|tj d}tj |}|  }tj j||jdddfdddddf}d||}|||S)a This method computes the marginal log-likelihood of the kernel hyperparameters given the training dataset (X, y). Assume that N = len(X) in this method. Mathematically, the closed form is given as: -0.5 * log((2*pi)**N * det(C)) - 0.5 * y.T @ inv(C) @ y = -0.5 * log(det(C)) - 0.5 * y.T @ inv(C) @ y + const, where C = cov_Y_Y = cov_fX_fX + noise_var * I and inv(...) is the inverse operator. We exploit the full advantages of the Cholesky decomposition (C = L @ L.T) in this method: 1. The determinant of a lower triangular matrix is the diagonal product, which can be computed with N flops where log(det(C)) = log(det(L.T @ L)) = 2 * log(det(L)). 2. Solving linear system L @ u = y, which yields u = inv(L) @ y, costs N**2 flops. Note that given `u = inv(L) @ y` and `inv(C) = inv(L @ L.T) = inv(L).T @ inv(L)`, y.T @ inv(C) @ y is calculated as (inv(L) @ y) @ (inv(L) @ y). In principle, we could invert the matrix C first, but in this case, it costs: 1. 1/3*N**3 flops for the determinant of inv(C). 2. 2*N**2-N flops to solve C @ alpha = y, which is alpha = inv(C) @ y. Since the Cholesky decomposition costs 1/3*N**3 flops and the matrix inversion costs 2/3*N**3 flops, the overall cost for the former is 1/3*N**3+N**2+N flops and that for the latter is N**3+2*N**2-N flops. rgdtypeNF)ri)r?rRmathlogpirPr:r eyerErTrUrpsumrVr@)rHn_pointsconstrYL logdet_partinv_L_y quad_partrrrmarginal_log_likelihoods  ,  z#GPRegressor.marginal_log_likelihood log_prior%Callable[[GPRegressor], torch.Tensor] minimum_noisefloatdeterministic_objectivegtolc s&jjdttjtj tj dgg}dfdd }t t jj||d d d |id }Wdn1sOwY|js_td |jt|j}t|d_t|_ rtjtjdn t|d_ S)Nr$gGz? raw_paramsr rtuple[float, np.ndarray]cst|d}tLt|d_t|_r)tjtjdn t|d_  }| |j d}rQ|dksQJWdn1s[wY| |j fS)NTr|r$r)r rXrequires_grad_ enable_gradr&r8r9tensorrEr:rr/r-rSrJrKrL)rraw_params_tensorlossraw_noise_var_gradrrrn_paramsrHrr loss_funcs  z1GPRegressor._fit_kernel_params..loss_funcTzl-bfgs-br)jacmethodoptionszOptimization failed: r|)rr rr)r?rRr concatenaterr8rJrKrLr9rSr:rr optimizeminimizesuccess RuntimeErrormessager rXrer&rrEr\) rHrrrrinitial_raw_paramsrresraw_params_opt_tensorrrr_fit_kernel_paramss:    zGPRegressor._fit_kernel_params)r5r"r6r"r7r"r8r"r9r"r:r"rr;)rr )rr;)NN)r]r^r_r^rr")F)rer"rfrgrrh)rr") rrrrrrgrrrr4) r0r1r2rIpropertyrMr\rPrzrrrrrrr4]s     %#r4{Gz?XYr5rrrrrrg gpr_cacheGPRegressor | Nonerc stjjddtjdd fdd }|} |dur!|}d} || fD]6} z tttt| j| j| jdj ||||d WSt y]} z| } WYd} ~ q'd} ~ wwt d | d |} | | S) Nr$r{r|rr4csBttttdddddS)Nr<rkr5r6r7r8r9r:)r4r rXclonerrrdefault_kernel_paramsr5rr _default_gprKs  z'fit_kernel_params.._default_gprr)rrrrz/The optimization of kernel parameters failed: z< The default initial kernel parameters will be used instead.)rr4)r onesrRrEr4rXr8r9r:rrloggerwarningr\)rrr5rrrrrrdefault_gpr_cacheerrorgpr_cache_to_usee default_gprrrrfit_kernel_params>s@      r)r r rr )Nr)rr rr r5r rrrrrrgrrrrrr4)__doc__ __future__rr~typingrrrLr"optuna._gp.scipy_blas_thread_patchroptuna._warningsroptuna.loggingrcollections.abcrr r optuna._importsr r0rrautogradFunctionrr4rrrrrs0          i