Merge branch 'dev' into Dom_Dev

.travis.yml

@@ -33,7 +33,7 @@ before_install:
 
 # Install packages
 install:
-  - conda install --yes pip python=$TRAVIS_PYTHON_VERSION numpy scipy matplotlib cython ipython nose
+  - conda install --yes pip python=$TRAVIS_PYTHON_VERSION numpy scipy matplotlib cython ipython nose vtk
   - pip install nose-cov python-coveralls
 
   - git clone https://github.com/rowanc1/pymatsolver.git

SimPEG/Directives.py

@@ -206,6 +206,36 @@ class SaveOutputEveryIteration(_SaveEveryIteration):
         f.write(' %3d %1.4e %1.4e %1.4e %1.4e\n'%(self.opt.iter, self.invProb.beta, self.invProb.phi_d, self.invProb.phi_m, self.opt.f))
         f.close()
 
+class SaveOutputDictEveryIteration(_SaveEveryIteration):
+    """SaveOutputDictEveryIteration"""
+
+    def initialize(self):
+        print "SimPEG.SaveOutputDictEveryIteration will save your inversion progress as dictionary: '###-%s.npz'"%self.fileName
+
+    def endIter(self):
+        # Save the data.
+        ms = self.reg.Ws * ( self.reg.mapping * (self.invProb.curModel - self.reg.mref) )
+        phi_ms = 0.5*ms.dot(ms)
+        if self.reg.smoothModel == True:
+            mref = self.reg.mref
+        else:
+            mref = 0
+        mx = self.reg.Wx * ( self.reg.mapping * (self.invProb.curModel - mref) )
+        phi_mx = 0.5 * mx.dot(mx)
+        if self.prob.mesh.dim==2:
+            my = self.reg.Wy * ( self.reg.mapping * (self.invProb.curModel - mref) )
+            phi_my = 0.5 * my.dot(my)
+        else:
+            phi_my = 'NaN'
+        if self.prob.mesh.dim==3:
+            mz = self.reg.Wz * ( self.reg.mapping * (self.invProb.curModel - mref) )
+            phi_mz = 0.5 * mz.dot(mz)
+        else:
+            phi_mz = 'NaN'
+
+
+        # Save the file as a npz
+        np.savez('{:03d}-{:s}'.format(self.opt.iter,self.fileName), iter=self.opt.iter, beta=self.invProb.beta, phi_d=self.invProb.phi_d, phi_m=self.invProb.phi_m, phi_ms=phi_ms, phi_mx=phi_mx, phi_my=phi_my, phi_mz=phi_mz,f=self.opt.f, m=self.invProb.curModel,dpred=self.invProb.dpred)
 
 
 

SimPEG/EM/FDEM/FDEM.py

@@ -50,7 +50,7 @@ class BaseFDEMProblem(BaseEMProblem):
             Srcs = self.survey.getSrcByFreq(freq)
             ftype = self._fieldType + 'Solution'
             F[Srcs, ftype] = sol
-
+            Ainv.clean()
         return F
 
     def Jvec(self, m, v, f=None):
@@ -89,6 +89,7 @@ class BaseFDEMProblem(BaseEMProblem):
 
                     Jv[src, rx] = P(Df_Dm)
 
+            Ainv.clean()
         return Utils.mkvc(Jv)
 
     def Jtvec(self, m, v, f=None):
@@ -139,7 +140,8 @@ class BaseFDEMProblem(BaseEMProblem):
                         Jtv += - np.array(du_dmT,dtype=complex).real
                     else:
                         raise Exception('Must be real or imag')
-
+            
+            ATinv.clean()
         return Jtv
 
     def getSourceTerm(self, freq):

SimPEG/EM/TDEM/BaseTDEM.py

@@ -37,13 +37,21 @@ class BaseTDEMProblem(BaseTimeProblem, BaseEMProblem):
 
     _FieldsForward_pair = FieldsTDEM  #: used for the forward calculation only
 
+    waveformType = "STEPOFF"
+    current = None
+
+    def currentwaveform(self, wave):
+        self._timeSteps = np.diff(wave[:,0])
+        self.current = wave[:,1]
+        self.waveformType = "GENERAL"
+
     def fields(self, m):
         if self.verbose: print '%s\nCalculating fields(m)\n%s'%('*'*50,'*'*50)
         self.curModel = m
         # Create a fields storage object
         F = self._FieldsForward_pair(self.mesh, self.survey)
         for src in self.survey.srcList:
-            # Set the initial conditions            
+            # Set the initial conditions
             F[src,:,0] = src.getInitialFields(self.mesh)
         F = self.forward(m, self.getRHS, F=F)
         if self.verbose: print '%s\nDone calculating fields(m)\n%s'%('*'*50,'*'*50)

SimPEG/InvProblem.py

@@ -66,8 +66,8 @@ class BaseInvProblem(object):
         self.curModel = m0
 
         print """SimPEG.InvProblem is setting bfgsH0 to the inverse of the eval2Deriv.
-                    ***Done using same solver as the problem***"""
-        self.opt.bfgsH0 = self.prob.Solver(self.reg.eval2Deriv(self.curModel))
+                    ***Done using same Solver and solverOpts as the problem***"""
+        self.opt.bfgsH0 = self.prob.Solver(self.reg.eval2Deriv(self.curModel), **self.prob.solverOpts)
 
     @property
     def warmstart(self):

SimPEG/Maps.py

@@ -10,21 +10,25 @@ class IdentityMap(object):
     SimPEG Map
 
     """
-
     __metaclass__ = Utils.SimPEGMetaClass
 
-    mesh = None      #: A SimPEG Mesh
-
-    def __init__(self, mesh, **kwargs):
+    def __init__(self, mesh=None, nP=None, **kwargs):
         Utils.setKwargs(self, **kwargs)
+
+        if nP is not None:
+            assert type(nP) in [int, long], ' Number of parameters must be an integer.'
+
         self.mesh = mesh
+        self._nP  = nP
 
     @property
     def nP(self):
         """
             :rtype: int
-            :return: number of parameters in the model
+            :return: number of parameters that the mapping accepts
         """
+        if self._nP is not None:
+            return self._nP
         if self.mesh is None:
             return '*'
         return self.mesh.nC
@@ -32,11 +36,15 @@ class IdentityMap(object):
     @property
     def shape(self):
         """
-            The default shape is (mesh.nC, nP).
+            The default shape is (mesh.nC, nP) if the mesh is defined.
+            If this is a meshless mapping (i.e. nP is defined independently)
+            the shape will be the the shape (nP,nP).
 
             :rtype: (int,int)
             :return: shape of the operator as a tuple
         """
+        if self._nP is not None:
+            return (self.nP, self.nP)
         if self.mesh is None:
             return ('*', self.nP)
         return (self.mesh.nC, self.nP)
@@ -118,6 +126,7 @@ class IdentityMap(object):
     def __str__(self):
         return "%s(%s,%s)" % (self.__class__.__name__, self.shape[0], self.shape[1])
 
+
 class ComboMap(IdentityMap):
     """Combination of various maps."""
 
@@ -475,7 +484,7 @@ class ActiveCells(IdentityMap):
         else:
             self.valInactive = valInactive.copy()
         self.valInactive[self.indActive] = 0
-        
+
         inds = np.nonzero(self.indActive)[0]
         self.P = sp.csr_matrix((np.ones(inds.size),(inds, range(inds.size))), shape=(self.nC, self.nP))
 
@@ -708,7 +717,7 @@ class PolyMap(IdentityMap):
         Parameterize the model space using a polynomials in a wholespace.
 
         ..math::
-            
+
             y = \mathbf{V} c
 
         Define the model as:
@@ -752,10 +761,10 @@ class PolyMap(IdentityMap):
             else:
                 raise(Exception("Input for normal = X or Y or Z"))
         #3D
-        elif self.mesh.dim == 3: 
+        elif self.mesh.dim == 3:
             X = self.mesh.gridCC[:,0]
-            Y = self.mesh.gridCC[:,1]            
-            Z = self.mesh.gridCC[:,2]            
+            Y = self.mesh.gridCC[:,1]
+            Z = self.mesh.gridCC[:,2]
             if self.normal =='X':
                 f = polynomial.polyval2d(Y, Z, c.reshape((self.order[0]+1,self.order[1]+1))) - X
             elif self.normal =='Y':
@@ -766,43 +775,43 @@ class PolyMap(IdentityMap):
                 raise(Exception("Input for normal = X or Y or Z"))
         else:
             raise(Exception("Only supports 2D"))
-                    
+
 
         return sig1+(sig2-sig1)*(np.arctan(alpha*f)/np.pi+0.5)
-    
+
     def deriv(self, m):
         alpha = self.slope
         sig1,sig2, c = m[0],m[1],m[2:]
         if self.logSigma:
             sig1, sig2 = np.exp(sig1), np.exp(sig2)
         #2D
-        if self.mesh.dim == 2:            
+        if self.mesh.dim == 2:
             X = self.mesh.gridCC[:,0]
             Y = self.mesh.gridCC[:,1]
 
             if self.normal =='X':
                 f = polynomial.polyval(Y, c) - X
-                V = polynomial.polyvander(Y, len(c)-1)        
+                V = polynomial.polyvander(Y, len(c)-1)
             elif self.normal =='Y':
                 f = polynomial.polyval(X, c) - Y
-                V = polynomial.polyvander(X, len(c)-1)        
+                V = polynomial.polyvander(X, len(c)-1)
             else:
-                raise(Exception("Input for normal = X or Y or Z"))            
+                raise(Exception("Input for normal = X or Y or Z"))
         #3D
