Update Oceananigans to use non-uniform conformal cubed sphere grid from CubedSphere.jl (#4724)

* Generate non-uniform conformal mapping

* Add verbose flag to control console output

* Update cubed sphere test script

* Relocate primary functions to CubedSphere.jl

* Relocate spherical coordinate conversion functions

* Restore spacing

* Update function names

* Use cubed sphere test script from main branch

* Update cubed sphere test script

* Fix bug

* Make necessary updates

* Swap latitude and longitude argument order

* ci: force CubedSphere.jl to PR #41 branch

* Remove accidentally added debug code

* Fix bug

* Fix bug

* Refine .github/workflows/ci.yml

* ci: update Buildkite YAMLs to force CS branch

* ci: revert all YAML files to main branch

* Update Project.toml

* Fix bugs

* Fix kwarg

* Remove R_Earth

* Fix bug

* Make checkpointer & output filenames more specific

* Use correct angle units in coordinate rotation

* Disable lat bounds check in rotate_coordinates

* Bump CubedSphere version to 0.3.4

* Use radian-based trig functions for Reactant/XLA

* Add comments

* Bump version from 0.100.6 to 0.100.7

---------

Co-authored-by: Navid C. Constantinou <[email protected]>
Co-authored-by: Simone Silvestri <[email protected]>
Co-authored-by: Gregory L. Wagner <[email protected]>

Author: 4 people

Project.toml

@@ -1,7 +1,7 @@
 name = "Oceananigans"
 uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 authors = ["Climate Modeling Alliance and contributors"]
-version = "0.100.6"
+version = "0.100.7"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -66,7 +66,7 @@ Adapt = "4.1.1"
 CUDA = "=5.8.5, 5.9.1"
 ConstructionBase = "1"
 Crayons = "4"
-CubedSphere = "0.2, 0.3 - 0.3.2"
+CubedSphere = "0.3.4"
 Dates = "1.9"
 Distances = "0.10"
 DocStringExtensions = "0.8, 0.9"

src/Grids/grid_utils.jl

@@ -329,81 +329,6 @@ coordinate_summary(topo, Δ::Union{AbstractVector, AbstractMatrix}, name) =
 @inline column_depthᶜᶠᵃ(i, j, k, grid, η) = static_column_depthᶜᶠᵃ(i, j, grid)
 @inline column_depthᶠᶠᵃ(i, j, k, grid, η) = static_column_depthᶠᶠᵃ(i, j, grid)
 
-#####
-##### Spherical geometry
-#####
-
-"""
-    spherical_area_triangle(a::Number, b::Number, c::Number)
-
-Return the area of a spherical triangle on the unit sphere with sides `a`, `b`, and `c`.
-
-The area of a spherical triangle on the unit sphere is ``E = A + B + C - π``, where ``A``, ``B``, and ``C``
-are the triangle's inner angles.
-
-It has been known since the time of Euler and Lagrange that
-``\\tan(E/2) = P / (1 + \\cos a + \\cos b + \\cos c)``, where
-``P = (1 - \\cos²a - \\cos²b - \\cos²c + 2 \\cos a \\cos b \\cos c)^{1/2}``.
-
-References
-==========
-
-* Euler, L. (1778) De mensura angulorum solidorum, Opera omnia, 26, 204-233 (Orig. in Acta adac. sc. Petrop. 1778)
-* Lagrange,  J.-L. (1798) Solutions de quilquies problèmes relatifs au triangles sphéruques, Oeuvres, 7, 331-359.
-"""
-function spherical_area_triangle(a::Number, b::Number, c::Number)
-    cosa = cos(a)
-    cosb = cos(b)
-    cosc = cos(c)
-
-    tan½E = sqrt(1 - cosa^2 - cosb^2 - cosc^2 + 2cosa * cosb * cosc)
-    tan½E /= 1 + cosa + cosb + cosc
-
-    return 2atan(tan½E)
-end
-
-"""
-    spherical_area_triangle(a::AbstractVector, b::AbstractVector, c::AbstractVector)
-
-Return the area of a spherical triangle on the unit sphere with vertices given by the 3-vectors
-`a`, `b`, and `c` whose origin is the the center of the sphere. The formula was first given by
-Eriksson (1990).
-
-If we denote with ``A``, ``B``, and ``C`` the inner angles of the spherical triangle and with
-``a``, ``b``, and ``c`` the side of the triangle then, it has been known since Euler and Lagrange
-that ``\\tan(E/2) = P / (1 + \\cos a + \\cos b + \\cos c)``, where ``E = A + B + C - π`` is the
-triangle's excess and ``P = (1 - \\cos²a - \\cos²b - \\cos²c + 2 \\cos a \\cos b \\cos c)^{1/2}``.
-On the unit sphere, ``E`` is precisely the area of the spherical triangle. Erikkson (1990) showed
-that ``P`` above is the same as the volume defined by the vectors `a`, `b`, and `c`, that is
-``P = |𝐚 \\cdot (𝐛 \\times 𝐜)|``.
-
-References
-==========
-
-* Eriksson, F. (1990) On the measure of solid angles, Mathematics Magazine, 63 (3), 184-187, doi:10.1080/0025570X.1990.11977515
-"""
-function spherical_area_triangle(a₁::AbstractVector, a₂::AbstractVector, a₃::AbstractVector)
-    (sum(a₁.^2) ≈ 1 && sum(a₂.^2) ≈ 1 && sum(a₃.^2) ≈ 1) || error("a₁, a₂, a₃ must be unit vectors")
-
-    tan½E = abs(dot(a₁, cross(a₂, a₃)))
-    tan½E /= 1 + dot(a₁, a₂) + dot(a₂, a₃) + dot(a₁, a₃)
-
-    return 2atan(tan½E)
-end
-
-"""
-    spherical_area_quadrilateral(a₁, a₂, a₃, a₄)
-
-Return the area of a spherical quadrilateral on the unit sphere whose points are given by 3-vectors,
-`a`, `b`, `c`, and `d`. The area of the quadrilateral is given as the sum of the ares of the two
-non-overlapping triangles. To avoid having to pick the triangles appropriately ensuring they are not
-overlapping, we compute the area of the quadrilateral as the half the sum of the areas of all four potential
-triangles formed by `a₁`, `a₂`, `a₃`, and `a₄`.
-"""
-spherical_area_quadrilateral(a::AbstractVector, b::AbstractVector, c::AbstractVector, d::AbstractVector) =
-    1/2 * (spherical_area_triangle(a, b, c) + spherical_area_triangle(a, b, d) +
-           spherical_area_triangle(a, c, d) + spherical_area_triangle(b, c, d))
-
 """
     add_halos(data, loc, topo, sz, halo_sz; warnings=true)
 

src/Grids/orthogonal_spherical_shell_grid.jl

@@ -263,15 +263,6 @@ function fill_metric_halo_regions!(grid)
     return nothing
 end
 
-function lat_lon_to_cartesian(lat, lon, radius)
-    abs(lat) > 90 && error("lat must be within -90 ≤ lat ≤ 90")
-    return [lat_lon_to_x(lat, lon, radius), lat_lon_to_y(lat, lon, radius), lat_lon_to_z(lat, lon, radius)]
-end
-
-lat_lon_to_x(lat, lon, radius) = radius * cosd(lon) * cosd(lat)
-lat_lon_to_y(lat, lon, radius) = radius * sind(lon) * cosd(lat)
-lat_lon_to_z(lat, lon, radius) = radius * sind(lat)
-
 function on_architecture(arch::AbstractSerialArchitecture, grid::OrthogonalSphericalShellGrid)
     coordinates = (:λᶜᶜᵃ,
                    :λᶠᶜᵃ,

src/MultiRegion/cubed_sphere_grid.jl

@@ -5,6 +5,7 @@ using Oceananigans.Grids: halo_size,
                           topology
 
 using CubedSphere
+using CubedSphere.SphericalGeometry
 using Oceananigans.OrthogonalSphericalShellGrids: ConformalCubedSpherePanelGrid
 using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, has_active_cells_map, has_active_z_columns
 
@@ -31,7 +32,7 @@ const ConformalCubedSphereGridOfSomeKind{FT, TX, TY, TZ, CZ} =
                              z_topology = Bounded,
                              radius = Oceananigans.defaults.planet_radius,
                              non_uniform_conformal_mapping = false,
-                             spacing_type = "geometric",
+                             spacing = GeometricSpacing(),
                              provided_conformal_mapping = nothing,
                              partition = CubedSpherePartition(; R = 1))
 
@@ -194,7 +195,7 @@ function ConformalCubedSphereGrid(arch::AbstractArchitecture=CPU(),
                                   z_topology = Bounded,
                                   radius = Oceananigans.defaults.planet_radius,
                                   non_uniform_conformal_mapping = false,
-                                  spacing_type = "geometric",
+                                  spacing = GeometricSpacing(),
                                   provided_conformal_mapping = nothing,
                                   partition = CubedSpherePartition(; R = 1))
 
@@ -242,7 +243,7 @@ function ConformalCubedSphereGrid(arch::AbstractArchitecture=CPU(),
                                         halo = region_halo,
                                         rotation = region_rotation,
                                         non_uniform_conformal_mapping,
-                                        spacing_type,
+                                        spacing,
                                         provided_conformal_mapping)
 
     # Propagate the vertical coordinate type in the `MultiRegionGrid`.

src/OrthogonalSphericalShellGrids/OrthogonalSphericalShellGrids.jl

@@ -6,6 +6,8 @@ export TripolarGrid, RotatedLatitudeLongitudeGrid, ConformalCubedSpherePanelGrid
 import Oceananigans
 import Oceananigans.Architectures: on_architecture
 
+using CubedSphere.SphericalGeometry
+
 using Oceananigans
 using Oceananigans.Grids
 using Oceananigans.Architectures: device, on_architecture, AbstractArchitecture, CPU, GPU
@@ -15,8 +17,7 @@ using Oceananigans.Utils
 using Oceananigans.BoundaryConditions: Zipper
 using Oceananigans.Fields: index_binary_search, convert_to_0_360
 using Oceananigans.Grids: RightConnected
-using Oceananigans.Grids: halo_size, spherical_area_quadrilateral,
-                          lat_lon_to_cartesian, generate_coordinate, topology
+using Oceananigans.Grids: halo_size, generate_coordinate, topology
 using Oceananigans.Grids: total_length, add_halos, fill_metric_halo_regions!
 
 using Oceananigans.Operators