Skip to content

Commit b78429a

Browse files
committed
remove prescribed flux callback, add fixes to enable use on GPUs
1 parent 8cf4e0d commit b78429a

File tree

6 files changed

+228
-101
lines changed

6 files changed

+228
-101
lines changed

src/DataWrangling/OSPapa/OSPapa.jl

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ module OSPapa
33
export OSPapaPrescribedAtmosphere
44
export OSPapaPrescribedFluxes
55
export OSPapaPrescribedFluxBoundaryConditions
6-
export PrescribedFluxCallback
76
export OSPapaHourly
87

98
using Oceananigans

src/DataWrangling/OSPapa/OSPapa_ocean_observations.jl

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -306,7 +306,8 @@ function set!(target_field::Field, metadata::OSPapaMetadatum; kw...)
306306
interpolated ./= 100 # cm/s → m/s
307307
end
308308

309-
interior(target_field, 1, 1, :) .= interpolated
309+
arch = Oceananigans.Architectures.architecture(target_field)
310+
interior(target_field, 1, 1, :) .= Oceananigans.on_architecture(arch, interpolated)
310311

311312
return target_field
312313
end

src/DataWrangling/OSPapa/OSPapa_prescribed_atmosphere.jl

Lines changed: 20 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,16 @@ barometric pressure, shortwave and longwave radiation, and precipitation.
9595
Relative humidity is converted to specific humidity; units are converted to
9696
SI (Kelvin, Pa, kg/m²/s).
9797
98+
!!! note "Radiation and albedo"
99+
The buoy `SW` and `LW` variables are **downwelling** fluxes. When this
100+
atmosphere is used with `OceanOnlyModel`, ClimaOcean applies its own
101+
ocean albedo (default α = 0.05) to compute net absorbed shortwave, and
102+
computes upwelling longwave from the model SST via Stefan-Boltzmann. This
103+
means the resulting net heat flux will differ from the COARE-computed
104+
`QNET` available via [`OSPapaPrescribedFluxes`](@ref). If you need the
105+
exact observed net fluxes, use [`OSPapaPrescribedFluxBoundaryConditions`](@ref)
106+
instead.
107+
98108
Keyword Arguments
99109
=================
100110
- `start_date`: start of the time range (default: `first_date(OSPapaHourly(), :air_temperature)`)
@@ -112,6 +122,8 @@ function OSPapaPrescribedAtmosphere(architecture = CPU(), FT = Float32;
112122
surface_layer_height = 2.5,
113123
max_gap_hours = 72)
114124

125+
on_arch = arr -> Oceananigans.on_architecture(architecture, arr)
126+
115127
filepath = download_ospapa_file(dir)
116128

117129
ds = NCDataset(filepath)
@@ -162,14 +174,14 @@ function OSPapaPrescribedAtmosphere(architecture = CPU(), FT = Float32;
162174
surface_layer_height = convert(FT, surface_layer_height))
163175

164176
Nt = length(times_seconds)
165-
parent(atmosphere.velocities.u) .= reshape(FT.(uwnd), 1, 1, 1, Nt)
166-
parent(atmosphere.velocities.v) .= reshape(FT.(vwnd), 1, 1, 1, Nt)
167-
parent(atmosphere.tracers.T) .= reshape(FT.(T_K), 1, 1, 1, Nt)
168-
parent(atmosphere.tracers.q) .= reshape(FT.(q), 1, 1, 1, Nt)
169-
parent(atmosphere.pressure) .= reshape(FT.(P_Pa), 1, 1, 1, Nt)
170-
parent(atmosphere.downwelling_radiation.shortwave) .= reshape(FT.(sw), 1, 1, 1, Nt)
171-
parent(atmosphere.downwelling_radiation.longwave) .= reshape(FT.(lw), 1, 1, 1, Nt)
172-
parent(atmosphere.freshwater_flux.rain) .= reshape(FT.(rain_kgm2s), 1, 1, 1, Nt)
177+
parent(atmosphere.velocities.u) .= on_arch(reshape(FT.(uwnd), 1, 1, 1, Nt))
178+
parent(atmosphere.velocities.v) .= on_arch(reshape(FT.(vwnd), 1, 1, 1, Nt))
179+
parent(atmosphere.tracers.T) .= on_arch(reshape(FT.(T_K), 1, 1, 1, Nt))
180+
parent(atmosphere.tracers.q) .= on_arch(reshape(FT.(q), 1, 1, 1, Nt))
181+
parent(atmosphere.pressure) .= on_arch(reshape(FT.(P_Pa), 1, 1, 1, Nt))
182+
parent(atmosphere.downwelling_radiation.shortwave) .= on_arch(reshape(FT.(sw), 1, 1, 1, Nt))
183+
parent(atmosphere.downwelling_radiation.longwave) .= on_arch(reshape(FT.(lw), 1, 1, 1, Nt))
184+
parent(atmosphere.freshwater_flux.rain) .= on_arch(reshape(FT.(rain_kgm2s), 1, 1, 1, Nt))
173185

