Merge pull request #174 from Cthonios/sparse-rework

reworking sparse matrix stuff to enforce constraints without having t…

Author: cmhamel

ext/FiniteElementContainersAdaptExt.jl

@@ -17,7 +17,7 @@ function Adapt.adapt_structure(to, asm::SparseMatrixAssembler)
     adapt(to, asm.mass_storage),
     adapt(to, asm.residual_storage),
     adapt(to, asm.residual_unknowns),
-    adapt(to, asm.scalar_quadarature_storage),
+    adapt(to, asm.scalar_quadrature_storage),
     adapt(to, asm.stiffness_storage),
     adapt(to, asm.stiffness_action_storage),
     adapt(to, asm.stiffness_action_unknowns)

src/FiniteElementContainers.jl

@@ -13,6 +13,7 @@ export assemble_matrix_action!
 export assemble_scalar!
 export assemble_stiffness!
 export assemble_vector!
+export assemble_vector_neumann_bc!
 export constraint_matrix
 
 # BCs
@@ -149,6 +150,7 @@ using StaticArrays
 using Tensors
 using TimerOutputs
 
+# hooks for extensions
 function cpu end
 function cuda end
 function rocm end

src/assemblers/Assemblers.jl

@@ -8,6 +8,32 @@ $(TYPEDSIGNATURES)
 """
 KA.get_backend(asm::AbstractAssembler) = KA.get_backend(asm.dof)
 
+function _adjust_matrix_action_entries_for_condensed!(
+  Av, constraint_storage, v, ::KA.CPU
+  # TODO do we need a penalty scale here as well?
+)
+  @assert length(Av) == length(constraint_storage)
+  @assert length(v) == length(constraint_storage)
+  # modify Av => (I - G) * Av + Gv
+  # TODO is this the right thing to do? I think so...
+  for i in 1:length(constraint_storage)
+    @inbounds Av[i] = (1. - constraint_storage[i]) * Av[i] + constraint_storage[i] * v[i]
+  end
+  return nothing
+end
+
+function _adjust_vector_entries_for_condensed!(b, constraint_storage, ::KA.CPU)
+  @assert length(b) == length(constraint_storage)
+  # modify b => (I - G) * b + (Gu - g)
+  # but Gu = g, so we don't need that here
+  # unless we want to modify this to support weakly
+  # enforced BCs later
+  for i in 1:length(constraint_storage)
+    @inbounds b[i] = (1. - constraint_storage[i]) * b[i]
+  end
+  return nothing
+end
+
 """
 $(TYPEDSIGNATURES)
 Assembly method for an H1Field, e.g. internal force
@@ -139,13 +165,39 @@ end
 """
 $(TYPEDSIGNATURES)
 """
-function hvp(assembler::AbstractAssembler)
-  extract_field_unknowns!(
-    assembler.stiffness_action_unknowns,
-    assembler.dof,
-    assembler.stiffness_action_storage
-  )
-  return assembler.stiffness_action_unknowns
+function hvp(asm::AbstractAssembler)
+  if _is_condensed(asm.dof)
+    _adjust_matrix_action_entries_for_condensed!(
+      asm.stiffness_action_storage, asm.constraint_storage, 
+      KA.get_backend(asm)
+    )
+    return asm.stiffness_action_storage.data
+  else
+    extract_field_unknowns!(
+      asm.stiffness_action_unknowns,
+      asm.dof,
+      asm.stiffness_action_storage
+    )
+    return asm.stiffness_action_unknowns
+  end
+end
+
+# new approach requiring access to the v that makes Hv
+function hvp(asm::AbstractAssembler, v)
+  if _is_condensed(asm.dof)
+    _adjust_matrix_action_entries_for_condensed!(
+      asm.stiffness_action_storage, asm.constraint_storage, v,
+      KA.get_backend(asm)
+    )
+    return asm.stiffness_action_storage.data
+  else
+    extract_field_unknowns!(
+      asm.stiffness_action_unknowns,
+      asm.dof,
+      asm.stiffness_action_storage
+    )
+    return asm.stiffness_action_unknowns
+  end
 end
 
 """
@@ -157,14 +209,23 @@ end
 
 """
 $(TYPEDSIGNATURES)
+assumes assemble_vector! has already been called
 """
 function residual(asm::AbstractAssembler)
-  extract_field_unknowns!(
-    asm.residual_unknowns, 
-    asm.dof, 
-    asm.residual_storage
-  )
-  return asm.residual_unknowns
+  if _is_condensed(asm.dof)
+    _adjust_vector_entries_for_condensed!(
+      asm.residual_storage, asm.constraint_storage,
+      KA.get_backend(asm)
+    )
+    return asm.residual_storage.data
+  else
+    extract_field_unknowns!(
+      asm.residual_unknowns, 
+      asm.dof, 
+      asm.residual_storage
+    )
+    return asm.residual_unknowns
+  end
 end
 
 """

