Add support for custom CRS in StructuredGrid (#1051)

* Add support for custom CRS in 'StructuredGrid'

* More adjustments

* Fix tests

* Add more tests

* Move helpers to units.jl

Author: eliascarv

src/coarsening/regular.jl

@@ -21,7 +21,7 @@ end
 RegularCoarsening(factors::Vararg{Int,N}) where {N} = RegularCoarsening(factors)
 
 function coarsen(grid::CartesianGrid, method::RegularCoarsening)
-  factors = fitdims(method.factors, embeddim(grid))
+  factors = fitdims(method.factors, paramdim(grid))
   CartesianGrid(minimum(grid), maximum(grid), dims=size(grid) .÷ factors)
 end
 
@@ -34,13 +34,13 @@ function coarsen(grid::RectilinearGrid, method::RegularCoarsening)
   RectilinearGrid{manifold(grid),crs(grid)}(xyzₜ)
 end
 
-function coarsen(grid::StructuredGrid{Datum}, method::RegularCoarsening) where {Datum}
-  factors = fitdims(method.factors, embeddim(grid))
+function coarsen(grid::StructuredGrid, 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...], embeddim(grid))
-  StructuredGrid{Datum}(XYZₜ)
+  XYZₜ = ntuple(i -> XYZₛ[i][rngs...], paramdim(grid))
+  StructuredGrid{manifold(grid),crs(grid)}(XYZₜ)
 end
 
 coarsen(grid::TransformedGrid, method::RegularCoarsening) =

src/domains.jl

@@ -152,6 +152,6 @@ Base.convert(::Type{GeometrySet}, d::Domain) = GeometrySet(collect(d))
 
 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{StructuredGrid}, g::Grid) = StructuredGrid{manifold(g),crs(g)}(XYZ(g))
 
 Base.convert(::Type{RectilinearGrid}, g::RegularGrid) = RectilinearGrid{manifold(g),crs(g)}(xyz(g))

src/domains/meshes/rectilineargrid.jl

@@ -42,7 +42,7 @@ function RectilinearGrid{M,C}(xyz::NTuple{N,AbstractVector}, topology::GridTopol
     """))
   end
 
-  xyz′ = ntuple(i -> numconvert.(T, _withunit.(xyz[i], us[i])), nc)
+  xyz′ = ntuple(i -> numconvert.(T, withunit.(xyz[i], us[i])), nc)
   RectilinearGrid{M,C,N,typeof(xyz′)}(xyz′, topology)
 end
 
@@ -54,14 +54,10 @@ end
 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
+  L = promote_type(ntuple(i -> aslentype(eltype(xyz[i])), N)...)
+  M = 𝔼{N}
+  C = Cartesian{NoDatum,N,L}
+  RectilinearGrid{M,C}(xyz)
 end
 
 RectilinearGrid(xyz::AbstractVector...) = RectilinearGrid(xyz)
@@ -94,10 +90,3 @@ 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

@@ -67,7 +67,7 @@ function RegularGrid(
     """))
   end
 
-  sp = ntuple(i -> numconvert(T, _withunit(spacing[i], us[i])), nc)
+  sp = ntuple(i -> numconvert(T, withunit(spacing[i], us[i])), nc)
 
   RegularGrid{M,C,typeof(sp),N}(origin, sp, offset, topology)
 end
@@ -135,10 +135,3 @@ function Base.show(io::IO, ::MIME"text/plain", g::RegularGrid)
   println(io, "├─ maximum: ", maximum(g))
   print(io, "└─ spacing: ", spacing(g))
 end
-
-# -----------------
-# HELPER FUNCTIONS
-# -----------------
-
-_withunit(x::Number, u) = x * u
-_withunit(x::Quantity, u) = uconvert(u, x)

src/domains/meshes/structuredgrid.jl

@@ -4,10 +4,10 @@
 
 """
     StructuredGrid(X, Y, Z, ...)
-    StructuredGrid{Datum}(X, Y, Z, ...)
+    StructuredGrid{M,C}(X, Y, Z, ...)
 
 A structured 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,41 +20,78 @@ julia> Y = repeat([0.0, 0.1, 0.3, 0.7, 0.9, 1.0]', 6, 1)
 julia> StructuredGrid(X, Y)
 ```
 """
