Belle II Software light-2406-ragdoll
create_and_plot_magneticfield.py
1#!/usr/bin/env python3
2
3
10
11
15
16from matplotlib.backends.backend_pdf import PdfPages
17from matplotlib import pyplot as pl
18import os
19from basf2 import Path, process
20import ROOT as root
21import numpy as np
22
23import matplotlib as mpl
24# switch to non-interactive backend
25mpl.use("Agg")
26
27
30
31scan_type = "zx"
32stepsU = (1800, -400., 500.)
33stepsV = (1600, -400., 400.)
34fieldtype = "MagneticField"
35filename = f"fieldmap-{scan_type}.pdf"
36# if true, the magnetic field will be shown using a blue-white-red gradient,
37# white being no field.
38show_color = True
39# if contours is defined and non empty, contour lines at the specified field
40# strengths are drawn on top of the color map
41show_contours = None
42# colours for the chosen contours: can either be None, a single color name or a
43# list of color names matching the chosen contours
44contour_colors = 'k'
45# line width of contour lines. Can either be one number or an array of numbers
46# denoting the width of each line
47contour_width = 0.1
48
49# example: show 1.4 to 1.6 tesla contours in blue, cyan, green, magenta, red
50# show_contours = [1.4, 1.45, 1.5, 1.55, 1.6]
51# contour_colors = ['b', 'c', 'g', 'm', 'r']
52
53
56
57fieldmap_file = os.path.splitext(filename)[0] + "-fieldmap.root"
58material_file = os.path.splitext(filename)[0] + "-material.root"
59
60# set Materialscan plane to be consistent with magnetic field map
61if scan_type == "xy":
62 material_plane = [0, 0, -0.1, 1, 0, 0, 0, 1, 0]
63elif scan_type == "zx":
64 material_plane = [0, -0.1, 0, 0, 0, 1, 1, 0, 0]
65elif scan_type == "zy":
66 material_plane = [-0.1, 0, 0, 0, 0, 1, 0, 1, 0]
67
68main = Path()
69# create modules
70main.add_module('EventInfoSetter')
71main.add_module('Gearbox')
72geometry = main.add_module('Geometry')
73main.add_module('FullSim')
74# override field if we want something different than the default
75if fieldtype != "MagneticField":
76 geometry.param({
77 "useDB": False,
78 "excludedComponents": ["MagneticField"],
79 "additionalComponents": [fieldtype],
80 })
81
82# Field map creator
83main.add_module('CreateFieldMap', **{
84 'filename': fieldmap_file,
85 "type": scan_type,
86 "nU": stepsU[0],
87 "minU": stepsU[1],
88 "maxU": stepsU[2],
89 "nV": stepsV[0],
90 "minV": stepsV[1],
91 "maxV": stepsV[2],
92})
93
94# Material map creator
95main.add_module('MaterialScan', **{
96 'Filename': material_file,
97 'spherical': False,
98 'planar': True,
99 'planar.plane': 'custom',
100 'planar.custom': material_plane,
101 'planar.maxDepth': 0.2,
102 'planar.nU': stepsU[0],
103 'planar.minU': stepsU[1],
104 'planar.maxU': stepsU[2],
105 'planar.nV': stepsV[0],
106 'planar.minV': stepsV[1],
107 'planar.maxV': stepsV[2],
108 'planar.splitByMaterials': False,
109 'planar.ignored': ['Air', 'Vacuum', 'G4_AIR', 'ColdAir'],
110})
111
112# Process one event to perform the scans
113process(main)
114
115
118
119
120def get_hist_data(hist):
121 """Obtain the data and extent of a ROOT TH2 as a numpy array"""
122 nbinsX = hist.GetNbinsX()
123 nbinsY = hist.GetNbinsY()
124 data = np.frombuffer(hist.GetArray(), count=hist.GetSize())
125 # rehsape and strip overflow/underflow bins
126 data = data.reshape(nbinsY + 2, nbinsX + 2)[1:-1, 1:-1]
127
128 xmin = hist.GetXaxis().GetXmin()
129 xmax = hist.GetXaxis().GetXmax()
130 ymin = hist.GetYaxis().GetXmin()
131 ymax = hist.GetYaxis().GetXmax()
132
133 return data, (xmin, xmax, ymin, ymax)
134
135
136def plotMap(axes, data, extent, cmap, **kwargs):
137 """Plot either material or field map"""
138 return axes.imshow(data, extent=extent, origin="lower", cmap=cmap,
139 aspect="auto", interpolation="nearest", **kwargs)
140
141
142# open root files we just created
143rfieldmap_file = root.TFile(fieldmap_file)
144rmaterial_file = root.TFile(material_file)
145# get the material map
146mat_hist = rmaterial_file.Get("Planar/All_Regions_x0")
147mat_data, mat_extent = get_hist_data(mat_hist)
148# obtain nice colormaps
149cmap_fieldmap = mpl.cm.get_cmap('seismic')
150cmap_material = mpl.cm.get_cmap('binary')
151# create pdf file and fill with plots
152pdf = PdfPages(filename)
153for component in ["Bx", "By", "Bz", "B"]:
154 print(f"Create {component} fieldmap plot")
155 h = rfieldmap_file.Get(component)
156 data, extent = get_hist_data(h)
157 # we need the scale to be symmetric so find out what is the maximal
158 # absolute value and we choose a minimum scale
159 field_max = max(np.abs(data).max(), 0.1)
160 f = pl.figure()
161 a = f.add_subplot(111)
162 # plot material
163 lognorm = mpl.colors.LogNorm()
164 plotMap(a, mat_data, mat_extent, cmap_material, norm=lognorm)
165 # and overlay magnetic field
166 if show_color:
167 m = plotMap(a, data, extent, cmap_fieldmap, alpha=0.9,
168 vmin=-field_max, vmax=field_max)
169 # add a colorbar and set the label
170 cb = f.colorbar(m)
171 if len(component) == 1:
172 cb.set_label(r"$B/{\rm T}$")
173 else:
174 cb.set_label(r"$B_%s/{\rm T}$" % component[1])
175
176 if show_contours is not None and len(show_contours) > 0:
177 # get bin centers in x and y
178 x, xstep = np.linspace(extent[0], extent[1], h.GetNbinsX() + 1,
179 retstep=True)
180 x += xstep / 2
181 y, ystep = np.linspace(extent[2], extent[3], h.GetNbinsY() + 1,
182 retstep=True)
183 y += ystep / 2
184 x, y = np.meshgrid(x[:-1], y[:-1])
185 # draw contour plot and label it
186 cs = a.contour(x, y, data, levels=show_contours, colors=contour_colors,
187 linewidths=contour_width)
188 a.clabel(cs, fmt="%.2f T", fontsize=6)
189
190 # label the plot
191 a.set_title(h.GetTitle())
192 a.set_xlabel(r"$%s/{\rm cm}$" % scan_type[0])
193 a.set_ylabel(r"$%s/{\rm cm}$" % scan_type[1])
194 # findally, save the plot
195 pdf.savefig(f)
196
197pdf.close()