fill it in

Author: simone-silvestri

experiments/veros_forced_simulation.jl

@@ -1,6 +1,7 @@
 using ClimaOcean
 using PythonCall
 using Oceananigans
+using Printf
 
 VerosModule = Base.get_extension(ClimaOcean, :ClimaOceanPythonCallExt)
 VerosModule.remove_outputs(:global_4deg)
@@ -11,4 +12,29 @@ VerosModule.veros_settings_set!(ocean, "dt_tracer", 1800.0)
 atmos = JRA55PrescribedAtmosphere(; backend = JRA55NetCDFBackend(10))
 radiation = Radiation()
 coupled_model = OceanSeaIceModel(ocean, nothing; atmosphere=atmos, radiation)
+simulation = Simulation(coupled_model; Δt = 1800, stop_iteration = 100)
 
+wall_time = Ref(time_ns())
+
+function progress(sim)
+    ocean   = sim.model.ocean
+    umax = maximum(PyArray(ocean.setup.state.variables.u))
+    vmax = maximum(PyArray(ocean.setup.state.variables.v))
+    wmax = maximum(PyArray(ocean.setup.state.variables.w))
+
+    step_time = 1e-9 * (time_ns() - wall_time[])
+
+    msg1 = @sprintf("time: %s, iteration: %d, Δt: %s, ", prettytime(sim), iteration(sim), prettytime(sim.Δt))
+    msg5 = @sprintf("maximum(u): (%.2f, %.2f, %.2f) m/s, ", umax, vmax, wmax)
+    msg6 = @sprintf("wall time: %s \n", prettytime(step_time))
+
+    @info msg1 * msg5 * msg6
+
+     wall_time[] = time_ns()
+
+     return nothing
+end
+
+add_callback!(simulation, progress, IterationInterval(10))
+
+run!(simulation)

ext/ClimaOceanPythonCallExt/ClimaOceanPythonCallExt.jl

@@ -10,5 +10,6 @@ using Dates: DateTime
 
 include("clima_ocean_copernicus.jl")
 include("clima_ocean_veros.jl")
+include("veros_state_exchanger.jl")
 
 end # module ClimaOceanPythonCallExt

ext/ClimaOceanPythonCallExt/clima_ocean_veros.jl

@@ -3,9 +3,12 @@ using CondaPkg
 using ClimaOcean.OceanSeaIceModels: reference_density, heat_capacity, SeaIceSimulation
 
 import Oceananigans.Fields: set!
-import Oceananigans.TimeSteppers: time_step!
+import Oceananigans.TimeSteppers: time_step!, initialize!
 
