diff --git a/test/assemblers/test_apply_K_masked_rows.jl b/test/assemblers/test_apply_K_masked_rows.jl new file mode 100644 index 0000000..a0a2f08 --- /dev/null +++ b/test/assemblers/test_apply_K_masked_rows.jl @@ -0,0 +1,233 @@ +# This file is a part of JuliaFEM. +# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md + +""" +`apply_K_masked_rows!`: disjoint owned-row masks (vertex-field ownership) sum to +the global [`apply_K!`](@ref) when `x` is the full vector on each fake rank. +""" + +using Test +using JuliaFEM +using JuliaFEM: DOFBasedCOOAssembler, DOFBasedCOOCache +using JuliaFEM: apply_K!, apply_K_masked_rows!, apply_K_owned_rows! +using JuliaFEM: create_elements!, @DOFSet, DOF, Displacement, Vertex +using JuliaFEM: brick_hex_partition_slabs, validate_partition +using JuliaFEM: node_partition_owner_min!, mark_owned_vertex_field_dofs! +using Random +using Tensors + +@testset "apply_K_masked_rows!: partition masks sum to apply_K!" begin + Random.seed!(20260511) + nx, ny, nz = 3, 4, 2 + nparts = 3 + mesh = create_structured_box_mesh(Hex8; nx = nx, ny = ny, nz = nz) + material = LinearElastic(E = 210e9, ν = 0.3) + kernel = ContinuumKernel(ContinuumFormulation{FullThreeD}(), + material, Displacement{3}()) + S = @DOFSet{u::DOF{Displacement{3}, Vertex}} + elements, handler = create_elements!(mesh, Element{Hex8, Lagrange{1}, S}) + asm = DOFBasedCOOAssembler() + cache = DOFBasedCOOCache(elements, handler, mesh, kernel) + + layout = brick_hex_partition_slabs(nx, ny, nz, nparts; axis = :y) + validate_partition(layout, length(elements)) + + nnodes = length(mesh.nodes) + node_own = Vector{Int}(undef, nnodes) + node_partition_owner_min!(node_own, layout, mesh) + + n = cache.ndofs + x = randn(n) + y_ref = zeros(n) + apply_K!(y_ref, cache, asm, kernel, mesh, x) + + y_sum = zeros(n) + own = falses(n) + for part in 1:nparts + mark_owned_vertex_field_dofs!(own, handler, node_own, part) + y_part = zeros(n) + apply_K_masked_rows!(y_part, own, cache, asm, kernel, mesh, x) + y_sum .+= y_part + end + @test y_sum ≈ y_ref rtol = 1e-11 atol = 1e-11 +end + +@testset "apply_K_masked_rows!: DimensionMismatch on mask length" begin + nx, ny, nz = 2, 2, 2 + mesh = create_structured_box_mesh(Hex8; nx = nx, ny = ny, nz = nz) + material = LinearElastic(E = 210e9, ν = 0.3) + kernel = ContinuumKernel(ContinuumFormulation{FullThreeD}(), + material, Displacement{3}()) + S = @DOFSet{u::DOF{Displacement{3}, Vertex}} + elements, handler = create_elements!(mesh, Element{Hex8, Lagrange{1}, S}) + asm = DOFBasedCOOAssembler() + cache = DOFBasedCOOCache(elements, handler, mesh, kernel) + n = cache.ndofs + @test_throws DimensionMismatch apply_K_masked_rows!( + zeros(n), falses(n - 1), cache, asm, kernel, mesh, zeros(n)) +end + +@testset "apply_K_owned_rows!: matches apply_K! on owned rows" begin + Random.seed!(20260512) + nx, ny, nz = 3, 3, 2 + mesh = create_structured_box_mesh(Hex8; nx = nx, ny = ny, nz = nz) + material = LinearElastic(E = 210e9, ν = 0.3) + kernel = ContinuumKernel(ContinuumFormulation{FullThreeD}(), + material, Displacement{3}()) + S = @DOFSet{u::DOF{Displacement{3}, Vertex}} + elements, handler = create_elements!(mesh, Element{Hex8, Lagrange{1}, S}) + asm = DOFBasedCOOAssembler() + cache = DOFBasedCOOCache(elements, handler, mesh, kernel) + + layout = brick_hex_partition_slabs(nx, ny, nz, 2; axis = :x) + validate_partition(layout, length(elements)) + nnodes = length(mesh.nodes) + node_own = Vector{Int}(undef, nnodes) + node_partition_owner_min!(node_own, layout, mesh) + + n = cache.ndofs + owned = falses(n) + mark_owned_vertex_field_dofs!(owned, handler, node_own, 1) + + x = randn(n) + y_ref = zeros(n) + apply_K!(y_ref, cache, asm, kernel, mesh, x) + + y_own = zeros(n) + apply_K_owned_rows!(y_own, owned, cache, asm, kernel, mesh, x) + + @inbounds for i in 1:n + if owned[i] + @test y_own[i] ≈ y_ref[i] + else + @test y_own[i] == 0.0 + end + end +end + +@testset "apply_K_owned_rows!: partition sum equals apply_K!" begin + Random.seed!(20260513) + nx, ny, nz = 3, 4, 2 + nparts = 3 + mesh = create_structured_box_mesh(Hex8; nx = nx, ny = ny, nz = nz) + material = LinearElastic(E = 210e9, ν = 0.3) + kernel = ContinuumKernel(ContinuumFormulation{FullThreeD}(), + material, Displacement{3}()) + S = @DOFSet{u::DOF{Displacement{3}, Vertex}} + elements, handler = create_elements!(mesh, Element{Hex8, Lagrange{1}, S}) + asm = DOFBasedCOOAssembler() + cache = DOFBasedCOOCache(elements, handler, mesh, kernel) + + layout = brick_hex_partition_slabs(nx, ny, nz, nparts; axis = :y) + validate_partition(layout, length(elements)) + + nnodes = length(mesh.nodes) + node_own = Vector{Int}(undef, nnodes) + node_partition_owner_min!(node_own, layout, mesh) + + n = cache.ndofs + x = randn(n) + y_ref = zeros(n) + apply_K!(y_ref, cache, asm, kernel, mesh, x) + + y_sum = zeros(n) + own = falses(n) + for part in 1:nparts + mark_owned_vertex_field_dofs!(own, handler, node_own, part) + y_part = zeros(n) + apply_K_owned_rows!(y_part, own, cache, asm, kernel, mesh, x) + y_sum .+= y_part + end + @test y_sum ≈ y_ref rtol = 1e-11 atol = 1e-11 +end + +@testset "apply_K_owned_rows!: agrees with apply_K_masked_rows!" begin + nx, ny, nz = 2, 3, 2 + mesh = create_structured_box_mesh(Hex8; nx = nx, ny = ny, nz = nz) + material = LinearElastic(E = 210e9, ν = 0.3) + kernel = ContinuumKernel(ContinuumFormulation{FullThreeD}(), + material, Displacement{3}()) + S = @DOFSet{u::DOF{Displacement{3}, Vertex}} + elements, handler = create_elements!(mesh, Element{Hex8, Lagrange{1}, S}) + asm = DOFBasedCOOAssembler() + cache = DOFBasedCOOCache(elements, handler, mesh, kernel) + layout = brick_hex_partition_slabs(nx, ny, nz, 2; axis = :z) + validate_partition(layout, length(elements)) + nnodes = length(mesh.nodes) + node_own = Vector{Int}(undef, nnodes) + node_partition_owner_min!(node_own, layout, mesh) + n = cache.ndofs + keep = falses(n) + mark_owned_vertex_field_dofs!(keep, handler, node_own, 2) + x = randn(n) + y_m = zeros(n) + y_o = zeros(n) + apply_K_masked_rows!(y_m, keep, cache, asm, kernel, mesh, x) + apply_K_owned_rows!(y_o, keep, cache, asm, kernel, mesh, x) + @test y_m ≈ y_o +end + +@testset "apply_K_owned_rows!: DimensionMismatch on mask length" begin + nx, ny, nz = 2, 2, 2 + mesh = create_structured_box_mesh(Hex8; nx = nx, ny = ny, nz = nz) + material = LinearElastic(E = 210e9, ν = 0.3) + kernel = ContinuumKernel(ContinuumFormulation{FullThreeD}(), + material, Displacement{3}()) + S = @DOFSet{u::DOF{Displacement{3}, Vertex}} + elements, handler = create_elements!(mesh, Element{Hex8, Lagrange{1}, S}) + asm = DOFBasedCOOAssembler() + cache = DOFBasedCOOCache(elements, handler, mesh, kernel) + n = cache.ndofs + @test_throws DimensionMismatch apply_K_owned_rows!( + zeros(n), falses(n - 1), cache, asm, kernel, mesh, zeros(n)) +end + +@testset "apply_K_masked_rows!: zero allocations after warmup" begin + nx, ny, nz = 2, 2, 2 + mesh = create_structured_box_mesh(Hex8; nx = nx, ny = ny, nz = nz) + material = LinearElastic(E = 210e9, ν = 0.3) + kernel = ContinuumKernel(ContinuumFormulation{FullThreeD}(), + material, Displacement{3}()) + S = @DOFSet{u::DOF{Displacement{3}, Vertex}} + elements, handler = create_elements!(mesh, Element{Hex8, Lagrange{1}, S}) + asm = DOFBasedCOOAssembler() + cache = DOFBasedCOOCache(elements, handler, mesh, kernel) + layout = brick_hex_partition_slabs(nx, ny, nz, 2; axis = :x) + validate_partition(layout, length(elements)) + nnodes = length(mesh.nodes) + node_own = Vector{Int}(undef, nnodes) + node_partition_owner_min!(node_own, layout, mesh) + n = cache.ndofs + keep = falses(n) + mark_owned_vertex_field_dofs!(keep, handler, node_own, 1) + y = zeros(n) + x = randn(n) + apply_K_masked_rows!(y, keep, cache, asm, kernel, mesh, x) + a = @allocated apply_K_masked_rows!(y, keep, cache, asm, kernel, mesh, x) + @test a == 0 +end + +@testset "apply_K_owned_rows!: zero allocations after warmup" begin + nx, ny, nz = 2, 2, 2 + mesh = create_structured_box_mesh(Hex8; nx = nx, ny = ny, nz = nz) + material = LinearElastic(E = 210e9, ν = 0.3) + kernel = ContinuumKernel(ContinuumFormulation{FullThreeD}(), + material, Displacement{3}()) + S = @DOFSet{u::DOF{Displacement{3}, Vertex}} + elements, handler = create_elements!(mesh, Element{Hex8, Lagrange{1}, S}) + asm = DOFBasedCOOAssembler() + cache = DOFBasedCOOCache(elements, handler, mesh, kernel) + layout = brick_hex_partition_slabs(nx, ny, nz, 2; axis = :x) + validate_partition(layout, length(elements)) + nnodes = length(mesh.nodes) + node_own = Vector{Int}(undef, nnodes) + node_partition_owner_min!(node_own, layout, mesh) + n = cache.ndofs + keep = falses(n) + mark_owned_vertex_field_dofs!(keep, handler, node_own, 1) + y = zeros(n) + x = randn(n) + apply_K_owned_rows!(y, keep, cache, asm, kernel, mesh, x) + a = @allocated apply_K_owned_rows!(y, keep, cache, asm, kernel, mesh, x) + @test a == 0 +end