
    e                     j   d dl Zd dlmZmZmZmZ d dlZddl	m
ZmZ ddlmZmZmZ d dlZej                            dej        ej        fd           d	             Zd
 Zd Zd Zej                            dej        ej        ej        fd           d             Zd Zej                            d ej         d          d           d             Z!ddZ"d Z#ej                            de          ej                            dej        ej        fd           d                         Z$dS )    N)assert_equalassert_array_equalassert_array_almost_equalassert_raises   )
unit_makerunit_maker_argparse)arraysarray_orderDTYPESfuncc                     | j         S N__name__xs    Dlib/python3.11/site-packages/bottleneck/tests/nonreduce_axis_test.py<lambda>r          1:     )idsc                    d}|dz  }| j         }t          d|z            }t          j                            g d          }t          t          |                    D ]a\  }}|j        dk    s|j        dk    s|j        dk    r(t          t          d|j                            dgz   D ]}||j        d	z
  }n|j        |         d	z
  }|d	k     r'|                    |          }	 |||	|          }
 | ||	|          }|d
k    r%t          |
||	|          }
t          |||	|          }n"t          |
|	|          }
t          ||	|          }|dt          |          z   t          |j                  t          |j                  |	t          |          t#          |          |f}||z  }t%          ||
|           cdS )ztest partition or argpartitionz@
func %s | input %s (%s) | shape %s | n %d | axis %s | order %s

Input array:
%s

bn.slow.%s)r         r   r   Nr   argpartitiona)r   evalnprandomRandomState	enumerater
   ndimsizelistrangeshaperandintcomplete_the_argpartitioncomplete_the_partitionstrdtyper   r   )r   msgnamefunc0rsir    axisnmaxns0s1tuperr_msgs                 r   test_partition_and_argpartitionr<      s    OC!!C=D$%%E			yyy	)	)B&,,'' 0 016Q;;!&A++!r16**++tf4 	0 	0D|vzwt}q(axx

4  Aq!T""BaD!!B~%%.r1a>>.r1a>>+B488+B488c!ffAGAGD		A	C CiGr2w////7	00 0r   c                     d }|                                  } | j        }||dk    rt          d          d}n |dk     r||z  }|dk     rt          d          t          j        ||| |          } | S )Nc                     t          j        | d |                   | d |<   t          j        | |dz   d                    | |dz   d <   | S )Nr   )r"   sort)r    r7   s     r   func1dz&complete_the_partition.<locals>.func1dD   sK    "1""1"WQq1uwwZ((!a%''
r   r   z `a` must be 1d when axis is Noner   `axis` out of range)copyr&   
ValueErrorr"   apply_along_axis)r    r7   r5   r@   r&   s        r   r-   r-   C   s      
 	
A6D|199?@@@	!882333
FD!Q//AHr   c           
         |                                 }|j        }|4| j        dk    rt          d          d}d}|                    d          }n |dk     r||z  }|dk     rt          d          |dk    r
