def plot(dataset,
ax,
cax=None,
vmin=None,
vmax=None,
cmap="jet",
label=None,
mesh_filename="visdump_mesh.h5",
directory="."):
"""Draws a dataset on an ax."""
if vmin is None:
vmin = dataset.min()
if vmax is None:
vmax = dataset.max()
# get the mesh and collapse to 2D
etype, coords, conn = mesh.meshElemXYZ(filename=mesh_filename,
directory=directory)
if etype is not 'HEX':
raise RuntimeError("Only works for Hexs")
coords2 = np.array([[coords[i][0::2] for i in c if coords[i][1] == 0.0]
for c in conn])
assert coords2.shape[2] == 2
assert coords2.shape[1] == 4
polygons = matplotlib.collections.PolyCollection(coords2,
edgecolor='k',
cmap=cmap)
polygons.set_array(dataset)
polygons.set_clim(vmin, vmax)
ax.add_collection(polygons)
xmin = min(c[0] for c in coords.itervalues())
xmax = max(c[0] for c in coords.itervalues())
zmin = min(c[2] for c in coords.itervalues())
zmax = max(c[2] for c in coords.itervalues())
ax.set_xlim(xmin, xmax)
ax.set_ylim(zmin, zmax)
if cax is not None:
cb = plt.colorbar(polygons, cax=cax)
if label is not None:
cb.set_label(label)
return ((xmin, xmax), (zmin, zmax))
def plot(dataset,
ax,
cax=None,
vmin=None,
vmax=None,
cmap="jet",
label=None,
mesh_filename="visdump_mesh.h5",
directory=".",
y_coord=0.0,
linewidths=1):
"""Draws a dataset on an ax."""
import matplotlib.collections
from matplotlib import pyplot as plt
if vmin is None:
vmin = dataset.min()
if vmax is None:
vmax = dataset.max()
# get the mesh and collapse to 2D
etype, coords, conn = mesh.meshElemXYZ(filename=mesh_filename,
directory=directory)
if etype is not 'HEX':
raise RuntimeError("Only works for Hexs")
coords2 = np.array([[
coords[i][0::2] for i in c[1:] if abs(coords[i][1] - y_coord) < 1.e-8
] for c in conn])
try:
assert coords2.shape[2] == 2
assert coords2.shape[1] == 4
except AssertionError:
print coords2.shape
for c in conn:
if len(c) != 9:
print c
raise RuntimeError("what is a conn?")
coords3 = np.array([
coords[i][:] for i in c[1:]
if abs(coords[i][1] - y_coord) < 1.e-8
])
if coords3.shape[0] != 4:
print coords
raise RuntimeError("Unable to squash to 2D")
# reorder anti-clockwise
for i, c in enumerate(coords2):
centroid = c.mean(axis=0)
def angle(p1, p2):
a1 = np.arctan2((p1[1] - centroid[1]), (p1[0] - centroid[0]))
a2 = np.arctan2((p2[1] - centroid[1]), (p2[0] - centroid[0]))
if a1 < a2:
return -1
elif a2 < a1:
return 1
else:
return 0
c2 = np.array(sorted(c, angle))
coords2[i] = c2
polygons = matplotlib.collections.PolyCollection(coords2,
edgecolor='k',
cmap=cmap,
linewidths=linewidths)
polygons.set_array(dataset)
polygons.set_clim(vmin, vmax)
ax.add_collection(polygons)
xmin = min(c[0] for c in coords.itervalues())
xmax = max(c[0] for c in coords.itervalues())
zmin = min(c[2] for c in coords.itervalues())
zmax = max(c[2] for c in coords.itervalues())
ax.set_xlim(xmin, xmax)
ax.set_ylim(zmin, zmax)
if cax is not None:
cb = plt.colorbar(polygons, cax=cax)
if label is not None:
cb.set_label(label)
return ((xmin, xmax), (zmin, zmax))