
    biH                     X   d Z ddlZddlmZmZ ddlmZmZmZm	Z	m
Z
 ddlZ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 e
d
ej2                  defd       Ze
d
ej8                  defd       Zd
e	ej2                  ej8                  f   defdZdeddfdZdededdfdZ de!de!de!dededdfdZ"de!de!de!de!dedededdfdZ#de!dejH                  dejH                  deddf
d Z%ded!eddfd"Z&dededdfd#Z' ejP                  ejR                        jT                  Z+dZ,d$e!de!d%e!ded&eded'ed(ed)edejZ                  fd*Z.d+e!d,e!d-e!d.e!deded(eddfd/Z/d+e!d,e!d.e!deded(eddfd0Z0d1eeef   d2ejZ                  d3e!deeef   fd4Z1 G d5 d6ejd                  jf                        Z4y)7zTests for updaters.    N)partialupdate_wrapper)AnyDictListUnionoverload)is_pd_cat_dtype   )DataIter   )CatItermodelreturnc                      y N r   s    R/home/cdr/jupyterlab/.venv/lib/python3.12/site-packages/xgboost/testing/updater.pyget_basescorer      s    14    c                      y r   r   r   s    r   r   r      s    03r   c                     t        | t        j                        r| j                         } t	        t        j                  | j                               d   d   d         }|S )z1Get base score from an XGBoost sklearn estimator.learnerlearner_model_param
base_score)
isinstancexgbXGBModelget_boosterfloatjsonloadssave_config)r   r   s     r   r   r      sU    %&!!#

