U
    2dUB                     @   sT  d Z ddlmZmZ ddlmZ ddlmZ ddlZddl	m
Z
 ddlmZmZ ddlmZ dd	lmZmZmZ dd
lmZ ddlmZ ddgddgddgddgddgddggZddddddgZddgddgddgddgddgddggZddddddgZddddddgZdd Zdd Zdd Zdd Zdd Z dd Z!dd Z"d d! Z#d"d# Z$dS )$zG
Testing for export functions of decision trees (sklearn.tree.export).
    )finditersearch)dedent)RandomStateN)is_classifier)DecisionTreeClassifierDecisionTreeRegressor)GradientBoostingClassifier)export_graphviz	plot_treeexport_text)StringIO)NotFittedError         g      ?c               
   C   s  t ddddd} | tt t| d d}d}||ks8tt| ddgd d	}d
}||ksZtt| ddgd d}d}||ks|tt| dddddd dd}d}||kstt| ddd d}d}||kstt| ddd dd}d}||kstt ddddd} | jtttd} t| ddd d}d}||ks(ttddddd} | tt t| ddd dddd}d}||ksltt dd} | tt	 t| dd d }d!}d S )"Nr   r   gini	max_depthmin_samples_split	criterionrandom_stateout_filea  digraph Tree {
node [shape=box, fontname="helvetica"] ;
edge [fontname="helvetica"] ;
0 [label="x[0] <= 0.0\ngini = 0.5\nsamples = 6\nvalue = [3, 3]"] ;
1 [label="gini = 0.0\nsamples = 3\nvalue = [3, 0]"] ;
0 -> 1 [labeldistance=2.5, labelangle=45, headlabel="True"] ;
2 [label="gini = 0.0\nsamples = 3\nvalue = [0, 3]"] ;
0 -> 2 [labeldistance=2.5, labelangle=-45, headlabel="False"] ;
}Zfeature0Zfeature1)feature_namesr   a  digraph Tree {
node [shape=box, fontname="helvetica"] ;
edge [fontname="helvetica"] ;
0 [label="feature0 <= 0.0\ngini = 0.5\nsamples = 6\nvalue = [3, 3]"] ;
1 [label="gini = 0.0\nsamples = 3\nvalue = [3, 0]"] ;
0 -> 1 [labeldistance=2.5, labelangle=45, headlabel="True"] ;
2 [label="gini = 0.0\nsamples = 3\nvalue = [0, 3]"] ;
0 -> 2 [labeldistance=2.5, labelangle=-45, headlabel="False"] ;
}yesno)class_namesr   a  digraph Tree {
node [shape=box, fontname="helvetica"] ;
edge [fontname="helvetica"] ;
0 [label="x[0] <= 0.0\ngini = 0.5\nsamples = 6\nvalue = [3, 3]\nclass = yes"] ;
1 [label="gini = 0.0\nsamples = 3\nvalue = [3, 0]\nclass = yes"] ;
0 -> 1 [labeldistance=2.5, labelangle=45, headlabel="True"] ;
2 [label="gini = 0.0\nsamples = 3\nvalue = [0, 3]\nclass = no"] ;
0 -> 2 [labeldistance=2.5, labelangle=-45, headlabel="False"] ;
}TFZsans)filledimpurity
proportionZspecial_charactersroundedr   fontnamea  digraph Tree {
node [shape=box, style="filled, rounded", color="black", fontname="sans"] ;
edge [fontname="sans"] ;
0 [label=<x<SUB>0</SUB> &le; 0.0<br/>samples = 100.0%<br/>value = [0.5, 0.5]>, fillcolor="#ffffff"] ;
1 [label=<samples = 50.0%<br/>value = [1.0, 0.0]>, fillcolor="#e58139"] ;
0 -> 1 [labeldistance=2.5, labelangle=45, headlabel="True"] ;
2 [label=<samples = 50.0%<br/>value = [0.0, 1.0]>, fillcolor="#399de5"] ;
0 -> 2 [labeldistance=2.5, labelangle=-45, headlabel="False"] ;
}r   )r   r   r   zdigraph Tree {
node [shape=box, fontname="helvetica"] ;
edge [fontname="helvetica"] ;
0 [label="x[0] <= 0.0\ngini = 0.5\nsamples = 6\nvalue = [3, 3]\nclass = y[0]"] ;
1 [label="(...)"] ;
0 -> 1 ;
2 [label="(...)"] ;
0 -> 2 ;
})r   r    r   Znode_idsa;  digraph Tree {
node [shape=box, style="filled", color="black", fontname="helvetica"] ;
edge [fontname="helvetica"] ;
0 [label="node #0\nx[0] <= 0.0\ngini = 0.5\nsamples = 6\nvalue = [3, 3]", fillcolor="#ffffff"] ;
1 [label="(...)", fillcolor="#C0C0C0"] ;
0 -> 1 ;
2 [label="(...)", fillcolor="#C0C0C0"] ;
0 -> 2 ;
})Zsample_weight)r    r!   r   a  digraph Tree {
node [shape=box, style="filled", color="black", fontname="helvetica"] ;
edge [fontname="helvetica"] ;
0 [label="x[0] <= 0.0\nsamples = 6\nvalue = [[3.0, 1.5, 0.0]\n[3.0, 1.0, 0.5]]", fillcolor="#ffffff"] ;
1 [label="samples = 3\nvalue = [[3, 0, 0]\n[3, 0, 0]]", fillcolor="#e58139"] ;
0 -> 1 [labeldistance=2.5, labelangle=45, headlabel="True"] ;
2 [label="x[0] <= 1.5\nsamples = 3\nvalue = [[0.0, 1.5, 0.0]\n[0.0, 1.0, 0.5]]", fillcolor="#f1bd97"] ;
0 -> 2 [labeldistance=2.5, labelangle=-45, headlabel="False"] ;
3 [label="samples = 2\nvalue = [[0, 1, 0]\n[0, 1, 0]]", fillcolor="#e58139"] ;
2 -> 3 ;
4 [label="samples = 1\nvalue = [[0.0, 0.5, 0.0]\n[0.0, 0.0, 0.5]]", fillcolor="#e58139"] ;
2 -> 4 ;
}Zsquared_error)r    Zleaves_parallelr   rotater#   r$   aT  digraph Tree {
node [shape=box, style="filled, rounded", color="black", fontname="sans"] ;
graph [ranksep=equally, splines=polyline] ;
edge [fontname="sans"] ;
rankdir=LR ;
0 [label="x[0] <= 0.0\nsquared_error = 1.0\nsamples = 6\nvalue = 0.0", fillcolor="#f2c09c"] ;
1 [label="squared_error = 0.0\nsamples = 3\nvalue = -1.0", fillcolor="#ffffff"] ;
0 -> 1 [labeldistance=2.5, labelangle=-45, headlabel="True"] ;
2 [label="squared_error = 0.0\nsamples = 3\nvalue = 1.0", fillcolor="#e58139"] ;
0 -> 2 [labeldistance=2.5, labelangle=45, headlabel="False"] ;
{rank=same ; 0} ;
{rank=same ; 1; 2} ;
}r   )r    r   zdigraph Tree {
node [shape=box, style="filled", color="black", fontname="helvetica"] ;
edge [fontname="helvetica"] ;
0 [label="gini = 0.0\nsamples = 6\nvalue = 6.0", fillcolor="#ffffff"] ;
})
r   fitXyr
   AssertionErrory2wr   
