Skip to content

Commit 2e8e66c

Browse files
charleskawczynskicharleskawczynski
committed
Remove coriolis forces from the cache
1 parent 8cdd3f8 commit 2e8e66c

File tree

13 files changed

+177
-116
lines changed

13 files changed

+177
-116
lines changed

‎src/ClimaAtmos.jl

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ include(joinpath("utils", "read_gcm_driven_scm_data.jl"))
2121
include(joinpath("utils", "AtmosArtifacts.jl"))
2222
import .AtmosArtifacts as AA
2323

24+
include(joinpath("core", "core_quantities.jl"))
2425
include(
2526
joinpath("parameterized_tendencies", "radiation", "radiation_utilities.jl"),
2627
)

‎src/cache/cache.jl

Lines changed: 5 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ struct AtmosCache{
3333
"""ClimaAtmosParameters that have to be used"""
3434
params::CAP
3535

36-
"""Variables that are used generally, such as ᶜΦ"""
36+
"""Variables that are used generally"""
3737
core::COR
3838

3939
"""Used by update_surface_conditions! in set_precomputed_quantities! and coupler"""
@@ -89,10 +89,8 @@ function build_cache(Y, atmos, params, surface_setup, sim_info, aerosol_names)
8989
ᶜcoord = Fields.local_geometry_field(Y.c).coordinates
9090
ᶠcoord = Fields.local_geometry_field(Y.f).coordinates
9191
grav = FT(CAP.grav(params))
92-
ᶜΦ = grav .* ᶜcoord.z
93-
ᶠΦ = grav .* ᶠcoord.z
94-
95-
(; ᶜf³, ᶠf¹²) = compute_coriolis(ᶜcoord, ᶠcoord, params)
92+
ᶜz = ᶜcoord.z
93+
ᶠz = ᶠcoord.z
9694

9795
ghost_buffer =
9896
!do_dss(axes(Y.c)) ? (;) :
@@ -121,11 +119,8 @@ function build_cache(Y, atmos, params, surface_setup, sim_info, aerosol_names)
121119
Fields.level(Fields.local_geometry_field(Y.f), Fields.half)
122120

123121
core = (
124-
ᶜΦ,
125-
ᶠgradᵥ_ᶜΦ = ᶠgradᵥ.(ᶜΦ),
126-
ᶜgradᵥ_ᶠΦ = ᶜgradᵥ.(ᶠΦ),
127-
ᶜf³,
128-
ᶠf¹²,
122+
ᶠgradᵥ_ᶜΦ = ᶠgradᵥ.(Φ.(grav, ᶜz)),
123+
ᶜgradᵥ_ᶠΦ = ᶜgradᵥ.(Φ.(grav, ᶠz)),
129124
# Used by diagnostics such as hfres, evspblw
130125
surface_ct3_unit = CT3.(
131126
unit_basis_vector_data.(CT3, sfc_local_geometry)
@@ -185,32 +180,3 @@ function build_cache(Y, atmos, params, surface_setup, sim_info, aerosol_names)
185180

186181
return AtmosCache{map(typeof, args)...}(args...)
187182
end
188-
189-
190-
function compute_coriolis(ᶜcoord, ᶠcoord, params)
191-
if eltype(ᶜcoord) <: Geometry.LatLongZPoint
192-
Ω = CAP.Omega(params)
193-
global_geom = Spaces.global_geometry(axes(ᶜcoord))
194-
if global_geom isa Geometry.DeepSphericalGlobalGeometry
195-
@info "using deep atmosphere"
196-
coriolis_deep(coord::Geometry.LatLongZPoint) = Geometry.LocalVector(
197-
Geometry.Cartesian123Vector(zero(Ω), zero(Ω), 2 * Ω),
198-
global_geom,
199-
coord,
200-
)
201-
ᶜf³ = @. CT3(CT123(coriolis_deep(ᶜcoord)))
202-
ᶠf¹² = @. CT12(CT123(coriolis_deep(ᶠcoord)))
203-
else
204-
coriolis_shallow(coord::Geometry.LatLongZPoint) =
205-
Geometry.WVector(2 * Ω * sind(coord.lat))
206-
ᶜf³ = @. CT3(coriolis_shallow(ᶜcoord))
207-
ᶠf¹² = nothing
208-
end
209-
else
210-
f = CAP.f_plane_coriolis_frequency(params)
211-
coriolis_f_plane(coord) = Geometry.WVector(f)
212-
ᶜf³ = @. CT3(coriolis_f_plane(ᶜcoord))
213-
ᶠf¹² = nothing
214-
end
215-
return (; ᶜf³, ᶠf¹²)
216-
end

‎src/cache/diagnostic_edmf_precomputed_quantities.jl

Lines changed: 16 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -34,13 +34,14 @@ NVTX.@annotate function set_diagnostic_edmfx_draft_quantities_level!(
3434
mse_level,
3535
q_tot_level,
3636
p_level,
37-
Φ_level,
37+
z_level,
3838
)
3939
FT = eltype(thermo_params)
40+
grav = TDP.grav(thermo_params)
4041
@. ts_level = TD.PhaseEquil_phq(
4142
thermo_params,
4243
p_level,
43-
mse_level - Φ_level,
44+
mse_level - Φ(grav, z_level),
4445
q_tot_level,
4546
8,
4647
FT(0.0003),
@@ -92,7 +93,6 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_bottom_bc!(
9293
(; turbconv_model) = p.atmos
9394
FT = eltype(Y)
9495
n = n_mass_flux_subdomains(turbconv_model)
95-
(; ᶜΦ) = p.core
9696
(; ᶜp, ᶠu³, ᶜh_tot, ᶜK) = p.precomputed
9797
(; q_tot) = p.precomputed.ᶜspecific
9898
(; ustar, obukhov_length, buoyancy_flux, ρ_flux_h_tot, ρ_flux_q_tot) =
@@ -112,7 +112,6 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_bottom_bc!(
112112
q_tot_int_level = Fields.field_values(Fields.level(q_tot, 1))
113113

114114
p_int_level = Fields.field_values(Fields.level(ᶜp, 1))
115-
Φ_int_level = Fields.field_values(Fields.level(ᶜΦ, 1))
116115

117116
local_geometry_int_level =
118117
Fields.field_values(Fields.level(Fields.local_geometry_field(Y.c), 1))
@@ -187,7 +186,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_bottom_bc!(
187186
mseʲ_int_level,
188187
q_totʲ_int_level,
189188
p_int_level,
190-
Φ_int_level,
189+
z_int_level,
191190
)
192191
@. ρaʲ_int_level = ρʲ_int_level * FT(turbconv_params.surface_area)
193192
end
@@ -296,7 +295,6 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
296295
ᶜdz = Fields.Δz_field(axes(Y.c))
297296
(; params) = p
298297
(; dt) = p
299-
(; ᶜΦ) = p.core
300298
(; ᶜp, ᶠu³, ᶜts, ᶜh_tot, ᶜK) = p.precomputed
301299
(; q_tot) = p.precomputed.ᶜspecific
302300
(;
@@ -324,13 +322,17 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
324322
microphys_1m_params = CAP.microphysics_1m_params(params)
325323
turbconv_params = CAP.turbconv_params(params)
326324

327-
ᶠΦ = p.scratch.ᶠtemp_scalar
328-
@. ᶠΦ = CAP.grav(params) * ᶠz
325+
g = CAP.grav(params)
326+
ᶜcoords = Fields.coordinate_field(Y.c)
327+
ᶜz = Fields.coordinate_field(Y.c).z
328+
ᶠz = Fields.coordinate_field(Y.f).z
329+
global_geom = Spaces.global_geometry(axes(ᶜcoord))
330+
329331
ᶜ∇Φ³ = p.scratch.ᶜtemp_CT3
330-
@. ᶜ∇Φ³ = CT3(ᶜgradᵥ(ᶠΦ))
331-
@. ᶜ∇Φ³ += CT3(gradₕ(ᶜΦ))
332+
@. ᶜ∇Φ³ = CT3(ᶜgradᵥ(Φ(g, ᶠz)))
333+
@. ᶜ∇Φ³ += CT3(gradₕ(Φ(g, ᶜz)))
332334
ᶜ∇Φ₃ = p.scratch.ᶜtemp_C3
333-
@. ᶜ∇Φ₃ = ᶜgradᵥ(ᶠΦ)
335+
@. ᶜ∇Φ₃ = ᶜgradᵥ(Φ(g, ᶠz))
334336

335337
z_sfc_halflevel =
336338
Fields.field_values(Fields.level(Fields.coordinate_field(Y.f).z, half))
@@ -344,7 +346,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
344346
h_tot_level = Fields.field_values(Fields.level(ᶜh_tot, i))
345347
q_tot_level = Fields.field_values(Fields.level(q_tot, i))
346348
p_level = Fields.field_values(Fields.level(ᶜp, i))
347-
Φ_level = Fields.field_values(Fields.level(ᶜΦ, i))
349+
z_level = Fields.field_values(Fields.level(ᶜz, i))
348350
local_geometry_level = Fields.field_values(
349351
Fields.level(Fields.local_geometry_field(Y.c), i),
350352
)
@@ -355,7 +357,6 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
355357
end_index = fieldcount(eltype(∂x∂ξ_level)) # This will be 4 in 2D and 9 in 3D.
356358
∂x³∂ξ³_level = ∂x∂ξ_level.:($end_index)
357359

358-
Φ_prev_level = Fields.field_values(Fields.level(ᶜΦ, i - 1))
359360
∇Φ³_prev_level = Fields.field_values(Fields.level(ᶜ∇Φ³, i - 1))
360361
∇Φ³_data_prev_level = ∇Φ³_prev_level.components.data.:1
361362
∇Φ₃_prev_level = Fields.field_values(Fields.level(ᶜ∇Φ₃, i - 1))
@@ -555,7 +556,6 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
555556
q_rai_prev_level,
556557
q_sno_prev_level,
557558
tsʲ_prev_level,
558-
Φ_prev_level,
559559
dt,
560560
microphys_1m_params,
561561
thermo_params,
@@ -702,7 +702,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
702702
e_tot_0M_precipitation_sources_helper(
703703
thermo_params,
704704
tsʲ_prev_level,
705-
Φ_prev_level,
705+
Φ(g, z_prev_level),
706706
)
707707
)
708708
)
@@ -807,7 +807,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
807807
mseʲ_level,
808808
q_totʲ_level,
809809
p_level,
810-
Φ_level,
810+
Φ(g, z_level),
811811
)
812812
end
813813
ρaʲs_level = Fields.field_values(Fields.level(ᶜρaʲs, i))
@@ -1072,7 +1072,6 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_precipita
10721072
ᶜqᵣ,
10731073
ᶜqₛ,
10741074
ᶜts,
1075-
p.core.ᶜΦ,
10761075
p.dt,
10771076
microphys_1m_params,
10781077
thermo_params,

‎src/cache/precomputed_quantities.jl

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -465,9 +465,10 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t)
465465
thermo_params = CAP.thermodynamics_params(p.params)
466466
n = n_mass_flux_subdomains(turbconv_model)
467467
thermo_args = (thermo_params, moisture_model)
468-
(; ᶜΦ) = p.core
469468
(; ᶜspecific, ᶜu, ᶠu³, ᶜK, ᶜts, ᶜp) = p.precomputed
470469
ᶠuₕ³ = p.scratch.ᶠtemp_CT3
470+
ᶜz = Fields.coordinate_field(axes(Y.c)).z
471+
grav = TDP.grav(thermo_params)
471472

472473
@. ᶜspecific = specific_gs(Y.c)
473474
set_ᶠuₕ³!(ᶠuₕ³, Y)
@@ -493,7 +494,7 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t)
493494
# @. ᶜK += Y.c.sgs⁰.ρatke / Y.c.ρ
494495
# TODO: We should think more about these increments before we use them.
495496
end
496-
@. ᶜts = ts_gs(thermo_args..., ᶜspecific, ᶜK, ᶜΦ, Y.c.ρ)
497+
@. ᶜts = ts_gs(thermo_args..., ᶜspecific, ᶜK, Φ(grav, ᶜz), Y.c.ρ)
497498
@. ᶜp = TD.air_pressure(thermo_params, ᶜts)
498499

499500
if turbconv_model isa AbstractEDMF

‎src/cache/prognostic_edmf_precomputed_quantities.jl

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,8 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!(
2020

2121
thermo_params = CAP.thermodynamics_params(p.params)
2222
(; turbconv_model) = p.atmos
23-
(; ᶜΦ,) = p.core
23+
ᶜz = Fields.coordinate_field(axes(Y.c)).z
24+
grav = TDP.grav(thermo_params)
2425
(; ᶜp, ᶜh_tot, ᶜK) = p.precomputed
2526
(; ᶜtke⁰, ᶜρa⁰, ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶜts⁰, ᶜρ⁰, ᶜmse⁰, ᶜq_tot⁰) =
2627
p.precomputed
@@ -44,7 +45,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!(
4445
set_sgs_ᶠu₃!(u₃⁰, ᶠu₃⁰, Y, turbconv_model)
4546
set_velocity_quantities!(ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶠu₃⁰, Y.c.uₕ, ᶠuₕ³)
4647
# @. ᶜK⁰ += ᶜtke⁰
47-
@. ᶜts⁰ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmse⁰ - ᶜΦ, ᶜq_tot⁰)
48+
@. ᶜts⁰ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmse⁰ - Φ(grav, ᶜz), ᶜq_tot⁰)
4849
@. ᶜρ⁰ = TD.air_density(thermo_params, ᶜts⁰)
4950
return nothing
5051
end
@@ -72,7 +73,8 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_draft_and_bc!
7273
thermo_params = CAP.thermodynamics_params(params)
7374
turbconv_params = CAP.turbconv_params(params)
7475

75-
(; ᶜΦ,) = p.core
76+
grav = TDP.grav(thermo_params)
77+
ᶜz = Fields.coordinate_field(Y.c).z
7678
(; ᶜspecific, ᶜp, ᶜh_tot, ᶜK) = p.precomputed
7779
(; ᶜuʲs, ᶠu³ʲs, ᶜKʲs, ᶠKᵥʲs, ᶜtsʲs, ᶜρʲs) = p.precomputed
7880
(; ustar, obukhov_length, buoyancy_flux) = p.precomputed.sfc_conditions
@@ -90,7 +92,8 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_draft_and_bc!
9092

9193
set_velocity_quantities!(ᶜuʲ, ᶠu³ʲ, ᶜKʲ, ᶠu₃ʲ, Y.c.uₕ, ᶠuₕ³)
9294
@. ᶠKᵥʲ = (adjoint(CT3(ᶠu₃ʲ)) * ᶠu₃ʲ) / 2
93-
@. ᶜtsʲ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmseʲ - ᶜΦ, ᶜq_totʲ)
95+
@. ᶜtsʲ =
96+
TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmseʲ - Φ(grav, ᶜz), ᶜq_totʲ)
9497
@. ᶜρʲ = TD.air_density(thermo_params, ᶜtsʲ)
9598

