Makie JuAFEM Plot

gistfile1.txt

using Makie
using JuAFEM, SparseArrays
using LinearAlgebra

#
# Poisson example from JuAFEM docs
#
grid = generate_grid(Triangle, (20, 20));

dim = 2
ip = Lagrange{dim, RefTetrahedron, 1}()
qr = QuadratureRule{dim, RefTetrahedron}(2)
cellvalues = CellScalarValues(qr, ip);

dh = DofHandler(grid)
push!(dh, :u, 1)
close!(dh);

K = create_sparsity_pattern(dh);

fill!(K.nzval, 1.0)
# using UnicodePlots
# spy(K; height = 15)

ch = ConstraintHandler(dh);

∂Ω = union(getfaceset.((grid, ), ["left", "right", "top", "bottom"])...);

dbc = Dirichlet(:u, ∂Ω, (x, t) -> 0)
add!(ch, dbc);

close!(ch)
JuAFEM.update!(ch, 0.0);

function doassemble(cellvalues::CellScalarValues{dim}, K::SparseMatrixCSC, dh::DofHandler) where {dim}

    n_basefuncs = getnbasefunctions(cellvalues)
    Ke = zeros(n_basefuncs, n_basefuncs)
    fe = zeros(n_basefuncs)

    f = zeros(ndofs(dh))
    assembler = start_assemble(K, f)

    @inbounds for cell in CellIterator(dh)

        fill!(Ke, 0)
        fill!(fe, 0)

        reinit!(cellvalues, cell)

        for q_point in 1:getnquadpoints(cellvalues)
            dΩ = getdetJdV(cellvalues, q_point)

            for i in 1:n_basefuncs
                v  = shape_value(cellvalues, q_point, i)
                ∇v = shape_gradient(cellvalues, q_point, i)
                fe[i] += v * dΩ
                for j in 1:n_basefuncs
                    ∇u = shape_gradient(cellvalues, q_point, j)
                    Ke[i, j] += (∇v ⋅ ∇u) * dΩ
                end
            end
        end

        assemble!(assembler, celldofs(cell), fe, Ke)
    end
    return K, f
end

K, f = doassemble(cellvalues, K, dh);

apply!(K, f, ch)
u = K \ f;

# 
# Now here begins the actual plotting
# 


N = length(CellIterator(dh))
coords = Matrix{Float64}(undef, 3N, 2)
vals = Vector{Float64}(undef, 3N)
faces = Matrix{Int}(undef, N, 3)
i = 0
for (icell, cell) in enumerate(CellIterator(dh))
    cdofs = celldofs(cell)
    faces[icell, :] .= [i+1, i+2, i+3]
    for j in 1:3
        i += 1
        vals[i] = u[cdofs[j]]
        inode = getnodes(cell)[j]
        coords[i,:] .= getcoordinates(cell)[j]
    end
end

scene = mesh(coords, faces, color=normalize(vals, Inf), shading=false)
display(wireframe!(scene[end][1]))