<commit_before>#!/usr/bin/python
import argparse
import matplotlib.pyplot as plt
import pandas as pd
import matplotlib.animation as animation
import numpy as np
from itertools import tee

parser = argparse.ArgumentParser(description='Plot data from output of the black hole simulation.')
parser.add_argument('--video', action='store_true',
                    help='save a video (MUCH slower)')
parser.add_argument('--dont-loop', action='store_false', dest='loop', default=True,
                    help='prevent from looping infinitely')
parser.add_argument('--video-file', metavar='file', default='evolution.mp4',
                    help='save the video as (default: %(default)s).')
parser.add_argument('--plot', metavar='dump_id', nargs='*', type=int, default=[],
                    help='dump_ids to plot.')
parser.add_argument('--c-points', metavar='file',
                    help='Critical point file (default: %(default)s).',
                    default='critical_points/file.dat')
parser.add_argument('--i-conditions', metavar='file',
                    help='Initial conditions file (default: %(default)s).', default='CI.dat')
parser.add_argument('--s-curves', metavar='n', nargs='+', type=int,
                    help='S curves to plot (default: %(default)s).', default=[1, 10, 20, 255])
parser.add_argument('--s-curves-dir', metavar='dir', nargs=1,
                    help='S curves directory (default: %(default)s).', default='s_curves')
parser.add_argument('--reread', action='store_true', default=False,
                    help='Reread file when fully read.')
parser.add_argument('--interval', default=1000, type=int,
                    help='Interval (in ms) between each frames (default %(default)s)')
parser.add_argument('--plot-Qs', default=False, action='store_true', dest='plot_Qs',
                    help='Plot Q⁺ and Q⁻ instead of Mdot')
args = parser.parse_args()

fig, ((ax11, ax12), (ax21, ax22)) = plt.subplots(2, 2)

class Data:
    ''' Class that opens filename and reads it. You can directly iterate over it.'''
    def __init__(self, filename):
        self.filename = filename # holds the filename
        self.data = {}           # holds the file data
        self.times = {}          # holds the data
        self.indexes = []        # holds a list of each snapshot
        self.restrictTo = []     # list of indexes to restrict to

        self.file = open(self.filename, 'r')
        # read ahead one line
        self.lastLineRead = self.file.readline()
        self.pause = False       # bool flag if paused (activates left<>right navigation
        self.reread = False      #
        self.offset = 0          # Number of indexes to jump back or forth

    def getChunk(self):
        '''Get a chunk (niter, time, headers + data)
        starting at last line read.
        In the end, reloads the last line read and not
        returned into self.lastlineread'''

        f = self.file

        ###########################################
        # Check not end of file
        ###########################################
        if self.lastLineRead == '':
            return -1, -1, [], [], True

        ###########################################
        # preload first lines
        ###########################################
        niter = int(self.lastLineRead.split('#')[1])
        time = float(f.readline().split('#')[1])
        headers = f.readline().split()
        eof = False # flag to set if end of file

        ###########################################
        # read until reaching a commented line
        ###########################################
        data = []
        line = f.readline()
        while '#' not in line:
            # if eof reached
            if line == '':
                eof = True
                break
            # remove last character ('\n')
            data.append(line[:-1].split())

            line = f.readline()

        pdData = pd.DataFrame(data=data, columns=headers, dtype=float)
        # store the last line read
        self.lastLineRead = line
        return niter, time, headers, pdData, eof

    def getPosition(self, _from, offset):
        ''' Given an initial position, returns the
        index of the n'th position after and the effective offset applied.'''
        indexOfNiter = self.indexes.index(_from)

        tmpIndex = indexOfNiter - 1 + offset
        if tmpIndex >= len(self.indexes):
            effectiveOffset = len(self.indexes) - indexOfNiter
        elif tmpIndex < 0:
            effectiveOffset = -indexOfNiter + 1
        else:
            effectiveOffset = offset

        index = indexOfNiter - 1 + effectiveOffset
        return effectiveOffset, self.indexes[index]

    def __iter__(self):
        ''' Infinite iterator over the whole file'''
        self.restrictTo.sort()

        # Convert the restriction into a boolean
        if len(self.restrictTo) > 0:
            noRestriction = False
        else:
            noRestriction = True

        ###########################################
        # First read the whole file chunk by chunk
        ###########################################
        i = 0
        niter, time, headers, data, eof = self.getChunk()
        while not eof:
            if self.pause:
                self.offset, i = self.getPosition(niter, self.offset)
                print('i:', self.indexes, self.offset)

                yield (i, self.times[i], self.data[i])
            elif noRestriction or niter in self.restrictTo:
                yield (niter, time, data)

            ###########################################
            # load the missing lines ahead
            ###########################################
            readAhead = 0
            if self.pause and self.offset:
                readAhead += self.offset
            elif not self.pause:
                readAhead = 1

            while readAhead > 0 and not eof:
                print('reading ahead')
                # if we have a restriction list and the current iteration
                # is after the last accepted item, stop reading the file
                if not noRestriction and niter > self.restrictTo[-1]:
                    break

                niter, time, headers, data, eof = self.getChunk()
                self.data[niter] = data
                self.times[niter] = time
                self.indexes.append(niter)
                self.indexes.sort()

                self.offset = 0

                readAhead -= 1
                # if reread flag activated, close and reopen the file
                if eof and self.reread:
                    self.file.close()
                    self.file = open(self.filename, 'r')
                    self.lastLineRead = self.file.readline()
                    eof = False

        ###########################################
        # If looping, pop back the data already read
        ###########################################
        if self.loop:
            # dataKeys holds the keys where to find the data
            if noRestriction:
                dataKeys = list(self.data.keys())
            else:
                dataKeys = [key for key in self.data.keys() if key in self.restrictTo]

            dataKeys.sort()
            while True:
                prevKey = dataKeys[-1]
                for key in dataKeys:
                    if self.pause:
                        self.offset, i = self.getPosition(niter, self.offset)
                        yield (i, self.times[i], self.data[i])
                    else:
                        yield (key, self.times[key], self.data[key])
                    prevKey = key

    def get(self):
        ''' Get a chunk of the file without moving the position of the reader in the file '''
        pos = self.file.tell()
        niter, time, headers, data, eof = self.getChunk()
        self.file.seek(pos)
        return niter, time, headers, data

    def __del__(self):
        ''' Close the file reader'''
        self.file.close()

lines = {'r-T': 0,
         'r-Sigma': 0,
         'Sigma-T': []}

class colorLooper:
    def __init__(self, colors= ['red', 'green', 'blue', 'grey', 'purple', 'brown', 'pink']):
        self.colors = colors
        self.i = 0

    def __next__(self):
        self.i += 1
        return self.colors[self.i - 1 % len(self.colors)]

    def __iter__(self):
        self.i = 0
        while self.i < len(self.colors):
            yield self.__next__()

    def reset(self):
        self.i = 0

def init(ic, crit_pts, s_curves, initial_data):
    ''' Plot the initial conditions, the critical points and the s_curves'''
    ## Top left panel
    ax11.plot(ic['r'], ic['T'], '--', label='Initial conditions')
    ax11.plot(ic['r'], crit_pts['Temp_thick'], '--', label='critical')

    ax11.set_xlabel('$r\ (cm)$')
    ax11.set_ylabel('$T\ (K)$')

    lines['r-T'] = ax11.plot(initial_data['r'], initial_data['T'])[0]

    ax11.set_yscale('log')
    ax11.grid()
    # ax11.legend()

    # add ticks corresponding to the s_curve
    ticks      = list(ax11.get_xticks())
    ticklabels = list(ax11.get_xticklabels())
    for ind, s_curve in s_curves:
        ticks.append(ic['r'][ind])
        ticklabels.append('$r_{'+str(ind)+'}$')

    ax11.set_xticks(ticks)
    ax11.set_xticklabels(ticklabels)

    ## Top right panel
    ax12.set_xlabel('$r\ (\mathrm{cm})$')
    ax12.set_ylabel('$\Sigma\ (\mathrm{g.cm^{-2}})$')
    ax12.plot(ic['r'], ic['Sigma'], '--')
    ax12.plot(ic['r'], crit_pts['Sigma_thick'], '--')

    lines['r-Sigma'] = ax12.plot(initial_data['r'], initial_data['Sigma'])[0]

    ax12.grid()
    ax12.set_yscale('log')

    if args.plot_Qs:
        lines['r-Q+'] = ax21.plot(initial_data['r'], initial_data['Q_+'], label='$Q^+$')[0]
        lines['r-Q-'] = ax21.plot(initial_data['r'], initial_data['Q_-'], label='$Q^-$')[0]
        ax21.legend()
    else:
        lines['r-Mdot'] = ax21.plot(initial_data['r'], initial_data['M_dot'])[0]

    # add ticks corresponding to the s_curve
    ticks      = list(ax12.get_xticks())
    ticklabels = list(ax12.get_xticklabels())
    for ind, s_curve in s_curves:
        ticks.append(ic['r'][ind])
        ticklabels.append('$r_{'+str(ind)+'}$')

    ax12.set_xticks(ticks)
    ax12.set_xticklabels(ticklabels)

    ## Bottom left panel
    ax21.set_xlabel('$r\ (\mathrm{cm})$')
    if args.plot_Qs:
        ax21.set_ylabel('$Q$')
    else:
        ax21.set_ylabel('$\dot{M}\ (\mathrm{g.s^{-1}})$')
    ax21.grid()

    ## Bottom right panel
    colorsIter = colorLooper()
    for ind, s_curve in s_curves:
        ax22.plot(s_curve['Surface_density'], s_curve['Temperature'],
                  label='$r_{'+str(ind)+'}$',
                  c=colorsIter.__next__())

    # add the initial data on the s_curve plot
    colorsIter.reset()
    lines['Sigma-T'] = [ (ind, ax22.plot(initial_data['Sigma'][ind],
                                         initial_data['T'][ind], 'o',
                                         c=colorsIter.__next__())[0])
                     for ind, foo in s_curves]

    ax22.set_xlabel('$\Sigma\ (\mathrm{g.cm^{-2}})$')
    ax22.set_ylabel('$T\ (\mathrm{K})$')
    ax22.legend(loc='upper right')
    ax22.grid()

prevTime = 0
prevIndex = 0
prevDt = 0

def plotData(plotArgs):
    ''' Update with the new data. Parameters:
    args: array that contains
      index: the indexes of the datadump
      data : the data
    '''
    global prevDt
    global prevTime

    index, time, data = plotArgs

    print('Plotting {}'.format(index))
    lines['r-T'].set_ydata(data['T'])
    # lines['r-T'].set_label('$t = {:.2}s$'.format(time))

    if args.plot_Qs:
        lines['r-Q+'].set_ydata(data['Q_+'])
        lines['r-Q-'].set_ydata(data['Q_-'])
    else:
        lines['r-Mdot'].set_ydata(data['M_dot'])

    lines['r-Sigma'].set_ydata(data['Sigma'])

    newDt = time - prevTime
    prevTime = time
    if newDt == 0:
        dt = prevDt
    else:
        dt = newDt
        prevDt = dt

    print('dt: ', dt)
    for ind, line in lines['Sigma-T']:
        x = data['Sigma'][ind]
        y = data['T'][ind]
        line.set_xdata(x)
        line.set_ydata(y)

        print(ind, dt * (data['Q_+'][ind] - data['Q_-'][ind]) / data['Cv'][ind])

    # ax11.legend()
    fig.suptitle('$t = {:.2f}s$, iteration {}'.format(time, index))