-        elif self.mesh.dim == 3: 
+        elif self.mesh.dim == 3:
             X = self.mesh.gridCC[:,0]
             Y = self.mesh.gridCC[:,1]
             Z = self.mesh.gridCC[:,2]
 
             if self.normal =='X':
                 f = polynomial.polyval2d(Y, Z, c.reshape((self.order[0]+1,self.order[1]+1))) - X
-                V = polynomial.polyvander2d(Y, Z, self.order)        
+                V = polynomial.polyvander2d(Y, Z, self.order)
             elif self.normal =='Y':
                 f = polynomial.polyval2d(X, Z, c.reshape((self.order[0]+1,self.order[1]+1))) - Y
-                V = polynomial.polyvander2d(X, Z, self.order)        
+                V = polynomial.polyvander2d(X, Z, self.order)
             elif self.normal =='Z':
                 f = polynomial.polyval2d(X, Y, c.reshape((self.order[0]+1,self.order[1]+1))) - Z
-                V = polynomial.polyvander2d(X, Y, self.order)        
+                V = polynomial.polyvander2d(X, Y, self.order)
             else:
                 raise(Exception("Input for normal = X or Y or Z"))
 
@@ -815,16 +824,16 @@ class PolyMap(IdentityMap):
 
         g3 = Utils.sdiag(alpha*(sig2-sig1)/(1.+(alpha*f)**2)/np.pi)*V
 
-        return sp.csr_matrix(np.c_[g1,g2,g3])  
+        return sp.csr_matrix(np.c_[g1,g2,g3])
 
 class SplineMap(IdentityMap):
 
     """SplineMap
 
-        Parameterize the boundary of two geological units using a spline interpolation 
+        Parameterize the boundary of two geological units using a spline interpolation
 
         ..math::
-            
+
             g = f(x)-y
 
         Define the model as:
@@ -849,7 +858,7 @@ class SplineMap(IdentityMap):
     def nP(self):
         if self.mesh.dim == 2:
             return np.size(self.pts)+2
-        elif self.mesh.dim == 3: 
+        elif self.mesh.dim == 3:
             return np.size(self.pts)*2+2
         else:
             raise(Exception("Only supports 2D and 3D"))
@@ -866,28 +875,28 @@ class SplineMap(IdentityMap):
             X = self.mesh.gridCC[:,0]
             Y = self.mesh.gridCC[:,1]
             self.spl = UnivariateSpline(self.pts, c, k=self.order, s=0)
-            if self.normal =='X':                
+            if self.normal =='X':
                 f = self.spl(Y) - X
             elif self.normal =='Y':
                 f = self.spl(X) - Y
             else:
                 raise(Exception("Input for normal = X or Y or Z"))
 
-        # 3D: 
-        # Comments: 
+        # 3D:
+        # Comments:
         # Make two spline functions and link them using linear interpolation.
         # This is not quite direct extension of 2D to 3D case
         # Using 2D interpolation  is possible
 
-        elif self.mesh.dim == 3: 
+        elif self.mesh.dim == 3:
             X = self.mesh.gridCC[:,0]
-            Y = self.mesh.gridCC[:,1]            
+            Y = self.mesh.gridCC[:,1]
             Z = self.mesh.gridCC[:,2]
 
-            npts = np.size(self.pts)            
+            npts = np.size(self.pts)
             if np.mod(c.size, 2):
                 raise(Exception("Put even points!"))
-            
+
             self.spl = {"splb":UnivariateSpline(self.pts, c[:npts], k=self.order, s=0),
                         "splt":UnivariateSpline(self.pts, c[npts:], k=self.order, s=0)}
 
@@ -902,7 +911,7 @@ class SplineMap(IdentityMap):
                 raise(Exception("Input for normal = X or Y or Z"))
         else:
             raise(Exception("Only supports 2D and 3D"))
-                    
+
 
         return sig1+(sig2-sig1)*(np.arctan(alpha*f)/np.pi+0.5)
 
@@ -912,7 +921,7 @@ class SplineMap(IdentityMap):
         if self.logSigma:
             sig1, sig2 = np.exp(sig1), np.exp(sig2)
         #2D
-        if self.mesh.dim == 2:            
+        if self.mesh.dim == 2:
             X = self.mesh.gridCC[:,0]
             Y = self.mesh.gridCC[:,1]
 
@@ -921,9 +930,9 @@ class SplineMap(IdentityMap):
             elif self.normal =='Y':
                 f = self.spl(X) - Y
             else:
-                raise(Exception("Input for normal = X or Y or Z"))            
+                raise(Exception("Input for normal = X or Y or Z"))
         #3D
-        elif self.mesh.dim == 3: 
+        elif self.mesh.dim == 3:
             X = self.mesh.gridCC[:,0]
             Y = self.mesh.gridCC[:,1]
             Z = self.mesh.gridCC[:,2]
@@ -931,7 +940,7 @@ class SplineMap(IdentityMap):
                 zb = self.ptsv[0]
                 zt = self.ptsv[1]
                 flines = (self.spl["splt"](Y)-self.spl["splb"](Y))*(Z-zb)/(zt-zb) + self.spl["splb"](Y)
-                f = flines - X            
+                f = flines - X
             # elif self.normal =='Y':
             # elif self.normal =='Z':
             else:
@@ -944,7 +953,7 @@ class SplineMap(IdentityMap):
             g1 = -(np.arctan(alpha*f)/np.pi + 0.5) + 1.0
             g2 = (np.arctan(alpha*f)/np.pi + 0.5)
 
-        
+
         if self.mesh.dim ==2:
             g3 = np.zeros((self.mesh.nC, self.npts))
             if self.normal =='Y':
@@ -958,7 +967,7 @@ class SplineMap(IdentityMap):
                     cb = c.copy()
                     dy = self.mesh.hy[ind]*1.5
                     ca[i] = ctemp+dy
-                    cb[i] = ctemp-dy                
+                    cb[i] = ctemp-dy
                     spla = UnivariateSpline(self.pts, ca, k=self.order, s=0)
                     splb = UnivariateSpline(self.pts, cb, k=self.order, s=0)
                     fderiv = (spla(X)-splb(X))/(2*dy)
@@ -968,7 +977,7 @@ class SplineMap(IdentityMap):
             g3 = np.zeros((self.mesh.nC, self.npts*2))
             if self.normal =='X':
                 # Here we use perturbation to compute sensitivity
-                for i in range(self.npts*2):                    
+                for i in range(self.npts*2):
                     ctemp = c[i]
                     ind = np.argmin(abs(self.mesh.vectorCCy-ctemp))
                     ca = c.copy()
@@ -982,20 +991,20 @@ class SplineMap(IdentityMap):
                         splbb = UnivariateSpline(self.pts, cb[:self.npts], k=self.order, s=0)
                         flinesa = (self.spl["splt"](Y)-splba(Y))*(Z-zb)/(zt-zb) + splba(Y) - X
                         flinesb = (self.spl["splt"](Y)-splbb(Y))*(Z-zb)/(zt-zb) + splbb(Y) - X
-                    #treat top boundary                        
+                    #treat top boundary
                     else:
                         splta = UnivariateSpline(self.pts, ca[self.npts:], k=self.order, s=0)
                         spltb = UnivariateSpline(self.pts, ca[self.npts:], k=self.order, s=0)
                         flinesa = (self.spl["splt"](Y)-splta(Y))*(Z-zb)/(zt-zb) + splta(Y) - X
-                        flinesb = (self.spl["splt"](Y)-spltb(Y))*(Z-zb)/(zt-zb) + spltb(Y) - X                        
-                    fderiv = (flinesa-flinesb)/(2*dy)                                
+                        flinesb = (self.spl["splt"](Y)-spltb(Y))*(Z-zb)/(zt-zb) + spltb(Y) - X
+                    fderiv = (flinesa-flinesb)/(2*dy)
                     g3[:,i] = Utils.sdiag(alpha*(sig2-sig1)/(1.+(alpha*f)**2)/np.pi)*fderiv
         else :
             raise(Exception("Not Implemented for Y and Z, your turn :)"))
-        return sp.csr_matrix(np.c_[g1,g2,g3])  
+        return sp.csr_matrix(np.c_[g1,g2,g3])
+
 
 
 
 
 
-            

SimPEG/Mesh/MeshIO.py

@@ -0,0 +1,416 @@
+import numpy as np, os
+from SimPEG import Utils
+
+class TensorMeshIO(object):
+
+    @classmethod
+    def readUBC(TensorMesh, fileName):
+        """
+            Read UBC GIF 3DTensor mesh and generate 3D Tensor mesh in simpegTD
+
+            Input:
+            :param fileName, path to the UBC GIF mesh file
+
+            Output:
+            :param SimPEG TensorMesh object
+        """
+
+        # Interal function to read cell size lines for the UBC mesh files.
+        def readCellLine(line):
+            for seg in line.split():
+                if '*' in seg:
+                    st = seg
+                    sp = seg.split('*')
+                    re = np.array(sp[0],dtype=int)*(' ' + sp[1])
+                    line = line.replace(st,re.strip())
+            return np.array(line.split(),dtype=float)
+
+        # Read the file as line strings, remove lines with comment = !
+        msh = np.genfromtxt(fileName,delimiter='\n',dtype=np.str,comments='!')
+
+        # Fist line is the size of the model
+        sizeM = np.array(msh[0].split(),dtype=float)
+        # Second line is the South-West-Top corner coordinates.
+        x0 = np.array(msh[1].split(),dtype=float)
+        # Read the cell sizes
+        h1 = readCellLine(msh[2])
+        h2 = readCellLine(msh[3])
+        h3temp = readCellLine(msh[4])
+        h3 = h3temp[::-1] # Invert the indexing of the vector to start from the bottom.
+        # Adjust the reference point to the bottom south west corner
+        x0[2] = x0[2] - np.sum(h3)
+        # Make the mesh
+        tensMsh = TensorMesh([h1,h2,h3],x0)
+        return tensMsh
+
+    @classmethod
+    def readVTK(TensorMesh, fileName):
+        """
+            Read VTK Rectilinear (vtr xml file) and return SimPEG Tensor mesh and model
+
+            Input:
+            :param vtrFileName, path to the vtr model file to write to
+
+            Output:
+            :return SimPEG TensorMesh object
+            :return SimPEG model dictionary
+
+        """
+        # Import
+        from vtk import vtkXMLRectilinearGridReader as vtrFileReader
+        from vtk.util.numpy_support import vtk_to_numpy
+
+        # Read the file
+        vtrReader = vtrFileReader()
+        vtrReader.SetFileName(fileName)
+        vtrReader.Update()
+        vtrGrid = vtrReader.GetOutput()
+        # Sort information
+        hx = np.abs(np.diff(vtk_to_numpy(vtrGrid.GetXCoordinates())))
+        xR = vtk_to_numpy(vtrGrid.GetXCoordinates())[0]
+        hy = np.abs(np.diff(vtk_to_numpy(vtrGrid.GetYCoordinates())))
+        yR = vtk_to_numpy(vtrGrid.GetYCoordinates())[0]
+        zD = np.diff(vtk_to_numpy(vtrGrid.GetZCoordinates()))
+        # Check the direction of hz
+        if np.all(zD < 0):
+            hz = np.abs(zD[::-1])
+            zR = vtk_to_numpy(vtrGrid.GetZCoordinates())[-1]
+        else:
+            hz = np.abs(zD)
+            zR = vtk_to_numpy(vtrGrid.GetZCoordinates())[0]
+        x0 = np.array([xR,yR,zR])
+
+        # Make the SimPEG object
+        tensMsh = TensorMesh([hx,hy,hz],x0)
+
+        # Grap the models
+        models = {}
+        for i in np.arange(vtrGrid.GetCellData().GetNumberOfArrays()):
+            modelName = vtrGrid.GetCellData().GetArrayName(i)
+            if np.all(zD < 0):
+                modFlip = vtk_to_numpy(vtrGrid.GetCellData().GetArray(i))
+                tM = tensMsh.r(modFlip,'CC','CC','M')
+                modArr = tensMsh.r(tM[:,:,::-1],'CC','CC','V')
+            else:
+                modArr = vtk_to_numpy(vtrGrid.GetCellData().GetArray(i))
+            models[modelName] = modArr
+
+        # Return the data
+        return tensMsh, models
+
+    def writeVTK(mesh, fileName, models=None):
+        """
+        Makes and saves a VTK rectilinear file (vtr) for a simpeg Tensor mesh and model.
+
+        Input:
+        :param str, path to the output vtk file
+        :param mesh, SimPEG TensorMesh object - mesh to be transfer to VTK
+        :param models, dictionary of numpy.array - Name('s) and array('s). Match number of cells
+
+        """
+        # Import
+        from vtk import vtkRectilinearGrid as rectGrid, vtkXMLRectilinearGridWriter as rectWriter, VTK_VERSION
+        from vtk.util.numpy_support import numpy_to_vtk
+
+        # Deal with dimensionalities
+        if mesh.dim >= 1:
+            vX = mesh.vectorNx
+            xD = mesh.nNx
+            yD,zD = 1,1
+            vY, vZ = np.array([0,0])
+        if mesh.dim >= 2:
+            vY = mesh.vectorNy
+            yD = mesh.nNy
+        if mesh.dim == 3:
+            vZ = mesh.vectorNz
+            zD = mesh.nNz
+        # Use rectilinear VTK grid.
+        # Assign the spatial information.
+        vtkObj = rectGrid()
+        vtkObj.SetDimensions(xD,yD,zD)
+        vtkObj.SetXCoordinates(numpy_to_vtk(vX,deep=1))
+        vtkObj.SetYCoordinates(numpy_to_vtk(vY,deep=1))
+        vtkObj.SetZCoordinates(numpy_to_vtk(vZ,deep=1))
+
+        # Assign the model('s) to the object
+        if models is not None:
+            for item in models.iteritems():
+                # Convert numpy array
+                vtkDoubleArr = numpy_to_vtk(item[1],deep=1)
+                vtkDoubleArr.SetName(item[0])
+                vtkObj.GetCellData().AddArray(vtkDoubleArr)
+            # Set the active scalar
+            vtkObj.GetCellData().SetActiveScalars(models.keys()[0])
+        # vtkObj.Update()
+
+        # Check the extension of the fileName
+        ext = os.path.splitext(fileName)[1]
+        if ext is '':
+            fileName = fileName + '.vtr'
+        elif ext not in '.vtr':
+            raise IOError('{:s} is an incorrect extension, has to be .vtr')
+        # Write the file.
+        vtrWriteFilter = rectWriter()
+        if float(VTK_VERSION.split('.')[0]) >=6:
+            vtrWriteFilter.SetInputData(vtkObj)
+        else:
+            vtuWriteFilter.SetInput(vtuObj)
+        vtrWriteFilter.SetFileName(fileName)
+        vtrWriteFilter.Update()
+
+
+    def readModelUBC(mesh, fileName):
+        """
+            Read UBC 3DTensor mesh model and generate 3D Tensor mesh model in simpeg
+
+            Input:
+            :param fileName, path to the UBC GIF mesh file to read
+            :param mesh, TensorMesh object, mesh that coresponds to the model
+
+            Output:
+            :return numpy array, model with TensorMesh ordered
+        """
+        f = open(fileName, 'r')
+        model = np.array(map(float, f.readlines()))
+        f.close()
+        model = np.reshape(model, (mesh.nCz, mesh.nCx, mesh.nCy), order = 'F')
+        model = model[::-1,:,:]
+        model = np.transpose(model, (1, 2, 0))
+        model = Utils.mkvc(model)
+        return model
+
+    def writeModelUBC(mesh, fileName, model):
+        """
+            Writes a model associated with a SimPEG TensorMesh
+            to a UBC-GIF format model file.
+
+            :param str fileName: File to write to
+            :param simpeg.Mesh.TensorMesh mesh: The mesh
+            :param numpy.ndarray model: The model
+        """
+
+        # Reshape model to a matrix
+        modelMat = mesh.r(model,'CC','CC','M')
+        # Transpose the axes
+        modelMatT = modelMat.transpose((2,0,1))
+        # Flip z to positive down
+        modelMatTR = Utils.mkvc(modelMatT[::-1,:,:])
+
+        np.savetxt(fileName, modelMatTR.ravel())
+
+    def writeUBC(mesh, fileName, models=None):
+        """
+            Writes a SimPEG TensorMesh to a UBC-GIF format mesh file.
+
+            :param str fileName: File to write to
+            :param simpeg.Mesh.TensorMesh mesh: The mesh
+
+        """
+        assert mesh.dim == 3
+        s = ''
+        s += '%i %i %i\n' %tuple(mesh.vnC)
+        origin = mesh.x0 + np.array([0,0,mesh.hz.sum()]) # Have to it in the same operation or use mesh.x0.copy(), otherwise the mesh.x0 is updated.
+        origin.dtype = float
+
+        s += '%.2f %.2f %.2f\n' %tuple(origin)
+        s += ('%.2f '*mesh.nCx+'\n')%tuple(mesh.hx)
+        s += ('%.2f '*mesh.nCy+'\n')%tuple(mesh.hy)
+        s += ('%.2f '*mesh.nCz+'\n')%tuple(mesh.hz[::-1])
+        f = open(fileName, 'w')
+        f.write(s)
+        f.close()
+
+        if models is None: return
+        assert type(models) is dict, 'models must be a dict'
+        for key in models:
+            assert type(key) is str, 'The dict key is a file name'
+            mesh.writeModelUBC(key, models[key])
+
+class TreeMeshIO(object):
+
+    def writeUBC(mesh, fileName, models=None):
+        """
+            Write UBC ocTree mesh and model files from a simpeg ocTree mesh and model.
+
+            :param str fileName: File to write to
+            :param simpeg.Mesh.TreeMesh mesh: The mesh
+            :param dictionary models: The models in a dictionary, where the keys is the name of the of the model file
+        """
+
+        # Calculate information to write in the file.
+        # Number of cells in the underlying mesh
+        nCunderMesh = np.array([h.size for h in mesh.h],dtype=np.int64)
+        # The top-south-west most corner of the mesh
+        tswCorn = mesh.x0 + np.array([0,0,np.sum(mesh.h[2])])
+        # Smallest cell size
+        smallCell = np.array([h.min() for h in mesh.h])
+        # Number of cells
+        nrCells = mesh.nC
+
+        ## Extract iformation about the cells.
+        # cell pointers
+        cellPointers = np.array([c._pointer for c in mesh])
+        # cell with
+        cellW = np.array([ mesh._levelWidth(i) for i in cellPointers[:,-1] ])
+        # Need to shift the pointers to work with UBC indexing
+        # UBC Octree indexes always the top-left-close (top-south-west) corner first and orders the cells in z(top-down),x,y vs x,y,z(bottom-up).
+        # Shift index up by 1
+        ubcCellPt = cellPointers[:,0:-1].copy() + np.array([1.,1.,1.])
+        # Need reindex the z index to be from the top-left-close corner and to be from the global top.
+        ubcCellPt[:,2] = ( nCunderMesh[-1] + 2) - (ubcCellPt[:,2] + cellW)
+
+        # Reorder the ubcCellPt
+        ubcReorder = np.argsort(ubcCellPt.view(','.join(3*['float'])),axis=0,order=['f2','f1','f0'])[:,0]
+        # Make a array with the pointers and the withs, that are order in the ubc ordering
+        indArr = np.concatenate((ubcCellPt[ubcReorder,:],cellW[ubcReorder].reshape((-1,1)) ),axis=1)
+
+        ## Write the UBC octree mesh file
+        with open(fileName,'w') as mshOut:
+            mshOut.write('{:.0f} {:.0f} {:.0f}\n'.format(nCunderMesh[0],nCunderMesh[1],nCunderMesh[2]))
+            mshOut.write('{:.4f} {:.4f} {:.4f}\n'.format(tswCorn[0],tswCorn[1],tswCorn[2]))
+            mshOut.write('{:.3f} {:.3f} {:.3f}\n'.format(smallCell[0],smallCell[1],smallCell[2]))
+            mshOut.write('{:.0f} \n'.format(nrCells))
+            np.savetxt(mshOut,indArr,fmt='%i')
+
+        ## Print the models
+        # Assign the model('s) to the object
+        if models is not None:
+            # indUBCvector = np.argsort(cX0[np.argsort(np.concatenate((cX0[:,0:2],cX0[:,2:3].max() - cX0[:,2:3]),axis=1).view(','.join(3*['float'])),axis=0,order=('f2','f1','f0'))[:,0]].view(','.join(3*['float'])),axis=0,order=('f2','f1','f0'))[:,0]
+            for item in models.iteritems():
+                # Save the data
+                np.savetxt(item[0],item[1][ubcReorder],fmt='%3.5e')
+
+    @classmethod
+    def readUBC(TreeMesh, meshFile):
+        """
+            Read UBC 3D OcTree mesh and/or modelFiles
+
+            Input:
+            :param str meshFile: path to the UBC GIF OcTree mesh file to read
+
+            Output:
+            :return SimPEG.Mesh.TreeMesh mesh: The octree mesh
+            :return list of ndarray's: models as a list of numpy array's
+        """
+
+        ## Read the file lines
+        fileLines = np.genfromtxt(meshFile,dtype=str,delimiter='\n')
+        # Extract the data
+        nCunderMesh = np.array(fileLines[0].split(),dtype=float)
+        # I think this is the case?
+        if np.unique(nCunderMesh).size >1:
+            raise Exception('SimPEG TreeMeshes have the same number of cell in all directions')
+        tswCorn = np.array(fileLines[1].split(),dtype=float)
+        smallCell = np.array(fileLines[2].split(),dtype=float)
+        nrCells = np.array(fileLines[3].split(),dtype=float)
+        # Read the index array
+        indArr = np.genfromtxt(fileLines[4::],dtype=np.int)
+
+        ## Calculate simpeg parameters
+        h1,h2,h3 = [np.ones(nr)*sz for nr,sz in zip(nCunderMesh,smallCell)]
+        x0 = tswCorn - np.array([0,0,np.sum(h3)])
+        # Need to convert the index array to a points list that complies with SimPEG TreeMesh.
+        # Shift to start at 0
+        simpegCellPt = indArr[:,0:-1].copy()
+        simpegCellPt[:,2] = ( nCunderMesh[-1] + 2) - (simpegCellPt[:,2] + indArr[:,3])
+        # Need reindex the z index to be from the bottom-left-close corner and to be from the global bottom.
+        simpegCellPt = simpegCellPt - np.array([1.,1.,1.])
+
+        # Calculate the cell level
+        simpegLevel = np.log2(np.min(nCunderMesh)) - np.log2(indArr[:,3])
+        # Make a pointer matrix
+        simpegPointers = np.concatenate((simpegCellPt,simpegLevel.reshape((-1,1))),axis=1)
+
+        ## Make the tree mesh
+        mesh = TreeMesh([h1,h2,h3],x0)
+        mesh._cells = set([mesh._index(p) for p in simpegPointers.tolist()])
+
+        # Figure out the reordering
+        mesh._simpegReorderUBC = np.argsort(np.array([mesh._index(i) for i in simpegPointers.tolist()]))
+        # mesh._simpegReorderUBC = np.argsort((np.array([[1,1,1,-1]])*simpegPointers).view(','.join(4*['float'])),axis=0,order=['f3','f2','f1','f0'])[:,0]
+
+        return mesh
+
+
+    def readModelUBC(mesh, fileName):
+        """
+            Read UBC OcTree model and get vector
+
+            Input:
+            :param fileName, path to the UBC GIF model file to read
+
+            Output:
+            :return numpy array, OcTree model
+        """
+
+        if type(fileName) is list:
+            out = {}
+            for f in fileName:
+                out[f] = mesh.readModelUBC(f)
+            return out
+
+        assert hasattr(mesh, '_simpegReorderUBC'), 'The file must have been loaded from a UBC format.'
+        assert mesh.dim == 3
+
+        modList = []
+        modArr = np.loadtxt(fileName)
+        if len(modArr.shape) == 1:
+            modList.append(modArr[mesh._simpegReorderUBC])
+        else:
+            modList.append(modArr[mesh._simpegReorderUBC,:])
+        return modList
+
+    def writeVTK(mesh, fileName, models=None):
+        """
+        Function to write a VTU file from a SimPEG TreeMesh and model.
+        """
+        import vtk
+        from vtk import vtkXMLUnstructuredGridWriter as Writer, VTK_VERSION
+        from vtk.util.numpy_support import numpy_to_vtk, numpy_to_vtkIdTypeArray
+
+        if str(type(mesh)).split()[-1][1:-2] not in 'SimPEG.Mesh.TreeMesh.TreeMesh':
+            raise IOError('mesh is not a SimPEG TreeMesh.')
+
+        # Make the data parts for the vtu object
+        # Points
+        mesh.number()
+        ptsMat = mesh._gridN + mesh.x0
+
+        vtkPts = vtk.vtkPoints()
+        vtkPts.SetData(numpy_to_vtk(ptsMat,deep=True))
+        # Cells
+        cellConn = np.array([c.nodes for c in mesh],dtype=np.int64)
+
+        cellsMat = np.concatenate((np.ones((cellConn.shape[0],1),dtype=np.int64)*cellConn.shape[1],cellConn),axis=1).ravel()
+        cellsArr = vtk.vtkCellArray()
+        cellsArr.SetNumberOfCells(cellConn.shape[0])
+        cellsArr.SetCells(cellConn.shape[0],numpy_to_vtkIdTypeArray(cellsMat,deep=True))
+
+        # Make the object
+        vtuObj = vtk.vtkUnstructuredGrid()
+        vtuObj.SetPoints(vtkPts)
+        vtuObj.SetCells(vtk.VTK_VOXEL,cellsArr)
+        # Add the level of refinement as a cell array
+        cellSides = np.array([np.array(vtuObj.GetCell(i).GetBounds()).reshape((3,2)).dot(np.array([-1, 1])) for i in np.arange(vtuObj.GetNumberOfCells())])
+        uniqueLevel, indLevel = np.unique(np.prod(cellSides,axis=1),return_inverse=True)
+        refineLevelArr = numpy_to_vtk(indLevel.max() - indLevel,deep=1)
+        refineLevelArr.SetName('octreeLevel')
+        vtuObj.GetCellData().AddArray(refineLevelArr)
+        # Assign the model('s) to the object
+        if models is not None:
+            for item in models.iteritems():
+                # Convert numpy array
+                vtkDoubleArr = numpy_to_vtk(item[1],deep=1)
+                vtkDoubleArr.SetName(item[0])
+                vtuObj.GetCellData().AddArray(vtkDoubleArr)
+
+        # Make the writer
+        vtuWriteFilter = Writer()
+        if float(VTK_VERSION.split('.')[0]) >=6:
+            vtuWriteFilter.SetInputData(vtuObj)
+        else:
+            vtuWriteFilter.SetInput(vtuObj)
+        vtuWriteFilter.SetFileName(fileName)
+        # Write the file
+        vtuWriteFilter.Update()
+

SimPEG/Mesh/TreeMesh.py

@@ -100,11 +100,12 @@ except Exception, e:
 
 from InnerProducts import InnerProducts
 from TensorMesh import TensorMesh, BaseTensorMesh
+from MeshIO import TreeMeshIO
 import time
 
 MAX_BITS = 20
 
-class TreeMesh(BaseTensorMesh, InnerProducts):
+class TreeMesh(BaseTensorMesh, InnerProducts, TreeMeshIO):
 
     _meshType = 'TREE'
 
@@ -564,15 +565,18 @@ class TreeMesh(BaseTensorMesh, InnerProducts):
         return [p - (p % mod) for p in pointer[:-1]] + [pointer[-1]-1]
 
     def _cellN(self, p):
+        """Node location [x,y(,z)] of a single cell, closest to origin, given a pointer."""
         p = self._asPointer(p)
         return [hi[:p[ii]].sum() for ii, hi in enumerate(self.h)]
 
     def _cellH(self, p):
+        """Widths of a single cell given a pointer."""
         p = self._asPointer(p)
         w = self._levelWidth(p[-1])
         return [hi[p[ii]:p[ii]+w].sum() for ii, hi in enumerate(self.h)]
 
     def _cellC(self, p):
+        """Cell center of a single cell (without origin correction), given a pointer."""
         return (np.array(self._cellH(p))/2.0 + self._cellN(p)).tolist()
 
     def _levelWidth(self, level):
@@ -827,8 +831,10 @@ class TreeMesh(BaseTensorMesh, InnerProducts):
     def _numberCells(self, force=False):
         if not self.__dirtyCells__ and not force: return
         self._cc2i = dict()
+        self._i2cc = dict()
         for ii, c in enumerate(sorted(self._cells)):
             self._cc2i[c] = ii
+            self._i2cc[ii] = c
         self.__dirtyCells__ = False
 
     def _numberNodes(self, force=False):
@@ -1704,9 +1710,9 @@ class TreeMesh(BaseTensorMesh, InnerProducts):
         "Construct the averaging operator on cell faces to cell centers."
         if getattr(self, '_aveF2CC', None) is None:
             if self.dim == 2:
-                self._aveF2CC = 1./self.dim*sp.hstack([self.aveFx2CC, self.aveFy2CC])
+                self._aveF2CC = 1./self.dim*sp.hstack([self.aveFx2CC, self.aveFy2CC]).tocsr()
             elif self.dim == 3:
-                self._aveF2CC = 1./self.dim*sp.hstack([self.aveFx2CC, self.aveFy2CC, self.aveFz2CC])
+                self._aveF2CC = 1./self.dim*sp.hstack([self.aveFx2CC, self.aveFy2CC, self.aveFz2CC]).tocsr()
         return self._aveF2CC
 
     @property
@@ -1714,9 +1720,9 @@ class TreeMesh(BaseTensorMesh, InnerProducts):
         "Construct the averaging operator on cell faces to cell centers."
         if getattr(self, '_aveF2CCV', None) is None:
             if self.dim == 2:
-                self._aveF2CCV = sp.block_diag([self.aveFx2CC, self.aveFy2CC])
+                self._aveF2CCV = sp.block_diag([self.aveFx2CC, self.aveFy2CC]).tocsr()
             elif self.dim == 3:
-                self._aveF2CCV = sp.block_diag([self.aveFx2CC, self.aveFy2CC, self.aveFz2CC])
+                self._aveF2CCV = sp.block_diag([self.aveFx2CC, self.aveFy2CC, self.aveFz2CC]).tocsr()
         return self._aveF2CCV
 
     @property
@@ -2218,6 +2224,25 @@ class TreeMesh(BaseTensorMesh, InnerProducts):
         if showIt: plt.show()
         return tuple(out)
 
+    def __len__(self): return self.nC
+
+    def __getitem__(self, key):
+        if isinstance( key, slice ) :
+            #Get the start, stop, and step from the slice
+            return [self[ii] for ii in xrange(*key.indices(len(self)))]
+        elif isinstance( key, int ) :
+            if key < 0 : #Handle negative indices
+                key += len( self )
+            if key >= len( self ) :
+                raise IndexError, "The index (%d) is out of range."%key
+
+            self._numberCells() # no-op if numbered
+            index   = self._i2cc[key]
+            pointer = self._asPointer(index)
+            return Cell(self, index, pointer)
+        else:
+            raise TypeError, "Invalid argument type."
+
 
 class Cell(object):
     def __init__(self, mesh, index, pointer):
@@ -2225,6 +2250,35 @@ class Cell(object):
         self._index   = index
         self._pointer = pointer
 
+    @property
+    def nodes(self):
+        """The node index in _gridN (this may include hanging nodes)."""
+        M = self.mesh
+        M._numberNodes()
+        p = self._pointer
+        i = self._index
+        w = M._levelWidth(p[-1])
+
+        if M.dim == 2:
+            n = [
+                    i,
+                    M._index([ p[0] + w, p[1]   , p[2]]),
+                    M._index([ p[0]    , p[1]+ w, p[2]]),
+                    M._index([ p[0] + w, p[1]+ w, p[2]]),
+                ]
+        elif self.dim == 3:
+            n = [
+                    i,
+                    M._index([ p[0] + w, p[1]    , p[2]    ,p[3]]),
+                    M._index([ p[0]    , p[1] + w, p[2]    ,p[3]]),
+                    M._index([ p[0] + w, p[1] + w, p[2]    ,p[3]]),
+                    M._index([ p[0]    , p[1]    , p[2] + w,p[3]]),
+                    M._index([ p[0] + w, p[1]    , p[2] + w,p[3]]),
+                    M._index([ p[0]    , p[1] + w, p[2] + w,p[3]]),
+                    M._index([ p[0] + w, p[1] + w, p[2] + w,p[3]]),
+                ]
+        return [M._n2i[_] for _ in n]
+
     @property
     def center(self):
         if getattr(self, '_center', None) is None:
@@ -2282,121 +2336,3 @@ class NotBalancedException(TreeException):
     pass
 class CellLookUpException(TreeException):
     pass
-
-if __name__ == '__main__':
-
-
-    import matplotlib.pyplot as plt
-    import matplotlib
-    from mpl_toolkits.mplot3d import Axes3D
-    import matplotlib.colors as colors
-    import matplotlib.cm as cmx
-
-    def topo(x):
-        return np.sin(x*(2.*np.pi))*0.3 + 0.5
-
-    def function(cell):
-        r = cell.center - np.array([0.5]*len(cell.center))
-        dist = np.sqrt(r.dot(r))
-        # dist2 = np.abs(cell.center[-1] - topo(cell.center[0]))
-
-        # dist = min([dist1,dist2])
-        # if dist < 0.05:
-        #     return 5
-        if dist < 0.1:
-            return 5
-        if dist < 0.2:
-            return 4
-        if dist < 0.4:
-            return 3
-        return 2
-
-    # T = TreeMesh([[(1,128)],[(1,128)],[(1,128)]],levels=7)
-    # T = TreeMesh([128,128,128])
-    # T = TreeMesh([64,64],levels=6)
-    T = TreeMesh([4,4,4])
-    # T = TreeMesh([[(1,128)],[(1,128)]],levels=7)
-    # T.refine(lambda xc:2, balance=False)
-    # T._index([0,0,0])
-    # T._pointer(0)
-
-
-    # tic = time.time()
-    T.refine(function)#, balance=False)
-    # print time.time() - tic
-    # print T.nC
-    T.plotSlice(np.log(T.vol))#np.random.rand(T.nC))
-
-    plt.show()
-    blah
-
-    # T.plotImage(np.arange(len(T.vol)),showIt=True)
-
-    # print T.getFaceInnerProduct()
-    # print T.gridFz
-
-
-    # T._refineCell([8,0,1])
-    # T._refineCell([8,0,2])
-    # T._refineCell([12,0,2])
-    # T._refineCell([8,4,2])
-    # T._refineCell([6,0,3])
-    # T._refineCell([8,8,1])
-    # T._refineCell([0,0,0,1])
-    # T.__dirty__ = True
-
-
-    # print T.gridFx.shape[0], T.nFx
-
-
-
-    ax = plt.subplot(211)
-    ax.spy(T.edgeCurl)
-
-    # print Mesh.TensorMesh([2,2,2]).edgeCurl.todense()
-    # print T.edgeCurl.todense()
-    # print Mesh.TensorMesh([2,2,2]).edgeCurl.todense() - T.edgeCurl.todense()
-    # print T.gridEy - Mesh.TensorMesh([2,2,2]).gridEy
-
-    # print T.edge
-    # T.plotGrid(ax=ax)
-
-    # R = deflationMatrix(T._facesX, T._hangingFx, T._fx2i)
-    # print R
-
-    ax = plt.subplot(212)#, projection='3d')
-    ax.spy(Mesh.TensorMesh([2,2,2]).edgeCurl)
-
-    # ax = plt.subplot(313)
-    # ax.spy(T.faceDiv[:,:T.nFx] * R)
-
-
-    # T.balance()
-    # T.plotGrid(ax=ax)
-
-    # cx = T._getNextCell([0,0,1],direction=0,positive=True)
-    # print cx
-    # # print [T._asPointer(_) for _ in cx]
-    # cx = T._getNextCell([8,0,3],direction=0,positive=False)
-    # print T._asPointer(cx)
-    # cx = T._getNextCell([8,8,1],direction=1,positive=False)
-    # print cx, #[T._asPointer(_) for _ in cx]
-    # cm = T._getNextCell([64,80,4],direction=0,positive=False)
-    # cy = T._getNextCell([64,80,4],direction=1,positive=True)
-    # cp = T._getNextCell([64,80,4],direction=1,positive=False)
-
-    # ax.plot( T._cellN([4,0,1])[0],T._cellN([4,0,1])[1], 'yd')
-    # ax.plot( T._cellN(cx)[0],T._cellN(cx)[1], 'ys')
-    # ax.plot( T._cellN(cm)[0],T._cellN(cm)[1], 'ys')
-    # ax.plot( T._cellN(cy)[0],T._cellN(cy)[1], 'ys')
-    # ax.plot( T._cellN(cp[0])[0],T._cellN(cp[0])[1], 'ys')
-    # ax.plot( T._cellN(cp[1])[0],T._cellN(cp[1])[1], 'ys')
-
-
-
-
-
-    # print T.nN
-
-    plt.show()
-

SimPEG/Problem.py

@@ -32,8 +32,8 @@ class BaseProblem(object):
             val._assertMatchesPair(self.mapPair)
             self._mapping = val
         else:
-            self._mapping = self.PropMap(val) 
-    
+            self._mapping = self.PropMap(val)
+
     def __init__(self, mesh, mapping=None, **kwargs):
         Utils.setKwargs(self, **kwargs)
         assert isinstance(mesh, Mesh.BaseMesh), "mesh must be a SimPEG.Mesh object."
@@ -158,9 +158,6 @@ class BaseProblem(object):
 
 class BaseTimeProblem(BaseProblem):
     """Sets up that basic needs of a time domain problem."""
-    
-    waveformType = "STEPOFF"    
-    current = None
 
     @property
     def timeSteps(self):
@@ -187,11 +184,6 @@ class BaseTimeProblem(BaseProblem):
         self._timeSteps = Utils.meshTensor(value)
         del self.timeMesh
 
-    def currentwaveform(self, wave):
-        self._timeSteps = np.diff(wave[:,0])
-        self.current = wave[:,1]
-        self.waveformType = "GENERAL"
-
     @property
     def nT(self):
         "Number of time steps."

SimPEG/Regularization.py

@@ -20,12 +20,13 @@ class BaseRegularization(object):
     mesh    = None    #: A SimPEG.Mesh instance.
     mref = None       #: Reference model.
 
-    def __init__(self, mesh, mapping=None, **kwargs):
+    def __init__(self, mesh, mapping=None, indActive=None, **kwargs):
         Utils.setKwargs(self, **kwargs)
         self.mesh = mesh
         assert isinstance(mesh, Mesh.BaseMesh), "mesh must be a SimPEG.Mesh object."
         self.mapping = mapping or self.mapPair(mesh)
         self.mapping._assertMatchesPair(self.mapPair)
+        self.indActive = indActive
 
     @property
     def parent(self):
@@ -112,8 +113,6 @@ class BaseRegularization(object):
         return mD.T * ( self.W.T * ( self.W * ( mD * v) ) )
 
 
-
-
 class Tikhonov(BaseRegularization):
     """
     """
@@ -126,14 +125,18 @@ class Tikhonov(BaseRegularization):
     alpha_yy = Utils.dependentProperty('_alpha_yy', 0.0, ['_W', '_Wyy'], "Weight for the second derivative in the y direction")
     alpha_zz = Utils.dependentProperty('_alpha_zz', 0.0, ['_W', '_Wzz'], "Weight for the second derivative in the z direction")
 
-    def __init__(self, mesh, mapping=None, **kwargs):
+    def __init__(self, mesh, mapping=None, indActive = None, **kwargs):
         BaseRegularization.__init__(self, mesh, mapping=mapping, **kwargs)
+        self.indActive = indActive
 
     @property
     def Ws(self):
         """Regularization matrix Ws"""
         if getattr(self,'_Ws', None) is None:
-                self._Ws = Utils.sdiag((self.mesh.vol*self.alpha_s)**0.5)
+            self._Ws = Utils.sdiag((self.mesh.vol*self.alpha_s)**0.5)
+            if self.indActive is not None:
+                Pac = Utils.speye(self.mesh.nC)[:,self.indActive]
+                self._Ws = Pac.T * self._Ws * Pac
         return self._Ws
 
     @property
@@ -142,6 +145,13 @@ class Tikhonov(BaseRegularization):
         if getattr(self, '_Wx', None) is None:
             Ave_x_vol = self.mesh.aveF2CC[:,:self.mesh.nFx].T*self.mesh.vol
             self._Wx = Utils.sdiag((Ave_x_vol*self.alpha_x)**0.5)*self.mesh.cellGradx
+
+            if self.indActive is not None:
+                indActive_Fx = (self.mesh.aveFx2CC.T * self.indActive) == 1
+                Pac = Utils.speye(self.mesh.nC)[:,self.indActive]
+                Pafx = Utils.speye(self.mesh.nFx)[:,indActive_Fx]
+                self._Wx = Pafx.T*self._Wx*Pac
+
         return self._Wx
 
     @property
@@ -150,6 +160,13 @@ class Tikhonov(BaseRegularization):
         if getattr(self, '_Wy', None) is None:
             Ave_y_vol = self.mesh.aveF2CC[:,self.mesh.nFx:np.sum(self.mesh.vnF[:2])].T*self.mesh.vol
             self._Wy = Utils.sdiag((Ave_y_vol*self.alpha_y)**0.5)*self.mesh.cellGrady
+
+            if self.indActive is not None:
+                indActive_Fy = (self.mesh.aveFy2CC.T * self.indActive) == 1
+                Pac = Utils.speye(self.mesh.nC)[:,self.indActive]
+                Pafy = Utils.speye(self.mesh.nFy)[:,indActive_Fy]
+                self._Wy = Pafy.T*self._Wy*Pac
+
         return self._Wy
 
     @property
@@ -158,6 +175,13 @@ class Tikhonov(BaseRegularization):
         if getattr(self, '_Wz', None) is None:
             Ave_z_vol = self.mesh.aveF2CC[:,np.sum(self.mesh.vnF[:2]):].T*self.mesh.vol
             self._Wz = Utils.sdiag((Ave_z_vol*self.alpha_z)**0.5)*self.mesh.cellGradz
+
+            if self.indActive is not None:
+                indActive_Fz = (self.mesh.aveFz2CC.T * self.indActive) == 1
+                Pac = Utils.speye(self.mesh.nC)[:,self.indActive]
+                Pafz = Utils.speye(self.mesh.nFz)[:,indActive_Fz]
+                self._Wz = Pafz.T*self._Wz*Pac
+
         return self._Wz
 
     @property
@@ -165,6 +189,11 @@ class Tikhonov(BaseRegularization):
         """Regularization matrix Wxx"""
         if getattr(self, '_Wxx', None) is None:
             self._Wxx = Utils.sdiag((self.mesh.vol*self.alpha_xx)**0.5)*self.mesh.faceDivx*self.mesh.cellGradx
+
+            if self.indActive is not None:
+                Pac = Utils.speye(self.mesh.nC)[:,self.indActive]
+                self._Wxx = Pac.T*self._Wxx*Pac
+
         return self._Wxx
 
     @property
@@ -172,6 +201,11 @@ class Tikhonov(BaseRegularization):
         """Regularization matrix Wyy"""
         if getattr(self, '_Wyy', None) is None:
             self._Wyy = Utils.sdiag((self.mesh.vol*self.alpha_yy)**0.5)*self.mesh.faceDivy*self.mesh.cellGrady
+
+            if self.indActive is not None:
+                Pac = Utils.speye(self.mesh.nC)[:,self.indActive]
+                self._Wyy = Pac.T*self._Wyy*Pac
+
         return self._Wyy
 
     @property
@@ -179,6 +213,11 @@ class Tikhonov(BaseRegularization):
         """Regularization matrix Wzz"""
         if getattr(self, '_Wzz', None) is None:
             self._Wzz = Utils.sdiag((self.mesh.vol*self.alpha_zz)**0.5)*self.mesh.faceDivz*self.mesh.cellGradz
+
+            if self.indActive is not None:
+                Pac = Utils.speye(self.mesh.nC)[:,self.indActive]
+                self._Wzz = Pac.T*self._Wzz*Pac
+
         return self._Wzz
 
     @property
@@ -461,4 +500,4 @@ class SparseRegularization(Simple):
         # Scaling to assure equal contribution of all norms
         eta = (self.eps**(1-p/2))**0.5
         R = Utils.sdiag( eta / ((self.mapping *((G*m)**2+self.eps**2)**(1-p/2) ) )**0.5 )
-        return R
+        return R

SimPEG/Utils/SolverUtils.py

@@ -26,7 +26,14 @@ def SolverWrapD(fun, factorize=True, checkAccuracy=True, accuracyTol=1e-6):
 
     def __init__(self, A, **kwargs):
         self.A = A.tocsc()
+
+        self.checkAccuracy = kwargs.get("checkAccuracy", checkAccuracy)
+        if kwargs.has_key("checkAccuracy"): del kwargs["checkAccuracy"]
+        self.accuracyTol = kwargs.get("accuracyTol", accuracyTol)
+        if kwargs.has_key("accuracyTol"): del kwargs["accuracyTol"]
+
         self.kwargs = kwargs
+
         if factorize:
             self.solver = fun(self.A, **kwargs)
 
@@ -57,8 +64,8 @@ def SolverWrapD(fun, factorize=True, checkAccuracy=True, accuracyTol=1e-6):
                 else:
                     X[:,i] = fun(self.A, b[:,i], **self.kwargs)
 
-        if checkAccuracy:
-            _checkAccuracy(self.A, b, X, accuracyTol)
+        if self.checkAccuracy:
+            _checkAccuracy(self.A, b, X, self.accuracyTol)
         return X
 
     def clean(self):
@@ -81,6 +88,12 @@ def SolverWrapI(fun, checkAccuracy=True, accuracyTol=1e-5):
 
     def __init__(self, A, **kwargs):
         self.A = A
+
+        self.checkAccuracy = kwargs.get("checkAccuracy", checkAccuracy)
+        if kwargs.has_key("checkAccuracy"): del kwargs["checkAccuracy"]
+        self.accuracyTol = kwargs.get("accuracyTol", accuracyTol)
+        if kwargs.has_key("accuracyTol"): del kwargs["accuracyTol"]
+
         self.kwargs = kwargs
 
     def __mul__(self, b):
@@ -108,8 +121,8 @@ def SolverWrapI(fun, checkAccuracy=True, accuracyTol=1e-5):
                 else:
                     X[:,i] = out
 
-        if checkAccuracy:
-            _checkAccuracy(self.A, b, X, accuracyTol)
+        if self.checkAccuracy:
+            _checkAccuracy(self.A, b, X, self.accuracyTol)
         return X
 
     def clean(self):

SimPEG/Utils/__init__.py

@@ -1,6 +1,6 @@
 from matutils import *
 from codeutils import *
-from meshutils import exampleLrmGrid, meshTensor, closestPoints, readUBCTensorMesh, writeUBCTensorMesh, writeUBCTensorModel, readVTRFile, writeVTRFile
+from meshutils import *
 from curvutils import volTetra, faceInfo, indexCube
 from interputils import interpmat
 from CounterUtils import *

SimPEG/Utils/codeutils.py

@@ -17,7 +17,7 @@ def memProfileWrapper(towrap, *funNames):
 
         For example::
 
-            foo_mem = memProfile(foo,'my_func')
+            foo_mem = memProfileWrapper(foo,['my_func'])
             fooi = foo_mem()
             for i in range(5):
                 fooi.my_func()

SimPEG/Utils/meshutils.py

@@ -102,223 +102,6 @@ def closestPoints(mesh, pts, gridLoc='CC'):
 
     return nodeInds
 
-def readUBCTensorMesh(fileName):
-    """
-        Read UBC GIF 3DTensor mesh and generate 3D Tensor mesh in simpegTD
-
-        Input:
-        :param fileName, path to the UBC GIF mesh file
-
-        Output:
-        :param SimPEG TensorMesh object
-        :return
-    """
-
-    # Interal function to read cell size lines for the UBC mesh files.
-    def readCellLine(line):
-        for seg in line.split():
-            if '*' in seg:
-                st = seg
-                sp = seg.split('*')
-                re = np.array(sp[0],dtype=int)*(' ' + sp[1])
-                line = line.replace(st,re.strip())
-        return np.array(line.split(),dtype=float)
-
-    # Read the file as line strings, remove lines with comment = !
-    msh = np.genfromtxt(fileName,delimiter='\n',dtype=np.str,comments='!')
-
-    # Fist line is the size of the model
-    sizeM = np.array(msh[0].split(),dtype=float)
-    # Second line is the South-West-Top corner coordinates.
-    x0 = np.array(msh[1].split(),dtype=float)
-    # Read the cell sizes
-    h1 = readCellLine(msh[2])
-    h2 = readCellLine(msh[3])
-    h3temp = readCellLine(msh[4])
-    h3 = h3temp[::-1] # Invert the indexing of the vector to start from the bottom.
-    # Adjust the reference point to the bottom south west corner
-    x0[2] = x0[2] - np.sum(h3)
-    # Make the mesh
-    from SimPEG import Mesh
-    tensMsh = Mesh.TensorMesh([h1,h2,h3],x0)
-    return tensMsh
-
-def readUBCTensorModel(fileName, mesh):
-    """
-        Read UBC 3DTensor mesh model and generate 3D Tensor mesh model in simpeg
-
-        Input:
-        :param fileName, path to the UBC GIF mesh file to read
-        :param mesh, TensorMesh object, mesh that coresponds to the model
-
-        Output:
-        :return numpy array, model with TensorMesh ordered
-    """
-    f = open(fileName, 'r')
-    model = np.array(map(float, f.readlines()))
-    f.close()
-    model = np.reshape(model, (mesh.nCz, mesh.nCx, mesh.nCy), order = 'F')
-    model = model[::-1,:,:]
-    model = np.transpose(model, (1, 2, 0))
-    model = mkvc(model)
-
-    return model
-
-def writeUBCTensorMesh(fileName, mesh):
-    """
-        Writes a SimPEG TensorMesh to a UBC-GIF format mesh file.
-
-        :param str fileName: File to write to
-        :param simpeg.Mesh.TensorMesh mesh: The mesh
-
-    """
-    assert mesh.dim == 3
-    s = ''
-    s += '%i %i %i\n' %tuple(mesh.vnC)
-    origin = mesh.x0 + np.array([0,0,mesh.hz.sum()]) # Have to it in the same operation or use mesh.x0.copy(), otherwise the mesh.x0 is updated.
-    origin.dtype = float
-
-    s += '%.2f %.2f %.2f\n' %tuple(origin)
-    s += ('%.2f '*mesh.nCx+'\n')%tuple(mesh.hx)
-    s += ('%.2f '*mesh.nCy+'\n')%tuple(mesh.hy)
-    s += ('%.2f '*mesh.nCz+'\n')%tuple(mesh.hz[::-1])
-    f = open(fileName, 'w')
-    f.write(s)
-    f.close()
-
-def writeUBCTensorModel(fileName, mesh, model):
-    """
-        Writes a model associated with a SimPEG TensorMesh
-        to a UBC-GIF format model file.
-
-        :param str fileName: File to write to
-        :param simpeg.Mesh.TensorMesh mesh: The mesh
-        :param numpy.ndarray model: The model
-    """
-
-    # Reshape model to a matrix
-    modelMat = mesh.r(model,'CC','CC','M')
-    # Transpose the axes
-    modelMatT = modelMat.transpose((2,0,1))
-    # Flip z to positive down
-    modelMatTR = mkvc(modelMatT[::-1,:,:])
-
-    np.savetxt(fileName, modelMatTR.ravel())
-
-
-def readVTRFile(fileName):
-    """
-        Read VTK Rectilinear (vtr xml file) and return SimPEG Tensor mesh and model
-
-        Input:
-        :param vtrFileName, path to the vtr model file to write to
-
-        Output:
-        :return SimPEG TensorMesh object
-        :return SimPEG model dictionary
-
-    """
-    # Import
-    from vtk import vtkXMLRectilinearGridReader as vtrFileReader
-    from vtk.util.numpy_support import vtk_to_numpy
-
-    # Read the file
-    vtrReader = vtrFileReader()
-    vtrReader.SetFileName(fileName)
-    vtrReader.Update()
-    vtrGrid = vtrReader.GetOutput()
-    # Sort information
-    hx = np.abs(np.diff(vtk_to_numpy(vtrGrid.GetXCoordinates())))
-    xR = vtk_to_numpy(vtrGrid.GetXCoordinates())[0]
-    hy = np.abs(np.diff(vtk_to_numpy(vtrGrid.GetYCoordinates())))
-    yR = vtk_to_numpy(vtrGrid.GetYCoordinates())[0]
-    zD = np.diff(vtk_to_numpy(vtrGrid.GetZCoordinates()))
-    # Check the direction of hz
-    if np.all(zD < 0):
-        hz = np.abs(zD[::-1])
-        zR = vtk_to_numpy(vtrGrid.GetZCoordinates())[-1]
-    else:
-        hz = np.abs(zD)
-        zR = vtk_to_numpy(vtrGrid.GetZCoordinates())[0]
-    x0 = np.array([xR,yR,zR])
-
-    # Make the SimPEG object
-    from SimPEG import Mesh
-    tensMsh = Mesh.TensorMesh([hx,hy,hz],x0)
-
-    # Grap the models
-    modelDict = {}
-    for i in np.arange(vtrGrid.GetCellData().GetNumberOfArrays()):
-        modelName = vtrGrid.GetCellData().GetArrayName(i)
-        if np.all(zD < 0):
-            modFlip = vtk_to_numpy(vtrGrid.GetCellData().GetArray(i))
-            tM = tensMsh.r(modFlip,'CC','CC','M')
-            modArr = tensMsh.r(tM[:,:,::-1],'CC','CC','V')
-        else:
-            modArr = vtk_to_numpy(vtrGrid.GetCellData().GetArray(i))
-        modelDict[modelName] = modArr
-
-    # Return the data
-    return tensMsh, modelDict
-
-def writeVTRFile(fileName,mesh,model=None):
-    """
-    Makes and saves a VTK rectilinear file (vtr) for a simpeg Tensor mesh and model.
-
-    Input:
-    :param str, path to the output vtk file
-    :param mesh, SimPEG TensorMesh object - mesh to be transfer to VTK
-    :param model, dictionary of numpy.array - Name('s) and array('s). Match number of cells
-
-    """
-    # Import
-    from vtk import vtkRectilinearGrid as rectGrid, vtkXMLRectilinearGridWriter as rectWriter
-    from vtk.util.numpy_support import numpy_to_vtk
-
-    # Deal with dimensionalities
-    if mesh.dim >= 1:
-        vX = mesh.vectorNx
-        xD = mesh.nNx
-        yD,zD = 1,1
-        vY, vZ = np.array([0,0])
-    if mesh.dim >= 2:
-        vY = mesh.vectorNy
-        yD = mesh.nNy
-    if mesh.dim == 3:
-        vZ = mesh.vectorNz
-        zD = mesh.nNz
-    # Use rectilinear VTK grid.
-    # Assign the spatial information.
-    vtkObj = rectGrid()
-    vtkObj.SetDimensions(xD,yD,zD)
-    vtkObj.SetXCoordinates(numpy_to_vtk(vX,deep=1))
-    vtkObj.SetYCoordinates(numpy_to_vtk(vY,deep=1))
-    vtkObj.SetZCoordinates(numpy_to_vtk(vZ,deep=1))
-
-    # Assign the model('s) to the object
-    for item in model.iteritems():
-        # Convert numpy array
-        vtkDoubleArr = numpy_to_vtk(item[1],deep=1)
-        vtkDoubleArr.SetName(item[0])
-        vtkObj.GetCellData().AddArray(vtkDoubleArr)
-    # Set the active scalar
-    vtkObj.GetCellData().SetActiveScalars(model.keys()[0])
-    vtkObj.Update()
-
-
-    # Check the extension of the fileName
-    ext = os.path.splitext(fileName)[1]
-    if ext is '':
-        fileName = fileName + '.vtr'
-    elif ext not in '.vtr':
-        raise IOError('{:s} is an incorrect extension, has to be .vtr')
-    # Write the file.
-    vtrWriteFilter = rectWriter()
-    vtrWriteFilter.SetInput(vtkObj)
-    vtrWriteFilter.SetFileName(fileName)
-    vtrWriteFilter.Update()
-
-
 def ExtractCoreMesh(xyzlim, mesh, meshType='tensor'):
     """
         Extracts Core Mesh from Global mesh

docs/flow/index.rst

@@ -41,7 +41,7 @@ Here we reproduce the results from Celia et al. (1990):
 Richards
 ========
 
-.. automodule:: simpegFLOW.Richards.Empirical
+.. automodule:: SimPEG.FLOW.Richards.Empirical
     :show-inheritance:
     :members:
     :undoc-members:

tests/base/test_regularization.py

@@ -4,11 +4,17 @@ from SimPEG import *
 from scipy.sparse.linalg import dsolve
 import inspect
 
+TOL = 1e-20
 
 class RegularizationTests(unittest.TestCase):
 
     def setUp(self):
-        self.mesh2 = Mesh.TensorMesh([3, 2])
+        hx, hy, hz = np.random.rand(10), np.random.rand(9), np.random.rand(8)
+        hx, hy, hz = hx/hx.sum(), hy/hy.sum(), hz/hz.sum()
+        mesh1 = Mesh.TensorMesh([hx])
+        mesh2 = Mesh.TensorMesh([hx, hy])
+        mesh3 = Mesh.TensorMesh([hx, hy, hz])
+        self.meshlist = [mesh1,mesh2, mesh3]
 
     def test_regularization(self):
         for R in dir(Regularization):
@@ -16,18 +22,63 @@ class RegularizationTests(unittest.TestCase):
             if not inspect.isclass(r): continue
             if not issubclass(r, Regularization.BaseRegularization):
                 continue
-            # if 'Regularization' not in R: continue
-            mapping = r.mapPair(self.mesh2)
-            reg = r(self.mesh2, mapping=mapping)
-            m = np.random.rand(mapping.nP)
-            reg.mref = m[:]*np.mean(m)
 
-            print 'Check:', R
-            passed = Tests.checkDerivative(lambda m : [reg.eval(m), reg.evalDeriv(m)], m, plotIt=False)
-            self.assertTrue(passed)
-            print 'Check 2 Deriv:', R
-            passed = Tests.checkDerivative(lambda m : [reg.evalDeriv(m), reg.eval2Deriv(m)], m, plotIt=False)
-            self.assertTrue(passed)
+            for i, mesh in enumerate(self.meshlist):
+
+                print 'Testing %iD'%mesh.dim
+
+                mapping = r.mapPair(mesh)
+                reg = r(mesh, mapping=mapping)
+                m = np.random.rand(mapping.nP)
+                reg.mref = np.ones_like(m)*np.mean(m)
+
+                print 'Check: phi_m (mref) = %f' %reg.eval(reg.mref)
+                passed = reg.eval(reg.mref) < TOL
+                self.assertTrue(passed)
+
+                print 'Check:', R
+                passed = Tests.checkDerivative(lambda m : [reg.eval(m), reg.evalDeriv(m)], m, plotIt=False)
+                self.assertTrue(passed)
+
+                print 'Check 2 Deriv:', R
+                passed = Tests.checkDerivative(lambda m : [reg.evalDeriv(m), reg.eval2Deriv(m)], m, plotIt=False)
+                self.assertTrue(passed)
+
+    def test_regularization_ActiveCells(self):
+        for R in dir(Regularization):
+            r = getattr(Regularization, R)
+            if not inspect.isclass(r): continue
+            if not issubclass(r, Regularization.BaseRegularization):
+                continue
+
+            for i, mesh in enumerate(self.meshlist):
+
+                print 'Testing Active Cells %iD'%(mesh.dim)
+
+                if mesh.dim == 1:
+                    indAct = Utils.mkvc(mesh.gridCC <= 0.8)
+                elif mesh.dim == 2:
+                    indAct = Utils.mkvc(mesh.gridCC[:,-1] <= 2*np.sin(2*np.pi*mesh.gridCC[:,0])+0.5)
+                elif mesh.dim == 3:
+                    indAct = Utils.mkvc(mesh.gridCC[:,-1] <= 2*np.sin(2*np.pi*mesh.gridCC[:,0])+0.5 * 2*np.sin(2*np.pi*mesh.gridCC[:,1])+0.5)
+
+                mapping = Maps.IdentityMap(nP=indAct.nonzero()[0].size)
+
+                reg = r(mesh, mapping=mapping, indActive=indAct)
+                m = np.random.rand(mesh.nC)[indAct]
+                reg.mref = np.ones_like(m)*np.mean(m)
+
+                print 'Check: phi_m (mref) = %f' %reg.eval(reg.mref)
+                passed = reg.eval(reg.mref) < TOL
+                self.assertTrue(passed)
+
+                print 'Check:', R
+                passed = Tests.checkDerivative(lambda m : [reg.eval(m), reg.evalDeriv(m)], m, plotIt=False)
+                self.assertTrue(passed)
+
+                print 'Check 2 Deriv:', R
+                passed = Tests.checkDerivative(lambda m : [reg.evalDeriv(m), reg.eval2Deriv(m)], m, plotIt=False)
+                self.assertTrue(passed)
 
 
 if __name__ == '__main__':

