Modularize assembly internals and add SparseMatrixCSR extension

Project.toml

@@ -18,15 +18,18 @@ 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"
 Metis = "1.3"
+SparseMatricesCSR = "0.6"
 NearestNeighbors = "0.4"
 OrderedCollections = "1"
 Preferences = "1"
@@ -49,8 +52,9 @@ Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SHA = "ea8e919c-243c-51af-8825-aaa63cd721ce"
+SparseMatricesCSR = "a0a7dd2c-ebf4-11e9-1f05-cf50bc540ca1"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 
 [targets]
-test = ["BlockArrays", "Downloads", "FerriteGmsh", "ForwardDiff", "Gmsh", "IterativeSolvers", "Metis", "Pkg", "NBInclude", "ProgressMeter", "Random", "SHA", "Test", "TimerOutputs", "Logging"]
+test = ["BlockArrays", "Downloads", "FerriteGmsh", "ForwardDiff", "Gmsh", "IterativeSolvers", "Metis", "Pkg", "NBInclude", "ProgressMeter", "Random", "SHA", "SparseMatricesCSR", "Test", "TimerOutputs", "Logging"]

docs/src/devdocs/assembler.md

@@ -0,0 +1,37 @@
+# [Assembler](@id devdocs-assembler)
+
+The assembler handles the insertion of the element matrices and element vectors into the system matrix and right hand side. 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 assembler [`AssemblerSparsityPattern`](@ref) should be able to handle most cases in practice. To support a custom format users have to dispatch the following functions:
+
+```@docs
+Ferrite._assemble_inner!
+Ferrite.zero_out_rows!
+Ferrite.zero_out_columns!
+```
+
+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
+Ferrite.matrix_handle
+Ferrite.vector_handle
+start_assemble!
+finish_assemble!
+assemble!
+```
+
+For local elimination support the following functions might also need custom dispatches
+
+```@docs
+Ferrite._condense_local!
+```
+
+## Custom Matrix Type Sparsity Pattern
+
+TODO `create_sparsity_pattern`

docs/src/devdocs/index.md

@@ -5,5 +5,5 @@ developing the library.
 
 ```@contents
 Depth = 1
