PRSONA/prsona/scripts/make_graphs.py

#!/usr/bin/env python3

import os
import sys
import json
from math import sqrt
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit 
from contextlib import contextmanager

PLOT_OPTIONS = {
    "workload": {
        "all": {
            "legend": 'Covert Security, Workload: "all"',
            "marker": '^',
            "color": "red"
        }, 
        "half": {
            "legend": 'Covert Security, Workload: "half"',
            "marker": 'X',
            "color": "green"
        }, 
        "no": {
            "legend": 'Covert Security, Workload: "no"',
            "marker": '*',
            "color": "blue"
        },
        "hbcall": {
            "legend": 'HBC Security, Workload: "all"',
            "marker": '^',
            "color": "red"
        }, 
        "hbchalf": {
            "legend": 'HBC Security, Workload: "half"',
            "marker": 'X',
            "color": "green"
        }, 
        "hbcno": {
            "legend": 'HBC Security, Workload: "no"',
            "marker": '*',
            "color": "blue"
        }
    },
    "numServers": {
        2: {
            "legend": "2 servers",
            "marker": '^',
            "color": "red"
        },
        3: {
            "legend": "3 servers",
            "marker": 'X',
            "color": "blue"
        },
        4: {
            "legend": "4 servers",
            "marker": '*',
            "color": "green"
        },
        5: {
            "legend": "5 servers",
            "marker": 'h',
            "color": "orange"
        }
    },
    "numClients": {
        5: {
            "legend": "5 clients",
            "marker": "^",
            "color": "red"
        },
        10: {
            "legend": "10 clients",
            "marker": "v",
            "color": "green"
        },
        15: {
            "legend": "15 clients",
            "marker": ">",
            "color": "blue"
        },
        20: {
            "legend": "20 clients",
            "marker": "<",
            "color": "orange"
        },
        25: {
            "legend": "25 clients",
            "marker": "X",
            "color": "magenta"
        },
        30: {
            "legend": "30 clients",
            "marker": "*",
            "color": "pink"
        },
        40: {
            "legend": "40 clients",
            "marker": "h",
            "color": "cyan"
        },
        50: {
            "legend": "50 clients",
            "marker": ".",
            "color": "black"
        }
    },
    "lambda": {
        40: {
            "legend": "Lambda: 40",
            "marker": "^",
            "color": "red"
        },
        50: {
            "legend": "Lambda: 50",
            "marker": "X",
            "color": "green"
        },
        64: {
            "legend": "Lambda: 64",
            "marker": "*",
            "color": "blue"
        }
    }
}

##
# This functionality allows us to temporarily change our working directory
#
# @input newdir - the new directory (relative to our current position) we want to be in
@contextmanager
def cd(newdir, makenew):
    prevdir = os.getcwd()
    directory = os.path.expanduser(newdir)
    if not os.path.exists(directory) and makenew:
        os.makedirs(directory)
    os.chdir(directory)
    try:
        yield
    finally:
        os.chdir(prevdir)

def genericCube(x, a, b, c, d):
    return a * (x * x * x) + b * (x * x) + c * x + d

