Allow setting component names in PVTK files (#154)

Author: jipolanco

src/gridtypes/pvtk_grid.jl

@@ -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
@@ -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.
@@ -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
 

src/write_data.jl

@@ -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.
@@ -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)
@@ -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.
@@ -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.
@@ -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
 
 """

test/checksums.sha1

@@ -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

test/pvtk_grid.jl

@@ -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