<commit_before>"""Plotting functions for linear models (broadly construed)."""
from __future__ import division
import copy
import itertools
import warnings
import numpy as np
import pandas as pd
from scipy.spatial import distance
import matplotlib as mpl
import matplotlib.pyplot as plt

try:
    import statsmodels.api as sm
    import statsmodels.formula.api as sf
    _has_statsmodels = True
except ImportError:
    _has_statsmodels = False

from .external.six import string_types
from .external.six.moves import range

from . import utils
from . import algorithms as algo
from .palettes import color_palette
from .axisgrid import FacetGrid


class _LinearPlotter(object):
    """Base class for plotting relational data in tidy format.

    To get anything useful done you'll have to inherit from this, but setup
    code that can be abstracted out should be put here.

    """
    def establish_variables(self, data, **kws):
        """Extract variables from data or use directly."""
        self.data = data

        # Validate the inputs
        any_strings = any([isinstance(v, string_types) for v in kws.values()])
        if any_strings and data is None:
            raise ValueError("Must pass `data` if using named variables.")

        # Set the variables
        for var, val in kws.items():
            if isinstance(val, string_types):
                setattr(self, var, data[val])
            else:
                setattr(self, var, val)

    def dropna(self, *vars):
        """Remove observations with missing data."""
        vals = [getattr(self, var) for var in vars]
        vals = [v for v in vals if v is not None]
        not_na = np.all(np.column_stack([pd.notnull(v) for v in vals]), axis=1)
        for var in vars:
            val = getattr(self, var)
            if val is not None:
                setattr(self, var, val[not_na])

    def plot(self, ax):
        raise NotImplementedError


