Source code for itm.grid.plot
@author: klingshi
'''
import matplotlib.patches
import matplotlib.pyplot as plt
import matplotlib.collections as collections
import numpy
# Note: due to some issues with UAL and matplotlib, first import matplotlib
import base
class Plot2d():
[docs] PLOT_GRID = 2**0
PLOT_NODE_LABELS = 2**1
PLOT_EDGE_LABELS = 2**2
PLOT_CELL_LABELS = 2**3
def __init__(self, grid, figure=None):
assert isinstance(grid, base.Grid)
[docs] self.grid = grid
if figure is None:
self.fig = create_default_figure()
else:
self.fig = figure
self.default_cmap = matplotlib.cm.cool
def plot_grid_for_class(self, cls):
assert self.grid.class_exists(cls)
[docs]
sg = base.SubGrid(self.grid, cls)
self.plot_subgrid(sg)
def plot_subgrid(self, sgindex, mode=PLOT_GRID):
sg = self.grid.subgrid(sgindex)
[docs] assert (sg.ndim == 1) or (sg.ndim == 2), "Currently plotting only for 1d or 2d objects"
ax = self.fig.gca();
if sg.ndim == 1:
edges = list()
for obj1d in sg:
edges.append(obj1d)
elif sg.ndim == 2:
edges = set()
for obj2d in sg:
composing_objs = obj2d.composing_objs
edges |= set(composing_objs)
# 2d cell label
if ( mode & self.PLOT_CELL_LABELS ):
xy = numpy.array((0, 0));
for edge in composing_objs:
xy = xy + edge.nodes_coord.sum(axis=0)
xy = xy / ( len(composing_objs) * 2 )
# Avoid output of (x,) for a len=1 tuple
#if ( len(obj2d.ind) == 1 ):
# ax.annotate(str(obj2d.ind[0]), xy)
#else:
# ax.annotate(str(obj2d.ind), xy)
label = str(obj2d).replace(' ', '')
xy[0] = xy[0] - 0.0025
xy[1] = xy[1] - 0.001
ax.annotate(label, xy, size=16)
if ( mode & self.PLOT_GRID ):
lines = list()
for edge in edges:
lines.append(edge.nodes_coord)
#print "Plotting: ", edge.nodes_coord
lcol = collections.LineCollection(lines)
#lcol.set_linewidth(1)
ax.add_collection(lcol)
if ( mode & self.PLOT_EDGE_LABELS ):
for edge in edges:
nodes_coord = edge.nodes_coord
xy = nodes_coord.sum(axis=0) / nodes_coord.shape[0]
if ( len(edge.ind) == 1 ):
ax.annotate(str(edge.ind[0]), xy)
else:
ax.annotate(str(edge.ind), xy)
def plot_patch(self, data, cmap = None, norm = None):
grid = self.grid
[docs] assert (data.subgrid > 0) and (data.subgrid <= grid.nsubgrid)
sg = grid.subgrid(data.subgrid)
# FIXME: make more generic
assert sg.ndim == 2, "Plot2d.plot_patch: only for 2d objects."
assert len(data.scalar) == len(grid.subgrid(data.subgrid)), "Length of data vector must match number of objects in subgrid."
if cmap is None:
cmap = self.default_cmap
if norm is None:
norm = matplotlib.colors.Normalize()
norm.autoscale(data.scalar)
patches = list()
for i, obj2d in enumerate(sg):
coords = sort_edges(obj2d.composing_objs)[0]
patches.append( matplotlib.patches.Polygon(coords, closed = True, \
fc = cmap(norm(data.scalar[i])), ec = cmap(norm(data.scalar[i])) ) )
#self.fig.gca().add_patch( matplotlib.patches.Polygon(sort_edges(obj2d.composing_objs)[0], \
# fc = cmap(norm(data[i])), ec = cmap(norm(data[i])) ) )
#print "added all patches"
# Patch collection plotting disabled until matplotlib updated on the gateway
patchCol = collections.PatchCollection(patches, match_original = True)
self.fig.gca().add_collection(patchCol)
#print "Used PatchCollection"
def sort_edges(edges):
(c, cInd, closed) = sort_edges_index(edges)
return (c, closed)
def sort_edges_index(edges):
grid = edges[0].grid
# Build a dictionary of nodes and which edges they are connected to
edgeDict = dict()
for e in edges:
nodes = e.composing_objs
for n in nodes:
if n in edgeDict:
edgeDict[n].add(e)
assert len(edgeDict[n]) <= 2, \
"sort_edges: more than two edges connected to a node, cannot sort them"
else:
edgeDict[n] = set((e,))
c = numpy.zeros([len(edges)+1, grid.ndim])
cInd = [0] * (len(edges) + 1)
# Get a start node
cnode = edges[0].composing_objs[0]
c[0, :] = grid.node_coord(cnode)
cInd[0] = cnode.global_ind
# process as many edges as given
for i in xrange(len(edges)):
# sanity check on edge counts for the current node
if i == 0:
assert len(edgeDict[cnode]) in xrange(1,3)
else:
assert len(edgeDict[cnode]) == 1
# get the next edge
cedge = edgeDict[cnode].pop()
# get the other node of the edge as next node
cenodes = cedge.composing_objs
for n in cenodes:
if n != cnode:
nnode = n
break
# cnode is current node, nnode is next node
c[i+1, :] = grid.node_coord(nnode)
cInd[i+1] = nnode.global_ind
# remove current edge (which is now done) from next node set
edgeDict[nnode].remove(cedge)
# move to next node
cnode = nnode
# sanity check: end node set and start node set should be empty now
assert not edgeDict[cnode] and not edgeDict[edges[0].composing_objs[0]]
# check if closed
closed = ( cnode == edges[0].composing_objs[0] )
return (c, cInd, closed)
def sort_edges_2(edges):
# Build array with all node indices
#nodes = numpy.array('i', )
# Pick a an edge and one of its nodes as current edge and node
# Put coordinates of starting node to list
# For as many times as there are edges:
# -get other node of current edge
# -find an edge in the list that has the same node (excluding the current edge)
# -switch current edge and node to the new objects
grid = edges[0].grid
# Build a dictionary of nodes and which edges they are connected to
edgeDict = dict()
for e in edges:
nodes = e.composing_objs
for n in nodes:
if n in edgeDict:
edgeDict[n].add(e)
assert len(edgeDict[n]) <= 2, \
"sort_edges: more than two edges connected to a node, cannot sort them"
else:
edgeDict[n] = set((e,))
c = numpy.zeros([len(edges)+1, grid.ndim])
# Get a start node
cnode = edges[0].composing_objs[0]
c[0, :] = grid.node_coord(cnode)
# process as many edges as given
for i in xrange(len(edges)):
# sanity check on edge counts for the current node
if i == 0:
assert len(edgeDict[cnode]) in xrange(1,3)
else:
assert len(edgeDict[cnode]) == 1
# get the next edge
cedge = edgeDict[cnode].pop()
# get the other node of the edge as next node
cenodes = cedge.composing_objs
for n in cenodes:
if n != cnode:
nnode = n
break
# cnode is current node, nnode is next node
c[i+1, :] = grid.node_coord(nnode)
# remove current edge (which is now done) from next node set
edgeDict[nnode].remove(cedge)
# move to next node
cnode = nnode
# sanity check: end node set and start node set should be empty now
assert not edgeDict[cnode] and not edgeDict[edges[0].composing_objs[0]]
# check if closed
closed = ( cnode == edges[0].composing_objs[0] )
return (c, closed)
def create_default_figure():
fig = plt.figure()
plt.axis('equal')
plt.xlabel('R')
plt.ylabel('Z')
return fig