Select Git revision
function.jl
function.jl 4.06 KiB
using Statistics: mean
export FeSpace, Mapper, FeFunction, P1, DP0, DP1
export interpolate!, sample, bind!, evaluate
# Finite Elements
struct P1 end
struct DP0 end
struct DP1 end
# FIXME: should be static vectors
# evaluate all (1-dim) local basis functions against x
evaluate_basis(::P1, x) = [1 - x[1] - x[2], x[1], x[2]]
evaluate_basis(::DP0, x) = [x]
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)
size::S
end
function FeSpace(mesh, el::P1, size_=(1,))
# special case for P1: dofs correspond to vertices
rdims = prod(size_)
ncells = size(mesh.cells, 2)
dofmap = Array{Int, 3}(undef, rdims, 3, ncells)
for i in CartesianIndices((rdims, 3, ncells))
dofmap[i] = rdims * (mesh.cells[i[2], i[3]] - 1) + i[1]
end
return FeSpace(mesh, el, dofmap, rdims * size(mesh.vertices, 2), size_)
end
function FeSpace(mesh, el::DP0, size_=(1,))
# special case for P1: dofs correspond to vertices
rdims = prod(size_)
ncells = size(mesh.cells, 2)
dofmap = Array{Int, 3}(undef, rdims, 1, ncells)
for i in CartesianIndices((rdims, 1, ncells))
dofmap[i] = rdims * (i[3] - 1) + i[1]
end
return FeSpace(mesh, el, dofmap, rdims * ncells, size_)
end
Base.show(io::IO, ::MIME"text/plain", x::FeSpace) =
print("$(nameof(typeof(x))), $(nameof(typeof(x.element))) elements, size $(x.size), $(x.ndofs) dofs")
# evaluate at local point
function evaluate(space::FeSpace, ldofs, xloc)
bv = evaluate_basis(space.element, xloc)
v = reshape(ldofs, size(space.dofmap)[1:2]) * bv
return reshape(v, space.size)
end
# dof ordering for vector valued functions:
# (ldof, fdims...)
# Array-valued function
struct FeFunction{Sp}
space::Sp
data::Vector{Float64} # gdof -> data
name::String
ldata::Vector{Float64} # ldof -> data
end
function FeFunction(space; name=string(gensym("f")))
data = Vector{Float64}(undef, space.ndofs)
ldata = Vector{Float64}(undef, prod(size(space.dofmap)[1:2]))
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, src::Function) = interpolate!(dst, dst.space.element, src)
function interpolate!(dst::FeFunction, ::P1, src::Function)
space = dst.space
mesh = space.mesh
for cell in axes(mesh.cells, 2)
for eldof in axes(mesh.cells, 1)
xid = mesh.cells[eldof, cell]
x = mesh.vertices[:, xid]
gdofs = space.dofmap[:, eldof, cell]
dst.data[gdofs] .= vec(src(x))
end
end
end
function interpolate!(dst::FeFunction, ::DP0, src::Function)
space = dst.space
mesh = space.mesh
for cell in axes(mesh.cells, 2)
vertices = mesh.vertices[:, mesh.cells[:, cell]]
centroid = mean(vertices, dims = 2)
gdofs = space.dofmap[:, 1, cell]
dst.data[gdofs] .= vec(src(centroid))
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)
struct CellFunction
f::FeFunction
dofs::Vector{Float64}
end
function sample(f::FeFunction)
mesh = f.mapper.mesh
for cell in axes(mesh.cells, 2)
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 * [1 - x[1] - x[2], x[1], x[2]], [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
# eval point
end
end
end
end
append_data!(vtkfile, f::FeFunction) = append_data!(vtkfile, f, f.space.element)
function append_data!(vtkfile, f::FeFunction, ::P1)
vtk_point_data(vtkfile, f.data, f.name)
end
function append_data!(vtkfile, f::FeFunction, ::DP0)
vtk_cell_data(vtkfile, f.data, f.name)
end