<commit_before>#!/usr/bin/env python
import os
import sys
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import argparse

def get_args():
    """Gets the command line arguments"""
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument(
        "results_file",
        help="The input results spreadsheet")
    parser.add_argument(
        "--worksheet_name",
        type=str,
        default="results",
        help="Worksheet containing the results to plot")
    parser.add_argument(
        "--compute_element_name",
        type=str,
        default="Threads",
        help="Compute element name that will be used on plot labels")
    parser.add_argument(
        "--walltime_units",
        default="Minutes",
        type=str,
        help="Walltime units name")
    parser.add_argument(
        "--filter_column",
        default="",
        type=str,
        help="""Optional filter column name. If given, the column contents will
        be equated to a boolean value. Any result group with a member filter
        that equates to False will be excluded from the plot.""")
    parser.add_argument(
        "--title_prefix",
        type=str,
        help="optional prefix that will be prepended to the standard plot titles")
    parser.add_argument(
        "--file_prefix",
        type=str,
        help="optional prefix that will be prepended to the standard file names")
    parser.add_argument(
        "--window",
        default=False,
        help="prints plots to a window instead of to files",
        action="store_true")
    parser.add_argument(
        "--plot_width",
        default=10,
        type=int,
        help="Plot width in inches")
    parser.add_argument(
        "--speedup_max",
        default=16.1,
        type=float,
        help="Max Y axis scale for speedup")

    return parser.parse_args()

def save(path, ext='png', close=True, verbose=True):
    """Save a figure from pyplot.

    Parameters
    ----------
    path : string
        The path (and filename, without the extension) to save the
        figure to.

    ext : string (default='png')
        The file extension. This must be supported by the active
        matplotlib backend (see matplotlib.backends module).  Most
        backends support 'png', 'pdf', 'ps', 'eps', and 'svg'.

    close : boolean (default=True)
        Whether to close the figure after saving.  If you want to save
        the figure multiple times (e.g., to multiple formats), you
        should NOT close it in between saves or you will have to
        re-plot it.

    verbose : boolean (default=True)
        Whether to print information about when and where the image
        has been saved.

    """

    # Extract the directory and filename from the given path
    directory = os.path.split(path)[0]
    filename = "%s.%s" % (os.path.split(path)[1], ext)
    if directory == '':
        directory = '.'

    # If the directory does not exist, create it
    if not os.path.exists(directory):
        os.makedirs(directory)

    # The final path to save to
    savepath = os.path.join(directory, filename)

    if verbose:
        print("Saving figure to '%s'..." % savepath),

    # Actually save the figure
    plt.savefig(savepath)

    # Close it
    if close:
        plt.close()

    if verbose:
        print("Done")

def read_dataframe_from_excel(
        filename,
        worksheet,
        **kwargs):
    """Reads a pandas dataframe from an Excel file
    filename: the excel file
    worksheet: worksheet name within the excel file
    The remaining named arguments are passed to pandas.read_excel
    """
    results = pd.read_excel(filename, worksheet, **kwargs)
    return results

def calculate_speedup_and_efficiency(
        rdf,
        compute_element_col_index,
        time_col_index):
    """Calculates the speedup and efficiency and
    adds them as new columns to a dataframe

    rdf: the results dataframe
    compute_element_col_index: compute elements column name
    time_col_index: column name containing the computation times

    Returns: a tuple containing the speedup and efficiency series
    """
    # get the t1 reference value
    t1 = float(rdf[rdf[compute_element_col_index] == 1][time_col_index])
    # speedup n = t1 / tn
    speedup = t1 / rdf[time_col_index]
    # efficiency n = t1 / (tn * n)
    efficiency = t1 / (rdf[time_col_index] * rdf[compute_element_col_index])

    return (speedup, efficiency)

def plot_walltime(
        series,
        colours,
        series_names,
        compute_elements,
        ymax,
        plot_size=None,
        group_width=0.8,
        xlabel='Processes',
        ylabel='Minutes',
        title='Walltime',
        file_name='walltime',
        file_extension='png',
        y_log_scale=False,
        show=False):
    """creates a bar plot as a new pyplot figure"""

    # the x locations for the groups.
    x_ind = np.arange(len(compute_elements))

    # create the figure and axes
    fig = plt.figure(figsize=plot_size)
    ax = fig.add_subplot(111)

    # create the individual bars
    bars_per_group = len(series)
    bar_width = group_width / bars_per_group
    bar_left = x_ind - (group_width / 2)  # centre the group on the ticks
    plot_left = x_ind[0] - group_width
    plot_right = x_ind[-1] + group_width

    for bar_values, bar_colour in zip(series, colours):
        ax.bar(
            bar_left,
            bar_values,
            width=bar_width,
            color=bar_colour,
            log=y_log_scale)
        bar_left += bar_width

    # apply the labels and formatting
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.set_xlim(plot_left, plot_right)
    ax.set_ylim(0, ymax)
    ax.set_xticks(x_ind)
    ax.set_xticklabels(compute_elements)
    ax.legend(series_names)
    # ax.legend().get_frame().set_facecolor('white')

    if show:
        plt.show()
    else:
        save(file_name, file_extension, close=True, verbose=True)

def plot_speedup(
        series,
        colours,
        series_names,
        compute_elements,
        xmax,
        ymax,
        line_width=1,
        plot_size=None,
        xlabel='Processes',
        ylabel='Speedup',
        title='Speedup',
        file_name='speedup',
        file_extension='png',
        show=False):
    """creates a speedup plot"""

    # create the figure and axes
    fig = plt.figure(figsize=plot_size)
    ax = fig.add_subplot(111)

    # define the sizes and locations of things
    x_ticks = list(compute_elements)
    x_ticks.append(xmax)
    plot_left = x_ticks[0]
    plot_right = x_ticks[-1]
    ax.set_xlim(plot_left, plot_right)
    ax.set_ylim(plot_left, ymax)

    # create the ideal speedup reference line
    ax.plot(
        x_ticks,
        x_ticks,
        color='red',
        linewidth=line_width,
        label='_nolegend_')

    # plot each series
    for values, colour, name in zip(series, colours, series_names):
        ax.plot(
            compute_elements,
            values,
            label=name,
            color=colour,
            linewidth=line_width,
            marker='o')

    # apply the labels and formatting
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.set_xticks(compute_elements)
    ax.set_xticklabels(compute_elements)
    ax.legend()
    ax.legend().get_frame().set_facecolor('white')
    # can also use get_frame().set_xy(x,y) to give a custom location

    if show:
        plt.show()
    else:
        save(file_name, file_extension, close=True, verbose=True)

def plot_efficiency(
        series,
        colours,
        series_names,
        compute_elements,
        xmax,
        ymax=1.2,
        line_width=1,
        plot_size=None,
        xlabel='Processes',
        ylabel='Efficiency',
        title='Efficiency',
        file_name='efficiency',
        file_extension='png',
        show=False):
    """creates an efficiency plot"""

    # create the figure and axes
    fig = plt.figure(figsize=plot_size)
    ax = fig.add_subplot(111)

    # define the sizes and locations of things
    x_ticks = list(compute_elements)
    x_ticks.append(xmax)
    plot_left = x_ticks[0]
    plot_right = x_ticks[-1]
    ax.set_xlim(plot_left, plot_right)
    ax.set_ylim(0, ymax)

    # create a horizontal ideal efficiency reference line
    ax.hlines(
        1,
        plot_left,
        plot_right,
        colors='red',
        linestyles='solid',
        linewidth=line_width,
        label='_nolegend_')

    # plot each series
    for values, colour, name in zip(series, colours, series_names):
        ax.plot(
            compute_elements,
            values,
            label=name,
            color=colour,
            linewidth=line_width,
            marker='o')

    # apply the labels and formatting
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.set_xticks(compute_elements)
    ax.set_xticklabels(compute_elements)
    ax.legend(shadow=True)
    ax.legend().get_frame().set_facecolor('white')

    if show:
        plt.show()
    else:
        save(file_name, file_extension, close=True, verbose=True)

def add_optional_prefix(base_string, prefix, separator):
    result = str(base_string)
    if prefix:
        result = prefix + separator + base_string
    return result

