Use multiple dispatch for get_sparsity_pattern

Author: amontoison

docs/src/sparse.md

@@ -77,8 +77,11 @@ lcon = -0.5 * ones(T, ncon)
 ucon = 0.5 * ones(T, ncon)
 
 nlp = ADNLPModel!(f, x0, lvar, uvar, c!, lcon, ucon)
-
+```
+```@example
 J = get_sparsity_pattern(nlp, :jacobian)
+```
+```@example
 H = get_sparsity_pattern(nlp, :hessian)
 ```
 

src/sparsity_pattern.jl

@@ -74,42 +74,58 @@ Supported backends include `SparseADJacobian`, `SparseADHessian`, and `SparseRev
 * `S`: A sparse matrix of type `SparseMatrixCSC{Bool,Int}` indicating the sparsity pattern of the requested derivative.
 """
 function get_sparsity_pattern(model::ADModel, derivative::Symbol)
-  if (derivative != :jacobian) && (derivative != :hessian)
-    if model isa AbstractADNLPModel
-      error("The only supported sparse derivatives for an AbstractADNLPModel are `:jacobian` and `:hessian`.")
-    elseif (model isa AbstractADNLSModel) && (derivative != :jacobian_residual) && (derivative != :hessian_resiual)
-      error("The only supported sparse derivatives for an AbstractADNLSModel are `:jacobian`, `:jacobian_residual`, `:hessian` and `:hessian_resiual`.")
-    end
-  end
-  if (derivative == :jacobian) || (derivative == :jacobian_residual)
-    backend = derivative == :jacobian ? model.adbackend.jacobian_backend : model.adbackend.jacobian_residual_backend
-    if backend isa SparseADJacobian
-      m = derivative == :jacobian ? model.meta.ncon : model.nls_meta.nequ
-      n = model.meta.nvar
-      colptr = backend.colptr
-      rowval = backend.rowval
-      nnzJ = length(rowval)
-      nzval = ones(Bool, nnzJ)
-      J = SparseMatrixCSC(m, n, colptr, rowval, nzval)
-      return J
-    else
-      B = typeof(backend)
-      error("The current backend ($B) doesn't compute a sparse Jacobian.")
-    end
-  end
-  if (derivative == :hessian) || (derivative == :hessian_residual)
-    backend = derivative == :hessian ? model.adbackend.hessian_backend : model.adbackend.hessian_residual_backend
-    if (backend isa SparseADHessian) || (backend isa SparseReverseADHessian)
-      n = model.meta.nvar
-      colptr = backend.colptr
-      rowval = backend.rowval
-      nnzH = length(rowval)
-      nzval = ones(Bool, nnzH)
-      H = SparseMatrixCSC(n, n, colptr, rowval, nzval)
-      return H
-    else
-      B = typeof(backend)
-      error("The current backend ($B) doesn't compute a sparse Hessian.")
+    get_sparsity_pattern(model, Val(derivative))
+end
+
+function get_sparsity_pattern(model::ADModel, ::Val{:jacobian})
+    backend = model.adbackend.jacobian_backend
+    validate_sparse_backend(backend, SparseADJacobian, "Jacobian")
+    m = model.meta.ncon
+    n = model.meta.nvar
+    colptr = backend.colptr
+    rowval = backend.rowval
+    nnzJ = length(rowval)
+    nzval = ones(Bool, nnzJ)
+    SparseMatrixCSC(m, n, colptr, rowval, nzval)
+end
+
+function get_sparsity_pattern(model::ADModel, ::Val{:hessian})
+    backend = model.adbackend.hessian_backend
+    validate_sparse_backend(backend, Union{SparseADHessian, SparseReverseADHessian}, "Hessian")
+    n = model.meta.nvar
+    colptr = backend.colptr
+    rowval = backend.rowval
+    nnzH = length(rowval)
+    nzval = ones(Bool, nnzH)
+    SparseMatrixCSC(n, n, colptr, rowval, nzval)
+end
+
+function get_sparsity_pattern(model::AbstractADNLSModel, ::Val{:jacobian_residual})
+    backend = model.adbackend.jacobian_residual_backend
+    validate_sparse_backend(backend, SparseADJacobian, "Jacobian of the residual")
+    m = model.nls_meta.nequ
+    n = model.meta.nvar
+    colptr = backend.colptr
+    rowval = backend.rowval
+    nnzJ = length(rowval)
+    nzval = ones(Bool, nnzJ)
+    SparseMatrixCSC(m, n, colptr, rowval, nzval)
+end
+
+function get_sparsity_pattern(model::AbstractADNLSModel, ::Val{:hessian_residual})
+    backend = model.adbackend.hessian_residual_backend
+    validate_sparse_backend(backend, Union{SparseADHessian, SparseReverseADHessian}, "Hessian of the residual")
+    n = model.meta.nvar
+    colptr = backend.colptr
+    rowval = backend.rowval
+    nnzH = length(rowval)
+    nzval = ones(Bool, nnzH)
+    SparseMatrixCSC(n, n, colptr, rowval, nzval)
+end
+
+function validate_sparse_backend(backend::ADBackend, expected_type, derivative_name::String)
+    if !(backend isa expected_type)
+        B = typeof(backend)
+        error("The current backend $B doesn't compute a sparse $derivative_name.")
     end
-  end
 end

test/sparse_hessian.jl

@@ -65,7 +65,7 @@ dt = (Float32, Float64)
 
   if (backend == ADNLPModels.SparseADHessian) || (backend == ADNLPModels.SparseReverseADHessian)
     H_sp = get_sparsity_pattern(nlp, :hessian)
-    @test H_sp == SparseMatrixCSC{Bool,Int}(
+    @test H_sp == SparseMatrixCSC{Bool, Int}(
       [ 1 0 ;
         1 1 ]
     )

test/sparse_hessian_nls.jl

@@ -52,7 +52,7 @@ dt = (Float32, Float64)
 
   if (backend == ADNLPModels.SparseADHessian) || (backend == ADNLPModels.SparseReverseADHessian)
     H_sp = get_sparsity_pattern(nls, :hessian_residual)
-    @test H_sp == SparseMatrixCSC{Bool,Int}(
+    @test H_sp == SparseMatrixCSC{Bool, Int}(
       [ 1 0 ;
         0 0 ]
     )

test/sparse_jacobian.jl

@@ -53,7 +53,7 @@ dt = (Float32, Float64)
 
   if backend == ADNLPModels.SparseADJacobian
     J_sp = get_sparsity_pattern(nlp, :jacobian)
-    @test J_sp == SparseMatrixCSC{Bool,Int}(
+    @test J_sp == SparseMatrixCSC{Bool, Int}(
       [ 1 0 ;
         1 1 ;
         0 1 ]

test/sparse_jacobian_nls.jl

@@ -45,7 +45,7 @@ dt = (Float32, Float64)
 
   if backend == ADNLPModels.SparseADJacobian
     J_sp = get_sparsity_pattern(nls, :jacobian_residual)
-    @test J_sp == SparseMatrixCSC{Bool,Int}(
+    @test J_sp == SparseMatrixCSC{Bool, Int}(
       [ 1 0 ;
         1 1 ;
         0 1 ]