-struct StructuredGrid{Datum,Dim,ℒ<:Len,A<:AbstractArray{ℒ}} <: Grid{𝔼{Dim},Cartesian{Datum,Dim,ℒ},Dim}
-  XYZ::NTuple{Dim,A}
-  topology::GridTopology{Dim}
-
-  function StructuredGrid{Datum}(XYZ::NTuple{Dim,<:AbstractArray{<:Len}}, topology::GridTopology{Dim}) where {Datum,Dim}
-    coords = float.(XYZ)
-    A = eltype(coords)
-    new{Datum,Dim,eltype(A),A}(coords, topology)
+struct StructuredGrid{M<:Manifold,C<:CRS,N,X<:NTuple{N,AbstractArray}} <: Grid{M,C,N}
+  XYZ::X
+  topology::GridTopology{N}
+  StructuredGrid{M,C,N,X}(XYZ, topology) where {M<:Manifold,C<:CRS,N,X<:NTuple{N,AbstractArray}} = new(XYZ, topology)
+end
+
+function StructuredGrid{M,C}(XYZ::NTuple{N,AbstractArray}, 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 structured 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)
+  StructuredGrid{M,C,N,typeof(XYZ′)}(XYZ′, topology)
 end
 
-StructuredGrid{Datum}(XYZ::NTuple{Dim,<:AbstractArray}, topology::GridTopology{Dim}) where {Datum,Dim} =
-  StructuredGrid{Datum}(addunit.(XYZ, u"m"), topology)
+function StructuredGrid{M,C}(XYZ::NTuple{N,AbstractArray}) where {M<:Manifold,C<:CRS,N}
+  if !allequal(size(X) for X in XYZ)
+    throw(ArgumentError("all coordinate arrays must be the same size"))
+  end
+
+  nd = ndims(first(XYZ))
 
-function StructuredGrid{Datum}(XYZ::Tuple) where {Datum}
-  coords = promote(XYZ...)
-  topology = GridTopology(size(first(coords)) .- 1)
-  StructuredGrid{Datum}(coords, topology)
+  if nd ≠ N
+    throw(ArgumentError("""
+    A $N-dimensional structured grid requires coordinate arrays with $N dimensions.
+    The provided coordinate arrays have $nd dimensions.
+    """))
+  end
+
+  topology = GridTopology(size(first(XYZ)) .- 1)
+  StructuredGrid{M,C}(XYZ, topology)
 end
 
-StructuredGrid{Datum}(XYZ...) where {Datum} = StructuredGrid{Datum}(XYZ)
+StructuredGrid{M,C}(XYZ::AbstractArray...) where {M<:Manifold,C<:CRS} = StructuredGrid{M,C}(XYZ)
 
-StructuredGrid(args...) = StructuredGrid{NoDatum}(args...)
+function StructuredGrid(XYZ::NTuple{N,AbstractArray}) where {N}
+  L = promote_type(ntuple(i -> aslentype(eltype(XYZ[i])), N)...)
+  M = 𝔼{N}
+  C = Cartesian{NoDatum,N,L}
+  StructuredGrid{M,C}(XYZ)
+end
 
-vertex(g::StructuredGrid{Datum}, ijk::Dims) where {Datum} =
-  Point(Cartesian{Datum}(ntuple(d -> g.XYZ[d][ijk...], embeddim(g))))
+StructuredGrid(XYZ::AbstractArray...) = StructuredGrid(XYZ)
+
+function vertex(g::StructuredGrid, ijk::Dims)
+  ctor = CoordRefSystems.constructor(crs(g))
+  Point(ctor(ntuple(d -> g.XYZ[d][ijk...], paramdim(g))...))
+end
 
 XYZ(g::StructuredGrid) = g.XYZ
 
-function Base.getindex(g::StructuredGrid{Datum}, I::CartesianIndices) where {Datum}
-  @boundscheck _checkbounds(g, I)
-  dims = size(I)
-  cinds = first(I):CartesianIndex(Tuple(last(I)) .+ 1)
-  XYZ = ntuple(i -> g.XYZ[i][cinds], embeddim(g))
-  StructuredGrid{Datum}(XYZ, GridTopology(dims))
+@generated function Base.getindex(g::StructuredGrid{M,C,N}, I::CartesianIndices) where {M,C,N}
+  exprs = ntuple(i -> :(g.XYZ[$i][cinds]), N)
+
+  quote
+    @boundscheck _checkbounds(g, I)
+    dims = size(I)
+    cinds = first(I):CartesianIndex(Tuple(last(I)) .+ 1)
+    XYZ = ($(exprs...),)
+    StructuredGrid{M,C}(XYZ, GridTopology(dims))
+  end
 end
 
 function Base.summary(io::IO, g::StructuredGrid)

src/refinement/regular.jl

@@ -21,7 +21,7 @@ end
 RegularRefinement(factors::Vararg{Int,N}) where {N} = RegularRefinement(factors)
 
 function refine(grid::CartesianGrid, method::RegularRefinement)
-  factors = fitdims(method.factors, embeddim(grid))
+  factors = fitdims(method.factors, paramdim(grid))
   CartesianGrid(minimum(grid), maximum(grid), dims=size(grid) .* factors)
 end
 
