WIP: Export of discontinuous fields to VTK

src/Export/VTK.jl

@@ -1,3 +1,4 @@
+include("vtk_discontinuous.jl")
 
 """
     VTKFile(filename::AbstractString, grid::AbstractGrid; kwargs...)
@@ -28,13 +29,21 @@ end
 """
 struct VTKFile{VTK<:WriteVTK.VTKFile}
     vtk::VTK
+    cellnodes::Vector{UnitRange{Int}}
 end
 function VTKFile(filename::String, dh::DofHandler; kwargs...)
+    for sdh in dh.subdofhandlers
+        for ip in sdh.field_interpolations
+            if is_discontinuous(ip)
+                return VTKFile(filename, get_grid(dh); write_discontinuous=true, kwargs...)
+            end
+        end
+    end
     return VTKFile(filename, get_grid(dh); kwargs...)
 end
-function VTKFile(filename::String, grid::AbstractGrid; kwargs...)
-    vtk = create_vtk_grid(filename, grid; kwargs...)
-    return VTKFile(vtk)
+function VTKFile(filename::String, grid::AbstractGrid; write_discontinuous=false, kwargs...)
+    vtk, cellnodes = create_vtk_grid(filename, grid, write_discontinuous; kwargs...)
+    return VTKFile(vtk, cellnodes)
 end
 # Makes it possible to use the `do`-block syntax
 function VTKFile(f::Function, args...; kwargs...)
@@ -46,6 +55,8 @@ function VTKFile(f::Function, args...; kwargs...)
     end
 end
 
+write_discontinuous(vtk::VTKFile) = length(vtk.cellnodes) > 0
+
 Base.close(vtk::VTKFile) = WriteVTK.vtk_save(vtk.vtk)
 
 function Base.show(io::IO, ::MIME"text/plain", vtk::VTKFile)
@@ -189,9 +200,14 @@ function create_vtk_griddata(grid::Grid{dim,C,T}) where {dim,C,T}
     return coords, cls
 end
 
-function create_vtk_grid(filename::AbstractString, grid::Grid{dim,C,T}; kwargs...) where {dim,C,T}
-    coords, cls = create_vtk_griddata(grid)
-    return WriteVTK.vtk_grid(filename, coords, cls; kwargs...)
+function create_vtk_grid(filename::AbstractString, grid::Grid{dim,C,T}, write_discontinuous; kwargs...) where {dim,C,T}
+    if write_discontinuous
+        coords, cls, cellnodes = create_discontinuous_vtk_griddata(grid)
+    else
+        coords, cls = create_vtk_griddata(grid)
+        cellnodes = Vector{UnitRange{Int}}(undef, 0)
+    end
+    return WriteVTK.vtk_grid(filename, coords, cls; kwargs...), cellnodes
 end
 
 function toparaview!(v, x::Vec{D}) where D
@@ -250,7 +266,11 @@ sorted by the nodes in the grid.
 function write_solution(vtk::VTKFile, dh::AbstractDofHandler, u::Vector, suffix="")
     fieldnames = Ferrite.getfieldnames(dh)  # all primary fields
     for name in fieldnames
