def spsolve(self, vec):
"""For a matrix A, solves the equation "Ax = vec" for x.
"""
if len(vec) == 1:
return vec / self.diag
(_, _, _, x, _) = lapack.dgtsv(self.down, self.diag, self.up, vec)
return x
def solve_tridiag(a, b, c, d):
"""
Solves a tridiagonal matrix system with diagonals a, b, c and RHS vector d.
"""
assert a.shape == b.shape and a.shape == c.shape and a.shape == d.shape
a[..., 0] = c[..., -1] = 0 # remove couplings between slices
return lapack.dgtsv(a.flatten()[1:], b.flatten(),
c.flatten()[:-1], d.flatten())[3].reshape(a.shape)
def cyclic(a, b, c, alpha, beta, r, n, bb, u, x, z):
start = time.time()
gamma = -b[0]
bb[:] = b[:]
end = time.time()
print("1")
print(end-start)
start = time.time()
bb[0] = b[0]-gamma
bb[-1] = b[-1]-alpha*beta/gamma
end = time.time()
print("2")
print(end-start)
start = time.time()
x = sll.dgtsv(a,bb,c,r)[3]
end = time.time()
print("3")
print(end-start)
start = time.time()
u[0] = gamma
u[-1] = alpha
end = time.time()
print("4")
print(end-start)
start = time.time()
z = sll.dgtsv(a,bb,c,u)[3]
end = time.time()
print("5")
print(end-start)
start = time.time()
fact = (x[0]+beta*x[-1]/gamma)/(1.0+z[0]+beta*z[-1]/gamma)
x = x - fact*z
end = time.time()
print("6")
print(end-start)
return x
def compactFilter_Fast(self, f):
if self.typeBC == "PERIODIC":
return cyclic2(self.aF, self.bF, self.bbF, self.cF,
self.cyclicCornerF, self.cyclicCornerF,
spspar.csr_matrix.dot(self.RHF, f), self.N,
self.work1, self.work2, self.work3)
elif self.typeBC == "DIRICHLET":
return sll.dgtsv(self.aF, self.bF, self.cF,
spspar.csr_matrix.dot(self.RHF, f))[3]
def compact2ndDeriv_Fast(self, f):
if self.typeBC == "PERIODIC":
return cyclic2(self.a2, self.b2, self.bb2, self.c2,
self.cyclicCorner2, self.cyclicCorner2,
spspar.csr_matrix.dot(self.RH2D, f), self.N,
self.work1, self.work2, self.work3)
elif self.typeBC == "DIRICHLET":
return sll.dgtsv(self.a2, self.b2, self.c2,
spspar.csr_matrix.dot(self.RH2D, f))[3]
def cyclic2(a, b, bb, c, alpha, beta, r, n, u, x, z):
gamma = -b[0]
bb[0] = b[0] - gamma
bb[-1] = b[-1] - alpha * beta / gamma
x = sll.dgtsv(a, bb, c, r)[3]
u[:] = 0.0
u[0] = gamma
u[-1] = alpha
z = sll.dgtsv(a, bb, c, u)[3]
fact = (x[0] + beta * x[-1] / gamma) / (1.0 + z[0] + beta * z[-1] / gamma)
x = x - fact * z
return x
def solve_tridiag(pyom, a, b, c, d):
"""
Solves a tridiagonal matrix system with diagonals a, b, c and RHS vector d.
Uses LAPACK when running with NumPy, and otherwise the Thomas algorithm iterating over the
last axis of the input arrays.
"""
assert a.shape == b.shape and a.shape == c.shape and a.shape == d.shape
try:
return np.linalg.solve_tridiagonal(a,b,c,d)
except AttributeError:
return lapack.dgtsv(a.flatten()[1:],b.flatten(),c.flatten()[:-1],d.flatten())[3].reshape(a.shape)
def solve_tridiag(vs, a, b, c, d):
"""
Solves a tridiagonal matrix system with diagonals a, b, c and RHS vector d.
Uses LAPACK when running with NumPy, and otherwise the Thomas algorithm iterating over the
last axis of the input arrays.
"""
assert a.shape == b.shape and a.shape == c.shape and a.shape == d.shape
if rs.backend == 'bohrium' and rst.vector_engine in ('opencl', 'openmp'):
return np.linalg.solve_tridiagonal(a, b, c, d)
# fall back to scipy
from scipy.linalg import lapack
a[..., 0] = c[..., -1] = 0 # remove couplings between slices
return lapack.dgtsv(a.flatten()[1:], b.flatten(),
c.flatten()[:-1], d.flatten())[3].reshape(a.shape)
def solve_tridiag(vs, a, b, c, d):
"""
Solves a tridiagonal matrix system with diagonals a, b, c and RHS vector d.
Uses LAPACK when running with NumPy, and otherwise the Thomas algorithm iterating over the
last axis of the input arrays.
"""
assert a.shape == b.shape and a.shape == c.shape and a.shape == d.shape
if vs.backend_name == "numpy":
a[..., 0] = c[..., -1] = 0 # remove couplings between slices
return lapack.dgtsv(a.flatten()[1:], b.flatten(),
c.flatten()[:-1], d.flatten())[3].reshape(a.shape)
if vs.vector_engine == "opencl":
return tdma_opencl.tdma(a, b, c, d)
return np.linalg.solve_tridiagonal(a, b, c, d)
N = 100
a = np.ones(N-1,dtype=np.double)
b = -2*np.ones(N,dtype=np.double)
c = np.ones(N-1,dtype=np.double)
d = np.ones(N,dtype=np.double)
start = time.time()
for i in range(10000):
x = sll.dgtsv(a,b,c,d)
end = time.time()
print(end-start)
lmat = spspar.lil_matrix((N,N))
lmat.setdiag(np.ones(N-1),-1)
lmat.setdiag(-2*np.ones(N),0)
lmat.setdiag(np.ones(N-1),1)
lmat = lmat.tocsr()
start = time.time()
for i in range(10000):
xsparse = spsparlin.spsolve(lmat,d)
end = time.time()
print(end-start)
