Warning, /epic/bin/g4MaterialScan_raw_plot_2d is written in an unsupported language. File is not indexed.
0001 #!/usr/bin/env python3
0002
0003 # SPDX-License-Identifier: LGPL-3.0-or-later
0004 # Copyright (C) 2024 Chao Peng
0005 '''
0006 A script to plot raw data output from the script g4MaterialScan_to_csv
0007 '''
0008
0009 import os
0010 import re
0011 import argparse
0012 import pandas as pd
0013 import numpy as np
0014
0015 import matplotlib as mpl
0016 from matplotlib import pyplot as plt
0017 from matplotlib.collections import LineCollection, PatchCollection
0018 from matplotlib.patches import Wedge
0019 from mpl_toolkits.axes_grid1 import make_axes_locatable
0020
0021
0022 '''
0023 A helper function to convert a string (<min>[:<max>[:<step>]]) to an array
0024 '''
0025 def args_array(arg, step=1, include_end=True):
0026 vals = [float(x.strip()) for x in arg.split(':')]
0027 # empty or only one value
0028 if len(vals) < 2:
0029 return np.array(vals)
0030 # has step input
0031 if len(vals) > 2:
0032 step = vals[2]
0033 # inclusion of the endpoint (max)
0034 if include_end:
0035 vals[1] += step
0036 return np.arange(vals[0], vals[1], step)
0037
0038
0039 def eta2theta(v_eta):
0040 return 2.*np.arctan(np.exp(-v_eta))
0041
0042
0043 if __name__ == '__main__':
0044 parser = argparse.ArgumentParser(
0045 prog='g4MaterialScan_raw_plot',
0046 description = 'A python script to draw 2d plot of material thickness from raw outputs of g4MaterialScan_to_csv.'
0047 )
0048 parser.add_argument(
0049 '--path-format', default=r'scan_raw_eta={eta:.3f}_phi={phi:.3f}.csv',
0050 help='modify the path format, default is \"scan_raw_eta={eta:.3f}_phi={phi:.3f}.csv\".'
0051 )
0052 parser.add_argument(
0053 '--eta', default='-4.0:4.0:0.1',
0054 help='Eta range, in the format of \"<min>[:<max>[:<step>]]\".'
0055 )
0056 parser.add_argument(
0057 '--eta-values', default=None,
0058 help='a list of eta values, separated by \",\", this option overwrites --eta.'
0059 )
0060 parser.add_argument(
0061 '--phi', default='0', type=float,
0062 help='A single phi value to lookup for the data files.'
0063 )
0064 parser.add_argument(
0065 '--path-lengths', default="0, 180, 600",
0066 help='path length points, separated by \",\".'
0067 )
0068 parser.add_argument(
0069 '--sep', default='\t',
0070 help='Seperator for the CSV file.'
0071 )
0072 parser.add_argument(
0073 '--z-max', type=float, default=600,
0074 help='maximum z (x-axis) of the plot.'
0075 )
0076 parser.add_argument(
0077 '--z-min', type=float, default=-600,
0078 help='minimum z (x-axis) of the plot.'
0079 )
0080 parser.add_argument(
0081 '--r-max', type=float, default=600,
0082 help='maximum r (y-axis) of the plot.'
0083 )
0084 parser.add_argument(
0085 '--r-min', type=float, default=0,
0086 help='minimum r (y-axis) of the plot.'
0087 )
0088 parser.add_argument(
0089 '--x0-max', type=float, default=None,
0090 help='maximum x0 of the plot.'
0091 )
0092 parser.add_argument(
0093 '--x0-min', type=float, default=None,
0094 help='minimum x0 of the plot.'
0095 )
0096 parser.add_argument(
0097 '--lambda-max', type=float, default=None,
0098 help='maximum lambda of the plot.'
0099 )
0100 parser.add_argument(
0101 '--lambda-min', type=float, default=None,
0102 help='minimum lambda of the plot.'
0103 )
0104 parser.add_argument(
0105 '--plot-scale', type=str, default='log',
0106 help='only support \"log\" or \"linear\" scale.'
0107 )
0108 parser.add_argument(
0109 '--output-path', type=str, default='mat_scan_2D.{output_format}',
0110 help='path of the output plot, it supports the string format with inputs of the arguments.'
0111 )
0112 parser.add_argument(
0113 '--output-format', type=str, default='png',
0114 help='format of the plots, default is png, eps is also recommended (vectorized graphics).'
0115 )
0116 args = parser.parse_args()
0117
0118 # get the path length points
0119 pls = np.array([float(x.strip()) for x in args.path_lengths.split(',')])
0120 if len(pls) < 2:
0121 print('Need at least two points in --path-lengths')
0122 exit(1)
0123
0124 if args.eta_values is not None:
0125 etas = np.array([float(xval.strip()) for xval in args.eta_values.split(',')])
0126 else:
0127 etas = args_array(args.eta)
0128
0129 norm = None
0130 if args.plot_scale.lower() == 'linear':
0131 norm = mpl.colors.Normalize
0132 elif args.plot_scale.lower() == 'log':
0133 norm = mpl.colors.LogNorm
0134 else:
0135 print('Error: unsupported plot scale {}, please choose it from [log, linear].'.format(args.plot_scale))
0136 exit(1)
0137
0138 phi = args.phi
0139 zmin, zmax = args.z_min, args.z_max
0140 rmin, rmax = args.r_min, args.r_max
0141
0142 # read raw output data, collect information
0143 patches, x0_array, lmd_array = [], [], []
0144 thetas = eta2theta(etas)
0145 th_diffs = np.hstack([0., -np.diff(thetas)/2., 0.])
0146 th_mins = thetas - th_diffs[:-1]
0147 th_maxes = thetas + th_diffs[1:]
0148
0149 # iterate every eta (theta) scan data
0150 for i, (eta, theta, th_min, th_max) in enumerate(zip(etas, thetas, th_mins, th_maxes)):
0151 # read data file
0152 dpath = args.path_format.format(eta=eta, phi=phi)
0153 if not os.path.exists(dpath):
0154 print('Error: cannot find data file \"{}\", please check the path.'.format(dpath))
0155 exit(1)
0156 df = pd.read_csv(args.path_format.format(eta=eta, phi=phi), sep=args.sep, index_col=0)
0157 pls = df['path_length'].values
0158 x0s = df['X0'].cumsum().values
0159 lmds = df['lambda'].cumsum().values
0160 # a virtual bin size for the scan (fill all the 2D phase space)
0161 angle_min = th_min/np.pi*180.
0162 angle_max = th_max/np.pi*180.
0163 # determine if the lines are in range
0164 # segments of each scan (assuming start from 0)
0165 for seg_start, seg_end, x0, lmd in zip(np.hstack([0., pls[:-1]]), pls, x0s, lmds):
0166 # start point is already out of the plot range
0167 z0, r0 = np.cos(theta)*seg_start, np.sin(theta)*seg_start
0168 in_range = (z0 <= zmax) & (z0 >= zmin) & (r0 <= rmax) & (r0 >= rmin)
0169 if not in_range:
0170 continue
0171 x0_array.append(x0)
0172 lmd_array.append(lmd)
0173 width = seg_end - seg_start
0174 patches.append(Wedge((0., 0.), seg_end, angle_min, angle_max, width=width))
0175
0176 # generate plots
0177 cmap = mpl.colormaps['viridis']
0178 plot_meta = [
0179 ('Cumulative X0', x0_array, dict(cmap=cmap, norm=norm(vmin=args.x0_min, vmax=args.x0_max))),
0180 ('Cumulative $\Lambda$', lmd_array, dict(cmap=cmap, norm=norm(vmin=args.lambda_min, vmax=args.lambda_max))),
0181 ]
0182
0183 fig, axs = plt.subplots(len(plot_meta), 1, figsize=(10, 4*len(plot_meta)), dpi=600)
0184 for ax, (zlabel, data, p_kwargs) in zip(axs.flat, plot_meta):
0185 ax.set_xlim(zmin, zmax)
0186 ax.set_ylim(rmin, rmax)
0187 p = PatchCollection(patches, **p_kwargs)
0188 p.set_array(data)
0189 ax.add_collection(p)
0190 ax.set_xlabel('Z [cm]')
0191 ax.set_ylabel('R [cm] ($\phi = {}^{{\circ}}$)'.format(phi))
0192 # color bar
0193 divider = make_axes_locatable(ax)
0194 cax = divider.append_axes('right', size='3%', pad=0.05)
0195 cbar = fig.colorbar(p, cax=cax, orientation='vertical')
0196 cbar.ax.set_ylabel(zlabel, rotation=90)
0197 fig.savefig(args.output_path.format(**vars(args)), format=args.output_format)
