Skip to content

Commit e7a1439

Browse files
ray-chewray-chew
committed
Fix for #2358 
1 parent 450b39c commit e7a1439

File tree

5 files changed

+162
-70
lines changed

5 files changed

+162
-70
lines changed

‎src/Grids/level.jl

Lines changed: 20 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,19 @@
11
struct LevelGrid{
22
G <: AbstractExtrudedFiniteDifferenceGrid,
33
L <: Union{Int, PlusHalf{Int}},
4+
Q <: Quadratures.QuadratureStyle,
45
} <: AbstractGrid
56
full_grid::G
67
level::L
8+
quadrature_style::Q
79
end
810

9-
quadrature_style(levelgrid::LevelGrid) =
10-
quadrature_style(levelgrid.full_grid.horizontal_grid)
11+
quadrature_style(levelgrid::LevelGrid) = levelgrid.quadrature_style
1112

1213
level(
1314
grid::AbstractExtrudedFiniteDifferenceGrid,
1415
level::Union{Int, PlusHalf{Int}},
15-
) = LevelGrid(grid, level)
16+
) = LevelGrid(grid, level, quadrature_style(grid))
1617

1718
topology(levelgrid::LevelGrid) = topology(levelgrid.full_grid)
1819

@@ -21,22 +22,28 @@ topology(levelgrid::LevelGrid) = topology(levelgrid.full_grid)
2122
# need to extract the weights at a particular level.
2223
dss_weights(levelgrid::LevelGrid, _) = dss_weights(levelgrid.full_grid, nothing)
2324

24-
local_geometry_type(::Type{LevelGrid{G, L}}) where {G, L} =
25+
local_geometry_type(::Type{LevelGrid{G, L, Q}}) where {G, L, Q} =
2526
local_geometry_type(G)
26-
local_geometry_data(levelgrid::LevelGrid{<:Any, Int}, ::Nothing) = level(
27+
28+
local_geometry_data(levelgrid::LevelGrid{<:Any, Int, <:Any}, ::Nothing) = level(
2729
local_geometry_data(levelgrid.full_grid, CellCenter()),
2830
levelgrid.level,
2931
)
30-
local_geometry_data(levelgrid::LevelGrid{<:Any, PlusHalf{Int}}, ::Nothing) =
31-
level(
32-
local_geometry_data(levelgrid.full_grid, CellFace()),
33-
levelgrid.level + half,
34-
)
35-
global_geometry(levlgrid::LevelGrid) = global_geometry(levlgrid.full_grid)
32+
local_geometry_data(
33+
levelgrid::LevelGrid{<:Any, PlusHalf{Int}, <:Any},
34+
::Nothing,
35+
) = level(
36+
local_geometry_data(levelgrid.full_grid, CellFace()),
37+
levelgrid.level + half,
38+
)
39+
global_geometry(levelgrid::LevelGrid) = global_geometry(levelgrid.full_grid)
3640

3741
## GPU compatibility
38-
Adapt.adapt_structure(to, grid::LevelGrid) =
39-
LevelGrid(Adapt.adapt(to, grid.full_grid), grid.level)
42+
Adapt.adapt_structure(to, grid::LevelGrid) = LevelGrid(
43+
Adapt.adapt(to, grid.full_grid),
44+
grid.level,
45+
quadrature_style(grid),
46+
)
4047

4148
## aliases
4249
const LevelCubedSphereSpectralElementGrid2D =

‎src/InputOutput/readers.jl

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -522,7 +522,24 @@ function read_grid_new(reader, name)
522522
else
523523
level = attrs(group)["level_half"] + half
524524
end
525-
return Grids.LevelGrid(full_grid, level)
525+
526+
# Check if quadrature attributes exist in current group, otherwise use horizontal_grid
527+
if haskey(attrs(group), "quadrature_num_points") &&
528+
haskey(attrs(group), "quadrature_type")
529+
group = reader.file["grids/horizontal_grid"]
530+
npts = attrs(group)["quadrature_num_points"]
531+
quadrature_style =
532+
_scan_quadrature_style(attrs(group)["quadrature_type"], npts)
533+
else
534+
horizontal_group = reader.file["grids/horizontal_grid"]
535+
npts = attrs(horizontal_group)["quadrature_num_points"]
536+
quadrature_style = _scan_quadrature_style(
537+
attrs(horizontal_group)["quadrature_type"],
538+
npts,
539+
)
540+
end
541+
542+
return Grids.LevelGrid(full_grid, level, quadrature_style)
526543
else
527544
error("Unsupported grid type $type")
528545
end

‎test/Fields/unit_field.jl

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -138,7 +138,7 @@ end
138138
)
139139

140140
@test begin
141-
@. c.∑ab = mapreduce(*, +, c.a, c.b)
141+
@. c.∑ab = mapreduce(*,+,c.a,c.b)
142142
true
143143
end broken = dev isa ClimaComms.CUDADevice
144144
end
@@ -689,6 +689,31 @@ end
689689
end
690690
end
691691

