From 89d564d9502a94c56cf42b940bdee48926699c62 Mon Sep 17 00:00:00 2001 From: Jukka Aho Date: Sat, 9 May 2026 18:37:22 +0300 Subject: [PATCH] test(assemblers): add partitioned DOF matvec regression --- .../test_dof_based_partitioned_matvec.jl | 186 ++++++++++++++++++ 1 file changed, 186 insertions(+) create mode 100644 test/assemblers/test_dof_based_partitioned_matvec.jl diff --git a/test/assemblers/test_dof_based_partitioned_matvec.jl b/test/assemblers/test_dof_based_partitioned_matvec.jl new file mode 100644 index 0000000..3eb46ce --- /dev/null +++ b/test/assemblers/test_dof_based_partitioned_matvec.jl @@ -0,0 +1,186 @@ +# This file is a part of JuliaFEM. +# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md + +""" +Reference tests for distributed-style matrix-free matvec **without MPI**: + +`apply_K!` matches the sum of `apply_K_contributions!` over disjoint element +sets that partition the mesh, and [`prepare_multiply_workspace!`](@ref) / +[`MeshPartitionLayout`](@ref) behave as documented. +""" + +using Test +using JuliaFEM +using JuliaFEM: DOFBasedCOOAssembler, DOFBasedCOOCache +using JuliaFEM: apply_K!, apply_K_contributions! +using JuliaFEM: create_elements!, @DOFSet, DOF, Displacement, Vertex +using JuliaFEM: MeshPartitionLayout, uniform_single_partition +using JuliaFEM: element_indices_for_part, validate_partition +using JuliaFEM: LocalMultiplyLayout, prepare_multiply_workspace! +using JuliaFEM: MatrixFreeOperator +using LinearAlgebra +using Random +using Tensors + +function _build_hex8_box(nx::Int, ny::Int, nz::Int) + nodes = Vec{3,Float64}[] + nidx(i, j, k) = (i - 1) + (j - 1) * (nx + 1) + (k - 1) * (nx + 1) * (ny + 1) + 1 + for k in 1:(nz + 1), j in 1:(ny + 1), i in 1:(nx + 1) + push!(nodes, Vec{3}((Float64(i - 1) / nx, + Float64(j - 1) / ny, + Float64(k - 1) / nz))) + end + conns = NTuple{8,UInt32}[] + for k in 1:nz, j in 1:ny, i in 1:nx + n1 = nidx(i, j, k) + n2 = nidx(i + 1, j, k) + n3 = nidx(i + 1, j + 1, k) + n4 = nidx(i, j + 1, k) + n5 = nidx(i, j, k + 1) + n6 = nidx(i + 1, j, k + 1) + n7 = nidx(i + 1, j + 1, k + 1) + n8 = nidx(i, j + 1, k + 1) + push!(conns, (UInt32(n1), UInt32(n2), UInt32(n3), UInt32(n4), + UInt32(n5), UInt32(n6), UInt32(n7), UInt32(n8))) + end + return Mesh{8,Hexahedron{8}}(nodes, conns) +end + +function _build_single_tet4() + nodes = Vec{3,Float64}[ + Vec{3}((0.0, 0.0, 0.0)), + Vec{3}((1.0, 0.0, 0.0)), + Vec{3}((0.5, 1.0, 0.0)), + Vec{3}((0.5, 0.5, 1.0)), + ] + conns = [(UInt32(1), UInt32(2), UInt32(3), UInt32(4))] + return Mesh{Tetrahedron{4}}(nodes, conns) +end + +function _setup(mesh, ::Type{Topo}) where {Topo} + material = LinearElastic(E = 210e9, ν = 0.3) + kernel = ContinuumKernel(ContinuumFormulation{FullThreeD}(), + material, Displacement{3}()) + S = @DOFSet{u::DOF{Displacement{3}, Vertex}} + elements, dof_mgr = create_elements!(mesh, Element{Topo, Lagrange{1}, S}) + asm = DOFBasedCOOAssembler() + cache = DOFBasedCOOCache(elements, dof_mgr, mesh, kernel) + return cache, asm, kernel, mesh, elements +end + +"""Two-block element partition: ids 1:half → part 1, rest → part 2.""" +function _two_block_layout(nelements::Int) + nelements ≥ 1 || throw(ArgumentError("need at least one element")) + half = nelements ÷ 2 + ids = ones(Int, nelements) + @inbounds for i in (half + 1):nelements + ids[i] = 2 + end + return MeshPartitionLayout(ids), 1:half, (half + 1):nelements +end + +@testset "partitioned matvec reference (two fake ranks)" begin + Random.seed!(20260509) + + @testset "single Tet4 — empty block + full block" begin + mesh = _build_single_tet4() + cache, asm, kernel, m, elements = _setup(mesh, Tetrahedron{4}) + ne = length(elements) + @test ne == 1 + layout, r1, r2 = _two_block_layout(ne) + validate_partition(layout, ne) + @test isempty(r1) + @test collect(r2) == [1] + + e1 = element_indices_for_part(layout, 1) + e2 = element_indices_for_part(layout, 2) + @test isempty(e1) + @test e2 == [1] + + n = cache.ndofs + x = randn(n) + y_full = zeros(n) + apply_K!(y_full, cache, asm, kernel, m, x) + + ya = zeros(n) + yb = zeros(n) + apply_K_contributions!(ya, cache, asm, kernel, m, x, e1) + apply_K_contributions!(yb, cache, asm, kernel, m, x, e2) + @test ya + yb ≈ y_full rtol = 1e-12 atol = 1e-12 + end + + @testset "Hex8 box — two non-empty blocks" begin + mesh = _build_hex8_box(3, 2, 2) + cache, asm, kernel, m, elements = _setup(mesh, Hexahedron{8}) + ne = length(elements) + layout, r1, r2 = _two_block_layout(ne) + validate_partition(layout, ne) + @test sort!(collect(union(Set(r1), Set(r2)))) == collect(1:ne) + + e1 = element_indices_for_part(layout, 1) + e2 = element_indices_for_part(layout, 2) + @test sort!(vcat(e1, e2)) == collect(1:ne) + + n = cache.ndofs + for trial in 1:6 + x = randn(n) + y_full = zeros(n) + apply_K!(y_full, cache, asm, kernel, m, x) + + ya = zeros(n) + yb = zeros(n) + apply_K_contributions!(ya, cache, asm, kernel, m, x, e1) + apply_K_contributions!(yb, cache, asm, kernel, m, x, e2) + @test ya + yb ≈ y_full rtol = 1e-11 atol = 1e-11 + end + end + + @testset "uniform_single_partition matches apply_K!" begin + mesh = _build_hex8_box(2, 2, 2) + cache, asm, kernel, m, elements = _setup(mesh, Hexahedron{8}) + ne = length(elements) + layout = uniform_single_partition(ne) + validate_partition(layout, ne) + eall = element_indices_for_part(layout, 1) + @test length(eall) == ne + + n = cache.ndofs + x = randn(n) + y_full = zeros(n) + apply_K!(y_full, cache, asm, kernel, m, x) + yc = zeros(n) + apply_K_contributions!(yc, cache, asm, kernel, m, x, eall) + @test yc ≈ y_full rtol = 1e-12 atol = 1e-12 + end + + @testset "invalid element id" begin + mesh = _build_single_tet4() + cache, asm, kernel, m, _ = _setup(mesh, Tetrahedron{4}) + n = cache.ndofs + x = zeros(n) + y = zeros(n) + @test_throws ArgumentError apply_K_contributions!(y, cache, asm, kernel, m, x, [0]) + @test_throws ArgumentError apply_K_contributions!(y, cache, asm, kernel, m, x, [2]) + end +end + +@testset "multiply workspace layout + MatrixFreeOperator default" begin + mesh = _build_hex8_box(2, 2, 2) + cache, asm, kernel, m, _ = _setup(mesh, Hexahedron{8}) + n = cache.ndofs + x = randn(n) + work = zeros(n) + prepare_multiply_workspace!(work, x, LocalMultiplyLayout()) + @test work ≈ x + + work2 = copy(x) + prepare_multiply_workspace!(work2, x, LocalMultiplyLayout()) + @test work2 ≈ x + + op = MatrixFreeOperator(cache, asm, kernel, m) + y = zeros(n) + mul!(y, op, x) + yref = zeros(n) + apply_K!(yref, cache, asm, kernel, m, x) + @test y ≈ yref rtol = 1e-12 atol = 1e-12 +end