mirror of
https://github.com/wassname/simpeg.git
synced 2026-07-10 01:56:12 +08:00
Merge branch 'master' of https://bitbucket.org/rcockett/simpeg
This commit is contained in:
Lars Ruthotto
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
File diff suppressed because it is too large
Load Diff
Binary file not shown.
Binary file not shown.
|
After Width: | Height: | Size: 80 KiB |
Binary file not shown.
@@ -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