Package madgraph :: Package iolibs :: Module drawing_eps
[hide private]
[frames] | no frames]

Source Code for Module madgraph.iolibs.drawing_eps

  1  ################################################################################ 
  2  # 
  3  # Copyright (c) 2009 The MadGraph5_aMC@NLO Development team and Contributors 
  4  # 
  5  # This file is a part of the MadGraph5_aMC@NLO project, an application which  
  6  # automatically generates Feynman diagrams and matrix elements for arbitrary 
  7  # high-energy processes in the Standard Model and beyond. 
  8  # 
  9  # It is subject to the MadGraph5_aMC@NLO license which should accompany this  
 10  # distribution. 
 11  # 
 12  # For more information, visit madgraph.phys.ucl.ac.be and amcatnlo.web.cern.ch 
 13  # 
 14  ################################################################################ 
 15   
 16  """This files contains class for creating files or object representing a  
 17  diagram or a set of diagrams. 
 18   
 19  class structure: 
 20    
 21  DrawDiagram:  
 22      In principle ALL routines representing a diagram in Any format SHOULD derive 
 23      from this class. This is a (nearly empty) frameworks to draw a diagram  
 24      in any type format.   
 25   
 26      This frameworks defines in particular  
 27          - function to convert the input diagram in the correct object.  
 28              [convert_diagram] 
 29          - main loop to draw a diagram in a line-by-line method 
 30              [draw - draw_diagram] 
 31           
 32  DrawDiagramEPS: 
 33      This contains all the routines to represent one diagram in Encapsuled  
 34      PostScript (EPS) 
 35       
 36  DrawDiagramsEPS: 
 37      This contains all the routines to represent a set of diagrams in Encapsuled  
 38      PostScript (EPS).""" 
 39   
 40  from __future__ import division 
 41   
 42  import os 
 43  import math 
 44  import madgraph.core.drawing as draw 
 45  import madgraph.core.base_objects as base_objects 
 46  import madgraph.loop.loop_base_objects as loop_objects 
 47   
 48  _file_path = os.path.split(os.path.dirname(os.path.realpath(__file__)))[0] + '/' 
 49   
 50  #=============================================================================== 
 51  # DrawDiagramEps 
 52  #=============================================================================== 
