Warning, /epic/bin/g4MaterialScan_to_csv 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) 2023 Chao Peng
0005 '''
0006 A script to run a 2D (eta, phi) scan using Geant4 kernel.
0007 Scan results are collected in a CSV format file.
0008 '''
0009
0010 import os
0011 import sys
0012 import errno
0013 import argparse
0014 import pandas as pd
0015 import numpy as np
0016 from io import StringIO
0017 from collections import OrderedDict as odict
0018 from wurlitzer import pipes
0019
0020 import DDG4
0021 import g4units
0022
0023 # pd.set_option('display.max_rows', 1000)
0024 PROGRESS_STEP = 1
0025
0026
0027 # a class to simplify the material scan, avoid the use of multipe instances of it
0028 class g4MaterialScanner:
0029 def __init__(self, compact, phy_list='QGSP_BERT'):
0030 kernel = DDG4.Kernel()
0031 kernel.loadGeometry(compact)
0032 DDG4.Core.setPrintFormat(str("%-32s %6s %s"))
0033 self.kernel = kernel
0034
0035 # use DDG4 Geant4 wrapper to simplify the setup
0036 geant4 = DDG4.Geant4(kernel)
0037 geant4.setupCshUI(ui=None)
0038 for i in geant4.description.detectors():
0039 o = DDG4.DetElement(i.second.ptr())
0040 sd = geant4.description.sensitiveDetector(o.name())
0041 if sd.isValid():
0042 typ = sd.type()
0043 if typ in geant4.sensitive_types:
0044 geant4.setupDetector(o.name(), geant4.sensitive_types[typ])
0045 else:
0046 logger.error('+++ %-32s type:%-12s --> Unknown Sensitive type: %s', o.name(), typ, typ)
0047 sys.exit(errno.EINVAL)
0048 geant4.setupPhysics(phy_list)
0049
0050 # save the gun for micromanagement later
0051 self.gun = geant4.setupGun('Scanner',
0052 Standalone=True,
0053 particle='geantino',
0054 energy=20*g4units.GeV,
0055 position='(0, 0, 0)',
0056 direction='(0, 0, 1)',
0057 multiplicity=1,
0058 isotrop=False)
0059
0060 scan = DDG4.SteppingAction(self.kernel, 'Geant4MaterialScanner/MaterialScan')
0061 self.kernel.steppingAction().adopt(scan)
0062 self.kernel.configure()
0063 self.kernel.initialize()
0064 self.kernel.NumEvents = 1
0065
0066 def __del__(self):
0067 self.kernel.terminate()
0068
0069 def scan(self, position, direction, phy_list='QGSP_BERT'):
0070 self.gun.position = '({},{},{})'.format(*position)
0071 self.gun.direction = '({},{},{})'.format(*direction)
0072 with pipes() as (out, err):
0073 self.kernel.run()
0074 return self.parse_scan_output(out.getvalue())
0075
0076 # A parser function to convert the output from Gean4MaterialScanner to a pandas dataframe
0077 @staticmethod
0078 def parse_scan_output(output):
0079 # find material scan lines
0080 lines = []
0081 add_line = False
0082
0083 for l in output.split('\n'):
0084 if add_line:
0085 lines.append(l.strip())
0086 if l.strip().startswith('MaterialScan'):
0087 add_line = True
0088
0089
0090 # NOTE: the code below depends on the output from DDG4::Geant4MaterialScanner
0091 # it is valid on 09/15/2023
0092 scans = []
0093 first_digit = False
0094 for i, l in enumerate(lines):
0095 line = l.strip('| ')
0096 if not first_digit and not line[:1].isdigit():
0097 continue
0098 first_digit = True
0099 if not line[:1].isdigit():
0100 break
0101 # break coordinates for endpoint, which has a format of (x, y, z)
0102 scans.append(line.strip('| ').translate({ord(i): None for i in '()'}).replace(',', ' '))
0103
0104 cols = [
0105 'material', 'Z', 'A', 'density',
0106 'rad_length', 'int_length', 'thickness', 'path_length',
0107 'int_X0', 'int_lambda', 'end_x', 'end_y', 'end_z'
0108 ]
0109
0110 dft = pd.read_csv(StringIO('\n'.join(scans)), sep='\s+', header=None, index_col=0, names=cols)
0111 print(dft)
0112 return dft.astype({key: np.float64 for key in cols[1:]})
0113
0114
0115 '''
0116 A helper function to convert a string (<min>[:<max>[:<step>]]) to an array
0117 '''
0118 def args_array(arg, step=1, include_end=True):
0119 vals = [float(x.strip()) for x in arg.split(':')]
0120 # empty or only one value
0121 if len(vals) < 2:
0122 return np.array(vals)
0123 # has step input
0124 if len(vals) > 2:
0125 step = vals[2]
0126 # inclusion of the endpoint (max)
0127 if include_end:
0128 vals[1] += step
0129 return np.arange(vals[0], vals[1], step)
0130
0131
0132 if __name__ == '__main__':
0133 parser = argparse.ArgumentParser(
0134 prog='g4MaterialScan_to_csv',
0135 description = 'A python script to scan materials of a DD4Hep geometry in geant4 kernel.'
0136 + '\n ' # 7 spaces for "usage: "
0137 + 'The scanned results are saved in a CSV file.'
0138 )
0139 parser.add_argument(
0140 '-c', '--compact', dest='compact', required=True,
0141 help='Top-level xml file of the detector description.'
0142 )
0143 parser.add_argument(
0144 '-o', '--output', default='g4mat_scan.csv',
0145 help='Path of the output csv file. Support python formatting of args, e.g., g4mat_scan_{eta}_{phi}.csv.'
0146 )
0147 parser.add_argument(
0148 '--start-point', default='0,0,0',
0149 help='Start point of the scan, use the format \"x,y,z\", unit is cm.'
0150 )
0151 parser.add_argument(
0152 '--eta', default='-4.0:4.0:0.1',
0153 help='Eta range, in the format of \"<min>[:<max>[:<step>]]\".'
0154 )
0155 parser.add_argument(
0156 '--eta-values', default=None,
0157 help='a list of eta values, separated by \",\", this option overwrites --eta.'
0158 )
0159 parser.add_argument(
0160 '--phi', default='0:30:1',
0161 help='Phi angle range, in the format of \"<min>[:<max>[:<step>]]\" (degree).'
0162 )
0163 parser.add_argument(
0164 '--phi-values', default=None,
0165 help='a list of phi values, separated by \",\", this option overwrites --phi.'
0166 )
0167 parser.add_argument(
0168 '--mat-buffer-size', type=int, default=50,
0169 help='Maximum number of materials included in the aggregated output.'
0170 )
0171 parser.add_argument(
0172 '--raw-output', action='store_true',
0173 help='Turn on to save the raw outputs from scan.'
0174 )
0175 parser.add_argument(
0176 '--sep', default='\t',
0177 help='Seperator for the CSV file.'
0178 )
0179 args = parser.parse_args()
0180
0181 if not os.path.exists(args.compact):
0182 print('Cannot find compact file {}'.format(args.compact))
0183 exit(-1)
0184
0185 start_point = np.array([float(v.strip()) for v in args.start_point.split(',')])
0186 if args.eta_values is not None:
0187 etas = np.array([float(xval.strip()) for xval in args.eta_values.split(',')])
0188 else:
0189 etas = args_array(args.eta)
0190
0191 if args.phi_values is not None:
0192 phis = np.array([float(xval.strip()) for xval in args.phi_values.split(',')])
0193 else:
0194 phis = args_array(args.phi)
0195 # sanity check
0196 if not len(phis):
0197 print('No phi values from the input {}, aborted!'.format(args.phi))
0198 exit(-1)
0199
0200 eta_phi = []
0201 mats_indices = odict()
0202 # should exist in the scan output as int_{value_type}
0203 value_types = ['X0', 'lambda']
0204 # a data buffer for the X0 values of ((eta, phi), materials, (X0, Lambda))
0205 data = np.zeros(shape=(len(etas)*len(phis), args.mat_buffer_size, len(value_types)))
0206
0207 # scan over eta and phi
0208 scanner = g4MaterialScanner(args.compact)
0209 print('Scanning {:d} eta values in [{:.2f}, {:.2f}] and {:d} phi values in [{:.2f}, {:.2f}] (degree).'\
0210 .format(len(etas), etas[0], etas[-1], len(phis), phis[0], phis[-1]))
0211 for i, eta in enumerate(etas):
0212 for j, phi in enumerate(phis):
0213
0214 nlines = i*len(phis) + j
0215 if nlines % PROGRESS_STEP == 0:
0216 print('Scanned {:d}/{:d} lines.'.format(nlines, len(etas)*len(phis)), end='\r', flush=True)
0217
0218 # scan
0219 direction = (np.cos(phi/180.*np.pi), np.sin(phi/180.*np.pi), np.sinh(eta))
0220 dfa = scanner.scan(start_point, direction)
0221 for vt in value_types:
0222 dfa.loc[:, vt] = dfa['int_{}'.format(vt)].diff(1).fillna(dfa['int_{}'.format(vt)])
0223
0224 if args.raw_output:
0225 dfa.to_csv('scan_raw_eta={:.3f}_phi={:.3f}.csv'.format(eta, phi), sep=args.sep, float_format='%g')
0226 # group by materials
0227 single_scan = dfa.groupby('material')[value_types].sum()
0228 # print(single_scan)
0229 for mat, xvals in single_scan.iterrows():
0230 if mat not in mats_indices:
0231 if len(mats_indices) >= args.mat_buffer_size:
0232 print('Number of materials exceeds MAT_BUFFER_SIZE ({:d}), dropped material {}.'.format(args.mat_buffer_size, mat))
0233 print('Hint: increase the buffer size with --mat-buffer-size.')
0234 continue
0235 mats_indices[mat] = len(mats_indices)
0236 k = mats_indices.get(mat)
0237 data[len(eta_phi), k] = xvals.values
0238 eta_phi.append((eta, phi))
0239 print('Scanned {:d}/{:d} lines.'.format(len(etas)*len(phis), len(etas)*len(phis)))
0240
0241 # save results
0242 result = dict()
0243 for i, vt in enumerate(value_types):
0244 indices = pd.MultiIndex.from_tuples(eta_phi, names=['eta', 'phi'])
0245 result[vt] = pd.DataFrame(columns=mats_indices.keys(), index=indices, data=data[:, :len(mats_indices), i])
0246 out_path = args.output.format(**vars(args))
0247 pd.concat(result, names=['value_type']).to_csv(out_path, sep=args.sep, float_format='%g')
0248 print('Scanned results saved to \"{}\"'.format(out_path))
