<commit_before>#! /usr/bin/env python
# -*- coding: utf-8 -*-

"""
@author: <NAME>
"""

from __future__ import print_function, division
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.figure import figaspect
from matplotlib.font_manager import FontProperties
import matplotlib.gridspec as matgrid
import re


def initmpl():
    """
    Initialize mpes plot style
    """

    mpl.rcParams['text.usetex'] = True
    mpl.rcParams['axes.linewidth'] = 2
    mpl.rcParams['axes.labelsize'] = 20
    mpl.rcParams['text.latex.preamble'] = [r'\usepackage{helvet}']
    mpl.rcParams['axes.edgecolor'] = 'black'
    mpl.rcParams['savefig.dpi'] = 100
    mpl.rcParams['xtick.major.pad'] = 5
    mpl.rcParams['ytick.major.pad'] = 5
    mpl.rcParams['xtick.major.width'] = 2
    mpl.rcParams['xtick.minor.width'] = 2
    mpl.rcParams['xtick.major.size'] = 8
    mpl.rcParams['xtick.minor.size'] = 6
    mpl.rcParams['ytick.major.width'] = 2
    mpl.rcParams['ytick.minor.width'] = 2
    mpl.rcParams['ytick.major.size'] = 8
    mpl.rcParams['ytick.minor.size'] = 6
    mpl.rcParams['xtick.labelsize'] = 15
    mpl.rcParams['ytick.labelsize'] = 15


#==========#
# 1D plots #
#==========#

def stackedlineplot(datamat, axis=0, interval=0, binning=None, **kwds):
    """
    Stacked line plots (used for visualizing energy or momentum dispersion curves)
    
    **Parameters**
    
    data : numeric 2D array
        the 2D data to plot
    axis : int
        the axis to cut along
    interval : float
        the interval between plots
    binning : int
        number of binned rows/columns
    **kwds : keyword arguments
        =============  ==========  ===================================
        keyword        data type   meaning
        =============  ==========  ===================================
        figsize        tuple/list  (horizontal_size, vertical_size)
        xlabel         str         x axis label
        ylabel         str         y axis label
        axislabelsize  int         font size of axis text labels
        ticklabelsize  int         font size of axis tick labels
        =============  ==========  ===================================
    **Return**
    
    ax : axes object
        handle for the plot axes
    """
    
    figuresize = kwds.pop('figsize', '')
    try:
        fw, fh = numFormatConversion(figuresize)
    except:
        fw, fh = 2 * figaspect(data)
    f, ax = plt.subplots(figsize=(fw, fh))
    
    datamat = np.rollaxis(np.asarray(datamat), axis)
    
    nr, nc = datamat.shape
    x = kwds.pop('x', range(nc))
    y = kwds.pop('y', range(nr))
    
    for i in y:
        
        yval = datamat[i,:]
        #print(list(x))
        ax.plot(x, yval + i*interval, linewidth=2)
        
    xlabel = kwds.pop('xlabel', '')
    ylabel = kwds.pop('ylabel', '')
    
    axislabelsize = kwds.pop('ax_labelsize', 12)
    ticklabelsize = kwds.pop('tk_labelsize', 10)
    
    ax.set_xlabel(xlabel, fontsize=axislabelsize)
    ax.set_ylabel(ylabel, fontsize=axislabelsize)
    ax.tick_params(labelsize=ticklabelsize)
    plt.tight_layout()
    
    return ax


#==========#
# 2D plots #
#==========#

