mirror of
https://github.com/wassname/simpeg.git
synced 2026-06-29 01:59:48 +08:00
Put utils in their own folder. Increases organization, and everything will be available under SimPEG.utils
NOTE: if you add a new function (or file), you must register it in __init__.py for it to be available under: e.g. ... from utils import mkvc, … ... from MYNEWFILE import MYNEWFUNCTION
This commit is contained in:
Rowan Cockett
@@ -1,7 +1,6 @@
|
||||
import numpy as np
|
||||
from scipy import sparse as sp
|
||||
from sputils import sdiag, speye, kron3, spzeros
|
||||
from utils import mkvc
|
||||
from utils import mkvc, sdiag, speye, kron3, spzeros
|
||||
|
||||
|
||||
def ddx(n):
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
from scipy import sparse as sp
|
||||
from sputils import sdiag, inv3X3BlockDiagonal, inv2X2BlockDiagonal
|
||||
from utils import sub2ind, ndgrid, mkvc, getSubArray
|
||||
from utils import sub2ind, ndgrid, mkvc, getSubArray, sdiag, inv3X3BlockDiagonal, inv2X2BlockDiagonal
|
||||
import numpy as np
|
||||
|
||||
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
from TensorMesh import TensorMesh
|
||||
from LogicallyOrthogonalMesh import LogicallyOrthogonalMesh
|
||||
import utils
|
||||
from exampleGrid import exampleLomGird
|
||||
|
||||
@@ -1,25 +0,0 @@
|
||||
import numpy as np
|
||||
from utils import ndgrid
|
||||
|
||||
|
||||
def exampleLomGird(nC, exType):
|
||||
assert type(nC) == list, "nC must be a list containing the number of nodes"
|
||||
assert len(nC) == 2 or len(nC) == 3, "nC must either two or three dimensions"
|
||||
exType = exType.lower()
|
||||
|
||||
possibleTypes = ['rect', 'rotate']
|
||||
assert exType in possibleTypes, "Not a possible example type."
|
||||
|
||||
if exType == 'rect':
|
||||
return ndgrid([np.cumsum(np.r_[0, np.ones(nx)/nx]) for nx in nC], vector=False)
|
||||
elif exType == 'rotate':
|
||||
if len(nC) == 2:
|
||||
X, Y = ndgrid([np.cumsum(np.r_[0, np.ones(nx)/nx]) for nx in nC], vector=False)
|
||||
amt = 0.5-np.sqrt((X - 0.5)**2 + (Y - 0.5)**2)
|
||||
amt[amt < 0] = 0
|
||||
return X + (-(Y - 0.5))*amt, Y + (+(X - 0.5))*amt
|
||||
elif len(nC) == 3:
|
||||
X, Y, Z = ndgrid([np.cumsum(np.r_[0, np.ones(nx)/nx]) for nx in nC], vector=False)
|
||||
amt = 0.5-np.sqrt((X - 0.5)**2 + (Y - 0.5)**2 + (Z - 0.5)**2)
|
||||
amt[amt < 0] = 0
|
||||
return X + (-(Y - 0.5))*amt, Y + (-(Z - 0.5))*amt, Z + (-(X - 0.5))*amt
|
||||
@@ -1,98 +0,0 @@
|
||||
from scipy import sparse as sp
|
||||
from utils import mkvc
|
||||
|
||||
|
||||
def sdiag(h):
|
||||
"""Sparse diagonal matrix"""
|
||||
return sp.spdiags(mkvc(h), 0, h.size, h.size, format="csr")
|
||||
|
||||
|
||||
def speye(n):
|
||||
"""Sparse identity"""
|
||||
return sp.identity(n, format="csr")
|
||||
|
||||
|
||||
def kron3(A, B, C):
|
||||
"""Three kron prods"""
|
||||
return sp.kron(sp.kron(A, B), C, format="csr")
|
||||
|
||||
|
||||
def spzeros(n1, n2):
|
||||
"""spzeros"""
|
||||
return sp.coo_matrix((n1, n2)).tocsr()
|
||||
|
||||
|
||||
def inv3X3BlockDiagonal(a11, a12, a13, a21, a22, a23, a31, a32, a33):
|
||||
""" B = inv3X3BlockDiagonal(a11, a12, a13, a21, a22, a23, a31, a32, a33)
|
||||
|
||||
inverts a stack of 3x3 matrices
|
||||
|
||||
Input:
|
||||
A - a11, a12, a13, a21, a22, a23, a31, a32, a33
|
||||
|
||||
Output:
|
||||
B - inverse
|
||||
"""
|
||||
|
||||
a11 = mkvc(a11)
|
||||
a12 = mkvc(a12)
|
||||
a13 = mkvc(a13)
|
||||
a21 = mkvc(a21)
|
||||
a22 = mkvc(a22)
|
||||
a23 = mkvc(a23)
|
||||
a31 = mkvc(a31)
|
||||
a32 = mkvc(a32)
|
||||
a33 = mkvc(a33)
|
||||
|
||||
detA = a31*a12*a23 - a31*a13*a22 - a21*a12*a33 + a21*a13*a32 + a11*a22*a33 - a11*a23*a32
|
||||
|
||||
b11 = +(a22*a33 - a23*a32)/detA
|
||||
b12 = -(a12*a33 - a13*a32)/detA
|
||||
b13 = +(a12*a23 - a13*a22)/detA
|
||||
|
||||
b21 = +(a31*a23 - a21*a33)/detA
|
||||
b22 = -(a31*a13 - a11*a33)/detA
|
||||
b23 = +(a21*a13 - a11*a23)/detA
|
||||
|
||||
b31 = -(a31*a22 - a21*a32)/detA
|
||||
b32 = +(a31*a12 - a11*a32)/detA
|
||||
b33 = -(a21*a12 - a11*a22)/detA
|
||||
|
||||
B = sp.vstack((sp.hstack((sdiag(b11), sdiag(b12), sdiag(b13))),
|
||||
sp.hstack((sdiag(b21), sdiag(b22), sdiag(b23))),
|
||||
sp.hstack((sdiag(b31), sdiag(b32), sdiag(b33)))))
|
||||
|
||||
return B
|
||||
|
||||
|
||||
def inv2X2BlockDiagonal(a11, a12, a21, a22):
|
||||
""" B = inv2X2BlockDiagonal(a11, a12, a21, a22)
|
||||
|
||||
Inverts a stack of 2x2 matrices by using the inversion formula
|
||||
|
||||
inv(A) = (1/det(A)) * cof(A)^T
|
||||
|
||||
Input:
|
||||
A - a11, a12, a13, a21, a22, a23, a31, a32, a33
|
||||
|
||||
Output:
|
||||
B - inverse
|
||||
"""
|
||||
|
||||
a11 = mkvc(a11)
|
||||
a12 = mkvc(a12)
|
||||
a21 = mkvc(a21)
|
||||
a22 = mkvc(a22)
|
||||
|
||||
# compute inverse of the determinant.
|
||||
detAinv = 1./(a11*a22 - a21*a12)
|
||||
|
||||
b11 = +detAinv*a22
|
||||
b12 = -detAinv*a12
|
||||
b21 = -detAinv*a21
|
||||
b22 = +detAinv*a11
|
||||
|
||||
B = sp.vstack((sp.hstack((sdiag(b11), sdiag(b12))),
|
||||
sp.hstack((sdiag(b21), sdiag(b22)))))
|
||||
|
||||
return B
|
||||
@@ -1,6 +1,6 @@
|
||||
import sys
|
||||
sys.path.append('../../')
|
||||
from SimPEG import TensorMesh, utils, LogicallyOrthogonalMesh, exampleLomGird
|
||||
from SimPEG import TensorMesh, utils, LogicallyOrthogonalMesh
|
||||
import numpy as np
|
||||
import unittest
|
||||
|
||||
@@ -100,10 +100,10 @@ class OrderTest(unittest.TestCase):
|
||||
else:
|
||||
raise Exception('Unexpected meshType')
|
||||
if self.meshDimension == 2:
|
||||
X, Y = exampleLomGird([nc, nc], kwrd)
|
||||
X, Y = utils.exampleLomGird([nc, nc], kwrd)
|
||||
self.M = LogicallyOrthogonalMesh([X, Y])
|
||||
if self.meshDimension == 3:
|
||||
X, Y, Z = exampleLomGird([nc, nc, nc], kwrd)
|
||||
X, Y, Z = utils.exampleLomGird([nc, nc, nc], kwrd)
|
||||
self.M = LogicallyOrthogonalMesh([X, Y, Z])
|
||||
return 1./nc
|
||||
|
||||
|
||||
@@ -1,7 +1,5 @@
|
||||
import numpy as np
|
||||
import unittest
|
||||
import sys
|
||||
sys.path.append('../')
|
||||
from OrderTest import OrderTest
|
||||
|
||||
MESHTYPES = ['uniformTensorMesh', 'uniformLOM', 'rotateLOM']
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
from utils import getSubArray, mkvc, ndgrid, ind2sub, sub2ind
|
||||
from sputils import spzeros, kron3, speye, sdiag
|
||||
from lomUtils import volTetra, faceInfo, inv2X2BlockDiagonal, inv3X3BlockDiagonal, indexCube, exampleLomGird
|
||||
@@ -1,101 +1,30 @@
|
||||
import numpy as np
|
||||
from scipy import sparse as sp
|
||||
from utils import mkvc, ndgrid, sub2ind
|
||||
from sputils import sdiag
|
||||
|
||||
|
||||
def mkvc(x, numDims=1):
|
||||
"""Creates a vector with the number of dimension specified
|
||||
def exampleLomGird(nC, exType):
|
||||
assert type(nC) == list, "nC must be a list containing the number of nodes"
|
||||
assert len(nC) == 2 or len(nC) == 3, "nC must either two or three dimensions"
|
||||
exType = exType.lower()
|
||||
|
||||
e.g.:
|
||||
possibleTypes = ['rect', 'rotate']
|
||||
assert exType in possibleTypes, "Not a possible example type."
|
||||
|
||||
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]
|
||||
if exType == 'rect':
|
||||
return ndgrid([np.cumsum(np.r_[0, np.ones(nx)/nx]) for nx in nC], vector=False)
|
||||
elif exType == 'rotate':
|
||||
if len(nC) == 2:
|
||||
X, Y = ndgrid([np.cumsum(np.r_[0, np.ones(nx)/nx]) for nx in nC], vector=False)
|
||||
amt = 0.5-np.sqrt((X - 0.5)**2 + (Y - 0.5)**2)
|
||||
amt[amt < 0] = 0
|
||||
return X + (-(Y - 0.5))*amt, Y + (+(X - 0.5))*amt
|
||||
elif len(nC) == 3:
|
||||
X, Y, Z = ndgrid([np.cumsum(np.r_[0, np.ones(nx)/nx]) for nx in nC], vector=False)
|
||||
amt = 0.5-np.sqrt((X - 0.5)**2 + (Y - 0.5)**2 + (Z - 0.5)**2)
|
||||
amt[amt < 0] = 0
|
||||
return X + (-(Y - 0.5))*amt, Y + (-(Z - 0.5))*amt, Z + (-(X - 0.5))*amt
|
||||
|
||||
|
||||
def volTetra(xyz, A, B, C, D):
|
||||
@@ -288,43 +217,77 @@ def faceInfo(xyz, A, B, C, D, average=True, normalizeNormals=True):
|
||||
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))
|
||||
def inv3X3BlockDiagonal(a11, a12, a13, a21, a22, a23, a31, a32, a33):
|
||||
""" B = inv3X3BlockDiagonal(a11, a12, a13, a21, a22, a23, a31, a32, a33)
|
||||
|
||||
sub = []
|
||||
inverts a stack of 3x3 matrices
|
||||
|
||||
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
|
||||
Input:
|
||||
A - a11, a12, a13, a21, a22, a23, a31, a32, a33
|
||||
|
||||
Output:
|
||||
B - inverse
|
||||
"""
|
||||
|
||||
a11 = mkvc(a11)
|
||||
a12 = mkvc(a12)
|
||||
a13 = mkvc(a13)
|
||||
a21 = mkvc(a21)
|
||||
a22 = mkvc(a22)
|
||||
a23 = mkvc(a23)
|
||||
a31 = mkvc(a31)
|
||||
a32 = mkvc(a32)
|
||||
a33 = mkvc(a33)
|
||||
|
||||
detA = a31*a12*a23 - a31*a13*a22 - a21*a12*a33 + a21*a13*a32 + a11*a22*a33 - a11*a23*a32
|
||||
|
||||
b11 = +(a22*a33 - a23*a32)/detA
|
||||
b12 = -(a12*a33 - a13*a32)/detA
|
||||
b13 = +(a12*a23 - a13*a22)/detA
|
||||
|
||||
b21 = +(a31*a23 - a21*a33)/detA
|
||||
b22 = -(a31*a13 - a11*a33)/detA
|
||||
b23 = +(a21*a13 - a11*a23)/detA
|
||||
|
||||
b31 = -(a31*a22 - a21*a32)/detA
|
||||
b32 = +(a31*a12 - a11*a32)/detA
|
||||
b33 = -(a21*a12 - a11*a22)/detA
|
||||
|
||||
B = sp.vstack((sp.hstack((sdiag(b11), sdiag(b12), sdiag(b13))),
|
||||
sp.hstack((sdiag(b21), sdiag(b22), sdiag(b23))),
|
||||
sp.hstack((sdiag(b31), sdiag(b32), sdiag(b33)))))
|
||||
|
||||
return B
|
||||
|
||||
|
||||
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)
|
||||
def inv2X2BlockDiagonal(a11, a12, a21, a22):
|
||||
""" B = inv2X2BlockDiagonal(a11, a12, a21, a22)
|
||||
|
||||
idx = np.dot((subs), (mult))
|
||||
return idx
|
||||
Inverts a stack of 2x2 matrices by using the inversion formula
|
||||
|
||||
inv(A) = (1/det(A)) * cof(A)^T
|
||||
|
||||
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"
|
||||
Input:
|
||||
A - a11, a12, a13, a21, a22, a23, a31, a32, a33
|
||||
|
||||
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.")
|
||||
Output:
|
||||
B - inverse
|
||||
"""
|
||||
|
||||
a11 = mkvc(a11)
|
||||
a12 = mkvc(a12)
|
||||
a21 = mkvc(a21)
|
||||
a22 = mkvc(a22)
|
||||
|
||||
# compute inverse of the determinant.
|
||||
detAinv = 1./(a11*a22 - a21*a12)
|
||||
|
||||
b11 = +detAinv*a22
|
||||
b12 = -detAinv*a12
|
||||
b21 = -detAinv*a21
|
||||
b22 = +detAinv*a11
|
||||
|
||||
B = sp.vstack((sp.hstack((sdiag(b11), sdiag(b12))),
|
||||
sp.hstack((sdiag(b21), sdiag(b22)))))
|
||||
|
||||
return B
|
||||
@@ -0,0 +1,22 @@
|
||||
from scipy import sparse as sp
|
||||
from utils import mkvc
|
||||
|
||||
|
||||
def sdiag(h):
|
||||
"""Sparse diagonal matrix"""
|
||||
return sp.spdiags(mkvc(h), 0, h.size, h.size, format="csr")
|
||||
|
||||
|
||||
def speye(n):
|
||||
"""Sparse identity"""
|
||||
return sp.identity(n, format="csr")
|
||||
|
||||
|
||||
def kron3(A, B, C):
|
||||
"""Three kron prods"""
|
||||
return sp.kron(sp.kron(A, B), C, format="csr")
|
||||
|
||||
|
||||
def spzeros(n1, n2):
|
||||
"""spzeros"""
|
||||
return sp.coo_matrix((n1, n2)).tocsr()
|
||||
@@ -0,0 +1,140 @@
|
||||
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 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.")
|
||||
Reference in New Issue
Block a user