Introduce isconstrained bitvector in ConstraintHandler (#1310)

* Introduce isconstrained bitvector in ch to speed up zero_out_rows

Also used to make close simpler and slightly faster, and make it easier to check if dof is constrained

Author: KnutAM

ext/FerriteSparseMatrixCSR.jl

@@ -67,10 +67,11 @@ function Ferrite.zero_out_rows!(K::SparseMatrixCSR, ch::ConstraintHandler)
 end
 
 function Ferrite.zero_out_columns!(K::SparseMatrixCSR, ch::ConstraintHandler)
+    @boundscheck checkbounds(ch.isconstrained, axes(K, 2))
     colval = K.colval
     nzval = K.nzval
-    return @inbounds for i in eachindex(colval, nzval)
-        if haskey(ch.dofmapping, colval[i])
+    return @inbounds for (i, col) in pairs(colval)
+        if ch.isconstrained[col]
             nzval[i] = 0
         end
     end

src/Dofs/ConstraintHandler.jl

@@ -166,6 +166,7 @@ mutable struct ConstraintHandler{DH <: AbstractDofHandler, T}
     const dofcoefficients::Vector{Union{Nothing, DofCoefficients{T}}}
     # global dof -> index into dofs and inhomogeneities and dofcoefficients
     const dofmapping::Dict{Int, Int}
+    const isconstrained::BitVector # Fast check if dof is constrained or not
     const bcvalues::Vector{BCValues{T}}
     const dh::DH
     closed::Bool
@@ -177,7 +178,7 @@ function ConstraintHandler(::Type{T}, dh::AbstractDofHandler) where {T <: Number
     @assert isclosed(dh)
     return ConstraintHandler(
         Dirichlet[], ProjectedDirichlet[], Int[], Int[], T[], Union{Nothing, T}[],
-        Union{Nothing, DofCoefficients{T}}[], Dict{Int, Int}(), BCValues{T}[], dh, false,
+        Union{Nothing, DofCoefficients{T}}[], Dict{Int, Int}(), BitVector(), BCValues{T}[], dh, false,
     )
 end
 
@@ -259,21 +260,15 @@ isclosed(ch::ConstraintHandler) = ch.closed
 free_dofs(ch::ConstraintHandler) = ch.free_dofs
 prescribed_dofs(ch::ConstraintHandler) = ch.prescribed_dofs
 
-# Equivalent to `copy!(out, setdiff(1:n_entries, diff))`, but requires that
-# `issorted(diff)` and that `1 ≤ diff[1] ≤ diff[end] ≤ n_entries`
-function _sorted_setdiff!(out::Vector{Int}, n_entries::Int, diff::Vector{Int})
-    n_diff = length(diff)
-    resize!(out, n_entries - n_diff)
-    diff_ind = out_ind = 1
-    for i in 1:n_entries
-        if diff_ind ≤ n_diff && i == diff[diff_ind]
-            diff_ind += 1
-        else
-            out[out_ind] = i
-            out_ind += 1
-        end
+function _set_freedofs!(free_dofs::Vector{Int}, isconstrained::BitVector, num_dofs::Int, num_prescribed::Int)
+    resize!(free_dofs, num_dofs - num_prescribed)
+    j = 1
+    for i in 1:num_dofs
+        isconstrained[i] && continue
+        free_dofs[j] = i
+        j += 1
     end
-    return out
+    return free_dofs
 end
 
 """
@@ -291,7 +286,12 @@ function close!(ch::ConstraintHandler)
     ch.affine_inhomogeneities .= ch.affine_inhomogeneities[I]
     ch.dofcoefficients .= ch.dofcoefficients[I]
 
-    _sorted_setdiff!(ch.free_dofs, ndofs(ch.dh), ch.prescribed_dofs)
+    resize!(ch.isconstrained, ndofs(ch.dh))
+    fill!(ch.isconstrained, false)
+    for dof in ch.prescribed_dofs
+        ch.isconstrained[dof] = true
+    end
+    _set_freedofs!(ch.free_dofs, ch.isconstrained, ndofs(ch.dh), length(ch.prescribed_dofs))
 
     for i in 1:length(ch.prescribed_dofs)
         ch.dofmapping[ch.prescribed_dofs[i]] = i
@@ -913,11 +913,12 @@ function zero_out_columns!(K::AbstractSparseMatrixCSC, ch::ConstraintHandler) #
 end
 
 function zero_out_rows!(K::AbstractSparseMatrixCSC, ch::ConstraintHandler)
+    @boundscheck checkbounds(ch.isconstrained, axes(K, 1))
     rowval = K.rowval
     nzval = K.nzval
-    @inbounds for i in eachindex(rowval, nzval)
-        if haskey(ch.dofmapping, rowval[i])
-            nzval[i] = 0
+    @inbounds for (i, row) in pairs(rowval)
+        if ch.isconstrained[row]
+            nzval[i] = zero(eltype(K))
         end
     end
     return
@@ -1765,7 +1766,7 @@ function _apply_local!(
     # 3. Condense any affine constraints
     if has_nontrivial_affine_constraints
         # Condense this constraint locally if possible, and otherwise modifies the global arrays.
-        _condense_local!(local_matrix, local_vector, global_matrix, global_vector, global_dofs, ch.dofmapping, ch.dofcoefficients)
+        _condense_local!(local_matrix, local_vector, global_matrix, global_vector, global_dofs, ch.dofmapping, ch.dofcoefficients, ch.isconstrained)
     end
     # 4. Zero out columns/rows of local matrix and replace diagonal entries with the mean
     if has_constraints
@@ -1794,15 +1795,17 @@ end
 """
     _condense_local!(local_matrix::AbstractMatrix, local_vector::AbstractVector,
                     global_matrix#=::SparseMatrixCSC=#, global_vector#=::Vector=#,
-                    global_dofs::AbstractVector, dofmapping::Dict, dofcoefficients::Vector)
+                    global_dofs::AbstractVector, dofmapping::Dict, dofcoefficients::Vector,
+                    isconstrained::BitVector)
 
 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=#,
-        global_dofs::AbstractVector, dofmapping::Dict, dofcoefficients::Vector
+        global_dofs::AbstractVector, dofmapping::Dict, dofcoefficients::Vector,
+        isconstrained::BitVector,
     )
     @assert axes(local_matrix, 1) == axes(local_matrix, 2) ==
         axes(local_vector, 1) == axes(global_dofs, 1)
@@ -1821,7 +1824,7 @@ function _condense_local!(
                     if local_mrow === nothing
                         # Only modify the global array if this isn't prescribed since we
                         # can't zero it out later like with the local matrix.
-                        if !haskey(dofmapping, global_col) && !haskey(dofmapping, global_mrow)
+                        if !isconstrained[global_col] && !isconstrained[global_mrow]
                             has_global_arrays || missing_global()
                             addindex!(global_matrix, mw, global_mrow, global_col)
                         end