Add assemble_lumped_mass! primitive

src/FiniteElementContainers.jl

@@ -9,6 +9,7 @@ export to_backend
 # Assemblers
 export SparseMatrixAssembler
 export create_assembler_cache
+export assemble_lumped_mass!
 export assemble_mass!
 export assemble_matrix!
 export assemble_matrix_action!
@@ -138,6 +139,7 @@ export assemble_diagonal!
 export diagonal
 export energy
 export hvp
+export lumped_mass
 export mass
 export mass!
 export mass_action

src/assemblers/Assemblers.jl

@@ -473,6 +473,7 @@ include("SparseMatrixAssembler.jl")
 
 # methods
 include("Diagonal.jl")
+include("LumpedMass.jl")
 include("Matrix.jl")
 include("MatrixAction.jl")
 include("QuadratureQuantity.jl")

src/assemblers/LumpedMass.jl

@@ -0,0 +1,80 @@
+"""
+$(TYPEDSIGNATURES)
+Assemble a partition-of-unity (row-sum) lumped mass vector.
+
+`func` is a physics-level kernel with the same signature as `residual`
+that returns an `SVector{NDOF, T}` of element-local lumped-mass
+contributions per DOF.  For a partition-of-unity basis the contribution
+is `density * N[a] * JxW`, identical in each spatial direction at node `a`.
+
+The result is scattered into `assembler.residual_storage` (reused as
+scratch — overwrites any previously-assembled residual) and retrieved
+via `lumped_mass(assembler)`.
+
+Distinct from:
+  * `assemble_mass!`   — full consistent sparse mass matrix `M_ab = ρ ∫ N_a N_b dV`.
+  * `assemble_diagonal!(asm, mass, ...)` — diagonal of the consistent
+    matrix `M_aa = ρ ∫ N_a^2 dV`.  Not equal to the row-sum lumped mass
+    in general (they differ by a factor depending on the element type).
+  * `assemble_matrix_action!(asm, mass, U_zeros, ones_free, p)` —
+    `M_red * 1_free`, which under-counts contributions from columns
+    corresponding to constrained DOFs, breaking partition of unity near
+    Dirichlet boundaries.
+
+The row-sum lumped mass is the correct quantity for explicit central
+difference dynamics (`a = M^{-1} f`) and for mass-diagonal Jacobi
+preconditioning, because it preserves partition of unity (total mass =
+density × volume) by construction.
+"""
+function assemble_lumped_mass!(
+  assembler, func::F, Uu, p
+) where F <: Function
+  storage = assembler.residual_storage
+  fill!(storage, zero(eltype(storage)))
+  dof = assembler.dof
+  fspace = function_space(dof)
+  X = coordinates(p)
+  t = current_time(p)
+  Δt = time_step(p)
+  U = p.field
+  U_old = p.field_old
+  _update_for_assembly!(p, dof, Uu)
+  return_type = AssembledVector()
+  conns = fspace.elem_conns
+  foreach_block(fspace, p) do physics, props, ref_fe, b
+    _assemble_block!(
+      storage,
+      conns.data, conns.offsets[b],
+      func,
+      physics, ref_fe,
+      X, t, Δt,
+      U, U_old,
+      block_view(p.state_old, b), block_view(p.state_new, b), props,
+      return_type
+    )
+  end
+  return nothing
+end
+
+"""
+$(TYPEDSIGNATURES)
+Return the assembled lumped mass vector after a call to
+`assemble_lumped_mass!`.  Extracts the free-DOF subset (non-condensed
+mode) or returns the full vector (condensed mode).
+
+Note: shares backing storage with `residual(asm)`.  Callers that need
+both quantities must `copy` this result before any subsequent
+`assemble_vector!` / `assemble_lumped_mass!` / `assemble_diagonal!` call.
+"""
+function lumped_mass(asm::AbstractAssembler)
+  if _is_condensed(asm.dof)
+    return asm.residual_storage.data
+  else
+    extract_field_unknowns!(
+      asm.residual_unknowns,
+      asm.dof,
+      asm.residual_storage
+    )
+    return asm.residual_unknowns
+  end
+end

test/TestAssemblers.jl

@@ -336,6 +336,97 @@ end
   end
 end
 
