Skip to content

Commit 90c5707

Browse files
simone-silvestrisimone-silvestri
committed
separate more stuff out
1 parent 3c3c6f3 commit 90c5707

File tree

3 files changed

+114
-95
lines changed

3 files changed

+114
-95
lines changed

‎src/OceanSeaIceModels/InterfaceComputations/InterfaceComputations.jl‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,8 @@ include("compute_interface_state.jl")
7373
include("similarity_theory_turbulent_fluxes.jl")
7474
include("coefficient_based_turbulent_fluxes.jl")
7575

76+
# State exchanger and interfaces
77+
include("state_exchanger.jl")
7678
include("component_interfaces.jl")
7779
include("atmosphere_ocean_fluxes.jl")
7880
include("atmosphere_sea_ice_fluxes.jl")

‎src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl‎

Lines changed: 0 additions & 95 deletions
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,6 @@ end
4848
@inline get_possibly_zero_flux(fluxes, name) = getfield(fluxes, name)
4949
@inline get_possibly_zero_flux(::Nothing, name) = ZeroField()
5050

51-
5251
mutable struct ComponentInterfaces{AO, ASI, SIO, C, AP, OP, SIP, EX, P}
5352
atmosphere_ocean_interface :: AO
5453
atmosphere_sea_ice_interface :: ASI
@@ -61,100 +60,6 @@ mutable struct ComponentInterfaces{AO, ASI, SIO, C, AP, OP, SIP, EX, P}
6160
properties :: P
6261
end
6362

64-
mutable struct StateExchanger{G, AST, AEX}
65-
exchange_grid :: G
66-
exchange_atmosphere_state :: AST
67-
atmosphere_exchanger :: AEX
68-
end
69-
70-
mutable struct ExchangeAtmosphereState{F}
71-
u :: F
72-
v :: F
73-
T :: F
74-
q :: F
75-
p :: F
76-
Qs :: F
77-
Qℓ :: F
78-
Mp :: F
79-
end
80-
81-
ExchangeAtmosphereState(grid) = ExchangeAtmosphereState(Field{Center, Center, Nothing}(grid),
82-
Field{Center, Center, Nothing}(grid),
83-
Field{Center, Center, Nothing}(grid),
84-
Field{Center, Center, Nothing}(grid),
85-
Field{Center, Center, Nothing}(grid),
86-
Field{Center, Center, Nothing}(grid),
87-
Field{Center, Center, Nothing}(grid),
88-
Field{Center, Center, Nothing}(grid))
89-
90-
# Note that Field location can also affect fractional index type.
91-
# Here we assume that we know the location of Fields that will be interpolated.
92-
fractional_index_type(FT, Topo) = FT
93-
fractional_index_type(FT, ::Flat) = Nothing
94-
95-
StateExchanger(ocean::Simulation, ::Nothing) = nothing
96-
97-
function StateExchanger(ocean::Simulation, atmosphere)
98-
# TODO: generalize this
99-
exchange_grid = ocean.model.grid
100-
exchange_atmosphere_state = ExchangeAtmosphereState(exchange_grid)
101-
exchanger = atmosphere_exchanger(atmosphere, exchange_grid, exchange_atmosphere_state)
102-
103-
return StateExchanger(ocean.model.grid, exchange_atmosphere_state, exchanger)
104-
end
105-
106-
function atmosphere_exchanger(atmosphere::PrescribedAtmosphere, exchange_grid, exchange_atmosphere_state)
107-
atmos_grid = atmosphere.grid
108-
arch = architecture(exchange_grid)
109-
Nx, Ny, Nz = size(exchange_grid)
110-
111-
# Make a NamedTuple of fractional indices
112-
# Note: we could use an array of FractionalIndices. Instead, for compatbility
113-
# with Reactant we construct FractionalIndices on the fly in `interpolate_atmospheric_state`.
114-
FT = eltype(atmos_grid)
115-
TX, TY, TZ = topology(exchange_grid)
116-
fi = TX() isa Flat ? nothing : Field{Center, Center, Nothing}(exchange_grid, FT)
117-
fj = TY() isa Flat ? nothing : Field{Center, Center, Nothing}(exchange_grid, FT)
118-
frac_indices = (i=fi, j=fj) # no k needed, only horizontal interpolation
119-
120-
return frac_indices
121-
end
122-
123-
initialize!(exchanger::StateExchanger, ::Nothing) = nothing
124-
125-
function initialize!(exchanger::StateExchanger, atmosphere::PrescribedAtmosphere)
126-
atmos_grid = atmosphere.grid
127-
exchange_grid = exchanger.exchange_grid
128-
arch = architecture(exchange_grid)
129-
frac_indices = exchanger.atmosphere_exchanger
130-
kernel_parameters = interface_kernel_parameters(exchange_grid)
131-
launch!(arch, exchange_grid, kernel_parameters,
132-
_compute_fractional_indices!, frac_indices, exchange_grid, atmos_grid)
133-
return nothing
134-
end
135-
136-
@kernel function _compute_fractional_indices!(indices_tuple, exchange_grid, atmos_grid)
137-
i, j = @index(Global, NTuple)
138-
kᴺ = size(exchange_grid, 3) # index of the top ocean cell
139-
X = _node(i, j, kᴺ + 1, exchange_grid, c, c, f)
140-
if topology(atmos_grid) == (Flat, Flat, Flat)
141-
fractional_indices_ij = FractionalIndices(nothing, nothing, nothing)
142-
else
143-
fractional_indices_ij = FractionalIndices(X, atmos_grid, c, c, c)
144-
end
145-
fi = indices_tuple.i
146-
fj = indices_tuple.j
147-
@inbounds begin
148-
if !isnothing(fi)
149-
fi[i, j, 1] = fractional_indices_ij.i
150-
end
151-
152-
if !isnothing(fj)
153-
fj[i, j, 1] = fractional_indices_ij.j
154-
end
155-
end
156-
end
157-
15863
# Possible units for temperature and salinity
15964
struct DegreesCelsius end
16065
struct DegreesKelvin end

