Skip to content

Run prescribed SST example on GPU #387

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
navidcy wants to merge 7 commits into
base: main
Choose a base branch
from
Open

Run prescribed SST example on GPU #387

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
970d385
Run prescribed SST example on GPU
navidcy Jan 6, 2026
af3235e
Import CUDA for enhanced performance
glwagner Jan 6, 2026
416f9de
temporary adapt for BoundaryConditionKernelFunction
navidcy Jan 6, 2026
4998a35
fix adapt
navidcy Jan 6, 2026
7564e06
Update examples/prescribed_sea_surface_temperature.jl
navidcy Jan 6, 2026
91120e3
Apply suggestions from code review
navidcy Jan 6, 2026
ab402f3
Reintroduce GPU adaptation for BoundaryConditionKernelFunction
navidcy Jan 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions docs/Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
[deps]
Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
AtmosphericProfilesLibrary = "86bc3604-9858-485a-bdbe-831ec50de11d"
Breeze = "660aa2fb-d4c8-4359-a52c-9c057bc511da"
CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
Expand Down
19 changes: 16 additions & 3 deletions examples/prescribed_sea_surface_temperature.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ using Oceananigans
using Oceananigans.Units
using Oceananigans.Models: BoundaryConditionOperation
using Printf
using CUDA
using CairoMakie

# ## Grid setup
Expand All @@ -36,7 +37,7 @@ using CairoMakie
# energy-containing scales of convective turbulence while remaining computationally
# tractable for this demonstration.

grid = RectilinearGrid(size = (128, 128), halo = (5, 5),
grid = RectilinearGrid(GPU(), size = (128, 128), halo = (5, 5),
x = (-10kilometers, 10kilometers),
z = (0, 10kilometers),
topology = (Periodic, Flat, Bounded))
Expand Down Expand Up @@ -111,9 +112,9 @@ Uᵍ = 1e-2 # Minimum wind speed (m/s)
#
# For `BulkVaporFlux`, the saturation specific humidity is computed from the surface
# temperature. Surface temperature can be provided as a `Field`, a `Function`, or a `Number`.
#
#
# In this example, we specify the sea surface temperature as a top hat function
# i.e. representing a pair of ocean fronts in a periodic domain, with a
# i.e. representing a pair of ocean fronts in a periodic domain, with a
# difference of 4 degrees K,

ΔT = 4 # K
Expand Down Expand Up @@ -188,6 +189,18 @@ qᵗ = model.specific_moisture
# values from the boundary conditions. These 1D fields (varying only in x)
# represent the actual flux values applied at the ocean-atmosphere interface.
#
# We need to adapt the `BoundaryConditionKernelFunction` for the GPU.
# Until [Oceananigans PR#5100](https://github.com/CliMA/Oceananigans.jl/pull/5100) is included
# in a tagged release, we can use Adapt here directly.

using Adapt: Adapt

function Adapt.adapt_structure(to, bckf::Oceananigans.Models.BoundaryConditionKernelFunction{Side}) where Side
bc = Adapt.adapt(to, bckf.bc)
BC = typeof(bc)
return Oceananigans.Models.BoundaryConditionKernelFunction{Side, BC}(bc)
end

# The surface fluxes are:
#
# - ``τˣ``: momentum flux (stress), in kg m⁻¹ s⁻²
Expand Down