o
    ×tBhÀB  ã                   @   sD  d dl Zd dlZd dlmZ d dlmZ d dlmZ d dlm	Z	 d dl
mZ d dlmZ ej dd	d
g¡ej deddƒ¡dd„ ƒƒZej dejdg¡dd„ ƒZej dejdg¡dd„ ƒZej dejdg¡dd„ ƒZej dejdg¡dd„ ƒZej dejdg¡dd„ ƒZej dejdg¡dd„ ƒZej dejdg¡dd„ ƒZej dejdg¡d d!„ ƒZd"d#„ Zej d$dd g¡ej ddejg¡ej d%¡d&d'„ ƒƒƒZej ddejg¡ej dd	d
g¡d(d)„ ƒƒZej ddejg¡d*d+„ ƒZej d$dd g¡ej ddejg¡d,d-„ ƒƒZ ej d.d/ejd0fg¡d1d2„ ƒZ!dS )3é    N)Úconfig_context)Ú
KNNImputer)Únan_euclidean_distances)Úpairwise_distances)ÚKNeighborsRegressor)Úassert_allcloseÚweightsÚuniformÚdistanceÚn_neighborsé   é   c                 C   sL   d}d}t j ||¡}t j|d< t|| d}| |¡}|j||fks$J ‚d S )Né
   é   )r   r   )r   r   )ÚnpÚrandomÚrandÚnanr   Úfit_transformÚshape)r   r   Ún_rowsÚn_colsÚXÚimputerÚ	X_imputed© r   út/var/www/html/riverr-enterprise-integrations-main/venv/lib/python3.10/site-packages/sklearn/impute/tests/test_knn.pyÚtest_knn_imputer_shape   s   

r   Únaéÿÿÿÿc                 C   s€  t  t jddd| gg d¢g d¢| ddddg| d	dd	d
gg d¢g¡}tjtdd t| d |¡ W d   ƒ n1 s;w   Y  t  t jddd| gg d¢g d¢| ddddg| d	dd	d
gg d¢g¡}t  dddd| gg d¢g d¢| ddddg| d	dd	d
gg d¢g¡}t| d |¡}tjtdd | |¡ W d   ƒ n1 s w   Y  tjtdd t| dd |¡ W d   ƒ n1 sÀw   Y  tdddd}t  t j	ddddgt j	ddt j	dgt j	ddddgt j	ddddgg¡}d}tjt|d | |¡ W d   ƒ n	1 sw   Y  t  ddgt j	dgg¡}tdd}d}tjt|d | |¡ W d   ƒ d S 1 s9w   Y  d S )Nr   r   )r   r   r   r   é   )r    r   r    r    é   r   r   é   é   é   r!   )r   r   r   r"   r$   zInput X contains (infinity|NaN))Úmatch©Úmissing_valueszExpected n_neighbors > 0©r'   r   r	   ©r'   r   r   r    zInput X contains NaNÚ	euclidean)Úmetricz/The selected metric does not support NaN values)
r   ÚarrayÚinfÚpytestÚraisesÚ
ValueErrorr   ÚfitÚ	transformr   )r   r   ÚX_fitr   ÚmsgÚbad_metric_msgr   r   r   Ú+test_knn_imputer_default_with_invalid_input   sv   úÿ
ÿúÿúÿ
ÿÿüÿÿþÿ
$ÿr6   c                 C   sº   t  dd| dddgdd| dddgdd| dd| gdd| | ddgg¡}t| dd |¡}| |¡}t  |¡ ¡ r8J ‚|jdks?J ‚t  d	d
¡ 	dd¡}| |¡}t
|d d …g d¢f |ƒ d S )Nr   ç      ð?r   r    é   r   r(   )r8   r"   r   é   ©r   r   r    r8   r"   )r   r,   r   r1   r2   ÚisnanÚanyr   ÚarangeÚreshaper   )r   r   ÚknnÚX_transformÚX_testr   r   r   Ú(test_knn_imputer_removes_all_na_featuresg   s   üÿ

rB   c              
   C   s¾   t  g d¢g d¢g d¢g d¢g¡}t  d| dddgg d¢dddd| gdd| ddgg¡}t  g d	¢g d¢g d
