From ead0549e69e8d13bf40d1cb12d8fbdd2ac6f4f9c Mon Sep 17 00:00:00 2001 From: Qian Bao Date: Thu, 27 Apr 2023 15:44:23 +0800 Subject: [PATCH 1/4] Rename taichi_cg into matrixfree_cg. --- python/taichi/linalg/__init__.py | 2 +- python/taichi/linalg/{taichi_cg.py => matrixfree_cg.py} | 2 +- tests/python/{test_taichi_cg.py => test_matrixfree_cg.py} | 6 +++--- 3 files changed, 5 insertions(+), 5 deletions(-) rename python/taichi/linalg/{taichi_cg.py => matrixfree_cg.py} (97%) rename tests/python/{test_taichi_cg.py => test_matrixfree_cg.py} (91%) diff --git a/python/taichi/linalg/__init__.py b/python/taichi/linalg/__init__.py index 2bb594ffe273f..3c49063c11a2a 100644 --- a/python/taichi/linalg/__init__.py +++ b/python/taichi/linalg/__init__.py @@ -3,4 +3,4 @@ from taichi.linalg.cg import CG from taichi.linalg.sparse_matrix import * from taichi.linalg.sparse_solver import SparseSolver -from taichi.linalg.taichi_cg import * +from taichi.linalg.matrixfree_cg import * diff --git a/python/taichi/linalg/taichi_cg.py b/python/taichi/linalg/matrixfree_cg.py similarity index 97% rename from python/taichi/linalg/taichi_cg.py rename to python/taichi/linalg/matrixfree_cg.py index bb39c0af734f0..d997dc2e61bd3 100644 --- a/python/taichi/linalg/taichi_cg.py +++ b/python/taichi/linalg/matrixfree_cg.py @@ -16,7 +16,7 @@ def matvec(self, x, Ax): self._matvec(x, Ax) -def taichi_cg_solver(A, b, x, tol=1e-6, maxiter=5000, quiet=True): +def MatrixFreeCG(A, b, x, tol=1e-6, maxiter=5000, quiet=True): if b.dtype != x.dtype: raise TaichiTypeError(f"Dtype mismatch b.dtype({b.dtype}) != x.dtype({x.dtype}).") if str(b.dtype) == "f32": diff --git a/tests/python/test_taichi_cg.py b/tests/python/test_matrixfree_cg.py similarity index 91% rename from tests/python/test_taichi_cg.py rename to tests/python/test_matrixfree_cg.py index 122669931aa62..733a0ea24c0d2 100644 --- a/tests/python/test_taichi_cg.py +++ b/tests/python/test_matrixfree_cg.py @@ -1,7 +1,7 @@ import math import pytest -from taichi.linalg import LinearOperator, taichi_cg_solver +from taichi.linalg import LinearOperator, MatrixFreeCG import taichi as ti from tests import test_utils @@ -11,7 +11,7 @@ @pytest.mark.parametrize("ti_dtype", [ti.f32, ti.f64]) @test_utils.test(arch=[ti.cpu, ti.cuda, ti.vulkan], exclude=[vk_on_mac]) -def test_taichi_cg(ti_dtype): +def test_matrixfree_cg(ti_dtype): GRID = 32 Ax = ti.field(dtype=ti_dtype, shape=(GRID, GRID)) x = ti.field(dtype=ti_dtype, shape=(GRID, GRID)) @@ -47,7 +47,7 @@ def check_solution(sol: ti.template(), ans: ti.template(), tol: ti_dtype) -> boo A = LinearOperator(compute_Ax) init() - taichi_cg_solver(A, b, x, maxiter=10 * GRID * GRID, tol=1e-18, quiet=True) + MatrixFreeCG(A, b, x, maxiter=10 * GRID * GRID, tol=1e-18, quiet=True) compute_Ax(x, Ax) # `tol` can't be < 1e-6 for ti.f32 because of accumulating round-off error; # see https://en.wikipedia.org/wiki/Conjugate_gradient_method#cite_note-6 From 3371aaef96488b8f643a167a54e6a4fed67cd1ed Mon Sep 17 00:00:00 2001 From: Qian Bao Date: Thu, 27 Apr 2023 15:48:28 +0800 Subject: [PATCH 2/4] Rename cg into sparse_cg. --- python/taichi/linalg/__init__.py | 2 +- python/taichi/linalg/{cg.py => sparse_cg.py} | 2 +- tests/python/{test_cg.py => test_sparse_cg.py} | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) rename python/taichi/linalg/{cg.py => sparse_cg.py} (99%) rename tests/python/{test_cg.py => test_sparse_cg.py} (93%) diff --git a/python/taichi/linalg/__init__.py b/python/taichi/linalg/__init__.py index 3c49063c11a2a..c4e28878b1b59 100644 --- a/python/taichi/linalg/__init__.py +++ b/python/taichi/linalg/__init__.py @@ -1,6 +1,6 @@ """Taichi support module for sparse matrix operations. """ -from taichi.linalg.cg import CG +from taichi.linalg.sparse_cg import SparseCG from taichi.linalg.sparse_matrix import * from taichi.linalg.sparse_solver import SparseSolver from taichi.linalg.matrixfree_cg import * diff --git a/python/taichi/linalg/cg.py b/python/taichi/linalg/sparse_cg.py similarity index 99% rename from python/taichi/linalg/cg.py rename to python/taichi/linalg/sparse_cg.py index 21973641e4db3..fc0611fbff660 100644 --- a/python/taichi/linalg/cg.py +++ b/python/taichi/linalg/sparse_cg.py @@ -6,7 +6,7 @@ from taichi.types import f32, f64 -class CG: +class SparseCG: def __init__(self, A, b, x0=None, max_iter=50, atol=1e-6): self.dtype = A.dtype self.ti_arch = get_runtime().prog.config().arch diff --git a/tests/python/test_cg.py b/tests/python/test_sparse_cg.py similarity index 93% rename from tests/python/test_cg.py rename to tests/python/test_sparse_cg.py index bd0eb38445f20..ad6afa4c071f4 100644 --- a/tests/python/test_cg.py +++ b/tests/python/test_sparse_cg.py @@ -28,7 +28,7 @@ def fill( fill(Abuilder, A_psd, b) A = Abuilder.build(dtype=ti_dtype) - cg = ti.linalg.CG(A, b, x0, max_iter=50, atol=1e-6) + cg = ti.linalg.SparseCG(A, b, x0, max_iter=50, atol=1e-6) x, exit_code = cg.solve() res = np.linalg.solve(A_psd, b.to_numpy()) assert exit_code == True @@ -59,7 +59,7 @@ def fill( fill(Abuilder, A_psd, b) A = Abuilder.build(dtype=ti_dtype) - cg = ti.linalg.CG(A, b, x0, max_iter=50, atol=1e-6) + cg = ti.linalg.SparseCG(A, b, x0, max_iter=50, atol=1e-6) x, exit_code = cg.solve() res = np.linalg.solve(A_psd, b.to_numpy()) assert exit_code == True From 9ffb94d2ca9149bbd69b8919d14169a942b4702f Mon Sep 17 00:00:00 2001 From: houkensjtu Date: Thu, 4 May 2023 19:42:43 +0800 Subject: [PATCH 3/4] Add docstring to SparseCG and MatrixFreeCG. --- python/taichi/linalg/matrixfree_cg.py | 14 ++++++++++++++ python/taichi/linalg/sparse_cg.py | 12 ++++++++++++ 2 files changed, 26 insertions(+) diff --git a/python/taichi/linalg/matrixfree_cg.py b/python/taichi/linalg/matrixfree_cg.py index d997dc2e61bd3..6af5ef177c387 100644 --- a/python/taichi/linalg/matrixfree_cg.py +++ b/python/taichi/linalg/matrixfree_cg.py @@ -17,6 +17,20 @@ def matvec(self, x, Ax): def MatrixFreeCG(A, b, x, tol=1e-6, maxiter=5000, quiet=True): + """Matrix-free conjugate-gradient solver. + + Use conjugate-gradient method to solve the linear system Ax = b, where A is implicitly + represented as a LinearOperator. + + Args: + A (LinearOperator): The coefficient matrix A of the linear system. + b (Field): The right-hand side of the linear system. + x (Field): The initial guess for the solution. + maxiter (int): Maximum number of iterations. + atol: Tolerance(absolute) for convergence. + quiet (bool): Switch to turn on/off iteration log. + """ + if b.dtype != x.dtype: raise TaichiTypeError(f"Dtype mismatch b.dtype({b.dtype}) != x.dtype({x.dtype}).") if str(b.dtype) == "f32": diff --git a/python/taichi/linalg/sparse_cg.py b/python/taichi/linalg/sparse_cg.py index fc0611fbff660..5ebdf663f5ca9 100644 --- a/python/taichi/linalg/sparse_cg.py +++ b/python/taichi/linalg/sparse_cg.py @@ -7,6 +7,18 @@ class SparseCG: + """Conjugate-gradient solver built for SparseMatrix. + + Use conjugate-gradient method to solve the linear system Ax = b, where A is SparseMatrix. + + Args: + A (SparseMatrix): The coefficient matrix A of the linear system. + b (numpy ndarray, taichi Ndarray): The right-hand side of the linear system. + x0 (numpy ndarray, taichi Ndarray): The initial guess for the solution. + max_iter (int): Maximum number of iterations. + atol: Tolerance(absolute) for convergence. + """ + def __init__(self, A, b, x0=None, max_iter=50, atol=1e-6): self.dtype = A.dtype self.ti_arch = get_runtime().prog.config().arch From 130c5d1c1e68e776d5e6540be4ff6f5f38a152a5 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 4 May 2023 11:44:06 +0000 Subject: [PATCH 4/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- python/taichi/linalg/matrixfree_cg.py | 4 ++-- python/taichi/linalg/sparse_cg.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/python/taichi/linalg/matrixfree_cg.py b/python/taichi/linalg/matrixfree_cg.py index 6af5ef177c387..0d333f4bbcea0 100644 --- a/python/taichi/linalg/matrixfree_cg.py +++ b/python/taichi/linalg/matrixfree_cg.py @@ -18,7 +18,7 @@ def matvec(self, x, Ax): def MatrixFreeCG(A, b, x, tol=1e-6, maxiter=5000, quiet=True): """Matrix-free conjugate-gradient solver. - + Use conjugate-gradient method to solve the linear system Ax = b, where A is implicitly represented as a LinearOperator. @@ -30,7 +30,7 @@ def MatrixFreeCG(A, b, x, tol=1e-6, maxiter=5000, quiet=True): atol: Tolerance(absolute) for convergence. quiet (bool): Switch to turn on/off iteration log. """ - + if b.dtype != x.dtype: raise TaichiTypeError(f"Dtype mismatch b.dtype({b.dtype}) != x.dtype({x.dtype}).") if str(b.dtype) == "f32": diff --git a/python/taichi/linalg/sparse_cg.py b/python/taichi/linalg/sparse_cg.py index 5ebdf663f5ca9..e0aba8112c690 100644 --- a/python/taichi/linalg/sparse_cg.py +++ b/python/taichi/linalg/sparse_cg.py @@ -8,7 +8,7 @@ class SparseCG: """Conjugate-gradient solver built for SparseMatrix. - + Use conjugate-gradient method to solve the linear system Ax = b, where A is SparseMatrix. Args: @@ -18,7 +18,7 @@ class SparseCG: max_iter (int): Maximum number of iterations. atol: Tolerance(absolute) for convergence. """ - + def __init__(self, A, b, x0=None, max_iter=50, atol=1e-6): self.dtype = A.dtype self.ti_arch = get_runtime().prog.config().arch