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

mesh.jl

Blame
  • mesh.jl 12.47 KiB
    export init_grid, init_hgrid, save, refine!, refine, cells, vertices,
        ndims_domain, ndims_space, nvertices, nvertices_cell, ncells
    
    using LinearAlgebra: norm
    
    using WriteVTK
    using StaticArrays: SVector
    
    struct HMesh
        vertices::Vector{SVector{2, Float64}}
        cells::Vector{NTuple{3, Int}}
        levels::Vector{Int}
    end
    
    Base.show(io::IO, ::MIME"text/plain", x::HMesh) =
        print("$(nameof(typeof(x))), $(ncells(x)) cells")
    
    ndims_domain(::HMesh) = 2
    ndims_space(::HMesh) = 2
    # TODO: move to simplex mesh interface
    nvertices_cell(::HMesh) = 3
    # TODO: do we even these two?
    nvertices(x::HMesh) = length(x.vertices)
    ncells(x::HMesh) = length(x.cells)
    
    cells(mesh::HMesh) = setdiff(axes(mesh.cells, 1), findall(>(0), mesh.levels))
    vertices(mesh::HMesh, cell) = mesh.cells[cell]
    
    function init_hgrid(m::Int, n::Int = m, v0 = (0., 0.), v1 = (1., 1.))
        r1 = LinRange(v0[1], v1[1], m + 1)
        r2 = LinRange(v0[2], v1[2], n + 1)
        coords = collect(Iterators.product(r1, r2))
    
        vertices = SVector{2, Float64}.(vec(coords))
        cells = Vector{NTuple{3, Int}}()
    
        vidx = LinearIndices((m + 1, n + 1))
        e1 = CartesianIndex(1, 0)
        e2 = CartesianIndex(0, 1)
        k = 0
        for I in CartesianIndices((m, n))
    	push!(cells, (vidx[I], vidx[I + e1], vidx[I + e1 + e2]))
    	push!(cells, (vidx[I], vidx[I + e1 + e2], vidx[I + e2]))
        end
    
        return HMesh(vertices, cells, zeros(Int, axes(cells)))
    end
    
    function init_hgrid(img::Array{<:Any, 2}; type=:vertex)
        s = (size(img, 2), size(img, 1))
        type == :vertex ?
            init_hgrid((s .- 1)..., (1.0, 1.0), s) :
    	init_hgrid(s..., (0.5, 0.5), s .- (0.5, 0.5))
    end
    
    function sub_mesh(hmesh::HMesh, subcells = cells(hmesh))
        cells = collect(reshape(reinterpret(Int, hmesh.cells[subcells]), 3, :))
        # TODO: restrict vertices too
        vertices = collect(reshape(reinterpret(Float64, hmesh.vertices), 2, :))
        return Mesh(vertices, cells)
    end
    
    
    # 2d, simplex grid
    struct Mesh
        vertices::Array{Float64, 2}
        cells::Array{Int, 2}
    end
    
    # needs to be after Mesh definition
    function HMesh(mesh::Mesh)
        vertices = Vector{SVector{2, Float64}}()
        for v in axes(mesh.vertices, 2)
    	push!(vertices, mesh.vertices[:, v])
        end
        cells = Vector{NTuple{3, Int}}()
        for c in axes(mesh.cells, 2)
    	push!(cells, NTuple{3}(mesh.cells[:, c]))
        end
        levels = zeros(Int, axes(cells))
    
        return HMesh(vertices, cells, levels)
    end
    
    Base.show(io::IO, ::MIME"text/plain", x::Mesh) =
        print("$(nameof(typeof(x))), $(ncells(x)) cells")
    
    ndims_domain(::Mesh) = 2
    ndims_space(::Mesh) = 2
    nvertices_cell(::Mesh) = 3
    nvertices(x::Mesh) = size(x.vertices, 2)
    ncells(x::Mesh) = size(x.cells, 2)
    
    cells(mesh::Mesh) = axes(mesh.cells, 2)
    vertices(mesh::Mesh) = axes(mesh.vertices, 2)
    vertices(mesh::Mesh, cell) = ntuple(i -> mesh.cells[i, cell], nvertices_cell(mesh))
    
    #function Base.getproperty(obj::Mesh, sym::Symbol)
    #    if sym === :vertices
    #	return reinterpret(SVector{ndims_space(obj), Float64}, vec(obj.vertices))
    #    else
    #	return getfield(obj, sym)
    #    end
    #end
    
    function init_grid_old(m::Int, n::Int = m, v0 = (0., 0.), v1 = (1., 1.))
        r1 = LinRange(v0[1], v1[1], m + 1)
        r2 = LinRange(v0[2], v1[2], n + 1)
        coords = collect(Iterators.product(r1, r2))
    
        vertices = [x[i] for i in 1:2, x in vec(coords)]
        cells = Array{Int, 2}(undef, 3, 2 * m * n)
    
        vidx = LinearIndices((m + 1, n + 1))
        e1 = CartesianIndex(1, 0)
        e2 = CartesianIndex(0, 1)
        k = 0
        for I in CartesianIndices((m, n))
    	cells[:, k += 1] .= (vidx[I], vidx[I + e1], vidx[I + e1 + e2])
    	cells[:, k += 1] .= (vidx[I], vidx[I + e1 + e2], vidx[I + e2])
        end
    
        return Mesh(vertices, cells)
    end
    
    function init_grid(m::Int, n::Int = m, v0 = (0., 0.), v1 = (1., 1.))
        r1 = LinRange(v0[1], v1[1], m + 1)
        r2 = LinRange(v0[2], v1[2], n + 1)
        coords = collect(Iterators.product(r1, r2))
    
        vertices = [x[i] for i in 1:2, x in vec(coords)]
        cells = Array{Int, 2}(undef, 3, 2 * m * n)
    
        vidx = zeros(Int, m + 1, n + 1)
        k = 0
        for j = 1 : n + 1
    	for i = 1 : m + 1
    	    if iseven(i + j)
    		vidx[i, j] = k += 1
    	    end
    	end
        end
        for j = 1 : n + 1
    	for i = 1 : m + 1
    	    if isodd(i + j)
    		vidx[i, j] = k += 1
    	    end
    	end
        end
        vidxinv = reshape(invperm(vec(vidx)), m + 1, n + 1)
        vertices = reshape(vertices[:, vidxinv], 2, :)
    
        e1 = CartesianIndex(1, 0)
        e2 = CartesianIndex(0, 1)
        k = 0
        for I in CartesianIndices((m, n))
    	if iseven(I[1] + I[2])
    	    cells[:, k += 1] .= (vidx[I], vidx[I + e1], vidx[I + e1 + e2])
    	    cells[:, k += 1] .= (vidx[I], vidx[I + e1 + e2], vidx[I + e2])
    	else
    	    cells[:, k += 1] .= (vidx[I], vidx[I + e1], vidx[I + e2])
    	    cells[:, k += 1] .= (vidx[I + e1], vidx[I + e1 + e2], vidx[I + e2])
    	end
        end
    
        return Mesh(vertices, cells)
    end
    
    function init_grid(img::Array{<:Any, 2}; type=:vertex)
        s = (size(img, 2), size(img, 1))
        type == :vertex ?
            init_grid((s .- 1)..., (1.0, 1.0), s) :
    	init_grid(s..., (0.5, 0.5), s .- (0.5, 0.5))
    end
    
    function init_grid(img::Array{<:Any, 2}, m::Int, n::Int = m; type=:vertex)
        s = (size(img, 2), size(img, 1))
        type == :vertex ?
            init_grid(((m, n) .- 1)..., (1.0, 1.0), s) :
            init_grid((m, n)..., (0.5, 0.5), s .- (0.5, 0.5))
    end
    
    # horribly implemented, please don't curse me
    function bisect!(mesh::HMesh, marked_cells::Set)
        refined_cells = Pair{Int, NTuple{2, Int}}[]
        # assemble edge -> cells map
        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)
        end
    
        #return edgemap
    
        function refine_cell(c1)
    	c2 = -1
    	# c1 -> c11 + c12
    	# c2 -> c21 + c22
    	c1_vs = sort(SVector(vertices(mesh, c1)))
    	c2_arr = setdiff(edgemap[(c1_vs[1], c1_vs[2])], c1)
    
    	if !isempty(c2_arr)
    	    @assert(length(c2_arr) == 1)
    	    c2 = c2_arr[begin]
    	    c2_vs = sort(SVector(vertices(mesh, c2)))
    
    	    if c1_vs[1:2] != c2_vs[1:2]
    		# cannot refine `cellop` compatibly, recurse
    		refine_cell(c2)
    		# refetch c2 because topology has changed
    		c2_arr = setdiff(edgemap[(c1_vs[1], c1_vs[2])], c1)
    		@assert(!isempty(c2_arr))
    		c2 = c2_arr[begin]
    		c2_vs = sort(SVector(vertices(mesh, c2)))
    	    end
    	    # c2 exists and compatibility is guaranteed
    	    @assert(c1_vs[1:2] == c2_vs[1:2])
    	end
    
    	# create bisection vertex
    	xbisect = (mesh.vertices[c1_vs[1]] + mesh.vertices[c1_vs[2]]) / 2
    	push!(mesh.vertices, xbisect)
    	vbisect = lastindex(mesh.vertices)
    
    	# take care to produce positively oriented cells
    	push!(mesh.cells, Tuple(replace(SVector(vertices(mesh, c1)), c1_vs[1] => vbisect)))
    	push!(mesh.levels, 0)
    	c3 = lastindex(mesh.cells)
    	@assert(length(setdiff(c1_vs, c1_vs[1])) == 2)
    	replace!(edgemap[NTuple{2}(setdiff(c1_vs, c1_vs[1]))], c1 => c3)
    
    	# take care to produce positively oriented cells
    	push!(mesh.cells, Tuple(replace(SVector(vertices(mesh, c1)), c1_vs[2] => vbisect)))
    	push!(mesh.levels, 0)
    	c4 = lastindex(mesh.cells)
    	@assert(length(setdiff(c1_vs, c1_vs[2])) == 2)
    	replace!(edgemap[NTuple{2}(setdiff(c1_vs, c1_vs[2]))], c1 => c4)
    
            # flip is correct
    	edgemap[(c1_vs[1], vbisect)] = [c4]
    	edgemap[(c1_vs[2], vbisect)] = [c3]
    	edgemap[(c1_vs[3], vbisect)] = [c3, c4]
    
    	mesh.levels[c1] += 1
    	delete!(marked_cells, c1)
    	push!(refined_cells, c1 => (c3, c4))
    
    	if c2 > 0
    	    # take care to produce positively oriented cells
    	    push!(mesh.cells, Tuple(replace(SVector(vertices(mesh, c2)), c1_vs[1] => vbisect))) # c1_vs is correct
    	    push!(mesh.levels, 0)
    	    c5 = lastindex(mesh.cells)
    	    @assert(length(setdiff(c2_vs, c1_vs[1])) == 2)
    	    replace!(edgemap[NTuple{2}(setdiff(c2_vs, c1_vs[1]))], c2 => c5)
    
    	    # take care to produce positively oriented cells
    	    push!(mesh.cells, Tuple(replace(SVector(vertices(mesh, c2)), c1_vs[2] => vbisect))) # c1_vs is correct
    	    push!(mesh.levels, 0)
    	    c6 = lastindex(mesh.cells)
    	    @assert(length(setdiff(c2_vs, c1_vs[2])) == 2)
    	    replace!(edgemap[NTuple{2}(setdiff(c2_vs, c1_vs[2]))], c2 => c6)
    
    	    # flip is correct
    	    push!(edgemap[(c1_vs[1], vbisect)], c6)
    	    push!(edgemap[(c1_vs[2], vbisect)], c5)
    	    edgemap[(c2_vs[3], vbisect)] = [c5, c6]
    
    	    mesh.levels[c2] += 1
    	    delete!(marked_cells, c2)
    	    push!(refined_cells, c2 => (c5, c6))
    	end
        end
    
        while !isempty(marked_cells)
    	refine_cell(first(marked_cells))
        end
    
        return refined_cells
        #deleteat!(mesh.cells, sort(collect(refined_cells)))
    end
    
    # horribly implemented, please don't curse me
    function refine!(hmesh::HMesh, marked_cells::Set; fs...)
        old_mesh = sub_mesh(hmesh)
        refined_cells = bisect!(hmesh, marked_cells)
        extended_mesh = sub_mesh(hmesh, axes(hmesh.cells, 1))
        new_mesh = sub_mesh(hmesh)
        removed_cells = map(x -> first(x), refined_cells)
    
        # extended functions onto newly created cells
        extended_fs = map(NamedTuple(fs)) do f
    	space = FeSpace(extended_mesh, f.space.element, f.space.size)
    	return FeFunction(space; f.name)
        end
        # copy over previous data for unmodified cells
        for (f, extended_f) in zip(NamedTuple(fs), extended_fs)
    	copyto!(extended_f.data, f.data)
        end
        # prolong data for refined cells
        for (old_cell, extended_cells) in refined_cells
    	for extended_f in extended_fs
    	    prolong!(extended_f, old_cell, extended_cells)
    	end
        end
    
        # retain only non-refined cells
        new_fs = map(NamedTuple(extended_fs)) do f
    	space = FeSpace(new_mesh, f.space.element, f.space.size)
    	return FeFunction(space; f.name)
        end
        retained_cells = setdiff(cells(extended_mesh), removed_cells)
        @assert(retained_cells == cells(hmesh))
        for (new_cell, old_cell) in enumerate(retained_cells)
    	for (f, new_f) in zip(extended_fs, new_fs)
    	    gdofs = f.space.dofmap[:, :, old_cell]
    	    new_gdofs = new_f.space.dofmap[:, :, new_cell]
    	    new_f.data[new_gdofs] .= f.data[gdofs]
    	end
        end
        return new_fs
    end
    
    "refine by creating temporary hierarchical mesh on the fly"
    function refine(mesh::Mesh, marked_cells; fs...)
        hmesh = HMesh(mesh)
        fs_new = refine!(hmesh, Set(marked_cells); fs...)
        mesh_new = sub_mesh(hmesh)
        return mesh_new, fs_new
    end
    
    function geo_tolocal(A, v)
        J = jacobian(x -> A * [1 - x[1] - x[2], x[1], x[2]], [0., 0.])
        return J \ (v - A[:, 1])
    end
    
    function geo_contains(A, v)
        J = jacobian(x -> A * [1 - x[1] - x[2], x[1], x[2]], [0., 0.])
        λ = J \ (v - A[:, 1])
        return all(λ .>= 0) && sum(λ) <= 1
    end
    
    #function cell_contains(mesh, cell, v)
    #    geo = mesh.vertices[:, mesh.cells[:, cell]]
    #    J = jacobian(x -> geo * [1 - x[1] - x[2], x[1], x[2]], [0., 0.])
    #    λ = J \ (v - geo[:, 1])
    #    return all(λ .>= 0) && sum(λ) <= 1
    #end
    
    function vtk_mesh(filename, mesh::Mesh)
        cells = [MeshCell(VTKCellTypes.VTK_TRIANGLE, 3*(i-1)+1:3*(i-1)+3)
    	for i in axes(mesh.cells, 2)]
        vertices = reshape(mesh.vertices[:, mesh.cells], 2, :)
    
        return vtk_grid(filename, vertices, cells)
    end
    
    "convenience function for saving to vtk"
    function save(filename::String, mesh::Mesh, fs...)
        vtk = vtk_mesh(filename, mesh)
        for f in fs
    	f.space.mesh == mesh ||
    	    throw(ArgumentError("meshes do not match"))
            vtk_append!(vtk, f)
        end
        vtk_save(vtk)
    end
    
    function area(mesh, cell)
        A = SArray{Tuple{ndims_space(mesh), nvertices_cell(mesh)}}(
    	view(mesh.vertices, :, view(mesh.cells, :, cell)))
        return det(A[:, 2:3] .- A[:, 1]) / 2
    end
    
    area(mesh) = mapreduce(cell -> area(mesh, cell), +, cells(mesh))
    
    function diam(mesh, cell)
        A = SArray{Tuple{ndims_space(mesh), nvertices_cell(mesh)}}(
    	view(mesh.vertices, :, view(mesh.cells, :, cell)))
        return max(
    	norm(A[:, 1] - A[:, 2]),
    	norm(A[:, 2] - A[:, 3]),
    	norm(A[:, 3] - A[:, 1]))
    end
    
    mesh_size(mesh) = mapreduce(cell -> diam(mesh, cell), max, cells(mesh))
    
    function elmap(mesh, cell)
        # TODO: can be improved
        A = SArray{Tuple{ndims_space(mesh), nvertices_cell(mesh)}}(
    	view(mesh.vertices, :, view(mesh.cells, :, cell)))
        return x -> A * SA[1 - x[1] - x[2], x[1], x[2]]
    end
    
    function test_mesh(mesh)
        # assemble edge -> cells map
        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)
        end
        for (edge, cells) in edgemap
    	@assert(1 <= length(cells) <= 2)
        end
        # are cells positively oriented?
        for cell in cells(mesh)
    	delmap = jacobian(elmap(mesh, cell), SA[0., 0.])
    	@assert(det(delmap) > 0)
        end
    end