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 return dft.astype({key: np.float64 for key in cols[1:]})
0112
0113
0114 '''
0115 A helper function to convert a string (<min>[:<max>[:<step>]]) to an array
0116 '''
0117 def args_array(arg, step=1, include_end=True):
0118 vals = [float(x.strip()) for x in arg.split(':')]
0119 # empty or only one value
0120 if len(vals) < 2:
0121 return np.array(vals)
0122 # has step input
0123 if len(vals) > 2:
0124 step = vals[2]
0125 # inclusion of the endpoint (max)
0126 if include_end:
0127 vals[1] += step
0128 return np.arange(vals[0], vals[1], step)
0129
0130
0131 if __name__ == '__main__':
0132 parser = argparse.ArgumentParser(
0133 prog='g4MaterialScan_to_csv',
0134 description = 'A python script to scan materials of a DD4Hep geometry in geant4 kernel.'
0135 + '\n ' # 7 spaces for "usage: "
0136 + 'The scanned results are saved in a CSV file.'
0137 )
0138 parser.add_argument(
0139 '-c', '--compact', dest='compact', required=True,
0140 help='Top-level xml file of the detector description.'
0141 )
0142 parser.add_argument(
0143 '-o', '--output', default='g4mat_scan.csv',
0144 help='Path of the output csv file. Support python formatting of args, e.g., g4mat_scan_{eta}_{phi}.csv.'
0145 )
0146 parser.add_argument(
0147 '--start-point', default='0,0,0',
0148 help='Start point of the scan, use the format \"x,y,z\", unit is cm.'
0149 )
0150 parser.add_argument(
0151 '--eta', default='-4.0:4.0:0.1',
0152 help='Eta range, in the format of \"<min>[:<max>[:<step>]]\".'
0153 )
0154 parser.add_argument(
0155 '--phi', default='0:30:1',
0156 help='Phi angle range, in the format of \"<min>[:<max>[:<step>]]\" (degree).'
0157 )
0158 parser.add_argument(
0159 '--mat-buffer-size', type=int, default=50,
0160 help='Material buffer size.'
0161 )
0162 parser.add_argument(
0163 '--sep', default='\t',
0164 help='Seperator for the CSV file.'
0165 )
0166 args = parser.parse_args()
0167
0168 if not os.path.exists(args.compact):
0169 print('Cannot find compact file {}'.format(args.compact))
0170 exit(-1)
0171
0172 start_point = np.array([float(v.strip()) for v in args.start_point.split(',')])
0173 etas = args_array(args.eta)
0174 phis = args_array(args.phi)
0175 # sanity check
0176 if not len(phis):
0177 print('No phi values from the input {}, aborted!'.format(args.phi))
0178 exit(-1)
0179
0180 eta_phi = []
0181 mats_indices = odict()
0182 # should exist in the scan output as int_{value_type}
0183 value_types = ['X0', 'lambda']
0184 # a data buffer for the X0 values of ((eta, phi), materials, (X0, Lambda))
0185 data = np.zeros(shape=(len(etas)*len(phis), args.mat_buffer_size, len(value_types)))
0186
0187 # scan over eta and phi
0188 scanner = g4MaterialScanner(args.compact)
0189 print('Scanning {:d} eta values in [{:.2f}, {:.2f}] and {:d} phi values in [{:.2f}, {:.2f}] (degree).'\
0190 .format(len(etas), etas[0], etas[-1], len(phis), phis[0], phis[-1]))
0191 for i, eta in enumerate(etas):
0192 for j, phi in enumerate(phis):
0193
0194 nlines = i*len(phis) + j
0195 if nlines % PROGRESS_STEP == 0:
0196 print('Scanned {:d}/{:d} lines.'.format(nlines, len(etas)*len(phis)), end='\r', flush=True)
0197
0198 # scan
0199 direction = (np.cos(phi/180.*np.pi), np.sin(phi/180.*np.pi), np.sinh(eta))
0200 dfa = scanner.scan(start_point, direction)
0201 for vt in value_types:
0202 dfa.loc[:, vt] = dfa['int_{}'.format(vt)].diff(1).fillna(dfa['int_{}'.format(vt)])
0203
0204 # group by materials
0205 single_scan = dfa.groupby('material')[value_types].sum()
0206 # print(single_scan)
0207 for mat, xvals in single_scan.iterrows():
0208 if mat not in mats_indices:
0209 if len(mats_indices) >= args.mat_buffer_size:
0210 print('Number of materials exceeds MAT_BUFFER_SIZE ({:d}), dropped material {}.'.format(args.mat_buffer_size, mat))
0211 print('Hint: increase the buffer size with --mat-buffer-size.')
0212 continue
0213 mats_indices[mat] = len(mats_indices)
0214 k = mats_indices.get(mat)
0215 data[len(eta_phi), k] = xvals.values
0216 eta_phi.append((eta, phi))
0217 print('Scanned {:d}/{:d} lines.'.format(len(etas)*len(phis), len(etas)*len(phis)))
0218
0219 # save results
0220 result = dict()
0221 for i, vt in enumerate(value_types):
0222 indices = pd.MultiIndex.from_tuples(eta_phi, names=['eta', 'phi'])
0223 result[vt] = pd.DataFrame(columns=mats_indices.keys(), index=indices, data=data[:, :len(mats_indices), i])
0224 out_path = args.output.format(**vars(args))
0225 pd.concat(result, names=['value_type']).to_csv(out_path, sep=args.sep, float_format='%g')
0226 print('Scanned results saved to \"{}\"'.format(out_path))