174186
return atmosphere
175187
end

src/DataWrangling/OSPapa/OSPapa_prescribed_fluxes.jl

Lines changed: 7 additions & 90 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
using Oceananigans.Units
2-
using Oceananigans.OutputReaders: TimeInterpolator, Cyclical
32

43
const ERDDAP_BASE = "https://data.pmel.noaa.gov/pmel/erddap/tabledap"
54
const ERDDAP_FLUX_VARS = "time,QLAT,QSEN,QNET,LWNET,SWNET,TAU,TAUX,TAUY,RAIN,EVAP,EMP,TSK"
@@ -96,93 +95,6 @@ function OSPapaPrescribedFluxes(FT = Float64;
9695
start_date = times_datetime[1])
9796
end
9897

99-
"""
100-
PrescribedFluxCallback(ocean, fluxes; ρ₀=reference_density(ocean), cₚ=heat_capacity(ocean))
101-
102-
Create a `Callback` that applies prescribed air-sea fluxes to an ocean simulation
103-
at every time step. The callback interpolates flux time series to the current
104-
model time and writes the result into the ocean's boundary condition fields.
105-
106-
Arguments
107-
=========
108-
- `ocean`: an `Oceananigans.Simulation` or ocean model returned by `ocean_simulation`
109-
- `fluxes`: a `NamedTuple` returned by `OSPapaPrescribedFluxes`
110-
111-
Keyword Arguments
112-
=================
113-
- `ρ₀`: reference ocean density (default: from ocean's equation of state, typically 1020 kg/m³)
114-
- `cₚ`: ocean heat capacity (default: from ocean's equation of state, typically 3991.87 J/(kg·K))
115-
116-
The callback converts:
117-
- Wind stress (N/m²) → kinematic stress (m²/s²) by dividing by ρ₀
118-
- Net heat flux (W/m²) → temperature flux (K·m/s) by dividing by ρ₀·cₚ
119-
- Freshwater flux is applied as a salinity flux using surface salinity
120-
121-
!!! note "Sign conventions"
122-
ERDDAP QNET is positive into the ocean (warming). The callback converts
123-
to the Oceananigans convention where a positive top flux is *out of* the
124-
domain, so the sign is flipped for heat.
125-
"""
126-
function PrescribedFluxCallback(ocean, fluxes; ρ₀=reference_density(ocean), cₚ=heat_capacity(ocean))
127-
# Pre-extract BC fields (ocean_simulation returns a Simulation)
128-
model = ocean.model
129-
τˣ = model.velocities.u.boundary_conditions.top.condition
130-
τʸ = model.velocities.v.boundary_conditions.top.condition
131-
Jᵀ = model.tracers.T.boundary_conditions.top.condition
132-
= model.tracers.S.boundary_conditions.top.condition
133-
134-
flux_times = fluxes.times
135-
Nt = length(flux_times)
136-
Δt = flux_times[end] - flux_times[end-1]
137-
period = flux_times[end] - flux_times[1] + Δt
138-
time_indexing = Cyclical(period)
139-
140-
function update_fluxes!(sim)
141-
t = sim.model.clock.time
142-
143-
# Use Oceananigans' TimeInterpolator for consistency with PrescribedAtmosphere
144-
time_interpolator = TimeInterpolator(time_indexing, flux_times, t)
145-
n₁ = time_interpolator.first_index
146-
n₂ = time_interpolator.second_index
147-
α = time_interpolator.fractional_index
148-
149-
interp(field) = field[n₁] + α * (field[n₂] - field[n₁])
150-
151-
# Interpolate fluxes
152-
τx_now = interp(fluxes.τx)
153-
τy_now = interp(fluxes.τy)
154-
Qnet_now = interp(fluxes.Qnet)
155-
156-
# Momentum: stress (N/m²) → kinematic stress (m²/s²)
157-
# ERDDAP: TAUX > 0 = eastward stress ON ocean (downward momentum flux INTO domain)
158-
# Oceananigans FluxBoundaryCondition: positive top flux = OUT of domain
159-
# Therefore negate: downward flux into ocean → negative top BC
160-
τˣ[1, 1, 1] = -τx_now / ρ₀
161-
τʸ[1, 1, 1] = -τy_now / ρ₀
162-
163-
# Heat: Qnet (W/m²) → temperature flux (K·m/s)
164-
# ERDDAP Qnet positive = into ocean (warming)
165-
# Oceananigans FluxBoundaryCondition: positive = out of domain
166-
Jᵀ[1, 1, 1] = -Qnet_now / (ρ₀ * cₚ)
167-
168-
# Salinity: EMP (mm/hr ≡ kg/m²/hr) → salinity flux
169-
# EMP > 0 means net evaporation (ocean loses freshwater, salinity increases)
170-
# Convert mass flux to Boussinesq volume flux: divide by (ρ₀ * 3600)
171-
EMP_now = interp(fluxes.EMP)
172-
EMP_ms = EMP_now / (ρ₀ * 3600) # kg/m²/hr → m/s (Boussinesq volume flux)
173-
174-
# Salinity flux: following assemble_net_ocean_fluxes.jl: Jˢ = -S * ΣFao
175-
# EMP > 0 = net evaporation = salinity should increase
176-
# Negative Jˢ at top = salt INTO domain = S increases
177-
S_surface = model.tracers.S[1, 1, size(model.tracers.S, 3)]
178-
Jˢ[1, 1, 1] = -S_surface * EMP_ms
179-
180-
return nothing
181-
end
182-
183-
return Callback(update_fluxes!, IterationInterval(1))
184-
end
185-
18698
using Oceananigans.OutputReaders: interpolating_time_indices
18799

