<commit_before>"""Plotting functions for linear models (broadly construed)."""
from __future__ import division
import itertools
import numpy as np
import pandas as pd
from scipy import stats
from scipy.spatial import distance
import statsmodels.api as sm
import statsmodels.formula.api as sf
import matplotlib as mpl
import matplotlib.pyplot as plt
import moss

from six.moves import range

from .utils import color_palette, ci_to_errsize
from .distributions import distplot
from .axisgrid import FacetGrid


def lmplot(x, y, data, color=None, row=None, col=None, col_wrap=None,
           x_estimator=None, x_ci=95, x_bins=None, n_boot=5000, fit_reg=True,
           order=1, ci=95, logistic=False, truncate=False,
           x_partial=None, y_partial=None, x_jitter=None, y_jitter=None,
           sharex=True, sharey=True, palette="husl", size=None,
           scatter_kws=None, line_kws=None, palette_kws=None):
    """Plot a linear model with faceting, color binning, and other options.

    Parameters
    ----------
    x, y : strings
        Column names in `data` DataFrame for x and y variables.
    data : DataFrame
        Dource of data for the model.
    color : string, optional
        DataFrame column name to group the model by color.
    row, col : strings, optional
        DataFrame column names to make separate plot facets.
    col_wrap : int, optional
        Wrap col variable at this width - cannot be used with row facet.
    x_estimator : callable, optional
        Interpret X values as factor labels and use this function
        to plot the point estimate and bootstrapped CI.
    x_ci : int optional
        Size of confidence interval for x_estimator error bars.
    x_bins : sequence of floats, optional
        Bin the x variable with these values. Implies that x_estimator is
        mean, unless otherwise provided.
    n_boot : int, optional
        Number of bootstrap iterations to perform.
    fit_reg : bool, optional
        If True fit a regression model by color/row/col and plot.
    order : int, optional
        Order of the regression polynomial to fit.
    ci : int, optional
        Confidence interval for the regression line.
    logistic : bool, optional
        Fit the regression line with logistic regression.
    truncate : bool, optional
        If True, only fit line from data min to data max.
    {x, y}_partial : string or list of strings, optional
        Regress these variables out of the factors before plotting.
    {x, y}_jitter : float, optional
        Parameters for uniformly distributed random noise added to positions.
    sharex, sharey : bools, optional
        Only relevant if faceting; passed to plt.subplots.
    palette : seaborn color palette argument
        If using separate plots by color, draw with this color palette.
    size : float, optional
        Size (plots are square) for each plot facet.
    {scatter, line}_kws : dictionary
        Keyword arguments to pass to the underlying plot functions.
    palette_kws : dictionary
        Keyword arguments for seaborn.color_palette.

    """
    # TODO
    # - legend when fit_line is False

    # First sort out the general figure layout
    if size is None:
        size = mpl.rcParams["figure.figsize"][1]

    if col is None and col_wrap is not None:
        raise ValueError("Need column facet variable for `col_wrap`")
    if row is not None and col_wrap is not None:
        raise ValueError("Cannot facet rows when using `col_wrap`")

    nrow = 1 if row is None else len(data[row].unique())
    ncol = 1 if col is None else len(data[col].unique())

    if col_wrap is not None:
        ncol = col_wrap
        nrow = int(np.ceil(len(data[col].unique()) / col_wrap))

    f, axes = plt.subplots(nrow, ncol, sharex=sharex, sharey=sharey,
                           figsize=(size * ncol, size * nrow))
    axes = np.atleast_2d(axes).reshape(nrow, ncol)

    if nrow == 1 or col_wrap is not None:
        row_masks = [np.repeat(True, len(data))]
    else:
        row_vals = np.sort(data[row].unique())
        row_masks = [data[row] == val for val in row_vals]

    if ncol == 1:
        col_masks = [np.repeat(True, len(data))]
    else:
        col_vals = np.sort(data[col].unique())
        col_masks = [data[col] == val for val in col_vals]

    if x_bins is not None:
        x_estimator = np.mean if x_estimator is None else x_estimator
        x_bins = np.c_[x_bins]

    if x_partial is not None:
        if not isinstance(x_partial, list):
            x_partial = [x_partial]
    if y_partial is not None:
        if not isinstance(y_partial, list):
            y_partial = [y_partial]

    if palette_kws is None:
        palette_kws = {}

    # Sort out the plot colors
    color_factor = color
    if color is None:
        hue_masks = [np.repeat(True, len(data))]
        colors = ["#222222"]
    else:
        hue_vals = np.sort(data[color].unique())
        hue_masks = [data[color] == val for val in hue_vals]
        colors = color_palette(palette, len(hue_masks), **palette_kws)

    # Default keyword arguments for plot components
    if scatter_kws is None:
        scatter_kws = {}
    if line_kws is None:
        line_kws = {}

    # First walk through the facets and plot the scatters
    scatter_ms = scatter_kws.pop("ms", 4)
    scatter_mew = mew = scatter_kws.pop("mew", 0)
    scatter_alpha = mew = scatter_kws.pop("alpha", .77)
    for row_i, row_mask in enumerate(row_masks):
        for col_j, col_mask in enumerate(col_masks):
            if col_wrap is not None:
                f_row = col_j // ncol
                f_col = col_j % ncol
            else:
                f_row, f_col = row_i, col_j
            ax = axes[f_row, f_col]
            if f_row + 1 == nrow:
                ax.set_xlabel(x)
            if f_col == 0:
                ax.set_ylabel(y)

            # Title the plot if we are faceting
            title = ""
            if row is not None:
                title += "%s = %s" % (row, row_vals[row_i])
            if row is not None and col is not None:
                title += " | "
            if col is not None:
                title += "%s = %s" % (col, col_vals[col_j])
            if size < 3:
                title = title.replace(" | ", "\n")
            ax.set_title(title)

            for hue_k, hue_mask in enumerate(hue_masks):
                color = colors[hue_k]
                data_ijk = data[row_mask & col_mask & hue_mask]

                if x_estimator is not None:
                    ms = scatter_kws.pop("ms", 7)
                    mew = scatter_kws.pop("mew", 0)
                    if x_bins is None:
                        x_vals = data_ijk[x].unique()
                        x_data = data_ijk[x]
                    else:
                        dist = distance.cdist(np.c_[data_ijk[x]], x_bins)
                        x_vals = x_bins.ravel()
                        x_data = x_bins[np.argmin(dist, axis=1)].ravel()

                    y_vals = data_ijk[y]

                    if y_partial is not None:
                        for var in y_partial:
                            conf = data_ijk[var]
                            conf -= conf.mean()
                            y_mean = y_vals.mean()
                            y_vals = moss.vector_reject(y_vals - y_mean, conf)
                            y_vals += y_mean

                    y_grouped = [np.array(y_vals[x_data == v])
                                 for v in x_vals]

                    y_est = [x_estimator(y_i) for y_i in y_grouped]
                    y_boots = [moss.bootstrap(np.array(y_i),
                                              func=x_estimator,
                                              n_boot=n_boot)
                               for y_i in y_grouped]
                    ci_lims = [50 - x_ci / 2., 50 + x_ci / 2.]
                    y_ci = [moss.percentiles(y_i, ci_lims) for y_i in y_boots]
                    y_error = ci_to_errsize(np.transpose(y_ci), y_est)

                    ax.plot(x_vals, y_est, "o", mew=mew, ms=ms,
                            color=color, **scatter_kws)
                    ax.errorbar(x_vals, y_est, y_error,
                                fmt=None, ecolor=color)
                else:
                    x_ = data_ijk[x]
                    y_ = data_ijk[y]

                    if x_partial is not None:
                        for var in x_partial:
                            conf = data_ijk[var]
                            conf -= conf.mean()
                            x_mean = x_.mean()
                            x_ = moss.vector_reject(x_ - x_mean, conf)
                            x_ += x_mean
                    if y_partial is not None:
                        for var in y_partial:
                            conf = data_ijk[var]
                            conf -= conf.mean()
                            y_mean = y_.mean()
                            y_ = moss.vector_reject(y_ - y_mean, conf)
                            y_ += y_mean

                    if x_jitter is not None:
                        x_ += np.random.uniform(-x_jitter, x_jitter, x_.shape)
                    if y_jitter is not None:
                        y_ += np.random.uniform(-y_jitter, y_jitter, y_.shape)
                    ax.plot(x_, y_, "o", color=color, alpha=scatter_alpha,
                            mew=scatter_mew, ms=scatter_ms, **scatter_kws)

    for ax_i in np.ravel(axes):
        ax_i.set_xmargin(.05)
        ax_i.autoscale_view()

    # Now walk through again and plot the regression estimate
    # and a confidence interval for the regression line
    if fit_reg:
        for row_i, row_mask in enumerate(row_masks):
            for col_j, col_mask in enumerate(col_masks):
                if col_wrap is not None:
                    f_row = col_j // ncol
                    f_col = col_j % ncol
                else:
                    f_row, f_col = row_i, col_j
                ax = axes[f_row, f_col]
                xlim = ax.get_xlim()

                for hue_k, hue_mask in enumerate(hue_masks):
                    color = colors[hue_k]
                    data_ijk = data[row_mask & col_mask & hue_mask]
                    x_vals = np.array(data_ijk[x])
                    y_vals = np.array(data_ijk[y])
                    if not len(x_vals):
                        continue

                    # Sort out the limit of the fit
                    if truncate:
                        xx = np.linspace(x_vals.min(),
                                         x_vals.max(), 100)
                    else:
                        xx = np.linspace(xlim[0], xlim[1], 100)
                    xx_ = sm.add_constant(xx, prepend=True)

                    # Inner function to bootstrap the regression
                    def _regress(x, y):
                        if logistic:
                            x_ = sm.add_constant(x, prepend=True)
                            fit = sm.GLM(y, x_,
                                         family=sm.families.Binomial()).fit()
                            reg = fit.predict(xx_)
                        else:
                            fit = np.polyfit(x, y, order)
                            reg = np.polyval(fit, xx)
                        return reg

                    # Remove nuisance variables with vector rejection
                    if x_partial is not None:
                        for var in x_partial:
                            conf = data_ijk[var]
                            conf -= conf.mean()
                            x_mean = x_vals.mean()
                            x_vals = moss.vector_reject(x_vals - x_mean, conf)
                            x_vals += x_mean
                    if y_partial is not None:
                        for var in y_partial:
                            conf = data_ijk[var]
                            conf -= conf.mean()
                            y_mean = y_vals.mean()
                            y_vals = moss.vector_reject(y_vals - y_mean, conf)
                            y_vals += y_mean

                    # Regression line confidence interval
                    if ci is not None:
                        ci_lims = [50 - ci / 2., 50 + ci / 2.]
                        boots = moss.bootstrap(x_vals, y_vals,
                                               func=_regress,
                                               n_boot=n_boot)
                        ci_band = moss.percentiles(boots, ci_lims, axis=0)
                        ax.fill_between(xx, *ci_band, color=color, alpha=.15)

                    # Regression line
                    reg = _regress(x_vals, y_vals)
                    if color_factor is None:
                        label = ""
                    else:
                        label = hue_vals[hue_k]
                    ax.plot(xx, reg, color=color,
                            label=str(label), **line_kws)
                    ax.set_xlim(xlim)

    # Plot the legend on the upper left facet and adjust the layout
    if color_factor is not None and color_factor not in [row, col]:
        axes[0, 0].legend(loc="best", title=color_factor)
    plt.tight_layout()


