diff --git a/SimPEG/Survey.py b/SimPEG/Survey.py index 8abb547a..b6a17f28 100644 --- a/SimPEG/Survey.py +++ b/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): diff --git a/SimPEG/Tests/test_SurveyAndData.py b/SimPEG/Tests/test_SurveyAndData.py index bb6b4645..6feccc71 100644 --- a/SimPEG/Tests/test_SurveyAndData.py +++ b/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()