if __name__ == '__main__':

    # get and process the arguments
    args = get_args()

    # create plots in a 4:3 aspect ratio
    # matplotlib works in inches
    plot_width = args.plot_width
    plot_size = (plot_width, plot_width * 3 / 4)

    show_instead_of_save = args.window

    walltime_title = add_optional_prefix('Walltime', args.title_prefix, ' - ')
    speedup_title = add_optional_prefix('Speedup', args.title_prefix, ' - ')
    efficiency_title = add_optional_prefix(
        'Strong Scaling Efficiency',
        args.title_prefix,
        ' - ')

    walltime_file = add_optional_prefix('walltime', args.file_prefix, '-')
    speedup_file = add_optional_prefix('speedup', args.file_prefix, '-')
    efficiency_file = add_optional_prefix('efficiency', args.file_prefix, '-')

    compute_element_name = args.compute_element_name
    walltime_units = args.walltime_units

    # CSIRO colours
    colours = [
        '#00a9ce',  # midday blue
        '#78be20',  # light forest
        '#DF1995',  # fuschia
        '#E87722',  # orange
        '#E4002B',  # vermillion
        '#00616c',  # midnight blue
        '#FFB81C',  # gold
        '#6D2077',  # plum
        '#1E22AA',  # blueberry
    ]

    usecols = [
        'group',
        'compute_elements',
        'walltime']
    if args.filter_column:
        usecols.append(args.filter_column)

    results = read_dataframe_from_excel(
        args.results_file,
        worksheet=args.worksheet_name,
        usecols=usecols)

    # filter out incomplete results
    results = results[results.group.notnull() & (results.walltime > 0)]

    # apply the optional filter column
    if args.filter_column:
        results = results[results[args.filter_column] > 0]

    # if there are multiple times for each (group,compute_element) tuple,
    # calculate the mean
    results = results.groupby(['group','compute_elements']).mean().reset_index()

    # create new dataframes for each group of results
    groups = results.group.unique()
    group_dataframes = {}
    for group in groups:
        df = results[(results.group == group)]
        if len(df):
            group_dataframes[group] = df.sort(
                'compute_elements',
                ascending=True)

    # calculate speedup and efficiency based on the total walltime, and on
    # the average time per iteration (which excludes a lot of the serial
    # code, and gives a closer look at just the model calculations)
    for data in group_dataframes.values():
        speedup, efficiency = calculate_speedup_and_efficiency(
            data,
            'compute_elements',
            'walltime')
        data['speedup'] = speedup
        data['efficiency'] = efficiency

    # extract the compute element names for grouping and labelling the charts
    compute_elements = np.sort(results.compute_elements.unique())

    # build the collections of values for the charts
    walltimes = []
    efficiencies = []
    speedups = []
    series_names = []
    max_walltime = 0
    for name, data in group_dataframes.items():
        walltimes.append(data.walltime)
        efficiencies.append(data.efficiency)
        speedups.append(data.speedup)
        series_names.append(name)
        local_max_walltime = max(data.walltime)
        if local_max_walltime > max_walltime:
            max_walltime = local_max_walltime

    # finally make the plots

    plot_walltime(
        walltimes,
        colours,
        series_names,
        compute_elements,
        ymax=max_walltime * 1.2,
        group_width=0.83,
        plot_size=plot_size,
        title=walltime_title,
        xlabel=compute_element_name,
        ylabel=walltime_units,
        show=show_instead_of_save,
        file_name=walltime_file,
        file_extension='png')

    plot_efficiency(
        efficiencies,
        colours,
        series_names,
        compute_elements,
        line_width=1.5,
        xmax=compute_elements.max() * 1.1,
        ymax=1.3,
        plot_size=plot_size,
        xlabel=compute_element_name,
        title=efficiency_title,
        show=show_instead_of_save,
        file_name=efficiency_file,
        file_extension='png')

    plot_speedup(
        speedups,
        colours,
        series_names,
        compute_elements,
        line_width=1.5,
        xmax=compute_elements.max() * 1.05,
        ymax=args.speedup_max,
        plot_size=plot_size,
        xlabel=compute_element_name,
        title=speedup_title,
        show=show_instead_of_save,
        file_name=speedup_file,
        file_extension='png')
<commit_msg>adding weak scaling support<commit_after>#!/usr/bin/env python
import os
import sys
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import argparse

def get_args():
    """Gets the command line arguments"""
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument(
        "results_file",
        help="The input results spreadsheet")
    parser.add_argument(
        "--worksheet_name",
        type=str,
        default="results",
        help="Worksheet containing the results to plot")
    parser.add_argument(
        "--compute_element_name",
        type=str,
        default="Threads",
        help="Compute element name that will be used on plot labels")
    parser.add_argument(
        "--walltime_units",
        default="Minutes",
        type=str,
        help="Walltime units name")
    parser.add_argument(
        "--filter_column",
        default="",
        type=str,
        help="""Optional filter column name. If given, the column contents will
        be equated to a boolean value. Any result group with a member filter
        that equates to False will be excluded from the plot.""")
    parser.add_argument(
        "--title_prefix",
        type=str,
        help="optional prefix that will be prepended to the standard plot titles")
    parser.add_argument(
        "--file_prefix",
        type=str,
        help="optional prefix that will be prepended to the standard file names")
    parser.add_argument(
        "--window",
        default=False,
        help="prints plots to a window instead of to files",
        action="store_true")
    parser.add_argument(
        "--plot_width",
        default=10,
        type=int,
        help="Plot width in inches")
    parser.add_argument(
        "--speedup_max",
        default=16.1,
        type=float,
        help="Max Y axis scale for speedup")
    parser.add_argument(
        "--weak",
        default=False,
        help="Generate plots for weak scaling",
        action="store_true")

    return parser.parse_args()

def save(path, ext='png', close=True, verbose=True):
    """Save a figure from pyplot.

    Parameters
    ----------
    path : string
        The path (and filename, without the extension) to save the
        figure to.

    ext : string (default='png')
        The file extension. This must be supported by the active
        matplotlib backend (see matplotlib.backends module).  Most
        backends support 'png', 'pdf', 'ps', 'eps', and 'svg'.

    close : boolean (default=True)
        Whether to close the figure after saving.  If you want to save
        the figure multiple times (e.g., to multiple formats), you
        should NOT close it in between saves or you will have to
        re-plot it.

    verbose : boolean (default=True)
        Whether to print information about when and where the image
        has been saved.

    """

    # Extract the directory and filename from the given path
    directory = os.path.split(path)[0]
    filename = "%s.%s" % (os.path.split(path)[1], ext)
    if directory == '':
        directory = '.'

    # If the directory does not exist, create it
    if not os.path.exists(directory):
        os.makedirs(directory)

    # The final path to save to
    savepath = os.path.join(directory, filename)

    if verbose:
        print("Saving figure to '%s'..." % savepath),

    # Actually save the figure
    plt.savefig(savepath)

    # Close it
    if close:
        plt.close()

    if verbose:
        print("Done")

def read_dataframe_from_excel(
        filename,
        worksheet,
        **kwargs):
    """Reads a pandas dataframe from an Excel file
    filename: the excel file
    worksheet: worksheet name within the excel file
    The remaining named arguments are passed to pandas.read_excel
    """
    results = pd.read_excel(filename, worksheet, **kwargs)
    return results

def calculate_speedup_and_efficiency(
        rdf,
        compute_element_col_index,
        time_col_index):
    """Calculates the speedup and efficiency and
    adds them as new columns to a dataframe

    rdf: the results dataframe
    compute_element_col_index: compute elements column name
    time_col_index: column name containing the computation times

    Returns: a tuple containing the speedup, strong scaling efficiency, and
    weak scaling efficiency series
    """
    # get the t1 reference value
    t1 = float(rdf.loc[rdf[compute_element_col_index] == 1][time_col_index])

    # speedup n = t1 / tn
    speedup = t1 / rdf[time_col_index]

    # strong scaling efficiency n = t1 / (tn * n)
    strong_efficiency = t1 / (rdf[time_col_index] * rdf[compute_element_col_index])

    # weak scaling efficiency n = t1 / tn
    weak_efficiency = t1 / rdf[time_col_index]

    return (speedup, strong_efficiency, weak_efficiency)

def plot_walltime(
        series,
        colours,
        series_names,
        compute_elements,
        ymax,
        plot_size=None,
        group_width=0.8,
        xlabel='Processes',
        ylabel='Minutes',
        title='Walltime',
        file_name='walltime',
        file_extension='png',
        y_log_scale=False,
        show=False):
    """creates a bar plot as a new pyplot figure"""

    # the x locations for the groups.
    x_ind = np.arange(len(compute_elements))

    # create the figure and axes
    fig = plt.figure(figsize=plot_size)
    ax = fig.add_subplot(111)

    # create the individual bars
    bars_per_group = len(series)
    bar_width = group_width / bars_per_group
    bar_left = x_ind - (group_width / 2)  # centre the group on the ticks
    plot_left = x_ind[0] - group_width
    plot_right = x_ind[-1] + group_width

    for bar_values, bar_colour in zip(series, colours):
        ax.bar(
            bar_left,
            bar_values,
            width=bar_width,
            color=bar_colour,
            log=y_log_scale)
        bar_left += bar_width

    # apply the labels and formatting
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.set_xlim(plot_left, plot_right)
    ax.set_ylim(0, ymax)
    ax.set_xticks(x_ind)
    ax.set_xticklabels(compute_elements)
    ax.legend(series_names)
    # ax.legend().get_frame().set_facecolor('white')

    if show:
        plt.show()
    else:
        save(file_name, file_extension, close=True, verbose=True)

def plot_speedup(
        series,
        colours,
        series_names,
        compute_elements,
        xmax,
        ymax,
        line_width=1,
        plot_size=None,
        xlabel='Processes',
        ylabel='Speedup',
        title='Speedup',
        file_name='speedup',
        file_extension='png',
        show=False):
    """creates a speedup plot"""

    # create the figure and axes
    fig = plt.figure(figsize=plot_size)
    ax = fig.add_subplot(111)

    # define the sizes and locations of things
    x_ticks = list(compute_elements)
    x_ticks.append(xmax)
    plot_left = x_ticks[0]
    plot_right = x_ticks[-1]
    ax.set_xlim(plot_left, plot_right)
    ax.set_ylim(plot_left, ymax)

    # create the ideal speedup reference line
    ax.plot(
        x_ticks,
        x_ticks,
        color='red',
        linewidth=line_width,
        label='_nolegend_')

    # plot each series
    for values, colour, name in zip(series, colours, series_names):
        ax.plot(
            compute_elements,
            values,
            label=name,
            color=colour,
            linewidth=line_width,
            marker='o')

    # apply the labels and formatting
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.set_xticks(compute_elements)
    ax.set_xticklabels(compute_elements)
    ax.legend()
    ax.legend().get_frame().set_facecolor('white')
    # can also use get_frame().set_xy(x,y) to give a custom location

    if show:
        plt.show()
    else:
        save(file_name, file_extension, close=True, verbose=True)

def plot_efficiency(
        series,
        colours,
        series_names,
        compute_elements,
        xmax,
        ymax=1.2,
        line_width=1,
        plot_size=None,
        xlabel='Processes',
        ylabel='Efficiency',
        title='Efficiency',
        file_name='efficiency',
        file_extension='png',
        show=False):
    """creates an efficiency plot"""

    # create the figure and axes
    fig = plt.figure(figsize=plot_size)
    ax = fig.add_subplot(111)

    # define the sizes and locations of things
    x_ticks = list(compute_elements)
    x_ticks.append(xmax)
    plot_left = x_ticks[0]
    plot_right = x_ticks[-1]
    ax.set_xlim(plot_left, plot_right)
    ax.set_ylim(0, ymax)

    # create a horizontal ideal efficiency reference line
    ax.hlines(
        1,
        plot_left,
        plot_right,
        colors='red',
        linestyles='solid',
        linewidth=line_width,
        label='_nolegend_')

    # plot each series
    for values, colour, name in zip(series, colours, series_names):
        ax.plot(
            compute_elements,
            values,
            label=name,
            color=colour,
            linewidth=line_width,
            marker='o')

    # apply the labels and formatting
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.set_xticks(compute_elements)
    ax.set_xticklabels(compute_elements)
    ax.legend(shadow=True)
    ax.legend().get_frame().set_facecolor('white')

    if show:
        plt.show()
    else:
        save(file_name, file_extension, close=True, verbose=True)

def add_optional_prefix(base_string, prefix, separator):
    result = str(base_string)
    if prefix:
        result = prefix + separator + base_string
    return result

if __name__ == '__main__':

    # get and process the arguments
    args = get_args()

    # create plots in a 4:3 aspect ratio
    # matplotlib works in inches
    plot_width = args.plot_width
    plot_size = (plot_width, plot_width * 3 / 4)

    show_instead_of_save = args.window

    walltime_title = add_optional_prefix('Walltime', args.title_prefix, ' - ')
    speedup_title = add_optional_prefix('Speedup', args.title_prefix, ' - ')

    efficiency_title = add_optional_prefix(
        'Weak Scaling Efficiency' if args.weak else 'Strong Scaling Efficiency',
        args.title_prefix,
        ' - ')

    walltime_file = add_optional_prefix('walltime', args.file_prefix, '-')
    speedup_file = add_optional_prefix('speedup', args.file_prefix, '-')
    efficiency_file = add_optional_prefix(
        'weak-efficiency' if args.weak else 'strong-efficiency',
        args.file_prefix,
        '-')

    compute_element_name = args.compute_element_name
    walltime_units = args.walltime_units

    # CSIRO colours
    colours = [
        '#00a9ce',  # midday blue
        '#78be20',  # light forest
        '#DF1995',  # fuschia
        '#E87722',  # orange
        '#E4002B',  # vermillion
        '#00616c',  # midnight blue
        '#FFB81C',  # gold
        '#6D2077',  # plum
        '#1E22AA',  # blueberry
    ]

    usecols = [
        'group',
        'compute_elements',
        'walltime']
    if args.filter_column:
        usecols.append(args.filter_column)

    results = read_dataframe_from_excel(
        args.results_file,
        worksheet=args.worksheet_name,
        usecols=usecols)

    # filter out incomplete results
    results = results[results.group.notnull() & (results.walltime > 0)]

    # apply the optional filter column
    if args.filter_column:
        results = results[results[args.filter_column] > 0]

    # if there are multiple times for each (group,compute_element) tuple,
    # calculate the mean
    results = results.groupby(['group','compute_elements']).mean().reset_index()

    # create new dataframes for each group of results
    groups = results.group.unique()
    group_dataframes = {}
    for group in groups:
        df = results[(results.group == group)]
        if len(df):
            group_dataframes[group] = df.sort(
                'compute_elements',
                ascending=True)

    # calculate speedup and efficiency
    for data in group_dataframes.values():
        data['speedup'], data['strong_efficiency'], data['weak_efficiency'] = \
            calculate_speedup_and_efficiency(
                data,
                'compute_elements',
                'walltime')

    # extract the compute element names for grouping and labelling the charts
    compute_elements = np.sort(results.compute_elements.unique())

    # build the collections of values for the charts
    walltimes = []
    strong_efficiencies = []
    weak_efficiencies = []
    speedups = []
    series_names = []
    max_walltime = 0
    for name, data in group_dataframes.items():
        walltimes.append(data.walltime)
        strong_efficiencies.append(data.strong_efficiency)
        weak_efficiencies.append(data.weak_efficiency)
        speedups.append(data.speedup)
        series_names.append(name)
        local_max_walltime = max(data.walltime)
        if local_max_walltime > max_walltime:
            max_walltime = local_max_walltime

    # finally make the plots

    plot_walltime(
        walltimes,
        colours,
        series_names,
        compute_elements,
        ymax=max_walltime * 1.2,
        group_width=0.83,
        plot_size=plot_size,
        title=walltime_title,
        xlabel=compute_element_name,
        ylabel=walltime_units,
        show=show_instead_of_save,
        file_name=walltime_file,
        file_extension='png')

    plot_efficiency(
        weak_efficiencies if args.weak else strong_efficiencies,
        colours,
        series_names,
        compute_elements,
        line_width=1.5,
        xmax=compute_elements.max() * 1.1,
        ymax=1.3,
        plot_size=plot_size,
        xlabel=compute_element_name,
        title=efficiency_title,
        show=show_instead_of_save,
        file_name=efficiency_file,
        file_extension='png')

    # The Speedup plot makes no sense for weak scaling.
    # I could create a new plot for weak scaling that shows the percentage
    # increase in walltime as the compute elements increase
    if not args.weak:
        plot_speedup(
            speedups,
            colours,
            series_names,
            compute_elements,
            line_width=1.5,
            xmax=compute_elements.max() * 1.05,
            ymax=args.speedup_max,
            plot_size=plot_size,
            xlabel=compute_element_name,
            title=speedup_title,
            show=show_instead_of_save,
            file_name=speedup_file,
            file_extension='png')
<|endoftext|>


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

__author__ = 'ievans3024'
__version__ = '0.0.1'

from os import getcwd
from tkinter import BOTH, E, filedialog, LEFT, Menu, N, NW, PhotoImage, S, Tk, W
from tkinter.ttk import Button, Frame, Menubutton, Separator, Style, Treeview

# Tk instantiated outside of __main__ to allow use of Image classes for icon definitions
root = Tk()

