<commit_before>"""Methods for plotting 2D histograms with optional 1D projected histograms for each axis.

Functions
---------
-    plot2DHistProj          - Plot a 2D histogram with projections of one or both axes.
-    plot2DHist              - Plot the given 2D histogram of data.

-    _constructFigure        - Add the required axes to the given figure object.

To-do
-----
-   `plot2DHistProj` and `plot2DHist` use different styles of input data, specifically
    `plot2DHist` requires a 2D histogram of data to be passed, while `plot2DHistProj` accepts
    separate arrays of x and y data to be histogrammed using ``scipy.stats.binned_statistic_2d``.

"""
from __future__ import absolute_import, division, print_function, unicode_literals
from six.moves import xrange

import numpy as np
import scipy as sp
import scipy.stats
import matplotlib.pyplot as plt

import zcode.math as zmath
from . import plot_core

__all__ = ['plot2DHist', 'plot2DHistProj']

_LEFT = 0.09
_RIGHT = 0.92     # Location of right of plots
_BOTTOM = 0.09
_TOP = 0.90       # Location of top of plots

_CB_WID = 0.02
_CB_WPAD = 0.1


def plot2DHistProj(xvals, yvals, weights=None, statistic=None, bins=10, filter=None, extrema=None,
                   fig=None, xproj=True, yproj=True, hratio=0.7, wratio=0.7, pad=0.0, alpha=1.0,
                   cmap=None, smap=None, type='hist', scale_to_cbar=True,
                   fs=12, scale='log', histScale='log', labels=None, cbar=True, write_counts=False,
                   left=_LEFT, bottom=_BOTTOM, right=_RIGHT, top=_TOP, lo=None, hi=None):
    """Plot a 2D histogram with projections of one or both axes.

    Arguments
    ---------
    xvals : (N,) array_like,
        Values corresponding to the x-points of the given data
    yvals : (N,) array_like,
        Values corresponding to the y-points of the given data
    weights : (N,) array_like or `None`,
        Weights used to create histograms.  If `None`, then counts are used.
    statistic : str or `None`,
        Type of statistic to be calculated, passed to ``scipy.stats.binned_statistic``.
        e.g. {'count', 'sum', 'mean'}.
        If `None`, then either 'sum' or 'count' is used depending on if `weights` are
        provieded or not.
    bins : int or [int, int] or array_like or [array, array],
        Specification for bin sizes.  integer values are treated as the number of bins to use,
        while arrays are used as the bin edges themselves.  If a tuple of two values is given, it
        is assumed that the first is for the x-axis and the second for the y-axis.
    filter : str or `None`, or [2,] tuple of str or `None`, or [3,] tubple of str or `None`
        String specifying how to filter the input `data` relative to zero.
        If this is a single value, it is applies to both `xvals` and `yvals`.
        If this is a tuple/list of two values, they correspond to `xvals` and `yvals` respectively.
        If `weights` are provided, the tuple/list should have three values.
    extrema :
    fig : ``matplotlib.figure.figure``,
        Figure instance to which axes are added for plotting.  One is created if not given.
    xproj : bool,
        Whether to also plot the projection of the x-axis (i.e. histogram ignoring y-values).
    yproj : bool,
        Whether to also plot the projection of the y-axis (i.e. histogram ignoring x-values).
    hratio : float,
        Fraction of the total available height-space to use for the primary axes object (2D hist)
    wratio : float,
        Fraction of the total available width-space to use for the primary axes object (2D hist)
    pad : float,
        Padding between central axis and the projected ones.
    cmap : ``matplotlib.colors.Colormap`` object
        Matplotlib colormap to use for coloring histogram.
        Overridden if `smap` is provided.
    smap : `matplotlib.cm.ScalarMappable` object or `None`
        A scalar-mappable object to use for colormaps, or `None` for one to be created.
    type : str, {'hist', 'scatter'}
        What type of plot should be in the center, a 2D Histogram or a scatter-plot.
    scale_to_cbar : 
    fs : int,
        Font-size
    scale : str or [str, str],
        Specification for the axes scaling {'log','lin'}.  If two values are given, the first is
        used for the x-axis and the second for the y-axis.
    histScale : str,
        Scaling to use for the histograms {'log','lin'}-- the color scale on the 2D histogram,
        or the Counts axis on the 1D histograms.
    labels : (2,) str
    cbar : bool,
        Add a colorbar.
    write_counts : bool
        Print a str on each bin writing the number of values included in that bin.
    left : float {0.0, 1.0}
        Location of the left edge of axes relative to the figure.
    bottom : float {0.0, 1.0}
        Location of the bottom edge of axes relative to the figure.
    right : float {0.0, 1.0}
        Location of the right edge of axes relative to the figure.
    top : float {0.0, 1.0}
        Location of the top edge of axes relative to the figure.
    lo : scalar or 'None'
        When autocalculating `extrema`, ignore histogram entries below this value.
    hi : scalar or 'None'
        When autocalculating `extrema`, ignore histogram entries above this value.

    Returns
    -------
    fig : matplotlib.figure.Figure
        Figure object containing plots.

    """
    # Make sure shapes of input arrays are valid
    if np.shape(xvals) != np.shape(yvals):
        raise ValueError("Shape of `xvals` ({}) must match `yvals` ({}).".format(
            np.shape(xvals), np.shape(yvals)))
    if weights is not None and np.shape(weights) != np.shape(xvals):
        raise ValueError("Shape of `weights` ({}) must match `xvals` and `yvals` ({}).".format(
            np.shape(weights), np.shape(xvals)))

    # Make sure the given `scale` is valid
    if np.size(scale) == 1:
        scale = [scale, scale]
    elif np.size(scale) != 2:
        raise ValueError("`scale` must be one or two scaling specifications!")

    # Check the `labels`
    if labels is None: labels = ['', '', '']
    elif np.size(labels) == 2: labels = [labels[0], labels[1], '']

    if np.size(labels) != 3:
        raise ValueError("`labels` = '{}' is invalid.".format(labels))

    # Make sure scale strings are matplotlib-compliant
    scale = [plot_core._clean_scale(sc) for sc in scale]

    # Determine type of central plot
    if type.startswith('hist'):
        type_hist = True
    elif type.startswith('scat'):
        type_hist = False
        cblabel = str(labels[2])
        labels[2] = 'Count'
    else:
        raise ValueError("`type` = '{}', must be either 'hist', or 'scatter'.".format(type))

    # Infer default statistic
    if statistic is None:
        if weights is None: statistic = 'count'
        else:               statistic = 'sum'

    if filter is None and histScale.startswith('log'):
        filter = 'g'

    # Filter input data
    if filter is not None:
        # Make sure `filter` is an iterable pair
        if weights is None: num = 2
        else:               num = 3

        if not np.iterable(filter): filter = num*[filter]
        elif len(filter) == 1: filter = num*[filter[0]]

        if len(filter) != num:
            raise ValueError("If `weights` are provided, number of `filter` values must match.")

        # Filter `xvals`
        if filter[0] is not None:
            inds = zmath.comparison_filter(xvals, filter[0], inds=True)
            xvals = xvals[inds]
            yvals = yvals[inds]
            if weights is not None:
                weights = weights[inds]
        # Filter `yvals`
        if filter[1] is not None:
            inds = zmath.comparison_filter(yvals, filter[1], inds=True)
            xvals = xvals[inds]
            yvals = yvals[inds]
            if weights is not None:
                weights = weights[inds]

        if weights is not None and filter[2] is not None:
            inds = zmath.comparison_filter(yvals, filter[2], inds=True)
            xvals = xvals[inds]
            yvals = yvals[inds]
            weights = weights[inds]

    # Create and initializae figure and axes
    fig, prax, xpax, ypax, cbax = _constructFigure(fig, xproj, yproj, hratio, wratio, pad,
                                                   scale, histScale, labels, cbar,
                                                   left, bottom, right, top, fs=fs)

    # Create bins
    # -----------
    #     `bins` is a single scalar value -- apply to both
    if np.isscalar(bins):
        xbins = bins
        ybins = bins
    else:
        #     `bins` is a pair of bin specifications, separate and apply
        if len(bins) == 2:
            xbins = bins[0]
            ybins = bins[1]
        #     `bins` is a single array -- apply to both
        elif len(bins) > 2:
            xbins = bins
            ybins = bins
        #     unrecognized option -- error
        else:
            raise ValueError("Unrecognized shape of ``bins`` = %s" % (str(np.shape(bins))))

    # If a number of bins is given, create an appropriate spacing
    if np.ndim(xbins) == 0:
        xbins = zmath.spacing(xvals, num=xbins+1, scale=scale[0])

    if np.ndim(ybins) == 0:
        ybins = zmath.spacing(yvals, num=ybins+1, scale=scale[1])

    # Make sure bins look okay
    for arr, name in zip([xbins, ybins], ['xbins', 'ybins']):
        delta = np.diff(arr)
        if np.any(~np.isfinite(delta) | (delta == 0.0)):
            raise ValueError("Error constructing `{}` = {}, delta = {}".format(name, arr, delta))

    # Calculate Histograms
    # --------------------
    #    2D
    hist_2d, xedges_2d, yedges_2d, binnums = sp.stats.binned_statistic_2d(
        xvals, yvals, weights, statistic=statistic, bins=[xbins, ybins])
    hist_2d = np.nan_to_num(hist_2d)
    #    X-projection (ignore Y)
    hist_xp, edges_xp, bins_xp = sp.stats.binned_statistic(
        xvals, weights, statistic=statistic, bins=xbins)
    #    Y-projection (ignore X)
    hist_yp, edges_yp, bins_yp = sp.stats.binned_statistic(
        yvals, weights, statistic=statistic, bins=ybins)

    # Calculate Extrema - Preserve input extrema if given, otherwise calculate
    extrema = _set_extrema(extrema, [hist_2d, hist_xp, hist_yp], filter=filter[2], lo=lo, hi=hi)
    # Create scalar-mappable if needed
    if smap is None:
        smap = plot_core.colormap(extrema, cmap=cmap, scale=histScale)

    # Plot Histograms and Projections
    # -------------------------------
    # Plot 2D Histogram
    if type_hist:
        counts = None
        # If we should overlay strings labeling the num values in each bin, calculate those `counts`
        if write_counts:
            counts, xedges_2d, yedges_2d, binnums = sp.stats.binned_statistic_2d(
                xvals, yvals, weights, statistic='count', bins=[xbins, ybins])

        pcm, smap = plot2DHist(prax, xedges_2d, yedges_2d, hist_2d, cscale=histScale, cbax=cbax,
                               labels=labels, counts=counts, cmap=cmap, smap=smap, extrema=extrema,
                               fs=fs)

        # Colors
        # X-projection
        colhist_xp = np.array(hist_xp)
        # Enforce positive values for colors in log-plots.
        if smap.log:
            tmin, tmax = zmath.minmax(colhist_xp, filter='g')
            colhist_xp = np.maximum(colhist_xp, tmin)
        colors_xp = smap.to_rgba(colhist_xp)
        colors_yp = smap.to_rgba(hist_yp)

    # Scatter Plot
    else:
        colors = smap.to_rgba(weights)
        prax.scatter(xvals, yvals, c=colors, alpha=alpha)

        cbar = plt.colorbar(smap, cax=cbax)
        cbar.set_label(cblabel, fontsize=fs)
        cbar.ax.tick_params(labelsize=fs)

        # Make projection colors all grey
        colors_xp = '0.8'
        colors_yp = '0.8'

    # Plot projection of the x-axis (i.e. ignore 'y')
    if xpax:
        islog = scale[0].startswith('log')

        xpax.bar(edges_xp[:-1], hist_xp, color=colors_xp, log=islog, width=np.diff(edges_xp))
        # set tick-labels to the top
        plt.setp(xpax.get_yticklabels(), fontsize=fs)
        xpax.xaxis.tick_top()
        # set bounds to bin edges
        plot_core.setLim(xpax, 'x', data=xedges_2d)
        # Set axes limits to match those of colorbar
        if scale_to_cbar:
            # extrema_y = [zmath.floor_log(extrema[0]), zmath.ceil_log(extrema[1])]
            extrema_y = zmath.minmax(extrema, round=0)
            xpax.set_ylim(extrema_y)

    # Plot projection of the y-axis (i.e. ignore 'x')
    if ypax:
        islog = plot_core._scale_to_log_flag(histScale)
        ypax.barh(edges_yp[:-1], hist_yp, color=colors_yp, log=islog, height=np.diff(edges_yp))
        # set tick-labels to the top
        plt.setp(ypax.get_yticklabels(), fontsize=fs, rotation=90)
        ypax.yaxis.tick_right()
        # set bounds to bin edges
        plot_core.setLim(ypax, 'y', data=yedges_2d)
        try:
            ypax.locator_params(axis='x', tight=True, nbins=4)
        except:
            ypax.locator_params(axis='x', tight=True)

        # Set axes limits to match those of colorbar
        if scale_to_cbar:
            extrema_x = zmath.minmax(extrema, round=0)
            ypax.set_xlim(extrema_x)

    prax.set(xlim=zmath.minmax(xedges_2d), ylim=zmath.minmax(yedges_2d))

    return fig


