o
    }oi
                     @  s\   d Z ddlmZ ddlZddlmZ ddlmZ dddZ	dddZ
dddZdddZdS )a  
The functions in this file are mostly based on BoTorch v0.13.0,
but they are refactored significantly from the original version.

For ``_get_upper_bound_set``, look at:
    * https://github.com/pytorch/botorch/blob/v0.13.0/botorch/utils/multi_objective/box_decompositions/utils.py#L101-L160

For ``_get_box_bounds``, look at:
    * https://github.com/pytorch/botorch/blob/v0.13.0/botorch/utils/multi_objective/box_decompositions/utils.py#L163-L193

For ``_get_non_dominated_box_bounds``, look at:
    * https://github.com/pytorch/botorch/blob/v0.13.0/botorch/utils/multi_objective/box_decompositions/non_dominated.py#L395-L430

The preprocessing for four or fewer objectives, we use the algorithm proposed by:
    Title: A Box Decomposition Algorithm to Compute the Hypervolume Indicator
    Authors: Renaud Lacour, Kathrin Klamroth, and Carlos M. Fonseca
    URL: https://arxiv.org/abs/1510.01963
We refer this paper as Lacour17 in this file.

    )annotationsN)optuna_warn)_is_pareto_frontsorted_pareto_sols
np.ndarray	ref_pointreturntuple[np.ndarray, np.ndarray]c                   s   | j \} t tj tdddddf< d fdd}t|g}td  ftj }||df< | D ]
}||||\}}q<||fS )a!  
    This function follows Algorithm 2 of Lacour17.

    Args:
        sorted_pareto_sols: Pareto solutions sorted with respect to the first objective.
        ref_point: The reference point.

    Returns:
        upper_bound_set: The upper bound set, which is ``U(N)`` in the paper. The shape is
        ``(n_bounds, n_objectives)``.
        def_points: The defining points of each vector in ``U(N)``. The shape is
        ``(n_bounds, n_objectives, n_objectives)``.

    NOTE:
        ``pareto_sols`` corresponds to ``N`` and ``upper_bound_set`` corresponds to ``U(N)`` in the
        paper.
        ``def_points`` (the shape is ``(n_bounds, n_objectives, n_objectives)``) is not well
        explained in the paper, but basically, ``def_points[i, j] = z[j]`` of
        ``upper_bound_set[i]``.
    )dtypeTNr   solr   ubsdpsr   r	   c                   s   t j| |k dd}t|s||fS || }|jd }| t jt t j |ddk}t t |d d t j	f  | }t |df| }|j
|j
ksNJ t |j
}	|| }
| |
|	|f< || | }| | ||	|f< t ||  |gt ||  |
gfS )Naxisr      )npallanyshapemaxwhereinftilearangenewaxissizevstack)r   r   r   is_dominateddominated_dpsn_boundsupdateubs_indices_to_updatedimensions_to_updateindices_for_sweepingnew_dpsnew_ubsn_objectivesobjective_indicesskip_ineq_judge Y/home/ubuntu/.local/lib/python3.10/site-packages/optuna/_hypervolume/box_decomposition.pyr"   ;   s   
 $(z$_get_upper_bound_set.<locals>.updater   )r   r   r   r   r   r   r   r	   )r   r   r   eyeboolasarrayfullr   )r   r   _r"   upper_bound_set
def_pointssolutionr,   r(   r-   _get_upper_bound_set   s   

r6   r3   r4   c                 C  s   | j d }|dksJ dtdg| j R }|d d ddf |dd d df< |d |dd d df< t|d \}}tjj|ddd d ||d f |dd d dd f< | d d dd f |dd d dd f< tj|d |d kdd }|d d |f S )Nr   r   8This function is used only for multi-objective problems.   r   r   r   )r   r   emptydiag_indicesmaximum
accumulater   )r3   r4   r   r)   boundsrowcol	not_emptyr,   r,   r-   _get_box_bounds`   s   
 4&rA   c           	      C  sb   t | |d  }tj|dd}t|tj}t |t|dd |d\}}t||| \}}||fS )Nr   r   Tassume_unique_lexsorted)r   r   )r6   r   unique	full_liker   r   rA   )	r   r   neg_upper_bound_setsorted_neg_upper_bound_setpoint_at_infinityneg_lower_bound_setneg_def_pointsbox_upper_boundsbox_lower_boundsr,   r,   r-   _get_non_dominated_box_boundsp   s   


rM   	loss_valsc                 C  sj   t t | sJ dt j| dd}|t|dd }| jd }|dks(J d|d	kr0td
 t||S )Nz3loss_vals must be clipped before box decomposition.r   r   TrB   r   r   r7      zBox decomposition (typically used by `GPSampler`) might be significantly slow for n_objectives > 4. Please consider using another sampler instead.)r   r   isfiniterD   r   r   r   rM   )rN   r   unique_lexsorted_loss_valsr   r)   r,   r,   r-   get_non_dominated_box_bounds   s   


rR   )r   r   r   r   r   r	   )r3   r   r4   r   r   r   r   r   )rN   r   r   r   r   r	   )__doc__
__future__r   numpyr   optuna._warningsr   optuna.study._multi_objectiver   r6   rA   rM   rR   r,   r,   r,   r-   <module>   s    

B