class _DiscretePlotter(_LinearPlotter):
    """Plotter for data with discrete independent variable(s).

    This will be used by the `barplot` and `pointplot` functions, and
    thus indirectly by the `factorplot` function. It can produce plots
    where some statistic for a dependent measure is estimated within
    subsets of the data, which can be hierarchically structured at up to two
    levels (`x` and `hue`). The plots can be drawn with a few different
    visual representations of the same underlying data (`bar`, and `point`,
    with `box` doing something similar but skipping the estimation).

    """
    def __init__(self, x, y=None, hue=None, data=None, units=None,
                 x_order=None, hue_order=None, color=None, palette=None,
                 kind="auto", markers=None, linestyles=None, dodge=0,
                 join=True, hline=None, estimator=np.mean, ci=95,
                 n_boot=1000, dropna=True):

        # This implies we have a single bar/point for each level of `x`
        # but that the different levels should be mapped with a palette
        self.x_palette = hue is None and palette is not None

        # Set class attributes based on inputs
        self.estimator = len if y is None else estimator
        self.ci = None if y is None else ci
        self.join = join
        self.n_boot = n_boot
        self.hline = hline

        # Other attributes that are hardcoded for now
        self.bar_widths = .8
        self.err_color = "#444444"
        self.lw = mpl.rcParams["lines.linewidth"] * 1.8

        # Once we've set the above values, if `y` is None we want the actual
        # y values to be the x values so we can count them
        self.y_count = y is None
        if y is None:
            y = x

        # Ascertain which values will be associated with what values
        self.establish_variables(data, x=x, y=y, hue=hue, units=units)

        # Figure out the order of the variables on the x axis
        x_sorted = np.sort(pd.unique(self.x))
        self.x_order = x_sorted if x_order is None else x_order
        if self.hue is not None:
            hue_sorted = np.sort(pd.unique(self.hue))
            self.hue_order = hue_sorted if hue_order is None else hue_order
        else:
            self.hue_order = [None]

        # Handle the other hue-mapped attributes
        if markers is None:
            self.markers = ["o"] * len(self.hue_order)
        else:
            if len(markers) != len(self.hue_order):
                raise ValueError("Length of marker list must equal "
                                 "number of hue levels")
            self.markers = markers
        if linestyles is None:
            self.linestyles = ["-"] * len(self.hue_order)
        else:
            if len(linestyles) != len(self.hue_order):
                raise ValueError("Length of linestyle list must equal "
                                 "number of hue levels")
            self.linestyles = linestyles

        # Drop null observations
        if dropna:
            self.dropna("x", "y", "hue", "units")

        # Settle whe kind of plot this is going to be
        self.establish_plot_kind(kind)

        # Determine the color palette
        self.establish_palette(color, palette)

        # Figure out where the data should be drawn
        self.establish_positions(dodge)

    def establish_palette(self, color, palette):
        """Set a list of colors for each plot element."""
        n_hues = len(self.x_order) if self.x_palette else len(self.hue_order)
        hue_names = self.x_order if self.x_palette else self.hue_order
        if self.hue is None and not self.x_palette:
            if color is None:
                color = color_palette()[0]
            palette = [color for _ in self.x_order]
        elif palette is None:
            palette = color_palette(n_colors=n_hues)
        elif isinstance(palette, dict):
            palette = [palette[k] for k in hue_names]
            palette = color_palette(palette, n_hues)
        else:
            palette = color_palette(palette, n_hues)
        self.palette = palette

        if self.kind == "point":
            self.err_palette = palette
        else:
            # TODO make this smarter
            self.err_palette = [self.err_color] * len(palette)

    def establish_positions(self, dodge):
        """Make list of center values for each x and offset for each hue."""
        self.positions = np.arange(len(self.x_order))

        # If there's no hue variable kind is irrelevant
        if self.hue is None:
            n_hues = 1
            width = self.bar_widths
            offset = np.zeros(n_hues)
        else:
            n_hues = len(self.hue_order)

            # Bar offset is set by hardcoded bar width
            if self.kind in ["bar", "box"]:
                width = self.bar_widths / n_hues
                offset = np.linspace(0, self.bar_widths - width, n_hues)
                if self.kind == "box":
                    width *= .95
                self.bar_widths = width

            # Point offset is set by `dodge` parameter
            elif self.kind == "point":
                offset = np.linspace(0, dodge, n_hues)
            offset -= offset.mean()

        self.offset = offset

    def establish_plot_kind(self, kind):
        """Use specified kind of apply heuristics to decide automatically."""
        if kind == "auto":
            y = self.y

            # Walk through some heuristics to automatically assign a kind
            if self.y_count:
                kind = "bar"
            elif y.max() <= 1:
                kind = "point"
            elif (y.mean() / y.std()) < 2.5:
                kind = "bar"
            else:
                kind = "point"
            self.kind = kind
        elif kind in ["bar", "point", "box"]:
            self.kind = kind
        else:
            raise ValueError("%s is not a valid kind of plot" % kind)

    @property
    def estimate_data(self):
        """Generator to yield x, y, and ci data for each hue subset."""
        # First iterate through the hues, as plots are drawn for all
        # positions of a given hue at the same time
        for i, hue in enumerate(self.hue_order):

            # Build intermediate lists of the values for each drawing
            pos = []
            height = []
            ci = []
            for j, x in enumerate(self.x_order):

                pos.append(self.positions[j] + self.offset[i])

                # Focus on the data for this specific bar/point
                current_data = (self.x == x) & (self.hue == hue)
                y_data = self.y[current_data]
                if self.units is None:
                    unit_data = None
                else:
                    unit_data = self.units[current_data]

                # This is where the main computation happens
                height.append(self.estimator(y_data))
                if self.ci is not None:
                    boots = algo.bootstrap(y_data, func=self.estimator,
                                           n_boot=self.n_boot,
                                           units=unit_data)
                    ci.append(utils.ci(boots, self.ci))

            yield pos, height, ci

    @property
    def binned_data(self):
        """Generator to yield entire subsets of data for each bin."""
        # First iterate through the hues, as plots are drawn for all
        # positions of a given hue at the same time
        for i, hue in enumerate(self.hue_order):

            # Build intermediate lists of the values for each drawing
            pos = []
            data = []
            for j, x in enumerate(self.x_order):

                pos.append(self.positions[j] + self.offset[i])
                current_data = (self.x == x) & (self.hue == hue)
                data.append(self.y[current_data])

            yield pos, data

    def plot(self, ax):
        """Plot based on the stored value for kind of plot."""
        plotter = getattr(self, self.kind + "plot")
        plotter(ax)

        # Set the plot attributes (these are shared across plot kinds
        if self.hue is not None:
            leg = ax.legend(loc="best", scatterpoints=1)
            if hasattr(self.hue, "name"):
                leg.set_title(self.hue.name),

                # Set the title size a roundabout way to maintain
                # compatability with matplotlib 1.1
                titlesize = mpl.rcParams["axes.labelsize"]
                prop = mpl.font_manager.FontProperties(size=titlesize)
                leg._legend_title_box._text.set_font_properties(prop)

        ax.xaxis.grid(False)
        ax.set_xticks(self.positions)
        ax.set_xticklabels(self.x_order)
        if hasattr(self.x, "name"):
            ax.set_xlabel(self.x.name)
        if self.y_count:
            ax.set_ylabel("count")
        else:
            if hasattr(self.y, "name"):
                ax.set_ylabel(self.y.name)

        if self.hline is not None:
            ymin, ymax = ax.get_ylim()
            if self.hline > ymin and self.hline < ymax:
                ax.axhline(self.hline, c="#666666")

    def barplot(self, ax):
        """Draw the plot with a bar representation."""
        for i, (pos, height, ci) in enumerate(self.estimate_data):

            color = self.palette if self.x_palette else self.palette[i]
            ecolor = self.err_palette[i]
            label = self.hue_order[i]

            # The main plot
            ax.bar(pos, height, self.bar_widths, color=color,
                   label=label, align="center")

            # The error bars
            for x, (low, high) in zip(pos, ci):
                ax.plot([x, x], [low, high], linewidth=self.lw, color=ecolor)

        # Set the x limits
        offset = .5
        xlim = self.positions.min() - offset, self.positions.max() + offset
        ax.set_xlim(xlim)

    def boxplot(self, ax):
        """Draw the plot with a bar representation."""
        from .distributions import boxplot
        for i, (pos, data) in enumerate(self.binned_data):

            color = self.palette if self.x_palette else self.palette[i]
            label = self.hue_order[i]

            # The main plot
            boxplot(data, widths=self.bar_widths, color=color,
                    positions=pos, label=label, ax=ax)

        # Set the x limits
        offset = .5
        xlim = self.positions.min() - offset, self.positions.max() + offset
        ax.set_xlim(xlim)

    def pointplot(self, ax):
        """Draw the plot with a point representation."""
        for i, (pos, height, ci) in enumerate(self.estimate_data):

            color = self.palette if self.x_palette else self.palette[i]
            err_palette = self.err_palette
            label = self.hue_order[i]
            marker = self.markers[i]
            markersize = np.pi * np.square(self.lw) * 2
            linestyle = self.linestyles[i]
            z = i + 1

            # The error bars
            for j, (x, (low, high)) in enumerate(zip(pos, ci)):
                ecolor = err_palette[j] if self.x_palette else err_palette[i]
                ax.plot([x, x], [low, high], linewidth=self.lw,
                        color=ecolor, zorder=z)

            # The main plot
            ax.scatter(pos, height, s=markersize, color=color, label=label,
                       marker=marker, zorder=z)

            # The join line
            if self.join:
                ax.plot(pos, height, color=color,
                        linewidth=self.lw, linestyle=linestyle, zorder=z)

        # Set the x limits
        xlim = (self.positions.min() + self.offset.min() - .3,
                self.positions.max() + self.offset.max() + .3)
        ax.set_xlim(xlim)


