ClimaComms @time fix

CliMA · Jun 11, 2024 · 6e528b4 · 6e528b4
1 parent 497c99f
commit 6e528b4
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 7 deletions.
diff --git a/experiments/standalone/Bucket/global_bucket_function.jl b/experiments/standalone/Bucket/global_bucket_function.jl
@@ -164,8 +164,13 @@ updatefunc = ClimaLand.make_update_drivers(bucket_atmos, bucket_rad)
 driver_cb = ClimaLand.DriverUpdateCallback(updateat, updatefunc)
 cb = SciMLBase.CallbackSet(driver_cb, saving_cb)
 
-ClimaComms.@time ClimaComms.device() sol =
-    SciMLBase.solve(prob, ode_algo; dt = Δt, saveat = saveat, callback = cb);
+sol = ClimaComms.@time ClimaComms.device() SciMLBase.solve(
+    prob,
+    ode_algo;
+    dt = Δt,
+    saveat = saveat,
+    callback = cb,
+);
 
 # Interpolate to grid
 space = axes(coords.surface)

diff --git a/experiments/standalone/Bucket/global_bucket_staticmap.jl b/experiments/standalone/Bucket/global_bucket_staticmap.jl
@@ -247,8 +247,13 @@ updatefunc = ClimaLand.make_update_drivers(bucket_atmos, bucket_rad)
 driver_cb = ClimaLand.DriverUpdateCallback(updateat, updatefunc)
 cb = SciMLBase.CallbackSet(driver_cb, saving_cb)
 
-ClimaComms.@time ClimaComms.device() sol =
-    SciMLBase.solve(prob, ode_algo; dt = Δt, saveat = saveat, callback = cb);
+sol = ClimaComms.@time ClimaComms.device() SciMLBase.solve(
+    prob,
+    ode_algo;
+    dt = Δt,
+    saveat = saveat,
+    callback = cb,
+);
 
 # Interpolate to grid
 space = axes(coords.surface)

diff --git a/experiments/standalone/Bucket/global_bucket_temporalmap.jl b/experiments/standalone/Bucket/global_bucket_temporalmap.jl
@@ -192,8 +192,13 @@ prob, cb, saveat, saved_values = setup_prob(t0, tf, Δt);
 timestepper = CTS.RK4()
 ode_algo = CTS.ExplicitAlgorithm(timestepper)
 
-ClimaComms.@time ClimaComms.device() sol =
-    SciMLBase.solve(prob, ode_algo; dt = Δt, saveat = saveat, callback = cb)
+sol = ClimaComms.@time ClimaComms.device() SciMLBase.solve(
+    prob,
+    ode_algo;
+    dt = Δt,
+    saveat = saveat,
+    callback = cb,
+)
 if PROFILING
     # Now that we compiled, solve again but collect profiling information
 

diff --git a/experiments/standalone/Soil/richards_runoff.jl b/experiments/standalone/Soil/richards_runoff.jl
@@ -256,7 +256,7 @@ updateat = Array(t0:dt:tf)
 updatefunc = ClimaLand.make_update_drivers(atmos, nothing)
 driver_cb = ClimaLand.DriverUpdateCallback(updateat, updatefunc)
 cb = SciMLBase.CallbackSet(driver_cb, saving_cb)
-@time sol = SciMLBase.solve(prob, ode_algo; dt = dt, saveat = dt, callback = cb)
+sol = @time SciMLBase.solve(prob, ode_algo; dt = dt, saveat = dt, callback = cb)
 
 # Make plots on CPU
 if context.device isa ClimaComms.CPUSingleThreaded