Files
simpeg/code/utils.py
T

124 lines
2.8 KiB
Python

from numpy import *
def diff(A,d):
end = -1
if(d==1):
return A[1:,0:,0:] - A[0:end,0:,0:]
elif(d==2):
return A[0:,1:,0:] - A[0:,0:end,0:]
else:
return A[0:,0:,1:] - A[0:,0:,0:end]
#else:
# print('d must be 1,2 or 3')
def diffp(A, d1, d2):
end = -1
if(d1 == 1 and d2 == 2 ):
return A[1:,1:, 0:] - A[0:end,0:end,0:]
elif(d1 == 1 and d2 == 3):
return A[1:,0:,1:] - A[0:end,0:,0:end]
else:
return A[0:,1:,1:] - A[0:,0:end,0:end]
def diffm(A, d1, d2):
end = -1
if(d1 == 3 and d2 == 2 ):
return A[:,0:end,1:] - A[:,1:,0:end]
elif(d1 == 1 and d2 == 3):
return A[1:, :, 0:end] - A[0:end,:,1:]
elif(d1 == 2 and d2 == 1):
return A[0:end, 1:, :] - A[1:, 0:end, :]
else:
print('d must be 1,2 or 3')
def ave(A,d):
end = 0
if(d==1):
return 0.5*(A[1:,:,:] + A[0:end-1,0:,:])
elif(d==2):
return 0.5*(A[:,1:,:] + A[0:,0:end-1,:])
elif(d==3):
return 0.5*(A[:,:,1:] + A[0:,0:,0:end-1])
else:
print('d must be 1,2 or 3')
def reshapeF(sp,d):
return reshape(sp,d,'F')
def mkvc(A):
return reshape(A,[size(A),1],'F').flatten()
def ndgrid(x,y,z):
n1 = size(x)
n2 = size(y)
n3 = size(z)
X = zeros([n1,n2,n3])
Y = zeros([n1,n2,n3])
Z = zeros([n1,n2,n3])
for i in range(0, n2):
for j in range(0,n3):
X[:,i,j] = x
for i in range(0, n1):
for j in range(0,n3):
Y[i,:,j] = y
for i in range(0, n1):
for j in range(0,n2):
Z[i,j,:] = z
return (X,Y,Z)
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]
t = ave(X, 1)