class MainWindow(QtWidgets.QMainWindow): stop_plot_signal = QtCore.pyqtSignal() def __init__(self, settings): super(MainWindow, self).__init__() self.central_widget = QtWidgets.QWidget(self) self.main_layout = QtWidgets.QVBoxLayout(self.central_widget) self.image_label = QtWidgets.QLabel(self.central_widget) self.status_bar = QtWidgets.QStatusBar(self) self.image_slider = QtWidgets.QSlider(self.central_widget, orientation=QtCore.Qt.Horizontal) self.bottom_layout = QtWidgets.QHBoxLayout() self.move_year_left_button = QtWidgets.QPushButton() self.move_year_right_button = QtWidgets.QPushButton() self.button_layout = QtWidgets.QVBoxLayout() self.plot_button = QtWidgets.QPushButton() self.stop_button = QtWidgets.QPushButton(enabled=False) self.sdate_animate_layout = QtWidgets.QVBoxLayout() self.start_date_layout = QtWidgets.QHBoxLayout() self.start_year = QtWidgets.QSpinBox() self.start_month = QtWidgets.QComboBox() self.animate_button = QtWidgets.QPushButton(enabled=False) self.edate_pref_layout = QtWidgets.QVBoxLayout() self.end_date_layout = QtWidgets.QHBoxLayout() self.end_year = QtWidgets.QSpinBox() self.end_month = QtWidgets.QComboBox() self.preferences_button = QtWidgets.QPushButton() self.animate_timer = QtCore.QTimer() self.image_count = 0 self.settings = settings def setup_ui(self): self.save_default_settings() self.setWindowFlag(QtCore.Qt.MSWindowsFixedSizeDialogHint) self.setCentralWidget(self.central_widget) self.setStatusBar(self.status_bar) self.status_bar.setSizeGripEnabled(False) self.retranslate_ui() self.set_ranges_values() self.connect_signals() self.set_shortcuts() spacer = QtWidgets.QSpacerItem(1, 1, QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Expanding) self.start_date_layout.addWidget(self.start_month) self.start_date_layout.addWidget(self.start_year) self.sdate_animate_layout.addLayout(self.start_date_layout) self.sdate_animate_layout.addWidget(self.animate_button) self.button_layout.addWidget(self.plot_button, alignment=QtCore.Qt.AlignCenter) self.button_layout.addWidget(self.stop_button, alignment=QtCore.Qt.AlignCenter) self.end_date_layout.addWidget(self.end_month) self.end_date_layout.addWidget(self.end_year) self.edate_pref_layout.addLayout(self.end_date_layout) self.edate_pref_layout.addWidget(self.preferences_button) self.bottom_layout.addLayout(self.sdate_animate_layout) self.bottom_layout.addSpacerItem(spacer) self.bottom_layout.addWidget(self.move_year_left_button) self.bottom_layout.addLayout(self.button_layout) self.bottom_layout.addWidget(self.move_year_right_button) self.bottom_layout.addSpacerItem(spacer) self.bottom_layout.addLayout(self.edate_pref_layout) self.main_layout.addWidget(self.image_label, alignment=QtCore.Qt.AlignCenter) self.main_layout.addSpacerItem(spacer) self.main_layout.addWidget(self.image_slider) self.main_layout.addLayout(self.bottom_layout) self.show() self.preferences_button.pressed.connect(self.show_options) self.set_sizes() def set_shortcuts(self): year_move_right = QtWidgets.QShortcut( QtGui.QKeySequence(QtCore.Qt.CTRL + QtCore.Qt.Key_Right), self) year_move_right.activated.connect(lambda: self.move_slider(self.year_step)) year_move_left = QtWidgets.QShortcut( QtGui.QKeySequence(QtCore.Qt.CTRL + QtCore.Qt.Key_Left), self) year_move_left.activated.connect(lambda: self.move_slider(-self.year_step)) month_move_right = QtWidgets.QShortcut(QtGui.QKeySequence(QtCore.Qt.Key_Right), self) month_move_right.activated.connect(lambda: self.move_slider(1)) month_move_left = QtWidgets.QShortcut(QtGui.QKeySequence(QtCore.Qt.Key_Left), self) month_move_left.activated.connect(lambda: self.move_slider(-1)) def set_sizes(self): self.setFixedSize(850, 650) self.image_label.setFixedSize(QtCore.QSize(796, 552)) # set year skip buttons to be square and 5 pixels wider than text in them font = QtGui.QFont() self.move_year_left_button.setFixedWidth( QtGui.QFontMetrics(font).boundingRect(self.move_year_left_button.text()).width() + 5) self.move_year_right_button.setFixedWidth( QtGui.QFontMetrics(font).boundingRect(self.move_year_right_button.text()).width() + 5) self.move_year_left_button.setFixedHeight(self.move_year_left_button.width()) self.move_year_right_button.setFixedHeight(self.move_year_right_button.width()) def set_ranges_values(self): months = ("January", "February", "March", "April", "May", "June", "July", "August", "September", "October", "November", "December") self.start_month.addItems(months) self.end_month.addItems(months) self.image_slider.setRange(0, 0) self.image_slider.setValue(0) self.start_year.setRange(1850, 2010) self.end_year.setRange(1980, 2019) self.start_year.setValue(1980) self.end_year.setValue(2010) def connect_signals(self): self.image_slider.valueChanged.connect(self.change_image) # ensure only valid dates can be entered self.start_month.currentIndexChanged.connect(self.date_changed) self.end_year.valueChanged.connect(self.date_changed) self.start_year.valueChanged.connect(self.date_changed) self.animate_button.pressed.connect(self.animate) self.plot_button.pressed.connect(self.plot) self.stop_button.pressed.connect(self.quit_current_tasks) self.move_year_left_button.pressed.connect(lambda: self.move_slider(-self.year_step)) self.move_year_right_button.pressed.connect(lambda: self.move_slider(self.year_step)) def retranslate_ui(self): self.setWindowTitle("World Temperature Anomaly Map") self.plot_button.setText("Plot") self.stop_button.setText("Stop") self.animate_button.setText("Play") self.preferences_button.setText("Preferences") self.move_year_right_button.setText("-->") self.move_year_left_button.setText("<--") self.move_year_left_button.setToolTip("Skip year") self.move_year_right_button.setToolTip("Skip year") def show_options(self): self.preferences_button.setEnabled(False) w = SettingsPop(self.settings, self) w.setup_ui() w.settings_signal.connect(self.refresh_settings) w.close_signal.connect(lambda: self.preferences_button.setEnabled(True)) def save_default_settings(self): if not self.settings.value("Plot step"): self.settings.setValue("Playback FPS", 5) self.settings.setValue("Plot step", 1) self.settings.setValue("Color map", "seismic") self.settings.sync() self.plot_step = self.settings.value("Plot step", type=int) self.play_fps = self.settings.value("Playback FPS", type=int) self.color_map = self.settings.value("Color map") self.year_step = max(int(12 / self.plot_step), 1) def refresh_settings(self, values): self.play_fps = int(values[0]) self.plot_step = int(values[1]) self.color_map = values[2] self.year_step = max(int(12 / self.plot_step), 1) def set_status(self, message): self.status_bar.showMessage(message) def plot(self): self.image_count = 0 QtGui.QPixmapCache.clear() # clear qt image cache self.stop_button.setEnabled(True) self.plot_button.setEnabled(False) self.animate_button.setEnabled(False) # send dates in decimal format to worker start_date = self.start_year.value() + (1 + self.start_month.currentIndex() * 2) / 24 end_date = self.end_year.value() + (1 + self.end_month.currentIndex() * 2) / 24 self.worker = Plotter(start_date, end_date, self.plot_step, self.color_map, self) self.worker.image_increment_signal.connect(self.add_image) self.worker.finished.connect(self.del_worker) self.worker.status_signal.connect(self.set_status) self.worker.start() def del_worker(self): self.worker.quit() self.stop_button.setEnabled(False) self.plot_button.setEnabled(True) self.animate_button.setEnabled(True) def add_image(self): self.image_slider.setMaximum(self.image_count) # move slider to max value if it was at max before if self.image_slider.value() == self.image_slider.maximum() - 1: self.move_slider(1) self.image_count += 1 def move_slider(self, amount: int): """ move image_slider by value""" self.image_slider.setValue(self.image_slider.value() + amount) def change_image(self, index): pixmap = QtGui.QPixmap(f"{Plotter.PLOTS_DIR}plot{index}") self.image_label.setPixmap(pixmap) def date_changed(self): """Ensure only valid dates can be entered if start and end years match, block out months above chosen start month in end months if year is 2019 allow only January-May range""" for item_index in range(0, 12): self.start_month.model().item(item_index).setEnabled(True) self.end_month.model().item(item_index).setEnabled(True) if self.start_year.value() == self.end_year.value(): for item_index in range(0, self.start_month.currentIndex()): self.end_month.model().item(item_index).setEnabled(False) if self.end_month.currentIndex() < self.start_month.currentIndex(): self.end_month.setCurrentIndex(self.start_month.currentIndex()) if self.start_year.value() == 2019: for item_index in range(5, 12): self.start_month.model().item(item_index).setEnabled(False) if self.start_month.currentIndex() > 4: self.start_month.setCurrentIndex(4) if self.end_year.value() == 2019: for item_index in range(5, 12): self.end_month.model().item(item_index).setEnabled(False) if self.end_month.currentIndex() > 4: self.end_month.setCurrentIndex(4) self.start_year.setRange(1850, self.end_year.value()) self.end_year.setRange(self.start_year.value(), 2019) def animate(self): self.image_slider.setValue(1) self.stop_button.setEnabled(True) self.animate_button.setEnabled(False) self.animate_timer.timeout.connect(self.animation) self.animate_timer.start(int(1000 / self.play_fps)) def animation(self): self.move_slider(1) if self.image_slider.value() == self.image_slider.maximum(): self.stop_animation() self.stop_button.setEnabled(False) def stop_animation(self): self.animate_timer.stop() try: self.animate_timer.timeout.disconnect() except TypeError: pass self.animate_button.setEnabled(True) def quit_current_tasks(self): self.stop_plot_signal.emit() self.stop_animation() self.stop_button.setEnabled(False) def closeEvent(self, *args, **kwargs): super(QtWidgets.QMainWindow, self).closeEvent(*args, **kwargs) try: self.worker.clear_plots() except AttributeError: pass
from plot import Plotter import matplotlib.pyplot as plt import numpy as np import time from pyqtgraph.Qt import QtCore, QtGui if __name__ == "__main__": initial_state = np.array([0., 0., 0., 0., 1e-2, 0., 0., 0., 0., 0.]) p = Plotter(0, 0, 0, -0.1, 0.1, np.array([0, 0, -9.8])) r = Rocket(*initial_state) def update(): r.propagate_state() p.update(r.x, r.y, r.z, r.theta1, r.theta2, r.c) timer = QtCore.QTimer() timer.timeout.connect(update) timer.start(10) states = np.ones((100, 10)) p.start() # labels = ['x', 'y', 'z', 'theta1', 'theta2', 'xdot', 'ydot', 'zdot', # 'theta1dot', 'theta2dot'] # colors = ['r', 'g', 'b', 'c', 'k', 'r', 'g', 'b', 'c', 'k'] # alphas = [1, 1, 1, 1, 1, 0.6, 0.6, 0.6, 0.6, 0.6] # x = np.linspace(0, 100, 100) # for label, data, color, alpha in zip(labels, states.T, colors, alphas): # plt.plot(x, data, label=label, c=color, alpha=alpha) # plt.legend() # plt.show()