mirror of
https://github.com/wassname/simpeg.git
synced 2026-06-29 00:55:56 +08:00
Rowan Cockett
331 lines
9.1 KiB
Python
331 lines
9.1 KiB
Python
import numpy as np
|
|
|
|
|
|
def mkvc(x, numDims=1):
|
|
"""Creates a vector with the number of dimension specified
|
|
|
|
e.g.:
|
|
|
|
a = np.array([1, 2, 3])
|
|
|
|
mkvc(a, 1).shape
|
|
> (3, )
|
|
|
|
mkvc(a, 2).shape
|
|
> (3, 1)
|
|
|
|
mkvc(a, 3).shape
|
|
> (3, 1, 1)
|
|
|
|
"""
|
|
if type(x) == np.matrix:
|
|
x = np.array(x)
|
|
|
|
assert type(x) == np.ndarray, "Vector must be a numpy array"
|
|
|
|
if numDims == 1:
|
|
return x.flatten(order='F')
|
|
elif numDims == 2:
|
|
return x.flatten(order='F')[:, np.newaxis]
|
|
elif numDims == 3:
|
|
return x.flatten(order='F')[:, np.newaxis, np.newaxis]
|
|
|
|
|
|
def ndgrid(*args, **kwargs):
|
|
"""
|
|
Form tensorial grid for 1, 2, or 3 dimensions.
|
|
|
|
Returns as column vectors by default.
|
|
|
|
To return as matrix input:
|
|
|
|
ndgrid(..., vector=False)
|
|
|
|
The inputs can be a list or separate arguments.
|
|
|
|
e.g.
|
|
|
|
a = np.array([1, 2, 3])
|
|
b = np.array([1, 2])
|
|
|
|
XY = ndgrid(a, b)
|
|
> [[1 1]
|
|
[2 1]
|
|
[3 1]
|
|
[1 2]
|
|
[2 2]
|
|
[3 2]]
|
|
|
|
X, Y = ndgrid(a, b, vector=False)
|
|
> X = [[1 1]
|
|
[2 2]
|
|
[3 3]]
|
|
> Y = [[1 2]
|
|
[1 2]
|
|
[1 2]]
|
|
|
|
"""
|
|
|
|
# Read the keyword arguments, and only accept a vector=True/False
|
|
vector = kwargs.pop('vector', True)
|
|
assert type(vector) == bool, "'vector' keyword must be a bool"
|
|
assert len(kwargs) == 0, "Only 'vector' keyword accepted"
|
|
|
|
# you can either pass a list [x1, x2, x3] or each seperately
|
|
if type(args[0]) == list:
|
|
xin = args[0]
|
|
else:
|
|
xin = args
|
|
|
|
# Each vector needs to be a numpy array
|
|
assert np.all([type(x) == np.ndarray for x in xin]), "All vectors must be numpy arrays."
|
|
|
|
if len(xin) == 1:
|
|
return xin[0]
|
|
elif len(xin) == 2:
|
|
XY = np.broadcast_arrays(mkvc(xin[1], 1), mkvc(xin[0], 2))
|
|
if vector:
|
|
X2, X1 = [mkvc(x) for x in XY]
|
|
return np.c_[X1, X2]
|
|
else:
|
|
return XY[1], XY[0]
|
|
elif len(xin) == 3:
|
|
XYZ = np.broadcast_arrays(mkvc(xin[2], 1), mkvc(xin[1], 2), mkvc(xin[0], 3))
|
|
if vector:
|
|
X3, X2, X1 = [mkvc(x) for x in XYZ]
|
|
return np.c_[X1, X2, X3]
|
|
else:
|
|
return XYZ[2], XYZ[1], XYZ[0]
|
|
|
|
|
|
def volTetra(xyz, A, B, C, D):
|
|
"""
|
|
Returns the volume for tetrahedras volume specified by the indexes A to D.
|
|
|
|
|
|
Input:
|
|
xyz - X,Y,Z vertex vector
|
|
A,B,C,D - vert index of the tetrahedra
|
|
|
|
Output:
|
|
V - volume
|
|
|
|
Algorithm: http://en.wikipedia.org/wiki/Tetrahedron#Volume
|
|
|
|
V = 1/3 A * h
|
|
|
|
V = 1/6 | ( a - d ) o ( ( b - d ) X ( c - d ) ) |
|
|
|
|
"""
|
|
|
|
AD = xyz[A, :] - xyz[D, :]
|
|
BD = xyz[B, :] - xyz[D, :]
|
|
CD = xyz[C, :] - xyz[D, :]
|
|
|
|
V = (BD[:, 0]*CD[:, 1] - BD[:, 1]*CD[:, 0])*AD[:, 2] - (BD[:, 0]*CD[:, 2] - BD[:, 2]*CD[:, 0])*AD[:, 1] + (BD[:, 1]*CD[:, 2] - BD[:, 2]*CD[:, 1])*AD[:, 0]
|
|
return V/6
|
|
|
|
|
|
def indexCube(nodes, gridSize, n=None):
|
|
"""
|
|
Returns the index of nodes on the mesh.
|
|
|
|
|
|
Input:
|
|
nodes - string of which nodes to return. e.g. 'ABCD'
|
|
gridSize - size of the nodal grid
|
|
n - number of nodes each i,j,k direction: [ni,nj,nk]
|
|
|
|
|
|
Output:
|
|
index - index in the order asked e.g. 'ABCD' --> (A,B,C,D)
|
|
|
|
TWO DIMENSIONS:
|
|
|
|
node(i,j) node(i,j+1)
|
|
A -------------- B
|
|
| |
|
|
| cell(i,j) |
|
|
| I |
|
|
| |
|
|
D -------------- C
|
|
node(i+1,j) node(i+1,j+1)
|
|
|
|
|
|
THREE DIMENSIONS:
|
|
|
|
node(i,j,k+1) node(i,j+1,k+1)
|
|
E --------------- F
|
|
/| / |
|
|
/ | / |
|
|
/ | / |
|
|
node(i,j,k) node(i,j+1,k)
|
|
A -------------- B |
|
|
| H ----------|---- G
|
|
| /cell(i,j) | /
|
|
| / I | /
|
|
| / | /
|
|
D -------------- C
|
|
node(i+1,j,k) node(i+1,j+1,k)
|
|
|
|
|
|
@author Rowan Cockett
|
|
|
|
Last modified on: 2013/07/26
|
|
"""
|
|
|
|
assert type(nodes) == str, "Nodes must be a str variable: e.g. 'ABCD'"
|
|
assert type(gridSize) == np.ndarray, "Number of nodes must be an ndarray"
|
|
nodes = nodes.upper()
|
|
# Make sure that we choose from the possible nodes.
|
|
possibleNodes = 'ABCD' if gridSize.size == 2 else 'ABCDEFGH'
|
|
for node in nodes:
|
|
assert node in possibleNodes, "Nodes must be chosen from: '%s'" % possibleNodes
|
|
dim = gridSize.size
|
|
if n is None:
|
|
n = gridSize - 1
|
|
|
|
if dim == 2:
|
|
ij = ndgrid(np.arange(n[0]), np.arange(n[1]))
|
|
i, j = ij[:, 0], ij[:, 1]
|
|
elif dim == 3:
|
|
ijk = ndgrid(np.arange(n[0]), np.arange(n[1]), np.arange(n[2]))
|
|
i, j, k = ijk[:, 0], ijk[:, 1], ijk[:, 2]
|
|
else:
|
|
raise Exception('Only 2 and 3 dimensions supported.')
|
|
|
|
nodeMap = {'A': [0, 0, 0], 'B': [0, 1, 0], 'C': [1, 1, 0], 'D': [1, 0, 0],
|
|
'E': [0, 0, 1], 'F': [0, 1, 1], 'G': [1, 1, 1], 'H': [1, 0, 1]}
|
|
out = ()
|
|
for node in nodes:
|
|
shift = nodeMap[node]
|
|
if dim == 2:
|
|
out += (sub2ind(gridSize, np.c_[i+shift[0], j+shift[1]]).flatten(), )
|
|
elif dim == 3:
|
|
out += (sub2ind(gridSize, np.c_[i+shift[0], j+shift[1], k+shift[2]]).flatten(), )
|
|
|
|
return out
|
|
|
|
|
|
def faceInfo(xyz, A, B, C, D, average=True, normalizeNormals=True):
|
|
"""
|
|
function [N] = faceInfo(y,A,B,C,D)
|
|
|
|
Returns the averaged normal, area, and edge lengths for a given set of faces.
|
|
|
|
If average option is FALSE then N is a cell array {nA,nB,nC,nD}
|
|
|
|
|
|
Input:
|
|
xyz - X,Y,Z vertex vector
|
|
A,B,C,D - vert index of the face (counter clockwize)
|
|
|
|
Options:
|
|
average - [true]/false, toggles returning all normals or the average
|
|
|
|
Output:
|
|
N - average face normal or {nA,nB,nC,nD} if average = false
|
|
area - average face area
|
|
edgeLengths - exact edge Lengths, 4 column vector [AB, BC, CD, DA]
|
|
|
|
see also testFaceNormal testFaceArea
|
|
|
|
@author Rowan Cockett
|
|
|
|
Last modified on: 2013/07/26
|
|
|
|
"""
|
|
assert type(average) is bool, 'average must be a boolean'
|
|
assert type(normalizeNormals) is bool, 'normalizeNormals must be a boolean'
|
|
# compute normal that is pointing away from you.
|
|
#
|
|
# A -------A-B------- B
|
|
# | |
|
|
# | |
|
|
# D-A (X) B-C
|
|
# | |
|
|
# | |
|
|
# D -------C-D------- C
|
|
|
|
AB = xyz[B, :] - xyz[A, :]
|
|
BC = xyz[C, :] - xyz[B, :]
|
|
CD = xyz[D, :] - xyz[C, :]
|
|
DA = xyz[A, :] - xyz[D, :]
|
|
|
|
def cross(X, Y):
|
|
return np.c_[X[:, 1]*Y[:, 2] - X[:, 2]*Y[:, 1],
|
|
X[:, 2]*Y[:, 0] - X[:, 0]*Y[:, 2],
|
|
X[:, 0]*Y[:, 1] - X[:, 1]*Y[:, 0]]
|
|
|
|
nA = cross(AB, DA)
|
|
nB = cross(BC, AB)
|
|
nC = cross(CD, BC)
|
|
nD = cross(DA, CD)
|
|
|
|
length = lambda x: np.sqrt(x[:, 0]**2 + x[:, 1]**2 + x[:, 2]**2)
|
|
normalize = lambda x: x/np.kron(np.ones((1, x.shape[1])), mkvc(length(x), 2))
|
|
if average:
|
|
# average the normals at each vertex.
|
|
N = (nA + nB + nC + nD)/4 # this is intrinsically weighted by area
|
|
# normalize
|
|
N = normalize(N)
|
|
else:
|
|
if normalizeNormals:
|
|
N = [normalize(nA), normalize(nB), normalize(nC), normalize(nD)]
|
|
else:
|
|
N = [nA, nB, nC, nD]
|
|
|
|
# Area calculation
|
|
#
|
|
# Approximate by 4 different triangles, and divide by 2.
|
|
# Each triangle is one half of the length of the cross product
|
|
#
|
|
# So also could be viewed as the average parallelogram.
|
|
#
|
|
# WARNING: This does not compute correctly for concave quadrilaterals
|
|
area = (length(nA)+length(nB)+length(nC)+length(nD))/4
|
|
|
|
return N, area
|
|
|
|
|
|
def ind2sub(shape, ind):
|
|
"""From the given shape, returns the subscrips of the given index"""
|
|
revshp = []
|
|
revshp.extend(shape)
|
|
mult = [1]
|
|
for i in range(0, len(revshp)-1):
|
|
mult.extend([mult[i]*revshp[i]])
|
|
mult = np.array(mult).reshape(len(mult))
|
|
|
|
sub = []
|
|
|
|
for i in range(0, len(shape)):
|
|
sub.extend([np.math.floor(ind / mult[i])])
|
|
ind = ind - (np.math.floor(ind/mult[i]) * mult[i])
|
|
return sub
|
|
|
|
|
|
def sub2ind(shape, subs):
|
|
"""From the given shape, returns the index of the given subscript"""
|
|
revshp = list(shape)
|
|
mult = [1]
|
|
for i in range(0, len(revshp)-1):
|
|
mult.extend([mult[i]*revshp[i]])
|
|
mult = np.array(mult).reshape(len(mult), 1)
|
|
|
|
idx = np.dot((subs), (mult))
|
|
return idx
|
|
|
|
|
|
def getSubArray(A, ind):
|
|
"""subArray"""
|
|
assert type(ind) == list, "ind must be a list of vectors"
|
|
assert len(A.shape) == len(ind), "ind must have the same length as the dimension of A"
|
|
|
|
if len(A.shape) == 2:
|
|
return A[ind[0], :][:, ind[1]]
|
|
elif len(A.shape) == 3:
|
|
return A[ind[0], :, :][:, ind[1], :][:, :, ind[2]]
|
|
else:
|
|
raise Exception("getSubArray does not support dimension asked.")
|