Glassvision/JuliaFEM.jl
1
0
Fork 0
mirror of https://github.com/JuliaFEM/JuliaFEM.jl.git synced 2026-09-19 17:58:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(elements): centralize interpolation code generation

Browse Source 
Consolidate the four `@generated` interpolators onto one Expr toolkit so
scalar/vector branching and flat-DOF indexing stay consistent, and fix the
broken vector arm in `interpolate_field_value`.

- Add compile-time Expr helpers (`_classify_field_quantity`,
  `_per_node_dof_exprs`, `_value_expr`, `_grad_expr`, `_field_dof_count`,
  `_classify_dofset_field`) shared by `interpolate_fields`,
  `interpolate_field`, `interpolate_field_value`, and
  `interpolate_local_fields`.
- Fix `interpolate_field_value` vector dispatch: the generated branch
  compared `quantity_type` as if it were a type value instead of using the
  field quantity from `quantity_type(field_spec)`.
- Refresh module docs (four entry points + helper layering); drop a long
  redundant doc example under `interpolate_local_fields`.
This commit is contained in:
Jukka Aho
2026-05-09 17:05:34 +03:00
parent 676825a730
commit 37abd55c34
1 changed files with 181 additions and 316 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/elements/interpolate.jl
+181 -316
View File
@@ -4,14 +4,132 @@
"""
Field interpolation at quadrature points.
Given element DOFs and a point in reference coordinates, interpolate field values
and gradients. Returns a NamedTuple with interpolated quantities.
Given element DOFs and a point in reference coordinates, interpolate field
values, gradients, and rates. The four entry points (`interpolate_fields`,
`interpolate_field`, `interpolate_field_value`, `interpolate_local_fields`)
are all `@generated` functions; their per-field expansion is built from a
single small set of expression-level helpers defined at the top of this
file.
See `src/elements/README.md` for usage examples.
"""
using Tensors
# ---------------------------------------------------------------------------
# Compile-time expression helpers shared by every @generated entry point.
#
# These are *plain* functions that return `Expr` values. The `@generated`
# bodies below call them while the specialization is being constructed, so
# the produced expressions are inlined into the final method body and the
# helpers themselves never run at execution time.
# ---------------------------------------------------------------------------
"""
_classify_field_quantity(Q) -> (kind::Symbol, vec_dim::Int)
Classify a quantity type `Q` (as produced by `quantity_type(field_spec)`).
- Returns `(:scalar, 1)` for `Float64`.
- Returns `(:vector, D)` for first-order `Tensor` types of dimension `D`
(e.g. `Vec{3}`).
Throws an error for any other quantity type.
"""
function _classify_field_quantity(Q)
if Q === Float64
return (:scalar, 1)
elseif Q isa UnionAll && Q.body <: Tensor && Q.body.parameters[1] == 1
return (:vector, Q.body.parameters[2])
else
error("Unsupported quantity type: $Q")
end
end
"""
_per_node_dof_exprs(varname, offset, n_nodes, kind, vec_dim) -> Vector{Expr}
Build the per-node DOF expressions read out of `varname[elem.dof_indices[...]]`
for a field that starts at the given flat `offset`.
For a scalar field each entry is just `varname[elem.dof_indices[i]]`.
For a vector field of dimension `vec_dim` each entry packs the `vec_dim`
component reads into a `Vec{vec_dim}(...)` literal.
"""
function _per_node_dof_exprs(
varname::Symbol, offset::Int, n_nodes::Int, kind::Symbol, vec_dim::Int,
)
if kind === :scalar
return [:($(varname)[elem.dof_indices[$(offset + i)]]) for i in 1:n_nodes]
else # :vector
out = Vector{Expr}(undef, n_nodes)
for node in 0:(n_nodes - 1)
comps = [
:($(varname)[elem.dof_indices[$(offset + node * vec_dim + comp)]])
for comp in 1:vec_dim
]
out[node + 1] = :(Vec{$vec_dim}($(Expr(:tuple, comps...))))
end
return out
end
end
"""
_value_expr(per_node_exprs, basis_var) -> Expr
Sum-of-products `∑ᵢ basis_var[i] * uᵢ` where `uᵢ` is the i-th expression
in `per_node_exprs`.
"""
function _value_expr(per_node_exprs::Vector{Expr}, basis_var::Symbol)
terms = [:($(basis_var)[$i] * $(per_node_exprs[i])) for i in eachindex(per_node_exprs)]
return Expr(:call, :+, terms...)
end
"""
_grad_expr(per_node_exprs, dN_var, kind) -> Expr
Sum-of-products `∑ᵢ dN_var[i] op uᵢ`. The product operator depends on the
field kind: scalar fields use `*` (Vec × Float → Vec), vector fields use
`⊗` (Vec ⊗ Vec → Tensor{2}).
"""
function _grad_expr(per_node_exprs::Vector{Expr}, dN_var::Symbol, kind::Symbol)
op = kind === :scalar ? :* : :⊗
terms = [
Expr(:call, op, :($(dN_var)[$i]), per_node_exprs[i])
for i in eachindex(per_node_exprs)
]
return Expr(:call, :+, terms...)
end
"""
_field_dof_count(kind, n_nodes, vec_dim) -> Int
Number of flat DOF slots consumed by one field block.
"""
_field_dof_count(kind::Symbol, n_nodes::Int, vec_dim::Int) =
kind === :scalar ? n_nodes : n_nodes * vec_dim
"""
_classify_dofset_field(S, fname) -> (kind, vec_dim)
Read the quantity classification for the named field of a `DOFSet`. Also
asserts the field lives on `Vertex` entities, which is the only entity
type the interpolators currently support.
"""
function _classify_dofset_field(S, fname::Symbol)
field_spec = fieldtype(S, fname)
entity_type = field_spec.parameters[2]
if entity_type !== Vertex
error("Unsupported entity type: $entity_type (only Vertex supported for now)")
end
Q = quantity_type(field_spec)
return _classify_field_quantity(Q)
end
# ---------------------------------------------------------------------------
# Generated entry points
# ---------------------------------------------------------------------------
"""
interpolate_fields(elem::Element{K,P,S,N}, u_global::AbstractVector, ξ::Vec) → NamedTuple
@@ -45,101 +163,31 @@ happens at compile time.
@generated function interpolate_fields(
elem::Element{K,P,S,N},
u_global::AbstractVector,
ξ::Vec
ξ::Vec,
) where {K,P,S<:DOFSet,N}
field_names = fieldnames(S)
topology = K()
basis = P()
n_nodes = nnodes(topology)
# Build expressions for each field interpolation
field_exprs = Expr[]
offset = 0 # Track position in flat dof_indices tuple
for fname in field_names
field_spec = fieldtype(S, fname)
field_type = field_spec.parameters[1] # Displacement{3}
entity_type = field_spec.parameters[2]
# Extract quantity type via trait
Q = quantity_type(field_spec) # Vec{3} or Float64
if entity_type === Vertex
# Standard nodal basis
if Q === Float64
# Scalar field interpolation
# value = ∑ Nᵢ(ξ) * uᵢ
# gradient = ∑ ∇Nᵢ(ξ) * uᵢ
value_terms = Expr[]
grad_terms = Expr[]
for i in 1:n_nodes
push!(value_terms, :(Nvals[$i] * u_global[elem.dof_indices[$(offset+i)]]))
push!(grad_terms, :(dN[$i] * u_global[elem.dof_indices[$(offset+i)]]))
end
value_expr = Expr(:call, :+, value_terms...)
grad_expr = Expr(:call, :+, grad_terms...)
# Add field value and gradient
push!(field_exprs, Expr(:(=), fname, value_expr))
push!(field_exprs, Expr(:(=), Symbol("∇", fname), grad_expr))
offset += n_nodes
elseif Q isa UnionAll && Q.body <: Tensor && Q.body.parameters[1] == 1
# Vector field interpolation
# value = ∑ Nᵢ(ξ) * uᵢ (each uᵢ is a Vec)
# gradient = ∑ ∇Nᵢ(ξ) ⊗ uᵢ (tensor product)
vec_dim = Q.body.parameters[2]
value_terms = Expr[]
grad_terms = Expr[]
for node in 0:(n_nodes-1)
# Extract vector components for this node from flat tuple
vec_comps = [:(u_global[elem.dof_indices[$(offset+node*vec_dim+comp)]]) for comp in 1:vec_dim]
u_node = :(Vec{$vec_dim}($(Expr(:tuple, vec_comps...))))
node_idx = node + 1
# value += N_i * u_i
push!(value_terms, :(Nvals[$node_idx] * $u_node))
# gradient += ∇N_i ⊗ u_i
push!(grad_terms, :(dN[$node_idx] ⊗ $u_node))
end
value_expr = Expr(:call, :+, value_terms...)
grad_expr = Expr(:call, :+, grad_terms...)
# Add field value and gradient
push!(field_exprs, Expr(:(=), fname, value_expr))
push!(field_exprs, Expr(:(=), Symbol("∇", fname), grad_expr))
offset += n_nodes * vec_dim
else
error("Unsupported quantity type: $Q")
end
else
error("Unsupported entity type: $entity_type (only Vertex supported for now)")
end
offset = 0
for fname in fieldnames(S)
kind, vec_dim = _classify_dofset_field(S, fname)
per_node = _per_node_dof_exprs(:u_global, offset, n_nodes, kind, vec_dim)
push!(field_exprs, Expr(:(=), fname, _value_expr(per_node, :Nvals)))
push!(field_exprs, Expr(:(=), Symbol("∇", fname), _grad_expr(per_node, :dN, kind)))
offset += _field_dof_count(kind, n_nodes, vec_dim)
end
# Build complete function body
# 1. Evaluate basis functions and derivatives
# 2. Compute all interpolations
# 3. Return NamedTuple
nt_expr = Expr(:tuple, field_exprs...)
return quote
@inbounds begin
# Evaluate basis functions once
Nvals = get_basis_functions($topology, $basis, ξ)
dN = get_basis_derivatives($topology, $basis, ξ)
# Return interpolated values
return $nt_expr
end
end
@@ -166,85 +214,35 @@ val, grad = interpolate_field(elem, u_global, :T, Vec((0.25, 0.25, 0.25)))
elem::Element{K,P,S,N},
u_global::AbstractVector,
field::Symbol,
ξ::Vec
ξ::Vec,
) where {K,P,S<:DOFSet,N}
field_names = fieldnames(S)
topology = K()
basis = P()
n_nodes = nnodes(topology)
# Generate separate branches for each field
branches = Expr[]
offset = 0
for fname in field_names
field_spec = fieldtype(S, fname)
field_type = field_spec.parameters[1] # Displacement{3}
entity_type = field_spec.parameters[2]
# Extract quantity type via trait
Q = quantity_type(field_spec) # Vec{3} or Float64
if entity_type === Vertex
if Q === Float64
# Scalar field
value_terms = Expr[]
grad_terms = Expr[]
for i in 1:n_nodes
push!(value_terms, :(Nvals[$i] * u_global[elem.dof_indices[$(offset+i)]]))
push!(grad_terms, :(dN[$i] * u_global[elem.dof_indices[$(offset+i)]]))
end
value_expr = Expr(:call, :+, value_terms...)
grad_expr = Expr(:call, :+, grad_terms...)
push!(branches, quote
if field === $(QuoteNode(fname))
value = $value_expr
grad = $grad_expr
return (value, grad)
end
end)
offset += n_nodes
elseif Q isa UnionAll && Q.body <: Tensor && Q.body.parameters[1] == 1
# Vector field
vec_dim = Q.body.parameters[2]
value_terms = Expr[]
grad_terms = Expr[]
for node in 0:(n_nodes-1)
vec_comps = [:(u_global[elem.dof_indices[$(offset+node*vec_dim+comp)]]) for comp in 1:vec_dim]
u_node = :(Vec{$vec_dim}($(Expr(:tuple, vec_comps...))))
node_idx = node + 1
push!(value_terms, :(Nvals[$node_idx] * $u_node))
push!(grad_terms, :(dN[$node_idx] ⊗ $u_node))
end
value_expr = Expr(:call, :+, value_terms...)
grad_expr = Expr(:call, :+, grad_terms...)
push!(branches, quote
if field === $(QuoteNode(fname))
value = $value_expr
grad = $grad_expr
return (value, grad)
end
end)
offset += n_nodes * vec_dim
for fname in fieldnames(S)
kind, vec_dim = _classify_dofset_field(S, fname)
per_node = _per_node_dof_exprs(:u_global, offset, n_nodes, kind, vec_dim)
value_expr = _value_expr(per_node, :Nvals)
grad_expr = _grad_expr(per_node, :dN, kind)
push!(branches, quote
if field === $(QuoteNode(fname))
value = $value_expr
grad = $grad_expr
return (value, grad)
end
end
end)
offset += _field_dof_count(kind, n_nodes, vec_dim)
end
# Add error case
push!(branches, :(error("Field ", field, " not found in element type $S")))
# Build complete function
return quote
@inbounds begin
Nvals = get_basis_functions($topology, $basis, ξ)
@@ -271,9 +269,8 @@ u_val = interpolate_field_value(elem, u_global, :u, ξ) # Returns Vec{3}
elem::Element{K,P,S,N},
u_global::AbstractVector,
field::Symbol,
ξ::Vec
ξ::Vec,
) where {K,P,S<:DOFSet,N}
field_names = fieldnames(S)
topology = K()
basis = P()
n_nodes = nnodes(topology)
@@ -281,51 +278,19 @@ u_val = interpolate_field_value(elem, u_global, :u, ξ) # Returns Vec{3}
branches = Expr[]
offset = 0
for fname in field_names
field_spec = fieldtype(S, fname)
field_type = field_spec.parameters[1] # Displacement{3}
entity_type = field_spec.parameters[2]
for fname in fieldnames(S)
kind, vec_dim = _classify_dofset_field(S, fname)
per_node = _per_node_dof_exprs(:u_global, offset, n_nodes, kind, vec_dim)
# Extract quantity type via trait
Q = quantity_type(field_spec) # Vec{3} or Float64
value_expr = _value_expr(per_node, :Nvals)
if entity_type === Vertex
if Q === Float64
value_terms = Expr[]
for i in 1:n_nodes
push!(value_terms, :(Nvals[$i] * u_global[elem.dof_indices[$(offset+i)]]))
end
value_expr = Expr(:call, :+, value_terms...)
push!(branches, quote
if field === $(QuoteNode(fname))
return $value_expr
end
end)
offset += n_nodes
elseif quantity_type isa UnionAll && quantity_type.body <: Tensor && quantity_type.body.parameters[1] == 1
vec_dim = quantity_type.body.parameters[2]
value_terms = Expr[]
for node in 0:(n_nodes-1)
vec_comps = [:(u_global[elem.dof_indices[$(offset+node*vec_dim+comp)]]) for comp in 1:vec_dim]
u_node = :(Vec{$vec_dim}($(Expr(:tuple, vec_comps...))))
node_idx = node + 1
push!(value_terms, :(Nvals[$node_idx] * $u_node))
end
value_expr = Expr(:call, :+, value_terms...)
push!(branches, quote
if field === $(QuoteNode(fname))
return $value_expr
end
end)
offset += n_nodes * vec_dim
push!(branches, quote
if field === $(QuoteNode(fname))
return $value_expr
end
end
end)
offset += _field_dof_count(kind, n_nodes, vec_dim)
end
push!(branches, :(error("Field ", field, " not found in element type $S")))
@@ -366,37 +331,18 @@ Returns a NamedTuple where each field is a LocalField containing:
# Unified Dynamic/Quasi-Static Treatment
**Quasi-static:**
Quasi-static:
```julia
local_fields = interpolate_local_fields(elem, u_new, u_old, zero(u_new), Δt, ξ)
# rate = 0, but gradient_rate computed from (∇u_new - ∇u_old)/Δt
```
**Dynamic:**
Dynamic:
```julia
local_fields = interpolate_local_fields(elem, u_new, u_old, u_rate, Δt, ξ)
# rate = u̇, gradient_rate from increments (more accurate than ∇(u̇))
```
# Example
```julia
S = @DOFSet{u::DOF{Displacement{3},Vertex}}
elem = Element{Tetrahedron, Lagrange{Tetrahedron,1}, S}(...)
# Quasi-static loading
u_new = [...] # Current configuration
u_old = [...] # Previous load step
Δt = 1.0
ξ = Vec((0.25, 0.25, 0.25))
local_fields = interpolate_local_fields(elem, u_new, u_old, zero(u_new), Δt, ξ)
# → (u = LocalField(u_val, ∇u, zero(Vec{3}), ∇u_rate), ...)
# Extract strain for material evaluation
ε = extract_strain(local_fields.u.gradient)
ε̇ = extract_strain_rate(local_fields.u.gradient_rate)
```
# Performance
Zero-allocation @generated function. All field access happens at compile time.
"""
@@ -406,121 +352,40 @@ Zero-allocation @generated function. All field access happens at compile time.
u_old::AbstractVector,
u_rate::AbstractVector,
Δt::Float64,
ξ::Vec
ξ::Vec,
) where {K,P,S<:DOFSet,N}
field_names = fieldnames(S)
topology = K()
basis = P()
n_nodes = nnodes(topology)
# Build expressions for LocalField creation for each field
field_exprs = Expr[]
offset = 0
for fname in field_names
field_spec = fieldtype(S, fname)
field_type = field_spec.parameters[1] # Displacement{3}
entity_type = field_spec.parameters[2]
for fname in fieldnames(S)
kind, vec_dim = _classify_dofset_field(S, fname)
# Extract quantity type via trait
Q = quantity_type(field_spec) # Vec{3} or Float64
per_node_new = _per_node_dof_exprs(:u_global, offset, n_nodes, kind, vec_dim)
per_node_old = _per_node_dof_exprs(:u_old, offset, n_nodes, kind, vec_dim)
per_node_rate = _per_node_dof_exprs(:u_rate, offset, n_nodes, kind, vec_dim)
if entity_type === Vertex
if Q === Float64
# Scalar field interpolation
value_terms = Expr[]
grad_terms = Expr[]
value_old_terms = Expr[]
grad_old_terms = Expr[]
rate_terms = Expr[]
value_expr = _value_expr(per_node_new, :Nvals)
grad_expr = _grad_expr(per_node_new, :dN, kind)
grad_old_expr = _grad_expr(per_node_old, :dN, kind)
rate_expr = _value_expr(per_node_rate, :Nvals)
grad_rate_expr = :(($grad_expr - $grad_old_expr) / Δt)
for i in 1:n_nodes
idx = offset + i
# Current value and gradient
push!(value_terms, :(Nvals[$i] * u_global[elem.dof_indices[$idx]]))
push!(grad_terms, :(dN[$i] * u_global[elem.dof_indices[$idx]]))
# Old value and gradient (for gradient_rate)
push!(value_old_terms, :(Nvals[$i] * u_old[elem.dof_indices[$idx]]))
push!(grad_old_terms, :(dN[$i] * u_old[elem.dof_indices[$idx]]))
# Rate
push!(rate_terms, :(Nvals[$i] * u_rate[elem.dof_indices[$idx]]))
end
local_field_expr = :(LocalField($value_expr, $grad_expr, $rate_expr, $grad_rate_expr))
push!(field_exprs, Expr(:(=), fname, local_field_expr))
value_expr = Expr(:call, :+, value_terms...)
grad_expr = Expr(:call, :+, grad_terms...)
grad_old_expr = Expr(:call, :+, grad_old_terms...)
rate_expr = Expr(:call, :+, rate_terms...)
# Gradient rate from increment
grad_rate_expr = :(($grad_expr - $grad_old_expr) / Δt)
# Create LocalField
local_field_expr = :(LocalField($value_expr, $grad_expr, $rate_expr, $grad_rate_expr))
push!(field_exprs, Expr(:(=), fname, local_field_expr))
offset += n_nodes
elseif Q isa UnionAll && Q.body <: Tensor && Q.body.parameters[1] == 1
# Vector field interpolation
vec_dim = Q.body.parameters[2]
value_terms = Expr[]
grad_terms = Expr[]
value_old_terms = Expr[]
grad_old_terms = Expr[]
rate_terms = Expr[]
for node in 0:(n_nodes-1)
node_idx = node + 1
# Current values
vec_comps = [:(u_global[elem.dof_indices[$(offset+node*vec_dim+comp)]]) for comp in 1:vec_dim]
u_node = :(Vec{$vec_dim}($(Expr(:tuple, vec_comps...))))
push!(value_terms, :(Nvals[$node_idx] * $u_node))
push!(grad_terms, :(dN[$node_idx] ⊗ $u_node))
# Old values (for gradient_rate)
vec_comps_old = [:(u_old[elem.dof_indices[$(offset+node*vec_dim+comp)]]) for comp in 1:vec_dim]
u_node_old = :(Vec{$vec_dim}($(Expr(:tuple, vec_comps_old...))))
push!(value_old_terms, :(Nvals[$node_idx] * $u_node_old))
push!(grad_old_terms, :(dN[$node_idx] ⊗ $u_node_old))
# Rate values
vec_comps_rate = [:(u_rate[elem.dof_indices[$(offset+node*vec_dim+comp)]]) for comp in 1:vec_dim]
u_node_rate = :(Vec{$vec_dim}($(Expr(:tuple, vec_comps_rate...))))
push!(rate_terms, :(Nvals[$node_idx] * $u_node_rate))
end
value_expr = Expr(:call, :+, value_terms...)
grad_expr = Expr(:call, :+, grad_terms...)
grad_old_expr = Expr(:call, :+, grad_old_terms...)
rate_expr = Expr(:call, :+, rate_terms...)
# Gradient rate from increment
grad_rate_expr = :(($grad_expr - $grad_old_expr) / Δt)
# Create LocalField
local_field_expr = :(LocalField($value_expr, $grad_expr, $rate_expr, $grad_rate_expr))
push!(field_exprs, Expr(:(=), fname, local_field_expr))
offset += n_nodes * vec_dim
else
error("Unsupported quantity type: $Q")
end
else
error("Unsupported entity type: $entity_type (only Vertex supported for now)")
end
offset += _field_dof_count(kind, n_nodes, vec_dim)
end
nt_expr = Expr(:tuple, field_exprs...)
return quote
@inbounds begin
# Evaluate basis functions once
Nvals = get_basis_functions($topology, $basis, ξ)
dN = get_basis_derivatives($topology, $basis, ξ)
# Return NamedTuple of LocalField
return $nt_expr
end
end