¢g d¢g¡}tdddd}t| ddd}t| |¡|ƒ t| |¡| |¡ƒ d S )N)r   r   r   r   r7   )r   r   r   r   r   )r    r    r    r    r   )r   r   r   r   r   r   r7   r    r   )r   ç      @r   r   r7   )r    r    r    r    ç      ø?)r   r   rC   r   r   r   r   r	   r)   ©r   r,   r   r   r   )r   ÚX_zeroÚX_nanr   Úimputer_zeroÚimputer_nanr   r   r   Ú*test_knn_imputer_zero_nan_imputes_the_same|   s6   üÿ	üÿ	üÿ	ÿrJ   c                 C   sè  t  g d¢ddd| gddd| g| dddgd| ddgg d	¢g d
¢g¡}t  g d¢g d¢g d¢g d¢g d¢g d	¢g d
¢g¡}t| d}t| |¡|ƒ t  ddd| gddd| gddd| gddd| gddd| gddd| gg d¢g d¢g¡}d}t  ddd|gddd|gddd|gddd|gddd|gddd|gg d¢g d¢g¡}t| d}t| |¡|ƒ t  ddg| dgddgddgddgddgddgg¡}t  ddgddgd| gg¡}d}t  ddgddgd|gg¡}t| d}t| |¡ |¡|ƒ d S )N©r   r   r   r   r   r   r    r8   r"   r   r$   )r!   r!   r!   r!   )é   é   é   é   )r   r   r   r!   )r    r   r    r!   )r8   r8   r"   r"   )r   r    r   r$   r&   r   r!   )é   rP   rP   rP   )é   rQ   rQ   rQ   g      5@é	   é   rL   g333333@)r   r,   r   r   r   r1   r2   )r   r   r   r   ÚX_impute_valueÚX1ÚX_2_1Ú
X1_imputedr   r   r   Útest_knn_imputer_verify¤   sj   



ùÿùÿ






øÿ





øÿ
4
rX   c              
   C   sˆ   t  ddg| dgddgd| gddg| dgdd	gg¡}t  ddgddgddgddgddgddgdd	gg¡}td
| d}t| |¡|ƒ d S )Nr   r   r8   r    r"   r$   r!   é   r#   r   ©r   r'   rE   )r   r   r   r   r   r   r   Ú test_knn_imputer_one_n_neighborsî   s   44r[   c              
   C   s¼   t  ddg| dgddgd| gddg| dgdd	gg¡}t  ddgd
dgddgddgddgd
dgdd	gg¡}|jd d }t|| d}t| |¡|ƒ |jd }t|| d}t| |¡|ƒ d S )Nr   r   r8   r    r"   r$   r!   rY   r#   r   g      @r   rZ   )r   r,   r   r   r   r   )r   r   r   r   r   Úimputer_plus1r   r   r   Ú*test_knn_imputer_all_samples_are_neighborsú   s   44
r]   c              
   C   sÐ   t  ddg| dgddgddgddgdd	gd
