From 446360274ce4cf9d99b37ad046abafb62b16711e Mon Sep 17 00:00:00 2001 From: Jukka Aho Date: Sat, 9 May 2026 18:37:30 +0300 Subject: [PATCH] test(assemblers): add matrix-free operator regression --- test/assemblers/test_matrix_free_operator.jl | 232 +++++++++++++++++++ 1 file changed, 232 insertions(+) create mode 100644 test/assemblers/test_matrix_free_operator.jl diff --git a/test/assemblers/test_matrix_free_operator.jl b/test/assemblers/test_matrix_free_operator.jl new file mode 100644 index 0000000..eeb1d9c --- /dev/null +++ b/test/assemblers/test_matrix_free_operator.jl @@ -0,0 +1,232 @@ +# This file is a part of JuliaFEM. +# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md + +""" +Regression tests for the typed `MatrixFreeOperator` / +`MatrixFreeMassOperator`. The contract is: + + 1. `mul!(y, op, x)` matches `K * x` (or `M * x`) to round-off, and + also matches the closure-style `op(y, x)` invocation. + 2. `op * x` allocates an output vector and produces the same answer. + 3. `eltype(op) == Float64` and `size(op) == (n, n)`. + 4. The operator plugs into `LinearOperators.LinearOperator(...)` and + drives `IterativeSolvers.cg!` to a correct solution. + 5. A `MatrixFreeOperator` built with `dirichlet = PenaltyDirichlet` + reproduces `K + λ · diag(eᵈ)` row-by-row. + 6. `mul!` after the warmup is allocation-free. +""" + +using Test +using JuliaFEM +using JuliaFEM: DOFBasedCOOAssembler, DOFBasedCOOCache +using JuliaFEM: MatrixFreeOperator, MatrixFreeMassOperator, matrix_free_op +using JuliaFEM: PenaltyDirichlet +using JuliaFEM: extract_system, apply_K!, apply_M! +using JuliaFEM: ContinuumKernel, ContinuumFormulation, FullThreeD +using JuliaFEM: HeatKernel, HeatConductivity +using JuliaFEM: LinearElastic, Displacement, Vertex, Temperature +using JuliaFEM: create_elements!, @DOFSet, DOF +using LinearAlgebra +using SparseArrays +using Tensors +using Random +using LinearOperators +using IterativeSolvers + +# ---------------------------------------------------------------------------- +# Mesh helper +# ---------------------------------------------------------------------------- + +function _hex8_unit_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 _setup_elasticity(nx::Int, ny::Int, nz::Int) + mesh = _hex8_unit_box(nx, ny, nz) + 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{Hexahedron{8}, Lagrange{1}, S}) + asm = DOFBasedCOOAssembler() + cache = DOFBasedCOOCache(elements, dof_mgr, mesh, kernel) + return cache, asm, kernel, mesh +end + +function _setup_heat(nx::Int, ny::Int, nz::Int) + mesh = _hex8_unit_box(nx, ny, nz) + cond = HeatConductivity(k = 5.0) + kernel = HeatKernel(ContinuumFormulation{FullThreeD}(), cond) + S = @DOFSet{T::DOF{Temperature, Vertex}} + elements, dof_mgr = create_elements!(mesh, Element{Hexahedron{8}, Lagrange{1}, S}) + asm = DOFBasedCOOAssembler() + cache = DOFBasedCOOCache(elements, dof_mgr, mesh, kernel) + return cache, asm, kernel, mesh +end + +# ---------------------------------------------------------------------------- +# Tests +# ---------------------------------------------------------------------------- + +@testset "MatrixFreeOperator: contract and behaviour" begin + Random.seed!(20260508) + + @testset "Stiffness operator: mul! matches K * x" begin + cache, asm, kernel, mesh = _setup_elasticity(3, 2, 2) + assemble!(cache, asm, kernel, mesh) + K, _ = extract_system(cache) + n = size(K, 1) + + op = MatrixFreeOperator(cache, asm, kernel, mesh) + + @test eltype(op) == Float64 + @test size(op) == (n, n) + @test size(op, 1) == n + @test size(op, 2) == n + + for trial in 1:5 + x = randn(n) + y_ref = K * x + + y_mul = zeros(n) + mul!(y_mul, op, x) + @test norm(y_mul - y_ref) / norm(y_ref) < 1e-12 + + y_call = zeros(n) + op(y_call, x) + @test y_call == y_mul + + y_mat = op * x + @test norm(y_mat - y_ref) / norm(y_ref) < 1e-12 + end + end + + @testset "matrix_free_op factory returns MatrixFreeOperator" begin + cache, asm, kernel, mesh = _setup_elasticity(2, 2, 2) + op = matrix_free_op(cache, asm, kernel, mesh) + @test op isa MatrixFreeOperator + end + + @testset "Mass operator: mul! matches M * x" begin + # Use heat kernel because the mass term is non-zero when the + # heat capacity is enabled; ContinuumKernel/HeatConductivity + # default mass entries are zero unless density is configured. + cache, asm, kernel, mesh = _setup_heat(2, 2, 2) + + # Build the matrix-free mass operator and reference M via + # apply_M! (rather than the full assemble_M!) so we are + # comparing the matrix-free path against itself column-wise. + op_M = MatrixFreeMassOperator(cache, asm, kernel, mesh) + n = size(op_M, 1) + + @test size(op_M) == (n, n) + @test eltype(op_M) == Float64 + + # Build a dense reference by mat-vec on canonical basis. + M_ref = zeros(n, n) + e_j = zeros(n) + for j in 1:n + fill!(e_j, 0.0); e_j[j] = 1.0 + apply_M!(view(M_ref, :, j) |> collect, cache, asm, kernel, mesh, e_j) + tmp = zeros(n) + apply_M!(tmp, cache, asm, kernel, mesh, e_j) + M_ref[:, j] .= tmp + end + + for trial in 1:3 + x = randn(n) + y_ref = M_ref * x + y_mul = zeros(n) + mul!(y_mul, op_M, x) + @test norm(y_mul - y_ref) / max(norm(y_ref), 1e-30) < 1e-10 + + y_seed = randn(n) + α = 1.7 + β = -0.25 + y_fused = copy(y_seed) + mul!(y_fused, op_M, x, α, β) + @test norm(y_fused - (α * y_ref + β * y_seed)) / + max(norm(y_ref), 1e-30) < 1e-10 + end + end + + @testset "PenaltyDirichlet folded into mat-vec" begin + cache, asm, kernel, mesh = _setup_heat(3, 2, 2) + assemble!(cache, asm, kernel, mesh) + K, _ = extract_system(cache) + n = size(K, 1) + + fixed = [1, 4, 7] + vals = [1.0, 2.0, 3.0] + λ = 1.0e10 + c = PenaltyDirichlet(fixed, vals; penalty = λ) + + K_pen = copy(K) + for d in fixed + K_pen[d, d] += λ + end + + op = MatrixFreeOperator(cache, asm, kernel, mesh; dirichlet = c) + x = randn(n) + y_ref = K_pen * x + y_mf = zeros(n) + mul!(y_mf, op, x) + @test norm(y_mf - y_ref) / norm(y_ref) < 1e-10 + end + + @testset "Plugs into LinearOperators + IterativeSolvers.cg!" begin + cache, asm, kernel, mesh = _setup_heat(3, 3, 3) + assemble!(cache, asm, kernel, mesh) + K, _ = extract_system(cache) + n = size(K, 1) + + # Constrain the boundary so K is invertible. + fixed = collect(1:n) |> dofs -> filter(d -> d in [1, n], dofs) + c = PenaltyDirichlet(fixed, fill(0.0, length(fixed)); penalty = 1.0e12) + + op = MatrixFreeOperator(cache, asm, kernel, mesh; dirichlet = c) + linop = LinearOperator(Float64, size(op, 1), size(op, 2), true, true, op) + + K_pen = copy(K) + for d in fixed + K_pen[d, d] += 1.0e12 + end + b = randn(n) + u_ref = K_pen \ b + + u_cg = zeros(n) + cg!(u_cg, linop, b; abstol = 1e-12, reltol = 1e-12, maxiter = 4n) + @test norm(u_cg - u_ref) / norm(u_ref) < 1e-6 + end + + @testset "mul! is allocation-free after warmup" begin + cache, asm, kernel, mesh = _setup_elasticity(2, 2, 2) + op = MatrixFreeOperator(cache, asm, kernel, mesh) + n = size(op, 1) + x = randn(n) + y = zeros(n) + + # Warmup (compile + first-call alloc on Pass 1). + mul!(y, op, x) + mul!(y, op, x) + + allocs = @allocated mul!(y, op, x) + @test allocs == 0 + end +end