Skip to content

Commit f880322

Browse files
SbozzoloSbozzolo
authored
Merge pull request #2324 from CliMA/gb/bump
Add `DebugOnly.allow_mismatched_spaces_unsafe`
2 parents 1344d58 + 3219408 commit f880322

5 files changed

Lines changed: 263 additions & 21 deletions

File tree

‎NEWS.md‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,18 @@ ClimaCore.jl Release Notes
44
main
55
-------
66

7-
- Latitude and longitudes were flipped in the default target coordiantes
7+
- Latitude and longitudes were flipped in the default target coordinates
88
in the remapper, leading to non-sense values.
99

10+
### ![][badge-✨feature/enhancement] Add option to disable check on spaces. PR [2322](https://github.com/CliMA/ClimaCore.jl/pull/2322)
11+
12+
`ClimaCore.DebugOnly` now contains a function `allow_mismatched_spaces_unsafe`. When
13+
this function is overridden to return `true`, `ClimaCore` will skip consistency
14+
checks on spaces in broadcasted expressions. This is useful when debugging with
15+
`deepcopy`. See
16+
[documentation](https://clima.github.io/ClimaCore.jl/dev/debugging/#Faster-explorations-when-the-initialization-is-expensive)
17+
for more information.
18+
1019
v0.14.33
1120
-------
1221

‎docs/src/debugging.md‎

Lines changed: 169 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,23 @@
11
# Debugging
22

3+
`ClimaCore` comes with a set of debugging tools for a variety of situations.
4+
These tools are collected in the [`ClimaCore.DebugOnly`](@ref) modules and are
5+
presented in this page.
6+
7+
## Finding where `NaN` arise
8+
39
One of the most challenging tasks that users have is: debug a large simulation
410
that is breaking, e.g., yielding `NaN`s somewhere. This is especially complex
511
for large models with many terms and implicit time-stepping with all the bells
612
and whistles that the CliMA ecosystem offers.
713

8-
ClimaCore has a module, [`ClimaCore.DebugOnly`](@ref), which contains tools for
9-
debugging simulations for these complicated situations.
10-
1114
Because so much data (for example, the solution state, and many cached fields)
12-
is typically contained in ClimaCore data structures, we offer a hook to inspect
13-
this data after any operation that ClimaCore performs.
15+
is typically contained in `ClimaCore` data structures, we offer a hook to
16+
inspect this data after any operation that `ClimaCore` performs.
1417

15-
## Example
18+
### Example
1619

17-
### Print `NaNs` when they are found
20+
#### Print `NaNs` when they are found
1821

1922
In this example, we add a callback that simply prints `NaNs found` every
2023
instance when they are detected in a `ClimaCore` operation.
@@ -46,10 +49,10 @@ If needed, multiple methods of `post_op_callback` can be defined, but here, we
4649
define a general method that checks if `NaN`s are in the given object.
4750

4851
Note that we need `post_op_callback` to be called for a wide variety of inputs
49-
because it is called by many many different functions with many different
50-
objects. Therefore, we recommend that you define `post_op_callback` with a very
51-
general method signature, like the one above and perhaps use `Infiltrator` to
52-
inspect the arguemtns.
52+
because it is called by many different functions with many different objects.
53+
Therefore, we recommend that you define `post_op_callback` with a very general
54+
method signature, like the one above and perhaps use `Infiltrator` to inspect
55+
the arguments.
5356

5457
Now, let us put everything together and demonstrate a complete example:
5558

@@ -72,14 +75,128 @@ ClimaCore.DebugOnly.call_post_op_callback() = false # hide
7275
```
7376
This example should print `NaN` on your standard output.
7477

75-
### Infiltrating
78+
As you see, this only tells us that a `NaN` was found, but not which function
79+
triggered the `NaN`. A simple way to find that is to extend this example with
80+
`Infiltrator`.
81+
82+
#### Infiltrating and Exfiltrating
7683

7784
[Infiltrator.jl](https://github.com/JuliaDebug/Infiltrator.jl) is a simple
7885
debugging tool for Julia packages.
7986

80-
Here is an example, where we can use Infiltrator.jl to find where NaNs is coming
87+
Here is an example, where we can use Infiltrator.jl to find where `NaN`s is coming
8188
from interactively.
8289

90+
##### Infiltrating
91+
92+
Suppose you have a `NaN` in your simulation and what to look at the variables
93+
right before the expression caused the `NaN`. One of the challenges in this is
94+
that you probably have a tower of function being called, many of which belong to
95+
other packages.
96+
97+
Let us simulate this case (note this is a toy example optimized for clarity and
98+
not for performance: the functions below have unnecessary allocations)
99+
```julia
100+
import ClimaCore
101+
using ClimaCore.CommonSpaces
102+
space = CubedSphereSpace(; radius = 10, n_quad_points = 4, h_elem = 10)
103+
myrho = ones(space)
104+
myP = ones(space)
105+
myu = ones(space)
106+
107+
kinetic_energy(field) = field .* field ./ 2
108+
other_energy(field) = field ./ sum(field)
109+
offset_by_one(field) = field .- 1
110+
111+
function specific_energy(rho, P, u)
112+
density_without_restmass = offset_by_one(rho)
113+
return (kinetic_energy(u) .+ other_energy(P)) ./ density_without_restmass
114+
end
115+
116+
function renormalized_energy(rho, P, u)
117+
energy = specific_energy(rho, P, u)
118+
return energy ./ sum(energy)
119+
end
120+
121+
any(isnan, renormalized_energy(myrho, myP, myu)) # true
122+
```
123+
124+
To debug this, we first need to identify where the first `NaN` is produced. We
125+
use `DebugOnly.call_post_op_callback` and infiltrate.
126+
```julia
127+
import Infiltrator # must be in your default environment
128+
ClimaCore.DebugOnly.call_post_op_callback() = true
129+
function ClimaCore.DebugOnly.post_op_callback(result, args...; kwargs...)
130+
has_nans = result isa Number ? isnan(result) : any(isnan, parent(result))
131+
has_inf = result isa Number ? isinf(result) : any(isinf, parent(result))
132+
@infiltrate has_nans || has_inf
133+
end
134+
```
135+
136+
"Infiltrating" means being dropped into a new REPL where in the scope of the
137+
`@infiltrate` macro. `@infiltrate condition` means that we want to infiltrate
138+
only when the `condition` is true (in this case `has_nans || has_inf`).
139+
140+
Now, when we run our example, we will see
141+
```julia
142+
julia> renormalized_energy(myrho, myP, myu)
143+
Infiltrating post_op_callback(::ClimaCore.DataLayouts.IJFH{Float64, 4, Array{Float64, 4}}, ::ClimaCore.DataLayouts.IJFH{Float64, 4, Array{Float64, 4}}, ::Vararg{Any}; kwargs::@Kwargs{})
144+
at REPL[40]:4
145+
infil>
146+
```
147+
Here, we are dropped into a new REPL with full access to the variables in the scope where the `NaN` occurred. However, because of how `post_op_callback`, this is at a low level within `ClimaCore`, which is typically not useful. Hence, the next step is to type `@trace`, which prints out
148+
```julia
149+
[1] post_op_callback(::ClimaCore.DataLayouts.IJFH{…}, ::ClimaCore.DataLayouts.IJFH{…}, ::Vararg{…}; kwargs::@Kwargs{})
150+
at REPL[40]:4
151+
[2] post_op_callback
152+
at REPL[40]:1
153+
[3] copyto!
154+
at ClimaCore.jl/src/DataLayouts/copyto.jl:18
155+
[4] copyto!
156+
at ClimaCore.jl/src/Fields/broadcast.jl:190
157+
[5] copy
158+
at ClimaCore.jl/src/Fields/broadcast.jl:97
159+
[6] materialize
160+
at .julia/juliaup/julia-1.11.4+0.x64.linux.gnu/share/julia/base/broadcast.jl:872
161+
[7] specific_energy(rho::ClimaCore.Fields.Field{…}, P::ClimaCore.Fields.Field{…}, u::ClimaCore.Fields.Field{…})
162+
at REPL[31]:2
163+
[8] renormalized_energy(rho::ClimaCore.Fields.Field{…}, P::ClimaCore.Fields.Field{…}, u::ClimaCore.Fields.Field{…})
164+
at REPL[36]:2
165+
[9] top-level scope
166+
```
167+
168+
`@trace` returns a type-limited stacktrace that we can read backwards until we
169+
see our functions. In this case, we see that the first `NaN` is in
170+
`specific_energy`, so we will investigate that function. We leave the
171+
Infiltrator REPL with `@exit`, disable the `call_post_op_callback`, and move our
172+
`infiltrate` call within the target function:
173+
```julia
174+
ClimaCore.DebugOnly.call_post_op_callback() = false
175+
function specific_energy(rho, P, u)
176+
@infiltrate
177+
density_without_restmass = offset_by_one(rho)
178+
return (kinetic_energy(u) .+ other_energy(P)) ./ density_without_restmass
179+
end
180+
```
181+
Now, when we evaluate our problematic expression (the one at the top level, in this case `renormalized_energy(myrho, myP, myu)`), we will be dropped in a REPL inside `specific_energy`. Here, we have access to `density_without_restmass`, and we notice that it can be zero, leading to the `NaN`.
182+
183+
!!! tip
184+
185+
The infiltrator REPL is different from the normal Julia repl. Type `?` for
186+
some useful commands. You can fetch objects defined in the main REPL by
187+
prepending their name with `Main`. Similarly, if you want to infiltrate inside a
188+
module, prepend `@infiltrate` with `Main` (`Main.@infiltrate`).
189+
190+
Looking at this code, we could have probably guess that the `NaN` comes from a
191+
division from 0, and the real question is how do we get that?. In more complex
192+
cases, the functions are spread across different packages and involve several
193+
different variables. This approach allows one to systematically identify where
194+
things go wrong.
195+
196+
##### Exfiltrating and StructuredPrinting
197+
198+
Let's now see a different way to use Infiltrator, where we move the variables in specific scope to the Main scope in the REPL and do some analysis on it.
199+
83200
```julia
84201
import ClimaCore
85202
import Infiltrator # must be in your default environment
@@ -162,7 +279,7 @@ bc.axes.5::Base.OneTo{Int64}
162279
bc.axes.5.stop::Int64
163280
```
164281

165-
### Caveats
282+
#### Caveats
166283

167284
!!! warn
168285

@@ -183,3 +300,41 @@ bc.axes.5.stop::Int64
183300
as Test.jl may continue running through code execution, until all of the
184301
tests in a given `@testset` are complete, and the result will be that you
185302
will get the _last_ observed instance of `NaN` or `Inf`.
303+
304+
#### Faster explorations when the initialization is expensive
305+
306+
Sometimes, we want to start from a given simulation state and explore different
307+
ideas. For example, we want to run a simulation for 10 days, and then test how
308+
different approaches affect its stability. Sometimes, checkpoints offer a way to
309+
do this, but not everything can be checkpointed.
310+
311+
A simple way to "checkpoint" a simulation is to `deepcopy` its state. This
312+
allows one to step the copy instead of the original one, which can be re-used to
313+
make new copies, allowing for various explorations. `ClimaCore` does not support
314+
this workflow out-of-the-box. The reason for this is that `ClimaCore` uses
315+
pointers to perform certain safety checks, and deepcopies return new pointers
316+
(by definition). To enable this, override the
317+
`DebugOnly.allow_mismatched_spaces_unsafe` function so that it returns true. When
318+
`DebugOnly.allow_mismatched_spaces_unsafe` returns true, `ClimaCore` can mix fields
319+
defined on space that are not identically the same.
320+
321+
Let us look at an example of this.
322+
```julia
323+
import ClimaCore
324+
using ClimaCore.CommonSpaces
325+
other_space = deepcopy(space)
326+
327+
one = ones(space)
328+
other_one = ones(other_space)
329+
330+
one .+ other_one # This throws an error
331+
332+
ClimaCore.DebugOnly.allow_mismatched_spaces_unsafe() = true
333+
one .+ other_one # Now it's fine!
334+
```
335+
336+
!!! warn
337+
338+
`ClimaCore` checks for consistency of spaces to protect you from non-sense
339+
results. If you disable this check, you are responsible to ensure that the
340+
results make sense.

‎src/DebugOnly/DebugOnly.jl‎

Lines changed: 74 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
"""
2-
DebugOnly
2+
DebugOnly
33
44
A module for debugging tools. Note that any tools in here are subject to sudden
55
changes without warning. So, please, do _not_ use any of these tools in
@@ -98,4 +98,77 @@ function print_depth_limited_stack_trace(
9898
end
9999
end
100100

101+
102+
"""
103+
allow_mismatched_spaces_unsafe()
104+
105+
When returning `true`, disable check for consistency among spaces in broadcasted
106+
operations.
107+
108+
By default, `ClimaCore` checks that broadcasted in-place expressions use
109+
consistent spaces (ie, the destination space is the same as the space that the
110+
expression returns). Sometimes, when debugging, it is convenient to disable this
111+
check.
112+
113+
The most common case for this is to allow combining spaces that were
114+
`deepcopied`, given that the consistency check in performed by comparing the
115+
pointer of the spaces, not their contents. In other words, allowing for
116+
mismatched spaces allows one to work with spaces that are identical, but not the
117+
same.
118+
119+
To allow combining mismatched spaces, override this function so that it returns
120+
`true`.
121+
122+
!!! warn
123+
124+
`ClimaCore` checks for consistency of spaces to protect you from non-sense
125+
results. If you disable this check, you are responsible to ensure that the
126+
results make sense.
127+
128+
Example
129+
=======
130+
131+
```julia
132+
julia> import ClimaCore;
133+
134+
julia> using ClimaCore.CommonSpaces;
135+
136+
julia> space = ExtrudedCubedSphereSpace(;
137+
z_elem = 10,
138+
z_min = 0,
139+
z_max = 1,
140+
radius = 10,
141+
h_elem = 10,
142+
n_quad_points = 4,
143+
staggering = CellCenter()
144+
);
145+
146+
julia> other_space = deepcopy(space);
147+
148+
julia> other_space == space
149+
false
150+
151+
julia> one = ones(space);
152+
153+
julia> other_one = ones(other_space);
154+
155+
julia> one .+ other_one
156+
ERROR: Broacasted spaces are not the same.
157+
Stacktrace:
158+
[1] error(s::String)
159+
@ Base ./error.jl:35
160+
[2] error_mismatched_spaces(space1::Type, space2::Type)
161+
@ ClimaCore.Fields ~/repos/ClimaCore.jl/src/Fields/broadcast.jl:227
162+
163+
# Turning `allow_mismatched_spaces_unsafe` on
164+
julia> ClimaCore.DebugOnly.allow_mismatched_spaces_unsafe() = true;
165+
166+
julia> one .+ other_one
167+
Float64-valued Field:
168+
[2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0 … 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0]
169+
"""
170+
function allow_mismatched_spaces_unsafe()
171+
return false
172+
end
173+
101174
end

‎src/Fields/broadcast.jl‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
import ..DebugOnly: allow_mismatched_spaces_unsafe
2+
13
"""
24
AbstractFieldStyle
35
@@ -232,7 +234,7 @@ end
232234
space1::AbstractSpace,
233235
space2::AbstractSpace,
234236
)
235-
if space1 !== space2
237+
if space1 !== space2 && !allow_mismatched_spaces_unsafe()
236238
if Spaces.issubspace(space2, space1)
237239
return space1
238240
elseif Spaces.issubspace(space1, space2)
@@ -266,7 +268,11 @@ end
266268
space1::AbstractSpace,
267269
space2::AbstractSpace,
268270
)
269-
if space1 !== space2
271+
# When DebugOnly.allow_mismatched_spaces_unsafe() returns true, we ignore the check
272+
# and blindly return `space1`. Responsibility is left up to the user to
273+
# handle this correctly. This is useful to work with spaces that are == but
274+
# not ===, e.g., deepcopied spaces.
275+
if space1 !== space2 && !allow_mismatched_spaces_unsafe()
270276
if Spaces.issubspace(space2, space1) ||
271277
Spaces.issubspace(space1, space2)
272278
nothing

