Plot vector field with Nedelec element

[learning-chip]

In order to visualize some Maxwell simulation results (without writing to VTK), I hacked some code to plot vector field with zeroth-order Nedelec element on quadrilateral mesh. Would it be useful to add to GridapMakie.jl?

Currently, plotting a field with Nedelec element seems unsupported:

using Gridap
using GridapMakie
using CairoMakie

n_1d = 16
domain = (0, 1, 0, 1)
partition = (n_1d, n_1d)
model = CartesianDiscreteModel(domain, partition) #|> simplexify
# model = simplexify(model)  # leads to `This function is not yet implemented`

reffe = ReferenceFE(nedelec, 0)
V = FESpace(model, reffe)
uh = interpolate(x->VectorValue(x[1], x[2]), V)
plot(uh)  # empty figure, nothing shows up

• Without the simplexify call, plot() returns empty figure. Probably because the plotting logic assumes unstructured mesh instead of Cartesian mesh?
• With the simplexify call, the later line V = FESpace(model, reffe) leads to This function is not yet implemented. Not sure if related to Nedelec elements for triangular/tetrahedral meshes? Gridap.jl#413

Error message:

This function is not yet implemented

Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:33
  [2] macro expansion
    @ ~/.julia/packages/Gridap/775Gn/src/Helpers/Macros.jl:21 [inlined]
  [3] _Nedelec_nodes_and_moments(#unused#::Type{Float64}, p::Gridap.ReferenceFEs.ExtrusionPolytope{2}, order::Int64)
    @ Gridap.ReferenceFEs ~/.julia/packages/Gridap/775Gn/src/ReferenceFEs/NedelecRefFEs.jl:73
  [4] NedelecRefFE(#unused#::Type{Float64}, p::Gridap.ReferenceFEs.ExtrusionPolytope{2}, order::Int64)
    @ Gridap.ReferenceFEs ~/.julia/packages/Gridap/775Gn/src/ReferenceFEs/NedelecRefFEs.jl:20
  [5] ReferenceFE
    @ ~/.julia/packages/Gridap/775Gn/src/ReferenceFEs/NedelecRefFEs.jl:45 [inlined]
  [6] #91
    @ ~/.julia/packages/Gridap/775Gn/src/Geometry/DiscreteModels.jl:386 [inlined]
  [7] iterate
    @ ./generator.jl:47 [inlined]
  [8] _collect(c::Vector{Gridap.ReferenceFEs.ExtrusionPolytope{2}}, itr::Base.Generator{Vector{Gridap.ReferenceFEs.ExtrusionPolytope{2}}, Gridap.Geometry.var"#91#92"{Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, Tuple{Nedelec, Int64}}}, #unused#::Base.EltypeUnknown, isz::Base.HasShape{1})
    @ Base ./array.jl:695
  [9] collect_similar(cont::Vector{Gridap.ReferenceFEs.ExtrusionPolytope{2}}, itr::Base.Generator{Vector{Gridap.ReferenceFEs.ExtrusionPolytope{2}}, Gridap.Geometry.var"#91#92"{Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, Tuple{Nedelec, Int64}}})
    @ Base ./array.jl:606
 [10] map(f::Function, A::Vector{Gridap.ReferenceFEs.ExtrusionPolytope{2}})
    @ Base ./abstractarray.jl:2294
 [11] ReferenceFE(::Gridap.Geometry.UnstructuredDiscreteModel{2, 2, Float64, Gridap.Geometry.Oriented}, ::Nedelec, ::Vararg{Any, N} where N; kwargs::Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ Gridap.Geometry ~/.julia/packages/Gridap/775Gn/src/Geometry/DiscreteModels.jl:386
 [12] ReferenceFE
    @ ~/.julia/packages/Gridap/775Gn/src/Geometry/DiscreteModels.jl:384 [inlined]
 [13] FESpace(model::Gridap.Geometry.UnstructuredDiscreteModel{2, 2, Float64, Gridap.Geometry.Oriented}, reffe::Tuple{Nedelec, Tuple{Int64}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}; kwargs::Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ Gridap.FESpaces ~/.julia/packages/Gridap/775Gn/src/FESpaces/FESpaceFactories.jl:125
 [14] FESpace(model::Gridap.Geometry.UnstructuredDiscreteModel{2, 2, Float64, Gridap.Geometry.Oriented}, reffe::Tuple{Nedelec, Tuple{Int64}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}})
    @ Gridap.FESpaces ~/.julia/packages/Gridap/775Gn/src/FESpaces/FESpaceFactories.jl:124
 [15] top-level scope
    @ In[4]:2
 [16] eval
    @ ./boot.jl:360 [inlined]
 [17] include_string(mapexpr::typeof(REPL.softscope), mod::Module, code::String, filename::String)
    @ Base ./loading.jl:1116

Manually plot vector field with Nedelec element

This looks hacky and only works for zeroth-order Nedelec element + quadrilateral mesh. 2D and 3D are both okay. For high-order elements and other meshes, need a more generic way to find the location of each DoF and the orientation of each edge.

using Gridap
using CairoMakie

n_1d = 16
domain = (0, 1, 0, 1)
partition = (n_1d, n_1d)
model = CartesianDiscreteModel(domain, partition)

# get x and y coordinates of each node
node_coords = model.grid.node_coords
node_x = map(coords -> coords[1], node_coords)
node_y = map(coords -> coords[2], node_coords)

reffe = ReferenceFE(nedelec, 0)
V = FESpace(model, reffe)
uh = interpolate(x -> VectorValue(x[1] + x[2], x[2] - x[1]), V)

# `uh.cell_dof_values` is an array of size-4 vectors, each corresponds to an edge DoF of a cell 
# `map(v -> v[i], uh.cell_dof_values)` extracts the i-th edge of all cells, i=1,2,3,4
u_comps = [map(v -> v[i], uh.cell_dof_values) for i=1:4]  # extract all edge DoF values
u_comps_2d = map(u -> reshape(u, (n_1d, n_1d)), u_comps)  # reshape to 2D grid structure

# The order of edges surrounding a cell should be:
# - 2 -
# 3   4
# - 1 -
# which can be verified by
@assert u_comps_2d[1][:, 2:end] == u_comps_2d[2][:, 1:end-1]
@assert u_comps_2d[3][2:end,:] == u_comps_2d[4][1:end-1,:]

x_1d = node_x[1:end-1, 1]  # remove boundary node
y_1d = node_y[1, 1:end-1]

# ignore edges on east and north boundaries:
ux = u_comps_2d[1]
uy = u_comps_2d[3]
x_1d = node_x[1:end-1, 1]
y_1d = node_y[1, 1:end-1]

heatmap(ux)  # plot a component

strength = sqrt.(ux.^2 + uy.^2)
heatmap(strength)  # plot vector norm

# plot vector field
# https://makie.juliaplots.org/stable/examples/plotting_functions/arrows/
strength_1d = vec(strength)
arrows(
    x_1d, y_1d, ux, uy,
    lengthscale=0.05,
    arrowcolor=strength_1d, linecolor=strength_1d
)

The resulting plot:

Makie arrows() also works for 3D case.

An additional hack, useful when using external linear solver to get approximate solution x, and to visualize such solution (useful for checking the results of an approximate/iterative solver).

# assume `op` is an assembled `AffineFEOperator` using `ReferenceFE(nedelec, 0)`
uh = solve(op)  # reference solution, just need "metadata" here, not actual values...
A = get_matrix(op)
b = get_vector(op)
x = A \ b  # or any other linear solver output
uh.free_values[:] = x  # overwrite free values (without duplicate values)
uh.cell_dof_values  # cell DoFs are changed accordingly, and can be used in the above script for visualization

A better way would be straightly figure out the location and orientation of each edge corresponding to free values, without querying uh.cell_dof_values. But I haven't figured out how to do that...

I'd like to clean up and improve those hacks, and hopefully contribute a more generic solution for vector plotting. Suggestions are welcome :)

[learning-chip]

A slightly cleaned-up script:

using Gridap
using CairoMakie


"""Extract node coordinates as 1D arrays from 2D regular mesh"""
function get_node_coords_2d(model)
    node_coords = model.grid.node_coords
    node_x = map(coords -> coords[1], node_coords)
    node_y = map(coords -> coords[2], node_coords)
    x_1d = node_x[1:end-1, 1]  # remove boundary node
    y_1d = node_y[1, 1:end-1]
    return x_1d, y_1d
end


"""Extract vector components from 2D nedelec field"""
function get_nedelec_comp_2d(uh)
    u_comps = [map(v -> v[i], uh.cell_dof_values) for i=1:4];
    u_comps_2d = map(u -> reshape(u, (n_1d, n_1d)), u_comps);
    ux = u_comps_2d[1]
    uy = u_comps_2d[3]
    return ux, uy
end


"""Plot separately two components of 2D nedelec field"""
function plot_nedelec_comp_2d(uh)
    model = uh.cell_field.trian.model
    x_1d, y_1d = get_node_coords_2d(model)
    ux, uy = get_nedelec_comp_2d(uh)

    fig = Figure(resolution=(600, 300))
    heatmap(fig[1, 1], x_1d, y_1d, ux)
    heatmap(fig[1, 2], x_1d, y_1d, uy)
    # TODO: add colorbar
    return fig
end


"""Plot arrows for 2D nedelec field """
function plot_nedelec_vector(uh; lengthscale=0.05)
    model = uh.cell_field.trian.model
    x_1d, y_1d = get_node_coords_2d(model)
    ux, uy = get_nedelec_comp_2d(uh)
    strength = sqrt.(ux.^2 + uy.^2)
    strength_1d = vec(strength)
    arrows(
        x_1d, y_1d, ux, uy,
        lengthscale=lengthscale,
        arrowcolor=strength_1d,
        linecolor=strength_1d
    )
end


"""Convert plain array of free values to visualizable field object"""
function free_to_field(x_free, V)
    xh = interpolate(x -> VectorValue(0, 0), V)
    xh.free_values[:] = x_free
    return xh
end

Example usage:

n_1d = 16
domain = (0, 1, 0, 1)
partition = (n_1d, n_1d)
model = CartesianDiscreteModel(domain, partition)
reffe = ReferenceFE(nedelec, 0)
V = FESpace(model, reffe)
uh = interpolate(x -> VectorValue(x[1] + x[2], x[2] - x[1]), V)

plot_nedelec_comp_2d(uh)  # plot separate components
plot_nedelec_vector(uh)  # plot vector field

# plot plain array of free values
x = rand(length(uh.free_values))  # some array, possibly from external solver
xh = free_to_field(x, V)
plot_nedelec_vector(xh)  # works exactly the same as for `uh` above

Still hacky with unnecessary duplicated calculations, but at least usable😅

[fverdugo]

Hi @learning-chip Thanks for sharing this nice example!

Which part of the code do you think would make sense to add to GridapMakie?

As a general rule, I would add code that is general rather than code for particular examples. I.e., code that is general enough and can be used by others in different contexts.

[learning-chip]

add code that is general

Then I would probably wait until I figure out a way to plot higher-order Nedelec elements on more general meshes. Just leave the current code here in case someone needs it