53 -class EpsDiagramDrawer(draw.DiagramDrawer):
54 """Class to write a EPS file containing the asked diagram 55 This class follows the DrawDiagram Frameworks. 56 57 The main routine to draw a diagram is 'draw' which call 58 1) initialize: setup things for the diagram (usually open a file) 59 2) convert_diagram : Update the diagram in the correct format if needed 60 3) draw_diagram : Perform diagram dependent operation 61 4) conclude : finish the operation. 62 """ 63 64 #some page information 65 width = 450 66 height = 450 67 npage = 1 68 69 # Define where to put the diagram in the page. This is the coordinate of 70 #the two opposites point of the drawing area 71 x_min = 150 72 y_min = 450 73 x_max = 450 74 y_max = 750 75 76 blob_size = 1.5 77
78 - def initialize(self):
79 """Operation done before starting to create diagram specific EPS content 80 First open the file in write mode then write in it the header and the 81 library of particle type.""" 82 83 # Open file 84 super(EpsDiagramDrawer, self).initialize() 85 86 # File Header 87 text = "%!PS-Adobe-2.0\n" 88 text += "%%" + "BoundingBox: -20 -20 %s %s \n" % \ 89 (self.width, self.height) 90 text += "%%DocumentFonts: Helvetica\n" 91 text += "%%" + "Pages: %s \n" % self.npage 92 self.file.writelines(text) 93 94 # Import the definition of the different way to represent a line 95 self.file.writelines(open(os.path.join(_file_path, \ 96 'iolibs/template_files/drawing_eps_header.inc')).read())
97 98
99 - def conclude(self):
100 """Operation to perform when all code related to a specific diagram are 101 finish. Operation : 102 - Add the 'end of page' code 103 - write unwritten text and close the file. [DrawDiagram.conclude]""" 104 105 # Add an 'end of page statement' 106 self.text = 'showpage\n' 107 self.text += '%%trailer\n' 108 109 #write the diagram file 110 super(EpsDiagramDrawer, self).conclude()
111 112
113 - def rescale(self, x, y):
114 """All coordinates belongs to [0,1]. So that in order to have a visible 115 graph we need to re-scale the graph. This method distort the square in 116 a oblong. Deformation are linear.""" 117 118 # Use the information for the graph position. 'self.x_???,self.y_???? 119 #are the coordinate of the two opposites point of the drawing area. 120 x = self.x_min + (self.x_max - self.x_min) * x 121 y = self.y_min + (self.y_max - self.y_min) * y 122 123 return x, y
124 125
126 - def line_format(self, x1, y1, x2, y2, name):
127 """Specify the text format of a specific Particles. 128 EPS format for Particle is either [X Y X Y NAME] or [X Y X Y NUM NAME]. 129 In this routine we will consider only the first format. The second can 130 be matched by redefining name in [NUM NAME].""" 131 132 # Compute real position for starting/ending point 133 x1, y1 = self.rescale(x1, y1) 134 x2, y2 = self.rescale(x2, y2) 135 136 #return the line in correct format 137 return " %s %s %s %s %s \n" % (x1, y1, x2, y2, name)
138
139 - def draw_vertex(self, vertex, bypass = ['QED','QCD'] ):
140 """Add blob in case on non QED-QCD information""" 141 142 interaction = self.model.get_interaction(vertex.id) 143 if interaction: 144 order = interaction.get('orders') 145 order = [key for key in order.keys() if order[key] and \ 146 key not in bypass] 147 148 if order: 149 x1, y1 = self.rescale(vertex.pos_x, vertex.pos_y) 150 self.text += " %s %s %s 1.0 Fblob \n" % (x1, y1, self.blob_size)
151 152 153
154 - def draw_straight(self, line):
155 """ADD the EPS code for this fermion line.""" 156 157 #add the code in the correct format 158 self.text += self.line_format(line.begin.pos_x, line.begin.pos_y, 159 line.end.pos_x, line.end.pos_y, 'Ffermion')
160
161 - def draw_curved_dashed(self, line, cercle):
162 """ADD the EPS code for this fermion line.""" 163 164 if not cercle: 165 curvature = 0.4 166 else: 167 curvature = 1 168 169 if (line.begin.pos_x, line.begin.pos_y) == self.curved_part_start: 170 curvature *= -1 171 172 #add the code in the correct format 173 x1, y1 = self.rescale(line.begin.pos_x, line.begin.pos_y) 174 self.text += ' %s %s moveto \n' % (x1, y1) 175 self.text += self.line_format(line.begin.pos_x, line.begin.pos_y, 176 line.end.pos_x, line.end.pos_y, '%s Fhiggsl' %\ 177 curvature)
178
179 - def draw_curved_straight(self, line, cercle):
180 """ADD the EPS code for this fermion line.""" 181 182 if not cercle: 183 curvature = 0.4 184 else: 185 curvature = 1 186 187 if (line.begin.pos_x, line.begin.pos_y) == self.curved_part_start: 188 curvature *= -1 189 190 #add the code in the correct format 191 self.text += self.line_format(line.begin.pos_x, line.begin.pos_y, 192 line.end.pos_x, line.end.pos_y, '%s Ffermionl' %\ 193 curvature)
194
195 - def draw_circled_straight(self, line, cercle):
196 """ADD the EPS code for this fermion line.""" 197 198 if not cercle: 199 curvature = 4 200 else: 201 curvature = 5 202 203 #add the code in the correct format 204 self.text += self.line_format(line.begin.pos_x, line.begin.pos_y, 205 line.end.pos_x+0.01, line.end.pos_y+0.01, '%s Ffermionl' %\ 206 curvature)
207 - def draw_dashed(self, line):
208 """ADD the EPS code for this Higgs line.""" 209 210 #add the code in the correct format 211 self.text += self.line_format(line.begin.pos_x, line.begin.pos_y, 212 line.end.pos_x, line.end.pos_y, 'Fhiggs')
213 214
215 - def draw_circled_dashed(self, line,cercle):
216 """ADD the EPS code for this Higgs line.""" 217 if not cercle: 218 curvature = 4 219 else: 220 curvature = 5 221 222 #add the code in the correct format 223 self.text += self.line_format(line.begin.pos_x, line.begin.pos_y, 224 line.end.pos_x+0.01, line.end.pos_y+0.01, '%s Fhiggsl'% curvature)
225
226 - def draw_dotted(self,line):
227 """ADD the EPS code for the ghost line.""" 228 229 #add the code in the correct format 230 self.text += self.line_format(line.begin.pos_x, line.begin.pos_y,\ 231 line.end.pos_x, line.end.pos_y, 'Fghost')
232
233 - def draw_curved_dotted(self, line, cercle):
234 """ADD the EPS code for the ghost line.""" 235 if not cercle: 236 curvature = 0.4 237 else: 238 curvature = 1 239 240 if (line.begin.pos_x, line.begin.pos_y) == self.curved_part_start: 241 curvature *= -1 242 #add the code in the correct format 243 self.text += self.line_format(line.begin.pos_x, line.begin.pos_y,\ 244 line.end.pos_x, line.end.pos_y, '%s Fghostl'% curvature)
245
246 - def draw_circled_dotted(self, line, cercle):
247 """ADD the EPS code for the ghost line.""" 248 if not cercle: 249 curvature = 4 250 else: 251 curvature = 5 252 253 #add the code in the correct format 254 self.text += self.line_format(line.begin.pos_x, line.begin.pos_y,\ 255 line.end.pos_x+0.01, line.end.pos_y+0.01, '%s Fghostl'% curvature)
256
257 - def draw_wavy(self, line, opt=0, type=''):
258 """ADD the EPS code for this photon line.""" 259 260 #add the code in the correct format 261 self.text += self.line_format(line.begin.pos_x, line.begin.pos_y, 262 line.end.pos_x, line.end.pos_y, '%d Fphoton%s' % (opt,type))
263
264 - def draw_curved_wavy(self, line, cercle, opt=0, type=''):
265 """ADD the EPS code for this photon line.""" 266 if not cercle: 267 curvature = 0.4 268 else: 269 curvature = 1 270 if (line.begin.pos_x, line.begin.pos_y) == self.curved_part_start: 271 curvature *= -1 272 #add the code in the correct format 273 self.text += self.line_format(line.begin.pos_x, line.begin.pos_y, 274 line.end.pos_x, line.end.pos_y, '%d %s Fphotonl%s' % (opt,curvature,type))
275
276 - def draw_circled_wavy(self, line, cercle, opt=0, type=''):
277 """ADD the EPS code for this photon line.""" 278 if not cercle: 279 curvature = 4 280 else: 281 curvature = 5 282 283 #add the code in the correct format 284 self.text += self.line_format(line.begin.pos_x, line.begin.pos_y, 285 line.end.pos_x+0.01, line.end.pos_y+0.01, '%d %s Fphotonl%s' % (opt,curvature,type))
286
287 - def draw_curly(self, line, type=''):
288 """ADD the EPS code for this gluon line.""" 289 290 # Due to the asymmetry in the way to draw the gluon (everything is draw 291 #upper or below the line joining the points). We have to put conditions 292 #in order to have nice diagram. 293 if (line.begin.pos_x < line.end.pos_x) or \ 294 (line.begin.pos_x == line.end.pos_x and \ 295 line.begin.pos_y > line.end.pos_y): 296 self.text += self.line_format(line.begin.pos_x, 297 line.begin.pos_y, line.end.pos_x, 298 line.end.pos_y, '0 Fgluon%s' % type) 299 else: 300 self.text += self.line_format(line.end.pos_x, 301 line.end.pos_y, line.begin.pos_x, 302 line.begin.pos_y, '0 Fgluon%s' % type)
303
304 - def draw_curved_curly(self, line, cercle, type=''):
305 """ADD the EPS code for this gluon line.""" 306 307 dist = math.sqrt((line.begin.pos_x-line.end.pos_x)**2 + \ 308 (line.begin.pos_y-line.end.pos_y)**2) 309 if not cercle or dist > 0.3: 310 curvature = 0.4 311 else: 312 curvature = 1 313 314 315 # Due to the asymmetry in the way to draw the gluon (everything is draw 316 #upper or below the line joining the points). We have to put conditions 317 #in order to have nice diagram. 318 319 if (line.begin.pos_x, line.begin.pos_y) == self.curved_part_start: 320 curvature *= -1 321 322 self.text += self.line_format(line.end.pos_x, 323 line.end.pos_y, line.begin.pos_x, 324 line.begin.pos_y, '0 %s Fgluonl%s' % (-1*curvature, type))
325 326 327
328 - def draw_scurly(self, line):
329 """ADD the EPS code for this gluino line.""" 330 self.draw_curly(line, type='r' 331 ) 332 self.draw_straight(line)
333
334 - def draw_swavy(self, line):
335 """ADD the EPS code for this neutralino line.""" 336 self.draw_wavy(line, type='r') 337 self.draw_straight(line)
338
339 - def draw_double(self, line, type='r'):
340 """ADD the EPS code for this neutralino line.""" 341 342 343 length = math.sqrt((line.end.pos_y - line.begin.pos_y)**2 + (line.end.pos_x - line.begin.pos_x) **2) 344 c1 = (line.end.pos_x - line.begin.pos_x)/length 345 c2 = (line.end.pos_y - line.begin.pos_y)/length 346 347 gap = 0.013 348 start2_x = line.begin.pos_x + gap * c1 349 start2_y = line.begin.pos_y + gap * c2 350 stop1_x = line.end.pos_x - gap * c1 351 stop1_y = line.end.pos_y - gap * c2 352 353 354 self.text += self.line_format(line.begin.pos_x, line.begin.pos_y, 355 stop1_x, stop1_y, '0 Fphoton%s' % (type)) 356 #add the code in the correct format 357 self.text += self.line_format(start2_x, start2_y, 358 line.end.pos_x, line.end.pos_y, '0 Fphoton%s' % (type))
359 360
361 - def put_diagram_number(self, number=0):
362 """ADD the comment 'diagram [number]' just below the diagram.""" 363 364 # Position of the text in [0,1] square 365 x = 0.2 366 y = -0.17 367 # Compute the EPS coordinate 368 x, y = self.rescale(x, y) 369 #write the text 370 self.text += ' %s %s moveto \n' % (x, y) 371 372 if hasattr(self, 'diagram_type'): 373 self.text += '(%s diagram %s ) show\n' % (self.diagram_type, number + 1) # +1 python 374 #starts to count at 375 #zero. 376 else: 377 self.text += '(diagram %s ) show\n' % (number + 1) # +1 python 378 #starts to count at 379 #zero. 380 381 mystr = " (%s)" % ", ".join(["%s=%d" % (key, self.diagram.diagram['orders'][key]) \ 382 for key in sorted(self.diagram.diagram['orders'].keys()) \ 383 if key != 'WEIGHTED']) 384 385 x = 0.6 386 y = -0.17 387 x, y = self.rescale(x, y) 388 #write the text 389 self.text += ' %s %s moveto \n' % (x, y) 390 self.text += '%s show\n' % (mystr)
391 392 393
394 - def associate_number(self, line, number):
395 """Write in the EPS figure the MadGraph5_aMC@NLO number associate to the line. 396 Note that this routine is called only for external particle.""" 397 398 # find the external vertex associate to the line 399 if line.begin.is_external(): 400 vertex = line.begin 401 else: 402 vertex = line.end 403 404 # find the position of this vertex 405 x = vertex.pos_x 406 y = vertex.pos_y 407 408 # Move slightly the position to avoid overlapping 409 if x == 0: 410 x = -0.04 411 else: 412 x += 0.04 413 y = line._has_ordinate(x) 414 415 # Re-scale x,y in order to pass in EPS coordinate 416 x, y = self.rescale(x, y) 417 # Write the EPS text associate 418 self.text += ' %s %s moveto \n' % (x, y) 419 self.text += '(%s) show\n' % (number)
420
421 - def associate_name(self, line, name, loop=False, reverse=False):
422 """ADD the EPS code associate to the name of the particle. Place it near 423 to the center of the line. 424 """ 425 426 # Put alias for vertex positions 427 x1, y1 = line.begin.pos_x, line.begin.pos_y 428 x2, y2 = line.end.pos_x, line.end.pos_y 429 430 d = line.get_length() 431 432 # compute gap from middle point 433 if abs(x1 - x2) < 1e-3: 434 dx = 0.015 435 dy = -0.01 436 elif abs(y1 - y2) < 1e-3: 437 dx = -0.01 438 dy = 0.025 439 elif ((x1 < x2) == (y1 < y2)): 440 dx = -0.03 * len(name) 441 dy = 0.02 * len(name) #d * 0.12 442 else: 443 dx = 0.01 #0.05 444 dy = 0.02 #d * 0.12 445 if loop: 446 dx, dy = 1.5* dx, dy 447 if x1 == x2: 448 if y1 < y2: 449 dx, dy = -dx, -dy 450 elif y1 == y2: 451 if x1 >x2: 452 dx, dy = -dx, -dy 453 elif x1 < x2: 454 dx, dy = -dx, -dy 455 if reverse: 456 dx, dy = -dx, -dy 457 458 459 460 # Assign position 461 x_pos = (x1 + x2) / 2 + dx 462 y_pos = (y1 + y2) / 2 + dy 463 464 # Pass in EPS coordinate 465 x_pos, y_pos = self.rescale(x_pos, y_pos) 466 #write EPS code 467 self.text += ' %s %s moveto \n' % (x_pos, y_pos) 468 self.text += '(' + name + ') show\n'
469 470 471 #=============================================================================== 472 # DrawDiagramsEps 473 #===============================================================================
474 -class MultiEpsDiagramDrawer(EpsDiagramDrawer):
475 """Class to write a EPS file containing the asked set of diagram 476 This class follows the DrawDiagram Frameworks. 477 478 The main routine to draw a diagram is 'draw' which call 479 1) initialize: setup things for the diagram (usually open a file) 480 2) convert_diagram : Update the diagram in the correct format if needed 481 3) draw_diagram : Perform diagram dependent operation 482 4) conclude : finish the operation. 483 """ 484 485 # Define where to put the diagrams in the page. This is the coordinate of 486 #the lower left corner of the drawing area of the first graph. and the 487 #dimension associate to this drawing area. 488 489 x_min = 75 490 x_size = 200 491 y_min = 560 492 y_size = 150 493 # Define distances between two drawing area 494 x_gap = 75 495 y_gap = 70 496 497 #define font 498 font=9 499 500 #Defines the number of line-column in a EPS page 501 nb_line = 3 502 nb_col = 2 503 504 blob_size = 1.5 505 506 lower_scale = 5 507 second_scale ={'x_min': 40, 'x_size':150,'y_min':620,'y_size':100, 508 'x_gap':42,'y_gap':30,'font':6,'nb_line':5,'nb_col':3, 509 'blob_size':0.9} 510
511 - def __init__(self, diagramlist=None, filename='diagram.eps', \ 512 model=None, amplitude=None, legend='',diagram_type=''):
513 """Define basic variable and store some global information 514 all argument are optional 515 diagramlist : are the list of object to draw. item should inherit 516 from either base_objects.Diagram or drawing_lib.FeynmanDiagram 517 filename: filename of the file to write 518 model: model associate to the diagram. In principle use only if diagram 519 inherit from base_objects.Diagram 520 amplitude: amplitude associate to the diagram. NOT USE for the moment. 521 In future you could pass the amplitude associate to the object in 522 order to adjust fermion flow in case of Majorana fermion.""" 523 524 #use standard initialization but without any diagram 525 super(MultiEpsDiagramDrawer, self).__init__(None, filename , model, \ 526 amplitude) 527 self.legend = legend 528 #additional information 529 self.block_nb = 0 # keep track of the number of diagram already written 530 self.curr_page = 0 # keep track of the page position 531 self.block_in_page = 0 #ckeep track of the block in a page 532 #compute the number of pages 533 self.npage = 1 534 self.diagram_type = diagram_type 535 536 diagramlist = [d for d in diagramlist if not (isinstance(d, loop_objects.LoopUVCTDiagram) or \ 537 (isinstance(d, loop_objects.LoopDiagram) and d.get('type') < 0))] 538 diagramlist = base_objects.DiagramList(diagramlist) 539 540 limit = self.lower_scale * self.nb_col * self.nb_line 541 if len(diagramlist) < limit: 542 self.npage += (len(diagramlist)-1) // (self.nb_col * self.nb_line) 543 else: 544 add = (len(diagramlist) - limit -1) // \ 545 (self.second_scale['nb_col'] * self.second_scale['nb_line']) 546 self.npage += self.lower_scale + add 547 548 if diagramlist: 549 # diagramlist Argument should be a DiagramList object 550 assert(isinstance(diagramlist, base_objects.DiagramList)) 551 self.diagramlist = diagramlist 552 else: 553 self.diagramlist = None
554
555 - def rescale(self, x, y):
556 """All coordinates belongs to [0,1]. So that in order to have a visible 557 graph we need to re-scale the graph. This method distort the square in 558 a oblong. Deformation are linear.""" 559 560 # Compute the current line and column 561 block_pos = self.block_in_page 562 line_pos = block_pos // self.nb_col 563 col_pos = block_pos % self.nb_col 564 565 # Compute the coordinate of the drawing area associate to this line 566 #and column. 567 x_min = self.x_min + (self.x_size + self.x_gap) * col_pos 568 x_max = self.x_min + self.x_gap * (col_pos) + self.x_size * \ 569 (col_pos + 1) 570 y_min = self.y_min - (self.y_size + self.y_gap) * line_pos 571 y_max = self.y_min - self.y_gap * (line_pos) - self.y_size * \ 572 (line_pos - 1) 573 574 # Re-scale the coordinate in that box 575 x = x_min + (x_max - x_min) * x 576 y = y_min + (y_max - y_min) * y 577 578 return x, y
579
580 - def draw_diagram(self, diagram):
581 """Creates the representation in EPS format associate to a specific 582 diagram.""" 583 584 # Standard method 585 super(MultiEpsDiagramDrawer, self).draw_diagram(diagram, self.block_nb) 586 # But keep track how many diagrams are already drawn 587 588 self.block_nb += 1 589 self.block_in_page +=1
590 591
592 - def draw(self, diagramlist='', opt=None):
593 """Creates the representation in EPS format associate to a specific 594 diagram. 'opt' keeps track of possible option of drawing. Those option 595 are used if we need to convert diagram to Drawing Object. 596 opt is an DrawOption object containing all the possible option on how 597 draw a diagram.""" 598 599 if diagramlist == '': 600 diagramlist = self.diagramlist 601 602 # Initialize some variable before starting to draw the diagram 603 # This creates the header-library of the output file 604 self.initialize() 605 self.text += '/Helvetica findfont %s scalefont setfont\n' % self.font 606 self.text += ' 50 770 moveto\n' 607 self.text += ' (%s) show\n' % self.legend 608 self.text += ' 525 770 moveto\n' 609 self.text += ' (page %s/%s) show\n' % (self.curr_page + 1, self.npage) 610 self.text += ' 260 50 moveto\n' 611 self.text += ' (Diagrams made by MadGraph5_aMC@NLO) show\n' 612 # Loop on all diagram 613 for i,diagram in enumerate(diagramlist): 614 # Check if they need to be convert in correct format 615 diagram = self.convert_diagram(diagram, self.model, self.amplitude, opt) 616 if diagram==None: 617 continue 618 # Write the code associate to this diagram 619 self.draw_diagram(diagram) 620 621 # Check if the current page is full or not 622 if self.block_in_page % (self.nb_col * self.nb_line) == 0: 623 #if full initialize a new page 624 self.pass_to_next_page() 625 626 #finish operation 627 self.conclude()
628
629 - def pass_to_next_page(self):
630 """Insert text in order to pass to next EPS page.""" 631 632 self.curr_page += 1 633 self.block_in_page = 0 634 if self.curr_page == self.lower_scale: 635 for key, value in self.second_scale.items(): 636 setattr(self, key, value) 637 638 639 self.text += 'showpage\n' 640 self.text += '%%' + 'Page: %s %s \n' % (self.curr_page+1, self.curr_page+1) 641 self.text += '%%PageBoundingBox:-20 -20 600 800\n' 642 self.text += '%%PageFonts: Helvetica\n' 643 self.text += '/Helvetica findfont %s scalefont setfont\n' % self.font 644 self.text += ' 50 770 moveto\n' 645 self.text += ' (%s) show\n' % self.legend 646 self.text += ' 525 770 moveto\n' 647 self.text += ' (page %s/%s) show\n' % (self.curr_page + 1, self.npage) 648 self.text += ' 260 40 moveto\n' 649 self.text += ' (Diagrams made by MadGraph5) show\n'
650