def getPointCloseToCenter(self, distance):
edgePath = QPainterPath(self.beginPoint)
edgePath.cubicTo(self.curvePoint1, self.curvePoint2, self.endPoint)
if edgePath.length() > 0:
percent = (edgePath.length() / 2 + distance) / edgePath.length()
else:
percent = 0
#Snap to begin/end point when the edge is too small
if percent < 0:
percent = 0
elif percent > 1:
percent = 1
return self.getPointOnEdge(percent)
def getPointCloseToCenter(self, distance):
edgePath = QPainterPath(self.beginPoint)
edgePath.cubicTo(self.curvePoint1, self.curvePoint2, self.endPoint)
if edgePath.length() > 0:
percent = (edgePath.length() / 2 + distance) / edgePath.length()
else:
percent = 0
#Snap to begin/end point when the edge is too small
if percent < 0:
percent = 0
elif percent > 1:
percent = 1
return self.getPointOnEdge(percent)
class CustomRectsItem(QGraphicsItem):
def __init__(self, parent=None):
QGraphicsItem.__init__(self, parent)
self.path = QPainterPath()
self.data_size = 1000
def paint(self,
painter: QtGui.QPainter,
option: 'QStyleOptionGraphicsItem',
widget: typing.Optional[QWidget] = ...):
if not self.path.length():
path = QPainterPath()
path.moveTo(0, 0)
rand = np.random.rand(self.data_size) * 100
for i in np.arange(self.data_size):
path.lineTo(i, rand[i])
self.path = path
painter.save()
pen = QPen()
pen.setCosmetic(True)
pen.setColor(Qt.darkBlue)
painter.setPen(pen)
painter.drawPath(self.path)
painter.restore()
def boundingRect(self):
return QRectF(0, 0, self.data_size, 100)
def __get_bezier_middle(self, bezier: QPainterPath) -> QPointF:
"""
Get the point in the middle of a bezier curve.
"""
percent_middle = bezier.percentAtLength(
bezier.length() / 2)
return bezier.pointAtPercent(percent_middle)
def __get_bezier_middle_angle(self, bezier: QPainterPath):
"""
Get the angle in the middle of a bezier curve.
"""
percent_middle = bezier.percentAtLength(
bezier.length() / 2)
return -math.radians(
bezier.angleAtPercent(percent_middle))
class Environment(QWidget):
def __init__(self, parent=None):
super(QWidget, self).__init__(parent)
self.resize(SETTINGS["MAP_WIDTH"], SETTINGS["MAP_HEIGHT"])
self.map = Map(SETTINGS["MAP_WIDTH"], SETTINGS["MAP_HEIGHT"])
self.robot = Robot(SETTINGS["ROBOT_RADIUS"], self.map)
# SET UP ROBOT FOR POTENTIAL EXPERIMENT
if SETTINGS["EXPERIMENT_MODE"]:
self.robot.v = SETTINGS["ROBOT_START_V"]
self.robot.w = SETTINGS["ROBOT_START_W"]
self.experiment_length = SETTINGS["EXPERIMENT_LENGTH"]
self.filter = Kalman(self.robot)
self.trace = QPainterPath()
self.trace.moveTo(QPoint(self.robot.x, self.robot.y))
self.estimated_trace = QPainterPath()
self.estimated_trace.moveTo(QPoint(self.robot.x, self.robot.y))
self.covariance_ellipses = []
for x, y in itertools.product(range(200, SETTINGS["MAP_WIDTH"], 200),
range(200, SETTINGS["MAP_HEIGHT"], 200)):
self.map.add_beacon(Beacon(x, y))
# self.map.add_beacon(Beacon(SETTINGS["MAP_WIDTH"]*0.5, SETTINGS["MAP_HEIGHT"]*0.5))
self.trace_smooth_level = SETTINGS["TRACE_SMOOTHING"]
self.setAutoFillBackground(True)
p = self.palette()
p.setColor(self.backgroundRole(), SETTINGS["COLOR_BACKGROUND"])
self.setPalette(p)
# STATISTICS
self.steps = 0
self.total_squared_correction_error = 0
self.total_squared_prediction_error = 0
self.total_squared_correction_error_trajectory = []
self.total_squared_prediction_error_trajectory = []
def paintEvent(self, e):
qp = QPainter()
qp.begin(self)
qp.setRenderHint(QPainter.Antialiasing)
draw_beacon_indicators(qp, self.map, self.robot)
draw_beacons(qp, self.robot)
draw_trace(qp, self.trace)
draw_filter_trace(qp, self.estimated_trace)
for ell in self.covariance_ellipses:
draw_filter_covariance_ellipse(qp, *ell)
draw_robot(qp, self.robot)
qp.end()
def keyPressEvent(self, e):
# INCREMENT v
if e.key() == QtCore.Qt.Key_W:
self.robot.increment_v()
# INCREMENT w
if e.key() == QtCore.Qt.Key_D:
self.robot.increment_w()
# DECREMENT v
if e.key() == QtCore.Qt.Key_S:
self.robot.decrement_v()
# DECREMENT w
if e.key() == QtCore.Qt.Key_A:
self.robot.decrement_w()
# STOP
if e.key() == QtCore.Qt.Key_X:
self.robot.stop()
# STOP
if e.key() == QtCore.Qt.Key_Escape:
self.finish()
def animate(self):
self.robot.velocity_based_model()
prediction = self.filter.prediction()
correction = self.filter.correction()
estimation = correction
# print(f"{self.robot.x - estimation[0,0]}\n{self.robot.y - estimation[1,0]}\n")
# TRACES
if self.trace_smooth_level == SETTINGS["TRACE_SMOOTHING"]:
self.trace.lineTo(QPoint(self.robot.x, self.robot.y))
self.estimated_trace.lineTo(
QPoint(estimation[0, 0], estimation[1, 0]))
elif self.trace_smooth_level == 0:
self.trace_smooth_level = SETTINGS["TRACE_SMOOTHING"]
else:
self.trace_smooth_level -= 1
# COVARIANCE CIRCLES
if self.trace.length() - (
len(self.covariance_ellipses) *
SETTINGS["DISTANCE_BETWEEN_COV_CIRCLES"]
) >= SETTINGS["DISTANCE_BETWEEN_COV_CIRCLES"]:
self.covariance_ellipses.append(
(self.filter.mu[0], self.filter.mu[1], self.filter.sigma[0, 0],
self.filter.sigma[1, 1], self.filter.mu[2]))
self.update()
# STATISTICS
correction_squared_error = numpy.linalg.norm(
numpy.asarray(correction[0:2].transpose()).squeeze() -
numpy.array((self.robot.x, self.robot.y)))
prediction_squared_error = numpy.linalg.norm(
numpy.asarray(prediction[0:2].transpose()).squeeze() -
numpy.array((self.robot.x, self.robot.y)))
self.total_squared_correction_error += correction_squared_error
self.total_squared_prediction_error += prediction_squared_error
self.total_squared_correction_error_trajectory.append(
self.total_squared_correction_error)
self.total_squared_prediction_error_trajectory.append(
self.total_squared_prediction_error)
self.steps += 1
if SETTINGS[
"EXPERIMENT_MODE"] and self.experiment_length is not None and self.steps >= self.experiment_length:
self.finish()
def finish(self):
with open(f"experiments/data.txt", "w") as f:
f.write(
f"{self.total_squared_correction_error/self.steps}; {self.total_squared_prediction_error/self.steps}\n"
f"{self.total_squared_correction_error_trajectory}\n"
f"{self.total_squared_prediction_error_trajectory}\n"
f"{self.filter.K_trace}")
p = QPixmap(self.size())
self.render(p)
p.save("experiments/screenshot.jpg", "jpg")
exit()
def drawTextAlongCubic(self, lay, painter, filename, msg):
fs = lay["defaultFontSize"]
font = QFont('Right Chalk', fs)
defaultMessage = msg
if len(msg) == 0:
return
c1 = QPointF(lay["x1"], lay["y1"])
c2 = QPointF(lay["x2"], lay["y2"])
c3 = QPointF(lay["x3"], lay["y3"])
c4 = QPointF(lay["x4"], lay["y4"])
path = QPainterPath(c1)
path.cubicTo(c2, c3, c4)
# painter.drawPath(path)
pathLength = path.length()
textMetricLength = QFontMetrics(font).width(defaultMessage)
fsn = int(fs * pathLength / (textMetricLength) * .95)
if fsn > 70:
fsn = 70
font = QFont('Right Chalk', fsn)
textMetricLength = QFontMetrics(font).width(defaultMessage)
messageSpacing = []
defaultMessageM = []
sumMessageSpacing = 0.0
for i in range(len(defaultMessage)):
messageSpacing.append(QFontMetrics(font).width(defaultMessage[i]))
sumMessageSpacing += messageSpacing[i]
for i in range(len(defaultMessage)):
messageSpacing[i] = messageSpacing[i] / sumMessageSpacing
steps = 0
painter.setFont(font)
for i in range(len(defaultMessage)):
steps += messageSpacing[i] / 2
point = QPointF(path.pointAtPercent(steps))
angle = path.angleAtPercent(steps)
painter.save()
painter.translate(point)
painter.rotate(-angle)
x = -QFontMetrics(font).width(defaultMessage[i]) / 2
y = -QFontMetrics(font).height() / 2
w = QFontMetrics(font).width(defaultMessage[i])
h = QFontMetrics(font).height()
r = QRectF(x, y, w, h)
painter.setPen(QPen(Qt.white, 2))
painter.drawText(r, defaultMessage[i])
if i % 2 == 0:
painter.setPen(QPen(Qt.red, 2))
else:
painter.setPen(QPen(Qt.green, 2))
# painter.drawRect(r)
painter.restore()
steps += messageSpacing[i] / 2
class CandleStickItem(QGraphicsItem):
def __init__(self, model, parent=None):
QGraphicsItem.__init__(self, parent)
self.model: Model = model
self.max_x_range = self.model.current_x_range() # 2 <
self.plus_line_path = QPainterPath()
self.minus_line_path = QPainterPath()
self.plus_bar_path = QPainterPath()
self.plus_bar_path.setFillRule(Qt.WindingFill)
self.minus_bar_path = QPainterPath()
self.minus_bar_path.setFillRule(Qt.WindingFill)
self.cache = {
'plus_line_path': QPainterPath(),
'minus_line_path': QPainterPath(),
'plus_bar_path': QPainterPath(),
'minus_bar_path': QPainterPath()
}
self.thread = Thread()
self.make_path()
def paint(self,
painter: QtGui.QPainter,
option: QStyleOptionGraphicsItem,
widget: typing.Optional[QWidget] = ...):
# Set level of detail
# print(option.levelOfDetailFromTransform(painter.worldTransform()))
# draw plus line
painter.save()
pen = QPen()
pen.setColor(QColor("#496856"))
pen.setCosmetic(True)
painter.setPen(pen)
# painter.setRenderHint(painter.Antialiasing)
self.print_cache_path(painter.drawPath, 'plus_line_path', 1000)
# draw minus line
pen.setColor(QColor("#6F3541"))
painter.setPen(pen)
self.print_cache_path(painter.drawPath, 'minus_line_path', 1000)
# draw plus bar
self.print_cache_path(painter.fillPath, 'plus_bar_path', 1000,
QColor('#7BB888'))
# draw minus bar
self.print_cache_path(painter.fillPath, 'minus_bar_path', 1000,
QColor('#CC4E5C'))
painter.restore()
def print_cache_path(self, paint_fn, name, cache_len, *args):
r = self.model.current_x_range()
l = self.cache[name].length()
if 0 < r <= cache_len:
if l > 0:
paint_fn(self.cache[name], *args) # draw cache
else:
paint_fn(getattr(self, name), *args)
else:
if l == 0:
self.cache[name] = QPainterPath(getattr(self, name))
paint_fn(getattr(self, name), *args) # draw non-cache
def wheelEvent(self, event: 'QGraphicsSceneWheelEvent'):
super().wheelEvent(event)
self.make_path()
def keyPressEvent(self, event: QKeyEvent):
super().keyPressEvent(event)
self.make_path()
def make_path(self):
df = self.model.current_data()
plus_cond = 'close > open'
if self.plus_bar_path.length() == 0.0: # path.length == 0.0
# draw new initial path
print('draw new initial path')
plus_df = df[df.eval(plus_cond)]
minus_df = df[~df.eval(plus_cond)]
self.create_in_thread(self.draw_path, plus_df, self.plus_line_path)
self.create_in_thread(self.draw_path, minus_df,
self.minus_line_path)
self.create_in_thread(self.draw_rect, plus_df, self.plus_bar_path)
self.create_in_thread(self.draw_rect, minus_df,
self.minus_bar_path)
else:
if len(df) > self.max_x_range: # 3 > 2
self.max_x_range += self.model.next_x_range() # 2 += 500
next_df = self.model.next_data()
plus_df = next_df[next_df.eval(plus_cond)]
minus_df = next_df[~next_df.eval(plus_cond)]
print('draw next path')
self.create_in_thread(self.draw_path, plus_df,
self.plus_line_path)
self.create_in_thread(self.draw_path, minus_df,
self.minus_line_path)
self.create_in_thread(self.draw_rect, plus_df,
self.plus_bar_path)
self.create_in_thread(self.draw_rect, minus_df,
self.minus_bar_path)
def draw_rect(self, data, path):
draw_rect(data, path)
def draw_path(self, data, path):
draw_path(data, path)
def create_in_thread(self, fn, *args):
w = self.thread.make_worker(fn, *args)
w.sig.finished.connect(self.update)
self.thread.start(w)
def boundingRect(self):
return self.model.make_scene_rect()