Fixes for flux computation + add fluxes regression tests (#271)

* avoid extra step

* Update atmosphere_ocean_fluxes.jl

* put atmosphere rotation in interpolation

* some improvement

* try adding Manifest

* run the tests

* try new manifest

* correct fluxes

* add tests

* correct regression

* remove Manifest

* fix tests

* add units

* fix tests and update Oceananigans

* update oceananigans

* run tests

* now it should work

Author: simone-silvestri

Project.toml

@@ -42,7 +42,7 @@ JLD2 = "0.4, 0.5"
 KernelAbstractions = "0.9"
 MPI = "0.20"
 NCDatasets = "0.12, 0.13, 0.14"
-Oceananigans = "0.94.2 - 0.99"
+Oceananigans = "0.94.3 - 0.99"
 OffsetArrays = "1.14"
 OrthogonalSphericalShellGrids = "0.1.8"
 Scratch = "1"

src/DataWrangling/ECCO/ECCO.jl

@@ -191,7 +191,12 @@ function ECCO_field(metadata::ECCOMetadata;
     # data by 180 degrees in longitude
     if metadata.version isa ECCO4Monthly 
         Nx = size(data, 1)
-        data = circshift(data, (Nx ÷ 2, 0, 0))
+        if variable_is_three_dimensional(metadata)
+            shift = (Nx ÷ 2, 0, 0)
+        else
+            shift = (Nx ÷ 2, 0)
+        end
+        data = circshift(data, shift)
     end
 
     set!(field, data)

src/DataWrangling/ECCO/ECCO_metadata.jl

@@ -166,7 +166,8 @@ ECCO4_short_names = Dict(
     :u_velocity            => "EVEL",
     :v_velocity            => "NVEL",
     :sea_ice_thickness     => "SIheff",
-    :sea_ice_area_fraction => "SIarea"
+    :sea_ice_area_fraction => "SIarea",
+    :net_heat_flux         => "oceQnet"
 )
 
 ECCO2_short_names = Dict(
@@ -175,14 +176,16 @@ ECCO2_short_names = Dict(
     :u_velocity            => "UVEL",
     :v_velocity            => "VVEL",
     :sea_ice_thickness     => "SIheff",
-    :sea_ice_area_fraction => "SIarea"
+    :sea_ice_area_fraction => "SIarea",
+    :net_heat_flux         => "oceQnet"
 )
 
 ECCO_location = Dict(
     :temperature           => (Center, Center, Center),
     :salinity              => (Center, Center, Center),
     :sea_ice_thickness     => (Center, Center, Nothing),
     :sea_ice_area_fraction => (Center, Center, Nothing),
+    :net_heat_flux         => (Center, Center, Nothing),
     :u_velocity            => (Face,   Center, Center),
     :v_velocity            => (Center, Face,   Center),
 )

src/OceanSeaIceModels/CrossRealmFluxes/atmosphere_ocean_fluxes.jl

@@ -365,9 +365,9 @@ end
     time = Time(clock.time)
 
     @inbounds begin
-        Tₒ = ocean_temperature[i, j, 1]
+        Tₒ = ocean_temperature[i, j, kᴺ]
         Tₒ = convert_to_kelvin(ocean_temperature_units, Tₒ)
-        Sₒ = ocean_salinity[i, j, 1]
+        Sₒ = ocean_salinity[i, j, kᴺ]
 
         Qs = downwelling_radiation.shortwave[i, j, 1]
         Qℓ = downwelling_radiation.longwave[i, j, 1]

src/OceanSeaIceModels/CrossRealmFluxes/similarity_theory_turbulent_fluxes.jl

@@ -310,9 +310,9 @@ end
 # Iterating condition for the characteristic scales solvers
 @inline function iterating(Σ★, iteration, maxiter, solver)
     havent_started = iteration == 0
-    not_converged = norm(Σ★) > solver.tolerance
-    havent_reached_maxiter = iteration < maxiter
-    return havent_started | not_converged | havent_reached_maxiter
+    converged = norm(Σ★) < solver.tolerance
+    reached_maxiter = iteration ≥ maxiter
+    return !(converged | reached_maxiter) | havent_started
 end
 
 """

src/OceanSeaIceModels/OceanSeaIceModels.jl

@@ -62,4 +62,8 @@ const NoAtmosphereModel = OceanSeaIceModel{<:Any, Nothing}
 
 compute_atmosphere_ocean_fluxes!(coupled_model::NoAtmosphereModel) = nothing
 
+const NoSeaIceModel = OceanSeaIceModel{Nothing}
+
+compute_sea_ice_ocean_fluxes!(cm::NoSeaIceModel) = nothing
+
 end # module

src/OceanSeaIceModels/time_step_ocean_sea_ice_model.jl

@@ -29,8 +29,7 @@ function time_step!(coupled_model::OceanSeaIceModel, Δt; callbacks=[], compute_
     end
 
     # TODO after ice time-step:
-    #   - Adjust ocean heat flux if the ice completely melts?
-
+    #  - Adjust ocean heat flux if the ice completely melts?
     ocean.Δt = Δt
     time_step!(ocean)
 

test/test_surface_fluxes.jl

@@ -10,6 +10,10 @@ using ClimaOcean.OceanSeaIceModels.CrossRealmFluxes:
 using Thermodynamics
 using CUDA
 using KernelAbstractions: @kernel, @index
+using Oceananigans.TimeSteppers: update_state!
+using Oceananigans.Units: hours, days
+
+using Statistics: mean, std
 
 import ClimaOcean.OceanSeaIceModels.CrossRealmFluxes: water_saturation_specific_humidity
 
@@ -196,4 +200,76 @@ _fractional_indices(at_node, grid, ::Nothing, ::Nothing, ::Nothing) = (nothing,
     end
 end
 
+@testset "Fluxes regression" begin
+    for arch in test_architectures
+        @info "Testing fluxes regression..."
+
+        grid = LatitudeLongitudeGrid(arch; 
+                                     size = (20, 20, 20), 
+                                 latitude = (-60, 60), 
+                                longitude = (0, 360),
+                                        z = (-5000, 0))
+
+        # Speed up compilation by removing all the unnecessary stuff
+        momentum_advection = nothing
+        tracer_advection   = nothing
+        tracers  = (:T, :S)
+        buoyancy = nothing
+        closure  = nothing
+        coriolis = nothing
+
+        ocean = ocean_simulation(grid; momentum_advection, tracer_advection, closure, tracers, coriolis)
+
+        T_metadata = ECCOMetadata(:temperature)
+        S_metadata = ECCOMetadata(:salinity)
+
+        set!(ocean.model; T=T_metadata, S=S_metadata)
+
+        atmosphere = JRA55_prescribed_atmosphere(1:10; grid, architecture = arch, backend = InMemory())
+        radiation  = Radiation(ocean_albedo=0.1, ocean_emissivity=1.0)
+        sea_ice    = nothing
+
+        coupled_model = OceanSeaIceModel(ocean, sea_ice; atmosphere, radiation)
+
+        times = 0:1hours:1days
+        Ntimes = length(times)
+
+        # average the fluxes over one day
+        Jᵀ = interior(ocean.model.tracers.T.boundary_conditions.top.condition, :, :, 1) ./ Ntimes
+        Jˢ = interior(ocean.model.tracers.S.boundary_conditions.top.condition, :, :, 1) ./ Ntimes
+        τˣ = interior(ocean.model.velocities.u.boundary_conditions.top.condition, :, :, 1) ./ Ntimes
+        τʸ = interior(ocean.model.velocities.v.boundary_conditions.top.condition, :, :, 1) ./ Ntimes
+
+        for time in times[2:end]
+            coupled_model.clock.time = time
+            update_state!(coupled_model)
+            Jᵀ .+= interior(ocean.model.tracers.T.boundary_conditions.top.condition, :, :, 1) ./ Ntimes
+            Jˢ .+= interior(ocean.model.tracers.S.boundary_conditions.top.condition, :, :, 1) ./ Ntimes
+            τˣ .+= interior(ocean.model.velocities.u.boundary_conditions.top.condition, :, :, 1) ./ Ntimes
+            τʸ .+= interior(ocean.model.velocities.v.boundary_conditions.top.condition, :, :, 1) ./ Ntimes
+        end
+
+        Jᵀ_mean = mean(Jᵀ) 
+        Jˢ_mean = mean(Jˢ) 
+        τˣ_mean = mean(τˣ) 
+        τʸ_mean = mean(τʸ) 
+
+        Jᵀ_std = std(Jᵀ)
+        Jˢ_std = std(Jˢ)
+        τˣ_std = std(τˣ)
+        τʸ_std = std(τʸ)
+
+        # Regression test
+        @test Jᵀ_mean ≈ -3.526464713488678e-5
+        @test Jˢ_mean ≈ 1.1470078542716042e-6
+        @test τˣ_mean ≈ -1.0881334225579832e-5
+        @test τʸ_mean ≈ 5.653281786086694e-6
+            
+        @test Jᵀ_std ≈ 7.477575901188957e-5
+        @test Jˢ_std ≈ 3.7416720607945508e-6
+        @test τˣ_std ≈ 0.00011349625113971719
+        @test τʸ_std ≈ 7.627885224680635e-5
+    end
+end
+