def factorplot(x, y=None, data=None, color=None, row=None, col=None,
               col_wrap=None, kind="auto", estimator=np.mean, ci=95, size=5,
               aspect=1, palette=None, dodge=0, link=True, legend=True,
               legend_out=True):

    if color is not None and palette is None:
        palette = "husl"
    elif color is None and palette is not None:
        color = x

    facet = FacetGrid(data, row, col, color, col_wrap=col_wrap, size=size,
                      aspect=aspect, legend=legend, legend_out=legend_out,
                      palette=palette)

    if kind == "auto":
        if y is None:
            kind = "bar"
        elif (data[y].mean() / data[y].std()) < 2.5:
            kind = "bar"
        else:
            kind = "point"

    mask_gen = facet._iter_masks()

    x_order = sorted(data[x].unique())

    if y is None:
        estimator = len
        ci = None
        y = x
        y_count = True
    else:
        y_count = False

    if color is None:
        colors = ["#777777" if kind == "bar" else "#333333"]
        n_colors = 1
        width = .8
        pos_adjust = [0]
    else:
        color_order = sorted(data[color].unique())
        n_colors = len(color_order)
        colors = color_palette(palette, n_colors)
        if color == row or color == col or color == x:
            width = .8
            pos_adjust = [0] * n_colors
        else:
            width = .8 / n_colors
            if kind == "bar":
                pos_adjust = np.linspace(0, .8 - width, n_colors)
                pos_adjust -= pos_adjust.mean()
            elif kind == "point" and dodge:
                pos_adjust = np.linspace(0, dodge, n_colors)
                pos_adjust -= pos_adjust.mean()
            else:
                pos_adjust = [0] * n_colors

    draw_legend = color is not None and color not in [x, row, col]

    for (row_i, col_j, hue_k), data_ijk in mask_gen:

        ax = facet._axes[row_i, col_j]

        grouped = {g: data[y] for g, data in data_ijk.groupby(x)}

        plot_pos = []
        plot_heights = []
        plot_cis = []
        for pos, var in enumerate(x_order):

            if var not in grouped:
                continue

            plot_pos.append(pos + pos_adjust[hue_k])
            plot_heights.append(estimator(grouped[var]))
            if ci is not None:
                boots = moss.bootstrap(grouped[var].values, func=estimator)
                plot_cis.append(moss.ci(boots, ci))

        kwargs = {}
        if draw_legend:
            kwargs["label"] = facet._hue_vals[hue_k]

        if kind == "bar":

            ax.bar(plot_pos, plot_heights, width, align="center",
                   color=colors[hue_k], **kwargs)
            facet._update_legend_data(ax)
            if ci is not None:
                for pos, ci_ in zip(plot_pos, plot_cis):
                    ax.plot([pos, pos], ci_, linewidth=2.5, color="#222222")

        elif kind == "point":

            hue = colors[hue_k]
            ls = "-" if link else ""
            ax.plot(plot_pos, plot_heights, color=hue, marker="o", ms=9,
                    ls=ls, lw=3, **kwargs)
            facet._update_legend_data(ax)
            if ci is not None:
                for pos, ci_ in zip(plot_pos, plot_cis):
                    ax.plot([pos, pos], ci_, linewidth=2.5, color=hue)

    n_x = len(x_order)
    facet.set(xticks=range(n_x), xticklabels=x_order, xlim=(-.5, n_x - .5))

    for ax in facet._axes.flat:
        ax.xaxis.grid(False)

    y_label = "count" if y_count else y
    facet._set_axis_labels(x, y_label)
    facet._set_title()
    if draw_legend:
        facet._make_legend()

    return facet