-        data = _evaluate_at_grid_nodes(dh, u, name, #=vtk=# Val(true))
+        if write_discontinuous(vtk)
+            data = evaluate_at_discontinuous_vtkgrid_nodes(dh, u, name, vtk.cellnodes)
+        else
+            data = _evaluate_at_grid_nodes(dh, u, name, #=vtk=# Val(true))
+        end
         _vtk_write_node_data(vtk.vtk, data, string(name, suffix))
     end
     return vtk

src/Export/vtk_discontinuous.jl

@@ -0,0 +1,81 @@
+function create_discontinuous_vtk_griddata(grid::Grid{dim, C, T}) where {dim, C, T}
+    cls = Vector{WriteVTK.MeshCell}(undef, getncells(grid))
+    cellnodes = Vector{UnitRange{Int}}(undef, getncells(grid))
+    ncoords = sum(nnodes, getcells(grid))
+    coords = zeros(T, dim, ncoords)
+    icoord = 0
+    for cell in CellIterator(grid)
+        CT = getcelltype(grid, cellid(cell))
+        vtk_celltype = Ferrite.cell_to_vtkcell(CT)
+        cell_coords = getcoordinates(cell)
+        n = length(cell_coords)
+        cellnodes[cellid(cell)] = (1:n) .+ icoord
+        vtk_cellnodes = nodes_to_vtkorder(CT((ntuple(i->i+icoord, n))))
+        cls[cellid(cell)] = WriteVTK.MeshCell(vtk_celltype, vtk_cellnodes)
+        for x in cell_coords
+            icoord += 1
+            coords[:, icoord] = x
+        end
+    end
+    return coords, cls, cellnodes
+end
+
+function evaluate_at_discontinuous_vtkgrid_nodes(dh::DofHandler, u::Vector{T}, fieldname::Symbol, cellnodes) where T
+    # Make sure the field exists
+    fieldname ∈ getfieldnames(dh) || error("Field $fieldname not found.")
+    # Figure out the return type (scalar or vector)
+    field_idx = find_field(dh, fieldname)
+    ip = getfieldinterpolation(dh, field_idx)
+    RT = ip isa ScalarInterpolation ? T : Vec{n_components(ip),T}
+    n_c = n_components(ip)
+    vtk_dim = n_c == 2 ? 3 : n_c # VTK wants vectors padded to 3D
+    n_vtk_nodes = maximum(maximum, cellnodes)
+    data = fill(NaN * zero(T), vtk_dim, n_vtk_nodes)
+    # Loop over the subdofhandlers
+    for sdh in dh.subdofhandlers
+        # Check if this sdh contains this field, otherwise continue to the next
+        field_idx = _find_field(sdh, fieldname)
+        field_idx === nothing && continue
+
+        # Set up CellValues with the local node coords as quadrature points
+        CT = getcelltype(sdh)
+        ip = getfieldinterpolation(sdh, field_idx)
+        ip_geo = default_interpolation(CT)
+        local_node_coords = reference_coordinates(ip_geo)
+        qr = QuadratureRule{getrefshape(ip)}(zeros(length(local_node_coords)), local_node_coords)
+        if ip isa VectorizedInterpolation
+            # TODO: Remove this hack when embedding works...
+            cv = CellValues(qr, ip.ip, ip_geo)
+        else
+            cv = CellValues(qr, ip, ip_geo)
+        end
+        drange = dof_range(sdh, field_idx)
+        # Function barrier
+        _evaluate_at_discontinuous_vtkgrid_nodes!(data, sdh, u, cv, drange, RT, cellnodes)
+    end
+    return data
+end
+
+function _evaluate_at_discontinuous_vtkgrid_nodes!(data::Matrix, sdh::SubDofHandler,
+    u::Vector{T}, cv::CellValues, drange::UnitRange, ::Type{RT}, cellnodes) where {T, RT}
+    ue = zeros(T, length(drange))
+    # TODO: Remove this hack when embedding works...
+    if RT <: Vec && function_interpolation(cv) isa ScalarInterpolation
+        uer = reinterpret(RT, ue)
+    else
+        uer = ue
+    end
+    for cell in CellIterator(sdh)
+        # Note: We are only using the shape functions: no reinit!(cv, cell) necessary
+        @assert getnquadpoints(cv) == length(cell.nodes)
+        for (i, I) in pairs(drange)
+            ue[i] = u[cell.dofs[I]]
+        end
+        for (qp, nodeid) in pairs(cellnodes[cellid(cell)])
+            val = function_value(cv, qp, uer)
+            data[1:length(val), nodeid] .= val
+            data[(length(val)+1):end, nodeid] .= 0 # purge the NaN
+        end
+    end
+    return data
+end

src/exports.jl

@@ -165,7 +165,6 @@ export
     write_node_data,
     VTKFileCollection,
     addstep!,
-
 # L2 Projection
     project,
     L2Projector,

test/test_vtk_export.jl

@@ -41,4 +41,61 @@
             @test bytes2hex(open(SHA.sha1, fname*".vtu")) == "31b506bd9729b11992f8bcb79a2191eb65d223bf"
         end
     end
+
+    @testset "discontinuous" begin
+        # Produce a u such that the overall shape is f(x, xc) = 2 * (x[1]^2 - x[2]^2) - (xc[1]^2 - xc[2]^2)
+        # where xc is the center point of the cell. To avoid floating point issues for the hash,
+        # we test that all values are approximately an integer, and round to integers before storing.
+        function calculate_u(dh)
+            f(z) = z[1]^2 - z[2]^2
+            u = zeros(ndofs(dh))
+            ip = Ferrite.getfieldinterpolation(dh, (1, 1)) # Only one subdofhandler and one field.
+            cv = CellValues(QuadratureRule{RefQuadrilateral}(:lobatto, 2), ip)
+            for cell in CellIterator(dh)
+                reinit!(cv, cell)
+                # Cell center
+                xc = sum(getcoordinates(cell)) / getnquadpoints(cv)
+                for q_point in 1:getnquadpoints(cv)
+                    x = spatial_coordinate(cv, q_point, getcoordinates(cell))
+                    for i in 1:getnbasefunctions(cv)
+                        δu = shape_value(cv, q_point, i)
+                        val = δu * (f(x) * 2 - f(xc))
+                        intval = round(Int, val)
+                        # Ensure output unaffected by floating point errors,
+                        # as we will compare vtk output with a hash
+                        @assert abs(val - intval) < sqrt(eps())
+                        u[celldofs(cell)[i]] += intval
+                    end
+                end
+            end
+            return u
+        end
+
+        mktempdir() do tmp
+            nel = 20 # Dimensions assure integer coordinates at nodes and quad cell centers
+            xcorner = nel * ones(Vec{2})
+            grid = generate_grid(Quadrilateral, (nel, nel), -xcorner, xcorner)
+            # Good to keep for comparison:
+            # dh_cont = close!(add!(DofHandler(grid), :u, Lagrange{RefQuadrilateral,1}()))
+            # u_cont = calculate_u(dh_cont)
+            dh_dg = close!(add!(DofHandler(grid), :u, DiscontinuousLagrange{RefQuadrilateral,1}()))
+
+            u_dg = calculate_u(dh_dg)
+
+            testhash = "daf0cbe26ff709705f338526b19881ef5758f16b"
+
+            fname1 = joinpath(tmp, "discont_kwarg")
+            VTKFile(fname1, grid; write_discontinuous=true) do vtk
+                write_solution(vtk, dh_dg, u_dg)
+            end
+            @test bytes2hex(open(SHA.sha1, fname1*".vtu")) == testhash
+
+            fname2 = joinpath(tmp, "discont_auto")
+            VTKFile(fname2, dh_dg) do vtk
+                write_solution(vtk, dh_dg, u_dg)
+            end
+            @test bytes2hex(open(SHA.sha1, fname2*".vtu")) == testhash
+        end
+
+    end
 end