+@testitem "Assemblers - test_lumped_mass_mechanics" setup=[AssemblerHelperMechanics] begin
+  # The contract on assemble_lumped_mass!:
+  #   (a) On a fully-free DOF space, it equals the row sums of the
+  #       full consistent mass matrix (partition of unity).
+  #   (b) On a BC-restricted DOF space, it equals (a) restricted to
+  #       the free DOFs — i.e., the per-element scatter is independent
+  #       of which DOFs are subsequently extracted.
+  #   (c) It differs from the legacy "M_red * 1_free" approach at
+  #       free DOFs adjacent to constrained ones, because the latter
+  #       drops contributions from the columns corresponding to
+  #       constrained DOFs (this is the bug being fixed).
+  include("TestUtils.jl")
+  backends = _get_backends()
+  for dev in backends
+    # ---- (a) Partition-of-unity on a fully-free DOF space ----
+    asm_full = SparseMatrixAssembler(
+      u;
+      sparse_matrix_type = :csc,
+      use_condensed = false,
+      use_inplace_methods = false,
+    )
+    p_full = create_parameters(mesh, asm_full, physics, props;
+                                dirichlet_bcs = DirichletBC[], times = times)
+    FiniteElementContainers.initialize!(p_full)
+    if dev != cpu
+      asm_full = asm_full |> dev
+      p_full   = p_full |> dev
+    end
+    Uu_full = create_unknowns(asm_full)
+
+    assemble_mass!(asm_full, mass, Uu_full, p_full)
+    M_full_host = mass(asm_full) |> copy
+    if dev != cpu
+      M_full_host = M_full_host |> cpu
+    end
+    m_ref_full = vec(sum(M_full_host; dims = 2))
+
+    assemble_lumped_mass!(asm_full, lumped_mass, Uu_full, p_full)
+    m_lumped_full = lumped_mass(asm_full) |> copy
+    if dev != cpu
+      m_lumped_full = m_lumped_full |> cpu
+    end
+
+    @test length(m_lumped_full) == length(m_ref_full)
+    @test all(isapprox.(m_lumped_full, m_ref_full, rtol = 1e-12, atol = 1e-14))
+
+    # ---- (b) BC-restricted lumped mass equals (a) at free DOFs ----
+    asm_bc = SparseMatrixAssembler(
+      u;
+      sparse_matrix_type = :csc,
+      use_condensed = false,
+      use_inplace_methods = false,
+    )
+    p_bc = create_parameters(mesh, asm_bc, physics, props;
+                              dirichlet_bcs = dbcs, times = times)
+    FiniteElementContainers.initialize!(p_bc)
+    if dev != cpu
+      asm_bc = asm_bc |> dev
+      p_bc   = p_bc |> dev
+    end
+    Uu_bc = create_unknowns(asm_bc)
+
+    assemble_lumped_mass!(asm_bc, lumped_mass, Uu_bc, p_bc)
+    m_lumped_bc = lumped_mass(asm_bc) |> copy
+    if dev != cpu
+      m_lumped_bc = m_lumped_bc |> cpu
+    end
+
+    unk_host = asm_bc.dof.unknown_dofs |> copy
+    if dev != cpu
+      unk_host = unk_host |> cpu
+    end
+
+    @test length(m_lumped_bc) == length(unk_host)
+    @test all(isapprox.(m_lumped_bc, m_ref_full[unk_host], rtol = 1e-12, atol = 1e-14))
+
+    # ---- (c) Regression: differs from M_red * 1_free at the boundary ----
+    assemble_mass!(asm_bc, mass, Uu_bc, p_bc)
+    M_red_host = mass(asm_bc) |> copy
+    if dev != cpu
+      M_red_host = M_red_host |> cpu
+    end
+    m_buggy = vec(sum(M_red_host; dims = 2))
+    # They must agree at interior DOFs (not adjacent to any boundary node)
+    # and disagree at boundary-adjacent free DOFs.  The strong test is
+    # simply that some entries disagree — for the standard test mesh,
+    # several entries do.
+    @test !all(isapprox.(m_lumped_bc, m_buggy, rtol = 1e-10, atol = 1e-14))
+  end
+end
+
 @testitem "Assemblers - test_enzyme_safe_consistency" setup=[AssemblerHelperPoisson] begin
   asm = SparseMatrixAssembler(u)
   p = create_parameters(mesh, asm, physics, props; dirichlet_bcs = dbcs)

test/mechanics/TestMechanicsCommon.jl

@@ -86,7 +86,7 @@ end
 
 @inline function FiniteElementContainers.mass!(
   storage, e,
-  physics::Mechanics, t, dt, props_el, 
+  physics::Mechanics, t, dt, props_el,
   state_old_q, state_new_q,
   conn, interps, x_el, u_el, u_el_old
 )
@@ -95,6 +95,36 @@ end
   scatter_with_values_and_values!(storage, physics.formulation, e, conn, N, JxW * props_el[1])
 end
 
+@inline function FiniteElementContainers.lumped_mass(
+  physics::Mechanics,
+  interps, x_el,
+  t, dt,
+  u_el, u_el_old,
+  state_old_q, state_new_q,
+  props_el,
+)
+  cell = map_interpolants(interps, x_el)
+  (; N, JxW) = cell
+
+  # Row sum of the element consistent mass at node a, direction d:
+  #   m_el[NDIM*(a-1)+d] = ρ * JxW * Σ_b N[a] * N[b]
+  #                     = ρ * JxW * N[a]      (partition of unity: Σ_b N[b] = 1)
+  # Identical in every spatial direction.  Interleaved DOF ordering to match
+  # the rest of the assembly pipeline.
+  N_nodes = size(N, 1)
+  NDIM    = num_fields(physics.formulation)
+  NDOF    = NDIM * N_nodes
+  ρJxW    = JxW * props_el[1]
+  tup = zeros(SVector{NDOF, eltype(N)})
+  for a in 1:N_nodes
+    m_a = ρJxW * N[a]
+    for d in 1:NDIM
+      tup = setindex(tup, m_a, NDIM * (a - 1) + d)
+    end
+  end
+  return tup
+end
+
 @inline function FiniteElementContainers.mass_action!(
   storage, e,
   physics::Mechanics, t, dt, props_el,