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

function.jl

Blame
  • function.jl 8.92 KiB
    using Statistics: mean
    using StaticArrays: SA, SArray, MVector
    export FeSpace, Mapper, FeFunction, P1, DP0, DP1
    export interpolate!, sample, bind!, evaluate, nabla
    
    # Finite Elements
    struct P1 end
    struct DP0 end
    struct DP1 end
    
    ndofs(::P1) = 3
    ndofs(::DP0) = 1
    ndofs(::DP1) = 3
    
    # FIXME: should be static vectors
    # 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
    
    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
    
    	vertices = mesh.vertices[:, mesh.cells[:, cell]]
    	for eldof in 1:nvertices_cell(mesh)
    	    x = SArray{Tuple{ndims_domain(mesh)}}(vertices[:, eldof])
    	    xloc = SA[0. 1. 0.; 0. 0. 1.][:, eldof]
    
    	    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
    
    	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...))
        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 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)
    
    # 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}
        f::F
    end
    
    Base.show(io::IO, ::MIME"text/plain", x::Derivative) =
        print("$(nameof(typeof(x))) of $(typeof(x.f))")
    
    nabla(f) = Derivative(f)
    
    bind!(df::Derivative, cell) = bind!(df.f, cell)
    function evaluate(df::Derivative, x)
        jac = jacobian(x -> evaluate(df.f.space, df.f.ldata, x), x)
        return SArray{Tuple{df.f.space.size..., length(x)}}(jac)
    end
    
    
    function sample(f::FeFunction)
        mesh = f.mapper.mesh
        for cell in cells(mesh)
    	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 * SA[1 - x[1] - x[2], x[1], x[2]], SA[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
    		# TODO: eval point
    	    end
    	end
        end
    end
    
    
    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())