Belle II Software  release-05-01-25
create_and_plot_magneticfield.py
1 #!/usr/bin/env python3
2 # -*- coding: utf-8 -*-
3 
4 
8 
9 import os
10 from basf2 import Path, process
11 import ROOT as root
12 import numpy as np
13 
14 import matplotlib as mpl
15 # switch to non-interactive backend
16 mpl.use("Agg")
17 from matplotlib import pyplot as pl
18 from matplotlib.backends.backend_pdf import PdfPages
19 
20 
23 
24 scan_type = "zx"
25 stepsU = (1800, -400., 500.)
26 stepsV = (1600, -400., 400.)
27 fieldtype = "MagneticField"
28 filename = "fieldmap-%s.pdf" % scan_type
29 # if true, the magnetic field will be shown using a blue-white-red gradient,
30 # white being no field.
31 show_color = True
32 # if contours is defined and non empty, contour lines at the specified field
33 # strengths are drawn on top of the color map
34 show_contours = None
35 # colours for the chosen contours: can either be None, a single color name or a
36 # list of color names matching the chosen contours
37 contour_colors = 'k'
38 # line width of contour lines. Can either be one number or an array of numbers
39 # denoting the width of each line
40 contour_width = 0.1
41 
42 # example: show 1.4 to 1.6 tesla contours in blue, cyan, green, magenta, red
43 # show_contours = [1.4, 1.45, 1.5, 1.55, 1.6]
44 # contour_colors = ['b', 'c', 'g', 'm', 'r']
45 
46 
49 
50 fieldmap_file = os.path.splitext(filename)[0] + "-fieldmap.root"
51 material_file = os.path.splitext(filename)[0] + "-material.root"
52 
53 # set Materialscan plane to be consistent with magnetic field map
54 if scan_type == "xy":
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]
60 
61 main = Path()
62 # create modules
63 main.add_module('EventInfoSetter')
64 main.add_module('Gearbox')
65 geometry = main.add_module('Geometry')
66 main.add_module('FullSim')
67 # override field if we want something different than the default
68 if fieldtype != "MagneticField":
69  geometry.param({
70  "useDB": False,
71  "excludedComponents": ["MagneticField"],
72  "additionalComponents": [fieldtype],
73  })
74 
75 # Field map creator
76 main.add_module('CreateFieldMap', **{
77  'filename': fieldmap_file,
78  "type": scan_type,
79  "nU": stepsU[0],
80  "minU": stepsU[1],
81  "maxU": stepsU[2],
82  "nV": stepsV[0],
83  "minV": stepsV[1],
84  "maxV": stepsV[2],
85 })
86 
87 # Material map creator
88 main.add_module('MaterialScan', **{
89  'Filename': material_file,
90  'spherical': False,
91  'planar': True,
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'],
103 })
104 
105 # Process one event to perform the scans
106 process(main)
107 
108 
111 
112 
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())
118  # rehsape and strip overflow/underflow bins
119  data = data.reshape(nbinsY + 2, nbinsX + 2)[1:-1, 1:-1]
120 
121  xmin = hist.GetXaxis().GetXmin()
122  xmax = hist.GetXaxis().GetXmax()
123  ymin = hist.GetYaxis().GetXmin()
124  ymax = hist.GetYaxis().GetXmax()
125 
126  return data, (xmin, xmax, ymin, ymax)
127 
128 
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)
133 
134 
135 # open root files we just created
136 rfieldmap_file = root.TFile(fieldmap_file)
137 rmaterial_file = root.TFile(material_file)
138 # get the material map
139 mat_hist = rmaterial_file.Get("Planar/All_Regions_x0")
140 mat_data, mat_extent = get_hist_data(mat_hist)
141 # obtain nice colormaps
142 cmap_fieldmap = mpl.cm.get_cmap('seismic')
143 cmap_material = mpl.cm.get_cmap('binary')
144 # create pdf file and fill with plots
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)
150  # we need the scale to be symmetric so find out what is the maximal
151  # absolute value and we choose a minimum scale
152  field_max = max(np.abs(data).max(), 0.1)
153  f = pl.figure()
154  a = f.add_subplot(111)
155  # plot material
156  lognorm = mpl.colors.LogNorm()
157  plotMap(a, mat_data, mat_extent, cmap_material, norm=lognorm)
158  # and overlay magnetic field
159  if show_color:
160  m = plotMap(a, data, extent, cmap_fieldmap, alpha=0.9,
161  vmin=-field_max, vmax=field_max)
162  # add a colorbar and set the label
163  cb = f.colorbar(m)
164  if len(component) == 1:
165  cb.set_label(r"$B/{\rm T}$")
166  else:
167  cb.set_label(r"$B_%s/{\rm T}$" % component[1])
168 
169  if show_contours is not None and len(show_contours) > 0:
170  # get bin centers in x and y
171  x, xstep = np.linspace(extent[0], extent[1], h.GetNbinsX() + 1,
172  retstep=True)
173  x += xstep / 2
174  y, ystep = np.linspace(extent[2], extent[3], h.GetNbinsY() + 1,
175  retstep=True)
176  y += ystep / 2
177  x, y = np.meshgrid(x[:-1], y[:-1])
178  # draw contour plot and label it
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)
182 
183  # label the plot
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])
187  # findally, save the plot
188  pdf.savefig(f)
189 
190 pdf.close()