Source code for pyfeyn2.render.pyx.diagrams

"""Classes for the actual diagram containers."""

import pyx

from pyfeyn2.render.pyx import config
from pyfeyn2.render.pyx.utils import Visible


## Diagram class
[docs]class FeynDiagram: """The main PyFeyn diagram class.""" currentDiagram = None
[docs] def __init__(self, objects=None, canvas=None): """Objects for holding a set of Feynman diagram components.""" self.__objs = objects if self.__objs is None: self.__objs = [] self.highestautolayer = 0 if canvas is None: self.currentCanvas = pyx.canvas.canvas() else: self.currentCanvas = canvas FeynDiagram.currentDiagram = self
[docs] def add(self, *objs): """Add an object to the diagram.""" for obj in objs: if config.getOptions().DEBUG: print("#objs = %d" % len(self.__objs)) offset = 0 if "layeroffset" in obj.__dict__: # print "offset =", obj.layeroffset offset = obj.layeroffset self.highestautolayer += 1 if isinstance(obj, Visible): obj.setDepth(self.highestautolayer + offset) if config.getOptions().DEBUG: print( "Object %s layer = %d + %d = %d" % ( obj.__class__, self.highestautolayer, offset, self.highestautolayer + offset, ) ) self.__objs.append(obj)
[docs] def drawToCanvas(self): """Draw the components of this diagram in a well-defined order.""" if config.getOptions().DEBUG: print("Final #objs = %d" % len(self.__objs)) if config.getOptions().VDEBUG: print("Running in visual debug mode") ## Sort drawing objects by layer drawingobjs = self.__objs try: drawingobjs.sort() except Exception: pass ## Draw each object for obj in drawingobjs: if config.getOptions().DEBUG: print("Depth = ", obj.getDepth()) obj.draw(self.currentCanvas) return self.currentCanvas
[docs] def draw(self, outfile, enlargement=0): """Draw the diagram to a file, with the filetype (EPS or PDF) derived from the file extension.""" c = self.drawToCanvas() if c is not None and outfile is not None: c.writetofile(outfile, page_bbox=c.bbox().enlarged(enlargement))