Back to home page

EIC code displayed by LXR

 
 

    


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))