Work on MT, fixing bugs

SimPEG/MT/SrcMT.py

@@ -28,7 +28,7 @@ class BaseMTSrc(FDEMBaseSrc):
     def __init__(self, rxList, freq):
 
         self.freq = float(freq)
-        Survey.BaseSrc.__init__(self, rxList)
+        FDEMBaseSrc.__init__(self, rxList)
 
 # 1D sources
 class polxy_1DhomotD(BaseMTSrc):

SimPEG/MT/SurveyMT.py

@@ -5,7 +5,7 @@ from scipy.constants import mu_0
 from numpy.lib import recfunctions as recFunc
 from Sources import homo1DModelSource
 from Utils import rec2ndarr
-
+import SrcMT
 import sys
 
 #################
@@ -322,7 +322,6 @@ class Survey(SimPEGsurvey.BaseSurvey):
         :param list srcList: List of sources associated with the survey
 
     """
-    import SrcMT
     srcPair = SrcMT.BaseMTSrc
 
     def __init__(self, srcList, **kwargs):

SimPEG/Utils/meshutils.py

@@ -303,7 +303,6 @@ def writeVTRFile(fileName,mesh,model=None):
         vtkObj.GetCellData().AddArray(vtkDoubleArr)
     # Set the active scalar
     vtkObj.GetCellData().SetActiveScalars(model.keys()[0])
-    vtkObj.Update()
 
 
     # Check the extension of the fileName
@@ -314,7 +313,7 @@ def writeVTRFile(fileName,mesh,model=None):
         raise IOError('{:s} is an incorrect extension, has to be .vtr')
     # Write the file.
     vtrWriteFilter = rectWriter()
-    vtrWriteFilter.SetInput(vtkObj)
+    vtrWriteFilter.SetInputData(vtkObj)
     vtrWriteFilter.SetFileName(fileName)
     vtrWriteFilter.Update()
 

tests/mesh/test_MeshIO.py

@@ -0,0 +1,50 @@
+import numpy as np
+import unittest, os
+import SimPEG as simpeg
+from SimPEG.Mesh import TensorMesh, TreeMesh
+
+
+class TestOcTreeIO(unittest.TestCase):
+
+    def setUp(self):
+        h = np.ones(16)
+        mesh = simpeg.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 aand read
+        simpeg.Utils.meshutils.writeUBCocTreeFiles('temp.msh',mesh,{'arange.txt':vec})
+        meshUBC, vecUBC = simpeg.Utils.meshutils.readUBCocTreeFiles('temp.msh',['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)
+        try:
+            simpeg.Utils.meshutils.writeVTUFile('temp.vtu',mesh,{'arange':vec})
+            run = True
+        except:
+            run = False
+        assert run
+        print 'Writing of VTU files is working'
+        os.remove('temp.vtu')
+
+
+
+if __name__ == '__main__':
+    unittest.main()