Fix missing copyto! methods for triangular matrix types

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

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