Persist FractionalIndices with two separate arrays for i, j (#371)

* one missing allowscalar

* small cleanup

* bugfix

* New way of constructing fractional indices

* add topoology

* Tuple of fractional indices

* Bugfix

* Bump ClimaSeaICe

* Fix bug in atmos interpolation

* Try to generalize to doubly-Flat grids

* Bug in StateExchanger construcotr

* Bump Oceananigans compat

* Update Project.toml

* Update Project.toml

* fix test simulation

* Fix interpolation for Flat topologies?

* Ohhh

* Fix test bug

* Better name

* Name mangle

* Fix test

* Tiny change in MixedLayerDepth diagnostic

---------

Co-authored-by: Gregory Wagner <[email protected]>
Co-authored-by: Gregory Wagner <[email protected]>

Author: simone-silvestri

Project.toml

@@ -33,7 +33,7 @@ Thermodynamics = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c"
 Adapt = "4"
 CFTime = "0.1"
 CUDA = "4, 5"
-ClimaSeaIce = "0.2.0"
+ClimaSeaIce = "0.2.1"
 CubicSplines = "0.2"
 DataDeps = "0.7"
 Downloads = "1.6"
@@ -42,7 +42,7 @@ JLD2 = "0.4, 0.5"
 KernelAbstractions = "0.9"
 MPI = "0.20"
 NCDatasets = "0.12, 0.13, 0.14"
-Oceananigans = "0.95.17 - 0.99"
+Oceananigans = "0.95.18 - 0.99"
 OffsetArrays = "1.14"
 OrthogonalSphericalShellGrids = "0.2.2"
 Scratch = "1"

src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl

@@ -2,7 +2,6 @@ using StaticArrays
 using Thermodynamics
 using SurfaceFluxes
 using OffsetArrays
-using CUDA: @allowscalar
 
 using ..OceanSeaIceModels: reference_density,
                            heat_capacity,
@@ -15,7 +14,7 @@ using ..OceanSeaIceModels: reference_density,
 using ClimaSeaIce: SeaIceModel
 
 using Oceananigans: HydrostaticFreeSurfaceModel, architecture
-using Oceananigans.Grids: inactive_node, node
+using Oceananigans.Grids: inactive_node, node, topology
 using Oceananigans.BoundaryConditions: fill_halo_regions!
 
 using Oceananigans.Fields: ConstantField, interpolate, FractionalIndices
@@ -60,6 +59,11 @@ struct StateExchanger{G, AST, AEX}
     atmosphere_exchanger :: AEX
 end
 
+# Note that Field location can also affect fractional index type.
+# Here we assume that we know the location of Fields that will be interpolated.
+fractional_index_type(FT, Topo) = FT
+fractional_index_type(FT, ::Flat) = Nothing
+
 function StateExchanger(ocean::Simulation, atmosphere)
     # TODO: generalize this
     exchange_grid = ocean.model.grid
@@ -80,13 +84,14 @@ function StateExchanger(ocean::Simulation, atmosphere)
     arch = architecture(exchange_grid)
     Nx, Ny, Nz = size(exchange_grid)
 
-    # Make an array of FractionalIndices
-    kᴺ = size(exchange_grid, 3)
-    @allowscalar X1 = _node(1, 1, kᴺ + 1, exchange_grid, c, c, f)
-    i1 = FractionalIndices(X1, atmosphere.grid, c, c, nothing)
-    frac_indices_data = [deepcopy(i1) for i=1:Nx+2, j=1:Ny+2, k=1:1]
-    frac_indices = OffsetArray(frac_indices_data, -1, -1, 0)
-    frac_indices = on_architecture(arch, frac_indices)
+    # Make a NamedTuple of fractional indices
+    # Note: we could use an array of FractionalIndices. Instead, for compatbility
+    # with Reactant we construct FractionalIndices on the fly in `interpolate_atmospheric_state`.
+    FT = eltype(atmos_grid)
+    TX, TY, TZ = topology(exchange_grid)
+    fi = TX() isa Flat ? nothing : Field{Center, Center, Nothing}(exchange_grid, FT)
+    fj = TY() isa Flat ? nothing : Field{Center, Center, Nothing}(exchange_grid, FT)
+    frac_indices = (i=fi, j=fj) # no k needed, only horizontal interpolation
 
     kernel_parameters = interface_kernel_parameters(exchange_grid)
     launch!(arch, exchange_grid, kernel_parameters,
@@ -95,11 +100,22 @@ function StateExchanger(ocean::Simulation, atmosphere)
     return StateExchanger(ocean.model.grid, exchange_atmosphere_state, frac_indices)
 end
 
-@kernel function _compute_fractional_indices!(frac_indices, exchange_grid, atmos_grid)
+@kernel function _compute_fractional_indices!(indices_tuple, exchange_grid, atmos_grid)
     i, j = @index(Global, NTuple)
     kᴺ = size(exchange_grid, 3) # index of the top ocean cell
     X = _node(i, j, kᴺ + 1, exchange_grid, c, c, f)
-    @inbounds frac_indices[i, j, 1] = FractionalIndices(X, atmos_grid, c, c, nothing)
+    fractional_indices_ij = FractionalIndices(X, atmos_grid, c, c, nothing)
+    fi = indices_tuple.i
+    fj = indices_tuple.j
+    @inbounds begin
+        if !isnothing(fi)
+            fi[i, j, 1] = fractional_indices_ij.i
+        end
+
+        if !isnothing(fj)
+            fj[i, j, 1] = fractional_indices_ij.j
+        end
+    end
 end
 
 const PATP = PrescribedAtmosphereThermodynamicsParameters

src/OceanSeaIceModels/InterfaceComputations/interpolate_atmospheric_state.jl

@@ -1,5 +1,6 @@
 using Oceananigans.Operators: intrinsic_vector
 using Oceananigans.Grids: _node
+using Oceananigans.Fields: FractionalIndices
 using Oceananigans.OutputReaders: TimeInterpolator
 
 using ...OceanSimulations: forcing_barotropic_potential
@@ -109,6 +110,9 @@ function interpolate_atmospheric_state!(coupled_model)
         parent(barotropic_potential) .= parent(atmosphere_data.p) ./ ρₒ
     end
 end
+
+@inline get_fractional_index(i, j, ::Nothing) = nothing
+@inline get_fractional_index(i, j, frac) = @inbounds frac[i, j, 1]
 
 @kernel function _interpolate_primary_atmospheric_state!(surface_atmos_state,
                                                          space_fractional_indices,
@@ -124,28 +128,33 @@ end
 
     i, j = @index(Global, NTuple)
 
-    @inbounds begin
-        x_itp = space_fractional_indices[i, j, 1]
-        t_itp = time_interpolator
-        atmos_args = (x_itp, t_itp, atmos_backend, atmos_time_indexing)
+    ii = space_fractional_indices.i
+    jj = space_fractional_indices.j
+    fi = get_fractional_index(i, j, ii)
+    fj = get_fractional_index(i, j, jj)
 
-        uₐ = interp_atmos_time_series(atmos_velocities.u, atmos_args...)
-        vₐ = interp_atmos_time_series(atmos_velocities.v, atmos_args...)
-        Tₐ = interp_atmos_time_series(atmos_tracers.T,    atmos_args...)
-        qₐ = interp_atmos_time_series(atmos_tracers.q,    atmos_args...)
-        pₐ = interp_atmos_time_series(atmos_pressure,     atmos_args...)
+    x_itp = FractionalIndices(fi, fj, nothing)
+    t_itp = time_interpolator
+    atmos_args = (x_itp, t_itp, atmos_backend, atmos_time_indexing)
 
-        Qs = interp_atmos_time_series(downwelling_radiation.shortwave, atmos_args...)
-        Qℓ = interp_atmos_time_series(downwelling_radiation.longwave,  atmos_args...)
+    uₐ = interp_atmos_time_series(atmos_velocities.u, atmos_args...)
+    vₐ = interp_atmos_time_series(atmos_velocities.v, atmos_args...)
+    Tₐ = interp_atmos_time_series(atmos_tracers.T,    atmos_args...)
+    qₐ = interp_atmos_time_series(atmos_tracers.q,    atmos_args...)
+    pₐ = interp_atmos_time_series(atmos_pressure,     atmos_args...)
 
-        # Usually precipitation
-        Mh = interp_atmos_time_series(prescribed_freshwater_flux, atmos_args...)
+    Qs = interp_atmos_time_series(downwelling_radiation.shortwave, atmos_args...)
+    Qℓ = interp_atmos_time_series(downwelling_radiation.longwave,  atmos_args...)
 
-        # Convert atmosphere velocities (usually defined on a latitude-longitude grid) to 
-        # the frame of reference of the native grid
-        kᴺ = size(exchange_grid, 3) # index of the top ocean cell
-        uₐ, vₐ = intrinsic_vector(i, j, kᴺ, exchange_grid, uₐ, vₐ)
+    # Usually precipitation
+    Mh = interp_atmos_time_series(prescribed_freshwater_flux, atmos_args...)
 
+    # Convert atmosphere velocities (usually defined on a latitude-longitude grid) to 
+    # the frame of reference of the native grid
+    kᴺ = size(exchange_grid, 3) # index of the top ocean cell
+    uₐ, vₐ = intrinsic_vector(i, j, kᴺ, exchange_grid, uₐ, vₐ)
+
+    @inbounds begin
         surface_atmos_state.u[i, j, 1] = uₐ
         surface_atmos_state.v[i, j, 1] = vₐ
         surface_atmos_state.T[i, j, 1] = Tₐ

test/test_diagnostics.jl

@@ -42,11 +42,11 @@ using ClimaOcean.Diagnostics: MixedLayerDepthField, MixedLayerDepthOperand
         @test h.operand.buoyancy_perturbation isa KernelFunctionOperation
 
         compute!(h)
-        @test @allowscalar h[1, 1, 1] ≈ 16.255836 # m
+        @test @allowscalar h[1, 1, 1] ≈ 16.2558363 # m
 
         tracers = (T=Tt[2], S=St[2])
         h.operand.buoyancy_perturbation = buoyancy(sb, grid, tracers)
         compute!(h)
-        @test @allowscalar h[1, 1, 1] ≈ 9.295890287 # m
+        @test @allowscalar h[1, 1, 1] ≈ 9.2957298 # m
     end
 end

test/test_simulations.jl

@@ -5,7 +5,9 @@ using OrthogonalSphericalShellGrids
 using ClimaOcean.OceanSeaIceModels: above_freezing_ocean_temperature!
 using ClimaSeaIce.SeaIceThermodynamics: melting_temperature
 
-@testset "GPU time stepping test" begin
+@inline kernel_melting_temperature(i, j, k, grid, liquidus, S) = @inbounds melting_temperature(liquidus, S[i, j, k])
+
+@testset "Time stepping test" begin
 
     for arch in test_architectures
 
@@ -24,9 +26,9 @@ using ClimaSeaIce.SeaIceThermodynamics: melting_temperature
                                             interpolation_passes = 20,
                                             major_basins = 1)
 
-        grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height); active_cells_map = true)
+        grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height); active_cells_map=true)
 