-import ClimaOcean.OceanSeaIceModels: OceanSeaIceModel
+import ClimaOcean.OceanSeaIceModels: OceanSeaIceModel, default_nan_checker
+import Oceananigans.Architectures: architecture
+
+import Base: eltype
 
 """
     install_veros()
@@ -24,7 +27,12 @@ struct VerosOceanSimulation{S}
     setup :: S
 end
 
-time_step!(sim::VerosOceanSimulation, Δt) = sim.setup.step()
+default_nan_checker(model::OceanSeaIceModel{<:Any, <:Any, <:VerosOceanSimulation}) = nothing
+
+initialize!(::ClimaOceanPythonCallExt.VerosOceanSimulation{Py}) = nothing
+time_step!(ocean::VerosOceanSimulation, Δt) = ocean.setup.step(ocean.setup.state)
+architecture(model::OceanSeaIceModel{<:Any, <:Any, <:VerosOceanSimulation}) = CPU()
+eltype(model::OceanSeaIceModel{<:Any, <:Any, <:VerosOceanSimulation}) = Float64
 
 function remove_outputs(setup::Symbol)
     rm("$(setup).averages.nc", force=true)
@@ -34,15 +42,31 @@ function remove_outputs(setup::Symbol)
     return nothing
 end
 
-const Field2D = Field{<:Any, <:Any, <:Nothing}
+const CCField2D = Field{<:Center, <:Center, <:Nothing}
+const FCField2D = Field{<:Face, <:Center,   <:Nothing}
+const CFField2D = Field{<:Center, <:Face,   <:Nothing}
 
-function set!(field::Field2D, pyarray::Py, k=pyconvert(Int, pyarray.shape[2]))
+function set!(field::CCField2D, pyarray::Py, k=pyconvert(Int, pyarray.shape[2]))
     array = PyArray(pyarray)
     Nx, Ny, Nz = size(array)
     set!(field, view(array, 3:Nx-2, 3:Ny-2, k, 1))
     return field
 end
 
+function set!(field::FCField2D, pyarray::Py, k=pyconvert(Int, pyarray.shape[2]))
+    array = PyArray(pyarray)
+    Nx, Ny, Nz = size(array)
+    set!(field, view(array, 3:Nx-2, 3:Ny-2, k, 1))
+    return field
+end
+
+function set!(field::CFField2D, pyarray::Py, k=pyconvert(Int, pyarray.shape[2]))
+    array = PyArray(pyarray)
+    Nx, Ny, Nz = size(array)
+    set!(field, view(array, 3:Nx-2, 2:Ny-2, k, 1))
+    return field
+end
+
 function veros_ocean_simulation(setup, setup_name)
     setups = pyimport("veros.setups." * setup)
     setup  = @eval $setups.$setup_name()
@@ -53,13 +77,13 @@ function veros_ocean_simulation(setup, setup_name)
     return VerosOceanSimulation(setup) 
 end
 
-function surface_grid(setup::VerosOceanSimulation)
+function surface_grid(ocean::VerosOceanSimulation)
 
-    xf = Array(PyArray(setup.state.variables.xu))
-    yf = Array(PyArray(setup.state.variables.yu))
+    xf = Array(PyArray(ocean.setup.state.variables.xu))
+    yf = Array(PyArray(ocean.setup.state.variables.yu))
 
-    xc = Array(PyArray(setup.state.variables.xt))
-    yc = Array(PyArray(setup.state.variables.yt))
+    xc = Array(PyArray(ocean.setup.state.variables.xt))
+    yc = Array(PyArray(ocean.setup.state.variables.yt))
 
     xf = xf[2:end-2]
     yf = yf[2:end-2]
@@ -79,7 +103,7 @@ function surface_grid(setup::VerosOceanSimulation)
     Nx = length(xc) 
     Ny = length(yc) 
 
-    return LatitudeLongitudeGrid(size=(Nx, Ny), longitude=xf, latitude=yf, topology=(TX, Bounded, Flat))
+    return LatitudeLongitudeGrid(size=(Nx, Ny), longitude=xf, latitude=yf, topology=(TX, Bounded, Flat), halo=(2, 2))
 end
 
 function veros_set!(ocean::VerosOceanSimulation, v, x)

ext/ClimaOceanPythonCallExt/veros_state_exchanger.jl

@@ -1,3 +1,10 @@
+using Oceananigans.Models: initialization_update_state!
+
+using ClimaOcean.OceanSeaIceModels.InterfaceComputations: ExchangeAtmosphereState, 
+                                                          atmosphere_exchanger, 
+                                                          SimilarityTheoryFluxes, 
+                                                          Radiation
+
 import ClimaOcean.OceanSeaIceModels.InterfaceComputations: 
     state_exchanger, 
     sea_ice_ocean_interface, 
@@ -8,7 +15,6 @@ import ClimaOcean.OceanSeaIceModels.InterfaceComputations:
     get_radiative_forcing,
     fill_up_net_fluxes!
 
-
 mutable struct VerosStateExchanger{G, OST, AST, AEX}
     exchange_grid :: G
     exchange_ocean_state :: OST
@@ -33,12 +39,12 @@ function state_exchanger(ocean::VerosOceanSimulation, atmosphere)
     exchange_ocean_state = ExchangeOceanState(exchange_grid)
     exchange_atmosphere_state = ExchangeAtmosphereState(exchange_grid)
 
-    atmosphere_exchanger = AtmosphereExchanger(atmosphere, exchange_grid)
+    atmos_exchanger = atmosphere_exchanger(atmosphere, exchange_grid)
 
     return VerosStateExchanger(exchange_grid,
                                exchange_ocean_state,
                                exchange_atmosphere_state,
-                               atmosphere_exchanger)
+                               atmos_exchanger)
 end
 
 atmosphere_ocean_interface(ocean::VerosOceanSimulation, args...) = 
@@ -77,8 +83,19 @@ end
 @inline get_radiative_forcing(ocean::VerosOceanSimulation) = nothing
 
 function fill_up_net_fluxes!(ocean::VerosOceanSimulation, net_ocean_fluxes)
-    veros_set!(ocean, "surface_taux", net_ocean_fluxes.u)
-    veros_set!(ocean, "surface_tauy", net_ocean_fluxes.v)
-    
+    nx = pyconvert(Int, ocean.state.settings.nx)
+    ny = pyconvert(Int, ocean.state.settings.ny)
+    t1 = parent(net_ocean_fluxes.u)[:, 1:44, 1]
+    t2 = parent(net_ocean_fluxes.v)[:, 1:44, 1]
+
+    ta = zeros(size(t1)..., 12)
+    tb = zeros(size(t2)..., 12)
+    for t in 1:12
+        ta[:, :, t] .= t1
+        tb[:, :, t] .= t2
+    end
+
+    veros_set!(ocean, "taux", ta)
+    veros_set!(ocean, "tauy", tb)
     return nothing
 end

src/OceanSeaIceModels/InterfaceComputations/assemble_net_fluxes.jl

@@ -36,7 +36,7 @@ fill_up_net_fluxes!(ocean, net_ocean_fluxes) = nothing
 
 function compute_net_ocean_fluxes!(coupled_model)
     sea_ice = coupled_model.sea_ice
-    grid = coupled_model.exchanger.exchange_grid
+    grid = coupled_model.interfaces.exchanger.exchange_grid
     arch = architecture(grid)
     clock = coupled_model.clock
 
@@ -60,13 +60,13 @@ function compute_net_ocean_fluxes!(coupled_model)
     freshwater_flux = atmosphere_fields.Mp.data
 
     ice_concentration = sea_ice_concentration(sea_ice)
-    ocean_salinity = ocean.model.tracers.S
+    ocean_salinity = get_ocean_state(coupled_model.ocean, coupled_model.interfaces.exchanger).S
     atmos_ocean_properties = coupled_model.interfaces.atmosphere_ocean_interface.properties
     ocean_properties = coupled_model.interfaces.ocean_properties
     kernel_parameters = interface_kernel_parameters(grid)
 
     ocean_surface_temperature = coupled_model.interfaces.atmosphere_ocean_interface.temperature
-    penetrating_radiation = get_radiative_forcing(ocean)
+    penetrating_radiation = get_radiative_forcing(coupled_model.ocean)
 
     launch!(arch, grid, kernel_parameters,
             _assemble_net_ocean_fluxes!,
@@ -84,7 +84,7 @@ function compute_net_ocean_fluxes!(coupled_model)
             atmos_ocean_properties,
             ocean_properties)
 
-    fill_up_net_fluxes!(ocean, net_ocean_fluxes)
+    fill_up_net_fluxes!(coupled_model.ocean, net_ocean_fluxes)
 
     return nothing
 end

src/OceanSeaIceModels/InterfaceComputations/interpolate_atmospheric_state.jl

@@ -26,7 +26,7 @@ function interpolate_atmosphere_state!(interfaces, atmosphere::PrescribedAtmosph
     atmosphere_grid = atmosphere.grid
 
     # Basic model properties
-    grid = ocean.model.grid
+    grid = coupled_model.interfaces.exchanger.exchange_grid
     arch = architecture(grid)
     clock = coupled_model.clock
 
@@ -125,8 +125,8 @@ function interpolate_atmosphere_state!(interfaces, atmosphere::PrescribedAtmosph
     #
     # TODO: find a better design for this that doesn't have redundant
     # arrays for the barotropic potential
-    u_potential = forcing_barotropic_potential(ocean.model.forcing.u)
-    v_potential = forcing_barotropic_potential(ocean.model.forcing.v)
+    u_potential = forcing_barotropic_potential(ocean)
+    v_potential = forcing_barotropic_potential(ocean)
     ρₒ = coupled_model.interfaces.ocean_properties.reference_density
 
     if !isnothing(u_potential)

src/OceanSeaIceModels/ocean_sea_ice_model.jl

@@ -45,7 +45,7 @@ function Base.show(io::IO, cm::OSIM)
     end
 
     print(io, summary(cm), "\n")
-    print(io, "├── ocean: ", summary(cm.ocean.model), "\n")
+    print(io, "├── ocean: ", summary(cm.ocean), "\n")
     print(io, "├── atmosphere: ", summary(cm.atmosphere), "\n")
     print(io, "├── sea_ice: ", sea_ice_summary, "\n")
     print(io, "└── interfaces: ", summary(cm.interfaces))