def plot2DHist(ax, xvals, yvals, hist,
               cax=None, cbax=None, cscale='log', cmap=None, smap=None, extrema=None,
               contours=None, clabel={}, fs=12, rasterized=True,
               title=None, labels=None, counts=None, **kwargs):
    """Plot the given 2D histogram of data.

    Use with (e.g.) ``numpy.histogram2d``,

    Arguments
    ---------
    ax : ``matplotlib.axes.Axes`` object
        Axes object on which to plot.
    xvals : (N,) array of scalars
        Positions (edges) of x values, assumed to be the same for all rows of
        the input data `hist`.<commit_msg>ENH: improve projected axes bounds for log-scales (round to nearest order of magnitude).<commit_after>"""Methods for plotting 2D histograms with optional 1D projected histograms for each axis.

Functions
---------
-    plot2DHistProj          - Plot a 2D histogram with projections of one or both axes.
-    plot2DHist              - Plot the given 2D histogram of data.

-    _constructFigure        - Add the required axes to the given figure object.

To-do
-----
-   `plot2DHistProj` and `plot2DHist` use different styles of input data, specifically
    `plot2DHist` requires a 2D histogram of data to be passed, while `plot2DHistProj` accepts
    separate arrays of x and y data to be histogrammed using ``scipy.stats.binned_statistic_2d``.

