Select Git revision
operator.jl
operator.jl 7.51 KiB
using SparseArrays: sparse
using LinearAlgebra: det, dot
using StaticArrays: SA, SArray, MArray
using ForwardDiff: jacobian
export Poisson, L2Projection, init_point!, assemble, project_img, project_img2
abstract type Operator end
Base.show(io::IO, ::MIME"text/plain", x::Operator) =
print("$(nameof(typeof(x)))")
struct L2Projection{Sp, F} <: Operator
space::Sp
f::F # target
end
L2Projection(f) = L2Projection(f.space, f)
params(op::L2Projection) = (f = op.f,)
a(op::L2Projection, xloc, u, du, v, dv; f) = dot(u, v)
l(op::L2Projection, xloc, v, dv; f) = dot(f, v)
struct Poisson{Sp} <: Operator
space::Sp
end
a(op::Poisson, xloc, u, du, v, dv) = dot(du, dv)
# degree 2 quadrature on triangle
quadrature() = SA[1/6, 1/6, 1/6], SA[1/6 4/6 1/6; 1/6 1/6 4/6]
assemble(op::Operator) = assemble(op.space,
(x...; y...) -> a(op, x...; y...),
(x...; y...) -> l(op, x...; y...);
params(op)...)
function assemble(space::FeSpace, a, l; params...)
mesh = space.mesh
opparams = NamedTuple(params)
# precompute basis at quadrature points
qw, qx = quadrature()
d = size(qx, 1) # domain dimension
nrdims = prod(space.size)
nldofs = ndofs(space.element) # number of element dofs (i.e. local dofs not counting range dimensions)
nqpts = length(qw) # number of quadrature points
qphi_ = zeros(nrdims, nrdims, nldofs, nqpts)
dqphi_ = zeros(nrdims, d, nrdims, nldofs, nqpts)
for r in 1:nrdims
for k in axes(qx, 2)
qphi_[r, r, :, k] .= evaluate_basis(space.element, qx[:, k])
dqphi_[r, :, r, :, k] .= transpose(jacobian(x -> evaluate_basis(space.element, x), qx[:, k]))
end
end
qphi = SArray{Tuple{nrdims, nrdims, nldofs, nqpts}}(qphi_)
dqphi = SArray{Tuple{nrdims, d, nrdims, nldofs, nqpts}}(dqphi_)
I = Float64[]
J = Float64[]
V = Float64[]
b = zeros(ndofs(space))
gdof = LinearIndices((nrdims, ndofs(space)))
# mesh cells
for cell in cells(mesh)
foreach(f -> bind!(f, cell), opparams)
# cell map is assumed to be constant per cell
delmap = jacobian(elmap(mesh, cell), SA[0., 0.])
delmapinv = inv(delmap)
intel = abs(det(delmap))
# quadrature points
for k in axes(qx, 2)
xhat = qx[:, k]::SArray
x = elmap(mesh, cell)(xhat)
opvalues = map(f -> evaluate(f, xhat), opparams)
# local test-function dofs
for jdim in 1:nrdims, ldofj in 1:nldofs
gdofj = space.dofmap[jdim, ldofj, cell]
phij = SArray{Tuple{space.size...}}(qphi[:, jdim, ldofj, k])
dphij = SArray{Tuple{space.size..., d}}(dqphi[:, :, jdim, ldofj, k] * delmapinv)
lv = qw[k] * l(x, phij, dphij; opvalues...) * intel
b[gdofj] += lv
# local trial-function dofs
for idim in 1:nrdims, ldofi in 1:nldofs
gdofi = space.dofmap[idim, ldofi, cell]
phii = SArray{Tuple{space.size...}}(qphi[:, idim, ldofi, k])
dphii = SArray{Tuple{space.size..., d}}(dqphi[:, :, idim, ldofi, k] * delmapinv)
av = qw[k] * a(x, phii, dphii, phij, dphij; opvalues...) * intel
push!(I, gdofi)
push!(J, gdofj)
push!(V, av)
end
end
end
end
ngdofs = ndofs(space)
A = sparse(I, J, V, ngdofs, ngdofs)
return A, b
end
project_img(space::FeSpace, img) =
(u = FeFunction(space); project_img!(u, img))
# composite midpoint quadrature on lagrange point lattice
function quadrature_composite_lagrange_midpoint(p)
d_ = 2
n = binomial(p + 2, 2)
weights = Vector{Float64}(undef, n)
points = Matrix{Float64}(undef, 2, n)
k = 0
for I in Iterators.product(ntuple(_ -> 0:p, d_)...)
sum(Tuple(I)) > p && continue
k += 1
weights[k] = 1 / n
points[1, k] = I[1] / p
points[2, k] = I[2] / p
end
return weights, points
end
function project_img!(u::FeFunction, img)
d = 2 # domain dimension
space = u.space
mesh = space.mesh
f = ImageFunction(mesh, img)
opparams = (; f)
nrdims = prod(space.size)
# number of element dofs (i.e. local dofs not counting range dimensions)
nldofs = ndofs(space.element)
a(xloc, u, du, v, dv; f) = dot(u, v)
l(xloc, v, dv; f) = dot(f, v)
I = Float64[]
J = Float64[]
V = Float64[]
b = zeros(ndofs(space))
# mesh cells
for cell in cells(mesh)
foreach(f -> bind!(f, cell), opparams)
# cell map is assumed to be constant per cell
delmap = jacobian(elmap(mesh, cell), SA[0., 0.])
delmapinv = inv(delmap)
intel = abs(det(delmap))
p = ceil(Int, diam(mesh, cell))
qw, qx = quadrature_composite_lagrange_midpoint(p)
nqpts = length(qw) # number of quadrature points
qphi = zeros(nrdims, nrdims, nldofs, nqpts)
dqphi = zeros(nrdims, d, nrdims, nldofs, nqpts)
for k in axes(qx, 2)
for r in 1:nrdims
qphi[r, r, :, k] .= evaluate_basis(space.element, SVector{d}(view(qx, :, k)))
dqphi[r, :, r, :, k] .= transpose(jacobian(x -> evaluate_basis(space.element, x), SVector{d}(view(qx, :, k))))
end
end
# quadrature points
for k in axes(qx, 2)
xhat = SVector{d}(view(qx, :, k))
x = elmap(mesh, cell)(xhat)
opvalues = map(f -> evaluate(f, xhat), opparams)
# local test-function dofs
for jdim in 1:nrdims, ldofj in 1:nldofs
gdofj = space.dofmap[jdim, ldofj, cell]
phij = SArray{Tuple{space.size...}}(view(qphi, :, jdim, ldofj, k))
dphij = SArray{Tuple{space.size..., d}}(
SArray{Tuple{nrdims, d}}(view(dqphi, :, :, jdim, ldofj, k)) * delmapinv)
lv = qw[k] * l(x, phij, dphij; opvalues...) * intel
b[gdofj] += lv
# local trial-function dofs
for idim in 1:nrdims, ldofi in 1:nldofs
gdofi = space.dofmap[idim, ldofi, cell]
phii = SArray{Tuple{space.size...}}(view(qphi, :, idim, ldofi, k))
dphii = SArray{Tuple{space.size..., d}}(
SArray{Tuple{nrdims, d}}(view(dqphi, :, :, idim, ldofi, k)) * delmapinv)
av = qw[k] * a(x, phii, dphii, phij, dphij; opvalues...) * intel
push!(I, gdofi)
push!(J, gdofj)
push!(V, av)
end
end
end
end
ngdofs = ndofs(space)
A = sparse(I, J, V, ngdofs, ngdofs)
u.data .= A \ b
return u
end
project_img2(space::FeSpace, img) =
(u = FeFunction(space); project_img2!(u, img))
# L1-stable quasi-interpolation operator
# (https://arxiv.org/pdf/1505.06931.pdf)
#
# FIXME: currently only approximate quadrature by sampling on lagrange lattice
# based on element size. Exact evaluation is tricky to implement (lots of
# intersection handling).
function project_img2!(u::FeFunction, img)
u.space.element == P1() ||
throw(ArgumentError("element unsupported"))
d = 2 # domain dimension
space = u.space
mesh = space.mesh
f = ImageFunction(mesh, img)
nrdims = prod(space.size)
# number of element dofs (i.e. local dofs not counting range dimensions)
nldofs = ndofs(space.element)
u.data .= 0
# count contributions to respective dof for subsequent averaging
gdofcount = zeros(Int, size(u.data))
# mesh cells
for cell in cells(mesh)
bind!(f, cell)
p = ceil(Int, diam(mesh, cell))
# we actually only use the lagrange lattice points
qw, qx = quadrature_composite_lagrange_midpoint(p)
nqpts = length(qw) # number of quadrature points
qphi = zeros(nldofs, nqpts)
for k in axes(qx, 2)
xhat = SVector{d}(view(qx, :, k))
x = elmap(mesh, cell)(xhat)
# size: nrdims
value = evaluate(f, xhat)
# size: nldofs
phix = evaluate_basis(space.element, xhat)
# size: nrdims × nldofs
gdofs = u.space.dofmap[:, :, cell]
u.data[gdofs] .+= value * (12 .* phix' .- 3)
gdofcount[gdofs] .+= 1
end
end
u.data ./= gdofcount
return u
end