692+
@testset "Levels of nonlocal Fields and nonlocal Field broadcasts" begin
693+
FT = Float64
694+
gradh = Operators.Gradient()
695+
# Todo: Make this work over all spaces; currently broken for everything else.
696+
for space in (
697+
TU.CenterExtrudedFiniteDifferenceSpace(FT),
698+
TU.FaceExtrudedFiniteDifferenceSpace(FT),
699+
)
700+
TU.levelable(space) || continue
701+
field = fill((; x = FT(1)), space)
702+
703+
level_of_field = Fields.Field(
704+
Spaces.level(Fields.field_values(gradh.(field.x)), 1),
705+
Spaces.level(space, TU.fc_index(1, space)),
706+
)
707+
708+
@test level_of_field ==
709+
gradh.(Spaces.level((field.x), TU.fc_index(1, space)))
710+
711+
@test level_of_field == Base.materialize(
712+
gradh.(Spaces.level(lazy.(field.x), TU.fc_index(1, space))),
713+
)
714+
end
715+
end
716+
692717
@testset "Columns of Fields and Field broadcasts" begin
693718
FT = Float64
694719
for space in TU.all_spaces(FT)

‎test/InputOutput/unit_hybrid3dcubedsphere_topography.jl

Lines changed: 58 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,11 @@ using ClimaCore:
1313
InputOutput,
1414
Grids
1515

16+
@isdefined(TU) || include(
17+
joinpath(pkgdir(ClimaCore), "test", "TestUtilities", "TestUtilities.jl"),
18+
);
19+
import .TestUtilities as TU;
20+
1621
using ClimaComms
1722
const comms_ctx = ClimaComms.context(ClimaComms.CPUSingleThreaded())
1823
pid, nprocs = ClimaComms.init(comms_ctx)
@@ -49,13 +54,12 @@ end
4954
)
5055

5156

52-
z_surface =
53-
Geometry.ZPoint.(
54-
z_max / 8 .* (
55-
cosd.(Fields.coordinate_field(h_space).lat) .+
56-
cosd.(Fields.coordinate_field(h_space).long) .+ 1
57-
)
58-
)
57+
z_surface = Geometry.ZPoint.(
58+
z_max / 8 .* (
59+
cosd.(Fields.coordinate_field(h_space).lat) .+
60+
cosd.(Fields.coordinate_field(h_space).long) .+ 1
61+
),
62+
)
5963

6064
z_mesh = Meshes.IntervalMesh(z_domain, nelems = z_elem)
6165
z_topology = Topologies.IntervalTopology(device, z_mesh)
@@ -90,3 +94,50 @@ end
9094
end
9195
end
9296
end
97+
98+
99+
@testset "HDF5 restart test for a Named Tuple of Levels of a 3D hybrid cubed sphere for deep" begin
100+
# This I/O is used for the computation of the topographic drag
101+
FT = Float32
102+
103+
for space in (
104+
TU.CenterExtrudedFiniteDifferenceSpace(FT, context = comms_ctx),
105+
TU.FaceExtrudedFiniteDifferenceSpace(FT, context = comms_ctx),
106+
)
107+
TU.levelable(space) || continue
108+
field = fill((; x = FT(1)), space)
109+
110+
level_of_field = Fields.Field(
111+
Spaces.level(Fields.field_values(field.x), 1),
112+
Spaces.level(space, TU.fc_index(1, space)),
113+
)
114+
115+
fake_drag = fill(
116+
(;
117+
t11 = FT(0.0),
118+
t12 = FT(0.0),
119+
t21 = FT(0.0),
120+
t22 = FT(0.0),
121+
hmin = FT(0.0),
122+
hmax = FT(0.0),
123+
),
124+
axes(level_of_field),
125+
)
126+
127+
# write field vector to hdf5 file
128+
InputOutput.HDF5Writer(filename, comms_ctx) do writer
129+
InputOutput.write!(writer, fake_drag, "fake_drag")
130+
end
131+
132+
InputOutput.HDF5Reader(filename, comms_ctx) do reader
133+
restart_fake_drag = InputOutput.read_field(reader, "fake_drag") # read fieldvector from hdf5 file
134+
135+
# The underlying space is of a different instance, so we cannot use == to check for equivalence.
136+
# Instead, we make sure that the values and types are the same.
137+
@test typeof(restart_fake_drag) == typeof(fake_drag)
138+
@test maximum(
139+
abs.(parent(fake_drag.t21) .- parent(restart_fake_drag.t21)),
140+
) == 0.0f0
141+
end
142+
end
143+
end

‎test/Operators/hybrid/extruded_3dbox_cuda.jl

Lines changed: 40 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -71,32 +71,28 @@ end
7171
coords_gpu = Fields.coordinate_field(hv_center_space_gpu)
7272

