
    M/Ph                     P    d Z ddlmZ ddlZddlmZ ddlmc m	Z	 	 	 	 d
dZ
d	 ZdS )z@
Authors:    Josef Perktold, Skipper Seabold, Denis A. Engemann
    )lrangeN)rainbowmeanbbestc           	         ddl m} t          j        |          \  }}|pt	          |dd          }t	          |dt          |                    }| d| }|pt	          | dd          }|pt	          |dd          }|                    |           |                    |           d	x}}t          | d         t
                    red
 t          j
        |           D             }t          t          |                    }t          | t          t          ||                              }  |t          | ||                    }|                    ddg                              |                                          }t          |d         
                                          }|
dg|z  n|
}
|	dg|z  n|	}	|t'          |          n|}t          |
          |k    rt)          d          t          |	          |k    rt)          d          t          |          |k    rt)          d          |dk    s|dk    rt+          |                    d                    D ]Z\  }\  }}t          |d         j        d                   } |j        |d         |d         f||         |	|         ||
|         d| [n|dk    s|dk    rwt+          |                    d                    D ]S\  }\  }}t          |d         j        d                   } |j        |d         |d         f||         ||
|         d| Tn|dk    s|dk    rwt+          |                    d                    D ]S\  }\  }}t          |d         j        d                   } |j        |d         |d         f||         ||	|         d| Tnt)          d|z            |                    ||           |                    d           t7          ||g          r*|                    |           |                    |           |S )a  
    Interaction plot for factor level statistics.

    Note. If categorial factors are supplied levels will be internally
    recoded to integers. This ensures matplotlib compatibility. Uses
    a DataFrame to calculate an `aggregate` statistic for each level of the
    factor or group given by `trace`.

    Parameters
    ----------
    x : array_like
        The `x` factor levels constitute the x-axis. If a `pandas.Series` is
        given its name will be used in `xlabel` if `xlabel` is None.
    trace : array_like
        The `trace` factor levels will be drawn as lines in the plot.
        If `trace` is a `pandas.Series` its name will be used as the
        `legendtitle` if `legendtitle` is None.
    response : array_like
        The reponse or dependent variable. If a `pandas.Series` is given
        its name will be used in `ylabel` if `ylabel` is None.
    func : function
        Anything accepted by `pandas.DataFrame.aggregate`. This is applied to
        the response variable grouped by the trace levels.
    ax : axes, optional
        Matplotlib axes instance
    plottype : str {'line', 'scatter', 'both'}, optional
        The type of plot to return. Can be 'l', 's', or 'b'
    xlabel : str, optional
        Label to use for `x`. Default is 'X'. If `x` is a `pandas.Series` it
        will use the series names.
    ylabel : str, optional
        Label to use for `response`. Default is 'func of response'. If
        `response` is a `pandas.Series` it will use the series names.
    colors : list, optional
        If given, must have length == number of levels in trace.
    markers : list, optional
        If given, must have length == number of levels in trace
    linestyles : list, optional
        If given, must have length == number of levels in trace.
    legendloc : {None, str, int}
        Location passed to the legend command.
    legendtitle : {None, str}
        Title of the legend.
    **kwargs
        These will be passed to the plot command used either plot or scatter.
        If you want to control the overall plotting options, use kwargs.

    Returns
    -------
    Figure
        The figure given by `ax.figure` or a new instance.

    Examples
    --------
    >>> import numpy as np
    >>> np.random.seed(12345)
    >>> weight = np.random.randint(1,4,size=60)
    >>> duration = np.random.randint(1,3,size=60)
    >>> days = np.log(np.random.randint(1,30, size=60))
    >>> fig = interaction_plot(weight, duration, days,
    ...             colors=['red','blue'], markers=['D','^'], ms=10)
    >>> import matplotlib.pyplot as plt
    >>> plt.show()

    .. plot::

       import numpy as np
       from statsmodels.graphics.factorplots import interaction_plot
       np.random.seed(12345)
       weight = np.random.randint(1,4,size=60)
       duration = np.random.randint(1,3,size=60)
       days = np.log(np.random.randint(1,30, size=60))
       fig = interaction_plot(weight, duration, days,
                   colors=['red','blue'], markers=['D','^'], ms=10)
       import matplotlib.pyplot as plt
       #plt.show()
    r   )	DataFramenameresponse__name__z of XTraceNc                     g | ]}|S  r   ).0ls     `/var/www/html/test/jupyter/venv/lib/python3.11/site-packages/statsmodels/graphics/factorplots.py
<listcomp>z$interaction_plot.<locals>.<listcomp>l   s    ,,,!A,,,    )xtracer   r   r   -.z(Must be a linestyle for each trace levelz%Must be a marker for each trace levelz$Must be a color for each trace levelbothr   )colormarkerlabel	linestyleliner   )r   r   r   scatters)r   r   r   zPlot type %s not understood)loctitleg?)pandasr	   utilscreate_mpl_axgetattrstr
set_ylabel
set_xlabel
isinstancenpuniquer   len_recodedictzipgroupby	aggregatereset_indexr   
ValueError	enumeratevaluesplotr    legendmarginsall
set_xticksset_xticklabels)r   r   r   funcaxplottypexlabelylabelcolorsmarkers
linestyles	legendloclegendtitlekwargsr	   figresponse_name	func_namex_valuesx_levelsdata	plot_datan_traceir7   groupr   s                              r   interaction_plotrS      s   d !     !"%%GCCgh
CCMj#d))44I..}..F.wq&#..F@!@!@KMM&MM&Hx!A$ 6,,ry||,,,#h--((AtC(3344559TAUX>>>??Dgs^,,66t<<HHJJI )G$++--..G$.$6#JJ!(segoogG!'WWVF
:'!!CDDD
7||w@AAA
6{{g?@@@6X__"+I,=,=g,F,F"G"G 	7 	7Ag-a011EBGE#Jj 1 7"1:U(m7 7/57 7 7 7	7 
V		x3"+I,=,=g,F,F"G"G 	D 	DAg-a011EBGE#Jj 1 D:a=D D<BD D D D	D
 
Y		(c//"+I,=,=g,F,F"G"G 	> 	>Ag-a011EBJuSz5#4 >F1I
> >6<> > > >	> 6ABBBII);I///JJrNNN
Hh   %
h
8$$$Jr   c                    ddl m} d}d}t          | |          r| j        }| j        }| j        } | j        j        t          j	        t          j
        fvrt          d          t          |t                    st          d          t          j        |           t          j        t          |                                                    k                                    st          d          t          j        | j        d         t&                    }|                                D ]\  }}||| |k    <   |r ||||          }|S )	a8   Recode categorial data to int factor.

    Parameters
    ----------
    x : array_like
        array like object supporting with numpy array methods of categorially
        coded data.
    levels : dict
        mapping of labels to integer-codings

    Returns
    -------
    out : instance numpy.ndarray
    r   )SeriesNz<This is not a categorial factor. Array of str type required.z4This is not a valid value for levels. Dict required.z)The levels do not match the array values.)dtype)r
   index)r$   rU   r+   r
   rW   r7   rV   typer,   str_object_r5   r0   r-   listkeysr;   emptyshapeintitems)r   levelsrU   r
   rW   outlevelcodings           r   r/   r/      sV    DE!V vHw|BGRZ000 8 9 9 	9 %%  + , , 	, illbiV[[]](;(;<<<AACC DEEE hqwqz---#\\^^ 	% 	%ME6$CU
OO 	6&4u555C
r   )
r   Nr   NNNNNr   N)__doc__statsmodels.compat.pythonr   numpyr,   statsmodels.graphics.plottoolsr   statsmodels.graphics.utilsgraphicsr%   rS   r/   r   r   r   <module>rk      s     - , , , , ,     2 2 2 2 2 2 * * * * * * * * * ILDHDHR R R Rj+ + + + +r   