CellMultiValues (new attempt)

docs/src/literate-tutorials/incompressible_elasticity.jl

@@ -23,7 +23,7 @@
 #md # The full program, without comments, can be found in the next
 #md # [section](@ref incompressible_elasticity-plain-program).
 using Ferrite, Tensors
-using BlockArrays, SparseArrays, LinearAlgebra
+using SparseArrays, LinearAlgebra
 
 # First we generate a simple grid, specifying the 4 corners of Cooks membrane.
 function create_cook_grid(nx, ny)
@@ -40,18 +40,17 @@ end;
 
 # Next we define a function to set up our cell- and facevalues.
 function create_values(interpolation_u, interpolation_p)
-    ## quadrature rules
+    ## Quadrature rules
     qr      = QuadratureRule{RefTriangle}(3)
     face_qr = FaceQuadratureRule{RefTriangle}(3)
 
-    ## cell and facevalues for u
-    cellvalues_u = CellValues(qr, interpolation_u)
+    ## Cell values, for both fields
+    cellvalues = MultiCellValues(qr, (u=interpolation_u, p=interpolation_p))    
+
+    ## Face values (only for the displacement, u)
     facevalues_u = FaceValues(face_qr, interpolation_u)
 
-    ## cellvalues for p
-    cellvalues_p = CellValues(qr, interpolation_p)
-
-    return cellvalues_u, cellvalues_p, facevalues_u
+    return cellvalues, facevalues_u
 end;
 
 
@@ -85,28 +84,30 @@ end
 # use a `PseudoBlockArray` from `BlockArrays.jl`.
 
 function doassemble(
-    cellvalues_u::CellValues,
-    cellvalues_p::CellValues,
-    facevalues_u::FaceValues,
-    K::SparseMatrixCSC, grid::Grid, dh::DofHandler, mp::LinearElasticity
-)
+        cellvalues::MultiCellValues, facevalues_u::FaceValues,
+        K::SparseMatrixCSC, grid::Grid, dh::DofHandler, mp::LinearElasticity
+        )
+
     f = zeros(ndofs(dh))
     assembler = start_assemble(K, f)
-    nu = getnbasefunctions(cellvalues_u)
-    np = getnbasefunctions(cellvalues_p)
 
-    fe = PseudoBlockArray(zeros(nu + np), [nu, np]) # local force vector
-    ke = PseudoBlockArray(zeros(nu + np, nu + np), [nu, np], [nu, np]) # local stiffness matrix
+    n = ndofs_per_cell(dh)
+    fe = zeros(n)    # local force vector
+    ke = zeros(n, n) # local stiffness matrix
 
     ## traction vector
     t = Vec{2}((0.0, 1 / 16))
     ## cache ɛdev outside the element routine to avoid some unnecessary allocations
-    ɛdev = [zero(SymmetricTensor{2, 2}) for i in 1:getnbasefunctions(cellvalues_u)]
+    ɛdev = [zero(SymmetricTensor{2, 2}) for i in 1:getnbasefunctions(cellvalues[:u])]
+
+    ## local dof ranges for each field
+    dofrange_u = dof_range(dh, :u)
+    dofrange_p = dof_range(dh, :p)
 
     for cell in CellIterator(dh)
         fill!(ke, 0)
         fill!(fe, 0)
-        assemble_up!(ke, fe, cell, cellvalues_u, cellvalues_p, facevalues_u, grid, mp, ɛdev, t)
+        assemble_up!(ke, fe, cell, cellvalues, facevalues_u, grid, mp, ɛdev, t, dofrange_u, dofrange_p)
         assemble!(assembler, celldofs(cell), fe, ke)
     end
 
@@ -116,37 +117,29 @@ end;
 # The element routine integrates the local stiffness and force vector for all elements.
 # Since the problem results in a symmetric matrix we choose to only assemble the lower part,
 # and then symmetrize it after the loop over the quadrature points.
-function assemble_up!(Ke, fe, cell, cellvalues_u, cellvalues_p, facevalues_u, grid, mp, ɛdev, t)
-
-    n_basefuncs_u = getnbasefunctions(cellvalues_u)
-    n_basefuncs_p = getnbasefunctions(cellvalues_p)
-    u▄, p▄ = 1, 2
-    reinit!(cellvalues_u, cell)
-    reinit!(cellvalues_p, cell)
-
+function assemble_up!(Ke, fe, cell, cellvalues, facevalues_u, grid, mp, ɛdev, t, dofrange_u, dofrange_p)
+    reinit!(cellvalues, cell)
     ## We only assemble lower half triangle of the stiffness matrix and then symmetrize it.
