diff --git a/beamline_editor/nodes.py b/beamline_editor/nodes.py index d9de588..2804245 100644 --- a/beamline_editor/nodes.py +++ b/beamline_editor/nodes.py @@ -51,6 +51,7 @@ _STYLES: dict[str, dict] = { "Vacuum": {"top": "#00838F", "bot": "#006064", "w": 64, "h": 36}, "Target": {"top": "#B71C1C", "bot": "#7B0000", "w": 36, "h": 36}, "Alignment": {"top": "#1565C0", "bot": "#0D3B8C", "w": 36, "h": 36}, + "Branch": {"top": "#F9A825", "bot": "#C17900", "w": 44, "h": 44}, "Start": {"top": "#00695C", "bot": "#004D40", "w": 54, "h": 40}, "Case": {"top": "#E65100", "bot": "#BF360C", "w": 60, "h": 72}, "Merge": {"top": "#1A237E", "bot": "#0D1452", "w": 60, "h": 72}, @@ -446,6 +447,28 @@ class AlignmentNode(BeamlineNode): p.drawLine(QPointF(cx, cy), QPointF(cx + dx, cy)) p.drawLine(QPointF(cx, cy), QPointF(cx, cy + dy)) +class BranchNode(BeamlineNode): + """Branch marker — 1 in, 1 out. Amber diamond with a fork glyph.""" + def __init__(self, pos): super().__init__("Branch", pos) + def _draw_body(self, p): + cx, cy = self.W / 2, self.H / 2 + diamond = QPolygonF([ + QPointF(cx, 2), + QPointF(self.W - 2, cy), + QPointF(cx, self.H - 2), + QPointF(2, cy), + ]) + p.setBrush(QBrush(self._gradient())) + p.setPen(QPen(QColor("#C17900"), 1.5)) + p.drawPolygon(diamond) + # decorative fork glyph (visual only) + p.setPen(QPen(QColor("#FFECB3"), 1.4)) + mid_x = cx - 4 + p.drawLine(QPointF(mid_x, cy), QPointF(mid_x + 8, cy - 7)) + p.drawLine(QPointF(mid_x, cy), QPointF(mid_x + 8, cy + 7)) + p.drawLine(QPointF(cx - 10, cy), QPointF(mid_x, cy)) + + class StartNode(BeamlineNode): """ Start marker: no in-port, 1 out. @@ -454,7 +477,6 @@ class StartNode(BeamlineNode): """ def __init__(self, pos): super().__init__("Start", pos) - # Second-row label for collapsed state — hidden by default self._collapsed_label = QGraphicsTextItem("(collapsed)", self) self._collapsed_label.setDefaultTextColor(QColor("#A5D6A7")) self._collapsed_label.setFont(QFont("Segoe UI", 6)) @@ -463,7 +485,6 @@ class StartNode(BeamlineNode): self._layout_labels() def _layout_labels(self): - """Stack name and (collapsed) label vertically, centred in the node.""" br1 = self._label.boundingRect() br2 = self._collapsed_label.boundingRect() if self._collapsed_label.isVisible(): @@ -478,7 +499,6 @@ class StartNode(BeamlineNode): self._collapsed_label.setPos((self.W - br2.width()) / 2, y2) def set_display_name(self, name: str): - """Override to re-run the two-row layout after the name changes.""" super().set_display_name(name) self._layout_labels() @@ -491,9 +511,8 @@ class StartNode(BeamlineNode): p.setBrush(QBrush(self._gradient())) p.setPen(QPen(QColor("#00897B"), 1.5)) p.drawRoundedRect(QRectF(0, 0, self.W, self.H), 6, 6) - # Arrow-head at the right edge, pointing right toward the out-port - cy = self.H / 2 - tip_x = self.W - 4 # flush with right edge + cy = self.H / 2 + tip_x = self.W - 4 base_x = tip_x - 12 half_h = 8 pts = QPolygonF([ @@ -506,6 +525,7 @@ class StartNode(BeamlineNode): p.setPen(Qt.NoPen) p.drawPolygon(pts) + class CaseNode(BeamlineNode): """ Case marker — 1 in, 3 out. @@ -588,6 +608,7 @@ _NODE_MAP: dict[str, type] = { "Reference": ReferenceNode, "Target": TargetNode, "Alignment": AlignmentNode, + "Branch": BranchNode, "Overlap": OverlapNode, "Start": StartNode, "Case": CaseNode, diff --git a/beamline_editor/palette.py b/beamline_editor/palette.py index d49f972..fbb3abc 100644 --- a/beamline_editor/palette.py +++ b/beamline_editor/palette.py @@ -18,7 +18,7 @@ ELEMENT_GROUPS: list[tuple[str, list[str]]] = [ ("Elements", ["Dipole", "Quadrupole", "Sextupole", "Solenoid", "Corrector", "RF", "Undulator", "Diagnostics", "Vacuum"]), - ("Markers", ["Reference", "Start", "Target", "Alignment", "Case", "Merge"]), + ("Markers", ["Reference", "Start", "Target", "Alignment", "Branch", "Case", "Merge"]), ("Containers", ["Line", "Variable Line"]), ] diff --git a/layoutexporter.py b/layoutexporter.py new file mode 100644 index 0000000..7183aae --- /dev/null +++ b/layoutexporter.py @@ -0,0 +1,110 @@ +from containers import VariableContainer +from elements import * + +class Exporter: + def __init__(self): + self.path = 'onlinemodel/exportedlayouts/layout.py' + self.fid = None + self.lattice='11.1.2' + + def openFile(self): + print(self.path) + self.fid = open(self.path,'w') + self.printHeader() + + def closeFile(self): + if self.fid is None: + return + self.fid.close() + self.fid = None + + def printHeader(self): + if self.fid is None: + return + fid = self.fid + fid.write('from onlinemodel.core.omtype import *\n\n') + fid.write('class SwissFEL:\n') + fid.write('\tdef __init__(self,alt = 0):\n') + fid.write('\t\tself.alt=alt\n') + fid.write('\t\tself.Version="%s"\n\n' % self.lattice) + fid.write('\tdef initialize(self,ElementDB):\n') + fid.write('\t\treturn\n\n') + fid.write('\tdef build(self):\n') + fid.write('\t\tif self.alt == 0:\n') + fid.write('\t\t\treturn self.build_current()\n') + fid.write('\t\telif self.alt == 1:\n') + fid.write('\t\t\treturn self.build_planned()\n') + fid.write('\t\telse:\n') + fid.write('\t\t\treturn self.build_final()\n\n') + + def buildCaseHeader(self,case): + cases=['_current','_planned','_final'] + self.fid.write('\tdef build_%s(self):\n' % cases[case]) + + + def writeDomain(self,Reference,Name,Length,Elements): + self.fid.write('\n\t\t%sLine=[' % Reference) + for ie,ele in enumerate(Elements): + self.fid.write('%s' % ele) + if ie < len(Elements)-1: + self.fid.write(',') + + self.fid.write(']\n') + self.fid.write('\t\t%s=LineContainer("%s",%f)\n' % (Reference,Name,Length)) + self.fid.write('\t\t%s.append(%sLine)\n\n' % (Reference,Reference)) + + + def writeSection(self,Name,Length): + print('Write Line Container',Name) + self.fid.write('\n\t\t%s=LineContainer("%s",%f)\n' % (Name,Name,Length)) + + def writeElement(self,Container,Element): + + prop, typename,off, relative= self.getProp(Element) + if prop is None: + return + self.fid.write('\t\tprop=%s\n' % prop) + self.fid.write('\t\t%s.append(%s(prop),%f,"%s")\n' % (Container,typename,off,relative)) + + + def getProp(self,Element): + if isinstance(Element, Marker) or isinstance(Element, VariableContainer): + return None,1,1,1 + prop= '{"Tag":"%s","index":%d,"Baugruppe":"%s","Length":%f' % (Element.Tag,Element.index,Element.Baugruppe, Element.Length) + if Element.LengthRes>Element.Length: + prop+=',"LengthRes":%f' % Element.LengthRes + if Element.Tilt>0: + prop+=',"Tilt":%f' % (Element.Tilt*np.pi/180.) + if isinstance(Element, RF): + prop+=',"Band":%s' % Element.Band + if isinstance(Element, Diagnostic): + if Element.Seval > 0: + prop+=',"Seval":%f' % Element.seval + if isinstance(Element,Undulator): + prop+=',"K"=%f,"ku":%f' % (Element.K,2*np.pi/Element.lambdau) + prop+='}' + pos = 'absolute' + if Element.relOff: + pos = 'relative' + offs=Element.__dict__.get('OffsetS',0) + typename=Element.__class__.__name__ + return prop,typename,offs,pos + + + + + def buildCaseFooter(self): + self.fid.write('\t\tPartsList = [[Injector, -1, 0, "SwissFEL Injector"],\n') + self.fid.write('\t\t [InjectorDump, 0, 1, "SwissFEL Injector beam dump"]]\n') + self.fid.write('\t\treturn PartsList\n\n') + +# PartsList = [[Injector, -1, 0, 'SwissFEL Injector'], # 0 +# [Linac, 0, -1, 'SwissFEL Linac'], # 1 +# [Aramis, 1, -1, 'Aramis beamline'], # 2 +# [Athos, 1, 2, 'Athos beamline'], # 3 +# [InjectorDump, 0, 1, 'SwissFEL Injector beam dump'], # 4 +# [Linac1Dump, 1, 1, 'Linac1 beam dump'], # 5 +# [AramisDump, 2, 1, 'Aramis beam dump'], # 6 +# [AthosDump, 3, 1, 'Athos beam dump'], # 7 +# [Porthos, 1, 3, 'Porthos beamline']] # 8 + diff --git a/main.py b/main.py index 524d600..ffe1faf 100644 --- a/main.py +++ b/main.py @@ -15,6 +15,7 @@ from elements import * from containers import * from propertytypes import getPropertyList from protoliste import ProtoListe +from layoutexporter import Exporter class BeamlineEditor(QtWidgets.QMainWindow, Ui_BeamlineGUI): def __init__(self): @@ -45,7 +46,7 @@ class BeamlineEditor(QtWidgets.QMainWindow, Ui_BeamlineGUI): self.actionLoad.triggered.connect(self.load) self.actionCompareLayout.triggered.connect(self.compareLayout) - + self.actionExportOnlineModel.triggered.connect(self.exportOnlineModel) # elements, lines and branch points self.elementDB={} @@ -99,6 +100,47 @@ class BeamlineEditor(QtWidgets.QMainWindow, Ui_BeamlineGUI): self.savefile=fileName self.editor.loadLayout(self.savefile.replace('.json','_gui.json')) + def exportOnlineModel(self): + print('Exporting Online Model Layout File...') + self.exporter= Exporter() + self.exporter.openFile() + for alt in range(3): + self.exportOnlineModelCase(alt) + self.exporter.closeFile() + + def exportOnlineModelCase(self,alt): + self.exporter.buildCaseHeader(alt) + secs=['IN','SINBD01'] + references={'IN':'Injector','SINBD01':'InjectorDump'} + for sec in secs: + # find element in database + id0 = None + for id in self.elementDB.keys(): + ele = self.elementDB[id] + if isinstance(ele, LineContainer) and ele.Name == sec and ele.firstElementID in self.elementDB.keys(): + id0=id + if id0: + print('Found Section',sec) + else: + continue #skip to next entry + root = self.elementDB[id0].firstElementID + domain=self.elementDB[root] + tree = self.editor.flatten(root, alt+1) + subsec = [] + for ele in tree.path: + if isinstance(self.elementDB[ele], LineContainer): + container=self.elementDB[ele].Name + self.exporter.writeSection(container,self.elementDB[ele].LengthRes) + startid = self.elementDB[ele].firstElementID + if startid in self.elementDB.keys(): + subtree=self.editor.flatten(startid, alt+1) + for atomic in subtree.path: + self.exporter.writeElement(container,self.elementDB[atomic]) + subsec.append(container) + self.exporter.writeDomain(references[sec],sec,domain.LengthRes,subsec) + self.exporter.buildCaseFooter() + + def compareLayout(self): if len(self.plrecs) == 0: return diff --git a/propertytypes.py b/propertytypes.py index e9cfc43..c3ffd9f 100644 --- a/propertytypes.py +++ b/propertytypes.py @@ -16,6 +16,8 @@ def getPropertyList(element_type: str | None) -> dict: return {'Name':'Name/Name','LengthRes':'Position/Reserved Length','firstElementID':'Anchor/Reference'} elif element_type == 'Start' or element_type == 'Target': return {'Name': 'Name/Name'} + elif element_type == 'Branch': + return {'Name': 'Name/Name','Tag':'Name/Tag','index':'Name/Index'} elif element_type == 'Variable Line': return {'L':'Position/Total Length','Loff':'Position/Constant Offset','ID':'Anchor/Reference'}