diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 58f6a67b2..e11ae4cc8 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -28,7 +28,7 @@ jobs: fail-fast: false matrix: version: - - '1.10' + - '1.12' os: - ubuntu-latest - windows-latest @@ -37,7 +37,7 @@ jobs: include: - os: macOS-latest arch: aarch64 - version: '1.10' + version: '1.12' steps: - uses: actions/checkout@v5 - uses: julia-actions/setup-julia@latest @@ -65,7 +65,7 @@ jobs: fail-fast: false matrix: version: - - '1.10' + - '1.12' os: - ubuntu-latest # - windows-latest @@ -74,7 +74,7 @@ jobs: include: - os: macOS-latest arch: aarch64 - version: '1.10' + version: '1.12' steps: - uses: actions/checkout@v5 - uses: julia-actions/setup-julia@latest @@ -128,4 +128,39 @@ jobs: with: files: lcov.info token: ${{ secrets.CODECOV_TOKEN }} + + speedy_weather: + name: SpeedyWeather extension - Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ github.event_name }} + runs-on: ${{ matrix.os }} + timeout-minutes: 80 + strategy: + fail-fast: false + matrix: + version: + - '1.12' + os: + - ubuntu-latest + arch: + - x64 + include: + - os: macOS-latest + arch: aarch64 + version: '1.12' + steps: + - uses: actions/checkout@v5 + - uses: julia-actions/setup-julia@latest + with: + version: ${{ matrix.version }} + arch: ${{ matrix.arch }} + - uses: julia-actions/cache@v2 + - uses: julia-actions/julia-buildpkg@v1 + - uses: julia-actions/julia-runtest@v1 + env: + TEST_GROUP: "speedy_weather" + GPU_TEST: "false" + - uses: julia-actions/julia-processcoverage@v1 + - uses: codecov/codecov-action@v5 + with: + files: lcov.info + token: ${{ secrets.CODECOV_TOKEN }} diff --git a/CondaPkg.toml b/CondaPkg.toml deleted file mode 100644 index dde757656..000000000 --- a/CondaPkg.toml +++ /dev/null @@ -1,7 +0,0 @@ - -[pip.deps] -copernicusmarine = ">=2.0.0" -xarray = ">=2024.7.0" -numpy = ">=2.0.0" -jax = ">=0.6" -tensorflow = ">=2.17" diff --git a/Project.toml b/Project.toml index f0e5118e7..ad4f625bb 100644 --- a/Project.toml +++ b/Project.toml @@ -35,10 +35,13 @@ ZipFile = "a5390f91-8eb1-5f08-bee0-b1d1ffed6cea" [weakdeps] CopernicusMarine = "cd43e856-93a3-40c8-bc9e-6146cdce14fa" Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" +SpeedyWeather = "9e226e20-d153-4fed-8a5b-493def4f21a9" +XESMF = "2e0b0046-e7a1-486f-88de-807ee8ffabe5" [extensions] ClimaOceanCopernicusMarineExt = "CopernicusMarine" ClimaOceanReactantExt = "Reactant" +ClimaOceanSpeedyWeatherExt = ["SpeedyWeather", "XESMF"] [compat] Adapt = "4" @@ -66,6 +69,7 @@ SeawaterPolynomials = "0.3.5" StaticArrays = "1" Statistics = "1.9" Thermodynamics = "0.14, 0.15" +XESMF = "0.1.6" ZipFile = "0.10" julia = "1.10" @@ -76,4 +80,4 @@ MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Coverage", "Test", "MPIPreferences", "CUDA_Runtime_jll", "Reactant", "CopernicusMarine"] +test = ["Coverage", "Test", "MPIPreferences", "CUDA_Runtime_jll", "Reactant", "CopernicusMarine", "XESMF", "SpeedyWeather"] diff --git a/docs/Project.toml b/docs/Project.toml index fea2a11e0..d12aa129c 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -2,14 +2,21 @@ CFTime = "179af706-886a-5703-950a-314cd64e0468" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0" +ClimaOcean = "0376089a-ecfe-4b0e-a64f-9c555d74d754" DataDeps = "124859b0-ceae-595e-8997-d05f6a7a8dfe" +Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" DocumenterCitations = "daee34ce-89f3-4625-b898-19384cb65244" Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306" +NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab" Oceananigans = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09" SeawaterPolynomials = "d496a93d-167e-4197-9f49-d3af4ff8fe40" +SpeedyWeather = "9e226e20-d153-4fed-8a5b-493def4f21a9" +XESMF = "2e0b0046-e7a1-486f-88de-807ee8ffabe5" [compat] Documenter = "1" DocumenterCitations = "1.3" +Oceananigans = "0.101" +XESMF = "0.1.6" Literate = "2.2" diff --git a/docs/make.jl b/docs/make.jl index 99afa9e46..03c694967 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 @@ -44,9 +45,10 @@ 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", + # "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/atmosphere_ocean_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..6c8006d52 --- /dev/null +++ b/examples/atmosphere_ocean_simulation.jl @@ -0,0 +1,225 @@ +# # Global coupled atmosphere and ocean--sea ice simulation +# +# This example configures a global ocean--sea ice simulation at 1.5ᵒ 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). +# The XESMF.jl package is used to regrid fields between the atmosphere and ocean--sea ice components. + +using Oceananigans, SpeedyWeather, XESMF, ClimaOcean +using NCDatasets, CairoMakie +using Oceananigans.Units +using Printf, Statistics, Dates + +# ## Ocean and sea-ice model configuration +# The ocean and sea-ice are a simplified version of 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 = 240 +Ny = 120 +Nz = 10 + +r_faces = ExponentialDiscretization(Nz, -2000, 0) +grid = TripolarGrid(Oceananigans.CPU(); size=(Nx, Ny, Nz), z=r_faces, halo=(6, 6, 5)) +nothing # hide + +# Regridding the bathymetry... + +bottom_height = regrid_bathymetry(grid; major_basins=1, interpolation_passes=15) +grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height); active_cells_map=true) +nothing # hide + +# Now we can specify the numerical details and closures for the ocean simulation. + +momentum_advection = VectorInvariant() +tracer_advection = WENO(order=5) +free_surface = SplitExplicitFreeSurface(grid; substeps=40) + +catke_closure = ClimaOcean.OceanSimulations.default_ocean_closure() +viscous_closure = Oceananigans.TurbulenceClosures.HorizontalScalarBiharmonicDiffusivity(ν=1e12) +eddy_closure = Oceananigans.TurbulenceClosures.IsopycnalSkewSymmetricDiffusivity(κ_skew=1e3, κ_symmetric=1e3) +closures = (catke_closure, eddy_closure, viscous_closure, VerticalScalarDiffusivity(ν=1e-4)) +nothing # hide + +# 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. + +sea_ice = sea_ice_simulation(grid, ocean; 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. + +nlayers = 4 +spectral_grid = SpectralGrid(; trunc=63, nlayers, Grid=FullClenshawGrid) +atmosphere = atmosphere_simulation(spectral_grid; output=true) +atmosphere.model.output.output_dt = Hour(3) +nothing # hide + +# ## 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) +nothing # hide + +# 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=30days) +nothing # hide + +# ## Running the simulation +# We can now run the simulation. +# We add callbacks to write outputs to disk every 6 hours. + +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 * 6), + filename="ocean_free_surface.jld2") + +ocean.output_writers[:surface] = JLD2Writer(ocean.model, outputs; + overwrite_existing=true, + schedule=TimeInterval(3600 * 6), + 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 * 6), + 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") + +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. SpeedyWeather outputs are stored in a NetCDF file located in the `run_0001` folder, +# while ocean and sea-ice outputs are stored in JLD2 files that can be read by Oceananigans.jl using the `FieldTimeSeries` type. + +SWO = Dataset("run_0001/output.nc") + +Ta = reverse(SWO["temp"][:, :, nlayers, :], dims=2) +ua = reverse(SWO["u"][:, :, nlayers, :], dims=2) +va = reverse(SWO["v"][:, :, nlayers, :], 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") + +Nt = min(length(sp[1,1,:]), length(Qcao)) + +uotmp = Oceananigans.Field{Face, Center, Nothing}(SST.grid) +votmp = Oceananigans.Field{Center, Face, Nothing}(SST.grid) + +uitmp = Oceananigans.Field{Face, Center, Nothing}(SST.grid) +vitmp = Oceananigans.Field{Center, Face, Nothing}(SST.grid) +atmp = Oceananigans.Field{Center, Center, Nothing}(SST.grid) + +sotmp = Oceananigans.Field(sqrt(uotmp^2 + votmp^2)) +sitmp = Oceananigans.Field(sqrt(uitmp^2 + vitmp^2) * atmp) + +iter = Observable(1) +san = @lift sp[:, :, $iter] +son = @lift begin + Oceananigans.set!(uotmp, SSU[$iter]) + Oceananigans.set!(votmp, SSV[$iter]) + Oceananigans.compute!(sotmp) + Oceananigans.interior(sotmp, :, :, 1) +end + +ssn = @lift begin + Oceananigans.set!(uitmp, SIU[$iter]) + Oceananigans.set!(vitmp, SIV[$iter]) + Oceananigans.set!(atmp, SIA[$iter]) + Oceananigans.compute!(sitmp) + Oceananigans.interior(sitmp, :, :, 1) +end + +fig = Figure(size = (800, 1800)) +ax2 = Axis(fig[1, 1], title = "Surface speed, atmosphere (m/s)") +hm2 = heatmap!(ax2, san; colormap = :deep) +ax1 = Axis(fig[2, 1], title = "Surface speed, ocean (m/s)") +hm = heatmap!(ax1, son; colormap = :deep) +ax3 = Axis(fig[3, 1], title = "Surface speed, sea-ice (m/s)") +hm = heatmap!(ax3, ssn; colormap = :deep) + +record(fig, "surface_speeds.mp4", 1:Nt, framerate=8) do i + iter[] = i +end +nothing #hide + +# ![](surface_speeds.mp4) + +Tan = @lift Ta[:, :, $iter] +Ton = @lift interior(SST[$iter], :, :, 1) +Qcn = @lift interior(Qcao[$iter], :, :, 1) +Qvn = @lift interior(Qvao[$iter], :, :, 1) + +fig = Figure(size = (800, 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:Nt, framerate=8) do i + iter[] = i +end +nothing #hide + +# ![](surface_temperature_and_heat_flux.mp4) diff --git a/examples/single_column_os_papa_simulation.jl b/examples/single_column_os_papa_simulation.jl index 9f07663fc..5c59b4a82 100644 --- a/examples/single_column_os_papa_simulation.jl +++ b/examples/single_column_os_papa_simulation.jl @@ -153,7 +153,7 @@ Q = ρₒ * cₚ * JT ρτx = ρₒ * τx ρτy = ρₒ * τy N² = buoyancy_frequency(ocean.model) -κc = ocean.model.diffusivity_fields.κc +κc = ocean.model.closure_fields.κc fluxes = (; ρτx, ρτy, E, Js, Qv, Qc) auxiliary_fields = (; N², κc) diff --git a/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl b/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl new file mode 100644 index 000000000..302202655 --- /dev/null +++ b/ext/ClimaOceanSpeedyWeatherExt/ClimaOceanSpeedyWeatherExt.jl @@ -0,0 +1,16 @@ +module ClimaOceanSpeedyWeatherExt + +using OffsetArrays +using KernelAbstractions +using Statistics + +import SpeedyWeather +import ClimaOcean +import Oceananigans +import SpeedyWeather.RingGrids + +include("speedy_atmosphere_simulations.jl") +include("speedy_regridder.jl") +include("speedy_weather_exchanger.jl") + +end # module ClimaOceanSpeedyWeatherExt diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl new file mode 100644 index 000000000..c45904c1e --- /dev/null +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl @@ -0,0 +1,85 @@ +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} +const SpeedyNoSeaIceCoupledModel = ClimaOcean.OceanSeaIceModel{<:Union{Nothing, ClimaOcean.FreezingLimitedOceanTemperature}, <: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_regridder.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_regridder.jl new file mode 100644 index 000000000..85b2db32c --- /dev/null +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_regridder.jl @@ -0,0 +1,49 @@ +using Oceananigans.Grids: AbstractGrid +using Oceananigans + +import XESMF: Regridder, xesmf_coordinates + +const Grids = Union{SpeedyWeather.SpectralGrid, AbstractGrid} + +function Regridder(src::Grids, dst::Grids; method::String="bilinear", periodic=true) + src_coords = xesmf_coordinates(src, Center(), Center(), Center()) + dst_coords = xesmf_coordinates(dst, Center(), Center(), Center()) + + return XESMF.Regridder(src_coords, dst_coords; method, periodic) +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 + +xesmf_coordinates(grid::SpeedyWeather.SpectralGrid, args...) = xesmf_coordinates(grid.grid, args...) +xesmf_coordinates(grid::SpeedyWeather.RingGrids.AbstractGrid, args...) = + throw(ArgumentError("xesmf_coordinates not implemented for grid type $(typeof(grid)), maybe you meant to pass a FullGrid?")) + +function xesmf_coordinates(grid::SpeedyWeather.RingGrids.AbstractFullGrid, args...) + 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) + + # Python's xESMF expects 2D arrays with (x, y) coordinates + # in which y varies in dim=1 and x varies in dim=2 + # therefore we transpose the coordinate matrices + coords_dictionary = Dict("lat" => permutedims(lat, (2, 1)), # φ is latitude + "lon" => permutedims(lon, (2, 1)), # λ is longitude + "lat_b" => permutedims(lat_b, (2, 1)), + "lon_b" => permutedims(lon_b, (2, 1))) + + return coords_dictionary +end diff --git a/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl new file mode 100644 index 000000000..5c3a6dab2 --- /dev/null +++ b/ext/ClimaOceanSpeedyWeatherExt/speedy_weather_exchanger.jl @@ -0,0 +1,148 @@ +using Oceananigans +using Oceananigans.BoundaryConditions +using Oceananigans.Grids: architecture +using Oceananigans.Utils: launch! +using Oceananigans.Operators: intrinsic_vector +using XESMF + +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) + + # TODO: Implement a conservative regridder when ready + ocean_atmosphere_regridder = XESMF.Regridder(spectral_grid, exchange_grid) + atmosphere_ocean_regridder = XESMF.Regridder(exchange_grid, spectral_grid) + exchanger = (; ocean_atmosphere_regridder, atmosphere_ocean_regridder) + + return exchanger +end + +@inline (regrid!::XESMF.Regridder)(field::Oceananigans.Field, data::AbstractArray) = regrid!(vec(interior(field)), data) +@inline (regrid!::XESMF.Regridder)(data::AbstractArray, field::Oceananigans.Field) = regrid!(data, vec(interior(field))) + +# Regrid the atmospheric state on the exchange grid +function interpolate_atmosphere_state!(interfaces, atmos::SpeedySimulation, coupled_model) + atmosphere_exchanger = interfaces.exchanger.atmosphere_exchanger + regrid! = atmosphere_exchanger.ocean_atmosphere_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 + Qℓa = atmos.diagnostic_variables.physics.surface_longwave_down.data + Mpa = atmos.diagnostic_variables.physics.total_precipitation_rate.data + + regrid!(exchange_state.u, ua) + regrid!(exchange_state.v, va) + regrid!(exchange_state.T, Ta) + regrid!(exchange_state.q, qa) + regrid!(exchange_state.p, pa) + regrid!(exchange_state.Qs, Qsa) + regrid!(exchange_state.Qℓ, Qℓa) + regrid!(exchange_state.Mp, Mpa) + + arch = architecture(exchange_grid) + + u = exchange_state.u + v = exchange_state.v + + launch!(arch, exchange_grid, :xy, _rotate_winds!, u, v, exchange_grid) + + 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, v, grid) + 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, atmos::SpeedySimulation) + regrid! = coupled_model.interfaces.exchanger.atmosphere_exchanger.atmosphere_ocean_regridder + 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 + ℵ = interior(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 + # TODO: This does not work on GPUs!! + regrid!(Qca, vec(interior(Qco) .* (1 .- ℵ) .+ ℵ .* interior(Qci))) + regrid!(Mva, vec(interior(Mvo) .* (1 .- ℵ) .+ ℵ .* interior(Mvi))) + regrid!(sst, vec(interior(To) .* (1 .- ℵ) .+ ℵ .* interior(Ti) .+ 273.15)) + + return nothing +end + +# Simple case -> there is no sea ice! +function compute_net_atmosphere_fluxes!(coupled_model::SpeedyNoSeaIceCoupledModel, atmos::SpeedySimulation) + regrid! = coupled_model.interfaces.exchanger.atmosphere_exchanger.atmosphere_ocean_regridder + ao_fluxes = coupled_model.interfaces.atmosphere_ocean_interface.fluxes + Qco = ao_fluxes.sensible_heat + Mvo = ao_fluxes.water_vapor + + # 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 + + # TODO: Figure out how we are going to deal with upwelling radiation + # TODO: This does not work on GPUs!! + regrid!(Qca, vec(interior(Qco))) + regrid!(Mva, vec(interior(Mvo))) + regrid!(sst, vec(interior(To) .+ 273.15)) + + return nothing +end diff --git a/src/ClimaOcean.jl b/src/ClimaOcean.jl index 7ac5f01cb..3759bca57 100644 --- a/src/ClimaOcean.jl +++ b/src/ClimaOcean.jl @@ -42,6 +42,7 @@ export DatasetRestoring, ocean_simulation, sea_ice_simulation, + atmosphere_simulation, initialize! using Oceananigans @@ -78,6 +79,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/InterfaceComputations.jl b/src/OceanSeaIceModels/InterfaceComputations/InterfaceComputations.jl index 10f8c71be..72d8d8af4 100644 --- a/src/OceanSeaIceModels/InterfaceComputations/InterfaceComputations.jl +++ b/src/OceanSeaIceModels/InterfaceComputations/InterfaceComputations.jl @@ -23,6 +23,14 @@ using ..OceanSeaIceModels: default_gravitational_acceleration, import ClimaOcean: stateindex +import ..OceanSeaIceModels: + compute_net_atmosphere_fluxes!, + compute_net_sea_ice_fluxes!, + compute_net_ocean_fluxes!, + compute_atmosphere_ocean_fluxes!, + compute_atmosphere_sea_ice_fluxes!, + compute_sea_ice_ocean_fluxes! + ##### ##### Utilities ##### @@ -70,6 +78,7 @@ include("interpolate_atmospheric_state.jl") include("atmosphere_ocean_fluxes.jl") include("atmosphere_sea_ice_fluxes.jl") include("sea_ice_ocean_fluxes.jl") -include("assemble_net_fluxes.jl") +include("assemble_net_ocean_fluxes.jl") +include("assemble_net_sea_ice_fluxes.jl") end # module diff --git a/src/OceanSeaIceModels/InterfaceComputations/assemble_net_fluxes.jl b/src/OceanSeaIceModels/InterfaceComputations/assemble_net_ocean_fluxes.jl similarity index 61% rename from src/OceanSeaIceModels/InterfaceComputations/assemble_net_fluxes.jl rename to src/OceanSeaIceModels/InterfaceComputations/assemble_net_ocean_fluxes.jl index 1532a3be7..2292fbf1f 100644 --- a/src/OceanSeaIceModels/InterfaceComputations/assemble_net_fluxes.jl +++ b/src/OceanSeaIceModels/InterfaceComputations/assemble_net_ocean_fluxes.jl @@ -2,14 +2,7 @@ using Printf using Oceananigans.Operators: ℑxᶠᵃᵃ, ℑyᵃᶠᵃ using Oceananigans.Forcings: MultipleForcings -using ClimaOcean.OceanSeaIceModels: sea_ice_concentration - -@inline computed_sea_ice_ocean_fluxes(interface) = interface.fluxes -@inline computed_sea_ice_ocean_fluxes(::Nothing) = (interface_heat = ZeroField(), - frazil_heat = ZeroField(), - salt = ZeroField(), - x_momentum = ZeroField(), - y_momentum = ZeroField()) +using ClimaOcean.OceanSeaIceModels: sea_ice_concentration, NoAtmosphereModel, NoSeaIceModel @inline shortwave_radiative_forcing(i, j, grid, Fᵀ, Qts, ocean_properties) = Qts @@ -29,8 +22,17 @@ function get_radiative_forcing(FT::MultipleForcings) return nothing end -function compute_net_ocean_fluxes!(coupled_model) - ocean = coupled_model.ocean +@inline τᶜᶜᶜ(i, j, k, grid, ρₒ⁻¹, ℵ, ρτᶜᶜᶜ) = @inbounds ρₒ⁻¹ * (1 - ℵ[i, j, k]) * ρτᶜᶜᶜ[i, j, k] + +##### +##### Generic flux assembler +##### + +computed_sea_ice_ocean_fluxes(sea_ice_ocean_interface) = sea_ice_ocean_interface.fluxes +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 arch = architecture(grid) @@ -83,7 +85,8 @@ function compute_net_ocean_fluxes!(coupled_model) return nothing end -@inline τᶜᶜᶜ(i, j, k, grid, ρₒ⁻¹, ℵ, ρτᶜᶜᶜ) = @inbounds ρₒ⁻¹ * (1 - ℵ[i, j, k]) * ρτᶜᶜᶜ[i, j, k] +@inline get_possibly_zero_flux(fluxes, name) = getfield(fluxes, name) +@inline get_possibly_zero_flux(::Nothing, name) = ZeroField() @kernel function _assemble_net_ocean_fluxes!(net_ocean_fluxes, penetrating_radiation, @@ -104,8 +107,8 @@ end time = Time(clock.time) ρτxao = atmos_ocean_fluxes.x_momentum # atmosphere - ocean zonal momentum flux ρτyao = atmos_ocean_fluxes.y_momentum # atmosphere - ocean meridional momentum flux - ρτxio = sea_ice_ocean_fluxes.x_momentum # sea_ice - ocean zonal momentum flux - ρτyio = sea_ice_ocean_fluxes.y_momentum # sea_ice - ocean meridional momentum flux + ρτxio = get_possibly_zero_flux(sea_ice_ocean_fluxes, :x_momentum) # sea_ice - ocean zonal momentum flux + ρτyio = get_possibly_zero_flux(sea_ice_ocean_fluxes, :y_momentum) # sea_ice - ocean meridional momentum flux @inbounds begin ℵᵢ = sea_ice_concentration[i, j, 1] @@ -117,15 +120,15 @@ 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) σ = atmos_ocean_properties.radiation.σ α = atmos_ocean_properties.radiation.α ϵ = atmos_ocean_properties.radiation.ϵ - Qu = emitted_longwave_radiation(i, j, kᴺ, grid, time, Tₛ, σ, ϵ) + Qu = emitted_longwave_radiation(i, j, kᴺ, grid, time, Tₛ, σ, ϵ) Qaℓ = absorbed_longwave_radiation(i, j, kᴺ, grid, time, ϵ, Qℓ) # Compute the interior + surface absorbed shortwave radiation @@ -167,12 +170,12 @@ end cₒ = ocean_properties.heat_capacity @inbounds begin - Qio = sea_ice_ocean_fluxes.interface_heat[i, j, 1] + Qio = get_possibly_zero_flux(sea_ice_ocean_fluxes, :interface_heat)[i, j, 1] Jᵀao = ΣQao * ρₒ⁻¹ / cₒ - Jˢao = - Sₒ * ΣFao + Jˢao = - Sₒ * ΣFao # salinity flux > 0 extracts salinity from the ocean --- the opposite of a water vapor flux Jᵀio = Qio * ρₒ⁻¹ / cₒ - Jˢio = sea_ice_ocean_fluxes.salt[i, j, 1] * ℵᵢ + Jˢio = get_possibly_zero_flux(sea_ice_ocean_fluxes, :salt)[i, j, 1] * ℵᵢ τxao = ℑxᶠᵃᵃ(i, j, 1, grid, τᶜᶜᶜ, ρₒ⁻¹, ℵ, ρτxao) τyao = ℑyᵃᶠᵃ(i, j, 1, grid, τᶜᶜᶜ, ρₒ⁻¹, ℵ, ρτyao) @@ -189,107 +192,78 @@ end end end -function compute_net_sea_ice_fluxes!(coupled_model) - sea_ice = coupled_model.sea_ice +##### +##### No atmosphere implementation +##### - if !(sea_ice isa SeaIceSimulation) - return nothing - end - - ocean = coupled_model.ocean - grid = ocean.model.grid - arch = architecture(grid) +# A generic ocean flux assembler for a coupled model with both an atmosphere and sea ice +function compute_net_ocean_fluxes!(coupled_model::NoAtmosphereModel, ocean) + sea_ice = coupled_model.sea_ice + grid = ocean.model.grid + arch = architecture(grid) clock = coupled_model.clock - top_fluxes = coupled_model.interfaces.net_fluxes.sea_ice_top - bottom_heat_flux = coupled_model.interfaces.net_fluxes.sea_ice_bottom.heat + net_ocean_fluxes = coupled_model.interfaces.net_fluxes.ocean_surface sea_ice_ocean_fluxes = coupled_model.interfaces.sea_ice_ocean_interface.fluxes - atmosphere_sea_ice_fluxes = coupled_model.interfaces.atmosphere_sea_ice_interface.fluxes - - # Simplify NamedTuple to reduce parameter space consumption. - # See https://github.com/CliMA/ClimaOcean.jl/issues/116. - atmosphere_fields = coupled_model.interfaces.exchanger.exchange_atmosphere_state - downwelling_radiation = (Qs = atmosphere_fields.Qs.data, - Qℓ = atmosphere_fields.Qℓ.data) - - freshwater_flux = atmosphere_fields.Mp.data - - atmos_sea_ice_properties = coupled_model.interfaces.atmosphere_sea_ice_interface.properties - sea_ice_properties = coupled_model.interfaces.sea_ice_properties + # We remove the heat flux since does not need to be assembled and bloats the parameter space. + net_ocean_fluxes = (u = net_ocean_fluxes.u, + v = net_ocean_fluxes.v, + T = net_ocean_fluxes.T, + S = net_ocean_fluxes.S) + ice_concentration = sea_ice.model.ice_concentration + ocean_properties = coupled_model.interfaces.ocean_properties kernel_parameters = interface_kernel_parameters(grid) - sea_ice_surface_temperature = coupled_model.interfaces.atmosphere_sea_ice_interface.temperature - ice_concentration = sea_ice_concentration(sea_ice) - launch!(arch, grid, kernel_parameters, - _assemble_net_sea_ice_fluxes!, - top_fluxes, - bottom_heat_flux, + _assemble_no_atmosphere_net_ocean_fluxes!, + net_ocean_fluxes, grid, clock, - atmosphere_sea_ice_fluxes, sea_ice_ocean_fluxes, - freshwater_flux, ice_concentration, - sea_ice_surface_temperature, - downwelling_radiation, - sea_ice_properties, - atmos_sea_ice_properties) + ocean_properties) return nothing end -@kernel function _assemble_net_sea_ice_fluxes!(top_fluxes, - bottom_heat_flux, - grid, - clock, - atmosphere_sea_ice_fluxes, - sea_ice_ocean_fluxes, - freshwater_flux, # Where do we add this one? - ice_concentration, - surface_temperature, - downwelling_radiation, - sea_ice_properties, - atmos_sea_ice_properties) +@kernel function _assemble_no_atmosphere_net_ocean_fluxes!(net_ocean_fluxes, + grid, + clock, + sea_ice_ocean_fluxes, + sea_ice_concentration, + ocean_properties) i, j = @index(Global, NTuple) kᴺ = size(grid, 3) - time = Time(clock.time) - - @inbounds begin - Ts = surface_temperature[i, j, kᴺ] - Ts = convert_to_kelvin(sea_ice_properties.temperature_units, Ts) - ℵi = ice_concentration[i, j, 1] - - Qs = downwelling_radiation.Qs[i, j, 1] - Qℓ = downwelling_radiation.Qℓ[i, j, 1] - Qc = atmosphere_sea_ice_fluxes.sensible_heat[i, j, 1] # sensible or "conductive" heat flux - Qv = atmosphere_sea_ice_fluxes.latent_heat[i, j, 1] # latent heat flux - Qf = sea_ice_ocean_fluxes.frazil_heat[i, j, 1] # frazil heat flux - Qi = sea_ice_ocean_fluxes.interface_heat[i, j, 1] # interfacial heat flux - end + ρτxio = sea_ice_ocean_fluxes.x_momentum # sea_ice - ocean zonal momentum flux + ρτyio = sea_ice_ocean_fluxes.y_momentum # sea_ice - ocean meridional momentum flux - ρτx = atmosphere_sea_ice_fluxes.x_momentum # zonal momentum flux - ρτy = atmosphere_sea_ice_fluxes.y_momentum # meridional momentum flux + @inbounds ℵᵢ = sea_ice_concentration[i, j, 1] - # Compute radiation fluxes - σ = atmos_sea_ice_properties.radiation.σ - α = atmos_sea_ice_properties.radiation.α - ϵ = atmos_sea_ice_properties.radiation.ϵ - Qu = emitted_longwave_radiation(i, j, kᴺ, grid, time, Ts, σ, ϵ) - Qs = transmitted_shortwave_radiation(i, j, kᴺ, grid, time, α, Qs) - Qℓ = absorbed_longwave_radiation(i, j, kᴺ, grid, time, ϵ, Qℓ) + # Compute fluxes for u, v, T, and S from momentum, heat, and freshwater fluxes + τx = net_ocean_fluxes.u + τy = net_ocean_fluxes.v + Jᵀ = net_ocean_fluxes.T + Jˢ = net_ocean_fluxes.S + ℵ = sea_ice_concentration + ρₒ⁻¹ = 1 / ocean_properties.reference_density + cₒ = ocean_properties.heat_capacity - ΣQt = (Qs + Qℓ + Qu + Qc + Qv) * (ℵi > 0) # If ℵi == 0 there is no heat flux from the top! - ΣQb = Qf + Qi + @inbounds begin + Qio = sea_ice_ocean_fluxes.interface_heat[i, j, 1] + Jᵀio = Qio * ρₒ⁻¹ / cₒ + Jˢio = sea_ice_ocean_fluxes.salt[i, j, 1] * ℵᵢ + τxio = ρτxio[i, j, 1] * ρₒ⁻¹ * ℑxᶠᵃᵃ(i, j, 1, grid, ℵ) + τyio = ρτyio[i, j, 1] * ρₒ⁻¹ * ℑyᵃᶠᵃ(i, j, 1, grid, ℵ) - # Mask fluxes over land for convenience - inactive = inactive_node(i, j, kᴺ, grid, c, c, c) + # Stresses + τx[i, j, 1] = τxio + τy[i, j, 1] = τyio - @inbounds top_fluxes.heat[i, j, 1] = ifelse(inactive, zero(grid), ΣQt) - @inbounds top_fluxes.u[i, j, 1] = ifelse(inactive, zero(grid), ℑxᶠᵃᵃ(i, j, 1, grid, ρτx)) - @inbounds top_fluxes.v[i, j, 1] = ifelse(inactive, zero(grid), ℑyᵃᶠᵃ(i, j, 1, grid, ρτy)) - @inbounds bottom_heat_flux[i, j, 1] = ifelse(inactive, zero(grid), ΣQb) + # Tracer fluxes + Jᵀ[i, j, 1] = Jᵀio # Jᵀao is already multiplied by the sea ice concentration + Jˢ[i, j, 1] = Jˢio + end end diff --git a/src/OceanSeaIceModels/InterfaceComputations/assemble_net_sea_ice_fluxes.jl b/src/OceanSeaIceModels/InterfaceComputations/assemble_net_sea_ice_fluxes.jl new file mode 100644 index 000000000..eed1c4c94 --- /dev/null +++ b/src/OceanSeaIceModels/InterfaceComputations/assemble_net_sea_ice_fluxes.jl @@ -0,0 +1,98 @@ +function compute_net_sea_ice_fluxes!(coupled_model, sea_ice::SeaIceSimulation) + ocean = coupled_model.ocean + grid = ocean.model.grid + arch = architecture(grid) + clock = coupled_model.clock + + top_fluxes = coupled_model.interfaces.net_fluxes.sea_ice_top + bottom_heat_flux = coupled_model.interfaces.net_fluxes.sea_ice_bottom.heat + sea_ice_ocean_fluxes = coupled_model.interfaces.sea_ice_ocean_interface.fluxes + atmosphere_sea_ice_fluxes = coupled_model.interfaces.atmosphere_sea_ice_interface.fluxes + + # Simplify NamedTuple to reduce parameter space consumption. + # See https://github.com/CliMA/ClimaOcean.jl/issues/116. + atmosphere_fields = coupled_model.interfaces.exchanger.exchange_atmosphere_state + + downwelling_radiation = (Qs = atmosphere_fields.Qs.data, + Qℓ = atmosphere_fields.Qℓ.data) + + freshwater_flux = atmosphere_fields.Mp.data + + atmos_sea_ice_properties = coupled_model.interfaces.atmosphere_sea_ice_interface.properties + sea_ice_properties = coupled_model.interfaces.sea_ice_properties + + kernel_parameters = interface_kernel_parameters(grid) + + sea_ice_surface_temperature = coupled_model.interfaces.atmosphere_sea_ice_interface.temperature + ice_concentration = sea_ice_concentration(sea_ice) + + launch!(arch, grid, kernel_parameters, + _assemble_net_sea_ice_fluxes!, + top_fluxes, + bottom_heat_flux, + grid, + clock, + atmosphere_sea_ice_fluxes, + sea_ice_ocean_fluxes, + freshwater_flux, + ice_concentration, + sea_ice_surface_temperature, + downwelling_radiation, + sea_ice_properties, + atmos_sea_ice_properties) + + return nothing +end + +@kernel function _assemble_net_sea_ice_fluxes!(top_fluxes, + bottom_heat_flux, + grid, + clock, + atmosphere_sea_ice_fluxes, + sea_ice_ocean_fluxes, + freshwater_flux, # Where do we add this one? + ice_concentration, + surface_temperature, + downwelling_radiation, + sea_ice_properties, + atmos_sea_ice_properties) + + i, j = @index(Global, NTuple) + kᴺ = size(grid, 3) + time = Time(clock.time) + + @inbounds begin + Ts = surface_temperature[i, j, kᴺ] + Ts = convert_to_kelvin(sea_ice_properties.temperature_units, Ts) + ℵi = ice_concentration[i, j, 1] + + Qs = downwelling_radiation.Qs[i, j, 1] + Qℓ = downwelling_radiation.Qℓ[i, j, 1] + Qc = atmosphere_sea_ice_fluxes.sensible_heat[i, j, 1] # sensible or "conductive" heat flux + Qv = atmosphere_sea_ice_fluxes.latent_heat[i, j, 1] # latent heat flux + Qf = sea_ice_ocean_fluxes.frazil_heat[i, j, 1] # frazil heat flux + Qi = sea_ice_ocean_fluxes.interface_heat[i, j, 1] # interfacial heat flux + end + + ρτx = atmosphere_sea_ice_fluxes.x_momentum # zonal momentum flux + ρτy = atmosphere_sea_ice_fluxes.y_momentum # meridional momentum flux + + # Compute radiation fluxes + σ = atmos_sea_ice_properties.radiation.σ + α = atmos_sea_ice_properties.radiation.α + ϵ = atmos_sea_ice_properties.radiation.ϵ + Qu = emitted_longwave_radiation(i, j, kᴺ, grid, time, Ts, σ, ϵ) + Qs = transmitted_shortwave_radiation(i, j, kᴺ, grid, time, α, Qs) + Qℓ = absorbed_longwave_radiation(i, j, kᴺ, grid, time, ϵ, Qℓ) + + ΣQt = (Qs + Qℓ + Qu + Qc + Qv) * (ℵi > 0) # If ℵi == 0 there is no heat flux from the top! + ΣQb = Qf + Qi + + # Mask fluxes over land for convenience + inactive = inactive_node(i, j, kᴺ, grid, c, c, c) + + @inbounds top_fluxes.heat[i, j, 1] = ifelse(inactive, zero(grid), ΣQt) + @inbounds top_fluxes.u[i, j, 1] = ifelse(inactive, zero(grid), ℑxᶠᵃᵃ(i, j, 1, grid, ρτx)) + @inbounds top_fluxes.v[i, j, 1] = ifelse(inactive, zero(grid), ℑyᵃᶠᵃ(i, j, 1, grid, ρτy)) + @inbounds bottom_heat_flux[i, j, 1] = ifelse(inactive, zero(grid), ΣQb) +end 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 592794f0e..215a9bcaa 100644 --- a/src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl +++ b/src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl @@ -93,12 +93,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) @@ -117,7 +117,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/OceanSeaIceModels.jl b/src/OceanSeaIceModels/OceanSeaIceModels.jl index 49a6ad740..a5226076b 100644 --- a/src/OceanSeaIceModels/OceanSeaIceModels.jl +++ b/src/OceanSeaIceModels/OceanSeaIceModels.jl @@ -31,6 +31,18 @@ using ClimaOcean: stateindex using KernelAbstractions: @kernel, @index using KernelAbstractions.Extras.LoopInfo: @unroll +import Thermodynamics as AtmosphericThermodynamics + +# Simulations interface +import Oceananigans: fields, prognostic_fields +import Oceananigans.Architectures: architecture +import Oceananigans.Fields: set! +import Oceananigans.Models: timestepper, NaNChecker, default_nan_checker, initialization_update_state! +import Oceananigans.OutputWriters: default_included_properties +import Oceananigans.Simulations: reset!, initialize!, iteration +import Oceananigans.TimeSteppers: time_step!, update_state!, time +import Oceananigans.Utils: prettytime + function downwelling_radiation end function freshwater_flux end function reference_density end @@ -39,7 +51,9 @@ function heat_capacity end const default_gravitational_acceleration = Oceananigans.defaults.gravitational_acceleration const default_freshwater_density = 1000 # kg m⁻³ +# Our default ocean and sea ice models const SeaIceSimulation = Simulation{<:SeaIceModel} +const OceananigansSimulation = Simulation{<:HydrostaticFreeSurfaceModel} sea_ice_thickness(::Nothing) = ZeroField() sea_ice_thickness(sea_ice::SeaIceSimulation) = sea_ice.model.ice_thickness @@ -47,13 +61,51 @@ sea_ice_thickness(sea_ice::SeaIceSimulation) = sea_ice.model.ice_thickness sea_ice_concentration(::Nothing) = ZeroField() sea_ice_concentration(sea_ice::SeaIceSimulation) = sea_ice.model.ice_concentration +mutable struct OceanSeaIceModel{I, A, O, F, C, Arch} <: AbstractModel{Nothing, Arch} + architecture :: Arch + clock :: C + atmosphere :: A + sea_ice :: I + ocean :: O + interfaces :: F +end + +struct FreezingLimitedOceanTemperature{L} + liquidus :: L +end + +const OSIM = OceanSeaIceModel +const NoAtmosphereModel = OceanSeaIceModel{<:Any, Nothing} +const NoSeaIceModel = Union{OceanSeaIceModel{Nothing}, OceanSeaIceModel{<:FreezingLimitedOceanTemperature}} + ##### ##### Some implementation ##### # Atmosphere interface interpolate_atmosphere_state!(interfaces, atmosphere, coupled_model) = nothing -compute_net_atmosphere_fluxes!(coupled_model) = nothing + +# Compute net fluxes: +compute_net_sea_ice_fluxes!(coupled_model, ::Nothing) = nothing +compute_net_ocean_fluxes!(coupled_model, ::Nothing) = nothing +compute_net_atmosphere_fluxes!(coupled_model, ::Nothing) = nothing + +# "No atmosphere" implementation +compute_atmosphere_ocean_fluxes!(::NoAtmosphereModel) = nothing +compute_atmosphere_sea_ice_fluxes!(::NoAtmosphereModel) = nothing + +# "No sea ice" implementation +compute_sea_ice_ocean_fluxes!(::OceanSeaIceModel{Nothing}) = nothing +compute_atmosphere_sea_ice_fluxes!(::NoSeaIceModel) = nothing + +# "Only ocean" implementation +const OnlyOceanModel = Union{OceanSeaIceModel{Nothing, Nothing}, OceanSeaIceModel{<:FreezingLimitedOceanTemperature, Nothing}} + +compute_atmosphere_sea_ice_fluxes!(::OnlyOceanModel) = nothing +compute_sea_ice_ocean_fluxes!(::OnlyOceanModel) = nothing +compute_net_ocean_fluxes!(::OnlyOceanModel, ocean) = nothing + +include("freezing_limited_ocean_temperature.jl") # TODO: import this last include("PrescribedAtmospheres.jl") @@ -67,35 +119,7 @@ include("InterfaceComputations/InterfaceComputations.jl") using .InterfaceComputations -import .InterfaceComputations: - compute_atmosphere_ocean_fluxes!, - compute_atmosphere_sea_ice_fluxes!, - compute_net_ocean_fluxes!, - compute_sea_ice_ocean_fluxes! - include("ocean_sea_ice_model.jl") -include("freezing_limited_ocean_temperature.jl") include("time_step_ocean_sea_ice_model.jl") -# "No atmosphere" implementation -const NoAtmosphereModel = OceanSeaIceModel{<:Any, Nothing} -compute_atmosphere_ocean_fluxes!(::NoAtmosphereModel) = nothing -compute_atmosphere_sea_ice_fluxes!(::NoAtmosphereModel) = nothing - -const PrescribedAtmosphereModel = OceanSeaIceModel{<:Any, <:PrescribedAtmosphere} -compute_net_atmosphere_fluxes!(::PrescribedAtmosphereModel) = nothing - -# "No sea ice" implementation -const NoSeaIceModel = Union{OceanSeaIceModel{Nothing}, FreezingLimitedCoupledModel} -compute_sea_ice_ocean_fluxes!(::OceanSeaIceModel{Nothing}) = nothing -compute_atmosphere_sea_ice_fluxes!(::NoSeaIceModel) = nothing - -# "Only ocean" implementation -const OnlyOceanModel = Union{OceanSeaIceModel{Nothing, Nothing}, - OceanSeaIceModel{<:FreezingLimitedOceanTemperature, Nothing}} - -compute_atmosphere_sea_ice_fluxes!(::OnlyOceanModel) = nothing -compute_sea_ice_ocean_fluxes!(::OnlyOceanModel) = nothing -compute_net_ocean_fluxes!(::OnlyOceanModel) = nothing - end # module diff --git a/src/OceanSeaIceModels/PrescribedAtmospheres.jl b/src/OceanSeaIceModels/PrescribedAtmospheres.jl index e7a3ca18c..49d72d32d 100644 --- a/src/OceanSeaIceModels/PrescribedAtmospheres.jl +++ b/src/OceanSeaIceModels/PrescribedAtmospheres.jl @@ -42,7 +42,8 @@ import Thermodynamics.Parameters: import ..OceanSeaIceModels: downwelling_radiation, - freshwater_flux + freshwater_flux, + compute_net_atmosphere_fluxes! ##### ##### Atmospheric thermodynamics parameters @@ -376,6 +377,9 @@ end @inline surface_layer_height(atmos::PrescribedAtmosphere) = atmos.surface_layer_height @inline boundary_layer_height(atmos::PrescribedAtmosphere) = atmos.boundary_layer_height +# No need to compute anything here... +compute_net_atmosphere_fluxes!(coupled_model, ::PrescribedAtmosphere) = nothing + """ PrescribedAtmosphere(grid, times=[zero(grid)]; clock = Clock{Float64}(time = 0), diff --git a/src/OceanSeaIceModels/freezing_limited_ocean_temperature.jl b/src/OceanSeaIceModels/freezing_limited_ocean_temperature.jl index b9475c895..d64bb917c 100644 --- a/src/OceanSeaIceModels/freezing_limited_ocean_temperature.jl +++ b/src/OceanSeaIceModels/freezing_limited_ocean_temperature.jl @@ -4,10 +4,6 @@ using ClimaSeaIce.SeaIceThermodynamics: LinearLiquidus ##### A workaround when you don't have a sea ice model ##### -struct FreezingLimitedOceanTemperature{L} - liquidus :: L -end - """ FreezingLimitedOceanTemperature(FT=Float64; liquidus=LinearLiquidus(FT)) @@ -31,6 +27,9 @@ reference_density(::FreezingLimitedOceanTemperature) = 0 heat_capacity(::FreezingLimitedOceanTemperature) = 0 time_step!(::FreezingLimitedOceanTemperature, Δt) = nothing +# No need to compute fluxes for this "sea ice model" +compute_net_sea_ice_fluxes!(coupled_model, ::FreezingLimitedOceanTemperature) = nothing + function compute_sea_ice_ocean_fluxes!(cm::FreezingLimitedCoupledModel) ocean = cm.ocean liquidus = cm.sea_ice.liquidus diff --git a/src/OceanSeaIceModels/ocean_sea_ice_model.jl b/src/OceanSeaIceModels/ocean_sea_ice_model.jl index 6442d578a..06f849793 100644 --- a/src/OceanSeaIceModels/ocean_sea_ice_model.jl +++ b/src/OceanSeaIceModels/ocean_sea_ice_model.jl @@ -5,29 +5,6 @@ using ClimaSeaIce.SeaIceThermodynamics: melting_temperature using KernelAbstractions: @kernel, @index using SeawaterPolynomials: TEOS10EquationOfState -import Thermodynamics as AtmosphericThermodynamics - -# Simulations interface -import Oceananigans: fields, prognostic_fields -import Oceananigans.Architectures: architecture -import Oceananigans.Fields: set! -import Oceananigans.Models: timestepper, NaNChecker, default_nan_checker, initialization_update_state! -import Oceananigans.OutputWriters: default_included_properties -import Oceananigans.Simulations: reset!, initialize!, iteration -import Oceananigans.TimeSteppers: time_step!, update_state!, time -import Oceananigans.Utils: prettytime - -mutable struct OceanSeaIceModel{I, A, O, F, C, Arch} <: AbstractModel{Nothing, Arch} - architecture :: Arch - clock :: C - atmosphere :: A - sea_ice :: I - ocean :: O - interfaces :: F -end - -const OSIM = OceanSeaIceModel - function Base.summary(model::OSIM) A = nameof(typeof(architecture(model))) return string("OceanSeaIceModel{$A}", diff --git a/src/OceanSeaIceModels/time_step_ocean_sea_ice_model.jl b/src/OceanSeaIceModels/time_step_ocean_sea_ice_model.jl index 0c8d3cc47..1428f88e2 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 @@ -36,18 +36,24 @@ end function update_state!(coupled_model::OceanSeaIceModel, callbacks=[]; compute_tendencies=true) + # The three components + ocean = coupled_model.ocean + sea_ice = coupled_model.sea_ice + atmosphere = coupled_model.atmosphere + + # TODO: add a `interpolate_ocean_state!` function if needed # This function needs to be specialized to allow different atmospheric models - interpolate_atmosphere_state!(coupled_model.interfaces, coupled_model.atmosphere, coupled_model) + interpolate_atmosphere_state!(coupled_model.interfaces, atmosphere, coupled_model) # Compute interface states compute_atmosphere_ocean_fluxes!(coupled_model) compute_atmosphere_sea_ice_fluxes!(coupled_model) compute_sea_ice_ocean_fluxes!(coupled_model) - # This function needs to be specialized to allow different atmospheric models - compute_net_atmosphere_fluxes!(coupled_model) - compute_net_ocean_fluxes!(coupled_model) - compute_net_sea_ice_fluxes!(coupled_model) + # This function needs to be specialized to allow different component models + compute_net_atmosphere_fluxes!(coupled_model, atmosphere) + compute_net_ocean_fluxes!(coupled_model, ocean) + compute_net_sea_ice_fluxes!(coupled_model, sea_ice) return nothing end diff --git a/src/OceanSimulations/ocean_simulation.jl b/src/OceanSimulations/ocean_simulation.jl index 9a34ccfbf..cd2094cae 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/runtests.jl b/test/runtests.jl index 5e7a0000f..18d9e8c54 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -114,3 +114,7 @@ end if test_group == :reactant || test_group == :all include("test_reactant.jl") end + +if test_group == :speedy_weather || test_group == :all + include("test_speedy_coupling.jl") +end \ No newline at end of file diff --git a/test/test_speedy_coupling.jl b/test/test_speedy_coupling.jl new file mode 100644 index 000000000..1c9e77993 --- /dev/null +++ b/test/test_speedy_coupling.jl @@ -0,0 +1,25 @@ +using SpeedyWeather, XESMF +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, Grid=FullClenshawGrid) +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) +Oceananigans.set!(ocean.model, T=EN4Metadatum(:temperature), S=EN4Metadatum(:salinity)) + +atmos = ClimaOcean.atmosphere_simulation(spectral_grid) + +radiation = Radiation(ocean_emissivity=0.0, sea_ice_emissivity=0.0) +earth_model = OceanSeaIceModel(ocean; atmosphere=atmos, radiation) + +Qca = atmos.prognostic_variables.ocean.sensible_heat_flux.data +Mva = atmos.prognostic_variables.ocean.surface_humidity_flux.data + +@test !(all(Qca .== 0.0)) +@test !(all(Mva .== 0.0)) \ No newline at end of file