def test_cg_with_tridiag(self):
size = 10
matrix = torch.randn(size, size, dtype=torch.float64)
matrix = matrix.matmul(matrix.transpose(-1, -2))
matrix.div_(matrix.norm())
matrix.add_(torch.eye(matrix.size(-1), dtype=torch.float64).mul_(1e-1))
rhs = torch.randn(size, 50, dtype=torch.float64)
solves, t_mats = linear_cg(matrix.matmul,
rhs=rhs,
n_tridiag=5,
max_tridiag_iter=10,
max_iter=size,
tolerance=0,
eps=1e-15)
# Check cg
matrix_chol = torch.linalg.cholesky(matrix)
actual = torch.cholesky_solve(rhs, matrix_chol)
self.assertTrue(torch.allclose(solves, actual, atol=1e-3, rtol=1e-4))
# Check tridiag
eigs = torch.linalg.eigvalsh(matrix)
for i in range(5):
approx_eigs = torch.linalg.eigvalsh(t_mats[i])
self.assertTrue(
torch.allclose(eigs, approx_eigs, atol=1e-3, rtol=1e-4))
def test_batch_cg_with_tridiag(self):
batch = 5
size = 10
matrix = torch.DoubleTensor(batch, size, size).normal_()
matrix = matrix.matmul(matrix.transpose(-1, -2))
matrix.div_(matrix.norm())
matrix.add_(torch.DoubleTensor(matrix.size(-1)).fill_(1e-1).diag())
rhs = torch.DoubleTensor(batch, size, 50).normal_()
solves, t_mats = linear_cg(
matrix.matmul,
rhs=rhs,
n_tridiag=8,
max_iter=size,
tolerance=0,
)
# Check cg
matrix_chol = torch.cat(
[matrix[i].potrf().unsqueeze(0) for i in range(5)])
actual = torch.cat([
torch.potrs(rhs[i], matrix_chol[i]).unsqueeze(0) for i in range(5)
])
self.assertTrue(approx_equal(solves, actual))
# Check tridiag
for i in range(5):
eigs = matrix[i].symeig()[0]
for j in range(8):
approx_eigs = t_mats[j, i].symeig()[0]
self.assertLess(
torch.mean(torch.abs((eigs - approx_eigs) / eigs)),
0.05,
)
def test_cg_with_tridiag(self):
size = 10
matrix = torch.DoubleTensor(size, size).normal_()
matrix = matrix.matmul(matrix.transpose(-1, -2))
matrix.div_(matrix.norm())
matrix.add_(torch.DoubleTensor(matrix.size(-1)).fill_(1e-1).diag())
rhs = torch.DoubleTensor(size, 50).normal_()
solves, t_mats = linear_cg(
matrix.matmul,
rhs=rhs,
n_tridiag=5,
max_iter=size,
tolerance=0,
)
# Check cg
matrix_chol = matrix.potrf()
actual = torch.potrs(rhs, matrix_chol)
self.assertTrue(approx_equal(solves, actual))
# Check tridiag
eigs = matrix.symeig()[0]
for i in range(5):
approx_eigs = t_mats[i].symeig()[0]
self.assertTrue(approx_equal(eigs, approx_eigs))
def test_batch_cg_with_tridiag(self):
batch = 5
size = 10
matrix = torch.randn(batch, size, size, dtype=torch.float64)
matrix = matrix.matmul(matrix.transpose(-1, -2))
matrix.div_(matrix.norm())
matrix.add_(torch.eye(matrix.size(-1), dtype=torch.float64).mul_(1e-1))
rhs = torch.randn(batch, size, 50, dtype=torch.float64)
solves, t_mats = linear_cg(matrix.matmul,
rhs=rhs,
n_tridiag=8,
max_iter=size,
max_tridiag_iter=10,
tolerance=0)
# Check cg
matrix_chol = batch_potrf(matrix)
actual = batch_potrs(rhs, matrix_chol)
self.assertTrue(approx_equal(solves, actual))
# Check tridiag
for i in range(5):
eigs = matrix[i].symeig()[0]
for j in range(8):
approx_eigs = t_mats[j, i].symeig()[0]
self.assertLess(
torch.mean(torch.abs((eigs - approx_eigs) / eigs)), 0.05)
def test_cg_with_tridiag(self):
size = 10
matrix = torch.randn(size, size, dtype=torch.float64)
matrix = matrix.matmul(matrix.transpose(-1, -2))
matrix.div_(matrix.norm())
matrix.add_(torch.eye(matrix.size(-1), dtype=torch.float64).mul_(1e-1))
rhs = torch.randn(size, 50, dtype=torch.float64)
solves, t_mats = linear_cg(matrix.matmul,
rhs=rhs,
n_tridiag=5,
max_tridiag_iter=10,
max_iter=size,
tolerance=0)
# Check cg
matrix_chol = matrix.potrf()
actual = torch.potrs(rhs, matrix_chol)
self.assertTrue(approx_equal(solves, actual))
# Check tridiag
eigs = matrix.symeig()[0]
for i in range(5):
approx_eigs = t_mats[i].symeig()[0]
self.assertTrue(approx_equal(eigs, approx_eigs))
def test_batch_cg_with_tridiag(self):
batch = 5
size = 10
matrix = torch.randn(batch, size, size, dtype=torch.float64)
matrix = matrix.matmul(matrix.transpose(-1, -2))
matrix.div_(matrix.norm())
matrix.add_(torch.eye(matrix.size(-1), dtype=torch.float64).mul_(1e-1))
rhs = torch.randn(batch, size, 10, dtype=torch.float64)
solves, t_mats = linear_cg(matrix.matmul,
rhs=rhs,
n_tridiag=8,
max_iter=size,
max_tridiag_iter=10,
tolerance=0,
eps=1e-30)
# Check cg
matrix_chol = torch.cholesky(matrix)
actual = torch.cholesky_solve(rhs, matrix_chol)
self.assertTrue(torch.allclose(solves, actual, atol=1e-3, rtol=1e-4))
# Check tridiag
for i in range(5):
eigs = matrix[i].symeig()[0]
for j in range(8):
approx_eigs = t_mats[j, i].symeig()[0]
self.assertTrue(
torch.allclose(eigs, approx_eigs, atol=1e-3, rtol=1e-4))
def test_cg(self):
size = 100
matrix = torch.DoubleTensor(size, size).normal_()
matrix = matrix.matmul(matrix.transpose(-1, -2))
matrix.div_(matrix.norm())
matrix.add_(torch.DoubleTensor(matrix.size(-1)).fill_(1e-1).diag())
rhs = torch.DoubleTensor(size, 50).normal_()
solves = linear_cg(matrix.matmul, rhs=rhs, max_iter=size)
# Check cg
matrix_chol = matrix.potrf()
actual = torch.potrs(rhs, matrix_chol)
self.assertTrue(approx_equal(solves, actual))
def test_cg(self):
size = 100
matrix = torch.randn(size, size, dtype=torch.float64)
matrix = matrix.matmul(matrix.transpose(-1, -2))
matrix.div_(matrix.norm())
matrix.add_(torch.eye(matrix.size(-1), dtype=torch.float64).mul_(1e-1))
rhs = torch.randn(size, 50, dtype=torch.float64)
solves = linear_cg(matrix.matmul, rhs=rhs, max_iter=size)
# Check cg
matrix_chol = matrix.cholesky()
actual = torch.cholesky_solve(rhs, matrix_chol)
self.assertTrue(torch.allclose(solves, actual, atol=1e-3, rtol=1e-4))
def test_cg(self):
size = 100
matrix = torch.randn(size, size, dtype=torch.float64)
matrix = matrix.matmul(matrix.transpose(-1, -2))
matrix.div_(matrix.norm())
matrix.add_(torch.eye(matrix.size(-1), dtype=torch.float64).mul_(1e-1))
rhs = torch.randn(size, 50, dtype=torch.float64)
solves = linear_cg(matrix.matmul, rhs=rhs, max_iter=size)
# Check cg
matrix_chol = matrix.cholesky(upper=True)
actual = torch.potrs(rhs, matrix_chol)
self.assertTrue(approx_equal(solves, actual))
def test_batch_cg():
batch = 5
size = 100
matrix = torch.DoubleTensor(batch, size, size).normal_()
matrix = matrix.matmul(matrix.transpose(-1, -2))
matrix.div_(matrix.norm())
matrix.add_(torch.DoubleTensor(matrix.size(-1)).fill_(1e-1).diag())
rhs = torch.DoubleTensor(batch, size, 50).normal_()
solves = linear_cg(matrix.matmul, rhs=rhs, max_iter=size)
# Check cg
matrix_chol = torch.cat([matrix[i].potrf().unsqueeze(0) for i in range(5)])
actual = torch.cat(
[torch.potrs(rhs[i], matrix_chol[i]).unsqueeze(0) for i in range(5)])
assert approx_equal(solves, actual)
