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

function.jl

Blame
  • function.jl 4.06 KiB
    using Statistics: mean
    export FeSpace, Mapper, FeFunction, P1, DP0, DP1
    export interpolate!, sample, bind!, evaluate
    
    # Finite Elements
    struct P1 end
    struct DP0 end
    struct DP1 end
    
    # FIXME: should be static vectors
    # evaluate all (1-dim) local basis functions against x
    evaluate_basis(::P1, x) = [1 - x[1] - x[2], x[1], x[2]]
    evaluate_basis(::DP0, x) = [x]
    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)
        size::S
    end
    
    function FeSpace(mesh, el::P1, size_=(1,))
        # special case for P1: dofs correspond to vertices
        rdims = prod(size_)
        ncells = size(mesh.cells, 2)
        dofmap = Array{Int, 3}(undef, rdims, 3, ncells)
        for i in CartesianIndices((rdims, 3, ncells))
    	dofmap[i] = rdims * (mesh.cells[i[2], i[3]] - 1) + i[1]
        end
        return FeSpace(mesh, el, dofmap, rdims * size(mesh.vertices, 2), size_)
    end
    
    function FeSpace(mesh, el::DP0, size_=(1,))
        # special case for P1: dofs correspond to vertices
        rdims = prod(size_)
        ncells = size(mesh.cells, 2)
        dofmap = Array{Int, 3}(undef, rdims, 1, ncells)
        for i in CartesianIndices((rdims, 1, ncells))
    	dofmap[i] = rdims * (i[3] - 1) + i[1]
        end
        return FeSpace(mesh, el, dofmap, rdims * ncells, size_)
    end
    
    Base.show(io::IO, ::MIME"text/plain", x::FeSpace) =
        print("$(nameof(typeof(x))), $(nameof(typeof(x.element))) elements, size $(x.size), $(x.ndofs) dofs")
    
    # evaluate at local point
    function evaluate(space::FeSpace, ldofs, xloc)
        bv = evaluate_basis(space.element, xloc)
        v = reshape(ldofs, size(space.dofmap)[1:2]) * bv
        return reshape(v, space.size)
    end
    
    
    
    # dof ordering for vector valued functions:
    #   (ldof, fdims...)
    
    # Array-valued function
    struct FeFunction{Sp}
        space::Sp
        data::Vector{Float64} # gdof -> data
        name::String
        ldata::Vector{Float64} # ldof -> data
    end
    
    function FeFunction(space; name=string(gensym("f")))
        data = Vector{Float64}(undef, space.ndofs)
        ldata = Vector{Float64}(undef, prod(size(space.dofmap)[1:2]))
        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, src::Function) = interpolate!(dst, dst.space.element, src)
    
    function interpolate!(dst::FeFunction, ::P1, src::Function)
        space = dst.space
        mesh = space.mesh
        for cell in axes(mesh.cells, 2)
    	for eldof in axes(mesh.cells, 1)
    	    xid = mesh.cells[eldof, cell]
    	    x = mesh.vertices[:, xid]
    
    	    gdofs = space.dofmap[:, eldof, cell]
    	    dst.data[gdofs] .= vec(src(x))
    	end
        end
    end
    
    function interpolate!(dst::FeFunction, ::DP0, src::Function)
        space = dst.space
        mesh = space.mesh
        for cell in axes(mesh.cells, 2)
    	vertices = mesh.vertices[:, mesh.cells[:, cell]]
    	centroid = mean(vertices, dims = 2)
    
    	gdofs = space.dofmap[:, 1, cell]
    	dst.data[gdofs] .= vec(src(centroid))
        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)
    
    struct CellFunction
        f::FeFunction
        dofs::Vector{Float64}
    end
    
    function sample(f::FeFunction)
        mesh = f.mapper.mesh
        for cell in axes(mesh.cells, 2)
    	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 * [1 - x[1] - x[2], x[1], x[2]], [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
    		# eval point
    	    end
    	end
        end
    end
    
    
    append_data!(vtkfile, f::FeFunction) = append_data!(vtkfile, f, f.space.element)
    
    function append_data!(vtkfile, f::FeFunction, ::P1)
        vtk_point_data(vtkfile, f.data, f.name)
    end
    
    function append_data!(vtkfile, f::FeFunction, ::DP0)
        vtk_cell_data(vtkfile, f.data, f.name)
    end