3D plots working!

SimPEG/TensorView.py

@@ -13,10 +13,13 @@ class TensorView(object):
     def __init__(self):
         pass
 
-    def plotImage(self, I, imageType='CC', figNum=1,ax=None):
+    def plotImage(self, I, imageType='CC', figNum=1,ax=None,direction='z',numbering=True):
 
         assert type(I) == np.ndarray, "I must be a numpy array"
+        assert type(numbering) == bool, "numbering must be a bool"
         assert imageType in ["CC", "N"], "imageType must be 'CC' or 'N'"
+        assert direction in ["x", "y","z"], "direction must be either x,y, or z"
+
 
         if imageType == 'CC':
             assert I.size == self.nC, "Incorrect dimensions for CC."
@@ -30,7 +33,7 @@ class TensorView(object):
         else:
             assert isinstance(ax,matplotlib.axes.Axes), "ax must be an Axes!"
             fig = ax.figure
-       
+
         if self.dim == 1:
             if imageType == 'CC':
                 ph = ax.plot(self.vectorCCx, I, '-ro')
@@ -52,12 +55,51 @@ class TensorView(object):
             ax.set_ylabel("y")
             ax.set_xticks(self.vectorNx)
             ax.set_yticks(self.vectorNy)
-            
+
+        elif self.dim == 3:
+            if direction == 'z':
+                nX = np.ceil(np.sqrt(self.nCz))
+                nY = np.ceil(self.nCz/nX)
+                C = np.zeros((nX*self.nCx, nY*self.nCy))
+
+                Ic = I[:].reshape(self.n, order='F')
+
+                nCx = self.nCx
+                nCy = self.nCy
+                for iy in range(int(nY)):
+                    for ix in range(int(nX)):
+                        iz = ix + iy*nX
+                        if iz < self.nCz:
+                            C[ix*nCx:(ix+1)*nCx, iy*nCy:(iy+1)*nCy] = Ic[:, :, iz]
+                        else:
+                            C[ix*nCx:(ix+1)*nCx, iy*nCy:(iy+1)*nCy] = np.nan
+
+                C = np.ma.masked_where(np.isnan(C), C)
+                xx = np.r_[0, np.cumsum(np.kron(np.ones((nX, 1)), self.hx).ravel())]
+                yy = np.r_[0, np.cumsum(np.kron(np.ones((nY, 1)), self.hy).ravel())]
+                ph = ax.pcolormesh(xx, yy, C.T)
+                # Plot the lines
+                gx = np.r_[0, np.cumsum(np.kron(np.ones((nX, 1)), np.sum(self.hy)).ravel())]
+                gy = np.r_[0, np.cumsum(np.kron(np.ones((nY, 1)), np.sum(self.hx)).ravel())]
+                # Repeat and seperate with NaN
+                gxX = np.c_[gx, gx, gx+np.nan].ravel()
+                gxY = np.kron(np.ones((nX+1, 1)), np.array([0, sum(self.hy)*nY, np.nan])).ravel()
+                gyX = np.kron(np.ones((nY+1, 1)), np.array([0, sum(self.hx)*nX, np.nan])).ravel()
+                gyY = np.c_[gy, gy, gy+np.nan].ravel()
+                ax.plot(gxX, gxY, 'w-', linewidth=2)
+                ax.plot(gyX, gyY, 'w-', linewidth=2)
+
+                if numbering:
+                    pad = np.sum(self.hx)*0.04
+                    for iy in range(int(nY)):
+                        for ix in range(int(nX)):
+                            iz = ix + iy*nX
+                            ax.text((ix+1)*self.vectorNx[-1]-pad,(iy)*self.vectorNy[-1]+pad,
+                                     '#%i'%iz,color='w',verticalalignment='bottom',horizontalalignment='right',size='x-large')
 
         fig.show()
         return ph
 
-
     def plotGrid(self):
         """Plot the nodal, cell-centered and staggered grids for 1,2 and 3 dimensions."""
         if self.dim == 1: