simpeg/SimPEG/Tests/test_SurveyAndData.py

import unittest
from SimPEG import *

class TestData(unittest.TestCase):

    def setUp(self):
        mesh = Mesh.TensorMesh([np.ones(n)*5 for n in [10,11,12]],[0,0,-30])
        x = np.linspace(5,10,3)
        XYZ = Utils.ndgrid(x,x,np.r_[0.])
        srcLoc = np.r_[0,0,0.]
        rxList0 = Survey.BaseRx(XYZ, 'exi')
        Src0 = Survey.BaseSrc([rxList0], loc=srcLoc)
        rxList1 = Survey.BaseRx(XYZ, 'bxi')
        Src1 = Survey.BaseSrc([rxList1], loc=srcLoc)
        rxList2 = Survey.BaseRx(XYZ, 'bxi')
        Src2 = Survey.BaseSrc([rxList2], loc=srcLoc)
        rxList3 = Survey.BaseRx(XYZ, 'bxi')
        Src3 = Survey.BaseSrc([rxList3], loc=srcLoc)
        Src4 = Survey.BaseSrc([rxList0, rxList1, rxList2, rxList3], loc=srcLoc)
        srcList = [Src0,Src1,Src2,Src3,Src4]
        survey = Survey.BaseSurvey(srcList=srcList)
        self.D = Survey.Data(survey)

    def test_data(self):
        V = []
        for src in self.D.survey.srcList:
            for rx in src.rxList:
                v = np.random.rand(rx.nD)
                V += [v]
                self.D[src, rx] = v
                self.assertTrue(np.all(v == self.D[src, rx]))
        V = np.concatenate(V)
        self.assertTrue(np.all(V == Utils.mkvc(self.D)))

        D2 = Survey.Data(self.D.survey, V)
        self.assertTrue(np.all(Utils.mkvc(D2) == Utils.mkvc(self.D)))

    def test_uniqueSrcs(self):
        srcs = self.D.survey.srcList
        srcs += [srcs[0]]
        self.assertRaises(AssertionError, Survey.BaseSurvey, srcList=srcs)


if __name__ == '__main__':
    unittest.main()