def onClick(data, event):
    data.pause ^= True
    data.offset = 0
    print(data.pause)

def onKey(data, event):
    def mod(amount):
        data.offset += amount

    behaviour = {
        'down': lambda: mod(-1),
        'up': lambda: mod(+1),
        ' ': lambda: onClick(data, ''),
        'ctrl+down': lambda: mod(-10),
        'ctrl+up': lambda: mod(+10),
        'shift+ctrl+down': lambda: mod(-1000),
        'shift+ctrl+up': lambda: mod(+1000),
        }
    if event.key in behaviour:
        behaviour[event.key]()
    else:
        print('Unsupported key "{}"'.format(event.key))

if __name__ == '__main__':
    print('Reading initial conditions')
    ic = pd.read_csv('CI.dat', delim_whitespace=True)

    print('Reading output file')
    simulData = Data('output.dat')

    print('Reading critical points')
    crit_pts = pd.read_csv(args.c_points, delim_whitespace=True)

    print('Reading S_curves')
    s_curves = [ (ind,
                  pd.read_csv(args.s_curves_dir + '/Temperature_Sigma_{:0>5}_tot.dat'.format(ind),
                              delim_whitespace=True, dtype=float, header=0))
                 for ind in args.s_curves]

    data0 = simulData.get()[-1]
    initFun = lambda: init(ic, crit_pts, s_curves, data0)

    if len(args.plot) == 1:
        initFun()
        plotData((0, 0, data0))
        plt.show()
    else:
        simulData.loop = args.loop
        simulData.reRead = args.reread

        if (len(args.plot) > 0):
            simulData.restrictTo = args.plot

        data = simulData

        onClickHandler = lambda event: onClick(data, event)
        onKeyHandler = lambda event: onKey(data, event)

        # fig.canvas.mpl_connect('button_press_event', onClickHandler)
        fig.canvas.mpl_connect('key_press_event', onKeyHandler)
        ani = animation.FuncAnimation(fig, plotData, data, init_func=initFun,
                                      interval=args.interval)
        if args.video:
            # temporaly deactivate looping for the video
            simulData.loop = False
            ani.save(args.video_file, writer='ffmpeg', fps=10, bitrate=10000, dpi=180)
            simulData.loop = args.loop

        # clear the axes
        for ax in (ax11, ax12, ax21, ax22):
            ax.axes.cla()

        plt.show()
<commit_msg>plot_data: better support for python2<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-
# necessary for Python2
from __future__ import print_function
import argparse
import matplotlib.pyplot as plt
import pandas as pd
import matplotlib.animation as animation
import numpy as np
from itertools import tee

parser = argparse.ArgumentParser(description='Plot data from output of the black hole simulation.')
parser.add_argument('--video', action='store_true',
                    help='save a video (MUCH slower)')
parser.add_argument('--dont-loop', action='store_false', dest='loop', default=True,
                    help='prevent from looping infinitely')
parser.add_argument('--video-file', metavar='file', default='evolution.mp4',
                    help='save the video as (default: %(default)s).')
parser.add_argument('--plot', metavar='dump_id', nargs='*', type=int, default=[],
                    help='dump_ids to plot.')
parser.add_argument('--c-points', metavar='file',
                    help='Critical point file (default: %(default)s).',
                    default='critical_points/file.dat')
parser.add_argument('--i-conditions', metavar='file',
                    help='Initial conditions file (default: %(default)s).', default='CI.dat')
parser.add_argument('--s-curves', metavar='n', nargs='+', type=int,
                    help='S curves to plot (default: %(default)s).', default=[1, 10, 20, 255])
parser.add_argument('--s-curves-dir', metavar='dir', nargs=1,
                    help='S curves directory (default: %(default)s).', default='s_curves')
parser.add_argument('--reread', action='store_true', default=False,
                    help='Reread file when fully read.')
parser.add_argument('--interval', default=1000, type=int,
                    help='Interval (in ms) between each frames (default %(default)s)')
parser.add_argument('--plot-Qs', default=False, action='store_true', dest='plot_Qs',
                    help='Plot Q⁺ and Q⁻ instead of Mdot')
args = parser.parse_args()

fig, ((ax11, ax12), (ax21, ax22)) = plt.subplots(2, 2)

class Data:
    ''' Class that opens filename and reads it. You can directly iterate over it.'''
    def __init__(self, filename):
        self.filename = filename # holds the filename
        self.data = {}           # holds the file data
        self.times = {}          # holds the data
        self.indexes = []        # holds a list of each snapshot
        self.restrictTo = []     # list of indexes to restrict to

        self.file = open(self.filename, 'r')
        # read ahead one line
        self.lastLineRead = self.file.readline()
        self.pause = False       # bool flag if paused (activates left<>right navigation
        self.reread = False      #
        self.offset = 0          # Number of indexes to jump back or forth

    def getChunk(self):
        '''Get a chunk (niter, time, headers + data)
        starting at last line read.
        In the end, reloads the last line read and not
        returned into self.lastlineread'''

        f = self.file

        ###########################################
        # Check not end of file
        ###########################################
        if self.lastLineRead == '':
            return -1, -1, [], [], True

        ###########################################
        # preload first lines
        ###########################################
        niter = int(self.lastLineRead.split('#')[1])
        time = float(f.readline().split('#')[1])
        headers = f.readline().split()
        eof = False # flag to set if end of file

        ###########################################
        # read until reaching a commented line
        ###########################################
        data = []
        line = f.readline()
        while '#' not in line:
            # if eof reached
            if line == '':
                eof = True
                break
            # remove last character ('\n')
            data.append(line[:-1].split())

            line = f.readline()

        pdData = pd.DataFrame(data=data, columns=headers, dtype=float)
        # store the last line read
        self.lastLineRead = line
        return niter, time, headers, pdData, eof

    def getPosition(self, _from, offset):
        ''' Given an initial position, returns the
        index of the n'th position after and the effective offset applied.'''
        indexOfNiter = self.indexes.index(_from)

        tmpIndex = indexOfNiter - 1 + offset
        if tmpIndex >= len(self.indexes):
            effectiveOffset = len(self.indexes) - indexOfNiter
        elif tmpIndex < 0:
            effectiveOffset = -indexOfNiter + 1
        else:
            effectiveOffset = offset

        index = indexOfNiter - 1 + effectiveOffset
        return effectiveOffset, self.indexes[index]

    def __iter__(self):
        ''' Infinite iterator over the whole file'''
        self.restrictTo.sort()

        # Convert the restriction into a boolean
        if len(self.restrictTo) > 0:
            noRestriction = False
        else:
            noRestriction = True

        ###########################################
        # First read the whole file chunk by chunk
        ###########################################
        i = 0
        niter, time, headers, data, eof = self.getChunk()
        while not eof:
            if self.pause:
                self.offset, i = self.getPosition(niter, self.offset)
                print('i:', self.indexes, self.offset)

                yield (i, self.times[i], self.data[i])
            elif noRestriction or niter in self.restrictTo:
                yield (niter, time, data)

            ###########################################
            # load the missing lines ahead
            ###########################################
            readAhead = 0
            if self.pause and self.offset:
                readAhead += self.offset
            elif not self.pause:
                readAhead = 1

            while readAhead > 0 and not eof:
                print('reading ahead')
                # if we have a restriction list and the current iteration
                # is after the last accepted item, stop reading the file
                if not noRestriction and niter > self.restrictTo[-1]:
                    break

                niter, time, headers, data, eof = self.getChunk()
                self.data[niter] = data
                self.times[niter] = time
                self.indexes.append(niter)
                self.indexes.sort()

                self.offset = 0

                readAhead -= 1
                # if reread flag activated, close and reopen the file
                if eof and self.reread:
                    self.file.close()
                    self.file = open(self.filename, 'r')
                    self.lastLineRead = self.file.readline()
                    eof = False

        ###########################################
        # If looping, pop back the data already read
        ###########################################
        if self.loop:
            # dataKeys holds the keys where to find the data
            if noRestriction:
                dataKeys = list(self.data.keys())
            else:
                dataKeys = [key for key in self.data.keys() if key in self.restrictTo]

            dataKeys.sort()
            while True:
                prevKey = dataKeys[-1]
                for key in dataKeys:
                    if self.pause:
                        self.offset, i = self.getPosition(niter, self.offset)
                        yield (i, self.times[i], self.data[i])
                    else:
                        yield (key, self.times[key], self.data[key])
                    prevKey = key

    def get(self):
        ''' Get a chunk of the file without moving the position of the reader in the file '''
        pos = self.file.tell()
        niter, time, headers, data, eof = self.getChunk()
        self.file.seek(pos)
        return niter, time, headers, data

    def __del__(self):
        ''' Close the file reader'''
        self.file.close()

lines = {'r-T': 0,
         'r-Sigma': 0,
         'Sigma-T': []}

class colorLooper:
    def __init__(self, colors= ['red', 'green', 'blue', 'grey', 'purple', 'brown', 'pink']):
        self.colors = colors
        self.i = 0

    def __next__(self):
        self.i += 1
        return self.colors[self.i - 1 % len(self.colors)]

    def __iter__(self):
        self.i = 0
        while self.i < len(self.colors):
            yield self.__next__()

    def reset(self):
        self.i = 0

def init(ic, crit_pts, s_curves, initial_data):
    ''' Plot the initial conditions, the critical points and the s_curves'''
    ## Top left panel
    ax11.plot(ic['r'], ic['T'], '--', label='Initial conditions')
    ax11.plot(ic['r'], crit_pts['Temp_thick'], '--', label='critical')

    ax11.set_xlabel('$r\ (cm)$')
    ax11.set_ylabel('$T\ (K)$')

    lines['r-T'] = ax11.plot(initial_data['r'], initial_data['T'])[0]

    ax11.set_yscale('log')
    ax11.grid()
    # ax11.legend()

    # add ticks corresponding to the s_curve
    ticks      = list(ax11.get_xticks())
    ticklabels = list(ax11.get_xticklabels())
    for ind, s_curve in s_curves:
        ticks.append(ic['r'][ind])
        ticklabels.append('$r_{'+str(ind)+'}$')

    ax11.set_xticks(ticks)
    ax11.set_xticklabels(ticklabels)

    ## Top right panel
    ax12.set_xlabel('$r\ (\mathrm{cm})$')
    ax12.set_ylabel('$\Sigma\ (\mathrm{g.cm^{-2}})$')
    ax12.plot(ic['r'], ic['Sigma'], '--')
    ax12.plot(ic['r'], crit_pts['Sigma_thick'], '--')

    lines['r-Sigma'] = ax12.plot(initial_data['r'], initial_data['Sigma'])[0]

    ax12.grid()
    ax12.set_yscale('log')

    if args.plot_Qs:
        lines['r-Q+'] = ax21.plot(initial_data['r'], initial_data['Q_+'], label='$Q^+$')[0]
        lines['r-Q-'] = ax21.plot(initial_data['r'], initial_data['Q_-'], label='$Q^-$')[0]
        ax21.legend()
    else:
        lines['r-Mdot'] = ax21.plot(initial_data['r'], initial_data['M_dot'])[0]

    # add ticks corresponding to the s_curve
    ticks      = list(ax12.get_xticks())
    ticklabels = list(ax12.get_xticklabels())
    for ind, s_curve in s_curves:
        ticks.append(ic['r'][ind])
        ticklabels.append('$r_{'+str(ind)+'}$')

    ax12.set_xticks(ticks)
    ax12.set_xticklabels(ticklabels)

    ## Bottom left panel
    ax21.set_xlabel('$r\ (\mathrm{cm})$')
    if args.plot_Qs:
        ax21.set_ylabel('$Q$')
    else:
        ax21.set_ylabel('$\dot{M}\ (\mathrm{g.s^{-1}})$')
    ax21.grid()

    ## Bottom right panel
    colorsIter = colorLooper()
    for ind, s_curve in s_curves:
        ax22.plot(s_curve['Surface_density'], s_curve['Temperature'],
                  label='$r_{'+str(ind)+'}$',
                  c=colorsIter.__next__())

    # add the initial data on the s_curve plot
    colorsIter.reset()
    lines['Sigma-T'] = [ (ind, ax22.plot(initial_data['Sigma'][ind],
                                         initial_data['T'][ind], 'o',
                                         c=colorsIter.__next__())[0])
                     for ind, foo in s_curves]

    ax22.set_xlabel('$\Sigma\ (\mathrm{g.cm^{-2}})$')
    ax22.set_ylabel('$T\ (\mathrm{K})$')
    ax22.legend(loc='upper right')
    ax22.grid()

