trophPlot v3

from matplotlib import pyplot as plt
import numpy as np
import os
from glob import glob
import csv
import matplotlib.cm as cm
from collections import OrderedDict
import ast
import re
import fnmatch

# data decription for plot titles
dataSetName = "ScavengeTroph2017-05-28_"

plot_title = "Trophallaxis Activity: Trophallaxis_Generous_Random & Levy"
'''
# the path to the data from where the current file is saved
dirname = os.path.abspath(os.path.join(
    os.path.dirname( __file__ ),
    '../../../..',
    "self_energy_robot"
    + "/" + "steptaken_infinitefood"
    ))

path = dirname + "/" + dataSetName
#print(path)
extension   = "*Summary.csv"
files = []

for dir,_,_ in os.walk(path):
    files.extend(glob(os.path.join(dir,extension)))
'''

data_to_import = ['count',
        'food_patches',
        'trophallaxis']

# 'Boundary' or 'Food'
findType = 'Boundary'
# use the following variable if you want to group by Boundary#
ngroups = 5
# use the following variable if you want to group by Food##
group_types = ['25', '81']

# ->
# 1 for Default avg area plot
# 2 for Normalised area plot (Avg / (food/bots))
drawY = 1

# ->
# 0 don't draw error bars
# 1 draw error bars
drawError = 0

#
path = []
root = []
testNumber = 'test14'
#Levy_No_Trophallaxis Levy_Trophallaxis_Generous Random_No_Trophallaxis Random_Trophallaxis_Generous
UTfolder = ['Random_Trophallaxis_Generous', 'Levy_Trophallaxis_Generous']
trophFolder = 'test14Finitefood'
#cwd='/Users/ammacpro1/AnacondaProjects/self_energy_robot/'+testNumber+'/'+UTfolder+'/'+trophFolder
cwd='/Users/ammacpro1/simulation_results/self_energy_robot/'+testNumber
#cwd = os.getcwd()
#cwd = cwd + "/" + testNumber
#root = cwd+'/'+UTfolder[0]+'/'+trophFolder
for f in UTfolder:
    path.append(cwd+'/'+f+'/'+trophFolder+'/Bot*/*Timestep.csv')
    root.append(cwd+'/'+f+'/'+trophFolder)
#init_food_per_groups = [32, 40, 48, 56, 64]
#

fig, ax = plt.subplots()

def findGroupsB(root, ngroups):
    # currently works only for a single UTfolder at once
    #groups = {}
    groups = []
    for root_ut in root:
        for g in range(ngroups):
            #groups["Group" + str(g+1)] = glob(root + '/*Boundary' + str(g+1) + 'Max*')
            #groups.append(glob(root + '/*Boundary' + str(g+1) + 'Max*/*Timestep.csv'))
            groups.append(root_ut + '/*Boundary' + str(g+1) + '*/*Timestep.csv')
    #print(groups)
    return groups

def findGroupsF(root, group_types):
    global ngroups
    ngroups = len(group_types)
    # currently works only for a single UTfolder at once
    #groups = {}
    groups = []
    for root_ut in root:
        for g in group_types:
            #groups["Group" + str(g+1)] = glob(root + '/*Boundary' + str(g+1) + 'Max*')
            #groups.append(glob(root + '/*Boundary' + str(g+1) + 'Max*/*Timestep.csv'))
            groups.append(root_ut + '/*Food' + g + '*/*Timestep.csv')
    #print(groups)
    return groups

