VTK: Add support for saving and loading views (#35)

* VTK: Add support for saving and loading views

* Apply suggestions

src/GeoIO.jl

@@ -9,6 +9,7 @@ using Tables
 using GeoTables
 using StaticArrays
 using PrettyTables
+using Meshes: SubDomain
 
 # image formats
 import FileIO

src/extra/vtkread.jl

@@ -14,8 +14,8 @@ const GEOMTYPE = Dict(
   VTKCellTypes.VTK_PYRAMID => Pyramid
 )
 
-function vtkread(fname)
-  if endswith(fname, ".vtu")
+function vtkread(fname; mask=:MASK)
+  gtb = if endswith(fname, ".vtu")
     vturead(fname)
   elseif endswith(fname, ".vtp")
     vtpread(fname)
@@ -28,6 +28,15 @@ function vtkread(fname)
   else
     error("unsupported VTK file format")
   end
+
+  names = propertynames(gtb)
+  if mask ∈ names
+    inds = findall(==(1), gtb[:, mask])
+    other = setdiff(names, [mask, :geometry])
+    view(gtb[:, other], inds)
+  else
+    gtb
+  end
 end
 
 function vturead(fname)

src/extra/vtkwrite.jl

@@ -3,9 +3,7 @@
 # ------------------------------------------------------------------
 
 function vtkwrite(fname, geotable)
-  dom = domain(geotable)
-  etable = values(geotable)
-  vtable = values(geotable, 0)
+  dom, etable, vtable = _extractvals(geotable)
 
   if endswith(fname, ".vtu")
     vtuwrite(fname, dom, etable, vtable)
@@ -78,6 +76,34 @@ end
 # UTILS
 #-------
 
+_extractvals(gtb) = _extractvals(domain(gtb), values(gtb), values(gtb, 0))
+_extractvals(dom::Domain, etable, vtable) = dom, etable, vtable
+function _extractvals(subdom::SubDomain, etable, vtable)
+  dom = parent(subdom)
+  inds = parentindices(subdom)
+  nelems = nelements(dom)
+
+  cols = Tables.columns(etable)
+  names = Tables.columnnames(cols)
+  pairs = map(names) do name
+    x = Tables.getcolumn(cols, name)
+    y = fill(NaN, nelems)
+    y[inds] .= x
+    name => y
+  end
+
+  mask = :MASK
+  # make unique
+  while mask ∈ names
+    mask = Symbol(mask, :_)
+  end
+  maskcol = zeros(UInt8, nelems)
+  maskcol[inds] .= 1
+
+  newtable = (; pairs..., mask => maskcol)
+  dom, newtable, nothing
+end
+
 _vtktype(::Type{<:Segment}) = VTKCellTypes.VTK_LINE
 _vtktype(::Type{<:Triangle}) = VTKCellTypes.VTK_TRIANGLE
 _vtktype(::Type{<:Quadrangle}) = VTKCellTypes.VTK_QUAD

src/load.jl

@@ -30,7 +30,7 @@ function load(fname; layer=0, fix=true, kwargs...)
 
   # VTK formats
   if any(ext -> endswith(fname, ext), VTKEXTS)
-    return vtkread(fname)
+    return vtkread(fname; kwargs...)
   end
 
   # Common Data Model formats

test/runtests.jl

@@ -500,6 +500,62 @@ end
       @test vertices(table2.geometry) == vertices(table1.geometry)
       @test values(table2) == values(table1)
 
+      # save views
+      grid = CartesianGrid(10, 10)
+      mesh = convert(SimpleMesh, grid)
+      rgrid = convert(RectilinearGrid, grid)
+      sgrid = convert(StructuredGrid, grid)
+      etable = (; a=rand(100))
+
+      gtb = georef(etable, mesh)
+      file = joinpath(savedir, "unstructured_view.vtu")
+      GeoIO.save(file, view(gtb, 1:25))
+      vgtb = GeoIO.load(file)
+      @test vgtb.a == view(gtb.a, 1:25)
+      @test parent(vgtb.geometry) isa SimpleMesh
+      @test vgtb.geometry == view(gtb.geometry, 1:25)
+
+      gtb = georef(etable, mesh)
+      file = joinpath(savedir, "polydata_view.vtp")
+      GeoIO.save(file, view(gtb, 1:25))
+      vgtb = GeoIO.load(file)
+      @test vgtb.a == view(gtb.a, 1:25)
+      @test parent(vgtb.geometry) isa SimpleMesh
+      @test vgtb.geometry == view(gtb.geometry, 1:25)
+
+      gtb = georef(etable, rgrid)
+      file = joinpath(savedir, "rectilinear_view.vtr")
+      GeoIO.save(file, view(gtb, 1:25))
+      vgtb = GeoIO.load(file)
+      @test vgtb.a == view(gtb.a, 1:25)
+      @test parent(vgtb.geometry) isa RectilinearGrid
+      @test vgtb.geometry == view(gtb.geometry, 1:25)
+
+      gtb = georef(etable, sgrid)
+      file = joinpath(savedir, "structured_view.vts")
+      GeoIO.save(file, view(gtb, 1:25))
+      vgtb = GeoIO.load(file)
+      @test vgtb.a == view(gtb.a, 1:25)
+      @test parent(vgtb.geometry) isa StructuredGrid
+      @test vgtb.geometry == view(gtb.geometry, 1:25)
+
+      gtb = georef(etable, grid)
+      file = joinpath(savedir, "imagedata_view.vti")
+      GeoIO.save(file, view(gtb, 1:25))
+      vgtb = GeoIO.load(file)
+      @test vgtb.a == view(gtb.a, 1:25)
+      @test parent(vgtb.geometry) isa CartesianGrid
+      @test vgtb.geometry == view(gtb.geometry, 1:25)
+
+      # mask column with different name
+      gtb = georef((; MASK=rand(100)), grid)
+      file = joinpath(savedir, "imagedata_view.vti")
+      GeoIO.save(file, view(gtb, 1:25))
+      vgtb = GeoIO.load(file, mask=:MASK_)
+      @test vgtb.MASK == view(gtb.MASK, 1:25)
+      @test parent(vgtb.geometry) isa CartesianGrid
+      @test vgtb.geometry == view(gtb.geometry, 1:25)
+
       # throw: the vtr format does not support structured grids
       table = GeoIO.load(joinpath(datadir, "structured.vts"))
       @test_throws ErrorException GeoIO.save(joinpath(savedir, "structured.vtr"), table)