@@ -32,10 +32,10 @@ function refine(grid::RectilinearGrid, method::RegularRefinement)
   RectilinearGrid{manifold(grid),crs(grid)}(xyzₜ)
 end
 
-function refine(grid::StructuredGrid{Datum}, method::RegularRefinement) where {Datum}
-  factors = fitdims(method.factors, embeddim(grid))
+function refine(grid::StructuredGrid, method::RegularRefinement)
+  factors = fitdims(method.factors, paramdim(grid))
   XYZ′ = _XYZ(grid, factors)
-  StructuredGrid{Datum}(XYZ′)
+  StructuredGrid{manifold(grid),crs(grid)}(XYZ′)
 end
 
 refine(grid::TransformedGrid, method::RegularRefinement) =
@@ -49,7 +49,7 @@ function _refinedims(x, f)
   x′
 end
 
-_XYZ(grid::StructuredGrid, factors::Dims) = _XYZ(grid, Val(embeddim(grid)), factors)
+_XYZ(grid::StructuredGrid, factors::Dims) = _XYZ(grid, Val(paramdim(grid)), factors)
 
 function _XYZ(grid::StructuredGrid, ::Val{2}, factors::Dims{2})
   T = numtype(lentype(grid))

src/transforms/lengthunit.jl

@@ -34,7 +34,7 @@ applycoord(t::LengthUnit, lens::NTuple{Dim,Len}) where {Dim} = uconvert.(t.unit,
 
 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)))
+applycoord(t::LengthUnit, g::StructuredGrid) = TransformedGrid(g, t)
 
 # -----------------
 # HELPER FUNCTIONS

src/transforms/scale.jl

@@ -79,5 +79,5 @@ end
 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))
+applycoord(t::Scale, g::StructuredGrid) =
+  StructuredGrid{manifold(g),crs(g)}(ntuple(i -> t.factors[i] * XYZ(g)[i], paramdim(g)))

src/transforms/shadow.jl

@@ -65,6 +65,8 @@ apply(t::Shadow, ct::CylindricalTrajectory) = apply(t, GeometrySet(collect(ct)))
 
 apply(t::Shadow, g::RectilinearGrid) = TransformedGrid(g, t), nothing
 
+apply(t::Shadow, g::StructuredGrid) = TransformedGrid(g, t), nothing
+
 # -----------------
 # HELPER FUNCTIONS
 # -----------------
@@ -99,12 +101,6 @@ function _shadow(g::CartesianGrid, dims)
   CartesianGrid(sz, or, sp, of)
 end
 
-function _shadow(g::StructuredGrid, dims)
-  ndims = length(size(g))
-  inds = ntuple(i -> ifelse(i ∈ dims, :, 1), ndims)
-  StructuredGrid{datum(crs(g))}(map(X -> X[inds...], XYZ(g)[dims]))
-end
-
 # apply shadow transform recursively
 @generated function _shadow(g::G, dims) where {G<:GeometryOrDomain}
   ctor = constructor(G)

src/transforms/translate.jl

@@ -42,7 +42,7 @@ apply(t::Translate, g::RectilinearGrid) =
 
 revert(t::Translate, g::RectilinearGrid, c) = first(apply(inverse(t), g))
 
-apply(t::Translate, g::StructuredGrid{Datum}) where {Datum} =
-  StructuredGrid{Datum}(ntuple(i -> XYZ(g)[i] .+ t.offsets[i], embeddim(g))), nothing
+apply(t::Translate, g::StructuredGrid) =
+  StructuredGrid{manifold(g),crs(g)}(ntuple(i -> XYZ(g)[i] .+ t.offsets[i], paramdim(g))), nothing
 
 revert(t::Translate, g::StructuredGrid, c) = first(apply(inverse(t), g))

src/units.jl

@@ -25,3 +25,22 @@ addunit(::AbstractArray{<:Quantity}, _) = throw(ArgumentError("invalid units, pl
 Converts the number type of quantity `x` to `T`.
 """
 numconvert(::Type{T}, x::Quantity{S,D,U}) where {T,S,D,U} = convert(Quantity{T,D,U}, x)
+
+"""
+    withunit(x, u)
+
+Adds the unit if the argument is not a quantity, 
+otherwise, converts the unit of `x` to `u`.
+"""
+withunit(x::Number, u) = x * u
+withunit(x::Quantity, u) = uconvert(u, x)
+
+"""
+    aslentype(T)
+
+If `T` has length unit, return it. If `T` is number, return it with meter unit.
+Otherwise, throw an error.
+"""
+aslentype(::Type{T}) where {T<:Len} = T
+aslentype(::Type{T}) where {T<:Number} = Met{T}
+aslentype(::Type{<:Quantity}) = throw(ArgumentError("invalid units, please use a valid length unit"))