diff --git a/SimPEG/Problem.py b/SimPEG/Problem.py index 63fc9c1c..094616fe 100644 --- a/SimPEG/Problem.py +++ b/SimPEG/Problem.py @@ -1,6 +1,279 @@ import Utils, Survey, Models, numpy as np, scipy.sparse as sp import Maps, Mesh + +class Fields(object): + """Fancy Field Storage + + u[:,'phi'] = phi + print u[tx0,'phi'] + + """ + + knownFields = None #: Known fields, a dict with locations, e.g. {"e": "E", "phi": "CC"} + aliasFields = None #: Aliased fields, a dict with [alias, location, function], e.g. {"b":["e","F",lambda(F,e,ind)]} + dtype = float #: dtype is the type of the storage matrix. This can be a dictionary. + + def __init__(self, mesh, survey, **kwargs): + self.survey = survey + self.mesh = mesh + Utils.setKwargs(self, **kwargs) + self._fields = {} + + if self.knownFields is None: + raise Exception('knownFields cannot be set to None') + if self.aliasFields is None: + self.aliasFields = {} + + allFields = [k for k in self.knownFields] + [a for a in self.aliasFields] + assert len(allFields) == len(set(allFields)), 'Aliased fields and Known Fields have overlapping definitions.' + self.startup() + + def startup(self): + pass + + @property + def approxSize(self): + """The approximate cost to storing all of the known fields.""" + sz = 0.0 + for f in self.knownFields: + loc =self.knownFields[f] + sz += np.array(self._storageShape(loc)).prod()*8.0/(1024**2) + return "%e MB"%sz + + def _storageShape(self, loc): + nTx = self.survey.nTx + + nP = {'CC': self.mesh.nC, + 'N': self.mesh.nN, + 'F': self.mesh.nF, + 'E': self.mesh.nE}[loc] + + return (nP, nTx) + + def _initStore(self, name): + if name in self._fields: + return self._fields[name] + + assert name in self.knownFields, 'field name is not known.' + + loc = self.knownFields[name] + + if type(self.dtype) is dict: + dtype = self.dtype[name] + else: + dtype = self.dtype + field = np.zeros(self._storageShape(loc), dtype=dtype) + + self._fields[name] = field + + return field + + def _txIndex(self, txTestList): + if type(txTestList) is slice: + ind = txTestList + else: + if type(txTestList) is not list: + txTestList = [txTestList] + for txTest in txTestList: + if txTest not in self.survey.txList: + raise KeyError('Invalid Transmitter, not in survey list.') + + ind = np.in1d(self.survey.txList, txTestList) + return ind + + def _nameIndex(self, name, accessType): + + if type(name) is slice: + assert name == slice(None,None,None), 'Fancy field name slicing is not supported... yet.' + name = None + + if name is None: + return + if accessType=='set' and name not in self.knownFields: + if name in self.aliasFields: + raise KeyError("Invalid field name (%s) for setter, you can't set an aliased property"%name) + else: + raise KeyError('Invalid field name (%s) for setter'%name) + + elif accessType=='get' and (name not in self.knownFields and name not in self.aliasFields): + raise KeyError('Invalid field name (%s) for getter'%name) + return name + + def _indexAndNameFromKey(self, key, accessType): + if type(key) is not tuple: + key = (key,) + if len(key) == 1: + key += (None,) + + assert len(key) == 2, 'must be [Tx, fieldName]' + + txTestList, name = key + name = self._nameIndex(name, accessType) + ind = self._txIndex(txTestList) + return ind, name + + def __setitem__(self, key, value): + ind, name = self._indexAndNameFromKey(key, 'set') + if name is None: + freq = key + assert type(value) is dict, 'New fields must be a dictionary, if field is not specified.' + newFields = value + elif name in self.knownFields: + newFields = {name: value} + else: + raise Exception('Unknown setter') + + for name in newFields: + field = self._initStore(name) + self._setField(field, newFields[name], name, ind) + + def __getitem__(self, key): + ind, name = self._indexAndNameFromKey(key, 'get') + if name is None: + out = {} + for name in self._fields: + out[name] = self._getField(name, ind) + return out + return self._getField(name, ind) + + def _setField(self, field, val, name, ind): + if isinstance(val, np.ndarray) and (field.shape[1] == 1 or val.ndim == 1): + val = Utils.mkvc(val,2) + field[:,ind] = val + + def _getField(self, name, ind): + if name in self._fields: + out = self._fields[name][:,ind] + else: + # Aliased fields + alias, loc, func = self.aliasFields[name] + + txII = np.array(self.survey.txList)[ind] + if isinstance(txII, np.ndarray): + txII = txII.tolist() + if len(txII) == 1: + txII = txII[0] + + if type(func) is str: + assert hasattr(self, func), 'The alias field function is a string, but it does not exist in the Fields class.' + func = getattr(self, func) + out = func(self._fields[alias][:,ind], txII) + + if out.shape[1] == 1: + out = Utils.mkvc(out) + return out + + def __contains__(self, other): + if other in self.aliasFields: + other = self.aliasFields[other][0] + return self._fields.__contains__(other) + + +class TimeFields(Fields): + """Fancy Field Storage for time domain problems + + u[:,'phi', timeInd] = phi + print u[tx0,'phi'] + + """ + + def _storageShape(self, loc): + nP = {'CC': self.mesh.nC, + 'N': self.mesh.nN, + 'F': self.mesh.nF, + 'E': self.mesh.nE}[loc] + nTx = self.survey.nTx + nT = self.survey.prob.nT + 1 + return (nP, nTx, nT) + + def _indexAndNameFromKey(self, key, accessType): + if type(key) is not tuple: + key = (key,) + if len(key) == 1: + key += (None,) + if len(key) == 2: + key += (slice(None,None,None),) + + assert len(key) == 3, 'must be [Tx, fieldName, times]' + + txTestList, name, timeInd = key + + name = self._nameIndex(name, accessType) + txInd = self._txIndex(txTestList) + + return (txInd, timeInd), name + + def _correctShape(self, name, ind, deflate=False): + txInd, timeInd = ind + if name in self.knownFields: + loc = self.knownFields[name] + else: + loc = self.aliasFields[name][1] + nP, total_nTx, total_nT = self._storageShape(loc) + nTx = np.ones(total_nTx, dtype=bool)[txInd].sum() + nT = np.ones(total_nT, dtype=bool)[timeInd].sum() + shape = nP, nTx, nT + if deflate: + shape = tuple([s for s in shape if s > 1]) + return shape + + def _setField(self, field, val, name, ind): + txInd, timeInd = ind + shape = self._correctShape(name, ind) + if Utils.isScalar(val): + field[:,txInd,timeInd] = val + return + if val.size != np.array(shape).prod(): + raise ValueError('Incorrect size for data.') + correctShape = field[:,txInd,timeInd].shape + field[:,txInd,timeInd] = val.reshape(correctShape, order='F') + + def _getField(self, name, ind): + txInd, timeInd = ind + + if name in self._fields: + out = self._fields[name][:,txInd,timeInd] + else: + # Aliased fields + alias, loc, func = self.aliasFields[name] + if type(func) is str: + assert hasattr(self, func), 'The alias field function is a string, but it does not exist in the Fields class.' + func = getattr(self, func) + pointerFields = self._fields[alias][:,txInd,timeInd] + pointerShape = self._correctShape(alias, ind) + pointerFields = pointerFields.reshape(pointerShape, order='F') + + timeII = np.arange(self.survey.prob.nT + 1)[timeInd] + txII = np.array(self.survey.txList)[txInd] + if isinstance(txII, np.ndarray): + txII = txII.tolist() + if len(txII) == 1: + txII = txII[0] + + if timeII.size == 1: + pointerShapeDeflated = self._correctShape(alias, ind, deflate=True) + pointerFields = pointerFields.reshape(pointerShapeDeflated, order='F') + out = func(pointerFields, txII, timeII) + else: #loop over the time steps + nT = pointerShape[2] + out = range(nT) + for i, TIND_i in enumerate(timeII): + fieldI = pointerFields[:,:,i] + if fieldI.ndim == 2 and fieldI.shape[1] == 1: + fieldI = Utils.mkvc(fieldI) + out[i] = func(fieldI, txII, TIND_i) + if out[i].ndim == 1: + out[i] = out[i][:,np.newaxis,np.newaxis] + elif out[i].ndim == 2: + out[i] = out[i][:,:,np.newaxis] + out = np.concatenate(out, axis=2) + + shape = self._correctShape(name, ind, deflate=True) + return out.reshape(shape, order='F') + + + class BaseProblem(object): """ Problem is the base class for all geophysical forward problems in SimPEG. diff --git a/SimPEG/Survey.py b/SimPEG/Survey.py index 29fbafca..fdf8ed34 100644 --- a/SimPEG/Survey.py +++ b/SimPEG/Survey.py @@ -211,276 +211,6 @@ class Data(object): indBot += rx.nD -class Fields(object): - """Fancy Field Storage - - u[:,'phi'] = phi - print u[tx0,'phi'] - - """ - - knownFields = None #: Known fields, a dict with locations, e.g. {"e": "E", "phi": "CC"} - aliasFields = None #: Aliased fields, a dict with [alias, location, function], e.g. {"b":["e","F",lambda(F,e,ind)]} - dtype = float #: dtype is the type of the storage matrix. This can be a dictionary. - - def __init__(self, mesh, survey, **kwargs): - self.survey = survey - self.mesh = mesh - Utils.setKwargs(self, **kwargs) - self._fields = {} - - if self.knownFields is None: - raise Exception('knownFields cannot be set to None') - if self.aliasFields is None: - self.aliasFields = {} - - allFields = [k for k in self.knownFields] + [a for a in self.aliasFields] - assert len(allFields) == len(set(allFields)), 'Aliased fields and Known Fields have overlapping definitions.' - self.startup() - - def startup(self): - pass - - @property - def approxSize(self): - """The approximate cost to storing all of the known fields.""" - sz = 0.0 - for f in self.knownFields: - loc =self.knownFields[f] - sz += np.array(self._storageShape(loc)).prod()*8.0/(1024**2) - return "%e MB"%sz - - def _storageShape(self, loc): - nTx = self.survey.nTx - - nP = {'CC': self.mesh.nC, - 'N': self.mesh.nN, - 'F': self.mesh.nF, - 'E': self.mesh.nE}[loc] - - return (nP, nTx) - - def _initStore(self, name): - if name in self._fields: - return self._fields[name] - - assert name in self.knownFields, 'field name is not known.' - - loc = self.knownFields[name] - - if type(self.dtype) is dict: - dtype = self.dtype[name] - else: - dtype = self.dtype - field = np.zeros(self._storageShape(loc), dtype=dtype) - - self._fields[name] = field - - return field - - def _txIndex(self, txTestList): - if type(txTestList) is slice: - ind = txTestList - else: - if type(txTestList) is not list: - txTestList = [txTestList] - for txTest in txTestList: - if txTest not in self.survey.txList: - raise KeyError('Invalid Transmitter, not in survey list.') - - ind = np.in1d(self.survey.txList, txTestList) - return ind - - def _nameIndex(self, name, accessType): - - if type(name) is slice: - assert name == slice(None,None,None), 'Fancy field name slicing is not supported... yet.' - name = None - - if name is None: - return - if accessType=='set' and name not in self.knownFields: - if name in self.aliasFields: - raise KeyError("Invalid field name (%s) for setter, you can't set an aliased property"%name) - else: - raise KeyError('Invalid field name (%s) for setter'%name) - - elif accessType=='get' and (name not in self.knownFields and name not in self.aliasFields): - raise KeyError('Invalid field name (%s) for getter'%name) - return name - - def _indexAndNameFromKey(self, key, accessType): - if type(key) is not tuple: - key = (key,) - if len(key) == 1: - key += (None,) - - assert len(key) == 2, 'must be [Tx, fieldName]' - - txTestList, name = key - name = self._nameIndex(name, accessType) - ind = self._txIndex(txTestList) - return ind, name - - def __setitem__(self, key, value): - ind, name = self._indexAndNameFromKey(key, 'set') - if name is None: - freq = key - assert type(value) is dict, 'New fields must be a dictionary, if field is not specified.' - newFields = value - elif name in self.knownFields: - newFields = {name: value} - else: - raise Exception('Unknown setter') - - for name in newFields: - field = self._initStore(name) - self._setField(field, newFields[name], name, ind) - - def __getitem__(self, key): - ind, name = self._indexAndNameFromKey(key, 'get') - if name is None: - out = {} - for name in self._fields: - out[name] = self._getField(name, ind) - return out - return self._getField(name, ind) - - def _setField(self, field, val, name, ind): - if isinstance(val, np.ndarray) and (field.shape[1] == 1 or val.ndim == 1): - val = Utils.mkvc(val,2) - field[:,ind] = val - - def _getField(self, name, ind): - if name in self._fields: - out = self._fields[name][:,ind] - else: - # Aliased fields - alias, loc, func = self.aliasFields[name] - - txII = np.array(self.survey.txList)[ind] - if isinstance(txII, np.ndarray): - txII = txII.tolist() - if len(txII) == 1: - txII = txII[0] - - if type(func) is str: - assert hasattr(self, func), 'The alias field function is a string, but it does not exist in the Fields class.' - func = getattr(self, func) - out = func(self._fields[alias][:,ind], txII) - - if out.shape[1] == 1: - out = Utils.mkvc(out) - return out - - def __contains__(self, other): - if other in self.aliasFields: - other = self.aliasFields[other][0] - return self._fields.__contains__(other) - -class TimeFields(Fields): - """Fancy Field Storage for time domain problems - - u[:,'phi', timeInd] = phi - print u[tx0,'phi'] - - """ - - def _storageShape(self, loc): - nP = {'CC': self.mesh.nC, - 'N': self.mesh.nN, - 'F': self.mesh.nF, - 'E': self.mesh.nE}[loc] - nTx = self.survey.nTx - nT = self.survey.prob.nT + 1 - return (nP, nTx, nT) - - def _indexAndNameFromKey(self, key, accessType): - if type(key) is not tuple: - key = (key,) - if len(key) == 1: - key += (None,) - if len(key) == 2: - key += (slice(None,None,None),) - - assert len(key) == 3, 'must be [Tx, fieldName, times]' - - txTestList, name, timeInd = key - - name = self._nameIndex(name, accessType) - txInd = self._txIndex(txTestList) - - return (txInd, timeInd), name - - def _correctShape(self, name, ind, deflate=False): - txInd, timeInd = ind - if name in self.knownFields: - loc = self.knownFields[name] - else: - loc = self.aliasFields[name][1] - nP, total_nTx, total_nT = self._storageShape(loc) - nTx = np.ones(total_nTx, dtype=bool)[txInd].sum() - nT = np.ones(total_nT, dtype=bool)[timeInd].sum() - shape = nP, nTx, nT - if deflate: - shape = tuple([s for s in shape if s > 1]) - return shape - - def _setField(self, field, val, name, ind): - txInd, timeInd = ind - shape = self._correctShape(name, ind) - if Utils.isScalar(val): - field[:,txInd,timeInd] = val - return - if val.size != np.array(shape).prod(): - raise ValueError('Incorrect size for data.') - correctShape = field[:,txInd,timeInd].shape - field[:,txInd,timeInd] = val.reshape(correctShape, order='F') - - def _getField(self, name, ind): - txInd, timeInd = ind - - if name in self._fields: - out = self._fields[name][:,txInd,timeInd] - else: - # Aliased fields - alias, loc, func = self.aliasFields[name] - if type(func) is str: - assert hasattr(self, func), 'The alias field function is a string, but it does not exist in the Fields class.' - func = getattr(self, func) - pointerFields = self._fields[alias][:,txInd,timeInd] - pointerShape = self._correctShape(alias, ind) - pointerFields = pointerFields.reshape(pointerShape, order='F') - - timeII = np.arange(self.survey.prob.nT + 1)[timeInd] - txII = np.array(self.survey.txList)[txInd] - if isinstance(txII, np.ndarray): - txII = txII.tolist() - if len(txII) == 1: - txII = txII[0] - - if timeII.size == 1: - pointerShapeDeflated = self._correctShape(alias, ind, deflate=True) - pointerFields = pointerFields.reshape(pointerShapeDeflated, order='F') - out = func(pointerFields, txII, timeII) - else: #loop over the time steps - nT = pointerShape[2] - out = range(nT) - for i, TIND_i in enumerate(timeII): - fieldI = pointerFields[:,:,i] - if fieldI.ndim == 2 and fieldI.shape[1] == 1: - fieldI = Utils.mkvc(fieldI) - out[i] = func(fieldI, txII, TIND_i) - if out[i].ndim == 1: - out[i] = out[i][:,np.newaxis,np.newaxis] - elif out[i].ndim == 2: - out[i] = out[i][:,:,np.newaxis] - out = np.concatenate(out, axis=2) - - shape = self._correctShape(name, ind, deflate=True) - return out.reshape(shape, order='F') - - class BaseSurvey(object): """Survey holds the observed data, and the standard deviations.""" diff --git a/SimPEG/Tests/test_Survey.py b/SimPEG/Tests/test_Survey.py index 6c702252..5a5b8fc2 100644 --- a/SimPEG/Tests/test_Survey.py +++ b/SimPEG/Tests/test_Survey.py @@ -20,14 +20,14 @@ class DataAndFieldsTest(unittest.TestCase): txList = [Tx0,Tx1,Tx2,Tx3,Tx4] survey = Survey.BaseSurvey(txList=txList) self.D = Survey.Data(survey) - self.F = Survey.Fields(mesh, survey, knownFields={'phi':'CC','e':'E','b':'F'}, dtype={"phi":float,"e":complex,"b":complex}) + self.F = Problem.Fields(mesh, survey, knownFields={'phi':'CC','e':'E','b':'F'}, dtype={"phi":float,"e":complex,"b":complex}) self.Tx0 = Tx0 self.Tx1 = Tx1 self.mesh = mesh self.XYZ = XYZ def test_overlappingFields(self): - self.assertRaises(AssertionError, Survey.Fields, self.F.mesh, self.F.survey, + self.assertRaises(AssertionError, Problem.Fields, self.F.mesh, self.F.survey, knownFields={'b':'F'}, aliasFields={'b':['b',(lambda F, b, ind: b)]}) @@ -133,7 +133,7 @@ class FieldsTest_Alias(unittest.TestCase): txList = [Tx0,Tx1,Tx2,Tx3,Tx4] survey = Survey.BaseSurvey(txList=txList) self.D = Survey.Data(survey) - self.F = Survey.Fields(mesh, survey, knownFields={'e':'E'}, aliasFields={'b':['e','F',(lambda e, ind: self.F.mesh.edgeCurl * e)]}) + self.F = Problem.Fields(mesh, survey, knownFields={'e':'E'}, aliasFields={'b':['e','F',(lambda e, ind: self.F.mesh.edgeCurl * e)]}) self.Tx0 = Tx0 self.Tx1 = Tx1 self.mesh = mesh @@ -168,7 +168,7 @@ class FieldsTest_Alias(unittest.TestCase): def alias(e, ind): self.assertTrue(ind is self.Tx0) return self.F.mesh.edgeCurl * e - F = Survey.Fields(self.F.mesh, self.F.survey, knownFields={'e':'E'}, aliasFields={'b':['e','F',alias]}) + F = Problem.Fields(self.F.mesh, self.F.survey, knownFields={'e':'E'}, aliasFields={'b':['e','F',alias]}) e = np.random.rand(F.mesh.nE,1) F[self.Tx0, 'e'] = e F[self.Tx0, 'b'] @@ -179,7 +179,7 @@ class FieldsTest_Alias(unittest.TestCase): self.assertTrue(ind[0] is self.Tx0) self.assertTrue(ind[1] is self.Tx1) return self.F.mesh.edgeCurl * e - F = Survey.Fields(self.F.mesh, self.F.survey, knownFields={'e':'E'}, aliasFields={'b':['e','F',alias]}) + F = Problem.Fields(self.F.mesh, self.F.survey, knownFields={'e':'E'}, aliasFields={'b':['e','F',alias]}) e = np.random.rand(F.mesh.nE,2) F[[self.Tx0, self.Tx1], 'e'] = e F[[self.Tx0, self.Tx1], 'b'] @@ -205,7 +205,7 @@ class FieldsTest_Time(unittest.TestCase): survey = Survey.BaseSurvey(txList=txList) prob = Problem.BaseTimeProblem(mesh, timeSteps=[(10.,3), (20.,2)]) survey.pair(prob) - self.F = Survey.TimeFields(mesh, survey, knownFields={'phi':'CC','e':'E','b':'F'}) + self.F = Problem.TimeFields(mesh, survey, knownFields={'phi':'CC','e':'E','b':'F'}) self.Tx0 = Tx0 self.Tx1 = Tx1 self.mesh = mesh @@ -309,7 +309,7 @@ class FieldsTest_Time_Aliased(unittest.TestCase): survey.pair(prob) def alias(b, txInd, timeInd): return self.F.mesh.edgeCurl.T * b + timeInd - self.F = Survey.TimeFields(mesh, survey, knownFields={'b':'F'}, aliasFields={'e':['b','E',alias]}) + self.F = Problem.TimeFields(mesh, survey, knownFields={'b':'F'}, aliasFields={'e':['b','E',alias]}) self.Tx0 = Tx0 self.Tx1 = Tx1 self.mesh = mesh @@ -364,7 +364,7 @@ class FieldsTest_Time_Aliased(unittest.TestCase): count[0] += 1 self.assertTrue(txInd is self.Tx0) return self.F.mesh.edgeCurl * e - F = Survey.TimeFields(self.F.mesh, self.F.survey, knownFields={'e':'E'}, aliasFields={'b':['e','F',alias]}) + F = Problem.TimeFields(self.F.mesh, self.F.survey, knownFields={'e':'E'}, aliasFields={'b':['e','F',alias]}) e = np.random.rand(F.mesh.nE,1,nT) F[self.Tx0, 'e', :] = e F[self.Tx0, 'b', :] @@ -382,7 +382,7 @@ class FieldsTest_Time_Aliased(unittest.TestCase): self.assertTrue(txInd[0] is self.Tx0) self.assertTrue(txInd[1] is self.Tx1) return self.F.mesh.edgeCurl * e - F = Survey.TimeFields(self.F.mesh, self.F.survey, knownFields={'e':'E'}, aliasFields={'b':['e','F',alias]}) + F = Problem.TimeFields(self.F.mesh, self.F.survey, knownFields={'e':'E'}, aliasFields={'b':['e','F',alias]}) e = np.random.rand(F.mesh.nE,2, nT) F[[self.Tx0, self.Tx1], 'e', :] = e count[0] = 0