diff --git a/Project.toml b/Project.toml index 6ea3089..40ff7d1 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "FiniteElementContainers" uuid = "d08262e4-672f-4e7f-a976-f2cea5767631" -version = "0.9.3" +version = "0.9.4" authors = ["Craig M. Hamel and contributors"] [deps] diff --git a/ext/FiniteElementContainersAMDGPUExt.jl b/ext/FiniteElementContainersAMDGPUExt.jl index ab71a56..17ca128 100644 --- a/ext/FiniteElementContainersAMDGPUExt.jl +++ b/ext/FiniteElementContainersAMDGPUExt.jl @@ -14,7 +14,7 @@ FiniteElementContainers.rocm(x) = Adapt.adapt_structure(ROCArray, x) function AMDGPU.rocSPARSE.ROCSparseMatrixCSC(asm::SparseMatrixAssembler) # Not sure if the line below is right on AMD @assert typeof(get_backend(asm)) <: ROCBackend "Assembler is not on a AMDGPU device" - @assert length(asm.pattern.cscnzval) > 0 "Need to assemble the assembler once with SparseArrays.sparse!(assembler)" + @assert length(asm.matrix_pattern.cscnzval) > 0 "Need to assemble the assembler once with SparseArrays.sparse!(assembler)" # below assertion isn't corect # @assert all(x -> x != zero(eltype(asm.pattern.cscnzval)), asm.pattern.cscnzval) "Need to assemble the assembler once with SparseArrays.sparse!(assembler)" @@ -28,14 +28,14 @@ function AMDGPU.rocSPARSE.ROCSparseMatrixCSC(asm::SparseMatrixAssembler) end return AMDGPU.rocSPARSE.ROCSparseMatrixCSC( - asm.pattern.csccolptr, - asm.pattern.cscrowval, - asm.pattern.cscnzval, + asm.matrix_pattern.csccolptr, + asm.matrix_pattern.cscrowval, + asm.matrix_pattern.cscnzval, (n_dofs, n_dofs) ) end -function FiniteElementContainers._stiffness(asm::SparseMatrixAssembler, backend::ROCBackend) +function FiniteElementContainers._stiffness(asm::SparseMatrixAssembler, ::ROCBackend) stiffness = AMDGPU.rocSPARSE.ROCSparseMatrixCSC(asm) return stiffness end diff --git a/ext/FiniteElementContainersCUDAExt.jl b/ext/FiniteElementContainersCUDAExt.jl index 761bbe5..d6fed26 100644 --- a/ext/FiniteElementContainersCUDAExt.jl +++ b/ext/FiniteElementContainersCUDAExt.jl @@ -23,9 +23,9 @@ function CUDA.CUSPARSE.CuSparseMatrixCSC(asm::SparseMatrixAssembler) end return CUDA.CUSPARSE.CuSparseMatrixCSC( - asm.pattern.csccolptr, - asm.pattern.cscrowval, - asm.pattern.cscnzval, + asm.matrix_pattern.csccolptr, + asm.matrix_pattern.cscrowval, + asm.matrix_pattern.cscnzval, (n_dofs, n_dofs) ) end diff --git a/src/assemblers/Assemblers.jl b/src/assemblers/Assemblers.jl index f39805c..96f982d 100644 --- a/src/assemblers/Assemblers.jl +++ b/src/assemblers/Assemblers.jl @@ -88,32 +88,32 @@ function _assemble_element!(global_val::H1Field, local_val, conn, e, b) return nothing end -function _check_backends(assembler, U, X, state_old, state_new, conns) - backend = KA.get_backend(assembler) - # TODO add get_backend method of ref_fe - @assert backend == KA.get_backend(U) - @assert backend == KA.get_backend(X) - @assert backend == KA.get_backend(conns) - @assert backend == KA.get_backend(state_old) - @assert backend == KA.get_backend(state_new) - # props will be complicated... - # TODO - return backend -end +# function _check_backends(assembler, U, X, state_old, state_new, conns) +# backend = KA.get_backend(assembler) +# # TODO add get_backend method of ref_fe +# @assert backend == KA.get_backend(U) +# @assert backend == KA.get_backend(X) +# @assert backend == KA.get_backend(conns) +# @assert backend == KA.get_backend(state_old) +# @assert backend == KA.get_backend(state_new) +# # props will be complicated... +# # TODO +# return backend +# end -function _check_backends(assembler, U, V, X, state_old, state_new, conns) - backend = KA.get_backend(assembler) - # TODO add get_backend method of ref_fe - @assert backend == KA.get_backend(U) - @assert backend == KA.get_backend(V) - @assert backend == KA.get_backend(X) - @assert backend == KA.get_backend(conns) - @assert backend == KA.get_backend(state_old) - @assert backend == KA.get_backend(state_new) - # props will be complicated... - # TODO - return backend -end +# function _check_backends(assembler, U, V, X, state_old, state_new, conns) +# backend = KA.get_backend(assembler) +# # TODO add get_backend method of ref_fe +# @assert backend == KA.get_backend(U) +# @assert backend == KA.get_backend(V) +# @assert backend == KA.get_backend(X) +# @assert backend == KA.get_backend(conns) +# @assert backend == KA.get_backend(state_old) +# @assert backend == KA.get_backend(state_new) +# # props will be complicated... +# # TODO +# return backend +# end """ $(TYPEDSIGNATURES) @@ -147,6 +147,13 @@ $(TYPEDSIGNATURES) return u_el end +""" +$(TYPEDSIGNATURES) +""" +@inline function _element_level_properties(props::AbstractArray, ::Int) + return props +end + """ $(TYPEDSIGNATURES) """ @@ -261,7 +268,7 @@ function stiffness(assembler::AbstractAssembler) end # some utilities -include("SparsityPattern.jl") +include("SparsityPatterns.jl") # types include("MatrixFreeAssembler.jl") diff --git a/src/assemblers/Matrix.jl b/src/assemblers/Matrix.jl index a724d80..0890f5e 100644 --- a/src/assemblers/Matrix.jl +++ b/src/assemblers/Matrix.jl @@ -2,7 +2,7 @@ function assemble_mass!( assembler, func::F, Uu, p ) where F <: Function assemble_matrix!( - assembler.mass_storage, assembler.pattern, assembler.dof, + assembler.mass_storage, assembler.matrix_pattern, assembler.dof, func, Uu, p ) end @@ -11,7 +11,7 @@ function assemble_stiffness!( assembler, func::F, Uu, p ) where F <: Function assemble_matrix!( - assembler.stiffness_storage, assembler.pattern, assembler.dof, + assembler.stiffness_storage, assembler.matrix_pattern, assembler.dof, func, Uu, p ) end @@ -71,8 +71,9 @@ function _assemble_block_matrix!( F2 <: AbstractField, F3 <: AbstractField, Func <: Function, - P <: Union{<:SVector, <:L2ElementField}, - Patt <: SparsityPattern, + # P <: Union{<:SVector, <:L2ElementField}, + P <: AbstractArray, + Patt <: SparseMatrixPattern, Phys <: AbstractPhysics, R <: ReferenceFE, S <: L2QuadratureField, @@ -113,8 +114,9 @@ KA.@kernel function _assemble_block_matrix_kernel!( F2 <: AbstractField, F3 <: AbstractField, Func <: Function, - P <: Union{<:SVector, <:L2ElementField}, - Patt <: SparsityPattern, + # P <: Union{<:SVector, <:L2ElementField}, + P <: AbstractArray, + Patt <: SparseMatrixPattern, Phys <: AbstractPhysics, R <: ReferenceFE, S <: L2QuadratureField, @@ -166,8 +168,9 @@ function _assemble_block_matrix!( F2 <: AbstractField, F3 <: AbstractField, Func <: Function, - P <: Union{<:SVector, <:L2ElementField}, - Patt <: SparsityPattern, + # P <: Union{<:SVector, <:L2ElementField}, + P <: AbstractArray, + Patt <: SparseMatrixPattern, Phys <: AbstractPhysics, R <: ReferenceFE, S <: L2QuadratureField, diff --git a/src/assemblers/MatrixAction.jl b/src/assemblers/MatrixAction.jl index 81bba70..86e07e1 100644 --- a/src/assemblers/MatrixAction.jl +++ b/src/assemblers/MatrixAction.jl @@ -1,12 +1,30 @@ +""" +$(TYPEDSIGNATURES) +""" +function assemble_matrix_action!( + assembler, func::F, Uu, Vu, p +) where F <: Function + assemble_matrix_action!( + assembler.stiffness_action_storage, + assembler.dof, + func, Uu, Vu, p + ) + return nothing +end + +""" +$(TYPEDSIGNATURES) +""" function assemble_matrix_action!( - assembler, func, Uu, Vu, p + # assembler, func, Uu, Vu, p + storage, dof, func, Uu, Vu, p ) - storage = assembler.stiffness_action_storage + # storage = assembler.stiffness_action_storage fill!(storage, zero(eltype(storage))) - fspace = function_space(assembler.dof) + fspace = function_space(dof) t = current_time(p.times) Δt = time_step(p.times) - _update_for_assembly!(p, assembler.dof, Uu, Vu) + _update_for_assembly!(p, dof, Uu, Vu) for (b, (conns, block_physics, state_old, state_new, props)) in enumerate(zip( values(fspace.elem_conns), values(p.physics), @@ -14,7 +32,8 @@ function assemble_matrix_action!( values(p.properties) )) ref_fe = values(fspace.ref_fes)[b] - backend = _check_backends(assembler, p.h1_field, p.h1_hvp_scratch_field, p.h1_coords, state_old, state_new, conns) + # backend = _check_backends(assembler, p.h1_field, p.h1_hvp_scratch_field, p.h1_coords, state_old, state_new, conns) + backend = KA.get_backend(p.h1_field) _assemble_block_matrix_action!( storage, block_physics, ref_fe, p.h1_field, p.h1_field_old, p.h1_hvp_scratch_field, p.h1_coords, state_old, state_new, props, t, Δt, @@ -44,7 +63,8 @@ function _assemble_block_matrix_action!( F2 <: AbstractField, F3 <: AbstractField, F4 <: AbstractField, - P <: Union{<:SVector, <:L2ElementField}, + # P <: Union{<:SVector, <:L2ElementField}, + P <: AbstractArray, Func <: Function, Phys <: AbstractPhysics, R <: ReferenceFE, @@ -93,7 +113,8 @@ KA.@kernel function _assemble_block_matrix_action_kernel!( F3 <: AbstractField, F4 <: AbstractField, Func <: Function, - P <: Union{<:SVector, <:L2ElementField}, + # P <: Union{<:SVector, <:L2ElementField}, + P <: AbstractArray, Phys <: AbstractPhysics, R <: ReferenceFE, S <: L2QuadratureField, @@ -149,7 +170,8 @@ function _assemble_block_matrix_action!( F3 <: AbstractField, F4 <: AbstractField, Func <: Function, - P <: Union{<:SVector, <:L2ElementField}, + # P <: Union{<:SVector, <:L2ElementField}, + P <: AbstractArray, Phys <: AbstractPhysics, R <: ReferenceFE, S <: L2QuadratureField, diff --git a/src/assemblers/NeumannBC.jl b/src/assemblers/NeumannBC.jl index 2f5b8c5..77511a0 100644 --- a/src/assemblers/NeumannBC.jl +++ b/src/assemblers/NeumannBC.jl @@ -1,16 +1,31 @@ -# # below method implicitly will not zero out arrays """ $(TYPEDSIGNATURES) """ function assemble_vector_neumann_bc!( assembler, Uu, p ) - storage = assembler.residual_storage + assemble_vector_neumann_bc!( + assembler.residual_storage, + assembler.dof, + Uu, p + ) + return nothing +end + +# # below method implicitly will not zero out arrays +""" +$(TYPEDSIGNATURES) +""" +function assemble_vector_neumann_bc!( + # assembler, Uu, p + storage, dof, Uu, p +) + # storage = assembler.residual_storage # do not zero! t = current_time(p.times) # TODO should below 2 methods calls be assumed to have # been conducted previously? - _update_for_assembly!(p, assembler.dof, Uu) + _update_for_assembly!(p, dof, Uu) for bc in values(p.neumann_bcs.bc_caches) backend = KA.get_backend(bc) _assemble_block_vector_neumann_bc!( diff --git a/src/assemblers/QuadratureQuantity.jl b/src/assemblers/QuadratureQuantity.jl index f419248..2e25138 100644 --- a/src/assemblers/QuadratureQuantity.jl +++ b/src/assemblers/QuadratureQuantity.jl @@ -60,7 +60,8 @@ function _assemble_block_quadrature_quantity!( F1 <: AbstractMatrix, F2 <: AbstractField, F3 <: AbstractField, - P <: Union{<:SVector, <:L2ElementField}, + # P <: Union{<:SVector, <:L2ElementField}, + P <: AbstractArray, Func <: Function, Phys <: AbstractPhysics, R <: ReferenceFE, @@ -104,7 +105,8 @@ KA.@kernel function _assemble_block_quadrature_quantity_kernel!( F2 <: AbstractField, F3 <: AbstractField, Func <: Function, - P <: Union{<:SVector, <:L2ElementField}, + # P <: Union{<:SVector, <:L2ElementField}, + P <: AbstractArray, Phys <: AbstractPhysics, R <: ReferenceFE, S <: L2QuadratureField, @@ -147,7 +149,8 @@ function _assemble_block_quadrature_quantity!( F2 <: AbstractField, F3 <: AbstractField, Func <: Function, - P <: Union{<:SVector, <:L2ElementField}, + # P <: Union{<:SVector, <:L2ElementField}, + P <: AbstractArray, Phys <: AbstractPhysics, R <: ReferenceFE, S <: L2QuadratureField, diff --git a/src/assemblers/SparseMatrixAssembler.jl b/src/assemblers/SparseMatrixAssembler.jl index 1fff8cf..5d2c63b 100644 --- a/src/assemblers/SparseMatrixAssembler.jl +++ b/src/assemblers/SparseMatrixAssembler.jl @@ -5,24 +5,29 @@ General sparse matrix assembler that can handle first or second order problems in time. """ struct SparseMatrixAssembler{ - Dof <: DofManager, - Pattern <: SparsityPattern, - Storage1 <: AbstractArray{<:Number, 1}, - Storage2 <: AbstractField, - Storage3 <: NamedTuple -} <: AbstractAssembler{Dof} - dof::Dof - pattern::Pattern - constraint_storage::Storage1 - damping_storage::Storage1 - hessian_storage::Storage1 - mass_storage::Storage1 - residual_storage::Storage2 - residual_unknowns::Storage1 - scalar_quadrature_storage::Storage3 # useful for energy like calculations - stiffness_storage::Storage1 - stiffness_action_storage::Storage2 - stiffness_action_unknowns::Storage1 + Condensed, + NumArrDims, + NumFields, + BlockPattern <: NamedTuple, + IV <: AbstractArray{Int, 1}, + RV <: AbstractArray{Float64, 1}, + Var <: AbstractFunction, + FieldStorage <: AbstractField{Float64, NumArrDims, RV, NumFields}, # should we make 2 not hardcoded? E.g. for + QuadratureStorage <: NamedTuple +} <: AbstractAssembler{DofManager{Condensed, Int, IV, Var}} + dof::DofManager{Condensed, Int, IV, Var} + matrix_pattern::SparseMatrixPattern{IV, BlockPattern, RV} + vector_pattern::SparseVectorPattern{IV} + constraint_storage::RV + damping_storage::RV + hessian_storage::RV + mass_storage::RV + residual_storage::FieldStorage + residual_unknowns::RV + scalar_quadrature_storage::QuadratureStorage # useful for energy like calculations + stiffness_storage::RV + stiffness_action_storage::FieldStorage + stiffness_action_unknowns::RV end # TODO this will not work for other than single H1 spaces @@ -33,22 +38,20 @@ e.g. ```H1Field```. Can be used to create block arrays for mixed FEM problems. """ function SparseMatrixAssembler(dof::DofManager) - pattern = SparsityPattern(dof) - # constraint_storage = zeros(length(dof)) - # ND, NN = num_dofs_per_node(dof), num_nodes(dof) + matrix_pattern = SparseMatrixPattern(dof) + vector_pattern = SparseVectorPattern(dof) + ND, NN = size(dof) n_total_dofs = ND * NN constraint_storage = zeros(n_total_dofs) - # constraint_storage = zeros(_dof_manager_vars(dof, type)) constraint_storage[dof.dirichlet_dofs] .= 1. - # fill!(constraint_storage, ) - # residual_storage = zeros(length(dof)) - damping_storage = zeros(num_entries(pattern)) - hessian_storage = zeros(num_entries(pattern)) - mass_storage = zeros(num_entries(pattern)) + + damping_storage = zeros(num_entries(matrix_pattern)) + hessian_storage = zeros(num_entries(matrix_pattern)) + mass_storage = zeros(num_entries(matrix_pattern)) residual_storage = create_field(dof) residual_unknowns = create_unknowns(dof) - stiffness_storage = zeros(num_entries(pattern)) + stiffness_storage = zeros(num_entries(matrix_pattern)) stiffness_action_storage = create_field(dof) stiffness_action_unknowns = create_unknowns(dof) @@ -59,15 +62,13 @@ function SparseMatrixAssembler(dof::DofManager) for (key, val) in pairs(fspace.ref_fes) NQ = ReferenceFiniteElements.num_quadrature_points(val) NE = size(getfield(fspace.elem_conns, key), 2) - syms = map(x -> Symbol("quadrature_field_$x"), 1:NQ) field = L2ElementField(zeros(Float64, NQ, NE)) push!(scalar_quadarature_storage, field) - # push!(scalar_quadarature_storage, zeros(Float64, NQ, NE)) end scalar_quadarature_storage = NamedTuple{keys(fspace.ref_fes)}(tuple(scalar_quadarature_storage...)) return SparseMatrixAssembler( - dof, pattern, + dof, matrix_pattern, vector_pattern, constraint_storage, damping_storage, hessian_storage, @@ -87,7 +88,8 @@ end function Adapt.adapt_structure(to, asm::SparseMatrixAssembler) return SparseMatrixAssembler( adapt(to, asm.dof), - adapt(to, asm.pattern), + adapt(to, asm.matrix_pattern), + adapt(to, asm.vector_pattern), adapt(to, asm.constraint_storage), adapt(to, asm.damping_storage), adapt(to, asm.hessian_storage), @@ -111,7 +113,7 @@ $(TYPEDSIGNATURES) Specialization of of ```_assemble_element!``` for ```SparseMatrixAssembler```. """ function _assemble_element!( - pattern::SparsityPattern, storage, K_el::SMatrix, el_id::Int, block_id::Int + pattern::SparseMatrixPattern, storage, K_el::SMatrix, el_id::Int, block_id::Int ) # figure out ids needed to update block_start_index = values(pattern.block_start_indices)[block_id] @@ -214,7 +216,7 @@ function _hessian(assembler::SparseMatrixAssembler, ::KA.CPU) end function _mass(assembler::SparseMatrixAssembler, ::KA.CPU) - M = SparseArrays.sparse!(assembler.pattern, assembler.mass_storage) + M = SparseArrays.sparse!(assembler.matrix_pattern, assembler.mass_storage) if _is_condensed(assembler.dof) _adjust_matrix_entries_for_constraints!(M, assembler.constraint_storage, KA.get_backend(assembler)) @@ -224,7 +226,7 @@ function _mass(assembler::SparseMatrixAssembler, ::KA.CPU) end function _stiffness(assembler::SparseMatrixAssembler, ::KA.CPU) - K = SparseArrays.sparse!(assembler.pattern, assembler.stiffness_storage) + K = SparseArrays.sparse!(assembler.matrix_pattern, assembler.stiffness_storage) if _is_condensed(assembler.dof) _adjust_matrix_entries_for_constraints!(K, assembler.constraint_storage, KA.get_backend(assembler)) @@ -282,7 +284,7 @@ function _update_dofs!(assembler::SparseMatrixAssembler, dirichlet_dofs::T) wher resize!(assembler.residual_unknowns, length(assembler.dof.unknown_dofs)) resize!(assembler.stiffness_action_unknowns, length(assembler.dof.unknown_dofs)) - _update_dofs!(assembler.pattern, assembler.dof, dirichlet_dofs) + _update_dofs!(assembler.matrix_pattern, assembler.dof, dirichlet_dofs) return nothing end diff --git a/src/assemblers/SparsityPattern.jl b/src/assemblers/SparsityPatterns.jl similarity index 76% rename from src/assemblers/SparsityPattern.jl rename to src/assemblers/SparsityPatterns.jl index c58173a..4d364a8 100644 --- a/src/assemblers/SparsityPattern.jl +++ b/src/assemblers/SparsityPatterns.jl @@ -5,7 +5,7 @@ Book-keeping struct for sparse matrices in FEM settings. This has all the information to construct a sparse matrix for either case where you want to eliminate fixed-dofs or not. """ -struct SparsityPattern{ +struct SparseMatrixPattern{ I <: AbstractArray{Int, 1}, B, R <: AbstractArray{Float64, 1} @@ -27,20 +27,15 @@ struct SparsityPattern{ end # TODO won't work for H(div) or H(curl) yet -function SparsityPattern(dof::DofManager) +function SparseMatrixPattern(dof::DofManager) # get number of dofs for creating cache arrays # TODO this line needs to be specialized for aribitrary fields # it's hardcoded for H1 firght now. - # ND, NN = num_dofs_per_node(dof), num_nodes(dof) ND, NN = size(dof) n_total_dofs = NN * ND - # vars = _dof_manager_vars(dof, type) - # n_blocks = length(vars[1].fspace.ref_fes) - # n_blocks = length(dof.H1_vars[1].fspace.ref_fes) - fspace = function_space(dof) n_blocks = length(fspace.ref_fes) @@ -51,18 +46,13 @@ function SparsityPattern(dof::DofManager) # for (n, conn) in enumerate(values(dof.H1_vars[1].fspace.elem_conns)) start_carry = 1 + ids = reshape(1:n_total_dofs, ND, NN) for (n, conn) in enumerate(values(fspace.elem_conns)) - ids = reshape(1:n_total_dofs, ND, NN) + # ids = reshape(1:n_total_dofs, ND, NN) # TODO do we need this operation? conn = reshape(ids[:, conn], ND * size(conn, 1), size(conn, 2)) n_entries += size(conn, 1)^2 * size(conn, 2) - # block_start_indices[n] = size(conn, 2) - # if n == 1 - # block_start_indices[n] = 1 - # else - # block_start_indices[n] = carry - # end block_start_indices[n] = start_carry start_carry = start_carry + size(conn, 1)^2 * size(conn, 2) @@ -82,7 +72,7 @@ function SparsityPattern(dof::DofManager) # now loop over function spaces and elements n = 1 for conn in values(fspace.elem_conns) - ids = reshape(1:n_total_dofs, ND, NN) + # ids = reshape(1:n_total_dofs, ND, NN) # TODO do we need this? block_conn = reshape(ids[:, conn], ND * size(conn, 1), size(conn, 2)) @@ -107,7 +97,7 @@ function SparsityPattern(dof::DofManager) cscrowval = Vector{Int64}(undef, 0) cscnzval = Vector{Float64}(undef, 0) - return SparsityPattern( + return SparseMatrixPattern( Is, Js, unknown_dofs, block_start_indices, block_el_level_sizes, @@ -118,7 +108,7 @@ function SparsityPattern(dof::DofManager) ) end -function Adapt.adapt_structure(to, asm::SparsityPattern) +function Adapt.adapt_structure(to, asm::SparseMatrixPattern) Is = adapt(to, asm.Is) Js = adapt(to, asm.Js) unknown_dofs = adapt(to, asm.unknown_dofs) @@ -133,7 +123,7 @@ function Adapt.adapt_structure(to, asm::SparsityPattern) csccolptr = adapt(to, asm.csccolptr) cscrowval = adapt(to, asm.cscrowval) cscnzval = adapt(to, asm.cscnzval) - return SparsityPattern( + return SparseMatrixPattern( Is, Js, unknown_dofs, block_start_indices, block_el_level_sizes, klasttouch, csrrowptr, csrcolval, csrnzval, @@ -141,7 +131,7 @@ function Adapt.adapt_structure(to, asm::SparsityPattern) ) end -function SparseArrays.sparse!(pattern::SparsityPattern, storage) +function SparseArrays.sparse!(pattern::SparseMatrixPattern, storage) return @views SparseArrays.sparse!( pattern.Is, pattern.Js, storage[pattern.unknown_dofs], length(pattern.klasttouch), length(pattern.klasttouch), +, pattern.klasttouch, @@ -150,13 +140,13 @@ function SparseArrays.sparse!(pattern::SparsityPattern, storage) ) end -num_entries(s::SparsityPattern) = length(s.Is) +num_entries(s::SparseMatrixPattern) = length(s.Is) # NOTE this methods assumes that dof is up to date # NOTE this method also only resizes unknown_dofs # in the pattern object, that means that things # like Is, Js, etc. need to be viewed into or sliced -function _update_dofs!(pattern::SparsityPattern, dof, dirichlet_dofs) +function _update_dofs!(pattern::SparseMatrixPattern, dof, dirichlet_dofs) n_total_dofs = length(dof) - length(dirichlet_dofs) # remove me @@ -199,3 +189,53 @@ function _update_dofs!(pattern::SparsityPattern, dof, dirichlet_dofs) return nothing end + +struct SparseVectorPattern{ + I <: AbstractArray{Int, 1} +} + Is::I + unknown_dofs::I +end + +function SparseVectorPattern(dof::DofManager) + ND, NN = size(dof) + n_total_dofs = NN * ND + ids = reshape(1:n_total_dofs, ND, NN) + + fspace = function_space(dof) + + n_entries = 0 + for (n, conn) in enumerate(values(fspace.elem_conns)) + # TODO do we need this operation? + conn = reshape(ids[:, conn], ND * size(conn, 1), size(conn, 2)) + n_entries += size(conn, 1)^2 * size(conn, 2) + end + + # setup pre-allocated arrays based on number of entries + Is = Vector{Int64}(undef, n_entries) + unknown_dofs = Vector{Int64}(undef, n_entries) + + n = 1 + for conn in values(fspace.elem_conns) + block_conn = reshape(ids[:, conn], ND * size(conn, 1), size(conn, 2)) + + for e in axes(block_conn, 2) + conn = @views block_conn[:, e] + for temp in conn + Is[n] = temp + unknown_dofs[n] = n + n += 1 + end + end + end + + return SparseVectorPattern(Is, unknown_dofs) +end + +function Adapt.adapt_structure(to, pattern::SparseVectorPattern) + return SparseVectorPattern( + adapt(to, pattern.Is), + adapt(to, pattern.unknown_dofs) + ) +end + diff --git a/src/assemblers/Vector.jl b/src/assemblers/Vector.jl index fbd172f..cacf54f 100644 --- a/src/assemblers/Vector.jl +++ b/src/assemblers/Vector.jl @@ -4,12 +4,24 @@ $(TYPEDSIGNATURES) function assemble_vector!( assembler, func::F, Uu, p ) where F <: Function - storage = assembler.residual_storage + assemble_vector!( + assembler.residual_storage, assembler.dof, + func, Uu, p + ) + return nothing +end + +""" +$(TYPEDSIGNATURES) +""" +function assemble_vector!( + storage, dof, func::F, Uu, p +) where F <: Function fill!(storage, zero(eltype(storage))) - fspace = function_space(assembler.dof) + fspace = function_space(dof) t = current_time(p.times) Δt = time_step(p.times) - _update_for_assembly!(p, assembler.dof, Uu) + _update_for_assembly!(p, dof, Uu) for (b, (conns, block_physics, state_old, state_new, props)) in enumerate(zip( values(fspace.elem_conns), values(p.physics), @@ -17,7 +29,7 @@ function assemble_vector!( values(p.properties) )) ref_fe = values(fspace.ref_fes)[b] - backend = _check_backends(assembler, p.h1_field, p.h1_coords, state_old, state_new, conns) + backend = KA.get_backend(p.h1_field) _assemble_block_vector!( storage, block_physics, ref_fe, p.h1_field, p.h1_field_old, p.h1_coords, state_old, state_new, props, t, Δt, @@ -50,7 +62,8 @@ function _assemble_block_vector!( F1 <: AbstractField, F2 <: AbstractField, F3 <: AbstractField, - P <: Union{<:SVector, <:L2ElementField}, + # P <: Union{<:SVector, <:L2ElementField}, + P <: AbstractArray, Func <: Function, Phys <: AbstractPhysics, R <: ReferenceFE, @@ -99,7 +112,8 @@ KA.@kernel function _assemble_block_vector_kernel!( F2 <: AbstractField, F3 <: AbstractField, Func <: Function, - P <: Union{<:SVector, <:L2ElementField}, + # P <: Union{<:SVector, <:L2ElementField}, + P <: AbstractArray, Phys <: AbstractPhysics, R <: ReferenceFE, S <: L2QuadratureField, @@ -154,7 +168,8 @@ function _assemble_block_vector!( F2 <: AbstractField, F3 <: AbstractField, Func <: Function, - P <: Union{<:SVector, <:L2ElementField}, + # P <: Union{<:SVector, <:L2ElementField}, + P <: AbstractArray, Phys <: AbstractPhysics, R <: ReferenceFE, S <: L2QuadratureField,