def readData(dataDirectory):
    serverData = {}
    clientData = {}

    realDirectory = os.path.expanduser(dataDirectory)
    for test in os.listdir(realDirectory):
        if not test.startswith('.') and not test.endswith('.tar.gz') and test.find("default") == -1:
            testParts = test.split("-")
            
            if not testParts[0] in serverData:
                serverData[testParts[0]] = {}

            if not testParts[0] in clientData:
                clientData[testParts[0]] = {}
            
            if not testParts[1] in serverData[testParts[0]]:
                serverData[testParts[0]][testParts[1]] = {}

            if not testParts[1] in clientData[testParts[0]]:
                clientData[testParts[0]][testParts[1]] = {}
            
            if not testParts[2] in serverData[testParts[0]][testParts[1]]:
                serverData[testParts[0]][testParts[1]][testParts[2]] = {}

            if not testParts[2] in clientData[testParts[0]][testParts[1]]:
                clientData[testParts[0]][testParts[1]][testParts[2]] = {}

            if not testParts[3] in serverData[testParts[0]][testParts[1]][testParts[2]]:
                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]] = {}

                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['overall_epoch_wall'] = []
                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['overall_epoch_cpu'] = []
                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['overall_epoch_recv'] = []
                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['overall_epoch_sent'] = []
                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['overall_epoch_total'] = []

                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_up_wall'] = []
                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_up_cpu'] = []
                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_up_recv'] = []
                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_up_sent'] = []
                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_up_total'] = []

                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_down_wall'] = []
                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_down_cpu'] = []
                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_down_recv'] = []
                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_down_sent'] = []
                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_down_total'] = []

                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_update_wall'] = []
                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_update_cpu'] = []
                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_update_recv'] = []
                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_update_sent'] = []
                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_update_total'] = []


            if not testParts[3] in clientData[testParts[0]][testParts[1]][testParts[2]]:
                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]] = {}

                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_wall'] = []
                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_cpu'] = []
                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_recv'] = []
                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_sent'] = []
                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_total'] = []

                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_prove_wall'] = []
                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_prove_cpu'] = []
                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_prove_recv'] = []
                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_prove_sent'] = []
                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_prove_total'] = []

                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_verify_wall'] = []
                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_verify_cpu'] = []
                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_verify_recv'] = []
                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_verify_sent'] = []
                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_verify_total'] = []

    for test in os.listdir(realDirectory):
        if not test.startswith('.') and not test.endswith('.tar.gz') and test.find("default") == -1:
            testParts = test.split("-")

            for whichEntity in os.listdir(os.path.join(realDirectory, test)):
                if whichEntity.startswith('s') or whichEntity.startswith('d'):
                    try:
                        with open(os.path.expanduser(os.path.join(realDirectory, test, whichEntity, 'overallEpoch.out')), 'r') as overallEpochFile:
                            for line in overallEpochFile:
                                lineParts = line.rstrip().split(',')
                                wallTime = float(lineParts[0])
                                cpuTime = float(lineParts[1])
                                dataRecv = float(lineParts[2])
                                dataSent = float(lineParts[3])

                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['overall_epoch_wall'].append(wallTime)
                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['overall_epoch_cpu'].append(cpuTime)
                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['overall_epoch_recv'].append(dataRecv)
                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['overall_epoch_sent'].append(dataSent)
                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['overall_epoch_total'].append(dataRecv + dataSent)
                    except FileNotFoundError as e:
                        pass

                    try:
                        with open(os.path.expanduser(os.path.join(realDirectory, test, whichEntity, 'epochUp.out')), 'r') as epochUpFile:
                            for line in epochUpFile:
                                lineParts = line.rstrip().split(',')
                                wallTime = float(lineParts[0])
                                cpuTime = float(lineParts[1])
                                dataRecv = float(lineParts[2])
                                dataSent = float(lineParts[3])

                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_up_wall'].append(wallTime)
                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_up_cpu'].append(cpuTime)
                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_up_recv'].append(dataRecv)
                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_up_sent'].append(dataSent)
                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_up_total'].append(dataRecv + dataSent)
                    except FileNotFoundError as e:
                        pass

                    try:
                        with open(os.path.expanduser(os.path.join(realDirectory, test, whichEntity, 'epochDown.out')), 'r') as epochDownFile:
                            for line in epochDownFile:
                                lineParts = line.rstrip().split(',')
                                wallTime = float(lineParts[0])
                                cpuTime = float(lineParts[1])
                                dataRecv = float(lineParts[2])
                                dataSent = float(lineParts[3])

                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_down_wall'].append(wallTime)
                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_down_cpu'].append(cpuTime)
                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_down_recv'].append(dataRecv)
                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_down_sent'].append(dataSent)
                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['epoch_down_total'].append(dataRecv + dataSent)
                    except FileNotFoundError as e:
                        pass

                    try:
                        with open(os.path.expanduser(os.path.join(realDirectory, test, whichEntity, 'voteUpdate.out')), 'r') as voteUpdateFile:
                            for line in voteUpdateFile:
                                lineParts = line.rstrip().split(',')
                                wallTime = float(lineParts[0])
                                cpuTime = float(lineParts[1])
                                dataRecv = float(lineParts[2])
                                dataSent = float(lineParts[3])

                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_update_wall'].append(wallTime)
                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_update_cpu'].append(cpuTime)
                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_update_recv'].append(dataRecv)
                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_update_sent'].append(dataSent)
                                serverData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_update_total'].append(dataRecv + dataSent)
                    except FileNotFoundError as e:
                        pass

                elif whichEntity.startswith('c'):
                    try:
                        with open(os.path.expanduser(os.path.join(realDirectory, test, whichEntity, 'vote.out')), 'r') as repVerifierFile:
                            for line in repVerifierFile:
                                lineParts = line.rstrip().split(',')
                                wallTime = float(lineParts[0])
                                cpuTime = float(lineParts[1])
                                dataRecv = float(lineParts[2])
                                dataSent = float(lineParts[3])

                                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_wall'].append(wallTime)
                                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_cpu'].append(cpuTime)
                                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_recv'].append(dataRecv)
                                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_sent'].append(dataSent)
                                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['vote_total'].append(dataRecv + dataSent)
                    except FileNotFoundError as e:
                        pass

                    try:
                        with open(os.path.expanduser(os.path.join(realDirectory, test, whichEntity, 'repProver.out')), 'r') as repProverFile:
                            for line in repProverFile:
                                lineParts = line.rstrip().split(',')
                                
                                if not ('*' in lineParts[0]):
                                    wallTime = float(lineParts[0])
                                    cpuTime = float(lineParts[1])
                                    dataRecv = float(lineParts[2])
                                    dataSent = float(lineParts[3])

                                    clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_prove_wall'].append(wallTime)
                                    clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_prove_cpu'].append(cpuTime)
                                    clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_prove_recv'].append(dataRecv)
                                    clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_prove_sent'].append(dataSent)
                                    clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_prove_total'].append(dataRecv + dataSent)

                    except FileNotFoundError as e:
                        pass

                    try:
                        with open(os.path.expanduser(os.path.join(realDirectory, test, whichEntity, 'repVerifier.out')), 'r') as repVerifierFile:
                            for line in repVerifierFile:
                                lineParts = line.rstrip().split(',')
                                wallTime = float(lineParts[0])
                                cpuTime = float(lineParts[1])
                                dataRecv = float(lineParts[2])
                                dataSent = float(lineParts[3])

                                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_verify_wall'].append(wallTime)
                                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_verify_cpu'].append(cpuTime)
                                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_verify_recv'].append(dataRecv)
                                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_verify_sent'].append(dataSent)
                                clientData[testParts[0]][testParts[1]][testParts[2]][testParts[3]]['rep_verify_total'].append(dataRecv + dataSent)
                    except FileNotFoundError as e:
                        pass

    return serverData, clientData

