from ovito.vis import ViewportOverlayInterface
from ovito.data import DataCollection
from ovito.qt_compat import QtCore, QtGui
from traits.api import Range
class HighlightParticleOverlay(ViewportOverlayInterface):
# Adjustable user parameter that selects which particle to highlight:
particle_index = Range(value=0, low=0, label='Particle index')
def render(self, canvas: ViewportOverlayInterface.Canvas, data: DataCollection, **kwargs):
# Determine world-space radius of the particle.
radius = 0.0
if 'Radius' in data.particles:
radius = data.particles['Radius'][self.particle_index]
if radius <= 0 and data.particles.particle_types is not None:
particle_type = data.particles.particle_types[self.particle_index]
radius = data.particles.particle_types.type_by_id(particle_type).radius
if radius <= 0:
radius = data.particles.vis.radius
# Project center of the particle to screen space.
positions = data.particles.positions
xy = canvas.project_location(positions[self.particle_index])
if xy is None:
return
# Calculate screen-space size of the particle as a fraction of the canvas height.
screen_radius = canvas.project_length(positions[self.particle_index], radius)
# Convert everything to logical pixel coordinates used by the QPainter.
x = xy[0] * canvas.logical_size[0]
y = (1 - xy[1]) * canvas.logical_size[1]
screen_radius *= canvas.logical_size[1]
# Start drawing using a QPainter.
with canvas.qt_painter() as painter:
# Draw a dashed circle around the particle.
pen = QtGui.QPen(QtCore.Qt.DashLine)
pen.setWidth(3)
pen.setColor(QtGui.QColor(0,0,255))
painter.setPen(pen)
painter.drawEllipse(QtCore.QPointF(x, y), screen_radius, screen_radius)
# Draw an arrow pointing at the particle.
arrow_shape = QtGui.QPolygonF()
arrow_shape.append(QtCore.QPointF(0,0))
arrow_shape.append(QtCore.QPointF(10,10))
arrow_shape.append(QtCore.QPointF(10,5))
arrow_shape.append(QtCore.QPointF(40,5))
arrow_shape.append(QtCore.QPointF(40,-5))
arrow_shape.append(QtCore.QPointF(10,-5))
arrow_shape.append(QtCore.QPointF(10,-10))
painter.setPen(QtGui.QPen())
painter.setBrush(QtGui.QBrush(QtGui.QColor(255,0,0)))
painter.translate(QtCore.QPointF(x, y))
painter.rotate(-45.0)
painter.translate(QtCore.QPointF(screen_radius,0))
painter.scale(2,2)
painter.drawPolygon(arrow_shape)
