Merge branch 'main' into ts/checkpointer-for-coupled-model

Author: navidcy

Project.toml

@@ -2,7 +2,7 @@ name = "ClimaOcean"
 uuid = "0376089a-ecfe-4b0e-a64f-9c555d74d754"
 license = "MIT"
 authors = ["Climate Modeling Alliance and contributors"]
-version = "0.5.1"
+version = "0.5.3"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -20,6 +20,7 @@ MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
 Oceananigans = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
+PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Scratch = "6c6a2e73-6563-6170-7368-637461726353"
 SeawaterPolynomials = "d496a93d-167e-4197-9f49-d3af4ff8fe40"
@@ -41,8 +42,9 @@ JLD2 = "0.4, 0.5"
 KernelAbstractions = "0.9"
 MPI = "0.20"
 NCDatasets = "0.12, 0.13, 0.14"
-Oceananigans = "0.95.26 - 0.99"
+Oceananigans = "0.95.27 - 0.99"
 OffsetArrays = "1.14"
+PrecompileTools = "1"
 Scratch = "1"
 SeawaterPolynomials = "0.3.4"
 StaticArrays = "1"

README.md

@@ -5,10 +5,10 @@
 
 <!-- description -->
 <p align="center">
-  <strong>🌎 A framework for realistic ocean-only and coupled ocean + sea-ice simulations driven by prescribed atmospheres and based on <a href=https://github.com/CliMA/Oceananigans.jl>Oceananigans</a> and <a href=https://github.com/CliMA/ClimaSeaIce.jl>ClimaSeaIce</a></strong>. 
+  <strong>🌎 A framework for realistic ocean-only and coupled ocean + sea-ice simulations driven by prescribed atmospheres and based on <a href=https://github.com/CliMA/Oceananigans.jl>Oceananigans</a> and <a href=https://github.com/CliMA/ClimaSeaIce.jl>ClimaSeaIce</a></strong>.
 </p>
 
