From 56faa1d05bd2f8882c49997aef7f2361dfe91298 Mon Sep 17 00:00:00 2001
From: Elias Carvalho <eliascarvdev@gmail.com>
Date: Tue, 24 Sep 2024 13:35:54 -0300
Subject: [PATCH] Add 'lenunit' option to VTK loading

---
 src/extra/img.jl     |  4 ++--
 src/extra/vtkread.jl | 20 ++++++++++++--------
 test/io/vtk.jl       | 17 +++++++++++++++++
 3 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/src/extra/img.jl b/src/extra/img.jl
index 1c7ad98..49d926f 100644
--- a/src/extra/img.jl
+++ b/src/extra/img.jl
@@ -8,9 +8,9 @@ function imgread(fname; lenunit)
   values = (; color=vec(data))
   # translation followed by rotation is faster
   transform = Translate(-dims[1], 0) → Rotate(-π / 2)
-  # construct the grid
+  # construct grid
   u = lenunit
-  origin = Point(ntuple(i -> 0.0u, length(dims)))
+  origin = ntuple(i -> 0.0u, length(dims))
   spacing = ntuple(i -> 1.0u, length(dims))
   domain = CartesianGrid(dims, origin, spacing) |> transform
   georef(values, domain)
diff --git a/src/extra/vtkread.jl b/src/extra/vtkread.jl
index f201b7c..f67daa2 100644
--- a/src/extra/vtkread.jl
+++ b/src/extra/vtkread.jl
@@ -44,7 +44,7 @@ function vturead(fname; lenunit)
   vtk = ReadVTK.VTKFile(fname)
 
   # construct mesh
-  points = _points(vtk)
+  points = _points(vtk, lenunit)
   connec = _vtuconnec(vtk)
   mesh = SimpleMesh(points, connec)
 
@@ -59,7 +59,7 @@ function vtpread(fname; lenunit)
   vtk = ReadVTK.VTKFile(fname)
 
   # construct mesh
-  points = _points(vtk)
+  points = _points(vtk, lenunit)
   connec = _vtpconnec(vtk)
   mesh = SimpleMesh(points, connec)
 
@@ -74,9 +74,11 @@ function vtrread(fname; lenunit)
   vtk = ReadVTK.VTKFile(fname)
 
   # construct grid
+  u = lenunit
   coords = ReadVTK.get_coordinates(vtk)
   inds = map(!allequal, coords) |> collect
-  grid = RectilinearGrid(coords[inds]...)
+  xyz = map(x -> x * u, coords[inds])
+  grid = RectilinearGrid(xyz...)
 
   # extract data
   vtable, etable = _datatables(vtk)
@@ -89,10 +91,11 @@ function vtsread(fname; lenunit)
   vtk = ReadVTK.VTKFile(fname)
 
   # construct grid
+  u = lenunit
   coords = ReadVTK.get_coordinates(vtk)
   inds = map(!allequal, coords) |> collect
   dims = findall(!, inds) |> Tuple
-  XYZ = map(A -> dropdims(A; dims), coords[inds])
+  XYZ = map(A -> dropdims(A; dims) * u, coords[inds])
   grid = StructuredGrid(XYZ...)
 
   # extract data
@@ -106,6 +109,7 @@ function vtiread(fname; lenunit)
   vtk = ReadVTK.VTKFile(fname)
 
   # construct grid
+  u = lenunit
   ext = ReadVTK.get_whole_extent(vtk)
   # the get_origin and get_spacing functions drop the z dimension if it is empty, 
   # but the get_whole_extent function does not
@@ -115,8 +119,8 @@ function vtiread(fname; lenunit)
     (ext[2] - ext[1], ext[4] - ext[3], ext[6] - ext[5])
   end
   inds = findall(!iszero, dims)
-  origin = ReadVTK.get_origin(vtk) |> Tuple
-  spacing = ReadVTK.get_spacing(vtk) |> Tuple
+  origin = Tuple(ReadVTK.get_origin(vtk)) .* u
+  spacing = Tuple(ReadVTK.get_spacing(vtk)) .* u
   grid = CartesianGrid(dims[inds], origin[inds], spacing[inds])
 
   # extract data
@@ -130,10 +134,10 @@ end
 # UTILS
 #-------
 
-function _points(vtk)
+function _points(vtk, u)
   coords = ReadVTK.get_points(vtk)
   inds = map(!allequal, eachrow(coords))
-  [Point(Tuple(c)) for c in eachcol(coords[inds, :])]
+  [Point(Tuple(c) .* u) for c in eachcol(coords[inds, :])]
 end
 
 function _vtuconnec(vtk)
diff --git a/test/io/vtk.jl b/test/io/vtk.jl
index 6a9121c..b8b9c21 100644
--- a/test/io/vtk.jl
+++ b/test/io/vtk.jl
@@ -53,6 +53,23 @@
     vtable = values(gtb, 0)
     @test size(eltype(vtable.myVector)) == (2,)
     @test eltype(eltype(vtable.myVector)) <: Float32
+
+    # custom lenunit
+    file = ReadVTK.get_example_file("celldata_appended_binary_compressed.vtu", output_directory=savedir)
+    gtb = GeoIO.load(file, lenunit=cm)
+    @test unit(Meshes.lentype(crs(gtb.geometry))) == cm
+    file = joinpath(datadir, "spiral.vtp")
+    gtb = GeoIO.load(file, lenunit=cm)
+    @test unit(Meshes.lentype(crs(gtb.geometry))) == cm
+    file = joinpath(datadir, "rectilinear.vtr")
+    gtb = GeoIO.load(file, lenunit=cm)
+    @test unit(Meshes.lentype(crs(gtb.geometry))) == cm
+    file = joinpath(datadir, "structured.vts")
+    gtb = GeoIO.load(file, lenunit=cm)
+    @test unit(Meshes.lentype(crs(gtb.geometry))) == cm
+    file = joinpath(datadir, "imagedata.vti")
+    gtb = GeoIO.load(file, lenunit=cm)
+    @test unit(Meshes.lentype(crs(gtb.geometry))) == cm
   end
 
   @testset "save" begin