def _regress_fast(grid, x, y, ci, n_boot):

    X = np.c_[np.ones(len(x)), x]
    grid = np.c_[np.ones(len(grid)), grid]
    reg_func = lambda _x, _y: np.linalg.pinv(_x).dot(_y)
    y_hat = grid.dot(reg_func(X, y))
    beta_boots = moss.bootstrap(X, y, func=reg_func, n_boot=n_boot).T
    y_hat_boots = grid.dot(beta_boots).T
    err_bands = moss.ci(y_hat_boots, ci, axis=0)
    return y_hat, err_bands


def _regress_poly(grid, x, y, order, ci, n_boot):

    reg_func = lambda _x, _y: np.polyval(np.polyfit(_x, _y, order), grid)
    y_hat = reg_func(x, y)
    y_hat_boots = moss.bootstrap(x, y, func=reg_func, n_boot=n_boot)
    err_bands = moss.ci(y_hat_boots, ci, axis=0)
    return y_hat, err_bands


def _regress_statsmodels(grid, x, y, model, kwargs, ci, n_boot):

    X = np.c_[np.ones(len(x)), x]
    grid = np.c_[np.ones(len(grid)), grid]
    reg_func = lambda _x, _y: model(_y, _x, **kwargs).fit().predict(grid)
    y_hat = reg_func(X, y)
    y_hat_boots = moss.bootstrap(X, y, func=reg_func, n_boot=n_boot)
    err_bands = moss.ci(y_hat_boots, ci, axis=0)
    return y_hat, err_bands