-### 
+###
 
 [![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.7677442.svg?style=flat-square)](https://doi.org/10.5281/zenodo.7677442)
 [![Build status](https://badge.buildkite.com/3113cca353b83df3b5855d3f0d69827124614aef7017c835d2.svg?style=flat-square)](https://buildkite.com/clima/climaocean-ci)
@@ -32,7 +32,7 @@ For more information, see the [documentation for `Pkg.jl`](https://pkgdocs.julia
 
 ## Why? What's the difference between ClimaOcean and [Oceananigans](https://github.com/CliMA/Oceananigans.jl)?
 
-Oceananigans is a general-purpose library for ocean-flavored fluid dynamics. 
+Oceananigans is a general-purpose library for ocean-flavored fluid dynamics.
 ClimaOcean implements a framework for driving realistic Oceananigans simulations with prescribed atmospheres, and coupling them to prognostic sea ice simulations.
 
 ### A core abstraction: `ClimaOcean.OceanSeaIceModel`
@@ -59,9 +59,10 @@ grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bathymetry))
 
 # Build an ocean simulation initialized to the ECCO state estimate on Jan 1, 1993
 ocean = ClimaOcean.ocean_simulation(grid)
-dates = DateTimeProlepticGregorian(1993, 1, 1)
-set!(ocean.model, T = ClimaOcean.ECCOMetadata(:temperature; dates),
-                  S = ClimaOcean.ECCOMetadata(:salinity; dates))
+dates = DateTime(1993, 1, 1) : Month(1) : DateTime(1994, 1, 1)
+set!(ocean.model,
+     T=ClimaOcean.Metadata(:temperature; dates=first(dates), dataset=ClimaOcean.ECCO4Monthly()),
+     S=ClimaOcean.Metadata(:salinity;    dates=first(dates), dataset=ClimaOcean.ECCO4Monthly()))
 
 # Build and run an OceanSeaIceModel (with no sea ice component) forced by JRA55 reanalysis
 atmosphere = ClimaOcean.JRA55PrescribedAtmosphere(arch)

docs/src/index.md

@@ -25,7 +25,7 @@ julia> Pkg.add("ClimaOcean")
 
 !!! compat "Julia 1.10 is required"
     ClimaOcean requires Julia 1.10 or later.
-    
+
 ## Quick start
 
 The following script implements a near-global ocean simulation initialized from the [ECCO state estimate](https://gmd.copernicus.org/articles/8/3071/2015/) and coupled to a prescribed atmosphere derived from the [JRA55-do reanalysis](https://www.sciencedirect.com/science/article/pii/S146350031830235X):
@@ -49,9 +49,10 @@ grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bathymetry))
 
 # Build an ocean simulation initialized to the ECCO state estimate on Jan 1, 1993
 ocean = ClimaOcean.ocean_simulation(grid)
-date  = DateTime(1993, 1, 1)
-set!(ocean.model, T = ClimaOcean.ECCOMetadata(:temperature; date),
-                  S = ClimaOcean.ECCOMetadata(:salinity; date))
+dates = DateTime(1993, 1, 1) : Month(1) : DateTime(1994, 1, 1)
+set!(ocean.model,
+     T=ClimaOcean.Metadata(:temperature; dates=first(dates), dataset=ClimaOcean.ECCO4Monthly()),
+     S=ClimaOcean.Metadata(:salinity;    dates=first(dates), dataset=ClimaOcean.ECCO4Monthly()))
 
 # Build and run an OceanSeaIceModel (with no sea ice component) forced by JRA55 reanalysis
 atmosphere = ClimaOcean.JRA55PrescribedAtmosphere(arch)

src/ClimaOcean.jl

@@ -91,5 +91,18 @@ using ClimaOcean.OceanSeaIceModels: PrescribedAtmosphere
 using ClimaOcean.DataWrangling.JRA55: JRA55PrescribedAtmosphere, JRA55NetCDFBackend
 using ClimaOcean.DataWrangling.ECCO
 
+using PrecompileTools: @setup_workload, @compile_workload
+
+@setup_workload begin
+    Nx, Ny, Nz = 32, 32, 10
+    @compile_workload begin
+        z = exponential_z_faces(Nz=Nz, depth=6000, h=34)
+        grid = Oceananigans.OrthogonalSphericalShellGrids.TripolarGrid(CPU(); size=(Nx, Ny, Nz), halo=(7, 7, 7), z)
+        grid = ImmersedBoundaryGrid(grid, GridFittedBottom((x, y) -> -5000))
+        ocean = ocean_simulation(grid)
+        model = OceanSeaIceModel(ocean)
+    end
+end
+
 end # module
 

src/OceanSeaIceModels/InterfaceComputations/atmosphere_sea_ice_fluxes.jl

@@ -75,8 +75,7 @@ end
 
     i, j = @index(Global, NTuple)
     kᴺ   = size(grid, 3) # index of the top ocean cell
-    time = Time(clock.time)
-    FT = eltype(grid)
+    FT   = eltype(grid)
 
     @inbounds begin
         uₐ = atmosphere_state.u[i, j, 1]

src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl

@@ -28,25 +28,27 @@ using Oceananigans.Operators: ℑxᶜᵃᵃ, ℑyᵃᶜᵃ, ℑxᶠᵃᵃ, ℑy
 
 using KernelAbstractions: @kernel, @index
 
+import Oceananigans.Simulations: initialize!
+
 #####
 ##### Container for organizing information related to fluxes
 #####
 
-struct AtmosphereInterface{J, F, ST, P}
+mutable struct AtmosphereInterface{J, F, ST, P}
     fluxes :: J
     flux_formulation :: F
     temperature :: ST
     properties :: P
 end
 
-struct SeaIceOceanInterface{J, P, H, A}
+mutable struct SeaIceOceanInterface{J, P, H, A}
     fluxes :: J
     properties :: P
     previous_ice_thickness :: H
     previous_ice_concentration :: A
 end
 
-struct ComponentInterfaces{AO, ASI, SIO, C, AP, OP, SIP, EX}
+mutable struct ComponentInterfaces{AO, ASI, SIO, C, AP, OP, SIP, EX}
     atmosphere_ocean_interface :: AO
     atmosphere_sea_ice_interface :: ASI
     sea_ice_ocean_interface :: SIO
@@ -57,30 +59,44 @@ struct ComponentInterfaces{AO, ASI, SIO, C, AP, OP, SIP, EX}
     net_fluxes :: C
 end
 
-struct StateExchanger{G, AST, AEX}
+mutable struct StateExchanger{G, AST, AEX}
     exchange_grid :: G
     exchange_atmosphere_state :: AST
     atmosphere_exchanger :: AEX
 end
 
+mutable struct ExchangeAtmosphereState{F}
+    u  :: F
+    v  :: F
+    T  :: F
+    q  :: F
+    p  :: F
+    Qs :: F
+    Qℓ :: F
+    Mp :: F
+end
+
+ExchangeAtmosphereState(grid) = ExchangeAtmosphereState(Field{Center, Center, Nothing}(grid),
+                                                        Field{Center, Center, Nothing}(grid),
+                                                        Field{Center, Center, Nothing}(grid),
+                                                        Field{Center, Center, Nothing}(grid),
+                                                        Field{Center, Center, Nothing}(grid),
+                                                        Field{Center, Center, Nothing}(grid),
+                                                        Field{Center, Center, Nothing}(grid),
+                                                        Field{Center, Center, Nothing}(grid))
+
 # 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
 
+
+StateExchanger(ocean::Simulation, ::Nothing) = nothing
+
 function StateExchanger(ocean::Simulation, atmosphere)
     # TODO: generalize this
     exchange_grid = ocean.model.grid
-
-    exchange_atmosphere_state = (u  = Field{Center, Center, Nothing}(exchange_grid),
-                                 v  = Field{Center, Center, Nothing}(exchange_grid),
-                                 T  = Field{Center, Center, Nothing}(exchange_grid),
-                                 q  = Field{Center, Center, Nothing}(exchange_grid),
-                                 p  = Field{Center, Center, Nothing}(exchange_grid),
-                                 Qs = Field{Center, Center, Nothing}(exchange_grid),
-                                 Qℓ = Field{Center, Center, Nothing}(exchange_grid),
-                                 Mp = Field{Center, Center, Nothing}(exchange_grid))
-
+    exchange_atmosphere_state = ExchangeAtmosphereState(exchange_grid)
     exchanger = atmosphere_exchanger(atmosphere, exchange_grid)
 
     return StateExchanger(ocean.model.grid, exchange_atmosphere_state, exchanger)
@@ -100,11 +116,20 @@ function atmosphere_exchanger(atmosphere::PrescribedAtmosphere, exchange_grid)
     fj = TY() isa Flat ? nothing : Field{Center, Center, Nothing}(exchange_grid, FT)
     frac_indices = (i=fi, j=fj) # no k needed, only horizontal interpolation
 
+    return frac_indices
+end
+
+initialize!(exchanger::StateExchanger, ::Nothing) = nothing
+
+function initialize!(exchanger::StateExchanger, atmosphere)
+    atmos_grid = atmosphere.grid
+    exchange_grid = exchanger.exchange_grid
+    arch = architecture(exchange_grid)
+    frac_indices = exchanger.atmosphere_exchanger
     kernel_parameters = interface_kernel_parameters(exchange_grid)
     launch!(arch, exchange_grid, kernel_parameters,
             _compute_fractional_indices!, frac_indices, exchange_grid, atmos_grid)
-
-    return frac_indices
+    return nothing
 end
 
 @kernel function _compute_fractional_indices!(indices_tuple, exchange_grid, atmos_grid)
@@ -141,8 +166,11 @@ const celsius_to_kelvin = 273.15
 Base.summary(crf::ComponentInterfaces) = "ComponentInterfaces"
 Base.show(io::IO, crf::ComponentInterfaces) = print(io, summary(crf))
 
-function atmosphere_ocean_interface(ocean,
-                                    radiation,
+atmosphere_ocean_interface(::Nothing, args...) = nothing
+
+function atmosphere_ocean_interface(atmos,
+                                    ocean, 
+                                    radiation, 
                                     ao_flux_formulation,
                                     temperature_formulation, 
                                     velocity_formulation,
@@ -170,10 +198,14 @@ function atmosphere_ocean_interface(ocean,
     return AtmosphereInterface(ao_fluxes, ao_flux_formulation, interface_temperature, ao_properties)
 end
 
-atmosphere_sea_ice_interface(sea_ice, args...) = nothing
+atmosphere_sea_ice_interface(atmos, sea_ice, args...) = nothing
+atmosphere_sea_ice_interface(::Nothing, args...) = nothing
+atmosphere_sea_ice_interface(::Nothing, ::Nothing, args...) = nothing
+atmosphere_sea_ice_interface(::Nothing, ::SeaIceSimulation, args...) = nothing
 
-function atmosphere_sea_ice_interface(sea_ice::SeaIceSimulation,
-                                      radiation,
+function atmosphere_sea_ice_interface(atmos,
+                                      sea_ice::SeaIceSimulation, 
+                                      radiation, 
                                       ai_flux_formulation,
                                       temperature_formulation,
                                       velocity_formulation)
@@ -295,7 +327,8 @@ function ComponentInterfaces(atmosphere, ocean, sea_ice=nothing;
                         freshwater_density = freshwater_density,
                         temperature_units  = ocean_temperature_units)
 
-    ao_interface = atmosphere_ocean_interface(ocean,
+    ao_interface = atmosphere_ocean_interface(atmosphere,
+                                              ocean,
                                               radiation,
                                               atmosphere_ocean_flux_formulation,
                                               atmosphere_ocean_interface_temperature,
@@ -304,7 +337,8 @@ function ComponentInterfaces(atmosphere, ocean, sea_ice=nothing;
 
     io_interface = sea_ice_ocean_interface(sea_ice, ocean)
 
-    ai_interface = atmosphere_sea_ice_interface(sea_ice,
+    ai_interface = atmosphere_sea_ice_interface(atmosphere,
+                                                sea_ice,
                                                 radiation,
                                                 atmosphere_sea_ice_flux_formulation,
                                                 atmosphere_sea_ice_interface_temperature,

src/OceanSeaIceModels/InterfaceComputations/interpolate_atmospheric_state.jl

@@ -63,13 +63,14 @@ function interpolate_atmosphere_state!(interfaces, atmosphere::PrescribedAtmosph
 
     # Assumption, should be generalized
     ua = atmosphere.velocities.u
-    time_interpolator = TimeInterpolator(ua, clock.time)
 
     launch!(arch, grid, kernel_parameters,
             _interpolate_primary_atmospheric_state!,
             atmosphere_data,
             space_fractional_indices,
-            time_interpolator,
+            clock.time,
+            ua.times,
+            ua.time_indexing,
             exchange_grid,
             atmosphere_velocities,
             atmosphere_tracers,
@@ -119,7 +120,9 @@ end
 
 @kernel function _interpolate_primary_atmospheric_state!(surface_atmos_state,
                                                          space_fractional_indices,
-                                                         time_interpolator,
+                                                         time,
+                                                         atmos_times,
+                                                         time_indexing,
                                                          exchange_grid,
                                                          atmos_velocities,
                                                          atmos_tracers,
@@ -137,7 +140,7 @@ end
     fj = get_fractional_index(i, j, jj)
 
     x_itp = FractionalIndices(fi, fj, nothing)
-    t_itp = time_interpolator
+    t_itp = TimeInterpolator(time_indexing, atmos_times, time)
     atmos_args = (x_itp, t_itp, atmos_backend, atmos_time_indexing)
 
     uₐ = interp_atmos_time_series(atmos_velocities.u, atmos_args...)