Merge pull request #391 from ChrisRackauckas-Claude/fix-lu-componentmatrix

Make lu/qr of ComponentMatrix preserve the wrapper on factors

Author: ChrisRackauckas

Project.toml

@@ -1,7 +1,7 @@
 name = "ComponentArrays"
 uuid = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 authors = ["Jonnie Diegelman <47193959+jonniedie@users.noreply.github.com>"]
-version = "0.15.41"
+version = "0.15.42"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/ComponentArrays.jl

@@ -5,8 +5,8 @@ import StaticArrayInterface, ArrayInterface, Functors
 import ConstructionBase
 import Adapt
 
-using LinearAlgebra: LinearAlgebra, Adjoint, Cholesky, Diagonal, I, LU, Transpose,
-    UniformScaling, axpby!, axpy!, ldiv!
+using LinearAlgebra: LinearAlgebra, Adjoint, Cholesky, ColumnNorm, Diagonal, I, LU,
+    NoPivot, RowMaximum, RowNonZero, Transpose, UniformScaling, axpby!, axpy!, ldiv!
 using StaticArraysCore: StaticArray, SArray, SVector
 
 const FlatIdx = Union{Integer, CartesianIndex, CartesianIndices, AbstractArray{<:Integer}}

src/linear_algebra.jl

@@ -8,6 +8,40 @@ function ArrayInterface.lu_instance(x::ComponentMatrix)
     return LU{luT}(similar(x), ipiv, info)
 end
 
+# The generic `lu`/`qr` copy through `similar(A, T, dims)`, which drops the axes
+# and returns a plain `Matrix`, whereas the in-place `lu!`/`qr!` (ComponentMatrix
+# is strided, so they hit LAPACK directly) keep the ComponentMatrix wrapper on the
+# factors. Make the out-of-place versions and the `ArrayInterface` instance
+# functions preserve the wrapper too, so callers that type a cache from the
+# instance functions and later store either variant (e.g. LinearSolve) don't hit
+# a cache type mismatch.
+function LinearAlgebra.lu(
+        A::ComponentMatrix, pivot::Union{NoPivot, RowNonZero, RowMaximum} = RowMaximum();
+        kwargs...,
+    )
+    F = LinearAlgebra.lu(getdata(A), pivot; kwargs...)
+    return LU(ComponentArray(F.factors, getaxes(A)), F.ipiv, F.info)
+end
+
+function LinearAlgebra.qr(
+        A::ComponentMatrix, pivot::Union{NoPivot, ColumnNorm} = NoPivot();
+        kwargs...,
+    )
+    F = LinearAlgebra.qr(getdata(A), pivot; kwargs...)
+    return _rewrap_qr(F, getaxes(A))
+end
+
+function ArrayInterface.qr_instance(A::ComponentMatrix, pivot = NoPivot())
+    F = ArrayInterface.qr_instance(getdata(A), pivot)
+    return _rewrap_qr(F, getaxes(A))
+end
+
+function _rewrap_qr(F::LinearAlgebra.QRCompactWY, ax)
+    return LinearAlgebra.QRCompactWY(ComponentArray(F.factors, ax), F.T)
+end
+_rewrap_qr(F::LinearAlgebra.QRPivoted, ax) = LinearAlgebra.QRPivoted(ComponentArray(F.factors, ax), F.τ, F.jpvt)
+_rewrap_qr(F, ax) = F
+
 # Helpers for dealing with adjoints and such
 _first_axis(x::AbstractComponentVecOrMat) = getaxes(x)[1]
 

test/math_tests.jl

@@ -143,4 +143,25 @@ s2_D = out2 * in2'
     (FlatAxis(), Axis(u2 = 1))
 
 @test ComponentArrays.ArrayInterface.lu_instance(cmat).factors isa ComponentMatrix
+
+# `lu`, `lu!`, and `lu_instance` (same for qr) must all keep the ComponentMatrix
+# wrapper on the factors, or caches typed from the instance functions break when
+# a `lu`/`qr` result is stored in them (LinearSolve's default solver does this;
+# see the PR #390 downstream CI failure).
+let A = cmat + I
+    F = lu(A)
+    @test F.factors isa ComponentMatrix
+    @test typeof(F) === typeof(ComponentArrays.ArrayInterface.lu_instance(A))
+    @test typeof(F) === typeof(lu!(copy(A)))
+    @test getdata(F.factors) == lu(getdata(A)).factors
+    @test F \ getdata(ca) ≈ getdata(A) \ getdata(ca)
+
+    Fqr = qr(A)
+    @test Fqr.factors isa ComponentMatrix
+    @test typeof(Fqr) === typeof(ComponentArrays.ArrayInterface.qr_instance(A, NoPivot()))
+    @test typeof(Fqr) === typeof(qr!(copy(A)))
+    @test Fqr \ getdata(ca) ≈ getdata(A) \ getdata(ca)
+    @test typeof(qr(A, ColumnNorm())) === typeof(qr!(copy(A), ColumnNorm())) ===
+        typeof(ComponentArrays.ArrayInterface.qr_instance(A, ColumnNorm()))
+end
 @test ComponentArrays.ArrayInterface.parent_type(cmat) === Matrix{Float64}