Address mesh IO #212

tests/mesh/test_MeshIO.py

@@ -4,11 +4,65 @@ import SimPEG as simpeg
 from SimPEG.Mesh import TensorMesh, TreeMesh
 
 
-class TestOcTreeIO(unittest.TestCase):
+class TestTensorMeshIO(unittest.TestCase):
 
     def setUp(self):
         h = np.ones(16)
-        mesh = simpeg.Mesh.TreeMesh([h,2*h,3*h])
+        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])
@@ -19,9 +73,10 @@ class TestOcTreeIO(unittest.TestCase):
         mesh = self.mesh
         # Make a vector
         vec = np.arange(mesh.nC)
-        # Write aand read
-        simpeg.Utils.meshutils.writeUBCocTreeFiles('temp.msh',mesh,{'arange.txt':vec})
-        meshUBC, vecUBC = simpeg.Utils.meshutils.readUBCocTreeFiles('temp.msh',['arange.txt'])
+        # 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__()
@@ -35,7 +90,7 @@ class TestOcTreeIO(unittest.TestCase):
     def test_VTUfiles(self):
         mesh = self.mesh
         vec = np.arange(mesh.nC)
-        simpeg.Utils.meshutils.writeVTUFile('temp.vtu',mesh,{'arange':vec})
+        mesh.writeVTK('temp.vtu',{'arange':vec})
         print 'Writing of VTU files is working'
         os.remove('temp.vtu')