def plotComparison(data, dataParts, xVariable, lineVariable, whichGraph, **kwargs):
    z_star = kwargs['z_star'] if 'z_star' in kwargs else 1.96
    title = kwargs['title'] if 'title' in kwargs else ''
    xLabel = kwargs['xLabel'] if 'xLabel' in kwargs else ''
    yLabel = kwargs['yLabel'] if 'yLabel' in kwargs else ''
    yUnit = kwargs['yUnit'] if 'yUnit' in kwargs else ''
    fileNameStr = kwargs['fileNameStr'] if 'fileNameStr' in kwargs else f"basic-{xVariable}-{lineVariable}-{whichGraph}"
    legendLoc = kwargs['legendLoc'] if 'legendLoc' in kwargs else "best"
    legendBBoxAnchor = kwargs['legendBBoxAnchor'] if 'legendBBoxAnchor' in kwargs else (0, 1)
    extraFit = kwargs['extraFit'] if 'extraFit' in kwargs else False
    loglog = kwargs['loglog'] if 'loglog' in kwargs else False
    yLim = kwargs['yLim'] if 'yLim' in kwargs else False
    aspect = kwargs['aspect'] if 'aspect' in kwargs else None
    ignoreWorkload = kwargs['ignoreWorkload'] if 'ignoreWorkload' in kwargs else False

    fig = plt.figure()
    ax = fig.gca()

    whichLines = 0
    lineSelection = []
    legendStart = ""
    legendEnd = ""
    if lineVariable == 'workload':
        whichLines = 0
        lineSelection = [str(x) for x in data.keys() if x != "rep" and x != "vote"]
        lineSelection.sort()
        ignoreWorkload = False
    elif lineVariable == 'numServers':
        whichLines = 1
        lineSelection = [int(x) for x in data['no'].keys()]
        lineSelection.sort()
    elif lineVariable == 'numClients':
        whichLines = 2
        lineSelection = [int(x) for x in data['all']['2'].keys()]
        lineSelection.sort()
    elif lineVariable == 'lambda':
        whichLines = 3
        lineSelection = [int(x) for x in data['all']['2']['5'].keys()]
        lineSelection.sort()

    whichX = 0
    xSelection = []
    if xVariable == 'workload':
        whichX = 0
        xSelection = data.keys()
        ignoreWorkload = False
    elif xVariable == 'numServers':
        whichX = 1
        xSelection = [int(x) for x in data['all'].keys()]
        xSelection.sort()
    elif xVariable == 'numClients':
        whichX = 2
        xSelection = [int(x) for x in data['all']['2'].keys()]
        xSelection.sort()
    elif xVariable == 'lambda':
        whichX = 3
        xSelection = [int(x) for x in data['all']['2']['5'].keys()]
        xSelection.sort()

    for selection in lineSelection:
        xs = []
        xTicks = []
        ys = []
        additionalYs = []
        yErrs = []
        additionalYErrs = []

        legend = PLOT_OPTIONS[lineVariable][selection]['legend']
        marker = PLOT_OPTIONS[lineVariable][selection]['marker']
        color = PLOT_OPTIONS[lineVariable][selection]['color']

        dataParts[whichLines] = str(selection)
        
        for x in xSelection:
            dataParts[whichX] = str(x)

            try:
                curr_data = []
                if ignoreWorkload == False:
                    curr_data = data[dataParts[0]][dataParts[1]][dataParts[2]][dataParts[3]][whichGraph]
                else:
                    for workload in ['all', 'half', 'no']:
                        try:
                            curr_data.extend(data[workload][dataParts[1]][dataParts[2]][dataParts[3]][whichGraph])
                        except KeyError as e:
                            pass

                if len(curr_data) == 0:
                    continue

                dividing_factor = 1
                if yUnit == 'KB':
                    dividing_factor = 1024.0
                if yUnit == 'MB':
                    dividing_factor = 1024.0 * 1024.0

                if lineVariable == 'numServers':
                    dividing_factor = dividing_factor * selection
                else:
                    dividing_factor = dividing_factor * int(dataParts[1])

                used_data = [x / dividing_factor for x in curr_data]

                mean = np.mean(used_data)
                std = np.std(used_data)
                sqrt_len = sqrt(len(used_data))

                xs.append(x)
                ys.append(mean)
                yErrs.append(z_star * std / sqrt_len)
            except KeyError as e:
                pass

        if len(xs) > 1:
            line, _, _ = ax.errorbar(xs, ys, yerr=yErrs, capsize=7.0, label=legend, marker=marker, linestyle='-', color=color)

            if extraFit:
                popt, pcov = curve_fit(genericCube, xs, ys)
                beyondXs = np.linspace(xs[-1], 100, 50)
                ax.plot(beyondXs, genericCube(beyondXs, *popt), linestyle='--', color=color)
        
    ax.set_title(title, fontsize='x-large')
    ax.set_xlabel(xLabel, fontsize='large')
    ax.set_ylabel(yLabel, fontsize='large')

    if loglog:
        ax.set_yscale("log")
        ax.set_xscale("log")
    else:
        bottom, top = ax.get_ylim()
        bottom = (0 if bottom > 0 else bottom)
        ax.set_ylim(bottom=bottom)
        if top > 100000:
            yTickLabels = ['{:2g}'.format(x) for x in ax.get_yticks().tolist()]
            ax.set_yticklabels(yTickLabels)

    if yLim:
        ax.set_ylim(bottom=yLim[0], top=yLim[1])

    if aspect:
        ax.set_aspect(aspect, adjustable='box')

    legend = ax.legend(loc=legendLoc, bbox_to_anchor=legendBBoxAnchor, fontsize='large')

    with cd('../plt/', True):
        fig.savefig(f"{fileNameStr}.pdf", bbox_inches='tight')
    plt.close(fig)

