diff --git a/test/elements/test_extract_element_dofs.jl b/test/elements/test_extract_element_dofs.jl new file mode 100644 index 0000000..8eabf4e --- /dev/null +++ b/test/elements/test_extract_element_dofs.jl @@ -0,0 +1,218 @@ +# This file is a part of JuliaFEM. +# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md + +using Test +using JuliaFEM +using Tensors + +@testset "DOF Extraction: Single-field Scalar" begin + # Heat equation: 1 scalar per node + S = @DOFSet{T::DOF{Temperature, Vertex}} + elem = Element{Triangle{3}, Lagrange{1}, S}( + UInt(1), + (UInt64(10), UInt64(20), UInt64(30)) # Flat tuple of DOF indices + ) + + # Global solution vector (100 DOFs total) + u_global = collect(1.0:100.0) + + # Extract as flat tuple (NamedTuple wrapper for consistency) + dofs_flat = extract_element_dofs(elem, u_global) + @test dofs_flat isa NamedTuple + @test keys(dofs_flat) == (:T,) + @test dofs_flat.T == (10.0, 20.0, 30.0) + + # Extract structured (same for scalars) + dofs_struct = extract_element_dofs_structured(elem, u_global) + @test dofs_struct isa NamedTuple + @test keys(dofs_struct) == (:T,) + @test dofs_struct.T == (10.0, 20.0, 30.0) +end + +@testset "DOF Extraction: Single-field Vector" begin + # 2D elasticity: Vec{2} per node, 3 nodes + S = @DOFSet{u::DOF{Displacement{2}, Vertex}} + elem = Element{Triangle{3}, Lagrange{1}, S}( + UInt(1), + (UInt64(1), UInt64(2), UInt64(3), UInt64(4), UInt64(5), UInt64(6)) # Flat tuple + ) + + u_global = collect(1.0:100.0) + + # Flat extraction: all scalars (NamedTuple wrapper) + dofs_flat = extract_element_dofs(elem, u_global) + @test dofs_flat isa NamedTuple + @test keys(dofs_flat) == (:u,) + @test length(dofs_flat.u) == 6 + @test dofs_flat.u == (1.0, 2.0, 3.0, 4.0, 5.0, 6.0) + + # Structured extraction: Vec{2} instances (NamedTuple wrapper) + dofs_struct = extract_element_dofs_structured(elem, u_global) + @test dofs_struct isa NamedTuple + @test keys(dofs_struct) == (:u,) + @test length(dofs_struct.u) == 3 # 3 nodes + @test dofs_struct.u[1] == Vec(1.0, 2.0) + @test dofs_struct.u[2] == Vec(3.0, 4.0) + @test dofs_struct.u[3] == Vec(5.0, 6.0) +end + +@testset "DOF Extraction: Single-field Vector 3D" begin + # 3D elasticity: Vec{3} per node, 4 nodes (tetrahedron) + S = @DOFSet{u::DOF{Displacement{3}, Vertex}} + elem = Element{Tetrahedron{4}, Lagrange{1}, S}( + UInt(1), + tuple(UInt64.(1:12)...) # Flat tuple: 12 DOFs + ) + + u_global = collect(1.0:100.0) + + # Flat extraction: Returns NamedTuple with flat tuple + dofs_flat = extract_element_dofs(elem, u_global) + @test dofs_flat isa NamedTuple + @test keys(dofs_flat) == (:u,) + @test length(dofs_flat.u) == 12 + @test dofs_flat.u == (1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0) + + # Structured extraction: Returns NamedTuple with Vec{3} instances (one per vertex) + dofs_struct = extract_element_dofs_structured(elem, u_global) + @test dofs_struct isa NamedTuple + @test keys(dofs_struct) == (:u,) + @test length(dofs_struct.u) == 4 # 4 vertices + @test dofs_struct.u[1] == Vec(1.0, 2.0, 3.0) + @test dofs_struct.u[2] == Vec(4.0, 5.0, 6.0) + @test dofs_struct.u[3] == Vec(7.0, 8.0, 9.0) + @test dofs_struct.u[4] == Vec(10.0, 11.0, 12.0) +end + +@testset "DOF Extraction: Multi-field (Temperature + Displacement)" begin + # Thermoelasticity: Float64 (T) + Vec{3} (u) at vertices + S = @DOFSet{T::DOF{Temperature, Vertex}, u::DOF{Displacement{3}, Vertex}} + elem = Element{Tetrahedron{4}, Lagrange{1}, S}( + UInt(1), + (UInt64(1), UInt64(2), UInt64(3), UInt64(4), UInt64(10), UInt64(11), UInt64(12), UInt64(13), UInt64(14), UInt64(15), UInt64(16), UInt64(17), UInt64(18), UInt64(19), UInt64(20), UInt64(21)) # Flat: 4 T + 12 u = 16 DOFs + ) + + u_global = collect(1.0:100.0) + + # Flat extraction: NamedTuple of tuples + dofs_flat = extract_element_dofs(elem, u_global) + @test dofs_flat isa NamedTuple + @test keys(dofs_flat) == (:T, :u) + @test dofs_flat.T == (1.0, 2.0, 3.0, 4.0) + @test dofs_flat.u == (10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0, 17.0, 18.0, 19.0, 20.0, 21.0) + + # Structured extraction: NamedTuple with Vec{3} for displacement + dofs_struct = extract_element_dofs_structured(elem, u_global) + @test dofs_struct isa NamedTuple + @test keys(dofs_struct) == (:T, :u) + @test dofs_struct.T == (1.0, 2.0, 3.0, 4.0) + @test length(dofs_struct.u) == 4 # 4 nodes + @test dofs_struct.u[1] == Vec(10.0, 11.0, 12.0) + @test dofs_struct.u[2] == Vec(13.0, 14.0, 15.0) + @test dofs_struct.u[3] == Vec(16.0, 17.0, 18.0) + @test dofs_struct.u[4] == Vec(19.0, 20.0, 21.0) +end + +@testset "DOF Extraction: Non-contiguous indices" begin + # Test with scattered global DOF indices + S = @DOFSet{u::DOF{Displacement{2}, Vertex}} + elem = Element{Triangle{3}, Lagrange{1}, S}( + UInt(1), + (UInt64(5), UInt64(12), UInt64(23), UInt64(34), UInt64(45), UInt64(56)) # Non-sequential indices + ) + + u_global = collect(1.0:100.0) + + # Flat extraction + dofs_flat = extract_element_dofs(elem, u_global) + @test dofs_flat isa NamedTuple + @test keys(dofs_flat) == (:u,) + @test dofs_flat.u == (5.0, 12.0, 23.0, 34.0, 45.0, 56.0) + + # Structured extraction + dofs_struct = extract_element_dofs_structured(elem, u_global) + @test dofs_struct isa NamedTuple + @test keys(dofs_struct) == (:u,) + @test dofs_struct.u[1] == Vec(5.0, 12.0) + @test dofs_struct.u[2] == Vec(23.0, 34.0) + @test dofs_struct.u[3] == Vec(45.0, 56.0) +end + +@testset "Zero Allocation: All extraction functions" begin + # Single-field scalar + S_scalar = @DOFSet{T::DOF{Temperature, Vertex}} + elem_scalar = Element{Triangle{3}, Lagrange{1}, S_scalar}( + UInt(1), + (UInt64(10), UInt64(20), UInt64(30)) + ) + u_global = collect(1.0:100.0) + + # Warm up + extract_element_dofs(elem_scalar, u_global) + extract_element_dofs_structured(elem_scalar, u_global) + + # Check zero allocation + alloc_flat = @allocated extract_element_dofs(elem_scalar, u_global) + alloc_struct = @allocated extract_element_dofs_structured(elem_scalar, u_global) + @test alloc_flat == 0 + @test alloc_struct == 0 + + # Single-field vector 2D + S_vec2d = @DOFSet{u::DOF{Displacement{2}, Vertex}} + elem_vec2d = Element{Triangle{3}, Lagrange{1}, S_vec2d}( + UInt(1), + tuple(UInt64.(1:6)...) + ) + + # Warm up + extract_element_dofs(elem_vec2d, u_global) + extract_element_dofs_structured(elem_vec2d, u_global) + + # Check zero allocation + alloc_flat = @allocated extract_element_dofs(elem_vec2d, u_global) + alloc_struct = @allocated extract_element_dofs_structured(elem_vec2d, u_global) + @test alloc_flat == 0 + @test alloc_struct == 0 + + # Single-field vector 3D + S_vec3d = @DOFSet{u::DOF{Displacement{3}, Vertex}} + elem_vec3d = Element{Tetrahedron{4}, Lagrange{1}, S_vec3d}( + UInt(1), + tuple(UInt64.(1:12)...) + ) + + # Warm up + extract_element_dofs(elem_vec3d, u_global) + extract_element_dofs_structured(elem_vec3d, u_global) + + # Check zero allocation + alloc_flat = @allocated extract_element_dofs(elem_vec3d, u_global) + alloc_struct = @allocated extract_element_dofs_structured(elem_vec3d, u_global) + @test alloc_flat == 0 + @test alloc_struct == 0 + + # Multi-field (Temperature + Displacement) + S_multi = @DOFSet{T::DOF{Temperature,Vertex}, u::DOF{Displacement{3},Vertex}} + elem_multi = Element{Tetrahedron{4}, Lagrange{1}, S_multi}( + UInt(1), + tuple(UInt64.([1, 2, 3, 4, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21])...) + ) + + # Warm up + extract_element_dofs(elem_multi, u_global) + extract_element_dofs_structured(elem_multi, u_global) + + # Check zero allocation + alloc_flat = @allocated extract_element_dofs(elem_multi, u_global) + alloc_struct = @allocated extract_element_dofs_structured(elem_multi, u_global) + @test alloc_flat == 0 + @test alloc_struct == 0 +end + +println("✓ All DOF extraction tests passed!") +println(" - Single-field scalar extraction works") +println(" - Single-field vector extraction works") +println(" - Multi-field extraction works") +println(" - Non-contiguous indices work") +println(" - Flat and structured variants both work") +println(" - Zero allocations verified for all cases")