Merge pull request #101 from Cthonios/physics/block-by-block

Physics/block by block

Author: cmhamel

ext/FiniteElementContainersAdaptExt.jl

@@ -106,6 +106,11 @@ function Adapt.adapt_structure(to, field::L2ElementField{T, NF, V, S}) where {T,
   return L2ElementField{T, NF, typeof(vals), S}(vals)
 end
 
+function Adapt.adapt_structure(to, field::L2QuadratureField{T, NF, NQ, V, S}) where {T, NF, NQ, V, S}
+  vals = adapt(to, field.vals)
+  return L2QuadratureField{T, NF, NQ, typeof(vals), S}(vals)
+end
+
 function Adapt.adapt_structure(to, field::H1Field{T, NF, V, S}) where {T, NF, V, S}
   vals = adapt(to, field.vals)
   return H1Field{T, NF, typeof(vals), S}(vals)

src/Parameters.jl

@@ -1,5 +1,6 @@
 abstract type AbstractParameters end
 
+# TODO need to break up bcs to different field types
 struct Parameters{D, N, T, Phys, Props, S, V, H1} <: AbstractParameters
   # actual parameter fields
   # TODO add boundary condition stuff and time stepping stuff
@@ -12,7 +13,7 @@ struct Parameters{D, N, T, Phys, Props, S, V, H1} <: AbstractParameters
   state_old::S
   state_new::S
   # scratch fields
-  h1_dbcs::V
+  h1_dbcs::V # TODO remove this
   h1_field::H1
 end
 
@@ -34,8 +35,6 @@ function Parameters(
 
   # TODO
   properties = nothing
-  state_old = nothing
-  state_new = nothing
 
   # for mixed spaces we'll need to do this more carefully
   if isa(physics, AbstractPhysics)
@@ -47,6 +46,25 @@ function Parameters(
     # TODO re-arrange physics tuple to match fspaces when appropriate
   end
 
+  # state_old = Array{Float64, 3}[]
+  state_old = L2QuadratureField[]
+  for (key, val) in pairs(physics)
+    NS = num_states(val)
+    NQ = ReferenceFiniteElements.num_quadrature_points(
+      getfield(values(assembler.dof.H1_vars)[1].fspace.ref_fes, key)
+    )
+    NE = size(
+      getfield(values(assembler.dof.H1_vars)[1].fspace.elem_conns, key),
+      2
+    )
+    syms = tuple(map(x -> Symbol("state_variable_$x"), 1:NS)...)
+    state_old_temp = L2QuadratureField(zeros(NS, NQ, NE), syms)
+    push!(state_old, state_old_temp)
+  end
+  state_new = copy(state_old)
+  state_old = NamedTuple{keys(physics)}(tuple(state_old...))
+  state_new = NamedTuple{keys(physics)}(tuple(state_new...))
+
   if dirichlet_bcs !== nothing
     syms = map(x -> Symbol("dirichlet_bc_$x"), 1:length(dirichlet_bcs))
     # dbcs = NamedTuple{tuple(syms...)}(tuple(dbcs...))
@@ -82,6 +100,11 @@ function Parameters(
 
   update_dofs!(assembler, p)
 
+  # assemble the stiffness at least once for 
+  # making easier to use on GPU
+  assemble!(assembler, H1Field, p, :stiffness)
+  K = stiffness(assembler)
+
   return p
 end
 

src/assemblers/Assemblers.jl

@@ -38,39 +38,107 @@ function _assemble_element!(global_val::H1Field, local_val, conn, e, b)
   return nothing
 end
 
-_assemble_block_method_from_sym(::Val{:mass}) = _assemble_block_mass!
-_assemble_block_method_from_sym(::Val{:residual}) = _assemble_block_residual!
-_assemble_block_method_from_sym(::Val{:residual_and_stiffness}) = _assemble_block_residual_and_stiffness!
-_assemble_block_method_from_sym(::Val{:stiffness}) = _assemble_block_stiffness!
-
-function _check_backends(assembler, U, X, conns)
+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)
   return backend
 end
 
+function assemble!(assembler, ::Type{H1Field}, p, val_sym::Val{:mass})
+  _zero_storage(assembler, val_sym)
+  fspace = assembler.dof.H1_vars[1].fspace
+  for (b, (conns, block_physics, state_old, state_new)) in enumerate(zip(
+    values(fspace.elem_conns), 
+    values(p.physics),
+    values(p.state_old), values(p.state_new)
+  ))
+    ref_fe = values(fspace.ref_fes)[b]
+    backend = _check_backends(assembler, p.h1_field, fspace.coords, state_old, state_new, conns)
+    _assemble_block_mass!(
+      assembler, block_physics, ref_fe, 
+      p.h1_field, fspace.coords, state_old, state_new,
+      conns, b, 
+      backend
+    )
+  end
+end
+
+function assemble!(assembler, ::Type{H1Field}, p, val_sym::Val{:residual})
+  _zero_storage(assembler, val_sym)
+  fspace = assembler.dof.H1_vars[1].fspace
+  for (b, (conns, block_physics, state_old, state_new)) in enumerate(zip(
+    values(fspace.elem_conns), 
+    values(p.physics),
+    values(p.state_old), values(p.state_new)
+  ))
+    ref_fe = values(fspace.ref_fes)[b]
+    backend = _check_backends(assembler, p.h1_field, fspace.coords, state_old, state_new, conns)
+    _assemble_block_residual!(
+      assembler, block_physics, ref_fe, 
+      p.h1_field, fspace.coords, state_old, state_new,
+      conns, b, 
+      backend
+    )
+  end
+end
+
+function assemble!(assembler, ::Type{H1Field}, p, val_sym::Val{:residual_and_stiffness})
+  _zero_storage(assembler, val_sym)
+  fspace = assembler.dof.H1_vars[1].fspace
+  for (b, (conns, block_physics, state_old, state_new)) in enumerate(zip(
+    values(fspace.elem_conns), 
+    values(p.physics),
+    values(p.state_old), values(p.state_new)
+  ))
+    ref_fe = values(fspace.ref_fes)[b]
+    backend = _check_backends(assembler, p.h1_field, fspace.coords, state_old, state_new, conns)
+    _assemble_block_residual_and_stiffness!(
+      assembler, block_physics, ref_fe, 
+      p.h1_field, fspace.coords, state_old, state_new,
+      conns, b, 
+      backend
+    )
+  end
+end
+
+function assemble!(assembler, ::Type{H1Field}, p, val_sym::Val{:stiffness})
+  _zero_storage(assembler, val_sym)
+  fspace = assembler.dof.H1_vars[1].fspace
+  for (b, (conns, block_physics, state_old, state_new)) in enumerate(zip(
+    values(fspace.elem_conns), 
+    values(p.physics),
+    values(p.state_old), values(p.state_new)
+  ))
+    ref_fe = values(fspace.ref_fes)[b]
+    backend = _check_backends(assembler, p.h1_field, fspace.coords, state_old, state_new, conns)
+    _assemble_block_stiffness!(
+      assembler, block_physics, ref_fe, 
+      p.h1_field, fspace.coords, state_old, state_new,
+      conns, b, 
+      backend
+    )
+  end
+end
+
 """
 $(TYPEDSIGNATURES)
 Top level assembly method for ```H1Field``` that loops over blocks and dispatches
 to appropriate kernels based on sym.
 
 TODO need to make sure at setup time that physics and elem_conns have the same
 values order. Otherwise, shenanigans.
+
+TODO figure out how to do generated functions
+
+creates one type instability from the Val
 """
-function assemble!(assembler, physics, U::H1Field, sym)
-  val_sym = Val{sym}()
-  _assemble_block_method! = _assemble_block_method_from_sym(val_sym)
-  _zero_storage(assembler, val_sym)
-  fspace = assembler.dof.H1_vars[1].fspace
-  X = fspace.coords
-  for (b, (conns, block_physics)) in enumerate(zip(values(fspace.elem_conns), values(physics)))
-    ref_fe = values(fspace.ref_fes)[b]
-    backend = _check_backends(assembler, U, X, conns)
-    _assemble_block_method!(assembler, block_physics, ref_fe, U, X, conns, b, backend)
-  end
+function assemble!(assembler, type::Type{H1Field}, p, sym::Symbol)
+  assemble!(assembler, type, p, Val{sym}())
 end
 
 """

src/assemblers/CPUGeneral.jl

@@ -33,7 +33,11 @@ with no threading.
 TODO add state variables and physics properties
 TODO remove Float64 typing below for eventual unitful use
 """
-function _assemble_block_residual!(assembler, physics, ref_fe, U, X, conns, block_id, ::KA.CPU)
+function _assemble_block_residual!(
+  assembler, physics, ref_fe, 
+  U, X, state_old, state_new,
+  conns, block_id, ::KA.CPU
+)
   ND = size(U, 1)
   NNPE = ReferenceFiniteElements.num_vertices(ref_fe)
   NxNDof = NNPE * ND

src/assemblers/GPUGeneral.jl

@@ -4,7 +4,11 @@
 $(TYPEDSIGNATURES)
 Kernel for residual block assembly
 """
-KA.@kernel function _assemble_block_residual_kernel!(assembler, physics, ref_fe, U, X, conns, block_id)
+KA.@kernel function _assemble_block_residual_kernel!(
+  assembler, physics, ref_fe, 
+  U, X, state_old, state_new,
+  conns, block_id
+)
   E = KA.@index(Global)
 
   ND = size(U, 1)
@@ -39,13 +43,25 @@ using KernelAbstractions and Atomix for eliminating race conditions
 
 TODO add state variables and physics properties
 """
-function _assemble_block_residual!(assembler, physics, ref_fe, U, X, conns, block_id, backend::KA.Backend)
+function _assemble_block_residual!(
+  assembler, physics, ref_fe, 
+  U, X, state_old, state_new,
+  conns, block_id, backend::KA.Backend
+)
   kernel! = _assemble_block_residual_kernel!(backend)
-  kernel!(assembler, physics, ref_fe, U, X, conns, block_id, ndrange=size(conns, 2))
+  kernel!(
+    assembler, physics, ref_fe, 
+    U, X, state_old, state_new,
+    conns, block_id, ndrange=size(conns, 2)
+  )
   return nothing
 end
 
-KA.@kernel function _assemble_block_stiffness_kernel!(assembler, physics, ref_fe, U, X, conns, block_id)
+KA.@kernel function _assemble_block_stiffness_kernel!(
+  assembler, physics, ref_fe, 
+  U, X, state_old, state_new,
+  conns, block_id
+)
   E = KA.@index(Global)
 
   ND = size(U, 1)
@@ -80,9 +96,17 @@ using KernelAbstractions and Atomix for eliminating race conditions
 
 TODO add state variables and physics properties
 """
-function _assemble_block_stiffness!(assembler, physics, ref_fe, U, X, conns, block_id, backend::KA.Backend)
+function _assemble_block_stiffness!(
+  assembler, physics, ref_fe, 
+  U, X, state_old, state_new,
+  conns, block_id, backend::KA.Backend
+)
   kernel! = _assemble_block_stiffness_kernel!(backend)
-  kernel!(assembler, physics, ref_fe, U, X, conns, block_id, ndrange=size(conns, 2))
+  kernel!(
+    assembler, physics, ref_fe, 
+    U, X, state_old, state_new, 
+    conns, block_id, ndrange=size(conns, 2)
+  )
   return nothing
 end
 

src/assemblers/SparseMatrixAssembler.jl

@@ -98,7 +98,11 @@ with no threading.
 
 TODO add state variables and physics properties
 """
-function _assemble_block_mass!(assembler, physics, ref_fe, U, X, conns, block_id, ::KA.CPU)
+function _assemble_block_mass!(
+  assembler, physics, ref_fe, 
+  U, X, state_old, state_new,
+  conns, block_id, ::KA.CPU
+)
   ND = size(U, 1)
   NNPE = ReferenceFiniteElements.num_vertices(ref_fe)
   NxNDof = NNPE * ND
@@ -125,7 +129,11 @@ with no threading.
 
 TODO add state variables and physics properties
 """
-function _assemble_block_stiffness!(assembler, physics, ref_fe, U, X, conns, block_id, ::KA.CPU)
+function _assemble_block_stiffness!(
+  assembler, physics, ref_fe, 
+  U, X, state_old, state_new,
+  conns, block_id, ::KA.CPU
+)
   ND = size(U, 1)
   NNPE = ReferenceFiniteElements.num_vertices(ref_fe)
   NxNDof = NNPE * ND
@@ -153,7 +161,11 @@ with no threading.
 TODO add state variables and physics properties
 TODO remove Float64 typing below for eventual unitful use
 """
-function _assemble_block_residual_and_stiffness!(assembler, physics, ref_fe, U, X, conns, block_id, ::KA.CPU)
+function _assemble_block_residual_and_stiffness!(
+  assembler, physics, ref_fe, 
+  U, X, state_old, state_new,
+  conns, block_id, ::KA.CPU
+)
   ND = size(U, 1)
   NNPE = ReferenceFiniteElements.num_vertices(ref_fe)
   NxNDof = NNPE * ND

src/fields/L2QuadratureField.jl

@@ -7,6 +7,9 @@ struct L2QuadratureField{T, NF, NQ, Vals <: AbstractArray{T, 1}, SymIDMap} <: Ab
   vals::Vals
 end
 
+# currently has shenanigans when vals has zero fields
+# will create a field that is 0 x NQ x 0 when it should
+# be 0 x NQ x NF. Likely due to the vec(vals) call
 function L2QuadratureField(vals::A, syms) where A <: AbstractArray{<:Number, 3}
   NF, NQ = size(vals, 1), size(vals, 2)
   @assert length(syms) == NF
@@ -21,7 +24,11 @@ end
 Base.IndexStyle(::Type{<:L2QuadratureField}) = IndexLinear()
 
 function Base.axes(field::L2QuadratureField{T, NF, NQ, V, SymIDMap}) where {T, NF, NQ, V, SymIDMap}
-  NE = length(field.vals) ÷ NF ÷ NQ
+  if NF == 0
+    NE = length(field.vals) ÷ NQ
+  else
+    NE = length(field.vals) ÷ NF ÷ NQ
+  end
   return (Base.OneTo(NF), Base.OneTo(NQ), Base.OneTo(NE))
 end
 
@@ -50,7 +57,14 @@ Base.setindex!(field::L2QuadratureField, v, n::Int) = setindex!(field.vals, v, n
 function Base.setindex!(field::L2QuadratureField{T, NF, NQ, A, S}, v, n::Int, q::Int, e::Int) where {T, NF, NQ, A, S}
   setindex!(field.vals, v, (e - 1) * NQ + (q - 1) * NF + n)
 end
-Base.size(field::L2QuadratureField{T, NF, NQ, A, S}) where {T, NF, NQ, A, S} = (NF, NQ, length(field.vals) ÷ NF ÷ NQ)
+
+function Base.size(field::L2QuadratureField{T, NF, NQ, A, S}) where {T, NF, NQ, A, S} 
+  if NF == 0
+    (NF, NQ, length(field.vals) ÷ NQ)
+  else
+    (NF, NQ, length(field.vals) ÷ NF ÷ NQ)
+  end
+end
 
 function Base.view(field::L2QuadratureField, sym::Symbol)
   d = _sym_id_map(field, sym)