def _bootstrap_reg(func, x, y, grid, n_boot, ci):

    boots = moss.bootstrap(x, y, func=func, n_boot=n_boot)
    y_hats = grid.dot(boots.reshape(-1, n_boot))
    ci = moss.ci(y_hats, ci, axis=1).squeeze()
    return ci


def regplot(x, y, x_estimator=None, x_ci=95, x_bins=None, n_boot=5000,
            fit_reg=True, order=1, ci=95, logistic=False, truncate=False,
            x_partial=None, y_partial=None, x_jitter=None, y_jitter=None,
            color=None, scatter_kws=None, line_kws=None, ax=None):

    if ax is None:
        ax = plt.gca()

    if color is None:
        color = mpl.rcParams["axes.color_cycle"][0]

    x = np.asarray(x)
    y = np.asarray(y)

    ax.scatter(x, y, color=color, alpha=.8)

    x_min, x_max = ax.get_xlim()
    grid = np.linspace(x_min, x_max, 100)

    if order > 1:
        y_hat, err_bands = _regress_poly(grid, x, y, order, ci, n_boot)
    elif logistic:
        kwargs = dict(family=sm.families.Binomial())
        y_hat, err_bands = _regress_statsmodels(grid, x, y, sm.GLM, kwargs,
                                                ci, n_boot)
    #elif robust:
    #    kwargs = dict()
    #    y_hat, err_bands = _regress_statsmodels(grid, x, y, sm.RLM, kwargs,
    #                                            ci, n_boot)
    else:
        y_hat, err_bands = _regress_fast(grid, x, y, ci, n_boot)

    ax.plot(grid, y_hat, color=color, lw=2)
    ax.fill_between(grid, *err_bands, color=color, alpha=.15)

    ax.set_xlim(x_min, x_max)