7373
# Define CPU & GPU fields
74-
x_cpu =
75-
Geometry.UVWVector.(
76-
sin.(coords_cpu.x .+ 2 .* coords_cpu.y .+ 3 .* coords_cpu.z),
77-
0.0,
78-
0.0,
79-
)
74+
x_cpu = Geometry.UVWVector.(
75+
sin.(coords_cpu.x .+ 2 .* coords_cpu.y .+ 3 .* coords_cpu.z),
76+
0.0,
77+
0.0,
78+
)
8079
f_cpu = sin.(coords_cpu.x .+ 2 .* coords_cpu.y .+ 3 .* coords_cpu.z)
81-
g_cpu =
82-
Geometry.UVWVector.(
83-
sin.(coords_cpu.x),
84-
2 .* cos.(coords_cpu.y .+ coords_cpu.x .+ coords_cpu.z),
85-
cos.(coords_cpu.z),
86-
)
87-
x_gpu =
88-
Geometry.UVWVector.(
89-
sin.(coords_gpu.x .+ 2 .* coords_gpu.y .+ 3 .* coords_gpu.z),
90-
0.0,
91-
0.0,
92-
)
80+
g_cpu = Geometry.UVWVector.(
81+
sin.(coords_cpu.x),
82+
2 .* cos.(coords_cpu.y .+ coords_cpu.x .+ coords_cpu.z),
83+
cos.(coords_cpu.z),
84+
)
85+
x_gpu = Geometry.UVWVector.(
86+
sin.(coords_gpu.x .+ 2 .* coords_gpu.y .+ 3 .* coords_gpu.z),
87+
0.0,
88+
0.0,
89+
)
9390
f_gpu = sin.(coords_gpu.x .+ 2 .* coords_gpu.y .+ 3 .* coords_gpu.z)
94-
g_gpu =
95-
Geometry.UVWVector.(
96-
sin.(coords_gpu.x),
97-
2 .* cos.(coords_gpu.y .+ coords_gpu.x .+ coords_gpu.z),
98-
cos.(coords_gpu.z),
99-
)
91+
g_gpu = Geometry.UVWVector.(
92+
sin.(coords_gpu.x),
93+
2 .* cos.(coords_gpu.y .+ coords_gpu.x .+ coords_gpu.z),
94+
cos.(coords_gpu.z),
95+
)
10096

10197
# Test grad operator
10298
grad = Operators.Gradient()
@@ -132,32 +128,28 @@ end
132128
coords_gpu = Fields.coordinate_field(hv_center_space_gpu)
133129

134130
# Define CPU & GPU fields
135-
x_cpu =
136-
Geometry.UVWVector.(
137-
sin.(coords_cpu.x .+ 2 .* coords_cpu.y .+ 3 .* coords_cpu.z),
138-
0.0,
139-
0.0,
140-
)
131+
x_cpu = Geometry.UVWVector.(
132+
sin.(coords_cpu.x .+ 2 .* coords_cpu.y .+ 3 .* coords_cpu.z),
133+
0.0,
134+
0.0,
135+
)
141136
f_cpu = sin.(coords_cpu.x .+ 2 .* coords_cpu.y .+ 3 .* coords_cpu.z)
142-
g_cpu =
143-
Geometry.UVWVector.(
144-
sin.(coords_cpu.x),
145-
2 .* cos.(coords_cpu.y .+ coords_cpu.x .+ coords_cpu.z),
146-
cos.(coords_cpu.z),
147-
)
148-
x_gpu =
149-
Geometry.UVWVector.(
150-
sin.(coords_gpu.x .+ 2 .* coords_gpu.y .+ 3 .* coords_gpu.z),
151-
0.0,
152-
0.0,
153-
)
137+
g_cpu = Geometry.UVWVector.(
138+
sin.(coords_cpu.x),
139+
2 .* cos.(coords_cpu.y .+ coords_cpu.x .+ coords_cpu.z),
140+
cos.(coords_cpu.z),
141+
)
142+
x_gpu = Geometry.UVWVector.(
143+
sin.(coords_gpu.x .+ 2 .* coords_gpu.y .+ 3 .* coords_gpu.z),
144+
0.0,
145+
0.0,
146+
)
154147
f_gpu = sin.(coords_gpu.x .+ 2 .* coords_gpu.y .+ 3 .* coords_gpu.z)
155-
g_gpu =
156-
Geometry.UVWVector.(
157-
sin.(coords_gpu.x),
158-
2 .* cos.(coords_gpu.y .+ coords_gpu.x .+ coords_gpu.z),
159-
cos.(coords_gpu.z),
160-
)
148+
g_gpu = Geometry.UVWVector.(
149+
sin.(coords_gpu.x),
150+
2 .* cos.(coords_gpu.y .+ coords_gpu.x .+ coords_gpu.z),
151+
cos.(coords_gpu.z),
152+
)
161153

162154
# Test weak grad operator
163155
wgrad = Operators.WeakGradient()

0 commit comments

Comments
 (0)