dgg¡}t  ddgddgddgddgddgdd	gd
dgg¡}td| d}t| |¡|ƒ dd„ }t|| d}t| |¡|ƒ dd„ }t|| d}t| |¡|ƒ d S )Nr   r   r8   r    r"   r   r$   rR   r!   rS   r   r	   ©r   r'   c                 S   s   d S ©Nr   ©Údistr   r   r   Ú	no_weight  s   z2test_knn_imputer_weight_uniform.<locals>.no_weightc                 S   s
   t  | ¡S r_   )r   Ú	ones_liker`   r   r   r   Úuniform_weight  s   
z7test_knn_imputer_weight_uniform.<locals>.uniform_weightrE   )r   r   ÚX_imputed_uniformr   rb   rd   r   r   r   Útest_knn_imputer_weight_uniform
  s   4,ÿrf   c           $      C   sÂ  t  ddg| dgddgddgddgdd	gd
dgg¡}tddd}g d¢}| ||dd …f ||df ¡ | |dd…dd …f ¡d }g d¢}t|dd…d d …f || d}d|d d …|f  ¡  }t j||df |d}t  ddg|dgddgddgddgdd	gd
dgg¡}	t  ddg|dgddgddgddgdd	gd
dgg¡}
td| d}t	| 
|¡|	ƒ t	| 
|¡|
ƒ t  | ddgg d¢g d¢g d¢g¡}t  d¡}t  d¡}t jddgd| d| gd}t  |ddgg d¢g d¢g d¢g¡}tdd| d}t	| 
|¡|ƒ t  g d¢d| d| gddd| gg d¢g d¢g d¢g d¢g¡}t|| d}|dg d¢f }|dg d ¢f }d| }d| }|dg d ¢f }d| }t j |g d¢df ¡ ¡ }t j |g d ¢df ¡ ¡ }t jj||d}t jj||d}t jj||d}t  g d¢d|d|gddd|gg d¢g d¢g d¢g d¢g¡}td| d}t	| 
|¡|ƒ t  ddd| gddd| gdd| dgg d!¢g d"¢g d#¢g d$¢| dddgg¡}t|d%d&| d'}d(|ddd)…f  }d(|ddd)…f  }d(|d*  }d(|ddd…f  }t j|dd)…d)f |d}t j|dd)…d)f |d} t j|d+ |d}!t j|dd…df |d}"t  ddd|gddd| gdd|!dgg d!¢g d"¢g d#¢g d$¢|"dddgg¡}t| dd,}#t	|# 
|¡|ƒ d S )-Nr   r   r8   r    r"   r   r$   rR   r!   rS   r   r*   r
   )r+   r   )r   r   r    r8   r"   r   r   )r   r   r    r8   r"   r&   ©r   r^   )r   r   r   )r    r   r    )r8   r"   r"   g      @g     €3@)r   r   r'   rK   )r   r   r   r   )r   r   r   r   )r   r   r   r   )r   r   r   r   )r   r    r8   r"   r   ©r    r    r    r    ©r8   r8   r8   r8   ©r"   r"   r"   r"   ©r   r   r   r   Únan_euclideanF©r+   Úsquaredr'   r7   r   )r   r:   )r:   r   )r'   r   )r   r,   r   r1   Úpredictr   ÚravelÚaverager   r   r   ÚsqrtÚmaÚmasked_invalidÚcopyr   )$r   r   ÚnnÚ
X_rows_idxÚknn_imputed_valueÚX_neighbors_idxra   r   Úmanual_imputed_valueÚX_imputed_distance1ÚX_imputed_distance2r   Údist_0_1Údist_0_2Úimputed_valuer   Úr1c1_nbor_distsÚr1c3_nbor_distsÚr1c1_nbor_wtÚr1c3_nbor_wtÚr2c3_nbor_distsÚr2c3_nbor_wtÚcol1_donor_valuesÚcol3_donor_valuesÚr1c1_impÚr1c3_impÚr2c3_impÚr0c3_wÚr1c3_wÚr2c2_wÚr7c0_wÚr0c3Úr1c3Úr2c2Úr7c0Úimputer_comp_wtr   r   r   Ú test_knn_imputer_weight_distance&  sÔ   4 ,ÿ,ÿüÿ


üÿ	

ùÿ

ùÿ



øÿÿ



øÿr”   c               	   C   s   t jdfdd„} t  dddt jgg d¢g d¢t jdd	d
gg¡}d}d}t  ddd|gg d¢g d¢|dd	d
gg¡}td| d}t| |¡|ƒ d S )NFc                 S   sD   t jj| t  | ¡d} t jj|t  |¡d}t  t  | | ¡¡}|S )N)Úmask)r   rs   r,   r;   ÚnansumÚabs)ÚxÚyr'   rn   ra   r   r   r   Úcustom_callable¾  s   z9test_knn_imputer_callable_metric.<locals>.custom_callabler8   r    )r   rR   r   rR   )r8   r!   r   rR   rR   rS   g      $@g      "@g      @r   )r   r+   )r   r   r,   r   r   r   )rš   r   ÚX_0_3ÚX_3_0r   r   r   r   r   Ú test_knn_imputer_callable_metric»  s   ."ÿr   Úworking_memoryzignore:adhere to working_memoryc           
      C   sL  t  d| d| gddd| gdd| dgg d¢g d¢g d¢g d¢| dddgg¡}t  |dd	…df ¡}t  |dd
…d
f ¡}t  |dd
…d
f ¡}t  |g d¢df ¡}t  |dd
…df ¡}t  d|d|gddd|gdd|dgg d¢g d¢g d¢g d¢|dddgg¡}t|d t| d}	t|	 |¡|ƒ W d   ƒ d S 1 sŸw   Y  d S )Nr   r   r   rh   ri   rj   rk   r$   r   r   r:   ©rž   r&   )r   r,   Úmeanr   r   r   r   )
r   rž   r   Úr0c1r   r   r‘   r’   r   Úimputer_compr   r   r   Ú$test_knn_imputer_with_simple_exampleÐ  s>   



øÿ



øÿ
"þr£   c                 C   sŠ   t  | dg| dgd| gg¡}t  ddgddgddgg¡}t| d|d}t| |¡|ƒ t  d| gg¡}t  ddgg¡}t| |¡|ƒ d S )NrS   r   r    r   r)   r8   ©r   r,   r   r   r   r2   )r   r   rU   rW   r?   ÚX2Ú
X2_imputedr   r   r   Ú+test_knn_imputer_not_enough_valid_distancesü  s   r§   c                 C   s€   t  | dg| dgg¡}t| dd}t  dgdgg¡}t| |¡|ƒ t  ddgd| gg¡}t  dgdgg¡}t| |¡|ƒ d S )Nr   r   r(   r    rD   r¤   )r   rU   r?   ÚX1_expectedr¥   ÚX2_expectedr   r   r   Ú'test_knn_imputer_drops_all_nan_features  s   rª   c              	   C   s€  t  d| gd| g| dgddgddg| dgg¡}t|dd| d	}t j|dd…d
f d
|ddd…f  d}t j|dd…d
f d
|d
dd…f  d}t j|dd…df d
|ddd…f  d}t j|dd…df d
|ddd…f  d}t  d|gd|g|dgddgddg|dgg¡}t|d' t| ddd}	t|	 |¡|ƒ t| ddd}
t|
 |¡|ƒ W d   ƒ d S 1 s¹w   Y  d S )Nr    r   r8   r"   r   r!   rl   Frm   r   r   rg   rŸ   r
   r)   )r   r,   r   rq   r   r   r   r   )r   rž   r   ra   ÚX_01ÚX_11ÚX_20ÚX_50Ú
X_expectedÚknn_3Úknn_4r   r   r   Ú7test_knn_imputer_distance_weighted_not_enough_neighbors  s   .ÿ****."ûr²   zna, allow_nan)r   FTc                 C   s"   t | d}| ¡ d |ksJ ‚d S )Nr&   Ú	allow_nan)r   Ú	_get_tags)r   r³   r?   r   r   r   Útest_knn_tags/  s   
rµ   )"Únumpyr   r.   Úsklearnr   Úsklearn.imputer   Úsklearn.metrics.pairwiser   r   Úsklearn.neighborsr   Úsklearn.utils._testingr   ÚmarkÚparametrizeÚranger   r   r6   rB   rJ   rX   r[   r]   rf   r”   r   Úfilterwarningsr£   r§   rª   r²   rµ   r   r   r   r   Ú<module>   sV    
K

'
I



 
'
