Add pivot tolerance option (#76)

* Add option to set pivot tolerance

* Test pivot tolerance option

* Bump minor version

Author: hersle

Project.toml

@@ -1,7 +1,7 @@
 name = "PureKLU"
 uuid = "0c0d3e7f-3a8b-4f7e-b6f1-9a4d2e7c1f01"
 authors = ["Chris Rackauckas <accounts@chrisrackauckas.com> and contributors"]
-version = "1.3.0"
+version = "1.4.0"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/PureKLU.jl

@@ -515,6 +515,10 @@ Compute a KLU sparse LU factorization.
 - `fully_preallocated = nothing`: select automatic workspace preallocation;
   pass `true` or `false` to override it.
 - `detect_banded = true`: detect narrow BTF blocks and use natural ordering.
+- `tol = 0.001`: Pivot on a column's diagonal instead of largest entry if it is
+  at least `tol` times larger in magnitude. Set `tol = 1.0` for partial pivoting,
+  and `tol = 0.0` to always use the diagonal. Only applies to the numeric
+  factorization; `klu!` reuses the existing pivot ordering.
 
 # Returns
 - A `KLUFactorization` that implements `LinearAlgebra.Factorization`.
@@ -536,11 +540,12 @@ function klu(
         check::Bool = true, allowsingular::Bool = false,
         full_factor::Bool = true, use_fma = true,
         fully_preallocated::Union{Bool, Nothing} = nothing,
-        detect_banded::Bool = true,
+        detect_banded::Bool = true, tol::Float64 = 0.001,
     ) where {Ti <: KLUITypes, Tv <: KLUGenericTypes}
     K = KLUFactorization(n, colptr, rowval, nzval)
     K.common.use_fma = _as_val(use_fma)
     K.common.detect_banded = detect_banded
+    K.common.tol = tol
     if fully_preallocated isa Bool
         K.common.fully_preallocated = fully_preallocated
         return full_factor ? klu_factor!(K; check, allowsingular) : klu_analyze!(K; check)
@@ -554,13 +559,14 @@ function klu(
         A::SparseMatrixCSC{Tv, Ti}; check::Bool = true,
         allowsingular::Bool = false, full_factor::Bool = true,
         use_fma = true, fully_preallocated::Union{Bool, Nothing} = nothing,
-        detect_banded::Bool = true,
+        detect_banded::Bool = true, tol::Float64 = 0.001,
     ) where {Tv <: KLUGenericTypes, Ti <: KLUITypes}
     n = size(A, 1)
     n == size(A, 2) || throw(DimensionMismatch())
     return klu(
         n, decrement(A.colptr), decrement(A.rowval), A.nzval;
-        check, allowsingular, full_factor, use_fma, fully_preallocated, detect_banded
+        check, allowsingular, full_factor, use_fma, fully_preallocated,
+        detect_banded, tol
     )
 end
 

test/test_tol.jl

@@ -0,0 +1,33 @@
+# Trivial test that the `tol` kwarg (diagonal-vs-partial pivoting threshold)
+# is actually honored, not just accepted and ignored.
+using Test
+using SparseArrays
+using LinearAlgebra
+using PureKLU
+
+@testset "tol kwarg: stored on common, solves correctly" begin
+    A = sparse([2.0 1.0; 1.0 3.0])
+    b = [1.0, 2.0]
+    for tol in (0.0, 0.001, 0.5, 1.0)
+        K = PureKLU.klu(A; tol)
+        @test K.common.tol == tol
+        @test A * (K \ b) ≈ b
+    end
+end
+
+@testset "tol kwarg: extreme values change pivot choice" begin
+    # Column 1 has a tiny diagonal entry (0.001) and a much larger
+    # off-diagonal entry (1.0) in row 2. tol = 0.0 always pivots on the
+    # diagonal; tol = 1.0 is full partial pivoting and must pick the larger
+    # entry instead, so the two runs choose different row pivots.
+    A = sparse([0.001 1.0; 1.0 1.0])
+    b = [1.0, 1.0]
+
+    K_diag = PureKLU.klu(A; tol = 0.0, detect_banded = false)
+    K_partial = PureKLU.klu(A; tol = 1.0, detect_banded = false)
+
+    @test K_diag.p == [1, 2]      # diagonal pivot kept despite tiny magnitude
+    @test K_partial.p == [2, 1]   # partial pivoting swaps to the larger entry
+    @test A * (K_diag \ b) ≈ b
+    @test A * (K_partial \ b) ≈ b
+end