mirror of
https://github.com/wassname/simpeg.git
synced 2026-06-30 21:29:40 +08:00
rowanc1
349 lines
11 KiB
Python
349 lines
11 KiB
Python
from SimPEG import Survey, Utils, np, sp
|
|
|
|
class RxFDEM(Survey.BaseRx):
|
|
|
|
knownRxTypes = {
|
|
'exr':['e', 'Ex', 'real'],
|
|
'eyr':['e', 'Ey', 'real'],
|
|
'ezr':['e', 'Ez', 'real'],
|
|
'exi':['e', 'Ex', 'imag'],
|
|
'eyi':['e', 'Ey', 'imag'],
|
|
'ezi':['e', 'Ez', 'imag'],
|
|
|
|
'bxr':['b', 'Fx', 'real'],
|
|
'byr':['b', 'Fy', 'real'],
|
|
'bzr':['b', 'Fz', 'real'],
|
|
'bxi':['b', 'Fx', 'imag'],
|
|
'byi':['b', 'Fy', 'imag'],
|
|
'bzi':['b', 'Fz', 'imag'],
|
|
}
|
|
|
|
def __init__(self, locs, rxType):
|
|
Survey.BaseRx.__init__(self, locs, rxType)
|
|
|
|
self._Ps = {}
|
|
self._Ps[self.projGLoc] = {}
|
|
|
|
@property
|
|
def projField(self):
|
|
"""Field Type projection (e.g. e b ...)"""
|
|
return self.knownRxTypes[self.rxType][0]
|
|
|
|
@property
|
|
def projGLoc(self):
|
|
"""Grid Location projection (e.g. Ex Fy ...)"""
|
|
return self.knownRxTypes[self.rxType][1]
|
|
|
|
@property
|
|
def projComp(self):
|
|
"""Component projection (real/imag)"""
|
|
return self.knownRxTypes[self.rxType][2]
|
|
|
|
@property
|
|
def nD(self):
|
|
"""Number of data in the receiver."""
|
|
return self.locs.shape[0]
|
|
|
|
def getP(self, mesh):
|
|
"""
|
|
Returns the projection matrices as a
|
|
list for all components collected by
|
|
the receivers.
|
|
|
|
.. note::
|
|
|
|
Projection matrices are stored as a nested dict,
|
|
First gridLocation, then mesh.
|
|
"""
|
|
gloc = self.projGLoc
|
|
if mesh not in self._Ps[gloc]:
|
|
self._Ps[gloc][mesh] = mesh.getInterpolationMat(self.locs, gloc)
|
|
P = self._Ps[gloc][mesh]
|
|
return P
|
|
|
|
def projectFields(self, tx, mesh, u):
|
|
P = self.getP(mesh)
|
|
u_part_complex = u[tx, self.projField]
|
|
# get the real or imag component
|
|
real_or_imag = self.projComp
|
|
u_part = getattr(u_part_complex, real_or_imag)
|
|
return P*u_part
|
|
|
|
def projectFieldsDeriv(self, tx, mesh, u, v, adjoint=False):
|
|
P = self.getP(mesh)
|
|
|
|
if not adjoint:
|
|
Pv_complex = P * v
|
|
real_or_imag = self.projComp
|
|
Pv = getattr(Pv_complex, real_or_imag)
|
|
elif adjoint:
|
|
Pv_real = P.T * v
|
|
|
|
real_or_imag = self.projComp
|
|
if real_or_imag == 'imag':
|
|
Pv = 1j*Pv_real
|
|
elif real_or_imag == 'real':
|
|
Pv = Pv_real.astype(complex)
|
|
else:
|
|
raise NotImplementedError('must be real or imag')
|
|
|
|
return Pv
|
|
|
|
|
|
class TxFDEM(Survey.BaseTx):
|
|
|
|
freq = None #: Frequency (float)
|
|
|
|
rxListPair = RxFDEM
|
|
|
|
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.vnD.sum()
|
|
|
|
@property
|
|
def vnD(self):
|
|
"""Vector number of data"""
|
|
return np.array([rx.nD for rx in self.rxList])
|
|
|
|
|
|
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 name in newFields:
|
|
field = self._initStore(name)
|
|
if field[freq].shape[1] == 1:
|
|
newFields[name] = Utils.mkvc(newFields[name],2)
|
|
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, v=None):
|
|
self.survey = survey
|
|
self._dataDict = {}
|
|
for tx in self.survey.txList:
|
|
self._dataDict[tx] = {}
|
|
if v is not None:
|
|
self.fromvec(v)
|
|
|
|
def _ensureCorrectKey(self, key):
|
|
if type(key) is tuple:
|
|
if len(key) is not 2:
|
|
raise KeyError('Key must be [Tx, Rx]')
|
|
if key[0] not in self.survey.txList:
|
|
raise KeyError('Tx Key must be a transmitter in the survey.')
|
|
if key[1] not in key[0].rxList:
|
|
raise KeyError('Rx Key must be a receiver for the transmitter.')
|
|
return key
|
|
elif isinstance(key, self.survey.txPair):
|
|
if key not in self.survey.txList:
|
|
raise KeyError('Key must be a transmitter in the survey.')
|
|
return key, None
|
|
else:
|
|
raise KeyError('Key must be [Tx] or [Tx,Rx]')
|
|
|
|
def __setitem__(self, key, value):
|
|
tx, rx = self._ensureCorrectKey(key)
|
|
assert rx is not None, 'set data using [Tx, Rx]'
|
|
assert type(value) == np.ndarray, 'value must by ndarray'
|
|
assert value.size == rx.nD, "value must have the same number of data as the transmitter."
|
|
self._dataDict[tx][rx] = Utils.mkvc(value)
|
|
|
|
def __getitem__(self, key):
|
|
tx, rx = self._ensureCorrectKey(key)
|
|
if rx is not None:
|
|
if rx not in self._dataDict[tx]:
|
|
raise Exception('Data for receiver has not yet been set.')
|
|
return self._dataDict[tx][rx]
|
|
|
|
return np.concatenate([self[tx,rx] for rx in tx.rxList])
|
|
|
|
def tovec(self):
|
|
return np.concatenate([self[tx] for tx in self.survey.txList])
|
|
|
|
def fromvec(self, v):
|
|
v = Utils.mkvc(v)
|
|
assert v.size == self.survey.nD, 'v must have the correct number of data.'
|
|
indBot, indTop = 0, 0
|
|
for tx in self.survey.txList:
|
|
for rx in tx.rxList:
|
|
indTop += rx.nD
|
|
self[tx, rx] = v[indBot:indTop]
|
|
indBot += rx.nD
|
|
|
|
|
|
class SurveyFDEM(Survey.BaseSurvey):
|
|
"""
|
|
docstring for SurveyFDEM
|
|
"""
|
|
|
|
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 nD(self):
|
|
"""Number of data"""
|
|
return np.array([tx.nD for tx in self.txList]).sum()
|
|
|
|
@property
|
|
def freqs(self):
|
|
"""Frequencies"""
|
|
return self._freqs
|
|
|
|
@property
|
|
def nFreq(self):
|
|
"""Number of frequencies"""
|
|
return len(self._freqDict)
|
|
|
|
@property
|
|
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):
|
|
data = DataFDEM(self)
|
|
for tx in self.txList:
|
|
for rx in tx.rxList:
|
|
data[tx, rx] = rx.projectFields(tx, self.mesh, u)
|
|
return data
|
|
|
|
def projectFieldsDeriv(self, u):
|
|
raise Exception('Use Transmitters to project fields deriv.')
|