Merge pull request #171 from Cthonios/dofmanager/condensed-bcs-initial-stab

cmhamel · web-flow · commit 2673552be39c · 2025-09-12T20:20:46.000-04:00
initial stab at condensed bcs (maybe we need a better name) where the…
diff --git a/ext/FiniteElementContainersAdaptExt.jl b/ext/FiniteElementContainersAdaptExt.jl
@@ -86,11 +86,13 @@ function Adapt.adapt_structure(to, bc::FiniteElementContainers.NeumannBCContaine
 end
 
 # DofManagers
-function Adapt.adapt_structure(to, dof::DofManager)
+function Adapt.adapt_structure(to, dof::DofManager{C, IT, IDs, Var}) where {C, IT, IDs, Var}
   dirichlet_dofs = adapt(to, dof.dirichlet_dofs)
   unknowns = adapt(to, dof.unknown_dofs)
   var = adapt(to, dof.var)
-  return DofManager(dirichlet_dofs, unknowns, var)
+  return DofManager{
+    C, IT, typeof(dirichlet_dofs), typeof(var) 
+  }(dirichlet_dofs, unknowns, var)
 end
 
 # Fields
diff --git a/src/DofManagers.jl b/src/DofManagers.jl
@@ -4,6 +4,8 @@ abstract type AbstractDofManager{
 } end
 
 struct DofManager{
+    Condensed, # boolean flag for whether or not to seperate data between unknowns and constrained dofs
+    # when creating unknowns
     IT, 
     IDs <: AbstractArray{IT, 1},
     Var <: AbstractFunction
@@ -15,13 +17,19 @@ end
 
 # method that initializes dof manager
 # with all dofs unknown
-function DofManager(var::AbstractFunction)
+function DofManager(var::AbstractFunction; use_condensed::Bool = false)
     dirichlet_dofs = zeros(Int, 0)
     unknown_dofs = 1:size(var.fspace.coords, 2) * length(names(var)) |> collect
-    return DofManager(dirichlet_dofs, unknown_dofs, var)
+    return DofManager{
+        use_condensed,
+        eltype(dirichlet_dofs),
+        typeof(dirichlet_dofs),
+        typeof(var)
+    }(dirichlet_dofs, unknown_dofs, var)
 end
 
 _field_type(dof::DofManager) = eval(typeof(dof.var.fspace.coords).name.name)
+_is_condensed(dof::DofManager{C, IT, IDs, V}) where {C, IT, IDs, V} = C
 
 Base.length(dof::DofManager) = length(dof.dirichlet_dofs) + length(dof.unknown_dofs)
 
@@ -57,29 +65,69 @@ function create_field(dof::DofManager)
     return _field_type(dof)(field)
 end
 
-function create_unknowns(dof::DofManager)
+function create_unknowns(dof::DofManager{false, IT, IDs, Var}) where {IT, IDs, Var}
     backend = KA.get_backend(dof)
     return KA.zeros(backend, Float64, length(dof.unknown_dofs))
 end
 
+function create_unknowns(dof::DofManager{true, IT, IDs, Var}) where {IT, IDs, Var}
+    backend = KA.get_backend(dof)
+    return KA.zeros(backend, Float64, length(dof))
+end
+
 # COV_EXCL_START
-KA.@kernel function _extract_field_unknowns_kernel!(Uu::V, dof::DofManager, U::AbstractField) where V <: AbstractVector{<:Number}
+KA.@kernel function _extract_field_unknowns_kernel!(
+    Uu::V, 
+    dof::DofManager{false, IT, IDs, Var}, 
+    U::AbstractField
+) where {V <: AbstractVector{<:Number}, IT, IDs, Var}
     N = KA.@index(Global)
     @inbounds Uu[N] = U[dof.unknown_dofs[N]]
 end
 # COV_EXCL_STOP
 
-function _extract_field_unknowns!(Uu::V, dof::DofManager, U::AbstractField, backend::KA.Backend) where V <: AbstractVector{<:Number}
+# COV_EXCL_START
+KA.@kernel function _extract_field_unknowns_kernel!(
+    Uu::V, 
+    dof::DofManager{true, IT, IDs, Var}, 
+    U::AbstractField
+) where {V <: AbstractVector{<:Number}, IT, IDs, Var}
+    N = KA.@index(Global)
+    @inbounds Uu[dof.unknown_dofs[N]] = U[dof.unknown_dofs[N]]
+end
+# COV_EXCL_STOP
+
+function _extract_field_unknowns!(
+    Uu::V, 
+    dof::DofManager, 
+    U::AbstractField, 
+    backend::KA.Backend
+) where V <: AbstractVector{<:Number}
     kernel! = _extract_field_unknowns_kernel!(backend)
     kernel!(Uu, dof, U, ndrange = length(Uu))
     return nothing
 end
 
-function _extract_field_unknowns!(Uu::V, dof::DofManager, U::AbstractField, ::KA.CPU) where V <: AbstractVector{<:Number}
+function _extract_field_unknowns!(
+    Uu::V, 
+    dof::DofManager{false, IT, IDs, Var}, 
+    U::AbstractField, 
+    ::KA.CPU
+) where {V <: AbstractVector{<:Number}, IT, IDs, Var}
     @views Uu .= U[dof.unknown_dofs]
     return nothing
 end
 
+function _extract_field_unknowns!(
+    Uu::V, 
+    dof::DofManager{true, IT, IDs, Var}, 
+    U::AbstractField, 
+    ::KA.CPU
+) where {V <: AbstractVector{<:Number}, IT, IDs, Var}
+    @views Uu[dof.unknown_dofs] .= U[dof.unknown_dofs]
+    return nothing
+end
+
 function extract_field_unknowns!(
     Uu::V,
     dof::DofManager,
@@ -96,7 +144,7 @@ function function_space(dof::DofManager)
     return dof.var.fspace
 end
 
-function update_dofs!(dof::DofManager, dirichlet_dofs)
+function update_dofs!(dof::DofManager, dirichlet_dofs::V) where V <: AbstractArray{<:Integer, 1}
     ND, NI = size(dof)
     Base.resize!(dof.dirichlet_dofs, length(dirichlet_dofs))
     Base.resize!(dof.unknown_dofs, ND * NI)
@@ -114,24 +162,59 @@ function update_dofs!(dof::DofManager, dirichlet_dofs)
 end
 
 # COV_EXCL_START
-KA.@kernel function _update_field_unknowns_kernel!(U::AbstractField, dof::DofManager, Uu::V) where V <: AbstractVector{<:Number}
+KA.@kernel function _update_field_unknowns_kernel!(
+    U::AbstractField, 
+    dof::DofManager{false, IT, IDs, Var}, 
+    Uu::V
+) where {V <: AbstractVector{<:Number}, IT, IDs, Var}
     N = KA.@index(Global)
     @inbounds U[dof.unknown_dofs[N]] = Uu[N]
 end
 # COV_EXCL_STOP
+
+# COV_EXCL_START
+KA.@kernel function _update_field_unknowns_kernel!(
+    U::AbstractField, 
+    dof::DofManager{true, IT, IDs, Var}, 
+    Uu::V
+) where {V <: AbstractVector{<:Number}, IT, IDs, Var}
+    N = KA.@index(Global)
+    @inbounds U[dof.unknown_dofs[N]] = Uu[dof.unknown_dofs[N]]
+end
+# COV_EXCL_STOP
   
-function _update_field_unknowns!(U::AbstractField, dof::DofManager, Uu::T, backend::KA.Backend) where T <: AbstractVector{<:Number}
+function _update_field_unknowns!(
+    U::AbstractField, 
+    dof::DofManager, 
+    Uu::T, 
+    backend::KA.Backend
+) where T <: AbstractVector{<:Number}
     kernel! = _update_field_unknowns_kernel!(backend)
     kernel!(U, dof, Uu, ndrange = length(Uu))
     return nothing
 end
   
 # Need a seperate CPU method since CPU is basically busted in KA
-function _update_field_unknowns!(U::AbstractField, dof::DofManager, Uu::T, ::KA.CPU) where T <: AbstractVector{<:Number}
+function _update_field_unknowns!(
+    U::AbstractField, 
+    dof::DofManager{false, IT, IDs, Var}, 
+    Uu::T, 
+    ::KA.CPU
+) where {T <: AbstractVector{<:Number}, IT, IDs, Var}
     U[dof.unknown_dofs] .= Uu
     return nothing
 end
 
+function _update_field_unknowns!(
+    U::AbstractField, 
+    dof::DofManager{true, IT, IDs, Var}, 
+    Uu::T, 
+    ::KA.CPU
+) where {T <: AbstractVector{<:Number}, IT, IDs, Var}
+    @views U[dof.unknown_dofs] .= Uu[dof.unknown_dofs]
+    return nothing
+end
+
 function update_field_unknowns!(
     U::AbstractField, 
     dof::DofManager,
@@ -140,6 +223,7 @@ function update_field_unknowns!(
     backend = KA.get_backend(dof)
     @assert KA.get_backend(U) == backend
     @assert KA.get_backend(Uu) == backend
+    # @assert length(dof.unknown_dofs) == length(Uu)
     _update_field_unknowns!(U, dof, Uu, backend)
     return nothing
 end
diff --git a/src/assemblers/SparseMatrixAssembler.jl b/src/assemblers/SparseMatrixAssembler.jl
@@ -79,13 +79,8 @@ function SparseMatrixAssembler(dof::DofManager)
   )
 end
 
-# function SparseMatrixAssembler(::Type{<:H1Field}, vars...)
-#   dof = DofManager(vars...)
-#   return SparseMatrixAssembler(dof, H1Field)
-# end
-
-function SparseMatrixAssembler(var::AbstractFunction)
-  dof = DofManager(var)
+function SparseMatrixAssembler(var::AbstractFunction; use_condensed::Bool = false)
+  dof = DofManager(var; use_condensed=use_condensed)
   return SparseMatrixAssembler(dof)
 end
 
@@ -119,15 +114,27 @@ end
 $(TYPEDSIGNATURES)
 TODO add symbol to interface
 """
-function SparseArrays.sparse!(assembler::SparseMatrixAssembler, ::Val{:stiffness})
+function SparseArrays.sparse!(
+  assembler::SparseMatrixAssembler, ::Val{:stiffness};
+  unknowns_only=true
+)
   pattern = assembler.pattern
   storage = assembler.stiffness_storage
-  return @views SparseArrays.sparse!(
-    pattern.Is, pattern.Js, storage[assembler.pattern.unknown_dofs],
-    length(pattern.klasttouch), length(pattern.klasttouch), +, pattern.klasttouch,
-    pattern.csrrowptr, pattern.csrcolval, pattern.csrnzval,
-    pattern.csccolptr, pattern.cscrowval, pattern.cscnzval
-  )
+  if unknowns_only
+    return @views SparseArrays.sparse!(
+      pattern.Is, pattern.Js, storage[assembler.pattern.unknown_dofs],
+      length(pattern.klasttouch), length(pattern.klasttouch), +, pattern.klasttouch,
+      pattern.csrrowptr, pattern.csrcolval, pattern.csrnzval,
+      pattern.csccolptr, pattern.cscrowval, pattern.cscnzval
+    )
+  # else
+  #   return SparseArrays.sparse!(
+  #     pattern.Is, pattern.Js, storage[assembler.pattern.unknown_dofs],
+  #     length(pattern.klasttouch), length(pattern.klasttouch), +, pattern.klasttouch,
+  #     pattern.csrrowptr, pattern.csrcolval, pattern.csrnzval,
+  #     pattern.csccolptr, pattern.cscrowval, pattern.cscnzval
+  #   )
+  end
 end
 
 """
@@ -170,15 +177,9 @@ end
 # the residual and stiffness appropriately without having to reshape, Is, Js, etc.
 # when we want to change BCs which is slow
 
-function update_dofs!(assembler::SparseMatrixAssembler, dirichlet_bcs; use_condensed=false)
-  # vars = assembler.dof.H1_vars
-  var = assembler.dof.var
+function update_dofs!(assembler::SparseMatrixAssembler, dirichlet_bcs)
+  use_condensed = _is_condensed(assembler.dof)
 
-  # if length(vars) != 1
-  #   @assert false "multiple fspace not supported yet"
-  # end
-
-  # dirichlet_dofs = dirichlet_bcs.bookkeeping.dofs
   if length(dirichlet_bcs) > 0
     dirichlet_dofs = mapreduce(x -> x.bookkeeping.dofs, vcat, dirichlet_bcs)
     dirichlet_dofs = unique(sort(dirichlet_dofs))
@@ -203,8 +204,8 @@ end
 # when we want to change BCs which is slow
 
 function _update_dofs_condensed!(assembler::SparseMatrixAssembler)
-  assembler.constraint_storage[assembler.dof.unknown_dofs] .= 1.
-  assembler.constraint_storage[assembler.dof.dirichlet_bcs] .= 0.
+  assembler.constraint_storage[assembler.dof.unknown_dofs] .= 0.
+  assembler.constraint_storage[assembler.dof.dirichlet_dofs] .= 1.
   return nothing
 end
 
@@ -214,55 +215,10 @@ end
 function _update_dofs!(assembler::SparseMatrixAssembler, dirichlet_dofs::T) where T <: AbstractArray{<:Integer, 1}
 
   # resize the resiual unkowns
-  # n_total_H1_dofs = num_nodes(assembler.dof) * num_dofs_per_node(assembler.dof)
-  # n_unknown_dofs = length(assembler.dof.unknown_dofs)
-  # resize!(assembler.residual_unknowns, length(assembler.dof.H1_unknown_dofs))
-  # resize!(assembler.stiffness_action_unknowns, length(assembler.dof.H1_unknown_dofs))
   resize!(assembler.residual_unknowns, length(assembler.dof.unknown_dofs))
   resize!(assembler.stiffness_action_unknowns, length(assembler.dof.unknown_dofs))
 
-  # n_total_dofs = length(assembler.dof) - length(dirichlet_dofs)
-  n_total_dofs = length(assembler.dof) - length(dirichlet_dofs)
-  # n_total_dofs = n_unknown_dofs - length(dirichlet_dofs)
-
-  # TODO change to a good sizehint!
-  resize!(assembler.pattern.Is, 0)
-  resize!(assembler.pattern.Js, 0)
-  resize!(assembler.pattern.unknown_dofs, 0)
-
-  # ND, NN = num_dofs_per_node(assembler.dof), num_nodes(assembler.dof)
-  ND, NN = size(assembler.dof)
-  # ids = reshape(1:length(assembler.dof), ND, NN)
-  ids = reshape(1:length(assembler.dof), ND, NN)
-
-  # TODO
-  # vars = assembler.dof.H1_vars
-  # fspace = vars[1].fspace
-  fspace = function_space(assembler.dof)
-
-  n = 1
-  for conns in values(fspace.elem_conns)
-    dof_conns = @views reshape(ids[:, conns], ND * size(conns, 1), size(conns, 2))
-
-    for e in 1:size(conns, 2)
-      conn = @views dof_conns[:, e]
-      for temp in Iterators.product(conn, conn)
-        if insorted(temp[1], dirichlet_dofs) || insorted(temp[2], dirichlet_dofs)
-          # really do nothing here
-        else
-          push!(assembler.pattern.Is, temp[1] - count(x -> x < temp[1], dirichlet_dofs))
-          push!(assembler.pattern.Js, temp[2] - count(x -> x < temp[2], dirichlet_dofs))
-          push!(assembler.pattern.unknown_dofs, n)
-        end
-        n += 1
-      end
-    end
-  end
-
-  resize!(assembler.pattern.klasttouch, n_total_dofs)
-  resize!(assembler.pattern.csrrowptr, n_total_dofs + 1)
-  resize!(assembler.pattern.csrcolval, length(assembler.pattern.Is))
-  resize!(assembler.pattern.csrnzval, length(assembler.pattern.Is))
+  _update_dofs!(assembler.pattern, assembler.dof, dirichlet_dofs)
 
   return nothing
 end
diff --git a/src/assemblers/SparsityPattern.jl b/src/assemblers/SparsityPattern.jl
@@ -109,3 +109,51 @@ function SparsityPattern(dof::DofManager)
 end
 
 num_entries(s::SparsityPattern) = length(s.Is)
+
+# NOTE this methods assumes that dof is up to date
+# NOTE this method also only resizes unknown_dofs
+# in the pattern object, that means that things
+# like Is, Js, etc. need to be viewed into or sliced
+function _update_dofs!(pattern::SparsityPattern, dof, dirichlet_dofs)
+  n_total_dofs = length(dof) - length(dirichlet_dofs)
+
+  # remove me
+  resize!(pattern.Is, 0)
+  resize!(pattern.Js, 0)
+  # end remove me
+  resize!(pattern.unknown_dofs, 0)
+  ND, NN = size(dof)
+  ids = reshape(1:length(dof), ND, NN)
+  fspace = function_space(dof)
+
+  n = 1
+  for conns in values(fspace.elem_conns)
+    dof_conns = @views reshape(ids[:, conns], ND * size(conns, 1), size(conns, 2))
+
+    for e in 1:size(conns, 2)
+      conn = @views dof_conns[:, e]
+      for temp in Iterators.product(conn, conn)
+        if insorted(temp[1], dirichlet_dofs) || insorted(temp[2], dirichlet_dofs)
+          # really do nothing here
+        else
+          # remove me
+          push!(pattern.Is, temp[1] - count(x -> x < temp[1], dirichlet_dofs))
+          push!(pattern.Js, temp[2] - count(x -> x < temp[2], dirichlet_dofs))
+          # end remove me
+          push!(pattern.unknown_dofs, n)
+        end
+        n += 1
+      end
+    end
+  end
+
+  resize!(pattern.klasttouch, n_total_dofs)
+  resize!(pattern.csrrowptr, n_total_dofs + 1)
+  # TODO Not sure about below 2 sizes
+  # resize!(assembler.pattern.csrcolval, length(assembler.pattern.Is))
+  # resize!(assembler.pattern.csrnzval, length(assembler.pattern.Is))
+  resize!(pattern.csrcolval, length(pattern.Is))
+  resize!(pattern.csrnzval, length(pattern.Is))
+
+  return nothing
+end
diff --git a/test/poisson/TestPoissonCondensedBCs.jl b/test/poisson/TestPoissonCondensedBCs.jl
diff --git a/test/runtests.jl b/test/runtests.jl