class _RegressionPlotter(_LinearPlotter):
    """Plotter for numeric independent variables with regression model.

    This does the computations and drawing for the `regplot` function, and
    is thus also used indirectly by `lmplot`. It is generally similar to
    the `_DiscretePlotter`, but it's intended for use when the independent
    variable is numeric (continuous or discrete), and its primary advantage
    is that a regression model can be fit to the data and visualized, allowing
    extrapolations beyond the observed datapoints.

    """
    def __init__(self, x, y, data=None, x_estimator=None, x_bins=None,
                 x_ci="ci", scatter=True, fit_reg=True, ci=95, n_boot=1000,
                 units=None, order=1, logistic=False, lowess=False,
                 robust=False, x_partial=None, y_partial=None,
                 truncate=False, dropna=True, x_jitter=None, y_jitter=None,
                 color=None, label=None):

        # Set member attributes
        self.x_estimator = x_estimator
        self.ci = ci
        self.x_ci = ci if x_ci == "ci" else x_ci
        self.n_boot = n_boot
        self.scatter = scatter
        self.fit_reg = fit_reg
        self.order = order
        self.logistic = logistic
        self.lowess = lowess
        self.robust = robust
        self.truncate = truncate
        self.x_jitter = x_jitter
        self.y_jitter = y_jitter
        self.color = color
        self.label = label

        # Validate the regression options:
        if sum((order > 1, logistic, robust, lowess)) > 1:
            raise ValueError("Mutually exclusive regression options.")

        # Extract the data vals from the arguments or passed dataframe
        self.establish_variables(data, x=x, y=y, units=units,
                                 x_partial=x_partial, y_partial=y_partial)

        # Drop null observations
        if dropna:
            self.dropna("x", "y", "units", "x_partial", "y_partial")

        # Regress nuisance variables out of the data
        if self.x_partial is not None:
            self.x = self.regress_out(self.x, self.x_partial)
        if self.y_partial is not None:
            self.y = self.regress_out(self.y, self.y_partial)

        # Possibly bin the predictor variable, which implies a point estimate
        if x_bins is not None:
            self.x_estimator = np.mean if x_estimator is None else x_estimator
            x_discrete, x_bins = self.bin_predictor(x_bins)
            self.x_discrete = x_discrete
        else:
            self.x_discrete = self.x

        # Save the range of the x variable for the grid later
        self.x_range = self.x.min(), self.x.max()

    @property
    def scatter_data(self):
        """Data where each observation is a point."""
        x_j = self.x_jitter
        if x_j is None:
            x = self.x
        else:
            x = self.x + np.random.uniform(-x_j, x_j, len(self.x))

        y_j = self.y_jitter
        if y_j is None:
            y = self.y
        else:
            y = self.y + np.random.uniform(-y_j, y_j, len(self.y))

        return x, y

    @property
    def estimate_data(self):
        """Data with a point estimate and CI for each discrete x value."""
        x, y = self.x_discrete, self.y
        vals = sorted(np.unique(x))
        points, cis = [], []

        for val in vals:

            # Get the point estimate of the y variable
            _y = y[x == val]
            est = self.x_estimator(_y)
            points.append(est)

            # Compute the confidence interval for this estimate
            if self.x_ci is None:
                cis.append(None)
            else:
                units = None
                if self.units is not None:
                    units = self.units[x == val]
                boots = algo.bootstrap(_y, func=self.x_estimator,
                                       n_boot=self.n_boot, units=units)
                _ci = utils.ci(boots, self.x_ci)
                cis.append(_ci)

        return vals, points, cis

    def fit_regression(self, ax=None, x_range=None, grid=None):
        """Fit the regression model."""
        # Create the grid for the regression
        if grid is None:
            if self.truncate:
                x_min, x_max = self.x_range
            else:
                if ax is None:
                    x_min, x_max = x_range
                else:
                    x_min, x_max = ax.get_xlim()
            grid = np.linspace(x_min, x_max, 100)
        ci = self.ci

        # Fit the regression
        if self.order > 1:
            yhat, yhat_boots = self.fit_poly(grid, self.order)
        elif self.logistic:
            from statsmodels.api import GLM, families
            yhat, yhat_boots = self.fit_statsmodels(grid, GLM,
                                                    family=families.Binomial())
        elif self.lowess:
            ci = None
            grid, yhat = self.fit_lowess()
        elif self.robust:
            from statsmodels.api import RLM
            yhat, yhat_boots = self.fit_statsmodels(grid, RLM)
        else:
            yhat, yhat_boots = self.fit_fast(grid)

        # Compute the confidence interval at each grid point
        if ci is None:
            err_bands = None
        else:
            err_bands = utils.ci(yhat_boots, ci, axis=0)

        return grid, yhat, err_bands

    def fit_fast(self, grid):
        """Low-level regression and prediction using linear algebra."""
        X, y = np.c_[np.ones(len(self.x)), self.x], self.y
        grid = np.c_[np.ones(len(grid)), grid]
        reg_func = lambda _x, _y: np.linalg.pinv(_x).dot(_y)
        yhat = grid.dot(reg_func(X, y))
        if self.ci is None:
            return yhat, None

        beta_boots = algo.bootstrap(X, y, func=reg_func,
                                    n_boot=self.n_boot, units=self.units).T
        yhat_boots = grid.dot(beta_boots).T
        return yhat, yhat_boots

    def fit_poly(self, grid, order):
        """Regression using numpy polyfit for higher-order trends."""
        x, y = self.x, self.y
        reg_func = lambda _x, _y: np.polyval(np.polyfit(_x, _y, order), grid)
        yhat = reg_func(x, y)
        if self.ci is None:
            return yhat, None

        yhat_boots = algo.bootstrap(x, y, func=reg_func,
                                    n_boot=self.n_boot, units=self.units)
        return yhat, yhat_boots

    def fit_statsmodels(self, grid, model, **kwargs):
        """More general regression function using statsmodels objects."""
        X, y = np.c_[np.ones(len(self.x)), self.x], self.y
        grid = np.c_[np.ones(len(grid)), grid]
        reg_func = lambda _x, _y: model(_y, _x, **kwargs).fit().predict(grid)
        yhat = reg_func(X, y)
        if self.ci is None:
            return yhat, None

        yhat_boots = algo.bootstrap(X, y, func=reg_func,
                                    n_boot=self.n_boot, units=self.units)
        return yhat, yhat_boots

    def fit_lowess(self):
        """Fit a locally-weighted regression, which returns its own grid."""
        from statsmodels.api import nonparametric
        grid, yhat = nonparametric.lowess(self.y, self.x).T
        return grid, yhat

    def bin_predictor(self, bins):
        """Discretize a predictor by assigning value to closest bin."""
        x = self.x
        if np.isscalar(bins):
            percentiles = np.linspace(0, 100, bins + 2)[1:-1]
            bins = np.c_[utils.percentiles(x, percentiles)]
        else:
            bins = np.c_[np.ravel(bins)]

        dist = distance.cdist(np.c_[x], bins)
        x_binned = bins[np.argmin(dist, axis=1)].ravel()

        return x_binned, bins.ravel()

    def regress_out(self, a, b):
        """Regress b from a keeping a's original mean."""
        a_mean = a.mean()
        a = a - a_mean
        b = b - b.mean()
        b = np.c_[b]
        a_prime = a - b.dot(np.linalg.pinv(b).dot(a))
        return (a_prime + a_mean).reshape(a.shape)

    def plot(self, ax, scatter_kws, line_kws):
        """Draw the full plot."""
        # Insert the plot label into the correct set of keyword arguments
        if self.fit_reg:
            line_kws["label"] = self.label
        else:
            scatter_kws["label"] = self.label

        # Use the current color cycle state as a default
        if self.color is None:
            lines, = plt.plot(self.x.mean(), self.y.mean())
            color = lines.get_color()
            lines.remove()
        else:
            color = self.color

        # Let color in keyword arguments override overall plot color
        scatter_kws.setdefault("color", color)
        line_kws.setdefault("color", color)

        # Draw the constituent plots
        if self.scatter:
            self.scatterplot(ax, scatter_kws)
        if self.fit_reg:
            self.lineplot(ax, line_kws)

        # Label the axes
        if hasattr(self.x, "name"):
            ax.set_xlabel(self.x.name)
        if hasattr(self.y, "name"):
            ax.set_ylabel(self.y.name)

    def scatterplot(self, ax, kws):
        """Draw the data."""
        if self.x_estimator is None:
            kws.setdefault("alpha", .8)
            x, y = self.scatter_data
            ax.scatter(x, y, **kws)
        else:
            # TODO abstraction
            ci_kws = {"color": kws["color"]}
            ci_kws["linewidth"] = mpl.rcParams["lines.linewidth"] * 1.75
            kws.setdefault("s", 50)

            xs, ys, cis = self.estimate_data
            if [ci for ci in cis if ci is not None]:
                for x, ci in zip(xs, cis):<commit_msg>fixed alpha setting bug in _RegressionPlotter scatter<commit_after>"""Plotting functions for linear models (broadly construed)."""