"""
from __future__ import absolute_import, division, print_function, unicode_literals
from six.moves import xrange

import numpy as np
import scipy as sp
import scipy.stats
import matplotlib.pyplot as plt

import zcode.math as zmath
from . import plot_core

__all__ = ['plot2DHist', 'plot2DHistProj']

_LEFT = 0.09
_RIGHT = 0.92     # Location of right of plots
_BOTTOM = 0.09
_TOP = 0.90       # Location of top of plots

_CB_WID = 0.02
_CB_WPAD = 0.1
_BAR_ALPHA = 0.8


def plot2DHistProj(xvals, yvals, weights=None, statistic=None, bins=10, filter=None, extrema=None,
                   fig=None, xproj=True, yproj=True, hratio=0.7, wratio=0.7, pad=0.0, alpha=1.0,
                   cmap=None, smap=None, type='hist', scale_to_cbar=True,
                   fs=12, scale='log', histScale='log', labels=None, cbar=True, write_counts=False,
                   left=_LEFT, bottom=_BOTTOM, right=_RIGHT, top=_TOP, lo=None, hi=None):
    """Plot a 2D histogram with projections of one or both axes.

    Arguments
    ---------
    xvals : (N,) array_like,
        Values corresponding to the x-points of the given data
    yvals : (N,) array_like,
        Values corresponding to the y-points of the given data
    weights : (N,) array_like or `None`,
        Weights used to create histograms.  If `None`, then counts are used.
    statistic : str or `None`,
        Type of statistic to be calculated, passed to ``scipy.stats.binned_statistic``.
        e.g. {'count', 'sum', 'mean'}.
        If `None`, then either 'sum' or 'count' is used depending on if `weights` are
        provieded or not.
    bins : int or [int, int] or array_like or [array, array],
        Specification for bin sizes.  integer values are treated as the number of bins to use,
        while arrays are used as the bin edges themselves.  If a tuple of two values is given, it
        is assumed that the first is for the x-axis and the second for the y-axis.
    filter : str or `None`, or [2,] tuple of str or `None`, or [3,] tubple of str or `None`
        String specifying how to filter the input `data` relative to zero.
        If this is a single value, it is applies to both `xvals` and `yvals`.
        If this is a tuple/list of two values, they correspond to `xvals` and `yvals` respectively.
        If `weights` are provided, the tuple/list should have three values.
    extrema :
    fig : ``matplotlib.figure.figure``,
        Figure instance to which axes are added for plotting.  One is created if not given.
    xproj : bool,
        Whether to also plot the projection of the x-axis (i.e. histogram ignoring y-values).
    yproj : bool,
        Whether to also plot the projection of the y-axis (i.e. histogram ignoring x-values).
    hratio : float,
        Fraction of the total available height-space to use for the primary axes object (2D hist)
    wratio : float,
        Fraction of the total available width-space to use for the primary axes object (2D hist)
    pad : float,
        Padding between central axis and the projected ones.
    cmap : ``matplotlib.colors.Colormap`` object
        Matplotlib colormap to use for coloring histogram.
        Overridden if `smap` is provided.
    smap : `matplotlib.cm.ScalarMappable` object or `None`
        A scalar-mappable object to use for colormaps, or `None` for one to be created.
    type : str, {'hist', 'scatter'}
        What type of plot should be in the center, a 2D Histogram or a scatter-plot.
    scale_to_cbar : 
    fs : int,
        Font-size
    scale : str or [str, str],
        Specification for the axes scaling {'log','lin'}.  If two values are given, the first is
        used for the x-axis and the second for the y-axis.
    histScale : str,
        Scaling to use for the histograms {'log','lin'}-- the color scale on the 2D histogram,
        or the Counts axis on the 1D histograms.
    labels : (2,) str
    cbar : bool,
        Add a colorbar.
    write_counts : bool
        Print a str on each bin writing the number of values included in that bin.
    left : float {0.0, 1.0}
        Location of the left edge of axes relative to the figure.
    bottom : float {0.0, 1.0}
        Location of the bottom edge of axes relative to the figure.
    right : float {0.0, 1.0}
        Location of the right edge of axes relative to the figure.
    top : float {0.0, 1.0}
        Location of the top edge of axes relative to the figure.
    lo : scalar or 'None'
        When autocalculating `extrema`, ignore histogram entries below this value.
    hi : scalar or 'None'
        When autocalculating `extrema`, ignore histogram entries above this value.

    Returns
    -------
    fig : matplotlib.figure.Figure
        Figure object containing plots.

    """
    # Make sure shapes of input arrays are valid
    if np.shape(xvals) != np.shape(yvals):
        raise ValueError("Shape of `xvals` ({}) must match `yvals` ({}).".format(
            np.shape(xvals), np.shape(yvals)))
    if weights is not None and np.shape(weights) != np.shape(xvals):
        raise ValueError("Shape of `weights` ({}) must match `xvals` and `yvals` ({}).".format(
            np.shape(weights), np.shape(xvals)))

    # Make sure the given `scale` is valid
    if np.size(scale) == 1:
        scale = [scale, scale]
    elif np.size(scale) != 2:
        raise ValueError("`scale` must be one or two scaling specifications!")

    # Check the `labels`
    if labels is None: labels = ['', '', '']
    elif np.size(labels) == 2: labels = [labels[0], labels[1], '']

    if np.size(labels) != 3:
        raise ValueError("`labels` = '{}' is invalid.".format(labels))

    # Make sure scale strings are matplotlib-compliant
    scale = [plot_core._clean_scale(sc) for sc in scale]

    # Determine type of central plot
    if type.startswith('hist'):
        type_hist = True
    elif type.startswith('scat'):
        type_hist = False
        cblabel = str(labels[2])
        labels[2] = 'Count'
    else:
        raise ValueError("`type` = '{}', must be either 'hist', or 'scatter'.".format(type))

    # Infer default statistic
    if statistic is None:
        if weights is None: statistic = 'count'
        else:               statistic = 'sum'

    if filter is None and histScale.startswith('log'):
        filter = 'g'

    # Filter input data
    if filter is not None:
        # Make sure `filter` is an iterable pair
        if weights is None: num = 2
        else:               num = 3

        if not np.iterable(filter): filter = num*[filter]
        elif len(filter) == 1: filter = num*[filter[0]]

        if len(filter) != num:
            raise ValueError("If `weights` are provided, number of `filter` values must match.")

        # Filter `xvals`
        if filter[0] is not None:
            inds = zmath.comparison_filter(xvals, filter[0], inds=True)
            xvals = xvals[inds]
            yvals = yvals[inds]
            if weights is not None:
                weights = weights[inds]
        # Filter `yvals`
        if filter[1] is not None:
            inds = zmath.comparison_filter(yvals, filter[1], inds=True)
            xvals = xvals[inds]
            yvals = yvals[inds]
            if weights is not None:
                weights = weights[inds]

        if weights is not None and filter[2] is not None:
            inds = zmath.comparison_filter(yvals, filter[2], inds=True)
            xvals = xvals[inds]
            yvals = yvals[inds]
            weights = weights[inds]

    # Create and initializae figure and axes
    fig, prax, xpax, ypax, cbax = _constructFigure(fig, xproj, yproj, hratio, wratio, pad,
                                                   scale, histScale, labels, cbar,
                                                   left, bottom, right, top, fs=fs)

    # Create bins
    # -----------
    #     `bins` is a single scalar value -- apply to both
    if np.isscalar(bins):
        xbins = bins
        ybins = bins
    else:
        #     `bins` is a pair of bin specifications, separate and apply
        if len(bins) == 2:
            xbins = bins[0]
            ybins = bins[1]
        #     `bins` is a single array -- apply to both
        elif len(bins) > 2:
            xbins = bins
            ybins = bins
        #     unrecognized option -- error
        else:
            raise ValueError("Unrecognized shape of ``bins`` = %s" % (str(np.shape(bins))))

    # If a number of bins is given, create an appropriate spacing
    if np.ndim(xbins) == 0:
        xbins = zmath.spacing(xvals, num=xbins+1, scale=scale[0])

    if np.ndim(ybins) == 0:
        ybins = zmath.spacing(yvals, num=ybins+1, scale=scale[1])

    # Make sure bins look okay
    for arr, name in zip([xbins, ybins], ['xbins', 'ybins']):
        delta = np.diff(arr)
        if np.any(~np.isfinite(delta) | (delta == 0.0)):
            raise ValueError("Error constructing `{}` = {}, delta = {}".format(name, arr, delta))

    # Calculate Histograms
    # --------------------
    #    2D
    hist_2d, xedges_2d, yedges_2d, binnums = sp.stats.binned_statistic_2d(
        xvals, yvals, weights, statistic=statistic, bins=[xbins, ybins])
    hist_2d = np.nan_to_num(hist_2d)
    #    X-projection (ignore Y)
    hist_xp, edges_xp, bins_xp = sp.stats.binned_statistic(
        xvals, weights, statistic=statistic, bins=xbins)
    #    Y-projection (ignore X)
    hist_yp, edges_yp, bins_yp = sp.stats.binned_statistic(
        yvals, weights, statistic=statistic, bins=ybins)

    # Calculate Extrema - Preserve input extrema if given, otherwise calculate
    extrema = _set_extrema(extrema, [hist_2d, hist_xp, hist_yp], filter=filter[2], lo=lo, hi=hi)
    # Create scalar-mappable if needed
    if smap is None:
        smap = plot_core.colormap(extrema, cmap=cmap, scale=histScale)

    # Plot Histograms and Projections
    # -------------------------------
    # Plot 2D Histogram
    if type_hist:
        counts = None
        # If we should overlay strings labeling the num values in each bin, calculate those `counts`
        if write_counts:
            counts, xedges_2d, yedges_2d, binnums = sp.stats.binned_statistic_2d(
                xvals, yvals, weights, statistic='count', bins=[xbins, ybins])

        pcm, smap = plot2DHist(prax, xedges_2d, yedges_2d, hist_2d, cscale=histScale, cbax=cbax,
                               labels=labels, counts=counts, cmap=cmap, smap=smap, extrema=extrema,
                               fs=fs)

        # Colors
        # X-projection
        colhist_xp = np.array(hist_xp)
        # Enforce positive values for colors in log-plots.
        if smap.log:
            tmin, tmax = zmath.minmax(colhist_xp, filter='g')
            colhist_xp = np.maximum(colhist_xp, tmin)
        colors_xp = smap.to_rgba(colhist_xp)
        colors_yp = smap.to_rgba(hist_yp)

    # Scatter Plot
    else:
        colors = smap.to_rgba(weights)
        prax.scatter(xvals, yvals, c=colors, alpha=alpha)

        cbar = plt.colorbar(smap, cax=cbax)
        cbar.set_label(cblabel, fontsize=fs)
        cbar.ax.tick_params(labelsize=fs)

        # Make projection colors all grey
        colors_xp = '0.8'
        colors_yp = '0.8'

    hist_log = plot_core._scale_to_log_flag(histScale)

    # Plot projection of the x-axis (i.e. ignore 'y')
    if xpax:
        islog = scale[0].startswith('log')

        xpax.bar(edges_xp[:-1], hist_xp, color=colors_xp, log=islog, width=np.diff(edges_xp),
                 alpha=_BAR_ALPHA)
        # set tick-labels to the top
        plt.setp(xpax.get_yticklabels(), fontsize=fs)
        xpax.xaxis.tick_top()
        # set bounds to bin edges
        plot_core.setLim(xpax, 'x', data=xedges_2d)
        # Set axes limits to match those of colorbar
        if scale_to_cbar:
            # extrema_y = [zmath.floor_log(extrema[0]), zmath.ceil_log(extrema[1])]
            round = 0
            if hist_log: round = -1
            extrema_y = zmath.minmax(extrema, round=round)
            xpax.set_ylim(extrema_y)

    # Plot projection of the y-axis (i.e. ignore 'x')
    if ypax:
        ypax.barh(edges_yp[:-1], hist_yp, color=colors_yp, log=hist_log, height=np.diff(edges_yp),
                  alpha=_BAR_ALPHA)
        # set tick-labels to the top
        plt.setp(ypax.get_yticklabels(), fontsize=fs, rotation=90)
        ypax.yaxis.tick_right()
        # set bounds to bin edges
        plot_core.setLim(ypax, 'y', data=yedges_2d)
        try:
            ypax.locator_params(axis='x', tight=True, nbins=4)
        except:
            ypax.locator_params(axis='x', tight=True)

        # Set axes limits to match those of colorbar
        if scale_to_cbar:
            round = 0
            if hist_log: round = -1
            extrema_x = zmath.minmax(extrema, round=round)
            ypax.set_xlim(extrema_x)

    prax.set(xlim=zmath.minmax(xedges_2d), ylim=zmath.minmax(yedges_2d))

    return fig


def plot2DHist(ax, xvals, yvals, hist,
               cax=None, cbax=None, cscale='log', cmap=None, smap=None, extrema=None,
               contours=None, clabel={}, fs=12, rasterized=True,
               title=None, labels=None, counts=None, **kwargs):
    """Plot the given 2D histogram of data.

    Use with (e.g.) ``numpy.histogram2d``,

    Arguments
    ---------
    ax : ``matplotlib.axes.Axes`` object
        Axes object on which to plot.
    xvals : (N,) array of scalars
        Positions (edges) of x values, assumed to be the same for all rows of
        the input data `hist`.
{'fraction_non_alphanumeric': 0.10734322846388751, 'fraction_numerical': 0.010013493359846602, 'mean_word_length': 2.30823446493598, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1738054', 'n_tokens_mistral': 9258, 'n_tokens_neox': 8802, 'n_words': 2905}