Add support for custom CRS in RectilinearGrid (#1049)

* Add support for custom CRS in 'RectilinearGrid'

* Update tests

* Format

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Fix typo

* Rename variables

* Apply suggestions

* Fix LengthUnit

* Fix tests

* Fix more tests

* Remove unused code

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: eliascarv

src/coarsening/regular.jl

@@ -25,13 +25,13 @@ function coarsen(grid::CartesianGrid, method::RegularCoarsening)
   CartesianGrid(minimum(grid), maximum(grid), dims=size(grid) .÷ factors)
 end
 
-function coarsen(grid::RectilinearGrid{Datum}, method::RegularCoarsening) where {Datum}
-  factors = fitdims(method.factors, embeddim(grid))
+function coarsen(grid::RectilinearGrid, method::RegularCoarsening)
+  factors = fitdims(method.factors, paramdim(grid))
   dims = vsize(grid)
-  rngs = ntuple(i -> 1:factors[i]:dims[i], embeddim(grid))
+  rngs = ntuple(i -> 1:factors[i]:dims[i], paramdim(grid))
   xyzₛ = xyz(grid)
-  xyzₜ = ntuple(i -> xyzₛ[i][rngs[i]], embeddim(grid))
-  RectilinearGrid{Datum}(xyzₜ)
+  xyzₜ = ntuple(i -> xyzₛ[i][rngs[i]], paramdim(grid))
+  RectilinearGrid{manifold(grid),crs(grid)}(xyzₜ)
 end
 
 function coarsen(grid::StructuredGrid{Datum}, method::RegularCoarsening) where {Datum}

src/domains.jl

@@ -154,4 +154,4 @@ Base.convert(::Type{SimpleMesh}, m::Mesh) = SimpleMesh(vertices(m), topology(m))
 
 Base.convert(::Type{StructuredGrid}, g::Grid) = StructuredGrid{datum(crs(g))}(XYZ(g))
 
-Base.convert(::Type{RectilinearGrid}, g::CartesianGrid) = RectilinearGrid{datum(crs(g))}(xyz(g))
+Base.convert(::Type{RectilinearGrid}, g::RegularGrid) = RectilinearGrid{manifold(g),crs(g)}(xyz(g))

src/domains/meshes/rectilineargrid.jl

@@ -4,10 +4,10 @@
 
 """
     RectilinearGrid(x, y, z, ...)
-    RectilinearGrid{Datum}(x, y, z, ...)
+    RectilinearGrid{M,C}(x, y, z, ...)
 
 A rectilinear grid with vertices at sorted coordinates `x`, `y`, `z`, ...,
-and a given `Datum` (default to `NoDatum`).
+manifold `M` (default to `𝔼`) and CRS type `C` (default to `Cartesian`).
 
 ## Examples
 
@@ -20,49 +20,84 @@ julia> y = [0.0, 0.1, 0.3, 0.7, 0.9, 1.0]
 julia> RectilinearGrid(x, y)
 ```
 """
-struct RectilinearGrid{Datum,Dim,ℒ<:Len,V<:AbstractVector{ℒ}} <: Grid{𝔼{Dim},Cartesian{Datum,Dim,ℒ},Dim}
-  xyz::NTuple{Dim,V}
-  topology::GridTopology{Dim}
-
-  function RectilinearGrid{Datum}(
-    xyz::NTuple{Dim,<:AbstractVector{<:Len}},
-    topology::GridTopology{Dim}
-  ) where {Datum,Dim}
-    coords = float.(xyz)
-    V = eltype(coords)
-    new{Datum,Dim,eltype(V),V}(coords, topology)
-  end
+struct RectilinearGrid{M<:Manifold,C<:CRS,N,X<:NTuple{N,AbstractVector}} <: Grid{M,C,N}
+  xyz::X
+  topology::GridTopology{N}
+  RectilinearGrid{M,C,N,X}(xyz, topology) where {M<:Manifold,C<:CRS,N,X<:NTuple{N,AbstractVector}} = new(xyz, topology)
 end
 
-RectilinearGrid{Datum}(xyz::NTuple{Dim,<:AbstractVector}, topology::GridTopology{Dim}) where {Datum,Dim} =
-  RectilinearGrid{Datum}(addunit.(xyz, u"m"), topology)
+function RectilinearGrid{M,C}(xyz::NTuple{N,AbstractVector}, topology::GridTopology{N}) where {M<:Manifold,C<:CRS,N}
+  if M <: 🌐 && !(C <: LatLon)
+    throw(ArgumentError("rectilinear grid on `🌐` requires `LatLon` coordinates"))
+  end
+
+  T = CoordRefSystems.mactype(C)
+  nc = CoordRefSystems.ncoords(C)
+  us = CoordRefSystems.units(C)
+
+  if N ≠ nc
+    throw(ArgumentError("""
+    A $N-dimensional rectilinear grid requires a CRS with $N coordinates.
+    The provided CRS has $nc coordinates.
+    """))
+  end
+
+  xyz′ = ntuple(i -> numconvert.(T, _withunit.(xyz[i], us[i])), nc)
+  RectilinearGrid{M,C,N,typeof(xyz′)}(xyz′, topology)
+end
 
-function RectilinearGrid{Datum}(xyz::Tuple) where {Datum}
-  coords = promote(collect.(xyz)...)
-  topology = GridTopology(length.(coords) .- 1)
-  RectilinearGrid{Datum}(coords, topology)
+function RectilinearGrid{M,C}(xyz::NTuple{N,AbstractVector}) where {M<:Manifold,C<:CRS,N}
+  topology = GridTopology(length.(xyz) .- 1)
+  RectilinearGrid{M,C}(xyz, topology)
 end
 
-RectilinearGrid{Datum}(xyz...) where {Datum} = RectilinearGrid{Datum}(xyz)
+RectilinearGrid{M,C}(xyz::AbstractVector...) where {M<:Manifold,C<:CRS} = RectilinearGrid{M,C}(xyz)
+
+function RectilinearGrid(xyz::NTuple{N,AbstractVector}) where {N}
+  try
+    L = promote_type(ntuple(i -> _lentype(eltype(xyz[i])), N)...)
+    M = 𝔼{N}
+    C = Cartesian{NoDatum,N,L}
+    RectilinearGrid{M,C}(xyz)
+  catch
+    throw(ArgumentError("invalid units for cartesian coordinates"))
+  end
+end
 
-RectilinearGrid(args...) = RectilinearGrid{NoDatum}(args...)
+RectilinearGrid(xyz::AbstractVector...) = RectilinearGrid(xyz)
 
-vertex(g::RectilinearGrid{Datum}, ijk::Dims) where {Datum} = Point(Cartesian{Datum}(getindex.(g.xyz, ijk)))
+function vertex(g::RectilinearGrid, ijk::Dims)
+  ctor = CoordRefSystems.constructor(crs(g))
+  Point(ctor(getindex.(g.xyz, ijk)...))
+end
 
 xyz(g::RectilinearGrid) = g.xyz
 
 XYZ(g::RectilinearGrid) = XYZ(xyz(g))
 
-function Base.getindex(g::RectilinearGrid{Datum}, I::CartesianIndices) where {Datum}
-  @boundscheck _checkbounds(g, I)
-  dims = size(I)
-  start = Tuple(first(I))
-  stop = Tuple(last(I)) .+ 1
-  xyz = ntuple(i -> g.xyz[i][start[i]:stop[i]], embeddim(g))
-  RectilinearGrid{Datum}(xyz, GridTopology(dims))
+@generated function Base.getindex(g::RectilinearGrid{M,C,N}, I::CartesianIndices) where {M,C,N}
+  exprs = ntuple(N) do i
+    :(g.xyz[$i][start[$i]:stop[$i]])
+  end
+
+  quote
+    @boundscheck _checkbounds(g, I)
+    dims = size(I)
+    start = Tuple(first(I))
+    stop = Tuple(last(I)) .+ 1
+    xyz = ($(exprs...),)
+    RectilinearGrid{M,C}(xyz, GridTopology(dims))
+  end
 end
 
 function Base.summary(io::IO, g::RectilinearGrid)
   join(io, size(g), "×")
   print(io, " RectilinearGrid")
 end
+
+# -----------------
+# HELPER FUNCTIONS
+# -----------------
+
+_lentype(::Type{T}) where {T<:Len} = T
+_lentype(::Type{T}) where {T<:Number} = Met{T}

src/domains/meshes/regulargrid.jl

@@ -17,7 +17,7 @@ A regular grid with dimensions `dims`, with lower left corner of element
 
 ```
 RegularGrid((10, 20), Point(LatLon(30.0°, 60.0°)), (1.0, 1.0)) # add coordinate units to spacing
-RegularGrid((10, 20), Point(Polar(0.0mm, 0.0rad)), (10.0mm, 1.0rad)) # convert spacing units to coordinate units
+RegularGrid((10, 20), Point(Polar(0.0cm, 0.0rad)), (10.0mm, 1.0rad)) # convert spacing units to coordinate units
 RegularGrid((10, 20), Point(Marcator(0.0, 0.0)), (1.5, 1.5))
 RegularGrid((10, 20, 30), Point(Cylindrical(0.0, 0.0, 0.0)), (3.0, 2.0, 1.0))
 ```

src/refinement/regular.jl

@@ -25,11 +25,11 @@ function refine(grid::CartesianGrid, method::RegularRefinement)
   CartesianGrid(minimum(grid), maximum(grid), dims=size(grid) .* factors)
 end
 
-function refine(grid::RectilinearGrid{Datum}, method::RegularRefinement) where {Datum}
-  factors = fitdims(method.factors, embeddim(grid))
+function refine(grid::RectilinearGrid, method::RegularRefinement)
+  factors = fitdims(method.factors, paramdim(grid))
   xyzₛ = xyz(grid)
-  xyzₜ = ntuple(i -> _refinedims(xyzₛ[i], factors[i]), embeddim(grid))
-  RectilinearGrid{Datum}(xyzₜ)
+  xyzₜ = ntuple(i -> _refinedims(xyzₛ[i], factors[i]), paramdim(grid))
+  RectilinearGrid{manifold(grid),crs(grid)}(xyzₜ)
 end
 
 function refine(grid::StructuredGrid{Datum}, method::RegularRefinement) where {Datum}

src/transforms/lengthunit.jl

@@ -32,7 +32,7 @@ applycoord(t::LengthUnit, len::Len) = uconvert(t.unit, len)
 
 applycoord(t::LengthUnit, lens::NTuple{Dim,Len}) where {Dim} = uconvert.(t.unit, lens)
 
-applycoord(t::LengthUnit, g::RectilinearGrid) = RectilinearGrid{datum(crs(g))}(map(x -> uconvert.(t.unit, x), xyz(g)))
+applycoord(t::LengthUnit, g::RectilinearGrid) = TransformedGrid(g, t)
 
 applycoord(t::LengthUnit, g::StructuredGrid) = StructuredGrid{datum(crs(g))}(map(X -> uconvert.(t.unit, X), XYZ(g)))
 

src/transforms/scale.jl

@@ -76,8 +76,8 @@ function applycoord(t::Scale, g::CartesianGrid)
   CartesianGrid(dims, orig, spac, offs)
 end
 
-applycoord(t::Scale{Dim}, g::RectilinearGrid{Datum,Dim}) where {Datum,Dim} =
-  RectilinearGrid{Datum}(ntuple(i -> t.factors[i] * xyz(g)[i], Dim))
+applycoord(t::Scale, g::RectilinearGrid) =
+  RectilinearGrid{manifold(g),crs(g)}(ntuple(i -> t.factors[i] * xyz(g)[i], paramdim(g)))
 
 applycoord(t::Scale{Dim}, g::StructuredGrid{Datum,Dim}) where {Datum,Dim} =
   StructuredGrid{Datum}(ntuple(i -> t.factors[i] * XYZ(g)[i], Dim))

src/transforms/shadow.jl

@@ -63,6 +63,8 @@ apply(t::Shadow, tr::Torus) = TransformedGeometry(tr, t), nothing
 
 apply(t::Shadow, ct::CylindricalTrajectory) = apply(t, GeometrySet(collect(ct))), nothing
 
+apply(t::Shadow, g::RectilinearGrid) = TransformedGrid(g, t), nothing
+
 # -----------------
 # HELPER FUNCTIONS
 # -----------------
@@ -97,8 +99,6 @@ function _shadow(g::CartesianGrid, dims)
   CartesianGrid(sz, or, sp, of)
 end
 
-_shadow(g::RectilinearGrid, dims) = RectilinearGrid{datum(crs(g))}(xyz(g)[dims])
-
 function _shadow(g::StructuredGrid, dims)
   ndims = length(size(g))
   inds = ntuple(i -> ifelse(i ∈ dims, :, 1), ndims)

src/transforms/translate.jl

@@ -37,8 +37,8 @@ applycoord(::Translate, v::Vec) = v
 # SPECIAL CASES
 # --------------
 
-apply(t::Translate, g::RectilinearGrid{Datum}) where {Datum} =
-  RectilinearGrid{Datum}(ntuple(i -> xyz(g)[i] .+ t.offsets[i], embeddim(g))), nothing
+apply(t::Translate, g::RectilinearGrid) =
+  RectilinearGrid{manifold(g),crs(g)}(ntuple(i -> xyz(g)[i] .+ t.offsets[i], paramdim(g))), nothing
 
 revert(t::Translate, g::RectilinearGrid, c) = first(apply(inverse(t), g))
 

test/meshes.jl

@@ -431,10 +431,31 @@ end
   @test vertex(grid, 1) == cart(0, 0)
   @test vertex(grid, 121) == cart(10, 10)
 
-  # constructor with datum & datum propagation
-  grid = RectilinearGrid{WGS84Latest}(x, y)
-  @test datum(crs(grid)) === WGS84Latest
-  @test datum(crs(centroid(grid))) === WGS84Latest
+  # constructor with manifold and CRS
+  x = range(zero(T), stop=one(T), length=6)
+  y = T[0.0, 0.1, 0.3, 0.7, 0.9, 1.0]
+  C = typeof(Mercator(T(0), T(0)))
+  grid = RectilinearGrid{𝔼{2},C}(x, y)
+  @test manifold(grid) === 𝔼{2}
+  @test crs(grid) === C
+  @test crs(grid[1, 1]) === C
+  @test crs(centroid(grid)) === C
+  C = typeof(LatLon(T(0), T(0)))
+  grid = RectilinearGrid{🌐,C}(x, y)
+  @test manifold(grid) === 🌐
+  @test crs(grid) === C
+  @test crs(grid[1, 1]) === C
+  @test crs(centroid(grid)) === C
+
+  # units
+  x = range(zero(T), stop=one(T), length=6) * u"mm"
+  y = T[0.0, 0.1, 0.3, 0.7, 0.9, 1.0] * u"cm"
+  grid = RectilinearGrid(x, y)
+  @test unit(Meshes.lentype(grid)) == u"m"
+  # error: invalid units for cartesian coordinates
+  x = range(zero(T), stop=one(T), length=6) * u"m"
+  y = T[0.0, 0.1, 0.3, 0.7, 0.9, 1.0] * u"°"
+  @test_throws ArgumentError RectilinearGrid(x, y)
 
   # conversion
   cg = cartgrid(10, 10)
@@ -453,6 +474,20 @@ end
   @test topology(rg) == topology(cg)
   @test vertices(rg) == vertices(cg)
 
+  # type stability
+  x = range(zero(T), stop=one(T), length=6) * u"mm"
+  y = T[0.0, 0.1, 0.3, 0.7, 0.9, 1.0] * u"cm"
+  ρ = range(zero(T), stop=one(T), length=6)
+  ϕ = range(zero(T), stop=T(2π), length=6)
+  C = typeof(Polar(T(0), T(0)))
+  grid = RectilinearGrid{𝔼{2},C}(ρ, ϕ)
+  @inferred RectilinearGrid(x, y)
+  @inferred RectilinearGrid{𝔼{2},C}(ρ, ϕ)
+  @inferred vertex(grid, (1, 1))
+  @inferred grid[1, 1]
+  @inferred grid[1:2, 1:2]
+  @inferred Meshes.XYZ(grid)
+
   x = range(zero(T), stop=one(T), length=6)
   y = T[0.0, 0.1, 0.3, 0.7, 0.9, 1.0]
   grid = RectilinearGrid(x, y)