-        free_surface = SplitExplicitFreeSurface(grid; substeps = 20)
+        free_surface = SplitExplicitFreeSurface(grid; substeps=20)
         ocean = ocean_simulation(grid; free_surface) 
 
         backend = JRA55NetCDFBackend(4)
@@ -52,24 +54,22 @@ using ClimaSeaIce.SeaIceThermodynamics: melting_temperature
 
         # Set the ocean temperature and salinity
         set!(ocean.model, T=temperature_metadata[1], S=salinity_metadata[1])
-
         above_freezing_ocean_temperature!(ocean, sea_ice)
 
         # Test that ocean temperatures are above freezing
-        T = on_architecture(CPU(), ocean.model.T)
-        S = on_architecture(CPU(), ocean.model.S)
-
-        @inline pointwise_melting_T(i, j, k, grid, liquidus, S) = @inbounds melting_temperature(liquidus, S[i, j, k])
+        T = on_architecture(CPU(), ocean.model.tracers.T)
+        S = on_architecture(CPU(), ocean.model.tracers.S)
 
-        Tm = KernelFunctionOperation{Center, Center, Center}(pointwise_melting_T, grid, S)
-
-        @test all(T .> Tm)
+        Tm = KernelFunctionOperation{Center, Center, Center}(kernel_melting_temperature, grid, liquidus, S)
+        @test all(T .>= Tm)
 
         # Fluxes are computed when the model is constructed, so we just test that this works.
         # And that we can time step with sea ice
         @test begin
             coupled_model = OceanSeaIceModel(ocean, sea_ice; atmosphere, radiation)
+            time_step!(coupled_model, 1)
             true
         end
     end
 end
+