‎src/OceanSeaIceModels/InterfaceComputations/state_exchanger.jl‎

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
2+
mutable struct StateExchanger{G, AS, OS, SS, AEX, OEX, SIE}
3+
exchange_grid :: G
4+
exchange_atmosphere_state :: AST
5+
exchange_ocean_state :: OS
6+
exchange_sea_ice_state :: SS
7+
atmosphere_exchanger :: AEX
8+
ocean_exchanger :: OEX
9+
sea_ice_exchanger :: SIE
10+
end
11+
12+
mutable struct ExchangeAtmosphereState{FC, CF, CC}
13+
u :: FC
14+
v :: CF
15+
T :: CC
16+
q :: CC
17+
p :: CC
18+
Qs :: CC
19+
Qℓ :: CC
20+
Mp :: CC
21+
end
22+
23+
mutable struct ExchangeOceanState{FC, CF, CC}
24+
u :: FC
25+
v :: CF
26+
T :: CC
27+
S :: CC
28+
end
29+
30+
mutable struct ExchangeSeaIceState{FC, CF, CC}
31+
u :: FC
32+
v :: CF
33+
h :: CC
34+
ℵ :: CC
35+
end
36+
37+
ExchangeAtmosphereState(grid) = ExchangeAtmosphereState(Field{Center, Center, Nothing}(grid),
38+
Field{Center, Center, Nothing}(grid),
39+
Field{Center, Center, Nothing}(grid),
40+
Field{Center, Center, Nothing}(grid),
41+
Field{Center, Center, Nothing}(grid),
42+
Field{Center, Center, Nothing}(grid),
43+
Field{Center, Center, Nothing}(grid),
44+
Field{Center, Center, Nothing}(grid))
45+
46+
# Note that Field location can also affect fractional index type.
47+
# Here we assume that we know the location of Fields that will be interpolated.
48+
fractional_index_type(FT, Topo) = FT
49+
fractional_index_type(FT, ::Flat) = Nothing
50+
51+
StateExchanger(ocean::Simulation, ::Nothing) = nothing
52+
53+
function StateExchanger(ocean::Simulation, atmosphere)
54+
# TODO: generalize this
55+
exchange_grid = ocean.model.grid
56+
exchange_atmosphere_state = ExchangeAtmosphereState(exchange_grid)
57+
exchanger = atmosphere_exchanger(atmosphere, exchange_grid, exchange_atmosphere_state)
58+
59+
return StateExchanger(ocean.model.grid, exchange_atmosphere_state, exchanger)
60+
end
61+
62+
function atmosphere_exchanger(atmosphere::PrescribedAtmosphere, exchange_grid, exchange_atmosphere_state)
63+
atmos_grid = atmosphere.grid
64+
arch = architecture(exchange_grid)
65+
Nx, Ny, Nz = size(exchange_grid)
66+
67+
# Make a NamedTuple of fractional indices
68+
# Note: we could use an array of FractionalIndices. Instead, for compatbility
69+
# with Reactant we construct FractionalIndices on the fly in `interpolate_atmospheric_state`.
70+
FT = eltype(atmos_grid)
71+
TX, TY, TZ = topology(exchange_grid)
72+
fi = TX() isa Flat ? nothing : Field{Center, Center, Nothing}(exchange_grid, FT)
73+
fj = TY() isa Flat ? nothing : Field{Center, Center, Nothing}(exchange_grid, FT)
74+
frac_indices = (i=fi, j=fj) # no k needed, only horizontal interpolation
75+
76+
return frac_indices
77+
end
78+
79+
initialize!(exchanger::StateExchanger, ::Nothing) = nothing
80+
81+
function initialize!(exchanger::StateExchanger, atmosphere::PrescribedAtmosphere)
82+
atmos_grid = atmosphere.grid
83+
exchange_grid = exchanger.exchange_grid
84+
arch = architecture(exchange_grid)
85+
frac_indices = exchanger.atmosphere_exchanger
86+
kernel_parameters = interface_kernel_parameters(exchange_grid)
87+
launch!(arch, exchange_grid, kernel_parameters,
88+
_compute_fractional_indices!, frac_indices, exchange_grid, atmos_grid)
89+
return nothing
90+
end
91+
92+
@kernel function _compute_fractional_indices!(indices_tuple, exchange_grid, atmos_grid)
93+
i, j = @index(Global, NTuple)
94+
kᴺ = size(exchange_grid, 3) # index of the top ocean cell
95+
X = _node(i, j, kᴺ + 1, exchange_grid, c, c, f)
96+
if topology(atmos_grid) == (Flat, Flat, Flat)
97+
fractional_indices_ij = FractionalIndices(nothing, nothing, nothing)
98+
else
99+
fractional_indices_ij = FractionalIndices(X, atmos_grid, c, c, c)
100+
end
101+
fi = indices_tuple.i
102+
fj = indices_tuple.j
103+
@inbounds begin
104+
if !isnothing(fi)
105+
fi[i, j, 1] = fractional_indices_ij.i
106+
end
107+
108+
if !isnothing(fj)
109+
fj[i, j, 1] = fractional_indices_ij.j
110+
end
111+
end
112+
end

0 commit comments

Comments
 (0)