Select Git revision
function.jl
function.jl 8.92 KiB
using Statistics: mean
using StaticArrays: SA, SArray, MVector
export FeSpace, Mapper, FeFunction, P1, DP0, DP1
export interpolate!, sample, bind!, evaluate, nabla
# Finite Elements
struct P1 end
struct DP0 end
struct DP1 end
ndofs(::P1) = 3
ndofs(::DP0) = 1
ndofs(::DP1) = 3
# FIXME: should be static vectors
# evaluate all (1-dim) local basis functions against x
evaluate_basis(::P1, x) = SA[1 - x[1] - x[2], x[1], x[2]]
evaluate_basis(::DP0, x) = SA[1]
evaluate_basis(::DP1, x) = evaluate_basis(P1(), x)
# non-mixed
# scalar elements
struct FeSpace{M, Fe, S}
mesh::M
element::Fe
dofmap::Array{Int, 3} # (rdim, eldof, cell) -> gdof
ndofs::Int # = maximum(dofmap)
end
function FeSpace(mesh, el::P1, size_=(1,))
# special case for P1: dofs correspond to vertices
rdims = prod(size_)
dofmap = Array{Int, 3}(undef, rdims, 3, ncells(mesh))
for i in CartesianIndices((rdims, 3, ncells(mesh)))
dofmap[i] = rdims * (mesh.cells[i[2], i[3]] - 1) + i[1]
end
return FeSpace{typeof(mesh), typeof(el), size_}(
mesh, el, dofmap, rdims * nvertices(mesh))
end
function FeSpace(mesh, el::DP0, size_=(1,))
# special case for DP1: dofs correspond to cells
rdims = prod(size_)
dofmap = Array{Int, 3}(undef, rdims, 1, ncells(mesh))
for i in CartesianIndices((rdims, 1, ncells(mesh)))
dofmap[i] = rdims * (i[3] - 1) + i[1]
end
return FeSpace{typeof(mesh), typeof(el), size_}(
mesh, el, dofmap, rdims * ncells(mesh))
end
function FeSpace(mesh, el::DP1, size_=(1,))
# special case for DP1: dofs correspond to vertices
rdims = prod(size_)
dofmap = Array{Int, 3}(undef, rdims, 3, ncells(mesh))
for i in LinearIndices((rdims, 3, ncells(mesh)))
dofmap[i] = i
end
return FeSpace{typeof(mesh), typeof(el), size_}(
mesh, el, dofmap, rdims * 3 * ncells(mesh))
end
Base.show(io::IO, ::MIME"text/plain", x::FeSpace) =
print("$(nameof(typeof(x))), $(nameof(typeof(x.element))) elements, size $(x.size), $(ndofs(x)) dofs")
function Base.getproperty(obj::FeSpace{<:Any, <:Any, S}, sym::Symbol) where S
if sym === :size
return S
else
return getfield(obj, sym)
end
end
ndofs(space::FeSpace) = space.ndofs
# evaluate at local point
@inline function evaluate(space::FeSpace, ldofs, xloc)
bv = evaluate_basis(space.element, xloc)
ldofs_ = SArray{Tuple{prod(space.size), ndofs(space.element)}}(ldofs)
v = ldofs_ * bv
return SArray{Tuple{space.size...}}(v)
end
# dof ordering for vector valued functions:
# (rsize..., eldofs)
# Array-valued function
struct FeFunction{Sp,Ld}
space::Sp
data::Vector{Float64} # gdof -> data
name::String
ldata::Ld # ldof -> data
end
function FeFunction(space; name=string(gensym("f")))
data = Vector{Float64}(undef, ndofs(space))
ldata = zero(MVector{prod(space.size) * ndofs(space.element)})
return FeFunction(space, data, name, ldata)
end
Base.show(io::IO, ::MIME"text/plain", f::FeFunction) =
print("$(nameof(typeof(f))), size $(f.space.size) with $(length(f.data)) dofs")
interpolate!(dst::FeFunction, expr::Function; params...) =
interpolate!(dst, dst.space.element, expr; params...)
myvec(x) = vec(x)
myvec(x::Number) = x
function interpolate!(dst::FeFunction, ::P1, expr::Function; params...)
params = NamedTuple(params)
space = dst.space
mesh = space.mesh
for cell in cells(mesh)
for f in params
bind!(f, cell)
end
vertices = mesh.vertices[:, mesh.cells[:, cell]]
for eldof in 1:nvertices_cell(mesh)
x = SArray{Tuple{ndims_domain(mesh)}}(vertices[:, eldof])
xloc = SA[0. 1. 0.; 0. 0. 1.][:, eldof]
opvalues = map(f -> evaluate(f, xloc), params)
gdofs = space.dofmap[:, eldof, cell]
dst.data[gdofs] .= myvec(expr(x; opvalues...))
end
end
end
function interpolate!(dst::FeFunction, ::DP0, expr::Function; params...)
params = NamedTuple(params)
space = dst.space
mesh = space.mesh
for cell in cells(mesh)
for f in params
bind!(f, cell)
end
vertices = mesh.vertices[:, mesh.cells[:, cell]]
centroid = SArray{Tuple{ndims_domain(mesh)}}(mean(vertices, dims = 2))
lcentroid = SA[1/3, 1/3]
opvalues = map(f -> evaluate(f, lcentroid), params)
gdofs = space.dofmap[:, 1, cell]
dst.data[gdofs] .= myvec(expr(centroid; opvalues...))
end
end
interpolate!(dst::FeFunction, ::DP1, expr::Function; params...) =
interpolate!(dst, P1(), expr; params...)
project!(dst::FeFunction, expr::Function; params...) =
project!(dst, dst.space.element, expr; params...)
function project!(dst::FeFunction, ::DP0, expr; params...)
# same as interpolate!, but scaling by cell size
params = NamedTuple(params)
space = dst.space
mesh = space.mesh
for cell in cells(mesh)
delmap = jacobian(elmap(mesh, cell), SA[0., 0.])
intel = abs(det(delmap))
for f in params
bind!(f, cell)
end
vertices = mesh.vertices[:, mesh.cells[:, cell]]
centroid = SArray{Tuple{ndims_domain(mesh)}}(mean(vertices, dims = 2))
lcentroid = SA[1/3, 1/3]
opvalues = map(f -> evaluate(f, lcentroid), params)
gdofs = space.dofmap[:, 1, cell]
dst.data[gdofs] .= myvec(expr(centroid; opvalues...)) .* intel
end
end
function bind!(f::FeFunction, cell)
f.ldata .= vec(f.data[f.space.dofmap[:, :, cell]])
return f
end
# evaluate at local point (needs bind! call before)
evaluate(f::FeFunction, x) = evaluate(f.space, f.ldata, x)
# allow any non-function to act as a constant function
bind!(c, cell) = c
evaluate(c, xloc) = c
# TODO: inherit from some abstract function type
struct Derivative{F}
f::F
end
Base.show(io::IO, ::MIME"text/plain", x::Derivative) =
print("$(nameof(typeof(x))) of $(typeof(x.f))")
nabla(f) = Derivative(f)
bind!(df::Derivative, cell) = bind!(df.f, cell)
function evaluate(df::Derivative, x)
jac = jacobian(x -> evaluate(df.f.space, df.f.ldata, x), x)
return SArray{Tuple{df.f.space.size..., length(x)}}(jac)
end
function sample(f::FeFunction)
mesh = f.mapper.mesh
for cell in cells(mesh)
vertices = mesh.vertices[:, mesh.cells[:, cell]]
I0 = floor.(Int, vec(minimum(vertices, dims = 2)))
I1 = ceil.(Int, vec(maximum(vertices, dims = 2)))
J = jacobian(x -> vertices * SA[1 - x[1] - x[2], x[1], x[2]], SA[0., 0.])
xloc = J \ (v - vertices[:, 1])
for I in CartesianIndex(I0[1], I0[2]):CartesianIndex(I1[1], I1[2])
if all(xloc .>= 0) && sum(xloc) <= 1
# TODO: eval point
end
end
end
end
prolong!(new_f, old_cell, new_cells) =
prolong!(new_f, old_cell, new_cells, new_f.space.element)
prolong!(new_f, old_cell, new_cells, ::DP1) =
prolong!(new_f, old_cell, new_cells, P1())
function prolong!(new_f, old_cell, new_cells, ::P1)
old_f = new_f
old_cell_vs = collect(vertices(old_f.space.mesh, old_cell))
new_cell1_vs = collect(vertices(new_f.space.mesh, new_cells[1]))
new_cell2_vs = collect(vertices(new_f.space.mesh, new_cells[2]))
# copy over data for common vertices
common_vs = intersect(old_cell_vs, new_cell1_vs)
old_ldofs = indexin(common_vs, old_cell_vs)
old_gdofs = old_f.space.dofmap[:, old_ldofs, old_cell]
new_ldofs = indexin(common_vs, new_cell1_vs)
new_gdofs = new_f.space.dofmap[:, new_ldofs, new_cells[1]]
new_f.data[new_gdofs] .= old_f.data[old_gdofs]
common_vs = intersect(old_cell_vs, new_cell2_vs)
old_ldofs = indexin(common_vs, old_cell_vs)
old_gdofs = old_f.space.dofmap[:, old_ldofs, old_cell]
new_ldofs = indexin(common_vs, new_cell2_vs)
new_gdofs = new_f.space.dofmap[:, new_ldofs, new_cells[2]]
new_f.data[new_gdofs] .= old_f.data[old_gdofs]
# vertices of bisection edge
avg_vs = symdiff(new_cell1_vs, new_cell2_vs)
old_ldofs = indexin(avg_vs, old_cell_vs)
old_gdofs = old_f.space.dofmap[:, old_ldofs, old_cell]
avg_data = (old_f.data[old_gdofs[:, 1]] .+ old_f.data[old_gdofs[:, 2]]) ./ 2
_, newldof = findmax(new_cell1_vs)
new_gdofs = new_f.space.dofmap[:, newldof, new_cells[1]]
new_f.data[new_gdofs] .= avg_data
_, newldof = findmax(new_cell2_vs)
new_gdofs = new_f.space.dofmap[:, newldof, new_cells[2]]
new_f.data[new_gdofs] .= avg_data
end
function prolong!(new_f, old_cell, new_cells, ::DP0)
old_f = new_f
# simply copy over the data
old_gdofs = old_f.space.dofmap[:, 1, old_cell]
new_gdofs = new_f.space.dofmap[:, 1, new_cells[1]]
new_f.data[new_gdofs] .= old_f.data[old_gdofs]
new_gdofs = new_f.space.dofmap[:, 1, new_cells[2]]
new_f.data[new_gdofs] .= old_f.data[old_gdofs]
end
vtk_append!(vtkfile, f::FeFunction, fs::FeFunction...) =
(vtk_append!(vtkfile, f); vtk_append!(vtkfile, fs...))
vtk_append!(vtkfile, f::FeFunction) =
vtk_append!(vtkfile, f, f.space.element)
function vtk_append!(vtkfile, f::FeFunction, ::P1)
# separate out data per cell since different cells have distinct vertices
# in the output vtk
fd = vec(f.data[f.space.dofmap])
vtk_point_data(vtkfile, fd, f.name)
end
function vtk_append!(vtkfile, f::FeFunction, ::DP0)
vtk_cell_data(vtkfile, f.data, f.name)
end
vtk_append!(vtkfile, f::FeFunction, ::DP1) =
vtk_append!(vtkfile, f, P1())