Source code for compas_ags.viewers.viewer


from __future__ import print_function
from __future__ import absolute_import
from __future__ import division

from math import fabs

from compas_plotters.core import draw_xpoints_xy
from compas_plotters.core import draw_xlines_xy
from compas_plotters.core import draw_xarrows_xy
from compas_plotters.core import draw_xpolygons_xy

from compas.utilities import color_to_colordict
from compas.utilities import is_color_light

from compas_plotters.core import size_to_sizedict

import matplotlib.pyplot as plt


__author__ = ['Tom Van Mele']
__email__ = 'vanmelet@ethz.ch'


__all__ = ['Viewer']


[docs]class Viewer(object): """"""
[docs] def __init__(self, form, force, delay_setup=True, figsize=(16, 9)): self.form = form self.force = force self.fig = None self.ax1 = None self.ax2 = None self.default_facecolor = '#ffffff' self.default_edgecolor = '#000000' self.default_vertexcolor = '#ffffff' self.default_vertexsize = 0.1 self.default_edgewidth = 1.0 self.default_compressioncolor = '#0000ff' self.default_tensioncolor = '#ff0000' self.default_externalforcecolor = '#00ff00' self.default_externalforcewidth = 2.0 self.default_textcolor = '#000000' self.default_fontsize = 8 self.default_pointsize = 0.1 self.default_linewidth = 1.0 self.default_pointcolor = '#ffffff' self.default_linecolor = '#000000' self.default_linestyle = '-' self.figsize = figsize if not delay_setup: self.setup()
[docs] def setup(self): self.fig = plt.figure(figsize=self.figsize, tight_layout=True, dpi=96) self.ax1 = self.fig.add_subplot('121') self.ax2 = self.fig.add_subplot('122') self.ax1.set_aspect('equal') self.ax2.set_aspect('equal') self.ax1.set_xticks([]) self.ax1.set_yticks([]) self.ax1.set_xmargin(1.0) self.ax1.set_ymargin(1.0) self.ax1.set_xlim(-1.0, 11.0) self.ax1.set_ylim(-1.0, 11.0) self.ax2.set_xticks([]) self.ax2.set_yticks([]) self.ax2.set_xmargin(1.0) self.ax2.set_ymargin(1.0) self.ax2.set_xlim(-1.0, 11.0) self.ax2.set_ylim(-1.0, 11.0) self.fig.patch.set_facecolor('white') self.ax1.axis('off') self.ax2.axis('off')
[docs] def draw_form(self, vertices_on=True, edges_on=True, faces_on=False, forces_on=True, external_on=True, arrows_on=False, vertexcolor=None, edgecolor=None, facecolor=None, edgelabel=None, vertexlabel=None, facelabel=None, vertexsize=None, forcescale=1.0, lines=None, points=None): """""" # preprocess vertexlabel = vertexlabel or {} edgelabel = edgelabel or {} facelabel = facelabel or {} vertexsize = size_to_sizedict(vertexsize, self.form.vertices(), self.default_vertexsize) vertexcolor = color_to_colordict(vertexcolor, self.form.vertices(), self.default_vertexcolor) edgecolor = color_to_colordict(edgecolor, self.form.edges(), self.default_edgecolor) facecolor = color_to_colordict(facecolor, self.form.faces(), self.default_facecolor) # scale and position x = self.form.vertices_attribute('x') y = self.form.vertices_attribute('y') if lines: x += [line['start'][0] for line in lines] x += [line['end'][0] for line in lines] y += [line['start'][1] for line in lines] y += [line['end'][1] for line in lines] xmin, ymin = min(x), min(y) xmax, ymax = max(x), max(y) dx, dy = -xmin, -ymin scale = max(fabs(xmax - xmin) / 10.0, fabs(ymax - ymin) / 10.0) # vertices if vertices_on: _points = [] for key, attr in self.form.vertices(True): bgcolor = vertexcolor[key] _points.append({ 'pos': [(attr['x'] + dx) / scale, (attr['y'] + dy) / scale], 'radius': vertexsize[key], 'facecolor': vertexcolor[key], 'edgecolor': self.default_edgecolor, 'linewidth': self.default_edgewidth * 0.5, 'text': None if key not in vertexlabel else str(vertexlabel[key]), 'textcolor': '#000000' if is_color_light(bgcolor) else '#ffffff', 'fontsize': self.default_fontsize, }) draw_xpoints_xy(_points, self.ax1) # edges if edges_on: leaves = set(self.form.leaves()) _lines = [] _arrows = [] for (u, v), attr in self.form.edges(True): sp, ep = self.form.edge_coordinates(u, v, 'xy') sp = ((sp[0] + dx) / scale, (sp[1] + dy) / scale) ep = ((ep[0] + dx) / scale, (ep[1] + dy) / scale) if external_on: if u in leaves or v in leaves: text = None if (u, v) not in edgelabel else str(edgelabel[(u, v)]) _arrows.append({ 'start': sp, 'end': ep, 'width': self.default_externalforcewidth if not attr['is_ind'] else self.default_edgewidth * 3, 'color': self.default_externalforcecolor if not attr['is_ind'] else '#000000', 'text': text, 'fontsize': self.default_fontsize }) else: if forces_on: width = forcescale * fabs(attr['f']) color = self.default_tensioncolor if attr['f'] > 0 else self.default_compressioncolor text = None if (u, v) not in edgelabel else str(edgelabel[(u, v)]) _lines.append({ 'start': sp, 'end': ep, 'width': width, 'color': color, 'text': text, 'fontsize': self.default_fontsize }) _arrows.append({ 'start': sp, 'end': ep, 'width': self.default_edgewidth }) else: if forces_on: width = forcescale * fabs(attr['f']) color = self.default_tensioncolor if attr['f'] > 0 else self.default_compressioncolor text = None if (u, v) not in edgelabel else str(edgelabel[(u, v)]) _lines.append({ 'start': sp, 'end': ep, 'width': width, 'color': color, 'text': text, 'fontsize': self.default_fontsize }) _arrows.append({ 'start': sp, 'end': ep, 'width': self.default_edgewidth if not attr['is_ind'] else self.default_edgewidth * 3 }) if _arrows: if arrows_on: draw_xarrows_xy(_arrows, self.ax1) else: draw_xlines_xy(_arrows, self.ax1) if _lines: draw_xlines_xy(_lines, self.ax1, alpha=0.5) # faces if faces_on: _face_polygons = [] for fkey in self.form.faces(): vkeys = [vkey for vkey in self.form.face_vertices(fkey)] polygon_vertices = [self.form.vertex_coordinates(vkey, axes='xy') for vkey in vkeys] polygon_vertices = [[(x + dx) / scale, (y + dy) / scale] for (x, y) in polygon_vertices] # scale the polygon _face_polygons.append({ 'points': polygon_vertices, 'facecolor': '#e5e5e5', 'edgecolor': '#ffffff', 'edgewidth': 10.0, 'text': str(fkey) if fkey not in facelabel else str(facelabel[fkey]), 'fontsize': self.default_fontsize * 2, # TEMP! TO DIFFER FROM OTHER LABELS 'textcolor': self.default_textcolor, }) if _face_polygons: draw_xpolygons_xy(_face_polygons, self.ax1) # points if points: _points = [] for point in points: x, y, _ = point['pos'] _points.append({ 'pos': [(x + dx) / scale, (y + dy) / scale], 'text': point.get('text', ''), 'radius': point.get('size', self.default_pointsize), 'textcolor': point.get('textcolor', self.default_textcolor), 'facecolor': point.get('facecolor', self.default_pointcolor), 'edgecolor': point.get('edgecolor', self.default_linecolor), 'fontsize': self.default_fontsize }) draw_xpoints_xy(_points, self.ax1) # lines if lines: _lines = {} style = lines[0].get('style', self.default_linestyle) for line in lines: temp = line.get('style', self.default_linestyle) if temp == style: if temp not in _lines: _lines[temp] = [] else: _lines[temp] = [] style = temp _lines[temp].append({ 'start': [(line['start'][0] + dx) / scale, (line['start'][1] + dy) / scale], 'end': [(line['end'][0] + dx) / scale, (line['end'][1] + dy) / scale], 'width': line.get('width', self.default_linewidth), 'color': line.get('color', self.default_linecolor), 'text': line.get('text', ''), 'textcolor': line.get('textcolor', self.default_textcolor), 'fontsize': self.default_fontsize }) for style in _lines: draw_xlines_xy(_lines[style], self.ax1, linestyle=style)
[docs] def draw_force(self, vertices_on=True, edges_on=True, faces_on=False, forces_on=True, arrows_on=False, vertexcolor=None, edgecolor=None, facecolor=None, edgelabel=None, facelabel=None, vertexlabel=None, vertexsize=None, lines=None, points=None): """""" # preprocess vertexlabel = vertexlabel or {} edgelabel = edgelabel or {} facelabel = facelabel or {} vertexsize = size_to_sizedict(vertexsize, self.force.vertices(), self.default_vertexsize) vertexcolor = color_to_colordict(vertexcolor, self.force.vertices(), self.default_vertexcolor) edgecolor = color_to_colordict(edgecolor, self.force.edges(), self.default_edgecolor) # scale and position x = self.force.vertices_attribute('x') y = self.force.vertices_attribute('y') if lines: x += [line['start'][0] for line in lines] x += [line['end'][0] for line in lines] y += [line['start'][1] for line in lines] y += [line['end'][1] for line in lines] xmin, ymin = min(x), min(y) xmax, ymax = max(x), max(y) dx, dy = -xmin, -ymin scale = max(fabs(xmax - xmin) / 10.0, fabs(ymax - ymin) / 10.0) # vertices if vertices_on: _points = [] for key, attr in self.force.vertices(True): bgcolor = vertexcolor[key] _points.append({ 'pos': ((attr['x'] + dx) / scale, (attr['y'] + dy) / scale), 'radius': vertexsize[key], 'facecolor': bgcolor, 'edgecolor': self.default_edgecolor, 'linewidth': self.default_edgewidth * 0.5, 'text': None if key not in vertexlabel else str(vertexlabel[key]), 'textcolor': '#000000' if is_color_light(bgcolor) else '#ffffff', 'fontsize': self.default_fontsize }) draw_xpoints_xy(_points, self.ax2) # edges if edges_on: leaves = set(self.form.leaves()) _arrows = [] for (u, v), attr in self.force.edges(True): sp, ep = self.force.edge_coordinates(u, v, 'xy') sp = ((sp[0] + dx) / scale, (sp[1] + dy) / scale) ep = ((ep[0] + dx) / scale, (ep[1] + dy) / scale) form_u, form_v = self.force.dual_edge((u, v)) text = None if (u, v) not in edgelabel else str(edgelabel[(u, v)]) if form_u in leaves or form_v in leaves: _arrows.append({ 'start': sp, 'end': ep, 'color': self.default_externalforcecolor if not self.form.edge_attribute((form_u, form_v), 'is_ind') else '#000000', 'width': self.default_externalforcewidth if not self.form.edge_attribute((form_u, form_v), 'is_ind') else self.default_edgewidth * 3, 'text': text, 'fontsize': self.default_fontsize, }) else: _arrows.append({ 'start': sp, 'end': ep, 'color': self.default_edgecolor, 'width': self.default_edgewidth, 'text': text, 'fontsize': self.default_fontsize, }) if arrows_on: draw_xarrows_xy(_arrows, self.ax2) else: draw_xlines_xy(_arrows, self.ax2) # lines if lines: _lines = {} style = lines[0].get('style', self.default_linestyle) for line in lines: temp = line.get('style', self.default_linestyle) if temp == style: if temp not in _lines: _lines[temp] = [] else: _lines[temp] = [] style = temp _lines[temp].append({ 'start': [(line['start'][0] + dx) / scale, (line['start'][1] + dy) / scale], 'end': [(line['end'][0] + dx) / scale, (line['end'][1] + dy) / scale], 'width': line.get('width', self.default_linewidth), 'color': line.get('color', self.default_linecolor), 'text': line.get('text', ''), 'textcolor': line.get('textcolor', self.default_textcolor), 'fontsize': self.default_fontsize }) for style in _lines: draw_xlines_xy(_lines[style], self.ax2, linestyle=style) # faces if faces_on: _face_polygons = [] for fkey in self.force.faces(): vkeys = [vkey for vkey in self.force.face_vertices(fkey)] polygon_vertices = [self.force.vertex_coordinates(vkey, axes='xy') for vkey in vkeys] polygon_vertices = [[(x + dx) / scale, (y + dy) / scale] for (x, y) in polygon_vertices] # scale the polygon _face_polygons.append({ 'points': polygon_vertices, 'facecolor': '#e5e5e5', 'edgecolor': '#ffffff', 'edgewidth': 10.0, 'text': str(fkey) if fkey not in facelabel else str(facelabel[fkey]), 'fontsize': self.default_fontsize * 2, # TEMP! TO DIFFER FROM OTHER LABELS 'textcolor': self.default_textcolor, }) if _face_polygons: draw_xpolygons_xy(_face_polygons, self.ax2)
[docs] def show(self): plt.show()
[docs] def save(self, filepath, **kwargs): """Saves the plot to a file. Parameters ---------- filepath : str Full path of the file. Notes ----- For an overview of all configuration options, see [1]_. References ---------- .. [1] https://matplotlib.org/2.0.2/api/pyplot_api.html#matplotlib.pyplot.savefig """ plt.savefig(filepath, **kwargs)
# ============================================================================== # Main # ============================================================================== if __name__ == '__main__': pass