Skip to content

Commit a95ce1b

Browse files
committed
rm bycolumn
1 parent ab243e2 commit a95ce1b

12 files changed

+122
-237
lines changed

docs/src/interfacer.md

+1-1
Original file line numberDiff line numberDiff line change
@@ -151,7 +151,7 @@ following properties:
151151
| `turbulent_moisture_flux` | aerodynamic turbulent surface fluxes of energy (evaporation) | kg m^-2 s^-1 |
152152

153153
### SurfaceModelSimulation - optional functions
154-
- `update_turbulent_fluxes_point!(::ComponentModelSimulation, fields::NamedTuple, colidx)`:
154+
- `update_turbulent_fluxes_point!(::ComponentModelSimulation, fields::NamedTuple)`:
155155
This function updates the turbulent fluxes of the component model simulation
156156
at this point in horizontal space. The values are updated using the energy
157157
and moisture turbulent fluxes stored in fields which are calculated by the

experiments/ClimaEarth/components/atmosphere/climaatmos_extra_diags.jl

+3-10
Original file line numberDiff line numberDiff line change
@@ -65,16 +65,9 @@ CAD.add_diagnostic_variable!(
6565
(; C_E) = cache.atmos.vert_diff
6666
interior_uₕ = CC.Fields.level(state.c.uₕ, 1)
6767
ᶠp = ᶠK_E = cache.scratch.ᶠtemp_scalar
68-
CC.Fields.bycolumn(axes(ᶜp)) do colidx
69-
@. ᶠp[colidx] = CAD.ᶠinterp(ᶜp[colidx])
70-
ᶜΔz_surface = CC.Fields.Δz_field(interior_uₕ)
71-
@. ᶠK_E[colidx] = CA.eddy_diffusivity_coefficient(
72-
C_E,
73-
CA.norm(interior_uₕ[colidx]),
74-
ᶜΔz_surface[colidx] / 2,
75-
ᶠp[colidx],
76-
)
77-
end
68+
@. ᶠp = CAD.ᶠinterp(ᶜp)
69+
ᶜΔz_surface = CC.Fields.Δz_field(interior_uₕ)
70+
@. ᶠK_E = CA.eddy_diffusivity_coefficient(C_E, CA.norm(interior_uₕ), ᶜΔz_surface / 2, ᶠp)
7871
if isnothing(out)
7972
return CAD.ᶜinterp.(ᶠK_E)
8073
else

experiments/ClimaEarth/components/land/climaland_bucket.jl

+8-13
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,8 @@ Interfacer.name(::BucketSimulation) = "BucketSimulation"
2929
"""
3030
get_new_cache(p, Y, energy_check)
3131
32-
Returns a new `p` with an updated field to store e_per_area if energy conservation
33-
checks are turned on.
32+
Returns a new `p` with an updated field to store e_per_area if energy conservation
33+
checks are turned on.
3434
"""
3535
function get_new_cache(p, Y, energy_check)
3636
if energy_check
@@ -216,16 +216,12 @@ Interfacer.step!(sim::BucketSimulation, t) = Interfacer.step!(sim.integrator, t
216216
Interfacer.reinit!(sim::BucketSimulation) = Interfacer.reinit!(sim.integrator)
217217

218218
# extensions required by FluxCalculator (partitioned fluxes)
219-
function FluxCalculator.update_turbulent_fluxes_point!(
220-
sim::BucketSimulation,
221-
fields::NamedTuple,
222-
colidx::CC.Fields.ColumnIndex,
223-
)
219+
function FluxCalculator.update_turbulent_fluxes_point!(sim::BucketSimulation, fields::NamedTuple)
224220
(; F_turb_energy, F_turb_moisture) = fields
225221
turbulent_fluxes = sim.integrator.p.bucket.turbulent_fluxes
226-
turbulent_fluxes.shf[colidx] .= F_turb_energy
222+
turbulent_fluxes.shf .= F_turb_energy
227223
earth_params = sim.model.parameters.earth_param_set
228-
turbulent_fluxes.vapor_flux[colidx] .= F_turb_moisture ./ LP.ρ_cloud_liq(earth_params)
224+
turbulent_fluxes.vapor_flux .= F_turb_moisture ./ LP.ρ_cloud_liq(earth_params)
229225
return nothing
230226
end
231227

@@ -234,15 +230,14 @@ function FluxCalculator.surface_thermo_state(
234230
sim::BucketSimulation,
235231
thermo_params::TD.Parameters.ThermodynamicsParameters,
236232
thermo_state_int,
237-
colidx::CC.Fields.ColumnIndex,
238233
)
239234

240-
T_sfc = Interfacer.get_field(sim, Val(:surface_temperature), colidx)
235+
T_sfc = Interfacer.get_field(sim, Val(:surface_temperature))
241236
# Note that the surface air density, ρ_sfc, is computed using the atmospheric state at the first level and making ideal gas
242237
# and hydrostatic balance assumptions. The land model does not compute the surface air density so this is
243238
# a reasonable stand-in.
244-
ρ_sfc = Interfacer.get_field(sim, Val(:air_density), colidx)
245-
q_sfc = Interfacer.get_field(sim, Val(:surface_humidity), colidx) # already calculated in rhs! (cache)
239+
ρ_sfc = Interfacer.get_field(sim, Val(:air_density))
240+
q_sfc = Interfacer.get_field(sim, Val(:surface_humidity)) # already calculated in rhs! (cache)
246241
@. TD.PhaseEquil_ρTq.(thermo_params, ρ_sfc, T_sfc, q_sfc)
247242
end
248243

experiments/ClimaEarth/components/ocean/eisenman_seaice.jl

+8-15
Original file line numberDiff line numberDiff line change
@@ -140,8 +140,8 @@ end
140140
function Interfacer.update_field!(sim::EisenmanIceSimulation, ::Val{:area_fraction}, field::CC.Fields.Field)
141141
sim.integrator.p.area_fraction .= field
142142
end
143-
function Interfacer.update_field!(sim::EisenmanIceSimulation, ::Val{:∂F_turb_energy∂T_sfc}, field, colidx)
144-
sim.integrator.p.Ya.∂F_turb_energy∂T_sfc[colidx] .= field
143+
function Interfacer.update_field!(sim::EisenmanIceSimulation, ::Val{:∂F_turb_energy∂T_sfc}, field)
144+
sim.integrator.p.Ya.∂F_turb_energy∂T_sfc .= field
145145
end
146146
function Interfacer.update_field!(sim::EisenmanIceSimulation, ::Val{:radiative_energy_flux_sfc}, field)
147147
parent(sim.integrator.p.Ya.F_rad) .= parent(field)
@@ -155,13 +155,9 @@ Interfacer.step!(sim::EisenmanIceSimulation, t) = Interfacer.step!(sim.integrato
155155
Interfacer.reinit!(sim::EisenmanIceSimulation) = Interfacer.reinit!(sim.integrator)
156156

157157
# extensions required by FluxCalculator (partitioned fluxes)
158-
function FluxCalculator.update_turbulent_fluxes_point!(
159-
sim::EisenmanIceSimulation,
160-
fields::NamedTuple,
161-
colidx::CC.Fields.ColumnIndex,
162-
)
158+
function FluxCalculator.update_turbulent_fluxes_point!(sim::EisenmanIceSimulation, fields::NamedTuple)
163159
(; F_turb_energy) = fields
164-
@. sim.integrator.p.Ya.F_turb[colidx] = F_turb_energy
160+
@. sim.integrator.p.Ya.F_turb = F_turb_energy
165161
end
166162

167163
"""
@@ -194,9 +190,8 @@ function FluxCalculator.differentiate_turbulent_fluxes!(
194190
fluxes;
195191
δT_sfc = 0.1,
196192
)
197-
(; thermo_state_int, surface_params, surface_scheme, colidx) = input_args
198-
thermo_state_sfc_dT =
199-
FluxCalculator.surface_thermo_state(sim, thermo_params, thermo_state_int, colidx, δT_sfc = δT_sfc)
193+
(; thermo_state_int, surface_params, surface_scheme) = input_args
194+
thermo_state_sfc_dT = FluxCalculator.surface_thermo_state(sim, thermo_params, thermo_state_int, δT_sfc = δT_sfc)
200195
input_args = merge(input_args, (; thermo_state_sfc = thermo_state_sfc_dT))
201196

202197
# set inputs based on whether the surface_scheme is `MoninObukhovScheme` or `BulkScheme`
@@ -210,7 +205,7 @@ function FluxCalculator.differentiate_turbulent_fluxes!(
210205
# calculate the derivative
211206
∂F_turb_energy∂T_sfc = @. (F_shf_δT_sfc + F_lhf_δT_sfc - F_shf - F_lhf) / δT_sfc
212207

213-
Interfacer.update_field!(sim, Val(:∂F_turb_energy∂T_sfc), ∂F_turb_energy∂T_sfc, colidx)
208+
Interfacer.update_field!(sim, Val(:∂F_turb_energy∂T_sfc), ∂F_turb_energy∂T_sfc)
214209
end
215210

216211
###
@@ -354,9 +349,7 @@ function ∑tendencies(dY, Y, cache, _)
354349
@. dY.T_sfc = -Y.T_sfc / Δt
355350
@. dY.q_sfc = -Y.q_sfc / Δt
356351

357-
CC.Fields.bycolumn(axes(Y.T_sfc)) do colidx
358-
solve_eisenman_model!(Y[colidx], Ya[colidx], p, thermo_params, Δt)
359-
end
352+
solve_eisenman_model!(Y, Ya, p, thermo_params, Δt)
360353

361354
# Get dY/dt
362355
@. dY.T_ml += Y.T_ml / Δt

experiments/ClimaEarth/components/ocean/prescr_seaice.jl

+2-6
Original file line numberDiff line numberDiff line change
@@ -131,13 +131,9 @@ Interfacer.step!(sim::PrescribedIceSimulation, t) = Interfacer.step!(sim.integra
131131
Interfacer.reinit!(sim::PrescribedIceSimulation) = Interfacer.reinit!(sim.integrator)
132132

133133
# extensions required by FluxCalculator (partitioned fluxes)
134-
function FluxCalculator.update_turbulent_fluxes_point!(
135-
sim::PrescribedIceSimulation,
136-
fields::NamedTuple,
137-
colidx::CC.Fields.ColumnIndex,
138-
)
134+
function FluxCalculator.update_turbulent_fluxes_point!(sim::PrescribedIceSimulation, fields::NamedTuple)
139135
(; F_turb_energy) = fields
140-
@. sim.integrator.p.F_turb_energy[colidx] = F_turb_energy
136+
@. sim.integrator.p.F_turb_energy = F_turb_energy
141137
end
142138

143139
"""

experiments/ClimaEarth/components/ocean/slab_ocean.jl

+2-6
Original file line numberDiff line numberDiff line change
@@ -143,13 +143,9 @@ Interfacer.step!(sim::SlabOceanSimulation, t) = Interfacer.step!(sim.integrator,
143143
Interfacer.reinit!(sim::SlabOceanSimulation) = Interfacer.reinit!(sim.integrator)
144144

145145
# extensions required by FluxCalculator (partitioned fluxes)
146-
function FluxCalculator.update_turbulent_fluxes_point!(
147-
sim::SlabOceanSimulation,
148-
fields::NamedTuple,
149-
colidx::CC.Fields.ColumnIndex,
150-
)
146+
function FluxCalculator.update_turbulent_fluxes_point!(sim::SlabOceanSimulation, fields::NamedTuple)
151147
(; F_turb_energy) = fields
152-
@. sim.integrator.p.F_turb_energy[colidx] = F_turb_energy
148+
@. sim.integrator.p.F_turb_energy = F_turb_energy
153149
end
154150

155151
"""

src/FluxCalculator.jl

+71-85
Original file line numberDiff line numberDiff line change
@@ -142,70 +142,65 @@ function partitioned_turbulent_fluxes!(
142142
csf.F_turb_energy .*= FT(0)
143143
csf.F_turb_moisture .*= FT(0)
144144

145-
# iterate over all columns (when regridding, this will need to change)
146-
CC.Fields.bycolumn(boundary_space) do colidx
147-
# atmos state of center level 1
148-
z_int = Interfacer.get_field(atmos_sim, Val(:height_int), colidx)
149-
uₕ_int = Interfacer.get_field(atmos_sim, Val(:uv_int), colidx)
150-
thermo_state_int = Interfacer.get_field(atmos_sim, Val(:thermo_state_int), colidx)
151-
z_sfc = Interfacer.get_field(atmos_sim, Val(:height_sfc), colidx)
152-
153-
for sim in model_sims
154-
# iterate over all surface models with non-zero area fractions
155-
if sim isa Interfacer.SurfaceModelSimulation
156-
# get area fraction (min = 0, max = 1)
157-
area_fraction = Interfacer.get_field(sim, Val(:area_fraction), colidx)
158-
# get area mask [0, 1], where area_mask = 1 if area_fraction > 0
159-
area_mask = Regridder.binary_mask.(area_fraction)
160-
161-
if !iszero(parent(area_mask))
162-
163-
thermo_state_sfc = surface_thermo_state(sim, thermo_params, thermo_state_int, colidx)
164-
165-
# set inputs based on whether the surface_scheme is `MoninObukhovScheme` or `BulkScheme`
166-
surface_params = get_surface_params(atmos_sim)
167-
scheme_properties = get_scheme_properties(surface_scheme, sim, colidx)
168-
input_args = (;
169-
thermo_state_sfc,
170-
thermo_state_int,
171-
uₕ_int,
172-
z_int,
173-
z_sfc,
174-
scheme_properties,
175-
surface_params,
176-
surface_scheme,
177-
colidx,
178-
)
179-
inputs = surface_inputs(surface_scheme, input_args)
180-
181-
# calculate the surface fluxes
182-
fluxes = get_surface_fluxes_point!(inputs, surface_params)
183-
(; F_turb_ρτxz, F_turb_ρτyz, F_shf, F_lhf, F_turb_moisture) = fluxes
184-
185-
# perform additional diagnostics if required
186-
differentiate_turbulent_fluxes!(sim, (thermo_params, input_args, fluxes))
187-
188-
# update fluxes in the coupler
189-
fields = (;
190-
F_turb_ρτxz = F_turb_ρτxz,
191-
F_turb_ρτyz = F_turb_ρτyz,
192-
F_turb_energy = F_shf .+ F_lhf,
193-
F_turb_moisture = F_turb_moisture,
194-
)
195-
196-
# update the fluxes of this surface model
197-
update_turbulent_fluxes_point!(sim, fields, colidx)
198-
199-
# add the flux contributing from this surface to the coupler field
200-
@. csf.F_turb_ρτxz[colidx] += F_turb_ρτxz * area_fraction * area_mask
201-
@. csf.F_turb_ρτyz[colidx] += F_turb_ρτyz * area_fraction * area_mask
202-
@. csf.F_turb_energy[colidx] += (F_shf .+ F_lhf) * area_fraction * area_mask
203-
@. csf.F_turb_moisture[colidx] += F_turb_moisture * area_fraction * area_mask
204-
205-
end
145+
# atmos state of center level 1
146+
z_int = Interfacer.get_field(atmos_sim, Val(:height_int))
147+
uₕ_int = Interfacer.get_field(atmos_sim, Val(:uv_int))
148+
thermo_state_int = Interfacer.get_field(atmos_sim, Val(:thermo_state_int))
149+
z_sfc = Interfacer.get_field(atmos_sim, Val(:height_sfc))
150+
151+
for sim in model_sims
152+
# iterate over all surface models with non-zero area fractions
153+
if sim isa Interfacer.SurfaceModelSimulation
154+
# get area fraction (min = 0, max = 1)
155+
area_fraction = Interfacer.get_field(sim, Val(:area_fraction))
156+
# get area mask [0, 1], where area_mask = 1 if area_fraction > 0
157+
area_mask = Regridder.binary_mask.(area_fraction)
158+
159+
if !iszero(parent(area_mask))
160+
161+
thermo_state_sfc = surface_thermo_state(sim, thermo_params, thermo_state_int)
162+
163+
# set inputs based on whether the surface_scheme is `MoninObukhovScheme` or `BulkScheme`
164+
surface_params = get_surface_params(atmos_sim)
165+
scheme_properties = get_scheme_properties(surface_scheme, sim)
166+
input_args = (;
167+
thermo_state_sfc,
168+
thermo_state_int,
169+
uₕ_int,
170+
z_int,
171+
z_sfc,
172+
scheme_properties,
173+
surface_params,
174+
surface_scheme,
175+
)
176+
inputs = surface_inputs(surface_scheme, input_args)
177+
178+
# calculate the surface fluxes
179+
fluxes = get_surface_fluxes_point!(inputs, surface_params)
180+
(; F_turb_ρτxz, F_turb_ρτyz, F_shf, F_lhf, F_turb_moisture) = fluxes
181+
182+
# perform additional diagnostics if required
183+
differentiate_turbulent_fluxes!(sim, (thermo_params, input_args, fluxes))
184+
185+
# update fluxes in the coupler
186+
fields = (;
187+
F_turb_ρτxz = F_turb_ρτxz,
188+
F_turb_ρτyz = F_turb_ρτyz,
189+
F_turb_energy = F_shf .+ F_lhf,
190+
F_turb_moisture = F_turb_moisture,
191+
)
192+
193+
# update the fluxes of this surface model
194+
update_turbulent_fluxes_point!(sim, fields)
195+
196+
# add the flux contributing from this surface to the coupler field
197+
@. csf.F_turb_ρτxz += F_turb_ρτxz * area_fraction * area_mask
198+
@. csf.F_turb_ρτyz += F_turb_ρτyz * area_fraction * area_mask
199+
@. csf.F_turb_energy += (F_shf .+ F_lhf) * area_fraction * area_mask
200+
@. csf.F_turb_moisture += F_turb_moisture * area_fraction * area_mask
201+
206202
end
207203
end
208-
209204
end
210205

211206
# TODO: add allowable bounds here, check explicitly that all fluxes are equal
@@ -217,25 +212,21 @@ struct BulkScheme <: AbstractSurfaceFluxScheme end
217212
struct MoninObukhovScheme <: AbstractSurfaceFluxScheme end
218213

219214
"""
220-
get_scheme_properties(scheme::AbstractSurfaceFluxScheme, sim::Interfacer.SurfaceModelSimulation, colidx::CC.Fields.ColumnIndex)
215+
get_scheme_properties(scheme::AbstractSurfaceFluxScheme, sim::Interfacer.SurfaceModelSimulation)
221216
222217
Returns the scheme-specific properties for the surface model simulation `sim`.
223218
"""
224-
function get_scheme_properties(::BulkScheme, sim::Interfacer.SurfaceModelSimulation, colidx::CC.Fields.ColumnIndex)
225-
Ch = Interfacer.get_field(sim, Val(:heat_transfer_coefficient), colidx)
226-
Cd = Interfacer.get_field(sim, Val(:beta), colidx)
227-
beta = Interfacer.get_field(sim, Val(:drag_coefficient), colidx)
219+
function get_scheme_properties(::BulkScheme, sim::Interfacer.SurfaceModelSimulation)
220+
Ch = Interfacer.get_field(sim, Val(:heat_transfer_coefficient))
221+
Cd = Interfacer.get_field(sim, Val(:beta))
222+
beta = Interfacer.get_field(sim, Val(:drag_coefficient))
228223
FT = eltype(Ch)
229224
return (; z0b = FT(0), z0m = FT(0), Ch = Ch, Cd = Cd, beta = beta, gustiness = FT(1))
230225
end
231-
function get_scheme_properties(
232-
::MoninObukhovScheme,
233-
sim::Interfacer.SurfaceModelSimulation,
234-
colidx::CC.Fields.ColumnIndex,
235-
)
236-
z0m = Interfacer.get_field(sim, Val(:roughness_momentum), colidx)
237-
z0b = Interfacer.get_field(sim, Val(:roughness_buoyancy), colidx)
238-
beta = Interfacer.get_field(sim, Val(:beta), colidx)
226+
function get_scheme_properties(::MoninObukhovScheme, sim::Interfacer.SurfaceModelSimulation)
227+
z0m = Interfacer.get_field(sim, Val(:roughness_momentum))
228+
z0b = Interfacer.get_field(sim, Val(:roughness_buoyancy))
229+
beta = Interfacer.get_field(sim, Val(:beta))
239230
FT = eltype(z0m)
240231
return (; z0b = z0b, z0m = z0m, Ch = FT(0), Cd = FT(0), beta = beta, gustiness = FT(1))
241232
end
@@ -287,21 +278,20 @@ function surface_inputs(::MoninObukhovScheme, input_args::NamedTuple)
287278
end
288279

289280
"""
290-
surface_thermo_state(sim::Interfacer.SurfaceModelSimulation, thermo_params::TD.Parameters.ThermodynamicsParameters, thermo_state_int, colidx::CC.Fields.ColumnIndex)
281+
surface_thermo_state(sim::Interfacer.SurfaceModelSimulation, thermo_params::TD.Parameters.ThermodynamicsParameters, thermo_state_int)
291282
292283
Returns the surface parameters for the surface model simulation `sim`. The default is assuming saturated surfaces, unless an extension is defined for the given `SurfaceModelSimulation`.
293284
"""
294285
function surface_thermo_state(
295286
sim::Interfacer.SurfaceModelSimulation,
296287
thermo_params::TD.Parameters.ThermodynamicsParameters,
297-
thermo_state_int,
298-
colidx::CC.Fields.ColumnIndex;
288+
thermo_state_int;
299289
δT_sfc = 0,
300290
)
301291
FT = eltype(parent(thermo_state_int))
302292
@warn("Simulation " * Interfacer.name(sim) * " uses the default thermo (saturated) surface state", maxlog = 10)
303293
# get surface temperature (or perturbed surface temperature for differentiation)
304-
T_sfc = Interfacer.get_field(sim, Val(:surface_temperature), colidx) .+ FT(δT_sfc)
294+
T_sfc = Interfacer.get_field(sim, Val(:surface_temperature)) .+ FT(δT_sfc)
305295
ρ_sfc = extrapolate_ρ_to_sfc.(thermo_params, thermo_state_int, T_sfc) # ideally the # calculate elsewhere, here just getter...
306296
q_sfc = TD.q_vap_saturation_generic.(thermo_params, T_sfc, ρ_sfc, TD.Liquid()) # default: saturated liquid surface
307297
@. TD.PhaseEquil_ρTq.(thermo_params, ρ_sfc, T_sfc, q_sfc)
@@ -367,15 +357,11 @@ function get_surface_params(atmos_sim::Interfacer.AtmosModelSimulation)
367357
end
368358

369359
"""
370-
update_turbulent_fluxes_point!(sim::Interfacer.SurfaceModelSimulation, fields::NamedTuple, colidx::CC.Fields.ColumnIndex)
360+
update_turbulent_fluxes_point!(sim::Interfacer.SurfaceModelSimulation, fields::NamedTuple)
371361
372362
Updates the fluxes in the surface model simulation `sim` with the fluxes in `fields`.
373363
"""
374-
function update_turbulent_fluxes_point!(
375-
sim::Interfacer.SurfaceModelSimulation,
376-
fields::NamedTuple,
377-
colidx::CC.Fields.ColumnIndex,
378-
)
364+
function update_turbulent_fluxes_point!(sim::Interfacer.SurfaceModelSimulation, fields::NamedTuple)
379365
return error(
380366
"update_turbulent_fluxes_point! is required to be dispatched on" *
381367
Interfacer.name(sim) *

src/Interfacer.jl

-13
Original file line numberDiff line numberDiff line change
@@ -168,19 +168,6 @@ get_field(sim::ComponentModelSimulation, val::Val) = get_field_error(sim, val)
168168

169169
get_field_error(sim, val::Val{X}) where {X} = error("undefined field `$X` for " * name(sim))
170170

171-
"""
172-
get_field(::ComponentModelSimulation, ::Val, colidx::CC.Fields.ColumnIndex)
173-
174-
Extension of `get_field(::ComponentModelSimulation, ::Val)`, indexing into the specified colum index.
175-
"""
176-
function get_field(sim::ComponentModelSimulation, val::Val, colidx::CC.Fields.ColumnIndex)
177-
if get_field(sim, val) isa AbstractFloat
178-
get_field(sim, val)
179-
else
180-
get_field(sim, val)[colidx]
181-
end
182-
end
183-
184171
"""
185172
update_field!(::AtmosModelSimulation, ::Val, _...)
186173

0 commit comments

Comments
 (0)