diff --git a/src/elements/interpolate.jl b/src/elements/interpolate.jl index 8a7593b..efdcc25 100644 --- a/src/elements/interpolate.jl +++ b/src/elements/interpolate.jl @@ -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