From d6213db486fb7ed81d1d3c6f6d849260d1f6b91e Mon Sep 17 00:00:00 2001 From: rowanc1 Date: Fri, 7 Feb 2014 14:03:37 -0800 Subject: [PATCH 1/2] Initial commit for slicing in matplotlib --- SimPEG/Mesh/BaseMesh.py | 2 +- SimPEG/Mesh/TensorView.py | 86 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 87 insertions(+), 1 deletion(-) diff --git a/SimPEG/Mesh/BaseMesh.py b/SimPEG/Mesh/BaseMesh.py index 2dd82808..57d6f482 100644 --- a/SimPEG/Mesh/BaseMesh.py +++ b/SimPEG/Mesh/BaseMesh.py @@ -44,7 +44,7 @@ class BaseMesh(object): For example, you have a face variable, and you want the x component of it reshaped to a 3D matrix. - Mesh.r can fulfil your dreams:: + Mesh.r can fulfill your dreams:: mesh.r(V, 'F', 'Fx', 'M') | | | { How: 'M' or ['V'] for a matrix (ndgrid style) or a vector (n x dim) } diff --git a/SimPEG/Mesh/TensorView.py b/SimPEG/Mesh/TensorView.py index 7446aecd..8e44ee10 100644 --- a/SimPEG/Mesh/TensorView.py +++ b/SimPEG/Mesh/TensorView.py @@ -238,6 +238,92 @@ class TensorView(object): if showIt: plt.show() return ph + def plotSlice(self, v, vType='CC', + normal='Z', ind=None, grid=False, view='real', + ax=None, clim=None, showIt=False, + gridOpts={'color':'k'} + ): + viewOpts = ['real','imag','abs','vec'] + normalOpts = ['X', 'Y', 'Z'] + vTypeOpts = ['CC','F','E'] + + # Some user error checking + assert vType in vTypeOpts, "vType must be in ['%s']" % "','".join(vTypeOpts) + assert self.dim == 3, 'Must be a 3D mesh.' + assert view in viewOpts, "view must be in ['%s']" % "','".join(viewOpts) + assert normal in normalOpts, "normal must be in ['%s']" % "','".join(normalOpts) + assert type(grid) is bool, 'grid must be a boolean' + + szSliceDim = getattr(self, 'nC'+normal.lower()) #: Size of the sliced dimension + if ind is None: ind = int(szSliceDim/2) + assert type(ind) in [int, long], 'ind must be an integer' + + if ax is None: + fig = plt.figure(1) + fig.clf() + ax = plt.subplot(111) + else: + assert isinstance(ax, matplotlib.axes.Axes), "ax must be an matplotlib.axes.Axes" + fig = ax.figure + + # The slicing and plotting code!! + + def getIndSlice(v): + if normal == 'X': v = v[ind,:,:] + elif normal == 'Y': v = v[:,ind,:] + elif normal == 'Z': v = v[:,:,ind] + return v + + def doSlice(v): + if vType == 'CC': + return getIndSlice(self.r(v,'CC','CC','M')) + # Now just deal with 'F' and 'E' + aveOp = 'ave' + vType + ('2CCV' if view == 'vec' else '2CC') + v = getattr(self,aveOp)*v # average to cell centers (might be a vector) + if view == 'vec': + v = self.r(v.reshape((self.nC,3),order='F'),'CC','CC','M') + outSlice = [] + if 'X' not in normal: outSlice.append(getIndSlice(v[0])) + if 'Y' not in normal: outSlice.append(getIndSlice(v[1])) + if 'Z' not in normal: outSlice.append(getIndSlice(v[2])) + return outSlice + else: + return getIndSlice(self.r(v,'CC','CC','M')) + + h2d = [] + if 'X' not in normal: h2d.append(self.hx) + if 'Y' not in normal: h2d.append(self.hy) + if 'Z' not in normal: h2d.append(self.hz) + tM = Mesh.TensorMesh(h2d) #: Temp Mesh + + out = () + if view in ['real','imag','abs']: + v = getattr(np,view)(v) # e.g. np.real(v) + v = doSlice(v) + if clim is None: + clim = [v.min(),v.max()] + out += (ax.pcolormesh(tM.vectorNx, tM.vectorNy, v.T, vmin=clim[0], vmax=clim[1]),) + elif view in ['vec']: + U, V = doSlice(v) + if clim is None: + clim = [v.min(),v.max()] + out += (ax.pcolormesh(tM.vectorNx, tM.vectorNy, 0.5*(U+V).T, vmin=clim[0], vmax=clim[1]),) + out += (quiver( tM.gridCC[:,0], tM.gridCC[:,1], U, V),) + + if grid: + xXGrid = np.c_[tM.vectorNx,tM.vectorNx,np.nan*np.ones(tM.nNx)].flatten() + xYGrid = np.c_[tM.vectorNy[0]*np.ones(tM.nNx),tM.vectorNy[-1]*np.ones(tM.nNx),np.nan*np.ones(tM.nNx)].flatten() + yXGrid = np.c_[tM.vectorNx[0]*np.ones(tM.nNy),tM.vectorNx[-1]*np.ones(tM.nNy),np.nan*np.ones(tM.nNy)].flatten() + yYGrid = np.c_[tM.vectorNy,tM.vectorNy,np.nan*np.ones(tM.nNy)].flatten() + out += (ax.plot(np.r_[xXGrid,yXGrid],np.r_[xYGrid,yYGrid],**gridOpts)[0],) + + ax.set_xlabel('y' if normal == 'X' else 'x') + ax.set_ylabel('y' if normal == 'Z' else 'z') + ax.set_title('Slice %d' % ind) + + if showIt: plt.show() + return out + def plotGrid(self, nodes=False, faces=False, centers=False, edges=False, lines=True, showIt=False): """Plot the nodal, cell-centered and staggered grids for 1,2 and 3 dimensions. From 1c7c79344ba811f51118fdc8e13942060ebff237 Mon Sep 17 00:00:00 2001 From: rowanc1 Date: Fri, 7 Feb 2014 14:54:15 -0800 Subject: [PATCH 2/2] Option for normalized vector plot --- SimPEG/Mesh/TensorView.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/SimPEG/Mesh/TensorView.py b/SimPEG/Mesh/TensorView.py index 8e44ee10..b0ae3915 100644 --- a/SimPEG/Mesh/TensorView.py +++ b/SimPEG/Mesh/TensorView.py @@ -243,7 +243,7 @@ class TensorView(object): ax=None, clim=None, showIt=False, gridOpts={'color':'k'} ): - viewOpts = ['real','imag','abs','vec'] + viewOpts = ['real','imag','abs','vec','|vec|'] normalOpts = ['X', 'Y', 'Z'] vTypeOpts = ['CC','F','E'] @@ -258,6 +258,11 @@ class TensorView(object): if ind is None: ind = int(szSliceDim/2) assert type(ind) in [int, long], 'ind must be an integer' + normalizeVector = False + if view == '|vec|': + view = 'vec' + normalizeVector = True + if ax is None: fig = plt.figure(1) fig.clf() @@ -308,6 +313,9 @@ class TensorView(object): if clim is None: clim = [v.min(),v.max()] out += (ax.pcolormesh(tM.vectorNx, tM.vectorNy, 0.5*(U+V).T, vmin=clim[0], vmax=clim[1]),) + if normalizeVector: + length = np.sqrt(U**2+V**2) + U, V = U/length, V/length out += (quiver( tM.gridCC[:,0], tM.gridCC[:,1], U, V),) if grid: