readUBC: loading either 2D or 3D mesh

discretize/MeshIO.py

@@ -7,8 +7,8 @@
 class TensorMeshIO(object):
 
     @classmethod
-    def readUBC(TensorMesh, fileName):
-        """Read UBC GIF 3D tensor mesh and generate 3D TensorMesh.
+    def _readUBC_3DMesh(TensorMesh, fileName):
+        """Read UBC GIF 3D tensor mesh and generate same dimension TensorMesh.
 
         :param string fileName: path to the UBC GIF mesh file
         :rtype: TensorMesh
@@ -23,13 +23,12 @@ def readCellLine(line):
                     sp = seg.split('*')
                     seg_arr = np.ones((int(sp[0]),)) * float(sp[1])
                 else:
-                    seg_arr = np.array([float(seg)],float)
+                    seg_arr = np.array([float(seg)], float)
                 line_list.append(seg_arr)
             return np.concatenate(line_list)
 
         # 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.
@@ -38,13 +37,97 @@ def readCellLine(line):
         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.
+        # Invert the indexing of the vector to start from the bottom.
+        h3 = h3temp[::-1]
         # 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 _readUBC_2DMesh(TensorMesh, fileName):
+        """Read UBC GIF 2DTensor mesh and generate 2D Tensor mesh in simpeg
+
+        :param string fileName: path to the UBC GIF mesh file
+        :rtype: TensorMesh
+        :return: SimPEG TensorMesh 2D object
+        """
+
+        fopen = open(fileName, 'r')
+
+        # Read down the file and unpack dx vector
+        def unpackdx(fid, nrows):
+            for ii in range(nrows):
+                line = fid.readline()
+                var = np.array(line.split(), dtype=float)
+                if ii == 0:
+                    x0 = var[0]
+                    xvec = np.ones(int(var[2])) * (var[1] - var[0]) / int(var[2])
+                    xend = var[1]
+                else:
+                    xvec = np.hstack((xvec, np.ones(int(var[1])) * (var[0] - xend) / int(var[1])))
+                    xend = var[0]
+            return x0, xvec
+
+        # Start with dx block
+        # First line specifies the number of rows for x-cells
+        line = fopen.readline()
+        # Strip comments lines
+        while line.startswith("!"):
+            line = fopen.readline()
+        nl = np.array(line.split(), dtype=int)
+        [x0, dx] = unpackdx(fopen, nl[0])
+        # Move down the file until reaching the z-block
+        line = fopen.readline()
+        if not line:
+            line = fopen.readline()
+        # End with dz block
+        # First line specifies the number of rows for z-cells
+        line = fopen.readline()
+        nl = np.array(line.split(), dtype=int)
+        [z0, dz] = unpackdx(fopen, nl[0])
+        # Flip z0 to be the bottom of the mesh for SimPEG
+        z0 = z0 - sum(dz)
+        dz = dz[::-1]
+        # Make the mesh
+        tensMsh = TensorMesh([dx, dz], (x0, z0))
+
+        fopen.close()
+
+        return tensMsh
+
+    @classmethod
+    def readUBC(TensorMesh, fileName, meshdim=None):
+        """Wrapper to Read UBC GIF 2D  and 3D tensor mesh and generate same dimension TensorMesh.
+
+        :param string fileName: path to the UBC GIF mesh file
+        :param int meshdim: expected dimension of the mesh, if unknown the default argument is None
+        :rtype: TensorMesh
+        :return: The tensor mesh for the fileName.
+        """
+        # Check the expected mesh dimensions
+        if meshdim == None:
+            # Read the file as line strings, remove lines with comment = !
+            msh = np.genfromtxt(fileName, delimiter='\n', dtype=np.str, comments='!', max_rows=1)
+            # Fist line is the size of the model
+            sizeM = np.array(msh.ravel()[0].split(), dtype=float)
+            # Check if the mesh is a UBC 2D mesh
+            if sizeM.shape[0] == 1:
+                Tnsmsh = TensorMesh._readUBC_2DMesh(fileName)
+            # Check if the mesh is a UBC 3D mesh
+            elif sizeM.shape[0] == 3:
+                Tnsmsh = TensorMesh._readUBC_3DMesh(fileName)
+            else:
+                raise Exception('File format not recognized')
+        # expected dimension is 2
+        elif meshdim == 2:
+            Tnsmsh = TensorMesh._readUBC_2DMesh(fileName)
+        # expected dimension is 3
+        elif meshdim == 3:
+            Tnsmsh = TensorMesh._readUBC_3DMesh(fileName)
+        return Tnsmsh
+
     @classmethod
     def readVTK(TensorMesh, fileName):
         """Read VTK Rectilinear (vtr xml file) and return Tensor mesh and model
@@ -245,10 +328,10 @@ def writeUBC(mesh, fileName, models=None):
         origin = mesh.x0 + np.array([0, 0, mesh.hz.sum()])
         origin.dtype = float
 
-        s += '{0:.2f} {1:.2f} {2:.2f}\n'.format(*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])
+        s += '{0:.6f} {1:.6f} {2:.6f}\n'.format(*tuple(origin))
+        s += ('%.6f '*mesh.nCx+'\n')%tuple(mesh.hx)
+        s += ('%.6f '*mesh.nCy+'\n')%tuple(mesh.hy)
+        s += ('%.6f '*mesh.nCz+'\n')%tuple(mesh.hz[::-1])
         f = open(fileName, 'w')
         f.write(s)
         f.close()
@@ -439,7 +522,7 @@ def writeVTK(mesh, fileName, models=None):
         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))
+        cellsArr.SetCells(cellConn.shape[0], numpy_to_vtkIdTypeArray(cellsMat, deep =True))
 
         # Make the object
         vtuObj = vtk.vtkUnstructuredGrid()

tests/base/test_tensor_io.py

@@ -81,6 +81,18 @@ def test_VTKfiles(self):
             print('IO of VTR tensor mesh files is working')
             os.remove('temp.vtr')
 
+    def test_read_ubc_2Dmesh(self):
+        fname = os.path.join(os.path.split(__file__)[0], 'ubc_DC2D_tensor_mesh.msh')
+        mesh = discretize.TensorMesh.readUBC(fname)
+        assert mesh.nCx == 178
+        assert mesh.nCy == 67
+        # spot check a few things in the file
+        assert mesh.hx[0] == 600.
+        # The x0 is in a different place (-z)
+        assert mesh.x0[-1] == - np.sum(mesh.hy)
+        # the z axis is flipped
+        assert mesh.hy[0] == 600.
+        assert mesh.hy[-1] == 10.
 
 if __name__ == '__main__':
     unittest.main()

tests/base/ubc_DC2D_tensor_mesh.msh

@@ -0,0 +1,35 @@
+! comment line
+! again
+! and again
+! and again and again
+19
+-2000 -1400 1
+-1000 1
+-700 1
+-500 1
+-350 1
+-250 1
+-200 1
+-150 1
+-125 1
+2675 160
+2700 1
+2750 1
+2800 1
+2900 1
+3050 1
+3250 1
+3550 1
+3950 1
+4550 1
+
+9
+-0 400 40
+800 20
+850 1
+950 1
+1100 1
+1300 1
+1600 1
+2000 1
+2600 1

tests/base/ubc_tensor_mesh.msh

@@ -1,3 +1,7 @@
+! comment line
+! again
+! and again
+! and again and again
 78 50 51
 -200000.000000 -200000.000000 3000.000000
 55000 45000 40000 35000 25000 10*20000 3*10000 43*10000 2*10000 10*20000 25000 35000 40000 45000 55000