||          }n|dk    r|dk    r:t	          |j        d                   D ]}|| d d |f         |f         |d d |f<   n|dk    r2t	          |j        d                   D ]}||| |         f         ||<   nQt          d          |dk    r,|dk    rYt	          |j        d                   D ]=}t	          |j        d                   D ] }|| d d ||f         ||f         |d d ||f<   !>n|dk    rYt	          |j        d                   D ]=}t	          |j        d                   D ] }||| |d d |f         |f         ||d d |f<   !>n}|dk    rYt	          |j        d                   D ]=}t	          |j        d                   D ] }|||| ||d d f         f         |||d d f<   !>nt          d          t          d          t          |||          }|S )	Nr   z$`index` must be 1d when axis is Noner   r   rA   r   r   z`a.ndim` must be 1, 2, or 3)rB   r&   rC   reshaper)   r*   r-   )indexr    r7   r5   r&   r4   js          r   r,   r,   W   s   	A6D|:??CDDDIIbMM	!882333qyyeH	199171:&& , ,E!!!Q$KN+!!!Q$,QYY171:&& & &E!H~!& 2333	199171:&& 9 9qwqz** 9 9A!"5Aq>1a#7!8AaaaAgJJ99 QYY171:&& 9 9qwqz** 9 9A!"1eAqqq!Gna#7!8AaAgJJ99 QYY171:&& 9 9qwqz** 9 9A!"1aq!QQQw#7!8AaAAAgJJ99 23336777q!T**AHr   c                     t          j        d                              dd          } t          j        | j        dd          j        }t          j                            | j        dd          j        }t          ||d            dS )zpartition transpose test      r   r   r   N)r"   arangerF   bn	partitionTslowr   )r    actualdesireds      r   test_transposerS      sn    
	"a##A\!#q"%%'FgQ++-G"<=====r   c                     | j         S r   r   r   s    r   r   r      s    !* r   c                      t          |           S )zTest nonreduce axis functions)reduce_unit_maker)r   s    r   test_nonreduce_axisrW      s    
 T"""r   c                     d} t          j        t           j        ddt           j        t           j        t           j        t           j        dt           j        g	          }| D ]~}t          j        |                                |          }t          j                            |                                |          }t          ||dt          |          z             dS )z	Test push)r   r   r   r   rK      Nr   r   r   )r7   zfailed on n=%sN)	r"   arraynanrM   pushrB   rP   r   r.   )nsr    r7   rQ   rR   s        r   	test_pushr^      s    	!B
"&!Q26JKKA G GQ'''',,qvvxx1,--67,<s1vv,EFFFFG Gr   nonreduce_axisc                     | j         S r   r   r   s    r   r   r      s    aj r   c                     | j         }|dv rt          |           S |dv rt          |           S |dv rt          |           S d}t          ||z            )z'test argument parsing in nonreduce_axisrN   r   r\   )rankdatananrankdataz+``%s` is an unknown nonreduce_axis function)r   unit_maker_parseunit_maker_parse_rankdatarC   unit_maker_raises)r   r1   fmts      r   test_arg_parsingri      so    
 =D,,,%%%	&		%%%	,	,	,(...;t$$$r   rY   c                 ^   | j         }t          d|z            }t          j        g d          }d| z  }|dz  }|d|z  z  } | |d          } ||d          }|dz  }t	          ||||            | |dd	          } ||dd	          }|d
z  }t	          ||||           |dk    ru | |dd          } ||dd          }|dz  }t	          ||||            | |dd	          } ||dd	          }|dz  }t	          ||||           |ddf}	i }
 | |	i |
 dS |df}	i }
 | |	i |
 dS )ztest argument parsing.r   g      ?r   r   z
%sz%s
r   r   z(a, 1)r   r5   z(a, 1, axis=0)r\   r   Nz(a, 2, None)z(a, 1, axis=None)r   )r   r!   r"   rZ   r   )r   decimalr1   r2   r    rh   rQ   rR   r;   argskwargss              r   re   re      s    =D$%%E
A
4-C6MC!A%%CT!QZZFeAqkkGHnGfgw@@@T!QQFeAqq!!!G$$Gfgw@@@v~~aD!!%1d##&!&'7GDDDa&&&%14(((++!&'7GDDD 1bzdf
 1vdfr   c           
      j   t          j        g d          }t          t          |            t          t          | |           t          t          | |dd           t          t          | |d|           | j        dv r5t          t          | |ddddd           t          t          | |d           dS dS )	z.test argument parsing raises in nonreduce_axisrk   rl   r   )r5   extra)r5   r    rb   0N)r"   rZ   r   	TypeErrorr   )r   r    s     r   rg   rg      s    
A)T""")T****)T11A6666)T112222}555iq!Q1a888iqs333333 65r   r/   c                     | j         S r   r   r   s    r   r   r      r   r   c                 @   t          j        d|          }dD ]}t          j        t          d          5   | |d|           d d d            n# 1 swxY w Y   t          j        t          d          5   | |d|           d d d            n# 1 swxY w Y   d S )	N)
   rv   )r/   )Nr   r   r   zmust be between)matchi  rl   r   )r"   onespytestraisesrC   )r   r/   rZ   r5   s       r   test_out_of_bounds_raisesr{      s8   
 GHE***E  ' ']:->??? 	) 	)D4((((	) 	) 	) 	) 	) 	) 	) 	) 	) 	) 	) 	) 	) 	) 	) ]:->??? 	' 	'D&&&&	' 	' 	' 	' 	' 	' 	' 	' 	' 	' 	' 	' 	' 	' 	'	' 's#   AA	A	7BB	B	)rY   )%numpyr"   numpy.testingr   r   r   r   
bottleneckrM   reduce_testr   rV   r	   rf   utilr
   r   r   ry   markparametrizerN   r   r<   r-   r,   rS   rc   rd   r\   rW   r^   get_functionsri   re   rg   r{    r   r   <module>r      sf                             . - - - - - - - - -  
R\2?+1E1E   (0 (0 (0V  (* * *Z> > > 
R["."'28L8L   # # #
G G G 
B-..4H4H   # # #, , , ,^	4 	4 	4 &))
R\2?+1E1E   ' '  *)' ' 'r   