5$$&'	23HI,WJ r   tree_methodc                 p    ddl m}m}m} dt        j
                  dt        j
                  ddf fd} |dd	
      \  }} |||        |ddd	      \  }} |||       dt        j
                  dt        j
                  ddf fd} |dd	
      \  }} |||        |dddd	      \  }} |||       y)zTest for init estimation.r   )make_classificationmake_multilabel_classificationmake_regressionXyr   Nc                    t        j                  dd      }|j                  | || |fg       t        |      }|j	                         d   d   d   }t        j                  ddd      }|j                  | || |fg       t        |      }|j	                         d   d   d   }t        j                  ||      rJ ||k  sJ y )Nr   r%   	max_depthn_estimatorseval_setvalidation_0rmser   r%   r.   r/   boost_from_average)r   XGBRegressorfitr   evals_resultnpisclose)r*   r+   regbase_score_0score_0base_score_1score_1r%   s          r   run_regz&check_init_estimation.<locals>.run_reg-   s    ;!RST1Ax($S)""$^4V<Q?#qqUV
 	1Ax($S)""$^4V<Q?::lL999   r         )	n_samplesrandom_state   )rC   	n_targetsrD   c                    t        j                  dd      }|j                  | || |fg       t        |      }|j	                         d   d   d   }t        j                  ddd      }|j                  | || |fg       t        |      }|j	                         d   d   d   }t        j                  ||      rJ ||k  sJ y )Nr   r-   r0   r2   loglossr   r4   )r   XGBClassifierr7   r   r8   r9   r:   )r*   r+   clfr<   r=   r>   r?   r%   s          r   run_clfz&check_init_estimation.<locals>.run_clfC   s    K1STU1Ax($S)""$^4Y?B#qqUV
 	1Ax($S)""$^4Y?B::lL999   r      )rC   n_labels	n_classesrD   )sklearn.datasetsr'   r(   r)   r9   ndarray)r%   r'   r(   r)   r@   r*   r+   rK   s   `       r   check_init_estimationrQ   %   s     !2:: !"** ! !  T;DAqAqMTQRHDAqAqM!2:: !"** ! !  B?DAqAqM)abDAq AqMr   weightedc                 
   ddl m} ddlm} ddlm} d}d}d}d}t        j                  j                  d      }	 ||||		      \  }
}|r|	j                  |
      }nd}t        j                  |
||      }t        j                  ddg      }i }t        j                  d| ||d|||dfg|      }|j                  |d      }t        j                  |d   d         sJ |d   d   d   dk  sJ t!        |j"                        D cg c]#  } |t%        t'        ||||         |            % }}t        j(                  |j*                        }t!        |j*                  d         D ]  }||   }t        j                  d| ||d|||dfg||   |      }t        j                  |d   d         sJ |d   d   d   dk  sJ t        j,                  j/                  t        j                  |d   d         t        j                  |d   d         dd       |j                  |      |dd|f<    t!        |j*                  d         D ]0  }t        j,                  j/                  |dd|f   |dd|f          2 yc c}w )zTest for quantile loss.r   )r)   )mean_pinball_loss)_metric_decoratorrA              i  )rC   
n_featuresrD   )sizeNweightg?      ?reg:quantileerror)	objectiver%   quantile_alphar   Trainnum_boost_roundevalsr8   T)strict_shapequantileg      4@)sample_weightalpha)rb   rc   custom_metricr8   g      >@rT   gư>)atolrtol)rO   r)   sklearn.metricsrT   xgboost.sklearnrU   r9   randomRandomStater   QuantileDMatrixarraytrainpredicttmnon_increasingrangerY   r   r   emptyshapetestingassert_allclose)r%   rR   r)   rT   rU   rC   rX   r/   r   rngr*   r+   r[   Xyrh   r8   booster_multipredt_multiimetricspredtsa	booster_is                          r   check_quantile_lossr   \   s   011IJL J
))


%CDAq
 +			Q&	1BHHc3Z E$&LII,&#$		
 	$G}o!M  '''>K\'2:>??? ,R04777 uzz"  	)uQxP!	
G  XXk''(F5;;q>" -!HII0*"#(	 (=/!!*%
	   g!6z!BCCCG$Z04t;;;


""HH\'*:67HH\'*+>?@	 	# 	
 !((,q!t1-4 5;;q>" D


""6!Q$<QT1BCDKs   	(J n_samples_per_batchrX   	n_batchesdevicec           	         t        j                  t        j                  | |||dk7        ddd}t        j                  |      }|d|ddgd}t        j
                  ||      }|j                         \  }	}
}t        j                  |	|
|	      }t        j
                  ||      }|j                  |      }|j                  |      }t        j                  j                  ||       y
)z2Check external memory with the quantile objective.cpucacheFr   on_hostr]   g?皙?)r%   r^   r   r_   rZ   N)rt   IteratorForTestmake_batchesr   DMatrixrr   	as_arraysrs   r9   ry   rz   )r   rX   r   r%   r   itXy_itparams
booster_itr*   r+   wr|   boosterpredt_itpredts                   r   check_quantile_loss_extmemr      s     
			,j)Vu_	U
B
 KKOE"(*	F 65)JllnGAq!	Q!	$Bii#G!!%(HOOBEJJuh/r   n_binsr   is_catc                   |rt        | ||dddd|d	      }n0t        j                  t        j                  | |||dk7        d|d	}t	        j
                  |||
      }t        j                  t        d      5  t	        j                  |d|d|d      }	ddd       t	        j                  ||d|d      }	|rt        | ||dddd|d	      }n/t        j                  t        j                  | |||dk7        ddi}t	        j                  |||
      }
t	        j                  ||d|
d      }|j                         }|
j                         }t        j                  j                  |d   |d          t        j                  j                  |d   |d          |	j                  |      }|j                  |
      }t        j                  j                  ||       y# 1 sw Y   MxY w)z+Basic test for the `ExtMemQuantileDMatrix`.rL   rW   r\   Fr   )	r   rX   r   n_catssparsity	cat_ratioonehotr   r   r   use_cupyr   max_binenable_categoricalzOnly the `hist`)matchapprox)r   r%   r   rV   )rb   N)r   r   r   r   )r   rt   r   r   r   ExtMemQuantileDMatrixpytestraises
ValueErrorrr   rp   get_quantile_cutr9   ry   rz   rs   )r   rX   r   r   r   r   r   r   r   r   r|   r   cut_itcutr   r   s                   r   check_extmem_qdmr      s     3!

 __#ZVu_ 
 %%b&VTE	z):	; 
YYh6J


 f-uaJ  3!

 __#ZVu_
 	
 
		RF	KBii6f=rSTUG##%F



CJJvay#a&1JJvay#a&1!!%(HOOBEJJx/O
 
s   ;G##G-	n_entriesindptrdatadtypesc                    |j                   d   |d   k(  sJ |j                   d   | k(  sJ |j                  t        j                  k(  sJ t	        d|j
                        D ]t  }t        ||dz
           }t        ||         }t	        |dz   |      D ]A  }||   ||dz
     kD  sJ t        |j                  |dz
           s/||   ||dz
     dz   k(  rAJ  v y)zCheck the cut values.r   rf   r   N)	rx   dtyper9   uint64rv   rY   intr
   iloc)r   r   r   r   r   begendjs           r   	check_cutr     s     ::a=F2J&&&::a=I%%%<<299$$$1fkk" 2&Q- &)nsQw$ 	2A7T!a%[(((v{{1q512Aw$q1u+/111	22r   r   c                    ddl }d}d}d}|j                  t        j                  g|z        }t	        j
                  |||      \  }}}	t        j                  |||	|      }
|
j                         \  }}t        |dz   |z  |||       t        j                  |||		      }
t        j                  | |d
|
       |
j                         \  }}t        |dz   |z  |||       d}d}t	        j                  t	        j                  ||||      ddd}t        j                  |      }t        j                  | |d
|
       |
j                         \  }}t        |dz   |z  |||       d}t	        j                  |||dd      \  }}|r*ddl}ddl}|j#                  |      }|j%                  |      }t        j                  |||d      }|j                         \  }}t        ||z  |||j&                         t        j                  ||d      }t        j                  | |d
|       |j                         \  }}t        ||z  |||j&                         t	        j                  |||ddd      \  }}t)        |j&                  D cg c]  }t+        |      sd c}      }||z
  }||z  |dz   |z  z   }t        j                  |||d      }|j                         \  }}t        ||||j&                         t        j                  ||d      }t        j                  | |d
|       |j                         \  }}t        ||||j&                         yc c}w )zCheck with optional cupy.r   Ni         r   )r[   r   r   rZ   )r%   r   rE      r   Fr       r   )r   r   Tr   r   r\   )r   r   r   )pandasSeriesr9   float32rt   r)   r   rp   r   r   r   rr   r   r   make_categoricalcudfcupyfrom_pandasrq   r   lenr
   )r%   r   pdrC   rX   r   r   r*   r+   r   Xywr   r   r   r   r   r|   n_categoriesr   cpr   n_cat_featuresn_num_featuresr   s                           r   check_get_quantile_cut_devicer   /  s!   IJGYY

|j01F   JJGAq!**1a7KC'')LFDw{j(&$?
++a1
%CIIkg>D'')LFDw{j(&$?I				,j)X	N
B
 kk"oBIIkg>D'')LFDw{j(&$? L:|ECDAq QHHQK			Q7t	LB&&(LFDlZ'qxx@	Qd	3BIIkg>C&&(LFDlZ'qxx@ :|ECSVDAq L_U5K!LMN.0N~-10NNI			Q7t	LB&&(LFDiqxx0	Qd	3BIIkg>C&&(LFDiqxx0 Ms   	MMc                 \    |j                  d      }t        | d       |rt        | d       yy)zCheck the quantile cut getter.cudaFTN)
startswithr   )r%   r   r   s      r   check_get_quantile_cutr   w  s0       (H!+u5%k48 r   rC   r   r   r   extmemr   c                X   t        t        d|       d      }	t        | |	z  ||	|d||||rdnd 	      }
|rK|dk(  rt        j                  |
|      }|S |dk(  rt        j
                  |
|      }|S t        d	| d
      |
j                         \  }}t        j
                  |||      }|S )Nr   r         ?r   )r   r   r   r   r   r   r   histr   r   ztree_method z not supported.)maxminr   r   r   r   r   xy)rC   rX   r   r   r   r%   r   r   r   r   r   r|   catlabels                 r   _create_dmatrixr     s     C9%q)I	YgT

B & !77'9B I H$R4FGB I	 |K=HIIUUW
U[[e8JKIr   rowscolsroundscatsc                    i }i }|t         |d}	t        | |||dd||d	      }
t        j                  |	|
||
dfg|       t        | |||d|d|d	      }t        j                  |	|||dfg|       t        j
                  j                  t	        j                  |d   d	         t	        j                  |d   d	         d
       t        j                  |d   d	         sJ i }t        |	d<   d|	d<   t        j                  |	|||dfg|       |d   d	   }|d   d	   }t        ||      D ]  \  }}||k\  rJ  d|	d<   i }t        j                  |	|d|dfg|       t        j                  |d   d	         sJ |       y)z0Test for one-hot encoding with categorical data.)r%   max_cat_to_onehotr   rW   TF)r   r   r   r   r%   r   r   r`   ra   )r   r   r   r%   r   r   r   r3   gMbP?rk   r   r   
