Optional SparseMatrixCSR support via SparseMatricesCSR extension

Project.toml

@@ -1,19 +1,20 @@
 name = "SparseColumnPivotedQR"
 uuid = "a57abbd0-fea5-4d57-96be-5e525945e8e4"
 authors = ["Chris Rackauckas <accounts@chrisrackauckas.com> and contributors"]
-version = "0.1.0"
+version = "0.2.0"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
-SparseMatricesCSR = "a0a7dd2c-ebf4-11e9-1f05-cf50bc540ca1"
 
 [weakdeps]
 AMD = "14f7f29c-3bd6-536c-9a0b-7339e30b5a3e"
+SparseMatricesCSR = "a0a7dd2c-ebf4-11e9-1f05-cf50bc540ca1"
 
 [extensions]
 SparseColumnPivotedQRAMDExt = "AMD"
+SparseColumnPivotedQRSparseMatricesCSRExt = "SparseMatricesCSR"
 
 [compat]
 AMD = "0.5"
@@ -27,7 +28,8 @@ julia = "1.10"
 AMD = "14f7f29c-3bd6-536c-9a0b-7339e30b5a3e"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+SparseMatricesCSR = "a0a7dd2c-ebf4-11e9-1f05-cf50bc540ca1"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["AMD", "ForwardDiff", "Random", "Test"]
+test = ["AMD", "ForwardDiff", "Random", "SparseMatricesCSR", "Test"]

ext/SparseColumnPivotedQRSparseMatricesCSRExt.jl

@@ -0,0 +1,67 @@
+module SparseColumnPivotedQRSparseMatricesCSRExt
+
+using SparseColumnPivotedQR
+using LinearAlgebra
+using SparseArrays
+using SparseMatricesCSR
+using PrecompileTools
+
+import SparseColumnPivotedQR: csr_qr, csr_analyze, csr_factor, csr_refactor!,
+    CSRQRSymbolic, CSRQRFactorization
+
+# Convert a `SparseMatrixCSR` to the `SparseMatrixCSC` the core operates on.
+# This is the same CSR -> CSC conversion the kernel performed internally before
+# the CSC-native refactor; it now lives here so the core never depends on
+# `SparseMatricesCSR`.
+@inline _to_csc(A::SparseMatrixCSR) = SparseMatrixCSC(A)
+
+function csr_analyze(A::SparseMatrixCSR; ordering::Symbol = :default)
+    return csr_analyze(_to_csc(A); ordering = ordering)
+end
+
+function csr_factor(A::SparseMatrixCSR, sym::CSRQRSymbolic; kwargs...)
+    return csr_factor(_to_csc(A), sym; kwargs...)
+end
+
+function csr_qr(A::SparseMatrixCSR; kwargs...)
+    return csr_qr(_to_csc(A); kwargs...)
+end
+
+function csr_refactor!(F::CSRQRFactorization, A::SparseMatrixCSR; kwargs...)
+    return csr_refactor!(F, _to_csc(A); kwargs...)
+end
+
+# Keep the CSR entry points specialized in the package image. Mirrors the
+# core workload, but on `SparseMatrixCSR` inputs through the conversion path.
+@setup_workload begin
+    @compile_workload begin
+        for T in (Float64, Float32, ComplexF64, ComplexF32)
+            for Ti in (Int32, Int64)
+                rows = Ti[1, 2, 3, 4, 5, 6, 1, 2, 3, 4]
+                cols = Ti[1, 2, 3, 4, 5, 6, 2, 3, 4, 5]
+                vals = T[4, 4, 4, 4, 4, 4, 1, 1, 1, 1]
+                A = sparsecsr(rows, cols, vals, 6, 6)
+                b = ones(T, 6)
+
+                F = csr_qr(A; ordering = :natural)
+                F \ b
+                rank(F)
+
+                sym = csr_analyze(A; ordering = :natural)
+                G = csr_factor(A, sym)
+                csr_refactor!(G, A)
+                G \ b
+
+                drows = Ti[1, 2, 3, 4, 5, 1, 2, 3, 4, 6]
+                dcols = Ti[1, 2, 3, 4, 5, 2, 3, 4, 5, 1]
+                dvals = T[4, 4, 4, 4, 4, 1, 1, 1, 1, 4]
+                Ad = sparsecsr(drows, dcols, dvals, 6, 6)
+                Fd = csr_qr(Ad; ordering = :natural)
+                Fd \ b
+                rank(Fd)
+            end
+        end
+    end
+end
+
+end # module