prevTime = 0
prevIndex = 0
prevDt = 0

def plotData(plotArgs):
    ''' Update with the new data. Parameters:
    args: array that contains
      index: the indexes of the datadump
      data : the data
    '''
    global prevDt
    global prevTime

    index, time, data = plotArgs

    print('Plotting {}'.format(index))
    lines['r-T'].set_ydata(data['T'])
    # lines['r-T'].set_label('$t = {:.2}s$'.format(time))

    if args.plot_Qs:
        lines['r-Q+'].set_ydata(data['Q_+'])
        lines['r-Q-'].set_ydata(data['Q_-'])
    else:
        lines['r-Mdot'].set_ydata(data['M_dot'])

    lines['r-Sigma'].set_ydata(data['Sigma'])

    newDt = time - prevTime
    prevTime = time
    if newDt == 0:
        dt = prevDt
    else:
        dt = newDt
        prevDt = dt

    print('dt: ', dt)
    for ind, line in lines['Sigma-T']:
        x = data['Sigma'][ind]
        y = data['T'][ind]
        line.set_xdata(x)
        line.set_ydata(y)

        print(ind, dt * (data['Q_+'][ind] - data['Q_-'][ind]) / data['Cv'][ind])

    # ax11.legend()
    fig.suptitle('$t = {:.2f}s$, iteration {}'.format(time, index))

def onClick(data, event):
    data.pause ^= True
    data.offset = 0
    print(data.pause)

def onKey(data, event):
    def mod(amount):
        data.offset += amount

    behaviour = {
        'down': lambda: mod(-1),
        'up': lambda: mod(+1),
        ' ': lambda: onClick(data, ''),
        'ctrl+down': lambda: mod(-10),
        'ctrl+up': lambda: mod(+10),
        'shift+ctrl+down': lambda: mod(-1000),
        'shift+ctrl+up': lambda: mod(+1000),
        }
    if event.key in behaviour:
        behaviour[event.key]()
    else:
        print('Unsupported key "{}"'.format(event.key))

if __name__ == '__main__':
    print('Reading initial conditions')
    ic = pd.read_csv('CI.dat', delim_whitespace=True)

    print('Reading output file')
    simulData = Data('output.dat')

    print('Reading critical points')
    crit_pts = pd.read_csv(args.c_points, delim_whitespace=True)

    print('Reading S_curves')
    s_curves = [ (ind,
                  pd.read_csv(args.s_curves_dir + '/Temperature_Sigma_{:0>5}_tot.dat'.format(ind),
                              delim_whitespace=True, dtype=float, header=0))
                 for ind in args.s_curves]

    data0 = simulData.get()[-1]
    initFun = lambda: init(ic, crit_pts, s_curves, data0)

    if len(args.plot) == 1:
        initFun()
        plotData((0, 0, data0))
        plt.show()
    else:
        simulData.loop = args.loop
        simulData.reRead = args.reread

        if (len(args.plot) > 0):
            simulData.restrictTo = args.plot

        data = simulData

        onClickHandler = lambda event: onClick(data, event)
        onKeyHandler = lambda event: onKey(data, event)

        # fig.canvas.mpl_connect('button_press_event', onClickHandler)
        fig.canvas.mpl_connect('key_press_event', onKeyHandler)
        ani = animation.FuncAnimation(fig, plotData, data, init_func=initFun,
                                      interval=args.interval)
        if args.video:
            # temporaly deactivate looping for the video
            simulData.loop = False
            ani.save(args.video_file, writer='ffmpeg', fps=10, bitrate=10000, dpi=180)
            simulData.loop = args.loop

        # clear the axes
        for ax in (ax11, ax12, ax21, ax22):
            ax.axes.cla()

        plt.show()
<|endoftext|>


<commit_before>#!/usr/bin/python
import argparse
import matplotlib.pyplot as plt
import pandas as pd
import matplotlib.animation as animation
import numpy as np
from itertools import tee

parser = argparse.ArgumentParser(description='Plot data from output of the black hole simulation.')
parser.add_argument('--video', action='store_true',
                    help='save a video (MUCH slower)')
parser.add_argument('--dont-loop', action='store_false', dest='loop', default=True,
                    help='prevent from looping infinitely')
parser.add_argument('--video-file', metavar='file', default='evolution.mp4',
                    help='save the video as (default: %(default)s).')
parser.add_argument('--plot', metavar='dump_id', nargs='*', type=int, default=[],
                    help='dump_ids to plot.')
parser.add_argument('--c-points', metavar='file',
                    help='Critical point file (default: %(default)s).',
                    default='critical_points/file.dat')
parser.add_argument('--i-conditions', metavar='file',
                    help='Initial conditions file (default: %(default)s).', default='CI.dat')
parser.add_argument('--s-curves', metavar='n', nargs='+', type=int,
                    help='S curves to plot (default: %(default)s).', default=[1, 10, 20, 255])
parser.add_argument('--s-curves-dir', metavar='dir', nargs=1,
                    help='S curves directory (default: %(default)s).', default='s_curves')
parser.add_argument('--reread', action='store_true', default=False,
                    help='Reread file when fully read.')
parser.add_argument('--interval', default=0, type=int,
                    help='Interval (in ms) between each frames (default %(default)s)')
parser.add_argument('--plot-Qs', default=False, action='store_true', dest='plot_Qs',
                    help='Plot Q⁺ and Q⁻ instead of Mdot')
args = parser.parse_args()

fig, ((ax11, ax12), (ax21, ax22)) = plt.subplots(2, 2)

class Data:
    ''' Class that opens filename and reads it. You can directly iterate over it.'''
    def __init__(self, filename):
        self.filename = filename # holds the filename
        self.data = {}           # holds the file data
        self.times = {}          # holds the data
        self.indexes = []        # holds a list of each snapshot
        self.restrictTo = []     # list of indexes to restrict to

        self.file = open(self.filename, 'r')
        # read ahead one line
        self.lastLineRead = self.file.readline()
        self.pause = False       # bool flag if paused (activates left<>right navigation
        self.reread = False      #
        self.offset = None       # Number of indexes to jump back or forth

    def getChunk(self):
        '''Get a chunk (niter, time, headers + data)
        starting at last line read.
        In the end, reloads the last line read and not
        returned into self.lastlineread'''

        f = self.file

        ###########################################
        # Check not end of file
        ###########################################
        if self.lastLineRead == '':
            return -1, -1, [], [], True

        ###########################################
        # preload first lines
        ###########################################
        niter = int(self.lastLineRead.split('#')[1])
        time = float(f.readline().split('#')[1])
        headers = f.readline().split()
        eof = False # flag to set if end of file

        ###########################################
        # read until reaching a commented line
        ###########################################
        data = []
        line = f.readline()
        while '#' not in line:
            # if eof reached
            if line == '':
                eof = True
                break
            # remove last character ('\n')
            data.append(line[:-1].split())

            line = f.readline()

        pdData = pd.DataFrame(data=data, columns=headers, dtype=float)
        # store the last line read
        self.lastLineRead = line
        return niter, time, headers, pdData, eof

    def getPosition(self, _from, offset):
        ''' Given an initial position, returns the
        index of the n'th position after and the effective offset applied.'''
        indexOfNiter = self.indexes.index(_from)

        tmpIndex = indexOfNiter - 1 + offset
        if tmpIndex >= len(self.indexes):
            effectiveOffset = len(self.indexes) - indexOfNiter
        elif tmpIndex < 0:
            effectiveOffset = -indexOfNiter + 1
        else:
            effectiveOffset = offset

        index = indexOfNiter - 1 + effectiveOffset
        return effectiveOffset, self.indexes[index]

    def __iter__(self):
        ''' Infinite iterator over the whole file'''
        self.restrictTo.sort()

        # Convert the restriction into a boolean
        if len(self.restrictTo) > 0:
            noRestriction = False
        else:
            noRestriction = True

        ###########################################
        # First read the whole file chunk by chunk
        ###########################################
        i = 0
        niter, time, headers, data, eof = self.getChunk()
        while not eof:
            if self.pause:
                self.offset, i = self.getPosition(niter, self.offset)
                print('i:', self.indexes, self.offset)

                yield (i, self.times[i], self.data[i])
            elif noRestriction or niter in self.restrictTo:
                yield (niter, time, data)

            ###########################################
            # load the missing lines ahead
            ###########################################
            readAhead = 0
            if self.pause and self.offset:
                readAhead += self.offset
            elif not self.pause:
                readAhead = 1

            while readAhead > 0 and not eof:
                print('reading ahead')
                # if we have a restriction list and the current iteration
                # is after the last accepted item, stop reading the file
                if not noRestriction and niter > self.restrictTo[-1]:
                    break

                niter, time, headers, data, eof = self.getChunk()
                self.data[niter] = data
                self.times[niter] = time
                self.indexes.append(niter)
                self.indexes.sort()

                self.offset = 0

                readAhead -= 1
                # if reread flag activated, close and reopen the file
                if eof and self.reread:
                    self.file.close()
                    self.file = open(self.filename, 'r')
                    self.lastLineRead = self.file.readline()
                    eof = False

        ###########################################
        # If looping, pop back the data already read
        ###########################################
        if self.loop:
            # dataKeys holds the keys where to find the data
            if noRestriction:
                dataKeys = list(self.data.keys())
            else:
                dataKeys = [key for key in self.data.keys() if key in self.restrictTo]

            dataKeys.sort()
            while True:
                prevKey = dataKeys[-1]
                for key in dataKeys:
                    if self.pause:
                        self.offset, i = self.getPosition(niter, self.offset)
                        yield (i, self.times[i], self.data[i])
                    else:
                        yield (key, self.times[key], self.data[key])
                    prevKey = key

    def get(self):
        ''' Get a chunk of the file without moving the position of the reader in the file '''
        pos = self.file.tell()
        niter, time, headers, data, eof = self.getChunk()
        self.file.seek(pos)
        return niter, time, headers, data

    def __del__(self):
        ''' Close the file reader'''
        self.file.close()

lines = {'r-T': 0,
         'r-Sigma': 0,
         'Sigma-T': []}

class colorLooper:
    def __init__(self, colors= ['red', 'green', 'blue', 'grey', 'purple', 'brown', 'pink']):
        self.colors = colors
        self.i = 0

    def __next__(self):
        self.i += 1
        return self.colors[self.i - 1 % len(self.colors)]

    def __iter__(self):
        self.i = 0
        while self.i < len(self.colors):
            yield self.__next__()

    def reset(self):
        self.i = 0

