init root renderer

Author: APN-Pucky

pyfeyn2/render/root/pyfeyn.py

@@ -0,0 +1,240 @@
+# Adapted code from https://github.com/aminnj/pyfeyn/blob/main/pyfeyn.py
+import math
+
+
+class Label(object):
+    def __init__(
+        self,
+        text="",
+        x1=0.0,
+        y1=0.0,
+        offsetx=0,
+        offsety=0,
+        textsize=0.08,
+        textalign=22,
+        roman=False,
+    ):
+        self.x1 = x1
+        self.y1 = y1
+        self.text = text
+        self.offsetx = offsetx
+        self.offsety = offsety
+        self.textsize = textsize
+        self.textalign = textalign
+        self.roman = roman
+
+    def set_location(self, x, y):
+        self.x1 = x
+        self.y1 = y
+
+    def transform_text(self, text):
+        # ROOT needs one of these characters to put in a $ and go into mathmode
+        # otherwise we do it explicitly
+        if self.roman or any([x in text for x in "#{}^"]):
+            return text
+        text = "${0}$".format(text)
+        return text
+
+    def draw(self):
+        import ROOT
+
+        if not self.text:
+            return
+        t = ROOT.TLatex()
+        t.SetTextAlign(self.textalign)
+        t.SetTextSize(self.textsize)
+        t.DrawLatex(
+            self.x1 + self.offsetx,
+            self.y1 + self.offsety,
+            self.transform_text(self.text),
+        )
+
+
+class Marker(object):
+    def __init__(self, color=None, radius=2, linewidth=0):
+        self.x = None
+        self.y = None
+        self.color = color
+        self.radius = radius
+        self.linewidth = linewidth
+
+    def set_location(self, x, y):
+        self.x = x
+        self.y = y
+
+    def draw(self):
+        import ROOT
+
+        if self.color is None:
+            return
+        m = ROOT.TEllipse(self.x, self.y, self.radius)
+        m.SetFillColor(self.color)
+        m.SetLineWidth(self.linewidth)
+        m.Draw()
+
+
+class Vertex(object):
+    def __init__(self, x1, y1, label=Label(), marker=Marker(), autolabel=True):
+        self.x1 = x1
+        self.y1 = y1
+        self.label = label
+        self.marker = marker
+
+        self.marker.set_location(x1, y1)
+
+        if autolabel:
+            self.label.set_location(self.x1, self.y1)
+
+    def draw(self, _nodelete=[]):
+        self.label.draw()
+        self.marker.draw()
+
+
+class Propagator(object):
+    def __init__(
+        self,
+        v1,
+        v2,
+        typ="line",
+        label=Label(),
+        autolabel=True,
+        linewidth=2,
+        linecolor=None,
+        fliparrow=False,
+        noarrow=False,
+    ):
+        import ROOT
+
+        if linecolor is None:
+            linecolor = ROOT.kBlack
+        self.v1 = v1
+        self.v2 = v2
+        self.typ = typ
+        self.label = label
+        self.linewidth = linewidth
+        self.linecolor = linecolor
+        self.fliparrow = fliparrow
+        self.noarrow = noarrow
+
+        if autolabel:
+            self.label.set_location(
+                0.5 * (self.v1.x1 + self.v2.x1), 0.5 * (self.v1.y1 + self.v2.y1)
+            )
+
+    def draw(self, _nodelete=[]):
+        import ROOT as r
+
+        prop1, prop2 = None, None
+        drawopt = ""
+        if self.typ == "line":
+            prop1 = r.TLine(self.v1.x1, self.v1.y1, self.v2.x1, self.v2.y1)
+        if self.typ == "dashedline":
+            prop1 = r.TLine(self.v1.x1, self.v1.y1, self.v2.x1, self.v2.y1)
+            r.gStyle.SetLineStyleString(11, "50 30")
+            prop1.SetLineStyle(11)
+        if self.typ == "dottedline":
+            prop1 = r.TLine(self.v1.x1, self.v1.y1, self.v2.x1, self.v2.y1)
+            r.gStyle.SetLineStyleString(11, "20 50")
+            prop1.SetLineStyle(11)
+        elif self.typ == "curlyline":
+            prop1 = r.TCurlyLine(self.v1.x1, self.v1.y1, self.v2.x1, self.v2.y1)
+            prop1.SetWaveLength(prop1.GetWaveLength() * 1.6)
+            prop1.SetAmplitude(prop1.GetAmplitude() * 1.4)
+        elif self.typ == "wavyline":
+            prop1 = r.TCurlyLine(self.v1.x1, self.v1.y1, self.v2.x1, self.v2.y1)
+            prop1.SetWavy()
+            prop1.SetWaveLength(prop1.GetWaveLength() * 1.6)
+            prop1.SetAmplitude(prop1.GetAmplitude() * 1.4)
+        elif self.typ == "wavystraightline":
+            prop1 = r.TCurlyLine(self.v1.x1, self.v1.y1, self.v2.x1, self.v2.y1)
+            prop1.SetWavy()
+            prop1.SetWaveLength(prop1.GetWaveLength() * 1.6)
+            prop1.SetAmplitude(prop1.GetAmplitude() * 1.4)
+            prop2 = r.TLine(self.v1.x1, self.v1.y1, self.v2.x1, self.v2.y1)
+        elif self.typ == "curlystraightline":
+            prop1 = r.TCurlyLine(self.v1.x1, self.v1.y1, self.v2.x1, self.v2.y1)
+            prop1.SetWaveLength(prop1.GetWaveLength() * 1.6)
+            prop1.SetAmplitude(prop1.GetAmplitude() * 1.4)
+            prop2 = r.TLine(self.v1.x1, self.v1.y1, self.v2.x1, self.v2.y1)
+        elif self.typ.startswith("wavyarc"):
+            # wavyarc(180,0) -> phimin = 180, phimax = 0
+            phimin, phimax = list(
+                map(float, self.typ.split("(", 1)[1].split(")", 1)[0].split(","))
+            )
+            xc = 0.5 * (self.v1.x1 + self.v2.x1)
+            yc = 0.5 * (self.v1.y1 + self.v2.y1)
+            radius = 0.5 * (self.v2.x1 - self.v1.x1)
+            prop1 = r.TCurlyArc(xc, yc, radius, phimin, phimax)
+            prop1.SetWavy()
+            prop1.SetWaveLength(prop1.GetWaveLength() * 1.6)
+            prop1.SetAmplitude(prop1.GetAmplitude() * 1.4)
+            prop1.SetFillColorAlpha(0, 0.0)
+            drawopt = "only"
+        elif self.typ.startswith("arc"):
+            # arc(180,0) -> phimin = 180, phimax = 0
+            phimin, phimax = list(
+                map(float, self.typ.split("(", 1)[1].split(")", 1)[0].split(","))
+            )
+            xc = 0.5 * (self.v1.x1 + self.v2.x1)
+            yc = 0.5 * (self.v1.y1 + self.v2.y1)
+            radius = 0.5 * (self.v2.x1 - self.v1.x1)
+            prop1 = r.TArc(xc, yc, radius, phimin, phimax)
+            prop1.SetFillColorAlpha(0, 0.0)
+            drawopt = "only"
+
+        if prop1:
+            prop1.SetLineColor(self.linecolor)
+            prop1.SetLineWidth(self.linewidth)
+            prop1.Draw(drawopt)
+            _nodelete.append(prop1)
+        if prop2:
+            prop2.SetLineColor(self.linecolor)
+            prop2.SetLineWidth(self.linewidth)
+            prop2.Draw(drawopt)
+            _nodelete.append(prop2)
+
+        if not self.noarrow:
+            if self.typ in ["line", "dashedline"]:
+                c1 = self.v1.x1, self.v1.y1
+                c2 = self.v2.x1, self.v2.y1
+                if self.fliparrow:
+                    c1, c2 = c2, c1
+                mult = 0.54
+                awidth = 0.025
+                a1 = r.TArrow(
+                    c1[0],
+                    c1[1],
+                    (1.0 - mult) * c1[0] + mult * c2[0],
+                    (1.0 - mult) * c1[1] + mult * c2[1],
+                    awidth,
+                    "|>",
+                )
+                a1.SetLineWidth(0)
+                a1.SetFillColor(self.linecolor)
+                a1.SetAngle(40)
+                a1.Draw()
+                _nodelete.append(a1)
+            elif self.typ.startswith("arc"):
+                phimin, phimax = list(
+                    map(float, self.typ.split("(", 1)[1].split(")", 1)[0].split(","))
+                )
+                phi = (0.5 * (phimax - phimin) % 360) * math.pi / 180
+                radius = 0.5 * (self.v2.x1 - self.v1.x1)
+                xc = 0.5 * (self.v1.x1 + self.v2.x1) + radius * math.cos(phi)
+                yc = 0.5 * (self.v1.y1 + self.v2.y1) + radius * math.sin(phi)
+                dx, dy = -radius * 0.2 * math.sin(phi), -radius * 0.2 * math.cos(phi)
+                if self.fliparrow:
+                    dx, dy = -dx, -dy
+                awidth = 0.025
+                a1 = r.TArrow(
+                    xc - dx / 2, yc - dy / 2, xc + dx / 2, yc + dy / 2, awidth, "|>"
+                )
+                a1.SetLineWidth(0)
+                a1.SetFillColor(self.linecolor)
+                a1.SetAngle(40)
+                a1.Draw()
+                _nodelete.append(a1)
+
+        self.v1.draw()
+        self.v2.draw()
+        self.label.draw()

