Select Git revision
operator.jl
operator.jl 5.70 KiB
using SparseArrays: sparse
using LinearAlgebra: det, dot
using StaticArrays: SA, SArray, MArray
using ForwardDiff: jacobian
export Poisson, L2Projection, init_point!, assemble, project_img
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
function project_img(space::FeSpace, img)
d = 2 # domain dimension
mesh = space.mesh
f = ImageFunction(mesh, img)
opparams = (; f)
nrdims = prod(space.size)
nldofs = ndofs(space.element) # number of element dofs (i.e. local dofs not counting range dimensions)
a(xloc, u, du, v, dv; f) = dot(u, v)
l(xloc, v, dv; f) = dot(f, v)
# composite midpoint quadrature on lagrange point lattice
function quadrature(p)
k = Iterators.filter(x -> sum(x) == p,
Iterators.product((0:p for _ in 1:d+1)...)) |> collect
weights = [1 / length(k) for _ in axes(k, 1)]
points = [x[i] / p for i in 1:2, x in k]
return weights::Vector{Float64}, points::Matrix{Float64}
end
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))
p = ceil(Int, diam(mesh, cell))
qw, qx = quadrature(p)
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), SVector{d}(qx[:, k])))
end
end
# quadrature points
for k in axes(qx, 2)
xhat = SVector{d}(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 = FeFunction(space)
u.data .= A \ b
return u
end