def init(ic, crit_pts, s_curves, initial_data):
    ''' Plot the initial conditions, the critical points and the s_curves'''
    ## Top left panel
    ax11.plot(ic['r'], ic['T'], '--', label='Initial conditions')
    ax11.plot(ic['r'], crit_pts['Temp_thick'], '--', label='critical')

    ax11.set_xlabel('$r\ (cm)$')
    ax11.set_ylabel('$T\ (K)$')

    lines['r-T'] = ax11.plot(initial_data['r'], initial_data['T'])[0]

    ax11.set_yscale('log')
    ax11.grid()
    # ax11.legend()

    # add ticks corresponding to the s_curve
    ticks      = list(ax11.get_xticks())
    ticklabels = list(ax11.get_xticklabels())
    for ind, s_curve in s_curves:
        ticks.append(ic['r'][ind])
        ticklabels.append('$r_{'+str(ind)+'}$')

    ax11.set_xticks(ticks)
    ax11.set_xticklabels(ticklabels)

    ## Top right panel
    ax12.set_xlabel('$r\ (\mathrm{cm})$')
    ax12.set_ylabel('$\Sigma\ (\mathrm{g.cm^{-2}})$')
    ax12.plot(ic['r'], ic['Sigma'], '--')
    ax12.plot(ic['r'], crit_pts['Sigma_thick'], '--')

    lines['r-Sigma'] = ax12.plot(initial_data['r'], initial_data['Sigma'])[0]

    ax12.grid()
    ax12.set_yscale('log')

    if args.plot_Qs:
        lines['r-Q+'] = ax21.plot(initial_data['r'], initial_data['Q_+'], label='$Q^+$')[0]
        lines['r-Q-'] = ax21.plot(initial_data['r'], initial_data['Q_-'], label='$Q^-$')[0]
        ax21.legend()
    else:
        lines['r-Mdot'] = ax21.plot(initial_data['r'], initial_data['M_dot'])[0]

    # add ticks corresponding to the s_curve
    ticks      = list(ax12.get_xticks())
    ticklabels = list(ax12.get_xticklabels())
    for ind, s_curve in s_curves:
        ticks.append(ic['r'][ind])
        ticklabels.append('$r_{'+str(ind)+'}$')

    ax12.set_xticks(ticks)
    ax12.set_xticklabels(ticklabels)

    ## Bottom left panel
    ax21.set_xlabel('$r\ (\mathrm{cm})$')
    if args.plot_Qs:
        ax21.set_ylabel('$Q$')
    else:
        ax21.set_ylabel('$\dot{M}\ (\mathrm{g.s^{-1}})$')
    ax21.grid()

    ## Bottom right panel
    colorsIter = colorLooper()
    for ind, s_curve in s_curves:
        ax22.plot(s_curve['Surface_density'], s_curve['Temperature'],
                  label='$r_{'+str(ind)+'}$',
                  c=colorsIter.__next__())

    # add the initial data on the s_curve plot
    colorsIter.reset()
    lines['Sigma-T'] = [ (ind, ax22.plot(initial_data['Sigma'][ind],
                                         initial_data['T'][ind], 'o',
                                         c=colorsIter.__next__())[0])
                     for ind, foo in s_curves]

    ax22.set_xlabel('$\Sigma\ (\mathrm{g.cm^{-2}})$')
    ax22.set_ylabel('$T\ (\mathrm{K})$')
    ax22.legend(loc='upper right')
    ax22.grid()

prevTime = 0
prevIndex = 0
prevDt = 0

def plotData(plotArgs):
    ''' Update with the new data. Parameters:
    args: array that contains
      index: the indexes of the datadump
      data : the data
    '''
    global prevDt
    global prevTime

    index, time, data = plotArgs

    print('Plotting {}'.format(index))
    lines['r-T'].set_ydata(data['T'])
    # lines['r-T'].set_label('$t = {:.2}s$'.format(time))

    if args.plot_Qs:
        lines['r-Q+'].set_ydata(data['Q_+'])
        lines['r-Q-'].set_ydata(data['Q_-'])
    else:
        lines['r-Mdot'].set_ydata(data['M_dot'])

    lines['r-Sigma'].set_ydata(data['Sigma'])

    newDt = time - prevTime
    prevTime = time
    if newDt == 0:
        dt = prevDt
    else:
        dt = newDt
        prevDt = dt

    print('dt: ', dt)
    for ind, line in lines['Sigma-T']:
        x = data['Sigma'][ind]
        y = data['T'][ind]
        line.set_xdata(x)
        line.set_ydata(y)

        print(ind, dt * (data['Q_+'][ind] - data['Q_-'][ind]) / data['Cv'][ind])

    # ax11.legend()
    fig.suptitle('$t = {:.2f}s$, iteration {}'.format(time, index))

def onClick(data, event):
    data.pause ^= True
    data.offset = 0
    print(data.pause)

def onKey(data, event):
    def mod(amount):
        data.offset += amount

    behaviour = {
        'down': lambda: mod(-1),
        'up': lambda: mod(+1),
        ' ': lambda: onClick(data, ''),
        'ctrl+down': lambda: mod(-10),
        'ctrl+up': lambda: mod(+10),
        'shift+ctrl+down': lambda: mod(-1000),
        'shift+ctrl+up': lambda: mod(+1000),
        }
    if event.key in behaviour:
        behaviour[event.key]()
    else:
        print('Unsupported key "{}"'.format(event.key))

if __name__ == '__main__':
    print('Reading initial conditions')
    ic = pd.read_csv('CI.dat', delim_whitespace=True)

    print('Reading output file')
    simulData = Data('output.dat')

    print('Reading critical points')
    crit_pts = pd.read_csv(args.c_points, delim_whitespace=True)

    print('Reading S_curves')
    s_curves = [ (ind,
                  pd.read_csv(args.s_curves_dir + '/Temperature_Sigma_{:0>5}_tot.dat'.format(ind),
                              delim_whitespace=True, dtype=float, header=0))
                 for ind in args.s_curves]

    data0 = simulData.get()[-1]
    initFun = lambda: init(ic, crit_pts, s_curves, data0)

    if len(args.plot) == 1:
        initFun()
        plotData((0, 0, data0))
        plt.show()
    else:
        simulData.loop = args.loop
        simulData.reRead = args.reread

        if (len(args.plot) > 0):
            simulData.restrictTo = args.plot

        data = simulData

        onClickHandler = lambda event: onClick(data, event)
        onKeyHandler = lambda event: onKey(data, event)

        # fig.canvas.mpl_connect('button_press_event', onClickHandler)
        fig.canvas.mpl_connect('key_press_event', onKeyHandler)
        ani = animation.FuncAnimation(fig, plotData, data, init_func=initFun,
                                      interval=args.interval)
        if args.video:
            # temporaly deactivate looping for the video
            simulData.loop = False
            ani.save(args.video_file, writer='ffmpeg', fps=10, bitrate=10000, dpi=180)
            simulData.loop = args.loop

        # clear the axes
        for ax in (ax11, ax12, ax21, ax22):
            ax.axes.cla()

        plt.show()
<commit_msg>plot_data: fix incompatible type<commit_after>#!/usr/bin/python
import argparse
import matplotlib.pyplot as plt
import pandas as pd
import matplotlib.animation as animation
import numpy as np
from itertools import tee

parser = argparse.ArgumentParser(description='Plot data from output of the black hole simulation.')
parser.add_argument('--video', action='store_true',
                    help='save a video (MUCH slower)')
parser.add_argument('--dont-loop', action='store_false', dest='loop', default=True,
                    help='prevent from looping infinitely')
parser.add_argument('--video-file', metavar='file', default='evolution.mp4',
                    help='save the video as (default: %(default)s).')
parser.add_argument('--plot', metavar='dump_id', nargs='*', type=int, default=[],
                    help='dump_ids to plot.')
parser.add_argument('--c-points', metavar='file',
                    help='Critical point file (default: %(default)s).',
                    default='critical_points/file.dat')
parser.add_argument('--i-conditions', metavar='file',
                    help='Initial conditions file (default: %(default)s).', default='CI.dat')
parser.add_argument('--s-curves', metavar='n', nargs='+', type=int,
                    help='S curves to plot (default: %(default)s).', default=[1, 10, 20, 255])
parser.add_argument('--s-curves-dir', metavar='dir', nargs=1,
                    help='S curves directory (default: %(default)s).', default='s_curves')
parser.add_argument('--reread', action='store_true', default=False,
                    help='Reread file when fully read.')
parser.add_argument('--interval', default=0, type=int,
                    help='Interval (in ms) between each frames (default %(default)s)')
parser.add_argument('--plot-Qs', default=False, action='store_true', dest='plot_Qs',
                    help='Plot Q⁺ and Q⁻ instead of Mdot')
args = parser.parse_args()

fig, ((ax11, ax12), (ax21, ax22)) = plt.subplots(2, 2)

class Data:
    ''' Class that opens filename and reads it. You can directly iterate over it.'''
    def __init__(self, filename):
        self.filename = filename # holds the filename
        self.data = {}           # holds the file data
        self.times = {}          # holds the data
        self.indexes = []        # holds a list of each snapshot
        self.restrictTo = []     # list of indexes to restrict to

        self.file = open(self.filename, 'r')
        # read ahead one line
        self.lastLineRead = self.file.readline()
        self.pause = False       # bool flag if paused (activates left<>right navigation
        self.reread = False      #
        self.offset = 0          # Number of indexes to jump back or forth

    def getChunk(self):
        '''Get a chunk (niter, time, headers + data)
        starting at last line read.
        In the end, reloads the last line read and not
        returned into self.lastlineread'''

        f = self.file

        ###########################################
        # Check not end of file
        ###########################################
        if self.lastLineRead == '':
            return -1, -1, [], [], True

        ###########################################
        # preload first lines
        ###########################################
        niter = int(self.lastLineRead.split('#')[1])
        time = float(f.readline().split('#')[1])
        headers = f.readline().split()
        eof = False # flag to set if end of file

        ###########################################
        # read until reaching a commented line
        ###########################################
        data = []
        line = f.readline()
        while '#' not in line:
            # if eof reached
            if line == '':
                eof = True
                break
            # remove last character ('\n')
            data.append(line[:-1].split())

            line = f.readline()

        pdData = pd.DataFrame(data=data, columns=headers, dtype=float)
        # store the last line read
        self.lastLineRead = line
        return niter, time, headers, pdData, eof

    def getPosition(self, _from, offset):
        ''' Given an initial position, returns the
        index of the n'th position after and the effective offset applied.'''
        indexOfNiter = self.indexes.index(_from)

        tmpIndex = indexOfNiter - 1 + offset
        if tmpIndex >= len(self.indexes):
            effectiveOffset = len(self.indexes) - indexOfNiter
        elif tmpIndex < 0:
            effectiveOffset = -indexOfNiter + 1
        else:
            effectiveOffset = offset

        index = indexOfNiter - 1 + effectiveOffset
        return effectiveOffset, self.indexes[index]

    def __iter__(self):
        ''' Infinite iterator over the whole file'''
        self.restrictTo.sort()

        # Convert the restriction into a boolean
        if len(self.restrictTo) > 0:
            noRestriction = False
        else:
            noRestriction = True

        ###########################################
        # First read the whole file chunk by chunk
        ###########################################
        i = 0
        niter, time, headers, data, eof = self.getChunk()
        while not eof:
            if self.pause:
                self.offset, i = self.getPosition(niter, self.offset)
                print('i:', self.indexes, self.offset)

                yield (i, self.times[i], self.data[i])
            elif noRestriction or niter in self.restrictTo:
                yield (niter, time, data)

            ###########################################
            # load the missing lines ahead
            ###########################################
            readAhead = 0
            if self.pause and self.offset:
                readAhead += self.offset
            elif not self.pause:
                readAhead = 1

            while readAhead > 0 and not eof:
                print('reading ahead')
                # if we have a restriction list and the current iteration
                # is after the last accepted item, stop reading the file
                if not noRestriction and niter > self.restrictTo[-1]:
                    break

                niter, time, headers, data, eof = self.getChunk()
                self.data[niter] = data
                self.times[niter] = time
                self.indexes.append(niter)
                self.indexes.sort()

                self.offset = 0

                readAhead -= 1
                # if reread flag activated, close and reopen the file
                if eof and self.reread:
                    self.file.close()
                    self.file = open(self.filename, 'r')
                    self.lastLineRead = self.file.readline()
                    eof = False

        ###########################################
        # If looping, pop back the data already read
        ###########################################
        if self.loop:
            # dataKeys holds the keys where to find the data
            if noRestriction:
                dataKeys = list(self.data.keys())
            else:
                dataKeys = [key for key in self.data.keys() if key in self.restrictTo]

            dataKeys.sort()
            while True:
                prevKey = dataKeys[-1]
                for key in dataKeys:
                    if self.pause:
                        self.offset, i = self.getPosition(niter, self.offset)
                        yield (i, self.times[i], self.data[i])
                    else:
                        yield (key, self.times[key], self.data[key])
                    prevKey = key

    def get(self):
        ''' Get a chunk of the file without moving the position of the reader in the file '''
        pos = self.file.tell()
        niter, time, headers, data, eof = self.getChunk()
        self.file.seek(pos)
        return niter, time, headers, data

    def __del__(self):
        ''' Close the file reader'''
        self.file.close()

lines = {'r-T': 0,
         'r-Sigma': 0,
         'Sigma-T': []}

class colorLooper:
    def __init__(self, colors= ['red', 'green', 'blue', 'grey', 'purple', 'brown', 'pink']):
        self.colors = colors
        self.i = 0

    def __next__(self):
        self.i += 1
        return self.colors[self.i - 1 % len(self.colors)]

    def __iter__(self):
        self.i = 0
        while self.i < len(self.colors):
            yield self.__next__()

    def reset(self):
        self.i = 0

def init(ic, crit_pts, s_curves, initial_data):
    ''' Plot the initial conditions, the critical points and the s_curves'''
    ## Top left panel
    ax11.plot(ic['r'], ic['T'], '--', label='Initial conditions')
    ax11.plot(ic['r'], crit_pts['Temp_thick'], '--', label='critical')

    ax11.set_xlabel('$r\ (cm)$')
    ax11.set_ylabel('$T\ (K)$')

    lines['r-T'] = ax11.plot(initial_data['r'], initial_data['T'])[0]

    ax11.set_yscale('log')
    ax11.grid()
    # ax11.legend()

    # add ticks corresponding to the s_curve
    ticks      = list(ax11.get_xticks())
    ticklabels = list(ax11.get_xticklabels())
    for ind, s_curve in s_curves:
        ticks.append(ic['r'][ind])
        ticklabels.append('$r_{'+str(ind)+'}$')

    ax11.set_xticks(ticks)
    ax11.set_xticklabels(ticklabels)

    ## Top right panel
    ax12.set_xlabel('$r\ (\mathrm{cm})$')
    ax12.set_ylabel('$\Sigma\ (\mathrm{g.cm^{-2}})$')
    ax12.plot(ic['r'], ic['Sigma'], '--')
    ax12.plot(ic['r'], crit_pts['Sigma_thick'], '--')

    lines['r-Sigma'] = ax12.plot(initial_data['r'], initial_data['Sigma'])[0]

    ax12.grid()
    ax12.set_yscale('log')

    if args.plot_Qs:
        lines['r-Q+'] = ax21.plot(initial_data['r'], initial_data['Q_+'], label='$Q^+$')[0]
        lines['r-Q-'] = ax21.plot(initial_data['r'], initial_data['Q_-'], label='$Q^-$')[0]
        ax21.legend()
    else:
        lines['r-Mdot'] = ax21.plot(initial_data['r'], initial_data['M_dot'])[0]

    # add ticks corresponding to the s_curve
    ticks      = list(ax12.get_xticks())
    ticklabels = list(ax12.get_xticklabels())
    for ind, s_curve in s_curves:
        ticks.append(ic['r'][ind])
        ticklabels.append('$r_{'+str(ind)+'}$')

    ax12.set_xticks(ticks)
    ax12.set_xticklabels(ticklabels)

    ## Bottom left panel
    ax21.set_xlabel('$r\ (\mathrm{cm})$')
    if args.plot_Qs:
        ax21.set_ylabel('$Q$')
    else:
        ax21.set_ylabel('$\dot{M}\ (\mathrm{g.s^{-1}})$')
    ax21.grid()

    ## Bottom right panel
    colorsIter = colorLooper()
    for ind, s_curve in s_curves:
        ax22.plot(s_curve['Surface_density'], s_curve['Temperature'],
                  label='$r_{'+str(ind)+'}$',
                  c=colorsIter.__next__())

    # add the initial data on the s_curve plot
    colorsIter.reset()
    lines['Sigma-T'] = [ (ind, ax22.plot(initial_data['Sigma'][ind],
                                         initial_data['T'][ind], 'o',
                                         c=colorsIter.__next__())[0])
                     for ind, foo in s_curves]

    ax22.set_xlabel('$\Sigma\ (\mathrm{g.cm^{-2}})$')
    ax22.set_ylabel('$T\ (\mathrm{K})$')
    ax22.legend(loc='upper right')
    ax22.grid()

prevTime = 0
prevIndex = 0
prevDt = 0

def plotData(plotArgs):
    ''' Update with the new data. Parameters:
    args: array that contains
      index: the indexes of the datadump
      data : the data
    '''
    global prevDt
    global prevTime

    index, time, data = plotArgs

    print('Plotting {}'.format(index))
    lines['r-T'].set_ydata(data['T'])
    # lines['r-T'].set_label('$t = {:.2}s$'.format(time))

    if args.plot_Qs:
        lines['r-Q+'].set_ydata(data['Q_+'])
        lines['r-Q-'].set_ydata(data['Q_-'])
    else:
        lines['r-Mdot'].set_ydata(data['M_dot'])

    lines['r-Sigma'].set_ydata(data['Sigma'])

    newDt = time - prevTime
    prevTime = time
    if newDt == 0:
        dt = prevDt
    else:
        dt = newDt
        prevDt = dt

    print('dt: ', dt)
    for ind, line in lines['Sigma-T']:
        x = data['Sigma'][ind]
        y = data['T'][ind]
        line.set_xdata(x)
        line.set_ydata(y)

        print(ind, dt * (data['Q_+'][ind] - data['Q_-'][ind]) / data['Cv'][ind])

    # ax11.legend()
    fig.suptitle('$t = {:.2f}s$, iteration {}'.format(time, index))

def onClick(data, event):
    data.pause ^= True
    data.offset = 0
    print(data.pause)

def onKey(data, event):
    def mod(amount):
        data.offset += amount

    behaviour = {
        'down': lambda: mod(-1),
        'up': lambda: mod(+1),
        ' ': lambda: onClick(data, ''),
        'ctrl+down': lambda: mod(-10),
        'ctrl+up': lambda: mod(+10),
        'shift+ctrl+down': lambda: mod(-1000),
        'shift+ctrl+up': lambda: mod(+1000),
        }
    if event.key in behaviour:
        behaviour[event.key]()
    else:
        print('Unsupported key "{}"'.format(event.key))

if __name__ == '__main__':
    print('Reading initial conditions')
    ic = pd.read_csv('CI.dat', delim_whitespace=True)

    print('Reading output file')
    simulData = Data('output.dat')

    print('Reading critical points')
    crit_pts = pd.read_csv(args.c_points, delim_whitespace=True)

    print('Reading S_curves')
    s_curves = [ (ind,
                  pd.read_csv(args.s_curves_dir + '/Temperature_Sigma_{:0>5}_tot.dat'.format(ind),
                              delim_whitespace=True, dtype=float, header=0))
                 for ind in args.s_curves]

    data0 = simulData.get()[-1]
    initFun = lambda: init(ic, crit_pts, s_curves, data0)

    if len(args.plot) == 1:
        initFun()
        plotData((0, 0, data0))
        plt.show()
    else:
        simulData.loop = args.loop
        simulData.reRead = args.reread

        if (len(args.plot) > 0):
            simulData.restrictTo = args.plot

        data = simulData

        onClickHandler = lambda event: onClick(data, event)
        onKeyHandler = lambda event: onKey(data, event)

        # fig.canvas.mpl_connect('button_press_event', onClickHandler)
        fig.canvas.mpl_connect('key_press_event', onKeyHandler)
        ani = animation.FuncAnimation(fig, plotData, data, init_func=initFun,
                                      interval=args.interval)
        if args.video:
            # temporaly deactivate looping for the video
            simulData.loop = False
            ani.save(args.video_file, writer='ffmpeg', fps=10, bitrate=10000, dpi=180)
            simulData.loop = args.loop

        # clear the axes
        for ax in (ax11, ax12, ax21, ax22):
            ax.axes.cla()

        plt.show()
<|endoftext|>


<commit_before>import numpy
from matplotlib import pyplot
from matplotlib.backends.backend_pdf import PdfPages
import os, sys, re
import argparse
import common
from nltk.stem.porter import *

"""Read the program similarity result files and plot histograms"""

def compute_method_text_similarity(m1_full_str, m2_full_str, re_prog):
    return

def compile_camel_case_re_pattern():
	return re.compile(r".+?(?:(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])|$)")

def compile_method_re_pattern():
	return re.compile(r"<[\w\d_$\.]+\s*:\s+[\w\d_$]+\s+([\w\d_$]+)\([\w\d_$\.\s]*\)>\s")

def get_method_name_only(method_full_str, re_prog):
	#Example: <org.dyn4j.dynamics.joint.RevoluteJoint: void setMotorEnabled(boolean)>
	m = re_prog.match(method_full_str)
	if m:
		return m.group(1)
	else:
		print("Should always find a method name. The fully qualitified method name was:")
		print(method_full_str)
		sys.exit(0)

def create_stemmer():
	return SnowballStemmer('english')

def stem_word_lst(stemmer, word_lst):
	return [stemmer.stem(w) for w in word_lst]

def parse_result_file(result_file):
	"""
	file format:
	path_to_dotA:
	path_to_similar_dot1 , score
	...
	path_to_similar_dot5 , score

	path_to_dotB:
	...
	"""
	dot_score_lst = []
	dot_sim_result = {}

	count = 0
	score = 0.0
	current_dot = None
	with open(result_file, "r") as fi:
		for line in fi:
			line = line.rstrip('\n')
                        if len(line)>0 and line[-1]==":":
                                current_dot = line[:-1]
                                dot_sim_result[current_dot] = []
                        else:
                                linarr = line.split(" , ")
                                if linarr[0][-3:]=="dot":
                                        count += 1
                                        score += float(linarr[1])
                                        dot_sim_result[current_dot].append((linarr[0],linarr[1]))
                                        if count==5:
                                                dot_score_lst.append((current_dot, score/count))
                                                count = 0
                                                score = 0.0
        return (dot_score_lst, dot_sim_result)

