"""
Components to draw a 2D floor plan of a given MAD-X lattice.
"""
# TODO: improve display of vertically oriented elements
# TODO: adjust display for custom entry/exit pole faces
# TODO: load styles from config
# TODO: rotate/place scene according to space requirements
__all__ = [
'FloorCoords',
'FloorPlanWidget',
'LatticeFloorPlan',
]
from collections import namedtuple
from math import cos, sin, sqrt, pi, atan2, floor, log10
import numpy as np
from PyQt5.QtCore import QMarginsF, QPointF, QRectF, Qt
from PyQt5.QtGui import (
QBrush, QColor, QFont, QFontMetrics, QPainter, QPainterPath,
QPen, QPolygonF)
from PyQt5.QtWidgets import (
QGraphicsItem, QGraphicsScene, QGraphicsView,
QPushButton, QWidget)
from madgui.util.layout import VBoxLayout
FloorCoords = namedtuple('FloorCoords', ['x', 'y', 'z', 'theta', 'phi', 'psi'])
ELEMENT_COLOR = {
'E_GUN': 'purple',
'SBEND': 'red',
'QUADRUPOLE': 'blue',
'DRIFT': 'black',
'LCAVITY': 'green',
'RFCAVITY': 'green',
'SEXTUPOLE': 'yellow',
'WIGGLER': 'orange',
}
ELEMENT_WIDTH = {
'E_GUN': 1.0,
'LCAVITY': 0.4,
'RFCAVITY': 0.4,
'SBEND': 0.6,
'QUADRUPOLE': 0.4,
'SEXTUPOLE': 0.5,
'DRIFT': 0.1,
}
rot90 = np.array([[0, -1], [1, 0]])
def Rotation2(phi):
c, s = cos(phi), sin(phi)
return lambda x, y: (c*x - s*y, c*y + s*x)
def Rotation3(theta, phi, psi, *, Rotation2=Rotation2):
ry = Rotation2(theta)
rx = Rotation2(-phi)
rz = Rotation2(psi)
def rotate(x, y, z):
x, y = rz(x, y)
y, z = rx(y, z)
z, x = ry(z, x)
return x, y, z
return rotate
def Projection(ax1, ax2):
ax1 = np.array(ax1) / np.dot(ax1, ax1)
ax2 = np.array(ax2) / np.dot(ax2, ax2)
return np.array([ax1, ax2])
def normalize(vec):
if np.allclose(vec, 0):
return np.zeros(2)
return vec / sqrt(np.dot(vec, vec))
[docs]class LatticeFloorPlan(QGraphicsView):
"""
Graphics widget to draw 2D floor plan of given lattice.
"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.setInteractive(True)
self.setDragMode(QGraphicsView.ScrollHandDrag)
self.setBackgroundBrush(QBrush(Qt.white, Qt.SolidPattern))
self.setProjection(-pi/2, pi/2)
self.selection = set()
[docs] def setProjection(self, theta, phi, psi=0):
phi = np.clip(phi, -pi/8, pi/2)
rot = Rotation3(theta, phi, psi)
ax1 = np.array(list(rot(1, 0, 0)))
ax2 = np.array(list(rot(0, 1, 0)))
self.theta, self.phi, self.psi = theta, phi, psi
self.projection = Projection(ax1, -ax2)
if self.replay is not None:
self.scene().clear()
self.setElements(*self.replay)
session = None
model = None
[docs] def set_session(self, session):
self.session = session
self.session.model.changed.connect(self.set_model)
self.set_model(session.model())
[docs] def set_model(self, model):
# TODO: only update when SBEND/MULTIPOLE/SROTATION etc changes?
self.selection = self.session.selected_elements
if self.model:
self.model.updated.disconnect(self._updateSurvey)
self.model = model
if model:
self.model.updated.connect(self._updateSurvey)
self._updateSurvey()
def _updateSurvey(self):
survey = self.model.survey()
array = np.array([survey[key] for key in FloorCoords._fields])
floor = [FloorCoords(*row) for row in array.T]
self.setElements(self.model.elements, floor, self.selection)
replay = None
[docs] def setElements(self, elements, survey, selection):
self.replay = elements, survey, selection
self.setScene(QGraphicsScene(self))
survey = [FloorCoords(0, 0, 0, 0, 0, 0)] + survey
for element, coords in zip(elements, zip(survey, survey[1:])):
self.scene().addItem(
ElementGraphicsItem(self, element, coords, selection))
self.coordinate_axes = CoordinateAxes(self)
self.scale_indicator = ScaleIndicator(self)
self.scene().addItem(self.coordinate_axes)
self.scene().addItem(self.scale_indicator)
self.setViewRect(self._sceneRect())
if hasattr(selection, 'update_finished'):
selection.update_finished.connect(self._update_selection)
def _sceneRect(self):
rect = self.scene().sceneRect()
return rect.marginsAdded(QMarginsF(
0.05*rect.width(), 0.05*rect.height(),
0.05*rect.width(), 0.05*rect.height(),
))
[docs] def resizeEvent(self, event):
"""Maintain visible region on resize."""
self.setViewRect(self.view_rect)
super().resizeEvent(event)
[docs] def mapRectToScene(self, rect):
"""
Map topleft/botright rect from viewport to scene coordinates.
This assumes there is no rotation/shearing.
"""
return QRectF(
self.mapToScene(rect.topLeft()),
self.mapToScene(rect.bottomRight()))
[docs] def setViewRect(self, rect):
"""
Fit the given scene rectangle into the visible view.
This assumes there is no rotation/shearing.
"""
cur = self.mapRectToScene(self.viewport().rect())
new = rect.intersected(self._sceneRect())
self.zoom(min(cur.width()/new.width(),
cur.height()/new.height()))
self.view_rect = new
self.coordinate_axes.update()
self.scale_indicator.update()
[docs] def zoom(self, scale):
"""Scale the figure uniformly along both axes."""
self.scale(scale, scale)
self.view_rect = self.mapRectToScene(self.viewport().rect())
[docs] def wheelEvent(self, event):
"""Handle mouse wheel as zoom."""
delta = event.angleDelta().y()
self.zoom(1.0 + delta/1000.0)
[docs] def mousePressEvent(self, event):
self.last_mouse_position = event.pos()
super().mousePressEvent(event)
[docs] def mouseMoveEvent(self, event):
if event.buttons() == Qt.RightButton:
delta = event.pos() - self.last_mouse_position
theta = self.theta + delta.x()/100
phi = self.phi + delta.y()/100
self.setProjection(theta, phi)
self.last_mouse_position = event.pos()
event.accept()
return
super().mouseMoveEvent(event)
def _update_selection(self, slice, old_values, new_values):
insert = set(new_values) - set(old_values)
delete = set(old_values) - set(new_values)
for item in self.scene().items():
if item.el_id in insert:
item.setSelected(True)
if item.el_id in delete:
item.setSelected(False)
class ElementGraphicsItem(QGraphicsItem):
"""Base class for element graphics items."""
outline_pen = {'width': 1}
orbit_pen = {'style': Qt.DashLine}
select_pen = {'style': Qt.DashLine,
'color': 'green',
'width': 4}
def __init__(self, plan, element, coords, selection):
super().__init__()
self.plan = plan
self.coords = coords
self.rotate = (Rotation3(coords[0].theta, coords[0].phi, coords[0].psi),
Rotation3(coords[1].theta, coords[1].phi, coords[1].psi))
self.element = element
self.length = element.length
self.angle = float(element.get('angle', 0.0))
self.width = getElementWidth(element)
self.color = getElementColor(element)
self.walls = (0.5*self.width, 0.5*self.width) # inner/outer wall widths
self.selection = selection
self._outline = self.outline()
self._orbit = self.orbit()
self.setFlag(QGraphicsItem.ItemIsSelectable, True)
self.setSelected(self.el_id in selection)
@property
def el_id(self):
return self.element.index
def itemChange(self, change, value):
if change == QGraphicsItem.ItemSelectedHasChanged:
self._on_select(value)
return value
def _on_select(self, select):
is_selected = self.el_id in self.selection
if select and not is_selected:
# TODO: incorporate whether shift is clicked
self.selection.add(self.el_id)
elif is_selected and not select:
self.selection.remove(self.el_id)
def shape(self):
return self._outline
def boundingRect(self):
return self._outline.boundingRect()
def paint(self, painter, option, widget):
"""Paint element + orbit + selection frame."""
painter.setRenderHint(QPainter.Antialiasing)
painter.setBrush(Qt.NoBrush)
# draw element outline:
painter.setPen(createPen(**self.outline_pen))
painter.fillPath(self._outline, self.color)
painter.drawPath(self._outline)
# draw beam orbit:
painter.setPen(createPen(**self.orbit_pen))
painter.drawPath(self._orbit)
# highlight selected elements:
if self.isSelected():
painter.setPen(createPen(**self.select_pen))
painter.drawPath(self._outline)
def endpoints(self):
proj2D = self.plan.projection.dot
p0, p1 = self.coords
return (proj2D([p0.x, p0.y, p0.z]),
proj2D([p1.x, p1.y, p1.z]))
def outline(self):
"""Return a QPainterPath that outlines the element."""
r1, r2 = self.walls
p0, p1 = self.endpoints()
proj2D = self.plan.projection.dot
vec0 = normalize(np.dot(rot90, proj2D(list(self.rotate[0](0, 0, 1)))))
vec1 = normalize(np.dot(rot90, proj2D(list(self.rotate[1](0, 0, 1)))))
path = QPainterPath()
path.moveTo(*(p0 - r2*vec0))
path.lineTo(*(p1 - r2*vec1))
path.lineTo(*(p1 + r1*vec1))
path.lineTo(*(p0 + r1*vec0))
path.closeSubpath()
return path
def orbit(self):
"""Return a QPainterPath that shows the beam orbit."""
a, b = self.endpoints()
path = QPainterPath()
path.moveTo(*a)
path.lineTo(*b)
return path
def getElementColor(element, default='black'):
return QColor(ELEMENT_COLOR.get(element.base_name.upper(), default))
def getElementWidth(element, default=0.2):
return ELEMENT_WIDTH.get(element.base_name.upper(), default)
def createPen(style=Qt.SolidLine, color='black', width=1):
"""
Use this function to conveniently create a cosmetic pen with specified
width. Integer widths create cosmetic pens (default) and float widths
create scaling pens (this way you can set the figure style by changing a
number).
This is particularly important on PyQt5 where the default pen blacks out
large areas of the figure if not being careful.
"""
pen = QPen(style)
pen.setColor(QColor(color))
if isinstance(width, int):
pen.setWidth(width)
pen.setCosmetic(True)
else:
pen.setWidthF(width)
pen.setCosmetic(False)
return pen
class CoordinateAxes(QGraphicsItem):
"""Display axes of coordinates."""
el_id = None
pen = {'style': Qt.SolidLine,
'color': 'orange',
'width': 1}
def __init__(self, plan):
super().__init__()
self.plan = plan
self.setFlag(QGraphicsItem.ItemIgnoresTransformations, True)
def update(self):
self._path = self.draw_path()
def shape(self):
return self._path
def boundingRect(self):
# Ignore this item when calculating the scene rect:
return QRectF()
def paint(self, painter, option, widget):
pen = createPen(**self.pen)
painter.setRenderHint(QPainter.Antialiasing)
painter.setPen(pen)
painter.setBrush(QBrush(pen.color(), Qt.SolidPattern))
painter.drawPath(self._path)
def draw_path(self):
l, s, d = 45, 10, 10
proj = self.plan.projection.dot
orig = np.array([0, 0])
axes = QPainterPath()
axes.addPath(self.axis_arrow("x", orig, orig+l*proj([1, 0, 0]), s))
axes.addPath(self.axis_arrow("y", orig, orig+l*proj([0, 1, 0]), s))
axes.addPath(self.axis_arrow("z", orig, orig+l*proj([0, 0, 1]), s))
tran = self.deviceTransform(self.plan.viewportTransform()).inverted()[0]
view = tran.mapRect(QRectF(self.plan.viewport().rect()))
rect = axes.boundingRect()
axes.translate(view.left() + view.width()/15 - rect.left(),
view.bottom() - view.height()/15 - rect.bottom())
path = QPainterPath()
path.addPath(axes)
path.addEllipse(-d/2, -d/2, d, d)
return path
def axis_arrow(self, label, x0, x1, arrow_size):
path = arrow(x0, x1, arrow_size)
if not path:
return QPainterPath()
plan = self.plan
font = QFont(plan.font())
font.setPointSize(14)
metr = QFontMetrics(font)
rect = metr.boundingRect(label)
rect.setHeight(metr.xHeight())
tran = self.deviceTransform(self.plan.viewportTransform()).inverted()[0]
size = tran.mapRect(QRectF(rect)).size()
w, h = size.width(), size.height()
dir_ = (x1 - x0) / np.linalg.norm(x1 - x0)
offs = [-w/2, +h/2] + dir_ * max(w, h)
path.addText(QPointF(*(x1 + offs)), font, label)
return path
def arrow(x0, x1, arrow_size=0.3, arrow_angle=pi/5):
dx, dy = x1 - x0
if dy**2 + dx**2 < arrow_size**2:
return None
path = QPainterPath()
path.moveTo(*x0)
path.lineTo(*x1)
angle = atan2(dy, dx)
p1 = x1 + [cos(angle + pi + arrow_angle) * arrow_size,
sin(angle + pi + arrow_angle) * arrow_size]
p2 = x1 + [cos(angle + pi - arrow_angle) * arrow_size,
sin(angle + pi - arrow_angle) * arrow_size]
path.addPolygon(QPolygonF([
QPointF(*x1),
QPointF(*p1),
QPointF(*p2),
QPointF(*x1),
]))
return path
class ScaleIndicator(QGraphicsItem):
"""Display small scale indicator."""
el_id = None
pen = {'style': Qt.SolidLine,
'color': 'orange',
'width': 2}
def __init__(self, plan):
super().__init__()
self.plan = plan
self.setFlag(QGraphicsItem.ItemIgnoresTransformations, True)
def update(self):
self._path = self.draw_path()
def shape(self):
return self._path
def boundingRect(self):
# Ignore this item when calculating the scene rect:
return QRectF()
def paint(self, painter, option, widget):
pen = createPen(**self.pen)
painter.setRenderHint(QPainter.Antialiasing)
painter.setPen(pen)
painter.setBrush(Qt.NoBrush)
painter.drawPath(self._path)
def draw_path(self):
plan = self.plan
rect = plan.mapRectToScene(plan.viewport().rect())
rect.setWidth(10**floor(log10(rect.width()/2)))
text = "{} m".format(round(rect.width()))
tran = self.deviceTransform(plan.viewportTransform()).inverted()[0]
width = tran.mapRect(QRectF(
plan.mapFromScene(rect.topLeft()),
plan.mapFromScene(rect.bottomRight()))).width()
view = tran.mapRect(plan.viewport().rect())
x0 = QPointF(view.right() - view.width()/15,
view.bottom() - view.height()/15)
x1 = x0 - QPointF(width, 0)
head = QPointF(0, 8)
path = QPainterPath()
path.moveTo(x0)
path.lineTo(x1)
path.moveTo(x0 + head)
path.lineTo(x0 - head)
path.moveTo(x1 + head)
path.lineTo(x1 - head)
# add label
font = QFont(plan.font())
font.setPointSize(14)
rect = QFontMetrics(font).boundingRect(text)
size = tran.mapRect(QRectF(rect)).size()
w, h = size.width(), size.height()
offs = [-w/2, -h/2]
path.addText((x0+x1)/2 + QPointF(*offs), font, text)
return path