# Icon definitions are necessary because widgets do not save data provided to the 'image' kwarg to the instance.
# To prevent image data from being garbage collected, they must exist in an external possible reference.
icons = {
    'about': PhotoImage(file='icons/gnome-about-logo.png'),
    'actions': {
        'chunk_find': PhotoImage(file='icons/workspace-switcher.png'),
        'copy': PhotoImage(file='icons/edit-copy.png'),
        'cut': PhotoImage(file='icons/edit-cut.png'),
        'delete': PhotoImage(file='icons/list-remove.png'),
        'edit': PhotoImage(file='icons/gtk-edit.png'),
        'exit': PhotoImage(file='icons/system-log-out.png'),
        'move': {
            'down': PhotoImage(file='icons/go-down.png'),
            'right': PhotoImage(file='icons/go-next.png'),
            'up': PhotoImage(file='icons/go-up.png')
        },
        'new': {
            'generic': PhotoImage(file='icons/list-add.png'),
            'tag': {
                'byte': PhotoImage(file='icons/tag-byte-new.png'),
                'byte_array': PhotoImage(file='icons/tag-byte-array-new.png'),
                'compound': PhotoImage(file='icons/tag-compound-new.png'),
                'double': PhotoImage(file='icons/tag-double-new.png'),
                'float': PhotoImage(file='icons/tag-float-new.png'),
                'int': PhotoImage(file='icons/tag-int-new.png'),
                'int_array': PhotoImage(file='icons/tag-int-array-new.png'),
                'list': PhotoImage(file='icons/tag-list-new.png'),
                'long': PhotoImage(file='icons/tag-long-new.png'),
                'short': PhotoImage(file='icons/tag-short-new.png'),
                'string': PhotoImage(file='icons/tag-string-new.png'),
            }
        },
        'open': {
            'file': PhotoImage(file='icons/document-open.png'),
            'folder': PhotoImage(file='icons/folder-move.png'),
            'mc_folder': PhotoImage(file='icons/go-home.png')
        },
        'paste': PhotoImage(file='icons/edit-paste.png'),
        'refresh': PhotoImage(file='icons/view-refresh.png'),
        'rename': PhotoImage(file='icons/format-text-strikethrough.png'),
        'replace': PhotoImage(file='icons/edit-find-replace.png'),
        'save': PhotoImage(file='icons/media-floppy.png'),
        'search': PhotoImage(file='icons/edit-find.png')
    },
    'tags': {
        'byte': PhotoImage(file='icons/tag-byte.png'),
        'byte_array': PhotoImage(file='icons/tag-byte-array.png'),
        'compound': PhotoImage(file='icons/tag-compound.png'),
        'double': PhotoImage(file='icons/tag-double.png'),
        'float': PhotoImage(file='icons/tag-float.png'),
        'int': PhotoImage(file='icons/tag-int.png'),
        'int_array': PhotoImage(file='icons/tag-int-array.png'),
        'list': PhotoImage(file='icons/tag-list.png'),
        'long': PhotoImage(file='icons/tag-long.png'),
        'short': PhotoImage(file='icons/tag-short.png'),
        'string': PhotoImage(file='icons/text-plain.png'),
    }
}

OPEN_FILES = lambda: filedialog.askopenfilenames(initialdir=getcwd())
OPEN_FOLDER = lambda: filedialog.askdirectory(initialdir=getcwd())

class MainMenu(Menu):

    def __init__(self, parent):

        Menu.__init__(self, parent)

        self.parent = parent

        menu_file = Menu(self, tearoff=0)
        menu_file.add_command(compound='left', image=icons['actions']['new']['generic'], label='New...')
        menu_file.add_command(compound='left',
                              image=icons['actions']['open']['file'], label='Open...', command=OPEN_FILES)
        menu_file.add_command(compound='left',
                              image=icons['actions']['open']['folder'], label='Open Folder...', command=OPEN_FOLDER)
        menu_file.add_command(compound='left',
                              image=icons['actions']['open']['mc_folder'], label='Open Minecraft Save Folder...')
        menu_file.add_separator()
        menu_file.add_command(compound='left', image=icons['actions']['save'], label='Save...')
        menu_file.add_command(compound='left', image=icons['actions']['refresh'], label='Refresh')
        menu_file.add_separator()
        menu_file.add_command(compound='left', image=icons['actions']['exit'], label='Exit', command=root.quit)

        menu_edit = Menu(self, tearoff=0)
        menu_edit.add_command(compound='left', image=icons['actions']['copy'], label='Copy')
        menu_edit.add_command(compound='left', image=icons['actions']['cut'], label='Cut')
        menu_edit.add_command(compound='left', image=icons['actions']['paste'], label='Paste')
        menu_edit.add_separator()
        menu_edit.add_command(compound='left', image=icons['actions']['rename'], label='Rename')
        menu_edit.add_command(compound='left', image=icons['actions']['edit'], label='Edit')
        menu_edit.add_command(compound='left', image=icons['actions']['delete'], label='Delete')
        menu_edit.add_separator()
        menu_edit.add_command(compound='left', image=icons['actions']['move']['up'], label='Move Up')
        menu_edit.add_command(compound='left', image=icons['actions']['move']['down'], label='Move Down')

        menu_search = Menu(self, tearoff=0)
        menu_search.add_command(compound='left', image=icons['actions']['search'], label='Find')
        menu_search.add_command(compound='left', image=icons['actions']['move']['right'], label='Find Next')
        menu_search.add_separator()
        menu_search.add_command(compound='left', image=icons['actions']['replace'], label='Replace')
        menu_search.add_separator()
        menu_search.add_command(compound='left', image=icons['actions']['chunk_find'], label='Chunk Finder')

        menu_help = Menu(self, tearoff=0)
        menu_help.add_command(compound='left', image=icons['about'], label='About')

        self.add_cascade(label='File', menu=menu_file)
        self.add_cascade(label='Edit', menu=menu_edit)
        self.add_cascade(label='Search', menu=menu_search)
        self.add_cascade(label='Help', menu=menu_help)

class ToolBar(Frame):

    def __init__(self, parent):

        Frame.__init__(self, parent)

        self.parent = parent

        elements = [
            Button(self, image=icons['actions']['open']['file'], command=OPEN_FILES),
            Button(self, image=icons['actions']['open']['folder'], command=OPEN_FOLDER),
            Button(self, image=icons['actions']['save']),
            Button(self, image=icons['actions']['refresh']),
            Separator(self, orient='vertical'),
            Button(self, image=icons['actions']['cut']),
            Button(self, image=icons['actions']['copy']),
            Button(self, image=icons['actions']['paste']),
            Separator(self, orient='vertical'),
            Button(self, image=icons['actions']['rename']),
            Button(self, image=icons['actions']['edit']),
            Button(self, image=icons['actions']['delete']),
            Separator(self, orient='vertical'),
            Button(self, image=icons['actions']['new']['tag']['byte']),
            Button(self, image=icons['actions']['new']['tag']['short']),
            Button(self, image=icons['actions']['new']['tag']['int']),
            Button(self, image=icons['actions']['new']['tag']['long']),
            Button(self, image=icons['actions']['new']['tag']['float']),
            Button(self, image=icons['actions']['new']['tag']['double']),
            Button(self, image=icons['actions']['new']['tag']['byte_array']),
            Button(self, image=icons['actions']['new']['tag']['int_array']),
            Button(self, image=icons['actions']['new']['tag']['string']),
            Button(self, image=icons['actions']['new']['tag']['list']),
            Button(self, image=icons['actions']['new']['tag']['compound']),
            Separator(self, orient='vertical'),
            Button(self, image=icons['actions']['search'])
        ]

        for e in elements:
            if isinstance(e, Separator):
                e.grid(row=0, column=elements.index(e), sticky=N+S, padx=6)
            else:
                e.grid(row=0, column=elements.index(e), padx=1)

class TreeDisplay(Treeview):

    class Node(object):

        def __init__(self):
            pass

        def add(self):
            pass

        def remove(self):
            pass

    def __init__(self, parent, **kwargs):

        # TODO: Store non-treeview kwargs locally, then strip out and pass the rest to Treeview.__init__()
        Treeview.__init__(self, parent, **kwargs)

        self.parent = parent

class MainWindow(Frame):

    def __init__(self, parent):

        Frame.__init__(self, parent)

        self.parent = parent
        self.parent.title('PyNBTx %s' % __version__)

        self.style = Style()
        self.style.theme_use('default')
        self.style.configure('TButton', padding=0)

        self.pack(fill=BOTH, expand=1)

        self.ui_init()

    def ui_init(self):

        menu = MainMenu(root)
        root['menu'] = menu

        toolbar = ToolBar(self)
        toolbar.pack(anchor=NW, padx=4, pady=4)

        tree = Treeview(self, height=20)
        tree.column('#0', width=300)
        tree.insert('', 'end', 'root', text='root')
        tree.insert('root', '0', 'one', text='one')
        tree.insert('one', '0', 'two', text='two')
        tree.insert('two', '0', 'three', text='three')
        tree.insert('three', '0', 'four', text='four')
        tree.insert('four', '0', 'five', text='five')
        tree.insert('five', '0', 'six', text='six')
        tree.insert('six', '0', 'seven', text='seven')
        tree.insert('seven', '0', 'eight', text='eight')
        tree.insert('eight', '0', 'nine', text='nine')
        tree.insert('nine', '0', 'ten', text='ten')
        tree.pack(fill=BOTH, anchor=NW, padx=4, pady=4)

