From 7a61d18e16a3f666d8a2c312223dc6ec6ae58e6b Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Tue, 11 Jun 2024 18:10:31 +0530 Subject: [PATCH] Fix missing copyto! methods for triangular matrix types --- stdlib/LinearAlgebra/src/triangular.jl | 39 +++++++--- stdlib/LinearAlgebra/test/triangular.jl | 96 +++++++++++++++++++------ 2 files changed, 101 insertions(+), 34 deletions(-) diff --git a/stdlib/LinearAlgebra/src/triangular.jl b/stdlib/LinearAlgebra/src/triangular.jl index 37c31cdf31dcc..c45db3e90fab2 100644 --- a/stdlib/LinearAlgebra/src/triangular.jl +++ b/stdlib/LinearAlgebra/src/triangular.jl @@ -499,41 +499,58 @@ tr(A::UnitLowerTriangular) = size(A, 1) * oneunit(eltype(A)) tr(A::UpperTriangular) = tr(A.data) tr(A::UnitUpperTriangular) = size(A, 1) * oneunit(eltype(A)) -@propagate_inbounds function copyto!(dest::UpperOrLowerTriangular, U::UpperOrLowerTriangular) - if axes(dest) != axes(U) - @invoke copyto!(dest::AbstractArray, U::AbstractArray) - else - _copyto!(dest, U) +for T in (:UpperOrUnitUpperTriangular, :LowerOrUnitLowerTriangular) + @eval @propagate_inbounds function copyto!(dest::$T, U::$T) + if axes(dest) != axes(U) + @invoke copyto!(dest::AbstractArray, U::AbstractArray) + else + _copyto!(dest, U) + end + return dest end - return dest end # copy and scale -for T in (:UpperTriangular, :LowerTriangular) +for (T, UT) in ((:UpperTriangular, :UnitUpperTriangular), (:LowerTriangular, :UnitLowerTriangular)) @eval @inline function _copyto!(A::$T, B::$T) @boundscheck checkbounds(A, axes(B)...) copytrito!(parent(A), parent(B), uplo_char(A)) return A end + @eval @inline function _copyto!(A::$UT, B::$T) + for dind in diagind(A, IndexStyle(A)) + if A[dind] != B[dind] + throw_nononeerror(typeof(A), B[dind], Tuple(dind)...) + end + end + _copyto!($T(parent(A)), B) + return A + end end -@inline function _copyto!(A::UnitUpperTriangular, B::UnitUpperTriangular) +@inline function _copyto!(A::UpperOrUnitUpperTriangular, B::UnitUpperTriangular) @boundscheck checkbounds(A, axes(B)...) n = size(B,1) B2 = Base.unalias(A, B) for j = 1:n for i = 1:j-1 - @inbounds A[i,j] = B2[i,j] + @inbounds parent(A)[i,j] = parent(B2)[i,j] + end + if A isa UpperTriangular # copy diagonal + @inbounds parent(A)[j,j] = B2[j,j] end end return A end -@inline function _copyto!(A::UnitLowerTriangular, B::UnitLowerTriangular) +@inline function _copyto!(A::LowerOrUnitLowerTriangular, B::UnitLowerTriangular) @boundscheck checkbounds(A, axes(B)...) n = size(B,1) B2 = Base.unalias(A, B) for j = 1:n + if A isa LowerTriangular # copy diagonal + @inbounds parent(A)[j,j] = B2[j,j] + end for i = j+1:n - @inbounds A[i,j] = B2[i,j] + @inbounds parent(A)[i,j] = parent(B2)[i,j] end end return A diff --git a/stdlib/LinearAlgebra/test/triangular.jl b/stdlib/LinearAlgebra/test/triangular.jl index 93077c5631390..5ee8143e3f4bb 100644 --- a/stdlib/LinearAlgebra/test/triangular.jl +++ b/stdlib/LinearAlgebra/test/triangular.jl @@ -1081,32 +1081,82 @@ end end end -@testset "copyto! with aliasing (#39460)" begin - M = Matrix(reshape(1:36, 6, 6)) - @testset for T in (UpperTriangular, LowerTriangular, UnitUpperTriangular, UnitLowerTriangular) - A = T(view(M, 1:5, 1:5)) - A2 = copy(A) - B = T(view(M, 2:6, 2:6)) - @test copyto!(B, A) == A2 +@testset "copyto! tests" begin + @testset "copyto! with aliasing (#39460)" begin + M = Matrix(reshape(1:36, 6, 6)) + @testset for T in (UpperTriangular, LowerTriangular, UnitUpperTriangular, UnitLowerTriangular) + A = T(view(M, 1:5, 1:5)) + A2 = copy(A) + B = T(view(M, 2:6, 2:6)) + @test copyto!(B, A) == A2 + end end -end -@testset "copyto! with different sizes" begin - Ap = zeros(3,3) - Bp = rand(2,2) - @testset for T in (UpperTriangular, LowerTriangular, UnitUpperTriangular, UnitLowerTriangular) - A = T(Ap) - B = T(Bp) - @test_throws ArgumentError copyto!(A, B) + @testset "copyto! with different matrix types" begin + M1 = Matrix(reshape(1:36, 6, 6)) + M2 = similar(M1) + # these copies always work + @testset for (Tdest, Tsrc) in ( + (UpperTriangular, UnitUpperTriangular), + (UpperTriangular, UpperTriangular), + (LowerTriangular, UnitLowerTriangular), + (LowerTriangular, LowerTriangular), + (UnitUpperTriangular, UnitUpperTriangular), + (UnitLowerTriangular, UnitLowerTriangular) + ) + + M2 .= 0 + copyto!(Tdest(M2), Tsrc(M1)) + @test Tdest(M2) == Tsrc(M1) + end + # these copies only work if the source has a unit diagonal + M3 = copy(M1) + M3[diagind(M3)] .= 1 + @testset for (Tdest, Tsrc) in ( + (UnitUpperTriangular, UpperTriangular), + (UnitLowerTriangular, LowerTriangular), + ) + + M2 .= 0 + copyto!(Tdest(M2), Tsrc(M3)) + @test Tdest(M2) == Tsrc(M3) + @test_throws ArgumentError copyto!(Tdest(M2), Tsrc(M1)) + end + # these copies work even when the parent of the source isn't initialized along the diagonal + @testset for (T, TU) in ((UpperTriangular, UnitUpperTriangular), + (LowerTriangular, UnitLowerTriangular)) + M1 = Matrix{BigFloat}(undef, 3, 3) + M2 = similar(M1) + if TU == UnitUpperTriangular + M1[1,2] = M1[1,3] = M1[2,3] = 2 + else + M1[2,1] = M1[3,1] = M1[3,2] = 2 + end + for TD in (T, TU) + M2 .= 0 + copyto!(T(M2), TU(M1)) + @test T(M2) == TU(M1) + end + end end - @testset "error message" begin - A = UpperTriangular(Ap) - B = UpperTriangular(Bp) - @test_throws "cannot set index in the lower triangular part" copyto!(A, B) - - A = LowerTriangular(Ap) - B = LowerTriangular(Bp) - @test_throws "cannot set index in the upper triangular part" copyto!(A, B) + + @testset "copyto! with different sizes" begin + Ap = zeros(3,3) + Bp = rand(2,2) + @testset for T in (UpperTriangular, LowerTriangular, UnitUpperTriangular, UnitLowerTriangular) + A = T(Ap) + B = T(Bp) + @test_throws ArgumentError copyto!(A, B) + end + @testset "error message" begin + A = UpperTriangular(Ap) + B = UpperTriangular(Bp) + @test_throws "cannot set index in the lower triangular part" copyto!(A, B) + + A = LowerTriangular(Ap) + B = LowerTriangular(Bp) + @test_throws "cannot set index in the upper triangular part" copyto!(A, B) + end end end