from __future__ import division
import copy
import itertools
import warnings
import numpy as np
import pandas as pd
from scipy.spatial import distance
import matplotlib as mpl
import matplotlib.pyplot as plt

try:
    import statsmodels.api as sm
    import statsmodels.formula.api as sf
    _has_statsmodels = True
except ImportError:
    _has_statsmodels = False

from .external.six import string_types
from .external.six.moves import range

from . import utils
from . import algorithms as algo
from .palettes import color_palette
from .axisgrid import FacetGrid


class _LinearPlotter(object):

    """Base class for plotting relational data in tidy format.

    To get anything useful done you'll have to inherit from this, but setup
    code that can be abstracted out should be put here.

    """

    def establish_variables(self, data, **kws):
        """Extract variables from data or use directly."""
        self.data = data

        # Validate the inputs
        any_strings = any([isinstance(v, string_types) for v in kws.values()])
        if any_strings and data is None:
            raise ValueError("Must pass `data` if using named variables.")

        # Set the variables
        for var, val in kws.items():
            if isinstance(val, string_types):
                setattr(self, var, data[val])
            else:
                setattr(self, var, val)

    def dropna(self, *vars):
        """Remove observations with missing data."""
        vals = [getattr(self, var) for var in vars]
        vals = [v for v in vals if v is not None]
        not_na = np.all(np.column_stack([pd.notnull(v) for v in vals]), axis=1)
        for var in vars:
            val = getattr(self, var)
            if val is not None:
                setattr(self, var, val[not_na])

    def plot(self, ax):
        raise NotImplementedError