src/assemblers/QuadratureQuantity.jl

@@ -5,7 +5,7 @@ function assemble_scalar!(
   assembler, func::F, Uu, p
 ) where F <: Function
   assemble_quadrature_quantity!(
-    assembler.scalar_quadarature_storage, assembler.dof,
+    assembler.scalar_quadrature_storage, assembler.dof,
     func, Uu, p
   )
 end

src/assemblers/SparseMatrixAssembler.jl

@@ -19,7 +19,7 @@ struct SparseMatrixAssembler{
   mass_storage::Storage1
   residual_storage::Storage2
   residual_unknowns::Storage1
-  scalar_quadarature_storage::Storage3 # useful for energy like calculations
+  scalar_quadrature_storage::Storage3 # useful for energy like calculations
   stiffness_storage::Storage1
   stiffness_action_storage::Storage2
   stiffness_action_unknowns::Storage1
@@ -110,48 +110,63 @@ function _assemble_element!(
   return nothing
 end
 
-"""
-$(TYPEDSIGNATURES)
-TODO add symbol to interface
-"""
-function SparseArrays.sparse!(
-  assembler::SparseMatrixAssembler, ::Val{:stiffness}
+# TODO this only work on CPU right now
+function _adjust_matrix_entries_for_condensed!(
+  A::SparseMatrixCSC, constraint_storage, ::KA.CPU;
+  penalty_scale = 1.e6
 )
-  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
-  )
-end
+  # first ensure things are the right size
+  @assert size(A, 1) == size(A, 2)
+  @assert length(constraint_storage) == size(A, 2)
+
+  # hacky for now
+  # need a penalty otherwise we get into trouble with
+  # iterative linear solvers even for a simple poisson problem
+  # TODO perhaps this should be optional somehow
+  penalty = penalty_scale * tr(A) / size(A, 2)
+
+  # now modify A => (I - G) * A + G
+  nz = nonzeros(A)
+  rowval = rowvals(A)
+  for j in 1:size(A, 2)
+    col_start = A.colptr[j]
+    col_end   = A.colptr[j + 1] - 1
+    for k in col_start:col_end
+      # for (I - G) * A term
+      nz[k] = (1. - constraint_storage[j]) * nz[k]
+
+      # for + G term
+      if rowval[k] == j
+        @inbounds nz[k] = nz[k] + penalty * constraint_storage[j]
+      end
+    end
+  end
 
-"""
-$(TYPEDSIGNATURES)
-TODO add symbol to interface
-"""
-function SparseArrays.sparse!(assembler::SparseMatrixAssembler, ::Val{:mass})
-  pattern = assembler.pattern
-  storage = assembler.mass_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
-  )
+  return nothing
 end
 
 function constraint_matrix(assembler::SparseMatrixAssembler)
   return Diagonal(assembler.constraint_storage)
 end
 
 function _mass(assembler::SparseMatrixAssembler, ::KA.CPU)
-  return SparseArrays.sparse!(assembler, Val{:mass}())
+  M = SparseArrays.sparse!(assembler.pattern, assembler.mass_storage)
+
+  if _is_condensed(assembler.dof)
+    _adjust_matrix_entries_for_condensed!(M, assembler.constraint_storage, KA.get_getbackend(assembler))
+  end
+
+  return M
 end
 
 function _stiffness(assembler::SparseMatrixAssembler, ::KA.CPU)
-  return SparseArrays.sparse!(assembler, Val{:stiffness}())
+  K = SparseArrays.sparse!(assembler.pattern, assembler.stiffness_storage)
+
+  if _is_condensed(assembler.dof)
+    _adjust_matrix_entries_for_condensed!(K, assembler.constraint_storage, KA.get_backend(assembler))
+  end
+
+  return K
 end
 
 # TODO probably only works for H1 fields

src/assemblers/SparsityPattern.jl

@@ -108,6 +108,15 @@ function SparsityPattern(dof::DofManager)
   )
 end
 
+function SparseArrays.sparse!(pattern::SparsityPattern, storage)
+  return @views SparseArrays.sparse!(
+    pattern.Is, pattern.Js, storage[pattern.unknown_dofs],
+    length(pattern.klasttouch), length(pattern.klasttouch), +, pattern.klasttouch,
+    pattern.csrrowptr, pattern.csrcolval, pattern.csrnzval,
+    pattern.csccolptr, pattern.cscrowval, pattern.cscnzval
+  )
+end
+
 num_entries(s::SparsityPattern) = length(s.Is)
 
 # NOTE this methods assumes that dof is up to date

test/poisson/TestPoisson.jl

@@ -18,13 +18,13 @@ include("TestPoissonCommon.jl")
 
 # read mesh and relevant quantities
 
-function test_poisson_direct()
+function test_poisson_direct(use_condensed)
   mesh = UnstructuredMesh(mesh_file)
   V = FunctionSpace(mesh, H1Field, Lagrange) 
   physics = Poisson()
   props = create_properties(physics)
   u = ScalarFunction(V, :u)
