Initial attempt at online downscaling via RingGrids.jl

Project.toml

@@ -43,6 +43,6 @@ IntervalSets = "0.7"
 KernelAbstractions = "0.9"
 Oceananigans = "0.100, 0.101, 0.102, 0.103, 0.104, 0.105, 0.106"
 RingGrids = "0.1"
-SpeedyWeatherInternals = "0.1"
+SpeedyWeatherInternals = "0.2"
 Unitful = "1"
 julia = "1.10"

examples/Project.toml

@@ -19,8 +19,5 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Terrarium = "80418d68-07fa-499d-ae2b-c12e531f5cd8"
 
-[compat]
-SpeedyWeather = "0.18.1"
-
 [sources.Terrarium]
 path = ".."

src/grids/column_ring_grid.jl

@@ -100,17 +100,17 @@ get_field_grid(grid::ColumnRingGrid) = grid.grid
 Converts the given Oceananigans `Field` to a `RingGrids.Field` with a ring grid matching that of the given `ColumnRingGrid`.
 """
 RingGrids.Field(field::Field{LX, LY, LZ}, grid::ColumnRingGrid; fill_value = NaN) where {LX, LY, LZ} = RingGrids.Field(architecture(field), dropdims(interior(field), dims = 2), grid; fill_value)
+RingGrids.Field(arch::AbstractArchitecture, field::Field{LX, LY, LZ}, grid::ColumnRingGrid; fill_value = NaN) where {LX, LY, LZ} = RingGrids.Field(arch, dropdims(interior(field), dims = 2), grid; fill_value)
 RingGrids.Field(field::AbstractArray, grid::ColumnRingGrid; fill_value = NaN) = RingGrids.Field(architecture(grid), field, grid; fill_value)
 function RingGrids.Field(arch::AbstractArchitecture, field::AbstractArray, grid::ColumnRingGrid; fill_value = NaN)
     # need to be on CPU to do the copying
     grid = on_architecture(arch, grid)
     field = on_architecture(arch, field)
     # create new RingGrids field initialized with fill_value
-    ring_field = RingGrids.Field(grid.rings, size(field)[2:end]...)
+    ring_field = RingGrids.Field(grid.rings, size(field, 2))
     fill!(ring_field, fill_value)
     # need to access underlying data arrays directly to avoid scalar indexing
-    colons = (Colon() for d in size(field)[2:end])
-    ring_field.data[grid.mask.data, colons...] .= field
+    ring_field.data[grid.mask.data, :] .= field
     return ring_field
 end
 

src/state_variables.jl

@@ -74,9 +74,10 @@ function Oceananigans.TimeSteppers.update_state!(state::StateVariables, model::A
     update_inputs!(state, inputs)
     fill_halo_regions!(state)
     compute_auxiliary!(state, model)
-    return if compute_tendencies
+    if compute_tendencies
         compute_tendencies!(state, model)
     end
+    return nothing
 end
 
 """

test/benchmarks/gpu/soil_heat_hydrology_global.jl

@@ -83,23 +83,54 @@ exit()
 # Interactive plotting
 
 using DataFrames, CSV
-using GLMakie, GeoMakie
+using CairoMakie, GeoMakie
+using CairoMakie.GeometryBasics
+using RingGrids
 
-GLMakie.activate!(inline = true)
+# GLMakie.activate!(inline = true)
 
 cpu_data = DataFrame(CSV.File("outputs/benchmarks/soil_heat_hydrology_benchmark_cpu_nthreads=32.csv"))
 gpu_data = DataFrame(CSV.File("outputs/benchmarks/soil_heat_hydrology_benchmark_gpu_nthreads=32.csv"))
 
-let fig = Figure(size = (800, 400)),
-        ax = Axis(fig[1, 1], title = "Terrarium.jl GPU vs. CPU scaling", ylabel = "log simulated years per day (SYPD)", xlabel = "log number of grid cells (Nₕ)", xscale = log10, yscale = log10)
-    # data is in ms / sim hr x 1 d / (1000*24*3600 ms) x 24 sim hr / sim day ->
-    cpu_sypd = 1000 * 24 * 3600 ./ (24 * cpu_data.mid_time)
-    gpu_sypd = 1000 * 24 * 3600 ./ (24 * gpu_data.mid_time)
-    scatterlines!(ax, cpu_data.npoints, cpu_sypd, label = "CPU")
-    scatterlines!(ax, gpu_data.npoints, gpu_sypd, label = "GPU")
-    axislegend(ax)
-    Makie.save("plots/terrarium_cpu_vs_gpu_scaling.svg", fig)
-    fig
+ring_grid_5deg = FullGaussianGrid(14)
+ring_grid_3deg = FullGaussianGrid(24)
+ring_grid_1deg = FullGaussianGrid(72)
+ring_grid_halfdeg = FullGaussianGrid(72 * 2)
+ring_grid_qtrdeg = FullGaussianGrid(72 * 4)
+ring_grid_10km = FullGaussianGrid(72 * 10)
+ring_grid_1km = FullGaussianGrid(72 * 100)
+ref_grids = [ring_grid_5deg, ring_grid_1deg, ring_grid_halfdeg, ring_grid_qtrdeg, ring_grid_10km]
+ref_npoints = [get_npoints(rg) for rg in ref_grids]
+
+Makie.with_theme(fontsize = 18) do
+    let fig = Figure(size = (800, 400))
+        ax = Axis(
+            fig[1, 1],
+            title = "Terrarium.jl GPU vs. CPU scaling",
+            ylabel = "Simulated years per day (SYPD)",
+            xlabel = "Number of grid cells",
+            xscale = log10,
+            yscale = log10
+        )
+        # data is in ms / sim hr x 1 d / (1000*24*3600 ms) x 24 sim hr / sim day ->
+        cpu_sypd = 1000 * 24 * 3600 ./ (24 * cpu_data.mid_time)
+        gpu_sypd = 1000 * 24 * 3600 ./ (24 * gpu_data.mid_time)
+        scatterlines!(ax, cpu_data.npoints, cpu_sypd, label = "CPU")
+        scatterlines!(ax, gpu_data.npoints, gpu_sypd, label = "GPU")
+        # Draw vertical reference lines and place a short annotation next to each one.
+        # Choose human-friendly labels for the reference grids.
+        ref_labels = ["5°", "1°", "0.5°", "0.25°", "0.1°"]
+        ys_max = maximum(vcat(cpu_sypd, gpu_sypd))
+        for (x, lbl) in zip(ref_npoints, ref_labels)
+            vlines!(ax, [x], linestyle = :dash, color = :black)
+            # place label slightly to the right of the line at the top of the plotted range
+            text!(ax, GeometryBasics.Point2(x * 1.2, 200); text = lbl, align = (:left, :top), fontsize = 18)
+        end
+        axislegend(ax)
+        Makie.save("plots/terrarium_cpu_vs_gpu_scaling.svg", fig)
+        Makie.save("plots/terrarium_cpu_vs_gpu_scaling.pdf", fig)
+        fig
+    end
 end
 
 using SpeedyWeather, CUDA