diff --git a/simpegEM/FDEM/FDEM.py b/simpegEM/FDEM/FDEM.py index d64c150e..260d0e06 100644 --- a/simpegEM/FDEM/FDEM.py +++ b/simpegEM/FDEM/FDEM.py @@ -1,8 +1,10 @@ from SimPEG import Problem, Solver, Utils, np, sp from scipy.constants import mu_0 -from FieldsFDEM import FieldsFDEM -from SurveyFDEM import SurveyFDEM +from SurveyFDEM import SurveyFDEM, DataFDEM, FieldsFDEM +def omega(freq): + """Change frequency to angular frequency, omega""" + return 2.*np.pi*freq class ProblemFDEM_e(Problem.BaseProblem): """ @@ -21,15 +23,10 @@ class ProblemFDEM_e(Problem.BaseProblem): storeTheseFields = 'e' surveyPair = SurveyFDEM + dataPair = DataFDEM solveOpts = {'factorize':False, 'backend':'scipy'} - j_s = None - - def getFieldsObject(self): - return FieldsFDEM(self.mesh, self.survey.nTx, - self.survey.nFreq, store=self.storeTheseFields) - #################################################### # Mass Matrices @@ -61,19 +58,18 @@ class ProblemFDEM_e(Problem.BaseProblem): # Internal Methods #################################################### - def getA(self, freqInd): + def getA(self, freq): """ :param int fInd: Frequency index :rtype: scipy.sparse.csr_matrix :return: A """ - omega = self.survey.omega[freqInd] - return self.mesh.edgeCurl.T*self.MfMui*self.mesh.edgeCurl + 1j*omega*self.MeSigma + return self.mesh.edgeCurl.T*self.MfMui*self.mesh.edgeCurl + 1j*omega(freq)*self.MeSigma - def getRHS(self, freqInd): - omega = self.survey.omega[freqInd] + def getRHS(self, freq): #TODO: this needs to also depend on your transmitter! - return -1j*omega*self.Me*self.j_s + + return -1j*omega(freq)*self.Me*self.j_s def fields(self, m, useThisRhs=None): @@ -81,18 +77,17 @@ class ProblemFDEM_e(Problem.BaseProblem): self.makeMassMatrices(m) - F = self.getFieldsObject() + F = FieldsFDEM(self.mesh, self.survey) - for freqInd in range(self.survey.nFreq): - A = self.getA(freqInd) - b = self.getRHS(freqInd) + for freq in self.survey.freqs: + A = self.getA(freq) + b = self.getRHS(freq) e = Solver(A, options=self.solveOpts).solve(b) - F.set_e(e, freqInd) - omega = self.survey.omega[freqInd] + F[freq, 'e'] = e #TODO: check if mass matrices needed: - b = -1./(1j*omega)*self.mesh.edgeCurl*e - F.set_b(b, freqInd) + b = -1./(1j*omega(freq))*self.mesh.edgeCurl*e + F[freq, 'b'] = b return F @@ -103,19 +98,23 @@ class ProblemFDEM_e(Problem.BaseProblem): if u is None: u = self.fields(m) - Jvs = range(self.survey.nFreq) - P = self.survey.projectFieldsDeriv(u) + Jv = self.dataPair(self.survey) - for i, freqInd in enumerate(range(self.survey.nFreq)): - e = u.get_e(freqInd) - omega = self.survey.omega[freqInd] - # for txInd in self.survey.nTx - dMe_dsig = self.mesh.getEdgeInnerProductDeriv(m, v=e) - dsig_dm = self.model.transformDeriv(m) - b = 1j*omega * ( dMe_dsig * ( dsig_dm * v ) ) - A = self.getA(freqInd) - Ab = Solver(A, options=self.solveOpts).solve(b) - Jvs[i] = -P*Ab + for i, freq in enumerate(self.survey.freqs): + e = u[freq, 'e'] + A = self.getA(freq) + solver = Solver(A, options=self.solveOpts) + + for tx in self.survey.getTransmitters(freq): + dMe_dsig = self.mesh.getEdgeInnerProductDeriv(m, v=e) + dsig_dm = self.model.transformDeriv(m) + b = 1j*omega(freq) * ( dMe_dsig * ( dsig_dm * v ) ) + Ab = solver.solve(b) + + #TODO: look at Rx for this... + P = self.survey.projectFieldsDeriv(u) + + Jv[tx] = -P*Ab Jv = np.concatenate(Jvs) @@ -156,10 +155,10 @@ if __name__ == '__main__': 'rxType':'bz', 'freq': np.logspace(0,3,4), } - dat = EM.FDEM.DataFDEM(**opts) + survey = EM.FDEM.SurveyFDEM(**opts) prb = EM.FDEM.ProblemFDEM_e(mesh, model) - prb.pair(dat) + prb.pair(survey) sigma = np.log(np.ones(mesh.nC)*1e-3) diff --git a/simpegEM/FDEM/FieldsFDEM.py b/simpegEM/FDEM/FieldsFDEM.py deleted file mode 100644 index 93add159..00000000 --- a/simpegEM/FDEM/FieldsFDEM.py +++ /dev/null @@ -1,52 +0,0 @@ -import numpy as np - - -class FieldsFDEM(object): - """docstring for FieldsFDEM""" - - phi = None #: Electric potential - A = None #: Magnetic vector potential - e = None #: Electric field - b = None #: Magnetic flux density - j = None #: Current density - h = None #: Magnetic field - - def __init__(self, mesh, nTx, nFreq, store='e'): - - self.nFreq = nFreq #: Number of times - self.nTx = nTx #: Number of transmitters - self.mesh = mesh - - def update(self, newFields, fInd): - self.set_b(newFields['b'], fInd) - self.set_e(newFields['e'], fInd) - - #################################################### - # Get Methods - #################################################### - - def get_b(self, ind): - return self.b[ind,:,:] - - def get_e(self, ind): - return self.e[ind,:,:] - - #################################################### - # Set Methods - #################################################### - - def set_b(self, b, ind): - if self.b is None: - self.b = np.zeros((self.nFreq, np.sum(self.mesh.nF), self.nTx), dtype=complex) - self.b[:] = np.nan - if len(b.shape) == 1: - b = b[:, np.newaxis] - self.b[ind,:,:] = b - - def set_e(self, e, ind): - if self.e is None: - self.e = np.zeros((self.nFreq, np.sum(self.mesh.nE), self.nTx), dtype=complex) - self.e[:] = np.nan - if len(e.shape) == 1: - e = e[:, np.newaxis] - self.e[ind,:,:] = e diff --git a/simpegEM/FDEM/SurveyFDEM.py b/simpegEM/FDEM/SurveyFDEM.py index 920ab00d..0fadb93f 100644 --- a/simpegEM/FDEM/SurveyFDEM.py +++ b/simpegEM/FDEM/SurveyFDEM.py @@ -1,35 +1,238 @@ from SimPEG import Survey, Utils, np, sp -from FieldsFDEM import FieldsFDEM -class SurveyFDEM(Survey.BaseSurvey): +class RxListFDEM(Survey.BaseRxList): + + knownRxTypes = ['Ex', 'Ey', 'Ez'] + + def __init__(self, locs, rxType): + Survey.BaseRxList.__init__(self, locs, rxType) + + self._Ps = {} + + def getP(self, mesh): + if mesh not in self._Ps: + self._Ps[mesh] = mesh.getInterpolationMat(self.locs, self.rxType) + return self._Ps[mesh] + + +class TxFDEM(Survey.BaseTx): + + freq = None #: Frequency (float) + + rxListPair = RxListFDEM + + knownTxTypes = ['VMD'] + + def __init__(self, loc, txType, freq, rxList): + self.freq = float(freq) + Survey.BaseTx.__init__(self, loc, txType, rxList) + + @property + def nD(self): + """Number of data""" + return self.rxList.locs.shape[0] + + + def projectFields(self, mesh, u): + + P = self.rxList.getP(mesh) + + u_part = u[self] + + Pu = P*u_part + return Pu + + def projectFieldsDeriv(self, mesh, u): + pass + + +class FieldsFDEM(object): + """Fancy Field Storage for a FDEM survey.""" + + knownFields = {'b': 'F', 'e': 'E'} + + def __init__(self, mesh, survey): + self.survey = survey + self.mesh = mesh + self._fields = {} + + 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] + + nP = {'CC': self.mesh.nC, + 'F': self.mesh.nF, + 'E': self.mesh.nE}[loc] + + field = {} + for freq in self.survey.freqs: + nTx_f = len(self.survey.getTransmitters(freq)) + field[freq] = np.empty((nP, nTx_f)) + + self._fields[name] = field + + return field + + def _ensureCorrectKey(self, key): + if type(key) is tuple: + assert len(key) == 2, 'must be [freq, fieldName]' + freqTest, name = key + + if name not in self.knownFields: + raise KeyError('Invalid field name') + + if type(freqTest) is float: + freq = freqTest + elif isinstance(freqTest, TxFDEM): + freq = freqTest.freq + if freqTest not in self.survey.txList: + raise KeyError('Invalid Transmitter') + else: + raise KeyError('Invalid Frequency Key') + + elif type(key) is float: + freq = key + elif isinstance(key, TxFDEM): + freq = key.freq + if key not in self.survey.txList: + raise KeyError('Invalid Transmitter') + else: + raise KeyError('Unexpected key use [freq, fieldName]') + if freq not in self.survey.freqs: + raise KeyError('Invalid frequency') + + def __setitem__(self, key, value): + self._ensureCorrectKey(key) + if type(key) is tuple: + freq, name = key + assert type(freq) is float, 'Frequency must be a float for setter.' + assert type(value) is np.ndarray, 'Must be set to a numpy array' + newFields = {name: value} + elif type(key) is float: + freq = key + assert type(value) is dict, 'New fields must be a dictionary' + newFields = value + elif isinstance(key, TxFDEM): + raise Exception('Cannot set one transmitter at a time.') + + for field in newFields: + field = self._initStore(name) + assert field[freq].shape == newFields[name].shape, 'Must be correct shape (n%s x nTx[freq])' % self.knownFields[name] + field[freq] = newFields[name] + + def __getitem__(self, key): + self._ensureCorrectKey(key) + if type(key) is tuple: + freqTest, name = key + if type(freqTest) is float: + return self._fields[name][freqTest] + elif isinstance(freqTest, TxFDEM): + key = freqTest + ind = np.array([tx is key for tx in self.survey.getTransmitters(key.freq)]) + return Utils.mkvc(self._fields[name][key.freq][:,ind]) + elif type(key) is float: + freq = key + out = {} + for name in self._fields: + out[name] = self._fields[name][freq] + return out + elif isinstance(key, TxFDEM): + freq = key.freq + ind = np.array([tx is key for tx in self.survey.getTransmitters(freq)]) + out = {} + for name in self._fields: + out[name] = Utils.mkvc(self._fields[name][freq][:,ind]) + return out + + + def __contains__(self, key): + return key in self.children + + +class DataFDEM(object): + """docstring for DataFDEM""" + def __init__(self, survey): + self.survey = survey + self._dataDict = {} + + def _ensureCorrectKey(self, key): + if key not in self.survey.txList: + raise KeyError('Key must be a transmitter in the survey.') + + def __setitem__(self, key, value): + self._ensureCorrectKey(key) + assert type(value) == np.ndarray, 'value must by ndarray' + assert value.size == key.nD, "value must have the same number of data as the transmitter." + self._dataDict[key] = Utils.mkvc(value) + + def __getitem__(self, key): + self._ensureCorrectKey(key) + return self._dataDict[key] + + def toarray(self): + D = self._dataDict + return np.concatenate([D[k] for k in D]) + + + +class SurveyFDEM(Survey.MixinFancyProjection, Survey.BaseSurvey): """ docstring for SurveyFDEM """ - txLoc = None #: txLoc - txType = None #: txType - nTx = 1 #: Number of transmitters - rxLoc = None #: rxLoc - rxType = None #: rxType - freq = None #: freq + txPair = TxFDEM + def __init__(self, txList, **kwargs): + assert type(txList) is list, 'txList must be a list' + for tx in txList: + assert isinstance(tx, self.txPair), 'txList must be a %s'%self.txPair.__name__ + + assert len(set(txList)) == len(txList), 'The txList must be unique' + # Sort these by frequency + _freqDict = {} + for tx in txList: + if tx.freq not in _freqDict: + _freqDict[tx.freq] = [] + _freqDict[tx.freq] += [tx] + + self._txList = txList + self._freqDict = _freqDict + self._freqs = sorted([f for f in self._freqDict]) + + Survey.BaseSurvey.__init__(self, **kwargs) @property - def omega(self): - return 2*np.pi*self.freq + def freqs(self): + """Frequencies""" + return self._freqs @property def nFreq(self): """Number of frequencies""" - return self.freq.size - - def __init__(self, **kwargs): - Survey.BaseSurvey.__init__(self, **kwargs) - Utils.setKwargs(self, **kwargs) + return len(self._freqDict) @property - def nRx(self): - return self.rxLoc.shape[0] + def txList(self): + """Transmitter List""" + return self._txList + + @property + def nTx(self): + if getattr(self, '_nTx', None) is None: + self._nTx = {} + for freq in self.freqs: + self._nTx[freq] = len(self.getTransmitters(freq)) + return self._nTx + + + def getTransmitters(self, freq): + """Returns the transmitters associated with a specific frequency.""" + assert freq in self._freqDict, "The requested frequency is not in this survey." + return self._freqDict[freq] def projectFields(self, u): P = sp.identity(self.prob.mesh.nE) @@ -41,12 +244,19 @@ class SurveyFDEM(Survey.BaseSurvey): Pe = np.concatenate(Pes) return Pe - def projectFieldsDeriv(self, u): - # TODO : more general - return sp.identity(self.prob.mesh.nE) + + def projectFieldsDerivVec(self, u, v=None): + raise NotImplemented('projectFieldsDerivVec is not yet implemented') + + def projectAdjointFieldsDerivVec(self, u, v=None): + raise NotImplemented('projectAdjointFieldsDerivVec is not yet implemented') + #################################################### # Interpolation Matrices #################################################### - + @Utils.requires('prob') + def getP(self, Tx): + # TODO: store these in a mesh lookup table by transmitter? + pass diff --git a/simpegEM/FDEM/__init__.py b/simpegEM/FDEM/__init__.py index 4d562918..f6583225 100644 --- a/simpegEM/FDEM/__init__.py +++ b/simpegEM/FDEM/__init__.py @@ -1,3 +1,2 @@ -from FieldsFDEM import FieldsFDEM from SurveyFDEM import * from FDEM import ProblemFDEM_e diff --git a/simpegEM/Tests/test_FDEM.py b/simpegEM/Tests/test_FDEM.py index 11aac751..f40e5dd6 100644 --- a/simpegEM/Tests/test_FDEM.py +++ b/simpegEM/Tests/test_FDEM.py @@ -2,7 +2,7 @@ import unittest from SimPEG import * import simpegEM as EM -class TDEM_bDerivTests(unittest.TestCase): +class FDEM_bDerivTests(unittest.TestCase): def setUp(self): diff --git a/simpegEM/Tests/test_FieldsObject.py b/simpegEM/Tests/test_FieldsObject.py new file mode 100644 index 00000000..302c4882 --- /dev/null +++ b/simpegEM/Tests/test_FieldsObject.py @@ -0,0 +1,69 @@ +import unittest +from SimPEG import * +import simpegEM as EM + +class FieldsTest(unittest.TestCase): + + def setUp(self): + x = np.linspace(5,10,3) + XYZ = Utils.ndgrid(x,x,np.r_[0]) + rxList = EM.FDEM.RxListFDEM(XYZ, 'Ex') + Tx0 = EM.FDEM.TxFDEM(None, 'VMD', 3, rxList) + Tx1 = EM.FDEM.TxFDEM(None, 'VMD', 3, rxList) + Tx2 = EM.FDEM.TxFDEM(None, 'VMD', 2, rxList) + Tx3 = EM.FDEM.TxFDEM(None, 'VMD', 1, rxList) + txList = [Tx0,Tx1,Tx2,Tx3] + mesh = Mesh.TensorMesh([np.ones(n)*5 for n in [10,11,12]],[0,0,-30]) + survey = EM.FDEM.SurveyFDEM(txList) + self.F = EM.FDEM.FieldsFDEM(mesh, survey) + self.Tx0 = Tx0 + self.Tx1 = Tx1 + + def test_SetGet(self): + F = self.F + for freq in F.survey.freqs: + e = np.random.rand(F.mesh.nE, F.survey.nTx[freq]) + F[freq, 'e'] = e + b = np.random.rand(F.mesh.nF, F.survey.nTx[freq]) + F[freq, 'b'] = b + self.assertTrue(np.all(F[freq, 'e'] == e)) + self.assertTrue(np.all(F[freq, 'b'] == b)) + + lastFreq = F[freq] + self.assertTrue(type(lastFreq) is dict) + self.assertTrue(sorted([k for k in lastFreq]) == ['b','e']) + self.assertTrue(np.all(lastFreq['b'] == b)) + self.assertTrue(np.all(lastFreq['e'] == e)) + + self.assertTrue(F[3.,'b'].shape == (F.mesh.nF, 2)) + + b = np.random.rand(F.mesh.nF, 2) + F[self.Tx0.freq,'b'] = b + self.assertTrue(F[self.Tx0]['b'].shape == (F.mesh.nF,)) + self.assertTrue(F[self.Tx0,'b'].shape == (F.mesh.nF,)) + self.assertTrue(np.all(F[self.Tx0,'b'] == b[:,0])) + self.assertTrue(np.all(F[self.Tx1,'b'] == b[:,1])) + + def test_assertions(self): + freq = self.F.survey.freqs[0] + bWrongSize = np.random.rand(self.F.mesh.nE, self.F.survey.nTx[freq]) + def fun(): self.F[freq, 'b'] = bWrongSize + self.assertRaises(AssertionError, fun) + def fun(): self.F[-999.] + self.assertRaises(KeyError, fun) + def fun(): self.F['notRight'] + self.assertRaises(KeyError, fun) + def fun(): self.F[freq,'notThere'] + self.assertRaises(KeyError, fun) + + def test_uniqueTxs(self): + txs = self.F.survey.txList + txs += [txs[0]] + self.assertRaises(AssertionError, EM.FDEM.SurveyFDEM, txs) + + + + + +if __name__ == '__main__': + unittest.main()