getSourceIndex

SimPEG/Survey.py

@@ -1,4 +1,4 @@
-import Utils, numpy as np, scipy.sparse as sp
+import Utils, numpy as np, scipy.sparse as sp, uuid
 
 
 class BaseRx(object):
@@ -13,6 +13,7 @@ class BaseRx(object):
     storeProjections = True #: Store calls to getP (organized by mesh)
 
     def __init__(self, locs, rxType, **kwargs):
+        self.uid = str(uuid.uuid4())
         self.locs = locs
         self.rxType = rxType
         self._Ps = {}
@@ -124,7 +125,7 @@ class BaseSrc(object):
         for rx in rxList:
             assert isinstance(rx, self.rxPair), 'rxList must be a %s'%self.rxPair.__name__
         assert len(set(rxList)) == len(rxList), 'The rxList must be unique'
-
+        self.uid = str(uuid.uuid4())
         self.rxList = rxList
         Utils.setKwargs(self, **kwargs)
 
@@ -144,6 +145,7 @@ class Data(object):
     """Fancy data storage by Src and Rx"""
 
     def __init__(self, survey, v=None):
+        self.uid = str(uuid.uuid4())
         self.survey = survey
         self._dataDict = {}
         for src in self.survey.srcList:
@@ -225,6 +227,14 @@ class BaseSurvey(object):
         assert np.all([isinstance(src, self.srcPair) for src in value]), 'All sources must be instances of %s' % self.srcPair.__name__
         assert len(set(value)) == len(value), 'The srcList must be unique'
         self._srcList = value
+        self._sourceOrder = dict()
+        [self._sourceOrder.setdefault(src.uid, ii) for ii, src in enumerate(self._srcList)]
+
+    def getSourceIndex(self, sources):
+        inds = map(lambda src: self._sourceOrder.get(src.uid, None), sources)
+        if None in inds:
+            raise KeyError('Some of the sources specified are not in this survey. %s'%str(inds))
+        return inds
 
     @property
     def prob(self):

SimPEG/Tests/test_SurveyAndData.py

@@ -40,6 +40,16 @@ class TestData(unittest.TestCase):
         srcs += [srcs[0]]
         self.assertRaises(AssertionError, Survey.BaseSurvey, srcList=srcs)
 
+    def test_sourceIndex(self):
+        survey = self.D.survey
+        srcs = survey.srcList
+        assert survey.getSourceIndex([srcs[1],srcs[0]]) == [1,0]
+        assert survey.getSourceIndex([srcs[1],srcs[2],srcs[2]]) == [1,2,2]
+        SrcNotThere = Survey.BaseSrc(srcs[0].rxList, loc=np.r_[0,0,0])
+        self.assertRaises(KeyError, survey.getSourceIndex, [SrcNotThere])
+        self.assertRaises(KeyError, survey.getSourceIndex, [srcs[1],srcs[2],SrcNotThere])
+
+
 
 if __name__ == '__main__':
     unittest.main()