diff --git a/examples/hybrid/driver.jl b/examples/hybrid/driver.jl index 9efbbeb16b8..f1deee1f96a 100644 --- a/examples/hybrid/driver.jl +++ b/examples/hybrid/driver.jl @@ -17,7 +17,7 @@ using Statistics: mean import ClimaAtmos.Parameters as CAP import Thermodynamics as TD import ClimaComms -using OrdinaryDiffEq +using SciMLBase using PrettyTables using DiffEqCallbacks using JLD2 diff --git a/perf/benchmark.jl b/perf/benchmark.jl index 9a23e918418..cef75aeb021 100644 --- a/perf/benchmark.jl +++ b/perf/benchmark.jl @@ -6,9 +6,9 @@ integrator = CA.get_integrator(config) (; parsed_args) = config -import OrdinaryDiffEq as ODE +import SciMLBase import ClimaTimeSteppers as CTS -ODE.step!(integrator) # compile first +SciMLBase.step!(integrator) # compile first (; sol, u, p, dt, t) = integrator @@ -41,7 +41,7 @@ trials["implicit_tendency!"] = get_trial(implicit_fun(integrator), implicit_args trials["remaining_tendency!"] = get_trial(remaining_fun(integrator), remaining_args(integrator), "remaining_tendency!"); trials["additional_tendency!"] = get_trial(CA.additional_tendency!, (X, u, p, t), "additional_tendency!"); trials["hyperdiffusion_tendency!"] = get_trial(CA.hyperdiffusion_tendency!, (X, u, p, t), "hyperdiffusion_tendency!"); -trials["step!"] = get_trial(ODE.step!, (integrator, ), "step!"); +trials["step!"] = get_trial(SciMLBase.step!, (integrator, ), "step!"); #! format: on table_summary = OrderedCollections.OrderedDict() diff --git a/perf/flame.jl b/perf/flame.jl index bfdd11b8ed3..caae2d7c185 100644 --- a/perf/flame.jl +++ b/perf/flame.jl @@ -7,8 +7,8 @@ integrator = CA.get_integrator(config) # The callbacks flame graph is very expensive, so only do 2 steps. @info "running step" -import OrdinaryDiffEq -OrdinaryDiffEq.step!(integrator) # compile first +import SciMLBase +SciMLBase.step!(integrator) # compile first CA.call_all_callbacks!(integrator) # compile callbacks import Profile, ProfileCanvas (; output_dir, job_id) = integrator.p.simulation @@ -18,7 +18,7 @@ mkpath(output_dir) @info "collect profile" Profile.clear() -prof = Profile.@profile OrdinaryDiffEq.step!(integrator) +prof = Profile.@profile SciMLBase.step!(integrator) results = Profile.fetch() Profile.clear() @@ -32,7 +32,7 @@ ProfileCanvas.html_file(joinpath(output_dir, "flame.html"), results) # use new allocation profiler @info "collecting allocations" Profile.Allocs.clear() -Profile.Allocs.@profile sample_rate = 0.01 OrdinaryDiffEq.step!(integrator) +Profile.Allocs.@profile sample_rate = 0.01 SciMLBase.step!(integrator) results = Profile.Allocs.fetch() Profile.Allocs.clear() profile = ProfileCanvas.view_allocs(results) @@ -49,8 +49,8 @@ buffer = occursin("threaded", job_id) ? 1.4 : 1 ## old allocation profiler (TODO: remove this) -allocs = @allocated OrdinaryDiffEq.step!(integrator) -@timev OrdinaryDiffEq.step!(integrator) +allocs = @allocated SciMLBase.step!(integrator) +@timev SciMLBase.step!(integrator) @info "`allocs ($job_id)`: $(allocs)" allocs_limit = Dict() diff --git a/perf/jet.jl b/perf/jet.jl index 90ea9d65d7b..07e4e5f0ce2 100644 --- a/perf/jet.jl +++ b/perf/jet.jl @@ -8,6 +8,6 @@ integrator = CA.get_integrator(config) import JET -import OrdinaryDiffEq -OrdinaryDiffEq.step!(integrator) # Make sure no errors -JET.@test_opt OrdinaryDiffEq.step!(integrator) +import SciMLBase +SciMLBase.step!(integrator) # Make sure no errors +JET.@test_opt SciMLBase.step!(integrator) diff --git a/perf/jet_report_nfailures.jl b/perf/jet_report_nfailures.jl index 97615409642..ec2d6271965 100644 --- a/perf/jet_report_nfailures.jl +++ b/perf/jet_report_nfailures.jl @@ -10,7 +10,7 @@ macro n_failures(ex) ) end -import OrdinaryDiffEq +import SciMLBase import ClimaAtmos as CA n = Dict() Y = integrator.u; @@ -20,7 +20,7 @@ t = integrator.t; Yₜ = similar(Y); ref_Y = similar(Y); #! format: off -n["step!"] = @n_failures OrdinaryDiffEq.step!(integrator); +n["step!"] = @n_failures SciMLBase.step!(integrator); n["limited_tendency!"] = @n_failures CA.limited_tendency!(Yₜ, Y, p, t); n["horizontal_advection_tendency!"] = @n_failures CA.horizontal_advection_tendency!(Yₜ, Y, p, t); n["horizontal_tracer_advection_tendency!"] = @n_failures CA.horizontal_tracer_advection_tendency!(Yₜ, Y, p, t); diff --git a/perf/jet_test_nfailures.jl b/perf/jet_test_nfailures.jl index af75272e08c..4a8c05e6ff0 100644 --- a/perf/jet_test_nfailures.jl +++ b/perf/jet_test_nfailures.jl @@ -6,8 +6,8 @@ integrator = CA.get_integrator(config) import JET -import OrdinaryDiffEq -OrdinaryDiffEq.step!(integrator) # Make sure no errors +import SciMLBase +SciMLBase.step!(integrator) # Make sure no errors # Suggested in: https://github.com/aviatesk/JET.jl/issues/455 macro n_failures(ex) @@ -20,7 +20,7 @@ end using Test @testset "Test N-jet failures" begin - n = @n_failures OrdinaryDiffEq.step!(integrator) + n = @n_failures SciMLBase.step!(integrator) # This test is intended to provide some friction when we # add code to our tendency function that results in degraded # inference. By increasing this counter, we acknowledge that diff --git a/src/callbacks/callback_helpers.jl b/src/callbacks/callback_helpers.jl index 3957fdd80e8..b6bb9ae5e01 100644 --- a/src/callbacks/callback_helpers.jl +++ b/src/callbacks/callback_helpers.jl @@ -1,4 +1,4 @@ -import DiffEqCallbacks +import SciMLBase ##### ##### Callback helpers ##### @@ -7,7 +7,7 @@ function call_every_n_steps(f!, n = 1; skip_first = false, call_at_end = false) previous_step = Ref(0) @assert n ≠ Inf "Adding callback that never gets called!" cb! = AtmosCallback(f!, EveryNSteps(n)) - return ODE.DiscreteCallback( + return SciMLBase.DiscreteCallback( (u, t, integrator) -> (previous_step[] += 1) % n == 0 || (call_at_end && t == integrator.sol.prob.tspan[2]), @@ -31,7 +31,7 @@ function call_every_dt(f!, dt; skip_first = false, call_at_end = false) next_t[] = min(next_t[], t_end) end end - return ODE.DiscreteCallback( + return SciMLBase.DiscreteCallback( (u, t, integrator) -> t >= next_t[], affect!; initialize = (cb, u, t, integrator) -> begin @@ -56,7 +56,7 @@ function callback_from_affect(affect!) end error("Callback not found in $(affect!)") end -function atmos_callbacks(cbs::ODE.CallbackSet) +function atmos_callbacks(cbs::SciMLBase.CallbackSet) all_cbs = [cbs.continuous_callbacks..., cbs.discrete_callbacks...] callback_objs = map(cb -> callback_from_affect(cb.affect!), all_cbs) filter!(x -> !(x isa DiffEqCallbacks.SavedValues), callback_objs) @@ -64,7 +64,7 @@ function atmos_callbacks(cbs::ODE.CallbackSet) end n_measured_calls(integrator) = n_measured_calls(integrator.callback) -n_measured_calls(cbs::ODE.CallbackSet) = +n_measured_calls(cbs::SciMLBase.CallbackSet) = map(x -> x.n_measured_calls, atmos_callbacks(cbs)) n_expected_calls(integrator) = n_expected_calls( @@ -72,12 +72,12 @@ n_expected_calls(integrator) = n_expected_calls( integrator.dt, integrator.sol.prob.tspan, ) -n_expected_calls(cbs::ODE.CallbackSet, dt, tspan) = +n_expected_calls(cbs::SciMLBase.CallbackSet, dt, tspan) = map(x -> n_expected_calls(x, dt, tspan), atmos_callbacks(cbs)) n_steps_per_cycle(integrator) = n_steps_per_cycle(integrator.callback, integrator.dt) -function n_steps_per_cycle(cbs::ODE.CallbackSet, dt) +function n_steps_per_cycle(cbs::SciMLBase.CallbackSet, dt) nspc = n_steps_per_cycle_per_cb(cbs, dt) return isempty(nspc) ? 1 : lcm(nspc) end @@ -85,7 +85,7 @@ end n_steps_per_cycle_per_cb(integrator) = n_steps_per_cycle_per_cb(integrator.callback, integrator.dt) -function n_steps_per_cycle_per_cb(cbs::ODE.CallbackSet, dt) +function n_steps_per_cycle_per_cb(cbs::SciMLBase.CallbackSet, dt) return map(atmos_callbacks(cbs)) do cb cbf = callback_frequency(cb) if cbf isa EveryΔt diff --git a/src/callbacks/callbacks.jl b/src/callbacks/callbacks.jl index 57cd7445cef..285a6d42461 100644 --- a/src/callbacks/callbacks.jl +++ b/src/callbacks/callbacks.jl @@ -6,7 +6,7 @@ import ClimaCore.Fields import ClimaComms import ClimaCore as CC import ClimaCore.Spaces -import OrdinaryDiffEq as ODE +import SciMLBase import ClimaAtmos.Parameters as CAP import DiffEqCallbacks as DEQ import ClimaCore: InputOutput @@ -70,7 +70,7 @@ function turb_conv_affect_filter!(integrator) # paying for an additional `∑tendencies!` call, which is required # to support supplying a continuous representation of the # solution. - ODE.u_modified!(integrator, false) + SciMLBase.u_modified!(integrator, false) end function rrtmgp_model_callback!(integrator) diff --git a/src/initial_conditions/initial_conditions.jl b/src/initial_conditions/initial_conditions.jl index 7941af7befe..3a8cedb4b08 100644 --- a/src/initial_conditions/initial_conditions.jl +++ b/src/initial_conditions/initial_conditions.jl @@ -544,8 +544,8 @@ function hydrostatic_pressure_profile(; dp_dz(p, _, z) = -grav * TD.air_density(thermo_params, ts(p, z, T, θ, q_tot)) - prob = ODE.ODEProblem(dp_dz, p_0, (FT(0), z_max)) - return ODE.solve(prob, ODE.Tsit5(), reltol = 10eps(FT), abstol = 10eps(FT)) + prob = SciMLBase.ODEProblem(dp_dz, p_0, (FT(0), z_max)) + return SciMLBase.solve(prob, ODE.Tsit5(), reltol = 10eps(FT), abstol = 10eps(FT)) end """ diff --git a/src/solver/solve.jl b/src/solver/solve.jl index f8d746509af..ea08f573eed 100644 --- a/src/solver/solve.jl +++ b/src/solver/solve.jl @@ -16,7 +16,7 @@ walltime_in_days(es::EfficiencyStats) = es.walltime * (1 / (24 * 3600)) #=second function timed_solve!(integrator) walltime = @elapsed begin s = @timed_str begin - sol = ODE.solve!(integrator) + sol = SciMLBase.solve!(integrator) end end @info "solve!: $s" @@ -52,7 +52,7 @@ function solve_atmos!(integrator) @info "Running" job_id = p.simulation.job_id output_dir = p.simulation.output_dir tspan comms_ctx = ClimaComms.context(axes(integrator.u.c)) - ODE.step!(integrator) + SciMLBase.step!(integrator) precompile_callbacks(integrator) GC.gc() try @@ -111,7 +111,7 @@ for the flags outlined in a table. """ function benchmark_step!(integrator, Y₀, n_steps = 10) for i in 1:n_steps - ODE.step!(integrator) + SciMLBase.step!(integrator) integrator.u .= Y₀ # temporary hack to simplify performance benchmark. end return nothing @@ -131,7 +131,7 @@ into account. function cycle!(integrator; n_cycles = 1) n_steps = n_steps_per_cycle(integrator) * n_cycles for i in 1:n_steps - ODE.step!(integrator) + SciMLBase.step!(integrator) end return nothing end @@ -150,7 +150,7 @@ Precompiles `step!` and all callbacks in the `integrator`. """ function precompile_atmos(integrator) - B = Base.precompile(ODE.step!, (typeof(integrator),)) + B = Base.precompile(SciMLBase.step!, (typeof(integrator),)) @assert B precompile_callbacks(integrator) return nothing diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl index bd1cb42bb3a..3d88fd5bc0a 100644 --- a/src/solver/type_getters.jl +++ b/src/solver/type_getters.jl @@ -359,7 +359,7 @@ is_ordinary_diffeq_newton(alg_or_tableau) = } is_imex_CTS_algo(::CTS.IMEXAlgorithm) = true -is_imex_CTS_algo(::DiffEqBase.AbstractODEAlgorithm) = false +is_imex_CTS_algo(::SciMLBase.AbstractODEAlgorithm) = false is_implicit(::ODE.OrdinaryDiffEqImplicitAlgorithm) = true is_implicit(::ODE.OrdinaryDiffEqAdaptiveImplicitAlgorithm) = true @@ -367,11 +367,11 @@ is_implicit(ode_algo) = is_imex_CTS_algo(ode_algo) is_rosenbrock(::ODE.Rosenbrock23) = true is_rosenbrock(::ODE.Rosenbrock32) = true -is_rosenbrock(::DiffEqBase.AbstractODEAlgorithm) = false +is_rosenbrock(::SciMLBase.AbstractODEAlgorithm) = false use_transform(ode_algo) = !(is_imex_CTS_algo(ode_algo) || is_rosenbrock(ode_algo)) -additional_integrator_kwargs(::DiffEqBase.AbstractODEAlgorithm) = (; +additional_integrator_kwargs(::SciMLBase.AbstractODEAlgorithm) = (; adaptive = false, progress = isinteractive(), progress_steps = isinteractive() ? 1 : 1000, @@ -382,7 +382,7 @@ additional_integrator_kwargs(::CTS.DistributedODEAlgorithm) = (; # TODO: enable progress bars in ClimaTimeSteppers ) -is_cts_algo(::DiffEqBase.AbstractODEAlgorithm) = false +is_cts_algo(::SciMLBase.AbstractODEAlgorithm) = false is_cts_algo(::CTS.DistributedODEAlgorithm) = true jacobi_flags(::TotalEnergy) = (; ∂ᶜ𝔼ₜ∂ᶠ𝕄_mode = :no_∂ᶜp∂ᶜK)