def main():

    app = MainWindow(root)
    root.mainloop()

if __name__ == '__main__':
    main()<commit_msg>* version increment * changed TreeDisplay to fit new purpose + created non-functional file open/save functions + added some (probably unhelpful) docstrings<commit_after>#!/usr/bin/env python
# coding: utf-8

__author__ = 'ievans3024'
__version__ = '0.0.2'

from os import getcwd, getenv, walk
from os.path import expanduser, join
from sys import platform
from tkinter import BOTH, filedialog, Menu, N, NW, PhotoImage, S, Tk
from tkinter.ttk import Button, Frame, Separator, Style, Treeview

import nbt

# Tk instantiated outside of __main__ to allow use of Image classes for icon definitions
root = Tk()

# Icon definitions are necessary because widgets do not save data provided to the 'image' kwarg to the instance.
# To prevent image data from being garbage collected, they must exist in an external possible reference.
icons = {
    'about': PhotoImage(file='icons/gnome-about-logo.png'),
    'actions': {
        'chunk_find': PhotoImage(file='icons/workspace-switcher.png'),
        'copy': PhotoImage(file='icons/edit-copy.png'),
        'cut': PhotoImage(file='icons/edit-cut.png'),
        'delete': PhotoImage(file='icons/list-remove.png'),
        'edit': PhotoImage(file='icons/gtk-edit.png'),
        'exit': PhotoImage(file='icons/system-log-out.png'),
        'move': {
            'down': PhotoImage(file='icons/go-down.png'),
            'right': PhotoImage(file='icons/go-next.png'),
            'up': PhotoImage(file='icons/go-up.png')
        },
        'new': {
            'generic': PhotoImage(file='icons/list-add.png'),
            'tag': {
                'byte': PhotoImage(file='icons/tag-byte-new.png'),
                'byte_array': PhotoImage(file='icons/tag-byte-array-new.png'),
                'compound': PhotoImage(file='icons/tag-compound-new.png'),
                'double': PhotoImage(file='icons/tag-double-new.png'),
                'float': PhotoImage(file='icons/tag-float-new.png'),
                'int': PhotoImage(file='icons/tag-int-new.png'),
                'int_array': PhotoImage(file='icons/tag-int-array-new.png'),
                'list': PhotoImage(file='icons/tag-list-new.png'),
                'long': PhotoImage(file='icons/tag-long-new.png'),
                'short': PhotoImage(file='icons/tag-short-new.png'),
                'string': PhotoImage(file='icons/tag-string-new.png'),
            }
        },
        'open': {
            'file': PhotoImage(file='icons/document-open.png'),
            'folder': PhotoImage(file='icons/folder-move.png'),
            'mc_folder': PhotoImage(file='icons/go-home.png')
        },
        'paste': PhotoImage(file='icons/edit-paste.png'),
        'refresh': PhotoImage(file='icons/view-refresh.png'),
        'rename': PhotoImage(file='icons/format-text-strikethrough.png'),
        'replace': PhotoImage(file='icons/edit-find-replace.png'),
        'save': PhotoImage(file='icons/media-floppy.png'),
        'search': PhotoImage(file='icons/edit-find.png')
    },
    'tags': {
        'byte': PhotoImage(file='icons/tag-byte.png'),
        'byte_array': PhotoImage(file='icons/tag-byte-array.png'),
        'compound': PhotoImage(file='icons/tag-compound.png'),
        'double': PhotoImage(file='icons/tag-double.png'),
        'float': PhotoImage(file='icons/tag-float.png'),
        'int': PhotoImage(file='icons/tag-int.png'),
        'int_array': PhotoImage(file='icons/tag-int-array.png'),
        'list': PhotoImage(file='icons/tag-list.png'),
        'long': PhotoImage(file='icons/tag-long.png'),
        'short': PhotoImage(file='icons/tag-short.png'),
        'string': PhotoImage(file='icons/text-plain.png'),
    }
}

loaded_files = []

def open_files(get_files=()):
    global loaded_files
    if not get_files:
        get_files = filedialog.askopenfilenames(initialdir=getcwd())
    # TODO: make this iterate over get_files, open and append valid nbt files to loaded_files

def open_folder():
    global loaded_files
    filetree = walk(filedialog.askdirectory(initialdir=getcwd()))
    # TODO: make this examine ever file in the tree, append valid nbt files to loaded_files list

def open_mc_dir():
    global loaded_files
    if str(platform) == 'windows':
        mc_parent = getenv('APPDATA')
    else:
        mc_parent = expanduser('~')

    mc_dir = join(mc_parent, '.minecraft', 'saves')
    get_files = filedialog.askopenfilename(initialdir=mc_dir)
    # TODO: decide if I want to open all savegame nbt files or just level.dat

def save_files():
    global loaded_files
    for f in loaded_files:
        if isinstance(f, nbt.nbt.NBTFile):
            f.write_file()

class MainMenu(Menu):
    """
    Macro class for main window menus
    """
    def __init__(self, parent):

        Menu.__init__(self, parent)

        self.parent = parent

        menu_file = Menu(self, tearoff=0)
        menu_file.add_command(compound='left', image=icons['actions']['new']['generic'], label='New...')
        menu_file.add_command(compound='left',
                              image=icons['actions']['open']['file'], label='Open...', command=open_files)
        menu_file.add_command(compound='left',
                              image=icons['actions']['open']['folder'], label='Open Folder...', command=open_folder)
        menu_file.add_command(compound='left',
                              image=icons['actions']['open']['mc_folder'],
                              label='Open Minecraft Save Folder...', command=open_mc_dir)
        menu_file.add_separator()
        menu_file.add_command(compound='left', image=icons['actions']['save'], label='Save...', command=save_files)
        menu_file.add_command(compound='left', image=icons['actions']['refresh'], label='Refresh')
        menu_file.add_separator()
        menu_file.add_command(compound='left', image=icons['actions']['exit'], label='Exit', command=root.quit)

        menu_edit = Menu(self, tearoff=0)
        menu_edit.add_command(compound='left', image=icons['actions']['copy'], label='Copy')
        menu_edit.add_command(compound='left', image=icons['actions']['cut'], label='Cut')
        menu_edit.add_command(compound='left', image=icons['actions']['paste'], label='Paste')
        menu_edit.add_separator()
        menu_edit.add_command(compound='left', image=icons['actions']['rename'], label='Rename')
        menu_edit.add_command(compound='left', image=icons['actions']['edit'], label='Edit')
        menu_edit.add_command(compound='left', image=icons['actions']['delete'], label='Delete')
        menu_edit.add_separator()
        menu_edit.add_command(compound='left', image=icons['actions']['move']['up'], label='Move Up')
        menu_edit.add_command(compound='left', image=icons['actions']['move']['down'], label='Move Down')

        menu_search = Menu(self, tearoff=0)
        menu_search.add_command(compound='left', image=icons['actions']['search'], label='Find')
        menu_search.add_command(compound='left', image=icons['actions']['move']['right'], label='Find Next')
        menu_search.add_separator()
        menu_search.add_command(compound='left', image=icons['actions']['replace'], label='Replace')
        menu_search.add_separator()
        menu_search.add_command(compound='left', image=icons['actions']['chunk_find'], label='Chunk Finder')

        menu_help = Menu(self, tearoff=0)
        menu_help.add_command(compound='left', image=icons['about'], label='About')

        self.add_cascade(label='File', menu=menu_file)
        self.add_cascade(label='Edit', menu=menu_edit)
        self.add_cascade(label='Search', menu=menu_search)
        self.add_cascade(label='Help', menu=menu_help)

