Skip to content

Commit

Permalink
Allow setting component names in PVTK files (#154)
Browse files Browse the repository at this point in the history
  • Loading branch information
jipolanco authored Oct 14, 2024
1 parent 79ed09b commit f625b5d
Show file tree
Hide file tree
Showing 4 changed files with 56 additions and 30 deletions.
27 changes: 19 additions & 8 deletions src/gridtypes/pvtk_grid.jl
Original file line number Diff line number Diff line change
Expand Up @@ -137,19 +137,21 @@ end

# Add point and cell data as usual
function Base.setindex!(
pvtk::PVTKFile, data, name::AbstractString, loc::AbstractFieldData,
pvtk::PVTKFile, data, name::AbstractString, loc::AbstractFieldData;
component_names::Union{Nothing, ListOfStrings} = nothing
)
pvtk.vtk[name, loc] = data
setindex!(pvtk.vtk, data, name, loc; component_names = component_names)
end

# In the case of field data, also add it to the pvtk file.
# Otherwise field data (typically the time / "TimeValue") is not seen by ParaView.
function Base.setindex!(
pvtk::PVTKFile, data, name::AbstractString, loc::VTKFieldData,
pvtk::PVTKFile, data, name::AbstractString, loc::VTKFieldData;
component_names::Union{Nothing, ListOfStrings} = nothing,
)
pvtk.vtk[name, loc] = data
setindex!(pvtk.vtk, data, name, loc; component_names = component_names)
if pvtk.pvtkargs.ismain
add_field_data(pvtk, data, name, loc)
add_field_data(pvtk, data, name, loc; component_names = component_names)
end
data
end
Expand All @@ -164,10 +166,11 @@ end

# If `loc` was not passed, try to guess which kind of data was passed.
function Base.setindex!(
pvtk::PVTKFile, data, name::AbstractString,
pvtk::PVTKFile, data, name::AbstractString;
component_names::Union{Nothing, ListOfStrings} = nothing
)
loc = guess_data_location(data, pvtk.vtk) :: AbstractFieldData
pvtk[name, loc] = data
loc = guess_data_location(data, pvtk.vtk)::AbstractFieldData
setindex!(pvtk, data, name, loc; component_names = component_names)
end

# Write XML attribute.
Expand Down Expand Up @@ -332,10 +335,18 @@ function _update_pvtk!(pvtk::PVTKFile)
t = attribute(vtk_data_array, "type")
name = attribute(vtk_data_array, "Name")
Nc = attribute(vtk_data_array, "NumberOfComponents")
Nc_val = parse(Int, Nc)
pvtk_pdata_array = new_child(pvtk_ppoint_data, "PDataArray")
set_attribute(pvtk_pdata_array, "type", t)
set_attribute(pvtk_pdata_array, "Name", name)
set_attribute(pvtk_pdata_array, "NumberOfComponents", Nc)
for n in 1:Nc_val
attr = "ComponentName$(n - 1)"
component_name = attribute(vtk_data_array, attr)
if component_name !== nothing
set_attribute(pvtk_pdata_array, attr, component_name)
end
end
end
end

Expand Down
43 changes: 28 additions & 15 deletions src/write_data.jl
Original file line number Diff line number Diff line change
Expand Up @@ -152,7 +152,7 @@ xml_data_array_name(::Union{String,ListOfStrings}) = "Array"
data_to_xml(
vtk::DatasetFile, xParent::XMLElement, data,
name::AbstractString, Nc::Union{Int,AbstractFieldData} = 1;
component_names::Union{AbstractVector, Nothing} = nothing
component_names = nothing
)
Add numerical data to VTK XML file.
Expand All @@ -164,8 +164,8 @@ In the latter case, the number of components will be deduced from the data
dimensions and the type of field data.
"""
function data_to_xml(vtk, xParent, data, name,
Nc::Union{Int,AbstractFieldData}=1;
component_names::Union{AbstractVector,Nothing}=nothing)
Nc::Union{Int,AbstractFieldData} = 1;
component_names = nothing)
xDA = new_child(xParent, xml_data_array_name(data))
set_attribute(xDA, "type", datatype_str(data))
set_attribute(xDA, "Name", name)
Expand Down Expand Up @@ -322,14 +322,16 @@ function data_to_xml_ascii(vtk::VTKFile, xDA::XMLElement, data)
end

"""
add_field_data(vtk::DatasetFile, data,
name::AbstractString, loc::AbstractFieldData)
add_field_data(
vtk::DatasetFile, data, name::AbstractString, loc::AbstractFieldData;
component_names = nothing,
)
Add either point or cell data to VTK file.
"""
function add_field_data(vtk::VTKFile, data, name::AbstractString,
loc::AbstractFieldData;
component_names::Union{AbstractVector, Nothing}=nothing)
component_names = nothing)
xbase = find_base_xml_node_to_add_field(vtk, loc)

# Find or create "nodetype" (PointData, CellData or FieldData) node.
Expand Down Expand Up @@ -360,19 +362,23 @@ vtk_cell_data(args...; kwargs...) = add_field_data(args..., VTKCellData(); kwarg
vtk_field_data(args...; kwargs...) = add_field_data(args..., VTKFieldData(); kwargs...)

"""
setindex!(vtk::DatasetFile, data, name::AbstractString, [field_type])
setindex!(vtk::DatasetFile, data, name::AbstractString, [field_type]; [component_names])
Add a new dataset to VTK file.
The number of components of the dataset (e.g. for scalar or vector fields) is
determined automatically from the input data dimensions.
The optional argument `field_type` should be an instance of `VTKPointData`,
The optional `field_type` argument should be an instance of `VTKPointData`,
`VTKCellData` or `VTKFieldData`.
It determines whether the data should be associated to grid points, cells or
none.
If not given, this is guessed from the input data size and the grid dimensions.
The optional `component_names` keyword argument allows to override the default component
names when writing vector or tensor fields. It should be a tuple or a vector of strings (see
below for an example).
# Example
Add "velocity" dataset and time scalar to VTK file.
Expand All @@ -382,20 +388,27 @@ vel = rand(3, 12, 14, 42) # vector field
time = 42.0
vtk = vtk_grid(...)
vtk["velocity", VTKPointData()] = vel
vtk["velocity", VTKPointData(), component_names = ("Ux", "Uy", "Uz")] = vel
vtk["time", VTKFieldData()] = time
# This also works, and will generally give the same result:
vtk["velocity"] = vel
vtk["velocity", component_names = ("Ux", "Uy", "Uz")] = vel
vtk["time"] = time
```
"""
Base.setindex!(vtk::DatasetFile, data, name::AbstractString,
loc::AbstractFieldData; kwargs...) = add_field_data(vtk, data, name, loc; kwargs...)
function Base.setindex!(
vtk::DatasetFile, data, name::AbstractString, loc::AbstractFieldData;
component_names::Union{Nothing, ListOfStrings} = nothing,
)
add_field_data(vtk, data, name, loc; component_names = component_names)
end

function Base.setindex!(vtk::DatasetFile, data, name::AbstractString; kwargs...)
loc = guess_data_location(data, vtk) :: AbstractFieldData
setindex!(vtk, data, name, loc; kwargs...)
function Base.setindex!(
vtk::DatasetFile, data, name::AbstractString;
component_names::Union{Nothing, ListOfStrings} = nothing
)
loc = guess_data_location(data, vtk)::AbstractFieldData
setindex!(vtk, data, name, loc; component_names = component_names)
end

"""
Expand Down
14 changes: 7 additions & 7 deletions test/checksums.sha1
Original file line number Diff line number Diff line change
Expand Up @@ -68,13 +68,13 @@ a898a771d8e0962628118d73a216a3e4660a5295 pimage/pimage_2.vti
4badae7ec804773110e662e105531e5e3d4fd288 pimage/pimage_4.vti
99843bc0f87c436d9ed385b40707ee1e0327f709 pimage/pimage_5.vti
4e37ea3120610664425712254b88d88f78d4e45d pimage/pimage_6.vti
9c0d3e7ee95c6e7bfc2f0554972d31fb9d832172 prectilinear.pvtr
e7ff55d7c8bb2c1739a2dc73c0c1fdeff2f7d536 prectilinear/prectilinear_1.vtr
6b0b2f8e5d381ec4f9644905677be8f29f432af2 prectilinear/prectilinear_2.vtr
68f065b66a0ff9ba80bc18776912f248ebc83390 prectilinear/prectilinear_3.vtr
b8b456b99f45f45e2e0fcc0de56963185b3f24ec prectilinear/prectilinear_4.vtr
a6fd8a0367efc31e7616dd7046f048c424aa8ae0 prectilinear/prectilinear_5.vtr
b1ef0b3a9e3c25781c31e0d2ceb552c78dab23f8 prectilinear/prectilinear_6.vtr
c70c5786fe01d09417c7a1772441910e0eb7ceda prectilinear.pvtr
8f1d50d1da766a64a9d702752209f8c2c84b30b0 prectilinear/prectilinear_1.vtr
548d706be5ac47cb64727f8de23bc39fbb4bf1b9 prectilinear/prectilinear_2.vtr
ca43f918aa4807323f39029fc31dd4e80428f324 prectilinear/prectilinear_3.vtr
ee05657f2065b5134f4763ac17e22c2ed0d0af2a prectilinear/prectilinear_4.vtr
d507180d6ada6da6831f4375dc47db1d4320d4c7 prectilinear/prectilinear_5.vtr
54dea6c985f95580167d768527dbd791ae60af90 prectilinear/prectilinear_6.vtr
d9c0b2367bf7ba42534a134b3ef4dace753560a9 pvtk_structured_output/pstructured.pvts
92ff18a7c1c631800de3adc3b0bd4631e16be845 pvtk_structured_output/pstructured/pstructured_1.vts
9bb94bd6b2374156e60df995340e3b395f2ae6b9 pvtk_structured_output/pstructured/pstructured_2.vts
Expand Down
2 changes: 2 additions & 0 deletions test/pvtk_grid.jl
Original file line number Diff line number Diff line change
Expand Up @@ -91,13 +91,15 @@ function pvtk_rectilinear()
@spawn begin
xs = map(is -> [sqrt(i) for i is], extent) # local grid
point_data = map(sum, Iterators.product(xs...))
vecdata = (1, 2, 3) .* Ref(point_data)
processid = fill(n, length.(xs) .- 1) # cell data
filenames[n] = pvtk_grid(
"prectilinear", xs...;
part = n, extents = extents,
append = false, compress = false,
) do vtk
vtk["point_data"] = point_data
vtk["vector", component_names = ("Ux", "Uy", "Uz")] = vecdata
vtk["process_id"] = processid
vtk["TimeValue"] = 3.4
end
Expand Down

0 comments on commit f625b5d

Please sign in to comment.