9699
# EDMFX boundary condition:
@@ -153,12 +156,11 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_draft_and_bc!
153156
)
154157

155158
# Then overwrite the prognostic variables at first inetrior point.
156-
ᶜΦ_int_val = Fields.field_values(Fields.level(ᶜΦ, 1))
157159
ᶜtsʲ_int_val = Fields.field_values(Fields.level(ᶜtsʲ, 1))
158160
@. ᶜtsʲ_int_val = TD.PhaseEquil_phq(
159161
thermo_params,
160162
ᶜp_int_val,
161-
ᶜmseʲ_int_val - ᶜΦ_int_val,
163+
ᶜmseʲ_int_val - Φ(grav, ᶜz_int_val),
162164
ᶜq_totʲ_int_val,
163165
)
164166
sgsʲs_ρ_int_val = Fields.field_values(Fields.level(ᶜρʲs.:($j), 1))
@@ -424,7 +426,6 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation
424426
@assert !(p.atmos.moisture_model isa DryModel)
425427

426428
(; params, dt) = p
427-
(; ᶜΦ,) = p.core
428429
thp = CAP.thermodynamics_params(params)
429430
cmp = CAP.microphysics_1m_params(params)
430431

@@ -451,7 +452,6 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation
451452
ᶜqᵣ,
452453
ᶜqₛ,
453454
ᶜtsʲs.:($j),
454-
ᶜΦ,
455455
dt,
456456
cmp,
457457
thp,
@@ -470,7 +470,6 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation
470470
ᶜqᵣ,
471471
ᶜqₛ,
472472
ᶜts⁰,
473-
ᶜΦ,
474473
dt,
475474
cmp,
476475
thp,

