From 609512b3787778cfb2673ec42bcb37c30b919547 Mon Sep 17 00:00:00 2001
From: "Alberto F. Martin" <alberto.martin@monash.edu>
Date: Mon, 26 Jun 2023 13:30:27 +1000
Subject: [PATCH] Revisited visualization module changes so that one does not
 actually need to pass the array with the part/rank identifiers

---
 src/Visualization.jl | 24 +++++++++++++++++-------
 1 file changed, 17 insertions(+), 7 deletions(-)

diff --git a/src/Visualization.jl b/src/Visualization.jl
index 7d090a33..592bf400 100644
--- a/src/Visualization.jl
+++ b/src/Visualization.jl
@@ -143,13 +143,6 @@ end
 
 # Vtk related
 
-function Visualization.writevtk(parts,args...;kwargs...)
-  map(visualization_data(args...;kwargs...)) do visdata
-    write_vtk_file(
-    parts,visdata.grid,visdata.filebase,celldata=visdata.celldata,nodaldata=visdata.nodaldata)
-  end
-end
-
 function Visualization.write_vtk_file(
   parts::AbstractArray,
   grid::AbstractArray{<:Grid}, filebase; celldata, nodaldata)
@@ -171,6 +164,23 @@ function Visualization.create_vtk_file(
   end
 end
 
+function Visualization.writevtk(arg::DistributedGridapType,args...;kwargs...)
+  parts=get_parts(arg)
+  map(visualization_data(arg,args...;kwargs...)) do visdata
+    write_vtk_file(
+    parts,visdata.grid,visdata.filebase,celldata=visdata.celldata,nodaldata=visdata.nodaldata)
+  end
+end
+
+function Visualization.createvtk(arg::DistributedGridapType,args...;kwargs...)
+  v = visualization_data(arg,args...;kwargs...)
+  parts=get_parts(arg)
+  @notimplementedif length(v) != 1
+  visdata = first(v)
+  Visualization.create_vtk_file(
+    parts,visdata.grid,visdata.filebase,celldata=visdata.celldata,nodaldata=visdata.nodaldata)
+end
+
 struct DistributedPvd{T<:AbstractArray}
   pvds::T
 end