tests/mesh/test_MeshIO.py

@@ -0,0 +1,100 @@
+import numpy as np
+import unittest, os
+import SimPEG as simpeg
+from SimPEG.Mesh import TensorMesh, TreeMesh
+
+
+class TestTensorMeshIO(unittest.TestCase):
+
+    def setUp(self):
+        h = np.ones(16)
+        mesh = TensorMesh([h,2*h,3*h])
+        self.mesh = mesh
+
+    def test_UBCfiles(self):
+
+        mesh = self.mesh
+        # Make a vector
+        vec = np.arange(mesh.nC)
+        # Write and read
+        mesh.writeUBC('temp.msh', {'arange.txt':vec})
+        meshUBC = TensorMesh.readUBC('temp.msh')
+        vecUBC  = meshUBC.readModelUBC('arange.txt')
+
+        # The mesh
+        assert mesh.__str__() == meshUBC.__str__()
+        assert np.sum(mesh.gridCC - meshUBC.gridCC) == 0
+        assert np.sum(vec - vecUBC) == 0
+        assert np.all(np.array(mesh.h) - np.array(meshUBC.h) == 0)
+
+
+        vecUBC  = mesh.readModelUBC('arange.txt')
+        assert np.sum(vec - vecUBC) == 0
+
+        mesh.writeModelUBC('arange2.txt', vec + 1)
+        vec2UBC  = mesh.readModelUBC('arange2.txt')
+        assert np.sum(vec + 1 - vec2UBC) == 0
+
+        print 'IO of UBC tensor mesh files is working'
+        os.remove('temp.msh')
+        os.remove('arange.txt')
+        os.remove('arange2.txt')
+
+    def test_VTKfiles(self):
+        mesh = self.mesh
+        vec = np.arange(mesh.nC)
+
+        mesh.writeVTK('temp.vtr', {'arange.txt':vec})
+        meshVTR, models = TensorMesh.readVTK('temp.vtr')
+
+        assert mesh.__str__() == meshVTR.__str__()
+        assert np.all(np.array(mesh.h) - np.array(meshVTR.h) == 0)
+
+        assert 'arange.txt' in models
+        vecVTK = models['arange.txt']
+        assert np.sum(vec - vecVTK) == 0
+
+        print 'IO of VTR tensor mesh files is working'
+        os.remove('temp.vtr')
+
+
+class TestOcTreeMeshIO(unittest.TestCase):
+
+    def setUp(self):
+        h = np.ones(16)
+        mesh = TreeMesh([h,2*h,3*h])
+        mesh.refine(3)
+        mesh._refineCell([0,0,0,3])
+        mesh._refineCell([0,2,0,3])
+        self.mesh = mesh
+
+    def test_UBCfiles(self):
+
+        mesh = self.mesh
+        # Make a vector
+        vec = np.arange(mesh.nC)
+        # Write and read
+        mesh.writeUBC('temp.msh', {'arange.txt':vec})
+        meshUBC = TreeMesh.readUBC('temp.msh')
+        vecUBC  = meshUBC.readModelUBC('arange.txt')
+
+        # The mesh
+        assert mesh.__str__() == meshUBC.__str__()
+        assert np.sum(mesh.gridCC - meshUBC.gridCC) == 0
+        assert np.sum(vec - vecUBC) == 0
+        assert np.all(np.array(mesh.h) - np.array(meshUBC.h) == 0)
+        print 'IO of UBC octree files is working'
+        os.remove('temp.msh')
+        os.remove('arange.txt')
+
+    def test_VTUfiles(self):
+        mesh = self.mesh
+        vec = np.arange(mesh.nC)
+        mesh.writeVTK('temp.vtu',{'arange':vec})
+        print 'Writing of VTU files is working'
+        os.remove('temp.vtu')
+
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/mesh/test_TreeMesh.py

