python/compareFigureOfMerit.py at master · ComPair/python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
#!/usr/bin/env python                                                          #
#                                                                              #
# Autor: Michela Negro, NASA-GSFC.                                             #
# On behalf of the AMEGO Team.                                                 #
#                                                                              #
#------------------------------------------------------------------------------#

"""Compare IRFs from different files
"""

import os
import re
import ast
import argparse
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
sns.set()
sns.set_style("ticks")


__description__ = 'Produce IRFs comparizon plots'


"""Command-line switches.
"""
formatter = argparse.ArgumentDefaultsHelpFormatter
PARSER = argparse.ArgumentParser(description=__description__,
                                 formatter_class=formatter)
PARSER.add_argument('-f', '--infiles', type=str, required=True, nargs='*',
                    help='input .txt file(s) (generated with plotFigureOfMwerit.py)')
PARSER.add_argument('-xs', '--xscale', type=str, choices=['linear', 'log'],
                    default='log', help='specify the scale of the x axis')
PARSER.add_argument('-ys', '--yscale', type=str, choices=['linear', 'log'],
                    default='linear', help='specify the scale of the y axis')
PARSER.add_argument('-ym', '--ymin', type=float, default=None,
                    help='Minimum y-axis value')
PARSER.add_argument('-yM', '--ymax', type=float, default=None,
                    help='Maximum y-axis value')
PARSER.add_argument('-xm', '--xmin', type=float, default=0.1,
                    help='Minimum x-axis value')
PARSER.add_argument('-xM', '--xmax', type=float, default=500,
                    help='Maximum x-axis value')
PARSER.add_argument('-t', '--title', type=str, default=' ',
                    help='Plot title')
PARSER.add_argument('-yl', '--ylabel', type=str, default='Figure of Merit',
                    help='labe to display on the y axis')
PARSER.add_argument('-s', '--save', type=str, default='', help='Save the plot; argument=filename')


FLAG_STR = '(?<=\_)\w+(?=\.txt)'
LABEL_STR = '^\S+(?=\_[AE])'
ANG_STR = '(?<=Cos)[0-9]\.[0-9]'
TC_color = np.concatenate((sns.color_palette("Greens_r")[:-2], sns.color_palette("Greens_r")[:-2][-1:None:-1]))
UC_color = np.concatenate((sns.color_palette("Blues_r")[:-2], sns.color_palette("Blues_r")[:-2][-1:None:-1]))
P_color  = np.concatenate((sns.color_palette("Reds_r")[:-2], sns.color_palette("Reds_r")[:-2][-1:None:-1]))

lines = ['-', '--', ':', '-.', '-', '--', ':', '-.']


def parse_figureofmerit_file(fom_file):
    print('Parsing %s ...'%fom_file)
    file_basename = os.path.basename(fom_file)
    m = re.search(r'%s'%FLAG_STR, file_basename)
    flag = m.group(0)
    m1 = re.search(r'%s'%LABEL_STR, file_basename)
    label = m1.group(0)
    m3 = re.search(r'%s'%ANG_STR, file_basename)
    ang = m3.group(0)
    en, fom = np.loadtxt(fom_file).T
    return en, fom, flag, float(ang), label


def compare_figureofmerit(**kwargs):
    """plot
    """

    file_list = kwargs['infiles']

    TC_count = 0
    UC_count = 0
    P_count = 0

    plt.figure(figsize=(10, 7), facecolor='white')
    plt.title(kwargs['title'], size=20)
    for i, f in enumerate(file_list):
        en_, fom_, flag, ang, label = parse_figureofmerit_file(f)
        if flag == 'TC':
            legend_label = f'{label} Cos{ang:.1f} Tracked'
            c = TC_color[TC_count]
            l = lines[TC_count]
        elif flag == 'UC':
            legend_label = f'{label} Cos{ang:.1f} Untracked'
            c = UC_color[UC_count]
            l = lines[UC_count]
        else:
            legend_label = f'{label} Cos{ang:.1f} Pair'
            c = P_color[P_count]
            l = lines[P_count]
        plt.plot(en_, fom_, 'o', color=c, linestyle=l, label=legend_label)
        plt.xlabel('Energy [MeV]', size=18)
        plt.ylabel(kwargs['ylabel'], size=18)
        plt.xlim(kwargs['xmin'], kwargs['xmax'])
        left, right = plt.ylim()
        if kwargs['ymin'] is not None and kwargs['ymax'] is not None:
            plt.ylim(kwargs['ymin'], kwargs['ymax'])
        elif kwargs['ymin'] is not None and kwargs['ymax'] is None:
            plt.ylim(kwargs['ymin'], right)
        elif kwargs['ymin'] is None and kwargs['ymax'] is not None:
            plt.ylim(left, kwargs['ymax'])
        plt.xscale(kwargs['xscale'])
        plt.yscale(kwargs['yscale'])
        plt.legend(fontsize=15)
        plt.grid(False)

        if flag == 'TC':
            TC_count += 1
        elif flag == 'UC':
            UC_count += 1
        else:
            P_count += 1

    plt.tight_layout()
    if kwargs['save'] != '':
        filename = kwargs['save']
        plt.savefig(f'{filename}')
        print(f"Created {filename} ...!")

    plt.show()

if __name__ == '__main__':
    args = PARSER.parse_args()
    compare_figureofmerit(**args.__dict__)