add IO functionalities to write DCIP2D mesh and models

discretize/MeshIO.py

@@ -300,7 +300,56 @@ 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)
+
+        temp = np.array(obsfile[1].split(), dtype=float)
+
+        if len(temp) > 1:
+            model = np.zeros((dim[0], dim[1]))
+
+            for ii in range(len(obsfile)-1):
+                mm = np.array(obsfile[ii+1].split(), dtype=float)
+                model[:, ii] = mm
+
+            model = model[:, ::-1]
+
+        else:
+
+            if len(obsfile[1:]) == 1:
+                mm = np.array(obsfile[1:].split(), dtype=float)
+
+            else:
+                mm = np.array(obsfile[1:], dtype=float)
+
+            # Permute the second dimension to flip the order
+            model = mm.reshape(dim[1], dim[0])
+
+            model = model[::-1, :]
+            model = np.transpose(model, (1, 0))
+
+        model = utils.mkvc(model)
+
+        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,53 +359,156 @@ 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 readModelUBC(mesh, fileName):
+        """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
+        """
+        if mesh.dim == 3:
+            model = mesh._readModelUBC_3D(fileName)
+        elif mesh.dim == 2:
+            model = mesh._readModelUBC_2D(fileName)
+        else:
+            raise Exception('mesh must be a Tensor Mesh 2D or 3D')
+        return model
+
     def writeModelUBC(mesh, fileName, model):
         """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
         """
+        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(fileName, modelMatTR.ravel())
+
+        elif mesh.dim == 2:
+            modelMat = mesh.r(model, 'CC', 'CC', 'M').T[::-1]
+            f = open(fileName, 'w')
+            f.write('{:d} {:d}\n'.format(mesh.nCx, mesh.nCy))
+            f.close()
+            f = open(fileName, '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, 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
+        :param str comment_lines: comment lines preceded with '!' to add
+        """
+        if mesh.dim == 3:
+            mesh._writeUBC_3DMesh(fileName, comment_lines=comment_lines)
+        elif mesh.dim == 2:
+            mesh._writeUBC_2DMesh(fileName, 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'