scripts/subdetector_tests/material_scan.py

0001 # SPDX-License-Identifier: LGPL-3.0-or-later
0002 # Copyright (C) 2023 Chao Peng
0003 '''
0004     A script to scan the materials thickness (X0 or Lambda) over eta
0005
0006     It uses dd4hep::rec::MaterialManager::placementsBetween to check the material layers, and then uses the detector
0007     alignment to transform world coordinates to local coordiantes, and then assigns the materials to a detector based
0008     on TGeoVolume::Contains
0009     Take a grain of salt about the materials->detector assignment, especially when the detector geometry is very complex
0010 '''
0011
0012 import os
0013 import math
0014 import json
0015 import fnmatch
0016 import argparse
0017 import subprocess
0018 import pandas as pd
0019 import numpy as np
0020 from io import StringIO
0021 from collections import OrderedDict
0022 from matplotlib import pyplot as plt
0023 from matplotlib.ticker import MultipleLocator
0024 import matplotlib.backends.backend_pdf
0025
0026 import ROOT
0027 import dd4hep
0028 import DDRec
0029
0030 pd.set_option('display.max_rows', 1000)
0031 PROGRESS_STEP = 10
0032 OTHERS_NAME = 'Others'
0033 # maximum number of detectors/materials to plot
0034 COLORS = ['royalblue', 'indianred', 'forestgreen', 'gold', 'darkviolet', 'orange', 'darkturquoise']
0035 # a specified color for Others, this should not be included in COLORS
0036 OTHERS_COLOR = 'silver'
0037
0038
0039 # FIXME: this is a work-around to an issue of dd4hep::rec::MaterialManager::placementsBetween
0040 # As of 04/21/2023, at negative eta (eta < 0.), the scan will miss the first material layer (vacuum, 2.8 cm)
0041 # To reporduce this issue, check the difference between:
0042 #   materialScan epic_craterlake.xml 0 0 0 100  40  -0.01 | grep Vacuum
0043 #   materialScan epic_craterlake.xml 0 0 0 100  40  0.01 | grep Vacuum
0044 # It is GEOMETRY DEPENDENT (the thickness of the first vacuum material layer)
0045 # This helper class can be removed once the issue is fixed
0046 # Plan to file a bug report for this (as of 04/21/2023)
0047 class ThicknessCorrector:
0048     def __init__(self, desc=None):
0049         self.missing = 0.
0050         self.scanned = False
0051         if desc is not None:
0052             self.scan_missing_thickness(desc)
0053
0054     def scan_missing_thickness(self, desc, pi=(0.,0.,0.), pf=(100.,40.,0.), dz=0.01, epsilon=1e-3):
0055         # assume negative eta is missed, maybe can do a more detailed check?
0056         pf1 = np.array(pf) + np.array([0., 0., dz])
0057         dft1 = material_scan(desc, pi, pf1, epsilon)
0058         th1 = dft1['path_length'].iloc[0]
0059         pf2 = np.array(pf) + np.array([0., 0., -dz])
0060         dft2 = material_scan(desc, pi, pf2, epsilon)
0061         th2 = dft2['path_length'].iloc[0]
0062         self.scanned = True
0063         # print(dft1.head(3))
0064         # print(dft2.head(3))
0065         if th1 != th2:
0066             self.missing = th1
0067
0068     def correct_path_length(self, direction, pl=0.):
0069         if direction[2] < 0.:
0070             pl2 = pl + np.sqrt(1./(direction[0]**2 + direction[1]**2))*self.missing
0071             # print('correcting path length {:.2f} -> {:.2f}'.format(pl, pl2))
0072             return pl2
0073         return pl
0074
0075
0076
0077 '''
0078     Re-implementation of MaterialScan::Print in DD4Hep::DDRec
0079     MaterialScan does not have a python interface and that function is relatively simple, so here it is
0080     desc: DD4hep::Detector
0081     start: 3D vector for start point
0082     end: 3D vector for end point
0083     epsilon: step size
0084 '''
0085 def material_scan(desc, start, end, epsilon, int_dets=None, thickness_corrector=None):
0086     mat_mng = DDRec.MaterialManager(desc.worldVolume())
0087     # only use the top-level detectors
0088     if int_dets is None:
0089         dets_list = [d for n, d in desc.world().children()]
0090     else:
0091         dets_list = [d for n, d in desc.world().children() if n in int_dets]
0092     # rvec = ROOT.Math.XYZVector()
0093     # id_conv = DDRec.CellIDPositionConverter(desc)
0094     # det_dict = {d.id(): n for n, d in desc.world().children()}
0095
0096     p0 = np.array(start)
0097     p1 = np.array(end)
0098     direction = (p1 - p0)/np.linalg.norm(p1 - p0)
0099     # print(p0, p1, direction)
0100     placements = mat_mng.placementsBetween(tuple(p0), tuple(p1), epsilon);
0101
0102     # calculate material layer by layer
0103     int_x0 = 0
0104     int_lambda = 0
0105     path_length = 0.
0106     if thickness_corrector is not None:
0107         path_length = thickness_corrector.correct_path_length(direction, path_length)
0108
0109     res = []
0110     for pv, l in placements:
0111         # print(pv, l)
0112         path_length += l
0113         pcurr = p0 + path_length*direction
0114         mat = pv.GetMedium().GetMaterial()
0115         radl = mat.GetRadLen()
0116         intl = mat.GetIntLen()
0117         x0 = l/radl
0118         lmd = l/intl
0119         d0 = OTHERS_NAME
0120         for d in dets_list:
0121             local = d.nominal().worldToLocal(pcurr)
0122             local = np.array([local.X(), local.Y(), local.Z()])
0123             if d.volume().Contains(local):
0124                 d0 = d.GetName()
0125         res.append([
0126             d0,
0127             # det_dict.get(det_id, 'Unknown'),
0128             mat.GetName(), mat.GetZ(), mat.GetA(), mat.GetDensity(),
0129             radl, intl, l, path_length,
0130             x0, lmd,
0131             pcurr[0], pcurr[1], pcurr[2],
0132             # local[0], local[1], local[2],
0133             ])
0134     cols = [
0135         'detector',
0136         'material', 'Z', 'A', 'density',
0137         'radl', 'intl', 'thickness', 'path_length',
0138         'X0', 'lambda',
0139         'x', 'y', 'z',
0140         # 'local_x', 'local_y', 'local_z'
0141         ]
0142     dft = pd.DataFrame(data=res, columns=cols)
0143     # print(dft.groupby('detector')['X0'].sum())
0144     return dft
0145
0146
0147 # the allowed value types {name: [col_name, label_name, major_step, minor_step]}
0148 ALLOWED_VALUE_TYPES = {
0149     'X0': ['X0', '$X_0$', 10, 5],
0150     'lambda': ['lambda', '$\lambda$', 5, 1],
0151 }
0152 # execute the script
0153 if __name__ == '__main__':
0154     parser = argparse.ArgumentParser()
0155     parser.add_argument(
0156             dest='compact',
0157             help='Top-level xml file of the detector description.'
0158             )
0159     parser.add_argument(
0160             '--path-r', type=float, default=400., # 120.,
0161             help='R_xy (cm) where the scan stops.'
0162             )
0163     parser.add_argument(
0164             '--start-point', default='0,0,0',
0165             help='Start point of the scan, use the format \"x,y,z\", unit is cm.'
0166             )
0167     parser.add_argument(
0168             '--eta-min', type=float, default=-2.0,
0169             help='Minimum eta for the scan.'
0170             )
0171     parser.add_argument(
0172             '--eta-max', type=float, default=2.0,
0173             help='Minimum eta for the scan.'
0174             )
0175     parser.add_argument(
0176             '--eta-nbins', type=int, default=401,
0177             help='Number of bins for the eta scan.'
0178             )
0179     parser.add_argument(
0180             '--phi', type=float, default=20.,
0181             help='Phi angle of the scan, unit is degree.'
0182             )
0183     parser.add_argument(
0184             '--epsilon', type=float, default=1e-3,
0185             help='Step size for each material scan.'
0186             )
0187     parser.add_argument(
0188             '--value-type', default='X0',
0189             help='Choose one in {}.'.format(list(ALLOWED_VALUE_TYPES.keys()))
0190             )
0191     parser.add_argument(
0192             '--detectors',
0193             # default='all',
0194             # ordering the detectors
0195             default='EcalBarrelScFi,EcalEndcapN,EcalEndcapP,SolenoidBarrel,HcalBarrel,HcalEndcapN,HcalEndcapP',
0196             help='Names of the interested detectors, separated by \",\". Use \"all\" for all detectors.'
0197             )
0198     args = parser.parse_args()
0199
0200     if not os.path.exists(args.compact):
0201         print('Cannot find {}'.format(args.compact))
0202         exit(-1)
0203
0204     if args.value_type not in ALLOWED_VALUE_TYPES.keys():
0205         print('Cannot find value {}, choose one in {}'.format(args.value_type, list(ALLOWED_VALUE_TYPES.keys())))
0206         exit(-1)
0207
0208     # scan parameters
0209     path_r = args.path_r
0210     phi = args.phi
0211     etas = np.linspace(args.eta_min, args.eta_max, args.eta_nbins)
0212     start_point = np.array([float(v.strip()) for v in args.start_point.split(',')])
0213     if len(start_point) != 3:
0214         print('Error: expecting three values (x,y,z) from --start-point, getting {:d} instead.'.format(len(start_point)))
0215         exit(-1)
0216
0217     # geometry initialization
0218     desc = dd4hep.Detector.getInstance()
0219     desc.fromXML(args.compact)
0220
0221     # check detector list (case sensitive)
0222     all_dets = [n for n, _ in desc.world().children()]
0223     dets = []
0224     missing_dets = []
0225     sort_dets = False
0226     for det in args.detectors.split(','):
0227         det = det.strip()
0228         if det.lower() == 'all':
0229             dets = all_dets
0230             # sort detectors by material thickness
0231             sort_dets = True
0232             break
0233         if det in all_dets:
0234             dets.append(det)
0235         else:
0236             missing_dets.append(det)
0237     if len(missing_dets) > 0:
0238         print('Warning: detectors {} were missing from the description.'.format(missing_dets))
0239         print('A full list of detectors are:')
0240         for d in all_dets:
0241             print('   --- {}'.format(d))
0242
0243
0244     # FIXME: work-around for dd4hep material scan issue, check ThicknessCorrector
0245     th_corr = ThicknessCorrector()
0246     th_corr.scan_missing_thickness(desc, start_point, np.array(start_point) + np.array([100., 50., 0.]),
0247                                    dz=0.0001, epsilon=args.epsilon)
0248
0249     # array for detailed data
0250     # number of materials cannot be pre-determined, so just assign a large number to be safe
0251     vals = np.zeros(shape=(len(etas), len(dets) + 1, 50))
0252     dets2 = dets + [OTHERS_NAME]
0253     det_mats = {d: [] for d in dets2}
0254     vt = ALLOWED_VALUE_TYPES[args.value_type]
0255
0256     for i, eta in enumerate(etas):
0257         if i % PROGRESS_STEP == 0:
0258             print('Scanned {:d}/{:d} for {:.2f} <= eta <= {:.2f}'.format(i, len(etas), etas[0], etas[-1]),
0259                   end='\r', flush=True)
0260         # scan material layers
0261         end_x = path_r*np.cos(phi/180.*np.pi)
0262         end_y = path_r*np.sin(phi/180.*np.pi)
0263         end_z = path_r*np.sinh(eta)
0264         # print('({:.2f}, {:.2f}, {:.2f})'.format(end_x, end_y, end_z))
0265         dfr = material_scan(desc, start_point, (end_x, end_y, end_z), epsilon=args.epsilon, int_dets=dets, thickness_corrector=th_corr)
0266
0267         # aggregated values for detectors
0268         x0_vals = dfr.groupby('detector')[vt[0]].sum().to_dict()
0269         for j, det in enumerate(dets2):
0270             vals[i, j, 0] = x0_vals.get(det, 0.)
0271             # update material dict
0272             dfd = dfr[dfr['detector'] == det]
0273             x0_mats = dfd.groupby('material')[vt[0]].sum()
0274             for mat in x0_mats.index:
0275                 if mat not in det_mats[det]:
0276                     det_mats[det] = det_mats[det] + [mat]
0277             for k, mat in enumerate(det_mats[det]):
0278                 vals[i, j, k + 1] = x0_mats.to_dict().get(mat, 0.)
0279
0280     print('Scanned {:d}/{:d} lines for {:.2f} < eta < {:.2f}'.format(len(etas), len(etas), etas[0], etas[-1]))
0281
0282     # aggregated data for plots
0283     dfa = pd.DataFrame(data=vals[:, :, 0], columns=dets2, index=etas)
0284     # sort columns by the total thickness (over eta range)
0285     sdets = dets
0286     if sort_dets:
0287         sdets = [d for d in dfa.sum().sort_values(ascending=False).index if d != OTHERS_NAME]
0288         dfa = dfa[sdets + [OTHERS_NAME]]
0289     dfa.to_csv('material_scan_agg.csv')
0290
0291     # plots
0292     pdf = matplotlib.backends.backend_pdf.PdfPages('material_scan_details.pdf')
0293     fig, ax = plt.subplots(figsize=(16, 8), dpi=160,
0294                            gridspec_kw={'top': 0.9, 'bottom': 0.2, 'left': 0.08, 'right': 0.98})
0295
0296     # group some detectors into Others because not enough colors for them
0297     plot_dets = sdets
0298     if len(sdets) > len(COLORS):
0299         group_dets = sdets[len(COLORS):]
0300         dfa.loc[:, OTHERS_NAME] += dfa.loc[:, group_dets].sum(axis=1)
0301         print('Detectors {} are grouped into {} due to limited number of colors.'.format(group_dets, OTHERS_NAME))
0302         plot_dets = sdets[:len(COLORS)]
0303
0304     # plot
0305     bottom = np.zeros(len(dfa))
0306     width = np.mean(np.diff(dfa.index))
0307     for col, c in zip(plot_dets, COLORS):
0308         ax.fill_between(dfa.index, bottom, dfa[col].values + bottom, label=col, step='mid', color=c)
0309         bottom += dfa[col].values
0310     # only plot Others if there are any
0311     if dfa[OTHERS_NAME].sum() > 0.:
0312         ax.fill_between(dfa.index, bottom, dfa[OTHERS_NAME].values + bottom, label=OTHERS_NAME, step='mid', color=OTHERS_COLOR)
0313
0314     # formatting
0315     ax.tick_params(which='both', direction='in', labelsize=22)
0316     ax.set_xlabel('$\eta$', fontsize=22)
0317     ax.set_ylabel('{} ($R_{{xy}} \leq {:.1f}$ cm)'.format(vt[1], args.path_r), fontsize=22)
0318     ax.xaxis.set_major_locator(MultipleLocator(0.5))
0319     ax.xaxis.set_minor_locator(MultipleLocator(0.1))
0320     ax.yaxis.set_major_locator(MultipleLocator(vt[2]))
0321     ax.yaxis.set_minor_locator(MultipleLocator(vt[3]))
0322     ax.grid(ls=':', which='both')
0323     ax.set_axisbelow(False)
0324     ax.set_xlim(args.eta_min, args.eta_max)
0325     ax.set_ylim(0., ax.get_ylim()[1]*1.1)
0326     ax.legend(bbox_to_anchor=(0.0, 0.9, 1.0, 0.1), ncol=4, loc="upper center", fontsize=22,
0327           borderpad=0.2, labelspacing=0.2, columnspacing=0.6, borderaxespad=0.05, handletextpad=0.4)
0328     # ax.set_yscale('log')
0329     fig.savefig('material_scan.png')
0330     pdf.savefig(fig)
0331     plt.close(fig)
0332
0333     # detailed plot for each detector
0334     for j, det in enumerate(dets):
0335         dmats = det_mats[det]
0336         dfa = pd.DataFrame(data=vals[:, j, 1:len(dmats)+1], columns=dmats, index=etas)
0337         if not len(dmats):
0338             print('No material found for detector {} in this scan, skipped it.'.format(det))
0339             continue
0340         # sort columns by the total thickness (over eta range)
0341         mats = [d for d in dfa.sum().sort_values(ascending=False).index]
0342         dfa = dfa[mats]
0343         plot_mats = mats
0344         if len(mats) > len(COLORS):
0345             group_mats = mats[len(COLORS):]
0346             dfa.loc[:, OTHERS_NAME] = dfa.loc[:, group_mats].sum(axis=1)
0347             print('For detector {}, materials {} are grouped into {} due to limited number of colors.'.format(det, group_mats, OTHERS_NAME))
0348             plot_mats = mats[:len(COLORS)]
0349
0350         fig, ax = plt.subplots(figsize=(16, 8), dpi=160,
0351                            gridspec_kw={'top': 0.9, 'bottom': 0.2, 'left': 0.08, 'right': 0.98})
0352         bottom = np.zeros(len(dfa))
0353         width = np.mean(np.diff(dfa.index))
0354         for col, c in zip(plot_mats, COLORS):
0355             ax.fill_between(dfa.index, bottom, dfa[col].values + bottom, label=col, step='mid', color=c)
0356             bottom += dfa[col].values
0357         if OTHERS_NAME in dfa.columns and dfa[OTHERS_NAME].sum() > 0.:
0358             ax.fill_between(dfa.index, bottom, dfa[OTHERS_NAME].values + bottom, label=OTHERS_NAME, step='mid', color=OTHERS_COLOR)
0359
0360         ax.legend(bbox_to_anchor=(0.0, 0.9, 1.0, 0.1), ncol=6, loc="upper left", fontsize=18,
0361                   borderpad=0.3, labelspacing=0.2, columnspacing=0.8, borderaxespad=0.1, handletextpad=0.4)
0362         ax.tick_params(which='both', direction='in', labelsize=22)
0363         ax.set_xlabel('$\eta$', fontsize=22)
0364         ax.set_ylabel('{} ($R_{{xy}} \leq {:.1f}$ cm)'.format(vt[1], args.path_r), fontsize=22)
0365         ax.set_title(det, fontsize=22)
0366         ax.xaxis.set_major_locator(MultipleLocator(0.5))
0367         ax.xaxis.set_minor_locator(MultipleLocator(0.1))
0368         ax.yaxis.set_major_locator(MultipleLocator(vt[2]))
0369         ax.yaxis.set_minor_locator(MultipleLocator(vt[3]))
0370         ax.grid(ls=':', which='both')
0371         ax.set_axisbelow(False)
0372         ax.set_xlim(args.eta_min, args.eta_max)
0373         ax.set_ylim(0., ax.get_ylim()[1]*1.1)
0374         pdf.savefig(fig)
0375         plt.close(fig)
0376     # close PDF
0377     pdf.close()