def __init__(self, h_in, x0=None):
assert type(h_in) is list, 'h_in must be a list'
assert len(h_in) > 1, "len(h_in) must be greater than 1"
h = range(len(h_in))
for i, h_i in enumerate(h_in):
if type(h_i) in [int, long, float]:
# This gives you something over the unit cube.
h_i = np.ones(int(h_i))/int(h_i)
assert isinstance(h_i, np.ndarray), ("h[%i] is not a numpy array." % i)
assert len(h_i.shape) == 1, ("h[%i] must be a 1D numpy array." % i)
h[i] = h_i[:] # make a copy.
self.h = h
if x0 is None:
x0 = np.zeros(self.dim)
else:
assert type(x0) in [list, np.ndarray], 'x0 must be a numpy array or a list'
x0 = np.array(x0, dtype=float)
assert len(x0) == self.dim, 'x0 must have the same dimensions as the mesh'
# TODO: this has a lot of stuff which doesn't work for this style of mesh...
BaseMesh.__init__(self, np.array([x.size for x in h]), x0)
# set the sets for holding the cells, nodes, faces, and edges
self.cells = set()
self.nodes = set()
self.faces = set()
self.facesX = set()
self.facesY = set()
if self.dim == 3:
self.facesZ = set()
self.edges = set()
self.edgesX = set()
self.edgesY = set()
self.edgesZ = set()
self.children = np.empty([hi.size for hi in h],dtype=TreeCell)
if self.dim == 2:
for i in range(h[0].size):
for j in range(h[1].size):
fXm = None if i is 0 else self.children[i-1][j].fXp
fYm = None if j is 0 else self.children[i][j-1].fYp
x0i = (np.r_[x0[0], h[0][:i]]).sum()
x0j = (np.r_[x0[1], h[1][:j]]).sum()
self.children[i][j] = TreeCell(self, x0=[x0i, x0j], depth=0, sz=[h[0][i], h[1][j]], fXm=fXm, fYm=fYm)
elif self.dim == 3:
for i in range(h[0].size):
for j in range(h[1].size):
for k in range(h[2].size):
fXm = None if i is 0 else self.children[i-1][j][k].fXp
fYm = None if j is 0 else self.children[i][j-1][k].fYp
fZm = None if k is 0 else self.children[i][j][k-1].fZp
x0i = (np.r_[x0[0], h[0][:i]]).sum()
x0j = (np.r_[x0[1], h[1][:j]]).sum()
x0k = (np.r_[x0[2], h[2][:k]]).sum()
self.children[i][j][k] = TreeCell(self, x0=[x0i, x0j, x0k], depth=0, sz=[h[0][i], h[1][j], h[2][k]], fXm=fXm, fYm=fYm, fZm=fZm)
def toRecArray(self,returnType='RealImag'):
'''
Function that returns a numpy.recarray for a SimpegMT impedance data object.
:param str returnType: Switches between returning a rec array where the impedance is split to real and imaginary ('RealImag') or is a complex ('Complex')
'''
# Define the record fields
dtRI = [('freq',float),('x',float),('y',float),('z',float),('zxxr',float),('zxxi',float),('zxyr',float),('zxyi',float),
('zyxr',float),('zyxi',float),('zyyr',float),('zyyi',float),('tzxr',float),('tzxi',float),('tzyr',float),('tzyi',float)]
dtCP = [('freq',float),('x',float),('y',float),('z',float),('zxx',complex),('zxy',complex),('zyx',complex),('zyy',complex),('tzx',complex),('tzy',complex)]
impList = ['zxxr','zxxi','zxyr','zxyi','zyxr','zyxi','zyyr','zyyi']
for src in self.survey.srcList:
# Temp array for all the receivers of the source.
# Note: needs to be written more generally, using diffterent rxTypes and not all the data at the locaitons
# Assume the same locs for all RX
locs = src.rxList[0].locs
if locs.shape[1] == 1:
locs = np.hstack((np.array([[0.0,0.0]]),locs))
elif locs.shape[1] == 2:
locs = np.hstack((np.array([[0.0]]),locs))
tArrRec = np.concatenate((src.freq*np.ones((locs.shape[0],1)),locs,np.nan*np.ones((locs.shape[0],12))),axis=1).view(dtRI)
# np.array([(src.freq,rx.locs[0,0],rx.locs[0,1],rx.locs[0,2],np.nan ,np.nan ,np.nan ,np.nan ,np.nan ,np.nan ,np.nan ,np.nan ) for rx in src.rxList],dtype=dtRI)
# Get the type and the value for the DataMT object as a list
typeList = [[rx.rxType.replace('z1d','zyx'),self[src,rx]] for rx in src.rxList]
# Insert the values to the temp array
for nr,(key,val) in enumerate(typeList):
tArrRec[key] = mkvc(val,2)
# Masked array
mArrRec = np.ma.MaskedArray(rec2ndarr(tArrRec),mask=np.isnan(rec2ndarr(tArrRec))).view(dtype=tArrRec.dtype)
# Unique freq and loc of the masked array
uniFLmarr = np.unique(mArrRec[['freq','x','y','z']]).copy()
try:
outTemp = recFunc.stack_arrays((outTemp,mArrRec))
#outTemp = np.concatenate((outTemp,dataBlock),axis=0)
except NameError as e:
outTemp = mArrRec
if 'RealImag' in returnType:
outArr = outTemp
elif 'Complex' in returnType:
# Add the real and imaginary to a complex number
outArr = np.empty(outTemp.shape,dtype=dtCP)
for comp in ['freq','x','y','z']:
outArr[comp] = outTemp[comp].copy()
for comp in ['zxx','zxy','zyx','zyy','tzx','tzy']:
outArr[comp] = outTemp[comp+'r'].copy() + 1j*outTemp[comp+'i'].copy()
else:
raise NotImplementedError('{:s} is not implemented, as to be RealImag or Complex.')
# Return
return outArr
def _refine2D(self):
self.children = np.empty(2,dtype=TreeFace)
# Create refined x0's
x0r_0 = self.x0
x0r_1 = self.x0+0.5*self.tangent0*self.sz
self.children[0] = TreeFace(self.mesh, x0=x0r_0, faceType=self.faceType, sz=0.5*self.sz, depth=self.depth+1, node0=self.node0)
self.children[1] = TreeFace(self.mesh, x0=x0r_1, faceType=self.faceType, sz=0.5*self.sz, depth=self.depth+1, node0=self.children[0].node1, node1=self.node1)
self.mesh.faces.remove(self)
if self.faceType is 'x':
self.mesh.facesX.remove(self)
elif self.faceType is 'y':
self.mesh.facesY.remove(self)
def _refine3D(self):
#
# 2_______________3 _______________
# | e1--> | | | |
# ^ | | ^ | (0,1) | (1,1) |
# | | | | | | |
# | | x | | ---> |-------+-------|
# e2 | | e3 | | |
# | | | (0,0) | (1,0) |
# |_______________| |_______|_______|
# 0 e0--> 1
order = [{'c':[0,0],
'e0': ('p', 'e0', [0]), 'e1': 'new' ,
'e2': ('p', 'e2', [0]), 'e3': 'new' },
{'c':[1,0],
'e0': ('p', 'e0', [1]), 'e1': 'new' ,
'e2': ('c', 'e3', [0,0]), 'e3': ('p', 'e3', [0])},
{'c':[0,1],
'e0': ('c', 'e1', [0,0]), 'e1': ('p', 'e1', [0]),
'e2': ('p', 'e2', [1]), 'e3': 'new' },
{'c':[1,1],
'e0': ('c', 'e1', [1,0]), 'e1': ('p', 'e1', [1]),
'e2': ('c', 'e3', [0,1]), 'e3': ('p', 'e3', [1])}]
def getEdge(pointer):
if pointer is 'new': return
if pointer[0] == 'p':
return getattr(self, 'edg' + pointer[1]).children[pointer[2][0]]
if pointer[0] == 'c':
f = self.children[pointer[2][0],pointer[2][1]]
return getattr(f, 'edg' + pointer[1])
self.children = np.empty((2,2), dtype=TreeFace)
for edge in [self.edge0, self.edge1, self.edge2, self.edge3]:
edge.refine()
for O in order:
i, j = O['c']
x0r = self.x0 + 0.5*i*self.tangent0*self.sz[0] + 0.5*j*self.tangent1*self.sz[1]
e0, e1, e2, e3 = getEdge(O['e0']), getEdge(O['e1']), getEdge(O['e2']), getEdge(O['e3'])
self.children[i,j] = TreeFace(self.mesh, x0=x0r, faceType=self.faceType, depth=self.depth+1, sz=0.5*self.sz, edge0=e0, edge1=e1, edge2=e2, edge3=e3)
self.mesh.faces.remove(self)
if self.faceType is 'x':
self.mesh.facesX.remove(self)
elif self.faceType is 'y':
self.mesh.facesY.remove(self)
elif self.faceType is 'z':
self.mesh.facesZ.remove(self)
def _grid(self, key):
self.number()
sObjs = {'CC':self.sortedCells,
'N':self.sortedNodes,
'Fx': self.sortedFaceX,
'Fy': self.sortedFaceY,
'Fz': getattr(self,'sortedFaceZ', None),
'Ex': getattr(self,'sortedEdgeX', self.sortedFaceY),
'Ey': getattr(self,'sortedEdgeY', self.sortedFaceX),
'Ez': getattr(self,'sortedEdgeZ', None)}[key]
G = np.empty((len(sObjs),self.dim))
for ii, obj in enumerate(sObjs):
G[ii,:] = obj.center
return G
def refine(self):
if not self.isleaf: return
self.mesh.isNumbered = False
self.children = np.empty(2,dtype=TreeFace)
# Create refined x0's
x0r_0 = self.x0
x0r_1 = self.x0+0.5*self.tangent*self.sz
self.children[0] = TreeEdge(self.mesh, x0=x0r_0, edgeType=self.edgeType, sz=0.5*self.sz, depth=self.depth+1, node0=self.node0)
self.children[1] = TreeEdge(self.mesh, x0=x0r_1, edgeType=self.edgeType, sz=0.5*self.sz, depth=self.depth+1, node0=self.children[0].node1, node1=self.node1)
self.mesh.edges.remove(self)
if self.edgeType is 'x':
self.mesh.edgesX.remove(self)
elif self.edgeType is 'y':
self.mesh.edgesY.remove(self)
elif self.edgeType is 'z':
self.mesh.edgesZ.remove(self)
def _refine2D(self):
self.mesh.isNumbered = False
self.children = np.empty((2,2), dtype=TreeCell)
x0, sz = self.x0, self.sz
for face in self.faceList:
face.refine()
order = [{'c':[0,0],
'fXm': ('p', 'fXm', [0]), 'fXp': 'new' ,
'fYm': ('p', 'fYm', [0]), 'fYp': 'new' },
{'c':[1,0],
'fXm': ('c', 'fXp', [0,0]), 'fXp': ('p', 'fXp', [0]),
'fYm': ('p', 'fYm', [1]), 'fYp': 'new' },
{'c':[0,1],
'fXm': ('p', 'fXm', [1]), 'fXp': 'new' ,
'fYm': ('c', 'fYp', [0,0]), 'fYp': ('p', 'fYp', [0])},
{'c':[1,1],
'fXm': ('c', 'fXp', [0,1]), 'fXp': ('p', 'fXp', [1]),
'fYm': ('c', 'fYp', [1,0]), 'fYp': ('p', 'fYp', [1])}]
def getFace(pointer):
if pointer is 'new': return None
if pointer[0] == 'p':
return self.faceDict[pointer[1]].children[pointer[2][0],]
if pointer[0] == 'c':
return self.children[pointer[2][0],pointer[2][1]].faceDict[pointer[1]]
for O in order:
i, j = O['c']
x0r = np.r_[x0[0] + 0.5*i*sz[0], x0[1] + 0.5*j*sz[1]]
fXm, fXp, fYm, fYp = getFace(O['fXm']), getFace(O['fXp']), getFace(O['fYm']), getFace(O['fYp'])
self.children[i,j] = TreeCell(self.mesh, x0=x0r, depth=self.depth+1, sz=0.5*sz, fXm=fXm, fXp=fXp, fYm=fYm, fYp=fYp)
self.mesh.cells.remove(self)
def MagneticDipoleVectorPotential(srcLoc, obsLoc, component, moment=1.,
orientation=np.r_[0., 0., 1.],
mu=mu_0):
"""
Calculate the vector potential of a set of magnetic dipoles
at given locations 'ref. <http://en.wikipedia.org/wiki/Dipole#Magnetic_vector_potential>'
:param numpy.ndarray srcLoc: Location of the source(s) (x, y, z)
:param numpy.ndarray,SimPEG.Mesh obsLoc: Where the potentials will be
calculated (x, y, z) or a
SimPEG Mesh
:param str,list component: The component to calculate - 'x', 'y', or
'z' if an array, or grid type if mesh, can
be a list
:param numpy.ndarray orientation: The vector dipole moment
:rtype: numpy.ndarray
:return: The vector potential each dipole at each observation location
"""
# TODO: break this out!
if isinstance(orientation, str):
orientation = orientationDict[orientation]
assert np.linalg.norm(np.array(orientation), 2) == 1., ("orientation must "
"be a unit vector")
if type(component) in [list, tuple]:
out = range(len(component))
for i, comp in enumerate(component):
out[i] = MagneticDipoleVectorPotential(srcLoc, obsLoc, comp,
orientation=orientation,
mu=mu)
return np.concatenate(out)
if isinstance(obsLoc, Mesh.BaseMesh):
mesh = obsLoc
assert component in ['Ex', 'Ey', 'Ez', 'Fx', 'Fy', 'Fz'], ("Components"
"must be in: ['Ex','Ey','Ez','Fx','Fy','Fz']")
return MagneticDipoleVectorPotential(srcLoc, getattr(mesh, 'grid' +
component),
component[1],
orientation=orientation)
if component == 'x':
dimInd = 0
elif component == 'y':
dimInd = 1
elif component == 'z':
dimInd = 2
else:
raise ValueError('Invalid component')
srcLoc = np.atleast_2d(srcLoc)
obsLoc = np.atleast_2d(obsLoc)
orientation = np.atleast_2d(orientation)
nObs = obsLoc.shape[0]
nSrc = srcLoc.shape[0]
m = moment*np.array(orientation).repeat(nObs, axis=0)
A = np.empty((nObs, nSrc))
for i in range(nSrc):
dR = obsLoc - srcLoc[i, np.newaxis].repeat(nObs, axis=0)
mCr = np.cross(m, dR)
r = np.sqrt((dR**2).sum(axis=1))
A[:, i] = +(mu/(4*np.pi)) * mCr[:, dimInd]/(r**3)
if nSrc == 1:
return A.flatten()
return A
def MagneticDipoleVectorPotential(srcLoc, obsLoc, component, moment=1.,
orientation=np.r_[0., 0., 1.],
mu=mu_0):
"""
Calculate the vector potential of a set of magnetic dipoles
at given locations 'ref. <http://en.wikipedia.org/wiki/Dipole#Magnetic_vector_potential>'
:param numpy.ndarray srcLoc: Location of the source(s) (x, y, z)
:param numpy.ndarray,SimPEG.Mesh obsLoc: Where the potentials will be
calculated (x, y, z) or a
SimPEG Mesh
:param str,list component: The component to calculate - 'x', 'y', or
'z' if an array, or grid type if mesh, can
be a list
:param numpy.ndarray orientation: The vector dipole moment
:rtype: numpy.ndarray
:return: The vector potential each dipole at each observation location
"""
# TODO: break this out!
if isinstance(orientation, str):
orientation = orientationDict[orientation]
assert np.linalg.norm(np.array(orientation), 2) == 1., ("orientation must "
"be a unit vector")
if type(component) in [list, tuple]:
out = list(range(len(component)))
for i, comp in enumerate(component):
out[i] = MagneticDipoleVectorPotential(srcLoc, obsLoc, comp,
orientation=orientation,
mu=mu)
return np.concatenate(out)
if isinstance(obsLoc, Mesh.BaseMesh):
mesh = obsLoc
assert component in ['Ex', 'Ey', 'Ez', 'Fx', 'Fy', 'Fz'], ("Components"
"must be in: ['Ex','Ey','Ez','Fx','Fy','Fz']")
return MagneticDipoleVectorPotential(srcLoc, getattr(mesh, 'grid' +
component),
component[1],
orientation=orientation)
if component == 'x':
dimInd = 0
elif component == 'y':
dimInd = 1
elif component == 'z':
dimInd = 2
else:
raise ValueError('Invalid component')
srcLoc = np.atleast_2d(srcLoc)
obsLoc = np.atleast_2d(obsLoc)
orientation = np.atleast_2d(orientation)
nObs = obsLoc.shape[0]
nSrc = srcLoc.shape[0]
m = moment*np.array(orientation).repeat(nObs, axis=0)
A = np.empty((nObs, nSrc))
for i in range(nSrc):
dR = obsLoc - srcLoc[i, np.newaxis].repeat(nObs, axis=0)
mCr = np.cross(m, dR)
r = np.sqrt((dR**2).sum(axis=1))
A[:, i] = +(mu/(4*np.pi)) * mCr[:, dimInd]/(r**3)
if nSrc == 1:
return A.flatten()
return A
def _refine3D(self):
# .----------------.----------------.
# /| /| /|
# / | / | / |
# / | 011 / | 111 / |
# / | / | / |
# .----------------.----+-----------. |
# /| . ---------/|----.----------/|----.
# / | /| / | /| / | /|
# / | / | 001 / | / | 101 / | / |
# / | / | / | / | / | / |
# . -------------- .----------------. |/ |
# | . ---+------|----.----+------|----. |
# | /| .______|___/|____.______|___/|____.
# | / | / 010 | / | / 110| / | /
# | / | / | / | / | / | /
# . ---+---------- . ---+---------- . | /
# | |/ | |/ | |/ z
# | . ----------|----.-----------|----. ^ y
# | / 000 | / 100 | / | /
# | / | / | / | /
# | / | / | / o----> x
# . -------------- . -------------- .
#
#
# Face Refinement:
#
# 2_______________3 _______________
# | | | | |
# ^ | | | (0,1) | (1,1) |
# | | | | | |
# | | x | ---> |-------+-------|
# t1 | | | | |
# | | | (0,0) | (1,0) |
# |_______________| |_______|_______|
# 0 t0--> 1
order = [{'c':[0,0,0],
'fXm': ('p', 'fXm', [0,0]), 'fXp': 'new' ,
'fYm': ('p', 'fYm', [0,0]), 'fYp': 'new' ,
'fZm': ('p', 'fZm', [0,0]), 'fZp': 'new' ,},
{'c':[1,0,0],
'fXm': ('c', 'fXp', [0,0,0]), 'fXp': ('p', 'fXp', [0,0]),
'fYm': ('p', 'fYm', [1,0]), 'fYp': 'new' ,
'fZm': ('p', 'fZm', [1,0]), 'fZp': 'new' },
{'c':[0,1,0],
'fXm': ('p', 'fXm', [1,0]), 'fXp': 'new' ,
'fYm': ('c', 'fYp', [0,0,0]), 'fYp': ('p', 'fYp', [0,0]),
'fZm': ('p', 'fZm', [0,1]), 'fZp': 'new' },
{'c':[1,1,0],
'fXm': ('c', 'fXp', [0,1,0]), 'fXp': ('p', 'fXp', [1,0]),
'fYm': ('c', 'fYp', [1,0,0]), 'fYp': ('p', 'fYp', [1,0]),
'fZm': ('p', 'fZm', [1,1]), 'fZp': 'new' },
{'c':[0,0,1],
'fXm': ('p', 'fXm', [0,1]), 'fXp': 'new' ,
'fYm': ('p', 'fYm', [0,1]), 'fYp': 'new' ,
'fZm': ('c', 'fZp', [0,0,0]), 'fZp': ('p', 'fZp', [0,0])},
{'c':[1,0,1],
'fXm': ('c', 'fXp', [0,0,1]), 'fXp': ('p', 'fXp', [0,1]),
'fYm': ('p', 'fYm', [1,1]), 'fYp': 'new' ,
'fZm': ('c', 'fZp', [1,0,0]), 'fZp': ('p', 'fZp', [1,0])},
{'c':[0,1,1],
'fXm': ('p', 'fXm', [1,1]), 'fXp': 'new' ,
'fYm': ('c', 'fYp', [0,0,1]), 'fYp': ('p', 'fYp', [0,1]),
'fZm': ('c', 'fZp', [0,1,0]), 'fZp': ('p', 'fZp', [0,1])},
{'c':[1,1,1],
'fXm': ('c', 'fXp', [0,1,1]), 'fXp': ('p', 'fXp', [1,1]),
'fYm': ('c', 'fYp', [1,0,1]), 'fYp': ('p', 'fYp', [1,1]),
'fZm': ('c', 'fZp', [1,1,0]), 'fZp': ('p', 'fZp', [1,1])}]
self.mesh.isNumbered = False
self.children = np.empty((2,2,2), dtype=TreeCell)
x0, sz = self.x0, self.sz
for face in self.faceList:
face.refine()
def getFace(pointer):
if pointer is 'new': return None
if pointer[0] == 'p':
return self.faceDict[pointer[1]].children[pointer[2][0],pointer[2][1]]
if pointer[0] == 'c':
return self.children[pointer[2][0],pointer[2][1],pointer[2][2]].faceDict[pointer[1]]
for O in order:
i, j, k = O['c']
x0r = np.r_[x0[0] + 0.5*i*sz[0], x0[1] + 0.5*j*sz[1], x0[2] + 0.5*k*sz[2]]
fXm, fXp, fYm, fYp, fZm, fZp = getFace(O['fXm']), getFace(O['fXp']), getFace(O['fYm']), getFace(O['fYp']), getFace(O['fZm']), getFace(O['fZp'])
self.children[i,j,k] = TreeCell(self.mesh, x0=x0r, depth=self.depth+1, sz=0.5*sz, fXm=fXm, fXp=fXp, fYm=fYm, fYp=fYp, fZm=fZm, fZp=fZp)
self.mesh.cells.remove(self)
