diff --git a/src/WriteVTK.jl b/src/WriteVTK.jl index 87ce3b98..8cd6c36e 100644 --- a/src/WriteVTK.jl +++ b/src/WriteVTK.jl @@ -11,6 +11,7 @@ export paraview_collection, collection_add_timestep, paraview_collection_load export vtk_write_array export VTKPointData, VTKCellData, VTKFieldData export PolyData +export pvtk_grid import CodecZlib: ZlibCompressorStream import TranscodingStreams @@ -147,6 +148,9 @@ include("gridtypes/unstructured/polydata.jl") include("gridtypes/array.jl") +# Parallel DataSet Formats +include("gridtypes/pvtk_grid.jl") + # This allows using do-block syntax for generation of VTK files. for func in (:vtk_grid, :vtk_multiblock, :paraview_collection, :paraview_collection_load) diff --git a/src/gridtypes/pvtk_grid.jl b/src/gridtypes/pvtk_grid.jl new file mode 100644 index 00000000..f8b622f5 --- /dev/null +++ b/src/gridtypes/pvtk_grid.jl @@ -0,0 +1,316 @@ + +struct ParallelDatasetFile <: VTKFile + pvtkargs::Dict + xdoc::XMLDocument + dataset::DatasetFile + path::AbstractString + function ParallelDatasetFile( + pvtkargs::Dict, + xdoc::XMLDocument, + dataset::DatasetFile, + path::AbstractString) + new(pvtkargs,xdoc,dataset,path) + end +end + +function _default_pvtkargs(pvtkargs) + mandatory = [:part,:nparts] + optional = [:ismain,:ghost_level] + allkeys = vcat(mandatory,optional) + + d = Dict{Symbol,Any}(pvtkargs) + for k in keys(d) + if !(k in allkeys) + throw(ArgumentError("$k is not a valid key in pvtkargs.")) + end + end + for k in mandatory + if !(haskey(d,k)) + throw(ArgumentError("$k is a mandatory key in pvtkargs.")) + end + end + if ! haskey(d,:ismain) + d[:ismain] = (d[:part] == 1) + end + if ! haskey(d,:ghost_level) + d[:ghost_level] = 0 + end + d +end + +function vtk_save(vtk::ParallelDatasetFile) + if isopen(vtk) + _pvtk_grid_body(vtk) + if vtk.pvtkargs[:ismain] + save_file(vtk.xdoc, vtk.path) + end + vtk_save(vtk.dataset) + close(vtk) + end + return [vtk.path] :: Vector{String} +end + +function new_pdata_array(xParent, + type::Type{<:VTKDataType}, + name::AbstractString, + Nc=nothing) + xDA = new_child(xParent, "PDataArray") + set_attribute(xDA, "type", datatype_str(type)) + set_attribute(xDA, "Name", name) + if Nc != nothing + set_attribute(xDA, "NumberOfComponents", Nc) + end +end + +function get_extension(filename::AbstractString) + path, ext = splitext(filename) + ext +end + +function get_path(filename::AbstractString) + path, ext = splitext(filename) + path +end + +function parallel_file_extension(g::AbstractVTKDataset) + ext=file_extension(g) + replace(ext,"."=>".p") +end + +function parallel_xml_name(g::AbstractVTKDataset) + "P"*xml_name(g) +end + +function parallel_xml_write_header(pvtx::XMLDocument,dtype::AbstractVTKDataset) + xroot = create_root(pvtx, "VTKFile") + set_attribute(xroot, "type", parallel_xml_name(dtype)) + set_attribute(xroot, "version", "1.0") + if IS_LITTLE_ENDIAN + set_attribute(xroot, "byte_order", "LittleEndian") + else + set_attribute(xroot, "byte_order", "BigEndian") + end + xroot +end + +function add_pieces!(grid_xml_node, + prefix::AbstractString, + extension::AbstractString, + num_pieces::Int) + for i=1:num_pieces + piece=new_child(grid_xml_node,"Piece") + set_attribute(piece,"Source",prefix*string(i)*extension) + end + grid_xml_node +end + +function add_ppoints!(grid_xml_node; + type::Type{<:VTKDataType}=Float64) + ppoints=new_child(grid_xml_node,"PPoints") + new_pdata_array(ppoints,type,"Points",3) +end + +function add_ppoint_data!(pdataset::ParallelDatasetFile, + name::AbstractString; + type::Type{<:VTKDataType}=Float64, + Nc::Integer=1) + xroot=root(pdataset.xdoc) + dtype = xml_name_to_VTK_Dataset(pdataset.dataset.grid_type) + grid_xml_node=find_element(xroot,parallel_xml_name(dtype)) + @assert grid_xml_node != nothing + grid_xml_node_ppoint=find_element(grid_xml_node,"PPointData") + if (grid_xml_node_ppoint==nothing) + grid_xml_node_ppoint=new_child(grid_xml_node,"PPointData") + end + new_pdata_array(grid_xml_node_ppoint,type,name,Nc) +end + +function add_pcell_data!(pdataset::ParallelDatasetFile, + name::AbstractString; + type::Type{<:VTKDataType}=Float64, + Nc=nothing) + xroot=root(pdataset.xdoc) + dtype = xml_name_to_VTK_Dataset(pdataset.dataset.grid_type) + grid_xml_node=find_element(xroot,parallel_xml_name(dtype)) + @assert grid_xml_node != nothing + grid_xml_node_pcell=find_element(grid_xml_node,"PCellData") + if (grid_xml_node_pcell==nothing) + grid_xml_node_pcell=new_child(grid_xml_node,"PCellData") + end + new_pdata_array(grid_xml_node_pcell,type,name,Nc) +end + +function Base.setindex!(pvtk::ParallelDatasetFile, + data, + name::AbstractString, + loc::AbstractFieldData) + pvtk.dataset[name,loc]=data +end + +function Base.setindex!(vtk::ParallelDatasetFile, data, name::AbstractString) + vtk.dataset[name]=data +end + +function get_dataset_extension(dataset::DatasetFile) + path, ext = splitext(dataset.path) + @assert ext != "" + ext +end + +function get_dataset_path(dataset::DatasetFile) + path, ext = splitext(dataset.path) + @assert ext != "" + path +end + +function xml_name_to_VTK_Dataset(xml_name::AbstractString) + if (xml_name=="ImageData") + VTKImageData() + elseif (xml_name=="RectilinearGrid") + VTKRectilinearGrid() + elseif (xml_name=="PolyData") + VTKPolyData() + elseif (xml_name=="StructuredGrid") + VTKStructuredGrid() + elseif (xml_name=="UnstructuredGrid") + VTKUnstructuredGrid() + else + @assert false + end +end + +function num_children(element) + length(collect(child_elements(element))) +end + +function get_child_attribute(element,child,attr) + children=collect(child_elements(element)) + attribute(children[child],attr) +end + +function get_dataset_xml_type(dataset::DatasetFile) + r=root(dataset.xdoc) + attribute(r,"type") +end + +function get_dataset_xml_grid_element(dataset::DatasetFile, + element::AbstractString) + r=root(dataset.xdoc) + uns=find_element(r,get_dataset_xml_type(dataset)) + piece=find_element(uns,"Piece") + find_element(piece,element) +end + + +const string_to_VTKDataType = Dict("Int8"=>Int8, + "UInt8"=>UInt8, + "Int16"=>Int16, + "UInt16"=>UInt16, + "Int32"=>Int32, + "UInt32"=>UInt32, + "Int64"=>Int64, + "UInt64"=>UInt64, + "Float32"=>Float32, + "Float64"=>Float64, + "String"=>String) + +""" + pvtk_grid(args...;pvtkargs,kwargs...) + +Return a handler representing a parallel vtk file, which can be +eventually written to file with `vtk_save`. + +Positional and keyword arguments in `args` and `kwargs` +are passed to `vtk_grid` verbatim (except file names that are augmented with the +corresponding part id). + +The extra keyword argument `pvtkargs` contains +extra options (as a `Dict{Symbol,Any}` or a `Vector{Pair{Symbol,Any}}`) +that only apply for parallel vtk file formats. + +Mandatory keys in `pvtkargs` are: + +- `:part` current (1-based) part id +- `:nparts` total number of parts + +Default key/value pairs in `pvtkargs` are +- `:ismain=>part==1` True if the current part id `part` is the main (the only one that will write the .pvtk file). +- `:ghost_level=>0` Ghost level. +""" +function pvtk_grid(filename::AbstractString,args...;pvtkargs,kwargs...) + _pvtkargs = _default_pvtkargs(pvtkargs) + path, ext = splitext(filename) + bname=basename(path) + path = mkpath(path) + part = _pvtkargs[:part] + nparts = _pvtkargs[:nparts] + p = lpad(part,ceil(Int,log10(nparts)),'0') + fn=joinpath(path,bname*"_$p"*ext) + vtk=vtk_grid(fn,args...;kwargs...) + _pvtk_grid_header(_pvtkargs,vtk,path) +end + +function _pvtk_grid_header( + pvtkargs::Dict, + dataset::DatasetFile, + filename::AbstractString) + + pvtx = XMLDocument() + dtype = xml_name_to_VTK_Dataset(dataset.grid_type) + xroot = parallel_xml_write_header(pvtx,dtype) + + # Generate parallel grid node + grid_xml_node=new_child(xroot,"P"*dataset.grid_type) + set_attribute(grid_xml_node, "GhostLevel", string(pvtkargs[:ghost_level])) + + prefix=joinpath(filename,basename(filename)) + extension=file_extension(dtype) + add_pieces!(grid_xml_node,prefix,extension,pvtkargs[:nparts]) + + pextension = parallel_file_extension(dtype) + pfilename = add_extension(filename,pextension) + pdataset=ParallelDatasetFile(pvtkargs,pvtx,dataset,pfilename) + + # Generate PPoints + points=get_dataset_xml_grid_element(dataset,"Points") + type=get_child_attribute(points,1,"type") + add_ppoints!(grid_xml_node,type=string_to_VTKDataType[type]) + + pdataset + end + +function _pvtk_grid_body(pdataset::ParallelDatasetFile) + dataset=pdataset.dataset + # Generate PPointData + pointdata=get_dataset_xml_grid_element(dataset,"PointData") + if pointdata != nothing + if (num_children(pointdata)>0) + for child=1:num_children(pointdata) + name=get_child_attribute(pointdata,child,"Name") + type=get_child_attribute(pointdata,child,"type") + Nc=get_child_attribute(pointdata,child,"NumberOfComponents") + add_ppoint_data!(pdataset, + name; + type=string_to_VTKDataType[type], + Nc=parse(Int64,Nc)) + end + end + end + + # Generate PCellData + celldata=get_dataset_xml_grid_element(dataset,"CellData") + if celldata!=nothing + if (num_children(celldata)>0) + for child=1:num_children(celldata) + name=get_child_attribute(celldata,child,"Name") + type=get_child_attribute(celldata,child,"type") + Nc=get_child_attribute(celldata,child,"NumberOfComponents") + add_pcell_data!(pdataset, + name; + type=string_to_VTKDataType[type], + Nc=(Nc==nothing ? nothing : parse(Int64,Nc))) + end + end + end + pdataset + end diff --git a/test/checksums.sha1 b/test/checksums.sha1 index fca5cbdd..0790650b 100644 --- a/test/checksums.sha1 +++ b/test/checksums.sha1 @@ -46,3 +46,4 @@ b20bf14b2531aaf323b891e090d0615d0e943ef6 collection_03.vtr 3cf1f5421c36377f38af90b2dc09083726128159 arraydata_2D.vti e0180e5cf18bb7c6f00df0bb9ded769f90b2574e arraydata_3D.vti 9e21f61ebf8641fdcc325836d4bd707e52932441 arrays.vti +bc0ce0fc82137c78a808f4c2a89bc0ca65e348c5 simulation.pvtu diff --git a/test/pvtk_grid.jl b/test/pvtk_grid.jl new file mode 100644 index 00000000..b0109180 --- /dev/null +++ b/test/pvtk_grid.jl @@ -0,0 +1,19 @@ +using WriteVTK + +function main() + # Suppose that the mesh is made of 5 points: + cells = [MeshCell(VTKCellTypes.VTK_TRIANGLE, [1, 4, 2]), + MeshCell(VTKCellTypes.VTK_QUAD, [2, 4, 3, 5])] + x=rand(5) + y=rand(5) + + @time pvtk = pvtk_grid( + "simulation", x, y, cells; + pvtkargs=[:part=>1,:nparts=>1]) # 2D + pvtk["Pressure"] = x + pvtk["Processor"] = rand(2) + @time outfiles = vtk_save(pvtk) + println("Saved: ", join(outfiles, " ")) + outfiles +end +main() diff --git a/test/runtests.jl b/test/runtests.jl index 1b2f8d57..9f044afa 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -15,6 +15,7 @@ const tests = [ "bezier.jl", "pvdCollection.jl", "array.jl", + "pvtk_grid.jl" ] # Only toggle to generate new checksums, if new tests are added. @@ -59,4 +60,3 @@ for test in tests outfiles = evalfile(test)::Vector{String} println() end -