diff --git a/.buildkite/examples_build.yml b/.buildkite/examples_build.yml index 86c282d22..a78ec7d20 100644 --- a/.buildkite/examples_build.yml +++ b/.buildkite/examples_build.yml @@ -30,7 +30,7 @@ steps: key: "build_documentation" commands: - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia --color=yes -O0 --project=docs/ -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate(); Pkg.precompile(; strict=true)'" - - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia --color=yes -O0 --project=docs/ docs/make.jl" + - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia --color=yes -t auto --check-bounds=no --project=docs/ docs/make.jl" agents: queue: ClimaOcean-docs diff --git a/CondaPkg.toml b/CondaPkg.toml index dde757656..e214d0360 100644 --- a/CondaPkg.toml +++ b/CondaPkg.toml @@ -1,7 +1,2 @@ - -[pip.deps] -copernicusmarine = ">=2.0.0" -xarray = ">=2024.7.0" -numpy = ">=2.0.0" -jax = ">=0.6" -tensorflow = ">=2.17" +[deps] +xesmf = "0.8.10" diff --git a/Project.toml b/Project.toml index 3d739e6c1..86dee0e9a 100644 --- a/Project.toml +++ b/Project.toml @@ -33,18 +33,22 @@ ZipFile = "a5390f91-8eb1-5f08-bee0-b1d1ffed6cea" [weakdeps] CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" +SpeedyWeather = "9e226e20-d153-4fed-8a5b-493def4f21a9" [extensions] ClimaOceanPythonCallExt = ["PythonCall", "CondaPkg"] ClimaOceanReactantExt = "Reactant" +ClimaOceanSpeedyWeatherExt = "SpeedyWeather" [compat] Adapt = "4" CFTime = "0.1, 0.2" CUDA = "4, 5" -ClimaSeaIce = "0.3.5" +ClimaSeaIce = "0.3.7" CondaPkg = "0.2.28" CubicSplines = "0.2" DataDeps = "0.7" @@ -55,10 +59,9 @@ JLD2 = "0.4, 0.5, 0.6" KernelAbstractions = "0.9" MPI = "0.20" NCDatasets = "0.12, 0.13, 0.14" -Oceananigans = "0.97.6, 0.98" +Oceananigans = "0.99" OffsetArrays = "1.14" PrecompileTools = "1" -PythonCall = "0.9" Reactant = "0.2.45" Scratch = "1" SeawaterPolynomials = "0.3.5" @@ -76,4 +79,4 @@ MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Coverage", "Test", "MPIPreferences", "CUDA_Runtime_jll", "Reactant", "PythonCall", "CondaPkg"] +test = ["Coverage", "Test", "MPIPreferences", "CUDA_Runtime_jll", "PythonCall", "CondaPkg", "SpeedyWeather", "Reactant"] diff --git a/docs/make.jl b/docs/make.jl index 99afa9e46..c19715d34 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -20,7 +20,8 @@ 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" + "near_global_ocean_simulation.jl", + "atmosphere_ocean_simulation.jl", ] for file in to_be_literated @@ -47,6 +48,7 @@ pages = [ "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", + "Coupled atmosphere--ocean--sea ice simulation" => "literated/coupled_simulation.md", ], "Vertical grids" => "vertical_grids.md", diff --git a/examples/atmosphere_ocean_simulation.jl b/examples/atmosphere_ocean_simulation.jl new file mode 100644 index 000000000..691264345 --- /dev/null +++ b/examples/atmosphere_ocean_simulation.jl @@ -0,0 +1,243 @@ +# # Global coupled atmosphere and ocean--sea ice simulation +# +# This example configures a global ocean--sea ice simulation at 1ᵒ horizontal resolution with +# realistic bathymetry and a few closures including the "Gent-McWilliams" `IsopycnalSkewSymmetricDiffusivity`. +# The atmosphere is represented by a SpeecdyWeather model at T63 resolution (approximately 1.875ᵒ). +# and initialized by temperature, salinity, sea ice concentration, and sea ice thickness +# from the ECCO state estimate. +# +# For this example, we need Oceananigans.HydrostaticFreeSurfaceModel (the ocean), ClimaSeaIce.SeaIceModel (the sea ice) and +# SpeedyWeather.PrimitiveWetModel (the atmosphere), coupled and orchestrated by ClimaOcean.OceanSeaIceModel (the coupled system). + +using Oceananigans, ClimaSeaIce, SpeedyWeather, ClimaOcean +using NCDatasets, CairoMakie +using Oceananigans.Units +using Printf, Statistics, Dates + +# ## Ocean and sea-ice model configuration +# The ocean and sea-ice are configured in accordance with the "one_degree_simulation" example +# https://clima.github.io/ClimaOceanDocumentation/dev/literated/one_degree_simulation/ +# The first step is to create the grid with realistic bathymetry. + +Nx = 360 +Ny = 180 +Nz = 30 + +r_faces = ExponentialCoordinate(Nz, -4000, 0) +grid = TripolarGrid(CPU(); size=(Nx, Ny, Nz), z=r_faces, halo=(6, 6, 5)) + +# Regridding the bathymetry... + +bottom_height = regrid_bathymetry(grid; major_basins=1, interpolation_passes=15) +grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height); active_cells_map=true) + +# Now we can specify the numerical details and closures for the ocean simulation. + +momentum_advection = WENOVectorInvariant(order=5) +tracer_advection = WENO(order=5) +free_surface = SplitExplicitFreeSurface(grid; substeps=80) + +catke_closure = ClimaOcean.OceanSimulations.default_ocean_closure() +eddy_closure = Oceananigans.TurbulenceClosures.IsopycnalSkewSymmetricDiffusivity(κ_skew=1e3, κ_symmetric=1e3) +closures = (catke_closure, eddy_closure, VerticalScalarDiffusivity(κ=1e-5, ν=1e-4)) + +# The ocean simulation, complete with initial conditions for temperature and salinity from ECCO. + +ocean = ocean_simulation(grid; + momentum_advection, + tracer_advection, + free_surface, + timestepper = :SplitRungeKutta3, + closure = closures) + +Oceananigans.set!(ocean.model, T=Metadatum(:temperature, dataset=ECCO4Monthly()), + S=Metadatum(:salinity, dataset=ECCO4Monthly())) + + +# The sea-ice simulation, complete with initial conditions for sea-ice thickness and concentration from ECCO. + +dynamics = ClimaOcean.SeaIceSimulations.sea_ice_dynamics(grid, ocean; solver=ClimaSeaIce.SeaIceDynamics.SplitExplicitSolver(100)) +sea_ice = sea_ice_simulation(grid, ocean; dynamics, advection=WENO(order=7)) + +Oceananigans.set!(sea_ice.model, h=Metadatum(:sea_ice_thickness, dataset=ECCO4Monthly()), + ℵ=Metadatum(:sea_ice_concentration, dataset=ECCO4Monthly())) + +# ## Atmosphere model configuration +# The atmosphere is provided by SpeedyWeather.jl. Here we configure a T63L8 model with a 3 hour output interval. +# The `atmosphere_simulation` function takes care of building an atmosphere model with appropriate +# hooks for ClimaOcean to compute intercomponent fluxes. We also set the output interval to 3 hours. + +spectral_grid = SpectralGrid(trunc=63, nlayers=8, Grid=FullClenshawGrid) +atmosphere = atmosphere_simulation(spectral_grid; output=true) +atmosphere.model.output.output_dt = Hour(3) + +# ## The coupled model +# Now we can build the coupled model. We need to specify the time step for the coupled model. +# We decide to step the global model every 2 atmosphere time steps. (i.e. the ocean and the +# sea-ice will be stepped every two atmosphere time steps). + +Δt = 2 * convert(eltype(grid), atmosphere.model.time_stepping.Δt_sec) + +# We build the complete model. Since radiation is idealized in this example, we set the emissivities to zero. + +radiation = Radiation(ocean_emissivity=0.0, sea_ice_emissivity=0.0) +earth_model = OceanSeaIceModel(ocean, sea_ice; atmosphere, radiation) +earth = Oceananigans.Simulation(earth_model; Δt, stop_time=360days) + +# ## Running the simulation +# We can now run the simulation. We add a callback to monitor the progress of the simulation +# and write outputs to disk every 3 hours. + +wall_time = Ref(time_ns()) + +function progress(sim) + sea_ice = sim.model.sea_ice + ocean = sim.model.ocean + hmax = maximum(sea_ice.model.ice_thickness) + ℵmax = maximum(sea_ice.model.ice_concentration) + uimax = maximum(abs, sea_ice.model.velocities.u) + vimax = maximum(abs, sea_ice.model.velocities.v) + uomax = maximum(abs, ocean.model.velocities.u) + vomax = maximum(abs, ocean.model.velocities.v) + + step_time = 1e-9 * (time_ns() - wall_time[]) + + msg1 = @sprintf("time: %s, iteration: %d, Δt: %s, ", prettytime(sim), iteration(sim), prettytime(sim.Δt)) + msg2 = @sprintf("max(h): %.2e m, max(ℵ): %.2e ", hmax, ℵmax) + msg3 = @sprintf("max uᵢ: (%.2f, %.2f) m s⁻¹, ", uimax, vimax) + msg4 = @sprintf("max uₒ: (%.2f, %.2f) m s⁻¹, ", uomax, vomax) + msg5 = @sprintf("wall time: %s \n", prettytime(step_time)) + + @info msg1 * msg2 * msg3 * msg4 * msg5 + + wall_time[] = time_ns() + + return nothing +end + +outputs = merge(ocean.model.velocities, ocean.model.tracers) +sea_ice_fields = merge(sea_ice.model.velocities, sea_ice.model.dynamics.auxiliaries.fields, + (; h=sea_ice.model.ice_thickness, ℵ=sea_ice.model.ice_concentration)) + +ocean.output_writers[:free_surf] = JLD2Writer(ocean.model, (; η=ocean.model.free_surface.η); + overwrite_existing=true, + schedule=TimeInterval(3600 * 3), + filename="ocean_free_surface.jld2") + +ocean.output_writers[:surface] = JLD2Writer(ocean.model, outputs; + overwrite_existing=true, + schedule=TimeInterval(3600 * 3), + filename="ocean_surface_fields.jld2", + indices=(:, :, grid.Nz)) + +sea_ice.output_writers[:fields] = JLD2Writer(sea_ice.model, sea_ice_fields; + overwrite_existing=true, + schedule=TimeInterval(3600 * 3), + filename="sea_ice_fields.jld2") + +Qcao = earth.model.interfaces.atmosphere_ocean_interface.fluxes.sensible_heat +Qvao = earth.model.interfaces.atmosphere_ocean_interface.fluxes.latent_heat +τxao = earth.model.interfaces.atmosphere_ocean_interface.fluxes.x_momentum +τyao = earth.model.interfaces.atmosphere_ocean_interface.fluxes.y_momentum +Qcai = earth.model.interfaces.atmosphere_sea_ice_interface.fluxes.sensible_heat +Qvai = earth.model.interfaces.atmosphere_sea_ice_interface.fluxes.latent_heat +τxai = earth.model.interfaces.atmosphere_sea_ice_interface.fluxes.x_momentum +τyai = earth.model.interfaces.atmosphere_sea_ice_interface.fluxes.y_momentum +Qoi = earth.model.interfaces.net_fluxes.sea_ice_bottom.heat +Soi = earth.model.interfaces.sea_ice_ocean_interface.fluxes.salt +fluxes = (; Qcao, Qvao, τxao, τyao, Qcai, Qvai, τxai, τyai, Qoi, Soi) + +earth.output_writers[:fluxes] = JLD2Writer(earth.model.ocean.model, fluxes; + overwrite_existing=true, + schedule=TimeInterval(3600 * 3), + filename="intercomponent_fluxes.jld2") + +add_callback!(earth, progress, IterationInterval(100)) + +Oceananigans.run!(earth) + +# ## Visualizing the results +# We can visualize some of the results. Here we plot the surface speeds in the atmosphere, ocean, and sea-ice +# as well as the 2m temperature in the atmosphere, the sea surface temperature, and the sensible and latent heat +# fluxes at the atmosphere-ocean interface. + +SWO = Dataset("run_0001/output.nc") + +Ta = reverse(SWO["temp"][:, :, 8, :], dims=2) +ua = reverse(SWO["u"][:, :, 8, :], dims=2) +va = reverse(SWO["v"][:, :, 8, :], dims=2) +sp = sqrt.(ua.^2 + va.^2) + +SST = FieldTimeSeries("ocean_surface_fields.jld2", "T") +SSU = FieldTimeSeries("ocean_surface_fields.jld2", "u") +SSV = FieldTimeSeries("ocean_surface_fields.jld2", "v") + +SIU = FieldTimeSeries("sea_ice_fields.jld2", "u") +SIV = FieldTimeSeries("sea_ice_fields.jld2", "v") +SIA = FieldTimeSeries("sea_ice_fields.jld2", "ℵ") + +Qcao = FieldTimeSeries("intercomponent_fluxes.jld2", "Qcao") +Qvao = FieldTimeSeries("intercomponent_fluxes.jld2", "Qvao") + +uotmp = Field{Face, Center, Nothing}(SST.grid) +votmp = Field{Center, Face, Nothing}(SST.grid) + +uitmp = Field{Face, Center, Nothing}(SST.grid) +vitmp = Field{Center, Face, Nothing}(SST.grid) +atmp = Field{Center, Center, Nothing}(SST.grid) + +sotmp = Field(sqrt(uotmp^2 + votmp^2)) +sitmp = Field(sqrt(uitmp^2 + vitmp^2) * atmp) + +iter = Observable(1) +san = @lift sp[:, :, $iter] +son = @lift begin + set!(uotmp, SSU[$iter * 2]) + set!(votmp, SSV[$iter * 2]) + compute!(sotmp) + interior(sotmp, :, :, 1) +end + +ssn = @lift begin + set!(uitmp, SIU[$iter * 2]) + set!(vitmp, SIV[$iter * 2]) + set!(atmp, SIA[$iter * 2]) + compute!(sitmp) + interior(sitmp, :, :, 1) +end + +fig = Figure(size = (1200, 800)) +ax2 = Axis(fig[1, 1], title = "Surface speed, atmosphere (m/s)") +hm2 = heatmap!(ax2, san; colormap = :deep) +ax1 = Axis(fig[1, 2], title = "Surface speed, ocean (m/s)") +hm = heatmap!(ax1, son; colormap = :deep) +ax3 = Axis(fig[1, 3], title = "Surface speed, sea-ice (m/s)") +hm = heatmap!(ax3, ssn; colormap = :deep) + +record(fig, "surface_speeds.mp4", iter, framerate = 15) do i + iter[] = i +endnothing #hide + +# ![](surface_speeds.mp4) + +Tan = @lift Ta[:, :, $iter] +Ton = @lift interior(SST[$iter * 2], :, :, 1) +Qcn = @lift interior(Qcao[$iter * 2], :, :, 1) +Qvn = @lift interior(Qvao[$iter * 2], :, :, 1) + +fig = Figure(size = (1200, 800)) +ax1 = Axis(fig[1, 1], title = "2m Temperature, atmosphere (K)") +hm = heatmap!(ax1, Tan; colormap = :plasma) +ax2 = Axis(fig[1, 2], title = "Sea Surface Temperature (C)") +hm2 = heatmap!(ax2, Ton; colormap = :plasma) +ax3 = Axis(fig[2, 1], title = "Sensible heat flux (W/m²)") +hm3 = heatmap!(ax3, Qcn; colormap = :balance, colorrange = (-200, 200)) +ax4 = Axis(fig[2, 2], title = "Latent heat flux (W/m²)") +hm4 = heatmap!(ax4, Qvn; colormap = :balance, colorrange = (-200, 200)) + +record(fig, "surface_temperature_and_heat_flux.mp4", 1:length(sp[1, 1, :])) do i + iter[] = i +end +nothing #hide + +# ![](surface_temperature_and_heat_flux.mp4) diff --git a/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl b/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl new file mode 100644 index 000000000..ed214dbb2 --- /dev/null +++ b/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl @@ -0,0 +1,19 @@ +module ClimaOceanSpeedyWeatherExt + +using OffsetArrays +using KernelAbstractions +using Statistics + +import SpeedyWeather +import ClimaOcean +import Oceananigans +import SpeedyWeather.RingGrids + +include("speedy_atmosphere_simulations.jl") + +# TODO: Remove this when we have a proper interface +# for regridding between tripolar and latitude-longitude grids +include("bilinear_interpolator.jl") +include("speedy_weather_exchanger.jl") + +end # module ClimaOceanSpeedyWeatherExt \ No newline at end of file diff --git a/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl b/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl new file mode 100644 index 000000000..e5d5eaaf9 --- /dev/null +++ b/ext/ClimaOceanSpeedyWeatherExt/bilinear_interpolator.jl @@ -0,0 +1,163 @@ +using PythonCall +using CondaPkg +using Oceananigans.Grids: AbstractGrid +using SparseArrays +using LinearAlgebra + +const Grids = Union{SpeedyWeather.SpectralGrid, AbstractGrid} + +struct BilinearInterpolator{W1, W2} + set1 :: W1 + set2 :: W2 +end + +function BilinearInterpolator(grid1::Grids, grid2::Grids) + W1 = regridder_weights(grid1, grid2; method="bilinear") + W2 = regridder_weights(grid2, grid1; method="bilinear") + return BilinearInterpolator(W1, W2) +end + +regrid!(dst, weights, scr) = LinearAlgebra.mul!(vec(dst), weights, vec(scr)) + +function install_package(package) + @info "Installing the $package package..." + CondaPkg.add(package; channel = "conda-forge") + return nothing +end + +# Import and store as constants for submodules +function get_package(package) + try + return pyimport(package) + catch + install_package(package) + return pyimport(package) + end +end + +two_dimensionalize(lat::Matrix, lon::Matrix) = lat, lon + +function two_dimensionalize(lat::AbstractVector, lon::AbstractVector) + Nx = length(lon) + Ny = length(lat) + lat = repeat(lat', Nx) + lon = repeat(lon, 1, Ny) + return lat, lon +end + +coordinate_dataset(grid::SpeedyWeather.SpectralGrid) = coordinate_dataset(grid.grid) + +function coordinate_dataset(grid::SpeedyWeather.RingGrids.AbstractGrid) + numpy = get_package("numpy") + lon, lat = RingGrids.get_londlatds(grid) + return numpy.array(Dict("lat" => lat, "lon" => lon)) +end + +function coordinate_dataset(grid::SpeedyWeather.RingGrids.AbstractFullGrid) + lon = RingGrids.get_lond(grid) + lat = RingGrids.get_latd(grid) + dlon = lon[2] - lon[1] + + lat_b = [90, 0.5 .* (lat[1:end-1] .+ lat[2:end])..., -90] + lon_b = [lon[1] - dlon / 2, lon .+ dlon / 2...] + + lat, lon = two_dimensionalize(lat, lon) + lat_b, lon_b = two_dimensionalize(lat_b, lon_b) + + return structured_coordinate_dataset(lat, lon, lat_b, lon_b) +end + +function coordinate_dataset(grid::AbstractGrid) + lat = Array(φnodes(grid, Center(), Center(), Center())) + lon = Array(λnodes(grid, Center(), Center(), Center())) + + lat_b = Array(φnodes(grid, Face(), Face(), Center())) + lon_b = Array(λnodes(grid, Face(), Face(), Center())) + + lat, lon = two_dimensionalize(lat, lon) + lat_b, lon_b = two_dimensionalize(lat_b, lon_b) + + return structured_coordinate_dataset(lat, lon, lat_b, lon_b) +end + +function structured_coordinate_dataset(lat, lon, lat_b, lon_b) + numpy = get_package("numpy") + xarray = get_package("xarray") + + lat = Array(lat') + lon = Array(lon') + lat_b = Array(lat_b') + lon_b = Array(lon_b') + + lat = numpy.array(lat) + lon = numpy.array(lon) + + lat_b = numpy.array(lat_b) + lon_b = numpy.array(lon_b) + + ds_lat = xarray.DataArray( + lat, + dims=["y", "x"], + coords=Dict( + "lat" => (["y", "x"], lat), + "lon" => (["y", "x"], lon) + ), + name="latitude" + ) + + ds_lon = xarray.DataArray( + lon, + dims=["y", "x"], + coords=Dict( + "lat" => (["y", "x"], lat), + "lon" => (["y", "x"], lon) + ), + name="longitude" + ) + + ds_lat_b = xarray.DataArray( + lat_b, + dims=["y_b", "x_b"], + coords=Dict( + "lat_b" => (["y_b", "x_b"], lat_b), + "lon_b" => (["y_b", "x_b"], lon_b) + ), + ) + + ds_lon_b = xarray.DataArray( + lon_b, + dims=["y_b", "x_b"], + coords=Dict( + "lat_b" => (["y_b", "x_b"], lat_b), + "lon_b" => (["y_b", "x_b"], lon_b) + ), + ) + + return xarray.Dataset( + Dict("lat" => ds_lat, + "lon" => ds_lon, + "lat_b" => ds_lat_b, + "lon_b" => ds_lon_b) + ) +end + +function regridder_weights(dst::Grids, src::Grids; method::String="bilinear") + + src_ds = coordinate_dataset(src) + dst_ds = coordinate_dataset(dst) + + regridder = get_package("xesmf").Regridder(src_ds, dst_ds, method, periodic=true) + + # Move back to Julia + # Convert the regridder weights to a Julia sparse matrix + coords = regridder.weights.data + shape = pyconvert(Tuple{Int, Int}, coords.shape) + vals = pyconvert(Array{Float64}, coords.data) + coords = pyconvert(Array{Float64}, coords.coords) + rows = coords[1,:].+1 + cols = coords[2,:].+1 + + W = sparse(rows, cols, vals, shape[1], shape[2]) + + return W +end diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl new file mode 100644 index 000000000..5713db86a --- /dev/null +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl @@ -0,0 +1,84 @@ +import ClimaOcean: atmosphere_simulation + +# Make sure the atmospheric parameters from SpeedyWeather can be used in the compute fluxes function +import ClimaOcean.OceanSeaIceModels.PrescribedAtmospheres: + thermodynamics_parameters, + boundary_layer_height, + surface_layer_height + +const SpeedySimulation = SpeedyWeather.Simulation +const SpeedyCoupledModel = ClimaOcean.OceanSeaIceModel{<:Any, <:SpeedySimulation} +Base.summary(::SpeedySimulation) = "SpeedyWeather.Simulation" + +# Take one time-step or more depending on the global timestep +function Oceananigans.TimeSteppers.time_step!(atmos::SpeedySimulation, Δt) + Δt_atmos = atmos.model.time_stepping.Δt_sec + nsteps = ceil(Int, Δt / Δt_atmos) + + if (Δt / Δt_atmos) % 1 != 0 + @warn "ClimaOcean only supports atmosphere timesteps that are integer divisors of the ESM timesteps" + end + + for _ in 1:nsteps + SpeedyWeather.timestep!(atmos) + end +end + +# The height of near-surface variables used in the turbulent flux solver +function surface_layer_height(s::SpeedySimulation) + T = s.model.atmosphere.temp_ref + g = s.model.planet.gravity + Φ = s.model.geopotential.Δp_geopot_full + return Φ[end] * T / g +end + +# This is a parameter that is used in the computation of the fluxes, +# It probably should not be here but in the similarity theory type. +boundary_layer_height(atmos::SpeedySimulation) = 600 + +# This is a _hack_!! The parameters should be consistent with what is specified in SpeedyWeather +thermodynamics_parameters(atmos::SpeedySimulation) = + ClimaOcean.OceanSeaIceModels.AtmosphereThermodynamicsParameters(Float32) + +function initialize_atmospheric_state!(simulation::SpeedyWeather.Simulation) + progn, diagn, model = SpeedyWeather.unpack(simulation) + (; time) = progn.clock # current time + + # set the tendencies back to zero for accumulation + fill!(diagn.tendencies, 0, typeof(model)) + + if model.physics + SpeedyWeather.parameterization_tendencies!(diagn, progn, time, model) + end + + return nothing +end + +function atmosphere_simulation(spectral_grid::SpeedyWeather.SpectralGrid; output=false) + # Surface fluxes + humidity_flux_ocean = SpeedyWeather.PrescribedOceanHumidityFlux(spectral_grid) + humidity_flux_land = SpeedyWeather.SurfaceLandHumidityFlux(spectral_grid) + surface_humidity_flux = SpeedyWeather.SurfaceHumidityFlux(ocean=humidity_flux_ocean, land=humidity_flux_land) + + ocean_heat_flux = SpeedyWeather.PrescribedOceanHeatFlux(spectral_grid) + land_heat_flux = SpeedyWeather.SurfaceLandHeatFlux(spectral_grid) + surface_heat_flux = SpeedyWeather.SurfaceHeatFlux(ocean=ocean_heat_flux, land=land_heat_flux) + + # The atmospheric model + atmosphere_model = SpeedyWeather.PrimitiveWetModel(spectral_grid; + surface_heat_flux, + surface_humidity_flux, + ocean = nothing, + sea_ice = nothing) # This is provided by ClimaSeaIce + + # Construct the simulation + atmosphere = SpeedyWeather.initialize!(atmosphere_model) + + # Initialize the simulation + SpeedyWeather.initialize!(atmosphere; output) + + # Fill in prognostic fields + initialize_atmospheric_state!(atmosphere) + + return atmosphere +end diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl new file mode 100644 index 000000000..2812354ca --- /dev/null +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl @@ -0,0 +1,147 @@ +using Oceananigans +using Oceananigans.BoundaryConditions +using Oceananigans.Grids: architecture +using Oceananigans.Utils: launch! +using Oceananigans.Operators: intrinsic_vector + +using ClimaOcean.OceanSeaIceModels: sea_ice_concentration + +# TODO: Implement conservative regridding when ready +# using ConservativeRegridding +# using GeoInterface: Polygon, LinearRing +import ClimaOcean.OceanSeaIceModels: + compute_net_atmosphere_fluxes! + +import ClimaOcean.OceanSeaIceModels.InterfaceComputations: + atmosphere_exchanger, + initialize!, + StateExchanger, + interpolate_atmosphere_state! + +# For the moment the workflow is: +# 1. Perform the regridding on the CPU +# 2. Eventually copy the regridded fields to the GPU +# If this work we can +# 1. Copy speedyweather gridarrays to the GPU +# 2. Perform the regridding on the GPU +function atmosphere_exchanger(atmosphere::SpeedySimulation, exchange_grid, exchange_atmosphere_state) + + # Figure this out: + spectral_grid = atmosphere.model.spectral_grid + FT = eltype(exchange_atmosphere_state.u) + + # sparse matrix multiply requires contiguous memory to speed up the computation + wrk = zeros(FT, size(exchange_grid, 1) * size(exchange_grid, 2)) + + # TODO: Implement a conservative regridder when ready + regridder = BilinearInterpolator(exchange_grid, spectral_grid) + exchanger = (; wrk, regridder) + + return exchanger +end + +# Regrid the atmospheric state on the exchange grid +function interpolate_atmosphere_state!(interfaces, atmos::SpeedySimulation, coupled_model) + atmosphere_exchanger = interfaces.exchanger.atmosphere_exchanger + regridder = atmosphere_exchanger.regridder + exchange_grid = interfaces.exchanger.exchange_grid + exchange_state = interfaces.exchanger.exchange_atmosphere_state + surface_layer = atmos.model.spectral_grid.nlayers + + ua = RingGrids.field_view(atmos.diagnostic_variables.grid.u_grid, :, surface_layer).data + va = RingGrids.field_view(atmos.diagnostic_variables.grid.v_grid, :, surface_layer).data + Ta = RingGrids.field_view(atmos.diagnostic_variables.grid.temp_grid, :, surface_layer).data + qa = RingGrids.field_view(atmos.diagnostic_variables.grid.humid_grid, :, surface_layer).data + pa = exp.(atmos.diagnostic_variables.grid.pres_grid.data) + Qsa = atmos.diagnostic_variables.physics.surface_shortwave_down.data + Qla = atmos.diagnostic_variables.physics.surface_longwave_down.data + Mpa = atmos.diagnostic_variables.physics.total_precipitation_rate.data + wrk = atmosphere_exchanger.wrk + + regrid!(wrk, regridder.set1, ua) + Oceananigans.set!(exchange_state.u, reshape(wrk, size(exchange_state.u))) + + regrid!(wrk, regridder.set1, va) + Oceananigans.set!(exchange_state.v, reshape(wrk, size(exchange_state.v))) + + regrid!(wrk, regridder.set1, Ta) + Oceananigans.set!(exchange_state.T, reshape(wrk, size(exchange_state.T))) + + regrid!(wrk, regridder.set1, qa) + Oceananigans.set!(exchange_state.q, reshape(wrk, size(exchange_state.q))) + + regrid!(wrk, regridder.set1, pa) + Oceananigans.set!(exchange_state.p, reshape(wrk, size(exchange_state.p))) + + regrid!(wrk, regridder.set1, Qsa) + Oceananigans.set!(exchange_state.Qs, reshape(wrk, size(exchange_state.Qs))) + + regrid!(wrk, regridder.set1, Qla) + Oceananigans.set!(exchange_state.Qℓ, reshape(wrk, size(exchange_state.Qℓ))) + + regrid!(wrk, regridder.set1, Mpa) + Oceananigans.set!(exchange_state.Mp, reshape(wrk, size(exchange_state.Mp))) + + arch = architecture(exchange_grid) + + u = exchange_state.u + v = exchange_state.v + + launch!(arch, exchange_grid, :xy, _rotate_winds!, u, exchange_grid, v) + + fill_halo_regions!((u, v)) + fill_halo_regions!(exchange_state.T) + fill_halo_regions!(exchange_state.q) + fill_halo_regions!(exchange_state.p) + fill_halo_regions!(exchange_state.Qs) + fill_halo_regions!(exchange_state.Qℓ) + fill_halo_regions!(exchange_state.Mp) + + return nothing +end + +@kernel function _rotate_winds!(u, grid, v) + i, j = @index(Global, NTuple) + kᴺ = size(grid, 3) + uₑ, vₑ = intrinsic_vector(i, j, kᴺ, grid, u, v) + @inbounds u[i, j, kᴺ] = uₑ + @inbounds v[i, j, kᴺ] = vₑ +end + +# TODO: Fix the coupling with the sea ice model and make sure that +# the this function works also for sea_ice=nothing and on GPUs without +# needing to allocate memory. +function compute_net_atmosphere_fluxes!(coupled_model::SpeedyCoupledModel) + atmos = coupled_model.atmosphere + grid = coupled_model.interfaces.exchanger.exchange_grid + regridder = coupled_model.interfaces.exchanger.atmosphere_exchanger.regridder + wrk = coupled_model.interfaces.exchanger.atmosphere_exchanger.wrk + + ao_fluxes = coupled_model.interfaces.atmosphere_ocean_interface.fluxes + ai_fluxes = coupled_model.interfaces.atmosphere_sea_ice_interface.fluxes + + Qco = ao_fluxes.sensible_heat + Qci = ai_fluxes.sensible_heat + Mvo = ao_fluxes.water_vapor + Mvi = ai_fluxes.water_vapor + ℵ = sea_ice_concentration(coupled_model.sea_ice) + + # All the location of these fluxes will change + Qca = atmos.prognostic_variables.ocean.sensible_heat_flux.data + Mva = atmos.prognostic_variables.ocean.surface_humidity_flux.data + sst = atmos.prognostic_variables.ocean.sea_surface_temperature.data + + To = coupled_model.interfaces.atmosphere_ocean_interface.temperature + Ti = coupled_model.interfaces.atmosphere_sea_ice_interface.temperature + + # TODO: Figure out how we are going to deal with upwelling radiation + # TODO: regrid longwave rather than a mixed surface temperature + copyto!(wrk, interior(Qco) .* (1 - ℵ) .+ ℵ .* interior(Qci)) + regrid!(Qca, regridder.set2, wrk) + copyto!(wrk, interior(Mvo) .* (1 - ℵ) .+ ℵ .* interior(Mvi)) + regrid!(Mva, regridder.set2, wrk) + copyto!(wrk, interior(To) .* (1 - ℵ) .+ ℵ .* interior(Ti) .+ 273.15) + regrid!(sst, regridder.set2, wrk) + + return nothing +end diff --git a/src/ClimaOcean.jl b/src/ClimaOcean.jl index 7f4edb40e..bd6da806c 100644 --- a/src/ClimaOcean.jl +++ b/src/ClimaOcean.jl @@ -41,6 +41,7 @@ export DatasetRestoring, ocean_simulation, sea_ice_simulation, + atmosphere_simulation, initialize! using Oceananigans @@ -77,6 +78,8 @@ end return NamedTuple{names}(vals) end +function atmosphere_simulation end + include("OceanSimulations/OceanSimulations.jl") include("SeaIceSimulations.jl") include("OceanSeaIceModels/OceanSeaIceModels.jl") diff --git a/src/InitialConditions/InitialConditions.jl b/src/InitialConditions/InitialConditions.jl index 6a815f056..8214be4c7 100644 --- a/src/InitialConditions/InitialConditions.jl +++ b/src/InitialConditions/InitialConditions.jl @@ -17,7 +17,6 @@ using JLD2 # Implementation of 3-dimensional regridding # TODO: move all the following to Oceananigans! - using Oceananigans.Fields: regrid!, interpolate! using Oceananigans.Grids: cpu_face_constructor_x, cpu_face_constructor_y, diff --git a/src/OceanSeaIceModels/InterfaceComputations/assemble_net_fluxes.jl b/src/OceanSeaIceModels/InterfaceComputations/assemble_net_fluxes.jl index 1532a3be7..e34aaddba 100644 --- a/src/OceanSeaIceModels/InterfaceComputations/assemble_net_fluxes.jl +++ b/src/OceanSeaIceModels/InterfaceComputations/assemble_net_fluxes.jl @@ -117,8 +117,8 @@ end Qs = downwelling_radiation.Qs[i, j, 1] # Downwelling shortwave radiation Qℓ = downwelling_radiation.Qℓ[i, j, 1] # Downwelling longwave radiation Qc = atmos_ocean_fluxes.sensible_heat[i, j, 1] # sensible or "conductive" heat flux - Qv = atmos_ocean_fluxes.latent_heat[i, j, 1] # latent heat flux - Mv = atmos_ocean_fluxes.water_vapor[i, j, 1] # mass flux of water vapor + Qv = atmos_ocean_fluxes.latent_heat[i, j, 1] # latent heat flux + Mv = atmos_ocean_fluxes.water_vapor[i, j, 1] # mass flux of water vapor end # Compute radiation fluxes (radiation is multiplied by the fraction of ocean, 1 - sea ice concentration) @@ -170,7 +170,7 @@ end Qio = sea_ice_ocean_fluxes.interface_heat[i, j, 1] Jᵀao = ΣQao * ρₒ⁻¹ / cₒ - Jˢao = - Sₒ * ΣFao + Jˢao = - Sₒ * ΣFao # salinity flux is the opposite of a water vapor flux Jᵀio = Qio * ρₒ⁻¹ / cₒ Jˢio = sea_ice_ocean_fluxes.salt[i, j, 1] * ℵᵢ diff --git a/src/OceanSeaIceModels/InterfaceComputations/atmosphere_sea_ice_fluxes.jl b/src/OceanSeaIceModels/InterfaceComputations/atmosphere_sea_ice_fluxes.jl index d12ae69b3..25a0fcfee 100644 --- a/src/OceanSeaIceModels/InterfaceComputations/atmosphere_sea_ice_fluxes.jl +++ b/src/OceanSeaIceModels/InterfaceComputations/atmosphere_sea_ice_fluxes.jl @@ -29,7 +29,7 @@ function compute_atmosphere_sea_ice_fluxes!(coupled_model) Qs = atmosphere_fields.Qs.data, Qℓ = atmosphere_fields.Qℓ.data, Mp = atmosphere_fields.Mp.data, - h_bℓ = atmosphere.boundary_layer_height) + h_bℓ = boundary_layer_height(atmosphere)) flux_formulation = coupled_model.interfaces.atmosphere_sea_ice_interface.flux_formulation interface_fluxes = coupled_model.interfaces.atmosphere_sea_ice_interface.fluxes diff --git a/src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl b/src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl index d32f21093..848cfe933 100644 --- a/src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl +++ b/src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl @@ -97,12 +97,12 @@ function StateExchanger(ocean::Simulation, atmosphere) # TODO: generalize this exchange_grid = ocean.model.grid exchange_atmosphere_state = ExchangeAtmosphereState(exchange_grid) - exchanger = atmosphere_exchanger(atmosphere, exchange_grid) + exchanger = atmosphere_exchanger(atmosphere, exchange_grid, exchange_atmosphere_state) return StateExchanger(ocean.model.grid, exchange_atmosphere_state, exchanger) end -function atmosphere_exchanger(atmosphere::PrescribedAtmosphere, exchange_grid) +function atmosphere_exchanger(atmosphere::PrescribedAtmosphere, exchange_grid, exchange_atmosphere_state) atmos_grid = atmosphere.grid arch = architecture(exchange_grid) Nx, Ny, Nz = size(exchange_grid) @@ -121,7 +121,7 @@ end initialize!(exchanger::StateExchanger, ::Nothing) = nothing -function initialize!(exchanger::StateExchanger, atmosphere) +function initialize!(exchanger::StateExchanger, atmosphere::PrescribedAtmosphere) atmos_grid = atmosphere.grid exchange_grid = exchanger.exchange_grid arch = architecture(exchange_grid) diff --git a/src/OceanSeaIceModels/time_step_ocean_sea_ice_model.jl b/src/OceanSeaIceModels/time_step_ocean_sea_ice_model.jl index 0c8d3cc47..0ea66d4f5 100644 --- a/src/OceanSeaIceModels/time_step_ocean_sea_ice_model.jl +++ b/src/OceanSeaIceModels/time_step_ocean_sea_ice_model.jl @@ -20,10 +20,10 @@ function time_step!(coupled_model::OceanSeaIceModel, Δt; callbacks=[], compute_ # TODO after ice time-step: # - Adjust ocean heat flux if the ice completely melts? - time_step!(ocean, Δt) + !isnothing(sea_ice) && time_step!(ocean, Δt) # Time step the atmosphere - time_step!(atmosphere, Δt) + !isnothing(sea_ice) && time_step!(atmosphere, Δt) # TODO: # - Store fractional ice-free / ice-covered _time_ for more diff --git a/src/OceanSimulations/ocean_simulation.jl b/src/OceanSimulations/ocean_simulation.jl index f823440fc..4944f0bd5 100644 --- a/src/OceanSimulations/ocean_simulation.jl +++ b/src/OceanSimulations/ocean_simulation.jl @@ -26,8 +26,8 @@ using Statistics: mean @inline v_quadratic_bottom_drag(i, j, grid, c, Φ, μ) = @inbounds - μ * Φ.v[i, j, 1] * spᶜᶠᶜ(i, j, 1, grid, Φ) # Keep a constant linear drag parameter independent on vertical level -@inline u_immersed_bottom_drag(i, j, k, grid, clock, fields, μ) = @inbounds - μ * fields.u[i, j, k] * spᶠᶜᶜ(i, j, k, grid, fields) -@inline v_immersed_bottom_drag(i, j, k, grid, clock, fields, μ) = @inbounds - μ * fields.v[i, j, k] * spᶜᶠᶜ(i, j, k, grid, fields) +@inline u_immersed_bottom_drag(i, j, k, grid, clock, Φ, μ) = @inbounds - μ * Φ.u[i, j, k] * spᶠᶜᶜ(i, j, k, grid, Φ) +@inline v_immersed_bottom_drag(i, j, k, grid, clock, Φ, μ) = @inbounds - μ * Φ.v[i, j, k] * spᶜᶠᶜ(i, j, k, grid, Φ) ##### ##### Defaults diff --git a/test/test_speedy_coupling.jl b/test/test_speedy_coupling.jl new file mode 100644 index 000000000..5c2f15436 --- /dev/null +++ b/test/test_speedy_coupling.jl @@ -0,0 +1,25 @@ +using SpeedyWeather +using ClimaOcean +using Oceananigans +using Dates +using Test + +ClimaOceanSpeedyWeatherExt = Base.get_extension(ClimaOcean, :ClimaOceanSpeedyWeatherExt) +@test !isnothing(ClimaOceanSpeedyWeatherExt) + +spectral_grid = SpeedyWeather.SpectralGrid(trunc=51, nlayers=3) +oceananigans_grid = LatitudeLongitudeGrid(Oceananigans.CPU(); size=(200, 100, 1), latitude=(-80, 80), longitude=(0, 360), z = (0, 1)) + +ocean = ClimaOcean.OceanSimulations.ocean_simulation(oceananigans_grid; momentum_advection=nothing, tracer_advection=nothing, closure=nothing) +atmos = ClimaOcean.atmosphere_simulation(spectral_grid) + +# We initialize the atmosphere simulation with some speed and some temperature +SpeedyWeather.initialize!(atmos) + +earth = OceanSeaIceModel(ocean; atmosphere=atmos) +fluxes = earth.interfaces.atmosphere_ocean_interface.fluxes + +Oceananigans.set!(fluxes.sensible_heat, (x, y) -> exp(- y^2 / 10^2)) + +# Regridding to speedy +ClimaOcean.OceanSeaIceModel.InterfaceComputations.compute_net_atmosphere_fluxes!(earth) \ No newline at end of file