Source code for pyEMA.pole_picking

import os
import sys
import glob

import numpy as np
import matplotlib.pyplot as plt

try:
    import tkinter as tk
except:
    pass

from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg, NavigationToolbar2Tk
from matplotlib.figure import Figure

from . import stabilization

[docs]class SelectPoles:
[docs] def __init__(self, Model): """ Plot the measured Frequency Response Functions and computed poles. Picking the poles is done with pressing the SHIFT key + left mouse button. To unselect last pole: SHIFT + right mouse button. To unselect closest pole: SHIFT + middle mouse button. For more information check the HELP menu tab in the chart window. param model: object of pyEMA.Model """ self.Model = Model self.shift_is_held = False self.chart_type = 0 # 0 - stability chart, 1 - cluster diagram self.show_legend = 0 self.frf_plot_type = 'abs' self.Model.nat_freq = [] self.Model.nat_xi = [] self.Model.pole_ind = [] self.root = tk.Tk() self.root.title('Stability Chart') self.fig = Figure(figsize=(20, 8)) # Create axes self.ax2 = self.fig.add_subplot(111) self.ax1 = self.ax2.twinx() self.ax1.grid(True) # Tkinter menu menubar = tk.Menu(self.root) filemenu = tk.Menu(menubar, tearoff=0) filemenu.add_command(label='Save figure', command=self.save_this_figure) menubar.add_cascade(label='File', menu=filemenu) chartmenu = tk.Menu(menubar, tearoff=0) chartmenu.add_command(label='Stability chart', command=lambda: self.toggle_chart_type(0)) chartmenu.add_command(label='Cluster diagram', command=lambda: self.toggle_chart_type(1)) menubar.add_cascade(label="Chart type", menu=chartmenu) mifmenu = tk.Menu(menubar, tearoff=0) mifmenu.add_command(label='Plot mean abs', command=lambda: self.toggle_mif_frf('abs')) mifmenu.add_command(label='Plot all FRFs', command=lambda: self.toggle_mif_frf('all')) menubar.add_cascade(label="FRF plot type", menu=mifmenu) legendmenu = tk.Menu(menubar, tearoff=0) legendmenu.add_command(label='Show legend', command=lambda: self.toggle_legend(1)) legendmenu.add_command(label='Hide legend', command=lambda: self.toggle_legend(0)) menubar.add_cascade(label="Legend", menu=legendmenu) helpmenu = tk.Menu(menubar, tearoff=0) helpmenu.add_command(label='Help', command=self.show_help) menubar.add_cascade(label='Help', menu=helpmenu) self.root.config(menu=menubar) # Program execution self.plot_frf(initial=True) self.get_stability() self.plot_stability() # Integrate matplotlib figure canvas = FigureCanvasTkAgg(self.fig, self.root) canvas.get_tk_widget().pack(side='top', fill='both', expand=1) NavigationToolbar2Tk(canvas, self.root) # Connecting functions to event manager self.fig.canvas.mpl_connect('key_press_event', lambda x: self.on_key_press(x)) self.fig.canvas.mpl_connect('key_release_event', lambda x: self.on_key_release(x)) self.fig.canvas.mpl_connect('button_press_event', lambda x: self.on_click(x)) self.root.protocol("WM_DELETE_WINDOW", lambda: self.on_closing()) self.root.mainloop()
[docs] def plot_frf(self, initial=False): """Reconstruct and plot the Frequency Response Function. This is done on the fly. :param initial: if True, the frf is not computed, only the measured FRFs are shown. """ self.ax2.clear() if self.frf_plot_type == 'abs': self.ax2.semilogy(self.Model.freq, np.average( np.abs(self.Model.frf), axis=0), alpha=0.7, color='k') elif self.frf_plot_type == 'all': self.ax2.semilogy(self.Model.freq, np.abs(self.Model.frf.T), alpha=0.3, color='k') if not initial and len(self.Model.nat_freq) > 0: self.H, self.A = self.Model.get_constants(whose_poles='own', FRF_ind='all', least_squares_type='new') if self.frf_plot_type == 'abs': self.ax2.semilogy(self.Model.freq, np.average( np.abs(self.H), axis=0), color='r', lw=2) elif self.frf_plot_type == 'all': self.ax2.semilogy(self.Model.freq, np.abs(self.H.T), color='r', lw=1) else: x_position = (self.Model.lower + self.Model.upper) / 2 y_position = np.max(np.abs(self.Model.frf)) message = [ 'Select a pole: SHIFT + LEFT mouse button', 'Deselect a pole: SHIFT + RIGHT mouse button' ] self.ax2.text(x_position, y_position, '\n'.join(message), fontsize=12, verticalalignment='top', horizontalalignment='center', bbox=dict(facecolor='lightgreen', edgecolor='lightgreen')) self.ax1.set_xlim([self.Model.lower, self.Model.upper]) self.fig.canvas.draw()
[docs] def get_stability(self, fn_temp=0.001, xi_temp=0.05): """Get the stability matrix. :param fn_temp: Natural frequency stability criterion. :param xi_temp: Damping stability criterion. """ Nmax = self.Model.pol_order_high self.fn_temp, self.xi_temp, self.test_fn, self.test_xi = stabilization._stabilization( self.Model.all_poles, Nmax, err_fn=fn_temp, err_xi=xi_temp)
[docs] def plot_stability(self, update_ticks=False): if not update_ticks: self.ax1.clear() self.ax1.grid(True) # stable eigenfrequencies, unstable damping ratios a = np.argwhere((self.test_fn > 0) & ((self.test_xi == 0) | (self.xi_temp <= 0))) # stable eigenfrequencies, stable damping ratios b = np.argwhere((self.test_fn > 0) & ((self.test_xi > 0) & (self.xi_temp > 0))) # unstable eigenfrequencies, unstable damping ratios c = np.argwhere((self.test_fn == 0) & ((self.test_xi == 0) | (self.xi_temp <= 0))) # unstable eigenfrequencies, stable damping ratios d = np.argwhere((self.test_fn == 0) & ((self.test_xi > 0) & (self.xi_temp > 0))) p1 = self.ax1.plot(self.fn_temp[a[:, 0], a[:, 1]], 1+a[:, 1], 'bx', markersize=4, label="stable frequency, unstable damping") p2 = self.ax1.plot(self.fn_temp[b[:, 0], b[:, 1]], 1+b[:, 1], 'gx', markersize=7, label="stable frequency, stable damping") p3 = self.ax1.plot(self.fn_temp[c[:, 0], c[:, 1]], 1+c[:, 1], 'r.', markersize=4, label="unstable frequency, unstable damping") p4 = self.ax1.plot(self.fn_temp[d[:, 0], d[:, 1]], 1+d[:, 1], 'r*', markersize=4, label="unstable frequency, stable damping") self.line, = self.ax1.plot(self.Model.nat_freq, np.repeat( self.Model.pol_order_high, len(self.Model.nat_freq)), 'kv', markersize=8) self.selected, = self.ax1.plot([self.Model.pole_freq[p[0]][p[1]] for p in self.Model.pole_ind], [p[0] for p in self.Model.pole_ind], 'ko') if self.show_legend: self.pole_legend = self.ax1.legend(loc='upper center', ncol=2, frameon=True) self.ax1.set_title('Stability chart') self.ax1.set_ylabel('Polynomial order') plt.tight_layout() else: self.line.set_xdata(np.asarray(self.Model.nat_freq)) # update data self.line.set_ydata(np.repeat(self.Model.pol_order_high*1.04, len(self.Model.nat_freq))) self.selected.set_xdata([self.Model.pole_freq[p[0]][p[1]] for p in self.Model.pole_ind]) # update data self.selected.set_ydata([p[0] for p in self.Model.pole_ind]) plt.tight_layout() self.ax1.set_ylim([0, self.Model.pol_order_high+5]) self.fig.canvas.draw()
[docs] def plot_cluster(self, update_ticks=False): """Plot clusters - damping with respect to frequency. :param update_ticks: if True, only ticks are updated, not the whole plot. """ b1 = np.argwhere(((self.test_fn > 0) & ((self.test_xi > 0) & (self.xi_temp > 0))) & ((self.fn_temp > self.Model.lower) & (self.fn_temp < self.Model.upper))) if not update_ticks: self.ax1.clear() self.ax1.grid(True) # stable eigenfrequencies, unstable damping ratios a = np.argwhere((self.test_fn > 0) & ((self.test_xi == 0) | (self.xi_temp <= 0))) # stable eigenfrequencies, stable damping ratios b = np.argwhere((self.test_fn > 0) & ((self.test_xi > 0) & (self.xi_temp > 0))) # unstable eigenfrequencies, unstable damping ratios c = np.argwhere((self.test_fn == 0) & ((self.test_xi == 0) | (self.xi_temp <= 0))) # unstable eigenfrequencies, stable damping ratios d = np.argwhere((self.test_fn == 0) & ((self.test_xi > 0) & (self.xi_temp > 0))) p1 = self.ax1.plot(self.fn_temp[a[:, 0], a[:, 1]], self.xi_temp[a[:, 0], a[:, 1]], 'bx', markersize=4, label="stable frequency, unstable damping") p2 = self.ax1.plot(self.fn_temp[b[:, 0], b[:, 1]], self.xi_temp[b[:, 0], b[:, 1]], 'gx', markersize=7, label="stable frequency, stable damping") p3 = self.ax1.plot(self.fn_temp[c[:, 0], c[:, 1]], self.xi_temp[c[:, 0], c[:, 1]], 'r.', markersize=4, label="unstable frequency, unstable damping") p4 = self.ax1.plot(self.fn_temp[d[:, 0], d[:, 1]], self.xi_temp[d[:, 0], d[:, 1]], 'r*', markersize=4, label="unstable frequency, stable damping") self.line, = self.ax1.plot(self.Model.nat_freq, np.repeat( 1.05*np.max(self.xi_temp[b1[:, 0], b1[:, 1]]), len(self.Model.nat_freq)), 'kv', markersize=8) self.selected, = self.ax1.plot([self.Model.pole_freq[p[0]][p[1]] for p in self.Model.pole_ind], [self.Model.pole_xi[p[0]][p[1]] for p in self.Model.pole_ind], 'ko') if self.show_legend: self.pole_legend = self.ax1.legend(loc='upper center', ncol=2, frameon=True) self.ax1.set_title('Cluster diagram') self.ax1.set_ylabel('Damping ratio') plt.tight_layout() else: self.line.set_xdata(np.asarray(self.Model.nat_freq)) # update data self.line.set_ydata(np.repeat(1.05*np.max(self.xi_temp[b1[:, 0], b1[:, 1]]), len(self.Model.nat_freq))) self.selected.set_xdata([self.Model.pole_freq[p[0]][p[1]] for p in self.Model.pole_ind]) # update data self.selected.set_ydata([self.Model.pole_xi[p[0]][p[1]] for p in self.Model.pole_ind]) to_lim = self.xi_temp[b1[:, 0], b1[:, 1]] up_lim = min(np.max(to_lim), np.mean(to_lim)+2*np.std(to_lim)) self.ax1.set_ylim([np.mean(to_lim)-np.std(to_lim), up_lim]) self.fig.canvas.draw()
[docs] def get_closest_poles_stability(self): """ On-the-fly selection of the closest poles. """ y_ind = int(np.argmin(np.abs(np.arange(0, len(self.Model.pole_freq) )-self.y_data_pole))) # Find closest pole order # Find cloeset frequency sel = np.argmin(np.abs(self.Model.pole_freq[y_ind] - self.x_data_pole)) self.Model.pole_ind.append([y_ind, sel]) self.Model.nat_freq.append(self.Model.pole_freq[y_ind][sel]) self.Model.nat_xi.append(self.Model.pole_xi[y_ind][sel]) self.sort_selected_poles()
[docs] def get_closest_poles_cluster(self): """ On-the-fly selection of the closest poles. """ min_ = [] for y_ind in range(len(self.Model.pole_freq)): for sel in range(len(self.Model.pole_freq[y_ind])): min_.append([y_ind, sel, np.abs(self.Model.pole_freq[y_ind][sel]-self.x_data_pole), np.abs(self.Model.pole_xi[y_ind][sel]-self.y_data_pole[0])]) min_a = np.asarray(min_) # select the poles that have a frequency within 2% of the observed range mask = min_a[:, 2] < (self.Model.upper - self.Model.lower) * 0.02 min_ind_x = np.argmin(min_a[mask][:, 3]) y_ind = int(min_a[mask][min_ind_x, 0]) sel = int(min_a[mask][min_ind_x, 1]) self.Model.pole_ind.append([y_ind, sel]) self.Model.nat_freq.append(self.Model.pole_freq[y_ind][sel]) self.Model.nat_xi.append(self.Model.pole_xi[y_ind][sel]) self.sort_selected_poles()
[docs] def on_click(self, event): # on button 1 press (left mouse button) + SHIFT is held if event.button == 1 and self.shift_is_held: self.y_data_pole = [event.ydata] self.x_data_pole = event.xdata if self.chart_type == 0: self.get_closest_poles_stability() elif self.chart_type == 1: self.get_closest_poles_cluster() self.plot_frf() # On button 3 press (left mouse button) elif event.button == 3 and self.shift_is_held: try: del self.Model.nat_freq[-1] # delete last point del self.Model.nat_xi[-1] del self.Model.pole_ind[-1] self.plot_frf() except: pass elif event.button == 2 and self.shift_is_held: i = np.argmin(np.abs(self.Model.nat_freq - event.xdata)) try: del self.Model.nat_freq[i] del self.Model.nat_xi[i] del self.Model.pole_ind[i] self.plot_frf() except: pass if self.shift_is_held: if self.chart_type == 0: self.plot_stability(update_ticks=True) elif self.chart_type == 1: self.plot_cluster(update_ticks=True)
[docs] def on_key_press(self, event): """Function triggered on key press (SHIFT).""" if event.key == 'shift': self.shift_is_held = True
[docs] def on_key_release(self, event): """Function triggered on key release (SHIFT).""" if event.key == 'shift': self.shift_is_held = False
[docs] def on_closing(self): self.root.destroy()
[docs] def toggle_legend(self, x): if x: self.show_legend = 1 else: self.show_legend = 0 if self.chart_type == 0: self.plot_stability() elif self.chart_type == 1: self.plot_cluster()
[docs] def toggle_chart_type(self, x): if x == 0: self.chart_type = 0 self.plot_stability() elif x == 1: self.chart_type = 1 self.plot_cluster()
[docs] def toggle_mif_frf(self, x): self.frf_plot_type = x self.plot_frf()
[docs] def sort_selected_poles(self): _ = np.argsort(self.Model.nat_freq) self.Model.pole_ind = list(np.array(self.Model.pole_ind)[_]) self.Model.nat_freq = list(np.array(self.Model.nat_freq)[_]) self.Model.nat_xi = list(np.array(self.Model.nat_xi)[_])
[docs] def show_help(self): lines = [ 'Pole selection help', ' ', '- Select a pole: SHIFT + LEFT mouse button', '- Deselect a pole: SHIFT + RIGHT mouse button', '- Deselect the closest pole (frequency wise): SHIFT + MIDDLE mouse button', ' ', '- Two different types of charts are currently avaliable:', ' 1. stability chart, where pole frequencies are plotted against polynomial order', ' 2. cluster diagram, where pole frequencies are plotted against damping ratios' ] tk.messagebox.showinfo('Picking poles', '\n'.join(lines))
[docs] def save_this_figure(self): filename = 'pole_chart_' directory = 'pole_figures' if not os.path.exists(directory): os.mkdir(directory) files = glob.glob(directory + '/*.png') i = 1 while True: f = os.path.join(directory, filename + f'{i:0>3}.png') if f not in files: break i += 1 self.fig.savefig(f)