add IO functionalities to write DCIP2D mesh and models

discretize/MeshIO.py

@@ -118,7 +118,7 @@ def unpackdx(fid, nrows):
         return tensMsh
 
     @classmethod
-    def readUBC(TensorMesh, fileName, meshdim=None):
+    def readUBC(TensorMesh, fileName, meshdim=None, folder=''):
         """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
@@ -127,25 +127,26 @@ def readUBC(TensorMesh, fileName, meshdim=None):
         :return: The tensor mesh for the fileName.
         """
         # Check the expected mesh dimensions
+        fname = os.path.join(folder, fileName)
         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)
+            msh = np.genfromtxt(fname, 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)
+                Tnsmsh = TensorMesh._readUBC_2DMesh(fname)
             # Check if the mesh is a UBC 3D mesh
             elif sizeM.shape[0] == 3:
-                Tnsmsh = TensorMesh._readUBC_3DMesh(fileName)
+                Tnsmsh = TensorMesh._readUBC_3DMesh(fname)
             else:
                 raise Exception('File format not recognized')
         # expected dimension is 2
         elif meshdim == 2:
-            Tnsmsh = TensorMesh._readUBC_2DMesh(fileName)
+            Tnsmsh = TensorMesh._readUBC_2DMesh(fname)
         # expected dimension is 3
         elif meshdim == 3:
-            Tnsmsh = TensorMesh._readUBC_3DMesh(fileName)
+            Tnsmsh = TensorMesh._readUBC_3DMesh(fname)
         return Tnsmsh
 
     @classmethod
@@ -300,7 +301,34 @@ def _toVTRObj(mesh, models=None):
             vtkObj.GetCellData().SetActiveScalars(models.keys()[0])
         return vtkObj
 
