#!/usr/bin/env python
# This file is part of MAUS: http://micewww.pp.rl.ac.uk:8080/projects/maus
#
# MAUS is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# MAUS is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with MAUS. If not, see .
"""
Geometry Validation tool is used for building a 2D map of the materials that
particles see as they traverse the beamline and then plotting it. Assumes a
straight geometry along z.
The algorithm is to:
* Force keep_steps = True so that MAUS reports position and material at every
step
* User specified initial z position and a line of points in x, y
* Force physics_processes = None so no energy loss, scattering etc can occur
* Fire a very high energy gamma down the beamline for each initial (x, y, z) in
the z direction
* Plot as a function of x, y, z the resultant map of materials in the beamline
** Colours are chosen automatically (this can sometimes be not great)
* Plot as a function of z a 1D graph of materials on the beamline axis
Uses the "geometry_validation" dictionary of datacards, with dictionary keys:
- "file_name", (string) name of the file that will be used to store the tracking
data
- "will_plot", (bool) set to boolean true to make plots of the particles
- "will_track" (bool) set to boolean true to do tracking, else attempts to reuse
old tracking
- "z_start" (double, mm) start z-position for tracking and plotting
- "z_end" (double, mm) end z-position for plotting (tracking will keep going
until it reaches the edge of the geometry)
- "x_start" (double, mm) x-position of the first particle to be tracked
- "x_step" (double, mm) increment in x-position for each subsequent particle
- "y_start" (double, mm) y-position of the first particle to be tracked
- "y_step" (double, mm) increment in y-position for each subsequent particle
- "n_steps" (int) number of particles to track
- "plot_formats" (string) list of plot formats, e.g. ["png", "eps"]
would plot in png and encapsulated postscript format
- "geometry_validation_materials_1d" (string) name of the 1D plot to be printed
by MAUS. The appropriate suffix will be appended for each format.
- "geometry_validation_materials_2d" (string) name of the 2D plot to be printed
by MAUS. The appropriate suffix will be appended for each format.
- "geometry_validation_volumes_1d" (string) name of the 1D plot to be printed
by MAUS giving z positions of volumes. The appropriate suffix will
be appended for each format.
- "geometry_validation_volumes_2d" (string) name of the 2D plot to be printed
by MAUS showing volume bounding rectangles and name labels. The
appropriate suffix will be appended for each format.
- "2d_volume_plot_label_size" (float) between 0.0 and 1.0 that tells how big the
labels should be in the 2d volume plot. Labels are stacked vertically, so if
the label size is small more labels can be put on the plot
- "volume_bounding_box_dump" (string) file name for dumping bounding boxes to.
Format is json, going something something like:
{
: {
"z_min": ,
"z_max": ,
"r_min": ,
"r_max": ,
"material":
},
...
}
e.g.
{
"Vol14": {
"z_max": 320.0,
"r_min": 0.0,
"material": "G4_Al",
"r_max": 2.23606806280132,
"z_min": 280.0
},
...
}
"""
# pylint: disable=E1101
import os
import sys
import json
import ROOT
import xboa.Common as common
import Configuration
import maus_cpp.simulation
import maus_cpp.globals
class Plotting(object): # pylint: disable=R0902
"""Plotting class handles reading and plotting of the tracking data"""
def __init__(self, configuration):
"""
Initialise the plotting class
- configuration should be a valid datacard set in string representation
"""
py_config = json.loads(configuration)["geometry_validation"]
print json.dumps(py_config, indent=2)
self.file_name = py_config["file_name"]
self.z_start = py_config["z_start"]
self.z_end = py_config["z_end"]
self.will_plot = py_config["will_plot"]
self.n_steps = py_config["n_steps"]
self.formats = py_config["plot_formats"]
self.mat_1d = os.path.expandvars(py_config["1d_material_plot"])
self.mat_2d = os.path.expandvars(py_config["2d_material_plot"])
self.vol_1d = os.path.expandvars(py_config["1d_volume_plot"])
self.vol_2d = os.path.expandvars(py_config["2d_volume_plot"])
self.vol_2d_label_size = py_config["2d_volume_plot_label_size"]
self.volume_bb_dump = py_config["volume_bounding_box_dump"]
self.material_data = {}
self.volume_data = {}
self.volume_label_index = 0
def do_plotting(self):
"""Load the data and plot"""
if not self.will_plot:
return
print "Doing plots"
# step_data should be member
step_data = self._load_data()
for key in step_data:
print "Found", len(step_data[key]), \
"steps for material '"+str(key)+"'"
self._bb_dump()
material_dict = self.get_material_dict(step_data)
data_range = self.get_range(step_data)
x_range = data_range[0]
z_range = data_range[1]
z_range[0] = max(z_range[0], self.z_start)
z_range[1] = min(z_range[1], self.z_end)
print "data range r:", x_range, "z:", z_range
self._draw_materials_1d(step_data, z_range)
self._draw_materials_2d(step_data, material_dict, x_range, z_range)
self._draw_volumes_1d(z_range)
self._draw_volumes_2d(x_range, z_range)
def _readline_generator(self):
"""
Generator that reads a file line-by-line, with a progress bar
Python builtin readlines() reads the file in one go, making big memory
usage for potentially large files
"""
bar_length = 72
sys.stdout.write("["+" "*(bar_length+1)+"]")
sys.stdout.flush()
sys.stdout.write("\b"*(bar_length+1))
old_done = 0
fin = open(self.file_name)
line = None
i = 0
while True:
line = fin.readline()
if line == "":
break
yield line
done = float(i)/(self.n_steps-1)*bar_length
while done-old_done > 1.:
old_done += 1
sys.stdout.write("=")
sys.stdout.flush()
i += 1
fin.close()
print
def _load_data(self):
"""Load the data"""
step_data = {}
print "Loading data"
for line in self._readline_generator():
json_data = json.loads(line)
for event in json_data:
for track in event["tracks"]:
for step in track["steps"]:
z_pos = step["position"]["z"]
if z_pos < self.z_start or z_pos > self.z_end:
continue
r_pos = self._get_r(step["position"])
mat = step["material"]
if not mat in step_data.keys():
step_data[mat] = []
step_data[mat].append({"r":r_pos, "z":z_pos})
vol = step["volume"]
if not vol in self.volume_data.keys():
self.volume_data[vol] = {
"r_min":r_pos,
"r_max":r_pos,
"z_min":z_pos,
"z_max":z_pos,
"material":mat,
}
else:
if r_pos < self.volume_data[vol]["r_min"]:
self.volume_data[vol]["r_min"] = r_pos
elif r_pos > self.volume_data[vol]["r_max"]:
self.volume_data[vol]["r_max"] = r_pos
if z_pos < self.volume_data[vol]["z_min"]:
self.volume_data[vol]["z_min"] = z_pos
elif z_pos > self.volume_data[vol]["z_max"]:
self.volume_data[vol]["z_max"] = z_pos
return step_data
def _bb_dump(self):
"""Dump bounding box data to disk"""
fout = open(self.volume_bb_dump, "w")
print >> fout, json.dumps(self.volume_data, indent=2)
def _draw_materials_1d(self, step_data, z_range):
"""Make a 1D plot of position vs material along the z axis"""
# we plot G4_AIR first, then G4_GALACTIC?
draw_list = step_data.keys()
draw_list = self.move(draw_list, "G4_AIR", 0)
draw_list = self.move(draw_list, "G4_Galactic", 1)
black_x_list, black_y_list = [], []
grey_x_list, grey_y_list = [], []
canvas = common.make_root_canvas("Material vs z")
hist = ROOT.TH2D("name", "Material vs z", 30000, z_range[0], z_range[1],
len(draw_list)+1, -0.5, len(draw_list)+0.5)
common._hist_persistent.append(hist) #pylint: disable=W0212
for material in draw_list:
hist.Fill(z_range[0], material, 1.)
hist.SetBit(ROOT.TH1.kCanRebin)
hist.SetStats(0)
hist.LabelsDeflate("Y")
hist.LabelsOption("v")
hist.Draw()
for i, material in enumerate(draw_list):
grey_x_list += [i+0.1 for step in step_data[material]]
grey_y_list += [step["z"] for step in step_data[material]]
black_x_list += [i for step in step_data[material] \
if step["r"] < 1e-9]
black_y_list += [step["z"] for step in step_data[material] \
if step["r"] < 1e-9]
if len(grey_x_list) == 0:
return
hist, graph = common.make_root_graph("", grey_y_list, "z [mm]",
grey_x_list, "material index")
graph.SetMarkerColor(ROOT.TColor.GetColor(0.5, 0.5, 0.5))
graph.SetMarkerStyle(7)
graph.Draw('p')
if len(black_x_list) != 0:
hist, graph = common.make_root_graph("", black_y_list, "z [mm]",
black_x_list, "material index")
graph.SetMarkerStyle(7)
graph.Draw('p')
canvas.Update()
for _format in self.formats:
canvas.Print(self.mat_1d+"."+_format)
def _draw_volumes_1d(self, z_range):
"""Make a 1D plot of position vs volume along the z axis"""
draw_list = sorted(self.volume_data.keys(), self._sort_cmp)
x_list = []
canvas = common.make_root_canvas("Volume Name vs z")
# set up axis labels
print "VOLUME PLOT", z_range
hist = ROOT.TH2D("name", "Volume Name vs z", 30000,
z_range[0], z_range[1],
len(draw_list)+1, -0.5, len(draw_list)+0.5)
common._hist_persistent.append(hist) #pylint: disable=W0212
for material in draw_list:
hist.Fill(z_range[0], material, 1.)
hist.SetBit(ROOT.TH1.kCanRebin)
hist.SetStats(0)
hist.LabelsDeflate("Y")
hist.LabelsOption("v")
hist.Draw()
# now draw graphs
for i, volume in enumerate(draw_list):
x_list = [self.volume_data[volume]["z_min"],
self.volume_data[volume]["z_max"]]
y_list = [i+0.1, i+0.1]
hist, graph = common.make_root_graph("", x_list, "z [mm]",
y_list, "volume index")
graph.SetLineColor(ROOT.TColor.GetColor(0.5, 0.5, 0.5))
graph.Draw('l')
if self.volume_data[volume]["r_min"] < 1e-9:
y_list = [i, i]
hist, graph = common.make_root_graph("", x_list, "z [mm]",
y_list, "material index")
graph.Draw('l')
canvas.Update()
for _format in self.formats:
canvas.Print(self.vol_1d+"."+_format)
def _sort_cmp(self, vol_key_1, vol_key_2):
"""Compare for sorting volume keys by z_min"""
return cmp(self.volume_data[vol_key_1]["z_min"],
self.volume_data[vol_key_2]["z_min"])
def _draw_volume_label(self, x_position, volume, x_range, color):
"""Draw a volume label at x_position"""
self.labels.append(ROOT.TLatex())
delta = abs(x_range[1]-x_range[0])/5.
n_vert = int(1./self.vol_2d_label_size)
scale = min(self.vol_2d_label_size*4, 1.0)
text = "#scale["+str(scale)+"]{#color["+str(color)+"]{"+volume+"}}"
vertical = float(self.volume_label_index%n_vert)/n_vert
vert = x_range[1]+delta*(1.0*vertical-0.4)
self.labels[-1].DrawLatex(x_position, vert, text)
self.volume_label_index += 1
return self.labels[-1]
def _draw_volumes_2d(self, x_range, z_range):
"""Make a 2D plot of position vs volume"""
canvas = common.make_root_canvas("Volume Name vs position")
hist, graph = common.make_root_graph("Volume Name vs position",
[1e9], "z [mm]",
[1e9], "x [mm]",
xmin=z_range[0], xmax=z_range[1],
ymin=x_range[0], ymax=x_range[1])
hist.Draw()
draw_list = sorted(self.volume_data.keys(), self._sort_cmp)
n_colors = 5
for i, vol in enumerate(draw_list):
data = self.volume_data[vol]
r_list = [data["r_min"], data["r_max"], data["r_max"],
data["r_min"], data["r_min"]]
z_list = [data["z_min"], data["z_min"], data["z_max"],
data["z_max"], data["z_min"]]
hist, graph = common.make_root_graph(vol,
z_list, "z [mm]",
r_list, "r [mm]", sort = False)
color = self.color(i % n_colors, n_colors)
graph.SetLineColor(color)
graph.Draw('l')
self._draw_volume_label(data["z_min"], vol, x_range, color)
canvas.Update()
for _format in self.formats:
canvas.Print(self.vol_2d+"."+_format)
def _draw_materials_2d(self, step_data, material_dict, x_range, z_range):
"""Make a 2D plot of position vs material in r, z"""
# we plot G4_AIR first, then G4_GALACTIC?
canvas = common.make_root_canvas("Material vs position")
hist, graph = common.make_root_graph("Material vs position",
[1e9], "z [mm]",
[1e9], "r [mm]",
xmin=z_range[0], xmax=z_range[1],
ymin=x_range[0], ymax=x_range[1])
hist.Draw()
draw_list = step_data.keys()
draw_list = self.move(draw_list, "G4_AIR", 0)
draw_list = self.move(draw_list, "G4_Galactic", 1)
graph_list = []
for material in draw_list:
x_list = [step["r"] for step in step_data[material]]
z_list = [step["z"] for step in step_data[material]]
hist, graph = common.make_root_graph(material,
z_list, "z [mm]",
x_list, "r [mm]")
graph.SetMarkerColor(material_dict[material])
graph.SetMarkerStyle(7)
graph.SetLineColor(10)
graph.Draw('p')
graph_list.append(graph)
common.make_root_legend(canvas, graph_list)
canvas.Update()
for _format in self.formats:
canvas.Print(self.mat_2d+"."+_format)
@classmethod
def _get_r(cls, step):
"""Get r for step"""
return (step["x"]**2+step["y"]**2)**0.5
@classmethod
def color(cls, ith_colour, n_colours):
"""Get the colour for the ith material of n total materials"""
if n_colours == 0:
return 0
ith_colour = float(ith_colour)
n_colours = float(n_colours)
n_lim = n_colours/3.
red, green, blue = 0., 0., 0.
if ith_colour < n_lim:
red = (n_lim-ith_colour)/n_lim
green = ith_colour/n_lim
if ith_colour >= n_lim and ith_colour < 2.*n_lim:
ith_colour -= n_lim
green = (n_lim-ith_colour)/n_lim
blue = ith_colour/n_lim
if ith_colour >= 2.*n_lim:
ith_colour -= 2.*n_lim
blue = (n_lim-ith_colour)/n_lim
red = ith_colour/n_lim
color = ROOT.TColor.GetColor(red, green, blue)
return color
@classmethod
def get_material_dict(cls, step_data):
"""
Get a dictionary that maps material to colour. Air and Galactic are
always grey.
"""
n_items = len(step_data)
material_dict = {}
for i, key in enumerate(step_data.keys()):
material_dict[key] = cls.color(i, n_items)
material_dict["G4_AIR"] = ROOT.TColor.GetColor(0.1, 0.1, 0.1)
material_dict["G4_Galactic"] = ROOT.TColor.GetColor(0.3, 0.3, 0.3)
return material_dict
@classmethod
def get_range(cls, step_data):
"""
Get axis range in r, z
"""
r_data, z_data = [], []
for steps in step_data.values():
r_data += [pos["r"] for pos in steps]
z_data += [pos["z"] for pos in steps]
r_data = sorted(r_data)
z_data = sorted(z_data)
data_range = []
for data in r_data, z_data:
delta = data[-1]-data[0]
if delta < 1e-9:
data_range.append([-1., 1.])
else:
data_range.append([data[0]-delta/10., data[-1]+delta/10.])
return data_range
@classmethod
def move(cls, a_list, value, index):
"""
Move item "value" in list "a_list" to position "index".
If "value" is not in "a_list", well never mind
"""
if value in a_list:
a_list.remove(value)
a_list.insert(index, value)
return a_list
labels = []
class Tracking(object): # pylint: disable=R0902
"""Plotting class handles generation of particles and tracking"""
def __init__(self, configuration):
"""
Initialise the tracking class
- configuration should be a valid datacard set in string representation
"""
self.config = json.loads(configuration)
py_config = self.config["geometry_validation"]
self.file_name = py_config["file_name"]
self.n_steps = py_config["n_steps"]
self.x_start = py_config["x_start"]
self.x_step = py_config["x_step"]
self.y_start = py_config["y_start"]
self.y_step = py_config["y_step"]
self.z_start = py_config["z_start"]
self.will_track = py_config["will_track"]
def do_tracking(self):
"""Do the tracking"""
if not self.will_track:
return
print "Force stepping to verbose"
output_file = open(self.file_name, "w")
self._fiddle_configuration()
print "Tracking", self.n_steps, "events"
bar_length = 72
sys.stdout.write("["+" "*(bar_length+1)+"]")
sys.stdout.flush()
sys.stdout.write("\b"*(bar_length+1))
old_done = 0
for i in range(self.n_steps):
done = i/float(self.n_steps-1)*bar_length
while done-old_done > 1.:
old_done += 1
sys.stdout.write("=")
sys.stdout.flush()
x_pos = self.x_start+i*self.x_step
y_pos = self.y_start+i*self.y_step
event = self.track_particle(x_pos, y_pos)
print >> output_file, json.dumps(event)
print
def _fiddle_configuration(self):
"""Switch off physics processes and force keep steps to True"""
self.config["keep_steps"] = True
self.config["physics_processes"] = "none"
maus_cpp.globals.birth(json.dumps(self.config))
def _get_primary(self, x_pos, y_pos):
"""Build a primary with initial position (x, y, z_start) and large
momentum in the z-direction"""
return {'primary':
{
'position':{'x':x_pos, 'y':y_pos, 'z':self.z_start},
'momentum':{'x':0., 'y':0., 'z':1.},
'particle_id':22,
'energy':1e9,
'time':0.,
'random_seed':0
}
}
def track_particle(self, x_pos, y_pos):
"""Track a particle with initial position x,y"""
primary = json.dumps([self._get_primary(x_pos, y_pos)])
event = json.loads(maus_cpp.simulation.track_particles(primary))
return event
def main():
"""Set up the datacards and call tracking and plotting routines"""
configuration = Configuration.Configuration().\
getConfigJSON(command_line_args=True)
tracking = Tracking(configuration)
plotting = Plotting(configuration)
tracking.do_tracking()
plotting.do_plotting()
print "Press to finish"
raw_input()
if __name__ == "__main__":
main()