-
Notifications
You must be signed in to change notification settings - Fork 191
/
run.jl
224 lines (164 loc) · 6.9 KB
/
run.jl
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
using Oceananigans.Fields: set!
using Oceananigans.OutputWriters: WindowedTimeAverage, checkpoint_superprefix
using Oceananigans.TimeSteppers: QuasiAdamsBashforth2TimeStepper, RungeKutta3TimeStepper, update_state!, next_time, unit_time
using Oceananigans: AbstractModel, run_diagnostic!, write_output!
import Oceananigans: initialize!
import Oceananigans.OutputWriters: checkpoint_path, set!
import Oceananigans.TimeSteppers: time_step!
import Oceananigans.Utils: aligned_time_step
# Simulations are for running
#####
##### Time-step "alignment" with output and callbacks scheduled on TimeInterval
#####
function collect_scheduled_activities(sim)
writers = values(sim.output_writers)
callbacks = values(sim.callbacks)
return tuple(writers..., callbacks...)
end
function schedule_aligned_Δt(sim, aligned_Δt)
clock = sim.model.clock
activities = collect_scheduled_activities(sim)
for activity in activities
aligned_Δt = aligned_time_step(activity.schedule, clock, aligned_Δt)
end
return aligned_Δt
end
"""
aligned_time_step(sim, Δt)
Return a time step 'aligned' with `sim.stop_time`, output writer schedules,
and callback schedules. Alignment with `sim.stop_time` takes precedence.
"""
function aligned_time_step(sim::Simulation, Δt)
clock = sim.model.clock
aligned_Δt = Δt
# Align time step with output writing and callback execution
aligned_Δt = schedule_aligned_Δt(sim, aligned_Δt)
# Align time step with simulation stop time
aligned_Δt = min(aligned_Δt, unit_time(sim.stop_time - clock.time))
# Temporary fix for https://github.com/CliMA/Oceananigans.jl/issues/1280
aligned_Δt = aligned_Δt <= 0 ? Δt : aligned_Δt
return aligned_Δt
end
"""
run!(simulation; pickup=false)
Run a `simulation` until one of `simulation.stop_criteria` evaluates `true`.
The simulation will then stop.
# Picking simulations up from a checkpoint
Simulations are "picked up" from a checkpoint if `pickup` is either `true`, a `String`, or an
`Integer` greater than 0.
Picking up a simulation sets field and tendency data to the specified checkpoint,
leaving all other model properties unchanged.
Possible values for `pickup` are:
* `pickup=true` picks a simulation up from the latest checkpoint associated with
the `Checkpointer` in `simulation.output_writers`.
* `pickup=iteration::Int` picks a simulation up from the checkpointed file associated
with `iteration` and the `Checkpointer` in `simulation.output_writers`.
* `pickup=filepath::String` picks a simulation up from checkpointer data in `filepath`.
Note that `pickup=true` and `pickup=iteration` fails if `simulation.output_writers` contains
more than one checkpointer.
"""
function run!(sim; pickup=false)
if we_want_to_pickup(pickup)
checkpoint_file_path = checkpoint_path(pickup, sim.output_writers)
set!(sim.model, checkpoint_file_path)
end
sim.initialized = false
sim.running = true
sim.run_wall_time = 0.0
while sim.running
time_step!(sim)
end
return nothing
end
const ModelCallsite = Union{TendencyCallsite, UpdateStateCallsite}
""" Step `sim`ulation forward by one time step. """
function time_step!(sim::Simulation)
start_time_step = time_ns()
model_callbacks = Tuple(cb for cb in values(sim.callbacks) if cb.callsite isa ModelCallsite)
if !(sim.initialized) # execute initialization step
initialize!(sim)
initialize!(sim.model)
if sim.running # check that initialization didn't stop time-stepping
if sim.verbose
@info "Executing initial time step..."
start_time = time_ns()
end
Δt = aligned_time_step(sim, sim.Δt)
time_step!(sim.model, Δt, callbacks=model_callbacks)
if sim.verbose
elapsed_initial_step_time = prettytime(1e-9 * (time_ns() - start_time))
@info " ... initial time step complete ($elapsed_initial_step_time)."
end
else
@warn "Simulation stopped during initialization."
end
else # business as usual...
Δt = aligned_time_step(sim, sim.Δt)
time_step!(sim.model, Δt, callbacks=model_callbacks)
end
# Callbacks and callback-like things
for diag in values(sim.diagnostics)
diag.schedule(sim.model) && run_diagnostic!(diag, sim.model)
end
for callback in values(sim.callbacks)
callback.callsite isa TimeStepCallsite && callback.schedule(sim.model) && callback(sim)
end
for writer in values(sim.output_writers)
writer.schedule(sim.model) && write_output!(writer, sim.model)
end
end_time_step = time_ns()
# Increment the wall clock
sim.run_wall_time += 1e-9 * (end_time_step - start_time_step)
return nothing
end
#####
##### Simulation initialization
#####
add_dependency!(diagnostics, output) = nothing # fallback
add_dependency!(diags, wta::WindowedTimeAverage) = wta ∈ values(diags) || push!(diags, wta)
add_dependencies!(diags, writer) = [add_dependency!(diags, out) for out in values(writer.outputs)]
add_dependencies!(sim, ::Checkpointer) = nothing # Checkpointer does not have "outputs"
we_want_to_pickup(pickup::Bool) = pickup
we_want_to_pickup(pickup::Integer) = true
we_want_to_pickup(pickup::String) = true
we_want_to_pickup(pickup) = throw(ArgumentError("Cannot run! with pickup=$pickup"))
"""
initialize!(sim::Simulation, pickup=false)
Initialize a simulation:
- Update the auxiliary state of the simulation (filling halo regions, computing auxiliary fields)
- Evaluate all diagnostics, callbacks, and output writers if sim.model.clock.iteration == 0
- Add diagnostics that "depend" on output writers
"""
function initialize!(sim::Simulation)
if sim.verbose
@info "Initializing simulation..."
start_time = time_ns()
end
model = sim.model
clock = model.clock
update_state!(model)
# Output and diagnostics initialization
[add_dependencies!(sim.diagnostics, writer) for writer in values(sim.output_writers)]
# Reset! the model time-stepper, evaluate all diagnostics, and write all output at first iteration
if clock.iteration == 0
reset!(timestepper(sim.model))
# Initialize schedules and run diagnostics, callbacks, and output writers
for diag in values(sim.diagnostics)
diag.schedule(sim.model)
run_diagnostic!(diag, model)
end
for callback in values(sim.callbacks)
callback.callsite isa TimeStepCallsite && initialize!(callback, sim)
end
for writer in values(sim.output_writers)
writer.schedule(sim.model)
write_output!(writer, model)
end
end
sim.initialized = true
if sim.verbose
initialization_time = prettytime(1e-9 * (time_ns() - start_time))
@info " ... simulation initialization complete ($initialization_time)"
end
return nothing
end