Skip to content

Commit 105484c

Browse files
nefrathenricinefrathenrici
committed
Use CommonGrids to create spaces
1 parent a419813 commit 105484c

File tree

3 files changed

+566
-126
lines changed

3 files changed

+566
-126
lines changed

‎src/ClimaAtmos.jl‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,7 @@ import .Diagnostics as CAD
144144
include(joinpath("callbacks", "get_callbacks.jl"))
145145

146146
include(joinpath("simulation", "AtmosSimulations.jl"))
147+
include(joinpath("simulation", "grids.jl"))
147148

148149
include(joinpath("solver", "model_getters.jl")) # high-level (using parsed_args) model getters
149150
include(joinpath("solver", "type_getters.jl"))

‎src/simulation/grids.jl‎

Lines changed: 386 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,386 @@
1+
"""
2+
Grid constructors that replace the domain-based approach with direct ClimaCore.CommonGrids usage.
3+
4+
These constructors handle topography integration and provide the same interface as the previous
5+
domain system but use ClimaCore grids directly.
6+
"""
7+
8+
using ClimaCore: Geometry, Hypsography, Fields, Spaces, Meshes
9+
using ClimaCore.CommonGrids: ExtrudedCubedSphereGrid, ColumnGrid, Box3DGrid, SliceXZGrid
10+
using ClimaUtilities: SpaceVaryingInputs.SpaceVaryingInput
11+
import .AtmosArtifacts as AA
12+
import ClimaComms
13+
14+
export SphereGrid, ColGrid, BoxGrid, PlaneGrid
15+
16+
abstract type MeshWarpType end
17+
18+
struct LinearWarp <: MeshWarpType end
19+
struct SLEVEWarp <: MeshWarpType end
20+
21+
"""
22+
SphereGrid(FT, params, comms_ctx; kwargs...)
23+
24+
Create an ExtrudedCubedSphereGrid with topography support.
25+
26+
# Arguments
27+
- `FT`: Floating point type
28+
- `params`: ClimaAtmos parameters
29+
- `comms_ctx`: Communications context
30+
31+
# Keyword Arguments
32+
- `z_elem::Int`: Number of vertical elements
33+
- `z_max::Real`: Maximum height
34+
- `z_stretch::Bool`: Whether to use vertical stretching
35+
- `dz_bottom::Real`: Bottom layer thickness for stretching
36+
- `radius::Real`: Earth radius
37+
- `h_elem::Int`: Number of horizontal elements per panel
38+
- `nh_poly::Int`: Polynomial order
39+
- `bubble::Bool`: Enable bubble correction
40+
- `deep_atmosphere::Bool`: Enable deep atmosphere
41+
- `topography::AbstractTopography`: Topography type
42+
- `topography_damping_factor::Real`: Topography damping
43+
- `mesh_warp_type`: MeshWarpType: Mesh warping type ("SLEVE" or "Linear")
44+
- `sleve_eta::Real`: SLEVE parameter
45+
- `sleve_s::Real`: SLEVE parameter
46+
- `topo_smoothing::Bool`: Apply topography smoothing
47+
"""
48+
function SphereGrid(
49+
FT,
50+
params,
51+
comms_ctx;
52+
z_elem::Int = 10,
53+
z_max::Real = 30000.0,
54+
z_stretch::Bool = true,
55+
dz_bottom::Real = 500.0,
56+
radius::Real = 6.371229e6,
57+
h_elem::Int = 6,
58+
nh_poly::Int = 3,
59+
bubble::Bool = false,
60+
deep_atmosphere::Bool = true,
61+
topography::AbstractTopography = NoTopography(),
62+
topography_damping_factor::Real = 5.0,
63+
mesh_warp_type::MeshWarpType = SLEVEWarp(),
64+
sleve_eta::Real = 0.7,
65+
sleve_s::Real = 10.0,
66+
topo_smoothing::Bool = false,
67+
)
68+
n_quad_points = nh_poly + 1
69+
70+
stretch =
71+
z_stretch ? Meshes.HyperbolicTangentStretching{FT}(dz_bottom) : Meshes.Uniform()
72+
73+
hypsography_fun = hypsography_function_from_topography(
74+
topography, topography_damping_factor, mesh_warp_type,
75+
sleve_eta, sleve_s, topo_smoothing, comms_ctx,
76+
)
77+
78+
global_geometry = if deep_atmosphere
79+
Geometry.DeepSphericalGlobalGeometry{FT}(radius)
80+
else
81+
Geometry.ShallowSphericalGlobalGeometry{FT}(radius)
82+
end
83+
84+
grid = ExtrudedCubedSphereGrid(
85+
FT;
86+
z_elem, z_min = 0, z_max, radius, h_elem,
87+
n_quad_points,
88+
device = ClimaComms.device(comms_ctx),
89+
context = comms_ctx,
90+
stretch,
91+
hypsography_fun,
92+
global_geometry,
93+
enable_bubble = bubble,
94+
)
95+
96+
return grid
97+
end
98+
99+
"""
100+
ColGrid(FT, params, comms_ctx; kwargs...)
101+
102+
Create a ColumnGrid.
103+
104+
# Arguments
105+
- `FT`: Floating point type
106+
- `params`: ClimaAtmos parameters
107+
- `comms_ctx`: Communications context
108+
109+
# Keyword Arguments
110+
- `z_elem::Int`: Number of vertical elements
111+
- `z_max::Real`: Maximum height
112+
- `z_stretch::Bool`: Whether to use vertical stretching
113+
- `dz_bottom::Real`: Bottom layer thickness for stretching
114+
"""
115+
function ColGrid(
116+
FT,
117+
params,
118+
comms_ctx;
119+
z_elem::Int = 10,
120+
z_max::Real = 30000.0,
121+
z_stretch::Bool = true,
122+
dz_bottom::Real = 500.0,
123+
)
124+
stretch =
125+
z_stretch ? Meshes.HyperbolicTangentStretching{FT}(dz_bottom) : Meshes.Uniform()
126+
127+
grid = ColumnGrid(
128+
FT;
129+
z_elem, z_min = 0, z_max,
130+
device = ClimaComms.device(comms_ctx),
131+
context = comms_ctx,
132+
stretch,
133+
)
134+
135+
return grid
136+
end
137+
138+
"""
139+
BoxGrid(FT, params, comms_ctx; kwargs...)
140+
141+
Create a Box3DGrid with topography support.
142+
143+
# Arguments
144+
- `FT`: Floating point type
145+
- `params`: ClimaAtmos parameters
146+
- `comms_ctx`: Communications context
147+
148+
# Keyword Arguments
149+
- `x_elem::Int`: Number of x elements
150+
- `x_max::Real`: Maximum x coordinate
151+
- `y_elem::Int`: Number of y elements
152+
- `y_max::Real`: Maximum y coordinate
153+
- `z_elem::Int`: Number of vertical elements
154+
- `z_max::Real`: Maximum height
155+
- `nh_poly::Int`: Polynomial order
156+
- `z_stretch::Bool`: Whether to use vertical stretching
157+
- `dz_bottom::Real`: Bottom layer thickness for stretching
158+
- `bubble::Bool`: Enable bubble correction
159+
- `deep_atmosphere::Bool`: Enable deep atmosphere
160+
- `periodic_x::Bool`: Periodic in x direction
161+
- `periodic_y::Bool`: Periodic in y direction
162+
- `topography::AbstractTopography`: Topography type
163+
- `topography_damping_factor::Real`: Topography damping
164+
- `mesh_warp_type::String`: Mesh warping type ("SLEVE" or "Linear")
165+
- `sleve_eta::Real`: SLEVE parameter
166+
- `sleve_s::Real`: SLEVE parameter
167+
- `topo_smoothing::Bool`: Apply topography smoothing
168+
"""
169+
function BoxGrid(
170+
FT,
171+
params,
172+
comms_ctx;
173+
x_elem::Int = 6,
174+
x_max::Real = 300000.0,
175+
y_elem::Int = 6,
176+
y_max::Real = 300000.0,
177+
z_elem::Int = 10,
178+
z_max::Real = 30000.0,
179+
nh_poly::Int = 3,
180+
z_stretch::Bool = true,
181+
dz_bottom::Real = 500.0,
182+
bubble::Bool = false,
183+
deep_atmosphere::Bool = true,
184+
periodic_x::Bool = true,
185+
periodic_y::Bool = true,
186+
topography::AbstractTopography = NoTopography(),
187+
topography_damping_factor::Real = 5.0,
188+
mesh_warp_type::MeshWarpType = LinearWarp(),
189+
sleve_eta::Real = 0.7,
190+
sleve_s::Real = 10.0,
191+
topo_smoothing::Bool = false,
192+
)
193+
n_quad_points = nh_poly + 1
194+
195+
stretch =
196+
z_stretch ? Meshes.HyperbolicTangentStretching{FT}(dz_bottom) : Meshes.Uniform()
197+
198+
hypsography_fun = hypsography_function_from_topography(
199+
topography, topography_damping_factor, mesh_warp_type,
200+
sleve_eta, sleve_s, topo_smoothing, comms_ctx,
201+
)
202+
203+
grid = Box3DGrid(
204+
FT;
205+
z_elem, x_min = 0, x_max, y_min = 0, y_max, z_min = 0, z_max,
206+
periodic_x, periodic_y, n_quad_points, x_elem, y_elem,
207+
device = ClimaComms.device(comms_ctx),
208+
context = comms_ctx,
209+
stretch,
210+
hypsography_fun,
211+
global_geometry = Geometry.CartesianGlobalGeometry(),
212+
enable_bubble = bubble,
213+
)
214+
215+
return grid
216+
end
217+
218+
"""
219+
PlaneGrid(FT, params, comms_ctx; kwargs...)
220+
221+
Create a SliceXZGrid with topography support.
222+
223+
# Arguments
224+
- `FT`: Floating point type
225+
- `params`: ClimaAtmos parameters
226+
- `comms_ctx`: Communications context
227+
228+
# Keyword Arguments
229+
- `x_elem::Int`: Number of x elements
230+
- `x_max::Real`: Maximum x coordinate
231+
- `z_elem::Int`: Number of vertical elements
232+
- `z_max::Real`: Maximum height
233+
- `nh_poly::Int`: Polynomial order
234+
- `z_stretch::Bool`: Whether to use vertical stretching
235+
- `dz_bottom::Real`: Bottom layer thickness for stretching
236+
- `bubble::Bool`: Enable bubble correction
237+
- `deep_atmosphere::Bool`: Enable deep atmosphere
238+
- `periodic_x::Bool`: Periodic in x direction
239+
- `topography::AbstractTopography`: Topography type
240+
- `topography_damping_factor::Real`: Topography damping
241+
- `mesh_warp_type::String`: Mesh warping type ("SLEVE" or "Linear")
242+
- `sleve_eta::Real`: SLEVE parameter
243+
- `sleve_s::Real`: SLEVE parameter
244+
- `topo_smoothing::Bool`: Apply topography smoothing
245+
"""
246+
function PlaneGrid(
247+
FT,
248+
params,
249+
comms_ctx;
250+
x_elem::Int = 6,
251+
x_max::Real = 300000.0,
252+
z_elem::Int = 10,
253+
z_max::Real = 30000.0,
254+
nh_poly::Int = 3,
255+
z_stretch::Bool = true,
256+
dz_bottom::Real = 500.0,
257+
bubble::Bool = false,
258+
deep_atmosphere::Bool = true,
259+
periodic_x::Bool = true,
260+
topography::AbstractTopography = NoTopography(),
261+
topography_damping_factor::Real = 5.0,
262+
mesh_warp_type::MeshWarpType = LinearWarp(),
263+
sleve_eta::Real = 0.7,
264+
sleve_s::Real = 10.0,
265+
topo_smoothing::Bool = false,
266+
)
267+
n_quad_points = nh_poly + 1
268+
269+
stretch =
270+
z_stretch ? Meshes.HyperbolicTangentStretching{FT}(dz_bottom) : Meshes.Uniform()
271+
272+
hypsography_fun = hypsography_function_from_topography(
273+
topography, topography_damping_factor, mesh_warp_type,
274+
sleve_eta, sleve_s, topo_smoothing, comms_ctx,
275+
)
276+
277+
grid = SliceXZGrid(
278+
FT;
279+
z_elem, x_elem, x_min = 0, x_max, z_min = 0, z_max,
280+
periodic_x,
281+
n_quad_points,
282+
device = ClimaComms.device(comms_ctx),
283+
context = comms_ctx,
284+
stretch,
285+
hypsography_fun,
286+
global_geometry = Geometry.CartesianGlobalGeometry(),
287+
)
288+
289+
return grid
290+
end
291+
292+
"""
293+
hypsography_function_from_topography(
294+
topography, topography_damping_factor, mesh_warp_type,
295+
sleve_eta, sleve_s, topo_smoothing, comms_ctx)
296+
297+
Create a hypsography function that handles topography integration.
298+
This replaces the topography handling logic from make_hybrid_spaces.
299+
"""
300+
function hypsography_function_from_topography(
301+
topography::AbstractTopography,
302+
topography_damping_factor::Real,
303+
mesh_warp_type,
304+
sleve_eta::Real,
305+
sleve_s::Real,
306+
topo_smoothing::Bool,
307+
comms_ctx,
308+
)
309+
# TODO: de-duplicate this with common_spaces.jl
310+
return function (h_grid, z_grid)
311+
FT = eltype(h_grid)
312+
# Create horizontal space to work with topography
313+
h_topology = Spaces.topology(Spaces.SpectralElementSpace2D(h_grid))
314+
h_space = Spaces.SpectralElementSpace2D(h_topology, h_grid.quadrature_style)
315+
316+
if topography isa NoTopography
317+
z_surface = zeros(h_space)
318+
elseif topography isa EarthTopography
319+
z_surface = SpaceVaryingInput(
320+
AA.earth_orography_file_path(;
321+
context = ClimaComms.context(h_space),
322+
),
323+
"z",
324+
h_space,
325+
)
326+
@info "Remapping Earth orography from ETOPO2022 data onto horizontal space"
327+
else
328+
z_surface = SpaceVaryingInput(topography_function(topography), h_space)
329+
@info "Using $(nameof(typeof(topography))) orography"
330+
end
331+
332+
if topography isa NoTopography
333+
hypsography = Hypsography.Flat()
334+
elseif topography isa EarthTopography
335+
mask(x::FT) where {FT} = x * FT(x > 0)
336+
z_surface = @. mask(z_surface)
337+
# diff_cfl = νΔt/Δx²
338+
diff_courant = 0.05 # Arbitrary example value.
339+
Δh_scale = Spaces.node_horizontal_length_scale(h_space)
340+
κ = FT(diff_courant * Δh_scale^2)
341+
maxiter = Int(round(log(topography_damping_factor) / diff_courant))
342+
Hypsography.diffuse_surface_elevation!(
343+
z_surface;
344+
κ,
345+
dt = FT(1),
346+
maxiter,
347+
)
348+
# Coefficient for horizontal diffusion may alternatively be
349+
# determined from the empirical parameters suggested by
350+
# E3SM v1/v2 Topography documentation found here:
351+
# https://acme-climate.atlassian.net/wiki/spaces/DOC/pages/1456603764/V1+Topography+GLL+grids
352+
z_surface = @. mask(z_surface)
353+
if mesh_warp_type isa SLEVEWarp
354+
@info "SLEVE mesh warp"
355+
hypsography = Hypsography.SLEVEAdaption(
356+
Geometry.ZPoint.(z_surface),
357+
FT(sleve_eta),
358+
FT(sleve_s),
359+
)
360+
elseif mesh_warp_type isa LinearWarp
361+
@info "Linear mesh warp"
362+
hypsography =
363+
Hypsography.LinearAdaption(Geometry.ZPoint.(z_surface))
364+
else
365+
@error "Undefined mesh-warping option"
366+
end
367+
else
368+
if topo_smoothing
369+
Hypsography.diffuse_surface_elevation!(z_surface)
370+
end
371+
if mesh_warp_type isa SLEVEWarp
372+
@info "SLEVE mesh warp"
373+
hypsography = Hypsography.SLEVEAdaption(
374+
Geometry.ZPoint.(z_surface),
375+
FT(sleve_eta),
376+
FT(sleve_s),
377+
)
378+
elseif mesh_warp_type isa LinearWarp
379+
@info "Linear mesh warp"
380+
hypsography =
381+
Hypsography.LinearAdaption(Geometry.ZPoint.(z_surface))
382+
end
383+
end
384+
return hypsography
385+
end
386+
end

0 commit comments

Comments
 (0)