reg_lambdar   r   N)
USE_ONEHOTr   r   rr   r9   ry   rz   rq   rt   ru   USE_PARTzip)r   r   r   r   r   r%   r   by_etl_resultsby_builtin_results
parameters	Xy_onehotXy_catby_groupingrmse_oh
rmse_groupr   bs                    r   check_categorical_oher     s    9;N<> #'	"J   
I II7#$# 
F II !' JJ
(01
#G,V45  
 /8@AAA57K&.J"# J|II !  !)&1GW%f-JGZ( 1Avv  #J|KII !  [1&9:GKG:r   c                   	 ||d	t        | ||d||d|d	      j                         dt        ddf	fd	} |t                |t               y)
z+Check categorical data with missing values.)r%   r   r\   FT)r   r   r   r%   r   r   r   r   r   Nc                 ,   | d<   i }t        j                  ddfg|      }t        j                  |d   d         sJ |j	                        }t        j
                  |      }t        j                  j                  ||d   d   d   d       y )	Nr   r   r`   ra   r3   rf   gh㈵>r   )	r   rr   rt   ru   rs   root_mean_squarer9   ry   rz   )r   r8   r   y_predtr3   r|   r   r   s        r   runz&check_categorical_missing.<locals>.run(  s    *;
&'(*))=/%
   g!6v!>???//"%""5'2


""4g)>v)Fr)JQU"Vr   )r   	get_labelr   r   r   )
r   r   r   r   r%   r   r   r|   r   r   s
          @@@r   check_categorical_missingr     si     2=!OJ	

B LLNEWs Wt W& 
O Mr   paramdmat
num_roundsc                 "   i }t        j                  | |||dfgd|      }|j                         |j                         k(  sJ |j	                         |k(  sJ |j
                  |j
                  k(  sJ |j                  |j                  k(  sJ |S )z-Get training result from parameters and data.rr   F)rc   verbose_evalr8   )r   rr   num_featuresnum_colnum_boosted_roundsfeature_namesfeature_types)r   r   r   resultr   s        r   train_resultr  A  s      FiigG !T\\^333%%':555  D$6$6666  D$6$6666Mr   c                   <    e Zd ZdZdej
                  dededefdZ	y)ResetStrategyz"Callback for testing multi-output.r   epoch	evals_logr   c                 b    |dz  dk(  r|j                  ddi       y|j                  ddi       y)Nr   r   multi_strategymulti_output_treeone_output_per_treeF)	set_param)selfr   r  r  s       r   after_iterationzResetStrategy.after_iterationY  s=    19>OO-/BCD  OO-/DEFr   N)
__name__
__module____qualname____doc__r   Boosterr   dictboolr  r   r   r   r
  r
  V  s+    ,S[[   RV r   r
  )5r  r"   	functoolsr   r   typingr   r   r   r   r	   numpyr9   r   xgboostr   xgboost.testingry   rt   xgboost.datar
   corer   	data_iterr   r   r!   r   r  strrQ   r  r   r   r   r   rP   r   r   r   iinfoint32r   r   r   r   r   r   r   r  callbackTrainingCallbackr
  r   r   r   <module>r(     sO     - 3 3     (   
 4 4% 4 
 4 
 3 3 3 
 3s||S[[89 e 3s 3t 3nTDS TDD TDT TDn000 0 	0
 0 
0>I0I0I0 	I0
 I0 I0 I0 I0 
I0X22JJ2.0jj2BE2	2"E1s E1d E1t E1P9 9S 9T 9 RXXbhh##
$$$ 	$
 $ $ $ $ $ $ 	[[$NbH
bH bH 	bH
 bH bH bH bH 
bHJ.
.
. .
 . . . 
.bS>!$:=	#s(^*CLL11 r   