def show_improvement(proj, dot_score_lst_nc, dot_score_lst_c, dot_sim_res_nc, dot_sim_res_c, dot_method_map, topk):
	total = 0.0
	impr_lst = []
	assert len(dot_score_lst_nc)==len(dot_score_lst_c), "Should have the same number of methods with or without clustering."
	for i in range(len(dot_score_lst_nc)):
		assert dot_score_lst_nc[i][0]==dot_score_lst_c[i][0], "Should be comparing the same dot."
		impr_score = dot_score_lst_c[i][1] - dot_score_lst_nc[i][1]
		#assert impr_score+0.00001>=0.0, "Clustering should not degrade the performance of similar program identification."
		impr_lst.append((dot_score_lst_nc[i][0], impr_score))
		total += impr_score
        impr_lst.sort(key=lambda x: x[1], reverse=True)
	print("\n***************************\n")
	print("{0}:\n".format(proj))
	print("Total similarity score improvement: {0}.".format(total))
	print("Average similarity score improvement per method: {0}.\n".format(total/len(dot_score_lst_nc)))
	print("The top {0} most improved methods are:\n".format(topk))
	for i in range(topk):
   		dot_name = impr_lst[i][0]
                print(dot_method_map[dot_name]+" : average similarity score improved by " + str(impr_lst[i][1]))
                print("Before clustering:")
                nc_lst = dot_sim_res_nc[dot_name]
                for j in range(len(nc_lst)):
                        print(dot_method_map[nc_lst[j][0]]+" , "+nc_lst[j][1])
                print("After clustering:")
                c_lst = dot_sim_res_c[dot_name]
                for j in range(len(c_lst)):
                        print(dot_method_map[c_lst[j][0]]+" , "+c_lst[j][1])
                print("\n")
	print("\n***************************\n")

def plot_hist(x, xlabel, y, ylabel, fig_file, perf_stat_lst):
	bins = numpy.linspace(0.0, 2.0, 100)
	#pyplot.hist(x, bins, alpha=0.5, label=xlabel)
	#pyplot.hist(y, bins, alpha=0.5, label=ylabel)
	data = numpy.vstack([x, y]).T
	pyplot.hist(data, bins, alpha=0.7, label=[xlabel, ylabel]+perf_stat_lst)
	pyplot.legend(loc="upper right")
	#pyplot.show()
	pp = PdfPages(fig_file+".pdf")
	pyplot.savefig(pp, format='pdf')
	pp.close()

def get_dot_method_map(proj_lst):
	dot_method_map = {}
	for proj in proj_lst:
		output_dir_lst = common.ALL_DOT_DIR[proj]
		for output_dir in output_dir_lst:
			method_file = common.get_method_path(proj, output_dir)
			with open(method_file, "r") as mf:
				for line in mf:
					line = line.rstrip()
					items = line.split("\t")
					method_name = items[0]
					method_dot = items[1]
					method_dot_path = common.get_dot_path(proj, output_dir, method_dot)
					dot_method_map[method_dot_path] = method_name
	return dot_method_map

def main():

	parser = argparse.ArgumentParser()
	parser.add_argument("-nc", "--nocluster", required=True, type=str, help="path to the result folder without relabeling")
	parser.add_argument("-c", "--cluster", required=True, type=str, help="path to the result folder with relabeling")
	parser.add_argument("-f", "--fig", required=True, type=str, help="path to the figure folder")
	parser.add_argument("-s", "-strategy", required=True, type=str, help="name of the strategy")
	parser.add_argument("-k", "--topk", type=int, help="top k most improved methods")
	parser.add_argument("-a", "--all", action="store_true", help="set to merge results from all benchmark projects in a single histogram")
	args = parser.parse_args()

	proj_lst = common.LIMITED_PROJECT_LIST
	common.mkdir(args.fig)

	dot_method_map = get_dot_method_map(proj_lst)

	topk = 10
	if args.topk:
		topk = args.topk

	all_score_lst_nc = []
	all_score_lst_c = []

	for proj in proj_lst:
	    proj_result_file_name = proj + "_result.txt"
	    (dot_lst_nc, dot_res_nc) = parse_result_file(os.path.join(args.nocluster, proj_result_file_name))
	    (dot_lst_c, dot_res_c) = parse_result_file(os.path.join(args.cluster, proj_result_file_name))
	    score_lst_nc = [x[1] for x in dot_lst_nc]
	    score_lst_c = [x[1] for x in dot_lst_c]
	    if args.all:
	    	all_score_lst_nc += score_lst_nc
	    	all_score_lst_c += score_lst_c
	    else:
                plot_hist(score_lst_nc, "w/o clustering", score_lst_c, args.strategy, os.path.join(args.fig, proj))
                show_improvement(proj, dot_lst_nc, dot_lst_c, dot_res_nc, dot_res_c, dot_method_map, topk)
                print("\n")
        if args.all:
            plot_hist(all_score_lst_nc, "w/o clustering", all_score_lst_c, args.strategy, os.path.join(args.fig, "all"))

if __name__ == "__main__":
    main()
<commit_msg>better plots<commit_after>import numpy
from matplotlib import pyplot
from matplotlib.backends.backend_pdf import PdfPages
import os, sys, re
import argparse
import common
from nltk.stem.porter import *

"""Read the program similarity result files and plot histograms"""

def compute_method_text_similarity(m1_full_str, m2_full_str, re_prog):
    return

def compile_camel_case_re_pattern():
	return re.compile(r".+?(?:(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])|$)")

def compile_method_re_pattern():
	return re.compile(r"<[\w\d_$\.]+\s*:\s+[\w\d_$]+\s+([\w\d_$]+)\([\w\d_$\.\s]*\)>\s")

def get_method_name_only(method_full_str, re_prog):
	#Example: <org.dyn4j.dynamics.joint.RevoluteJoint: void setMotorEnabled(boolean)>
	m = re_prog.match(method_full_str)
	if m:
		return m.group(1)
	else:
		print("Should always find a method name. The fully qualitified method name was:")
		print(method_full_str)
		sys.exit(0)

def create_stemmer():
	return SnowballStemmer('english')

def stem_word_lst(stemmer, word_lst):
	return [stemmer.stem(w) for w in word_lst]

def parse_result_file(result_file):
	"""
	file format:
	path_to_dotA:
	path_to_similar_dot1 , score
	...
	path_to_similar_dot5 , score

	path_to_dotB:
	...
	"""
	dot_score_lst = []
	dot_sim_result = {}

	count = 0
	score = 0.0
	current_dot = None
	with open(result_file, "r") as fi:
		for line in fi:
			line = line.rstrip('\n')
                        if len(line)>0 and line[-1]==":":
                                current_dot = line[:-1]
                                dot_sim_result[current_dot] = []
                        else:
                                linarr = line.split(" , ")
                                if linarr[0][-3:]=="dot":
                                        count += 1
                                        score += float(linarr[1])
                                        dot_sim_result[current_dot].append((linarr[0],linarr[1]))
                                        if count==5:
                                                dot_score_lst.append((current_dot, score/count))
                                                count = 0
                                                score = 0.0
        return (dot_score_lst, dot_sim_result)

def show_improvement(proj, dot_score_lst_nc, dot_score_lst_c, dot_sim_res_nc, dot_sim_res_c, dot_method_map, topk):
	nc_total = 0.0
	c_total = 0.0
	largest_impr = 0.0
	impr_lst = []
	assert len(dot_score_lst_nc)==len(dot_score_lst_c), "Should have the same number of methods with or without clustering."
	for i in range(len(dot_score_lst_nc)):
		assert dot_score_lst_nc[i][0]==dot_score_lst_c[i][0], "Should be comparing the same dot."
		nc_total += dot_score_lst_nc[i][1]
		c_total += dot_score_lst_c[i][1]
		impr = dot_score_lst_c[i][1] - ot_score_lst_nc[i][1]
		if impr > largest_impr:
			largest_impr = impr
		#assert impr_score+0.00001>=0.0, "Clustering should not degrade the performance of similar program identification."
		impr_lst.append((dot_score_lst_nc[i][0], impr_score))
        impr_lst.sort(key=lambda x: x[1], reverse=True)
    total_impr = c_total - nc_total
	print("\n***************************\n")
	print("{0}:\n".format(proj))
	print("Average similarity score improvement per method: {0}.\n".format(total_impr/len(dot_score_lst_nc)))
	print("Percentage score improvement: {0}.\n".format(total_impr/nc_total))
	print("Largest score improvement for a single method: {0}.\n".format(largest_impr))
	print("The top {0} most improved methods are:\n".format(topk))
	for i in range(topk):
   		dot_name = impr_lst[i][0]
                print(dot_method_map[dot_name]+" : average similarity score improved by " + str(impr_lst[i][1]))
                print("Before clustering:")
                nc_lst = dot_sim_res_nc[dot_name]
                for j in range(len(nc_lst)):
                        print(dot_method_map[nc_lst[j][0]]+" , "+nc_lst[j][1])
                print("After clustering:")
                c_lst = dot_sim_res_c[dot_name]
                for j in range(len(c_lst)):
                        print(dot_method_map[c_lst[j][0]]+" , "+c_lst[j][1])
                print("\n")
	print("\n***************************\n")
	# output some stats
	return (avg_impr, percent_impr, largest_impr)

def plot_hist(x, xlabel, y, ylabel, fig_file, avg_impr, percent_impr, largest_impr):
	perf_labels = ["Avg score improvement: ".format(avg_impr), "\% score improvement: ".format(percent_impr*100), "Largest single improvement: ".format(largest_impr)]
	bins = numpy.linspace(0.0, 2.0, 100)
	#pyplot.hist(x, bins, alpha=0.5, label=xlabel)
	#pyplot.hist(y, bins, alpha=0.5, label=ylabel)
	data = numpy.vstack([x, y]).T
	pyplot.hist(data, bins, alpha=0.7, label=[xlabel, ylabel]+perf_labels)
	pyplot.legend(loc="upper right")
	#pyplot.show()
	pyplot.xlim(0.0, 1.0)
	pp = PdfPages(fig_file+".pdf")
	pyplot.savefig(pp, format='pdf')
	pp.close()

def get_dot_method_map(proj_lst):
	dot_method_map = {}
	for proj in proj_lst:
		output_dir_lst = common.ALL_DOT_DIR[proj]
		for output_dir in output_dir_lst:
			method_file = common.get_method_path(proj, output_dir)
			with open(method_file, "r") as mf:
				for line in mf:
					line = line.rstrip()
					items = line.split("\t")
					method_name = items[0]
					method_dot = items[1]
					method_dot_path = common.get_dot_path(proj, output_dir, method_dot)
					dot_method_map[method_dot_path] = method_name
	return dot_method_map

