Skip to content

Commit 1d8f874

Browse files
ArnoStrouwenArnoStrouwen
committed
reapply formatter
1 parent 082e7c3 commit 1d8f874

32 files changed

+391
-263
lines changed

‎.JuliaFormatter.toml

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,3 @@
11
style = "sciml"
2-
format_markdown = true
2+
format_markdown = true
3+
format_docstrings = true

‎docs/pages.jl

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,15 @@
22

33
pages = ["index.md",
44
"getting_started.md",
5-
"Tutorials" => Any["GPU Ensembles" => Any["tutorials/gpu_ensemble_basic.md",
5+
"Tutorials" => Any[
6+
"GPU Ensembles" => Any["tutorials/gpu_ensemble_basic.md",
67
"tutorials/parallel_callbacks.md",
78
"tutorials/multigpu.md",
89
"tutorials/lower_level_api.md",
910
"tutorials/weak_order_conv_sde.md"],
1011
"Within-Method GPU" => Any["tutorials/within_method_gpu.md"]],
11-
"Examples" => Any["GPU Ensembles" => Any["examples/sde.md",
12+
"Examples" => Any[
13+
"GPU Ensembles" => Any["examples/sde.md",
1214
"examples/ad.md",
1315
"examples/reductions.md"],
1416
"Within-Method GPU" => Any["examples/reaction_diffusion.md",
@@ -17,5 +19,5 @@ pages = ["index.md",
1719
"manual/ensemblegpuarray.md",
1820
"manual/backends.md",
1921
"manual/optimal_trajectories.md",
20-
"manual/choosing_ensembler.md"],
22+
"manual/choosing_ensembler.md"]
2123
]

‎docs/src/examples/sde.md

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ prob = SDEProblem(lorenz, multiplicative_noise, u0, tspan, p)
2727
const pre_p = [rand(Float32, 3) for i in 1:10_000]
2828
prob_func = (prob, i, repeat) -> remake(prob, p = pre_p[i] .* p)
2929
monteprob = EnsembleProblem(prob, prob_func = prob_func)
30-
sol = solve(monteprob, SOSRI(), EnsembleGPUArray(CUDA.CUDABackend()), trajectories = 10_000,
30+
sol = solve(
31+
monteprob, SOSRI(), EnsembleGPUArray(CUDA.CUDABackend()), trajectories = 10_000,
3132
saveat = 1.0f0)
3233
```

‎docs/src/tutorials/gpu_ensemble_basic.md

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,8 @@ Changing this to being GPU-parallelized is as simple as changing the ensemble me
2525
`EnsembleGPUArray`:
2626

2727
```@example lorenz
28-
sol = solve(monteprob, Tsit5(), EnsembleGPUArray(CUDA.CUDABackend()), trajectories = 10_000,
28+
sol = solve(
29+
monteprob, Tsit5(), EnsembleGPUArray(CUDA.CUDABackend()), trajectories = 10_000,
2930
saveat = 1.0f0);
3031
```
3132

‎docs/src/tutorials/lower_level_api.md

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -89,12 +89,14 @@ end
8989
9090
## Finally use the lower API for faster solves! (Fixed time-stepping)
9191
92-
@time CUDA.@sync sol = DiffEqGPU.vectorized_map_solve(probs, Tsit5(), EnsembleGPUArray(0.0),
92+
@time CUDA.@sync sol = DiffEqGPU.vectorized_map_solve(
93+
probs, Tsit5(), EnsembleGPUArray(0.0),
9394
batch, false, dt = 0.001f0,
9495
save_everystep = false, dense = false)
9596
9697
## Adaptive time-stepping (Notice the boolean argument)
97-
@time CUDA.@sync sol = DiffEqGPU.vectorized_map_solve(probs, Tsit5(), EnsembleGPUArray(0.0),
98+
@time CUDA.@sync sol = DiffEqGPU.vectorized_map_solve(
99+
probs, Tsit5(), EnsembleGPUArray(0.0),
98100
batch, true, dt = 0.001f0,
99101
save_everystep = false, dense = false)
100102
```

‎src/algorithms.jl

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ struct EnsembleCPUArray <: EnsembleArrayAlgorithm end
1616

1717
"""
1818
```julia
19-
EnsembleGPUArray(backend,cpu_offload = 0.2)
19+
EnsembleGPUArray(backend, cpu_offload = 0.2)
2020
```
2121
2222
An `EnsembleArrayAlgorithm` which utilizes the GPU kernels to parallelize each ODE solve
@@ -73,13 +73,14 @@ function lorenz(du, u, p, t)
7373
du[3] = u[1] * u[2] - p[3] * u[3]
7474
end
7575
76-
u0 = Float32[1.0;0.0;0.0]
77-
tspan = (0.0f0,100.0f0)
78-
p = [10.0f0,28.0f0,8/3f0]
79-
prob = ODEProblem(lorenz,u0,tspan,p)
80-
prob_func = (prob,i,repeat) -> remake(prob,p=rand(Float32,3).*p)
81-
monteprob = EnsembleProblem(prob, prob_func = prob_func, safetycopy=false)
82-
@time sol = solve(monteprob,Tsit5(),EnsembleGPUArray(CUDADevice()),trajectories=10_000,saveat=1.0f0)
76+
u0 = Float32[1.0; 0.0; 0.0]
77+
tspan = (0.0f0, 100.0f0)
78+
p = [10.0f0, 28.0f0, 8 / 3.0f0]
79+
prob = ODEProblem(lorenz, u0, tspan, p)
80+
prob_func = (prob, i, repeat) -> remake(prob, p = rand(Float32, 3) .* p)
81+
monteprob = EnsembleProblem(prob, prob_func = prob_func, safetycopy = false)
82+
@time sol = solve(monteprob, Tsit5(), EnsembleGPUArray(CUDADevice()),
83+
trajectories = 10_000, saveat = 1.0f0)
8384
```
8485
"""
8586
struct EnsembleGPUArray{Backend} <: EnsembleArrayAlgorithm
@@ -89,7 +90,7 @@ end
8990

9091
"""
9192
```julia
92-
EnsembleGPUKernel(backend,cpu_offload = 0.2)
93+
EnsembleGPUKernel(backend, cpu_offload = 0.2)
9394
```
9495
9596
A massively-parallel ensemble algorithm which generates a unique GPU kernel for the entire
@@ -146,7 +147,7 @@ prob_func = (prob, i, repeat) -> remake(prob, p = (@SVector rand(Float32, 3)) .*
146147
monteprob = EnsembleProblem(prob, prob_func = prob_func, safetycopy = false)
147148
148149
@time sol = solve(monteprob, GPUTsit5(), EnsembleGPUKernel(), trajectories = 10_000,
149-
adaptive = false, dt = 0.1f0)
150+
adaptive = false, dt = 0.1f0)
150151
```
151152
"""
152153
struct EnsembleGPUKernel{Dev} <: EnsembleKernelAlgorithm

‎src/ensemblegpuarray/lowerlevel_solve.jl

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -3,23 +3,22 @@ Lower level API for `EnsembleArrayAlgorithm`. Avoids conversion of solution to C
33
44
```julia
55
vectorized_map_solve(probs, alg,
6-
ensemblealg::Union{EnsembleArrayAlgorithm}, I,
7-
adaptive)
6+
ensemblealg::Union{EnsembleArrayAlgorithm}, I,
7+
adaptive)
88
```
99
1010
## Arguments
1111
12-
- `probs`: the GPU-setup problems generated by the ensemble.
13-
- `alg`: the kernel-based differential equation solver. Most of the solvers from OrdinaryDiffEq.jl
14-
are supported.
15-
- `ensemblealg`: The `EnsembleGPUArray()` algorithm.
16-
- `I`: The iterator argument. Can be set to for e.g. 1:10_000 to simulate 10,000 trajectories.
17-
- `adaptive`: The Boolean argument for time-stepping. Use `true` to enable adaptive time-stepping.
12+
- `probs`: the GPU-setup problems generated by the ensemble.
13+
- `alg`: the kernel-based differential equation solver. Most of the solvers from OrdinaryDiffEq.jl
14+
are supported.
15+
- `ensemblealg`: The `EnsembleGPUArray()` algorithm.
16+
- `I`: The iterator argument. Can be set to for e.g. 1:10_000 to simulate 10,000 trajectories.
17+
- `adaptive`: The Boolean argument for time-stepping. Use `true` to enable adaptive time-stepping.
1818
1919
## Keyword Arguments
2020
2121
Only a subset of the common solver arguments are supported.
22-
2322
"""
2423
function vectorized_map_solve end
2524

‎src/ensemblegpukernel/callbacks.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ struct GPUContinuousCallback{F1, F2, F3, F4, F5, F6, T, T2, T3, I, R} <:
4646
interp_points, save_positions::F6, dtrelax::R, abstol::T,
4747
reltol::T2,
4848
repeat_nudge::T3) where {F1, F2, F3, F4, F5, F6, T, T2,
49-
T3, I, R,
49+
T3, I, R
5050
}
5151
if save_positions != (false, false)
5252
error("Callback `save_positions` are incompatible with kernel-based GPU ODE solvers due requiring static sizing. Please ensure `save_positions = (false,false)` is set in all callback definitions used with such solvers.")

‎src/ensemblegpukernel/integrators/integrator_utils.jl

Lines changed: 47 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -10,11 +10,12 @@ function build_adaptive_controller_cache(alg::A, ::Type{T}) where {A, T}
1010
return beta1, beta2, qmax, qmin, gamma, qoldinit, qold
1111
end
1212

13-
@inline function savevalues!(integrator::DiffEqBase.AbstractODEIntegrator{
13+
@inline function savevalues!(
14+
integrator::DiffEqBase.AbstractODEIntegrator{
1415
AlgType,
1516
IIP,
1617
S,
17-
T,
18+
T
1819
}, ts,
1920
us,
2021
force = false) where {AlgType <: GPUODEAlgorithm, IIP, S, T}
@@ -45,23 +46,25 @@ end
4546
saved, savedexactly
4647
end
4748

48-
@inline function DiffEqBase.terminate!(integrator::DiffEqBase.AbstractODEIntegrator{AlgType,
49+
@inline function DiffEqBase.terminate!(
50+
integrator::DiffEqBase.AbstractODEIntegrator{AlgType,
4951
IIP, S,
5052
T},
5153
retcode = ReturnCode.Terminated) where {
5254
AlgType <:
5355
GPUODEAlgorithm,
5456
IIP,
5557
S,
56-
T,
58+
T
5759
}
5860
integrator.retcode = retcode
5961
end
6062

61-
@inline function apply_discrete_callback!(integrator::DiffEqBase.AbstractODEIntegrator{
63+
@inline function apply_discrete_callback!(
64+
integrator::DiffEqBase.AbstractODEIntegrator{
6265
AlgType,
6366
IIP,
64-
S, T,
67+
S, T
6568
},
6669
ts, us,
6770
callback::GPUDiscreteCallback) where {
@@ -79,10 +82,11 @@ end
7982
integrator.u_modified, saved_in_cb
8083
end
8184

82-
@inline function apply_discrete_callback!(integrator::DiffEqBase.AbstractODEIntegrator{
85+
@inline function apply_discrete_callback!(
86+
integrator::DiffEqBase.AbstractODEIntegrator{
8387
AlgType,
8488
IIP,
85-
S, T,
89+
S, T
8690
},
8791
ts, us,
8892
callback::GPUDiscreteCallback,
@@ -93,10 +97,11 @@ end
9397
args...)
9498
end
9599

96-
@inline function apply_discrete_callback!(integrator::DiffEqBase.AbstractODEIntegrator{
100+
@inline function apply_discrete_callback!(
101+
integrator::DiffEqBase.AbstractODEIntegrator{
97102
AlgType,
98103
IIP,
99-
S, T,
104+
S, T
100105
},
101106
ts, us,
102107
discrete_modified::Bool,
@@ -110,10 +115,11 @@ end
110115
discrete_modified || bool, saved_in_cb || saved_in_cb2
111116
end
112117

113-
@inline function apply_discrete_callback!(integrator::DiffEqBase.AbstractODEIntegrator{
118+
@inline function apply_discrete_callback!(
119+
integrator::DiffEqBase.AbstractODEIntegrator{
114120
AlgType,
115121
IIP,
116-
S, T,
122+
S, T
117123
},
118124
ts, us,
119125
discrete_modified::Bool,
@@ -126,11 +132,12 @@ end
126132
discrete_modified || bool, saved_in_cb || saved_in_cb2
127133
end
128134

129-
@inline function interpolate(integrator::DiffEqBase.AbstractODEIntegrator{
135+
@inline function interpolate(
136+
integrator::DiffEqBase.AbstractODEIntegrator{
130137
AlgType,
131138
IIP,
132139
S,
133-
T,
140+
T
134141
},
135142
t) where {AlgType <: GPUODEAlgorithm, IIP, S, T}
136143
θ = (t - integrator.tprev) / integrator.dt
@@ -142,11 +149,12 @@ end
142149
b7θ * integrator.k7)
143150
end
144151

145-
@inline function _change_t_via_interpolation!(integrator::DiffEqBase.AbstractODEIntegrator{
152+
@inline function _change_t_via_interpolation!(
153+
integrator::DiffEqBase.AbstractODEIntegrator{
146154
AlgType,
147155
IIP,
148156
S,
149-
T,
157+
T
150158
},
151159
t,
152160
modify_save_endpoint::Type{Val{T1}}) where {
@@ -155,7 +163,7 @@ end
155163
IIP,
156164
S,
157165
T,
158-
T1,
166+
T1
159167
}
160168
# Can get rid of an allocation here with a function
161169
# get_tmp_arr(integrator.cache) which gives a pointer to some
@@ -169,11 +177,12 @@ end
169177
#integrator.dt = integrator.t - integrator.tprev
170178
end
171179
end
172-
@inline function DiffEqBase.change_t_via_interpolation!(integrator::DiffEqBase.AbstractODEIntegrator{
180+
@inline function DiffEqBase.change_t_via_interpolation!(
181+
integrator::DiffEqBase.AbstractODEIntegrator{
173182
AlgType,
174183
IIP,
175184
S,
176-
T,
185+
T
177186
},
178187
t,
179188
modify_save_endpoint::Type{Val{T1}} = Val{
@@ -184,12 +193,13 @@ end
184193
IIP,
185194
S,
186195
T,
187-
T1,
196+
T1
188197
}
189198
_change_t_via_interpolation!(integrator, t, modify_save_endpoint)
190199
end
191200

192-
@inline function apply_callback!(integrator::DiffEqBase.AbstractODEIntegrator{AlgType, IIP,
201+
@inline function apply_callback!(
202+
integrator::DiffEqBase.AbstractODEIntegrator{AlgType, IIP,
193203
S, T},
194204
callback::GPUContinuousCallback,
195205
cb_time, prev_sign, event_idx, ts,
@@ -219,7 +229,8 @@ end
219229
true, saved_in_cb
220230
end
221231

222-
@inline function handle_callbacks!(integrator::DiffEqBase.AbstractODEIntegrator{AlgType,
232+
@inline function handle_callbacks!(
233+
integrator::DiffEqBase.AbstractODEIntegrator{AlgType,
223234
IIP, S, T},
224235
ts, us) where {AlgType <: GPUODEAlgorithm, IIP, S, T}
225236
discrete_callbacks = integrator.callback.discrete_callbacks
@@ -232,7 +243,8 @@ end
232243
if !(continuous_callbacks isa Tuple{})
233244
event_occurred = false
234245

235-
time, upcrossing, event_occurred, event_idx, idx, counter = DiffEqBase.find_first_continuous_callback(integrator,
246+
time, upcrossing, event_occurred, event_idx, idx, counter = DiffEqBase.find_first_continuous_callback(
247+
integrator,
236248
continuous_callbacks...)
237249

238250
if event_occurred
@@ -256,16 +268,18 @@ end
256268
return false, saved_in_cb
257269
end
258270

259-
@inline function DiffEqBase.find_callback_time(integrator::DiffEqBase.AbstractODEIntegrator{
271+
@inline function DiffEqBase.find_callback_time(
272+
integrator::DiffEqBase.AbstractODEIntegrator{
260273
AlgType,
261274
IIP,
262275
S,
263-
T,
276+
T
264277
},
265278
callback::DiffEqGPU.GPUContinuousCallback,
266279
counter) where {AlgType <: GPUODEAlgorithm,
267280
IIP, S, T}
268-
event_occurred, interp_index, prev_sign, prev_sign_index, event_idx = DiffEqBase.determine_event_occurance(integrator,
281+
event_occurred, interp_index, prev_sign, prev_sign_index, event_idx = DiffEqBase.determine_event_occurance(
282+
integrator,
269283
callback,
270284
counter)
271285

@@ -321,15 +335,16 @@ end
321335
GPUODEAlgorithm,
322336
IIP,
323337
S,
324-
T,
338+
T
325339
}
326340
return nothing
327341
end
328342

329-
@inline function DiffEqBase.get_condition(integrator::DiffEqBase.AbstractODEIntegrator{
343+
@inline function DiffEqBase.get_condition(
344+
integrator::DiffEqBase.AbstractODEIntegrator{
330345
AlgType,
331346
IIP,
332-
S, T,
347+
S, T
333348
},
334349
callback,
335350
abst) where {AlgType <: GPUODEAlgorithm, IIP, S, T
@@ -345,11 +360,12 @@ end
345360
end
346361

347362
# interp_points = 0 or equivalently nothing
348-
@inline function DiffEqBase.determine_event_occurance(integrator::DiffEqBase.AbstractODEIntegrator{
363+
@inline function DiffEqBase.determine_event_occurance(
364+
integrator::DiffEqBase.AbstractODEIntegrator{
349365
AlgType,
350366
IIP,
351367
S,
352-
T,
368+
T
353369
},
354370
callback::DiffEqGPU.GPUContinuousCallback,
355371
counter) where {

0 commit comments

Comments
 (0)