diff --git a/Project.toml b/Project.toml index ee97cb825..3ac4eb0f3 100644 --- a/Project.toml +++ b/Project.toml @@ -1,8 +1,8 @@ name = "ClimaOcean" uuid = "0376089a-ecfe-4b0e-a64f-9c555d74d754" license = "MIT" -authors = ["Climate Modeling Alliance and contributors"] version = "0.8.7" +authors = ["Climate Modeling Alliance and contributors"] [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" @@ -33,6 +33,8 @@ Thermodynamics = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c" ZipFile = "a5390f91-8eb1-5f08-bee0-b1d1ffed6cea" [weakdeps] +PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" +CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab" CopernicusMarine = "cd43e856-93a3-40c8-bc9e-6146cdce14fa" Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" SpeedyWeather = "9e226e20-d153-4fed-8a5b-493def4f21a9" @@ -40,6 +42,7 @@ XESMF = "2e0b0046-e7a1-486f-88de-807ee8ffabe5" [extensions] ClimaOceanCopernicusMarineExt = "CopernicusMarine" +ClimaOceanPythonCallExt = ["PythonCall", "CondaPkg"] ClimaOceanReactantExt = "Reactant" ClimaOceanSpeedyWeatherExt = ["SpeedyWeather", "XESMF"] @@ -48,6 +51,7 @@ Adapt = "4" CFTime = "0.1, 0.2" CUDA = "4, 5" ClimaSeaIce = "0.3.9" +CondaPkg = "0.2.33" CopernicusMarine = "0.1.1" CubicSplines = "0.2" DataDeps = "0.7" @@ -63,6 +67,7 @@ NCDatasets = "0.12, 0.13, 0.14" Oceananigans = "0.101" OffsetArrays = "1.14" PrecompileTools = "1" +PythonCall = "0.9.28" Reactant = "0.2.45" Scratch = "1" SeawaterPolynomials = "0.3.5" @@ -81,4 +86,4 @@ MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Coverage", "Test", "MPIPreferences", "CUDA_Runtime_jll", "Reactant", "CopernicusMarine", "XESMF", "SpeedyWeather"] +test = ["Coverage", "Test", "MPIPreferences", "CUDA_Runtime_jll", "Reactant", "CopernicusMarine", "XESMF", "SpeedyWeather" "PythonCall", "CondaPkg"] diff --git a/docs/Project.toml b/docs/Project.toml index 20652b262..730d84707 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -10,6 +10,7 @@ Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306" NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab" Oceananigans = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09" SeawaterPolynomials = "d496a93d-167e-4197-9f49-d3af4ff8fe40" +PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" SpeedyWeather = "9e226e20-d153-4fed-8a5b-493def4f21a9" XESMF = "2e0b0046-e7a1-486f-88de-807ee8ffabe5" diff --git a/docs/make.jl b/docs/make.jl index 162ad232e..46c35624a 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -18,10 +18,11 @@ const EXAMPLES_DIR = joinpath(@__DIR__, "..", "examples") const OUTPUT_DIR = joinpath(@__DIR__, "src/literated") to_be_literated = [ - "single_column_os_papa_simulation.jl", - "one_degree_simulation.jl", - "near_global_ocean_simulation.jl", - "global_climate_simulation.jl", + # "single_column_os_papa_simulation.jl", + # "one_degree_simulation.jl", + # "near_global_ocean_simulation.jl", + "python_ocean_forced_simulation.jl", + # "global_climate_simulation.jl", ] for file in to_be_literated @@ -45,10 +46,11 @@ pages = [ "Home" => "index.md", "Examples" => [ - "Single-column ocean simulation" => "literated/single_column_os_papa_simulation.md", - "One-degree ocean--sea ice simulation" => "literated/one_degree_simulation.md", - "Near-global ocean simulation" => "literated/near_global_ocean_simulation.md", - "Global climate simulation" => "literated/global_climate_simulation.md", + # "Single-column ocean simulation" => "literated/single_column_os_papa_simulation.md", + # "One-degree ocean--sea ice simulation" => "literated/one_degree_simulation.md", + # "Near-global ocean simulation" => "literated/near_global_ocean_simulation.md", + "Python ocean forced simulation" => "literated/python_ocean_forced_simulation.md",], + # "Global climate simulation" => "literated/global_climate_simulation.md", ], "Vertical grids" => "vertical_grids.md", diff --git a/examples/python_ocean_forced_simulation.jl b/examples/python_ocean_forced_simulation.jl new file mode 100644 index 000000000..f62be558d --- /dev/null +++ b/examples/python_ocean_forced_simulation.jl @@ -0,0 +1,125 @@ +# # A Python Ocean Simulation at 4ᵒ Resolution Forced by JRA55 Reanalysis and Initialized from ECCO +# +# This example showcases the use of ClimaOcean's PythonCall extension to run a +# near-global ocean simulation at 4-degree resolution using the Veros ocean model. +# The ocean is forced by the JRA55 reanalysis data and initialized from the ECCO +# state estimate. +# +# For this example, we need Oceananigans, ClimaOcean, Dates, CUDA, and +# CairoMakie to visualize the simulation. + +using ClimaOcean +using PythonCall +using Oceananigans +using CairoMakie +using Printf + +# We import the Veros 4 degree ocean simulation setup, which consists of a near-global ocean +# with a uniform resolution of 4 degrees in both latitude and longitude and a latitude range spanning +# from 80S to 80N. The setup is defined in the `veros.setups.global_4deg` module. + +# Before importing the setup, we need to ensure that the Veros module is installed and loaded +# and that every output is removed to avoid conflicts. + +VerosModule = Base.get_extension(ClimaOcean, :ClimaOceanPythonCallExt) + +VerosModule.install_veros() +VerosModule.remove_outputs(:global_4deg) + +# Actually loading and instantiating the Veros setup in the variable `ocean`. +# This setup uses by default a different time-step for tracers and momentum, +# so we set it to the same value (1800 seconds) for both. + +ocean = VerosModule.VerosOceanSimulation("global_4deg", :GlobalFourDegreeSetup) + +set!(ocean, "dt_tracer", 1800.0; path=:settings) +set!(ocean, "dt_mom", 1800.0; path=:settings) + +# We force the 4-degree setup with a prescribed atmosphere based on the JRA-55 reanalysis data. +# This includes 2-meter wind velocity, temperature, humidity, downwelling longwave and shortwave +# radiation, as well as freshwater fluxes. + +atmos = JRA55PrescribedAtmosphere(; backend = JRA55NetCDFBackend(10)) + + +# The coupled ocean--atmosphere model. +# We use the default radiation model and we do not couple an ice model for simplicity. + +radiation = Radiation() +coupled_model = OceanSeaIceModel(ocean, nothing; atmosphere=atmos, radiation) +simulation = Simulation(coupled_model; Δt = 1800, stop_time = 60days) + +# We set up a progress callback that will print the current time, iteration, and maximum velocities +# at every 5 iterations. It also collects the surface velocity fields and the net fluxes +# into the arrays `s`, `tx`, and `ty` for later visualization. + +wall_time = Ref(time_ns()) + +s = [] +tx = [] +ty = [] + +us = coupled_model.interfaces.exchanger.exchange_ocean_state.u +vs = coupled_model.interfaces.exchanger.exchange_ocean_state.v + +stmp = Field(sqrt(us^2 + vs^2)) + +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() + + compute!(stmp) + push!(s, deepcopy(interior(stmp, :, :, 1))) + push!(tx, deepcopy(interior(coupled_model.interfaces.net_fluxes.ocean_surface.u, :, :, 1) .* 1020)) + push!(ty, deepcopy(interior(coupled_model.interfaces.net_fluxes.ocean_surface.v, :, :, 1) .* 1020)) + + return nothing +end + +add_callback!(simulation, progress, IterationInterval(5)) + +# Let's run the simulation! + +run!(simulation) + +# We can now visualize the surface speed and wind stress at the ocean surface +# over the course of the simulation. + +iter = Observable(1) +si = @lift(s[$iter]) +txi = @lift(tx[$iter]) +tyi = @lift(ty[$iter]) +Nt = length(s) + +fig = Figure(resolution = (1200, 300)) +ax1 = Axis(fig[1, 1]; title = "Surface speed (m/s)", xlabel = "Longitude", ylabel = "Latitude") +ax2 = Axis(fig[1, 2]; title = "Zonal wind stress (N/m²)", xlabel = "Longitude") +ax3 = Axis(fig[1, 3]; title = "Meridional wind stress (N/m²)", xlabel = "Longitude") + +grid = coupled_model.interfaces.exchanger.exchange_grid + +λ = λnodes(grid, Center()) +φ = φnodes(grid, Center()) + +heatmap!(ax1, λ, φ, si, colormap = :ice, colorrange = (0, 0.15)) +heatmap!(ax2, λ, φ, txi, colormap = :bwr, colorrange = (-0.2, 0.2)) +heatmap!(ax3, λ, φ, tyi, colormap = :bwr, colorrange = (-0.2, 0.2)) + +CairoMakie.record(fig, "veros_ocean_surface.mp4", 1:Nt, framerate = 8) do nn + iter[] = nn +end +nothing #hide + +# ![](veros_ocean_surface.mp4) diff --git a/ext/ClimaOceanPythonCallExt/ClimaOceanPythonCallExt.jl b/ext/ClimaOceanPythonCallExt/ClimaOceanPythonCallExt.jl new file mode 100644 index 000000000..c241d1ba8 --- /dev/null +++ b/ext/ClimaOceanPythonCallExt/ClimaOceanPythonCallExt.jl @@ -0,0 +1,14 @@ +module ClimaOceanVerosExt + +using ClimaOcean +using CondaPkg +using PythonCall +using Oceananigans +using Oceananigans.DistributedComputations: @root + +using Dates: DateTime + +include("VerosOceanSimulations/veros_ocean_simulation.jl") +include("veros_state_exchanger.jl") + +end # module ClimaOceanVerosExt diff --git a/ext/ClimaOceanPythonCallExt/VerosOceanSimulations/VerosOceanSimulations.jl b/ext/ClimaOceanPythonCallExt/VerosOceanSimulations/VerosOceanSimulations.jl new file mode 100644 index 000000000..2aafc0057 --- /dev/null +++ b/ext/ClimaOceanPythonCallExt/VerosOceanSimulations/VerosOceanSimulations.jl @@ -0,0 +1,32 @@ +module VerosOceanSimulations + +using CondaPkg + +using Oceananigans.Grids: topology +using ClimaOcean.OceanSeaIceModels: reference_density, heat_capacity, SeaIceSimulation + +import Oceananigans.Fields: set! +import Oceananigans.TimeSteppers: time_step!, initialize! + +import ClimaOcean.OceanSeaIceModels: OceanSeaIceModel, default_nan_checker +import Oceananigans.Architectures: architecture + +import Base: eltype + +""" + install_veros() + +Install the Veros ocean model Marine CLI using CondaPkg. +Returns a NamedTuple containing package information if successful. +""" +function install_veros() + CondaPkg.add_pip("veros") + cli = CondaPkg.which("veros") + @info "... the veros CLI has been installed at $(cli)." + return cli +end + +include("veros_ocean_simulation.jl") +include("veros_state_exchanger.jl") + +end # module VerosOceanSimulations diff --git a/ext/ClimaOceanPythonCallExt/VerosOceanSimulations/flexible_veros_simulation.py b/ext/ClimaOceanPythonCallExt/VerosOceanSimulations/flexible_veros_simulation.py new file mode 100644 index 000000000..aa9b99037 --- /dev/null +++ b/ext/ClimaOceanPythonCallExt/VerosOceanSimulations/flexible_veros_simulation.py @@ -0,0 +1,349 @@ +#!/usr/bin/env python +import os +import h5netcdf +import scipy.ndimage + +from veros import veros_routine, veros_kernel, KernelOutput, VerosSetup, runtime_settings as rs, runtime_state as rst +from veros.variables import Variable, allocate +from veros.core.utilities import enforce_boundaries +from veros.core.operators import numpy as npx, update, at +import veros.tools +import veros.time + +BASE_PATH = os.path.dirname(os.path.realpath(__file__)) +DATA_FILES = veros.tools.get_assets("global_flexible", os.path.join(BASE_PATH, "assets.json")) + +class GlobalFlexibleResolutionSetup(VerosSetup): + """ + Global model with flexible resolution. + """ + + # global settings + min_depth = 10.0 + max_depth = 5400.0 + equatorial_grid_spacing_factor = 0.5 + polar_grid_spacing_factor = None + + @veros_routine + def set_parameter(self, state): + settings = state.settings + + settings.identifier = "global_flexible" + settings.description = "Global model with flexible resolution" + + settings.nx = 360 + settings.ny = 160 + settings.nz = 60 + settings.dt_mom = settings.dt_tracer = 900 + settings.runlen = 86400 * 10 + + settings.x_origin = 90.0 + settings.y_origin = -80.0 + + settings.coord_degree = True + settings.enable_cyclic_x = True + + # friction + settings.enable_hor_friction = True + settings.A_h = 5e4 + settings.enable_hor_friction_cos_scaling = True + settings.hor_friction_cosPower = 1 + settings.enable_tempsalt_sources = True + settings.enable_implicit_vert_friction = True + + settings.eq_of_state_type = 5 + + # isoneutral + settings.enable_neutral_diffusion = True + settings.K_iso_0 = 1000.0 + settings.K_iso_steep = 50.0 + settings.iso_dslope = 0.005 + settings.iso_slopec = 0.005 + settings.enable_skew_diffusion = True + + # tke + settings.enable_tke = True + settings.c_k = 0.1 + settings.c_eps = 0.7 + settings.alpha_tke = 30.0 + settings.mxl_min = 1e-8 + settings.tke_mxl_choice = 2 + settings.kappaM_min = 2e-4 + settings.kappaH_min = 2e-5 + settings.enable_kappaH_profile = True + settings.enable_tke_superbee_advection = True + + # eke + settings.enable_eke = True + settings.eke_k_max = 1e4 + settings.eke_c_k = 0.4 + settings.eke_c_eps = 0.5 + settings.eke_cross = 2.0 + settings.eke_crhin = 1.0 + settings.eke_lmin = 100.0 + settings.enable_eke_superbee_advection = True + settings.enable_eke_isopycnal_diffusion = True + + # idemix + settings.enable_idemix = False + settings.enable_eke_diss_surfbot = True + settings.eke_diss_surfbot_frac = 0.2 + settings.enable_idemix_superbee_advection = True + settings.enable_idemix_hor_diffusion = True + + def _get_data(self, var, idx=None): + if idx is None: + idx = Ellipsis + else: + idx = idx[::-1] + + kwargs = {} + if rst.proc_num > 1: + kwargs.update( + driver="mpio", + comm=rs.mpi_comm, + ) + + with h5netcdf.File(DATA_FILES["forcing"], "r", **kwargs) as forcing_file: + var_obj = forcing_file.variables[var] + return npx.array(var_obj[idx]).T + + @veros_routine(dist_safe=False, local_variables=["dxt", "dyt", "dzt"]) + def set_grid(self, state): + vs = state.variables + settings = state.settings + + if settings.ny % 2: + raise ValueError("ny has to be an even number of grid cells") + + vs.dxt = update(vs.dxt, at[...], 360.0 / settings.nx) + + if self.equatorial_grid_spacing_factor is not None: + eq_spacing = self.equatorial_grid_spacing_factor * 160.0 / settings.ny + else: + eq_spacing = None + + if self.polar_grid_spacing_factor is not None: + polar_spacing = self.polar_grid_spacing_factor * 160.0 / settings.ny + else: + polar_spacing = None + + vs.dyt = update( + vs.dyt, + at[2:-2], + veros.tools.get_vinokur_grid_steps( + settings.ny, 160.0, eq_spacing, upper_stepsize=polar_spacing, two_sided_grid=True + ), + ) + vs.dzt = veros.tools.get_vinokur_grid_steps(settings.nz, self.max_depth, self.min_depth, refine_towards="lower") + + @veros_routine + def set_coriolis(self, state): + vs = state.variables + settings = state.settings + vs.coriolis_t = update( + vs.coriolis_t, at[...], 2 * settings.omega * npx.sin(vs.yt[npx.newaxis, :] / 180.0 * settings.pi) + ) + + def _shift_longitude_array(self, vs, lon, arr): + wrap_i = npx.where((lon[:-1] < vs.xt.min()) & (lon[1:] >= vs.xt.min()))[0][0] + new_lon = npx.concatenate((lon[wrap_i:-1], lon[:wrap_i] + 360.0)) + new_arr = npx.concatenate((arr[wrap_i:-1, ...], arr[:wrap_i, ...])) + return new_lon, new_arr + + @veros_routine(dist_safe=False, local_variables=["kbot", "xt", "yt", "zt"]) + def set_topography(self, state): + vs = state.variables + settings = state.settings + + with h5netcdf.File(DATA_FILES["topography"], "r") as topography_file: + topo_x, topo_y, topo_z = (npx.array(topography_file.variables[k], dtype="float").T for k in ("x", "y", "z")) + + topo_z = npx.minimum(topo_z, 0.0) + + # smooth topography to match grid resolution + gaussian_sigma = (0.5 * len(topo_x) / settings.nx, 0.5 * len(topo_y) / settings.ny) + topo_z_smoothed = scipy.ndimage.gaussian_filter(topo_z, sigma=gaussian_sigma) + topo_z_smoothed = npx.where(topo_z >= -1, 0, topo_z_smoothed) + + topo_x_shifted, topo_z_shifted = self._shift_longitude_array(vs, topo_x, topo_z_smoothed) + coords = (vs.xt[2:-2], vs.yt[2:-2]) + z_interp = allocate(state.dimensions, ("xt", "yt"), local=False) + z_interp = update( + z_interp, + at[2:-2, 2:-2], + veros.tools.interpolate((topo_x_shifted, topo_y), topo_z_shifted, coords, kind="nearest", fill=False), + ) + + depth_levels = 1 + npx.argmin(npx.abs(z_interp[:, :, npx.newaxis] - vs.zt[npx.newaxis, npx.newaxis, :]), axis=2) + vs.kbot = update(vs.kbot, at[2:-2, 2:-2], npx.where(z_interp < 0.0, depth_levels, 0)[2:-2, 2:-2]) + vs.kbot = npx.where(vs.kbot < settings.nz, vs.kbot, 0) + vs.kbot = enforce_boundaries(vs.kbot, settings.enable_cyclic_x, local=True) + + # remove marginal seas + # (dilate to close 1-cell passages, fill holes, undo dilation) + marginal = scipy.ndimage.binary_erosion( + scipy.ndimage.binary_fill_holes(scipy.ndimage.binary_dilation(vs.kbot == 0)) + ) + + vs.kbot = npx.where(marginal, 0, vs.kbot) + + @veros_routine + def set_initial_conditions(self, state): + vs = state.variables + settings = state.settings + + rpart_shortwave = 0.58 + efold1_shortwave = 0.35 + efold2_shortwave = 23.0 + + t_grid = (vs.xt[2:-2], vs.yt[2:-2], vs.zt) + xt_forc, yt_forc, zt_forc = (self._get_data(k) for k in ("xt", "yt", "zt")) + zt_forc = zt_forc[::-1] + + # coordinates must be monotonous for this to work + assert npx.diff(xt_forc).all() > 0 + assert npx.diff(yt_forc).all() > 0 + + # determine slice to read from forcing file + data_subset = ( + slice( + max(0, int(npx.argmax(xt_forc >= vs.xt.min())) - 1), + len(xt_forc) - max(0, int(npx.argmax(xt_forc[::-1] <= vs.xt.max())) - 1), + ), + slice( + max(0, int(npx.argmax(yt_forc >= vs.yt.min())) - 1), + len(yt_forc) - max(0, int(npx.argmax(yt_forc[::-1] <= vs.yt.max())) - 1), + ), + Ellipsis, + ) + + xt_forc = xt_forc[data_subset[0]] + yt_forc = yt_forc[data_subset[1]] + + # initial conditions + temp_raw = self._get_data("temperature", idx=data_subset)[..., ::-1] + temp_data = veros.tools.interpolate((xt_forc, yt_forc, zt_forc), temp_raw, t_grid) + vs.temp = update(vs.temp, at[2:-2, 2:-2, :, :], (temp_data * vs.maskT[2:-2, 2:-2, :])[..., npx.newaxis]) + + salt_raw = self._get_data("salinity", idx=data_subset)[..., ::-1] + salt_data = veros.tools.interpolate((xt_forc, yt_forc, zt_forc), salt_raw, t_grid) + vs.salt = update(vs.salt, at[2:-2, 2:-2, :, :], (salt_data * vs.maskT[2:-2, 2:-2, :])[..., npx.newaxis]) + + # wind stress on MIT grid + time_grid = (vs.xt[2:-2], vs.yt[2:-2], npx.arange(12)) + taux_raw = self._get_data("tau_x", idx=data_subset) + taux_data = veros.tools.interpolate((xt_forc, yt_forc, npx.arange(12)), taux_raw, time_grid) + vs.taux = update(vs.taux, at[2:-2, 2:-2, :], taux_data) + + tauy_raw = self._get_data("tau_y", idx=data_subset) + tauy_data = veros.tools.interpolate((xt_forc, yt_forc, npx.arange(12)), tauy_raw, time_grid) + vs.tauy = update(vs.tauy, at[2:-2, 2:-2, :], tauy_data) + + vs.taux = enforce_boundaries(vs.taux, settings.enable_cyclic_x) + vs.tauy = enforce_boundaries(vs.tauy, settings.enable_cyclic_x) + + # Qnet and dQ/dT and Qsol + qnet_raw = self._get_data("q_net", idx=data_subset) + qnet_data = veros.tools.interpolate((xt_forc, yt_forc, npx.arange(12)), qnet_raw, time_grid) + vs.qnet = update(vs.qnet, at[2:-2, 2:-2, :], -qnet_data * vs.maskT[2:-2, 2:-2, -1, npx.newaxis]) + + qnec_raw = self._get_data("dqdt", idx=data_subset) + qnec_data = veros.tools.interpolate((xt_forc, yt_forc, npx.arange(12)), qnec_raw, time_grid) + vs.qnec = update(vs.qnec, at[2:-2, 2:-2, :], qnec_data * vs.maskT[2:-2, 2:-2, -1, npx.newaxis]) + + qsol_raw = self._get_data("swf", idx=data_subset) + qsol_data = veros.tools.interpolate((xt_forc, yt_forc, npx.arange(12)), qsol_raw, time_grid) + vs.qsol = update(vs.qsol, at[2:-2, 2:-2, :], -qsol_data * vs.maskT[2:-2, 2:-2, -1, npx.newaxis]) + + if settings.enable_idemix: + tidal_energy_raw = self._get_data("tidal_energy", idx=data_subset) + tidal_energy_data = veros.tools.interpolate((xt_forc, yt_forc), tidal_energy_raw, t_grid[:-1]) + mask_x, mask_y = (i + 2 for i in npx.indices((vs.nx, vs.ny))) + mask_z = npx.maximum(0, vs.kbot[2:-2, 2:-2] - 1) + tidal_energy_data[:, :] *= vs.maskW[mask_x, mask_y, mask_z] / vs.rho_0 + vs.forc_iw_bottom[2:-2, 2:-2] = tidal_energy_data + + """ + Initialize penetration profile for solar radiation and store divergence in divpen + note that pen is set to 0.0 at the surface instead of 1.0 to compensate for the + shortwave part of the total surface flux + """ + swarg1 = vs.zw / efold1_shortwave + swarg2 = vs.zw / efold2_shortwave + pen = rpart_shortwave * npx.exp(swarg1) + (1.0 - rpart_shortwave) * npx.exp(swarg2) + pen = update(pen, at[-1], 0.0) + vs.divpen_shortwave = update(vs.divpen_shortwave, at[1:], (pen[1:] - pen[:-1]) / vs.dzt[1:]) + vs.divpen_shortwave = update(vs.divpen_shortwave, at[0], pen[0] / vs.dzt[0]) + + @veros_routine + def set_forcing(self, state): + vs = state.variables + vs.update(set_forcing_kernel(state)) + + @veros_routine + def set_diagnostics(self, state): + settings = state.settings + diagnostics = state.diagnostics + + diagnostics["cfl_monitor"].output_frequency = settings.dt_tracer * 100 + diagnostics["tracer_monitor"].output_frequency = settings.dt_tracer * 100 + diagnostics["snapshot"].output_frequency = 30 * 86400.0 + diagnostics["overturning"].output_frequency = 360 * 86400 + diagnostics["overturning"].sampling_frequency = 86400.0 + diagnostics["energy"].output_frequency = 360 * 86400 + diagnostics["energy"].sampling_frequency = 10 * settings.dt_tracer + diagnostics["averages"].output_frequency = 30 * 86400 + diagnostics["averages"].sampling_frequency = settings.dt_tracer + + average_vars = [ + "surface_taux", + "surface_tauy", + "forc_temp_surface", + "forc_salt_surface", + "psi", + "temp", + "salt", + "u", + "v", + "w", + "Nsqr", + "Hd", + "rho", + "kappaH", + ] + if settings.enable_skew_diffusion: + average_vars += ["B1_gm", "B2_gm"] + if settings.enable_TEM_friction: + average_vars += ["kappa_gm", "K_diss_gm"] + if settings.enable_tke: + average_vars += ["tke", "Prandtlnumber", "mxl", "tke_diss", "forc_tke_surface", "tke_surf_corr"] + if settings.enable_idemix: + average_vars += ["E_iw", "forc_iw_surface", "iw_diss", "c0", "v0"] + if settings.enable_eke: + average_vars += ["eke", "K_gm", "L_rossby", "L_rhines"] + diagnostics["averages"].output_variables = average_vars + + @veros_routine + def after_timestep(self, state): + pass + +@veros_kernel +def set_forcing_kernel(state): + vs = state.variables + settings = state.settings + + t_rest = 30.0 * 86400.0 + cp_0 = 3991.86795711963 # J/kg /K + + year_in_seconds = veros.time.convert_time(1.0, "years", "seconds") + (n1, f1), (n2, f2) = veros.tools.get_periodic_interval(vs.time, year_in_seconds, year_in_seconds / 12.0, 12) + + return KernelOutput( + surface_taux=vs.surface_taux, + surface_tauy=vs.surface_tauy, + temp_source=vs.temp_source, + forc_tke_surface=vs.forc_tke_surface, + forc_temp_surface=vs.forc_temp_surface, + forc_salt_surface=vs.forc_salt_surface, + ) \ No newline at end of file diff --git a/ext/ClimaOceanPythonCallExt/VerosOceanSimulations/veros_ocean_simulation.jl b/ext/ClimaOceanPythonCallExt/VerosOceanSimulations/veros_ocean_simulation.jl new file mode 100644 index 000000000..ee149976d --- /dev/null +++ b/ext/ClimaOceanPythonCallExt/VerosOceanSimulations/veros_ocean_simulation.jl @@ -0,0 +1,177 @@ + +struct VerosOceanSimulation{S} + setup :: S +end + +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) + rm("$(setup).energy.nc", force=true) + rm("$(setup).overturning.nc", force=true) + rm("$(setup).snapshot.nc", force=true) + return nothing +end + +const CCField2D = Field{<:Center, <:Center, <:Nothing} +const FCField2D = Field{<:Face, <:Center, <:Nothing} +const CFField2D = Field{<:Center, <:Face, <:Nothing} + +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) + TX, TY, _ = topology(field.grid) + i_indices = TX == Periodic ? UnitRange(3, Nx-2) : UnitRange(2, Nx-2) + set!(field, view(array, i_indices, 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 + +""" + VerosOceanSimulation(setup, setup_name::Symbol) + +Creates and initializes a preconfigured Veros ocean simulation using the +specified setup module and setup name. + +Arguments +========== +- `setup::AbstractString`: The name of the Veros setup module to import (e.g., `"global_4deg"`). +- `setup_name::Symbol`: The name of the setup class or function within the module to instantiate (e.g., `:GlobalFourDegreeSetup`). +""" +function VerosOceanSimulation(setup, setup_name::Symbol) + setups = pyimport("veros.setups." * setup) + setup = @eval $setups.$setup_name() + + # instantiate the setup + setup.setup() + + return VerosOceanSimulation(setup) +end + +""" + surface_grid(ocean::VerosOceanSimulation) + +Constructs a `LatitudeLongitudeGrid` representing the surface grid of the given `VerosOceanSimulation` object. +Notes: Veros always uses a LatitudeLongitudeGrid with 2 halos in both the latitude and longitude directions. +Both latitude and longitude can be either stretched or uniform, depending on the setup, and while the meridional +direction (latitude) is always Bounded, the zonal direction (longitude) can be either Periodic or Bounded. + +Arguments +========== +- `ocean::VerosOceanSimulation`: The ocean simulation object containing the grid state variables. +""" +function surface_grid(ocean::VerosOceanSimulation) + + xf = Array(PyArray(ocean.setup.state.variables.xu)) + yf = Array(PyArray(ocean.setup.state.variables.yu)) + + 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] + + xc = xc[3:end-2] + yc = yc[3:end-2] + + xf[1] = xf[2] - 2xc[1] + yf[1] = sign(yf[2]) * (yf[2] - 2yc[1]) + + TX = if xf[1] == 0 && xf[end] == 360 + Periodic + else + Bounded + end + + Nx = length(xc) + Ny = length(yc) + + return LatitudeLongitudeGrid(size=(Nx, Ny), longitude=xf, latitude=yf, topology=(TX, Bounded, Flat), halo=(2, 2)) +end + +""" + set!(ocean, v, x; path = :variable) + +Set the `v` variable in the `ocean` model to the value of `x`. +the path corresponds to the path inside the class where to locate the +variable `v` to set. It can be either `:variables` or `:settings`. +""" +function set!(ocean::VerosOceanSimulation, v, x; path = :variables) + setup = ocean.setup + if path == :variables + pyexec(""" + with setup.state.variables.unlock(): + setup.state.variables.__setattr__(y, t) + """, Main, (y=v, t=x, setup=setup)) + elseif path == :settings + pyexec(""" + with setup.state.settings.unlock(): + setup.state.settings.__setattr__(y, t) + """, Main, (y=v, t=x, setup=setup)) + else + error("path must be either :variable or :settings.") + end +end + +function OceanSeaIceModel(ocean::VerosOceanSimulation, sea_ice=nothing; + atmosphere = nothing, + radiation = Radiation(), + clock = Clock(time=0), + ocean_reference_density = 1020.0, + ocean_heat_capacity = 3998.0, + sea_ice_reference_density = reference_density(sea_ice), + sea_ice_heat_capacity = heat_capacity(sea_ice), + interfaces = nothing) + + if sea_ice isa SeaIceSimulation + if !isnothing(sea_ice.callbacks) + pop!(sea_ice.callbacks, :stop_time_exceeded, nothing) + pop!(sea_ice.callbacks, :stop_iteration_exceeded, nothing) + pop!(sea_ice.callbacks, :wall_time_limit_exceeded, nothing) + pop!(sea_ice.callbacks, :nan_checker, nothing) + end + end + + # Contains information about flux contributions: bulk formula, prescribed fluxes, etc. + if isnothing(interfaces) && !(isnothing(atmosphere) && isnothing(sea_ice)) + interfaces = ComponentInterfaces(atmosphere, ocean, sea_ice; + ocean_reference_density, + ocean_heat_capacity, + sea_ice_reference_density, + sea_ice_heat_capacity, + radiation) + end + + arch = CPU() + + ocean_sea_ice_model = OceanSeaIceModel(arch, + clock, + atmosphere, + sea_ice, + ocean, + interfaces) + + # Make sure the initial temperature of the ocean + # is not below freezing and above melting near the surface + initialization_update_state!(ocean_sea_ice_model) + + return ocean_sea_ice_model +end diff --git a/ext/ClimaOceanPythonCallExt/VerosOceanSimulations/veros_state_exchanger.jl b/ext/ClimaOceanPythonCallExt/VerosOceanSimulations/veros_state_exchanger.jl new file mode 100644 index 000000000..c8bef5412 --- /dev/null +++ b/ext/ClimaOceanPythonCallExt/VerosOceanSimulations/veros_state_exchanger.jl @@ -0,0 +1,104 @@ +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, + atmosphere_ocean_interface, + initialize!, + get_ocean_state, + ocean_surface_fluxes, + get_radiative_forcing, + fill_net_fluxes! + +mutable struct VerosStateExchanger{G, OST, AST, AEX} + exchange_grid :: G + exchange_ocean_state :: OST + exchange_atmosphere_state :: AST + atmosphere_exchanger :: AEX +end + +mutable struct ExchangeOceanState{FC, CF, CC} + u :: FC + v :: CF + T :: CC + S :: CC +end + +ExchangeOceanState(grid) = ExchangeOceanState(Field{Face, Center, Nothing}(grid), + Field{Center, Face, Nothing}(grid), + Field{Center, Center, Nothing}(grid), + Field{Center, Center, Nothing}(grid)) + +function state_exchanger(ocean::VerosOceanSimulation, atmosphere) + exchange_grid = surface_grid(ocean) + exchange_ocean_state = ExchangeOceanState(exchange_grid) + exchange_atmosphere_state = ExchangeAtmosphereState(exchange_grid) + + atmos_exchanger = atmosphere_exchanger(atmosphere, exchange_grid) + + return VerosStateExchanger(exchange_grid, + exchange_ocean_state, + exchange_atmosphere_state, + atmos_exchanger) +end + +atmosphere_ocean_interface(ocean::VerosOceanSimulation, args...) = + atmosphere_ocean_interface(surface_grid(ocean), args...) + +sea_ice_ocean_interface(ocean::VerosOceanSimulation, args...) = + sea_ice_ocean_interface(surface_grid(ocean), args...) + +initialize!(exchanger::VerosStateExchanger, atmosphere) = + initialize!(exchanger.atmosphere_exchanger, exchanger.exchange_grid, atmosphere) + +@inline function get_ocean_state(ocean::VerosOceanSimulation, exchanger::VerosStateExchanger) + u = exchanger.exchange_ocean_state.u + v = exchanger.exchange_ocean_state.v + T = exchanger.exchange_ocean_state.T + S = exchanger.exchange_ocean_state.S + + set!(u, ocean.setup.state.variables.u) + set!(v, ocean.setup.state.variables.v) + set!(T, ocean.setup.state.variables.temp) + set!(S, ocean.setup.state.variables.salt) + + return (; u, v, T, S) +end + +@inline function ocean_surface_fluxes(ocean::VerosOceanSimulation, ρₒ, cₒ) + grid = surface_grid(ocean) + u = Field{Face, Center, Nothing}(grid) + v = Field{Center, Face, Nothing}(grid) + T = Field{Center, Center, Nothing}(grid) + S = Field{Center, Center, Nothing}(grid) + Q = ρₒ * cₒ * T + + return (; u, v, T, S, Q) +end + +@inline get_radiative_forcing(ocean::VerosOceanSimulation) = nothing + +function fill_net_fluxes!(ocean::VerosOceanSimulation, net_ocean_fluxes) + nx = pyconvert(Int, ocean.setup.state.settings.nx) + 4 + ny = pyconvert(Int, ocean.setup.state.settings.ny) + 4 + + ρₒ = pyconvert(eltype(ocean), ocean.setup.state.settings.rho_0) + taux = view(parent(net_ocean_fluxes.u), 1:nx, 1:ny, 1) .* ρₒ + tauy = view(parent(net_ocean_fluxes.v), 1:nx, 1:ny, 1) .* ρₒ + + set!(ocean, "surface_taux", tx; path=:variables) + set!(ocean, "surface_tauy", ty; path=:variables) + + temp_flux = view(parent(net_ocean_fluxes.T), 1:nx, 1:ny, 1) + salt_flux = view(parent(net_ocean_fluxes.S), 1:nx, 1:ny, 1) + + set!(ocean, "forc_temp_surface", temp_flux; path=:variables) + set!(ocean, "forc_salt_surface", salt_flux; path=:variables) + + return nothing +end diff --git a/src/OceanSeaIceModels/InterfaceComputations/assemble_net_ocean_fluxes.jl b/src/OceanSeaIceModels/InterfaceComputations/assemble_net_ocean_fluxes.jl index 2292fbf1f..46fdf8f94 100644 --- a/src/OceanSeaIceModels/InterfaceComputations/assemble_net_ocean_fluxes.jl +++ b/src/OceanSeaIceModels/InterfaceComputations/assemble_net_ocean_fluxes.jl @@ -14,14 +14,19 @@ using ClimaOcean.OceanSeaIceModels: sea_ice_concentration, NoAtmosphereModel, No return zero(Iˢʷ) end -get_radiative_forcing(FT) = FT -function get_radiative_forcing(FT::MultipleForcings) +@inline get_radiative_forcing(ocean::OceananigansSimulation) = get_radiative_forcing(ocean.model.forcing.T) +@inline get_radiative_forcing(FT) = FT + +@inline function get_radiative_forcing(FT::MultipleForcings) for forcing in FT.forcings forcing isa TwoColorRadiation && return forcing end return nothing end +# No need to do this for an Oceananigans Simulation +fill_net_fluxes!(ocean, net_ocean_fluxes) = nothing + @inline τᶜᶜᶜ(i, j, k, grid, ρₒ⁻¹, ℵ, ρτᶜᶜᶜ) = @inbounds ρₒ⁻¹ * (1 - ℵ[i, j, k]) * ρτᶜᶜᶜ[i, j, k] ##### @@ -34,7 +39,7 @@ computed_sea_ice_ocean_fluxes(::Nothing) = nothing # A generic ocean flux assembler for a coupled model with both an atmosphere and sea ice function compute_net_ocean_fluxes!(coupled_model, ocean) sea_ice = coupled_model.sea_ice - grid = ocean.model.grid + grid = coupled_model.interfaces.exchanger.exchange_grid arch = architecture(grid) clock = coupled_model.clock @@ -58,13 +63,14 @@ function compute_net_ocean_fluxes!(coupled_model, ocean) freshwater_flux = atmosphere_fields.Mp.data ice_concentration = sea_ice_concentration(sea_ice) - ocean_salinity = ocean.model.tracers.S + ocean_state = get_ocean_state(coupled_model.ocean, coupled_model.interfaces.exchanger) + ocean_salinity = ocean_state.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.model.forcing.T) + penetrating_radiation = get_radiative_forcing(coupled_model.ocean) launch!(arch, grid, kernel_parameters, _assemble_net_ocean_fluxes!, @@ -82,6 +88,8 @@ function compute_net_ocean_fluxes!(coupled_model, ocean) atmos_ocean_properties, ocean_properties) + fill_net_fluxes!(coupled_model.ocean, net_ocean_fluxes) + return nothing end diff --git a/src/OceanSeaIceModels/InterfaceComputations/atmosphere_ocean_fluxes.jl b/src/OceanSeaIceModels/InterfaceComputations/atmosphere_ocean_fluxes.jl index 319ba5366..041e81a43 100644 --- a/src/OceanSeaIceModels/InterfaceComputations/atmosphere_ocean_fluxes.jl +++ b/src/OceanSeaIceModels/InterfaceComputations/atmosphere_ocean_fluxes.jl @@ -4,19 +4,21 @@ using ClimaOcean.OceanSeaIceModels.PrescribedAtmospheres: thermodynamics_paramet surface_layer_height, boundary_layer_height -function compute_atmosphere_ocean_fluxes!(coupled_model) - ocean = coupled_model.ocean - atmosphere = coupled_model.atmosphere - grid = ocean.model.grid - arch = architecture(grid) - clock = coupled_model.clock - - ocean_state = (u = ocean.model.velocities.u, +@inline get_ocean_state(ocean::OceananigansSimulation, coupled_model) = + (u = ocean.model.velocities.u, v = ocean.model.velocities.v, T = ocean.model.tracers.T, S = ocean.model.tracers.S) - atmosphere_fields = coupled_model.interfaces.exchanger.exchange_atmosphere_state +function compute_atmosphere_ocean_fluxes!(coupled_model) + ocean = coupled_model.ocean + atmosphere = coupled_model.atmosphere + exchanger = coupled_model.interfaces.exchanger + grid = exchanger.exchange_grid + arch = architecture(grid) + clock = coupled_model.clock + ocean_state = get_ocean_state(ocean, exchanger) + atmosphere_fields = exchanger.exchange_atmosphere_state # Simplify NamedTuple to reduce parameter space consumption. # See https://github.com/CliMA/ClimaOcean.jl/issues/116. diff --git a/src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl b/src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl index 215a9bcaa..68391b5fd 100644 --- a/src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl +++ b/src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl @@ -8,7 +8,8 @@ using ..OceanSeaIceModels: reference_density, sea_ice_thickness, downwelling_radiation, freshwater_flux, - SeaIceSimulation + SeaIceSimulation, + OceananigansSimulation using ..OceanSeaIceModels.PrescribedAtmospheres: PrescribedAtmosphere, @@ -17,7 +18,7 @@ using ..OceanSeaIceModels.PrescribedAtmospheres: using ClimaSeaIce: SeaIceModel using Oceananigans: HydrostaticFreeSurfaceModel, architecture -using Oceananigans.Grids: inactive_node, node, topology +using Oceananigans.Grids: inactive_node, node, topology, AbstractGrid using Oceananigans.BoundaryConditions: fill_halo_regions! using Oceananigans.Fields: ConstantField, interpolate, FractionalIndices using Oceananigans.Utils: launch!, Time, KernelParameters @@ -86,10 +87,9 @@ ExchangeAtmosphereState(grid) = ExchangeAtmosphereState(Field{Center, Center, No fractional_index_type(FT, Topo) = FT fractional_index_type(FT, ::Flat) = Nothing +state_exchanger(ocean::Simulation, ::Nothing) = nothing -StateExchanger(ocean::Simulation, ::Nothing) = nothing - -function StateExchanger(ocean::Simulation, atmosphere) +function state_exchanger(ocean::Simulation, atmosphere) # TODO: generalize this exchange_grid = ocean.model.grid exchange_atmosphere_state = ExchangeAtmosphereState(exchange_grid) @@ -116,12 +116,11 @@ function atmosphere_exchanger(atmosphere::PrescribedAtmosphere, exchange_grid, e end initialize!(exchanger::StateExchanger, ::Nothing) = nothing +initialize!(exchanger::StateExchanger, atmosphere) = initialize!(exchanger.atmosphere_exchanger, exchanger.exchange_grid, atmosphere) -function initialize!(exchanger::StateExchanger, atmosphere::PrescribedAtmosphere) +function initialize!(frac_indices, exchange_grid::AbstractGrid, 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) @@ -165,26 +164,28 @@ Base.summary(crf::ComponentInterfaces) = "ComponentInterfaces" Base.show(io::IO, crf::ComponentInterfaces) = print(io, summary(crf)) atmosphere_ocean_interface(::Nothing, args...) = nothing +atmosphere_ocean_interface(ocean::OceananigansSimulation, args...) = + atmosphere_ocean_interface(ocean.model.grid, args...) -function atmosphere_ocean_interface(atmos, - ocean, +function atmosphere_ocean_interface(grid::AbstractGrid, + atmos, radiation, ao_flux_formulation, temperature_formulation, velocity_formulation, specific_humidity_formulation) - water_vapor = Field{Center, Center, Nothing}(ocean.model.grid) - latent_heat = Field{Center, Center, Nothing}(ocean.model.grid) - sensible_heat = Field{Center, Center, Nothing}(ocean.model.grid) - x_momentum = Field{Center, Center, Nothing}(ocean.model.grid) - y_momentum = Field{Center, Center, Nothing}(ocean.model.grid) - friction_velocity = Field{Center, Center, Nothing}(ocean.model.grid) - temperature_scale = Field{Center, Center, Nothing}(ocean.model.grid) - water_vapor_scale = Field{Center, Center, Nothing}(ocean.model.grid) - upwelling_longwave = Field{Center, Center, Nothing}(ocean.model.grid) - downwelling_longwave = Field{Center, Center, Nothing}(ocean.model.grid) - downwelling_shortwave = Field{Center, Center, Nothing}(ocean.model.grid) + water_vapor = Field{Center, Center, Nothing}(grid) + latent_heat = Field{Center, Center, Nothing}(grid) + sensible_heat = Field{Center, Center, Nothing}(grid) + x_momentum = Field{Center, Center, Nothing}(grid) + y_momentum = Field{Center, Center, Nothing}(grid) + friction_velocity = Field{Center, Center, Nothing}(grid) + temperature_scale = Field{Center, Center, Nothing}(grid) + water_vapor_scale = Field{Center, Center, Nothing}(grid) + upwelling_longwave = Field{Center, Center, Nothing}(grid) + downwelling_longwave = Field{Center, Center, Nothing}(grid) + downwelling_shortwave = Field{Center, Center, Nothing}(grid) ao_fluxes = (; latent_heat, sensible_heat, @@ -208,7 +209,7 @@ function atmosphere_ocean_interface(atmos, temperature_formulation, velocity_formulation) - interface_temperature = Field{Center, Center, Nothing}(ocean.model.grid) + interface_temperature = Field{Center, Center, Nothing}(grid) return AtmosphereInterface(ao_fluxes, ao_flux_formulation, interface_temperature, ao_properties) end @@ -250,16 +251,18 @@ function atmosphere_sea_ice_interface(atmos, return AtmosphereInterface(fluxes, ai_flux_formulation, interface_temperature, properties) end -sea_ice_ocean_interface(sea_ice, ocean) = nothing +sea_ice_ocean_interface(ocean, sea_ice) = nothing +sea_ice_ocean_interface(ocean, sea_ice::SeaIceSimulation; kw...) = + sea_ice_ocean_interface(ocean.grid, sea_ice; kw...) -function sea_ice_ocean_interface(sea_ice::SeaIceSimulation, ocean; +function sea_ice_ocean_interface(grid::AbstractGrid, sea_ice::SeaIceSimulation; characteristic_melting_speed = 1e-5) - io_bottom_heat_flux = Field{Center, Center, Nothing}(ocean.model.grid) - io_frazil_heat_flux = Field{Center, Center, Nothing}(ocean.model.grid) - io_salt_flux = Field{Center, Center, Nothing}(ocean.model.grid) - x_momentum = Field{Face, Center, Nothing}(ocean.model.grid) - y_momentum = Field{Center, Face, Nothing}(ocean.model.grid) + io_bottom_heat_flux = Field{Center, Center, Nothing}(grid) + io_frazil_heat_flux = Field{Center, Center, Nothing}(grid) + io_salt_flux = Field{Center, Center, Nothing}(grid) + x_momentum = Field{Face, Center, Nothing}(grid) + y_momentum = Field{Center, Face, Nothing}(grid) @assert io_frazil_heat_flux isa Field{Center, Center, Nothing} @assert io_bottom_heat_flux isa Field{Center, Center, Nothing} @@ -284,12 +287,23 @@ function default_ai_temperature(sea_ice::SeaIceSimulation) end function default_ao_specific_humidity(ocean) - FT = eltype(ocean.model.grid) + FT = eltype(ocean) phase = AtmosphericThermodynamics.Liquid() x_H₂O = convert(FT, 0.98) return ImpureSaturationSpecificHumidity(phase, x_H₂O) end +function ocean_surface_fluxes(ocean::OceananigansSimulation, ρₒ, cₒ) + τx = surface_flux(ocean.model.velocities.u) + τy = surface_flux(ocean.model.velocities.v) + tracers = ocean.model.tracers + Qₒ = ρₒ * cₒ * surface_flux(ocean.model.tracers.T) + net_ocean_surface_fluxes = (u=τx, v=τy, Q=Qₒ) + + ocean_surface_tracer_fluxes = NamedTuple(name => surface_flux(tracers[name]) for name in keys(tracers)) + return merge(ocean_surface_tracer_fluxes, net_ocean_surface_fluxes) +end + """ ComponentInterfaces(atmosphere, ocean, sea_ice=nothing; radiation = Radiation(), @@ -310,8 +324,8 @@ end function ComponentInterfaces(atmosphere, ocean, sea_ice=nothing; radiation = Radiation(), freshwater_density = default_freshwater_density, - atmosphere_ocean_fluxes = SimilarityTheoryFluxes(eltype(ocean.model.grid)), - atmosphere_sea_ice_fluxes = SimilarityTheoryFluxes(eltype(ocean.model.grid)), + atmosphere_ocean_fluxes = SimilarityTheoryFluxes(), + atmosphere_sea_ice_fluxes = SimilarityTheoryFluxes(), atmosphere_ocean_interface_temperature = BulkTemperature(), atmosphere_ocean_velocity_difference = RelativeVelocity(), atmosphere_ocean_interface_specific_humidity = default_ao_specific_humidity(ocean), @@ -325,8 +339,7 @@ function ComponentInterfaces(atmosphere, ocean, sea_ice=nothing; sea_ice_heat_capacity = heat_capacity(sea_ice), gravitational_acceleration = default_gravitational_acceleration) - ocean_grid = ocean.model.grid - FT = eltype(ocean_grid) + FT = eltype(ocean) ocean_reference_density = convert(FT, ocean_reference_density) ocean_heat_capacity = convert(FT, ocean_heat_capacity) @@ -342,8 +355,8 @@ function ComponentInterfaces(atmosphere, ocean, sea_ice=nothing; freshwater_density = freshwater_density, temperature_units = ocean_temperature_units) - ao_interface = atmosphere_ocean_interface(atmosphere, - ocean, + ao_interface = atmosphere_ocean_interface(ocean, + atmosphere, radiation, atmosphere_ocean_fluxes, atmosphere_ocean_interface_temperature, @@ -384,23 +397,14 @@ function ComponentInterfaces(atmosphere, ocean, sea_ice=nothing; net_bottom_sea_ice_fluxes = nothing end - τx = surface_flux(ocean.model.velocities.u) - τy = surface_flux(ocean.model.velocities.v) - tracers = ocean.model.tracers - ρₒ = ocean_reference_density - cₒ = ocean_heat_capacity - Qₒ = ρₒ * cₒ * surface_flux(ocean.model.tracers.T) - net_ocean_surface_fluxes = (u=τx, v=τy, Q=Qₒ) - - ocean_surface_tracer_fluxes = NamedTuple(name => surface_flux(tracers[name]) for name in keys(tracers)) - net_ocean_surface_fluxes = merge(ocean_surface_tracer_fluxes, net_ocean_surface_fluxes) + net_ocean_surface_fluxes = ocean_surface_fluxes(ocean, ocean_reference_density, ocean_heat_capacity) # Total interface fluxes net_fluxes = (ocean_surface = net_ocean_surface_fluxes, sea_ice_top = net_top_sea_ice_fluxes, sea_ice_bottom = net_bottom_sea_ice_fluxes) - exchanger = StateExchanger(ocean, atmosphere) + exchanger = state_exchanger(ocean, atmosphere) properties = (; gravitational_acceleration) diff --git a/src/OceanSeaIceModels/InterfaceComputations/interpolate_atmospheric_state.jl b/src/OceanSeaIceModels/InterfaceComputations/interpolate_atmospheric_state.jl index d79a9caf9..63baf19f2 100644 --- a/src/OceanSeaIceModels/InterfaceComputations/interpolate_atmospheric_state.jl +++ b/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) diff --git a/src/OceanSeaIceModels/InterfaceComputations/sea_ice_ocean_fluxes.jl b/src/OceanSeaIceModels/InterfaceComputations/sea_ice_ocean_fluxes.jl index 1cc1f7cb8..8d5d805a0 100644 --- a/src/OceanSeaIceModels/InterfaceComputations/sea_ice_ocean_fluxes.jl +++ b/src/OceanSeaIceModels/InterfaceComputations/sea_ice_ocean_fluxes.jl @@ -23,10 +23,13 @@ function compute_sea_ice_ocean_fluxes!(sea_ice_ocean_fluxes, ocean, sea_ice, mel ℵᵢ = sea_ice.model.ice_concentration hᵢ = sea_ice.model.ice_thickness Gh = sea_ice.model.ice_thermodynamics.thermodynamic_tendency + Δt = sea_ice.Δt + + ocean_state = get_ocean_state(ocean, coupled_model) liquidus = sea_ice.model.ice_thermodynamics.phase_transitions.liquidus - grid = ocean.model.grid - clock = ocean.model.clock + grid = sea_ice.model.grid + clock = sea_ice.model.clock arch = architecture(grid) uᵢ, vᵢ = sea_ice.model.velocities @@ -41,7 +44,7 @@ function compute_sea_ice_ocean_fluxes!(sea_ice_ocean_fluxes, ocean, sea_ice, mel # What about the latent heat removed from the ocean when ice forms? # Is it immediately removed from the ocean? Or is it stored in the ice? launch!(arch, grid, :xy, _compute_sea_ice_ocean_fluxes!, - sea_ice_ocean_fluxes, grid, clock, hᵢ, ℵᵢ, Sᵢ, Gh, Tₒ, Sₒ, uᵢ, vᵢ, + sea_ice_ocean_fluxes, grid, clock, hᵢ, ℵᵢ, Sᵢ, Gh, ocean_state, uᵢ, vᵢ, τs, liquidus, ocean_properties, melting_speed, Δt) return nothing @@ -54,8 +57,7 @@ end ice_concentration, ice_salinity, thermodynamic_tendency, - ocean_temperature, - ocean_salinity, + ocean_state, sea_ice_u_velocity, sea_ice_v_velocity, sea_ice_ocean_stresses, @@ -74,8 +76,8 @@ end τy = sea_ice_ocean_fluxes.y_momentum uᵢ = sea_ice_u_velocity vᵢ = sea_ice_v_velocity - Tₒ = ocean_temperature - Sₒ = ocean_salinity + Tₒ = ocean_state.T + Sₒ = ocean_state.S Sᵢ = ice_salinity hᵢ = ice_thickness ℵᵢ = ice_concentration diff --git a/src/OceanSeaIceModels/OceanSeaIceModels.jl b/src/OceanSeaIceModels/OceanSeaIceModels.jl index a5226076b..84803a823 100644 --- a/src/OceanSeaIceModels/OceanSeaIceModels.jl +++ b/src/OceanSeaIceModels/OceanSeaIceModels.jl @@ -55,6 +55,8 @@ const default_freshwater_density = 1000 # kg m⁻³ const SeaIceSimulation = Simulation{<:SeaIceModel} const OceananigansSimulation = Simulation{<:HydrostaticFreeSurfaceModel} +Base.eltype(ocean::OceananigansSimulation) = eltype(ocean.model.grid) + sea_ice_thickness(::Nothing) = ZeroField() sea_ice_thickness(sea_ice::SeaIceSimulation) = sea_ice.model.ice_thickness diff --git a/src/OceanSeaIceModels/PrescribedAtmospheres.jl b/src/OceanSeaIceModels/PrescribedAtmospheres.jl index 49d72d32d..1d6d4d57e 100644 --- a/src/OceanSeaIceModels/PrescribedAtmospheres.jl +++ b/src/OceanSeaIceModels/PrescribedAtmospheres.jl @@ -5,6 +5,7 @@ using Oceananigans.Fields: Center using Oceananigans.Grids: grid_name using Oceananigans.OutputReaders: FieldTimeSeries, update_field_time_series!, extract_field_time_series using Oceananigans.TimeSteppers: Clock, tick! +using Oceananigans.Simulations: TimeStepWizard using Oceananigans.Utils: prettysummary, Time using Adapt @@ -372,6 +373,8 @@ end return nothing end +(wizard::TimeStepWizard)(atmos::PrescribedAtmosphere) = Inf + @inline thermodynamics_parameters(atmos::Nothing) = nothing @inline thermodynamics_parameters(atmos::PrescribedAtmosphere) = atmos.thermodynamics_parameters @inline surface_layer_height(atmos::PrescribedAtmosphere) = atmos.surface_layer_height diff --git a/src/OceanSeaIceModels/ocean_sea_ice_model.jl b/src/OceanSeaIceModels/ocean_sea_ice_model.jl index 3e09b751e..1f2b9d2ca 100644 --- a/src/OceanSeaIceModels/ocean_sea_ice_model.jl +++ b/src/OceanSeaIceModels/ocean_sea_ice_model.jl @@ -20,7 +20,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))