Modularize assembly internals and add SparseMatrixCSR extension

Project.toml

@@ -18,17 +18,20 @@ WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
 [weakdeps]
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
 Metis = "2679e427-3c69-5b7f-982b-ece356f1e94b"
+SparseMatricesCSR = "a0a7dd2c-ebf4-11e9-1f05-cf50bc540ca1"
 
 [extensions]
 FerriteBlockArrays = "BlockArrays"
 FerriteMetis = "Metis"
+FerriteSparseMatrixCSR = "SparseMatricesCSR"
 
 [compat]
 BlockArrays = "0.16, 1"
 EnumX = "1"
 HCubature = "1.7.0"
 ForwardDiff = "0.10, 1"
 Metis = "1.3"
+SparseMatricesCSR = "0.6"
 NearestNeighbors = "0.4"
 OrderedCollections = "1"
 Preferences = "1"
@@ -53,9 +56,10 @@ ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 HCubature = "19dc6840-f33b-545b-b366-655c7e3ffd49"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SHA = "ea8e919c-243c-51af-8825-aaa63cd721ce"
+SparseMatricesCSR = "a0a7dd2c-ebf4-11e9-1f05-cf50bc540ca1"
 TaskLocalValues = "ed4db957-447d-4319-bfb6-7fa9ae7ecf34"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 
 [targets]
-test = ["BlockArrays", "Downloads", "FerriteGmsh", "ForwardDiff", "Gmsh", "IterativeSolvers", "Metis", "NBInclude", "OhMyThreads", "ProgressMeter", "Random", "SHA", "TaskLocalValues", "Test", "TimerOutputs", "Logging", "HCubature"]
+test = ["BlockArrays", "Downloads", "FerriteGmsh", "ForwardDiff", "Gmsh", "IterativeSolvers", "Metis", "NBInclude", "OhMyThreads", "ProgressMeter", "Random", "SHA", "SparseMatricesCSR", "TaskLocalValues", "Test", "TimerOutputs", "Logging", "HCubature"]

docs/src/devdocs/assembly.md

@@ -1,10 +1,48 @@
 # [Assembly](@id devdocs-assembly)
 
+The assembler handles the insertion of the element matrices and element vectors into the system matrix and right hand side.
+
+## Custom matrix formats
+While the CSC and CSR formats are the most common sparse matrix formats in practice, users might want to have optimized custom matrix formats for their specific use-case. The default assemblers [`Ferrite.CSCAssembler`](@ref) and [`Ferrite.CSRAssembler`](@ref) should be able to handle most cases in practice. To support a custom format users have to dispatch the following functions on their matrix type. There is the public interface
+
+```@docs; canonical=false
+Ferrite.allocate_matrix
+```
+
+the internal interface
+```@docs
+Ferrite.zero_out_rows!
+Ferrite.zero_out_columns!
+Ferrite._condense!
+```
+
+and the `AbstractSparseMatrix` interface for their custom matrix type. Optional dispatches to speed up operations might be
+
+```@docs
+Ferrite.add_inhomogeneities!
+```
+
+## Custom Assembler
+
+In case the default assembler is insufficient, users can implement a custom assemblers. For this, they can create a custom type and dispatch the following functions.
+
+```@docs; canonical=false
+start_assemble
+assemble!
+```
+
+For local elimination support the following functions might also need custom dispatches
+
+```@docs
+Ferrite._condense_local!
+```
+
 ## Type definitions
 
 ```@docs
 Ferrite.COOAssembler
 Ferrite.CSCAssembler
+Ferrite.CSRAssembler
 Ferrite.SymmetricCSCAssembler
 ```
 

ext/FerriteSparseMatrixCSR.jl