pyfeyn2/render/root/rootpdfrender.py

@@ -0,0 +1,104 @@
+import tempfile
+from typing import List
+
+from pylatex import Command
+from pylatex.utils import NoEscape
+
+from pyfeyn2.render.latex.latex import LatexRender
+from pyfeyn2.render.render import Render
+from pyfeyn2.render.root.pyfeyn import Propagator, Vertex
+from pyfeyn2.render.root.rootrender import ROOTRender
+
+
+def root_to_latex(root_canvas) -> str:
+    # create a tmp tex file
+    with tempfile.NamedTemporaryFile(
+        suffix=".tex", delete=True, mode="w+"
+    ) as temp_file:
+        root_canvas.SaveAs(temp_file.name)
+        # read the file
+        tex_src = temp_file.read()
+        return tex_src
+    raise RuntimeError("Failed to create temporary file")
+
+
+class ROOTPDFRender(LatexRender, ROOTRender):
+    def __init__(
+        self,
+        fd=None,
+        documentclass="standalone",
+        document_options=None,
+        *args,
+        **kwargs,
+    ):
+        if document_options is None:
+            document_options = ["preview", "crop", "tikz"]
+        super().__init__(
+            *args,
+            fd=fd,
+            documentclass=documentclass,
+            document_options=document_options,
+            **kwargs,
+        )
+        self.preamble.append(Command("RequirePackage", "luatex85"))
+        self.preamble.append(
+            Command("usepackage", NoEscape("tikz-feynman"), "compat=1.1.0")
+        )
+
+    def render(
+        self,
+        file=None,
+        show=True,
+        resolution=100,
+        width=None,
+        height=None,
+        clean_up=True,
+    ):
+        # First need to convert feynman to ROOT through ROOTRender
+        ROOTRender.render(
+            self,
+            file=None,
+            show=False,
+            resolution=resolution,
+            width=width,
+            height=height,
+        )
+        # Then convert ROOT to LaTeX
+        self.set_src_diag(NoEscape(root_to_latex(self.get_src_root())))
+        LatexRender.render(self, file, show, resolution, width, height, clean_up)
+
+    @classmethod
+    def valid_styles(cls) -> bool:
+        return super().valid_styles() + [
+            "line",
+            "symbol",
+            "symbol-size",
+            "color",
+            "opacity",
+            "bend-direction",
+            "bend-in",
+            "bend-out",
+            "bend-loop",
+            "bend-min-distance",
+            "momentum-arrow",
+            "momentum-arrow-sense",
+            "double-distance",
+            "label-side",
+        ]
+
+    @classmethod
+    def valid_attributes(cls) -> List[str]:
+        return super().valid_attributes() + [
+            "x",
+            "y",
+            "label",
+            "style",
+        ]
+
+    @classmethod
+    def valid_types(cls) -> List[str]:
+        return super().valid_types() + []
+
+    @classmethod
+    def valid_shapes(cls) -> List[str]:
+        return super().valid_types() + []