mirror of
https://github.com/JuliaFEM/JuliaFEM.jl.git
synced 2026-08-06 04:21:33 +00:00
test(assemblers): add matrix-free operator regression
This commit is contained in:
Jukka Aho
@@ -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
|
||||
Reference in New Issue
Block a user