@@ -0,0 +1,110 @@
+module FerriteSparseMatrixCSR
+
+using Ferrite, SparseArrays, SparseMatricesCSR
+import Ferrite: AbstractSparsityPattern, CSRAssembler
+import Base: @propagate_inbounds
+
+# Could be generalized if https://github.com/JuliaSparse/SparseArrays.jl/pull/546 is merged
+function Ferrite.start_assemble(K::SparseMatrixCSR{<:Any, T}, f::Vector = T[]; fillzero::Bool = true, maxcelldofs_hint::Int = 0) where {T}
+    fillzero && (Ferrite.fillzero!(K); Ferrite.fillzero!(f))
+    return CSRAssembler(K, f, zeros(Int, maxcelldofs_hint), zeros(Int, maxcelldofs_hint))
+end
+
+@propagate_inbounds function Ferrite._assemble_inner!(K::SparseMatrixCSR, Ke::AbstractMatrix, dofs::AbstractVector, sorteddofs::AbstractVector, permutation::AbstractVector, sym::Bool)
+    current_row = 1
+    ld = length(dofs)
+    return @inbounds for Krow in sorteddofs
+        maxlookups = sym ? current_row : ld
+        Kerow = permutation[current_row]
+        ci = 1 # col index pointer for the local matrix
+        Ci = 1 # col index pointer for the global matrix
+        nzr = nzrange(K, Krow)
+        while Ci <= length(nzr) && ci <= maxlookups
+            C = nzr[Ci]
+            Kcol = K.colval[C]
+            Kecol = permutation[ci]
+            val = Ke[Kerow, Kecol]
+            if Kcol == dofs[Kecol]
+                # Match: add the value (if non-zero) and advance the pointers
+                if !iszero(val)
+                    K.nzval[C] += val
+                end
+                ci += 1
+                Ci += 1
+            elseif Kcol < dofs[Kecol]
+                # No match yet: advance the global matrix row pointer
+                Ci += 1
+            else # Kcol > dofs[Kecol]
+                # No match: no entry exist in the global matrix for this row. This is
+                # allowed as long as the value which would have been inserted is zero.
+                iszero(val) || Ferrite._missing_sparsity_pattern_error(Krow, Kcol)
+                # Advance the local matrix row pointer
+                ci += 1
+            end
+        end
+        # Make sure that remaining entries in this column of the local matrix are all zero
+        for i in ci:maxlookups
+            if !iszero(Ke[Kerow, permutation[i]])
+                Ferrite._missing_sparsity_pattern_error(Krow, sorteddofs[i])
+            end
+        end
+        current_row += 1
+    end
+end
+
+function Ferrite.zero_out_rows!(K::SparseMatrixCSR, ch::ConstraintHandler)
+    @debug @assert issorted(ch.prescribed_dofs)
+    for row in ch.prescribed_dofs
+        r = nzrange(K, row)
+        K.nzval[r] .= 0.0
+    end
+    return
+end
+
+function Ferrite.zero_out_columns!(K::SparseMatrixCSR, ch::ConstraintHandler)
+    colval = K.colval
+    nzval = K.nzval
+    return @inbounds for i in eachindex(colval, nzval)
+        if haskey(ch.dofmapping, colval[i])
+            nzval[i] = 0
+        end
+    end
+end
+
+function Ferrite.allocate_matrix(::Type{SparseMatrixCSR}, sp::AbstractSparsityPattern)
+    return Ferrite.allocate_matrix(SparseMatrixCSR{1, Float64, Int64}, sp)
+end
+
+function Ferrite.allocate_matrix(::Type{SparseMatrixCSR{1, Tv, Ti}}, sp::AbstractSparsityPattern) where {Tv, Ti}
+    return _allocate_matrix(SparseMatrixCSR{1, Tv, Ti}, sp, false)
+end
+
+function _allocate_matrix(::Type{SparseMatrixCSR{1, Tv, Ti}}, sp::AbstractSparsityPattern, sym::Bool) where {Tv, Ti}
+    # 1. Setup rowptr
+    rowptr = zeros(Ti, Ferrite.getnrows(sp) + 1)
+    rowptr[1] = 1
+    for (row, colidxs) in enumerate(Ferrite.eachrow(sp))
+        for col in colidxs
+            sym && row > col && continue
+            rowptr[row + 1] += 1
+        end
+    end
+    cumsum!(rowptr, rowptr)
+    nnz = rowptr[end] - 1
+    # 2. Allocate colval and nzval now that nnz is known
+    colval = Vector{Ti}(undef, nnz)
+    nzval = zeros(Tv, nnz)
+    # 3. Populate colval.
+    k = 1
+    for (row, colidxs) in zip(1:Ferrite.getnrows(sp), Ferrite.eachrow(sp)) # pairs(eachrow(sp))
+        for col in colidxs
+            sym && row > col && continue
+            colval[k] = col
+            k += 1
+        end
+    end
+    S = SparseMatrixCSR{1}(Ferrite.getnrows(sp), Ferrite.getncols(sp), rowptr, colval, nzval)
+    return S
+end
+
+end

src/Dofs/ConstraintHandler.jl

