feat(assemblers): Add node-based COO assembly implementation

New file: src/assemblers/node_based_coo.jl Features: - assemble! implementation for nodal assembly - Assembles contributions node-by-node instead of element-by-element - Uses node_to_elements connectivity - Accumulates blocks for all elements touching each node Architecture: - Outer loop over nodes (not elements) - Inner loop over elements containing each node - Natural for contact mechanics (contact is nodal) Status: Experimental, proof-of-concept implementation. Not yet optimized like element-based assembly.
2026-09-12 22:33:19 +00:00 · 2025-11-20 16:56:39 +02:00
parent 9c62148525
commit 08f940f3db
1 changed files with 484 additions and 0 deletions
@@ -0,0 +1,484 @@
+# This file is a part of JuliaFEM.
+# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md
+
+"""
+Node-based COO assembly using block integration.
+
+**NODAL ASSEMBLY PARADIGM**: Loop over nodes, not elements!
+
+Each node:
+1. Finds all elements touching it (via inverse connectivity)
+2. For each touching element:
+   - Prepares element geometry once (PreparedElement)
+   - Computes only needed 3×3 blocks (compute_block!)
+3. Scatters blocks to COO triplets
+
+# Key Differences from Element-Based Assembly
+
+**Element-Based (traditional):**
+```julia
+for element in elements
+    K_e = compute_element_stiffness(element)  # Full N×N matrix of 3×3 blocks
+    scatter(K_e)                               # Scatter all entries
+end
+```
+
+**Node-Based (this file):**
+```julia
+for node_i in nodes
+    for element in elements_touching(node_i)
+        prepared = prepare_element(element)    # Geometry preprocessing
+        for node_j in element.nodes
+            K_ij = compute_block!(prepared, i, j)  # Single 3×3 block
+            scatter(K_ij, i, j)                    # Scatter one block
+        end
+    end
+end
+```
+
+# Advantages
+
+1. **GPU-friendly**: One thread per node, no race conditions
+2. **Contact-ready**: Contact is naturally node-based
+3. **Matrix-free ready**: Can compute K*v without forming K
+4. **Cache-friendly**: Reuses PreparedElement for multiple blocks
+5. **Adaptive-ready**: Easy to refine/coarsen at node level
+
+# Performance Expectations
+
+- **CPU Single-thread**: ~1.5-2x slower than element-based (more kernel calls)
+- **CPU Multi-thread**: ~1.5-2x faster (better parallelization)
+- **GPU**: ~10-50x faster (massive parallelization, no atomics needed)
+
+# References
+
+- Golden standard: `docs/src/book/multigpu_nodal_assembly.md`
+- PreparedElement: `src/domains/continuum/integration.jl`
+- Block kernel: `src/domains/continuum/kernel.jl`
+
+# Example
+
+```julia
+# Setup
+mesh = create_cantilever_mesh(50, 10, 10)
+material = LinearElastic(E=210e9, ν=0.3)
+kernel = ContinuumKernel(
+    ContinuumFormulation{FullThreeD}(),
+    material,
+    Displacement{3}()
+)
+
+# Create node-based assembler and cache
+assembler = NodeBasedCOOAssembler()
+cache = create_cache(assembler, mesh, kernel)
+
+# Assemble (zero allocations after warmup!)
+assemble!(cache, assembler, kernel, mesh)
+
+# Extract system
+K, f = extract_system(cache)
+
+# Solve
+apply_dirichlet_bcs!(K, f, kernel, mesh, bc_dirichlet)
+u = K \\ f
+```
+"""
+
+using SparseArrays
+using Tensors
+
+"""
+    NodeBasedCOOCache
+
+Pre-allocated cache for node-based COO assembly.
+
+Similar to COOCache but includes inverse connectivity mapping.
+
+# Fields
+- `I::Vector{Int}`: Row indices (COO format)
+- `J::Vector{Int}`: Column indices (COO format)
+- `V::Vector{Float64}`: Values (COO format)
+- `f::Vector{Float64}`: Global force vector
+- `counter::Ref{Int}`: Current triplet count
+- `capacity::Int`: Maximum triplet capacity
+- `node_to_elements::NodeToElementsMap`: Inverse connectivity
+- `element_cache::ElementCache`: Cache for element operations
+- `ndofs::Int`: Total DOFs in system
+"""
+struct NodeBasedCOOCache{T<:AbstractTopology,B<:AbstractBasis,IPS}
+    I::Vector{Int}
+    J::Vector{Int}
+    V::Vector{Float64}
+    f::Vector{Float64}
+    counter::Ref{Int}
+    capacity::Int
+    node_to_elements::NodeToElementsMap
+    element_cache::ElementCache{T,B,IPS}
+    ndofs::Int
+end
+
+"""
+    NodeBasedCOOCache(mesh::AbstractMesh, kernel::ContinuumKernel)
+
+Create cache for node-based assembly.
+
+Builds inverse connectivity and allocates buffers.
+
+# Arguments
+- `mesh`: Finite element mesh
+- `kernel`: Continuum kernel
+
+# Returns
+- Pre-allocated node-based COO cache
+"""
+function NodeBasedCOOCache(mesh::AbstractMesh, kernel::ContinuumKernel)
+    # Build inverse connectivity
+    node_to_elements = NodeToElementsMap(mesh.connectivity)
+
+    # Estimate triplet count (same as element-based)
+    ndofs_per_node = dofs_per_node(kernel)
+    nnodes = length(mesh.nodes)
+    ndofs = ndofs_per_node * nnodes
+
+    # Estimate: For each node, sum over touching elements
+    # Each element contributes N blocks (N = nodes per element)
+    # Each block = 3×3 = 9 triplets
+    avg_elements_per_node = node_to_elements.nelements / nnodes
+    N = length(first(mesh.connectivity))  # Nodes per element
+    estimated_triplets = Int(ceil(1.2 * nnodes * avg_elements_per_node * N * 9))
+
+    # Allocate triplet arrays
+    I = Vector{Int}(undef, estimated_triplets)
+    J = Vector{Int}(undef, estimated_triplets)
+    V = Vector{Float64}(undef, estimated_triplets)
+    f = zeros(Float64, ndofs)
+    counter = Ref(0)
+
+    # Create element cache (for prepare_element! and compute_block!)
+    element_cache = ElementCache(mesh, kernel)
+
+    return NodeBasedCOOCache(I, J, V, f, counter, estimated_triplets,
+        node_to_elements, element_cache, ndofs)
+end
+
+"""
+    reset!(cache::NodeBasedCOOCache)
+
+Reset cache for new assembly (zero force vector, reset counter).
+
+Does NOT clear inverse connectivity (that's permanent structure).
+"""
+function reset!(cache::NodeBasedCOOCache)
+    fill!(cache.f, 0.0)
+    cache.counter[] = 0
+    return nothing
+end
+
+"""
+    assemble!(
+        cache::NodeBasedCOOCache,
+        assembler::NodeBasedCOOAssembler,
+        kernel::ContinuumKernel,
+        mesh::AbstractMesh
+    ) -> Nothing
+
+Assemble global system using **node-based traversal**.
+
+# Algorithm
+
+```julia
+for node_i in 1:nnodes
+    # Get all elements touching this node
+    for elem_info in node_to_elements[node_i]
+        element_id = elem_info.element_id
+        local_i = elem_info.local_node_idx
+        
+        # Prepare element geometry ONCE
+        prepared = prepare_element!(cache.element_cache, kernel, element_id, mesh)
+        
+        # Compute blocks for all nodes in this element
+        for local_j in 1:N
+            global_j = connectivity[element_id][local_j]
+            
+            # Compute single 3×3 block
+            K_ij = compute_block!(prepared, kernel.material, local_i, local_j)
+            
+            # Scatter to triplets
+            scatter_block_to_triplets!(cache, K_ij, node_i, global_j)
+        end
+    end
+end
+```
+
+# Key Operations
+
+1. **prepare_element!** - Precompute geometry (Jacobian, gradients) once per element
+2. **compute_block!** - Compute single 3×3 stiffness block using prepared geometry
+3. **scatter_block_to_triplets!** - Add 9 triplets (i,j,value) for 3×3 block
+
+# Zero-Allocation (After Warmup)
+
+All arrays pre-allocated. Element preparation reuses cache buffers.
+
+# Arguments
+- `cache`: Pre-allocated node-based COO cache
+- `assembler`: Node-based COO assembler
+- `kernel`: Continuum kernel
+- `mesh`: Finite element mesh
+"""
+function assemble!(
+    cache::NodeBasedCOOCache,
+    assembler::NodeBasedCOOAssembler,
+    kernel::ContinuumKernel,
+    mesh::AbstractMesh
+)
+    # Reset cache
+    reset!(cache)
+
+    nnodes = length(mesh.nodes)
+    ndofs_per_node = dofs_per_node(kernel)
+
+    # NODAL LOOP: One iteration per node (GPU: one thread per node!)
+    for node_i in 1:nnodes
+        # Get all elements touching this node
+        touching_elements = cache.node_to_elements.node_to_elements[node_i]
+
+        # Loop over touching elements
+        for elem_info in touching_elements
+            element_id = elem_info.element_id
+            local_i = elem_info.local_node_idx  # Position of node_i in element
+
+            # Prepare element geometry ONCE (reuses cache.element_cache)
+            prepared = prepare_element!(cache.element_cache, kernel, element_id, mesh)
+
+            # Get element connectivity
+            conn = mesh.connectivity[element_id]
+            N = length(conn)  # Nodes per element
+
+            # Compute blocks for all nodes j in this element
+            for local_j in 1:N
+                global_j = conn[local_j]
+
+                # Compute single 3×3 block K[i,j]
+                # This is THE KEY OPERATION: block-based integration
+                K_ij = compute_block!(
+                    prepared,
+                    kernel.material,
+                    local_i,
+                    local_j
+                )
+
+                # Scatter 3×3 block to triplets (adds 9 entries)
+                scatter_block_to_triplets!(
+                    cache,
+                    K_ij,
+                    node_i,
+                    global_j,
+                    ndofs_per_node
+                )
+            end
+        end
+    end
+
+    return nothing
+end
+
+"""
+    scatter_block_to_triplets!(
+        cache::NodeBasedCOOCache,
+        K_block::Tensor{2,3},
+        node_i::Int,
+        node_j::Int,
+        ndofs_per_node::Int
+    )
+
+Scatter single 3×3 block to COO triplets **in-place**.
+
+Maps block[α,β] → triplet at DOF indices:
+- Row: 3*(node_i-1) + α
+- Col: 3*(node_j-1) + β
+- Val: K_block[α,β]
+
+# Arguments
+- `cache`: Node-based COO cache
+- `K_block`: 3×3 stiffness block (Tensor{2,3})
+- `node_i`: Global row node index
+- `node_j`: Global column node index
+- `ndofs_per_node`: DOFs per node (typically 3)
+
+# Zero-Allocation
+
+Writes to pre-allocated triplet arrays, updates counter.
+"""
+function scatter_block_to_triplets!(
+    cache::NodeBasedCOOCache,
+    K_block::Tensor{2,3,Float64},
+    node_i::Int,
+    node_j::Int,
+    ndofs_per_node::Int
+)
+    counter = cache.counter[]
+
+    # Check capacity
+    new_triplets = ndofs_per_node * ndofs_per_node  # 3×3 = 9
+    if counter + new_triplets > cache.capacity
+        error("Node-based COO cache overflow: need $(counter + new_triplets) triplets, " *
+              "capacity is $(cache.capacity). Increase cache size.")
+    end
+
+    # DOF offsets for nodes i and j
+    row_offset = ndofs_per_node * (node_i - 1)
+    col_offset = ndofs_per_node * (node_j - 1)
+
+    # Scatter 3×3 block to triplets
+    for β in 1:ndofs_per_node  # Column (node j DOF)
+        j_global = col_offset + β
+        for α in 1:ndofs_per_node  # Row (node i DOF)
+            i_global = row_offset + α
+            counter += 1
+            cache.I[counter] = i_global
+            cache.J[counter] = j_global
+            cache.V[counter] = K_block[α, β]
+        end
+    end
+
+    cache.counter[] = counter
+    return nothing
+end
+
+"""
+    extract_system(cache::NodeBasedCOOCache) -> (K, f)
+
+Build sparse matrix from triplets and return system.
+
+Calls `sparse(I, J, V)` to build CSC matrix. Duplicates are summed automatically.
+
+# Arguments
+- `cache`: Assembled node-based COO cache
+
+# Returns
+- `K::SparseMatrixCSC`: Global stiffness matrix
+- `f::Vector`: Global force vector
+
+# Allocation
+
+Allocates sparse matrix structure (CSC format). This is the only allocation
+outside cache construction.
+"""
+function extract_system(cache::NodeBasedCOOCache)
+    ntriplets = cache.counter[]
+
+    # Build sparse matrix (duplicates are summed automatically)
+    I_used = @view cache.I[1:ntriplets]
+    J_used = @view cache.J[1:ntriplets]
+    V_used = @view cache.V[1:ntriplets]
+
+    K = sparse(I_used, J_used, V_used, cache.ndofs, cache.ndofs)
+
+    return K, cache.f
+end
+
+# ============================================================================
+# HELPER FUNCTIONS
+# ============================================================================
+
+"""
+    create_cache(
+        assembler::NodeBasedCOOAssembler,
+        mesh::AbstractMesh,
+        kernel::ContinuumKernel
+    ) -> NodeBasedCOOCache
+
+Create pre-allocated cache for node-based COO assembly.
+
+Convenience function that wraps `NodeBasedCOOCache(mesh, kernel)`.
+
+# Example
+
+```julia
+assembler = NodeBasedCOOAssembler()
+cache = create_cache(assembler, mesh, kernel)
+assemble!(cache, assembler, kernel, mesh)
+K, f = extract_system(cache)
+```
+"""
+function create_cache(
+    assembler::NodeBasedCOOAssembler,
+    mesh::AbstractMesh,
+    kernel::ContinuumKernel
+)
+    return NodeBasedCOOCache(mesh, kernel)
+end
+
+"""
+    dofs_per_node(kernel::ContinuumKernel) -> Int
+
+Return DOFs per node for continuum kernel (always 3 for displacement).
+
+Dispatches on kernel field dimension.
+"""
+function dofs_per_node(kernel::ContinuumKernel{Theory,Mat}) where {Theory,Mat}
+    field = kernel.field
+    return field.dim  # Displacement{3} → 3
+end
+
+# ============================================================================
+# PERFORMANCE NOTES
+# ============================================================================
+
+#=
+# CPU Performance Comparison (Estimated)
+
+**Element-Based Assembly:**
+- Elements: 1000 Tet4
+- Operations: 1000 elements × 4×4 blocks × 3×3 entries = 48,000 block computations
+- Time: ~5ms (baseline)
+
+**Node-Based Assembly:**
+- Nodes: 500 nodes
+- Operations: 500 nodes × 8 elements/node × 4 blocks/element = 16,000 block computations
+- But: 3× more kernel calls due to overlaps
+- Time: ~7-10ms (1.5-2× slower single-threaded)
+
+**Why slower on CPU?**
+- Each block computed once in element assembly
+- Each block computed 2× on average in nodal assembly (shared between 2 elements)
+- More function call overhead
+
+**Why faster on GPU?**
+- Element assembly: Sequential (can't parallelize over elements efficiently)
+- Nodal assembly: Massive parallelism (one thread per node)
+- GPU speedup: ~10-50× depending on problem size
+
+**Multi-threaded CPU (Threads.@threads):**
+- Can parallelize outer node loop
+- Expected speedup: 1.5-2× over element-based
+- No race conditions (each node writes different triplets)
+
+# Memory Comparison
+
+**Element-Based:**
+- Triplet storage: ~50 KB per 1000 elements
+- Element cache: ~2 KB per thread
+
+**Node-Based:**
+- Triplet storage: Same (~50 KB)
+- Element cache: ~2 KB per thread
+- Inverse connectivity: ~10-20 KB (one-time)
+
+→ Nearly identical memory usage!
+
+# When to Use Node-Based Assembly
+
+**Use when:**
+- ✅ GPU acceleration needed
+- ✅ Contact mechanics (naturally nodal)
+- ✅ Matrix-free methods (K*v without forming K)
+- ✅ Adaptive refinement (local node operations)
+- ✅ Multi-threading on CPU
+
+**Don't use when:**
+- ❌ Single-threaded CPU only
+- ❌ Simple problems (< 1000 nodes)
+- ❌ Prototyping/debugging (element-based is clearer)
+=#