class ToolBar(Frame):
    """
    Macro class for main window toolbar
    """
    def __init__(self, parent):

        Frame.__init__(self, parent)

        self.parent = parent

        elements = [
            Button(self, image=icons['actions']['open']['file'], command=open_files),
            Button(self, image=icons['actions']['open']['folder'], command=open_folder),
            Button(self, image=icons['actions']['save']),
            Button(self, image=icons['actions']['refresh']),
            Separator(self, orient='vertical'),
            Button(self, image=icons['actions']['cut']),
            Button(self, image=icons['actions']['copy']),
            Button(self, image=icons['actions']['paste']),
            Separator(self, orient='vertical'),
            Button(self, image=icons['actions']['rename']),
            Button(self, image=icons['actions']['edit']),
            Button(self, image=icons['actions']['delete']),
            Separator(self, orient='vertical'),
            Button(self, image=icons['actions']['new']['tag']['byte']),
            Button(self, image=icons['actions']['new']['tag']['short']),
            Button(self, image=icons['actions']['new']['tag']['int']),
            Button(self, image=icons['actions']['new']['tag']['long']),
            Button(self, image=icons['actions']['new']['tag']['float']),
            Button(self, image=icons['actions']['new']['tag']['double']),
            Button(self, image=icons['actions']['new']['tag']['byte_array']),
            Button(self, image=icons['actions']['new']['tag']['int_array']),
            Button(self, image=icons['actions']['new']['tag']['string']),
            Button(self, image=icons['actions']['new']['tag']['list']),
            Button(self, image=icons['actions']['new']['tag']['compound']),
            Separator(self, orient='vertical'),
            Button(self, image=icons['actions']['search'])
        ]

        for e in elements:
            if isinstance(e, Separator):
                e.grid(row=0, column=elements.index(e), sticky=N+S, padx=6)
            else:
                e.grid(row=0, column=elements.index(e), padx=1)

class TreeDisplay(object):
    """
    Acts as an intermediary between a treeview object and one or more nbt objects
    """

    def __init__(self, treeview, nbts=[], **kwargs):

        self.treeview = treeview
        self.nbts = nbts

class MainWindow(Frame):
    """
    Macro class for the main program window
    """
    def __init__(self, parent):

        Frame.__init__(self, parent)

        self.parent = parent
        self.parent.title('PyNBTx %s' % __version__)

        self.style = Style()
        self.style.theme_use('default')
        self.style.configure('TButton', padding=0)

        self.pack(fill=BOTH, expand=1)

        self.ui_init()

    def ui_init(self):

        menu = MainMenu(root)
        root['menu'] = menu

        toolbar = ToolBar(self)
        toolbar.pack(anchor=NW, padx=4, pady=4)

        tree = Treeview(self, height=20)
        tree.column('#0', width=300)
        tree.insert('', 'end', 'root', text='root')
        tree.insert('root', '0', 'one', text='one')
        tree.insert('one', '0', 'two', text='two')
        tree.insert('two', '0', 'three', text='three')
        tree.insert('three', '0', 'four', text='four')
        tree.insert('four', '0', 'five', text='five')
        tree.insert('five', '0', 'six', text='six')
        tree.insert('six', '0', 'seven', text='seven')
        tree.insert('seven', '0', 'eight', text='eight')
        tree.insert('eight', '0', 'nine', text='nine')
        tree.insert('nine', '0', 'ten', text='ten')
        tree.pack(fill=BOTH, anchor=NW, padx=4, pady=4)

def main():

    app = MainWindow(root)
    root.mainloop()

if __name__ == '__main__':
    main()<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""pyjstat is a python module for JSON-stat formatted data manipulation.

This module allows reading and writing JSON-stat [1]_ format with python,
using data frame structures provided by the widely accepted
pandas library [2]_. The JSON-stat format is a simple lightweight JSON format
for data dissemination. Pyjstat is inspired in rjstat [3]_, a library to read
and write JSON-stat with R, by ajschumacher.

pyjstat is written and maintained by `Miguel Expósito Martín
<https://twitter.com/predicador37>`_ and is distributed under the Apache 2.0
License (see LICENSE file).

.. [1] http://json-stat.org/ for JSON-stat information
.. [2] http://pandas.pydata.org for Python Data Analysis Library information
.. [3] https://github.com/ajschumacher/rjstat for rjstat library information

Example:
  Importing a JSON-stat file into a pandas data frame can be done as follows::

    import urllib2
    import json
    import pyjstat
    results = pyjstat.from_json_stat(json.load(urllib2.urlopen(
    'http://json-stat.org/samples/oecd-canada.json')))
    print results

