Added root flattening

This commit is contained in:
2026-08-07 16:55:47 +02:00
parent 8b98bf9a73
commit 02bb260fea
6 changed files with 85 additions and 12 deletions
+58
View File
@@ -49,6 +49,8 @@ _STYLES: dict[str, dict] = {
"Undulator": {"top": "#558B2F", "bot": "#33691E", "w": 88, "h": 36},
"Diagnostics": {"top": "#4527A0", "bot": "#311B92", "w": 64, "h": 44},
"Vacuum": {"top": "#00838F", "bot": "#006064", "w": 64, "h": 36},
"Photonics": {"top": "#EF6C00", "bot": "#BF360C", "w": 66, "h": 40},
"Dechirper": {"top": "#37474F", "bot": "#1C262B", "w": 78, "h": 38},
"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},
@@ -417,6 +419,60 @@ class VacuumNode(BeamlineNode):
p.drawEllipse(QPointF(cx, cy), r, r)
p.drawLine(QPointF(cx, cy - r - 4), QPointF(cx, cy + r + 4))
class PhotonicsNode(BeamlineNode):
"""Photonics element — orange rect with a light-ray / lens symbol."""
def __init__(self, pos): super().__init__("Photonics", pos)
def _draw_body(self, p):
p.setBrush(QBrush(self._gradient()))
p.setPen(QPen(QColor("#E65100"), 1.5))
p.drawRoundedRect(QRectF(0, 0, self.W, self.H), 5, 5)
cy = self.H / 2
# convex lens in the middle
p.setBrush(Qt.NoBrush)
p.setPen(QPen(QColor("#FFCC80"), 1.5))
lens = QPainterPath()
lx = self.W / 2
lens.moveTo(lx, cy - 11)
lens.quadTo(lx + 7, cy, lx, cy + 11)
lens.quadTo(lx - 7, cy, lx, cy - 11)
p.drawPath(lens)
# incoming / outgoing rays
p.setPen(QPen(QColor("#FFE0B2"), 1.2))
for dy in (-6, 0, 6):
p.drawLine(QPointF(6, cy + dy), QPointF(lx - 5, cy + dy))
p.drawLine(QPointF(lx + 5, cy + dy * 0.4),
QPointF(self.W - 6, cy + dy * 0.4))
class DechirperNode(BeamlineNode):
"""Dechirper — dark slate rect with corrugated-plate symbol."""
def __init__(self, pos): super().__init__("Dechirper", pos)
def _draw_body(self, p):
p.setBrush(QBrush(self._gradient()))
p.setPen(QPen(QColor("#546E7A"), 1.5))
p.drawRect(QRectF(0, 0, self.W, self.H))
# corrugated plates: square wave along top and bottom
p.setPen(QPen(QColor("#B0BEC5"), 1.3))
step = 6
amp = 4
gap = 7 # distance of plate from the mid-plane
cy = self.H / 2
for sign in (-1, 1):
base_y = cy + sign * gap
path = QPainterPath()
path.moveTo(5, base_y)
x = 5.0
up = True
while x < self.W - 5:
nx = min(x + step, self.W - 5)
y = base_y + (sign * amp if up else 0)
path.lineTo(x, y)
path.lineTo(nx, y)
x = nx
up = not up
p.setBrush(Qt.NoBrush)
p.drawPath(path)
class TargetNode(BeamlineNode):
def __init__(self, pos): super().__init__("Target", pos)
def _draw_body(self, p):
@@ -605,6 +661,8 @@ _NODE_MAP: dict[str, type] = {
"Undulator": UndulatorNode,
"Diagnostics": DiagnosticsNode,
"Vacuum": VacuumNode,
"Photonics": PhotonicsNode,
"Dechirper": DechirperNode,
"Reference": ReferenceNode,
"Target": TargetNode,
"Alignment": AlignmentNode,
+1 -1
View File
@@ -17,7 +17,7 @@ from .nodes import _STYLES
ELEMENT_GROUPS: list[tuple[str, list[str]]] = [
("Elements", ["Dipole", "Quadrupole", "Sextupole",
"Solenoid", "Corrector", "RF", "Undulator", "Diagnostics",
"Vacuum"]),
"Vacuum", "Photonics", "Dechirper"]),
("Markers", ["Reference", "Start", "Target", "Alignment", "Branch", "Case", "Merge"]),
("Containers", ["Line", "Variable Line"]),
]
+7 -1
View File
@@ -333,6 +333,8 @@ class BeamlineEditor(QtWidgets.QMainWindow, Ui_BeamlineGUI):
ele = Diagnostic({})
case "Vacuum":
ele = Vacuum({})
case "Photonics":
ele = Photonic({})
case "Reference":
ele = Marker({"Tag":"QREF"})
case "Start":
@@ -486,6 +488,7 @@ class BeamlineEditor(QtWidgets.QMainWindow, Ui_BeamlineGUI):
if ref in self.elementDB.keys():
secondBranch = self.unwrap(name + name1, ref, subline,case_num)
if secondBranch:
print('Was branching - should continue with second branch, returning from',name+name1)
break # do not continue with any other elements
else:
if self.verbose:
@@ -525,14 +528,17 @@ class BeamlineEditor(QtWidgets.QMainWindow, Ui_BeamlineGUI):
self.elementDB[ele].angle = self.elementDB[ele].design_angle
line.hasBranchPoint = True
secondBranch=True
break
return secondBranch
else:
self.elementDB[ele].angle=0
line.hasBranchPoint = False
if secondBranch:
print("There should be now following the second branch from container:", name,tree.secondary)
if tree.secondary:
print("Has secondary tree", name)
self.unwrap(name, tree.secondary.node_id, line, case_num)
else:
print("Has no secondary tree", name,'- returning to parent container')
return secondBranch
else:
if tree.primary:
+3
View File
@@ -23,6 +23,7 @@ def getPropertyList(element_type: str | None) -> dict:
elif element_type == 'Variable Line':
return {'L':'Position/Total Length','Loff':'Position/Constant Offset','ID':'Anchor/Reference'}
fields={}
for i , fld in enumerate(eName):
fields[fld]='Name/'+Name[i]
@@ -59,6 +60,8 @@ def getPropertyList(element_type: str | None) -> dict:
fields['helical'] = 'Setting/Helical'
case "Diagnostic":
fields['Seval'] = 'Position/Evaluation Point'
case "Photonics":
fields['Seval'] = 'Position/Evaluation Point'
case "Start":
fields['RefID'] = 'Setting/Reference Element ID'
case "Target":
+8 -5
View File
@@ -223,13 +223,16 @@ class Ui_BeamlineGUI(object):
self.menuLAYOUT.addSeparator()
self.menuLAYOUT.addSeparator()
self.menuLAYOUT.addAction(self.actionExport_as_Proto_List)
self.menuLAYOUT.addAction(self.actionCompareLayout)
self.menuLAYOUT.addSeparator()
self.menuLAYOUT.addAction(self.actionExportMadX)
self.menuLAYOUT.addAction(self.actionExportElegant)
self.menuFLATTEN.addAction(self.actionCase_1)
self.menuFLATTEN.addAction(self.actionCase_2)
self.menuFLATTEN.addAction(self.actionCase_3)
self.menuFLATTEN.addSeparator()
self.menuFLATTEN.addAction(self.actionCompareLayout)
self.menuFLATTEN.addSeparator()
self.menuFLATTEN.addSeparator()
self.menubar.addAction(self.menuFile.menuAction())
self.menubar.addAction(self.menuFLATTEN.menuAction())
self.menubar.addAction(self.menuLAYOUT.menuAction())
@@ -283,7 +286,7 @@ class Ui_BeamlineGUI(object):
item.setText(_translate("BeamlineGUI", "SATBD01"))
self.tabWidget.setTabText(self.tabWidget.indexOf(self.tab_4), _translate("BeamlineGUI", "Settings"))
self.menuFile.setTitle(_translate("BeamlineGUI", "FILE"))
self.menuLAYOUT.setTitle(_translate("BeamlineGUI", "LAYOUT"))
self.menuLAYOUT.setTitle(_translate("BeamlineGUI", "EXPORT"))
self.menuFLATTEN.setTitle(_translate("BeamlineGUI", "FLATTEN"))
self.actionNew_Layout.setText(_translate("BeamlineGUI", "New Layout"))
self.actionSave.setText(_translate("BeamlineGUI", "Save"))
@@ -300,9 +303,9 @@ class Ui_BeamlineGUI(object):
self.actionExportMadX.setText(_translate("BeamlineGUI", "Export as MadX Lattice..."))
self.actionExportElegant.setText(_translate("BeamlineGUI", "Export as ELegant Lattice..."))
self.actionExport_as_Proto_List.setText(_translate("BeamlineGUI", "Export as Proto List..."))
self.actionCase_1.setText(_translate("BeamlineGUI", "Case 1"))
self.actionCase_1.setText(_translate("BeamlineGUI", "Case 1 (Current)"))
self.actionCase_1.setShortcut(_translate("BeamlineGUI", "Ctrl+1"))
self.actionCase_2.setText(_translate("BeamlineGUI", "Case 2"))
self.actionCase_2.setText(_translate("BeamlineGUI", "Case 2 (Planned)"))
self.actionCase_2.setShortcut(_translate("BeamlineGUI", "Ctrl+2"))
self.actionCase_3.setText(_translate("BeamlineGUI", "Case 3"))
self.actionCase_3.setText(_translate("BeamlineGUI", "Case 3 (FInal)"))
self.actionCase_3.setShortcut(_translate("BeamlineGUI", "Ctrl+3"))
+8 -5
View File
@@ -324,13 +324,12 @@ background-color: rgb(13, 15, 26);</string>
</widget>
<widget class="QMenu" name="menuLAYOUT">
<property name="title">
<string>LAYOUT</string>
<string>EXPORT</string>
</property>
<addaction name="actionExportOnlineModel"/>
<addaction name="separator"/>
<addaction name="separator"/>
<addaction name="actionExport_as_Proto_List"/>
<addaction name="actionCompareLayout"/>
<addaction name="separator"/>
<addaction name="actionExportMadX"/>
<addaction name="actionExportElegant"/>
@@ -342,6 +341,10 @@ background-color: rgb(13, 15, 26);</string>
<addaction name="actionCase_1"/>
<addaction name="actionCase_2"/>
<addaction name="actionCase_3"/>
<addaction name="separator"/>
<addaction name="actionCompareLayout"/>
<addaction name="separator"/>
<addaction name="separator"/>
</widget>
<addaction name="menuFile"/>
<addaction name="menuFLATTEN"/>
@@ -419,7 +422,7 @@ background-color: rgb(13, 15, 26);</string>
</action>
<action name="actionCase_1">
<property name="text">
<string>Case 1</string>
<string>Case 1 (Current)</string>
</property>
<property name="shortcut">
<string>Ctrl+1</string>
@@ -427,7 +430,7 @@ background-color: rgb(13, 15, 26);</string>
</action>
<action name="actionCase_2">
<property name="text">
<string>Case 2</string>
<string>Case 2 (Planned)</string>
</property>
<property name="shortcut">
<string>Ctrl+2</string>
@@ -435,7 +438,7 @@ background-color: rgb(13, 15, 26);</string>
</action>
<action name="actionCase_3">
<property name="text">
<string>Case 3</string>
<string>Case 3 (FInal)</string>
</property>
<property name="shortcut">
<string>Ctrl+3</string>