GeoMakie extension and interactive 3D grid visualisation on the sphere

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- GeoMakie weak dependecy, globe function for 3D data visualisation [#600](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/600)
+
 ## v0.12.0
 
 - OctaminimalGaussianArray/Grid to start with 4 points around the poles [#595](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/595)

Project.toml

@@ -29,11 +29,13 @@ UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228"
 
 [weakdeps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+GeoMakie = "db073c08-6b98-4ee5-b6a4-5efafb3259c6"
 JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
 
 [extensions]
 SpeedyWeatherCUDAExt = "CUDA"
+SpeedyWeatherGeoMakieExt = "GeoMakie"
 SpeedyWeatherJLArraysExt = "JLArrays"
 SpeedyWeatherMakieExt = "Makie"
 
@@ -50,6 +52,7 @@ FFTW = "1"
 FastGaussQuadrature = "0.4, 0.5, 1"
 GPUArrays = "10"
 GenericFFT = "0.1"
+GeoMakie = "0.7.5"
 JLArrays = "0.1.4"
 JLD2 = "0.4, 0.5"
 KernelAbstractions = "0.9"

docs/Project.toml

@@ -6,8 +6,8 @@ NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
 UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228"
 
 [compat]
-CairoMakie = "0.11"
+CairoMakie = "0.11, 0.12"
 Documenter = "0.26, 0.27, 1"
-GeoMakie = "0.6"
+GeoMakie = "0.7.5"
 NCDatasets = "0.12, 0.13, 0.14"
-UnicodePlots = "3.3"
+UnicodePlots = "3"

ext/SpeedyWeatherGeoMakieExt/SpeedyWeatherGeoMakieExt.jl

@@ -0,0 +1,59 @@
+module SpeedyWeatherGeoMakieExt
+
+using SpeedyWeather
+using GeoMakie
+
+include("faces.jl")
+
+SpeedyWeather.globe(Grid::Type{<:AbstractGridArray}, nlat_half::Integer) = SpeedyWeather.globe(SpectralGrid(; Grid, nlat_half))
+SpeedyWeather.globe(SG::SpectralGrid) = globe(Geometry(SG))
+
+function SpeedyWeather.globe(geometry::Geometry{NF, Grid}) where {NF, Grid}
+
+    faces, facesr = _faces(geometry)
+
+    transf = GeoMakie.Geodesy.ECEFfromLLA(GeoMakie.Geodesy.WGS84())
+
+    fig = Figure(size=(800, 800));
+    ax = LScene(fig[1,1], show_axis=false);
+
+    # background image
+    bg = meshimage!(ax, -180..180, -90..90, rotr90(GeoMakie.earth()); npoints = 100, z_level = -10_000);
+    bg.transformation.transform_func[] = transf
+
+    # cell centers
+    color = :black
+    s = scatter!(ax, geometry.londs, geometry.latds, markersize=5; color)
+    s.transformation.transform_func[] = transf
+
+    # cell faces
+    for i in 1:geometry.nlon_max÷4
+        for offset in [0, 90, 180, 270]
+            lp1 = lines!(ax, facesr[i, :, 1] .+ offset, facesr[i, :, 2]; color)
+            lp2 = lines!(ax, offset .- faces[i, :, 1], faces[i, :, 2]; color)
+
+            lp1.transformation.transform_func[] = transf
+            lp2.transformation.transform_func[] = transf
+        end
+    end
+
+    if RingGrids.nlat_odd(Grid) == false
+        # add equator
+        lpe = lines!(ax, 0:360, zeros(361); color)
+        lpe.transformation.transform_func[] = transf
+    end
+
+    # coastlines
+    lpc = lines!(GeoMakie.coastlines(50); color, linewidth=1)
+    lpc.transformation.transform_func[] = transf
+
+    # Makie stuff
+    cc = cameracontrols(ax.scene)
+    cc.settings.mouse_translationspeed[] = 0.0
+    cc.settings.zoom_shift_lookat[] = false
+    Makie.update_cam!(ax.scene, cc)
+
+    return fig
+end
+
+end # module

ext/SpeedyWeatherGeoMakieExt/faces.jl

@@ -0,0 +1,65 @@
+function _faces(geometry::Geometry{NF, Grid}) where {NF, Grid<:OctaminimalGaussianArray}
+
+    faces = zeros(geometry.nlon_max÷4+1, geometry.nlat+2, 2)
+    fill!(faces, NaN)
+    facesr = deepcopy(faces)
+
+    for (j, latd) in enumerate(geometry.latd)
+        nlon = RingGrids.get_nlon_per_ring(Grid, geometry.nlat_half, j)
+        lond = RingGrids.get_lon_per_ring(Grid, geometry.nlat_half, j) * (360/2π)
+        nlon ÷= 4
+        dϕ = 90/nlon
+
+        for i in 1:nlon
+            faces[i,  j+1,  1] = lond[i] + dϕ/2
+            faces[i,  j+1,  2] = latd
+            facesr[i, j+1,  :] = faces[i, j+1, :]
+        end
+
+        facesr[nlon+1, 1, :] .= NaN      # remove to avoid overlapping grid lines at edges
+
+        faces[end-1, j+1, 1] = 90
+        faces[end-1, j+1, 2] = latd
+    end
+
+    # continue to the poles
+    faces[:, 1, 1] = faces[:, 2, 1]
+    faces[:, 1, 2] .= 90
+
+    faces[:, end, 1] = faces[:, end-1, 1]
+    faces[:, end, 2] .= -90
+
+    return faces, facesr
+end
+
+function _faces(geometry::Geometry{NF, Grid}) where {NF, Grid<:Union{OctahedralGaussianArray, OctahedralClenshawArray}}
+
+    faces = zeros(geometry.nlon_max÷4+1, geometry.nlat+2, 2)
+    fill!(faces, NaN)
+    facesr = deepcopy(faces)
+
+    for (j, latd) in enumerate(geometry.latd)
+        nlon = RingGrids.get_nlon_per_ring(Grid, geometry.nlat_half, j)
+        lond = RingGrids.get_lon_per_ring(Grid, geometry.nlat_half, j) * (360/2π)
+        nlon ÷= 4
+        dϕ = 90/nlon
+
+        for i in 1:nlon+1
+            faces[i,  j+1,  1] = lond[i] + dϕ/2
+            faces[i,  j+1,  2] = latd
+            facesr[i, j+1,  :] = faces[i, j+1, :]
+        end
+
+        facesr[nlon+1, 1, :] .= NaN      # remove to avoid overlapping grid lines at edges
+        facesr[nlon+1, j+1, :] .= NaN    # remove to avoid overlapping grid lines at edges
+    end
+
+    # continue to the poles
+    faces[:, 1, 1] = faces[:, 2, 1]
+    faces[:, 1, 2] .= 90
+
+    faces[:, end, 1] = faces[:, end-1, 1]
+    faces[:, end, 2] .= -90
+
+    return faces, facesr
+end

src/SpeedyWeather.jl

@@ -76,6 +76,10 @@ include("SpeedyTransforms/SpeedyTransforms.jl")
 using .SpeedyTransforms
 import .SpeedyTransforms: prettymemory
 
+# to be defined in GeoMakie extension
+export globe
+function globe end
+
 # Utility for GPU / KernelAbstractions
 include("gpu.jl")                               
 

src/dynamics/geometry.jl

@@ -7,14 +7,11 @@ Construct Geometry struct containing parameters and arrays describing an iso-lat
 and the vertical levels. Pass on `SpectralGrid` to calculate the following fields
 $(TYPEDFIELDS)
 """
-@kwdef struct Geometry{NF<:AbstractFloat} <: AbstractGeometry
+@kwdef struct Geometry{NF<:AbstractFloat, Grid<:AbstractGrid} <: AbstractGeometry
 
     "SpectralGrid that defines spectral and grid resolution"
     spectral_grid::SpectralGrid
 
-    "grid of the dynamical core"
-    Grid::Type{<:AbstractGrid} = spectral_grid.Grid
-
     "resolution parameter nlat_half of Grid, # of latitudes on one hemisphere (incl Equator)"
     nlat_half::Int = spectral_grid.nlat_half      
 
@@ -32,7 +29,7 @@ $(TYPEDFIELDS)
     "number of vertical levels"
     nlayers::Int = spectral_grid.nlayers
 
-    "total number of grid points"
+    "total number of horizontal grid points"
     npoints::Int = spectral_grid.npoints
 
     "Planet's radius [m]"
@@ -99,11 +96,11 @@ end
 """
 $(TYPEDSIGNATURES)
 Generator function for `Geometry` struct based on `spectral_grid`."""
-function Geometry(spectral_grid::SpectralGrid)
-    error_message = "nlayers=$(spectral_grid.nlayers) does not match length nlayers="*
-        "$(spectral_grid.vertical_coordinates.nlayers) in spectral_grid.vertical_coordinates."
-    @assert spectral_grid.nlayers == spectral_grid.vertical_coordinates.nlayers error_message
-    return Geometry{spectral_grid.NF}(; spectral_grid)
+function Geometry(SG::SpectralGrid)
+    error_message = "nlayers=$(SG.nlayers) does not match length nlayers="*
+        "$(SG.vertical_coordinates.nlayers) in spectral_grid.vertical_coordinates."
+    @assert SG.nlayers == SG.vertical_coordinates.nlayers error_message
+    return Geometry{SG.NF, SG.Grid}(; spectral_grid=SG)
 end
 
 function Base.show(io::IO, G::Geometry)

src/dynamics/prognostic_variables.jl

@@ -290,8 +290,8 @@ function set!(var::LowerTriangularArray, grids::AbstractGridArray, geometry::Uni
 end
 
 # set LTA <- func 
-function set!(var::LowerTriangularArray, f::Function, geometry::Geometry{NF}, S::Union{SpectralTransform, Nothing}=nothing; add::Bool) where NF
-    grid = ndims(var) == 1 ? zeros(geometry.Grid{NF}, geometry.nlat_half) : zeros(geometry.Grid{NF}, geometry.nlat_half, geometry.nlayers)
+function set!(var::LowerTriangularArray, f::Function, geometry::Geometry{NF, Grid}, S::Union{SpectralTransform, Nothing}=nothing; add::Bool) where {NF, Grid}
+    grid = ndims(var) == 1 ? zeros(Grid{NF}, geometry.nlat_half) : zeros(Grid{NF}, geometry.nlat_half, geometry.nlayers)
     set!(grid, f, geometry, S; add=false)
     set!(var, grid, geometry, S; add)
 end

src/physics/tendencies.jl

@@ -15,14 +15,12 @@ function parameterization_tendencies!(
     # TODO move into shortwave radiation code?
     cos_zenith!(diagn, time, model)
 
-    G = model.geometry
-    rings = eachring(G.Grid, G.nlat_half)
-
-    for ij in eachgridpoint(diagn)       # loop over all horizontal grid points
+    rings = eachring(diagn.grid.vor_grid)       # indices on every latitude ring
+    for ij in eachgridpoint(diagn)              # loop over all horizontal grid points
 
         thread_id = Threads.threadid()          # not two threads should use the same ColumnVariables
         column = diagn.columns[thread_id]
-        jring = whichring(ij, rings)             # ring index gridpoint ij is on
+        jring = whichring(ij, rings)            # ring index gridpoint ij is on
 
         # extract current column for contiguous memory access
         reset_column!(column)                   # set accumulators back to zero for next grid point

test/spectral_gradients.jl

@@ -94,7 +94,6 @@ end
                         HEALPixGrid)
 
             SG = SpectralGrid(; NF, Grid, nlayers=1)
-            G = Geometry(SG)
 
             A = Grid(randn(NF, SG.npoints))
             B = copy(A)