diff --git a/demos/nodal_assembly_cpu.jl b/demos/nodal_assembly_cpu.jl new file mode 100644 index 0000000..8e95f89 --- /dev/null +++ b/demos/nodal_assembly_cpu.jl @@ -0,0 +1,387 @@ +""" +Nodal Assembly CPU Reference Implementation + +This implements the TWO-PHASE nodal assembly approach that will be ported to GPU: + +Phase 1: Compute integration point data (stresses, material states) +Phase 2: Nodal assembly (matrix-free, no atomics) + +Uses Tensors.jl throughout for natural tensor operations. +""" + +using Tensors +using LinearAlgebra +using Printf + +# Material state for plasticity +mutable struct PlasticState + ε_p::SymmetricTensor{2,3,Float64,6} # Plastic strain tensor + α::Float64 # Accumulated plastic strain +end + +# Material properties +struct Material + E::Float64 # Young's modulus + ν::Float64 # Poisson's ratio + σ_y::Float64 # Yield stress +end + +# Node-to-elements connectivity (CSR format) +struct NodeToElementsMap + ptr::Vector{Int} # Length: n_nodes + 1 + data::Vector{Int} # Length: total connections +end + +""" +Build CSR map: which elements touch each node? +""" +function build_node_to_elems(elements::Vector{NTuple{4,Int}}, n_nodes::Int) + # Count connections per node + counts = zeros(Int, n_nodes) + for elem in elements + for node in elem + counts[node] += 1 + end + end + + # Build CSR structure + ptr = cumsum([1; counts]) + data = Vector{Int}(undef, sum(counts)) + + # Fill data array + offset = copy(ptr[1:end-1]) + for (elem_idx, elem) in enumerate(elements) + for node in elem + data[offset[node]] = elem_idx + offset[node] += 1 + end + end + + return NodeToElementsMap(ptr, data) +end + +""" +Tet4 shape function derivatives in reference coordinates (constant!) +""" +function tet4_shape_derivatives() + return ( + Vec{3}((-1.0, -1.0, -1.0)), # dN1/dξ + Vec{3}((1.0, 0.0, 0.0)), # dN2/dξ + Vec{3}((0.0, 1.0, 0.0)), # dN3/dξ + Vec{3}((0.0, 0.0, 1.0)) # dN4/dξ + ) +end + +""" +Return mapping for von Mises perfect plasticity (using Tensors.jl!) +""" +function return_mapping_tensor(ε_total::SymmetricTensor{2,3,T}, + state_old::PlasticState, + mat::Material) where T + # Elastic strain + ε_e = ε_total - state_old.ε_p + + # Elastic predictor + λ = mat.E * mat.ν / ((1 + mat.ν) * (1 - 2mat.ν)) + μ = mat.E / (2(1 + mat.ν)) + I = one(ε_e) + σ_trial = λ * tr(ε_e) * I + 2μ * ε_e + + # Deviatoric stress + σ_dev = dev(σ_trial) + σ_eq = sqrt(3 / 2 * (σ_dev ⊡ σ_dev)) + + # Yield function + f = σ_eq - mat.σ_y + + if f <= 0.0 + # Elastic + return (σ_trial, state_old) + else + # Plastic - radial return + Δγ = f / (3μ) + n = σ_dev / σ_eq + + σ = σ_trial - 2μ * Δγ * n + + # Update plastic state + Δε_p = Δγ * n + ε_p_new = state_old.ε_p + Δε_p + α_new = state_old.α + Δγ + + state_new = PlasticState(ε_p_new, α_new) + + return (σ, state_new) + end +end + +""" +Phase 1: Compute integration point data (stresses and material states) + +One "thread" per integration point (in CPU version, just a loop) +""" +function compute_gp_data!( + σ_gp::Vector{SymmetricTensor{2,3,Float64,6}}, + states_new::Vector{PlasticState}, + u::Vector{Float64}, + nodes::Matrix{Float64}, # Shape: 3 × n_nodes + elements::Vector{NTuple{4,Int}}, + states_old::Vector{PlasticState}, + mat::Material +) + n_elems = length(elements) + n_gps_per_elem = 4 # 4 Gauss points for Tet4 + + dN_dxi = tet4_shape_derivatives() + + for elem_idx in 1:n_elems + # Extract element nodes + n1, n2, n3, n4 = elements[elem_idx] + + X1 = Vec{3}((nodes[1, n1], nodes[2, n1], nodes[3, n1])) + X2 = Vec{3}((nodes[1, n2], nodes[2, n2], nodes[3, n2])) + X3 = Vec{3}((nodes[1, n3], nodes[2, n3], nodes[3, n3])) + X4 = Vec{3}((nodes[1, n4], nodes[2, n4], nodes[3, n4])) + + u1 = Vec{3}((u[3*n1-2], u[3*n1-1], u[3*n1])) + u2 = Vec{3}((u[3*n2-2], u[3*n2-1], u[3*n2])) + u3 = Vec{3}((u[3*n3-2], u[3*n3-1], u[3*n3])) + u4 = Vec{3}((u[3*n4-2], u[3*n4-1], u[3*n4])) + + # Jacobian (using tensor products!) + J = dN_dxi[1] ⊗ X1 + dN_dxi[2] ⊗ X2 + dN_dxi[3] ⊗ X3 + dN_dxi[4] ⊗ X4 + invJ = inv(J) + + # Physical derivatives + dN1_dx = invJ ⋅ dN_dxi[1] + dN2_dx = invJ ⋅ dN_dxi[2] + dN3_dx = invJ ⋅ dN_dxi[3] + dN4_dx = invJ ⋅ dN_dxi[4] + + # Loop over Gauss points (for Tet4, same strain at all GPs since linear) + # In real code, would have different GP locations + for local_gp in 1:n_gps_per_elem + gp_idx = (elem_idx - 1) * n_gps_per_elem + local_gp + + # Strain (using tensor products!) + ε = symmetric(dN1_dx ⊗ u1 + dN2_dx ⊗ u2 + dN3_dx ⊗ u3 + dN4_dx ⊗ u4) + + # Material state update + state_old = states_old[gp_idx] + σ, state_new = return_mapping_tensor(ε, state_old, mat) + + # Store results + σ_gp[gp_idx] = σ + states_new[gp_idx] = state_new + end + end +end + +""" +Phase 2: Nodal assembly (matrix-free, no atomics!) + +One "thread" per node (in CPU version, just a loop) +""" +function nodal_assembly!( + r::Vector{Float64}, + σ_gp::Vector{SymmetricTensor{2,3,Float64,6}}, + nodes::Matrix{Float64}, + elements::Vector{NTuple{4,Int}}, + node_to_elems::NodeToElementsMap +) + n_nodes = size(nodes, 2) + n_gps_per_elem = 4 + + dN_dxi = tet4_shape_derivatives() + + # Gauss weights for Tet4 (standard 4-point quadrature) + gauss_weights = (1 / 24, 1 / 24, 1 / 24, 1 / 24) + + fill!(r, 0.0) + + for node_idx in 1:n_nodes + f_node = zero(Vec{3,Float64}) + + # Get elements touching this node (CSR traversal) + elem_start = node_to_elems.ptr[node_idx] + elem_end = node_to_elems.ptr[node_idx+1] - 1 + + # Loop over touching elements + for elem_offset in elem_start:elem_end + elem_idx = node_to_elems.data[elem_offset] + elem_nodes = elements[elem_idx] + + # Find local node index in element + local_node = findfirst(==(node_idx), elem_nodes) + @assert local_node !== nothing "Node not found in element!" + + # Recompute geometry (matrix-free approach!) + n1, n2, n3, n4 = elem_nodes + X1 = Vec{3}((nodes[1, n1], nodes[2, n1], nodes[3, n1])) + X2 = Vec{3}((nodes[1, n2], nodes[2, n2], nodes[3, n2])) + X3 = Vec{3}((nodes[1, n3], nodes[2, n3], nodes[3, n3])) + X4 = Vec{3}((nodes[1, n4], nodes[2, n4], nodes[3, n4])) + + J = dN_dxi[1] ⊗ X1 + dN_dxi[2] ⊗ X2 + dN_dxi[3] ⊗ X3 + dN_dxi[4] ⊗ X4 + detJ = det(J) + invJ = inv(J) + + # Physical derivative for this node + dN_dx = invJ ⋅ dN_dxi[local_node] + + # Loop over Gauss points + for local_gp in 1:n_gps_per_elem + gp_idx = (elem_idx - 1) * n_gps_per_elem + local_gp + + # Get stress at this GP + σ = σ_gp[gp_idx] + + # Gauss weight + w = gauss_weights[local_gp] + + # Accumulate force (using tensor contraction!) + f_node += (dN_dx ⋅ σ) * (w * detJ) + end + end + + # Write result (in GPU version, no atomics needed - this node is ours!) + r[3*node_idx-2] = f_node[1] + r[3*node_idx-1] = f_node[2] + r[3*node_idx] = f_node[3] + end +end + +""" +Complete residual computation (two-phase approach) +""" +function compute_residual!( + r::Vector{Float64}, + u::Vector{Float64}, + nodes::Matrix{Float64}, + elements::Vector{NTuple{4,Int}}, + states_old::Vector{PlasticState}, + mat::Material, + node_to_elems::NodeToElementsMap +) + n_gp = length(states_old) + + # Storage for integration point data + σ_gp = Vector{SymmetricTensor{2,3,Float64,6}}(undef, n_gp) + states_new = Vector{PlasticState}(undef, n_gp) + + # Phase 1: Compute integration point data + compute_gp_data!(σ_gp, states_new, u, nodes, elements, states_old, mat) + + # Phase 2: Nodal assembly + nodal_assembly!(r, σ_gp, nodes, elements, node_to_elems) + + return r, states_new +end + +# ============================================================================ +# Test Setup +# ============================================================================ + +function main() + println("\n" * "="^70) + println("Nodal Assembly CPU Reference - Using Tensors.jl") + println("="^70) + + # Single Tet4 element + nodes = Float64[ + 0.0 1.0 0.0 0.0; # X coordinates + 0.0 0.0 1.0 0.0; # Y coordinates + 0.0 0.0 0.0 1.0 # Z coordinates + ] + + elements = [(1, 2, 3, 4)] + n_nodes = 4 + n_elems = 1 + n_gps = n_elems * 4 # 4 GPs per Tet4 + + # Material + mat = Material( + 210e3, # E = 210 GPa (steel) + 0.3, # ν = 0.3 + 250.0 # σ_y = 250 MPa + ) + + # Displacement (apply tension) + u = zeros(12) + u[4] = 0.01 # Move node 2 in X-direction + + # Initial states (all elastic) + states_old = [PlasticState(zero(SymmetricTensor{2,3,Float64}), 0.0) for _ in 1:n_gps] + + # Build node-to-elements map + println("\nBuilding node-to-elements map (CSR format)...") + node_to_elems = build_node_to_elems(elements, n_nodes) + + println("CSR ptr: ", node_to_elems.ptr) + println("CSR data: ", node_to_elems.data) + + # Verify each node touches exactly 1 element + for node_idx in 1:n_nodes + elem_start = node_to_elems.ptr[node_idx] + elem_end = node_to_elems.ptr[node_idx+1] - 1 + n_touching = elem_end - elem_start + 1 + touching_elems = node_to_elems.data[elem_start:elem_end] + println("Node $node_idx touches $n_touching element(s): $touching_elems") + end + + # Compute residual (two-phase approach) + println("\n" * "-"^70) + println("Computing residual (two-phase nodal assembly)...") + println("-"^70) + + r = zeros(12) + r, states_new = compute_residual!(r, u, nodes, elements, states_old, mat, node_to_elems) + + println("\nResidual vector (internal forces):") + for i in 1:n_nodes + rx = r[3*i-2] + ry = r[3*i-1] + rz = r[3*i] + @printf("Node %d: [%12.6e, %12.6e, %12.6e]\n", i, rx, ry, rz) + end + + println("\nResidual norm: ", norm(r)) + + # Check material states + println("\n" * "-"^70) + println("Material States at Gauss Points:") + println("-"^70) + + for (gp_idx, state) in enumerate(states_new) + elem_idx = (gp_idx - 1) ÷ 4 + 1 + local_gp = (gp_idx - 1) % 4 + 1 + + status = state.α > 0.0 ? "Plastic" : "Elastic" + @printf("Elem %d, GP %d: %s (α = %.6e)\n", elem_idx, local_gp, status, state.α) + end + + # Test force balance (should sum to zero for internal forces) + println("\n" * "-"^70) + println("Force Balance Check:") + println("-"^70) + + f_total = sum(reshape(r, 3, :), dims=2) + @printf("Sum of forces: [%.6e, %.6e, %.6e]\n", f_total[1], f_total[2], f_total[3]) + @printf("Should be ≈ zero for internal forces (tol: 1e-10)\n") + + if norm(f_total) < 1e-10 + println("✅ Force balance: PASSED") + else + println("❌ Force balance: FAILED") + end + + println("\n" * "="^70) + println("✅ Nodal assembly CPU reference complete!") + println("="^70) + println("\nNext step: Port this to GPU with CUDA.jl") + println(" Phase 1: @cuda compute_gp_data_kernel!(...)") + println(" Phase 2: @cuda nodal_assembly_kernel!(...)") + println("="^70 * "\n") +end + +main()