diff --git a/examples/cooks_membrane/generate_meshes.sh b/examples/cooks_membrane/generate_meshes.sh new file mode 100755 index 0000000..59acbf6 --- /dev/null +++ b/examples/cooks_membrane/generate_meshes.sh @@ -0,0 +1,2 @@ +gmsh geometry.geo -2 -setnumber element_order 1 +gmsh geometry.geo -2 -setnumber element_order 2 diff --git a/examples/cooks_membrane/geometry.geo b/examples/cooks_membrane/geometry.geo new file mode 100644 index 0000000..3b1e1bc --- /dev/null +++ b/examples/cooks_membrane/geometry.geo @@ -0,0 +1,53 @@ +//////////////////////////////////////////////////////////// +// Cook's membrane (Q1 quadrilateral mesh) +// +// Geometry: +// (0,44) ----------- (48,60) +// | | +// | | +// (0,0) ------------ (48,44) +//////////////////////////////////////////////////////////// + +SetFactory("OpenCASCADE"); + +// Default if nothing is supplied on the command line +DefineConstant[ + element_order = 1 +]; + +// Geometry +Point(1) = {0, 0, 0}; +Point(2) = {48,44, 0}; +Point(3) = {48,60, 0}; +Point(4) = {0,44, 0}; + +Line(1) = {1,2}; +Line(2) = {2,3}; +Line(3) = {3,4}; +Line(4) = {4,1}; + +Curve Loop(1) = {1,2,3,4}; +Plane Surface(1) = {1}; + +// Structured mesh +nx = 32; +ny = 32; + +Transfinite Curve{1,3} = nx + 1; +Transfinite Curve{2,4} = ny + 1; + +Transfinite Surface{1}; +Recombine Surface{1}; + +// Physical groups +Physical Surface("Domain") = {1}; + +Physical Curve("Left") = {4}; +Physical Curve("Right") = {2}; +Physical Curve("Bottom") = {1}; +Physical Curve("Top") = {3}; + +Mesh.ElementOrder = element_order; +Mesh.SecondOrderIncomplete = 0; + +Mesh 2; \ No newline at end of file diff --git a/examples/cooks_membrane/script.jl b/examples/cooks_membrane/script.jl new file mode 100644 index 0000000..eff7bbf --- /dev/null +++ b/examples/cooks_membrane/script.jl @@ -0,0 +1,172 @@ +import FiniteElementContainers as FEC +using FiniteElementContainers +using Gmsh +using StaticArrays +using Tensors + +struct TwoFieldSolidMechanics{NF, NP, NS} <: AbstractPhysics{NF, NP, NS} +end + +struct Displ <: AbstractPhysics{2, 3, 0} +end + +struct Pressure <: AbstractPhysics{1, 3, 0} +end + +function FiniteElementContainers.create_properties(::TwoFieldSolidMechanics) + ρ = 1e3 + K = 1.e9 + G = 1.e6 + return SVector{3, Float64}(ρ, K, G) +end + +function jacobian(∇u) + return det(∇u + one(∇u)) +end + +function pk1_stress_iso(props, ∇u, p) + κ, μ = props[2], props[3] + F = ∇u + one(∇u) + J = det(F) + J_m_13 = 1. / cbrt(J) + J_m_23 = J_m_13 * J_m_13 + I_1 = tr(tdot(F)) + F_inv_T = inv(F)' + P_iso = μ * J_m_23 * (F - (1. / 3.) * I_1 * F_inv_T) + return P_iso +end + +function pk1_stress_vol(props, ∇u, p) + κ, μ = props[2], props[3] + F = ∇u + one(∇u) + J = det(F) + F_inv_T = inv(F)' + # P_vol = 0.5 * κ * (J * J - 1.) * F_inv_T + P_vol = p * J * F_inv_T + return P_vol +end + +material_tangent_iso(props, ∇u, p) = Tensors.gradient(z -> pk1_stress_iso(props, z, p), ∇u) +material_tangent_vol(props, ∇u, p) = Tensors.gradient(z -> pk1_stress_vol(props, z, p), ∇u) + +@inline function FiniteElementContainers.residual( + physics::TwoFieldSolidMechanics, interps, x_el, t, dt, + u_el, u_el_old, state_old_q, state_new_q, props_el +) + u_el, p_el = u_el + interps_u, interps_p = interps + x_el_u, x_el_p = x_el + interps_u = map_interpolants(interps_u, x_el_u) + interps_p = map_interpolants(interps_p, x_el_p) + JxW_u = interps_u.JxW + JxW_p = interps_p.JxW + ∇u_q = interpolate_field_gradients(Displ(), interps_u, u_el) + ∇u_q = modify_field_gradients(PlaneStrain(), ∇u_q) + p_q = interpolate_field_values(Pressure(), interps_p, p_el) + + # constitutive + P_iso = pk1_stress_iso(props, ∇u_q, p_q[1]) + P_vol = pk1_stress_vol(props, ∇u_q, p_q[1]) + J = jacobian(∇u_q) + + P_q = extract_stress(PlaneStrain(), P_iso + P_vol) + G_q = discrete_gradient(PlaneStrain(), interps_u.∇N_X) + R_u = JxW_u * G_q * P_q + R_p = JxW_p * (J - one(J)) * interps_p.N + return R_u, R_p +end + +@inline function FiniteElementContainers.stiffness( + physics::TwoFieldSolidMechanics, interps, x_el, t, dt, + u_el, u_el_old, state_old_q, state_new_q, props_el +) + u_el, p_el = u_el + interps_u, interps_p = interps + x_el_u, x_el_p = x_el + interps_u = map_interpolants(interps_u, x_el_u) + interps_p = map_interpolants(interps_p, x_el_p) + JxW_u = interps_u.JxW + JxW_p = interps_p.JxW + ∇u_q = interpolate_field_gradients(Displ(), interps_u, u_el) + ∇u_q = modify_field_gradients(PlaneStrain(), ∇u_q) + p_q = interpolate_field_values(Pressure(), interps_p, p_el) + J_q = jacobian(∇u_q) + F_q = ∇u_q + one(∇u_q) + F_inv_T_q = inv(F_q)' + dPdp_q = extract_stress(PlaneStrain(), J_q * F_inv_T_q) + A_iso = material_tangent_iso(props, ∇u_q, p_q[1]) + A_vol = material_tangent_vol(props, ∇u_q, p_q[1]) + G_q = discrete_gradient(PlaneStrain(), interps_u.∇N_X) + G_pu_x = J_q .* ( + F_inv_T_q[1, 1] * interps_u.∇N_X[: ,1] + + F_inv_T_q[2, 1] * interps_u.∇N_X[: ,2] + ) + + G_pu_y = J_q * ( + F_inv_T_q[1, 2] * interps_u.∇N_X[:, 1] + + F_inv_T_q[2, 2] * interps_u.∇N_X[:, 2] + ) + Nd = length(G_pu_x) + + tup = MVector{2 * Nd, eltype(G_pu_x)}(undef) + + for i in 1:Nd + tup[2i-1] = G_pu_x[i] + tup[2i] = G_pu_y[i] + end + + G_pu_q = SVector{2 * Nd, eltype(G_pu_x)}(tup) + K_uu = JxW_u * G_q * extract_stiffness(PlaneStrain(), A_iso + A_vol) * G_q' + K_up = JxW_u * G_q * dPdp_q * interps_p.N' + K_pu = JxW_p * interps_p.N * G_pu_q' + K_pp = zero(SMatrix{length(p_el), length(p_el), Float64, length(p_el)^2}) + return ( + (K_uu, K_up), + (K_pu, K_pp) + ) +end + +mesh_u = UnstructuredMesh(Base.source_dir() * "/geometry_q2.geo") +mesh_p = UnstructuredMesh(Base.source_dir() * "/geometry_q1.geo") + +V_u = FunctionSpace(mesh_u, H1Field, Lagrange) +V_p = FunctionSpace(mesh_p, H1Field, Lagrange) +# V_J = FunctionSpace(mesh_p, H1Field, Lagrange) + +u = VectorFunction(V_u, "displ") +p = ScalarFunction(V_p, "pressure") +# J = ScalarFunction(V_J, "jacobian") + +zero_func(_, _) = 0.0 +displ_func(_, t) = 0.1 * t +dbcs_u = DirichletBC[ + DirichletBC("displ_x", zero_func; nodeset_name = "Left") + DirichletBC("displ_y", zero_func; nodeset_name = "Left") + DirichletBC("displ_x", zero_func; nodeset_name = "Right") + DirichletBC("displ_y", displ_func; nodeset_name = "Right") +] +physics = TwoFieldSolidMechanics{3, 0, 0}() +props = create_properties(physics) +times = TimeStepper(0.0, 1.0, 20) + +dof_u, dof_p = DofManager(u), DofManager(p) +dof = (dof_u, dof_p) +# dof_u, dof_p, dof_J = DofManager(u), DofManager(p), DofManager(J) +# dof = (dof_u, dof_p, dof_J) +asm = FEC.BlockSparseMatrixAssembler(dof) + +p_u = create_parameters(mesh_u, SparseMatrixAssembler(dof_u), physics, props; dirichlet_bcs = dbcs_u, times = times) +p_p = create_parameters(mesh_p, SparseMatrixAssembler(dof_p), physics, props; times = times) +params = (p_u, p_p) + +FEC.update_dofs!( + asm, + (p_u.dirichlet_bcs, p_p.dirichlet_bcs), + (p_u.periodic_bcs, p_p.periodic_bcs) +) + +Uu = create_unknowns(asm) +U = create_field(asm) + +assemble_vector!(asm, residual, Uu, params) +assemble_matrix!(asm, stiffness, Uu, params) \ No newline at end of file diff --git a/src/Enzyme.jl b/src/Enzyme.jl index 85e6653..fb6d6a9 100644 --- a/src/Enzyme.jl +++ b/src/Enzyme.jl @@ -144,9 +144,7 @@ function _assemble_vector_block_enzyme_safe!( for e in 1:conns_all.nelems[b] conn = connectivity(ref_fe, conns, e, coffset) x_el, u_el, u_el_old = element_level_fields(ref_fe, conn, X, U, U_old) - props_el = properties(props, e, b) - # # val_el = _element_scratch(return_type, ref_fe, U) for q in 1:num_cell_quadrature_points(ref_fe) interps = _cell_interpolants(ref_fe, q) diff --git a/src/FiniteElementContainers.jl b/src/FiniteElementContainers.jl index 77d56bc..f418339 100644 --- a/src/FiniteElementContainers.jl +++ b/src/FiniteElementContainers.jl @@ -222,6 +222,7 @@ include("Constraints.jl") include("InitialConditions.jl") include("Formulations.jl") +include("Gather.jl") include("Physics.jl") include("assemblers/Assemblers.jl") # diff --git a/src/FunctionSpaces.jl b/src/FunctionSpaces.jl index 11c0329..5a0b036 100644 --- a/src/FunctionSpaces.jl +++ b/src/FunctionSpaces.jl @@ -42,6 +42,8 @@ const _el_name_to_juliac_safe_id = Dict{String, Int}( "TETRA4" => 6, "TETRA10" => 7 ) +# Lagrange elements +# defaulting to fully integrated elements for now const _juliac_safe_ref_fes = ( ReferenceFE(Hex{Lagrange, 1}(), GaussLobattoLegendre(2, 2)), # HEX8 for Lagrange ReferenceFE(Quad{Lagrange, 1}(), GaussLobattoLegendre(2, 2)), # QUAD4 for Lagrange @@ -137,7 +139,6 @@ struct FunctionSpace{ block_to_ref_fe_id::BTRE coords::H1Field{RT, RV, ND} elem_conns::Connectivity{IT, IV} - # need to remove this mabye? elem_id_maps::Vector{Vector{IT}} # TODO create new type for ID map similar to connectivity elem_to_facets::Union{Nothing, Connectivity{IT, IV}} facet_orientation::Union{Nothing, Connectivity{IT, IV}} @@ -146,7 +147,7 @@ struct FunctionSpace{ ref_fes::RefFEs function FunctionSpace{is_juliac_safe, FT}( - block_names, block_to_ref_fe_id, coords, elem_conns, + block_names, block_to_ref_fe_id, coords, elem_conns, elem_id_maps, elem_to_facets, facet_orientation, node_id_map, ref_fes ) where {is_juliac_safe, FT} new{ @@ -272,7 +273,7 @@ function block_entity_size(fspace::FunctionSpace, b::Int) return (num_entities_per_element(fspace, b), num_elements(fspace, b)) end -function block_reference_element(fspace::FunctionSpace{false, FT, I, V, BTRE, C, R}, block_id::Int) where {I, FT, V, BTRE, C, R} +function block_reference_element(fspace::FunctionSpace{false, FT, I, V, BTRE, C, R}, block_id::Int) where {FT, I, V, BTRE, C, R} return fspace.ref_fes[block_id] end @@ -360,7 +361,6 @@ function num_entities(fspace::FunctionSpace) end function num_entities_per_element(fspace::FunctionSpace, b::Int) - return num_entities_per_element(fspace.elem_conns, b) if _field_type(fspace) == L2Field return block_quadrature_size(fspace, b)[1] else diff --git a/src/Gather.jl b/src/Gather.jl new file mode 100644 index 0000000..fb76223 --- /dev/null +++ b/src/Gather.jl @@ -0,0 +1,115 @@ +# gathers and interpolates field u +# at quadrature points on block b, +# element e, quadrature point q. +# TODO call in shape_function_values +# and write method like shape_function_values(ref_fe, n, q) +# to get a specific basis function at a specific q point +@generated function interpolate_values( + u::H1Field{T, D, NF}, x::H1Field{T, D, ND}, + conns, ref_fe, + q, e, b, +) where {T, D, NF, ND} + # create NF accumulator values + sums = [Symbol(:u_, i) for i = 1:NF] + # set them all to zero with identical type to u + init = [:($(s) = zero(T)) for s in sums] + # create update methods + updates = [ + :($(sums[i]) += Ni * u[$i, idx]) + for i = 1:NF + ] + ret = :(SVector{$NF, T}($(sums...))) + + quote + $(init...) + N = ref_fe.cell_interps.values + @inbounds for n in axes(N, 2) + idx = connectivity(conns, n, e, b) + Ni = N[n, q] + $(updates...) + end + $ret + end +end + +@generated function mapping_jacobian( + x::H1Field{T, D, ND}, + conns, ref_fe, + q, e, b, +) where {T, D, ND} + + entries = Expr[] + for i in 1:ND, j in 1:ND + push!(entries, + quote + let s = zero(T) + @inbounds for n in axes(dN, 2) + idx = connectivity(conns, n, e, b) + s += x[$i, idx] * dN[$j, n, q] + end + s + end + end) + end + + quote + dN = ref_fe.cell_interps.gradients + SMatrix{$ND, $ND, T, $(ND * ND)}( + $(entries...) + ) + end +end + +# TODO not finished... +# need to write jacobian mapping operator +# @generated function interpolate_gradients( +# u::H1Field{T, D, NF}, x::H1Field{T, D, NF}, +# conns, ref_fe, +# q, e, b +# ) where {T, D, NF} +# dim = ReferenceFiniteElements.dimension(ref_fe) +# sums = [Symbol(:u_, j, :_, i) for j = 1:dim, i = 1:NF] +# init = [:($(s) = zero(T)) for s in sums] +# updates = [ +# :($(sums[j, i]) += ∇N_ji * u[$i, idx]) +# for j = 1:dim, i = 1:NF +# ] +# ret = :(SMatrix{$dim, $NF, T $dim * $NF})($(sums...)) +# quote +# $(init...) +# ∇N_ξ +# end +# end + +@inline function interpolate_gradient( + u::H1Field{T, D, NF}, x::H1Field{T, D, ND}, + conns, ref_fe, + q, e, b, +) where {T, D, NF, ND} + dN = ref_fe.cell_interps.gradients + J = mapping_jacobian(x, conns, ref_fe, q, e, b) + # inverse Jacobian (computed once) + invJ = inv(J) + + # result: ND x NF matrix + return SMatrix{ND, NF, T}(ntuple(Val(ND * NF)) do k + i = (k - 1) ÷ NF + 1 # spatial direction + f = (k - 1) % NF + 1 # field component + + # compute reference gradient component ∂u_f / ∂ξ_i + du_dξ = zero(T) + + @inbounds for n in axes(dN, 2) + idx = connectivity(conns, n, e, b) + du_dξ += dN[i, n, q] * u[f, idx] + end + + # map to physical space: (J^{-1})_{i,j} ∂u/∂ξ_j + du_dx = zero(T) + @inbounds for j in 1:ND + du_dx += invJ[j, i] * du_dξ + end + + du_dx + end) +end diff --git a/src/assemblers/Assemblers.jl b/src/assemblers/Assemblers.jl index 9a34da8..8cadf0d 100644 --- a/src/assemblers/Assemblers.jl +++ b/src/assemblers/Assemblers.jl @@ -2,7 +2,7 @@ $(TYPEDEF) $(TYPEDFIELDS) """ -abstract type AbstractAssembler{Dof <: DofManager} end +abstract type AbstractAssembler end """ $(TYPEDSIGNATURES) """ @@ -60,7 +60,7 @@ end # that the provided storage came from either a sparse matrix or vector # if you pass a vector expecting it to act like a dense vector, it WON'T! # for IndexedAssembledReturnType, does nothing -function _assemble_element!( +@inline function _assemble_element!( storage, val_q, conns, # all connectivities for this element ::Int @@ -69,7 +69,7 @@ function _assemble_element!( end # assembled vector where storage is a field -function _assemble_element!( +@inline function _assemble_element!( storage::AbstractField, R_el::SVector{NDOF, T}, conns, # all connectivities for this element ::Int @@ -88,7 +88,7 @@ end # sparse vector attempt # this one is good for pbcs -function _assemble_element!( +@inline function _assemble_element!( storage::AbstractVector, R_el::SVector{NDOF, T}, conns, el_id::Int @@ -106,19 +106,17 @@ end # TODO we'll need a regular matrix implementation # as well (Can we live with 1?) # sparse matrix -function _assemble_element!( +@inline function _assemble_element!( storage, K_el::SMatrix{NDOF1, NDOF2, T, NDOF1xNDOF2}, conns, # all connectivities for this element el_id::Int ) where {NDOF1, NDOF2, T, NDOF1xNDOF2} - # figure out ids needed to update - start_id = (el_id - 1) * NDOF1xNDOF2 + 1 - end_id = start_id + NDOF1xNDOF2 - 1 - ids = start_id:end_id - - # get appropriate storage and update values - for (i, id) in enumerate(ids) - storage[id] = K_el.data[i] + base = (el_id - 1) * NDOF1xNDOF2 + for j in 1:NDOF2 + for i in 1:NDOF1 + idx = base + (j - 1) * NDOF1 + i + storage[idx] = K_el[i, j] + end end return nothing end @@ -137,34 +135,14 @@ create_unknowns(asm::AbstractAssembler) = create_unknowns(asm.dof) return @views conns[:, e] end -""" -$(TYPEDSIGNATURES) -""" -@inline function _element_level_fields(U::H1Field{T, D, NF}, ref_fe, conns, e) where {T, D, NF} - NNPE = ReferenceFiniteElements.num_vertices(ref_fe) - NxNDof = NNPE * NF - u_el = @views SMatrix{NNPE, NF, eltype(U), NxNDof}(U[:, conns[:, e]]) - return u_el -end - @inline function _element_level_fields(U::H1Field{T, D, NF}, ref_fe, conns) where {T, D, NF} NNPE = ReferenceFiniteElements.num_cell_dofs(ref_fe) NxNDof = NNPE * NF u_el = @views SMatrix{NF, NNPE, eltype(U), NxNDof}(U[:, conns]) - # u_el = @views SMatrix{NNPE, ND, eltype(U), NxNDof}(U[:, conns]) - return u_el -end - -""" -$(TYPEDSIGNATURES) -""" -@inline function _element_level_fields_flat(U::H1Field{T, D, NF}, ref_fe, conns, e) where {T, D, NF} - NNPE = ReferenceFiniteElements.num_cell_dofs(ref_fe) - NxNDof = NNPE * NF - u_el = @views SVector{NxNDof, eltype(U)}(U[:, conns[:, e]]) return u_el end +# Used by source currently @inline function _element_level_fields_flat(U::H1Field{T, D, NF}, ref_fe, conns) where {T, D, NF} NNPE = ReferenceFiniteElements.num_cell_dofs(ref_fe) NxNDof = NNPE * NF @@ -172,6 +150,13 @@ end return u_el end +# @inline function _element_level_fields_flat(U::H1Field{T, D, NF}, ref_fe, conns, e) where {T, D, NF} +# NNPE = ReferenceFiniteElements.num_cell_dofs(ref_fe) +# NxNDof = NNPE * NF +# u_el = @views SVector{NxNDof, eltype(U)}(U[:, conns]) +# return u_el +# end + @inline function element_level_fields(ref_fe, conn, X, U, U_old) x_el = _element_level_fields_flat(X, ref_fe, conn) u_el = _element_level_fields_flat(U, ref_fe, conn) @@ -197,6 +182,18 @@ $(TYPEDSIGNATURES) return K_el end +@inline function _element_scratch( + ::AssembledMatrix, + ref_fe_row, + U_row::H1Field{T, D, NFr}, + ref_fe_col, + U_col::H1Field{T, D, NFc}, +) where {T, D, NFr, NFc} + nr = num_cell_dofs(ref_fe_row) * NFr + nc = num_cell_dofs(ref_fe_col) * NFc + return zeros(SMatrix{nr, nc, T, nr * nc}) +end + """ $(TYPEDSIGNATURES) """ @@ -480,6 +477,7 @@ include("SparsityPatterns.jl") # types include("MatrixFreeAssembler.jl") include("SparseMatrixAssembler.jl") +include("BlockMatrixAssemblers.jl") # methods include("Diagonal.jl") diff --git a/src/assemblers/BlockMatrixAssemblers.jl b/src/assemblers/BlockMatrixAssemblers.jl new file mode 100644 index 0000000..c349fd6 --- /dev/null +++ b/src/assemblers/BlockMatrixAssemblers.jl @@ -0,0 +1,225 @@ +abstract type AbstractBlockAssembler <: AbstractAssembler end + +function create_field(asm::AbstractBlockAssembler) + return create_field(asm.dof) +end + +function create_unknowns(asm::AbstractBlockAssembler) + return create_unknowns(asm.dof) +end + +struct BlockSparseMatrixAssembler{ + I <: AbstractVector{Int}, + R <: AbstractVector{Float64}, + D, + F, + U <: BlockedVector, + S +} <: AbstractBlockAssembler + dof::D + matrix_patterns::Matrix{SparseMatrixPattern{I, R}} + vector_patterns::Vector{SparseVectorPattern{I}} + residual_storage::F + residual_unknowns::U + stiffness_storage::S +end + +function BlockSparseMatrixAssembler(dof::Tuple) + matrix_patterns = Matrix{SparseMatrixPattern{Vector{Int}, Vector{Float64}}}(undef, length(dof), length(dof)) + vector_patterns = Vector{SparseVectorPattern{Vector{Int}}}(undef, length(dof)) + for i in 1:length(dof) + for j in 1:length(dof) + if i == 1 + pattern = SparseMatrixPattern(dof[i]) + else + pattern = SparseMatrixPattern(dof[i], dof[j]) + end + matrix_patterns[i, j] = pattern + end + vector_patterns[i] = SparseVectorPattern(dof[i]) + end + # n_matrix_entries = matrix_free ? 0 : num_entries(matrix_pattern) + n_matrix_entries = map(num_entries, matrix_patterns) + residual = create_field(dof) + residual_unknowns = create_unknowns(dof) + # stiffness_storage = zeros(n_matrix_entries) + stiffness_storage = map(zeros, n_matrix_entries) + return BlockSparseMatrixAssembler( + dof, matrix_patterns, vector_patterns, + residual, residual_unknowns, + stiffness_storage + ) +end + +function Base.show(io::IO, asm::BlockSparseMatrixAssembler) + sz = size(asm.matrix_patterns) + println(io, "BlockSparseMatrixAssembler:") + println(io, " Block layout = $(sz[1]) x $(sz[2])") + for dof in asm.dof + show(io, dof; pad = " ") + println(io) + end + println(" Matrix sizes:") + for i in axes(asm.matrix_patterns, 1) + string = " " + for j in axes(asm.matrix_patterns, 2) + string = string * "($(length(asm.dof[i].unknown_dofs)), $(length(asm.dof[j].unknown_dofs)))" + if j < size(asm.matrix_patterns, 2) + string = string * ", " + end + end + println(io, string) + end + # println(io, " Block variables = ") +end + +function assemble_matrix!( + assembler::BlockSparseMatrixAssembler, func::F, Uu, p +) where F <: Function + @assert length(assembler.dof) == 2 "Only two spaces supported currently" + storage = assembler.residual_storage + map(x -> fill!(x, zero(eltype(x))), storage) + fspace = map(function_space, assembler.dof) + X = map(coordinates, p) + # should we do a check that all times, and time steps are consistent? + t = current_time(p[1]) + Δt = time_step(p[1]) + U = map(x -> x.field, p) + U_old = map(x -> x.field_old, p) + + for sol_id in 1:length(fspace) + _update_for_assembly!(p[sol_id], assembler.dof[sol_id], Uu[BlockArrays.Block(sol_id)]) + end + + return_type = AssembledMatrix() + conns = map(x -> x.elem_conns.data, fspace) + coffsets = map(x -> x.elem_conns.offsets, fspace) + physics = p[1].physics + props = p[1].properties + for b in 1:num_blocks(fspace[1]) + block_physics = values(physics)[b] + ref_fe = map(x -> block_reference_element(x, b), fspace) + num_q_pts = map(num_cell_quadrature_points, ref_fe) + @assert all(==(num_q_pts[1]), num_q_pts) + num_q_pts = num_q_pts[1] + state_old = block_view(p[1].state_old, b) + state_new = block_view(p[1].state_new, b) + for e in 1:block_entity_size(fspace[1], b)[2] + conn = map((r, c, co) -> connectivity(r, c, e, co[b]), ref_fe, conns, coffsets) + out = map((r, c, x, u, u_old) -> element_level_fields(r, c, e, x, u, u_old), ref_fe, conn, X, U, U_old) + x_el = map(x -> x[1], out) + u_el = map(x -> x[2], out) + u_el_old = map(x -> x[3], out) + props_el = _element_level_properties(values(props)[b], e) + # val_el = map((r, u) -> _element_scratch(return_type, r, u), ref_fe, U) + nfields = length(U) + + val_el = ntuple(i -> ntuple(j->begin + _element_scratch( + return_type, + ref_fe[i], U[i], + ref_fe[j], U[j] + ) + end, nfields), nfields) + for q in 1:num_q_pts + interps = map(r -> _cell_interpolants(r, q), ref_fe) + state_old_q = _quadrature_level_state(state_old, q, e) + state_new_q = _quadrature_level_state(state_new, q, e) + val_q = func(block_physics, interps, x_el, t, Δt, u_el, u_el_old, state_old_q, state_new_q, props_el) + # val_el = map((f, vq, ve) -> _accumulate_q_value(return_type, f, vq, ve, q, e), U, val_q, val_el) + val_el = map( + (vq1, ve1) -> + map((f, vq, ve) -> _accumulate_q_value(return_type, f, vq, ve, q, e), U, vq1, ve1), + val_q, val_el + ) + end + # map((f, v, c) -> _assemble_element!(f, v, c, e), U, val_el, conn, e) + for i in 1:nfields + for j in 1:nfields + _assemble_element!(assembler.stiffness_storage[i, j], val_el[i][j], conn[i], e) + end + end + end + end +end + +function assemble_vector!( + assembler::BlockSparseMatrixAssembler, func::F, Uu, p +) where F <: Function + @assert length(assembler.dof) == 2 "Only two spaces supported currently" + storage = assembler.residual_storage + map(x -> fill!(x, zero(eltype(x))), storage) + fspace = map(function_space, assembler.dof) + X = map(coordinates, p) + # should we do a check that all times, and time steps are consistent? + t = current_time(p[1]) + display(t) + Δt = time_step(p[1]) + U = map(x -> x.field, p) + U_old = map(x -> x.field_old, p) + + for sol_id in 1:length(fspace) + _update_for_assembly!(p[sol_id], assembler.dof[sol_id], Uu[BlockArrays.Block(sol_id)]) + end + + return_type = AssembledVector() + conns = map(x -> x.elem_conns.data, fspace) + coffsets = map(x -> x.elem_conns.offsets, fspace) + physics = p[1].physics + props = p[1].properties + for b in 1:num_blocks(fspace[1]) + block_physics = values(physics)[b] + ref_fe = map(x -> block_reference_element(x, b), fspace) + num_q_pts = map(num_cell_quadrature_points, ref_fe) + @assert all(==(num_q_pts[1]), num_q_pts) + num_q_pts = num_q_pts[1] + state_old = block_view(p[1].state_old, b) + state_new = block_view(p[1].state_new, b) + for e in 1:block_entity_size(fspace[1], b)[2] + conn = map((r, c, co) -> connectivity(r, c, e, co[b]), ref_fe, conns, coffsets) + out = map((r, c, x, u, u_old) -> element_level_fields(r, c, e, x, u, u_old), ref_fe, conn, X, U, U_old) + x_el = map(x -> x[1], out) + u_el = map(x -> x[2], out) + u_el_old = map(x -> x[3], out) + props_el = _element_level_properties(values(props)[b], e) + val_el = map((r, u) -> _element_scratch(return_type, r, u), ref_fe, U) + for q in 1:num_q_pts + interps = map(r -> _cell_interpolants(r, q), ref_fe) + state_old_q = _quadrature_level_state(state_old, q, e) + state_new_q = _quadrature_level_state(state_new, q, e) + val_q = func(block_physics, interps, x_el, t, Δt, u_el, u_el_old, state_old_q, state_new_q, props_el) + val_el = map((f, vq, ve) -> _accumulate_q_value(return_type, f, vq, ve, q, e), U, val_q, val_el) + end + # out = map((f, v, c) -> _assemble_element!(f, v, c, e), U, val_el, conn) + for sol_id in 1:length(U) + _assemble_element!(U[sol_id], val_el[sol_id], conn[sol_id], e) + end + end + end +end + +# this won't work with condensed right now +function update_dofs!( + assembler::BlockSparseMatrixAssembler, dirichlet_bcs, periodic_bcs +) + ddofs = map(dirichlet_dofs, dirichlet_bcs) + pdofs = map(periodic_dofs, periodic_bcs) + pdofs_side_a = map(x -> x[1], pdofs) + pdofs_side_b = map(x -> x[2], pdofs) + + # update dof managers first + for (n, dof) in enumerate(assembler.dof) + update_dofs!(dof, ddofs[n], pdofs_side_a[n], pdofs_side_b[n]) + end + + # now update sparsity patterns + for i in axes(assembler.matrix_patterns, 1) + for j in axes(assembler.matrix_patterns, 2) + _update_dofs!( + assembler.matrix_patterns[i, j], + assembler.dof[i], ddofs[i], pdofs_side_b[i], + assembler.dof[j], ddofs[j], pdofs_side_b[j] + ) + end + end +end diff --git a/src/assemblers/MatrixFreeAssembler.jl b/src/assemblers/MatrixFreeAssembler.jl index 1834601..ca0e1ae 100644 --- a/src/assemblers/MatrixFreeAssembler.jl +++ b/src/assemblers/MatrixFreeAssembler.jl @@ -5,7 +5,7 @@ struct MatrixFreeAssembler{ RV <: AbstractArray{Float64, 1}, Var <: AbstractFunction, FieldStorage <: AbstractField{Float64, NumArrDims, RV} -} <: AbstractAssembler{DofManager{Condensed, Int, IV, Var}} +} <: AbstractAssembler dof::DofManager{Condensed, Int, IV, Var} vector_pattern::SparseVectorPattern{IV} constraint_storage::RV diff --git a/src/assemblers/SparseMatrixAssembler.jl b/src/assemblers/SparseMatrixAssembler.jl index 4367a80..2ea9ab5 100644 --- a/src/assemblers/SparseMatrixAssembler.jl +++ b/src/assemblers/SparseMatrixAssembler.jl @@ -13,7 +13,7 @@ struct SparseMatrixAssembler{ RV <: AbstractArray{Float64, 1}, Var <: AbstractFunction, FieldStorage -} <: AbstractAssembler{DofManager{Condensed, Int, IV, Var}} +} <: AbstractAssembler dof::DofManager{Condensed, Int, IV, Var} matrix_pattern::SparseMatrixPattern{IV, RV} vector_pattern::SparseVectorPattern{IV} @@ -301,7 +301,6 @@ function update_dofs!(assembler::AbstractAssembler, dirichlet_bcs::DirichletBCs, assembler.dof, ddofs, pdofs_side_b, assembler.dof, ddofs, pdofs_side_b ) - end _update_dofs!(assembler.vector_pattern, assembler.dof, ddofs, pdofs_side_b) end diff --git a/src/meshes/Exodus.jl b/src/meshes/Exodus.jl index b08b61e..4c10b74 100644 --- a/src/meshes/Exodus.jl +++ b/src/meshes/Exodus.jl @@ -183,7 +183,7 @@ function PostProcessor( all_el_var_names = element_var_names append!(all_el_var_names, quadrature_var_names) - # TODO need to add all the quadrature values labelled by block id + # TODO need to add all the quadrature values labelled by Exodus.Block id exo = ExodusDatabase(file_name, "rw") @@ -270,9 +270,9 @@ end function write_field(pp::PostProcessor, time_index::Int, field_names, field::NamedTuple) @assert length(field_names) == length(field) field_names = String.(field_names) - for (block, val) in field + for (Exodus.Block, val) in field for name in field_names - # write_values(pp.field_output_db, ElementVariable, time_index, block, name, val) + # write_values(pp.field_output_db, ElementVariable, time_index, Exodus.Block, name, val) end end end diff --git a/test/TestAssemblers.jl b/test/TestAssemblers.jl index 3df96cf..1a2e37e 100644 --- a/test/TestAssemblers.jl +++ b/test/TestAssemblers.jl @@ -2,6 +2,7 @@ using Adapt if "--test-amdgpu" in ARGS @eval using AMDGPU end if "--test-cuda" in ARGS @eval using CUDA end + using KernelAbstractions using LinearAlgebra using SparseArrays using SparseMatricesCSR @@ -41,6 +42,7 @@ end @testsnippet AssemblerHelperPoisson begin if "--test-amdgpu" in ARGS @eval using AMDGPU end if "--test-cuda" in ARGS @eval using CUDA end + using KernelAbstractions using LinearAlgebra using StaticArrays include("poisson/TestPoissonCommon.jl") @@ -60,6 +62,7 @@ end @testsnippet AssemblerHelperMechanics begin if "--test-amdgpu" in ARGS @eval using AMDGPU end if "--test-cuda" in ARGS @eval using CUDA end + using KernelAbstractions using StaticArrays using Tensors include("mechanics/TestMechanicsCommon.jl") @@ -121,7 +124,9 @@ end # test vector consistency assemble_vector!(asm_1, residual, U_1, p_1) + KA.synchronize(KA.get_backend(asm_1)) assemble_vector!(asm_2, residual!, U_2, p_2) + KA.synchronize(KA.get_backend(asm_2)) R_1 = residual(asm_1) |> copy R_2 = residual(asm_2) |> copy if dev != cpu