Resolve materialize! ambiguity with LazyArrays 2.13 padded ldiv

src/FastAlmostBandedMatrices.jl

@@ -6,10 +6,11 @@ import ArrayInterface
 import ArrayLayouts
 import ArrayLayouts: LayoutMatrix, LayoutVector, Ldiv, Lmul, TriangularLayout
 import ConcreteStructs: @concrete
+import LazyArrays
 import LazyArrays: LazyArray, Mul
 import LinearAlgebra
-import LinearAlgebra: LowerTriangular, NoPivot, UnitLowerTriangular, UnitUpperTriangular,
-    UpperTriangular, diagind, lmul!, lu, qr, rank, triu!
+import LinearAlgebra: LowerTriangular, NoPivot, SingularException, UnitLowerTriangular,
+    UnitUpperTriangular, UpperTriangular, diagind, lmul!, lu, qr, rank, triu!
 import MatrixFactorizations
 
 # The BandedMatrices.jl surface that FastAlmostBandedMatrices reexports (see the second
@@ -813,6 +814,47 @@ function Base.materialize!(M::MatLdivVec{TriangularLayout{'U', 'U', AlmostBanded
     return x
 end
 
+# A padded right-hand side is zero past its data, so an upper-triangular solve only
+# needs the leading square block; the `Zeros` tail cannot be written into anyway.
+# Back-substitute on that prefix in-place (O(n) storage). Densifying the leading
+# block would be O(n²), and the existing almost-banded upper `materialize!` path
+# cannot be reused here: it goes through `__original_almostbandedrank` / a
+# `SubArray` of `AlmostBandedMatrix` and is already `@test_broken` for a dense fill.
+@views function _padded_upper_ldiv!(Tri, R, x)
+    A = triangulardata(R)
+    b = view(x, colsupport(x))
+    n = length(b)
+    iszero(n) && return x
+    unit = Tri === UnitUpperTriangular
+    Av = view(A, 1:n, 1:n)
+    rnk = almostbandedrank(Av)
+    _, u = almostbandwidths(Av)
+    for k in n:-1:1
+        s = b[k]
+        # Superdiagonals only: `k + u` can be < `k` when `u` is negative.
+        jmax_band = min(n, k + u)
+        for j in (k + 1):jmax_band
+            s -= Av[k, j] * b[j]
+        end
+        if k <= rnk
+            for j in max(jmax_band + 1, k + 1):n
+                s -= Av[k, j] * b[j]
+            end
+        end
+        if unit
+            b[k] = s
+        else
+            akk = Av[k, k]
+            iszero(akk) && throw(SingularException(k))
+            b[k] = s / akk
+        end
+    end
+    return x
+end
+
+Base.materialize!(M::MatLdivVec{TriangularLayout{'U', 'N', AlmostBandedLayout}, <:LazyArrays.PaddedColumns}) = _padded_upper_ldiv!(UpperTriangular, M.A, M.B)
+Base.materialize!(M::MatLdivVec{TriangularLayout{'U', 'U', AlmostBandedLayout}, <:LazyArrays.PaddedColumns}) = _padded_upper_ldiv!(UnitUpperTriangular, M.A, M.B)
+
 function Base.materialize!(M::MatLdivVec{TriangularLayout{'L', 'N', AlmostBandedLayout}})
     R, x = M.A, M.B
     A = triangulardata(R)

test/core_tests.jl

@@ -194,6 +194,43 @@ end
     @test_broken UnitUpperTriangular(Matrix(A)) \ b ≈ UnitUpperTriangular(A) \ b
     @test LowerTriangular(Matrix(A)) \ b ≈ LowerTriangular(A) \ b
     @test UnitLowerTriangular(Matrix(A)) \ b ≈ UnitLowerTriangular(A) \ b
+
+    LA = FastAlmostBandedMatrices.LazyArrays
+    # In-place `ldiv!` on `Vcat(v, Zeros)` uses these methods for all LazyArrays
+    # 2.x; out-of-place `\` only keeps a padded RHS from LazyArrays 2.13.
+    bpad = LA.Vcat(b[1:5], LA.Zeros(n - 5))
+    @test ldiv!(UpperTriangular(A), copy(bpad)) ≈
+        UpperTriangular(Matrix(A)) \ Vector(bpad)
+    @test ldiv!(UnitUpperTriangular(A), copy(bpad)) ≈
+        UnitUpperTriangular(Matrix(A)) \ Vector(bpad)
+    if Base.pkgversion(LA) >= v"2.13"
+        @test UpperTriangular(A) \ bpad ≈ UpperTriangular(Matrix(A)) \ Vector(bpad)
+        @test UnitUpperTriangular(A) \ bpad ≈
+            UnitUpperTriangular(Matrix(A)) \ Vector(bpad)
+    end
+
+    # Negative upper bandwidths: only strictly-upper columns may be subtracted.
+    for u in (-1, -2)
+        F = [
+            2.0 0.1 0.2 0.3
+            0.1 2.0 0.3 0.2
+            0.2 0.3 2.0 0.1
+            0.3 0.2 0.1 2.0
+        ]
+        Aneg = AlmostBandedMatrix(BandedMatrix(zeros(4, 4), (3, u)), F)
+        bneg = LA.Vcat([1.0, 2.0], LA.Zeros(2))
+        for Tri in (UpperTriangular, UnitUpperTriangular)
+            @test ldiv!(Tri(Aneg), copy(bneg)) ≈ Tri(Matrix(Aneg)) \ Vector(bneg)
+        end
+    end
+
+    Asing = AlmostBandedMatrix(BandedMatrix(zeros(4, 4), (1, 1)), zeros(1, 4))
+    Asing[band(0)] .= 2
+    Asing[2, 2] = 0
+    bsing = LA.Vcat([1.0, 2.0], LA.Zeros(2))
+    # Package path must throw SingularException; dense LAPACK may throw
+    # LAPACKException instead (Julia 1.10) so is not asserted here.
+    @test_throws SingularException(2) ldiv!(UpperTriangular(Asing), copy(bsing))
 end
 
 # https://github.com/SciML/FastAlmostBandedMatrices.jl/issues/19

test/qa/qa.jl

@@ -32,6 +32,8 @@ run_qa(
             ignore = (
                 :OneTo, :array_summary, :dims2string, :inds2string, :materialize!,
                 :QRPackedQ, :arguments, :fast_scalar_indexing, :qr_instance,
+                # needed to name LazyArrays' padded layout to resolve the ambiguity; not public in LazyArrays as of the tested version
+                :PaddedColumns,
             ),
         ),
     ),