From 83b805f0ece78be7871d22498a5e317000ac79a8 Mon Sep 17 00:00:00 2001 From: Alexandru Fikl Date: Fri, 4 Sep 2026 16:32:45 +0300 Subject: [PATCH 1/3] feat: add a classic solve_lu --- sumpy/symbolic.py | 97 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 97 insertions(+) diff --git a/sumpy/symbolic.py b/sumpy/symbolic.py index 3a0c552f..851c1863 100644 --- a/sumpy/symbolic.py +++ b/sumpy/symbolic.py @@ -48,6 +48,8 @@ if TYPE_CHECKING: + from collections.abc import Callable, Sequence + from pymbolic.typing import ArithmeticExpression, Expression from pytools import T from pytools.obj_array import ObjectArray1D @@ -258,6 +260,101 @@ def checked_cse(exprs, symbols=None): # }}} +# {{{ solve_lu + +def _solve_lu_forward_substitution( + L: Matrix, # ruff: ignore[invalid-argument-name] + b: Matrix, + postprocess: Callable[[Basic], Basic], + ) -> Matrix: + """Solve the lower triangular system :math:`L x = b`. + + :arg postprocess: a callable that is applied on each expression at the end + of division calls. + """ + n = len(b) + x = b + + for i in range(n): + for j in range(i): + x[i] -= L[i, j] * x[j] + x[i] = postprocess(x[i] / L[i, i]) + + return x + + +def _solve_lu_backward_substitution( + U: Matrix, # ruff: ignore[invalid-argument-name] + b: Matrix, + postprocess: Callable[[Basic], Basic], + ) -> Matrix: + """Solve the upper triangular system :math:`U x = b`. + + :arg postprocess: a callable that is applied on each expression at the end + of division calls. + """ + n = len(b) + x = b + + for i in range(n - 1, -1, -1): + for j in range(n - 1, i, -1): + x[i] -= U[i, j] * x[j] + x[i] = postprocess(x[i] / U[i, i]) + + return x + + +def solve_lu( + L: Matrix, # ruff: ignore[invalid-argument-name] + U: Matrix, # ruff: ignore[invalid-argument-name] + permutation: Sequence[tuple[int, int]], + b: Matrix, *, + postprocess: Callable[[Basic], Basic] | None = None, + ) -> Matrix: + """Solve the system :math:`L U x = P b`. + + This function uses standard forward and backward substitution to solve the + system. The matrix *L* is assumed to be unit lower triangular, *U* is + assumed to be upper triangular and *permutation* is a sequence of row swaps + for the LU decomposition. + + :arg postprocess: a callable that is applied on each expression at the end + of division calls in both forward and backward substitution. + """ + + if postprocess is None: + def default_postprocess(x: Basic) -> Basic: + return x + + postprocess = default_postprocess + + if L.shape[0] != U.shape[1]: + raise ValueError( + f"system matrix is not square: L is {L.shape} and U is {U.shape}" + ) + + if L.shape[1] != U.shape[0]: + raise ValueError( + f"incorrect LU decomposition shapes: L is {L.shape} and U is {U.shape}" + ) + + if b.shape != (U.shape[1], 1): + raise ValueError(f"'b' is not a column vector matching U: {b.shape}") + + b = Matrix(b) + + # Permute first + for p, q in permutation: + b[p], b[q] = b[q], b[p] + + y = _solve_lu_forward_substitution(L, b, postprocess=postprocess) + x = _solve_lu_backward_substitution(U, y, postprocess=postprocess) + + return x + +# }}} + + # {{{ pymbolic expressions def sym_real_norm_2(x: Matrix) -> Expr: From 376e7c1af08bd06991792b24182af0cc6bd6ac7c Mon Sep 17 00:00:00 2001 From: Alexandru Fikl Date: Fri, 4 Sep 2026 16:45:14 +0300 Subject: [PATCH 2/3] test: add test for solve_lu --- sumpy/test/test_misc.py | 76 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 76 insertions(+) diff --git a/sumpy/test/test_misc.py b/sumpy/test/test_misc.py index 21fbc629..703a783f 100644 --- a/sumpy/test/test_misc.py +++ b/sumpy/test/test_misc.py @@ -433,6 +433,82 @@ def test_cse_matvec(): # }}} +# {{{ test_solve_lu + +def _make_lu_system( + n: int, *, + seed: int, + force_pivot: bool = True, + symbolic_rhs: bool = False, + ) -> tuple[sym.Matrix, sym.Matrix, list[tuple[int, int]], sym.Matrix, sym.Matrix]: + rng = np.random.default_rng(seed) + + while True: + a_mat = sym.Matrix(n, n, [sym.Integer(rng.integers(-5, 6)) for _ in range(n*n)]) + if force_pivot: + a_mat[0, 0] = 0 + + if a_mat.det() != 0: + break + + if symbolic_rhs: + s = sym.Symbol("s") + x_true = sym.Matrix([ + s if i == 0 else (s**2 if i == n - 1 else sym.Integer(i + 1)) + for i in range(n) + ]) + else: + x_true = sym.Matrix([sym.Integer(1)] * n) + + l_mat, u_mat, permutation = a_mat.LUdecomposition() + return l_mat, u_mat, permutation, a_mat * x_true, x_true + + +@pytest.mark.parametrize("n", [3, 4]) +@pytest.mark.parametrize("seed", [0, 1, 2]) +@pytest.mark.parametrize("force_pivot", [False, True]) +def test_solve_lu(n: int, seed: int, force_pivot: bool) -> None: + l_mat, u_mat, permutation, rhs, x_true = _make_lu_system( + n, seed=seed, force_pivot=force_pivot) + + rhs_orig = sym.Matrix(rhs) + got = sym.solve_lu(l_mat, u_mat, permutation, rhs) + + assert all((got[i] - x_true[i]).expand() == 0 for i in range(n)) + assert rhs == rhs_orig + + # postprocess is applied after each division and may change the form of + # the solution without changing its value + got = sym.solve_lu(l_mat, u_mat, permutation, rhs, + postprocess=lambda e: e.expand()) + assert all((got[i] - x_true[i]).expand() == 0 for i in range(n)) + + +@pytest.mark.parametrize("n", [3, 4]) +def test_solve_lu_symbolic_rhs(n: int) -> None: + l_mat, u_mat, permutation, rhs, x_true = _make_lu_system( + n, seed=42, symbolic_rhs=True) + + got = sym.solve_lu(l_mat, u_mat, permutation, rhs, + postprocess=lambda e: e.expand()) + assert all((got[i] - x_true[i]).expand() == 0 for i in range(n)) + + +def test_solve_lu_raises_on_bad_shapes() -> None: + l_mat, u_mat, permutation, _, _ = _make_lu_system(3, seed=0) + + # wrong-sized b + with pytest.raises(ValueError): + sym.solve_lu(l_mat, u_mat, permutation, sym.Matrix([1, 2, 3, 4])) + + # non-square L + bad_l = sym.Matrix([[1, 0], [2, 1], [0, 0]]) + with pytest.raises(ValueError): + sym.solve_lu(bad_l, u_mat, permutation, sym.Matrix([[1], [2], [3]])) + +# }}} + + # {{{ test_diff_op_stokes def test_diff_op_stokes(): From a66e9431fd398896486b2e5e148d435dfb4ece25 Mon Sep 17 00:00:00 2001 From: Alexandru Fikl Date: Fri, 4 Sep 2026 16:52:50 +0300 Subject: [PATCH 3/3] chore: update baseline --- .basedpyright/baseline.json | 184 ++++++++++++++++++++++++++++++++++++ 1 file changed, 184 insertions(+) diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index a159110e..f6e48746 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -13355,6 +13355,54 @@ "lineCount": 1 } }, + { + "code": "reportOperatorIssue", + "range": { + "startColumn": 20, + "endColumn": 34, + "lineCount": 1 + } + }, + { + "code": "reportOperatorIssue", + "range": { + "startColumn": 27, + "endColumn": 41, + "lineCount": 1 + } + }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 27, + "endColumn": 41, + "lineCount": 1 + } + }, + { + "code": "reportOperatorIssue", + "range": { + "startColumn": 20, + "endColumn": 34, + "lineCount": 1 + } + }, + { + "code": "reportOperatorIssue", + "range": { + "startColumn": 27, + "endColumn": 41, + "lineCount": 1 + } + }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 27, + "endColumn": 41, + "lineCount": 1 + } + }, { "code": "reportUnknownArgumentType", "range": { @@ -20805,6 +20853,142 @@ "lineCount": 1 } }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 11, + "endColumn": 20, + "lineCount": 1 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 32, + "endColumn": 53, + "lineCount": 1 + } + }, + { + "code": "reportOperatorIssue", + "range": { + "startColumn": 15, + "endColumn": 35, + "lineCount": 1 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 15, + "endColumn": 42, + "lineCount": 1 + } + }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 15, + "endColumn": 67, + "lineCount": 1 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 45, + "endColumn": 53, + "lineCount": 1 + } + }, + { + "code": "reportUnknownLambdaType", + "range": { + "startColumn": 45, + "endColumn": 55, + "lineCount": 1 + } + }, + { + "code": "reportAttributeAccessIssue", + "range": { + "startColumn": 47, + "endColumn": 53, + "lineCount": 1 + } + }, + { + "code": "reportOperatorIssue", + "range": { + "startColumn": 15, + "endColumn": 35, + "lineCount": 1 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 15, + "endColumn": 42, + "lineCount": 1 + } + }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 15, + "endColumn": 67, + "lineCount": 1 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 45, + "endColumn": 53, + "lineCount": 1 + } + }, + { + "code": "reportUnknownLambdaType", + "range": { + "startColumn": 45, + "endColumn": 55, + "lineCount": 1 + } + }, + { + "code": "reportAttributeAccessIssue", + "range": { + "startColumn": 47, + "endColumn": 53, + "lineCount": 1 + } + }, + { + "code": "reportOperatorIssue", + "range": { + "startColumn": 15, + "endColumn": 35, + "lineCount": 1 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 15, + "endColumn": 42, + "lineCount": 1 + } + }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 15, + "endColumn": 67, + "lineCount": 1 + } + }, { "code": "reportAny", "range": {