def numFormatConversion(seq, form='int', **kwds):
    """
    When length keyword is not specified as an argument, the function
    returns a format-converted sequence of numbers
    <commit_msg>Extended stackedlineplot()<commit_after>#! /usr/bin/env python
# -*- coding: utf-8 -*-

"""
@author: <NAME>
"""

from __future__ import print_function, division
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.figure import figaspect
from matplotlib.font_manager import FontProperties
import matplotlib.gridspec as matgrid
import re


def initmpl():
    """
    Initialize mpes plot style
    """

    mpl.rcParams['text.usetex'] = True
    mpl.rcParams['axes.linewidth'] = 2
    mpl.rcParams['axes.labelsize'] = 20
    mpl.rcParams['text.latex.preamble'] = [r'\usepackage{helvet}']
    mpl.rcParams['axes.edgecolor'] = 'black'
    mpl.rcParams['savefig.dpi'] = 100
    mpl.rcParams['xtick.major.pad'] = 5
    mpl.rcParams['ytick.major.pad'] = 5
    mpl.rcParams['xtick.major.width'] = 2
    mpl.rcParams['xtick.minor.width'] = 2
    mpl.rcParams['xtick.major.size'] = 8
    mpl.rcParams['xtick.minor.size'] = 6
    mpl.rcParams['ytick.major.width'] = 2
    mpl.rcParams['ytick.minor.width'] = 2
    mpl.rcParams['ytick.major.size'] = 8
    mpl.rcParams['ytick.minor.size'] = 6
    mpl.rcParams['xtick.labelsize'] = 15
    mpl.rcParams['ytick.labelsize'] = 15


#==========#
# 1D plots #
#==========#

def stackedlineplot(datamat, axis=0, interval=0, binning=1, **kwds):
    """
    Stacked line plots (used for visualizing energy or momentum dispersion curves)
    
    **Parameters**
    
    data : numeric 2D array
        the 2D data to plot
    axis : int | 0
        the axis to cut along
    interval : float | 0
        the interval between plots
    binning : int | 1 (no binning)
        number of binned rows/columns
    **kwds : keyword arguments
        =============  ==========  ===================================
        keyword        data type   meaning
        =============  ==========  ===================================
        figsize        tuple/list  (horizontal_size, vertical_size)
        x              1D array    x axis values
        xlabel         str         x axis label
        ylabel         str         y axis label
        cmap           str         `matplotlib colormap string <https://matplotlib.org/users/colormaps.html>`_
        axislabelsize  int         font size of axis text labels
        ticklabelsize  int         font size of axis tick labels
        =============  ==========  ===================================
    **Return**
    
    ax : axes object
        handle for the plot axes
    """
    
    # Check binning input
    binning = round(binning)
    if binning < 1:
        binning = 1
    
    # Determine figure size
    figuresize = kwds.pop('figsize', '')
    try:
        fw, fh = numFormatConversion(figuresize)
    except:
        fw, fh = 2 * figaspect(datamat)
    f, ax = plt.subplots(figsize=(fw, fh))
    
    datamat = np.rollaxis(np.asarray(datamat), axis)
    nr, nc = datamat.shape
    x = kwds.pop('x', range(0, nc))
    y = range(0, nr, binning)
    ny = len(y)
    colormap = kwds.pop('cmap', '')
    lw = kwds.pop('linewidth', 2)
    
    # Plot lines
    for ind, i in enumerate(y):
        
        yval = np.mean(datamat[i:i+binning,:], axis=0)
        line = ax.plot(x, yval + i*interval, linewidth=lw)
        # Set line color
        if colormap:
            line[0].set_color(eval('plt.cm.' + colormap + '(ind/ny)'))
        
    xlabel = kwds.pop('xlabel', '')
    ylabel = kwds.pop('ylabel', '')
    
    axislabelsize = kwds.pop('ax_labelsize', 12)
    ticklabelsize = kwds.pop('tk_labelsize', 10)
    
    ax.set_xlabel(xlabel, fontsize=axislabelsize)
    ax.set_ylabel(ylabel, fontsize=axislabelsize)
    ax.tick_params(labelsize=ticklabelsize)
    plt.tight_layout()
    
    return ax


#==========#
# 2D plots #
#==========#

def numFormatConversion(seq, form='int', **kwds):
    """
    When length keyword is not specified as an argument, the function
    returns a format-converted sequence of numbers
    
{'fraction_non_alphanumeric': 0.17157263559553013, 'fraction_numerical': 0.010493322431180159, 'mean_word_length': 2.6081612586037366, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 6, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6138105', 'n_tokens_mistral': 2526, 'n_tokens_neox': 2334, 'n_words': 558}