"""

import json
import sys
import pandas as pd
import numpy as np
from collections import OrderedDict

class NumpyEncoder(json.JSONEncoder):
    """Custom JSON encoder class for Numpy data types.

    """
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        elif isinstance(obj, str):
            return obj.encode('utf-8')
        else:
            return super(NumpyEncoder, self).default(obj)

def make_unicode(text):
    """Convert string to unicode depending on python version.

    Args:
      text (string): a string.

    Returns:
       text(string): a utf-8 enconded string.

    """
    if sys.version_info < (3,):
        return unicode(text)
    else:
        return str(text)

def to_int(variable):
    """Convert variable to integer or string depending on the case.

    Args:
      variable (string): a string containing a real string or an integer.

    Returns:
      variable(int, string): an integer or a string, depending on the content\
                             of variable.

    """
    try:
        return int(variable)
    except ValueError:
        return variable.encode('utf-8')

def parse_value(value):
    """Remove trailing zeroes from float values if they can be represented\
        as integers.

    Args:
      value (object): a python object (hopefully a number).

    Returns:
      value(int, object): an integer or the same input object, depending on\
                          the content of value.

    """

    if isinstance(value, float):
        if value.is_integer():
            return int(value)
        else:
            return value
    else:
        return value

def check_input(naming):
    """Check and validate input params.

    Args:
      naming (string): a string containing the naming type (label or id).

    Returns:
      Nothing

    Raises:
      ValueError: if the parameter is not in the allowed list.

    """

    if naming not in ['label', 'id']:
        raise ValueError('naming must be "label" or "id"')

def get_dimensions(js_dict, naming):
    """Get dimensions from input data.

    Args:
      js_dict (dict): dictionary containing dataset data and metadata.
      naming (string, optional): dimension naming. Possible values: 'label' \
                                 or 'id'.

    Returns:
      dimensions (list): list of pandas data frames with dimension \
                         category data.
      dim_names (list): list of strings with dimension names.
    """

    dimensions = []
    dim_names = []
    for dim in js_dict['dimension']['id']:
        dim_name = js_dict['dimension'][dim]['label']
        if not dim_name:
            dim_name = dim
        if naming == 'label':
            dim_label = get_dim_label(js_dict, dim)
            dimensions.append(dim_label)
            dim_names.append(dim_name)
        else:
            dim_index = get_dim_index(js_dict, dim)
            dimensions.append(dim_index)
            dim_names.append(dim)
    return dimensions, dim_names

def get_dim_label(js_dict, dim):
    """Get label from a given dimension.

    Args:
      js_dict (dict): dictionary containing dataset data and metadata.
      dim (string): dimension name obtained from JSON file.

    Returns:
      dim_label(pandas.DataFrame): DataFrame with label-based dimension data.

    """

    try:
        dim_label = js_dict['dimension'][dim]['category']['label']
    except KeyError:
        dim_index = get_dim_index(js_dict, dim)
        dim_label = pd.concat([dim_index['id'],
                               dim_index['id']],
                              axis=1)
        dim_label.columns = ['id', 'label']
    else:
        dim_label = pd.DataFrame(list(zip(dim_label.keys(),
                                          dim_label.values())),
                                 index=dim_label.keys(),
                                 columns=['id', 'label'])
    return dim_label

def get_dim_index(js_dict, dim):
    """Get index from a given dimension.

    Args:
      js_dict (dict): dictionary containing dataset data and metadata.
      dim (string): dimension name obtained from JSON file.

    Returns:
      dim_index (pandas.DataFrame): DataFrame with index-based dimension data.

    """

    try:
        dim_index = js_dict['dimension'][dim]['category']['index']
    except KeyError:
        dim_label = get_dim_label(js_dict, dim)
        dim_index = pd.DataFrame(list(zip([dim_label['id'][0]], [0])),
                                 index=[0],
                                 columns=['id', 'index'])
    else:
        if type(dim_index) is list:
            dim_index = pd.DataFrame(list(zip(dim_index,
                                              range(0, len(dim_index)))),
                                     index=dim_index, columns=['id', 'index'])
        else:
            dim_index = pd.DataFrame(list(zip(dim_index.keys(),
                                              dim_index.values())),
                                     index=dim_index.keys(),
                                     columns=['id', 'index'])
    return dim_index

def get_values(js_dict, value='value'):
    """Get values from input data.

    Args:
      js_dict (dict): dictionary containing dataset data and metadata.
      value (string, optional): name of the value column. Defaults to 'value'.

    Returns:
      values (list): list of dataset values.

    """

    values = js_dict[value]
    if type(values) is list:
        if type(values[0]) is not dict or tuple:
            return values
    # being not a list of dicts or tuples leaves us with a dict...
    values = {int(key): value for (key, value) in values.items()}
    max_val = max(values.keys(), key=int) + 1
    vals = []
    for i in range(0, max_val):
        if values[i]:
            vals.append(values[i])
        else:
            vals.append(None)
    values = vals
    return values

def get_df_row(dimensions, naming='label', i=0, record=None):
    """Generate row dimension values for a pandas dataframe.

    Args:
      dimensions (list): list of pandas dataframes with dimension labels \
                         generated by get_dim_label or get_dim_index methods.
      naming (string, optional): dimension naming. Possible values: 'label' \
                                 or 'id'.
      i (int): dimension list iteration index. Default is 0, it's used in the \
                         recursive calls to the method.
      record (list): list of values representing a pandas dataframe row, \
                     except for the value column. Default is empty, it's used \
                     in the recursive calls to the method.

    Yields:
      list: list with pandas dataframe column values except for value column

    """

    check_input(naming)
    if i == 0 or record is None:
        record = []
    for dimension in dimensions[i][naming]:
        record.append(dimension)
        if len(record) == len(dimensions):
            yield record

        if i + 1 < len(dimensions):
            for row in get_df_row(dimensions, naming, i + 1, record):
                yield row
        if len(record) == i + 1:
            record.pop()

def uniquify(seq):
    """Return unique values in a list in the original order. See:  \
        http://www.peterbe.com/plog/uniqifiers-benchmark

    Args:
      seq (list): original list.

    Returns:
      list: list without duplicates preserving original order.

    """

    seen = set()
    seen_add = seen.add
    return [x for x in seq if x not in seen and not seen_add(x)]

def generate_df(js_dict, naming, value="value"):
    """Decode JSON-stat dict into pandas.DataFrame object. Helper method \
       that should be called inside from_json_stat().

    Args:
      js_dict(OrderedDict): OrderedDict with data in JSON-stat format, \
                            previously deserialized into a python object by \
                            json.load() or json.loads(), for example.
      naming(string): dimension naming. Possible values: 'label' or 'id.'
      value (string, optional): name of the value column. Defaults to 'value'.

    Returns:
      output(DataFrame): pandas.DataFrame with converted data.

    """

    values = []

    dimensions, dim_names = get_dimensions(js_dict, naming)
    values = get_values(js_dict, value=value)
    output = pd.DataFrame(columns=dim_names + [value],
                          index=range(0, len(values)))
    for i, category in enumerate(get_df_row(dimensions, naming)):
        output.loc[i] = category + [values[i]]
    output = output.convert_objects(convert_numeric=True)
    return output

def from_json_stat(datasets, naming='label', value='value'):
    """Decode JSON-stat formatted data into pandas.DataFrame object.

    Args:
      datasets(OrderedDict, list): data in JSON-stat format, previously \
                                   deserialized to a python object by \
                                   json.load() or json.loads(), for example.\
                                   Both List and OrderedDict are accepted \
                                   as inputs.
      naming(string, optional): dimension naming. Possible values: 'label'
                                or 'id.'
      value (string, optional): name of the value column. Defaults to 'value'.

    Returns:
      results(list): list of pandas.DataFrame with imported data.

    """

    check_input(naming)
    results = []
    if type(datasets) is list:
        for idx, element in enumerate(datasets):
            for dataset in element:
                js_dict = datasets[idx][dataset]
                results.append(generate_df(js_dict, naming, value))
    elif isinstance(datasets, OrderedDict) or type(datasets) is dict:
        if 'class' in datasets:
            if datasets['class'] == 'dataset':
                js_dict = datasets
                results.append(generate_df(js_dict, naming, value))
        else:  # 1.00 bundle type
            for dataset in datasets:
                js_dict = datasets[dataset]
                results.append(generate_df(js_dict, naming, value))
    return results

def to_json_stat(input_df, value='value', output='list'):
    """Encode pandas.DataFrame object into JSON-stat format. The DataFrames
       must have exactly one value column.

    Args:
      df(pandas.DataFrame): pandas data frame (or list of data frames) to
      encode.
       value (string, optional): name of the value column. Defaults to 'value'.
      output(string): accepts two values: 'list' or 'dict'. Produce list of\
                      dicts or dict of dicts as output.

    Returns:
      output(string): String with JSON-stat object.

    """

    data = []
    if output == 'list':
        result = []
    elif output == 'dict':
        result = {}
    if isinstance(input_df, pd.DataFrame):
        data.append(input_df)
    else:
        data = input_df
    for row, dataframe in enumerate(data):
        dims = data[row].filter([item for item in data[row].columns.values
                                 if item not in value])
        if len(dims.columns.values) != len(set(dims.columns.values)):
            raise ValueError('Non-value columns must constitute a unique ID')
        dim_names = list(dims)
        categories = [{to_int(i):
                       {"label": to_int(i),
                        "category":
                            {"index":
                             OrderedDict([(make_unicode(j), to_int(k))
                                          for k, j in enumerate(
                                              uniquify(dims[i]))]),
                             "label":
                                 OrderedDict([(to_int(j), make_unicode(j))
                                              for k, j in enumerate(
                                                  uniquify(dims[i]))])}}}
                      for i in dims.columns.values]
        dataset = {"dataset" + str(row + 1):
                   {"dimension": OrderedDict(),
                    value: [parse_value(x) for x in dataframe[value].where(
                        pd.notnull(dataframe[value]), None).values]}}
        for category in categories:
            dataset["dataset" + str(row + 1)]["dimension"].update(category)
        dataset[
            "dataset" + str(row + 1)]["dimension"].update({"id": dim_names})
        dataset["dataset" + str(row + 1)]["dimension"].update(
            {"size": [len(dims[i].unique()) for i in dims.columns.values]})
        for category in categories:
            dataset["dataset" + str(row + 1)]["dimension"].update(category)
        if output == 'list':
            result.append(dataset)
        elif output == 'dict':
            result.update(dataset)
        else:
            result = None
    return json.dumps(result, cls=NumpyEncoder)
<commit_msg>Removed unnecessary json encoding<commit_after># -*- coding: utf-8 -*-
"""pyjstat is a python module for JSON-stat formatted data manipulation.

This module allows reading and writing JSON-stat [1]_ format with python,
using data frame structures provided by the widely accepted
pandas library [2]_. The JSON-stat format is a simple lightweight JSON format
for data dissemination. Pyjstat is inspired in rjstat [3]_, a library to read
and write JSON-stat with R, by ajschumacher.

pyjstat is written and maintained by `Miguel Expósito Martín
<https://twitter.com/predicador37>`_ and is distributed under the Apache 2.0
License (see LICENSE file).

.. [1] http://json-stat.org/ for JSON-stat information
.. [2] http://pandas.pydata.org for Python Data Analysis Library information
.. [3] https://github.com/ajschumacher/rjstat for rjstat library information

Example:
  Importing a JSON-stat file into a pandas data frame can be done as follows::

    import urllib2
    import json
    import pyjstat
    results = pyjstat.from_json_stat(json.load(urllib2.urlopen(
    'http://json-stat.org/samples/oecd-canada.json')))
    print results

