ATLAS Offline Software
draw_obj.py
Go to the documentation of this file.
1 # Copyright (C) 2002-2025 CERN for the benefit of the ATLAS collaboration
2 
3 #
4 # File: draw_obj.py
5 # Created: sss, 2004.
6 # Purpose: Helpers for drawing ROOT objects.
7 #
8 
9 """Helpers for drawing ROOT objects.
10 The main function in here is draw_obj(obj [, options]).
11 This is similar to obj->Draw(options), but some extra
12 functionality is added:
13 
14  - Automatic zone handling. You can split the canvas into
15  subpads (zones) with zone(). Each call to draw_obj() will
16  plot into the next sequential pad (unless SAME was specified).
17  Selecting a pad by clicking with button 2 will change
18  the sequence so that the next objects is drawn in the
19  selected pad. draw_obj also takes an optional padnum
20  argument to explicitly specify the pad number to use.
21 
22  - If the SAME option is given, then the new object will be
23  superimposed on an existing one in the current pad.
24  In the case of 1D histograms, we change the line
25  type of the new histogram, and also rescale the vertical
26  axis, if needed to make the new histogram fit.
27 
28  - The option MERGE is like SAME, except that we step through
29  the pads as normal. The new plot will be made on top of
30  whatever happens to be on that pad. By default, the line
31  type used will be the same as one would get the first time
32  one overlays a plot with SAME; to override this, put a number
33  directly after the option, like `MERGE2'.
34 
35  - The option NORM draws the histogram normalized to 1.
36 
37  - The option LOGY draws the histogram using a logarithmic
38  y-axis, and LOGX draws it with a logarithmic x-axis.
39 
40  - The line type will recycle after four histograms have
41  been superimposed. If the LINECOLORS option is given, then
42  the new histograms will be drawn in a different color.
43 
44  - draw_obj takes optional min and max parameters to set the
45  minimum and maximum values for the y-axis of the histogram
46  being drawn.
47 
48 Besides draw_obj, a few other functions are available.
49 zone(nx,ny) was already mentioned.
50 printeps(fname) is a shortcut to print an eps file from the current canvas.
51 get_canvas() returns the canvas currently being used for drawing.
52 get_pad() returns the pad currently being used for drawing.
53 """
54 from __future__ import division
55 
56 
57 from ROOT import gROOT, TCanvas, TVirtualPad, TH1, TH2, TObject
58 
59 
60 # The number of histograms we've superimposed on the current pad with SAME.
61 _samecount = 0
62 
63 # The number of the last subpad we drew in.
64 # Will be 0 if there is only one pad.
65 _lastpad = 0
66 
67 # The number of the next subpad to draw in, starting with 1.
68 _nextpad = 1
69 
70 # Total number of subpads.
71 _npads = 1
72 
73 
74 
76  """Helper class for parsing options."""
77  def __init__ (self, options):
78  self.merge = 0
79  self.same = 0
80  self.norm = 0
81  self.logx = 0
82  self.logy = 0
83  self.fill = -1
84  self.linetype = -1
85  self.color = -1
86  self.linecolors = 0
87  self.other = ""
88 
89  options = options.replace (',', ' ')
90  for o in options.split():
91  lo = o.lower()
92  if lo in ["merge", "same", "norm", "logx", "logy", 'linecolors']:
93  setattr (self, lo, 1)
94  elif (self._optmatch (lo, "fill") or
95  self._optmatch (lo, "linetype") or
96  self._optmatch (lo, "color")):
97  pass
98  else:
99  self.other += o
100  return
101 
102 
103  def _optmatch (self, lo, pat):
104  if lo.startswith (pat+'='):
105  setattr (self, pat, int (lo[len (pat)+1:]))
106  return 1
107  return 0
108 
109 
110 def draw_obj (obj, options = "", padnum = -1, pad = None, min=None, max=None):
111  """Draw the root object OBJ in the next available pad.
112 
113 Inputs:
114  obj - The object to draw.
115  options - Drawing options.
116  These are passed through to the root Draw
117  method, except that we have special
118  handling for the SAME option, and add a new MERGE option.
119  See the header for details.
120  padnum - If this is a non-negative integer, then this specifies
121  the pad in which to draw, overriding
122  other specifications except for PAD. Note: subpad numbers
123  start with 1.
124  pad - Explicitly specify the pad to use for drawing.
125  Overrides all other specifications.
126 
127 Returns:
128  The object that we drew (may not be the same as OBJ if we
129  made a copy).
130 """
131  global _samecount
132 
133  if min is not None:
134  obj.SetMinimum (min)
135  if max is not None:
136  obj.SetMaximum (max)
137 
138  op = _options (options)
139 
140  # Should we advance to the next pad?
141  advance_p = 0
142 
143  # Should we do vertical axis rescaling?
144  rescale_p = 0
145 
146  # For SAME and MERGE, we have to pass SAME on to root.
147  if op.same or op.merge:
148  op.other += "SAME"
149 
150  if not op.same:
151  # No SAME option. Reset the count.
152  _samecount = 0
153 
154  # Advance to the next pad.
155  advance_p = 1
156  else:
157  # SAME was specified. Keep count of the number of such.
158  _samecount += 1
159  rescale_p = 1
160 
161  # Handle the MERGE option.
162  if op.merge:
163  rescale_p = 1
164  advance_p = 1
165 
166  if pad:
167  pad.cd()
168  else:
169  pad = get_pad (advance_p, padnum)
170 
171  if not op.merge and not op.same:
172  pad.SetLogx (not not op.logx)
173  pad.SetLogy (not not op.logy)
174 
175  if isinstance (obj, TH1) and not isinstance (obj, TH2):
176  h = obj
177  if op.norm:
178  h = h.Clone()
179  intg = h.Integral()
180  if intg == 0:
181  intg = 1
182  h.Scale (1. / intg)
183  if max is not None:
184  h.SetMaximum (max)
185 
186  # Special handling for 1D histograms.
187  # If SAME was specified, rescale the vertical axis, if needed.
188  if rescale_p and max is None:
189  # Find the first hist already plotted.
190  hfirst = None
191  prims = pad.GetListOfPrimitives()
192  # Avoids RecursiveRemove crash...
193  prims.ResetBit(TObject.kMustCleanup)
194  for obj in prims:
195  if isinstance (obj, TH1):
196  hfirst = obj
197  break
198 
199  # If the new hist's maximum is larger than the first one,
200  # adjust the maximum of the first.
201  if hfirst and h.GetMaximum() > hfirst.GetMaximum():
202  hfirst.SetMaximum (h.GetMaximum() * 1.1)
203 
204  if hfirst and h.GetMinimum() < hfirst.GetMinimum():
205  hfirst.SetMinimum (h.GetMinimum())
206 
207  # Draw a copy of the histogram.
208  # Adjust the line style.
209  hh = h.DrawCopy (op.other)
210  if op.linetype >= 0:
211  hh.SetLineStyle (op.linetype)
212  else:
213  hh.SetLineStyle ((_samecount%4)+1)
214  if op.color >= 0:
215  hh.SetLineColor (op.color)
216  elif op.linecolors and _samecount >= 4:
217  hh.SetLineColor (_samecount//4 + 1)
218  if op.fill >= 0:
219  hh.SetFillColor (op.fill)
220  obj = hh
221  else:
222  # Not a 1-D histogram. Just draw it.
223  obj.Draw (op.other)
224 
225  return obj
226 
227 
228 def get_pad (advance_p = 1, padnum = -1):
229  """Advance to the next pad, if requested.
230 Allow clicking on the pads (button 2) to change the next pad.
231 """
232  global _lastpad, _nextpad, _npads
233 
234  c1 = get_canvas()
235  pad = TVirtualPad.Pad()
236 
237  if advance_p:
238  if (pad and
239  pad.GetCanvasID() == c1.GetCanvasID() and
240  pad.GetNumber() != _lastpad and pad.GetNumber() != 0):
241  _nextpad = pad.GetNumber()
242  if _nextpad > _npads:
243  _nextpad = 1
244 
245  # Set up to draw on nextpad.
246  # Bump it for the next go around.
247  if _npads > 1:
248  _lastpad = _nextpad
249  _nextpad += 1
250  if _nextpad > _npads:
251  _nextpad = 1
252  else:
253  _lastpad = 0
254 
255  # Select the pad.
256  c1.cd (_lastpad)
257 
258  # Handle padnum.
259  if padnum >= 0:
260  _lastpad = padnum
261  if _npads > 0:
262  _nextpad = _lastpad + 1
263  if _nextpad > _npads:
264  _nextpad = 0
265  else:
266  _nextpad = 0
267 
268  c1.cd (_lastpad)
269 
270  # Refetch the pad.
271  return TVirtualPad.Pad()
272 
273 
274 _canvas = None
275 def get_canvas (cname = "c1"):
276  """Return the canvas named CNAME.
277 Create it if it doesn't exist.
278 """
279  global _canvas
280  _canvas = gROOT.FindObject (cname)
281  if not _canvas:
282  _canvas = TCanvas (cname, cname, 700, 600)
283  _canvas.SetLeftMargin (0.15)
284  _canvas.SetBottomMargin (0.15)
285  _canvas.SetLogx (0)
286  _canvas.SetLogy (0)
287  return _canvas
288 
289 
290 
291 def zone (nx, ny):
292  """Divide the canvas into subpads.
293 NX and NY are the number of subpads in x and y, respectively.
294 """
295  global _npads, _nextpad, _lastpad
296  c1 = get_canvas()
297  c1.Clear()
298  c1.Divide (nx, ny)
299  _npads = nx*ny
300  _nextpad = 1
301  _lastpad = 0
302  return
303 
304 
305 def printeps (s = "out.eps"):
306  """Print the current canvas as an eps file."""
307 
308  c1 = get_canvas()
309  c1.Print (s, "eps,Portrait")
310  return
python.draw_obj._options.color
color
Definition: draw_obj.py:85
python.draw_obj.printeps
def printeps(s="out.eps")
Definition: draw_obj.py:305
python.draw_obj._options.logx
logx
Definition: draw_obj.py:81
python.draw_obj._options.merge
merge
Definition: draw_obj.py:78
python.draw_obj.draw_obj
def draw_obj(obj, options="", padnum=-1, pad=None, min=None, max=None)
Definition: draw_obj.py:110
python.draw_obj._options.linetype
linetype
Definition: draw_obj.py:84
python.draw_obj._options._optmatch
def _optmatch(self, lo, pat)
Definition: draw_obj.py:103
python.draw_obj.get_canvas
def get_canvas(cname="c1")
Definition: draw_obj.py:275
python.draw_obj._options.__init__
def __init__(self, options)
Definition: draw_obj.py:77
python.draw_obj._options.other
other
Definition: draw_obj.py:87
python.draw_obj._options
Definition: draw_obj.py:75
python.draw_obj.zone
def zone(nx, ny)
Definition: draw_obj.py:291
python.draw_obj._options.norm
norm
Definition: draw_obj.py:80
python.draw_obj.get_pad
def get_pad(advance_p=1, padnum=-1)
Definition: draw_obj.py:228
pickleTool.object
object
Definition: pickleTool.py:29
python.draw_obj._options.logy
logy
Definition: draw_obj.py:82
python.draw_obj._options.fill
fill
Definition: draw_obj.py:83
python.draw_obj._options.linecolors
linecolors
Definition: draw_obj.py:86
python.draw_obj._options.same
same
Definition: draw_obj.py:79