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 
55 
56 from ROOT import gROOT, TCanvas, TVirtualPad, TH1, TH2, TObject
57 
58 
59 # The number of histograms we've superimposed on the current pad with SAME.
60 _samecount = 0
61 
62 # The number of the last subpad we drew in.
63 # Will be 0 if there is only one pad.
64 _lastpad = 0
65 
66 # The number of the next subpad to draw in, starting with 1.
67 _nextpad = 1
68 
69 # Total number of subpads.
70 _npads = 1
71 
72 
73 
75  """Helper class for parsing options."""
76  def __init__ (self, options):
77  self.merge = 0
78  self.same = 0
79  self.norm = 0
80  self.logx = 0
81  self.logy = 0
82  self.fill = -1
83  self.linetype = -1
84  self.color = -1
85  self.linecolors = 0
86  self.other = ""
87 
88  options = options.replace (',', ' ')
89  for o in options.split():
90  lo = o.lower()
91  if lo in ["merge", "same", "norm", "logx", "logy", 'linecolors']:
92  setattr (self, lo, 1)
93  elif (self._optmatch (lo, "fill") or
94  self._optmatch (lo, "linetype") or
95  self._optmatch (lo, "color")):
96  pass
97  else:
98  self.other += o
99  return
100 
101 
102  def _optmatch (self, lo, pat):
103  if lo.startswith (pat+'='):
104  setattr (self, pat, int (lo[len (pat)+1:]))
105  return 1
106  return 0
107 
108 
109 def draw_obj (obj, options = "", padnum = -1, pad = None, min=None, max=None):
110  """Draw the root object OBJ in the next available pad.
111 
112 Inputs:
113  obj - The object to draw.
114  options - Drawing options.
115  These are passed through to the root Draw
116  method, except that we have special
117  handling for the SAME option, and add a new MERGE option.
118  See the header for details.
119  padnum - If this is a non-negative integer, then this specifies
120  the pad in which to draw, overriding
121  other specifications except for PAD. Note: subpad numbers
122  start with 1.
123  pad - Explicitly specify the pad to use for drawing.
124  Overrides all other specifications.
125 
126 Returns:
127  The object that we drew (may not be the same as OBJ if we
128  made a copy).
129 """
130  global _samecount
131 
132  if min is not None:
133  obj.SetMinimum (min)
134  if max is not None:
135  obj.SetMaximum (max)
136 
137  op = _options (options)
138 
139  # Should we advance to the next pad?
140  advance_p = 0
141 
142  # Should we do vertical axis rescaling?
143  rescale_p = 0
144 
145  # For SAME and MERGE, we have to pass SAME on to root.
146  if op.same or op.merge:
147  op.other += "SAME"
148 
149  if not op.same:
150  # No SAME option. Reset the count.
151  _samecount = 0
152 
153  # Advance to the next pad.
154  advance_p = 1
155  else:
156  # SAME was specified. Keep count of the number of such.
157  _samecount += 1
158  rescale_p = 1
159 
160  # Handle the MERGE option.
161  if op.merge:
162  rescale_p = 1
163  advance_p = 1
164 
165  if pad:
166  pad.cd()
167  else:
168  pad = get_pad (advance_p, padnum)
169 
170  if not op.merge and not op.same:
171  pad.SetLogx (not not op.logx)
172  pad.SetLogy (not not op.logy)
173 
174  if isinstance (obj, TH1) and not isinstance (obj, TH2):
175  h = obj
176  if op.norm:
177  h = h.Clone()
178  intg = h.Integral()
179  if intg == 0:
180  intg = 1
181  h.Scale (1. / intg)
182  if max is not None:
183  h.SetMaximum (max)
184 
185  # Special handling for 1D histograms.
186  # If SAME was specified, rescale the vertical axis, if needed.
187  if rescale_p and max is None:
188  # Find the first hist already plotted.
189  hfirst = None
190  prims = pad.GetListOfPrimitives()
191  # Avoids RecursiveRemove crash...
192  prims.ResetBit(TObject.kMustCleanup)
193  for obj in prims:
194  if isinstance (obj, TH1):
195  hfirst = obj
196  break
197 
198  # If the new hist's maximum is larger than the first one,
199  # adjust the maximum of the first.
200  if hfirst and h.GetMaximum() > hfirst.GetMaximum():
201  hfirst.SetMaximum (h.GetMaximum() * 1.1)
202 
203  if hfirst and h.GetMinimum() < hfirst.GetMinimum():
204  hfirst.SetMinimum (h.GetMinimum())
205 
206  # Draw a copy of the histogram.
207  # Adjust the line style.
208  hh = h.DrawCopy (op.other)
209  if op.linetype >= 0:
210  hh.SetLineStyle (op.linetype)
211  else:
212  hh.SetLineStyle ((_samecount%4)+1)
213  if op.color >= 0:
214  hh.SetLineColor (op.color)
215  elif op.linecolors and _samecount >= 4:
216  hh.SetLineColor (_samecount//4 + 1)
217  if op.fill >= 0:
218  hh.SetFillColor (op.fill)
219  obj = hh
220  else:
221  # Not a 1-D histogram. Just draw it.
222  obj.Draw (op.other)
223 
224  return obj
225 
226 
227 def get_pad (advance_p = 1, padnum = -1):
228  """Advance to the next pad, if requested.
229 Allow clicking on the pads (button 2) to change the next pad.
230 """
231  global _lastpad, _nextpad, _npads
232 
233  c1 = get_canvas()
234  pad = TVirtualPad.Pad()
235 
236  if advance_p:
237  if (pad and
238  pad.GetCanvasID() == c1.GetCanvasID() and
239  pad.GetNumber() != _lastpad and pad.GetNumber() != 0):
240  _nextpad = pad.GetNumber()
241  if _nextpad > _npads:
242  _nextpad = 1
243 
244  # Set up to draw on nextpad.
245  # Bump it for the next go around.
246  if _npads > 1:
247  _lastpad = _nextpad
248  _nextpad += 1
249  if _nextpad > _npads:
250  _nextpad = 1
251  else:
252  _lastpad = 0
253 
254  # Select the pad.
255  c1.cd (_lastpad)
256 
257  # Handle padnum.
258  if padnum >= 0:
259  _lastpad = padnum
260  if _npads > 0:
261  _nextpad = _lastpad + 1
262  if _nextpad > _npads:
263  _nextpad = 0
264  else:
265  _nextpad = 0
266 
267  c1.cd (_lastpad)
268 
269  # Refetch the pad.
270  return TVirtualPad.Pad()
271 
272 
273 _canvas = None
274 def get_canvas (cname = "c1"):
275  """Return the canvas named CNAME.
276 Create it if it doesn't exist.
277 """
278  global _canvas
279  _canvas = gROOT.FindObject (cname)
280  if not _canvas:
281  _canvas = TCanvas (cname, cname, 700, 600)
282  _canvas.SetLeftMargin (0.15)
283  _canvas.SetBottomMargin (0.15)
284  _canvas.SetLogx (0)
285  _canvas.SetLogy (0)
286  return _canvas
287 
288 
289 
290 def zone (nx, ny):
291  """Divide the canvas into subpads.
292 NX and NY are the number of subpads in x and y, respectively.
293 """
294  global _npads, _nextpad, _lastpad
295  c1 = get_canvas()
296  c1.Clear()
297  c1.Divide (nx, ny)
298  _npads = nx*ny
299  _nextpad = 1
300  _lastpad = 0
301  return
302 
303 
304 def printeps (s = "out.eps"):
305  """Print the current canvas as an eps file."""
306 
307  c1 = get_canvas()
308  c1.Print (s, "eps,Portrait")
309  return
python.draw_obj._options.color
color
Definition: draw_obj.py:84
python.draw_obj.printeps
def printeps(s="out.eps")
Definition: draw_obj.py:304
python.draw_obj._options.logx
logx
Definition: draw_obj.py:80
python.draw_obj._options.merge
merge
Definition: draw_obj.py:77
python.draw_obj.draw_obj
def draw_obj(obj, options="", padnum=-1, pad=None, min=None, max=None)
Definition: draw_obj.py:109
python.draw_obj._options.linetype
linetype
Definition: draw_obj.py:83
python.draw_obj._options._optmatch
def _optmatch(self, lo, pat)
Definition: draw_obj.py:102
python.draw_obj.get_canvas
def get_canvas(cname="c1")
Definition: draw_obj.py:274
python.draw_obj._options.__init__
def __init__(self, options)
Definition: draw_obj.py:76
python.draw_obj._options.other
other
Definition: draw_obj.py:86
python.draw_obj._options
Definition: draw_obj.py:74
python.draw_obj.zone
def zone(nx, ny)
Definition: draw_obj.py:290
python.draw_obj._options.norm
norm
Definition: draw_obj.py:79
python.draw_obj.get_pad
def get_pad(advance_p=1, padnum=-1)
Definition: draw_obj.py:227
pickleTool.object
object
Definition: pickleTool.py:29
python.draw_obj._options.logy
logy
Definition: draw_obj.py:81
python.draw_obj._options.fill
fill
Definition: draw_obj.py:82
python.draw_obj._options.linecolors
linecolors
Definition: draw_obj.py:85
python.draw_obj._options.same
same
Definition: draw_obj.py:78