def regplot_off(x, y, data=None, corr_func=stats.pearsonr, func_name=None,
                xlabel="", ylabel="", ci=95, size=None, annotloc=None, color=None,
                reg_kws=None, scatter_kws=None, dist_kws=None, text_kws=None):
    """Scatterplot with regresion line, marginals, and correlation value.

    Parameters
    ----------
    x : sequence or string
        Independent variable.
    y : sequence or string
        Dependent variable.
    data : dataframe, optional
        If dataframe is given, x, and y are interpreted as string keys
        for selecting to dataframe column names.
    corr_func : callable, optional
        Correlation function; expected to take two arrays and return a
        numeric or (statistic, pval) tuple.<commit_msg>Adding robust and partial to new regplot<commit_after>"""Plotting functions for linear models (broadly construed)."""
from __future__ import division
import itertools
import numpy as np
import pandas as pd
from scipy import stats
from scipy.spatial import distance
import statsmodels.api as sm
import statsmodels.formula.api as sf
import matplotlib as mpl
import matplotlib.pyplot as plt
import moss

from six.moves import range

from .utils import color_palette, ci_to_errsize
from .distributions import distplot
from .axisgrid import FacetGrid


def lmplot(x, y, data, color=None, row=None, col=None, col_wrap=None,
           x_estimator=None, x_ci=95, x_bins=None, n_boot=5000, fit_reg=True,
           order=1, ci=95, logistic=False, truncate=False,
           x_partial=None, y_partial=None, x_jitter=None, y_jitter=None,
           sharex=True, sharey=True, palette="husl", size=None,
           scatter_kws=None, line_kws=None, palette_kws=None):
    """Plot a linear model with faceting, color binning, and other options.

    Parameters
    ----------
    x, y : strings
        Column names in `data` DataFrame for x and y variables.
    data : DataFrame
        Dource of data for the model.
    color : string, optional
        DataFrame column name to group the model by color.
    row, col : strings, optional
        DataFrame column names to make separate plot facets.
    col_wrap : int, optional
        Wrap col variable at this width - cannot be used with row facet.
    x_estimator : callable, optional
        Interpret X values as factor labels and use this function
        to plot the point estimate and bootstrapped CI.
    x_ci : int optional
        Size of confidence interval for x_estimator error bars.
    x_bins : sequence of floats, optional
        Bin the x variable with these values. Implies that x_estimator is
        mean, unless otherwise provided.
    n_boot : int, optional
        Number of bootstrap iterations to perform.
    fit_reg : bool, optional
        If True fit a regression model by color/row/col and plot.
    order : int, optional
        Order of the regression polynomial to fit.
    ci : int, optional
        Confidence interval for the regression line.
    logistic : bool, optional
        Fit the regression line with logistic regression.
    truncate : bool, optional
        If True, only fit line from data min to data max.
    {x, y}_partial : string or list of strings, optional
        Regress these variables out of the factors before plotting.
    {x, y}_jitter : float, optional
        Parameters for uniformly distributed random noise added to positions.
    sharex, sharey : bools, optional
        Only relevant if faceting; passed to plt.subplots.
    palette : seaborn color palette argument
        If using separate plots by color, draw with this color palette.
    size : float, optional
        Size (plots are square) for each plot facet.
    {scatter, line}_kws : dictionary
        Keyword arguments to pass to the underlying plot functions.
    palette_kws : dictionary
        Keyword arguments for seaborn.color_palette.

    """
    # TODO
    # - legend when fit_line is False

    # First sort out the general figure layout
    if size is None:
        size = mpl.rcParams["figure.figsize"][1]

    if col is None and col_wrap is not None:
        raise ValueError("Need column facet variable for `col_wrap`")
    if row is not None and col_wrap is not None:
        raise ValueError("Cannot facet rows when using `col_wrap`")

    nrow = 1 if row is None else len(data[row].unique())
    ncol = 1 if col is None else len(data[col].unique())

    if col_wrap is not None:
        ncol = col_wrap
        nrow = int(np.ceil(len(data[col].unique()) / col_wrap))

    f, axes = plt.subplots(nrow, ncol, sharex=sharex, sharey=sharey,
                           figsize=(size * ncol, size * nrow))
    axes = np.atleast_2d(axes).reshape(nrow, ncol)

    if nrow == 1 or col_wrap is not None:
        row_masks = [np.repeat(True, len(data))]
    else:
        row_vals = np.sort(data[row].unique())
        row_masks = [data[row] == val for val in row_vals]

    if ncol == 1:
        col_masks = [np.repeat(True, len(data))]
    else:
        col_vals = np.sort(data[col].unique())
        col_masks = [data[col] == val for val in col_vals]

    if x_bins is not None:
        x_estimator = np.mean if x_estimator is None else x_estimator
        x_bins = np.c_[x_bins]

    if x_partial is not None:
        if not isinstance(x_partial, list):
            x_partial = [x_partial]
    if y_partial is not None:
        if not isinstance(y_partial, list):
            y_partial = [y_partial]

    if palette_kws is None:
        palette_kws = {}

    # Sort out the plot colors
    color_factor = color
    if color is None:
        hue_masks = [np.repeat(True, len(data))]
        colors = ["#222222"]
    else:
        hue_vals = np.sort(data[color].unique())
        hue_masks = [data[color] == val for val in hue_vals]
        colors = color_palette(palette, len(hue_masks), **palette_kws)

    # Default keyword arguments for plot components
    if scatter_kws is None:
        scatter_kws = {}
    if line_kws is None:
        line_kws = {}

    # First walk through the facets and plot the scatters
    scatter_ms = scatter_kws.pop("ms", 4)
    scatter_mew = mew = scatter_kws.pop("mew", 0)
    scatter_alpha = mew = scatter_kws.pop("alpha", .77)
    for row_i, row_mask in enumerate(row_masks):
        for col_j, col_mask in enumerate(col_masks):
            if col_wrap is not None:
                f_row = col_j // ncol
                f_col = col_j % ncol
            else:
                f_row, f_col = row_i, col_j
            ax = axes[f_row, f_col]
            if f_row + 1 == nrow:
                ax.set_xlabel(x)
            if f_col == 0:
                ax.set_ylabel(y)

            # Title the plot if we are faceting
            title = ""
            if row is not None:
                title += "%s = %s" % (row, row_vals[row_i])
            if row is not None and col is not None:
                title += " | "
            if col is not None:
                title += "%s = %s" % (col, col_vals[col_j])
            if size < 3:
                title = title.replace(" | ", "\n")
            ax.set_title(title)

            for hue_k, hue_mask in enumerate(hue_masks):
                color = colors[hue_k]
                data_ijk = data[row_mask & col_mask & hue_mask]

                if x_estimator is not None:
                    ms = scatter_kws.pop("ms", 7)
                    mew = scatter_kws.pop("mew", 0)
                    if x_bins is None:
                        x_vals = data_ijk[x].unique()
                        x_data = data_ijk[x]
                    else:
                        dist = distance.cdist(np.c_[data_ijk[x]], x_bins)
                        x_vals = x_bins.ravel()
                        x_data = x_bins[np.argmin(dist, axis=1)].ravel()

                    y_vals = data_ijk[y]

                    if y_partial is not None:
                        for var in y_partial:
                            conf = data_ijk[var]
                            conf -= conf.mean()
                            y_mean = y_vals.mean()
                            y_vals = moss.vector_reject(y_vals - y_mean, conf)
                            y_vals += y_mean

                    y_grouped = [np.array(y_vals[x_data == v])
                                 for v in x_vals]

                    y_est = [x_estimator(y_i) for y_i in y_grouped]
                    y_boots = [moss.bootstrap(np.array(y_i),
                                              func=x_estimator,
                                              n_boot=n_boot)
                               for y_i in y_grouped]
                    ci_lims = [50 - x_ci / 2., 50 + x_ci / 2.]
                    y_ci = [moss.percentiles(y_i, ci_lims) for y_i in y_boots]
                    y_error = ci_to_errsize(np.transpose(y_ci), y_est)

                    ax.plot(x_vals, y_est, "o", mew=mew, ms=ms,
                            color=color, **scatter_kws)
                    ax.errorbar(x_vals, y_est, y_error,
                                fmt=None, ecolor=color)
                else:
                    x_ = data_ijk[x]
                    y_ = data_ijk[y]

                    if x_partial is not None:
                        for var in x_partial:
                            conf = data_ijk[var]
                            conf -= conf.mean()
                            x_mean = x_.mean()
                            x_ = moss.vector_reject(x_ - x_mean, conf)
                            x_ += x_mean
                    if y_partial is not None:
                        for var in y_partial:
                            conf = data_ijk[var]
                            conf -= conf.mean()
                            y_mean = y_.mean()
                            y_ = moss.vector_reject(y_ - y_mean, conf)
                            y_ += y_mean

                    if x_jitter is not None:
                        x_ += np.random.uniform(-x_jitter, x_jitter, x_.shape)
                    if y_jitter is not None:
                        y_ += np.random.uniform(-y_jitter, y_jitter, y_.shape)
                    ax.plot(x_, y_, "o", color=color, alpha=scatter_alpha,
                            mew=scatter_mew, ms=scatter_ms, **scatter_kws)

    for ax_i in np.ravel(axes):
        ax_i.set_xmargin(.05)
        ax_i.autoscale_view()

    # Now walk through again and plot the regression estimate
    # and a confidence interval for the regression line
    if fit_reg:
        for row_i, row_mask in enumerate(row_masks):
            for col_j, col_mask in enumerate(col_masks):
                if col_wrap is not None:
                    f_row = col_j // ncol
                    f_col = col_j % ncol
                else:
                    f_row, f_col = row_i, col_j
                ax = axes[f_row, f_col]
                xlim = ax.get_xlim()

                for hue_k, hue_mask in enumerate(hue_masks):
                    color = colors[hue_k]
                    data_ijk = data[row_mask & col_mask & hue_mask]
                    x_vals = np.array(data_ijk[x])
                    y_vals = np.array(data_ijk[y])
                    if not len(x_vals):
                        continue

                    # Sort out the limit of the fit
                    if truncate:
                        xx = np.linspace(x_vals.min(),
                                         x_vals.max(), 100)
                    else:
                        xx = np.linspace(xlim[0], xlim[1], 100)
                    xx_ = sm.add_constant(xx, prepend=True)

                    # Inner function to bootstrap the regression
                    def _regress(x, y):
                        if logistic:
                            x_ = sm.add_constant(x, prepend=True)
                            fit = sm.GLM(y, x_,
                                         family=sm.families.Binomial()).fit()
                            reg = fit.predict(xx_)
                        else:
                            fit = np.polyfit(x, y, order)
                            reg = np.polyval(fit, xx)
                        return reg

                    # Remove nuisance variables with vector rejection
                    if x_partial is not None:
                        for var in x_partial:
                            conf = data_ijk[var]
                            conf -= conf.mean()
                            x_mean = x_vals.mean()
                            x_vals = moss.vector_reject(x_vals - x_mean, conf)
                            x_vals += x_mean
                    if y_partial is not None:
                        for var in y_partial:
                            conf = data_ijk[var]
                            conf -= conf.mean()
                            y_mean = y_vals.mean()
                            y_vals = moss.vector_reject(y_vals - y_mean, conf)
                            y_vals += y_mean

                    # Regression line confidence interval
                    if ci is not None:
                        ci_lims = [50 - ci / 2., 50 + ci / 2.]
                        boots = moss.bootstrap(x_vals, y_vals,
                                               func=_regress,
                                               n_boot=n_boot)
                        ci_band = moss.percentiles(boots, ci_lims, axis=0)
                        ax.fill_between(xx, *ci_band, color=color, alpha=.15)

                    # Regression line
                    reg = _regress(x_vals, y_vals)
                    if color_factor is None:
                        label = ""
                    else:
                        label = hue_vals[hue_k]
                    ax.plot(xx, reg, color=color,
                            label=str(label), **line_kws)
                    ax.set_xlim(xlim)

    # Plot the legend on the upper left facet and adjust the layout
    if color_factor is not None and color_factor not in [row, col]:
        axes[0, 0].legend(loc="best", title=color_factor)
    plt.tight_layout()


def factorplot(x, y=None, data=None, color=None, row=None, col=None,
               col_wrap=None, kind="auto", estimator=np.mean, ci=95, size=5,
               aspect=1, palette=None, dodge=0, link=True, legend=True,
               legend_out=True):

    if color is not None and palette is None:
        palette = "husl"
    elif color is None and palette is not None:
        color = x

    facet = FacetGrid(data, row, col, color, col_wrap=col_wrap, size=size,
                      aspect=aspect, legend=legend, legend_out=legend_out,
                      palette=palette)

    if kind == "auto":
        if y is None:
            kind = "bar"
        elif (data[y].mean() / data[y].std()) < 2.5:
            kind = "bar"
        else:
            kind = "point"

    mask_gen = facet._iter_masks()

    x_order = sorted(data[x].unique())

    if y is None:
        estimator = len
        ci = None
        y = x
        y_count = True
    else:
        y_count = False

    if color is None:
        colors = ["#777777" if kind == "bar" else "#333333"]
        n_colors = 1
        width = .8
        pos_adjust = [0]
    else:
        color_order = sorted(data[color].unique())
        n_colors = len(color_order)
        colors = color_palette(palette, n_colors)
        if color == row or color == col or color == x:
            width = .8
            pos_adjust = [0] * n_colors
        else:
            width = .8 / n_colors
            if kind == "bar":
                pos_adjust = np.linspace(0, .8 - width, n_colors)
                pos_adjust -= pos_adjust.mean()
            elif kind == "point" and dodge:
                pos_adjust = np.linspace(0, dodge, n_colors)
                pos_adjust -= pos_adjust.mean()
            else:
                pos_adjust = [0] * n_colors

    draw_legend = color is not None and color not in [x, row, col]

    for (row_i, col_j, hue_k), data_ijk in mask_gen:

        ax = facet._axes[row_i, col_j]

        grouped = {g: data[y] for g, data in data_ijk.groupby(x)}

        plot_pos = []
        plot_heights = []
        plot_cis = []
        for pos, var in enumerate(x_order):

            if var not in grouped:
                continue

            plot_pos.append(pos + pos_adjust[hue_k])
            plot_heights.append(estimator(grouped[var]))
            if ci is not None:
                boots = moss.bootstrap(grouped[var].values, func=estimator)
                plot_cis.append(moss.ci(boots, ci))

        kwargs = {}
        if draw_legend:
            kwargs["label"] = facet._hue_vals[hue_k]

        if kind == "bar":

            ax.bar(plot_pos, plot_heights, width, align="center",
                   color=colors[hue_k], **kwargs)
            facet._update_legend_data(ax)
            if ci is not None:
                for pos, ci_ in zip(plot_pos, plot_cis):
                    ax.plot([pos, pos], ci_, linewidth=2.5, color="#222222")

        elif kind == "point":

            hue = colors[hue_k]
            ls = "-" if link else ""
            ax.plot(plot_pos, plot_heights, color=hue, marker="o", ms=9,
                    ls=ls, lw=3, **kwargs)
            facet._update_legend_data(ax)
            if ci is not None:
                for pos, ci_ in zip(plot_pos, plot_cis):
                    ax.plot([pos, pos], ci_, linewidth=2.5, color=hue)

    n_x = len(x_order)
    facet.set(xticks=range(n_x), xticklabels=x_order, xlim=(-.5, n_x - .5))

    for ax in facet._axes.flat:
        ax.xaxis.grid(False)

    y_label = "count" if y_count else y
    facet._set_axis_labels(x, y_label)
    facet._set_title()
    if draw_legend:
        facet._make_legend()

    return facet


def _regress_fast(grid, x, y, ci, n_boot):

    X = np.c_[np.ones(len(x)), x]
    grid = np.c_[np.ones(len(grid)), grid]
    reg_func = lambda _x, _y: np.linalg.pinv(_x).dot(_y)
    y_hat = grid.dot(reg_func(X, y))
    if ci is None:
        return y_hat, None
    beta_boots = moss.bootstrap(X, y, func=reg_func, n_boot=n_boot).T
    y_hat_boots = grid.dot(beta_boots).T
    return y_hat, y_hat_boots


def _regress_poly(grid, x, y, order, ci, n_boot):

    reg_func = lambda _x, _y: np.polyval(np.polyfit(_x, _y, order), grid)
    y_hat = reg_func(x, y)
    if ci is None:
        return y_hat, None
    y_hat_boots = moss.bootstrap(x, y, func=reg_func, n_boot=n_boot)
    return y_hat, y_hat_boots


def _regress_statsmodels(grid, x, y, model, kwargs, ci, n_boot):

    X = np.c_[np.ones(len(x)), x]
    grid = np.c_[np.ones(len(grid)), grid]
    reg_func = lambda _x, _y: model(_y, _x, **kwargs).fit().predict(grid)
    y_hat = reg_func(X, y)
    if ci is None:
        return y_hat, None
    y_hat_boots = moss.bootstrap(X, y, func=reg_func, n_boot=n_boot)
    return y_hat, y_hat_boots


def regplot(x, y, x_estimator=None, x_ci=95, x_bins=None,
            fit_reg=True, ci=95, n_boot=5000,
            order=1, logistic=False, robust=False, partial=None,
            truncate=False, x_jitter=None, y_jitter=None,
            color=None, scatter_kws=None, line_kws=None,
            ax=None):

    if ax is None:
        ax = plt.gca()

    x = np.asarray(x)
    y = np.asarray(y)

    if color is None:
        lines, = plt.plot(x.mean(), y.mean())
        color = lines.get_color()
        lines.remove()

    if partial is None:
        y_scatter = y
    else:
        X = np.c_[np.ones_like(x), x, partial]
        y_scatter = y - X[:, 2].dot(np.linalg.pinv(X).dot(y)[2])

    ax.scatter(x, y_scatter, color=color, alpha=.8)

    x_min, x_max = ax.get_xlim()
    grid = np.linspace(x_min, x_max, 100)

    if sum((order > 1, logistic, robust, partial is not None)) > 1:
        raise ValueError("`order` > 1 and `logistic` are mutally exclusive")

    if order > 1:
        y_hat, y_hat_boots = _regress_poly(grid, x, y, order, ci, n_boot)

    elif logistic:
        kwargs = dict(family=sm.families.Binomial())
        y_hat, y_hat_boots = _regress_statsmodels(grid, x, y, sm.GLM, kwargs,
                                                  ci, n_boot)
    elif robust:
        kwargs = dict()
        y_hat, y_hat_boots = _regress_statsmodels(grid, x, y, sm.RLM, kwargs,
                                                  ci, n_boot)
    elif partial is not None:
        _x = np.c_[x, partial]
        _grid = np.c_[grid, np.zeros((grid.size, partial.shape[1]))]
        kwargs = dict()
        y_hat, y_hat_boots = _regress_statsmodels(_grid, _x, y, sm.OLS, kwargs,
                                                  ci, n_boot)
    else:
        y_hat, y_hat_boots = _regress_fast(grid, x, y, ci, n_boot)

    if ci is not None:
        err_bands = moss.ci(y_hat_boots, ci, axis=0)

    ax.plot(grid, y_hat, color=color, lw=2)
    ax.fill_between(grid, *err_bands, color=color, alpha=.15)

    ax.set_xlim(x_min, x_max)


def regplot_off(x, y, data=None, corr_func=stats.pearsonr, func_name=None,
                xlabel="", ylabel="", ci=95, size=None, annotloc=None, color=None,
                reg_kws=None, scatter_kws=None, dist_kws=None, text_kws=None):
    """Scatterplot with regresion line, marginals, and correlation value.

    Parameters
    ----------
    x : sequence or string
        Independent variable.
    y : sequence or string
        Dependent variable.
    data : dataframe, optional
        If dataframe is given, x, and y are interpreted as string keys
        for selecting to dataframe column names.
    corr_func : callable, optional
        Correlation function; expected to take two arrays and return a
        numeric or (statistic, pval) tuple.
{'fraction_non_alphanumeric': 0.09414603236498342, 'fraction_numerical': 0.00616591928251121, 'mean_word_length': 1.5930864509605662, 'pattern_counts': {'":': 8, '<': 7, '<?xml version=': 0, '>': 8, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 40, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3479649', 'n_tokens_mistral': 13187, 'n_tokens_neox': 12444, 'n_words': 3512}