y_degraded)clfZ	contents1Z	contents2 r/   B/tmp/pip-unpacked-wheel-zrfo1fqw/sklearn/tree/tests/test_export.pytest_graphviz_toy   s                   

r1   c               	   C   sR  t ddd} t }tt t| | W 5 Q R X | tt d}tjt	|d t| d dgd W 5 Q R X d}tjt	|d t| d dd	d
gd W 5 Q R X d}tjt
|d t| ttj W 5 Q R X t }tt t| |g d W 5 Q R X t }tjt	dd t| |dd W 5 Q R X tjt	dd t| |dd W 5 Q R X d S )Nr   r   )r   r   z?Length of feature_names, 1 does not match number of features, 2matchar   z?Length of feature_names, 3 does not match number of features, 2bczis not an estimator instance)r   zshould be greater or equalr   )	precisionzshould be an integer1)r   r   pytestraisesr   r
   r'   r(   r)   
ValueError	TypeErrorZtree_
IndexError)r.   outmessager/   r/   r0   test_graphviz_errors   s,    rA   c                  C   s   t ddd} | tt t }t| |d tddd} | tt | jD ]}t|d |d qHtd|	 D ]}d|
 ksltqld S )Nfriedman_mser   )r   r   r   r   )Zn_estimatorsr   z\[.*?samples.*?\])r   r'   r(   r)   r   r
   r	   Zestimators_r   getvaluegroupr*   )r.   dot_dataZ	estimatorfindingr/   r/   r0   test_friedman_mse_in_graphviz#  s    
rG   c            	      C   s8  t d} t d}t| d|df| d|jdddftdd	d
dtd
d	dfD ]\}}}||| dD ]}t|d |dd}td|D ]&}t	t
d|  |d
 kstqt|rd}nd}t||D ]&}t	t
d|  |d
 kstqtd|D ]*}t	t