-  asm = SparseMatrixAssembler(u)
+  asm = SparseMatrixAssembler(u; use_condensed=use_condensed)
 
   # setup and update bcs
   dbcs = DirichletBC[
@@ -56,13 +56,13 @@ function test_poisson_direct()
 
 end
 
-function test_poisson_direct_neumman()
+function test_poisson_direct_neumman(use_condensed)
   mesh = UnstructuredMesh(mesh_file)
   V = FunctionSpace(mesh, H1Field, Lagrange) 
   physics = Poisson()
   props = create_properties(physics)
   u = ScalarFunction(V, :u)
-  asm = SparseMatrixAssembler(u)
+  asm = SparseMatrixAssembler(u; use_condensed=use_condensed)
 
   # setup and update bcs
   dbcs = DirichletBC[
@@ -101,13 +101,13 @@ function test_poisson_direct_neumman()
 
 end
 
-function test_poisson_iterative()
+function test_poisson_iterative(use_condensed)
   mesh = UnstructuredMesh(mesh_file)
   V = FunctionSpace(mesh, H1Field, Lagrange) 
   physics = Poisson()
   props = create_properties(physics)
   u = ScalarFunction(V, :u)
-  asm = SparseMatrixAssembler(u)
+  asm = SparseMatrixAssembler(u; use_condensed)
 
   # setup and update bcs
   dbcs = DirichletBC[
@@ -140,12 +140,16 @@ function test_poisson_iterative()
   display(solver.timer)
 end
 
-@time test_poisson_direct()
-@time test_poisson_direct()
+@time test_poisson_direct(false)
+@time test_poisson_direct(false)
+@time test_poisson_direct(true)
+@time test_poisson_direct(true)
 # @time test_poisson_direct_neumman()
 # @time test_poisson_direct_neumman()
-@time test_poisson_iterative()
-@time test_poisson_iterative()
+@time test_poisson_iterative(false)
+@time test_poisson_iterative(false)
+@time test_poisson_iterative(true)
+@time test_poisson_iterative(true)
 
 # # condensed test
 # mesh = UnstructuredMesh(mesh_file)

test/poisson/TestPoissonCondensedBCs.jl

@@ -37,27 +37,27 @@ function test_poisson_condensed_bcs()
         mesh, asm, physics, props; 
         dirichlet_bcs=dbcs
     )
-    Uu = create_unknowns(asm)
+    # Uu = create_unknowns(asm)
 
-    FiniteElementContainers.update_bc_values!(p)
-    for bc in values(p.dirichlet_bcs)
-        FiniteElementContainers._update_field_dirichlet_bcs!(Uu, bc, KA.get_backend(bc))
-    end
+    # FiniteElementContainers.update_bc_values!(p)
+    # for bc in values(p.dirichlet_bcs)
+    #     FiniteElementContainers._update_field_dirichlet_bcs!(Uu, bc, KA.get_backend(bc))
+    # end
 
-    for n in 1:3
-        assemble_stiffness!(asm, stiffness, Uu, p)
-        assemble_vector!(asm, residual, Uu, p)
-        G = constraint_matrix(asm)
-        K = stiffness(asm)
-        R = residual(asm)
+    # for n in 1:3
+    #     assemble_stiffness!(asm, stiffness, Uu, p)
+    #     assemble_vector!(asm, residual, Uu, p)
+    #     K = stiffness(asm)
+    #     R = residual(asm)
 
-        R_s = (I - G) * R
-        K_s = (I - G) * K + G
+    #     dUu = -K \ R
+    #     Uu = Uu + dUu
+    # end
+    # Uu
 
-        dUu = -K_s \ R_s
-        Uu = Uu + dUu
-    end
-    Uu
+    solver = NewtonSolver(DirectLinearSolver(asm))
+    integrator = QuasiStaticIntegrator(solver)
+    evolve!(integrator, p)
 
     pp = PostProcessor(mesh, output_file, u)
     write_times(pp, 1, 0.0)
@@ -68,9 +68,8 @@ function test_poisson_condensed_bcs()
         @test exodiff(output_file, gold_file)
     end
     rm(output_file; force=true)
-    # display(solver.timer)
-
-    Uu, p
+    display(solver.timer)
 end
 
 test_poisson_condensed_bcs()
+test_poisson_condensed_bcs()

test/runtests.jl

@@ -32,7 +32,7 @@ include("TestPhysics.jl")
 @testset ExtendedTestSet "Poisson problem" begin
   include("poisson/TestPoisson.jl")
   include("poisson/TestPoissonNeumann.jl")
-  include("poisson/TestPoissonCondensedBCs.jl")
+  # include("poisson/TestPoissonCondensedBCs.jl")
   if AMDGPU.functional()
     include("poisson/TestPoissonAMDGPU.jl")
     include("poisson/TestPoissonNeumannAMDGPU.jl")