‎src/core/core_quantities.jl

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
import ClimaCore.Geometry
2+
3+
f³(lg, global_geom, params) = f³(lg, lg.coordinates, global_geom, params)
4+
5+
f¹²(lg, coord, global_geom, params) =
6+
f¹²(lg, lg.coordinates, global_geom, params)
7+
8+
function f³(
9+
lg,
10+
coord::Geometry.LatLongZPoint,
11+
global_geom::Geometry.DeepSphericalGlobalGeometry,
12+
params,
13+
)
14+
Ω = CAP.Omega(params)
15+
CT3(
16+
CT123(
17+
Geometry.LocalVector(
18+
Geometry.Cartesian123Vector(zero(Ω), zero(Ω), 2 * Ω),
19+
global_geom,
20+
coord,
21+
),
22+
lg,
23+
),
24+
lg,
25+
)
26+
end
27+
28+
function f¹²(
29+
lg,
30+
coord::Geometry.LatLongZPoint,
31+
global_geom::Geometry.DeepSphericalGlobalGeometry,
32+
params,
33+
)
34+
Ω = CAP.Omega(params)
35+
CT12(
36+
CT123(
37+
Geometry.LocalVector(
38+
Geometry.Cartesian123Vector(zero(Ω), zero(Ω), 2 * Ω),
39+
global_geom,
40+
coord,
41+
),
42+
lg,
43+
),
44+
lg,
45+
)
46+
end
47+
48+
# Shallow sphere
49+
function f³(lg, coord::Geometry.LatLongZPoint, global_geom, params)
50+
Ω = CAP.Omega(params)
51+
CT3(Geometry.WVector(2 * Ω * sind(coord.lat), lg), lg)
52+
end
53+
54+
# Shallow cartesian
55+
function f³(lg, coord, global_geom, params)
56+
f = CAP.f_plane_coriolis_frequency(params)
57+
CT3(Geometry.WVector(f, lg), lg)
58+
end
59+
60+
f¹²(lg, coord::Geometry.LatLongZPoint, global_geom, params) =
61+
error("Not supported for $coord, $global_geom.")
62+
f¹²(lg, coord, global_geom, p) =
63+
error("Not supported for $coord, $global_geom.")
64+
65+
Φ(g, z) = g * z

‎src/diagnostics/Diagnostics.jl

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import ClimaCore.Utilities: half
1212
import Thermodynamics as TD
1313

1414
import ..AtmosModel
15+
import ..Φ
1516
import ..AtmosCallback
1617
import ..EveryNSteps
1718

0 commit comments

Comments
 (0)