-Pages = ["reference_cells.md", "interpolations.md", "elements.md", "FEValues.md", "dofhandler.md", "performance.md"]
+Pages = ["reference_cells.md", "interpolations.md", "elements.md", "FEValues.md", "dofhandler.md", "assembler.md", "performance.md"]
 ```

ext/FerriteSparseMatrixCSR.jl

@@ -0,0 +1,79 @@
+module FerriteSparseMatrixCSR
+
+using Ferrite, SparseArrays, SparseMatricesCSR
+import Ferrite: AbstractSparsityPattern
+import Base: @propagate_inbounds
+
+@propagate_inbounds function Ferrite._assemble_inner!(K::SparseMatrixCSR, Ke::AbstractMatrix, dofs::AbstractVector, sorteddofs::AbstractVector, permutation::AbstractVector, sym::Bool)
+    current_row = 1
+    ld = length(dofs)
+    @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) # can be removed in 0.7 with #24711 merged
+    @debug @assert issorted(ch.prescribed_dofs)
+    for row in ch.prescribed_dofs
+        r = nzrange(K, row)
+        K.nzval[r] .= 0.0
+    end
+end
+
+function Ferrite.zero_out_columns!(K::SparseMatrixCSR, ch::ConstraintHandler)
+    colval = K.colval
+    nzval = K.nzval
+    @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 allocate_matrix(SparseMatrixCSR{1,Float64,Int}, sp)
+end
+
+function Ferrite.allocate_matrix(::Type{MatrixType}, sp::AbstractSparsityPattern) where {Bi, Tv, Ti, MatrixType <: SparseMatrixCSR{Bi, Tv, Ti}}
+    # Allocate CSC first ...
+    K = allocate_matrix(SparseMatrixCSC{Tv, Ti}, sp)
+    # ... and transform to SparseMatrixCSR
+    return SparseMatrixCSR{Bi}(transpose(sparse(transpose(K))))
+end
+
+end

src/Dofs/ConstraintHandler.jl

@@ -495,7 +495,7 @@ end
 
 """
 
-    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.
@@ -587,11 +587,11 @@ function _apply_v(v::AbstractVector, ch::ConstraintHandler, apply_zero::Bool)
     return v
 end
 
-function apply!(K::Union{SparseMatrixCSC,Symmetric}, ch::ConstraintHandler)
+function apply!(K::Union{AbstractSparseMatrix,Symmetric}, ch::ConstraintHandler)
     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)
     apply!(K, f, ch, true)
 end
 
@@ -600,7 +600,7 @@ end
 const APPLY_TRANSPOSE = ApplyStrategy.Transpose
 const APPLY_INPLACE = ApplyStrategy.Inplace
 
-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;
                 strategy::ApplyStrategy.T=ApplyStrategy.Transpose)
     @assert isclosed(ch)
     sym = isa(KK, Symmetric)
@@ -612,43 +612,14 @@ function apply!(KK::Union{SparseMatrixCSC,Symmetric}, f::AbstractVector, ch::Con
     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!(KK, f, ch.inhomogeneities, ch.prescribed_dofs, ch.dofmapping)
 
-    # 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)
@@ -661,6 +632,50 @@ function apply!(KK::Union{SparseMatrixCSC,Symmetric}, f::AbstractVector, ch::Con
     end
 end
 
+"""
+    add_inhomogeneities!(K, f::AbstractVector, inhomogeneities::AbstractVector, prescribed_dofs::AbstractVector{<:Integer}, dofmapping::Dict{Int,Int})
+
+Compute "f -= K*inhomogeneities".
+By default this is a generic version via SpMSpV kernel.
+"""
+function add_inhomogeneities!(K, f::AbstractVector, inhomogeneities::AbstractVector, prescribed_dofs::AbstractVector{<:Integer}, dofmapping)
+    f .-= K*sparsevec(prescribed_dofs, inhomogeneities, size(K,2))
+end
+
+# Optimized version for SparseMatrixCSC
+add_inhomogeneities!(K::SparseMatrixCSC, f::AbstractVector, inhomogeneities::AbstractVector, prescribed_dofs::AbstractVector{<:Integer}, dofmapping) = add_inhomogeneities_csc!(K, f, inhomogeneities, prescribed_dofs, dofmapping, false)
+add_inhomogeneities!(K::Symmetric{<:Any,<:SparseMatrixCSC}, f::AbstractVector, inhomogeneities::AbstractVector, prescribed_dofs::AbstractVector{<:Integer}, dofmapping) = add_inhomogeneities_csc!(K.data, f, inhomogeneities, prescribed_dofs, dofmapping, true)
+function add_inhomogeneities_csc!(K::SparseMatrixCSC, f::AbstractVector, inhomogeneities::AbstractVector, prescribed_dofs::AbstractVector{<:Integer}, dofmapping, sym::Bool)
+    @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
+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)
@@ -669,6 +684,13 @@ end
     return dofcoeffs[idx]
 end
 
+
+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
 
@@ -779,19 +801,19 @@ function create_constraint_matrix(ch::ConstraintHandler{dh,T}) where {dh,T}
 end
 
 # 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::SparseArrays.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
 end
 
-function zero_out_rows!(K, dofmapping)
+function zero_out_rows!(K::SparseArrays.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

src/Ferrite.jl

@@ -14,7 +14,7 @@ using NearestNeighbors:
 using OrderedCollections:
     OrderedSet
 using SparseArrays:
-    SparseArrays, SparseMatrixCSC, nonzeros, nzrange, rowvals, sparse
+    AbstractSparseMatrix, SparseArrays, SparseMatrixCSC, nonzeros, nzrange, rowvals, sparse, sparsevec
 using StaticArrays:
     StaticArrays, MArray, MMatrix, SArray, SMatrix, SVector
 using WriteVTK:

src/arrayutils.jl

@@ -77,11 +77,11 @@ function addindex!(A::SparseMatrixCSC{Tv}, v::Tv, i::Int, j::Int) where Tv
     end
 end
 
-function fillzero!(A::SparseMatrixCSC{T}) where T
+function fillzero!(A::AbstractSparseMatrix{T}) where T
     fill!(nonzeros(A), zero(T))
     return A
 end
-function fillzero!(A::Symmetric{T,<:SparseMatrixCSC}) where T
+function fillzero!(A::Symmetric{T,<:AbstractSparseMatrix}) where T
     fillzero!(A.data)
     return A
 end