‎src/Operators/finitedifference.jl‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import ..Utilities: PlusHalf, half, unionall_type
2+
import ..DebugOnly: allow_mismatched_spaces_unsafe
23
import UnrolledUtilities: unrolled_map
34

45
const AllFiniteDifferenceSpace =
@@ -3889,8 +3890,6 @@ function Base.Broadcast.broadcasted(
38893890
StencilBroadcasted{Style}(promote_bcs(op, FT), args)
38903891
end
38913892

3892-
allow_mismatched_fd_spaces() = false
3893-
38943893
# check that inferred output field space is equal to dest field space
38953894
@noinline inferred_stencil_spaces_error(
38963895
dest_space_type::Type,
@@ -3905,7 +3904,7 @@ function Base.Broadcast.materialize!(
39053904
bc::Base.Broadcast.Broadcasted{Style},
39063905
) where {Style <: AbstractStencilStyle}
39073906
dest_space, result_space = axes(dest), axes(bc)
3908-
if result_space !== dest_space && !allow_mismatched_fd_spaces()
3907+
if result_space !== dest_space && !allow_mismatched_spaces_unsafe()
39093908
# TODO: we pass the types here to avoid stack copying data
39103909
# but this could lead to a confusing error message (same space type but different instances)
39113910
inferred_stencil_spaces_error(typeof(dest_space), typeof(result_space))

0 commit comments

Comments
 (0)