d|  |d
 kstqqpqVd S )Nr      )   r   )     )rI   )rJ   )sizerB   r   r   )r   r   r   r   r   )rK   r   T)r   r8   r"   zvalue = \d+\.\d+z\.\d+zgini = \d+\.\d+zfriedman_mse = \d+\.\d+z<= \d+\.\d+)r   zipZrandom_samplerandintr   r   r'   r
   r   lenr   rD   r*   r   )	Zrng_regZrng_clfr(   r)   r.   r8   rE   rF   patternr/   r/   r0   test_precision2  s<      
   $$rR   c               	   C   s   t ddd} | tt d}tjt|d t| dd W 5 Q R X d}tjt|d t| d	gd
 W 5 Q R X d}tjt|d t| dd W 5 Q R X d}tjt|d t| dd W 5 Q R X d S )Nr   r   rM   z max_depth bust be >= 0, given -1r2   r   r&   z,feature_names must contain 2 elements, got 1r4   r5   zdecimals must be >= 0, given -1decimalszspacing must be > 0, given 0spacing)r   r'   r(   r)   r:   r;   r<   r   )r.   err_msgr/   r/   r0   test_export_text_errors^  s    rX   c                  C   sV  t ddd} | tt td }t| |ks4tt| dd|ksHtt| dd|ks\ttd }t| dd	gd
|ksttd }t| dd|ksttd }t| dd|kstddgddgddgddgddgddgddgg}dddddddg}t ddd} | || td }t| dd|ks:tddgddgddgddgddgddgg}ddgddgddgddgddgddgg}tddd}||| td }t|dd|kstt|ddd|kstdgdgdgdgdgdgg}tddd}||| td }t|ddgd|ks6tt|dddgd|ksRtd S )Nr   r   rM   zh
    |--- feature_1 <= 0.00
    |   |--- class: -1
    |--- feature_1 >  0.00
    |   |--- class: 1
    r&   
   zX
    |--- b <= 0.00
    |   |--- class: -1
    |--- b >  0.00
    |   |--- class: 1
    r4   r6   r5   z
    |--- feature_1 <= 0.00
    |   |--- weights: [3.00, 0.00] class: -1
    |--- feature_1 >  0.00
    |   |--- weights: [0.00, 3.00] class: 1
    T)show_weightsz\
    |- feature_1 <= 0.00
    | |- class: -1
    |- feature_1 >  0.00
    | |- class: 1
    r   rU   r   r   rK   z{
    |--- feature_1 <= 0.00
    |   |--- class: -1
    |--- feature_1 >  0.00
    |   |--- truncated branch of depth 2
    zy
    |--- feature_1 <= 0.0
    |   |--- value: [-1.0, -1.0]
    |--- feature_1 >  0.0
    |   |--- value: [1.0, 1.0]
    rS   )rT   rZ   zq
    |--- first <= 0.0
    |   |--- value: [-1.0, -1.0]
    |--- first >  0.0
    |   |--- value: [1.0, 1.0]
    first)rT   r   )rT   rZ   r   )	r   r'   r(   r)   r   lstripr   r*   r   )r.   Zexpected_reportZX_lZy_lZX_moZy_moregZX_singler/   r/   r0   test_export_textp  s`    	.((r^   c                 C   s   t ddddd}|tt ddg}t||d}t|dks@t|d  d	ksTt|d
  dksht|d  dks|td S )Nr   r   Zentropyr   
first featsepal_widthr5   r   z:first feat <= 0.0
entropy = 1.0
samples = 6
value = [3, 3]r   z(entropy = 0.0
samples = 3
value = [3, 0]z(entropy = 0.0
samples = 3
value = [0, 3]r   r'   r(   r)   r   rP   r*   Zget_textpyplotr.   r   Znodesr/   r/   r0   test_plot_tree_entropy  s        
rd   c                 C   s   t ddddd}|tt ddg}t||d}t|dks@t|d  d	ksTt|d
  dksht|d  dks|td S )Nr   r   r   r   r_   r`   r5   r   z7first feat <= 0.0
gini = 0.5
samples = 6
value = [3, 3]r   z%gini = 0.0
samples = 3
value = [3, 0]z%gini = 0.0
samples = 3
value = [0, 3]ra   rb   r/   r/   r0   test_plot_tree_gini  s        
re   c              	   C   s(   t  }tt t| W 5 Q R X d S )N)r   r:   r;   r   r   )rc   r.   r/   r/   r0   test_not_fitted_tree  s    rf   )%__doc__rer   r   textwrapr   Znumpy.randomr   r:   Zsklearn.baser   Zsklearn.treer   r   Zsklearn.ensembler	   r
   r   r   ior   Zsklearn.exceptionsr   r(   r)   r+   r,   r-   r1   rA   rG   rR   rX   r^   rd   re   rf   r/   r/   r/   r0   <module>   s2   (( d',c