mirror of
https://github.com/wassname/simpeg.git
synced 2026-07-08 18:56:57 +08:00
Merge branch 'master' of https://bitbucket.org/rcockett/simpeg
This commit is contained in:
Lars Ruthotto
@@ -0,0 +1,87 @@
|
||||
from numpy import *
|
||||
import numpy as np
|
||||
|
||||
|
||||
def diff(A, d):
|
||||
if(d == 1):
|
||||
return A[1:, :, :] - A[:-1, :, :]
|
||||
elif(d == 2):
|
||||
return A[:, 1:, :] - A[:, :-1, :]
|
||||
else:
|
||||
return A[:, :, 1:] - A[:, :, :-1]
|
||||
#else:
|
||||
# print('d must be 1,2 or 3')
|
||||
|
||||
|
||||
def diffp(A, d1, d2):
|
||||
if(d1 == 1 and d2 == 2):
|
||||
return A[1:, 1:, :] - A[:-1, :-1, :]
|
||||
elif(d1 == 1 and d2 == 3):
|
||||
return A[1:, :, 1:] - A[:-1, :, :-1]
|
||||
else:
|
||||
return A[:, 1:, 1:] - A[:, :-1, :-1]
|
||||
|
||||
|
||||
def diffm(A, d1, d2):
|
||||
if(d1 == 3 and d2 == 2):
|
||||
return A[:, :-1, 1:] - A[:, 1:, :-1]
|
||||
elif(d1 == 1 and d2 == 3):
|
||||
return A[1:, :, :-1] - A[:-1, :, 1:]
|
||||
elif(d1 == 2 and d2 == 1):
|
||||
return A[:-1, 1:, :] - A[1:, :-1, :]
|
||||
else:
|
||||
print('d must be 1, 2 or 3')
|
||||
|
||||
|
||||
def ave(A, d):
|
||||
if(d == 1):
|
||||
return 0.5*(A[1:, :, :] + A[:-1, :, :])
|
||||
elif(d == 2):
|
||||
return 0.5*(A[:, 1:, :] + A[:, :-1, :])
|
||||
elif(d == 3):
|
||||
return 0.5*(A[:, :, 1:] + A[:, :, :-1])
|
||||
else:
|
||||
print('d must be 1,2 or 3')
|
||||
|
||||
|
||||
def mkmat(x):
|
||||
return reshape(matrix(x), (size(x), 1), 'F')
|
||||
|
||||
|
||||
def hstack3(a, b, c):
|
||||
a = mkvc(a)
|
||||
b = mkvc(b)
|
||||
c = mkvc(c)
|
||||
a = mkmat(a)
|
||||
b = mkmat(b)
|
||||
c = mkmat(c)
|
||||
return hstack((hstack((a, b)), c))
|
||||
|
||||
|
||||
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 = array(mult).reshape(len(mult))
|
||||
|
||||
sub = []
|
||||
|
||||
for i in range(0, len(shape)):
|
||||
sub.extend([math.floor(ind / mult[i])])
|
||||
ind = ind - (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 = array(mult).reshape(len(mult), 1)
|
||||
|
||||
idx = dot((subs), (mult))
|
||||
return idx
|
||||
@@ -1,144 +0,0 @@
|
||||
import numpy as np
|
||||
from utils import ndgrid
|
||||
|
||||
|
||||
class TensorGrid(object):
|
||||
"""
|
||||
Define nodal, cell-centered and staggered tensor grids for 1, 2 and 3
|
||||
dimensions.
|
||||
|
||||
This class is inherited by TensorMesh
|
||||
"""
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def vectorNx():
|
||||
doc = "Nodal grid vector (1D) in the x direction."
|
||||
fget = lambda self: np.r_[0., self.hx.cumsum()] + self.x0[0]
|
||||
return locals()
|
||||
vectorNx = property(**vectorNx())
|
||||
|
||||
def vectorNy():
|
||||
doc = "Nodal grid vector (1D) in the y direction."
|
||||
fget = lambda self: None if self.dim < 2 else np.r_[0., self.hy.cumsum()] + self.x0[1]
|
||||
return locals()
|
||||
vectorNy = property(**vectorNy())
|
||||
|
||||
def vectorNz():
|
||||
doc = "Nodal grid vector (1D) in the z direction."
|
||||
fget = lambda self: None if self.dim < 3 else np.r_[0., self.hz.cumsum()] + self.x0[2]
|
||||
return locals()
|
||||
vectorNz = property(**vectorNz())
|
||||
|
||||
def vectorCCx():
|
||||
doc = "Cell-centered grid vector (1D) in the x direction."
|
||||
fget = lambda self: np.r_[0, self.hx[:-1].cumsum()] + self.hx*0.5 + self.x0[0]
|
||||
return locals()
|
||||
vectorCCx = property(**vectorCCx())
|
||||
|
||||
def vectorCCy():
|
||||
doc = "Cell-centered grid vector (1D) in the y direction."
|
||||
fget = lambda self: None if self.dim < 2 else np.r_[0, self.hy[:-1].cumsum()] + self.hy*0.5 + self.x0[1]
|
||||
return locals()
|
||||
vectorCCy = property(**vectorCCy())
|
||||
|
||||
def vectorCCz():
|
||||
doc = "Cell-centered grid vector (1D) in the z direction."
|
||||
fget = lambda self: None if self.dim < 3 else np.r_[0, self.hz[:-1].cumsum()] + self.hz*0.5 + self.x0[2]
|
||||
return locals()
|
||||
vectorCCz = property(**vectorCCz())
|
||||
|
||||
def gridCC():
|
||||
doc = "Cell-centered grid."
|
||||
|
||||
def fget(self):
|
||||
if self._gridCC is None:
|
||||
self._gridCC = ndgrid([x for x in [self.vectorCCx, self.vectorCCy, self.vectorCCz] if not x is None])
|
||||
return self._gridCC
|
||||
return locals()
|
||||
_gridCC = None # Store grid by default
|
||||
gridCC = property(**gridCC())
|
||||
|
||||
def gridN():
|
||||
doc = "Nodal grid."
|
||||
|
||||
def fget(self):
|
||||
if self._gridN is None:
|
||||
self._gridN = ndgrid([x for x in [self.vectorNx, self.vectorNy, self.vectorNz] if not x is None])
|
||||
return self._gridN
|
||||
return locals()
|
||||
_gridN = None # Store grid by default
|
||||
gridN = property(**gridN())
|
||||
|
||||
def gridFx():
|
||||
doc = "Face staggered grid in the x direction."
|
||||
|
||||
def fget(self):
|
||||
if self._gridFx is None:
|
||||
self._gridFx = ndgrid([x for x in [self.vectorNx, self.vectorCCy, self.vectorCCz] if not x is None])
|
||||
return self._gridFx
|
||||
return locals()
|
||||
_gridFx = None # Store grid by default
|
||||
gridFx = property(**gridFx())
|
||||
|
||||
def gridFy():
|
||||
doc = "Face staggered grid in the y direction."
|
||||
|
||||
def fget(self):
|
||||
if self._gridFy is None:
|
||||
self._gridFy = ndgrid([x for x in [self.vectorCCx, self.vectorNy, self.vectorCCz] if not x is None])
|
||||
return self._gridFy
|
||||
return locals()
|
||||
_gridFy = None # Store grid by default
|
||||
gridFy = property(**gridFy())
|
||||
|
||||
def gridFz():
|
||||
doc = "Face staggered grid in the z direction."
|
||||
|
||||
def fget(self):
|
||||
if self._gridFz is None:
|
||||
self._gridFz = ndgrid([x for x in [self.vectorCCx, self.vectorCCy, self.vectorNz] if not x is None])
|
||||
return self._gridFz
|
||||
return locals()
|
||||
_gridFz = None # Store grid by default
|
||||
gridFz = property(**gridFz())
|
||||
|
||||
def gridEx():
|
||||
doc = "Edge staggered grid in the x direction."
|
||||
|
||||
def fget(self):
|
||||
if self._gridEx is None:
|
||||
self._gridEx = ndgrid([x for x in [self.vectorCCx, self.vectorNy, self.vectorNz] if not x is None])
|
||||
return self._gridEx
|
||||
return locals()
|
||||
_gridEx = None # Store grid by default
|
||||
gridEx = property(**gridEx())
|
||||
|
||||
def gridEy():
|
||||
doc = "Edge staggered grid in the y direction."
|
||||
|
||||
def fget(self):
|
||||
if self._gridEy is None:
|
||||
self._gridEy = ndgrid([x for x in [self.vectorNx, self.vectorCCy, self.vectorNz] if not x is None])
|
||||
return self._gridEy
|
||||
return locals()
|
||||
_gridEy = None # Store grid by default
|
||||
gridEy = property(**gridEy())
|
||||
|
||||
def gridEz():
|
||||
doc = "Edge staggered grid in the z direction."
|
||||
|
||||
def fget(self):
|
||||
if self._gridEz is None:
|
||||
self._gridEz = ndgrid([x for x in [self.vectorNx, self.vectorNy, self.vectorCCz] if not x is None])
|
||||
return self._gridEz
|
||||
return locals()
|
||||
_gridEz = None # Store grid by default
|
||||
gridEz = property(**gridEz())
|
||||
|
||||
def getBoundaryIndex(self, gridType):
|
||||
"""Needed for faces edges and cells"""
|
||||
pass
|
||||
|
||||
def getCellNumbering(self):
|
||||
pass
|
||||
@@ -2,6 +2,7 @@ import numpy as np
|
||||
from BaseMesh import BaseMesh
|
||||
from TensorGrid import TensorGrid
|
||||
from TensorView import TensorView
|
||||
from utils import ndgrid
|
||||
|
||||
|
||||
class TensorMesh(BaseMesh, TensorGrid, TensorView):
|
||||
@@ -55,6 +56,137 @@ class TensorMesh(BaseMesh, TensorGrid, TensorView):
|
||||
return locals()
|
||||
hz = property(**hz())
|
||||
|
||||
def vectorNx():
|
||||
doc = "Nodal grid vector (1D) in the x direction."
|
||||
fget = lambda self: np.r_[0., self.hx.cumsum()] + self.x0[0]
|
||||
return locals()
|
||||
vectorNx = property(**vectorNx())
|
||||
|
||||
def vectorNy():
|
||||
doc = "Nodal grid vector (1D) in the y direction."
|
||||
fget = lambda self: None if self.dim < 2 else np.r_[0., self.hy.cumsum()] + self.x0[1]
|
||||
return locals()
|
||||
vectorNy = property(**vectorNy())
|
||||
|
||||
def vectorNz():
|
||||
doc = "Nodal grid vector (1D) in the z direction."
|
||||
fget = lambda self: None if self.dim < 3 else np.r_[0., self.hz.cumsum()] + self.x0[2]
|
||||
return locals()
|
||||
vectorNz = property(**vectorNz())
|
||||
|
||||
def vectorCCx():
|
||||
doc = "Cell-centered grid vector (1D) in the x direction."
|
||||
fget = lambda self: np.r_[0, self.hx[:-1].cumsum()] + self.hx*0.5 + self.x0[0]
|
||||
return locals()
|
||||
vectorCCx = property(**vectorCCx())
|
||||
|
||||
def vectorCCy():
|
||||
doc = "Cell-centered grid vector (1D) in the y direction."
|
||||
fget = lambda self: None if self.dim < 2 else np.r_[0, self.hy[:-1].cumsum()] + self.hy*0.5 + self.x0[1]
|
||||
return locals()
|
||||
vectorCCy = property(**vectorCCy())
|
||||
|
||||
def vectorCCz():
|
||||
doc = "Cell-centered grid vector (1D) in the z direction."
|
||||
fget = lambda self: None if self.dim < 3 else np.r_[0, self.hz[:-1].cumsum()] + self.hz*0.5 + self.x0[2]
|
||||
return locals()
|
||||
vectorCCz = property(**vectorCCz())
|
||||
|
||||
def gridCC():
|
||||
doc = "Cell-centered grid."
|
||||
|
||||
def fget(self):
|
||||
if self._gridCC is None:
|
||||
self._gridCC = ndgrid([x for x in [self.vectorCCx, self.vectorCCy, self.vectorCCz] if not x is None])
|
||||
return self._gridCC
|
||||
return locals()
|
||||
_gridCC = None # Store grid by default
|
||||
gridCC = property(**gridCC())
|
||||
|
||||
def gridN():
|
||||
doc = "Nodal grid."
|
||||
|
||||
def fget(self):
|
||||
if self._gridN is None:
|
||||
self._gridN = ndgrid([x for x in [self.vectorNx, self.vectorNy, self.vectorNz] if not x is None])
|
||||
return self._gridN
|
||||
return locals()
|
||||
_gridN = None # Store grid by default
|
||||
gridN = property(**gridN())
|
||||
|
||||
def gridFx():
|
||||
doc = "Face staggered grid in the x direction."
|
||||
|
||||
def fget(self):
|
||||
if self._gridFx is None:
|
||||
self._gridFx = ndgrid([x for x in [self.vectorNx, self.vectorCCy, self.vectorCCz] if not x is None])
|
||||
return self._gridFx
|
||||
return locals()
|
||||
_gridFx = None # Store grid by default
|
||||
gridFx = property(**gridFx())
|
||||
|
||||
def gridFy():
|
||||
doc = "Face staggered grid in the y direction."
|
||||
|
||||
def fget(self):
|
||||
if self._gridFy is None:
|
||||
self._gridFy = ndgrid([x for x in [self.vectorCCx, self.vectorNy, self.vectorCCz] if not x is None])
|
||||
return self._gridFy
|
||||
return locals()
|
||||
_gridFy = None # Store grid by default
|
||||
gridFy = property(**gridFy())
|
||||
|
||||
def gridFz():
|
||||
doc = "Face staggered grid in the z direction."
|
||||
|
||||
def fget(self):
|
||||
if self._gridFz is None:
|
||||
self._gridFz = ndgrid([x for x in [self.vectorCCx, self.vectorCCy, self.vectorNz] if not x is None])
|
||||
return self._gridFz
|
||||
return locals()
|
||||
_gridFz = None # Store grid by default
|
||||
gridFz = property(**gridFz())
|
||||
|
||||
def gridEx():
|
||||
doc = "Edge staggered grid in the x direction."
|
||||
|
||||
def fget(self):
|
||||
if self._gridEx is None:
|
||||
self._gridEx = ndgrid([x for x in [self.vectorCCx, self.vectorNy, self.vectorNz] if not x is None])
|
||||
return self._gridEx
|
||||
return locals()
|
||||
_gridEx = None # Store grid by default
|
||||
gridEx = property(**gridEx())
|
||||
|
||||
def gridEy():
|
||||
doc = "Edge staggered grid in the y direction."
|
||||
|
||||
def fget(self):
|
||||
if self._gridEy is None:
|
||||
self._gridEy = ndgrid([x for x in [self.vectorNx, self.vectorCCy, self.vectorNz] if not x is None])
|
||||
return self._gridEy
|
||||
return locals()
|
||||
_gridEy = None # Store grid by default
|
||||
gridEy = property(**gridEy())
|
||||
|
||||
def gridEz():
|
||||
doc = "Edge staggered grid in the z direction."
|
||||
|
||||
def fget(self):
|
||||
if self._gridEz is None:
|
||||
self._gridEz = ndgrid([x for x in [self.vectorNx, self.vectorNy, self.vectorCCz] if not x is None])
|
||||
return self._gridEz
|
||||
return locals()
|
||||
_gridEz = None # Store grid by default
|
||||
gridEz = property(**gridEz())
|
||||
|
||||
def getBoundaryIndex(self, gridType):
|
||||
"""Needed for faces edges and cells"""
|
||||
pass
|
||||
|
||||
def getCellNumbering(self):
|
||||
pass
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
print('Welcome to tensor mesh!')
|
||||
|
||||
@@ -0,0 +1,75 @@
|
||||
import numpy as np
|
||||
from scipy import sparse
|
||||
from utils import mkvc
|
||||
from sputils import ddx, sdiag, speye, kron3
|
||||
|
||||
def getvol(h):
|
||||
|
||||
# Cell sizes in each direction
|
||||
h1 = h[0]
|
||||
h2 = h[1]
|
||||
h3 = h[2]
|
||||
|
||||
# Compute cell volumes
|
||||
v12 = h1.T*h2
|
||||
V = mkvc(v12.reshape(-1,1)*h3)
|
||||
|
||||
return V
|
||||
|
||||
def getarea(h):
|
||||
|
||||
# Cell sizes in each direction
|
||||
h1 = h[0]
|
||||
h2 = h[1]
|
||||
h3 = h[2]
|
||||
|
||||
# The number of cell centers in each direction
|
||||
n1 = np.size(h1)
|
||||
n2 = np.size(h2)
|
||||
n3 = np.size(h3)
|
||||
# Compute areas of cell faces
|
||||
area1 = np.ones((n1+1,1))*mkvc(h2.T*h3)
|
||||
area2 = h1.T*mkvc(np.ones((n2+1,1))*h3)
|
||||
area3 = h1.T*mkvc(h2.T*np.ones(n3+1))
|
||||
area = np.hstack((np.hstack((mkvc(area1), mkvc(area2))), mkvc(area3)))
|
||||
|
||||
return area
|
||||
|
||||
def getDivMatrix(h):
|
||||
"""Consturct the 3D divergence operator on Faces."""
|
||||
|
||||
# Cell sizes in each direction
|
||||
h1 = h[0]
|
||||
h2 = h[1]
|
||||
h3 = h[2]
|
||||
|
||||
# The number of cell centers in each direction
|
||||
n1 = np.size(h1)
|
||||
n2 = np.size(h2)
|
||||
n3 = np.size(h3)
|
||||
|
||||
# Compute areas of cell faces
|
||||
#area1 = np.ones((n1+1,1))*mkvc(h2.T*h3)
|
||||
#area2 = h1.T*mkvc(np.ones((n2+1,1))*h3)
|
||||
#area3 = h1.T*mkvc(h2.T*np.ones(n3+1))
|
||||
#area = np.hstack((np.hstack((mkvc(area1), mkvc(area2))), mkvc(area3)))
|
||||
area = getarea(h)
|
||||
|
||||
S = sdiag(area)
|
||||
|
||||
# Compute cell volumes
|
||||
#v12 = h1.T*h2
|
||||
#V = mkvc(v12.reshape(-1,1)*h3)
|
||||
V = getvol(h)
|
||||
|
||||
# Compute divergence operator on faces
|
||||
d1 = ddx(n1)
|
||||
d2 = ddx(n2)
|
||||
d3 = ddx(n3)
|
||||
D1 = kron3(speye(n3), speye(n2), d1)
|
||||
D2 = kron3(speye(n3), d2, speye(n1))
|
||||
D3 = kron3(d3, speye(n2), speye(n1))
|
||||
|
||||
D = sparse.hstack((sparse.hstack((D1, D2)), D3))
|
||||
return sdiag(1/V)*D*S
|
||||
|
||||
@@ -0,0 +1,18 @@
|
||||
import numpy as np
|
||||
from scipy import sparse
|
||||
|
||||
def ddx(n):
|
||||
"""Define 1D derivatives"""
|
||||
return sparse.spdiags((np.ones((n+1,1))*[-1,1]).T, [0,1], n, n+1)
|
||||
|
||||
def sdiag(h):
|
||||
"""Sparse diagonal matrix"""
|
||||
return sparse.spdiags(h, 0, np.size(h), np.size(h))
|
||||
|
||||
def speye(n):
|
||||
"""Sparse identity"""
|
||||
return sparse.identity(n)
|
||||
|
||||
def kron3(A, B, C):
|
||||
"""Two kron prods"""
|
||||
return sparse.kron(sparse.kron(A, B), C)
|
||||
@@ -0,0 +1,45 @@
|
||||
import numpy as np
|
||||
|
||||
import sys
|
||||
sys.path.append('../')
|
||||
from TensorMesh import TensorMesh
|
||||
from getDIV import getDivMatrix, getarea, getvol
|
||||
|
||||
# Define the mesh
|
||||
err=0.
|
||||
for i in range(4):
|
||||
icount=i+1;
|
||||
nc = 2*icount;
|
||||
h1 = np.pi/nc*np.ones((1,nc))
|
||||
h2 = np.pi/nc*np.ones((1,nc))
|
||||
h3 = np.pi/nc*np.ones((1,nc))
|
||||
h = [h1, h2, h3]
|
||||
x0 = -np.pi/2*np.ones((3, 1))
|
||||
M = TensorMesh(h, x0)
|
||||
#n = M.plotGrid()
|
||||
|
||||
# Generate DIV matrix
|
||||
DIV = getDivMatrix(h)
|
||||
|
||||
#Test function
|
||||
fun = lambda x: np.sin(x)
|
||||
Fx = fun(M.gridFx[:,0])
|
||||
Fy = fun(M.gridFy[:,1])
|
||||
Fz = fun(M.gridFz[:,2])
|
||||
|
||||
F = np.concatenate((Fx,Fy,Fz))
|
||||
divF = DIV*F
|
||||
sol = lambda x, y, z: (np.cos(x)+np.cos(y)+np.cos(z))
|
||||
divF_anal = sol(M.gridCC[:,0], M.gridCC[:,1], M.gridCC[:,2])
|
||||
|
||||
area = getarea(h)
|
||||
vol = getvol(h)
|
||||
err = np.linalg.norm((divF-divF_anal)*np.sqrt(vol), 2)
|
||||
if icount == 1:
|
||||
err1 = err
|
||||
print 'h | 2 norm | error ratio'
|
||||
print '---------------------------------------'
|
||||
print '%6.4f | %8.2e |'% (h1[0,0], err)
|
||||
else:
|
||||
print '%6.4f | %8.2e | %6.4f' % (h1[0,0], err, err1/err)
|
||||
|
||||
+56
-104
@@ -1,49 +1,6 @@
|
||||
from numpy import *
|
||||
import numpy as np
|
||||
|
||||
|
||||
def diff(A, d):
|
||||
if(d == 1):
|
||||
return A[1:, :, :] - A[:-1, :, :]
|
||||
elif(d == 2):
|
||||
return A[:, 1:, :] - A[:, :-1, :]
|
||||
else:
|
||||
return A[:, :, 1:] - A[:, :, :-1]
|
||||
#else:
|
||||
# print('d must be 1,2 or 3')
|
||||
|
||||
|
||||
def diffp(A, d1, d2):
|
||||
if(d1 == 1 and d2 == 2):
|
||||
return A[1:, 1:, :] - A[:-1, :-1, :]
|
||||
elif(d1 == 1 and d2 == 3):
|
||||
return A[1:, :, 1:] - A[:-1, :, :-1]
|
||||
else:
|
||||
return A[:, 1:, 1:] - A[:, :-1, :-1]
|
||||
|
||||
|
||||
def diffm(A, d1, d2):
|
||||
if(d1 == 3 and d2 == 2):
|
||||
return A[:, :-1, 1:] - A[:, 1:, :-1]
|
||||
elif(d1 == 1 and d2 == 3):
|
||||
return A[1:, :, :-1] - A[:-1, :, 1:]
|
||||
elif(d1 == 2 and d2 == 1):
|
||||
return A[:-1, 1:, :] - A[1:, :-1, :]
|
||||
else:
|
||||
print('d must be 1, 2 or 3')
|
||||
|
||||
|
||||
def ave(A, d):
|
||||
if(d == 1):
|
||||
return 0.5*(A[1:, :, :] + A[:-1, :, :])
|
||||
elif(d == 2):
|
||||
return 0.5*(A[:, 1:, :] + A[:, :-1, :])
|
||||
elif(d == 3):
|
||||
return 0.5*(A[:, :, 1:] + A[:, :, :-1])
|
||||
else:
|
||||
print('d must be 1,2 or 3')
|
||||
|
||||
|
||||
def reshapeF(x, size):
|
||||
return np.reshape(x, size, order='F')
|
||||
|
||||
@@ -53,7 +10,7 @@ def mkvc(x, numDims=1):
|
||||
|
||||
e.g.:
|
||||
|
||||
a = np.array(1,2,3)
|
||||
a = np.array([1, 2, 3])
|
||||
|
||||
mkvc(a, 1).shape
|
||||
> (3, )
|
||||
@@ -75,8 +32,45 @@ def mkvc(x, numDims=1):
|
||||
return x.flatten(order='F')[:, np.newaxis, np.newaxis]
|
||||
|
||||
|
||||
def ndgrid(*args):
|
||||
"""Form tensorial grid for 1, 2 and 3 dimensions. Return X1,X2,X3 arrays depending on the dimension"""
|
||||
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:
|
||||
@@ -84,64 +78,22 @@ def ndgrid(*args):
|
||||
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
|
||||
return xin[0]
|
||||
elif len(xin) == 2:
|
||||
X2, X1 = [mkvc(x) for x in np.broadcast_arrays(mkvc(xin[1], 1), mkvc(xin[0], 2))]
|
||||
return np.c_[X1, X2]
|
||||
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:
|
||||
X3, X2, X1 = [mkvc(x) for x in np.broadcast_arrays(mkvc(xin[2], 1), mkvc(xin[1], 2), mkvc(xin[0], 3))]
|
||||
return np.c_[X1, X2, X3]
|
||||
|
||||
|
||||
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 = array(mult).reshape(len(mult))
|
||||
|
||||
sub = []
|
||||
|
||||
for i in range(0, len(shape)):
|
||||
sub.extend([math.floor(ind / mult[i])])
|
||||
ind = ind - (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 = array(mult).reshape(len(mult), 1)
|
||||
|
||||
idx = dot((subs), (mult))
|
||||
return idx
|
||||
|
||||
|
||||
def mkmat(x):
|
||||
return reshape(matrix(x), (size(x), 1), 'F')
|
||||
|
||||
|
||||
def hstack3(a, b, c):
|
||||
a = mkvc(a)
|
||||
b = mkvc(b)
|
||||
c = mkvc(c)
|
||||
a = mkmat(a)
|
||||
b = mkmat(b)
|
||||
c = mkmat(c)
|
||||
return hstack((hstack((a, b)), c))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
X, Y, Z = mgrid[0:4, 0:5, 0:6]
|
||||
|
||||
print Z
|
||||
|
||||
t = ave(X, 1)
|
||||
print t
|
||||
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]
|
||||
|
||||
Reference in New Issue
Block a user