def main(dataDirectory, plotOptionsFile):
    serverData, clientData = readData(dataDirectory)

    plotOptions = []
    with open(plotOptionsFile, 'r') as options:
        plotOptions = json.load(options)

    for option in plotOptions:
        try:
            data = serverData if (option['data'].lower() == "server" or option['data'].lower() == "s") else clientData
            dataParts = option['dataParts']
            xVariable = option['xVariable']
            lineVariable = option['lineVariable']
            whichGraph = option['whichGraph']
        except KeyError as e:
            continue
        
        kwargs = {}
        if "z_star" in option:
            kwargs["z_star"] = option["z_star"]
        if "title" in option:
            kwargs["title"] = option["title"]
        if "xLabel" in option:
            kwargs["xLabel"] = option["xLabel"]
        if "yLabel" in option:
            kwargs["yLabel"] = option["yLabel"]
        if "yUnit" in option:
            kwargs["yUnit"] = option["yUnit"]
        if "extraFit" in option:
            kwargs["extraFit"] = option["extraFit"]
        if "fileNameStr" in option:
            kwargs["fileNameStr"] = option["fileNameStr"]
        if "legendLoc" in option:
            kwargs["legendLoc"] = option["legendLoc"]
        if "legendBBoxAnchor" in option:
            anchor = (option["legendBBoxAnchor"][0], option["legendBBoxAnchor"][1])
            kwargs["legendBBoxAnchor"] = anchor
        if "loglog" in option:
            kwargs["loglog"] = option["loglog"]
        if "yLim" in option:
            kwargs["yLim"] = option["yLim"]
        if "aspect" in option:
            kwargs["aspect"] = option["aspect"]
        if "ignoreWorkload" in option:
            kwargs["ignoreWorkload"] = option["ignoreWorkload"]

        plotComparison(data, dataParts, xVariable, lineVariable, whichGraph, **kwargs)

if __name__ == "__main__":
    main("../out", "../plt/plots.json")