From 115c6fc7c85e7782408c41c71ce26f8793b9fa6c Mon Sep 17 00:00:00 2001 From: "Craig M. Hamel" Date: Wed, 1 Jul 2026 21:08:35 -0400 Subject: [PATCH 1/2] Some work towards mixed space assembly. Its really hacky right now and using maximum code re-use. Definitely not efficient are ready for prime time though. --- Project.toml | 5 +- examples/cooks_membrane/geometry_q1.geo | 48 ++++ examples/cooks_membrane/geometry_q2.geo | 44 ++++ examples/cooks_membrane/script.jl | 172 +++++++++++++ ext/GmshExt.jl | 190 +++++++++++++-- src/DofManagers.jl | 57 +++-- src/Enzyme.jl | 4 +- src/FiniteElementContainers.jl | 1 + src/FunctionSpaces.jl | 308 ++++++++++-------------- src/assemblers/Assemblers.jl | 80 +++--- src/assemblers/BlockMatrixAssemblers.jl | 225 +++++++++++++++++ src/assemblers/MatrixAction.jl | 6 +- src/assemblers/MatrixFreeAssembler.jl | 2 +- src/assemblers/SparseMatrixAssembler.jl | 8 +- src/assemblers/SparsityPatterns.jl | 153 ++++++++++-- src/bcs/PeriodicBCs.jl | 2 +- src/meshes/Exodus.jl | 10 +- src/meshes/Meshes.jl | 6 +- src/meshes/Toplogy.jl | 123 ++++++++++ test/poisson/TestPoisson.jl | 59 +++++ 20 files changed, 1201 insertions(+), 302 deletions(-) create mode 100644 examples/cooks_membrane/geometry_q1.geo create mode 100644 examples/cooks_membrane/geometry_q2.geo create mode 100644 examples/cooks_membrane/script.jl create mode 100644 src/assemblers/BlockMatrixAssemblers.jl create mode 100644 src/meshes/Toplogy.jl diff --git a/Project.toml b/Project.toml index 18375ea8..530607ac 100644 --- a/Project.toml +++ b/Project.toml @@ -1,12 +1,13 @@ name = "FiniteElementContainers" uuid = "d08262e4-672f-4e7f-a976-f2cea5767631" -version = "0.14.2" +version = "0.14.3" authors = ["Craig M. Hamel and contributors"] [deps] AcceleratedKernels = "6a4ca0a5-0e36-4168-a932-d9be78d558f1" Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" Atomix = "a9b6321e-bd34-4604-b9c9-b65b8de01458" +BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" DynamicExpressions = "a40a106e-89c9-4ca8-8020-a735e8728b6b" Exodus = "f57ae99e-f805-4780-bdca-96e224be1e5a" @@ -26,7 +27,6 @@ TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f" [weakdeps] AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" AbaqusReader = "bc6b9049-e460-56d6-94b4-a597b2c0390d" -BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" Gmsh = "705231aa-382f-11e9-3f0c-b7cb4346fdeb" Metis = "2679e427-3c69-5b7f-982b-ece356f1e94b" @@ -35,7 +35,6 @@ PartitionedArrays = "5a9dfac6-5c52-46f7-8278-5e2210713be9" [extensions] AMDGPUExt = "AMDGPU" AbaqusReaderExt = "AbaqusReader" -BlockArraysExt = "BlockArrays" CUDAExt = "CUDA" GmshExt = "Gmsh" MetisExt = "Metis" diff --git a/examples/cooks_membrane/geometry_q1.geo b/examples/cooks_membrane/geometry_q1.geo new file mode 100644 index 00000000..0b298b65 --- /dev/null +++ b/examples/cooks_membrane/geometry_q1.geo @@ -0,0 +1,48 @@ +//////////////////////////////////////////////////////////// +// Cook's membrane (Q1 quadrilateral mesh) +// +// Geometry: +// (0,44) ----------- (48,60) +// | | +// | | +// (0,0) ------------ (48,44) +//////////////////////////////////////////////////////////// + +SetFactory("OpenCASCADE"); + +// 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 = 1; +Mesh.SecondOrderIncomplete = 0; + +Mesh 2; \ No newline at end of file diff --git a/examples/cooks_membrane/geometry_q2.geo b/examples/cooks_membrane/geometry_q2.geo new file mode 100644 index 00000000..986b3f59 --- /dev/null +++ b/examples/cooks_membrane/geometry_q2.geo @@ -0,0 +1,44 @@ +//////////////////////////////////////////////////////////// +// Cook's membrane (Q2 quadrilateral mesh) +//////////////////////////////////////////////////////////// + +SetFactory("OpenCASCADE"); + +// 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}; + +// Second-order quadrilateral elements (9-node) +Mesh.ElementOrder = 2; +Mesh.SecondOrderIncomplete = 0; +Mesh.HighOrderOptimize = 1; + +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 00000000..eff7bbf8 --- /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/ext/GmshExt.jl b/ext/GmshExt.jl index 11bd502f..5deb71fd 100644 --- a/ext/GmshExt.jl +++ b/ext/GmshExt.jl @@ -14,7 +14,32 @@ const GmshFile = FileMesh{Nothing, GmshMesh} # TODO flesh this out more # should we really do anything beyond linear elements? const gmsh_to_exo = Dict{String, String}( - "Triangle 3" => "Tri3" + "Quadrilateral 4" => "Quad4", + "Quadrilateral 9" => "Quad9", + "Triangle 3" => "Tri3", + "Triangle 6" => "Tri6", + "Tetrahedron 4" => "Tet4", + "Hexahedron 8" => "Hex8", +) +const gmsh_local_sides = Dict{String, Vector{Vector{Int}}}( + # 2D: edges of Tri3 / Quad4 -- matches FiniteElementContainers._create_local_edges + "Quadrilateral 4" => [[1, 2], [2, 3], [3, 4], [4, 1]], + "Quadrilateral 9" => [[1, 2, 5], [2, 3, 6], [3, 4, 7], [4, 1, 8]], + "Triangle 3" => [[1, 2], [2, 3], [3, 1]], + "Triangle 6" => [[1, 2, 4], [2, 3, 5], [3, 1, 6]], + # 3D: faces of Tet4 / Hex8, standard Exodus face/node ordering + "Hexahedron 8" => [[1, 2, 6, 5], [2, 3, 7, 6], [3, 4, 8, 7], + [1, 5, 8, 4], [1, 4, 3, 2], [5, 6, 7, 8]], + "Tetrahedron 4" => [[1, 2, 4], [2, 3, 4], [1, 4, 3], [1, 3, 2]], + +) + +# gmsh boundary element type names that can appear as the *boundary* +# entities of a given parent dimension. Used only to know how many nodes +# to expect / sanity check; the actual matching is purely by node-set. +const gmsh_boundary_types_by_dim = Dict{Int, Vector{String}}( + 1 => ["Line 2", "Line 3"], + 2 => ["Triangle 3", "Quadrilateral 4"], ) """ @@ -57,21 +82,26 @@ function FiniteElementContainers.element_blocks(mesh::GmshFile) for (block_id, name) in zip(block_ids, names) temp_el_types = String[] temp_conns = Matrix{Int}[] + temp_elem_tags = Vector{Int}[] for ent in gmsh.model.getEntitiesForPhysicalGroup(dim, block_id) ent_type, elem_tags, node_tags = gmsh.model.mesh.getElements(dim, ent) @assert length(ent_type) == 1 "Not supporting mixed element blocks right now" node_tags = convert(Vector{Int}, node_tags[1]) + elem_tags = convert(Vector{Int}, elem_tags[1]) el_type, _, _, num_nodes, _, _ = gmsh.model.mesh.getElementProperties(ent_type[1]) conn = reshape(node_tags, Int(num_nodes), length(node_tags) ÷ num_nodes) push!(temp_conns, conn) push!(temp_el_types, el_type) + push!(temp_elem_tags, elem_tags) end @assert length(temp_conns) == 1 && length(temp_el_types) == 1 conns[name] = temp_conns[1] el_types[name] = gmsh_to_exo[temp_el_types[1]] - # TODO fix this up - el_id_maps[name] = Vector{Int}(undef, 0) + # global element tags for this block, in the same column order as + # conns[name] -- this is what _setup_sideset uses via indexin to + # convert global sideset element tags into local block-column indices + el_id_maps[name] = temp_elem_tags[1] end block_ids = convert(Vector{Int}, block_ids) names_map = Dict(zip(block_ids, names)) @@ -138,39 +168,155 @@ function FiniteElementContainers.num_dimensions(::GmshFile) return Int64(gmsh.model.getDimension()) end +""" +$(TYPEDSIGNATURES) + +Builds a map from a canonical (sorted) set of boundary-entity nodes to the +parent element tag and the local side number of that element, for *all* +`dim`-dimensional elements in the mesh. This lets us match a lower +dimensional boundary element (from a sideset physical group) back to the +volume/area element it bounds, and the side number on that element, the +way Exodus side sets are defined. + +Works for both: +- `dim == 2`: boundary = edges of Tri3 / Quad4 +- `dim == 3`: boundary = faces of Tet4 / Hex8 + +The node ordering within a side only matters for orientation-sensitive uses +elsewhere (e.g. `_create_local_edges`); here we key on the *sorted* node +tuple, since a boundary entity and its parent's matching side are made up +of exactly the same set of nodes regardless of traversal order. + +Higher order or otherwise unsupported parent element types are skipped +with a warning rather than silently mismatched. +""" +function _build_boundary_to_parent_map(dim::Int) + side_map = Dict{Vector{Int}, Tuple{Int, Int}}() + elem_types, elem_tags_by_type, node_tags_by_type = gmsh.model.mesh.getElements(dim, -1) + + for (et, tags, nds) in zip(elem_types, elem_tags_by_type, node_tags_by_type) + el_type, _, _, num_nodes, _, _ = gmsh.model.mesh.getElementProperties(et) + local_sides = get(gmsh_local_sides, el_type, nothing) + if local_sides === nothing + @warn "Element type \"$el_type\" is not supported for sideset construction; boundary elements bordering this type will not be matched" + continue + end + + tags = convert(Vector{Int}, tags) + nds = convert(Vector{Int}, nds) + nn = Int(num_nodes) + + for (e, tag) in enumerate(tags) + conn = @views nds[(e - 1) * nn + 1:e * nn] + for (side_num, local_side) in enumerate(local_sides) + side_nodes = sort(conn[local_side]) + side_map[side_nodes] = (tag, side_num) + end + end + end + + return side_map +end + """ $(TYPEDEF) + +Constructs side sets (for e.g. Neumann BCs) from gmsh physical groups +of dimension `num_dimensions(mesh) - 1`. Each boundary element in the +physical group is matched against its parent volume/area element via +shared nodes, using `_build_boundary_to_parent_map`, to recover the +(element, side) pairs Exodus-style side sets require. + +Supports: +- 2D meshes (boundary = linear edges of Tri3 / Quad4) +- 3D meshes (boundary = linear faces of Tet4 / Hex8) + +Higher-order parent elements (e.g. Quad9) are not yet supported and will +produce a warning + skipped boundary elements rather than a hard failure. """ function FiniteElementContainers.sidesets(mesh::GmshFile) - @warn "Sidesets are currently not supported with Gmsh meshes" - dim = num_dimensions(mesh) - 1 + dim = num_dimensions(mesh) + bdim = dim - 1 + + if dim != 2 && dim != 3 + @warn "Gmsh sidesets are only supported for 2D (edge) and 3D (face) meshes. Returning empty sidesets." + return ( + Dict{String, Vector{Int}}(), Dict{Int, String}(), + Dict{String, Vector{Int}}(), Dict{String, Vector{Int}}(), + Dict{String, Matrix{Int}}() + ) + end + phys_groups = gmsh.model.getPhysicalGroups() - temp = filter(pg -> pg[1] == dim, phys_groups) + temp = filter(pg -> pg[1] == bdim, phys_groups) sset_ids = map(x -> x[2], temp) - names = map(x -> gmsh.model.getPhysicalName(dim, x), sset_ids) + names = map(x -> gmsh.model.getPhysicalName(bdim, x), sset_ids) - # pre-allocate stuff elems = Dict{String, Vector{Int}}() nodes = Dict{String, Vector{Int}}() sides = Dict{String, Vector{Int}}() side_nodes = Dict{String, Matrix{Int}}() - # for (name, sset_id) in zip(names, sset_ids) - # for ent in gmsh.model.getEntitiesForPhysicalGroup(dim, sset_id) - # ent_type, elem_tags, node_tags = gmsh.model.mesh.getElements(dim, ent) - # @assert length(ent_type) == 1 "Not supporting mixed element blocks right now" - # elem_tags = convert(Vector{Int}, elem_tags[1]) - # node_tags = convert(Vector{Int}, node_tags[1]) - # el_type, _, _, num_nodes, _, _ = gmsh.model.mesh.getElementProperties(ent_type[1]) - # @show elem_tags - # @show node_tags - # @show gmsh.model.mesh.getNodes(dim, sset_id) - # end - # end + side_map = _build_boundary_to_parent_map(dim) + + for (sset_id, name) in zip(sset_ids, names) + temp_elems = Int[] + temp_sides = Int[] + temp_side_nodes = Vector{Vector{Int}}(undef, 0) + temp_nodes = Int[] + + for ent in gmsh.model.getEntitiesForPhysicalGroup(bdim, sset_id) + ent_types, ent_elem_tags, ent_node_tags = gmsh.model.mesh.getElements(bdim, ent) + for (et, tags, nds) in zip(ent_types, ent_elem_tags, ent_node_tags) + bnd_type, _, _, num_nodes_elem, _, _ = gmsh.model.mesh.getElementProperties(et) + nn = Int(num_nodes_elem) + + tags = convert(Vector{Int}, tags) + nds = convert(Vector{Int}, nds) + + for e in 1:length(tags) + conn = nds[(e - 1) * nn + 1:e * nn] + key = sort(conn) + + if !haskey(side_map, key) + @warn "Could not find a parent element for boundary $(bnd_type) with nodes $conn in sideset \"$name\"; skipping this boundary element" + continue + end + + parent_tag, side_num = side_map[key] + push!(temp_elems, parent_tag) + push!(temp_sides, side_num) + push!(temp_side_nodes, conn) + append!(temp_nodes, conn) + end + end + end + + perm = sortperm(temp_elems) + elems[name] = temp_elems[perm] + sides[name] = temp_sides[perm] + + if length(temp_elems) == 0 + side_nodes[name] = Matrix{Int}(undef, 0, 0) + else + num_nodes_per_side = length(temp_side_nodes[1]) + # all boundary elements in a physical group should share the + # same face/edge node count; warn (don't crash) if not + if !all(x -> length(x) == num_nodes_per_side, temp_side_nodes) + @warn "Mixed boundary element node-counts found in sideset \"$name\"; side_nodes matrix may be malformed" + end + sn = reduce(hcat, temp_side_nodes) + # match the (1, N) flattened convention used in the Exodus ext + side_nodes[name] = reshape(sn[:, perm], 1, length(temp_nodes)) + end + + nodes[name] = sort(unique(temp_nodes)) + end + sset_ids = convert(Vector{Int}, sset_ids) - names = Dict(zip(sset_ids, names)) + names_map = Dict(zip(sset_ids, names)) - return elems, names, nodes, sides, side_nodes + return elems, names_map, nodes, sides, side_nodes end end # module diff --git a/src/DofManagers.jl b/src/DofManagers.jl index 7507f929..19463f47 100644 --- a/src/DofManagers.jl +++ b/src/DofManagers.jl @@ -1,5 +1,4 @@ -const CONSTRAINED_DOF = -1 - +const CONSTRAINED_DOF = -1 const DIRICHLET_DOF = -1 const PERIODIC_SIDE_B_DOF = -2 @@ -44,11 +43,11 @@ struct DofManager{ function DofManager{Condensed}(var::AbstractFunction) where Condensed dirichlet_dofs = zeros(Int, 0) - dof_to_unknown = 1:size(var.fspace.coords, 2) * length(names(var)) |> collect + dof_to_unknown = 1:num_entities(var.fspace) * length(names(var)) |> collect periodic_side_a_dofs = zeros(Int, 0) periodic_side_b_dofs = zeros(Int, 0) periodic_side_b_to_side_a_unknown = zeros(Int, 0) - unknown_dofs = 1:size(var.fspace.coords, 2) * length(names(var)) |> collect + unknown_dofs = 1:num_entities(var.fspace) * length(names(var)) |> collect return DofManager{Condensed}( dirichlet_dofs, dof_to_unknown, periodic_side_a_dofs, periodic_side_b_dofs, @@ -128,15 +127,15 @@ function Base.size(dof::DofManager, i::Int) end end -Base.show(io::IO, dof::DofManager) = -print(io, "DofManager\n", - " Number of entities = $(size(dof, 2))\n", - " Number of dofs per entity = $(size(dof, 1))\n", - " Number of total dofs = $(length(dof))\n", - " Number of dirichlet dofs = $(length(dof.dirichlet_dofs))\n", - " Number of periodic dofs = $(length(dof.periodic_side_b_dofs))\n", - " Number of unknown dofs = $(length(dof.unknown_dofs))\n", - " Storage type = $(typeof(dof.unknown_dofs))") +Base.show(io::IO, dof::DofManager; pad::String = "") = +print(io, "$(pad)DofManager\n", + "$(pad) Number of entities = $(size(dof, 2))\n", + "$(pad) Number of dofs per entity = $(size(dof, 1))\n", + "$(pad) Number of total dofs = $(length(dof))\n", + "$(pad) Number of dirichlet dofs = $(length(dof.dirichlet_dofs))\n", + "$(pad) Number of periodic dofs = $(length(dof.periodic_side_b_dofs))\n", + "$(pad) Number of unknown dofs = $(length(dof.unknown_dofs))\n", + "$(pad) Storage type = $(typeof(dof.unknown_dofs))") KA.get_backend(dof::DofManager) = KA.get_backend(dof.dirichlet_dofs) @@ -148,15 +147,25 @@ based on the variable ```dof``` was created with function create_field(dof::DofManager, float_type::Type{RT} = Float64) where RT <: Number backend = KA.get_backend(dof) data = fec_dense_array(backend, float_type, length(dof)) - coords = dof.var.fspace.coords - if isa(coords, H1Field) - field = H1Field{RT, typeof(data), num_fields(dof.var)}(data) - else - @assert false "Finish me" + var = dof.var + fspace = var.fspace + if _field_type(fspace) == H1Field + field = H1Field{RT, typeof(data), num_fields(var)}(data) + elseif _field_type(fspace) == HcurlField + field = HcurlField{RT, typeof(data), num_fields(var)}(data) + elseif _field_type(fspace) == HdivField + field = HdivField{RT, typeof(data), num_fields(var)}(data) + elseif _field_type(fspace) == L2Field + field = L2Field{RT, typeof(data), num_fields(var)}(undef, block_quadrature_sizes(fspace)) + fill!(field, zero(RT)) end return field end +function create_field(dof::Tuple) + return map(create_field, dof) +end + """ $(TYPEDSIGNATURES) Creates a vector of unknown dofs. @@ -174,6 +183,13 @@ function create_unknowns(dof::DofManager{Flag, IT, IDs, Var}, float_type = Float end end +function create_unknowns(dof::Tuple) + # data = BlockedArray{Float64}(undef, map(length, dof) |> collect) + data = BlockedArray{Float64}(undef, map(x -> length(x.unknown_dofs), dof) |> collect) + fill!(data, zero(eltype(data))) + return data +end + """ Map a single global DOF index to its reduced unknown index. @@ -186,7 +202,6 @@ by global DOF, so this is a pure array lookup with no allocation — safe to call inside GPU kernels. """ @inline function dof_to_unknown_index(dof::DofManager, n::Int) - # return dof_to_unknown_index(dof, n) dtu = dof.dof_to_unknown[n] if dtu == DIRICHLET_DOF return DIRICHLET_DOF @@ -401,9 +416,9 @@ function update_field_unknowns!( @assert KA.get_backend(Uu) == backend if _is_condensed(dof) - @assert length(Uu) == length(U) + @assert length(Uu) == length(U) "Unknown size $(length(Uu)), field size $(size(U))" else - @assert length(Uu) == length(dof.unknown_dofs) + @assert length(Uu) == length(dof.unknown_dofs) "Unknown size $(length(Uu)), field size $(size(U))" end _update_field_unknowns!(KA.get_backend(U), U, dof, Uu) diff --git a/src/Enzyme.jl b/src/Enzyme.jl index c8fb19f3..48bfd7ca 100644 --- a/src/Enzyme.jl +++ b/src/Enzyme.jl @@ -78,7 +78,7 @@ function _assemble_block_enzyme_safe!( for e in axes(state_old, 3) conn = connectivity(ref_fe, conns, e, coffset) - x_el, u_el, u_el_old = element_level_fields(ref_fe, conn, X, U, U_old) + x_el, u_el, u_el_old = element_level_fields(ref_fe, conn, e, X, U, U_old) props_el = _element_level_properties(props, e) val_el = _element_scratch(return_type, ref_fe, U) @@ -124,7 +124,7 @@ KA.@kernel function _assemble_block_enzyme_safe_kernel!( } E = KA.@index(Global) conn = connectivity(ref_fe, conns, E, coffset) - x_el, u_el, u_el_old = element_level_fields(ref_fe, conn, X, U, U_old) + x_el, u_el, u_el_old = element_level_fields(ref_fe, conn, E, X, U, U_old) props_el = _element_level_properties(props, E) val_el = _element_scratch(return_type, ref_fe, U) for q in 1:num_cell_quadrature_points(ref_fe) diff --git a/src/FiniteElementContainers.jl b/src/FiniteElementContainers.jl index da33f4f9..ac89e324 100644 --- a/src/FiniteElementContainers.jl +++ b/src/FiniteElementContainers.jl @@ -180,6 +180,7 @@ import AcceleratedKernels as AK import KernelAbstractions as KA using Adapt using Atomix +using BlockArrays using DocStringExtensions using Exodus using ForwardDiff diff --git a/src/FunctionSpaces.jl b/src/FunctionSpaces.jl index 4234b9ea..b7c11222 100644 --- a/src/FunctionSpaces.jl +++ b/src/FunctionSpaces.jl @@ -27,23 +27,6 @@ const _default_q = Dict{String, Int}( "TETRA4" => 2, "TETRA10" => 2 ) - -function _get_p_degree(el_type::String, p_degree::Union{Int, Nothing} = nothing) - if p_degree === nothing - return _default_p[el_type] - else - return p_degree - end -end - -function _get_q_degree(el_type::String, q_degree::Union{Int, Nothing} = nothing) - if q_degree === nothing - return _default_q[el_type] - else - return q_degree - end -end - const _el_name_to_juliac_safe_id = Dict{String, Int}( "HEX" => 1, "HEX8" => 1, @@ -59,22 +42,19 @@ 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 + ReferenceFE(Quad{Lagrange, 2}(), GaussLobattoLegendre(2, 2)), # QUAD9 for Lagrange + ReferenceFE(Tri{Lagrange, 1}(), GaussLobattoLegendre(2, 2)), # Tri3 for Lagrange + ReferenceFE(Tri{Lagrange, 2}(), GaussLobattoLegendre(2, 2)), # Tri3 for Lagrange + ReferenceFE(Tet{Lagrange, 1}(), GaussLobattoLegendre(2, 2)), # Tri3 for Lagrange + ReferenceFE(Tet{Lagrange, 2}(), GaussLobattoLegendre(2, 2)) # Tri3 for Lagrange +) const MAX_BLOCKS = 16 -function _setup_block_to_ref_fe_id(mesh::AbstractMesh, is_juliac_safe::Bool) - if is_juliac_safe - block_ids = Vector{Int}(undef, 0) - for block_name in mesh.element_block_names - el_type = mesh.element_types[block_name] - push!(block_ids, _el_name_to_juliac_safe_id[el_type]) - end - return block_ids - else - return 1:length(mesh.element_types) |> collect - end -end - function _setup_juliac_safe_block_to_ref_fe_id(mesh::AbstractMesh) names = mesh.element_block_names el_types = map(x -> _el_name_to_juliac_safe_id[mesh.element_types[x]], names) @@ -86,17 +66,13 @@ function _setup_block_to_ref_fe_id(mesh::AbstractMesh) return 1:length(mesh.element_types) |> collect end -# 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 - ReferenceFE(Quad{Lagrange, 2}(), GaussLobattoLegendre(2, 2)), # QUAD9 for Lagrange - ReferenceFE(Tri{Lagrange, 1}(), GaussLobattoLegendre(2, 2)), # Tri3 for Lagrange - ReferenceFE(Tri{Lagrange, 2}(), GaussLobattoLegendre(2, 2)), # Tri3 for Lagrange - ReferenceFE(Tet{Lagrange, 1}(), GaussLobattoLegendre(2, 2)), # Tri3 for Lagrange - ReferenceFE(Tet{Lagrange, 2}(), GaussLobattoLegendre(2, 2)) # Tri3 for Lagrange -) +function _setup_block_to_ref_fe_id(mesh::AbstractMesh, ::Val{is_juliac_safe}) where is_juliac_safe + if is_juliac_safe + return _setup_juliac_safe_block_to_ref_fe_id(mesh) + else + return _setup_block_to_ref_fe_id(mesh) + end +end """ default code path that sets up ref fes as a namedtuple @@ -125,21 +101,14 @@ function _setup_ref_fes( return ref_fes end -function _setup_quad_coords(mesh, X, conns, ref_fe) - NE = size(conns, 2) - NNPE = size(conns, 1) - NQ = num_quadrature_points(ref_fe) - ND = num_fields(mesh.nodal_coords) - coords_temp = zeros(ND, NQ, NE) - - for e in axes(X, 3) - X_el = SMatrix{ND, NNPE, Float64, ND * NNPE}(@views X[:, :, e]) - for q in 1:NQ - X_q = X_el * shape_function_value(ref_fe, q) - coords_temp[:, q, e] .= X_q - end +function _setup_ref_fes( + mesh, interp_type, p_degree, q_type, q_degree, ::Val{is_juliac_safe} +) where is_juliac_safe + if is_juliac_safe + return _juliac_safe_ref_fes + else + return _setup_ref_fes(mesh, interp_type, p_degree, q_type, q_degree) end - return coords_temp end """ @@ -157,38 +126,49 @@ $(TYPEDFIELDS) """ struct FunctionSpace{ IsJuliaCSafe, + FieldType, IT <: Integer, IV <: AbstractVector{IT}, + RT <: Number, + RV <: AbstractVector{RT}, + ND, BTRE, - Coords, + # Coords, RefFEs } <: AbstractFunctionSpace block_names::Vector{String} block_to_ref_fe_id::BTRE - coords::Coords + 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}} + # need to remove this mabye? node_id_map::IV ref_fes::RefFEs - function FunctionSpace{is_juliac_safe}( - block_names, block_to_ref_fe_id, coords, conns, elem_id_maps, node_id_map, ref_fes - ) where is_juliac_safe + function FunctionSpace{is_juliac_safe, FT}( + 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{ - is_juliac_safe, eltype(conns.data), typeof(conns.data), - typeof(block_to_ref_fe_id), typeof(coords), typeof(ref_fes) - }(block_names, block_to_ref_fe_id, coords, conns, elem_id_maps, node_id_map, ref_fes) + is_juliac_safe, FT, eltype(elem_conns.data), typeof(elem_conns.data), + eltype(coords), typeof(coords.data), size(coords, 1), + typeof(block_to_ref_fe_id), typeof(ref_fes) + }( + block_names, block_to_ref_fe_id, coords, elem_conns, elem_id_maps, + elem_to_facets, facet_orientation, node_id_map, ref_fes + ) end end function FunctionSpace( - mesh::AbstractMesh, field_type::Type{IT}, interp_type, + mesh::AbstractMesh, field_type::Type{FT}, interp_type, ::Type{QT} = GaussLobattoLegendre; is_juliac_safe::Bool = false, p_degree::Union{Int, Nothing} = nothing, q_degree::Union{Int, Nothing} = nothing, -) where {IT, QT <: ReferenceFiniteElements.AbstractQuadratureType} +) where {FT, QT <: ReferenceFiniteElements.AbstractQuadratureType} return FunctionSpace{is_juliac_safe}(mesh, field_type, interp_type, QT; p_degree = p_degree, q_degree = q_degree) end @@ -198,41 +178,38 @@ function FunctionSpace{is_juliac_safe}( p_degree::Union{Int, Nothing} = nothing, q_degree::Union{Int, Nothing} = nothing, ) where {is_juliac_safe, QT <: ReferenceFiniteElements.AbstractQuadratureType} - # TODO move to some common function so we can use it across - # all constructors - if p_degree !== nothing - # some error checking on p_degree input - if p_degree < 0 - @assert false "Bad polynomial degree $(p_degree)" - elseif p_degree == 0 - @assert false "TODO 0 order elements" - else - # TODO all of this needs TLC - ref_fes = _setup_ref_fes(mesh, interp_type, p_degree, q_type, q_degree) - coords, conns = create_higher_order_mesh(mesh, H1Field, interp_type, p_degree) - conns = Connectivity([val for val in values(conns)]) - end - else - if is_juliac_safe - # ref_fes = _setup_juliac_safe_ref_fes(interp_type, q_type) - ref_fes = _juliac_safe_ref_fes - else - ref_fes = _setup_ref_fes(mesh, interp_type, nothing, q_type, q_degree) - end - coords = mesh.nodal_coords - conns = Connectivity([val for val in values(mesh.element_conns)]) - end + ref_fes = _setup_ref_fes(mesh, interp_type, p_degree, q_type, q_degree, Val{is_juliac_safe}()) + coords = mesh.nodal_coords + conns = Connectivity([val for val in values(mesh.element_conns)]) elem_id_maps = [val for val in values(mesh.element_id_maps)] - if is_juliac_safe - block_to_ref_fe_id = _setup_juliac_safe_block_to_ref_fe_id(mesh) - else - # block_to_ref_fe_id = _setup_block_to_ref_fe_id(mesh, is_juliac_afe) - block_to_ref_fe_id = _setup_block_to_ref_fe_id(mesh) - end + block_to_ref_fe_id = _setup_block_to_ref_fe_id(mesh, Val{is_juliac_safe}()) + + return FunctionSpace{is_juliac_safe, H1Field}( + mesh.element_block_names, block_to_ref_fe_id, coords, + conns, elem_id_maps, nothing, nothing, mesh.node_id_map, ref_fes + ) +end + +function FunctionSpace{is_juliac_safe}( + mesh::AbstractMesh, ::Type{HdivField}, interp_type, + q_type::Type{QT} = GaussLobattoLegendre; + p_degree::Union{Int, Nothing} = nothing, + q_degree::Union{Int, Nothing} = nothing, +) where {is_juliac_safe, QT <: ReferenceFiniteElements.AbstractQuadratureType} + ref_fes = _setup_ref_fes(mesh, interp_type, p_degree, q_type, q_degree, Val{is_juliac_safe}()) + coords = mesh.nodal_coords + conns = Connectivity([val for val in values(mesh.element_conns)]) + elem_id_maps = [val for val in values(mesh.element_id_maps)] + block_to_ref_fe_id = _setup_block_to_ref_fe_id(mesh, Val{is_juliac_safe}()) + + # TODO only will work for likely linear elements right now + topology = UnstructuredTopology(mesh) + elem_to_facets = Connectivity([val for val in values(topology.elem_to_facets)]) + facet_orientation = Connectivity([val for val in values(topology.facet_orientation)]) - return FunctionSpace{is_juliac_safe}( + return FunctionSpace{is_juliac_safe, HdivField}( mesh.element_block_names, block_to_ref_fe_id, coords, - conns, elem_id_maps, mesh.node_id_map, ref_fes + conns, elem_id_maps, elem_to_facets, facet_orientation, mesh.node_id_map, ref_fes ) end @@ -243,39 +220,40 @@ function FunctionSpace{is_juliac_safe}( p_degree = nothing, q_degree = nothing ) where {is_juliac_safe, QT <: ReferenceFiniteElements.AbstractQuadratureType} - if is_juliac_safe - ref_fes = _setup_juliac_safe_ref_fes(interp_type, q_type) - else - ref_fes = _setup_ref_fes(mesh, interp_type, p_degree, q_type, q_degree) + if interp_type == Lagrange && p_degree !== nothing && q_degree === nothing + q_degree = p_degree + 1 end - + ref_fes = _setup_ref_fes(mesh, interp_type, p_degree, q_type, q_degree, Val{is_juliac_safe}()) + conns = Connectivity([val for val in values(mesh.element_conns)]) + # coords = L2Field(map(x -> mesh.nodal_coords[:, x], [values(mesh.element_conns)...])) + coords = mesh.nodal_coords + # new_conns = Array{Int, 2}[] + # offset = 1 + # for name in keys(mesh.element_conns) + # conn = mesh.element_conns[name] + # push!(new_conns, reshape(offset:offset + length(conn) - 1, size(conn)...)) + # offset += size(conn, 1) * size(conn, 2) + # end + # conns = Connectivity(new_conns) conns = Connectivity([val for val in values(mesh.element_conns)]) - coords = L2Field(map(x -> mesh.nodal_coords[:, x], [values(mesh.element_conns)...])) - - new_conns = Array{Int, 2}[] - offset = 1 - for name in keys(mesh.element_conns) - conn = mesh.element_conns[name] - push!(new_conns, reshape(offset:offset + length(conn) - 1, size(conn)...)) - offset += size(conn, 1) * size(conn, 2) - end - conns = Connectivity(new_conns) elem_id_maps = [val for val in values(mesh.element_id_maps)] - block_to_ref_fe_id = _setup_block_to_ref_fe_id(mesh, is_juliac_safe) + block_to_ref_fe_id = _setup_block_to_ref_fe_id(mesh) - return FunctionSpace{is_juliac_safe}( + return FunctionSpace{is_juliac_safe, L2Field}( mesh.element_block_names, block_to_ref_fe_id, coords, - conns, elem_id_maps, mesh.node_id_map, ref_fes + conns, elem_id_maps, nothing, nothing, mesh.node_id_map, ref_fes ) end function Adapt.adapt_structure(to, fspace::FunctionSpace) - return FunctionSpace{_is_juliac_safe(fspace)}( + return FunctionSpace{_is_juliac_safe(fspace), _field_type(fspace)}( fspace.block_names, fspace.block_to_ref_fe_id, adapt(to, fspace.coords), adapt(to, fspace.elem_conns), # fspace.elem_id_maps, map(x -> adapt(to, x), fspace.elem_id_maps), + adapt(to, fspace.elem_to_facets), + adapt(to, fspace.facet_orientation), adapt(to, fspace.node_id_map), map(x -> adapt(to, x), fspace.ref_fes) ) @@ -290,7 +268,11 @@ function Base.show(io::IO, fspace::FunctionSpace) end end -function _is_juliac_safe(::FunctionSpace{B, I, V, BTRE, C, R}) where {B, I, V, BTRE, C, R} +function _field_type(::FunctionSpace{B, FT, I, V, BTRE, C, R}) where {B, FT, I, V, BTRE, C, R} + return FT +end + +function _is_juliac_safe(::FunctionSpace{B, FT, I, V, BTRE, C, R}) where {B, FT, I, V, BTRE, C, R} return B end @@ -298,46 +280,38 @@ 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, I, V, BTRE, C, R}, block_id::Int) where {I, 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 @generated function block_reference_element( - fspace::FunctionSpace{true, IT, IV, BTRE, C, R}, + fspace::FunctionSpace{true, FT, IT, IV, BTRE, C, R}, block_id::Int -) where {IT, IV, BTRE, C, R} +) where {FT, IT, IV, BTRE, C, R} n_refs = length(BTRE.parameters) cases = map(1:n_refs) do j - quote - if block_id == $j - return fspace.ref_fes[$j] - end + quote + if block_id == $j + return fspace.ref_fes[$j] end + end end quote - id = fspace.block_to_ref_fe_id[block_id] # runtime value (OK) + id = fspace.block_to_ref_fe_id[block_id] # runtime value (OK) - if id == -1 - error("Inactive block ", block_id) - end + if id == -1 + error("Inactive block ", block_id) + end - $(cases...) # ONLY depends on compile-time j + $(cases...) # ONLY depends on compile-time j - error("Invalid ref_fe id ", id) + error("Invalid ref_fe id ", id) end end -# function block_reference_element( -# fspace::FunctionSpace{true, IT, IV, BTRE, C, R}, -# block_id::Int -# ) where {IT, IV, BTRE, C, R} -# index = fspace.block_to_ref_fe_id[block_id] -# return _block_reference_element(fspace, Val{index}()) -# end - # this does not work behind juliac function block_quadrature_size(fspace::FunctionSpace{false}, b::Int) ref_fe = block_reference_element(fspace, b) @@ -382,8 +356,23 @@ function num_elements(fspace::FunctionSpace, b::Int) return num_elements(fspace.elem_conns, b) end +function num_entities(fspace::FunctionSpace) + if _field_type(fspace) == H1Field + return size(fspace.coords, 2) + elseif _field_type(fspace) == L2Field + return mapreduce( + x -> block_quadrature_size(fspace, x)[1] * block_quadrature_size(fspace, x)[2], + prod, 1:num_blocks(fspace) + ) + end +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 + return num_entities_per_element(fspace.elem_conns, b) + end end # function num_q_points(fspace::FunctionSpace, b::Int) @@ -394,40 +383,3 @@ end function unsafe_connectivity(fspace::FunctionSpace, e::Int, b::Int) return unsafe_connectivity(fspace.elem_conns, e, b) end - -# # this is an H1 method, need to specialize -# function _create_linear_edges(fspace::FunctionSpace, ::Type{<:H1Field}, ::Type{<:Lagrange}) -# # need to assert this is already a linear function space -# for re in fspace.ref_fes -# if dimension(re) == 2 -# @assert polynomial_degree(re) == 1 -# else -# @assert false "Unsupported dimension = $(dimension(re))" -# end -# end - -# # maps canonical edges -# # to a vector of triplets where each element of that -# # vector corresponds to (el_id, local_edge_num, orientation) -# edge2elem = Dict{NTuple{2, Int}, Vector{NTuple{3, Int}}}() - -# for b in 1:num_blocks(fspace) -# conn = connectivity(fspace, b) -# el_ids = fspace.elem_id_maps[b] -# re = values(fspace.ref_fes)[b] -# for e in axes(conn, 2) -# local_edges = _create_local_edges(conn, re, e) -# el_id = el_ids[e] -# for (le_num, le) in enumerate(local_edges) -# ce = _canonical_edge(le...) -# orientation = le[1] < le[2] ? 1 : -1 -# push!( -# get!(edge2elem, ce, Vector{NTuple{3, Int}}()), -# (el_id, le_num, orientation) -# ) -# end -# end -# end -# return edge2elem -# end - diff --git a/src/assemblers/Assemblers.jl b/src/assemblers/Assemblers.jl index 92e0bbfc..4d6513a5 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,53 +135,40 @@ 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} +# 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 - u_el = @views SVector{NxNDof, eltype(U)}(U[:, conns[:, e]]) + u_el = @views SVector{NxNDof, eltype(U)}(U[:, conns]) return u_el end -@inline function _element_level_fields_flat(U::H1Field{T, D, NF}, ref_fe, conns) where {T, D, NF} +@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) - u_el_old = _element_level_fields_flat(U_old, ref_fe, conn) +@inline function element_level_fields(ref_fe, conn, e, X, U, U_old) + x_el = _element_level_fields_flat(X, ref_fe, conn, e) + u_el = _element_level_fields_flat(U, ref_fe, conn, e) + u_el_old = _element_level_fields_flat(U_old, ref_fe, conn, e) return x_el, u_el, u_el_old end -@inline function element_level_fields(ref_fe, conn, X, U, U_old, V) - x_el = _element_level_fields_flat(X, ref_fe, conn) - u_el = _element_level_fields_flat(U, ref_fe, conn) - u_el_old = _element_level_fields_flat(U_old, ref_fe, conn) - v_el = _element_level_fields_flat(V, ref_fe, conn) +@inline function element_level_fields(ref_fe, conn, e, X, U, U_old, V) + x_el = _element_level_fields_flat(X, ref_fe, conn, e) + u_el = _element_level_fields_flat(U, ref_fe, conn, e) + u_el_old = _element_level_fields_flat(U_old, ref_fe, conn, e) + v_el = _element_level_fields_flat(V, ref_fe, conn, e) return x_el, u_el, u_el_old, v_el end @@ -211,6 +196,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) """ @@ -417,7 +414,7 @@ function _assemble_block!( } fec_foraxes(state_old, 3) do e conn = connectivity(ref_fe, conns, e, coffset) - x_el, u_el, u_el_old = element_level_fields(ref_fe, conn, X, U, U_old) + x_el, u_el, u_el_old = element_level_fields(ref_fe, conn, e, X, U, U_old) props_el = _element_level_properties(props, e) val_el = _element_scratch(return_type, ref_fe, U) for q in 1:num_cell_quadrature_points(ref_fe) @@ -448,7 +445,7 @@ function _assemble_block!( } fec_foraxes(state_old, 3) do e conn = connectivity(ref_fe, conns, e, coffset) - x_el, u_el, u_el_old = element_level_fields(ref_fe, conn, X, U, U_old) + x_el, u_el, u_el_old = element_level_fields(ref_fe, conn, e, X, U, U_old) props_el = _element_level_properties(props, e) for q in 1:num_cell_quadrature_points(ref_fe) interps = _cell_interpolants(ref_fe, q) @@ -493,6 +490,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 00000000..c349fd67 --- /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/MatrixAction.jl b/src/assemblers/MatrixAction.jl index 0ac446a5..2e3172a1 100644 --- a/src/assemblers/MatrixAction.jl +++ b/src/assemblers/MatrixAction.jl @@ -62,7 +62,7 @@ function _assemble_block_matrix_free_action!( } fec_foraxes(state_old, 3) do e conn = connectivity(ref_fe, conns, e, coffset) - x_el, u_el, u_el_old, v_el = element_level_fields(ref_fe, conn, X, U, U_old, V) + x_el, u_el, u_el_old, v_el = element_level_fields(ref_fe, conn, e, X, U, U_old, V) props_el = _element_level_properties(props, e) Kv_el = _element_scratch(AssembledVector(), ref_fe, U) for q in 1:num_cell_quadrature_points(ref_fe) @@ -221,7 +221,7 @@ function _assemble_block_matrix_action!( } fec_foraxes(state_old, 3) do e conn = connectivity(ref_fe, conns, e, coffset) - x_el, u_el, u_el_old, v_el = element_level_fields(ref_fe, conn, X, U, U_old, V) + x_el, u_el, u_el_old, v_el = element_level_fields(ref_fe, conn, e, X, U, U_old, V) props_el = _element_level_properties(props, e) K_el = _element_scratch(AssembledMatrix(), ref_fe, U) for q in 1:num_cell_quadrature_points(ref_fe) @@ -254,7 +254,7 @@ function _assemble_block_matrix_action!( } fec_foraxes(state_old, 3) do e conn = connectivity(ref_fe, conns, e, coffset) - x_el, u_el, u_el_old, v_el = element_level_fields(ref_fe, conn, X, U, U_old, V) + x_el, u_el, u_el_old, v_el = element_level_fields(ref_fe, conn, e, X, U, U_old, V) props_el = _element_level_properties(props, e) for q in 1:num_cell_quadrature_points(ref_fe) interps = _cell_interpolants(ref_fe, q) diff --git a/src/assemblers/MatrixFreeAssembler.jl b/src/assemblers/MatrixFreeAssembler.jl index 1834601e..ca0e1ae0 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 74035b2d..71dbe2bd 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} @@ -259,7 +259,11 @@ function update_dofs!(assembler::AbstractAssembler, dirichlet_bcs::DirichletBCs, if assembler isa SparseMatrixAssembler && _is_matrix_free(assembler) # no-op: matrix pattern stays empty else - _update_dofs!(assembler.matrix_pattern, assembler.dof, ddofs, pdofs_side_b) + _update_dofs!( + assembler.matrix_pattern, + 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/assemblers/SparsityPatterns.jl b/src/assemblers/SparsityPatterns.jl index 565eecb9..25745c06 100644 --- a/src/assemblers/SparsityPatterns.jl +++ b/src/assemblers/SparsityPatterns.jl @@ -20,6 +20,26 @@ function _setup_block_sizes(dof::DofManager, ndims::Int) return block_start_indices, n_entries end +# below method only makes sense for sparse matrix patterns +function _setup_block_sizes(dof_1::DofManager, dof_2::DofManager) + ND1, ND2 = size(dof_1, 1), size(dof_2, 1) + fspace_1, fspace_2 = function_space(dof_1), function_space(dof_2) + @assert num_blocks(fspace_1) == num_blocks(fspace_2) "Got different number of blocks in the two function spaces." + n_blocks = num_blocks(fspace_1) + block_start_indices = Vector{Int64}(undef, n_blocks) + start_carry = 1 + for b in 1:n_blocks + NEPE1, NE1 = block_entity_size(fspace_1, b) + NEPE2, NE2 = block_entity_size(fspace_2, b) + @assert NE1 == NE2 "Got different numbers of elements in the two function spaces." + n_dofs_per_el = (ND1 * NEPE1) * (ND2 * NEPE2) + block_start_indices[b] = start_carry + start_carry = start_carry + n_dofs_per_el * NE1 + end + n_entries = start_carry - 1 + return block_start_indices, n_entries +end + # """ # $(TYPEDEF) # $(TYPEDFIELDS) @@ -116,6 +136,79 @@ function SparseMatrixPattern(dof::DofManager) return pattern end +function SparseMatrixPattern(dof_1::DofManager, dof_2::DofManager) + ND1, NN1 = size(dof_1) + ND2, NN2 = size(dof_2) + n_total_dofs_1 = NN1 * ND1 + n_total_dofs_2 = NN2 * ND2 + + fspace_1, fspace_2 = function_space(dof_1), function_space(dof_2) + n_blocks = num_blocks(fspace_1) + block_start_indices, n_entries = _setup_block_sizes(dof_1, dof_2) + + # setup pre-allocated arrays based on number of entries found above + Is = Vector{Int64}(undef, n_entries) + Js = Vector{Int64}(undef, n_entries) + unknown_dofs = Vector{Int64}(undef, n_entries) + + ids_1 = reshape(1:n_total_dofs_1, ND1, NN1) + ids_2 = reshape(1:n_total_dofs_2, ND2, NN2) + n = 1 + for b in 1:n_blocks + for e in 1:num_elements(fspace_1, b) + conn_1 = unsafe_connectivity(fspace_1, e, b) + conn_2 = unsafe_connectivity(fspace_2, e, b) + dof_conn_1 = @views reshape(ids_1[:, conn_1], ND1 * num_entities_per_element(fspace_1, b)) + dof_conn_2 = @views reshape(ids_2[:, conn_2], ND2 * num_entities_per_element(fspace_2, b)) + for i in axes(dof_conn_1, 1) + for j in axes(dof_conn_2, 1) + Is[n] = dof_conn_1[i] + Js[n] = dof_conn_2[j] + unknown_dofs[n] = n + n += 1 + end + end + end + end + + # create caches + # TODO not sure what to do for klasttouch and csrrowptr + # so making it the maximum size so we don't hit index out of bounds + # errors + # Is this right though? + n_total_dofs = max(n_total_dofs_1, n_total_dofs_2) + # hack above + klasttouch = zeros(Int64, n_total_dofs) + csrrowptr = zeros(Int64, n_total_dofs + 1) + csrcolval = zeros(Int64, length(Is)) + csrnzval = zeros(Float64, length(Is)) + + csccolptr = Vector{Int64}(undef, 0) + cscrowval = Vector{Int64}(undef, 0) + cscnzval_mass = Vector{Float64}(undef, 0) + cscnzval_stiffness = Vector{Float64}(undef, 0) + + # set permutation — pack (row, col) into one Int64 key so sortperm takes the + # integer radix-sort path instead of an O(N log N) tuple comparison sort. + # DOF indices are << 2^32, so the packed key preserves (i, j) lexicographic + # order and (sortperm being stable) yields an identical permutation. + ks = map((i, j) -> (Int64(i) << 32) | Int64(j), Is, Js) + permutation = sortperm(ks) + + pattern = SparseMatrixPattern( + Is, Js, + unknown_dofs, + block_start_indices, [n_entries], + # cache arrays + klasttouch, csrrowptr, csrcolval, csrnzval, + # additional cache arrays + csccolptr, cscrowval, cscnzval_mass, cscnzval_stiffness, + permutation + ) + + return pattern +end + function Adapt.adapt_structure(to, asm::SparseMatrixPattern) return SparseMatrixPattern( adapt(to, asm.Is), @@ -157,22 +250,36 @@ num_entries(s::SparseMatrixPattern) = length(s.Is) # NOTE assumes that dirichlet_dofs comes in sorted and uniqued # in the pattern object, that means that things # like Is, Js, etc. need to be viewed into or sliced -function _update_dofs!(pattern::SparseMatrixPattern, dof, dirichlet_dofs, periodic_side_b_dofs) - n_total_dofs = length(dof) - length(dirichlet_dofs) - length(periodic_side_b_dofs) - fspace = function_space(dof) +function _update_dofs!( + pattern::SparseMatrixPattern, + dof_1, dirichlet_dofs_1, periodic_side_b_dofs_1, + dof_2, dirichlet_dofs_2, periodic_side_b_dofs_2 +) + n_total_dofs_1 = length(dof_1) - length(dirichlet_dofs_1) - length(periodic_side_b_dofs_1) + n_total_dofs_2 = length(dof_2) - length(dirichlet_dofs_2) - length(periodic_side_b_dofs_2) + fspace_1, fspace_2 = function_space(dof_1), function_space(dof_2) + @assert num_blocks(fspace_1) == num_blocks(fspace_2) # figure out total number of entries - ND, NN = size(dof) - ids = reshape(1:length(dof), ND, NN) + ND1, NN1 = size(dof_1) + ND2, NN2 = size(dof_2) + + ids_1 = reshape(1:length(dof_1), ND1, NN1) + ids_2 = reshape(1:length(dof_2), ND2, NN2) + n_entries = 0 - for b in 1:num_blocks(fspace) - for e in 1:num_elements(fspace, b) - conns = unsafe_connectivity(fspace, e, b) - dof_conns = @views reshape(ids[:, conns], ND * num_entities_per_element(fspace, b)) - for i in axes(dof_conns, 1) - ri = dof_to_unknown_index(dof, dof_conns[i]) - for j in axes(dof_conns, 1) - rj = dof_to_unknown_index(dof, dof_conns[j]) + for b in 1:num_blocks(fspace_1) + # @assert num_elements(fspace_1) == num_elements(fspace_2) + @assert block_entity_size(fspace_1, b)[2] == block_entity_size(fspace_2, b)[2] + for e in 1:num_elements(fspace_1, b) + conns_1 = unsafe_connectivity(fspace_1, e, b) + conns_2 = unsafe_connectivity(fspace_2, e, b) + dof_conns_1 = @views reshape(ids_1[:, conns_1], ND1 * num_entities_per_element(fspace_1, b)) + dof_conns_2 = @views reshape(ids_2[:, conns_2], ND2 * num_entities_per_element(fspace_2, b)) + for i in axes(dof_conns_1, 1) + ri = dof_to_unknown_index(dof_1, dof_conns_1[i]) + for j in axes(dof_conns_2, 1) + rj = dof_to_unknown_index(dof_2, dof_conns_2[j]) if ri > 0 && rj > 0 n_entries += 1 end @@ -192,14 +299,16 @@ function _update_dofs!(pattern::SparseMatrixPattern, dof, dirichlet_dofs, period dof_num = 1 n = 1 - for b in 1:num_blocks(fspace) - for e in 1:num_elements(fspace, b) - conns = unsafe_connectivity(fspace, e, b) - dof_conns = @views reshape(ids[:, conns], ND * num_entities_per_element(fspace, b)) - for i in axes(dof_conns, 1) - ri = dof_to_unknown_index(dof, dof_conns[i]) - for j in axes(dof_conns, 1) - rj = dof_to_unknown_index(dof, dof_conns[j]) + for b in 1:num_blocks(fspace_1) + for e in 1:num_elements(fspace_1, b) + conns_1 = unsafe_connectivity(fspace_1, e, b) + conns_2 = unsafe_connectivity(fspace_2, e, b) + dof_conns_1 = @views reshape(ids_1[:, conns_1], ND1 * num_entities_per_element(fspace_1, b)) + dof_conns_2 = @views reshape(ids_2[:, conns_2], ND2 * num_entities_per_element(fspace_2, b)) + for i in axes(dof_conns_1, 1) + ri = dof_to_unknown_index(dof_1, dof_conns_1[i]) + for j in axes(dof_conns_2, 1) + rj = dof_to_unknown_index(dof_2, dof_conns_2[j]) if ri > 0 && rj > 0 pattern.Is[n] = ri pattern.Js[n] = rj @@ -212,6 +321,8 @@ function _update_dofs!(pattern::SparseMatrixPattern, dof, dirichlet_dofs, period end end + # TODO check this sizing + n_total_dofs = max(n_total_dofs_1, n_total_dofs_2) resize!(pattern.klasttouch, n_total_dofs) resize!(pattern.csrrowptr, n_total_dofs + 1) # TODO Not sure about below 2 sizes diff --git a/src/bcs/PeriodicBCs.jl b/src/bcs/PeriodicBCs.jl index 798b1b6b..5bef8646 100644 --- a/src/bcs/PeriodicBCs.jl +++ b/src/bcs/PeriodicBCs.jl @@ -223,7 +223,7 @@ Base.length(bcs::PeriodicBCs) = length(bcs.bc_funcs) function Base.show(io::IO, bcs::PeriodicBCs) # for (n, (cache, func)) in enumerate(zip(bcs.bc_cache, bcs.bc_funcs)) show(io, "PeriodicBC:") - show(io, bcs.bc_cache) + show(io, bcs.bc_caches) # show(io, bcs.bc_lengths) # show(io, func) show(io, "\n") diff --git a/src/meshes/Exodus.jl b/src/meshes/Exodus.jl index ef328f35..4c10b746 100644 --- a/src/meshes/Exodus.jl +++ b/src/meshes/Exodus.jl @@ -27,12 +27,12 @@ end # minimum interface function element_blocks(mesh::FileMesh{<:ExodusDatabase, ExodusMesh}) - blocks = read_sets(mesh.mesh_obj, Block) + blocks = read_sets(mesh.mesh_obj, Exodus.Block) block_ids = convert(Vector{Int}, map(x -> x.id, blocks)) conns = map(x -> convert(Matrix{Int}, x.conn), blocks) # el_id_maps = element_block_id_map.((mesh,), block_ids) el_id_maps = map(x -> convert(Vector{Int}, Exodus.read_block_id_map(mesh.mesh_obj, x)), block_ids) - names = Exodus.read_names(mesh.mesh_obj, Block) + names = Exodus.read_names(mesh.mesh_obj, Exodus.Block) types = map(x -> x.elem_type, blocks) conns = Dict(zip(names, conns)) @@ -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/src/meshes/Meshes.jl b/src/meshes/Meshes.jl index d8d2601e..d1f92cfe 100644 --- a/src/meshes/Meshes.jl +++ b/src/meshes/Meshes.jl @@ -272,7 +272,7 @@ function write_to_file(mesh::AbstractMesh, file_name::String; force::Bool = fals # write block names # block_names = map(String, values(mesh.element_block_names)) block_names = mesh.element_block_names - write_names(exo, Block, block_names) + write_names(exo, Exodus.Block, block_names) # TODO write block id maps # for (n, block_name) in mesh.element_block_names @@ -513,4 +513,6 @@ end include("AMRMesh.jl") include("StructuredMesh.jl") -include("UnstructuredMesh.jl") \ No newline at end of file +include("UnstructuredMesh.jl") + +include("Toplogy.jl") \ No newline at end of file diff --git a/src/meshes/Toplogy.jl b/src/meshes/Toplogy.jl new file mode 100644 index 00000000..ce7a6b5c --- /dev/null +++ b/src/meshes/Toplogy.jl @@ -0,0 +1,123 @@ +function _canonical_facet(nodes::NTuple{N, I}) where {N, I <: Integer} + return sort(nodes) +end + +function _get_local_facets(el_type, conn, e) + # eltype = block_element_type(block) + + if el_type in ("HEX","HEX8") + n = conn[:, e] + return ( + (n[1], n[2], n[3], n[4]), + (n[5], n[6], n[7], n[8]), + (n[1], n[2], n[6], n[5]), + (n[2], n[3], n[7], n[6]), + (n[3], n[4], n[8], n[7]), + (n[4], n[1], n[5], n[8]), + ) + elseif el_type in ("QUAD", "QUAD4") + n1, n2, n3, n4 = conn[:, e] + return ( + (n1, n2), + (n2, n3), + (n3, n4), + (n4, n1), + ) + elseif el_type in ("TET","TETRA","TET4") + n1, n2, n3, n4 = conn[:, e] + return ( + (n1, n2, n3), + (n1, n2, n4), + (n2, n3, n4), + (n1, n3, n4), + ) + elseif el_type in ("TRI", "TRI3") + n1, n2, n3 = conn[:, e] + return ( + (n1, n2), + (n2, n3), + (n3, n1), + ) + else + error("unsupported 2D element") + end +end + +# minimum working version for likely 2d only +function _is_oriented_positive(nodes) + return nodes[1] == minimum(nodes) +end + +struct UnstructuredTopology{I <: Integer} + elem_to_facets::Dict{String, Matrix{I}} + facets::Vector{Vector{I}} + facet_map::Dict{Vector{I}, I} + # facet_orientation::Dict{Tuple{String, I, I}, Int} + facet_orientation::Dict{String, Matrix{I}} + facet_to_elem::Vector{Vector{I}} +end + +function UnstructuredTopology(mesh::UnstructuredMesh) + facets = Vector{Vector{Int}}() + facet_map = Dict{Vector{Int}, Int}() + facet_to_elem = Vector{Vector{Int}}() + elem_to_facets = Dict{String, Matrix{Int}}() + facet_orientation = Dict{Tuple{String, Int, Int}, Int}() + + adj = Dict{Vector{Int}, Vector{Tuple{String, Int, Int}}}() + + for (block, conn) in mesh.element_conns + ne = size(conn, 2) + el_type = mesh.element_types[block] + nfacets = length(_get_local_facets(el_type, conn, 1)) + elem_to_facets[block] = Matrix{Int}(undef, nfacets, ne) + + for e in 1:ne + local_facets = _get_local_facets(el_type, conn, e) + + for (lf, nodes_tuple) in enumerate(local_facets) + + key = _canonical_facet(nodes_tuple) |> collect + fid = get!(facet_map, key) do + push!(facets, collect(key)) + push!(facet_to_elem, Int[]) + length(facets) + end + + elem_to_facets[block][lf, e] = fid + + # orientation (simple but extensible) + facet_orientation[(block, e, lf)] = + _is_oriented_positive(nodes_tuple) ? 1 : -1 + + push!(get!(adj, collect(key), Tuple{String, Int, Int}[]), + (block, e, lf)) + end + end + end + + for (fkey, uses) in adj + fid = facet_map[fkey] + facet_to_elem[fid] = [u[2] for u in uses] + end + + # better pack orientations + facet_orientation_new = Dict{String, Matrix{Int}}() + for (block, conn) in mesh.element_conns + temp = Matrix{Int}(undef, size(conn)) + for e in axes(conn, 2) + for n in axes(conn, 1) + temp[n, e] = facet_orientation[(block, e, n)] + end + end + facet_orientation_new[block] = temp + end + + return UnstructuredTopology( + elem_to_facets, + facets, + facet_map, + facet_orientation_new, + facet_to_elem + ) +end \ No newline at end of file diff --git a/test/poisson/TestPoisson.jl b/test/poisson/TestPoisson.jl index 6211b055..3b22f756 100644 --- a/test/poisson/TestPoisson.jl +++ b/test/poisson/TestPoisson.jl @@ -321,6 +321,65 @@ end end end +@testitem "Regression test - test_poisson_dirichlet_with_sidesets_gmsh_msh_tri3" setup=[PoissonRegressionHelper] tags=[:gmsh] begin + using Gmsh + msh_file_tri3 = dirname(Base.source_dir()) * "/gmsh/square_meshed_with_tris.msh" + output_file = "poisson_test_gmsh_sidesets_2.e" + mesh = UnstructuredMesh(msh_file_tri3) + V = FunctionSpace(mesh, H1Field, Lagrange) + physics = Poisson(f) + props = create_properties(physics) + u = ScalarFunction(V, "u") + for dev in backends + for use_condensed in use_condenseds + asm = SparseMatrixAssembler( + u; + sparse_matrix_type = :csc, + use_condensed = use_condensed, + use_inplace_methods = false + ) + # setup and update bcs + dbcs = DirichletBC[ + DirichletBC("u", bc_func; sideset_name = "top"), + DirichletBC("u", bc_func; sideset_name = "left") + ] + + nbcs = NeumannBC[ + NeumannBC("u", bc_func_neumann, "bottom"), + NeumannBC("u", bc_func_neumann, "right") + ] + + # direct solver test + # setup the parameters + p = create_parameters( + mesh, asm, physics, props; + dirichlet_bcs=dbcs, + neumann_bcs=nbcs + ) + + if dev != cpu + p = p |> dev + asm = asm |> dev + end + + # setup solver and integrator + solver = nlsolver(lsolver(asm)) + integrator = QuasiStaticIntegrator(solver) + evolve!(integrator, p) + + if dev != cpu + p = p |> cpu + end + display(solver.timer) + + pp = PostProcessor(mesh, output_file, u) + write_times(pp, 1, 0.0) + write_field(pp, 1, ("u",), p.field) + close(pp) + end + end +end + @testitem "Regression test - test_poisson_neumann" setup=[PoissonRegressionHelper] begin output_file = "poisson_test_5.e" for dev in backends From d46731208572076705a048949cbd1de757f3828a Mon Sep 17 00:00:00 2001 From: "Craig M. Hamel" Date: Fri, 10 Jul 2026 17:57:14 -0400 Subject: [PATCH 2/2] stashing some stuff for gather operations. --- src/Fields.jl | 41 ++++++++++++ src/FiniteElementContainers.jl | 1 + src/Gather.jl | 115 +++++++++++++++++++++++++++++++++ 3 files changed, 157 insertions(+) create mode 100644 src/Gather.jl diff --git a/src/Fields.jl b/src/Fields.jl index 102ba16d..975ad0a3 100644 --- a/src/Fields.jl +++ b/src/Fields.jl @@ -226,6 +226,47 @@ function unsafe_connectivity(conn::Connectivity, e::Int, b::Int) return view(conn.data, start:finish) end +###################################################################################################### +# Connectivity - new +###################################################################################################### +""" +$(TYPEDEF) +""" +struct Connectivity_v2{ + T <: Integer, + D <: AbstractVector{T} +} + data::D + nblocks::T + nepes::D + nelems::D + offsets::D +end + +function Connectivity_v2(conn::Connectivity) + return Connectivity_v2(conn.data, conn.nblocks, conn.nepes, conn.nelems, conn.offsets) +end + +function Adapt.adapt_structure(to, conn::Connectivity_v2) + return Connectivity_v2( + adapt(to, conn.data), + conn.nblocks, + adapt(to, conn.nepes), + adapt(to, conn.nelems), + adapt(to, conn.offsets) + ) +end + +@inline function connectivity(conn::Connectivity_v2, n::Int, e::Int, b::Int) + # idx = conn.nepes[b] + idx = conn.offsets[b] + conn.nepes[b] * (e - 1) + n - 1 + return conn.data[idx] +end +# function connectivity(data, offset::Int, nepe::Int, n::Int, e::Int) +# idx = offset + nepe * (e - 1) + n - 1 +# return data[idx] +# end + ###################################################################################################### # H1Field ###################################################################################################### diff --git a/src/FiniteElementContainers.jl b/src/FiniteElementContainers.jl index ac89e324..ec1fe916 100644 --- a/src/FiniteElementContainers.jl +++ b/src/FiniteElementContainers.jl @@ -221,6 +221,7 @@ include("Constraints.jl") include("InitialConditions.jl") include("Formulations.jl") +include("Gather.jl") include("Physics.jl") include("Properties.jl") include("assemblers/Assemblers.jl") diff --git a/src/Gather.jl b/src/Gather.jl new file mode 100644 index 00000000..b0b7854c --- /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