def main():

	parser = argparse.ArgumentParser()
	parser.add_argument("-nc", "--nocluster", required=True, type=str, help="path to the result folder without relabeling")
	parser.add_argument("-c", "--cluster", required=True, type=str, help="path to the result folder with relabeling")
	parser.add_argument("-f", "--fig", required=True, type=str, help="path to the figure folder")
	parser.add_argument("-s", "-strategy", required=True, type=str, help="name of the strategy")
	parser.add_argument("-k", "--topk", type=int, help="top k most improved methods")
	parser.add_argument("-a", "--all", action="store_true", help="set to merge results from all benchmark projects in a single histogram")
	args = parser.parse_args()

	proj_lst = common.LIMITED_PROJECT_LIST
	common.mkdir(args.fig)

	dot_method_map = get_dot_method_map(proj_lst)

	topk = 10
	if args.topk:
		topk = args.topk

	all_score_lst_nc = []
	all_score_lst_c = []

	for proj in proj_lst:
	    proj_result_file_name = proj + "_result.txt"
	    (dot_lst_nc, dot_res_nc) = parse_result_file(os.path.join(args.nocluster, proj_result_file_name))
	    (dot_lst_c, dot_res_c) = parse_result_file(os.path.join(args.cluster, proj_result_file_name))
	    score_lst_nc = [x[1] for x in dot_lst_nc]
	    score_lst_c = [x[1] for x in dot_lst_c]
	    if args.all:
	    	all_score_lst_nc += score_lst_nc
	    	all_score_lst_c += score_lst_c
	    else:
			show_improvement(proj, dot_lst_nc, dot_lst_c, dot_res_nc, dot_res_c, dot_method_map, topk)
			plot_hist(score_lst_nc, "w/o clustering", score_lst_c, args.strategy, os.path.join(args.fig, proj))
			print("\n")
	if args.all:
		plot_hist(all_score_lst_nc, "w/o clustering", all_score_lst_c, args.strategy, os.path.join(args.fig, "all"))

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


<commit_before>import numpy
from matplotlib import pyplot
from matplotlib.backends.backend_pdf import PdfPages
import os, sys
import argparse
import common

"""Read the program similarity result files and plot histograms"""

def parse_result_file(result_file):
	"""
	file format:
	path_to_dotA:
	path_to_similar_dot1 , score
	...
	path_to_similar_dot5 , score

	path_to_dotB:
	...
	"""
	avg_score_lst = []
        dot_lst = []

	count = 0
	score = 0.0
        dot_name = None
	with open(result_file, "r") as fi:
		for line in fi:
			line = line.rstrip('\n')
                        if len(line)>0 and line[-1]==":":
                                dot_name = line[:-1]
                        else:
                                linarr = line.split(" , ")
                                if linarr[0][-3:]=="dot":
                                        count += 1
                                        score += float(linarr[1])
                                        if count==5:
                                                avg_score_lst.append(score/count)
                                                dot_lst.append(dot_name)
                                                count = 0
                                                score = 0.0
	return (dot_lst, avg_score_lst)

def show_improvement(avg_score_lst_nc, avg_score_lst_c, dot_lst, topk=10):
	total = 0.0
        impr_lst = []
	assert len(avg_score_lst_nc)==len(avg_score_lst_c), "Should have the same number of methods with or without clustering."
	for i in range(len(avg_score_lst_nc)):
		#assert avg_score_lst_nc[i]<=avg_score_lst_c[i], "Clustering cannot degrade the performance of similar program identification."
                impr_lst.append(avg_score_lst_c[i] - avg_score_lst_nc[i])
		total += avg_score_lst_c[i] - avg_score_lst_nc[i]
        impr_pairs = list(zip(dot_lst, impr_lst))
        impr_pairs.sort(key=lambda x: x[1], reverse=True)
	print("Total similarity score improvement: {0}.".format(total))
	print("Average similarity score improvement per method: {0}.".format(total/len(avg_score_lst_c)))
        print("The top {0} most improved methods are:".format(topk))
        for i in range(topk):
                print(impr_pairs[i][0]+" improved by " + str(impr_pairs[i][1]))

def plot_hist(x, xlabel, y, ylabel, fig_file):
	bins = numpy.linspace(0.0, 2.0, 100)
	#pyplot.hist(x, bins, alpha=0.5, label=xlabel)
	#pyplot.hist(y, bins, alpha=0.5, label=ylabel)
	data = numpy.vstack([x, y]).T
	pyplot.hist(data, bins, alpha=0.7, label=[xlabel, ylabel])
	pyplot.legend(loc="upper right")
	#pyplot.show()
	pp = PdfPages(fig_file+".pdf")
	pyplot.savefig(pp, format='pdf')
	pp.close()

def main():

	parser = argparse.ArgumentParser()
	parser.add_argument("-nc", "--nocluster", required=True, type=str, help="path to the result folder without relabeling")
	parser.add_argument("-c", "--cluster", required=True, type=str, help="path to the result folder with relabeling")
	parser.add_argument("-f", "--fig", required=True, type=str, help="path to the figure folder")
	args = parser.parse_args()

        proj_lst = common.LIMITED_PROJECT_LIST
        common.mkdir(args.fig)

        for proj in proj_lst:
                proj_result_file_name = proj + "_result.txt"
                (dot_lst1, score1) = parse_result_file(os.path.join(args.nocluster, proj_result_file_name))
                (dot_lst2, score2) = parse_result_file(os.path.join(args.cluster, proj_result_file_name))
                plot_hist(score1, "no cluster", score2, "cluster", os.path.join(args.fig, proj))
                show_improvement(score1, score2, dot_lst1)
                print()

if __name__ == "__main__":
	main()
<commit_msg>improve plot_hist to show method instead dot<commit_after>import numpy
from matplotlib import pyplot
from matplotlib.backends.backend_pdf import PdfPages
import os, sys
import argparse
import common

"""Read the program similarity result files and plot histograms"""

def parse_result_file(result_file):
	"""
	file format:
	path_to_dotA:
	path_to_similar_dot1 , score
	...
	path_to_similar_dot5 , score

	path_to_dotB:
	...
	"""
	avg_score_lst = []
        dot_lst = []

	count = 0
	score = 0.0
        dot_name = None
	with open(result_file, "r") as fi:
		for line in fi:
			line = line.rstrip('\n')
                        if len(line)>0 and line[-1]==":":
                                dot_name = line[:-1]
                        else:
                                linarr = line.split(" , ")
                                if linarr[0][-3:]=="dot":
                                        count += 1
                                        score += float(linarr[1])
                                        if count==5:
                                                avg_score_lst.append(score/count)
                                                dot_lst.append(dot_name)
                                                count = 0
                                                score = 0.0
	return (dot_lst, avg_score_lst)

def show_improvement(avg_score_lst_nc, avg_score_lst_c, dot_lst, dot_method_map, topk=10):
	total = 0.0
        impr_lst = []
	assert len(avg_score_lst_nc)==len(avg_score_lst_c), "Should have the same number of methods with or without clustering."
	for i in range(len(avg_score_lst_nc)):
		#assert avg_score_lst_nc[i]<=avg_score_lst_c[i], "Clustering cannot degrade the performance of similar program identification."
                impr_lst.append(avg_score_lst_c[i] - avg_score_lst_nc[i])
		total += avg_score_lst_c[i] - avg_score_lst_nc[i]
        impr_pairs = list(zip(dot_lst, impr_lst))
        impr_pairs.sort(key=lambda x: x[1], reverse=True)
	print("Total similarity score improvement: {0}.".format(total))
	print("Average similarity score improvement per method: {0}.".format(total/len(avg_score_lst_c)))
        print("The top {0} most improved methods are:".format(topk))
        for i in range(topk):
                print(dot_method_map[impr_pairs[i][0]]+" : score improved by " + str(impr_pairs[i][1]))

def plot_hist(x, xlabel, y, ylabel, fig_file):
	bins = numpy.linspace(0.0, 2.0, 100)
	#pyplot.hist(x, bins, alpha=0.5, label=xlabel)
	#pyplot.hist(y, bins, alpha=0.5, label=ylabel)
	data = numpy.vstack([x, y]).T
	pyplot.hist(data, bins, alpha=0.7, label=[xlabel, ylabel])
	pyplot.legend(loc="upper right")
	#pyplot.show()
	pp = PdfPages(fig_file+".pdf")
	pyplot.savefig(pp, format='pdf')
	pp.close()

def get_dot_method_map(proj_lst):
	dot_method_map = {}
	for proj in proj_lst:
		output_dir = comon.DOT_DIR[proj]
		method_file = common.get_method_path(proj, output_dir)
		with open(method_file, "r") as mf:
			for line in mf:
				line = line.rstrip()
				items = line.split("\t")
				method_name = items[0]
				method_dot = items[1]
				method_dot_path = common.get_dot_path(proj, output_dir, method_dot)
				dot_method_map[method_dot_path] = method_name
	return dot_method_map

def main():

	parser = argparse.ArgumentParser()
	parser.add_argument("-nc", "--nocluster", required=True, type=str, help="path to the result folder without relabeling")
	parser.add_argument("-c", "--cluster", required=True, type=str, help="path to the result folder with relabeling")
	parser.add_argument("-f", "--fig", required=True, type=str, help="path to the figure folder")
	args = parser.parse_args()

    proj_lst = common.LIMITED_PROJECT_LIST
    common.mkdir(args.fig)

    dot_method_map = get_dot_method_map(proj_lst)

    for proj in proj_lst:
    	proj_result_file_name = proj + "_result.txt"
        (dot_lst1, score1) = parse_result_file(os.path.join(args.nocluster, proj_result_file_name))
       	(dot_lst2, score2) = parse_result_file(os.path.join(args.cluster, proj_result_file_name))
        plot_hist(score1, "no cluster", score2, "cluster", os.path.join(args.fig, proj))
        show_improvement(score1, score2, dot_lst1, dot_method_map)
        print()

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


<commit_before>#!/usr/bin/env python

""" homebrew-pypi-poet

Invoked like "poet foo" for some package foo **which is presently
installed in sys.path**, determines which packages foo and its dependents
depend on, downloads them from pypi and computes their checksums, and
spits out Homebrew resource stanzas.
"""

from __future__ import print_function
import argparse
import codecs
from collections import OrderedDict
from contextlib import closing
from hashlib import sha256
import json
import logging
import os
import sys
from textwrap import dedent
import warnings

from jinja2 import Template
import pkg_resources

from .version import __version__

try:
    # Python 2.x
    from urllib2 import urlopen
except ImportError:
    # Python 3.x
    from urllib.request import urlopen

# Show warnings and greater by default
logging.basicConfig(level=int(os.environ.get("POET_DEBUG", 30)))

FORMULA_TEMPLATE = Template(dedent("""\
    class {{ package.name|capitalize }} < Formula
      include Language::Python::Virtualenv

      desc "Shiny new formula"
      homepage "{{ package.homepage }}"
      url "{{ package.url }}"
      sha256 "{{ package.checksum }}"

    {% if python == "python" %}
      depends_on :python if MacOS.version <= :snow_leopard
    {% else %}
      depends_on :python3
    {% endif %}

    {% if resources %}
    {%   for resource in resources %}
    {%     include ResourceTemplate %}

    {%   endfor %}
    {% endif %}
      def install
    {% if python == "python3" %}
        virtualenv_create(libexec, "python3")
    {% endif %}
        virtualenv_install_with_resources
      end

      test do
        false
      end
    end
    """), trim_blocks=True)

RESOURCE_TEMPLATE = Template("""\
  resource "{{ resource.name }}" do
    url "{{ resource.url }}"
    {{ resource.checksum_type }} "{{ resource.checksum }}"
  end
""")

class PackageNotInstalledWarning(UserWarning):
    pass

class PackageVersionNotFoundWarning(UserWarning):
    pass

def recursive_dependencies(package):
    discovered = {package}
    visited = set()

    def walk(package):
        if package in visited:
            return
        visited.add(package)
        extras = ("security",) if package == "requests" else ()
        try:
            reqs = pkg_resources.get_distribution(package).requires(extras)
        except pkg_resources.DistributionNotFound:
            return
        discovered.update(req.project_name.lower() for req in reqs)
        for req in reqs:
            walk(req)

    walk(package)
    return sorted(discovered)