-    def readModelUBC(mesh, fileName):
+    def _readModelUBC_2D(mesh, fileName):
+        """
+        Read UBC GIF 2DTensor model and generate 2D Tensor model in simpeg
+
+        Input:
+        :param string fileName: path to the UBC GIF 2D model file
+
+        Output:
+        :rtype: numpy.ndarray
+        :return: model with TensorMesh ordered
+        """
+
+        # Open fileand skip header... assume that we know the mesh already
+        obsfile = np.genfromtxt(
+            fileName, delimiter=' \n',
+            dtype=np.str, comments='!'
+        )
+
+        dim = np.array(obsfile[0].split(), dtype=int)
+        if not np.all([mesh.nCx, mesh.nCy] == dim):
+            raise Exception('Dimension of the model and mesh mismatch')
+
+        model = [line.split() for line in obsfile[1:]]
+        model = utils.mkvc(np.array(model, dtype=float)[::-1].T)
+
+        return model
+
+    def _readModelUBC_3D(mesh, fileName):
         """Read UBC 3DTensor mesh model and generate 3D Tensor mesh model
 
         :param string fileName: path to the UBC GIF mesh file to read
@@ -310,57 +338,164 @@ def readModelUBC(mesh, fileName):
         f = open(fileName, 'r')
         model = np.array(list(map(float, f.readlines())))
         f.close()
-        model = np.reshape(model, (mesh.nCz, mesh.nCx, mesh.nCy), order = 'F')
+        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):
+    def readModelUBC(mesh, fileName, folder=''):
+        """Read UBC 2D or 3D Tensor mesh model
+            and generate Tensor mesh model
+
+        :param string fileName: path to the UBC GIF mesh file to read
+        :rtype: numpy.ndarray
+        :return: model with TensorMesh ordered
+        """
+        fname = os.path.join(folder, fileName)
+        if mesh.dim == 3:
+            model = mesh._readModelUBC_3D(fname)
+        elif mesh.dim == 2:
+            model = mesh._readModelUBC_2D(fname)
+        else:
+            raise Exception('mesh must be a Tensor Mesh 2D or 3D')
+        return model
+
+    def writeModelUBC(mesh, fileName, model, folder=''):
         """Writes a model associated with a TensorMesh
         to a UBC-GIF format model file.
 
         :param string fileName: File to write to
         :param numpy.ndarray model: The model
         """
+        fname = os.path.join(folder, fileName)
+        if mesh.dim == 3:
+            # 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(fname, modelMatTR.ravel())
+
+        elif mesh.dim == 2:
+            modelMat = mesh.r(model, 'CC', 'CC', 'M').T[::-1]
+            f = open(fname, 'w')
+            f.write('{:d} {:d}\n'.format(mesh.nCx, mesh.nCy))
+            f.close()
+            f = open(fname, 'ab')
+            np.savetxt(f, modelMat)
+            f.close()
 
-        # 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())
+        else:
+            raise Exception('mesh must be a Tensor Mesh 2D or 3D')
 
-    def writeUBC(mesh, fileName, models=None):
+    def _writeUBC_3DMesh(mesh, fileName, comment_lines=''):
         """Writes a TensorMesh to a UBC-GIF format mesh file.
 
         :param string fileName: File to write to
         :param dict models: A dictionary of the models
 
         """
         assert mesh.dim == 3
-        s = ''
+        s = comment_lines
         s += '{0:d} {1:d} {2:d}\n'.format(*tuple(mesh.vnC))
         # Have to it in the same operation or use mesh.x0.copy(),
         # otherwise the mesh.x0 is updated.
         origin = mesh.x0 + np.array([0, 0, mesh.hz.sum()])
         origin.dtype = float
 
         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])
+        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()
 
-        if models is None: return
+    def _writeUBC_2DMesh(mesh, fileName, comment_lines=''):
+        """Writes a TensorMesh to a UBC-GIF format mesh file.
+
+        :param string fileName: File to write to
+        :param dict models: A dictionary of the models
+
+        """
+        assert mesh.dim == 2
+
+        def writeF(fx, outStr=''):
+            # Init
+            i = 0
+            origin = True
+            x0 = fx[i]
+            f = fx[i]
+            number_segment = 0
+            auxStr = ''
+
+            while True:
+                i = i + 1
+                if i >= fx.size:
+                    break
+                dx = -f + fx[i]
+                f = fx[i]
+                n = 1
+
+                for j in range(i+1, fx.size):
+                    if -f + fx[j] == dx:
+                        n += 1
+                        i += 1
+                        f = fx[j]
+                    else:
+                        break
+
+                number_segment += 1
+                if origin:
+                    auxStr += '{:.10f} {:.10f} {:d} \n'.format(x0, f, n)
+                    origin = False
+                else:
+                    auxStr += '{:.10f} {:d} \n'.format(f, n)
+
+            auxStr = '{:d}\n'.format(number_segment) + auxStr
+            outStr += auxStr
+
+            return outStr
+
+        # Grab face coordinates
+        fx = mesh.vectorNx
+        fz = -mesh.vectorNy[::-1]
+
+        # Create the string
+        outStr = comment_lines
+        outStr = writeF(fx, outStr=outStr)
+        outStr += '\n'
+        outStr = writeF(fz, outStr=outStr)
+
+        # Write file
+        f = open(fileName, 'w')
+        f.write(outStr)
+        f.close()
+
+    def writeUBC(mesh, fileName, models=None, folder='', comment_lines=''):
+        """Writes a TensorMesh to a UBC-GIF format mesh file.
+
+        :param str fileName: File to write to
+        :param str folder: folder where to save model
+        :param dict models: A dictionary of the models
+        :param str comment_lines: comment lines preceded with '!' to add
+        """
+        fname = os.path.join(folder, fileName)
+        if mesh.dim == 3:
+            mesh._writeUBC_3DMesh(fname, comment_lines=comment_lines)
+        elif mesh.dim == 2:
+            mesh._writeUBC_2DMesh(fname, comment_lines=comment_lines)
+        else:
+            raise Exception('mesh must be a Tensor Mesh 2D or 3D')
+
+        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])
+            mesh.writeModelUBC(key, models[key], folder=folder)
 
 
 class TreeMeshIO(object):

tests/base/test_tensor_io.py

@@ -18,36 +18,43 @@ def setUp(self):
         h = np.ones(16)
         mesh = discretize.TensorMesh([h, 2*h, 3*h])
         self.mesh = mesh
+        self.basePath = os.path.expanduser('~/TestIO')
 
     def test_UBCfiles(self):
+        if not os.path.exists(self.basePath):
+            os.mkdir(self.basePath)
 
         mesh = self.mesh
         # Make a vector
         vec = np.arange(mesh.nC)
         # Write and read
-        mesh.writeUBC('temp.msh', {'arange.txt': vec})
-        meshUBC = discretize.TensorMesh.readUBC('temp.msh')
-        vecUBC = meshUBC.readModelUBC('arange.txt')
+        fname = 'temp.msh'
+        modelfname = 'arange.txt'
+        mesh.writeUBC('temp.msh', {modelfname: vec}, folder=self.basePath)
+        meshUBC = discretize.TensorMesh.readUBC(fname, folder=self.basePath)
+        vecUBC = meshUBC.readModelUBC(modelfname, folder=self.basePath)
 
         # 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')
+        vecUBC = mesh.readModelUBC(modelfname, folder=self.basePath)
         assert np.sum(vec - vecUBC) == 0
 
-        mesh.writeModelUBC('arange2.txt', vec + 1)
-        vec2UBC = mesh.readModelUBC('arange2.txt')
+        modelfname1 = 'arange2.txt'
+        mesh.writeModelUBC(modelfname1, vec + 1, folder=self.basePath)
+        vec2UBC = mesh.readModelUBC(modelfname1, folder=self.basePath)
         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')
+        os.remove(os.path.join(self.basePath, fname))
+        os.remove(os.path.join(self.basePath, modelfname))
+        os.remove(os.path.join(self.basePath, modelfname1))
+        os.rmdir(self.basePath)
 
-    def test_read_ubc_mesh(self):
+    def test_write_read_ubc_mesh(self):
         fname = os.path.join(os.path.split(__file__)[0], 'ubc_tensor_mesh.msh')
         mesh = discretize.TensorMesh.readUBC(fname)
         assert mesh.nCx == 78
@@ -62,7 +69,6 @@ def test_read_ubc_mesh(self):
         assert mesh.hz[0] == 20000
         assert mesh.hz[-1] == 250
 
-
     if has_vtk:
         def test_VTKfiles(self):
             mesh = self.mesh
@@ -81,18 +87,69 @@ 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')
+    def test_write_read_ubc_2Dmesh(self):
+        if not os.path.exists(self.basePath):
+            os.mkdir(self.basePath)
+
+        fname = 'ubc_DC2D_tensor_mesh.msh'
+        # Create 2D Mesh
+        # Cells size
+        csx, csz = 0.25, 0.25
+        # Number of core cells in each direction
+        ncx,  ncz = 123,  41
+        # Number of padding cells and expansion rate
+        npad, expansion = 6, 1.5
+        # Vectors of cell lengthts in each direction
+        hx = [(csx, npad,  -expansion), (csx, ncx), (csx, npad,  expansion)]
+        hz = [(csz, npad, -expansion), (csz, ncz)]
+        mesh = discretize.TensorMesh([hx,  hz], x0="CN")
+
+        # Create 2-cylinders Model Creation
+        # Spheres parameters
+        x0,  z0,  r0 = -6.,  -5.,  3.
+        x1,  z1,  r1 = 6.,  -5.,  3.
+        background = -5.
+        sphere0 = -3.
+        sphere1 = -6.
+
+        # Create Model
+        model = background*np.ones(mesh.nC)
+        csph = (np.sqrt((mesh.gridCC[:, 1]-z0)**2.+(mesh.gridCC[:, 0]-x0)**2.)) < r0
+        model[csph] = sphere0*np.ones_like(model[csph])
+        # Define the sphere limit
+        rsph = (np.sqrt((mesh.gridCC[:, 1]-z1)**2.+(mesh.gridCC[:, 0]-x1)**2.)) < r1
+        model[rsph] = sphere1*np.ones_like(model[rsph])
+        modeldict = {'2d_2cyl_model': model}
+
+        # Write Mesh and model
+        comment_lines = '!comment line\n'+'!again\n'+'!and again\n'
+        mesh.writeUBC(
+            fname, models=modeldict,
+            folder=self.basePath, comment_lines=comment_lines
+        )
+
+        # Read back mesh and model
+        fname = os.path.sep.join([self.basePath, 'ubc_DC2D_tensor_mesh.msh'])
         mesh = discretize.TensorMesh.readUBC(fname)
-        assert mesh.nCx == 178
-        assert mesh.nCy == 67
+        modelfname = os.path.sep.join([self.basePath, '2d_2cyl_model'])
+        readmodel = mesh.readModelUBC(modelfname)
+        assert mesh.nCx == 135
+        assert mesh.nCy == 47
         # spot check a few things in the file
-        assert mesh.hx[0] == 600.
+        assert mesh.hx[0] == 2.84765625
         # 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.
+        assert mesh.hy[0] == 2.84765625
+        assert mesh.hy[-1] == csz
+        assert mesh.dim == 2
+        assert np.all(model == readmodel)
+
+        # Clean up the working directory
+        os.remove(os.path.join(self.basePath, fname))
+        os.remove(os.path.join(self.basePath, modelfname))
+        os.rmdir(self.basePath)
+
 
 if __name__ == '__main__':
     unittest.main()