-    for q_point in 1:getnquadpoints(cellvalues_u)
-        for i in 1:n_basefuncs_u
-            ɛdev[i] = dev(symmetric(shape_gradient(cellvalues_u, q_point, i)))
+    for q_point in 1:getnquadpoints(cellvalues)
+        for i in keys(dofrange_u)
+            ɛdev[i] = dev(symmetric(shape_gradient(cellvalues[:u], q_point, i)))
         end
-        dΩ = getdetJdV(cellvalues_u, q_point)
-        for i in 1:n_basefuncs_u
-            divδu = shape_divergence(cellvalues_u, q_point, i)
-            δu = shape_value(cellvalues_u, q_point, i)
-            for j in 1:i
-                Ke[BlockIndex((u▄, u▄), (i, j))] += 2 * mp.G * ɛdev[i] ⊡ ɛdev[j] * dΩ
+        dΩ = getdetJdV(cellvalues, q_point)
+        for (iᵤ, Iᵤ) in pairs(dofrange_u)
+            for (jᵤ, Jᵤ) in pairs(dofrange_u[1:iᵤ])
+                Ke[Iᵤ, Jᵤ] += 2 * mp.G * ɛdev[iᵤ] ⊡ ɛdev[jᵤ] * dΩ
             end
         end
 
-        for i in 1:n_basefuncs_p
-            δp = shape_value(cellvalues_p, q_point, i)
-            for j in 1:n_basefuncs_u
-                divδu = shape_divergence(cellvalues_u, q_point, j)
-                Ke[BlockIndex((p▄, u▄), (i, j))] += -δp * divδu * dΩ
+        for (iₚ, Iₚ) in pairs(dofrange_p)
+            δp = shape_value(cellvalues[:p], q_point, iₚ)
+            for (jᵤ, Jᵤ) in pairs(dofrange_u)
+                divδu = shape_divergence(cellvalues[:u], q_point, jᵤ)
+                Ke[Iₚ, Jᵤ] += -δp * divδu * dΩ
             end
-            for j in 1:i
-                p = shape_value(cellvalues_p, q_point, j)
-                Ke[BlockIndex((p▄, p▄), (i, j))] += - 1 / mp.K * δp * p * dΩ
+            for (jₚ, Jₚ) in pairs(dofrange_p[1:iₚ])
+                p = shape_value(cellvalues[:p], q_point, jₚ)
+                Ke[Iₚ, Jₚ] += - 1 / mp.K * δp * p * dΩ
             end
 
         end
