From f1990e2e3d8d31087b23288f640ba16909f7ec7b Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Fri, 18 Oct 2024 02:49:19 +0530 Subject: [PATCH] Fix triu/tril for partly initialized matrices (#55312) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This fixes ```julia julia> using LinearAlgebra, StaticArrays julia> M = Matrix{BigInt}(undef, 2, 2); M[1,1] = M[2,2] = M[1,2] = 3; julia> S = SizedMatrix{2,2}(M) 2×2 SizedMatrix{2, 2, BigInt, 2, Matrix{BigInt}} with indices SOneTo(2)×SOneTo(2): 3 3 #undef 3 julia> triu(S) ERROR: UndefRefError: access to undefined reference Stacktrace: [1] getindex @ ./essentials.jl:907 [inlined] [2] getindex @ ~/.julia/packages/StaticArrays/MSJcA/src/SizedArray.jl:92 [inlined] [3] copyto_unaliased! @ ./abstractarray.jl:1086 [inlined] [4] copyto!(dest::SizedMatrix{2, 2, BigInt, 2, Matrix{BigInt}}, src::SizedMatrix{2, 2, BigInt, 2, Matrix{BigInt}}) @ Base ./abstractarray.jl:1066 [5] copymutable @ ./abstractarray.jl:1200 [inlined] [6] triu(M::SizedMatrix{2, 2, BigInt, 2, Matrix{BigInt}}) @ LinearAlgebra ~/.julia/juliaup/julia-nightly/share/julia/stdlib/v1.12/LinearAlgebra/src/generic.jl:413 [7] top-level scope @ REPL[11]:1 ``` After this PR: ```julia julia> triu(S) 2×2 SizedMatrix{2, 2, BigInt, 2, Matrix{BigInt}} with indices SOneTo(2)×SOneTo(2): 3 3 0 3 ``` Only the indices that need to be copied are accessed, and the others are written to without being read. --- stdlib/LinearAlgebra/src/generic.jl | 80 ++++++++++------------------ stdlib/LinearAlgebra/test/generic.jl | 55 +++++++++++++++++++ 2 files changed, 83 insertions(+), 52 deletions(-) diff --git a/stdlib/LinearAlgebra/src/generic.jl b/stdlib/LinearAlgebra/src/generic.jl index e5f23b4981616..6c65c49add74b 100644 --- a/stdlib/LinearAlgebra/src/generic.jl +++ b/stdlib/LinearAlgebra/src/generic.jl @@ -389,55 +389,7 @@ function cross(a::AbstractVector, b::AbstractVector) end """ - triu(M) - -Upper triangle of a matrix. - -# Examples -```jldoctest -julia> a = fill(1.0, (4,4)) -4×4 Matrix{Float64}: - 1.0 1.0 1.0 1.0 - 1.0 1.0 1.0 1.0 - 1.0 1.0 1.0 1.0 - 1.0 1.0 1.0 1.0 - -julia> triu(a) -4×4 Matrix{Float64}: - 1.0 1.0 1.0 1.0 - 0.0 1.0 1.0 1.0 - 0.0 0.0 1.0 1.0 - 0.0 0.0 0.0 1.0 -``` -""" -triu(M::AbstractMatrix) = triu!(copymutable(M)) - -""" - tril(M) - -Lower triangle of a matrix. - -# Examples -```jldoctest -julia> a = fill(1.0, (4,4)) -4×4 Matrix{Float64}: - 1.0 1.0 1.0 1.0 - 1.0 1.0 1.0 1.0 - 1.0 1.0 1.0 1.0 - 1.0 1.0 1.0 1.0 - -julia> tril(a) -4×4 Matrix{Float64}: - 1.0 0.0 0.0 0.0 - 1.0 1.0 0.0 0.0 - 1.0 1.0 1.0 0.0 - 1.0 1.0 1.0 1.0 -``` -""" -tril(M::AbstractMatrix) = tril!(copymutable(M)) - -""" - triu(M, k::Integer) + triu(M, k::Integer = 0) Return the upper triangle of `M` starting from the `k`th superdiagonal. @@ -465,10 +417,22 @@ julia> triu(a,-3) 1.0 1.0 1.0 1.0 ``` """ -triu(M::AbstractMatrix,k::Integer) = triu!(copymutable(M),k) +function triu(M::AbstractMatrix, k::Integer = 0) + d = similar(M) + A = triu!(d,k) + if iszero(k) + copytrito!(A, M, 'U') + else + for col in axes(A,2) + rows = firstindex(A,1):min(col-k, lastindex(A,1)) + A[rows, col] = @view M[rows, col] + end + end + return A +end """ - tril(M, k::Integer) + tril(M, k::Integer = 0) Return the lower triangle of `M` starting from the `k`th superdiagonal. @@ -496,7 +460,19 @@ julia> tril(a,-3) 1.0 0.0 0.0 0.0 ``` """ -tril(M::AbstractMatrix,k::Integer) = tril!(copymutable(M),k) +function tril(M::AbstractMatrix,k::Integer=0) + d = similar(M) + A = tril!(d,k) + if iszero(k) + copytrito!(A, M, 'L') + else + for col in axes(A,2) + rows = max(firstindex(A,1),col-k):lastindex(A,1) + A[rows, col] = @view M[rows, col] + end + end + return A +end """ triu!(M) diff --git a/stdlib/LinearAlgebra/test/generic.jl b/stdlib/LinearAlgebra/test/generic.jl index e0a1704913f78..2bf9c75141700 100644 --- a/stdlib/LinearAlgebra/test/generic.jl +++ b/stdlib/LinearAlgebra/test/generic.jl @@ -18,6 +18,9 @@ using .Main.DualNumbers isdefined(Main, :FillArrays) || @eval Main include(joinpath($(BASE_TEST_PATH), "testhelpers", "FillArrays.jl")) using .Main.FillArrays +isdefined(Main, :SizedArrays) || @eval Main include(joinpath($(BASE_TEST_PATH), "testhelpers", "SizedArrays.jl")) +using .Main.SizedArrays + Random.seed!(123) n = 5 # should be odd @@ -725,4 +728,56 @@ end @test det(A) == det(M) end +@testset "tril/triu" begin + @testset "with partly initialized matrices" begin + function test_triu(M, k=nothing) + M[1,1] = M[2,2] = M[1,2] = M[1,3] = M[2,3] = 3 + if isnothing(k) + MU = triu(M) + else + MU = triu(M, k) + end + @test iszero(MU[2,1]) + @test MU[1,1] == MU[2,2] == MU[1,2] == MU[1,3] == MU[2,3] == 3 + end + test_triu(Matrix{BigInt}(undef, 2, 3)) + test_triu(Matrix{BigInt}(undef, 2, 3), 0) + test_triu(SizedArrays.SizedArray{(2,3)}(Matrix{BigInt}(undef, 2, 3))) + test_triu(SizedArrays.SizedArray{(2,3)}(Matrix{BigInt}(undef, 2, 3)), 0) + + function test_tril(M, k=nothing) + M[1,1] = M[2,2] = M[2,1] = 3 + if isnothing(k) + ML = tril(M) + else + ML = tril(M, k) + end + @test ML[1,2] == ML[1,3] == ML[2,3] == 0 + @test ML[1,1] == ML[2,2] == ML[2,1] == 3 + end + test_tril(Matrix{BigInt}(undef, 2, 3)) + test_tril(Matrix{BigInt}(undef, 2, 3), 0) + test_tril(SizedArrays.SizedArray{(2,3)}(Matrix{BigInt}(undef, 2, 3))) + test_tril(SizedArrays.SizedArray{(2,3)}(Matrix{BigInt}(undef, 2, 3)), 0) + end + + @testset "block arrays" begin + for nrows in 0:3, ncols in 0:3 + M = [randn(2,2) for _ in 1:nrows, _ in 1:ncols] + Mu = triu(M) + for col in axes(M,2) + rowcutoff = min(col, size(M,1)) + @test @views Mu[1:rowcutoff, col] == M[1:rowcutoff, col] + @test @views Mu[rowcutoff+1:end, col] == zero.(M[rowcutoff+1:end, col]) + end + Ml = tril(M) + for col in axes(M,2) + @test @views Ml[col:end, col] == M[col:end, col] + rowcutoff = min(col-1, size(M,1)) + @test @views Ml[1:rowcutoff, col] == zero.(M[1:rowcutoff, col]) + end + end + end +end + end # module TestGeneric