Skip to content

Add a classic solve_lu #304

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
inducer merged 3 commits into from
Sep 4, 2026
+357 −0
Merged

Add a classic solve_lu #304

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
83b805f
feat: add a classic solve_lu
alexfikl Sep 4, 2026
376e7c1
test: add test for solve_lu
alexfikl Sep 4, 2026
a66e943
chore: update baseline
alexfikl Sep 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
184 changes: 184 additions & 0 deletions .basedpyright/baseline.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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": {
Expand Down Expand Up @@ -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": {
Expand Down
97 changes: 97 additions & 0 deletions sumpy/symbolic.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down
76 changes: 76 additions & 0 deletions sumpy/test/test_misc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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():
Expand Down
Loading