def research_package(name, version=None):
    f = urlopen("https://pypi.io/pypi/{}/json".format(name))
    reader = codecs.getreader("utf-8")
    pkg_data = json.load(reader(f))
    f.close()
    d = {}
    d['name'] = pkg_data['info']['name']
    d['homepage'] = pkg_data['info'].get('home_page', '')
    artefact = None
    if version:
        for pypi_version in pkg_data['releases']:
            if pkg_resources.safe_version(pypi_version) == version:
                for version_artefact in pkg_data['releases'][pypi_version]:
                    if version_artefact['packagetype'] == 'sdist':
                        artefact = version_artefact
                        break
        if artefact is None:
            warnings.warn("Could not find an exact version match for "
                          "{} version {}; using newest instead".
                          format(name, version), PackageVersionNotFoundWarning)

    if artefact is None:  # no version given or exact match not found
        for url in pkg_data['urls']:
            if url['packagetype'] == 'sdist':
                artefact = url
                break

    d['url'] = artefact['url']
    if 'digests' in artefact and 'sha256' in artefact['digests']:
        logging.debug("Using provided checksum for %s", name)
        d['checksum'] = artefact['digests']['sha256']
    else:
        logging.debug("Fetching sdist to compute checksum for %s", name)
        with closing(urlopen(artefact['url'])) as f:
            d['checksum'] = sha256(f.read()).hexdigest()
        logging.debug("Done fetching %s", name)
    d['checksum_type'] = 'sha256'
    f.close()
    return d

def make_graph(pkg):
    ignore = ['argparse', 'pip', 'setuptools', 'wsgiref']
    pkg_deps = recursive_dependencies(pkg)

    dependencies = {key: {} for key in pkg_deps if key not in ignore}
    installed_packages = pkg_resources.working_set
    versions = {package.key: package.version for package in installed_packages}
    for package in dependencies:
        try:
            dependencies[package]['version'] = versions[package]
        except KeyError:
            warnings.warn("{} is not installed so we cannot compute "
                          "resources for its dependencies.".format(package),
                          PackageNotInstalledWarning)
            dependencies[package]['version'] = None

    for package in dependencies:
        package_data = research_package(package, dependencies[package]['version'])
        dependencies[package].update(package_data)

    return OrderedDict(
        [(package, dependencies[package]) for package in sorted(dependencies.keys())]
    )

def formula_for(package):
    nodes = make_graph(package)
    resources = [value for key, value in nodes.items()
                 if key.lower() != package.lower()]

    if package in nodes:
        root = nodes[package]
    elif package.lower() in nodes:
        root = nodes[package.lower()]
    else:
        raise "Could not find package {} in nodes {}".format(package, nodes.keys())

    python = "python" if sys.version_info.major == 2 else "python3"
    return FORMULA_TEMPLATE.render(package=root,
                                   resources=resources,
                                   python=python,
                                   ResourceTemplate=RESOURCE_TEMPLATE)

def resources_for(package):
    nodes = make_graph(package)
    return '\n\n'.join([RESOURCE_TEMPLATE.render(resource=node)
                        for node in nodes.values()])

def main():
    parser = argparse.ArgumentParser(
        description='Generate Homebrew resource stanzas for pypi packages '
                    'and their dependencies.')
    actions = parser.add_mutually_exclusive_group()
    actions.add_argument(
        '--single', '-s', metavar='package', nargs='+',
        help='Generate a resource stanza for one or more packages, '
             'without considering dependencies.')
    actions.add_argument(
        '--formula', '-f', metavar='package',
        help='Generate a complete formula for a pypi package with its '
             'recursive pypi dependencies as resources.')
    actions.add_argument(
        '--resources', '-r', metavar='package',
        help='Generate resource stanzas for a package and its recursive '
             'dependencies (default).')
    parser.add_argument('package', help=argparse.SUPPRESS, nargs='?')
    parser.add_argument(
        '-V', '--version', action='version',
        version='homebrew-pypi-poet {}'.format(__version__))
    args = parser.parse_args()

    if (args.formula or args.resources) and args.package:
        print('--formula and --resources take a single argument.',
              file=sys.stderr)
        parser.print_usage(sys.stderr)
        return 1

    if args.formula:
        print(formula_for(args.formula))
    elif args.single:
        for i, package in enumerate(args.single):
            data = research_package(package)
            print(RESOURCE_TEMPLATE.render(resource=data))
            if i != len(args.single)-1:
                print()
    else:
        package = args.resources or args.package
        if not package:
            parser.print_usage(sys.stderr)
            return 1
        print(resources_for(package))
    return 0

if __name__ == '__main__':
    sys.exit(main())
<commit_msg>Use real exceptions<commit_after>#!/usr/bin/env python

""" homebrew-pypi-poet

Invoked like "poet foo" for some package foo **which is presently
installed in sys.path**, determines which packages foo and its dependents
depend on, downloads them from pypi and computes their checksums, and
spits out Homebrew resource stanzas.
"""

from __future__ import print_function
import argparse
import codecs
from collections import OrderedDict
from contextlib import closing
from hashlib import sha256
import json
import logging
import os
import sys
from textwrap import dedent
import warnings

from jinja2 import Template
import pkg_resources

from .version import __version__

try:
    # Python 2.x
    from urllib2 import urlopen
except ImportError:
    # Python 3.x
    from urllib.request import urlopen

# Show warnings and greater by default
logging.basicConfig(level=int(os.environ.get("POET_DEBUG", 30)))

FORMULA_TEMPLATE = Template(dedent("""\
    class {{ package.name|capitalize }} < Formula
      include Language::Python::Virtualenv

      desc "Shiny new formula"
      homepage "{{ package.homepage }}"
      url "{{ package.url }}"
      sha256 "{{ package.checksum }}"

    {% if python == "python" %}
      depends_on :python if MacOS.version <= :snow_leopard
    {% else %}
      depends_on :python3
    {% endif %}

    {% if resources %}
    {%   for resource in resources %}
    {%     include ResourceTemplate %}

    {%   endfor %}
    {% endif %}
      def install
    {% if python == "python3" %}
        virtualenv_create(libexec, "python3")
    {% endif %}
        virtualenv_install_with_resources
      end

      test do
        false
      end
    end
    """), trim_blocks=True)

RESOURCE_TEMPLATE = Template("""\
  resource "{{ resource.name }}" do
    url "{{ resource.url }}"
    {{ resource.checksum_type }} "{{ resource.checksum }}"
  end
""")

class PackageNotInstalledWarning(UserWarning):
    pass

class PackageVersionNotFoundWarning(UserWarning):
    pass

def recursive_dependencies(package):
    discovered = {package}
    visited = set()

    def walk(package):
        if package in visited:
            return
        visited.add(package)
        extras = ("security",) if package == "requests" else ()
        try:
            reqs = pkg_resources.get_distribution(package).requires(extras)
        except pkg_resources.DistributionNotFound:
            return
        discovered.update(req.project_name.lower() for req in reqs)
        for req in reqs:
            walk(req)

    walk(package)
    return sorted(discovered)

def research_package(name, version=None):
    f = urlopen("https://pypi.io/pypi/{}/json".format(name))
    reader = codecs.getreader("utf-8")
    pkg_data = json.load(reader(f))
    f.close()
    d = {}
    d['name'] = pkg_data['info']['name']
    d['homepage'] = pkg_data['info'].get('home_page', '')
    artefact = None
    if version:
        for pypi_version in pkg_data['releases']:
            if pkg_resources.safe_version(pypi_version) == version:
                for version_artefact in pkg_data['releases'][pypi_version]:
                    if version_artefact['packagetype'] == 'sdist':
                        artefact = version_artefact
                        break
        if artefact is None:
            warnings.warn("Could not find an exact version match for "
                          "{} version {}; using newest instead".
                          format(name, version), PackageVersionNotFoundWarning)

    if artefact is None:  # no version given or exact match not found
        for url in pkg_data['urls']:
            if url['packagetype'] == 'sdist':
                artefact = url
                break

    d['url'] = artefact['url']
    if 'digests' in artefact and 'sha256' in artefact['digests']:
        logging.debug("Using provided checksum for %s", name)
        d['checksum'] = artefact['digests']['sha256']
    else:
        logging.debug("Fetching sdist to compute checksum for %s", name)
        with closing(urlopen(artefact['url'])) as f:
            d['checksum'] = sha256(f.read()).hexdigest()
        logging.debug("Done fetching %s", name)
    d['checksum_type'] = 'sha256'
    f.close()
    return d

def make_graph(pkg):
    ignore = ['argparse', 'pip', 'setuptools', 'wsgiref']
    pkg_deps = recursive_dependencies(pkg)

    dependencies = {key: {} for key in pkg_deps if key not in ignore}
    installed_packages = pkg_resources.working_set
    versions = {package.key: package.version for package in installed_packages}
    for package in dependencies:
        try:
            dependencies[package]['version'] = versions[package]
        except KeyError:
            warnings.warn("{} is not installed so we cannot compute "
                          "resources for its dependencies.".format(package),
                          PackageNotInstalledWarning)
            dependencies[package]['version'] = None

    for package in dependencies:
        package_data = research_package(package, dependencies[package]['version'])
        dependencies[package].update(package_data)

    return OrderedDict(
        [(package, dependencies[package]) for package in sorted(dependencies.keys())]
    )

def formula_for(package):
    nodes = make_graph(package)
    resources = [value for key, value in nodes.items()
                 if key.lower() != package.lower()]

    if package in nodes:
        root = nodes[package]
    elif package.lower() in nodes:
        root = nodes[package.lower()]
    else:
        raise Exception("Could not find package {} in nodes {}".format(package, nodes.keys()))

    python = "python" if sys.version_info.major == 2 else "python3"
    return FORMULA_TEMPLATE.render(package=root,
                                   resources=resources,
                                   python=python,
                                   ResourceTemplate=RESOURCE_TEMPLATE)

def resources_for(package):
    nodes = make_graph(package)
    return '\n\n'.join([RESOURCE_TEMPLATE.render(resource=node)
                        for node in nodes.values()])

def main():
    parser = argparse.ArgumentParser(
        description='Generate Homebrew resource stanzas for pypi packages '
                    'and their dependencies.')
    actions = parser.add_mutually_exclusive_group()
    actions.add_argument(
        '--single', '-s', metavar='package', nargs='+',
        help='Generate a resource stanza for one or more packages, '
             'without considering dependencies.')
    actions.add_argument(
        '--formula', '-f', metavar='package',
        help='Generate a complete formula for a pypi package with its '
             'recursive pypi dependencies as resources.')
    actions.add_argument(
        '--resources', '-r', metavar='package',
        help='Generate resource stanzas for a package and its recursive '
             'dependencies (default).')
    parser.add_argument('package', help=argparse.SUPPRESS, nargs='?')
    parser.add_argument(
        '-V', '--version', action='version',
        version='homebrew-pypi-poet {}'.format(__version__))
    args = parser.parse_args()

    if (args.formula or args.resources) and args.package:
        print('--formula and --resources take a single argument.',
              file=sys.stderr)
        parser.print_usage(sys.stderr)
        return 1

    if args.formula:
        print(formula_for(args.formula))
    elif args.single:
        for i, package in enumerate(args.single):
            data = research_package(package)
            print(RESOURCE_TEMPLATE.render(resource=data))
            if i != len(args.single)-1:
                print()
    else:
        package = args.resources or args.package
        if not package:
            parser.print_usage(sys.stderr)
            return 1
        print(resources_for(package))
    return 0

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