Skip to content

Commit

Permalink
Merge pull request #106 from fverdugo/writevtk_extension_module
Browse files Browse the repository at this point in the history 
Add WriteVTK extension module
  • Loading branch information
@fverdugo
fverdugo authored Oct 11, 2024
2 parents 5f6f5c4 + 33a11d5 commit f440211
Show file tree
Hide file tree
Showing 5 changed files with 351 additions and 313 deletions.
12 changes: 7 additions & 5 deletions Project.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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"]
326 changes: 326 additions & 0 deletions ext/GalerkinToolkitWriteVTKExt.jl
View file
Original file line number Diff line number Diff line change
@@ -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
1 change: 0 additions & 1 deletion src/GalerkinToolkit.jl
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
module GalerkinToolkit
const GT = GalerkinToolkit

using WriteVTK
using FastGaussQuadrature
using StaticArrays
using LinearAlgebra
Expand Down
Loading

0 comments on commit f440211

Please sign in to comment.