Permalink
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
100 lines (79 sloc) 3.27 KB
from .mesh import MeshVisual
from ..geometry import create_grid_mesh, MeshData
class GridMeshVisual(MeshVisual):
"""Displays a mesh in a Cartesian grid about x,y,z coordinates.
This makes it simple to generate a mesh from e.g. the output
of numpy.meshgrid.
All arguments are optional, though they can be changed
individually later with the set_data method.
Parameters
----------
xs : ndarray
A 2d array of x coordinates for the vertices of the mesh. Must
have the same dimensions as ys and zs.
ys : ndarray
A 2d array of y coordinates for the vertices of the mesh. Must
have the same dimensions as xs and zs.
zs : ndarray
A 2d array of z coordinates for the vertices of the mesh. Must
have the same dimensions as xs and ys.
colors : ndarray | None
The colors of the points of the mesh. Should be either a
(width, height, 4) array of rgba colors at each grid point or
a (width, height, 3) array of rgb colors at each grid point.
Defaults to None, in which case the default color of a
MeshVisual is used.
shading : str | None
Same as for the `MeshVisual` class. Defaults to 'smooth'.
**kwargs :
Other arguments are passed directly to MeshVisual.
"""
def __init__(self, xs, ys, zs, colors=None, shading='smooth',
**kwargs):
if xs is None or ys is None or zs is None:
raise ValueError('All of xs, ys and zs must be initialised '
'with arrays.')
self._xs = None
self._ys = None
self._zs = None
self.__vertices = None
self.__meshdata = MeshData()
MeshVisual.__init__(self, shading=shading, **kwargs)
self.set_data(xs, ys, zs, colors)
def set_data(self, xs=None, ys=None, zs=None, colors=None):
'''Update the mesh data.
Parameters
----------
xs : ndarray | None
A 2d array of x coordinates for the vertices of the mesh.
ys : ndarray | None
A 2d array of y coordinates for the vertices of the mesh.
zs : ndarray | None
A 2d array of z coordinates for the vertices of the mesh.
colors : ndarray | None
The color at each point of the mesh. Must have shape
(width, height, 4) or (width, height, 3) for rgba or rgb
color definitions respectively.
'''
if xs is None:
xs = self._xs
self.__vertices = None
if ys is None:
ys = self._ys
self.__vertices = None
if zs is None:
zs = self._zs
self.__vertices = None
if self.__vertices is None:
vertices, indices = create_grid_mesh(xs, ys, zs)
self._xs = xs
self._ys = ys
self._zs = zs
if self.__vertices is None:
vertices, indices = create_grid_mesh(self._xs, self._ys, self._zs)
self.__meshdata.set_vertices(vertices)
self.__meshdata.set_faces(indices)
if colors is not None:
self.__meshdata.set_vertex_colors(colors.reshape(
colors.shape[0] * colors.shape[1], colors.shape[2]))
MeshVisual.set_data(self, meshdata=self.__meshdata)