Select Git revision
plot_ce_ecal_placement.py
Chao Peng authored
plot_ce_ecal_placement.py 11.26 KiB
'''
A script to calcualte placement of ecal endcap modules
lxml is not included in container, get it by simply typing 'pip install lxml'
Author: Chao Peng (ANL)
Date: 06/17/2021
'''
import os
import numpy as np
import argparse
import DDG4
from lxml import etree as ET
from matplotlib import pyplot as plt
from matplotlib.collections import PatchCollection
from matplotlib.patches import Rectangle, Circle
CRYSTAL_SIZE = (20., 20., 200.) # mm
CRYSTAL_GAP = 0.5 # mm
CRYSTAL_ALIGNMENT = [
(5, 21), (5, 21), (5, 21), (4, 22),
(3, 23), (0, 26), (0, 24), (0, 24),
(0, 24), (0, 24), (0, 24), (0, 24),
(0, 22), (0, 22), (0, 20), (0, 20),
(0, 18), (0, 18), (0, 16), (0, 16),
(0, 14), (0, 14), (0, 12), (0, 12),
(0, 6), (0, 6),
]
GLASS_SIZE = (40., 40., 400.) # mm
GLASS_GAP = 1.0 # mm
GLASS_ALIGNMENT = [
(13, 10), (13, 10), (13, 10), (12, 10),
(12, 10), (12, 10), (11, 11), (10, 11),
(9, 12), (8, 13), (7, 13), (6, 14),
(3, 16), (0, 18), (0, 18), (0, 16),
(0, 16), (0, 14), (0, 13), (0, 11),
(0, 10), (0, 7), (0, 3),
]
# calculate positions of modules with a quad-alignment and module size
def individual_placement(alignment, module_x, module_y, gap=0.):
placements = []
for row, (start, num) in enumerate(alignment):
for col in np.arange(start, start + num):
placements.append(((col + 0.5)*(module_y + gap), (row + 0.5)*(module_x + gap)))
placements = np.asarray(placements)
return np.vstack((placements,
np.vstack((placements.T[0]*-1., placements.T[1])).T,
np.vstack((placements.T[0], placements.T[1]*-1.)).T,
np.vstack((placements.T[0]*-1., placements.T[1]*-1.)).T))
def draw_placement(axis, colors=('teal'), module_alignment=((CRYSTAL_SIZE, CRYSTAL_GAP, CRYSTAL_ALIGNMENT))):
xmin, ymin, xmax, ymax = 0., 0., 0., 0.
patches = []
numbers = []
for color, (mod_size, mod_gap, alignment) in zip(colors, module_alignment):
placements = individual_placement(alignment, *mod_size[:2], mod_gap)
boxes = [Rectangle((x - (mod_size[0] + mod_gap)/2., y - (mod_size[1] + mod_gap)/2.), mod_size[0], mod_size[1])
for x, y in placements]
patches.append(Rectangle((0., 0.), *mod_size[:2], facecolor=color, alpha=0.5, edgecolor='k'))
numbers.append(len(placements))
pc = PatchCollection(boxes, facecolor=color, alpha=0.5, edgecolor='k')
xmin = min(xmin, placements.T[0].min() - 8.*(mod_size[0] + mod_gap))
ymin = min(ymin, placements.T[1].min() - 8.*(mod_size[1] + mod_gap))
xmax = max(xmax, placements.T[0].max() + 8.*(mod_size[0] + mod_gap))
ymax = max(ymax, placements.T[1].max() + 8.*(mod_size[1] + mod_gap))
# Add collection to axes
axis.add_collection(pc)
axis.set_xlim(xmin, xmax)
axis.set_ylim(ymin, ymax)
return axis, patches, numbers
def compact_constants(path, names):
if not os.path.exists(path):
print('Cannot find compact file \"{}\".'.format(path))
return []
kernel = DDG4.Kernel()
description = kernel.detectorDescription()
kernel.loadGeometry("file:{}".format(path))
try:
vals = [description.constantAsDouble(n) for n in names]
except:
print('Fail to extract values from {}, return empty.'.format(names))
vals = []
kernel.terminate()
return vals
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('-s', '--save', default='ce_ecal_placement_test.xml',
help='path to save compact file.')
parser.add_argument('-c', '--compact', default='',
help='compact file to get contant to plot')
parser.add_argument('--radii-constants', dest='radii', default='EcalBarrel_rmin',
help='constant names in compact file to plot, seprate by \",\"')
parser.add_argument('--individual', dest='indiv', action='store_true',
help='individual block placements instead of line placements')
args = parser.parse_args()
data = ET.Element('lccdd')
defines = ET.SubElement(data, 'define')
# constants: name, value
CONSTANTS = [
('CrystalModule_sx', '{:.2f}*mm'.format(CRYSTAL_SIZE[0])),
('CrystalModule_sy', '{:.2f}*mm'.format(CRYSTAL_SIZE[1])),
('CrystalModule_sz', '{:.2f}*mm'.format(CRYSTAL_SIZE[2])),
('CrystalModule_wrap', '{:.2f}*mm'.format(CRYSTAL_GAP)),
('GlassModule_sx', '{:.2f}*mm'.format(GLASS_SIZE[0])),
('GlassModule_sy', '{:.2f}*mm'.format(GLASS_SIZE[1])),
('GlassModule_sz', '{:.2f}*mm'.format(GLASS_SIZE[2])),
('GlassModule_wrap', '{:.2f}*mm'.format(GLASS_GAP)),
('CrystalModule_z0', '10.*cm'),
('GlassModule_z0', '0.0*cm'),
('EcalEndcapN_z0', '-EcalEndcapN_zmin-max(CrystalModule_sz,GlassModule_sz)/2.'),
('CrystalModule_dx', 'CrystalModule_sx + CrystalModule_wrap'),
('CrystalModule_dy', 'CrystalModule_sy + CrystalModule_wrap'),
('GlassModule_dx', 'GlassModule_sx + GlassModule_wrap'),
('GlassModule_dy', 'GlassModule_sy + GlassModule_wrap'),
]
# line-by-line alignment start pos, total number of blocks
for name, value in CONSTANTS:
constant = ET.SubElement(defines, 'constant')
constant.set('name', name)
constant.set('value', value)
# this value will be used multiple times, so define it here
readout_name = 'EcalEndcapNHits'
# detector and its dimension/position/rotation
dets = ET.SubElement(data, 'detectors')
cmt = ET.SubElement(dets, 'comment')
cmt.text = ' Backwards Endcap EM Calorimeter, placements generated by script '
det = ET.SubElement(dets, 'detector')
det.set('id', 'ECalEndcapN_ID')
det.set('name', 'EcalEndcapN')
det.set('type', 'HomogeneousCalorimeter')
det.set('readout', readout_name)
pos = ET.SubElement(det, 'position')
pos.set('x', '0')
pos.set('y', '0')
pos.set('z', 'EcalEndcapN_z0')
rot = ET.SubElement(det, 'rotation')
rot.set('x', '0')
rot.set('y', '0')
rot.set('z', '0')
# placements of modules
plm = ET.SubElement(det, 'placements')
pltype = 'individuals' if args.indiv else 'lines'
# crystal
crystal = ET.SubElement(plm, pltype)
crystal.set('sector', '1')
crystal_mod = ET.SubElement(crystal, 'module')
crystal_mod.set('sizex', 'CrystalModule_sx')
crystal_mod.set('sizey', 'CrystalModule_sy')
crystal_mod.set('sizez', 'CrystalModule_sz')
crystal_mod.set('material', 'PbWO4')
crystal_mod.set('vis', 'AnlTeal')
crystal_wrap = ET.SubElement(crystal, 'wrapper')
crystal_wrap.set('thickness', 'CrystalModule_wrap')
crystal_wrap.set('material', 'Epoxy')
crystal_wrap.set('vis', 'WhiteVis')
# crystal placements (for individuals)
if args.indiv:
for m, (x, y) in enumerate(individual_placement(CRYSTAL_ALIGNMENT, *CRYSTAL_SIZE[:2], CRYSTAL_GAP)):
module = ET.SubElement(crystal, 'placement')
module.set('x', '{:.3f}*mm'.format(x))
module.set('y', '{:.3f}*mm'.format(y))
module.set('z', 'CrystalModule_z0')
module.set('id', '{:d}'.format(m))
# crystal placements (for lines)
else:
crystal.set('mirrorx', 'true')
crystal.set('mirrory', 'true')
for row, (begin, nmods) in enumerate(CRYSTAL_ALIGNMENT):
line = ET.SubElement(crystal, 'line')
line.set('axis', 'x')
line.set('x', 'CrystalModule_dx/2.')
line.set('y', 'CrystalModule_dy*{:d}/2.'.format(row*2 + 1))
line.set('z', 'CrystalModule_z0')
line.set('begin', '{:d}'.format(begin))
line.set('nmods', '{:d}'.format(nmods))
# glass
glass = ET.SubElement(plm, pltype)
glass.set('sector', '2')
glass_mod = ET.SubElement(glass, 'module')
glass_mod.set('sizex', 'GlassModule_sx')
glass_mod.set('sizey', 'GlassModule_sy')
glass_mod.set('sizez', 'GlassModule_sz')
# TODO: change glass material
glass_mod.set('material', 'PbGlass')
glass_mod.set('vis', 'AnlBlue')
glass_wrap = ET.SubElement(glass, 'wrapper')
glass_wrap.set('thickness', 'GlassModule_wrap')
glass_wrap.set('material', 'Epoxy')
glass_wrap.set('vis', 'WhiteVis')
# crystal placements (for individuals)
if args.indiv:
for m, (x, y) in enumerate(individual_placement(GLASS_ALIGNMENT, *GLASS_SIZE[:2], GLASS_GAP)):
module = ET.SubElement(glass, 'placement')
module.set('x', '{:.3f}*mm'.format(x))
module.set('y', '{:.3f}*mm'.format(y))
module.set('z', 'GlassModule_z0')
module.set('id', '{:d}'.format(m))
# crystal placements (for lines)
else:
glass.set('mirrorx', 'true')
glass.set('mirrory', 'true')
for row, (begin, nmods) in enumerate(GLASS_ALIGNMENT):
line = ET.SubElement(glass, 'line')
line.set('axis', 'x')
line.set('x', 'GlassModule_dx/2.')
line.set('y', 'GlassModule_dy*{:d}/2.'.format(row*2 + 1))
line.set('z', 'GlassModule_z0')
line.set('begin', '{:d}'.format(begin))
line.set('nmods', '{:d}'.format(nmods))
# readout
readouts = ET.SubElement(data, 'readouts')
cmt = ET.SubElement(readouts, 'comment')
cmt.text = 'Effectively no segmentation, the segmentation is used to provide cell dimension info'
readout = ET.SubElement(readouts, 'readout')
readout.set('name', readout_name)
seg = ET.SubElement(readout, 'segmentation')
# need segmentation to provide cell dimension info
# seg.set('type', 'NoSegmentation')
seg.set('type', 'MultiSegmentation')
seg.set('key', 'sector')
crystal_seg = ET.SubElement(seg, 'segmentation')
crystal_seg.set('name', 'CrystalSeg')
crystal_seg.set('key_value', '1')
crystal_seg.set('type', 'CartesianGridXY')
crystal_seg.set('grid_size_x', 'CrystalModule_dx')
crystal_seg.set('grid_size_y', 'CrystalModule_dy')
glass_seg = ET.SubElement(seg, 'segmentation')
glass_seg.set('name', 'GlassSeg')
glass_seg.set('key_value', '2')
glass_seg.set('type', 'CartesianGridXY')
glass_seg.set('grid_size_x', 'GlassModule_dx')
glass_seg.set('grid_size_y', 'GlassModule_dy')
rid = ET.SubElement(readout, 'id')
rid.text = 'system:8,sector:4,module:20,x:32:-16,y:-16'
text = ET.tostring(data, pretty_print=True)
with open(args.save, 'wb') as f:
f.write(text)
fig, ax = plt.subplots(figsize=(12, 12), dpi=160)
ax, patches, nblocks = draw_placement(ax, ['teal', 'royalblue'],
[(CRYSTAL_SIZE, CRYSTAL_GAP, CRYSTAL_ALIGNMENT), (GLASS_SIZE, GLASS_GAP, GLASS_ALIGNMENT)])
ax.set_xlabel('x (mm)', fontsize=24)
ax.set_ylabel('y (mm)', fontsize=24)
ax.tick_params(direction='in', labelsize=22, which='both')
ax.set_axisbelow(True)
ax.grid(linestyle=':', which='both')
ax.legend(patches, ['{} {}'.format(num, name) for num, name in zip(nblocks, ['PbWO$_4$', 'SciGlass'])], fontsize=24)
if args.compact and args.radii:
names = [c.strip() for c in args.radii.split(',')]
radii = compact_constants(args.compact, names)
for name, radius in zip(names, radii):
ax.add_patch(Circle((0, 0), radius*10., facecolor='none', edgecolor='k', linewidth=2))
ax.annotate(name, xy=(radius*10/1.4, radius*10/1.4), fontsize=22)
fig.savefig('ce_ecal_placement.png')