GalerkinToolkit · fverdugo · Oct 11, 2024 · Oct 9, 2024 · Oct 9, 2024
diff --git a/Project.toml b/Project.toml
@@ -18,13 +18,14 @@ Metis = "2679e427-3c69-5b7f-982b-ece356f1e94b"
 PartitionedArrays = "5a9dfac6-5c52-46f7-8278-5e2210713be9"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
-WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
 
 [weakdeps]
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
+WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
 
 [extensions]
 GalerkinToolkitMakieExt = "Makie"
+GalerkinToolkitWriteVTKExt = "WriteVTK"
 
 [compat]
 Combinatorics = "1"
@@ -33,16 +34,17 @@ ForwardDiff = "0.10"
 Gmsh = "0.3"
 IterativeSolvers = "0.9"
 MPI = "0.20"
+Makie = "0.21"
+WriteVTK = "1.21"
 Metis = "1"
 PartitionedArrays = "0.5.4"
 StaticArrays = "1"
-WriteVTK = "1"
-Makie = "0.21"
 julia = "1"
 
 [extras]
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a"
+WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "GLMakie"]
+test = ["Test", "GLMakie", "WriteVTK"]
diff --git a/ext/GalerkinToolkitWriteVTKExt.jl b/ext/GalerkinToolkitWriteVTKExt.jl
@@ -0,0 +1,326 @@
+module GalerkinToolkitWriteVTKExt
+
+using GalerkinToolkit
+import GalerkinToolkit: plot, plot!, translate_vtk_data, translate_vtk_data!, vtk_close_impl, vtk_close_impl!
+import GalerkinToolkit: vtk_points, vtk_points!, vtk_cells, vtk_cells!, vtk_args, vtk_args!
+import GalerkinToolkit: vtk_physical_faces, vtk_physical_faces!, vtk_physical_nodes, vtk_physical_nodes!
+import GalerkinToolkit: vtk_mesh_cell, vtk_mesh_cell!
+using PartitionedArrays
+using StaticArrays
+using WriteVTK
+
+# Visualization
+
+function translate_vtk_data(v)
+    v
+end
+
+function translate_vtk_data(v::AbstractVector{<:SVector{2}})
+    z = zero(eltype(eltype(v)))
+    map(vi->SVector((vi...,z)),v)
+end
+
+function plot!(field,plt::GalerkinToolkit.VTKPlot;label)
+    plot!(plt.plot,field;label)
+end
+
+function plot!(plt::GalerkinToolkit.VTKPlot,field;label)
+    plot!(plt.plot,field;label)
+end
+
+function WriteVTK.vtk_grid(filename,plt::GalerkinToolkit.Plot;kwargs...)
+    mesh = plt.mesh
+    D = GalerkinToolkit.num_dims(mesh)
+    vtk = vtk_grid(filename,GalerkinToolkit.vtk_args(mesh)...;kwargs...)
+    for (k,v) in GalerkinToolkit.node_data(plt)
+        vtk[k,WriteVTK.VTKPointData()] = translate_vtk_data(v)
+    end
+    for (k,v) in GalerkinToolkit.face_data(plt;merge_dims=true)
+        vtk[k,WriteVTK.VTKCellData()] = translate_vtk_data(v)
+    end
+    GalerkinToolkit.VTKPlot(plt,vtk)
+end
+
+function WriteVTK.vtk_grid(f::Function,filename,plt::Union{GalerkinToolkit.Plot,GalerkinToolkit.PPlot};kwargs...)
+    vtk = vtk_grid(filename,plt)
+    files = nothing
+    try
+        f(vtk)
+    finally
+        files = close(vtk)
+    end
+    files
+end
+
+function WriteVTK.vtk_grid(filename,pplt::GalerkinToolkit.PPlot;kwargs...)
+    plts = partition(pplt)
+    parts = linear_indices(plts)
+    nparts = length(parts)
+    vtks  = map(plts,parts) do plt,part
+        mesh = plt.mesh
+        vtk = pvtk_grid(filename,GalerkinToolkit.vtk_args(mesh)...;part,nparts,kwargs...)
+        for (k,v) in GalerkinToolkit.node_data(plt)
+            vtk[k,WriteVTK.VTKPointData()] = translate_vtk_data(v)
+        end
+        for (k,v) in GalerkinToolkit.face_data(plt;merge_dims=true)
+            vtk[k,WriteVTK.VTKCellData()] = translate_vtk_data(v)
+        end
+        vtk
+    end
+    plts = GalerkinToolkit.PPlot(plts,pplt.cache)
+    GalerkinToolkit.VTKPlot(plts,vtks)
+end
+
+function WriteVTK.close(plt::GalerkinToolkit.VTKPlot)
+    vtk_close_impl(plt.plot,plt.vtk)
+end
+
+function vtk_close_impl(plt::GalerkinToolkit.Plot,vtk)
+    WriteVTK.close(vtk)
+end
+
+function vtk_close_impl(plt::GalerkinToolkit.PPlot,vtks)
+    map(WriteVTK.close,vtks)
+end
+
+function WriteVTK.vtk_grid(filename,mesh::Union{GalerkinToolkit.AbstractMesh,GalerkinToolkit.PMesh};plot_params=(;),vtk_grid_params...)
+    plt = plot(mesh;plot_params...)
+    WriteVTK.vtk_grid(filename,plt;vtk_grid_params...)
+end
+
+function WriteVTK.vtk_grid(filename,dom::GalerkinToolkit.AbstractDomain;plot_params=(;),vtk_grid_params...)
+    plt = plot(dom;plot_params...)
+    WriteVTK.vtk_grid(filename,plt;vtk_grid_params...)
+end
+
+function WriteVTK.vtk_grid(f::Function,filename,mesh::Union{GalerkinToolkit.AbstractMesh,GalerkinToolkit.PMesh};plot_params=(;),vtk_grid_params...)
+    plt = plot(mesh;plot_params...)
+    WriteVTK.vtk_grid(f,filename,plt;vtk_grid_params...)
+end
+
+function WriteVTK.vtk_grid(f::Function,filename,dom::GalerkinToolkit.AbstractDomain;plot_params=(;),vtk_grid_params...)
+    plt = plot(dom;plot_params...)
+    WriteVTK.vtk_grid(f,filename,plt;vtk_grid_params...)
+end
+
+function WriteVTK.close(pvd::GalerkinToolkit.PVD)
+    WriteVTK.close(pvd.pvd)
+end
+
+function WriteVTK.close(ppvd::GalerkinToolkit.PPVD)
+    map_main(WriteVTK.close,ppvd.partition)
+end
+
+function Base.setindex!(a::GalerkinToolkit.PVD,plt::GalerkinToolkit.VTKPlot,time)
+    a.pvd[time] = plt.vtk
+end
+
+function Base.setindex!(a::GalerkinToolkit.PPVD,plt::GalerkinToolkit.VTKPlot,time)
+    map_main(a.partition,plt.vtk) do a,vtk
+        a[time] = vtk
+    end
+end
+
+function WriteVTK.paraview_collection(filename,pplt::GalerkinToolkit.Plot;kwargs...)
+    WriteVTK.paraview_collection(filename;kwargs...) |> GalerkinToolkit.PVD
+end
+
+function WriteVTK.paraview_collection(filename,pplt::GalerkinToolkit.PPlot;kwargs...)
+    pvds = map_main(partition(pplt)) do _
+        WriteVTK.paraview_collection(filename;kwargs...)
+    end
+    GalerkinToolkit.PPVD(pvds)
+end
+
+function WriteVTK.paraview_collection(f::Function,filename,plt::Union{GalerkinToolkit.Plot,GalerkinToolkit.PPlot};kwargs...)
+    vtk = WriteVTK.paraview_collection(filename,plt;kwargs...)
+    files = nothing
+    try
+        f(vtk)
+    finally
+        files = close(vtk)
+    end
+    files
+end
+
+function WriteVTK.paraview_collection(filename,mesh::Union{GalerkinToolkit.AbstractMesh,GalerkinToolkit.PMesh};plot_params=(;),vtk_grid_params...)
+    plt = GalerkinToolkit.plot(mesh;plot_params...)
+    WriteVTK.paraview_collection(filename,plt;vtk_grid_params...)
+end
+
+function WriteVTK.paraview_collection(filename,dom::GalerkinToolkit.AbstractDomain;plot_params=(;),vtk_grid_params...)
+    plt = GalerkinToolkit.plot(dom;plot_params...)
+    WriteVTK.paraview_collection(filename,plt;vtk_grid_params...)
+end
+
+function WriteVTK.paraview_collection(f::Function,filename,mesh::Union{GalerkinToolkit.AbstractMesh,GalerkinToolkit.PMesh};plot_params=(;),vtk_grid_params...)
+    plt = GalerkinToolkit.plot(mesh;plot_params...)
+    WriteVTK.paraview_collection(f,filename,plt;vtk_grid_params...)
+end
+
+function WriteVTK.paraview_collection(f::Function,filename,dom::GalerkinToolkit.AbstractDomain;plot_params=(;),vtk_grid_params...)
+    plt = GalerkinToolkit.plot(dom;plot_params...)
+    WriteVTK.paraview_collection(f,filename,plt;vtk_grid_params...)
+end
+
+# Mesh
+
+function vtk_points(mesh)
+    function barrirer(coords)
+        nnodes = length(coords)
+        points = zeros(3,nnodes)
+        for node in 1:nnodes
+            coord = coords[node]
+            for i in 1:length(coord)
+                points[i,node] = coord[i]
+            end
+        end
+        points
+    end
+    coords = GalerkinToolkit.node_coordinates(mesh)
+    barrirer(coords)
+end
+
+function vtk_cells(mesh,d)
+    function barrirer(face_to_refid,face_to_nodes,refid_mesh_cell)
+        cells = map(face_to_refid,face_to_nodes) do refid, nodes
+            mesh_cell = refid_mesh_cell[refid]
+            if mesh_cell === nothing
+                msg = """
+                Not enough information to visualize this mesh via vtk:
+                vtk_mesh_cell returns nothing for the reference face in position $refid in dimension $d.
+                """
+                error(msg)
+            end
+            mesh_cell(nodes)
+        end
+        cells
+    end
+    face_to_nodes = GalerkinToolkit.face_nodes(mesh,d)
+    face_to_refid = GalerkinToolkit.face_reference_id(mesh,d)
+    refid_refface = GalerkinToolkit.reference_faces(mesh,d)
+    refid_mesh_cell = map(vtk_mesh_cell,refid_refface)
+    barrirer(face_to_refid,face_to_nodes,refid_mesh_cell)
+end
+
+"""
+    args = vtk_args(mesh[,d])
+
+Return the arguments `args` to be passed in final position
+to functions like `WriteVTK.vtk_grid`.
+"""
+function vtk_args(mesh,d)
+    points = vtk_points(mesh)
+    cells = vtk_cells(mesh,d)
+    points, cells
+end
+
+function vtk_args(mesh)
+    points = vtk_points(mesh)
+    D = GalerkinToolkit.num_dims(mesh)
+    allcells = [vtk_cells(mesh,d) for d in 0:D if GalerkinToolkit.num_faces(mesh,d) != 0]
+    cells = reduce(vcat,allcells)
+    points, cells
+end
+
+"""
+"""
+function vtk_physical_faces!(vtk,mesh,d;physical_faces=GalerkinToolkit.physical_faces(mesh,d))
+    ndfaces = GalerkinToolkit.num_faces(mesh,d)
+    for group in physical_faces
+        name,faces = group
+        face_mask = zeros(Int,ndfaces)
+        face_mask[faces] .= 1
+        vtk[name,WriteVTK.VTKCellData()] = face_mask
+    end
+    vtk
+end
+
+function vtk_physical_faces!(vtk,mesh;physical_faces=GalerkinToolkit.physical_faces(mesh))
+    nfaces = sum(GalerkinToolkit.num_faces(mesh))
+    offsets = GalerkinToolkit.face_offset(mesh)
+    D = GalerkinToolkit.num_dims(mesh)
+    data = Dict{String,Vector{Int}}()
+    for d in 0:D
+        for group in physical_faces[d+1]
+            name, = group
+            if !haskey(data,name)
+                face_mask = zeros(Int,nfaces)
+                data[name] = face_mask
+            end
+        end
+    end
+    for d in 0:D
+        for group in physical_faces[d+1]
+            offset = offsets[d+1]
+            name,faces = group
+            face_mask = data[name]
+            face_mask[faces.+offset] .= 1
+        end
+    end
+    for (name,face_mask) in data
+        vtk[name,WriteVTK.VTKCellData()] = face_mask
+    end
+    vtk
+end
+
+function vtk_physical_nodes!(vtk,mesh,d;physical_nodes=GalerkinToolkit.physical_nodes(mesh,d))
+    nnodes = GalerkinToolkit.num_nodes(mesh)
+    for group in physical_nodes
+        name,nodes = group
+        nodes_mask = zeros(Int,nnodes)
+        nodes_mask[nodes] .= 1
+        vtk[name,WriteVTK.VTKPointData()] = nodes_mask
+    end
+    vtk
+end
+
+function vtk_physical_nodes!(vtk,mesh;physical_nodes=GalerkinToolkit.physical_nodes(mesh))
+    D = GalerkinToolkit.num_dims(mesh)
+    for d in 0:D
+        vtk_physical_nodes!(vtk,mesh,d,physical_nodes=physical_nodes[d+1])
+    end
+    vtk
+end
+
+"""
+"""
+function vtk_mesh_cell(ref_face)
+    geom = GalerkinToolkit.geometry(ref_face)
+    d = GalerkinToolkit.num_dims(geom)
+    nnodes = GalerkinToolkit.num_nodes(ref_face)
+    lib_to_user = GalerkinToolkit.lib_to_user_nodes(ref_face)
+    if d == 0 && nnodes == 1
+        cell_type = WriteVTK.VTKCellTypes.VTK_VERTEX
+        vtk_to_lib = [1]
+    elseif d == 1 && (GalerkinToolkit.is_simplex(geom) || GalerkinToolkit.is_n_cube(geom)) && nnodes == 2
+        cell_type = WriteVTK.VTKCellTypes.VTK_LINE
+        vtk_to_lib = [1,2]
+    elseif d == 1 && (GalerkinToolkit.is_simplex(geom) || GalerkinToolkit.is_n_cube(geom)) && nnodes == 3
+        cell_type = WriteVTK.VTKCellTypes.VTK_QUADRATIC_EDGE
+        vtk_to_lib = [1,3,2]
+    elseif d == 2 && GalerkinToolkit.is_n_cube(geom) && nnodes == 4
+        cell_type = WriteVTK.VTKCellTypes.VTK_QUAD
+        vtk_to_lib = [1,2,4,3]
+    elseif d == 2 && GalerkinToolkit.is_simplex(geom) && nnodes == 3
+        cell_type = WriteVTK.VTKCellTypes.VTK_TRIANGLE
+        vtk_to_lib = [1,2,3]
+    elseif d == 2 && GalerkinToolkit.is_simplex(geom) && nnodes == 6
+        cell_type = WriteVTK.VTKCellTypes.VTK_QUADRATIC_TRIANGLE
+        vtk_to_lib = [1,3,6,2,5,4]
+    elseif d == 3 && GalerkinToolkit.is_n_cube(geom) && nnodes == 8
+        cell_type = WriteVTK.VTKCellTypes.VTK_HEXAHEDRON
+        vtk_to_lib = [1,2,4,3,5,6,8,7]
+    elseif d == 3 && GalerkinToolkit.is_simplex(geom) && nnodes == 4
+        cell_type = WriteVTK.VTKCellTypes.VTK_TETRA
+        vtk_to_lib = [1,2,3,4]
+    elseif d == 3 && GalerkinToolkit.is_simplex(geom) && nnodes == 10
+        cell_type = WriteVTK.VTKCellTypes.VTK_QUADRATIC_TETRA
+        vtk_to_lib = [1,3,6,10,2,5,4,7,8,9]
+    else
+        return nothing
+    end
+    nodes -> WriteVTK.MeshCell(cell_type,(nodes[lib_to_user])[vtk_to_lib])
+end
+
+end # module
diff --git a/src/GalerkinToolkit.jl b/src/GalerkinToolkit.jl
@@ -1,7 +1,6 @@
 module GalerkinToolkit
 const GT = GalerkinToolkit
 
-using WriteVTK
 using FastGaussQuadrature
 using StaticArrays
 using LinearAlgebra