Select Git revision
function.jl
Stephan Hilb authored
Previously I thought it'd be nice to accept images in their own coordinate system and convert on the fly. Sadly this leads to more confusion than it is worth. It is better to talk of arrays instead of images and not do any conversion. Please permute your inputs in advance if you care about the internal coordinate system used (e.g. for vtk output data).
function.jl 11.31 KiB
using Statistics: mean
using StaticArrays: SA, SArray, MVector, MMatrix, SUnitRange
using DataFrames, CSV
export FeSpace, Mapper, FeFunction, P1, DP0, DP1
export interpolate!, sample, bind!, evaluate, integrate, nabla, save_csv
# Finite Elements
struct P1 end
struct DP0 end
struct DP1 end
ndofs(::P1) = 3
ndofs(::DP0) = 1
ndofs(::DP1) = 3
# 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
# TODO: size should probably be a type argument to make the constructor type
# safe
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
F = elmap(mesh, cell)
vloc = SA[0. 1. 0.; 0. 0. 1.]
for eldof in 1:nvertices_cell(mesh)
xloc = vloc[:, eldof]
x = F(xloc)::SArray
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
F = elmap(mesh, cell)
lcentroid = SA[1/3, 1/3]
centroid = F(lcentroid)
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 integrate(mesh::Mesh, expr; params...)
opparams = NamedTuple(params)
qw, qx = quadrature()
# only for type inference
opvalues = map(f -> evaluate(f, SA[0., 0.]), opparams)
val = zero(expr(SA[0., 0.]; opvalues...))
for cell in cells(mesh)
foreach(f -> bind!(f, cell), opparams)
delmap = jacobian(elmap(mesh, cell), SA[0., 0.])
intel = abs(det(delmap))
# quadrature points
for k in axes(qx, 2)
xhat = qx[:, k]
x = elmap(mesh, cell)(xhat)
opvalues = map(f -> evaluate(f, xhat), opparams)
val += qw[k] * expr(x; opvalues...) * intel
end
end
return val
end
function bind!(f::FeFunction, cell)
gdofs_view = view(f.space.dofmap, :, :, cell)
gdofs = SArray{Tuple{prod(f.space.size) * ndofs(f.space.element)}}(gdofs_view)
f.ldata .= f.data[gdofs]
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, M}
f::F
delmapinv::M
end
Base.show(io::IO, ::MIME"text/plain", x::Derivative) =
print("$(nameof(typeof(x))) of $(typeof(x.f))")
# TODO: should be called "derivative"
function nabla(f)
d = ndims_domain(f.space.mesh)
Derivative(f, zero(MMatrix{d, d, Float64}))
end
function bind!(df::Derivative, cell)
bind!(df.f, cell)
delmap = jacobian(elmap(df.f.space.mesh, cell), SA[0., 0.])
df.delmapinv .= inv(delmap)
end
function evaluate(df::Derivative, x)
# size: (:, d)
jac = jacobian(x -> evaluate(df.f.space, df.f.ldata, x), x)
jac *= df.delmapinv
return SArray{Tuple{df.f.space.size..., length(x)}}(jac)
end
struct FacetDivergence{F}
f::F
cell::Base.RefValue{Int}
edgemap::Dict{NTuple{2, Int}, Vector{Int}}
end
# TODO: unfinished!
function FacetDivergence(f::FeFunction)
f.space.element == DP0() || throw(ArgumentError("unimplemented"))
mesh = f.space.mesh
edgedivergence = Dict{NTuple{2, Int}, Float64}()
edgemap = Dict{NTuple{2, Int}, Vector{Int}}()
for cell in cells(mesh)
vs = sort(SVector(vertices(mesh, cell)))
e1 = (vs[1], vs[2])
e2 = (vs[1], vs[3])
e3 = (vs[2], vs[3])
edgemap[e1] = push!(get!(edgemap, e1, []), cell)
edgemap[e2] = push!(get!(edgemap, e2, []), cell)
edgemap[e3] = push!(get!(edgemap, e3, []), cell)
bind!(f, cell)
p = evaluate(f, SA[0., 0.])
A = SArray{Tuple{ndims_space(mesh), nvertices_cell(mesh)}}(
view(mesh.vertices, :, view(mesh.cells, :, cell)))
v1 = (A[:,2] - A[:,1])
v2 = (A[:,3] - A[:,2])
v3 = (A[:,1] - A[:,3])
p1 = norm(v1) * (p[1] * v1[2] - p[2] * v1[1])
p2 = norm(v2) * (p[1] * v2[2] - p[2] * v2[1])
p3 = norm(v3) * (p[1] * v3[2] - p[2] * v3[1])
edgedivergence[e1] = get!(edgedivergence, e1, 0.) + p1
edgedivergence[e3] = get!(edgedivergence, e3, 0.) + p2
edgedivergence[e2] = get!(edgedivergence, e2, 0.) + p3
end
return FacetDivergence(f, Ref(1), edgemap)
end
bind!(f::FacetDivergence, cell) = f.cell[] = cell
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())
function save_csv(path, f::FeFunction)
f.space.element::P1
f.space.size == (1,) ||
throw(ArgumentError("non-scalar functions unsupported"))
mesh = f.space.mesh
df = DataFrame()
X = SA[0 1 0; 0 0 1] # corners in local coordinates
for cell in cells(mesh)
bind!(f, cell)
A = SArray{Tuple{ndims_space(mesh), nvertices_cell(mesh)}}(
view(mesh.vertices, :, view(mesh.cells, :, cell)))
for k in axes(A, 2)
push!(df, (x = A[1, k], y = A[2, k],
f_xy = evaluate(f, X[:, k])[1]))
end
end
CSV.write(path, df)
end