test(assemblers): add partitioned DOF matvec regression

This commit is contained in:
Jukka Aho
2026-05-09 18:37:22 +03:00
parent 9dd24ce2ce
commit 89d564d950
@@ -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