1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201
|
import numpy as np
from ..MeshData import MeshData
from .GLMeshItem import GLMeshItem
from .GLLinePlotItem import GLLinePlotItem
from OpenGL import GL as ogl
__all__ = ['GLSurfacePlotItem']
class GLSurfacePlotItem(GLMeshItem):
"""
**Bases:** :class:`GLMeshItem <pyqtgraph.opengl.GLMeshItem>`
Displays a surface plot on a regular x,y grid with optional wireframe overlay.
"""
mesh_keys = ('x', 'y', 'z', 'colors')
grid_keys = ('showGrid', 'lineColor', 'lineWidth', 'lineAntialias')
def __init__(self, parentItem=None, **kwds):
"""
The x, y, z, colors, showGrid, lineColor, lineWidth and lineAntialias
arguments are passed to setData().
All other keyword arguments are passed to GLMeshItem.__init__().
"""
self._x = None
self._y = None
self._z = None
self._color = None
self._showGrid = False
self._lineColor = (0, 0, 0, 1)
self._lineWidth = 1.0
self._lineAntialias = False
self._vertexes = None
self._meshdata = MeshData()
# splitout GLSurfacePlotItem from kwds
surface_keys = self.mesh_keys + self.grid_keys
surface_kwds = {}
for arg in surface_keys:
if arg in kwds:
surface_kwds[arg] = kwds.pop(arg)
super().__init__(meshdata=self._meshdata, **kwds)
self.lineplot = GLLinePlotItem(parentItem=self, mode='lines', glOptions='translucent')
# in GLViewWidget.drawItemTree(), at the same depth value, child items
# come before the parent. make it such that our grid lines get drawn
# after the surface mesh.
self.lineplot.setDepthValue(self.depthValue() + 1)
self.setParentItem(parentItem)
self.setData(**surface_kwds)
def setData(self, **kwds):
"""
Update the data in this surface plot.
============== =====================================================================
**Arguments:**
x,y 1D arrays of values specifying the x,y positions of vertexes in the
grid. If these are omitted, then the values will be assumed to be
integers.
z 2D array of height values for each grid vertex.
colors (width, height, 4) array of vertex colors.
showGrid Show the grid lines.
lineColor Color of the grid lines.
lineWidth Width of the grid lines.
lineAntialias Enable antialiasing for the grid lines.
============== =====================================================================
All arguments are optional.
Note that if vertex positions are updated, the normal vectors for each triangle must
be recomputed. This is somewhat expensive if the surface was initialized with smooth=False
and very expensive if smooth=True. For faster performance, initialize with
computeNormals=False and use per-vertex colors or a normal-independent shader program.
"""
for arg in self.grid_keys:
if arg in kwds:
setattr(self, '_' + arg, kwds[arg])
x, y, z, colors = map(kwds.get, self.mesh_keys)
if x is not None:
if self._x is None or len(x) != len(self._x):
self._vertexes = None
self._x = x
if y is not None:
if self._y is None or len(y) != len(self._y):
self._vertexes = None
self._y = y
if z is not None:
if self._x is not None and z.shape[0] != len(self._x):
raise Exception('Z values must have shape (len(x), len(y))')
if self._y is not None and z.shape[1] != len(self._y):
raise Exception('Z values must have shape (len(x), len(y))')
self._z = z
if self._vertexes is not None and self._z.shape != self._vertexes.shape[:2]:
self._vertexes = None
if colors is not None:
self._colors = colors
self._meshdata.setVertexColors(colors)
if self._z is None:
return
updateMesh = False
newVertexes = False
## Generate vertex and face array
if self._vertexes is None:
newVertexes = True
self._vertexes = np.empty((self._z.shape[0], self._z.shape[1], 3), dtype=np.float32)
self.generateFaces()
self._meshdata.setFaces(self._faces)
updateMesh = True
## Copy x, y, z data into vertex array
if newVertexes or x is not None:
if x is None:
if self._x is None:
x = np.arange(self._z.shape[0])
else:
x = self._x
self._vertexes[:, :, 0] = x.reshape(len(x), 1)
updateMesh = True
if newVertexes or y is not None:
if y is None:
if self._y is None:
y = np.arange(self._z.shape[1])
else:
y = self._y
self._vertexes[:, :, 1] = y.reshape(1, len(y))
updateMesh = True
if newVertexes or z is not None:
self._vertexes[...,2] = self._z
updateMesh = True
## Update MeshData
if updateMesh:
self._meshdata.setVertexes(self._vertexes.reshape(self._vertexes.shape[0]*self._vertexes.shape[1], 3))
self.meshDataChanged()
# rebuild grid whenever mesh or parent changes
self._update_grid()
def paint(self):
if self._showGrid:
ogl.glEnable(ogl.GL_POLYGON_OFFSET_FILL)
ogl.glPolygonOffset(1.0, 1.0)
super().paint()
if self._showGrid:
ogl.glDisable(ogl.GL_POLYGON_OFFSET_FILL)
ogl.glPolygonOffset(0.0, 0.0)
def generateFaces(self):
cols = self._z.shape[1]-1
rows = self._z.shape[0]-1
faces = np.empty((cols*rows*2, 3), dtype=np.uint32)
rowtemplate1 = np.arange(cols).reshape(cols, 1) + np.array([[0, 1, cols+1]])
rowtemplate2 = np.arange(cols).reshape(cols, 1) + np.array([[cols+1, 1, cols+2]])
for row in range(rows):
start = row * cols * 2
faces[start:start+cols] = rowtemplate1 + row * (cols+1)
faces[start+cols:start+(cols*2)] = rowtemplate2 + row * (cols+1)
self._faces = faces
def _update_grid(self):
if not self._showGrid or self._z is None:
return
opts = {
'antialias': self._lineAntialias,
'color': self._lineColor,
'width': self._lineWidth,
}
z = self._z.astype(np.float32)
rows, cols = z.shape
x = (self._x if self._x is not None else np.arange(rows, dtype=z.dtype))
y = (self._y if self._y is not None else np.arange(cols, dtype=z.dtype))
xvals, yvals = np.meshgrid(x, y, indexing='ij') # shape (rows, cols)
verts_flat = np.column_stack((xvals.ravel(), yvals.ravel(), z.ravel()))
idx = np.arange(z.size, dtype=np.int32).reshape(rows, cols)
h = np.column_stack((idx[:, :-1].ravel(), idx[:, 1:].ravel()))
v = np.column_stack((idx[:-1, :].ravel(), idx[1:, :].ravel()))
edges = np.vstack((h, v))
pts = verts_flat[edges].reshape(-1, 3)
self.lineplot.setData(pos=pts, **opts)
|