def analyseFolders(folder, curr_p, drawY, drawError):
    files = glob(folder)
    data = {}
    for _data in data_to_import:
        data[_data] = []
    #indexesOfInterest = [0, 2, 7] # count, food_patches, trophallaxis
    for file in files:
        #print(file)
        with open(file, 'rt') as f:
            reader = csv.reader(f, delimiter=',')
            for row in reader:
                #print(row)
                if row[0].isdigit():
                    #currentIndex = 0
                    #for _data in data_to_import:
                        ##if row[0] in [_data]:
                        ##    data[_data].append(row[1])
                        #data[_data].append(row[currentIndex])
                        #currentIndex += 1
                    data[data_to_import[0]].append(row[0])
                    data[data_to_import[1]].append(row[2])
                    data[data_to_import[2]].append(row[7])

    '''
    print(data["trophallaxis"])
    print(len(data))
    print(len(data[data_to_import[0]]))
    print(len(data[data_to_import[1]]))
    print(len(data[data_to_import[2]]))
    '''

    sorted_data = dict.fromkeys(data.keys(),[])
    data_to_plot = {}

    keys = [key for key, value in data.items()]
    #print(keys)

    timeSteps = sorted(list(set(data["count"])), key=int)
    foodPatchSize = sorted(list(set(data["food_patches"])), key=int)

    #timeSteps = [int(x) for x in timeSteps]
    #foodPatchSize= [int(x) for x in foodPatchSize]
    #print(foodPatchSize)
    #input("Press Enter to terminate.")

    '''
    for SSize in timeSteps:
        for PSize in foodPatchSize:
            data_to_plot["trophallaxis_ts" + SSize + "_food" + PSize] = \
                    [float(i)
                            for i, j, k in zip(
                                data["trophallaxis"],
                                data["count"],
                                data["food_patches"])
                            if j == SSize and k == PSize]
    '''
    for SSize in timeSteps:
        data_to_plot["trophallaxis_ts" + SSize] = \
                [float(i)
                        for i, j in zip(
                            data["trophallaxis"],
                            data["count"])
                        if j == SSize]

    #print(data_to_plot)
    #input("Press Enter to terminate.")

    #colors = iter(cm.Greys(np.linspace(0.4, 1, len(foodPatchSize))))
    #colors = cm.rainbow(np.linspace(0.4, 1, len(UTfolder)))
    #widths = np.linspace(2, 0.5, len(UTfolder))
    colors = cm.rainbow(np.linspace(0.4, 1, ngroups))
    widths = np.linspace(2, 0.5, ngroups)
    c = colors[curr_p % ngroups]
    w = widths[curr_p % ngroups]

    #fig, ax = plt.subplots()
    #data_dict = {k:v for (k,v) in data_to_plot.items() if _data in k}

    Markers = ['+','s','^','o','p']
    LineTypes = ['-','--','-.',':']
    LineStyles = [LineTypes[r] for r in (range(len(root) % len(LineTypes))) for i in range(ngroups)]
    #MarkerIndex = 0;

    data_mean = []
    data_sd = []

    '''
    for SSize in timeSteps:
        for PSize in foodPatchSize:
            #food_dict = {k: v for (k, v) in data_dict.items() if PSize in k}
            vals = []
            sd = []
            if len(data_to_plot["trophallaxis_ts" + SSize + "_food" + PSize]) > 0:
                vals = np.mean(data_to_plot["trophallaxis_ts" + SSize + "_food" + PSize])
                if (drawError):
                    sd = np.std(data_to_plot["trophallaxis_ts" + SSize + "_food" + PSize])
                else:
                    sd = 0
            else:
                vals = 0
                sd = 0
            #if (drawY == 2):
            #    vals /= (int(PSize)/int(SSize))
            data_mean.append(vals)
            data_sd.append(sd)
    '''

    for SSize in timeSteps:
        vals = []
        sd = []
        if len(data_to_plot["trophallaxis_ts" + SSize]) > 0:
            vals = np.mean(data_to_plot["trophallaxis_ts" + SSize])
            if (drawError):
                sd = np.std(data_to_plot["trophallaxis_ts" + SSize])
            else:
                sd = 0
        else:
            vals = 0
            sd = 0
        #if (drawY == 2):
        #    vals /= (int(PSize)/int(SSize))
        data_mean.append(vals)
        data_sd.append(sd)

    timeSteps = [int(x) for x in timeSteps]
    foodPatchSize= [int(x) for x in foodPatchSize]

    #print(timeSteps)
    #print(foodPatchSize)
    #input("Press Enter to terminate.")

    ax.errorbar(timeSteps,
            data_mean,
            data_sd,
            #sd, fmt = 'o',
            #capsize = 2,
            linestyle = LineStyles[curr_p % len(LineStyles)],
            #linestyle = "None",
            linewidth = w,
            marker = Markers[curr_p % len(Markers)],
            #marker='o',
            color=c,
            #label="Number of food patches = " + str(PSize))
            #label=UTfolder[curr_p])
            label=("Group " + str((curr_p // ngroups) +1) + "-" + str((curr_p % ngroups) + 1)))

    ax.set_title(plot_title)
    ax.legend(frameon=False, loc="best", prop={'size': 10})
    ax.set_xlabel("Timesteps")
    if (drawY == 1):
        ax.set_ylabel("Average Trophallaxis")
    #elif (drawY == 2):
    #    ax.set_ylabel("Normalised Average Area: Avg/(food/bots)")
    #ax.set_title(plot_title)

    #plotname = path + '/' + plot_title + "_" + _data
    #fig.savefig(plotname + '.pdf', orientation='landscape')
    #fig.savefig(plotname + '.png', orientation='landscape')
    #plt.show()

curr_p = 0
if findType == 'Boundary':
    groups = findGroupsB(root, ngroups)
elif findType == 'Food':
    groups = findGroupsF(root, group_types)
#print(groups)
for group in groups:
    #for p in path:
    analyseFolders(group, curr_p, drawY, drawError)
    curr_p += 1

plt.show()