Add GPU tests, plus tests for bathymetry module (#147)

* add test bathymetry * some tests * test bathymetry * Update test/test_bathymetry.jl Co-authored-by: Gregory L. Wagner <[email protected]> * Update test/test_bathymetry.jl Co-authored-by: Gregory L. Wagner <[email protected]> * adding gpu tests * Update test/test_bathymetry.jl Co-authored-by: Gregory L. Wagner <[email protected]> * more understandable * correct keys * correct the pipeline * add a small change * fixing the tests * test fix * Update test/test_bathymetry.jl Co-authored-by: Gregory L. Wagner <[email protected]> * what is the problem here? * fixed bathymetry tests * comment * retest bathymetry * save fts on the CPU * fix tests * this shouldn't be failing... * using on_architecture * hopefully final fix * why would test simulations not find the cuda driver? --------- Co-authored-by: Gregory L. Wagner <[email protected]>
CliMA · Sep 17, 2024 · ddea939 · ddea939
1 parent b3b44b7
commit ddea939
Show file tree

Hide file tree

Showing 9 changed files with 210 additions and 63 deletions.
diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
@@ -34,40 +34,119 @@ steps:
 
   - wait
 
-  - label: "Run JRA55 tests"
-    key: "test_jra55"
+  - label: "CPU JRA55 tests"
+    key: "cpu_jra55_tests"
     env:
-      CUDA_VISIBLE_DEVICES: "-1"
       TEST_GROUP: "jra55"
     commands:
       - "julia --project -e 'using Pkg; Pkg.test()'"
+    agents:
+      slurm_mem: 32G
+      slurm_cpus_per_task: 8
+      slurm_ntasks: 1
 
-  - label: "Run ECCO tests"
-    key: "test_ecco"
+  - label: "CPU ECCO tests"
+    key: "cpu_ecco_tests"
     env:
-      CUDA_VISIBLE_DEVICES: "-1"
       TEST_GROUP: "ecco"
     commands:
       - "julia --project -e 'using Pkg; Pkg.test()'"
+    agents:
+      slurm_mem: 32G
+      slurm_cpus_per_task: 8
+      slurm_ntasks: 1
 
-  - label: "Run surface flux tests"
-    key: "test_fluxes"
+  - label: "CPU surface flux tests"
+    key: "cpu_fluxes_tests"
     env:
-      CUDA_VISIBLE_DEVICES: "-1"
       TEST_GROUP: "fluxes"
     commands:
       - "julia --project -e 'using Pkg; Pkg.test()'"
+    agents:
+      slurm_mem: 32G
+      slurm_cpus_per_task: 8
+      slurm_ntasks: 1
+
+  - label: "CPU bathymetry tests"
+    key: "cpu_bathymetry_tests"
+    env:
+      TEST_GROUP: "bathymetry"
+    commands:
+      - "julia --project -e 'using Pkg; Pkg.test()'"
+    agents:
+      slurm_mem: 32G
+      slurm_cpus_per_task: 8
+      slurm_ntasks: 1
+
+  - label: "GPU JRA55 tests"
+    key: "gpu_jra55_tests"
+    env:
+      TEST_GROUP: "jra55"
+      GPU_TEST: "true"
+    commands:
+      - "julia --project -e 'using Pkg; Pkg.test()'"
+    agents:
+      slurm_mem: 32G
+      slurm_gpus: 1
+      slurm_cpus_per_task: 8
+      slurm_ntasks: 1
+      slurm_gpus_per_task: 1
 
+  - label: "GPU ECCO tests"
+    key: "gpu_ecco_tests"
+    env:
+      TEST_GROUP: "ecco"
+      GPU_TEST: "true"
+    commands:
+      - "julia --project -e 'using Pkg; Pkg.test()'"
+    agents:
+      slurm_mem: 32G
+      slurm_gpus: 1
+      slurm_cpus_per_task: 8
+      slurm_ntasks: 1
+      slurm_gpus_per_task: 1
+
+  - label: "GPU surface flux tests"
+    key: "gpu_fluxes_tests"
+    env:
+      TEST_GROUP: "fluxes"
+      GPU_TEST: "true"
+    commands:
+      - "julia --project -e 'using Pkg; Pkg.test()'"
+    agents:
+      slurm_mem: 32G
+      slurm_gpus: 1
+      slurm_cpus_per_task: 8
+      slurm_ntasks: 1
+      slurm_gpus_per_task: 1
+
+  - label: "GPU bathymetry tests"
+    key: "gpu_bathymetry_tests"
+    env:
+      TEST_GROUP: "bathymetry"
+      GPU_TEST: "true"
+    commands:
+      - "julia --project -e 'using Pkg; Pkg.test()'"
+    agents:
+      slurm_mem: 32G
+      slurm_gpus: 1
+      slurm_cpus_per_task: 8
+      slurm_ntasks: 1
+      slurm_gpus_per_task: 1
+
   - label: "Run simulation tests"
     key: "test_simulations"
     env:
       TEST_GROUP: "simulations"
+      GPU_TEST: "true"
     commands:
       - "julia --project -e 'using Pkg; Pkg.test()'"
     agents:
-      slurm_mem: 120G
+      slurm_mem: 100G
       slurm_gpus: 1
       slurm_cpus_per_task: 8
+      slurm_ntasks: 1
+      slurm_gpus_per_task: 1
 
   - wait: ~
     continue_on_failure: true

diff --git a/Project.toml b/Project.toml
@@ -39,7 +39,7 @@ ImageMorphology = "0.4"
 JLD2 = "0.4"
 KernelAbstractions = "0.9"
 NCDatasets = "0.12, 0.13, 0.14"
-Oceananigans = "0.91.8"
+Oceananigans = "0.91.11"
 OrthogonalSphericalShellGrids = "0.1.2"
 Scratch = "1"
 SeawaterPolynomials = "0.3.4"

diff --git a/src/Bathymetry.jl b/src/Bathymetry.jl
@@ -210,9 +210,9 @@ function interpolate_bathymetry_in_passes(native_z, target_grid;
     if Nλt > Nλn || Nφt > Nφn
         target_z = Field{Center, Center, Nothing}(target_grid)
         interpolate!(target_z, native_z)
-        @info string("Skipping passes for interplating bathymetry of size $Nn", '\n',
-                     "to target grid of size $Nt. Interpolation passes may only", '\n',
-                     "be used to refine bathymetryand requires that the bathymetry", '\n',
+        @info string("Skipping passes for interplating bathymetry of size $Nn ", '\n',
+                     "to target grid of size $Nt. Interpolation passes may only ", '\n',
+                     "be used to refine bathymetryand requires that the bathymetry ", '\n',
                      "is larger than the target grid in both horizontal directions.")
         return target_z
     end
@@ -304,28 +304,29 @@ function remove_minor_basins!(Z, keep_major_basins)
     end
 
     mm_basins = [] # major basins indexes
-    for m = 1:keep_major_basins
+    m = 1
+
+    # We add basin indexes until we reach the specified number (m == keep_major_basins) or
+    # we run out of basins to keep -> isempty(total_elements) 
+    while (m <= keep_major_basins) && !isempty(total_elements) 
         next_maximum = findfirst(x -> x == maximum(total_elements), total_elements)
         push!(mm_basins, label_elements[next_maximum])
-        total_elements = filter(x -> x != total_elements[next_maximum], total_elements)
-        label_elements = filter(x -> x != label_elements[next_maximum], label_elements)
+        deleteat!(total_elements, next_maximum)
+        deleteat!(label_elements, next_maximum)
+        m += 1
     end
-        
+
     labels = map(Float64, labels)
 
     for ℓ = 1:maximum(labels)
         remove_basin = all(ℓ != m for m in mm_basins)
         if remove_basin
-            labels[labels .== ℓ] .= 1e10 # large number
+            labels[labels .== ℓ] .= NaN # Regions to remove
         end
     end
 
-    # Flatten land
-    labels[labels .<  1e10] .= 0 
-    labels[labels .== 1e10] .= NaN
-
-    Z .+= labels
-    Z[isnan.(Z)] .= 0
+    # Flatten minor basins, corresponding to regions where `labels == NaN`
+    Z[isnan.(labels)] .= 0
 
     return nothing
 end

diff --git a/src/DataWrangling/DataWrangling.jl b/src/DataWrangling/DataWrangling.jl
@@ -4,7 +4,7 @@ using Oceananigans
 using Downloads
 using Printf
 
-using Oceananigans.Architectures: architecture
+using Oceananigans.Architectures: architecture, on_architecture
 using Oceananigans.Grids: node
 using Oceananigans.BoundaryConditions: fill_halo_regions!
 using Oceananigans.Fields: interpolate
@@ -54,8 +54,10 @@ end
 
 function save_field_time_series!(fts; path, name, overwrite_existing=false)
     overwrite_existing && rm(path; force=true)
-    times = fts.times
-    grid = fts.grid
+
+    times = on_architecture(CPU(), fts.times)
+    grid  = on_architecture(CPU(), fts.grid)
+
     LX, LY, LZ = location(fts)
     ondisk_fts = FieldTimeSeries{LX, LY, LZ}(grid, times;
                                              backend = OnDisk(), path, name)

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -27,6 +27,10 @@ if test_group == :turbulent_fluxes || test_group == :all
     include("test_surface_fluxes.jl")
 end
 
+if test_group == :bathymetry || test_group == :all
+    include("test_bathymetry.jl")
+end
+
 if test_group == :simulations || test_group == :all
     include("test_simulations.jl")
 end

diff --git a/test/runtests_setup.jl b/test/runtests_setup.jl
@@ -8,9 +8,10 @@ using ClimaOcean.ECCO
 using ClimaOcean.JRA55
 using ClimaOcean.JRA55: JRA55_field_time_series
 
-using Oceananigans.Architectures: architecture
+using Oceananigans.Architectures: architecture, on_architecture
 using Oceananigans.OutputReaders: interpolate!
 
 using ClimaOcean
 
-test_architectures = [CPU()]
+gpu_test = parse(Bool, get(ENV, "GPU_TEST", "false"))
+test_architectures = gpu_test ? [GPU()] : [CPU()]
diff --git a/test/test_bathymetry.jl b/test/test_bathymetry.jl
@@ -0,0 +1,61 @@
+using Oceananigans
+using ClimaOcean
+using ClimaOcean.Bathymetry: remove_minor_basins!
+using Statistics
+
+@testset "Availability of Bathymetry" begin
+    @info "Testing Bathymetry utils..."
+    for arch in test_architectures
+        grid = LatitudeLongitudeGrid(arch;
+                                     size = (100, 100, 10), 
+                                     longitude = (0, 100), 
+                                     latitude = (0, 50),
+                                     z = (-6000, 0))
+
+        # Test that remove_minor_basins!(Z, Inf) does nothing
+        control_bottom_height = regrid_bathymetry(grid)        
+        bottom_height = deepcopy(control_bottom_height)
+        @test_throws ArgumentError remove_minor_basins!(bottom_height, Inf)
+
+        # A fictitiously large number which should presumably keep all the basins
+        remove_minor_basins!(bottom_height, 10000000) 
+        @test parent(bottom_height) == parent(control_bottom_height)
+
+        # Test that remove_minor_basins!(Z, 2) remove the correct number of Basins
+        bottom_height = Field{Center, Center, Nothing}(grid)
+        control_bottom_height = Field{Center, Center, Nothing}(grid)
+
+        # A two basins bathymetry
+        bottom(x, y) = - 1000 * Int((x < 10) | (x > 50)) 
+
+        set!(bottom_height, bottom)
+        set!(control_bottom_height, bottom)
+
+        # This should have not changed anything
+        remove_minor_basins!(bottom_height, 2)
+        @test parent(bottom_height) == parent(control_bottom_height)
+
+        # This should have removed the left basin
+        remove_minor_basins!(bottom_height, 1)
+
+        # The remaning bottom cells that are not immersed should be only on the right hand side
+        # The left half of the domain should be fully immersed i.e. bottom == 0
+        @test sum(view(bottom_height, 1:50, :, 1)) == 0
+
+        # While the right side should be not immersed, with a mean bottom depth 
+        # of -1000 meters
+        @test mean(view(bottom_height, 51:100, :, 1)) == -1000
+
+        grid = LatitudeLongitudeGrid(arch;
+                                     size = (200, 200, 10), 
+                                     longitude = (0, 2), 
+                                     latitude = (-10, 50),
+                                     z = (-6000, 0))
+
+        control_bottom_height = regrid_bathymetry(grid)
+        interpolated_bottom_height = regrid_bathymetry(grid; interpolation_passes = 100)
+
+        # Testing that multiple passes do not change the solution when refining the grid
+        @test parent(control_bottom_height) == parent(interpolated_bottom_height)
+    end
+end 
diff --git a/test/test_jra55.jl b/test/test_jra55.jl
@@ -92,16 +92,15 @@ using ClimaOcean.OceanSeaIceModels: PrescribedAtmosphere
 
         # Test that we can load the data back
         Qswt = FieldTimeSeries(filepath, "Qsw")
-        @show size(Qswt.data), size(target_fts)
-        @test parent(Qswt) == parent(target_fts)    
+        @test on_architecture(CPU(), parent(Qswt.data)) == on_architecture(CPU(), parent(target_fts.data))
         @test Qswt.times == target_fts.times
         rm(filepath)
 
         #####
         ##### JRA55 prescribed atmosphere
         #####
 
-        backend    = JRA55NetCDFBackend(1) 
+        backend    = JRA55NetCDFBackend(2) 
         atmosphere = JRA55_prescribed_atmosphere(arch; backend, include_rivers_and_icebergs=false)
         @test atmosphere isa PrescribedAtmosphere
         @test isnothing(atmosphere.auxiliary_freshwater_flux)

diff --git a/test/test_simulations.jl b/test/test_simulations.jl
@@ -4,37 +4,37 @@ using OrthogonalSphericalShellGrids
 
 @testset "Parameter space test" begin
 
-    arch = GPU() 
-    grid = TripolarGrid(arch; 
-                        size = (50, 50, 10), 
-                        halo = (7, 7, 7), 
-                        z = (-6000, 0), 
-                        first_pole_longitude = 75,
-                        north_poles_latitude = 55)
-
-    bottom_height = retrieve_bathymetry(grid; 
-                                        minimum_depth = 10,
-                                        dir = "./",
-                                        interpolation_passes = 20,
-                                        major_basins = 1)
-
-    grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height); active_cells_map = true)                           
-
-    free_surface = SplitExplicitFreeSurface(grid; substeps = 20)
-    ocean = ocean_simulation(grid; free_surface) 
-
-    backend    = JRA55NetCDFBackend(4) 
-    atmosphere = JRA55_prescribed_atmosphere(arch; backend)
-    radiation  = Radiation(arch)
-
-    sea_ice = ClimaOcean.OceanSeaIceModels.MinimumTemperatureSeaIce()
-
-    # Fluxes are computed when the model is constructed, so we just test that this works.
-    @test begin
-        coupled_model = OceanSeaIceModel(ocean, sea_ice; atmosphere, radiation)
-        true
+    for arch in test_architectures
+        grid = TripolarGrid(arch; 
+                            size = (50, 50, 10), 
+                            halo = (7, 7, 7), 
+                            z = (-6000, 0), 
+                            first_pole_longitude = 75,
+                            north_poles_latitude = 55)
+
+        bottom_height = retrieve_bathymetry(grid; 
+                                            minimum_depth = 10,
+                                            dir = "./",
+                                            interpolation_passes = 20,
+                                            major_basins = 1)
+
+        grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height); active_cells_map = true)                           
+
+        free_surface = SplitExplicitFreeSurface(grid; substeps = 20)
+        ocean = ocean_simulation(grid; free_surface) 
+
+        backend    = JRA55NetCDFBackend(4) 
+        atmosphere = JRA55_prescribed_atmosphere(arch; backend)
+        radiation  = Radiation(arch)
+
+        sea_ice = ClimaOcean.OceanSeaIceModels.MinimumTemperatureSeaIce()
+
+        # Fluxes are computed when the model is constructed, so we just test that this works.
+        @test begin
+            coupled_model = OceanSeaIceModel(ocean, sea_ice; atmosphere, radiation)
+            true
+        end
     end
-
 end