@@ -26,6 +26,27 @@ class TestSimpleQuadTree(unittest.TestCase):
 
         assert np.allclose(np.r_[M._areaFxFull, M._areaFyFull], M._deflationMatrix('F') * M.area)
 
+    def test_getitem(self):
+        M = Mesh.TreeMesh([4,4])
+        M.refine(1)
+        assert M.nC == 4
+        assert len(M) == M.nC
+        assert np.allclose(M[0].center, [0.25,0.25])
+        actual = [[0,0],[0.5,0],[0,0.5],[0.5,0.5]]
+        for i, n in enumerate(M[0].nodes):
+            assert np.allclose(M._gridN[n,:], actual[i])
+
+    def test_getitem3D(self):
+        M = Mesh.TreeMesh([4,4,4])
+        M.refine(1)
+        assert M.nC == 8
+        assert len(M) == M.nC
+        assert np.allclose(M[0].center, [0.25,0.25,0.25])
+        actual = [[0,0,0],[0.5,0,0],[0,0.5,0],[0.5,0.5,0],
+                  [0,0,0.5],[0.5,0,0.5],[0,0.5,0.5],[0.5,0.5,0.5]]
+        for i, n in enumerate(M[0].nodes):
+            assert np.allclose(M._gridN[n,:], actual[i])
+
     def test_refine(self):
         M = Mesh.TreeMesh([4,4,4])
         M.refine(1)