Skip to content
Snippets Groups Projects
Select Git revision
  • c2e24364b3c6b77044854bc79b45e6ce54daa728
  • master default protected
  • andreas/paper2
  • v1.0
  • v0.1
5 results

function.jl

Blame
  • Stephan Hilb's avatar
    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).
    c2e24364
    History
    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