"""

import json
import sys
import pandas as pd
import numpy as np
from collections import OrderedDict

class NumpyEncoder(json.JSONEncoder):
    """Custom JSON encoder class for Numpy data types.

    """
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        else:
            return super(NumpyEncoder, self).default(obj)

def make_unicode(text):
    """Convert string to unicode depending on python version.

    Args:
      text (string): a string.

    Returns:
       text(string): a utf-8 enconded string.

    """
    if sys.version_info < (3,):
        return unicode(text)
    else:
        return str(text)

def to_int(variable):
    """Convert variable to integer or string depending on the case.

    Args:
      variable (string): a string containing a real string or an integer.

    Returns:
      variable(int, string): an integer or a string, depending on the content\
                             of variable.

    """
    try:
        return int(variable)
    except ValueError:
        return variable.encode('utf-8')

def parse_value(value):
    """Remove trailing zeroes from float values if they can be represented\
        as integers.

    Args:
      value (object): a python object (hopefully a number).

    Returns:
      value(int, object): an integer or the same input object, depending on\
                          the content of value.

    """

    if isinstance(value, float):
        if value.is_integer():
            return int(value)
        else:
            return value
    else:
        return value

def check_input(naming):
    """Check and validate input params.

    Args:
      naming (string): a string containing the naming type (label or id).

    Returns:
      Nothing

    Raises:
      ValueError: if the parameter is not in the allowed list.

    """

    if naming not in ['label', 'id']:
        raise ValueError('naming must be "label" or "id"')

def get_dimensions(js_dict, naming):
    """Get dimensions from input data.

    Args:
      js_dict (dict): dictionary containing dataset data and metadata.
      naming (string, optional): dimension naming. Possible values: 'label' \
                                 or 'id'.

    Returns:
      dimensions (list): list of pandas data frames with dimension \
                         category data.
      dim_names (list): list of strings with dimension names.
    """

    dimensions = []
    dim_names = []
    for dim in js_dict['dimension']['id']:
        dim_name = js_dict['dimension'][dim]['label']
        if not dim_name:
            dim_name = dim
        if naming == 'label':
            dim_label = get_dim_label(js_dict, dim)
            dimensions.append(dim_label)
            dim_names.append(dim_name)
        else:
            dim_index = get_dim_index(js_dict, dim)
            dimensions.append(dim_index)
            dim_names.append(dim)
    return dimensions, dim_names

def get_dim_label(js_dict, dim):
    """Get label from a given dimension.

    Args:
      js_dict (dict): dictionary containing dataset data and metadata.
      dim (string): dimension name obtained from JSON file.

    Returns:
      dim_label(pandas.DataFrame): DataFrame with label-based dimension data.

    """

    try:
        dim_label = js_dict['dimension'][dim]['category']['label']
    except KeyError:
        dim_index = get_dim_index(js_dict, dim)
        dim_label = pd.concat([dim_index['id'],
                               dim_index['id']],
                              axis=1)
        dim_label.columns = ['id', 'label']
    else:
        dim_label = pd.DataFrame(list(zip(dim_label.keys(),
                                          dim_label.values())),
                                 index=dim_label.keys(),
                                 columns=['id', 'label'])
    return dim_label

def get_dim_index(js_dict, dim):
    """Get index from a given dimension.

    Args:
      js_dict (dict): dictionary containing dataset data and metadata.
      dim (string): dimension name obtained from JSON file.

    Returns:
      dim_index (pandas.DataFrame): DataFrame with index-based dimension data.

    """

    try:
        dim_index = js_dict['dimension'][dim]['category']['index']
    except KeyError:
        dim_label = get_dim_label(js_dict, dim)
        dim_index = pd.DataFrame(list(zip([dim_label['id'][0]], [0])),
                                 index=[0],
                                 columns=['id', 'index'])
    else:
        if type(dim_index) is list:
            dim_index = pd.DataFrame(list(zip(dim_index,
                                              range(0, len(dim_index)))),
                                     index=dim_index, columns=['id', 'index'])
        else:
            dim_index = pd.DataFrame(list(zip(dim_index.keys(),
                                              dim_index.values())),
                                     index=dim_index.keys(),
                                     columns=['id', 'index'])
    return dim_index

def get_values(js_dict, value='value'):
    """Get values from input data.

    Args:
      js_dict (dict): dictionary containing dataset data and metadata.
      value (string, optional): name of the value column. Defaults to 'value'.

    Returns:
      values (list): list of dataset values.

    """

    values = js_dict[value]
    if type(values) is list:
        if type(values[0]) is not dict or tuple:
            return values
    # being not a list of dicts or tuples leaves us with a dict...
    values = {int(key): value for (key, value) in values.items()}
    max_val = max(values.keys(), key=int) + 1
    vals = []
    for i in range(0, max_val):
        if values[i]:
            vals.append(values[i])
        else:
            vals.append(None)
    values = vals
    return values

def get_df_row(dimensions, naming='label', i=0, record=None):
    """Generate row dimension values for a pandas dataframe.

    Args:
      dimensions (list): list of pandas dataframes with dimension labels \
                         generated by get_dim_label or get_dim_index methods.
      naming (string, optional): dimension naming. Possible values: 'label' \
                                 or 'id'.
      i (int): dimension list iteration index. Default is 0, it's used in the \
                         recursive calls to the method.
      record (list): list of values representing a pandas dataframe row, \
                     except for the value column. Default is empty, it's used \
                     in the recursive calls to the method.

    Yields:
      list: list with pandas dataframe column values except for value column

    """

    check_input(naming)
    if i == 0 or record is None:
        record = []
    for dimension in dimensions[i][naming]:
        record.append(dimension)
        if len(record) == len(dimensions):
            yield record

        if i + 1 < len(dimensions):
            for row in get_df_row(dimensions, naming, i + 1, record):
                yield row
        if len(record) == i + 1:
            record.pop()

def uniquify(seq):
    """Return unique values in a list in the original order. See:  \
        http://www.peterbe.com/plog/uniqifiers-benchmark

    Args:
      seq (list): original list.

    Returns:
      list: list without duplicates preserving original order.

    """

    seen = set()
    seen_add = seen.add
    return [x for x in seq if x not in seen and not seen_add(x)]

def generate_df(js_dict, naming, value="value"):
    """Decode JSON-stat dict into pandas.DataFrame object. Helper method \
       that should be called inside from_json_stat().

    Args:
      js_dict(OrderedDict): OrderedDict with data in JSON-stat format, \
                            previously deserialized into a python object by \
                            json.load() or json.loads(), for example.
      naming(string): dimension naming. Possible values: 'label' or 'id.'
      value (string, optional): name of the value column. Defaults to 'value'.

    Returns:
      output(DataFrame): pandas.DataFrame with converted data.

    """

    values = []

    dimensions, dim_names = get_dimensions(js_dict, naming)
    values = get_values(js_dict, value=value)
    output = pd.DataFrame(columns=dim_names + [value],
                          index=range(0, len(values)))
    for i, category in enumerate(get_df_row(dimensions, naming)):
        output.loc[i] = category + [values[i]]
    output = output.convert_objects(convert_numeric=True)
    return output

def from_json_stat(datasets, naming='label', value='value'):
    """Decode JSON-stat formatted data into pandas.DataFrame object.

    Args:
      datasets(OrderedDict, list): data in JSON-stat format, previously \
                                   deserialized to a python object by \
                                   json.load() or json.loads(), for example.\
                                   Both List and OrderedDict are accepted \
                                   as inputs.
      naming(string, optional): dimension naming. Possible values: 'label'
                                or 'id.'
      value (string, optional): name of the value column. Defaults to 'value'.

    Returns:
      results(list): list of pandas.DataFrame with imported data.

    """

    check_input(naming)
    results = []
    if type(datasets) is list:
        for idx, element in enumerate(datasets):
            for dataset in element:
                js_dict = datasets[idx][dataset]
                results.append(generate_df(js_dict, naming, value))
    elif isinstance(datasets, OrderedDict) or type(datasets) is dict:
        if 'class' in datasets:
            if datasets['class'] == 'dataset':
                js_dict = datasets
                results.append(generate_df(js_dict, naming, value))
        else:  # 1.00 bundle type
            for dataset in datasets:
                js_dict = datasets[dataset]
                results.append(generate_df(js_dict, naming, value))
    return results

def to_json_stat(input_df, value='value', output='list'):
    """Encode pandas.DataFrame object into JSON-stat format. The DataFrames
       must have exactly one value column.

    Args:
      df(pandas.DataFrame): pandas data frame (or list of data frames) to
      encode.
       value (string, optional): name of the value column. Defaults to 'value'.
      output(string): accepts two values: 'list' or 'dict'. Produce list of\
                      dicts or dict of dicts as output.

    Returns:
      output(string): String with JSON-stat object.

    """

    data = []
    if output == 'list':
        result = []
    elif output == 'dict':
        result = {}
    if isinstance(input_df, pd.DataFrame):
        data.append(input_df)
    else:
        data = input_df
    for row, dataframe in enumerate(data):
        dims = data[row].filter([item for item in data[row].columns.values
                                 if item not in value])
        if len(dims.columns.values) != len(set(dims.columns.values)):
            raise ValueError('Non-value columns must constitute a unique ID')
        dim_names = list(dims)
        categories = [{to_int(i):
                       {"label": to_int(i),
                        "category":
                            {"index":
                             OrderedDict([(make_unicode(j), to_int(k))
                                          for k, j in enumerate(
                                              uniquify(dims[i]))]),
                             "label":
                                 OrderedDict([(to_int(j), make_unicode(j))
                                              for k, j in enumerate(
                                                  uniquify(dims[i]))])}}}
                      for i in dims.columns.values]
        dataset = {"dataset" + str(row + 1):
                   {"dimension": OrderedDict(),
                    value: [parse_value(x) for x in dataframe[value].where(
                        pd.notnull(dataframe[value]), None).values]}}
        for category in categories:
            dataset["dataset" + str(row + 1)]["dimension"].update(category)
        dataset[
            "dataset" + str(row + 1)]["dimension"].update({"id": dim_names})
        dataset["dataset" + str(row + 1)]["dimension"].update(
            {"size": [len(dims[i].unique()) for i in dims.columns.values]})
        for category in categories:
            dataset["dataset" + str(row + 1)]["dimension"].update(category)
        if output == 'list':
            result.append(dataset)
        elif output == 'dict':
            result.update(dataset)
        else:
            result = None
    return json.dumps(result, cls=NumpyEncoder)
<|endoftext|>