Merge branch 'main' into ss/optimize-masking

CliMA · Jan 13, 2025 · 4b83486 · 4b83486
2 parents 6decaf7 + 31e62c6
commit 4b83486
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 26 deletions.
diff --git a/src/Grids/Grids.jl b/src/Grids/Grids.jl
@@ -118,7 +118,7 @@ struct ZDirection <: AbstractDirection end
 struct NegativeZDirection <: AbstractDirection end
 
 include("abstract_grid.jl")
-include("vertical_coordinate.jl")
+include("vertical_coordinates.jl")
 include("grid_utils.jl")
 include("nodes_and_spacings.jl")
 include("zeros_and_ones.jl")

diff --git a/src/Grids/vertical_coordinate.jl → src/Grids/vertical_coordinates.jl b/src/Grids/vertical_coordinate.jl → src/Grids/vertical_coordinates.jl
@@ -3,48 +3,54 @@
 ####
 
 # This file implements everything related to vertical coordinates in Oceananigans.
-# Vertical coordinates are independent of the underlying grid type as we support grids that are 
-# "unstructured" or "curvilinear" only in the horizontal direction. 
-# For this reason the vertical coodinate is _special_, and it can be implemented once for all grid types.
+# Vertical coordinates are independent of the underlying grid type since only grids that are 
+# "unstructured" or "curvilinear" in the horizontal directions are supported in Oceananigans. 
+# Thus the vertical coordinate is _special_, and it can be implemented once for all grid types.
 
 abstract type AbstractVerticalCoordinate end
 
-# Represents a static one-dimensional vertical coordinate.
-#
-# # Fields
-# - `cᶜ::C`: Cell-centered coordinate.
-# - `cᶠ::C`: Face-centered coordinate.
-# - `Δᶜ::D`: Cell-centered grid spacing.
-# - `Δᶠ::D`: Face-centered grid spacing.
-struct StaticVerticalCoordinate{C, D} <: AbstractVerticalCoordinate
+"""
+    struct StaticVerticalCoordinate{C, D, E, F} <: AbstractVerticalCoordinate
+
+Represent a static one-dimensional vertical coordinate.
+
+Fields
+======
+
+- `cᶜ::C`: Cell-centered coordinate.
+- `cᶠ::D`: Face-centered coordinate.
+- `Δᶜ::E`: Cell-centered grid spacing.
+- `Δᶠ::F`: Face-centered grid spacing.
+"""
+struct StaticVerticalCoordinate{C, D, E, F} <: AbstractVerticalCoordinate
     cᵃᵃᶠ :: C
-    cᵃᵃᶜ :: C
-    Δᵃᵃᶠ :: D
-    Δᵃᵃᶜ :: D
+    cᵃᵃᶜ :: D
+    Δᵃᵃᶠ :: E
+    Δᵃᵃᶜ :: F
 end
 
 ####
-#### Some usefull aliases
+#### Some useful aliases
 ####
 
-const RegularVerticalCoordinate = StaticVerticalCoordinate{<:Any, <:Number}
+const RegularVerticalCoordinate = StaticVerticalCoordinate{<:Any, <:Any, <:Number, <:Number}
 const RegularVerticalGrid = AbstractUnderlyingGrid{<:Any, <:Any, <:Any, <:Any, <:RegularVerticalCoordinate}
 
 ####
 #### Adapt and on_architecture
 ####
 
-Adapt.adapt_structure(to, coord::StaticVerticalCoordinate) = 
-   StaticVerticalCoordinate(Adapt.adapt(to, coord.cᵃᵃᶠ),
-                            Adapt.adapt(to, coord.cᵃᵃᶜ),
-                            Adapt.adapt(to, coord.Δᵃᵃᶠ),
-                            Adapt.adapt(to, coord.Δᵃᵃᶜ))
+Adapt.adapt_structure(to, coord::StaticVerticalCoordinate) =
+    StaticVerticalCoordinate(Adapt.adapt(to, coord.cᵃᵃᶠ),
+                             Adapt.adapt(to, coord.cᵃᵃᶜ),
+                             Adapt.adapt(to, coord.Δᵃᵃᶠ),
+                             Adapt.adapt(to, coord.Δᵃᵃᶜ))
 
 on_architecture(arch, coord::StaticVerticalCoordinate) = 
-   StaticVerticalCoordinate(on_architecture(arch, coord.cᵃᵃᶠ),
-                            on_architecture(arch, coord.cᵃᵃᶜ),
-                            on_architecture(arch, coord.Δᵃᵃᶠ),
-                            on_architecture(arch, coord.Δᵃᵃᶜ))
+    StaticVerticalCoordinate(on_architecture(arch, coord.cᵃᵃᶠ),
+                             on_architecture(arch, coord.cᵃᵃᶜ),
+                             on_architecture(arch, coord.Δᵃᵃᶠ),
+                             on_architecture(arch, coord.Δᵃᵃᶜ))
 
 #####
 ##### Nodes and spacings (common to every grid)...