from __future__ import division from matplotlib import colors import matplotlib.cm as colormaps import matplotlib.pyplot as plt import numpy as np from gui import GUI from utils import xy_to_id, id_to_xy class StressGUI(GUI): """ Graphics user interface for drawing the outputs a topology optimization problem """ def __init__(self, nelx, nely, stress_calculator, nu, title=""): """Initialize plot and plot the initial design""" super(StressGUI, self).__init__(nelx, nely, title) self.stress_im = self.ax.imshow( np.swapaxes(np.zeros((nelx, nely, 4)), 0, 1), norm=colors.Normalize(vmin=0, vmax=1), cmap='jet') self.fig.colorbar(self.stress_im) self.stress_calculator = stress_calculator self.nu = nu self.myColorMap = colormaps.ScalarMappable( norm=colors.Normalize(vmin=0, vmax=1), cmap=colormaps.jet) def update(self, xPhys, u, title=None): """Plot to screen""" self.im.set_array(-xPhys.reshape((self.nelx, self.nely)).T) stress = self.stress_calculator.calculate_stress(xPhys, u, self.nu) # self.stress_calculator.calculate_fdiff_stress(xPhys, u, self.nu) self.myColorMap.set_norm(colors.Normalize(vmin=0, vmax=max(stress))) stress_rgba = self.myColorMap.to_rgba(stress) stress_rgba[:, :, 3] = xPhys.reshape(-1, 1) self.stress_im.set_array(np.swapaxes( stress_rgba.reshape((self.nelx, self.nely, 4)), 0, 1)) self.fig.canvas.draw() self.fig.canvas.flush_events() if title is not None: plt.title(title) else: plt.xlabel("Max stress = {:.2f}".format(max(stress)[0])) plt.pause(0.01)