class _DiscretePlotter(_LinearPlotter):

    """Plotter for data with discrete independent variable(s).

    This will be used by the `barplot` and `pointplot` functions, and
    thus indirectly by the `factorplot` function. It can produce plots
    where some statistic for a dependent measure is estimated within
    subsets of the data, which can be hierarchically structured at up to two
    levels (`x` and `hue`). The plots can be drawn with a few different
    visual representations of the same underlying data (`bar`, and `point`,
    with `box` doing something similar but skipping the estimation).

    """

    def __init__(self, x, y=None, hue=None, data=None, units=None,
                 x_order=None, hue_order=None, color=None, palette=None,
                 kind="auto", markers=None, linestyles=None, dodge=0,
                 join=True, hline=None, estimator=np.mean, ci=95,
                 n_boot=1000, dropna=True):

        # This implies we have a single bar/point for each level of `x`
        # but that the different levels should be mapped with a palette
        self.x_palette = hue is None and palette is not None

        # Set class attributes based on inputs
        self.estimator = len if y is None else estimator
        self.ci = None if y is None else ci
        self.join = join
        self.n_boot = n_boot
        self.hline = hline

        # Other attributes that are hardcoded for now
        self.bar_widths = .8
        self.err_color = "#444444"
        self.lw = mpl.rcParams["lines.linewidth"] * 1.8

        # Once we've set the above values, if `y` is None we want the actual
        # y values to be the x values so we can count them
        self.y_count = y is None
        if y is None:
            y = x

        # Ascertain which values will be associated with what values
        self.establish_variables(data, x=x, y=y, hue=hue, units=units)

        # Figure out the order of the variables on the x axis
        x_sorted = np.sort(pd.unique(self.x))
        self.x_order = x_sorted if x_order is None else x_order
        if self.hue is not None:
            hue_sorted = np.sort(pd.unique(self.hue))
            self.hue_order = hue_sorted if hue_order is None else hue_order
        else:
            self.hue_order = [None]

        # Handle the other hue-mapped attributes
        if markers is None:
            self.markers = ["o"] * len(self.hue_order)
        else:
            if len(markers) != len(self.hue_order):
                raise ValueError("Length of marker list must equal "
                                 "number of hue levels")
            self.markers = markers
        if linestyles is None:
            self.linestyles = ["-"] * len(self.hue_order)
        else:
            if len(linestyles) != len(self.hue_order):
                raise ValueError("Length of linestyle list must equal "
                                 "number of hue levels")
            self.linestyles = linestyles

        # Drop null observations
        if dropna:
            self.dropna("x", "y", "hue", "units")

        # Settle whe kind of plot this is going to be
        self.establish_plot_kind(kind)

        # Determine the color palette
        self.establish_palette(color, palette)

        # Figure out where the data should be drawn
        self.establish_positions(dodge)

    def establish_palette(self, color, palette):
        """Set a list of colors for each plot element."""
        n_hues = len(self.x_order) if self.x_palette else len(self.hue_order)
        hue_names = self.x_order if self.x_palette else self.hue_order
        if self.hue is None and not self.x_palette:
            if color is None:
                color = color_palette()[0]
            palette = [color for _ in self.x_order]
        elif palette is None:
            palette = color_palette(n_colors=n_hues)
        elif isinstance(palette, dict):
            palette = [palette[k] for k in hue_names]
            palette = color_palette(palette, n_hues)
        else:
            palette = color_palette(palette, n_hues)
        self.palette = palette

        if self.kind == "point":
            self.err_palette = palette
        else:
            # TODO make this smarter
            self.err_palette = [self.err_color] * len(palette)

    def establish_positions(self, dodge):
        """Make list of center values for each x and offset for each hue."""
        self.positions = np.arange(len(self.x_order))

        # If there's no hue variable kind is irrelevant
        if self.hue is None:
            n_hues = 1
            width = self.bar_widths
            offset = np.zeros(n_hues)
        else:
            n_hues = len(self.hue_order)

            # Bar offset is set by hardcoded bar width
            if self.kind in ["bar", "box"]:
                width = self.bar_widths / n_hues
                offset = np.linspace(0, self.bar_widths - width, n_hues)
                if self.kind == "box":
                    width *= .95
                self.bar_widths = width

            # Point offset is set by `dodge` parameter
            elif self.kind == "point":
                offset = np.linspace(0, dodge, n_hues)
            offset -= offset.mean()

        self.offset = offset

    def establish_plot_kind(self, kind):
        """Use specified kind of apply heuristics to decide automatically."""
        if kind == "auto":
            y = self.y

            # Walk through some heuristics to automatically assign a kind
            if self.y_count:
                kind = "bar"
            elif y.max() <= 1:
                kind = "point"
            elif (y.mean() / y.std()) < 2.5:
                kind = "bar"
            else:
                kind = "point"
            self.kind = kind
        elif kind in ["bar", "point", "box"]:
            self.kind = kind
        else:
            raise ValueError("%s is not a valid kind of plot" % kind)

    @property
    def estimate_data(self):
        """Generator to yield x, y, and ci data for each hue subset."""
        # First iterate through the hues, as plots are drawn for all
        # positions of a given hue at the same time
        for i, hue in enumerate(self.hue_order):

            # Build intermediate lists of the values for each drawing
            pos = []
            height = []
            ci = []
            for j, x in enumerate(self.x_order):

                pos.append(self.positions[j] + self.offset[i])

                # Focus on the data for this specific bar/point
                current_data = (self.x == x) & (self.hue == hue)
                y_data = self.y[current_data]
                if self.units is None:
                    unit_data = None
                else:
                    unit_data = self.units[current_data]

                # This is where the main computation happens
                height.append(self.estimator(y_data))
                if self.ci is not None:
                    boots = algo.bootstrap(y_data, func=self.estimator,
                                           n_boot=self.n_boot,
                                           units=unit_data)
                    ci.append(utils.ci(boots, self.ci))

            yield pos, height, ci

    @property
    def binned_data(self):
        """Generator to yield entire subsets of data for each bin."""
        # First iterate through the hues, as plots are drawn for all
        # positions of a given hue at the same time
        for i, hue in enumerate(self.hue_order):

            # Build intermediate lists of the values for each drawing
            pos = []
            data = []
            for j, x in enumerate(self.x_order):

                pos.append(self.positions[j] + self.offset[i])
                current_data = (self.x == x) & (self.hue == hue)
                data.append(self.y[current_data])

            yield pos, data

    def plot(self, ax):
        """Plot based on the stored value for kind of plot."""
        plotter = getattr(self, self.kind + "plot")
        plotter(ax)

        # Set the plot attributes (these are shared across plot kinds
        if self.hue is not None:
            leg = ax.legend(loc="best", scatterpoints=1)
            if hasattr(self.hue, "name"):
                leg.set_title(self.hue.name),

                # Set the title size a roundabout way to maintain
                # compatability with matplotlib 1.1
                titlesize = mpl.rcParams["axes.labelsize"]
                prop = mpl.font_manager.FontProperties(size=titlesize)
                leg._legend_title_box._text.set_font_properties(prop)

        ax.xaxis.grid(False)
        ax.set_xticks(self.positions)
        ax.set_xticklabels(self.x_order)
        if hasattr(self.x, "name"):
            ax.set_xlabel(self.x.name)
        if self.y_count:
            ax.set_ylabel("count")
        else:
            if hasattr(self.y, "name"):
                ax.set_ylabel(self.y.name)

        if self.hline is not None:
            ymin, ymax = ax.get_ylim()
            if self.hline > ymin and self.hline < ymax:
                ax.axhline(self.hline, c="#666666")

    def barplot(self, ax):
        """Draw the plot with a bar representation."""
        for i, (pos, height, ci) in enumerate(self.estimate_data):

            color = self.palette if self.x_palette else self.palette[i]
            ecolor = self.err_palette[i]
            label = self.hue_order[i]

            # The main plot
            ax.bar(pos, height, self.bar_widths, color=color,
                   label=label, align="center")

            # The error bars
            for x, (low, high) in zip(pos, ci):
                ax.plot([x, x], [low, high], linewidth=self.lw, color=ecolor)

        # Set the x limits
        offset = .5
        xlim = self.positions.min() - offset, self.positions.max() + offset
        ax.set_xlim(xlim)

    def boxplot(self, ax):
        """Draw the plot with a bar representation."""
        from .distributions import boxplot
        for i, (pos, data) in enumerate(self.binned_data):

            color = self.palette if self.x_palette else self.palette[i]
            label = self.hue_order[i]

            # The main plot
            boxplot(data, widths=self.bar_widths, color=color,
                    positions=pos, label=label, ax=ax)

        # Set the x limits
        offset = .5
        xlim = self.positions.min() - offset, self.positions.max() + offset
        ax.set_xlim(xlim)

    def pointplot(self, ax):
        """Draw the plot with a point representation."""
        for i, (pos, height, ci) in enumerate(self.estimate_data):

            color = self.palette if self.x_palette else self.palette[i]
            err_palette = self.err_palette
            label = self.hue_order[i]
            marker = self.markers[i]
            markersize = np.pi * np.square(self.lw) * 2
            linestyle = self.linestyles[i]
            z = i + 1

            # The error bars
            for j, (x, (low, high)) in enumerate(zip(pos, ci)):
                ecolor = err_palette[j] if self.x_palette else err_palette[i]
                ax.plot([x, x], [low, high], linewidth=self.lw,
                        color=ecolor, zorder=z)

            # The main plot
            ax.scatter(pos, height, s=markersize, color=color, label=label,
                       marker=marker, zorder=z)

            # The join line
            if self.join:
                ax.plot(pos, height, color=color,
                        linewidth=self.lw, linestyle=linestyle, zorder=z)

        # Set the x limits
        xlim = (self.positions.min() + self.offset.min() - .3,
                self.positions.max() + self.offset.max() + .3)
        ax.set_xlim(xlim)


