10 from basf2
import Path, process
14 import matplotlib
as mpl
17 from matplotlib
import pyplot
as pl
18 from matplotlib.backends.backend_pdf
import PdfPages
25 stepsU = (1800, -400., 500.)
26 stepsV = (1600, -400., 400.)
27 fieldtype =
"MagneticField"
28 filename =
"fieldmap-%s.pdf" % scan_type
50 fieldmap_file = os.path.splitext(filename)[0] +
"-fieldmap.root"
51 material_file = os.path.splitext(filename)[0] +
"-material.root"
55 material_plane = [0, 0, -0.1, 1, 0, 0, 0, 1, 0]
56 elif scan_type ==
"zx":
57 material_plane = [0, -0.1, 0, 0, 0, 1, 1, 0, 0]
58 elif scan_type ==
"zy":
59 material_plane = [-0.1, 0, 0, 0, 0, 1, 0, 1, 0]
63 main.add_module(
'EventInfoSetter')
64 main.add_module(
'Gearbox')
65 geometry = main.add_module(
'Geometry')
66 main.add_module(
'FullSim')
68 if fieldtype !=
"MagneticField":
71 "excludedComponents": [
"MagneticField"],
72 "additionalComponents": [fieldtype],
76 main.add_module(
'CreateFieldMap', **{
77 'filename': fieldmap_file,
88 main.add_module(
'MaterialScan', **{
89 'Filename': material_file,
92 'planar.plane':
'custom',
93 'planar.custom': material_plane,
94 'planar.maxDepth': 0.2,
95 'planar.nU': stepsU[0],
96 'planar.minU': stepsU[1],
97 'planar.maxU': stepsU[2],
98 'planar.nV': stepsV[0],
99 'planar.minV': stepsV[1],
100 'planar.maxV': stepsV[2],
101 'planar.splitByMaterials':
False,
102 'planar.ignored': [
'Air',
'Vacuum',
'G4_AIR',
'ColdAir'],
113 def get_hist_data(hist):
114 """Obtain the data and extent of a ROOT TH2 as a numpy array"""
115 nbinsX = hist.GetNbinsX()
116 nbinsY = hist.GetNbinsY()
117 data = np.frombuffer(hist.GetArray(), count=hist.GetSize())
119 data = data.reshape(nbinsY + 2, nbinsX + 2)[1:-1, 1:-1]
121 xmin = hist.GetXaxis().GetXmin()
122 xmax = hist.GetXaxis().GetXmax()
123 ymin = hist.GetYaxis().GetXmin()
124 ymax = hist.GetYaxis().GetXmax()
126 return data, (xmin, xmax, ymin, ymax)
129 def plotMap(axes, data, extent, cmap, **kwargs):
130 """Plot either material or field map"""
131 return axes.imshow(data, extent=extent, origin=
"lower", cmap=cmap,
132 aspect=
"auto", interpolation=
"nearest", **kwargs)
136 rfieldmap_file = root.TFile(fieldmap_file)
137 rmaterial_file = root.TFile(material_file)
139 mat_hist = rmaterial_file.Get(
"Planar/All_Regions_x0")
140 mat_data, mat_extent = get_hist_data(mat_hist)
142 cmap_fieldmap = mpl.cm.get_cmap(
'seismic')
143 cmap_material = mpl.cm.get_cmap(
'binary')
145 pdf = PdfPages(filename)
146 for component
in [
"Bx",
"By",
"Bz",
"B"]:
147 print(
"Create %s fieldmap plot" % component)
148 h = rfieldmap_file.Get(component)
149 data, extent = get_hist_data(h)
152 field_max = max(np.abs(data).max(), 0.1)
154 a = f.add_subplot(111)
156 lognorm = mpl.colors.LogNorm()
157 plotMap(a, mat_data, mat_extent, cmap_material, norm=lognorm)
160 m = plotMap(a, data, extent, cmap_fieldmap, alpha=0.9,
161 vmin=-field_max, vmax=field_max)
164 if len(component) == 1:
165 cb.set_label(
r"$B/{\rm T}$")
167 cb.set_label(
r"$B_%s/{\rm T}$" % component[1])
169 if show_contours
is not None and len(show_contours) > 0:
171 x, xstep = np.linspace(extent[0], extent[1], h.GetNbinsX() + 1,
174 y, ystep = np.linspace(extent[2], extent[3], h.GetNbinsY() + 1,
177 x, y = np.meshgrid(x[:-1], y[:-1])
179 cs = a.contour(x, y, data, levels=show_contours, colors=contour_colors,
180 linewidths=contour_width)
181 a.clabel(cs, fmt=
"%.2f T", fontsize=6)
184 a.set_title(h.GetTitle())
185 a.set_xlabel(
r"$%s/{\rm cm}$" % scan_type[0])
186 a.set_ylabel(
r"$%s/{\rm cm}$" % scan_type[1])