@@ -582,7 +582,7 @@
 
 """
 
-    apply!(K::SparseMatrixCSC, rhs::AbstractVector, ch::ConstraintHandler)
+    apply!(K::AbstractSparseMatrix, rhs::AbstractVector, ch::ConstraintHandler)
 
 Adjust the matrix `K` and right hand side `rhs` to account for the Dirichlet boundary
 conditions specified in `ch` such that `K \\ rhs` gives the expected solution.
@@ -674,15 +674,15 @@
     return v
 end
 
-function apply!(K::Union{SparseMatrixCSC, Symmetric}, ch::ConstraintHandler)
+function apply!(K::Union{AbstractSparseMatrix, Symmetric}, ch::ConstraintHandler)
     return apply!(K, eltype(K)[], ch, true)
 end
 
-function apply_zero!(K::Union{SparseMatrixCSC, Symmetric}, f::AbstractVector, ch::ConstraintHandler)
+function apply_zero!(K::Union{AbstractSparseMatrix, Symmetric}, f::AbstractVector, ch::ConstraintHandler)
     return apply!(K, f, ch, true)
 end
 
-function apply!(KK::Union{SparseMatrixCSC, Symmetric}, f::AbstractVector, ch::ConstraintHandler, applyzero::Bool = false)
+function apply!(KK::Union{AbstractSparseMatrix, Symmetric{<:Any, <:AbstractSparseMatrix}}, f::AbstractVector, ch::ConstraintHandler, applyzero::Bool = false)
     @assert isclosed(ch)
     sym = isa(KK, Symmetric)
     K = sym ? KK.data : KK
@@ -693,43 +693,14 @@
     m = meandiag(K) # Use the mean of the diagonal here to not ruin things for iterative solver
 
     # Add inhomogeneities to f: (f - K * ch.inhomogeneities)
-    if !applyzero
-        @inbounds for i in 1:length(ch.inhomogeneities)
-            d = ch.prescribed_dofs[i]
-            v = ch.inhomogeneities[i]
-            if v != 0
-                for j in nzrange(K, d)
-                    r = K.rowval[j]
-                    sym && r > d && break # don't look below diagonal
-                    f[r] -= v * K.nzval[j]
-                end
-            end
-        end
-        if sym
-            # In the symmetric case, for a constrained dof `d`, we handle the contribution
-            # from `K[1:d, d]` in the loop above, but we are still missing the contribution
-            # from `K[(d+1):size(K,1), d]`. These values are not stored, but since the
-            # matrix is symmetric we can instead use `K[d, (d+1):size(K,1)]`. Looping over
-            # rows is slow, so loop over all columns again, and check if the row is a
-            # constrained row.
-            @inbounds for col in 1:size(K, 2)
-                for ri in nzrange(K, col)
-                    row = K.rowval[ri]
-                    row >= col && break
-                    if (i = get(ch.dofmapping, row, 0); i != 0)
-                        f[col] -= ch.inhomogeneities[i] * K.nzval[ri]
-                    end
-                end
-            end
-        end
-    end
+    !applyzero && add_inhomogeneities!(f, KK, ch)
 
-    # Condense K (C' * K * C) and f (C' * f)
+    # Condense K := (C' * K * C) and f := (C' * f)
     _condense!(K, f, ch.dofcoefficients, ch.dofmapping, sym)
 
     # Remove constrained dofs from the matrix
-    zero_out_columns!(K, ch.prescribed_dofs)
-    zero_out_rows!(K, ch.dofmapping)
+    zero_out_columns!(K, ch)
+    zero_out_rows!(K, ch)
 
     # Add meandiag to constraint dofs
     @inbounds for i in 1:length(ch.inhomogeneities)
@@ -743,6 +714,53 @@
     return
 end
 
+"""
+    add_inhomogeneities!(f::AbstractVector, K::AbstractMatrix, ch::ConstraintHandler)
+
+Compute "f -= K*inhomogeneities".
+By default this is a generic version via SpMSpV kernel.
+"""
+function add_inhomogeneities!(f::AbstractVector, K::AbstractMatrix, ch::ConstraintHandler)
+    return mul!(f, K, sparsevec(ch.prescribed_dofs, ch.inhomogeneities, size(K, 2)), -1, 1)
+end
+
+# Optimized version for SparseMatrixCSC
+add_inhomogeneities!(f::AbstractVector, K::SparseMatrixCSC, ch::ConstraintHandler) = add_inhomogeneities_csc!(f, K, ch, false)
+add_inhomogeneities!(f::AbstractVector, K::Symmetric{<:Any, <:SparseMatrixCSC}, ch::ConstraintHandler) = add_inhomogeneities_csc!(f, K.data, ch, true)
+function add_inhomogeneities_csc!(f::AbstractVector, K::SparseMatrixCSC, ch::ConstraintHandler, sym::Bool)
+    (; inhomogeneities, prescribed_dofs, dofmapping) = ch
+
+    @inbounds for i in 1:length(inhomogeneities)
+        d = prescribed_dofs[i]
+        v = inhomogeneities[i]
+        if v != 0
+            for j in nzrange(K, d)
+                r = K.rowval[j]
+                sym && r > d && break # don't look below diagonal
+                f[r] -= v * K.nzval[j]
+            end
+        end
+    end
+    if sym
+        # In the symmetric case, for a constrained dof `d`, we handle the contribution
+        # from `K[1:d, d]` in the loop above, but we are still missing the contribution
+        # from `K[(d+1):size(K,1), d]`. These values are not stored, but since the
+        # matrix is symmetric we can instead use `K[d, (d+1):size(K,1)]`. Looping over
+        # rows is slow, so loop over all columns again, and check if the row is a
+        # constrained row.
+        @inbounds for col in 1:size(K, 2)
+            for ri in nzrange(K, col)
+                row = K.rowval[ri]
+                row >= col && break
+                if (i = get(dofmapping, row, 0); i != 0)
+                    f[col] -= inhomogeneities[i] * K.nzval[ri]
+                end
+            end
+        end
+    end
+    return f
+end
+
 # Fetch dof coefficients for a dof prescribed by an affine constraint. Return nothing if the
 # dof is not prescribed, or prescribed by DBC.
 @inline function coefficients_for_dof(dofmapping, dofcoeffs, dof)
@@ -751,6 +769,17 @@
     return dofcoeffs[idx]
 end
 
+"""
+    _condense!(K::AbstractSparseMatrix, f::AbstractVector, dofcoefficients::Vector{Union{Nothing, DofCoefficients{T}}}, dofmapping::Dict{Int, Int}, sym::Bool = false)
+
+Condenses affine constraints K := C'*K*C and f := C'*f in-place, assuming the sparsity pattern is correct.
+"""
+function _condense!(K::AbstractSparseMatrix, f::AbstractVector, dofcoefficients::Vector{Union{Nothing, DofCoefficients{T}}}, dofmapping::Dict{Int, Int}, sym::Bool = false) where {T}
+    # Return early if there are no non-trivial affine constraints
+    any(i -> !(i === nothing || isempty(i)), dofcoefficients) || return
+    error("condensation of ::$(typeof(K)) matrix not supported")
+end
+
 # Condenses K and f: C'*K*C, C'*f, in-place assuming the sparsity pattern is correct
 function _condense!(K::SparseMatrixCSC, f::AbstractVector, dofcoefficients::Vector{Union{Nothing, DofCoefficients{T}}}, dofmapping::Dict{Int, Int}, sym::Bool = false) where {T}
 
@@ -861,22 +890,33 @@
 
     return C, g
 end
+"""
+    zero_out_columns!(K::AbstractMatrix, ch::ConstraintHandler)
+Set the values of all columns associated with constrained dofs to zero.
+"""
+zero_out_columns!
+
+"""
+    zero_out_rows!(K::AbstractMatrix, ch::ConstraintHandler)
+Set the values of all rows associated with constrained dofs to zero.
+"""
+zero_out_rows!
 
 # columns need to be stored entries, this is not checked
-function zero_out_columns!(K, dofs::Vector{Int}) # can be removed in 0.7 with #24711 merged
-    @debug @assert issorted(dofs)
-    for col in dofs
+function zero_out_columns!(K::AbstractSparseMatrixCSC, ch::ConstraintHandler) # can be removed in 0.7 with #24711 merged
+    @debug @assert issorted(ch.prescribed_dofs)
+    for col in ch.prescribed_dofs
         r = nzrange(K, col)
         K.nzval[r] .= 0.0
     end
     return
 end
 
-function zero_out_rows!(K, dofmapping)
+function zero_out_rows!(K::AbstractSparseMatrixCSC, ch::ConstraintHandler)
     rowval = K.rowval
     nzval = K.nzval
     @inbounds for i in eachindex(rowval, nzval)
-        if haskey(dofmapping, rowval[i])
+        if haskey(ch.dofmapping, rowval[i])
             nzval[i] = 0
         end
     end
@@ -1749,9 +1789,16 @@
     return
 end
 
-# Condensation of affine constraints on element level. If possible this function only
-# modifies the local arrays.
 @noinline missing_global() = error("can not condense constraint without the global matrix and vector")
+
+"""
+    _condense_local!(local_matrix::AbstractMatrix, local_vector::AbstractVector,
+                    global_matrix#=::SparseMatrixCSC=#, global_vector#=::Vector=#,
+                    global_dofs::AbstractVector, dofmapping::Dict, dofcoefficients::Vector)
+
+Condensation of affine constraints on element level. If possible this function only
+modifies the local arrays.
+"""
 function _condense_local!(
         local_matrix::AbstractMatrix, local_vector::AbstractVector,
         global_matrix #=::SparseMatrixCSC=#, global_vector #=::Vector=#,