@@ -158,13 +151,13 @@ function assemble_up!(Ke, fe, cell, cellvalues_u, cellvalues_p, facevalues_u, gr
     ## We loop over all the faces in the cell, then check if the face
     ## is in our `"traction"` faceset.
     for face in 1:nfaces(cell)
-        if onboundary(cell, face) && (cellid(cell), face) ∈ getfaceset(grid, "traction")
+        if (cellid(cell), face) ∈ getfaceset(grid, "traction")
             reinit!(facevalues_u, cell, face)
             for q_point in 1:getnquadpoints(facevalues_u)
                 dΓ = getdetJdV(facevalues_u, q_point)
-                for i in 1:n_basefuncs_u
-                    δu = shape_value(facevalues_u, q_point, i)
-                    fe[i] += (δu ⋅ t) * dΓ
+                for (iᵤ, Iᵤ) in pairs(dofrange_u)
+                    δu = shape_value(facevalues_u, q_point, iᵤ)
+                    fe[Iᵤ] += (δu ⋅ t) * dΓ
                 end
             end
         end
@@ -201,8 +194,7 @@ end;
 # this formulation. Therefore we expand the strain to a 3D tensor, and then compute the (3D)
 # stress tensor.
 
-function compute_stresses(cellvalues_u::CellValues, cellvalues_p::CellValues,
-                          dh::DofHandler, mp::LinearElasticity, a::Vector)
+function compute_stresses(cellvalues::MultiCellValues, dh::DofHandler, mp::LinearElasticity, a::Vector)
     ae = zeros(ndofs_per_cell(dh)) # local solution vector
     u_range = dof_range(dh, :u)    # local range of dofs corresponding to u
     p_range = dof_range(dh, :p)    # local range of dofs corresponding to p
@@ -211,20 +203,19 @@ function compute_stresses(cellvalues_u::CellValues, cellvalues_p::CellValues,
     ## Loop over the cells and compute the cell-average stress
     for cc in CellIterator(dh)
         ## Update cellvalues
-        reinit!(cellvalues_u, cc)
-        reinit!(cellvalues_p, cc)
+        reinit!(cellvalues, cc)
         ## Extract the cell local part of the solution
         for (i, I) in pairs(cc.dofs)
             ae[i] = a[I]
         end
         ## Loop over the quadrature points
         σΩi = zero(SymmetricTensor{2, 3}) # stress integrated over the cell
         Ωi = 0.0                          # cell volume (area)
-        for qp in 1:getnquadpoints(cellvalues_u)
-            dΩ = getdetJdV(cellvalues_u, qp)
+        for qp in 1:getnquadpoints(cellvalues)
+            dΩ = getdetJdV(cellvalues, qp)
             ## Evaluate the strain and the pressure
-            ε = function_symmetric_gradient(cellvalues_u, qp, ae, u_range)
-            p = function_value(cellvalues_p, qp, ae, p_range)
+            ε = function_symmetric_gradient(cellvalues[:u], qp, ae, u_range)
+            p = function_value(cellvalues[:p], qp, ae, p_range)
             ## Expand strain to 3D
             ε3D = SymmetricTensor{2, 3}((i, j) -> i < 3 && j < 3 ? ε[i, j] : 0.0)
             ## Compute the stress in this quadrature point
@@ -255,16 +246,16 @@ function solve(ν, interpolation_u, interpolation_p)
     dbc = create_bc(dh)
 
     ## CellValues
-    cellvalues_u, cellvalues_p, facevalues_u = create_values(interpolation_u, interpolation_p)
+    cellvalues, facevalues_u = create_values(interpolation_u, interpolation_p)
 
     ## Assembly and solve
     K = create_sparsity_pattern(dh)
-    K, f = doassemble(cellvalues_u, cellvalues_p, facevalues_u, K, grid, dh, mp)
+    K, f = doassemble(cellvalues, facevalues_u, K, grid, dh, mp)
     apply!(K, f, dbc)
     u = K \ f
 
     ## Compute the stress
-    σ = compute_stresses(cellvalues_u, cellvalues_p, dh, mp, u)
+    σ = compute_stresses(cellvalues, dh, mp, u)
     σvM = map(x -> √(3/2 * dev(x) ⊡ dev(x)), σ) # von Mise effective stress
 
     ## Export the solution and the stress
@@ -276,7 +267,7 @@ function solve(ν, interpolation_u, interpolation_p)
             σij = [x[i, j] for x in σ]
             vtk_cell_data(vtkfile, σij, "sigma_$(i)$(j)")
         end
-        vtk_cell_data(vtkfile, σvM, "sigma von Mise")
+        vtk_cell_data(vtkfile, σvM, "sigma von Mises")
     end
     return u
 end

src/FEValues/CellValues.jl

@@ -127,4 +127,110 @@
     print(io, "\n Function interpolation: "); show(io, d, ip_fun)
     print(io, "\nGeometric interpolation: "); 
     sdim === nothing ? show(io, d, ip_geo) : show(io, d, ip_geo^sdim)
-end
+end
+
+"""
+    MultiCellValues([::Type{T},] quad_rule::QuadratureRule, func_interpols::NamedTuple, [geom_interpol::Interpolation])
+
+A `mcv::MultiCellValues` object generalizes the `CellValues` object to multiple fields. In general, functions applicable to 
+a `CellValues` associated with the function interpolation with `key` can be called on `mcv[key]`, while other functions 
+relating to geometric properties and quadrature rules are called directly on `mcv`. 
+
+**Arguments:**
+* `T`: an optional argument (default to `Float64`) to determine the type the internal data is stored as.
+* `quad_rule`: an instance of a [`QuadratureRule`](@ref)
+* `func_interpols`: A named tuple with entires of type ``Interpolation``, used to interpolate the approximated function identified by the key in `func_interpols`
+* `geom_interpol`: an optional instance of a [`Interpolation`](@ref) which is used to interpolate the geometry.
+  By default linear Lagrange interpolation is used. For embedded elements the geometric interpolations should
+  be vectorized to the spatial dimension.
+"""
+MultiCellValues
+
+struct MultiCellValues{FVS, GM, QR, detT, FVT} <: AbstractCellValues
+    fun_values::FVS         # FunctionValues collected in a named tuple (not necessarily unique)
+    fun_values_tuple::FVT   # FunctionValues collected in a tuple (each unique)
+    geo_mapping::GM         # GeometryMapping
+    qr::QR                  # QuadratureRule
+    detJdV::detT            # AbstractVector{<:Number} or Nothing
+end
+
+function MultiCellValues(::Type{T}, qr::QuadratureRule, ip_funs::NamedTuple, ip_geo::VectorizedInterpolation;
+    update_gradients::Bool = true, update_detJdV::Bool = true) where T 
+
+    FunDiffOrder = convert(Int, update_gradients) # Logic must change when supporting update_hessian kwargs
+    GeoDiffOrder = max(maximum(ip_fun -> required_geo_diff_order(mapping_type(ip_fun), FunDiffOrder), values(ip_funs)), update_detJdV)
+    geo_mapping = GeometryMapping{GeoDiffOrder}(T, ip_geo.ip, qr)
+    unique_ips = unique(values(ip_funs))
+    fun_values_tuple = tuple((FunctionValues{FunDiffOrder}(T, ip_fun, qr, ip_geo) for ip_fun in unique_ips)...)
+    fun_values = NamedTuple((key => fun_values_tuple[findfirst(unique_ip -> ip == unique_ip, unique_ips)] for (key, ip) in pairs(ip_funs)))
+    detJdV = update_detJdV ? fill(T(NaN), length(getweights(qr))) : nothing
+    return MultiCellValues(fun_values, fun_values_tuple, geo_mapping, qr, detJdV)
+end
+
+MultiCellValues(qr::QuadratureRule, ip_funs::NamedTuple, args...; kwargs...) = MultiCellValues(Float64, qr, ip_funs, args...; kwargs...)
+function MultiCellValues(::Type{T}, qr, ip_funs::NamedTuple, ip_geo::ScalarInterpolation=default_geometric_interpolation(first(ip_funs)); kwargs...) where T
+    return MultiCellValues(T, qr, ip_funs, VectorizedInterpolation(ip_geo); kwargs...)
+end
+
+function Base.copy(cv::MultiCellValues)
+    fun_values_tuple = map(copy, cv.fun_values_tuple)
+    fun_values = NamedTuple((key => fun_values_tuple[findfirst(fv === named_fv for fv in cv.fun_values_tuple)] for (key, named_fv) in pairs(cv.fun_values)))
+    return MultiCellValues(fun_values, fun_values_tuple, copy(cv.geo_mapping), copy(cv.qr), _copy_or_nothing(cv.detJdV))
+end
+
+# Access geometry values
+@propagate_inbounds getngeobasefunctions(cv::MultiCellValues) = getngeobasefunctions(cv.geo_mapping)
+@propagate_inbounds geometric_value(cv::MultiCellValues, args...) = geometric_value(cv.geo_mapping, args...)
+geometric_interpolation(cv::MultiCellValues) = geometric_interpolation(cv.geo_mapping)
+
+function getdetJdV(cv::MultiCellValues, q_point::Int)
+    cv.detJdV === nothing && throw(ArgumentError("detJdV calculation was not requested at construction"))
+    return cv.detJdV[q_point]
+end
+
+# No accessors for function values, just ability to get the stored `FunctionValues` which can be called directly. 
+@inline Base.getindex(cv::MultiCellValues, key::Symbol) = cv.fun_values[key]
+
+# Access quadrature rule values 
+getnquadpoints(cv::MultiCellValues) = getnquadpoints(cv.qr)
+
+@inline function reinit!(cv::MultiCellValues, x::AbstractVector)
+    return reinit!(cv, nothing, x)
+end
+
+function reinit!(cv::MultiCellValues, cell::Union{AbstractCell, Nothing}, x::AbstractVector{<:Vec})
+    geo_mapping = cv.geo_mapping
+    fun_values = cv.fun_values_tuple
+    n_geom_basefuncs = getngeobasefunctions(geo_mapping)
+
+    map(fv -> check_reinit_sdim_consistency(:MultiCellValues, shape_gradient_type(fv), eltype(x)), fun_values)
+    if cell === nothing && !all(map(fv -> isa(mapping_type(fv), IdentityMapping), fun_values))
+        throw(ArgumentError("The cell::AbstractCell input is required to reinit! non-identity function mappings"))
+    end
+    if !checkbounds(Bool, x, 1:n_geom_basefuncs) || length(x) != n_geom_basefuncs
+        throw_incompatible_coord_length(length(x), n_geom_basefuncs)
+    end
+
+    @inbounds for (q_point, w) in enumerate(getweights(cv.qr))
+        mapping = calculate_mapping(geo_mapping, q_point, x)
+        _update_detJdV!(cv.detJdV, q_point, w, mapping)
+        _apply_mappings!(fun_values, q_point, mapping, cell)
+    end
+    return nothing
+end
+
+@inline function _apply_mappings!(fun_values::Tuple, q_point, mapping, cell)
+    map(fv -> (@inbounds apply_mapping!(fv, q_point, mapping, cell)), fun_values)
+end
+
+# Slightly faster for unknown reason to write out each call, only worth it for a few unique function values. 
+@inline function _apply_mappings!(fun_values::NTuple{1, <:FunctionValues}, q_point, mapping, cell)
+    @inbounds apply_mapping!(fun_values[1], q_point, mapping, cell)
+end
+
+@inline function _apply_mappings!(fun_values::NTuple{2, <:FunctionValues}, q_point, mapping, cell)
+    @inbounds begin
+        apply_mapping!(fun_values[1], q_point, mapping, cell)
+        apply_mapping!(fun_values[2], q_point, mapping, cell)
+    end
+end

src/FEValues/FunctionValues.jl

@@ -33,7 +33,7 @@ for both the reference cell (precalculated) and the real cell (updated in `reini
 """
 FunctionValues
 
-struct FunctionValues{DiffOrder, IP, N_t, dNdx_t, dNdξ_t}
+struct FunctionValues{DiffOrder, IP, N_t, dNdx_t, dNdξ_t} <: AbstractValues
     ip::IP          # ::Interpolation
     Nx::N_t         # ::AbstractMatrix{Union{<:Tensor,<:Number}}
     Nξ::N_t         # ::AbstractMatrix{Union{<:Tensor,<:Number}}
@@ -83,6 +83,7 @@ function Base.copy(v::FunctionValues)
 end
 
 getnbasefunctions(funvals::FunctionValues) = size(funvals.Nx, 1)
+getnquadpoints(funvals::FunctionValues) = size(funvals.Nx, 2)
 @propagate_inbounds shape_value(funvals::FunctionValues, q_point::Int, base_func::Int) = funvals.Nx[base_func, q_point]
 @propagate_inbounds shape_gradient(funvals::FunctionValues, q_point::Int, base_func::Int) = funvals.dNdx[base_func, q_point]
 @propagate_inbounds shape_symmetric_gradient(funvals::FunctionValues, q_point::Int, base_func::Int) = symmetric(shape_gradient(funvals, q_point, base_func))

src/FEValues/common_values.jl

@@ -26,8 +26,8 @@ end
 end
 
 """
-    reinit!(cv::CellValues, cell::AbstractCell, x::Vector)
-    reinit!(cv::CellValues, x::Vector)
+    reinit!(cv::Union{CellValues, MultiCellValues}, cell::AbstractCell, x::Vector)
+    reinit!(cv::Union{CellValues, MultiCellValues}, x::Vector)
     reinit!(fv::FaceValues, cell::AbstractCell, x::Vector, face::Int)
     reinit!(fv::FaceValues, x::Vector, face::Int)
 

src/exports.jl

@@ -28,6 +28,7 @@ export
     AbstractCellValues,
     AbstractFaceValues,
     CellValues,
+    MultiCellValues,
     FaceValues,
     InterfaceValues,
     reinit!,

src/iterators.jl

@@ -89,7 +89,7 @@ function reinit!(cc::CellCache, i::Int)
 end
 
 # reinit! FEValues with CellCache
-reinit!(cv::CellValues, cc::CellCache) = reinit!(cv, cc.coords)
+reinit!(cv::Union{CellValues, MultiCellValues}, cc::CellCache) = reinit!(cv, cc.coords)
 reinit!(fv::FaceValues, cc::CellCache, f::Int) = reinit!(fv, cc.coords, f) # TODO: Deprecate?
 
 # Accessor functions (TODO: Deprecate? We are so inconsistent with `getxx` vs `xx`...)