188100
"""
@@ -234,12 +146,17 @@ Keyword Arguments
234146
Each correction function must have the signature `(i, j, grid, clock, model_fields, p)` and return a value
235147
in the same units as the corresponding flux boundary condition.
236148
237-
Example
238-
=======
149+
Examples
150+
========
239151
```julia
152+
# Basic usage on GPU:
240153
fluxes = OSPapaPrescribedFluxes(; start_date, end_date)
241154
bcs = OSPapaPrescribedFluxBoundaryConditions(fluxes, GPU())
242155
ocean = ocean_simulation(grid; Δt=10minutes, boundary_conditions=bcs)
156+
157+
# With a uniform heat flux correction of +5 W/m² to close the heat budget:
158+
heat_correction = (i, j, grid, clock, model_fields, p) -> 5.0 / (p.ρ₀ * p.cₚ)
159+
bcs = OSPapaPrescribedFluxBoundaryConditions(fluxes, GPU(); T_correction=heat_correction)
243160
```
244161
"""
245162
function OSPapaPrescribedFluxBoundaryConditions(fluxes, architecture=CPU();

src/NumericalEarth.jl

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,6 @@ export
2929
OSPapaPrescribedAtmosphere,
3030
OSPapaPrescribedFluxes,
3131
OSPapaPrescribedFluxBoundaryConditions,
32-
PrescribedFluxCallback,
3332
OSPapaHourly,
3433
JRA55NetCDFBackend,
3534
regrid_bathymetry,

test/test_ospapa.jl

Lines changed: 199 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,199 @@
1+
include("runtests_setup.jl")
2+
3+
using NumericalEarth.OSPapa
4+
using NumericalEarth.Atmospheres: PrescribedAtmosphere
5+
using Oceananigans.BoundaryConditions: BoundaryCondition, Flux
6+
using CUDA: @allowscalar
7+
8+
const OSPAPA_TEST_START = DateTime(2012, 10, 1)
9+
const OSPAPA_TEST_END = DateTime(2012, 10, 3)
10+
11+
@testset "OS Papa Prescribed Atmosphere" begin
12+
for arch in test_architectures
13+
A = typeof(arch)
14+
@info "Testing OSPapaPrescribedAtmosphere on $A..."
15+
16+
atmosphere = OSPapaPrescribedAtmosphere(arch;
17+
start_date = OSPAPA_TEST_START,
18+
end_date = OSPAPA_TEST_END)
19+
20+
@test atmosphere isa PrescribedAtmosphere
21+
22+
# All expected fields are present
23+
@test haskey(atmosphere.velocities, :u)
24+
@test haskey(atmosphere.velocities, :v)
25+
@test haskey(atmosphere.tracers, :T)
26+
@test haskey(atmosphere.tracers, :q)
27+
@test !isnothing(atmosphere.pressure)
28+
@test !isnothing(atmosphere.downwelling_radiation)
29+
@test haskey(atmosphere.freshwater_flux, :rain)
30+
31+
# Radiation sanity checks
32+
ℐꜜˢʷ = atmosphere.downwelling_radiation.shortwave
33+
ℐꜜˡʷ = atmosphere.downwelling_radiation.longwave
34+
35+
@allowscalar begin
36+
sw_data = interior(ℐꜜˢʷ)
37+
lw_data = interior(ℐꜜˡʷ)
38+
39+
@test all(sw_data .>= 0) # downwelling SW is non-negative
40+
@test all(lw_data .>= 0) # downwelling LW is non-negative
41+
@test maximum(sw_data) < 1500 # below solar constant
42+
@test maximum(lw_data) < 600 # reasonable atmospheric LW
43+
@test maximum(lw_data) > 50 # not all zero
44+
45+
# Air temperature in physical range (K)
46+
T_data = interior(atmosphere.tracers.T)
47+
@test all(T_data .>= 240)
48+
@test all(T_data .<= 320)
49+
end
50+
end
51+
end
52+
53+
@testset "OS Papa Prescribed Fluxes" begin
54+
@info "Testing OSPapaPrescribedFluxes on CPU..."
55+
56+
fluxes = OSPapaPrescribedFluxes(; start_date = OSPAPA_TEST_START,
57+
end_date = OSPAPA_TEST_END)
58+
59+
# NamedTuple structure
60+
@test haskey(fluxes, :Qnet)
61+
@test haskey(fluxes, :τx)
62+
@test haskey(fluxes, :τy)
63+
@test haskey(fluxes, :EMP)
64+
@test haskey(fluxes, :times)
65+
@test haskey(fluxes, :start_date)
66+
67+
# Times are monotonically increasing
68+
@test issorted(fluxes.times)
69+
@test length(fluxes.times) > 1
70+
71+
# No NaN after gap filling (short window with no large gaps)
72+
@test all(isfinite, fluxes.Qnet)
73+
@test all(isfinite, fluxes.τx)
74+
@test all(isfinite, fluxes.τy)
75+
@test all(isfinite, fluxes.EMP)
76+
end
77+
78+
@testset "OS Papa Prescribed Flux BCs" begin
79+
for arch in test_architectures
80+
A = typeof(arch)
81+
@info "Testing OSPapaPrescribedFluxBoundaryConditions on $A..."
82+
83+
fluxes = OSPapaPrescribedFluxes(; start_date = OSPAPA_TEST_START,
84+
end_date = OSPAPA_TEST_END)
85+
86+
bcs = OSPapaPrescribedFluxBoundaryConditions(fluxes, arch)
87+
88+
# Returns BCs for all four fields
89+
@test haskey(bcs, :u)
90+
@test haskey(bcs, :v)
91+
@test haskey(bcs, :T)
92+
@test haskey(bcs, :S)
93+
94+
# Each has a FluxBoundaryCondition at the top
95+
@test bcs.u.top isa BoundaryCondition{<:Flux}
96+
@test bcs.v.top isa BoundaryCondition{<:Flux}
97+
@test bcs.T.top isa BoundaryCondition{<:Flux}
98+
@test bcs.S.top isa BoundaryCondition{<:Flux}
99+
end
100+
end
101+
102+
@testset "OS Papa ocean profile set!" begin
103+
for arch in test_architectures
104+
A = typeof(arch)
105+
@info "Testing OSPapaMetadatum set! on $A..."
106+
107+
grid = RectilinearGrid(arch;
108+
size = 20,
109+
x = -144.9, y = 50.1,
110+
z = (-200, 0),
111+
topology = (Flat, Flat, Bounded))
112+
113+
T_field = CenterField(grid)
114+
S_field = CenterField(grid)
115+
116+
@test begin
117+
set!(T_field, Metadatum(:temperature, dataset=OSPapaHourly(), date=OSPAPA_TEST_START))
118+
set!(S_field, Metadatum(:salinity, dataset=OSPapaHourly(), date=OSPAPA_TEST_START))
119+
true
120+
end
121+
122+
# Values should be finite and physically reasonable
123+
@allowscalar begin
124+
T_interior = Array(interior(T_field, 1, 1, :))
125+
S_interior = Array(interior(S_field, 1, 1, :))
126+
@test all(isfinite, T_interior)
127+
@test all(isfinite, S_interior)
128+
@test all(T_interior .> -2) # above freezing
129+
@test all(T_interior .< 35) # below boiling
130+
@test all(S_interior .> 20) # saline ocean
131+
@test all(S_interior .< 45) # not hypersaline
132+
end
133+
end
134+
end
135+
136+
@testset "OS Papa flux BC simulation" begin
137+
for arch in test_architectures
138+
A = typeof(arch)
139+
@info "Testing short simulation with OSPapaPrescribedFluxBoundaryConditions on $A..."
140+
141+
fluxes = OSPapaPrescribedFluxes(; start_date = OSPAPA_TEST_START,
142+
end_date = OSPAPA_TEST_END)
143+
144+
bcs = OSPapaPrescribedFluxBoundaryConditions(fluxes, arch)
145+
146+
grid = RectilinearGrid(arch;
147+
size = 10,
148+
x = -144.9, y = 50.1,
149+
z = (-200, 0),
150+
topology = (Flat, Flat, Bounded))
151+
152+
ocean = ocean_simulation(grid; Δt = 10minutes,
153+
coriolis = FPlane(latitude=50.1),
154+
boundary_conditions = bcs)
155+
156+
set!(ocean.model,
157+
T = Metadatum(:temperature, dataset=OSPapaHourly(), date=OSPAPA_TEST_START),
158+
S = Metadatum(:salinity, dataset=OSPapaHourly(), date=OSPAPA_TEST_START))
159+
160+
ocean.stop_iteration = 2
161+
162+
@test begin
163+
run!(ocean)
164+
true
165+
end
166+
end
167+
end
168+
169+
@testset "OS Papa prescribed atmosphere simulation" begin
170+
for arch in test_architectures
171+
A = typeof(arch)
172+
@info "Testing OceanOnlyModel with OSPapaPrescribedAtmosphere on $A..."
173+
174+
grid = RectilinearGrid(arch;
175+
size = 10,
176+
x = -144.9, y = 50.1,
177+
z = (-200, 0),
178+
topology = (Flat, Flat, Bounded))
179+
180+
ocean = ocean_simulation(grid; Δt = 10minutes,
181+
coriolis = FPlane(latitude=50.1))
182+
183+
set!(ocean.model,
184+
T = Metadatum(:temperature, dataset=OSPapaHourly(), date=OSPAPA_TEST_START),
185+
S = Metadatum(:salinity, dataset=OSPapaHourly(), date=OSPAPA_TEST_START))
186+
187+
atmosphere = OSPapaPrescribedAtmosphere(arch;
188+
start_date = OSPAPA_TEST_START,
189+
end_date = OSPAPA_TEST_END)
190+
191+
coupled_model = OceanOnlyModel(ocean; atmosphere, radiation=Radiation(arch))
192+
simulation = Simulation(coupled_model; Δt=ocean.Δt, stop_iteration=2)
193+
194+
@test begin
195+
run!(simulation)
196+
true
197+
end
198+
end
199+
end

0 commit comments

Comments
 (0)