class _RegressionPlotter(_LinearPlotter):

    """Plotter for numeric independent variables with regression model.

    This does the computations and drawing for the `regplot` function, and
    is thus also used indirectly by `lmplot`. It is generally similar to
    the `_DiscretePlotter`, but it's intended for use when the independent
    variable is numeric (continuous or discrete), and its primary advantage
    is that a regression model can be fit to the data and visualized, allowing
    extrapolations beyond the observed datapoints.

    """

    def __init__(self, x, y, data=None, x_estimator=None, x_bins=None,
                 x_ci="ci", scatter=True, fit_reg=True, ci=95, n_boot=1000,
                 units=None, order=1, logistic=False, lowess=False,
                 robust=False, x_partial=None, y_partial=None,
                 truncate=False, dropna=True, x_jitter=None, y_jitter=None,
                 color=None, label=None):

        # Set member attributes
        self.x_estimator = x_estimator
        self.ci = ci
        self.x_ci = ci if x_ci == "ci" else x_ci
        self.n_boot = n_boot
        self.scatter = scatter
        self.fit_reg = fit_reg
        self.order = order
        self.logistic = logistic
        self.lowess = lowess
        self.robust = robust
        self.truncate = truncate
        self.x_jitter = x_jitter
        self.y_jitter = y_jitter
        self.color = color
        self.label = label

        # Validate the regression options:
        if sum((order > 1, logistic, robust, lowess)) > 1:
            raise ValueError("Mutually exclusive regression options.")

        # Extract the data vals from the arguments or passed dataframe
        self.establish_variables(data, x=x, y=y, units=units,
                                 x_partial=x_partial, y_partial=y_partial)

        # Drop null observations
        if dropna:
            self.dropna("x", "y", "units", "x_partial", "y_partial")

        # Regress nuisance variables out of the data
        if self.x_partial is not None:
            self.x = self.regress_out(self.x, self.x_partial)
        if self.y_partial is not None:
            self.y = self.regress_out(self.y, self.y_partial)

        # Possibly bin the predictor variable, which implies a point estimate
        if x_bins is not None:
            self.x_estimator = np.mean if x_estimator is None else x_estimator
            x_discrete, x_bins = self.bin_predictor(x_bins)
            self.x_discrete = x_discrete
        else:
            self.x_discrete = self.x

        # Save the range of the x variable for the grid later
        self.x_range = self.x.min(), self.x.max()

    @property
    def scatter_data(self):
        """Data where each observation is a point."""
        x_j = self.x_jitter
        if x_j is None:
            x = self.x
        else:
            x = self.x + np.random.uniform(-x_j, x_j, len(self.x))

        y_j = self.y_jitter
        if y_j is None:
            y = self.y
        else:
            y = self.y + np.random.uniform(-y_j, y_j, len(self.y))

        return x, y

    @property
    def estimate_data(self):
        """Data with a point estimate and CI for each discrete x value."""
        x, y = self.x_discrete, self.y
        vals = sorted(np.unique(x))
        points, cis = [], []

        for val in vals:

            # Get the point estimate of the y variable
            _y = y[x == val]
            est = self.x_estimator(_y)
            points.append(est)

            # Compute the confidence interval for this estimate
            if self.x_ci is None:
                cis.append(None)
            else:
                units = None
                if self.units is not None:
                    units = self.units[x == val]
                boots = algo.bootstrap(_y, func=self.x_estimator,
                                       n_boot=self.n_boot, units=units)
                _ci = utils.ci(boots, self.x_ci)
                cis.append(_ci)

        return vals, points, cis

    def fit_regression(self, ax=None, x_range=None, grid=None):
        """Fit the regression model."""
        # Create the grid for the regression
        if grid is None:
            if self.truncate:
                x_min, x_max = self.x_range
            else:
                if ax is None:
                    x_min, x_max = x_range
                else:
                    x_min, x_max = ax.get_xlim()
            grid = np.linspace(x_min, x_max, 100)
        ci = self.ci

        # Fit the regression
        if self.order > 1:
            yhat, yhat_boots = self.fit_poly(grid, self.order)
        elif self.logistic:
            from statsmodels.api import GLM, families
            yhat, yhat_boots = self.fit_statsmodels(grid, GLM,
                                                    family=families.Binomial())
        elif self.lowess:
            ci = None
            grid, yhat = self.fit_lowess()
        elif self.robust:
            from statsmodels.api import RLM
            yhat, yhat_boots = self.fit_statsmodels(grid, RLM)
        else:
            yhat, yhat_boots = self.fit_fast(grid)

        # Compute the confidence interval at each grid point
        if ci is None:
            err_bands = None
        else:
            err_bands = utils.ci(yhat_boots, ci, axis=0)

        return grid, yhat, err_bands

    def fit_fast(self, grid):
        """Low-level regression and prediction using linear algebra."""
        X, y = np.c_[np.ones(len(self.x)), self.x], self.y
        grid = np.c_[np.ones(len(grid)), grid]
        reg_func = lambda _x, _y: np.linalg.pinv(_x).dot(_y)
        yhat = grid.dot(reg_func(X, y))
        if self.ci is None:
            return yhat, None

        beta_boots = algo.bootstrap(X, y, func=reg_func,
                                    n_boot=self.n_boot, units=self.units).T
        yhat_boots = grid.dot(beta_boots).T
        return yhat, yhat_boots

    def fit_poly(self, grid, order):
        """Regression using numpy polyfit for higher-order trends."""
        x, y = self.x, self.y
        reg_func = lambda _x, _y: np.polyval(np.polyfit(_x, _y, order), grid)
        yhat = reg_func(x, y)
        if self.ci is None:
            return yhat, None

        yhat_boots = algo.bootstrap(x, y, func=reg_func,
                                    n_boot=self.n_boot, units=self.units)
        return yhat, yhat_boots

    def fit_statsmodels(self, grid, model, **kwargs):
        """More general regression function using statsmodels objects."""
        X, y = np.c_[np.ones(len(self.x)), self.x], self.y
        grid = np.c_[np.ones(len(grid)), grid]
        reg_func = lambda _x, _y: model(_y, _x, **kwargs).fit().predict(grid)
        yhat = reg_func(X, y)
        if self.ci is None:
            return yhat, None

        yhat_boots = algo.bootstrap(X, y, func=reg_func,
                                    n_boot=self.n_boot, units=self.units)
        return yhat, yhat_boots

    def fit_lowess(self):
        """Fit a locally-weighted regression, which returns its own grid."""
        from statsmodels.api import nonparametric
        grid, yhat = nonparametric.lowess(self.y, self.x).T
        return grid, yhat

    def bin_predictor(self, bins):
        """Discretize a predictor by assigning value to closest bin."""
        x = self.x
        if np.isscalar(bins):
            percentiles = np.linspace(0, 100, bins + 2)[1:-1]
            bins = np.c_[utils.percentiles(x, percentiles)]
        else:
            bins = np.c_[np.ravel(bins)]

        dist = distance.cdist(np.c_[x], bins)
        x_binned = bins[np.argmin(dist, axis=1)].ravel()

        return x_binned, bins.ravel()

    def regress_out(self, a, b):
        """Regress b from a keeping a's original mean."""
        a_mean = a.mean()
        a = a - a_mean
        b = b - b.mean()
        b = np.c_[b]
        a_prime = a - b.dot(np.linalg.pinv(b).dot(a))
        return (a_prime + a_mean).reshape(a.shape)

    def plot(self, ax, scatter_kws, line_kws):
        """Draw the full plot."""
        # Insert the plot label into the correct set of keyword arguments
        if self.fit_reg:
            line_kws["label"] = self.label
        else:
            scatter_kws["label"] = self.label

        # Use the current color cycle state as a default
        if self.color is None:
            lines, = plt.plot(self.x.mean(), self.y.mean())
            color = lines.get_color()
            lines.remove()
        else:
            color = self.color

        # Let color in keyword arguments override overall plot color
        scatter_kws.setdefault("color", color)
        line_kws.setdefault("color", color)

        # Draw the constituent plots
        if self.scatter:
            self.scatterplot(ax, scatter_kws)
        if self.fit_reg:
            self.lineplot(ax, line_kws)

        # Label the axes
        if hasattr(self.x, "name"):
            ax.set_xlabel(self.x.name)
        if hasattr(self.y, "name"):
            ax.set_ylabel(self.y.name)

    def scatterplot(self, ax, kws):
        """Draw the data."""
        if self.x_estimator is None:
            if not hasattr(kws['color'], 'shape') or kws['color'].shape[1] < 4:
                kws.setdefault("alpha", .8)
            x, y = self.scatter_data
            ax.scatter(x, y, **kws)
        else:
            # TODO abstraction
            ci_kws = {"color": kws["color"]}
            ci_kws["linewidth"] = mpl.rcParams["lines.linewidth"] * 1.75
            kws.setdefault("s", 50)

            xs, ys, cis = self.estimate_data
            if [ci for ci in cis if ci is not None]:
                for x, ci in zip(xs, cis):
{'fraction_non_alphanumeric': 0.09520342612419701, 'fraction_numerical': 0.0029978586723768737, 'mean_word_length': 1.803517829271221, 'pattern_counts': {'":': 10, '<': 10, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 26, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3675702', 'n_tokens_mistral': 14280, 'n_tokens_neox': 13564, 'n_words': 4208}