Source code for pyfeyn2.feynmandiagram

from dataclasses import dataclass, field
from decimal import Decimal
from typing import List, Optional

from pyfeyn2.particles import get_name

# from pyfeyn2.propagator import Propagator
# from pyfeyn2.vertex import Vertex


@dataclass
[docs]class PDG:
[docs]    pdgid: Optional[int] = field(
        default=21, metadata={"xml_attribute": True, "type": "Attribute"}
    )
[docs]    type: Optional[str] = field(
        default="", metadata={"xml_attribute": True, "type": "Attribute"}
    )
[docs]    latexname: Optional[str] = field(
        default="", metadata={"xml_attribute": True, "type": "Attribute"}
    )

[docs]    def _sync_latexname(self):
        """Sync the latexname with the pdgid"""
        if self.pdgid is not None:
            self.latexname = get_name(self.pdgid)

[docs]    def __post_init__(self):
        self._sync_latexname()

[docs]    def set_pdgid(self, pdgid):
        self.pdgid = pdgid
        self._sync_latexname()
        return self

[docs]    def set_type(self, type):
        self.type = type
        return self


[docs]id = 0


@dataclass
[docs]class Identifiable:
[docs]    id: Optional[str] = field(
        default="", metadata={"xml_attribute": True, "type": "Attribute"}
    )

[docs]    def __post_init__(self):
        global id
        if self.id == "":
            # use some global counter to generate unique id
            self.id = self.__class__.__name__ + str(id)
            id = id + 1


@dataclass
[docs]class Labeled:
[docs]    label: Optional[str] = field(
        default="", metadata={"xml_attribute": True, "type": "Attribute"}
    )

[docs]    def set_label(self, label):
        self.label = label
        return self


@dataclass
[docs]class Point:
[docs]    x: Optional[Decimal] = field(
        default=None, metadata={"xml_attribute": True, "type": "Attribute"}
    )
[docs]    y: Optional[Decimal] = field(
        default=None, metadata={"xml_attribute": True, "type": "Attribute"}
    )
[docs]    z: Optional[Decimal] = field(
        default=None, metadata={"xml_attribute": True, "type": "Attribute"}
    )

[docs]    def set_xy(self, x, y):
        self.x = float(x)
        self.y = float(y)
        return self

[docs]    def set_xyz(self, x, y):
        self.x = float(x)
        self.y = float(y)
        self.z = float(z)
        return self


@dataclass
[docs]class Styled:
[docs]    style: Optional[str] = field(
        default=None, metadata={"xml_attribute": True, "type": "Attribute"}
    )


@dataclass
[docs]class Bending:
[docs]    bend: Optional[str] = field(
        default=None, metadata={"xml_attribute": True, "type": "Attribute"}
    )


@dataclass
[docs]class Targeting:
[docs]    target: Optional[str] = field(
        default="", metadata={"xml_attribute": True, "type": "Attribute"}
    )

[docs]    def set_target(self, target):
        self.target = target.id
        return self


@dataclass
[docs]class Sourcing:
[docs]    source: Optional[str] = field(
        default="", metadata={"xml_attribute": True, "type": "Attribute"}
    )

[docs]    def set_source(self, source):
        self.source = source.id
        return self


@dataclass
[docs]class Line(Targeting, Sourcing):
[docs]    def connect(self, source, target):
        self.set_source(source)
        self.set_target(target)
        return self


@dataclass
[docs]class Vertex(Labeled, Styled, Point, Identifiable):
    pass


@dataclass
[docs]class Leg(Labeled, Styled, PDG, Bending, Point, Targeting, Identifiable):
[docs]    sense: str = field(
        default="", metadata={"xml_attribute": True, "type": "Attribute"}
    )

[docs]    def set_incoming(self):
        self.sense = "incoming"
        return self

[docs]    def set_outgoing(self):
        self.sense = "outgoing"
        return self


@dataclass
[docs]class Propagator(Labeled, Styled, PDG, Bending, Line, Identifiable):
    pass


@dataclass
[docs]class FeynmanDiagram:
[docs]    class Meta:
[docs]        name = "feynmandiagram"

[docs]    propagators: List[Propagator] = field(
        default_factory=list,
        metadata={"name": "propagator", "type": "Element", "namespace": ""},
    )
[docs]    vertices: List[Vertex] = field(
        default_factory=list,
        metadata={"name": "vertex", "type": "Element", "namespace": ""},
    )
[docs]    legs: List[Leg] = field(
        default_factory=list,
        metadata={"name": "leg", "type": "Element", "namespace": ""},
    )

[docs]    def get_point(self, id):
        for v in self.vertices:
            if v.id == id:
                return v
        for l in self.legs:
            if l.id == id:
                return l
        return None