o
    Ú­jH  ã                   @   s¢   d Z ddlmZmZ ddlZddlmZ ddl	m
Z
 ddlmZmZmZmZ dd	lmZ d
eddfdd„Zd
eddfdd„Zd
eddfdd„ZG dd„ deƒZdS )z*Tests related to the `DataIter` interface.é    )ÚCallableÚOptionalN)Útestingé   )Úimport_cupy)ÚDataIterÚDMatrixÚExtMemQuantileDMatrixÚQuantileDMatrixé   )Úpredictor_equalÚdeviceÚreturnc                    s  t  ddd¡\}}}|  d¡rt  dddd¡\}}n
t  dddd¡\}}t  ddd	d¡\}}|||g}|||g}	|  d¡rItƒ ‰ ‡ fd
d„|D ƒ}t j||	dddd�}
t|
ƒ}t  dddd¡\}}|||g}|||g}	tj|dd�}tj|	dd�}t||ƒ}t	||ƒs…J ‚dS )zCheck QDM with mixed batches.é€   é   FÚcudaé   gš™™™™™¹?Té   gÍÌÌÌÌÌì?c                    s   g | ]}ˆ   |¡‘qS © )Úarray)Ú.0Úbatch©Úcpr   úV/var/www/html/CropPilot/venv/lib/python3.10/site-packages/xgboost/testing/data_iter.pyÚ
<listcomp>   s    z&run_mixed_sparsity.<locals>.<listcomp>N©ÚcacheÚon_hostr   )Úaxis)
ÚtmÚmake_regressionÚ
startswithÚmake_sparse_regressionr   ÚIteratorForTestr
   ÚnpÚconcatenater   )r   ÚX_0Úy_0Ú_ÚX_1Úy_1ÚX_2Úy_2ÚXÚyÚitÚXy_0ÚX_arrÚy_arrÚXy_1r   r   r   Úrun_mixed_sparsity   s&   






r5   c                    sÖ   G ‡ fdd„dt ƒ}|ƒ }ddl}|jtdd�� t|dd� W d  ƒ n1 s)w   Y  |jtdd�� t|dd� W d  ƒ n1 sFw   Y  |jtd	d�� t|dd� W d  ƒ dS 1 sdw   Y  dS )
z3Check error message for inconsistent feature types.c                       s>   e Zd Zd
‡ fdd„Zdedef‡fdd„Zd
dd	„Z‡  ZS )z2check_invalid_cat_batches.<locals>._InvalidCatIterr   Nc                    s   t ƒ jd d� d| _d S )N©Úcache_prefixr   )ÚsuperÚ__init__Ú_it©Úself©Ú	__class__r   r   r9   /   s   
z;check_invalid_cat_batches.<locals>._InvalidCatIter.__init__Ú
input_datac                    s€   | j dkrdS tjddddd| j dkrdndd	�\}}ˆ d
kr1dd l}dd l}| |¡}| |¡}|||d� |  j d7  _ dS )Nr   Fé@   é   é   g      à?r   g      ð?)ÚonehotÚsparsityÚ	cat_ratior   ©ÚdataÚlabelr   T)r:   r    Úmake_categoricalÚcudfÚcupyÚ	DataFramer   ©r<   r?   r.   r/   rJ   rK   ©r   r   r   Únext3   s$   

ú

z7check_invalid_cat_batches.<locals>._InvalidCatIter.nextc                 S   ó
   d| _ d S ©Nr   ©r:   r;   r   r   r   ÚresetI   ó   
z8check_invalid_cat_batches.<locals>._InvalidCatIter.reset©r   N)	Ú__name__Ú
__module__Ú__qualname__r9   r   ÚboolrO   rS   Ú__classcell__r   rN   r=   r   Ú_InvalidCatIter.   s    r[   r   Nz)Inconsistent number of categories between)ÚmatchT)Úenable_categoricalzInconsistent feature types)r   ÚpytestÚraisesÚ
ValueErrorr	   r
   r   )r   r[   r0   r^   r   rN   r   Úcheck_invalid_cat_batches+   s   ÿÿ"ÿra   c                    s²   ‡ fdd„dD ƒ}t t|Ž ƒ}tj|d |d dddd	�}t|ƒ}| ¡ d
ks)J ‚| ¡ tdd„ |d D ƒƒks:J ‚t|ƒ}| ¡ d
ksFJ ‚| ¡ tdd„ |d D ƒƒksWJ ‚dS )z'Tests for having irregular data shapes.c                    s    g | ]}t j|d ˆ dkd�‘qS )r   r   )Úuse_cupy)r    r!   )r   Ú	n_samplesrN   r   r   r   [   s    ÿÿz&check_uneven_sizes.<locals>.<listcomp>)r   r   i   r   r   Nr   Tr   r   c                 s   ó   � | ]}|j d  V  qdS ©r   N©Úshape©r   Úxr   r   r   Ú	<genexpr>d   ó   € z%check_uneven_sizes.<locals>.<genexpr>c                 s   rd   re   rf   rh   r   r   r   rj   h   rk   )	ÚlistÚzipr    r$   r   Únum_colÚnum_rowÚsumr	   )r   ÚbatchesÚunzipr0   ÚXyr   rN   r   Úcheck_uneven_sizesY   s   
þ"&rt   c                       sx   e Zd ZdZdededededededed	ed
ee ddf‡ fdd„Z	de
fdd„Zdedefdd„Zddd„Z‡  ZS )ÚCatIterz-An iterator for testing categorical features.Ún_samples_per_batchÚ
n_featuresÚ	n_batchesÚn_catsrD   rE   rC   r   r   r   Nc                   sÊ   t ƒ j|	d� || _|| _|| }
tj|
|||||d�\}}g g }}d}t|ƒD ]*}t||
| ƒ}|j||| …d d …f }| 	|¡ | 	|||| … ¡ ||7 }q)|| _
|| _|| _|| _d| _d S )Nr6   )Ún_categoriesrC   rE   rD   r   )r8   r9   rx   r   r    rI   ÚrangeÚminÚilocÚappendÚxsÚysri   r/   r:   )r<   rv   rw   rx   ry   rD   rE   rC   r   r   rc   Úcatr/   r   r€   Úprevr)   Únr.   r=   r   r   r9   n   s2   
ú



zCatIter.__init__c                 C   s   | j | jfS )zReturn the concatenated data.)ri   r/   r;   r   r   r   Úxyš   s   z
CatIter.xyr?   c                 C   sv   | j | jkrdS | j| j  | j| j  }}| jdkr,dd l}dd l}| |¡}| |¡}|||d� |  j d7  _ dS )NFr   r   rF   r   T)	r:   rx   r   r€   r   rJ   rK   rL   r   rM   r   r   r   rO   ž   s   


zCatIter.nextc                 C   rP   rQ   rR   r;   r   r   r   rS   ­   rT   zCatIter.resetrU   )rV   rW   rX   Ú__doc__ÚintÚfloatrY   Ústrr   r9   Útupler„   r   rO   rS   rZ   r   r   r=   r   ru   k   s4    þýûúùø	÷
öõô,ru   )r…   Útypingr   r   Únumpyr%   Úxgboostr   r    Úcompatr   Úcorer   r   r	   r
   Úutilsr   rˆ   r5   ra   rt   ru   r   r   r   r   Ú<module>   s    .